P2.1 Spherical potentials: backend-agnostic namespace-swap

galpy/potential/BurkertPotential.py

@@ -1,9 +1,9 @@
 ###############################################################################
 #   BurkertPotential.py: Potential with a Burkert density
 ###############################################################################
-import numpy
-from scipy import special
+import math
 
+from ..backend import get_namespace
 from ..util import conversion
 from .SphericalPotential import SphericalPotential
 
@@ -58,15 +58,28 @@ def __init__(self, amp=1.0, a=2.0, normalize=False, ro=None, vo=None):
 
     def _revaluate(self, r, t=0.0):
         """Potential as a function of r and time"""
+        xp = get_namespace(r)
         x = r / self.a
+        # special.xlogy(2/x, 1+x**2) == (2/x)*log(1+x**2), but with the convention
+        # that it is 0 where the prefactor is 0 (i.e. as x -> infty, where the bare
+        # product would be 0*inf = NaN). Reproduce that backend-agnostically: the
+        # prefactor 2/x only vanishes at x == infty, so guard exactly that point.
+        pref = 2.0 / x
+        # safe argument so the (dead) finite branch cannot make log(inf) at x==inf
+        safe_x2 = xp.where(xp.isinf(x), xp.ones_like(x * 1.0), x**2.0)
+        xlogy_term = xp.where(
+            xp.isinf(x),
+            xp.zeros_like(x * 1.0),
+            pref * xp.log(1.0 + safe_x2),
+        )
         return (
             -(self.a**2.0)
-            * numpy.pi
+            * math.pi
             * (
-                -numpy.pi / x
-                + 2.0 * (1.0 / x + 1) * numpy.arctan(1 / x)
-                + (1.0 / x + 1) * numpy.log((1.0 + 1.0 / x) ** 2.0 / (1.0 + 1 / x**2.0))
-                + special.xlogy(2.0 / x, 1.0 + x**2.0)
+                -math.pi / x
+                + 2.0 * (1.0 / x + 1) * xp.arctan(1 / x)
+                + (1.0 / x + 1) * xp.log((1.0 + 1.0 / x) ** 2.0 / (1.0 + 1 / x**2.0))
+                + xlogy_term
             )
         )
 
@@ -76,23 +89,24 @@ def _revaluate(self, r, t=0.0):
     #                                +(1.-x)*numpy.log(1.+x**2.))
 
     def _rforce(self, r, t=0.0):
+        xp = get_namespace(r)
         x = r / self.a
         return (
             self.a
-            * numpy.pi
+            * math.pi
             / x**2.0
             * (
-                numpy.pi
-                - 2.0 * numpy.arctan(1.0 / x)
-                - 2.0 * numpy.log(1.0 + x)
-                - numpy.log(1.0 + x**2.0)
+                math.pi
+                - 2.0 * xp.arctan(1.0 / x)
+                - 2.0 * xp.log(1.0 + x)
+                - xp.log(1.0 + x**2.0)
             )
         )
 
     def _r2deriv(self, r, t=0.0):
         x = r / self.a
         return (
-            4.0 * numpy.pi / (1.0 + x**2.0) / (1.0 + x)
+            4.0 * math.pi / (1.0 + x**2.0) / (1.0 + x)
             + 2.0 * self._rforce(r) / x / self.a
         )
 
@@ -101,34 +115,41 @@ def _rdens(self, r, t=0.0):
         return 1.0 / (1.0 + x) / (1.0 + x**2.0)
 
     def _surfdens(self, R, z, phi=0.0, t=0.0):
-        r = numpy.sqrt(R**2.0 + z**2.0)
+        xp = get_namespace(R, z)
+        r = xp.sqrt(R**2.0 + z**2.0)
         x = r / self.a
-        Rpa = numpy.sqrt(R**2.0 + self.a**2.0)
-        Rma = numpy.sqrt(R**2.0 - self.a**2.0 + 0j)
-        if Rma == 0:
-            za = z / self.a
-            return (
-                self.a**2.0
-                / 2.0
-                * (
-                    (
-                        2.0
-                        - 2.0 * numpy.sqrt(za**2.0 + 1)
-                        + numpy.sqrt(2.0) * za * numpy.arctan(za / numpy.sqrt(2.0))
-                    )
-                    / z
-                    + numpy.sqrt(2 * za**2.0 + 2.0)
-                    * numpy.arctanh(za / numpy.sqrt(2.0 * (za**2.0 + 1)))
-                    / numpy.sqrt(self.a**2.0 + z**2.0)
+        Rpa = xp.sqrt(R**2.0 + self.a**2.0)
+        # R == a is a removable singularity of the generic (Rma != 0) branch:
+        # there Rma -> 0 and the arctan/Rma terms blow up, so we use a separate
+        # closed-form limit. xp.where evaluates BOTH branches, so the generic
+        # branch must stay NaN-free at the edge: build Rma from a safe argument
+        # that is never zero there (so 1/Rma, arctan(z/x/Rma) etc. stay finite).
+        at_edge = R == self.a
+        d2 = R**2.0 - self.a**2.0
+        safe_d2 = xp.where(at_edge, xp.ones_like(d2 * 1.0), d2)
+        Rma = xp.sqrt(xp.astype(safe_d2, xp.complex128))
+        # Edge (R == a) branch
+        za = z / self.a
+        edge = (
+            self.a**2.0
+            / 2.0
+            * (
+                (
+                    2.0
+                    - 2.0 * xp.sqrt(za**2.0 + 1)
+                    + 2.0**0.5 * za * xp.arctan(za / 2.0**0.5)
                 )
+                / z
+                + xp.sqrt(2 * za**2.0 + 2.0)
+                * xp.arctanh(za / xp.sqrt(2.0 * (za**2.0 + 1)))
+                / xp.sqrt(self.a**2.0 + z**2.0)
             )
-        else:
-            return (
-                self.a**2.0
-                * (
-                    numpy.arctan(z / x / Rma) / Rma
-                    + numpy.arctanh(z / x / Rpa) / Rpa
-                    - numpy.arctan(z / Rma) / Rma
-                    + numpy.arctan(z / Rpa) / Rpa
-                ).real
-            )
+        )
+        # Generic (R != a) branch; .real of the complex combination
+        generic = self.a**2.0 * xp.real(
+            xp.arctan(z / x / Rma) / Rma
+            + xp.arctanh(z / x / Rpa) / Rpa
+            - xp.arctan(z / Rma) / Rma
+            + xp.arctan(z / Rpa) / Rpa
+        )
+        return xp.where(at_edge, edge, generic)

galpy/potential/HomogeneousSpherePotential.py

@@ -1,8 +1,9 @@
 ###############################################################################
 #   HomogeneousSpherePotential.py: The potential of a homogeneous sphere
 ###############################################################################
-import numpy
+import math
 
+from ..backend import get_namespace
 from ..util import conversion
 from .Potential import Potential
 
@@ -55,50 +56,65 @@ def __init__(self, amp=1.0, R=1.1, normalize=False, ro=None, vo=None):
         self.hasC_dens = True
 
     def _evaluate(self, R, z, phi=0.0, t=0.0):
+        xp = get_namespace(R, z)
         r2 = R**2.0 + z**2.0
-        if r2 < self._R2:
-            return r2 - 3.0 * self._R2
-        else:
-            return -2.0 * self._R3 / numpy.sqrt(r2)
+        inside = r2 < self._R2
+        # safe denominator so the (dead) outside branch cannot produce a
+        # 0/0 NaN at r2 == 0 that would poison reverse-mode gradients
+        safe = xp.where(inside, xp.ones_like(r2 * 1.0), r2)
+        return xp.where(inside, r2 - 3.0 * self._R2, -2.0 * self._R3 / xp.sqrt(safe))
 
     def _Rforce(self, R, z, phi=0.0, t=0.0):
+        xp = get_namespace(R, z)
         r2 = R**2.0 + z**2.0
-        if r2 < self._R2:
-            return -2.0 * R
-        else:
-            return -2.0 * self._R3 * R / r2**1.5
+        inside = r2 < self._R2
+        safe = xp.where(inside, xp.ones_like(r2 * 1.0), r2)
+        return xp.where(inside, -2.0 * R, -2.0 * self._R3 * R / safe**1.5)
 
     def _zforce(self, R, z, phi=0.0, t=0.0):
+        xp = get_namespace(R, z)
         r2 = R**2.0 + z**2.0
-        if r2 < self._R2:
-            return -2.0 * z
-        else:
-            return -2.0 * self._R3 * z / r2**1.5
+        inside = r2 < self._R2
+        safe = xp.where(inside, xp.ones_like(r2 * 1.0), r2)
+        return xp.where(inside, -2.0 * z, -2.0 * self._R3 * z / safe**1.5)
 
     def _R2deriv(self, R, z, phi=0.0, t=0.0):
+        xp = get_namespace(R, z)
         r2 = R**2.0 + z**2.0
-        if r2 < self._R2:
-            return 2.0
-        else:
-            return 2.0 * self._R3 / r2**1.5 - 6.0 * self._R3 * R**2.0 / r2**2.5
+        inside = r2 < self._R2
+        safe = xp.where(inside, xp.ones_like(r2 * 1.0), r2)
+        return xp.where(
+            inside,
+            2.0 * xp.ones_like(r2 * 1.0),
+            2.0 * self._R3 / safe**1.5 - 6.0 * self._R3 * R**2.0 / safe**2.5,
+        )
 
     def _z2deriv(self, R, z, phi=0.0, t=0.0):
+        xp = get_namespace(R, z)
         r2 = R**2.0 + z**2.0
-        if r2 < self._R2:
-            return 2.0
-        else:
-            return 2.0 * self._R3 / r2**1.5 - 6.0 * self._R3 * z**2.0 / r2**2.5
+        inside = r2 < self._R2
+        safe = xp.where(inside, xp.ones_like(r2 * 1.0), r2)
+        return xp.where(
+            inside,
+            2.0 * xp.ones_like(r2 * 1.0),
+            2.0 * self._R3 / safe**1.5 - 6.0 * self._R3 * z**2.0 / safe**2.5,
+        )
 
     def _Rzderiv(self, R, z, phi=0.0, t=0.0):
+        xp = get_namespace(R, z)
         r2 = R**2.0 + z**2.0
-        if r2 < self._R2:
-            return 0.0
-        else:
-            return -6.0 * self._R3 * R * z / r2**2.5
+        inside = r2 < self._R2
+        safe = xp.where(inside, xp.ones_like(r2 * 1.0), r2)
+        return xp.where(
+            inside,
+            xp.zeros_like(r2 * 1.0),
+            -6.0 * self._R3 * R * z / safe**2.5,
+        )
 
     def _dens(self, R, z, phi=0.0, t=0.0):
+        xp = get_namespace(R, z)
         r2 = R**2.0 + z**2.0
-        if r2 < self._R2:
-            return 1.5 / numpy.pi
-        else:
-            return 0.0
+        inside = r2 < self._R2
+        return xp.where(
+            inside, 1.5 / math.pi * xp.ones_like(r2 * 1.0), xp.zeros_like(r2 * 1.0)
+        )

galpy/potential/PowerSphericalPotential.py

@@ -6,9 +6,12 @@
 #                          rho(r)= ---------
 #                                   r^\alpha
 ###############################################################################
+import math
+
 import numpy
 from scipy import special
 
+from ..backend import get_namespace
 from ..util import conversion
 from .Potential import Potential
 
@@ -83,16 +86,20 @@ def _evaluate(self, R, z, phi=0.0, t=0.0):
         -----
         - Started: 2010-07-10 by Bovy (NYU)
         """
+        xp = get_namespace(R, z)
         r2 = R**2.0 + z**2.0
         if self.alpha == 2.0:
-            return numpy.log(r2) / 2.0
-        elif isinstance(r2, (float, int)) and r2 == 0 and self.alpha > 2:
-            return -numpy.inf
+            return xp.log(r2) / 2.0
+        elif self.alpha > 2:
+            # potential -> -inf at the center; use a safe r2 so the singular
+            # branch (0**negative) cannot produce a NaN/inf that poisons
+            # reverse-mode gradients, then patch in -inf where r2 == 0.
+            bad = r2 == 0.0
+            safe = xp.where(bad, xp.ones_like(r2 * 1.0), r2)
+            out = -(safe ** (1.0 - self.alpha / 2.0)) / (self.alpha - 2.0)
+            return xp.where(bad, -math.inf, out)
         else:
-            out = -(r2 ** (1.0 - self.alpha / 2.0)) / (self.alpha - 2.0)
-            if isinstance(r2, numpy.ndarray) and self.alpha > 2:
-                out[r2 == 0] = -numpy.inf
-            return out
+            return -(r2 ** (1.0 - self.alpha / 2.0)) / (self.alpha - 2.0)
 
     def _Rforce(self, R, z, phi=0.0, t=0.0):
         """
@@ -271,8 +278,9 @@ def _dens(self, R, z, phi=0.0, t=0.0):
         -----
         - 2013-01-09 - Written - Bovy (IAS)
         """
-        r = numpy.sqrt(R**2.0 + z**2.0)
-        return (3.0 - self.alpha) / 4.0 / numpy.pi / r**self.alpha
+        xp = get_namespace(R, z)
+        r = xp.sqrt(R**2.0 + z**2.0)
+        return (3.0 - self.alpha) / 4.0 / math.pi / r**self.alpha
 
     def _ddensdr(self, r, t=0.0):
         """

galpy/potential/PowerSphericalPotentialwCutoff.py

@@ -8,6 +8,7 @@
 import numpy
 from scipy import special
 
+from ..backend import get_namespace
 from ..util import conversion
 from ..util._optional_deps import _JAX_LOADED
 from .Potential import Potential, kms_to_kpcGyrDecorator
@@ -146,18 +147,20 @@ def _rforce_jax(self, r):
         )
 
     def _ddensdr(self, r, t=0.0):
+        xp = get_namespace(r)
         return (
             -self._amp
             * r ** (-1.0 - self.alpha)
-            * numpy.exp(-((r / self.rc) ** 2.0))
+            * xp.exp(-((r / self.rc) ** 2.0))
             * (2.0 * r**2.0 / self.rc**2.0 + self.alpha)
         )
 
     def _d2densdr2(self, r, t=0.0):
+        xp = get_namespace(r)
         return (
             self._amp
             * r ** (-2.0 - self.alpha)
-            * numpy.exp(-((r / self.rc) ** 2))
+            * xp.exp(-((r / self.rc) ** 2))
             * (
                 self.alpha**2.0
                 + self.alpha
@@ -199,8 +202,9 @@ def _ddenstwobetadr(self, r, beta=0):
         )
 
     def _dens(self, R, z, phi=0.0, t=0.0):
-        r = numpy.sqrt(R**2.0 + z**2.0)
-        return 1.0 / r**self.alpha * numpy.exp(-((r / self.rc) ** 2.0))
+        xp = get_namespace(R, z)
+        r = xp.sqrt(R**2.0 + z**2.0)
+        return 1.0 / r**self.alpha * xp.exp(-((r / self.rc) ** 2.0))
 
     def _mass(self, R, z=None, t=0.0):
         if z is not None: