diff --git a/checks/resolution/plot_pull.py b/checks/resolution/plot_pull.py new file mode 100644 index 00000000..7792a931 --- /dev/null +++ b/checks/resolution/plot_pull.py @@ -0,0 +1,43 @@ +import glob +import json + +import matplotlib.pyplot as plt +import numpy as np + +all_params = [] +for i in glob.glob("results/*.json"): + with open(i) as f: + data = json.load(f) + all_params.append(data) + +with open("toy_params.json") as f: + ref_params = json.load(f) + +x = np.linspace(-5, 5, 1000) +for i in all_params[0]["error"]: + plt.clf() + v = np.array([j["value"][i] for j in all_params]) + e = np.array([j["error"][i] for j in all_params]) + pull = (v - ref_params[i]) / e + mu = np.mean(pull) + sigma = np.std(pull) + label = "$\\mu={:.2f}\\pm{:.2f}$\n$\\sigma={:.2f}\\pm{:.2f}$".format( + mu, + sigma / np.sqrt(pull.shape[0]), + sigma, + sigma / np.sqrt(2 * pull.shape[0]), + ) + plt.hist(pull, bins=20, range=(-5, 5), label="toy") + plt.plot( + x, + np.exp(-(((x - mu) / sigma) ** 2) / 2) + / np.sqrt(2 * np.pi) + / sigma + * 100 + / 20 + * 10, + label=label, + ) + plt.legend() + plt.title(i) + plt.savefig(f"results/{i}.png") diff --git a/checks/resolution/sample.py b/checks/resolution/sample.py index b44aac3f..b674bd4f 100644 --- a/checks/resolution/sample.py +++ b/checks/resolution/sample.py @@ -307,7 +307,7 @@ def main(): decay_chain = config.get_decay(False).get_decay_chain(results.particle) for name in ["data_rec", "bg_rec"]: - data = config.get_data("deta_rec") + data = config.get_data(name) if data is None: continue for i, toy in enumerate(data): @@ -321,11 +321,11 @@ def main(): config, decay_chain, toy, smear_method=results.method ) w = toy.get_weight() * w - save_name = config.data.dic["data"] + save_name = config.data.dic[name[:-4]] if isinstance(save_name, list): save_name = save_name[i] np.savetxt(save_name, np.stack(p4).reshape((-1, 4))) - save_name = config.data.dic["data_weight"] + save_name = config.data.dic[name[:-4] + "_weight"] if isinstance(save_name, list): save_name = save_name[i] np.savetxt(save_name, np.transpose(w).reshape((-1,))) diff --git a/docs/conf.py b/docs/conf.py index d24124f7..a1be622c 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -121,10 +121,12 @@ def gen_particle_model(): name_list = model_params[v]["name"] name = ", ".join(name_list) doc_i = model_params[v]["doc"] - particle_model_doc += ( + particle_model_head = ( f"\n{idx+1}. :code:`{name}`" - f" (`~{v.__module__}.{v.__qualname__}`)\n\n" + f" (`~{v.__module__}.{v.__qualname__}`)\n" ) + particle_model_doc += particle_model_head + particle_model_doc += "^" * (len(particle_model_head) - 2) + "\n\n" idx += 1 particle_model_doc += add_indent(doc_i) + "\n\n" @@ -170,10 +172,12 @@ def gen_decay_model(): name_list = model_params[v]["name"] name = ", ".join(name_list) doc_i = model_params[v]["doc"] - decay_model_doc += ( + decay_model_head = ( f"\n{idx+1}. :code:`{name}`" - f" (`~{v.__module__}.{v.__qualname__}`)\n\n" + f" (`~{v.__module__}.{v.__qualname__}`)\n" ) + decay_model_doc += decay_model_head + decay_model_doc += "^" * (len(decay_model_head) - 2) + "\n\n" idx += 1 decay_model_doc += add_indent(doc_i) + "\n\n" diff --git a/tf_pwa/amp/base.py b/tf_pwa/amp/base.py index 265193ee..03e5846c 100644 --- a/tf_pwa/amp/base.py +++ b/tf_pwa/amp/base.py @@ -466,6 +466,35 @@ def get_amp(self, data, _data_c=None, **kwargs): return tf.exp(r) +@regist_particle("poly") +class ParticlePoly(Particle): + """ + .. math:: + R(m) = \\sum c_i (m-m_0)^{n-i} + + lineshape when :math:`c_0=1, c_1=c_2=0` + + .. plot:: + + >>> import matplotlib.pyplot as plt + >>> plt.clf() + >>> from tf_pwa.utils import plot_particle_model + >>> axis = plot_particle_model("poly", params={"n_order": 2}, plot_params={"R_BC_c_1r": 0., "R_BC_c_2r": 0., "R_BC_c_1i": 0., "R_BC_c_2i": 0.}) + + """ + + def init_params(self): + self.n_order = getattr(self, "n_order", 3) + self.pi = self.add_var("c", shape=(self.n_order + 1,), is_complex=True) + self.pi.set_fix_idx(fix_idx=0, fix_vals=(1.0, 0.0)) + + def get_amp(self, data, _data_c=None, **kwargs): + mass = data["m"] - self.get_mass() + pi = list(self.pi()) + mass = tf.complex(mass, tf.zeros_like(mass)) + return tf.math.polyval(pi, mass) + + @regist_decay("particle-decay") class ParticleDecay(HelicityDecay): def get_ls_amp(self, data, data_p, **kwargs): @@ -697,3 +726,70 @@ def get_ls_amp(self, data, data_p, **kwargs): else: m_dep = g_ls return m_dep + + +@regist_decay("gls_reduce_h0") +class HelicityDecayReduceH0(HelicityDecay): + """ + decay model that remove helicity =0 for massless particles + """ + + def init_params(self): + self.d = 3.0 + + all_hel, remove_hel = self.get_helicity_list2() + ls = self.get_ls_list() + + self.g_ls = self.add_var( + "g_ls", + is_complex=True, + shape=(len(ls) - len(remove_hel),), + is_cp=True, + ) + try: + self.g_ls.set_fix_idx(fix_idx=0, fix_vals=(1.0, 0.0)) + except Exception as e: + print(e, self, self.get_ls_list()) + + all_matrix = self.get_cg_matrix() + print(all_hel, remove_hel) + + matrix = [] + for i, j in remove_hel: + idx_i = _spin_int(i + self.outs[0].J) + idx_j = _spin_int(j + self.outs[1].J) + matrix.append(all_matrix[:, idx_i, idx_j]) + # m g = h + matrix = np.stack(matrix) + # m_{zero,last} g_{last} + m_{zero, head} g_{head} = 0 + # m_{zero,last} g_{last} = - m_{zero,last}^{-1} m_{zero, head} g_{head} + matrix_inv = np.linalg.inv(matrix[:, -len(remove_hel) :]) + self.trans_matrix = ( + -np.dot(matrix_inv, matrix[:, : -len(remove_hel)]) + 0.0j + ) + + def get_helicity_list2(self): + all_hel = [] + for i in _spin_range(-self.outs[0].J, self.outs[0].J): + for j in _spin_range(-self.outs[1].J, self.outs[1].J): + if abs(i - j) <= self.core.J: + if self.p_break or (-i, -j) not in all_hel: + all_hel.append((i, j)) + reduce_item = [] + for hi in all_hel: + flag = False + for p, k in zip(self.outs, hi): + if p.get_mass() == 0 and k == 0: + flag = True + if flag: + reduce_item.append(hi) + return all_hel, reduce_item + + def get_g_ls(self, charge=1): + gls = self.g_ls(charge) + gls = tf.stack(gls) + gls_last = tf.linalg.matvec(self.trans_matrix, gls) + gls = list(tf.unstack(gls)) + list(tf.unstack(gls_last)) + if self.ls_index is None: + return tf.stack(gls) + return tf.stack([gls[k] for k in self.ls_index]) diff --git a/tf_pwa/amp/core.py b/tf_pwa/amp/core.py index 8cfd14b3..43f9a147 100644 --- a/tf_pwa/amp/core.py +++ b/tf_pwa/amp/core.py @@ -353,6 +353,8 @@ class Particle(BaseParticle, AmpBase): .. math:: R(m) = \\frac{1}{m_0^2 - m^2 - i m_0 \\Gamma(m)} + Argand diagram + .. plot:: >>> import matplotlib.pyplot as plt @@ -360,6 +362,15 @@ class Particle(BaseParticle, AmpBase): >>> from tf_pwa.utils import plot_particle_model >>> axis = plot_particle_model("BWR") + Pole position + + .. plot:: + + >>> import matplotlib.pyplot as plt + >>> plt.clf() + >>> from tf_pwa.utils import plot_pole_function + >>> axis = plot_pole_function("BWR") + """ def __init__( @@ -491,6 +502,14 @@ def get_sympy_dom(self, m, m0, g0, m1=None, m2=None, sheet=0): self.bw_l = min(decay.get_l_list()) return BWR_dom(m, m0, g0, self.bw_l, m1, m2) + def pole_function(self, sheet=0, modules="numpy"): + from tf_pwa.formula import create_numpy_function + + var = self.get_sympy_var() + f = self.get_sympy_dom(*var, sheet=sheet) + val = self.get_num_var() + return create_numpy_function(f, var[1:], val, var[0], modules=modules) + @regist_particle("x") class ParticleX(Particle): diff --git a/tf_pwa/formula.py b/tf_pwa/formula.py index 129d4b35..7eda5b7f 100644 --- a/tf_pwa/formula.py +++ b/tf_pwa/formula.py @@ -134,3 +134,10 @@ def _f(*y): return tf.complex(z[0], z[1]) return _f + + +def create_numpy_function(f, var, val, x, modules="numpy"): + f_var = _flatten(var) + f_val = _flatten(val) + f = f.subs(dict(zip(f_var, f_val))) + return sympy.lambdify(x, f, modules=modules) diff --git a/tf_pwa/tests/test_amp.py b/tf_pwa/tests/test_amp.py index 8e451b7d..09dc341f 100644 --- a/tf_pwa/tests/test_amp.py +++ b/tf_pwa/tests/test_amp.py @@ -452,3 +452,18 @@ def test_split_ls2(): data = cal_angle_from_momentum(p, dg) amp1 = amp(data) a(np.array([5.9, 6.0, 6.1])) + + +def test_gls_reduce_h0(): + a = get_particle("a", J=1, P=1, mass=6.0, width=0.02) + b = get_particle("b", J=1, P=-1, mass=0.0) + c = get_particle("c", J=0, P=-1, mass=0.1) + d = get_particle("d", J=0, P=-1) + t = get_particle("t", J=0, P=-1) + dec = get_decay(a, [b, c], model="gls_reduce_h0") + dec2 = get_decay(t, [a, d], p_break=True) + dg = DecayGroup([DecayChain([dec, dec2])]) + amp = AmplitudeModel(dg) + p = dict(zip([b, c, d], test_data[0])) + data = cal_angle_from_momentum(p, dg) + amp1 = amp(data) diff --git a/tf_pwa/utils.py b/tf_pwa/utils.py index fd7106f4..1d951c7e 100644 --- a/tf_pwa/utils.py +++ b/tf_pwa/utils.py @@ -413,3 +413,79 @@ def plot_particle_model( ax0.scatter(special_points, np.imag(at)) ax0.set_xlabel("mass") return [ax0, ax1, ax2, ax3] + + +def plot_pole_function( + model_name, params={}, plot_params={}, axis=None, **kwargs +): + import matplotlib.pyplot as plt + import numpy as np + + config = create_test_config(model_name, params, plot_params) + p = config.get_decay().get_particle("R_BC") + f = p.pole_function() + mx = np.linspace(0.2, 0.9 - 1e-12, 2000) + my = np.linspace(-0.1, 0.1, 2000) + a = np.abs(1 / f(mx + 1.0j * my[:, None])) ** 2 + if axis is None: + axis = plt.gca() + axis.contour(mx, my, a, linewidths=3) + axis.set_xlabel("Re$\\sqrt{s}$") + axis.set_ylabel("Im$\\sqrt{s}$") + return axis + + +def search_interval(px, cl=0.6826894921370859, xrange=(0, 1)): + """ + Search interval (a, b) that satisfly :math:`p(a)=p(b)` and + + .. math:: + \\frac{\\int_{a}^{b} p(x) dx}{\\int_{x_{min]}^{x_{max}} p(x) dx} = cl + + >>> x = np.linspace(-10, 10, 10000) + >>> a, b = search_interval(np.exp(-x**2/2), xrange=(-10, 10)) + >>> assert abs(a+1) < 0.01 + >>> assert abs(b-1) < 0.01 + + """ + n = px.shape[0] + px = px / np.sum(px) + idx = np.argsort(px)[::-1] + y = px[idx] + f = np.cumsum(y) + cut = f < cl + inner = idx[cut] + left = (np.min(inner) - 1) / px.shape[0] + right = (np.max(inner) + 1) / px.shape[0] + a, b = xrange + return a + (b - a) * left, a + (b - a) * right + + +def combine_asym_error(errs, N=10000): + """ + + combine asymmetry uncertanties using convolution + + >>> a, b = combine_asym_error([[-0.4, 0.4], 0.3]) + >>> assert abs(a+0.5) < 0.01 + >>> assert abs(b-0.5) < 0.01 + + """ + from scipy.signal import convolve + + errs = [[i, i] if isinstance(i, float) else i for i in errs] + errs = np.abs(np.stack(errs)) + max_range = np.max(np.sqrt(np.sum(errs**2, axis=0))) + xrange = -max_range * 10, max_range * 10 + + x = np.linspace(*xrange, N) + + y = np.exp(-(x**2) / 2 / np.where(x > 0, errs[-1, 1], errs[-1, 0]) ** 2) + y = y / np.sum(y) + for i in range(errs.shape[0] - 1): + tmp = np.exp( + -(x**2) / 2 / np.where(x > 0, errs[i, 1], errs[i, 0]) ** 2 + ) + tmp = tmp / np.sum(tmp) + y = convolve(y, tmp, mode="same") + return search_interval(y, xrange=xrange)