Modify amplitude model¶

Tip

Submit a feature request if you would like to see these procedures provided by AmpForm as a function!

import attr
import graphviz
import qrules as q
import sympy as sp

from ampform import get_builder

Since a HelicityModel.expression is simply a sympy.Expr, it’s relatively easy to modify it. In this notebook, we show some examples using the following example decay:

result = q.generate_transitions(
    initial_state=("J/psi(1S)", [-1, +1]),
    final_state=["gamma", "pi0", "pi0"],
    allowed_intermediate_particles=["f(0)(980)", "f(0)(1500)"],
    allowed_interaction_types=["strong", "EM"],
    formalism="helicity",
)
model_builder = get_builder(result)
original_model = model_builder.formulate()

dot = q.io.asdot(result, collapse_graphs=True)
graphviz.Source(dot)

Parameter substitution¶

Let’s say we want to express all coefficients as a product of magnitude \(C\) with their phase \(\phi\).

original_coefficients = [
    par
    for par in original_model.parameter_defaults
    if par.name.startswith("C")
]
original_coefficients

[C_{J/\psi(1S) \to f_{0}(980)_{0} \gamma_{+1}; f_{0}(980) \to \pi^{0}_{0} \pi^{0}_{0}},
 C_{J/\psi(1S) \to f_{0}(1500)_{0} \gamma_{+1}; f_{0}(1500) \to \pi^{0}_{0} \pi^{0}_{0}}]

There are two things to note now:

These parameters appear in HelicityModel.expression, its parameter_defaults, and its components, so both these attributes should be modified accordingly.
A HelicityModel is immutable, so we cannot directly replace its attributes. Instead, we should create a new HelicityModel with substituted attributes using attr.evolve():

The following snippet shows how to do all this. It’s shown in full, because it could well be you want to perform some completely different substitutions (can be any kinds of subs()). The overall procedure is comparable, however.

new_expression = original_model.expression
new_parameter_defaults = dict(original_model.parameter_defaults)  # copy!
new_components = dict(original_model.components)  # copy!

for coefficient in original_coefficients:
    decay_description = coefficient.name[3:-1]
    magnitude = sp.Symbol(  # coefficient with same name, but real, not complex
        coefficient.name,
        real=True,
        positive=True,
    )
    phase = sp.Symbol(
        fR"\phi_{{{decay_description}}}",
        real=True,
    )
    replacement = magnitude * sp.exp(sp.I * phase)
    display(replacement)
    # replace parameter defaults
    del new_parameter_defaults[coefficient]
    new_parameter_defaults[magnitude] = 1.0
    new_parameter_defaults[phase] = 0.0
    # replace parameters in expression
    new_expression = new_expression.subs(
        coefficient, replacement, simultaneous=True
    )
    # replace parameters in each component
    new_components = {
        key: old_expression.subs(coefficient, replacement, simultaneous=True)
        for key, old_expression in new_components.items()
    }

# create new model from the old
new_model = attr.evolve(
    original_model,
    expression=new_expression,
    parameter_defaults=new_parameter_defaults,
    components=new_components,
)

\[\displaystyle C_{J/\psi(1S) \to f_{0}(980)_{0} \gamma_{+1}; f_{0}(980) \to \pi^{0}_{0} \pi^{0}_{0}} e^{i \phi_{J/\psi(1S) \to f_{0}(980)_{0} \gamma_{+1}; f_{0}(980) \to \pi^{0}_{0} \pi^{0}_{0}}}\]

\[\displaystyle C_{J/\psi(1S) \to f_{0}(1500)_{0} \gamma_{+1}; f_{0}(1500) \to \pi^{0}_{0} \pi^{0}_{0}} e^{i \phi_{J/\psi(1S) \to f_{0}(1500)_{0} \gamma_{+1}; f_{0}(1500) \to \pi^{0}_{0} \pi^{0}_{0}}}\]

assert new_model != original_model

As can be seen, the parameter_defaults have bene updated, as have the components:

new_model.parameter_defaults

{C_{J/\psi(1S) \to f_{0}(980)_{0} \gamma_{+1}; f_{0}(980) \to \pi^{0}_{0} \pi^{0}_{0}}: 1.0,
 \phi_{J/\psi(1S) \to f_{0}(980)_{0} \gamma_{+1}; f_{0}(980) \to \pi^{0}_{0} \pi^{0}_{0}}: 0.0,
 C_{J/\psi(1S) \to f_{0}(1500)_{0} \gamma_{+1}; f_{0}(1500) \to \pi^{0}_{0} \pi^{0}_{0}}: 1.0,
 \phi_{J/\psi(1S) \to f_{0}(1500)_{0} \gamma_{+1}; f_{0}(1500) \to \pi^{0}_{0} \pi^{0}_{0}}: 0.0}

new_model.components[
    R"A_{J/\psi(1S)_{-1} \to f_{0}(980)_{0} \gamma_{-1}; f_{0}(980)_{0} \to \pi^{0}_{0} \pi^{0}_{0}}"
]

\[\displaystyle C_{J/\psi(1S) \to f_{0}(980)_{0} \gamma_{+1}; f_{0}(980) \to \pi^{0}_{0} \pi^{0}_{0}} e^{i \phi_{J/\psi(1S) \to f_{0}(980)_{0} \gamma_{+1}; f_{0}(980) \to \pi^{0}_{0} \pi^{0}_{0}}} D^{0}_{0,0}\left(- \phi_{1,1+2},\theta_{1,1+2},0\right) D^{1}_{-1,1}\left(- \phi_{1+2},\theta_{1+2},0\right)\]

Also note that the new model reduces to the old once we replace the parameters with their suggested default values:

evaluated_expr = new_model.expression.subs(new_model.parameter_defaults).doit()
evaluated_expr

\[\displaystyle 8.0 \left(\frac{\cos{\left(\theta_{1+2} \right)}}{2} - \frac{1}{2}\right)^{2} + 8.0 \left(\frac{\cos{\left(\theta_{1+2} \right)}}{2} + \frac{1}{2}\right)^{2}\]

assert (
    original_model.expression.subs(original_model.parameter_defaults).doit()
    == evaluated_expr
)

Couple parameters¶

Similarly to Parameter substitution, we can couple parameters by substituting pairs of parameters with a new Symbol:

parameter_couplings = [
    (
        sp.Symbol(
            R"C_{J/\psi(1S) \to f_{0}(980)_{0} \gamma_{+1}; f_{0}(980) \to \pi^{0}_{0} \pi^{0}_{0}}"
        ),
        sp.Symbol(
            R"C_{J/\psi(1S) \to f_{0}(1500)_{0} \gamma_{+1}; f_{0}(1500) \to \pi^{0}_{0} \pi^{0}_{0}}"
        ),
    )
]

new_expression = original_model.expression
new_parameter_defaults = dict(original_model.parameter_defaults)  # copy!
new_components = dict(original_model.components)  # copy!

for i, (par1, par2) in enumerate(parameter_couplings):
    # construct new symbol for the coupled parameter pair
    original_assumptions = {
        **par1.assumptions0,
        **par2.assumptions0,
    }
    new_par = sp.Symbol(f"C{i}", **original_assumptions)
    # replace parameter defaults
    value = new_parameter_defaults[par1]
    new_parameter_defaults[new_par] = value
    del new_parameter_defaults[par1]
    del new_parameter_defaults[par2]
    # replace parameters in expression
    new_expression = new_expression.subs(par1, new_par, simultaneous=True)
    new_expression = new_expression.subs(par2, new_par, simultaneous=True)
    # replace parameters in each component
    new_components = {
        key: old_expression.subs(
            {
                par1: new_par,
                par2: new_par,
            },
            simultaneous=True,
        )
        for key, old_expression in new_components.items()
    }

# create new model from the old
new_model = attr.evolve(
    original_model,
    expression=new_expression,
    parameter_defaults=new_parameter_defaults,
    components=new_components,
)

new_model.parameter_defaults

{C0: (1+0j)}

new_model.components[
    R"A_{J/\psi(1S)_{-1} \to f_{0}(980)_{0} \gamma_{-1}; f_{0}(980)_{0} \to \pi^{0}_{0} \pi^{0}_{0}}"
]

\[\displaystyle C_{0} D^{0}_{0,0}\left(- \phi_{1,1+2},\theta_{1,1+2},0\right) D^{1}_{-1,1}\left(- \phi_{1+2},\theta_{1+2},0\right)\]