SCF

Signed-off-by: nstarman <[email protected]>

pyproject.toml

@@ -44,6 +44,7 @@ scf = [
   "sympy2jax",
 ]
 test = [
+  "gala",
   "hypothesis[numpy]",
   "pytest >=6",
   "pytest-cov >=3",
@@ -159,6 +160,7 @@ ignore = [
   "F821",    # undefined name  <- jaxtyping
   "FIX002",  # Line contains TODO, consider resolving the issue
   "N80",     # Naming conventions.
+  "N816",    # Variable in global scope should not be mixedCase
   "PD",      # pandas-vet
   "PLR",     # Design related pylint codes
   "PYI041",  # Use `float` instead of `int | float` <- beartype is more strict

src/galax/dynamics/_core.py

@@ -51,7 +51,7 @@ def shape(self) -> tuple[int, ...]:
     @property
     @partial_jit()
     def qp(self) -> BatchVec6:
-        """Return as a single Array[(*batch, Q + P),]."""
+        """Return as a single Array[float, (*batch, Q + P),]."""
         batch_shape, component_shapes = self._shape_tuple
         q = xp.broadcast_to(self.q, batch_shape + component_shapes[0:1])
         p = xp.broadcast_to(self.p, batch_shape + component_shapes[1:2])
@@ -99,7 +99,7 @@ def _shape_tuple(self) -> tuple[tuple[int, ...], tuple[int, int, int]]:
     @property
     @partial_jit()
     def w(self) -> BatchVec7:
-        """Return as a single Array[(*batch, Q + P + T),]."""
+        """Return as a single Array[float, (*batch, Q + P + T)]."""
         batch_shape, component_shapes = self._shape_tuple
         q = xp.broadcast_to(self.q, batch_shape + component_shapes[0:1])
         p = xp.broadcast_to(self.p, batch_shape + component_shapes[1:2])

src/galax/dynamics/mockstream/_core.py

@@ -38,7 +38,7 @@ def _shape_tuple(self) -> tuple[tuple[int, ...], tuple[int, int, int]]:
     @property
     @partial_jit()
     def w(self) -> BatchVec7:
-        """Return as a single Array[(*batch, Q + P + T),]."""
+        """Return as a single Array[float, (*batch, Q + P + T)]."""
         batch_shape, component_shapes = self._shape_tuple
         q = xp.broadcast_to(self.q, batch_shape + component_shapes[0:1])
         p = xp.broadcast_to(self.p, batch_shape + component_shapes[1:2])

src/galax/dynamics/mockstream/_df/fardal.py

@@ -117,7 +117,7 @@ def dphidr(potential: AbstractPotentialBase, x: Vec3, t: FloatScalar) -> Vec3:
     ----------
     potential: AbstractPotentialBase
         The gravitational potential.
-    x: Array[(3,), Any]
+    x: Array[Any, (3,)]
         3d position (x, y, z) in [kpc]
     t: Numeric
         Time in [Myr]
@@ -143,7 +143,7 @@ def d2phidr2(
     ----------
     potential: AbstractPotentialBase
         The gravitational potential.
-    x: Array[(3,), Any]
+    x: Array[Any, (3,)]
         3d position (x, y, z) in [kpc]
     t: Numeric
         Time in [Myr]
@@ -172,9 +172,9 @@ def orbital_angular_velocity(x: Vec3, v: Vec3, /) -> Vec3:
 
     Arguments:
     ---------
-    x: Array[(3,), Any]
+    x: Array[Any, (3,)]
         3d position (x, y, z) in [length]
-    v: Array[(3,), Any]
+    v: Array[Any, (3,)]
         3d velocity (v_x, v_y, v_z) in [length/time]
 
     Returns:
@@ -199,9 +199,9 @@ def orbital_angular_velocity_mag(x: Vec3, v: Vec3, /) -> FloatScalar:
 
     Arguments:
     ---------
-    x: Array[(3,), Any]
+    x: Array[Any, (3,)]
         3d position (x, y, z) in [kpc]
-    v: Array[(3,), Any]
+    v: Array[Any, (3,)]
         3d velocity (v_x, v_y, v_z) in [kpc/Myr]
 
     Returns:

src/galax/potential/_potential/__init__.py

@@ -7,11 +7,12 @@
 from .composite import *
 from .core import *
 from .param import *
+from .scf import SCFPotential
 
 __all__: list[str] = []
 __all__ += base.__all__
 __all__ += core.__all__
 __all__ += composite.__all__
 __all__ += param.__all__
 __all__ += builtin.__all__
-__all__ += ["scf"]
+__all__ += ["scf", "SCFPotential"]

src/galax/potential/_potential/param/core.py

@@ -43,14 +43,14 @@ def __call__(self, t: FloatScalar, **kwargs: Any) -> FloatArrayAnyShape:
 
         Parameters
         ----------
-        t : Array
+        t : float | Array[float, ()]
             The time(s) at which to compute the parameter value.
-        **kwargs
+        **kwargs : Any
             Additional parameters to pass to the parameter function.
 
         Returns
         -------
-        Array
+        Array[float, "*shape"]
             The parameter value at times ``t``.
         """
         ...
@@ -74,7 +74,7 @@ def __call__(
 
         Parameters
         ----------
-        t : Array, optional
+        t : float | Array[float, ()], optional
             This is ignored and is thus optional.
             Note that for most :class:`~galax.potential.AbstractParameter`
             the time is required.
@@ -83,7 +83,7 @@ def __call__(
 
         Returns
         -------
-        Array
+        Array[float, "*shape"]
             The constant parameter value.
         """
         # Vectorization to enable broadcasting over the time dimension.
@@ -107,14 +107,14 @@ def __call__(self, t: FloatScalar, **kwargs: Any) -> FloatArrayAnyShape:
 
         Parameters
         ----------
-        t : Array
+        t : float | Array[float, ()]
             Time(s) at which to compute the parameter value.
         **kwargs : Any
             Additional parameters to pass to the parameter function.
 
         Returns
         -------
-        Array
+        Array[float, "*shape"]
             Parameter value(s) at the given time(s).
         """
         ...

src/galax/potential/_potential/scf/__init__.py

@@ -1,7 +1,11 @@
-from __future__ import annotations
-
-from . import gegenbauer
-from .gegenbauer import *
+from . import bfe, bfe_helper, coeffs, coeffs_helper
+from .bfe import *
+from .bfe_helper import *
+from .coeffs import *
+from .coeffs_helper import *
 
 __all__: list[str] = []
-__all__ += gegenbauer.__all__
+__all__ += bfe.__all__
+__all__ += bfe_helper.__all__
+__all__ += coeffs.__all__
+__all__ += coeffs_helper.__all__