Addresssing #257

starsim/distributions.py

@@ -21,7 +21,7 @@
 
 class ScipyDistribution():
     def __init__(self, gen, rng=None):
-        self.gen = None
+        self._gen = gen
         class starsim_gen(type(gen.dist)):
             def __init__(self, *args, **kwargs):
                 super().__init__(*args, **kwargs)
@@ -198,16 +198,19 @@ def __deepcopy__(self, memo):
         for k, v in self.__dict__.items():
             setattr(result, k, deepcopy(v, memo))
         return result
+
+    def __getstate__(self):
+        dct = self.__dict__.copy()
+        dct.pop('gen')
+        return dct
+
+    def __setstate__(self, state):
+        self.__init__(state['_gen'], state['rng'])
+        return
 
     def __getattr__(self, attr):
         # Returns wrapped generator.(attr) if not a property
-        if attr == '__getstate__':
-            # Must be from pickle, return a callable function that returns None
-            return lambda: None
-        #elif attr == '__deepcopy__':
-        #    return self
-        elif attr in ['__setstate__', '__await__']:
-            # Must be from pickle, async programming, copy
+        if attr in ['__await__']:
             return None
         try:
             return self.__getattribute__(attr)

tests/test_simple.py

@@ -126,10 +126,10 @@ def test_components():
 def test_parallel():
     """ Test running two identical sims in parallel """
     pars = make_sim_pars()
-    s1 = ss.Sim(pars)
-    s2 = ss.Sim(pars)
-    # s1, s2 = ss.parallel(s1, s2).sims
-    # assert np.allclose(s1.summary[:], s2.summary[:], rtol=0, atol=0, equal_nan=True)
+    sims = ss.MultiSim([ss.Sim(pars, label='Sim1'), ss.Sim(pars, label='Sim2')])
+    sims.run()
+    s1, s2 = sims.sims
+    assert np.allclose(s1.summary[:], s2.summary[:], rtol=0, atol=0, equal_nan=True)
     return s1, s2