Fix pickling of states and distributions

[daniel-klein]

Not pretty, but seems to work. Would appreciate a review. Should close #257 and now we can run MultiSim!

[daniel-klein]

I understand that it's not great form to call __init__ from __setstate__, in case init changes in some unexpected way. Better practice could be to copy the relevant guts out of init to a separate function that's call both by __init__ and __setstate__, but meh?

[daniel-klein]

This is basically reimplementing __getstate__... because I don't know how to call the default state getter. Ideas?

[daniel-klein]

Again, basically just the default state setter because I'm not sure what else to do.

[daniel-klein]

Switched from parallel to MultiSim, now that we have that functionality.

[cliffckerr]

Looks good to me, but let's see if @RomeshA or @kaa2102 have any thoughts on extra Python wizardry to do here :) Otherwise, I think we can refactor further if other things break in future. Working is a big step up from not working!

[RomeshA]

I was a bit unsure what the UIDArray.__getstate__ and UIDArray.__setstate__ are needed for, is there a test case that fails if they are omitted? Is it perhaps related to the new addition of UIDArray.__getattr__()? I think something isn't quite right with that, if I run

import starsim as ss
import sciris as sc
s = ss.Sim()
s.initialize()
s2 = sc.dcp(s)

then states like s2.people.female are turned into just plain arrays. @cliffckerr do we really need to override that function entirely? It feels like it might open us up to a wide range of weird side effects further down the track

[daniel-klein]

@RomeshA - these changes are intended to enable parallel processing, e.g. MultiSim. Initially, there was a problem with forking and collecting the ScipyDistributions that prevented us from even testing parallel processing. But with those issues addressed, I encountered all sorts of issues with states, for example when results are pickled in returning from a completed sim. Try running test_parallel in test_simple.py.

[cliffckerr]

@RomeshA To fix that particular issue, we could change it to

    def __getattr__(self, attr):
        """ Make it behave like a regular array mostly -- enables things like sum(), mean(), etc. """
        if attr in ['__deepcopy__', '__setstate__']:
            return self.__getattribute__(attr)
        else:
            return getattr(self.values, attr)

Buuuuut ... this is definitely ugly and probably not super performant, though not sure if this path is encountered enough for it to matter.

[cliffckerr]

NB: on main, the copied sim can't be run. If you comment out getattr or replace it with the code above, the sim runs, but produces bizarre results (where are the recovered people going?!):

import starsim as ss
import sciris as sc

s = ss.Sim(pars=dict(diseases='sir', networks='random'))
s.initialize()

s2 = sc.dcp(s)

s.run()
s2.run()

s.plot()
s2.plot()

So in any case, something needs to be fixed.

[cliffckerr]

@daniel-klein I pushed this change to getattr so you don't need the explicit getstate/setstate. I am baffled as to why it works, though, because I thought __getattr__() is only called if __getattribute__() fails, but this calls it for those special cases and it works?!?

[RomeshA]

Those changes above should fix @cliffckerr's example - the issue was that we were inadvertently deep-copying arrays that were supposed to be references to other existing arrays, so the discrepancy would occur because the first 10 time steps would update the inadvertently copied arrays, and then when dead agents are removed at t=10, the views were re-connected but because the proper arrays hadn't been used up to that point, they would still contain the original values (e.g., everyone susceptible). So the fix is to make sure that the arrays are re-linked when copied/unpickled

[RomeshA]

Seems to all work now with the latest set of changes!