Allow InterpSnapshotRZPotential instances to be pickled, fixes #509

galpy/potential/SnapshotRZPotential.py

@@ -1,5 +1,6 @@
 import hashlib
 from os import system
+import copy
 
 import numpy
 from scipy import interpolate
@@ -229,7 +230,7 @@ def __init__(self, s,
         self._interpverticalfreq = interpverticalfreq
 
         # make the potential accessible at points beyond the grid
-        self._origPot = SnapshotRZPotential(s, numcores)
+        self._origPot = SnapshotRZPotential(s,self._numcores)
 
         # setup the grid
         self._zsym = zsym
@@ -594,3 +595,31 @@ def denormalize(self) :
         if self._interpverticalfreq:
             self._z2interp = self._savedsplines['z2deriv']
             self._verticalfreqInterp = self._savedsplines['verticalfreq']
+
+    # Pickling functions
+    def __getstate__(self):
+        print("Here")
+        pdict= copy.copy(self.__dict__)
+        # Deconstruct _s
+        pdict['_pos']= self._s['pos']
+        pdict['_mass']= self._s['mass']
+        pdict['_eps']= self._s['eps']
+        # rm _s and _origPot,
+        del pdict['_s']
+        del pdict['_origPot']
+        return pdict
+
+    def __setstate__(self,pdict):
+        # Set up snapshot again for origPot
+        pdict['_s']= pynbody.new(star=len(pdict['_mass']))
+        pdict['_s']['pos']= pdict['_pos'] 
+        pdict['_s']['mass']= pdict['_mass'] 
+        pdict['_s']['eps']= pdict['_eps']
+        # Transfer __dict__
+        del pdict['_pos']
+        del pdict['_mass']
+        del pdict['_eps']
+        self.__dict__= pdict
+        # Now setup origPotnagain
+        self._origPot = SnapshotRZPotential(self._s,self._numcores)
+        return None

tests/test_potential.py

@@ -4501,6 +4501,24 @@ def test_diskscf_overflow():
     delta= estimateDeltaStaeckel(McMillan17,o17.R(use_physical=False),o17.z(use_physical=False))
     assert not numpy.isnan(delta), 'estimateDeltaStaeckel returns NaN due to overflow in DiskSCFPotential'
 
+def test_InterpSnapshotRZPotential_pickling():
+    # Test that InterpSnapshotRZPotential can be pickled (see #507, #509)
+    import pickle
+    import pynbody
+    from galpy.potential import InterpSnapshotRZPotential
+    # Set up simple snapshot: 1 star!
+    s= pynbody.new(star=1)
+    s['mass']= 1.
+    s['eps']= 0.
+    spi= InterpSnapshotRZPotential(s)
+    test= pickle.dumps(spi)
+    newspi= pickle.loads(test)
+    # Inside the grid
+    assert numpy.fabs(newspi(1.,0.)-spi(1.,0.)) < 1e-10, "Unpickled InterpSnapshotRZPotential does not return the same potential as original instance"
+    # Outside the grid, needs _origPot
+    assert numpy.fabs(newspi(1.,10.)-spi(1.,10.)) < 1e-10, "Unpickled InterpSnapshotRZPotential does not return the same potential as original instance"
+    return None
+
 def test_plotting():
     import tempfile