🐛 fix(potential): PowerLawCutoffPotential normalization + Multipole check-init (#782)

* 🐛 fix(potential): PowerLawCutoffPotential normalization. See adrn/gala#425
* 🐛 fix(potential): Multipole check-init
* ⬆️ dep-bump: gala v1.10+
* 📝 docs(potential): improve laplacian docs

Signed-off-by: nstarman <nstarman@users.noreply.github.com>

Author: nstarman

pyproject.toml

@@ -67,7 +67,7 @@
       "galax[interop-galpy]",
     ]
     interop-astropy = ["astropy>=6.1"]
-    interop-gala = ["gala>=1.9", "galax[interop-astropy]"]
+    interop-gala = ["gala>=1.10", "galax[interop-astropy]"]
     interop-galpy = ["galax[interop-astropy]", "galpy >= 1.8"]
     # TODO: agama
     plot-all        = ["galax[plot-matplotlib]"]
@@ -231,6 +231,7 @@ ignore-words-list = "Hart"  # in dynamics/_src/cluster/relax_time.py
     "ignore:auto\\-close\\(\\)ing of figures upon backend switching is deprecated:DeprecationWarning",
     "ignore:jax\\.core\\.pp_eqn_rules is deprecated:DeprecationWarning",
     "ignore:numpy\\.ndarray size changed:RuntimeWarning",
+    "ignore:The get_name method is deprecated:astropy.utils.exceptions.AstropyDeprecationWarning",     # from gala
     "ignore:unhashable type:FutureWarning",                                                            # TODO: from diffrax
     "ignore:jax\\.interpreters\\.xla\\.pytype_aval_mappings is deprecated:DeprecationWarning",  # from tensorflow-probability[jax]
   ]

src/galax/_interop/galax_interop_gala/potential.py

@@ -69,7 +69,7 @@ def convert_potential(
 
     >>> pot = gp.KeplerPotential(m_tot=1e11, units="galactic")
     >>> gp.io.convert_potential(gp.io.GalaLibrary, pot)
-    <KeplerPotential: m=1.00e+11 (kpc,Myr,solMass,rad)>
+    <KeplerPotential: m=1.00e+11 solMass (kpc,Myr,solMass,rad)>
 
     """  # noqa: E501
     return galax_to_gala(from_)
@@ -321,7 +321,7 @@ def galax_to_gala(pot: gp.BurkertPotential, /) -> galap.BurkertPotential:
 
         >>> pot = gp.BurkertPotential(m=1e11, r_s=20, units="galactic")
         >>> gp.io.convert_potential(gp.io.GalaLibrary, pot)
-        <BurkertPotential: rho=7.82e+06, r0=20.00 (kpc,Myr,solMass,rad)>
+        <BurkertPotential: rho=7.82e+06 solMass / kpc3, r0=20.00 kpc (kpc,Myr,solMass,rad)>
 
         .. skip: end
 
@@ -379,7 +379,7 @@ def galax_to_gala(
 
     >>> pot = gp.HarmonicOscillatorPotential(omega=u.Quantity(1, "1/Myr"), units="galactic")
     >>> gp.io.convert_potential(gp.io.GalaLibrary, pot)
-    <HarmonicOscillatorPotential: omega=[1.] (kpc,Myr,solMass,rad)>
+    <HarmonicOscillatorPotential: omega=[1.] 1 / Myr (kpc,Myr,solMass,rad)>
 
     """  # noqa: E501
     _error_if_not_all_constant_parameters(pot, *pot.parameters.keys())
@@ -433,7 +433,7 @@ def galax_to_gala(pot: gp.HernquistPotential, /) -> galap.HernquistPotential:
 
     >>> pot = gp.HernquistPotential(m_tot=1e11, r_s=20, units="galactic")
     >>> gp.io.convert_potential(gp.io.GalaLibrary, pot)
-    <HernquistPotential: m=1.00e+11, c=20.00 (kpc,Myr,solMass,rad)>
+    <HernquistPotential: m=1.00e+11 solMass, c=20.00 kpc (kpc,Myr,solMass,rad)>
 
     """  # noqa: E501
     _error_if_not_all_constant_parameters(pot, *pot.parameters.keys())
@@ -489,7 +489,7 @@ def galax_to_gala(pot: gp.IsochronePotential, /) -> galap.IsochronePotential:
 
     >>> pot = gp.IsochronePotential(m_tot=1e11, r_s=10, units="galactic")
     >>> gp.io.convert_potential(gp.io.GalaLibrary, pot)
-    <IsochronePotential: m=1.00e+11, b=10.00 (kpc,Myr,solMass,rad)>
+    <IsochronePotential: m=1.00e+11 solMass, b=10.00 kpc (kpc,Myr,solMass,rad)>
 
     """  # noqa: E501
     _error_if_not_all_constant_parameters(pot, *pot.parameters.keys())
@@ -549,7 +549,7 @@ def galax_to_gala(pot: gp.JaffePotential, /) -> galap.JaffePotential:
 
     >>> pot = gp.JaffePotential(m_tot=1e11, r_s=20, units="galactic")
     >>> gp.io.convert_potential(gp.io.GalaLibrary, pot)
-    <JaffePotential: m=1.00e+11, c=20.00 (kpc,Myr,solMass,rad)>
+    <JaffePotential: m=1.00e+11 solMass, c=20.00 kpc (kpc,Myr,solMass,rad)>
 
     """
     _error_if_not_all_constant_parameters(pot, *pot.parameters.keys())
@@ -603,7 +603,7 @@ def galax_to_gala(pot: gp.KeplerPotential, /) -> galap.KeplerPotential:
 
     >>> pot = gp.KeplerPotential(m_tot=1e11, units="galactic")
     >>> gp.io.convert_potential(gp.io.GalaLibrary, pot)
-    <KeplerPotential: m=1.00e+11 (kpc,Myr,solMass,rad)>
+    <KeplerPotential: m=1.00e+11 solMass (kpc,Myr,solMass,rad)>
 
     """  # noqa: E501
     _error_if_not_all_constant_parameters(pot, *pot.parameters.keys())
@@ -663,7 +663,7 @@ def galax_to_gala(pot: gp.KuzminPotential, /) -> galap.KuzminPotential:
 
     >>> pot = gp.KuzminPotential(m_tot=1e11, r_s=20, units="galactic")
     >>> gp.io.convert_potential(gp.io.GalaLibrary, pot)
-    <KuzminPotential: m=1.00e+11, a=20.00 (kpc,Myr,solMass,rad)>
+    <KuzminPotential: m=1.00e+11 solMass, a=20.00 kpc (kpc,Myr,solMass,rad)>
 
     """  # noqa: E501
     _error_if_not_all_constant_parameters(pot, *pot.parameters.keys())
@@ -735,7 +735,7 @@ def galax_to_gala(pot: gp.LongMuraliBarPotential, /) -> galap.LongMuraliBarPoten
     ...     units="galactic",
     ... )
     >>> gp.io.convert_potential(gp.io.GalaLibrary, pot)
-    <LongMuraliBarPotential: m=1.00e+11, a=20.00, b=10.00, c=5.00, alpha=0.10 (kpc,Myr,solMass,rad)>
+    <LongMuraliBarPotential: m=1.00e+11 solMass, a=20.00 kpc, b=10.00 kpc, c=5.00 kpc, alpha=0.10 rad (kpc,Myr,solMass,rad)>
 
     """  # noqa: E501
     _error_if_not_all_constant_parameters(pot, *pot.parameters.keys())
@@ -795,7 +795,7 @@ def galax_to_gala(pot: gp.MiyamotoNagaiPotential, /) -> galap.MiyamotoNagaiPoten
 
     >>> pot = gp.MiyamotoNagaiPotential(m_tot=1e11, a=6.5, b=0.26, units="galactic")
     >>> gp.io.convert_potential(gp.io.GalaLibrary, pot)
-    <MiyamotoNagaiPotential: m=1.00e+11, a=6.50, b=0.26 (kpc,Myr,solMass,rad)>
+    <MiyamotoNagaiPotential: m=1.00e+11 solMass, a=6.50 kpc, b=0.26 kpc (kpc,Myr,solMass,rad)>
 
     """  # noqa: E501
     _error_if_not_all_constant_parameters(pot, *pot.parameters.keys())
@@ -861,7 +861,7 @@ def galax_to_gala(
 
     >>> pot = gp.MN3ExponentialPotential(m_tot=1e11, h_R=3.0, h_z=0.2, units="galactic")
     >>> gp.io.convert_potential(gp.io.GalaLibrary, pot)
-    <MN3ExponentialDiskPotential: m=1.00e+11, h_R=3.00, h_z=0.20 (kpc,Myr,solMass,rad)>
+    <MN3ExponentialDiskPotential: m=1.00e+11 solMass, h_R=3.00 kpc, h_z=0.20 kpc (kpc,Myr,solMass,rad)>
 
     """  # noqa: E501
     _error_if_not_all_constant_parameters(pot, *pot.parameters.keys())
@@ -969,7 +969,7 @@ def galax_to_gala(pot: gp.PlummerPotential, /) -> galap.PlummerPotential:
 
     >>> pot = gp.PlummerPotential(m_tot=1e11, r_s=10, units="galactic")
     >>> gp.io.convert_potential(gp.io.GalaLibrary, pot)
-    <PlummerPotential: m=1.00e+11, b=10.00 (kpc,Myr,solMass,rad)>
+    <PlummerPotential: m=1.00e+11 solMass, b=10.00 kpc (kpc,Myr,solMass,rad)>
     """  # noqa: E501
     _error_if_not_all_constant_parameters(pot, *pot.parameters.keys())
 
@@ -1044,7 +1044,7 @@ def galax_to_gala(pot: gp.PowerLawCutoffPotential, /) -> galap.PowerLawCutoffPot
 
     >>> pot = gp.PowerLawCutoffPotential(m_tot=1e11, alpha=1.8, r_c=20, units="galactic")
     >>> gp.io.convert_potential(gp.io.GalaLibrary, pot)
-    <PowerLawCutoffPotential: m=1.00e+11, alpha=1.80, r_c=20.00 (kpc,Myr,solMass,rad)>
+    <PowerLawCutoffPotential: m=1.00e+11 solMass, alpha=1.80 , r_c=20.00 kpc (kpc,Myr,solMass,rad)>
 
     .. skip: end
     """  # noqa: E501
@@ -1106,7 +1106,7 @@ def galax_to_gala(pot: gp.SatohPotential, /) -> galap.SatohPotential:
 
     >>> pot = gp.SatohPotential(m_tot=1e11, a=20, b=10, units="galactic")
     >>> gp.io.convert_potential(gp.io.GalaLibrary, pot)
-    <SatohPotential: m=1.00e+11, a=20.00, b=10.00 (kpc,Myr,solMass,rad)>
+    <SatohPotential: m=1.00e+11 solMass, a=20.00 kpc, b=10.00 kpc (kpc,Myr,solMass,rad)>
 
     """
     _error_if_not_all_constant_parameters(pot, *pot.parameters.keys())
@@ -1163,7 +1163,7 @@ def galax_to_gala(pot: gp.StoneOstriker15Potential, /) -> galap.StonePotential:
 
     >>> pot = gp.StoneOstriker15Potential(m_tot=1e11, r_c=20, r_h=10, units="galactic")
     >>> gp.io.convert_potential(gp.io.GalaLibrary, pot)
-    <StonePotential: m=1.00e+11, r_c=20.00, r_h=10.00 (kpc,Myr,solMass,rad)>
+    <StonePotential: m=1.00e+11 solMass, r_c=20.00 kpc, r_h=10.00 kpc (kpc,Myr,solMass,rad)>
 
     """  # noqa: E501
     _error_if_not_all_constant_parameters(pot, *pot.parameters.keys())
@@ -1243,7 +1243,7 @@ def galax_to_gala(pot: gp.LogarithmicPotential, /) -> galap.LogarithmicPotential
 
     >>> pot = gp.LogarithmicPotential(v_c=u.Quantity(220, "km/s"), r_s=20, units="galactic")
     >>> gp.io.convert_potential(gp.io.GalaLibrary, pot)
-    <LogarithmicPotential: v_c=0.22, r_h=20.00, q1=1.00, q2=1.00, q3=1.00, phi=0 (kpc,Myr,solMass,rad)>
+    <LogarithmicPotential: v_c=0.22 kpc / Myr, r_h=20.00 kpc, q1=1.00 , q2=1.00 , q3=1.00 , phi=0 rad (kpc,Myr,solMass,rad)>
 
     """  # noqa: E501
     _error_if_not_all_constant_parameters(pot, *pot.parameters.keys())
@@ -1270,7 +1270,7 @@ def galax_to_gala(pot: gp.LMJ09LogarithmicPotential, /) -> galap.LogarithmicPote
     ...     units="galactic",
     ... )
     >>> gp.io.convert_potential(gp.io.GalaLibrary, pot)
-    <LogarithmicPotential: v_c=0.22, r_h=20.00, q1=1.00, q2=1.00, q3=1.00, phi=0 (kpc,Myr,solMass,rad)>
+    <LogarithmicPotential: v_c=0.22 kpc / Myr, r_h=20.00 kpc, q1=1.00 , q2=1.00 , q3=1.00 , phi=0 rad (kpc,Myr,solMass,rad)>
 
     """  # noqa: E501
     _error_if_not_all_constant_parameters(pot, *pot.parameters.keys())
@@ -1466,7 +1466,8 @@ def galax_to_gala(pot: gp.NFWPotential, /) -> galap.NFWPotential:
 
     >>> pot = gp.NFWPotential(m=1e12, r_s=20, units="galactic")
     >>> gp.io.convert_potential(gp.io.GalaLibrary, pot)
-    <NFWPotential: m=1.00e+12, r_s=20.00, a=1.00, b=1.00, c=1.00 (kpc,Myr,solMass,rad)>
+    <NFWPotential: m=1.00e+12 solMass, r_s=20.00 kpc, a=1.00 , b=1.00 , c=1.00
+     (kpc,Myr,solMass,rad)>
 
     """
     _error_if_not_all_constant_parameters(pot, *pot.parameters.keys())
@@ -1534,7 +1535,7 @@ def galax_to_gala(
     ...                                      a1=1, a2=0.9, a3=0.8,
     ...                                      units="galactic")
     >>> gp.io.convert_potential(gp.io.GalaLibrary, pot)
-    <LeeSutoTriaxialNFWPotential: v_c=0.47, r_s=20.00, a=1.00, b=0.90, c=0.80 (kpc,Myr,solMass,rad)>
+    <LeeSutoTriaxialNFWPotential: v_c=0.47 kpc / Myr, r_s=20.00 kpc, a=1.00 , b=0.90 , c=0.80  (kpc,Myr,solMass,rad)>
 
     """  # noqa: E501
     _error_if_not_all_constant_parameters(pot, *pot.parameters.keys())

src/galax/potential/_src/api.py

@@ -383,48 +383,48 @@ def laplacian(*args: Any, **kwargs: Any) -> u.Quantity["1/s^2"] | Array:
     >>> w = gc.PhaseSpaceCoordinate(q=u.Quantity([[1, 2, 3], [4, 5, 6]], "kpc"),
     ...                             p=u.Quantity([[4, 5, 6], [7, 8, 9]], "km/s"),
     ...                             t=u.Quantity([0, 1], "Gyr"))
-    >>> pot.laplacian(w)
-    Quantity(Array([2.77555756e-17, 8.67361738e-19], dtype=float64), unit='1 / Myr2')
+    >>> pot.laplacian(w).round(4)
+    Quantity(Array([0., 0.], dtype=float64), unit='1 / Myr2')
 
     This function is very flexible and can accept a broad variety of inputs. For
     example, instead of passing a
     `galax.coordinates.PhaseSpaceCoordinate`, we can instead pass a
     `~vector.FourVector`:
 
     >>> w = cx.FourVector(q=u.Quantity([1, 2, 3], "kpc"), t=u.Quantity(0, "Gyr"))
-    >>> pot.laplacian(w)
-    Quantity(Array(2.77555756e-17, dtype=float64), unit='1 / Myr2')
+    >>> pot.laplacian(w).round(4)
+    Quantity(Array(0., dtype=float64), unit='1 / Myr2')
 
     Or using a `~coordinax.AbstractPos3D` and time `unxt.Quantity` (which can be
     positional or a keyword argument):
 
     >>> q = cx.CartesianPos3D.from_([1, 2, 3], "kpc")
     >>> t = u.Quantity(0, "Gyr")
-    >>> pot.laplacian(q, t=t)
-    Quantity(Array(2.77555756e-17, dtype=float64), unit='1 / Myr2')
+    >>> pot.laplacian(q, t=t).round(4)
+    Quantity(Array(0., dtype=float64), unit='1 / Myr2')
 
     We can also compute the potential energy at multiple positions:
 
     >>> q = cx.CartesianPos3D.from_([[1, 2, 3], [4, 5, 6]], "kpc")
-    >>> pot.laplacian(q, t=t)
-    Quantity(Array([2.77555756e-17, 8.67361738e-19], dtype=float64), unit='1 / Myr2')
+    >>> pot.laplacian(q, t=t).round(4)
+    Quantity(Array([0., 0.], dtype=float64), unit='1 / Myr2')
 
     Instead of passing a `~coordinax.AbstractPos3D` (in this case a
     `~coordinax.CartesianPos3D`), we can instead pass a
     `unxt.Quantity`, which is interpreted as a Cartesian position:
 
     >>> q = u.Quantity([1., 2, 3], "kpc")
-    >>> pot.laplacian(q, t)
-    Quantity(Array(2.77555756e-17, dtype=float64), unit='1 / Myr2')
+    >>> pot.laplacian(q, t).round(4)
+    Quantity(Array(0., dtype=float64), unit='1 / Myr2')
 
     Again, this can be batched.  If the input position object has no units (i.e.
     is an `~jax.Array`), it is assumed to be in the same unit system as the
     potential.
 
     >>> q = jnp.asarray([[1, 2, 3], [4, 5, 6]])
     >>> t = 0
-    >>> pot.laplacian(q, t)
-    Array([2.77555756e-17, 8.67361738e-19], dtype=float64)
+    >>> pot.laplacian(q, t).round(4)
+    Array([0., 0.], dtype=float64)
 
     - - -
 
@@ -454,25 +454,25 @@ def laplacian(*args: Any, **kwargs: Any) -> u.Quantity["1/s^2"] | Array:
     We can also compute the potential energy at multiple positions:
 
     >>> q = apyc.CartesianRepresentation(apyu.Quantity([[1, 4], [2, 5], [3, 6]], "kpc"))
-    >>> pot.laplacian(q, t)
-    Quantity(Array([2.77555756e-17, 8.67361738e-19], dtype=float64), unit='1 / Myr2')
+    >>> pot.laplacian(q, t).round(4)
+    Quantity(Array([0., 0.], dtype=float64), unit='1 / Myr2')
 
     Instead of passing a `~coordinax.AbstractPos3D` (in this case a
     `~coordinax.CartesianPos3D`), we can instead pass a
     `unxt.Quantity`, which is interpreted as a Cartesian position:
 
     >>> q = [1., 2, 3] * apyu.kpc
-    >>> pot.laplacian(q, t)
-    Array(2.77555756e-17, dtype=float64)
+    >>> pot.laplacian(q, t).round(4)
+    Array(0., dtype=float64)
 
     Again, this can be batched.  If the input position object has no units (i.e.
     is a `~numpy.ndarray`), it is assumed to be in the same unit system
     as the potential.
 
     >>> import numpy as np
     >>> q = jnp.asarray([[1, 2, 3], [4, 5, 6]])
-    >>> pot.laplacian(q, t)
-    Array([2.77555756e-17, 8.67361738e-19], dtype=float64)
+    >>> pot.laplacian(q, t).round(4)
+    Array([0., 0.], dtype=float64)
 
     .. skip: end
 

src/galax/potential/_src/builtin/multipole.py

@@ -209,13 +209,11 @@ class MultipolePotential(AbstractMultipolePotential):
     def __check_init__(self) -> None:
         shape = (self.l_max + 1, self.l_max + 1)
         t = u.Quantity(0.0, "Gyr")
-        is_shape, it_shape = self.ISlm(t).shape, self.ITlm(t).shape
-        os_shape, ot_shape = self.OSlm(t).shape, self.OTlm(t).shape
-        # TODO: check shape across time.
+        iss, its = self.ISlm(t).shape, self.ITlm(t).shape
+        oss, ots = self.OSlm(t).shape, self.OTlm(t).shape
+        # Check shapes match expected
         msg = "I/OSlm and I/OTlm must have the shape (l_max + 1, l_max + 1)."
-        pred = jnp.any(
-            jnp.array([x != shape for x in (is_shape, it_shape, os_shape, ot_shape)])
-        )
+        pred = (iss != shape) or (its != shape) or (oss != shape) or (ots != shape)
         _ = eqx.error_if(t, pred, msg)
 
     @ft.partial(jax.jit)

src/galax/potential/_src/builtin/powerlawcutoff.py

@@ -88,15 +88,29 @@ def _potential(self, xyz: gt.BBtQuSz3, t: gt.BBtQuSz0, /) -> gt.BtSz0:
 def potential(p: gt.Params, r: gt.Sz0, /) -> gt.FloatSz0:
     r"""Potential for the power-law cutoff density profile.
 
-    $$ \Phi(r) = -\frac{G M}{2\pi \Gamma\left(\frac{3 - \alpha}{2}\right) r_c^3}
-    \left(\frac{r_c}{r}\right)^\alpha
-    \exp\left[-\left(\frac{r}{r_c}\right)^2\right] $$
+    Following Gala's implementation, we compute the potential and subtract
+    the asymptotic value to enforce Phi(infinity) = 0.
     """
-    a = p["alpha"] / 2
+    alpha_half = p["alpha"] / 2
     s2 = (r / p["r_c"]) ** 2
     GM = p["G"] * p["m_tot"]
 
-    return GM * (
-        (a - 1.5) * _safe_gamma_inc(1.5 - a, s2) / (r * jsp.gamma(2.5 - a))
-        + _safe_gamma_inc(1 - a, s2) / (p["r_c"] * jsp.gamma(1.5 - a))
+    # Compute potential terms
+    gamma_arg = 1.5 - alpha_half
+    term1 = (
+        GM
+        * _safe_gamma_inc(gamma_arg, s2)
+        * (alpha_half - 1.5)
+        / (r * jsp.gamma(2.5 - alpha_half))
     )
+    term2 = GM * _safe_gamma_inc(1 - alpha_half, s2) / (p["r_c"] * jsp.gamma(gamma_arg))
+
+    # Subtract asymptotic value to enforce Phi(infinity) = 0
+    # Only needed when gamma_arg > 0 (i.e., alpha < 3)
+    phi_infinity = jax.lax.cond(
+        gamma_arg > 0,
+        lambda: GM * jsp.gamma(1 - alpha_half) / (p["r_c"] * jsp.gamma(gamma_arg)),
+        lambda: 0.0,
+    )
+
+    return term1 + term2 - phi_infinity

tests/unit/potential/builtin/test_bovymwpotential2014.py

@@ -29,7 +29,7 @@ def pot_map(
     # ==========================================================================
 
     def test_potential(self, pot: gp.BovyMWPotential2014, x: gt.QuSz3) -> None:
-        expect = u.Quantity(-0.09550731, unit="kpc2 / Myr2")
+        expect = u.Quantity(-0.16359185, unit="kpc2 / Myr2")
         assert jnp.isclose(
             pot.potential(x, t=0), expect, atol=u.Quantity(1e-8, expect.unit)
         )

tests/unit/potential/builtin/test_multipole.py

@@ -261,7 +261,7 @@ def test_check_init(
         """Test the `MultipoleInnerPotential.__check_init__` method."""
         fields_["ISlm"] = fields_["ISlm"][::2]  # make it the wrong shape
         match = re.escape("I/OSlm and I/OTlm must have the shape")
-        with pytest.raises(eqx.EquinoxRuntimeError, match=match):
+        with pytest.raises(eqx.EquinoxTracetimeError, match=match):
             pot_cls(**fields_)
 
     # ==========================================================================

tests/unit/potential/builtin/test_powerlawcutoff.py

@@ -97,7 +97,7 @@ def fields_(
     # ==========================================================================
 
     def test_potential(self, pot: gp.PowerLawCutoffPotential, x: gt.QuSz3) -> None:
-        expect = u.Quantity(6.26573365, unit="kpc2 / Myr2")
+        expect = u.Quantity(-1.20227523, unit="kpc2 / Myr2")
         assert jnp.isclose(
             pot.potential(x, t=0), expect, atol=u.Quantity(1e-8, expect.unit)
         )