"""
This module provides functions to describe the lineshapes of the intermediate particles, namely generalized
Breit-Wigner function. Users can also define new lineshape using the function wrapper **regist_lineshape()**.
"""
import fractions
import functools
import math
import warnings
import sympy as sym
from .tensorflow_wrapper import tf
breit_wigner_dict = {}
[docs]def regist_lineshape(name=None):
"""
It will be used as a wrapper to define various Breit-Wigner functions
:param name: String name of the BW function
:return: A function used in a wrapper
"""
def fopt(f):
name_t = name
if name_t is None:
name_t = f.__name__
if name_t in breit_wigner_dict:
warnings.warn(
"override breit wigner function :", name
) # warning to users
breit_wigner_dict[name_t] = f # function
return f
return fopt
[docs]@regist_lineshape("one")
def one(*args):
"""
A uniform function
"""
return tf.complex(
1.0, 0.0
) # breit_wigner_dict["one"]==tf.complex(1.0,0.0)
[docs]@regist_lineshape("BW")
def BW(m, m0, g0, *args):
"""
Breit-Wigner function
.. math::
BW(m) = \\frac{1}{m_0^2 - m^2 - i m_0 \\Gamma_0 }
"""
m0 = tf.cast(m0, m.dtype)
gamma = tf.cast(g0, m.dtype)
x = m0 * m0 - m * m
y = m0 * gamma
s = x * x + y * y
ret = tf.complex(x / s, y / s)
return ret
[docs]@regist_lineshape("default") # 两个名字
@regist_lineshape("BWR") # BW with running width
def BWR(m, m0, g0, q, q0, L, d):
"""
Relativistic Breit-Wigner function (with running width). It's also set as the default lineshape.
.. math::
BW(m) = \\frac{1}{m_0^2 - m^2 - i m_0 \\Gamma(m)}
"""
gamma = Gamma(m, g0, q, q0, L, m0, d)
num = 1.0
m0 = tf.cast(m0, m.dtype)
x = m0 * m0 - m * m
y = m0 * gamma
s = x * x + y * y
ret = tf.complex(x / s, y / s)
return ret
# added by xiexh for GS model rho
[docs]def twoBodyCMmom(m_0, m_1, m_2):
"""relative momentum for 0 -> 1 + 2"""
M12S = m_1 + m_2
M12D = m_1 - m_2
if hasattr(M12S, "dtype"):
m_0 = tf.convert_to_tensor(m_0, dtype=M12S.dtype)
# m_eff = tf.where(m_0 > M12S, m_0, M12S)
# p = (m_eff - M12S) * (m_eff + M12S) * (m_eff - M12D) * (m_eff + M12D)
# if p is negative, which results from bad data, the return value is 0.0
# print("p", tf.where(p==0), m_0, m_1, m_2)
p = (m_0 - M12S) * (m_0 + M12S) * (m_0 - M12D) * (m_0 + M12D)
zeros = tf.zeros_like(m_0)
ret = tf.where(p > 0, tf.sqrt(p) / (2 * m_0), zeros)
return ret
[docs]def hFun(s, daug2Mass, daug3Mass):
_pi = 3.14159265359
_pi = tf.cast(_pi, s.dtype)
sm = daug2Mass + daug3Mass
sqrt_s = tf.sqrt(s)
k_s = twoBodyCMmom(tf.sqrt(s), daug2Mass, daug3Mass)
ret = (
(2.0 / _pi)
* (k_s / sqrt_s)
* tf.math.log((sqrt_s + 2.0 * k_s) / (sm), name="log")
)
ret = tf.cast(ret, s.dtype)
return ret
[docs]def dh_dsFun(s, daug2Mass, daug3Mass):
_pi = 3.14159265359
_pi = tf.cast(_pi, s.dtype)
k_s = twoBodyCMmom(tf.sqrt(s), daug2Mass, daug3Mass)
ret = hFun(s, daug2Mass, daug3Mass) * (
1.0 / (8.0 * tf.pow(k_s, 2)) - 1.0 / (2.0 * s)
) + 1.0 / (2.0 * _pi * s)
ret = tf.cast(ret, s.dtype)
return ret
[docs]def dFun(s, daug2Mass, daug3Mass):
_pi = 3.14159265359
_pi = tf.cast(_pi, s.dtype)
sm = daug2Mass + daug3Mass
sm24 = sm * sm / 4.0
m = tf.sqrt(s)
k_m2 = twoBodyCMmom(tf.sqrt(s), daug2Mass, daug3Mass)
ret = (
3.0
/ _pi
* sm24
/ tf.pow(k_m2, 2)
* tf.math.log((m + 2 * k_m2) / sm, name="log")
+ m / (2 * _pi * k_m2)
- sm24 * m / (_pi * tf.pow(k_m2, 3))
)
ret = tf.cast(ret, s.dtype)
return ret
[docs]def fsFun(s, m2, gam, daug2Mass, daug3Mass):
k_s = twoBodyCMmom(tf.sqrt(s), daug2Mass, daug3Mass)
k_Am2 = twoBodyCMmom(tf.sqrt(m2), daug2Mass, daug3Mass)
f = gam * m2 / tf.pow(k_Am2, 3)
f *= tf.pow(k_s, 2) * (
hFun(s, daug2Mass, daug3Mass) - hFun(m2, daug2Mass, daug3Mass)
) + (m2 - s) * tf.pow(k_Am2, 2) * dh_dsFun(m2, daug2Mass, daug3Mass)
f = tf.cast(f, s.dtype)
return f
# Gounaris-Sakurai model for rho
[docs]def GS(m, m0, g0, q, q0, L, d, c_daug2Mass=0.13957039, c_daug3Mass=0.1349768):
gamma = Gamma(m, g0, q, q0, L, m0, d)
c_daug2Mass = tf.cast(c_daug2Mass, m.dtype)
c_daug3Mass = tf.cast(c_daug3Mass, m.dtype)
D = 1.0 + dFun(m0 * m0, c_daug2Mass, c_daug3Mass) * g0 / m0
E = m0 * m0 - m * m + fsFun(m * m, m0 * m0, g0, c_daug2Mass, c_daug3Mass)
F = m0 * gamma
D /= E * E + F * F
ret = tf.complex(D * E, D * F)
return ret
# added by xiexh end
[docs]def BWR2(m, m0, g0, q2, q02, L, d):
"""
Relativistic Breit-Wigner function (with running width). Allow complex :math:`\\Gamma`.
.. math::
BW(m) = \\frac{1}{m_0^2 - m^2 - i m_0 \\Gamma(m)}
"""
gamma = Gamma2(m, g0, q2, q02, L, m0, d)
num = 1.0
m0 = tf.cast(m0, m.dtype)
x = tf.cast(m0 * m0 - m * m, gamma.dtype)
y = tf.cast(m0, gamma.dtype) * gamma
d = x - 1j * y
bw_x = tf.math.real(d)
bw_y = tf.math.imag(d)
bw_r2 = bw_x * bw_x + bw_y * bw_y
ret = tf.complex(bw_x / bw_r2, bw_y / bw_r2)
return ret
[docs]def BWR_normal(m, m0, g0, q2, q02, L, d):
"""
Relativistic Breit-Wigner function (with running width) with a normal factor.
.. math::
BW(m) = \\frac{\\sqrt{m_0 \\Gamma(m)}}{m_0^2 - m^2 - i m_0 \\Gamma(m)}
"""
gamma = Gamma2(m, g0, q2, q02, L, m0, d)
num = 1.0
m0 = tf.cast(m0, m.dtype)
x = tf.cast(m0 * m0 - m * m, gamma.dtype)
y = tf.cast(m0, gamma.dtype) * gamma
ret = tf.sqrt(tf.cast(m0, gamma.dtype) * gamma) / (x - 1j * y)
return ret
[docs]def Gamma(m, gamma0, q, q0, L, m0, d):
"""
Running width in the RBW
.. math::
\\Gamma(m) = \\Gamma_0 \\left(\\frac{q}{q_0}\\right)^{2L+1}\\frac{m_0}{m} B_{L}'^2(q,q_0,d)
"""
q0 = tf.cast(q0, q.dtype)
_epsilon = 1e-15
qq0 = tf.where(q0 > _epsilon, (q / q0) ** (2 * L + 1), 1.0)
mm0 = tf.cast(m0, m.dtype) / m
bp = Bprime(L, q, q0, d) ** 2
gammaM = gamma0 * qq0 * mm0 * tf.cast(bp, qq0.dtype)
return gammaM
[docs]def Gamma2(m, gamma0, q2, q02, L, m0, d):
"""
Running width in the RBW
.. math::
\\Gamma(m) = \\Gamma_0 \\left(\\frac{q}{q_0}\\right)^{2L+1}\\frac{m_0}{m} B_{L}'^2(q,q_0,d)
"""
q02 = tf.cast(q02, q2.dtype)
_epsilon = 1e-15
qq0 = q2 / q02
qq0 = tf.cast(qq0**L, tf.complex128) * tf.sqrt(
tf.cast(qq0, tf.complex128)
)
mm0 = tf.cast(m0, m.dtype) / m
z0 = q02 * d**2
z = q2 * d**2
bp = Bprime_polynomial(L, z0) / Bprime_polynomial(L, z)
gammaM = qq0 * tf.cast(gamma0 * bp * mm0, qq0.dtype)
return gammaM
[docs]def Bprime_q2(L, q2, q02, d):
"""
Blatt-Weisskopf barrier factors.
"""
q02 = tf.cast(q02, q2.dtype)
_epsilon = 1e-15
z0 = q02 * d**2
z = q2 * d**2
bp = Bprime_polynomial(L, z0) / Bprime_polynomial(L, z)
return tf.sqrt(tf.where(bp > 0, bp, 1.0))
[docs]def Bprime_num(L, q, d):
"""
The numerator (as well as the denominator) inside the square root in the barrier factor
"""
z = (q * d) ** 2
bp = Bprime_polynomial(L, z)
return tf.sqrt(bp)
[docs]def Bprime(L, q, q0, d):
"""
Blatt-Weisskopf barrier factors. E.g. the first three orders
=========== ===================================================
:math:`L` :math:`B_L'(q,q_0,d)`
=========== ===================================================
0 1
1 :math:`\\sqrt{\\frac{(q_0d)^2+1}{(qd)^2+1}}`
2 :math:`\\sqrt{\\frac{(q_0d)^4+3*(q_0d)^2+9}{(qd)^4+3*(qd)^2+9}}`
=========== ===================================================
:math:`d` is 3.0 by default.
"""
num = Bprime_num(L, q0, d)
denom = Bprime_num(L, q, d)
return tf.cast(num, denom.dtype) / denom
[docs]def barrier_factor(l, q, q0, d=3.0, axis=0): # cache q^l * B_l 只用于H里
"""
Barrier factor multiplied with :math:`q^L`, which is used as a combination in the amplitude expressions. The values
are cached for :math:`L` ranging from 0 to **l**.
"""
ret = []
for i in l:
tmp = q**i * tf.cast(Bprime(i, q, q0, d), q.dtype)
ret.append(tmp)
return tf.stack(ret)
[docs]def barrier_factor2(l, q, q0, d=3.0, axis=-1): # cache q^l * B_l 只用于H里
"""
???
"""
ret = []
for i in l:
tmp = q**i * tf.cast(Bprime(i, q, q0, d), q.dtype)
ret.append(tf.reshape(tmp, (-1, 1)))
return tf.concat(ret, axis=axis)
[docs]def Bprime_polynomial(l, z):
"""
It stores the Blatt-Weisskopf polynomial up to the fifth order (:math:`L=5`)
:param l: The order
:param z: The variable in the polynomial
:return: The calculated value
"""
coeff = {
0: [1.0],
1: [1.0, 1.0],
2: [1.0, 3.0, 9.0],
3: [1.0, 6.0, 45.0, 225.0],
4: [1.0, 10.0, 135.0, 1575.0, 11025.0],
5: [1.0, 15.0, 315.0, 6300.0, 99225.0, 893025.0],
}
l = int(l + 0.01)
if l not in coeff:
coeff[l] = [float(i) for i in get_bprime_coeff(l)]
# raise NotImplementedError
z = tf.convert_to_tensor(z)
cof = [tf.convert_to_tensor(i, z.dtype) for i in coeff[l]]
ret = tf.math.polyval(cof, z)
return ret
[docs]def reverse_bessel_polynomials(n, x):
"""Reverse Bessel polynomials.
.. math::
\\theta_{n}(x) = \\sum_{k=0}^{n} \\frac{(n+k)!}{(n-k)!k!} \\frac{x^{n-k}}{2^k}
"""
ret = 0
for k in range(n + 1):
c = fractions.Fraction(
math.factorial(n + k),
math.factorial(n - k) * math.factorial(k) * 2**k,
)
ret += c * x ** (n - k)
return ret
[docs]@functools.lru_cache()
def get_bprime_coeff(l):
"""The coefficients of polynomial in Bprime function.
.. math::
|\\theta_{l}(jw)|^2 = \\sum_{i=0}^{l} c_i w^{2 i}
"""
x = sym.Symbol("x")
theta = reverse_bessel_polynomials(l, x)
w = sym.Symbol("w", real=True)
Hjw = theta.subs({"x": sym.I * w})
Hjw2 = sym.Poly(Hjw * Hjw.conjugate(), w)
coeffs = Hjw2.as_dict()
ret = [coeffs.get((2 * l - 2 * i,), 0) for i in range(l + 1)]
return ret