diff --git a/src/galax/dynamics/__init__.py b/src/galax/dynamics/__init__.py index 1d4d46bd..3cd56fbc 100644 --- a/src/galax/dynamics/__init__.py +++ b/src/galax/dynamics/__init__.py @@ -15,9 +15,10 @@ "AbstractSolver", "DynamicsSolver", # mockstream + "StreamSimulator", "MockStreamArm", "MockStream", - "MockStreamGenerator", + "MockStreamGenerator", # TODO: deprecate # mockstream.df "AbstractStreamDF", "FardalStreamDF", @@ -47,6 +48,7 @@ MockStream, MockStreamArm, MockStreamGenerator, + StreamSimulator, ) from .solve import AbstractSolver, DynamicsSolver diff --git a/src/galax/dynamics/_src/mockstream/__init__.py b/src/galax/dynamics/_src/mockstream/__init__.py index 71ad0023..8b909bbb 100644 --- a/src/galax/dynamics/_src/mockstream/__init__.py +++ b/src/galax/dynamics/_src/mockstream/__init__.py @@ -5,9 +5,19 @@ """ __all__ = [ + "StreamSimulator", + # Coordinates "MockStream", "MockStreamArm", + # Phase-Space Distribution + "AbstractStreamDF", + "Fardal15StreamDF", + "Chen24StreamDF", ] from .arm import MockStreamArm from .core import MockStream +from .df_base import AbstractStreamDF +from .df_chen24 import Chen24StreamDF +from .df_fardal15 import Fardal15StreamDF +from .simulate import StreamSimulator diff --git a/src/galax/dynamics/_src/mockstream/arm.py b/src/galax/dynamics/_src/mockstream/arm.py index d59b466c..89a8b689 100644 --- a/src/galax/dynamics/_src/mockstream/arm.py +++ b/src/galax/dynamics/_src/mockstream/arm.py @@ -4,12 +4,14 @@ from typing import Any, ClassVar, Protocol, cast, final, runtime_checkable +import diffrax as dfx import equinox as eqx from plum import dispatch import coordinax as cx import quaxed.numpy as jnp import unxt as u +from unxt.quantity import BareQuantity import galax._custom_types as gt import galax.coordinates as gc @@ -67,6 +69,41 @@ def _shape_tuple(self) -> tuple[gt.Shape, gc.ComponentShapeTuple]: ##################################################################### + +@gc.AbstractPhaseSpaceObject.from_.dispatch # type: ignore[attr-defined,misc] +def from_( + cls: type[MockStreamArm], + soln: dfx.Solution, + /, + *, + release_time: gt.BBtQuSz0, + frame: cx.frames.AbstractReferenceFrame, + units: u.AbstractUnitSystem, # not dispatched on, but required + unbatch_time: bool = True, +) -> MockStreamArm: + """Create a new instance of the class.""" + # Reshape (*tbatch, T, *ybatch) to (*tbatch, *ybatch, T) + t = soln.ts # already in the shape (*tbatch, T) + n_tbatch = soln.t0.ndim + q = jnp.moveaxis(soln.ys[0], n_tbatch, -2) + p = jnp.moveaxis(soln.ys[1], n_tbatch, -2) + + # Reshape (*tbatch, *ybatch, T) to (*tbatch, *ybatch) if T == 1 + if unbatch_time and t.shape[-1] == 1: + t = t[..., -1] + q = q[..., -1, :] + p = p[..., -1, :] + + # Convert the solution to a phase-space position + return cls( + q=cx.CartesianPos3D.from_(q, units["length"]), + p=cx.CartesianVel3D.from_(p, units["speed"]), + t=BareQuantity(t, units["time"]), + release_time=release_time, + frame=frame, + ) + + # ========================================================= # `__getitem__` diff --git a/src/galax/dynamics/_src/mockstream/df_base.py b/src/galax/dynamics/_src/mockstream/df_base.py new file mode 100644 index 00000000..d4bc4460 --- /dev/null +++ b/src/galax/dynamics/_src/mockstream/df_base.py @@ -0,0 +1,55 @@ +"""Stream Distribution Functions for ejecting mock stream particles.""" + +__all__ = ["AbstractStreamDF"] + +import abc +from typing import TypeAlias + +import equinox as eqx +from jaxtyping import PRNGKeyArray + +import galax._custom_types as gt +import galax.potential as gp + +Carry: TypeAlias = tuple[gt.QuSz3, gt.QuSz3, gt.QuSz3, gt.QuSz3] + + +class AbstractStreamDF(eqx.Module, strict=True): # type: ignore[call-arg, misc] + """Abstract base class of Stream Distribution Functions.""" + + # TODO: keep units and PSP through this func + @abc.abstractmethod + def sample( + self, + key: PRNGKeyArray, + potential: gp.AbstractPotential, + x: gt.BBtQuSz3, + v: gt.BBtQuSz3, + prog_mass: gt.BBtFloatQuSz0, + t: gt.BBtFloatQuSz0, + ) -> tuple[gt.BtQuSz3, gt.BtQuSz3, gt.BtQuSz3, gt.BtQuSz3]: + """Generate stream particle initial conditions. + + Parameters + ---------- + rng : :class:`jaxtyping.PRNGKeyArray` + Pseudo-random number generator. + potential : :class:`galax.potential.AbstractPotential` + The potential of the host galaxy. + x : Quantity[float, (*#batch, 3), "length"] + 3d position (x, y, z) + v : Quantity[float, (*#batch, 3), "speed"] + 3d velocity (v_x, v_y, v_z) + prog_mass : Quantity[float, (*#batch), "mass"] + Mass of the progenitor. + t : Quantity[float, (*#batch), "time"] + The release time of the stream particles. + + Returns + ------- + x_lead, v_lead: Quantity[float, (*batch, 3), "length" | "speed"] + Position and velocity of the leading arm. + x_trail, v_trail : Quantity[float, (*batch, 3), "length" | "speed"] + Position and velocity of the trailing arm. + """ + ... diff --git a/src/galax/dynamics/_src/mockstream/df_chen24.py b/src/galax/dynamics/_src/mockstream/df_chen24.py new file mode 100644 index 00000000..b21eb028 --- /dev/null +++ b/src/galax/dynamics/_src/mockstream/df_chen24.py @@ -0,0 +1,136 @@ +"""galax: Galactic Dynamix in Jax.""" + +__all__ = ["Chen24StreamDF"] + + +import warnings +from functools import partial +from typing import final + +import jax +import jax.random as jr +from jaxtyping import PRNGKeyArray + +import coordinax as cx +import quaxed.numpy as jnp + +import galax._custom_types as gt +import galax.potential as gp +from .df_base import AbstractStreamDF +from galax.dynamics._src.cluster.radius import tidal_radius +from galax.dynamics._src.register_api import specific_angular_momentum + +# ============================================================ +# Constants + +mean = jnp.array([1.6, -30, 0, 1, 20, 0]) + +cov = jnp.array( + [ + [0.1225, 0, 0, 0, -4.9, 0], + [0, 529, 0, 0, 0, 0], + [0, 0, 144, 0, 0, 0], + [0, 0, 0, 0, 0, 0], + [-4.9, 0, 0, 0, 400, 0], + [0, 0, 0, 0, 0, 484], + ] +) + +# ============================================================ + + +@final +class Chen24StreamDF(AbstractStreamDF): + """Chen Stream Distribution Function. + + A class for representing the Chen+2024 distribution function for + generating stellar streams based on Chen et al. 2024 + https://ui.adsabs.harvard.edu/abs/2024arXiv240801496C/abstract + """ + + def __init__(self) -> None: + super().__init__() + warnings.warn( + 'Currently only the "no progenitor" version ' + "of the Chen+24 model is supported!", + RuntimeWarning, + stacklevel=1, + ) + + @partial(jax.jit, inline=True) + def sample( + self, + key: PRNGKeyArray, + potential: gp.AbstractPotential, + x: gt.BBtQuSz3, + v: gt.BBtQuSz3, + prog_mass: gt.BBtFloatQuSz0, + t: gt.BBtFloatQuSz0, + ) -> tuple[gt.BtQuSz3, gt.BtQuSz3, gt.BtQuSz3, gt.BtQuSz3]: + """Generate stream particle initial conditions.""" + # Random number generation + + # x_new-hat + r = jnp.linalg.vector_norm(x, axis=-1, keepdims=True) + x_new_hat = x / r + + # z_new-hat + L_vec = specific_angular_momentum(x, v) + z_new_hat = cx.vecs.normalize_vector(L_vec) + + # y_new-hat + phi_vec = v - jnp.sum(v * x_new_hat, axis=-1, keepdims=True) * x_new_hat + y_new_hat = cx.vecs.normalize_vector(phi_vec) + + r_tidal = tidal_radius(potential, x, v, mass=prog_mass, t=t) + + # Bill Chen: method="cholesky" doesn't work here! + posvel = jr.multivariate_normal( + key, mean, cov, shape=r_tidal.shape, method="svd" + ) + + Dr = posvel[:, 0] * r_tidal + + v_esc = jnp.sqrt(2 * potential.constants["G"] * prog_mass / Dr) + Dv = posvel[:, 3] * v_esc + + # convert degrees to radians + phi = posvel[:, 1] * 0.017453292519943295 + theta = posvel[:, 2] * 0.017453292519943295 + alpha = posvel[:, 4] * 0.017453292519943295 + beta = posvel[:, 5] * 0.017453292519943295 + + ctheta, stheta = jnp.cos(theta), jnp.sin(theta) + cphi, sphi = jnp.cos(phi), jnp.sin(phi) + calpha, salpha = jnp.cos(alpha), jnp.sin(alpha) + cbeta, sbeta = jnp.cos(beta), jnp.sin(beta) + + # Trailing arm + x_trail = ( + x + + (Dr * ctheta * cphi)[:, None] * x_new_hat + + (Dr * ctheta * sphi)[:, None] * y_new_hat + + (Dr * stheta)[:, None] * z_new_hat + ) + v_trail = ( + v + + (Dv * cbeta * calpha)[:, None] * x_new_hat + + (Dv * cbeta * salpha)[:, None] * y_new_hat + + (Dv * sbeta)[:, None] * z_new_hat + ) + + # Leading arm + x_lead = ( + x + - (Dr * ctheta * cphi)[:, None] * x_new_hat + - (Dr * ctheta * sphi)[:, None] * y_new_hat + + (Dr * stheta)[:, None] * z_new_hat + ) + v_lead = ( + v + - (Dv * cbeta * calpha)[:, None] * x_new_hat + - (Dv * cbeta * salpha)[:, None] * y_new_hat + + (Dv * sbeta)[:, None] * z_new_hat + ) + + return x_lead, v_lead, x_trail, v_trail diff --git a/src/galax/dynamics/_src/mockstream/df_fardal15.py b/src/galax/dynamics/_src/mockstream/df_fardal15.py new file mode 100644 index 00000000..cb9a2772 --- /dev/null +++ b/src/galax/dynamics/_src/mockstream/df_fardal15.py @@ -0,0 +1,93 @@ +"""galax: Galactic Dynamix in Jax.""" + +__all__ = ["Fardal15StreamDF"] + + +from functools import partial +from typing import final + +import jax +import jax.random as jr +from jaxtyping import PRNGKeyArray + +import coordinax as cx +import quaxed.numpy as jnp + +import galax._custom_types as gt +import galax.potential as gp +from .df_base import AbstractStreamDF +from galax.dynamics._src.api import omega +from galax.dynamics._src.cluster.radius import tidal_radius + +# ============================================================ +# Constants + +kr_bar = 2.0 +kvphi_bar = 0.3 + +kz_bar = 0.0 +kvz_bar = 0.0 + +sigma_kr = 0.5 # TODO: use actual Fardal values +sigma_kvphi = 0.5 # TODO: use actual Fardal values +sigma_kz = 0.5 +sigma_kvz = 0.5 + +# ============================================================ + + +@final +class Fardal15StreamDF(AbstractStreamDF): + """Fardal Stream Distribution Function. + + A class for representing the Fardal+2015 distribution function for + generating stellar streams based on Fardal et al. 2015 + https://ui.adsabs.harvard.edu/abs/2015MNRAS.452..301F/abstract + """ + + @partial(jax.jit) + def sample( + self, + key: PRNGKeyArray, + potential: gp.AbstractPotential, + x: gt.BBtQuSz3, + v: gt.BBtQuSz3, + prog_mass: gt.BBtFloatQuSz0, + t: gt.BBtFloatQuSz0, + ) -> tuple[gt.BtQuSz3, gt.BtQuSz3, gt.BtQuSz3, gt.BtQuSz3]: + """Generate stream particle initial conditions.""" + # Random number generation + key1, key2, key3, key4 = jr.split(key, 4) + + om = omega(x, v)[..., None] + + # r-hat + r_hat = cx.vecs.normalize_vector(x) + + r_tidal = tidal_radius(potential, x, v, mass=prog_mass, t=t)[..., None] + v_circ = om * r_tidal # relative velocity + + # z-hat + L_vec = jnp.linalg.cross(x, v) + z_hat = cx.vecs.normalize_vector(L_vec) + + # phi-hat + phi_vec = v - jnp.sum(v * r_hat, axis=-1, keepdims=True) * r_hat + phi_hat = cx.vecs.normalize_vector(phi_vec) + + # k vals + shape = r_tidal.shape + kr_samp = kr_bar + jr.normal(key1, shape) * sigma_kr + kvphi_samp = kr_samp * (kvphi_bar + jr.normal(key2, shape) * sigma_kvphi) + kz_samp = kz_bar + jr.normal(key3, shape) * sigma_kz + kvz_samp = kvz_bar + jr.normal(key4, shape) * sigma_kvz + + # Trailing arm + x_trail = x + r_tidal * (kr_samp * r_hat + kz_samp * z_hat) + v_trail = v + v_circ * (kvphi_samp * phi_hat + kvz_samp * z_hat) + + # Leading arm + x_lead = x - r_tidal * (kr_samp * r_hat - kz_samp * z_hat) + v_lead = v - v_circ * (kvphi_samp * phi_hat - kvz_samp * z_hat) + + return x_lead, v_lead, x_trail, v_trail diff --git a/src/galax/dynamics/_src/mockstream/simulate.py b/src/galax/dynamics/_src/mockstream/simulate.py new file mode 100644 index 00000000..6dd64dfa --- /dev/null +++ b/src/galax/dynamics/_src/mockstream/simulate.py @@ -0,0 +1,338 @@ +"""Simulate a mockstream. + +This is private API. + +""" + +__all__ = ["simulate_stream", "StreamSimulator"] + +from dataclasses import dataclass +from typing import Any + +import equinox as eqx +import jax.random as jr +from jaxtyping import PRNGKeyArray +from plum import convert + +import quaxed.numpy as jnp +import unxt as u +from dataclassish.converters import Unless + +import galax.coordinates as gc +import galax.potential as gp +from .arm import MockStreamArm +from .core import MockStream +from .df_base import AbstractStreamDF +from .df_fardal15 import Fardal15StreamDF +from galax.dynamics._src.cluster.fields import AbstractMassRateField +from galax.dynamics._src.cluster.sample import ReleaseTimeSampler +from galax.dynamics._src.cluster.solver import MassSolver +from galax.dynamics._src.dynamics.field_base import AbstractDynamicsField +from galax.dynamics._src.dynamics.field_hamiltonian import HamiltonianField +from galax.dynamics._src.dynamics.solver import DynamicsSolver +from galax.dynamics._src.orbit import Orbit, compute_orbit + +default_dynamics_solver = DynamicsSolver() +default_dynamics_release_model = Fardal15StreamDF() + +# default_ +# converter_dynamics_solver = lambda x: +converter_dynamics_field = Unless(AbstractDynamicsField, HamiltonianField) + + +@dataclass +class StreamSimulator: + """Simulate a mock stellar stream. + + Notes + ----- + - Support solving for the mass history of the progenitor using the + mass_history_solver, in which case a mass_release_model is required. + - Support pre-specified stripping times, in which case the + mass_history_solver is not required. The mass loss model can be given and + will be used to change the mass of the cluster. Note that this will not be + self-consistent in the mass history. + + """ + + #: Solver for the dynamics of the progenitor and stream stars. + dynamics_solver: DynamicsSolver = default_dynamics_solver + + #: The position & velocity release model for the stream stars. + dynamics_release_model: AbstractStreamDF = default_dynamics_release_model + + #: Solver for the mass history of the progenitor. If `None` (default), the + #: mass history is not solved. Instead the stream simulator requires an + #: array of stripping times. + mass_history_solver: MassSolver | None = None + + # TODO: figure out how to fold this into the mass history solver + #: The mass field to use for the mass history. If `None` (default), the + #: progenitor doee not lose mass. An error will be raised if the + #: `mass_history_solver` is not None. + mass_release_model: AbstractMassRateField | None = None + + def __check_init__(self) -> None: + if self.mass_history_solver is not None and self.mass_release_model is None: + msg = "`mass_history_solver` is not None, so a `mass_release_model` must be provided." # noqa: E501 + raise ValueError(msg) + + def _call_nomasshistory( + self, + dynamics_field: AbstractDynamicsField, + w0: gc.PhaseSpaceCoordinate, + M0: u.AbstractQuantity, # noqa: ARG002 + t0: u.AbstractQuantity, # noqa: ARG002 + t1: u.AbstractQuantity, # noqa: ARG002 + /, + stripping_times: u.AbstractQuantity | None = None, + **_: Any, + ) -> tuple[Orbit, MockStream]: + _ = compute_orbit( + dynamics_field, + w0, + stripping_times, + dense=False, + solver=self.dynamics_solver, + ) + + raise NotImplementedError + + def _call_masshistory( + self, + dynamics_field: AbstractDynamicsField, + w0: gc.PhaseSpaceCoordinate, + M0: u.AbstractQuantity, + t0: u.AbstractQuantity, + t1: u.AbstractQuantity, + /, + key: PRNGKeyArray, + n_stars: int, + mass_params: dict[str, Any] = {}, # noqa: B006 + **_: Any, + ) -> tuple[Orbit, MockStream]: + """Simulate a mock stellar stream with mass history.""" + mass_solver = eqx.error_if( + self.mass_history_solver, + self.mass_history_solver is None, + "`mass_history_solver` is None", + ) + mass_loss_rate = eqx.error_if( + self.mass_release_model, + self.mass_release_model is None, + "`mass_release_model` is None", + ) + + # Step 1) Solve the progenitor's orbit + ts = jnp.linspace(t0, t1, 2) + orbit = compute_orbit( + dynamics_field, w0, ts, dense=True, solver=self.dynamics_solver + ) + + # Step 2) Solve the mass history of the progenitor + mass_params = mass_params.copy() + mass_params["orbit"] = orbit + mass_history = mass_solver.solve( + mass_loss_rate, + M0, + t0, + t1, + args=mass_params, + dense=True, + vectorize_interpolation=True, + ) + + # Step 3) Sample release model. + # A) Release time from mass history + release_time_sampler = ReleaseTimeSampler( + mass_loss_rate, mass_history, orbit.potential.units + ) + key, subkey = jr.split(key) + release_times = release_time_sampler.sample( + subkey, t0, t1, n_stars=n_stars, mass_params=mass_params + ) + + # B) star positions + key, subkey = jr.split(key) + o_at_t = orbit(release_times) + x_lead, v_lead, x_trail, v_trail = default_dynamics_release_model.sample( + subkey, + orbit.potential, + convert(o_at_t.q, u.Quantity), + convert(o_at_t.p, u.Quantity), + M0, + release_times, + ) + wlead = gc.PhaseSpaceCoordinate( + q=x_lead, p=v_lead, t=release_times, frame=w0.frame + ) + wtrail = gc.PhaseSpaceCoordinate( + q=x_trail, p=v_trail, t=release_times, frame=w0.frame + ) + + # Step 4) Solve the orbits of the stream stars + # TODO: enable dense model + soln_lead = self.dynamics_solver.solve( + dynamics_field, + wlead, + t1, + dense=False, + vectorize_interpolation=True, + ) + stream_lead = MockStreamArm.from_( + soln_lead, release_time=release_times, frame=wlead.frame + ) + + soln_trail = self.dynamics_solver.solve( + dynamics_field, + wtrail, + t1, + dense=False, + vectorize_interpolation=True, + ) + stream_trail = MockStreamArm.from_( + soln_trail, release_time=release_times, frame=wtrail.frame + ) + + return orbit, MockStream(lead=stream_lead, trail=stream_trail) + + def __call__( + self, + dynamics_field: AbstractDynamicsField, + w0: gc.PhaseSpaceCoordinate, + M0: u.AbstractQuantity, + t0: u.AbstractQuantity, + t1: u.AbstractQuantity, + /, + stripping_times: u.AbstractQuantity | None = None, + **kw: Any, + ) -> Any: + # Step 0) Prep the inputs + # Sort the times + t0 = jnp.minimum(t0, t1) + t1 = jnp.maximum(t0, t1) + + # Step 1) Solve the progenitor's orbit + if self.mass_history_solver is None: + t_strip = eqx.error_if( + stripping_times, stripping_times is None, "`stripping_times` is None" + ) + out = self._call_nomasshistory( + dynamics_field, + w0, + M0, + t0, + t1, + stripping_times=t_strip, + ) + + else: + out = self._call_masshistory( + dynamics_field, + w0, + M0, + t0, + t1, + key=kw["key"], + n_stars=kw["n_stars"], + mass_params=kw["mass_params"], + ) + + return out + + +# =================================================================== + + +def simulate_stream( + key: PRNGKeyArray, + dynamics_field: AbstractDynamicsField | gp.AbstractPotential, + w0: gc.PhaseSpaceCoordinate, + t0: u.AbstractQuantity, # start time of the stream, e.g. -3 Gyr + t1: u.AbstractQuantity, # end time of the stream, e.g. 0 Gyr + n_stream_stars: int, # number of stream stars + Mc0: u.AbstractQuantity, # initial mass of the progenitor + mass_params: dict[str, Any], + *, + dynamics_solver: DynamicsSolver | None = None, + mass_history_solver: MassSolver | None = None, + mass_field: AbstractMassRateField, # TODO: or callable +) -> Any: + """Simulate a mock stellar stream.""" + # Step 0) Prep the inputs + # Sort the times + t0 = jnp.minimum(t0, t1) + t1 = jnp.maximum(t0, t1) + dynamics_solver = DynamicsSolver() if dynamics_solver is None else dynamics_solver + dynamics_field = converter_dynamics_field(dynamics_field) + mass_solver = MassSolver() if mass_history_solver is None else mass_history_solver + + # Step 1) Solve the progenitor's orbit + orbit = compute_orbit( + dynamics_field, w0, jnp.linspace(t0, t1, 2), dense=True, solver=dynamics_solver + ) + + # Step 2) Solve the mass history of the progenitor + mass_params = mass_params.copy() + mass_params["orbit"] = orbit + mass_history = mass_solver.solve( + mass_field, + Mc0, + t0, + t1, + args=mass_params, + dense=True, + vectorize_interpolation=True, + ) + + # Step 3) Sample release model. + # A) Release time from mass history + release_time_sampler = ReleaseTimeSampler( + mass_field, mass_history, orbit.potential.units + ) + key, subkey = jr.split(key) + release_times = release_time_sampler.sample( + subkey, t0, t1, n_stars=n_stream_stars, mass_params=mass_params + ) + + # B) star positions + key, subkey = jr.split(key) + o_at_t = orbit(release_times) + x_lead, v_lead, x_trail, v_trail = default_dynamics_release_model.sample( + subkey, + orbit.potential, + convert(o_at_t.q, u.Quantity), + convert(o_at_t.p, u.Quantity), + Mc0, + release_times, + ) + wlead = gc.PhaseSpaceCoordinate(q=x_lead, p=v_lead, t=release_times, frame=w0.frame) + wtrail = gc.PhaseSpaceCoordinate( + q=x_trail, p=v_trail, t=release_times, frame=w0.frame + ) + + # Step 4) Solve the orbits of the stream stars + # TODO: enable dense model + soln_lead = dynamics_solver.solve( + dynamics_field, + wlead, + t1, + dense=False, + vectorize_interpolation=True, + ) + orbit_lead = Orbit.from_( + soln_lead, frame=wlead.frame, potential=dynamics_field.potential + ) + + soln_trail = dynamics_solver.solve( + dynamics_field, + wtrail, + t1, + dense=False, + vectorize_interpolation=True, + ) + orbit_trail = Orbit.from_( + soln_trail, frame=wtrail.frame, potential=dynamics_field.potential + ) + + return orbit_lead, orbit_trail, release_times diff --git a/src/galax/dynamics/mockstream.py b/src/galax/dynamics/mockstream.py index aaf8f470..8b23ba02 100644 --- a/src/galax/dynamics/mockstream.py +++ b/src/galax/dynamics/mockstream.py @@ -1,10 +1,11 @@ """:mod:`galax.dynamics.mockstream`.""" __all__ = [ - # Legacy - "MockStreamGenerator", + "StreamSimulator", "MockStreamArm", "MockStream", + # Legacy + "MockStreamGenerator", "AbstractStreamDF", "FardalStreamDF", "ChenStreamDF", @@ -25,7 +26,7 @@ FardalStreamDF, ProgenitorMassCallable, ) - from ._src.mockstream import MockStream, MockStreamArm + from ._src.mockstream import MockStream, MockStreamArm, StreamSimulator # Cleanup del install_import_hook, RUNTIME_TYPECHECKER