From d2d121c3f02eb817b9f8b8728fbd9fead31b2cee Mon Sep 17 00:00:00 2001 From: Nathaniel Starkman Date: Wed, 13 Dec 2023 12:21:44 -0500 Subject: [PATCH] rename galdynamix to galax (#37) * rename galdynamix to galax --- .copier-answers.yml | 8 ++-- .github/CONTRIBUTING.md | 2 +- LICENSE | 2 +- README.md | 22 +++++----- docs/conf.py | 8 ++-- docs/index.md | 2 +- notebooks/example.ipynb | 41 +++++++++++++++---- noxfile.py | 4 +- pyproject.toml | 20 ++++----- src/{galdynamix => galax}/__init__.py | 6 +-- src/{galdynamix => galax}/_version.pyi | 0 .../dynamics/__init__.py | 2 +- src/{galdynamix => galax}/dynamics/_core.py | 14 +++---- src/{galdynamix => galax}/dynamics/_orbit.py | 10 ++--- .../dynamics/mockstream/__init__.py | 2 +- .../dynamics/mockstream/_core.py | 12 +++--- .../dynamics/mockstream/_df/__init__.py | 2 +- .../dynamics/mockstream/_df/base.py | 12 +++--- .../dynamics/mockstream/_df/fardal.py | 16 ++++---- .../mockstream/_mockstream_generator.py | 18 ++++---- src/galax/integrate/__init__.py | 10 +++++ src/{galdynamix => galax}/integrate/_base.py | 2 +- .../integrate/_builtin.py | 6 +-- .../potential/__init__.py | 2 +- .../potential/_potential/__init__.py | 2 +- .../potential/_potential/base.py | 20 ++++----- .../potential/_potential/builtin.py | 12 +++--- .../potential/_potential/composite.py | 8 ++-- .../potential/_potential/core.py | 4 +- .../potential/_potential/param/__init__.py | 0 .../potential/_potential/param/core.py | 8 ++-- .../potential/_potential/param/field.py | 4 +- .../potential/_potential/subhalo.py | 12 +++--- .../potential/_potential/utils.py | 4 +- src/{galdynamix => galax}/py.typed | 0 src/{galdynamix => galax}/typing.py | 2 +- src/{galdynamix => galax}/units.py | 0 src/{galdynamix => galax}/utils/__init__.py | 2 +- .../utils/_collections.py | 4 +- src/{galdynamix => galax}/utils/_jax.py | 2 +- src/{galdynamix => galax}/utils/_misc.py | 2 +- src/{galdynamix => galax}/utils/_shape.py | 12 +++--- .../utils/dataclasses.py | 2 +- src/galdynamix/integrate/__init__.py | 10 ----- tests/potential/builtin/test_bar.py | 2 +- tests/potential/builtin/test_common.py | 6 +-- tests/potential/builtin/test_isochrone.py | 2 +- tests/potential/builtin/test_miyamotonagai.py | 4 +- tests/potential/builtin/test_nfw.py | 6 +-- tests/potential/params/test_core.py | 26 ++++++------ tests/potential/params/test_field.py | 4 +- tests/potential/test_base.py | 10 ++--- tests/potential/test_core.py | 12 +++--- tests/potential/test_package.py | 6 +-- tests/potential/test_utils.py | 6 +-- tests/test_package.py | 4 +- tests/utils/test_collections.py | 2 +- tests/utils/test_dataclasses.py | 2 +- tests/utils/test_jax.py | 2 +- tests/utils/test_misc.py | 2 +- tests/utils/test_package.py | 4 +- tests/utils/test_shape.py | 8 ++-- 62 files changed, 232 insertions(+), 209 deletions(-) rename src/{galdynamix => galax}/__init__.py (71%) rename src/{galdynamix => galax}/_version.pyi (100%) rename src/{galdynamix => galax}/dynamics/__init__.py (83%) rename src/{galdynamix => galax}/dynamics/_core.py (92%) rename src/{galdynamix => galax}/dynamics/_orbit.py (80%) rename src/{galdynamix => galax}/dynamics/mockstream/__init__.py (81%) rename src/{galdynamix => galax}/dynamics/mockstream/_core.py (80%) rename src/{galdynamix => galax}/dynamics/mockstream/_df/__init__.py (75%) rename src/{galdynamix => galax}/dynamics/mockstream/_df/base.py (90%) rename src/{galdynamix => galax}/dynamics/mockstream/_df/fardal.py (95%) rename src/{galdynamix => galax}/dynamics/mockstream/_mockstream_generator.py (93%) create mode 100644 src/galax/integrate/__init__.py rename src/{galdynamix => galax}/integrate/_base.py (96%) rename src/{galdynamix => galax}/integrate/_builtin.py (92%) rename src/{galdynamix => galax}/potential/__init__.py (71%) rename src/{galdynamix => galax}/potential/_potential/__init__.py (88%) rename src/{galdynamix => galax}/potential/_potential/base.py (94%) rename src/{galdynamix => galax}/potential/_potential/builtin.py (92%) rename src/{galdynamix => galax}/potential/_potential/composite.py (94%) rename src/{galdynamix => galax}/potential/_potential/core.py (87%) rename src/{galdynamix => galax}/potential/_potential/param/__init__.py (100%) rename src/{galdynamix => galax}/potential/_potential/param/core.py (93%) rename src/{galdynamix => galax}/potential/_potential/param/field.py (98%) rename src/{galdynamix => galax}/potential/_potential/subhalo.py (87%) rename src/{galdynamix => galax}/potential/_potential/utils.py (88%) rename src/{galdynamix => galax}/py.typed (100%) rename src/{galdynamix => galax}/typing.py (98%) rename src/{galdynamix => galax}/units.py (100%) rename src/{galdynamix => galax}/utils/__init__.py (81%) rename src/{galdynamix => galax}/utils/_collections.py (96%) rename src/{galdynamix => galax}/utils/_jax.py (97%) rename src/{galdynamix => galax}/utils/_misc.py (83%) rename src/{galdynamix => galax}/utils/_shape.py (94%) rename src/{galdynamix => galax}/utils/dataclasses.py (99%) delete mode 100644 src/galdynamix/integrate/__init__.py diff --git a/.copier-answers.yml b/.copier-answers.yml index 9461176d..e7a33ce3 100644 --- a/.copier-answers.yml +++ b/.copier-answers.yml @@ -3,10 +3,10 @@ _commit: 2023.10.27 _src_path: gh:scientific-python/cookie backend: hatch email: nstarman@users.noreply.github.com -full_name: galdynamix maintainers +full_name: galax maintainers license: MIT -org: galdynamix -project_name: galdynamix +org: galax +project_name: galax project_short_description: Galactic Dynamix in Jax. -url: https://github.com/galdynamix/galdynamix +url: https://github.com/GalacticDynamics/galax vcs: true diff --git a/.github/CONTRIBUTING.md b/.github/CONTRIBUTING.md index 73de9c63..5a6162b2 100644 --- a/.github/CONTRIBUTING.md +++ b/.github/CONTRIBUTING.md @@ -71,7 +71,7 @@ pytest Use pytest-cov to generate coverage reports: ```bash -pytest --cov=galdynamix +pytest --cov=galax ``` # Building docs diff --git a/LICENSE b/LICENSE index 68e13639..90ca9c1c 100644 --- a/LICENSE +++ b/LICENSE @@ -1,4 +1,4 @@ -Copyright 2023 galdynamix maintainers +Copyright 2023 galax maintainers Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in diff --git a/README.md b/README.md index 474d342e..17163ecb 100644 --- a/README.md +++ b/README.md @@ -1,4 +1,4 @@ -# galdynamix +# galax [![Actions Status][actions-badge]][actions-link] [![Documentation Status][rtd-badge]][rtd-link] @@ -12,16 +12,16 @@ -[actions-badge]: https://github.com/galdynamix/galdynamix/workflows/CI/badge.svg -[actions-link]: https://github.com/galdynamix/galdynamix/actions -[conda-badge]: https://img.shields.io/conda/vn/conda-forge/galdynamix -[conda-link]: https://github.com/conda-forge/galdynamix-feedstock +[actions-badge]: https://github.com/GalacticDynamics/galax/workflows/CI/badge.svg +[actions-link]: https://github.com/GalacticDynamics/galax/actions +[conda-badge]: https://img.shields.io/conda/vn/conda-forge/galax +[conda-link]: https://github.com/conda-forge/galax-feedstock [github-discussions-badge]: https://img.shields.io/static/v1?label=Discussions&message=Ask&color=blue&logo=github -[github-discussions-link]: https://github.com/galdynamix/galdynamix/discussions -[pypi-link]: https://pypi.org/project/galdynamix/ -[pypi-platforms]: https://img.shields.io/pypi/pyversions/galdynamix -[pypi-version]: https://img.shields.io/pypi/v/galdynamix -[rtd-badge]: https://readthedocs.org/projects/galdynamix/badge/?version=latest -[rtd-link]: https://galdynamix.readthedocs.io/en/latest/?badge=latest +[github-discussions-link]: https://github.com/GalacticDynamics/galax/discussions +[pypi-link]: https://pypi.org/project/galax/ +[pypi-platforms]: https://img.shields.io/pypi/pyversions/galax +[pypi-version]: https://img.shields.io/pypi/v/galax +[rtd-badge]: https://readthedocs.org/projects/galax/badge/?version=latest +[rtd-link]: https://galax.readthedocs.io/en/latest/?badge=latest diff --git a/docs/conf.py b/docs/conf.py index ae68c78c..cf688468 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -2,10 +2,10 @@ import importlib.metadata -project = "galdynamix" -copyright = "2023, galdynamix maintainers" -author = "galdynamix maintainers" -version = release = importlib.metadata.version("galdynamix") +project = "galax" +copyright = "2023, galax maintainers" +author = "galax maintainers" +version = release = importlib.metadata.version("galax") extensions = [ "myst_parser", diff --git a/docs/index.md b/docs/index.md index 4be0aa88..c4ea60af 100644 --- a/docs/index.md +++ b/docs/index.md @@ -1,4 +1,4 @@ -# galdynamix +# galax ```{toctree} :maxdepth: 2 diff --git a/notebooks/example.ipynb b/notebooks/example.ipynb index a643746d..20a1981d 100644 --- a/notebooks/example.ipynb +++ b/notebooks/example.ipynb @@ -23,15 +23,16 @@ "\n", "import numpy as np\n", "\n", - "import galdynamix.potential as gpx\n", - "import galdynamix.dynamics as gdx\n", - "from galdynamix.units import UnitSystem\n", + "import galax.potential as gpx\n", + "import galax.dynamics as gdx\n", + "from galax.units import UnitSystem\n", "\n", "config.update(\"jax_enable_x64\", True)" ] }, { "cell_type": "markdown", + "id": "6a306379", "metadata": {}, "source": [ "## Potentials" @@ -40,6 +41,7 @@ { "cell_type": "code", "execution_count": 2, + "id": "480c8d57", "metadata": {}, "outputs": [], "source": [ @@ -61,7 +63,7 @@ " \"m_NFW\": 1.0e12,\n", " \"r_s_NFW\": 15.0,\n", "}\n", - "pot_disk = gpx.MiyamotoNagaiPotential(\n", + "pot_disk = gpx.MiyamotoNagaiDisk(\n", " m=params_global[\"m_disk\"],\n", " a=params_global[\"a_disk\"],\n", " b=params_global[\"b_disk\"],\n", @@ -83,6 +85,7 @@ { "cell_type": "code", "execution_count": 4, + "id": "70797f97", "metadata": {}, "outputs": [], "source": [ @@ -106,6 +109,7 @@ }, { "cell_type": "markdown", + "id": "4459de91", "metadata": {}, "source": [ "## Integrating Orbits" @@ -139,6 +143,7 @@ { "cell_type": "code", "execution_count": 6, + "id": "a241d227", "metadata": {}, "outputs": [], "source": [ @@ -148,6 +153,7 @@ }, { "cell_type": "markdown", + "id": "fc3a23b2", "metadata": {}, "source": [ "And with Gala" @@ -156,6 +162,7 @@ { "cell_type": "code", "execution_count": 7, + "id": "c0a8cc69", "metadata": {}, "outputs": [], "source": [ @@ -215,6 +222,7 @@ }, { "cell_type": "markdown", + "id": "ae730d05", "metadata": {}, "source": [ "## Mock Streams" @@ -282,6 +290,7 @@ }, { "cell_type": "markdown", + "id": "e2597d53", "metadata": {}, "source": [ "Now in gala" @@ -316,6 +325,7 @@ }, { "cell_type": "markdown", + "id": "36345058", "metadata": {}, "source": [ "Comparison plot" @@ -409,6 +419,7 @@ }, { "cell_type": "markdown", + "id": "d5826d57", "metadata": {}, "source": [ "## Demonstrating Gradients" @@ -417,6 +428,7 @@ { "cell_type": "code", "execution_count": 14, + "id": "9a884254", "metadata": {}, "outputs": [], "source": [ @@ -426,6 +438,7 @@ { "cell_type": "code", "execution_count": 15, + "id": "42110e83", "metadata": {}, "outputs": [], "source": [ @@ -436,10 +449,11 @@ { "cell_type": "code", "execution_count": 16, + "id": "3a956392", "metadata": {}, "outputs": [], "source": [ - "pot_disk = gpx.MiyamotoNagaiPotential(\n", + "pot_disk = gpx.MiyamotoNagaiDisk(\n", " m=params_global[\"m_disk\"],\n", " a=params_global[\"a_disk\"],\n", " b=params_global[\"b_disk\"],\n", @@ -462,12 +476,13 @@ { "cell_type": "code", "execution_count": 17, + "id": "8cb87962", "metadata": {}, "outputs": [], "source": [ "@jax.jit\n", "def gen_observed_stream(params, ts, w0, M_sat):\n", - " pot_disk = gpx.MiyamotoNagaiPotential(\n", + " pot_disk = gpx.MiyamotoNagaiDisk(\n", " m=params[\"m_disk\"], a=params[\"a_disk\"], b=params[\"b_disk\"], units=usys\n", " )\n", " pot_NFW = gpx.NFWPotential(m=params[\"m_NFW\"], r_s=params[\"r_s_NFW\"], units=usys)\n", @@ -487,9 +502,10 @@ "def compare_streams(params, ts, w0, M_sat):\n", " ## first generate data stream (no grad tracking)\n", " lead_arm, trail_arm = gen_observed_stream(params, ts, w0, M_sat)\n", - " lead_arm_obs, trail_arm_obs = jax.lax.stop_gradient(\n", - " lead_arm\n", - " ), jax.lax.stop_gradient(trail_arm)\n", + " lead_arm_obs, trail_arm_obs = (\n", + " jax.lax.stop_gradient(lead_arm),\n", + " jax.lax.stop_gradient(trail_arm),\n", + " )\n", " return -jnp.sum((lead_arm - lead_arm_obs) ** 2 + (trail_arm - trail_arm_obs) ** 2)\n", "\n", "\n", @@ -500,6 +516,7 @@ { "cell_type": "code", "execution_count": 18, + "id": "a9e65c06", "metadata": {}, "outputs": [], "source": [ @@ -511,6 +528,7 @@ { "cell_type": "code", "execution_count": 19, + "id": "27123f52", "metadata": {}, "outputs": [], "source": [ @@ -520,6 +538,7 @@ { "cell_type": "code", "execution_count": 20, + "id": "227aa50e", "metadata": {}, "outputs": [], "source": [ @@ -529,6 +548,7 @@ { "cell_type": "code", "execution_count": 21, + "id": "b4e0dfda", "metadata": {}, "outputs": [ { @@ -553,6 +573,7 @@ { "cell_type": "code", "execution_count": 22, + "id": "cf5f061a", "metadata": {}, "outputs": [], "source": [ @@ -564,6 +585,7 @@ { "cell_type": "code", "execution_count": 23, + "id": "5ae7c3bf", "metadata": {}, "outputs": [ { @@ -607,6 +629,7 @@ }, { "cell_type": "markdown", + "id": "c1579395", "metadata": {}, "source": [ "# Benchmark" diff --git a/noxfile.py b/noxfile.py index 2f704029..227ff1e0 100644 --- a/noxfile.py +++ b/noxfile.py @@ -28,7 +28,7 @@ def lint(session: nox.Session) -> None: # # This needs to be installed into the package environment, and is slower # # than a pre-commit check # session.install(".", "pylint") -# session.run("pylint", "galdynamix", *session.posargs) +# session.run("pylint", "galax", *session.posargs) @nox.session @@ -89,7 +89,7 @@ def build_api_docs(session: nox.Session) -> None: "--module-first", "--no-toc", "--force", - "../src/galdynamix", + "../src/galax", ) diff --git a/pyproject.toml b/pyproject.toml index 2a828084..71653688 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,9 +4,9 @@ build-backend = "hatchling.build" [project] -name = "galdynamix" +name = "galax" authors = [ - { name = "galdynamix maintainers", email = "nstarman@users.noreply.github.com" }, + { name = "galax maintainers", email = "nstarman@users.noreply.github.com" }, { name = "Jake Nibauer", email = "jnibauer@princeton.edu" }, { name = "Nathaniel Starkman", email = "nstarman@users.noreply.github.com" }, ] @@ -58,15 +58,15 @@ docs = [ ] [project.urls] -Homepage = "https://github.com/galdynamix/galdynamix" -"Bug Tracker" = "https://github.com/galdynamix/galdynamix/issues" -Discussions = "https://github.com/galdynamix/galdynamix/discussions" -Changelog = "https://github.com/galdynamix/galdynamix/releases" +Homepage = "https://github.com/GalacticDynamics/galax" +"Bug Tracker" = "https://github.com/GalacticDynamics/galax/issues" +Discussions = "https://github.com/GalacticDynamics/galax/discussions" +Changelog = "https://github.com/GalacticDynamics/galax/releases" [tool.hatch] version.source = "vcs" -build.hooks.vcs.version-file = "src/galdynamix/_version.py" +build.hooks.vcs.version-file = "src/galax/_version.py" envs.default.dependencies = [ "pytest", "pytest-cov", @@ -88,7 +88,7 @@ testpaths = [ [tool.coverage] -run.source = ["galdynamix"] +run.source = ["galax"] port.exclude_lines = [ 'pragma: no cover', '\.\.\.', @@ -107,7 +107,7 @@ disallow_untyped_defs = false disallow_incomplete_defs = false [[tool.mypy.overrides]] - module = "galdynamix.*" + module = "galax.*" disallow_untyped_defs = true disallow_incomplete_defs = true disable_error_code = ["name-defined"] # <- jaxtyping @@ -163,7 +163,7 @@ ignore = [ "F841", ] # Uncomment if using a _compat.typing backport -# typing-modules = ["galdynamix._compat.typing"] +# typing-modules = ["galax._compat.typing"] [tool.ruff.lint.per-file-ignores] "__init__.py" = ["F403"] diff --git a/src/galdynamix/__init__.py b/src/galax/__init__.py similarity index 71% rename from src/galdynamix/__init__.py rename to src/galax/__init__.py index 36555e01..4aa69086 100644 --- a/src/galdynamix/__init__.py +++ b/src/galax/__init__.py @@ -1,4 +1,4 @@ -"""Copyright (c) 2023 galdynamix maintainers. All rights reserved.""" +"""Copyright (c) 2023 galax maintainers. All rights reserved.""" # ruff:noqa: F401 __all__ = [ @@ -27,5 +27,5 @@ else: typechecker = None -with install_import_hook("galdynamix", typechecker): - from galdynamix import dynamics, integrate, potential, typing, units, utils +with install_import_hook("galax", typechecker): + from galax import dynamics, integrate, potential, typing, units, utils diff --git a/src/galdynamix/_version.pyi b/src/galax/_version.pyi similarity index 100% rename from src/galdynamix/_version.pyi rename to src/galax/_version.pyi diff --git a/src/galdynamix/dynamics/__init__.py b/src/galax/dynamics/__init__.py similarity index 83% rename from src/galdynamix/dynamics/__init__.py rename to src/galax/dynamics/__init__.py index 3a1fe489..81291305 100644 --- a/src/galdynamix/dynamics/__init__.py +++ b/src/galax/dynamics/__init__.py @@ -1,4 +1,4 @@ -"""galdynamix: Galactic Dynamix in Jax.""" +"""galax: Galactic Dynamix in Jax.""" from . import _core, _orbit, mockstream from ._core import * diff --git a/src/galdynamix/dynamics/_core.py b/src/galax/dynamics/_core.py similarity index 92% rename from src/galdynamix/dynamics/_core.py rename to src/galax/dynamics/_core.py index 68a93bf7..d18237b9 100644 --- a/src/galdynamix/dynamics/_core.py +++ b/src/galax/dynamics/_core.py @@ -1,4 +1,4 @@ -"""galdynamix: Galactic Dynamix in Jax.""" +"""galax: Galactic Dynamix in Jax.""" __all__ = ["AbstractPhaseSpacePosition", "PhaseSpacePosition"] @@ -8,13 +8,13 @@ import equinox as eqx import jax.numpy as xp -from galdynamix.typing import BatchFloatScalar, BatchVec3, BatchVec6, BatchVec7 -from galdynamix.utils import partial_jit -from galdynamix.utils._shape import atleast_batched, batched_shape -from galdynamix.utils.dataclasses import converter_float_array +from galax.typing import BatchFloatScalar, BatchVec3, BatchVec6, BatchVec7 +from galax.utils import partial_jit +from galax.utils._shape import atleast_batched, batched_shape +from galax.utils.dataclasses import converter_float_array if TYPE_CHECKING: - from galdynamix.potential._potential.base import AbstractPotentialBase + from galax.potential._potential.base import AbstractPotentialBase class AbstractPhaseSpacePositionBase(eqx.Module): # type: ignore[misc] @@ -153,7 +153,7 @@ def potential_energy( Parameters ---------- - potential : `galdynamix.potential.AbstractPotentialBase` + potential : `galax.potential.AbstractPotentialBase` The potential object to compute the energy from. Returns diff --git a/src/galdynamix/dynamics/_orbit.py b/src/galax/dynamics/_orbit.py similarity index 80% rename from src/galdynamix/dynamics/_orbit.py rename to src/galax/dynamics/_orbit.py index d1a420c6..c402efe2 100644 --- a/src/galdynamix/dynamics/_orbit.py +++ b/src/galax/dynamics/_orbit.py @@ -1,12 +1,12 @@ -"""galdynamix: Galactic Dynamix in Jax.""" +"""galax: Galactic Dynamix in Jax.""" __all__ = ["Orbit"] from typing_extensions import override -from galdynamix.potential._potential.base import AbstractPotentialBase -from galdynamix.typing import BatchFloatScalar -from galdynamix.utils._jax import partial_jit +from galax.potential._potential.base import AbstractPotentialBase +from galax.typing import BatchFloatScalar +from galax.utils._jax import partial_jit from ._core import AbstractPhaseSpacePosition @@ -38,7 +38,7 @@ def potential_energy( Parameters ---------- - potential : `galdynamix.potential.AbstractPotentialBase` + potential : `galax.potential.AbstractPotentialBase` The potential object to compute the energy from. Returns diff --git a/src/galdynamix/dynamics/mockstream/__init__.py b/src/galax/dynamics/mockstream/__init__.py similarity index 81% rename from src/galdynamix/dynamics/mockstream/__init__.py rename to src/galax/dynamics/mockstream/__init__.py index 0ddd57c8..9dcfc958 100644 --- a/src/galdynamix/dynamics/mockstream/__init__.py +++ b/src/galax/dynamics/mockstream/__init__.py @@ -1,4 +1,4 @@ -"""galdynamix: Galactic Dynamix in Jax.""" +"""galax: Galactic Dynamix in Jax.""" from . import _df, _mockstream_generator from ._df import * diff --git a/src/galdynamix/dynamics/mockstream/_core.py b/src/galax/dynamics/mockstream/_core.py similarity index 80% rename from src/galdynamix/dynamics/mockstream/_core.py rename to src/galax/dynamics/mockstream/_core.py index c5278f91..0238f413 100644 --- a/src/galdynamix/dynamics/mockstream/_core.py +++ b/src/galax/dynamics/mockstream/_core.py @@ -1,15 +1,15 @@ -"""galdynamix: Galactic Dynamix in Jax.""" +"""galax: Galactic Dynamix in Jax.""" __all__ = ["MockStream"] import equinox as eqx import jax.numpy as xp -from galdynamix.dynamics._core import AbstractPhaseSpacePositionBase -from galdynamix.typing import BatchFloatScalar, BatchVec7 -from galdynamix.utils import partial_jit -from galdynamix.utils._shape import atleast_batched, batched_shape -from galdynamix.utils.dataclasses import converter_float_array +from galax.dynamics._core import AbstractPhaseSpacePositionBase +from galax.typing import BatchFloatScalar, BatchVec7 +from galax.utils import partial_jit +from galax.utils._shape import atleast_batched, batched_shape +from galax.utils.dataclasses import converter_float_array class MockStream(AbstractPhaseSpacePositionBase): diff --git a/src/galdynamix/dynamics/mockstream/_df/__init__.py b/src/galax/dynamics/mockstream/_df/__init__.py similarity index 75% rename from src/galdynamix/dynamics/mockstream/_df/__init__.py rename to src/galax/dynamics/mockstream/_df/__init__.py index 96103f7a..37b4e11d 100644 --- a/src/galdynamix/dynamics/mockstream/_df/__init__.py +++ b/src/galax/dynamics/mockstream/_df/__init__.py @@ -1,4 +1,4 @@ -"""galdynamix: Galactic Dynamix in Jax.""" +"""galax: Galactic Dynamix in Jax.""" from . import base, fardal diff --git a/src/galdynamix/dynamics/mockstream/_df/base.py b/src/galax/dynamics/mockstream/_df/base.py similarity index 90% rename from src/galdynamix/dynamics/mockstream/_df/base.py rename to src/galax/dynamics/mockstream/_df/base.py index 7be7d427..a4580ecb 100644 --- a/src/galdynamix/dynamics/mockstream/_df/base.py +++ b/src/galax/dynamics/mockstream/_df/base.py @@ -1,4 +1,4 @@ -"""galdynamix: Galactic Dynamix in Jax.""" +"""galax: Galactic Dynamix in Jax.""" __all__ = ["AbstractStreamDF"] @@ -10,11 +10,11 @@ import jax import jax.numpy as xp -from galdynamix.dynamics._orbit import Orbit -from galdynamix.dynamics.mockstream._core import MockStream -from galdynamix.potential._potential.base import AbstractPotentialBase -from galdynamix.typing import BatchVec3, FloatScalar, IntLike, Vec3, Vec6 -from galdynamix.utils import partial_jit +from galax.dynamics._orbit import Orbit +from galax.dynamics.mockstream._core import MockStream +from galax.potential._potential.base import AbstractPotentialBase +from galax.typing import BatchVec3, FloatScalar, IntLike, Vec3, Vec6 +from galax.utils import partial_jit Wif: TypeAlias = tuple[Vec3, Vec3, Vec3, Vec3] Carry: TypeAlias = tuple[IntLike, Vec3, Vec3, Vec3, Vec3] diff --git a/src/galdynamix/dynamics/mockstream/_df/fardal.py b/src/galax/dynamics/mockstream/_df/fardal.py similarity index 95% rename from src/galdynamix/dynamics/mockstream/_df/fardal.py rename to src/galax/dynamics/mockstream/_df/fardal.py index d01a0479..089a3a69 100644 --- a/src/galdynamix/dynamics/mockstream/_df/fardal.py +++ b/src/galax/dynamics/mockstream/_df/fardal.py @@ -1,4 +1,4 @@ -"""galdynamix: Galactic Dynamix in Jax.""" +"""galax: Galactic Dynamix in Jax.""" __all__ = ["FardalStreamDF"] @@ -7,15 +7,15 @@ import jax import jax.numpy as xp -from galdynamix.potential._potential.base import AbstractPotentialBase -from galdynamix.typing import ( +from galax.potential._potential.base import AbstractPotentialBase +from galax.typing import ( FloatOrIntScalarLike, FloatScalar, IntLike, Vec3, Vec6, ) -from galdynamix.utils import partial_jit +from galax.utils import partial_jit from .base import AbstractStreamDF @@ -155,8 +155,8 @@ def d2phidr2( Examples -------- - >>> from galdynamix.potential import NFWPotential - >>> from galdynamix.units import galactic + >>> from galax.potential import NFWPotential + >>> from galax.units import galactic >>> pot = NFWPotential(m=1e12, r_s=20.0, units=galactic) >>> d2phidr2(pot, xp.array([8.0, 0.0, 0.0]), t=0) Array(-0.00017469, dtype=float64) @@ -250,8 +250,8 @@ def tidal_radius( Examples -------- - >>> from galdynamix.potential import NFWPotential - >>> from galdynamix.units import galactic + >>> from galax.potential import NFWPotential + >>> from galax.units import galactic >>> pot = NFWPotential(m=1e12, r_s=20.0, units=galactic) >>> x=xp.array([8.0, 0.0, 0.0]) >>> v=xp.array([8.0, 0.0, 0.0]) diff --git a/src/galdynamix/dynamics/mockstream/_mockstream_generator.py b/src/galax/dynamics/mockstream/_mockstream_generator.py similarity index 93% rename from src/galdynamix/dynamics/mockstream/_mockstream_generator.py rename to src/galax/dynamics/mockstream/_mockstream_generator.py index 67a705a8..dae333a7 100644 --- a/src/galdynamix/dynamics/mockstream/_mockstream_generator.py +++ b/src/galax/dynamics/mockstream/_mockstream_generator.py @@ -1,4 +1,4 @@ -"""galdynamix: Galactic Dynamix in Jax.""" +"""galax: Galactic Dynamix in Jax.""" __all__ = ["MockStreamGenerator"] @@ -11,19 +11,19 @@ import jax.numpy as xp from jaxtyping import Float -from galdynamix.dynamics._orbit import Orbit -from galdynamix.integrate._base import AbstractIntegrator -from galdynamix.integrate._builtin import DiffraxIntegrator -from galdynamix.potential._potential.base import AbstractPotentialBase -from galdynamix.typing import ( +from galax.dynamics._orbit import Orbit +from galax.integrate._base import AbstractIntegrator +from galax.integrate._builtin import DiffraxIntegrator +from galax.potential._potential.base import AbstractPotentialBase +from galax.typing import ( FloatScalar, IntScalar, TimeVector, Vec6, VecN, ) -from galdynamix.utils import partial_jit -from galdynamix.utils._collections import ImmutableDict +from galax.utils import partial_jit +from galax.utils._collections import ImmutableDict from ._df import AbstractStreamDF @@ -38,7 +38,7 @@ class MockStreamGenerator(eqx.Module): # type: ignore[misc] df: AbstractStreamDF """Distribution function for generating mock streams. - E.g. ``galdynamix.dynamics.mockstream.FardalStreamDF``. + E.g. ``galax.dynamics.mockstream.FardalStreamDF``. """ potential: AbstractPotentialBase diff --git a/src/galax/integrate/__init__.py b/src/galax/integrate/__init__.py new file mode 100644 index 00000000..eede9d3f --- /dev/null +++ b/src/galax/integrate/__init__.py @@ -0,0 +1,10 @@ +"""galax: Galactic Dynamix in Jax.""" + + +from galax.integrate import _base, _builtin +from galax.integrate._base import * +from galax.integrate._builtin import * + +__all__: list[str] = [] +__all__ += _base.__all__ +__all__ += _builtin.__all__ diff --git a/src/galdynamix/integrate/_base.py b/src/galax/integrate/_base.py similarity index 96% rename from src/galdynamix/integrate/_base.py rename to src/galax/integrate/_base.py index e16e59e7..263a79e3 100644 --- a/src/galdynamix/integrate/_base.py +++ b/src/galax/integrate/_base.py @@ -6,7 +6,7 @@ import equinox as eqx from jaxtyping import Array, Float -from galdynamix.typing import FloatScalar, Vec6, Vec7 +from galax.typing import FloatScalar, Vec6, Vec7 @runtime_checkable diff --git a/src/galdynamix/integrate/_builtin.py b/src/galax/integrate/_builtin.py similarity index 92% rename from src/galdynamix/integrate/_builtin.py rename to src/galax/integrate/_builtin.py index 0def152f..2b183832 100644 --- a/src/galdynamix/integrate/_builtin.py +++ b/src/galax/integrate/_builtin.py @@ -17,9 +17,9 @@ from diffrax import SaveAt as DiffraxSaveAt from jaxtyping import Array, Float -from galdynamix.integrate._base import AbstractIntegrator -from galdynamix.typing import FloatScalar, Vec6 -from galdynamix.utils import ImmutableDict +from galax.integrate._base import AbstractIntegrator +from galax.typing import FloatScalar, Vec6 +from galax.utils import ImmutableDict class DiffraxIntegrator(AbstractIntegrator): diff --git a/src/galdynamix/potential/__init__.py b/src/galax/potential/__init__.py similarity index 71% rename from src/galdynamix/potential/__init__.py rename to src/galax/potential/__init__.py index 361bffa2..aaebbfd1 100644 --- a/src/galdynamix/potential/__init__.py +++ b/src/galax/potential/__init__.py @@ -1,4 +1,4 @@ -"""galdynamix: Galactic Dynamix in Jax.""" +"""galax: Galactic Dynamix in Jax.""" from . import _potential diff --git a/src/galdynamix/potential/_potential/__init__.py b/src/galax/potential/_potential/__init__.py similarity index 88% rename from src/galdynamix/potential/_potential/__init__.py rename to src/galax/potential/_potential/__init__.py index 64495d01..d2605106 100644 --- a/src/galdynamix/potential/_potential/__init__.py +++ b/src/galax/potential/_potential/__init__.py @@ -1,4 +1,4 @@ -"""galdynamix: Galactic Dynamix in Jax.""" +"""galax: Galactic Dynamix in Jax.""" from . import base, builtin, composite, core, param diff --git a/src/galdynamix/potential/_potential/base.py b/src/galax/potential/_potential/base.py similarity index 94% rename from src/galdynamix/potential/_potential/base.py rename to src/galax/potential/_potential/base.py index 54d5655c..601debeb 100644 --- a/src/galdynamix/potential/_potential/base.py +++ b/src/galax/potential/_potential/base.py @@ -12,9 +12,9 @@ from jax import grad, hessian, jacfwd from jaxtyping import Array, Float -from galdynamix.integrate._base import AbstractIntegrator -from galdynamix.integrate._builtin import DiffraxIntegrator -from galdynamix.typing import ( +from galax.integrate._base import AbstractIntegrator +from galax.integrate._builtin import DiffraxIntegrator +from galax.typing import ( BatchableFloatOrIntScalarLike, BatchFloatScalar, BatchMatrix33, @@ -25,13 +25,13 @@ Vec3, Vec6, ) -from galdynamix.units import UnitSystem, dimensionless -from galdynamix.utils import partial_jit, vectorize_method -from galdynamix.utils._shape import batched_shape, expand_arr_dims, expand_batch_dims -from galdynamix.utils.dataclasses import ModuleMeta +from galax.units import UnitSystem, dimensionless +from galax.utils import partial_jit, vectorize_method +from galax.utils._shape import batched_shape, expand_arr_dims, expand_batch_dims +from galax.utils.dataclasses import ModuleMeta if TYPE_CHECKING: - from galdynamix.dynamics._orbit import Orbit + from galax.dynamics._orbit import Orbit class AbstractPotentialBase(eqx.Module, metaclass=ModuleMeta): # type: ignore[misc] @@ -62,7 +62,7 @@ def _init_units(self) -> None: G = 1 if self.units == dimensionless else _G.decompose(self.units).value object.__setattr__(self, "_G", G) - from galdynamix.potential._potential.param.field import ParameterField + from galax.potential._potential.param.field import ParameterField # Handle unit conversion for all fields, e.g. the parameters. for f in fields(self): @@ -302,7 +302,7 @@ def integrate_orbit( Integrator: type[AbstractIntegrator] | None = None, integrator_kw: Mapping[str, Any] | None = None, ) -> "Orbit": - from galdynamix.dynamics._orbit import Orbit + from galax.dynamics._orbit import Orbit integrator_cls = Integrator if Integrator is not None else DiffraxIntegrator diff --git a/src/galdynamix/potential/_potential/builtin.py b/src/galax/potential/_potential/builtin.py similarity index 92% rename from src/galdynamix/potential/_potential/builtin.py rename to src/galax/potential/_potential/builtin.py index 2d313801..bb02ecc4 100644 --- a/src/galdynamix/potential/_potential/builtin.py +++ b/src/galax/potential/_potential/builtin.py @@ -1,4 +1,4 @@ -"""galdynamix: Galactic Dynamix in Jax.""" +"""galax: Galactic Dynamix in Jax.""" __all__ = [ "MiyamotoNagaiPotential", @@ -12,9 +12,9 @@ import astropy.units as u import jax.numpy as xp -from galdynamix.potential._potential.core import AbstractPotential -from galdynamix.potential._potential.param import AbstractParameter, ParameterField -from galdynamix.typing import ( +from galax.potential._potential.core import AbstractPotential +from galax.potential._potential.param import AbstractParameter, ParameterField +from galax.typing import ( BatchableFloatOrIntScalarLike, BatchFloatScalar, BatchVec3, @@ -23,8 +23,8 @@ FloatScalar, Vec3, ) -from galdynamix.utils import partial_jit, vectorize_method -from galdynamix.utils.dataclasses import field +from galax.utils import partial_jit, vectorize_method +from galax.utils.dataclasses import field mass = u.get_physical_type("mass") length = u.get_physical_type("length") diff --git a/src/galdynamix/potential/_potential/composite.py b/src/galax/potential/_potential/composite.py similarity index 94% rename from src/galdynamix/potential/_potential/composite.py rename to src/galax/potential/_potential/composite.py index 50ed106b..bef09cf0 100644 --- a/src/galdynamix/potential/_potential/composite.py +++ b/src/galax/potential/_potential/composite.py @@ -9,14 +9,14 @@ import jax.numpy as xp from typing_extensions import override -from galdynamix.typing import ( +from galax.typing import ( BatchableFloatOrIntScalarLike, BatchFloatScalar, BatchVec3, ) -from galdynamix.units import UnitSystem -from galdynamix.utils import ImmutableDict, partial_jit -from galdynamix.utils._misc import first +from galax.units import UnitSystem +from galax.utils import ImmutableDict, partial_jit +from galax.utils._misc import first from .base import AbstractPotentialBase from .utils import converter_to_usys diff --git a/src/galdynamix/potential/_potential/core.py b/src/galax/potential/_potential/core.py similarity index 87% rename from src/galdynamix/potential/_potential/core.py rename to src/galax/potential/_potential/core.py index 4c89e3e5..62e0ec40 100644 --- a/src/galdynamix/potential/_potential/core.py +++ b/src/galax/potential/_potential/core.py @@ -6,7 +6,7 @@ import equinox as eqx -from galdynamix.units import UnitSystem +from galax.units import UnitSystem from .base import AbstractPotentialBase from .composite import CompositePotential @@ -25,7 +25,7 @@ def __add__(self, other: Any) -> CompositePotential: if not isinstance(other, AbstractPotentialBase): return NotImplemented - from galdynamix.potential._potential.composite import CompositePotential + from galax.potential._potential.composite import CompositePotential if isinstance(other, CompositePotential): return other.__ror__(self) diff --git a/src/galdynamix/potential/_potential/param/__init__.py b/src/galax/potential/_potential/param/__init__.py similarity index 100% rename from src/galdynamix/potential/_potential/param/__init__.py rename to src/galax/potential/_potential/param/__init__.py diff --git a/src/galdynamix/potential/_potential/param/core.py b/src/galax/potential/_potential/param/core.py similarity index 93% rename from src/galdynamix/potential/_potential/param/core.py rename to src/galax/potential/_potential/param/core.py index 80b44dbe..0887f79f 100644 --- a/src/galdynamix/potential/_potential/param/core.py +++ b/src/galax/potential/_potential/param/core.py @@ -9,15 +9,15 @@ import astropy.units as u import equinox as eqx -from galdynamix.typing import ( +from galax.typing import ( BatchableFloatOrIntScalarLike, FloatArrayAnyShape, FloatOrIntScalar, FloatScalar, Unit, ) -from galdynamix.utils import partial_jit, vectorize_method -from galdynamix.utils.dataclasses import converter_float_array +from galax.utils import partial_jit, vectorize_method +from galax.utils.dataclasses import converter_float_array class AbstractParameter(eqx.Module): # type: ignore[misc] @@ -78,7 +78,7 @@ def __call__( ---------- t : Array, optional This is ignored and is thus optional. - Note that for most :class:`~galdynamix.potential.AbstractParameter` + Note that for most :class:`~galax.potential.AbstractParameter` the time is required. **kwargs : Any This is ignored. diff --git a/src/galdynamix/potential/_potential/param/field.py b/src/galax/potential/_potential/param/field.py similarity index 98% rename from src/galdynamix/potential/_potential/param/field.py rename to src/galax/potential/_potential/param/field.py index 560def80..4be9c431 100644 --- a/src/galdynamix/potential/_potential/param/field.py +++ b/src/galax/potential/_potential/param/field.py @@ -6,8 +6,8 @@ import astropy.units as u import jax.numpy as xp -from galdynamix.potential._potential.core import AbstractPotential -from galdynamix.typing import Unit +from galax.potential._potential.core import AbstractPotential +from galax.typing import Unit from .core import AbstractParameter, ConstantParameter, ParameterCallable, UserParameter diff --git a/src/galdynamix/potential/_potential/subhalo.py b/src/galax/potential/_potential/subhalo.py similarity index 87% rename from src/galdynamix/potential/_potential/subhalo.py rename to src/galax/potential/_potential/subhalo.py index e9d3ad7e..9361ec01 100644 --- a/src/galdynamix/potential/_potential/subhalo.py +++ b/src/galax/potential/_potential/subhalo.py @@ -1,4 +1,4 @@ -"""galdynamix: Galactic Dynamix in Jax.""" +"""galax: Galactic Dynamix in Jax.""" __all__ = [ @@ -13,11 +13,11 @@ import jax.typing as jt from jax_cosmo.scipy.interpolate import InterpolatedUnivariateSpline -from galdynamix.potential._potential.builtin import IsochronePotential -from galdynamix.potential._potential.core import AbstractPotential -from galdynamix.potential._potential.param import AbstractParameter, ParameterField -from galdynamix.units import galactic -from galdynamix.utils import partial_jit +from galax.potential._potential.builtin import IsochronePotential +from galax.potential._potential.core import AbstractPotential +from galax.potential._potential.param import AbstractParameter, ParameterField +from galax.units import galactic +from galax.utils import partial_jit # ------------------------------------------------------------------- diff --git a/src/galdynamix/potential/_potential/utils.py b/src/galax/potential/_potential/utils.py similarity index 88% rename from src/galdynamix/potential/_potential/utils.py rename to src/galax/potential/_potential/utils.py index 79259556..b56c0f1d 100644 --- a/src/galdynamix/potential/_potential/utils.py +++ b/src/galax/potential/_potential/utils.py @@ -1,10 +1,10 @@ -"""galdynamix: Galactic Dynamix in Jax.""" +"""galax: Galactic Dynamix in Jax.""" from functools import singledispatch from typing import Any -from galdynamix.units import UnitSystem, dimensionless, galactic, solarsystem +from galax.units import UnitSystem, dimensionless, galactic, solarsystem @singledispatch diff --git a/src/galdynamix/py.typed b/src/galax/py.typed similarity index 100% rename from src/galdynamix/py.typed rename to src/galax/py.typed diff --git a/src/galdynamix/typing.py b/src/galax/typing.py similarity index 98% rename from src/galdynamix/typing.py rename to src/galax/typing.py index 000b5000..57714b63 100644 --- a/src/galdynamix/typing.py +++ b/src/galax/typing.py @@ -1,4 +1,4 @@ -"""Type hints for galdynamix.""" +"""Type hints for galax.""" # TODO: Finalize variable names and make everything public. __all__: list[str] = [] diff --git a/src/galdynamix/units.py b/src/galax/units.py similarity index 100% rename from src/galdynamix/units.py rename to src/galax/units.py diff --git a/src/galdynamix/utils/__init__.py b/src/galax/utils/__init__.py similarity index 81% rename from src/galdynamix/utils/__init__.py rename to src/galax/utils/__init__.py index cde25cfb..ae1b03b0 100644 --- a/src/galdynamix/utils/__init__.py +++ b/src/galax/utils/__init__.py @@ -1,4 +1,4 @@ -"""galdynamix: Galactic Dynamix in Jax.""" +"""galax: Galactic Dynamix in Jax.""" from . import _collections, _jax, dataclasses diff --git a/src/galdynamix/utils/_collections.py b/src/galax/utils/_collections.py similarity index 96% rename from src/galdynamix/utils/_collections.py rename to src/galax/utils/_collections.py index 2925abb4..9ab2a8e0 100644 --- a/src/galdynamix/utils/_collections.py +++ b/src/galax/utils/_collections.py @@ -1,4 +1,4 @@ -"""galdynamix: Galactic Dynamix in Jax.""" +"""galax: Galactic Dynamix in Jax.""" __all__ = ["ImmutableDict"] @@ -25,7 +25,7 @@ class ImmutableDict(Mapping[str, V]): Examples -------- - >>> from galdynamix.utils import ImmutableDict + >>> from galax.utils import ImmutableDict >>> d = ImmutableDict(a=1, b=2) >>> d ImmutableDict({'a': 1, 'b': 2}) diff --git a/src/galdynamix/utils/_jax.py b/src/galax/utils/_jax.py similarity index 97% rename from src/galdynamix/utils/_jax.py rename to src/galax/utils/_jax.py index 48f34b27..31bb013c 100644 --- a/src/galdynamix/utils/_jax.py +++ b/src/galax/utils/_jax.py @@ -1,4 +1,4 @@ -"""galdynamix: Galactic Dynamix in Jax.""" +"""galax: Galactic Dynamix in Jax.""" __all__ = [ diff --git a/src/galdynamix/utils/_misc.py b/src/galax/utils/_misc.py similarity index 83% rename from src/galdynamix/utils/_misc.py rename to src/galax/utils/_misc.py index f446b874..30d13dbc 100644 --- a/src/galdynamix/utils/_misc.py +++ b/src/galax/utils/_misc.py @@ -1,4 +1,4 @@ -"""galdynamix: Galactic Dynamix in Jax.""" +"""galax: Galactic Dynamix in Jax.""" __all__: list[str] = [] diff --git a/src/galdynamix/utils/_shape.py b/src/galax/utils/_shape.py similarity index 94% rename from src/galdynamix/utils/_shape.py rename to src/galax/utils/_shape.py index eceeaa36..9bf71ad9 100644 --- a/src/galdynamix/utils/_shape.py +++ b/src/galax/utils/_shape.py @@ -1,4 +1,4 @@ -"""galdynamix: Galactic Dynamix in Jax.""" +"""galax: Galactic Dynamix in Jax.""" __all__: list[str] = [] @@ -7,7 +7,7 @@ import jax.numpy as xp from jaxtyping import Array, ArrayLike -from galdynamix.typing import AnyScalar, ArrayAnyShape +from galax.typing import AnyScalar, ArrayAnyShape from ._jax import partial_jit @@ -47,7 +47,7 @@ def atleast_batched(*arys: Any) -> Array | tuple[Array, ...]: Examples -------- - >>> from galdynamix.utils._shape import atleast_batched + >>> from galax.utils._shape import atleast_batched >>> atleast_batched(0) Array([[0]], dtype=int64, ...) @@ -137,7 +137,7 @@ def batched_shape( Standard imports: >>> import jax.numpy as xp - >>> from galdynamix.utils._shape import batched_shape + >>> from galax.utils._shape import batched_shape Expecting a scalar: @@ -191,7 +191,7 @@ def expand_batch_dims(arr: ArrayAnyShape, /, ndim: int) -> ArrayAnyShape: Examples -------- >>> import jax.numpy as jnp - >>> from galdynamix.utils._shape import expand_batch_dims + >>> from galax.utils._shape import expand_batch_dims >>> expand_batch_dims(jnp.array(0), ndim=0).shape () @@ -226,7 +226,7 @@ def expand_arr_dims(arr: ArrayAnyShape, /, ndim: int) -> ArrayAnyShape: Examples -------- >>> import jax.numpy as jnp - >>> from galdynamix.utils._shape import expand_arr_dims + >>> from galax.utils._shape import expand_arr_dims >>> expand_arr_dims(jnp.array(0), ndim=0).shape () diff --git a/src/galdynamix/utils/dataclasses.py b/src/galax/utils/dataclasses.py similarity index 99% rename from src/galdynamix/utils/dataclasses.py rename to src/galax/utils/dataclasses.py index dae970ec..7ebe4900 100644 --- a/src/galdynamix/utils/dataclasses.py +++ b/src/galax/utils/dataclasses.py @@ -1,4 +1,4 @@ -"""galdynamix: Galactic Dynamix in Jax.""" +"""galax: Galactic Dynamix in Jax.""" __all__ = ["field"] diff --git a/src/galdynamix/integrate/__init__.py b/src/galdynamix/integrate/__init__.py deleted file mode 100644 index 6a02e1ba..00000000 --- a/src/galdynamix/integrate/__init__.py +++ /dev/null @@ -1,10 +0,0 @@ -"""galdynamix: Galactic Dynamix in Jax.""" - - -from galdynamix.integrate import _base, _builtin -from galdynamix.integrate._base import * -from galdynamix.integrate._builtin import * - -__all__: list[str] = [] -__all__ += _base.__all__ -__all__ += _builtin.__all__ diff --git a/tests/potential/builtin/test_bar.py b/tests/potential/builtin/test_bar.py index 657a1002..e63af621 100644 --- a/tests/potential/builtin/test_bar.py +++ b/tests/potential/builtin/test_bar.py @@ -3,7 +3,7 @@ import jax.numpy as xp import pytest -import galdynamix.potential as gp +import galax.potential as gp from ..test_core import TestAbstractPotential from .test_common import ( diff --git a/tests/potential/builtin/test_common.py b/tests/potential/builtin/test_common.py index 50f5a3a0..18191bcd 100644 --- a/tests/potential/builtin/test_common.py +++ b/tests/potential/builtin/test_common.py @@ -1,9 +1,9 @@ import astropy.units as u import pytest -import galdynamix.potential as gp -from galdynamix.potential import ConstantParameter -from galdynamix.units import galactic +import galax.potential as gp +from galax.potential import ConstantParameter +from galax.units import galactic from ..params.test_field import ParameterFieldMixin diff --git a/tests/potential/builtin/test_isochrone.py b/tests/potential/builtin/test_isochrone.py index a4434891..5856ca6b 100644 --- a/tests/potential/builtin/test_isochrone.py +++ b/tests/potential/builtin/test_isochrone.py @@ -3,7 +3,7 @@ import jax.numpy as xp import pytest -import galdynamix.potential as gp +import galax.potential as gp from ..test_core import TestAbstractPotential from .test_common import MassParameterMixin, ShapeAParameterMixin diff --git a/tests/potential/builtin/test_miyamotonagai.py b/tests/potential/builtin/test_miyamotonagai.py index 8676d7f6..c92ad6d6 100644 --- a/tests/potential/builtin/test_miyamotonagai.py +++ b/tests/potential/builtin/test_miyamotonagai.py @@ -3,7 +3,7 @@ import jax.numpy as xp import pytest -import galdynamix.potential as gp +import galax.potential as gp from ..test_core import TestAbstractPotential from .test_common import ( @@ -20,7 +20,7 @@ class TestMiyamotoNagaiPotential( ShapeAParameterMixin, ShapeBParameterMixin, ): - """Test the `galdynamix.potential.MiyamotoNagaiPotential` class.""" + """Test the `galax.potential.MiyamotoNagaiPotential` class.""" @pytest.fixture(scope="class") def pot_cls(self) -> type[gp.MiyamotoNagaiPotential]: diff --git a/tests/potential/builtin/test_nfw.py b/tests/potential/builtin/test_nfw.py index f50f200d..082731a4 100644 --- a/tests/potential/builtin/test_nfw.py +++ b/tests/potential/builtin/test_nfw.py @@ -4,9 +4,9 @@ import jax.numpy as xp import pytest -import galdynamix.potential as gp -from galdynamix.potential import ConstantParameter -from galdynamix.units import galactic +import galax.potential as gp +from galax.potential import ConstantParameter +from galax.units import galactic from ..params.test_field import ParameterFieldMixin from ..test_core import TestAbstractPotential diff --git a/tests/potential/params/test_core.py b/tests/potential/params/test_core.py index 680b14eb..68d626a7 100644 --- a/tests/potential/params/test_core.py +++ b/tests/potential/params/test_core.py @@ -1,4 +1,4 @@ -"""Test :mod:`galdynamix.potential._potential.param.core`.""" +"""Test :mod:`galax.potential._potential.param.core`.""" import astropy.units as u @@ -6,13 +6,13 @@ import jaxtyping import pytest -from galdynamix.potential import AbstractParameter, ConstantParameter, UserParameter -from galdynamix.potential._potential.param.core import ParameterCallable -from galdynamix.typing import Unit +from galax.potential import AbstractParameter, ConstantParameter, UserParameter +from galax.potential._potential.param.core import ParameterCallable +from galax.typing import Unit class TestAbstractParameter: - """Test the `galdynamix.potential.AbstractParameter` class.""" + """Test the `galax.potential.AbstractParameter` class.""" @pytest.fixture(scope="class") def param_cls(self) -> type[AbstractParameter]: @@ -35,19 +35,19 @@ def __call__(self, t, **kwargs): # ========================================================================== def test_init(self): - """Test init `galdynamix.potential.AbstractParameter` method.""" + """Test init `galax.potential.AbstractParameter` method.""" # Test that the abstract class cannot be instantiated with pytest.raises(TypeError): AbstractParameter() def test_call(self): - """Test `galdynamix.potential.AbstractParameter` call method.""" + """Test `galax.potential.AbstractParameter` call method.""" # Test that the abstract class cannot be instantiated with pytest.raises(TypeError): AbstractParameter()() def test_unit_field(self, param, field_unit): - """Test `galdynamix.potential.AbstractParameter` unit field.""" + """Test `galax.potential.AbstractParameter` unit field.""" assert param.unit == field_unit @@ -55,7 +55,7 @@ def test_unit_field(self, param, field_unit): class TestConstantParameter(TestAbstractParameter): - """Test the `galdynamix.potential.ConstantParameter` class.""" + """Test the `galax.potential.ConstantParameter` class.""" @pytest.fixture(scope="class") def param_cls(self) -> type[AbstractParameter]: @@ -72,7 +72,7 @@ def param(self, param_cls, field_unit, field_value) -> AbstractParameter: # ========================================================================== def test_call(self, param, field_value): - """Test `galdynamix.potential.ConstantParameter` call method.""" + """Test `galax.potential.ConstantParameter` call method.""" assert param(t=1.0) == field_value assert param(t=1.0 * u.s) == field_value assert xp.array_equal(param(t=xp.array([1.0, 2.0])), [field_value, field_value]) @@ -82,7 +82,7 @@ def test_call(self, param, field_value): class TestParameterCallable: - """Test the `galdynamix.potential.ParameterCallable` class.""" + """Test the `galax.potential.ParameterCallable` class.""" def test_issubclass(self): assert issubclass(AbstractParameter, ParameterCallable) @@ -98,7 +98,7 @@ def test_isinstance(self): class TestUserParameter(TestAbstractParameter): - """Test the `galdynamix.potential.UserParameter` class.""" + """Test the `galax.potential.UserParameter` class.""" @pytest.fixture(scope="class") def param_cls(self) -> type[AbstractParameter]: @@ -118,7 +118,7 @@ def param(self, param_cls, field_unit, field_func) -> AbstractParameter: # ========================================================================== def test_call(self, param): - """Test `galdynamix.potential.UserParameter` call method.""" + """Test `galax.potential.UserParameter` call method.""" assert param(t=1.0) == 1.0 assert param(t=1.0 * u.s) == 1.0 * u.s diff --git a/tests/potential/params/test_field.py b/tests/potential/params/test_field.py index c6d89c89..4db3d04d 100644 --- a/tests/potential/params/test_field.py +++ b/tests/potential/params/test_field.py @@ -1,5 +1,5 @@ -"""Test :mod:`galdynamix.potential._potential.param.core`.""" +"""Test :mod:`galax.potential._potential.param.core`.""" class ParameterFieldMixin: - """Test the `galdynamix.potential.ParameterField` class.""" + """Test the `galax.potential.ParameterField` class.""" diff --git a/tests/potential/test_base.py b/tests/potential/test_base.py index 1c76896c..4ca5dbd3 100644 --- a/tests/potential/test_base.py +++ b/tests/potential/test_base.py @@ -6,14 +6,14 @@ import pytest from jaxtyping import Array, Float -import galdynamix.dynamics as gd -import galdynamix.potential as gp -from galdynamix.typing import BatchableFloatOrIntScalarLike, BatchFloatScalar, BatchVec3 -from galdynamix.utils import partial_jit, vectorize_method +import galax.dynamics as gd +import galax.potential as gp +from galax.typing import BatchableFloatOrIntScalarLike, BatchFloatScalar, BatchVec3 +from galax.utils import partial_jit, vectorize_method class TestAbstractPotentialBase: - """Test the `galdynamix.potential.AbstractPotentialBase` class.""" + """Test the `galax.potential.AbstractPotentialBase` class.""" @pytest.fixture(scope="class") def pot_cls(self) -> type[gp.AbstractPotentialBase]: diff --git a/tests/potential/test_core.py b/tests/potential/test_core.py index d591a1c2..6440b6e1 100644 --- a/tests/potential/test_core.py +++ b/tests/potential/test_core.py @@ -5,18 +5,18 @@ import jax.numpy as xp import pytest -import galdynamix.potential as gp -from galdynamix.potential._potential.utils import converter_to_usys -from galdynamix.typing import BatchableFloatOrIntScalarLike, BatchFloatScalar, BatchVec3 -from galdynamix.units import UnitSystem, dimensionless, galactic -from galdynamix.utils import partial_jit, vectorize_method +import galax.potential as gp +from galax.potential._potential.utils import converter_to_usys +from galax.typing import BatchableFloatOrIntScalarLike, BatchFloatScalar, BatchVec3 +from galax.units import UnitSystem, dimensionless, galactic +from galax.utils import partial_jit, vectorize_method from .test_base import TestAbstractPotentialBase from .test_utils import FieldUnitSystemMixin class TestAbstractPotential(TestAbstractPotentialBase, FieldUnitSystemMixin): - """Test the `galdynamix.potential.AbstractPotentialBase` class.""" + """Test the `galax.potential.AbstractPotentialBase` class.""" @pytest.fixture(scope="class") def pot_cls(self) -> type[gp.AbstractPotentialBase]: diff --git a/tests/potential/test_package.py b/tests/potential/test_package.py index 9c635be4..4dfc09ff 100644 --- a/tests/potential/test_package.py +++ b/tests/potential/test_package.py @@ -1,9 +1,9 @@ -import galdynamix.potential as gp -from galdynamix.potential import _potential +import galax.potential as gp +from galax.potential import _potential def test_all(): - """Test the `galdynamix.potential` package contents.""" + """Test the `galax.potential` package contents.""" # Test correct dumping of contents assert gp.__all__ == _potential.__all__ diff --git a/tests/potential/test_utils.py b/tests/potential/test_utils.py index 61c885d5..2a809096 100644 --- a/tests/potential/test_utils.py +++ b/tests/potential/test_utils.py @@ -1,11 +1,11 @@ -"""Tests for `galdynamix.potential._potential.utils` package.""" +"""Tests for `galax.potential._potential.utils` package.""" from dataclasses import replace import astropy.units as u import pytest -from galdynamix.potential._potential.utils import ( +from galax.potential._potential.utils import ( UnitSystem, converter_to_usys, dimensionless, @@ -15,7 +15,7 @@ class TestConverterToUtils: - """Tests for `galdynamix.potential._potential.utils.converter_to_usys`.""" + """Tests for `galax.potential._potential.utils.converter_to_usys`.""" def test_invalid(self): """Test conversion from unsupported value.""" diff --git a/tests/test_package.py b/tests/test_package.py index 3b4c47a6..f5a6b674 100644 --- a/tests/test_package.py +++ b/tests/test_package.py @@ -1,7 +1,7 @@ from importlib.metadata import version -import galdynamix as pkg +import galax as pkg def test_version(): - assert version("galdynamix") == pkg.__version__ + assert version("galax") == pkg.__version__ diff --git a/tests/utils/test_collections.py b/tests/utils/test_collections.py index 0af11a2e..dcd6dac7 100644 --- a/tests/utils/test_collections.py +++ b/tests/utils/test_collections.py @@ -3,7 +3,7 @@ import pytest -from galdynamix.utils import ImmutableDict +from galax.utils import ImmutableDict class TestImmutableDict: diff --git a/tests/utils/test_dataclasses.py b/tests/utils/test_dataclasses.py index 01f7cae5..33f80786 100644 --- a/tests/utils/test_dataclasses.py +++ b/tests/utils/test_dataclasses.py @@ -2,7 +2,7 @@ import equinox as eqx -from galdynamix.utils.dataclasses import field +from galax.utils.dataclasses import field COMP_FIELDS: tuple[str, ...] = ( "name", diff --git a/tests/utils/test_jax.py b/tests/utils/test_jax.py index 29afb952..94ea7a98 100644 --- a/tests/utils/test_jax.py +++ b/tests/utils/test_jax.py @@ -2,7 +2,7 @@ import jax.numpy as jnp from jaxtyping import Array, Float -from galdynamix.utils import ( +from galax.utils import ( partial_jit, partial_vectorize, partial_vmap, diff --git a/tests/utils/test_misc.py b/tests/utils/test_misc.py index 702b0818..5ba4e3fb 100644 --- a/tests/utils/test_misc.py +++ b/tests/utils/test_misc.py @@ -1,7 +1,7 @@ import numpy as np import pytest -from galdynamix.utils._misc import first +from galax.utils._misc import first @pytest.mark.parametrize( diff --git a/tests/utils/test_package.py b/tests/utils/test_package.py index 4c182596..e2b2c110 100644 --- a/tests/utils/test_package.py +++ b/tests/utils/test_package.py @@ -1,10 +1,10 @@ import pytest -from galdynamix import utils +from galax import utils def test__all__(): - """Test that `galdynamix.utils` has the expected `__all__`.""" + """Test that `galax.utils` has the expected `__all__`.""" assert utils.__all__ == [ "dataclasses", *utils._jax.__all__, diff --git a/tests/utils/test_shape.py b/tests/utils/test_shape.py index 543a6409..4dfd2661 100644 --- a/tests/utils/test_shape.py +++ b/tests/utils/test_shape.py @@ -1,11 +1,11 @@ -"""Test the `galdynamix.utils._shape` module.""" +"""Test the `galax.utils._shape` module.""" import re import jax.numpy as xp import pytest -from galdynamix.utils._shape import atleast_batched, batched_shape +from galax.utils._shape import atleast_batched, batched_shape class TestAtleastBatched: @@ -54,7 +54,7 @@ def test_atleast_batched_multiple_args(self): class TestBatchedShape: - """Test the `galdynamix.utils._shape.batched_shape` function.""" + """Test the `galax.utils._shape.batched_shape` function.""" @pytest.mark.parametrize( ("arr", "expect_ndim", "expect"), @@ -67,7 +67,7 @@ class TestBatchedShape: ], ) def test_batched_shape(self, arr, expect_ndim, expect): - """Test the `galdynamix.utils._shape.batched_shape` function.""" + """Test the `galax.utils._shape.batched_shape` function.""" batch, shape = batched_shape(arr, expect_ndim=expect_ndim) assert batch == expect[0] assert shape == expect[1]