[DRAFT] Coherent random numbers

.gitignore

@@ -1,6 +1,8 @@
 # HPVsim-generated files
 *.obj
 files/
+tests/figs/*
+figs/*
 modules.rst
 
 # Other files

stisim/__init__.py

@@ -4,6 +4,7 @@
 from .utils         import *
 from .distributions import *
 from .states        import *
+from .streams       import *
 from .people        import *
 from .networks      import *
 from .modules       import *

stisim/analyzers.py

@@ -11,5 +11,4 @@
 class Analyzer(ss.Module):
 
     def update_results(self, sim):
-        raise NotImplementedError
-
+        raise NotImplementedError

stisim/demographics.py

@@ -241,6 +241,14 @@ def __init__(self, pars=None):
             'initial': 3,  # Number of women initially pregnant
         }, self.pars)
 
+        self.rng_female = ss.Stream(f'female_{self.name}')
+        self.female_dist = ss.bernoulli(p=0.5, rng=self.rng_female) # Replace 0.5 with sex ratio at birth
+
+        self.rng_conception = ss.Stream('conception')
+        self.rng_dead = ss.Stream(f'dead_{self.name}')
+
+        self.rng_uids = ss.Stream(f'uids_{self.name}')
+
         return
 
     def initialize(self, sim):
@@ -305,14 +313,20 @@ def make_pregnancies(self, sim):
         if self.pars.inci > 0:
             denom_conds = ppl.female & ppl.active & self.susceptible
             inds_to_choose_from = ss.true(denom_conds)
-            uids = ss.binomial_filter(self.pars.inci, inds_to_choose_from)
+            uids = self.rng_conception.bernoulli_filter(uids=inds_to_choose_from, prob=self.pars.inci)
 
             # Add UIDs for the as-yet-unborn agents so that we can track prognoses and transmission patterns
             n_unborn_agents = len(uids)
             if n_unborn_agents > 0:
+                new_slots = self.rng_uids.integers(low=sim.pars['n_agents'], high=sim.pars['slot_scale']*sim.pars['n_agents'], uids=uids, dtype=int)
+
                 # Grow the arrays and set properties for the unborn agents
-                new_uids = sim.people.grow(n_unborn_agents)
+                new_uids = sim.people.grow(len(new_slots))
+
                 sim.people.age[new_uids] = -self.pars.dur_pregnancy
+                #sim.people.female[new_uids] = self.rng_sex.bernoulli(uids=uids, prob=0.5) # Replace 0.5 with sex ratio at birth
+                sim.people.slot[new_uids] = new_slots # Before sampling female_dist
+                sim.people.female[new_uids] = self.female_dist.sample(uids=uids)
 
                 # Add connections to any vertical transmission layers
                 # Placeholder code to be moved / refactored. The maternal network may need to be
@@ -323,11 +337,11 @@ def make_pregnancies(self, sim):
                         layer.add_pairs(uids, new_uids, dur=durs)
 
                 # Set prognoses for the pregnancies
-                self.set_prognoses(sim, uids)
+                self.set_prognoses(sim, uids) # Could set from_uids to network partners?
 
         return
 
-    def set_prognoses(self, sim, uids):
+    def set_prognoses(self, sim, to_uids, from_uids=None):
         """
         Make pregnancies
         Add miscarriage/termination logic here
@@ -336,22 +350,22 @@ def set_prognoses(self, sim, uids):
         """
 
         # Change states for the newly pregnant woman
-        self.susceptible[uids] = False
-        self.pregnant[uids] = True
-        self.ti_pregnant[uids] = sim.ti
+        self.susceptible[to_uids] = False
+        self.pregnant[to_uids] = True
+        self.ti_pregnant[to_uids] = sim.ti
 
         # Outcomes for pregnancies
-        dur = np.full(len(uids), sim.ti + self.pars.dur_pregnancy / sim.dt)
-        dead = np.random.random(len(uids)) < self.pars.p_death
-        self.ti_delivery[uids] = dur  # Currently assumes maternal deaths still result in a live baby
-        dur_post_partum = np.full(len(uids), dur + self.pars.dur_postpartum / sim.dt)
-        self.ti_postpartum[uids] = dur_post_partum
+        dur = np.full(len(to_uids), sim.ti + self.pars.dur_pregnancy / sim.dt)
+        dead = self.rng_dead.bernoulli(uids=to_uids, prob=self.pars.p_death)
+        self.ti_delivery[to_uids] = dur  # Currently assumes maternal deaths still result in a live baby
+        dur_post_partum = np.full(len(to_uids), dur + self.pars.dur_postpartum / sim.dt)
+        self.ti_postpartum[to_uids] = dur_post_partum
 
         if np.count_nonzero(dead):
-            self.ti_dead[uids[dead]] = dur[dead]
+            self.ti_dead[to_uids[dead]] = dur[dead]
         return
 
     def update_results(self, sim):
         self.results['pregnancies'][sim.ti] = np.count_nonzero(self.ti_pregnant == sim.ti)
         self.results['births'][sim.ti] = np.count_nonzero(self.ti_delivery == sim.ti)
-        return
+        return

stisim/distributions.py

@@ -17,14 +17,25 @@
 import sciris as sc
 
 __all__ = [
-    'Distribution', 'uniform', 'choice', 'normal', 'normal_pos', 'normal_int', 'lognormal', 'lognormal_int',
+    'Distribution', 'bernoulli', 'uniform', 'choice', 'normal', 'normal_pos', 'normal_int', 'lognormal', 'lognormal_int',
     'poisson', 'neg_binomial', 'beta', 'gamma', 'from_data'
 ]
 
 
 class Distribution():
 
-    def __init__(self):
+    def __init__(self, rng=None):
+        if rng is None:
+            self.stream = np.random.default_rng() # Default to centralized random number generator
+        else:
+            self.stream = rng
+        return
+
+    def set_stream(self, stream):
+        """
+        Switch to a user-supplied random number stream
+        """
+        self.stream = stream
         return
 
     def mean(self):
@@ -33,10 +44,10 @@ def mean(self):
         """
         raise NotImplementedError
 
-    def __call__(self, n=1, **kwargs):
-        return self.sample(n, **kwargs)
+    def __call__(self, size=None, uids=None, **kwargs):
+        return self.sample(size=size, uids=uids, **kwargs)
 
-    def sample(cls, n=1, **kwargs):
+    def sample(cls, size=None, uids=None, **kwargs):
         """
         Return a specified number of samples from the distribution
         """
@@ -46,19 +57,20 @@ def sample(cls, n=1, **kwargs):
 class from_data(Distribution):
     """ Sample from data """
 
-    def __init__(self, vals, bins):
+    def __init__(self, vals, bins, **kwargs):
+        super().__init__(**kwargs)
         self.vals = vals
         self.bins = bins
 
     def mean(self):
         return
 
-    def sample(self, n=1):
+    def sample(self, size, **kwargs):
         """ Sample using CDF """
         bin_midpoints = self.bins[:-1] + np.diff(self.bins) / 2
         cdf = np.cumsum(self.vals)
         cdf = cdf / cdf[-1]
-        values = np.random.rand(n)
+        values = self.stream.rand(size)
         value_bins = np.searchsorted(cdf, values)
         return bin_midpoints[value_bins]
 
@@ -68,42 +80,60 @@ class uniform(Distribution):
     Uniform distribution
     """
 
-    def __init__(self, low, high):
+    def __init__(self, low, high, **kwargs):
+        super().__init__(**kwargs)
         self.low = low
         self.high = high
 
     def mean(self):
         return (self.low + self.high) / 2
 
-    def sample(self, n=1):
-        return np.random.uniform(low=self.low, high=self.high, size=n)
+    def sample(self, **kwargs):
+        return self.stream.uniform(low=self.low, high=self.high, **kwargs)
+
+class bernoulli(Distribution):
+    """
+    Bernoulli distribution, returns sequence of True or False from independent trials
+    """
+
+    def __init__(self, p, **kwargs):
+        super().__init__(**kwargs)
+        self.p = p
+
+    def mean(self):
+        return self.p
+
+    def sample(self, **kwargs):
+        return self.stream.bernoulli(prob=self.p, **kwargs)
 
 
 class choice(Distribution):
     """
     Choose from samples, optionally with specific probabilities
     """
 
-    def __init__(self, choices, probabilities=None, replace=True):
+    def __init__(self, choices, probabilities=None, replace=True, **kwargs):
+        super().__init__(**kwargs)
         self.choices = choices
         self.probabilities = probabilities
         self.replace = replace
 
-    def sample(self, n, replace=True):
-        return np.random.choice(a=self.choices, p=self.probabilities, replace=self.replace, size=n)
+    def sample(self, **kwargs):
+        return self.stream.choice(a=self.choices, p=self.probabilities, replace=self.replace, **kwargs)
 
 
 class normal(Distribution):
     """
     Normal distribution
     """
 
-    def __init__(self, mean, std):
+    def __init__(self, mean, std, **kwargs):
+        super().__init__(**kwargs)
         self.mean = mean
         self.std = std
 
-    def sample(self, n=1):
-        return np.random.normal(loc=self.mean, scale=self.std, size=n)
+    def sample(self, **kwargs):
+        return self.stream.normal(loc=self.mean, scale=self.std, **kwargs)
 
 
 class normal_pos(normal):
@@ -113,17 +143,17 @@ class normal_pos(normal):
     WARNING - this function came from hpvsim but confirm that the implementation is correct?
     """
 
-    def sample(self, n=1):
-        return np.abs(super().sample(n))
+    def sample(self, **kwargs):
+        return np.abs(super().sample(**kwargs))
 
 
 class normal_int(Distribution):
     """
     Normal distribution returning only integer values
     """
 
-    def sample(self, n=1):
-        return np.round(super().sample(n))
+    def sample(self, **kwargs):
+        return np.round(super().sample(**kwargs))
 
 
 class lognormal(Distribution):
@@ -135,42 +165,49 @@ class lognormal(Distribution):
     function assumes the user wants to specify the mean and std of the lognormal distribution.
     """
 
-    def __init__(self, mean, std):
+    def __init__(self, mean, std, **kwargs):
+        super().__init__(**kwargs)
         self.mean = mean
         self.std = std
         self.underlying_mean = np.log(mean ** 2 / np.sqrt(std ** 2 + mean ** 2))  # Computes the mean of the underlying normal distribution
         self.underlying_std = np.sqrt(np.log(std ** 2 / mean ** 2 + 1))  # Computes sigma for the underlying normal distribution
 
-    def sample(self, n=1):
+    def sample(self, **kwargs):
 
         if (sc.isnumber(self.mean) and self.mean > 0) or (sc.checktype(self.mean, 'arraylike') and (self.mean > 0).all()):
-            return np.random.lognormal(mean=self.underlying_mean, sigma=self.underlying_std, size=n)
+            return self.stream.lognormal(mean=self.underlying_mean, sigma=self.underlying_std, **kwargs)
         else:
-            return np.zeros(n)
+            if 'size' in kwargs:
+                return np.zeros(kwargs['size'])
+            elif 'uids' in kwargs:
+                return np.zeros(len(kwargs['uids']))
+            else:
+                raise Exception('When calling sample(), please provide either "size" or "uids".')
 
 
 class lognormal_int(lognormal):
     """
     Lognormal returning only integer values
     """
 
-    def sample(self, n=1):
-        return np.round(super().sample(n))
+    def sample(self, **kwargs):
+        return np.round(super().sample(**kwargs))
 
 
 class poisson(Distribution):
     """
     Poisson distribution
     """
 
-    def __init__(self, rate):
+    def __init__(self, rate, **kwargs):
+        super().__init__(**kwargs)
         self.rate = rate
 
     def mean(self):
         return self.rate
 
-    def sample(self, n=1):
-        return np.random.poisson(self.rate, n)
+    def sample(self, **kwargs):
+        return self.stream.poisson(self.rate, **kwargs)
 
 
 class neg_binomial(Distribution):
@@ -185,50 +222,51 @@ class neg_binomial(Distribution):
     the mean.
     """
 
-    def __init__(self, mean, dispersion, step=1):
+    def __init__(self, mean, dispersion, **kwargs):
         """
         mean (float): the rate of the process (same as Poisson)
         dispersion (float):  dispersion parameter; lower is more dispersion, i.e. 0 = infinite, ∞ = Poisson
         n (int): number of trials
-        step (float): the step size to use if non-integer outputs are desired
         """
+        super().__init__(**kwargs)
         self.mean = mean
         self.dispersion = dispersion
-        self.step = step
 
-    def sample(self, n=1):
+    def sample(self, **kwargs):
         nbn_n = self.dispersion
-        nbn_p = self.dispersion / (self.mean / self.step + self.dispersion)
-        return np.random.negative_binomial(n=nbn_n, p=nbn_p, size=n) * self.step
+        nbn_p = self.dispersion / (self.mean + self.dispersion)
+        return self.stream.negative_binomial(n=nbn_n, p=nbn_p, **kwargs)
 
 
 class beta(Distribution):
     """
     Beta distribution
     """
 
-    def __init__(self, alpha, beta):
+    def __init__(self, alpha, beta, **kwargs):
+        super().__init__(**kwargs)
         self.alpha = alpha
         self.beta = beta
 
     def mean(self):
         return self.alpha / (self.alpha + self.beta)
 
-    def sample(self, n=1):
-        return np.random.beta(a=self.alpha, b=self.beta, size=n)
+    def sample(self, **kwargs):
+        return self.stream.beta(a=self.alpha, b=self.beta, **kwargs)
 
 
 class gamma(Distribution):
     """
     Gamma distribution
     """
 
-    def __init__(self, shape, scale):
+    def __init__(self, shape, scale, **kwargs):
+        super().__init__(**kwargs)
         self.shape = shape
         self.scale = scale
 
     def mean(self):
         return self.shape * self.scale
 
-    def sample(self, n=1):
-        return np.random.gamma(shape=self.shape, scale=self.scale, size=n)
+    def sample(self, **kwargs):
+        return self.stream.gamma(shape=self.shape, scale=self.scale, **kwargs)