diff options
Diffstat (limited to 'utils')
-rwxr-xr-x | utils/dc | 30 | ||||
-rwxr-xr-x | utils/dc-closure-test | 30 | ||||
-rwxr-xr-x | utils/sddm/plot_energy.py | 6 |
3 files changed, 49 insertions, 17 deletions
@@ -96,7 +96,7 @@ flasher_r_udotr = Constraint(range(39,42)) # others must add up to less than 1. breakdown_r_udotr = Constraint(range(44,47)) -def make_nll(data, sacrifice, constraints): +def make_nll(data, sacrifice, constraints, fitted_fraction): def nll(x, grad=None, fill_value=1e9): if grad is not None and grad.size > 0: raise Exception("nll got passed grad!") @@ -183,7 +183,7 @@ def make_nll(data, sacrifice, constraints): [p_r_psi_z_udotr_muon_hilohilo, p_r_psi_z_udotr_muon_hilohihi]], \ [[p_r_psi_z_udotr_muon_hihilolo, p_r_psi_z_udotr_muon_hihilohi], \ [p_r_psi_z_udotr_muon_hihihilo, p_r_psi_z_udotr_muon_hihihihi]]]]) - expected_muon = p_muon*contamination_muon*mu_muon + sacrifice['muon']*mu_signal + expected_muon = p_muon*contamination_muon*mu_muon*fitted_fraction['muon'] + sacrifice['muon']*mu_signal nll -= fast_poisson_logpmf(data['muon'],expected_muon).sum() @@ -194,7 +194,7 @@ def make_nll(data, sacrifice, constraints): [p_z_udotr_noise_lolo,p_z_udotr_noise_lohi], [p_z_udotr_noise_hilo,p_z_udotr_noise_hihi]]) p_noise = p_r_noise[:,np.newaxis,np.newaxis,np.newaxis]*p_psi_noise[:,np.newaxis,np.newaxis]*p_z_udotr_noise - expected_noise = p_noise*contamination_noise*mu_noise + sacrifice['noise']*mu_signal + expected_noise = p_noise*contamination_noise*mu_noise*fitted_fraction['noise'] + sacrifice['noise']*mu_signal nll -= fast_poisson_logpmf(data['noise'],expected_noise).sum() @@ -207,9 +207,9 @@ def make_nll(data, sacrifice, constraints): [p_r_z_udotr_neck_hihilo, p_r_z_udotr_neck_hihihi]]]) p_psi_neck = np.array([p_psi_neck_lo,1-p_psi_neck_lo]) p_neck = p_r_z_udotr_neck[:,np.newaxis,:,:]*p_psi_neck[:,np.newaxis,np.newaxis] - expected_neck = p_neck*contamination_neck*mu_neck + sacrifice['neck']*mu_signal + expected_neck = p_neck*contamination_neck*mu_neck*fitted_fraction['neck'] + sacrifice['neck']*mu_signal # FIXME: pdf should be different for muon given neck - expected_neck += p_muon*p_neck_given_muon*mu_muon + expected_neck += p_muon*p_neck_given_muon*mu_muon*fitted_fraction['neck'] nll -= fast_poisson_logpmf(data['neck'],expected_neck).sum() @@ -220,7 +220,7 @@ def make_nll(data, sacrifice, constraints): p_psi_flasher = np.array([p_psi_flasher_lo,1-p_psi_flasher_lo]) p_z_flasher = np.array([p_z_flasher_lo,1-p_z_flasher_lo]) p_flasher = p_r_udotr_flasher[:,np.newaxis,np.newaxis,:]*p_psi_flasher[:,np.newaxis,np.newaxis]*p_z_flasher[:,np.newaxis] - expected_flasher = p_flasher*contamination_flasher*mu_flasher + sacrifice['flasher']*mu_signal + expected_flasher = p_flasher*contamination_flasher*mu_flasher*fitted_fraction['flasher'] + sacrifice['flasher']*mu_signal nll -= fast_poisson_logpmf(data['flasher'],expected_flasher).sum() @@ -231,7 +231,7 @@ def make_nll(data, sacrifice, constraints): p_psi_breakdown = np.array([p_psi_breakdown_lo,1-p_psi_breakdown_lo]) p_z_breakdown = np.array([p_z_breakdown_lo,1-p_z_breakdown_lo]) p_breakdown = p_r_udotr_breakdown[:,np.newaxis,np.newaxis,:]*p_psi_breakdown[:,np.newaxis,np.newaxis]*p_z_breakdown[:,np.newaxis] - expected_breakdown = p_breakdown*contamination_breakdown*mu_breakdown + sacrifice['breakdown']*mu_signal + expected_breakdown = p_breakdown*contamination_breakdown*mu_breakdown*fitted_fraction['breakdown'] + sacrifice['breakdown']*mu_signal nll -= fast_poisson_logpmf(data['breakdown'],expected_breakdown).sum() @@ -283,6 +283,20 @@ if __name__ == '__main__': ev = ev[ev.prompt] ev = ev[ev.nhit_cal > 100] + fitted_fraction = {} + for bg in ['signal','muon','noise','neck','flasher','breakdown']: + if np.count_nonzero(ev[bg]): + fitted_fraction[bg] = np.count_nonzero(ev[bg] & ~np.isnan(ev.fmin))/np.count_nonzero(ev[bg]) + print("Fitted fraction for %s: %.0f %%" % (bg,fitted_fraction[bg]*100)) + elif bg != 'signal': + print_warning("Warning: No %s events in sample, assuming 10%% fitted fraction" % bg) + fitted_fraction[bg] = 0.1 + else: + print_warning("Warning: No signal events in sample, assuming 100%% fitted fraction") + fitted_fraction[bg] = 1.0 + + ev = ev[~np.isnan(ev.fmin)] + # figure out bins for high level variables ev = radius_cut(ev) ev = psi_cut(ev) @@ -332,7 +346,7 @@ if __name__ == '__main__': sacrifice[bg] /= len(ev_mc) constraints = [flasher_r_udotr, breakdown_r_udotr,muon_r_psi_z_udotr,neck_r_z_udotr,noise_z_udotr] - nll = make_nll(data,sacrifice,constraints) + nll = make_nll(data,sacrifice,constraints,fitted_fraction) x0 = [] for bg in ['signal','muon','noise','neck','flasher','breakdown']: diff --git a/utils/dc-closure-test b/utils/dc-closure-test index 5b9f77d..ba5daf8 100755 --- a/utils/dc-closure-test +++ b/utils/dc-closure-test @@ -96,7 +96,7 @@ flasher_r_udotr = Constraint(range(39,42)) # others must add up to less than 1. breakdown_r_udotr = Constraint(range(44,47)) -def make_nll(data, sacrifice, constraints): +def make_nll(data, sacrifice, constraints, fitted_fraction): def nll(x, grad=None, fill_value=1e9): if grad is not None and grad.size > 0: raise Exception("nll got passed grad!") @@ -183,7 +183,7 @@ def make_nll(data, sacrifice, constraints): [p_r_psi_z_udotr_muon_hilohilo, p_r_psi_z_udotr_muon_hilohihi]], \ [[p_r_psi_z_udotr_muon_hihilolo, p_r_psi_z_udotr_muon_hihilohi], \ [p_r_psi_z_udotr_muon_hihihilo, p_r_psi_z_udotr_muon_hihihihi]]]]) - expected_muon = p_muon*contamination_muon*mu_muon + sacrifice['muon']*mu_signal + expected_muon = p_muon*contamination_muon*mu_muon*fitted_fraction['muon'] + sacrifice['muon']*mu_signal nll -= fast_poisson_logpmf(data['muon'],expected_muon).sum() @@ -194,7 +194,7 @@ def make_nll(data, sacrifice, constraints): [p_z_udotr_noise_lolo,p_z_udotr_noise_lohi], [p_z_udotr_noise_hilo,p_z_udotr_noise_hihi]]) p_noise = p_r_noise[:,np.newaxis,np.newaxis,np.newaxis]*p_psi_noise[:,np.newaxis,np.newaxis]*p_z_udotr_noise - expected_noise = p_noise*contamination_noise*mu_noise + sacrifice['noise']*mu_signal + expected_noise = p_noise*contamination_noise*mu_noise*fitted_fraction['noise'] + sacrifice['noise']*mu_signal nll -= fast_poisson_logpmf(data['noise'],expected_noise).sum() @@ -207,9 +207,9 @@ def make_nll(data, sacrifice, constraints): [p_r_z_udotr_neck_hihilo, p_r_z_udotr_neck_hihihi]]]) p_psi_neck = np.array([p_psi_neck_lo,1-p_psi_neck_lo]) p_neck = p_r_z_udotr_neck[:,np.newaxis,:,:]*p_psi_neck[:,np.newaxis,np.newaxis] - expected_neck = p_neck*contamination_neck*mu_neck + sacrifice['neck']*mu_signal + expected_neck = p_neck*contamination_neck*mu_neck*fitted_fraction['neck'] + sacrifice['neck']*mu_signal # FIXME: pdf should be different for muon given neck - expected_neck += p_muon*p_neck_given_muon*mu_muon + expected_neck += p_muon*p_neck_given_muon*mu_muon*fitted_fraction['neck'] nll -= fast_poisson_logpmf(data['neck'],expected_neck).sum() @@ -220,7 +220,7 @@ def make_nll(data, sacrifice, constraints): p_psi_flasher = np.array([p_psi_flasher_lo,1-p_psi_flasher_lo]) p_z_flasher = np.array([p_z_flasher_lo,1-p_z_flasher_lo]) p_flasher = p_r_udotr_flasher[:,np.newaxis,np.newaxis,:]*p_psi_flasher[:,np.newaxis,np.newaxis]*p_z_flasher[:,np.newaxis] - expected_flasher = p_flasher*contamination_flasher*mu_flasher + sacrifice['flasher']*mu_signal + expected_flasher = p_flasher*contamination_flasher*mu_flasher*fitted_fraction['flasher'] + sacrifice['flasher']*mu_signal nll -= fast_poisson_logpmf(data['flasher'],expected_flasher).sum() @@ -231,7 +231,7 @@ def make_nll(data, sacrifice, constraints): p_psi_breakdown = np.array([p_psi_breakdown_lo,1-p_psi_breakdown_lo]) p_z_breakdown = np.array([p_z_breakdown_lo,1-p_z_breakdown_lo]) p_breakdown = p_r_udotr_breakdown[:,np.newaxis,np.newaxis,:]*p_psi_breakdown[:,np.newaxis,np.newaxis]*p_z_breakdown[:,np.newaxis] - expected_breakdown = p_breakdown*contamination_breakdown*mu_breakdown + sacrifice['breakdown']*mu_signal + expected_breakdown = p_breakdown*contamination_breakdown*mu_breakdown*fitted_fraction['breakdown'] + sacrifice['breakdown']*mu_signal nll -= fast_poisson_logpmf(data['breakdown'],expected_breakdown).sum() @@ -262,7 +262,7 @@ def make_nll(data, sacrifice, constraints): def fit(data, sacrifice, steps): constraints = [flasher_r_udotr, breakdown_r_udotr,muon_r_psi_z_udotr,neck_r_z_udotr,noise_z_udotr] - nll = make_nll(data,sacrifice,constraints) + nll = make_nll(data,sacrifice,constraints,fitted_fraction) x0 = [] for bg in ['signal','muon','noise','neck','flasher','breakdown']: @@ -435,6 +435,20 @@ if __name__ == '__main__': ev = ev[ev.prompt] ev = ev[ev.nhit_cal > 100] + fitted_fraction = {} + for bg in ['signal','muon','noise','neck','flasher','breakdown']: + if np.count_nonzero(ev[bg]): + fitted_fraction[bg] = np.count_nonzero(ev[bg] & ~np.isnan(ev.fmin))/np.count_nonzero(ev[bg]) + print("Fitted fraction for %s: %.0f %%" % (bg,fitted_fraction[bg]*100)) + elif bg != 'signal': + print_warning("Warning: No %s events in sample, assuming 10%% fitted fraction" % bg) + fitted_fraction[bg] = 0.1 + else: + print_warning("Warning: No signal events in sample, assuming 100%% fitted fraction") + fitted_fraction[bg] = 1.0 + + ev = ev[~np.isnan(ev.fmin)] + # figure out bins for high level variables ev = radius_cut(ev) ev = psi_cut(ev) diff --git a/utils/sddm/plot_energy.py b/utils/sddm/plot_energy.py index b4882d9..5bb464e 100755 --- a/utils/sddm/plot_energy.py +++ b/utils/sddm/plot_energy.py @@ -598,7 +598,11 @@ def get_events(filenames, merge_fits=False): ev = ev.groupby('run',group_keys=False).apply(tag_michels) if merge_fits: - ev = pd.merge(fits,ev,how='inner',on=['run','gtid']) + ev = pd.merge(ev,fits,how='left',on=['run','gtid']) + # Reset n, id1, id2, and id3 columns to integers + for column in ['n','id1','id2','id3']: + ev[column] = ev[column].fillna(0) + ev[column] = ev[column].astype(np.int32) # Set the index to (run, gtid) so we can set columns from the single particle results ev = ev.set_index(['run','gtid']) ev_single_particle = ev[(ev.id2 == 0) & (ev.id3 == 0)] |