Skip to content

Commit

Permalink
ENH: cache Blatt–Weisskopf polynomials (#428)
Browse files Browse the repository at this point in the history
* DX: add test for Blatt–Weisskopf polynomials
* MAINT: move form factor tests to separate file
  • Loading branch information
redeboer authored May 22, 2024
1 parent effb0de commit 1ef2142
Show file tree
Hide file tree
Showing 3 changed files with 65 additions and 23 deletions.
37 changes: 26 additions & 11 deletions src/ampform/dynamics/form_factor.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
from __future__ import annotations

from functools import lru_cache
from typing import Any
from typing import Any, Callable

import sympy as sp

Expand Down Expand Up @@ -63,16 +63,31 @@ class BlattWeisskopfSquared(sp.Expr):
_latex_repr_ = R"B_{{{angular_momentum}}}^2\left({z}\right)"

def evaluate(self) -> sp.Expr:
ell = self.angular_momentum
z = sp.Dummy("z", nonnegative=True, real=True)
expr = (
sp.Abs(SphericalHankel1(ell, 1)) ** 2
/ sp.Abs(SphericalHankel1(ell, sp.sqrt(z))) ** 2
/ z
)
if not ell.free_symbols:
expr = expr.doit().simplify()
return expr.xreplace({z: self.z})
z, ell = self.args
if ell.free_symbols:
return _formulate_blatt_weisskopf(ell, z)
expr = _get_polynomial_blatt_weisskopf(ell)(z)
return sp.sympify(expr)


@lru_cache(maxsize=20)
def _get_polynomial_blatt_weisskopf(ell: int | sp.Integer) -> Callable[[Any], Any]:
"""Get the Blatt-Weisskopf factor as a fraction of polynomials.
See https://github.com/ComPWA/ampform/issues/426.
"""
z = sp.Symbol("z", nonnegative=True, real=True)
expr = _formulate_blatt_weisskopf(ell, z)
expr = expr.doit().simplify()
return sp.lambdify(z, expr, "math")


def _formulate_blatt_weisskopf(ell, z) -> sp.Expr:
return (
sp.Abs(SphericalHankel1(ell, 1)) ** 2
/ sp.Abs(SphericalHankel1(ell, sp.sqrt(z))) ** 2
/ z
)


@unevaluated
Expand Down
12 changes: 0 additions & 12 deletions tests/dynamics/test_dynamics.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,25 +12,13 @@
PhaseSpaceFactorSWave,
relativistic_breit_wigner_with_ff,
)
from ampform.dynamics.form_factor import BlattWeisskopfSquared

if TYPE_CHECKING:
from qrules import ParticleCollection

from ampform.helicity import HelicityModel


class TestBlattWeisskopfSquared:
def test_factorials(self):
z = sp.Symbol("z")
angular_momentum = sp.Symbol("L", integer=True)
form_factor = BlattWeisskopfSquared(z, angular_momentum)
form_factor_9 = form_factor.subs(angular_momentum, 8).evaluate()
factor, z_power, _ = form_factor_9.args
assert factor == 4392846440677
assert z_power == z**8


class TestEnergyDependentWidth:
@staticmethod
def test_init():
Expand Down
39 changes: 39 additions & 0 deletions tests/dynamics/test_form_factor.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,39 @@
import pytest
import sympy as sp

from ampform.dynamics.form_factor import _get_polynomial_blatt_weisskopf

z = sp.Symbol("z", nonnegative=True, real=True)


@pytest.mark.parametrize(
("ell", "expected"),
[
(0, 1),
(1, 2 * z / (z + 1)),
(2, 13 * z**2 / (z**2 + 3 * z + 9)),
(3, 277 * z**3 / (z**3 + 6 * z**2 + 45 * z + 225)),
(4, 12746 * z**4 / (z**4 + 10 * z**3 + 135 * z**2 + 1575 * z + 11025)),
(
10,
451873017324894386
* z**10
/ (
z**10
+ 55 * z**9
+ 4455 * z**8
+ 386100 * z**7
+ 33108075 * z**6
+ 2681754075 * z**5
+ 196661965500 * z**4
+ 12417798393000 * z**3
+ 628651043645625 * z**2
+ 22561587455281875 * z
+ 428670161650355625
),
),
],
)
def test_get_polynomial_blatt_weisskopf(ell: int, expected: sp.Expr):
expr = _get_polynomial_blatt_weisskopf(ell)(z)
assert expr == expected

0 comments on commit 1ef2142

Please sign in to comment.