Fix spherical DF sampling for divergent potentials/densities (#812)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: caverac

galpy/df/constantbetadf.py

@@ -8,7 +8,7 @@
 from ..potential.Potential import _evaluatePotentials
 from ..util import conversion, quadpack
 from ..util._optional_deps import _JAX_LOADED
-from .sphericaldf import anisotropicsphericaldf, sphericaldf
+from .sphericaldf import _handle_rmin, anisotropicsphericaldf, sphericaldf
 
 if _JAX_LOADED:
     from jax import grad, vmap
@@ -156,6 +156,7 @@ def __init__(
         beta=0.0,
         twobeta=None,
         rmax=None,
+        rmin=None,
         scale=None,
         ro=None,
         vo=None,
@@ -175,6 +176,8 @@ def __init__(
             twice the anisotropy parameter (useful for \beta = half-integer, which is a special case); has priority over beta
         rmax : float or Quantity, optional
             maximum radius to consider; DF is cut off at E = Phi(rmax)
+        rmin : float or Quantity, optional
+            Minimum radius to consider; the distribution function is cut off at E = Phi(rmin). For potentials that diverge at r=0 (e.g., PowerSphericalPotential with alpha > 2), this is automatically set to a small value if not specified.
         scale : float or Quantity, optional
             Characteristic scale radius to aid sampling calculations. Optional and will also be overridden by value from pot if available.
         ro : float or Quantity, optional
@@ -190,7 +193,7 @@ def __init__(
         if not _JAX_LOADED:  # pragma: no cover
             raise ImportError("galpy.df.constantbetadf requires the google/jax library")
         # Parse twobeta
-        if not twobeta is None:
+        if twobeta is not None:
             beta = twobeta / 2.0
         else:
             twobeta = 2.0 * beta
@@ -210,6 +213,11 @@ def __init__(
             ro=ro,
             vo=vo,
         )
+        # Handle rmin for divergent potentials
+        self._rmin = _handle_rmin(
+            rmin, self._pot, self._denspot, self._scale, self._ro, "constantbetadf"
+        )
+
         self._twobeta = twobeta
         self._halfint = False
         if isinstance(self._twobeta, int) and self._twobeta % 2 == 1:
@@ -252,9 +260,11 @@ def __init__(
         self._gradfunc = vmap(func)
         # Min and max energy
         self._potInf = _evaluatePotentials(self._pot, self._rmax, 0)
-        self._Emin = _evaluatePotentials(self._pot, 0.0, 0)
-        # Build interpolator r(pot)
-        self._rphi = self._setup_rphi_interpolator()
+        self._Emin = _evaluatePotentials(self._pot, self._rmin, 0)
+        # Build interpolator r(pot), starting at rmin for divergent potentials
+        self._rphi = self._setup_rphi_interpolator(
+            r_a_min=max(1e-6, self._rmin / self._scale)
+        )
         # Build interpolator for the lower limit of the integration (near the
         # 1/(Phi-E)^alpha divergence; at the end, we slightly adjust it up
         # to be sure to be above the point where things go haywire...
@@ -283,9 +293,19 @@ def __init__(
                 Es, numpy.log10(startt) + 10.0 / 3.0 * (1.0 - self._alpha), k=3
             )
 
-    def sample(self, R=None, z=None, phi=None, n=1, return_orbit=True, rmin=0.0):
+    def sample(self, R=None, z=None, phi=None, n=1, return_orbit=True, rmin=None):
         # Slight over-write of superclass method to first build f(E) interp
         # No docstring so superclass' is used
+        # Use self._rmin as default if rmin is not specified
+        if rmin is None:
+            rmin = self._rmin
+        self._ensure_fE_interp()
+        return sphericaldf.sample(
+            self, R=R, z=z, phi=phi, n=n, return_orbit=return_orbit, rmin=rmin
+        )
+
+    def _ensure_fE_interp(self):
+        """Build the f(E) interpolator if not already built."""
         if not hasattr(self, "_fE_interp"):
             Es4interp = numpy.hstack(
                 (
@@ -299,9 +319,6 @@ def sample(self, R=None, z=None, phi=None, n=1, return_orbit=True, rmin=0.0):
             self._fE_interp = interpolate.InterpolatedUnivariateSpline(
                 Es4interp[iindx], fE4interp[iindx], k=3, ext=3
             )
-        return sphericaldf.sample(
-            self, R=R, z=z, phi=phi, n=n, return_orbit=return_orbit, rmin=rmin
-        )
 
     def fE(self, E):
         """

galpy/df/eddingtondf.py

@@ -6,7 +6,7 @@
 from ..potential import CompositePotential, evaluateR2derivs
 from ..potential.Potential import _evaluatePotentials, _evaluateRforces
 from ..util import conversion
-from .sphericaldf import isotropicsphericaldf, sphericaldf
+from .sphericaldf import _handle_rmin, isotropicsphericaldf, sphericaldf
 
 
 class eddingtondf(isotropicsphericaldf):
@@ -19,7 +19,9 @@ class eddingtondf(isotropicsphericaldf):
     where :math:`\\Psi = -\\Phi+\\Phi(\\infty)` is the relative potential, :math:`\\mathcal{E} = \\Psi-v^2/2` is the relative (binding) energy, and :math:`\\rho` is the density of the tracer population (not necessarily the density corresponding to :math:`\\Psi` according to the Poisson equation). Note that the second term on the right-hand side is currently assumed to be zero in the code.
     """
 
-    def __init__(self, pot=None, denspot=None, rmax=1e4, scale=None, ro=None, vo=None):
+    def __init__(
+        self, pot=None, denspot=None, rmax=1e4, rmin=None, scale=None, ro=None, vo=None
+    ):
         """
         Initialize an isotropic distribution function computed using the Eddington inversion.
 
@@ -31,6 +33,10 @@ def __init__(self, pot=None, denspot=None, rmax=1e4, scale=None, ro=None, vo=Non
             Represents the density of the tracers (assumed to be spherical; if None, set equal to pot).
         rmax : float or Quantity, optional
             Maximum radius to consider. DF is cut off at E = Phi(rmax).
+        rmin : float or Quantity, optional
+            Minimum radius to consider. For divergent potentials (Phi(0) = -inf),
+            this sets the inner boundary for the energy range. Auto-detected if
+            not specified.
         scale : float or Quantity, optional
             Characteristic scale radius to aid sampling calculations. Optional and will also be overridden by value from pot if available.
         ro : float or Quantity, optional
@@ -46,6 +52,12 @@ def __init__(self, pot=None, denspot=None, rmax=1e4, scale=None, ro=None, vo=Non
         isotropicsphericaldf.__init__(
             self, pot=pot, denspot=denspot, rmax=rmax, scale=scale, ro=ro, vo=vo
         )
+
+        # Handle rmin for divergent potentials
+        self._rmin = _handle_rmin(
+            rmin, self._pot, self._denspot, self._scale, self._ro, "eddingtondf"
+        )
+
         self._dnudr = (
             self._denspot._ddensdr
             if not isinstance(self._denspot, CompositePotential)
@@ -57,7 +69,7 @@ def __init__(self, pot=None, denspot=None, rmax=1e4, scale=None, ro=None, vo=Non
             else lambda r: numpy.sum([p._d2densdr2(r) for p in self._denspot])
         )
         self._potInf = _evaluatePotentials(pot, self._rmax, 0)
-        self._Emin = _evaluatePotentials(pot, 0.0, 0)
+        self._Emin = _evaluatePotentials(pot, self._rmin, 0)
         # Current calculation of the boundary term uses r -> inf limit
         try:
             self._rInf = (
@@ -68,12 +80,23 @@ def __init__(self, pot=None, denspot=None, rmax=1e4, scale=None, ro=None, vo=Non
             )
         except ZeroDivisionError:
             self._rInf = 1e12
-        # Build interpolator r(pot)
-        self._rphi = self._setup_rphi_interpolator()
+        # Build interpolator r(pot), starting at rmin for divergent potentials
+        self._rphi = self._setup_rphi_interpolator(
+            r_a_min=max(1e-6, self._rmin / self._scale)
+        )
 
-    def sample(self, R=None, z=None, phi=None, n=1, return_orbit=True, rmin=0.0):
+    def sample(self, R=None, z=None, phi=None, n=1, return_orbit=True, rmin=None):
         # Slight over-write of superclass method to first build f(E) interp
         # No docstring so superclass' is used
+        if rmin is None:
+            rmin = self._rmin
+        self._ensure_fE_interp()
+        return sphericaldf.sample(
+            self, R=R, z=z, phi=phi, n=n, return_orbit=return_orbit, rmin=rmin
+        )
+
+    def _ensure_fE_interp(self):
+        """Build the f(E) interpolator if not already built."""
         if not hasattr(self, "_fE_interp"):
             Es4interp = numpy.hstack(
                 (
@@ -87,9 +110,6 @@ def sample(self, R=None, z=None, phi=None, n=1, return_orbit=True, rmin=0.0):
             self._fE_interp = interpolate.InterpolatedUnivariateSpline(
                 Es4interp[iindx], fE4interp[iindx], k=3, ext=3
             )
-        return sphericaldf.sample(
-            self, R=R, z=z, phi=phi, n=n, return_orbit=return_orbit, rmin=rmin
-        )
 
     def fE(self, E):
         """

galpy/df/osipkovmerrittdf.py

@@ -200,7 +200,15 @@ class osipkovmerrittdf(_osipkovmerrittdf):
     """
 
     def __init__(
-        self, pot=None, denspot=None, ra=1.4, rmax=1e4, scale=None, ro=None, vo=None
+        self,
+        pot=None,
+        denspot=None,
+        ra=1.4,
+        rmax=1e4,
+        rmin=None,
+        scale=None,
+        ro=None,
+        vo=None,
     ):
         """
         Initialize a DF with Osipkov-Merritt anisotropy.
@@ -215,6 +223,10 @@ def __init__(
             Anisotropy radius. Default: 1.4
         rmax : float or Quantity, optional
             Maximum radius to consider; DF is cut off at E = Phi(rmax). Default: None
+        rmin : float or Quantity, optional
+            Minimum radius to consider. For divergent potentials (Phi(0) = -inf),
+            this sets the inner boundary for the energy range. Auto-detected if
+            not specified.
         scale : float or Quantity, optional
             Characteristic scale radius to aid sampling calculations. Not necessary, and will also be overridden by value from pot if available. Default: None
         ro : float or Quantity, optional
@@ -233,8 +245,16 @@ def __init__(
         # using the augmented density rawdensx(1+r^2/ra^2), we use a helper
         # eddingtondf to do this integral, hacked to use the augmented density
         self._edf = eddingtondf(
-            pot=self._pot, denspot=self._denspot, scale=scale, rmax=rmax, ro=ro, vo=vo
+            pot=self._pot,
+            denspot=self._denspot,
+            scale=scale,
+            rmax=rmax,
+            rmin=rmin,
+            ro=ro,
+            vo=vo,
         )
+        # Copy rmin from the internal eddingtondf
+        self._rmin = self._edf._rmin
         self._edf._dnudr = (
             (
                 lambda r: self._denspot._ddensdr(r) * (1.0 + r**2.0 / self._ra2)
@@ -270,9 +290,18 @@ def __init__(
             )
         )
 
-    def sample(self, R=None, z=None, phi=None, n=1, return_orbit=True, rmin=0.0):
+    def sample(self, R=None, z=None, phi=None, n=1, return_orbit=True, rmin=None):
         # Slight over-write of superclass method to first build f(Q) interp
         # No docstring so superclass' is used
+        if rmin is None:
+            rmin = self._rmin
+        self._ensure_fQ_interp()
+        return sphericaldf.sample(
+            self, R=R, z=z, phi=phi, n=n, return_orbit=return_orbit, rmin=rmin
+        )
+
+    def _ensure_fQ_interp(self):
+        """Build the f(Q) interpolator if not already built."""
         if not hasattr(self, "_logfQ_interp"):
             Qs4interp = numpy.hstack(
                 (
@@ -288,9 +317,6 @@ def sample(self, R=None, z=None, phi=None, n=1, return_orbit=True, rmin=0.0):
             self._logfQ_interp = interpolate.InterpolatedUnivariateSpline(
                 Qs4interp[iindx], fQ4interp[iindx], k=3, ext=3
             )
-        return sphericaldf.sample(
-            self, R=R, z=z, phi=phi, n=n, return_orbit=return_orbit, rmin=rmin
-        )
 
     def fQ(self, Q):
         """