diff --git a/.gitignore b/.gitignore index 2b1aa21f..f5571df5 100644 --- a/.gitignore +++ b/.gitignore @@ -9,7 +9,6 @@ tf_pwa/version.py *.npy *.json *.log -*.yml *.png *.pdf *.pptx diff --git a/examples/ex2_particle_amplitude.py b/examples/ex2_particle_amplitude.py index eb2b3444..3061508d 100644 --- a/examples/ex2_particle_amplitude.py +++ b/examples/ex2_particle_amplitude.py @@ -15,8 +15,11 @@ "R1": {"J": 0, "P": 1, "mass": 1.0, "width": 0.07}, "R2": {"J": 1, "P": -1, "mass": 1.225, "width": 0.08}, } - -a, b, c, d = [get_particle(i, J=0, P=-1) for i in "ABCD"] +m_A, m_B, m_C, m_D = 1.8, 0.18, 0.18, 0.18 +a, b, c, d = [ + get_particle(i, J=0, P=-1, mass=m) + for i, m in zip("ABCD", [m_A, m_B, m_C, m_D]) +] r1, r2, r3 = [get_particle(i, **resonances[i]) for i in resonances.keys()] diff --git a/examples/ex_params_trans.py b/examples/ex_params_trans.py index 8fe0ba7f..c83fca46 100644 --- a/examples/ex_params_trans.py +++ b/examples/ex_params_trans.py @@ -78,23 +78,52 @@ print(a2_x.numpy(), pt.get_error(a2_x).numpy()) # %% -# We can also calculate some more complex examples, such as the ratio in mass range (0.75, 0.85) over full phace space. +# Uncertainties of fit fractions +# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +# We can also calculate some more complex examples, such as the fit fractions of all in `C+D`. # Even further, we can get the error of error in the meaning of error propagation. -from tf_pwa.data import data_mask - -m_R2 = phsp.get_mass("(B, D)") -cut_cond = (m_R2 < 0.85) & (m_R2 > 0.75) - amp = config.get_amplitude() with config.params_trans() as pt1: with config.params_trans() as pt: int_mc = tf.reduce_sum(amp(phsp)) - cut_phsp = data_mask(phsp, cut_cond) - cut_int_mc = tf.reduce_sum(amp(cut_phsp)) - ratio = cut_int_mc / int_mc + with amp.temp_used_res(["R1_a", "R1_b"]): + part_int_mc = tf.reduce_sum(amp(phsp)) + ratio = part_int_mc / int_mc error = pt.get_error(ratio) print(ratio.numpy(), "+/-", error.numpy()) print(error.numpy(), "+/-", pt1.get_error(error).numpy()) + +# %% +# For large data size it would be some problem named OOM (out of memory). +# TFPWA provide `vm.batch_sum_var` to do sum of large samples + +int_mc_v = config.vm.batch_sum_var(amp, phsp, batch=5000) + +with amp.temp_used_res(["R1_a", "R1_b"]): + part_int_mc_v = config.vm.batch_sum_var(amp, phsp, batch=5000) + +# %% +# It will store the pre-calculated gradients as + +print(int_mc_v.grad, part_int_mc_v.grad) + +# %% +# Then, we can use it as a function to do error propagation: + +with config.params_trans() as pt: + ratio = part_int_mc_v() / int_mc_v() +error = pt.get_error(ratio) + +print(ratio.numpy(), "+/-", error.numpy()) + +# %% +# Besides the error propagation, there would be some additional uncertainties. +# For example, the uncertainty from the integration sample size is often defined as the sum of square as + +with amp.temp_used_res(["R1_a", "R1_b"]): + int_square = tf.reduce_sum((amp(phsp) / int_mc) ** 2) + +print(ratio.numpy(), "+/-", error.numpy(), "+/-", tf.sqrt(int_square).numpy()) diff --git a/fit.py b/fit.py index 45bc90b3..6c119572 100644 --- a/fit.py +++ b/fit.py @@ -55,7 +55,7 @@ def print_fit_result_roofit(config, fit_result): value = fit_result.params params_name = config.vm.trainable_vars n_par = len(params_name) - name_size = max(len(i) for i in params_name) + name_size = max([5] + [len(i) for i in params_name]) # fcn = config.get_fcn() # _, grad = fcn.nll_grad(fit_result.params) # edm = np.dot(np.dot(config.inv_he, grad), grad) diff --git a/tf_pwa/amp/core.py b/tf_pwa/amp/core.py index ad303c63..109c727d 100644 --- a/tf_pwa/amp/core.py +++ b/tf_pwa/amp/core.py @@ -259,6 +259,11 @@ def add_var(self, names, is_complex=False, shape=(), **kwargs): """ if not hasattr(self, "_variables_map"): self._variables_map = {} + if True: + default_config = getattr(self, "default_params", {}).get(names, {}) + if isinstance(default_config, (float, int)): + default_config = {"value": default_config} + kwargs.update(default_config) name = self.get_variable_name(names) var = Variable(name, shape, is_complex, **kwargs) self._variables_map[names] = var diff --git a/tf_pwa/applications.py b/tf_pwa/applications.py index b9ebc871..86a27b90 100644 --- a/tf_pwa/applications.py +++ b/tf_pwa/applications.py @@ -48,17 +48,24 @@ def fit_fractions( :param mcdata: MCdata array. :param inv_he: The inverse of Hessian matrix. If it's not given, the errors will not be calculated. :return frac: Dictionary of fit fractions for each resonance. - :return err_frac: Dictionary of their errors. If ``inv_he`` is ``None``, it will be a dictionary of ``None``. + :return err_frac: Dictionary of their errors. If ``inv_he`` is ``None``, it will be a dictionary of ``{}``. """ if params is None: params = {} err_frac = {} if method == "old": - with amp.temp_params(params): - frac, grad = cal_fitfractions(amp, mcdata, res=res, batch=batch) if inv_he is not None: + with amp.temp_params(params): + frac, grad = cal_fitfractions( + amp, mcdata, res=res, batch=batch + ) for i in frac: err_frac[i] = np.sqrt(np.dot(np.dot(inv_he, grad[i]), grad[i])) + else: + with amp.temp_params(params): + frac = cal_fitfractions_no_grad( + amp, mcdata, res=res, batch=batch + ) return frac, err_frac else: ret = FitFractions(amp, res) @@ -289,6 +296,8 @@ def cal_hesse_error( :return hesse_error: List of errors. :return inv_he: The inverse Hessian matrix. """ + if len(fcn.vm.trainable_vars) == 0: + return {}, None t = time.time() nll, g, h = fcn.nll_grad_hessian( params diff --git a/tf_pwa/config_loader/config_loader.py b/tf_pwa/config_loader/config_loader.py index 1c14f4f2..020c38bf 100644 --- a/tf_pwa/config_loader/config_loader.py +++ b/tf_pwa/config_loader/config_loader.py @@ -262,8 +262,6 @@ def add_constraints(self, amp): constrains = {} self.add_decay_constraints(amp, constrains.get("decay", {})) self.add_particle_constraints(amp, constrains.get("particle", {})) - self.add_fix_var_constraints(amp, constrains.get("fix_var", {})) - self.add_free_var_constraints(amp, constrains.get("free_var", [])) self.add_var_range_constraints(amp, constrains.get("var_range", {})) self.add_var_equal_constraints(amp, constrains.get("var_equal", [])) self.add_pre_trans_constraints(amp, constrains.get("pre_trans", None)) @@ -273,6 +271,8 @@ def add_constraints(self, amp): self.add_gauss_constr_constraints( amp, constrains.get("gauss_constr", {}) ) + self.add_fix_var_constraints(amp, constrains.get("fix_var", {})) + self.add_free_var_constraints(amp, constrains.get("free_var", [])) for k, v in self.extra_constrains.items(): v(amp, constrains.get(k, {})) @@ -832,8 +832,8 @@ def get_params_error( params = {} if correct_params is None: correct_params = [] - if method is None: - method = "correct" + if len(correct_params) > 0 and method is None: + method = "correct" if hasattr(params, "params"): params = getattr(params, "params") if not using_cached: @@ -1016,10 +1016,10 @@ def set_params(self, params, neglect_params=None): neglect_params = self._neglect_when_set_params if len(neglect_params) != 0: for v in params: - if v in self._neglect_when_set_params: + if v in neglect_params: warnings.warn( "Neglect {} when setting params.".format( - neglect_params + [i for i in params if i in neglect_params] ) ) del ret[v] diff --git a/tf_pwa/data_trans/helicity_angle.py b/tf_pwa/data_trans/helicity_angle.py index 2d67ec90..219532fc 100644 --- a/tf_pwa/data_trans/helicity_angle.py +++ b/tf_pwa/data_trans/helicity_angle.py @@ -3,6 +3,7 @@ from tf_pwa.amp import get_relative_p from tf_pwa.angle import LorentzVector as lv +from tf_pwa.data import data_index class HelicityAngle1: @@ -73,7 +74,9 @@ def get_all_mass(self, replace_mass): for i in self.decay_chain: for j in [i.core] + list(i.outs): if j not in ms: - if str(j) in replace_mass: + if j in replace_mass: + ms[j] = replace_mass[j] + elif str(j) in replace_mass: ms[j] = tf.convert_to_tensor( replace_mass[str(j)], tf.float64 ) @@ -90,7 +93,9 @@ def generate_p_mass(self, name, m, random=False): for i in self.decay_chain: data[i] = {} data[i]["|p|"] = get_relative_p( - ms[i.core], ms[i.outs[0]], ms[i.outs[1]] + data_index(ms, i.core), + data_index(ms, i.outs[0]), + data_index(ms, i.outs[1]), ) if random: costheta = np.random.random(m.shape) * 2 - 1 @@ -108,10 +113,13 @@ def generate_p_mass(self, name, m, random=False): def build_data(self, ms, costheta, phi): """generate monmentum with M_name = m""" data = {} + ms = self.get_all_mass(ms) for j, i in enumerate(self.decay_chain): data[i] = {} data[i]["|p|"] = get_relative_p( - ms[i.core], ms[i.outs[0]], ms[i.outs[1]] + data_index(ms, i.core), + data_index(ms, i.outs[0]), + data_index(ms, i.outs[1]), ) costheta_i = costheta[j] phi_i = phi[j] diff --git a/tf_pwa/fit.py b/tf_pwa/fit.py index 6fc0734a..9e6496a6 100644 --- a/tf_pwa/fit.py +++ b/tf_pwa/fit.py @@ -278,7 +278,12 @@ def callback(x): f_g = Cached_FG(f_g, grad_scale=grad_scale) # print(f_g) x0 = np.array(fcn.vm.get_all_val(True)) - # print(x0, fcn.vm.get_all_dic()) + if len(x0) == 0: + min_nll, _ = f_g(x0) + params = fcn.get_params() + return FitResult( + params, fcn, min_nll, ndf=0, success=True, hess_inv=None + ) # s = minimize(f_g, x0, method='trust-constr', jac=True, hess=BFGS(), options={'gtol': 1e-4, 'disp': True}) if method == "test": try: diff --git a/tf_pwa/fitfractions.py b/tf_pwa/fitfractions.py index 993aa470..cd472f5c 100644 --- a/tf_pwa/fitfractions.py +++ b/tf_pwa/fitfractions.py @@ -6,20 +6,32 @@ from tf_pwa.data import LazyCall, data_split -def eval_integral( +def _eval_integral( f, data, var, weight=None, args=(), no_grad=False, kwargs=None ): kwargs = {} if kwargs is None else kwargs weight = 1.0 if weight is None else weight if no_grad: - ret = tf.reduce_sum(f(data, *args, **kwargs) * weight) + fx = f(data, *args, **kwargs) * weight + ret = tf.reduce_sum(fx) ret_grad = np.zeros((len(var),)) else: with tf.GradientTape() as tape: - ret = tf.reduce_sum(f(data, *args, **kwargs) * weight) + fx = f(data, *args, **kwargs) * weight + ret = tf.reduce_sum(fx) ret_grad = tape.gradient(ret, var, unconnected_gradients="zero") - ret_grad = np.stack([i.numpy() for i in ret_grad]) - return ret.numpy(), ret_grad + if len(ret_grad) == 0: + ret_grad = np.array([]) + else: + ret_grad = np.stack([i.numpy() for i in ret_grad]) + int_square = tf.reduce_sum(fx**2) + return ret.numpy(), ret_grad, int_square.numpy() + + +def force_list(x): + if isinstance(x, (list, tuple)): + return x + return [x] class FitFractions: @@ -30,6 +42,7 @@ def __init__(self, amp, res): self.res = res self.cached_int = {} self.cached_grad = {} + self.cached_square = {} self.cached_int_total = 0.0 self.cached_grad_total = np.zeros((self.n_var,)) self.error_matrix = np.diag(np.zeros((self.n_var,))) @@ -46,6 +59,7 @@ def init_res_table(self): else: name = (str(self.res[i]), str(self.res[j])) self.cached_int[name] = 0.0 + self.cached_square[name] = 0.0 self.cached_grad[name] = np.zeros_like((self.n_var,)) def integral(self, mcdata, *args, batch=None, no_grad=False, **kwargs): @@ -62,7 +76,7 @@ def append_int(self, mcdata, *args, weight=None, no_grad=False, **kwargs): mcdata = mcdata.eval() if weight is None: weight = mcdata.get("weight", 1.0) - int_mc, g_int_mc = eval_integral( + int_mc, g_int_mc, _ = _eval_integral( self.amp, mcdata, var=self.var, @@ -74,15 +88,16 @@ def append_int(self, mcdata, *args, weight=None, no_grad=False, **kwargs): self.cached_grad_total += g_int_mc cahced_res = self.amp.used_res amp_tmp = self.amp + fl = force_list for i in range(len(self.res)): for j in range(i, -1, -1): if i == j: name = str(self.res[i]) - amp_tmp.set_used_res([self.res[i]]) + amp_tmp.set_used_res(fl(self.res[i])) else: name = (str(self.res[i]), str(self.res[j])) - amp_tmp.set_used_res([self.res[i], self.res[j]]) - int_tmp, g_int_tmp = eval_integral( + amp_tmp.set_used_res(fl(self.res[i]) + fl(self.res[j])) + int_tmp, g_int_tmp, int_square = _eval_integral( amp_tmp, mcdata, var=self.var, @@ -92,6 +107,9 @@ def append_int(self, mcdata, *args, weight=None, no_grad=False, **kwargs): ) self.cached_int[name] = self.cached_int[name] + int_tmp self.cached_grad[name] = self.cached_grad[name] + g_int_tmp + self.cached_square[name] = ( + self.cached_square[name] + int_square + ) self.amp.set_used_res(cahced_res) @@ -130,10 +148,9 @@ def get_frac_grad(self, sum_diag=True): g_fit_frac["sum_diag"] = sum( [g_fit_frac[str(i)] for i in self.res] ) - print(fit_frac) return fit_frac, g_fit_frac - def get_frac(self, error_matrix=None, sum_diag=True): + def get_frac(self, error_matrix=None, sum_diag=True, add_int_error=False): if error_matrix is None: error_matrix = self.error_matrix fit_frac, g_fit_frac = self.get_frac_grad(sum_diag=sum_diag) @@ -143,6 +160,9 @@ def get_frac(self, error_matrix=None, sum_diag=True): for k, v in g_fit_frac.items(): e = np.sqrt(np.dot(np.dot(error_matrix, v), v)) fit_frac_err[k] = e + if add_int_error and k in self.cached_square: + scale = 1 / self.cached_int_total**2 + fit_frac_err[k] = np.sqrt(e**2 + self.cached_square[k] * scale) return fit_frac, fit_frac_err def __iter__(self): diff --git a/tf_pwa/tests/config_no_params.yml b/tf_pwa/tests/config_no_params.yml new file mode 100644 index 00000000..fc04a00f --- /dev/null +++ b/tf_pwa/tests/config_no_params.yml @@ -0,0 +1,49 @@ +data: + dat_order: [B, C, D] + extra_var: + c: + default: 1 + data: [["toy_data/data.dat"]] + bg: [["toy_data/bg.dat"]] + phsp: [["toy_data/PHSP.dat"]] + phsp_c: ["toy_data/phsp_eff_value.dat"] + bg_weight: [0.1] + random_z: False + r_boost: False + cached_amp: True + use_mix_likelihood: True + +decay: + A: + - [R_BC, D, ls_list: [[0, 1], [2, 1]]] + - [R_BD, C, l_list: [0, 2]] + - [R_CD, B] + R_BC: [B, C] + R_BD: [B, D] + R_CD: [C, D] + +particle: + $top: + A: { J: 1, P: -1, spins: [-1, 1] } + $finals: + B: { J: 1, P: -1 } + C: { J: 1, P: -1 } + D: { J: 0, P: -1 } + R_BC: [] # { J: 1, Par: 1, m0: 4.16, g0: 0.1, model: BWR2 } + R_BD: { J: 1, Par: 1, m0: 2.43, g0: 0.3, model: BW } + R_CD: [] # { J: 1, Par: 1, m0: 2.42, g0: 0.03 } + +constrains: + fix_var: + A->R_BD.C_g_ls_1r: 0.0 + A->R_BD.C_g_ls_1i: 0.0 + A->R_BD.C_g_ls_2r: 0.0 + A->R_BD.C_g_ls_2i: 0.0 + R_BD->B.D_g_ls_1r: 0.0 + R_BD->B.D_g_ls_1i: 0.0 + +plot: + mass: + R_BC: { display: "$M_{BC}$" } + R_BD: { display: "$M_{BD}$" } + R_CD: { display: "$M_{CD}$" } diff --git a/tf_pwa/tests/test_full.py b/tf_pwa/tests/test_full.py index e64f49f0..35f8ec64 100644 --- a/tf_pwa/tests/test_full.py +++ b/tf_pwa/tests/test_full.py @@ -144,6 +144,14 @@ def test_save_model(toy_config): config.save_tensorflow_model("toy_data/model") +def test_fit_no_params(gen_toy): + config = ConfigLoader(f"{this_dir}/config_no_params.yml") + config.fit() + config.get_params_error(method="default") + config.cal_fitfractions(method="old") + config.cal_fitfractions(method="new") + + def test_cfit(gen_toy): config = ConfigLoader(f"{this_dir}/config_cfit.yml") config.set_params(f"{this_dir}/gen_params.json") @@ -333,6 +341,9 @@ def pw_f(x, **kwargs): fit_result.save_as("toy_data/final_params.json") fit_frac, frac_err = toy_config.cal_fitfractions() fit_frac, frac_err = toy_config.cal_fitfractions(method="new") + fit_frac_obj = toy_config.cal_fitfractions( + method="new", res=["R_BC", ["R_BD", "R_CD"]] + ) save_frac_csv("toy_data/fit_frac.csv", fit_frac) with toy_config.params_trans() as pt: