Account for Pynbody 2.0 API change for how to use the direct gravity calculation

galpy/potential/SnapshotRZPotential.py

@@ -5,7 +5,7 @@
 import numpy
 from scipy import interpolate
 
-from ..util._optional_deps import _PYNBODY_LOADED
+from ..util._optional_deps import _PYNBODY_GE_20, _PYNBODY_LOADED
 from .interpRZPotential import (
     calc_2dsplinecoeffs_c,
     interpRZPotential,
@@ -16,7 +16,11 @@
 
 if _PYNBODY_LOADED:
     import pynbody
-    from pynbody import gravity
+
+    if _PYNBODY_GE_20:  # pragma: no cover
+        from pynbody import gravity as pynbody_gravity_calc
+    else:
+        from pynbody.gravity import calc as pynbody_gravity_calc
     from pynbody.units import NoUnit
 
 
@@ -111,7 +115,7 @@
                 ).T
 
             points_new = points.reshape(points.size // 3, 3)
-            pot, acc = gravity.calc.direct(
+            pot, acc = pynbody_gravity_calc.direct(
                 self._s, points_new, num_threads=self._num_threads
             )
 
@@ -393,7 +397,7 @@
 
         else:
             if self._interpPot:
-                pot, acc = gravity.calc.direct(
+                pot, acc = pynbody_gravity_calc.direct(
                     self._s, points_new, num_threads=self._numcores
                 )
 
@@ -426,7 +430,7 @@
 
             # first get the accelerations
             if self._interpverticalfreq:
-                zgrad_pot, zgrad_acc = gravity.calc.direct(
+                zgrad_pot, zgrad_acc = pynbody_gravity_calc.direct(
                     self._s, zgrad_points, num_threads=self._numcores
                 )
                 # each point from the points used above for pot and acc is straddled by
@@ -452,7 +456,7 @@
 
             # do the same for the radial component
             if self._interpepifreq:
-                rgrad_pot, rgrad_acc = gravity.calc.direct(
+                rgrad_pot, rgrad_acc = pynbody_gravity_calc.direct(
                     self._s, rgrad_points, num_threads=self._numcores
                 )
                 rgrad = numpy.zeros(len(points_new))

galpy/util/_optional_deps.py

@@ -1,4 +1,5 @@
 # Central place to process optional dependencies
+from packaging.version import Version
 from packaging.version import parse as parse_version
 
 # astropy
@@ -62,7 +63,12 @@
 
 # pynbody
 _PYNBODY_LOADED = True
+_PYNBODY_GE_20 = None
 try:
     import pynbody
 except ImportError:  # pragma: no cover
     _PYNBODY_LOADED = False
+else:
+    _PYNBODY_GE_20 = Version(
+        parse_version(pynbody.__version__).base_version
+    ) >= Version("2.0.0")

tests/test_snapshotpotential.py

@@ -493,6 +493,7 @@ def test_interpsnapshotKeplerPotential_normalize_units():
     s["mass"] = 4.0
     s["eps"] = 0.0
     s["pos"].units = "kpc"
+    s["eps"].units = "kpc"
     s["vel"].units = "km s**-1"
     sp = potential.InterpSnapshotRZPotential(
         s,