Potential parameters attribute

src/galax/potential/_potential/base.py

@@ -2,7 +2,8 @@
 
 import abc
 from dataclasses import KW_ONLY, fields, replace
-from typing import TYPE_CHECKING, Any
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Any, ClassVar
 
 import equinox as eqx
 import jax.experimental.array_api as xp
@@ -15,6 +16,8 @@
 
 from galax.integrate._base import Integrator
 from galax.integrate._builtin import DiffraxIntegrator
+from galax.potential._potential.param.attr import ParametersAttribute
+from galax.potential._potential.param.utils import all_parameters
 from galax.typing import (
     BatchableFloatOrIntScalarLike,
     BatchFloatScalar,
@@ -44,9 +47,20 @@
 class AbstractPotentialBase(eqx.Module, metaclass=ModuleMeta, strict=True):  # type: ignore[misc]
     """Abstract Potential Class."""
 
+    parameters: ClassVar = ParametersAttribute(MappingProxyType({}))
+
     _: KW_ONLY
     units: eqx.AbstractVar[UnitSystem]
 
+    def __init_subclass__(cls) -> None:
+        """Initialize the subclass."""
+        # Replace the ``parameters`` attribute with a mapping of the values
+        type(cls).__setattr__(
+            cls,
+            "parameters",
+            ParametersAttribute(MappingProxyType(all_parameters(cls))),
+        )
+
     ###########################################################################
     # Abstract methods that must be implemented by subclasses
 

src/galax/potential/_potential/io/gala.py

@@ -24,9 +24,9 @@
     KeplerPotential,
     MiyamotoNagaiPotential,
     NFWPotential,
-    NullPotential,
 )
 from galax.potential._potential.composite import CompositePotential
+from galax.potential._potential.core import AbstractPotential
 from galax.potential._potential.special import MilkyWayPotential
 
 ##############################################################################
@@ -68,52 +68,29 @@ def _gala_to_galax_composite(pot: GalaCompositePotential, /) -> CompositePotenti
     return CompositePotential(**{k: gala_to_galax(p) for k, p in pot.items()})
 
 
-# -----------------------------------------------------------------------------
-# Builtin potentials
+_GALA_TO_GALAX_REGISTRY: dict[type[GalaPotentialBase], type[AbstractPotential]] = {
+    GalaHernquistPotential: HernquistPotential,
+    GalaIsochronePotential: IsochronePotential,
+    GalaKeplerPotential: KeplerPotential,
+    GalaMiyamotoNagaiPotential: MiyamotoNagaiPotential,
+}
 
 
-@gala_to_galax.register
-def _gala_to_galax_hernquist(pot: GalaHernquistPotential, /) -> HernquistPotential:
+@gala_to_galax.register(GalaHernquistPotential)
+@gala_to_galax.register(GalaIsochronePotential)
+@gala_to_galax.register(GalaKeplerPotential)
+@gala_to_galax.register(GalaMiyamotoNagaiPotential)
+@gala_to_galax.register(GalaNullPotential)
+def _gala_to_galax_hernquist(pot: GalaPotentialBase, /) -> AbstractPotential:
     """Convert a Gala HernquistPotential to a Galax potential."""
     if not _static_at_origin(pot):
         msg = "Galax does not support rotating or offset potentials."
         raise TypeError(msg)
-    params = pot.parameters
-    return HernquistPotential(m=params["m"], c=params["c"], units=pot.units)
-
-
-@gala_to_galax.register
-def _gala_to_galax_isochrone(pot: GalaIsochronePotential, /) -> IsochronePotential:
-    """Convert a Gala IsochronePotential to a Galax potential."""
-    if not _static_at_origin(pot):
-        msg = "Galax does not support rotating or offset potentials."
-        raise TypeError(msg)
-    params = pot.parameters
-    return IsochronePotential(m=params["m"], b=params["b"], units=pot.units)
-
-
-@gala_to_galax.register
-def _gala_to_galax_kepler(pot: GalaKeplerPotential, /) -> KeplerPotential:
-    """Convert a Gala KeplerPotential to a Galax potential."""
-    if not _static_at_origin(pot):
-        msg = "Galax does not support rotating or offset potentials."
-        raise TypeError(msg)
-    params = pot.parameters
-    return KeplerPotential(m=params["m"], units=pot.units)
+    return _GALA_TO_GALAX_REGISTRY[type(pot)](**pot.parameters, units=pot.units)
 
 
-@gala_to_galax.register
-def _gala_to_galax_miyamotonagi(
-    pot: GalaMiyamotoNagaiPotential, /
-) -> MiyamotoNagaiPotential:
-    """Convert a Gala MiyamotoNagaiPotential to a Galax potential."""
-    if not _static_at_origin(pot):
-        msg = "Galax does not support rotating or offset potentials."
-        raise TypeError(msg)
-    params = pot.parameters
-    return MiyamotoNagaiPotential(
-        m=params["m"], a=params["a"], b=params["b"], units=pot.units
-    )
+# -----------------------------------------------------------------------------
+# Builtin potentials
 
 
 @gala_to_galax.register
@@ -128,15 +105,6 @@ def _gala_to_galax_nfw(pot: GalaNFWPotential, /) -> NFWPotential:
     )
 
 
-@gala_to_galax.register
-def _gala_to_galax_nullpotential(pot: GalaNullPotential, /) -> NullPotential:
-    """Convert a Gala NullPotential to a Galax potential."""
-    if not _static_at_origin(pot):
-        msg = "Galax does not support rotating or offset potentials."
-        raise TypeError(msg)
-    return NullPotential(units=pot.units)
-
-
 # -----------------------------------------------------------------------------
 # MW potentials
 

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

@@ -0,0 +1,65 @@
+"""Descriptor for a Parameters attributes."""
+
+__all__ = ["ParametersAttribute"]
+
+from dataclasses import dataclass, field
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Any, NoReturn
+
+from .field import ParameterField
+
+if TYPE_CHECKING:
+    from galax.potential._potential.base import AbstractPotentialBase
+
+
+@dataclass(frozen=True, slots=True)
+class ParametersAttribute:
+    """Mapping of the :class:`~galax.potential.ParameterField` values.
+
+    If accessed from the :class:`~galax.potential.AbstractPotentialBase` class,
+    this returns a mapping of the :class:`~galax.potential.AbstractParameter`
+    objects themselves.  If accessed from an instance, this returns a mapping of
+    the values of the Parameters.
+
+    This class is used to implement
+    :obj:`galax.potential.AbstractPotentialBase.parameters`.
+
+    Examples
+    --------
+    The normal usage of this class is the ``parameters`` attribute of
+    :class:`~galax.potential.AbstractPotentialBase`.
+
+        >>> from galax.potential import KeplerPotential
+
+        >>> KeplerPotential.parameters
+        mappingproxy({'m': ParameterField(...), ...})
+
+        >>> import astropy.units as u
+        >>> kepler = KeplerPotential(m=1e12 * u.solMass, units="galactic")
+        >>> kepler.parameters
+        mappingproxy({'m': ConstantParameter(unit=Unit("solMass"), value=f64[])})
+    """
+
+    parameters: "MappingProxyType[str, ParameterField]"  # TODO: specify type hint
+    """Class attribute name on Potential."""
+
+    _name: str = field(init=False)
+    """The name of the descriptor on the containing class."""
+
+    def __set_name__(self, _: Any, name: str) -> None:
+        object.__setattr__(self, "_name", name)
+
+    def __get__(
+        self,
+        instance: "AbstractPotentialBase | None",
+        owner: "type[AbstractPotentialBase] | None",
+    ) -> "MappingProxyType[str, ParameterField]":
+        # Called from the class
+        if instance is None:
+            return self.parameters
+        # Called from the instance
+        return MappingProxyType({n: getattr(instance, n) for n in self.parameters})
+
+    def __set__(self, instance: Any, _: Any) -> NoReturn:
+        msg = f"cannot set {self._name!r} of {instance!r}."
+        raise AttributeError(msg)

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

@@ -62,6 +62,8 @@ class ConstantParameter(AbstractParameter):
     # TODO: unit handling
     # TODO: link this shape to the return shape from __call__
     value: FloatArrayAnyShape = eqx.field(converter=converter_float_array)
+    _: KW_ONLY
+    unit: Unit = eqx.field(static=True, converter=u.Unit)
 
     # This is a workaround since vectorized methods don't support kwargs.
     @partial_jit()
@@ -132,6 +134,8 @@ class UserParameter(AbstractParameter):
 
     # TODO: unit handling
     func: ParameterCallable = eqx.field(static=True)
+    _: KW_ONLY
+    unit: Unit = eqx.field(static=True, converter=u.Unit)
 
     @partial_jit()
     def __call__(self, t: FloatOrIntScalar, **kwargs: Any) -> FloatArrayAnyShape:

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

@@ -1,18 +1,31 @@
 """Descriptor for a Potential Parameter."""
 
+from __future__ import annotations
+
 __all__ = ["ParameterField"]
 
 from dataclasses import KW_ONLY, dataclass, field, is_dataclass
-from typing import Annotated, Any, cast, get_args, get_origin, get_type_hints, overload
+from typing import (
+    TYPE_CHECKING,
+    Annotated,
+    Any,
+    cast,
+    get_args,
+    get_origin,
+    get_type_hints,
+    overload,
+)
 
 import astropy.units as u
 import jax.experimental.array_api as xp
 
-from galax.potential._potential.core import AbstractPotential
 from galax.typing import Unit
 
 from .core import AbstractParameter, ConstantParameter, ParameterCallable, UserParameter
 
+if TYPE_CHECKING:
+    from galax.potential._potential.base import AbstractPotentialBase
+
 
 @dataclass(frozen=True, slots=True)
 class ParameterField:
@@ -46,24 +59,28 @@
     # ===========================================
     # Descriptor
 
-    def __set_name__(self, owner: type[AbstractPotential], name: str) -> None:
+    def __set_name__(self, owner: "type[AbstractPotentialBase]", name: str) -> None:
         object.__setattr__(self, "name", name)
 
     # -----------------------------
 
     @overload  # TODO: use `Self` when beartype is happy
     def __get__(
-        self, instance: None, owner: type[AbstractPotential]
+        self, instance: None, owner: "type[AbstractPotentialBase]"
     ) -> "ParameterField":
         ...
 
     @overload
-    def __get__(self, instance: AbstractPotential, owner: None) -> AbstractParameter:
+    def __get__(
+        self, instance: "AbstractPotentialBase", owner: None
+    ) -> AbstractParameter:
         ...
 
     def __get__(  # TODO: use `Self` when beartype is happy
-        self, instance: AbstractPotential | None, owner: type[AbstractPotential] | None
-    ) -> "ParameterField | AbstractParameter":
+        self,
+        instance: "AbstractPotentialBase | None",
+        owner: "type[AbstractPotentialBase] | None",
+    ) -> ParameterField | AbstractParameter:
         # Get from class
         if instance is None:
             # If the Parameter is being set as part of a dataclass constructor,
@@ -79,7 +96,7 @@
 
     # -----------------------------
 
-    def _check_unit(self, potential: AbstractPotential, unit: Unit) -> None:
+    def _check_unit(self, potential: "AbstractPotentialBase", unit: Unit) -> None:
         """Check that the given unit is compatible with the parameter's."""
         # When the potential is being constructed, the units may not have been
         # set yet, so we don't check the unit.
@@ -99,7 +116,7 @@
 
     def __set__(
         self,
-        potential: AbstractPotential,
+        potential: "AbstractPotentialBase",
         value: AbstractParameter | ParameterCallable | Any,
     ) -> None:
         # Convert
@@ -130,6 +147,9 @@
         potential.__dict__[self.name] = value
 
 
+# -------------------------------------------
+
+
 def _get_unit_from_return_annotation(func: ParameterCallable) -> Unit:
     """Get the unit from the return annotation of a Parameter function.
 

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

@@ -0,0 +1,44 @@
+"""Descriptor for a Parameters attributes."""
+
+__all__: list[str] = []
+
+import functools
+import inspect
+import operator
+from dataclasses import Field
+from typing import TYPE_CHECKING, Any, cast
+
+if TYPE_CHECKING:
+    from galax.potential._potential.param.field import ParameterField
+
+
+def all_cls_vars(obj: object | type, /) -> dict[str, Any]:
+    """Return all variables in the whole class hierarchy."""
+    cls = obj if inspect.isclass(obj) else type(obj)
+    return functools.reduce(operator.__or__, map(vars, cls.mro()[::-1]))
+
+
+def all_parameters(obj: object | type, /) -> dict[str, "ParameterField"]:
+    """Get all fields of a dataclass, including those not-yet finalized.
+
+    Parameters
+    ----------
+    obj : object | type
+        A dataclass.
+
+    Returns
+    -------
+    dict[str, Field | Parameter]
+        All fields of the dataclass, including those not yet finalized in the class, if
+        it's still under construction, e.g. in ``__init_subclass__``.
+    """
+    from galax.potential._potential.param.field import ParameterField
+
+    return {
+        k: cast(ParameterField, v if isinstance(v, ParameterField) else v.default)
+        for k, v in all_cls_vars(obj).items()
+        if (
+            isinstance(v, ParameterField)
+            or (isinstance(v, Field) and isinstance(v.default, ParameterField))
+        )
+    }