TR-005

Symbolic K-matrix expressions

Implementation of this report is tracked through ampform#67.

To be implemented

K-matrix

This report investigates how to implement \(K\)-matrix dynamics with SymPy. We here describe only the version that is not Lorentz-invariant, because it is simplest and allows us to check whether the case \(n_R=1, n=1\) (single resonance, single channel) reduces to a Breit-Wigner function. We followed the physics as described by PDG2020, §Resonances and [Chung et al., 1995, Peters, 2004, Meyer, 2008]. For the Lorentz-invariant version, see TR-009.

A brief overview of the origin of the \(\boldsymbol{K}\)-matrix is given first. This overview follows [Chung et al., 1995], but skips over quite a few details, as this is only an attempt to provide some context of what is going on.

Show code cell content

Hide code cell content

from __future__ import annotations

import warnings
from pathlib import Path

import graphviz
import matplotlib.pyplot as plt
import mpl_interactions.ipyplot as iplt
import numpy as np
import symplot
import sympy as sp
from IPython.display import Math, display
from ipywidgets import widgets as ipywidgets
from mpl_interactions.controller import Controls

warnings.filterwarnings("ignore")

Physics

The \(\boldsymbol{K}\)-matrix formalism is used to describe coupled, two-body scattering processes of the type \(c_id_i \to R \to a_ib_i\), with \(i\) representing each separate channel and \(R\) a number of resonances that these channels have in common.

Show code cell source

Hide code cell source

dot = """
digraph {
    rankdir=LR;
    node [shape=point, width=0];
    edge [arrowhead=none];
    "Na" [shape=none, label="aᵢ"];
    "Nb" [shape=none, label="bᵢ"];
    "Nc" [shape=none, label="cᵢ"];
    "Nd" [shape=none, label="dᵢ"];
    { rank=same "Nc", "Nd" };
    { rank=same "Na", "Nb" };
    "Nc" -> "N0";
    "Nd" -> "N0";
    "N1" -> "Na";
    "N1" -> "Nb";
    "N0" -> "N1" [label="R"];
    "N0" [shape=none, label=""];
    "N1" [shape=none, label=""];
}
"""
graph = graphviz.Source(dot)
graph

Partial wave expansion

In amplitude analysis, the main aim is to express the differential cross section \(\frac{d\sigma}{d\Omega}\) (that is, the intensity distribution in each spherical direction \(\Omega=(\phi,\theta)\) as we can observe in experiments). This differential cross section can be expressed in terms of the scattering amplitude \(A\) as:

(1)\[ \frac{d\sigma}{d\Omega} = \left|A(\Omega)\right|^2 \]

We can now further express \(A\) in terms of partial wave amplitudes by splitting it up in terms of its angular momentum components \(J\):

(2)\[ A(\Omega) = \frac{1}{2q_i}\sum_J\left(2J+1\right) T^J(s) {D^J_{\lambda\mu}}^*\left(\phi,\theta,0\right) \]

with \(\lambda=\lambda_a-\lambda_b\) and \(\mu=\lambda_c-\lambda_d\) the helicity differences of the final and initial states \(ab,cd\).

The above sketch is just with one channel in mind, but the same holds true though for a number of channels \(n\), with the only difference that the \(T\) operator becomes a \(\boldsymbol{T}\)-matrix of rank \(n\).

Transition operator

The important point is that we have now expressed \(A\) in terms of an angular part (depending on \(\Omega\)) and a dynamical part \(\boldsymbol{T}\) that depends on the Mandelstam variable \(s\).

The dynamical part \(\boldsymbol{T}\) is usually called the transition operator. The reason is that it describes the interacting part of the scattering operator \(\boldsymbol{S}\), which describes the (complex) amplitude \(\langle f|\boldsymbol{S}|i\rangle\) of an initial state \(|i\rangle\) transitioning to a final state \(|f\rangle\). The scattering operator describes both the non-interacting amplitude and the transition amplitude, so it relates to the transition operator as:[1]

(3)\[ \boldsymbol{S} = \boldsymbol{I} + i\boldsymbol{T} \]

with \(\boldsymbol{I}\) the identity operator. With this in mind, there is an important restriction that the \(T\)-operator needs to comply with: unitarity. This means that \(\boldsymbol{S}\) should conserve probability, namely \(\boldsymbol{S}^\dagger\boldsymbol{S} = \boldsymbol{I}\).

K-matrix formalism

Now there is a trick to ensure unitarity of \(\boldsymbol{S}\). We can express \(\boldsymbol{S}\) in terms of an operator \(\boldsymbol{K}\) by applying a Cayley transformation:

(4)\[ \boldsymbol{S} = (\boldsymbol{I} + i\boldsymbol{K})(I - i\boldsymbol{K})^{-1} \]

Unitarity is conserved if \(K\) is real. Finally, the \(\boldsymbol{T}\)-matrix can be expressed in terms of \(\boldsymbol{K}\) as follows:

(5)\[ \boldsymbol{T} = \boldsymbol{K} \left(\boldsymbol{I} - i\boldsymbol{K}\right)^{-1} \]

Resonances

The challenge is now to choose a correct parametrization for the elements of \(\boldsymbol{K}\) so that it correctly describes the resonances we observe. There are several choices, but a common one is the following summation over the resonances \(R\):

(6)\[ K_{ij} = \sum_R\frac{g_{R,i}^*g_{R,j}}{m_R^2-m^2} \]

with \(g_{R,i}\) the residue functions that can be further expressed as

(7)\[ g_{R,i}=\gamma_{R,i}\sqrt{m_R\Gamma_R} \]

Implementation

The challenge is to generate a correct parametrization for an arbitrary number of coupled channels \(n\) and an arbitrary number of resonances \(n_R\). Our approach is to construct an \(n \times n\) sympy.Matrix with Symbols as its elements. We then use substitute these Symbols with certain parametrizations using subs(). In order to generate symbols for \(n_R\) resonances and \(n\) channels, we use indexed symbols.

This approach is less elegant and (theoretically) slower than using MatrixSymbols. That approach is explored in TR-007.

It would be nice to use a Symbol to represent the number of channels \(n\) and specify its value later.

n_channels = sp.Symbol("n", integer=True, positive=True)

Unfortunately, this does not work well in the Matrix class. We therefore set variables \(n\) to a specific int value and define some other Symbols for the rest of the implementation.[2] The value we choose in this example is n_channels=1, because we want to see if this reproduces a Breit-Wigner function.[3]

n_channels = 1
i, j, R, n_resonances = sp.symbols("i j R n_R", integer=True, negative=False)
m = sp.Symbol("m", real=True)
M = sp.IndexedBase("m", shape=(n_resonances,))
Gamma = sp.IndexedBase("Gamma", shape=(n_resonances,))
gamma = sp.IndexedBase("gamma", shape=(n_resonances, n_channels))

The parametrization of \(K_{ij}\) from Eq. (6) can be expressed as follows:

def Kij(
    m: sp.Symbol,
    M: sp.IndexedBase,
    Gamma: sp.IndexedBase,
    gamma: sp.IndexedBase,
    i: int,
    j: int,
    n_resonances: int | sp.Symbol,
) -> sp.Expr:
    g_i = gamma[R, i] * sp.sqrt(M[R] * Gamma[R])
    g_j = gamma[R, j] * sp.sqrt(M[R] * Gamma[R])
    parametrization = (g_i * g_j) / (M[R] ** 2 - m**2)
    return sp.Sum(parametrization, (R, 0, n_resonances - 1))

Show code cell source

Hide code cell source

n_R = sp.Symbol("n_R")
kij = Kij(m, M, Gamma, gamma, i, j, n_R)
Math("K_{ij} = " + f"{sp.latex(kij)}")

\[\displaystyle K_{ij} = \sum_{R=0}^{n_{R} - 1} \frac{{\Gamma}_{R} {\gamma}_{R,i} {\gamma}_{R,j} {m}_{R}}{- m^{2} + {m}_{R}^{2}}\]

We now define the \(\boldsymbol{K}\)-matrix in terms of a Matrix with IndexedBase instances as elements that can serve as Symbols. These Symbols will be substituted with the parametrization later. We could of course have inserted the parametrization directly, but this slows down matrix multiplication in the following steps.

K_symbol = sp.IndexedBase("K", shape=(n_channels, n_channels))
K = sp.Matrix([[K_symbol[i, j] for j in range(n_channels)] for i in range(n_channels)])
display(K_symbol[i, j], K)

\[\displaystyle {K}_{i,j}\]

\[\displaystyle \left[\begin{matrix}{K}_{0,0}\end{matrix}\right]\]

The \(\boldsymbol{T}\)-matrix can now be computed from Eq. (5):

T = K * (sp.eye(n_channels) - sp.I * K).inv()
T

\[\displaystyle \left[\begin{matrix}\frac{{K}_{0,0}}{- i {K}_{0,0} + 1}\end{matrix}\right]\]

Next, we need to substitute the elements \(K_{i,j}\) with the parametrization we defined above:

T_subs = T.subs({
    K[i, j]: Kij(m, M, Gamma, gamma, i, j, n_resonances)
    for i in range(n_channels)
    for j in range(n_channels)
})
T_subs

\[\displaystyle \left[\begin{matrix}\frac{\sum_{R=0}^{n_{R} - 1} \frac{{\Gamma}_{R} {\gamma}_{R,0}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}}}{- i \sum_{R=0}^{n_{R} - 1} \frac{{\Gamma}_{R} {\gamma}_{R,0}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}} + 1}\end{matrix}\right]\]

Warning

It is important to perform doit() after subs(), otherwise the Sum cannot be evaluated and there will be no warning of a failed substitution.

Now indeed, when taking \(n_R=1\), the resulting element from the \(\boldsymbol{T}\)-matrix looks like a Breit-Wigner function (compare relativistic_breit_wigner())!

Show code cell source

Hide code cell source

n_resonances_val = 1
rel_bw = T_subs[0, 0].subs(n_resonances, n_resonances_val).doit()
if n_resonances_val == 1 or n == 2:
    rel_bw = rel_bw.simplify()
rel_bw

\[\displaystyle - \frac{{\Gamma}_{0} {\gamma}_{0,0}^{2} {m}_{0}}{m^{2} + i {\Gamma}_{0} {\gamma}_{0,0}^{2} {m}_{0} - {m}_{0}^{2}}\]

Generalization

The above procedure has been condensed into a function that can handle an arbitrary number of resonances and an arbitrary number of channels.

def create_symbol_matrix(name: str, n: int) -> sp.Matrix:
    symbol = sp.IndexedBase("K", shape=(n, n))
    return sp.Matrix([[symbol[i, j] for j in range(n)] for i in range(n)])


def k_matrix(n_resonances: int, n_channels: int) -> sp.Matrix:
    # Define symbols
    m = sp.Symbol("m", real=True)
    M = sp.IndexedBase("m", shape=(n_resonances,))
    Gamma = sp.IndexedBase("Gamma", shape=(n_resonances,))
    gamma = sp.IndexedBase("gamma", shape=(n_resonances, n_channels))
    # Define K-matrix and T-matrix
    K = create_symbol_matrix("K", n_channels)
    T = K * (sp.eye(n_channels) - sp.I * K).inv()
    # Substitute elements
    return T.subs({
        K[i, j]: Kij(m, M, Gamma, gamma, i, j, n_resonances)
        for i in range(n_channels)
        for j in range(n_channels)
    })

Single channel, single resonance:

k_matrix(n_resonances=1, n_channels=1)[0, 0].doit().simplify()

\[\displaystyle - \frac{{\Gamma}_{0} {\gamma}_{0,0}^{2} {m}_{0}}{m^{2} + i {\Gamma}_{0} {\gamma}_{0,0}^{2} {m}_{0} - {m}_{0}^{2}}\]

Single channel, \(n_R\) resonances

k_matrix(n_resonances=sp.Symbol("n_R"), n_channels=1)[0, 0]

\[\displaystyle \frac{\sum_{R=0}^{n_{R} - 1} \frac{{\Gamma}_{R} {\gamma}_{R,0}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}}}{- i \sum_{R=0}^{n_{R} - 1} \frac{{\Gamma}_{R} {\gamma}_{R,0}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}} + 1}\]

Two channels, one resonance (Flatté function):

k_matrix(n_resonances=1, n_channels=2)[0, 0].doit().simplify()

\[\displaystyle - \frac{{\Gamma}_{0} {\gamma}_{0,0}^{2} {m}_{0}}{m^{2} + i {\Gamma}_{0} {\gamma}_{0,0}^{2} {m}_{0} + i {\Gamma}_{0} {\gamma}_{0,1}^{2} {m}_{0} - {m}_{0}^{2}}\]

Two channels, \(n_R\) resonances:

expr = k_matrix(n_resonances=sp.Symbol("n_R"), n_channels=2)[0, 0]
Math(sp.multiline_latex("", expr))

\[\begin{split}\displaystyle \begin{align*} \mathtt{\text{}} = & \frac{\left(i \sum_{R=0}^{n_{R} - 1} \frac{{\Gamma}_{R} {\gamma}_{R,1}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}} - 1\right) \sum_{R=0}^{n_{R} - 1} \frac{{\Gamma}_{R} {\gamma}_{R,0}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}}}{\left(\sum_{R=0}^{n_{R} - 1} \frac{{\Gamma}_{R} {\gamma}_{R,0}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}}\right) \sum_{R=0}^{n_{R} - 1} \frac{{\Gamma}_{R} {\gamma}_{R,1}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}} + i \sum_{R=0}^{n_{R} - 1} \frac{{\Gamma}_{R} {\gamma}_{R,0}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}} + i \sum_{R=0}^{n_{R} - 1} \frac{{\Gamma}_{R} {\gamma}_{R,1}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}} - \left(\sum_{R=0}^{n_{R} - 1} \frac{{\Gamma}_{R} {\gamma}_{R,0} {\gamma}_{R,1} {m}_{R}}{- m^{2} + {m}_{R}^{2}}\right)^{2} - 1} \\ & + \frac{i \left(\sum_{R=0}^{n_{R} - 1} \frac{{\Gamma}_{R} {\gamma}_{R,0} {\gamma}_{R,1} {m}_{R}}{- m^{2} + {m}_{R}^{2}}\right)^{2}}{- \left(\sum_{R=0}^{n_{R} - 1} \frac{{\Gamma}_{R} {\gamma}_{R,0}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}}\right) \sum_{R=0}^{n_{R} - 1} \frac{{\Gamma}_{R} {\gamma}_{R,1}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}} - i \sum_{R=0}^{n_{R} - 1} \frac{{\Gamma}_{R} {\gamma}_{R,0}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}} - i \sum_{R=0}^{n_{R} - 1} \frac{{\Gamma}_{R} {\gamma}_{R,1}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}} + \left(\sum_{R=0}^{n_{R} - 1} \frac{{\Gamma}_{R} {\gamma}_{R,0} {\gamma}_{R,1} {m}_{R}}{- m^{2} + {m}_{R}^{2}}\right)^{2} + 1} \end{align*}\end{split}\]

Visualization

Now, let’s use matplotlib, mpl_interactions, and symplot to visualize the \(\boldsymbol{K}\)-matrix for arbitrary \(n\) and \(n_R\).

TR-008 explains the need for symplot.substitute_indexed_symbols().

Show code cell source

Hide code cell source

def plot_k_matrix(
    n_channels: int,
    n_resonances: int,
    title: str = "",
) -> None:
    # Convert to Symbol: symplot cannot handle IndexedBase
    i = sp.Symbol("i", integer=True, negative=False)
    expr = k_matrix(n_resonances, n_channels)[i, i].doit()
    expr = symplot.substitute_indexed_symbols(expr)
    np_expr, sliders = symplot.prepare_sliders(expr, plot_symbol=m)
    symbol_to_arg = {symbol: arg for arg, symbol in sliders._arg_to_symbol.items()}

    # Set plot domain
    x_min, x_max = 1e-3, 3
    y_min, y_max = -0.5, +0.5

    plot_domain = np.linspace(x_min, x_max, num=500)
    x_values = np.linspace(x_min, x_max, num=160)
    y_values = np.linspace(y_min, y_max, num=80)
    X, Y = np.meshgrid(x_values, y_values)
    plot_domain_complex = X + Y * 1j

    # Set slider values and ranges
    m0_values = np.linspace(x_min, x_max, num=n_resonances + 2)
    m0_values = m0_values[1:-1]
    for R in range(n_resonances):
        for i in range(n_channels):
            sliders.set_ranges({
                "i": (0, n_channels - 1),
                f"m{R}": (0, 3, 100),
                f"Gamma{R}": (-1, 1, 100),
                Rf"\gamma_{{{R},{i}}}": (0, 2, 100),
            })
            sliders.set_values({
                f"m{R}": m0_values[R],
                f"Gamma{R}": (R + 1) * 0.1,
                Rf"\gamma_{{{R},{i}}}": 1 - 0.1 * R + 0.1 * i,
            })

    # Create interactive plots
    controls = Controls(**sliders)
    fig, (ax_2d, ax_3d) = plt.subplots(
        nrows=2,
        figsize=(8, 6),
        sharex=True,
        tight_layout=True,
    )

    fig.canvas.toolbar_visible = False
    fig.canvas.header_visible = False
    fig.canvas.footer_visible = False
    if not title:
        title = (
            Rf"${n_channels} \times {n_channels}$ $K$-matrix"
            f" with {n_resonances} resonances"
        )
    fig.suptitle(title)

    ax_2d.set_ylabel("$|T|^{2}$")
    ax_2d.set_yticks([])
    ax_3d.set_xlabel("Re $m$")
    ax_3d.set_ylabel("Im $m$")
    ax_3d.set_xticks([])
    ax_3d.set_yticks([])
    ax_3d.set_facecolor("white")

    ax_3d.axhline(0, linewidth=0.5, c="black", linestyle="dotted")

    # 2D plot
    def plot(channel: int):
        def wrapped(*args, **kwargs) -> sp.Expr:
            kwargs = {k: v.value for k, v in kwargs.items()}
            kwargs["i"] = channel
            return np.abs(np_expr(*args, **kwargs)) ** 2

        return wrapped

    for i in range(n_channels):
        iplt.plot(
            plot_domain,
            plot(i),
            ax=ax_2d,
            controls=controls,
            ylim="auto",
            label=f"channel {i}",
        )
    if n_channels > 1:
        ax_2d.legend(loc="upper right")
    mass_line_style = {
        "c": "red",
        "alpha": 0.3,
    }
    for name in controls.params:
        if not name.startswith("m"):
            continue
        iplt.axvline(sliders[name].value, ax=ax_2d, **mass_line_style)

    # 3D plot
    color_mesh = None
    resonances_indicators = []

    def plot3(*, z_cutoff, complex_rendering, **kwargs):
        nonlocal color_mesh
        Z = np_expr(plot_domain_complex, **kwargs)
        if complex_rendering == "imag":
            Z_values = Z.imag
            ax_title = "Re $T$"
        elif complex_rendering == "real":
            Z_values = Z.real
            ax_title = "Im $T$"
        elif complex_rendering == "abs":
            Z_values = np.abs(Z)
            ax_title = "$|T|$"
        else:
            raise NotImplementedError

        if n_channels == 1:
            ax_3d.set_title(ax_title)
        else:
            i = kwargs["i"]
            ax_3d.set_title(f"{ax_title}, channel {i}")

        if color_mesh is None:
            color_mesh = ax_3d.pcolormesh(X, Y, Z_values, cmap=plt.cm.coolwarm)
        else:
            color_mesh.set_array(Z_values)
        color_mesh.set_clim(vmin=-z_cutoff, vmax=+z_cutoff)

        if resonances_indicators:
            for R, (line, text) in enumerate(resonances_indicators):
                mass = kwargs[f"m{R}"]
                line.set_xdata(mass)
                text.set_x(mass + (x_max - x_min) * 0.008)
        else:
            for R in range(n_resonances):
                mass = kwargs[f"m{R}"]
                resonances_indicators.append(
                    (
                        ax_3d.axvline(mass, **mass_line_style),
                        ax_3d.text(
                            x=mass + (x_max - x_min) * 0.008,
                            y=0.95 * y_min,
                            s=f"$m_{R}$",
                            c="red",
                        ),
                    ),
                )

    # Create switch for imag/real/abs
    name = "complex_rendering"
    sliders._sliders[name] = ipywidgets.RadioButtons(
        options=["imag", "real", "abs"],
        description=R"\(s\)-plane plot",
    )
    sliders._arg_to_symbol[name] = name

    # Create cut-off slider for z-direction
    name = "z_cutoff"
    sliders._sliders[name] = ipywidgets.FloatSlider(
        value=1.5,
        min=0.01,
        max=10,
        step=0.1,
        description=R"\(z\)-cutoff",
    )
    sliders._arg_to_symbol[name] = name

    # Create GUI
    sliders_copy = dict(sliders)
    h_boxes = []
    for R in range(n_resonances):
        buttons = [
            sliders_copy.pop(f"m{R}"),
            sliders_copy.pop(f"Gamma{R}"),
        ]
        if n_channels == 1:
            dummy_name = symbol_to_arg[Rf"\gamma_{{{R},0}}"]
            buttons.append(sliders_copy.pop(dummy_name))
        h_box = ipywidgets.HBox(buttons)
        h_boxes.append(h_box)
    remaining_sliders = sorted(sliders_copy.values(), key=lambda s: s.description)
    if n_channels == 1:
        remaining_sliders.remove(sliders["i"])
    ui = ipywidgets.VBox(h_boxes + remaining_sliders)
    output = ipywidgets.interactive_output(plot3, controls=sliders)
    display(ui, output)

plot_k_matrix(n_resonances=3, n_channels=1)

Run this notebook on Binder or locally on Jupyter Lab to interactively modify the parameters.

plot_k_matrix(n_resonances=2, n_channels=2)

Run this notebook on Binder or locally on Jupyter Lab to interactively modify the parameters.

---

[1]

Some authors like [Chung et al., 1995] multiply the transition operator by a factor 2.

[2]

We use Symbols as indices, because that renders more nicely.

[3]

Of course, there is no need to work with matrices in this \(1 \times 1\) case. To keeps things general, however, we keep using matrices.