From a2dd29c3630662028ed324b1f59ba8c96ff58065 Mon Sep 17 00:00:00 2001 From: Daniel Klein Date: Fri, 1 Sep 2023 17:55:08 -0700 Subject: [PATCH 01/34] Work in progress with new multi-stream random number generation functionality. --- stisim/__init__.py | 1 + stisim/demographics.py | 7 +- stisim/gonorrhea.py | 8 +- stisim/hiv.py | 11 +- stisim/modules.py | 28 +- stisim/networks.py | 138 +++- stisim/people.py | 2 +- stisim/sim.py | 6 +- stisim/streams.py | 207 ++++++ stisim/utils.py | 174 +---- tests/rng_experiment.ipynb | 1381 ++++++++++++++++++++++++++++++++++++ tests/simple.py | 7 +- 12 files changed, 1768 insertions(+), 202 deletions(-) create mode 100644 stisim/streams.py create mode 100644 tests/rng_experiment.ipynb diff --git a/stisim/__init__.py b/stisim/__init__.py index 82e24983..3601b787 100644 --- a/stisim/__init__.py +++ b/stisim/__init__.py @@ -3,6 +3,7 @@ from .parameters import * from .utils import * from .states import * +from .streams import * from .people import * from .networks import * from .modules import * diff --git a/stisim/demographics.py b/stisim/demographics.py index e383a51e..780a05f7 100644 --- a/stisim/demographics.py +++ b/stisim/demographics.py @@ -32,6 +32,9 @@ def __init__(self, pars=None): 'initial': 3, # Number of women initially pregnant }, self.pars) + self.rng_sex = ss.Stream('sex_at_birth') + self.rng_conception = ss.Stream('conception') + return def initialize(self, sim): @@ -98,7 +101,7 @@ 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(prob=self.pars.inci, arr=inds_to_choose_from) # Add UIDs for the as-yet-unborn agents so that we can track prognoses and transmission patterns n_unborn_agents = len(uids) @@ -106,7 +109,7 @@ def make_pregnancies(self, sim): # Grow the arrays and set properties for the unborn agents new_uids = sim.people.grow(n_unborn_agents) sim.people.age[new_uids] = -self.pars.dur_pregnancy - sim.people.female[new_uids] = np.random.choice([True, False], size=n_unborn_agents) + sim.people.female[new_uids] = self.rng_sex.bernoulli(prob=0.5, arr=uids) # Replace 0.5 with sex ratio at birth # Add connections to any vertical transmission layers # Placeholder code to be moved / refactored. The maternal network may need to be diff --git a/stisim/gonorrhea.py b/stisim/gonorrhea.py index a25eff57..72c577a1 100644 --- a/stisim/gonorrhea.py +++ b/stisim/gonorrhea.py @@ -20,6 +20,9 @@ def __init__(self, pars=None): self.ti_recovered = ss.State('ti_recovered', float, 0) self.ti_dead = ss.State('ti_dead', float, np.nan) # Death due to gonorrhea + self.rng_prog = ss.Stream('prog_dur') + self.rng_dead = ss.Stream('dead') + self.pars = ss.omerge({ 'dur_inf': 3, # not modelling diagnosis or treatment explicitly here 'p_death': 0.2, @@ -27,7 +30,7 @@ def __init__(self, pars=None): 'eff_condoms': 0.7, }, self.pars) return - + def update_states(self, sim): # What if something in here should depend on another module? # I guess we could just check for it e.g., 'if HIV in sim.modules' or @@ -51,7 +54,8 @@ def set_prognoses(self, sim, uids): self.ti_infected[uids] = sim.ti dur = sim.ti + np.random.poisson(self.pars['dur_inf']/sim.pars.dt, len(uids)) - dead = np.random.random(len(uids)) < self.pars.p_death + dead = self.rng_dead.bernoulli(self.pars.p_death, uids) + self.ti_recovered[uids[~dead]] = dur[~dead] self.ti_dead[uids[dead]] = dur[dead] return diff --git a/stisim/hiv.py b/stisim/hiv.py index 400bb7c2..8319952e 100644 --- a/stisim/hiv.py +++ b/stisim/hiv.py @@ -64,8 +64,14 @@ def __init__(self, t: np.array, capacity: np.array): self.t = sc.promotetoarray(t) self.capacity = sc.promotetoarray(capacity) + self.rng_add_ART = ss.Stream('add_ART') + self.rng_remove_ART = ss.Stream('remove_ART') + + return + def initialize(self, sim): sim.hiv.results += ss.Result(self.name, 'n_art', sim.npts, dtype=int) + return def apply(self, sim): if sim.t < self.t[0]: @@ -80,12 +86,13 @@ def apply(self, sim): eligible = sim.people.alive & sim.people.hiv.infected & ~sim.people.hiv.on_art n_eligible = np.count_nonzero(eligible) if n_eligible: - inds = np.random.choice(ss.true(eligible), min(n_eligible, n_change), replace=False) + inds = self.rng_add_ART.bernoulli_filter(prob=min(n_eligible, n_change)/n_eligible, arr=eligible) sim.people.hiv.on_art[inds] = True elif n_change < 0: # Take some people off ART eligible = sim.people.alive & sim.people.hiv.infected & sim.people.hiv.on_art - inds = np.random.choice(ss.true(eligible), min(n_change), replace=False) + n_eligible = np.count_nonzero(eligible) + inds = self.rng_remove_ART.bernoulli_filter(prob=-n_change/n_eligible, arr=eligible) sim.people.hiv.on_art[inds] = False # Add result diff --git a/stisim/modules.py b/stisim/modules.py index 1b1c9839..97c78250 100644 --- a/stisim/modules.py +++ b/stisim/modules.py @@ -19,6 +19,12 @@ def __init__(self, pars=None, label=None, requires=None, *args, **kwargs): self.results = ss.Results() self.initialized = False self.finalized = False + + # Random number streams + self.rng_init_cases = ss.Stream('initial_cases') + self.rng_trans_ab = ss.Stream('trans_ab') + self.rng_trans_ba = ss.Stream('trans_ba') + return def __call__(self, *args, **kwargs): @@ -37,10 +43,14 @@ def check_requires(self, sim): def initialize(self, sim): self.check_requires(sim) - # Connect the states to the sim + # Connect the states to the people for state in self.states.values(): state.initialize(sim.people) + # Connect the streams to the sim + for stream in self.streams.values(): + stream.initialize(sim) + self.initialized = True return @@ -59,6 +69,10 @@ def name(self): def states(self): return ss.ndict({k:v for k,v in self.__dict__.items() if isinstance(v, ss.State)}) + @property + def streams(self): + return ss.ndict({k:v for k,v in self.__dict__.items() if isinstance(v, ss.Stream)}) + class Modules(ss.ndict): def __init__(self, *args, type=Module, **kwargs): @@ -100,7 +114,13 @@ def set_initial_states(self, sim): i.e., creating their dynamic array, linking them to a People instance. That should have already taken place by the time this method is called. """ - initial_cases = np.random.choice(sim.people.uid, self.pars['initial']) + if self.pars['initial'] <= 0: + return + + #initial_cases = np.random.choice(sim.people.uid, self.pars['initial']) + #rng = sim.rngs.get(f'initial_cases_{self.name}') + #initial_cases = ss.binomial_filter_rng(prob=self.pars['initial']/len(sim.people), arr=sim.people.uid, rng=rng, block_size=len(sim.people._uid_map)) + initial_cases = self.rng_init_cases.bernoulli_filter(prob=self.pars['initial']/len(sim.people), arr=sim.people.uid) self.set_prognoses(sim, initial_cases) return @@ -135,10 +155,10 @@ def make_new_cases(self, sim): if k in pars['beta']: rel_trans = (self.infected & sim.people.alive).astype(float) rel_sus = (self.susceptible & sim.people.alive).astype(float) - for a, b, beta in [[layer['p1'], layer['p2'], pars['beta'][k][0]], [layer['p2'], layer['p1'], pars['beta'][k][1]]]: + for a, b, beta, rng in [[layer['p1'], layer['p2'], pars['beta'][k][0], self.rng_trans_ab], [layer['p2'], layer['p1'], pars['beta'][k][1], self.rng_trans_ba]]: # probability of a->b transmission p_transmit = rel_trans[a] * rel_sus[b] * layer['beta'] * beta - new_cases = np.random.random(len(a)) < p_transmit + new_cases = rng.random(len(a)) < p_transmit # TODO: Convert to flat list of N probs and block sample if new_cases.any(): self.set_prognoses(sim, b[new_cases]) diff --git a/stisim/networks.py b/stisim/networks.py index de5c843c..8b0b9778 100644 --- a/stisim/networks.py +++ b/stisim/networks.py @@ -9,7 +9,7 @@ # Specify all externally visible functions this file defines -__all__ = ['Networks', 'Network', 'simple_sexual', 'hpv_network', 'maternal'] +__all__ = ['Networks', 'Network', 'simple_sexual', 'simple_embedding', 'hpv_network', 'maternal'] class Network(sc.objdict): """ @@ -80,8 +80,21 @@ def name(self): # The module name is a lower-case version of its class name return self.__class__.__name__.lower() - def initialize(self): + def initialize(self, sim): pass + # Auto initialization trick for streams does not work because objdict overwrites __dict__ + ''' + # Connect the streams to the sim + for stream in self.streams.values(): + stream.initialize(sim) + return + ''' + + ''' + @property + def streams(self): + return ss.ndict({k:v for k,v in self.__dict__.items() if isinstance(v, ss.Stream)}) + ''' def __len__(self): try: @@ -270,8 +283,23 @@ def __init__(self, mean_dur=5): # Set other parameters self.mean_dur = mean_dur - def initialize(self, people): - self.add_pairs(people, ti=0) + # Define random streams + self.rng_pair_12 = ss.Stream('pair_12') + self.rng_pair_21 = ss.Stream('pair_21') + self.rng_mean_dur = ss.Stream('mean_dur') + + return + + def initialize(self, sim): + super().initialize(sim) + + # Initialize random streams + self.rng_pair_12.initialize(sim) + self.rng_pair_21.initialize(sim) + self.rng_mean_dur.initialize(sim) + + self.add_pairs(sim.people, ti=0) + return def add_pairs(self, people, ti=None): # Find unpartnered males and females - could in principle check other contact layers too @@ -282,13 +310,13 @@ def add_pairs(self, people, ti=None): if len(available_m) <= len(available_f): p1 = available_m - p2 = np.random.choice(available_f, len(p1), replace=False) + p2 = self.rng_pair_12.choice(available_f, len(p1), replace=False) # TODO: Stream-ify else: p2 = available_f - p1 = np.random.choice(available_m, len(p2), replace=False) + p1 = self.rng_pair_21.choice(available_m, len(p2), replace=False) # TODO: Stream-ify beta = np.ones_like(p1) - dur = np.random.poisson(self.mean_dur, len(p1)) + dur = self.rng_mean_dur.poisson(self.mean_dur, len(p1)) # TODO: Stream-ify self['p1'] = np.concatenate([self['p1'], p1]) self['p2'] = np.concatenate([self['p2'], p2]) self['beta'] = np.concatenate([self['beta'], beta]) @@ -307,6 +335,77 @@ def update(self, people, dt=None): # Then add new relationships for unpartnered people self.add_pairs(people) +class simple_embedding(simple_sexual): + """ + A class holding a single network of contact edges (connections) between people. + This network is built by **randomly pairing** males and female with variable relationship durations. + """ + + def add_pairs(self, people, ti=None): + # Find unpartnered males and females - could in principle check other contact layers too + # by having the People object passed in here + + available_m = np.setdiff1d(people.uid[~people.female], self.members) + available_f = np.setdiff1d(people.uid[people.female], self.members) + + # Make p1 the shorter array + if len(available_f) < len(available_m): + p1 = available_f + p2 = available_m + else: + p1 = available_m + p2 = available_f + + loc1 = self.rng_pair_12.random(arr=p1) + loc2 = self.rng_pair_21.random(arr=p2) + + done = False + + p1v = np.tile(loc1, (len(loc2),1)) + p2v = np.tile(loc2[:,np.newaxis], len(loc1)) + d = np.absolute(p2v-p1v) + + pairs = [] + + while not done: + closest1 = d.argmin(axis=0) + # with axis=0, 0 at loc1[0] will be closest to closest1[0] at loc2[ closest1[0] ] + + ele = np.unique(closest1) + closest2 = d.argmin(axis=1) + + #loc1[ closest2[np.unique(closest1)]] is near loc2[ele] + # so pair p2=ele with p1=closest2[np.unique(closest1)] + pairs.append( (p1[closest2[np.unique(closest1)]], p2[ele]) ) + + # WIP! + + # remove pairs and repeat + unmatched_inds_1 = np.setdiff1d(np.arange(len(p1)), closest2[np.unique(closest1)]) + unmatched_inds_2 = np.setdiff1d(np.arange(len(p2)), ele) + #p1 = np.setdiff1d( p1, p1[closest2[np.unique(closest1)]]) + #p2 = np.setdiff1d(p2, p2[ele]) + p1 = p1[unmatched_inds_1] + p2 = p2[unmatched_inds_2] + + d = d[unmatched_inds_1, unmatched_inds_2] + + + + if len(available_m) <= len(available_f): + p1 = available_m + p2 = self.rng_pair_12.choice(available_f, len(p1), replace=False) # TODO: Stream-ify + else: + p2 = available_f + p1 = self.rng_pair_21.choice(available_m, len(p2), replace=False) # TODO: Stream-ify + + beta = np.ones_like(p1) + dur = self.rng_mean_dur.poisson(self.mean_dur, len(p1)) # TODO: Stream-ify + self['p1'] = np.concatenate([self['p1'], p1]) + self['p2'] = np.concatenate([self['p2'], p2]) + self['beta'] = np.concatenate([self['beta'], beta]) + self['dur'] = np.concatenate([self['dur'], dur]) + class hpv_network(Network): def __init__(self, pars=None): @@ -334,11 +433,22 @@ def __init__(self, pars=None): self.pars['participation'] = None # Incidence of partnership formation by age self.pars['mixing'] = None # Mixing matrices for storing age differences in partnerships + self.rng_partners = ss.Stream('partners') + self.rng_acts = ss.Stream('acts') + self.rng_dur_pship = ss.Stream('dur_pship') + self.update_pars(pars) self.get_layer_probs() - def initialize(self, people): - self.add_pairs(people, ti=0) + def initialize(self, sim): + super().initialize(sim) + + self.rng_partners.initialize(sim) + self.rng_acts.initialize(sim) + self.rng_dur_pship.initialize(sim) + + self.add_pairs(sim.people, ti=0) + return def update_pars(self, pars): if pars is not None: @@ -397,7 +507,7 @@ def add_pairs(self, people, ti=0): current_partners = np.zeros((len(people))) current_partners[f_partnered_inds] = f_partnered_counts current_partners[m_partnered_inds] = m_partnered_counts - partners = ss.sample(**self.pars['partners'], size=len(people)) + 1 + partners = ss.sample(self.rng_partners, **self.pars['partners'], size=len(people)) + 1 underpartnered = current_partners < partners # Indices of underpartnered people f_eligible = f_active & underpartnered m_eligible = m_active & underpartnered @@ -439,12 +549,12 @@ def add_pairs(self, people, ti=0): choices = [] fems = np.arange(len(f)) f_paired_bools = np.full(len(fems), True, dtype=bool) - np.random.shuffle(fems) + np.random.shuffle(fems) # TODO: Stream-ify for fem in fems: m_col = pair_probs[:, fem] if m_col.sum() > 0: m_col_norm = m_col / m_col.sum() - choice = np.random.choice(len(m_col_norm), 1, replace=False, p=m_col_norm) + choice = np.random.choice(len(m_col_norm), 1, replace=False, p=m_col_norm) # TODO: Stream-ify choices.append(choice) pair_probs[choice, :] = 0 # Once male partner is assigned, remove from eligible pool else: @@ -455,8 +565,8 @@ def add_pairs(self, people, ti=0): p1 = np.array(f) p2 = selected_males n_partnerships = len(p1) - dur = ss.sample(**self.pars['dur_pship'], size=n_partnerships) - acts = ss.sample(**self.pars['acts'], size=n_partnerships) + dur = ss.sample(self.rng_dur_pship, **self.pars['dur_pship'], size=n_partnerships) + acts = ss.sample(self.rng_acts, **self.pars['acts'], size=n_partnerships) age_p1 = people.age[p1] age_p2 = people.age[p2] diff --git a/stisim/people.py b/stisim/people.py index 7bb76641..93223590 100644 --- a/stisim/people.py +++ b/stisim/people.py @@ -170,7 +170,7 @@ def __init__(self, n, extra_states=None, networks=None): def initialize(self, popdict=None): """ Initialize people by setting their attributes """ if popdict is None: # TODO: update - self['age'][:] = np.random.random(size=len(self)) * 100 + self['age'][:] = np.random.random(size=len(self)) * 100 # These are okay for now as direct calls to the central rng self['female'][:] = np.random.choice([False, True], size=len(self)) else: # Use random defaults diff --git a/stisim/sim.py b/stisim/sim.py index eff3efbf..5888438e 100644 --- a/stisim/sim.py +++ b/stisim/sim.py @@ -42,6 +42,9 @@ def __init__(self, pars=None, label=None, people=None, modules=None, **kwargs): self.interventions = ss.Interventions() self.analyzers = ss.Analyzers() + # Decision container + self.streams = ss.Streams() + return @property @@ -246,7 +249,7 @@ def init_networks(self): else: layer_name = key network.label = layer_name - network.initialize(self.people) + network.initialize(self) self.people.networks[layer_name] = network return @@ -322,6 +325,7 @@ def step(self): raise AlreadyRunError('Simulation already complete (call sim.initialize() to re-run)') # Update states, modules, partnerships + self.streams.step(self.ti) self.people.update(self) self.apply_interventions() self.update_modules() diff --git a/stisim/streams.py b/stisim/streams.py new file mode 100644 index 00000000..3812437a --- /dev/null +++ b/stisim/streams.py @@ -0,0 +1,207 @@ +import numpy as np +import sciris as sc + +__all__ = ['Streams', 'Stream'] + + +class Streams: + """ + Class for managing a collection random number streams + """ + + def __init__(self): + self.streams = [] + return + + def add(self, stream): + """ Add a stream """ + + # Return value will be used as the seed offset for this stream + # Depends on order, which will become a problem - need a better solution + n = len(self.streams) + self.streams.append(stream) + + return n + + def step(self, ti): + for stream in self.streams: + stream.step(ti) + return + +class Stream(np.random.Generator): + """ + Class for tracking one random number stream associated with one decision per timestep + """ + + def __init__(self, name, seed_offset=0, **kwargs): + """ + Create a random number stream + + name: a name for this Stream, like "coin_flip" + uid: an identifier added to the name to make it uniquely identifiable, for example the name or id of the calling class + """ + + ''' + if 'bit_generator' not in kwargs: + kwargs['bit_generator'] = np.random.PCG64(seed=self.seed + self.seed_offset) + super().__init__(bit_generator=np.random.PCG64()) + ''' + + self.name = name + self.seed_offset = seed_offset + self.kwargs = kwargs + + self.seed = None + self.ppl = None # Needed for block size + + self.initialized = False + + return + + def initialize(self, sim): + if self.initialized: + # TODO: Raise warning + assert not self.initialized + return + + self.seed = sim.streams.add(self) + + if 'bit_generator' not in self.kwargs: + self.kwargs['bit_generator'] = np.random.PCG64(seed=self.seed + self.seed_offset) + super().__init__(**self.kwargs) + + #self.rng = np.random.default_rng(seed=self.seed + self.seed_offset) + self._init_state = self.bit_generator.state # Store the initial state + + self.ppl = sim.people + + self.initialized = True + return + + def step(self, ti): + """ Advance to time ti step by jumping """ + self.bit_generator.state = self._init_state # restore initial state + self.bit_generator.jumped(jumps=ti) # Take ti jumps + return + + @property + def block_size(self): + return len(self.ppl._uid_map) + + def random(self, arr): + return (super(Stream, self).random(self.block_size))[arr] + + def bernoulli(self, prob, arr): + #return super(Stream, self).choice([True, False], size=self.block_size, p=[prob, 1-prob]) # very slow + #return (super(Stream, self).binomial(n=1, p=prob, size=self.block_size))[arr].astype(bool) # pretty fast + return (self.random(self.block_size) < prob)[arr] # fastest + + def bernoulli_filter(self, prob, arr): + #return arr[self.bernoulli(prob, arr).nonzero()[0]] + return arr[self.bernoulli(prob, arr)] # Slightly faster on my machine for bernoulli to typecast + + def sample(self, dist=None, par1=None, par2=None, size=None, **kwargs): + """ + Draw a sample from the distribution specified by the input. The available + distributions are: + + - 'uniform' : uniform from low=par1 to high=par2; mean is equal to (par1+par2)/2 + - 'choice' : par1=array of choices, par2=probability of each choice + - 'normal' : normal with mean=par1 and std=par2 + - 'lognormal' : lognormal with mean=par1, std=par2 (parameters are for the lognormal, not the underlying normal) + - 'normal_pos' : right-sided normal (i.e. only +ve values), with mean=par1, std=par2 of the underlying normal + - 'normal_int' : normal distribution with mean=par1 and std=par2, returns only integer values + - 'lognormal_int' : lognormal distribution with mean=par1 and std=par2, returns only integer values + - 'poisson' : Poisson distribution with rate=par1 (par2 is not used); mean and variance are equal to par1 + - 'neg_binomial' : negative binomial distribution with mean=par1 and k=par2; converges to Poisson with k=∞ + - 'beta' : beta distribution with alpha=par1 and beta=par2; + - 'gamma' : gamma distribution with shape=par1 and scale=par2; + + Args: + self (Stream) : the random number generator stream + dist (str) : the distribution to sample from + par1 (float) : the "main" distribution parameter (e.g. mean) + par2 (float) : the "secondary" distribution parameter (e.g. std) + size (int) : the number of samples (default=1) + kwargs (dict) : passed to individual sampling functions + + Returns: + A length N array of samples + + **Examples**:: + + ss.sample() # returns Unif(0,1) + ss.sample(dist='normal', par1=3, par2=0.5) # returns Normal(μ=3, σ=0.5) + ss.sample(dist='lognormal_int', par1=5, par2=3) # returns lognormally distributed values with mean 5 and std 3 + + Notes: + Lognormal distributions are parameterized with reference to the underlying normal distribution (see: + https://docs.scipy.org/doc/numpy-1.14.0/reference/generated/numpy.random.lognormal.html), but this + function assumes the user wants to specify the mean and std of the lognormal distribution. + + Negative binomial distributions are parameterized with reference to the mean and dispersion parameter k + (see: https://en.wikipedia.org/wiki/Negative_binomial_distribution). The r parameter of the underlying + distribution is then calculated from the desired mean and k. For a small mean (~1), a dispersion parameter + of ∞ corresponds to the variance and standard deviation being equal to the mean (i.e., Poisson). For a + large mean (e.g. >100), a dispersion parameter of 1 corresponds to the standard deviation being equal to + the mean. + """ + + # Some of these have aliases, but these are the "official" names + choices = [ + 'uniform', + 'normal', + 'choice', + 'normal_pos', + 'normal_int', + 'lognormal', + 'lognormal_int', + 'poisson', + 'neg_binomial', + 'beta', + 'gamma', + ] + + # Ensure it's an integer + if size is not None and not isinstance(size, tuple): + size = int(size) + + # Compute distribution parameters and draw samples + # NB, if adding a new distribution, also add to choices above + if dist in ['unif', 'uniform']: + samples = self.uniform(low=par1, high=par2, size=size) + elif dist in ['choice']: + samples = self.choice(a=par1, p=par2, size=size, **kwargs) + elif dist in ['norm', 'normal']: + samples = self.normal(loc=par1, scale=par2, size=size) + elif dist == 'normal_pos': + samples = np.abs(self.normal(loc=par1, scale=par2, size=size)) + elif dist == 'normal_int': + samples = np.round(np.abs(self.normal(loc=par1, scale=par2, size=size))) + elif dist == 'poisson': + samples = self.poisson(rate=par1, n=size) # Use Numba version below for speed + elif dist == 'beta': + samples = self.beta(a=par1, b=par2, size=size) + elif dist == 'gamma': + samples = self.gamma(shape=par1, scale=par2, size=size) + elif dist in ['lognorm', 'lognormal', 'lognorm_int', 'lognormal_int']: + if (sc.isnumber(par1) and par1 > 0) or (sc.checktype(par1, 'arraylike') and (par1 > 0).all()): + mean = np.log( + par1 ** 2 / np.sqrt(par2 ** 2 + par1 ** 2)) # Computes the mean of the underlying normal distribution + sigma = np.sqrt(np.log(par2 ** 2 / par1 ** 2 + 1)) # Computes sigma for the underlying normal distribution + samples = self.lognormal(mean=mean, sigma=sigma, size=size) + else: + samples = np.zeros(size) + if '_int' in dist: + samples = np.round(samples) + # Calculate a and b using mean (par1) and variance (par2) + # https://stats.stackexchange.com/questions/12232/calculating-the-parameters-of-a-beta-distribution-using-the-mean-and-variance + elif dist == 'beta_mean': + a = ((1 - par1) / par2 - 1 / par1) * par1 ** 2 + b = a * (1 / par1 - 1) + samples = self.beta(a=a, b=b, size=size) + else: + errormsg = f'The selected distribution "{dist}" is not implemented; choices are: {sc.newlinejoin(choices)}' + raise NotImplementedError(errormsg) + + return samples \ No newline at end of file diff --git a/stisim/utils.py b/stisim/utils.py index e38fbb17..a2bb44d6 100644 --- a/stisim/utils.py +++ b/stisim/utils.py @@ -175,115 +175,7 @@ def find_contacts(p1, p2, inds): # pragma: no cover # %% Sampling and seed methods -__all__ += ['sample', 'set_seed'] - - -def sample(dist=None, par1=None, par2=None, size=None, **kwargs): - """ - Draw a sample from the distribution specified by the input. The available - distributions are: - - - 'uniform' : uniform from low=par1 to high=par2; mean is equal to (par1+par2)/2 - - 'choice' : par1=array of choices, par2=probability of each choice - - 'normal' : normal with mean=par1 and std=par2 - - 'lognormal' : lognormal with mean=par1, std=par2 (parameters are for the lognormal, not the underlying normal) - - 'normal_pos' : right-sided normal (i.e. only +ve values), with mean=par1, std=par2 of the underlying normal - - 'normal_int' : normal distribution with mean=par1 and std=par2, returns only integer values - - 'lognormal_int' : lognormal distribution with mean=par1 and std=par2, returns only integer values - - 'poisson' : Poisson distribution with rate=par1 (par2 is not used); mean and variance are equal to par1 - - 'neg_binomial' : negative binomial distribution with mean=par1 and k=par2; converges to Poisson with k=∞ - - 'beta' : beta distribution with alpha=par1 and beta=par2; - - 'gamma' : gamma distribution with shape=par1 and scale=par2; - - Args: - dist (str): the distribution to sample from - par1 (float): the "main" distribution parameter (e.g. mean) - par2 (float): the "secondary" distribution parameter (e.g. std) - size (int): the number of samples (default=1) - kwargs (dict): passed to individual sampling functions - - Returns: - A length N array of samples - - **Examples**:: - - ss.sample() # returns Unif(0,1) - ss.sample(dist='normal', par1=3, par2=0.5) # returns Normal(μ=3, σ=0.5) - ss.sample(dist='lognormal_int', par1=5, par2=3) # returns lognormally distributed values with mean 5 and std 3 - - Notes: - Lognormal distributions are parameterized with reference to the underlying normal distribution (see: - https://docs.scipy.org/doc/numpy-1.14.0/reference/generated/numpy.random.lognormal.html), but this - function assumes the user wants to specify the mean and std of the lognormal distribution. - - Negative binomial distributions are parameterized with reference to the mean and dispersion parameter k - (see: https://en.wikipedia.org/wiki/Negative_binomial_distribution). The r parameter of the underlying - distribution is then calculated from the desired mean and k. For a small mean (~1), a dispersion parameter - of ∞ corresponds to the variance and standard deviation being equal to the mean (i.e., Poisson). For a - large mean (e.g. >100), a dispersion parameter of 1 corresponds to the standard deviation being equal to - the mean. - """ - - # Some of these have aliases, but these are the "official" names - choices = [ - 'uniform', - 'normal', - 'choice', - 'normal_pos', - 'normal_int', - 'lognormal', - 'lognormal_int', - 'poisson', - 'neg_binomial', - 'beta', - 'gamma', - ] - - # Ensure it's an integer - if size is not None and not isinstance(size, tuple): - size = int(size) - - # Compute distribution parameters and draw samples - # NB, if adding a new distribution, also add to choices above - if dist in ['unif', 'uniform']: - samples = np.random.uniform(low=par1, high=par2, size=size) - elif dist in ['choice']: - samples = np.random.choice(a=par1, p=par2, size=size, **kwargs) - elif dist in ['norm', 'normal']: - samples = np.random.normal(loc=par1, scale=par2, size=size) - elif dist == 'normal_pos': - samples = np.abs(np.random.normal(loc=par1, scale=par2, size=size)) - elif dist == 'normal_int': - samples = np.round(np.abs(np.random.normal(loc=par1, scale=par2, size=size))) - elif dist == 'poisson': - samples = n_poisson(rate=par1, n=size) # Use Numba version below for speed - elif dist == 'neg_binomial': - samples = n_neg_binomial(rate=par1, dispersion=par2, n=size, **kwargs) # Use custom version below - elif dist == 'beta': - samples = np.random.beta(a=par1, b=par2, size=size) - elif dist == 'gamma': - samples = np.random.gamma(shape=par1, scale=par2, size=size) - elif dist in ['lognorm', 'lognormal', 'lognorm_int', 'lognormal_int']: - if (sc.isnumber(par1) and par1 > 0) or (sc.checktype(par1, 'arraylike') and (par1 > 0).all()): - mean = np.log( - par1 ** 2 / np.sqrt(par2 ** 2 + par1 ** 2)) # Computes the mean of the underlying normal distribution - sigma = np.sqrt(np.log(par2 ** 2 / par1 ** 2 + 1)) # Computes sigma for the underlying normal distribution - samples = np.random.lognormal(mean=mean, sigma=sigma, size=size) - else: - samples = np.zeros(size) - if '_int' in dist: - samples = np.round(samples) - # Calculate a and b using mean (par1) and variance (par2) - # https://stats.stackexchange.com/questions/12232/calculating-the-parameters-of-a-beta-distribution-using-the-mean-and-variance - elif dist == 'beta_mean': - a = ((1 - par1) / par2 - 1 / par1) * par1 ** 2 - b = a * (1 / par1 - 1) - samples = np.random.beta(a=a, b=b, size=size) - else: - errormsg = f'The selected distribution "{dist}" is not implemented; choices are: {sc.newlinejoin(choices)}' - raise NotImplementedError(errormsg) - - return samples +__all__ += ['set_seed'] def set_seed(seed=None): @@ -301,68 +193,6 @@ def set_seed(seed=None): return -# %% Probabilities -- mostly not jitted since performance gain is minimal - -__all__ += ['binomial_filter', 'n_poisson', 'n_neg_binomial'] - - - -def binomial_filter(prob, arr): - """ - Binomial "filter" -- the same as n_binomial, except return - the elements of arr that succeeded. - - Args: - prob (float): probability of each trial succeeding - arr (array): the array to be filtered - - Returns: - Subset of array for which trials succeeded - - **Example**:: - - inds = ss.binomial_filter(0.5, np.arange(20)**2) # Which values in the (arbitrary) array passed the coin flip - """ - return arr[(np.random.random(len(arr)) < prob).nonzero()[0]] - - - -def n_poisson(rate, n): - """ - An array of Poisson trials. - - Args: - rate (float): the rate of the Poisson process (mean) - n (int): number of trials - - **Example**:: - - outcomes = ss.n_poisson(100, 20) # 20 Poisson trials with mean 100 - """ - return np.random.poisson(rate, n) - - -def n_neg_binomial(rate, dispersion, n, step=1): # Numba not used due to incompatible implementation - """ - An array of negative binomial trials. See ss.sample() for more explanation. - - Args: - rate (float): the rate of the process (mean, 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 - - **Example**:: - - outcomes = ss.n_neg_binomial(100, 1, 50) # 50 negative binomial trials with mean 100 and dispersion roughly equal to mean (large-mean limit) - outcomes = ss.n_neg_binomial(1, 100, 20) # 20 negative binomial trials with mean 1 and dispersion still roughly equal to mean (approximately Poisson) - """ - nbn_n = dispersion - nbn_p = dispersion / (rate / step + dispersion) - samples = np.random.negative_binomial(n=nbn_n, p=nbn_p, size=n) * step - return samples - - # %% Simple array operations __all__ += ['true', 'false', 'defined', 'undefined'] @@ -422,4 +252,4 @@ def undefined(arr): inds = ss.defined(np.array([1,np.nan,0,np.nan,1,0,1])) """ - return np.isnan(arr).nonzero()[-1] + return np.isnan(arr).nonzero()[-1] \ No newline at end of file diff --git a/tests/rng_experiment.ipynb b/tests/rng_experiment.ipynb new file mode 100644 index 00000000..5c3af2f2 --- /dev/null +++ b/tests/rng_experiment.ipynb @@ -0,0 +1,1381 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 50, + "id": "55881074", + "metadata": { + "ExecuteTime": { + "end_time": "2023-08-27T07:52:23.276040Z", + "start_time": "2023-08-27T07:52:23.214481Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "The autoreload extension is already loaded. To reload it, use:\n", + " %reload_ext autoreload\n" + ] + } + ], + "source": [ + "%load_ext autoreload\n", + "%autoreload 2" + ] + }, + { + "cell_type": "code", + "execution_count": 51, + "id": "c60d109c", + "metadata": {}, + "outputs": [], + "source": [ + "import stisim as ss\n", + "import stisim.utils as ssu\n", + "import numpy as np\n", + "import pandas as pd" + ] + }, + { + "cell_type": "markdown", + "id": "0fb247d0", + "metadata": {}, + "source": [ + "IDEA 1: Draw more random numbers than needed on each time step and use by row of people arrays.\n", + "Set seed to seed + ti on each step?" + ] + }, + { + "cell_type": "code", + "execution_count": 52, + "id": "c73f0ce1", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Initializing sim (warning: not resetting sim.people) with 10000 agents\n" + ] + } + ], + "source": [ + "ppl = ss.People(10)\n", + "ppl.networks = ss.ndict(ss.simple_sexual(), ss.maternal())\n", + "\n", + "hiv = ss.HIV()\n", + "hiv.pars['beta'] = {'simple_sexual': [0.0008, 0.0004], 'maternal': [0.2, 0]}\n", + "\n", + "sim = ss.Sim(people=ppl, modules=[hiv, ss.Gonorrhea(), ss.Pregnancy()])\n", + "sim.initialize();\n", + "#sim.run()" + ] + }, + { + "cell_type": "code", + "execution_count": 53, + "id": "106ad6cf", + "metadata": {}, + "outputs": [], + "source": [ + "def dfp(vec):\n", + " display(pd.DataFrame({'result': vec}, index=pd.Index(sim.people.uid, name='uid')))" + ] + }, + { + "cell_type": "code", + "execution_count": 54, + "id": "00dab08a", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
result
uid
00
11
20
30
40
50
60
71
81
90
\n", + "
" + ], + "text/plain": [ + " result\n", + "uid \n", + "0 0\n", + "1 1\n", + "2 0\n", + "3 0\n", + "4 0\n", + "5 0\n", + "6 0\n", + "7 1\n", + "8 1\n", + "9 0" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
result
uid
0True
1False
2False
3True
4True
5False
6False
7False
8False
9True
\n", + "
" + ], + "text/plain": [ + " result\n", + "uid \n", + "0 True\n", + "1 False\n", + "2 False\n", + "3 True\n", + "4 True\n", + "5 False\n", + "6 False\n", + "7 False\n", + "8 False\n", + "9 True" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "ssu.set_seed(0)\n", + "dfp(np.random.binomial(n=1, p=len(sim.people)*[0.3]))\n", + "dfp(np.random.choice([True, False], len(sim.people)))" + ] + }, + { + "cell_type": "code", + "execution_count": 55, + "id": "99f417c3", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
result
uid
00
11
20
30
40
50
60
71
81
90
\n", + "
" + ], + "text/plain": [ + " result\n", + "uid \n", + "0 0\n", + "1 1\n", + "2 0\n", + "3 0\n", + "4 0\n", + "5 0\n", + "6 0\n", + "7 1\n", + "8 1\n", + "9 0" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
result
uid
0True
1False
2False
3True
4True
5False
6False
7False
8False
9True
\n", + "
" + ], + "text/plain": [ + " result\n", + "uid \n", + "0 True\n", + "1 False\n", + "2 False\n", + "3 True\n", + "4 True\n", + "5 False\n", + "6 False\n", + "7 False\n", + "8 False\n", + "9 True" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "# Check that results are repeatable\n", + "ssu.set_seed(0)\n", + "dfp(np.random.binomial(n=1, p=len(sim.people)*[0.3]))\n", + "dfp(np.random.choice([True, False], len(sim.people)))" + ] + }, + { + "cell_type": "code", + "execution_count": 56, + "id": "db057c24", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "array([1, 2, 3, 4, 5, 6, 7, 8, 9])" + ] + }, + "execution_count": 56, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "# Now remove agent with uid=0\n", + "sim.people.remove(0)\n", + "sim.people.uid" + ] + }, + { + "cell_type": "code", + "execution_count": 57, + "id": "4aa9220c", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
result
uid
10
21
30
40
50
60
70
81
91
\n", + "
" + ], + "text/plain": [ + " result\n", + "uid \n", + "1 0\n", + "2 1\n", + "3 0\n", + "4 0\n", + "5 0\n", + "6 0\n", + "7 0\n", + "8 1\n", + "9 1" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
result
uid
1True
2False
3True
4False
5False
6True
7True
8False
9False
\n", + "
" + ], + "text/plain": [ + " result\n", + "uid \n", + "1 True\n", + "2 False\n", + "3 True\n", + "4 False\n", + "5 False\n", + "6 True\n", + "7 True\n", + "8 False\n", + "9 False" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "ssu.set_seed(0)\n", + "dfp(np.random.binomial(n=1, p=len(sim.people)*[0.3]))\n", + "dfp(np.random.choice([True, False], len(sim.people)))" + ] + }, + { + "cell_type": "code", + "execution_count": 58, + "id": "221ea7c0", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
result
uid
11
20
30
40
50
60
71
81
90
\n", + "
" + ], + "text/plain": [ + " result\n", + "uid \n", + "1 1\n", + "2 0\n", + "3 0\n", + "4 0\n", + "5 0\n", + "6 0\n", + "7 1\n", + "8 1\n", + "9 0" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
result
uid
1False
2False
3True
4True
5False
6False
7False
8False
9True
\n", + "
" + ], + "text/plain": [ + " result\n", + "uid \n", + "1 False\n", + "2 False\n", + "3 True\n", + "4 True\n", + "5 False\n", + "6 False\n", + "7 False\n", + "8 False\n", + "9 True" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "# Now try sampling for all uids and keeping only uids in the current view\n", + "ssu.set_seed(0)\n", + "dfp(np.random.binomial(n=1, p=len(sim.people._uid_map)*[0.3])[sim.people.uid])\n", + "dfp(np.random.choice([True, False], len(sim.people._uid_map))[sim.people.uid])\n", + "\n", + "# ^^^ WORKS! (These are independent draws per-individual)" + ] + }, + { + "cell_type": "markdown", + "id": "773d934f", + "metadata": {}, + "source": [ + "Now working through utils towards integration" + ] + }, + { + "cell_type": "code", + "execution_count": 59, + "id": "25892203", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "array([10.48431215, 10.57914048, 9.81841743, 11.41020463, 9.62552831,\n", + " 10.27519832, 9.03924539, 10.37692697, 10.03343893])" + ] + }, + "execution_count": 59, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "ssu.sample(dist='normal', par1=10, par2=1, size=len(sim.people))" + ] + }, + { + "cell_type": "code", + "execution_count": 60, + "id": "a24b7dde", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "array([7, 6])" + ] + }, + "execution_count": 60, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "ssu.set_seed(0)\n", + "# Choice (potentially) adds correlation across the entire population\n", + "# The following is fine, because the draws are independent per agent\n", + "np.random.choice([True, False], size=5)\n", + "\n", + "# But choosing k of n agents is a centralized thing, right?\n", + "# Internally, this must use just two random numbers, not one per agent\n", + "np.random.choice(sim.people.uid, 2, replace=False)\n", + "\n", + "# Centralized choices in the model:\n", + "# * Choosing k initial cases (could be a prob per agent instead?)\n", + "# * Choosing k agents to start/stop ART\n", + "# ~~~ Very few! Just need to encourage programming away from this pattern." + ] + }, + { + "cell_type": "markdown", + "id": "7d20df90", + "metadata": {}, + "source": [ + "# Play with random number streams" + ] + }, + { + "cell_type": "code", + "execution_count": 61, + "id": "55e763ab", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "array([0.63696169, 0.26978671, 0.04097352, 0.01652764, 0.81327024])" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/plain": [ + "array([0.51182162, 0.9504637 , 0.14415961, 0.94864945, 0.31183145])" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "0.6369616873214543\n", + "0.2697867137638703\n", + "0.04097352393619469\n", + "0.016527635528529094\n", + "0.8132702392002724\n", + "\n", + "0.5118216247002567\n", + "0.9504636963259353\n", + "0.14415961271963373\n", + "0.9486494471372439\n", + "0.31183145201048545\n", + "\n", + "0.6369616873214543\n", + "0.5118216247002567\n", + "\n", + "0.2697867137638703\n", + "0.9504636963259353\n", + "\n", + "0.04097352393619469\n", + "0.14415961271963373\n", + "\n", + "0.016527635528529094\n", + "0.9486494471372439\n", + "\n", + "0.8132702392002724\n", + "0.31183145201048545\n" + ] + } + ], + "source": [ + "r1 = np.random.default_rng(seed=0)\n", + "display(r1.random(5))\n", + "\n", + "r2 = np.random.default_rng(seed=1)\n", + "display(r2.random(5))\n", + "\n", + "r1 = np.random.default_rng(seed=0)\n", + "r2 = np.random.default_rng(seed=1)\n", + "for i in range(5):\n", + " print(r1.random())\n", + "print()\n", + "for i in range(5):\n", + " print(r2.random())\n", + "\n", + "r1 = np.random.default_rng(seed=0)\n", + "r2 = np.random.default_rng(seed=1)\n", + "for i in range(5):\n", + " print()\n", + " print(r1.random())\n", + " print(r2.random())\n", + "\n", + "# So! Multiple random number streams are easy and work as expected." + ] + }, + { + "cell_type": "code", + "execution_count": 125, + "id": "a59dc80c", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "Generator(PCG64) at 0x17E474740" + ] + }, + "execution_count": 125, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "r1" + ] + }, + { + "cell_type": "markdown", + "id": "51a1d9ae", + "metadata": {}, + "source": [ + "# Let's do some performance evaluation" + ] + }, + { + "cell_type": "code", + "execution_count": 83, + "id": "2df72fc7", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Draw k: 0.316 ms\n", + "Draw n, keep k: 43.8 ms\n" + ] + } + ], + "source": [ + "# k=1 person in a population of n=100000 get infected. What difference does it make if we draw k vs n?\n", + "r1 = np.random.default_rng(seed=0)\n", + "\n", + "k = 100\n", + "n = 100_000\n", + "trials = 100\n", + "import sciris as sc\n", + "\n", + "t = sc.timer('Draw k')\n", + "for i in range(trials):\n", + " v = r1.random(size=k)\n", + "t.toc()\n", + "\n", + "t = sc.timer('Draw n, keep k')\n", + "for i in range(trials):\n", + " v = r1.random(size=n)[:k]\n", + "t.toc()\n", + "\n", + "# UGH, much slower, but still doesn't take all that much time compared to other calculations\n", + "# Sample things like prognosis (once) for each agent at creation...\n", + "# Okay for S->I... except not: in HIV, prognosis depends on age and possibly other factors\n", + "# Plus pure burn if nobody infected on a given step" + ] + }, + { + "cell_type": "code", + "execution_count": 121, + "id": "14ecb71a", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Draw k from separate streams [list comprehension]: 0.460 s\n", + "Draw k from separate streams [map]: 44.1 ms\n" + ] + } + ], + "source": [ + "# What about separate generators per-agent?\n", + "\n", + "r = [np.random.default_rng(seed=s) for s in range(n)]\n", + "t = sc.timer('Draw k from separate streams [list comprehension]')\n", + "for i in range(trials):\n", + " v = [z.choice([True, False]) for i,z in enumerate(r) if i < k]\n", + "t.toc()\n", + "\n", + "r = [np.random.default_rng(seed=s) for s in range(n)]\n", + "t = sc.timer('Draw k from separate streams [map]')\n", + "for i in range(trials):\n", + " #v = list(map(lambda x: x.random(), r[:k]))\n", + " v = list(map(lambda x: x.choice([True, False]), r[:k]))\n", + "t.toc()\n", + "\n", + "# Map is much faster!" + ] + }, + { + "cell_type": "markdown", + "id": "1ad30476", + "metadata": {}, + "source": [ + "# Binomial filtering" + ] + }, + { + "cell_type": "code", + "execution_count": 106, + "id": "008f9dac", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "array([1, 8])" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "text/plain": [ + "array([1, 8])" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "denom_conds = sim.people.female\n", + "uids_to_choose_from = ss.true(denom_conds)\n", + "\n", + "ssu.set_seed(0)\n", + "# Uses len(arr) random numbers:\n", + "u1 = ss.binomial_filter(0.7, uids_to_choose_from) # return arr[(np.random.random(len(arr)) < prob).nonzero()[0]]\n", + "display(u1)\n", + "\n", + "r1 = np.random.default_rng(seed=0)\n", + "def binomial_filter(prob, arr):\n", + " # Uses len(uid map) random numbers\n", + " return arr[(r1.random(len(sim.people._uid_map)) < prob)[arr].nonzero()[0]]\n", + "\n", + "u2 = binomial_filter(0.7, uids_to_choose_from)\n", + "display(u2)" + ] + }, + { + "cell_type": "markdown", + "id": "41813fbb", + "metadata": {}, + "source": [ + "## Point is greater comparability between two simulations.\n", + "\n", + "Can imagine all this working until the point at which the uid map grows in one simulation (baseline) but not another (intervention). Resetting the seed each timestep helps, if/when the (intervention) sim gets a similar uid growth. Oh but this is assuming the uids grow in big chunks. Probably need to be drawing enough random numbers for each decision to cover all (live) agents in a consistent manner.\n", + "\n", + "Need to draw not len(uid map) but rather len(people ever in simulation)\n", + "\n", + "NEED: equivalent of one stream per agent-decision-timestep\n", + "\n", + "Above concept uses \"slots\" where a agent-decision indexes into a block of rngs per timestep. Problem is that blocksize is tied to population size, which will cause offsets between simulations even with per-step seed resetting.\n", + "\n", + "* Fixed block size is possible, but inefficient.\n", + "* Can't use current population size because sub-blocks are selected from one RNG sequentially\n", + "* [ --SUBBLOCK_1-- ][ --SUBBLOCK_2-- ]...[ --SUBBLOCK_N-- ]\n", + "* If length of subblock (e.g. number of uids) grows, then start of next subblock is off\n", + "* --> Use different rngs for each SUBBLOCK, or fix seed at beginning of each SUBBLOCK?!\n", + "* --> Still need to draw full uid size and sample within so that decisions correctly map to agents\n", + "* --> Is RNG resetting fast?\n", + "\n", + "* At some point, per-agent RNG streams with slots per decision-timestep becomes more efficient" + ] + }, + { + "cell_type": "code", + "execution_count": 119, + "id": "277323c3", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "set_seed: 0.175 s\n" + ] + } + ], + "source": [ + "t = sc.timer('set_seed')\n", + "for i in range(100_000):\n", + " ssu.set_seed(i)\n", + "t.toc()\n", + "# Seems plenty fast" + ] + }, + { + "cell_type": "markdown", + "id": "114c06aa", + "metadata": {}, + "source": [ + "grow() happens early in a timestep, so\n", + "* TIME0 with pop N0-->N1: [seed 0, draw N0][seed 1, draw N0]...[seed d, draw N0][seed d+1, draw N1][seed d+2, draw N1]...[seed D, draw N1]\n", + "* Same for time 1 with seeds offset by a fixed amount more than D\n", + "* Works because the len of the uid_map (N0, N1, ...) only grows as new agents are added at the end of the population arrays\n", + "* Long running simulations, at the end, will still be drawing numbers for long-since-dead/removed agents - inefficient, but fixed problem of needing to know the final total number of agents\n", + "\n", + "What if timestep flow conditionally adds or removes a decision?\n", + "* Each decision could have a unique offset instead of going sequentially\n", + "* Offset per-timestep would be max of unique within-timestep offsets\n", + "\n", + "Use case is comparing two simulations with different networks, A vs B. The networks will use random numbers in forming partners, and each can have their own offset so that decisions that come later in the flow would not be affected by the network's use of random numbers. There WOULD be differences due to the different pairs created, however.\n", + "* If edge 1-2 exists in both A and B, being consistent about transmission on this edge between sims would be tricky.\n", + "* Could sample transmission from approx N^2 (per layer?!) (N*(N-1)/2 if heterosexual) and select down to k active edges...\n", + "* Do network layers get \"flattened\" before transmission?\n", + "* If transmission is not consistent... what's the point\n", + "* Any deviation in the network throws everything off?\n", + "\n", + "* Confirmed uses length(edges) random numbers for p1-->p2 and again length(edges) for p2-->p1\n", + "\n", + "* OR!!! Reset seed on a per-edge basis in some clever way? That would require one call per edge instead of a length(edges) vectorized call.\n", + "* Make an RNG per-edge for transmission related inquiries, seeded based on node properties?\n", + "\n", + "Switch to node-based view?\n", + "* Node i connected to k other nodes\n", + "* k chances to get infected on this step\n", + "* Combine probability of infection to a single RNG draw\n", + "* Use one slot per agent\n", + "* Do I get infected on this timestep?\n", + "* If yes, need to determine from whom - sample 1 of k proportional to prob, roulette --> or keep list of possible sources and tree later?\n", + "* NOTE also for each pathogen - so part of module" + ] + }, + { + "cell_type": "code", + "execution_count": 110, + "id": "7dcaf6e1", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "9\n", + "10\n" + ] + } + ], + "source": [ + "print(len(sim.people))\n", + "print(len(sim.people._uid_map))" + ] + }, + { + "cell_type": "code", + "execution_count": 111, + "id": "6d88258a", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "array([10, 11, 12])" + ] + }, + "execution_count": 111, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "sim.people.grow(3)" + ] + }, + { + "cell_type": "code", + "execution_count": 118, + "id": "c59846b6", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "12\n", + "13\n", + "array([9223372036854775807, 0, 1,\n", + " 2, 3, 4,\n", + " 5, 6, 7,\n", + " 8, 9, 10,\n", + " 11])\n" + ] + } + ], + "source": [ + "print(len(sim.people))\n", + "print(len(sim.people._uid_map))\n", + "print(sim.people._uid_map)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7d3bd037", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "IDM", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.12" + }, + "toc": { + "base_numbering": 1, + "nav_menu": {}, + "number_sections": true, + "sideBar": true, + "skip_h1_title": false, + "title_cell": "Table of Contents", + "title_sidebar": "Contents", + "toc_cell": false, + "toc_position": { + "height": "calc(100% - 180px)", + "left": "10px", + "top": "150px", + "width": "328.6px" + }, + "toc_section_display": true, + "toc_window_display": true + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/tests/simple.py b/tests/simple.py index e348d2e3..029bfb72 100644 --- a/tests/simple.py +++ b/tests/simple.py @@ -8,10 +8,10 @@ ppl = ss.People(10000) -ppl.networks = ss.ndict(ss.simple_sexual(), ss.maternal()) +ppl.networks = ss.ndict(ss.simple_embedding(), ss.maternal()) hiv = ss.HIV() -hiv.pars['beta'] = {'simple_sexual': [0.0008, 0.0004], 'maternal': [0.2, 0]} +hiv.pars['beta'] = {'simple_embedding': [0.0008, 0.0004], 'maternal': [0.2, 0]} sim = ss.Sim(people=ppl, modules=[hiv, ss.Gonorrhea(), ss.Pregnancy()]) sim.initialize() @@ -19,5 +19,4 @@ plt.figure() plt.plot(sim.tivec, sim.results.hiv.n_infected) -plt.title('HIV number of infections') - +plt.title('HIV number of infections') \ No newline at end of file From 58b517ff4544c2c2b23858d0c185f3ed6473a671 Mon Sep 17 00:00:00 2001 From: Daniel Klein Date: Fri, 1 Sep 2023 23:55:04 -0700 Subject: [PATCH 02/34] Getting closer with new "embedded" network algorithm --- stisim/modules.py | 10 +++++--- stisim/networks.py | 60 ++++++++++++++++++++++------------------------ stisim/sim.py | 2 +- stisim/streams.py | 14 +++++------ tests/simple.py | 12 ++++++---- 5 files changed, 52 insertions(+), 46 deletions(-) diff --git a/stisim/modules.py b/stisim/modules.py index 97c78250..b7d41ed7 100644 --- a/stisim/modules.py +++ b/stisim/modules.py @@ -150,6 +150,7 @@ def update_states(self, sim): def make_new_cases(self, sim): """ Add new cases of module, through transmission, incidence, etc. """ + n_new_cases = 0 # number of new cases made pars = sim.pars[self.name] for k, layer in sim.people.networks.items(): if k in pars['beta']: @@ -158,9 +159,12 @@ def make_new_cases(self, sim): for a, b, beta, rng in [[layer['p1'], layer['p2'], pars['beta'][k][0], self.rng_trans_ab], [layer['p2'], layer['p1'], pars['beta'][k][1], self.rng_trans_ba]]: # probability of a->b transmission p_transmit = rel_trans[a] * rel_sus[b] * layer['beta'] * beta - new_cases = rng.random(len(a)) < p_transmit # TODO: Convert to flat list of N probs and block sample - if new_cases.any(): - self.set_prognoses(sim, b[new_cases]) + #new_cases = rng.random(len(a)) < p_transmit # TODO: Convert to flat list of N probs and block sample + new_cases = rng.bernoulli_filter(p_transmit, b) + n_new_cases += len(new_cases) + if len(new_cases): + self.set_prognoses(sim, new_cases) + return n_new_cases def set_prognoses(self, sim, uids): pass diff --git a/stisim/networks.py b/stisim/networks.py index 8b0b9778..257b0180 100644 --- a/stisim/networks.py +++ b/stisim/networks.py @@ -350,62 +350,60 @@ def add_pairs(self, people, ti=None): # Make p1 the shorter array if len(available_f) < len(available_m): + switch = True # Want p1 as m in the end p1 = available_f p2 = available_m else: + switch = False # Want p1 as m in the end p1 = available_m p2 = available_f loc1 = self.rng_pair_12.random(arr=p1) loc2 = self.rng_pair_21.random(arr=p2) - done = False - p1v = np.tile(loc1, (len(loc2),1)) p2v = np.tile(loc2[:,np.newaxis], len(loc1)) - d = np.absolute(p2v-p1v) - - pairs = [] + d_full = np.absolute(p2v-p1v) + d = d_full.copy() - while not done: - closest1 = d.argmin(axis=0) - # with axis=0, 0 at loc1[0] will be closest to closest1[0] at loc2[ closest1[0] ] + unmatched_p1i = np.arange(len(p1)) + unmatched_p2i = np.arange(len(p2)) - ele = np.unique(closest1) - closest2 = d.argmin(axis=1) + pairs = [] - #loc1[ closest2[np.unique(closest1)]] is near loc2[ele] - # so pair p2=ele with p1=closest2[np.unique(closest1)] - pairs.append( (p1[closest2[np.unique(closest1)]], p2[ele]) ) + while len(unmatched_p1i)>0: + # Perhaps more efficient to change up directionality of matching at times? + p2i_closest_to_each_p1 = d.argmin(axis=0) + selected_up2i, selected_up1i = np.unique(p2i_closest_to_each_p1, return_index=True) + selected_p1i = unmatched_p1i[selected_up1i] + selected_p2i = unmatched_p2i[selected_up2i] - # WIP! + # loc1[selected_p1i] should be close to loc2[selected_p2i] + pairs.append( (p1[selected_p1i], p2[selected_p2i]) ) - # remove pairs and repeat - unmatched_inds_1 = np.setdiff1d(np.arange(len(p1)), closest2[np.unique(closest1)]) - unmatched_inds_2 = np.setdiff1d(np.arange(len(p2)), ele) - #p1 = np.setdiff1d( p1, p1[closest2[np.unique(closest1)]]) - #p2 = np.setdiff1d(p2, p2[ele]) - p1 = p1[unmatched_inds_1] - p2 = p2[unmatched_inds_2] + # Remove pairs and repeat + unmatched_p1i = np.setdiff1d(unmatched_p1i, selected_p1i) + unmatched_p2i = np.setdiff1d(unmatched_p2i, selected_p2i) - d = d[unmatched_inds_1, unmatched_inds_2] + # Trim distance matrix + d = d_full[np.ix_(unmatched_p2i, unmatched_p1i)] + print(f'Matching with {len(unmatched_p1i)} to go') + pairs = np.concatenate(pairs, axis=1) + n_pairs = pairs.shape[1] - if len(available_m) <= len(available_f): - p1 = available_m - p2 = self.rng_pair_12.choice(available_f, len(p1), replace=False) # TODO: Stream-ify - else: - p2 = available_f - p1 = self.rng_pair_21.choice(available_m, len(p2), replace=False) # TODO: Stream-ify - - beta = np.ones_like(p1) - dur = self.rng_mean_dur.poisson(self.mean_dur, len(p1)) # TODO: Stream-ify + (p1, p2) = (pairs[1], pairs[0]) if switch else (pairs[0], pairs[1]) self['p1'] = np.concatenate([self['p1'], p1]) self['p2'] = np.concatenate([self['p2'], p2]) + + beta = np.ones(n_pairs) + dur = self.rng_mean_dur.poisson(self.mean_dur, n_pairs) # TODO: Stream-ify self['beta'] = np.concatenate([self['beta'], beta]) self['dur'] = np.concatenate([self['dur'], dur]) + return n_pairs + class hpv_network(Network): def __init__(self, pars=None): diff --git a/stisim/sim.py b/stisim/sim.py index 5888438e..506a7fe0 100644 --- a/stisim/sim.py +++ b/stisim/sim.py @@ -325,7 +325,7 @@ def step(self): raise AlreadyRunError('Simulation already complete (call sim.initialize() to re-run)') # Update states, modules, partnerships - self.streams.step(self.ti) + self.streams.step(self.ti+1) # on first step, ti=0, but 0 used for initialization self.people.update(self) self.apply_interventions() self.update_modules() diff --git a/stisim/streams.py b/stisim/streams.py index 3812437a..543837a9 100644 --- a/stisim/streams.py +++ b/stisim/streams.py @@ -1,5 +1,6 @@ import numpy as np import sciris as sc +import stisim as ss __all__ = ['Streams', 'Stream'] @@ -10,7 +11,7 @@ class Streams: """ def __init__(self): - self.streams = [] + self._streams = ss.ndict() return def add(self, stream): @@ -18,13 +19,12 @@ def add(self, stream): # Return value will be used as the seed offset for this stream # Depends on order, which will become a problem - need a better solution - n = len(self.streams) - self.streams.append(stream) - + n = len(self._streams) + self._streams.append(stream) return n def step(self, ti): - for stream in self.streams: + for stream in self._streams.dict_values(): stream.step(ti) return @@ -89,12 +89,12 @@ def block_size(self): return len(self.ppl._uid_map) def random(self, arr): - return (super(Stream, self).random(self.block_size))[arr] + return super(Stream, self).random(self.block_size)[arr] def bernoulli(self, prob, arr): #return super(Stream, self).choice([True, False], size=self.block_size, p=[prob, 1-prob]) # very slow #return (super(Stream, self).binomial(n=1, p=prob, size=self.block_size))[arr].astype(bool) # pretty fast - return (self.random(self.block_size) < prob)[arr] # fastest + return super(Stream, self).random(self.block_size)[arr] < prob # fastest def bernoulli_filter(self, prob, arr): #return arr[self.bernoulli(prob, arr).nonzero()[0]] diff --git a/tests/simple.py b/tests/simple.py index 029bfb72..96d4b67d 100644 --- a/tests/simple.py +++ b/tests/simple.py @@ -11,12 +11,16 @@ ppl.networks = ss.ndict(ss.simple_embedding(), ss.maternal()) hiv = ss.HIV() -hiv.pars['beta'] = {'simple_embedding': [0.0008, 0.0004], 'maternal': [0.2, 0]} +#hiv.pars['beta'] = {'simple_embedding': [0.0008, 0.0004], 'maternal': [0.2, 0]} +hiv.pars['beta'] = {'simple_embedding': [0.02, 0.01], 'maternal': [0.2, 0]} sim = ss.Sim(people=ppl, modules=[hiv, ss.Gonorrhea(), ss.Pregnancy()]) sim.initialize() sim.run() -plt.figure() -plt.plot(sim.tivec, sim.results.hiv.n_infected) -plt.title('HIV number of infections') \ No newline at end of file +fig, axv = plt.subplots(2,1, sharex=True) +axv[0].plot(sim.tivec, sim.results.hiv.n_infected) +axv[0].set_title('HIV number of infections') +axv[1].plot(sim.tivec, sim.results.gonorrhea.n_infected) +axv[1].set_title('Gonorrhea number of infections') +print('Done') \ No newline at end of file From 1f424ee849e764b4e85b90f7cb8a765de02eda1f Mon Sep 17 00:00:00 2001 From: Daniel Klein Date: Tue, 5 Sep 2023 00:49:39 -0700 Subject: [PATCH 03/34] Now allowing modules to choose a seed_offset to as to go at the end if variable between simulations. Experimenting in ART. Added logic to do transmission by acquiring node rather than edge. Slow, but rng safe. Discovered np.random.poisson in Gonorrhea set_prognosis - likely needs to be passed index of infecting agent, which I do not currently calculate. TODO! --- stisim/gonorrhea.py | 6 +++++- stisim/hiv.py | 4 ++-- stisim/modules.py | 11 ++++++++--- stisim/sim.py | 3 ++- stisim/streams.py | 32 +++++++++++++++++++++++++------- tests/simple.py | 26 ++++++++++++++++++++------ 6 files changed, 62 insertions(+), 20 deletions(-) diff --git a/stisim/gonorrhea.py b/stisim/gonorrhea.py index 72c577a1..46f41bfa 100644 --- a/stisim/gonorrhea.py +++ b/stisim/gonorrhea.py @@ -29,6 +29,10 @@ def __init__(self, pars=None): 'initial': 3, 'eff_condoms': 0.7, }, self.pars) + + + self.rng_dur_inf = ss.Stream('dur_inf') + return def update_states(self, sim): @@ -53,7 +57,7 @@ def set_prognoses(self, sim, uids): self.infected[uids] = True self.ti_infected[uids] = sim.ti - dur = sim.ti + np.random.poisson(self.pars['dur_inf']/sim.pars.dt, len(uids)) + dur = sim.ti + self.rng.poisson(self.pars['dur_inf']/sim.pars.dt, len(uids)) # By whom infected from??? dead = self.rng_dead.bernoulli(self.pars.p_death, uids) self.ti_recovered[uids[~dead]] = dur[~dead] diff --git a/stisim/hiv.py b/stisim/hiv.py index 8319952e..772272da 100644 --- a/stisim/hiv.py +++ b/stisim/hiv.py @@ -64,8 +64,8 @@ def __init__(self, t: np.array, capacity: np.array): self.t = sc.promotetoarray(t) self.capacity = sc.promotetoarray(capacity) - self.rng_add_ART = ss.Stream('add_ART') - self.rng_remove_ART = ss.Stream('remove_ART') + self.rng_add_ART = ss.Stream('add_ART', seed_offset=100) + self.rng_remove_ART = ss.Stream('remove_ART', seed_offset=101) return diff --git a/stisim/modules.py b/stisim/modules.py index b7d41ed7..c5ecebec 100644 --- a/stisim/modules.py +++ b/stisim/modules.py @@ -158,9 +158,14 @@ def make_new_cases(self, sim): rel_sus = (self.susceptible & sim.people.alive).astype(float) for a, b, beta, rng in [[layer['p1'], layer['p2'], pars['beta'][k][0], self.rng_trans_ab], [layer['p2'], layer['p1'], pars['beta'][k][1], self.rng_trans_ba]]: # probability of a->b transmission - p_transmit = rel_trans[a] * rel_sus[b] * layer['beta'] * beta - #new_cases = rng.random(len(a)) < p_transmit # TODO: Convert to flat list of N probs and block sample - new_cases = rng.bernoulli_filter(p_transmit, b) + p_tran_edge = rel_trans[a] * rel_sus[b] * layer['beta'] * beta + + # Will need to be more efficient here - can maintain edge to node matrix + node_from_edge = np.zeros( (len(sim.people._uid_map), len(p_tran_edge)) ) + node_from_edge[b, np.arange(len(b))] = 1 + p_acq_node = np.dot(node_from_edge, p_tran_edge) # calculated for all nodes, including those outside b + new_cases = rng.bernoulli_filter(p_acq_node[b], b) + n_new_cases += len(new_cases) if len(new_cases): self.set_prognoses(sim, new_cases) diff --git a/stisim/sim.py b/stisim/sim.py index 506a7fe0..12716d0b 100644 --- a/stisim/sim.py +++ b/stisim/sim.py @@ -75,6 +75,7 @@ def initialize(self, popdict=None, reset=False, **kwargs): ss.set_seed(self.pars['rand_seed']) # Reset the random seed before the population is created # Initialize the core sim components + self.streams.initialize(self.pars['rand_seed'] + 1) # +1 ensures that the base seed is not reused self.init_people(popdict=popdict, reset=reset, **kwargs) # Create all the people (the heaviest step) self.init_networks() self.init_results() @@ -87,7 +88,7 @@ def initialize(self, popdict=None, reset=False, **kwargs): # Reset the random seed to the default run seed, so that if the simulation is run with # reset_seed=False right after initialization, it will still produce the same output - ss.set_seed(self.pars['rand_seed'] + 1) + ss.set_seed(self.pars['rand_seed'] + 1) # Hopefully not used now that we have streams # Final steps self.initialized = True diff --git a/stisim/streams.py b/stisim/streams.py index 543837a9..335bb687 100644 --- a/stisim/streams.py +++ b/stisim/streams.py @@ -12,16 +12,32 @@ class Streams: def __init__(self): self._streams = ss.ndict() + self.used_seeds = [] + return + + def initialize(self, base_seed): + self.base_seed = base_seed return def add(self, stream): - """ Add a stream """ + """ + Add a stream + + Can request an offset, will check for overlap + Otherwise, return value will be used as the seed offset for this stream + """ + + if stream.seed_offset is None: + seed = len(self._streams) # Put at end by default + elif stream.seed_offset in self.used_seeds: + raise Exception(f'Requested seed offset {stream.seed_offset} for stream {stream} has already been used.') + else: + seed = stream.seed_offset + self.used_seeds.append(seed) - # Return value will be used as the seed offset for this stream - # Depends on order, which will become a problem - need a better solution - n = len(self._streams) self._streams.append(stream) - return n + + return self.base_seed + seed def step(self, ti): for stream in self._streams.dict_values(): @@ -33,9 +49,11 @@ class Stream(np.random.Generator): Class for tracking one random number stream associated with one decision per timestep """ - def __init__(self, name, seed_offset=0, **kwargs): + def __init__(self, name, seed_offset=None, **kwargs): """ Create a random number stream + + seed_offset will be automatically assigned (sequentially in first-come order) if None name: a name for this Stream, like "coin_flip" uid: an identifier added to the name to make it uniquely identifiable, for example the name or id of the calling class @@ -67,7 +85,7 @@ def initialize(self, sim): self.seed = sim.streams.add(self) if 'bit_generator' not in self.kwargs: - self.kwargs['bit_generator'] = np.random.PCG64(seed=self.seed + self.seed_offset) + self.kwargs['bit_generator'] = np.random.PCG64(seed=self.seed) super().__init__(**self.kwargs) #self.rng = np.random.default_rng(seed=self.seed + self.seed_offset) diff --git a/tests/simple.py b/tests/simple.py index 96d4b67d..900b365d 100644 --- a/tests/simple.py +++ b/tests/simple.py @@ -7,20 +7,34 @@ import matplotlib.pyplot as plt -ppl = ss.People(10000) +ppl = ss.People(10_000) ppl.networks = ss.ndict(ss.simple_embedding(), ss.maternal()) hiv = ss.HIV() #hiv.pars['beta'] = {'simple_embedding': [0.0008, 0.0004], 'maternal': [0.2, 0]} hiv.pars['beta'] = {'simple_embedding': [0.02, 0.01], 'maternal': [0.2, 0]} -sim = ss.Sim(people=ppl, modules=[hiv, ss.Gonorrhea(), ss.Pregnancy()]) -sim.initialize() -sim.run() +sim1 = ss.Sim(people=ppl, modules=[hiv, ss.Gonorrhea(), ss.Pregnancy()]) +sim1.initialize() +sim1.run() + +''' +ppl = ss.People(10_000) +ppl.networks = ss.ndict(ss.simple_embedding(), ss.maternal()) +hiv = ss.HIV() +#hiv.pars['beta'] = {'simple_embedding': [0.0008, 0.0004], 'maternal': [0.2, 0]} +hiv.pars['beta'] = {'simple_embedding': [0.02, 0.01], 'maternal': [0.2, 0]} +hiv.pars['interventions'] = ss.hiv.ART(0, 10_000) +sim2 = ss.Sim(people=ppl, modules=[hiv, ss.Gonorrhea(), ss.Pregnancy()]) +sim2.initialize() +sim2.run() +''' fig, axv = plt.subplots(2,1, sharex=True) -axv[0].plot(sim.tivec, sim.results.hiv.n_infected) +axv[0].plot(sim1.tivec, sim1.results.hiv.n_infected) +#axv[0].plot(sim2.tivec, sim2.results.hiv.n_infected) axv[0].set_title('HIV number of infections') -axv[1].plot(sim.tivec, sim.results.gonorrhea.n_infected) +axv[1].plot(sim1.tivec, sim1.results.gonorrhea.n_infected) +#axv[1].plot(sim2.tivec, sim2.results.gonorrhea.n_infected) axv[1].set_title('Gonorrhea number of infections') print('Done') \ No newline at end of file From 1429c589a249f036619307d761c38967bb461d02 Mon Sep 17 00:00:00 2001 From: Daniel Klein Date: Wed, 6 Sep 2023 22:42:31 -0700 Subject: [PATCH 04/34] First proof of concept! (run simple.py) * Redo transmission and identification of source * Added a poisson stream * Initializing intervention streams * Now set prognosis based on infector UID --- stisim/demographics.py | 22 ++++++++++---------- stisim/gonorrhea.py | 35 ++++++++++++++++---------------- stisim/modules.py | 37 ++++++++++++++++++++++++--------- stisim/networks.py | 2 +- stisim/sim.py | 3 +++ stisim/streams.py | 3 +++ tests/simple.py | 46 +++++++++++++++++++++--------------------- 7 files changed, 85 insertions(+), 63 deletions(-) diff --git a/stisim/demographics.py b/stisim/demographics.py index 780a05f7..a2ec2670 100644 --- a/stisim/demographics.py +++ b/stisim/demographics.py @@ -120,11 +120,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 @@ -133,19 +133,19 @@ 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 = np.random.random(len(to_uids)) < 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): diff --git a/stisim/gonorrhea.py b/stisim/gonorrhea.py index 46f41bfa..40267bc4 100644 --- a/stisim/gonorrhea.py +++ b/stisim/gonorrhea.py @@ -14,14 +14,15 @@ class Gonorrhea(ss.Disease): def __init__(self, pars=None): super().__init__(pars) - self.susceptible = ss.State('susceptible', bool, True) - self.infected = ss.State('infected', bool, False) - self.ti_infected = ss.State('ti_infected', float, 0) - self.ti_recovered = ss.State('ti_recovered', float, 0) - self.ti_dead = ss.State('ti_dead', float, np.nan) # Death due to gonorrhea + self.susceptible = ss.State('susceptible', bool, True) + self.infected = ss.State('infected', bool, False) + self.ti_infected = ss.State('ti_infected', float, 0) + self.ti_recovered = ss.State('ti_recovered', float, 0) + self.ti_dead = ss.State('ti_dead', float, np.nan) # Death due to gonorrhea - self.rng_prog = ss.Stream('prog_dur') - self.rng_dead = ss.Stream('dead') + self.rng_prog = ss.Stream('prog_dur') + self.rng_dead = ss.Stream('dead') + self.rng_dur_inf = ss.Stream('dur_inf') self.pars = ss.omerge({ 'dur_inf': 3, # not modelling diagnosis or treatment explicitly here @@ -30,9 +31,6 @@ def __init__(self, pars=None): 'eff_condoms': 0.7, }, self.pars) - - self.rng_dur_inf = ss.Stream('dur_inf') - return def update_states(self, sim): @@ -52,14 +50,15 @@ def make_new_cases(self, sim): super(Gonorrhea, self).make_new_cases(sim) return - def set_prognoses(self, sim, uids): - self.susceptible[uids] = False - self.infected[uids] = True - self.ti_infected[uids] = sim.ti + def set_prognoses(self, sim, to_uids, from_uids=None): + self.susceptible[to_uids] = False + self.infected[to_uids] = True + self.ti_infected[to_uids] = sim.ti - dur = sim.ti + self.rng.poisson(self.pars['dur_inf']/sim.pars.dt, len(uids)) # By whom infected from??? - dead = self.rng_dead.bernoulli(self.pars.p_death, uids) + #dur = sim.ti + self.rng_dur_inf.poisson(self.pars['dur_inf']/sim.pars.dt, len(to_uids)) # By whom infected from??? TODO + dur = sim.ti + self.rng_dur_inf.poisson(self.pars['dur_inf']/sim.pars.dt, to_uids) # By whom infected from??? TODO + dead = self.rng_dead.bernoulli(self.pars.p_death, to_uids) - self.ti_recovered[uids[~dead]] = dur[~dead] - self.ti_dead[uids[dead]] = dur[dead] + self.ti_recovered[to_uids[~dead]] = dur[~dead] + self.ti_dead[to_uids[dead]] = dur[dead] return diff --git a/stisim/modules.py b/stisim/modules.py index c5ecebec..e5796cd9 100644 --- a/stisim/modules.py +++ b/stisim/modules.py @@ -24,6 +24,7 @@ def __init__(self, pars=None, label=None, requires=None, *args, **kwargs): self.rng_init_cases = ss.Stream('initial_cases') self.rng_trans_ab = ss.Stream('trans_ab') self.rng_trans_ba = ss.Stream('trans_ba') + self.rng_choose_infector = ss.Stream('choose_infector') return @@ -122,7 +123,7 @@ def set_initial_states(self, sim): #initial_cases = ss.binomial_filter_rng(prob=self.pars['initial']/len(sim.people), arr=sim.people.uid, rng=rng, block_size=len(sim.people._uid_map)) initial_cases = self.rng_init_cases.bernoulli_filter(prob=self.pars['initial']/len(sim.people), arr=sim.people.uid) - self.set_prognoses(sim, initial_cases) + self.set_prognoses(sim, initial_cases, from_uids=None) # TODO: sentinel value to indicate seeds? return def init_results(self, sim): @@ -155,23 +156,39 @@ def make_new_cases(self, sim): for k, layer in sim.people.networks.items(): if k in pars['beta']: rel_trans = (self.infected & sim.people.alive).astype(float) - rel_sus = (self.susceptible & sim.people.alive).astype(float) + rel_sus = self.rel_sus * (self.susceptible & sim.people.alive) for a, b, beta, rng in [[layer['p1'], layer['p2'], pars['beta'][k][0], self.rng_trans_ab], [layer['p2'], layer['p1'], pars['beta'][k][1], self.rng_trans_ba]]: # probability of a->b transmission - p_tran_edge = rel_trans[a] * rel_sus[b] * layer['beta'] * beta + if beta == 0: + continue + #p_tran_edge = rel_trans[a] * rel_sus[b] * layer['beta'] * beta # Will need to be more efficient here - can maintain edge to node matrix - node_from_edge = np.zeros( (len(sim.people._uid_map), len(p_tran_edge)) ) - node_from_edge[b, np.arange(len(b))] = 1 - p_acq_node = np.dot(node_from_edge, p_tran_edge) # calculated for all nodes, including those outside b - new_cases = rng.bernoulli_filter(p_acq_node[b], b) + node_from_edge = np.ones( (len(sim.people._uid_map), len(a)) ) + node_from_edge[b, np.arange(len(b))] = 1 - rel_trans[a] * rel_sus[b] * layer['beta'] * beta + #p_acq_node = np.dot(node_from_edge, p_tran_edge) # calculated for all nodes, including those outside b + p_acq_node = 1 - node_from_edge.prod(axis=1) # 1 - (1-p1)*(1-p2)*... + new_cases_bool = rng.bernoulli(p_acq_node[b], b) + new_cases = b[new_cases_bool] + + if not len(new_cases): + continue n_new_cases += len(new_cases) - if len(new_cases): - self.set_prognoses(sim, new_cases) + # Decide whom the infection came from + frm = np.zeros_like(new_cases) + + # Need one random number for each b <-- align by blocksize + r = self.rng_choose_infector.random(new_cases) # Align rng to b + prob = (1-node_from_edge[new_cases]) + cumsum = (prob / ((prob.sum(axis=1)[:,np.newaxis]))).cumsum(axis=1) + frm_idx = np.argmax( cumsum >= r[:,np.newaxis], axis=1) + frm = a[frm_idx] + self.set_prognoses(sim, new_cases, frm) + return n_new_cases - def set_prognoses(self, sim, uids): + def set_prognoses(self, sim, to_uids, from_uids): pass def update_results(self, sim): diff --git a/stisim/networks.py b/stisim/networks.py index 257b0180..dfb5d339 100644 --- a/stisim/networks.py +++ b/stisim/networks.py @@ -398,7 +398,7 @@ def add_pairs(self, people, ti=None): self['p2'] = np.concatenate([self['p2'], p2]) beta = np.ones(n_pairs) - dur = self.rng_mean_dur.poisson(self.mean_dur, n_pairs) # TODO: Stream-ify + dur = self.rng_mean_dur.poisson(self.mean_dur, p1) # TODO: Stream-ify... trying using p1 self['beta'] = np.concatenate([self['beta'], beta]) self['dur'] = np.concatenate([self['dur'], dur]) diff --git a/stisim/sim.py b/stisim/sim.py index 12716d0b..6dd0d91e 100644 --- a/stisim/sim.py +++ b/stisim/sim.py @@ -286,6 +286,9 @@ def init_interventions(self): errormsg = f'Intervention {intervention} does not seem to be a valid intervention: must be a function or Intervention subclass' raise TypeError(errormsg) + for stream in intervention.streams.values(): + stream.initialize(self) + return def init_analyzers(self): diff --git a/stisim/streams.py b/stisim/streams.py index 335bb687..71fb2548 100644 --- a/stisim/streams.py +++ b/stisim/streams.py @@ -109,6 +109,9 @@ def block_size(self): def random(self, arr): return super(Stream, self).random(self.block_size)[arr] + def poisson(self, lam, arr): + return super(Stream, self).poisson(lam=lam, size=self.block_size)[arr] + def bernoulli(self, prob, arr): #return super(Stream, self).choice([True, False], size=self.block_size, p=[prob, 1-prob]) # very slow #return (super(Stream, self).binomial(n=1, p=prob, size=self.block_size))[arr].astype(bool) # pretty fast diff --git a/tests/simple.py b/tests/simple.py index 900b365d..6383d271 100644 --- a/tests/simple.py +++ b/tests/simple.py @@ -5,36 +5,36 @@ # %% Imports and settings import stisim as ss import matplotlib.pyplot as plt +import sciris as sc +def run_sim(n=2_000, intervention=False): + ppl = ss.People(n) + ppl.networks = ss.ndict(ss.simple_embedding(), ss.maternal()) -ppl = ss.People(10_000) -ppl.networks = ss.ndict(ss.simple_embedding(), ss.maternal()) + hiv = ss.HIV() + #hiv.pars['beta'] = {'simple_embedding': [0.0008, 0.0004], 'maternal': [0.2, 0]} + hiv.pars['beta'] = {'simple_embedding': [0.02, 0.01], 'maternal': [0.2, 0]} -hiv = ss.HIV() -#hiv.pars['beta'] = {'simple_embedding': [0.0008, 0.0004], 'maternal': [0.2, 0]} -hiv.pars['beta'] = {'simple_embedding': [0.02, 0.01], 'maternal': [0.2, 0]} + gon = ss.Gonorrhea() + gon.pars['beta'] = 0.8 -sim1 = ss.Sim(people=ppl, modules=[hiv, ss.Gonorrhea(), ss.Pregnancy()]) -sim1.initialize() -sim1.run() + pars = dict(interventions = [ss.hiv.PrEP(0, 10_000)]) if intervention else {} + sim = ss.Sim(people=ppl, modules=[hiv, ss.Gonorrhea(), ss.Pregnancy()], pars=pars) + sim.initialize() + sim.run() -''' -ppl = ss.People(10_000) -ppl.networks = ss.ndict(ss.simple_embedding(), ss.maternal()) -hiv = ss.HIV() -#hiv.pars['beta'] = {'simple_embedding': [0.0008, 0.0004], 'maternal': [0.2, 0]} -hiv.pars['beta'] = {'simple_embedding': [0.02, 0.01], 'maternal': [0.2, 0]} -hiv.pars['interventions'] = ss.hiv.ART(0, 10_000) -sim2 = ss.Sim(people=ppl, modules=[hiv, ss.Gonorrhea(), ss.Pregnancy()]) -sim2.initialize() -sim2.run() -''' + return sim +sim1, sim2 = sc.parallelize(run_sim, iterkwargs=[{'intervention':False}, {'intervention':True}]) + +# Plot fig, axv = plt.subplots(2,1, sharex=True) -axv[0].plot(sim1.tivec, sim1.results.hiv.n_infected) -#axv[0].plot(sim2.tivec, sim2.results.hiv.n_infected) +axv[0].plot(sim1.tivec, sim1.results.hiv.n_infected, label='Baseline') +axv[0].plot(sim2.tivec, sim2.results.hiv.n_infected, ls=':', label='Intervention') axv[0].set_title('HIV number of infections') -axv[1].plot(sim1.tivec, sim1.results.gonorrhea.n_infected) -#axv[1].plot(sim2.tivec, sim2.results.gonorrhea.n_infected) +axv[1].plot(sim1.tivec, sim1.results.gonorrhea.n_infected, label='Baseline') +axv[1].plot(sim2.tivec, sim2.results.gonorrhea.n_infected, ls=':', label='Intervention') axv[1].set_title('Gonorrhea number of infections') +plt.legend() +plt.show() print('Done') \ No newline at end of file From d5e6d2e7ed195f110544c9ddb3ab1f5d7ff3d8ac Mon Sep 17 00:00:00 2001 From: Daniel Klein Date: Wed, 6 Sep 2023 22:43:15 -0700 Subject: [PATCH 05/34] * Missed hiv.py updates, which are critical to the "Hello World" proof of concept. Added a PrEP intervention that's like ART but modifies rel_sus. --- stisim/hiv.py | 66 ++++++++++++++++++++++++++++++++++++++++++++------- 1 file changed, 57 insertions(+), 9 deletions(-) diff --git a/stisim/hiv.py b/stisim/hiv.py index 772272da..35e554ca 100644 --- a/stisim/hiv.py +++ b/stisim/hiv.py @@ -18,6 +18,7 @@ def __init__(self, pars=None): self.infected = ss.State('infected', bool, False) self.ti_infected = ss.State('ti_infected', float, 0) self.on_art = ss.State('on_art', bool, False) + self.on_prep = ss.State('on_prep', bool, False) self.cd4 = ss.State('cd4', float, 500) self.pars = ss.omerge({ @@ -34,6 +35,8 @@ def update_states(self, sim): """ Update CD4 """ self.cd4[sim.people.alive & self.infected & self.on_art] += (self.pars.cd4_max - self.cd4[sim.people.alive & self.infected & self.on_art])/self.pars.cd4_rate self.cd4[sim.people.alive & self.infected & ~self.on_art] += (self.pars.cd4_min - self.cd4[sim.people.alive & self.infected & ~self.on_art])/self.pars.cd4_rate + + self.rel_sus[sim.people.alive & self.on_prep] = 0.1 return def init_results(self, sim): @@ -49,10 +52,10 @@ def make_new_cases(self, sim): super().make_new_cases(sim) return - def set_prognoses(self, sim, uids): - self.susceptible[uids] = False - self.infected[uids] = True - self.ti_infected[uids] = sim.ti + def set_prognoses(self, sim, to_uids, from_uids=None): + self.susceptible[to_uids] = False + self.infected[to_uids] = True + self.ti_infected[to_uids] = sim.ti # %% Interventions @@ -70,21 +73,22 @@ def __init__(self, t: np.array, capacity: np.array): return def initialize(self, sim): - sim.hiv.results += ss.Result(self.name, 'n_art', sim.npts, dtype=int) + sim.results.hiv += ss.Result(self.name, 'n_art', sim.npts, dtype=int) + self.initialized = True return def apply(self, sim): - if sim.t < self.t[0]: + if sim.ti < self.t[0]: return - capacity = self.capacity[np.where(self.t <= sim.t)[0][-1]] + capacity = self.capacity[np.where(self.t <= sim.ti)[0][-1]] on_art = sim.people.alive & sim.people.hiv.on_art n_change = capacity - np.count_nonzero(on_art) if n_change > 0: # Add more ART - eligible = sim.people.alive & sim.people.hiv.infected & ~sim.people.hiv.on_art - n_eligible = np.count_nonzero(eligible) + eligible = ss.true(sim.people.alive & sim.people.hiv.infected & ~sim.people.hiv.on_art) + n_eligible = len(eligible) #np.count_nonzero(eligible) if n_eligible: inds = self.rng_add_ART.bernoulli_filter(prob=min(n_eligible, n_change)/n_eligible, arr=eligible) sim.people.hiv.on_art[inds] = True @@ -100,6 +104,50 @@ def apply(self, sim): return +class PrEP(ss.Intervention): + + def __init__(self, t: np.array, capacity: np.array): + self.requires = HIV + self.t = sc.promotetoarray(t) + self.capacity = sc.promotetoarray(capacity) + + self.rng_add_PrEP = ss.Stream('add_PrEP', seed_offset=102) + self.rng_remove_PrEP = ss.Stream('remove_PrEP', seed_offset=103) + + return + + def initialize(self, sim): + sim.results.hiv += ss.Result(self.name, 'n_prep', sim.npts, dtype=int) + self.initialized = True + return + + def apply(self, sim): + if sim.ti < self.t[0]: + return + + capacity = self.capacity[np.where(self.t <= sim.ti)[0][-1]] + on_prep = sim.people.alive & sim.people.hiv.on_prep + + n_change = capacity - np.count_nonzero(on_prep) + if n_change > 0: + # Add more PrEP + eligible = ss.true(sim.people.alive & ~sim.people.hiv.infected & ~sim.people.hiv.on_prep) + n_eligible = len(eligible) #np.count_nonzero(eligible) + if n_eligible: + inds = self.rng_add_PrEP.bernoulli_filter(prob=min(n_eligible, n_change)/n_eligible, arr=eligible) + sim.people.hiv.on_prep[inds] = True + elif n_change < 0: + # Take some people off PrEP + eligible = sim.people.alive & sim.people.hiv.on_prep + n_eligible = np.count_nonzero(eligible) + inds = self.rng_remove_PrEP.bernoulli_filter(prob=-n_change/n_eligible, arr=eligible) + sim.people.hiv.on_prep[inds] = False + + # Add result + sim.results.hiv.n_prep = np.count_nonzero(sim.people.alive & sim.people.hiv.on_prep) + + return + #%% Analyzers From dc41908ecbe7707d4b5e40ca24cc2553be035389 Mon Sep 17 00:00:00 2001 From: Daniel Klein Date: Thu, 7 Sep 2023 00:05:37 -0700 Subject: [PATCH 06/34] * Adding HIV deaths and results * Breaking relationships on death --- stisim/gonorrhea.py | 3 +++ stisim/hiv.py | 19 +++++++++++++++---- stisim/modules.py | 7 ++++--- stisim/networks.py | 10 +++++++--- tests/simple.py | 22 ++++++++++++++++++---- 5 files changed, 47 insertions(+), 14 deletions(-) diff --git a/stisim/gonorrhea.py b/stisim/gonorrhea.py index 40267bc4..fc17d4fc 100644 --- a/stisim/gonorrhea.py +++ b/stisim/gonorrhea.py @@ -40,6 +40,9 @@ def update_states(self, sim): gonorrhea_deaths = self.ti_dead <= sim.ti sim.people.alive[gonorrhea_deaths] = False sim.people.ti_dead[gonorrhea_deaths] = sim.ti + + self.results.new_deaths += len(gonorrhea_deaths) + return def update_results(self, sim): diff --git a/stisim/hiv.py b/stisim/hiv.py index 35e554ca..e58a1c40 100644 --- a/stisim/hiv.py +++ b/stisim/hiv.py @@ -15,11 +15,13 @@ def __init__(self, pars=None): super().__init__(pars) self.susceptible = ss.State('susceptible', bool, True) - self.infected = ss.State('infected', bool, False) + self.infected = ss.State('infected', bool, False) self.ti_infected = ss.State('ti_infected', float, 0) - self.on_art = ss.State('on_art', bool, False) - self.on_prep = ss.State('on_prep', bool, False) - self.cd4 = ss.State('cd4', float, 500) + self.on_art = ss.State('on_art', bool, False) + self.on_prep = ss.State('on_prep', bool, False) + self.cd4 = ss.State('cd4', float, 500) + + self.rng_dead = ss.Stream('dead') self.pars = ss.omerge({ 'cd4_min': 100, @@ -37,6 +39,15 @@ def update_states(self, sim): self.cd4[sim.people.alive & self.infected & ~self.on_art] += (self.pars.cd4_min - self.cd4[sim.people.alive & self.infected & ~self.on_art])/self.pars.cd4_rate self.rel_sus[sim.people.alive & self.on_prep] = 0.1 + + hiv_death_prob = 0.1 / (self.pars.cd4_min - self.pars.cd4_max)**2 * (self.pars.cd4_max - self.cd4)**2 + can_die = ss.true(sim.people.alive & sim.people.hiv.infected) + hiv_deaths = self.rng_dead.bernoulli_filter(prob=hiv_death_prob[can_die], arr = can_die) + + sim.people.alive[hiv_deaths] = False + sim.people.ti_dead[hiv_deaths] = sim.ti + self.results['new_deaths'][sim.ti] = len(hiv_deaths) + return def init_results(self, sim): diff --git a/stisim/modules.py b/stisim/modules.py index e5796cd9..3d5f4bf0 100644 --- a/stisim/modules.py +++ b/stisim/modules.py @@ -21,9 +21,9 @@ def __init__(self, pars=None, label=None, requires=None, *args, **kwargs): self.finalized = False # Random number streams - self.rng_init_cases = ss.Stream('initial_cases') - self.rng_trans_ab = ss.Stream('trans_ab') - self.rng_trans_ba = ss.Stream('trans_ba') + self.rng_init_cases = ss.Stream('initial_cases') + self.rng_trans_ab = ss.Stream('trans_ab') + self.rng_trans_ba = ss.Stream('trans_ba') self.rng_choose_infector = ss.Stream('choose_infector') return @@ -134,6 +134,7 @@ def init_results(self, sim): self.results += ss.Result(self.name, 'n_infected', sim.npts, dtype=int) self.results += ss.Result(self.name, 'prevalence', sim.npts, dtype=float) self.results += ss.Result(self.name, 'new_infections', sim.npts, dtype=int) + self.results += ss.Result(self.name, 'new_deaths', sim.npts, dtype=int) return def update(self, sim): diff --git a/stisim/networks.py b/stisim/networks.py index dfb5d339..6e391f98 100644 --- a/stisim/networks.py +++ b/stisim/networks.py @@ -326,7 +326,7 @@ def update(self, people, dt=None): if dt is None: dt = people.dt # First remove any relationships due to end self['dur'] = self['dur'] - dt - active = self['dur'] > 0 + active = (self['dur'] > 0) & people.alive[self['p1']] & people.alive[self['p1']] self['p1'] = self['p1'][active] self['p2'] = self['p2'][active] self['beta'] = self['beta'][active] @@ -345,8 +345,12 @@ def add_pairs(self, people, ti=None): # Find unpartnered males and females - could in principle check other contact layers too # by having the People object passed in here - available_m = np.setdiff1d(people.uid[~people.female], self.members) - available_f = np.setdiff1d(people.uid[people.female], self.members) + available_m = np.setdiff1d(ss.true(~people.female), self.members) + available_f = np.setdiff1d(ss.true(people.female), self.members) + + # slow: + available_m = ss.true(people.alive[available_m]) + available_f = ss.true(people.alive[available_f]) # Make p1 the shorter array if len(available_f) < len(available_m): diff --git a/tests/simple.py b/tests/simple.py index 6383d271..802fca36 100644 --- a/tests/simple.py +++ b/tests/simple.py @@ -16,9 +16,14 @@ def run_sim(n=2_000, intervention=False): hiv.pars['beta'] = {'simple_embedding': [0.02, 0.01], 'maternal': [0.2, 0]} gon = ss.Gonorrhea() - gon.pars['beta'] = 0.8 - - pars = dict(interventions = [ss.hiv.PrEP(0, 10_000)]) if intervention else {} + gon.pars['beta'] = 0.6 + + pars = { + 'start': 1980, + 'end': 2010, + 'interventions': [ss.hiv.PrEP(0, 10_000)] if intervention else [] + #'interventions': [ss.hiv.ART(0, 10_000)] if intervention else [] + } sim = ss.Sim(people=ppl, modules=[hiv, ss.Gonorrhea(), ss.Pregnancy()], pars=pars) sim.initialize() sim.run() @@ -26,15 +31,24 @@ def run_sim(n=2_000, intervention=False): return sim sim1, sim2 = sc.parallelize(run_sim, iterkwargs=[{'intervention':False}, {'intervention':True}]) +#sim1 = run_sim(intervention=False) +#sim2 = run_sim(intervention=True) # Plot -fig, axv = plt.subplots(2,1, sharex=True) +fig, axv = plt.subplots(3,1, sharex=True) axv[0].plot(sim1.tivec, sim1.results.hiv.n_infected, label='Baseline') axv[0].plot(sim2.tivec, sim2.results.hiv.n_infected, ls=':', label='Intervention') axv[0].set_title('HIV number of infections') + axv[1].plot(sim1.tivec, sim1.results.gonorrhea.n_infected, label='Baseline') axv[1].plot(sim2.tivec, sim2.results.gonorrhea.n_infected, ls=':', label='Intervention') axv[1].set_title('Gonorrhea number of infections') + +axv[2].plot(sim1.tivec, sim1.results.hiv.new_deaths, label='Baseline') +axv[2].plot(sim2.tivec, sim2.results.hiv.new_deaths, ls=':', label='Intervention') +axv[2].set_title('HIV Deaths') + + plt.legend() plt.show() print('Done') \ No newline at end of file From 6853d7a97ab1c94db3178b0199fef4dda58512ec Mon Sep 17 00:00:00 2001 From: Daniel Klein Date: Sat, 16 Sep 2023 23:02:35 -0700 Subject: [PATCH 07/34] * Fixed random stream jumping * Now calling finalize * ART reduces transmission * In ART and PrEP, changed capacity to coverage * Bug fix in networks * Updated "simple.py" example, need more testing --- stisim/analyzers.py | 1 - stisim/hiv.py | 24 +++++----- stisim/networks.py | 10 ++++- stisim/sim.py | 11 ++++- stisim/streams.py | 3 +- tests/simple.py | 106 +++++++++++++++++++++++++++++++++++++++----- 6 files changed, 126 insertions(+), 29 deletions(-) diff --git a/stisim/analyzers.py b/stisim/analyzers.py index bbbd5148..6a751441 100644 --- a/stisim/analyzers.py +++ b/stisim/analyzers.py @@ -11,7 +11,6 @@ class Analyzer(ss.Module): pass - class Analyzers(ss.ndict): def __init__(self, *args, type=Analyzer, **kwargs): return super().__init__(self, *args, type=type, **kwargs) \ No newline at end of file diff --git a/stisim/hiv.py b/stisim/hiv.py index e58a1c40..368fe08d 100644 --- a/stisim/hiv.py +++ b/stisim/hiv.py @@ -38,12 +38,12 @@ def update_states(self, sim): self.cd4[sim.people.alive & self.infected & self.on_art] += (self.pars.cd4_max - self.cd4[sim.people.alive & self.infected & self.on_art])/self.pars.cd4_rate self.cd4[sim.people.alive & self.infected & ~self.on_art] += (self.pars.cd4_min - self.cd4[sim.people.alive & self.infected & ~self.on_art])/self.pars.cd4_rate - self.rel_sus[sim.people.alive & self.on_prep] = 0.1 + self.rel_sus[sim.people.alive & ~self.infected & self.on_prep] = 0.04 + self.rel_sus[sim.people.alive & self.infected & self.on_art] = 0.04 - hiv_death_prob = 0.1 / (self.pars.cd4_min - self.pars.cd4_max)**2 * (self.pars.cd4_max - self.cd4)**2 + hiv_death_prob = 0.1 / (self.pars.cd4_min - self.pars.cd4_max)**2 * (self.cd4 - self.pars.cd4_max)**2 can_die = ss.true(sim.people.alive & sim.people.hiv.infected) hiv_deaths = self.rng_dead.bernoulli_filter(prob=hiv_death_prob[can_die], arr = can_die) - sim.people.alive[hiv_deaths] = False sim.people.ti_dead[hiv_deaths] = sim.ti self.results['new_deaths'][sim.ti] = len(hiv_deaths) @@ -73,10 +73,10 @@ def set_prognoses(self, sim, to_uids, from_uids=None): class ART(ss.Intervention): - def __init__(self, t: np.array, capacity: np.array): + def __init__(self, t: np.array, coverage: np.array): self.requires = HIV self.t = sc.promotetoarray(t) - self.capacity = sc.promotetoarray(capacity) + self.coverage = sc.promotetoarray(coverage) self.rng_add_ART = ss.Stream('add_ART', seed_offset=100) self.rng_remove_ART = ss.Stream('remove_ART', seed_offset=101) @@ -92,10 +92,9 @@ def apply(self, sim): if sim.ti < self.t[0]: return - capacity = self.capacity[np.where(self.t <= sim.ti)[0][-1]] + coverage = self.coverage[np.where(self.t <= sim.ti)[0][-1]] on_art = sim.people.alive & sim.people.hiv.on_art - - n_change = capacity - np.count_nonzero(on_art) + n_change = np.round(coverage * sim.people.alive.sum() - np.count_nonzero(on_art)).astype(int) if n_change > 0: # Add more ART eligible = ss.true(sim.people.alive & sim.people.hiv.infected & ~sim.people.hiv.on_art) @@ -117,10 +116,10 @@ def apply(self, sim): class PrEP(ss.Intervention): - def __init__(self, t: np.array, capacity: np.array): + def __init__(self, t: np.array, coverage: np.array): self.requires = HIV self.t = sc.promotetoarray(t) - self.capacity = sc.promotetoarray(capacity) + self.coverage = sc.promotetoarray(coverage) self.rng_add_PrEP = ss.Stream('add_PrEP', seed_offset=102) self.rng_remove_PrEP = ss.Stream('remove_PrEP', seed_offset=103) @@ -136,10 +135,9 @@ def apply(self, sim): if sim.ti < self.t[0]: return - capacity = self.capacity[np.where(self.t <= sim.ti)[0][-1]] + coverage = self.coverage[np.where(self.t <= sim.ti)[0][-1]] on_prep = sim.people.alive & sim.people.hiv.on_prep - - n_change = capacity - np.count_nonzero(on_prep) + n_change = np.round(coverage * sim.people.alive.sum() - np.count_nonzero(on_prep)).astype(int) if n_change > 0: # Add more PrEP eligible = ss.true(sim.people.alive & ~sim.people.hiv.infected & ~sim.people.hiv.on_prep) diff --git a/stisim/networks.py b/stisim/networks.py index 6e391f98..a58d4b83 100644 --- a/stisim/networks.py +++ b/stisim/networks.py @@ -326,7 +326,7 @@ def update(self, people, dt=None): if dt is None: dt = people.dt # First remove any relationships due to end self['dur'] = self['dur'] - dt - active = (self['dur'] > 0) & people.alive[self['p1']] & people.alive[self['p1']] + active = (self['dur'] > 0) & people.alive[self['p1']] & people.alive[self['p2']] self['p1'] = self['p1'][active] self['p2'] = self['p2'][active] self['beta'] = self['beta'][active] @@ -352,6 +352,11 @@ def add_pairs(self, people, ti=None): available_m = ss.true(people.alive[available_m]) available_f = ss.true(people.alive[available_f]) + if not len(available_m) or not len(available_f): + if ss.options.verbose > 1: + print('No pairs to add') + return 0 + # Make p1 the shorter array if len(available_f) < len(available_m): switch = True # Want p1 as m in the end @@ -392,7 +397,8 @@ def add_pairs(self, people, ti=None): # Trim distance matrix d = d_full[np.ix_(unmatched_p2i, unmatched_p1i)] - print(f'Matching with {len(unmatched_p1i)} to go') + if ss.options.verbose > 1: + print(f'Matching with {len(unmatched_p1i)} to go') pairs = np.concatenate(pairs, axis=1) n_pairs = pairs.shape[1] diff --git a/stisim/sim.py b/stisim/sim.py index 6dd0d91e..31382c6a 100644 --- a/stisim/sim.py +++ b/stisim/sim.py @@ -303,7 +303,7 @@ def init_analyzers(self): raise TypeError(errormsg) self.analyzers += analyzer # Add it in - for analyzer in self.analyzers: + for analyzer in self.analyzers.values(): if isinstance(analyzer, ss.Analyzer): analyzer.initialize(self) @@ -433,6 +433,15 @@ def finalize(self, verbose=None): # otherwise the scale factor will be applied multiple times raise AlreadyRunError('Simulation has already been finalized') + for module in self.modules.values(): + module.finalize(self) + + for intervention in self.interventions.values(): + intervention.finalize(self) + + for analyzer in self.analyzers.values(): + analyzer.finalize(self) + # Final settings self.results_ready = True # Set this first so self.summary() knows to print the results self.ti -= 1 # During the run, this keeps track of the next step; restore this be the final day of the sim diff --git a/stisim/streams.py b/stisim/streams.py index 71fb2548..fb40b61b 100644 --- a/stisim/streams.py +++ b/stisim/streams.py @@ -99,7 +99,8 @@ def initialize(self, sim): def step(self, ti): """ Advance to time ti step by jumping """ self.bit_generator.state = self._init_state # restore initial state - self.bit_generator.jumped(jumps=ti) # Take ti jumps + # jumped returns a new bit_generator, use directly instead of setting state? + self.bit_generator.state = self.bit_generator.jumped(jumps=ti).state # Take ti jumps return @property diff --git a/tests/simple.py b/tests/simple.py index 802fca36..bd3a072b 100644 --- a/tests/simple.py +++ b/tests/simple.py @@ -5,34 +5,118 @@ # %% Imports and settings import stisim as ss import matplotlib.pyplot as plt -import sciris as sc +import pandas as pd +import networkx as nx +import seaborn as sns -def run_sim(n=2_000, intervention=False): +class Graph(): + def __init__(self, nodes, edges): + self.graph = nx.from_pandas_edgelist(df=edges, source='p1', target='p2', edge_attr=True) + self.graph.add_nodes_from(nodes.index) + nx.set_node_attributes(self.graph, nodes.transpose().to_dict()) + return + + def draw_nodes(self, filter, ax, **kwargs): + inds = [i for i,n in self.graph.nodes.data() if filter(n)] + nc = ['red' if nd['hiv'] else 'lightgray' for i, nd in self.graph.nodes.data() if i in inds] + ec = ['green' if nd['on_art'] or nd['on_prep'] else 'black' for i, nd in self.graph.nodes.data() if i in inds] + if inds: + nx.draw_networkx_nodes(self.graph, nodelist=inds, pos=pos, ax=ax, node_color=nc, edgecolors=ec, **kwargs) + + def plot(self, pos, ax=None): + kwargs = dict(node_shape='x', node_size=250, linewidths=2, ax=ax) + self.draw_nodes(lambda n: n['dead'], **kwargs) + + kwargs['node_shape'] = 'o' + self.draw_nodes(lambda n: not n['dead'] and n['female'], **kwargs) + + kwargs['node_shape'] = 's' + self.draw_nodes(lambda n: not n['dead'] and not n['female'], **kwargs) + + nx.draw_networkx_edges(self.graph, pos=pos, ax=ax) + nx.draw_networkx_labels(self.graph, labels={i:int(a['cd4']) for i,a in self.graph.nodes.data()}, pos=pos, ax=ax) + nx.draw_networkx_edge_labels(self.graph, edge_labels={(i,j): int(a['dur']) for i,j,a in self.graph.edges.data()}, pos=pos, ax=ax) + + return + +class rng_analyzer(ss.Analyzer): + ''' Simple analyzer to assess if random streams are working ''' + + def __init__(self, **kwargs): + super().__init__(**kwargs) # Initialize the Analyzer object + + self.graphs = {} + return + + def initialize(self, sim): + self.initialized = True + return + + def apply(self, sim): + nodes = pd.DataFrame({ + 'female': sim.people.female.values, + 'dead': sim.people.dead.values, + 'hiv': sim.people.hiv.infected.values, + 'on_art': sim.people.hiv.on_art.values, + 'cd4': sim.people.hiv.cd4.values, + 'on_prep': sim.people.hiv.on_prep.values, + #'ng': sim.people.gonorrhea.infected.values, + }) + + edges = pd.DataFrame(sim.people.networks['simple_embedding'].to_dict()) #sim.people.networks['simple_embedding'].to_df() #TODO: repr issues + + self.graphs[sim.ti] = Graph(nodes, edges) + return + + def finalize(self, sim): + super().finalize(sim) + return + + +def run_sim(n=10, intervention=False): ppl = ss.People(n) - ppl.networks = ss.ndict(ss.simple_embedding(), ss.maternal()) + ppl.networks = ss.ndict(ss.simple_embedding())#, ss.maternal()) hiv = ss.HIV() #hiv.pars['beta'] = {'simple_embedding': [0.0008, 0.0004], 'maternal': [0.2, 0]} - hiv.pars['beta'] = {'simple_embedding': [0.02, 0.01], 'maternal': [0.2, 0]} + hiv.pars['beta'] = {'simple_embedding': [0.06, 0.04]}#, 'maternal': [0.2, 0]} + hiv.pars['initial'] = n/4 gon = ss.Gonorrhea() - gon.pars['beta'] = 0.6 + gon.pars['beta'] = 0.3 pars = { 'start': 1980, 'end': 2010, - 'interventions': [ss.hiv.PrEP(0, 10_000)] if intervention else [] - #'interventions': [ss.hiv.ART(0, 10_000)] if intervention else [] + 'interventions': [ss.hiv.ART(0, 0.2)] if intervention else [], # ss.hiv.PrEP(0, 0.2), + 'rand_seed': 6, + 'analyzers': [rng_analyzer()], } - sim = ss.Sim(people=ppl, modules=[hiv, ss.Gonorrhea(), ss.Pregnancy()], pars=pars) + #sim = ss.Sim(people=ppl, modules=[hiv, ss.Gonorrhea(), ss.Pregnancy()], pars=pars, label=f'Sim with {n} agents and intv={intervention}') + sim = ss.Sim(people=ppl, modules=[hiv], pars=pars, label=f'Sim with {n} agents and intv={intervention}') sim.initialize() sim.run() return sim -sim1, sim2 = sc.parallelize(run_sim, iterkwargs=[{'intervention':False}, {'intervention':True}]) -#sim1 = run_sim(intervention=False) -#sim2 = run_sim(intervention=True) +sim1 = run_sim(intervention=False) +sim2 = run_sim(intervention=True) +#sc.save('sims.obj', [sim1, sim2]) + +# TODO: Parallelization does not work with the current snapshot analyzer +#sim1, sim2 = sc.parallelize(run_sim, iterkwargs=[{'intervention':False}, {'intervention':True}], die=True) + +g1 = sim1.analyzers[0].graphs +g2 = sim2.analyzers[0].graphs + +pos = nx.circular_layout(g1[0].graph) + +for ti in sim1.tivec: + fig, axv = plt.subplots(1,2, figsize=(10,5)) + g1[ti].plot(pos, ax=axv[0]) + g2[ti].plot(pos, ax=axv[1]) + fig.suptitle(f'Time is {ti}') + plt.show() # Plot fig, axv = plt.subplots(3,1, sharex=True) From 26e62fc8ff7fd6305b8241a8c3cf775745fffeb7 Mon Sep 17 00:00:00 2001 From: Daniel Klein Date: Sat, 16 Sep 2023 23:03:54 -0700 Subject: [PATCH 08/34] Removing unnecessary seaborn import. Forgot to mention on preview commit that simple.py now runs two simulations and visualizes them side by side for each time step. --- tests/simple.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/simple.py b/tests/simple.py index bd3a072b..cb90de05 100644 --- a/tests/simple.py +++ b/tests/simple.py @@ -7,7 +7,6 @@ import matplotlib.pyplot as plt import pandas as pd import networkx as nx -import seaborn as sns class Graph(): def __init__(self, nodes, edges): From baa362c881727cc538a091dd893f4e8ca23e7d6a Mon Sep 17 00:00:00 2001 From: Daniel Klein Date: Sat, 16 Sep 2023 23:09:02 -0700 Subject: [PATCH 09/34] Removing gonorrhea from simple plotting. --- tests/simple.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/tests/simple.py b/tests/simple.py index cb90de05..77d1e46c 100644 --- a/tests/simple.py +++ b/tests/simple.py @@ -118,20 +118,20 @@ def run_sim(n=10, intervention=False): plt.show() # Plot -fig, axv = plt.subplots(3,1, sharex=True) +fig, axv = plt.subplots(2,1, sharex=True) axv[0].plot(sim1.tivec, sim1.results.hiv.n_infected, label='Baseline') axv[0].plot(sim2.tivec, sim2.results.hiv.n_infected, ls=':', label='Intervention') axv[0].set_title('HIV number of infections') -axv[1].plot(sim1.tivec, sim1.results.gonorrhea.n_infected, label='Baseline') -axv[1].plot(sim2.tivec, sim2.results.gonorrhea.n_infected, ls=':', label='Intervention') -axv[1].set_title('Gonorrhea number of infections') - -axv[2].plot(sim1.tivec, sim1.results.hiv.new_deaths, label='Baseline') -axv[2].plot(sim2.tivec, sim2.results.hiv.new_deaths, ls=':', label='Intervention') -axv[2].set_title('HIV Deaths') - +axv[1].plot(sim1.tivec, sim1.results.hiv.new_deaths, label='Baseline') +axv[1].plot(sim2.tivec, sim2.results.hiv.new_deaths, ls=':', label='Intervention') +axv[1].set_title('HIV Deaths') +''' Gonorrhea removed for now +axv[2].plot(sim1.tivec, sim1.results.gonorrhea.n_infected, label='Baseline') +axv[2].plot(sim2.tivec, sim2.results.gonorrhea.n_infected, ls=':', label='Intervention') +axv[2].set_title('Gonorrhea number of infections') +''' plt.legend() plt.show() print('Done') \ No newline at end of file From 399c74530a230856d84e2b3a6ec2c78affc4c5a5 Mon Sep 17 00:00:00 2001 From: Daniel Klein Date: Tue, 19 Sep 2023 07:49:03 -0700 Subject: [PATCH 10/34] * Enabling rel_trans in modules * Making ART reduce transmission * Added a simple "ladder" network in which 1<-->2, 3<-->4, ... for testing purposes * Improved graph visualization in simple.py, but will likely revert this to a simple test later. * Adding simple_ladder.py to run and plot the ladder network --- stisim/hiv.py | 2 +- stisim/modules.py | 2 +- stisim/networks.py | 30 +++++++++- tests/simple.py | 49 +++++++++------- tests/test_ladder.py | 133 +++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 192 insertions(+), 24 deletions(-) create mode 100644 tests/test_ladder.py diff --git a/stisim/hiv.py b/stisim/hiv.py index 368fe08d..71d5f96d 100644 --- a/stisim/hiv.py +++ b/stisim/hiv.py @@ -39,7 +39,7 @@ def update_states(self, sim): self.cd4[sim.people.alive & self.infected & ~self.on_art] += (self.pars.cd4_min - self.cd4[sim.people.alive & self.infected & ~self.on_art])/self.pars.cd4_rate self.rel_sus[sim.people.alive & ~self.infected & self.on_prep] = 0.04 - self.rel_sus[sim.people.alive & self.infected & self.on_art] = 0.04 + self.rel_trans[sim.people.alive & self.infected & self.on_art] = 0.04 hiv_death_prob = 0.1 / (self.pars.cd4_min - self.pars.cd4_max)**2 * (self.cd4 - self.pars.cd4_max)**2 can_die = ss.true(sim.people.alive & sim.people.hiv.infected) diff --git a/stisim/modules.py b/stisim/modules.py index 3d5f4bf0..96b6e959 100644 --- a/stisim/modules.py +++ b/stisim/modules.py @@ -156,7 +156,7 @@ def make_new_cases(self, sim): pars = sim.pars[self.name] for k, layer in sim.people.networks.items(): if k in pars['beta']: - rel_trans = (self.infected & sim.people.alive).astype(float) + rel_trans = self.rel_trans * (self.infected & sim.people.alive) rel_sus = self.rel_sus * (self.susceptible & sim.people.alive) for a, b, beta, rng in [[layer['p1'], layer['p2'], pars['beta'][k][0], self.rng_trans_ab], [layer['p2'], layer['p1'], pars['beta'][k][1], self.rng_trans_ba]]: # probability of a->b transmission diff --git a/stisim/networks.py b/stisim/networks.py index a58d4b83..55050074 100644 --- a/stisim/networks.py +++ b/stisim/networks.py @@ -9,7 +9,7 @@ # Specify all externally visible functions this file defines -__all__ = ['Networks', 'Network', 'simple_sexual', 'simple_embedding', 'hpv_network', 'maternal'] +__all__ = ['Networks', 'Network', 'simple_sexual', 'simple_embedding', 'ladder', 'hpv_network', 'maternal'] class Network(sc.objdict): """ @@ -415,6 +415,34 @@ def add_pairs(self, people, ti=None): return n_pairs +class ladder(simple_sexual): + """ + Very simple network for debugging in which edges are: + 1-2, 3-4, 5-6, ... + """ + def __init__(self): + key_dict = { + 'p1': ss.int_, + 'p2': ss.int_, + 'beta': ss.float_, + } + + # Call init for the base class, which sets all the keys + super().__init__(mean_dur=np.iinfo(int).max) + return + + def initialize(self, sim): + n = len(sim.people._uid_map) + self['p1'] = np.arange(0,n,2) # EVEN + self['p2'] = np.arange(1,n,2) # ODD + self['beta'] = np.ones(len(self['p1'])) + self['dur'] = np.iinfo(int).max + return + + def update(self, people, dt=None): + pass + + class hpv_network(Network): def __init__(self, pars=None): diff --git a/tests/simple.py b/tests/simple.py index 77d1e46c..d242f3cf 100644 --- a/tests/simple.py +++ b/tests/simple.py @@ -6,8 +6,11 @@ import stisim as ss import matplotlib.pyplot as plt import pandas as pd +import numpy as np import networkx as nx +plot_graph = True + class Graph(): def __init__(self, nodes, edges): self.graph = nx.from_pandas_edgelist(df=edges, source='p1', target='p2', edge_attr=True) @@ -33,8 +36,8 @@ def plot(self, pos, ax=None): self.draw_nodes(lambda n: not n['dead'] and not n['female'], **kwargs) nx.draw_networkx_edges(self.graph, pos=pos, ax=ax) - nx.draw_networkx_labels(self.graph, labels={i:int(a['cd4']) for i,a in self.graph.nodes.data()}, pos=pos, ax=ax) - nx.draw_networkx_edge_labels(self.graph, edge_labels={(i,j): int(a['dur']) for i,j,a in self.graph.edges.data()}, pos=pos, ax=ax) + nx.draw_networkx_labels(self.graph, labels={i:int(a['cd4']) for i,a in self.graph.nodes.data()}, font_size=8, pos=pos, ax=ax) + nx.draw_networkx_edge_labels(self.graph, edge_labels={(i,j): int(a['dur']) for i,j,a in self.graph.edges.data()}, font_size=8, pos=pos, ax=ax) return @@ -72,50 +75,54 @@ def finalize(self, sim): return -def run_sim(n=10, intervention=False): +def run_sim(n=25, intervention=False, analyze=False): ppl = ss.People(n) ppl.networks = ss.ndict(ss.simple_embedding())#, ss.maternal()) hiv = ss.HIV() #hiv.pars['beta'] = {'simple_embedding': [0.0008, 0.0004], 'maternal': [0.2, 0]} hiv.pars['beta'] = {'simple_embedding': [0.06, 0.04]}#, 'maternal': [0.2, 0]} - hiv.pars['initial'] = n/4 + hiv.pars['initial'] = 3 - gon = ss.Gonorrhea() - gon.pars['beta'] = 0.3 + #gon = ss.Gonorrhea() + #gon.pars['beta'] = 0.3 pars = { 'start': 1980, 'end': 2010, - 'interventions': [ss.hiv.ART(0, 0.2)] if intervention else [], # ss.hiv.PrEP(0, 0.2), - 'rand_seed': 6, - 'analyzers': [rng_analyzer()], + 'interventions': [ss.hiv.ART(0, 0.1)] if intervention else [], # ss.hiv.PrEP(0, 0.2), + 'rand_seed': 0, + 'analyzers': [rng_analyzer()] if analyze else [], } - #sim = ss.Sim(people=ppl, modules=[hiv, ss.Gonorrhea(), ss.Pregnancy()], pars=pars, label=f'Sim with {n} agents and intv={intervention}') + #sim = ss.Sim(people=ppl, modules=[hiv, gon, ss.Pregnancy()], pars=pars, label=f'Sim with {n} agents and intv={intervention}') sim = ss.Sim(people=ppl, modules=[hiv], pars=pars, label=f'Sim with {n} agents and intv={intervention}') sim.initialize() sim.run() return sim -sim1 = run_sim(intervention=False) -sim2 = run_sim(intervention=True) +n = 10 +sim1 = run_sim(n, intervention=False, analyze=plot_graph) +sim2 = run_sim(n, intervention=True, analyze=plot_graph) #sc.save('sims.obj', [sim1, sim2]) # TODO: Parallelization does not work with the current snapshot analyzer #sim1, sim2 = sc.parallelize(run_sim, iterkwargs=[{'intervention':False}, {'intervention':True}], die=True) -g1 = sim1.analyzers[0].graphs -g2 = sim2.analyzers[0].graphs +if plot_graph: + g1 = sim1.analyzers[0].graphs + g2 = sim2.analyzers[0].graphs -pos = nx.circular_layout(g1[0].graph) + #pos = nx.circular_layout(g1[0].graph) + #n = len(sim1.people._uid_map) # Will eventually need this line due to population growth + pos = {i:(np.cos(2*np.pi*i/n), np.sin(2*np.pi*i/n)) for i in range(n)} -for ti in sim1.tivec: - fig, axv = plt.subplots(1,2, figsize=(10,5)) - g1[ti].plot(pos, ax=axv[0]) - g2[ti].plot(pos, ax=axv[1]) - fig.suptitle(f'Time is {ti}') - plt.show() + for ti in sim1.tivec: + fig, axv = plt.subplots(1,2, figsize=(10,5)) + g1[ti].plot(pos, ax=axv[0]) + g2[ti].plot(pos, ax=axv[1]) + fig.suptitle(f'Time is {ti}') + plt.show() # Plot fig, axv = plt.subplots(2,1, sharex=True) diff --git a/tests/test_ladder.py b/tests/test_ladder.py new file mode 100644 index 00000000..ed0ddcf1 --- /dev/null +++ b/tests/test_ladder.py @@ -0,0 +1,133 @@ +""" +Run test with ladder network +""" + +# %% Imports and settings +import stisim as ss +import matplotlib.pyplot as plt +import pandas as pd +import numpy as np +import networkx as nx + +plot_graph = True + +class Graph(): + def __init__(self, nodes, edges): + self.graph = nx.from_pandas_edgelist(df=edges, source='p1', target='p2', edge_attr=True) + self.graph.add_nodes_from(nodes.index) + nx.set_node_attributes(self.graph, nodes.transpose().to_dict()) + return + + def draw_nodes(self, filter, ax, **kwargs): + inds = [i for i,n in self.graph.nodes.data() if filter(n)] + nc = ['red' if nd['hiv'] else 'lightgray' for i, nd in self.graph.nodes.data() if i in inds] + ec = ['green' if nd['on_art'] or nd['on_prep'] else 'black' for i, nd in self.graph.nodes.data() if i in inds] + if inds: + nx.draw_networkx_nodes(self.graph, nodelist=inds, pos=pos, ax=ax, node_color=nc, edgecolors=ec, **kwargs) + + def plot(self, pos, ax=None): + kwargs = dict(node_shape='x', node_size=250, linewidths=2, ax=ax) + self.draw_nodes(lambda n: n['dead'], **kwargs) + + kwargs['node_shape'] = 'o' + self.draw_nodes(lambda n: not n['dead'] and n['female'], **kwargs) + + kwargs['node_shape'] = 's' + self.draw_nodes(lambda n: not n['dead'] and not n['female'], **kwargs) + + nx.draw_networkx_edges(self.graph, pos=pos, ax=ax) + nx.draw_networkx_labels(self.graph, labels={i:int(a['cd4']) for i,a in self.graph.nodes.data()}, font_size=8, pos=pos, ax=ax) + nx.draw_networkx_edge_labels(self.graph, edge_labels={(i,j): int(a['dur']) for i,j,a in self.graph.edges.data()}, font_size=8, pos=pos, ax=ax) + + return + +class rng_analyzer(ss.Analyzer): + ''' Simple analyzer to assess if random streams are working ''' + + def __init__(self, **kwargs): + super().__init__(**kwargs) # Initialize the Analyzer object + + self.graphs = {} + return + + def initialize(self, sim): + self.initialized = True + return + + def apply(self, sim): + nodes = pd.DataFrame({ + 'female': sim.people.female.values, + 'dead': sim.people.dead.values, + 'hiv': sim.people.hiv.infected.values, + 'on_art': sim.people.hiv.on_art.values, + 'cd4': sim.people.hiv.cd4.values, + 'on_prep': sim.people.hiv.on_prep.values, + #'ng': sim.people.gonorrhea.infected.values, + }) + + edges = pd.DataFrame(sim.people.networks[0].to_dict()) #sim.people.networks['simple_embedding'].to_df() #TODO: repr issues + + self.graphs[sim.ti] = Graph(nodes, edges) + return + + def finalize(self, sim): + super().finalize(sim) + return + + +def run_sim(n=25, intervention=False, analyze=False): + ppl = ss.People(n) + ppl.networks = ss.ndict(ss.ladder())#, ss.maternal()) + + hiv = ss.HIV() + hiv.pars['beta'] = {'ladder': [0.06, 0.04]} + hiv.pars['initial'] = 3 + + pars = { + 'start': 1980, + 'end': 2010, + 'interventions': [ss.hiv.ART(0, 0.1)] if intervention else [], # ss.hiv.PrEP(0, 0.2), + 'rand_seed': 0, + 'analyzers': [rng_analyzer()] if analyze else [], + } + sim = ss.Sim(people=ppl, modules=[hiv], pars=pars, label=f'Sim with {n} agents and intv={intervention}') + sim.initialize() + sim.run() + + return sim + +n = 10 +sim2 = run_sim(n, intervention=True, analyze=plot_graph) +sim1 = run_sim(n, intervention=False, analyze=plot_graph) + +if plot_graph: + g1 = sim1.analyzers[0].graphs + g2 = sim2.analyzers[0].graphs + + pos = {i:(np.cos(2*np.pi*i/n), np.sin(2*np.pi*i/n)) for i in range(n)} + + for ti in sim1.tivec: + fig, axv = plt.subplots(1,2, figsize=(10,5)) + g1[ti].plot(pos, ax=axv[0]) + g2[ti].plot(pos, ax=axv[1]) + fig.suptitle(f'Time is {ti}') + plt.show() + +# Plot +fig, axv = plt.subplots(2,1, sharex=True) +axv[0].plot(sim1.tivec, sim1.results.hiv.n_infected, label='Baseline') +axv[0].plot(sim2.tivec, sim2.results.hiv.n_infected, ls=':', label='Intervention') +axv[0].set_title('HIV number of infections') + +axv[1].plot(sim1.tivec, sim1.results.hiv.new_deaths, label='Baseline') +axv[1].plot(sim2.tivec, sim2.results.hiv.new_deaths, ls=':', label='Intervention') +axv[1].set_title('HIV Deaths') + +''' Gonorrhea removed for now +axv[2].plot(sim1.tivec, sim1.results.gonorrhea.n_infected, label='Baseline') +axv[2].plot(sim2.tivec, sim2.results.gonorrhea.n_infected, ls=':', label='Intervention') +axv[2].set_title('Gonorrhea number of infections') +''' +plt.legend() +plt.show() +print('Done') \ No newline at end of file From f66a8849c422975cc62792e179400d4858ff3dca Mon Sep 17 00:00:00 2001 From: Daniel Klein Date: Wed, 20 Sep 2023 22:44:34 -0400 Subject: [PATCH 11/34] Improved the ladder example: * One member of each dyad is initally HIV+ * 50% of HIV+ are on ART --- stisim/hiv.py | 3 ++- tests/test_ladder.py | 7 +++++-- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/stisim/hiv.py b/stisim/hiv.py index 71d5f96d..a9a4e2cf 100644 --- a/stisim/hiv.py +++ b/stisim/hiv.py @@ -94,7 +94,8 @@ def apply(self, sim): coverage = self.coverage[np.where(self.t <= sim.ti)[0][-1]] on_art = sim.people.alive & sim.people.hiv.on_art - n_change = np.round(coverage * sim.people.alive.sum() - np.count_nonzero(on_art)).astype(int) + infected = sim.people.alive & sim.people.hiv.infected + n_change = np.round(coverage * infected.sum() - on_art.sum()).astype(int) if n_change > 0: # Add more ART eligible = ss.true(sim.people.alive & sim.people.hiv.infected & ~sim.people.hiv.on_art) diff --git a/tests/test_ladder.py b/tests/test_ladder.py index ed0ddcf1..92e2fb26 100644 --- a/tests/test_ladder.py +++ b/tests/test_ladder.py @@ -81,17 +81,20 @@ def run_sim(n=25, intervention=False, analyze=False): hiv = ss.HIV() hiv.pars['beta'] = {'ladder': [0.06, 0.04]} - hiv.pars['initial'] = 3 + hiv.pars['initial'] = 0 pars = { 'start': 1980, 'end': 2010, - 'interventions': [ss.hiv.ART(0, 0.1)] if intervention else [], # ss.hiv.PrEP(0, 0.2), + 'interventions': [ss.hiv.ART(0, 0.5)] if intervention else [], # ss.hiv.PrEP(0, 0.2), 'rand_seed': 0, 'analyzers': [rng_analyzer()] if analyze else [], } sim = ss.Sim(people=ppl, modules=[hiv], pars=pars, label=f'Sim with {n} agents and intv={intervention}') sim.initialize() + + sim.modules['hiv'].set_prognoses(sim, np.arange(0,n,2), from_uids=None) + sim.run() return sim From 2ad911abdf72f87d622dd45e5ca01f969a729526 Mon Sep 17 00:00:00 2001 From: Daniel Klein Date: Thu, 21 Sep 2023 22:51:01 -0400 Subject: [PATCH 12/34] Renaming "ladder" to "stable monogamy" --- stisim/networks.py | 4 ++-- tests/{test_ladder.py => test_stable_monogamy.py} | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) rename tests/{test_ladder.py => test_stable_monogamy.py} (96%) diff --git a/stisim/networks.py b/stisim/networks.py index 55050074..ba65c8ab 100644 --- a/stisim/networks.py +++ b/stisim/networks.py @@ -9,7 +9,7 @@ # Specify all externally visible functions this file defines -__all__ = ['Networks', 'Network', 'simple_sexual', 'simple_embedding', 'ladder', 'hpv_network', 'maternal'] +__all__ = ['Networks', 'Network', 'simple_sexual', 'simple_embedding', 'stable_monogamy', 'hpv_network', 'maternal'] class Network(sc.objdict): """ @@ -415,7 +415,7 @@ def add_pairs(self, people, ti=None): return n_pairs -class ladder(simple_sexual): +class stable_monogamy(simple_sexual): """ Very simple network for debugging in which edges are: 1-2, 3-4, 5-6, ... diff --git a/tests/test_ladder.py b/tests/test_stable_monogamy.py similarity index 96% rename from tests/test_ladder.py rename to tests/test_stable_monogamy.py index 92e2fb26..9c516eec 100644 --- a/tests/test_ladder.py +++ b/tests/test_stable_monogamy.py @@ -1,5 +1,5 @@ """ -Run test with ladder network +Run test with stable_monogamy network """ # %% Imports and settings @@ -77,10 +77,10 @@ def finalize(self, sim): def run_sim(n=25, intervention=False, analyze=False): ppl = ss.People(n) - ppl.networks = ss.ndict(ss.ladder())#, ss.maternal()) + ppl.networks = ss.ndict(ss.stable_monogamy())#, ss.maternal()) hiv = ss.HIV() - hiv.pars['beta'] = {'ladder': [0.06, 0.04]} + hiv.pars['beta'] = {'stable_monogamy': [0.06, 0.04]} hiv.pars['initial'] = 0 pars = { From 970a4696008a3daed23eeb3005921b473dea4e6b Mon Sep 17 00:00:00 2001 From: Daniel Klein Date: Fri, 22 Sep 2023 00:25:32 -0400 Subject: [PATCH 13/34] * Tests for Stream * Added checking to ensure a stream is not called twice in a row without a reset() or a step() --- stisim/streams.py | 36 ++++++++++++++- tests/test_stream.py | 101 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 135 insertions(+), 2 deletions(-) create mode 100644 tests/test_stream.py diff --git a/stisim/streams.py b/stisim/streams.py index fb40b61b..77d635de 100644 --- a/stisim/streams.py +++ b/stisim/streams.py @@ -13,10 +13,12 @@ class Streams: def __init__(self): self._streams = ss.ndict() self.used_seeds = [] + self.initialized = False return def initialize(self, base_seed): self.base_seed = base_seed + self.initialized = True return def add(self, stream): @@ -26,6 +28,8 @@ def add(self, stream): Can request an offset, will check for overlap Otherwise, return value will be used as the seed offset for this stream """ + if not self.initialized: + raise Exception('Please call initialize before adding a stream to Streams.') if stream.seed_offset is None: seed = len(self._streams) # Put at end by default @@ -44,6 +48,20 @@ def step(self, ti): stream.step(ti) return +class NotResetException(Exception): + "Raised when stream is called when not ready." + pass + +def _pre_draw(func): + def check_ready(self, arr): + """ Validation before drawing """ + if not self.ready: + msg = f'Stream {self.name} has already been sampled on this timestep!' + raise NotResetException(msg) + self.ready = False + return func(self, arr) + return check_ready + class Stream(np.random.Generator): """ Class for tracking one random number stream associated with one decision per timestep @@ -73,6 +91,7 @@ def __init__(self, name, seed_offset=None, **kwargs): self.ppl = None # Needed for block size self.initialized = False + self.ready = True return @@ -94,19 +113,32 @@ def initialize(self, sim): self.ppl = sim.people self.initialized = True + self.ready = True + return + + def reset(self): + """ Restore initial state """ + self.bit_generator.state = self._init_state + self.ready = True return def step(self, ti): """ Advance to time ti step by jumping """ - self.bit_generator.state = self._init_state # restore initial state + + # First reset back to the initial state + self.reset() + + # Now take ti jumps # jumped returns a new bit_generator, use directly instead of setting state? - self.bit_generator.state = self.bit_generator.jumped(jumps=ti).state # Take ti jumps + self.bit_generator.state = self.bit_generator.jumped(jumps=ti).state + return @property def block_size(self): return len(self.ppl._uid_map) + @_pre_draw def random(self, arr): return super(Stream, self).random(self.block_size)[arr] diff --git a/tests/test_stream.py b/tests/test_stream.py new file mode 100644 index 00000000..81eb2668 --- /dev/null +++ b/tests/test_stream.py @@ -0,0 +1,101 @@ +""" +Test the Stream object from streams.py +""" + +# %% Imports and settings +import numpy as np +import sciris as sc +import stisim as ss +from stisim.streams import NotResetException + +def make_rng(name='Test'): + """ Create and initialize a stream """ + sim = ss.Sim() + sim.initialize() + rng = ss.Stream('Test') + rng.initialize(sim) + + return rng + + +# %% Define the tests + +def test_sample(n=5): + """ Simple sample """ + sc.heading('Testing stream object') + rng = make_rng(n) + uids = np.arange(n) + + draws = rng.random(uids) + print(f'\nSAMPLE({n}): {draws}') + + return draws + + +def test_reset(n=5): + """ Sample, reset, sample """ + sc.heading('Testing sample, reset, sample') + rng = make_rng(n) + uids = np.arange(n) + + draws1 = rng.random(uids) + print(f'\nSAMPLE({n}): {draws1}') + + print(f'\nRESET') + rng.reset() + + draws2 = rng.random(uids) + print(f'\nSAMPLE({n}): {draws2}') + + return np.all(draws2-draws1 == 0) + + +def test_step(n=5): + """ Sample, step, sample """ + sc.heading('Testing sample, step, sample') + rng = make_rng(n) + uids = np.arange(n) + + draws1 = rng.random(uids) + print(f'\nSAMPLE({n}): {draws1}') + + print(f'\nSTEP(1) - sample should change') + rng.step(1) + + draws2 = rng.random(uids) + print(f'\nSAMPLE({n}): {draws2}') + + return np.all(draws2-draws1 != 0) + + +def test_repeat(n=5): + """ Sample, sample - should raise and exception""" + sc.heading('Testing sample, sample - should raise an exception') + rng = make_rng(n) + uids = np.arange(n) + + draws1 = rng.random(uids) + print(f'\nSAMPLE({n}): {draws1}') + + print(f'\nSAMPLE({n}): [should raise an exception as neither reset() nor step() have been called]') + try: + rng.random(uids) + return False # Should not get here! + except NotResetException as e: + print(f'YAY! Got exception: {e}') + return True + + +# %% Run as a script +if __name__ == '__main__': + # Start timing + T = sc.tic() + + # Run tests + test_sample() + assert test_reset() + assert test_step() + assert test_repeat() + + sc.toc(T) + print('Done.') \ No newline at end of file From 31f8281f295a313fdb4684480880739e772ae5aa Mon Sep 17 00:00:00 2001 From: Daniel Klein Date: Sat, 23 Sep 2023 21:46:14 -0400 Subject: [PATCH 14/34] * Interactive plotting in test_stable_monogamy.py (use arrow keys) * Bug fix in test_stream.py --- tests/test_stable_monogamy.py | 31 +++++++++++++++++++++++++++---- tests/test_stream.py | 10 +++++----- 2 files changed, 32 insertions(+), 9 deletions(-) diff --git a/tests/test_stable_monogamy.py b/tests/test_stable_monogamy.py index 9c516eec..39920370 100644 --- a/tests/test_stable_monogamy.py +++ b/tests/test_stable_monogamy.py @@ -8,6 +8,7 @@ import pandas as pd import numpy as np import networkx as nx +import sys plot_graph = True @@ -24,6 +25,7 @@ def draw_nodes(self, filter, ax, **kwargs): ec = ['green' if nd['on_art'] or nd['on_prep'] else 'black' for i, nd in self.graph.nodes.data() if i in inds] if inds: nx.draw_networkx_nodes(self.graph, nodelist=inds, pos=pos, ax=ax, node_color=nc, edgecolors=ec, **kwargs) + return def plot(self, pos, ax=None): kwargs = dict(node_shape='x', node_size=250, linewidths=2, ax=ax) @@ -38,9 +40,9 @@ def plot(self, pos, ax=None): nx.draw_networkx_edges(self.graph, pos=pos, ax=ax) nx.draw_networkx_labels(self.graph, labels={i:int(a['cd4']) for i,a in self.graph.nodes.data()}, font_size=8, pos=pos, ax=ax) nx.draw_networkx_edge_labels(self.graph, edge_labels={(i,j): int(a['dur']) for i,j,a in self.graph.edges.data()}, font_size=8, pos=pos, ax=ax) - return + class rng_analyzer(ss.Analyzer): ''' Simple analyzer to assess if random streams are working ''' @@ -109,12 +111,33 @@ def run_sim(n=25, intervention=False, analyze=False): pos = {i:(np.cos(2*np.pi*i/n), np.sin(2*np.pi*i/n)) for i in range(n)} - for ti in sim1.tivec: - fig, axv = plt.subplots(1,2, figsize=(10,5)) + fig, axv = plt.subplots(1, 2, figsize=(10,5)) + global ti + ti = 0 + timax = sim1.tivec[-1] + def on_press(event): + print('press', event.key) + sys.stdout.flush() + global ti + if event.key == 'right': + ti = min(ti+1, timax) + elif event.key == 'left': + ti = max(ti-1, 0) + + # Clear + axv[0].clear() + axv[1].clear() + g1[ti].plot(pos, ax=axv[0]) g2[ti].plot(pos, ax=axv[1]) fig.suptitle(f'Time is {ti}') - plt.show() + fig.canvas.draw() + + fig.canvas.mpl_connect('key_press_event', on_press) + + g1[ti].plot(pos, ax=axv[0]) + g2[ti].plot(pos, ax=axv[1]) + plt.show() # Plot fig, axv = plt.subplots(2,1, sharex=True) diff --git a/tests/test_stream.py b/tests/test_stream.py index 81eb2668..5f18877b 100644 --- a/tests/test_stream.py +++ b/tests/test_stream.py @@ -12,7 +12,7 @@ def make_rng(name='Test'): """ Create and initialize a stream """ sim = ss.Sim() sim.initialize() - rng = ss.Stream('Test') + rng = ss.Stream(name) rng.initialize(sim) return rng @@ -23,7 +23,7 @@ def make_rng(name='Test'): def test_sample(n=5): """ Simple sample """ sc.heading('Testing stream object') - rng = make_rng(n) + rng = make_rng() uids = np.arange(n) draws = rng.random(uids) @@ -35,7 +35,7 @@ def test_sample(n=5): def test_reset(n=5): """ Sample, reset, sample """ sc.heading('Testing sample, reset, sample') - rng = make_rng(n) + rng = make_rng() uids = np.arange(n) draws1 = rng.random(uids) @@ -53,7 +53,7 @@ def test_reset(n=5): def test_step(n=5): """ Sample, step, sample """ sc.heading('Testing sample, step, sample') - rng = make_rng(n) + rng = make_rng() uids = np.arange(n) draws1 = rng.random(uids) @@ -71,7 +71,7 @@ def test_step(n=5): def test_repeat(n=5): """ Sample, sample - should raise and exception""" sc.heading('Testing sample, sample - should raise an exception') - rng = make_rng(n) + rng = make_rng() uids = np.arange(n) draws1 = rng.random(uids) From 284883aeae5f25030437e2b51f4f156c5284d027 Mon Sep 17 00:00:00 2001 From: Daniel Klein Date: Mon, 25 Sep 2023 11:56:31 -0400 Subject: [PATCH 15/34] * Adding test for Streams object. * Improving test for Stream object. * Improving streams.py: - Exception handling - reset() - Avoiding dependence of Stream on Sim object, now taking a Streams object and a block_size_object. --- stisim/modules.py | 2 +- stisim/sim.py | 4 +- stisim/streams.py | 36 +++++++-- tests/test_stream.py | 43 ++++++++--- tests/test_streams.py | 169 ++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 231 insertions(+), 23 deletions(-) create mode 100644 tests/test_streams.py diff --git a/stisim/modules.py b/stisim/modules.py index 96b6e959..d8972fad 100644 --- a/stisim/modules.py +++ b/stisim/modules.py @@ -50,7 +50,7 @@ def initialize(self, sim): # Connect the streams to the sim for stream in self.streams.values(): - stream.initialize(sim) + stream.initialize(sim.streams, sim.people._uid_map) self.initialized = True return diff --git a/stisim/sim.py b/stisim/sim.py index 31382c6a..2baa5aae 100644 --- a/stisim/sim.py +++ b/stisim/sim.py @@ -42,7 +42,7 @@ def __init__(self, pars=None, label=None, people=None, modules=None, **kwargs): self.interventions = ss.Interventions() self.analyzers = ss.Analyzers() - # Decision container + # Streams container self.streams = ss.Streams() return @@ -287,7 +287,7 @@ def init_interventions(self): raise TypeError(errormsg) for stream in intervention.streams.values(): - stream.initialize(self) + stream.initialize(self.streams, self.people._uid_map) return diff --git a/stisim/streams.py b/stisim/streams.py index 77d635de..1feffe3c 100644 --- a/stisim/streams.py +++ b/stisim/streams.py @@ -29,12 +29,12 @@ def add(self, stream): Otherwise, return value will be used as the seed offset for this stream """ if not self.initialized: - raise Exception('Please call initialize before adding a stream to Streams.') + raise NotInitializedException('Please call initialize before adding a stream to Streams.') if stream.seed_offset is None: seed = len(self._streams) # Put at end by default elif stream.seed_offset in self.used_seeds: - raise Exception(f'Requested seed offset {stream.seed_offset} for stream {stream} has already been used.') + raise SeedRepeatException(f'Requested seed offset {stream.seed_offset} for stream {stream} has already been used.') else: seed = stream.seed_offset self.used_seeds.append(seed) @@ -48,13 +48,33 @@ def step(self, ti): stream.step(ti) return + def reset(self): + for stream in self._streams.dict_values(): + stream.reset() + return + + class NotResetException(Exception): "Raised when stream is called when not ready." pass + +class NotInitializedException(Exception): + "Raised when stream is called when not initialized." + pass + + +class SeedRepeatException(Exception): + "Raised when stream is called when not initialized." + pass + + def _pre_draw(func): def check_ready(self, arr): """ Validation before drawing """ + if not self.initialized: + msg = f'Stream {self.name} has not been initialized!' + raise NotInitializedException(msg) if not self.ready: msg = f'Stream {self.name} has already been sampled on this timestep!' raise NotResetException(msg) @@ -62,6 +82,7 @@ def check_ready(self, arr): return func(self, arr) return check_ready + class Stream(np.random.Generator): """ Class for tracking one random number stream associated with one decision per timestep @@ -88,20 +109,19 @@ def __init__(self, name, seed_offset=None, **kwargs): self.kwargs = kwargs self.seed = None - self.ppl = None # Needed for block size + self.bso = None # Block size object, typically sim.people._uid_map self.initialized = False self.ready = True - return - def initialize(self, sim): + def initialize(self, streams, block_size_object): if self.initialized: # TODO: Raise warning assert not self.initialized return - self.seed = sim.streams.add(self) + self.seed = streams.add(self) if 'bit_generator' not in self.kwargs: self.kwargs['bit_generator'] = np.random.PCG64(seed=self.seed) @@ -110,7 +130,7 @@ def initialize(self, sim): #self.rng = np.random.default_rng(seed=self.seed + self.seed_offset) self._init_state = self.bit_generator.state # Store the initial state - self.ppl = sim.people + self.bso = block_size_object self.initialized = True self.ready = True @@ -136,7 +156,7 @@ def step(self, ti): @property def block_size(self): - return len(self.ppl._uid_map) + return len(self.bso) @_pre_draw def random(self, arr): diff --git a/tests/test_stream.py b/tests/test_stream.py index 5f18877b..48506373 100644 --- a/tests/test_stream.py +++ b/tests/test_stream.py @@ -8,12 +8,14 @@ import stisim as ss from stisim.streams import NotResetException -def make_rng(name='Test'): + +def make_rng(n, base_seed=1, name='Test'): """ Create and initialize a stream """ - sim = ss.Sim() - sim.initialize() + streams = ss.Streams() + streams.initialize(base_seed=base_seed) rng = ss.Stream(name) - rng.initialize(sim) + bso = np.arange(n) + rng.initialize(streams, bso) return rng @@ -23,8 +25,8 @@ def make_rng(name='Test'): def test_sample(n=5): """ Simple sample """ sc.heading('Testing stream object') - rng = make_rng() - uids = np.arange(n) + rng = make_rng(n) + uids = np.arange(0,n,2) # every other to make it interesting draws = rng.random(uids) print(f'\nSAMPLE({n}): {draws}') @@ -35,8 +37,8 @@ def test_sample(n=5): def test_reset(n=5): """ Sample, reset, sample """ sc.heading('Testing sample, reset, sample') - rng = make_rng() - uids = np.arange(n) + rng = make_rng(n) + uids = np.arange(0,n,2) # every other to make it interesting draws1 = rng.random(uids) print(f'\nSAMPLE({n}): {draws1}') @@ -53,8 +55,8 @@ def test_reset(n=5): def test_step(n=5): """ Sample, step, sample """ sc.heading('Testing sample, step, sample') - rng = make_rng() - uids = np.arange(n) + rng = make_rng(n) + uids = np.arange(0,n,2) # every other to make it interesting draws1 = rng.random(uids) print(f'\nSAMPLE({n}): {draws1}') @@ -68,11 +70,27 @@ def test_step(n=5): return np.all(draws2-draws1 != 0) +def test_seed(n=5): + """ Sample, step, sample """ + sc.heading('Testing sample with seeds 0 and 1') + uids = np.arange(0,n,2) # every other to make it interesting + + rng0 = make_rng(n, based_seed=0) + draws0 = rng0.random(uids) + print(f'\nSAMPLE({n}): {draws0}') + + rng1 = make_rng(n, based_seed=1) + draws1 = rng1.random(uids) + print(f'\nSAMPLE({n}): {draws1}') + + return np.all(draws1-draws0 != 0) + + def test_repeat(n=5): """ Sample, sample - should raise and exception""" sc.heading('Testing sample, sample - should raise an exception') - rng = make_rng() - uids = np.arange(n) + rng = make_rng(n) + uids = np.arange(0,n,2) # every other to make it interesting draws1 = rng.random(uids) print(f'\nSAMPLE({n}): {draws1}') @@ -95,6 +113,7 @@ def test_repeat(n=5): test_sample() assert test_reset() assert test_step() + assert test_seed() assert test_repeat() sc.toc(T) diff --git a/tests/test_streams.py b/tests/test_streams.py new file mode 100644 index 00000000..9a60defc --- /dev/null +++ b/tests/test_streams.py @@ -0,0 +1,169 @@ +""" +Test the Streams object from streams.py +""" + +# %% Imports and settings +import numpy as np +import sciris as sc +import stisim as ss +from stisim.streams import NotInitializedException, SeedRepeatException + + +# %% Define the tests + +def test_streams(n=5): + """ Simple sample """ + sc.heading('Testing streams object') + streams = ss.Streams() + streams.initialize(base_seed=10) + + rng = ss.Stream('stream1') + bso = np.arange(n) + rng.initialize(streams, bso) + + uids = np.arange(0,n,2) # every other to make it interesting + draws = rng.random(uids) + print(f'\nCREATED SEED AND SAMPLED: {draws}') + + return len(draws) == len(uids) + + +def test_seed(n=5): + """ Sample, reset, sample """ + sc.heading('Testing streams reset') + streams = ss.Streams() + streams.initialize(base_seed=10) + + bso = np.arange(n) + + rng0 = ss.Stream('stream0') + rng0.initialize(streams, bso) + + rng1 = ss.Stream('stream1') + rng1.initialize(streams, bso) + + return rng1.seed != rng0.seed + + +def test_reset(n=5): + """ Sample, step, sample """ + sc.heading('Testing sample, step, sample') + streams = ss.Streams() + streams.initialize(base_seed=10) + + bso = np.arange(n) + + rng = ss.Stream('stream0') + rng.initialize(streams, bso) + + uids = np.arange(0,n,2) # every other to make it interesting + s_before = rng.random(uids) + + streams.reset() # Return to step 0 + + s_after = rng.random(uids) + + return np.all(s_after == s_before) + + +def test_step(n=5): + """ Sample, step, sample """ + sc.heading('Testing sample, step, sample') + streams = ss.Streams() + streams.initialize(base_seed=10) + + bso = np.arange(n) + + rng = ss.Stream('stream0') + rng.initialize(streams, bso) + + uids = np.arange(0,n,2) # every other to make it interesting + s_before = rng.random(uids) + + streams.step(10) # 10 steps + + s_after = rng.random(uids) + + return np.all(s_after != s_before) + + +def test_initialize(n=5): + """ Sample without initializing, should raise exception """ + sc.heading('Testing without initializing, should raise exception.') + streams = ss.Streams() + #streams.initialize(base_seed=3) + + bso = np.arange(n) + + rng = ss.Stream('stream0') + + try: + rng.initialize(streams, bso) + return False # Should not get here! + except NotInitializedException as e: + print(f'YAY! Got exception: {e}') + return True + + +def test_seedrepeat(n=5): + """ Two streams with the same seed, should raise exception """ + sc.heading('Testing two streams with the same seed, should raise exception.') + streams = ss.Streams() + streams.initialize(base_seed=10) + + bso = np.arange(n) + + rng = ss.Stream('stream0', seed_offset=0) + rng.initialize(streams, bso) + + try: + rng1 = ss.Stream('stream1', seed_offset=0) + rng1.initialize(streams, bso) + return False # Should not get here! + except SeedRepeatException as e: + print(f'YAY! Got exception: {e}') + return True + +def test_samplingorder(n=5): + """ Ensure sampling from one stream doesn't affect another """ + sc.heading('Testing from multiple streams to test if sampling order matters') + streams = ss.Streams() + streams.initialize(base_seed=10) + + uids = np.arange(0,n,2) # every other to make it interesting + bso = np.arange(n) + + rng0 = ss.Stream('stream0') + rng0.initialize(streams, bso) + + rng1 = ss.Stream('stream1') + rng1.initialize(streams, bso) + + s_before = rng0.random(uids) + _ = rng1.random(uids) + + streams.reset() + + _ = rng1.random(uids) + s_after = rng0.random(uids) + + return np.all(s_before == s_after) + + + +# %% Run as a script +if __name__ == '__main__': + # Start timing + T = sc.tic() + + # Run tests + assert test_streams() + assert test_seed() + assert test_reset() + assert test_step() + assert test_initialize() + assert test_seedrepeat() + assert test_samplingorder() + + sc.toc(T) + print('Done.') \ No newline at end of file From 173262cf3be4857b9e3818db6235c08f91bca677 Mon Sep 17 00:00:00 2001 From: Daniel Klein Date: Mon, 25 Sep 2023 12:49:02 -0400 Subject: [PATCH 16/34] *Initialization fix in networks.py * module streams by layer, about to fix. * Checking for repeat stream names. --- stisim/modules.py | 43 +++--- stisim/networks.py | 6 +- stisim/streams.py | 8 + tests/simple.py | 144 ------------------ ..._stable_monogamy.py => stable_monogamy.py} | 0 tests/test_hiv.py | 46 ++++++ tests/test_streams.py | 39 ++++- 7 files changed, 121 insertions(+), 165 deletions(-) delete mode 100644 tests/simple.py rename tests/{test_stable_monogamy.py => stable_monogamy.py} (100%) create mode 100644 tests/test_hiv.py diff --git a/stisim/modules.py b/stisim/modules.py index d8972fad..94b91c36 100644 --- a/stisim/modules.py +++ b/stisim/modules.py @@ -21,10 +21,12 @@ def __init__(self, pars=None, label=None, requires=None, *args, **kwargs): self.finalized = False # Random number streams - self.rng_init_cases = ss.Stream('initial_cases') - self.rng_trans_ab = ss.Stream('trans_ab') - self.rng_trans_ba = ss.Stream('trans_ba') - self.rng_choose_infector = ss.Stream('choose_infector') + self.rng_init_cases = ss.Stream('initial_cases') + # The following random streams are dicts from layer key to rng + self.rng_trans_ab = {} + self.rng_trans_ba = {} + self.rng_choose_infector_ab = {} + self.rng_choose_infector_ba = {} return @@ -44,6 +46,13 @@ def check_requires(self, sim): def initialize(self, sim): self.check_requires(sim) + # Random number streams per network layer + for lkey in sim.people.networks.keys(): + self.rng_trans_ab[lkey] = ss.Stream(f'trans_ab: {lkey}') + self.rng_trans_ba[lkey] = ss.Stream(f'trans_ba: {lkey}') + self.rng_choose_infector_ab[lkey] = ss.Stream(f'choose_infector_ab: {lkey}') + self.rng_choose_infector_ba[lkey] = ss.Stream(f'choose_infector_ba: {lkey}') + # Connect the states to the people for state in self.states.values(): state.initialize(sim.people) @@ -154,33 +163,33 @@ def make_new_cases(self, sim): """ Add new cases of module, through transmission, incidence, etc. """ n_new_cases = 0 # number of new cases made pars = sim.pars[self.name] - for k, layer in sim.people.networks.items(): - if k in pars['beta']: + for lkey, layer in sim.people.networks.items(): + if lkey in pars['beta']: rel_trans = self.rel_trans * (self.infected & sim.people.alive) rel_sus = self.rel_sus * (self.susceptible & sim.people.alive) - for a, b, beta, rng in [[layer['p1'], layer['p2'], pars['beta'][k][0], self.rng_trans_ab], [layer['p2'], layer['p1'], pars['beta'][k][1], self.rng_trans_ba]]: - # probability of a->b transmission + + a_to_b = [layer['p1'], layer['p2'], pars['beta'][lkey][0], self.rng_trans_ab[lkey], self.rng_choose_infector_ab[lkey]] + b_to_a = [layer['p2'], layer['p1'], pars['beta'][lkey][1], self.rng_trans_ba[lkey], self.rng_choose_infector_ba[lkey]] + for a, b, beta, rng_trans, rng_chs_inf in [a_to_b, b_to_a]: if beta == 0: continue - #p_tran_edge = rel_trans[a] * rel_sus[b] * layer['beta'] * beta - # Will need to be more efficient here - can maintain edge to node matrix + # Check for new transmission from a --> b + # TODO: Will need to be more efficient here - can maintain edge to node matrix node_from_edge = np.ones( (len(sim.people._uid_map), len(a)) ) node_from_edge[b, np.arange(len(b))] = 1 - rel_trans[a] * rel_sus[b] * layer['beta'] * beta - #p_acq_node = np.dot(node_from_edge, p_tran_edge) # calculated for all nodes, including those outside b p_acq_node = 1 - node_from_edge.prod(axis=1) # 1 - (1-p1)*(1-p2)*... - new_cases_bool = rng.bernoulli(p_acq_node[b], b) + new_cases_bool = rng_trans.bernoulli(p_acq_node[b], b) new_cases = b[new_cases_bool] if not len(new_cases): continue n_new_cases += len(new_cases) - # Decide whom the infection came from - frm = np.zeros_like(new_cases) - # Need one random number for each b <-- align by blocksize - r = self.rng_choose_infector.random(new_cases) # Align rng to b + # Decide whom the infection came from using one random number for each b (aligned by block size) + frm = np.zeros_like(new_cases) + r = rng_chs_inf.random(new_cases) prob = (1-node_from_edge[new_cases]) cumsum = (prob / ((prob.sum(axis=1)[:,np.newaxis]))).cumsum(axis=1) frm_idx = np.argmax( cumsum >= r[:,np.newaxis], axis=1) @@ -199,4 +208,4 @@ def update_results(self, sim): self.results['new_infections'][sim.ti] = np.count_nonzero(self.ti_infected == sim.ti) def finalize_results(self, sim): - pass + pass \ No newline at end of file diff --git a/stisim/networks.py b/stisim/networks.py index ba65c8ab..afe5a3ce 100644 --- a/stisim/networks.py +++ b/stisim/networks.py @@ -294,9 +294,9 @@ def initialize(self, sim): super().initialize(sim) # Initialize random streams - self.rng_pair_12.initialize(sim) - self.rng_pair_21.initialize(sim) - self.rng_mean_dur.initialize(sim) + self.rng_pair_12.initialize(sim.streams, sim.people._uid_map) + self.rng_pair_21.initialize(sim.streams, sim.people._uid_map) + self.rng_mean_dur.initialize(sim.streams, sim.people._uid_map) self.add_pairs(sim.people, ti=0) return diff --git a/stisim/streams.py b/stisim/streams.py index 1feffe3c..5d24bc67 100644 --- a/stisim/streams.py +++ b/stisim/streams.py @@ -31,6 +31,9 @@ def add(self, stream): if not self.initialized: raise NotInitializedException('Please call initialize before adding a stream to Streams.') + if stream.name in self._streams: + raise RepeatNameException(f'A Stream with name {stream.name} has already been added.') + if stream.seed_offset is None: seed = len(self._streams) # Put at end by default elif stream.seed_offset in self.used_seeds: @@ -64,6 +67,11 @@ class NotInitializedException(Exception): pass +class RepeatNameException(Exception): + "Raised when adding a stream to streams when the stream name has already been used." + pass + + class SeedRepeatException(Exception): "Raised when stream is called when not initialized." pass diff --git a/tests/simple.py b/tests/simple.py deleted file mode 100644 index d242f3cf..00000000 --- a/tests/simple.py +++ /dev/null @@ -1,144 +0,0 @@ -""" -Run simplest tests -""" - -# %% Imports and settings -import stisim as ss -import matplotlib.pyplot as plt -import pandas as pd -import numpy as np -import networkx as nx - -plot_graph = True - -class Graph(): - def __init__(self, nodes, edges): - self.graph = nx.from_pandas_edgelist(df=edges, source='p1', target='p2', edge_attr=True) - self.graph.add_nodes_from(nodes.index) - nx.set_node_attributes(self.graph, nodes.transpose().to_dict()) - return - - def draw_nodes(self, filter, ax, **kwargs): - inds = [i for i,n in self.graph.nodes.data() if filter(n)] - nc = ['red' if nd['hiv'] else 'lightgray' for i, nd in self.graph.nodes.data() if i in inds] - ec = ['green' if nd['on_art'] or nd['on_prep'] else 'black' for i, nd in self.graph.nodes.data() if i in inds] - if inds: - nx.draw_networkx_nodes(self.graph, nodelist=inds, pos=pos, ax=ax, node_color=nc, edgecolors=ec, **kwargs) - - def plot(self, pos, ax=None): - kwargs = dict(node_shape='x', node_size=250, linewidths=2, ax=ax) - self.draw_nodes(lambda n: n['dead'], **kwargs) - - kwargs['node_shape'] = 'o' - self.draw_nodes(lambda n: not n['dead'] and n['female'], **kwargs) - - kwargs['node_shape'] = 's' - self.draw_nodes(lambda n: not n['dead'] and not n['female'], **kwargs) - - nx.draw_networkx_edges(self.graph, pos=pos, ax=ax) - nx.draw_networkx_labels(self.graph, labels={i:int(a['cd4']) for i,a in self.graph.nodes.data()}, font_size=8, pos=pos, ax=ax) - nx.draw_networkx_edge_labels(self.graph, edge_labels={(i,j): int(a['dur']) for i,j,a in self.graph.edges.data()}, font_size=8, pos=pos, ax=ax) - - return - -class rng_analyzer(ss.Analyzer): - ''' Simple analyzer to assess if random streams are working ''' - - def __init__(self, **kwargs): - super().__init__(**kwargs) # Initialize the Analyzer object - - self.graphs = {} - return - - def initialize(self, sim): - self.initialized = True - return - - def apply(self, sim): - nodes = pd.DataFrame({ - 'female': sim.people.female.values, - 'dead': sim.people.dead.values, - 'hiv': sim.people.hiv.infected.values, - 'on_art': sim.people.hiv.on_art.values, - 'cd4': sim.people.hiv.cd4.values, - 'on_prep': sim.people.hiv.on_prep.values, - #'ng': sim.people.gonorrhea.infected.values, - }) - - edges = pd.DataFrame(sim.people.networks['simple_embedding'].to_dict()) #sim.people.networks['simple_embedding'].to_df() #TODO: repr issues - - self.graphs[sim.ti] = Graph(nodes, edges) - return - - def finalize(self, sim): - super().finalize(sim) - return - - -def run_sim(n=25, intervention=False, analyze=False): - ppl = ss.People(n) - ppl.networks = ss.ndict(ss.simple_embedding())#, ss.maternal()) - - hiv = ss.HIV() - #hiv.pars['beta'] = {'simple_embedding': [0.0008, 0.0004], 'maternal': [0.2, 0]} - hiv.pars['beta'] = {'simple_embedding': [0.06, 0.04]}#, 'maternal': [0.2, 0]} - hiv.pars['initial'] = 3 - - #gon = ss.Gonorrhea() - #gon.pars['beta'] = 0.3 - - pars = { - 'start': 1980, - 'end': 2010, - 'interventions': [ss.hiv.ART(0, 0.1)] if intervention else [], # ss.hiv.PrEP(0, 0.2), - 'rand_seed': 0, - 'analyzers': [rng_analyzer()] if analyze else [], - } - #sim = ss.Sim(people=ppl, modules=[hiv, gon, ss.Pregnancy()], pars=pars, label=f'Sim with {n} agents and intv={intervention}') - sim = ss.Sim(people=ppl, modules=[hiv], pars=pars, label=f'Sim with {n} agents and intv={intervention}') - sim.initialize() - sim.run() - - return sim - -n = 10 -sim1 = run_sim(n, intervention=False, analyze=plot_graph) -sim2 = run_sim(n, intervention=True, analyze=plot_graph) -#sc.save('sims.obj', [sim1, sim2]) - -# TODO: Parallelization does not work with the current snapshot analyzer -#sim1, sim2 = sc.parallelize(run_sim, iterkwargs=[{'intervention':False}, {'intervention':True}], die=True) - -if plot_graph: - g1 = sim1.analyzers[0].graphs - g2 = sim2.analyzers[0].graphs - - #pos = nx.circular_layout(g1[0].graph) - #n = len(sim1.people._uid_map) # Will eventually need this line due to population growth - pos = {i:(np.cos(2*np.pi*i/n), np.sin(2*np.pi*i/n)) for i in range(n)} - - for ti in sim1.tivec: - fig, axv = plt.subplots(1,2, figsize=(10,5)) - g1[ti].plot(pos, ax=axv[0]) - g2[ti].plot(pos, ax=axv[1]) - fig.suptitle(f'Time is {ti}') - plt.show() - -# Plot -fig, axv = plt.subplots(2,1, sharex=True) -axv[0].plot(sim1.tivec, sim1.results.hiv.n_infected, label='Baseline') -axv[0].plot(sim2.tivec, sim2.results.hiv.n_infected, ls=':', label='Intervention') -axv[0].set_title('HIV number of infections') - -axv[1].plot(sim1.tivec, sim1.results.hiv.new_deaths, label='Baseline') -axv[1].plot(sim2.tivec, sim2.results.hiv.new_deaths, ls=':', label='Intervention') -axv[1].set_title('HIV Deaths') - -''' Gonorrhea removed for now -axv[2].plot(sim1.tivec, sim1.results.gonorrhea.n_infected, label='Baseline') -axv[2].plot(sim2.tivec, sim2.results.gonorrhea.n_infected, ls=':', label='Intervention') -axv[2].set_title('Gonorrhea number of infections') -''' -plt.legend() -plt.show() -print('Done') \ No newline at end of file diff --git a/tests/test_stable_monogamy.py b/tests/stable_monogamy.py similarity index 100% rename from tests/test_stable_monogamy.py rename to tests/stable_monogamy.py diff --git a/tests/test_hiv.py b/tests/test_hiv.py new file mode 100644 index 00000000..a41afeb7 --- /dev/null +++ b/tests/test_hiv.py @@ -0,0 +1,46 @@ +""" +Run a simple HIV simulation with random number coherence +""" + +# %% Imports and settings +import stisim as ss +import matplotlib.pyplot as plt +import sciris as sc + +n = 1_000 # Agents + +def run_sim(n=25, intervention=False, analyze=False): + ppl = ss.People(n) + ppl.networks = ss.ndict(ss.simple_embedding())#, ss.maternal()) + + hiv = ss.HIV() + hiv.pars['beta'] = {'simple_embedding': [0.10, 0.08]} + hiv.pars['initial'] = 10 + + pars = { + 'start': 1980, + 'end': 2010, + 'interventions': [ss.hiv.ART(t=[0, 1], coverage=[0, 0.9**3])] if intervention else [], + 'rand_seed': 0, + } + sim = ss.Sim(people=ppl, modules=[hiv], pars=pars, label=f'Sim with {n} agents and intv={intervention}') + sim.initialize() + sim.run() + + return sim + +sim1, sim2 = sc.parallelize(run_sim, kwargs={'n': n}, iterkwargs=[{'intervention':False}, {'intervention':True}], die=True) + +# Plot +fig, axv = plt.subplots(2,1, sharex=True) +axv[0].plot(sim1.tivec, sim1.results.hiv.n_infected, label='Baseline') +axv[0].plot(sim2.tivec, sim2.results.hiv.n_infected, ls=':', label='Intervention') +axv[0].set_title('HIV number of infections') + +axv[1].plot(sim1.tivec, sim1.results.hiv.new_deaths, label='Baseline') +axv[1].plot(sim2.tivec, sim2.results.hiv.new_deaths, ls=':', label='Intervention') +axv[1].set_title('HIV Deaths') + +plt.legend() +plt.show() +print('Done') \ No newline at end of file diff --git a/tests/test_streams.py b/tests/test_streams.py index 9a60defc..265e9dc0 100644 --- a/tests/test_streams.py +++ b/tests/test_streams.py @@ -6,7 +6,7 @@ import numpy as np import sciris as sc import stisim as ss -from stisim.streams import NotInitializedException, SeedRepeatException +from stisim.streams import NotInitializedException, SeedRepeatException, RepeatNameException # %% Define the tests @@ -139,6 +139,42 @@ def test_samplingorder(n=5): rng1 = ss.Stream('stream1') rng1.initialize(streams, bso) + s_before = rng0.random(uids) + _ = rng1.random(uids) + + streams.reset() + + _ = rng1.random(uids) + s_after = rng0.random(uids) + + return np.all(s_before == s_after) + + +def test_repeatname(n=5): + """ Test two streams with the same name """ + sc.heading('Testing if two streams with the same name are allowed') + streams = ss.Streams() + streams.initialize(base_seed=17) + + uids = np.arange(0,n,2) # every other to make it interesting + bso = np.arange(n) + + rng0 = ss.Stream('test') + rng0.initialize(streams, bso) + + rng1 = ss.Stream('test') + try: + rng1.initialize(streams, bso) + return False # Should not get here! + except RepeatNameException as e: + print(f'YAY! Got exception: {e}') + return True + + + + + + s_before = rng0.random(uids) _ = rng1.random(uids) @@ -164,6 +200,7 @@ def test_samplingorder(n=5): assert test_initialize() assert test_seedrepeat() assert test_samplingorder() + assert test_repeatname() sc.toc(T) print('Done.') \ No newline at end of file From 8df0662f0e2af121397b4d431d2380f58ab10967 Mon Sep 17 00:00:00 2001 From: Daniel Klein Date: Wed, 27 Sep 2023 18:08:30 -0700 Subject: [PATCH 17/34] LOTS OF GREAT IMPROVEMENTS: * Adding run_multistream.py to compare a simple HIV simulation with and without random stream coherence. * Stream block size now determined automatically * Stream seed offsets are now determined by hashing the stream name instead of sequentially. Thanks to @RomeshA for the suggestion. * New transmission calcs. Now determining acquisition probability across all layers and edge directions before using one random draw to determine acquisition. * This change makes it more challenging to determine whom the infection came from, that is currently a N^2 operation that needs optimization. * Adding ability to disable multiple streams through parameter named "multistream" * The ability switch from MultiStream (the default) to CentralizedStream adds some complexity. Consider removing this feature later. * As a result, everything now has parameters. * Fixed prevalence calculation to include only people who are alive in the denominator. * Additional error checking for streams, for example when the index array is boolean or empty. * The _pre_draw decorator is now applied to more MultiStream calls for error checking. * New CentralizedStream class that behaves like MultiStream but calls into a centralized random number generator. * Fixed a small bug in states.py regarding the size of grow(). * Improved stream and streams tests. --- stisim/demographics.py | 8 +- stisim/gonorrhea.py | 11 +- stisim/hiv.py | 24 ++-- stisim/modules.py | 90 ++++++++------- stisim/networks.py | 64 +++++------ stisim/parameters.py | 1 + stisim/sim.py | 2 +- stisim/states.py | 5 +- stisim/streams.py | 242 +++++++++++++++++++++++++++++++++++---- tests/run_multistream.py | 140 ++++++++++++++++++++++ tests/stable_monogamy.py | 24 ++-- tests/test_hiv.py | 46 -------- tests/test_stream.py | 112 +++++++++++++----- tests/test_streams.py | 81 +++++-------- 14 files changed, 596 insertions(+), 254 deletions(-) create mode 100644 tests/run_multistream.py delete mode 100644 tests/test_hiv.py diff --git a/stisim/demographics.py b/stisim/demographics.py index a2ec2670..31065448 100644 --- a/stisim/demographics.py +++ b/stisim/demographics.py @@ -32,8 +32,8 @@ def __init__(self, pars=None): 'initial': 3, # Number of women initially pregnant }, self.pars) - self.rng_sex = ss.Stream('sex_at_birth') - self.rng_conception = ss.Stream('conception') + self.rng_sex = ss.Stream(self.multistream)('sex_at_birth') + self.rng_conception = ss.Stream(self.multistream)('conception') return @@ -101,7 +101,7 @@ 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 = self.rng_conception.bernoulli_filter(prob=self.pars.inci, arr=inds_to_choose_from) + uids = self.rng_conception.bernoulli_filter(arr=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) @@ -109,7 +109,7 @@ def make_pregnancies(self, sim): # Grow the arrays and set properties for the unborn agents new_uids = sim.people.grow(n_unborn_agents) sim.people.age[new_uids] = -self.pars.dur_pregnancy - sim.people.female[new_uids] = self.rng_sex.bernoulli(prob=0.5, arr=uids) # Replace 0.5 with sex ratio at birth + sim.people.female[new_uids] = self.rng_sex.bernoulli(arr=uids, prob=0.5) # Replace 0.5 with sex ratio at birth # Add connections to any vertical transmission layers # Placeholder code to be moved / refactored. The maternal network may need to be diff --git a/stisim/gonorrhea.py b/stisim/gonorrhea.py index fc17d4fc..158a67b3 100644 --- a/stisim/gonorrhea.py +++ b/stisim/gonorrhea.py @@ -20,9 +20,9 @@ def __init__(self, pars=None): self.ti_recovered = ss.State('ti_recovered', float, 0) self.ti_dead = ss.State('ti_dead', float, np.nan) # Death due to gonorrhea - self.rng_prog = ss.Stream('prog_dur') - self.rng_dead = ss.Stream('dead') - self.rng_dur_inf = ss.Stream('dur_inf') + self.rng_prog = ss.Stream(self.multistream)('prog_dur') + self.rng_dead = ss.Stream(self.multistream)('dead') + self.rng_dur_inf = ss.Stream(self.multistream)('dur_inf') self.pars = ss.omerge({ 'dur_inf': 3, # not modelling diagnosis or treatment explicitly here @@ -58,9 +58,8 @@ def set_prognoses(self, sim, to_uids, from_uids=None): self.infected[to_uids] = True self.ti_infected[to_uids] = sim.ti - #dur = sim.ti + self.rng_dur_inf.poisson(self.pars['dur_inf']/sim.pars.dt, len(to_uids)) # By whom infected from??? TODO - dur = sim.ti + self.rng_dur_inf.poisson(self.pars['dur_inf']/sim.pars.dt, to_uids) # By whom infected from??? TODO - dead = self.rng_dead.bernoulli(self.pars.p_death, to_uids) + dur = sim.ti + self.rng_dur_inf.poisson(to_uids, self.pars['dur_inf']/sim.pars.dt) # By whom infected from??? TODO + dead = self.rng_dead.bernoulli(to_uids, self.pars.p_death) self.ti_recovered[to_uids[~dead]] = dur[~dead] self.ti_dead[to_uids[dead]] = dur[dead] diff --git a/stisim/hiv.py b/stisim/hiv.py index a9a4e2cf..dfdbe332 100644 --- a/stisim/hiv.py +++ b/stisim/hiv.py @@ -21,7 +21,7 @@ def __init__(self, pars=None): self.on_prep = ss.State('on_prep', bool, False) self.cd4 = ss.State('cd4', float, 500) - self.rng_dead = ss.Stream('dead') + self.rng_dead = ss.Stream(self.multistream)('dead') self.pars = ss.omerge({ 'cd4_min': 100, @@ -43,7 +43,7 @@ def update_states(self, sim): hiv_death_prob = 0.1 / (self.pars.cd4_min - self.pars.cd4_max)**2 * (self.cd4 - self.pars.cd4_max)**2 can_die = ss.true(sim.people.alive & sim.people.hiv.infected) - hiv_deaths = self.rng_dead.bernoulli_filter(prob=hiv_death_prob[can_die], arr = can_die) + hiv_deaths = self.rng_dead.bernoulli_filter(arr=can_die, prob=hiv_death_prob[can_die]) sim.people.alive[hiv_deaths] = False sim.people.ti_dead[hiv_deaths] = sim.ti self.results['new_deaths'][sim.ti] = len(hiv_deaths) @@ -73,13 +73,15 @@ def set_prognoses(self, sim, to_uids, from_uids=None): class ART(ss.Intervention): - def __init__(self, t: np.array, coverage: np.array): + def __init__(self, t: np.array, coverage: np.array, **kwargs): self.requires = HIV self.t = sc.promotetoarray(t) self.coverage = sc.promotetoarray(coverage) - self.rng_add_ART = ss.Stream('add_ART', seed_offset=100) - self.rng_remove_ART = ss.Stream('remove_ART', seed_offset=101) + super().__init__(**kwargs) + + self.rng_add_ART = ss.Stream(self.multistream)('add_ART', seed_offset=100) + self.rng_remove_ART = ss.Stream(self.multistream)('remove_ART', seed_offset=101) return @@ -101,13 +103,13 @@ def apply(self, sim): eligible = ss.true(sim.people.alive & sim.people.hiv.infected & ~sim.people.hiv.on_art) n_eligible = len(eligible) #np.count_nonzero(eligible) if n_eligible: - inds = self.rng_add_ART.bernoulli_filter(prob=min(n_eligible, n_change)/n_eligible, arr=eligible) + inds = self.rng_add_ART.bernoulli_filter(arr=eligible, prob=min(n_eligible, n_change)/n_eligible) sim.people.hiv.on_art[inds] = True elif n_change < 0: # Take some people off ART eligible = sim.people.alive & sim.people.hiv.infected & sim.people.hiv.on_art n_eligible = np.count_nonzero(eligible) - inds = self.rng_remove_ART.bernoulli_filter(prob=-n_change/n_eligible, arr=eligible) + inds = self.rng_remove_ART.bernoulli_filter(arr=eligible, prob=-n_change/n_eligible) sim.people.hiv.on_art[inds] = False # Add result @@ -122,8 +124,8 @@ def __init__(self, t: np.array, coverage: np.array): self.t = sc.promotetoarray(t) self.coverage = sc.promotetoarray(coverage) - self.rng_add_PrEP = ss.Stream('add_PrEP', seed_offset=102) - self.rng_remove_PrEP = ss.Stream('remove_PrEP', seed_offset=103) + self.rng_add_PrEP = ss.Stream(self.multistream)('add_PrEP', seed_offset=102) + self.rng_remove_PrEP = ss.Stream(self.multistream)('remove_PrEP', seed_offset=103) return @@ -144,13 +146,13 @@ def apply(self, sim): eligible = ss.true(sim.people.alive & ~sim.people.hiv.infected & ~sim.people.hiv.on_prep) n_eligible = len(eligible) #np.count_nonzero(eligible) if n_eligible: - inds = self.rng_add_PrEP.bernoulli_filter(prob=min(n_eligible, n_change)/n_eligible, arr=eligible) + inds = self.rng_add_PrEP.bernoulli_filter(arr=eligible, prob=min(n_eligible, n_change)/n_eligible) sim.people.hiv.on_prep[inds] = True elif n_change < 0: # Take some people off PrEP eligible = sim.people.alive & sim.people.hiv.on_prep n_eligible = np.count_nonzero(eligible) - inds = self.rng_remove_PrEP.bernoulli_filter(prob=-n_change/n_eligible, arr=eligible) + inds = self.rng_remove_PrEP.bernoulli_filter(arr=eligible, prob=-n_change/n_eligible) sim.people.hiv.on_prep[inds] = False # Add result diff --git a/stisim/modules.py b/stisim/modules.py index 94b91c36..27827e0b 100644 --- a/stisim/modules.py +++ b/stisim/modules.py @@ -13,20 +13,22 @@ class Module(sc.prettyobj): def __init__(self, pars=None, label=None, requires=None, *args, **kwargs): - self.pars = ss.omerge(pars) + + default_pars = {'multistream': True} + self.pars = sc.mergedicts(default_pars, pars) + self.label = label if label else '' self.requires = sc.mergelists(requires) self.results = ss.Results() self.initialized = False self.finalized = False + self.multistream = self.pars['multistream'] + # Random number streams - self.rng_init_cases = ss.Stream('initial_cases') - # The following random streams are dicts from layer key to rng - self.rng_trans_ab = {} - self.rng_trans_ba = {} - self.rng_choose_infector_ab = {} - self.rng_choose_infector_ba = {} + self.rng_init_cases = ss.Stream(self.multistream)(f'initial_cases_{self.name}') + self.rng_trans = ss.Stream(self.multistream)(f'trans_{self.name}') + self.rng_choose_infector = ss.Stream(self.multistream)(f'choose_infector_{self.name}') return @@ -46,20 +48,14 @@ def check_requires(self, sim): def initialize(self, sim): self.check_requires(sim) - # Random number streams per network layer - for lkey in sim.people.networks.keys(): - self.rng_trans_ab[lkey] = ss.Stream(f'trans_ab: {lkey}') - self.rng_trans_ba[lkey] = ss.Stream(f'trans_ba: {lkey}') - self.rng_choose_infector_ab[lkey] = ss.Stream(f'choose_infector_ab: {lkey}') - self.rng_choose_infector_ba[lkey] = ss.Stream(f'choose_infector_ba: {lkey}') - # Connect the states to the people for state in self.states.values(): state.initialize(sim.people) # Connect the streams to the sim for stream in self.streams.values(): - stream.initialize(sim.streams, sim.people._uid_map) + if not stream.initialized: + stream.initialize(sim.streams) self.initialized = True return @@ -81,7 +77,7 @@ def states(self): @property def streams(self): - return ss.ndict({k:v for k,v in self.__dict__.items() if isinstance(v, ss.Stream)}) + return ss.ndict({k:v for k,v in self.__dict__.items() if isinstance(v, (ss.MultiStream, ss.CentralizedStream))}) class Modules(ss.ndict): @@ -130,7 +126,7 @@ def set_initial_states(self, sim): #initial_cases = np.random.choice(sim.people.uid, self.pars['initial']) #rng = sim.rngs.get(f'initial_cases_{self.name}') #initial_cases = ss.binomial_filter_rng(prob=self.pars['initial']/len(sim.people), arr=sim.people.uid, rng=rng, block_size=len(sim.people._uid_map)) - initial_cases = self.rng_init_cases.bernoulli_filter(prob=self.pars['initial']/len(sim.people), arr=sim.people.uid) + initial_cases = self.rng_init_cases.bernoulli_filter(arr=sim.people.uid, prob=self.pars['initial']/len(sim.people)) self.set_prognoses(sim, initial_cases, from_uids=None) # TODO: sentinel value to indicate seeds? return @@ -161,42 +157,52 @@ def update_states(self, sim): def make_new_cases(self, sim): """ Add new cases of module, through transmission, incidence, etc. """ - n_new_cases = 0 # number of new cases made pars = sim.pars[self.name] + + # Probability of each node acquiring a case + # TODO: Just people that who are alive? + p_acq_node = np.zeros( len(sim.people._uid_map) ) + node_from_node = np.ones( (sim.people._uid_map.n, sim.people._uid_map.n) ) + for lkey, layer in sim.people.networks.items(): if lkey in pars['beta']: rel_trans = self.rel_trans * (self.infected & sim.people.alive) rel_sus = self.rel_sus * (self.susceptible & sim.people.alive) - a_to_b = [layer['p1'], layer['p2'], pars['beta'][lkey][0], self.rng_trans_ab[lkey], self.rng_choose_infector_ab[lkey]] - b_to_a = [layer['p2'], layer['p1'], pars['beta'][lkey][1], self.rng_trans_ba[lkey], self.rng_choose_infector_ba[lkey]] - for a, b, beta, rng_trans, rng_chs_inf in [a_to_b, b_to_a]: + a_to_b = [layer['p1'], layer['p2'], pars['beta'][lkey][0]] + b_to_a = [layer['p2'], layer['p1'], pars['beta'][lkey][1]] + for a, b, beta in [a_to_b, b_to_a]: if beta == 0: continue # Check for new transmission from a --> b # TODO: Will need to be more efficient here - can maintain edge to node matrix - node_from_edge = np.ones( (len(sim.people._uid_map), len(a)) ) - node_from_edge[b, np.arange(len(b))] = 1 - rel_trans[a] * rel_sus[b] * layer['beta'] * beta - p_acq_node = 1 - node_from_edge.prod(axis=1) # 1 - (1-p1)*(1-p2)*... - new_cases_bool = rng_trans.bernoulli(p_acq_node[b], b) - new_cases = b[new_cases_bool] - - if not len(new_cases): - continue - - n_new_cases += len(new_cases) - - # Decide whom the infection came from using one random number for each b (aligned by block size) - frm = np.zeros_like(new_cases) - r = rng_chs_inf.random(new_cases) - prob = (1-node_from_edge[new_cases]) - cumsum = (prob / ((prob.sum(axis=1)[:,np.newaxis]))).cumsum(axis=1) - frm_idx = np.argmax( cumsum >= r[:,np.newaxis], axis=1) - frm = a[frm_idx] - self.set_prognoses(sim, new_cases, frm) + node_from_edge = np.ones( (sim.people._uid_map.n, len(a)) ) + node_from_edge[b, np.arange(len(a))] = 1 - rel_trans[a] * rel_sus[b] * layer['beta'] * beta + p_not_acq_by_node_this_layer_b_from_a = node_from_edge.prod(axis=1) # (1-p1)*(1-p2)*... + p_acq_node = 1 - (1-p_acq_node) * p_not_acq_by_node_this_layer_b_from_a + + node_from_node_this_layer_b_from_a = np.ones( (sim.people._uid_map.n, sim.people._uid_map.n) ) + node_from_node_this_layer_b_from_a[b, a] = 1 - rel_trans[a] * rel_sus[b] * layer['beta'] * beta + node_from_node *= node_from_node_this_layer_b_from_a + + new_cases_bool = self.rng_trans.bernoulli(sim.people._uid_map._view, p_acq_node) + new_cases = sim.people._uid_map._view[new_cases_bool] + + if not len(new_cases): + return 0 + + # Decide whom the infection came from using one random number for each b (aligned by block size) + frm = np.zeros_like(new_cases) + r = self.rng_choose_infector.random(new_cases) + prob = (1-node_from_node[new_cases]) + cumsum = (prob / ((prob.sum(axis=1)[:,np.newaxis]))).cumsum(axis=1) + frm = np.argmax( cumsum >= r[:,np.newaxis], axis=1) + #frm = a[frm_idx] + self.set_prognoses(sim, new_cases, frm) + + return len(new_cases) # number of new cases made - return n_new_cases def set_prognoses(self, sim, to_uids, from_uids): pass @@ -204,7 +210,7 @@ def set_prognoses(self, sim, to_uids, from_uids): def update_results(self, sim): self.results['n_susceptible'][sim.ti] = np.count_nonzero(self.susceptible) self.results['n_infected'][sim.ti] = np.count_nonzero(self.infected) - self.results['prevalence'][sim.ti] = self.results.n_infected[sim.ti] / len(sim.people) + self.results['prevalence'][sim.ti] = self.results.n_infected[sim.ti] / sim.people.alive.sum() self.results['new_infections'][sim.ti] = np.count_nonzero(self.ti_infected == sim.ti) def finalize_results(self, sim): diff --git a/stisim/networks.py b/stisim/networks.py index afe5a3ce..e755d3e2 100644 --- a/stisim/networks.py +++ b/stisim/networks.py @@ -51,18 +51,24 @@ class Network(sc.objdict): network2 = ss.Network(**network, index=index, self_conn=self_conn, label=network.label) """ - def __init__(self, *args, key_dict=None, vertical=False, label=None, **kwargs): + def __init__(self, *args, pars=None, key_dict=None, vertical=False, label=None, **kwargs): default_keys = { 'p1': ss.int_, 'p2': ss.int_, 'beta': ss.float_, } + + default_pars = {'multistream': True} + self.pars = sc.mergedicts(default_pars, pars) + self.meta = sc.mergedicts(default_keys, key_dict) self.vertical = vertical # Whether transmission is bidirectional self.basekey = 'p1' # Assign a base key for calculating lengths and performing other operations self.label = label self.initialized = False + self.multistream = self.pars['multistream'] + # Handle args kwargs = sc.mergedicts(*args, kwargs) @@ -86,7 +92,7 @@ def initialize(self, sim): ''' # Connect the streams to the sim for stream in self.streams.values(): - stream.initialize(sim) + stream.initialize(sim.streams) return ''' @@ -269,7 +275,7 @@ class simple_sexual(Network): A class holding a single network of contact edges (connections) between people. This network is built by **randomly pairing** males and female with variable relationship durations. """ - def __init__(self, mean_dur=5): + def __init__(self, mean_dur=5, **kwargs): key_dict = { 'p1': ss.int_, 'p2': ss.int_, @@ -278,15 +284,15 @@ def __init__(self, mean_dur=5): } # Call init for the base class, which sets all the keys - super().__init__(key_dict=key_dict) + super().__init__(key_dict=key_dict, **kwargs) # Set other parameters self.mean_dur = mean_dur # Define random streams - self.rng_pair_12 = ss.Stream('pair_12') - self.rng_pair_21 = ss.Stream('pair_21') - self.rng_mean_dur = ss.Stream('mean_dur') + self.rng_pair_12 = ss.Stream(self.multistream)('pair_12') + self.rng_pair_21 = ss.Stream(self.multistream)('pair_21') + self.rng_mean_dur = ss.Stream(self.multistream)('mean_dur') return @@ -294,9 +300,9 @@ def initialize(self, sim): super().initialize(sim) # Initialize random streams - self.rng_pair_12.initialize(sim.streams, sim.people._uid_map) - self.rng_pair_21.initialize(sim.streams, sim.people._uid_map) - self.rng_mean_dur.initialize(sim.streams, sim.people._uid_map) + self.rng_pair_12.initialize(sim.streams) + self.rng_pair_21.initialize(sim.streams) + self.rng_mean_dur.initialize(sim.streams) self.add_pairs(sim.people, ti=0) return @@ -310,13 +316,13 @@ def add_pairs(self, people, ti=None): if len(available_m) <= len(available_f): p1 = available_m - p2 = self.rng_pair_12.choice(available_f, len(p1), replace=False) # TODO: Stream-ify + p2 = self.rng_pair_12.choice(len(p1), available_f, replace=False) # TODO: Stream-ify else: p2 = available_f - p1 = self.rng_pair_21.choice(available_m, len(p2), replace=False) # TODO: Stream-ify + p1 = self.rng_pair_21.choice(len(p2), available_m, replace=False) # TODO: Stream-ify beta = np.ones_like(p1) - dur = self.rng_mean_dur.poisson(self.mean_dur, len(p1)) # TODO: Stream-ify + dur = self.rng_mean_dur.poisson(len(p1), self.mean_dur) # TODO: Stream-ify self['p1'] = np.concatenate([self['p1'], p1]) self['p2'] = np.concatenate([self['p2'], p2]) self['beta'] = np.concatenate([self['beta'], beta]) @@ -408,7 +414,7 @@ def add_pairs(self, people, ti=None): self['p2'] = np.concatenate([self['p2'], p2]) beta = np.ones(n_pairs) - dur = self.rng_mean_dur.poisson(self.mean_dur, p1) # TODO: Stream-ify... trying using p1 + dur = self.rng_mean_dur.poisson(p1, self.mean_dur) self['beta'] = np.concatenate([self['beta'], beta]) self['dur'] = np.concatenate([self['dur'], dur]) @@ -420,15 +426,9 @@ class stable_monogamy(simple_sexual): Very simple network for debugging in which edges are: 1-2, 3-4, 5-6, ... """ - def __init__(self): - key_dict = { - 'p1': ss.int_, - 'p2': ss.int_, - 'beta': ss.float_, - } - + def __init__(self, **kwargs): # Call init for the base class, which sets all the keys - super().__init__(mean_dur=np.iinfo(int).max) + super().__init__(mean_dur=np.iinfo(int).max, **kwargs) return def initialize(self, sim): @@ -444,7 +444,7 @@ def update(self, people, dt=None): class hpv_network(Network): - def __init__(self, pars=None): + def __init__(self, pars=None, **kwargs): key_dict = { 'p1': ss.int_, @@ -456,22 +456,22 @@ def __init__(self, pars=None): } # Call init for the base class, which sets all the keys - super().__init__(key_dict=key_dict) + super().__init__(key_dict=key_dict, **kwargs) # Define default parameters self.pars = dict() self.pars['cross_layer'] = 0.05 # Proportion of agents who have concurrent cross-layer relationships - self.pars['partners'] = dict(dist='poisson', par1=0.01) # The number of concurrent sexual partners - self.pars['acts'] = dict(dist='neg_binomial', par1=80, par2=40) # The number of sexual acts per year + self.pars['partners'] = dict(dist='poisson', par1=0.01) # The number of concurrent sexual partners # TODO: Stream-ify + self.pars['acts'] = dict(dist='neg_binomial', par1=80, par2=40) # The number of sexual acts per year # TODO: Stream-ify self.pars['age_act_pars'] = dict(peak=30, retirement=100, debut_ratio=0.5, retirement_ratio=0.1) # Parameters describing changes in coital frequency over agent lifespans self.pars['condoms'] = 0.2 # The proportion of acts in which condoms are used - self.pars['dur_pship'] = dict(dist='normal_pos', par1=1, par2=1) # Duration of partnerships + self.pars['dur_pship'] = dict(dist='normal_pos', par1=1, par2=1) # Duration of partnerships # TODO: Stream-ify self.pars['participation'] = None # Incidence of partnership formation by age self.pars['mixing'] = None # Mixing matrices for storing age differences in partnerships - self.rng_partners = ss.Stream('partners') - self.rng_acts = ss.Stream('acts') - self.rng_dur_pship = ss.Stream('dur_pship') + self.rng_partners = ss.Stream(self.multistream)('partners') + self.rng_acts = ss.Stream(self.multistream)('acts') + self.rng_dur_pship = ss.Stream(self.multistream)('dur_pship') self.update_pars(pars) self.get_layer_probs() @@ -666,12 +666,12 @@ def update(self, people, ti=None, dt=None): class maternal(Network): - def __init__(self, key_dict=None, vertical=True): + def __init__(self, key_dict=None, vertical=True, **kwargs): """ Initialized empty and filled with pregnancies throughout the simulation """ key_dict = sc.mergedicts({'dur': ss.float_}, key_dict) - super().__init__(key_dict=key_dict, vertical=vertical) + super().__init__(key_dict=key_dict, vertical=vertical, **kwargs) return def update(self, people, dt=None): diff --git a/stisim/parameters.py b/stisim/parameters.py index b2267a33..34de00e6 100644 --- a/stisim/parameters.py +++ b/stisim/parameters.py @@ -41,6 +41,7 @@ def __init__(self, **kwargs): self.dt = 1.0 # Timestep (in years) self.dt_demog = 1.0 # Timestep for demographic updates (in years) self.rand_seed = 1 # Random seed, if None, don't reset + self.multistream = True # Run with multiple random number streams as opposed to one centralized random number generator self.verbose = ss.options.verbose # Whether or not to display information during the run -- options are 0 (silent), 0.1 (some; default), 1 (default), 2 (everything) # Events and interventions diff --git a/stisim/sim.py b/stisim/sim.py index 2baa5aae..6015847c 100644 --- a/stisim/sim.py +++ b/stisim/sim.py @@ -287,7 +287,7 @@ def init_interventions(self): raise TypeError(errormsg) for stream in intervention.streams.values(): - stream.initialize(self.streams, self.people._uid_map) + stream.initialize(self.streams) return diff --git a/stisim/states.py b/stisim/states.py index f76dcd8e..fb4ff39c 100644 --- a/stisim/states.py +++ b/stisim/states.py @@ -267,8 +267,9 @@ def initialize(self, n): def grow(self, n): if self.n + n > self._s: - n_new = max(n, int(self._s / 2)) # Minimum 50% growth - self._data = np.concatenate([self._data, self._new_items(n)], axis=0) + # Grow the underlying _data + num_to_add = max(n, int(self._s / 2)) # Minimum 50% growth + self._data = np.concatenate([self._data, self._new_items(num_to_add)], axis=0) self.n += n self._map_arrays() diff --git a/stisim/streams.py b/stisim/streams.py index 5d24bc67..c72f809f 100644 --- a/stisim/streams.py +++ b/stisim/streams.py @@ -2,7 +2,7 @@ import sciris as sc import stisim as ss -__all__ = ['Streams', 'Stream'] +__all__ = ['Streams', 'MultiStream', 'CentralizedStream', 'Stream'] class Streams: @@ -35,7 +35,7 @@ def add(self, stream): raise RepeatNameException(f'A Stream with name {stream.name} has already been added.') if stream.seed_offset is None: - seed = len(self._streams) # Put at end by default + seed = abs(hash(stream.name)) elif stream.seed_offset in self.used_seeds: raise SeedRepeatException(f'Requested seed offset {stream.seed_offset} for stream {stream} has already been used.') else: @@ -77,9 +77,28 @@ class SeedRepeatException(Exception): pass +def Stream(multistream=True): + """ + Class to choose a stream + """ + if multistream: + return MultiStream + + return CentralizedStream + + def _pre_draw(func): - def check_ready(self, arr): + def check_ready(self, *args, **kwargs): """ Validation before drawing """ + + # Check for zero length arr + if 'arr' in kwargs.keys(): + arr = kwargs['arr'] + else: + arr = args[0] + if len(arr) == 0: + return np.array([], dtype=int) # int dtype allows use as index, e.g. bernoulli_filter + if not self.initialized: msg = f'Stream {self.name} has not been initialized!' raise NotInitializedException(msg) @@ -87,11 +106,11 @@ def check_ready(self, arr): msg = f'Stream {self.name} has already been sampled on this timestep!' raise NotResetException(msg) self.ready = False - return func(self, arr) + return func(self, *args, **kwargs) return check_ready -class Stream(np.random.Generator): +class MultiStream(np.random.Generator): """ Class for tracking one random number stream associated with one decision per timestep """ @@ -117,13 +136,12 @@ def __init__(self, name, seed_offset=None, **kwargs): self.kwargs = kwargs self.seed = None - self.bso = None # Block size object, typically sim.people._uid_map self.initialized = False self.ready = True return - def initialize(self, streams, block_size_object): + def initialize(self, streams): if self.initialized: # TODO: Raise warning assert not self.initialized @@ -138,8 +156,6 @@ def initialize(self, streams, block_size_object): #self.rng = np.random.default_rng(seed=self.seed + self.seed_offset) self._init_state = self.bit_generator.state # Store the initial state - self.bso = block_size_object - self.initialized = True self.ready = True return @@ -162,25 +178,38 @@ def step(self, ti): return - @property - def block_size(self): - return len(self.bso) + def draw_size(self, arr): + """ Determine how many random numbers to draw for a given arr """ + + if isinstance(arr, ss.states.FusedArray): + v = arr.values + elif isinstance(arr, ss.states.DynamicView): + v = arr._view + else: + v = arr + + if v.dtype == bool: + return len(arr) + + return v.max()+1 @_pre_draw def random(self, arr): - return super(Stream, self).random(self.block_size)[arr] + return super(MultiStream, self).random(self.draw_size(arr))[arr] - def poisson(self, lam, arr): - return super(Stream, self).poisson(lam=lam, size=self.block_size)[arr] + @_pre_draw + def poisson(self, arr, lam): + return super(MultiStream, self).poisson(lam=lam, size=self.draw_size(arr))[arr] - def bernoulli(self, prob, arr): - #return super(Stream, self).choice([True, False], size=self.block_size, p=[prob, 1-prob]) # very slow - #return (super(Stream, self).binomial(n=1, p=prob, size=self.block_size))[arr].astype(bool) # pretty fast - return super(Stream, self).random(self.block_size)[arr] < prob # fastest + @_pre_draw + def bernoulli(self, arr, prob): + #return super(MultiStream, self).choice([True, False], size=arr.max()+1) # very slow + #return (super(MultiStream, self).binomial(n=1, p=prob, size=arr.max()+1))[arr].astype(bool) # pretty fast + return super(MultiStream, self).random(self.draw_size(arr))[arr] < prob # fastest - def bernoulli_filter(self, prob, arr): - #return arr[self.bernoulli(prob, arr).nonzero()[0]] - return arr[self.bernoulli(prob, arr)] # Slightly faster on my machine for bernoulli to typecast + # @_pre_draw <-- handled by call to self.bernoullli + def bernoulli_filter(self, arr, prob): + return arr[self.bernoulli(arr, prob)] # Slightly faster on my machine for bernoulli to typecast def sample(self, dist=None, par1=None, par2=None, size=None, **kwargs): """ @@ -286,4 +315,173 @@ def sample(self, dist=None, par1=None, par2=None, size=None, **kwargs): errormsg = f'The selected distribution "{dist}" is not implemented; choices are: {sc.newlinejoin(choices)}' raise NotImplementedError(errormsg) + return samples + + +class CentralizedStream(np.random.Generator): + """ + Class to imitate the behavior of a centralized random number generator + """ + + def __init__(self, name, seed_offset=None, **kwargs): + """ + Create a random number stream + + seed_offset will be automatically assigned (sequentially in first-come order) if None + + name: a name for this Stream, like "coin_flip" + uid: an identifier added to the name to make it uniquely identifiable, for example the name or id of the calling class + """ + self.name = name + self.initialized = False + self.seed_offset = None # Not used, so override to avoid potential seed collisions in Streams. + return + + def initialize(self, streams): + if self.initialized: + # TODO: Raise warning + assert not self.initialized + return + + streams.add(self) # Seed is returned, but not used here + self.initialized = True + return + + def reset(self): + pass + + def step(self, ti): + pass + + def draw_size(self, arr): + """ Determine how many random numbers to draw for a given arr """ + + if isinstance(arr, ss.states.FusedArray): + v = arr.values + elif isinstance(arr, ss.states.DynamicView): + v = arr._view + else: + v = arr + + if v.dtype == bool: + return arr.sum() + + return len(arr) + + def random(self, arr): + return np.random.random(self.draw_size(arr)) + + def poisson(self, arr, lam): + return np.random.poisson(lam=lam, size=self.draw_size(arr)) + + def bernoulli(self, arr, prob): + return np.random.random(self.draw_size(arr)) < prob + + def bernoulli_filter(self, arr, prob): + return arr[self.bernoulli(arr, prob)] # Slightly faster on my machine for bernoulli to typecast + + def sample(self, dist=None, par1=None, par2=None, size=None, **kwargs): + """ + Draw a sample from the distribution specified by the input. The available + distributions are: + + - 'uniform' : uniform from low=par1 to high=par2; mean is equal to (par1+par2)/2 + - 'choice' : par1=array of choices, par2=probability of each choice + - 'normal' : normal with mean=par1 and std=par2 + - 'lognormal' : lognormal with mean=par1, std=par2 (parameters are for the lognormal, not the underlying normal) + - 'normal_pos' : right-sided normal (i.e. only +ve values), with mean=par1, std=par2 of the underlying normal + - 'normal_int' : normal distribution with mean=par1 and std=par2, returns only integer values + - 'lognormal_int' : lognormal distribution with mean=par1 and std=par2, returns only integer values + - 'poisson' : Poisson distribution with rate=par1 (par2 is not used); mean and variance are equal to par1 + - 'neg_binomial' : negative binomial distribution with mean=par1 and k=par2; converges to Poisson with k=∞ + - 'beta' : beta distribution with alpha=par1 and beta=par2; + - 'gamma' : gamma distribution with shape=par1 and scale=par2; + + Args: + self (Stream) : the random number generator stream + dist (str) : the distribution to sample from + par1 (float) : the "main" distribution parameter (e.g. mean) + par2 (float) : the "secondary" distribution parameter (e.g. std) + size (int) : the number of samples (default=1) + kwargs (dict) : passed to individual sampling functions + + Returns: + A length N array of samples + + **Examples**:: + + ss.sample() # returns Unif(0,1) + ss.sample(dist='normal', par1=3, par2=0.5) # returns Normal(μ=3, σ=0.5) + ss.sample(dist='lognormal_int', par1=5, par2=3) # returns lognormally distributed values with mean 5 and std 3 + + Notes: + Lognormal distributions are parameterized with reference to the underlying normal distribution (see: + https://docs.scipy.org/doc/numpy-1.14.0/reference/generated/numpy.random.lognormal.html), but this + function assumes the user wants to specify the mean and std of the lognormal distribution. + + Negative binomial distributions are parameterized with reference to the mean and dispersion parameter k + (see: https://en.wikipedia.org/wiki/Negative_binomial_distribution). The r parameter of the underlying + distribution is then calculated from the desired mean and k. For a small mean (~1), a dispersion parameter + of ∞ corresponds to the variance and standard deviation being equal to the mean (i.e., Poisson). For a + large mean (e.g. >100), a dispersion parameter of 1 corresponds to the standard deviation being equal to + the mean. + """ + + # Some of these have aliases, but these are the "official" names + choices = [ + 'uniform', + 'normal', + 'choice', + 'normal_pos', + 'normal_int', + 'lognormal', + 'lognormal_int', + 'poisson', + 'neg_binomial', + 'beta', + 'gamma', + ] + + # Ensure it's an integer + if size is not None and not isinstance(size, tuple): + size = int(size) + + # Compute distribution parameters and draw samples + # NB, if adding a new distribution, also add to choices above + if dist in ['unif', 'uniform']: + samples = np.random.uniform(low=par1, high=par2, size=size) + elif dist in ['choice']: + samples = np.random.choice(a=par1, p=par2, size=size, **kwargs) + elif dist in ['norm', 'normal']: + samples = np.random.normal(loc=par1, scale=par2, size=size) + elif dist == 'normal_pos': + samples = np.abs(np.random.normal(loc=par1, scale=par2, size=size)) + elif dist == 'normal_int': + samples = np.round(np.abs(np.random.normal(loc=par1, scale=par2, size=size))) + elif dist == 'poisson': + samples = np.random.poisson(rate=par1, n=size) # Use Numba version below for speed + elif dist == 'beta': + samples = np.random.beta(a=par1, b=par2, size=size) + elif dist == 'gamma': + samples = np.random.gamma(shape=par1, scale=par2, size=size) + elif dist in ['lognorm', 'lognormal', 'lognorm_int', 'lognormal_int']: + if (sc.isnumber(par1) and par1 > 0) or (sc.checktype(par1, 'arraylike') and (par1 > 0).all()): + mean = np.log( + par1 ** 2 / np.sqrt(par2 ** 2 + par1 ** 2)) # Computes the mean of the underlying normal distribution + sigma = np.sqrt(np.log(par2 ** 2 / par1 ** 2 + 1)) # Computes sigma for the underlying normal distribution + samples = np.random.lognormal(mean=mean, sigma=sigma, size=size) + else: + samples = np.zeros(size) + if '_int' in dist: + samples = np.round(samples) + # Calculate a and b using mean (par1) and variance (par2) + # https://stats.stackexchange.com/questions/12232/calculating-the-parameters-of-a-beta-distribution-using-the-mean-and-variance + elif dist == 'beta_mean': + a = ((1 - par1) / par2 - 1 / par1) * par1 ** 2 + b = a * (1 / par1 - 1) + samples = np.random.beta(a=a, b=b, size=size) + else: + errormsg = f'The selected distribution "{dist}" is not implemented; choices are: {sc.newlinejoin(choices)}' + raise NotImplementedError(errormsg) + return samples \ No newline at end of file diff --git a/tests/run_multistream.py b/tests/run_multistream.py new file mode 100644 index 00000000..b3411937 --- /dev/null +++ b/tests/run_multistream.py @@ -0,0 +1,140 @@ +""" +Run a simple HIV simulation with random number coherence +""" + +# %% Imports and settings +import os +import stisim as ss +import matplotlib.pyplot as plt +import sciris as sc +import pandas as pd +import seaborn as sns + +n = 1_000 # Agents +n_rand_seeds = 250 +art_cov_levels = [0.025, 0.05, 0.10, 0.73] + [0] # Must include 0 as that's the baseline + +figdir = 'figs' +sc.path(figdir).mkdir(parents=True, exist_ok=True) + +def run_sim(n, art_cov, rand_seed, multistream): + ppl = ss.People(n) + + net_pars = {'multistream': multistream} + ppl.networks = ss.ndict(ss.simple_embedding(pars=net_pars), ss.maternal(pars=net_pars)) + + hiv_pars = { + 'beta': {'simple_embedding': [0.10, 0.08], 'maternal': [1,0]}, + 'initial': 10, + 'multistream': multistream, + } + hiv = ss.HIV(hiv_pars) + + preg_pars = {'multistream': multistream} + pregnancy = ss.Pregnancy(preg_pars) + + art_pars = {'multistream': multistream} + art = ss.hiv.ART(t=[0, 1], coverage=[0, art_cov], pars=art_pars) + pars = { + 'start': 1980, + 'end': 2010, + 'interventions': [art], + 'rand_seed': rand_seed, + 'multistream': multistream, + } + sim = ss.Sim(people=ppl, modules=[hiv, pregnancy], pars=pars, label=f'Sim with {n} agents and art_cov={art_cov}') + sim.initialize() + sim.run() + + df = pd.DataFrame( { + 'ti': sim.tivec, + #'hiv.n_infected': sim.results.hiv.n_infected, + 'hiv.prevalence': sim.results.hiv.prevalence, + 'hiv.cum_deaths': sim.results.hiv.new_deaths.cumsum(), + 'pregnancy.cum_births': sim.results.pregnancy.births.cumsum(), + }) + df['art_cov'] = art_cov + df['rand_seed'] = rand_seed + df['multistream'] = multistream + + return df + +def run_scenarios(): + results = [] + times = {} + for multistream in [True, False]: + cfgs = [] + for rs in range(n_rand_seeds): + for art_cov in art_cov_levels: + cfgs.append({'art_cov':art_cov, 'rand_seed':rs, 'multistream':multistream}) + T = sc.tic() + results += sc.parallelize(run_sim, kwargs={'n': n}, iterkwargs=cfgs, die=True) + times[f'Multistream={multistream}'] = sc.toc(T, output=True) + + print('Timings:', times) + + df = pd.concat(results) + df.to_csv(os.path.join(figdir, 'result.csv')) + return df + +def plot_scenarios(df): + d = pd.melt(df, id_vars=['ti', 'rand_seed', 'art_cov', 'multistream'], var_name='channel', value_name='Value') + d['baseline'] = d['art_cov']==0 + bl = d.loc[d['baseline']] + scn = d.loc[~d['baseline']] + bl = bl.set_index(['ti', 'channel', 'rand_seed', 'art_cov', 'multistream'])[['Value']].reset_index('art_cov') + scn = scn.set_index(['ti', 'channel', 'rand_seed', 'art_cov', 'multistream'])[['Value']].reset_index('art_cov') + mrg = scn.merge(bl, on=['ti', 'channel', 'rand_seed', 'multistream'], suffixes=('', '_ref')) + mrg['Value - Reference'] = mrg['Value'] - mrg['Value_ref'] + mrg = mrg.sort_index() + + fkw = {'sharey': False, 'sharex': 'col', 'margin_titles': True} + + ## TIMESERIES + g = sns.relplot(kind='line', data=d, x='ti', y='Value', hue='art_cov', col='channel', row='multistream', + height=5, aspect=1.2, palette='Set1', errorbar='sd', lw=2, facet_kws=fkw) + g.set_titles(col_template='{col_name}', row_template='Multistream: {row_name}') + g.set_xlabels(r'$t_i$') + g.fig.savefig(os.path.join(figdir, 'timeseries.png'), bbox_inches='tight', dpi=300) + + ## DIFF TIMESERIES + for ms, mrg_by_ms in mrg.groupby('multistream'): + g = sns.relplot(kind='line', data=mrg_by_ms, x='ti', y='Value - Reference', hue='art_cov', col='channel', row='art_cov', + height=3, aspect=1.0, palette='Set1', estimator=None, units='rand_seed', lw=0.5, facet_kws=fkw) #errorbar='sd', lw=2, + g.set_titles(col_template='{col_name}', row_template='ART cov: {row_name}') + g.fig.suptitle('Multiscale' if ms else 'Centralized') + g.fig.subplots_adjust(top=0.88) + g.set_xlabels(r'Value - Reference at $t_i$') + g.fig.savefig(os.path.join(figdir, 'diff_multistream.png' if ms else 'diff_centralized.png'), bbox_inches='tight', dpi=300) + + ## FINAL TIME + tf = df['ti'].max() + mtf = mrg.loc[tf] + g = sns.displot(data=mtf.reset_index(), kind='kde', fill=True, rug=True, cut=0, hue='art_cov', x='Value - Reference', col='channel', row='multistream', + height=5, aspect=1.2, facet_kws=fkw, palette='Set1') + g.set_titles(col_template='{col_name}', row_template='Multistream: {row_name}') + g.set_xlabels(f'Value - Reference at $t_i={{{tf}}}$') + g.fig.savefig(os.path.join(figdir, 'final.png'), bbox_inches='tight', dpi=300) + + print('Figures saved to:', os.path.join(os.getcwd(), figdir)) + + return + + +if __name__ == '__main__': + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument('-p', '--plot', help='plot from a cached CSV file', type=str) + args = parser.parse_args() + + if args.plot: + print('Reading CSV file', args.plot) + df = pd.read_csv(args.plot, index_col=0) + else: + print('Running scenarios') + df = run_scenarios() + + plot_scenarios(df) + + print('Done') \ No newline at end of file diff --git a/tests/stable_monogamy.py b/tests/stable_monogamy.py index 39920370..b7957e50 100644 --- a/tests/stable_monogamy.py +++ b/tests/stable_monogamy.py @@ -10,6 +10,7 @@ import networkx as nx import sys +multistream = True # Can set multistream to False for comparison plot_graph = True class Graph(): @@ -39,7 +40,7 @@ def plot(self, pos, ax=None): nx.draw_networkx_edges(self.graph, pos=pos, ax=ax) nx.draw_networkx_labels(self.graph, labels={i:int(a['cd4']) for i,a in self.graph.nodes.data()}, font_size=8, pos=pos, ax=ax) - nx.draw_networkx_edge_labels(self.graph, edge_labels={(i,j): int(a['dur']) for i,j,a in self.graph.edges.data()}, font_size=8, pos=pos, ax=ax) + #nx.draw_networkx_edge_labels(self.graph, edge_labels={(i,j): int(a['dur']) for i,j,a in self.graph.edges.data()}, font_size=8, pos=pos, ax=ax) return @@ -79,18 +80,27 @@ def finalize(self, sim): def run_sim(n=25, intervention=False, analyze=False): ppl = ss.People(n) - ppl.networks = ss.ndict(ss.stable_monogamy())#, ss.maternal()) - hiv = ss.HIV() - hiv.pars['beta'] = {'stable_monogamy': [0.06, 0.04]} - hiv.pars['initial'] = 0 + net_pars = {'multistream': multistream} + ppl.networks = ss.ndict(ss.stable_monogamy(pars=net_pars))#, ss.maternal(net_pars)) + + hiv_pars = { + 'beta': {'stable_monogamy': [0.06, 0.04]}, + 'initial': 0, + 'multistream': multistream, + } + hiv = ss.HIV(hiv_pars) + + art_pars = {'multistream': multistream} + art = ss.hiv.ART(0, 0.5, pars=art_pars) pars = { 'start': 1980, 'end': 2010, - 'interventions': [ss.hiv.ART(0, 0.5)] if intervention else [], # ss.hiv.PrEP(0, 0.2), + 'interventions': [art] if intervention else [], 'rand_seed': 0, - 'analyzers': [rng_analyzer()] if analyze else [], + 'analyzers': [rng_analyzer(pars={'multistream': multistream})] if analyze else [], + 'multistream': multistream, } sim = ss.Sim(people=ppl, modules=[hiv], pars=pars, label=f'Sim with {n} agents and intv={intervention}') sim.initialize() diff --git a/tests/test_hiv.py b/tests/test_hiv.py deleted file mode 100644 index a41afeb7..00000000 --- a/tests/test_hiv.py +++ /dev/null @@ -1,46 +0,0 @@ -""" -Run a simple HIV simulation with random number coherence -""" - -# %% Imports and settings -import stisim as ss -import matplotlib.pyplot as plt -import sciris as sc - -n = 1_000 # Agents - -def run_sim(n=25, intervention=False, analyze=False): - ppl = ss.People(n) - ppl.networks = ss.ndict(ss.simple_embedding())#, ss.maternal()) - - hiv = ss.HIV() - hiv.pars['beta'] = {'simple_embedding': [0.10, 0.08]} - hiv.pars['initial'] = 10 - - pars = { - 'start': 1980, - 'end': 2010, - 'interventions': [ss.hiv.ART(t=[0, 1], coverage=[0, 0.9**3])] if intervention else [], - 'rand_seed': 0, - } - sim = ss.Sim(people=ppl, modules=[hiv], pars=pars, label=f'Sim with {n} agents and intv={intervention}') - sim.initialize() - sim.run() - - return sim - -sim1, sim2 = sc.parallelize(run_sim, kwargs={'n': n}, iterkwargs=[{'intervention':False}, {'intervention':True}], die=True) - -# Plot -fig, axv = plt.subplots(2,1, sharex=True) -axv[0].plot(sim1.tivec, sim1.results.hiv.n_infected, label='Baseline') -axv[0].plot(sim2.tivec, sim2.results.hiv.n_infected, ls=':', label='Intervention') -axv[0].set_title('HIV number of infections') - -axv[1].plot(sim1.tivec, sim1.results.hiv.new_deaths, label='Baseline') -axv[1].plot(sim2.tivec, sim2.results.hiv.new_deaths, ls=':', label='Intervention') -axv[1].set_title('HIV Deaths') - -plt.legend() -plt.show() -print('Done') \ No newline at end of file diff --git a/tests/test_stream.py b/tests/test_stream.py index 48506373..2d93cfb1 100644 --- a/tests/test_stream.py +++ b/tests/test_stream.py @@ -6,38 +6,37 @@ import numpy as np import sciris as sc import stisim as ss -from stisim.streams import NotResetException +from stisim.streams import NotResetException, Stream -def make_rng(n, base_seed=1, name='Test'): +def make_rng(multistream=False, base_seed=1, name='Test'): """ Create and initialize a stream """ streams = ss.Streams() streams.initialize(base_seed=base_seed) - rng = ss.Stream(name) - bso = np.arange(n) - rng.initialize(streams, bso) + rng = Stream(multistream)(name) + rng.initialize(streams) return rng # %% Define the tests -def test_sample(n=5): +def test_sample(multistream=True, n=5): """ Simple sample """ sc.heading('Testing stream object') - rng = make_rng(n) + rng = make_rng(multistream) uids = np.arange(0,n,2) # every other to make it interesting draws = rng.random(uids) print(f'\nSAMPLE({n}): {draws}') - return draws + return len(draws) == len(uids) -def test_reset(n=5): +def test_reset(multistream=True, n=5): """ Sample, reset, sample """ sc.heading('Testing sample, reset, sample') - rng = make_rng(n) + rng = make_rng(multistream) uids = np.arange(0,n,2) # every other to make it interesting draws1 = rng.random(uids) @@ -49,13 +48,13 @@ def test_reset(n=5): draws2 = rng.random(uids) print(f'\nSAMPLE({n}): {draws2}') - return np.all(draws2-draws1 == 0) + return np.all(np.equal(draws1, draws2)) -def test_step(n=5): +def test_step(multistream=True, n=5): """ Sample, step, sample """ sc.heading('Testing sample, step, sample') - rng = make_rng(n) + rng = make_rng(multistream) uids = np.arange(0,n,2) # every other to make it interesting draws1 = rng.random(uids) @@ -67,29 +66,29 @@ def test_step(n=5): draws2 = rng.random(uids) print(f'\nSAMPLE({n}): {draws2}') - return np.all(draws2-draws1 != 0) + return np.all(np.equal(draws1, draws2)) -def test_seed(n=5): +def test_seed(multistream=True, n=5): """ Sample, step, sample """ sc.heading('Testing sample with seeds 0 and 1') uids = np.arange(0,n,2) # every other to make it interesting - rng0 = make_rng(n, based_seed=0) + rng0 = make_rng(multistream, base_seed=0) draws0 = rng0.random(uids) print(f'\nSAMPLE({n}): {draws0}') - rng1 = make_rng(n, based_seed=1) + rng1 = make_rng(multistream, base_seed=1) draws1 = rng1.random(uids) print(f'\nSAMPLE({n}): {draws1}') - return np.all(draws1-draws0 != 0) + return np.all(np.equal(draws0, draws1)) -def test_repeat(n=5): +def test_repeat(multistream=True, n=5): """ Sample, sample - should raise and exception""" sc.heading('Testing sample, sample - should raise an exception') - rng = make_rng(n) + rng = make_rng(multistream) uids = np.arange(0,n,2) # every other to make it interesting draws1 = rng.random(uids) @@ -104,17 +103,78 @@ def test_repeat(n=5): return True +def test_boolmask(multistream=True, n=5): + """ Simple sample with a boolean mask""" + sc.heading('Testing stream object') + rng = make_rng(multistream) + uids = np.arange(0,n,2) # every other to make it interesting + mask = np.full(n, False) + mask[uids] = True + + draws_bool = rng.random(mask) + print(f'\nSAMPLE({n}): {draws_bool}') + + rng.reset() + + draws_uids = rng.random(uids) + print(f'\nSAMPLE({n}): {draws_uids}') + + return np.all(np.equal(draws_bool, draws_uids)) + + +def test_empty(multistream=True): + """ Simple sample with a boolean mask""" + sc.heading('Testing empty draw') + rng = make_rng(multistream) + uids = np.array([]) # EMPTY + draws = rng.random(uids) + print(f'\nSAMPLE: {draws}') + + return len(draws) == 0 + + +def test_drawsize(): + """ Testing the draw_size function directly """ + sc.heading('Testing draw size') + + rng = make_rng(multistream) + + x = ss.states.FusedArray(values=np.array([1,3,9]), uid=np.array([0,1,2])) + ds_FA = rng.draw_size(x) == 10 # Should be 10 because max value of x is 9 + + x = ss.states.DynamicView(int, fill_value=0) + x.initialize(3) + x[0] = 9 + ds_DV = rng.draw_size(x) == 10 # Should be 10 because max value (representing uid) is 9 + + x = np.full(10, fill_value=True) + ds_bool = rng.draw_size(x) == 10 # Should be 10 because 10 objects + + x = np.array([9]) + ds_array = rng.draw_size(x) == 10 # Should be 10 because 10 objects + + return np.all([ds_FA, ds_DV, ds_bool, ds_array]) + + # %% Run as a script if __name__ == '__main__': # Start timing T = sc.tic() - # Run tests - test_sample() - assert test_reset() - assert test_step() - assert test_seed() - assert test_repeat() + n=5 + + for multistream in [True, False]: + print('Testing with multistream set to', multistream) + + # Run tests + print(test_sample(multistream, n)) + print(test_reset(multistream, n)) + print(test_step(multistream, n)) + print(test_seed(multistream, n)) + print(test_repeat(multistream, n)) + print(test_boolmask(multistream, n)) + print(test_empty(multistream)) + print(test_drawsize()) sc.toc(T) print('Done.') \ No newline at end of file diff --git a/tests/test_streams.py b/tests/test_streams.py index 265e9dc0..0d79058d 100644 --- a/tests/test_streams.py +++ b/tests/test_streams.py @@ -17,9 +17,8 @@ def test_streams(n=5): streams = ss.Streams() streams.initialize(base_seed=10) - rng = ss.Stream('stream1') - bso = np.arange(n) - rng.initialize(streams, bso) + rng = ss.MultiStream('stream1') + rng.initialize(streams) uids = np.arange(0,n,2) # every other to make it interesting draws = rng.random(uids) @@ -34,13 +33,12 @@ def test_seed(n=5): streams = ss.Streams() streams.initialize(base_seed=10) - bso = np.arange(n) - rng0 = ss.Stream('stream0') - rng0.initialize(streams, bso) + rng0 = ss.MultiStream('stream0') + rng0.initialize(streams) - rng1 = ss.Stream('stream1') - rng1.initialize(streams, bso) + rng1 = ss.MultiStream('stream1') + rng1.initialize(streams) return rng1.seed != rng0.seed @@ -48,13 +46,11 @@ def test_seed(n=5): def test_reset(n=5): """ Sample, step, sample """ sc.heading('Testing sample, step, sample') - streams = ss.Streams() + streams = ss.MultiStream() streams.initialize(base_seed=10) - bso = np.arange(n) - - rng = ss.Stream('stream0') - rng.initialize(streams, bso) + rng = ss.MultiStream('stream0') + rng.initialize(streams) uids = np.arange(0,n,2) # every other to make it interesting s_before = rng.random(uids) @@ -72,10 +68,8 @@ def test_step(n=5): streams = ss.Streams() streams.initialize(base_seed=10) - bso = np.arange(n) - - rng = ss.Stream('stream0') - rng.initialize(streams, bso) + rng = ss.MultiStream('stream0') + rng.initialize(streams) uids = np.arange(0,n,2) # every other to make it interesting s_before = rng.random(uids) @@ -93,12 +87,10 @@ def test_initialize(n=5): streams = ss.Streams() #streams.initialize(base_seed=3) - bso = np.arange(n) - - rng = ss.Stream('stream0') + rng = ss.MultiStream('stream0') try: - rng.initialize(streams, bso) + rng.initialize(streams) return False # Should not get here! except NotInitializedException as e: print(f'YAY! Got exception: {e}') @@ -111,19 +103,18 @@ def test_seedrepeat(n=5): streams = ss.Streams() streams.initialize(base_seed=10) - bso = np.arange(n) - - rng = ss.Stream('stream0', seed_offset=0) - rng.initialize(streams, bso) + rng = ss.MultiStream('stream0', seed_offset=0) + rng.initialize(streams) try: - rng1 = ss.Stream('stream1', seed_offset=0) - rng1.initialize(streams, bso) + rng1 = ss.MultiStream('stream1', seed_offset=0) + rng1.initialize(streams) return False # Should not get here! except SeedRepeatException as e: print(f'YAY! Got exception: {e}') return True + def test_samplingorder(n=5): """ Ensure sampling from one stream doesn't affect another """ sc.heading('Testing from multiple streams to test if sampling order matters') @@ -131,13 +122,12 @@ def test_samplingorder(n=5): streams.initialize(base_seed=10) uids = np.arange(0,n,2) # every other to make it interesting - bso = np.arange(n) - rng0 = ss.Stream('stream0') - rng0.initialize(streams, bso) + rng0 = ss.MultiStream('stream0') + rng0.initialize(streams) - rng1 = ss.Stream('stream1') - rng1.initialize(streams, bso) + rng1 = ss.MultiStream('stream1') + rng1.initialize(streams) s_before = rng0.random(uids) _ = rng1.random(uids) @@ -156,37 +146,18 @@ def test_repeatname(n=5): streams = ss.Streams() streams.initialize(base_seed=17) - uids = np.arange(0,n,2) # every other to make it interesting - bso = np.arange(n) - - rng0 = ss.Stream('test') - rng0.initialize(streams, bso) + rng0 = ss.MultiStream('test') + rng0.initialize(streams) - rng1 = ss.Stream('test') + rng1 = ss.MultiStream('test') try: - rng1.initialize(streams, bso) + rng1.initialize(streams) return False # Should not get here! except RepeatNameException as e: print(f'YAY! Got exception: {e}') return True - - - - - s_before = rng0.random(uids) - _ = rng1.random(uids) - - streams.reset() - - _ = rng1.random(uids) - s_after = rng0.random(uids) - - return np.all(s_before == s_after) - - - # %% Run as a script if __name__ == '__main__': # Start timing From a6890898764abe8d63ae1b65474d2fb93d0cd344 Mon Sep 17 00:00:00 2001 From: Daniel Klein Date: Wed, 27 Sep 2023 18:28:37 -0700 Subject: [PATCH 18/34] * Fixing random calls in simple_sexual, which is not rng safe. * renaming stable_monogamy.py to run_stable_monogamy.py * Adjusting test_base.py, it need some work to restore test_networks() --- stisim/networks.py | 12 ++++++------ tests/{stable_monogamy.py => run_stable_monogamy.py} | 0 tests/test_base.py | 10 ++++++++-- 3 files changed, 14 insertions(+), 8 deletions(-) rename tests/{stable_monogamy.py => run_stable_monogamy.py} (100%) diff --git a/stisim/networks.py b/stisim/networks.py index e755d3e2..c7c2e921 100644 --- a/stisim/networks.py +++ b/stisim/networks.py @@ -316,13 +316,13 @@ def add_pairs(self, people, ti=None): if len(available_m) <= len(available_f): p1 = available_m - p2 = self.rng_pair_12.choice(len(p1), available_f, replace=False) # TODO: Stream-ify + p2 = self.rng_pair_12.choice(available_f, size=len(p1), replace=False) # TODO: Stream-ify else: p2 = available_f - p1 = self.rng_pair_21.choice(len(p2), available_m, replace=False) # TODO: Stream-ify + p1 = self.rng_pair_21.choice(available_m, size=len(p2), replace=False) # TODO: Stream-ify beta = np.ones_like(p1) - dur = self.rng_mean_dur.poisson(len(p1), self.mean_dur) # TODO: Stream-ify + dur = self.rng_mean_dur.poisson(p1, self.mean_dur) # TODO: Stream-ify self['p1'] = np.concatenate([self['p1'], p1]) self['p2'] = np.concatenate([self['p2'], p2]) self['beta'] = np.concatenate([self['beta'], beta]) @@ -479,9 +479,9 @@ def __init__(self, pars=None, **kwargs): def initialize(self, sim): super().initialize(sim) - self.rng_partners.initialize(sim) - self.rng_acts.initialize(sim) - self.rng_dur_pship.initialize(sim) + self.rng_partners.initialize(sim.streams) + self.rng_acts.initialize(sim.streams) + self.rng_dur_pship.initialize(sim.streams) self.add_pairs(sim.people, ti=0) return diff --git a/tests/stable_monogamy.py b/tests/run_stable_monogamy.py similarity index 100% rename from tests/stable_monogamy.py rename to tests/run_stable_monogamy.py diff --git a/tests/test_base.py b/tests/test_base.py index 4d70ef5b..3c8fc5f6 100644 --- a/tests/test_base.py +++ b/tests/test_base.py @@ -41,16 +41,22 @@ def test_networks(): nw1 = ss.Network(p1=p1, p2=p2, beta=beta, label='rand') nw2 = ss.Network(dict(p1=p1, p2=p2, beta=beta), label='rand') # Alternate method + #TODO: The following two tests do not work because network initialization + # expecting a Sim object, not a People object + nw3 = None + nw4 = None + ''' # Make people, then make a dynamic sexual layer and update it ppl = ss.People(100) # BasePeople ppl.initialize() # This seems to be necessary, although not completely clear why... nw3 = ss.hpv_network() - nw3.initialize(ppl) + nw3.initialize(ppl) #TODO: Initialize is expecting a Sim object, not a People object nw3.update(ppl, ti=1, dt=1) # Update by providing a timestep & current time index nw4 = ss.maternal() - nw4.initialize(ppl) + nw4.initialize(ppl) #TODO: Initialize is expecting a Sim object, not a People object nw4.add_pairs(mother_inds=[1, 2, 3], unborn_inds=[100, 101, 102], dur=[1, 1, 1]) + ''' return nw1, nw2, nw3, nw4 From 104097371071660f534dc68b7f313979dd9e2330 Mon Sep 17 00:00:00 2001 From: Daniel Klein Date: Wed, 27 Sep 2023 22:47:24 -0700 Subject: [PATCH 19/34] Now using the "linear_sum_assignment" from scipy.optimize. It's faster and better! --- stisim/networks.py | 58 ++++++++++------------------------------------ 1 file changed, 12 insertions(+), 46 deletions(-) diff --git a/stisim/networks.py b/stisim/networks.py index c7c2e921..6c511813 100644 --- a/stisim/networks.py +++ b/stisim/networks.py @@ -6,6 +6,7 @@ import numpy as np import sciris as sc import stisim as ss +import scipy.optimize as spo # Specify all externally visible functions this file defines @@ -363,58 +364,23 @@ def add_pairs(self, people, ti=None): print('No pairs to add') return 0 - # Make p1 the shorter array - if len(available_f) < len(available_m): - switch = True # Want p1 as m in the end - p1 = available_f - p2 = available_m - else: - switch = False # Want p1 as m in the end - p1 = available_m - p2 = available_f - - loc1 = self.rng_pair_12.random(arr=p1) - loc2 = self.rng_pair_21.random(arr=p2) + ### NEW + loc_m = self.rng_pair_12.random(arr=available_m) + loc_f = self.rng_pair_21.random(arr=available_f) - p1v = np.tile(loc1, (len(loc2),1)) - p2v = np.tile(loc2[:,np.newaxis], len(loc1)) + p1v = np.tile(loc_m, (len(loc_f),1)) + p2v = np.tile(loc_f[:,np.newaxis], len(loc_m)) d_full = np.absolute(p2v-p1v) - d = d_full.copy() - - unmatched_p1i = np.arange(len(p1)) - unmatched_p2i = np.arange(len(p2)) - - pairs = [] - - while len(unmatched_p1i)>0: - # Perhaps more efficient to change up directionality of matching at times? - p2i_closest_to_each_p1 = d.argmin(axis=0) - selected_up2i, selected_up1i = np.unique(p2i_closest_to_each_p1, return_index=True) - selected_p1i = unmatched_p1i[selected_up1i] - selected_p2i = unmatched_p2i[selected_up2i] - # loc1[selected_p1i] should be close to loc2[selected_p2i] - pairs.append( (p1[selected_p1i], p2[selected_p2i]) ) + ind_f, ind_m = spo.linear_sum_assignment(d_full) + # loc_f[ind_f[0]] is close to loc_m[ind_m[0]] - # Remove pairs and repeat - unmatched_p1i = np.setdiff1d(unmatched_p1i, selected_p1i) - unmatched_p2i = np.setdiff1d(unmatched_p2i, selected_p2i) - - # Trim distance matrix - d = d_full[np.ix_(unmatched_p2i, unmatched_p1i)] - - if ss.options.verbose > 1: - print(f'Matching with {len(unmatched_p1i)} to go') - - pairs = np.concatenate(pairs, axis=1) - n_pairs = pairs.shape[1] - - (p1, p2) = (pairs[1], pairs[0]) if switch else (pairs[0], pairs[1]) - self['p1'] = np.concatenate([self['p1'], p1]) - self['p2'] = np.concatenate([self['p2'], p2]) + n_pairs = len(ind_f) + self['p1'] = np.concatenate([self['p1'], ind_m]) + self['p2'] = np.concatenate([self['p2'], ind_f]) beta = np.ones(n_pairs) - dur = self.rng_mean_dur.poisson(p1, self.mean_dur) + dur = self.rng_mean_dur.poisson(ind_m, self.mean_dur) self['beta'] = np.concatenate([self['beta'], beta]) self['dur'] = np.concatenate([self['dur'], dur]) From c63b98ff1935f1eb13114a451da6b1dccc200428 Mon Sep 17 00:00:00 2001 From: Daniel Klein Date: Wed, 27 Sep 2023 23:36:21 -0700 Subject: [PATCH 20/34] * Age-assortative mixing! * Now using scipy to compute the distance matrix, it's faster * Added a normal draw to streams. --- stisim/networks.py | 13 +++++-------- stisim/streams.py | 7 +++++++ tests/run_multistream.py | 4 ++-- 3 files changed, 14 insertions(+), 10 deletions(-) diff --git a/stisim/networks.py b/stisim/networks.py index 6c511813..81665460 100644 --- a/stisim/networks.py +++ b/stisim/networks.py @@ -7,6 +7,7 @@ import sciris as sc import stisim as ss import scipy.optimize as spo +import scipy.spatial as sps # Specify all externally visible functions this file defines @@ -364,15 +365,11 @@ def add_pairs(self, people, ti=None): print('No pairs to add') return 0 - ### NEW - loc_m = self.rng_pair_12.random(arr=available_m) - loc_f = self.rng_pair_21.random(arr=available_f) + loc_m = people.age[available_m].values - 5 + self.rng_pair_12.normal(arr=available_m, std=3) + loc_f = people.age[available_f].values + self.rng_pair_21.normal(arr=available_f, std=3) + dist_mat = sps.distance_matrix(loc_m[:, np.newaxis], loc_f[:, np.newaxis]) - p1v = np.tile(loc_m, (len(loc_f),1)) - p2v = np.tile(loc_f[:,np.newaxis], len(loc_m)) - d_full = np.absolute(p2v-p1v) - - ind_f, ind_m = spo.linear_sum_assignment(d_full) + ind_m, ind_f = spo.linear_sum_assignment(dist_mat) # loc_f[ind_f[0]] is close to loc_m[ind_m[0]] n_pairs = len(ind_f) diff --git a/stisim/streams.py b/stisim/streams.py index c72f809f..b346981e 100644 --- a/stisim/streams.py +++ b/stisim/streams.py @@ -201,6 +201,10 @@ def random(self, arr): def poisson(self, arr, lam): return super(MultiStream, self).poisson(lam=lam, size=self.draw_size(arr))[arr] + @_pre_draw + def normal(self, arr, mu=0, std=1): + return mu + std*super(MultiStream, self).normal(size=self.draw_size(arr))[arr] + @_pre_draw def bernoulli(self, arr, prob): #return super(MultiStream, self).choice([True, False], size=arr.max()+1) # very slow @@ -374,6 +378,9 @@ def random(self, arr): def poisson(self, arr, lam): return np.random.poisson(lam=lam, size=self.draw_size(arr)) + def normal(self, arr, mu=0, std=1): + return mu + std*np.random.normal(size=self.draw_size(arr), loc=mu, scale=std) + def bernoulli(self, arr, prob): return np.random.random(self.draw_size(arr)) < prob diff --git a/tests/run_multistream.py b/tests/run_multistream.py index b3411937..73c4d919 100644 --- a/tests/run_multistream.py +++ b/tests/run_multistream.py @@ -15,7 +15,7 @@ art_cov_levels = [0.025, 0.05, 0.10, 0.73] + [0] # Must include 0 as that's the baseline figdir = 'figs' -sc.path(figdir).mkdir(parents=True, exist_ok=True) +#sc.path(figdir).mkdir(parents=True, exist_ok=True) def run_sim(n, art_cov, rand_seed, multistream): ppl = ss.People(n) @@ -102,7 +102,7 @@ def plot_scenarios(df): g = sns.relplot(kind='line', data=mrg_by_ms, x='ti', y='Value - Reference', hue='art_cov', col='channel', row='art_cov', height=3, aspect=1.0, palette='Set1', estimator=None, units='rand_seed', lw=0.5, facet_kws=fkw) #errorbar='sd', lw=2, g.set_titles(col_template='{col_name}', row_template='ART cov: {row_name}') - g.fig.suptitle('Multiscale' if ms else 'Centralized') + g.fig.suptitle('Multistream' if ms else 'Centralized') g.fig.subplots_adjust(top=0.88) g.set_xlabels(r'Value - Reference at $t_i$') g.fig.savefig(os.path.join(figdir, 'diff_multistream.png' if ms else 'diff_centralized.png'), bbox_inches='tight', dpi=300) From 6e60267ce7758569ccec843dd27dba5a4c447e5e Mon Sep 17 00:00:00 2001 From: Daniel Klein Date: Thu, 28 Sep 2023 17:41:01 -0700 Subject: [PATCH 21/34] * Updated HIV interventions * run_multistream now compares multstrem on and off for either ART or PrEP * New run_sweep.py sweeps beta with multstream on or off --- stisim/hiv.py | 26 ++++---- tests/run_multistream.py | 45 +++++++------ tests/run_sweep.py | 137 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 176 insertions(+), 32 deletions(-) create mode 100644 tests/run_sweep.py diff --git a/stisim/hiv.py b/stisim/hiv.py index dfdbe332..c1879146 100644 --- a/stisim/hiv.py +++ b/stisim/hiv.py @@ -80,8 +80,8 @@ def __init__(self, t: np.array, coverage: np.array, **kwargs): super().__init__(**kwargs) - self.rng_add_ART = ss.Stream(self.multistream)('add_ART', seed_offset=100) - self.rng_remove_ART = ss.Stream(self.multistream)('remove_ART', seed_offset=101) + self.rng_add_ART = ss.Stream(self.multistream)('add_ART') + self.rng_remove_ART = ss.Stream(self.multistream)('remove_ART') return @@ -119,13 +119,15 @@ def apply(self, sim): class PrEP(ss.Intervention): - def __init__(self, t: np.array, coverage: np.array): + def __init__(self, t: np.array, coverage: np.array, **kwargs): self.requires = HIV self.t = sc.promotetoarray(t) self.coverage = sc.promotetoarray(coverage) - self.rng_add_PrEP = ss.Stream(self.multistream)('add_PrEP', seed_offset=102) - self.rng_remove_PrEP = ss.Stream(self.multistream)('remove_PrEP', seed_offset=103) + super().__init__(**kwargs) + + self.rng_add_PrEP = ss.Stream(self.multistream)('add_PrEP') + self.rng_remove_PrEP = ss.Stream(self.multistream)('remove_PrEP') return @@ -140,20 +142,22 @@ def apply(self, sim): coverage = self.coverage[np.where(self.t <= sim.ti)[0][-1]] on_prep = sim.people.alive & sim.people.hiv.on_prep - n_change = np.round(coverage * sim.people.alive.sum() - np.count_nonzero(on_prep)).astype(int) + sus = sim.people.alive & ~sim.people.hiv.infected + n_change = np.round(coverage * sus.sum() - on_prep.sum()).astype(int) if n_change > 0: # Add more PrEP eligible = ss.true(sim.people.alive & ~sim.people.hiv.infected & ~sim.people.hiv.on_prep) - n_eligible = len(eligible) #np.count_nonzero(eligible) + n_eligible = len(eligible) if n_eligible: inds = self.rng_add_PrEP.bernoulli_filter(arr=eligible, prob=min(n_eligible, n_change)/n_eligible) sim.people.hiv.on_prep[inds] = True elif n_change < 0: # Take some people off PrEP - eligible = sim.people.alive & sim.people.hiv.on_prep - n_eligible = np.count_nonzero(eligible) - inds = self.rng_remove_PrEP.bernoulli_filter(arr=eligible, prob=-n_change/n_eligible) - sim.people.hiv.on_prep[inds] = False + eligible = ss.true(sim.people.alive & ~sim.people.hiv.infected & sim.people.hiv.on_prep) + n_eligible = len(eligible) + if n_eligible: + inds = self.rng_remove_PrEP.bernoulli_filter(arr=eligible, prob=-n_change/n_eligible) + sim.people.hiv.on_prep[inds] = False # Add result sim.results.hiv.n_prep = np.count_nonzero(sim.people.alive & sim.people.hiv.on_prep) diff --git a/tests/run_multistream.py b/tests/run_multistream.py index 73c4d919..ccc502a2 100644 --- a/tests/run_multistream.py +++ b/tests/run_multistream.py @@ -5,26 +5,29 @@ # %% Imports and settings import os import stisim as ss -import matplotlib.pyplot as plt import sciris as sc import pandas as pd import seaborn as sns n = 1_000 # Agents n_rand_seeds = 250 -art_cov_levels = [0.025, 0.05, 0.10, 0.73] + [0] # Must include 0 as that's the baseline +intv_cov_levels = [0.025, 0.05, 0.10, 0.73] + [0] # Must include 0 as that's the baseline -figdir = 'figs' -#sc.path(figdir).mkdir(parents=True, exist_ok=True) +# Choose ART or PrEP +choice = 'ART' +intervention = {'ART': ss.hiv.ART, 'PrEP': ss.hiv.PrEP} -def run_sim(n, art_cov, rand_seed, multistream): +figdir = os.path.join(os.getcwd(), 'figs', choice) +sc.path(figdir).mkdir(parents=True, exist_ok=True) + +def run_sim(n, intv_cov, rand_seed, multistream): ppl = ss.People(n) net_pars = {'multistream': multistream} ppl.networks = ss.ndict(ss.simple_embedding(pars=net_pars), ss.maternal(pars=net_pars)) hiv_pars = { - 'beta': {'simple_embedding': [0.10, 0.08], 'maternal': [1,0]}, + 'beta': {'simple_embedding': [0.10, 0.08], 'maternal': [0.2,0]}, 'initial': 10, 'multistream': multistream, } @@ -33,16 +36,16 @@ def run_sim(n, art_cov, rand_seed, multistream): preg_pars = {'multistream': multistream} pregnancy = ss.Pregnancy(preg_pars) - art_pars = {'multistream': multistream} - art = ss.hiv.ART(t=[0, 1], coverage=[0, art_cov], pars=art_pars) + intv_pars = {'multistream': multistream} + intv = intervention[choice](t=[0, 1], coverage=[0, intv_cov], pars=intv_pars) pars = { 'start': 1980, 'end': 2010, - 'interventions': [art], + 'interventions': [intv], 'rand_seed': rand_seed, 'multistream': multistream, } - sim = ss.Sim(people=ppl, modules=[hiv, pregnancy], pars=pars, label=f'Sim with {n} agents and art_cov={art_cov}') + sim = ss.Sim(people=ppl, modules=[hiv, pregnancy], pars=pars, label=f'Sim with {n} agents and intv_cov={intv_cov}') sim.initialize() sim.run() @@ -53,7 +56,7 @@ def run_sim(n, art_cov, rand_seed, multistream): 'hiv.cum_deaths': sim.results.hiv.new_deaths.cumsum(), 'pregnancy.cum_births': sim.results.pregnancy.births.cumsum(), }) - df['art_cov'] = art_cov + df['intv_cov'] = intv_cov df['rand_seed'] = rand_seed df['multistream'] = multistream @@ -65,8 +68,8 @@ def run_scenarios(): for multistream in [True, False]: cfgs = [] for rs in range(n_rand_seeds): - for art_cov in art_cov_levels: - cfgs.append({'art_cov':art_cov, 'rand_seed':rs, 'multistream':multistream}) + for intv_cov in intv_cov_levels: + cfgs.append({'intv_cov':intv_cov, 'rand_seed':rs, 'multistream':multistream}) T = sc.tic() results += sc.parallelize(run_sim, kwargs={'n': n}, iterkwargs=cfgs, die=True) times[f'Multistream={multistream}'] = sc.toc(T, output=True) @@ -78,12 +81,12 @@ def run_scenarios(): return df def plot_scenarios(df): - d = pd.melt(df, id_vars=['ti', 'rand_seed', 'art_cov', 'multistream'], var_name='channel', value_name='Value') - d['baseline'] = d['art_cov']==0 + d = pd.melt(df, id_vars=['ti', 'rand_seed', 'intv_cov', 'multistream'], var_name='channel', value_name='Value') + d['baseline'] = d['intv_cov']==0 bl = d.loc[d['baseline']] scn = d.loc[~d['baseline']] - bl = bl.set_index(['ti', 'channel', 'rand_seed', 'art_cov', 'multistream'])[['Value']].reset_index('art_cov') - scn = scn.set_index(['ti', 'channel', 'rand_seed', 'art_cov', 'multistream'])[['Value']].reset_index('art_cov') + bl = bl.set_index(['ti', 'channel', 'rand_seed', 'intv_cov', 'multistream'])[['Value']].reset_index('intv_cov') + scn = scn.set_index(['ti', 'channel', 'rand_seed', 'intv_cov', 'multistream'])[['Value']].reset_index('intv_cov') mrg = scn.merge(bl, on=['ti', 'channel', 'rand_seed', 'multistream'], suffixes=('', '_ref')) mrg['Value - Reference'] = mrg['Value'] - mrg['Value_ref'] mrg = mrg.sort_index() @@ -91,7 +94,7 @@ def plot_scenarios(df): fkw = {'sharey': False, 'sharex': 'col', 'margin_titles': True} ## TIMESERIES - g = sns.relplot(kind='line', data=d, x='ti', y='Value', hue='art_cov', col='channel', row='multistream', + g = sns.relplot(kind='line', data=d, x='ti', y='Value', hue='intv_cov', col='channel', row='multistream', height=5, aspect=1.2, palette='Set1', errorbar='sd', lw=2, facet_kws=fkw) g.set_titles(col_template='{col_name}', row_template='Multistream: {row_name}') g.set_xlabels(r'$t_i$') @@ -99,9 +102,9 @@ def plot_scenarios(df): ## DIFF TIMESERIES for ms, mrg_by_ms in mrg.groupby('multistream'): - g = sns.relplot(kind='line', data=mrg_by_ms, x='ti', y='Value - Reference', hue='art_cov', col='channel', row='art_cov', + g = sns.relplot(kind='line', data=mrg_by_ms, x='ti', y='Value - Reference', hue='intv_cov', col='channel', row='intv_cov', height=3, aspect=1.0, palette='Set1', estimator=None, units='rand_seed', lw=0.5, facet_kws=fkw) #errorbar='sd', lw=2, - g.set_titles(col_template='{col_name}', row_template='ART cov: {row_name}') + g.set_titles(col_template='{col_name}', row_template='Coverage: {row_name}') g.fig.suptitle('Multistream' if ms else 'Centralized') g.fig.subplots_adjust(top=0.88) g.set_xlabels(r'Value - Reference at $t_i$') @@ -110,7 +113,7 @@ def plot_scenarios(df): ## FINAL TIME tf = df['ti'].max() mtf = mrg.loc[tf] - g = sns.displot(data=mtf.reset_index(), kind='kde', fill=True, rug=True, cut=0, hue='art_cov', x='Value - Reference', col='channel', row='multistream', + g = sns.displot(data=mtf.reset_index(), kind='kde', fill=True, rug=True, cut=0, hue='intv_cov', x='Value - Reference', col='channel', row='multistream', height=5, aspect=1.2, facet_kws=fkw, palette='Set1') g.set_titles(col_template='{col_name}', row_template='Multistream: {row_name}') g.set_xlabels(f'Value - Reference at $t_i={{{tf}}}$') diff --git a/tests/run_sweep.py b/tests/run_sweep.py new file mode 100644 index 00000000..c43494f2 --- /dev/null +++ b/tests/run_sweep.py @@ -0,0 +1,137 @@ +""" +Run a simple HIV simulation with random number coherence +""" + +# %% Imports and settings +import os +import stisim as ss +import sciris as sc +import pandas as pd +import seaborn as sns + +n = 1_000 # Agents +n_rand_seeds = 250 +x_beta_levels = [0.5, 0.95, 1.05, 2] + [1] # Must include 1 as that's the baseline + +figdir = os.path.join(os.getcwd(), 'figs', 'BetaSweep') +sc.path(figdir).mkdir(parents=True, exist_ok=True) + +def run_sim(n, x_beta, rand_seed, multistream): + ppl = ss.People(n) + + net_pars = {'multistream': multistream} + ppl.networks = ss.ndict(ss.simple_embedding(pars=net_pars), ss.maternal(pars=net_pars)) + + hiv_pars = { + 'beta': {'simple_embedding': [x_beta * 0.10, x_beta * 0.08], 'maternal': [x_beta * 0.2, 0]}, + 'initial': 10, + 'multistream': multistream, + } + hiv = ss.HIV(hiv_pars) + + preg_pars = {'multistream': multistream} + pregnancy = ss.Pregnancy(preg_pars) + + pars = { + 'start': 1980, + 'end': 2010, + 'rand_seed': rand_seed, + 'multistream': multistream, + } + sim = ss.Sim(people=ppl, modules=[hiv, pregnancy], pars=pars, label=f'Sim with {n} agents and x_beta={x_beta}') + sim.initialize() + sim.run() + + df = pd.DataFrame( { + 'ti': sim.tivec, + #'hiv.n_infected': sim.results.hiv.n_infected, + 'hiv.prevalence': sim.results.hiv.prevalence, + 'hiv.cum_deaths': sim.results.hiv.new_deaths.cumsum(), + 'pregnancy.cum_births': sim.results.pregnancy.births.cumsum(), + }) + df['x_beta'] = x_beta + df['rand_seed'] = rand_seed + df['multistream'] = multistream + + return df + +def run_scenarios(): + results = [] + times = {} + for multistream in [True, False]: + cfgs = [] + for rs in range(n_rand_seeds): + for x_beta in x_beta_levels: + cfgs.append({'x_beta':x_beta, 'rand_seed':rs, 'multistream':multistream}) + T = sc.tic() + results += sc.parallelize(run_sim, kwargs={'n': n}, iterkwargs=cfgs, die=True) + times[f'Multistream={multistream}'] = sc.toc(T, output=True) + + print('Timings:', times) + + df = pd.concat(results) + df.to_csv(os.path.join(figdir, 'result.csv')) + return df + + +def plot_scenarios(df): + d = pd.melt(df, id_vars=['ti', 'rand_seed', 'x_beta', 'multistream'], var_name='channel', value_name='Value') + d['baseline'] = d['x_beta']==1 + bl = d.loc[d['baseline']] + scn = d.loc[~d['baseline']] + bl = bl.set_index(['ti', 'channel', 'rand_seed', 'x_beta', 'multistream'])[['Value']].reset_index('x_beta') + scn = scn.set_index(['ti', 'channel', 'rand_seed', 'x_beta', 'multistream'])[['Value']].reset_index('x_beta') + mrg = scn.merge(bl, on=['ti', 'channel', 'rand_seed', 'multistream'], suffixes=('', '_ref')) + mrg['Value - Reference'] = mrg['Value'] - mrg['Value_ref'] + mrg = mrg.sort_index() + + fkw = {'sharey': False, 'sharex': 'col', 'margin_titles': True} + + ## TIMESERIES + g = sns.relplot(kind='line', data=d, x='ti', y='Value', hue='x_beta', col='channel', row='multistream', + height=5, aspect=1.2, palette='Set1', errorbar='sd', lw=2, facet_kws=fkw) + g.set_titles(col_template='{col_name}', row_template='Multistream: {row_name}') + g.set_xlabels(r'$t_i$') + g.fig.savefig(os.path.join(figdir, 'timeseries.png'), bbox_inches='tight', dpi=300) + + ## DIFF TIMESERIES + for ms, mrg_by_ms in mrg.groupby('multistream'): + g = sns.relplot(kind='line', data=mrg_by_ms, x='ti', y='Value - Reference', hue='x_beta', col='channel', row='x_beta', + height=3, aspect=1.0, palette='Set1', estimator=None, units='rand_seed', lw=0.5, facet_kws=fkw) #errorbar='sd', lw=2, + g.set_titles(col_template='{col_name}', row_template='Coverage: {row_name}') + g.fig.suptitle('Multistream' if ms else 'Centralized') + g.fig.subplots_adjust(top=0.88) + g.set_xlabels(r'Value - Reference at $t_i$') + g.fig.savefig(os.path.join(figdir, 'diff_multistream.png' if ms else 'diff_centralized.png'), bbox_inches='tight', dpi=300) + + ## FINAL TIME + tf = df['ti'].max() + mtf = mrg.loc[tf] + g = sns.displot(data=mtf.reset_index(), kind='kde', fill=True, rug=True, cut=0, hue='x_beta', x='Value - Reference', col='channel', row='multistream', + height=5, aspect=1.2, facet_kws=fkw, palette='Set1') + g.set_titles(col_template='{col_name}', row_template='Multistream: {row_name}') + g.set_xlabels(f'Value - Reference at $t_i={{{tf}}}$') + g.fig.savefig(os.path.join(figdir, 'final.png'), bbox_inches='tight', dpi=300) + + print('Figures saved to:', os.path.join(os.getcwd(), figdir)) + + return + + +if __name__ == '__main__': + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument('-p', '--plot', help='plot from a cached CSV file', type=str) + args = parser.parse_args() + + if args.plot: + print('Reading CSV file', args.plot) + df = pd.read_csv(args.plot, index_col=0) + else: + print('Running scenarios') + df = run_scenarios() + + plot_scenarios(df) + + print('Done') \ No newline at end of file From 31836c8637c9e4b7974722a2225680204fe9c6fc Mon Sep 17 00:00:00 2001 From: Daniel Klein Date: Thu, 28 Sep 2023 17:43:00 -0700 Subject: [PATCH 22/34] Updating comment. --- tests/run_multistream.py | 2 +- tests/run_sweep.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/run_multistream.py b/tests/run_multistream.py index ccc502a2..e449c1bb 100644 --- a/tests/run_multistream.py +++ b/tests/run_multistream.py @@ -1,5 +1,5 @@ """ -Run a simple HIV simulation with random number coherence +Exploring multisim (on vs off) for a model HIV system sweeping coverage of either ART or PrEP. """ # %% Imports and settings diff --git a/tests/run_sweep.py b/tests/run_sweep.py index c43494f2..991d79e7 100644 --- a/tests/run_sweep.py +++ b/tests/run_sweep.py @@ -1,5 +1,5 @@ """ -Run a simple HIV simulation with random number coherence +A simple beta sweep with multisim on and off using HIV as an example pathogen """ # %% Imports and settings From 96ab37bff611c0d90ddcad597f3d903c003e8257 Mon Sep 17 00:00:00 2001 From: Daniel Klein Date: Thu, 28 Sep 2023 18:28:42 -0700 Subject: [PATCH 23/34] Small improvements to sweeping --- tests/run_sweep.py | 25 +++++++++++++++++-------- 1 file changed, 17 insertions(+), 8 deletions(-) diff --git a/tests/run_sweep.py b/tests/run_sweep.py index 991d79e7..e1c79fbd 100644 --- a/tests/run_sweep.py +++ b/tests/run_sweep.py @@ -11,7 +11,7 @@ n = 1_000 # Agents n_rand_seeds = 250 -x_beta_levels = [0.5, 0.95, 1.05, 2] + [1] # Must include 1 as that's the baseline +x_beta_levels = [0.8, 0.92, 1.05, 1.2] + [1] # Must include 1 as that's the baseline | roughly np.logspace(np.log2(0.8), np.log2(1.2), 4, base=2) figdir = os.path.join(os.getcwd(), 'figs', 'BetaSweep') sc.path(figdir).mkdir(parents=True, exist_ok=True) @@ -55,7 +55,7 @@ def run_sim(n, x_beta, rand_seed, multistream): return df -def run_scenarios(): +def run_scenarios(figdir): results = [] times = {} for multistream in [True, False]: @@ -74,7 +74,7 @@ def run_scenarios(): return df -def plot_scenarios(df): +def plot_scenarios(df, figdir): d = pd.melt(df, id_vars=['ti', 'rand_seed', 'x_beta', 'multistream'], var_name='channel', value_name='Value') d['baseline'] = d['x_beta']==1 bl = d.loc[d['baseline']] @@ -98,7 +98,7 @@ def plot_scenarios(df): for ms, mrg_by_ms in mrg.groupby('multistream'): g = sns.relplot(kind='line', data=mrg_by_ms, x='ti', y='Value - Reference', hue='x_beta', col='channel', row='x_beta', height=3, aspect=1.0, palette='Set1', estimator=None, units='rand_seed', lw=0.5, facet_kws=fkw) #errorbar='sd', lw=2, - g.set_titles(col_template='{col_name}', row_template='Coverage: {row_name}') + g.set_titles(col_template='{col_name}', row_template='Beta: {row_name}') g.fig.suptitle('Multistream' if ms else 'Centralized') g.fig.subplots_adjust(top=0.88) g.set_xlabels(r'Value - Reference at $t_i$') @@ -113,6 +113,14 @@ def plot_scenarios(df): g.set_xlabels(f'Value - Reference at $t_i={{{tf}}}$') g.fig.savefig(os.path.join(figdir, 'final.png'), bbox_inches='tight', dpi=300) + ## FINAL TIME function of beta + dtf = d.set_index(['ti', 'rand_seed']).sort_index().loc[tf] + g = sns.relplot(kind='line', data=dtf.reset_index(), x='x_beta', y='Value', col='channel', row='multistream', + height=5, aspect=1.2, facet_kws=fkw, estimator=None, units='rand_seed', lw=0.25) + g.set_titles(col_template='{col_name}', row_template='Multistream: {row_name}') + g.set_ylabels(f'Value at $t_i={{{tf}}}$') + g.fig.savefig(os.path.join(figdir, 'final_beta.png'), bbox_inches='tight', dpi=300) + print('Figures saved to:', os.path.join(os.getcwd(), figdir)) return @@ -122,16 +130,17 @@ def plot_scenarios(df): import argparse parser = argparse.ArgumentParser() - parser.add_argument('-p', '--plot', help='plot from a cached CSV file', type=str) + parser.add_argument('-p', '--plot', help='Folder containing cached results.csv', type=str) args = parser.parse_args() if args.plot: print('Reading CSV file', args.plot) - df = pd.read_csv(args.plot, index_col=0) + df = pd.read_csv(os.path.join(args.plot, 'result.csv'), index_col=0) + figdir = args.plot else: print('Running scenarios') - df = run_scenarios() + df = run_scenarios(figdir) - plot_scenarios(df) + plot_scenarios(df, figdir) print('Done') \ No newline at end of file From 75431f87750814e7df6ff4ace0f310af72baf6cd Mon Sep 17 00:00:00 2001 From: Daniel Klein Date: Fri, 13 Oct 2023 18:04:15 -0700 Subject: [PATCH 24/34] Everything seems to be working, just verifying now. --- stisim/demographics.py | 14 ++-- stisim/distributions.py | 109 ++++++++++++++++++------------ stisim/gonorrhea.py | 5 +- stisim/hiv.py | 37 +++++++---- stisim/modules.py | 56 ++++++++++------ stisim/networks.py | 37 +++++------ stisim/parameters.py | 1 - stisim/people.py | 50 +++++++------- stisim/sim.py | 24 ++++--- stisim/states.py | 18 ++++- stisim/streams.py | 125 ++++++++++++++++++----------------- tests/run_multistream.py | 13 ++-- tests/run_stable_monogamy.py | 19 +++--- tests/run_sweep.py | 9 +-- tests/simple.py | 3 +- 15 files changed, 294 insertions(+), 226 deletions(-) diff --git a/stisim/demographics.py b/stisim/demographics.py index 1b753cf2..dabfcb25 100644 --- a/stisim/demographics.py +++ b/stisim/demographics.py @@ -241,8 +241,11 @@ def __init__(self, pars=None): 'initial': 3, # Number of women initially pregnant }, self.pars) - self.rng_sex = ss.Stream(self.multistream)('sex_at_birth') - self.rng_conception = ss.Stream(self.multistream)('conception') + 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}') return @@ -308,7 +311,7 @@ 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 = self.rng_conception.bernoulli_filter(arr=inds_to_choose_from, prob=self.pars.inci) + uids = self.rng_conception.bernoulli_filter(size=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) @@ -316,7 +319,8 @@ def make_pregnancies(self, sim): # Grow the arrays and set properties for the unborn agents new_uids = sim.people.grow(n_unborn_agents) sim.people.age[new_uids] = -self.pars.dur_pregnancy - sim.people.female[new_uids] = self.rng_sex.bernoulli(arr=uids, prob=0.5) # Replace 0.5 with sex ratio at birth + #sim.people.female[new_uids] = self.rng_sex.bernoulli(size=uids, prob=0.5) # Replace 0.5 with sex ratio at birth + sim.people.female[new_uids] = self.female_dist.sample(uids) # Add connections to any vertical transmission layers # Placeholder code to be moved / refactored. The maternal network may need to be @@ -346,7 +350,7 @@ def set_prognoses(self, sim, to_uids, from_uids=None): # Outcomes for pregnancies dur = np.full(len(to_uids), sim.ti + self.pars.dur_pregnancy / sim.dt) - dead = np.random.random(len(to_uids)) < self.pars.p_death + dead = self.rng_dead.bernoulli(size=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 diff --git a/stisim/distributions.py b/stisim/distributions.py index e2aecf0b..9b994e91 100644 --- a/stisim/distributions.py +++ b/stisim/distributions.py @@ -17,15 +17,18 @@ 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): - self.stream = np.random.defaut_rng() # Default to centralized random number generator + 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): @@ -41,10 +44,10 @@ def mean(self): """ raise NotImplementedError - def __call__(self, n=1, **kwargs): - return self.sample(n, **kwargs) + def __call__(self, size=1, **kwargs): + return self.sample(size, **kwargs) - def sample(cls, n=1, **kwargs): + def sample(cls, size=1, **kwargs): """ Return a specified number of samples from the distribution """ @@ -54,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=1): """ Sample using CDF """ bin_midpoints = self.bins[:-1] + np.diff(self.bins) / 2 cdf = np.cumsum(self.vals) cdf = cdf / cdf[-1] - values = self.stream.rand(n) + values = self.stream.rand(size) value_bins = np.searchsorted(cdf, values) return bin_midpoints[value_bins] @@ -76,15 +80,31 @@ 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 self.stream.uniform(low=self.low, high=self.high, size=n) + def sample(self, size=1): + return self.stream.uniform(low=self.low, high=self.high, size=size) + +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, size=1): + return self.stream.bernoulli(prob=self.p, size=size) class choice(Distribution): @@ -92,13 +112,14 @@ 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): - return self.stream.choice(a=self.choices, p=self.probabilities, replace=self.replace, size=n) + def sample(self, size): + return self.stream.choice(a=self.choices, p=self.probabilities, replace=self.replace, size=size) class normal(Distribution): @@ -106,12 +127,13 @@ 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 self.stream.normal(loc=self.mean, scale=self.std, size=n) + def sample(self, size=1): + return self.stream.normal(loc=self.mean, scale=self.std, size=size) class normal_pos(normal): @@ -121,8 +143,8 @@ 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, size=1): + return np.abs(super().sample(size)) class normal_int(Distribution): @@ -130,8 +152,8 @@ class normal_int(Distribution): Normal distribution returning only integer values """ - def sample(self, n=1): - return np.round(super().sample(n)) + def sample(self, size=1): + return np.round(super().sample(size)) class lognormal(Distribution): @@ -143,18 +165,19 @@ 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, size=1): if (sc.isnumber(self.mean) and self.mean > 0) or (sc.checktype(self.mean, 'arraylike') and (self.mean > 0).all()): - return self.stream.lognormal(mean=self.underlying_mean, sigma=self.underlying_std, size=n) + return self.stream.lognormal(mean=self.underlying_mean, sigma=self.underlying_std, size=size) else: - return np.zeros(n) + return np.zeros(size) class lognormal_int(lognormal): @@ -162,8 +185,8 @@ class lognormal_int(lognormal): Lognormal returning only integer values """ - def sample(self, n=1): - return np.round(super().sample(n)) + def sample(self, size=1): + return np.round(super().sample(size)) class poisson(Distribution): @@ -171,14 +194,15 @@ 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 self.stream.poisson(self.rate, n) + def sample(self, size=1): + return self.stream.poisson(self.rate, size) class neg_binomial(Distribution): @@ -193,21 +217,20 @@ 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, size=1): nbn_n = self.dispersion nbn_p = self.dispersion / (self.mean / self.step + self.dispersion) - return self.stream.negative_binomial(n=nbn_n, p=nbn_p, arr=n)# * self.step + return self.stream.negative_binomial(n=nbn_n, p=nbn_p, size=size) class beta(Distribution): @@ -215,15 +238,16 @@ 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 self.stream.beta(a=self.alpha, b=self.beta, size=n) + def sample(self, size=1): + return self.stream.beta(a=self.alpha, b=self.beta, size=size) class gamma(Distribution): @@ -231,12 +255,13 @@ 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 self.stream.gamma(shape=self.shape, scale=self.scale, size=n) + def sample(self, size=1): + return self.stream.gamma(shape=self.shape, scale=self.scale, size=size) \ No newline at end of file diff --git a/stisim/gonorrhea.py b/stisim/gonorrhea.py index 2fdd6239..e08de10d 100644 --- a/stisim/gonorrhea.py +++ b/stisim/gonorrhea.py @@ -20,9 +20,8 @@ def __init__(self, pars=None): self.ti_recovered = ss.State('ti_recovered', float, 0) self.ti_dead = ss.State('ti_dead', float, np.nan) # Death due to gonorrhea - self.rng_prog = ss.Stream(self.multistream)('prog_dur') - self.rng_dead = ss.Stream(self.multistream)('dead') - self.rng_dur_inf = ss.Stream(self.multistream)('dur_inf') + self.rng_dead = ss.Stream(f'dead_{self.name}') + self.rng_dur_inf = ss.Stream(f'dur_inf_{self.name}') self.pars = ss.omerge({ 'dur_inf': 3, # not modelling diagnosis or treatment explicitly here diff --git a/stisim/hiv.py b/stisim/hiv.py index c1879146..e7827ca5 100644 --- a/stisim/hiv.py +++ b/stisim/hiv.py @@ -16,12 +16,12 @@ def __init__(self, pars=None): self.susceptible = ss.State('susceptible', bool, True) self.infected = ss.State('infected', bool, False) - self.ti_infected = ss.State('ti_infected', float, 0) + self.ti_infected = ss.State('ti_infected', float, ss.INT_NAN) self.on_art = ss.State('on_art', bool, False) self.on_prep = ss.State('on_prep', bool, False) self.cd4 = ss.State('cd4', float, 500) - self.rng_dead = ss.Stream(self.multistream)('dead') + self.rng_dead = ss.Stream(f'dead_{self.name}') self.pars = ss.omerge({ 'cd4_min': 100, @@ -41,9 +41,9 @@ def update_states(self, sim): self.rel_sus[sim.people.alive & ~self.infected & self.on_prep] = 0.04 self.rel_trans[sim.people.alive & self.infected & self.on_art] = 0.04 - hiv_death_prob = 0.1 / (self.pars.cd4_min - self.pars.cd4_max)**2 * (self.cd4 - self.pars.cd4_max)**2 + hiv_death_prob = 0.2 / (self.pars.cd4_min - self.pars.cd4_max)**2 * (self.cd4 - self.pars.cd4_max)**2 can_die = ss.true(sim.people.alive & sim.people.hiv.infected) - hiv_deaths = self.rng_dead.bernoulli_filter(arr=can_die, prob=hiv_death_prob[can_die]) + hiv_deaths = self.rng_dead.bernoulli_filter(size=can_die, prob=hiv_death_prob[can_die]) sim.people.alive[hiv_deaths] = False sim.people.ti_dead[hiv_deaths] = sim.ti self.results['new_deaths'][sim.ti] = len(hiv_deaths) @@ -80,8 +80,8 @@ def __init__(self, t: np.array, coverage: np.array, **kwargs): super().__init__(**kwargs) - self.rng_add_ART = ss.Stream(self.multistream)('add_ART') - self.rng_remove_ART = ss.Stream(self.multistream)('remove_ART') + self.rng_add_ART = ss.Stream('add_ART') + self.rng_remove_ART = ss.Stream('remove_ART') return @@ -95,6 +95,16 @@ def apply(self, sim): return coverage = self.coverage[np.where(self.t <= sim.ti)[0][-1]] + recently_infected = ss.true((sim.people.hiv.ti_infected == sim.ti-2) & sim.people.alive) # 2 year (step) delay + inds = self.rng_add_ART.bernoulli_filter(size=recently_infected, prob=coverage) + sim.people.hiv.on_art[inds] = True + + # Add result + sim.results.hiv.n_art = np.count_nonzero(sim.people.alive & sim.people.hiv.on_art) + + return len(inds) + + ''' on_art = sim.people.alive & sim.people.hiv.on_art infected = sim.people.alive & sim.people.hiv.infected n_change = np.round(coverage * infected.sum() - on_art.sum()).astype(int) @@ -103,19 +113,20 @@ def apply(self, sim): eligible = ss.true(sim.people.alive & sim.people.hiv.infected & ~sim.people.hiv.on_art) n_eligible = len(eligible) #np.count_nonzero(eligible) if n_eligible: - inds = self.rng_add_ART.bernoulli_filter(arr=eligible, prob=min(n_eligible, n_change)/n_eligible) + inds = self.rng_add_ART.bernoulli_filter(size=eligible, prob=min(n_eligible, n_change)/n_eligible) sim.people.hiv.on_art[inds] = True elif n_change < 0: # Take some people off ART eligible = sim.people.alive & sim.people.hiv.infected & sim.people.hiv.on_art n_eligible = np.count_nonzero(eligible) - inds = self.rng_remove_ART.bernoulli_filter(arr=eligible, prob=-n_change/n_eligible) + inds = self.rng_remove_ART.bernoulli_filter(size=eligible, prob=-n_change/n_eligible) sim.people.hiv.on_art[inds] = False # Add result sim.results.hiv.n_art = np.count_nonzero(sim.people.alive & sim.people.hiv.on_art) return + ''' class PrEP(ss.Intervention): @@ -126,8 +137,8 @@ def __init__(self, t: np.array, coverage: np.array, **kwargs): super().__init__(**kwargs) - self.rng_add_PrEP = ss.Stream(self.multistream)('add_PrEP') - self.rng_remove_PrEP = ss.Stream(self.multistream)('remove_PrEP') + self.rng_add_PrEP = ss.Stream('add_PrEP') + self.rng_remove_PrEP = ss.Stream('remove_PrEP') return @@ -149,14 +160,14 @@ def apply(self, sim): eligible = ss.true(sim.people.alive & ~sim.people.hiv.infected & ~sim.people.hiv.on_prep) n_eligible = len(eligible) if n_eligible: - inds = self.rng_add_PrEP.bernoulli_filter(arr=eligible, prob=min(n_eligible, n_change)/n_eligible) + inds = self.rng_add_PrEP.bernoulli_filter(size=eligible, prob=min(n_eligible, n_change)/n_eligible) sim.people.hiv.on_prep[inds] = True elif n_change < 0: # Take some people off PrEP eligible = ss.true(sim.people.alive & ~sim.people.hiv.infected & sim.people.hiv.on_prep) n_eligible = len(eligible) if n_eligible: - inds = self.rng_remove_PrEP.bernoulli_filter(arr=eligible, prob=-n_change/n_eligible) + inds = self.rng_remove_PrEP.bernoulli_filter(size=eligible, prob=-n_change/n_eligible) sim.people.hiv.on_prep[inds] = False # Add result @@ -179,4 +190,4 @@ def initialize(self, sim): self.cd4 = np.zeros((sim.npts, sim.people.n), dtype=int) def apply(self, sim): - self.cd4[sim.t] = sim.people.hiv.cd4 + self.cd4[sim.t] = sim.people.hiv.cd4 \ No newline at end of file diff --git a/stisim/modules.py b/stisim/modules.py index 9d4fbe60..facc8452 100644 --- a/stisim/modules.py +++ b/stisim/modules.py @@ -13,21 +13,17 @@ class Module(sc.prettyobj): def __init__(self, pars=None, label=None, requires=None, *args, **kwargs): - default_pars = {'multistream': True} - self.pars = sc.mergedicts(default_pars, pars) - + self.pars = pars self.label = label if label else '' self.requires = sc.mergelists(requires) self.results = ss.ndict(type=ss.Result) self.initialized = False self.finalized = False - self.multistream = self.pars['multistream'] - # Random number streams - self.rng_init_cases = ss.Stream(self.multistream)(f'initial_cases_{self.name}') - self.rng_trans = ss.Stream(self.multistream)(f'trans_{self.name}') - self.rng_choose_infector = ss.Stream(self.multistream)(f'choose_infector_{self.name}') + self.rng_init_cases = ss.Stream(f'initial_cases_{self.name}') + self.rng_trans = ss.Stream(f'trans_{self.name}') + self.rng_choose_infector = ss.Stream(f'choose_infector_{self.name}') return @@ -52,7 +48,7 @@ def initialize(self, sim): state.initialize(sim.people) # Connect the streams to the sim - for stream in self.streams.values(): + for stream in self.streams: if not stream.initialized: stream.initialize(sim.streams) @@ -82,6 +78,15 @@ def states(self): """ return [x for x in self.__dict__.values() if isinstance(x, ss.State)] + @property + def streams(self): + """ + Return a flat collection of all streams, as with states above + + :return: + """ + return [x for x in self.__dict__.values() if isinstance(x, (ss.MultiStream, ss.CentralizedStream))] + class Disease(Module): """ Base module contains states/attributes that all modules have """ @@ -126,8 +131,8 @@ def set_initial_states(self, sim): #initial_cases = np.random.choice(sim.people.uid, self.pars['initial']) #rng = sim.rngs.get(f'initial_cases_{self.name}') - #initial_cases = ss.binomial_filter_rng(prob=self.pars['initial']/len(sim.people), arr=sim.people.uid, rng=rng, block_size=len(sim.people._uid_map)) - initial_cases = self.rng_init_cases.bernoulli_filter(arr=sim.people.uid, prob=self.pars['initial']/len(sim.people)) + #initial_cases = ss.binomial_filter_rng(prob=self.pars['initial']/len(sim.people), size=sim.people.uid, rng=rng, block_size=len(sim.people._uid_map)) + initial_cases = self.rng_init_cases.bernoulli_filter(size=sim.people.uid, prob=self.pars['initial']/len(sim.people)) self.set_prognoses(sim, initial_cases, from_uids=None) # TODO: sentinel value to indicate seeds? return @@ -159,11 +164,13 @@ def make_new_cases(self, sim): # Probability of each node acquiring a case # TODO: Just people that who are alive? - p_acq_node = np.zeros( len(sim.people._uid_map) ) - node_from_node = np.ones( (sim.people._uid_map.n, sim.people._uid_map.n) ) + n = len(sim.people.uid) + p_acq_node = np.zeros( n ) + node_from_node = np.ones( (n,n) ) for lkey, layer in sim.people.networks.items(): if lkey in pars['beta']: + # TODO: Likely no longer need alive here, at least not if dead people are removed rel_trans = self.rel_trans * (self.infected & sim.people.alive) rel_sus = self.rel_sus * (self.susceptible & sim.people.alive) @@ -175,17 +182,21 @@ def make_new_cases(self, sim): # Check for new transmission from a --> b # TODO: Will need to be more efficient here - can maintain edge to node matrix - node_from_edge = np.ones( (sim.people._uid_map.n, len(a)) ) - node_from_edge[b, np.arange(len(a))] = 1 - rel_trans[a] * rel_sus[b] * layer['beta'] * beta + node_from_edge = np.ones( (n, len(a)) ) + ai = sim.people._uid_map[a] # Indices of a and b (rather than uid) + bi = sim.people._uid_map[b] + p_not_acq = 1 - rel_trans[a] * rel_sus[b] * layer['beta'] * beta # Needs DT + + node_from_edge[bi, np.arange(len(a))] = p_not_acq p_not_acq_by_node_this_layer_b_from_a = node_from_edge.prod(axis=1) # (1-p1)*(1-p2)*... p_acq_node = 1 - (1-p_acq_node) * p_not_acq_by_node_this_layer_b_from_a - node_from_node_this_layer_b_from_a = np.ones( (sim.people._uid_map.n, sim.people._uid_map.n) ) - node_from_node_this_layer_b_from_a[b, a] = 1 - rel_trans[a] * rel_sus[b] * layer['beta'] * beta + node_from_node_this_layer_b_from_a = np.ones( (n,n) ) + node_from_node_this_layer_b_from_a[bi, ai] = p_not_acq node_from_node *= node_from_node_this_layer_b_from_a - new_cases_bool = self.rng_trans.bernoulli(sim.people._uid_map._view, p_acq_node) - new_cases = sim.people._uid_map._view[new_cases_bool] + new_cases_bool = self.rng_trans.bernoulli(size=sim.people.uid, prob=p_acq_node) + new_cases = sim.people.uid[new_cases_bool] if not len(new_cases): return 0 @@ -193,10 +204,11 @@ def make_new_cases(self, sim): # Decide whom the infection came from using one random number for each b (aligned by block size) frm = np.zeros_like(new_cases) r = self.rng_choose_infector.random(new_cases) - prob = (1-node_from_node[new_cases]) + new_cases_idx = new_cases_bool.nonzero()[0] + prob = (1-node_from_node[new_cases_idx]) # Prob of acquiring from each node | can constrain to just neighbors? cumsum = (prob / ((prob.sum(axis=1)[:,np.newaxis]))).cumsum(axis=1) - frm = np.argmax( cumsum >= r[:,np.newaxis], axis=1) - #frm = a[frm_idx] + frm_idx = np.argmax( cumsum >= r[:,np.newaxis], axis=1) + frm = sim.people.uid[frm_idx] self.set_prognoses(sim, new_cases, frm) return len(new_cases) # number of new cases made diff --git a/stisim/networks.py b/stisim/networks.py index fa43a71b..d2c200ed 100644 --- a/stisim/networks.py +++ b/stisim/networks.py @@ -60,17 +60,12 @@ def __init__(self, *args, pars=None, key_dict=None, vertical=False, label=None, 'beta': ss.float_, } - default_pars = {'multistream': True} - self.pars = sc.mergedicts(default_pars, pars) - self.meta = sc.mergedicts(default_keys, key_dict) self.vertical = vertical # Whether transmission is bidirectional self.basekey = 'p1' # Assign a base key for calculating lengths and performing other operations self.label = label self.initialized = False - self.multistream = self.pars['multistream'] - # Handle args kwargs = sc.mergedicts(*args, kwargs) @@ -276,9 +271,13 @@ def remove_uids(self, uids): is specifically used when removing agents from the simulation. """ + if len(uids) == 0: + return + keep = ~(np.isin(self.p1, uids) | np.isin(self.p2, uids)) for k in self.meta_keys(): self[k] = self[k][keep] + return class Networks(ss.ndict): @@ -306,9 +305,9 @@ def __init__(self, mean_dur=5, **kwargs): self.mean_dur = mean_dur # Define random streams - self.rng_pair_12 = ss.Stream(self.multistream)('pair_12') - self.rng_pair_21 = ss.Stream(self.multistream)('pair_21') - self.rng_mean_dur = ss.Stream(self.multistream)('mean_dur') + self.rng_pair_12 = ss.Stream('pair_12') + self.rng_pair_21 = ss.Stream('pair_21') + self.rng_mean_dur = ss.Stream('mean_dur') return @@ -332,10 +331,10 @@ def add_pairs(self, people, ti=None): if len(available_m) <= len(available_f): p1 = available_m - p2 = self.rng_pair_12.choice(available_f, size=len(p1), replace=False) # TODO: Stream-ify + p2 = self.rng_pair_12.choice(a=available_f, size=len(p1), replace=False) # TODO: Stream-ify else: p2 = available_f - p1 = self.rng_pair_21.choice(available_m, size=len(p2), replace=False) # TODO: Stream-ify + p1 = self.rng_pair_21.choice(a=available_m, size=len(p2), replace=False) # TODO: Stream-ify beta = np.ones_like(p1) dur = self.rng_mean_dur.poisson(p1, self.mean_dur) # TODO: Stream-ify @@ -368,28 +367,24 @@ def add_pairs(self, people, ti=None): # Find unpartnered males and females - could in principle check other contact layers too # by having the People object passed in here - available_m = np.setdiff1d(ss.true(~people.female), self.members) - available_f = np.setdiff1d(ss.true(people.female), self.members) - - # slow: - available_m = ss.true(people.alive[available_m]) - available_f = ss.true(people.alive[available_f]) + available_m = np.setdiff1d(ss.true(~people.female & people.alive), self.members) + available_f = np.setdiff1d(ss.true(people.female & people.alive), self.members) if not len(available_m) or not len(available_f): if ss.options.verbose > 1: print('No pairs to add') return 0 - loc_m = people.age[available_m].values - 5 + self.rng_pair_12.normal(arr=available_m, std=3) - loc_f = people.age[available_f].values + self.rng_pair_21.normal(arr=available_f, std=3) + loc_m = people.age[available_m].values - 5 + self.rng_pair_12.normal(size=available_m, std=3) + loc_f = people.age[available_f].values + self.rng_pair_21.normal(size=available_f, std=3) dist_mat = sps.distance_matrix(loc_m[:, np.newaxis], loc_f[:, np.newaxis]) ind_m, ind_f = spo.linear_sum_assignment(dist_mat) # loc_f[ind_f[0]] is close to loc_m[ind_m[0]] n_pairs = len(ind_f) - self['p1'] = np.concatenate([self['p1'], ind_m]) - self['p2'] = np.concatenate([self['p2'], ind_f]) + self['p1'] = np.concatenate([self['p1'], available_m[ind_m]]) + self['p2'] = np.concatenate([self['p2'], available_f[ind_f]]) beta = np.ones(n_pairs) dur = self.rng_mean_dur.poisson(ind_m, self.mean_dur) @@ -414,7 +409,7 @@ def initialize(self, sim): self['p1'] = np.arange(0,n,2) # EVEN self['p2'] = np.arange(1,n,2) # ODD self['beta'] = np.ones(len(self['p1'])) - self['dur'] = np.iinfo(int).max + self['dur'] = np.full(len(self['p1']), fill_value=np.iinfo(int).max, dtype=int) return def update(self, people, dt=None): diff --git a/stisim/parameters.py b/stisim/parameters.py index 229395a8..6b7f37bd 100644 --- a/stisim/parameters.py +++ b/stisim/parameters.py @@ -46,7 +46,6 @@ def __init__(self, **kwargs): self.dt = 1.0 # Timestep (in years) self.dt_demog = 1.0 # Timestep for demographic updates (in years) self.rand_seed = 1 # Random seed, if None, don't reset - self.multistream = True # Run with multiple random number streams as opposed to one centralized random number generator self.verbose = ss.options.verbose # Whether or not to display information during the run -- options are 0 (silent), 0.1 (some; default), 1 (default), 2 (everything) # Events and interventions diff --git a/stisim/people.py b/stisim/people.py index 3601a263..64efe9e7 100644 --- a/stisim/people.py +++ b/stisim/people.py @@ -77,6 +77,11 @@ def remove(self, uids_to_remove): :param uids_to_remove: An int/list/array containing the UID(s) to remove """ + + # Shortcut exit if nothing to do + if len(uids_to_remove) == 0: + return + # Calculate the *indices* to keep keep_uids = self.uid[~np.in1d(self.uid, uids_to_remove)] # Calculate UIDs to keep keep_inds = self._uid_map[keep_uids] # Calculate indices to keep @@ -140,7 +145,7 @@ class People(BasePeople): # %% Basic methods - def __init__(self, n, age_data=None, extra_states=None, networks=None): + def __init__(self, n, age_data=None, extra_states=None, networks=None, rand_seed=0): """ Initialize """ super().__init__(n) @@ -149,45 +154,44 @@ def __init__(self, n, age_data=None, extra_states=None, networks=None): self.version = ss.__version__ # Store version info self.rng_female = ss.Stream('female') - states = [ ss.State('age', float, 0), - ss.State('female', bool, ss.choice([True, False], self.rng_female)), + ss.State('female', bool, ss.bernoulli(0.5, rng=self.rng_female)), ss.State('debut', float), ss.State('alive', bool, True), ss.State('ti_dead', int, ss.INT_NAN), # Time index for death ss.State('scale', float, 1.0), ] states.extend(sc.promotetolist(extra_states)) - - self.states = ss.ndict() - self._initialize_states(states) + self.states = ss.ndict(states) self.networks = ss.ndict(networks) # Set initial age distribution - likely move this somewhere else later - age_data_dist = self.get_age_dist(age_data) - self.age[:] = age_data_dist.sample(len(self)) + self.rng_agedist = ss.Stream('agedist') + self.age_data_dist = self.get_age_dist(age_data, self.rng_agedist) return @staticmethod - def get_age_dist(age_data): + def get_age_dist(age_data, rng): """ Return an age distribution based on provided data """ - if age_data is None: return ss.uniform(0, 100) + if age_data is None: return ss.uniform(0, 100, rng=rng) if sc.checktype(age_data, pd.DataFrame): - return ss.from_data(vals=age_data['value'].values, bins=age_data['age'].values) + return ss.from_data(vals=age_data['value'].values, bins=age_data['age'].values, rng=rng) - def initialize(self): - """ Initialization - TBC what needs to go here """ - self.initialized = True - return + def initialize(self, sim): + """ Initialization """ + self.rng_female.initialize(sim.streams) + self.rng_agedist.initialize(sim.streams) - def _initialize_states(self, states): - for state in states: + for name, state in self.states.items(): self.add_state(state) # Register the state internally for dynamic growth - self.states.append(state) # Expose these states with their original names + #self.states.append(state) # Expose these states with their original names state.initialize(self) # Connect the state to this people instance - setattr(self, state.name, state) + setattr(self, name, state) + + self.age[:] = self.age_data_dist.sample(len(self)) + self.initialized = True return def add_module(self, module, force=False): @@ -229,8 +233,6 @@ def remove_dead(self, sim): """ uids_to_remove = ss.true(self.dead) self.remove(uids_to_remove) - for network in self.networks.values(): - network.remove_uids(uids_to_remove) return def update_post(self, sim): @@ -241,7 +243,7 @@ def update_post(self, sim): :return: """ self.age[self.alive] += self.dt - + return def resolve_deaths(self): """ @@ -251,6 +253,7 @@ def resolve_deaths(self): """ death_uids = ss.true(self.ti_dead <= self.ti) self.alive[death_uids] = False + return def update_networks(self): """ @@ -258,6 +261,7 @@ def update_networks(self): """ for network in self.networks.values(): network.update(self) + return @property def active(self): @@ -296,7 +300,7 @@ def request_death(self, uids): track of that internally. When the module is ready to cause the agent to die, it should call this method, and can update its own results for the cause of death. This way, if multiple modules request death on the same day, they can each record a death due to their - own cause., + own cause. The actual deaths are resolved after modules have all run, but before analyzers. That way, regardless of whether removing dead agents is enabled or not, analyzers will be able to diff --git a/stisim/sim.py b/stisim/sim.py index 12a555bd..01d2f894 100644 --- a/stisim/sim.py +++ b/stisim/sim.py @@ -207,7 +207,7 @@ def init_people(self, popdict=None, reset=False, verbose=None, **kwargs): # Any other initialization if not self.people.initialized: - self.people.initialize() + self.people.initialize(self) # Set time attributes self.people.ti = self.ti @@ -216,9 +216,10 @@ def init_people(self, popdict=None, reset=False, verbose=None, **kwargs): return self def init_demographics(self): - for module in self.demographics.values(): - module.initialize(self) - self.results[module.name] = module.results + for demog in self.demographics.values(): + demog.initialize(self) + self.results[demog.name] = demog.results + return def init_diseases(self): """ Initialize modules and connectors to be simulated """ @@ -271,7 +272,7 @@ def init_interventions(self): errormsg = f'Intervention {intervention} does not seem to be a valid intervention: must be a function or Intervention subclass' raise TypeError(errormsg) - for stream in intervention.streams.values(): + for stream in intervention.streams: stream.initialize(self.streams) return @@ -320,13 +321,13 @@ def step(self): self.people.remove_dead(self) # Update demographic modules (create new agents from births/immigration, schedule non-disease deaths and emigration) - for module in self.demographics.values(): - module.update(self) + for demog in self.demographics.values(): + demog.update(self) # Carry out autonomous state changes in the disease modules. This allows autonomous state changes/initializations # to be applied to newly created agents for disease in self.diseases.values(): - disease.update_pre(self) + disease.update_states(self) # Update networks - this takes place here in case autonomous state changes at this timestep # affect eligibility for contacts @@ -426,8 +427,11 @@ def finalize(self, verbose=None): # otherwise the scale factor will be applied multiple times raise AlreadyRunError('Simulation has already been finalized') - for module in self.modules.values(): - module.finalize(self) + for demog in self.demographics.values(): + demog.finalize(self) + + for disease in self.diseases.values(): + disease.finalize(self) for intervention in self.interventions.values(): intervention.finalize(self) diff --git a/stisim/states.py b/stisim/states.py index 7ee2eba2..2ee99287 100644 --- a/stisim/states.py +++ b/stisim/states.py @@ -267,6 +267,9 @@ def __repr__(self): return self._view.__repr__() def _new_items(self, n): + # Fill with empty values + return np.empty(n, dtype=self.dtype) + ''' # Create new arrays of the correct dtype and fill them based on the (optionally callable) fill value if callable(self.fill_value): new = np.empty(n, dtype=self.dtype) @@ -274,6 +277,16 @@ def _new_items(self, n): else: new = np.full(n, dtype=self.dtype, fill_value=self.fill_value) return new + ''' + + def _fill_data(self, n0, n1): + # Fill data + n = n1-n0 + if callable(self.fill_value): + self._data[n0:n1] = self.fill_value(n) + else: + self._data[n0:n1] = self.fill_value + return def initialize(self, n): self._data = self._new_items(n) @@ -294,7 +307,7 @@ def _trim(self, inds): # Note that these are indices, not UIDs! n = len(inds) self._data[:n] = self._data[inds] - self._data[n:self.n] = self.fill_value(self.n-n) if callable(self.fill_value) else self.fill_value + # DJK MOVING TO MAP: self._data[n:self.n] = self.fill_value(self.n-n) if callable(self.fill_value) else self.fill_value self.n = n self._map_arrays() @@ -305,6 +318,9 @@ def _map_arrays(self): This method should be called whenever the number of agents required changes (regardless of whether or not the underlying arrays have been resized) """ + n0 = len(self._view) if self._view is not None else 0 + if self.n > n0: + self._fill_data(n0, self.n) # Fill the new data self._view = self._data[:self.n] def __getitem__(self, key): diff --git a/stisim/streams.py b/stisim/streams.py index e528bd5c..6440ede3 100644 --- a/stisim/streams.py +++ b/stisim/streams.py @@ -1,5 +1,5 @@ +import hashlib import numpy as np -import sciris as sc import stisim as ss __all__ = ['Streams', 'MultiStream', 'CentralizedStream', 'Stream'] @@ -35,7 +35,7 @@ def add(self, stream): raise RepeatNameException(f'A Stream with name {stream.name} has already been added.') if stream.seed_offset is None: - seed = abs(hash(stream.name)) + seed = int(hashlib.sha256(stream.name.encode('utf-8')).hexdigest(), 16) % 10**8 #abs(hash(stream.name)) elif stream.seed_offset in self.used_seeds: raise SeedRepeatException(f'Requested seed offset {stream.seed_offset} for stream {stream} has already been used.') else: @@ -96,15 +96,18 @@ def _pre_draw(func): def check_ready(self, *args, **kwargs): """ Validation before drawing """ - # Check for zero length arr - if 'arr' in kwargs.keys(): - arr = kwargs['arr'] + # Check for zero length size + if 'size' in kwargs.keys(): + size = kwargs['size'] else: - arr = args[0] - if isinstance(arr, int): + size = args[0] + if isinstance(size, int): # If an integer, the user wants "n" samples - kwargs['arr'] = np.arange(arr) - elif len(arr) == 0: + if size == 0: + return np.array([], dtype=int) # int dtype allows use as index, e.g. bernoulli_filter + else: + kwargs['size'] = np.arange(size) + elif len(size) == 0: return np.array([], dtype=int) # int dtype allows use as index, e.g. bernoulli_filter if not self.initialized: @@ -185,52 +188,52 @@ def step(self, ti): return - def draw_size(self, arr): - """ Determine how many random numbers to draw for a given arr """ + def draw_size(self, size): + """ Determine how many random numbers to draw for a given size """ - if isinstance(arr, ss.states.FusedArray): - v = arr.values - elif isinstance(arr, ss.states.DynamicView): - v = arr._view + if isinstance(size, ss.states.FusedArray): + v = size.values + elif isinstance(size, ss.states.DynamicView): + v = size._view else: - v = arr + v = size if v.dtype == bool: - return len(arr) + return len(size) return v.max()+1 @_pre_draw - def random(self, arr): - return super(MultiStream, self).random(self.draw_size(arr))[arr] + def random(self, size): + return super(MultiStream, self).random(size=self.draw_size(size))[size] @_pre_draw - def uniform(self, arr, **kwargs): - return super(MultiStream, self).uniform(self.draw_size(arr), **kwargs)[arr] + def uniform(self, size, **kwargs): + return super(MultiStream, self).uniform(size=self.draw_size(size), **kwargs)[size] @_pre_draw - def poisson(self, arr, lam): - return super(MultiStream, self).poisson(lam=lam, size=self.draw_size(arr))[arr] + def poisson(self, size, lam): + return super(MultiStream, self).poisson(size=self.draw_size(size), lam=lam)[size] @_pre_draw - def normal(self, arr, mu=0, std=1): - return mu + std*super(MultiStream, self).normal(size=self.draw_size(arr))[arr] + def normal(self, size, mu=0, std=1): + return mu + std*super(MultiStream, self).normal(size=self.draw_size(size))[size] @_pre_draw - def negative_binomial(self, arr, **kwargs): #n=nbn_n, p=nbn_p, size=n) - return super(MultiStream, self).negative_binomial(**kwargs, size=self.draw_size(arr))[arr] + def negative_binomial(self, size, **kwargs): #n=nbn_n, p=nbn_p, size=n) + return super(MultiStream, self).negative_binomial(size=self.draw_size(size), **kwargs)[size] @_pre_draw - def bernoulli(self, arr, prob): - #return super(MultiStream, self).choice([True, False], size=arr.max()+1) # very slow - #return (super(MultiStream, self).binomial(n=1, p=prob, size=arr.max()+1))[arr].astype(bool) # pretty fast - return super(MultiStream, self).random(self.draw_size(arr))[arr] < prob # fastest + def bernoulli(self, size, prob): + #return super(MultiStream, self).choice([True, False], size=size.max()+1) # very slow + #return (super(MultiStream, self).binomial(n=1, p=prob, size=size.max()+1))[size].astype(bool) # pretty fast + return super(MultiStream, self).random(size=self.draw_size(size))[size] < prob # fastest # @_pre_draw <-- handled by call to self.bernoullli - def bernoulli_filter(self, arr, prob): - return arr[self.bernoulli(arr, prob)] # Slightly faster on my machine for bernoulli to typecast + def bernoulli_filter(self, size, prob): + return size[self.bernoulli(size, prob)] # Slightly faster on my machine for bernoulli to typecast - def choice(self, arr, prob): + def choice(self, size, a, **kwargs): # Consider raising a warning instead? raise NotStreamSafeException('The "choice" function is not MultiStream-safe.') @@ -270,43 +273,43 @@ def reset(self): def step(self, ti): pass - def draw_size(self, arr): - """ Determine how many random numbers to draw for a given arr """ + def draw_size(self, size): + """ Determine how many random numbers to draw for a given size """ - if isinstance(arr, int): - return arr - elif isinstance(arr, ss.states.FusedArray): - v = arr.values - elif isinstance(arr, ss.states.DynamicView): - v = arr._view + if isinstance(size, int): + return size + elif isinstance(size, ss.states.FusedArray): + v = size.values + elif isinstance(size, ss.states.DynamicView): + v = size._view else: - v = arr + v = size if v.dtype == bool: - return arr.sum() + return size.sum() - return len(arr) + return len(size) - def random(self, arr): - return np.random.random(self.draw_size(arr)) + def random(self, size): + return np.random.random(self.draw_size(size)) - def uniform(self, arr, **kwargs): - return np.random.uniform(self.draw_size(arr), **kwargs) + def uniform(self, size, **kwargs): + return np.random.uniform(size=self.draw_size(size), **kwargs) - def poisson(self, arr, lam): - return np.random.poisson(lam=lam, size=self.draw_size(arr)) + def poisson(self, size, lam): + return np.random.poisson(lam=lam, size=self.draw_size(size)) - def normal(self, arr, mu=0, std=1): - return mu + std*np.random.normal(size=self.draw_size(arr), loc=mu, scale=std) + def normal(self, size, mu=0, std=1): + return mu + std*np.random.normal(size=self.draw_size(size), loc=mu, scale=std) - def negative_binomial(self, arr, **kwargs): #n=nbn_n, p=nbn_p, size=n) - return np.random.negative_binomial(**kwargs, size=self.draw_size(arr)) + def negative_binomial(self, size, **kwargs): #n=nbn_n, p=nbn_p, size=n) + return np.random.negative_binomial(**kwargs, size=self.draw_size(size)) - def bernoulli(self, arr, prob): - return np.random.random(self.draw_size(arr)) < prob + def bernoulli(self, size, prob): + return np.random.random(self.draw_size(size)) < prob - def bernoulli_filter(self, arr, prob): - return arr[self.bernoulli(arr, prob)] # Slightly faster on my machine for bernoulli to typecast + def bernoulli_filter(self, size, prob): + return size[self.bernoulli(size, prob)] # Slightly faster on my machine for bernoulli to typecast - def choice(self, arr, **kwargs): - return self.stream.choice(size=self.draw_size(arr), **kwargs) \ No newline at end of file + def choice(self, size, a, **kwargs): + return np.random.choice(a, size=self.draw_size(size), **kwargs) \ No newline at end of file diff --git a/tests/run_multistream.py b/tests/run_multistream.py index f4e7c468..8b63551e 100644 --- a/tests/run_multistream.py +++ b/tests/run_multistream.py @@ -10,7 +10,7 @@ import seaborn as sns n = 1_000 # Agents -n_rand_seeds = 250 +n_rand_seeds = 100 intv_cov_levels = [0.025, 0.05, 0.10, 0.73] + [0] # Must include 0 as that's the baseline # Choose ART or PrEP @@ -28,7 +28,8 @@ def run_sim(n, intv_cov, rand_seed, multistream): ppl.networks = ss.ndict(ss.simple_embedding(), ss.maternal()) hiv_pars = { - 'beta': {'simple_embedding': [0.10, 0.08], 'maternal': [0.2,0]}, + #'beta': {'simple_embedding': [0.10, 0.08], 'maternal': [0.2,0]}, + 'beta': {'simple_embedding': [0.40, 0.35], 'maternal': [0.2,0]}, 'initial': 10, } hiv = ss.HIV(hiv_pars) @@ -38,11 +39,13 @@ def run_sim(n, intv_cov, rand_seed, multistream): intv = intervention[choice](t=[0, 1], coverage=[0, intv_cov]) pars = { 'start': 1980, - 'end': 2010, + 'end': 2020, 'interventions': [intv], 'rand_seed': rand_seed, + 'verbose': 0, + #'remove_dead': False, } - sim = ss.Sim(people=ppl, modules=[hiv, pregnancy], pars=pars, label=f'Sim with {n} agents and intv_cov={intv_cov}') + sim = ss.Sim(people=ppl, diseases=[hiv], demographics=[pregnancy], pars=pars, label=f'Sim with {n} agents and intv_cov={intv_cov}') sim.initialize() sim.run() @@ -104,7 +107,7 @@ def plot_scenarios(df): g.set_titles(col_template='{col_name}', row_template='Coverage: {row_name}') g.fig.suptitle('Multistream' if ms else 'Centralized') g.fig.subplots_adjust(top=0.88) - g.set_xlabels(r'Value - Reference at $t_i$') + g.set_xlabels(r'Timestep $t_i$') g.fig.savefig(os.path.join(figdir, 'diff_multistream.png' if ms else 'diff_centralized.png'), bbox_inches='tight', dpi=300) ## FINAL TIME diff --git a/tests/run_stable_monogamy.py b/tests/run_stable_monogamy.py index b7957e50..53bb33df 100644 --- a/tests/run_stable_monogamy.py +++ b/tests/run_stable_monogamy.py @@ -10,7 +10,7 @@ import networkx as nx import sys -multistream = True # Can set multistream to False for comparison +ss.options(multistream = True) # Can set multistream to False for comparison plot_graph = True class Graph(): @@ -57,7 +57,7 @@ def initialize(self, sim): self.initialized = True return - def apply(self, sim): + def update_results(self, sim): nodes = pd.DataFrame({ 'female': sim.people.female.values, 'dead': sim.people.dead.values, @@ -81,31 +81,28 @@ def finalize(self, sim): def run_sim(n=25, intervention=False, analyze=False): ppl = ss.People(n) - net_pars = {'multistream': multistream} - ppl.networks = ss.ndict(ss.stable_monogamy(pars=net_pars))#, ss.maternal(net_pars)) + ppl.networks = ss.ndict(ss.stable_monogamy())#, ss.maternal(net_pars)) hiv_pars = { 'beta': {'stable_monogamy': [0.06, 0.04]}, 'initial': 0, - 'multistream': multistream, } hiv = ss.HIV(hiv_pars) - art_pars = {'multistream': multistream} - art = ss.hiv.ART(0, 0.5, pars=art_pars) + art = ss.hiv.ART(0, 0.5) pars = { 'start': 1980, 'end': 2010, + 'remove_dead': False, # So we can see who dies 'interventions': [art] if intervention else [], 'rand_seed': 0, - 'analyzers': [rng_analyzer(pars={'multistream': multistream})] if analyze else [], - 'multistream': multistream, + 'analyzers': [rng_analyzer()] if analyze else [], } - sim = ss.Sim(people=ppl, modules=[hiv], pars=pars, label=f'Sim with {n} agents and intv={intervention}') + sim = ss.Sim(people=ppl, diseases=[hiv], pars=pars, label=f'Sim with {n} agents and intv={intervention}') sim.initialize() - sim.modules['hiv'].set_prognoses(sim, np.arange(0,n,2), from_uids=None) + sim.diseases['hiv'].set_prognoses(sim, np.arange(0,n,2), from_uids=None) sim.run() diff --git a/tests/run_sweep.py b/tests/run_sweep.py index a751891b..737fbb64 100644 --- a/tests/run_sweep.py +++ b/tests/run_sweep.py @@ -20,24 +20,20 @@ def run_sim(n, x_beta, rand_seed, multistream): ppl = ss.People(n) - net_pars = {'multistream': multistream} - ppl.networks = ss.ndict(ss.simple_embedding(pars=net_pars), ss.maternal(pars=net_pars)) + ppl.networks = ss.ndict(ss.simple_embedding(), ss.maternal()) hiv_pars = { 'beta': {'simple_embedding': [x_beta * 0.10, x_beta * 0.08], 'maternal': [x_beta * 0.2, 0]}, 'initial': 10, - 'multistream': multistream, } hiv = ss.HIV(hiv_pars) - preg_pars = {'multistream': multistream} - pregnancy = ss.Pregnancy(preg_pars) + pregnancy = ss.Pregnancy() pars = { 'start': 1980, 'end': 2010, 'rand_seed': rand_seed, - 'multistream': multistream, } sim = ss.Sim(people=ppl, modules=[hiv, pregnancy], pars=pars, label=f'Sim with {n} agents and x_beta={x_beta}') sim.initialize() @@ -60,6 +56,7 @@ def run_scenarios(figdir): results = [] times = {} for multistream in [True, False]: + ss.options(multistream=multistream) cfgs = [] for rs in range(n_rand_seeds): for x_beta in x_beta_levels: diff --git a/tests/simple.py b/tests/simple.py index 45c489c6..2c5fa2de 100644 --- a/tests/simple.py +++ b/tests/simple.py @@ -19,5 +19,4 @@ plt.figure() plt.plot(sim.tivec, sim.results.hiv.n_infected) -plt.title('HIV number of infections') - +plt.title('HIV number of infections') \ No newline at end of file From 3c15d380f457e94cdcb9bbf0b9633aec20bdaed5 Mon Sep 17 00:00:00 2001 From: Daniel Klein Date: Mon, 16 Oct 2023 09:00:20 -0700 Subject: [PATCH 25/34] New plotting in "stable monogamy" revealed a bug with births. --- stisim/hiv.py | 31 +---- stisim/modules.py | 14 +- stisim/networks.py | 2 +- tests/run_multistream.py | 21 +-- tests/run_stable_monogamy.py | 257 +++++++++++++++++++++++++++++------ 5 files changed, 239 insertions(+), 86 deletions(-) diff --git a/stisim/hiv.py b/stisim/hiv.py index e7827ca5..e41d5dc2 100644 --- a/stisim/hiv.py +++ b/stisim/hiv.py @@ -16,8 +16,9 @@ def __init__(self, pars=None): self.susceptible = ss.State('susceptible', bool, True) self.infected = ss.State('infected', bool, False) - self.ti_infected = ss.State('ti_infected', float, ss.INT_NAN) + self.ti_infected = ss.State('ti_infected', int, ss.INT_NAN) self.on_art = ss.State('on_art', bool, False) + self.ti_art = ss.State('ti_art', int, ss.INT_NAN) self.on_prep = ss.State('on_prep', bool, False) self.cd4 = ss.State('cd4', float, 500) @@ -41,7 +42,7 @@ def update_states(self, sim): self.rel_sus[sim.people.alive & ~self.infected & self.on_prep] = 0.04 self.rel_trans[sim.people.alive & self.infected & self.on_art] = 0.04 - hiv_death_prob = 0.2 / (self.pars.cd4_min - self.pars.cd4_max)**2 * (self.cd4 - self.pars.cd4_max)**2 + hiv_death_prob = 0.05 / (self.pars.cd4_min - self.pars.cd4_max)**2 * (self.cd4 - self.pars.cd4_max)**2 can_die = ss.true(sim.people.alive & sim.people.hiv.infected) hiv_deaths = self.rng_dead.bernoulli_filter(size=can_die, prob=hiv_death_prob[can_die]) sim.people.alive[hiv_deaths] = False @@ -81,7 +82,6 @@ def __init__(self, t: np.array, coverage: np.array, **kwargs): super().__init__(**kwargs) self.rng_add_ART = ss.Stream('add_ART') - self.rng_remove_ART = ss.Stream('remove_ART') return @@ -98,36 +98,13 @@ def apply(self, sim): recently_infected = ss.true((sim.people.hiv.ti_infected == sim.ti-2) & sim.people.alive) # 2 year (step) delay inds = self.rng_add_ART.bernoulli_filter(size=recently_infected, prob=coverage) sim.people.hiv.on_art[inds] = True + sim.people.hiv.ti_art[inds] = sim.ti # Add result sim.results.hiv.n_art = np.count_nonzero(sim.people.alive & sim.people.hiv.on_art) return len(inds) - ''' - on_art = sim.people.alive & sim.people.hiv.on_art - infected = sim.people.alive & sim.people.hiv.infected - n_change = np.round(coverage * infected.sum() - on_art.sum()).astype(int) - if n_change > 0: - # Add more ART - eligible = ss.true(sim.people.alive & sim.people.hiv.infected & ~sim.people.hiv.on_art) - n_eligible = len(eligible) #np.count_nonzero(eligible) - if n_eligible: - inds = self.rng_add_ART.bernoulli_filter(size=eligible, prob=min(n_eligible, n_change)/n_eligible) - sim.people.hiv.on_art[inds] = True - elif n_change < 0: - # Take some people off ART - eligible = sim.people.alive & sim.people.hiv.infected & sim.people.hiv.on_art - n_eligible = np.count_nonzero(eligible) - inds = self.rng_remove_ART.bernoulli_filter(size=eligible, prob=-n_change/n_eligible) - sim.people.hiv.on_art[inds] = False - - # Add result - sim.results.hiv.n_art = np.count_nonzero(sim.people.alive & sim.people.hiv.on_art) - - return - ''' - class PrEP(ss.Intervention): def __init__(self, t: np.array, coverage: np.array, **kwargs): diff --git a/stisim/modules.py b/stisim/modules.py index facc8452..cb07f8c9 100644 --- a/stisim/modules.py +++ b/stisim/modules.py @@ -20,11 +20,6 @@ def __init__(self, pars=None, label=None, requires=None, *args, **kwargs): self.initialized = False self.finalized = False - # Random number streams - self.rng_init_cases = ss.Stream(f'initial_cases_{self.name}') - self.rng_trans = ss.Stream(f'trans_{self.name}') - self.rng_choose_infector = ss.Stream(f'choose_infector_{self.name}') - return def __call__(self, *args, **kwargs): @@ -100,6 +95,11 @@ def __init__(self, pars=None, *args, **kwargs): self.infected = ss.State('infected', bool, False) self.ti_infected = ss.State('ti_infected', float, np.nan) + # Random number streams + self.rng_init_cases = ss.Stream(f'initial_cases_{self.name}') + self.rng_trans = ss.Stream(f'trans_{self.name}') + self.rng_choose_infector = ss.Stream(f'choose_infector_{self.name}') + return def initialize(self, sim): @@ -218,8 +218,8 @@ def set_prognoses(self, sim, to_uids, from_uids): pass def update_results(self, sim): - self.results['n_susceptible'][sim.ti] = np.count_nonzero(self.susceptible) - self.results['n_infected'][sim.ti] = np.count_nonzero(self.infected) + self.results['n_susceptible'][sim.ti] = np.count_nonzero(self.susceptible & sim.people.alive) + self.results['n_infected'][sim.ti] = np.count_nonzero(self.infected & sim.people.alive) self.results['prevalence'][sim.ti] = self.results.n_infected[sim.ti] / np.count_nonzero(sim.people.alive) self.results['new_infections'][sim.ti] = np.count_nonzero(self.ti_infected == sim.ti) diff --git a/stisim/networks.py b/stisim/networks.py index d2c200ed..620fad04 100644 --- a/stisim/networks.py +++ b/stisim/networks.py @@ -387,7 +387,7 @@ def add_pairs(self, people, ti=None): self['p2'] = np.concatenate([self['p2'], available_f[ind_f]]) beta = np.ones(n_pairs) - dur = self.rng_mean_dur.poisson(ind_m, self.mean_dur) + dur = self.rng_mean_dur.poisson(available_m[ind_m], self.mean_dur) self['beta'] = np.concatenate([self['beta'], beta]) self['dur'] = np.concatenate([self['dur'], dur]) diff --git a/tests/run_multistream.py b/tests/run_multistream.py index 8b63551e..def80e38 100644 --- a/tests/run_multistream.py +++ b/tests/run_multistream.py @@ -8,9 +8,10 @@ import sciris as sc import pandas as pd import seaborn as sns +import numpy as np n = 1_000 # Agents -n_rand_seeds = 100 +n_rand_seeds = 250 intv_cov_levels = [0.025, 0.05, 0.10, 0.73] + [0] # Must include 0 as that's the baseline # Choose ART or PrEP @@ -25,18 +26,21 @@ def run_sim(n, intv_cov, rand_seed, multistream): ss.options(multistream=multistream) ppl = ss.People(n) - ppl.networks = ss.ndict(ss.simple_embedding(), ss.maternal()) + ppl.networks = ss.ndict( + ss.simple_embedding(mean_dur=4), + ##################ss.maternal() + ) hiv_pars = { - #'beta': {'simple_embedding': [0.10, 0.08], 'maternal': [0.2,0]}, - 'beta': {'simple_embedding': [0.40, 0.35], 'maternal': [0.2,0]}, - 'initial': 10, + #'beta': {'simple_embedding': [0.10, 0.08], 'maternal': [0.2, 0]}, + 'beta': {'simple_embedding': [0.2, 0.15], 'maternal': [0.2, 0]}, + 'initial': int(np.maximum(10, np.ceil(0.05*n))), } hiv = ss.HIV(hiv_pars) pregnancy = ss.Pregnancy() - intv = intervention[choice](t=[0, 1], coverage=[0, intv_cov]) + intv = intervention[choice](t=[0, 10, 20], coverage=[0, intv_cov/3, intv_cov]) pars = { 'start': 1980, 'end': 2020, @@ -45,7 +49,8 @@ def run_sim(n, intv_cov, rand_seed, multistream): 'verbose': 0, #'remove_dead': False, } - sim = ss.Sim(people=ppl, diseases=[hiv], demographics=[pregnancy], pars=pars, label=f'Sim with {n} agents and intv_cov={intv_cov}') + ##############################sim = ss.Sim(people=ppl, diseases=[hiv], demographics=[pregnancy], pars=pars, label=f'Sim with {n} agents and intv_cov={intv_cov}') + sim = ss.Sim(people=ppl, diseases=[hiv], demographics=None, pars=pars, label=f'Sim with {n} agents and intv_cov={intv_cov}') sim.initialize() sim.run() @@ -54,7 +59,7 @@ def run_sim(n, intv_cov, rand_seed, multistream): #'hiv.n_infected': sim.results.hiv.n_infected, 'hiv.prevalence': sim.results.hiv.prevalence, 'hiv.cum_deaths': sim.results.hiv.new_deaths.cumsum(), - 'pregnancy.cum_births': sim.results.pregnancy.births.cumsum(), + #################################'pregnancy.cum_births': sim.results.pregnancy.births.cumsum(), }) df['intv_cov'] = intv_cov df['rand_seed'] = rand_seed diff --git a/tests/run_stable_monogamy.py b/tests/run_stable_monogamy.py index 53bb33df..f83db182 100644 --- a/tests/run_stable_monogamy.py +++ b/tests/run_stable_monogamy.py @@ -9,10 +9,17 @@ import numpy as np import networkx as nx import sys +import os +import argparse +import sciris as sc +import seaborn as sns ss.options(multistream = True) # Can set multistream to False for comparison plot_graph = True +figdir = os.path.join(os.getcwd(), 'figs', 'stable_monogamy') +sc.path(figdir).mkdir(parents=True, exist_ok=True) + class Graph(): def __init__(self, nodes, edges): self.graph = nx.from_pandas_edgelist(df=edges, source='p1', target='p2', edge_attr=True) @@ -20,7 +27,7 @@ def __init__(self, nodes, edges): nx.set_node_attributes(self.graph, nodes.transpose().to_dict()) return - def draw_nodes(self, filter, ax, **kwargs): + def draw_nodes(self, filter, pos, ax, **kwargs): inds = [i for i,n in self.graph.nodes.data() if filter(n)] nc = ['red' if nd['hiv'] else 'lightgray' for i, nd in self.graph.nodes.data() if i in inds] ec = ['green' if nd['on_art'] or nd['on_prep'] else 'black' for i, nd in self.graph.nodes.data() if i in inds] @@ -28,8 +35,8 @@ def draw_nodes(self, filter, ax, **kwargs): nx.draw_networkx_nodes(self.graph, nodelist=inds, pos=pos, ax=ax, node_color=nc, edgecolors=ec, **kwargs) return - def plot(self, pos, ax=None): - kwargs = dict(node_shape='x', node_size=250, linewidths=2, ax=ax) + def plot(self, pos, edge_labels=False, ax=None): + kwargs = dict(node_shape='x', node_size=250, linewidths=2, ax=ax, pos=pos) self.draw_nodes(lambda n: n['dead'], **kwargs) kwargs['node_shape'] = 'o' @@ -40,11 +47,12 @@ def plot(self, pos, ax=None): nx.draw_networkx_edges(self.graph, pos=pos, ax=ax) nx.draw_networkx_labels(self.graph, labels={i:int(a['cd4']) for i,a in self.graph.nodes.data()}, font_size=8, pos=pos, ax=ax) - #nx.draw_networkx_edge_labels(self.graph, edge_labels={(i,j): int(a['dur']) for i,j,a in self.graph.edges.data()}, font_size=8, pos=pos, ax=ax) + if edge_labels: + nx.draw_networkx_edge_labels(self.graph, edge_labels={(i,j): int(a['dur']) for i,j,a in self.graph.edges.data()}, font_size=8, pos=pos, ax=ax) return -class rng_analyzer(ss.Analyzer): +class GraphAnalyzer(ss.Analyzer): ''' Simple analyzer to assess if random streams are working ''' def __init__(self, **kwargs): @@ -55,10 +63,12 @@ def __init__(self, **kwargs): def initialize(self, sim): self.initialized = True + self.update_results(sim, init=True) return - def update_results(self, sim): + def update_results(self, sim, init=False): nodes = pd.DataFrame({ + 'age': sim.people.age.values, 'female': sim.people.female.values, 'dead': sim.people.dead.values, 'hiv': sim.people.hiv.infected.values, @@ -70,7 +80,8 @@ def update_results(self, sim): edges = pd.DataFrame(sim.people.networks[0].to_dict()) #sim.people.networks['simple_embedding'].to_df() #TODO: repr issues - self.graphs[sim.ti] = Graph(nodes, edges) + idx = sim.ti if not init else -1 + self.graphs[idx] = Graph(nodes, edges) return def finalize(self, sim): @@ -78,14 +89,16 @@ def finalize(self, sim): return + def run_sim(n=25, intervention=False, analyze=False): ppl = ss.People(n) - ppl.networks = ss.ndict(ss.stable_monogamy())#, ss.maternal(net_pars)) + ppl.networks = ss.ndict(ss.simple_embedding(mean_dur=5))#, ss.maternal()) hiv_pars = { - 'beta': {'stable_monogamy': [0.06, 0.04]}, - 'initial': 0, + #'beta': {'simple_embedding': [0.06, 0.04]}, + 'beta': {'simple_embedding': [0.3, 0.25]}, + 'initial': 0.25 * n, } hiv = ss.HIV(hiv_pars) @@ -93,34 +106,47 @@ def run_sim(n=25, intervention=False, analyze=False): pars = { 'start': 1980, - 'end': 2010, - 'remove_dead': False, # So we can see who dies + 'end': 2020, + 'remove_dead': False, # So we can see who dies, sim results should not change with True 'interventions': [art] if intervention else [], 'rand_seed': 0, - 'analyzers': [rng_analyzer()] if analyze else [], + 'analyzers': [GraphAnalyzer()] if analyze else [], } - sim = ss.Sim(people=ppl, diseases=[hiv], pars=pars, label=f'Sim with {n} agents and intv={intervention}') + sim = ss.Sim(people=ppl, diseases=[hiv], demographics=[ss.Pregnancy()], pars=pars, label=f'Sim with {n} agents and intv={intervention}') sim.initialize() - sim.diseases['hiv'].set_prognoses(sim, np.arange(0,n,2), from_uids=None) + #sim.diseases['hiv'].set_prognoses(sim, np.arange(0,n,2), from_uids=None) sim.run() return sim -n = 10 -sim2 = run_sim(n, intervention=True, analyze=plot_graph) -sim1 = run_sim(n, intervention=False, analyze=plot_graph) -if plot_graph: +def run_scenario(n=10, analyze=True): + #sim2 = run_sim(n, intervention=True, analyze=analyze) + #sim1 = run_sim(n, intervention=False, analyze=analyze) + sims = sc.parallelize(run_sim, kwargs={'n':n, 'analyze': analyze}, iterkwargs=[{'intervention':True}, {'intervention':False}], die=True) + + for i, sim in enumerate(sims): + sim.save(os.path.join(figdir, f'sim{i}.obj')) + + return sims + + +def plot_graph(sim1, sim2): g1 = sim1.analyzers[0].graphs g2 = sim2.analyzers[0].graphs - pos = {i:(np.cos(2*np.pi*i/n), np.sin(2*np.pi*i/n)) for i in range(n)} + n = len(g1[0].graph) + el = n <= 10 # Edge labels + #pos = {i:(np.cos(2*np.pi*i/n), np.sin(2*np.pi*i/n)) for i in range(n)} + pos = {i:(nd['age'], 2*nd['female']-1 + np.random.uniform(-0.3, 0.3)) for i, nd in g1[0].graph.nodes.data()} + #pos = nx.spring_layout(g1[0].graph, k=None, pos=None, fixed=None, iterations=50, threshold=0.0001, weight=None, scale=1, center=None, dim=2, seed=None) + #pos = nx.multipartite_layout(g1[0].graph, subset_key='female', align='vertical', scale=10, center=None) fig, axv = plt.subplots(1, 2, figsize=(10,5)) global ti - ti = 0 + ti = -1 # Initial state is -1 timax = sim1.tivec[-1] def on_press(event): print('press', event.key) @@ -129,14 +155,14 @@ def on_press(event): if event.key == 'right': ti = min(ti+1, timax) elif event.key == 'left': - ti = max(ti-1, 0) + ti = max(ti-1, -1) # Clear axv[0].clear() axv[1].clear() - g1[ti].plot(pos, ax=axv[0]) - g2[ti].plot(pos, ax=axv[1]) + g1[ti].plot(pos, edge_labels=el, ax=axv[0]) + g2[ti].plot(pos, edge_labels=el, ax=axv[1]) fig.suptitle(f'Time is {ti}') fig.canvas.draw() @@ -144,23 +170,168 @@ def on_press(event): g1[ti].plot(pos, ax=axv[0]) g2[ti].plot(pos, ax=axv[1]) - plt.show() + fig.suptitle(f'Time is {ti}') + + # Plot + fig, axv = plt.subplots(2,1, sharex=True) + axv[0].plot(sim1.tivec, sim1.results.hiv.n_infected, label='Baseline') + axv[0].plot(sim2.tivec, sim2.results.hiv.n_infected, ls=':', label='Intervention') + axv[0].set_title('HIV number of infections') + + axv[1].plot(sim1.tivec, sim1.results.hiv.new_deaths, label='Baseline') + axv[1].plot(sim2.tivec, sim2.results.hiv.new_deaths, ls=':', label='Intervention') + axv[1].set_title('HIV Deaths') + + ''' Gonorrhea removed for now + axv[2].plot(sim1.tivec, sim1.results.gonorrhea.n_infected, label='Baseline') + axv[2].plot(sim2.tivec, sim2.results.gonorrhea.n_infected, ls=':', label='Intervention') + axv[2].set_title('Gonorrhea number of infections') + ''' + plt.legend() + +def analyze_people(sim): + p = sim.people + years_lived = np.full(len(p), sim.ti) + years_lived[p.dead] = p.ti_dead[p.dead] + age_initial = p.age.values - years_lived + + df = pd.DataFrame({ + 'uid': p.uid._view, + 'age_initial': age_initial, + 'years_lived': years_lived, + 'ti_infected': p.hiv.ti_infected.values, + 'ti_art': p.hiv.ti_art.values, + 'ti_dead': p.ti_dead.values, + }) + df.replace(to_replace=ss.INT_NAN, value=np.nan, inplace=True) + df['age_infected'] = df['age_initial'] + df['ti_infected'] + df['age_art'] = df['age_initial'] + df['ti_art'] + df['age_dead'] = df['age_initial'] + df['ti_dead'] + return df + +def life_bars(data, **kwargs): + age_final = data['age_initial'] + data['years_lived'] + plt.barh(y=data.index, left=data['age_initial'], width=age_final-data['age_initial'], color='k') + + # Define bools + infected = ~data['ti_infected'].isna() + art = ~data['ti_art'].isna() + dead = ~data['ti_dead'].isna() + + # Infected + plt.barh(y=data.index[infected], left=data.loc[infected]['age_infected'], width=age_final[infected]-data.loc[infected]['age_infected'], color='r') + + # ART + plt.barh(y=data.index[art], left=data.loc[art]['age_art'], width=age_final[art]-data.loc[art]['age_art'], color='g') + + # Dead + #plt.barh(y=data.index[dead], left=data.loc[dead]['age_dead'], width=age_final[dead]-data.loc[dead]['age_dead'], color='k') + plt.scatter(y=data.index[dead], x=data.loc[dead]['age_dead'], color='k', marker='|') + + return + +def life_bars_nested(df): + + N = df['sim'].nunique() + height = 0.9/N + + fig, ax = plt.subplots(figsize=(10,6)) + + for n, (lbl, data) in enumerate(df.groupby('sim')): + ys = n/(N+1) # Leave space + + age_final = data['age_initial'] + data['years_lived'] + plt.barh(y=data.index + ys, left=data['age_initial'], width=age_final-data['age_initial'], color='k', height=height) + + # Define bools + infected = ~data['ti_infected'].isna() + art = ~data['ti_art'].isna() + dead = ~data['ti_dead'].isna() + + # Infected + plt.barh(y=data.index[infected] + ys, left=data.loc[infected]['age_infected'], width=age_final[infected]-data.loc[infected]['age_infected'], color='r', height=height) + + # ART + plt.barh(y=data.index[art] + ys, left=data.loc[art]['age_art'], width=age_final[art]-data.loc[art]['age_art'], color='g', height=height) + + # Dead + plt.scatter(y=data.index[dead] + ys, x=data.loc[dead]['age_dead'], color='k', marker='|') + + return fig + +def ti_bars_nested(df): -# Plot -fig, axv = plt.subplots(2,1, sharex=True) -axv[0].plot(sim1.tivec, sim1.results.hiv.n_infected, label='Baseline') -axv[0].plot(sim2.tivec, sim2.results.hiv.n_infected, ls=':', label='Intervention') -axv[0].set_title('HIV number of infections') - -axv[1].plot(sim1.tivec, sim1.results.hiv.new_deaths, label='Baseline') -axv[1].plot(sim2.tivec, sim2.results.hiv.new_deaths, ls=':', label='Intervention') -axv[1].set_title('HIV Deaths') - -''' Gonorrhea removed for now -axv[2].plot(sim1.tivec, sim1.results.gonorrhea.n_infected, label='Baseline') -axv[2].plot(sim2.tivec, sim2.results.gonorrhea.n_infected, ls=':', label='Intervention') -axv[2].set_title('Gonorrhea number of infections') -''' -plt.legend() -plt.show() -print('Done') \ No newline at end of file + N = df['sim'].nunique() + height = 0.5/N + + fig, ax = plt.subplots(figsize=(10,6)) + + for n, (lbl, data) in enumerate(df.groupby('sim')): + ys = n/(N+1) # Leave space + + ti_initial = np.maximum(-data['age_initial'], 0) + ti_final = data['ti_dead'].fillna(40) + plt.barh(y=data.index + ys, left=ti_initial, width=ti_final - ti_initial, color='k', height=height) + + # Define bools + infected = ~data['ti_infected'].isna() + art = ~data['ti_art'].isna() + dead = ~data['ti_dead'].isna() + + # Infected + plt.barh(y=data.index[infected] + ys, left=data.loc[infected]['ti_infected'], width=ti_final[infected]-data.loc[infected]['ti_infected'], color='r', height=height) + + # ART + plt.barh(y=data.index[art] + ys, left=data.loc[art]['ti_art'], width=ti_final[art]-data.loc[art]['ti_art'], color='g', height=height) + + # Dead + plt.scatter(y=data.index[dead] + ys, x=data.loc[dead]['ti_dead'], color='k', marker='|') + + return fig + + +def plot_longitudinal(sim1, sim2): + + df1 = analyze_people(sim1) + df1['sim'] = 'Baseline' + df2 = analyze_people(sim2) + df2['sim'] = 'With ART' + + df = pd.concat([df1, df2]).set_index('uid') + #age_vars = ['age_initial', 'age_infected', 'age_art', 'age_dead'] + #age_vars = ['age_end', 'age_dead', 'age_art', 'age_infected', 'age_initial'] + #dfm = df.melt(id_vars=['uid', 'sim'], value_vars=age_vars, var_name='event', value_name='age') + #g = sns.catplot(kind='bar', col='sim', data=dfm.reset_index(), y='uid', x='age', hue='event', orient='h', hue_order=age_vars, dodge=False, + # palette=sns.blend_palette(colors=['black', 'green', 'red', 'gray', 'white'], n_colors=len(age_vars))) + + #f = life_bars_nested(df) + f = ti_bars_nested(df) + + #g = sns.FacetGrid(data=df, col='sim') + #g.map_dataframe(life_bars) + + + return + + + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('-p', '--plot', help='Plot from a cached CSV file', type=str) + parser.add_argument('-n', help='Number of agents', type=int, default=100) + args = parser.parse_args() + + if args.plot: + print('Reading files', args.plot) + sim1 = sc.load(os.path.join(args.plot, 'sim1.obj')) + sim2 = sc.load(os.path.join(args.plot, 'sim2.obj')) + else: + print('Running scenarios') + [sim1, sim2] = run_scenario(n=args.n) + + #plot_graph(sim1, sim2) + plot_longitudinal(sim1, sim2) + + plt.show() + print('Done') \ No newline at end of file From 1f13cc85d0ad519d92c42b48cabdcab844f3528d Mon Sep 17 00:00:00 2001 From: Daniel Klein Date: Mon, 16 Oct 2023 22:22:22 -0700 Subject: [PATCH 26/34] About to make more changes, code is non-functional in current state. --- stisim/demographics.py | 9 +++++++- stisim/modules.py | 7 ++---- stisim/networks.py | 6 ++--- stisim/people.py | 43 ++++++++++++++++++++++++++++++++---- stisim/sim.py | 2 +- stisim/states.py | 14 ++++++------ stisim/streams.py | 22 ++++++++++++------ tests/run_multistream.py | 27 +++++++++++++--------- tests/run_stable_monogamy.py | 26 ++++++++++++---------- 9 files changed, 105 insertions(+), 51 deletions(-) diff --git a/stisim/demographics.py b/stisim/demographics.py index dabfcb25..d29f5f72 100644 --- a/stisim/demographics.py +++ b/stisim/demographics.py @@ -247,6 +247,8 @@ def __init__(self, pars=None): 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): @@ -316,10 +318,15 @@ def make_pregnancies(self, sim): # 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=10*sim.pars['n_agents'], size=uids.max()+1, dtype=int)[uids] + # 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(size=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) # Add connections to any vertical transmission layers diff --git a/stisim/modules.py b/stisim/modules.py index cb07f8c9..e4b83b5a 100644 --- a/stisim/modules.py +++ b/stisim/modules.py @@ -45,7 +45,7 @@ def initialize(self, sim): # Connect the streams to the sim for stream in self.streams: if not stream.initialized: - stream.initialize(sim.streams) + stream.initialize(sim.streams, sim.people.slot) self.initialized = True return @@ -129,9 +129,6 @@ def set_initial_states(self, sim): if self.pars['initial'] <= 0: return - #initial_cases = np.random.choice(sim.people.uid, self.pars['initial']) - #rng = sim.rngs.get(f'initial_cases_{self.name}') - #initial_cases = ss.binomial_filter_rng(prob=self.pars['initial']/len(sim.people), size=sim.people.uid, rng=rng, block_size=len(sim.people._uid_map)) initial_cases = self.rng_init_cases.bernoulli_filter(size=sim.people.uid, prob=self.pars['initial']/len(sim.people)) self.set_prognoses(sim, initial_cases, from_uids=None) # TODO: sentinel value to indicate seeds? @@ -203,7 +200,7 @@ def make_new_cases(self, sim): # Decide whom the infection came from using one random number for each b (aligned by block size) frm = np.zeros_like(new_cases) - r = self.rng_choose_infector.random(new_cases) + r = self.rng_choose_infector.random( new_cases ) new_cases_idx = new_cases_bool.nonzero()[0] prob = (1-node_from_node[new_cases_idx]) # Prob of acquiring from each node | can constrain to just neighbors? cumsum = (prob / ((prob.sum(axis=1)[:,np.newaxis]))).cumsum(axis=1) diff --git a/stisim/networks.py b/stisim/networks.py index 620fad04..dae32f40 100644 --- a/stisim/networks.py +++ b/stisim/networks.py @@ -315,9 +315,9 @@ def initialize(self, sim): super().initialize(sim) # Initialize random streams - self.rng_pair_12.initialize(sim.streams) - self.rng_pair_21.initialize(sim.streams) - self.rng_mean_dur.initialize(sim.streams) + self.rng_pair_12.initialize(sim.streams, sim.people.slot) + self.rng_pair_21.initialize(sim.streams, sim.people.slot) + self.rng_mean_dur.initialize(sim.streams, sim.people.slot) self.add_pairs(sim.people, ti=0) return diff --git a/stisim/people.py b/stisim/people.py index 64efe9e7..3151188a 100644 --- a/stisim/people.py +++ b/stisim/people.py @@ -71,6 +71,34 @@ def grow(self, n): return new_uids + def grow_uids(self, uids): + """ + Increase the number of agents given new uids to add + + :param uids: uids of agents to add, but need to check for collisions + """ + start_uid = len(self._uid_map) + ids_in_map = np.intersect1d(uids, self._uid_map, assume_unique=False, return_indices=True)[1] + if len(ids_in_map): + is_collision = ~np.isnan(self.age[uids[ids_in_map]]) + #collisions = ss.false(is_collision) + if is_collision.any(): + n_collisions = is_collision.sum() + uids[ids_in_map[is_collision]] = np.arange(start_uid, start_uid + n_collisions) # Put at end, will mess up rng coherence slightly + #raise Exception('no bueno') + print(f'Encountered {n_collisions} collisions') + + n = uids.max() - start_uid + 1 + if n > 0: + new_uids = self.grow(n) + self.alive[new_uids] = False # Override the default + + # Restore sensible defaults + self.age[uids] = 0 + self.alive[uids] = True + + return uids # MAY NEED TO MODIFY IF COLLISION + def remove(self, uids_to_remove): """ Reduce the number of agents @@ -155,10 +183,11 @@ def __init__(self, n, age_data=None, extra_states=None, networks=None, rand_seed self.rng_female = ss.Stream('female') states = [ - ss.State('age', float, 0), + ss.State('slot', int, ss.INT_NAN), # MUST BE FIRST + ss.State('age', float, np.nan), # NaN until conceived ss.State('female', bool, ss.bernoulli(0.5, rng=self.rng_female)), ss.State('debut', float), - ss.State('alive', bool, True), + ss.State('alive', bool, True), # Redundant with ti_dead == ss.INT_NAN ss.State('ti_dead', int, ss.INT_NAN), # Time index for death ss.State('scale', float, 1.0), ] @@ -181,14 +210,20 @@ def get_age_dist(age_data, rng): def initialize(self, sim): """ Initialization """ - self.rng_female.initialize(sim.streams) - self.rng_agedist.initialize(sim.streams) + + # Slot first... TODO + + self.rng_female.initialize(sim.streams, self.states['slot']) + self.rng_agedist.initialize(sim.streams, self.states['slot']) for name, state in self.states.items(): self.add_state(state) # Register the state internally for dynamic growth #self.states.append(state) # Expose these states with their original names state.initialize(self) # Connect the state to this people instance setattr(self, name, state) + if name == 'slot': + # Initialize here in case other states use random streams that depend on slots being initialized + self.slot[:] = np.copy(self.uid) # TODO: .values and [:] needed? self.age[:] = self.age_data_dist.sample(len(self)) self.initialized = True diff --git a/stisim/sim.py b/stisim/sim.py index 01d2f894..b082c3ae 100644 --- a/stisim/sim.py +++ b/stisim/sim.py @@ -273,7 +273,7 @@ def init_interventions(self): raise TypeError(errormsg) for stream in intervention.streams: - stream.initialize(self.streams) + stream.initialize(self.streams, self.people.slot) return diff --git a/stisim/states.py b/stisim/states.py index 2ee99287..b4af06e6 100644 --- a/stisim/states.py +++ b/stisim/states.py @@ -281,7 +281,7 @@ def _new_items(self, n): def _fill_data(self, n0, n1): # Fill data - n = n1-n0 + n = int(n1-n0) if callable(self.fill_value): self._data[n0:n1] = self.fill_value(n) else: @@ -293,25 +293,25 @@ def initialize(self, n): self.n = n self._map_arrays() - def grow(self, n): + def grow(self, n, fill=True): if self.n + n > self._s: # If the total number of agents exceeds the array size, extend the storage array n_new = max(n, int(self._s / 2)) # Minimum 50% growth self._data = np.concatenate([self._data, self._new_items(n_new)], axis=0) self.n += n # Increase the count of the number of agents by `n` (the requested number of new agents) - self._map_arrays() + self._map_arrays(fill) - def _trim(self, inds): + def _trim(self, inds, fill=True): # Keep only specified indices # Note that these are indices, not UIDs! n = len(inds) self._data[:n] = self._data[inds] # DJK MOVING TO MAP: self._data[n:self.n] = self.fill_value(self.n-n) if callable(self.fill_value) else self.fill_value self.n = n - self._map_arrays() + self._map_arrays(fill) - def _map_arrays(self): + def _map_arrays(self, fill=True): """ Set main simulation attributes to be views of the underlying data @@ -319,7 +319,7 @@ def _map_arrays(self): (regardless of whether or not the underlying arrays have been resized) """ n0 = len(self._view) if self._view is not None else 0 - if self.n > n0: + if fill and self.n > n0: self._fill_data(n0, self.n) # Fill the new data self._view = self._data[:self.n] diff --git a/stisim/streams.py b/stisim/streams.py index 6440ede3..f0d0a9fe 100644 --- a/stisim/streams.py +++ b/stisim/streams.py @@ -151,13 +151,14 @@ def __init__(self, name, seed_offset=None, **kwargs): self.ready = True return - def initialize(self, streams): + def initialize(self, streams, slots): if self.initialized: # TODO: Raise warning assert not self.initialized return self.seed = streams.add(self) + self.slots = slots # E.g. sim.people.slots (instead of using uid as the slots directly) if 'bit_generator' not in self.kwargs: self.kwargs['bit_generator'] = np.random.PCG64(seed=self.seed) @@ -205,32 +206,39 @@ def draw_size(self, size): @_pre_draw def random(self, size): - return super(MultiStream, self).random(size=self.draw_size(size))[size] + slots = self.slots.values[size] + return super(MultiStream, self).random(size=self.draw_size(slots))[slots] @_pre_draw def uniform(self, size, **kwargs): - return super(MultiStream, self).uniform(size=self.draw_size(size), **kwargs)[size] + slots = self.slots.values[size] + return super(MultiStream, self).uniform(size=self.draw_size(slots), **kwargs)[slots] @_pre_draw def poisson(self, size, lam): - return super(MultiStream, self).poisson(size=self.draw_size(size), lam=lam)[size] + slots = self.slots.values[size] + return super(MultiStream, self).poisson(size=self.draw_size(slots), lam=lam)[slots] @_pre_draw def normal(self, size, mu=0, std=1): - return mu + std*super(MultiStream, self).normal(size=self.draw_size(size))[size] + slots = self.slots.values[size] + return mu + std*super(MultiStream, self).normal(size=self.draw_size(slots))[slots] @_pre_draw def negative_binomial(self, size, **kwargs): #n=nbn_n, p=nbn_p, size=n) - return super(MultiStream, self).negative_binomial(size=self.draw_size(size), **kwargs)[size] + slots = self.slots.values[size] + return super(MultiStream, self).negative_binomial(size=self.draw_size(slots), **kwargs)[slots] @_pre_draw def bernoulli(self, size, prob): #return super(MultiStream, self).choice([True, False], size=size.max()+1) # very slow #return (super(MultiStream, self).binomial(n=1, p=prob, size=size.max()+1))[size].astype(bool) # pretty fast - return super(MultiStream, self).random(size=self.draw_size(size))[size] < prob # fastest + slots = self.slots.values[size] + return super(MultiStream, self).random(size=self.draw_size(slots))[slots] < prob # fastest # @_pre_draw <-- handled by call to self.bernoullli def bernoulli_filter(self, size, prob): + #slots = self.slots[size[:]] return size[self.bernoulli(size, prob)] # Slightly faster on my machine for bernoulli to typecast def choice(self, size, a, **kwargs): diff --git a/tests/run_multistream.py b/tests/run_multistream.py index def80e38..86b649ee 100644 --- a/tests/run_multistream.py +++ b/tests/run_multistream.py @@ -10,8 +10,8 @@ import seaborn as sns import numpy as np -n = 1_000 # Agents -n_rand_seeds = 250 +n = 250 # Agents +n_rand_seeds = 100 intv_cov_levels = [0.025, 0.05, 0.10, 0.73] + [0] # Must include 0 as that's the baseline # Choose ART or PrEP @@ -21,20 +21,22 @@ figdir = os.path.join(os.getcwd(), 'figs', choice) sc.path(figdir).mkdir(parents=True, exist_ok=True) -def run_sim(n, intv_cov, rand_seed, multistream): +def run_sim(n, idx, intv_cov, rand_seed, multistream): + + print(f'Starting sim {idx} with rand_seed={rand_seed} and intv_cov={intv_cov}, multistream={multistream}') ss.options(multistream=multistream) ppl = ss.People(n) ppl.networks = ss.ndict( ss.simple_embedding(mean_dur=4), - ##################ss.maternal() + ss.maternal() # PROBLEM ) hiv_pars = { #'beta': {'simple_embedding': [0.10, 0.08], 'maternal': [0.2, 0]}, 'beta': {'simple_embedding': [0.2, 0.15], 'maternal': [0.2, 0]}, - 'initial': int(np.maximum(10, np.ceil(0.05*n))), + 'initial': int(np.maximum(10, np.ceil(0.01*n))), } hiv = ss.HIV(hiv_pars) @@ -47,10 +49,11 @@ def run_sim(n, intv_cov, rand_seed, multistream): 'interventions': [intv], 'rand_seed': rand_seed, 'verbose': 0, - #'remove_dead': False, + 'remove_dead': False, # PROBLEM + 'n_agents': len(ppl), # TODO } - ##############################sim = ss.Sim(people=ppl, diseases=[hiv], demographics=[pregnancy], pars=pars, label=f'Sim with {n} agents and intv_cov={intv_cov}') - sim = ss.Sim(people=ppl, diseases=[hiv], demographics=None, pars=pars, label=f'Sim with {n} agents and intv_cov={intv_cov}') + sim = ss.Sim(people=ppl, diseases=[hiv], demographics=[pregnancy], pars=pars, label=f'Sim with {n} agents and intv_cov={intv_cov}') # PROBLEM + ########################sim = ss.Sim(people=ppl, diseases=[hiv], demographics=None, pars=pars, label=f'Sim with {n} agents and intv_cov={intv_cov}') # FIX sim.initialize() sim.run() @@ -59,12 +62,14 @@ def run_sim(n, intv_cov, rand_seed, multistream): #'hiv.n_infected': sim.results.hiv.n_infected, 'hiv.prevalence': sim.results.hiv.prevalence, 'hiv.cum_deaths': sim.results.hiv.new_deaths.cumsum(), - #################################'pregnancy.cum_births': sim.results.pregnancy.births.cumsum(), + 'pregnancy.cum_births': sim.results.pregnancy.births.cumsum(), # PROBLEM }) df['intv_cov'] = intv_cov df['rand_seed'] = rand_seed df['multistream'] = multistream + print(f'Finishing sim {idx} with rand_seed={rand_seed} and intv_cov={intv_cov}, multistream={multistream}') + return df def run_scenarios(): @@ -74,9 +79,9 @@ def run_scenarios(): cfgs = [] for rs in range(n_rand_seeds): for intv_cov in intv_cov_levels: - cfgs.append({'intv_cov':intv_cov, 'rand_seed':rs, 'multistream':multistream}) + cfgs.append({'intv_cov':intv_cov, 'rand_seed':rs, 'multistream':multistream, 'idx':len(cfgs)}) T = sc.tic() - results += sc.parallelize(run_sim, kwargs={'n': n}, iterkwargs=cfgs, die=True) + results += sc.parallelize(run_sim, kwargs={'n': n}, iterkwargs=cfgs, die=True, serial=True) times[f'Multistream={multistream}'] = sc.toc(T, output=True) print('Timings:', times) diff --git a/tests/run_stable_monogamy.py b/tests/run_stable_monogamy.py index f83db182..2d9fb400 100644 --- a/tests/run_stable_monogamy.py +++ b/tests/run_stable_monogamy.py @@ -111,6 +111,7 @@ def run_sim(n=25, intervention=False, analyze=False): 'interventions': [art] if intervention else [], 'rand_seed': 0, 'analyzers': [GraphAnalyzer()] if analyze else [], + 'n_agents': len(ppl), # TODO: Build into Sim } sim = ss.Sim(people=ppl, diseases=[hiv], demographics=[ss.Pregnancy()], pars=pars, label=f'Sim with {n} agents and intv={intervention}') sim.initialize() @@ -123,8 +124,6 @@ def run_sim(n=25, intervention=False, analyze=False): def run_scenario(n=10, analyze=True): - #sim2 = run_sim(n, intervention=True, analyze=analyze) - #sim1 = run_sim(n, intervention=False, analyze=analyze) sims = sc.parallelize(run_sim, kwargs={'n':n, 'analyze': analyze}, iterkwargs=[{'intervention':True}, {'intervention':False}], die=True) for i, sim in enumerate(sims): @@ -191,17 +190,19 @@ def on_press(event): def analyze_people(sim): p = sim.people + ever_alive = ss.false(np.isnan(p.age)) years_lived = np.full(len(p), sim.ti) years_lived[p.dead] = p.ti_dead[p.dead] - age_initial = p.age.values - years_lived + years_lived = years_lived[ever_alive] # Trim, could be more efficient + age_initial = p.age[ever_alive].values - years_lived df = pd.DataFrame({ - 'uid': p.uid._view, + 'uid': p.uid[ever_alive], # if slicing, don't need ._view, 'age_initial': age_initial, 'years_lived': years_lived, - 'ti_infected': p.hiv.ti_infected.values, - 'ti_art': p.hiv.ti_art.values, - 'ti_dead': p.ti_dead.values, + 'ti_infected': p.hiv.ti_infected[ever_alive].values, + 'ti_art': p.hiv.ti_art[ever_alive].values, + 'ti_dead': p.ti_dead[ever_alive].values, }) df.replace(to_replace=ss.INT_NAN, value=np.nan, inplace=True) df['age_infected'] = df['age_initial'] + df['ti_infected'] @@ -261,17 +262,18 @@ def life_bars_nested(df): def ti_bars_nested(df): + df['ypos'] = pd.factorize(df.index.values)[0] N = df['sim'].nunique() height = 0.5/N fig, ax = plt.subplots(figsize=(10,6)) for n, (lbl, data) in enumerate(df.groupby('sim')): - ys = n/(N+1) # Leave space + yp = data['ypos'] + n/(N+1) # Leave space ti_initial = np.maximum(-data['age_initial'], 0) ti_final = data['ti_dead'].fillna(40) - plt.barh(y=data.index + ys, left=ti_initial, width=ti_final - ti_initial, color='k', height=height) + plt.barh(y=yp, left=ti_initial, width=ti_final - ti_initial, color='k', height=height) # Define bools infected = ~data['ti_infected'].isna() @@ -279,13 +281,13 @@ def ti_bars_nested(df): dead = ~data['ti_dead'].isna() # Infected - plt.barh(y=data.index[infected] + ys, left=data.loc[infected]['ti_infected'], width=ti_final[infected]-data.loc[infected]['ti_infected'], color='r', height=height) + plt.barh(y=yp[infected], left=data.loc[infected]['ti_infected'], width=ti_final[infected]-data.loc[infected]['ti_infected'], color='r', height=height) # ART - plt.barh(y=data.index[art] + ys, left=data.loc[art]['ti_art'], width=ti_final[art]-data.loc[art]['ti_art'], color='g', height=height) + plt.barh(y=yp[art], left=data.loc[art]['ti_art'], width=ti_final[art]-data.loc[art]['ti_art'], color='g', height=height) # Dead - plt.scatter(y=data.index[dead] + ys, x=data.loc[dead]['ti_dead'], color='k', marker='|') + plt.scatter(y=yp[dead], x=data.loc[dead]['ti_dead'], color='k', marker='|') return fig From 4780b612d50521029a0281e625310822808b736b Mon Sep 17 00:00:00 2001 From: Daniel Klein Date: Tue, 17 Oct 2023 08:09:13 -0700 Subject: [PATCH 27/34] New slotting seems to work. Interface still a bit messy. Need to fix Centralized RNG now. --- stisim/demographics.py | 12 +-- stisim/distributions.py | 63 ++++++++------- stisim/hiv.py | 8 +- stisim/modules.py | 6 +- stisim/networks.py | 6 +- stisim/people.py | 2 +- stisim/streams.py | 150 +++++++++++++++++++++++++++-------- tests/run_multistream.py | 4 +- tests/run_stable_monogamy.py | 100 ++++++++--------------- 9 files changed, 203 insertions(+), 148 deletions(-) diff --git a/stisim/demographics.py b/stisim/demographics.py index d29f5f72..f08708a1 100644 --- a/stisim/demographics.py +++ b/stisim/demographics.py @@ -313,21 +313,21 @@ 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 = self.rng_conception.bernoulli_filter(size=inds_to_choose_from, prob=self.pars.inci) + 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=10*sim.pars['n_agents'], size=uids.max()+1, dtype=int)[uids] + # TODO: User configure the bounds + new_slots = self.rng_uids.integers(low=sim.pars['n_agents'], high=10*sim.pars['n_agents'], uids=uids, dtype=int) # Grow the arrays and set properties for the 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(size=uids, prob=0.5) # Replace 0.5 with sex ratio at birth + #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) + 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 @@ -357,7 +357,7 @@ def set_prognoses(self, sim, to_uids, from_uids=None): # Outcomes for pregnancies dur = np.full(len(to_uids), sim.ti + self.pars.dur_pregnancy / sim.dt) - dead = self.rng_dead.bernoulli(size=to_uids, prob=self.pars.p_death) + 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 diff --git a/stisim/distributions.py b/stisim/distributions.py index 9b994e91..54accc85 100644 --- a/stisim/distributions.py +++ b/stisim/distributions.py @@ -44,10 +44,10 @@ def mean(self): """ raise NotImplementedError - def __call__(self, size=1, **kwargs): - return self.sample(size, **kwargs) + def __call__(self, size=None, uids=None, **kwargs): + return self.sample(size=size, uids=uids, **kwargs) - def sample(cls, size=1, **kwargs): + def sample(cls, size=None, uids=None, **kwargs): """ Return a specified number of samples from the distribution """ @@ -65,7 +65,7 @@ def __init__(self, vals, bins, **kwargs): def mean(self): return - def sample(self, size=1): + def sample(self, size, **kwargs): """ Sample using CDF """ bin_midpoints = self.bins[:-1] + np.diff(self.bins) / 2 cdf = np.cumsum(self.vals) @@ -88,8 +88,8 @@ def __init__(self, low, high, **kwargs): def mean(self): return (self.low + self.high) / 2 - def sample(self, size=1): - return self.stream.uniform(low=self.low, high=self.high, size=size) + def sample(self, **kwargs): + return self.stream.uniform(low=self.low, high=self.high, **kwargs) class bernoulli(Distribution): """ @@ -103,8 +103,8 @@ def __init__(self, p, **kwargs): def mean(self): return self.p - def sample(self, size=1): - return self.stream.bernoulli(prob=self.p, size=size) + def sample(self, **kwargs): + return self.stream.bernoulli(prob=self.p, **kwargs) class choice(Distribution): @@ -118,8 +118,8 @@ def __init__(self, choices, probabilities=None, replace=True, **kwargs): self.probabilities = probabilities self.replace = replace - def sample(self, size): - return self.stream.choice(a=self.choices, p=self.probabilities, replace=self.replace, size=size) + def sample(self, **kwargs): + return self.stream.choice(a=self.choices, p=self.probabilities, replace=self.replace, **kwargs) class normal(Distribution): @@ -132,8 +132,8 @@ def __init__(self, mean, std, **kwargs): self.mean = mean self.std = std - def sample(self, size=1): - return self.stream.normal(loc=self.mean, scale=self.std, size=size) + def sample(self, **kwargs): + return self.stream.normal(loc=self.mean, scale=self.std, **kwargs) class normal_pos(normal): @@ -143,8 +143,8 @@ class normal_pos(normal): WARNING - this function came from hpvsim but confirm that the implementation is correct? """ - def sample(self, size=1): - return np.abs(super().sample(size)) + def sample(self, **kwargs): + return np.abs(super().sample(**kwargs)) class normal_int(Distribution): @@ -152,8 +152,8 @@ class normal_int(Distribution): Normal distribution returning only integer values """ - def sample(self, size=1): - return np.round(super().sample(size)) + def sample(self, **kwargs): + return np.round(super().sample(**kwargs)) class lognormal(Distribution): @@ -172,12 +172,17 @@ def __init__(self, mean, std, **kwargs): 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, size=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 self.stream.lognormal(mean=self.underlying_mean, sigma=self.underlying_std, size=size) + return self.stream.lognormal(mean=self.underlying_mean, sigma=self.underlying_std, **kwargs) else: - return np.zeros(size) + if 'size' in kwargs: + return np.zeros(kwargs['size']) + elif 'uids' in kwargs: + return np.zeros(len(kwargs['uids'])) + else: + raise Exception('TODO') class lognormal_int(lognormal): @@ -185,8 +190,8 @@ class lognormal_int(lognormal): Lognormal returning only integer values """ - def sample(self, size=1): - return np.round(super().sample(size)) + def sample(self, **kwargs): + return np.round(super().sample(**kwargs)) class poisson(Distribution): @@ -201,8 +206,8 @@ def __init__(self, rate, **kwargs): def mean(self): return self.rate - def sample(self, size=1): - return self.stream.poisson(self.rate, size) + def sample(self, **kwargs): + return self.stream.poisson(self.rate, **kwargs) class neg_binomial(Distribution): @@ -227,10 +232,10 @@ def __init__(self, mean, dispersion, **kwargs): self.mean = mean self.dispersion = dispersion - def sample(self, size=1): + def sample(self, **kwargs): nbn_n = self.dispersion nbn_p = self.dispersion / (self.mean / self.step + self.dispersion) - return self.stream.negative_binomial(n=nbn_n, p=nbn_p, size=size) + return self.stream.negative_binomial(n=nbn_n, p=nbn_p, **kwargs) class beta(Distribution): @@ -246,8 +251,8 @@ def __init__(self, alpha, beta, **kwargs): def mean(self): return self.alpha / (self.alpha + self.beta) - def sample(self, size=1): - return self.stream.beta(a=self.alpha, b=self.beta, size=size) + def sample(self, **kwargs): + return self.stream.beta(a=self.alpha, b=self.beta, **kwargs) class gamma(Distribution): @@ -263,5 +268,5 @@ def __init__(self, shape, scale, **kwargs): def mean(self): return self.shape * self.scale - def sample(self, size=1): - return self.stream.gamma(shape=self.shape, scale=self.scale, size=size) \ No newline at end of file + def sample(self, **kwargs): + return self.stream.gamma(shape=self.shape, scale=self.scale, **kwargs) \ No newline at end of file diff --git a/stisim/hiv.py b/stisim/hiv.py index e41d5dc2..d052d738 100644 --- a/stisim/hiv.py +++ b/stisim/hiv.py @@ -44,7 +44,7 @@ def update_states(self, sim): hiv_death_prob = 0.05 / (self.pars.cd4_min - self.pars.cd4_max)**2 * (self.cd4 - self.pars.cd4_max)**2 can_die = ss.true(sim.people.alive & sim.people.hiv.infected) - hiv_deaths = self.rng_dead.bernoulli_filter(size=can_die, prob=hiv_death_prob[can_die]) + hiv_deaths = self.rng_dead.bernoulli_filter(uids=can_die, prob=hiv_death_prob[can_die]) sim.people.alive[hiv_deaths] = False sim.people.ti_dead[hiv_deaths] = sim.ti self.results['new_deaths'][sim.ti] = len(hiv_deaths) @@ -96,7 +96,7 @@ def apply(self, sim): coverage = self.coverage[np.where(self.t <= sim.ti)[0][-1]] recently_infected = ss.true((sim.people.hiv.ti_infected == sim.ti-2) & sim.people.alive) # 2 year (step) delay - inds = self.rng_add_ART.bernoulli_filter(size=recently_infected, prob=coverage) + inds = self.rng_add_ART.bernoulli_filter(uids=recently_infected, prob=coverage) sim.people.hiv.on_art[inds] = True sim.people.hiv.ti_art[inds] = sim.ti @@ -137,14 +137,14 @@ def apply(self, sim): eligible = ss.true(sim.people.alive & ~sim.people.hiv.infected & ~sim.people.hiv.on_prep) n_eligible = len(eligible) if n_eligible: - inds = self.rng_add_PrEP.bernoulli_filter(size=eligible, prob=min(n_eligible, n_change)/n_eligible) + inds = self.rng_add_PrEP.bernoulli_filter(uids=eligible, prob=min(n_eligible, n_change)/n_eligible) sim.people.hiv.on_prep[inds] = True elif n_change < 0: # Take some people off PrEP eligible = ss.true(sim.people.alive & ~sim.people.hiv.infected & sim.people.hiv.on_prep) n_eligible = len(eligible) if n_eligible: - inds = self.rng_remove_PrEP.bernoulli_filter(size=eligible, prob=-n_change/n_eligible) + inds = self.rng_remove_PrEP.bernoulli_filter(uids=eligible, prob=-n_change/n_eligible) sim.people.hiv.on_prep[inds] = False # Add result diff --git a/stisim/modules.py b/stisim/modules.py index e4b83b5a..a41e6af4 100644 --- a/stisim/modules.py +++ b/stisim/modules.py @@ -129,7 +129,7 @@ def set_initial_states(self, sim): if self.pars['initial'] <= 0: return - initial_cases = self.rng_init_cases.bernoulli_filter(size=sim.people.uid, prob=self.pars['initial']/len(sim.people)) + initial_cases = self.rng_init_cases.bernoulli_filter(uids=sim.people.uid, prob=self.pars['initial']/len(sim.people)) self.set_prognoses(sim, initial_cases, from_uids=None) # TODO: sentinel value to indicate seeds? return @@ -192,7 +192,7 @@ def make_new_cases(self, sim): node_from_node_this_layer_b_from_a[bi, ai] = p_not_acq node_from_node *= node_from_node_this_layer_b_from_a - new_cases_bool = self.rng_trans.bernoulli(size=sim.people.uid, prob=p_acq_node) + new_cases_bool = self.rng_trans.bernoulli(uids=sim.people.uid, prob=p_acq_node) new_cases = sim.people.uid[new_cases_bool] if not len(new_cases): @@ -200,7 +200,7 @@ def make_new_cases(self, sim): # Decide whom the infection came from using one random number for each b (aligned by block size) frm = np.zeros_like(new_cases) - r = self.rng_choose_infector.random( new_cases ) + r = self.rng_choose_infector.random(uids=new_cases) new_cases_idx = new_cases_bool.nonzero()[0] prob = (1-node_from_node[new_cases_idx]) # Prob of acquiring from each node | can constrain to just neighbors? cumsum = (prob / ((prob.sum(axis=1)[:,np.newaxis]))).cumsum(axis=1) diff --git a/stisim/networks.py b/stisim/networks.py index dae32f40..1f0bb734 100644 --- a/stisim/networks.py +++ b/stisim/networks.py @@ -375,8 +375,8 @@ def add_pairs(self, people, ti=None): print('No pairs to add') return 0 - loc_m = people.age[available_m].values - 5 + self.rng_pair_12.normal(size=available_m, std=3) - loc_f = people.age[available_f].values + self.rng_pair_21.normal(size=available_f, std=3) + loc_m = people.age[available_m].values - 5 + self.rng_pair_12.normal(uids=available_m, std=3) + loc_f = people.age[available_f].values + self.rng_pair_21.normal(uids=available_f, std=3) dist_mat = sps.distance_matrix(loc_m[:, np.newaxis], loc_f[:, np.newaxis]) ind_m, ind_f = spo.linear_sum_assignment(dist_mat) @@ -387,7 +387,7 @@ def add_pairs(self, people, ti=None): self['p2'] = np.concatenate([self['p2'], available_f[ind_f]]) beta = np.ones(n_pairs) - dur = self.rng_mean_dur.poisson(available_m[ind_m], self.mean_dur) + dur = self.rng_mean_dur.poisson(uids=available_m[ind_m], lam=self.mean_dur) self['beta'] = np.concatenate([self['beta'], beta]) self['dur'] = np.concatenate([self['dur'], dur]) diff --git a/stisim/people.py b/stisim/people.py index 3151188a..c089e0f1 100644 --- a/stisim/people.py +++ b/stisim/people.py @@ -225,7 +225,7 @@ def initialize(self, sim): # Initialize here in case other states use random streams that depend on slots being initialized self.slot[:] = np.copy(self.uid) # TODO: .values and [:] needed? - self.age[:] = self.age_data_dist.sample(len(self)) + self.age[:] = self.age_data_dist.sample(size=len(self)) self.initialized = True return diff --git a/stisim/streams.py b/stisim/streams.py index f0d0a9fe..d41ac85f 100644 --- a/stisim/streams.py +++ b/stisim/streams.py @@ -5,6 +5,10 @@ __all__ = ['Streams', 'MultiStream', 'CentralizedStream', 'Stream'] +SIZE = 0 +UIDS = 1 +BOOLS = 2 + class Streams: """ Class for managing a collection random number streams @@ -93,22 +97,49 @@ def Stream(*args, **kwargs): def _pre_draw(func): - def check_ready(self, *args, **kwargs): + def check_ready(self, **kwargs): """ Validation before drawing """ + if 'size' in kwargs and 'uids' in kwargs and not ((kwargs['size'] is None) ^ (kwargs['uids'] is None)): + raise Exception('Specify either "uids" or "size", but not both.') + # Check for zero length size if 'size' in kwargs.keys(): - size = kwargs['size'] - else: - size = args[0] - if isinstance(size, int): - # If an integer, the user wants "n" samples + # size-based + size = kwargs.pop('size') + + if not isinstance(size, int): + raise Exception('Input "size" must be an integer') + + if size < 0: + raise Exception('Input "size" cannot be negative') + if size == 0: return np.array([], dtype=int) # int dtype allows use as index, e.g. bernoulli_filter + + basis = SIZE + + else: + # uid-based + uids = kwargs['uids'] + + if len(uids) == 0: + return np.array([], dtype=int) # int dtype allows use as index, e.g. bernoulli_filter + + if isinstance(uids, ss.states.FusedArray): + v = uids.values + elif isinstance(uids, ss.states.DynamicView): + v = uids._view else: - kwargs['size'] = np.arange(size) - elif len(size) == 0: - return np.array([], dtype=int) # int dtype allows use as index, e.g. bernoulli_filter + v = uids + + if v.dtype == bool: + size = len(uids) + basis = BOOLS + else: + #size = self.slots.values[v.max()] + 1 + size = self.slots.values[v].max() + 1 + basis = UIDS if not self.initialized: msg = f'Stream {self.name} has not been initialized!' @@ -117,7 +148,9 @@ def check_ready(self, *args, **kwargs): msg = f'Stream {self.name} has already been sampled on this timestep!' raise NotResetException(msg) self.ready = False - return func(self, *args, **kwargs) + + return func(self, basis=basis, size=size, **kwargs) + return check_ready @@ -205,43 +238,96 @@ def draw_size(self, size): return v.max()+1 @_pre_draw - def random(self, size): - slots = self.slots.values[size] - return super(MultiStream, self).random(size=self.draw_size(slots))[slots] + def random(self, size, basis, uids=None): + if basis==SIZE: + return super(MultiStream, self).random(size=size) + elif basis == UIDS: + slots = self.slots.values[uids] + return super(MultiStream, self).random(size=size)[slots] + elif basis == BOOLS: + return super(MultiStream, self).random(size=size)[uids] + else: + raise Exception('TODO BASISEXCPETION') @_pre_draw - def uniform(self, size, **kwargs): - slots = self.slots.values[size] - return super(MultiStream, self).uniform(size=self.draw_size(slots), **kwargs)[slots] + def uniform(self, size, basis, low, high, uids=None): + if basis == SIZE: + return super(MultiStream, self).uniform(size=size, low=low, high=high) + elif basis == UIDS: + slots = self.slots.values[uids] + return super(MultiStream, self).uniform(size=size, low=low, high=high)[slots] + elif basis == BOOLS: + return super(MultiStream, self).uniform(size=size, low=low, high=high)[uids] + else: + raise Exception('TODO BASISEXCPETION') @_pre_draw - def poisson(self, size, lam): - slots = self.slots.values[size] - return super(MultiStream, self).poisson(size=self.draw_size(slots), lam=lam)[slots] + def integers(self, size, basis, low, high, uids=None, **kwargs): + if basis == SIZE: + return super(MultiStream, self).integers(size=size, low=low, high=high, **kwargs) + elif basis == UIDS: + slots = self.slots.values[uids] + return super(MultiStream, self).integers(size=size, low=low, high=high, **kwargs)[slots] + elif basis == BOOLS: + return super(MultiStream, self).integers(size=size, low=low, high=high, **kwargs)[uids] + else: + raise Exception('TODO BASISEXCPETION') @_pre_draw - def normal(self, size, mu=0, std=1): - slots = self.slots.values[size] - return mu + std*super(MultiStream, self).normal(size=self.draw_size(slots))[slots] + def poisson(self, size, basis, lam, uids=None): + if basis == SIZE: + return super(MultiStream, self).poisson(size=size, lam=lam) + elif basis == UIDS: + slots = self.slots.values[uids] + return super(MultiStream, self).poisson(size=size, lam=lam)[slots] + elif basis == BOOLS: + return super(MultiStream, self).poisson(size=size, lam=lam)[uids] + else: + raise Exception('TODO BASISEXCPETION') @_pre_draw - def negative_binomial(self, size, **kwargs): #n=nbn_n, p=nbn_p, size=n) - slots = self.slots.values[size] - return super(MultiStream, self).negative_binomial(size=self.draw_size(slots), **kwargs)[slots] + def normal(self, size, basis, mu=0, std=1, uids=None): + if basis == SIZE: + return mu + std*super(MultiStream, self).normal(size=size) + elif basis == UIDS: + slots = self.slots.values[uids] + return mu + std*super(MultiStream, self).normal(size=size)[slots] + elif basis == BOOLS: + return mu + std*super(MultiStream, self).normal(size=size)[uids] + else: + raise Exception('TODO BASISEXCPETION') @_pre_draw - def bernoulli(self, size, prob): + def negative_binomial(self, size, basis, n, p, uids=None): #n=nbn_n, p=nbn_p, size=n) + if basis == SIZE: + return super(MultiStream, self).negative_binomial(size=size, n=n, p=p) + elif basis == UIDS: + slots = self.slots.values[uids] + return super(MultiStream, self).negative_binomial(size=size, n=n, p=p)[slots] + elif basis == BOOLS: + return super(MultiStream, self).negative_binomial(size=size, n=n, p=p)[uids] + else: + raise Exception('TODO BASISEXCPETION') + + @_pre_draw + def bernoulli(self, prob, size, basis, uids=None): #return super(MultiStream, self).choice([True, False], size=size.max()+1) # very slow #return (super(MultiStream, self).binomial(n=1, p=prob, size=size.max()+1))[size].astype(bool) # pretty fast - slots = self.slots.values[size] - return super(MultiStream, self).random(size=self.draw_size(slots))[slots] < prob # fastest + if basis == SIZE: + return super(MultiStream, self).random(size=size) < prob # fastest + elif basis == UIDS: + slots = self.slots.values[uids] + return super(MultiStream, self).random(size=size)[slots] < prob # fastest + elif basis == BOOLS: + return super(MultiStream, self).random(size=size)[uids] < prob # fastest + else: + raise Exception('TODO BASISEXCPETION') # @_pre_draw <-- handled by call to self.bernoullli - def bernoulli_filter(self, size, prob): - #slots = self.slots[size[:]] - return size[self.bernoulli(size, prob)] # Slightly faster on my machine for bernoulli to typecast + def bernoulli_filter(self, uids, prob): + return uids[self.bernoulli(uids=uids, prob=prob)] - def choice(self, size, a, **kwargs): + def choice(self, size, basis, a, **kwargs): # Consider raising a warning instead? raise NotStreamSafeException('The "choice" function is not MultiStream-safe.') diff --git a/tests/run_multistream.py b/tests/run_multistream.py index 86b649ee..a9ed8e25 100644 --- a/tests/run_multistream.py +++ b/tests/run_multistream.py @@ -10,8 +10,8 @@ import seaborn as sns import numpy as np -n = 250 # Agents -n_rand_seeds = 100 +n = 100 # Agents +n_rand_seeds = 10 intv_cov_levels = [0.025, 0.05, 0.10, 0.73] + [0] # Must include 0 as that's the baseline # Choose ART or PrEP diff --git a/tests/run_stable_monogamy.py b/tests/run_stable_monogamy.py index 2d9fb400..a8c44651 100644 --- a/tests/run_stable_monogamy.py +++ b/tests/run_stable_monogamy.py @@ -12,7 +12,6 @@ import os import argparse import sciris as sc -import seaborn as sns ss.options(multistream = True) # Can set multistream to False for comparison plot_graph = True @@ -89,8 +88,7 @@ def finalize(self, sim): return - -def run_sim(n=25, intervention=False, analyze=False): +def run_sim(n=25, rand_seed=0, intervention=False, analyze=False): ppl = ss.People(n) ppl.networks = ss.ndict(ss.simple_embedding(mean_dur=5))#, ss.maternal()) @@ -109,7 +107,7 @@ def run_sim(n=25, intervention=False, analyze=False): 'end': 2020, 'remove_dead': False, # So we can see who dies, sim results should not change with True 'interventions': [art] if intervention else [], - 'rand_seed': 0, + 'rand_seed': rand_seed, 'analyzers': [GraphAnalyzer()] if analyze else [], 'n_agents': len(ppl), # TODO: Build into Sim } @@ -123,8 +121,10 @@ def run_sim(n=25, intervention=False, analyze=False): return sim -def run_scenario(n=10, analyze=True): - sims = sc.parallelize(run_sim, kwargs={'n':n, 'analyze': analyze}, iterkwargs=[{'intervention':True}, {'intervention':False}], die=True) +def run_scenario(n=10, rand_seed=0, analyze=True): + sims = sc.parallelize(run_sim, + kwargs={'n':n, 'analyze': analyze, 'rand_seed': rand_seed}, + iterkwargs=[{'intervention':True}, {'intervention':False}], die=True) for i, sim in enumerate(sims): sim.save(os.path.join(figdir, f'sim{i}.obj')) @@ -191,18 +191,24 @@ def on_press(event): def analyze_people(sim): p = sim.people ever_alive = ss.false(np.isnan(p.age)) - years_lived = np.full(len(p), sim.ti) + years_lived = np.full(len(p), sim.ti+1) # Actually +1 dt here, I think years_lived[p.dead] = p.ti_dead[p.dead] years_lived = years_lived[ever_alive] # Trim, could be more efficient age_initial = p.age[ever_alive].values - years_lived + age_initial = age_initial.astype(np.float32) # For better hash comparability, there are small differences at float64 df = pd.DataFrame({ - 'uid': p.uid[ever_alive], # if slicing, don't need ._view, + #'uid': p.uid[ever_alive], # if slicing, don't need ._view, + 'id': [hash((p.slot[i], age_initial[i], p.female[i])) for i in ever_alive], # if slicing, don't need ._view, 'age_initial': age_initial, 'years_lived': years_lived, 'ti_infected': p.hiv.ti_infected[ever_alive].values, 'ti_art': p.hiv.ti_art[ever_alive].values, 'ti_dead': p.ti_dead[ever_alive].values, + + # Useful for debugging, but not needed for plotting + 'slot': p.slot[ever_alive].values, + 'female': p.female[ever_alive].values, }) df.replace(to_replace=ss.INT_NAN, value=np.nan, inplace=True) df['age_infected'] = df['age_initial'] + df['ti_infected'] @@ -231,36 +237,15 @@ def life_bars(data, **kwargs): return -def life_bars_nested(df): - - N = df['sim'].nunique() - height = 0.9/N - - fig, ax = plt.subplots(figsize=(10,6)) - - for n, (lbl, data) in enumerate(df.groupby('sim')): - ys = n/(N+1) # Leave space - - age_final = data['age_initial'] + data['years_lived'] - plt.barh(y=data.index + ys, left=data['age_initial'], width=age_final-data['age_initial'], color='k', height=height) - - # Define bools - infected = ~data['ti_infected'].isna() - art = ~data['ti_art'].isna() - dead = ~data['ti_dead'].isna() - - # Infected - plt.barh(y=data.index[infected] + ys, left=data.loc[infected]['age_infected'], width=age_final[infected]-data.loc[infected]['age_infected'], color='r', height=height) - - # ART - plt.barh(y=data.index[art] + ys, left=data.loc[art]['age_art'], width=age_final[art]-data.loc[art]['age_art'], color='g', height=height) - - # Dead - plt.scatter(y=data.index[dead] + ys, x=data.loc[dead]['age_dead'], color='k', marker='|') +def plot_longitudinal(sim1, sim2): - return fig + df1 = analyze_people(sim1) + df1['sim'] = 'Baseline' + df2 = analyze_people(sim2) + df2['sim'] = 'With ART' -def ti_bars_nested(df): + df = pd.concat([df1, df2]).set_index('id') + #f = ti_bars_nested(df) df['ypos'] = pd.factorize(df.index.values)[0] N = df['sim'].nunique() @@ -275,53 +260,32 @@ def ti_bars_nested(df): ti_final = data['ti_dead'].fillna(40) plt.barh(y=yp, left=ti_initial, width=ti_final - ti_initial, color='k', height=height) - # Define bools - infected = ~data['ti_infected'].isna() - art = ~data['ti_art'].isna() - dead = ~data['ti_dead'].isna() + # Infected before birth + vertical = data['age_infected']<0 + plt.barh(y=yp[vertical], left=data.loc[vertical]['ti_infected'], width=ti_final[vertical]-data.loc[vertical]['ti_infected'], color='m', height=height) # Infected - plt.barh(y=yp[infected], left=data.loc[infected]['ti_infected'], width=ti_final[infected]-data.loc[infected]['ti_infected'], color='r', height=height) + infected = ~data['ti_infected'].isna() + ai = data.loc[infected]['age_infected'].values # Adjust for vertical transmission + ai[~(ai<0)] = 0 + plt.barh(y=yp[infected], left=data.loc[infected]['ti_infected']-ai, width=ti_final[infected]-data.loc[infected]['ti_infected']+ai, color='r', height=height) # ART + art = ~data['ti_art'].isna() plt.barh(y=yp[art], left=data.loc[art]['ti_art'], width=ti_final[art]-data.loc[art]['ti_art'], color='g', height=height) # Dead + dead = ~data['ti_dead'].isna() plt.scatter(y=yp[dead], x=data.loc[dead]['ti_dead'], color='k', marker='|') return fig -def plot_longitudinal(sim1, sim2): - - df1 = analyze_people(sim1) - df1['sim'] = 'Baseline' - df2 = analyze_people(sim2) - df2['sim'] = 'With ART' - - df = pd.concat([df1, df2]).set_index('uid') - #age_vars = ['age_initial', 'age_infected', 'age_art', 'age_dead'] - #age_vars = ['age_end', 'age_dead', 'age_art', 'age_infected', 'age_initial'] - #dfm = df.melt(id_vars=['uid', 'sim'], value_vars=age_vars, var_name='event', value_name='age') - #g = sns.catplot(kind='bar', col='sim', data=dfm.reset_index(), y='uid', x='age', hue='event', orient='h', hue_order=age_vars, dodge=False, - # palette=sns.blend_palette(colors=['black', 'green', 'red', 'gray', 'white'], n_colors=len(age_vars))) - - #f = life_bars_nested(df) - f = ti_bars_nested(df) - - #g = sns.FacetGrid(data=df, col='sim') - #g.map_dataframe(life_bars) - - - return - - - - if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('-p', '--plot', help='Plot from a cached CSV file', type=str) parser.add_argument('-n', help='Number of agents', type=int, default=100) + parser.add_argument('-s', help='Rand seed', type=int, default=2) args = parser.parse_args() if args.plot: @@ -330,7 +294,7 @@ def plot_longitudinal(sim1, sim2): sim2 = sc.load(os.path.join(args.plot, 'sim2.obj')) else: print('Running scenarios') - [sim1, sim2] = run_scenario(n=args.n) + [sim1, sim2] = run_scenario(n=args.n, rand_seed=args.s) #plot_graph(sim1, sim2) plot_longitudinal(sim1, sim2) From efc1471dee64907bc01e38ab733e099d1dc2509d Mon Sep 17 00:00:00 2001 From: Daniel Klein Date: Tue, 17 Oct 2023 09:14:31 -0700 Subject: [PATCH 28/34] It works! --- stisim/streams.py | 113 +++++++++++++++++++++++---------------- tests/run_multistream.py | 10 ++-- 2 files changed, 72 insertions(+), 51 deletions(-) diff --git a/stisim/streams.py b/stisim/streams.py index d41ac85f..e71e27dd 100644 --- a/stisim/streams.py +++ b/stisim/streams.py @@ -222,21 +222,6 @@ def step(self, ti): return - def draw_size(self, size): - """ Determine how many random numbers to draw for a given size """ - - if isinstance(size, ss.states.FusedArray): - v = size.values - elif isinstance(size, ss.states.DynamicView): - v = size._view - else: - v = size - - if v.dtype == bool: - return len(size) - - return v.max()+1 - @_pre_draw def random(self, size, basis, uids=None): if basis==SIZE: @@ -332,6 +317,48 @@ def choice(self, size, basis, a, **kwargs): raise NotStreamSafeException('The "choice" function is not MultiStream-safe.') +def _pre_draw_centralized(func): + def check_ready(self, **kwargs): + """ Validation before drawing """ + + uids = None + if 'uids' in kwargs: + uids = kwargs.pop('uids') + + size = None + if 'size' in kwargs: + size = kwargs.pop('size') + + if not ((size is None) ^ (uids is None)): + raise Exception('Specify either "uids" or "size", but not both.') + + # Check for zero length size + if size is not None: + # size-based + if not isinstance(size, int): + raise Exception('Input "size" must be an integer') + + if size < 0: + raise Exception('Input "size" cannot be negative') + + if size == 0: + return np.array([], dtype=int) # int dtype allows use as index, e.g. bernoulli_filter + + else: + # uid-based + if len(uids) == 0: + return np.array([], dtype=int) # int dtype allows use as index, e.g. bernoulli_filter + + size = len(uids) + + if not self.initialized: + msg = f'Stream {self.name} has not been initialized!' + raise NotInitializedException(msg) + + return func(self, size=size, **kwargs) + + return check_ready + class CentralizedStream(np.random.Generator): """ Class to imitate the behavior of a centralized random number generator @@ -351,7 +378,7 @@ def __init__(self, name, seed_offset=None, **kwargs): self.seed_offset = None # Not used, so override to avoid potential seed collisions in Streams. return - def initialize(self, streams): + def initialize(self, streams, slots=None): if self.initialized: # TODO: Raise warning assert not self.initialized @@ -367,43 +394,37 @@ def reset(self): def step(self, ti): pass - def draw_size(self, size): - """ Determine how many random numbers to draw for a given size """ - - if isinstance(size, int): - return size - elif isinstance(size, ss.states.FusedArray): - v = size.values - elif isinstance(size, ss.states.DynamicView): - v = size._view - else: - v = size - - if v.dtype == bool: - return size.sum() - - return len(size) + @_pre_draw_centralized + def random(self, size, **kwargs): + return np.random.random(size=size, **kwargs) - def random(self, size): - return np.random.random(self.draw_size(size)) + @_pre_draw_centralized + def uniform(self, size, low, high, **kwargs): + return np.random.uniform(size=size, low=low, high=high, **kwargs) - def uniform(self, size, **kwargs): - return np.random.uniform(size=self.draw_size(size), **kwargs) + @_pre_draw_centralized + def integers(self, size, low, high, **kwargs): + return np.random.random_integers(size=size, low=low, high=high) - def poisson(self, size, lam): - return np.random.poisson(lam=lam, size=self.draw_size(size)) + @_pre_draw_centralized + def poisson(self, size, lam, **kwargs): + return np.random.poisson(lam=lam, size=size, **kwargs) + @_pre_draw_centralized def normal(self, size, mu=0, std=1): - return mu + std*np.random.normal(size=self.draw_size(size), loc=mu, scale=std) + return mu + std*np.random.normal(size=size, loc=mu, scale=std) - def negative_binomial(self, size, **kwargs): #n=nbn_n, p=nbn_p, size=n) - return np.random.negative_binomial(**kwargs, size=self.draw_size(size)) + @_pre_draw_centralized + def negative_binomial(self, size, n, p, **kwargs): + return np.random.negative_binomial(size=size, n=n, p=p, **kwargs) - def bernoulli(self, size, prob): - return np.random.random(self.draw_size(size)) < prob + @_pre_draw_centralized + def bernoulli(self, prob, size, **kwargs): + return np.random.random(size=size, **kwargs) < prob - def bernoulli_filter(self, size, prob): - return size[self.bernoulli(size, prob)] # Slightly faster on my machine for bernoulli to typecast + def bernoulli_filter(self, uids, prob): + return uids[self.bernoulli(uids=uids, prob=prob)] + @_pre_draw_centralized def choice(self, size, a, **kwargs): - return np.random.choice(a, size=self.draw_size(size), **kwargs) \ No newline at end of file + return np.random.choice(a, size=size, **kwargs) \ No newline at end of file diff --git a/tests/run_multistream.py b/tests/run_multistream.py index a9ed8e25..70ca0f76 100644 --- a/tests/run_multistream.py +++ b/tests/run_multistream.py @@ -10,8 +10,8 @@ import seaborn as sns import numpy as np -n = 100 # Agents -n_rand_seeds = 10 +n = 1000 # Agents +n_rand_seeds = 100 intv_cov_levels = [0.025, 0.05, 0.10, 0.73] + [0] # Must include 0 as that's the baseline # Choose ART or PrEP @@ -49,7 +49,7 @@ def run_sim(n, idx, intv_cov, rand_seed, multistream): 'interventions': [intv], 'rand_seed': rand_seed, 'verbose': 0, - 'remove_dead': False, # PROBLEM + 'remove_dead': False, 'n_agents': len(ppl), # TODO } sim = ss.Sim(people=ppl, diseases=[hiv], demographics=[pregnancy], pars=pars, label=f'Sim with {n} agents and intv_cov={intv_cov}') # PROBLEM @@ -75,13 +75,13 @@ def run_sim(n, idx, intv_cov, rand_seed, multistream): def run_scenarios(): results = [] times = {} - for multistream in [True, False]: + for multistream in [False, True]: cfgs = [] for rs in range(n_rand_seeds): for intv_cov in intv_cov_levels: cfgs.append({'intv_cov':intv_cov, 'rand_seed':rs, 'multistream':multistream, 'idx':len(cfgs)}) T = sc.tic() - results += sc.parallelize(run_sim, kwargs={'n': n}, iterkwargs=cfgs, die=True, serial=True) + results += sc.parallelize(run_sim, kwargs={'n': n}, iterkwargs=cfgs, die=True, serial=False) times[f'Multistream={multistream}'] = sc.toc(T, output=True) print('Timings:', times) From b9f9dfa32cd670ce03bc3ba76ad9c3935c6eb465 Mon Sep 17 00:00:00 2001 From: Daniel Klein Date: Tue, 17 Oct 2023 09:49:41 -0700 Subject: [PATCH 29/34] Small adjustment to slotting, all working now. --- stisim/streams.py | 24 ++++++++++++++++-------- tests/run_multistream.py | 17 ++++++++--------- 2 files changed, 24 insertions(+), 17 deletions(-) diff --git a/stisim/streams.py b/stisim/streams.py index e71e27dd..7c27df38 100644 --- a/stisim/streams.py +++ b/stisim/streams.py @@ -138,7 +138,8 @@ def check_ready(self, **kwargs): basis = BOOLS else: #size = self.slots.values[v.max()] + 1 - size = self.slots.values[v].max() + 1 + #size = self.slots.values[v].max() + 1 + size = self.slots[v].values.max() + 1 basis = UIDS if not self.initialized: @@ -227,7 +228,8 @@ def random(self, size, basis, uids=None): if basis==SIZE: return super(MultiStream, self).random(size=size) elif basis == UIDS: - slots = self.slots.values[uids] + #slots = self.slots.values[uids] + slots = self.slots[uids].values return super(MultiStream, self).random(size=size)[slots] elif basis == BOOLS: return super(MultiStream, self).random(size=size)[uids] @@ -239,7 +241,8 @@ def uniform(self, size, basis, low, high, uids=None): if basis == SIZE: return super(MultiStream, self).uniform(size=size, low=low, high=high) elif basis == UIDS: - slots = self.slots.values[uids] + #slots = self.slots.values[uids] + slots = self.slots[uids].values return super(MultiStream, self).uniform(size=size, low=low, high=high)[slots] elif basis == BOOLS: return super(MultiStream, self).uniform(size=size, low=low, high=high)[uids] @@ -251,7 +254,8 @@ def integers(self, size, basis, low, high, uids=None, **kwargs): if basis == SIZE: return super(MultiStream, self).integers(size=size, low=low, high=high, **kwargs) elif basis == UIDS: - slots = self.slots.values[uids] + #slots = self.slots.values[uids] + slots = self.slots[uids].values return super(MultiStream, self).integers(size=size, low=low, high=high, **kwargs)[slots] elif basis == BOOLS: return super(MultiStream, self).integers(size=size, low=low, high=high, **kwargs)[uids] @@ -263,7 +267,8 @@ def poisson(self, size, basis, lam, uids=None): if basis == SIZE: return super(MultiStream, self).poisson(size=size, lam=lam) elif basis == UIDS: - slots = self.slots.values[uids] + #slots = self.slots.values[uids] + slots = self.slots[uids].values return super(MultiStream, self).poisson(size=size, lam=lam)[slots] elif basis == BOOLS: return super(MultiStream, self).poisson(size=size, lam=lam)[uids] @@ -275,7 +280,8 @@ def normal(self, size, basis, mu=0, std=1, uids=None): if basis == SIZE: return mu + std*super(MultiStream, self).normal(size=size) elif basis == UIDS: - slots = self.slots.values[uids] + #slots = self.slots.values[uids] + slots = self.slots[uids].values return mu + std*super(MultiStream, self).normal(size=size)[slots] elif basis == BOOLS: return mu + std*super(MultiStream, self).normal(size=size)[uids] @@ -287,7 +293,8 @@ def negative_binomial(self, size, basis, n, p, uids=None): #n=nbn_n, p=nbn_p, si if basis == SIZE: return super(MultiStream, self).negative_binomial(size=size, n=n, p=p) elif basis == UIDS: - slots = self.slots.values[uids] + #slots = self.slots.values[uids] + slots = self.slots[uids].values return super(MultiStream, self).negative_binomial(size=size, n=n, p=p)[slots] elif basis == BOOLS: return super(MultiStream, self).negative_binomial(size=size, n=n, p=p)[uids] @@ -301,7 +308,8 @@ def bernoulli(self, prob, size, basis, uids=None): if basis == SIZE: return super(MultiStream, self).random(size=size) < prob # fastest elif basis == UIDS: - slots = self.slots.values[uids] + #slots = self.slots.values[uids] + slots = self.slots[uids].values return super(MultiStream, self).random(size=size)[slots] < prob # fastest elif basis == BOOLS: return super(MultiStream, self).random(size=size)[uids] < prob # fastest diff --git a/tests/run_multistream.py b/tests/run_multistream.py index 70ca0f76..3c8a1dc9 100644 --- a/tests/run_multistream.py +++ b/tests/run_multistream.py @@ -11,7 +11,7 @@ import numpy as np n = 1000 # Agents -n_rand_seeds = 100 +n_rand_seeds = 250 intv_cov_levels = [0.025, 0.05, 0.10, 0.73] + [0] # Must include 0 as that's the baseline # Choose ART or PrEP @@ -29,13 +29,13 @@ def run_sim(n, idx, intv_cov, rand_seed, multistream): ppl = ss.People(n) ppl.networks = ss.ndict( - ss.simple_embedding(mean_dur=4), - ss.maternal() # PROBLEM + ss.simple_embedding(mean_dur=5), + ss.maternal() ) hiv_pars = { #'beta': {'simple_embedding': [0.10, 0.08], 'maternal': [0.2, 0]}, - 'beta': {'simple_embedding': [0.2, 0.15], 'maternal': [0.2, 0]}, + 'beta': {'simple_embedding': [0.25, 0.20], 'maternal': [0.2, 0]}, 'initial': int(np.maximum(10, np.ceil(0.01*n))), } hiv = ss.HIV(hiv_pars) @@ -49,11 +49,10 @@ def run_sim(n, idx, intv_cov, rand_seed, multistream): 'interventions': [intv], 'rand_seed': rand_seed, 'verbose': 0, - 'remove_dead': False, + 'remove_dead': True, 'n_agents': len(ppl), # TODO } - sim = ss.Sim(people=ppl, diseases=[hiv], demographics=[pregnancy], pars=pars, label=f'Sim with {n} agents and intv_cov={intv_cov}') # PROBLEM - ########################sim = ss.Sim(people=ppl, diseases=[hiv], demographics=None, pars=pars, label=f'Sim with {n} agents and intv_cov={intv_cov}') # FIX + sim = ss.Sim(people=ppl, diseases=[hiv], demographics=[pregnancy], pars=pars, label=f'Sim with {n} agents and intv_cov={intv_cov}') sim.initialize() sim.run() @@ -62,7 +61,7 @@ def run_sim(n, idx, intv_cov, rand_seed, multistream): #'hiv.n_infected': sim.results.hiv.n_infected, 'hiv.prevalence': sim.results.hiv.prevalence, 'hiv.cum_deaths': sim.results.hiv.new_deaths.cumsum(), - 'pregnancy.cum_births': sim.results.pregnancy.births.cumsum(), # PROBLEM + 'pregnancy.cum_births': sim.results.pregnancy.births.cumsum(), }) df['intv_cov'] = intv_cov df['rand_seed'] = rand_seed @@ -75,7 +74,7 @@ def run_sim(n, idx, intv_cov, rand_seed, multistream): def run_scenarios(): results = [] times = {} - for multistream in [False, True]: + for multistream in [True, False]: cfgs = [] for rs in range(n_rand_seeds): for intv_cov in intv_cov_levels: From 5dd9e0cf46a5c4e596fc2a72e20c3de53c93c4d5 Mon Sep 17 00:00:00 2001 From: Daniel Klein Date: Thu, 26 Oct 2023 13:45:16 -0700 Subject: [PATCH 30/34] Various improvements and cleanup including * Catching edge cases * Update and cleanup of test scripts * Removing prep * Making art_efficacy a parameter * Removed step from negative binomial --- .gitignore | 2 + stisim/distributions.py | 2 +- stisim/hiv.py | 55 +------------ stisim/streams.py | 46 +++++------ tests/run_multistream.py | 32 ++++---- tests/run_stable_monogamy.py | 147 ++++++++++++++++++++++------------- tests/test_stream.py | 76 +++++++----------- tests/test_streams.py | 58 +++++++------- 8 files changed, 189 insertions(+), 229 deletions(-) diff --git a/.gitignore b/.gitignore index 34a00e22..900997cb 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,8 @@ # HPVsim-generated files *.obj files/ +tests/figs/* +figs/* modules.rst # Other files diff --git a/stisim/distributions.py b/stisim/distributions.py index 54accc85..e926ff61 100644 --- a/stisim/distributions.py +++ b/stisim/distributions.py @@ -234,7 +234,7 @@ def __init__(self, mean, dispersion, **kwargs): def sample(self, **kwargs): nbn_n = self.dispersion - nbn_p = self.dispersion / (self.mean / self.step + self.dispersion) + nbn_p = self.dispersion / (self.mean + self.dispersion) return self.stream.negative_binomial(n=nbn_n, p=nbn_p, **kwargs) diff --git a/stisim/hiv.py b/stisim/hiv.py index d052d738..2b5a5089 100644 --- a/stisim/hiv.py +++ b/stisim/hiv.py @@ -19,7 +19,6 @@ def __init__(self, pars=None): self.ti_infected = ss.State('ti_infected', int, ss.INT_NAN) self.on_art = ss.State('on_art', bool, False) self.ti_art = ss.State('ti_art', int, ss.INT_NAN) - self.on_prep = ss.State('on_prep', bool, False) self.cd4 = ss.State('cd4', float, 500) self.rng_dead = ss.Stream(f'dead_{self.name}') @@ -30,6 +29,7 @@ def __init__(self, pars=None): 'cd4_rate': 5, 'initial': 30, 'eff_condoms': 0.7, + 'art_efficacy': 0.96, }, self.pars) return @@ -39,8 +39,7 @@ def update_states(self, sim): self.cd4[sim.people.alive & self.infected & self.on_art] += (self.pars.cd4_max - self.cd4[sim.people.alive & self.infected & self.on_art])/self.pars.cd4_rate self.cd4[sim.people.alive & self.infected & ~self.on_art] += (self.pars.cd4_min - self.cd4[sim.people.alive & self.infected & ~self.on_art])/self.pars.cd4_rate - self.rel_sus[sim.people.alive & ~self.infected & self.on_prep] = 0.04 - self.rel_trans[sim.people.alive & self.infected & self.on_art] = 0.04 + self.rel_trans[sim.people.alive & self.infected & self.on_art] = 1 - self.pars['art_efficacy'] hiv_death_prob = 0.05 / (self.pars.cd4_min - self.pars.cd4_max)**2 * (self.cd4 - self.pars.cd4_max)**2 can_die = ss.true(sim.people.alive & sim.people.hiv.infected) @@ -95,7 +94,8 @@ def apply(self, sim): return coverage = self.coverage[np.where(self.t <= sim.ti)[0][-1]] - recently_infected = ss.true((sim.people.hiv.ti_infected == sim.ti-2) & sim.people.alive) # 2 year (step) delay + ti_delay = 1 # 1 time step delay + recently_infected = ss.true((sim.people.hiv.ti_infected == sim.ti-ti_delay) & sim.people.alive) inds = self.rng_add_ART.bernoulli_filter(uids=recently_infected, prob=coverage) sim.people.hiv.on_art[inds] = True sim.people.hiv.ti_art[inds] = sim.ti @@ -105,53 +105,6 @@ def apply(self, sim): return len(inds) -class PrEP(ss.Intervention): - - def __init__(self, t: np.array, coverage: np.array, **kwargs): - self.requires = HIV - self.t = sc.promotetoarray(t) - self.coverage = sc.promotetoarray(coverage) - - super().__init__(**kwargs) - - self.rng_add_PrEP = ss.Stream('add_PrEP') - self.rng_remove_PrEP = ss.Stream('remove_PrEP') - - return - - def initialize(self, sim): - sim.results.hiv += ss.Result(self.name, 'n_prep', sim.npts, dtype=int) - self.initialized = True - return - - def apply(self, sim): - if sim.ti < self.t[0]: - return - - coverage = self.coverage[np.where(self.t <= sim.ti)[0][-1]] - on_prep = sim.people.alive & sim.people.hiv.on_prep - sus = sim.people.alive & ~sim.people.hiv.infected - n_change = np.round(coverage * sus.sum() - on_prep.sum()).astype(int) - if n_change > 0: - # Add more PrEP - eligible = ss.true(sim.people.alive & ~sim.people.hiv.infected & ~sim.people.hiv.on_prep) - n_eligible = len(eligible) - if n_eligible: - inds = self.rng_add_PrEP.bernoulli_filter(uids=eligible, prob=min(n_eligible, n_change)/n_eligible) - sim.people.hiv.on_prep[inds] = True - elif n_change < 0: - # Take some people off PrEP - eligible = ss.true(sim.people.alive & ~sim.people.hiv.infected & sim.people.hiv.on_prep) - n_eligible = len(eligible) - if n_eligible: - inds = self.rng_remove_PrEP.bernoulli_filter(uids=eligible, prob=-n_change/n_eligible) - sim.people.hiv.on_prep[inds] = False - - # Add result - sim.results.hiv.n_prep = np.count_nonzero(sim.people.alive & sim.people.hiv.on_prep) - - return - #%% Analyzers diff --git a/stisim/streams.py b/stisim/streams.py index 7c27df38..1a9767b7 100644 --- a/stisim/streams.py +++ b/stisim/streams.py @@ -137,9 +137,7 @@ def check_ready(self, **kwargs): size = len(uids) basis = BOOLS else: - #size = self.slots.values[v.max()] + 1 - #size = self.slots.values[v].max() + 1 - size = self.slots[v].values.max() + 1 + size = self.slots[v].__array__().max() + 1 basis = UIDS if not self.initialized: @@ -192,13 +190,17 @@ def initialize(self, streams, slots): return self.seed = streams.add(self) - self.slots = slots # E.g. sim.people.slots (instead of using uid as the slots directly) + + if isinstance(slots, int): + # Handle edge case in which the user wants n sequential slots, as used in testing. + self.slots = np.arange(slots) + else: + self.slots = slots # E.g. sim.people.slots (instead of using uid as the slots directly) if 'bit_generator' not in self.kwargs: self.kwargs['bit_generator'] = np.random.PCG64(seed=self.seed) super().__init__(**self.kwargs) - #self.rng = np.random.default_rng(seed=self.seed + self.seed_offset) self._init_state = self.bit_generator.state # Store the initial state self.initialized = True @@ -228,8 +230,7 @@ def random(self, size, basis, uids=None): if basis==SIZE: return super(MultiStream, self).random(size=size) elif basis == UIDS: - #slots = self.slots.values[uids] - slots = self.slots[uids].values + slots = self.slots[uids].__array__() return super(MultiStream, self).random(size=size)[slots] elif basis == BOOLS: return super(MultiStream, self).random(size=size)[uids] @@ -241,8 +242,7 @@ def uniform(self, size, basis, low, high, uids=None): if basis == SIZE: return super(MultiStream, self).uniform(size=size, low=low, high=high) elif basis == UIDS: - #slots = self.slots.values[uids] - slots = self.slots[uids].values + slots = self.slots[uids].__array__() return super(MultiStream, self).uniform(size=size, low=low, high=high)[slots] elif basis == BOOLS: return super(MultiStream, self).uniform(size=size, low=low, high=high)[uids] @@ -254,8 +254,7 @@ def integers(self, size, basis, low, high, uids=None, **kwargs): if basis == SIZE: return super(MultiStream, self).integers(size=size, low=low, high=high, **kwargs) elif basis == UIDS: - #slots = self.slots.values[uids] - slots = self.slots[uids].values + slots = self.slots[uids].__array__() return super(MultiStream, self).integers(size=size, low=low, high=high, **kwargs)[slots] elif basis == BOOLS: return super(MultiStream, self).integers(size=size, low=low, high=high, **kwargs)[uids] @@ -267,8 +266,7 @@ def poisson(self, size, basis, lam, uids=None): if basis == SIZE: return super(MultiStream, self).poisson(size=size, lam=lam) elif basis == UIDS: - #slots = self.slots.values[uids] - slots = self.slots[uids].values + slots = self.slots[uids].__array__() return super(MultiStream, self).poisson(size=size, lam=lam)[slots] elif basis == BOOLS: return super(MultiStream, self).poisson(size=size, lam=lam)[uids] @@ -280,8 +278,7 @@ def normal(self, size, basis, mu=0, std=1, uids=None): if basis == SIZE: return mu + std*super(MultiStream, self).normal(size=size) elif basis == UIDS: - #slots = self.slots.values[uids] - slots = self.slots[uids].values + slots = self.slots[uids].__array__() return mu + std*super(MultiStream, self).normal(size=size)[slots] elif basis == BOOLS: return mu + std*super(MultiStream, self).normal(size=size)[uids] @@ -289,12 +286,11 @@ def normal(self, size, basis, mu=0, std=1, uids=None): raise Exception('TODO BASISEXCPETION') @_pre_draw - def negative_binomial(self, size, basis, n, p, uids=None): #n=nbn_n, p=nbn_p, size=n) + def negative_binomial(self, size, basis, n, p, uids=None): if basis == SIZE: return super(MultiStream, self).negative_binomial(size=size, n=n, p=p) elif basis == UIDS: - #slots = self.slots.values[uids] - slots = self.slots[uids].values + slots = self.slots[uids].__array__() return super(MultiStream, self).negative_binomial(size=size, n=n, p=p)[slots] elif basis == BOOLS: return super(MultiStream, self).negative_binomial(size=size, n=n, p=p)[uids] @@ -303,13 +299,10 @@ def negative_binomial(self, size, basis, n, p, uids=None): #n=nbn_n, p=nbn_p, si @_pre_draw def bernoulli(self, prob, size, basis, uids=None): - #return super(MultiStream, self).choice([True, False], size=size.max()+1) # very slow - #return (super(MultiStream, self).binomial(n=1, p=prob, size=size.max()+1))[size].astype(bool) # pretty fast if basis == SIZE: return super(MultiStream, self).random(size=size) < prob # fastest elif basis == UIDS: - #slots = self.slots.values[uids] - slots = self.slots[uids].values + slots = self.slots[uids].__array__() return super(MultiStream, self).random(size=size)[slots] < prob # fastest elif basis == BOOLS: return super(MultiStream, self).random(size=size)[uids] < prob # fastest @@ -357,7 +350,10 @@ def check_ready(self, **kwargs): if len(uids) == 0: return np.array([], dtype=int) # int dtype allows use as index, e.g. bernoulli_filter - size = len(uids) + if uids.dtype == bool: + size = uids.sum() + else: + size = len(uids) if not self.initialized: msg = f'Stream {self.name} has not been initialized!' @@ -367,7 +363,8 @@ def check_ready(self, **kwargs): return check_ready -class CentralizedStream(np.random.Generator): + +class CentralizedStream(): """ Class to imitate the behavior of a centralized random number generator """ @@ -380,7 +377,6 @@ def __init__(self, name, seed_offset=None, **kwargs): name: a name for this Stream, like "coin_flip" """ - super().__init__(bit_generator=np.random.PCG64()) self.name = name self.initialized = False self.seed_offset = None # Not used, so override to avoid potential seed collisions in Streams. diff --git a/tests/run_multistream.py b/tests/run_multistream.py index 3c8a1dc9..423998b0 100644 --- a/tests/run_multistream.py +++ b/tests/run_multistream.py @@ -10,55 +10,52 @@ import seaborn as sns import numpy as np -n = 1000 # Agents -n_rand_seeds = 250 -intv_cov_levels = [0.025, 0.05, 0.10, 0.73] + [0] # Must include 0 as that's the baseline +n = 100 # Agents +n_rand_seeds = 25 +intv_cov_levels = [0.01, 0.10, 0.25, 0.73] + [0] # Must include 0 as that's the baseline -# Choose ART or PrEP -choice = 'ART' -intervention = {'ART': ss.hiv.ART, 'PrEP': ss.hiv.PrEP} - -figdir = os.path.join(os.getcwd(), 'figs', choice) +figdir = os.path.join(os.getcwd(), 'figs', 'ART') sc.path(figdir).mkdir(parents=True, exist_ok=True) def run_sim(n, idx, intv_cov, rand_seed, multistream): print(f'Starting sim {idx} with rand_seed={rand_seed} and intv_cov={intv_cov}, multistream={multistream}') - ss.options(multistream=multistream) ppl = ss.People(n) ppl.networks = ss.ndict( - ss.simple_embedding(mean_dur=5), + ss.simple_embedding(mean_dur=4), ss.maternal() ) hiv_pars = { - #'beta': {'simple_embedding': [0.10, 0.08], 'maternal': [0.2, 0]}, - 'beta': {'simple_embedding': [0.25, 0.20], 'maternal': [0.2, 0]}, + 'beta': {'simple_embedding': [0.2, 0.15], 'maternal': [0.3, 0]}, 'initial': int(np.maximum(10, np.ceil(0.01*n))), + 'art_efficacy': 0.96, } hiv = ss.HIV(hiv_pars) pregnancy = ss.Pregnancy() - intv = intervention[choice](t=[0, 10, 20], coverage=[0, intv_cov/3, intv_cov]) pars = { 'start': 1980, - 'end': 2020, - 'interventions': [intv], + 'end': 2070, 'rand_seed': rand_seed, 'verbose': 0, 'remove_dead': True, 'n_agents': len(ppl), # TODO } + + if intv_cov > 0: + pars['interventions'] = [ ss.hiv.ART(t=[0, 10, 20], coverage=[0, intv_cov/3, intv_cov]) ] + sim = ss.Sim(people=ppl, diseases=[hiv], demographics=[pregnancy], pars=pars, label=f'Sim with {n} agents and intv_cov={intv_cov}') sim.initialize() sim.run() df = pd.DataFrame( { 'ti': sim.tivec, - #'hiv.n_infected': sim.results.hiv.n_infected, + #'hiv.n_infected': sim.results.hiv.n_infected, # Optional, but mostly redundant with prevalence 'hiv.prevalence': sim.results.hiv.prevalence, 'hiv.cum_deaths': sim.results.hiv.new_deaths.cumsum(), 'pregnancy.cum_births': sim.results.pregnancy.births.cumsum(), @@ -74,7 +71,8 @@ def run_sim(n, idx, intv_cov, rand_seed, multistream): def run_scenarios(): results = [] times = {} - for multistream in [True, False]: + for multistream in [False, True]: + ss.options(multistream=multistream) cfgs = [] for rs in range(n_rand_seeds): for intv_cov in intv_cov_levels: diff --git a/tests/run_stable_monogamy.py b/tests/run_stable_monogamy.py index a8c44651..4ceabfca 100644 --- a/tests/run_stable_monogamy.py +++ b/tests/run_stable_monogamy.py @@ -13,8 +13,15 @@ import argparse import sciris as sc +do_plot_graph = True +kind = ['radial', 'bipartite', 'spring', 'multipartite'][1] + +default_n_agents = 25 + +do_plot_longitudinal = True +do_plot_timeseries = True + ss.options(multistream = True) # Can set multistream to False for comparison -plot_graph = True figdir = os.path.join(os.getcwd(), 'figs', 'stable_monogamy') sc.path(figdir).mkdir(parents=True, exist_ok=True) @@ -29,7 +36,7 @@ def __init__(self, nodes, edges): def draw_nodes(self, filter, pos, ax, **kwargs): inds = [i for i,n in self.graph.nodes.data() if filter(n)] nc = ['red' if nd['hiv'] else 'lightgray' for i, nd in self.graph.nodes.data() if i in inds] - ec = ['green' if nd['on_art'] or nd['on_prep'] else 'black' for i, nd in self.graph.nodes.data() if i in inds] + ec = ['green' if nd['on_art'] else 'black' for i, nd in self.graph.nodes.data() if i in inds] if inds: nx.draw_networkx_nodes(self.graph, nodelist=inds, pos=pos, ax=ax, node_color=nc, edgecolors=ec, **kwargs) return @@ -73,8 +80,6 @@ def update_results(self, sim, init=False): 'hiv': sim.people.hiv.infected.values, 'on_art': sim.people.hiv.on_art.values, 'cd4': sim.people.hiv.cd4.values, - 'on_prep': sim.people.hiv.on_prep.values, - #'ng': sim.people.gonorrhea.infected.values, }) edges = pd.DataFrame(sim.people.networks[0].to_dict()) #sim.people.networks['simple_embedding'].to_df() #TODO: repr issues @@ -91,16 +96,16 @@ def finalize(self, sim): def run_sim(n=25, rand_seed=0, intervention=False, analyze=False): ppl = ss.People(n) - ppl.networks = ss.ndict(ss.simple_embedding(mean_dur=5))#, ss.maternal()) + ppl.networks = ss.ndict(ss.simple_embedding(mean_dur=5), ss.maternal()) hiv_pars = { - #'beta': {'simple_embedding': [0.06, 0.04]}, - 'beta': {'simple_embedding': [0.3, 0.25]}, + 'beta': {'simple_embedding': [0.3, 0.25], 'maternal': [0.2, 0]}, 'initial': 0.25 * n, + 'art_efficacy': 0.96, } hiv = ss.HIV(hiv_pars) - art = ss.hiv.ART(0, 0.5) + art = ss.hiv.ART(0, 0.4) # 40% coverage from day 0 pars = { 'start': 1980, @@ -132,25 +137,66 @@ def run_scenario(n=10, rand_seed=0, analyze=True): return sims +def getpos(ti, g1, g2, guess=None, kind='bipartite'): + + n1 = dict(g1[ti].graph.nodes.data()) + n2 = dict(g2[ti].graph.nodes.data()) + nodes = sc.mergedicts(n2, n1) + n = len(nodes) + + if kind == 'radial': + pos = {i:(np.cos(2*np.pi*i/n), np.sin(2*np.pi*i/n)) for i in range(n)} + if guess: + if len(guess) < n: + pos = {i:(np.cos(2*np.pi*i/n), np.sin(2*np.pi*i/n)) for i in range(n)} + + elif kind == 'spring': + pos = nx.spring_layout(g1[ti].graph, k=None, pos=guess, fixed=None, iterations=50, threshold=0.0001, weight=None, scale=1, center=None, dim=2, seed=None) + if guess: + pos = sc.mergedicts(pos, guess) + + elif kind == 'multipartite': + pos = nx.multipartite_layout(g1[ti].graph, subset_key='female', align='vertical', scale=10, center=None) + if guess: + pos = sc.mergedicts(pos, guess) + + if guess: + for i in guess.keys(): + pos[i] = (pos[i][0], guess[i][1]) # Keep new x but carry over y + + elif kind == 'bipartite': + pos = {i:(nd['age'], 2*nd['female']-1 + np.random.uniform(-0.3, 0.3)) for i, nd in nodes.items()} + + if guess: + for i in guess.keys(): + pos[i] = (pos[i][0], guess[i][1]) # Keep new x but carry over y + + return pos + + def plot_graph(sim1, sim2): g1 = sim1.analyzers[0].graphs g2 = sim2.analyzers[0].graphs n = len(g1[0].graph) el = n <= 10 # Edge labels - #pos = {i:(np.cos(2*np.pi*i/n), np.sin(2*np.pi*i/n)) for i in range(n)} - pos = {i:(nd['age'], 2*nd['female']-1 + np.random.uniform(-0.3, 0.3)) for i, nd in g1[0].graph.nodes.data()} - #pos = nx.spring_layout(g1[0].graph, k=None, pos=None, fixed=None, iterations=50, threshold=0.0001, weight=None, scale=1, center=None, dim=2, seed=None) - #pos = nx.multipartite_layout(g1[0].graph, subset_key='female', align='vertical', scale=10, center=None) fig, axv = plt.subplots(1, 2, figsize=(10,5)) global ti - ti = -1 # Initial state is -1 timax = sim1.tivec[-1] + + global pos + pos = {} + pos[-1] = getpos(0, g1, g2, kind=kind) + for ti in range(timax): + pos[ti] = getpos(ti, g1, g2, guess=pos[ti-1], kind=kind) + + ti = -1 # Initial state is -1, representing the state before the first step + def on_press(event): print('press', event.key) sys.stdout.flush() - global ti + global ti, pos if event.key == 'right': ti = min(ti+1, timax) elif event.key == 'left': @@ -160,18 +206,21 @@ def on_press(event): axv[0].clear() axv[1].clear() - g1[ti].plot(pos, edge_labels=el, ax=axv[0]) - g2[ti].plot(pos, edge_labels=el, ax=axv[1]) - fig.suptitle(f'Time is {ti}') + g1[ti].plot(pos[ti], edge_labels=el, ax=axv[0]) + g2[ti].plot(pos[ti], edge_labels=el, ax=axv[1]) + fig.suptitle(f'Time is {ti} (use the arrow keys to change)') fig.canvas.draw() fig.canvas.mpl_connect('key_press_event', on_press) - g1[ti].plot(pos, ax=axv[0]) - g2[ti].plot(pos, ax=axv[1]) - fig.suptitle(f'Time is {ti}') + g1[ti].plot(pos[ti], edge_labels=el, ax=axv[0]) + g2[ti].plot(pos[ti], edge_labels=el, ax=axv[1]) + fig.suptitle(f'Time is {ti} (use the arrow keys to change)') + return fig + - # Plot +def plot_ts(): + # Plot timeseries summary fig, axv = plt.subplots(2,1, sharex=True) axv[0].plot(sim1.tivec, sim1.results.hiv.n_infected, label='Baseline') axv[0].plot(sim2.tivec, sim2.results.hiv.n_infected, ls=':', label='Intervention') @@ -181,12 +230,9 @@ def on_press(event): axv[1].plot(sim2.tivec, sim2.results.hiv.new_deaths, ls=':', label='Intervention') axv[1].set_title('HIV Deaths') - ''' Gonorrhea removed for now - axv[2].plot(sim1.tivec, sim1.results.gonorrhea.n_infected, label='Baseline') - axv[2].plot(sim2.tivec, sim2.results.gonorrhea.n_infected, ls=':', label='Intervention') - axv[2].set_title('Gonorrhea number of infections') - ''' plt.legend() + return fig + def analyze_people(sim): p = sim.people @@ -198,7 +244,6 @@ def analyze_people(sim): age_initial = age_initial.astype(np.float32) # For better hash comparability, there are small differences at float64 df = pd.DataFrame({ - #'uid': p.uid[ever_alive], # if slicing, don't need ._view, 'id': [hash((p.slot[i], age_initial[i], p.female[i])) for i in ever_alive], # if slicing, don't need ._view, 'age_initial': age_initial, 'years_lived': years_lived, @@ -215,27 +260,7 @@ def analyze_people(sim): df['age_art'] = df['age_initial'] + df['ti_art'] df['age_dead'] = df['age_initial'] + df['ti_dead'] return df - -def life_bars(data, **kwargs): - age_final = data['age_initial'] + data['years_lived'] - plt.barh(y=data.index, left=data['age_initial'], width=age_final-data['age_initial'], color='k') - # Define bools - infected = ~data['ti_infected'].isna() - art = ~data['ti_art'].isna() - dead = ~data['ti_dead'].isna() - - # Infected - plt.barh(y=data.index[infected], left=data.loc[infected]['age_infected'], width=age_final[infected]-data.loc[infected]['age_infected'], color='r') - - # ART - plt.barh(y=data.index[art], left=data.loc[art]['age_art'], width=age_final[art]-data.loc[art]['age_art'], color='g') - - # Dead - #plt.barh(y=data.index[dead], left=data.loc[dead]['age_dead'], width=age_final[dead]-data.loc[dead]['age_dead'], color='k') - plt.scatter(y=data.index[dead], x=data.loc[dead]['age_dead'], color='k', marker='|') - - return def plot_longitudinal(sim1, sim2): @@ -258,25 +283,29 @@ def plot_longitudinal(sim1, sim2): ti_initial = np.maximum(-data['age_initial'], 0) ti_final = data['ti_dead'].fillna(40) - plt.barh(y=yp, left=ti_initial, width=ti_final - ti_initial, color='k', height=height) + plt.barh(y=yp, left=ti_initial, width=ti_final - ti_initial, color='k', height=height, label='Alive' if n==0 else None) # Infected before birth vertical = data['age_infected']<0 - plt.barh(y=yp[vertical], left=data.loc[vertical]['ti_infected'], width=ti_final[vertical]-data.loc[vertical]['ti_infected'], color='m', height=height) + plt.barh(y=yp[vertical], left=data.loc[vertical]['ti_infected'], width=ti_final[vertical]-data.loc[vertical]['ti_infected'], color='m', height=height, label='Infected before birth' if n==0 else None) # Infected infected = ~data['ti_infected'].isna() ai = data.loc[infected]['age_infected'].values # Adjust for vertical transmission ai[~(ai<0)] = 0 - plt.barh(y=yp[infected], left=data.loc[infected]['ti_infected']-ai, width=ti_final[infected]-data.loc[infected]['ti_infected']+ai, color='r', height=height) + plt.barh(y=yp[infected], left=data.loc[infected]['ti_infected']-ai, width=ti_final[infected]-data.loc[infected]['ti_infected']+ai, color='r', height=height, label='Infected' if n==0 else None) # ART art = ~data['ti_art'].isna() - plt.barh(y=yp[art], left=data.loc[art]['ti_art'], width=ti_final[art]-data.loc[art]['ti_art'], color='g', height=height) + plt.barh(y=yp[art], left=data.loc[art]['ti_art'], width=ti_final[art]-data.loc[art]['ti_art'], color='g', height=height, label='ART' if n==0 else None) # Dead dead = ~data['ti_dead'].isna() - plt.scatter(y=yp[dead], x=data.loc[dead]['ti_dead'], color='k', marker='|') + plt.scatter(y=yp[dead], x=data.loc[dead]['ti_dead'], color='c', marker='|', label='Death' if n==0 else None) + + ax.set_xlabel('Age (years)') + ax.set_ylabel('UID') + ax.legend() return fig @@ -284,8 +313,8 @@ def plot_longitudinal(sim1, sim2): if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('-p', '--plot', help='Plot from a cached CSV file', type=str) - parser.add_argument('-n', help='Number of agents', type=int, default=100) - parser.add_argument('-s', help='Rand seed', type=int, default=2) + parser.add_argument('-n', help='Number of agents', type=int, default=default_n_agents) + parser.add_argument('-s', help='Rand seed', type=int, default=0) args = parser.parse_args() if args.plot: @@ -296,8 +325,14 @@ def plot_longitudinal(sim1, sim2): print('Running scenarios') [sim1, sim2] = run_scenario(n=args.n, rand_seed=args.s) - #plot_graph(sim1, sim2) - plot_longitudinal(sim1, sim2) + if do_plot_longitudinal: + plot_longitudinal(sim1, sim2) + + if do_plot_graph: + plot_graph(sim1, sim2) + + if do_plot_timeseries: + plot_ts() plt.show() print('Done') \ No newline at end of file diff --git a/tests/test_stream.py b/tests/test_stream.py index 7bca0022..e70ce571 100644 --- a/tests/test_stream.py +++ b/tests/test_stream.py @@ -9,12 +9,12 @@ from stisim.streams import NotResetException, Stream -def make_rng(base_seed=1, name='Test'): +def make_rng(slots, base_seed=1, name='Test'): """ Create and initialize a stream """ streams = ss.Streams() streams.initialize(base_seed=base_seed) rng = Stream(name) - rng.initialize(streams) + rng.initialize(streams, slots=slots) return rng @@ -24,10 +24,10 @@ def make_rng(base_seed=1, name='Test'): def test_sample(n=5): """ Simple sample """ sc.heading('Testing stream object') - rng = make_rng() + rng = make_rng(n) uids = np.arange(0,n,2) # every other to make it interesting - draws = rng.random(uids) + draws = rng.random(uids=uids) print(f'\nSAMPLE({n}): {draws}') return len(draws) == len(uids) @@ -36,7 +36,7 @@ def test_sample(n=5): def test_neg_binomial(n=5): """ Negative Binomial """ sc.heading('Testing negative binomial') - rng = make_rng() + rng = make_rng(n) # Can call directly: #rng.negative_binomial(n=40, p=1/3, size=5) @@ -44,12 +44,12 @@ def test_neg_binomial(n=5): # Or through a Distribution nb = ss.neg_binomial(mean=80, dispersion=40) nb.set_stream(rng) - draws = nb.sample(n) + draws = nb.sample(size=n) rng.step(1) # Prepare to call again # Now try calling with UIDs instead of n uids = np.arange(0,n,2) # every other to make it interesting - draws_u = nb.sample(uids) + draws_u = nb.sample(uids=uids) print(f'\nSAMPLE({n}): {draws}') return len(draws) == n and len(draws_u) == len(uids) @@ -58,16 +58,16 @@ def test_neg_binomial(n=5): def test_reset(n=5): """ Sample, reset, sample """ sc.heading('Testing sample, reset, sample') - rng = make_rng() + rng = make_rng(n) uids = np.arange(0,n,2) # every other to make it interesting - draws1 = rng.random(uids) + draws1 = rng.random(uids=uids) print(f'\nSAMPLE({n}): {draws1}') print(f'\nRESET') rng.reset() - draws2 = rng.random(uids) + draws2 = rng.random(uids=uids) print(f'\nSAMPLE({n}): {draws2}') return np.all(np.equal(draws1, draws2)) @@ -76,16 +76,16 @@ def test_reset(n=5): def test_step(n=5): """ Sample, step, sample """ sc.heading('Testing sample, step, sample') - rng = make_rng() + rng = make_rng(n) uids = np.arange(0,n,2) # every other to make it interesting - draws1 = rng.random(uids) + draws1 = rng.random(uids=uids) print(f'\nSAMPLE({n}): {draws1}') print(f'\nSTEP(1) - sample should change') rng.step(1) - draws2 = rng.random(uids) + draws2 = rng.random(uids=uids) print(f'\nSAMPLE({n}): {draws2}') return np.all(np.equal(draws1, draws2)) @@ -96,12 +96,12 @@ def test_seed(n=5): sc.heading('Testing sample with seeds 0 and 1') uids = np.arange(0,n,2) # every other to make it interesting - rng0 = make_rng(base_seed=0) - draws0 = rng0.random(uids) + rng0 = make_rng(n, base_seed=0) + draws0 = rng0.random(uids=uids) print(f'\nSAMPLE({n}): {draws0}') - rng1 = make_rng(base_seed=1) - draws1 = rng1.random(uids) + rng1 = make_rng(n, base_seed=1) + draws1 = rng1.random(uids=uids) print(f'\nSAMPLE({n}): {draws1}') return np.all(np.equal(draws0, draws1)) @@ -110,15 +110,15 @@ def test_seed(n=5): def test_repeat(n=5): """ Sample, sample - should raise and exception""" sc.heading('Testing sample, sample - should raise an exception') - rng = make_rng() + rng = make_rng(n) uids = np.arange(0,n,2) # every other to make it interesting - draws1 = rng.random(uids) + draws1 = rng.random(uids=uids) print(f'\nSAMPLE({n}): {draws1}') print(f'\nSAMPLE({n}): [should raise an exception as neither reset() nor step() have been called]') try: - rng.random(uids) + rng.random(uids=uids) return False # Should not get here! except NotResetException as e: print(f'YAY! Got exception: {e}') @@ -128,17 +128,17 @@ def test_repeat(n=5): def test_boolmask(n=5): """ Simple sample with a boolean mask""" sc.heading('Testing stream object') - rng = make_rng() + rng = make_rng(n) uids = np.arange(0,n,2) # every other to make it interesting mask = np.full(n, False) mask[uids] = True - draws_bool = rng.random(mask) + draws_bool = rng.random(uids=mask) print(f'\nSAMPLE({n}): {draws_bool}') rng.reset() - draws_uids = rng.random(uids) + draws_uids = rng.random(uids=uids) print(f'\nSAMPLE({n}): {draws_uids}') return np.all(np.equal(draws_bool, draws_uids)) @@ -147,37 +147,14 @@ def test_boolmask(n=5): def test_empty(): """ Simple sample with a boolean mask""" sc.heading('Testing empty draw') - rng = make_rng() + rng = make_rng(n) uids = np.array([]) # EMPTY - draws = rng.random(uids) + draws = rng.random(uids=uids) print(f'\nSAMPLE: {draws}') return len(draws) == 0 -def test_drawsize(): - """ Testing the draw_size function directly """ - sc.heading('Testing draw size') - - rng = make_rng() - - x = ss.states.FusedArray(values=np.array([1,3,9]), uid=np.array([0,1,2])) - ds_FA = rng.draw_size(x) == 10 # Should be 10 because max value of x is 9 - - x = ss.states.DynamicView(int, fill_value=0) - x.initialize(3) - x[0] = 9 - ds_DV = rng.draw_size(x) == 10 # Should be 10 because max value (representing uid) is 9 - - x = np.full(10, fill_value=True) - ds_bool = rng.draw_size(x) == 10 # Should be 10 because 10 objects - - x = np.array([9]) - ds_array = rng.draw_size(x) == 10 # Should be 10 because 10 objects - - return np.all([ds_FA, ds_DV, ds_bool, ds_array]) - - # %% Run as a script if __name__ == '__main__': # Start timing @@ -189,7 +166,7 @@ def test_drawsize(): ss.options(multistream=multistream) print('Testing with multistream set to', multistream) - # Run tests + # Run tests - some will only pass if multistream is True print(test_sample(n)) print(test_neg_binomial(n)) print(test_reset(n)) @@ -198,7 +175,6 @@ def test_drawsize(): print(test_repeat(n)) print(test_boolmask(n)) print(test_empty()) - print(test_drawsize()) sc.toc(T) print('Done.') \ No newline at end of file diff --git a/tests/test_streams.py b/tests/test_streams.py index 22a7ec58..317f1c28 100644 --- a/tests/test_streams.py +++ b/tests/test_streams.py @@ -18,10 +18,10 @@ def test_streams(n=5): streams.initialize(base_seed=10) rng = ss.MultiStream('stream1') - rng.initialize(streams) + rng.initialize(streams, slots=n) uids = np.arange(0,n,2) # every other to make it interesting - draws = rng.random(uids) + draws = rng.random(uids=uids) print(f'\nCREATED SEED AND SAMPLED: {draws}') return len(draws) == len(uids) @@ -34,10 +34,10 @@ def test_seed(n=5): streams.initialize(base_seed=10) rng0 = ss.MultiStream('stream0') - rng0.initialize(streams) + rng0.initialize(streams, slots=n) rng1 = ss.MultiStream('stream1') - rng1.initialize(streams) + rng1.initialize(streams, slots=n) return rng1.seed != rng0.seed @@ -49,14 +49,14 @@ def test_reset(n=5): streams.initialize(base_seed=10) rng = ss.MultiStream('stream0') - rng.initialize(streams) + rng.initialize(streams, slots=n) uids = np.arange(0,n,2) # every other to make it interesting - s_before = rng.random(uids) + s_before = rng.random(uids=uids) streams.reset() # Return to step 0 - s_after = rng.random(uids) + s_after = rng.random(uids=uids) return np.all(s_after == s_before) @@ -68,14 +68,14 @@ def test_step(n=5): streams.initialize(base_seed=10) rng = ss.MultiStream('stream0') - rng.initialize(streams) + rng.initialize(streams, slots=n) uids = np.arange(0,n,2) # every other to make it interesting - s_before = rng.random(uids) + s_before = rng.random(uids=uids) streams.step(10) # 10 steps - s_after = rng.random(uids) + s_after = rng.random(uids=uids) return np.all(s_after != s_before) @@ -89,7 +89,7 @@ def test_initialize(n=5): rng = ss.MultiStream('stream0') try: - rng.initialize(streams) + rng.initialize(streams, slots=n) return False # Should not get here! except NotInitializedException as e: print(f'YAY! Got exception: {e}') @@ -103,11 +103,11 @@ def test_seedrepeat(n=5): streams.initialize(base_seed=10) rng = ss.MultiStream('stream0', seed_offset=0) - rng.initialize(streams) + rng.initialize(streams, slots=n) try: rng1 = ss.MultiStream('stream1', seed_offset=0) - rng1.initialize(streams) + rng1.initialize(streams, slots=n) return False # Should not get here! except SeedRepeatException as e: print(f'YAY! Got exception: {e}') @@ -123,18 +123,18 @@ def test_samplingorder(n=5): uids = np.arange(0,n,2) # every other to make it interesting rng0 = ss.MultiStream('stream0') - rng0.initialize(streams) + rng0.initialize(streams, slots=n) rng1 = ss.MultiStream('stream1') - rng1.initialize(streams) + rng1.initialize(streams, slots=n) - s_before = rng0.random(uids) - _ = rng1.random(uids) + s_before = rng0.random(uids=uids) + _ = rng1.random(uids=uids) streams.reset() - _ = rng1.random(uids) - s_after = rng0.random(uids) + _ = rng1.random(uids=uids) + s_after = rng0.random(uids=uids) return np.all(s_before == s_after) @@ -146,11 +146,11 @@ def test_repeatname(n=5): streams.initialize(base_seed=17) rng0 = ss.MultiStream('test') - rng0.initialize(streams) + rng0.initialize(streams, slots=n) rng1 = ss.MultiStream('test') try: - rng1.initialize(streams) + rng1.initialize(streams, slots=n) return False # Should not get here! except RepeatNameException as e: print(f'YAY! Got exception: {e}') @@ -163,14 +163,14 @@ def test_repeatname(n=5): T = sc.tic() # Run tests - assert test_streams() - assert test_seed() - assert test_reset() - assert test_step() - assert test_initialize() - assert test_seedrepeat() - assert test_samplingorder() - assert test_repeatname() + print(test_streams()) + print(test_seed()) + print(test_reset()) + print(test_step()) + print(test_initialize()) + print(test_seedrepeat()) + print(test_samplingorder()) + print(test_repeatname()) sc.toc(T) print('Done.') \ No newline at end of file From 9130a4424a1243f198a7725a8a22143724427050 Mon Sep 17 00:00:00 2001 From: Daniel Klein Date: Thu, 26 Oct 2023 13:54:19 -0700 Subject: [PATCH 31/34] Enabling the user to configure "slot_scale" - a new configuration parameter. --- stisim/demographics.py | 2 +- stisim/parameters.py | 4 ++++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/stisim/demographics.py b/stisim/demographics.py index f08708a1..9a3ec405 100644 --- a/stisim/demographics.py +++ b/stisim/demographics.py @@ -319,7 +319,7 @@ def make_pregnancies(self, sim): n_unborn_agents = len(uids) if n_unborn_agents > 0: # TODO: User configure the bounds - new_slots = self.rng_uids.integers(low=sim.pars['n_agents'], high=10*sim.pars['n_agents'], uids=uids, dtype=int) + 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(len(new_slots)) diff --git a/stisim/parameters.py b/stisim/parameters.py index 6b7f37bd..c3520564 100644 --- a/stisim/parameters.py +++ b/stisim/parameters.py @@ -46,6 +46,7 @@ def __init__(self, **kwargs): self.dt = 1.0 # Timestep (in years) self.dt_demog = 1.0 # Timestep for demographic updates (in years) self.rand_seed = 1 # Random seed, if None, don't reset + self.slot_scale = 5 # Random slots will be assigned to newborn agents between min=n_agents and max=slot_scale*n_agents. Choosing a larger value here will reduce the probability of two agents using the same slot (and hence random draws), but increase the number of random numbers that are required. self.verbose = ss.options.verbose # Whether or not to display information during the run -- options are 0 (silent), 0.1 (some; default), 1 (default), 2 (everything) # Events and interventions @@ -61,6 +62,9 @@ def __init__(self, **kwargs): # Update with any supplied parameter values and generate things that need to be generated self.update(kwargs) + if self.slot_scale < 1: + raise Exception('The value of the "slot_scale" parameter must be a number >= 1.0') + return def update_pars(self, pars=None, create=False, **kwargs): From 01fd0f4904d4f8d333b8af0bb2add80fb0bc992d Mon Sep 17 00:00:00 2001 From: Daniel Klein Date: Thu, 26 Oct 2023 13:55:00 -0700 Subject: [PATCH 32/34] Removing unnecessary comment that should have been bundled with the previous commit. --- stisim/demographics.py | 1 - 1 file changed, 1 deletion(-) diff --git a/stisim/demographics.py b/stisim/demographics.py index 9a3ec405..ee539162 100644 --- a/stisim/demographics.py +++ b/stisim/demographics.py @@ -318,7 +318,6 @@ def make_pregnancies(self, sim): # 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: - # TODO: User configure the bounds 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 From df65c2e10b5f528a48ea5629e1e848d35f848d97 Mon Sep 17 00:00:00 2001 From: Daniel Klein Date: Thu, 26 Oct 2023 14:56:31 -0700 Subject: [PATCH 33/34] Minor cleanup of comments. --- stisim/distributions.py | 2 +- stisim/gonorrhea.py | 2 +- stisim/modules.py | 5 ++--- 3 files changed, 4 insertions(+), 5 deletions(-) diff --git a/stisim/distributions.py b/stisim/distributions.py index e926ff61..ffc9baf6 100644 --- a/stisim/distributions.py +++ b/stisim/distributions.py @@ -182,7 +182,7 @@ def sample(self, **kwargs): elif 'uids' in kwargs: return np.zeros(len(kwargs['uids'])) else: - raise Exception('TODO') + raise Exception('When calling sample(), please provide either "size" or "uids".') class lognormal_int(lognormal): diff --git a/stisim/gonorrhea.py b/stisim/gonorrhea.py index e08de10d..fc2a8cb3 100644 --- a/stisim/gonorrhea.py +++ b/stisim/gonorrhea.py @@ -61,7 +61,7 @@ def set_prognoses(self, sim, to_uids, from_uids=None): self.infected[to_uids] = True self.ti_infected[to_uids] = sim.ti - dur = sim.ti + self.rng_dur_inf.poisson(to_uids, self.pars['dur_inf']/sim.pars.dt) # By whom infected from??? TODO + dur = sim.ti + self.rng_dur_inf.poisson(to_uids, self.pars['dur_inf']/sim.pars.dt) dead = self.rng_dead.bernoulli(to_uids, self.pars.p_death) self.ti_recovered[to_uids[~dead]] = dur[~dead] diff --git a/stisim/modules.py b/stisim/modules.py index a41e6af4..27833fa6 100644 --- a/stisim/modules.py +++ b/stisim/modules.py @@ -160,8 +160,7 @@ def make_new_cases(self, sim): pars = sim.pars[self.name] # Probability of each node acquiring a case - # TODO: Just people that who are alive? - n = len(sim.people.uid) + n = len(sim.people.uid) # TODO: possibly could be shortened to just the people who are alive p_acq_node = np.zeros( n ) node_from_node = np.ones( (n,n) ) @@ -178,7 +177,6 @@ def make_new_cases(self, sim): continue # Check for new transmission from a --> b - # TODO: Will need to be more efficient here - can maintain edge to node matrix node_from_edge = np.ones( (n, len(a)) ) ai = sim.people._uid_map[a] # Indices of a and b (rather than uid) bi = sim.people._uid_map[b] @@ -188,6 +186,7 @@ def make_new_cases(self, sim): p_not_acq_by_node_this_layer_b_from_a = node_from_edge.prod(axis=1) # (1-p1)*(1-p2)*... p_acq_node = 1 - (1-p_acq_node) * p_not_acq_by_node_this_layer_b_from_a + # TODO: Will need to be more efficient here - can maintain edge to node matrix node_from_node_this_layer_b_from_a = np.ones( (n,n) ) node_from_node_this_layer_b_from_a[bi, ai] = p_not_acq node_from_node *= node_from_node_this_layer_b_from_a From bdff4391b04a217d9fd0e86ca47e0ff11f24e406 Mon Sep 17 00:00:00 2001 From: Daniel Klein Date: Thu, 26 Oct 2023 22:10:02 -0700 Subject: [PATCH 34/34] *More efficient identification of transmission sources, no longer N^2. * In states.py, added check for DynamicView --- stisim/modules.py | 52 +++++++++++++++++++++++++++--------- stisim/states.py | 2 ++ tests/run_stable_monogamy.py | 6 ++--- 3 files changed, 44 insertions(+), 16 deletions(-) diff --git a/stisim/modules.py b/stisim/modules.py index 27833fa6..6e829658 100644 --- a/stisim/modules.py +++ b/stisim/modules.py @@ -129,7 +129,7 @@ def set_initial_states(self, sim): if self.pars['initial'] <= 0: return - initial_cases = self.rng_init_cases.bernoulli_filter(uids=sim.people.uid, prob=self.pars['initial']/len(sim.people)) + initial_cases = self.rng_init_cases.bernoulli_filter(uids=ss.true(sim.people.alive), prob=self.pars['initial']/len(sim.people)) self.set_prognoses(sim, initial_cases, from_uids=None) # TODO: sentinel value to indicate seeds? return @@ -162,7 +162,6 @@ def make_new_cases(self, sim): # Probability of each node acquiring a case n = len(sim.people.uid) # TODO: possibly could be shortened to just the people who are alive p_acq_node = np.zeros( n ) - node_from_node = np.ones( (n,n) ) for lkey, layer in sim.people.networks.items(): if lkey in pars['beta']: @@ -186,25 +185,52 @@ def make_new_cases(self, sim): p_not_acq_by_node_this_layer_b_from_a = node_from_edge.prod(axis=1) # (1-p1)*(1-p2)*... p_acq_node = 1 - (1-p_acq_node) * p_not_acq_by_node_this_layer_b_from_a - # TODO: Will need to be more efficient here - can maintain edge to node matrix - node_from_node_this_layer_b_from_a = np.ones( (n,n) ) - node_from_node_this_layer_b_from_a[bi, ai] = p_not_acq - node_from_node *= node_from_node_this_layer_b_from_a - new_cases_bool = self.rng_trans.bernoulli(uids=sim.people.uid, prob=p_acq_node) new_cases = sim.people.uid[new_cases_bool] if not len(new_cases): return 0 - # Decide whom the infection came from using one random number for each b (aligned by block size) frm = np.zeros_like(new_cases) r = self.rng_choose_infector.random(uids=new_cases) - new_cases_idx = new_cases_bool.nonzero()[0] - prob = (1-node_from_node[new_cases_idx]) # Prob of acquiring from each node | can constrain to just neighbors? - cumsum = (prob / ((prob.sum(axis=1)[:,np.newaxis]))).cumsum(axis=1) - frm_idx = np.argmax( cumsum >= r[:,np.newaxis], axis=1) - frm = sim.people.uid[frm_idx] + for i, uid in enumerate(new_cases): + p_acqs = [] + sources = [] + + for lkey, layer in sim.people.networks.items(): + if lkey in pars['beta']: + a_to_b = [layer['p1'], layer['p2'], pars['beta'][lkey][0]] + b_to_a = [layer['p2'], layer['p1'], pars['beta'][lkey][1]] + for a, b, beta in [a_to_b, b_to_a]: + if beta == 0: + continue + + inds = np.where(b==uid)[0] + if len(inds) == 0: + continue + frms = a[inds] + + # TODO: Likely no longer need alive here, at least not if dead people are removed + rel_trans = self.rel_trans[frms] * (self.infected[frms] & sim.people.alive[frms]) + rel_sus = self.rel_sus[uid] * (self.susceptible[uid] & sim.people.alive[uid]) + beta_combined = layer['beta'][inds] * beta + + # Check for new transmission from a --> b + # TODO: Remove zeros from this... + p_acqs.append((rel_trans * rel_sus * beta_combined).__array__()) # Needs DT + sources.append(frms.__array__()) + + p_acqs = np.concatenate(p_acqs) + sources = np.concatenate(sources) + + if len(sources) == 1: + frm[i] = sources[0] + else: + # Choose using draw r from above + cumsum = p_acqs / p_acqs.sum() + frm_idx = np.argmax( cumsum >= r[i]) + frm[i] = sources[frm_idx] + self.set_prognoses(sim, new_cases, frm) return len(new_cases) # number of new cases made diff --git a/stisim/states.py b/stisim/states.py index b4af06e6..2966526e 100644 --- a/stisim/states.py +++ b/stisim/states.py @@ -117,6 +117,8 @@ def __getitem__(self, key): else: # Access items by an array of integers. We do get a decent performance boost from using numba here values, uids, new_uid_map = self._get_vals_uids(self.values, key, self._uid_map.__array__()) + elif isinstance(key, DynamicView): + values, uids, new_uid_map = self._get_vals_uids(self.values, key.__array__(), self._uid_map.__array__()) elif isinstance(key, slice): if key.start is None and key.stop is None and key.step is None: return sc.dcp(self) diff --git a/tests/run_stable_monogamy.py b/tests/run_stable_monogamy.py index 4ceabfca..e294a6a3 100644 --- a/tests/run_stable_monogamy.py +++ b/tests/run_stable_monogamy.py @@ -13,11 +13,11 @@ import argparse import sciris as sc -do_plot_graph = True -kind = ['radial', 'bipartite', 'spring', 'multipartite'][1] - default_n_agents = 25 +do_plot_graph = True +kind = ['bipartite', 'radial', 'spring', 'multipartite'][0] + do_plot_longitudinal = True do_plot_timeseries = True