Add KuzminLikeWrapperPotential

HISTORY.txt

@@ -1,6 +1,16 @@
 v1.9.2 (expected around 2024-03-01)
 ===================
 
+- Added KuzminLikeWrapperPotential, a potential wrapper that allows
+  one to make a Kuzmin-like or Miyamoto-Nagai-like potential out of any
+  spherical or axisymmetric potential (evaluated in the plane, i.e.,
+  treated as a spherical potential). Kuzmin-like potentials are obtained by
+  replacing the spherical radius r with \sqrt(R^2 + (a + |z|^2)), while
+  Miyamoto-Nagai-like potentials are obtained by replacing the spherical
+  radius with \sqrt(R^2 + (a + \\sqrt(z^2 + b^2))^2). The standard KuzminDiskPotential
+  and MiyamotoNagaiPotential are obtained by applying this procedure to a point-mass
+  potential and the Kuzmin/Miyamoto-Nagai-like potentials generalize this to any
+  spherical potential.
 
 v1.9.1 (2023-11-06)
 ===================

README.dev

@@ -24,22 +24,19 @@ set up your new potential (in the 'parse_leapFuncArgs_Full' function)
 6) Edit the code in orbit/integrateFullOrbit.py to set up your
 new potential
 
-7) Edit the code in actionAngle/actionAngle_c_ext/actionAngle.c to
-parse the new potential
-
-8) Finally, add 'self.hasC= True' to the initialization of the
+7) Finally, add 'self.hasC= True' to the initialization of the
 potential in question (after the initialization of the super class)
 
-9) It should work now!
+8) It should work now!
 
-10) If you implement the second derivatives of the potential necessary
+9) If you implement the second derivatives of the potential necessary
 to integrate phase-space volumes, also set self.hasC_dxdv=True to the
 initialization of the potential in question.
 
-11) If you add a potential that gets passed to C as a list, you need
+10) If you add a potential that gets passed to C as a list, you need
 to edit orbit/integrateLinearOrbit.py and
 orbit/orbit_c_ext/integrateLinearPotential.c to parse it properly (for
 regular 3D potentials this works out of the box).
 
-12) If you add a 1D potential, do the steps above, but for
+11) If you add a 1D potential, do the steps above, but for
 integrateLinearOrbit.*

doc/source/reference/potential.rst

@@ -597,5 +597,6 @@ Specific wrappers
    potentialcorotwrapper.rst
    potentialdehnensmoothwrapper.rst
    potentialgaussampwrapper.rst
+   potentialkuzminlikewrapper.rst
    potentialsolidbodyrotationwrapper.rst
    potentialrotateandtiltwrapper.rst

doc/source/reference/potentialkuzminlikewrapper.rst

@@ -0,0 +1,5 @@
+Kuzmin-like wrapper potential
+=============================
+
+.. autoclass:: galpy.potential.KuzminLikeWrapperPotential
+   :members: __init__

galpy/orbit/integrateFullOrbit.py

@@ -466,6 +466,17 @@ def _parse_pot(pot, potforactions=False, potfortorus=False):
             pot_tfuncs.extend(wrap_pot_tfuncs)
             pot_args.append(p._amp)
             pot_tfuncs.append(p._A)
+        elif isinstance(p, potential.KuzminLikeWrapperPotential):
+            pot_type.append(-10)
+            # Not sure how to easily avoid this duplication
+            wrap_npot, wrap_pot_type, wrap_pot_args, wrap_pot_tfuncs = _parse_pot(
+                p._pot, potforactions=potforactions, potfortorus=potfortorus
+            )
+            pot_args.append(wrap_npot)
+            pot_type.extend(wrap_pot_type)
+            pot_args.extend(wrap_pot_args)
+            pot_tfuncs.extend(wrap_pot_tfuncs)
+            pot_args.extend([p._amp, p._a, p._b2])
     pot_type = numpy.array(pot_type, dtype=numpy.int32, order="C")
     pot_args = numpy.array(pot_args, dtype=numpy.float64, order="C")
     return (npot, pot_type, pot_args, pot_tfuncs)

galpy/orbit/integratePlanarOrbit.py

@@ -16,7 +16,7 @@
     planarPotentialFromFullPotential,
     planarPotentialFromRZPotential,
 )
-from ..potential.WrapperPotential import parentWrapperPotential
+from ..potential.WrapperPotential import WrapperPotential, parentWrapperPotential
 from ..util import _load_extension_libs, symplecticode
 from ..util._optional_deps import _NUMBA_LOADED, _TQDM_LOADED
 from ..util.leung_dop853 import dop853
@@ -51,9 +51,9 @@ def _parse_pot(pot):
                 isinstance(p, planarPotentialFromFullPotential)
                 or isinstance(p, planarPotentialFromRZPotential)
             )
-            and isinstance(p._Pot, parentWrapperPotential)
-        ) or isinstance(p, parentWrapperPotential):
-            if not isinstance(p, parentWrapperPotential):
+            and isinstance(p._Pot, (parentWrapperPotential, WrapperPotential))
+        ) or isinstance(p, (parentWrapperPotential, WrapperPotential)):
+            if not isinstance(p, (parentWrapperPotential, WrapperPotential)):
                 wrap_npot, wrap_pot_type, wrap_pot_args, wrap_pot_tfuncs = _parse_pot(
                     potential.toPlanarPotential(p._Pot._pot)
                 )
@@ -584,6 +584,22 @@ def _parse_pot(pot):
             pot_tfuncs.extend(wrap_pot_tfuncs)
             pot_args.append(p._amp)
             pot_tfuncs.append(p._A)
+        elif (
+            (
+                isinstance(p, planarPotentialFromFullPotential)
+                or isinstance(p, planarPotentialFromRZPotential)
+            )
+            and isinstance(p._Pot, potential.KuzminLikeWrapperPotential)
+        ) or isinstance(p, potential.KuzminLikeWrapperPotential):
+            if not isinstance(p, potential.KuzminLikeWrapperPotential):
+                p = p._Pot
+            pot_type.append(-10)
+            # wrap_pot_type, args, and npot obtained before this horrible if
+            pot_args.append(wrap_npot)
+            pot_type.extend(wrap_pot_type)
+            pot_args.extend(wrap_pot_args)
+            pot_tfuncs.extend(wrap_pot_tfuncs)
+            pot_args.extend([p._amp, p._a, p._b2])
     pot_type = numpy.array(pot_type, dtype=numpy.int32, order="C")
     pot_args = numpy.array(pot_args, dtype=numpy.float64, order="C")
     return (npot, pot_type, pot_args, pot_tfuncs)

galpy/orbit/orbit_c_ext/integrateFullOrbit.c

@@ -597,6 +597,14 @@ void parse_leapFuncArgs_Full(int npot,
       potentialArgs->ntfuncs= 1;
       potentialArgs->requiresVelocity= false;
       break;
+    case -10: // KuzminLikeWrapperPotential
+      potentialArgs->potentialEval= &KuzminLikeWrapperPotentialEval;
+      potentialArgs->Rforce= &KuzminLikeWrapperPotentialRforce;
+      potentialArgs->zforce= &KuzminLikeWrapperPotentialzforce;
+      potentialArgs->phitorque= &ZeroForce;
+      potentialArgs->nargs= 3;
+      potentialArgs->ntfuncs= 0;
+      potentialArgs->requiresVelocity= false;
     }
     int setupMovingObjectSplines = *(*pot_type-1) == -6 ? 1 : 0;
     int setupChandrasekharDynamicalFrictionSplines = *(*pot_type-1) == -7 ? 1 : 0;

galpy/orbit/orbit_c_ext/integratePlanarOrbit.c

@@ -555,6 +555,17 @@ void parse_leapFuncArgs(int npot,struct potentialArg * potentialArgs,
       potentialArgs->ntfuncs= 1;
       potentialArgs->requiresVelocity= false;
       break;
+    case -10: //KuzminLikeWrapperPotential
+      potentialArgs->potentialEval= &KuzminLikeWrapperPotentialEval;
+      potentialArgs->planarRforce= &KuzminLikeWrapperPotentialPlanarRforce;
+      potentialArgs->planarphitorque= &ZeroPlanarForce;
+      potentialArgs->planarR2deriv= &KuzminLikeWrapperPotentialPlanarR2deriv;
+      potentialArgs->planarphi2deriv= &ZeroPlanarForce;
+      potentialArgs->planarRphideriv= &ZeroPlanarForce;
+      potentialArgs->nargs= 3;
+      potentialArgs->ntfuncs= 0;
+      potentialArgs->requiresVelocity= false;
+      break;
     }
     int setupSplines = *(*pot_type-1) == -6 ? 1 : 0;
     if ( *(*pot_type-1) < 0) { // Parse wrapped potential for wrappers

galpy/potential/KuzminLikeWrapperPotential.py

@@ -0,0 +1,156 @@
+###############################################################################
+#   KuzminLikeWrapperPotential.py: Wrapper to convert a potential to a Kuzmin
+#   like potential (phi(r) --> phi(xi) where xi = sqrt(R^2 + (a + sqrt(z^2 + b^2))^2))
+###############################################################################
+import numpy
+
+from ..util import conversion
+from .Potential import (
+    _evaluatePotentials,
+    _evaluateRforces,
+    _isNonAxi,
+    evaluateR2derivs,
+)
+from .WrapperPotential import WrapperPotential
+
+
+# Only implement 3D wrapper
+class KuzminLikeWrapperPotential(WrapperPotential):
+    """Wrapper to convert a spherical potential to a Kuzmin-like potential
+
+    .. math::
+
+        \\Phi(r) \\rightarrow \\Phi(\\xi)\\,,
+
+    where
+
+    .. math::
+
+        \\xi = \\sqrt{R^2 + \\left(a + \\sqrt{z^2 + b^2}\\right)^2}\\,.
+
+    Applying this wrapper to a ``KeplerPotential`` results in the ``KuzminDiskPotential`` (for :math:`b=0`) or the ``MiyamotoNagaiPotential`` (for :math:`b \\neq 0`).
+    """
+
+    def __init__(
+        self,
+        amp=1.0,
+        a=1.1,
+        b=0.0,
+        pot=None,
+        ro=None,
+        vo=None,
+    ):
+        """
+        Initialize a KuzminLikeWrapperPotential
+
+        Parameters
+        ----------
+        amp : float, optional
+            Overall amplitude to apply to the potential. Default is 1.0.
+        a : float or Quantity, optional
+            Scale radius of the Kuzmin-like potential. Default is 1.1.
+        b : float or Quantity, optional
+            Scale height of the Kuzmin-like potential. Default is 0.0.
+        pot : Potential instance or list thereof
+            The potential to be wrapped.
+        ro : float or Quantity, optional
+            Distance scale for translation into internal units (default from configuration file).
+        vo : float or Quantity, optional
+            Velocity scale for translation into internal units (default from configuration file).
+
+        Notes
+        -----
+        - 2024-01-15 - Written - Bovy (UofT)
+        """
+        WrapperPotential.__init__(self, amp=amp, pot=pot, ro=ro, vo=vo, _init=True)
+        self._a = conversion.parse_length(a, ro=self._ro)
+        self._scale = self._a
+        self._b = conversion.parse_length(b, ro=self._ro)
+        self._b2 = self._b**2.0
+        if _isNonAxi(self._pot):
+            raise RuntimeError(
+                "KuzminLikeWrapperPotential only works for spherical or axisymmetric potentials"
+            )
+        self.hasC = True
+        self.hasC_dxdv = True
+        self.isNonAxi = False
+
+    def _xi(self, R, z):
+        return numpy.sqrt(R**2.0 + (self._a + numpy.sqrt(z**2.0 + self._b2)) ** 2.0)
+
+    def _dxidR(self, R, z):
+        return R / self._xi(R, z)
+
+    def _dxidz(self, R, z):
+        return (
+            (self._a + numpy.sqrt(z**2.0 + self._b2))
+            * z
+            / self._xi(R, z)
+            / numpy.sqrt(z**2.0 + self._b2)
+        )
+
+    def _d2xidR2(self, R, z):
+        return ((self._a + numpy.sqrt(z**2.0 + self._b2)) ** 2.0) / self._xi(
+            R, z
+        ) ** 3.0
+
+    def _d2xidz2(self, R, z):
+        return (
+            (
+                self._a**3.0 * self._b2
+                + 3.0 * self._a**2.0 * self._b2 * numpy.sqrt(self._b2 + z**2.0)
+                + self._a * self._b2 * (3.0 * self._b2 + R**2.0 + 3.0 * z**2.0)
+                + (self._b2 + R**2.0) * (self._b2 + z**2.0) ** (1.5)
+            )
+            / (self._b2 + z**2.0) ** 1.5
+            / self._xi(R, z) ** 3.0
+        )
+
+    def _d2xidRdz(self, R, z):
+        return -(R * z * (self._a + numpy.sqrt(self._b2 + z**2.0))) / (
+            numpy.sqrt(self._b2 + z**2.0)
+            * ((self._a + numpy.sqrt(self._b2 + z**2.0)) ** 2.0 + R**2.0) ** 1.5
+        )
+
+    def _evaluate(self, R, z, phi=0.0, t=0.0):
+        return _evaluatePotentials(self._pot, self._xi(R, z), 0.0, phi=phi, t=t)
+
+    def _Rforce(self, R, z, phi=0.0, t=0.0):
+        return _evaluateRforces(
+            self._pot, self._xi(R, z), 0.0, phi=phi, t=t
+        ) * self._dxidR(R, z)
+
+    def _zforce(self, R, z, phi=0.0, t=0.0):
+        return _evaluateRforces(
+            self._pot, self._xi(R, z), 0.0, phi=phi, t=t
+        ) * self._dxidz(R, z)
+
+    def _phitorque(self, R, z, phi=0.0, t=0.0):
+        return 0.0
+
+    def _R2deriv(self, R, z, phi=0.0, t=0.0):
+        return evaluateR2derivs(
+            self._pot, self._xi(R, z), 0.0, phi=phi, t=t
+        ) * self._dxidR(R, z) ** 2.0 - _evaluateRforces(
+            self._pot, self._xi(R, z), 0.0, phi=phi, t=t
+        ) * self._d2xidR2(
+            R, z
+        )
+
+    def _z2deriv(self, R, z, phi=0.0, t=0.0):
+        return evaluateR2derivs(
+            self._pot, self._xi(R, z), 0.0, phi=phi, t=t
+        ) * self._dxidz(R, z) ** 2.0 - _evaluateRforces(
+            self._pot, self._xi(R, z), 0.0, phi=phi, t=t
+        ) * self._d2xidz2(
+            R, z
+        )
+
+    def _Rzderiv(self, R, z, phi=0.0, t=0.0):
+        return evaluateR2derivs(
+            self._pot, self._xi(R, z), 0.0, phi=phi, t=t
+        ) * self._dxidR(R, z) * self._dxidz(R, z) - _evaluateRforces(
+            self._pot, self._xi(R, z), 0.0, phi=phi, t=t
+        ) * self._d2xidRdz(
+            R, z
+        )

galpy/potential/RotateAndTiltWrapperPotential.py

@@ -65,6 +65,8 @@ def __init__(
             3D vector specifying the direction of the rotated z axis.
         offset : numpy.ndarray or Quantity, optional
             Static offset in Cartesian coordinates.
+        pot : Potential instance or list thereof
+            The Potential instance or list thereof to rotate and tilt.
         ro : float or Quantity, optional
             Distance scale for translation into internal units (default from configuration file).
         vo : float or Quantity, optional

galpy/potential/__init__.py

@@ -23,6 +23,7 @@
     KingPotential,
     KuzminDiskPotential,
     KuzminKutuzovStaeckelPotential,
+    KuzminLikeWrapperPotential,
     LogarithmicHaloPotential,
     MiyamotoNagaiPotential,
     MN3ExponentialDiskPotential,
@@ -233,6 +234,7 @@
 TimeDependentAmplitudeWrapperPotential = (
     TimeDependentAmplitudeWrapperPotential.TimeDependentAmplitudeWrapperPotential
 )
+KuzminLikeWrapperPotential = KuzminLikeWrapperPotential.KuzminLikeWrapperPotential
 
 # MW potential models, now in galpy.potential.mwpotentials, but keep these two
 # for tests, backwards compatibility, and convenience