Skip to content

Commit b4e6574

Browse files
redeboerredeboer
committed
FEAT: implement higher angular momentum in widget
1 parent ac9818a commit b4e6574

File tree

2 files changed

+106
-19
lines changed

2 files changed

+106
-19
lines changed

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

Lines changed: 104 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,32 @@
6874
" canvas.toolbar_visible = False"
6975
]
7076
},
77+
{
78+
"cell_type": "code",
79+
"execution_count": null,
80+
"metadata": {
81+
"tags": [
82+
"hide-input",
83+
"full-width"
84+
]
85+
},
86+
"outputs": [],
87+
"source": [
88+
"UnevaluatableIntegral.dummify = False\n",
89+
"s, m1, m2, z = sp.symbols(\"s m1 m2 z\", nonnegative=True)\n",
90+
"ell = sp.Symbol(\"ell\", integer=True, nonnegative=True)\n",
91+
"cm = ChewMandelstamIntegral(s, m1, m2, ell)\n",
92+
"ff = FormFactor(s, m1, m2, ell)\n",
93+
"rho = PhaseSpaceFactor(s, m1, m2)\n",
94+
"bl = BlattWeisskopfSquared(z, ell)\n",
95+
"max_ell = 5\n",
96+
"src = aslatex({\n",
97+
" **{e: e.doit(deep=False) for e in (cm, ff, rho)},\n",
98+
" **{bl.subs(ell, i): bl.subs(ell, i).doit() for i in range(max_ell + 1)},\n",
99+
"})\n",
100+
"display(Math(src))"
101+
]
102+
},
71103
{
72104
"cell_type": "code",
73105
"execution_count": null,
@@ -88,36 +120,91 @@
88120
" s: npt.NDArray[np.float64],\n",
89121
" m1: float,\n",
90122
" m2: float,\n",
123+
" ell: int = 0,\n",
91124
" start_offset: float = 0,\n",
92125
" algorithm: Callable = quadax.quadcc,\n",
93126
" **configuration,\n",
94127
"):\n",
95128
" s_thr = (m1 + m2) ** 2\n",
96129
" if algorithm is quad_vec:\n",
97130
" integral, _ = algorithm(\n",
98-
" partial(integrand, s=s, m1=m1, m2=m2),\n",
131+
" partial(integrand, s=s, m1=m1, m2=m2, ell=ell),\n",
99132
" s_thr + start_offset,\n",
100133
" np.inf,\n",
101134
" **configuration,\n",
102135
" )\n",
103136
" else:\n",
104137
" integral, _ = algorithm(\n",
105-
" jax.tree_util.Partial(integrand, s=s, m1=m1, m2=m2),\n",
138+
" jax.tree_util.Partial(integrand, s=s, m1=m1, m2=m2, ell=ell),\n",
106139
" interval=[s_thr + start_offset, jnp.inf],\n",
107140
" **configuration,\n",
108141
" )\n",
109142
" return (s - s_thr) * integral / jnp.pi\n",
110143
"\n",
111144
"\n",
112145
"@jax.jit\n",
113-
"def integrand(sp, s, m1, m2):\n",
146+
"def integrand(sp, s, m1, m2, ell):\n",
114147
" s_thr = (m1 + m2) ** 2\n",
115-
" return rho(sp, m1, m2) / ((sp - s_thr) * (sp - s))\n",
148+
" return rho_func(sp, m1, m2) * n2(s, m1, m2, ell) / ((sp - s_thr) * (sp - s))\n",
149+
"\n",
150+
"\n",
151+
"@jax.jit\n",
152+
"def rho_func(s, m1, m2):\n",
153+
" return jnp.sqrt(s - (m1 - m2) ** 2) * jnp.sqrt(s - (m1 + m2) ** 2) / s\n",
154+
"\n",
155+
"\n",
156+
"def n2(s, m1, m2, ell):\n",
157+
" return blatt_weisskopf_squared(q(s, m1, m2), ell)\n",
158+
"\n",
159+
"\n",
160+
"def blatt_weisskopf_squared(z, ell):\n",
161+
" return jnp.select(\n",
162+
" [ell == 0, ell == 1, ell == 2, ell == 3, ell == 4, ell == 5],\n",
163+
" [\n",
164+
" 1,\n",
165+
" 2 * z / (z + 1),\n",
166+
" 13 * z**2 / (z**2 + 3 * z + 9),\n",
167+
" 277 * z**3 / (z**3 + 6 * z**2 + 45 * z + 225),\n",
168+
" 12746 * z**4 / (z**4 + 10 * z**3 + 135 * z**2 + 1575 * z + 11025),\n",
169+
" 998881\n",
170+
" * z**5\n",
171+
" / (z**5 + 15 * z**4 + 315 * z**3 + 6300 * z**2 + 99225 * z + 893025),\n",
172+
" ],\n",
173+
" default=jnp.nan,\n",
174+
" )\n",
116175
"\n",
117176
"\n",
118177
"@jax.jit\n",
119-
"def rho(s, m1, m2):\n",
120-
" return jnp.sqrt((s - (m1 - m2) ** 2) * (s - (m1 + m2) ** 2)) / s"
178+
"def q(s, m1, m2):\n",
179+
" return (\n",
180+
" jnp.sqrt(s - (m1 - m2) ** 2) * jnp.sqrt(s - (m1 + m2) ** 2) / (2 * jnp.sqrt(s))\n",
181+
" )"
182+
]
183+
},
184+
{
185+
"cell_type": "markdown",
186+
"metadata": {},
187+
"source": [
188+
"In the case of S-waves, we can compare the result of the integration to the analytical result.\n",
189+
"\n",
190+
"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`)."
191+
]
192+
},
193+
{
194+
"cell_type": "code",
195+
"execution_count": null,
196+
"metadata": {
197+
"tags": [
198+
"hide-input",
199+
"full-width"
200+
]
201+
},
202+
"outputs": [],
203+
"source": [
204+
"from ampform.dynamics.phasespace import ChewMandelstamSWave\n",
205+
"\n",
206+
"CM0 = ChewMandelstamSWave(s, m1, m2)\n",
207+
"Math(aslatex({CM0: CM0.doit(deep=False)}))"
121208
]
122209
},
123210
{
@@ -142,12 +229,7 @@
142229
" (2 * q(s, m1, m2) / jnp.sqrt(s))\n",
143230
" * jnp.log((m1**2 + m2**2 - s + 2 * q(s, m1, m2) * jnp.sqrt(s)) / (2 * m1 * m2))\n",
144231
" - (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))"
232+
" )"
151233
]
152234
},
153235
{
@@ -177,6 +259,7 @@
177259
" ),\n",
178260
" m1=w.FloatSlider(value=0.13, min=0.0, max=2.0, step=0.01, description=\"m₁\", **cont),\n",
179261
" m2=w.FloatSlider(value=0.98, min=0.0, max=2.0, step=0.01, description=\"m₂\", **cont),\n",
262+
" ell=w.IntSlider(value=0, min=0, max=5, description=\"ℓ\", **cont),\n",
180263
" y_lim=w.FloatRangeSlider(\n",
181264
" description=\"y range\",\n",
182265
" min=-5,\n",
@@ -313,8 +396,8 @@
313396
" tabs := w.Tab([\n",
314397
" w.HBox([\n",
315398
" physics_sliders[\"projection\"],\n",
316-
" w.VBox(list(physics_sliders.values())[1:4]),\n",
317-
" w.VBox(list(physics_sliders.values())[4:]),\n",
399+
" w.VBox(list(physics_sliders.values())[1:5]),\n",
400+
" w.VBox(list(physics_sliders.values())[5:]),\n",
318401
" ]),\n",
319402
" w.HBox([\n",
320403
" algorithm_sliders[\"algorithm_name\"],\n",
@@ -380,6 +463,7 @@
380463
" projection: Literal[\"real\", \"imag\", \"abs\"],\n",
381464
" m1: float,\n",
382465
" m2: float,\n",
466+
" ell: int,\n",
383467
" epsilon: float,\n",
384468
" start_offset: float,\n",
385469
" z_max: float,\n",
@@ -399,15 +483,17 @@
399483
" z_ana = sigma0(s, m1, m2)\n",
400484
" alg_kwargs = get_algorithm_options(algorithm_name, **alg_kwargs)\n",
401485
" start_time = time.perf_counter()\n",
402-
" z_num = integrate_numerically(s, m1, m2, start_offset, algorithm, **alg_kwargs)\n",
486+
" z_num = integrate_numerically(\n",
487+
" s, m1, m2, ell, start_offset, algorithm, **alg_kwargs\n",
488+
" )\n",
403489
" z_num.block_until_ready()\n",
404490
" end_time = time.perf_counter()\n",
405491
" z = z_exact, z_ana, z_num\n",
406492
" duration = end_time - start_time\n",
407493
" timer_box.value = f\"Computation time: <b>{format_time(duration)}</b> for {resolution:,} points\"\n",
408494
" if np.all(np.isnan(z_num)):\n",
409495
" timer_box.value += \" (<font color='red'>all values are NaN</font>)\"\n",
410-
" Z = rho(S, m1, m2)\n",
496+
" Z = rho_func(S, m1, m2) * n2(S, m1, m2, ell)\n",
411497
" if projection == \"abs\":\n",
412498
" Z = jnp.abs(Z)\n",
413499
" else:\n",
@@ -540,7 +626,7 @@
540626
" ax.spines[\"right\"].set_visible(False)\n",
541627
" ax.spines[\"top\"].set_visible(False)\n",
542628
"ax1.spines[\"bottom\"].set_visible(False)\n",
543-
"ax1.set_title(R\"$\\rho(s)$\")\n",
629+
"ax1.set_title(R\"$\\rho(s) \\, n_\\ell^2(s)$\")\n",
544630
"ax1.set_ylabel(\"Im $s$\")\n",
545631
"ax2.set_ylabel(R\"Im $\\Sigma_0(s)$\")\n",
546632
"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)