From 2d78942a57a5de672ebbad1736f4c7a019702163 Mon Sep 17 00:00:00 2001
From: Jo Bovy <bovy@astro.utoronto.ca>
Date: Fri, 13 Sep 2024 15:10:45 -0400
Subject: [PATCH] Account for Pynbody 2.0 API change for how to use the direct
 gravity calculation

---
 galpy/potential/SnapshotRZPotential.py | 16 ++++++++++------
 galpy/util/_optional_deps.py           |  6 ++++++
 tests/test_snapshotpotential.py        |  1 +
 3 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/galpy/potential/SnapshotRZPotential.py b/galpy/potential/SnapshotRZPotential.py
index af6e19046..87a307ae2 100644
--- a/galpy/potential/SnapshotRZPotential.py
+++ b/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 @@ def _setup_potential(self, R, z, use_pkdgrav=False):
                 ).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 @@ def _setup_potential(self, R, z, use_pkdgrav=False, dr=0.0001):
 
         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 @@ def _setup_potential(self, R, z, use_pkdgrav=False, dr=0.0001):
 
             # 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 @@ def _setup_potential(self, R, z, use_pkdgrav=False, dr=0.0001):
 
             # 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))
diff --git a/galpy/util/_optional_deps.py b/galpy/util/_optional_deps.py
index 8a6d7e7b8..74334110d 100644
--- a/galpy/util/_optional_deps.py
+++ b/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")
diff --git a/tests/test_snapshotpotential.py b/tests/test_snapshotpotential.py
index 07128e261..cec965098 100644
--- a/tests/test_snapshotpotential.py
+++ b/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,