ComPWA · redeboer · Aug 9, 2021 · Aug 9, 2021 · Aug 9, 2021 · Aug 9, 2021
diff --git a/docs/usage/dynamics/k-matrix.ipynb b/docs/usage/dynamics/k-matrix.ipynb
@@ -60,11 +60,11 @@
    "source": [
     "While {mod}`ampform` does not yet provide a generic way to formulate an amplitude model with $\\boldsymbol{K}$-matrix dynamics, the (experimental) {mod}`.kmatrix` module makes it fairly simple to produce a symbolic expression for a parameterized $\\boldsymbol{K}$-matrix with an arbitrary number of poles and channels and play around with it interactively. For more info on the $\\boldsymbol{K}$-matrix, see the classic paper by Chung {cite}`chungPartialWaveAnalysis1995`, {pdg-review}`Resonances`, or this instructive presentation {cite}`meyerMatrixTutorial2008`.\n",
     "\n",
-    "Section {ref}`usage/dynamics/k-matrix:Physics` provides a summary of {cite}`chungPartialWaveAnalysis1995`, so that the equations can be referenced to in the {mod}`.kmatrix` module. It also points out some subtleties and deviations.\n",
+    "Section {ref}`usage/dynamics/k-matrix:Physics` summarizes {cite}`chungPartialWaveAnalysis1995`, so that the {mod}`.kmatrix` module can reference to the equations. It also points out some subtleties and deviations.\n",
     "\n",
     ":::{note}\n",
     "\n",
-    "The $\\boldsymbol{K}$-matrix approach was worked originally worked out in {doc}`compwa-org:report/005`, {doc}`compwa-org:report/009`, and {doc}`compwa-org:report/010`. Those reports contained a few mistakes, which have been (and are being) addressed here.\n",
+    "The $\\boldsymbol{K}$-matrix approach was originally worked worked out in {doc}`compwa-org:report/005`, {doc}`compwa-org:report/009`, and {doc}`compwa-org:report/010`. Those reports contained a few mistakes, which have been addressed here.\n",
     "\n",
     ":::"
    ]
@@ -99,6 +99,7 @@
     "\n",
     "import symplot\n",
     "from ampform.dynamics import (\n",
+    "    BlattWeisskopfSquared,\n",
     "    PhaseSpaceFactor,\n",
     "    kmatrix,\n",
     "    relativistic_breit_wigner,\n",
@@ -107,7 +108,9 @@
     "logging.basicConfig()\n",
     "logging.getLogger().setLevel(logging.ERROR)\n",
     "\n",
-    "warnings.filterwarnings(\"ignore\")"
+    "warnings.filterwarnings(\"ignore\")\n",
+    "\n",
+    "BlattWeisskopfSquared.max_angular_momentum = 3"
    ]
   },
   {
@@ -194,7 +197,7 @@
     "\\frac{d\\sigma}{d\\Omega} = \\left|A(\\Omega)\\right|^2.\n",
     "$$ (differential cross section)\n",
     "\n",
-    "We can now further express $A$ in terms of **partial wave amplitudes** by splitting it up in terms of its angular momentum components:[^spin-formalisms]\n",
+    "We can now further express $A$ in terms of **partial wave amplitudes** by expanding it in terms of its angular momentum components:[^spin-formalisms]\n",
     "\n",
     "```{margin}\n",
     "{cite}`chungPartialWaveAnalysis1995` Eq. (2)\n",
@@ -208,7 +211,7 @@
     "\n",
     "[^spin-formalisms]: Further subtleties arise when taking spin into account, especially for sequential decays. This is where {doc}`spin formalisms </usage/formalism>` come in.\n",
     "\n",
-    "The above sketch is just with one channel in mind. The same holds true though for a number of channels $n$, with the only difference being that the $T$ operator becomes a $\\boldsymbol{T}$-matrix of rank $n$."
+    "The above sketch is just with one channel in mind. The same holds true though for a number of channels $n$, with the only difference being that the $T$ operator becomes an $n\\times n$ $\\boldsymbol{T}$-matrix."
    ]
   },
   {
@@ -1033,7 +1036,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Note that the $F$-vector approach introduces two additional coefficients $\\beta$. These can constants can be complex and can introduce phase differences."
+    "Note that the $F$-vector approach introduces additional $\\beta$-coefficients. These can constants can be complex and can introduce phase differences form the production process."
    ]
   },
   {
@@ -1488,8 +1491,11 @@
     "    # Set slider values and ranges\n",
     "    m0_values = np.linspace(x_min, x_max, num=n_poles + 2)\n",
     "    m0_values = m0_values[1:-1]\n",
+    "    max_angular_momentum = BlattWeisskopfSquared.max_angular_momentum\n",
+    "    if max_angular_momentum is None:\n",
+    "        max_angular_momentum = 8\n",
     "    if \"L\" in sliders:\n",
-    "        sliders.set_ranges(L=(0, 8))\n",
+    "        sliders.set_ranges(L=(0, max_angular_momentum))\n",
     "    sliders.set_ranges(i=(0, n_channels - 1))\n",
     "    for R in range(1, n_poles + 1):\n",
     "        for i in range(n_channels):\n",
@@ -1775,7 +1781,6 @@
    },
    "outputs": [],
    "source": [
-    "%%time\n",
     "plot(\n",
     "    kmatrix.RelativisticKMatrix,\n",
     "    n_poles=2,\n",
@@ -1822,7 +1827,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.10"
+   "version": "3.8.11"
   }
  },
  "nbformat": 4,

diff --git a/pytest.ini b/pytest.ini
@@ -17,6 +17,7 @@ filterwarnings =
 nb_diff_ignore =
     /cells/*/execution_count
     /cells/*/outputs
+    /metadata/language_info/version
     /metadata/widgets
 norecursedirs =
     _build

diff --git a/src/ampform/dynamics/__init__.py b/src/ampform/dynamics/__init__.py
@@ -8,7 +8,7 @@
 """
 
 import re
-from typing import Any, List, Optional, Sequence, Union
+from typing import Any, Dict, List, Optional, Sequence, Union
 
 import sympy as sp
 from sympy.printing.conventions import split_super_sub
@@ -55,6 +55,12 @@ class BlattWeisskopfSquared(UnevaluatedExpression):
     See also :ref:`usage/dynamics:Form factor`.
     """
     is_commutative = True
+    max_angular_momentum: Optional[int] = None
+    """Limit the maximum allowed angular momentum :math:`L`.
+
+    This improves performance when :math:`L` is a `~sympy.core.symbol.Symbol`
+    and you are note interested in higher angular momenta.
+    """
 
     def __new__(
         cls, angular_momentum: sp.Symbol, z: sp.Symbol, **hints: Any
@@ -63,96 +69,77 @@ def __new__(
 
     def evaluate(self) -> sp.Expr:
         angular_momentum, z = self.args
-        return sp.Piecewise(
-            (
-                1,
-                sp.Eq(angular_momentum, 0),
-            ),
-            (
-                2 * z / (z + 1),
-                sp.Eq(angular_momentum, 1),
-            ),
-            (
-                13 * z ** 2 / ((z - 3) * (z - 3) + 9 * z),
-                sp.Eq(angular_momentum, 2),
-            ),
-            (
-                (
-                    277
-                    * z ** 3
-                    / (z * (z - 15) * (z - 15) + 9 * (2 * z - 5) * (2 * z - 5))
-                ),
-                sp.Eq(angular_momentum, 3),
+        cases: Dict[int, sp.Expr] = {
+            0: 1,
+            1: 2 * z / (z + 1),
+            2: 13 * z ** 2 / ((z - 3) * (z - 3) + 9 * z),
+            3: (
+                277
+                * z ** 3
+                / (z * (z - 15) * (z - 15) + 9 * (2 * z - 5) * (2 * z - 5))
             ),
-            (
-                (
-                    12746
-                    * z ** 4
-                    / (
-                        (z ** 2 - 45 * z + 105) * (z ** 2 - 45 * z + 105)
-                        + 25 * z * (2 * z - 21) * (2 * z - 21)
-                    )
-                ),
-                sp.Eq(angular_momentum, 4),
-            ),
-            (
-                998881
-                * z ** 5
+            4: (
+                12746
+                * z ** 4
                 / (
-                    z ** 5
-                    + 15 * z ** 4
-                    + 315 * z ** 3
-                    + 6300 * z ** 2
-                    + 99225 * z
-                    + 893025
-                ),
-                sp.Eq(angular_momentum, 5),
+                    (z ** 2 - 45 * z + 105) * (z ** 2 - 45 * z + 105)
+                    + 25 * z * (2 * z - 21) * (2 * z - 21)
+                )
             ),
-            (
-                118394977
-                * z ** 6
-                / (
-                    z ** 6
-                    + 21 * z ** 5
-                    + 630 * z ** 4
-                    + 18900 * z ** 3
-                    + 496125 * z ** 2
-                    + 9823275 * z
-                    + 108056025
-                ),
-                sp.Eq(angular_momentum, 6),
+            5: 998881
+            * z ** 5
+            / (
+                z ** 5
+                + 15 * z ** 4
+                + 315 * z ** 3
+                + 6300 * z ** 2
+                + 99225 * z
+                + 893025
             ),
-            (
-                19727003738
-                * z ** 7
-                / (
-                    z ** 7
-                    + 28 * z ** 6
-                    + 1134 * z ** 5
-                    + 47250 * z ** 4
-                    + 1819125 * z ** 3
-                    + 58939650 * z ** 2
-                    + 1404728325 * z
-                    + 18261468225
-                ),
-                sp.Eq(angular_momentum, 7),
+            6: 118394977
+            * z ** 6
+            / (
+                z ** 6
+                + 21 * z ** 5
+                + 630 * z ** 4
+                + 18900 * z ** 3
+                + 496125 * z ** 2
+                + 9823275 * z
+                + 108056025
             ),
-            (
-                4392846440677
-                * z ** 8
-                / (
-                    z ** 8
-                    + 36 * z ** 7
-                    + 1890 * z ** 6
-                    + 103950 * z ** 5
-                    + 5457375 * z ** 4
-                    + 255405150 * z ** 3
-                    + 9833098275 * z ** 2
-                    + 273922023375 * z
-                    + 4108830350625
-                ),
-                sp.Eq(angular_momentum, 8),
+            7: 19727003738
+            * z ** 7
+            / (
+                z ** 7
+                + 28 * z ** 6
+                + 1134 * z ** 5
+                + 47250 * z ** 4
+                + 1819125 * z ** 3
+                + 58939650 * z ** 2
+                + 1404728325 * z
+                + 18261468225
+            ),
+            8: 4392846440677
+            * z ** 8
+            / (
+                z ** 8
+                + 36 * z ** 7
+                + 1890 * z ** 6
+                + 103950 * z ** 5
+                + 5457375 * z ** 4
+                + 255405150 * z ** 3
+                + 9833098275 * z ** 2
+                + 273922023375 * z
+                + 4108830350625
             ),
+        }
+        return sp.Piecewise(
+            *[
+                (expression, sp.Eq(angular_momentum, value))
+                for value, expression in cases.items()
+                if self.max_angular_momentum is None
+                or value <= self.max_angular_momentum
+            ]
         )
 
     def _latex(self, printer: LatexPrinter, *args: Any) -> str:

diff --git a/tests/test_dynamics.py b/tests/test_dynamics.py
@@ -7,10 +7,27 @@
 from qrules import ParticleCollection
 from sympy import preorder_traversal
 
-from ampform.dynamics import ComplexSqrt
+from ampform.dynamics import BlattWeisskopfSquared, ComplexSqrt
 from ampform.helicity import HelicityModel
 
 
+class TestBlattWeisskopfSquared:
+    def test_max_angular_momentum(self):
+        z = sp.Symbol("z")
+        angular_momentum = sp.Symbol("L", integer=True)
+        form_factor = BlattWeisskopfSquared(angular_momentum, z=z)
+        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
+        assert BlattWeisskopfSquared.max_angular_momentum is None
+        BlattWeisskopfSquared.max_angular_momentum = 1
+        assert form_factor.evaluate() == sp.Piecewise(
+            (1, sp.Eq(angular_momentum, 0)),
+            (2 * z / (z + 1), sp.Eq(angular_momentum, 1)),
+        )
+
+
 def test_generate(
     amplitude_model: Tuple[str, HelicityModel],
     particle_database: ParticleCollection,