"""Simulation properties for the population class.
This class contains the simulation properties, e.g. flow, steps and `max_time`.
"""
__authors__ = [
"Kyle Akira Rocha <kylerocha2024@u.northwestern.edu>",
"Konstantinos Kovlakas <Konstantinos.Kovlakas@unige.ch>",
"Jeffrey Andrews <jeffrey.andrews@northwestern.edu>",
"Simone Bavera <Simone.Bavera@unige.ch>",
"Nam Tran <tranhn03@gmail.com>",
"Seth Gossage <seth.gossage@northwestern.edu>"
]
import os
import time
import numpy as np
from posydon.binary_evol.track_match import TrackMatcher
from posydon.config import PATH_TO_POSYDON_DATA
from posydon.interpolation.interpolation import GRIDInterpolator
from posydon.popsyn.io import simprop_kwargs_from_ini
from posydon.utils.common_functions import convert_metallicity_to_string
from posydon.utils.constants import age_of_universe
from posydon.utils.posydonerror import GridError
from posydon.utils.posydonwarning import Pwarn
[docs]
class NullStep:
"""An evolution step that does nothing but is used to initialize."""
pass
[docs]
class SimulationProperties:
"""Class describing the properties of a population synthesis simulation."""
# each value in this dict represents the expected path for the respective grid.
# A user may specify their own full path to a custom grid in the [grid_paths]
# section of their .ini file. I.e., HMS-HMS_path = 'path/to/my_own_grid/' to
# override these defaults.
default_grid_paths = {"single_HMS_path": os.path.join(PATH_TO_POSYDON_DATA, "single_HMS"),
"single_HeMS_path": os.path.join(PATH_TO_POSYDON_DATA, "single_HeMS"),
"HMS_HMS_path": os.path.join(PATH_TO_POSYDON_DATA, "HMS-HMS"),
"CO_HMS_RLO_path": os.path.join(PATH_TO_POSYDON_DATA, "CO-HMS_RLO"),
"CO_HeMS_path": os.path.join(PATH_TO_POSYDON_DATA, "CO-HeMS"),
"CO_HeMS_RLO_path": os.path.join(PATH_TO_POSYDON_DATA, "CO-HeMS_RLO")}
def __init__(self, flow=({}, {}),
step_HMS_HMS = (NullStep(), {}),
step_CO_HeMS = (NullStep(), {}),
step_CO_HMS_RLO = (NullStep(), {}),
step_CO_HeMS_RLO = (NullStep(), {}),
step_detached = (NullStep(), {}),
step_disrupted = (NullStep(), {}),
step_merged = (NullStep(), {}),
step_initially_single = (NullStep(), {}),
step_dco = (NullStep(), {}),
step_SN = (NullStep(), {}),
step_CE = (NullStep(), {}),
step_end = (NullStep(), {}),
extra_hooks = [],
**kwargs):
"""Construct the simulation properties object.
Parameters
----------
flow_chart : dict
A POSYDON flow_chart dictionary.
step_HMS_HMS : tuple
A tuple whose first element is a MesaGridStep class handling
HMS-HMS evolution, like MS_MS_step. The second element is a
dictionary of kwargs for that step.
step_CO_HeMS : tuple
A tuple whose first element is a MesaGridStep class handling
CO-HeMS evolution, like CO_HeMS_step. The second element is a
dictionary of kwargs for that step.
step_CO_HMS_RLO : tuple
A tuple whose first element is a MesaGridStep class handling
CO-HMS-RLO evolution, like CO_HMS_RLO_step. The second element is a
dictionary of kwargs for that step.
step_CO_HeMS_RLO : tuple
A tuple whose first element is a MesaGridStep class handling
CO-HeMS-RLO evolution, like CO_HeMS_RLO_step. The second element is a
dictionary of kwargs for that step.
step_detached : tuple
A tuple whose first element is a detached_step class handling detached
evolution. The second element is a dictionary of kwargs for that step.
step_disrupted : tuple
A tuple whose first element is a DisruptedStep class handling
disrupted evolution. The second element is a dictionary of kwargs for
that step.
step_merged : tuple
A tuple whose first element is a MergedStep class handling
merged evolution. The second element is a dictionary of kwargs for
that step.
step_initially_single : tuple
A tuple whose first element is a InitiallySingleStep class handling
initially single evolution. The second element is a dictionary of kwargs for
that step.
step_dco : tuple
A tuple whose first element is a DoubleCO class handling
double CO evolution. The second element is a dictionary of kwargs for
that step.
step_SN : tuple
A tuple whose first element is a StepSN class handling
supernova evolution. The second element is a dictionary of kwargs for
that step.
step_CE : tuple
A tuple whose first element is a StepCEE class handling
common envelope evolution. The second element is a dictionary of kwargs for
that step.
step_end : tuple
A tuple whose first element is a step_end class handling
the end of evolution. The second element is a dictionary of kwargs for
that step.
extra_hooks : list of tuples
Each tuple contains a hooks class and kwargs or the extra step name
(e.g., 'extra_pre_evolve', 'extra_pre_step', 'extra_post_step',
'extra_post_evolve') and the corresponding function.
"""
self.kwargs = kwargs
verbose = self.kwargs.get('verbose', False)
# gather kwargs
self.kwargs["flow"] = flow
step_kwargs = {"step_HMS_HMS": step_HMS_HMS,
"step_CO_HeMS": step_CO_HeMS,
"step_CO_HMS_RLO": step_CO_HMS_RLO,
"step_CO_HeMS_RLO": step_CO_HeMS_RLO,
"step_detached": step_detached,
"step_disrupted": step_disrupted,
"step_merged": step_merged,
"step_initially_single": step_initially_single,
"step_dco": step_dco,
"step_SN": step_SN,
"step_CE": step_CE,
"step_end": step_end}
for key, step_tuple in step_kwargs.items():
step_class, _ = step_tuple
if verbose and isinstance(step_class, NullStep):
Pwarn(f"Step {key} not provided, skipping it.",
"StepWarning")
self.kwargs[key] = step_tuple
self.kwargs["extra_hooks"] = extra_hooks
self.default_hooks = EvolveHooks()
self.all_hooks_classes = [self.default_hooks]
for item in self.kwargs.get('extra_hooks', []):
if isinstance(item, tuple):
if isinstance(item[0], type) and isinstance(item[1], dict):
cls, params = item
self.all_hooks_classes.append(cls(**params))
elif isinstance(item[0], str):
# setting extra_pre/post_step/evolve methods
setattr(self, item[0], item[1])
else:
raise ValueError(
"`extra_hooks` must be list of tuples with either "
"(i) a class deriving from EvolveHooks and a kwargs dict, "
"or (ii) the name of the extra function and the callable.")
# Limits on simulation
if not hasattr(self, 'max_simulation_time'):
self.max_simulation_time = age_of_universe
if not hasattr(self, 'end_events'):
self.end_events = []
if not hasattr(self, 'end_states'):
self.end_states = []
# for debugging purposes
if not hasattr(self, 'max_n_steps_per_binary'):
self.max_n_steps_per_binary = 100
# Set functions for evolution
self.all_step_names = [] ## list of strings of all evolutionary steps
for key, val in self.kwargs.items():
if "step" not in key: # skip loading steps
setattr(self, key, val)
elif "step" in key:
self.all_step_names.append(key)
self.steps_loaded = False
self.preload_imports()
# To hold TrackMatcher objects per step, if needed.
# maybe get rid of this
self.track_matchers = {}
for grid_name in self.default_grid_paths:
try:
self.set_path(grid_name, self.kwargs[grid_name])
except KeyError as e:
Pwarn(f"{grid_name} is not set in the kwargs passed to SimulationProperties. "
f"Falling back to the default: {self.default_grid_paths[grid_name]}",
"ReplaceValueWarning")
self.set_path(grid_name, self.default_grid_paths[grid_name])
# These hold GRIDInterpolator objects
# and associated grid names for ea. metallicity
# (intended keys are metallicities):
self.grids_Hrich = {}
self.grids_strippedHe = {}
[docs]
def set_path(self, path_name, path_str):
"""
Set and normalize a grid path attribute that points to one of the
MESA grids needed for binary evolution. By default, these are the
grids inside of the directory name held in $PATH_TO_POSYDON_DATA.
For example, for the step_HMS_HMS, the grid would be
$PATH_TO_POSYDON_DATA/HMS-HMS/<metallicity>_Zsun.h5
by default. The grid HDF5 file names themselves are expected to
follow formats like so: 1e+00_Zsun.h5, 1e-04_Zsun.h5, etc.
If ``path_str`` is ``None``, a default path is assigned based on
``path_name`` using ``self.default_grid_paths``. If ``path_name`` is not
recognized, a ``GridError`` is raised listing the valid options.
The resulting path is converted to an absolute path before being stored
as an attribute of the instance.
Parameters
----------
path_name : str
Name of the grid path attribute to set. Must be a key in
``self.default_grid_paths`` if ``path_str`` is ``None``.
path_str : str or None
Path to assign. If ``None``, a default path corresponding to
``path_name`` is used.
Raises
------
GridError
If ``path_name`` is not recognized and no default path can be assigned.
Notes
-----
The path is not validated for existence here; only normalization to an
absolute path is performed.
"""
# construct path to *_Zsun.h5 files if not specified
if path_str is None:
if path_name in self.default_grid_paths:
path_str = self.default_grid_paths[path_name]
else:
valid_names = "\n".join(f"{k} = <your-path-or-None>" for k in self.default_grid_paths)
raise GridError(f'Trying to assign a grid path for "{path_name}".\n'
"This is an unrecognized path name. Please check "
"the [grid_paths] section of your .ini file.\n\n"
"Valid path variable names are:\n"
f"{valid_names}\n")
path_str = os.path.abspath(path_str)
setattr(self, path_name, path_str)
[docs]
def preload_imports(self):
"""
Preload the imports of detached_step and MesaGridStep to avoid
importing them when they are needed when `close()` is called. In
particular, detached_step imports sklearn, which in turn utilizes
loky, which invokes its own register.at_exit call. If this happens
during the `close()` call, which is invoked at shutdown, a
failure occurs, hence the need for something like this.
"""
from posydon.binary_evol.CE.step_CEE import StepCEE
from posydon.binary_evol.DT.step_detached import detached_step
from posydon.binary_evol.MESA.step_mesa import MesaGridStep
self._detached_step = detached_step
self._step_CE = StepCEE
self._MesaGridStep = MesaGridStep
[docs]
@classmethod
def from_ini(cls, path, metallicity = None, load_steps=False, RNG=np.random.default_rng(),
verbose=False, **override_sim_kwargs):
"""Create a SimulationProperties instance from an inifile.
Parameters
----------
path : str
Path to an inifile to load in.
metallicity : float
A metallicity (Z) may be provided to automatically assign
to steps as they are loaded. Should be one of e.g., 2.0, 1.0,
4.5e-1, 2e-1, 1e-1, 1e-2, 1e-3, 1e-4, corresponding to
metallicities available in your POSYDON_DATA grids.
load_steps : bool
Whether or not evolution steps should be automatically loaded.
RNG : numpy.random.Generator, optional
Random number generator used for any stochastic components of
the simulation. Defaults to a new NumPy Generator instance
created via ``np.random.default_rng()``.
verbose : bool
Print useful info.
**override_sim_kwargs
Additional keyword arguments that override values specified
in the .ini file when constructing the SimulationProperties
instance.
Returns
-------
SimulationProperties
A new instance of SimulationProperties.
"""
sim_kwargs = simprop_kwargs_from_ini(path)
sim_kwargs = {**sim_kwargs, **override_sim_kwargs}
new_instance = cls(**sim_kwargs)
if load_steps:
# Load the steps and required data
new_instance.load_steps(metallicity=metallicity,
RNG=RNG,
verbose=verbose)
return new_instance
[docs]
def load_steps(self, metallicity=None, RNG=np.random.default_rng(), verbose=False):
"""Instantiate all step classes and set as instance attributes.
Parameters
----------
metallicity : float
A metallicity (Z) may be provided to automatically assign
to steps as they are loaded. Should be one of e.g., 2.0, 1.0,
4.5e-1, 2e-1, 1e-1, 1e-2, 1e-3, 1e-4, corresponding to
metallicities available in your POSYDON_DATA grids.
verbose : bool
Print extra information.
Returns
-------
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)
self.load_a_step(name, tup, metallicity=metallicity, RNG=RNG, verbose=verbose)
if verbose:
if self.steps_loaded:
print("All steps loaded successfully.")
else:
print("Not all steps were loaded successfully. Check warnings for details.")
[docs]
def load_a_step(self, step_name, step_tup=(NullStep, {}), metallicity=None,
RNG=np.random.default_rng(), from_ini='', verbose=False):
"""
Instantiate and attach a simulation step to this object.
This method creates an instance of a step class and assigns it as an
attribute of SimulationProperties using ``step_name`` as the attribute
name. Step keyword arguments may be provided directly via ``step_tup``
or loaded from an `.ini` configuration file. Before instantiation,
step arguments are validated and augmented (e.g., assigning metallicity
and creating a TrackMatcher if required).
Parameters
----------
step_name : str
Name of the evolution step. The created step instance will be
attached to the object as ``self.<step_name>``. See
``SimulationProperties.__init__`` for the standard set of steps.
step_tup : tuple, optional
Tuple of the form ``(step_class, kwargs_dict)`` where:
- ``step_class`` is the class representing the step.
- ``kwargs_dict`` is a dictionary of keyword arguments used to
initialize the step.
Default is ``(NullStep, {})``.
metallicity : float, optional
Metallicity (Z) to assign to the step if required and not already
specified in the step keyword arguments. Default supported values
are: 2.0, 1.0, 4.5e-1, 2e-1, 1e-1, 1e-2, 1e-3, 1e-4.
from_ini : str, optional
Path to an `.ini` file containing step configuration. If provided
and the file exists, the step class and keyword arguments for
``step_name`` are loaded from this file and override ``step_tup``.
verbose : bool, optional
If True, print detailed information about step loading and the
keyword arguments used to instantiate the step.
Returns
-------
None
Notes
-----
- Step keyword arguments are processed by ``self.check_step`` before
instantiation. This may assign a metallicity and/or attach a
``TrackMatcher`` if required for the step.
- The instantiated step is stored as an attribute of SimulationProperties.
- After loading, ``self.steps_loaded`` is updated to indicate whether
all configured steps have been successfully attached.
"""
if verbose:
print(f"Loading {step_name}...")
# grab kwargs from ini file for given step
if os.path.isfile(from_ini):
step_tup = simprop_kwargs_from_ini(from_ini, only=step_name)[step_name]
if step_name != "flow":
# check to make sure the step has a...
# 1) metallicity assigned (if needed)
# 2) TrackMatcher assigned (if needed)
step_tup = self.check_step(metallicity, RNG, step_name,
step_tup, verbose)
step_func, step_kwargs = step_tup
# Try to load the step
try:
setattr(self, step_name, step_func(**step_kwargs))
if verbose:
print(f"Class: {step_func}")
if step_kwargs:
print("step_kwargs: ")
kw_list = [f"\t{key}: {val}" for key, val in step_kwargs.items()]
print("\n".join(kw_list))
print(f"{step_name} loaded successfully.\n")
except TypeError as e:
Pwarn(f"Error loading {step_name}: {e}", "StepWarning")
print(f"Loading {step_name} without arguments.")
setattr(self, step_name, step_func())
# check if all steps have been loaded
self.steps_loaded = all(hasattr(self, name)
for name, tup in self.kwargs.items()
if isinstance(tup, tuple))
[docs]
def check_step(self, metallicity, RNG, step_name, step_tup, verbose=False):
"""
Validate and update configuration for an evolution step.
This method ensures that a valid metallicity is assigned to the step
(unless the step is excluded from metallicity handling) and that a
corresponding TrackMatcher exists if the step requires track matching.
If a TrackMatcher for the `(metallicity, step_name)` combination does
not yet exist, it is created and stored.
Parameters
----------
metallicity : float or None
Default metallicity value to use for the step if not explicitly
provided in ``step_kwargs``.
step_name : str
Name of the pipeline step being checked.
step_kwargs : dict
Keyword arguments for the step. This dictionary may be modified
in-place to include validated metallicity and/or a TrackMatcher
instance.
verbose : bool, optional
If True, print the keyword arguments used to construct the
TrackMatcher.
Returns
-------
tuple
The step tuple containing the step function and the updated
``step_kwargs`` dictionary. This new step_kwargs contains a
validated ``metallicity`` entry and potentially a
``track_matcher`` object.
Notes
-----
- If metallicity is not provided for a step that requires it, a warning
is issued and a default value of ``Z = 1.0`` (solar metallicity) is used.
- TrackMatcher objects are stored in ``self.track_matchers`` and reused
for repeated `(metallicity, step_name)` combinations.
"""
step_func, step_kwargs = step_tup
# copy these to track what was originally passed before adding defaults
# in case we need to update a TrackMatcher for a step, this method only
# cares about updating what it is told to, not everything.
original_step_kwargs = step_kwargs.copy()
# check/assign metallicity for the step
if "metallicity" in step_func.DEFAULT_KWARGS:
metallicity = step_kwargs.get('metallicity', metallicity)
if metallicity is None:
Pwarn(f"{step_name} not assigned a metallicity. "
"Defaulting to Z = Zsun (solar).",
"ReplaceValueWarning")
metallicity = 1.0
step_kwargs['metallicity'] = float(metallicity)
# These steps need these grids:
step_grid_map = {"step_HMS_HMS": self.HMS_HMS_path,
"step_CO_HMS_RLO": self.CO_HMS_RLO_path,
"step_CO_HeMS": self.CO_HeMS_path,
"step_CO_HeMS_RLO": self.CO_HeMS_RLO_path}
if step_name in step_grid_map:
step_kwargs['grid_path'] = step_grid_map[step_name]
if "track_matcher" in step_func.DEFAULT_KWARGS:
# each metallicity/step combo could require
# a unique TrackMatcher, so check for that
step_kwargs = self._check_track_matcher(metallicity,
step_name,
step_kwargs,
original_step_kwargs,
verbose)
if "RNG" in step_func.DEFAULT_KWARGS:
step_kwargs['RNG'] = RNG
return (step_func, step_kwargs)
def _check_track_matcher(self, metallicity, step_name, step_kwargs,
original_step_kwargs, verbose=False):
"""
Validate, create, update, and assign the TrackMatcher for an evolution step.
A unique TrackMatcher is maintained for each ``(metallicity, step_name)``
combination. If a TrackMatcher for the requested combination does not
already exist in this SimulationProperties, one is created using the
supplied matcher-specific keyword arguments. If one already exists,
its explicitly supplied properties are compared against the existing
values and updated when necessary. If an updated property is one of the
TrackMatcher's training triggers, the TrackMatcher is recreated so that
it is retrained with the new settings.
TrackMatcher-specific keyword arguments are separated from the remaining
step keyword arguments using ``TrackMatcher.separate_kwargs``. The
resulting TrackMatcher instance is then added to ``step_kwargs`` under
the ``"track_matcher"`` key.
TrackMatchers for other metallicities associated with the same step are
removed after the requested matcher has been created or retrieved. This
limits the number of simultaneously stored TrackMatcher objects and
reduces memory usage.
Note that in multi-metallicity runs, each metallicity is assigned to a
unique BinaryPopulation, each of which has its own unique
SimulationProperties at that given metallicity..
Parameters
----------
metallicity : float
Stellar metallicity associated with the evolutionary step. This value
is used together with ``step_name`` to identify the required
TrackMatcher.
step_name : str
Name of the evolutionary step for which the TrackMatcher is required.
step_kwargs : dict
Keyword arguments for the step. TrackMatcher-specific arguments are
separated from this dictionary and the resulting TrackMatcher instance
is added under the ``"track_matcher"`` key.
original_step_kwargs : dict
Copy of the keyword arguments originally supplied to the evolution step
before defaults or other values were added. Only properties explicitly
present in this dictionary are considered when determining whether an
existing TrackMatcher needs to be updated.
verbose : bool, optional
If True, print the TrackMatcher keyword arguments used for the current
step.
Returns
-------
dict
The updated step keyword arguments containing the TrackMatcher-specific
arguments removed and the appropriate TrackMatcher instance assigned
to the ``"track_matcher"`` key.
Notes
-----
- TrackMatcher objects are stored in ``self.track_matchers`` using
``(metallicity, step_name)`` as the key.
- GRIDInterpolator objects required by the TrackMatcher are created and
cached by ``create_track_matcher``.
- Existing TrackMatchers are reused when their explicitly supplied
configuration has not changed.
- If a training-triggering TrackMatcher property changes, the existing
matcher is recreated using the updated configuration.
- TrackMatchers for other metallicities but the same ``step_name`` are
deleted to conserve memory.
"""
matcher_key = (metallicity, step_name)
matcher_needed = matcher_key not in self.track_matchers
step_kwargs, matcher_kwargs = TrackMatcher.separate_kwargs(step_kwargs)
if matcher_needed:
# Always just create a new TrackMatcher if it does not exist
self.create_track_matcher(metallicity, step_name, matcher_kwargs)
else:
track_matcher = self.track_matchers[matcher_key]
retrain = False
# subsequently check if any properties need to be updated,
# in case reloading for example
for k, v in matcher_kwargs.items():
# only care to update kwargs actually passed via load_step
if k in original_step_kwargs:
do_update = track_matcher.kwargs[k] != original_step_kwargs[k]
if do_update:
setattr(track_matcher, k, matcher_kwargs[k])
track_matcher.kwargs[k] = matcher_kwargs[k]
if k in track_matcher.TRAINING_TRIGGERS:
retrain = True
if retrain:
updated_kwargs = track_matcher.kwargs
self.create_track_matcher(metallicity, step_name, updated_kwargs)
# check for and delete any old TrackMatcher's this step has to save RAM
for mk in list(self.track_matchers):
this_met, this_stepn = mk
if this_met != metallicity and this_stepn == step_name:
del self.track_matchers[mk]
# delete associated cached grids to save RAM too,
# but only if no other TrackMatcher is using this
# metallicity. Otherwise, grid still needed.
other_steps_at_this_met = any(other_met == this_met
for other_met, _ in self.track_matchers)
if not other_steps_at_this_met:
del self.grids_Hrich[this_met]
del self.grids_strippedHe[this_met]
if verbose:
kw_list = [f"\t{key}: {val}" for key, val in matcher_kwargs.items()]
print(f"matcher_kwargs: \n" + "\n".join(kw_list))
# reference ths TrackMatcher in this step's kwargs
step_kwargs['track_matcher'] = self.track_matchers[matcher_key]
return step_kwargs
[docs]
def create_track_matcher(self, metallicity, step_name, matcher_kwargs):
"""
Create and store a TrackMatcher for a given metallicity and step.
This method ensures that the required stellar evolution grids
(H-rich and stripped-He) are loaded for the specified metallicity.
If the corresponding GRIDInterpolator objects do not yet exist,
they are created and cached. The interpolators are then passed to
a TrackMatcher instance, which is stored internally.
Parameters
----------
metallicity : float
Stellar metallicity used to select the appropriate grid files.
step_name : str
Identifier for the evolutionary step associated with this
TrackMatcher.
matcher_kwargs : dict
Keyword arguments used to initialize the TrackMatcher. This
dictionary will be updated in-place with the following keys:
'grid_Hrich' and 'grid_strippedHe'.
Notes
-----
- GRIDInterpolator objects are created only once per metallicity
and reused for subsequent TrackMatcher creations.
- The created TrackMatcher is stored in ``self.track_matchers``
using the key ``(metallicity, step_name)``.
"""
z_str = convert_metallicity_to_string(metallicity)
# set up GRIDInterpolator objects (for HMS and HeMS)
# (only if one hasn't been created already for a given metallicity)
if metallicity not in self.grids_Hrich:
grid_path_Hrich = os.path.join(self.single_HMS_path, f"{z_str}_Zsun.h5")
self.grids_Hrich[metallicity] = GRIDInterpolator(grid_path_Hrich)
if metallicity not in self.grids_strippedHe:
grid_path_strippedHe = os.path.join(self.single_HeMS_path, f"{z_str}_Zsun.h5")
self.grids_strippedHe[metallicity] = GRIDInterpolator(grid_path_strippedHe)
# Create TrackMatcher object as needed, passing GRIDInterpolator references
matcher_kwargs['grid_Hrich'] = self.grids_Hrich[metallicity]
matcher_kwargs['grid_strippedHe'] = self.grids_strippedHe[metallicity]
self.track_matchers[(metallicity, step_name)] = TrackMatcher(**matcher_kwargs)
[docs]
def close(self):
"""Close hdf5 files before exiting."""
all_step_funcs = [getattr(self, key) for key, val in
self.__dict__.items() if 'step_' in key]
for step_func in all_step_funcs:
if isinstance(step_func, self._MesaGridStep):
step_func.close()
for metallicity in self.grids_Hrich:
self.grids_Hrich[metallicity].close()
for metallicity in self.grids_strippedHe:
self.grids_strippedHe[metallicity].close()
[docs]
def pre_evolve(self, binary):
"""Functions called before a binary evolves.
Uses all extra hooks classes or extra functions.
Parameters
----------
binary : instance of <class, BinaryStar>
The binary before evolution starts.
Returns
-------
binary : instance of <class, BinaryStar>
"""
for hooks in self.all_hooks_classes:
hooks.pre_evolve(binary)
if hasattr(self, 'extra_pre_evolve'):
self.extra_pre_evolve(binary)
return binary
[docs]
def pre_step(self, binary, step_name):
"""Prepare for step.
The method is called before every evolution step; uses all extra hooks
classes or extra functions (except for undefined next step errors).
Parameters
----------
binary : instance of <class, BinaryStar>
The binary before evolution starts.
step_name : str
The name of the step about to be called (as defined in the flow).
Returns
-------
binary : instance of <class, BinaryStar>
"""
for hooks in self.all_hooks_classes:
hooks.pre_step(binary, step_name)
if hasattr(self, 'extra_pre_step'):
self.extra_pre_step(binary, step_name)
return binary
[docs]
def post_step(self, binary, step_name):
"""Finalize step.
The method is called after every evolution step; uses all extra hooks
classes or extra functions (except for undefined next step errors).
Parameters
----------
binary : instance of <class, BinaryStar>
The binary before evolution starts.
step_name : str
The name of the step about to be called (as defined in the flow).
Returns
-------
binary : instance of <class, BinaryStar>
"""
## do not call extra step hooks if history_verbose=False
if not binary.history_verbose and binary.event is not None:
if "redirect" in binary.event:
return binary
for hooks in self.all_hooks_classes:
hooks.post_step(binary, step_name)
if hasattr(self, 'extra_post_step'):
self.extra_post_step(binary, step_name)
return binary
[docs]
def post_evolve(self, binary):
"""Finalize the evolution of the binary.
The method is called after a binary exits the evolution loop.
Uses all extra hooks classes or extra functions.
Parameters
----------
binary : instance of <class, BinaryStar>
The binary after evolution is ended.
Returns
-------
binary : instance of <class, BinaryStar>
"""
for hooks in self.all_hooks_classes:
hooks.post_evolve(binary)
if hasattr(self, 'extra_post_evolve'):
self.extra_post_evolve(binary)
return binary
[docs]
class EvolveHooks:
"""Base class for hooking into binary evolution."""
def __init__(self):
"""
Add any new output columns to the hooks constructor.
Example for extra binary columns:
self.extra_binary_col_names = ["column_name_1", "column_name_2"]
Example for extra star columns:
self.extra_star_col_names = ["column_name_1", "column_name_2"]
"""
pass
[docs]
def pre_evolve(self, binary):
"""Perform actions before a binary evolves."""
return binary
[docs]
def pre_step(self, binary, step_name):
"""Perform actions before every evolution step."""
return binary
[docs]
def post_step(self, binary, step_name):
"""Perform acctions after every evolution step."""
return binary
[docs]
def post_evolve(self, binary):
"""Perform actions after a binary exits the evolution loop."""
return binary
[docs]
class TimingHooks(EvolveHooks):
"""Add history column 'step_times' (time taken by step) to each binary.
Example
-------
>>> pop.to_df(extra_columns={'step_times': float})
"""
def __init__(self):
self.extra_binary_col_names = ["step_times"]
[docs]
def pre_evolve(self, binary):
"""Initialize the step time to match history."""
if not hasattr(binary, 'step_times'):
binary.step_times = [0.0]
return binary
[docs]
def pre_step(self, binary, step_name):
"""Record the wall time before taking the step."""
self.step_start_time = time.time()
return binary
[docs]
def post_step(self, binary, step_name):
"""Record the duration of the step."""
binary.step_times.append(time.time() - self.step_start_time)
if len(binary.event_history) > len(binary.step_times):
diff = len(binary.event_history) - len(binary.step_times)
binary.step_times += [None] * (diff)
elif len(binary.event_history) < len(binary.step_times):
last_items = len(binary.event_history)
binary.step_times = binary.step_times[-(last_items - 1):]
return binary
[docs]
def post_evolve(self, binary):
"""Add None's to step_times to match history rows."""
if binary.event == 'END' or binary.event == 'FAILED':
diff = int(len(binary.event_history) - len(binary.step_times))
binary.step_times += [None] * diff
return binary
[docs]
class StepNamesHooks(EvolveHooks):
"""Add history column 'step_name' to each binary.
Name of evolutionary step as defined in SimulationProperties.
>>> pop.to_df(extra_columns={'step_names': str})
"""
def __init__(self):
self.extra_binary_col_names = ["step_names"]
[docs]
def pre_evolve(self, binary):
"""Initialize the step name to match history."""
if not hasattr(binary, 'step_names'):
binary.step_names = ['initial_cond']
return binary
[docs]
def pre_step(self, binary, step_name):
"""Do not do anything before the step."""
return binary
[docs]
def post_step(self, binary, step_name):
"""Record the step name."""
binary.step_names.append(step_name)
len_binary_hist = len(binary.event_history)
len_step_names = len(binary.step_names)
diff = len_binary_hist - len_step_names
if len_binary_hist > len_step_names:
binary.step_names += [None] * (diff)
elif len_binary_hist < len_step_names:
binary.step_names = binary.step_names[-(len_binary_hist - 1):]
return binary
[docs]
def post_evolve(self, binary):
"""Ensure None's are append to step_names to match rows in history."""
if binary.event == 'END' or binary.event == 'FAILED':
diff = int(len(binary.event_history) - len(binary.step_names))
binary.step_names += [None]*diff
return binary
[docs]
class PrintStepInfoHooks(EvolveHooks):
"""Simple example for adding extra print info."""
[docs]
def pre_step(self, binary, step_name):
"""Print the step name for each binary, before taking it."""
print(binary.index, step_name)
return binary
[docs]
def post_evolve(self, binary):
"""Report at the end of the evolution of each binary."""
print("End evol for binary {}".format(binary.index), end='\n'*2)
return binary