P-vector#

Hide code cell content
from __future__ import annotations

import os
import re
import warnings
from typing import TYPE_CHECKING

import matplotlib.pyplot as plt
import mpl_interactions.ipyplot as iplt
import numpy as np
import symplot
import sympy as sp
from ampform.dynamics import (
    BlattWeisskopfSquared,
    breakup_momentum_squared,
    coupled_width,
    phase_space_factor_complex,
)
from ampform.dynamics.decorator import (
    UnevaluatedExpression,
    create_expression,
    implement_doit_method,
)
from IPython.display import display
from ipywidgets import widgets as ipywidgets
from matplotlib import cm
from mpl_interactions.controller import Controls

if TYPE_CHECKING:
    from sympy.printing.latex import LatexPrinter

warnings.filterwarnings("ignore")
STATIC_WEB_PAGE = {"EXECUTE_NB", "READTHEDOCS"}.intersection(os.environ)

Physics#

As described in TR-005, the \(\boldsymbol{K}\)-matrix describes scattering processes of the type \(cd \to ab\). The \(P\)-vector approach is one of two generalizations for production processes of the type \(c \to ab\). For more details on this approach, [Chung et al., 1995] refers to [Aitchison, 1972].

If we take the production vector \(P\) to be:

(1)#\[ P_i = \sum_R \frac{\beta^0_R\,g_{R,i}(m)}{m_R^2-m^2} \]

and, in its invariant form,

(2)#\[ \hat{P}_i = \sum_R \frac{\beta^0_R\,g_{R,i}(m)}{\left(m_R^2-m^2\right)\sqrt{\rho_i}} \]

with \(g_{R,i}(m)\) given by Eq. (7) (possibly with coupled_width()), then the vector \(F\) describes the resulting amplitudes by

(3)#\[\begin{split} \begin{eqnarray} F & = & \left(\boldsymbol{I}-i\boldsymbol{K}\right)^{-1}P \\ \hat{F} & = & \left(\boldsymbol{I}-i\boldsymbol{\hat{K}\boldsymbol{\rho}}\right)^{-1}\hat{P} \end{eqnarray} \end{split}\]

with, from Eq. (4):

(4)#\[ \hat{\boldsymbol{K}} = \sqrt{\left(\boldsymbol{\rho}^\dagger\right)^{-1}} \boldsymbol{K} \sqrt{\boldsymbol{\rho}^{-1}} \]

Just like with the residue functions in TR-005 and TR-009, \(\beta\) is often expressed in terms of resonance mass and β€˜width’:

(5)#\[ \beta^0_R = \beta_R\sqrt{m_R\Gamma^0_R} \]

When in addition, we use a coupled_width(), the \(\hat{P}\)-vector becomes:

\[ \hat{P}_i = \sum_R \frac{\beta_R\gamma_{R,i}m_R\Gamma^0_R B_{R,i}(m)}{m_R^2-m^2} \]

with \(B_{R,i}(m)\) the ratio of Blatt-Weisskopf barrier factors (BlattWeisskopfSquared) for channel \(i\).

Implementation#

Hide code cell content
@implement_doit_method()
class PhaseSpaceFactor(UnevaluatedExpression):
    is_commutative = True

    def __new__(
        cls,
        s: sp.Symbol,
        m_a: sp.Symbol,
        m_b: sp.Symbol,
        i: int,
        **hints,
    ) -> PhaseSpaceFactor:
        return create_expression(cls, s, m_a, m_b, i, **hints)

    def evaluate(self) -> sp.Expr:
        s, m_a, m_b, *_ = self.args
        return phase_space_factor_complex(s, m_a, m_b)

    def _latex(self, printer: LatexPrinter, *args) -> str:
        s = printer._print(self.args[0])
        i = self.args[-1]
        return Rf"\rho_{{{i}}}({s})"


@implement_doit_method()
class CoupledWidth(UnevaluatedExpression):
    is_commutative = True

    def __new__(
        cls,
        s: sp.Symbol,
        mass0: sp.IndexedBase,
        gamma0: sp.IndexedBase,
        m_a: sp.IndexedBase,
        m_b: sp.IndexedBase,
        angular_momentum: int,
        R: int | sp.Symbol,
        i: int,
        **hints,
    ) -> CoupledWidth:
        return create_expression(
            cls, s, mass0, gamma0, m_a, m_b, angular_momentum, R, i, **hints
        )

    def evaluate(self) -> sp.Expr:
        s, mass0, gamma0, m_a, m_b, angular_momentum, R, i = self.args

        def phsp_factor(s, m_a, m_b):
            return PhaseSpaceFactor(s, m_a, m_b, i)

        return coupled_width(
            s,
            mass0[R],
            gamma0[R, i],
            m_a[i],
            m_b[i],
            angular_momentum=angular_momentum,
            meson_radius=1,
            phsp_factor=phsp_factor,
        )

    def _latex(self, printer: LatexPrinter, *args) -> str:
        s = printer._print(self.args[0])
        R = self.args[-2]
        i = self.args[-1]
        return Rf"{{\Gamma_{{{R},{i}}}}}({s})"
def Pi_relativistic(
    m: sp.Symbol,
    M: sp.IndexedBase,
    Gamma: sp.IndexedBase,
    gamma: sp.IndexedBase,
    beta: sp.IndexedBase,
    m_a: sp.IndexedBase,
    m_b: sp.IndexedBase,
    R: int | sp.Symbol,
    i: int,
    n_resonances: int | sp.Symbol,
    angular_momentum: int | sp.Symbol = 0,
) -> sp.Expr:
    q_squared = breakup_momentum_squared(m**2, m_a[i], m_b[i])
    ff2 = BlattWeisskopfSquared(z=q_squared, angular_momentum=angular_momentum)
    parametrization = (beta[R] * gamma[R, i] * M[R] * Gamma[R, i] * sp.sqrt(ff2)) / (
        M[R] ** 2 - m**2
    )
    return sp.Sum(parametrization, (R, 0, n_resonances - 1))


def Kij_relativistic(
    m: sp.Symbol,
    M: sp.IndexedBase,
    Gamma: sp.IndexedBase,
    gamma: sp.IndexedBase,
    m_a: sp.IndexedBase,
    m_b: sp.IndexedBase,
    R: sp.IndexedBase,
    i: int,
    j: int,
    n_resonances: int | sp.Symbol,
    angular_momentum: int | sp.Symbol = 0,
) -> sp.Expr:
    def residue_function(i):
        return gamma[R, i] * sp.sqrt(
            M[R] * CoupledWidth(m**2, M, Gamma, m_a, m_b, angular_momentum, R, i)
        )

    g_i = residue_function(i)
    g_j = residue_function(j)
    parametrization = (g_i * g_j) / (M[R] ** 2 - m**2)
    return sp.Sum(parametrization, (R, 0, n_resonances - 1))


def f_vector(
    n_resonances: int,
    n_channels: int,
    angular_momentum: int | sp.Symbol = 0,
) -> sp.Matrix:
    # Define symbols
    R = sp.Symbol("R")
    m = sp.Symbol("m", real=True)
    M = sp.IndexedBase("m", shape=(n_resonances,))
    Gamma = sp.IndexedBase("Gamma", shape=(n_resonances, n_channels))
    gamma = sp.IndexedBase("gamma", shape=(n_resonances, n_channels))
    beta = sp.IndexedBase("beta", shape=(n_resonances,))
    m_a = sp.IndexedBase("m_a", shape=(n_channels,))
    m_b = sp.IndexedBase("m_b", shape=(n_channels,))
    # Define phase space matrix
    rho = sp.zeros(n_channels, n_channels)
    for i in range(n_channels):
        rho[i, i] = PhaseSpaceFactor(m**2, m_a[i], m_b[i], i)
    # Define P-vector, K-matrix and T-matrix
    P = create_symbol_matrix("P", n_channels, 1)
    K = create_symbol_matrix("K", n_channels, n_channels)
    F = (sp.eye(n_channels) - sp.I * K * rho).inv() * P
    # Substitute elements
    return F.subs({
        K[i, j]: Kij_relativistic(
            m=m,
            M=M,
            Gamma=Gamma,
            gamma=gamma,
            m_a=m_a,
            m_b=m_b,
            i=i,
            j=j,
            R=R,
            n_resonances=n_resonances,
            angular_momentum=angular_momentum,
        )
        for i in range(n_channels)
        for j in range(n_channels)
    }).subs({
        P[i]: Pi_relativistic(
            m=m,
            M=M,
            Gamma=Gamma,
            gamma=gamma,
            beta=beta,
            i=i,
            m_a=m_a,
            m_b=m_b,
            R=R,
            n_resonances=n_resonances,
            angular_momentum=angular_momentum,
        )
        for i in range(n_channels)
    })


def create_symbol_matrix(name: str, m: int, n: int) -> sp.Matrix:
    symbol = sp.IndexedBase(name, shape=(m, n))
    return sp.Matrix([[symbol[i, j] for j in range(n)] for i in range(m)])
L = sp.Symbol("L", integer=True)
symplot.partial_doit(
    f_vector(n_resonances=1, n_channels=1, angular_momentum=L)[0, 0], sp.Sum
)
\[\displaystyle \frac{\sqrt{B_{L}^2\left(\frac{\left(m^{2} - \left({m_{a}}_{0} - {m_{b}}_{0}\right)^{2}\right) \left(m^{2} - \left({m_{a}}_{0} + {m_{b}}_{0}\right)^{2}\right)}{4 m^{2}}\right)} {\Gamma}_{0,0} {\beta}_{0} {\gamma}_{0,0} {m}_{0}}{\left(1 - \frac{i {\Gamma_{0,0}}(m^{2}) {\gamma}_{0,0}^{2} {m}_{0} \rho_{0}(m^{2})}{- m^{2} + {m}_{0}^{2}}\right) \left(- m^{2} + {m}_{0}^{2}\right)}\]

Visualization#

def plot_f_vector(
    n_channels: int,
    n_resonances: int,
    angular_momentum: int | sp.Symbol = 0,
    title: str = "",
) -> None:
    # Convert to Symbol: symplot cannot handle
    m = sp.Symbol("m", real=True)
    epsilon = sp.Symbol("epsilon", real=True)
    i = sp.Symbol("i", integer=True, negative=False)
    expr = f_vector(n_resonances, n_channels, angular_momentum=angular_momentum).doit()[
        i, 0
    ]
    expr = symplot.substitute_indexed_symbols(expr)
    expr = expr.subs(m, m + epsilon * sp.I)
    np_expr, sliders = symplot.prepare_sliders(expr, 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]

    def set_default_values():
        if "L" in sliders:
            sliders.set_ranges(L=(0, 8))
        sliders.set_ranges(
            i=(0, n_channels - 1),
            epsilon=(y_min * 0.2, y_max * 0.2, 0.01),
        )
        for R in range(n_resonances):
            for i in range(n_channels):
                sliders.set_ranges({
                    f"m{R}": (0, 3, 100),
                    f"beta{R}": (0, 5, 0.1),
                    Rf"\Gamma_{{{R},{i}}}": (-5, +5, 100),
                    Rf"\gamma_{{{R},{i}}}": (0, 20, 100),
                    f"m_a{i}": (0, 1, 0.01),
                    f"m_b{i}": (0, 1, 0.01),
                })
                sliders.set_values({
                    f"m{R}": m0_values[R],
                    f"beta{R}": 1,
                    Rf"\Gamma_{{{R},{i}}}": 3 * (0.4 + R * 0.2 - i * 0.3),
                    Rf"\gamma_{{{R},{i}}}": 0.2 * (10 - R + i),
                    f"m_a{i}": (i + 1) * 0.25,
                    f"m_b{i}": (i + 1) * 0.25,
                })

    set_default_values()

    # Create interactive plots
    controls = Controls(**sliders)
    nrows = 2  # set to 3 for imag+real
    fig, axes = plt.subplots(
        nrows=nrows,
        figsize=(8, nrows * 3.0),
        sharex=True,
        tight_layout=True,
    )
    fig.canvas.toolbar_visible = False
    fig.canvas.header_visible = False
    fig.canvas.footer_visible = False
    for ax in axes:
        ax.set_xlim(x_min, x_max)
    if not title:
        title = f"{n_channels}-channel $F$-vector with {n_resonances} resonances"
    fig.suptitle(title)

    # 2D plot
    axes[0].set_ylabel("$|T|^{2}$")
    axes[0].set_yticks([])

    def plot(channel: int):
        def wrapped(*args, **kwargs) -> sp.Expr:
            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=axes[0],
            controls=controls,
            ylim="auto",
            label=f"channel {i}",
        )
    if n_channels > 1:
        axes[0].legend(loc="upper right")
    mass_line_style = {
        "c": "red",
        "alpha": 0.3,
    }
    for name in controls.params:
        if not re.match(r"^m[0-9]+$", name):
            continue
        iplt.axvline(controls[name], ax=axes[0], **mass_line_style)

    # 3D plot
    def plot3(**kwargs):
        z_cutoff = kwargs.pop("z_cutoff")
        epsilon = kwargs["epsilon"]
        kwargs["epsilon"] = 0
        imag_real = kwargs.pop("imag_real")
        Z = np_expr(plot_domain_complex, **kwargs)
        if imag_real == "imag":
            Z_values = Z.imag
            ax_title = "Re $T$"
        elif imag_real == "real":
            Z_values = Z.real
            ax_title = "Im $T$"
        elif imag_real == "abs":
            Z_values = np.abs(Z)
            ax_title = "$|T|$"
        else:
            raise NotImplementedError
        for ax in axes[1:]:
            ax.clear()
        axes[-1].pcolormesh(
            X, Y, Z_values, cmap=cm.coolwarm, vmin=-z_cutoff, vmax=+z_cutoff
        )
        i = kwargs["i"]
        if n_channels == 1:
            axes[-1].set_title(ax_title)
        else:
            axes[-1].set_title(f"{ax_title}, channel {i}")
        for ax in axes[1:]:
            ax.axhline(0, linewidth=0.5, c="black", linestyle="dotted")
            if epsilon != 0.0:
                ax.axhline(
                    epsilon,
                    linewidth=0.5,
                    c="blue",
                    linestyle="dotted",
                    label=R"$\epsilon$",
                )
                axes[-1].text(
                    x=x_min + 0.008,
                    y=epsilon + 0.01,
                    s=R"$\epsilon$",
                    c="blue",
                )
            for R in range(n_resonances):
                mass = kwargs[f"m{R}"]
                ax.axvline(mass, **mass_line_style)
            if "m_a0" in kwargs:
                colors = cm.plasma(np.linspace(0, 1, n_channels))
                for i, color in enumerate(colors):
                    m_a = kwargs[f"m_a{i}"]
                    m_b = kwargs[f"m_b{i}"]
                    s_thr = m_a + m_b
                    ax.axvline(s_thr, c=color, linestyle="dotted")
                    ax.text(
                        x=s_thr,
                        y=0.95 * y_min,
                        s=f"$m_{{a{i}}}+m_{{b{i}}}$",
                        c=color,
                        rotation=-90,
                    )
                    m_diff = m_a - m_b
                    x_offset = (x_max - x_min) * 0.015
                    if m_diff > x_offset + 0.01 and s_thr - abs(m_diff) > x_offset:
                        ax.axvline(
                            m_diff,
                            c=color,
                            linestyle="dashed",
                            alpha=0.5,
                        )
                        ax.text(
                            x=m_diff - x_offset,
                            y=0.95 * y_min,
                            s=f"$m_{{a{i}}}-m_{{b{i}}}$",
                            c=color,
                            rotation=+90,
                        )
            ax.set_ylabel("Im $m$")
            ax.set_xticks([])
            ax.set_yticks([])
            ax.set_facecolor("white")
        for R in range(n_resonances):
            mass = kwargs[f"m{R}"]
            axes[-1].text(
                x=mass + (x_max - x_min) * 0.008,
                y=0.95 * y_min,
                s=f"$m_{R}$",
                c="red",
            )
        axes[-1].set_xlabel("Re $m$")
        fig.canvas.draw_idle()

    # Create switch for imag/real/abs
    name = "imag_real"
    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.IntSlider(
        value=10,
        min=+1,
        max=+50,
        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}")]
        if n_channels == 1:
            buttons.append(sliders_copy.pop(symbol_to_arg[Rf"\Gamma_{{{R},0}}"]))
            buttons.append(sliders_copy.pop(symbol_to_arg[Rf"\gamma_{{{R},0}}"]))
        h_box = ipywidgets.HBox(buttons)
        h_boxes.append(h_box)
    remaining_sliders = sorted(
        sliders_copy.values(), key=lambda s: (str(type(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_f_vector(
    n_resonances=2,
    n_channels=1,
    angular_momentum=L,
    title="Relativistic $F$-vector, single channel",
)