update infrastructure

.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: Nathaniel Starkman
-license: BSD
-org: nstarman
+full_name: galdynamix maintainers
+license: MIT
+org: galdynamix
 project_name: galdynamix
-project_short_description: Galactic Dynamix in Jax
-url: https://github.com/nstarman/galdynamix
+project_short_description: Galactic Dynamix in Jax.
+url: https://github.com/galdynamix/galdynamix
 vcs: true

LICENSE

@@ -1,29 +1,19 @@
-BSD 3-Clause License
+Copyright 2023 galdynamix maintainers
 
-Copyright (c) 2023, Nathaniel Starkman.
-All rights reserved.
+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
+the Software without restriction, including without limitation the rights to
+use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
+of the Software, and to permit persons to whom the Software is furnished to do
+so, subject to the following conditions:
 
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions are met:
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
 
-* Redistributions of source code must retain the above copyright notice, this
-  list of conditions and the following disclaimer.
-
-* Redistributions in binary form must reproduce the above copyright notice,
-  this list of conditions and the following disclaimer in the documentation
-  and/or other materials provided with the distribution.
-
-* Neither the name of the vector package developers nor the names of its
-  contributors may be used to endorse or promote products derived from
-  this software without specific prior written permission.
-
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
-FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
-DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
-SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
-CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
-OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -12,12 +12,12 @@
 <!-- SPHINX-START -->
 
 <!-- prettier-ignore-start -->
-[actions-badge]:            https://github.com/nstarman/galdynamix/workflows/CI/badge.svg
-[actions-link]:             https://github.com/nstarman/galdynamix/actions
+[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
 [github-discussions-badge]: https://img.shields.io/static/v1?label=Discussions&message=Ask&color=blue&logo=github
-[github-discussions-link]:  https://github.com/nstarman/galdynamix/discussions
+[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

docs/conf.py

@@ -1,10 +1,12 @@
+"""Sphinx configuration."""
+
 from __future__ import annotations
 
 import importlib.metadata
 
 project = "galdynamix"
-copyright = "2023, Nathaniel Starkman"
-author = "Nathaniel Starkman"
+copyright = "2023, galdynamix maintainers"
+author = "galdynamix maintainers"
 version = release = importlib.metadata.version("galdynamix")
 
 extensions = [

noxfile.py

@@ -1,3 +1,5 @@
+"""Configuration for Nox."""
+
 from __future__ import annotations
 
 import argparse
@@ -13,9 +15,7 @@
 
 @nox.session
 def lint(session: nox.Session) -> None:
-    """
-    Run the linter.
-    """
+    """Run the linter."""
     session.install("pre-commit")
     session.run(
         "pre-commit", "run", "--all-files", "--show-diff-on-failure", *session.posargs
@@ -35,19 +35,14 @@ def lint(session: nox.Session) -> None:
 
 @nox.session
 def tests(session: nox.Session) -> None:
-    """
-    Run the unit and regular tests.
-    """
+    """Run the unit and regular tests."""
     session.install(".[test]")
     session.run("pytest", *session.posargs)
 
 
 @nox.session(reuse_venv=True)
 def docs(session: nox.Session) -> None:
-    """
-    Build the docs. Pass "--serve" to serve. Pass "-b linkcheck" to check links.
-    """
-
+    """Build the docs. Pass "--serve" to serve. Pass "-b linkcheck" to check links."""
     parser = argparse.ArgumentParser()
     parser.add_argument("--serve", action="store_true", help="Serve after building")
     parser.add_argument(
@@ -86,10 +81,7 @@ def docs(session: nox.Session) -> None:
 
 @nox.session
 def build_api_docs(session: nox.Session) -> None:
-    """
-    Build (regenerate) API docs.
-    """
-
+    """Build (regenerate) API docs."""
     session.install("sphinx")
     session.chdir("docs")
     session.run(
@@ -105,10 +97,7 @@ def build_api_docs(session: nox.Session) -> None:
 
 @nox.session
 def build(session: nox.Session) -> None:
-    """
-    Build an SDist and wheel.
-    """
-
+    """Build an SDist and wheel."""
     build_path = DIR.joinpath("build")
     if build_path.exists():
         shutil.rmtree(build_path)

pyproject.toml

@@ -6,16 +6,18 @@ build-backend = "hatchling.build"
 [project]
 name = "galdynamix"
 authors = [
+  { name = "galdynamix maintainers", email = "nstarman@users.noreply.github.com" },
+  { name = "Jake Nibauer", email = "jnibauer@princeton.edu" },
   { name = "Nathaniel Starkman", email = "nstarman@users.noreply.github.com" },
 ]
-description = "Galactic Dynamix in Jax"
+description = "Galactic Dynamix in Jax."
 readme = "README.md"
 requires-python = ">=3.11"
 classifiers = [
   "Development Status :: 1 - Planning",
   "Intended Audience :: Science/Research",
   "Intended Audience :: Developers",
-  "License :: OSI Approved :: BSD License",
+  "License :: OSI Approved :: MIT License",
   "Operating System :: OS Independent",
   "Programming Language :: Python",
   "Programming Language :: Python :: 3",
@@ -53,10 +55,10 @@ docs = [
 ]
 
 [project.urls]
-Homepage = "https://github.com/nstarman/galdynamix"
-"Bug Tracker" = "https://github.com/nstarman/galdynamix/issues"
-Discussions = "https://github.com/nstarman/galdynamix/discussions"
-Changelog = "https://github.com/nstarman/galdynamix/releases"
+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"
 
 
 [tool.hatch]
@@ -114,32 +116,22 @@ ignore_missing_imports = true
 src = ["src"]
 
 [tool.ruff.lint]
-extend-select = [
-  "B",           # flake8-bugbear
-  "I",           # isort
-  "ARG",         # flake8-unused-arguments
-  "C4",          # flake8-comprehensions
-  "EM",          # flake8-errmsg
-  "ICN",         # flake8-import-conventions
-  "G",           # flake8-logging-format
-  "PGH",         # pygrep-hooks
-  "PIE",         # flake8-pie
-  "PL",          # pylint
-  "PT",          # flake8-pytest-style
-  "PTH",         # flake8-use-pathlib
-  "RET",         # flake8-return
-  "RUF",         # Ruff-specific
-  "SIM",         # flake8-simplify
-  "T20",         # flake8-print
-  "UP",          # pyupgrade
-  "YTT",         # flake8-2020
-  "EXE",         # flake8-executable
-  "NPY",         # NumPy specific rules
-]
+extend-select = ["ALL"]
 ignore = [
-  "PD",  # pandas-vet
-  "PLR",    # Design related pylint codes
-  # TODO! fix these
+  "ANN101",  # Missing type annotation for self in method
+  "COM812",  # Missing trailing comma in Python 3.6+
+  "D203",    # 1 blank line required before class docstring
+  "D105",    # Missing docstring in magic method
+  "D107",    # Missing docstring in `__init__`
+  "D213",    # Multi-line docstring summary should start at the second line
+  "FIX002",  # Line contains TODO, consider resolving the issue
+  "N80",     # Naming conventions.
+  "PD",      # pandas-vet
+  "PLR",     # Design related pylint codes
+  "TCH00",   # Move into a type-checking block
+  "TD002",   # Missing author in TODO
+  "TD003",   # Missing issue link on the line following this TODO
+  # TODO: fix these
   "ARG001",
   "ARG002",
   "F841",
@@ -149,8 +141,12 @@ isort.required-imports = ["from __future__ import annotations"]
 # typing-modules = ["galdynamix._compat.typing"]
 
 [tool.ruff.lint.per-file-ignores]
-"tests/**" = ["T20"]
-"noxfile.py" = ["T20"]
+"tests/**" = ["ANN", "D10", "INP001", "S101", "T20"]
+"noxfile.py" = ["ERA001", "T20"]
+"docs/conf.py" = [
+  "A001",  # Variable `copyright` is shadowing a Python builtin
+  "INP001",  # implicit namespace package
+]
 
 
 [tool.pylint]

src/galdynamix/__init__.py

@@ -1,4 +1,4 @@
-"""galdynamix: Galactic Dynamix in Jax"""
+"""Copyright (c) 2023 galdynamix maintainers. All rights reserved."""
 from __future__ import annotations
 
 __all__ = ["__version__"]
@@ -7,4 +7,4 @@
 
 from ._version import version as __version__
 
-config.update("jax_enable_x64", True)
+config.update("jax_enable_x64", True)  # noqa: FBT003

src/galdynamix/_version.pyi

@@ -1,4 +1,2 @@
-from __future__ import annotations
-
 version: str
 version_tuple: tuple[int, int, int] | tuple[int, int, int, str, str]

src/galdynamix/dynamics/__init__.py

@@ -1,4 +1,4 @@
-"""galdynamix: Galactic Dynamix in Jax"""
+"""galdynamix: Galactic Dynamix in Jax."""
 # ruff: noqa: F403
 
 from __future__ import annotations

src/galdynamix/dynamics/_orbit.py

@@ -1,4 +1,4 @@
-"""galdynamix: Galactic Dynamix in Jax"""
+"""galdynamix: Galactic Dynamix in Jax."""
 
 from __future__ import annotations
 
@@ -10,12 +10,14 @@
 import jax.typing as jt
 
 from galdynamix.potential._potential.base import AbstractPotentialBase
+from galdynamix.utils._jax import partial_jit
 
 
 class Orbit(eqx.Module):  # type: ignore[misc]
     """Orbit.
 
-    TODO:
+    Todo:
+    ----
     - Units stuff
     - GR stuff
     """
@@ -27,20 +29,36 @@ class Orbit(eqx.Module):  # type: ignore[misc]
     """Position of the stream particles (x, y, z) [kpc/Myr]."""
 
     t: jt.Array
-    """Release time of the stream particles [Myr]."""
+    """Array of times [Myr]."""
 
     potential: AbstractPotentialBase
     """Potential in which the orbit was integrated."""
 
-    def to_w(self) -> jt.Array:
+    @property
+    @partial_jit()
+    def qp(self) -> jt.Array:
+        """Return as a single Array[(N, Q + P),]."""
+        # Determine output shape
+        qd = self.q.shape[1]  # dimensionality of q
+        shape = (self.q.shape[0], qd + self.p.shape[1])
+        # Create output array (jax will fuse these ops)
+        out = xp.empty(shape)
+        out = out.at[:, :qd].set(self.q)
+        out = out.at[:, qd:].set(self.p)
+        return out  # noqa: RET504
+
+    @property
+    @partial_jit()
+    def w(self) -> jt.Array:
+        """Return as a single Array[(N, Q + P + T),]."""
+        qp = self.qp
+        qpd = qp.shape[1]  # dimensionality of qp
+        # Reshape t to (N, 1) if necessary
         t = self.t[:, None] if self.t.ndim == 1 else self.t
-        out = xp.empty(
-            (
-                self.q.shape[0],
-                self.q.shape[1] + self.p.shape[1] + t.shape[1],
-            )
-        )
-        out = out.at[:, : self.q.shape[1]].set(self.q)
-        out = out.at[:, self.q.shape[1] : -1].set(self.p)
-        out = out.at[:, -1:].set(t)
+        # Determine output shape
+        shape = (qp.shape[0], qpd + t.shape[1])
+        # Create output array (jax will fuse these ops)
+        out = xp.empty(shape)
+        out = out.at[:, :qpd].set(qp)
+        out = out.at[:, qpd:].set(t)
         return out  # noqa: RET504

src/galdynamix/dynamics/mockstream/__init__.py

@@ -1,4 +1,4 @@
-"""galdynamix: Galactic Dynamix in Jax"""
+"""galdynamix: Galactic Dynamix in Jax."""
 # ruff: noqa: F403
 
 from __future__ import annotations