Using composition

Using composition#

Hide code cell content 
from __future__ import annotations

import sympy as sp
from attrs import frozen
from helpers import (
    StateTransitionGraph,
    blatt_weisskopf,
    determine_attached_final_state,
    two_body_momentum_squared,
)

try:
    from typing import Protocol
except ImportError:
    from typing import Protocol

A frozen DynamicExpression class keeps track of variables, parameters, and the dynamics expression in which they should appear:

@frozen
class DynamicsExpression:
    variables: tuple[sp.Symbol, ...]
    parameters: tuple[sp.Symbol, ...]
    expression: sp.Expr

    def substitute(self) -> sp.Expr:
        return self.expression(*self.variables, *self.parameters)

The expression attribute can be formulated as a simple Python function that takes sympy.Symbols as arguments and returns a sympy.Expr:

def relativistic_breit_wigner(
    mass: sp.Symbol, mass0: sp.Symbol, gamma0: sp.Symbol
) -> sp.Expr:
    return gamma0 * mass0 / (mass0**2 - mass**2 - gamma0 * mass0 * sp.I)

def relativistic_breit_wigner_with_form_factor(
    mass: sp.Symbol,
    mass0: sp.Symbol,
    gamma0: sp.Symbol,
    m_a: sp.Symbol,
    m_b: sp.Symbol,
    angular_momentum: sp.Symbol,
    meson_radius: sp.Symbol,
) -> sp.Expr:
    q_squared = two_body_momentum_squared(mass, m_a, m_b)
    q0_squared = two_body_momentum_squared(mass0, m_a, m_b)
    ff2 = blatt_weisskopf(q_squared, meson_radius, angular_momentum)
    ff02 = blatt_weisskopf(q0_squared, meson_radius, angular_momentum)
    width = gamma0 * (mass0 / mass) * (ff2 / ff02) * sp.sqrt(q_squared / q0_squared)
    return relativistic_breit_wigner(mass, mass0, width) * mass0 * gamma0 * sp.sqrt(ff2)

The DynamicsExpression container class enables us to provide the expression with correctly named Symbols for the decay that is being described. Here, we use some naming scheme for an \(f_0(980)\) decaying to final state edges 3 and 4 (say \(\pi^0\pi^0\)):

bw_decay_f0 = DynamicsExpression(
    variables=sp.symbols("m_3+4", seq=True),
    parameters=sp.symbols(R"m_f(0)(980) \Gamma_f(0)(980)"),
    expression=relativistic_breit_wigner,
)
bw_decay_f0.substitute()
\[\displaystyle \frac{\Gamma_{f(0)(980)} m_{f(0)(980)}}{- i \Gamma_{f(0)(980)} m_{f(0)(980)} - m_{3+4}^{2} + m_{f(0)(980)}^{2}}\]

For each dynamics expression, we have to provide a ‘builder’ function that can create a DynamicsExpression for a specific edge within the StateTransitionGraph:

def relativistic_breit_wigner_from_graph(
    graph: StateTransitionGraph, edge_id: int
) -> DynamicsExpression:
    edge_ids = determine_attached_final_state(graph, edge_id)
    final_state_ids = map(str, edge_ids)
    mass = sp.Symbol(f"m_{{{"+".join(final_state_ids)}}}")
    particle, _ = graph.get_edge_props(edge_id)
    mass0 = sp.Symbol(f"m_{{{particle.latex}}}")
    gamma0 = sp.Symbol(Rf"\Gamma_{{{particle.latex}}}")
    return DynamicsExpression(
        variables=(mass),
        parameters=(mass0, gamma0),
        expression=relativistic_breit_wigner(mass, mass0, gamma0),
    )

def relativistic_breit_wigner_with_form_factor_from_graph(
    graph: StateTransitionGraph, edge_id: int
) -> DynamicsExpression:
    edge_ids = determine_attached_final_state(graph, edge_id)
    final_state_ids = map(str, edge_ids)
    mass = sp.Symbol(f"m_{{{"+".join(final_state_ids)}}}")
    particle, _ = graph.get_edge_props(edge_id)
    mass0 = sp.Symbol(f"m_{{{particle.latex}}}")
    gamma0 = sp.Symbol(Rf"\Gamma_{{{particle.latex}}}")
    m_a = sp.Symbol(f"m_{edge_ids[0]}")
    m_b = sp.Symbol(f"m_{edge_ids[1]}")
    angular_momentum = particle.spin  # helicity formalism only!
    meson_radius = sp.Symbol(Rf"R_{{{particle.latex}}}")
    return DynamicsExpression(
        variables=(mass),
        parameters=(
            mass0,
            gamma0,
            m_a,
            m_b,
            angular_momentum,
            meson_radius,
        ),
        expression=relativistic_breit_wigner_with_form_factor(
            mass,
            mass0,
            gamma0,
            m_a,
            m_b,
            angular_momentum,
            meson_radius,
        ),
    )

The fact that DynamicsExpression.expression is just a Python function, allows one to inspect the dynamics formulation of these functions independently, purely in terms of SymPy:

m, m0, w0 = sp.symbols(R"m m_0 \Gamma")
evaluated_bw = relativistic_breit_wigner(m, 1.0, 0.3)
relativistic_breit_wigner(m, m0, w0)
\[\displaystyle \frac{\Gamma m_{0}}{- i \Gamma m_{0} - m^{2} + m_{0}^{2}}\]
sp.plot(sp.Abs(evaluated_bw), (m, 0, 2), axis_center=(0, 0), ylim=(0, 1))
sp.plot(sp.arg(evaluated_bw), (m, 0, 2), axis_center=(0, 0), ylim=(0, sp.pi));

This closes the gap between the code and the theory that is being implemented.

Alternative signature#

An alternative way to specify the expression is:

def expression(
    variables: tuple[sp.Symbol, ...], parameters: tuple[sp.Symbol, ...]
) -> sp.Expr:
    pass

Here, one would however need to unpack the variables and parameters. The advantage is that the signature becomes more general.

Type checking#

There is no way to enforce the appropriate signature on the builder function, other than following a Protocol:

class DynamicsBuilder(Protocol):
    def __call__(
        self, graph: StateTransitionGraph, edge_id: int
    ) -> DynamicsExpression: ...

This DynamicsBuilder protocol would be used in the syntax proposed at Considered solutions.

It carries another subtle problem, though: a Protocol is only used in static type checking, while potential problems with the implementation of the builder and its respective expression only arrise at runtime.

Contents