Issue: TypeError: can't pickle _thread.RLock objects - when pickling InterpolatedSnapshotRZPotential Follow-Up

[smandhai]

Original Query (originally asked on Slack but now solved):

I was wondering if there's anyway to pickle/dump a InterpolatedSnapshotRZPotential object from a pynbody processed galaxy simulation?

I've tried this for other potentials like the preset MW galpy potential and those work fine (can be dumped and recovered). But for this particular interpolated potential, I get the traceback:
TypeError: can't pickle _thread.RLock objects.

Does this potential use multiprocessing/threading? If so, are there any arguments that can be used to bypass this?
Context: Building on my previous post, I've now created the potentials I need but to avoid recalculating the potentials each time, I'm attempting to dump them and then recover them in another script.

Suggested Solution (from @jobovy ):
Delete '_s' and '_origPot' from the potential instance using the following example -

import pynbody s= pynbody.new(star=1) s['mass']= 1. s['eps']= 0. spi= InterpSnapshotRZPotential(s) delattr(spi,'_s') delattr(spi,'_origPot') test= pickle.dumps(spi) # Works without error newspi= pickle.loads(test) newspi(1.,0.), spi(1.,0.)

This worked and solved the pickling issue!

Follow-up Query:
However, I'm not 100% sure how to restore the attributes once the potentials have been reloaded and unpickled. I imagine it requires reloading the galaxy simulation with pynbody, but I'm not too sure how to relink the "_origPot" attribute in particular.

[jobovy]

Thanks for posting! As you mentioned (but repeated here so it's easy to find for others), you can pickle InterpolatedSnapshotRZPotential instances by deleting _s and _origPot from the potential instance. For example,

import pynbody
s= pynbody.new(star=1)
s['mass']= 1.
s['eps']= 0.
spi= InterpSnapshotRZPotential(s)
delattr(spi,'_s')
delattr(spi,'_origPot')
test= pickle.dumps(spi) # Works without error 
newspi= pickle.loads(test)
print(newspi(1.,0.), spi(1.,0.))

Follow-up Query

Yes, you cannot fully restore the InterpSnapshotRZPotential in this way. To get _s and _origPot back, you'd have to reload the pynbody snapshot, but that will probably take a long time. Perhaps you can pickle it separately in a way that quickly allows you to re-setup _s. _origPot is just defined as

_origPot = SnapshotRZPotential(s, numcores)

so if you can restore s, then you can also re-initialize origPot.

[jobovy]

I should also note that as long as you don't venture outside of the grid boundaries, you don't need _s or _origPot. So if you set up the InterpSnapshotRZPotential instance with a wide enough grid, it should be fine to ignore those missing attributes (the code should crash if it needs them).

[smandhai]

Excellent, thank you so much, Jo! My current issue is that I'm orbiting systems within the galaxy interpolated potentials which sometimes escape beyond the gridsize (the particular systems I'm looking at can escape beyond ~few x 100 kpc from the galaxy in rare cases). At any rate, I'll try to pickle the necessary components and get them restored in the way that you've suggested.

[jobovy]

Let me know if you find anything that works to pickle those components. If you do, we might be able to just incorporate this into galpy.

A quick hacky fix could also be to just set _origPot to a KeplerPotential with the same mass as that of the total halo, because the potential is probably close to that of a point mass at those distances? Anyway, might be interesting to experiment with that a bit...

[jobovy]

Looking at the relevant pynbody code, all you really need for the gravity calculation are the 'pos', 'mass', and 'eps' attributes of the snapshot, so it should be possible to just pickle those and then after unpickling them do something like

s= {'pos': unpickled_pos, 'mass':unpickled_mass,'eps':unpickled_eps}
newspi ._s= s
newspi._origPot= SnapshotRZPotential(s)

[jobovy]

Okay, I decided to fix this, see #509. Pickling/unpickling InterpSnapshotRZPotential instances should now just work in the latest main branch.

[smandhai]

Thank you so much, Jo! Apologies for the late response, but I'll feed this into my scripts (following a package update) and report back if there are any issues but I suspect this should be okay now.

[smandhai]

Finally got around to updating everything to the latest Python + galpy - and I noticed that there's a rogue print("Here") in SnapshotRZPotential.py in the pickling routine. It's working great though - can confirm that the pickling is going smoothly without the need to manually load in dictionaries (which now saves storage space too) and redefining _s and _origPot.

[jobovy]

Great, thanks! And thanks for letting me know about the errant print statement, I'll remove it.