From 56dfcc60fdf60647d8e36499e4bebbccc0ec849a Mon Sep 17 00:00:00 2001
From: dermen <dermen@lbl.gov>
Date: Mon, 4 Oct 2021 14:09:58 -0700
Subject: [PATCH] frees memory after running prediction algorithm

---
 simtbx/command_line/hopper_process.py | 11 +++++++++--
 simtbx/diffBragg/hopper_utils.py      |  6 ++++--
 simtbx/modeling/forward_models.py     |  6 ++++++
 3 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/simtbx/command_line/hopper_process.py b/simtbx/command_line/hopper_process.py
index 354876093f..1cc38aca9a 100644
--- a/simtbx/command_line/hopper_process.py
+++ b/simtbx/command_line/hopper_process.py
@@ -45,6 +45,9 @@
   .type = bool
   .help = if True, save the data modelers after running refinement. The file includes model values, errors, and useful information
   .help = for examining the modeling results, can be loaded using modeler = np.load(fname, allow_pickle=True)[()]
+refspec = None
+  .type = str
+  .help = path to a reference .lam file to use as the spectra for each shot
 """
 import os
 from libtbx.phil import parse
@@ -69,10 +72,11 @@ def __init__(self, *args, **kwargs):
         logging.basicConfig(level=logging.DEBUG)
 
         if self.params.save_modelers:
+            self.modeler_dir = os.path.join(self.params.output.output_dir, "modelers")
             if COMM.rank == 0:
-                self.modeler_dir = os.path.join(self.params.output.output_dir, "modelers")
                 if not os.path.exists(self.modeler_dir):
                     os.makedirs(self.modeler_dir)
+            COMM.barrier()
 
     @property
     def device_id(self):
@@ -133,7 +137,10 @@ def refine(self, exps, ref):
             assert len(exps)==1
             # TODO MPI select GPU device
 
-            exp, ref, data_modeler, x = refine(exps[0], ref, self.params.diffBragg, gpu_device=self.device_id, return_modeler=True)
+            exp, ref, data_modeler, x = refine(exps[0], ref,
+                                               self.params.diffBragg,
+                                               spec=self.params.refspec,
+                                               gpu_device=self.device_id, return_modeler=True)
             orig_exp_name = os.path.abspath(self.params.output.refined_experiments_filename)
             refls_name = os.path.abspath(self.params.output.indexed_filename)
             self.params.diffBragg.outdir = self.params.output.output_dir
diff --git a/simtbx/diffBragg/hopper_utils.py b/simtbx/diffBragg/hopper_utils.py
index 5b02d6bb71..b057807d8e 100644
--- a/simtbx/diffBragg/hopper_utils.py
+++ b/simtbx/diffBragg/hopper_utils.py
@@ -45,7 +45,8 @@ class DataModeler:
     2- lists whose length is the number of shoeboxes being modeled
 
     for example if one is modeling 3 shoeboxes whose dimensions are 10x10, then
-    the objects below like self.all_data will have length 300, and other objects like self.selection_flags will have length 2
+    the objects below like self.all_data will have length 300, and other objects like self.selection_flags
+    will have length 3
     """
 
     def __init__(self, params):
@@ -76,12 +77,13 @@ def __init__(self, params):
         self.simple_weights = None  # not used
         self.refls_idx = None  # position of modeled spot in original refl array
         self.refls = None  # reflection table
+        self.sigma_rdout = None   # the value of the readout noise in photon units
 
         self.Hi = None  # miller index (P1)
         self.Hi_asu = None  # miller index (high symmetry)
 
         # which attributes to save when pickling a data modeler
-        self.saves = ["all_data", "all_background", "all_trusted", "best_model", "all_sigmas",
+        self.saves = ["all_data", "all_background", "all_trusted", "best_model", "sigma_rdout",
                       "rois", "pids", "tilt_abc", "selection_flags", "refls_idx", "pan_fast_slow",
                         "Hi", "Hi_asu", "roi_id", "params", "all_pid", "all_fast", "all_slow"]
 
diff --git a/simtbx/modeling/forward_models.py b/simtbx/modeling/forward_models.py
index a11bcdd556..992ff06510 100644
--- a/simtbx/modeling/forward_models.py
+++ b/simtbx/modeling/forward_models.py
@@ -260,6 +260,12 @@ def diffBragg_forward(CRYSTAL, DETECTOR, BEAM, Famp, energies, fluxes,
         S.D.printout_pixel_fastslow = f,s
         S.D.show_params()
         S.D.add_diffBragg_spots(printout_pix)
+
+    # free up memory
+    S.D.free_all()
+    S.D.free_Fhkl2()
+    if S.D.gpu_free is not None:
+        S.D.gpu_free()
     return data