Amplitude behavior investigation

Tip

▶️ Run this notebook interactively here.

Intensity function for two-pseudoscalar system:

(1)\[I(\Omega,\Phi) = 2\kappa\sum_{k}\left( (1-P_{\gamma})\left|\sum_{l,m}[l]_{m;k}^{(-)}\Re[Z_{l}^{m}(\Omega,\Phi)]\right|^2 +(1-P_{\gamma})\left|\sum_{l,m}[l]_{m;k}^{(+)}\Im[Z_{l}^{m}(\Omega,\Phi)]\right|^2 +(1+P_{\gamma})\left|\sum_{l,m}[l]_{m;k}^{(+)}\Re[Z_{l}^{m}(\Omega,\Phi)]\right|^2 +(1+P_{\gamma})\left|\sum_{l,m}[l]_{m;k}^{(-)}\Im[Z_{l}^{m}(\Omega,\Phi)]\right|^2 \right)\]

where \(Z_{l}^{m}(\Omega,\Phi)=Y_{l}^{m}(\Omega)e^{-i\Phi}\) is a phase-rotated spherical harmonic, \(\Omega\) is the solid angle, \(\Phi\) is the angle between the production and polarization planes, \(P_{\gamma}\) is the polarization magnitude, \([l]\) are the partial wave amplitudes, \(m\) is the associated m-projection, \(k\) refers to a spin flip (\(k=1\)) or non-flip (\(k=0\)) at the nucleon vertex, and \(\kappa\) is an overall phase space factor.

Import Python modules

Hide code cell source

from collections import defaultdict
from typing import Self

import ipywidgets as w
import matplotlib.pyplot as plt
import numpy as np
import sympy as sp
from ampform.io import aslatex
from ampform.sympy import (
    PoolSum,
    UnevaluatedExpression,
    create_expression,
    implement_doit_method,
)
from IPython.display import Math
from qrules.quantum_numbers import arange
from sympy.functions.special.spherical_harmonics import Ynm
from sympy.printing.precedence import PRECEDENCE
from sympy.printing.printer import Printer


def _print_Indexed_latex(self, printer, *args):  # noqa: N802
    base = printer._print(self.base)
    indices = ", ".join(map(printer._print, self.indices))
    return f"{base}_{{{indices}}}"


sp.Indexed.precedence = PRECEDENCE["Pow"] - 1
sp.Indexed._latex = _print_Indexed_latex
Printer._global_settings["gothic_re_im"] = True

Main intensity definition

We build the amplitude symbolically and step by step using SymPy. Of course, we can automate things here later, but defining everything step-by-step gives us most control. Let’s start by defining the main symbols that appear in amplitude expression in Eq. (1).

k, l, m = sp.symbols("k l m")
phi, theta, Phi = sp.symbols("phi theta Phi", real=True)
kappa = sp.Symbol("kappa")
Py = sp.Symbol("P_gamma")

Next, we define a special handler class for nicely rendering the \(Z_m^l(\theta,\phi,\Phi)\) function:

Define symbolic Zlm class

Hide code cell source

@implement_doit_method
class Znm(UnevaluatedExpression):
    is_commutative = True
    is_real = False

    def __new__(cls, n, m, theta, phi, Phi, **hints) -> Self:
        return create_expression(cls, n, m, theta, phi, Phi, **hints)

    def evaluate(self) -> sp.Mul:
        n, m, theta, phi, Phi = self.args
        return Ynm(n, m, theta, phi) * sp.exp(-sp.I * Phi)

    def _latex(self, printer, *args):
        n, m, theta, phi, Phi = map(printer._print, self.args)
        return Rf"Z_{n}^{m}({theta}, {phi}, {Phi})"

znm = Znm(l, m, theta, phi, Phi)
Math(aslatex({znm: znm.doit()}))

\[\begin{split}\displaystyle \begin{array}{rcl} Z_l^m(\theta, \phi, \Phi) &=& e^{- i \Phi} Y_{l}^{m}\left(\theta,\phi\right) \\ \end{array}\end{split}\]

We’re now ready to define the main intensity. Note the use of IndexedBase for generating the \([l]^{(\pm)}_{m;k}\) symbols and the use of PoolSum, which is more suited for nested sums than SymPy’s Sum class.

Show code cell source

Hide code cell source

def create_range(min_, max_, /) -> tuple[sp.Rational, ...]:
    return tuple(sp.Rational(x) for x in arange(min_, max_ + 1))


max_l = 2
k_values = [0]
l_range = [max_l]  # create_range(0, max_l)
m_range = create_range(-max_l, max_l)
lp = sp.IndexedBase("[l]^{(+)}", complex=True)
lm = sp.IndexedBase("[l]^{(-)}", complex=True)

intensity_expr = (
    2
    * kappa
    * PoolSum(
        (1 - Py)
        * sp.Abs(PoolSum(lm[m, k] * sp.re(znm), (l, l_range), (m, m_range))) ** 2
        + (1 - Py)
        * sp.Abs(PoolSum(lp[m, k] * sp.im(znm), (l, l_range), (m, m_range))) ** 2
        + (1 + Py)
        * sp.Abs(PoolSum(lp[m, k] * sp.re(znm), (l, l_range), (m, m_range))) ** 2
        + (1 + Py)
        * sp.Abs(PoolSum(lm[m, k] * sp.im(znm), (l, l_range), (m, m_range))) ** 2,
        (k, k_values),
    )
)
intensity_expr

\[\displaystyle 2 \kappa \sum_{k=0}{\left(1 - P_{\gamma}\right) \left|{\sum_{l=2} \sum_{m=-2}^{2}{\Re{\left(Z_l^m(\theta, \phi, \Phi)\right)} [l]^{(-)}_{m, k}}}\right|^{2} + \left(1 - P_{\gamma}\right) \left|{\sum_{l=2} \sum_{m=-2}^{2}{\Im{\left(Z_l^m(\theta, \phi, \Phi)\right)} [l]^{(+)}_{m, k}}}\right|^{2} + \left(P_{\gamma} + 1\right) \left|{\sum_{l=2} \sum_{m=-2}^{2}{\Re{\left(Z_l^m(\theta, \phi, \Phi)\right)} [l]^{(+)}_{m, k}}}\right|^{2} + \left(P_{\gamma} + 1\right) \left|{\sum_{l=2} \sum_{m=-2}^{2}{\Im{\left(Z_l^m(\theta, \phi, \Phi)\right)} [l]^{(-)}_{m, k}}}\right|^{2}}\]

SymPy can ‘unfold’ all definitons like \(Z^m_l\) and \(\sum_{m;k}\) into basic mathematical operations as follows. This is not only a nice mathematical exercise, but is crucial for when we use the expression to create a numerical function with lambdification.

sp.expand_func(intensity_expr.doit())

Show code cell output

Hide code cell output

\[\displaystyle 2 \kappa \left(\left(1 - P_{\gamma}\right) \left|{\frac{\sqrt{5} \left(\frac{3 \cos^{2}{\left(\theta \right)}}{2} - \frac{1}{2}\right) \sin{\left(\Phi \right)} [l]^{(+)}_{0, 0}}{2 \sqrt{\pi}} + \frac{\sqrt{30} \left(3 \cos^{2}{\left(\theta \right)} - 3\right) \Im{\left(e^{- i \Phi} e^{- 2 i \phi}\right)} [l]^{(+)}_{-2, 0}}{24 \sqrt{\pi}} + \frac{\sqrt{30} \left(3 \cos^{2}{\left(\theta \right)} - 3\right) \Im{\left(e^{- i \Phi} e^{2 i \phi}\right)} [l]^{(+)}_{2, 0}}{24 \sqrt{\pi}} - \frac{\sqrt{30} \sin{\left(\theta \right)} \cos{\left(\theta \right)} \Im{\left(e^{- i \Phi} e^{- i \phi}\right)} [l]^{(+)}_{-1, 0}}{4 \sqrt{\pi}} + \frac{\sqrt{30} \sin{\left(\theta \right)} \cos{\left(\theta \right)} \Im{\left(e^{- i \Phi} e^{i \phi}\right)} [l]^{(+)}_{1, 0}}{4 \sqrt{\pi}}}\right|^{2} + \left(1 - P_{\gamma}\right) \left|{- \frac{\sqrt{5} \left(\frac{3 \cos^{2}{\left(\theta \right)}}{2} - \frac{1}{2}\right) \cos{\left(\Phi \right)} [l]^{(-)}_{0, 0}}{2 \sqrt{\pi}} + \frac{\sqrt{30} \left(3 \cos^{2}{\left(\theta \right)} - 3\right) \Re{\left(e^{- i \Phi} e^{- 2 i \phi}\right)} [l]^{(-)}_{-2, 0}}{24 \sqrt{\pi}} + \frac{\sqrt{30} \left(3 \cos^{2}{\left(\theta \right)} - 3\right) \Re{\left(e^{- i \Phi} e^{2 i \phi}\right)} [l]^{(-)}_{2, 0}}{24 \sqrt{\pi}} - \frac{\sqrt{30} \sin{\left(\theta \right)} \cos{\left(\theta \right)} \Re{\left(e^{- i \Phi} e^{- i \phi}\right)} [l]^{(-)}_{-1, 0}}{4 \sqrt{\pi}} + \frac{\sqrt{30} \sin{\left(\theta \right)} \cos{\left(\theta \right)} \Re{\left(e^{- i \Phi} e^{i \phi}\right)} [l]^{(-)}_{1, 0}}{4 \sqrt{\pi}}}\right|^{2} + \left(P_{\gamma} + 1\right) \left|{\frac{\sqrt{5} \left(\frac{3 \cos^{2}{\left(\theta \right)}}{2} - \frac{1}{2}\right) \sin{\left(\Phi \right)} [l]^{(-)}_{0, 0}}{2 \sqrt{\pi}} + \frac{\sqrt{30} \left(3 \cos^{2}{\left(\theta \right)} - 3\right) \Im{\left(e^{- i \Phi} e^{- 2 i \phi}\right)} [l]^{(-)}_{-2, 0}}{24 \sqrt{\pi}} + \frac{\sqrt{30} \left(3 \cos^{2}{\left(\theta \right)} - 3\right) \Im{\left(e^{- i \Phi} e^{2 i \phi}\right)} [l]^{(-)}_{2, 0}}{24 \sqrt{\pi}} - \frac{\sqrt{30} \sin{\left(\theta \right)} \cos{\left(\theta \right)} \Im{\left(e^{- i \Phi} e^{- i \phi}\right)} [l]^{(-)}_{-1, 0}}{4 \sqrt{\pi}} + \frac{\sqrt{30} \sin{\left(\theta \right)} \cos{\left(\theta \right)} \Im{\left(e^{- i \Phi} e^{i \phi}\right)} [l]^{(-)}_{1, 0}}{4 \sqrt{\pi}}}\right|^{2} + \left(P_{\gamma} + 1\right) \left|{- \frac{\sqrt{5} \left(\frac{3 \cos^{2}{\left(\theta \right)}}{2} - \frac{1}{2}\right) \cos{\left(\Phi \right)} [l]^{(+)}_{0, 0}}{2 \sqrt{\pi}} + \frac{\sqrt{30} \left(3 \cos^{2}{\left(\theta \right)} - 3\right) \Re{\left(e^{- i \Phi} e^{- 2 i \phi}\right)} [l]^{(+)}_{-2, 0}}{24 \sqrt{\pi}} + \frac{\sqrt{30} \left(3 \cos^{2}{\left(\theta \right)} - 3\right) \Re{\left(e^{- i \Phi} e^{2 i \phi}\right)} [l]^{(+)}_{2, 0}}{24 \sqrt{\pi}} - \frac{\sqrt{30} \sin{\left(\theta \right)} \cos{\left(\theta \right)} \Re{\left(e^{- i \Phi} e^{- i \phi}\right)} [l]^{(+)}_{-1, 0}}{4 \sqrt{\pi}} + \frac{\sqrt{30} \sin{\left(\theta \right)} \cos{\left(\theta \right)} \Re{\left(e^{- i \Phi} e^{i \phi}\right)} [l]^{(+)}_{1, 0}}{4 \sqrt{\pi}}}\right|^{2}\right)\]

Analytic substitution

Before we go into numerical computations, let’s try what happens when we substitute certain parameters analytically with SymPy. First, we define a mapping of symbols in our expressions to values with which they should be substituted.

Define suggested parameter values

Hide code cell source

lp_defaults = {lp[m, k]: 0 for k in k_values for m in m_range}
lm_defaults = {lm[m, k]: 1 for k in k_values for m in m_range}
parameter_defaults = {
    Py: 1,
    kappa: sp.pi,
    **lp_defaults,
    **lm_defaults,
}

This mapping can then be used to substitute the symbols in our ‘unfolded’ expression analytically. Note the use of xreplace(), which is faster than subs():

substituted_expr = intensity_expr.doit().xreplace(parameter_defaults)
substituted_expr

\[\displaystyle 4 \pi \left(- \frac{\sqrt{5} \left(\frac{3 \cos^{2}{\left(\theta \right)}}{2} - \frac{1}{2}\right) \sin{\left(\Phi \right)}}{2 \sqrt{\pi}} + \Im{\left(e^{- i \Phi} Y_{2}^{1}\left(\theta,\phi\right)\right)} + \Im{\left(e^{- i \Phi} Y_{2}^{2}\left(\theta,\phi\right)\right)} + \Im{\left(e^{- i \Phi} e^{- 4 i \phi} Y_{2}^{2}\left(\theta,\phi\right)\right)} - \Im{\left(e^{- i \Phi} e^{- 2 i \phi} Y_{2}^{1}\left(\theta,\phi\right)\right)}\right)^{2}\]

The expression can be further unfolded and simplified as follows:

sp.expand_func(substituted_expr).simplify()

\[\displaystyle \frac{\left(2 \sqrt{5} \left(3 \cos^{2}{\left(\theta \right)} - 1\right) \sin{\left(\Phi \right)} - \frac{\sqrt{30} \left(\cos{\left(- \Phi + \phi + 2 \theta \right)} - \cos{\left(\Phi - \phi + 2 \theta \right)}\right)}{2} + \frac{\sqrt{30} \left(\cos{\left(\Phi + \phi - 2 \theta \right)} - \cos{\left(\Phi + \phi + 2 \theta \right)}\right)}{2} + \sqrt{30} \sin^{2}{\left(\theta \right)} \sin{\left(\Phi - 2 \phi \right)} + \sqrt{30} \sin^{2}{\left(\theta \right)} \sin{\left(\Phi + 2 \phi \right)}\right)^{2}}{16}\]

Widget for investigating amplitude behavior

SymPy is not suitable for evaluating expressions over large data samples. For this, we need to lambdify our main expression to a numerical function that uses faster backends, such as NumPy or JAX.

In this section, we use this mechanism to evaluate the intensity function over a uniform \((\theta,\phi)\) phase space grid sample. To get an idea of performance, we make the visualisation of the intensity distribution interactive with ipywidgets sliders that allow us to modify the parameters and amplitudes in our model, as well as chose specific values for \(\Phi\).

Our phase space grid sample is defined as follows. Note that the grid is defined in terms of \(\cos\theta\), but that our intensity expects \(\theta\) values.

cos_theta_array, phi_array = np.meshgrid(
    np.linspace(-1, +1, num=200),
    np.linspace(0, 2 * np.pi, num=400),
)
theta_array = np.arccos(cos_theta_array)

To lambdify to a numerical function, we need to completely unfold our main intensity expression and identify which symbols are to become arguments to our numerical function. These function arguments are then placeholds to data columns (here: \(\theta\) and \(\phi\)) and parameters (all remaining symbols).

free_symbols = sorted(intensity_expr.doit().free_symbols, key=str)
amp_symbols = [s for s in free_symbols if isinstance(s, sp.Indexed)]
arguments = (theta, phi, Phi, Py, *amp_symbols)
plotted_expr = intensity_expr.doit().xreplace({kappa: sp.pi}).expand(func=True)
intensity_func = sp.lambdify(arguments, plotted_expr, "numpy", cse=True)
intensity_func

<function _lambdifygenerated(theta, phi, Phi, P_gamma, _Dummy_43, _Dummy_42, _Dummy_41, _Dummy_40, _Dummy_39, _Dummy_38, _Dummy_37, _Dummy_36, _Dummy_35, _Dummy_34)>

The following code is the most complicated part, but is just there to create a nice user interface for the sliders as well as the interactive visualization…

Define widget sliders

Hide code cell source

# Create sliders for each symbol/argument
sliders = dict(
    Phi_val=w.FloatSlider(
        description="Φ",
        min=0,
        max=np.pi,
        value=1,
        step=np.pi / 20,
    ),
    Py_val=w.FloatSlider(
        description="Pγ",
        min=0,
        max=1,
        value=1,
    ),
)
for s in free_symbols:
    if not isinstance(s, sp.Indexed):
        continue
    if not s.name.startswith("[l]"):
        continue
    key = s.name
    sliders[f"{key}_mag"] = w.FloatSlider(
        description="magnitude",
        min=0,
        max=2,
        value=1 if "+" in s.name else 0,
    )
    sliders[f"{key}_phi"] = w.FloatSlider(
        description="phase",
        min=0,
        max=2,
    )

# Group sliders so they can be organised in a layout for the UI
l_sliders = defaultdict(list)
for key in sliders:
    if not key.startswith("[l]^{("):
        continue
    if not key.endswith("_mag"):
        continue
    l_str = key.strip("_mag")
    sign = l_str[6]
    indices = eval(l_str[9:])
    superscript = f"<sup>({sign})</sup>"
    subscript = f"<sub>{','.join(f'{i:+d}' if i else str(i) for i in indices)}</sub>"
    label = f"[𝑙]{superscript}{subscript}"
    row = w.HBox([w.HTML(label), sliders[f"{l_str}_mag"], sliders[f"{l_str}_phi"]])
    l_sliders[sign].append(row)

# Create the UI
vbox_items = [
    sliders["Phi_val"],
    sliders["Py_val"],
    w.Accordion(
        [w.VBox(l_sliders["+"]), w.VBox(l_sliders["-"])],
        selected_index=0,
        titles=["Positive [𝑙]", "Negative [𝑙]"],
    ),
]
UI = w.VBox(vbox_items)

Show code cell source

Hide code cell source

%matplotlib widget
fig, ax = plt.subplots(figsize=(7, 4))
ax.set_xlabel(R"$\phi$")
ax.set_ylabel(R"$\cos\theta$")
ax.set_yticks([-1, 0, +1])
ax.set_yticklabels(["-1", "0", "+1"])
ax.set_xticks([0, np.pi / 2, np.pi, 3 * np.pi / 2, 2 * np.pi])
ax.set_xticklabels(["0", R"$\frac{1}{2}\pi$", R"$\pi$", R"$\frac{3}{2}\pi$", R"$2\pi$"])
fig.canvas.toolbar_visible = False
fig.canvas.header_visible = False
fig.canvas.footer_visible = False

MESH = None


def plot(*, Phi_val, Py_val, **mag_phase):
    global MESH
    amplitudes = [
        mag * np.exp(mag_phase[f"{k[:-4]}_phi"] * 1j)
        for k, mag in mag_phase.items()
        if k.endswith("mag")
    ]
    Z = intensity_func(
        theta_array,
        phi_array,
        Phi_val,
        Py_val,
        *amplitudes,
    )

    if MESH is None:
        MESH = ax.pcolormesh(phi_array, cos_theta_array, Z, cmap=plt.cm.Reds)
        c_bar = fig.colorbar(MESH, ax=ax, pad=0.015)
        c_bar.ax.get_yaxis().labelpad = 15
        c_bar.ax.set_ylabel("Intensity (A.U.)", rotation=270)
    else:
        MESH.set_array(Z)
    Z_max = np.nanmax(Z)
    if Z_max > 0.1:
        MESH.set_clim(vmin=0, vmax=Z_max)
    fig.canvas.draw()


fig.tight_layout()
output = w.interactive_output(plot, controls=sliders)
w.VBox([UI, output])