diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f39758c7..6ada1d81 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -91,6 +91,7 @@ jobs: run: | conda install --file requirements-min.txt -y python -m pip install -e . --no-deps + conda install pylint -y conda install pre-commit -c conda-forge -y pre-commit install pre-commit run -a @@ -117,6 +118,7 @@ jobs: run: | conda install --file tensorflow_2_6_requirements.txt -c conda-forge -y python -m pip install -e . --no-deps + conda install pylint -y conda install pre-commit -c conda-forge -y pre-commit install pre-commit run -a diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 910f6883..6b8f26af 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -50,10 +50,13 @@ repos: hooks: - id: isort - - repo: https://github.com/pre-commit/mirrors-pylint - rev: v2.5.0 # Use the sha / tag you want to point at + - repo: local hooks: - id: pylint + name: pylint + entry: pylint + language: system + types: [python] args: ["--rcfile=.pylintrc", "--score=no"] - repo: https://github.com/myint/rstcheck diff --git a/setup.cfg b/setup.cfg index 2ca1ccab..ad54bf5e 100644 --- a/setup.cfg +++ b/setup.cfg @@ -94,6 +94,7 @@ test = dev = %(doc)s %(test)s + pylint pre-commit all = %(dev)s diff --git a/tf_pwa/amp/core.py b/tf_pwa/amp/core.py index 22074bee..1e83c08d 100644 --- a/tf_pwa/amp/core.py +++ b/tf_pwa/amp/core.py @@ -1409,15 +1409,48 @@ def get_angle_amp(self, data): @functools.lru_cache() def get_swap_factor(self, key): factor = 1.0 + used = [] for i, j in zip(self.identical_particles, key[1]): p = self.get_particle(i[0]) if int(p.J * 2) % 2 == 0: continue for m, n in zip(i, j): + if (m, n) in used or (n, m) in used: + continue + used.append((m, n)) if m != n: factor *= -1.0 return factor + @functools.lru_cache() + def get_id_swap_transpose(self, key, n): + _, change = key + # print(key) + old_order = [str(i) for i in self.outs] + trans = [] + for i, j in zip(self.identical_particles, change): + for k, l in zip(i, j): + trans.append((k, l)) + trans = tuple(trans) + return self.get_swap_transpose(trans, n) + + @functools.lru_cache() + def get_swap_transpose(self, trans, n): + trans = dict(trans) + # print(trans) + tmp = {v: k for k, v in trans.items()} + tmp.update(trans) + trans = tmp + # print(trans) + old_order = [str(i) for i in self.outs] + new_order = [] + for i in old_order: + new_order.append(trans.get(i, i)) + index_map = {k: i for i, k in enumerate(new_order)} + trans_order = [index_map[str(i)] for i in self.outs] + diff = n - len(trans_order) + return [i for i in range(diff)] + [i + diff for i in trans_order] + def get_amp2(self, data): amp = self.get_amp(data) id_swap = data.get("id_swap", {}) @@ -1426,6 +1459,9 @@ def get_amp2(self, data): factor = self.get_swap_factor(k) amp_swap = factor * self.get_amp(new_data) # print(k, amp, amp_swap) + swap_index = self.get_id_swap_transpose(k, len(amp_swap.shape)) + # print(swap_index) + amp_swap = tf.transpose(amp_swap, swap_index) amp = amp + amp_swap return amp @@ -1441,15 +1477,26 @@ def get_amp3(self, data): name_map = {str(i): i for i in self.outs} frac = 1.0 same_particle = [] + change = [] for a, b in cg: for i, j in zip(a, b): if i == j: same_particle.append(i) frac = frac * getattr(name_map[i], "C", -1) + else: + change.append((i, j)) + transpose = self.get_swap_transpose( + tuple(change), len(amp_swap.shape) + ) p_reverse = [Ellipsis] + [ slice(None, None, -1) for i in range(len(amp_swap.shape) - 1) ] - amp = amp + amp_swap.__getitem__(p_reverse) * frac + + amp = ( + amp + + tf.transpose(amp_swap, transpose).__getitem__(p_reverse) + * frac + ) return amp def sum_amp(self, data, cached=True): diff --git a/tf_pwa/amp/interpolation.py b/tf_pwa/amp/interpolation.py index b526b355..68c7f519 100644 --- a/tf_pwa/amp/interpolation.py +++ b/tf_pwa/amp/interpolation.py @@ -85,6 +85,43 @@ def get_bin_index(self, m): return bin_idx +@register_particle("linear_npy") +class InterpLinearNpy(InterpolationParticle): + def __init__(self, *args, **kwargs): + self.input_file = kwargs.get("file") + self.data = np.load(self.input_file) + points = self.data[:, 0] + kwargs["points"] = points + super().__init__(*args, **kwargs) + + def init_params(self): + pass + + def get_point_values(self): + v_r = np.concatenate([[0.0], self.data[:, 1], [0.0]]) + v_i = np.concatenate([[0.0], self.data[:, 1], [0.0]]) + return self.data[:, 0], v_r, v_i + + def interp(self, m): + x, p_r, p_i = self.get_point_values() + bin_idx = tf.raw_ops.Bucketize(input=m, boundaries=x) + bin_idx = (bin_idx + len(self.bound)) % len(self.bound) + ret_r_l = tf.gather(p_r[1:], bin_idx) + ret_i_l = tf.gather(p_r[1:], bin_idx) + ret_r_r = tf.gather(p_r[:-1], bin_idx) + ret_i_r = tf.gather(p_r[:-1], bin_idx) + delta = np.concatenate( + [[1.0], self.data[1:, 1] - self.data[:-1, 1], [1.0]] + ) + x_left = np.concatenate([[x[0] - 1], x]) + delta = tf.gather(delta, bin_idx) + x_left = tf.gather(x_left, bin_idx) + step = (m - x_left) / delta + a = step * (ret_r_l - ret_r_r) + b = step * (ret_i_l - ret_i_r) + return tf.complex(a, b) + + @register_particle("interp") class Interp(InterpolationParticle): """linear interpolation for real number""" diff --git a/tf_pwa/cal_angle.py b/tf_pwa/cal_angle.py index d5f1971c..6c898166 100644 --- a/tf_pwa/cal_angle.py +++ b/tf_pwa/cal_angle.py @@ -343,6 +343,60 @@ def cal_helicity_angle( return ret +def aligned_angle_ref_rule1(decay_group, decay_chain_struct, decay_data): + # calculate aligned angle of final particles in each decay chain + set_x = {} # reference particles + ref_matrix = {} + + # for particle from a the top rest frame + for idx, decay_chain in enumerate(decay_chain_struct): + for decay in decay_chain: + if decay.core == decay_group.top: + for i in decay.outs: + if (i not in set_x) and (i in decay_group.outs): + part_data2 = decay_data[decay_chain][decay] + set_x[i] = (decay_chain, part_data2[i]) + ref_matrix[i] = decay_chain + # or in the first chain + for i in decay_group.outs: + if i not in set_x: + decay_chain = next(iter(decay_chain_struct)) + for decay in decay_chain: + for j in decay.outs: + if i == j: + part_data2 = decay_data[decay_chain][decay] + set_x[i] = (decay_chain, part_data2[i]) + ref_matrix[i] = decay_chain + + ref_matrix_final = {} + for i in decay_group.outs: + ref_matrix_final[i] = { + "b_matrix": decay_data[ref_matrix[i]]["b_matrix"][i], + "r_matrix": decay_data[ref_matrix[i]]["r_matrix"][i], + } + + return set_x, ref_matrix_final + + +def aligned_angle_ref_rule2(decay_group, decay_chain_struct, decay_data): + # calculate aligned angle of final particles in each decay chain + set_x = {} # reference particles + ref_matrix = {} + + ref_matrix_final = {} + for i in decay_group.outs: + set_x[i] = ( + None, + {"x": np.array([[1.0, 0, 0]]), "z": np.array([[0.0, 0, 1]])}, + ) + ref_matrix_final[i] = { + "b_matrix": SU2M([[1, 0], [0, 1]]), + "r_matrix": SU2M([[1, 0], [0, 1]]), + } + + return set_x, ref_matrix_final + + def cal_angle_from_particle( data, decay_group: DecayGroup, @@ -350,6 +404,7 @@ def cal_angle_from_particle( random_z=True, r_boost=True, final_rest=True, + align_ref=None, # "center_mass", ): """ Calculate helicity angle for particle momentum, add aligned angle. @@ -376,39 +431,29 @@ def cal_angle_from_particle( for i in decay_chain_struct: data_i = cal_helicity_angle(data, i, base_z=base_z) decay_data[i] = data_i + if align_ref == "center_mass": + set_x, ref_matrix_final = aligned_angle_ref_rule2( + decay_group, decay_chain_struct, decay_data + ) + else: + set_x, ref_matrix_final = aligned_angle_ref_rule1( + decay_group, decay_chain_struct, decay_data + ) - # calculate aligned angle of final particles in each decay chain - set_x = {} # reference particles - ref_matrix = {} - # for particle from a the top rest frame - for idx, decay_chain in enumerate(decay_chain_struct): - for decay in decay_chain: - if decay.core == decay_group.top: - for i in decay.outs: - if (i not in set_x) and (i in decay_group.outs): - set_x[i] = (decay_chain, decay) - ref_matrix[i] = decay_chain - # or in the first chain - for i in decay_group.outs: - if i not in set_x: - decay_chain = next(iter(decay_chain_struct)) - for decay in decay_chain: - for j in decay.outs: - if i == j: - set_x[i] = (decay_chain, decay) - ref_matrix[i] = decay_chain for idx, decay_chain in enumerate(decay_chain_struct): for decay in decay_chain: part_data = decay_data[decay_chain][decay] for i in decay.outs: if i in decay_group.outs and decay_chain != set_x[i][0]: - idx2, decay2 = set_x[i] - part_data2 = decay_data[idx2][decay2] if r_boost: r_matrix = decay_data[decay_chain]["r_matrix"][i] b_matrix = decay_data[decay_chain]["b_matrix"][i] - r_matrix_ref = decay_data[ref_matrix[i]]["r_matrix"][i] - b_matrix_ref = decay_data[ref_matrix[i]]["b_matrix"][i] + r_matrix_ref = ref_matrix_final[i][ + "r_matrix" + ] # decay_data[ref_matrix[i]]["r_matrix"][i] + b_matrix_ref = ref_matrix_final[i][ + "b_matrix" + ] # decay_data[ref_matrix[i]]["b_matrix"][i] R = SU2M(r_matrix_ref["x"]) * SU2M.inv(r_matrix) # print(R) if final_rest: @@ -419,10 +464,12 @@ def cal_angle_from_particle( ) ang = R.get_euler_angle() else: + idx2, part_data2 = set_x[i] + # part_data2 = decay_data[idx2][decay2] x1 = part_data[i]["x"] - x2 = part_data2[i]["x"] + x2 = part_data2["x"] z1 = part_data[i]["z"] - z2 = part_data2[i]["z"] + z2 = part_data2["z"] ang = EulerAngle.angle_zx_zx(z1, x1, z2, x2) # ang = AlignmentAngle.angle_px_px(z1, x1, z2, x2) part_data[i]["aligned_angle"] = ang @@ -581,6 +628,7 @@ def cal_angle_from_momentum_base( r_boost=True, random_z=False, batch=65000, + align_ref=None, ) -> CalAngleData: """ Transform 4-momentum data in files for the amplitude model automatically via DecayGroup. @@ -591,13 +639,25 @@ def cal_angle_from_momentum_base( """ if data_shape(p) is None: return cal_angle_from_momentum_single( - p, decs, using_topology, center_mass, r_boost, random_z + p, + decs, + using_topology, + center_mass, + r_boost, + random_z, + align_ref=align_ref, ) ret = [] for i in split_generator(p, batch): ret.append( cal_angle_from_momentum_single( - i, decs, using_topology, center_mass, r_boost, random_z + i, + decs, + using_topology, + center_mass, + r_boost, + random_z, + align_ref=align_ref, ) ) return data_merge(*ret) @@ -612,6 +672,7 @@ def identical_particles_swap(id_particles): def identical_particles_swap_p(p4, id_particles): + old_order = tuple(tuple(i) for i in id_particles) for comb in identical_particles_swap(id_particles): all_keys = tuple(p4.keys()) name_map = {str(k): k for k in all_keys} @@ -620,7 +681,7 @@ def identical_particles_swap_p(p4, id_particles): for ci, pi in zip(c, p_list): swap_map[ci] = name_map[pi] new_order = tuple([swap_map.get(str(i), i) for i in all_keys]) - if all_keys == new_order: + if comb == old_order: continue yield (new_order, comb), dict(zip(new_order, p4.values())) @@ -645,6 +706,7 @@ def cal_angle_from_momentum_id_swap( r_boost=True, random_z=False, batch=65000, + align_ref=None, ) -> CalAngleData: ret = [] id_particles = decs.identical_particles @@ -657,7 +719,14 @@ def cal_angle_from_momentum_id_swap( data["id_swap"] = {} for i, pi in identical_particles_swap_p(p, id_particles): data["id_swap"][i] = cal_angle_from_momentum_base( - pi, decs, using_topology, center_mass, r_boost, random_z, batch + pi, + decs, + using_topology, + center_mass, + r_boost, + random_z, + batch, + align_ref=align_ref, ) return data @@ -670,6 +739,7 @@ def cal_angle_from_momentum( r_boost=True, random_z=False, batch=65000, + align_ref=None, ) -> CalAngleData: """ Transform 4-momentum data in files for the amplitude model automatically via DecayGroup. @@ -693,14 +763,28 @@ def cal_angle_from_momentum( id_particles = decs.identical_particles cp_particles = decs.cp_particles data = cal_angle_from_momentum_id_swap( - p, decs, using_topology, center_mass, r_boost, random_z, batch + p, + decs, + using_topology, + center_mass, + r_boost, + random_z, + batch, + align_ref=align_ref, ) if cp_particles is None or len(cp_particles) == 0: return data else: p2 = cp_swap_p(p, decs.outs, id_particles, cp_particles) data["cp_swap"] = cal_angle_from_momentum_id_swap( - p2, decs, using_topology, center_mass, r_boost, random_z, batch + p2, + decs, + using_topology, + center_mass, + r_boost, + random_z, + batch, + align_ref=align_ref, ) return data @@ -712,6 +796,7 @@ def cal_angle_from_momentum_single( center_mass=False, r_boost=True, random_z=True, + align_ref=None, ) -> CalAngleData: """ Transform 4-momentum data in files for the amplitude model automatically via DecayGroup. @@ -733,7 +818,12 @@ def cal_angle_from_momentum_single( # exit() data_p = add_mass(data_p, dec) data_d = cal_angle_from_particle( - data_p, decs, using_topology, r_boost=r_boost, random_z=random_z + data_p, + decs, + using_topology, + r_boost=r_boost, + random_z=random_z, + align_ref=align_ref, ) data = {"particle": data_p, "decay": data_d} add_relative_momentum(data) diff --git a/tf_pwa/config_loader/data.py b/tf_pwa/config_loader/data.py index 44a480da..b630d48c 100644 --- a/tf_pwa/config_loader/data.py +++ b/tf_pwa/config_loader/data.py @@ -153,12 +153,14 @@ def cal_angle(self, p4, charge=None): center_mass = self.dic.get("center_mass", False) r_boost = self.dic.get("r_boost", True) random_z = self.dic.get("random_z", True) + align_ref = self.dic.get("align_ref", None) data = cal_angle_from_momentum( p4, self.decay_struct, center_mass=center_mass, r_boost=r_boost, random_z=random_z, + align_ref=align_ref, ) if charge is not None: data["charge_conjugation"] = charge diff --git a/tf_pwa/config_loader/plot.py b/tf_pwa/config_loader/plot.py index 7263c65f..fe495547 100644 --- a/tf_pwa/config_loader/plot.py +++ b/tf_pwa/config_loader/plot.py @@ -962,10 +962,14 @@ def get_var(dic, tail): return dict(zip(var_index, ret)) var_x_f = sym.lambdify( - var_x.free_symbols | var_y.free_symbols, var_x, modules="numpy" + tuple(var_x.free_symbols | var_y.free_symbols), + var_x, + modules="numpy", ) var_y_f = sym.lambdify( - var_x.free_symbols | var_y.free_symbols, var_y, modules="numpy" + tuple(var_x.free_symbols | var_y.free_symbols), + var_y, + modules="numpy", ) data_1 = var_x_f(**get_var(data_dict, "")) diff --git a/tf_pwa/model/model.py b/tf_pwa/model/model.py index d3ceadd4..35f48703 100644 --- a/tf_pwa/model/model.py +++ b/tf_pwa/model/model.py @@ -43,8 +43,8 @@ def _batch_sum(f, data_i, weight_i, trans, resolution_size, args, kwargs): part_y = tf.reshape(part_y, (-1, resolution_size)) part_y = tf.reduce_sum(part_y, axis=-1) event_w = tf.reduce_sum(rw, axis=-1) - event_w = tf.where(event_w == 0, tf.ones_like(event_w), event_w) - part_y = part_y / event_w + dom_event_w = tf.where(event_w == 0, tf.ones_like(event_w), event_w) + part_y = part_y / dom_event_w # event_w part_y = trans(part_y) y_i = tf.reduce_sum(event_w * part_y) return y_i @@ -311,7 +311,8 @@ def nll(self, data, mcdata): amp_s2 = tf.reshape(amp_s2, (-1, self.resolution_size)) amp_s2 = tf.reduce_sum(amp_s2, axis=-1) weight = tf.reduce_sum(rw, axis=-1) - ln_data = clip_log(amp_s2 / weight) + dom_weight = tf.where(weight == 0, 1.0, weight) + ln_data = clip_log(amp_s2 / dom_weight) mc_weight = mcdata.get("weight", tf.ones((data_shape(mcdata),))) int_mc = tf.reduce_sum( mc_weight * self.signal(mcdata) diff --git a/tf_pwa/tests/config_toy2.yml b/tf_pwa/tests/config_toy2.yml index 07219322..048322ac 100644 --- a/tf_pwa/tests/config_toy2.yml +++ b/tf_pwa/tests/config_toy2.yml @@ -6,6 +6,7 @@ data: bg_weight: 0.1 random_z: False r_boost: False + align_ref: "center_mass" use_tf_function: True cached_int: True cached_data: "toy_data/data.npy" diff --git a/tf_pwa/tests/linear_npy.npy b/tf_pwa/tests/linear_npy.npy new file mode 100644 index 00000000..7e3ef408 Binary files /dev/null and b/tf_pwa/tests/linear_npy.npy differ diff --git a/tf_pwa/tests/test_amp.py b/tf_pwa/tests/test_amp.py index 1f040a84..1d8d3b41 100644 --- a/tf_pwa/tests/test_amp.py +++ b/tf_pwa/tests/test_amp.py @@ -169,6 +169,31 @@ def test_BW(): ) +def test_linear_npy(): + import os + + dir_name = os.path.dirname(__file__) + from tf_pwa.amp import interpolation + + a = get_particle( + "linear", + J=1, + P=-1, + model="linear_npy", + mass=3.6, + width=0.01, + file=os.path.join(dir_name, "linear_npy.npy"), + ) + b = [get_particle(i, J=0, P=-1) for i in "ac"] + get_decay(a, b) + a.init_params() + b = a.get_amp( + {"m": np.array(1.0)}, + {"|q|": np.array(1.0), "|q0|": np.array(1.0)}, + all_data={}, + ) + + def simple_run(name): a = get_particle(name, J=1, P=-1, model=name, mass=3.6, width=0.01) b = [get_particle(i, J=0, P=-1) for i in "ac"] diff --git a/tf_pwa/tests/test_formula.py b/tf_pwa/tests/test_formula.py index d930cb04..121c5f70 100644 --- a/tf_pwa/tests/test_formula.py +++ b/tf_pwa/tests/test_formula.py @@ -17,6 +17,17 @@ def simple_decay(name, **extra): return a +def nest_dict_zip(a, b): + ret = {} + for i, j in zip(a, b): + if isinstance(i, (list, tuple)): + tmp = nest_dict_zip(i, j) + else: + tmp = {i: j} + ret.update(tmp) + return ret + + def simple_test(model, **extra): p = simple_decay(model, **extra) p.solve_pole() @@ -24,7 +35,7 @@ def simple_test(model, **extra): a = 1 / p(m).numpy() var = p.get_sympy_var() f = p.get_sympy_dom(*var) - g = f.subs(dict(zip(var[1:], p.get_num_var()))) + g = f.subs(nest_dict_zip(var[1:], p.get_num_var())) b = np.array([complex(g.subs({var[0]: i}).evalf()) for i in m]) assert np.allclose(np.real(a), np.real(b)) assert np.allclose(np.imag(a), np.imag(b)) @@ -49,7 +60,7 @@ def run_BWR_LS(**extra): var = p.get_sympy_var() f = p.get_sympy_dom(*var) var_num = p.get_num_var() - g = f.subs(dict(zip(var[1:], var_num))) + g = f.subs(nest_dict_zip(var[1:], var_num)) b = np.array([complex(g.subs({var[0]: i}).evalf()) for i in m]) q = get_relative_p2(m, *var_num[-2:]) diff --git a/tf_pwa/variable.py b/tf_pwa/variable.py index c5bb1ceb..74193db6 100644 --- a/tf_pwa/variable.py +++ b/tf_pwa/variable.py @@ -987,7 +987,12 @@ def get_func(self): # init func string into sympy f(x) or f(y) """ x, a, b, y = sy.symbols("x a b y") f = sy.sympify(self.func) - f = f.subs({a: self.lower, b: self.upper}) + f = f.subs( + { + a: self.lower if self.lower is not None else -1e9, + b: self.upper if self.upper is not None else 1e9, + } + ) df = sy.diff(f, x) df2 = sy.diff(df, x) inv = sy.solve(f - y, x)