✨ feat(dynamics): add solve namespace

Also, moving some stuff out of the integrate namespace

Author: nstarman

src/galax/_interop/galax_interop_astropy/dynamics.py

@@ -25,7 +25,7 @@ def evaluate_orbit(
     t: APYQuantity,
     /,
     *,
-    solver: gd.integrate.DynamicsSolver | None = None,
+    solver: gd.DynamicsSolver | None = None,
     dense: Literal[True, False] = False,
 ) -> gd.Orbit:
     """Compute an orbit in a potential.

src/galax/dynamics/__init__.py

@@ -3,13 +3,16 @@
 __all__ = [
     # Modules
     "fields",
-    "integrate",
+    "solve",
+    "integrate",  # TODO: deprecate
     "mockstream",
     "plot",
     "cluster",
-    # integrate
+    # solve
     "evaluate_orbit",
     "Orbit",
+    "AbstractSolver",
+    "DynamicsSolver",
     # mockstream
     "MockStreamArm",
     "MockStream",
@@ -35,7 +38,6 @@
     from ._src.api import omega, specific_angular_momentum
     from ._src.cluster.funcs import lagrange_points, tidal_radius
     from ._src.orbit import Orbit
-    from .integrate import evaluate_orbit
     from .mockstream import (
         AbstractStreamDF,
         ChenStreamDF,
@@ -44,6 +46,7 @@
         MockStreamArm,
         MockStreamGenerator,
     )
+    from .solve import AbstractSolver, DynamicsSolver, evaluate_orbit
 
     #
     # isort: split

src/galax/dynamics/_src/dynamics/field_hamiltonian.py

@@ -77,7 +77,7 @@ class HamiltonianField(AbstractDynamicsField, strict=True):  # type: ignore[call
     Just to continue the example, we can use this field to integrate the
     equations of motion:
 
-    >>> solver = gd.integrate.DynamicsSolver()  # defaults to Dopri8
+    >>> solver = gd.DynamicsSolver()  # defaults to Dopri8
     >>> w0 = gc.PhaseSpacePosition(
     ...     q=u.Quantity([[8, 0, 9], [9, 0, 3]], "kpc"),
     ...     p=u.Quantity([0, 220, 0], "km/s"),
@@ -268,7 +268,7 @@ def terms(
 
     For completeness we'll integrate the EoM.
 
-    >>> dynamics_solver = gd.integrate.DynamicsSolver(solver)
+    >>> dynamics_solver = gd.DynamicsSolver(solver)
     >>> w0 = gc.PhaseSpacePosition(
     ...     q=u.Quantity([8., 0, 0], "kpc"),
     ...     p=u.Quantity([0, 220, 0], "km/s"),

src/galax/dynamics/_src/dynamics/field_nbody.py

@@ -40,7 +40,7 @@ class NBodyField(AbstractDynamicsField, strict=True):  # type: ignore[call-arg]
     >>> q = u.Quantity([[-1, 0, 0], [1, 0, 0]], "AU") / 2
     >>> p = u.Quantity([[0, -1, 0], [0, 1, 0]], "km/s") * 25
 
-    >>> solver = gd.integrate.DynamicsSolver()
+    >>> solver = gd.DynamicsSolver()
 
     >>> field = gd.fields.NBodyField(
     ...     masses=u.Quantity([1, 1], "Msun"),

src/galax/dynamics/_src/dynamics/register_gc.py

@@ -38,7 +38,7 @@ def from_(
     >>> pot = gp.HernquistPotential(m_tot=u.Quantity(1e12, "Msun"),
     ...    r_s=u.Quantity(5, "kpc"), units="galactic")
     >>> field = gd.fields.HamiltonianField(pot)
-    >>> solver = gd.integrate.DynamicsSolver()  # defaults to Dopri8
+    >>> solver = gd.DynamicsSolver()  # defaults to Dopri8
     >>> w0 = gc.PhaseSpacePosition(
     ...     q=u.Quantity([[8, 0, 9], [9, 0, 3]], "kpc"),
     ...     p=u.Quantity([0, 220, 0], "km/s"),

src/galax/dynamics/_src/dynamics/solver.py

@@ -50,7 +50,7 @@ class DynamicsSolver(AbstractSolver, strict=True):  # type: ignore[call-arg]
     >>> import galax.potential as gp
     >>> import galax.dynamics as gd
 
-    >>> solver = gd.integrate.DynamicsSolver()  # defaults to Dopri8
+    >>> solver = gd.DynamicsSolver()  # defaults to Dopri8
 
     Define the vector field. In this example it's to solve Hamilton's EoM in a
     gravitational potential.
@@ -92,9 +92,9 @@ class DynamicsSolver(AbstractSolver, strict=True):  # type: ignore[call-arg]
        setting the `diffrax.AbstractSolver`,
        `diffrax.AbstractStepSizeController`, etc.
 
-    >>> diffeqsolver = gd.integrate.DiffEqSolver(dfx.Dopri8(),
+    >>> diffeqsolver = gd.solve.DiffEqSolver(dfx.Dopri8(),
     ...     stepsize_controller=dfx.PIDController(rtol=1e-5, atol=1e-5))
-    >>> solver = gd.integrate.DynamicsSolver(diffeqsolver)
+    >>> solver = gd.DynamicsSolver(diffeqsolver)
     >>> solver
     DynamicsSolver(
       diffeqsolver=DiffEqSolver(
@@ -107,7 +107,7 @@ class DynamicsSolver(AbstractSolver, strict=True):  # type: ignore[call-arg]
     2. A `dict` of keyword arguments that are passed to
        `galax.dynamics.integrate.DiffEqSolver`.
 
-    >>> solver = gd.integrate.DynamicsSolver({
+    >>> solver = gd.DynamicsSolver({
     ...     "solver": dfx.Dopri8(), "stepsize_controller": dfx.ConstantStepSize()})
     >>> solver
     DynamicsSolver(
@@ -187,7 +187,7 @@ def solve(
         >>> import galax.potential as gp
         >>> import galax.dynamics as gd
 
-        >>> solver = gd.integrate.DynamicsSolver()
+        >>> solver = gd.DynamicsSolver()
 
         Specify the vector field.
 
@@ -658,7 +658,7 @@ def terms(
     >>> import galax.potential as gp
     >>> import galax.dynamics as gd
 
-    >>> solver = gd.integrate.DynamicsSolver(dfx.Dopri8())
+    >>> solver = gd.DynamicsSolver(dfx.Dopri8())
 
     >>> pot = gp.KeplerPotential(m_tot=u.Quantity(1e12, "Msun"), units="galactic")
     >>> field = gd.fields.HamiltonianField(pot)

src/galax/dynamics/_src/fields.py

@@ -43,7 +43,7 @@ def terms(
     >>> import galax.potential as gp
     >>> import galax.dynamics as gd
 
-    >>> solver = gd.integrate.DiffEqSolver(dfx.Dopri8())
+    >>> solver = gd.solve.DiffEqSolver(dfx.Dopri8())
 
     >>> pot = gp.KeplerPotential(m_tot=u.Quantity(1e12, "Msun"), units="galactic")
     >>> field = gd.fields.HamiltonianField(pot)

src/galax/dynamics/_src/solver.py

@@ -62,9 +62,9 @@ def from_(cls: type[AbstractSolver], solver: AbstractSolver) -> AbstractSolver:
 
     Examples
     --------
-    >>> from galax.dynamics.integrate import DynamicsSolver
-    >>> solver = DynamicsSolver()
-    >>> new_solver = DynamicsSolver.from_(solver)
+    >>> import galax.dynamics as gd
+    >>> solver = gd.DynamicsSolver()
+    >>> new_solver = gd.DynamicsSolver.from_(solver)
     >>> new_solver is solver
     True
 
@@ -91,7 +91,7 @@ def from_(cls: type[AbstractSolver], obj: Any) -> AbstractSolver:
     Examples
     --------
     >>> import diffrax as dfx
-    >>> from galax.dynamics.integrate import DynamicsSolver, DiffEqSolver
+    >>> from galax.dynamics.solve import DynamicsSolver, DiffEqSolver
 
     >>> DynamicsSolver.from_( DiffEqSolver(dfx.Dopri5()))
     DynamicsSolver(
@@ -118,16 +118,16 @@ def from_(cls: type[AbstractSolver], obj: Mapping[str, Any]) -> AbstractSolver:
     Examples
     --------
     >>> import diffrax as dfx
-    >>> from galax.dynamics.integrate import DynamicsSolver
+    >>> import galax.dynamics as gd
 
-    >>> DynamicsSolver.from_({})
+    >>> gd.DynamicsSolver.from_({})
     DynamicsSolver(
       diffeqsolver=DiffEqSolver(
         solver=Dopri8(scan_kind=None),
         stepsize_controller=PIDController( ... ),
         adjoint=RecursiveCheckpointAdjoint(checkpoints=None) ) )
 
-    >>> DynamicsSolver.from_({"diffeqsolver": dfx.Dopri5()})
+    >>> gd.DynamicsSolver.from_({"diffeqsolver": dfx.Dopri5()})
     DynamicsSolver(
       diffeqsolver=DiffEqSolver(
         solver=Dopri5(scan_kind=None),

src/galax/dynamics/integrate.py

@@ -5,29 +5,21 @@
     "Integrator",
     "Interpolant",
     "parse_time_specification",
-    "AbstractSolver",
-    "DynamicsSolver",
     "InterpolatedPhaseSpacePosition",
-    # Diffraxtra external library
-    "DiffEqSolver",
-    "VectorizedDenseInterpolation",
 ]
 
 from jaxtyping import install_import_hook
 
 from galax.setup_package import RUNTIME_TYPECHECKER
 
 with install_import_hook("galax.dynamics.integrate", RUNTIME_TYPECHECKER):
-    from diffraxtra import DiffEqSolver, VectorizedDenseInterpolation
-
-    from ._src.dynamics import DynamicsSolver, parse_time_specification
+    from ._src.dynamics import parse_time_specification
     from ._src.integrate import (
         Integrator,
         Interpolant,
         InterpolatedPhaseSpacePosition,
         evaluate_orbit,
     )
-    from ._src.solver import AbstractSolver
 
 # Cleanup
 del install_import_hook, RUNTIME_TYPECHECKER

src/galax/dynamics/solve.py

@@ -0,0 +1,26 @@
+""":mod:`galax.dynamics.solve`."""
+
+__all__ = [
+    "evaluate_orbit",
+    "parse_time_specification",
+    "AbstractSolver",
+    "DynamicsSolver",
+    "DiffEqSolver",
+    "VectorizedDenseInterpolation",
+    "Orbit",
+]
+
+from jaxtyping import install_import_hook
+
+from galax.setup_package import RUNTIME_TYPECHECKER
+
+with install_import_hook("galax.dynamics.integrate", RUNTIME_TYPECHECKER):
+    from diffraxtra import DiffEqSolver, VectorizedDenseInterpolation
+
+    from ._src.dynamics import DynamicsSolver, parse_time_specification
+    from ._src.integrate import evaluate_orbit
+    from ._src.orbit import Orbit
+    from ._src.solver import AbstractSolver
+
+# Cleanup
+del install_import_hook, RUNTIME_TYPECHECKER