def interp(self, m):
# q = data_extra[self.outs[0]]["|q|"]
# a = self.a()
p = self.point_value()
zeros = tf.zeros_like(m)
ones = tf.ones_like(m)
def poly_i(i, xi):
tmp = zeros
for j in range(i - 1, i + 1):
if j < 0 or j > self.interp_N - 1:
continue
r = ones
for k in range(j, j + 2):
if k == i:
continue
r = r * (m - xi[k]) / (xi[i] - xi[k])
r = tf.where((m >= xi[j]) & (m < xi[j + 1]), r, zeros)
tmp = tmp + r
return tmp
h = tf.stack(
[poly_i(i, self.points) for i in range(1, self.interp_N - 1)],
axis=-1,
)
h = tf.stop_gradient(h)
p_r = tf.math.real(p)
p_i = tf.math.imag(p)
ret_r = tf.reduce_sum(h * p_r, axis=-1)
ret_i = tf.reduce_sum(h * p_i, axis=-1)
return tf.complex(ret_r, ret_i)
def get_matrix_interp1d3_v2(x, xi):
N = len(xi) - 1
zeros = tf.zeros_like(x)
ones = tf.ones_like(x)
# @pysnooper.snoop()
def poly_i(i):
tmp = zeros
x_i = (xi[i] + xi[i - 1]) / 2
for j in range(i - 1, i + 3):
if j < 0 or j > N - 1:
continue
r = ones
for k in range(j - 1, j + 3):
if k == i or k < 1 or k > N:
continue
x_k = (xi[k] + xi[k - 1]) / 2
r = r * (x - x_k) / (x_i - x_k)
r = tf.where(
(x >= (xi[j] + xi[j - 1]) / 2) & (x < (xi[j] + xi[j + 1]) / 2),
r,
zeros,
)
tmp = tmp + r
return tmp
h = tf.stack([poly_i(i) for i in range(1, N)], axis=-1)
b = tf.zeros_like(x)
return h, b
def get_amp(self, data, data_c, **kwargs):
mass = self.get_mass()
width = self.get_width()
if width is None:
return tf.ones_like(data["m"])
if not self.running_width:
ret = BW(data["m"], mass, width)
else:
q = data_c["|q|"]
q0 = data_c["|q0|"]
if self.bw_l is None:
decay = self.decay[0]
self.bw_l = min(decay.get_l_list())
ret = GS(
data["m"],
mass,
width,
q,
q0,
self.bw_l,
self.d,
self.c_daug2Mass,
self.c_daug3Mass,
)
return ret
def interp(self, m):
p = self.point_value()
ones = tf.ones_like(m)
zeros = tf.zeros_like(m)
p_r = tf.math.real(p)
p_i = tf.math.imag(p)
h, b = get_matrix_interp1d3_v2(m, self.points)
h = tf.stop_gradient(h)
f = lambda x: tf.reshape(
tf.matmul(tf.cast(h, x.dtype), tf.reshape(x, (-1, 1))), b.shape
) + tf.cast(b, x.dtype)
ret_r = f(p_r)
ret_i = f(p_i)
return tf.complex(ret_r, ret_i)
def spline_x_matrix(x, xi):
"""build matrix of x for spline interpolation"""
ones = tf.ones_like(x)
x2 = x * x
x3 = x2 * x
x_p = tf.stack([ones, x, x2, x3], axis=-1)
x = tf.expand_dims(x, axis=-1)
zeros = tf.zeros_like(x)
def poly_i(i):
cut = (x >= xi[i]) & (x < xi[i + 1])
return tf.where(cut, x_p, zeros)
xs = [poly_i(i) for i in range(len(xi) - 1)]
return tf.stack(xs, axis=-2)
def get_amp(self, data, data_c, **kwargs):
mass = self.get_mass()
width = self.get_width()
if width is None:
return tf.ones_like(data["m"])
if not self.running_width:
ret = BW(data["m"], mass, width)
else:
q2 = data_c["|q|2"]
q02 = data_c["|q0|2"]
if self.bw_l is None:
decay = self.decay[0]
self.bw_l = min(decay.get_l_list())
ret = BWR2(data["m"], mass, width, q2, q02, self.bw_l, self.d)
return ret
def get_matrix_interp1d3(x, xi):
N = len(xi) - 1
zeros = tf.zeros_like(x)
ones = tf.ones_like(x)
# @pysnooper.snoop()
def poly_i(i):
tmp = zeros
for j in range(i - 1, i + 3):
if j < 0 or j > N - 1:
continue
r = ones
for k in range(j - 1, j + 3):
if k == i or k < 0 or k > N:
continue
r = r * (x - xi[k]) / (xi[i] - xi[k])
r = tf.where((x >= xi[j]) & (x < xi[j + 1]), r, zeros)
tmp = tmp + r
return tmp
h = tf.stack([poly_i(i) for i in range(1, N)], axis=-1)
b = tf.zeros_like(x)
return h, b
def interp(self, m):
p = self.point_value()
ones = tf.ones_like(m)
zeros = tf.zeros_like(m)
def add_f(x, bl, br):
return tf.where((x > bl) & (x <= br), ones, zeros)
x_bin = tf.stack(
[
add_f(
m,
(self.points[i] + self.points[i + 1]) / 2,
(self.points[i + 1] + self.points[i + 2]) / 2,
) for i in range(self.interp_N - 2)
],
axis=-1,
)
p_r = tf.math.real(p)
p_i = tf.math.imag(p)
x_bin = tf.stop_gradient(x_bin)
ret_r = tf.reduce_sum(x_bin * p_r, axis=-1)
ret_i = tf.reduce_sum(x_bin * p_i, axis=-1)
return tf.complex(ret_r, ret_i)
def get_amp(self, data, _data_c=None, **kwargs):
mass = data["m"]
zeros = tf.zeros_like(mass)
ones = tf.ones_like(mass)
return tf.complex(ones, zeros)
