From 29f7d21e23d7da01c45b93397a44a03167e25fda Mon Sep 17 00:00:00 2001 From: nstarman Date: Thu, 1 Feb 2024 16:38:15 -0500 Subject: [PATCH 1/2] feat: move PSP to coordinates module Signed-off-by: nstarman --- docs/getting_started.rst | 11 ++++++----- src/galax/coordinates/__init__.py | 11 +++++++++++ .../_dynamics/base.py => coordinates/_base.py} | 2 +- .../_dynamics/core.py => coordinates/_core.py} | 2 +- .../_dynamics/utils.py => coordinates/_utils.py} | 0 src/galax/dynamics/_dynamics/__init__.py | 8 ++------ src/galax/dynamics/_dynamics/mockstream/core.py | 4 ++-- src/galax/dynamics/_dynamics/orbit.py | 6 +++--- tests/smoke/coordinates/__init__.py | 0 tests/smoke/coordinates/test_package.py | 9 +++++++++ tests/smoke/dynamics/test_package.py | 10 ++-------- tests/unit/dynamics/test_core.py | 2 +- 12 files changed, 38 insertions(+), 27 deletions(-) create mode 100644 src/galax/coordinates/__init__.py rename src/galax/{dynamics/_dynamics/base.py => coordinates/_base.py} (99%) rename src/galax/{dynamics/_dynamics/core.py => coordinates/_core.py} (98%) rename src/galax/{dynamics/_dynamics/utils.py => coordinates/_utils.py} (100%) create mode 100644 tests/smoke/coordinates/__init__.py create mode 100644 tests/smoke/coordinates/test_package.py diff --git a/docs/getting_started.rst b/docs/getting_started.rst index d77e6cbc..e23a5fb7 100644 --- a/docs/getting_started.rst +++ b/docs/getting_started.rst @@ -86,13 +86,13 @@ Now that we have a potential model, if we want to compute an orbit, we need to specify a set of initial conditions to initialize the numerical orbit integration. In :mod:`galax`, initial conditions and other positions in phase-space (locations in position and velocity space) are defined using the -:class:`~galax.dynamics.PhaseSpacePosition` class. This class allows a number of +:class:`~galax.coordinates.PhaseSpacePosition` class. This class allows a number of possible inputs, but one of the most common inputs are Cartesian position and velocity vectors. As an example orbit, we will use a position and velocity that is close to the Sun's Galactocentric position and velocity:: - >>> import galax.dynamics as gd - >>> psp = gd.PhaseSpacePosition(q=[-8.1, 0, 0.02] * u.kpc, + >>> import galax.coordinates as gc + >>> psp = gc.PhaseSpacePosition(q=[-8.1, 0, 0.02] * u.kpc, ... p=[13, 245, 8.] * u.km/u.s) By convention, I typically use the variable ``w`` to represent phase-space @@ -124,7 +124,7 @@ on any Potential object through the By default, this method uses Leapfrog integration , which is a fast, symplectic integration scheme. The returned object is an instance of the :class:`~galax.dynamics.Orbit` class, which is similar to the -:class:`~galax.dynamics.PhaseSpacePosition` but represents a collection of +:class:`~galax.coordinates.PhaseSpacePosition` but represents a collection of phase-space positions at times:: >>> orbit @@ -144,11 +144,12 @@ performing common tasks, like plotting an orbit:: import astropy.units as u import matplotlib.pyplot as plt import numpy as np + import galax.coordinates as gc import galax.dynamics as gd import galax.potential as gp mw = gp.MilkyWayPotential() - psp = gd.PhaseSpacePosition(pos=[-8.1, 0, 0.02] * u.kpc, + psp = gc.PhaseSpacePosition(pos=[-8.1, 0, 0.02] * u.kpc, vel=[13, 245, 8.] * u.km/u.s) orbit = mw.integrate_orbit(psp.w(), dt=1*u.Myr, t1=0, t2=2*u.Gyr) diff --git a/src/galax/coordinates/__init__.py b/src/galax/coordinates/__init__.py new file mode 100644 index 00000000..c64ce4bf --- /dev/null +++ b/src/galax/coordinates/__init__.py @@ -0,0 +1,11 @@ +""":mod:`galax.coordinates`.""" + +from . import _base, _core, _utils +from ._base import * +from ._core import * +from ._utils import * + +__all__: list[str] = [] +__all__ += _base.__all__ +__all__ += _core.__all__ +__all__ += _utils.__all__ diff --git a/src/galax/dynamics/_dynamics/base.py b/src/galax/coordinates/_base.py similarity index 99% rename from src/galax/dynamics/_dynamics/base.py rename to src/galax/coordinates/_base.py index 806f0b62..a132c91f 100644 --- a/src/galax/dynamics/_dynamics/base.py +++ b/src/galax/coordinates/_base.py @@ -17,7 +17,7 @@ from galax.units import UnitSystem from galax.utils._shape import atleast_batched -from .utils import getitem_time_index +from ._utils import getitem_time_index if TYPE_CHECKING: from typing import Self diff --git a/src/galax/dynamics/_dynamics/core.py b/src/galax/coordinates/_core.py similarity index 98% rename from src/galax/dynamics/_dynamics/core.py rename to src/galax/coordinates/_core.py index 974e48e2..dfb4ee9b 100644 --- a/src/galax/dynamics/_dynamics/core.py +++ b/src/galax/coordinates/_core.py @@ -18,7 +18,7 @@ from galax.utils._shape import batched_shape, expand_batch_dims from galax.utils.dataclasses import converter_float_array -from .base import AbstractPhaseSpacePosition +from ._base import AbstractPhaseSpacePosition if TYPE_CHECKING: from galax.potential._potential.base import AbstractPotentialBase diff --git a/src/galax/dynamics/_dynamics/utils.py b/src/galax/coordinates/_utils.py similarity index 100% rename from src/galax/dynamics/_dynamics/utils.py rename to src/galax/coordinates/_utils.py diff --git a/src/galax/dynamics/_dynamics/__init__.py b/src/galax/dynamics/_dynamics/__init__.py index bfc29444..ee8f00a5 100644 --- a/src/galax/dynamics/_dynamics/__init__.py +++ b/src/galax/dynamics/_dynamics/__init__.py @@ -1,17 +1,13 @@ """galax: Galactic Dynamix in Jax.""" -from . import base, core, mockstream, orbit -from .base import * -from .core import * +from . import mockstream, orbit from .mockstream import * from .orbit import * __all__: list[str] = ["mockstream"] -__all__ += base.__all__ -__all__ += core.__all__ __all__ += orbit.__all__ __all__ += mockstream.__all__ # Cleanup -del base, core, orbit +del orbit diff --git a/src/galax/dynamics/_dynamics/mockstream/core.py b/src/galax/dynamics/_dynamics/mockstream/core.py index 43bea5e0..40d6260c 100644 --- a/src/galax/dynamics/_dynamics/mockstream/core.py +++ b/src/galax/dynamics/_dynamics/mockstream/core.py @@ -10,8 +10,8 @@ import jax import jax.numpy as jnp -from galax.dynamics._dynamics.base import AbstractPhaseSpacePosition -from galax.dynamics._dynamics.utils import getitem_vectime_index +from galax.coordinates import AbstractPhaseSpacePosition +from galax.coordinates._utils import getitem_vectime_index from galax.typing import BatchFloatScalar, BroadBatchVec3, VecTime from galax.utils._shape import batched_shape from galax.utils.dataclasses import converter_float_array diff --git a/src/galax/dynamics/_dynamics/orbit.py b/src/galax/dynamics/_dynamics/orbit.py index 5ad86e50..a12b45a2 100644 --- a/src/galax/dynamics/_dynamics/orbit.py +++ b/src/galax/dynamics/_dynamics/orbit.py @@ -11,15 +11,14 @@ import jax.numpy as jnp from astropy.units import Quantity +from galax.coordinates import AbstractPhaseSpacePosition +from galax.coordinates._utils import getitem_vectime_index from galax.integrate import DiffraxIntegrator, Integrator from galax.potential._potential.base import AbstractPotentialBase from galax.typing import BatchFloatScalar, BatchVec6, BroadBatchVec3, VecTime from galax.utils._shape import batched_shape from galax.utils.dataclasses import converter_float_array -from .base import AbstractPhaseSpacePosition -from .utils import getitem_vectime_index - if TYPE_CHECKING: from typing import Self @@ -118,6 +117,7 @@ def energy( ############################################################################## +# TODO: enable setting the default integrator _default_integrator: Integrator = DiffraxIntegrator() diff --git a/tests/smoke/coordinates/__init__.py b/tests/smoke/coordinates/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/smoke/coordinates/test_package.py b/tests/smoke/coordinates/test_package.py new file mode 100644 index 00000000..9ada440a --- /dev/null +++ b/tests/smoke/coordinates/test_package.py @@ -0,0 +1,9 @@ +"""Testing :mod:`galax.dynamics` module.""" + +import galax.coordinates as gc +from galax.coordinates import _base, _core, _utils + + +def test_all() -> None: + """Test the `galax.potential` API.""" + assert set(gc.__all__) == {*_base.__all__, *_core.__all__, *_utils.__all__} diff --git a/tests/smoke/dynamics/test_package.py b/tests/smoke/dynamics/test_package.py index d107dea9..3c73b836 100644 --- a/tests/smoke/dynamics/test_package.py +++ b/tests/smoke/dynamics/test_package.py @@ -1,15 +1,9 @@ """Testing :mod:`galax.dynamics` module.""" import galax.dynamics as gd -from galax.dynamics._dynamics import base, core, mockstream, orbit +from galax.dynamics._dynamics import mockstream, orbit def test_all() -> None: """Test the `galax.potential` API.""" - assert set(gd.__all__) == { - "mockstream", - *base.__all__, - *core.__all__, - *orbit.__all__, - *mockstream.__all__, - } + assert set(gd.__all__) == {"mockstream", *orbit.__all__, *mockstream.__all__} diff --git a/tests/unit/dynamics/test_core.py b/tests/unit/dynamics/test_core.py index e0d05403..0db01ccf 100644 --- a/tests/unit/dynamics/test_core.py +++ b/tests/unit/dynamics/test_core.py @@ -7,7 +7,7 @@ import pytest from jax import random -from galax.dynamics import PhaseSpacePosition +from galax.coordinates import PhaseSpacePosition from galax.units import galactic Shape: TypeAlias = tuple[int, ...] From 77e6eabf9f1ff7a9bf5e6a5764cf191aca2b3e47 Mon Sep 17 00:00:00 2001 From: nstarman Date: Fri, 2 Feb 2024 09:52:14 -0500 Subject: [PATCH 2/2] fix: docs ref to PSP Signed-off-by: nstarman --- src/galax/coordinates/_base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/galax/coordinates/_base.py b/src/galax/coordinates/_base.py index a132c91f..d5036e0e 100644 --- a/src/galax/coordinates/_base.py +++ b/src/galax/coordinates/_base.py @@ -215,7 +215,7 @@ def angular_momentum(self) -> BatchVec3: >>> import numpy as np >>> import astropy.units as u - >>> from galax.dynamics import PhaseSpacePosition + >>> from galax.coordinates import PhaseSpacePosition We can compute the angular momentum of a single object