Source code for ampform.dynamics.kmatrix
r"""Experimental, symbol :math:`\boldsymbol{K}`-matrix implementations.
.. seealso:: :doc:`/usage/dynamics/k-matrix`.
This module is an implementation of :doc:`compwa-org:report/005`,
:doc:`compwa-org:report/009`, and :doc:`compwa-org:report/010`. It works with classes to
keep the code organized and to enable caching of the matrix multiplications, but this
might change once these dynamics are implemented into the amplitude builder.
"""
from __future__ import annotations
import functools
from abc import ABC, abstractmethod
import sympy as sp
from ampform.dynamics import (
EnergyDependentWidth,
PhaseSpaceFactor,
PhaseSpaceFactorProtocol,
formulate_form_factor,
)
from ampform.sympy import create_symbol_matrix
[docs]class TMatrix(ABC):
[docs] @classmethod
@abstractmethod
def formulate(
cls, n_channels, n_poles, parametrize: bool = True, **kwargs
) -> sp.MutableDenseMatrix:
"""Formulate :math:`K`-matrix with its own parametrization."""
[docs]class RelativisticKMatrix(TMatrix):
@staticmethod
@functools.lru_cache(maxsize=None)
def _create_matrices(
n_channels, return_t_hat: bool = False
) -> tuple[sp.MutableDenseMatrix, sp.MutableDenseMatrix]:
# pylint: disable=no-member
rho = _create_rho_matrix(n_channels)
sqrt_rho: sp.MutableDenseMatrix = sp.sqrt(rho).doit()
sqrt_rho_conj = sp.conjugate(sqrt_rho)
k_matrix = create_symbol_matrix("K", n_channels, n_channels)
t_hat = k_matrix * (sp.eye(n_channels) - sp.I * rho * k_matrix).inv()
if return_t_hat:
return t_hat, k_matrix
t_matrix = sqrt_rho_conj * t_hat * sqrt_rho
return t_matrix, k_matrix
[docs] @classmethod
def formulate(
cls, n_channels, n_poles, parametrize: bool = True, **kwargs
) -> sp.MutableDenseMatrix:
r"""Implementation of :eq:`T-hat-in-terms-of-K-hat`.
Args:
n_channels: Number of coupled channels.
n_poles: Number of poles.
parametrize: Set to `False` if don't want to parametrize and only get
symbols for the matrix multiplication of :math:`\boldsymbol{K}` and
:math:`\boldsymbol{\rho}`.
return_t_hat: Set to `True` if you want to get the Lorentz-invariant
:math:`\boldsymbol{\hat{T}}`-matrix instead of the
:math:`\boldsymbol{T}`-matrix from Eq. :eq:`K-hat-and-T-hat`.
"""
return_t_hat: bool = kwargs.pop("return_t_hat", False)
t_matrix, k_matrix = cls._create_matrices(n_channels, return_t_hat)
if not parametrize:
return t_matrix
phsp_factor: PhaseSpaceFactorProtocol = kwargs.get(
"phsp_factor", PhaseSpaceFactor
)
s = sp.Symbol("s", nonnegative=True)
m_a = sp.IndexedBase("m_a", nonnegative=True)
m_b = sp.IndexedBase("m_b", nonnegative=True)
return t_matrix.xreplace(
{
k_matrix[i, j]: cls.parametrization(
i=i,
j=j,
s=s,
pole_position=sp.IndexedBase("m", nonnegative=True),
pole_width=sp.IndexedBase("Gamma", nonnegative=True),
m_a=m_a,
m_b=m_b,
residue_constant=sp.IndexedBase("gamma", nonnegative=True),
n_poles=n_poles,
pole_id=sp.Symbol("R", integer=True, positive=True),
angular_momentum=kwargs.get("angular_momentum", 0),
meson_radius=kwargs.get("meson_radius", 1),
phsp_factor=phsp_factor,
)
for i in range(n_channels)
for j in range(n_channels)
}
).xreplace(
{
sp.Symbol(f"rho{i}"): phsp_factor(s, m_a[i], m_b[i])
for i in range(n_channels)
}
)
[docs] @staticmethod
def parametrization( # pylint: disable=too-many-arguments, too-many-locals
i,
j,
s,
pole_position: sp.IndexedBase,
pole_width: sp.IndexedBase,
m_a: sp.IndexedBase,
m_b: sp.IndexedBase,
residue_constant: sp.IndexedBase,
n_poles,
pole_id,
angular_momentum=0,
meson_radius=1,
phsp_factor: PhaseSpaceFactorProtocol | None = None,
) -> sp.Expr:
def residue_function(pole_id, i) -> sp.Expr:
return residue_constant[pole_id, i] * sp.sqrt(
pole_position[pole_id]
* EnergyDependentWidth(
s=s,
mass0=pole_position[pole_id],
gamma0=pole_width[pole_id, i],
m_a=m_a[i],
m_b=m_b[i],
angular_momentum=angular_momentum,
meson_radius=meson_radius,
phsp_factor=phsp_factor,
)
)
g_i = residue_function(pole_id, i)
g_j = residue_function(pole_id, j)
parametrization = (g_i * g_j) / (pole_position[pole_id] ** 2 - s)
return sp.Sum(parametrization, (pole_id, 1, n_poles))
[docs]class NonRelativisticKMatrix(TMatrix):
@staticmethod
@functools.lru_cache(maxsize=None)
def _create_matrices(
n_channels,
) -> tuple[sp.MutableDenseMatrix, sp.MutableDenseMatrix]:
k_matrix = create_symbol_matrix("K", n_channels, n_channels)
t_matrix = k_matrix * (sp.eye(n_channels) - sp.I * k_matrix).inv()
return t_matrix, k_matrix
[docs] @classmethod
def formulate(
cls,
n_channels,
n_poles,
parametrize: bool = True,
**kwargs,
) -> sp.MutableDenseMatrix:
t_matrix, k_matrix = cls._create_matrices(n_channels)
if not parametrize:
return t_matrix
return t_matrix.xreplace(
{
k_matrix[i, j]: cls.parametrization(
i=i,
j=j,
s=sp.Symbol("s", nonnegative=True),
pole_position=sp.IndexedBase("m", nonnegative=True),
pole_width=sp.IndexedBase("Gamma", nonnegative=True),
residue_constant=sp.IndexedBase("gamma", nonnegative=True),
n_poles=n_poles,
pole_id=sp.Symbol("R", integer=True, positive=True),
)
for i in range(n_channels)
for j in range(n_channels)
}
)
[docs] @staticmethod
def parametrization( # pylint: disable=too-many-arguments
i,
j,
s,
pole_position: sp.IndexedBase,
pole_width: sp.IndexedBase,
residue_constant: sp.IndexedBase,
n_poles,
pole_id,
) -> sp.Expr:
def residue_function(pole_id, i) -> sp.Expr:
return residue_constant[pole_id, i] * sp.sqrt(
pole_position[pole_id] * pole_width[pole_id, i]
)
g_i = residue_function(pole_id, i)
g_j = residue_function(pole_id, j)
parametrization = (g_i * g_j) / (pole_position[pole_id] ** 2 - s)
return sp.Sum(parametrization, (pole_id, 1, n_poles))
[docs]class NonRelativisticPVector(TMatrix):
@staticmethod
@functools.lru_cache(maxsize=None)
def _create_matrices(
n_channels,
) -> tuple[sp.MutableDenseMatrix, sp.MutableDenseMatrix, sp.MutableDenseMatrix]:
k_matrix = create_symbol_matrix("K", m=n_channels, n=n_channels)
p_vector = create_symbol_matrix("P", m=n_channels, n=1)
f_vector = (sp.eye(n_channels) - sp.I * k_matrix).inv() * p_vector
return f_vector, k_matrix, p_vector
[docs] @classmethod
def formulate(
cls,
n_channels,
n_poles,
parametrize: bool = True,
**kwargs,
) -> sp.MutableDenseMatrix:
f_vector, k_matrix, p_vector = cls._create_matrices(n_channels)
if not parametrize:
return f_vector
s = sp.Symbol("s", nonnegative=True)
pole_position = sp.IndexedBase("m", nonnegative=True)
pole_width = sp.IndexedBase("Gamma", nonnegative=True)
residue_constant = sp.IndexedBase("gamma", nonnegative=True)
pole_id = sp.Symbol("R", integer=True, positive=True)
return f_vector.xreplace(
{
k_matrix[i, j]: NonRelativisticKMatrix.parametrization(
i=i,
j=j,
s=s,
pole_position=pole_position,
pole_width=pole_width,
residue_constant=residue_constant,
n_poles=n_poles,
pole_id=pole_id,
)
for i in range(n_channels)
for j in range(n_channels)
}
).xreplace(
{
p_vector[i]: cls.parametrization(
i=i,
s=sp.Symbol("s", nonnegative=True),
pole_position=pole_position,
pole_width=pole_width,
residue_constant=residue_constant,
beta_constant=sp.IndexedBase("beta", nonnegative=True),
n_poles=n_poles,
pole_id=pole_id,
)
for i in range(n_channels)
}
)
[docs] @staticmethod
def parametrization( # pylint: disable=too-many-arguments
i,
s,
pole_position: sp.IndexedBase,
pole_width: sp.IndexedBase,
residue_constant: sp.IndexedBase,
beta_constant: sp.IndexedBase,
n_poles,
pole_id,
) -> sp.Expr:
beta = beta_constant[pole_id]
gamma = residue_constant[pole_id, i]
mass = pole_position[pole_id]
width = pole_width[pole_id, i]
parametrization = beta * gamma * mass * width / (mass**2 - s)
return sp.Sum(parametrization, (pole_id, 1, n_poles))
[docs]class RelativisticPVector(TMatrix):
@staticmethod
@functools.lru_cache(maxsize=None)
def _create_matrices(
n_channels, return_f_hat: bool = False
) -> tuple[sp.MutableDenseMatrix, sp.MutableDenseMatrix, sp.MutableDenseMatrix]:
# pylint: disable=no-member
k_matrix = create_symbol_matrix("K", m=n_channels, n=n_channels)
rho = _create_rho_matrix(n_channels)
sqrt_rho: sp.MutableDenseMatrix = sp.sqrt(rho).doit()
sqrt_rho_conj: sp.MutableDenseMatrix = sp.conjugate(sqrt_rho)
k_matrix = create_symbol_matrix("K", n_channels, n_channels)
k_hat = sqrt_rho_conj.inv() * k_matrix * sqrt_rho.inv()
p_vector = create_symbol_matrix("P", m=n_channels, n=1)
f_hat = (sp.eye(n_channels) - sp.I * k_hat * rho).inv() * p_vector
if return_f_hat:
return f_hat, k_matrix, p_vector
f_vector = sqrt_rho * f_hat
return f_vector, k_matrix, p_vector
[docs] @classmethod
def formulate( # pylint: disable=too-many-locals
cls, n_channels, n_poles, parametrize: bool = True, **kwargs
) -> sp.MutableDenseMatrix:
r"""Implementation of :eq:`F-in-terms-of-P`.
Args:
n_channels: Number of coupled channels.
n_poles: Number of poles.
parametrize: Set to `False` if don't want to parametrize and only get
symbols for the matrix multiplication of :math:`\boldsymbol{K}` and
:math:`\boldsymbol{\rho}`.
return_f_hat: Set to `True` if you want to get theLorentz-invariant
:math:`\hat{F}`-vector instead of the :math:`T`-vector from Eq.
:eq:`invariant-vectors`.
"""
return_f_hat: bool = kwargs.pop("return_f_hat", False)
f_vector, k_matrix, p_vector = cls._create_matrices(n_channels, return_f_hat)
if not parametrize:
return f_vector
phsp_factor: PhaseSpaceFactorProtocol = kwargs.get(
"phsp_factor", PhaseSpaceFactor
)
s = sp.Symbol("s", nonnegative=True)
pole_position = sp.IndexedBase("m", nonnegative=True)
pole_width = sp.IndexedBase("Gamma", nonnegative=True)
residue_constant = sp.IndexedBase("gamma", nonnegative=True)
m_a = sp.IndexedBase("m_a", nonnegative=True)
m_b = sp.IndexedBase("m_b", nonnegative=True)
pole_id = sp.Symbol("R", integer=True, positive=True)
angular_momentum = kwargs.get("angular_momentum", 0)
meson_radius = kwargs.get("meson_radius", 1)
return (
f_vector.xreplace(
{
k_matrix[i, j]: RelativisticKMatrix.parametrization(
i=i,
j=j,
s=s,
pole_position=pole_position,
pole_width=pole_width,
m_a=m_a,
m_b=m_b,
residue_constant=residue_constant,
n_poles=n_poles,
pole_id=pole_id,
angular_momentum=angular_momentum,
meson_radius=meson_radius,
)
for i in range(n_channels)
for j in range(n_channels)
}
)
.xreplace(
{
p_vector[i]: cls.parametrization(
i=i,
s=s,
pole_position=pole_position,
pole_width=pole_width,
m_a=m_a,
m_b=m_b,
beta_constant=sp.IndexedBase("beta", nonnegative=True),
residue_constant=residue_constant,
n_poles=n_poles,
pole_id=pole_id,
angular_momentum=angular_momentum,
meson_radius=meson_radius,
)
for i in range(n_channels)
}
)
.xreplace(
{
sp.Symbol(f"rho{i}"): phsp_factor(s, m_a[i], m_b[i])
for i in range(n_channels)
}
)
)
[docs] @staticmethod
def parametrization( # pylint: disable=too-many-arguments, too-many-locals
i,
s,
pole_position: sp.IndexedBase,
pole_width: sp.IndexedBase,
m_a: sp.IndexedBase,
m_b: sp.IndexedBase,
beta_constant: sp.IndexedBase,
residue_constant: sp.IndexedBase,
n_poles,
pole_id,
angular_momentum=0,
meson_radius=1,
) -> sp.Expr:
beta = beta_constant[pole_id]
gamma = residue_constant[pole_id, i]
mass0 = pole_position[pole_id]
width = pole_width[pole_id, i]
form_factor = formulate_form_factor(
s, m_a[i], m_b[i], angular_momentum, meson_radius
)
return sp.Sum(
beta * gamma * mass0 * width * form_factor / (mass0**2 - s),
(pole_id, 1, n_poles),
)
def _create_rho_matrix(n_channels) -> sp.MutableDenseMatrix:
"""Create a phase space matrix with :code:`n_channels`.
>>> _create_rho_matrix(n_channels=2)
Matrix([
[rho0, 0],
[ 0, rho1]])
"""
rho_matrix: sp.MutableDenseMatrix = sp.zeros(n_channels, n_channels)
for i in range(n_channels):
rho_matrix[i, i] = sp.Symbol(f"rho{i}")
return rho_matrix