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)