reduce memory usage by creating weights dict early

2 files changed, 19 insertions, 19 deletions

diff --git a/utils/chi2 b/utils/chi2
index f4d94c4..ac4093a 100755
--- a/utils/chi2
+++ b/utils/chi2
@@ -406,13 +406,11 @@ def do_fit(data,muon,data_mc,weights,atmo_scale_factor,muon_scale_factor,bins,st
     xopt = opt.optimize(x0)
 
     # Get the total number of "universes" simulated in the GENIE reweight tool
-    nuniverses = weights['universe'].max()+1
-
-    weights_dict = dict(tuple(weights.groupby('universe')))
+    nuniverses = max(weights.keys())+1
 
     nlls = []
     for universe in range(nuniverses):
-        data_mc_with_weights = pd.merge(data_mc,weights_dict[universe],how='left',on=['run','unique_id'])
+        data_mc_with_weights = pd.merge(data_mc,weights[universe],how='left',on=['run','unique_id'])
         data_mc_with_weights.weight = data_mc_with_weights.weight.fillna(1.0)
 
         nll = make_nll(data,muon,data_mc_with_weights,atmo_scale_factor,muon_scale_factor,bins,reweight=True,print_nll=print_nll)
@@ -421,7 +419,7 @@ def do_fit(data,muon,data_mc,weights,atmo_scale_factor,muon_scale_factor,bins,st
     universe = np.argmin(nlls)
 
     if refit:
-        data_mc_with_weights = pd.merge(data_mc,weights[weights.universe == universe],how='left',on=['run','unique_id'])
+        data_mc_with_weights = pd.merge(data_mc,weights[universe],how='left',on=['run','unique_id'])
         data_mc_with_weights.weight = data_mc_with_weights.weight.fillna(1.0)
 
         # Create a new negative log likelihood function with the weighted Monte Carlo.
@@ -587,6 +585,8 @@ if __name__ == '__main__':
     # 15752  154       957 -0.006827
     weights = weights[weights.weight > 0]
 
+    weights = dict(tuple(weights.groupby('universe')))
+
     ev_mc = correct_energy_bias(ev_mc)
     muon_mc = correct_energy_bias(muon_mc)
 
@@ -651,8 +651,8 @@ if __name__ == '__main__':
 
         xtrue = truncnorm_scaled(PRIORS_LOW,PRIORS_HIGH,PRIORS,PRIOR_UNCERTAINTIES)
 
-        data_mc_with_weights = pd.merge(data_mc,weights[weights.universe == 0],how='left',on=['run','unique_id'])
-        data_atm_mc_with_weights = pd.merge(data_atm_mc,weights[weights.universe == 0],how='left',on=['run','unique_id'])
+        data_mc_with_weights = pd.merge(data_mc,weights[0],how='left',on=['run','unique_id'])
+        data_atm_mc_with_weights = pd.merge(data_atm_mc,weights[0],how='left',on=['run','unique_id'])
 
         data_mc_with_weights.weight *= data_mc_with_weights.flux_weight
         data_atm_mc_with_weights.weight *= data_atm_mc_with_weights.flux_weight
@@ -689,10 +689,10 @@ if __name__ == '__main__':
 
             xopt, universe, samples = do_fit(data,muon,data_mc,weights,atmo_scale_factor,muon_scale_factor,bins,args.steps,args.print_nll,args.walkers,args.thin)
 
-            data_mc_with_weights = pd.merge(data_mc,weights[weights.universe == universe],how='left',on=['run','unique_id'])
+            data_mc_with_weights = pd.merge(data_mc,weights[universe],how='left',on=['run','unique_id'])
             data_mc_with_weights.weight = data_mc_with_weights.weight.fillna(1.0)
 
-            data_atm_mc_with_weights = pd.merge(data_atm_mc,weights[weights.universe == universe],how='left',on=['run','unique_id'])
+            data_atm_mc_with_weights = pd.merge(data_atm_mc,weights[universe],how='left',on=['run','unique_id'])
             data_atm_mc_with_weights.weight = data_atm_mc_with_weights.weight.fillna(1.0)
 
             prob = get_prob(data,muon,data_mc_with_weights,atmo_scale_factor,muon_scale_factor,samples,bins,size=args.multinomial_prob_size)
@@ -821,10 +821,10 @@ if __name__ == '__main__':
 
     xopt, universe, samples = do_fit(data,muon,data_mc,weights,atmo_scale_factor,muon_scale_factor,bins,args.steps,args.print_nll,args.walkers,args.thin)
 
-    data_mc_with_weights = pd.merge(data_mc,weights[weights.universe == universe],how='left',on=['run','unique_id'])
+    data_mc_with_weights = pd.merge(data_mc,weights[universe],how='left',on=['run','unique_id'])
     data_mc_with_weights.weight = data_mc_with_weights.weight.fillna(1.0)
 
-    data_atm_mc_with_weights = pd.merge(data_atm_mc,weights[weights.universe == universe],how='left',on=['run','unique_id'])
+    data_atm_mc_with_weights = pd.merge(data_atm_mc,weights[universe],how='left',on=['run','unique_id'])
     data_atm_mc_with_weights.weight = data_atm_mc_with_weights.weight.fillna(1.0)
 
     prob = get_prob(data,muon,data_mc_with_weights,atmo_scale_factor,muon_scale_factor,samples,bins,size=args.multinomial_prob_size)
diff --git a/utils/dm-search b/utils/dm-search
index cd91eab..9e1a136 100755
--- a/utils/dm-search
+++ b/utils/dm-search
@@ -347,13 +347,11 @@ def do_fit(dm_particle_id,dm_mass,dm_energy,data,muon,data_mc,weights,atmo_scale
     xopt = opt.optimize(x0)
 
     # Get the total number of "universes" simulated in the GENIE reweight tool
-    nuniverses = weights['universe'].max()+1
-
-    weights_dict = dict(tuple(weights.groupby('universe')))
+    nuniverses = max(weights.keys())+1
 
     nlls = []
     for universe in range(nuniverses):
-        data_mc_with_weights = pd.merge(data_mc,weights_dict[universe],how='left',on=['run','unique_id'])
+        data_mc_with_weights = pd.merge(data_mc,weights[universe],how='left',on=['run','unique_id'])
         data_mc_with_weights.weight = data_mc_with_weights.weight.fillna(1.0)
 
         nll = make_nll(dm_particle_id,dm_mass,dm_energy,data,muon,data_mc_with_weights,atmo_scale_factor,muon_scale_factor,bins,reweight=True,print_nll=print_nll,dm_sample=dm_sample)
@@ -362,7 +360,7 @@ def do_fit(dm_particle_id,dm_mass,dm_energy,data,muon,data_mc,weights,atmo_scale
     universe = np.argmin(nlls)
 
     if refit:
-        data_mc_with_weights = pd.merge(data_mc,weights[weights.universe == universe],how='left',on=['run','unique_id'])
+        data_mc_with_weights = pd.merge(data_mc,weights[universe],how='left',on=['run','unique_id'])
         data_mc_with_weights.weight = data_mc_with_weights.weight.fillna(1.0)
 
         # Create a new negative log likelihood function with the weighted Monte Carlo.
@@ -470,7 +468,7 @@ def get_limits(dm_masses,data,muon,data_mc,atmo_scale_factor,muon_scale_factor,b
             dm_energy = dm_mass
             xopt, universe, samples = do_fit(dm_particle_id,dm_mass,dm_energy,data,muon,data_mc,weights,atmo_scale_factor,muon_scale_factor,bins,steps,print_nll,walkers,thin)
 
-            data_mc_with_weights = pd.merge(data_mc,weights[weights.universe == universe],how='left',on=['run','unique_id'])
+            data_mc_with_weights = pd.merge(data_mc,weights[universe],how='left',on=['run','unique_id'])
             data_mc_with_weights.weight = data_mc_with_weights.weight.fillna(1.0)
 
             limit = np.percentile(samples[:,6],90)
@@ -649,6 +647,8 @@ if __name__ == '__main__':
     # 15752  154       957 -0.006827
     weights = weights[weights.weight > 0]
 
+    weights = dict(tuple(weights.groupby('universe')))
+
     ev_mc = correct_energy_bias(ev_mc)
     muon_mc = correct_energy_bias(muon_mc)
 
@@ -785,8 +785,8 @@ if __name__ == '__main__':
         # Set the random seed so we get reproducible results here
         np.random.seed(0)
 
-        data_mc_with_weights = pd.merge(data_mc,weights[weights.universe == 0],how='left',on=['run','unique_id'])
-        data_atm_mc_with_weights = pd.merge(data_atm_mc,weights[weights.universe == 0],how='left',on=['run','unique_id'])
+        data_mc_with_weights = pd.merge(data_mc,weights[0],how='left',on=['run','unique_id'])
+        data_atm_mc_with_weights = pd.merge(data_atm_mc,weights[0],how='left',on=['run','unique_id'])
 
         discoveries = 0