[fix] Lazy access for PSyGrid initial/final values to save RAM

posydon/binary_evol/CE/step_CEE.py

@@ -74,7 +74,8 @@
                                               # assuming a final separation where the inner core RLOF starts.
          # "one_phase_variable_core_definition" for core_definition_H_fraction=0.01
          "metallicity": None,
-         "record_matching": False
+         "record_matching": False,
+         "track_matcher": None
          }
 
 
@@ -190,18 +191,19 @@ def __init__(
                         ["log_min_max", "min_max", "min_max"],
                         [0.1, 300], [0.0, None]
                     ]
-        self.track_matcher = TrackMatcher(grid_name_Hrich = None,
-                                    grid_name_strippedHe = None,
-                                    path=PATH_TO_POSYDON_DATA,
-                                    metallicity = self.metallicity,
-                                    matching_method = "minimize",
-                                    matching_tolerance=1e-2,
-                                    matching_tolerance_hard=1e-1,
-                                    list_for_matching_HMS = list_for_matching_HMS,
-                                    list_for_matching_HeStar = None,
-                                    list_for_matching_postMS = None,
-                                    record_matching = self.record_matching,
-                                    verbose = self.verbose)
+        #self.track_matcher = TrackMatcher(grid_name_Hrich = None,
+        #                            grid_name_strippedHe = None,
+        #                            path=PATH_TO_POSYDON_DATA,
+        #                            metallicity = self.metallicity,
+        #                            matching_method = "minimize",
+        #                            matching_tolerance=1e-2,
+        #                            matching_tolerance_hard=1e-1,
+        #                            list_for_matching_HMS = list_for_matching_HMS,
+        #                            list_for_matching_HeStar = None,
+        #                            list_for_matching_postMS = None,
+        #                            record_matching = self.record_matching,
+        #                            verbose = self.verbose)
+        self.track_matcher = kwargs.get('track_matcher', None)
     def __call__(self, binary):
         """Perform the CEE step for a BinaryStar object."""
         # Determine which star is the donor and which is the companion

posydon/binary_evol/DT/step_detached.py

@@ -36,7 +36,8 @@
     RVJ83_braking,
 )
 from posydon.binary_evol.DT.tides.default_tides import default_tides
-from posydon.binary_evol.DT.track_match import TrackMatcher
+
+#from posydon.binary_evol.DT.track_match import TrackMatcher
 from posydon.binary_evol.DT.winds.default_winds import (
     default_sep_from_winds,
     default_spin_from_winds,
@@ -237,7 +238,8 @@ def __init__(
             matching_tolerance_hard=1e-1,
             list_for_matching_HMS=None,
             list_for_matching_postMS=None,
-            list_for_matching_HeStar=None
+            list_for_matching_HeStar=None,
+            **kwargs
     ):
         """Initialize the step. See class documentation for details."""
         self.dt = dt
@@ -272,17 +274,18 @@ def __init__(
         self.KEYS = DEFAULT_TRANSLATED_KEYS
 
         # creating a track matching object
-        self.track_matcher = TrackMatcher(grid_name_Hrich = grid_name_Hrich,
-                                          grid_name_strippedHe = grid_name_strippedHe,
-                                          path=path, metallicity = metallicity,
-                                          matching_method = matching_method,
-                                          matching_tolerance=matching_tolerance,
-                                          matching_tolerance_hard=matching_tolerance_hard,
-                                          list_for_matching_HMS = list_for_matching_HMS,
-                                          list_for_matching_HeStar = list_for_matching_HeStar,
-                                          list_for_matching_postMS = list_for_matching_postMS,
-                                          record_matching = record_matching,
-                                          verbose = self.verbose)
+        #self.track_matcher = TrackMatcher(grid_name_Hrich = grid_name_Hrich,
+        #                                  grid_name_strippedHe = grid_name_strippedHe,
+        #                                  path=path, metallicity = metallicity,
+        #                                  matching_method = matching_method,
+        #                                  matching_tolerance=matching_tolerance,
+        #                                  matching_tolerance_hard=matching_tolerance_hard,
+        #                                  list_for_matching_HMS = list_for_matching_HMS,
+        #                                  list_for_matching_HeStar = list_for_matching_HeStar,
+        #                                  list_for_matching_postMS = list_for_matching_postMS,
+        #                                  record_matching = record_matching,
+        #                                  verbose = self.verbose)
+        self.track_matcher = kwargs.get('track_matcher', None)
 
         # create evolution handler object
         self.init_evo_kwargs()
@@ -308,6 +311,7 @@ def __repr__(self):
         """Return the name of evolution step."""
         return "Detached Step."
 
+    #@profile
     def __call__(self, binary):
         """
             Evolve the binary through detached evolution until RLO or

posydon/binary_evol/DT/track_match.py

@@ -252,7 +252,7 @@ class TrackMatcher:
     the previous step and the h5 track.
 
     """
-
+    #@profile
     def __init__(
             self,
             grid_name_Hrich,
@@ -424,6 +424,116 @@ def __init__(
 
         self.record_matching = record_matching
 
+
+        self.create_root0_h()
+        self.create_root0_he()
+        self.train_scalers()
+
+    #@profile
+    def train_scalers(self):
+
+       # ...if not, fit a new scaler, and store it for later use
+
+        lists_for_matching = [self.list_for_matching_HMS,
+                              self.list_for_matching_HeStar,
+                              self.list_for_matching_HMS_alternative,
+                              self.list_for_matching_HeStar_alternative,
+                              self.list_for_matching_postHeMS,
+                              self.list_for_matching_postHeMS_alternative,
+                              self.list_for_matching_postMS,
+                              self.list_for_matching_postMS_alternative]
+
+        #htracks = [True, False, False, False, False, False, True, False]
+
+        for list_for_matching in lists_for_matching:
+
+            match_attr_names = list_for_matching[0]
+            rescale_facs = list_for_matching[1]
+            scaler_methods = list_for_matching[2]
+            bnds = list_for_matching[3:]
+
+            if self.verbose:
+                print("Matching parameters and their normalizations:\n",
+                        match_attr_names, rescale_facs)
+            for htrack in [True, False]:
+                grid = self.grid_Hrich if htrack else self.grid_strippedHe
+                self.initial_mass = grid.grid_mass
+
+                # get (or train and get) scalers for attributes
+                # attributes are scaled to range (0, 1)
+                #scalers = []
+                for attr_name, method in zip(match_attr_names, scaler_methods):
+                    all_attributes = []
+                    # check that attributes are allowed as matching attributes
+                    if attr_name not in self.root_keys:
+                        raise AttributeError("Expected matching attribute "
+                                                f"{attr_name} not "
+                                                "added in root_keys list: "
+                                                f"{self.root_keys}")
+
+                    scaler_options = (attr_name, htrack, method)
+
+                    for mass in self.initial_mass:
+                        for i in grid.get(attr_name, mass):
+                            all_attributes.append(i)
+
+                    all_attributes = np.array(all_attributes)
+                    scaler = DataScaler()
+                    scaler.fit(all_attributes, method=method, lower=0.0, upper=1.0)
+                    self.stored_scalers[scaler_options] = scaler
+
+    #@profile
+    def create_root0_h(self):
+
+        # set which grid to search based on htrack condition
+        grid = self.grid_Hrich
+
+        # initial masses within grid (defined but never used? used in scale())
+        self.initial_mass = grid.grid_mass
+
+        # search across all initial masses and get max track length
+        max_track_length = 0
+        for mass in grid.grid_mass:
+            track_length = len(grid.get("age", mass))
+            max_track_length = max(max_track_length, track_length)
+
+        # intialize root matrix
+        # (DIM = [N(Mi), N(max_track_length), N(root_keys)])
+        self.rootm_h = np.inf * np.ones((len(grid.grid_mass),
+                                    max_track_length, len(self.root_keys)))
+
+        # for each mass, get matching metrics and store in matrix
+        for i, mass in enumerate(grid.grid_mass):
+            for j, key in enumerate(self.root_keys):
+                track = grid.get(key, mass)
+                self.rootm_h[i, : len(track), j] = track
+
+    #@profile
+    def create_root0_he(self):
+
+        # set which grid to search based on htrack condition
+        grid = self.grid_strippedHe
+
+        # initial masses within grid (defined but never used? used in scale())
+        self.initial_mass = grid.grid_mass
+
+        # search across all initial masses and get max track length
+        max_track_length = 0
+        for mass in grid.grid_mass:
+            track_length = len(grid.get("age", mass))
+            max_track_length = max(max_track_length, track_length)
+
+        # intialize root matrix
+        # (DIM = [N(Mi), N(max_track_length), N(root_keys)])
+        self.rootm_he = np.inf * np.ones((len(grid.grid_mass),
+                                    max_track_length, len(self.root_keys)))
+
+        # for each mass, get matching metrics and store in matrix
+        for i, mass in enumerate(grid.grid_mass):
+            for j, key in enumerate(self.root_keys):
+                track = grid.get(key, mass)
+                self.rootm_he[i, : len(track), j] = track
+
     def get_root0(self, attr_names, attr_vals, htrack, rescale_facs=None):
         """
             Get the stellar evolution track in the single star grid with values
@@ -462,29 +572,10 @@ def get_root0(self, attr_names, attr_vals, htrack, rescale_facs=None):
 
         """
 
+        rootm = self.rootm_h if htrack else self.rootm_he
         # set which grid to search based on htrack condition
         grid = self.grid_Hrich if htrack else self.grid_strippedHe
 
-        # initial masses within grid (defined but never used? used in scale())
-        self.initial_mass = grid.grid_mass
-
-        # search across all initial masses and get max track length
-        max_track_length = 0
-        for mass in grid.grid_mass:
-            track_length = len(grid.get("age", mass))
-            max_track_length = max(max_track_length, track_length)
-
-        # intialize root matrix
-        # (DIM = [N(Mi), N(max_track_length), N(root_keys)])
-        self.rootm = np.inf * np.ones((len(grid.grid_mass),
-                                    max_track_length, len(self.root_keys)))
-
-        # for each mass, get matching metrics and store in matrix
-        for i, mass in enumerate(grid.grid_mass):
-            for j, key in enumerate(self.root_keys):
-                track = grid.get(key, mass)
-                self.rootm[i, : len(track), j] = track
-
         # rescaling factors
         if rescale_facs is None:
             rescale_facs = np.ones_like(attr_names)
@@ -500,7 +591,7 @@ def get_root0(self, attr_names, attr_vals, htrack, rescale_facs=None):
         # Slice out just the matching metric data for all stellar tracks
         # grid_attr_vals now has shape
         # (N(Mi), N(max_track_len), N(matching_metrics))
-        grid_attr_vals = self.rootm[:, :, idx]
+        grid_attr_vals = rootm[:, :, idx]
 
         # For all stellar tracks in grid:
         # Take difference btwn. grid track and given star values...
@@ -521,7 +612,7 @@ def get_root0(self, attr_names, attr_vals, htrack, rescale_facs=None):
 
         # time and initial mass corresp. to track w/ minimum difference
         m0 = grid.grid_mass[mass_i]
-        t0 = self.rootm[mass_i][age_i][np.argmax("age" == self.root_keys)]
+        t0 = rootm[mass_i][age_i][np.argmax("age" == self.root_keys)]
 
         return m0, t0
 
@@ -607,22 +698,6 @@ def scale(self, attr_name, htrack, scaler_method):
 
         # find if the scaler has already been fitted and return it if so...
         scaler = self.stored_scalers.get(scaler_options, None)
-        if scaler is not None:
-            return scaler
-
-        # ...if not, fit a new scaler, and store it for later use
-        grid = self.grid_Hrich if htrack else self.grid_strippedHe
-        self.initial_mass = grid.grid_mass
-        all_attributes = []
-
-        for mass in self.initial_mass:
-            for i in grid.get(attr_name, mass):
-                all_attributes.append(i)
-
-        all_attributes = np.array(all_attributes)
-        scaler = DataScaler()
-        scaler.fit(all_attributes, method=scaler_method, lower=0.0, upper=1.0)
-        self.stored_scalers[scaler_options] = scaler
 
         return scaler
 

posydon/binary_evol/MESA/step_mesa.py

@@ -141,7 +141,8 @@ def __init__(
             stop_var_name=None,
             stop_value=None,
             stop_interpolate=True,
-            verbose=False):
+            verbose=False,
+            **kwargs):
         """Evolve a binary object given a MESA grid or interpolation object.
 
         Parameters
@@ -290,6 +291,7 @@ def close(self):
         # if hasattr(self, '_Interp'):
         #     self._Interp.close()
 
+    #@profile
     def get_final_MESA_step_time(self):
         """Infer the maximum MESA step time.
 
@@ -315,6 +317,7 @@ def get_final_MESA_step_time(self):
 
         return max_MESA_sim_time
 
+    #@profile
     def __call__(self, binary):
         """Evolve a binary using the MESA step."""
 

posydon/binary_evol/binarystar.py

@@ -231,6 +231,7 @@ def __init__(self, star_1=None, star_2=None, index=None, properties=None,
         else:
             self.properties = SimulationProperties()
 
+    #@profile
     def evolve(self):
         """Evolve a binary from start to finish."""
 
@@ -264,6 +265,7 @@ def evolve(self):
 
         self.properties.post_evolve(self)
 
+    #@profile
     def run_step(self):
         """Evolve the binary through one evolutionary step."""
         total_state = self.get_total_state()
@@ -280,7 +282,30 @@ def run_step(self):
                              "SimulationProperties.".format(next_step_name))
 
         self.properties.pre_step(self, next_step_name)
-        next_step(self)
+
+        # breaking up for profiling
+        if next_step_name == "step_HMS_HMS":
+            next_step(self)
+        elif next_step_name == "step_CO_HeMS":
+            next_step(self)
+        elif next_step_name == "step_CO_HMS_RLO":
+            next_step(self)
+        elif next_step_name == "step_CO_HeMS_RLO":
+            next_step(self)
+        elif next_step_name == "step_detached":
+            next_step(self)
+        elif next_step_name == "step_dco":
+            next_step(self)
+        elif next_step_name == "step_merged":
+            next_step(self)
+        elif next_step_name == "step_disrupted":
+            next_step(self)
+        elif next_step_name == "step_SN":
+            next_step(self)
+        elif next_step_name == "step_CE":
+            next_step(self)
+        else:
+            next_step(self)
 
         self.append_state()
         self.properties.post_step(self, next_step_name)