diff --git a/posydon/binary_evol/CE/step_CEE.py b/posydon/binary_evol/CE/step_CEE.py index a4c899d1a5..8a821abe51 100644 --- a/posydon/binary_evol/CE/step_CEE.py +++ b/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 diff --git a/posydon/binary_evol/DT/step_detached.py b/posydon/binary_evol/DT/step_detached.py index de47199265..f6f79f5f32 100644 --- a/posydon/binary_evol/DT/step_detached.py +++ b/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 diff --git a/posydon/binary_evol/DT/track_match.py b/posydon/binary_evol/DT/track_match.py index 07399cb708..23702a3670 100644 --- a/posydon/binary_evol/DT/track_match.py +++ b/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 diff --git a/posydon/binary_evol/MESA/step_mesa.py b/posydon/binary_evol/MESA/step_mesa.py index b6fb65343c..735f3e12fa 100644 --- a/posydon/binary_evol/MESA/step_mesa.py +++ b/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.""" diff --git a/posydon/binary_evol/binarystar.py b/posydon/binary_evol/binarystar.py index efad6b2067..0b442671fb 100644 --- a/posydon/binary_evol/binarystar.py +++ b/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) diff --git a/posydon/binary_evol/simulationproperties.py b/posydon/binary_evol/simulationproperties.py index adec0e8324..9df73a2eca 100644 --- a/posydon/binary_evol/simulationproperties.py +++ b/posydon/binary_evol/simulationproperties.py @@ -17,6 +17,8 @@ import os import time +from posydon.binary_evol.DT.track_match import TrackMatcher +from posydon.config import PATH_TO_POSYDON_DATA from posydon.popsyn.io import simprop_kwargs_from_ini from posydon.utils.constants import age_of_universe from posydon.utils.posydonwarning import Pwarn @@ -245,6 +247,7 @@ def from_ini(cls, path, metallicity = None, load_steps=False, verbose=False, **o return new_instance + #@profile def load_steps(self, metallicity=None, verbose=False): """Instantiate all step classes and set as instance attributes. @@ -266,11 +269,32 @@ def load_steps(self, metallicity=None, verbose=False): if verbose: print('STEP NAME'.ljust(20) + 'STEP FUNCTION'.ljust(25) + 'KWARGS') + + # creating a track matching object + self.track_matcher = None + # for every other step, give it a metallicity and load each step for name, tup in self.kwargs.items(): if isinstance(tup, tuple): step_kwargs = tup[1] metallicity = step_kwargs.get('metallicity', metallicity) + + if self.track_matcher is None and metallicity is not None: + self.track_matcher = TrackMatcher(grid_name_Hrich = None, + grid_name_strippedHe = None, + path=PATH_TO_POSYDON_DATA, metallicity = metallicity, + matching_method = "minimize", + matching_tolerance=1e-2, + matching_tolerance_hard=1e-1, + list_for_matching_HMS = None, + list_for_matching_HeStar = None, + list_for_matching_postMS = None, + record_matching = False, + verbose = False) + + if name not in ["flow", "step_SN", "step_end"]: + step_kwargs['track_matcher'] = self.track_matcher + self.load_a_step(name, tup, metallicity=metallicity, verbose=verbose) # track that all steps have been loaded diff --git a/posydon/grids/psygrid.py b/posydon/grids/psygrid.py index 8e57a26d1a..c9cb9a510e 100644 --- a/posydon/grids/psygrid.py +++ b/posydon/grids/psygrid.py @@ -1472,18 +1472,8 @@ def load(self, filepath=None): hdf5 = self.hdf5 # load initial/final_values self._say("\tLoading initial/final values...") - self.initial_values = hdf5['/grid/initial_values'][()] - self.final_values = hdf5['/grid/final_values'][()] - - # change ASCII to UNICODE in termination flags in `final_values` - new_dtype = [] - for dtype in self.final_values.dtype.descr: - if (dtype[0].startswith("termination_flag") or - (dtype[0] == "mt_history") or ("_type" in dtype[0]) or - ("_state" in dtype[0]) or ("_class" in dtype[0])): - dtype = (dtype[0], H5_REC_STR_DTYPE.replace("S", "U")) - new_dtype.append(dtype) - self.final_values = self.final_values.astype(new_dtype) + self.initial_values = hdf5['/grid/initial_values'] + self.final_values = hdf5['/grid/final_values'] # load MESA dirs self._say("\tAcquiring paths to MESA directories...") diff --git a/posydon/interpolation/interpolation.py b/posydon/interpolation/interpolation.py index e975033db5..53b16c91bb 100644 --- a/posydon/interpolation/interpolation.py +++ b/posydon/interpolation/interpolation.py @@ -331,7 +331,9 @@ def __init__(self, path, verbose=False): grid_mass = [] for i in range(len(grid)): - grid_mass.append(grid[i].history1['star_mass'][0]) + minit = grid[i].history1['star_mass'][0] + grid_mass.append(minit) + self.grid_mass = np.array(grid_mass) grid_age = [] @@ -502,6 +504,11 @@ def __init__(self, path, verbose=False): 'avg_charge_He' ) + # pre-load the grid data for all masses + for i in range(len(grid)): + minit = grid[i].history1['star_mass'][0] + self.load_grid(minit) + def load_grid(self, *args): """Load the requested data to `grid_data`. diff --git a/posydon/popsyn/binarypopulation.py b/posydon/popsyn/binarypopulation.py index 56aaa79f8b..23b8a01e02 100644 --- a/posydon/popsyn/binarypopulation.py +++ b/posydon/popsyn/binarypopulation.py @@ -271,6 +271,7 @@ def evolve(self, **kwargs): self.kwargs.update(params) self._safe_evolve(**self.kwargs) + #@profile def _safe_evolve(self, **kwargs): """Evolve binaries in a population, catching warnings/exceptions.""" if not self.population_properties.steps_loaded: