From 7778b326153b2d51e253155f850ac6d0ff81e5c3 Mon Sep 17 00:00:00 2001 From: jiangyi15 Date: Fri, 25 Aug 2023 23:43:24 +0800 Subject: [PATCH 1/5] feat: plot 2d pull directly --- tf_pwa/config_loader/plot.py | 220 +++++++++++++++++++++++++++++------ tf_pwa/tests/test_full.py | 9 ++ 2 files changed, 194 insertions(+), 35 deletions(-) diff --git a/tf_pwa/config_loader/plot.py b/tf_pwa/config_loader/plot.py index c4d8eeee..e359cc19 100644 --- a/tf_pwa/config_loader/plot.py +++ b/tf_pwa/config_loader/plot.py @@ -226,6 +226,7 @@ def plot_partial_wave( chains_id_method=None, phsp_rec=None, cut_function=lambda x: 1, + plot_function=None, **kwargs ): """ @@ -251,6 +252,8 @@ def plot_partial_wave( :param linestyle_file: legend linestyle configuration file name (YAML format), string (such as "legend.yml") """ + if plot_function is None: + plot_function = self._plot_partial_wave if params is None: params = {} @@ -329,13 +332,13 @@ def plot_partial_wave( cut_function=cut_function, **kwargs, ) - self._plot_partial_wave( + plot_function( data_dict, phsp_dict, bg_dict, - prefix, - plot_var_dic, - chain_property, + prefix=prefix, + plot_var_dic=plot_var_dic, + chain_property=chain_property, nll=nll, **kwargs, ) @@ -360,13 +363,13 @@ def plot_partial_wave( cut_function=cut_function, **kwargs, ) - self._plot_partial_wave( + plot_function( data_dict, phsp_dict, bg_dict, - prefix + "d{}_".format(i), - plot_var_dic, - chain_property, + prefix=prefix + "d{}_".format(i), + plot_var_dic=plot_var_dic, + chain_property=chain_property, nll=nll, **kwargs, ) @@ -414,13 +417,13 @@ def plot_partial_wave( phsps_dict[ct] = np.concatenate(phsps_dict[ct]) for ct in bgs_dict: bgs_dict[ct] = np.concatenate(bgs_dict[ct]) - self._plot_partial_wave( + plot_function( datas_dict, phsps_dict, bgs_dict, - prefix + "com_", - plot_var_dic, - chain_property, + prefix=prefix + "com_", + plot_var_dic=plot_var_dic, + chain_property=chain_property, nll=nll, **kwargs, ) @@ -970,6 +973,31 @@ def _plot_var_name(name): raise TypeError("not string or list") +def build_read_var_function(all_var, where={}): + vari = [sym.simplify(i) for i in all_var] + used_var = [] + var_index = [] + all_symbols = set() + for i in vari: + all_symbols = all_symbols | i.free_symbols + all_symbols = tuple(all_symbols) + + for i in all_symbols: + var_index.append(str(i)) + used_var.append(where.get(str(i), str(i))) + + used_var = [_plot_var_name(i) for i in used_var] + + def get_var(dic, tail): + ret = [] + for i in used_var: + ret.append(dic[i + tail]) + return dict(zip(var_index, ret)) + + var_f = [sym.lambdify(all_symbols, i, modules="numpy") for i in vari] + return var_f, get_var + + @ConfigLoader.register_function() def _2d_plot_v2( self, @@ -995,32 +1023,13 @@ def _2d_plot_v2( if "&" in k: continue assert ("x" in v) and ("y" in v) + var_x = sym.simplify(v["x"]) var_y = sym.simplify(v["y"]) where = v.get("where", {}) - used_var = [] - var_index = [] - for i in var_x.free_symbols | var_y.free_symbols: - var_index.append(str(i)) - used_var.append(where.get(str(i), str(i))) - - used_var = [_plot_var_name(i) for i in used_var] - - def get_var(dic, tail): - ret = [] - for i in used_var: - ret.append(dic[i + tail]) - return dict(zip(var_index, ret)) - - var_x_f = sym.lambdify( - tuple(var_x.free_symbols | var_y.free_symbols), - var_x, - modules="numpy", - ) - var_y_f = sym.lambdify( - tuple(var_x.free_symbols | var_y.free_symbols), - var_y, - modules="numpy", + + (var_x_f, var_y_f), get_var = build_read_var_function( + [var_x, var_y], where ) data_1 = var_x_f(**get_var(data_dict, "")) @@ -1124,6 +1133,147 @@ def plot_axis(): print("Finish plotting 2D fitted " + prefix + k) +@ConfigLoader.register_function() +def get_dalitz(config, a, b): + decay = config.get_decay(False) + da = decay.get_decay_chain(a) + db = decay.get_decay_chain(b) + pa = decay.get_particle(a) + pb = decay.get_particle(b) + + for i in da: + if pa in i.outs: + topa = i.core + if pa == i.core: + outs_a = i.outs + for i in db: + if pb in i.outs: + topb = i.core + if pb == i.core: + outs_b = i.outs + same_finals = [i for i in outs_a if i in db.outs] + p1 = [i for i in outs_a if i not in same_finals] + p3 = [i for i in outs_b if i not in same_finals] + check = ((topa == topb),) + check = check and len(same_finals) == 1 + check = check and len(p1) == 1 + check = check and len(p3) == 1 + if not check: + return None + p0, p1, p2, p3 = topa, p1[0], same_finals[0], p3[0] + p0, p1, p2, p3 = [ + config.get_decay().get_particle(str(i)) for i in [p0, p1, p2, p3] + ] + m0, m1, m2, m3 = map(lambda x: x.get_mass(), [p0, p1, p2, p3]) + return m0, m1, m2, m3 + + +@ConfigLoader.register_function() +def get_dalitz_boundary(config, a, b, N=1000): + dalitz = get_dalitz(config, a, b) + assert dalitz is not None, "not valid daliz plot" + m0, m1, m2, m3 = dalitz + # print(m0, m1, m2, m3) + from tf_pwa.angle import kine_min_max + + s12_min, s12_max = float(m1 + m2), float(m0 - m3) + s12 = np.linspace(s12_min**2, s12_max**2, N) + s23_min, s23_max = kine_min_max(s12, *map(float, [m0, m1, m2, m3])) + return s12, np.stack([s23_min, s23_max], axis=-1) + + +@ConfigLoader.register_function() +def plot_adaptive_2dpull( + config, var1, var2, binning=[[2, 2]] * 3, ax=plt, where={}, cut_zero=True +): + import matplotlib as mpl + import matplotlib.colors as mcolors + import matplotlib.patches as mpathes + + from tf_pwa.adaptive_bins import AdaptiveBound + + def plot_function_2dpull( + data_dict, phsp_dict, bg_dict, plot_var_dic, **kwargs + ): + nonlocal ax + if cut_zero: + cut = data_dict["data_weights"] != 0 + else: + cut = np.ones(data_dict["data_weights"].shape, dtype=np.bool) + (var_x_f, var_y_f), get_var = build_read_var_function( + [var1, var2], where=where + ) + x = var_x_f(**get_var(data_dict, ""))[cut] + y = var_y_f(**get_var(data_dict, ""))[cut] + w = data_dict["data_weights"][cut] + x_phsp = var_x_f(**get_var(phsp_dict, "_MC")) + y_phsp = var_y_f(**get_var(phsp_dict, "_MC")) + w_phsp = phsp_dict["MC_total_fit"] + data_cut = np.array([x, y]) + adapter = AdaptiveBound(data_cut, binning) + phsps = adapter.split_data(np.array([x_phsp, y_phsp, w_phsp])) + datas = adapter.split_data(np.array([x, y, w])) + if bg_dict != {}: + x_bg = var_x_f(**get_var(bg_dict, "_sideband")) + y_bg = var_y_f(**get_var(bg_dict, "_sideband")) + w_bg = bg_dict["sideband_weights"] + bgs = adapter.split_data(np.array([x_bg, y_bg, w_bg])) + bound = adapter.get_bounds() + numbers = [] + pulls = [] + int_norm = 1 + for i, bnd in enumerate(bound): + min_x, min_y = bnd[0] + max_x, max_y = bnd[1] + ndata = np.sum(datas[i][2]) + nmc = np.sum(phsps[i][2]) + if bg_dict != {}: + nmc += np.sum(bgs[i][2]) + numbers.append((ndata, nmc)) + pulls.append((ndata - nmc) / np.sqrt(nmc)) + + max_weight = max(np.max(np.abs(pulls)), 5) + + my_cmap = plt.get_cmap("jet") + if ax == plt: + ax = plt.gca() # fig, ax = plt.subplots() + ax.scatter(x, y, s=1, c="black") + for i, bnd in enumerate(bound): + min_x, min_y = bnd[0] + max_x, max_y = bnd[1] + # print(weights[i]) # max_weight) + rect = mpathes.Rectangle( + (min_x, min_y), + max_x - min_x, + max_y - min_y, + linewidth=1, + facecolor=my_cmap( + pulls[i] / max_weight / 2 + 0.5 + ), # max_weight), + edgecolor="none", # black", + zorder=-1, + ) # cmap(weights[i]/max_weight)) + ax.add_patch(rect) + + normal = mpl.colors.Normalize(vmin=-max_weight, vmax=max_weight) + im = mpl.cm.ScalarMappable(norm=normal, cmap=my_cmap) + # ax.colorbar(im) + ax.get_figure().colorbar(im) + ax.title( + "$\\chi^2/Nbins={:.2f}/{}$".format( + np.sum(np.abs(pulls) ** 2), len(bound) + ) + ) + ax.set_xlim([np.min(x_phsp), np.max(x_phsp)]) + ax.set_ylim([np.min(y_phsp), np.max(y_phsp)]) + ax.set_xlabel(var1) + ax.set_ylabel(var2) + + config.plot_partial_wave( + plot_function=plot_function_2dpull, combine_plot=True + ) + + def hist_error(data, bins=50, xrange=None, weights=1.0, kind="poisson"): if not hasattr(weights, "__len__"): weights = [weights] * data.__len__() diff --git a/tf_pwa/tests/test_full.py b/tf_pwa/tests/test_full.py index df36a65d..b7c6218d 100644 --- a/tf_pwa/tests/test_full.py +++ b/tf_pwa/tests/test_full.py @@ -357,3 +357,12 @@ def test_cp_particles(): config = ConfigLoader(f"{this_dir}/config_self_cp.yml") phsp = config.generate_phsp(100) config.get_amplitude()(phsp) + + +def test_plot_2dpull(toy_config): + import matplotlib.pyplot as plt + + toy_config.plot_adaptive_2dpull("m_R_BC**2", "m_R_CD**2") + a, b = toy_config.get_dalitz_boundary("R_BC", "R_CD") + plt.plot(a, b, color="red") + plt.savefig("adptive_2d.png") From d222f46488ca1617800e7f6825395f1681e2f30a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jiang=20Yi=20=28=E8=92=8B=E8=89=BA=29?= Date: Sat, 26 Aug 2023 00:20:32 +0800 Subject: [PATCH 2/5] ci: fix ci error --- tf_pwa/config_loader/plot.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tf_pwa/config_loader/plot.py b/tf_pwa/config_loader/plot.py index e359cc19..157f53e5 100644 --- a/tf_pwa/config_loader/plot.py +++ b/tf_pwa/config_loader/plot.py @@ -1259,7 +1259,7 @@ def plot_function_2dpull( im = mpl.cm.ScalarMappable(norm=normal, cmap=my_cmap) # ax.colorbar(im) ax.get_figure().colorbar(im) - ax.title( + ax.set_title( "$\\chi^2/Nbins={:.2f}/{}$".format( np.sum(np.abs(pulls) ** 2), len(bound) ) From 3523c2ec30556b3ffe85a3cedbdafd2336291e51 Mon Sep 17 00:00:00 2001 From: jiangyi15 Date: Sat, 26 Aug 2023 11:13:25 +0800 Subject: [PATCH 3/5] fix: ci error --- tf_pwa/config_loader/plot.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/tf_pwa/config_loader/plot.py b/tf_pwa/config_loader/plot.py index 157f53e5..7dd3e92d 100644 --- a/tf_pwa/config_loader/plot.py +++ b/tf_pwa/config_loader/plot.py @@ -1151,10 +1151,10 @@ def get_dalitz(config, a, b): topb = i.core if pb == i.core: outs_b = i.outs - same_finals = [i for i in outs_a if i in db.outs] + same_finals = [i for i in outs_a if i in outs_b] p1 = [i for i in outs_a if i not in same_finals] p3 = [i for i in outs_b if i not in same_finals] - check = ((topa == topb),) + check = topa == topb check = check and len(same_finals) == 1 check = check and len(p1) == 1 check = check and len(p3) == 1 @@ -1184,7 +1184,14 @@ def get_dalitz_boundary(config, a, b, N=1000): @ConfigLoader.register_function() def plot_adaptive_2dpull( - config, var1, var2, binning=[[2, 2]] * 3, ax=plt, where={}, cut_zero=True + config, + var1, + var2, + binning=[[2, 2]] * 3, + ax=plt, + where={}, + cut_zero=True, + **kwargs ): import matplotlib as mpl import matplotlib.colors as mcolors @@ -1270,7 +1277,7 @@ def plot_function_2dpull( ax.set_ylabel(var2) config.plot_partial_wave( - plot_function=plot_function_2dpull, combine_plot=True + plot_function=plot_function_2dpull, combine_plot=True, **kwargs ) From 6bf9ca4f8b34b6a4b8c184cd5f2ed68d4e2a4285 Mon Sep 17 00:00:00 2001 From: jiangyi15 Date: Sat, 26 Aug 2023 13:53:00 +0800 Subject: [PATCH 4/5] feat: pull in 2dplot --- tf_pwa/adaptive_bins.py | 8 +- tf_pwa/config_loader/plot.py | 361 ++++++++++++++++------------------- tf_pwa/tests/config_cfit.yml | 2 +- tf_pwa/tests/test_full.py | 2 + 4 files changed, 173 insertions(+), 200 deletions(-) diff --git a/tf_pwa/adaptive_bins.py b/tf_pwa/adaptive_bins.py index fc5170ac..192257a5 100644 --- a/tf_pwa/adaptive_bins.py +++ b/tf_pwa/adaptive_bins.py @@ -17,7 +17,8 @@ class AdaptiveBound(object): """adaptive bound cut for data value""" - def __init__(self, base_data, bins): + def __init__(self, base_data, bins, base_bound=None): + self._base_bound = base_bound if isinstance(bins, int): self._base_data = np.array([base_data]) self.bins = [[bins]] @@ -33,9 +34,10 @@ def __init__(self, base_data, bins): @functools.lru_cache() def get_bounds_data(self): """get split data bounds, and the data after splitting""" - base_bound = AdaptiveBound.base_bound(self._base_data) + if self._base_bound is None: + self._base_bound = AdaptiveBound.base_bound(self._base_data) bounds, datas = AdaptiveBound.loop_split_bound( - self._base_data, self.bins + self._base_data, self.bins, self._base_bound ) return bounds, datas diff --git a/tf_pwa/config_loader/plot.py b/tf_pwa/config_loader/plot.py index 7dd3e92d..d3057457 100644 --- a/tf_pwa/config_loader/plot.py +++ b/tf_pwa/config_loader/plot.py @@ -822,24 +822,6 @@ def _plot_partial_wave( style.save() - self._2d_plot( - data_dict, - phsp_dict, - bg_dict, - prefix, - plot_var_dic, - chain_property, - plot_delta=plot_delta, - plot_pull=plot_pull, - save_pdf=save_pdf, - bin_scale=bin_scale, - single_legend=single_legend, - format=format, - nll=nll, - smooth=smooth, - color_first=color_first, - **kwargs, - ) self._2d_plot_v2( data_dict, phsp_dict, @@ -860,102 +842,6 @@ def _plot_partial_wave( ) -@ConfigLoader.register_function() -def _2d_plot( - self, - data_dict, - phsp_dict, - bg_dict, - prefix, - plot_var_dic, - chain_property, - plot_delta=False, - plot_pull=False, - save_pdf=False, - bin_scale=3, - single_legend=False, - format="png", - nll=None, - smooth=True, - color_first=True, - **kwargs -): - twodplot = self.config["plot"].get("2Dplot", {}) - for k, i in twodplot.items(): - if "&" not in k: - continue - var1, var2 = k.split("&") - var1 = var1.rstrip() - var2 = var2.lstrip() - k = var1 + "_vs_" + var2 - display = i.get("display", k) - plot_figs = i["plot_figs"] - name1, name2 = display.split("vs") - name1 = name1.rstrip() - name2 = name2.lstrip() - range1 = plot_var_dic[var1]["range"] - data_1 = data_dict[var1] - phsp_1 = phsp_dict[var1 + "_MC"] - range2 = plot_var_dic[var2]["range"] - data_2 = data_dict[var2] - phsp_2 = phsp_dict[var2 + "_MC"] - - def plot_axis(): - plt.xlabel(name1) - plt.ylabel(name2) - plt.title(display, fontsize="xx-large") - plt.xlim(range1) - plt.ylim(range2) - - # data - if "data" in plot_figs: - plt.scatter(data_1, data_2, s=1, alpha=0.8, label="data") - plot_axis() - plt.legend() - plt.savefig(prefix + k + "_data") - plt.clf() - print("Finish plotting 2D data " + prefix + k) # data - if "data_hist" in plot_figs: - plt.hist2d( - data_1, - data_2, - bins=100, - weights=data_dict["data_weights"], - cmin=1e-12, - ) - plot_axis() - plt.colorbar() - plt.savefig(prefix + k + "_data_hist") - plt.clf() - print("Finish plotting 2D data " + prefix + k) - # sideband - if "sideband" in plot_figs: - if bg_dict: - bg_1 = bg_dict[var1 + "_sideband"] - bg_2 = bg_dict[var2 + "_sideband"] - plt.scatter( - bg_1, bg_2, s=1, c="g", alpha=0.8, label="sideband" - ) - plot_axis() - plt.legend() - plt.savefig(prefix + k + "_bkg") - plt.clf() - print("Finish plotting 2D sideband " + prefix + k) - else: - print("There's no bkg input") - # fit pdf - if "fitted" in plot_figs: - phsp_weights = phsp_dict["MC_total_fit"] - plt.hist2d( - phsp_1, phsp_2, bins=100, weights=phsp_weights, cmin=1e-12 - ) - plot_axis() - plt.colorbar() - plt.savefig(prefix + k + "_fitted") - plt.clf() - print("Finish plotting 2D fitted " + prefix + k) - - def _plot_var_name(name): if isinstance(name, (list, tuple)): sub = name[0] @@ -1019,6 +905,37 @@ def _2d_plot_v2( **kwargs ): twodplot = self.config["plot"].get("2Dplot", {}) + new_plot = {} + for k, v in twodplot.items(): + if "&" in k: + var1, var2 = k.split("&") + var1 = var1.rstrip() + var2 = var2.lstrip() + k = var1 + "_vs_" + var2 + v["x"] = var1 + v["y"] = var2 + if ( + "xrange" not in v + and var1 in plot_var_dic + and "range" in plot_var_dic[var1] + ): + v["xrange"] = plot_var_dic[var1]["range"] + if ( + "yrange" not in v + and var2 in plot_var_dic + and "range" in plot_var_dic[var2] + ): + v["yrange"] = plot_var_dic[var2]["range"] + if "vs" in v.get("display", ""): + name1, name2 = v.get("display", "").split("vs") + name1 = name1.rstrip() + name2 = name2.lstrip() + if "xlabel" not in v: + v["xlabel"] = name1 + if "ylabel" not in v: + v["ylabel"] = name2 + new_plot[k] = v + twodplot.update(new_plot) for k, v in twodplot.items(): if "&" in k: continue @@ -1048,6 +965,7 @@ def _2d_plot_v2( y_bins = v.get("ybins", 100) display = v.get("display", k) + title = display plot_figs = v.get("plot_figs", ["data", "sidbanand", "fitted"]) name1 = v.get("xlabel", str(var_x)) @@ -1056,7 +974,6 @@ def _2d_plot_v2( def plot_axis(): plt.xlabel(name1) plt.ylabel(name2) - plt.title(display, fontsize="xx-large") plt.xlim(x_range) plt.ylim(y_range) @@ -1064,6 +981,7 @@ def plot_axis(): if "data" in plot_figs: plt.scatter(data_1, data_2, s=1, alpha=0.8, label="data") plot_axis() + plt.title(title, fontsize="xx-large") plt.savefig(prefix + k + "_data") plt.clf() print("Finish plotting 2D data " + prefix + k) @@ -1077,6 +995,7 @@ def plot_axis(): cmin=1e-12, ) plot_axis() + plt.title(title, fontsize="xx-large") plt.colorbar() plt.savefig(prefix + k + "_data_hist") plt.clf() @@ -1090,6 +1009,7 @@ def plot_axis(): bg_1, bg_2, s=1, c="g", alpha=0.8, label="sideband" ) plot_axis() + plt.title(title, fontsize="xx-large") plt.savefig(prefix + k + "_bkg") plt.clf() print("Finish plotting 2D sideband " + prefix + k) @@ -1109,6 +1029,7 @@ def plot_axis(): cmin=1e-12, ) plot_axis() + plt.title(title, fontsize="xx-large") plt.colorbar() plt.savefig(prefix + k + "_bkg_hist") plt.clf() @@ -1127,10 +1048,27 @@ def plot_axis(): cmin=1e-12, ) plot_axis() + plt.title(title, fontsize="xx-large") plt.colorbar() plt.savefig(prefix + k + "_fitted") plt.clf() print("Finish plotting 2D fitted " + prefix + k) + if "pull" in plot_figs: + n = max(int(np.log(data_1.shape[0] / 50) / np.log(4)), 2) + pull_binning = v.get("adaptive_binning", [[2, 2]] * n) + plot_function_2dpull( + data_dict, + phsp_dict, + bg_dict, + var1=v["x"], + var2=v["y"], + where=where, + binning=pull_binning, + ) + plot_axis() + plt.savefig(prefix + k + "_pull") + plt.clf() + print("Finish plotting 2D pull " + prefix + k) @ConfigLoader.register_function() @@ -1182,15 +1120,18 @@ def get_dalitz_boundary(config, a, b, N=1000): return s12, np.stack([s23_min, s23_max], axis=-1) -@ConfigLoader.register_function() -def plot_adaptive_2dpull( - config, - var1, - var2, +def plot_function_2dpull( + data_dict, + phsp_dict, + bg_dict, + var1="x", + var2="y", binning=[[2, 2]] * 3, - ax=plt, where={}, + ax=plt, cut_zero=True, + plot_scatter=True, + scatter_style={"s": 1, "c": "black"}, **kwargs ): import matplotlib as mpl @@ -1199,85 +1140,113 @@ def plot_adaptive_2dpull( from tf_pwa.adaptive_bins import AdaptiveBound - def plot_function_2dpull( - data_dict, phsp_dict, bg_dict, plot_var_dic, **kwargs - ): - nonlocal ax - if cut_zero: - cut = data_dict["data_weights"] != 0 - else: - cut = np.ones(data_dict["data_weights"].shape, dtype=np.bool) - (var_x_f, var_y_f), get_var = build_read_var_function( - [var1, var2], where=where - ) - x = var_x_f(**get_var(data_dict, ""))[cut] - y = var_y_f(**get_var(data_dict, ""))[cut] - w = data_dict["data_weights"][cut] - x_phsp = var_x_f(**get_var(phsp_dict, "_MC")) - y_phsp = var_y_f(**get_var(phsp_dict, "_MC")) - w_phsp = phsp_dict["MC_total_fit"] - data_cut = np.array([x, y]) - adapter = AdaptiveBound(data_cut, binning) - phsps = adapter.split_data(np.array([x_phsp, y_phsp, w_phsp])) - datas = adapter.split_data(np.array([x, y, w])) + if cut_zero: + cut = data_dict["data_weights"] != 0 + else: + cut = np.ones(data_dict["data_weights"].shape, dtype=np.bool) + (var_x_f, var_y_f), get_var = build_read_var_function( + [var1, var2], where=where + ) + x = var_x_f(**get_var(data_dict, ""))[cut] + y = var_y_f(**get_var(data_dict, ""))[cut] + w = data_dict["data_weights"][cut] + x_phsp = var_x_f(**get_var(phsp_dict, "_MC")) + y_phsp = var_y_f(**get_var(phsp_dict, "_MC")) + w_phsp = phsp_dict["MC_total_fit"] + data_cut = np.array([x, y]) + phsp_cut = np.array([x_phsp, y_phsp]) + base_bound = ( + np.min(phsp_cut, axis=-1) - 1e-6, + np.max(phsp_cut, axis=-1) + 1e-6, + ) + adapter = AdaptiveBound(data_cut, binning, base_bound) + phsps = adapter.split_data(np.array([x_phsp, y_phsp, w_phsp])) + datas = adapter.split_data(np.array([x, y, w])) + if bg_dict != {}: + x_bg = var_x_f(**get_var(bg_dict, "_sideband")) + y_bg = var_y_f(**get_var(bg_dict, "_sideband")) + w_bg = bg_dict["sideband_weights"] + bgs = adapter.split_data(np.array([x_bg, y_bg, w_bg])) + bound = adapter.get_bounds() + numbers = [] + pulls = [] + int_norm = 1 + for i, bnd in enumerate(bound): + min_x, min_y = bnd[0] + max_x, max_y = bnd[1] + ndata = np.sum(datas[i][2]) + nmc = np.sum(phsps[i][2]) if bg_dict != {}: - x_bg = var_x_f(**get_var(bg_dict, "_sideband")) - y_bg = var_y_f(**get_var(bg_dict, "_sideband")) - w_bg = bg_dict["sideband_weights"] - bgs = adapter.split_data(np.array([x_bg, y_bg, w_bg])) - bound = adapter.get_bounds() - numbers = [] - pulls = [] - int_norm = 1 - for i, bnd in enumerate(bound): - min_x, min_y = bnd[0] - max_x, max_y = bnd[1] - ndata = np.sum(datas[i][2]) - nmc = np.sum(phsps[i][2]) - if bg_dict != {}: - nmc += np.sum(bgs[i][2]) - numbers.append((ndata, nmc)) - pulls.append((ndata - nmc) / np.sqrt(nmc)) - - max_weight = max(np.max(np.abs(pulls)), 5) - - my_cmap = plt.get_cmap("jet") - if ax == plt: - ax = plt.gca() # fig, ax = plt.subplots() - ax.scatter(x, y, s=1, c="black") - for i, bnd in enumerate(bound): - min_x, min_y = bnd[0] - max_x, max_y = bnd[1] - # print(weights[i]) # max_weight) - rect = mpathes.Rectangle( - (min_x, min_y), - max_x - min_x, - max_y - min_y, - linewidth=1, - facecolor=my_cmap( - pulls[i] / max_weight / 2 + 0.5 - ), # max_weight), - edgecolor="none", # black", - zorder=-1, - ) # cmap(weights[i]/max_weight)) - ax.add_patch(rect) - - normal = mpl.colors.Normalize(vmin=-max_weight, vmax=max_weight) - im = mpl.cm.ScalarMappable(norm=normal, cmap=my_cmap) - # ax.colorbar(im) - ax.get_figure().colorbar(im) - ax.set_title( - "$\\chi^2/Nbins={:.2f}/{}$".format( - np.sum(np.abs(pulls) ** 2), len(bound) - ) + nmc += np.sum(bgs[i][2]) + numbers.append((ndata, nmc)) + pulls.append((ndata - nmc) / np.sqrt(nmc)) + + max_weight = max(np.max(np.abs(pulls)), 5) + + my_cmap = plt.get_cmap("jet") + if ax == plt: + ax = plt.gca() # fig, ax = plt.subplots() + if plot_scatter: + ax.scatter(x, y, **scatter_style) + for i, bnd in enumerate(bound): + min_x, min_y = bnd[0] + max_x, max_y = bnd[1] + # print(weights[i]) # max_weight) + rect = mpathes.Rectangle( + (min_x, min_y), + max_x - min_x, + max_y - min_y, + linewidth=1, + facecolor=my_cmap(pulls[i] / max_weight / 2 + 0.5), # max_weight), + edgecolor="none", # black", + zorder=-1, + ) # cmap(weights[i]/max_weight)) + ax.add_patch(rect) + + normal = mpl.colors.Normalize(vmin=-max_weight, vmax=max_weight) + im = mpl.cm.ScalarMappable(norm=normal, cmap=my_cmap) + # ax.colorbar(im) + ax.get_figure().colorbar(im) + ax.set_title( + "$\\chi^2/Nbins={:.2f}/{}$".format( + np.sum(np.abs(pulls) ** 2), len(bound) ) - ax.set_xlim([np.min(x_phsp), np.max(x_phsp)]) - ax.set_ylim([np.min(y_phsp), np.max(y_phsp)]) - ax.set_xlabel(var1) - ax.set_ylabel(var2) + ) + ax.set_xlim([np.min(x_phsp), np.max(x_phsp)]) + ax.set_ylim([np.min(y_phsp), np.max(y_phsp)]) + ax.set_xlabel(var1) + ax.set_ylabel(var2) + + +@ConfigLoader.register_function() +def plot_adaptive_2dpull( + config, + var1, + var2, + binning=[[2, 2]] * 3, + ax=plt, + where={}, + cut_zero=True, + plot_scatter=True, + scatter_style={"s": 1, "c": "black"}, + **kwargs +): + + pull_kwargs = { + "var1": var1, + "var2": var2, + "binning": binning, + "where": where, + "plot_scatter": plot_scatter, + "scatter_style": scatter_style, + "cut_zero": cut_zero, + } + + def my_plot_function_2dpull(*args, **kwargs): + plot_function_2dpull(*args, **pull_kwargs, **kwargs) config.plot_partial_wave( - plot_function=plot_function_2dpull, combine_plot=True, **kwargs + plot_function=my_plot_function_2dpull, combine_plot=True, **kwargs ) diff --git a/tf_pwa/tests/config_cfit.yml b/tf_pwa/tests/config_cfit.yml index 81a25f53..29cff6c9 100644 --- a/tf_pwa/tests/config_cfit.yml +++ b/tf_pwa/tests/config_cfit.yml @@ -72,4 +72,4 @@ plot: costheta: [angle, R_BC, cos(beta)] xbins: 50 ybins: 50 - plot_figs: ["data", "sideband_hist", "fitted"] + plot_figs: ["data", "sideband_hist", "fitted", "pull"] diff --git a/tf_pwa/tests/test_full.py b/tf_pwa/tests/test_full.py index b7c6218d..d4803cbd 100644 --- a/tf_pwa/tests/test_full.py +++ b/tf_pwa/tests/test_full.py @@ -363,6 +363,8 @@ def test_plot_2dpull(toy_config): import matplotlib.pyplot as plt toy_config.plot_adaptive_2dpull("m_R_BC**2", "m_R_CD**2") + with pytest.raises(AssertionError): + a, b = toy_config.get_dalitz_boundary("R_BC", "R_BC") a, b = toy_config.get_dalitz_boundary("R_BC", "R_CD") plt.plot(a, b, color="red") plt.savefig("adptive_2d.png") From 2b6b5237818f177b819f9ae8bbc84562b679a49d Mon Sep 17 00:00:00 2001 From: jiangyi15 Date: Sat, 26 Aug 2023 14:03:06 +0800 Subject: [PATCH 5/5] misc: more option for pull --- tf_pwa/config_loader/plot.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tf_pwa/config_loader/plot.py b/tf_pwa/config_loader/plot.py index d3057457..641b2d57 100644 --- a/tf_pwa/config_loader/plot.py +++ b/tf_pwa/config_loader/plot.py @@ -1056,6 +1056,10 @@ def plot_axis(): if "pull" in plot_figs: n = max(int(np.log(data_1.shape[0] / 50) / np.log(4)), 2) pull_binning = v.get("adaptive_binning", [[2, 2]] * n) + pull_scatter_style = v.get( + "pull_scatter_style", {"c": "black", "s": 1} + ) + pull_cmap = v.get("pull_cmap", "jet") plot_function_2dpull( data_dict, phsp_dict, @@ -1064,6 +1068,7 @@ def plot_axis(): var2=v["y"], where=where, binning=pull_binning, + scatter_style=pull_scatter_style, ) plot_axis() plt.savefig(prefix + k + "_pull") @@ -1132,6 +1137,7 @@ def plot_function_2dpull( cut_zero=True, plot_scatter=True, scatter_style={"s": 1, "c": "black"}, + cmap="jet", **kwargs ): import matplotlib as mpl @@ -1183,7 +1189,7 @@ def plot_function_2dpull( max_weight = max(np.max(np.abs(pulls)), 5) - my_cmap = plt.get_cmap("jet") + my_cmap = plt.get_cmap(cmap) if ax == plt: ax = plt.gca() # fig, ax = plt.subplots() if plot_scatter: