Skip to content

Commit 6142f93

Browse files
redeboerredeboer
committed
FEAT: implement higher angular momentum in widget
1 parent a6b9deb commit 6142f93

File tree

2 files changed

+127
-19
lines changed

2 files changed

+127
-19
lines changed

‎docs/usage/dynamics/integration-algorithms.ipynb‎

Lines changed: 125 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -50,13 +50,19 @@
5050
"import numpy as np\n",
5151
"import numpy.typing as npt\n",
5252
"import quadax\n",
53+
"import sympy as sp\n",
5354
"from ipympl.backend_nbagg import Canvas\n",
54-
"from IPython.display import SVG, display\n",
55+
"from IPython.display import SVG, Math, display\n",
5556
"from matplotlib.axes import Axes\n",
5657
"from matplotlib.collections import LineCollection, QuadMesh\n",
5758
"from matplotlib.lines import Line2D\n",
5859
"from scipy.integrate import quad_vec\n",
5960
"\n",
61+
"from ampform.dynamics.form_factor import BlattWeisskopfSquared, FormFactor\n",
62+
"from ampform.dynamics.phasespace import ChewMandelstamIntegral, PhaseSpaceFactor\n",
63+
"from ampform.io import aslatex\n",
64+
"from ampform.sympy import UnevaluatableIntegral\n",
65+
"\n",
6066
"# cspell:disable-next-line\n",
6167
"Algorithm = Literal[\"quadcc\", \"quadgk\", \"quadts\", \"romberg\", \"rombergts\", \"quad_vec\"]\n",
6268
"jax.config.update(\"jax_enable_x64\", True)\n",
@@ -68,6 +74,53 @@
6874
" canvas.toolbar_visible = False"
6975
]
7076
},
77+
{
78+
"cell_type": "markdown",
79+
"metadata": {},
80+
"source": [
81+
"```{autolink-skip}\n",
82+
"\n",
83+
"```"
84+
]
85+
},
86+
{
87+
"cell_type": "code",
88+
"execution_count": null,
89+
"metadata": {
90+
"tags": [
91+
"remove-cell"
92+
]
93+
},
94+
"outputs": [],
95+
"source": [
96+
"UnevaluatableIntegral.dummify = False"
97+
]
98+
},
99+
{
100+
"cell_type": "code",
101+
"execution_count": null,
102+
"metadata": {
103+
"tags": [
104+
"hide-input",
105+
"full-width"
106+
]
107+
},
108+
"outputs": [],
109+
"source": [
110+
"s, m1, m2, z = sp.symbols(\"s m1 m2 z\", nonnegative=True)\n",
111+
"ell = sp.Symbol(\"ell\", integer=True, nonnegative=True)\n",
112+
"cm = ChewMandelstamIntegral(s, m1, m2, ell)\n",
113+
"ff = FormFactor(s, m1, m2, ell)\n",
114+
"rho = PhaseSpaceFactor(s, m1, m2)\n",
115+
"bl = BlattWeisskopfSquared(z, ell)\n",
116+
"max_ell = 5\n",
117+
"src = aslatex({\n",
118+
" **{e: e.doit(deep=False) for e in (cm, ff, rho)},\n",
119+
" **{bl.subs(ell, i): bl.subs(ell, i).doit() for i in range(max_ell + 1)},\n",
120+
"})\n",
121+
"Math(src)"
122+
]
123+
},
71124
{
72125
"cell_type": "code",
73126
"execution_count": null,
@@ -88,36 +141,91 @@
88141
" s: npt.NDArray[np.float64],\n",
89142
" m1: float,\n",
90143
" m2: float,\n",
144+
" ell: int = 0,\n",
91145
" start_offset: float = 0,\n",
92146
" algorithm: Callable = quadax.quadcc,\n",
93147
" **configuration,\n",
94148
"):\n",
95149
" s_thr = (m1 + m2) ** 2\n",
96150
" if algorithm is quad_vec:\n",
97151
" integral, _ = algorithm(\n",
98-
" partial(integrand, s=s, m1=m1, m2=m2),\n",
152+
" partial(integrand, s=s, m1=m1, m2=m2, ell=ell),\n",
99153
" s_thr + start_offset,\n",
100154
" np.inf,\n",
101155
" **configuration,\n",
102156
" )\n",
103157
" else:\n",
104158
" integral, _ = algorithm(\n",
105-
" jax.tree_util.Partial(integrand, s=s, m1=m1, m2=m2),\n",
159+
" jax.tree_util.Partial(integrand, s=s, m1=m1, m2=m2, ell=ell),\n",
106160
" interval=[s_thr + start_offset, jnp.inf],\n",
107161
" **configuration,\n",
108162
" )\n",
109163
" return (s - s_thr) * integral / jnp.pi\n",
110164
"\n",
111165
"\n",
112166
"@jax.jit\n",
113-
"def integrand(sp, s, m1, m2):\n",
167+
"def integrand(sp, s, m1, m2, ell):\n",
114168
" s_thr = (m1 + m2) ** 2\n",
115-
" return rho(sp, m1, m2) / ((sp - s_thr) * (sp - s))\n",
169+
" return rho_func(sp, m1, m2) * n2(s, m1, m2, ell) / ((sp - s_thr) * (sp - s))\n",
116170
"\n",
117171
"\n",
118172
"@jax.jit\n",
119-
"def rho(s, m1, m2):\n",
120-
" return jnp.sqrt((s - (m1 - m2) ** 2) * (s - (m1 + m2) ** 2)) / s"
173+
"def rho_func(s, m1, m2):\n",
174+
" return jnp.sqrt(s - (m1 - m2) ** 2) * jnp.sqrt(s - (m1 + m2) ** 2) / s\n",
175+
"\n",
176+
"\n",
177+
"def n2(s, m1, m2, ell):\n",
178+
" return blatt_weisskopf_squared(q(s, m1, m2), ell)\n",
179+
"\n",
180+
"\n",
181+
"def blatt_weisskopf_squared(z, ell):\n",
182+
" return jnp.select(\n",
183+
" [ell == 0, ell == 1, ell == 2, ell == 3, ell == 4, ell == 5],\n",
184+
" [\n",
185+
" 1,\n",
186+
" 2 * z / (z + 1),\n",
187+
" 13 * z**2 / (z**2 + 3 * z + 9),\n",
188+
" 277 * z**3 / (z**3 + 6 * z**2 + 45 * z + 225),\n",
189+
" 12746 * z**4 / (z**4 + 10 * z**3 + 135 * z**2 + 1575 * z + 11025),\n",
190+
" 998881\n",
191+
" * z**5\n",
192+
" / (z**5 + 15 * z**4 + 315 * z**3 + 6300 * z**2 + 99225 * z + 893025),\n",
193+
" ],\n",
194+
" default=jnp.nan,\n",
195+
" )\n",
196+
"\n",
197+
"\n",
198+
"@jax.jit\n",
199+
"def q(s, m1, m2):\n",
200+
" return (\n",
201+
" jnp.sqrt(s - (m1 - m2) ** 2) * jnp.sqrt(s - (m1 + m2) ** 2) / (2 * jnp.sqrt(s))\n",
202+
" )"
203+
]
204+
},
205+
{
206+
"cell_type": "markdown",
207+
"metadata": {},
208+
"source": [
209+
"In the case of S-waves, we can compare the result of the integration to the analytical result.\n",
210+
"\n",
211+
"Note that in this case, we define the break-up momentum in the same way as in {class}`.BreakupMomentum`, because in the widget below, we only evaluate this function along the real axis, where the cut structure doesn't matter and we prefer performance (see {ref}`usage/dynamics/analytic-continuation:Numerical precision and performance`)."
212+
]
213+
},
214+
{
215+
"cell_type": "code",
216+
"execution_count": null,
217+
"metadata": {
218+
"tags": [
219+
"hide-input",
220+
"full-width"
221+
]
222+
},
223+
"outputs": [],
224+
"source": [
225+
"from ampform.dynamics.phasespace import ChewMandelstamSWave\n",
226+
"\n",
227+
"CM0 = ChewMandelstamSWave(s, m1, m2)\n",
228+
"Math(aslatex({CM0: CM0.doit(deep=False)}))"
121229
]
122230
},
123231
{
@@ -142,12 +250,7 @@
142250
" (2 * q(s, m1, m2) / jnp.sqrt(s))\n",
143251
" * jnp.log((m1**2 + m2**2 - s + 2 * q(s, m1, m2) * jnp.sqrt(s)) / (2 * m1 * m2))\n",
144252
" - (m1**2 - m2**2) * (1 / s - 1 / (m1 + m2) ** 2) * jnp.log(m1 / m2)\n",
145-
" )\n",
146-
"\n",
147-
"\n",
148-
"@jax.jit\n",
149-
"def q(s, m1, m2):\n",
150-
" return jnp.sqrt((s - (m1 - m2) ** 2) * (s - (m1 + m2) ** 2)) / (2 * jnp.sqrt(s))"
253+
" )"
151254
]
152255
},
153256
{
@@ -177,6 +280,7 @@
177280
" ),\n",
178281
" m1=w.FloatSlider(value=0.13, min=0.0, max=2.0, step=0.01, description=\"m₁\", **cont),\n",
179282
" m2=w.FloatSlider(value=0.98, min=0.0, max=2.0, step=0.01, description=\"m₂\", **cont),\n",
283+
" ell=w.IntSlider(value=0, min=0, max=5, description=\"ℓ\", **cont),\n",
180284
" y_lim=w.FloatRangeSlider(\n",
181285
" description=\"y range\",\n",
182286
" min=-5,\n",
@@ -314,8 +418,8 @@
314418
" tabs := w.Tab([\n",
315419
" w.HBox([\n",
316420
" physics_sliders[\"projection\"],\n",
317-
" w.VBox(list(physics_sliders.values())[1:4]),\n",
318-
" w.VBox(list(physics_sliders.values())[4:]),\n",
421+
" w.VBox(list(physics_sliders.values())[1:5]),\n",
422+
" w.VBox(list(physics_sliders.values())[5:]),\n",
319423
" ]),\n",
320424
" w.HBox([\n",
321425
" algorithm_sliders[\"algorithm_name\"],\n",
@@ -381,6 +485,7 @@
381485
" projection: Literal[\"real\", \"imag\", \"abs\"],\n",
382486
" m1: float,\n",
383487
" m2: float,\n",
488+
" ell: int,\n",
384489
" epsilon: float,\n",
385490
" start_offset: float,\n",
386491
" z_max: float,\n",
@@ -400,15 +505,17 @@
400505
" z_ana = sigma0(s, m1, m2)\n",
401506
" alg_kwargs = get_algorithm_options(algorithm_name, **alg_kwargs)\n",
402507
" start_time = time.perf_counter()\n",
403-
" z_num = integrate_numerically(s, m1, m2, start_offset, algorithm, **alg_kwargs)\n",
508+
" z_num = integrate_numerically(\n",
509+
" s, m1, m2, ell, start_offset, algorithm, **alg_kwargs\n",
510+
" )\n",
404511
" z_num.block_until_ready()\n",
405512
" end_time = time.perf_counter()\n",
406513
" z = z_exact, z_ana, z_num\n",
407514
" duration = end_time - start_time\n",
408515
" timer_box.value = f\"Computation time: <b>{format_time(duration)}</b> for {resolution:,} points\"\n",
409516
" if np.all(np.isnan(z_num)):\n",
410517
" timer_box.value += \" (<font color='red'>all values are NaN</font>)\"\n",
411-
" Z = rho(S, m1, m2)\n",
518+
" Z = rho_func(S, m1, m2) * n2(S, m1, m2, ell)\n",
412519
" if projection == \"abs\":\n",
413520
" Z = jnp.abs(Z)\n",
414521
" else:\n",
@@ -541,7 +648,7 @@
541648
" ax.spines[\"right\"].set_visible(False)\n",
542649
" ax.spines[\"top\"].set_visible(False)\n",
543650
"ax1.spines[\"bottom\"].set_visible(False)\n",
544-
"ax1.set_title(R\"$\\rho(s)$\")\n",
651+
"ax1.set_title(R\"$\\rho(s) \\, n_\\ell^2(s)$\")\n",
545652
"ax1.set_ylabel(\"Im $s$\")\n",
546653
"ax2.set_ylabel(R\"Im $\\Sigma_0(s)$\")\n",
547654
"ax3.set_ylabel(R\"Re $\\Sigma_0(s)$\")\n",

‎pyproject.toml‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -394,7 +394,7 @@ docstring-code-format = true
394394
line-ending = "lf"
395395

396396
[tool.ruff.lint]
397-
allowed-confusables = ["µ", "–"]
397+
allowed-confusables = ["–"]
398398
ignore = [
399399
"ANN",
400400
"ARG00",
@@ -478,6 +478,7 @@ split-on-trailing-comma = false
478478
"TC00",
479479
]
480480
"**/docs/usage/dynamics.ipynb" = ["FURB118", "RUF027"]
481+
"**/docs/usage/dynamics/integration-algorithms.ipynb" = ["RUF001"]
481482
"**/docs/usage/dynamics/riemann-sheets.ipynb" = ["E741", "RUF027"]
482483
"**/docs/usage/sympy.ipynb" = ["E731"]
483484
"benchmarks/*" = [

0 commit comments

Comments
 (0)