FIX: correct serialized intensity (#209)

* DOC: remove obsolete serialization validation warning
* DX: assert full intensity checksum in notebook
* DX: test parity recoupling and form factor conventions
* FIX: correct serialized chain amplitude construction
* FIX: normalize serialized Blatt-Weisskopf form factors
* FIX: use per-state reference subsystems for alignment

Author: redeboer

docs/serialization.ipynb

@@ -773,9 +773,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    ":::{error}\n",
-    "The following serves as a numerical check on whether the amplitude model has been deserialized correctly. For now, this is not the case, see [ComPWA/ampform-dpd#133](https://github.com/ComPWA/ampform-dpd/issues/133) for updates.\n",
-    ":::"
+    "The following checks whether the full amplitude model has been deserialized correctly."
    ]
   },
   {
@@ -819,15 +817,49 @@
    },
    "outputs": [],
    "source": [
+    "theta31 = next(symbol for symbol in MODEL.variables if str(symbol) == \"theta_31\")\n",
+    "cos_theta31_expr = sp.cos(MODEL.variables[theta31]).xreplace(MODEL.parameter_defaults)\n",
+    "sigma1 = next(s for s in cos_theta31_expr.free_symbols if str(s) == \"sigma1\")\n",
+    "sigma2 = next(s for s in cos_theta31_expr.free_symbols if str(s) == \"sigma2\")\n",
+    "z = sp.Symbol(\"z\", real=True)\n",
+    "s1_expr = sp.solve(sp.Eq(cos_theta31_expr, z), sigma1)[0]\n",
     "array = []\n",
     "for point_name, expected in checksums.items():\n",
     "    parameters = checksum_points[point_name]\n",
-    "    s1 = parameters[\"m_31_2\"] ** 2\n",
     "    s2 = parameters[\"m_31\"] ** 2\n",
-    "    computed = intensity_funcs[3]({\"sigma1\": s1, \"sigma2\": s2})\n",
+    "    s1 = s1_expr.xreplace({z: parameters[\"cos_theta_31\"], sigma2: s2})\n",
+    "    computed = intensity_funcs[3]({\"sigma1\": float(s1), \"sigma2\": s2})\n",
     "    status = label_diff(expected - computed)\n",
     "    array.append((INTENSITY_NAME, point_name, computed, expected, status))\n",
-    "pd.DataFrame(array, columns=[\"Distribution\", \"Point\", \"Computed\", \"Expected\", \"Status\"])"
+    "intensity_validation = pd.DataFrame(\n",
+    "    array,\n",
+    "    columns=[\"Distribution\", \"Point\", \"Computed\", \"Expected\", \"Status\"],\n",
+    ")\n",
+    "intensity_validation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "jupyter": {
+     "source_hidden": true
+    },
+    "mystnb": {
+     "code_prompt_show": "Assert that the intensity matches the reference values"
+    },
+    "tags": [
+     "hide-input"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "np.testing.assert_allclose(\n",
+    "    intensity_validation[\"Computed\"].astype(float).to_numpy(),\n",
+    "    intensity_validation[\"Expected\"].to_numpy(),\n",
+    "    atol=1e-10,\n",
+    "    rtol=0,\n",
+    ")"
    ]
   },
   {
@@ -944,12 +976,8 @@
     "jupyter": {
      "source_hidden": true
     },
-    "mystnb": {
-     "code_prompt_show": "Dalitz plot is not yet correct"
-    },
     "tags": [
-     "hide-input",
-     "hide-output"
+     "hide-input"
     ]
    },
    "outputs": [],

src/ampform_dpd/io/serialization/amplitude.py

@@ -42,16 +42,25 @@
     ParityVertex,
     get_decay_chains,
     get_distribution_def,
-    get_reference_topology,
 )
 from ampform_dpd.spin import create_spin_range
 
 if TYPE_CHECKING:
     from collections.abc import Callable
 
-    from ampform_dpd.decay import FinalStateID
+    from ampform_dpd.decay import FinalStateID, State, StateID
     from ampform_dpd.io.serialization.format import ModelDefinition
 
+_REFERENCE_SUBSYSTEMS: dict[StateID, FinalStateID] = {0: 1, 1: 1, 2: 2, 3: 3}
+"""Reference subsystem to use for the alignment rotation of each state.
+
+Each final state is aligned with respect to its own subsystem and the initial state with
+respect to subsystem 1. This is the convention of `ThreeBodyDecays.jl
+<https://github.com/mmikhasenko/ThreeBodyDecays.jl>`_, which the serialization format
+follows, and it is used instead of the single reference topology that the serialized
+model declares.
+"""
+
 
 def formulate(  # ruff: ignore[too-many-locals]
     model: ModelDefinition,
@@ -124,36 +133,46 @@ def formulate_chain_amplitude(  # ruff: ignore[too-many-locals, too-many-positio
     to_latex: Callable[[str], str] = identity_function,
     additional_builders: dict[str, PropagatorDynamicsBuilder] | None = None,
 ) -> dict[sp.Basic, complex | float | sp.Expr]:
+    r"""Formulate the amplitude for one decay chain of a serialized model.
+
+    This is the serialization counterpart of
+    `.DalitzPlotDecompositionBuilder.formulate_subsystem_amplitude`: the couplings and
+    dynamics are read from the model definition instead of being generated, but the
+    phase conventions, the Kronecker delta over the production helicities, and the sum
+    over the resonance helicity :math:`\lambda_R` are the same. The two implementations
+    have to be kept in sync.
+    """
     chain_defs = get_decay_chains(model)
     chain_definition = chain_defs[chain_idx]
-    # -----------------------
     dynamics = formulate_dynamics(
         chain_definition, model, to_latex, additional_builders
     )
     for vertex in chain_definition["vertices"]:
         dynamics *= formulate_form_factor(vertex, model)
-    # -----------------------
     weight, weight_val = _get_weight(chain_definition, to_latex)
-    # -----------------------
-    (i, λi_val), (j, λj_val) = _get_decay_product_helicities(chain_definition)
+    i, j = _get_decay_product_ids(chain_definition)
     θij, θij_expr = formulate_scattering_angle(i, j)
     jR = sp.Rational(chain_definition["propagators"][0]["spin"])  # ruff: ignore[non-lowercase-variable-in-function]
-    R_node, λR_val = _get_resonance_helicity(chain_definition)  # ruff: ignore[non-lowercase-variable-in-function]
-    λR = _get_helicity_symbol(R_node)
-    # -----------------------
+    λR = _get_helicity_symbol(_get_resonance_node(chain_definition))
     A = _generate_amplitude_index_bases()
-    subsystem_id = get_spectator_id(chain_definition["topology"])
+    spectator_id = get_spectator_id(chain_definition["topology"])
+    states = get_states(model)
+    helicities = (λ0, λ1, λ2, λ3)
     h_prod = formulate_recoupling(model, chain_idx, vertex_idx=0)
     h_dec = formulate_recoupling(model, chain_idx, vertex_idx=1)
-    amplitude_expression = (
+    chain_amplitude = (
         weight
+        * sp.sqrt(2 * jR + 1)
+        * _formulate_phase_factor(states[spectator_id], helicities[spectator_id])
+        * _formulate_phase_factor(states[j], helicities[j])
+        * δ(λ0, λR - helicities[spectator_id])
         * h_prod
         * h_dec
-        * Wigner.d(jR, λR, λi_val - λj_val, θij)
+        * Wigner.d(jR, λR, helicities[i] - helicities[j], θij)
         * dynamics.expression
     )
-    amplitude_expression = amplitude_expression.subs({λR: λR_val})
-    amplitude_symbol = A[subsystem_id][λ0, λ1, λ2, λ3]
+    amplitude_expression = PoolSum(chain_amplitude, (λR, create_spin_range(jR)))
+    amplitude_symbol = A[spectator_id][λ0, λ1, λ2, λ3]
     return {
         amplitude_symbol: amplitude_expression,
         weight: weight_val,
@@ -162,6 +181,34 @@ def formulate_chain_amplitude(  # ruff: ignore[too-many-locals, too-many-positio
     }
 
 
+def _formulate_phase_factor(state: State, helicity: sp.Rational | sp.Symbol) -> sp.Expr:
+    r"""Formulate the :math:`(-1)^{j-\lambda}` phase factor of a state."""
+    return (-1) ** (state.spin - helicity)
+
+
+def _get_decay_product_ids(
+    chain_definition: DecayChain,
+) -> tuple[FinalStateID, FinalStateID]:
+    """Get the IDs of the two decay products, ignoring their serialized helicities.
+
+    The helicity values from `._get_decay_product_helicities` are not substituted into
+    the chain amplitude: it is summed over all allowed helicities instead.
+    """
+    (i, _), (j, _) = _get_decay_product_helicities(chain_definition)
+    return cast("FinalStateID", i), cast("FinalStateID", j)
+
+
+def _get_resonance_node(
+    chain_definition: DecayChain,
+) -> tuple[FinalStateID, FinalStateID]:
+    """Get the node of the resonance, ignoring its serialized helicity.
+
+    See `._get_decay_product_ids` for why the helicity value is discarded.
+    """
+    node, _ = _get_resonance_helicity(chain_definition)
+    return node
+
+
 def _get_decay_product_helicities(
     chain_definition: DecayChain,
 ) -> tuple[tuple[int, sp.Rational], tuple[int, sp.Rational]]:
@@ -187,28 +234,38 @@ def formulate_aligned_amplitude(
     λ2: sp.Rational | sp.Symbol,
     λ3: sp.Rational | sp.Symbol,
 ) -> tuple[PoolSum, dict[sp.Symbol, sp.Expr]]:
-    reference_topology = get_reference_topology(model)
-    reference_subsystem = get_spectator_id(reference_topology)
-    wigner_generator = _AlignmentWignerGenerator(reference_subsystem)
+    generators = {
+        subsystem_id: _AlignmentWignerGenerator(subsystem_id)
+        for subsystem_id in sorted(set(_REFERENCE_SUBSYSTEMS.values()))
+    }
+    wigner_generators = {
+        rotated_state: generators[subsystem_id]
+        for rotated_state, subsystem_id in _REFERENCE_SUBSYSTEMS.items()
+    }
     _λ0, _λ1, _λ2, _λ3 = sp.symbols(R"\lambda_(:4)^{\prime}", rational=True)
     states = get_states(model)
     j0, j1, j2, j3 = (states[i].spin for i in sorted(states))
     A = _generate_amplitude_index_bases()
     amp_expr = PoolSum(
         sum(
             A[k][_λ0, _λ1, _λ2, _λ3]
-            * wigner_generator(j0, λ0, _λ0, rotated_state=0, aligned_subsystem=k)
-            * wigner_generator(j1, _λ1, λ1, rotated_state=1, aligned_subsystem=k)
-            * wigner_generator(j2, _λ2, λ2, rotated_state=2, aligned_subsystem=k)
-            * wigner_generator(j3, _λ3, λ3, rotated_state=3, aligned_subsystem=k)
+            * wigner_generators[0](j0, λ0, _λ0, rotated_state=0, aligned_subsystem=k)
+            * wigner_generators[1](j1, _λ1, λ1, rotated_state=1, aligned_subsystem=k)
+            * wigner_generators[2](j2, _λ2, λ2, rotated_state=2, aligned_subsystem=k)
+            * wigner_generators[3](j3, _λ3, λ3, rotated_state=3, aligned_subsystem=k)
             for k in get_existing_subsystem_ids(model)
         ),
         (_λ0, create_spin_range(j0)),
         (_λ1, create_spin_range(j1)),
         (_λ2, create_spin_range(j2)),
         (_λ3, create_spin_range(j3)),
     )
-    return amp_expr, wigner_generator.angle_definitions
+    angle_definitions = {
+        symbol: expression
+        for generator in generators.values()
+        for symbol, expression in generator.angle_definitions.items()
+    }
+    return amp_expr, angle_definitions
 
 
 def _get_weight(
@@ -386,6 +443,8 @@ class ParityRecoupling(sp.Expr):
 
     def evaluate(self) -> sp.Expr:
         λa, λb, λa0, λb0, f = self.args
+        if λa0 == 0 and λb0 == 0:
+            return δ(λa, λa0) * δ(λb, λb0)
         return δ(λa, λa0) * δ(λb, λb0) + f * δ(λa, -λa0) * δ(λb, -λb0)  # ty: ignore[unsupported-operator]
 
 

src/ampform_dpd/io/serialization/dynamics.py

@@ -4,7 +4,7 @@
 from typing import TYPE_CHECKING, Protocol, TypeVar, cast
 
 import sympy as sp
-from ampform.dynamics.form_factor import FormFactor
+from ampform.dynamics.form_factor import FormFactor, SphericalHankel1
 from sympy.parsing.sympy_parser import parse_expr
 
 from ampform_dpd import DefinedExpression
@@ -82,16 +82,21 @@ def formulate_form_factor(vertex: Vertex, model: ModelDefinition) -> DefinedExpr
     function_type = function_definition["type"]
     if function_type == "BlattWeisskopf":
         node = vertex["node"]
-        s = to_mandelstam_symbol(node)
-        m1, m2 = (to_mass_symbol(i) for i in node)
         if all(isinstance(i, int) for i in node):
+            s = to_mandelstam_symbol(node)
+            m1, m2 = (to_mass_symbol(i) for i in node)
             meson_radius = sp.Symbol(R"R_\mathrm{res}", nonnegative=True)
         else:
+            parent_mass = to_mandelstam_symbol(node)
+            isobar_invariant, m2 = (to_mass_symbol(i) for i in node)
+            s = parent_mass**2
+            m1 = sp.sqrt(isobar_invariant)
             initial_state = get_initial_state(model)
             meson_radius = sp.Symbol(f"R_{{{initial_state.latex}}}", nonnegative=True)
         angular_momentum = int(function_definition["l"])
         return DefinedExpression(
-            expression=FormFactor(s, m1, m2, angular_momentum, meson_radius),  # ty: ignore[invalid-argument-type]
+            expression=FormFactor(s, m1, m2, angular_momentum, meson_radius)  # ty: ignore[invalid-argument-type]
+            / _blatt_weisskopf_normalization(angular_momentum),
             parameters={
                 meson_radius: function_definition["radius"],
             },
@@ -100,6 +105,24 @@ def formulate_form_factor(vertex: Vertex, model: ModelDefinition) -> DefinedExpr
     raise NotImplementedError(msg)
 
 
+def _blatt_weisskopf_normalization(angular_momentum: int) -> sp.Expr:
+    r"""Undo the normalization of AmpForm's `~ampform.dynamics.form_factor.FormFactor`.
+
+    AmpForm normalizes its Blatt--Weisskopf factor to one at :math:`z=1`, whereas the
+    serialization format uses the unnormalized convention, so the form factor has to be
+    divided by this value.
+
+    >>> _blatt_weisskopf_normalization(0)
+    1
+    >>> _blatt_weisskopf_normalization(1)
+    sqrt(2)
+    >>> _blatt_weisskopf_normalization(2)
+    sqrt(13)
+    """
+    hankel = SphericalHankel1(sp.Integer(angular_momentum), sp.Integer(1))
+    return sp.Abs(hankel.doit())
+
+
 def formulate_generic_function(
     propagator: Propagator, resonance: str, model: ModelDefinition
 ) -> DefinedExpression:

tests/io_serialization/test_amplitude.py

@@ -6,6 +6,7 @@
 import sympy as sp
 
 from ampform_dpd.io.serialization.amplitude import (
+    ParityRecoupling,
     _get_decay_product_helicities,
     _get_final_state_helicities,
     _get_resonance_helicity,
@@ -60,3 +61,8 @@ def test_get_weight(model_definition: ModelDefinition):
     symbol, value = _get_weight(chain_defs[0])
     assert symbol.name == R"c^{L1405[1/2]}_{\frac{1}{2}, 0, 0}"
     assert value == pytest.approx(7.38649400481717 + 1.971018433257411j)
+
+
+def test_parity_recoupling_does_not_duplicate_zero_helicities():
+    recoupling = ParityRecoupling(λa=0, λb=0, λa0=0, λb0=0, f=1)  # ty: ignore[unknown-argument]
+    assert recoupling.doit() == 1

tests/io_serialization/test_dynamics.py

@@ -2,8 +2,12 @@
 
 import pytest
 import sympy as sp
+from ampform.dynamics.form_factor import FormFactor
 
-from ampform_dpd.io.serialization.dynamics import formulate_dynamics
+from ampform_dpd.io.serialization.dynamics import (
+    formulate_dynamics,
+    formulate_form_factor,
+)
 from ampform_dpd.io.serialization.format import (
     DecayChain,
     GenericFunctionDefinition,
@@ -52,3 +56,17 @@ def test_formulate_bugg_lineshapes(model_definition: ModelDefinition, chain_id:
     dynamics = formulate_dynamics(chain, model_definition)
 
     assert dynamics.expression.free_symbols == {sp.Symbol("sigma1", nonnegative=True)}
+
+
+def test_formulate_form_factor_uses_serialized_normalization(
+    model_definition: ModelDefinition,
+):
+    vertex = get_decay_chains(model_definition)[2]["vertices"][0]
+    form_factor = formulate_form_factor(vertex, model_definition)
+    assert sp.sqrt(2) * form_factor.expression == FormFactor(
+        s=sp.Symbol("m0", nonnegative=True) ** 2,  # ty: ignore[unknown-argument]
+        m1=sp.sqrt(sp.Symbol("sigma2", nonnegative=True)),  # ty: ignore[unknown-argument]
+        m2=sp.Symbol("m2", nonnegative=True),  # ty: ignore[unknown-argument]
+        angular_momentum=1,  # ty: ignore[unknown-argument]
+        meson_radius=sp.Symbol("R_{Lc}", nonnegative=True),  # ty: ignore[unknown-argument]
+    )