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=None, **kwargs): mass = data["m"] zeros = tf.zeros_like(mass) a = tf.abs(self.a()) b = self.b() r = -tf.complex(a, b) * tf.complex(mass * mass, zeros) return tf.exp(r)
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 cal_monentum(m, ma, mb): mabp = ma + mb mabm = ma - mb s = m * m p2 = (s - mabp * mabp) * (s - mabm * mabm) / 4 / s zeros = tf.zeros_like(s) p_p = tf.complex(tf.sqrt(tf.abs(p2)), zeros) p_m = tf.complex(zeros, tf.sqrt(tf.abs(p2))) return tf.where(p2 > 0, p_p, p_m)
def interp(self, m): zeros = tf.zeros_like(m) p = self.point_value() xs = [] def poly_i(i): x = 1.0 for j in range(self.interp_N): if i == j: continue x = (x * (m - self.points[j]) / (self.points[i] - self.points[j])) return x xs = tf.stack([poly_i(i) for i in range(self.interp_N)], axis=-1) zeros = tf.zeros_like(xs) xs = tf.complex(xs, zeros) ret = tf.reduce_sum(xs[:, 1:-1] * p, axis=-1) return ret
def interp(self, m): zeros = tf.zeros_like(m) p = self.point_value() p_r = tf.math.real(p) p_i = tf.math.imag(p) xi_m = self.h_matrix x_m = spline_x_matrix(m, self.points) x_m = tf.expand_dims(x_m, axis=-1) m_xi = tf.reduce_sum(xi_m * x_m, axis=[-3, -2]) m_xi = tf.stop_gradient(m_xi) ret_r = tf.reduce_sum(tf.cast(m_xi, p_r.dtype) * p_r, axis=-1) ret_i = tf.reduce_sum(tf.cast(m_xi, p_i.dtype) * p_i, axis=-1) return tf.complex(ret_r, ret_i)
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 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 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 interp(self, m): # q = data_extra[self.outs[0]]["|q|"] # a = self.a() zeros = tf.zeros_like(m) p = tf.abs(self.point_value()) def add_f(x, bl, br, pl, pr): return tf.where( (x > bl) & (x <= br), (x - bl) / (br - bl) * (pr - pl) + pl, zeros, ) ret = [ add_f(m, self.points[i], self.points[i + 1], p[i], p[i + 1]) for i in range(self.interp_N - 1) ] return tf.complex(tf.reduce_sum(ret, axis=0), zeros)
def get_amp(self, *args, **kwargs): m = args[0]["m"] mass = self.get_mass() zeros = tf.zeros_like(m) delta_s = mass * mass - m * m m_c = mass / m rhos = [] for i, mab in enumerate(self.mass_list): ma, mb = mab pi = cal_monentum(m, ma, mb) # print(pi) m_rho_i = pi * tf.complex(zeros, self.g_value[i]() * m_c) rhos.append(m_rho_i) rho = sum(rhos) re = delta_s + tf.math.real(rho) im = tf.math.imag(rho) d = re * re + im * im ret = tf.complex(re / d, -im / d) return ret
def get_ls_amp(self, data, data_p, **kwargs): amp = super(ParticleDecay, self).get_ls_amp(data, data_p, **kwargs) a = self.core b = self.outs[0] c = self.outs[1] mass = a.get_mass() width = a.get_width() m = data_p[a]["m"] if width is None: ret = tf.zeros_like(m) ret = tf.complex(ret, ret) elif not a.running_width: ret = tf.reshape(BW(m, mass, width), (-1, 1)) else: q = data["|q|"] q0 = data["|q0|"] ret = [] for i in self.get_l_list(): bw = BWR(m, mass, width, q, q0, i, self.d) ret.append(tf.reshape(bw, (-1, 1))) ret = tf.concat(ret, axis=-1) return ret * amp
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) a = tf.abs(self.a()) return tf.complex(tf.exp(-a * mass), zeros)
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)