Add options to change galactic orbit integrator

cogsworth/_version.py

@@ -1 +1 @@
-__version__ = "2.1.0"
+__version__ = "2.1.2"

cogsworth/kicks.py

@@ -44,8 +44,9 @@ def get_kick_differential(delta_v_sys_xyz, phase=None, inclination=None):
     return kick_differential
 
 
-def integrate_orbit_with_events(w0, t1, t2, dt, potential=gp.MilkyWayPotential(), events=None,
-                                store_all=True, quiet=False):
+def integrate_orbit_with_events(w0, t1, t2, dt, potential=gp.MilkyWayPotential(), events=None, 
+                                store_all=True, quiet=False, 
+                                integrator=gi.DOPRI853Integrator, integrator_kwargs={}):
     """Integrate :class:`~gala.dynamics.PhaseSpacePosition` in a 
     :class:`Potential <gala.potential.potential.PotentialBase>` with events that occur at certain times
 
@@ -73,6 +74,8 @@ def integrate_orbit_with_events(w0, t1, t2, dt, potential=gp.MilkyWayPotential()
         PhaseSpacePosition will be stored - this cuts down on memory usage.
     quiet : `bool`, optional
         Whether to silence warning messages about failing orbits
+    integrator : :class:`~gala.integrate.Integrator`, optional
+        The integrator used by gala for evolving the orbits of binaries in the galactic potential
 
     Returns
     -------
@@ -82,7 +85,7 @@ def integrate_orbit_with_events(w0, t1, t2, dt, potential=gp.MilkyWayPotential()
     """
     # if there are no events then just integrate the whole thing
     if events is None:
-        full_orbit = potential.integrate_orbit(w0, t1=t1, t2=t2, dt=dt, Integrator=gi.DOPRI853Integrator)
+        full_orbit = potential.integrate_orbit(w0, t1=t1, t2=t2, dt=dt, Integrator=integrator, Integrator_kwargs=integrator_kwargs)
         # jettison everything but the final timestep if user says so
         if not store_all:
             full_orbit = full_orbit[-1:]
@@ -115,7 +118,8 @@ def integrate_orbit_with_events(w0, t1, t2, dt, potential=gp.MilkyWayPotential()
 
                     # integrate the orbit over these timesteps
                     orbit = potential.integrate_orbit(current_w0, t=matching_timesteps,
-                                                      Integrator=gi.DOPRI853Integrator)
+                                                      Integrator=integrator,
+                                                      Integrator_kwargs=integrator_kwargs)
 
                     # save the orbit data (minus the last timestep to avoid duplicates)
                     orbit_data.append(orbit.data[:-1])
@@ -143,7 +147,8 @@ def integrate_orbit_with_events(w0, t1, t2, dt, potential=gp.MilkyWayPotential()
                 # evolve the rest of the orbit out
                 matching_timesteps = timesteps[timesteps >= time_cursor]
                 orbit = potential.integrate_orbit(current_w0, t=matching_timesteps,
-                                                  Integrator=gi.DOPRI853Integrator)
+                                                  Integrator=integrator,
+                                                  Integrator_kwargs=integrator_kwargs)
                 orbit_data.append(orbit.data)
 
             data = coords.concatenate_representations(orbit_data) if len(orbit_data) > 1 else orbit_data[0]

cogsworth/pop.py

@@ -19,6 +19,7 @@
 from cosmic.utils import parse_inifile
 import gala.potential as gp
 import gala.dynamics as gd
+import gala.integrate as gi
 from gala.potential.potential.io import to_dict as potential_to_dict, from_dict as potential_from_dict
 
 from cogsworth import sfh
@@ -79,12 +80,18 @@ class Population():
         Whether to store the entire orbit for each binary, by default True. If not then only the final
         PhaseSpacePosition will be stored. This cuts down on both memory usage and disk space used if you
         save the Population (as well as how long it takes to reload the data).
+    integrator : :class:`~gala.integrate.Integrator`, optional
+        The integrator used by gala for evolving the orbits of binaries in the galactic potential
+        (default is :class:`~gala.integrate.DOPRI853Integrator`).
+    integrator_kwargs : `dict`, optional
+        Any additional keyword arguments to pass to the gala integrator, by default an empty dict
     """
     def __init__(self, n_binaries, processes=8, m1_cutoff=0, final_kstar1=list(range(16)),
                  final_kstar2=list(range(16)), sfh_model=sfh.Wagg2022, sfh_params={},
                  galactic_potential=gp.MilkyWayPotential(), v_dispersion=5 * u.km / u.s,
                  max_ev_time=12.0*u.Gyr, timestep_size=1 * u.Myr, BSE_settings={}, ini_file=None,
-                 sampling_params={}, bcm_timestep_conditions=[], store_entire_orbits=True):
+                 sampling_params={}, bcm_timestep_conditions=[], store_entire_orbits=True, 
+                 integrator=gi.DOPRI853Integrator, integrator_kwargs={"a_tol": 1e-10, "r_tol": 1e-10}):
 
         # require a sensible number of binaries if you are not targetting total mass
         if not ("sampling_target" in sampling_params and sampling_params["sampling_target"] == "total_mass"):
@@ -105,6 +112,8 @@ def __init__(self, n_binaries, processes=8, m1_cutoff=0, final_kstar1=list(range
         self.timestep_size = timestep_size
         self.pool = None
         self.store_entire_orbits = store_entire_orbits
+        self.integrator = integrator
+        self.integrator_kwargs = integrator_kwargs
 
         self._file = None
         self._initial_binaries = None
@@ -235,7 +244,9 @@ def __getitem__(self, ind):
                                  v_dispersion=self.v_dispersion, max_ev_time=self.max_ev_time,
                                  timestep_size=self.timestep_size, BSE_settings=self.BSE_settings,
                                  sampling_params=self.sampling_params,
-                                 store_entire_orbits=self.store_entire_orbits)
+                                 store_entire_orbits=self.store_entire_orbits,
+                                 integrator=self.integrator,
+                                 integrator_kwargs=self.integrator_kwargs)
         new_pop.n_binaries_match = new_pop.n_binaries
 
         # proxy for checking whether sampling has been done
@@ -1091,11 +1102,13 @@ def perform_galactic_evolution(self, quiet=False, progress_bar=True):
             # setup arguments to combine primary and secondaries into a single list
             primary_args = [(w0s[i], self.max_ev_time - self.initial_galaxy.tau[i], self.max_ev_time,
                              copy(self.timestep_size), self.galactic_potential,
-                             primary_events[i], self.store_entire_orbits, quiet)
+                             primary_events[i], self.store_entire_orbits, quiet, 
+                             self.integrator, self.integrator_kwargs)
                             for i in range(self.n_binaries_match)]
             secondary_args = [(w0s[i], self.max_ev_time - self.initial_galaxy.tau[i], self.max_ev_time,
                                copy(self.timestep_size), self.galactic_potential,
-                               secondary_events[i], self.store_entire_orbits, quiet)
+                               secondary_events[i], self.store_entire_orbits, quiet, 
+                               self.integrator, self.integrator_kwargs)
                               for i in range(self.n_binaries_match) if secondary_events[i] is not None]
             args = primary_args + secondary_args
 
@@ -1119,15 +1132,19 @@ def perform_galactic_evolution(self, quiet=False, progress_bar=True):
                                                           t1=self.max_ev_time - self.initial_galaxy.tau[i],
                                                           t2=self.max_ev_time, dt=copy(self.timestep_size),
                                                           events=primary_events[i], quiet=quiet,
-                                                          store_all=self.store_entire_orbits))
+                                                          store_all=self.store_entire_orbits),
+                                                          integrator=self.integrator,
+                                                          integrator_kwargs=self.integrator_kwargs)
             for i in range(self.n_binaries_match):
                 if secondary_events[i] is None:
                     continue
                 orbits.append(integrate_orbit_with_events(w0=w0s[i], potential=self.galactic_potential,
                                                           t1=self.max_ev_time - self.initial_galaxy.tau[i],
                                                           t2=self.max_ev_time, dt=copy(self.timestep_size),
                                                           events=secondary_events[i], quiet=quiet,
-                                                          store_all=self.store_entire_orbits))
+                                                          store_all=self.store_entire_orbits),
+                                                          integrator=self.integrator,
+                                                          integrator_kwargs=self.integrator_kwargs)
 
         # check for bad orbits
         bad_orbits = np.array([orbit is None for orbit in orbits])
@@ -1633,6 +1650,7 @@ def save(self, file_name, overwrite=False):
         with h5.File(file_name, "a") as f:
             f.attrs["potential_dict"] = yaml.dump(potential_to_dict(self.galactic_potential),
                                                   default_flow_style=None)
+
         if self._initial_galaxy is not None:
             self.initial_galaxy.save(file_name, key="initial_galaxy")
 
@@ -1682,6 +1700,11 @@ def save(self, file_name, overwrite=False):
             # save sampling params
             d = file.create_dataset("sampling_params", data=[])
             d.attrs["dict"] = yaml.dump(self.sampling_params, default_flow_style=None)
+
+            # save integrator settings
+            d = file.create_dataset("integrator_settings", data=[])
+            d.attrs["integrator"] = self.integrator.__name__
+            d.attrs["integrator_kwargs"] = yaml.dump(self.integrator_kwargs, default_flow_style=None)
 
 
 def load(file_name, parts=["initial_binaries", "initial_galaxy", "stellar_evolution"]):
@@ -1723,6 +1746,13 @@ def load(file_name, parts=["initial_binaries", "initial_galaxy", "stellar_evolut
 
         sampling_params = yaml.load(file["sampling_params"].attrs["dict"], Loader=yaml.Loader)
 
+        integrator_name = file.get("integrator", None)
+        integrator_kwargs = file.get("integrator_kwargs", {"a_tol": 1e-10, "r_tol": 1e-10})
+        if integrator_name is None:
+            integrator = gi.DOPRI853Integrator
+        else:
+            integrator = getattr(gi, integrator_name)
+
     with h5.File(file_name, 'r') as f:
         galactic_potential = potential_from_dict(yaml.load(f.attrs["potential_dict"], Loader=yaml.Loader))
 
@@ -1732,7 +1762,8 @@ def load(file_name, parts=["initial_binaries", "initial_galaxy", "stellar_evolut
                    v_dispersion=numeric_params[4] * u.km / u.s, max_ev_time=numeric_params[5] * u.Gyr,
                    timestep_size=numeric_params[6] * u.Myr, BSE_settings=BSE_settings,
                    sampling_params=sampling_params, store_entire_orbits=store_entire_orbits,
-                   bcm_timestep_conditions=bcm_tc)
+                   bcm_timestep_conditions=bcm_tc,
+                   integrator=integrator, integrator_kwargs=integrator_kwargs)
 
     p._file = file_name
     p.n_binaries_match = int(numeric_params[1])

docs/modules/changelog.rst

@@ -4,6 +4,10 @@ Full changelog
 
 This page tracks all of the changes that have been made to ``cogsworth``. We follow the standard versioning convention of A.B.C, where C is a patch/bugfix, B is a large bugfix or new feature and A is a major new breaking change. B/C are backwards compatible but A changes may be breaking.
 
+2.1.2
+=====
+- New feature: added optional key word arguments `integrator` and `integrator_kwargs` to `Population` class to allow users to specify the integrator settings for `gala`
+
 2.1.0
 =====