diff --git a/seirsplus/FARZ.py b/seirsplus/FARZ.py index f81acc4..4b3444f 100644 --- a/seirsplus/FARZ.py +++ b/seirsplus/FARZ.py @@ -5,12 +5,12 @@ import random import bisect import math -import os +import os def random_choice(values, weights=None , size = 1, replace = True): if weights is None: i = int(random.random() * len(values)) - else : + else: total = 0 cum_weights = [] for w in weights: @@ -18,31 +18,31 @@ def random_choice(values, weights=None , size = 1, replace = True): cum_weights.append(total) x = random.random() * total i = bisect.bisect(cum_weights, x) - if size <=1: - if len(values)>i: return values[i] + if size <=1: + if len(values)>i: return values[i] else: return None - else: + else: cval = [values[j] for j in range(len(values)) if replace or i!=j] - if weights is None: cwei=None + if weights is None: cwei=None else: cwei = [weights[j] for j in range(len(weights)) if replace or i!=j] tmp= random_choice(cval, cwei, size-1, replace) if not isinstance(tmp,list): tmp = [tmp] tmp.append(values[i]) - return tmp + return tmp class Comms: def __init__(self, k): self.k = k self.groups = [[] for i in range(k)] self.memberships = {} - + def add(self, cluster_id, i, s = 1): if i not in [m[0] for m in self.groups[cluster_id]]: - self.groups[cluster_id].append((i,s)) + self.groups[cluster_id].append((i,s)) if i in self.memberships: self.memberships[i].append((cluster_id,s)) else: - self.memberships[i] =[(cluster_id,s)] + self.memberships[i] =[(cluster_id,s)] def write_groups(self, path): with open(path, 'w') as f: for g in self.groups: @@ -50,8 +50,8 @@ def write_groups(self, path): f.write(str(i) + ' ') f.write('\n') - - + + class Graph: def __init__(self,directed=False, weighted=False): self.n = 0 @@ -59,27 +59,27 @@ def __init__(self,directed=False, weighted=False): self.max_degree = 0 self.directed = directed self.weighted = weighted - self.edge_list = [] + self.edge_list = [] self.edge_time = [] self.deg = [] self.neigh = [[]] - return + return def add_node(self): self.deg.append(0) self.neigh.append([]) self.n+=1 - + def weight(self, u, v): for i,w in self.neigh[u]: if i == v: return w return 0 - + def is_neigh(self, u, v): for i,_ in self.neigh[u]: if i == v: return True return False - + def add_edge(self, u, v, w=1): if u==v: return if not self.weighted : w =1 @@ -89,15 +89,15 @@ def add_edge(self, u, v, w=1): self.neigh[u].append((v,w)) self.deg[v]+=w if self.deg[v]>self.max_degree: self.max_degree = self.deg[v] - + if not self.directed: #if directed deg is indegree, outdegree = len(negh) self.neigh[v].append((u,w)) self.deg[u]+=w if self.deg[u]>self.max_degree: self.max_degree = self.deg[u] - return - - + return + + def to_nx(self): import networkx as nx G=nx.Graph() @@ -105,7 +105,7 @@ def to_nx(self): G.add_edge(u, v) # G.add_edges_from(self.edge_list) return G - + def to_nx(self, C): import networkx as nx G=nx.Graph() @@ -121,19 +121,19 @@ def to_nx(self, C): G.add_edge(u, v, weight=w, capacity=self.edge_time[i]) # G.add_edges_from(self.edge_list) return G - + def to_ig(self): G=ig.Graph() G.add_edges(self.edge_list) - return G - - + return G + + def write_edgelist(self, path): with open(path, 'w') as f: for i,j,w in self.edge_list: f.write(str(i) + '\t'+str(j) + '\n') - + def Q(G, C): q = 0.0 m = 2 * len(G.edge_list) @@ -148,7 +148,7 @@ def common_neighbour(i, G, normalize=True): p = {} for k,wik in G.neigh[i]: for j,wjk in G.neigh[k]: - if j in p: p[j]+=(wik * wjk) + if j in p: p[j]+=(wik * wjk) else: p[j]= (wik * wjk) if len(p)<=0 or not normalize: return p maxp = p[max(p, key = lambda i: p[i])] @@ -159,7 +159,7 @@ def choose_community(i, G, C, alpha, beta, gamma, epsilon): mids =[k for k,uik in C.memberships[i]] if random.random()< beta: #inside cids = mids - else: + else: cids = [j for j in range(len(C.groups)) if j not in mids] #: cids.append(j) return cids[ int(random.random()*len(cids))] if len(cids)>0 else None @@ -180,14 +180,14 @@ def combine (a,b,alpha,gamma): def choose_node(i,c, G, C, alpha, beta, gamma, epsilon): ids = [j for j,_ in C.groups[c] if j !=i ] # also remove nodes that are already connected from the candidate list - for k,_ in G.neigh[i]: - if k in ids: ids.remove(k) + for k,_ in G.neigh[i]: + if k in ids: ids.remove(k) norma = False cn = common_neighbour(i, G, normalize=norma) trim_ids = [id for id in ids if id in cn] dd = degree_similarity(i, trim_ids, G, gamma, normalize=norma) - + if random.random()beta)): G.add_edge(i,k,wjk*pj) - + def connect(i, b, G, C, alpha, beta, gamma, epsilon): #Choose community c = choose_community(i, G, C, alpha, beta, gamma, epsilon) @@ -219,12 +219,12 @@ def connect(i, b, G, C, alpha, beta, gamma, epsilon): #Choose node within community tmp = choose_node(i, c, G, C, alpha, beta, gamma, epsilon) if tmp is None: return - j, pj = tmp + j, pj = tmp G.add_edge(i,j,pj) connect_neighbor(i, j, pj , c, b, G, C, beta) - + def select_node(G, method = 'uniform'): - if method=='uniform': + if method=='uniform': return int(random.random() * G.n) # uniform else: if method == 'older_less_active': p = [(i+1) for i in range(G.n)] # older less active @@ -237,19 +237,19 @@ def assign(i, C, e=1, r=1, q = 0.5): id = random_choice(range(C.k),p ) C.add(id, i) for j in range(1,r): #todo add strength for fuzzy - if (random.random()1 else '') + name = net_name+( str(r+1) if repeat>1 else '') # G = write_to_file(G,C,path,name,format,farz_params) - + # print(len([memtup[0][0] for memtup in C.memberships.values()])) # import numpy # (unique, counts) = numpy.unique( [memtup[0][0] for memtup in C.memberships.values()], return_counts=True ) @@ -335,20 +335,20 @@ def get_range(s,e,i): # exit() node_communities = {node: [c[0] for c in comm_tup] for node, comm_tup in C.memberships.items()} - + return G, node_communities if arange ==None: arange = default_ranges[vari] - for i,var in enumerate(get_range(arange[0],arange[1],arange[2])): + for i,var in enumerate(get_range(arange[0],arange[1],arange[2])): for r in range(repeat): farz_params[vari] = var print('s',i+1, r+1, str(farz_params)) G, C =realize(**farz_params) - name = 'S'+str(i+1)+'-'+net_name+ (str(r+1) if repeat>1 else '') + name = 'S'+str(i+1)+'-'+net_name+ (str(r+1) if repeat>1 else '') write_to_file(G,C,path,name,format,farz_params) - + import sys @@ -356,7 +356,7 @@ def main(argv): import getopt FARZsetting = default_FARZ_setting.copy() batch_setting= default_batch_setting.copy() - try: + try: opts, args = getopt.getopt(argv,"ho:s:v:c:f:n:k:m:a:b:g:p:r:q:t:e:dw",\ ["output=","path=","repeat=","vary=",'range=','format=',"alpha=","beta=","gamma=",'phi=','overlap=','oProb=','epsilon=','cneigh=','directed','weighted']) except getopt.GetoptError: @@ -375,7 +375,7 @@ def main(argv): print('> python FARZ.py --path ./data -s 10 -v beta -c [0.5,1,0.05] -n 1000 -m 5 -k 4 \n') print('+ example 5: generate overlapping communities, where each node belongs to at most 3 communities and the portion of overlapping nodes varies') print('python FARZ.py -r 3 -v q --path ./datavrq -s 5 --format list\n') - + print('*** parameters:') print('-n: number of nodes, default (1000)') print('-m: half the average degree of nodes, default (5)') @@ -398,26 +398,26 @@ def main(argv): print('-c: the range to change the given parameter, should be in format of [s,e,inc]') #print('default FARZ parameters are :\n', default_FARZ_setting) #print('default batch generator parameters are :\n', default_batch_setting) - + sys.exit() - + elif opt in ("-o", "--output"): batch_setting['net_name'] = arg - elif opt in ("--path"): + elif opt in ("--path",): batch_setting['path'] = arg elif opt in ("-f", "--format"): if arg in supported_formats: batch_setting['format'] = arg else: print('Format not supported , choose from ',supported_formats,' or try -h to see the usage and options') - sys.exit(2) + sys.exit(2) elif opt in ("-s","--repeat"): - try: batch_setting['repeat'] = int(arg) + try: batch_setting['repeat'] = int(arg) except ValueError: print('Invalid Number , try -h to see the usage and options') - sys.exit(2) + sys.exit(2) elif opt in ("-v", "--vary"): - if (arg in list(default_ranges.keys())): + if arg in list(default_ranges.keys()): batch_setting['vari'] = arg else: print('Invalid variable, choose form :', list(default_ranges.keys()), ', try -h to see the usage and options') @@ -428,47 +428,47 @@ def main(argv): batch_setting['arange'] = arange except Error: print('Invalid range, should have the following form : [start,end,incrementBy], try -h to see the usage and options ') - sys.exit(2) - elif opt in ("-n"): - try: FARZsetting['n'] = int(arg) + sys.exit(2) + elif opt in ("-n",): + try: FARZsetting['n'] = int(arg) except ValueError: print('Invalid Number , try -h to see the usage and options') sys.exit(2) - elif opt in ("-k"): + elif opt in ("-k",): try: FARZsetting['k'] = int(arg) except ValueError: print('Invalid Number , try -h to see the usage and options') sys.exit(2) - elif opt in ("-m"): + elif opt in ("-m",): try: FARZsetting['m'] = int(arg) except ValueError: print('Invalid Number , try -h to see the usage and options') sys.exit(2) elif opt in ("-a","--alpha"): - try: FARZsetting['alpha'] = float(arg) + try: FARZsetting['alpha'] = float(arg) except ValueError: print('Invalid Number , try -h to see the usage and options') - sys.exit(2) + sys.exit(2) elif opt in ("-b","--beta"): try: FARZsetting['beta'] = float(arg) except ValueError: print('Invalid Number , try -h to see the usage and options') - sys.exit(2) + sys.exit(2) elif opt in ("-g","--gamma"): try: FARZsetting['gamma'] = float(arg) except ValueError: print('Invalid Number , try -h to see the usage and options') - sys.exit(2) + sys.exit(2) elif opt in ("-p","--phi"): try: FARZsetting['phi'] = int(arg) except ValueError: print('Invalid Number , try -h to see the usage and options') - sys.exit(2) + sys.exit(2) elif opt in ("-r","--overlap"): try: FARZsetting['r'] = int(arg) except ValueError: print('Invalid Number , try -h to see the usage and options') - sys.exit(2) + sys.exit(2) elif opt in ("-q","--oProb"): try: FARZsetting['q'] = float(arg) except ValueError: @@ -477,43 +477,43 @@ def main(argv): elif opt in ("-d","--directed"): FARZsetting['directed'] = True elif opt in ("-w","--wighted"): - FARZsetting['weighted'] = True + FARZsetting['weighted'] = True elif opt in ("-t","--cneigh"): try: FARZsetting['b'] = float(arg) except ValueError: print('Invalid Number , try -h to see the usage and options') - sys.exit(2) + sys.exit(2) elif opt in ("-e","--epsilon"): - try: FARZsetting['epsilon'] = float(arg) + try: FARZsetting['epsilon'] = float(arg) except ValueError: print('Invalid Number , try -h to see the usage and options') sys.exit(2) - + batch_setting['farz_params'] = FARZsetting print('generating FARZ benchmark(s) ... ') generate( **batch_setting) - + if __name__ == "__main__": main(sys.argv[1:]) -# generate(farz_params={"n":25000, -# "k":4, -# "m":5, +# generate(farz_params={"n":25000, +# "k":4, +# "m":5, # "alpha":0.5, -# "gamma":0.5, -# "beta":.8, -# "phi":1, -# "o":1, -# 'q':0.5, -# "b":0.0, -# "epsilon":0.0000001, -# 'directed':False, +# "gamma":0.5, +# "beta":.8, +# "phi":1, +# "o":1, +# 'q':0.5, +# "b":0.0, +# "epsilon":0.0000001, +# 'directed':False, # 'weighted':False}) - - + + # python FARZ.py --path ./dataVb55 -s 10 -v beta # python FARZ.py --path ./dataVb82 -s 10 -v beta --alpha 0.8 --gamma 0.2 # python FARZ.py --path ./dataVb5-5 -s 10 -v beta --alpha 0.5 --gamma -0.5 -# python FARZ.py --path ./dataVb2-8 -s 10 -v beta --alpha 0.2 --gamma -0.8 \ No newline at end of file +# python FARZ.py --path ./dataVb2-8 -s 10 -v beta --alpha 0.2 --gamma -0.8 diff --git a/seirsplus/legacy_models.py b/seirsplus/legacy_models.py index 23e86f8..d6d4e55 100644 --- a/seirsplus/legacy_models.py +++ b/seirsplus/legacy_models.py @@ -21,14 +21,14 @@ class SEIRSModel(): """ A class to simulate the Deterministic SEIRS Model =================================================== - Params: beta Rate of transmission (exposure) - sigma Rate of infection (upon exposure) - gamma Rate of recovery (upon infection) - xi Rate of re-susceptibility (upon recovery) - mu_I Rate of infection-related death - mu_0 Rate of baseline death + Params: beta Rate of transmission (exposure) + sigma Rate of infection (upon exposure) + gamma Rate of recovery (upon infection) + xi Rate of re-susceptibility (upon recovery) + mu_I Rate of infection-related death + mu_0 Rate of baseline death nu Rate of baseline birth - + beta_D Rate of transmission (exposure) for individuals with detected infections sigma_D Rate of infection (upon exposure) for individuals with detected infections gamma_D Rate of recovery (upon infection) for individuals with detected infections @@ -38,32 +38,32 @@ class SEIRSModel(): psi_E Probability of positive test results for exposed individuals psi_I Probability of positive test results for exposed individuals q Probability of quarantined individuals interacting with others - - initE Init number of exposed individuals - initI Init number of infectious individuals + + initE Init number of exposed individuals + initI Init number of infectious individuals initD_E Init number of detected infectious individuals - initD_I Init number of detected infectious individuals - initR Init number of recovered individuals + initD_I Init number of detected infectious individuals + initR Init number of recovered individuals initF Init number of infection-related fatalities - (all remaining nodes initialized susceptible) + (all remaining nodes initialized susceptible) """ def __init__(self, initN, beta, sigma, gamma, xi=0, mu_I=0, mu_0=0, nu=0, p=0, - beta_D=None, sigma_D=None, gamma_D=None, mu_D=None, + beta_D=None, sigma_D=None, gamma_D=None, mu_D=None, theta_E=0, theta_I=0, psi_E=0, psi_I=0, q=0, initE=0, initI=10, initD_E=0, initD_I=0, initR=0, initF=0): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Model Parameters: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - self.beta = beta - self.sigma = sigma - self.gamma = gamma - self.xi = xi - self.mu_I = mu_I - self.mu_0 = mu_0 - self.nu = nu - self.p = p + self.beta = beta + self.sigma = sigma + self.gamma = gamma + self.xi = xi + self.mu_I = mu_I + self.mu_0 = mu_0 + self.nu = nu + self.p = p # Testing-related parameters: self.beta_D = beta_D if beta_D is not None else self.beta @@ -82,9 +82,9 @@ def __init__(self, initN, beta, sigma, gamma, xi=0, mu_I=0, mu_0=0, nu=0, p=0, self.t = 0 self.tmax = 0 # will be set when run() is called self.tseries = numpy.array([0]) - + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Initialize Counts of inidividuals with each state: + # Initialize Counts of individuals with each state: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ self.N = numpy.array([int(initN)]) self.numE = numpy.array([int(initE)]) @@ -104,7 +104,7 @@ def __init__(self, initN, beta, sigma, gamma, xi=0, mu_I=0, mu_0=0, nu=0, p=0, def system_dfes(t, variables, beta, sigma, gamma, xi, mu_I, mu_0, nu, beta_D, sigma_D, gamma_D, mu_D, theta_E, theta_I, psi_E, psi_I, q): - S, E, I, D_E, D_I, R, F = variables # varibles is a list with compartment counts as elements + S, E, I, D_E, D_I, R, F = variables # variables is a list with compartment counts as elements N = S + E + I + D_E + D_I + R @@ -132,14 +132,14 @@ def run_epoch(self, runtime, dt=0.1): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Create a list of times at which the ODE solver should output system values. - # Append this list of times as the model's timeseries + # Append this list of times as the model's time series t_eval = numpy.arange(start=self.t, stop=self.t+runtime, step=dt) # Define the range of time values for the integration: t_span = (self.t, self.t+runtime) # Define the initial conditions as the system's current state: - # (which will be the t=0 condition if this is the first run of this model, + # (which will be the t=0 condition if this is the first run of this model, # else where the last sim left off) init_cond = [self.numS[-1], self.numE[-1], self.numI[-1], self.numD_E[-1], self.numD_I[-1], self.numR[-1], self.numF[-1]] @@ -149,7 +149,7 @@ def run_epoch(self, runtime, dt=0.1): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ solution = scipy.integrate.solve_ivp(lambda t, X: SEIRSModel.system_dfes(t, X, self.beta, self.sigma, self.gamma, self.xi, self.mu_I, self.mu_0, self.nu, self.beta_D, self.sigma_D, self.gamma_D, self.mu_D, self.theta_E, self.theta_I, self.psi_E, self.psi_I, self.q - ), + ), t_span=t_span, y0=init_cond, t_eval=t_eval ) @@ -173,37 +173,37 @@ def run_epoch(self, runtime, dt=0.1): def run(self, T, dt=0.1, checkpoints=None, verbose=False): - if(T>0): + if T > 0: self.tmax += T else: return False - + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Pre-process checkpoint values: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(checkpoints): + if checkpoints: numCheckpoints = len(checkpoints['t']) paramNames = ['beta', 'sigma', 'gamma', 'xi', 'mu_I', 'mu_0', 'nu', 'beta_D', 'sigma_D', 'gamma_D', 'mu_D', 'theta_E', 'theta_I', 'psi_E', 'psi_I', 'q'] for param in paramNames: - # For params that don't have given checkpoint values (or bad value given), + # For params that don't have given checkpoint values (or bad value given), # set their checkpoint values to the value they have now for all checkpoints. - if(param not in list(checkpoints.keys()) - or not isinstance(checkpoints[param], (list, numpy.ndarray)) + if (param not in list(checkpoints.keys()) + or not isinstance(checkpoints[param], (list, numpy.ndarray)) or len(checkpoints[param])!=numCheckpoints): checkpoints[param] = [getattr(self, param)]*numCheckpoints #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% # Run the simulation loop: #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% - if(not checkpoints): + if not checkpoints: self.run_epoch(runtime=self.tmax, dt=dt) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ print("t = %.2f" % self.t) - if(verbose): + if verbose: print("\t S = " + str(self.numS[-1])) print("\t E = " + str(self.numE[-1])) print("\t I = " + str(self.numI[-1])) @@ -211,7 +211,7 @@ def run(self, T, dt=0.1, checkpoints=None, verbose=False): print("\t D_I = " + str(self.numD_I[-1])) print("\t R = " + str(self.numR[-1])) print("\t F = " + str(self.numF[-1])) - + else: # checkpoints provided for checkpointIdx, checkpointTime in enumerate(checkpoints['t']): @@ -225,7 +225,7 @@ def run(self, T, dt=0.1, checkpoints=None, verbose=False): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ print("t = %.2f" % self.t) - if(verbose): + if verbose: print("\t S = " + str(self.numS[-1])) print("\t E = " + str(self.numE[-1])) print("\t I = " + str(self.numI[-1])) @@ -234,7 +234,7 @@ def run(self, T, dt=0.1, checkpoints=None, verbose=False): print("\t R = " + str(self.numR[-1])) print("\t F = " + str(self.numF[-1])) - if(self.t < self.tmax): + if self.t < self.tmax: self.run_epoch(runtime=self.tmax-self.t, dt=dt) return True @@ -243,10 +243,10 @@ def run(self, T, dt=0.1, checkpoints=None, verbose=False): #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_infections(self, t_idx=None): - if(t_idx is None): - return (self.numE[:] + self.numI[:] + self.numD_E[:] + self.numD_I[:]) + if t_idx is None: + return self.numE[:] + self.numI[:] + self.numD_E[:] + self.numD_I[:] else: - return (self.numE[t_idx] + self.numI[t_idx] + self.numD_E[t_idx] + self.numD_I[t_idx]) + return self.numE[t_idx] + self.numI[t_idx] + self.numD_E[t_idx] + self.numD_I[t_idx] #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ @@ -256,8 +256,8 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin plot_D_E='line', plot_D_I='line', combine_D=True, color_S='tab:green', color_E='orange', color_I='crimson', color_R='tab:blue', color_F='black', color_D_E='mediumorchid', color_D_I='mediumorchid', color_reference='#E0E0E0', - dashed_reference_results=None, dashed_reference_label='reference', - shaded_reference_results=None, shaded_reference_label='reference', + dashed_reference_results=None, dashed_reference_label='reference', + shaded_reference_results=None, shaded_reference_label='reference', vlines=[], vline_colors=[], vline_styles=[], vline_labels=[], ylim=None, xlim=None, legend=True, title=None, side_title=None, plot_percentages=True): @@ -266,7 +266,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Create an Axes object if None provided: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(not ax): + if not ax: fig, ax = pyplot.subplots() #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -279,16 +279,16 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin D_Iseries = self.numD_I/self.N if plot_percentages else self.numD_I Iseries = self.numI/self.N if plot_percentages else self.numI Rseries = self.numR/self.N if plot_percentages else self.numR - Sseries = self.numS/self.N if plot_percentages else self.numS + Sseries = self.numS/self.N if plot_percentages else self.numS #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Draw the reference data: + # Draw the reference data: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(dashed_reference_results): + if dashed_reference_results: dashedReference_tseries = dashed_reference_results.tseries[::int(self.N/100)] dashedReference_IDEstack = (dashed_reference_results.numI + dashed_reference_results.numD_I + dashed_reference_results.numD_E + dashed_reference_results.numE)[::int(self.N/100)] / (self.N if plot_percentages else 1) ax.plot(dashedReference_tseries, dashedReference_IDEstack, color='#E0E0E0', linestyle='--', label='$I+D+E$ ('+dashed_reference_label+')', zorder=0) - if(shaded_reference_results): + if shaded_reference_results: shadedReference_tseries = shaded_reference_results.tseries shadedReference_IDEstack = (shaded_reference_results.numI + shaded_reference_results.numD_I + shaded_reference_results.numD_E + shaded_reference_results.numE) / (self.N if plot_percentages else 1) ax.fill_between(shaded_reference_results.tseries, shadedReference_IDEstack, 0, color='#EFEFEF', label='$I+D+E$ ('+shaded_reference_label+')', zorder=0) @@ -298,102 +298,102 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin # Draw the stacked variables: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ topstack = numpy.zeros_like(self.tseries) - if(any(Fseries) and plot_F=='stacked'): + if (any(Fseries) and plot_F=='stacked'): ax.fill_between(numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, topstack+Fseries), topstack, color=color_F, alpha=0.5, label='$F$', zorder=2) ax.plot( numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, topstack+Fseries), color=color_F, zorder=3) topstack = topstack+Fseries - if(any(Eseries) and plot_E=='stacked'): + if (any(Eseries) and plot_E=='stacked'): ax.fill_between(numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, topstack+Eseries), topstack, color=color_E, alpha=0.5, label='$E$', zorder=2) ax.plot( numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, topstack+Eseries), color=color_E, zorder=3) topstack = topstack+Eseries - if(combine_D and plot_D_E=='stacked' and plot_D_I=='stacked'): + if (combine_D and plot_D_E=='stacked' and plot_D_I=='stacked'): ax.fill_between(numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, topstack+Dseries), topstack, color=color_D_E, alpha=0.5, label='$D_{all}$', zorder=2) ax.plot( numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, topstack+Dseries), color=color_D_E, zorder=3) topstack = topstack+Dseries else: - if(any(D_Eseries) and plot_D_E=='stacked'): + if (any(D_Eseries) and plot_D_E=='stacked'): ax.fill_between(numpy.ma.masked_where(D_Eseries<=0, self.tseries), numpy.ma.masked_where(D_Eseries<=0, topstack+D_Eseries), topstack, color=color_D_E, alpha=0.5, label='$D_E$', zorder=2) ax.plot( numpy.ma.masked_where(D_Eseries<=0, self.tseries), numpy.ma.masked_where(D_Eseries<=0, topstack+D_Eseries), color=color_D_E, zorder=3) topstack = topstack+D_Eseries - if(any(D_Iseries) and plot_D_I=='stacked'): + if (any(D_Iseries) and plot_D_I=='stacked'): ax.fill_between(numpy.ma.masked_where(D_Iseries<=0, self.tseries), numpy.ma.masked_where(D_Iseries<=0, topstack+D_Iseries), topstack, color=color_D_I, alpha=0.5, label='$D_I$', zorder=2) ax.plot( numpy.ma.masked_where(D_Iseries<=0, self.tseries), numpy.ma.masked_where(D_Iseries<=0, topstack+D_Iseries), color=color_D_I, zorder=3) topstack = topstack+D_Iseries - if(any(Iseries) and plot_I=='stacked'): + if (any(Iseries) and plot_I=='stacked'): ax.fill_between(numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, topstack+Iseries), topstack, color=color_I, alpha=0.5, label='$I$', zorder=2) ax.plot( numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, topstack+Iseries), color=color_I, zorder=3) topstack = topstack+Iseries - if(any(Rseries) and plot_R=='stacked'): + if (any(Rseries) and plot_R=='stacked'): ax.fill_between(numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, topstack+Rseries), topstack, color=color_R, alpha=0.5, label='$R$', zorder=2) ax.plot( numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, topstack+Rseries), color=color_R, zorder=3) topstack = topstack+Rseries - if(any(Sseries) and plot_S=='stacked'): + if (any(Sseries) and plot_S=='stacked'): ax.fill_between(numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, topstack+Sseries), topstack, color=color_S, alpha=0.5, label='$S$', zorder=2) ax.plot( numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, topstack+Sseries), color=color_S, zorder=3) topstack = topstack+Sseries - + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Draw the shaded variables: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(any(Fseries) and plot_F=='shaded'): + if (any(Fseries) and plot_F=='shaded'): ax.fill_between(numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, Fseries), 0, color=color_F, alpha=0.5, label='$F$', zorder=4) ax.plot( numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, Fseries), color=color_F, zorder=5) - if(any(Eseries) and plot_E=='shaded'): + if (any(Eseries) and plot_E=='shaded'): ax.fill_between(numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, Eseries), 0, color=color_E, alpha=0.5, label='$E$', zorder=4) ax.plot( numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, Eseries), color=color_E, zorder=5) - if(combine_D and (any(Dseries) and plot_D_E=='shaded' and plot_D_E=='shaded')): + if (combine_D and (any(Dseries) and plot_D_E=='shaded' and plot_D_E=='shaded')): ax.fill_between(numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, Dseries), 0, color=color_D_E, alpha=0.5, label='$D_{all}$', zorder=4) ax.plot( numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, Dseries), color=color_D_E, zorder=5) else: - if(any(D_Eseries) and plot_D_E=='shaded'): + if (any(D_Eseries) and plot_D_E=='shaded'): ax.fill_between(numpy.ma.masked_where(D_Eseries<=0, self.tseries), numpy.ma.masked_where(D_Eseries<=0, D_Eseries), 0, color=color_D_E, alpha=0.5, label='$D_E$', zorder=4) ax.plot( numpy.ma.masked_where(D_Eseries<=0, self.tseries), numpy.ma.masked_where(D_Eseries<=0, D_Eseries), color=color_D_E, zorder=5) - if(any(D_Iseries) and plot_D_I=='shaded'): + if (any(D_Iseries) and plot_D_I=='shaded'): ax.fill_between(numpy.ma.masked_where(D_Iseries<=0, self.tseries), numpy.ma.masked_where(D_Iseries<=0, D_Iseries), 0, color=color_D_I, alpha=0.5, label='$D_I$', zorder=4) ax.plot( numpy.ma.masked_where(D_Iseries<=0, self.tseries), numpy.ma.masked_where(D_Iseries<=0, D_Iseries), color=color_D_I, zorder=5) - if(any(Iseries) and plot_I=='shaded'): + if (any(Iseries) and plot_I=='shaded'): ax.fill_between(numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, Iseries), 0, color=color_I, alpha=0.5, label='$I$', zorder=4) ax.plot( numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, Iseries), color=color_I, zorder=5) - if(any(Sseries) and plot_S=='shaded'): + if (any(Sseries) and plot_S=='shaded'): ax.fill_between(numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, Sseries), 0, color=color_S, alpha=0.5, label='$S$', zorder=4) ax.plot( numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, Sseries), color=color_S, zorder=5) - if(any(Rseries) and plot_R=='shaded'): + if (any(Rseries) and plot_R=='shaded'): ax.fill_between(numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, Rseries), 0, color=color_R, alpha=0.5, label='$R$', zorder=4) ax.plot( numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, Rseries), color=color_R, zorder=5) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Draw the line variables: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(any(Fseries) and plot_F=='line'): + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + if (any(Fseries) and plot_F=='line'): ax.plot(numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, Fseries), color=color_F, label='$F$', zorder=6) - if(any(Eseries) and plot_E=='line'): + if (any(Eseries) and plot_E=='line'): ax.plot(numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, Eseries), color=color_E, label='$E$', zorder=6) - if(combine_D and (any(Dseries) and plot_D_E=='line' and plot_D_E=='line')): + if (combine_D and (any(Dseries) and plot_D_E=='line' and plot_D_E=='line')): ax.plot(numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, Dseries), color=color_D_E, label='$D_{all}$', zorder=6) else: - if(any(D_Eseries) and plot_D_E=='line'): + if (any(D_Eseries) and plot_D_E=='line'): ax.plot(numpy.ma.masked_where(D_Eseries<=0, self.tseries), numpy.ma.masked_where(D_Eseries<=0, D_Eseries), color=color_D_E, label='$D_E$', zorder=6) - if(any(D_Iseries) and plot_D_I=='line'): + if (any(D_Iseries) and plot_D_I=='line'): ax.plot(numpy.ma.masked_where(D_Iseries<=0, self.tseries), numpy.ma.masked_where(D_Iseries<=0, D_Iseries), color=color_D_I, label='$D_I$', zorder=6) - if(any(Iseries) and plot_I=='line'): + if (any(Iseries) and plot_I=='line'): ax.plot(numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, Iseries), color=color_I, label='$I$', zorder=6) - if(any(Sseries) and plot_S=='line'): + if (any(Sseries) and plot_S=='line'): ax.plot(numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, Sseries), color=color_S, label='$S$', zorder=6) - if(any(Rseries) and plot_R=='line'): + if (any(Rseries) and plot_R=='line'): ax.plot(numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, Rseries), color=color_R, label='$R$', zorder=6) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Draw the vertical line annotations: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(len(vlines)>0 and len(vline_colors)==0): + if (len(vlines)>0 and len(vline_colors)==0): vline_colors = ['gray']*len(vlines) - if(len(vlines)>0 and len(vline_labels)==0): + if (len(vlines)>0 and len(vline_labels)==0): vline_labels = [None]*len(vlines) - if(len(vlines)>0 and len(vline_styles)==0): + if (len(vlines)>0 and len(vline_styles)==0): vline_styles = [':']*len(vlines) for vline_x, vline_color, vline_style, vline_label in zip(vlines, vline_colors, vline_styles, vline_labels): - if(vline_x is not None): + if vline_x is not None: ax.axvline(x=vline_x, color=vline_color, linestyle=vline_style, alpha=1, label=vline_label) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -403,17 +403,17 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin ax.set_ylabel('percent of population' if plot_percentages else 'number of individuals') ax.set_xlim(0, (max(self.tseries) if not xlim else xlim)) ax.set_ylim(0, ylim) - if(plot_percentages): + if plot_percentages: ax.set_yticklabels(['{:,.0%}'.format(y) for y in ax.get_yticks()]) - if(legend): + if legend: legend_handles, legend_labels = ax.get_legend_handles_labels() ax.legend(legend_handles[::-1], legend_labels[::-1], loc='upper right', facecolor='white', edgecolor='none', framealpha=0.9, prop={'size': 8}) - if(title): + if title: ax.set_title(title, size=12) - if(side_title): + if side_title: ax.annotate(side_title, (0, 0.5), xytext=(-45, 0), ha='right', va='center', size=12, rotation=90, xycoords='axes fraction', textcoords='offset points') - + return ax @@ -424,8 +424,8 @@ def figure_basic(self, plot_S='line', plot_E='line', plot_I='line',plot_R='line' plot_D_E='line', plot_D_I='line', combine_D=True, color_S='tab:green', color_E='orange', color_I='crimson', color_R='tab:blue', color_F='black', color_D_E='mediumorchid', color_D_I='mediumorchid', color_reference='#E0E0E0', - dashed_reference_results=None, dashed_reference_label='reference', - shaded_reference_results=None, shaded_reference_label='reference', + dashed_reference_results=None, dashed_reference_label='reference', + shaded_reference_results=None, shaded_reference_label='reference', vlines=[], vline_colors=[], vline_styles=[], vline_labels=[], ylim=None, xlim=None, legend=True, title=None, side_title=None, plot_percentages=True, figsize=(12,8), use_seaborn=True, show=True): @@ -434,7 +434,7 @@ def figure_basic(self, plot_S='line', plot_E='line', plot_I='line',plot_R='line' fig, ax = pyplot.subplots(figsize=figsize) - if(use_seaborn): + if use_seaborn: import seaborn seaborn.set_style('ticks') seaborn.despine() @@ -443,12 +443,12 @@ def figure_basic(self, plot_S='line', plot_E='line', plot_I='line',plot_R='line' plot_D_E=plot_D_E, plot_D_I=plot_D_I, combine_D=combine_D, color_S=color_S, color_E=color_E, color_I=color_I, color_R=color_R, color_F=color_F, color_D_E=color_D_E, color_D_I=color_D_I, color_reference=color_reference, - dashed_reference_results=dashed_reference_results, dashed_reference_label=dashed_reference_label, - shaded_reference_results=shaded_reference_results, shaded_reference_label=shaded_reference_label, + dashed_reference_results=dashed_reference_results, dashed_reference_label=dashed_reference_label, + shaded_reference_results=shaded_reference_results, shaded_reference_label=shaded_reference_label, vlines=vlines, vline_colors=vline_colors, vline_styles=vline_styles, vline_labels=vline_labels, ylim=ylim, xlim=xlim, legend=legend, title=title, side_title=side_title, plot_percentages=plot_percentages) - if(show): + if show: pyplot.show() return fig, ax @@ -461,8 +461,8 @@ def figure_infections(self, plot_S=False, plot_E='stacked', plot_I='stacked',plo plot_D_E='stacked', plot_D_I='stacked', combine_D=True, color_S='tab:green', color_E='orange', color_I='crimson', color_R='tab:blue', color_F='black', color_D_E='mediumorchid', color_D_I='mediumorchid', color_reference='#E0E0E0', - dashed_reference_results=None, dashed_reference_label='reference', - shaded_reference_results=None, shaded_reference_label='reference', + dashed_reference_results=None, dashed_reference_label='reference', + shaded_reference_results=None, shaded_reference_label='reference', vlines=[], vline_colors=[], vline_styles=[], vline_labels=[], ylim=None, xlim=None, legend=True, title=None, side_title=None, plot_percentages=True, figsize=(12,8), use_seaborn=True, show=True): @@ -471,7 +471,7 @@ def figure_infections(self, plot_S=False, plot_E='stacked', plot_I='stacked',plo fig, ax = pyplot.subplots(figsize=figsize) - if(use_seaborn): + if use_seaborn: import seaborn seaborn.set_style('ticks') seaborn.despine() @@ -480,12 +480,12 @@ def figure_infections(self, plot_S=False, plot_E='stacked', plot_I='stacked',plo plot_D_E=plot_D_E, plot_D_I=plot_D_I, combine_D=combine_D, color_S=color_S, color_E=color_E, color_I=color_I, color_R=color_R, color_F=color_F, color_D_E=color_D_E, color_D_I=color_D_I, color_reference=color_reference, - dashed_reference_results=dashed_reference_results, dashed_reference_label=dashed_reference_label, - shaded_reference_results=shaded_reference_results, shaded_reference_label=shaded_reference_label, - vlines=vlines, vline_colors=vline_colors, vline_styles=vline_styles, vline_labels=vline_labels, + dashed_reference_results=dashed_reference_results, dashed_reference_label=dashed_reference_label, + shaded_reference_results=shaded_reference_results, shaded_reference_label=shaded_reference_label, + vlines=vlines, vline_colors=vline_colors, vline_styles=vline_styles, vline_labels=vline_labels, ylim=ylim, xlim=xlim, legend=legend, title=title, side_title=side_title, plot_percentages=plot_percentages) - if(show): + if show: pyplot.show() return fig, ax @@ -504,14 +504,14 @@ class SEIRSNetworkModel(): Params: G Network adjacency matrix (numpy array) or Networkx graph object. beta Rate of transmission (exposure) (global) beta_local Rate(s) of transmission (exposure) for adjacent individuals (optional) - sigma Rate of infection (upon exposure) - gamma Rate of recovery (upon infection) - xi Rate of re-susceptibility (upon recovery) - mu_I Rate of infection-related death - mu_0 Rate of baseline death + sigma Rate of infection (upon exposure) + gamma Rate of recovery (upon infection) + xi Rate of re-susceptibility (upon recovery) + mu_I Rate of infection-related death + mu_0 Rate of baseline death nu Rate of baseline birth p Probability of interaction outside adjacent nodes - + Q Quarantine adjacency matrix (numpy array) or Networkx graph object. beta_D Rate of transmission (exposure) for individuals with detected infections (global) beta_local Rate(s) of transmission (exposure) for adjacent individuals with detected infections (optional) @@ -525,14 +525,14 @@ class SEIRSNetworkModel(): psi_E Probability of positive test results for exposed individuals psi_I Probability of positive test results for exposed individuals q Probability of quarantined individuals interaction outside adjacent nodes - - initE Init number of exposed individuals - initI Init number of infectious individuals + + initE Init number of exposed individuals + initI Init number of infectious individuals initD_E Init number of detected infectious individuals - initD_I Init number of detected infectious individuals - initR Init number of recovered individuals + initD_I Init number of detected infectious individuals + initR Init number of recovered individuals initF Init number of infection-related fatalities - (all remaining nodes initialized susceptible) + (all remaining nodes initialized susceptible) """ def __init__(self, G, beta, sigma, gamma, xi=0, mu_I=0, mu_0=0, nu=0, beta_local=None, p=0, @@ -547,22 +547,22 @@ def __init__(self, G, beta, sigma, gamma, xi=0, mu_I=0, mu_0=0, nu=0, beta_local #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Setup Quarantine Adjacency matrix: - if(Q is None): + if Q is None: Q = G # If no Q graph is provided, use G in its place self.update_Q(Q) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Model Parameters: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - self.parameters = { 'beta':beta, 'sigma':sigma, 'gamma':gamma, 'xi':xi, 'mu_I':mu_I, 'mu_0':mu_0, 'nu':nu, - 'beta_D':beta_D, 'sigma_D':sigma_D, 'gamma_D':gamma_D, 'mu_D':mu_D, + self.parameters = { 'beta':beta, 'sigma':sigma, 'gamma':gamma, 'xi':xi, 'mu_I':mu_I, 'mu_0':mu_0, 'nu':nu, + 'beta_D':beta_D, 'sigma_D':sigma_D, 'gamma_D':gamma_D, 'mu_D':mu_D, 'beta_local':beta_local, 'beta_D_local':beta_D_local, 'p':p,'q':q, 'theta_E':theta_E, 'theta_I':theta_I, 'phi_E':phi_E, 'phi_I':phi_I, 'psi_E':psi_E, 'psi_I':psi_I } self.update_parameters() #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Each node can undergo up to 4 transitions (sans vitality/re-susceptibility returns to S state), - # so there are ~numNodes*4 events/timesteps expected; initialize numNodes*5 timestep slots to start + # so there are ~numNodes*4 events/timesteps expected; initialize numNodes*5 timestep slots to start # (will be expanded during run if needed) self.tseries = numpy.zeros(5*self.numNodes) self.numE = numpy.zeros(5*self.numNodes) @@ -581,7 +581,7 @@ def __init__(self, G, beta, sigma, gamma, xi=0, mu_I=0, mu_0=0, nu=0, beta_local self.tmax = 0 # will be set when run() is called self.tidx = 0 self.tseries[0] = 0 - + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Initialize Counts of inidividuals with each state: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -593,7 +593,7 @@ def __init__(self, G, beta, sigma, gamma, xi=0, mu_I=0, mu_0=0, nu=0, beta_local self.numF[0] = int(initF) self.numS[0] = self.numNodes - self.numE[0] - self.numI[0] - self.numD_E[0] - self.numD_I[0] - self.numR[0] - self.numF[0] self.N[0] = self.numS[0] + self.numE[0] + self.numI[0] + self.numD_E[0] + self.numD_I[0] + self.numR[0] - + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Node states: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -609,11 +609,11 @@ def __init__(self, G, beta, sigma, gamma, xi=0, mu_I=0, mu_0=0, nu=0, beta_local numpy.random.shuffle(self.X) self.store_Xseries = store_Xseries - if(store_Xseries): + if store_Xseries: self.Xseries = numpy.zeros(shape=(5*self.numNodes, self.numNodes), dtype='uint8') self.Xseries[0,:] = self.X.T - self.transitions = { + self.transitions = { 'StoE': {'currentState':self.S, 'newState':self.E}, 'EtoI': {'currentState':self.E, 'newState':self.I}, 'ItoR': {'currentState':self.I, 'newState':self.R}, @@ -631,7 +631,7 @@ def __init__(self, G, beta, sigma, gamma, xi=0, mu_I=0, mu_0=0, nu=0, beta_local # Initialize node subgroup data series: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ self.nodeGroupData = None - if(node_groups): + if node_groups: self.nodeGroupData = {} for groupName, nodeList in node_groups.items(): self.nodeGroupData[groupName] = {'nodes': numpy.array(nodeList), @@ -653,7 +653,7 @@ def __init__(self, G, beta, sigma, gamma, xi=0, mu_I=0, mu_0=0, nu=0, beta_local self.nodeGroupData[groupName]['numF'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.F) self.nodeGroupData[groupName]['N'][0] = self.nodeGroupData[groupName]['numS'][0] + self.nodeGroupData[groupName]['numE'][0] + self.nodeGroupData[groupName]['numI'][0] + self.nodeGroupData[groupName]['numD_E'][0] + self.nodeGroupData[groupName]['numD_I'][0] + self.nodeGroupData[groupName]['numR'][0] - + #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ @@ -673,7 +673,7 @@ def update_parameters(self): self.mu_0 = numpy.array(self.parameters['mu_0']).reshape((self.numNodes, 1)) if isinstance(self.parameters['mu_0'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['mu_0'], shape=(self.numNodes,1)) self.nu = numpy.array(self.parameters['nu']).reshape((self.numNodes, 1)) if isinstance(self.parameters['nu'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['nu'], shape=(self.numNodes,1)) self.p = numpy.array(self.parameters['p']).reshape((self.numNodes, 1)) if isinstance(self.parameters['p'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['p'], shape=(self.numNodes,1)) - + # Testing-related parameters: self.beta_D = (numpy.array(self.parameters['beta_D']).reshape((self.numNodes, 1)) if isinstance(self.parameters['beta_D'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['beta_D'], shape=(self.numNodes,1))) if self.parameters['beta_D'] is not None else self.beta self.sigma_D = (numpy.array(self.parameters['sigma_D']).reshape((self.numNodes, 1)) if isinstance(self.parameters['sigma_D'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['sigma_D'], shape=(self.numNodes,1))) if self.parameters['sigma_D'] is not None else self.sigma @@ -688,45 +688,45 @@ def update_parameters(self): self.q = numpy.array(self.parameters['q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['q'], shape=(self.numNodes,1)) #Local transmission parameters: - if(self.parameters['beta_local'] is not None): - if(isinstance(self.parameters['beta_local'], (list, numpy.ndarray))): - if(isinstance(self.parameters['beta_local'], list)): + if self.parameters['beta_local'] is not None: + if isinstance(self.parameters['beta_local'], (list, numpy.ndarray)): + if isinstance(self.parameters['beta_local'], list): self.beta_local = numpy.array(self.parameters['beta_local']) else: # is numpy.ndarray self.beta_local = self.parameters['beta_local'] - if(self.beta_local.ndim == 1): + if self.beta_local.ndim == 1: self.beta_local.reshape((self.numNodes, 1)) - elif(self.beta_local.ndim == 2): + elif self.beta_local.ndim == 2: self.beta_local.reshape((self.numNodes, self.numNodes)) else: self.beta_local = numpy.full_like(self.beta, fill_value=self.parameters['beta_local']) else: self.beta_local = self.beta #---------------------------------------- - if(self.parameters['beta_D_local'] is not None): - if(isinstance(self.parameters['beta_D_local'], (list, numpy.ndarray))): - if(isinstance(self.parameters['beta_D_local'], list)): + if self.parameters['beta_D_local'] is not None: + if isinstance(self.parameters['beta_D_local'], (list, numpy.ndarray)): + if isinstance(self.parameters['beta_D_local'], list): self.beta_D_local = numpy.array(self.parameters['beta_D_local']) else: # is numpy.ndarray self.beta_D_local = self.parameters['beta_D_local'] - if(self.beta_D_local.ndim == 1): + if self.beta_D_local.ndim == 1: self.beta_D_local.reshape((self.numNodes, 1)) - elif(self.beta_D_local.ndim == 2): + elif self.beta_D_local.ndim == 2: self.beta_D_local.reshape((self.numNodes, self.numNodes)) else: self.beta_D_local = numpy.full_like(self.beta_D, fill_value=self.parameters['beta_D_local']) else: self.beta_D_local = self.beta_D - + # Pre-multiply beta values by the adjacency matrix ("transmission weight connections") - if(self.beta_local.ndim == 1): + if self.beta_local.ndim == 1: self.A_beta = scipy.sparse.csr_matrix.multiply(self.A, numpy.tile(self.beta_local, (1,self.numNodes))).tocsr() - elif(self.beta_local.ndim == 2): + elif self.beta_local.ndim == 2: self.A_beta = scipy.sparse.csr_matrix.multiply(self.A, self.beta_local).tocsr() # Pre-multiply beta_D values by the quarantine adjacency matrix ("transmission weight connections") - if(self.beta_D_local.ndim == 1): + if self.beta_D_local.ndim == 1: self.A_Q_beta_D = scipy.sparse.csr_matrix.multiply(self.A_Q, numpy.tile(self.beta_D_local, (1,self.numNodes))).tocsr() - elif(self.beta_D_local.ndim == 2): + elif self.beta_D_local.ndim == 2: self.A_Q_beta_D = scipy.sparse.csr_matrix.multiply(self.A_Q, self.beta_D_local).tocsr() #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -780,9 +780,9 @@ def update_Q(self, new_Q): #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def update_scenario_flags(self): - self.testing_scenario = ( (numpy.any(self.psi_I) and (numpy.any(self.theta_I) or numpy.any(self.phi_I))) + self.testing_scenario = ( (numpy.any(self.psi_I) and (numpy.any(self.theta_I) or numpy.any(self.phi_I))) or (numpy.any(self.psi_E) and (numpy.any(self.theta_E) or numpy.any(self.phi_E))) ) - self.tracing_scenario = ( (numpy.any(self.psi_E) and numpy.any(self.phi_E)) + self.tracing_scenario = ( (numpy.any(self.psi_E) and numpy.any(self.phi_E)) or (numpy.any(self.psi_I) and numpy.any(self.phi_I)) ) self.vitality_scenario = (numpy.any(self.mu_0) and numpy.any(self.nu)) self.resusceptibility_scenario = (numpy.any(self.xi)) @@ -790,36 +790,36 @@ def update_scenario_flags(self): #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_infections(self, t_idx=None): - if(t_idx is None): - return (self.numE[:] + self.numI[:] + self.numD_E[:] + self.numD_I[:]) + if t_idx is None: + return self.numE[:] + self.numI[:] + self.numD_E[:] + self.numD_I[:] else: - return (self.numE[t_idx] + self.numI[t_idx] + self.numD_E[t_idx] + self.numD_I[t_idx]) + return self.numE[t_idx] + self.numI[t_idx] + self.numD_E[t_idx] + self.numD_I[t_idx] #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - + def calc_propensities(self): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Pre-calculate matrix multiplication terms that may be used in multiple propensity calculations, # and check to see if their computation is necessary before doing the multiplication transmissionTerms_I = numpy.zeros(shape=(self.numNodes,1)) - if(numpy.any(self.numI[self.tidx]) + if (numpy.any(self.numI[self.tidx]) and numpy.any(self.beta!=0)): transmissionTerms_I = numpy.asarray( scipy.sparse.csr_matrix.dot(self.A_beta, self.X==self.I) ) transmissionTerms_DI = numpy.zeros(shape=(self.numNodes,1)) - if(self.testing_scenario + if (self.testing_scenario and numpy.any(self.numD_I[self.tidx]) and numpy.any(self.beta_D)): transmissionTerms_DI = numpy.asarray( scipy.sparse.csr_matrix.dot(self.A_Q_beta_D, self.X==self.D_I) ) numContacts_D = numpy.zeros(shape=(self.numNodes,1)) - if(self.tracing_scenario + if (self.tracing_scenario and (numpy.any(self.numD_E[self.tidx]) or numpy.any(self.numD_I[self.tidx]))): numContacts_D = numpy.asarray( scipy.sparse.csr_matrix.dot( self.A, ((self.X==self.D_E)|(self.X==self.D_I)) ) ) - + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -849,9 +849,9 @@ def calc_propensities(self): propensities__toS = self.nu*(self.X!=self.F) - propensities = numpy.hstack([propensities_StoE, propensities_EtoI, - propensities_ItoR, propensities_ItoF, - propensities_EtoDE, propensities_ItoDI, propensities_DEtoDI, + propensities = numpy.hstack([propensities_StoE, propensities_EtoI, + propensities_ItoR, propensities_ItoF, + propensities_EtoDE, propensities_ItoDI, propensities_DEtoDI, propensities_DItoR, propensities_DItoF, propensities_RtoS, propensities__toS]) @@ -861,7 +861,7 @@ def calc_propensities(self): #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def increase_data_series_length(self): self.tseries= numpy.pad(self.tseries, [(0, 5*self.numNodes)], mode='constant', constant_values=0) @@ -874,10 +874,10 @@ def increase_data_series_length(self): self.numF = numpy.pad(self.numF, [(0, 5*self.numNodes)], mode='constant', constant_values=0) self.N = numpy.pad(self.N, [(0, 5*self.numNodes)], mode='constant', constant_values=0) - if(self.store_Xseries): + if self.store_Xseries: self.Xseries = numpy.pad(self.Xseries, [(0, 5*self.numNodes), (0,0)], mode='constant', constant_values=0) - if(self.nodeGroupData): + if self.nodeGroupData: for groupName in self.nodeGroupData: self.nodeGroupData[groupName]['numS'] = numpy.pad(self.nodeGroupData[groupName]['numS'], [(0, 5*self.numNodes)], mode='constant', constant_values=0) self.nodeGroupData[groupName]['numE'] = numpy.pad(self.nodeGroupData[groupName]['numE'], [(0, 5*self.numNodes)], mode='constant', constant_values=0) @@ -890,7 +890,7 @@ def increase_data_series_length(self): return None -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def finalize_data_series(self): self.tseries= numpy.array(self.tseries, dtype=float)[:self.tidx+1] @@ -903,10 +903,10 @@ def finalize_data_series(self): self.numF = numpy.array(self.numF, dtype=float)[:self.tidx+1] self.N = numpy.array(self.N, dtype=float)[:self.tidx+1] - if(self.store_Xseries): + if self.store_Xseries: self.Xseries = self.Xseries[:self.tidx+1, :] - if(self.nodeGroupData): + if self.nodeGroupData: for groupName in self.nodeGroupData: self.nodeGroupData[groupName]['numS'] = numpy.array(self.nodeGroupData[groupName]['numS'], dtype=float)[:self.tidx+1] self.nodeGroupData[groupName]['numE'] = numpy.array(self.nodeGroupData[groupName]['numE'], dtype=float)[:self.tidx+1] @@ -920,11 +920,11 @@ def finalize_data_series(self): return None #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def run_iteration(self): - if(self.tidx >= len(self.tseries)-1): + if self.tidx >= len(self.tseries) - 1: # Room has run out in the timeseries storage arrays; double the size of these arrays: self.increase_data_series_length() @@ -940,7 +940,7 @@ def run_iteration(self): propensities, transitionTypes = self.calc_propensities() # Terminate when probability of all events is 0: - if(propensities.sum() <= 0.0): + if propensities.sum() <= 0.0: self.finalize_data_series() return False @@ -971,7 +971,7 @@ def run_iteration(self): self.X[transitionNode] = self.transitions[transitionType]['newState'] self.tidx += 1 - + self.tseries[self.tidx] = self.t self.numS[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.S), a_min=0, a_max=self.numNodes) self.numE[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.E), a_min=0, a_max=self.numNodes) @@ -982,10 +982,10 @@ def run_iteration(self): self.numF[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.F), a_min=0, a_max=self.numNodes) self.N[self.tidx] = numpy.clip((self.numS[self.tidx] + self.numE[self.tidx] + self.numI[self.tidx] + self.numD_E[self.tidx] + self.numD_I[self.tidx] + self.numR[self.tidx]), a_min=0, a_max=self.numNodes) - if(self.store_Xseries): + if self.store_Xseries: self.Xseries[self.tidx,:] = self.X.T - if(self.nodeGroupData): + if self.nodeGroupData: for groupName in self.nodeGroupData: self.nodeGroupData[groupName]['numS'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.S) self.nodeGroupData[groupName]['numE'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.E) @@ -999,7 +999,7 @@ def run_iteration(self): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Terminate if tmax reached or num infectious and num exposed is 0: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(self.t >= self.tmax or (self.numI[self.tidx]<1 and self.numE[self.tidx]<1 and self.numD_E[self.tidx]<1 and self.numD_I[self.tidx]<1)): + if (self.t >= self.tmax or (self.numI[self.tidx]<1 and self.numE[self.tidx]<1 and self.numD_E[self.tidx]<1 and self.numD_I[self.tidx]<1)): self.finalize_data_series() return False @@ -1012,7 +1012,7 @@ def run_iteration(self): #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def run(self, T, checkpoints=None, print_interval=10, verbose='t'): - if(T>0): + if T > 0: self.tmax += T else: return False @@ -1020,14 +1020,14 @@ def run(self, T, checkpoints=None, print_interval=10, verbose='t'): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Pre-process checkpoint values: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(checkpoints): + if checkpoints: numCheckpoints = len(checkpoints['t']) for chkpt_param, chkpt_values in checkpoints.items(): assert(isinstance(chkpt_values, (list, numpy.ndarray)) and len(chkpt_values)==numCheckpoints), "Expecting a list of values with length equal to number of checkpoint times ("+str(numCheckpoints)+") for each checkpoint parameter." checkpointIdx = numpy.searchsorted(checkpoints['t'], self.t) # Finds 1st index in list greater than given val - if(checkpointIdx >= numCheckpoints): + if checkpointIdx >= numCheckpoints: # We are out of checkpoints, stop checking them: - checkpoints = None + checkpoints = None else: checkpointTime = checkpoints['t'][checkpointIdx] @@ -1042,36 +1042,36 @@ def run(self, T, checkpoints=None, print_interval=10, verbose='t'): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Handle checkpoints if applicable: - if(checkpoints): - if(self.t >= checkpointTime): - if(verbose is not False): + if checkpoints: + if self.t >= checkpointTime: + if verbose is not False: print("[Checkpoint: Updating parameters]") # A checkpoint has been reached, update param values: - if('G' in list(checkpoints.keys())): + if 'G' in list(checkpoints.keys()): self.update_G(checkpoints['G'][checkpointIdx]) - if('Q' in list(checkpoints.keys())): + if 'Q' in list(checkpoints.keys()): self.update_Q(checkpoints['Q'][checkpointIdx]) for param in list(self.parameters.keys()): - if(param in list(checkpoints.keys())): + if param in list(checkpoints.keys()): self.parameters.update({param: checkpoints[param][checkpointIdx]}) # Update parameter data structures and scenario flags: self.update_parameters() # Update the next checkpoint time: checkpointIdx = numpy.searchsorted(checkpoints['t'], self.t) # Finds 1st index in list greater than given val - if(checkpointIdx >= numCheckpoints): + if checkpointIdx >= numCheckpoints: # We are out of checkpoints, stop checking them: - checkpoints = None + checkpoints = None else: checkpointTime = checkpoints['t'][checkpointIdx] #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(print_interval): - if(print_reset and (int(self.t) % print_interval == 0)): - if(verbose=="t"): + if print_interval: + if (print_reset and (int(self.t) % print_interval == 0)): + if verbose=="t": print("t = %.2f" % self.t) - if(verbose==True): + if verbose==True: print("t = %.2f" % self.t) print("\t S = " + str(self.numS[self.tidx])) print("\t E = " + str(self.numE[self.tidx])) @@ -1081,7 +1081,7 @@ def run(self, T, checkpoints=None, print_interval=10, verbose='t'): print("\t R = " + str(self.numR[self.tidx])) print("\t F = " + str(self.numF[self.tidx])) print_reset = False - elif(not print_reset and (int(self.t) % 10 != 0)): + elif (not print_reset and (int(self.t) % 10 != 0)): print_reset = True return True @@ -1094,8 +1094,8 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin plot_D_E='line', plot_D_I='line', combine_D=True, color_S='tab:green', color_E='orange', color_I='crimson', color_R='tab:blue', color_F='black', color_D_E='mediumorchid', color_D_I='mediumorchid', color_reference='#E0E0E0', - dashed_reference_results=None, dashed_reference_label='reference', - shaded_reference_results=None, shaded_reference_label='reference', + dashed_reference_results=None, dashed_reference_label='reference', + shaded_reference_results=None, shaded_reference_label='reference', vlines=[], vline_colors=[], vline_styles=[], vline_labels=[], ylim=None, xlim=None, legend=True, title=None, side_title=None, plot_percentages=True): @@ -1104,7 +1104,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Create an Axes object if None provided: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(not ax): + if not ax: fig, ax = pyplot.subplots() #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -1117,16 +1117,16 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin D_Iseries = self.numD_I/self.numNodes if plot_percentages else self.numD_I Iseries = self.numI/self.numNodes if plot_percentages else self.numI Rseries = self.numR/self.numNodes if plot_percentages else self.numR - Sseries = self.numS/self.numNodes if plot_percentages else self.numS + Sseries = self.numS/self.numNodes if plot_percentages else self.numS #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Draw the reference data: + # Draw the reference data: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(dashed_reference_results): + if dashed_reference_results: dashedReference_tseries = dashed_reference_results.tseries[::int(self.numNodes/100)] dashedReference_IDEstack = (dashed_reference_results.numI + dashed_reference_results.numD_I + dashed_reference_results.numD_E + dashed_reference_results.numE)[::int(self.numNodes/100)] / (self.numNodes if plot_percentages else 1) ax.plot(dashedReference_tseries, dashedReference_IDEstack, color='#E0E0E0', linestyle='--', label='$I+D+E$ ('+dashed_reference_label+')', zorder=0) - if(shaded_reference_results): + if shaded_reference_results: shadedReference_tseries = shaded_reference_results.tseries shadedReference_IDEstack = (shaded_reference_results.numI + shaded_reference_results.numD_I + shaded_reference_results.numD_E + shaded_reference_results.numE) / (self.numNodes if plot_percentages else 1) ax.fill_between(shaded_reference_results.tseries, shadedReference_IDEstack, 0, color='#EFEFEF', label='$I+D+E$ ('+shaded_reference_label+')', zorder=0) @@ -1136,102 +1136,102 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin # Draw the stacked variables: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ topstack = numpy.zeros_like(self.tseries) - if(any(Fseries) and plot_F=='stacked'): + if (any(Fseries) and plot_F=='stacked'): ax.fill_between(numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, topstack+Fseries), topstack, color=color_F, alpha=0.5, label='$F$', zorder=2) ax.plot( numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, topstack+Fseries), color=color_F, zorder=3) topstack = topstack+Fseries - if(any(Eseries) and plot_E=='stacked'): + if (any(Eseries) and plot_E=='stacked'): ax.fill_between(numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, topstack+Eseries), topstack, color=color_E, alpha=0.5, label='$E$', zorder=2) ax.plot( numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, topstack+Eseries), color=color_E, zorder=3) topstack = topstack+Eseries - if(combine_D and plot_D_E=='stacked' and plot_D_I=='stacked'): + if (combine_D and plot_D_E=='stacked' and plot_D_I=='stacked'): ax.fill_between(numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, topstack+Dseries), topstack, color=color_D_E, alpha=0.5, label='$D_{all}$', zorder=2) ax.plot( numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, topstack+Dseries), color=color_D_E, zorder=3) topstack = topstack+Dseries else: - if(any(D_Eseries) and plot_D_E=='stacked'): + if (any(D_Eseries) and plot_D_E=='stacked'): ax.fill_between(numpy.ma.masked_where(D_Eseries<=0, self.tseries), numpy.ma.masked_where(D_Eseries<=0, topstack+D_Eseries), topstack, color=color_D_E, alpha=0.5, label='$D_E$', zorder=2) ax.plot( numpy.ma.masked_where(D_Eseries<=0, self.tseries), numpy.ma.masked_where(D_Eseries<=0, topstack+D_Eseries), color=color_D_E, zorder=3) topstack = topstack+D_Eseries - if(any(D_Iseries) and plot_D_I=='stacked'): + if (any(D_Iseries) and plot_D_I=='stacked'): ax.fill_between(numpy.ma.masked_where(D_Iseries<=0, self.tseries), numpy.ma.masked_where(D_Iseries<=0, topstack+D_Iseries), topstack, color=color_D_I, alpha=0.5, label='$D_I$', zorder=2) ax.plot( numpy.ma.masked_where(D_Iseries<=0, self.tseries), numpy.ma.masked_where(D_Iseries<=0, topstack+D_Iseries), color=color_D_I, zorder=3) topstack = topstack+D_Iseries - if(any(Iseries) and plot_I=='stacked'): + if (any(Iseries) and plot_I=='stacked'): ax.fill_between(numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, topstack+Iseries), topstack, color=color_I, alpha=0.5, label='$I$', zorder=2) ax.plot( numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, topstack+Iseries), color=color_I, zorder=3) topstack = topstack+Iseries - if(any(Rseries) and plot_R=='stacked'): + if (any(Rseries) and plot_R=='stacked'): ax.fill_between(numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, topstack+Rseries), topstack, color=color_R, alpha=0.5, label='$R$', zorder=2) ax.plot( numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, topstack+Rseries), color=color_R, zorder=3) topstack = topstack+Rseries - if(any(Sseries) and plot_S=='stacked'): + if (any(Sseries) and plot_S=='stacked'): ax.fill_between(numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, topstack+Sseries), topstack, color=color_S, alpha=0.5, label='$S$', zorder=2) ax.plot( numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, topstack+Sseries), color=color_S, zorder=3) topstack = topstack+Sseries - + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Draw the shaded variables: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(any(Fseries) and plot_F=='shaded'): + if (any(Fseries) and plot_F=='shaded'): ax.fill_between(numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, Fseries), 0, color=color_F, alpha=0.5, label='$F$', zorder=4) ax.plot( numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, Fseries), color=color_F, zorder=5) - if(any(Eseries) and plot_E=='shaded'): + if (any(Eseries) and plot_E=='shaded'): ax.fill_between(numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, Eseries), 0, color=color_E, alpha=0.5, label='$E$', zorder=4) ax.plot( numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, Eseries), color=color_E, zorder=5) - if(combine_D and (any(Dseries) and plot_D_E=='shaded' and plot_D_I=='shaded')): + if (combine_D and (any(Dseries) and plot_D_E=='shaded' and plot_D_I=='shaded')): ax.fill_between(numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, Dseries), 0, color=color_D_E, alpha=0.5, label='$D_{all}$', zorder=4) ax.plot( numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, Dseries), color=color_D_E, zorder=5) else: - if(any(D_Eseries) and plot_D_E=='shaded'): + if (any(D_Eseries) and plot_D_E=='shaded'): ax.fill_between(numpy.ma.masked_where(D_Eseries<=0, self.tseries), numpy.ma.masked_where(D_Eseries<=0, D_Eseries), 0, color=color_D_E, alpha=0.5, label='$D_E$', zorder=4) ax.plot( numpy.ma.masked_where(D_Eseries<=0, self.tseries), numpy.ma.masked_where(D_Eseries<=0, D_Eseries), color=color_D_E, zorder=5) - if(any(D_Iseries) and plot_D_I=='shaded'): + if (any(D_Iseries) and plot_D_I=='shaded'): ax.fill_between(numpy.ma.masked_where(D_Iseries<=0, self.tseries), numpy.ma.masked_where(D_Iseries<=0, D_Iseries), 0, color=color_D_I, alpha=0.5, label='$D_I$', zorder=4) ax.plot( numpy.ma.masked_where(D_Iseries<=0, self.tseries), numpy.ma.masked_where(D_Iseries<=0, D_Iseries), color=color_D_I, zorder=5) - if(any(Iseries) and plot_I=='shaded'): + if (any(Iseries) and plot_I=='shaded'): ax.fill_between(numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, Iseries), 0, color=color_I, alpha=0.5, label='$I$', zorder=4) ax.plot( numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, Iseries), color=color_I, zorder=5) - if(any(Sseries) and plot_S=='shaded'): + if (any(Sseries) and plot_S=='shaded'): ax.fill_between(numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, Sseries), 0, color=color_S, alpha=0.5, label='$S$', zorder=4) ax.plot( numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, Sseries), color=color_S, zorder=5) - if(any(Rseries) and plot_R=='shaded'): + if (any(Rseries) and plot_R=='shaded'): ax.fill_between(numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, Rseries), 0, color=color_R, alpha=0.5, label='$R$', zorder=4) ax.plot( numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, Rseries), color=color_R, zorder=5) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Draw the line variables: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(any(Fseries) and plot_F=='line'): + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + if (any(Fseries) and plot_F=='line'): ax.plot(numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, Fseries), color=color_F, label='$F$', zorder=6) - if(any(Eseries) and plot_E=='line'): + if (any(Eseries) and plot_E=='line'): ax.plot(numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, Eseries), color=color_E, label='$E$', zorder=6) - if(combine_D and (any(Dseries) and plot_D_E=='line' and plot_D_I=='line')): + if (combine_D and (any(Dseries) and plot_D_E=='line' and plot_D_I=='line')): ax.plot(numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, Dseries), color=color_D_E, label='$D_{all}$', zorder=6) else: - if(any(D_Eseries) and plot_D_E=='line'): + if (any(D_Eseries) and plot_D_E=='line'): ax.plot(numpy.ma.masked_where(D_Eseries<=0, self.tseries), numpy.ma.masked_where(D_Eseries<=0, D_Eseries), color=color_D_E, label='$D_E$', zorder=6) - if(any(D_Iseries) and plot_D_I=='line'): + if (any(D_Iseries) and plot_D_I=='line'): ax.plot(numpy.ma.masked_where(D_Iseries<=0, self.tseries), numpy.ma.masked_where(D_Iseries<=0, D_Iseries), color=color_D_I, label='$D_I$', zorder=6) - if(any(Iseries) and plot_I=='line'): + if (any(Iseries) and plot_I=='line'): ax.plot(numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, Iseries), color=color_I, label='$I$', zorder=6) - if(any(Sseries) and plot_S=='line'): + if (any(Sseries) and plot_S=='line'): ax.plot(numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, Sseries), color=color_S, label='$S$', zorder=6) - if(any(Rseries) and plot_R=='line'): + if (any(Rseries) and plot_R=='line'): ax.plot(numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, Rseries), color=color_R, label='$R$', zorder=6) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Draw the vertical line annotations: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(len(vlines)>0 and len(vline_colors)==0): + if (len(vlines)>0 and len(vline_colors)==0): vline_colors = ['gray']*len(vlines) - if(len(vlines)>0 and len(vline_labels)==0): + if (len(vlines)>0 and len(vline_labels)==0): vline_labels = [None]*len(vlines) - if(len(vlines)>0 and len(vline_styles)==0): + if (len(vlines)>0 and len(vline_styles)==0): vline_styles = [':']*len(vlines) for vline_x, vline_color, vline_style, vline_label in zip(vlines, vline_colors, vline_styles, vline_labels): - if(vline_x is not None): + if vline_x is not None: ax.axvline(x=vline_x, color=vline_color, linestyle=vline_style, alpha=1, label=vline_label) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -1241,17 +1241,17 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin ax.set_ylabel('percent of population' if plot_percentages else 'number of individuals') ax.set_xlim(0, (max(self.tseries) if not xlim else xlim)) ax.set_ylim(0, ylim) - if(plot_percentages): + if plot_percentages: ax.set_yticklabels(['{:,.0%}'.format(y) for y in ax.get_yticks()]) - if(legend): + if legend: legend_handles, legend_labels = ax.get_legend_handles_labels() ax.legend(legend_handles[::-1], legend_labels[::-1], loc='upper right', facecolor='white', edgecolor='none', framealpha=0.9, prop={'size': 8}) - if(title): + if title: ax.set_title(title, size=12) - if(side_title): + if side_title: ax.annotate(side_title, (0, 0.5), xytext=(-45, 0), ha='right', va='center', size=12, rotation=90, xycoords='axes fraction', textcoords='offset points') - + return ax @@ -1262,8 +1262,8 @@ def figure_basic(self, plot_S='line', plot_E='line', plot_I='line',plot_R='line' plot_D_E='line', plot_D_I='line', combine_D=True, color_S='tab:green', color_E='orange', color_I='crimson', color_R='tab:blue', color_F='black', color_D_E='mediumorchid', color_D_I='mediumorchid', color_reference='#E0E0E0', - dashed_reference_results=None, dashed_reference_label='reference', - shaded_reference_results=None, shaded_reference_label='reference', + dashed_reference_results=None, dashed_reference_label='reference', + shaded_reference_results=None, shaded_reference_label='reference', vlines=[], vline_colors=[], vline_styles=[], vline_labels=[], ylim=None, xlim=None, legend=True, title=None, side_title=None, plot_percentages=True, figsize=(12,8), use_seaborn=True, show=True): @@ -1272,7 +1272,7 @@ def figure_basic(self, plot_S='line', plot_E='line', plot_I='line',plot_R='line' fig, ax = pyplot.subplots(figsize=figsize) - if(use_seaborn): + if use_seaborn: import seaborn seaborn.set_style('ticks') seaborn.despine() @@ -1281,12 +1281,12 @@ def figure_basic(self, plot_S='line', plot_E='line', plot_I='line',plot_R='line' plot_D_E=plot_D_E, plot_D_I=plot_D_I, combine_D=combine_D, color_S=color_S, color_E=color_E, color_I=color_I, color_R=color_R, color_F=color_F, color_D_E=color_D_E, color_D_I=color_D_I, color_reference=color_reference, - dashed_reference_results=dashed_reference_results, dashed_reference_label=dashed_reference_label, - shaded_reference_results=shaded_reference_results, shaded_reference_label=shaded_reference_label, + dashed_reference_results=dashed_reference_results, dashed_reference_label=dashed_reference_label, + shaded_reference_results=shaded_reference_results, shaded_reference_label=shaded_reference_label, vlines=vlines, vline_colors=vline_colors, vline_styles=vline_styles, vline_labels=vline_labels, ylim=ylim, xlim=xlim, legend=legend, title=title, side_title=side_title, plot_percentages=plot_percentages) - if(show): + if show: pyplot.show() return fig, ax @@ -1299,8 +1299,8 @@ def figure_infections(self, plot_S=False, plot_E='stacked', plot_I='stacked',plo plot_D_E='stacked', plot_D_I='stacked', combine_D=True, color_S='tab:green', color_E='orange', color_I='crimson', color_R='tab:blue', color_F='black', color_D_E='mediumorchid', color_D_I='mediumorchid', color_reference='#E0E0E0', - dashed_reference_results=None, dashed_reference_label='reference', - shaded_reference_results=None, shaded_reference_label='reference', + dashed_reference_results=None, dashed_reference_label='reference', + shaded_reference_results=None, shaded_reference_label='reference', vlines=[], vline_colors=[], vline_styles=[], vline_labels=[], ylim=None, xlim=None, legend=True, title=None, side_title=None, plot_percentages=True, figsize=(12,8), use_seaborn=True, show=True): @@ -1309,7 +1309,7 @@ def figure_infections(self, plot_S=False, plot_E='stacked', plot_I='stacked',plo fig, ax = pyplot.subplots(figsize=figsize) - if(use_seaborn): + if use_seaborn: import seaborn seaborn.set_style('ticks') seaborn.despine() @@ -1318,12 +1318,12 @@ def figure_infections(self, plot_S=False, plot_E='stacked', plot_I='stacked',plo plot_D_E=plot_D_E, plot_D_I=plot_D_I, combine_D=combine_D, color_S=color_S, color_E=color_E, color_I=color_I, color_R=color_R, color_F=color_F, color_D_E=color_D_E, color_D_I=color_D_I, color_reference=color_reference, - dashed_reference_results=dashed_reference_results, dashed_reference_label=dashed_reference_label, - shaded_reference_results=shaded_reference_results, shaded_reference_label=shaded_reference_label, - vlines=vlines, vline_colors=vline_colors, vline_styles=vline_styles, vline_labels=vline_labels, + dashed_reference_results=dashed_reference_results, dashed_reference_label=dashed_reference_label, + shaded_reference_results=shaded_reference_results, shaded_reference_label=shaded_reference_label, + vlines=vlines, vline_colors=vline_colors, vline_styles=vline_styles, vline_labels=vline_labels, ylim=ylim, xlim=xlim, legend=legend, title=title, side_title=side_title, plot_percentages=plot_percentages) - if(show): + if show: pyplot.show() return fig, ax @@ -1332,4 +1332,4 @@ def figure_infections(self, plot_S=False, plot_E='stacked', plot_I='stacked',plo #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% -#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% \ No newline at end of file +#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% diff --git a/seirsplus/models.py b/seirsplus/models.py deleted file mode 100644 index 2b03171..0000000 --- a/seirsplus/models.py +++ /dev/null @@ -1,3227 +0,0 @@ -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import networkx as networkx -import numpy as numpy -import scipy as scipy -import scipy.integrate - - -######################################################## -#@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@# -#@ @# -#@ BASIC SEIRS MODELS @# -#@ @# -#@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@# -######################################################## - -class SEIRSModel(): - """ - A class to simulate the Deterministic SEIRS Model - =================================================== - Params: beta Rate of transmission (exposure) - sigma Rate of infection (upon exposure) - gamma Rate of recovery (upon infection) - xi Rate of re-susceptibility (upon recovery) - mu_I Rate of infection-related death - mu_0 Rate of baseline death - nu Rate of baseline birth - - beta_Q Rate of transmission (exposure) for individuals with detected infections - sigma_Q Rate of infection (upon exposure) for individuals with detected infections - gamma_Q Rate of recovery (upon infection) for individuals with detected infections - mu_Q Rate of infection-related death for individuals with detected infections - theta_E Rate of baseline testing for exposed individuals - theta_I Rate of baseline testing for infectious individuals - psi_E Probability of positive test results for exposed individuals - psi_I Probability of positive test results for exposed individuals - q Probability of quarantined individuals interacting with others - - initE Init number of exposed individuals - initI Init number of infectious individuals - initQ_E Init number of detected infectious individuals - initQ_I Init number of detected infectious individuals - initR Init number of recovered individuals - initF Init number of infection-related fatalities - (all remaining nodes initialized susceptible) - """ - - def __init__(self, initN, beta, sigma, gamma, xi=0, mu_I=0, mu_0=0, nu=0, p=0, - beta_Q=None, sigma_Q=None, gamma_Q=None, mu_Q=None, - theta_E=0, theta_I=0, psi_E=0, psi_I=0, q=0, - initE=0, initI=10, initQ_E=0, initQ_I=0, initR=0, initF=0): - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Model Parameters: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - self.beta = beta - self.sigma = sigma - self.gamma = gamma - self.xi = xi - self.mu_I = mu_I - self.mu_0 = mu_0 - self.nu = nu - self.p = p - - # Testing-related parameters: - self.beta_Q = beta_Q if beta_Q is not None else self.beta - self.sigma_Q = sigma_Q if sigma_Q is not None else self.sigma - self.gamma_Q = gamma_Q if gamma_Q is not None else self.gamma - self.mu_Q = mu_Q if mu_Q is not None else self.mu_I - self.theta_E = theta_E if theta_E is not None else self.theta_E - self.theta_I = theta_I if theta_I is not None else self.theta_I - self.psi_E = psi_E if psi_E is not None else self.psi_E - self.psi_I = psi_I if psi_I is not None else self.psi_I - self.q = q if q is not None else self.q - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Initialize Timekeeping: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - self.t = 0 - self.tmax = 0 # will be set when run() is called - self.tseries = numpy.array([0]) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Initialize Counts of inidividuals with each state: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - self.N = numpy.array([int(initN)]) - self.numE = numpy.array([int(initE)]) - self.numI = numpy.array([int(initI)]) - self.numQ_E = numpy.array([int(initQ_E)]) - self.numQ_I = numpy.array([int(initQ_I)]) - self.numR = numpy.array([int(initR)]) - self.numF = numpy.array([int(initF)]) - self.numS = numpy.array([self.N[-1] - self.numE[-1] - self.numI[-1] - self.numQ_E[-1] - self.numQ_I[-1] - self.numR[-1] - self.numF[-1]]) - assert(self.numS[0] >= 0), "The specified initial population size N must be greater than or equal to the initial compartment counts." - - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - @staticmethod - def system_dfes(t, variables, beta, sigma, gamma, xi, mu_I, mu_0, nu, - beta_Q, sigma_Q, gamma_Q, mu_Q, theta_E, theta_I, psi_E, psi_I, q): - - S, E, I, Q_E, Q_I, R, F = variables # varibles is a list with compartment counts as elements - - N = S + E + I + Q_E + Q_I + R - - dS = - (beta*S*I)/N - q*(beta_Q*S*Q_I)/N + xi*R + nu*N - mu_0*S - - dE = (beta*S*I)/N + q*(beta_Q*S*Q_I)/N - sigma*E - theta_E*psi_E*E - mu_0*E - - dI = sigma*E - gamma*I - mu_I*I - theta_I*psi_I*I - mu_0*I - - dDE = theta_E*psi_E*E - sigma_Q*Q_E - mu_0*Q_E - - dDI = theta_I*psi_I*I + sigma_Q*Q_E - gamma_Q*Q_I - mu_Q*Q_I - mu_0*Q_I - - dR = gamma*I + gamma_Q*Q_I - xi*R - mu_0*R - - dF = mu_I*I + mu_Q*Q_I - - return [dS, dE, dI, dDE, dDI, dR, dF] - - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def run_epoch(self, runtime, dt=0.1): - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Create a list of times at which the ODE solver should output system values. - # Append this list of times as the model's timeseries - t_eval = numpy.arange(start=self.t, stop=self.t+runtime, step=dt) - - # Define the range of time values for the integration: - t_span = [self.t, self.t+runtime] - - # Define the initial conditions as the system's current state: - # (which will be the t=0 condition if this is the first run of this model, - # else where the last sim left off) - - init_cond = [self.numS[-1], self.numE[-1], self.numI[-1], self.numQ_E[-1], self.numQ_I[-1], self.numR[-1], self.numF[-1]] - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Solve the system of differential eqns: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - solution = scipy.integrate.solve_ivp(lambda t, X: SEIRSModel.system_dfes(t, X, self.beta, self.sigma, self.gamma, self.xi, self.mu_I, self.mu_0, self.nu, - self.beta_Q, self.sigma_Q, self.gamma_Q, self.mu_Q, self.theta_E, self.theta_I, self.psi_E, self.psi_I, self.q - ), - t_span=t_span, y0=init_cond, t_eval=t_eval - ) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Store the solution output as the model's time series and data series: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - self.tseries = numpy.append(self.tseries, solution['t']) - self.numS = numpy.append(self.numS, solution['y'][0]) - self.numE = numpy.append(self.numE, solution['y'][1]) - self.numI = numpy.append(self.numI, solution['y'][2]) - self.numQ_E = numpy.append(self.numQ_E, solution['y'][3]) - self.numQ_I = numpy.append(self.numQ_I, solution['y'][4]) - self.numR = numpy.append(self.numR, solution['y'][5]) - self.numF = numpy.append(self.numF, solution['y'][6]) - - self.t = self.tseries[-1] - - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def run(self, T, dt=0.1, checkpoints=None, verbose=False): - - if(T>0): - self.tmax += T - else: - return False - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Pre-process checkpoint values: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(checkpoints): - numCheckpoints = len(checkpoints['t']) - paramNames = ['beta', 'sigma', 'gamma', 'xi', 'mu_I', 'mu_0', 'nu', - 'beta_Q', 'sigma_Q', 'gamma_Q', 'mu_Q', - 'theta_E', 'theta_I', 'psi_E', 'psi_I', 'q'] - for param in paramNames: - # For params that don't have given checkpoint values (or bad value given), - # set their checkpoint values to the value they have now for all checkpoints. - if(param not in list(checkpoints.keys()) - or not isinstance(checkpoints[param], (list, numpy.ndarray)) - or len(checkpoints[param])!=numCheckpoints): - checkpoints[param] = [getattr(self, param)]*numCheckpoints - - #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% - # Run the simulation loop: - #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% - if(not checkpoints): - self.run_epoch(runtime=self.tmax, dt=dt) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - print("t = %.2f" % self.t) - if(verbose): - print("\t S = " + str(self.numS[-1])) - print("\t E = " + str(self.numE[-1])) - print("\t I = " + str(self.numI[-1])) - print("\t Q_E = " + str(self.numQ_E[-1])) - print("\t Q_I = " + str(self.numQ_I[-1])) - print("\t R = " + str(self.numR[-1])) - print("\t F = " + str(self.numF[-1])) - - - else: # checkpoints provided - for checkpointIdx, checkpointTime in enumerate(checkpoints['t']): - # Run the sim until the next checkpoint time: - self.run_epoch(runtime=checkpointTime-self.t, dt=dt) - # Having reached the checkpoint, update applicable parameters: - print("[Checkpoint: Updating parameters]") - for param in paramNames: - setattr(self, param, checkpoints[param][checkpointIdx]) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - print("t = %.2f" % self.t) - if(verbose): - print("\t S = " + str(self.numS[-1])) - print("\t E = " + str(self.numE[-1])) - print("\t I = " + str(self.numI[-1])) - print("\t Q_E = " + str(self.numQ_E[-1])) - print("\t Q_I = " + str(self.numQ_I[-1])) - print("\t R = " + str(self.numR[-1])) - print("\t F = " + str(self.numF[-1])) - - if(self.t < self.tmax): - self.run_epoch(runtime=self.tmax-self.t, dt=dt) - - return True - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def total_num_susceptible(self, t_idx=None): - if(t_idx is None): - return (self.numS[:]) - else: - return (self.numS[t_idx]) - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def total_num_infected(self, t_idx=None): - if(t_idx is None): - return (self.numE[:] + self.numI[:] + self.numQ_E[:] + self.numQ_I[:]) - else: - return (self.numE[t_idx] + self.numI[t_idx] + self.numQ_E[t_idx] + self.numQ_I[t_idx]) - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def total_num_isolated(self, t_idx=None): - if(t_idx is None): - return (self.numQ_E[:] + self.numQ_I[:]) - else: - return (self.numQ_E[t_idx] + self.numQ_I[t_idx]) - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def total_num_recovered(self, t_idx=None): - if(t_idx is None): - return (self.numR[:]) - else: - return (self.numR[t_idx]) - - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='line', plot_F='line', - plot_Q_E='line', plot_Q_I='line', combine_Q=True, - color_S='tab:green', color_E='orange', color_I='crimson', color_R='tab:blue', color_F='black', - color_Q_E='mediumorchid', color_Q_I='mediumorchid', color_reference='#E0E0E0', - dashed_reference_results=None, dashed_reference_label='reference', - shaded_reference_results=None, shaded_reference_label='reference', - vlines=[], vline_colors=[], vline_styles=[], vline_labels=[], - ylim=None, xlim=None, legend=True, title=None, side_title=None, plot_percentages=True): - - import matplotlib.pyplot as pyplot - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Create an Axes object if None provided: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(not ax): - fig, ax = pyplot.subplots() - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Prepare data series to be plotted: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - Fseries = self.numF/self.N if plot_percentages else self.numF - Eseries = self.numE/self.N if plot_percentages else self.numE - Dseries = (self.numQ_E+self.numQ_I)/self.N if plot_percentages else (self.numQ_E+self.numQ_I) - Q_Eseries = self.numQ_E/self.N if plot_percentages else self.numQ_E - Q_Iseries = self.numQ_I/self.N if plot_percentages else self.numQ_I - Iseries = self.numI/self.N if plot_percentages else self.numI - Rseries = self.numR/self.N if plot_percentages else self.numR - Sseries = self.numS/self.N if plot_percentages else self.numS - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Draw the reference data: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(dashed_reference_results): - dashedReference_tseries = dashed_reference_results.tseries[::int(self.N/100)] - dashedReference_IDEstack = (dashed_reference_results.numI + dashed_reference_results.numQ_I + dashed_reference_results.numQ_E + dashed_reference_results.numE)[::int(self.N/100)] / (self.N if plot_percentages else 1) - ax.plot(dashedReference_tseries, dashedReference_IDEstack, color='#E0E0E0', linestyle='--', label='$I+D+E$ ('+dashed_reference_label+')', zorder=0) - if(shaded_reference_results): - shadedReference_tseries = shaded_reference_results.tseries - shadedReference_IDEstack = (shaded_reference_results.numI + shaded_reference_results.numQ_I + shaded_reference_results.numQ_E + shaded_reference_results.numE) / (self.N if plot_percentages else 1) - ax.fill_between(shaded_reference_results.tseries, shadedReference_IDEstack, 0, color='#EFEFEF', label='$I+D+E$ ('+shaded_reference_label+')', zorder=0) - ax.plot(shaded_reference_results.tseries, shadedReference_IDEstack, color='#E0E0E0', zorder=1) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Draw the stacked variables: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - topstack = numpy.zeros_like(self.tseries) - if(any(Fseries) and plot_F=='stacked'): - ax.fill_between(numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, topstack+Fseries), topstack, color=color_F, alpha=0.5, label='$F$', zorder=2) - ax.plot( numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, topstack+Fseries), color=color_F, zorder=3) - topstack = topstack+Fseries - if(any(Eseries) and plot_E=='stacked'): - ax.fill_between(numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, topstack+Eseries), topstack, color=color_E, alpha=0.5, label='$E$', zorder=2) - ax.plot( numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, topstack+Eseries), color=color_E, zorder=3) - topstack = topstack+Eseries - if(combine_Q and plot_Q_E=='stacked' and plot_Q_I=='stacked'): - ax.fill_between(numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, topstack+Dseries), topstack, color=color_Q_E, alpha=0.5, label='$Q_{all}$', zorder=2) - ax.plot( numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, topstack+Dseries), color=color_Q_E, zorder=3) - topstack = topstack+Dseries - else: - if(any(Q_Eseries) and plot_Q_E=='stacked'): - ax.fill_between(numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, topstack+Q_Eseries), topstack, color=color_Q_E, alpha=0.5, label='$Q_E$', zorder=2) - ax.plot( numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, topstack+Q_Eseries), color=color_Q_E, zorder=3) - topstack = topstack+Q_Eseries - if(any(Q_Iseries) and plot_Q_I=='stacked'): - ax.fill_between(numpy.ma.masked_where(Q_Iseries<=0, self.tseries), numpy.ma.masked_where(Q_Iseries<=0, topstack+Q_Iseries), topstack, color=color_Q_I, alpha=0.5, label='$Q_I$', zorder=2) - ax.plot( numpy.ma.masked_where(Q_Iseries<=0, self.tseries), numpy.ma.masked_where(Q_Iseries<=0, topstack+Q_Iseries), color=color_Q_I, zorder=3) - topstack = topstack+Q_Iseries - if(any(Iseries) and plot_I=='stacked'): - ax.fill_between(numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, topstack+Iseries), topstack, color=color_I, alpha=0.5, label='$I$', zorder=2) - ax.plot( numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, topstack+Iseries), color=color_I, zorder=3) - topstack = topstack+Iseries - if(any(Rseries) and plot_R=='stacked'): - ax.fill_between(numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, topstack+Rseries), topstack, color=color_R, alpha=0.5, label='$R$', zorder=2) - ax.plot( numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, topstack+Rseries), color=color_R, zorder=3) - topstack = topstack+Rseries - if(any(Sseries) and plot_S=='stacked'): - ax.fill_between(numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, topstack+Sseries), topstack, color=color_S, alpha=0.5, label='$S$', zorder=2) - ax.plot( numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, topstack+Sseries), color=color_S, zorder=3) - topstack = topstack+Sseries - - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Draw the shaded variables: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(any(Fseries) and plot_F=='shaded'): - ax.fill_between(numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, Fseries), 0, color=color_F, alpha=0.5, label='$F$', zorder=4) - ax.plot( numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, Fseries), color=color_F, zorder=5) - if(any(Eseries) and plot_E=='shaded'): - ax.fill_between(numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, Eseries), 0, color=color_E, alpha=0.5, label='$E$', zorder=4) - ax.plot( numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, Eseries), color=color_E, zorder=5) - if(combine_Q and (any(Dseries) and plot_Q_E=='shaded' and plot_Q_E=='shaded')): - ax.fill_between(numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, Dseries), 0, color=color_Q_E, alpha=0.5, label='$Q_{all}$', zorder=4) - ax.plot( numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, Dseries), color=color_Q_E, zorder=5) - else: - if(any(Q_Eseries) and plot_Q_E=='shaded'): - ax.fill_between(numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, Q_Eseries), 0, color=color_Q_E, alpha=0.5, label='$Q_E$', zorder=4) - ax.plot( numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, Q_Eseries), color=color_Q_E, zorder=5) - if(any(Q_Iseries) and plot_Q_I=='shaded'): - ax.fill_between(numpy.ma.masked_where(Q_Iseries<=0, self.tseries), numpy.ma.masked_where(Q_Iseries<=0, Q_Iseries), 0, color=color_Q_I, alpha=0.5, label='$Q_I$', zorder=4) - ax.plot( numpy.ma.masked_where(Q_Iseries<=0, self.tseries), numpy.ma.masked_where(Q_Iseries<=0, Q_Iseries), color=color_Q_I, zorder=5) - if(any(Iseries) and plot_I=='shaded'): - ax.fill_between(numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, Iseries), 0, color=color_I, alpha=0.5, label='$I$', zorder=4) - ax.plot( numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, Iseries), color=color_I, zorder=5) - if(any(Sseries) and plot_S=='shaded'): - ax.fill_between(numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, Sseries), 0, color=color_S, alpha=0.5, label='$S$', zorder=4) - ax.plot( numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, Sseries), color=color_S, zorder=5) - if(any(Rseries) and plot_R=='shaded'): - ax.fill_between(numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, Rseries), 0, color=color_R, alpha=0.5, label='$R$', zorder=4) - ax.plot( numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, Rseries), color=color_R, zorder=5) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Draw the line variables: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(any(Fseries) and plot_F=='line'): - ax.plot(numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, Fseries), color=color_F, label='$F$', zorder=6) - if(any(Eseries) and plot_E=='line'): - ax.plot(numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, Eseries), color=color_E, label='$E$', zorder=6) - if(combine_Q and (any(Dseries) and plot_Q_E=='line' and plot_Q_E=='line')): - ax.plot(numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, Dseries), color=color_Q_E, label='$Q_{all}$', zorder=6) - else: - if(any(Q_Eseries) and plot_Q_E=='line'): - ax.plot(numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, Q_Eseries), color=color_Q_E, label='$Q_E$', zorder=6) - if(any(Q_Iseries) and plot_Q_I=='line'): - ax.plot(numpy.ma.masked_where(Q_Iseries<=0, self.tseries), numpy.ma.masked_where(Q_Iseries<=0, Q_Iseries), color=color_Q_I, label='$Q_I$', zorder=6) - if(any(Iseries) and plot_I=='line'): - ax.plot(numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, Iseries), color=color_I, label='$I$', zorder=6) - if(any(Sseries) and plot_S=='line'): - ax.plot(numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, Sseries), color=color_S, label='$S$', zorder=6) - if(any(Rseries) and plot_R=='line'): - ax.plot(numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, Rseries), color=color_R, label='$R$', zorder=6) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Draw the vertical line annotations: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(len(vlines)>0 and len(vline_colors)==0): - vline_colors = ['gray']*len(vlines) - if(len(vlines)>0 and len(vline_labels)==0): - vline_labels = [None]*len(vlines) - if(len(vlines)>0 and len(vline_styles)==0): - vline_styles = [':']*len(vlines) - for vline_x, vline_color, vline_style, vline_label in zip(vlines, vline_colors, vline_styles, vline_labels): - if(vline_x is not None): - ax.axvline(x=vline_x, color=vline_color, linestyle=vline_style, alpha=1, label=vline_label) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Draw the plot labels: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - ax.set_xlabel('days') - ax.set_ylabel('percent of population' if plot_percentages else 'number of individuals') - ax.set_xlim(0, (max(self.tseries) if not xlim else xlim)) - ax.set_ylim(0, ylim) - if(plot_percentages): - ax.set_yticklabels(['{:,.0%}'.format(y) for y in ax.get_yticks()]) - if(legend): - legend_handles, legend_labels = ax.get_legend_handles_labels() - ax.legend(legend_handles[::-1], legend_labels[::-1], loc='upper right', facecolor='white', edgecolor='none', framealpha=0.9, prop={'size': 8}) - if(title): - ax.set_title(title, size=12) - if(side_title): - ax.annotate(side_title, (0, 0.5), xytext=(-45, 0), ha='right', va='center', - size=12, rotation=90, xycoords='axes fraction', textcoords='offset points') - - return ax - - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def figure_basic(self, plot_S='line', plot_E='line', plot_I='line',plot_R='line', plot_F='line', - plot_Q_E='line', plot_Q_I='line', combine_Q=True, - color_S='tab:green', color_E='orange', color_I='crimson', color_R='tab:blue', color_F='black', - color_Q_E='mediumorchid', color_Q_I='mediumorchid', color_reference='#E0E0E0', - dashed_reference_results=None, dashed_reference_label='reference', - shaded_reference_results=None, shaded_reference_label='reference', - vlines=[], vline_colors=[], vline_styles=[], vline_labels=[], - ylim=None, xlim=None, legend=True, title=None, side_title=None, plot_percentages=True, - figsize=(12,8), use_seaborn=True, show=True): - - import matplotlib.pyplot as pyplot - - fig, ax = pyplot.subplots(figsize=figsize) - - if(use_seaborn): - import seaborn - seaborn.set_style('ticks') - seaborn.despine() - - self.plot(ax=ax, plot_S=plot_S, plot_E=plot_E, plot_I=plot_I,plot_R=plot_R, plot_F=plot_F, - plot_Q_E=plot_Q_E, plot_Q_I=plot_Q_I, combine_Q=combine_Q, - color_S=color_S, color_E=color_E, color_I=color_I, color_R=color_R, color_F=color_F, - color_Q_E=color_Q_E, color_Q_I=color_Q_I, color_reference=color_reference, - dashed_reference_results=dashed_reference_results, dashed_reference_label=dashed_reference_label, - shaded_reference_results=shaded_reference_results, shaded_reference_label=shaded_reference_label, - vlines=vlines, vline_colors=vline_colors, vline_styles=vline_styles, vline_labels=vline_labels, - ylim=ylim, xlim=xlim, legend=legend, title=title, side_title=side_title, plot_percentages=plot_percentages) - - if(show): - pyplot.show() - - return fig, ax - - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def figure_infections(self, plot_S=False, plot_E='stacked', plot_I='stacked',plot_R=False, plot_F=False, - plot_Q_E='stacked', plot_Q_I='stacked', combine_Q=True, - color_S='tab:green', color_E='orange', color_I='crimson', color_R='tab:blue', color_F='black', - color_Q_E='mediumorchid', color_Q_I='mediumorchid', color_reference='#E0E0E0', - dashed_reference_results=None, dashed_reference_label='reference', - shaded_reference_results=None, shaded_reference_label='reference', - vlines=[], vline_colors=[], vline_styles=[], vline_labels=[], - ylim=None, xlim=None, legend=True, title=None, side_title=None, plot_percentages=True, - figsize=(12,8), use_seaborn=True, show=True): - - import matplotlib.pyplot as pyplot - - fig, ax = pyplot.subplots(figsize=figsize) - - if(use_seaborn): - import seaborn - seaborn.set_style('ticks') - seaborn.despine() - - self.plot(ax=ax, plot_S=plot_S, plot_E=plot_E, plot_I=plot_I,plot_R=plot_R, plot_F=plot_F, - plot_Q_E=plot_Q_E, plot_Q_I=plot_Q_I, combine_Q=combine_Q, - color_S=color_S, color_E=color_E, color_I=color_I, color_R=color_R, color_F=color_F, - color_Q_E=color_Q_E, color_Q_I=color_Q_I, color_reference=color_reference, - dashed_reference_results=dashed_reference_results, dashed_reference_label=dashed_reference_label, - shaded_reference_results=shaded_reference_results, shaded_reference_label=shaded_reference_label, - vlines=vlines, vline_colors=vline_colors, vline_styles=vline_styles, vline_labels=vline_labels, - ylim=ylim, xlim=xlim, legend=legend, title=title, side_title=side_title, plot_percentages=plot_percentages) - - if(show): - pyplot.show() - - return fig, ax - - -#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% -#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% -#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% -#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% - - - -class SEIRSNetworkModel(): - """ - A class to simulate the SEIRS Stochastic Network Model - ====================================================== - Params: - G Network adjacency matrix (numpy array) or Networkx graph object. - beta Rate of transmission (global interactions) - beta_local Rate(s) of transmission between adjacent individuals (optional) - sigma Rate of progression to infectious state (inverse of latent period) - gamma Rate of recovery (inverse of symptomatic infectious period) - mu_I Rate of infection-related death - xi Rate of re-susceptibility (upon recovery) - mu_0 Rate of baseline death - nu Rate of baseline birth - p Probability of individuals interacting with global population - - G_Q Quarantine adjacency matrix (numpy array) or Networkx graph object. - beta_Q Rate of transmission for isolated individuals (global interactions) - beta_Q_local Rate(s) of transmission (exposure) for adjacent isolated individuals (optional) - sigma_Q Rate of progression to infectious state for isolated individuals - gamma_Q Rate of recovery for isolated individuals - mu_Q Rate of infection-related death for isolated individuals - q Probability of isolated individuals interacting with global population - isolation_time Time to remain in isolation upon positive test, self-isolation, etc. - - theta_E Rate of random testing for exposed individuals - theta_I Rate of random testing for infectious individuals - phi_E Rate of testing when a close contact has tested positive for exposed individuals - phi_I Rate of testing when a close contact has tested positive for infectious individuals - psi_E Probability of positive test for exposed individuals - psi_I Probability of positive test for infectious individuals - - initE Initial number of exposed individuals - initI Initial number of infectious individuals - initR Initial number of recovered individuals - initF Initial number of infection-related fatalities - initQ_S Initial number of isolated susceptible individuals - initQ_E Initial number of isolated exposed individuals - initQ_I Initial number of isolated infectious individuals - initQ_R Initial number of isolated recovered individuals - (all remaining nodes initialized susceptible) - """ - def __init__(self, G, beta, sigma, gamma, - mu_I=0, alpha=1.0, xi=0, mu_0=0, nu=0, f=0, p=0, - beta_local=None, beta_pairwise_mode='infected', delta=None, delta_pairwise_mode=None, - G_Q=None, beta_Q=None, beta_Q_local=None, sigma_Q=None, gamma_Q=None, mu_Q=None, alpha_Q=None, delta_Q=None, - theta_E=0, theta_I=0, phi_E=0, phi_I=0, psi_E=1, psi_I=1, q=0, isolation_time=14, - initE=0, initI=0, initR=0, initF=0, initQ_E=0, initQ_I=0, - transition_mode='exponential_rates', node_groups=None, store_Xseries=False, seed=None): - - if(seed is not None): - numpy.random.seed(seed) - self.seed = seed - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Model Parameters: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - self.parameters = { 'G':G, 'G_Q':G_Q, - 'beta':beta, 'sigma':sigma, 'gamma':gamma, 'mu_I':mu_I, - 'xi':xi, 'mu_0':mu_0, 'nu':nu, 'f':f, 'p':p, - 'beta_local':beta_local, 'beta_pairwise_mode':beta_pairwise_mode, - 'alpha':alpha, 'delta':delta, 'delta_pairwise_mode':delta_pairwise_mode, - 'beta_Q':beta_Q, 'beta_Q_local':beta_Q_local, 'sigma_Q':sigma_Q, 'gamma_Q':gamma_Q, 'mu_Q':mu_Q, - 'alpha_Q':alpha_Q, 'delta_Q':delta_Q, - 'theta_E':theta_E, 'theta_I':theta_I, 'phi_E':phi_E, 'phi_I':phi_I, 'psi_E':psi_E, 'psi_I':psi_I, - 'q':q, 'isolation_time':isolation_time, - 'initE':initE, 'initI':initI, 'initR':initR, 'initF':initF, - 'initQ_E':initQ_E, 'initQ_I':initQ_I } - self.update_parameters() - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Each node can undergo 4-6 transitions (sans vitality/re-susceptibility returns to S state), - # so there are ~numNodes*6 events/timesteps expected; initialize numNodes*6 timestep slots to start - # (will be expanded during run if needed for some reason) - self.tseries = numpy.zeros(6*self.numNodes) - self.numS = numpy.zeros(6*self.numNodes) - self.numE = numpy.zeros(6*self.numNodes) - self.numI = numpy.zeros(6*self.numNodes) - self.numR = numpy.zeros(6*self.numNodes) - self.numF = numpy.zeros(6*self.numNodes) - self.numQ_E = numpy.zeros(6*self.numNodes) - self.numQ_I = numpy.zeros(6*self.numNodes) - self.N = numpy.zeros(6*self.numNodes) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Initialize Timekeeping: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - self.t = 0 - self.tmax = 0 # will be set when run() is called - self.tidx = 0 - self.tseries[0] = 0 - - # Vectors holding the time that each node has been in a given state or in isolation: - self.timer_state = numpy.zeros((self.numNodes,1)) - self.timer_isolation = numpy.zeros(self.numNodes) - self.isolationTime = isolation_time - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Initialize Counts of inidividuals with each state: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - self.numE[0] = int(initE) - self.numI[0] = int(initI) - self.numR[0] = int(initR) - self.numF[0] = int(initF) - self.numQ_E[0] = int(initQ_E) - self.numQ_I[0] = int(initQ_I) - self.numS[0] = (self.numNodes - self.numE[0] - self.numI[0] - self.numR[0] - self.numQ_E[0] - self.numQ_I[0] - self.numF[0]) - self.N[0] = self.numNodes - self.numF[0] - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Node states: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - self.S = 1 - self.E = 2 - self.I = 3 - self.R = 4 - self.F = 5 - self.Q_E = 6 - self.Q_I = 7 - - self.X = numpy.array( [self.S]*int(self.numS[0]) + [self.E]*int(self.numE[0]) + [self.I]*int(self.numI[0]) - + [self.R]*int(self.numR[0]) + [self.F]*int(self.numF[0]) - + [self.Q_E]*int(self.numQ_E[0]) + [self.Q_I]*int(self.numQ_I[0]) - ).reshape((self.numNodes,1)) - numpy.random.shuffle(self.X) - - self.store_Xseries = store_Xseries - if(store_Xseries): - self.Xseries = numpy.zeros(shape=(6*self.numNodes, self.numNodes), dtype='uint8') - self.Xseries[0,:] = self.X.T - - self.transitions = { - 'StoE': {'currentState':self.S, 'newState':self.E}, - 'EtoI': {'currentState':self.E, 'newState':self.I}, - 'ItoR': {'currentState':self.I, 'newState':self.R}, - 'ItoF': {'currentState':self.I, 'newState':self.F}, - 'RtoS': {'currentState':self.R, 'newState':self.S}, - 'EtoQE': {'currentState':self.E, 'newState':self.Q_E}, - 'ItoQI': {'currentState':self.I, 'newState':self.Q_I}, - 'QEtoQI': {'currentState':self.Q_E, 'newState':self.Q_I}, - 'QItoR': {'currentState':self.Q_I, 'newState':self.R}, - 'QItoF': {'currentState':self.Q_I, 'newState':self.F}, - '_toS': {'currentState':True, 'newState':self.S}, - } - - self.transition_mode = transition_mode - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Initialize other node metadata: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - self.tested = numpy.array([False]*self.numNodes).reshape((self.numNodes,1)) - self.positive = numpy.array([False]*self.numNodes).reshape((self.numNodes,1)) - self.numTested = numpy.zeros(6*self.numNodes) - self.numPositive = numpy.zeros(6*self.numNodes) - - self.testedInCurrentState = numpy.array([False]*self.numNodes).reshape((self.numNodes,1)) - - self.infectionsLog = [] - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Initialize node subgroup data series: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - self.nodeGroupData = None - if(node_groups): - self.nodeGroupData = {} - for groupName, nodeList in node_groups.items(): - self.nodeGroupData[groupName] = {'nodes': numpy.array(nodeList), - 'mask': numpy.isin(range(self.numNodes), nodeList).reshape((self.numNodes,1))} - self.nodeGroupData[groupName]['numS'] = numpy.zeros(6*self.numNodes) - self.nodeGroupData[groupName]['numE'] = numpy.zeros(6*self.numNodes) - self.nodeGroupData[groupName]['numI'] = numpy.zeros(6*self.numNodes) - self.nodeGroupData[groupName]['numR'] = numpy.zeros(6*self.numNodes) - self.nodeGroupData[groupName]['numF'] = numpy.zeros(6*self.numNodes) - self.nodeGroupData[groupName]['numQ_E'] = numpy.zeros(6*self.numNodes) - self.nodeGroupData[groupName]['numQ_I'] = numpy.zeros(6*self.numNodes) - self.nodeGroupData[groupName]['N'] = numpy.zeros(6*self.numNodes) - self.nodeGroupData[groupName]['numPositive'] = numpy.zeros(6*self.numNodes) - self.nodeGroupData[groupName]['numTested'] = numpy.zeros(6*self.numNodes) - self.nodeGroupData[groupName]['numS'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.S) - self.nodeGroupData[groupName]['numE'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.E) - self.nodeGroupData[groupName]['numI'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.I) - self.nodeGroupData[groupName]['numR'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.R) - self.nodeGroupData[groupName]['numF'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.F) - self.nodeGroupData[groupName]['numQ_E'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_E) - self.nodeGroupData[groupName]['numQ_I'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_I) - self.nodeGroupData[groupName]['N'][0] = self.numNodes - self.numF[0] - - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def update_parameters(self): - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Model graphs: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - self.G = self.parameters['G'] - # Adjacency matrix: - if type(self.G)==numpy.ndarray: - self.A = scipy.sparse.csr_matrix(self.G) - elif type(self.G)==networkx.classes.graph.Graph: - self.A = networkx.adj_matrix(self.G) # adj_matrix gives scipy.sparse csr_matrix - else: - raise BaseException("Input an adjacency matrix or networkx object only.") - self.numNodes = int(self.A.shape[1]) - self.degree = numpy.asarray(self.node_degrees(self.A)).astype(float) - #---------------------------------------- - if(self.parameters['G_Q'] is None): - self.G_Q = self.G # If no Q graph is provided, use G in its place - else: - self.G_Q = self.parameters['G_Q'] - # Quarantine Adjacency matrix: - if type(self.G_Q)==numpy.ndarray: - self.A_Q = scipy.sparse.csr_matrix(self.G_Q) - elif type(self.G_Q)==networkx.classes.graph.Graph: - self.A_Q = networkx.adj_matrix(self.G_Q) # adj_matrix gives scipy.sparse csr_matrix - else: - raise BaseException("Input an adjacency matrix or networkx object only.") - self.numNodes_Q = int(self.A_Q.shape[1]) - self.degree_Q = numpy.asarray(self.node_degrees(self.A_Q)).astype(float) - #---------------------------------------- - assert(self.numNodes == self.numNodes_Q), "The normal and quarantine adjacency graphs must be of the same size." - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Model parameters: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - self.beta = numpy.array(self.parameters['beta']).reshape((self.numNodes, 1)) if isinstance(self.parameters['beta'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['beta'], shape=(self.numNodes,1)) - self.sigma = numpy.array(self.parameters['sigma']).reshape((self.numNodes, 1)) if isinstance(self.parameters['sigma'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['sigma'], shape=(self.numNodes,1)) - self.gamma = numpy.array(self.parameters['gamma']).reshape((self.numNodes, 1)) if isinstance(self.parameters['gamma'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['gamma'], shape=(self.numNodes,1)) - self.mu_I = numpy.array(self.parameters['mu_I']).reshape((self.numNodes, 1)) if isinstance(self.parameters['mu_I'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['mu_I'], shape=(self.numNodes,1)) - self.alpha = numpy.array(self.parameters['alpha']).reshape((self.numNodes, 1)) if isinstance(self.parameters['alpha'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['alpha'], shape=(self.numNodes,1)) - self.xi = numpy.array(self.parameters['xi']).reshape((self.numNodes, 1)) if isinstance(self.parameters['xi'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['xi'], shape=(self.numNodes,1)) - self.mu_0 = numpy.array(self.parameters['mu_0']).reshape((self.numNodes, 1)) if isinstance(self.parameters['mu_0'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['mu_0'], shape=(self.numNodes,1)) - self.nu = numpy.array(self.parameters['nu']).reshape((self.numNodes, 1)) if isinstance(self.parameters['nu'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['nu'], shape=(self.numNodes,1)) - self.f = numpy.array(self.parameters['f']).reshape((self.numNodes, 1)) if isinstance(self.parameters['f'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['f'], shape=(self.numNodes,1)) - self.p = numpy.array(self.parameters['p']).reshape((self.numNodes, 1)) if isinstance(self.parameters['p'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['p'], shape=(self.numNodes,1)) - - self.rand_f = numpy.random.rand(self.f.shape[0], self.f.shape[1]) - - #---------------------------------------- - # Testing-related parameters: - #---------------------------------------- - self.beta_Q = (numpy.array(self.parameters['beta_Q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['beta_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['beta_Q'], shape=(self.numNodes,1))) if self.parameters['beta_Q'] is not None else self.beta - self.sigma_Q = (numpy.array(self.parameters['sigma_Q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['sigma_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['sigma_Q'], shape=(self.numNodes,1))) if self.parameters['sigma_Q'] is not None else self.sigma - self.gamma_Q = (numpy.array(self.parameters['gamma_Q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['gamma_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['gamma_Q'], shape=(self.numNodes,1))) if self.parameters['gamma_Q'] is not None else self.gamma - self.mu_Q = (numpy.array(self.parameters['mu_Q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['mu_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['mu_Q'], shape=(self.numNodes,1))) if self.parameters['mu_Q'] is not None else self.mu_I - self.alpha_Q = (numpy.array(self.parameters['alpha_Q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['alpha_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['alpha_Q'], shape=(self.numNodes,1))) if self.parameters['alpha_Q'] is not None else self.alpha - self.theta_E = numpy.array(self.parameters['theta_E']).reshape((self.numNodes, 1)) if isinstance(self.parameters['theta_E'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['theta_E'], shape=(self.numNodes,1)) - self.theta_I = numpy.array(self.parameters['theta_I']).reshape((self.numNodes, 1)) if isinstance(self.parameters['theta_I'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['theta_I'], shape=(self.numNodes,1)) - self.phi_E = numpy.array(self.parameters['phi_E']).reshape((self.numNodes, 1)) if isinstance(self.parameters['phi_E'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['phi_E'], shape=(self.numNodes,1)) - self.phi_I = numpy.array(self.parameters['phi_I']).reshape((self.numNodes, 1)) if isinstance(self.parameters['phi_I'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['phi_I'], shape=(self.numNodes,1)) - self.psi_E = numpy.array(self.parameters['psi_E']).reshape((self.numNodes, 1)) if isinstance(self.parameters['psi_E'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['psi_E'], shape=(self.numNodes,1)) - self.psi_I = numpy.array(self.parameters['psi_I']).reshape((self.numNodes, 1)) if isinstance(self.parameters['psi_I'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['psi_I'], shape=(self.numNodes,1)) - self.q = numpy.array(self.parameters['q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['q'], shape=(self.numNodes,1)) - - #---------------------------------------- - - self.beta_pairwise_mode = self.parameters['beta_pairwise_mode'] - - #---------------------------------------- - # Global transmission parameters: - #---------------------------------------- - if(self.beta_pairwise_mode == 'infected' or self.beta_pairwise_mode is None): - self.beta_global = numpy.full_like(self.beta, fill_value=numpy.mean(self.beta)) - self.beta_Q_global = numpy.full_like(self.beta_Q, fill_value=numpy.mean(self.beta_Q)) - elif(self.beta_pairwise_mode == 'infectee'): - self.beta_global = self.beta - self.beta_Q_global = self.beta_Q - elif(self.beta_pairwise_mode == 'min'): - self.beta_global = numpy.minimum(self.beta, numpy.mean(beta)) - self.beta_Q_global = numpy.minimum(self.beta_Q, numpy.mean(beta_Q)) - elif(self.beta_pairwise_mode == 'max'): - self.beta_global = numpy.maximum(self.beta, numpy.mean(beta)) - self.beta_Q_global = numpy.maximum(self.beta_Q, numpy.mean(beta_Q)) - elif(self.beta_pairwise_mode == 'mean'): - self.beta_global = (self.beta + numpy.full_like(self.beta, fill_value=numpy.mean(self.beta)))/2 - self.beta_Q_global = (self.beta_Q + numpy.full_like(self.beta_Q, fill_value=numpy.mean(self.beta_Q)))/2 - - #---------------------------------------- - # Local transmission parameters: - #---------------------------------------- - self.beta_local = self.beta if self.parameters['beta_local'] is None else numpy.array(self.parameters['beta_local']) if isinstance(self.parameters['beta_local'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['beta_local'], shape=(self.numNodes,1)) - self.beta_Q_local = self.beta_Q if self.parameters['beta_Q_local'] is None else numpy.array(self.parameters['beta_Q_local']) if isinstance(self.parameters['beta_Q_local'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['beta_Q_local'], shape=(self.numNodes,1)) - - #---------------------------------------- - if(self.beta_local.ndim == 2 and self.beta_local.shape[0] == self.numNodes and self.beta_local.shape[1] == self.numNodes): - self.A_beta_pairwise = self.beta_local - elif((self.beta_local.ndim == 1 and self.beta_local.shape[0] == self.numNodes) or (self.beta_local.ndim == 2 and (self.beta_local.shape[0] == self.numNodes or self.beta_local.shape[1] == self.numNodes))): - self.beta_local = self.beta_local.reshape((self.numNodes,1)) - # Pre-multiply beta values by the adjacency matrix ("transmission weight connections") - A_beta_pairwise_byInfected = scipy.sparse.csr_matrix.multiply(self.A, self.beta_local.T).tocsr() - A_beta_pairwise_byInfectee = scipy.sparse.csr_matrix.multiply(self.A, self.beta_local).tocsr() - #------------------------------ - # Compute the effective pairwise beta values as a function of the infected/infectee pair: - if(self.beta_pairwise_mode == 'infected'): - self.A_beta_pairwise = A_beta_pairwise_byInfected - elif(self.beta_pairwise_mode == 'infectee'): - self.A_beta_pairwise = A_beta_pairwise_byInfectee - elif(self.beta_pairwise_mode == 'min'): - self.A_beta_pairwise = scipy.sparse.csr_matrix.minimum(A_beta_pairwise_byInfected, A_beta_pairwise_byInfectee) - elif(self.beta_pairwise_mode == 'max'): - self.A_beta_pairwise = scipy.sparse.csr_matrix.maximum(A_beta_pairwise_byInfected, A_beta_pairwise_byInfectee) - elif(self.beta_pairwise_mode == 'mean' or self.beta_pairwise_mode is None): - self.A_beta_pairwise = (A_beta_pairwise_byInfected + A_beta_pairwise_byInfectee)/2 - else: - print("Unrecognized beta_pairwise_mode value (support for 'infected', 'infectee', 'min', 'max', and 'mean').") - else: - print("Invalid values given for beta_local (expected 1xN list/array or NxN 2d array)") - #---------------------------------------- - if(self.beta_Q_local.ndim == 2 and self.beta_Q_local.shape[0] == self.numNodes and self.beta_Q_local.shape[1] == self.numNodes): - self.A_Q_beta_Q_pairwise = self.beta_Q_local - elif((self.beta_Q_local.ndim == 1 and self.beta_Q_local.shape[0] == self.numNodes) or (self.beta_Q_local.ndim == 2 and (self.beta_Q_local.shape[0] == self.numNodes or self.beta_Q_local.shape[1] == self.numNodes))): - self.beta_Q_local = self.beta_Q_local.reshape((self.numNodes,1)) - # Pre-multiply beta_Q values by the isolation adjacency matrix ("transmission weight connections") - A_Q_beta_Q_pairwise_byInfected = scipy.sparse.csr_matrix.multiply(self.A_Q, self.beta_Q_local.T).tocsr() - A_Q_beta_Q_pairwise_byInfectee = scipy.sparse.csr_matrix.multiply(self.A_Q, self.beta_Q_local).tocsr() - #------------------------------ - # Compute the effective pairwise beta values as a function of the infected/infectee pair: - if(self.beta_pairwise_mode == 'infected'): - self.A_Q_beta_Q_pairwise = A_Q_beta_Q_pairwise_byInfected - elif(self.beta_pairwise_mode == 'infectee'): - self.A_Q_beta_Q_pairwise = A_Q_beta_Q_pairwise_byInfectee - elif(self.beta_pairwise_mode == 'min'): - self.A_Q_beta_Q_pairwise = scipy.sparse.csr_matrix.minimum(A_Q_beta_Q_pairwise_byInfected, A_Q_beta_Q_pairwise_byInfectee) - elif(self.beta_pairwise_mode == 'max'): - self.A_Q_beta_Q_pairwise = scipy.sparse.csr_matrix.maximum(A_Q_beta_Q_pairwise_byInfected, A_Q_beta_Q_pairwise_byInfectee) - elif(self.beta_pairwise_mode == 'mean' or self.beta_pairwise_mode is None): - self.A_Q_beta_Q_pairwise = (A_Q_beta_Q_pairwise_byInfected + A_Q_beta_Q_pairwise_byInfectee)/2 - else: - print("Unrecognized beta_pairwise_mode value (support for 'infected', 'infectee', 'min', 'max', and 'mean').") - else: - print("Invalid values given for beta_Q_local (expected 1xN list/array or NxN 2d array)") - #---------------------------------------- - - #---------------------------------------- - # Degree-based transmission scaling parameters: - #---------------------------------------- - self.delta_pairwise_mode = self.parameters['delta_pairwise_mode'] - with numpy.errstate(divide='ignore'): # ignore log(0) warning, then convert log(0) = -inf -> 0.0 - self.delta = numpy.log(self.degree)/numpy.log(numpy.mean(self.degree)) if self.parameters['delta'] is None else numpy.array(self.parameters['delta']) if isinstance(self.parameters['delta'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['delta'], shape=(self.numNodes,1)) - self.delta_Q = numpy.log(self.degree_Q)/numpy.log(numpy.mean(self.degree_Q)) if self.parameters['delta_Q'] is None else numpy.array(self.parameters['delta_Q']) if isinstance(self.parameters['delta_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['delta_Q'], shape=(self.numNodes,1)) - self.delta[numpy.isneginf(self.delta)] = 0.0 - self.delta_Q[numpy.isneginf(self.delta_Q)] = 0.0 - #---------------------------------------- - if(self.delta.ndim == 2 and self.delta.shape[0] == self.numNodes and self.delta.shape[1] == self.numNodes): - self.A_delta_pairwise = self.delta - elif((self.delta.ndim == 1 and self.delta.shape[0] == self.numNodes) or (self.delta.ndim == 2 and (self.delta.shape[0] == self.numNodes or self.delta.shape[1] == self.numNodes))): - self.delta = self.delta.reshape((self.numNodes,1)) - # Pre-multiply delta values by the adjacency matrix ("transmission weight connections") - A_delta_pairwise_byInfected = scipy.sparse.csr_matrix.multiply(self.A, self.delta.T).tocsr() - A_delta_pairwise_byInfectee = scipy.sparse.csr_matrix.multiply(self.A, self.delta).tocsr() - #------------------------------ - # Compute the effective pairwise delta values as a function of the infected/infectee pair: - if(self.delta_pairwise_mode == 'infected'): - self.A_delta_pairwise = A_delta_pairwise_byInfected - elif(self.delta_pairwise_mode == 'infectee'): - self.A_delta_pairwise = A_delta_pairwise_byInfectee - elif(self.delta_pairwise_mode == 'min'): - self.A_delta_pairwise = scipy.sparse.csr_matrix.minimum(A_delta_pairwise_byInfected, A_delta_pairwise_byInfectee) - elif(self.delta_pairwise_mode == 'max'): - self.A_delta_pairwise = scipy.sparse.csr_matrix.maximum(A_delta_pairwise_byInfected, A_delta_pairwise_byInfectee) - elif(self.delta_pairwise_mode == 'mean'): - self.A_delta_pairwise = (A_delta_pairwise_byInfected + A_delta_pairwise_byInfectee)/2 - elif(self.delta_pairwise_mode is None): - self.A_delta_pairwise = self.A - else: - print("Unrecognized delta_pairwise_mode value (support for 'infected', 'infectee', 'min', 'max', and 'mean').") - else: - print("Invalid values given for delta (expected 1xN list/array or NxN 2d array)") - #---------------------------------------- - if(self.delta_Q.ndim == 2 and self.delta_Q.shape[0] == self.numNodes and self.delta_Q.shape[1] == self.numNodes): - self.A_Q_delta_Q_pairwise = self.delta_Q - elif((self.delta_Q.ndim == 1 and self.delta_Q.shape[0] == self.numNodes) or (self.delta_Q.ndim == 2 and (self.delta_Q.shape[0] == self.numNodes or self.delta_Q.shape[1] == self.numNodes))): - self.delta_Q = self.delta_Q.reshape((self.numNodes,1)) - # Pre-multiply delta_Q values by the isolation adjacency matrix ("transmission weight connections") - A_Q_delta_Q_pairwise_byInfected = scipy.sparse.csr_matrix.multiply(self.A_Q, self.delta_Q).tocsr() - A_Q_delta_Q_pairwise_byInfectee = scipy.sparse.csr_matrix.multiply(self.A_Q, self.delta_Q.T).tocsr() - #------------------------------ - # Compute the effective pairwise delta values as a function of the infected/infectee pair: - if(self.delta_pairwise_mode == 'infected'): - self.A_Q_delta_Q_pairwise = A_Q_delta_Q_pairwise_byInfected - elif(self.delta_pairwise_mode == 'infectee'): - self.A_Q_delta_Q_pairwise = A_Q_delta_Q_pairwise_byInfectee - elif(self.delta_pairwise_mode == 'min'): - self.A_Q_delta_Q_pairwise = scipy.sparse.csr_matrix.minimum(A_Q_delta_Q_pairwise_byInfected, A_Q_delta_Q_pairwise_byInfectee) - elif(self.delta_pairwise_mode == 'max'): - self.A_Q_delta_Q_pairwise = scipy.sparse.csr_matrix.maximum(A_Q_delta_Q_pairwise_byInfected, A_Q_delta_Q_pairwise_byInfectee) - elif(self.delta_pairwise_mode == 'mean'): - self.A_Q_delta_Q_pairwise = (A_Q_delta_Q_pairwise_byInfected + A_Q_delta_Q_pairwise_byInfectee)/2 - elif(self.delta_pairwise_mode is None): - self.A_Q_delta_Q_pairwise = self.A - else: - print("Unrecognized delta_pairwise_mode value (support for 'infected', 'infectee', 'min', 'max', and 'mean').") - else: - print("Invalid values given for delta_Q (expected 1xN list/array or NxN 2d array)") - - #---------------------------------------- - # Pre-calculate the pairwise delta*beta values: - #---------------------------------------- - self.A_deltabeta = scipy.sparse.csr_matrix.multiply(self.A_delta_pairwise, self.A_beta_pairwise) - self.A_Q_deltabeta_Q = scipy.sparse.csr_matrix.multiply(self.A_Q_delta_Q_pairwise, self.A_Q_beta_Q_pairwise) - - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def node_degrees(self, Amat): - return Amat.sum(axis=0).reshape(self.numNodes,1) # sums of adj matrix cols - - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def total_num_susceptible(self, t_idx=None): - if(t_idx is None): - return (self.numS[:]) - else: - return (self.numS[t_idx]) - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def total_num_infected(self, t_idx=None): - if(t_idx is None): - return (self.numE[:] + self.numI[:] + self.numQ_E[:] + self.numQ_I[:]) - else: - return (self.numE[t_idx] + self.numI[t_idx] + self.numQ_E[t_idx] + self.numQ_I[t_idx]) - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def total_num_isolated(self, t_idx=None): - if(t_idx is None): - return (self.numQ_E[:] + self.numQ_I[:]) - else: - return (self.numQ_E[t_idx] + self.numQ_I[t_idx]) - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def total_num_tested(self, t_idx=None): - if(t_idx is None): - return (self.numTested[:]) - else: - return (self.numTested[t_idx]) - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def total_num_positive(self, t_idx=None): - if(t_idx is None): - return (self.numPositive[:]) - else: - return (self.numPositive[t_idx]) - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def total_num_recovered(self, t_idx=None): - if(t_idx is None): - return (self.numR[:]) - else: - return (self.numR[t_idx]) - - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def calc_propensities(self): - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Pre-calculate matrix multiplication terms that may be used in multiple propensity calculations, - # and check to see if their computation is necessary before doing the multiplication - #------------------------------------ - - self.transmissionTerms_I = numpy.zeros(shape=(self.numNodes,1)) - if(numpy.any(self.numI[self.tidx])): - self.transmissionTerms_I = numpy.asarray(scipy.sparse.csr_matrix.dot(self.A_deltabeta, self.X==self.I)) - - #------------------------------------ - - self.transmissionTerms_Q = numpy.zeros(shape=(self.numNodes,1)) - if(numpy.any(self.numQ_I[self.tidx])): - self.transmissionTerms_Q = numpy.asarray(scipy.sparse.csr_matrix.dot(self.A_Q_deltabeta_Q, self.X==self.Q_I)) - - #------------------------------------ - - numContacts_Q = numpy.zeros(shape=(self.numNodes,1)) - if(numpy.any(self.positive) and (numpy.any(self.phi_E) or numpy.any(self.phi_I))): - numContacts_Q = numpy.asarray(scipy.sparse.csr_matrix.dot(self.A, ((self.positive)&(self.X!=self.R)&(self.X!=self.F)))) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - propensities_StoE = (self.alpha * - (self.p*((self.beta_global*self.numI[self.tidx] + self.q*self.beta_Q_global*self.numQ_I[self.tidx])/self.N[self.tidx]) - + (1-self.p)*(numpy.divide(self.transmissionTerms_I, self.degree, out=numpy.zeros_like(self.degree), where=self.degree!=0) - +numpy.divide(self.transmissionTerms_Q, self.degree_Q, out=numpy.zeros_like(self.degree_Q), where=self.degree_Q!=0))) - )*(self.X==self.S) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - if(self.transition_mode == 'time_in_state'): - - propensities_EtoI = 1e5 * ((self.X==self.E) & numpy.greater(self.timer_state, 1/self.sigma)) - - propensities_ItoR = 1e5 * ((self.X==self.I) & numpy.greater(self.timer_state, 1/self.gamma) & numpy.greater_equal(self.rand_f, self.f)) - - propensities_ItoF = 1e5 * ((self.X==self.I) & numpy.greater(self.timer_state, 1/self.mu_I) & numpy.less(self.rand_f, self.f)) - - propensities_EtoQE = numpy.zeros_like(propensities_StoE) - - propensities_ItoQI = numpy.zeros_like(propensities_StoE) - - propensities_QEtoQI = 1e5 * ((self.X==self.Q_E) & numpy.greater(self.timer_state, 1/self.sigma_Q)) - - propensities_QItoR = 1e5 * ((self.X==self.Q_I) & numpy.greater(self.timer_state, 1/self.gamma_Q) & numpy.greater_equal(self.rand_f, self.f)) - - propensities_QItoF = 1e5 * ((self.X==self.Q_I) & numpy.greater(self.timer_state, 1/self.mu_Q) & numpy.less(self.rand_f, self.f)) - - propensities_RtoS = 1e5 * ((self.X==self.R) & numpy.greater(self.timer_state, 1/self.xi)) - - propensities__toS = 1e5 * ((self.X!=self.F) & numpy.greater(self.timer_state, 1/self.nu)) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - else: # exponential_rates - - propensities_EtoI = self.sigma * (self.X==self.E) - - propensities_ItoR = self.gamma * ((self.X==self.I) & (numpy.greater_equal(self.rand_f, self.f))) - - propensities_ItoF = self.mu_I * ((self.X==self.I) & (numpy.less(self.rand_f, self.f))) - - propensities_EtoQE = (self.theta_E + self.phi_E*numContacts_Q)*self.psi_E * (self.X==self.E) - - propensities_ItoQI = (self.theta_I + self.phi_I*numContacts_Q)*self.psi_I * (self.X==self.I) - - propensities_QEtoQI = self.sigma_Q * (self.X==self.Q_E) - - propensities_QItoR = self.gamma_Q * ((self.X==self.Q_I) & (numpy.greater_equal(self.rand_f, self.f))) - - propensities_QItoF = self.mu_Q * ((self.X==self.Q_I) & (numpy.less(self.rand_f, self.f))) - - propensities_RtoS = self.xi * (self.X==self.R) - - propensities__toS = self.nu * (self.X!=self.F) - - - - - - - - - - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - propensities = numpy.hstack([propensities_StoE, propensities_EtoI, - propensities_ItoR, propensities_ItoF, - propensities_EtoQE, propensities_ItoQI, propensities_QEtoQI, - propensities_QItoR, propensities_QItoF, - propensities_RtoS, propensities__toS]) - - columns = ['StoE', 'EtoI', 'ItoR', 'ItoF', 'EtoQE', 'ItoQI', 'QEtoQI', 'QItoR', 'QItoF', 'RtoS', '_toS'] - - return propensities, columns - - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def set_isolation(self, node, isolate): - # Move this node in/out of the appropriate isolation state: - if(isolate == True): - if(self.X[node] == self.E): - self.X[node] = self.Q_E - self.timer_state = 0 - elif(self.X[node] == self.I): - self.X[node] = self.Q_I - self.timer_state = 0 - elif(isolate == False): - if(self.X[node] == self.Q_E): - self.X[node] = self.E - self.timer_state = 0 - elif(self.X[node] == self.Q_I): - self.X[node] = self.I - self.timer_state = 0 - # Reset the isolation timer: - self.timer_isolation[node] = 0 - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def set_tested(self, node, tested): - self.tested[node] = tested - self.testedInCurrentState[node] = tested - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def set_positive(self, node, positive): - self.positive[node] = positive - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def introduce_exposures(self, num_new_exposures): - exposedNodes = numpy.random.choice(range(self.numNodes), size=num_new_exposures, replace=False) - for exposedNode in exposedNodes: - if(self.X[exposedNode]==self.S): - self.X[exposedNode] = self.E - - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def increase_data_series_length(self): - self.tseries = numpy.pad(self.tseries, [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.numS = numpy.pad(self.numS, [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.numE = numpy.pad(self.numE, [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.numI = numpy.pad(self.numI, [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.numR = numpy.pad(self.numR, [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.numF = numpy.pad(self.numF, [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.numQ_E = numpy.pad(self.numQ_E, [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.numQ_I = numpy.pad(self.numQ_I, [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.N = numpy.pad(self.N, [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.numTested = numpy.pad(self.numTested, [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.numPositive = numpy.pad(self.numPositive, [(0, 6*self.numNodes)], mode='constant', constant_values=0) - - if(self.store_Xseries): - self.Xseries = numpy.pad(self.Xseries, [(0, 6*self.numNodes), (0,0)], mode='constant', constant_values=0) - - if(self.nodeGroupData): - for groupName in self.nodeGroupData: - self.nodeGroupData[groupName]['numS'] = numpy.pad(self.nodeGroupData[groupName]['numS'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.nodeGroupData[groupName]['numE'] = numpy.pad(self.nodeGroupData[groupName]['numE'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.nodeGroupData[groupName]['numI'] = numpy.pad(self.nodeGroupData[groupName]['numI'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.nodeGroupData[groupName]['numR'] = numpy.pad(self.nodeGroupData[groupName]['numR'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.nodeGroupData[groupName]['numF'] = numpy.pad(self.nodeGroupData[groupName]['numF'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.nodeGroupData[groupName]['numQ_E'] = numpy.pad(self.nodeGroupData[groupName]['numQ_E'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.nodeGroupData[groupName]['numQ_I'] = numpy.pad(self.nodeGroupData[groupName]['numQ_I'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.nodeGroupData[groupName]['N'] = numpy.pad(self.nodeGroupData[groupName]['N'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.nodeGroupData[groupName]['numTested'] = numpy.pad(self.nodeGroupData[groupName]['numTested'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.nodeGroupData[groupName]['numPositive'] = numpy.pad(self.nodeGroupData[groupName]['numPositive'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) - - return None - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def finalize_data_series(self): - self.tseries = numpy.array(self.tseries, dtype=float)[:self.tidx+1] - self.numS = numpy.array(self.numS, dtype=float)[:self.tidx+1] - self.numE = numpy.array(self.numE, dtype=float)[:self.tidx+1] - self.numI = numpy.array(self.numI, dtype=float)[:self.tidx+1] - self.numR = numpy.array(self.numR, dtype=float)[:self.tidx+1] - self.numF = numpy.array(self.numF, dtype=float)[:self.tidx+1] - self.numQ_E = numpy.array(self.numQ_E, dtype=float)[:self.tidx+1] - self.numQ_I = numpy.array(self.numQ_I, dtype=float)[:self.tidx+1] - self.N = numpy.array(self.N, dtype=float)[:self.tidx+1] - self.numTested = numpy.array(self.numTested, dtype=float)[:self.tidx+1] - self.numPositive = numpy.array(self.numPositive, dtype=float)[:self.tidx+1] - - if(self.store_Xseries): - self.Xseries = self.Xseries[:self.tidx+1, :] - - if(self.nodeGroupData): - for groupName in self.nodeGroupData: - self.nodeGroupData[groupName]['numS'] = numpy.array(self.nodeGroupData[groupName]['numS'], dtype=float)[:self.tidx+1] - self.nodeGroupData[groupName]['numE'] = numpy.array(self.nodeGroupData[groupName]['numE'], dtype=float)[:self.tidx+1] - self.nodeGroupData[groupName]['numI'] = numpy.array(self.nodeGroupData[groupName]['numI'], dtype=float)[:self.tidx+1] - self.nodeGroupData[groupName]['numR'] = numpy.array(self.nodeGroupData[groupName]['numR'], dtype=float)[:self.tidx+1] - self.nodeGroupData[groupName]['numF'] = numpy.array(self.nodeGroupData[groupName]['numF'], dtype=float)[:self.tidx+1] - self.nodeGroupData[groupName]['numQ_E'] = numpy.array(self.nodeGroupData[groupName]['numQ_E'], dtype=float)[:self.tidx+1] - self.nodeGroupData[groupName]['numQ_I'] = numpy.array(self.nodeGroupData[groupName]['numQ_I'], dtype=float)[:self.tidx+1] - self.nodeGroupData[groupName]['N'] = numpy.array(self.nodeGroupData[groupName]['N'], dtype=float)[:self.tidx+1] - self.nodeGroupData[groupName]['numTested'] = numpy.array(self.nodeGroupData[groupName]['numTested'], dtype=float)[:self.tidx+1] - self.nodeGroupData[groupName]['numPositive'] = numpy.array(self.nodeGroupData[groupName]['numPositive'], dtype=float)[:self.tidx+1] - - return None - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def run_iteration(self): - - if(self.tidx >= len(self.tseries)-1): - # Room has run out in the timeseries storage arrays; double the size of these arrays: - self.increase_data_series_length() - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Generate 2 random numbers uniformly distributed in (0,1) - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - r1 = numpy.random.rand() - r2 = numpy.random.rand() - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Calculate propensities - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - propensities, transitionTypes = self.calc_propensities() - - if(propensities.sum() > 0): - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Calculate alpha - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - propensities_flat = propensities.ravel(order='F') - cumsum = propensities_flat.cumsum() - alpha = propensities_flat.sum() - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Compute the time until the next event takes place - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - tau = (1/alpha)*numpy.log(float(1/r1)) - self.t += tau - self.timer_state += tau - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Compute which event takes place - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - transitionIdx = numpy.searchsorted(cumsum,r2*alpha) - transitionNode = transitionIdx % self.numNodes - transitionType = transitionTypes[ int(transitionIdx/self.numNodes) ] - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Perform updates triggered by rate propensities: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - assert(self.X[transitionNode] == self.transitions[transitionType]['currentState'] and self.X[transitionNode]!=self.F), "Assertion error: Node "+str(transitionNode)+" has unexpected current state "+str(self.X[transitionNode])+" given the intended transition of "+str(transitionType)+"." - self.X[transitionNode] = self.transitions[transitionType]['newState'] - - self.testedInCurrentState[transitionNode] = False - - self.timer_state[transitionNode] = 0.0 - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - # Save information about infection events when they occur: - if(transitionType == 'StoE'): - transitionNode_GNbrs = list(self.G[transitionNode].keys()) - transitionNode_GQNbrs = list(self.G_Q[transitionNode].keys()) - self.infectionsLog.append({ 't': self.t, - 'infected_node': transitionNode, - 'infection_type': transitionType, - 'infected_node_degree': self.degree[transitionNode], - 'local_contact_nodes': transitionNode_GNbrs, - 'local_contact_node_states': self.X[transitionNode_GNbrs].flatten(), - 'isolation_contact_nodes': transitionNode_GQNbrs, - 'isolation_contact_node_states':self.X[transitionNode_GQNbrs].flatten() }) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - if(transitionType in ['EtoQE', 'ItoQI']): - self.set_positive(node=transitionNode, positive=True) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - else: - - tau = 0.01 - self.t += tau - self.timer_state += tau - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - self.tidx += 1 - - self.tseries[self.tidx] = self.t - self.numS[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.S), a_min=0, a_max=self.numNodes) - self.numE[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.E), a_min=0, a_max=self.numNodes) - self.numI[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.I), a_min=0, a_max=self.numNodes) - self.numF[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.F), a_min=0, a_max=self.numNodes) - self.numQ_E[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.Q_E), a_min=0, a_max=self.numNodes) - self.numQ_I[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.Q_I), a_min=0, a_max=self.numNodes) - self.numTested[self.tidx] = numpy.clip(numpy.count_nonzero(self.tested), a_min=0, a_max=self.numNodes) - self.numPositive[self.tidx] = numpy.clip(numpy.count_nonzero(self.positive), a_min=0, a_max=self.numNodes) - - self.N[self.tidx] = numpy.clip((self.numNodes - self.numF[self.tidx]), a_min=0, a_max=self.numNodes) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Update testing and isolation statuses - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - isolatedNodes = numpy.argwhere((self.X==self.Q_E)|(self.X==self.Q_I))[:,0].flatten() - self.timer_isolation[isolatedNodes] = self.timer_isolation[isolatedNodes] + tau - - nodesExitingIsolation = numpy.argwhere(self.timer_isolation >= self.isolationTime) - for isoNode in nodesExitingIsolation: - self.set_isolation(node=isoNode, isolate=False) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Store system states - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(self.store_Xseries): - self.Xseries[self.tidx,:] = self.X.T - - if(self.nodeGroupData): - for groupName in self.nodeGroupData: - self.nodeGroupData[groupName]['numS'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.S) - self.nodeGroupData[groupName]['numE'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.E) - self.nodeGroupData[groupName]['numI'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.I) - self.nodeGroupData[groupName]['numR'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.R) - self.nodeGroupData[groupName]['numF'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.F) - self.nodeGroupData[groupName]['numQ_E'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_E) - self.nodeGroupData[groupName]['numQ_I'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_I) - self.nodeGroupData[groupName]['N'][self.tidx] = numpy.clip((self.nodeGroupData[groupName]['numS'][0] + self.nodeGroupData[groupName]['numE'][0] + self.nodeGroupData[groupName]['numI'][0] + self.nodeGroupData[groupName]['numQ_E'][0] + self.nodeGroupData[groupName]['numQ_I'][0] + self.nodeGroupData[groupName]['numR'][0]), a_min=0, a_max=self.numNodes) - self.nodeGroupData[groupName]['numTested'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.tested) - self.nodeGroupData[groupName]['numPositive'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.positive) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Terminate if tmax reached or num infections is 0: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(self.t >= self.tmax or (self.total_num_infected(self.tidx) < 1 and self.total_num_isolated(self.tidx) < 1)): - self.finalize_data_series() - return False - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - return True - - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def run(self, T, checkpoints=None, print_interval=10, verbose='t'): - if(T>0): - self.tmax += T - else: - return False - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Pre-process checkpoint values: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(checkpoints): - numCheckpoints = len(checkpoints['t']) - for chkpt_param, chkpt_values in checkpoints.items(): - assert(isinstance(chkpt_values, (list, numpy.ndarray)) and len(chkpt_values)==numCheckpoints), "Expecting a list of values with length equal to number of checkpoint times ("+str(numCheckpoints)+") for each checkpoint parameter." - checkpointIdx = numpy.searchsorted(checkpoints['t'], self.t) # Finds 1st index in list greater than given val - if(checkpointIdx >= numCheckpoints): - # We are out of checkpoints, stop checking them: - checkpoints = None - else: - checkpointTime = checkpoints['t'][checkpointIdx] - - #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% - # Run the simulation loop: - #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% - print_reset = True - running = True - while running: - - running = self.run_iteration() - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Handle checkpoints if applicable: - if(checkpoints): - if(self.t >= checkpointTime): - if(verbose is not False): - print("[Checkpoint: Updating parameters]") - # A checkpoint has been reached, update param values: - for param in list(self.parameters.keys()): - if(param in list(checkpoints.keys())): - self.parameters.update({param: checkpoints[param][checkpointIdx]}) - # Update parameter data structures and scenario flags: - self.update_parameters() - # Update the next checkpoint time: - checkpointIdx = numpy.searchsorted(checkpoints['t'], self.t) # Finds 1st index in list greater than given val - if(checkpointIdx >= numCheckpoints): - # We are out of checkpoints, stop checking them: - checkpoints = None - else: - checkpointTime = checkpoints['t'][checkpointIdx] - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - if(print_interval): - if(print_reset and (int(self.t) % print_interval == 0)): - if(verbose=="t"): - print("t = %.2f" % self.t) - if(verbose==True): - print("t = %.2f" % self.t) - print("\t S = " + str(self.numS[self.tidx])) - print("\t E = " + str(self.numE[self.tidx])) - print("\t I = " + str(self.numI[self.tidx])) - print("\t R = " + str(self.numR[self.tidx])) - print("\t F = " + str(self.numF[self.tidx])) - print("\t Q_E = " + str(self.numQ_E[self.tidx])) - print("\t Q_I = " + str(self.numQ_I[self.tidx])) - print_reset = False - elif(not print_reset and (int(self.t) % 10 != 0)): - print_reset = True - - return True - - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='line', plot_F='line', - plot_Q_E='line', plot_Q_I='line', combine_D=True, - color_S='tab:green', color_E='orange', color_I='crimson', color_R='tab:blue', color_F='black', - color_Q_E='mediumorchid', color_Q_I='mediumorchid', color_reference='#E0E0E0', - dashed_reference_results=None, dashed_reference_label='reference', - shaded_reference_results=None, shaded_reference_label='reference', - vlines=[], vline_colors=[], vline_styles=[], vline_labels=[], - ylim=None, xlim=None, legend=True, title=None, side_title=None, plot_percentages=True): - - import matplotlib.pyplot as pyplot - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Create an Axes object if None provided: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(not ax): - fig, ax = pyplot.subplots() - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Prepare data series to be plotted: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - Fseries = self.numF/self.numNodes if plot_percentages else self.numF - Eseries = self.numE/self.numNodes if plot_percentages else self.numE - Dseries = (self.numQ_E+self.numQ_I)/self.numNodes if plot_percentages else (self.numQ_E+self.numQ_I) - Q_Eseries = self.numQ_E/self.numNodes if plot_percentages else self.numQ_E - Q_Iseries = self.numQ_I/self.numNodes if plot_percentages else self.numQ_I - Iseries = self.numI/self.numNodes if plot_percentages else self.numI - Rseries = self.numR/self.numNodes if plot_percentages else self.numR - Sseries = self.numS/self.numNodes if plot_percentages else self.numS - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Draw the reference data: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(dashed_reference_results): - dashedReference_tseries = dashed_reference_results.tseries[::int(self.numNodes/100)] - dashedReference_IDEstack = (dashed_reference_results.numI + dashed_reference_results.numQ_I + dashed_reference_results.numQ_E + dashed_reference_results.numE)[::int(self.numNodes/100)] / (self.numNodes if plot_percentages else 1) - ax.plot(dashedReference_tseries, dashedReference_IDEstack, color='#E0E0E0', linestyle='--', label='$I+D+E$ ('+dashed_reference_label+')', zorder=0) - if(shaded_reference_results): - shadedReference_tseries = shaded_reference_results.tseries - shadedReference_IDEstack = (shaded_reference_results.numI + shaded_reference_results.numQ_I + shaded_reference_results.numQ_E + shaded_reference_results.numE) / (self.numNodes if plot_percentages else 1) - ax.fill_between(shaded_reference_results.tseries, shadedReference_IDEstack, 0, color='#EFEFEF', label='$I+D+E$ ('+shaded_reference_label+')', zorder=0) - ax.plot(shaded_reference_results.tseries, shadedReference_IDEstack, color='#E0E0E0', zorder=1) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Draw the stacked variables: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - topstack = numpy.zeros_like(self.tseries) - if(any(Fseries) and plot_F=='stacked'): - ax.fill_between(numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, topstack+Fseries), topstack, color=color_F, alpha=0.5, label='$F$', zorder=2) - ax.plot( numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, topstack+Fseries), color=color_F, zorder=3) - topstack = topstack+Fseries - if(any(Eseries) and plot_E=='stacked'): - ax.fill_between(numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, topstack+Eseries), topstack, color=color_E, alpha=0.5, label='$E$', zorder=2) - ax.plot( numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, topstack+Eseries), color=color_E, zorder=3) - topstack = topstack+Eseries - if(combine_D and plot_Q_E=='stacked' and plot_Q_I=='stacked'): - ax.fill_between(numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, topstack+Dseries), topstack, color=color_Q_E, alpha=0.5, label='$Q_{all}$', zorder=2) - ax.plot( numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, topstack+Dseries), color=color_Q_E, zorder=3) - topstack = topstack+Dseries - else: - if(any(Q_Eseries) and plot_Q_E=='stacked'): - ax.fill_between(numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, topstack+Q_Eseries), topstack, color=color_Q_E, alpha=0.5, label='$Q_E$', zorder=2) - ax.plot( numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, topstack+Q_Eseries), color=color_Q_E, zorder=3) - topstack = topstack+Q_Eseries - if(any(Q_Iseries) and plot_Q_I=='stacked'): - ax.fill_between(numpy.ma.masked_where(Q_Iseries<=0, self.tseries), numpy.ma.masked_where(Q_Iseries<=0, topstack+Q_Iseries), topstack, color=color_Q_I, alpha=0.5, label='$Q_I$', zorder=2) - ax.plot( numpy.ma.masked_where(Q_Iseries<=0, self.tseries), numpy.ma.masked_where(Q_Iseries<=0, topstack+Q_Iseries), color=color_Q_I, zorder=3) - topstack = topstack+Q_Iseries - if(any(Iseries) and plot_I=='stacked'): - ax.fill_between(numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, topstack+Iseries), topstack, color=color_I, alpha=0.5, label='$I$', zorder=2) - ax.plot( numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, topstack+Iseries), color=color_I, zorder=3) - topstack = topstack+Iseries - if(any(Rseries) and plot_R=='stacked'): - ax.fill_between(numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, topstack+Rseries), topstack, color=color_R, alpha=0.5, label='$R$', zorder=2) - ax.plot( numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, topstack+Rseries), color=color_R, zorder=3) - topstack = topstack+Rseries - if(any(Sseries) and plot_S=='stacked'): - ax.fill_between(numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, topstack+Sseries), topstack, color=color_S, alpha=0.5, label='$S$', zorder=2) - ax.plot( numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, topstack+Sseries), color=color_S, zorder=3) - topstack = topstack+Sseries - - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Draw the shaded variables: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(any(Fseries) and plot_F=='shaded'): - ax.fill_between(numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, Fseries), 0, color=color_F, alpha=0.5, label='$F$', zorder=4) - ax.plot( numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, Fseries), color=color_F, zorder=5) - if(any(Eseries) and plot_E=='shaded'): - ax.fill_between(numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, Eseries), 0, color=color_E, alpha=0.5, label='$E$', zorder=4) - ax.plot( numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, Eseries), color=color_E, zorder=5) - if(combine_D and (any(Dseries) and plot_Q_E=='shaded' and plot_Q_I=='shaded')): - ax.fill_between(numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, Dseries), 0, color=color_Q_E, alpha=0.5, label='$Q_{all}$', zorder=4) - ax.plot( numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, Dseries), color=color_Q_E, zorder=5) - else: - if(any(Q_Eseries) and plot_Q_E=='shaded'): - ax.fill_between(numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, Q_Eseries), 0, color=color_Q_E, alpha=0.5, label='$Q_E$', zorder=4) - ax.plot( numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, Q_Eseries), color=color_Q_E, zorder=5) - if(any(Q_Iseries) and plot_Q_I=='shaded'): - ax.fill_between(numpy.ma.masked_where(Q_Iseries<=0, self.tseries), numpy.ma.masked_where(Q_Iseries<=0, Q_Iseries), 0, color=color_Q_I, alpha=0.5, label='$Q_I$', zorder=4) - ax.plot( numpy.ma.masked_where(Q_Iseries<=0, self.tseries), numpy.ma.masked_where(Q_Iseries<=0, Q_Iseries), color=color_Q_I, zorder=5) - if(any(Iseries) and plot_I=='shaded'): - ax.fill_between(numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, Iseries), 0, color=color_I, alpha=0.5, label='$I$', zorder=4) - ax.plot( numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, Iseries), color=color_I, zorder=5) - if(any(Sseries) and plot_S=='shaded'): - ax.fill_between(numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, Sseries), 0, color=color_S, alpha=0.5, label='$S$', zorder=4) - ax.plot( numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, Sseries), color=color_S, zorder=5) - if(any(Rseries) and plot_R=='shaded'): - ax.fill_between(numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, Rseries), 0, color=color_R, alpha=0.5, label='$R$', zorder=4) - ax.plot( numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, Rseries), color=color_R, zorder=5) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Draw the line variables: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(any(Fseries) and plot_F=='line'): - ax.plot(numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, Fseries), color=color_F, label='$F$', zorder=6) - if(any(Eseries) and plot_E=='line'): - ax.plot(numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, Eseries), color=color_E, label='$E$', zorder=6) - if(combine_D and (any(Dseries) and plot_Q_E=='line' and plot_Q_I=='line')): - ax.plot(numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, Dseries), color=color_Q_E, label='$Q_{all}$', zorder=6) - else: - if(any(Q_Eseries) and plot_Q_E=='line'): - ax.plot(numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, Q_Eseries), color=color_Q_E, label='$Q_E$', zorder=6) - if(any(Q_Iseries) and plot_Q_I=='line'): - ax.plot(numpy.ma.masked_where(Q_Iseries<=0, self.tseries), numpy.ma.masked_where(Q_Iseries<=0, Q_Iseries), color=color_Q_I, label='$Q_I$', zorder=6) - if(any(Iseries) and plot_I=='line'): - ax.plot(numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, Iseries), color=color_I, label='$I$', zorder=6) - if(any(Sseries) and plot_S=='line'): - ax.plot(numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, Sseries), color=color_S, label='$S$', zorder=6) - if(any(Rseries) and plot_R=='line'): - ax.plot(numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, Rseries), color=color_R, label='$R$', zorder=6) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Draw the vertical line annotations: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(len(vlines)>0 and len(vline_colors)==0): - vline_colors = ['gray']*len(vlines) - if(len(vlines)>0 and len(vline_labels)==0): - vline_labels = [None]*len(vlines) - if(len(vlines)>0 and len(vline_styles)==0): - vline_styles = [':']*len(vlines) - for vline_x, vline_color, vline_style, vline_label in zip(vlines, vline_colors, vline_styles, vline_labels): - if(vline_x is not None): - ax.axvline(x=vline_x, color=vline_color, linestyle=vline_style, alpha=1, label=vline_label) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Draw the plot labels: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - ax.set_xlabel('days') - ax.set_ylabel('percent of population' if plot_percentages else 'number of individuals') - ax.set_xlim(0, (max(self.tseries) if not xlim else xlim)) - ax.set_ylim(0, ylim) - if(plot_percentages): - ax.set_yticklabels(['{:,.0%}'.format(y) for y in ax.get_yticks()]) - if(legend): - legend_handles, legend_labels = ax.get_legend_handles_labels() - ax.legend(legend_handles[::-1], legend_labels[::-1], loc='upper right', facecolor='white', edgecolor='none', framealpha=0.9, prop={'size': 8}) - if(title): - ax.set_title(title, size=12) - if(side_title): - ax.annotate(side_title, (0, 0.5), xytext=(-45, 0), ha='right', va='center', - size=12, rotation=90, xycoords='axes fraction', textcoords='offset points') - - return ax - - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def figure_basic(self, plot_S='line', plot_E='line', plot_I='line',plot_R='line', plot_F='line', - plot_Q_E='line', plot_Q_I='line', combine_D=True, - color_S='tab:green', color_E='orange', color_I='crimson', color_R='tab:blue', color_F='black', - color_Q_E='mediumorchid', color_Q_I='mediumorchid', color_reference='#E0E0E0', - dashed_reference_results=None, dashed_reference_label='reference', - shaded_reference_results=None, shaded_reference_label='reference', - vlines=[], vline_colors=[], vline_styles=[], vline_labels=[], - ylim=None, xlim=None, legend=True, title=None, side_title=None, plot_percentages=True, - figsize=(12,8), use_seaborn=True, show=True): - - import matplotlib.pyplot as pyplot - - fig, ax = pyplot.subplots(figsize=figsize) - - if(use_seaborn): - import seaborn - seaborn.set_style('ticks') - seaborn.despine() - - self.plot(ax=ax, plot_S=plot_S, plot_E=plot_E, plot_I=plot_I,plot_R=plot_R, plot_F=plot_F, - plot_Q_E=plot_Q_E, plot_Q_I=plot_Q_I, combine_D=combine_D, - color_S=color_S, color_E=color_E, color_I=color_I, color_R=color_R, color_F=color_F, - color_Q_E=color_Q_E, color_Q_I=color_Q_I, color_reference=color_reference, - dashed_reference_results=dashed_reference_results, dashed_reference_label=dashed_reference_label, - shaded_reference_results=shaded_reference_results, shaded_reference_label=shaded_reference_label, - vlines=vlines, vline_colors=vline_colors, vline_styles=vline_styles, vline_labels=vline_labels, - ylim=ylim, xlim=xlim, legend=legend, title=title, side_title=side_title, plot_percentages=plot_percentages) - - if(show): - pyplot.show() - - return fig, ax - - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def figure_infections(self, plot_S=False, plot_E='stacked', plot_I='stacked',plot_R=False, plot_F=False, - plot_Q_E='stacked', plot_Q_I='stacked', combine_D=True, - color_S='tab:green', color_E='orange', color_I='crimson', color_R='tab:blue', color_F='black', - color_Q_E='mediumorchid', color_Q_I='mediumorchid', color_reference='#E0E0E0', - dashed_reference_results=None, dashed_reference_label='reference', - shaded_reference_results=None, shaded_reference_label='reference', - vlines=[], vline_colors=[], vline_styles=[], vline_labels=[], - ylim=None, xlim=None, legend=True, title=None, side_title=None, plot_percentages=True, - figsize=(12,8), use_seaborn=True, show=True): - - import matplotlib.pyplot as pyplot - - fig, ax = pyplot.subplots(figsize=figsize) - - if(use_seaborn): - import seaborn - seaborn.set_style('ticks') - seaborn.despine() - - self.plot(ax=ax, plot_S=plot_S, plot_E=plot_E, plot_I=plot_I,plot_R=plot_R, plot_F=plot_F, - plot_Q_E=plot_Q_E, plot_Q_I=plot_Q_I, combine_D=combine_D, - color_S=color_S, color_E=color_E, color_I=color_I, color_R=color_R, color_F=color_F, - color_Q_E=color_Q_E, color_Q_I=color_Q_I, color_reference=color_reference, - dashed_reference_results=dashed_reference_results, dashed_reference_label=dashed_reference_label, - shaded_reference_results=shaded_reference_results, shaded_reference_label=shaded_reference_label, - vlines=vlines, vline_colors=vline_colors, vline_styles=vline_styles, vline_labels=vline_labels, - ylim=ylim, xlim=xlim, legend=legend, title=title, side_title=side_title, plot_percentages=plot_percentages) - - if(show): - pyplot.show() - - return fig, ax - - -#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% -#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% -#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% -#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% - - - - -######################################################## -#@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@# -#@ @# -#@ EXTENDED SEIRS MODELS @# -#@ @# -#@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@# -######################################################## - - -class ExtSEIRSNetworkModel(): - """ - A class to simulate the Extended SEIRS Stochastic Network Model - =================================================== - Params: - G Network adjacency matrix (numpy array) or Networkx graph object. - beta Rate of transmission (global interactions) - beta_local Rate(s) of transmission between adjacent individuals (optional) - beta_asym Rate of transmission (global interactions) - beta_asym_local Rate(s) of transmission between adjacent individuals (optional) - sigma Rate of progression to infectious state (inverse of latent period) - lamda Rate of progression to infectious (a)symptomatic state (inverse of prodromal period) - eta Rate of progression to hospitalized state (inverse of onset-to-admission period) - gamma Rate of recovery for non-hospitalized symptomatic individuals (inverse of symptomatic infectious period) - gamma_asym Rate of recovery for asymptomatic individuals (inverse of asymptomatic infectious period) - gamma_H Rate of recovery for hospitalized symptomatic individuals (inverse of hospitalized infectious period) - mu_H Rate of death for hospitalized individuals (inverse of admission-to-death period) - xi Rate of re-susceptibility (upon recovery) - mu_0 Rate of baseline death - nu Rate of baseline birth - a Probability of an infected individual remaining asymptomatic - h Probability of a symptomatic individual being hospitalized - f Probability of death for hospitalized individuals (case fatality rate) - p Probability of individuals interacting with global population - - G_Q Quarantine adjacency matrix (numpy array) or Networkx graph object. - beta_Q Rate of transmission for isolated individuals (global interactions) - beta_Q_local Rate(s) of transmission (exposure) for adjacent isolated individuals (optional) - sigma_Q Rate of progression to infectious state for isolated individuals - lamda_Q Rate of progression to infectious (a)symptomatic state for isolated individuals - eta_Q Rate of progression to hospitalized state for isolated individuals - gamma_Q_sym Rate of recovery for non-hospitalized symptomatic individuals for isolated individuals - gamma_Q_asym Rate of recovery for asymptomatic individuals for isolated individuals - theta_E Rate of random testing for exposed individuals - theta_pre Rate of random testing for infectious pre-symptomatic individuals - theta_sym Rate of random testing for infectious symptomatic individuals - theta_asym Rate of random testing for infectious asymptomatic individuals - phi_E Rate of testing when a close contact has tested positive for exposed individuals - phi_pre Rate of testing when a close contact has tested positive for infectious pre-symptomatic individuals - phi_sym Rate of testing when a close contact has tested positive for infectious symptomatic individuals - phi_asym Rate of testing when a close contact has tested positive for infectious asymptomatic individuals - psi_E Probability of positive test for exposed individuals - psi_pre Probability of positive test for infectious pre-symptomatic individuals - psi_sym Probability of positive test for infectious symptomatic individuals - psi_asym Probability of positive test for infectious asymptomatic individuals - q Probability of isolated individuals interacting with global population - isolation_time Time to remain in isolation upon positive test, self-isolation, etc. - - initE Initial number of exposed individuals - initI_pre Initial number of infectious pre-symptomatic individuals - initI_sym Initial number of infectious symptomatic individuals - initI_asym Initial number of infectious asymptomatic individuals - initH Initial number of hospitalized individuals - initR Initial number of recovered individuals - initF Initial number of infection-related fatalities - initQ_S Initial number of isolated susceptible individuals - initQ_E Initial number of isolated exposed individuals - initQ_pre Initial number of isolated infectious pre-symptomatic individuals - initQ_sym Initial number of isolated infectious symptomatic individuals - initQ_asym Initial number of isolated infectious asymptomatic individuals - initQ_R Initial number of isolated recovered individuals - (all remaining nodes initialized susceptible) - """ - def __init__(self, G, beta, sigma, lamda, gamma, - gamma_asym=None, eta=0, gamma_H=None, mu_H=0, alpha=1.0, xi=0, mu_0=0, nu=0, a=0, h=0, f=0, p=0, - beta_local=None, beta_asym=None, beta_asym_local=None, beta_pairwise_mode='infected', delta=None, delta_pairwise_mode=None, - G_Q=None, beta_Q=None, beta_Q_local=None, sigma_Q=None, lamda_Q=None, eta_Q=None, gamma_Q_sym=None, gamma_Q_asym=None, alpha_Q=None, delta_Q=None, - theta_S=0, theta_E=0, theta_pre=0, theta_sym=0, theta_asym=0, phi_S=0, phi_E=0, phi_pre=0, phi_sym=0, phi_asym=0, - psi_S=0, psi_E=1, psi_pre=1, psi_sym=1, psi_asym=1, q=0, isolation_time=14, - initE=0, initI_pre=0, initI_sym=0, initI_asym=0, initH=0, initR=0, initF=0, - initQ_S=0, initQ_E=0, initQ_pre=0, initQ_sym=0, initQ_asym=0, initQ_R=0, - o=0, prevalence_ext=0, - transition_mode='exponential_rates', node_groups=None, store_Xseries=False, seed=None): - - if(seed is not None): - numpy.random.seed(seed) - self.seed = seed - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Model Parameters: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - self.parameters = { 'G':G, 'G_Q':G_Q, - 'beta':beta, 'sigma':sigma, 'lamda':lamda, 'gamma':gamma, - 'eta':eta, 'gamma_asym':gamma_asym, 'gamma_H':gamma_H, 'mu_H':mu_H, - 'xi':xi, 'mu_0':mu_0, 'nu':nu, 'a':a, 'h':h, 'f':f, 'p':p, - 'beta_local':beta_local, 'beta_asym':beta_asym, 'beta_asym_local':beta_asym_local, 'beta_pairwise_mode':beta_pairwise_mode, - 'alpha':alpha, 'delta':delta, 'delta_pairwise_mode':delta_pairwise_mode, - 'lamda_Q':lamda_Q, 'beta_Q':beta_Q, 'beta_Q_local':beta_Q_local, 'alpha_Q':alpha_Q, 'sigma_Q':sigma_Q, - 'eta_Q':eta_Q, 'gamma_Q_sym':gamma_Q_sym, 'gamma_Q_asym':gamma_Q_asym, 'delta_Q':delta_Q, - 'theta_S':theta_S, 'theta_E':theta_E, 'theta_pre':theta_pre, 'theta_sym':theta_sym, 'theta_asym':theta_asym, - 'phi_S':phi_S, 'phi_E':phi_E, 'phi_pre':phi_pre, 'phi_sym':phi_sym, 'phi_asym':phi_asym, - 'psi_S':psi_S, 'psi_E':psi_E, 'psi_pre':psi_pre, 'psi_sym':psi_sym, 'psi_asym':psi_asym, 'q':q, 'isolation_time':isolation_time, - 'initE':initE, 'initI_pre':initI_pre, 'initI_sym':initI_sym, 'initI_asym':initI_asym, - 'initH':initH, 'initR':initR, 'initF':initF, - 'initQ_S':initQ_S, 'initQ_E':initQ_E, 'initQ_pre':initQ_pre, - 'initQ_sym':initQ_sym, 'initQ_asym':initQ_asym, 'initQ_R':initQ_R, - 'o':o, 'prevalence_ext':prevalence_ext} - self.update_parameters() - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Each node can undergo 4-6 transitions (sans vitality/re-susceptibility returns to S state), - # so there are ~numNodes*6 events/timesteps expected; initialize numNodes*6 timestep slots to start - # (will be expanded during run if needed for some reason) - self.tseries = numpy.zeros(6*self.numNodes) - self.numS = numpy.zeros(6*self.numNodes) - self.numE = numpy.zeros(6*self.numNodes) - self.numI_pre = numpy.zeros(6*self.numNodes) - self.numI_sym = numpy.zeros(6*self.numNodes) - self.numI_asym = numpy.zeros(6*self.numNodes) - self.numH = numpy.zeros(6*self.numNodes) - self.numR = numpy.zeros(6*self.numNodes) - self.numF = numpy.zeros(6*self.numNodes) - self.numQ_S = numpy.zeros(6*self.numNodes) - self.numQ_E = numpy.zeros(6*self.numNodes) - self.numQ_pre = numpy.zeros(6*self.numNodes) - self.numQ_sym = numpy.zeros(6*self.numNodes) - self.numQ_asym = numpy.zeros(6*self.numNodes) - self.numQ_R = numpy.zeros(6*self.numNodes) - self.N = numpy.zeros(6*self.numNodes) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Initialize Timekeeping: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - self.t = 0 - self.tmax = 0 # will be set when run() is called - self.tidx = 0 - self.tseries[0] = 0 - - # Vectors holding the time that each node has been in a given state or in isolation: - self.timer_state = numpy.zeros((self.numNodes,1)) - self.timer_isolation = numpy.zeros(self.numNodes) - self.isolationTime = isolation_time - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Initialize Counts of inidividuals with each state: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - self.numE[0] = int(initE) - self.numI_pre[0] = int(initI_pre) - self.numI_sym[0] = int(initI_sym) - self.numI_asym[0] = int(initI_asym) - self.numH[0] = int(initH) - self.numR[0] = int(initR) - self.numF[0] = int(initF) - self.numQ_S[0] = int(initQ_S) - self.numQ_E[0] = int(initQ_E) - self.numQ_pre[0] = int(initQ_pre) - self.numQ_sym[0] = int(initQ_sym) - self.numQ_asym[0] = int(initQ_asym) - self.numQ_R[0] = int(initQ_R) - self.numS[0] = (self.numNodes - self.numE[0] - self.numI_pre[0] - self.numI_sym[0] - self.numI_asym[0] - self.numH[0] - self.numR[0] - - self.numQ_S[0] - self.numQ_E[0] - self.numQ_pre[0] - self.numQ_sym[0] - self.numQ_asym[0] - self.numQ_R[0] - self.numF[0]) - self.N[0] = self.numNodes - self.numF[0] - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Node states: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - self.S = 1 - self.E = 2 - self.I_pre = 3 - self.I_sym = 4 - self.I_asym = 5 - self.H = 6 - self.R = 7 - self.F = 8 - self.Q_S = 11 - self.Q_E = 12 - self.Q_pre = 13 - self.Q_sym = 14 - self.Q_asym = 15 - self.Q_R = 17 - - self.X = numpy.array( [self.S]*int(self.numS[0]) + [self.E]*int(self.numE[0]) - + [self.I_pre]*int(self.numI_pre[0]) + [self.I_sym]*int(self.numI_sym[0]) + [self.I_asym]*int(self.numI_asym[0]) - + [self.H]*int(self.numH[0]) + [self.R]*int(self.numR[0]) + [self.F]*int(self.numF[0]) - + [self.Q_S]*int(self.numQ_S[0]) + [self.Q_E]*int(self.numQ_E[0]) - + [self.Q_pre]*int(self.numQ_pre[0]) + [self.Q_sym]*int(self.numQ_sym[0]) + [self.Q_asym]*int(self.numQ_asym[0]) - + [self.Q_R]*int(self.numQ_R[0]) - ).reshape((self.numNodes,1)) - numpy.random.shuffle(self.X) - - self.store_Xseries = store_Xseries - if(store_Xseries): - self.Xseries = numpy.zeros(shape=(6*self.numNodes, self.numNodes), dtype='uint8') - self.Xseries[0,:] = self.X.T - - self.transitions = { - 'StoE': {'currentState':self.S, 'newState':self.E}, - 'StoQS': {'currentState':self.S, 'newState':self.Q_S}, - 'EtoIPRE': {'currentState':self.E, 'newState':self.I_pre}, - 'EtoQE': {'currentState':self.E, 'newState':self.Q_E}, - 'IPREtoISYM': {'currentState':self.I_pre, 'newState':self.I_sym}, - 'IPREtoIASYM': {'currentState':self.I_pre, 'newState':self.I_asym}, - 'IPREtoQPRE': {'currentState':self.I_pre, 'newState':self.Q_pre}, - 'ISYMtoH': {'currentState':self.I_sym, 'newState':self.H}, - 'ISYMtoR': {'currentState':self.I_sym, 'newState':self.R}, - 'ISYMtoQSYM': {'currentState':self.I_sym, 'newState':self.Q_sym}, - 'IASYMtoR': {'currentState':self.I_asym, 'newState':self.R}, - 'IASYMtoQASYM': {'currentState':self.I_asym, 'newState':self.Q_asym}, - 'HtoR': {'currentState':self.H, 'newState':self.R}, - 'HtoF': {'currentState':self.H, 'newState':self.F}, - 'RtoS': {'currentState':self.R, 'newState':self.S}, - 'QStoQE': {'currentState':self.Q_S, 'newState':self.Q_E}, - 'QEtoQPRE': {'currentState':self.Q_E, 'newState':self.Q_pre}, - 'QPREtoQSYM': {'currentState':self.Q_pre, 'newState':self.Q_sym}, - 'QPREtoQASYM': {'currentState':self.Q_pre, 'newState':self.Q_asym}, - 'QSYMtoH': {'currentState':self.Q_sym, 'newState':self.H}, - 'QSYMtoQR': {'currentState':self.Q_sym, 'newState':self.Q_R}, - 'QASYMtoQR': {'currentState':self.Q_asym, 'newState':self.Q_R}, - '_toS': {'currentState':True, 'newState':self.S}, - } - - self.transition_mode = transition_mode - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Initialize other node metadata: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - self.tested = numpy.array([False]*self.numNodes).reshape((self.numNodes,1)) - self.positive = numpy.array([False]*self.numNodes).reshape((self.numNodes,1)) - self.numTested = numpy.zeros(6*self.numNodes) - self.numPositive = numpy.zeros(6*self.numNodes) - - self.testedInCurrentState = numpy.array([False]*self.numNodes).reshape((self.numNodes,1)) - - self.infectionsLog = [] - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Initialize node subgroup data series: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - self.nodeGroupData = None - if(node_groups): - self.nodeGroupData = {} - for groupName, nodeList in node_groups.items(): - self.nodeGroupData[groupName] = {'nodes': numpy.array(nodeList), - 'mask': numpy.isin(range(self.numNodes), nodeList).reshape((self.numNodes,1))} - self.nodeGroupData[groupName]['numS'] = numpy.zeros(6*self.numNodes) - self.nodeGroupData[groupName]['numE'] = numpy.zeros(6*self.numNodes) - self.nodeGroupData[groupName]['numI_pre'] = numpy.zeros(6*self.numNodes) - self.nodeGroupData[groupName]['numI_sym'] = numpy.zeros(6*self.numNodes) - self.nodeGroupData[groupName]['numI_asym'] = numpy.zeros(6*self.numNodes) - self.nodeGroupData[groupName]['numH'] = numpy.zeros(6*self.numNodes) - self.nodeGroupData[groupName]['numR'] = numpy.zeros(6*self.numNodes) - self.nodeGroupData[groupName]['numF'] = numpy.zeros(6*self.numNodes) - self.nodeGroupData[groupName]['numQ_S'] = numpy.zeros(6*self.numNodes) - self.nodeGroupData[groupName]['numQ_E'] = numpy.zeros(6*self.numNodes) - self.nodeGroupData[groupName]['numQ_pre'] = numpy.zeros(6*self.numNodes) - self.nodeGroupData[groupName]['numQ_sym'] = numpy.zeros(6*self.numNodes) - self.nodeGroupData[groupName]['numQ_asym'] = numpy.zeros(6*self.numNodes) - self.nodeGroupData[groupName]['numQ_R'] = numpy.zeros(6*self.numNodes) - self.nodeGroupData[groupName]['N'] = numpy.zeros(6*self.numNodes) - self.nodeGroupData[groupName]['numPositive'] = numpy.zeros(6*self.numNodes) - self.nodeGroupData[groupName]['numTested'] = numpy.zeros(6*self.numNodes) - self.nodeGroupData[groupName]['numS'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.S) - self.nodeGroupData[groupName]['numE'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.E) - self.nodeGroupData[groupName]['numI_pre'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.I_pre) - self.nodeGroupData[groupName]['numI_sym'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.I_sym) - self.nodeGroupData[groupName]['numI_asym'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.I_asym) - self.nodeGroupData[groupName]['numH'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.H) - self.nodeGroupData[groupName]['numR'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.R) - self.nodeGroupData[groupName]['numF'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.F) - self.nodeGroupData[groupName]['numQ_S'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_E) - self.nodeGroupData[groupName]['numQ_E'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_E) - self.nodeGroupData[groupName]['numQ_pre'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_pre) - self.nodeGroupData[groupName]['numQ_I_sym'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_I_sym) - self.nodeGroupData[groupName]['numQ_I_asym'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_I_asym) - self.nodeGroupData[groupName]['numQ_R'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_E) - self.nodeGroupData[groupName]['N'][0] = self.numNodes - self.numF[0] - - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def update_parameters(self): - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Model graphs: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - self.G = self.parameters['G'] - # Adjacency matrix: - if type(self.G)==numpy.ndarray: - self.A = scipy.sparse.csr_matrix(self.G) - elif type(self.G)==networkx.classes.graph.Graph: - self.A = networkx.adj_matrix(self.G) # adj_matrix gives scipy.sparse csr_matrix - else: - raise BaseException("Input an adjacency matrix or networkx object only.") - self.numNodes = int(self.A.shape[1]) - self.degree = numpy.asarray(self.node_degrees(self.A)).astype(float) - #---------------------------------------- - if(self.parameters['G_Q'] is None): - self.G_Q = self.G # If no Q graph is provided, use G in its place - else: - self.G_Q = self.parameters['G_Q'] - # Quarantine Adjacency matrix: - if type(self.G_Q)==numpy.ndarray: - self.A_Q = scipy.sparse.csr_matrix(self.G_Q) - elif type(self.G_Q)==networkx.classes.graph.Graph: - self.A_Q = networkx.adj_matrix(self.G_Q) # adj_matrix gives scipy.sparse csr_matrix - else: - raise BaseException("Input an adjacency matrix or networkx object only.") - self.numNodes_Q = int(self.A_Q.shape[1]) - self.degree_Q = numpy.asarray(self.node_degrees(self.A_Q)).astype(float) - #---------------------------------------- - assert(self.numNodes == self.numNodes_Q), "The normal and quarantine adjacency graphs must be of the same size." - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Model parameters: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - self.beta = numpy.array(self.parameters['beta']).reshape((self.numNodes, 1)) if isinstance(self.parameters['beta'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['beta'], shape=(self.numNodes,1)) - self.beta_asym = (numpy.array(self.parameters['beta_asym']).reshape((self.numNodes, 1)) if isinstance(self.parameters['beta_asym'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['beta_asym'], shape=(self.numNodes,1))) if self.parameters['beta_asym'] is not None else self.beta - self.sigma = numpy.array(self.parameters['sigma']).reshape((self.numNodes, 1)) if isinstance(self.parameters['sigma'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['sigma'], shape=(self.numNodes,1)) - self.lamda = numpy.array(self.parameters['lamda']).reshape((self.numNodes, 1)) if isinstance(self.parameters['lamda'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['lamda'], shape=(self.numNodes,1)) - self.gamma = numpy.array(self.parameters['gamma']).reshape((self.numNodes, 1)) if isinstance(self.parameters['gamma'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['gamma'], shape=(self.numNodes,1)) - self.eta = numpy.array(self.parameters['eta']).reshape((self.numNodes, 1)) if isinstance(self.parameters['eta'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['eta'], shape=(self.numNodes,1)) - self.gamma_asym = (numpy.array(self.parameters['gamma_asym']).reshape((self.numNodes, 1)) if isinstance(self.parameters['gamma_asym'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['gamma_asym'], shape=(self.numNodes,1))) if self.parameters['gamma_asym'] is not None else self.gamma - self.gamma_H = (numpy.array(self.parameters['gamma_H']).reshape((self.numNodes, 1)) if isinstance(self.parameters['gamma_H'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['gamma_H'], shape=(self.numNodes,1))) if self.parameters['gamma_H'] is not None else self.gamma - self.mu_H = numpy.array(self.parameters['mu_H']).reshape((self.numNodes, 1)) if isinstance(self.parameters['mu_H'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['mu_H'], shape=(self.numNodes,1)) - self.alpha = numpy.array(self.parameters['alpha']).reshape((self.numNodes, 1)) if isinstance(self.parameters['alpha'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['alpha'], shape=(self.numNodes,1)) - self.xi = numpy.array(self.parameters['xi']).reshape((self.numNodes, 1)) if isinstance(self.parameters['xi'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['xi'], shape=(self.numNodes,1)) - self.mu_0 = numpy.array(self.parameters['mu_0']).reshape((self.numNodes, 1)) if isinstance(self.parameters['mu_0'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['mu_0'], shape=(self.numNodes,1)) - self.nu = numpy.array(self.parameters['nu']).reshape((self.numNodes, 1)) if isinstance(self.parameters['nu'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['nu'], shape=(self.numNodes,1)) - self.a = numpy.array(self.parameters['a']).reshape((self.numNodes, 1)) if isinstance(self.parameters['a'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['a'], shape=(self.numNodes,1)) - self.h = numpy.array(self.parameters['h']).reshape((self.numNodes, 1)) if isinstance(self.parameters['h'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['h'], shape=(self.numNodes,1)) - self.f = numpy.array(self.parameters['f']).reshape((self.numNodes, 1)) if isinstance(self.parameters['f'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['f'], shape=(self.numNodes,1)) - self.p = numpy.array(self.parameters['p']).reshape((self.numNodes, 1)) if isinstance(self.parameters['p'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['p'], shape=(self.numNodes,1)) - self.o = numpy.array(self.parameters['o']).reshape((self.numNodes, 1)) if isinstance(self.parameters['o'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['o'], shape=(self.numNodes,1)) - - self.rand_a = numpy.random.rand(self.a.shape[0], self.a.shape[1]) - self.rand_h = numpy.random.rand(self.h.shape[0], self.h.shape[1]) - self.rand_f = numpy.random.rand(self.f.shape[0], self.f.shape[1]) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # External infection introduction variables: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - self.prevalence_ext = numpy.array(self.parameters['prevalence_ext']).reshape((self.numNodes, 1)) if isinstance(self.parameters['prevalence_ext'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['prevalence_ext'], shape=(self.numNodes,1)) - - #---------------------------------------- - # Testing-related parameters: - #---------------------------------------- - self.beta_Q = (numpy.array(self.parameters['beta_Q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['beta_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['beta_Q'], shape=(self.numNodes,1))) if self.parameters['beta_Q'] is not None else self.beta - self.sigma_Q = (numpy.array(self.parameters['sigma_Q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['sigma_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['sigma_Q'], shape=(self.numNodes,1))) if self.parameters['sigma_Q'] is not None else self.sigma - self.lamda_Q = (numpy.array(self.parameters['lamda_Q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['lamda_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['lamda_Q'], shape=(self.numNodes,1))) if self.parameters['lamda_Q'] is not None else self.lamda - self.gamma_Q_sym = (numpy.array(self.parameters['gamma_Q_sym']).reshape((self.numNodes, 1)) if isinstance(self.parameters['gamma_Q_sym'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['gamma_Q_sym'], shape=(self.numNodes,1))) if self.parameters['gamma_Q_sym'] is not None else self.gamma - self.gamma_Q_asym = (numpy.array(self.parameters['gamma_Q_asym']).reshape((self.numNodes, 1)) if isinstance(self.parameters['gamma_Q_asym'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['gamma_Q_asym'], shape=(self.numNodes,1))) if self.parameters['gamma_Q_asym'] is not None else self.gamma - self.eta_Q = (numpy.array(self.parameters['eta_Q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['eta_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['eta_Q'], shape=(self.numNodes,1))) if self.parameters['eta_Q'] is not None else self.eta - self.alpha_Q = (numpy.array(self.parameters['alpha_Q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['alpha_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['alpha_Q'], shape=(self.numNodes,1))) if self.parameters['alpha_Q'] is not None else self.alpha - self.theta_S = numpy.array(self.parameters['theta_S']).reshape((self.numNodes, 1)) if isinstance(self.parameters['theta_S'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['theta_S'], shape=(self.numNodes,1)) - self.theta_E = numpy.array(self.parameters['theta_E']).reshape((self.numNodes, 1)) if isinstance(self.parameters['theta_E'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['theta_E'], shape=(self.numNodes,1)) - self.theta_pre = numpy.array(self.parameters['theta_pre']).reshape((self.numNodes, 1)) if isinstance(self.parameters['theta_pre'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['theta_pre'], shape=(self.numNodes,1)) - self.theta_sym = numpy.array(self.parameters['theta_sym']).reshape((self.numNodes, 1)) if isinstance(self.parameters['theta_sym'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['theta_sym'], shape=(self.numNodes,1)) - self.theta_asym = numpy.array(self.parameters['theta_asym']).reshape((self.numNodes, 1)) if isinstance(self.parameters['theta_asym'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['theta_asym'], shape=(self.numNodes,1)) - self.phi_S = numpy.array(self.parameters['phi_S']).reshape((self.numNodes, 1)) if isinstance(self.parameters['phi_S'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['phi_S'], shape=(self.numNodes,1)) - self.phi_E = numpy.array(self.parameters['phi_E']).reshape((self.numNodes, 1)) if isinstance(self.parameters['phi_E'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['phi_E'], shape=(self.numNodes,1)) - self.phi_pre = numpy.array(self.parameters['phi_pre']).reshape((self.numNodes, 1)) if isinstance(self.parameters['phi_pre'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['phi_pre'], shape=(self.numNodes,1)) - self.phi_sym = numpy.array(self.parameters['phi_sym']).reshape((self.numNodes, 1)) if isinstance(self.parameters['phi_sym'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['phi_sym'], shape=(self.numNodes,1)) - self.phi_asym = numpy.array(self.parameters['phi_asym']).reshape((self.numNodes, 1)) if isinstance(self.parameters['phi_asym'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['phi_asym'], shape=(self.numNodes,1)) - self.psi_S = numpy.array(self.parameters['psi_S']).reshape((self.numNodes, 1)) if isinstance(self.parameters['psi_S'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['psi_S'], shape=(self.numNodes,1)) - self.psi_E = numpy.array(self.parameters['psi_E']).reshape((self.numNodes, 1)) if isinstance(self.parameters['psi_E'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['psi_E'], shape=(self.numNodes,1)) - self.psi_pre = numpy.array(self.parameters['psi_pre']).reshape((self.numNodes, 1)) if isinstance(self.parameters['psi_pre'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['psi_pre'], shape=(self.numNodes,1)) - self.psi_sym = numpy.array(self.parameters['psi_sym']).reshape((self.numNodes, 1)) if isinstance(self.parameters['psi_sym'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['psi_sym'], shape=(self.numNodes,1)) - self.psi_asym = numpy.array(self.parameters['psi_asym']).reshape((self.numNodes, 1)) if isinstance(self.parameters['psi_asym'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['psi_asym'], shape=(self.numNodes,1)) - self.q = numpy.array(self.parameters['q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['q'], shape=(self.numNodes,1)) - - #---------------------------------------- - - self.beta_pairwise_mode = self.parameters['beta_pairwise_mode'] - - #---------------------------------------- - # Global transmission parameters: - #---------------------------------------- - if(self.beta_pairwise_mode == 'infected' or self.beta_pairwise_mode is None): - self.beta_global = numpy.full_like(self.beta, fill_value=numpy.mean(self.beta)) - self.beta_Q_global = numpy.full_like(self.beta_Q, fill_value=numpy.mean(self.beta_Q)) - self.beta_asym_global = numpy.full_like(self.beta_asym, fill_value=numpy.mean(self.beta_asym)) - elif(self.beta_pairwise_mode == 'infectee'): - self.beta_global = self.beta - self.beta_Q_global = self.beta_Q - self.beta_asym_global = self.beta_asym - elif(self.beta_pairwise_mode == 'min'): - self.beta_global = numpy.minimum(self.beta, numpy.mean(beta)) - self.beta_Q_global = numpy.minimum(self.beta_Q, numpy.mean(beta_Q)) - self.beta_asym_global = numpy.minimum(self.beta_asym, numpy.mean(beta_asym)) - elif(self.beta_pairwise_mode == 'max'): - self.beta_global = numpy.maximum(self.beta, numpy.mean(beta)) - self.beta_Q_global = numpy.maximum(self.beta_Q, numpy.mean(beta_Q)) - self.beta_asym_global = numpy.maximum(self.beta_asym, numpy.mean(beta_asym)) - elif(self.beta_pairwise_mode == 'mean'): - self.beta_global = (self.beta + numpy.full_like(self.beta, fill_value=numpy.mean(self.beta)))/2 - self.beta_Q_global = (self.beta_Q + numpy.full_like(self.beta_Q, fill_value=numpy.mean(self.beta_Q)))/2 - self.beta_asym_global = (self.beta_asym + numpy.full_like(self.beta_asym, fill_value=numpy.mean(self.beta_asym)))/2 - - #---------------------------------------- - # Local transmission parameters: - #---------------------------------------- - self.beta_local = self.beta if self.parameters['beta_local'] is None else numpy.array(self.parameters['beta_local']) if isinstance(self.parameters['beta_local'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['beta_local'], shape=(self.numNodes,1)) - self.beta_Q_local = self.beta_Q if self.parameters['beta_Q_local'] is None else numpy.array(self.parameters['beta_Q_local']) if isinstance(self.parameters['beta_Q_local'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['beta_Q_local'], shape=(self.numNodes,1)) - self.beta_asym_local = None if self.parameters['beta_asym_local'] is None else numpy.array(self.parameters['beta_asym_local']) if isinstance(self.parameters['beta_asym_local'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['beta_asym_local'], shape=(self.numNodes,1)) - #---------------------------------------- - if(self.beta_local.ndim == 2 and self.beta_local.shape[0] == self.numNodes and self.beta_local.shape[1] == self.numNodes): - self.A_beta_pairwise = self.beta_local - elif((self.beta_local.ndim == 1 and self.beta_local.shape[0] == self.numNodes) or (self.beta_local.ndim == 2 and (self.beta_local.shape[0] == self.numNodes or self.beta_local.shape[1] == self.numNodes))): - self.beta_local = self.beta_local.reshape((self.numNodes,1)) - # Pre-multiply beta values by the adjacency matrix ("transmission weight connections") - A_beta_pairwise_byInfected = scipy.sparse.csr_matrix.multiply(self.A, self.beta_local.T).tocsr() - A_beta_pairwise_byInfectee = scipy.sparse.csr_matrix.multiply(self.A, self.beta_local).tocsr() - #------------------------------ - # Compute the effective pairwise beta values as a function of the infected/infectee pair: - if(self.beta_pairwise_mode == 'infected'): - self.A_beta_pairwise = A_beta_pairwise_byInfected - elif(self.beta_pairwise_mode == 'infectee'): - self.A_beta_pairwise = A_beta_pairwise_byInfectee - elif(self.beta_pairwise_mode == 'min'): - self.A_beta_pairwise = scipy.sparse.csr_matrix.minimum(A_beta_pairwise_byInfected, A_beta_pairwise_byInfectee) - elif(self.beta_pairwise_mode == 'max'): - self.A_beta_pairwise = scipy.sparse.csr_matrix.maximum(A_beta_pairwise_byInfected, A_beta_pairwise_byInfectee) - elif(self.beta_pairwise_mode == 'mean' or self.beta_pairwise_mode is None): - self.A_beta_pairwise = (A_beta_pairwise_byInfected + A_beta_pairwise_byInfectee)/2 - else: - print("Unrecognized beta_pairwise_mode value (support for 'infected', 'infectee', 'min', 'max', and 'mean').") - else: - print("Invalid values given for beta_local (expected 1xN list/array or NxN 2d array)") - #---------------------------------------- - if(self.beta_Q_local.ndim == 2 and self.beta_Q_local.shape[0] == self.numNodes and self.beta_Q_local.shape[1] == self.numNodes): - self.A_Q_beta_Q_pairwise = self.beta_Q_local - elif((self.beta_Q_local.ndim == 1 and self.beta_Q_local.shape[0] == self.numNodes) or (self.beta_Q_local.ndim == 2 and (self.beta_Q_local.shape[0] == self.numNodes or self.beta_Q_local.shape[1] == self.numNodes))): - self.beta_Q_local = self.beta_Q_local.reshape((self.numNodes,1)) - # Pre-multiply beta_Q values by the isolation adjacency matrix ("transmission weight connections") - A_Q_beta_Q_pairwise_byInfected = scipy.sparse.csr_matrix.multiply(self.A_Q, self.beta_Q_local.T).tocsr() - A_Q_beta_Q_pairwise_byInfectee = scipy.sparse.csr_matrix.multiply(self.A_Q, self.beta_Q_local).tocsr() - #------------------------------ - # Compute the effective pairwise beta values as a function of the infected/infectee pair: - if(self.beta_pairwise_mode == 'infected'): - self.A_Q_beta_Q_pairwise = A_Q_beta_Q_pairwise_byInfected - elif(self.beta_pairwise_mode == 'infectee'): - self.A_Q_beta_Q_pairwise = A_Q_beta_Q_pairwise_byInfectee - elif(self.beta_pairwise_mode == 'min'): - self.A_Q_beta_Q_pairwise = scipy.sparse.csr_matrix.minimum(A_Q_beta_Q_pairwise_byInfected, A_Q_beta_Q_pairwise_byInfectee) - elif(self.beta_pairwise_mode == 'max'): - self.A_Q_beta_Q_pairwise = scipy.sparse.csr_matrix.maximum(A_Q_beta_Q_pairwise_byInfected, A_Q_beta_Q_pairwise_byInfectee) - elif(self.beta_pairwise_mode == 'mean' or self.beta_pairwise_mode is None): - self.A_Q_beta_Q_pairwise = (A_Q_beta_Q_pairwise_byInfected + A_Q_beta_Q_pairwise_byInfectee)/2 - else: - print("Unrecognized beta_pairwise_mode value (support for 'infected', 'infectee', 'min', 'max', and 'mean').") - else: - print("Invalid values given for beta_Q_local (expected 1xN list/array or NxN 2d array)") - #---------------------------------------- - if(self.beta_asym_local is None): - self.A_beta_asym_pairwise = None - elif(self.beta_asym_local.ndim == 2 and self.beta_asym_local.shape[0] == self.numNodes and self.beta_asym_local.shape[1] == self.numNodes): - self.A_beta_asym_pairwise = self.beta_asym_local - elif((self.beta_asym_local.ndim == 1 and self.beta_asym_local.shape[0] == self.numNodes) or (self.beta_asym_local.ndim == 2 and (self.beta_asym_local.shape[0] == self.numNodes or self.beta_asym_local.shape[1] == self.numNodes))): - self.beta_asym_local = self.beta_asym_local.reshape((self.numNodes,1)) - # Pre-multiply beta_asym values by the adjacency matrix ("transmission weight connections") - A_beta_asym_pairwise_byInfected = scipy.sparse.csr_matrix.multiply(self.A, self.beta_asym_local.T).tocsr() - A_beta_asym_pairwise_byInfectee = scipy.sparse.csr_matrix.multiply(self.A, self.beta_asym_local).tocsr() - #------------------------------ - # Compute the effective pairwise beta values as a function of the infected/infectee pair: - if(self.beta_pairwise_mode == 'infected'): - self.A_beta_asym_pairwise = A_beta_asym_pairwise_byInfected - elif(self.beta_pairwise_mode == 'infectee'): - self.A_beta_asym_pairwise = A_beta_asym_pairwise_byInfectee - elif(self.beta_pairwise_mode == 'min'): - self.A_beta_asym_pairwise = scipy.sparse.csr_matrix.minimum(A_beta_asym_pairwise_byInfected, A_beta_asym_pairwise_byInfectee) - elif(self.beta_pairwise_mode == 'max'): - self.A_beta_asym_pairwise = scipy.sparse.csr_matrix.maximum(A_beta_asym_pairwise_byInfected, A_beta_asym_pairwise_byInfectee) - elif(self.beta_pairwise_mode == 'mean' or self.beta_pairwise_mode is None): - self.A_beta_asym_pairwise = (A_beta_asym_pairwise_byInfected + A_beta_asym_pairwise_byInfectee)/2 - else: - print("Unrecognized beta_pairwise_mode value (support for 'infected', 'infectee', 'min', 'max', and 'mean').") - else: - print("Invalid values given for beta_asym_local (expected 1xN list/array or NxN 2d array)") - - #---------------------------------------- - # Degree-based transmission scaling parameters: - #---------------------------------------- - self.delta_pairwise_mode = self.parameters['delta_pairwise_mode'] - with numpy.errstate(divide='ignore'): # ignore log(0) warning, then convert log(0) = -inf -> 0.0 - self.delta = numpy.log(self.degree)/numpy.log(numpy.mean(self.degree)) if self.parameters['delta'] is None else numpy.array(self.parameters['delta']) if isinstance(self.parameters['delta'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['delta'], shape=(self.numNodes,1)) - self.delta_Q = numpy.log(self.degree_Q)/numpy.log(numpy.mean(self.degree_Q)) if self.parameters['delta_Q'] is None else numpy.array(self.parameters['delta_Q']) if isinstance(self.parameters['delta_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['delta_Q'], shape=(self.numNodes,1)) - self.delta[numpy.isneginf(self.delta)] = 0.0 - self.delta_Q[numpy.isneginf(self.delta_Q)] = 0.0 - #---------------------------------------- - if(self.delta.ndim == 2 and self.delta.shape[0] == self.numNodes and self.delta.shape[1] == self.numNodes): - self.A_delta_pairwise = self.delta - elif((self.delta.ndim == 1 and self.delta.shape[0] == self.numNodes) or (self.delta.ndim == 2 and (self.delta.shape[0] == self.numNodes or self.delta.shape[1] == self.numNodes))): - self.delta = self.delta.reshape((self.numNodes,1)) - # Pre-multiply delta values by the adjacency matrix ("transmission weight connections") - A_delta_pairwise_byInfected = scipy.sparse.csr_matrix.multiply(self.A, self.delta.T).tocsr() - A_delta_pairwise_byInfectee = scipy.sparse.csr_matrix.multiply(self.A, self.delta).tocsr() - #------------------------------ - # Compute the effective pairwise delta values as a function of the infected/infectee pair: - if(self.delta_pairwise_mode == 'infected'): - self.A_delta_pairwise = A_delta_pairwise_byInfected - elif(self.delta_pairwise_mode == 'infectee'): - self.A_delta_pairwise = A_delta_pairwise_byInfectee - elif(self.delta_pairwise_mode == 'min'): - self.A_delta_pairwise = scipy.sparse.csr_matrix.minimum(A_delta_pairwise_byInfected, A_delta_pairwise_byInfectee) - elif(self.delta_pairwise_mode == 'max'): - self.A_delta_pairwise = scipy.sparse.csr_matrix.maximum(A_delta_pairwise_byInfected, A_delta_pairwise_byInfectee) - elif(self.delta_pairwise_mode == 'mean'): - self.A_delta_pairwise = (A_delta_pairwise_byInfected + A_delta_pairwise_byInfectee)/2 - elif(self.delta_pairwise_mode is None): - self.A_delta_pairwise = self.A - else: - print("Unrecognized delta_pairwise_mode value (support for 'infected', 'infectee', 'min', 'max', and 'mean').") - else: - print("Invalid values given for delta (expected 1xN list/array or NxN 2d array)") - #---------------------------------------- - if(self.delta_Q.ndim == 2 and self.delta_Q.shape[0] == self.numNodes and self.delta_Q.shape[1] == self.numNodes): - self.A_Q_delta_Q_pairwise = self.delta_Q - elif((self.delta_Q.ndim == 1 and self.delta_Q.shape[0] == self.numNodes) or (self.delta_Q.ndim == 2 and (self.delta_Q.shape[0] == self.numNodes or self.delta_Q.shape[1] == self.numNodes))): - self.delta_Q = self.delta_Q.reshape((self.numNodes,1)) - # Pre-multiply delta_Q values by the isolation adjacency matrix ("transmission weight connections") - A_Q_delta_Q_pairwise_byInfected = scipy.sparse.csr_matrix.multiply(self.A_Q, self.delta_Q).tocsr() - A_Q_delta_Q_pairwise_byInfectee = scipy.sparse.csr_matrix.multiply(self.A_Q, self.delta_Q.T).tocsr() - #------------------------------ - # Compute the effective pairwise delta values as a function of the infected/infectee pair: - if(self.delta_pairwise_mode == 'infected'): - self.A_Q_delta_Q_pairwise = A_Q_delta_Q_pairwise_byInfected - elif(self.delta_pairwise_mode == 'infectee'): - self.A_Q_delta_Q_pairwise = A_Q_delta_Q_pairwise_byInfectee - elif(self.delta_pairwise_mode == 'min'): - self.A_Q_delta_Q_pairwise = scipy.sparse.csr_matrix.minimum(A_Q_delta_Q_pairwise_byInfected, A_Q_delta_Q_pairwise_byInfectee) - elif(self.delta_pairwise_mode == 'max'): - self.A_Q_delta_Q_pairwise = scipy.sparse.csr_matrix.maximum(A_Q_delta_Q_pairwise_byInfected, A_Q_delta_Q_pairwise_byInfectee) - elif(self.delta_pairwise_mode == 'mean'): - self.A_Q_delta_Q_pairwise = (A_Q_delta_Q_pairwise_byInfected + A_Q_delta_Q_pairwise_byInfectee)/2 - elif(self.delta_pairwise_mode is None): - self.A_Q_delta_Q_pairwise = self.A - else: - print("Unrecognized delta_pairwise_mode value (support for 'infected', 'infectee', 'min', 'max', and 'mean').") - else: - print("Invalid values given for delta_Q (expected 1xN list/array or NxN 2d array)") - - #---------------------------------------- - # Pre-calculate the pairwise delta*beta values: - #---------------------------------------- - self.A_deltabeta = scipy.sparse.csr_matrix.multiply(self.A_delta_pairwise, self.A_beta_pairwise) - self.A_Q_deltabeta_Q = scipy.sparse.csr_matrix.multiply(self.A_Q_delta_Q_pairwise, self.A_Q_beta_Q_pairwise) - if(self.A_beta_asym_pairwise is not None): - self.A_deltabeta_asym = scipy.sparse.csr_matrix.multiply(self.A_delta_pairwise, self.A_beta_asym_pairwise) - else: - self.A_deltabeta_asym = None - - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def node_degrees(self, Amat): - return Amat.sum(axis=0).reshape(self.numNodes,1) # sums of adj matrix cols - - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def total_num_susceptible(self, t_idx=None): - if(t_idx is None): - return (self.numS[:] + self.numQ_S[:]) - else: - return (self.numS[t_idx] + self.numQ_S[t_idx]) - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def total_num_infected(self, t_idx=None): - if(t_idx is None): - return (self.numE[:] + self.numI_pre[:] + self.numI_sym[:] + self.numI_asym[:] + self.numH[:] - + self.numQ_E[:] + self.numQ_pre[:] + self.numQ_sym[:] + self.numQ_asym[:]) - else: - return (self.numE[t_idx] + self.numI_pre[t_idx] + self.numI_sym[t_idx] + self.numI_asym[t_idx] + self.numH[t_idx] - + self.numQ_E[t_idx] + self.numQ_pre[t_idx] + self.numQ_sym[t_idx] + self.numQ_asym[t_idx]) - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def total_num_isolated(self, t_idx=None): - if(t_idx is None): - return (self.numQ_S[:] + self.numQ_E[:] + self.numQ_pre[:] + self.numQ_sym[:] + self.numQ_asym[:] + self.numQ_R[:]) - else: - return (self.numQ_S[t_idx] + self.numQ_E[t_idx] + self.numQ_pre[t_idx] + self.numQ_sym[t_idx] + self.numQ_asym[t_idx] + self.numQ_R[t_idx]) - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def total_num_tested(self, t_idx=None): - if(t_idx is None): - return (self.numTested[:]) - else: - return (self.numTested[t_idx]) - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def total_num_positive(self, t_idx=None): - if(t_idx is None): - return (self.numPositive[:]) - else: - return (self.numPositive[t_idx]) - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def total_num_recovered(self, t_idx=None): - if(t_idx is None): - return (self.numR[:] + self.numQ_R[:]) - else: - return (self.numR[t_idx] + self.numQ_R[t_idx]) - - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def calc_propensities(self): - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Pre-calculate matrix multiplication terms that may be used in multiple propensity calculations, - # and check to see if their computation is necessary before doing the multiplication - #------------------------------------ - - self.transmissionTerms_I = numpy.zeros(shape=(self.numNodes,1)) - if(numpy.any(self.numI_sym[self.tidx]) or numpy.any(self.numI_asym[self.tidx]) or numpy.any(self.numI_pre[self.tidx])): - if(self.A_deltabeta_asym is not None): - self.transmissionTerms_sym = numpy.asarray(scipy.sparse.csr_matrix.dot(self.A_deltabeta, self.X==self.I_sym)) - self.transmissionTerms_asym = numpy.asarray(scipy.sparse.csr_matrix.dot(self.A_deltabeta_asym, ((self.X==self.I_pre)|(self.X==self.I_asym)))) - self.transmissionTerms_I = self.transmissionTerms_sym+self.transmissionTerms_asym - else: - self.transmissionTerms_I = numpy.asarray(scipy.sparse.csr_matrix.dot(self.A_deltabeta, ((self.X==self.I_sym)|(self.X==self.I_pre)|(self.X==self.I_asym)))) - - #------------------------------------ - - self.transmissionTerms_Q = numpy.zeros(shape=(self.numNodes,1)) - if(numpy.any(self.numQ_pre[self.tidx]) or numpy.any(self.numQ_sym[self.tidx]) or numpy.any(self.numQ_asym[self.tidx])): - self.transmissionTerms_Q = numpy.asarray(scipy.sparse.csr_matrix.dot(self.A_Q_deltabeta_Q, ((self.X==self.Q_pre)|(self.X==self.Q_sym)|(self.X==self.Q_asym)))) - - #------------------------------------ - - self.transmissionTerms_IQ = numpy.zeros(shape=(self.numNodes,1)) - if(numpy.any(self.numQ_S[self.tidx]) and (numpy.any(self.numI_sym[self.tidx]) or numpy.any(self.numI_asym[self.tidx]) or numpy.any(self.numI_pre[self.tidx]))): - self.transmissionTerms_IQ = numpy.asarray(scipy.sparse.csr_matrix.dot(self.A_Q_deltabeta_Q, ((self.X==self.I_sym)|(self.X==self.I_pre)|(self.X==self.I_asym)))) - - #------------------------------------ - - numContacts_Q = numpy.zeros(shape=(self.numNodes,1)) - if(numpy.any(self.positive) and (numpy.any(self.phi_S) or numpy.any(self.phi_E) or numpy.any(self.phi_pre) or numpy.any(self.phi_sym) or numpy.any(self.phi_asym))): - numContacts_Q = numpy.asarray(scipy.sparse.csr_matrix.dot(self.A, ((self.positive)&(self.X!=self.R)&(self.X!=self.Q_R)&(self.X!=self.F)))) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - propensities_StoE = ( self.alpha * - (self.o*(self.beta_global*self.prevalence_ext) - + (1-self.o)*( - self.p*((self.beta_global*self.numI_sym[self.tidx] + self.beta_asym_global*(self.numI_pre[self.tidx] + self.numI_asym[self.tidx]) - + self.q*self.beta_Q_global*(self.numQ_pre[self.tidx] + self.numQ_sym[self.tidx] + self.numQ_asym[self.tidx]))/self.N[self.tidx]) - + (1-self.p)*(numpy.divide(self.transmissionTerms_I, self.degree, out=numpy.zeros_like(self.degree), where=self.degree!=0) - + numpy.divide(self.transmissionTerms_Q, self.degree_Q, out=numpy.zeros_like(self.degree_Q), where=self.degree_Q!=0)))) - )*(self.X==self.S) - - propensities_QStoQE = numpy.zeros_like(propensities_StoE) - if(numpy.any(self.X==self.Q_S)): - propensities_QStoQE = ( self.alpha_Q * - (self.o*(self.q*self.beta_global*self.prevalence_ext) - + (1-self.o)*( - self.p*(self.q*(self.beta_global*self.numI_sym[self.tidx] + self.beta_asym_global*(self.numI_pre[self.tidx] + self.numI_asym[self.tidx]) - + self.beta_Q_global*(self.numQ_pre[self.tidx] + self.numQ_sym[self.tidx] + self.numQ_asym[self.tidx]))/self.N[self.tidx]) - + (1-self.p)*(numpy.divide(self.transmissionTerms_IQ+self.transmissionTerms_Q, self.degree_Q, out=numpy.zeros_like(self.degree_Q), where=self.degree_Q!=0)))) - )*(self.X==self.Q_S) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - if(self.transition_mode == 'time_in_state'): - - propensities_EtoIPRE = 1e5 * ((self.X==self.E) & numpy.greater(self.timer_state, 1/self.sigma)) - - propensities_IPREtoISYM = 1e5 * ((self.X==self.I_pre) & numpy.greater(self.timer_state, 1/self.lamda) & numpy.greater_equal(self.rand_a, self.a)) - - propensities_IPREtoIASYM = 1e5 * ((self.X==self.I_pre) & numpy.greater(self.timer_state, 1/self.lamda) & numpy.less(self.rand_a, self.a)) - - propensities_ISYMtoR = 1e5 * ((self.X==self.I_sym) & numpy.greater(self.timer_state, 1/self.gamma) & numpy.greater_equal(self.rand_h, self.h)) - - propensities_ISYMtoH = 1e5 * ((self.X==self.I_sym) & numpy.greater(self.timer_state, 1/self.eta) & numpy.less(self.rand_h, self.h)) - - propensities_IASYMtoR = 1e5 * ((self.X==self.I_asym) & numpy.greater(self.timer_state, 1/self.gamma)) - - propensities_HtoR = 1e5 * ((self.X==self.H) & numpy.greater(self.timer_state, 1/self.gamma_H) & numpy.greater_equal(self.rand_f, self.f)) - - propensities_HtoF = 1e5 * ((self.X==self.H) & numpy.greater(self.timer_state, 1/self.mu_H) & numpy.less(self.rand_f, self.f)) - - propensities_StoQS = numpy.zeros_like(propensities_StoE) - - propensities_EtoQE = numpy.zeros_like(propensities_StoE) - - propensities_IPREtoQPRE = numpy.zeros_like(propensities_StoE) - - propensities_ISYMtoQSYM = numpy.zeros_like(propensities_StoE) - - propensities_IASYMtoQASYM = numpy.zeros_like(propensities_StoE) - - propensities_QEtoQPRE = 1e5 * ((self.X==self.Q_E) & numpy.greater(self.timer_state, 1/self.sigma_Q)) - - propensities_QPREtoQSYM = 1e5 * ((self.X==self.Q_pre) & numpy.greater(self.timer_state, 1/self.lamda_Q) & numpy.greater_equal(self.rand_a, self.a)) - - propensities_QPREtoQASYM = 1e5 * ((self.X==self.Q_pre) & numpy.greater(self.timer_state, 1/self.lamda_Q) & numpy.less(self.rand_a, self.a)) - - propensities_QSYMtoQR = 1e5 * ((self.X==self.Q_sym) & numpy.greater(self.timer_state, 1/self.gamma_Q_sym) & numpy.greater_equal(self.rand_h, self.h)) - - propensities_QSYMtoH = 1e5 * ((self.X==self.Q_sym) & numpy.greater(self.timer_state, 1/self.eta_Q) & numpy.less(self.rand_h, self.h)) - - propensities_QASYMtoQR = 1e5 * ((self.X==self.Q_asym) & numpy.greater(self.timer_state, 1/self.gamma_Q_asym)) - - propensities_RtoS = 1e5 * ((self.X==self.R) & numpy.greater(self.timer_state, 1/self.xi)) - - propensities__toS = 1e5 * ((self.X!=self.F) & numpy.greater(self.timer_state, 1/self.nu)) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - else: # exponential_rates - - propensities_EtoIPRE = self.sigma * (self.X==self.E) - - propensities_IPREtoISYM = self.lamda * ((self.X==self.I_pre) & (numpy.greater_equal(self.rand_a, self.a))) - - propensities_IPREtoIASYM = self.lamda * ((self.X==self.I_pre) & (numpy.less(self.rand_a, self.a))) - - propensities_ISYMtoR = self.gamma * ((self.X==self.I_sym) & (numpy.greater_equal(self.rand_h, self.h))) - - propensities_ISYMtoH = self.eta * ((self.X==self.I_sym) & (numpy.less(self.rand_h, self.h))) - - propensities_IASYMtoR = self.gamma_asym * (self.X==self.I_asym) - - propensities_HtoR = self.gamma_H * ((self.X==self.H) & (numpy.greater_equal(self.rand_f, self.f))) - - propensities_HtoF = self.mu_H * ((self.X==self.H) & (numpy.less(self.rand_f, self.f))) - - propensities_StoQS = (self.theta_S + self.phi_S*numContacts_Q)*self.psi_S * (self.X==self.S) - - propensities_EtoQE = (self.theta_E + self.phi_E*numContacts_Q)*self.psi_E * (self.X==self.E) - - propensities_IPREtoQPRE = (self.theta_pre + self.phi_pre*numContacts_Q)*self.psi_pre * (self.X==self.I_pre) - - propensities_ISYMtoQSYM = (self.theta_sym + self.phi_sym*numContacts_Q)*self.psi_sym * (self.X==self.I_sym) - - propensities_IASYMtoQASYM = (self.theta_asym + self.phi_asym*numContacts_Q)*self.psi_asym * (self.X==self.I_asym) - - propensities_QEtoQPRE = self.sigma_Q * (self.X==self.Q_E) - - propensities_QPREtoQSYM = self.lamda_Q * ((self.X==self.Q_pre) & (numpy.greater_equal(self.rand_a, self.a))) - - propensities_QPREtoQASYM = self.lamda_Q * ((self.X==self.Q_pre) & (numpy.less(self.rand_a, self.a))) - - propensities_QSYMtoQR = self.gamma_Q_sym * ((self.X==self.Q_sym) & (numpy.greater_equal(self.rand_h, self.h))) - - propensities_QSYMtoH = self.eta_Q * ((self.X==self.Q_sym) & (numpy.less(self.rand_h, self.h))) - - propensities_QASYMtoQR = self.gamma_Q_asym * (self.X==self.Q_asym) - - propensities_RtoS = self.xi * (self.X==self.R) - - propensities__toS = self.nu * (self.X!=self.F) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - propensities = numpy.hstack([propensities_StoE, propensities_EtoIPRE, propensities_IPREtoISYM, propensities_IPREtoIASYM, - propensities_ISYMtoR, propensities_ISYMtoH, propensities_IASYMtoR, propensities_HtoR, propensities_HtoF, - propensities_StoQS, propensities_EtoQE, propensities_IPREtoQPRE, propensities_ISYMtoQSYM, propensities_IASYMtoQASYM, - propensities_QStoQE, propensities_QEtoQPRE, propensities_QPREtoQSYM, propensities_QPREtoQASYM, - propensities_QSYMtoQR, propensities_QSYMtoH, propensities_QASYMtoQR, propensities_RtoS, propensities__toS]) - - columns = [ 'StoE', 'EtoIPRE', 'IPREtoISYM', 'IPREtoIASYM', - 'ISYMtoR', 'ISYMtoH', 'IASYMtoR', 'HtoR', 'HtoF', - 'StoQS', 'EtoQE', 'IPREtoQPRE', 'ISYMtoQSYM', 'IASYMtoQASYM', - 'QStoQE', 'QEtoQPRE', 'QPREtoQSYM', 'QPREtoQASYM', - 'QSYMtoQR', 'QSYMtoH', 'QASYMtoQR', 'RtoS', '_toS' ] - - return propensities, columns - - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def set_isolation(self, node, isolate): - # Move this node in/out of the appropriate isolation state: - if(isolate == True): - if(self.X[node] == self.S): - self.X[node] = self.Q_S - elif(self.X[node] == self.E): - self.X[node] = self.Q_E - elif(self.X[node] == self.I_pre): - self.X[node] = self.Q_pre - elif(self.X[node] == self.I_sym): - self.X[node] = self.Q_sym - elif(self.X[node] == self.I_asym): - self.X[node] = self.Q_asym - elif(self.X[node] == self.R): - self.X[node] = self.Q_R - elif(isolate == False): - if(self.X[node] == self.Q_S): - self.X[node] = self.S - elif(self.X[node] == self.Q_E): - self.X[node] = self.E - elif(self.X[node] == self.Q_pre): - self.X[node] = self.I_pre - elif(self.X[node] == self.Q_sym): - self.X[node] = self.I_sym - elif(self.X[node] == self.Q_asym): - self.X[node] = self.I_asym - elif(self.X[node] == self.Q_R): - self.X[node] = self.R - # Reset the isolation timer: - self.timer_isolation[node] = 0 - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def set_tested(self, node, tested): - self.tested[node] = tested - self.testedInCurrentState[node] = tested - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def set_positive(self, node, positive): - self.positive[node] = positive - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def introduce_exposures(self, num_new_exposures): - exposedNodes = numpy.random.choice(range(self.numNodes), size=num_new_exposures, replace=False) - for exposedNode in exposedNodes: - if(self.X[exposedNode]==self.S): - self.X[exposedNode] = self.E - elif(self.X[exposedNode]==self.Q_S): - self.X[exposedNode] = self.Q_E - - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def increase_data_series_length(self): - self.tseries = numpy.pad(self.tseries, [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.numS = numpy.pad(self.numS, [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.numE = numpy.pad(self.numE, [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.numI_pre = numpy.pad(self.numI_pre, [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.numI_sym = numpy.pad(self.numI_sym, [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.numI_asym = numpy.pad(self.numI_asym, [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.numH = numpy.pad(self.numH, [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.numR = numpy.pad(self.numR, [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.numF = numpy.pad(self.numF, [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.numQ_S = numpy.pad(self.numQ_S, [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.numQ_E = numpy.pad(self.numQ_E, [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.numQ_pre = numpy.pad(self.numQ_pre, [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.numQ_sym = numpy.pad(self.numQ_sym, [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.numQ_asym = numpy.pad(self.numQ_asym, [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.numQ_R = numpy.pad(self.numQ_R, [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.N = numpy.pad(self.N, [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.numTested = numpy.pad(self.numTested, [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.numPositive = numpy.pad(self.numPositive, [(0, 6*self.numNodes)], mode='constant', constant_values=0) - - if(self.store_Xseries): - self.Xseries = numpy.pad(self.Xseries, [(0, 6*self.numNodes), (0,0)], mode='constant', constant_values=0) - - if(self.nodeGroupData): - for groupName in self.nodeGroupData: - self.nodeGroupData[groupName]['numS'] = numpy.pad(self.nodeGroupData[groupName]['numS'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.nodeGroupData[groupName]['numE'] = numpy.pad(self.nodeGroupData[groupName]['numE'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.nodeGroupData[groupName]['numI_pre'] = numpy.pad(self.nodeGroupData[groupName]['numI_pre'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.nodeGroupData[groupName]['numI_sym'] = numpy.pad(self.nodeGroupData[groupName]['numI_sym'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.nodeGroupData[groupName]['numI_asym'] = numpy.pad(self.nodeGroupData[groupName]['numI_asym'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.nodeGroupData[groupName]['numH'] = numpy.pad(self.nodeGroupData[groupName]['numH'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.nodeGroupData[groupName]['numR'] = numpy.pad(self.nodeGroupData[groupName]['numR'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.nodeGroupData[groupName]['numF'] = numpy.pad(self.nodeGroupData[groupName]['numF'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.nodeGroupData[groupName]['numQ_S'] = numpy.pad(self.nodeGroupData[groupName]['numQ_S'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.nodeGroupData[groupName]['numQ_E'] = numpy.pad(self.nodeGroupData[groupName]['numQ_E'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.nodeGroupData[groupName]['numQ_pre'] = numpy.pad(self.nodeGroupData[groupName]['numQ_pre'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.nodeGroupData[groupName]['numQ_sym'] = numpy.pad(self.nodeGroupData[groupName]['numQ_sym'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.nodeGroupData[groupName]['numQ_asym'] = numpy.pad(self.nodeGroupData[groupName]['numQ_asym'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.nodeGroupData[groupName]['numQ_R'] = numpy.pad(self.nodeGroupData[groupName]['numQ_R'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.nodeGroupData[groupName]['N'] = numpy.pad(self.nodeGroupData[groupName]['N'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.nodeGroupData[groupName]['numTested'] = numpy.pad(self.nodeGroupData[groupName]['numTested'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) - self.nodeGroupData[groupName]['numPositive'] = numpy.pad(self.nodeGroupData[groupName]['numPositive'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) - - return None - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def finalize_data_series(self): - self.tseries = numpy.array(self.tseries, dtype=float)[:self.tidx+1] - self.numS = numpy.array(self.numS, dtype=float)[:self.tidx+1] - self.numE = numpy.array(self.numE, dtype=float)[:self.tidx+1] - self.numI_pre = numpy.array(self.numI_pre, dtype=float)[:self.tidx+1] - self.numI_sym = numpy.array(self.numI_sym, dtype=float)[:self.tidx+1] - self.numI_asym = numpy.array(self.numI_asym, dtype=float)[:self.tidx+1] - self.numH = numpy.array(self.numH, dtype=float)[:self.tidx+1] - self.numR = numpy.array(self.numR, dtype=float)[:self.tidx+1] - self.numF = numpy.array(self.numF, dtype=float)[:self.tidx+1] - self.numQ_S = numpy.array(self.numQ_S, dtype=float)[:self.tidx+1] - self.numQ_E = numpy.array(self.numQ_E, dtype=float)[:self.tidx+1] - self.numQ_pre = numpy.array(self.numQ_pre, dtype=float)[:self.tidx+1] - self.numQ_sym = numpy.array(self.numQ_sym, dtype=float)[:self.tidx+1] - self.numQ_asym = numpy.array(self.numQ_asym, dtype=float)[:self.tidx+1] - self.numQ_R = numpy.array(self.numQ_R, dtype=float)[:self.tidx+1] - self.N = numpy.array(self.N, dtype=float)[:self.tidx+1] - self.numTested = numpy.array(self.numTested, dtype=float)[:self.tidx+1] - self.numPositive = numpy.array(self.numPositive, dtype=float)[:self.tidx+1] - - if(self.store_Xseries): - self.Xseries = self.Xseries[:self.tidx+1, :] - - if(self.nodeGroupData): - for groupName in self.nodeGroupData: - self.nodeGroupData[groupName]['numS'] = numpy.array(self.nodeGroupData[groupName]['numS'], dtype=float)[:self.tidx+1] - self.nodeGroupData[groupName]['numE'] = numpy.array(self.nodeGroupData[groupName]['numE'], dtype=float)[:self.tidx+1] - self.nodeGroupData[groupName]['numI_pre'] = numpy.array(self.nodeGroupData[groupName]['numI_pre'], dtype=float)[:self.tidx+1] - self.nodeGroupData[groupName]['numI_sym'] = numpy.array(self.nodeGroupData[groupName]['numI_sym'], dtype=float)[:self.tidx+1] - self.nodeGroupData[groupName]['numI_asym'] = numpy.array(self.nodeGroupData[groupName]['numI_asym'], dtype=float)[:self.tidx+1] - self.nodeGroupData[groupName]['numR'] = numpy.array(self.nodeGroupData[groupName]['numR'], dtype=float)[:self.tidx+1] - self.nodeGroupData[groupName]['numF'] = numpy.array(self.nodeGroupData[groupName]['numF'], dtype=float)[:self.tidx+1] - self.nodeGroupData[groupName]['numQ_S'] = numpy.array(self.nodeGroupData[groupName]['numQ_S'], dtype=float)[:self.tidx+1] - self.nodeGroupData[groupName]['numQ_E'] = numpy.array(self.nodeGroupData[groupName]['numQ_E'], dtype=float)[:self.tidx+1] - self.nodeGroupData[groupName]['numQ_pre'] = numpy.array(self.nodeGroupData[groupName]['numQ_pre'], dtype=float)[:self.tidx+1] - self.nodeGroupData[groupName]['numQ_sym'] = numpy.array(self.nodeGroupData[groupName]['numQ_sym'], dtype=float)[:self.tidx+1] - self.nodeGroupData[groupName]['numQ_asym'] = numpy.array(self.nodeGroupData[groupName]['numQ_asym'], dtype=float)[:self.tidx+1] - self.nodeGroupData[groupName]['numQ_R'] = numpy.array(self.nodeGroupData[groupName]['numQ_R'], dtype=float)[:self.tidx+1] - self.nodeGroupData[groupName]['N'] = numpy.array(self.nodeGroupData[groupName]['N'], dtype=float)[:self.tidx+1] - self.nodeGroupData[groupName]['numTested'] = numpy.array(self.nodeGroupData[groupName]['numTested'], dtype=float)[:self.tidx+1] - self.nodeGroupData[groupName]['numPositive'] = numpy.array(self.nodeGroupData[groupName]['numPositive'], dtype=float)[:self.tidx+1] - - return None - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def run_iteration(self, max_dt=None): - - max_dt = self.tmax if max_dt is None else max_dt - - if(self.tidx >= len(self.tseries)-1): - # Room has run out in the timeseries storage arrays; double the size of these arrays: - self.increase_data_series_length() - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Generate 2 random numbers uniformly distributed in (0,1) - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - r1 = numpy.random.rand() - r2 = numpy.random.rand() - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Calculate propensities - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - propensities, transitionTypes = self.calc_propensities() - - if(propensities.sum() > 0): - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Calculate alpha - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - propensities_flat = propensities.ravel(order='F') - cumsum = propensities_flat.cumsum() - alpha = propensities_flat.sum() - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Compute the time until the next event takes place - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - tau = (1/alpha)*numpy.log(float(1/r1)) - - if(tau > max_dt): - # If the time to next event exceeds the max allowed interval, - # advance the system time by the max allowed interval, - # but do not execute any events (recalculate Gillespie interval/event next iteration) - self.t += max_dt - self.timer_state += max_dt - # Update testing and isolation timers/statuses - isolatedNodes = numpy.argwhere((self.X==self.Q_S)|(self.X==self.Q_E)|(self.X==self.Q_pre)|(self.X==self.Q_sym)|(self.X==self.Q_asym)|(self.X==self.Q_R))[:,0].flatten() - self.timer_isolation[isolatedNodes] = self.timer_isolation[isolatedNodes] + max_dt - nodesExitingIsolation = numpy.argwhere(self.timer_isolation >= self.isolationTime) - for isoNode in nodesExitingIsolation: - self.set_isolation(node=isoNode, isolate=False) - # return without any further event execution - return True - else: - self.t += tau - self.timer_state += tau - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Compute which event takes place - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - transitionIdx = numpy.searchsorted(cumsum,r2*alpha) - transitionNode = transitionIdx % self.numNodes - transitionType = transitionTypes[ int(transitionIdx/self.numNodes) ] - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Perform updates triggered by rate propensities: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - assert(self.X[transitionNode] == self.transitions[transitionType]['currentState'] and self.X[transitionNode]!=self.F), "Assertion error: Node "+str(transitionNode)+" has unexpected current state "+str(self.X[transitionNode])+" given the intended transition of "+str(transitionType)+"." - self.X[transitionNode] = self.transitions[transitionType]['newState'] - - self.testedInCurrentState[transitionNode] = False - - self.timer_state[transitionNode] = 0.0 - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - # Save information about infection events when they occur: - if(transitionType == 'StoE' or transitionType == 'QStoQE'): - transitionNode_GNbrs = list(self.G[transitionNode].keys()) - transitionNode_GQNbrs = list(self.G_Q[transitionNode].keys()) - self.infectionsLog.append({ 't': self.t, - 'infected_node': transitionNode, - 'infection_type': transitionType, - 'infected_node_degree': self.degree[transitionNode], - 'local_contact_nodes': transitionNode_GNbrs, - 'local_contact_node_states': self.X[transitionNode_GNbrs].flatten(), - 'isolation_contact_nodes': transitionNode_GQNbrs, - 'isolation_contact_node_states':self.X[transitionNode_GQNbrs].flatten() }) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - if(transitionType in ['EtoQE', 'IPREtoQPRE', 'ISYMtoQSYM', 'IASYMtoQASYM', 'ISYMtoH']): - self.set_positive(node=transitionNode, positive=True) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - else: - - tau = 0.01 - self.t += tau - self.timer_state += tau - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - self.tidx += 1 - - self.tseries[self.tidx] = self.t - self.numS[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.S), a_min=0, a_max=self.numNodes) - self.numE[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.E), a_min=0, a_max=self.numNodes) - self.numI_pre[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.I_pre), a_min=0, a_max=self.numNodes) - self.numI_sym[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.I_sym), a_min=0, a_max=self.numNodes) - self.numI_asym[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.I_asym), a_min=0, a_max=self.numNodes) - self.numH[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.H), a_min=0, a_max=self.numNodes) - self.numR[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.R), a_min=0, a_max=self.numNodes) - self.numF[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.F), a_min=0, a_max=self.numNodes) - self.numQ_S[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.Q_S), a_min=0, a_max=self.numNodes) - self.numQ_E[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.Q_E), a_min=0, a_max=self.numNodes) - self.numQ_pre[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.Q_pre), a_min=0, a_max=self.numNodes) - self.numQ_sym[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.Q_sym), a_min=0, a_max=self.numNodes) - self.numQ_asym[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.Q_asym), a_min=0, a_max=self.numNodes) - self.numQ_R[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.Q_R), a_min=0, a_max=self.numNodes) - self.numTested[self.tidx] = numpy.clip(numpy.count_nonzero(self.tested), a_min=0, a_max=self.numNodes) - self.numPositive[self.tidx] = numpy.clip(numpy.count_nonzero(self.positive), a_min=0, a_max=self.numNodes) - - self.N[self.tidx] = numpy.clip((self.numNodes - self.numF[self.tidx]), a_min=0, a_max=self.numNodes) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Update testing and isolation statuses - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - isolatedNodes = numpy.argwhere((self.X==self.Q_S)|(self.X==self.Q_E)|(self.X==self.Q_pre)|(self.X==self.Q_sym)|(self.X==self.Q_asym)|(self.X==self.Q_R))[:,0].flatten() - self.timer_isolation[isolatedNodes] = self.timer_isolation[isolatedNodes] + tau - - nodesExitingIsolation = numpy.argwhere(self.timer_isolation >= self.isolationTime) - for isoNode in nodesExitingIsolation: - self.set_isolation(node=isoNode, isolate=False) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Store system states - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(self.store_Xseries): - self.Xseries[self.tidx,:] = self.X.T - - if(self.nodeGroupData): - for groupName in self.nodeGroupData: - self.nodeGroupData[groupName]['numS'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.S) - self.nodeGroupData[groupName]['numE'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.E) - self.nodeGroupData[groupName]['numI_pre'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.I_pre) - self.nodeGroupData[groupName]['numI_sym'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.I_sym) - self.nodeGroupData[groupName]['numI_asym'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.I_asym) - self.nodeGroupData[groupName]['numH'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.H) - self.nodeGroupData[groupName]['numR'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.R) - self.nodeGroupData[groupName]['numF'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.F) - self.nodeGroupData[groupName]['numQ_S'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_S) - self.nodeGroupData[groupName]['numQ_E'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_E) - self.nodeGroupData[groupName]['numQ_pre'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_pre) - self.nodeGroupData[groupName]['numQ_sym'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_sym) - self.nodeGroupData[groupName]['numQ_asym'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_asym) - self.nodeGroupData[groupName]['numQ_R'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_R) - self.nodeGroupData[groupName]['N'][self.tidx] = numpy.clip((self.nodeGroupData[groupName]['numS'][0] + self.nodeGroupData[groupName]['numE'][0] + self.nodeGroupData[groupName]['numI'][0] + self.nodeGroupData[groupName]['numQ_E'][0] + self.nodeGroupData[groupName]['numQ_I'][0] + self.nodeGroupData[groupName]['numR'][0]), a_min=0, a_max=self.numNodes) - self.nodeGroupData[groupName]['numTested'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.tested) - self.nodeGroupData[groupName]['numPositive'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.positive) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Terminate if tmax reached or num infections is 0: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(self.t >= self.tmax or (self.total_num_infected(self.tidx) < 1 and self.total_num_isolated(self.tidx) < 1)): - self.finalize_data_series() - return False - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - return True - - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def run(self, T, checkpoints=None, max_dt=None, min_dt=None, print_interval=10, verbose='t'): - if(T>0): - self.tmax += T - else: - return False - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Pre-process checkpoint values: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(checkpoints): - numCheckpoints = len(checkpoints['t']) - for chkpt_param, chkpt_values in checkpoints.items(): - assert(isinstance(chkpt_values, (list, numpy.ndarray)) and len(chkpt_values)==numCheckpoints), "Expecting a list of values with length equal to number of checkpoint times ("+str(numCheckpoints)+") for each checkpoint parameter." - checkpointIdx = numpy.searchsorted(checkpoints['t'], self.t) # Finds 1st index in list greater than given val - if(checkpointIdx >= numCheckpoints): - # We are out of checkpoints, stop checking them: - checkpoints = None - else: - checkpointTime = checkpoints['t'][checkpointIdx] - - #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% - # Run the simulation loop: - #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% - print_reset = True - running = True - while running: - - running = self.run_iteration() - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Handle checkpoints if applicable: - if(checkpoints): - if(self.t >= checkpointTime): - if(verbose is not False): - print("[Checkpoint: Updating parameters]") - # A checkpoint has been reached, update param values: - for param in list(self.parameters.keys()): - if(param in list(checkpoints.keys())): - self.parameters.update({param: checkpoints[param][checkpointIdx]}) - # Update parameter data structures and scenario flags: - self.update_parameters() - # Update the next checkpoint time: - checkpointIdx = numpy.searchsorted(checkpoints['t'], self.t) # Finds 1st index in list greater than given val - if(checkpointIdx >= numCheckpoints): - # We are out of checkpoints, stop checking them: - checkpoints = None - else: - checkpointTime = checkpoints['t'][checkpointIdx] - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - - if(print_interval): - if(print_reset and (int(self.t) % print_interval == 0)): - if(verbose=="t"): - print("t = %.2f" % self.t) - if(verbose==True): - print("t = %.2f" % self.t) - print("\t S = " + str(self.numS[self.tidx])) - print("\t E = " + str(self.numE[self.tidx])) - print("\t I_pre = " + str(self.numI_pre[self.tidx])) - print("\t I_sym = " + str(self.numI_sym[self.tidx])) - print("\t I_asym = " + str(self.numI_asym[self.tidx])) - print("\t H = " + str(self.numH[self.tidx])) - print("\t R = " + str(self.numR[self.tidx])) - print("\t F = " + str(self.numF[self.tidx])) - print("\t Q_S = " + str(self.numQ_S[self.tidx])) - print("\t Q_E = " + str(self.numQ_E[self.tidx])) - print("\t Q_pre = " + str(self.numQ_pre[self.tidx])) - print("\t Q_sym = " + str(self.numQ_sym[self.tidx])) - print("\t Q_asym = " + str(self.numQ_asym[self.tidx])) - print("\t Q_R = " + str(self.numQ_R[self.tidx])) - - print_reset = False - elif(not print_reset and (int(self.t) % 10 != 0)): - print_reset = True - - return True - - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def plot(self, ax=None, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_sym='line', plot_I_asym='line', - plot_H='line', plot_R='line', plot_F='line', - plot_Q_E='line', plot_Q_pre='line', plot_Q_sym='line', plot_Q_asym='line', - plot_Q_S='line', plot_Q_R='line', combine_Q_infected=True, - color_S='tab:green', color_E='orange', color_I_pre='tomato', color_I_sym='crimson', color_I_asym='#F0909B', - color_H='violet', color_R='tab:blue', color_F='black', - color_Q_E='orange', color_Q_pre='tomato', color_Q_sym='crimson', color_Q_asym='#F0909B', - color_Q_S='tab:green', color_Q_R='tab:blue', color_Q_infected='tab:purple', - color_reference='#E0E0E0', - dashed_reference_results=None, dashed_reference_label='reference', - shaded_reference_results=None, shaded_reference_label='reference', - vlines=[], vline_colors=[], vline_styles=[], vline_labels=[], - ylim=None, xlim=None, legend=True, title=None, side_title=None, plot_percentages=True): - - import matplotlib.pyplot as pyplot - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Create an Axes object if None provided: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(not ax): - fig, ax = pyplot.subplots() - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Prepare data series to be plotted: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - Sseries = self.numS/self.numNodes if plot_percentages else self.numS - Eseries = self.numE/self.numNodes if plot_percentages else self.numE - I_preseries = self.numI_pre/self.numNodes if plot_percentages else self.numI_pre - I_symseries = self.numI_sym/self.numNodes if plot_percentages else self.numI_sym - I_asymseries = self.numI_asym/self.numNodes if plot_percentages else self.numI_asym - Rseries = self.numR/self.numNodes if plot_percentages else self.numR - Hseries = self.numH/self.numNodes if plot_percentages else self.numH - Fseries = self.numF/self.numNodes if plot_percentages else self.numF - Q_Sseries = self.numQ_S/self.numNodes if plot_percentages else self.numQ_S - Q_Eseries = self.numQ_E/self.numNodes if plot_percentages else self.numQ_E - Q_preseries = self.numQ_pre/self.numNodes if plot_percentages else self.numQ_pre - Q_asymseries = self.numQ_asym/self.numNodes if plot_percentages else self.numQ_asym - Q_symseries = self.numQ_sym/self.numNodes if plot_percentages else self.numQ_sym - Q_Rseries = self.numQ_R/self.numNodes if plot_percentages else self.numQ_R - Q_infectedseries = (Q_Eseries + Q_preseries + Q_asymseries + Q_symseries) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Draw the reference data: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(dashed_reference_results): - dashedReference_tseries = dashed_reference_results.tseries[::int(self.numNodes/100)] - dashedReference_infectedStack = dashed_reference_results.total_num_infected()[::int(self.numNodes/100)] / (self.numNodes if plot_percentages else 1) - ax.plot(dashedReference_tseries, dashedReference_infectedStack, color='#E0E0E0', linestyle='--', label='Total infections ('+dashed_reference_label+')', zorder=0) - if(shaded_reference_results): - shadedReference_tseries = shaded_reference_results.tseries - shadedReference_infectedStack = shaded_reference_results.total_num_infected() / (self.numNodes if plot_percentages else 1) - ax.fill_between(shaded_reference_results.tseries, shadedReference_infectedStack, 0, color='#EFEFEF', label='Total infections ('+shaded_reference_label+')', zorder=0) - ax.plot(shaded_reference_results.tseries, shadedReference_infectedStack, color='#E0E0E0', zorder=1) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Draw the stacked variables: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - topstack = numpy.zeros_like(self.tseries) - if(any(Fseries) and plot_F=='stacked'): - ax.fill_between(numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, topstack+Fseries), topstack, color=color_F, alpha=0.75, label='$F$', zorder=2) - ax.plot( numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, topstack+Fseries), color=color_F, zorder=3) - topstack = topstack+Fseries - - if(any(Hseries) and plot_H=='stacked'): - ax.fill_between(numpy.ma.masked_where(Hseries<=0, self.tseries), numpy.ma.masked_where(Hseries<=0, topstack+Hseries), topstack, color=color_H, alpha=0.75, label='$H$', zorder=2) - ax.plot( numpy.ma.masked_where(Hseries<=0, self.tseries), numpy.ma.masked_where(Hseries<=0, topstack+Hseries), color=color_H, zorder=3) - topstack = topstack+Hseries - - if(combine_Q_infected and any(Q_infectedseries) and plot_Q_E=='stacked' and plot_Q_pre=='stacked' and plot_Q_sym=='stacked' and plot_Q_asym=='stacked'): - ax.fill_between(numpy.ma.masked_where(Q_infectedseries<=0, self.tseries), numpy.ma.masked_where(Q_infectedseries<=0, topstack+Q_infectedseries), topstack, facecolor=color_Q_infected, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_{infected}$', zorder=2) - ax.plot( numpy.ma.masked_where(Q_infectedseries<=0, self.tseries), numpy.ma.masked_where(Q_infectedseries<=0, topstack+Q_infectedseries), color=color_Q_infected, zorder=3) - topstack = topstack+Q_infectedseries - - if(not combine_Q_infected and any(Q_Eseries) and plot_Q_E=='stacked'): - ax.fill_between(numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, topstack+Q_Eseries), topstack, facecolor=color_Q_E, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_E$', zorder=2) - ax.plot( numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, topstack+Q_Eseries), color=color_Q_E, zorder=3) - topstack = topstack+Q_Eseries - - if(any(Eseries) and plot_E=='stacked'): - ax.fill_between(numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, topstack+Eseries), topstack, color=color_E, alpha=0.75, label='$E$', zorder=2) - ax.plot( numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, topstack+Eseries), color=color_E, zorder=3) - topstack = topstack+Eseries - - if(not combine_Q_infected and any(Q_preseries) and plot_Q_pre=='stacked'): - ax.fill_between(numpy.ma.masked_where(Q_preseries<=0, self.tseries), numpy.ma.masked_where(Q_preseries<=0, topstack+Q_preseries), topstack, facecolor=color_Q_pre, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_{pre}$', zorder=2) - ax.plot( numpy.ma.masked_where(Q_preseries<=0, self.tseries), numpy.ma.masked_where(Q_preseries<=0, topstack+Q_preseries), color=color_Q_pre, zorder=3) - topstack = topstack+Q_preseries - - if(any(I_preseries) and plot_I_pre=='stacked'): - ax.fill_between(numpy.ma.masked_where(I_preseries<=0, self.tseries), numpy.ma.masked_where(I_preseries<=0, topstack+I_preseries), topstack, color=color_I_pre, alpha=0.75, label='$I_{pre}$', zorder=2) - ax.plot( numpy.ma.masked_where(I_preseries<=0, self.tseries), numpy.ma.masked_where(I_preseries<=0, topstack+I_preseries), color=color_I_pre, zorder=3) - topstack = topstack+I_preseries - - if(not combine_Q_infected and any(Q_symseries) and plot_Q_sym=='stacked'): - ax.fill_between(numpy.ma.masked_where(Q_symseries<=0, self.tseries), numpy.ma.masked_where(Q_symseries<=0, topstack+Q_symseries), topstack, facecolor=color_Q_sym, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_{sym}$', zorder=2) - ax.plot( numpy.ma.masked_where(Q_symseries<=0, self.tseries), numpy.ma.masked_where(Q_symseries<=0, topstack+Q_symseries), color=color_Q_sym, zorder=3) - topstack = topstack+Q_symseries - - if(any(I_symseries) and plot_I_sym=='stacked'): - ax.fill_between(numpy.ma.masked_where(I_symseries<=0, self.tseries), numpy.ma.masked_where(I_symseries<=0, topstack+I_symseries), topstack, color=color_I_sym, alpha=0.75, label='$I_{sym}$', zorder=2) - ax.plot( numpy.ma.masked_where(I_symseries<=0, self.tseries), numpy.ma.masked_where(I_symseries<=0, topstack+I_symseries), color=color_I_sym, zorder=3) - topstack = topstack+I_symseries - - if(not combine_Q_infected and any(Q_asymseries) and plot_Q_asym=='stacked'): - ax.fill_between(numpy.ma.masked_where(Q_asymseries<=0, self.tseries), numpy.ma.masked_where(Q_asymseries<=0, topstack+Q_asymseries), topstack, facecolor=color_Q_asym, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_{asym}$', zorder=2) - ax.plot( numpy.ma.masked_where(Q_asymseries<=0, self.tseries), numpy.ma.masked_where(Q_asymseries<=0, topstack+Q_asymseries), color=color_Q_asym, zorder=3) - topstack = topstack+Q_asymseries - - if(any(I_asymseries) and plot_I_asym=='stacked'): - ax.fill_between(numpy.ma.masked_where(I_asymseries<=0, self.tseries), numpy.ma.masked_where(I_asymseries<=0, topstack+I_asymseries), topstack, color=color_I_asym, alpha=0.75, label='$I_{asym}$', zorder=2) - ax.plot( numpy.ma.masked_where(I_asymseries<=0, self.tseries), numpy.ma.masked_where(I_asymseries<=0, topstack+I_asymseries), color=color_I_asym, zorder=3) - topstack = topstack+I_asymseries - - if(any(Q_Rseries) and plot_Q_R=='stacked'): - ax.fill_between(numpy.ma.masked_where(Q_Rseries<=0, self.tseries), numpy.ma.masked_where(Q_Rseries<=0, topstack+Q_Rseries), topstack, facecolor=color_Q_R, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_R$', zorder=2) - ax.plot( numpy.ma.masked_where(Q_Rseries<=0, self.tseries), numpy.ma.masked_where(Q_Rseries<=0, topstack+Q_Rseries), color=color_Q_R, zorder=3) - topstack = topstack+Q_Rseries - - if(any(Rseries) and plot_R=='stacked'): - ax.fill_between(numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, topstack+Rseries), topstack, color=color_R, alpha=0.75, label='$R$', zorder=2) - ax.plot( numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, topstack+Rseries), color=color_R, zorder=3) - topstack = topstack+Rseries - - if(any(Q_Sseries) and plot_Q_S=='stacked'): - ax.fill_between(numpy.ma.masked_where(Q_Sseries<=0, self.tseries), numpy.ma.masked_where(Q_Sseries<=0, topstack+Q_Sseries), topstack, facecolor=color_Q_S, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_S$', zorder=2) - ax.plot( numpy.ma.masked_where(Q_Sseries<=0, self.tseries), numpy.ma.masked_where(Q_Sseries<=0, topstack+Q_Sseries), color=color_Q_S, zorder=3) - topstack = topstack+Q_Sseries - - if(any(Sseries) and plot_S=='stacked'): - ax.fill_between(numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, topstack+Sseries), topstack, color=color_S, alpha=0.75, label='$S$', zorder=2) - ax.plot( numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, topstack+Sseries), color=color_S, zorder=3) - topstack = topstack+Sseries - - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Draw the shaded variables: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(any(Fseries) and plot_F=='shaded'): - ax.fill_between(numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, Fseries), 0, color=color_F, alpha=0.75, label='$F$', zorder=4) - ax.plot( numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, Fseries), color=color_F, zorder=5) - - if(any(Hseries) and plot_H=='shaded'): - ax.fill_between(numpy.ma.masked_where(Hseries<=0, self.tseries), numpy.ma.masked_where(Hseries<=0, Hseries), 0, color=color_H, alpha=0.75, label='$H$', zorder=4) - ax.plot( numpy.ma.masked_where(Hseries<=0, self.tseries), numpy.ma.masked_where(Hseries<=0, Hseries), color=color_H, zorder=5) - - if(combine_Q_infected and any(Q_infectedseries) and plot_Q_E=='shaded' and plot_Q_pre=='shaded' and plot_Q_sym=='shaded' and plot_Q_asym=='shaded'): - ax.fill_between(numpy.ma.masked_where(Q_infectedseries<=0, self.tseries), numpy.ma.masked_where(Q_infectedseries<=0, Q_infectedseries), 0, color=color_Q_infected, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_{infected}$', zorder=4) - ax.plot( numpy.ma.masked_where(Q_infectedseries<=0, self.tseries), numpy.ma.masked_where(Q_infectedseries<=0, Q_infectedseries), color=color_Q_infected, zorder=5) - - if(not combine_Q_infected and any(Q_Eseries) and plot_Q_E=='shaded'): - ax.fill_between(numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, Q_Eseries), 0, facecolor=color_Q_E, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_E$', zorder=4) - ax.plot( numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, Q_Eseries), color=color_Q_E, zorder=5) - - if(any(Eseries) and plot_E=='shaded'): - ax.fill_between(numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, Eseries), 0, color=color_E, alpha=0.75, label='$E$', zorder=4) - ax.plot( numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, Eseries), color=color_E, zorder=5) - - if(not combine_Q_infected and any(Q_preseries) and plot_Q_pre=='shaded'): - ax.fill_between(numpy.ma.masked_where(Q_preseries<=0, self.tseries), numpy.ma.masked_where(Q_preseries<=0, Q_preseries), 0, facecolor=color_Q_pre, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_{pre}$', zorder=4) - ax.plot( numpy.ma.masked_where(Q_preseries<=0, self.tseries), numpy.ma.masked_where(Q_preseries<=0, Q_preseries), color=color_Q_pre, zorder=5) - - if(any(I_preseries) and plot_I_pre=='shaded'): - ax.fill_between(numpy.ma.masked_where(I_preseries<=0, self.tseries), numpy.ma.masked_where(I_preseries<=0, I_preseries), 0, color=color_I_pre, alpha=0.75, label='$I_{pre}$', zorder=4) - ax.plot( numpy.ma.masked_where(I_preseries<=0, self.tseries), numpy.ma.masked_where(I_preseries<=0, I_preseries), color=color_I_pre, zorder=5) - - if(not combine_Q_infected and any(Q_symseries) and plot_Q_sym=='shaded'): - ax.fill_between(numpy.ma.masked_where(Q_symseries<=0, self.tseries), numpy.ma.masked_where(Q_symseries<=0, Q_symseries), 0, facecolor=color_Q_sym, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_{sym}$', zorder=4) - ax.plot( numpy.ma.masked_where(Q_symseries<=0, self.tseries), numpy.ma.masked_where(Q_symseries<=0, Q_symseries), color=color_Q_sym, zorder=5) - - if(any(I_symseries) and plot_I_sym=='shaded'): - ax.fill_between(numpy.ma.masked_where(I_symseries<=0, self.tseries), numpy.ma.masked_where(I_symseries<=0, I_symseries), 0, color=color_I_sym, alpha=0.75, label='$I_{sym}$', zorder=4) - ax.plot( numpy.ma.masked_where(I_symseries<=0, self.tseries), numpy.ma.masked_where(I_symseries<=0, I_symseries), color=color_I_sym, zorder=5) - - if(not combine_Q_infected and any(Q_asymseries) and plot_Q_asym=='shaded'): - ax.fill_between(numpy.ma.masked_where(Q_asymseries<=0, self.tseries), numpy.ma.masked_where(Q_asymseries<=0, Q_asymseries), 0, facecolor=color_Q_asym, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_{asym}$', zorder=4) - ax.plot( numpy.ma.masked_where(Q_asymseries<=0, self.tseries), numpy.ma.masked_where(Q_asymseries<=0, Q_asymseries), color=color_Q_asym, zorder=5) - - if(any(I_asymseries) and plot_I_asym=='shaded'): - ax.fill_between(numpy.ma.masked_where(I_asymseries<=0, self.tseries), numpy.ma.masked_where(I_asymseries<=0, I_asymseries), 0, color=color_I_asym, alpha=0.75, label='$I_{asym}$', zorder=4) - ax.plot( numpy.ma.masked_where(I_asymseries<=0, self.tseries), numpy.ma.masked_where(I_asymseries<=0, I_asymseries), color=color_I_asym, zorder=5) - - if(any(Q_Rseries) and plot_Q_R=='shaded'): - ax.fill_between(numpy.ma.masked_where(Q_Rseries<=0, self.tseries), numpy.ma.masked_where(Q_Rseries<=0, Q_Rseries), 0, facecolor=color_Q_R, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_R$', zorder=4) - ax.plot( numpy.ma.masked_where(Q_Rseries<=0, self.tseries), numpy.ma.masked_where(Q_Rseries<=0, Q_Rseries), color=color_Q_R, zorder=5) - - if(any(Rseries) and plot_R=='shaded'): - ax.fill_between(numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, Rseries), 0, color=color_R, alpha=0.75, label='$R$', zorder=4) - ax.plot( numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, Rseries), color=color_R, zorder=5) - - if(any(Q_Sseries) and plot_Q_S=='shaded'): - ax.fill_between(numpy.ma.masked_where(Q_Sseries<=0, self.tseries), numpy.ma.masked_where(Q_Sseries<=0, Q_Sseries), 0, facecolor=color_Q_S, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_S$', zorder=4) - ax.plot( numpy.ma.masked_where(Q_Sseries<=0, self.tseries), numpy.ma.masked_where(Q_Sseries<=0, Q_Sseries), color=color_Q_S, zorder=5) - - if(any(Sseries) and plot_S=='shaded'): - ax.fill_between(numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, Sseries), 0, color=color_S, alpha=0.75, label='$S$', zorder=4) - ax.plot( numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, Sseries), color=color_S, zorder=5) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Draw the line variables: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(any(Fseries) and plot_F=='line'): - ax.plot(numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, Fseries), color=color_F, label='$F$', zorder=6) - - if(any(Hseries) and plot_H=='line'): - ax.plot(numpy.ma.masked_where(Hseries<=0, self.tseries), numpy.ma.masked_where(Hseries<=0, Hseries), color=color_H, label='$H$', zorder=6) - - if(combine_Q_infected and any(Q_infectedseries) and plot_Q_E=='line' and plot_Q_pre=='line' and plot_Q_sym=='line' and plot_Q_asym=='line'): - ax.plot(numpy.ma.masked_where(Q_infectedseries<=0, self.tseries), numpy.ma.masked_where(Q_infectedseries<=0, Q_infectedseries), color=color_Q_infected, label='$Q_{infected}$', zorder=6) - - if(not combine_Q_infected and any(Q_Eseries) and plot_Q_E=='line'): - ax.plot(numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, Q_Eseries), color=color_Q_E, label='$Q_E$', zorder=6) - - if(any(Eseries) and plot_E=='line'): - ax.plot(numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, Eseries), color=color_E, label='$E$', zorder=6) - - if(not combine_Q_infected and any(Q_preseries) and plot_Q_pre=='line'): - ax.plot(numpy.ma.masked_where(Q_preseries<=0, self.tseries), numpy.ma.masked_where(Q_preseries<=0, Q_preseries), color=color_Q_pre, label='$Q_{pre}$', zorder=6) - - if(any(I_preseries) and plot_I_pre=='line'): - ax.plot(numpy.ma.masked_where(I_preseries<=0, self.tseries), numpy.ma.masked_where(I_preseries<=0, I_preseries), color=color_I_pre, label='$I_{pre}$', zorder=6) - - if(not combine_Q_infected and any(Q_symseries) and plot_Q_sym=='line'): - ax.plot(numpy.ma.masked_where(Q_symseries<=0, self.tseries), numpy.ma.masked_where(Q_symseries<=0, Q_symseries), color=color_Q_sym, label='$Q_{sym}$', zorder=6) - - if(any(I_symseries) and plot_I_sym=='line'): - ax.plot(numpy.ma.masked_where(I_symseries<=0, self.tseries), numpy.ma.masked_where(I_symseries<=0, I_symseries), color=color_I_sym, label='$I_{sym}$', zorder=6) - - if(not combine_Q_infected and any(Q_asymseries) and plot_Q_asym=='line'): - ax.plot(numpy.ma.masked_where(Q_asymseries<=0, self.tseries), numpy.ma.masked_where(Q_asymseries<=0, Q_asymseries), color=color_Q_asym, label='$Q_{asym}$', zorder=6) - - if(any(I_asymseries) and plot_I_asym=='line'): - ax.plot(numpy.ma.masked_where(I_asymseries<=0, self.tseries), numpy.ma.masked_where(I_asymseries<=0, I_asymseries), color=color_I_asym, label='$I_{asym}$', zorder=6) - - if(any(Q_Rseries) and plot_Q_R=='line'): - ax.plot(numpy.ma.masked_where(Q_Rseries<=0, self.tseries), numpy.ma.masked_where(Q_Rseries<=0, Q_Rseries), color=color_Q_R, linestyle='--', label='$Q_R$', zorder=6) - - if(any(Rseries) and plot_R=='line'): - ax.plot(numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, Rseries), color=color_R, label='$R$', zorder=6) - - if(any(Q_Sseries) and plot_Q_S=='line'): - ax.plot(numpy.ma.masked_where(Q_Sseries<=0, self.tseries), numpy.ma.masked_where(Q_Sseries<=0, Q_Sseries), color=color_Q_S, linestyle='--', label='$Q_S$', zorder=6) - - if(any(Sseries) and plot_S=='line'): - ax.plot(numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, Sseries), color=color_S, label='$S$', zorder=6) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Draw the vertical line annotations: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(len(vlines)>0 and len(vline_colors)==0): - vline_colors = ['gray']*len(vlines) - if(len(vlines)>0 and len(vline_labels)==0): - vline_labels = [None]*len(vlines) - if(len(vlines)>0 and len(vline_styles)==0): - vline_styles = [':']*len(vlines) - for vline_x, vline_color, vline_style, vline_label in zip(vlines, vline_colors, vline_styles, vline_labels): - if(vline_x is not None): - ax.axvline(x=vline_x, color=vline_color, linestyle=vline_style, alpha=1, label=vline_label) - - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Draw the plot labels: - #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - ax.set_xlabel('days') - ax.set_ylabel('percent of population' if plot_percentages else 'number of individuals') - ax.set_xlim(0, (max(self.tseries) if not xlim else xlim)) - ax.set_ylim(0, ylim) - if(plot_percentages): - ax.set_yticklabels(['{:,.0%}'.format(y) for y in ax.get_yticks()]) - if(legend): - legend_handles, legend_labels = ax.get_legend_handles_labels() - ax.legend(legend_handles[::-1], legend_labels[::-1], loc='upper right', facecolor='white', edgecolor='none', framealpha=0.9, prop={'size': 8}) - if(title): - ax.set_title(title, size=12) - if(side_title): - ax.annotate(side_title, (0, 0.5), xytext=(-45, 0), ha='right', va='center', - size=12, rotation=90, xycoords='axes fraction', textcoords='offset points') - - return ax - - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def figure_basic(self, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_sym='line', plot_I_asym='line', - plot_H='line', plot_R='line', plot_F='line', - plot_Q_E='line', plot_Q_pre='line', plot_Q_sym='line', plot_Q_asym='line', - plot_Q_S=False, plot_Q_R=False, combine_Q_infected=True, - color_S='tab:green', color_E='orange', color_I_pre='tomato', color_I_sym='crimson', color_I_asym='#F0909B', - color_H='violet', color_R='tab:blue', color_F='black', - color_Q_E='orange', color_Q_pre='tomato', color_Q_sym='crimson', color_Q_asym='#F0909B', - color_Q_S='tab:green', color_Q_R='tab:blue', color_Q_infected='tab:purple', - color_reference='#E0E0E0', - dashed_reference_results=None, dashed_reference_label='reference', - shaded_reference_results=None, shaded_reference_label='reference', - vlines=[], vline_colors=[], vline_styles=[], vline_labels=[], - ylim=None, xlim=None, legend=True, title=None, side_title=None, plot_percentages=True, - figsize=(12,8), use_seaborn=True, show=True): - - import matplotlib.pyplot as pyplot - - fig, ax = pyplot.subplots(figsize=figsize) - - if(use_seaborn): - import seaborn - seaborn.set_style('ticks') - seaborn.despine() - - self.plot(ax=ax, plot_S=plot_S, plot_E=plot_E, plot_I_pre=plot_I_pre, plot_I_sym=plot_I_sym, plot_I_asym=plot_I_asym, - plot_H=plot_H, plot_R=plot_R, plot_F=plot_F, - plot_Q_E=plot_Q_E, plot_Q_pre=plot_Q_pre, plot_Q_sym=plot_Q_sym, plot_Q_asym=plot_Q_asym, - plot_Q_S=plot_Q_S, plot_Q_R=plot_Q_R, combine_Q_infected=combine_Q_infected, - color_S=color_S, color_E=color_E, color_I_pre=color_I_pre, color_I_sym=color_I_sym, color_I_asym=color_I_asym, - color_H=color_H, color_R=color_R, color_F=color_F, - color_Q_E=color_Q_E, color_Q_pre=color_Q_pre, color_Q_sym=color_Q_sym, color_Q_asym=color_Q_asym, - color_Q_S=color_Q_S, color_Q_R=color_Q_R, color_Q_infected=color_Q_infected, - color_reference=color_reference, - dashed_reference_results=dashed_reference_results, dashed_reference_label=dashed_reference_label, - shaded_reference_results=shaded_reference_results, shaded_reference_label=shaded_reference_label, - vlines=vlines, vline_colors=vline_colors, vline_styles=vline_styles, vline_labels=vline_labels, - ylim=ylim, xlim=xlim, legend=legend, title=title, side_title=side_title, plot_percentages=plot_percentages) - - if(show): - pyplot.show() - - return fig, ax - - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - def figure_infections(self, plot_S=False, plot_E='stacked', plot_I_pre='stacked', plot_I_sym='stacked', plot_I_asym='stacked', - plot_H='stacked', plot_R=False, plot_F='stacked', - plot_Q_E='stacked', plot_Q_pre='stacked', plot_Q_sym='stacked', plot_Q_asym='stacked', - plot_Q_S=False, plot_Q_R=False, combine_Q_infected=True, - color_S='tab:green', color_E='orange', color_I_pre='tomato', color_I_sym='crimson', color_I_asym='#F0909B', - color_H='violet', color_R='tab:blue', color_F='black', - color_Q_E='orange', color_Q_pre='tomato', color_Q_sym='crimson', color_Q_asym='#F0909B', - color_Q_S='tab:green', color_Q_R='tab:blue', color_Q_infected='tab:purple', - color_reference='#E0E0E0', - dashed_reference_results=None, dashed_reference_label='reference', - shaded_reference_results=None, shaded_reference_label='reference', - vlines=[], vline_colors=[], vline_styles=[], vline_labels=[], - ylim=None, xlim=None, legend=True, title=None, side_title=None, plot_percentages=True, - figsize=(12,8), use_seaborn=True, show=True): - - import matplotlib.pyplot as pyplot - - fig, ax = pyplot.subplots(figsize=figsize) - - if(use_seaborn): - import seaborn - seaborn.set_style('ticks') - seaborn.despine() - - self.plot(ax=ax, plot_S=plot_S, plot_E=plot_E, plot_I_pre=plot_I_pre, plot_I_sym=plot_I_sym, plot_I_asym=plot_I_asym, - plot_H=plot_H, plot_R=plot_R, plot_F=plot_F, - plot_Q_E=plot_Q_E, plot_Q_pre=plot_Q_pre, plot_Q_sym=plot_Q_sym, plot_Q_asym=plot_Q_asym, - plot_Q_S=plot_Q_S, plot_Q_R=plot_Q_R, combine_Q_infected=combine_Q_infected, - color_S=color_S, color_E=color_E, color_I_pre=color_I_pre, color_I_sym=color_I_sym, color_I_asym=color_I_asym, - color_H=color_H, color_R=color_R, color_F=color_F, - color_Q_E=color_Q_E, color_Q_pre=color_Q_pre, color_Q_sym=color_Q_sym, color_Q_asym=color_Q_asym, - color_Q_S=color_Q_S, color_Q_R=color_Q_R, color_Q_infected=color_Q_infected, - color_reference=color_reference, - dashed_reference_results=dashed_reference_results, dashed_reference_label=dashed_reference_label, - shaded_reference_results=shaded_reference_results, shaded_reference_label=shaded_reference_label, - vlines=vlines, vline_colors=vline_colors, vline_styles=vline_styles, vline_labels=vline_labels, - ylim=ylim, xlim=xlim, legend=legend, title=title, side_title=side_title, plot_percentages=plot_percentages) - - if(show): - pyplot.show() - - return fig, ax - - -#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% -#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% - - - - - - diff --git a/seirsplus/models/__init__.py b/seirsplus/models/__init__.py new file mode 100644 index 0000000..d87ccb2 --- /dev/null +++ b/seirsplus/models/__init__.py @@ -0,0 +1,5 @@ +"""SEIRS and extended SEIRS models.""" + +from .extended_seirs_network_model import ExtSEIRSNetworkModel +from .seirs_model import SEIRSModel +from .seirs_network_model import SEIRSNetworkModel diff --git a/seirsplus/models/base_plotable_model.py b/seirsplus/models/base_plotable_model.py new file mode 100644 index 0000000..0f7003c --- /dev/null +++ b/seirsplus/models/base_plotable_model.py @@ -0,0 +1,262 @@ +"""Common model.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy + +class BasePlotableModel: + """ + Base model to abstract common plotting features. + """ + + numE = None + numF = None + numI = None + numQ_E = None + numQ_I = None + numR = None + numS = None + tseries = None + + plotting_number_property = None + """Property to access the number to base plotting on.""" + + def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='line', plot_F='line', + plot_Q_E='line', plot_Q_I='line', combine_Q=True, + color_S='tab:green', color_E='orange', color_I='crimson', color_R='tab:blue', color_F='black', + color_Q_E='mediumorchid', color_Q_I='mediumorchid', color_reference='#E0E0E0', + dashed_reference_results=None, dashed_reference_label='reference', + shaded_reference_results=None, shaded_reference_label='reference', + vlines=[], vline_colors=[], vline_styles=[], vline_labels=[], + ylim=None, xlim=None, legend=True, title=None, side_title=None, plot_percentages=True): + + import matplotlib.pyplot as pyplot + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Create an Axes object if None provided: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + if not ax: + fig, ax = pyplot.subplots() + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Prepare data series to be plotted: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + Fseries = self.numF/getattr(self, self.plotting_number_property) if plot_percentages else self.numF + Eseries = self.numE/getattr(self, self.plotting_number_property) if plot_percentages else self.numE + Dseries = (self.numQ_E+self.numQ_I)/getattr(self, self.plotting_number_property) if plot_percentages else (self.numQ_E+self.numQ_I) + Q_Eseries = self.numQ_E/getattr(self, self.plotting_number_property) if plot_percentages else self.numQ_E + Q_Iseries = self.numQ_I/getattr(self, self.plotting_number_property) if plot_percentages else self.numQ_I + Iseries = self.numI/getattr(self, self.plotting_number_property) if plot_percentages else self.numI + Rseries = self.numR/getattr(self, self.plotting_number_property) if plot_percentages else self.numR + Sseries = self.numS/getattr(self, self.plotting_number_property) if plot_percentages else self.numS + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Draw the reference data: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + if dashed_reference_results: + dashedReference_tseries = dashed_reference_results.tseries[::int(getattr(self, self.plotting_number_property)/100)] + dashedReference_IDEstack = (dashed_reference_results.numI + dashed_reference_results.numQ_I + dashed_reference_results.numQ_E + dashed_reference_results.numE)[::int(getattr(self, self.plotting_number_property)/100)] / (getattr(self, self.plotting_number_property) if plot_percentages else 1) + ax.plot(dashedReference_tseries, dashedReference_IDEstack, color='#E0E0E0', linestyle='--', label='$I+D+E$ ('+dashed_reference_label+')', zorder=0) + if shaded_reference_results: + shadedReference_tseries = shaded_reference_results.tseries + shadedReference_IDEstack = (shaded_reference_results.numI + shaded_reference_results.numQ_I + shaded_reference_results.numQ_E + shaded_reference_results.numE) / (getattr(self, self.plotting_number_property) if plot_percentages else 1) + ax.fill_between(shaded_reference_results.tseries, shadedReference_IDEstack, 0, color='#EFEFEF', label='$I+D+E$ ('+shaded_reference_label+')', zorder=0) + ax.plot(shaded_reference_results.tseries, shadedReference_IDEstack, color='#E0E0E0', zorder=1) + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Draw the stacked variables: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + topstack = numpy.zeros_like(self.tseries) + if (any(Fseries) and plot_F=='stacked'): + ax.fill_between(numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, topstack+Fseries), topstack, color=color_F, alpha=0.5, label='$F$', zorder=2) + ax.plot( numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, topstack+Fseries), color=color_F, zorder=3) + topstack = topstack+Fseries + if (any(Eseries) and plot_E=='stacked'): + ax.fill_between(numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, topstack+Eseries), topstack, color=color_E, alpha=0.5, label='$E$', zorder=2) + ax.plot( numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, topstack+Eseries), color=color_E, zorder=3) + topstack = topstack+Eseries + if (combine_Q and plot_Q_E=='stacked' and plot_Q_I=='stacked'): + ax.fill_between(numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, topstack+Dseries), topstack, color=color_Q_E, alpha=0.5, label='$Q_{all}$', zorder=2) + ax.plot( numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, topstack+Dseries), color=color_Q_E, zorder=3) + topstack = topstack+Dseries + else: + if (any(Q_Eseries) and plot_Q_E=='stacked'): + ax.fill_between(numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, topstack+Q_Eseries), topstack, color=color_Q_E, alpha=0.5, label='$Q_E$', zorder=2) + ax.plot( numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, topstack+Q_Eseries), color=color_Q_E, zorder=3) + topstack = topstack+Q_Eseries + if (any(Q_Iseries) and plot_Q_I=='stacked'): + ax.fill_between(numpy.ma.masked_where(Q_Iseries<=0, self.tseries), numpy.ma.masked_where(Q_Iseries<=0, topstack+Q_Iseries), topstack, color=color_Q_I, alpha=0.5, label='$Q_I$', zorder=2) + ax.plot( numpy.ma.masked_where(Q_Iseries<=0, self.tseries), numpy.ma.masked_where(Q_Iseries<=0, topstack+Q_Iseries), color=color_Q_I, zorder=3) + topstack = topstack+Q_Iseries + if (any(Iseries) and plot_I=='stacked'): + ax.fill_between(numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, topstack+Iseries), topstack, color=color_I, alpha=0.5, label='$I$', zorder=2) + ax.plot( numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, topstack+Iseries), color=color_I, zorder=3) + topstack = topstack+Iseries + if (any(Rseries) and plot_R=='stacked'): + ax.fill_between(numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, topstack+Rseries), topstack, color=color_R, alpha=0.5, label='$R$', zorder=2) + ax.plot( numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, topstack+Rseries), color=color_R, zorder=3) + topstack = topstack+Rseries + if (any(Sseries) and plot_S=='stacked'): + ax.fill_between(numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, topstack+Sseries), topstack, color=color_S, alpha=0.5, label='$S$', zorder=2) + ax.plot( numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, topstack+Sseries), color=color_S, zorder=3) + topstack = topstack+Sseries + + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Draw the shaded variables: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + if (any(Fseries) and plot_F=='shaded'): + ax.fill_between(numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, Fseries), 0, color=color_F, alpha=0.5, label='$F$', zorder=4) + ax.plot( numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, Fseries), color=color_F, zorder=5) + if (any(Eseries) and plot_E=='shaded'): + ax.fill_between(numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, Eseries), 0, color=color_E, alpha=0.5, label='$E$', zorder=4) + ax.plot( numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, Eseries), color=color_E, zorder=5) + if (combine_Q and (any(Dseries) and plot_Q_E=='shaded' and plot_Q_I=='shaded')): + ax.fill_between(numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, Dseries), 0, color=color_Q_E, alpha=0.5, label='$Q_{all}$', zorder=4) + ax.plot( numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, Dseries), color=color_Q_E, zorder=5) + else: + if (any(Q_Eseries) and plot_Q_E=='shaded'): + ax.fill_between(numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, Q_Eseries), 0, color=color_Q_E, alpha=0.5, label='$Q_E$', zorder=4) + ax.plot( numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, Q_Eseries), color=color_Q_E, zorder=5) + if (any(Q_Iseries) and plot_Q_I=='shaded'): + ax.fill_between(numpy.ma.masked_where(Q_Iseries<=0, self.tseries), numpy.ma.masked_where(Q_Iseries<=0, Q_Iseries), 0, color=color_Q_I, alpha=0.5, label='$Q_I$', zorder=4) + ax.plot( numpy.ma.masked_where(Q_Iseries<=0, self.tseries), numpy.ma.masked_where(Q_Iseries<=0, Q_Iseries), color=color_Q_I, zorder=5) + if (any(Iseries) and plot_I=='shaded'): + ax.fill_between(numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, Iseries), 0, color=color_I, alpha=0.5, label='$I$', zorder=4) + ax.plot( numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, Iseries), color=color_I, zorder=5) + if (any(Sseries) and plot_S=='shaded'): + ax.fill_between(numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, Sseries), 0, color=color_S, alpha=0.5, label='$S$', zorder=4) + ax.plot( numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, Sseries), color=color_S, zorder=5) + if (any(Rseries) and plot_R=='shaded'): + ax.fill_between(numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, Rseries), 0, color=color_R, alpha=0.5, label='$R$', zorder=4) + ax.plot( numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, Rseries), color=color_R, zorder=5) + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Draw the line variables: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + if (any(Fseries) and plot_F=='line'): + ax.plot(numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, Fseries), color=color_F, label='$F$', zorder=6) + if (any(Eseries) and plot_E=='line'): + ax.plot(numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, Eseries), color=color_E, label='$E$', zorder=6) + if (combine_Q and (any(Dseries) and plot_Q_E=='line' and plot_Q_I == 'line')): + ax.plot(numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, Dseries), color=color_Q_E, label='$Q_{all}$', zorder=6) + else: + if (any(Q_Eseries) and plot_Q_E=='line'): + ax.plot(numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, Q_Eseries), color=color_Q_E, label='$Q_E$', zorder=6) + if (any(Q_Iseries) and plot_Q_I=='line'): + ax.plot(numpy.ma.masked_where(Q_Iseries<=0, self.tseries), numpy.ma.masked_where(Q_Iseries<=0, Q_Iseries), color=color_Q_I, label='$Q_I$', zorder=6) + if (any(Iseries) and plot_I=='line'): + ax.plot(numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, Iseries), color=color_I, label='$I$', zorder=6) + if (any(Sseries) and plot_S=='line'): + ax.plot(numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, Sseries), color=color_S, label='$S$', zorder=6) + if (any(Rseries) and plot_R=='line'): + ax.plot(numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, Rseries), color=color_R, label='$R$', zorder=6) + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Draw the vertical line annotations: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + if (len(vlines)>0 and len(vline_colors)==0): + vline_colors = ['gray']*len(vlines) + if (len(vlines)>0 and len(vline_labels)==0): + vline_labels = [None]*len(vlines) + if (len(vlines)>0 and len(vline_styles)==0): + vline_styles = [':']*len(vlines) + for vline_x, vline_color, vline_style, vline_label in zip(vlines, vline_colors, vline_styles, vline_labels): + if vline_x is not None: + ax.axvline(x=vline_x, color=vline_color, linestyle=vline_style, alpha=1, label=vline_label) + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Draw the plot labels: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + ax.set_xlabel('days') + ax.set_ylabel('percent of population' if plot_percentages else 'number of individuals') + ax.set_xlim(0, (max(self.tseries) if not xlim else xlim)) + ax.set_ylim(0, ylim) + if plot_percentages: + ax.set_yticklabels(['{:,.0%}'.format(y) for y in ax.get_yticks()]) + if legend: + legend_handles, legend_labels = ax.get_legend_handles_labels() + ax.legend(legend_handles[::-1], legend_labels[::-1], loc='upper right', facecolor='white', edgecolor='none', framealpha=0.9, prop={'size': 8}) + if title: + ax.set_title(title, size=12) + if side_title: + ax.annotate(side_title, (0, 0.5), xytext=(-45, 0), ha='right', va='center', + size=12, rotation=90, xycoords='axes fraction', textcoords='offset points') + + return ax + + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def figure_basic(self, plot_S='line', plot_E='line', plot_I='line',plot_R='line', plot_F='line', + plot_Q_E='line', plot_Q_I='line', combine_Q=True, + color_S='tab:green', color_E='orange', color_I='crimson', color_R='tab:blue', color_F='black', + color_Q_E='mediumorchid', color_Q_I='mediumorchid', color_reference='#E0E0E0', + dashed_reference_results=None, dashed_reference_label='reference', + shaded_reference_results=None, shaded_reference_label='reference', + vlines=[], vline_colors=[], vline_styles=[], vline_labels=[], + ylim=None, xlim=None, legend=True, title=None, side_title=None, plot_percentages=True, + figsize=(12,8), use_seaborn=True, show=True): + + import matplotlib.pyplot as pyplot + + fig, ax = pyplot.subplots(figsize=figsize) + + if use_seaborn: + import seaborn + seaborn.set_style('ticks') + seaborn.despine() + + self.plot(ax=ax, plot_S=plot_S, plot_E=plot_E, plot_I=plot_I,plot_R=plot_R, plot_F=plot_F, + plot_Q_E=plot_Q_E, plot_Q_I=plot_Q_I, combine_Q=combine_Q, + color_S=color_S, color_E=color_E, color_I=color_I, color_R=color_R, color_F=color_F, + color_Q_E=color_Q_E, color_Q_I=color_Q_I, color_reference=color_reference, + dashed_reference_results=dashed_reference_results, dashed_reference_label=dashed_reference_label, + shaded_reference_results=shaded_reference_results, shaded_reference_label=shaded_reference_label, + vlines=vlines, vline_colors=vline_colors, vline_styles=vline_styles, vline_labels=vline_labels, + ylim=ylim, xlim=xlim, legend=legend, title=title, side_title=side_title, plot_percentages=plot_percentages) + + if show: + pyplot.show() + + return fig, ax + + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def figure_infections(self, plot_S=False, plot_E='stacked', plot_I='stacked',plot_R=False, plot_F=False, + plot_Q_E='stacked', plot_Q_I='stacked', combine_Q=True, + color_S='tab:green', color_E='orange', color_I='crimson', color_R='tab:blue', color_F='black', + color_Q_E='mediumorchid', color_Q_I='mediumorchid', color_reference='#E0E0E0', + dashed_reference_results=None, dashed_reference_label='reference', + shaded_reference_results=None, shaded_reference_label='reference', + vlines=[], vline_colors=[], vline_styles=[], vline_labels=[], + ylim=None, xlim=None, legend=True, title=None, side_title=None, plot_percentages=True, + figsize=(12,8), use_seaborn=True, show=True): + + import matplotlib.pyplot as pyplot + + fig, ax = pyplot.subplots(figsize=figsize) + + if use_seaborn: + import seaborn + seaborn.set_style('ticks') + seaborn.despine() + + self.plot(ax=ax, plot_S=plot_S, plot_E=plot_E, plot_I=plot_I,plot_R=plot_R, plot_F=plot_F, + plot_Q_E=plot_Q_E, plot_Q_I=plot_Q_I, combine_Q=combine_Q, + color_S=color_S, color_E=color_E, color_I=color_I, color_R=color_R, color_F=color_F, + color_Q_E=color_Q_E, color_Q_I=color_Q_I, color_reference=color_reference, + dashed_reference_results=dashed_reference_results, dashed_reference_label=dashed_reference_label, + shaded_reference_results=shaded_reference_results, shaded_reference_label=shaded_reference_label, + vlines=vlines, vline_colors=vline_colors, vline_styles=vline_styles, vline_labels=vline_labels, + ylim=ylim, xlim=xlim, legend=legend, title=title, side_title=side_title, plot_percentages=plot_percentages) + + if show: + pyplot.show() + + return fig, ax diff --git a/seirsplus/models/extended_seirs_network_model.py b/seirsplus/models/extended_seirs_network_model.py new file mode 100644 index 0000000..26e1a36 --- /dev/null +++ b/seirsplus/models/extended_seirs_network_model.py @@ -0,0 +1,1563 @@ +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import networkx +import numpy +import scipy + +from .base_plotable_model import BasePlotableModel + + +class ExtSEIRSNetworkModel(BasePlotableModel): + """ + A class to simulate the Extended SEIRS Stochastic Network Model + =================================================== + Params: + G Network adjacency matrix (numpy array) or Networkx graph object. + beta Rate of transmission (global interactions) + beta_local Rate(s) of transmission between adjacent individuals (optional) + beta_asym Rate of transmission (global interactions) + beta_asym_local Rate(s) of transmission between adjacent individuals (optional) + sigma Rate of progression to infectious state (inverse of latent period) + lamda Rate of progression to infectious (a)symptomatic state (inverse of prodromal period) + eta Rate of progression to hospitalized state (inverse of onset-to-admission period) + gamma Rate of recovery for non-hospitalized symptomatic individuals (inverse of symptomatic infectious period) + gamma_asym Rate of recovery for asymptomatic individuals (inverse of asymptomatic infectious period) + gamma_H Rate of recovery for hospitalized symptomatic individuals (inverse of hospitalized infectious period) + mu_H Rate of death for hospitalized individuals (inverse of admission-to-death period) + xi Rate of re-susceptibility (upon recovery) + mu_0 Rate of baseline death + nu Rate of baseline birth + a Probability of an infected individual remaining asymptomatic + h Probability of a symptomatic individual being hospitalized + f Probability of death for hospitalized individuals (case fatality rate) + p Probability of individuals interacting with global population + + G_Q Quarantine adjacency matrix (numpy array) or Networkx graph object. + beta_Q Rate of transmission for isolated individuals (global interactions) + beta_Q_local Rate(s) of transmission (exposure) for adjacent isolated individuals (optional) + sigma_Q Rate of progression to infectious state for isolated individuals + lamda_Q Rate of progression to infectious (a)symptomatic state for isolated individuals + eta_Q Rate of progression to hospitalized state for isolated individuals + gamma_Q_sym Rate of recovery for non-hospitalized symptomatic individuals for isolated individuals + gamma_Q_asym Rate of recovery for asymptomatic individuals for isolated individuals + theta_E Rate of random testing for exposed individuals + theta_pre Rate of random testing for infectious pre-symptomatic individuals + theta_sym Rate of random testing for infectious symptomatic individuals + theta_asym Rate of random testing for infectious asymptomatic individuals + phi_E Rate of testing when a close contact has tested positive for exposed individuals + phi_pre Rate of testing when a close contact has tested positive for infectious pre-symptomatic individuals + phi_sym Rate of testing when a close contact has tested positive for infectious symptomatic individuals + phi_asym Rate of testing when a close contact has tested positive for infectious asymptomatic individuals + psi_E Probability of positive test for exposed individuals + psi_pre Probability of positive test for infectious pre-symptomatic individuals + psi_sym Probability of positive test for infectious symptomatic individuals + psi_asym Probability of positive test for infectious asymptomatic individuals + q Probability of isolated individuals interacting with global population + isolation_time Time to remain in isolation upon positive test, self-isolation, etc. + + initE Initial number of exposed individuals + initI_pre Initial number of infectious pre-symptomatic individuals + initI_sym Initial number of infectious symptomatic individuals + initI_asym Initial number of infectious asymptomatic individuals + initH Initial number of hospitalized individuals + initR Initial number of recovered individuals + initF Initial number of infection-related fatalities + initQ_S Initial number of isolated susceptible individuals + initQ_E Initial number of isolated exposed individuals + initQ_pre Initial number of isolated infectious pre-symptomatic individuals + initQ_sym Initial number of isolated infectious symptomatic individuals + initQ_asym Initial number of isolated infectious asymptomatic individuals + initQ_R Initial number of isolated recovered individuals + (all remaining nodes initialized susceptible) + """ + + plotting_number_property = 'numNodes' + """Property to access the number to base plotting on.""" + + def __init__(self, G, beta, sigma, lamda, gamma, + gamma_asym=None, eta=0, gamma_H=None, mu_H=0, alpha=1.0, xi=0, mu_0=0, nu=0, a=0, h=0, f=0, p=0, + beta_local=None, beta_asym=None, beta_asym_local=None, beta_pairwise_mode='infected', delta=None, delta_pairwise_mode=None, + G_Q=None, beta_Q=None, beta_Q_local=None, sigma_Q=None, lamda_Q=None, eta_Q=None, gamma_Q_sym=None, gamma_Q_asym=None, alpha_Q=None, delta_Q=None, + theta_S=0, theta_E=0, theta_pre=0, theta_sym=0, theta_asym=0, phi_S=0, phi_E=0, phi_pre=0, phi_sym=0, phi_asym=0, + psi_S=0, psi_E=1, psi_pre=1, psi_sym=1, psi_asym=1, q=0, isolation_time=14, + initE=0, initI_pre=0, initI_sym=0, initI_asym=0, initH=0, initR=0, initF=0, + initQ_S=0, initQ_E=0, initQ_pre=0, initQ_sym=0, initQ_asym=0, initQ_R=0, + o=0, prevalence_ext=0, + transition_mode='exponential_rates', node_groups=None, store_Xseries=False, seed=None): + + if (seed is not None): + numpy.random.seed(seed) + self.seed = seed + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Model Parameters: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + self.parameters = { 'G':G, 'G_Q':G_Q, + 'beta':beta, 'sigma':sigma, 'lamda':lamda, 'gamma':gamma, + 'eta':eta, 'gamma_asym':gamma_asym, 'gamma_H':gamma_H, 'mu_H':mu_H, + 'xi':xi, 'mu_0':mu_0, 'nu':nu, 'a':a, 'h':h, 'f':f, 'p':p, + 'beta_local':beta_local, 'beta_asym':beta_asym, 'beta_asym_local':beta_asym_local, 'beta_pairwise_mode':beta_pairwise_mode, + 'alpha':alpha, 'delta':delta, 'delta_pairwise_mode':delta_pairwise_mode, + 'lamda_Q':lamda_Q, 'beta_Q':beta_Q, 'beta_Q_local':beta_Q_local, 'alpha_Q':alpha_Q, 'sigma_Q':sigma_Q, + 'eta_Q':eta_Q, 'gamma_Q_sym':gamma_Q_sym, 'gamma_Q_asym':gamma_Q_asym, 'delta_Q':delta_Q, + 'theta_S':theta_S, 'theta_E':theta_E, 'theta_pre':theta_pre, 'theta_sym':theta_sym, 'theta_asym':theta_asym, + 'phi_S':phi_S, 'phi_E':phi_E, 'phi_pre':phi_pre, 'phi_sym':phi_sym, 'phi_asym':phi_asym, + 'psi_S':psi_S, 'psi_E':psi_E, 'psi_pre':psi_pre, 'psi_sym':psi_sym, 'psi_asym':psi_asym, 'q':q, 'isolation_time':isolation_time, + 'initE':initE, 'initI_pre':initI_pre, 'initI_sym':initI_sym, 'initI_asym':initI_asym, + 'initH':initH, 'initR':initR, 'initF':initF, + 'initQ_S':initQ_S, 'initQ_E':initQ_E, 'initQ_pre':initQ_pre, + 'initQ_sym':initQ_sym, 'initQ_asym':initQ_asym, 'initQ_R':initQ_R, + 'o':o, 'prevalence_ext':prevalence_ext} + self.update_parameters() + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Each node can undergo 4-6 transitions (sans vitality/re-susceptibility returns to S state), + # so there are ~numNodes*6 events/timesteps expected; initialize numNodes*6 timestep slots to start + # (will be expanded during run if needed for some reason) + self.tseries = numpy.zeros(6*self.numNodes) + self.numS = numpy.zeros(6*self.numNodes) + self.numE = numpy.zeros(6*self.numNodes) + self.numI_pre = numpy.zeros(6*self.numNodes) + self.numI_sym = numpy.zeros(6*self.numNodes) + self.numI_asym = numpy.zeros(6*self.numNodes) + self.numH = numpy.zeros(6*self.numNodes) + self.numR = numpy.zeros(6*self.numNodes) + self.numF = numpy.zeros(6*self.numNodes) + self.numQ_S = numpy.zeros(6*self.numNodes) + self.numQ_E = numpy.zeros(6*self.numNodes) + self.numQ_pre = numpy.zeros(6*self.numNodes) + self.numQ_sym = numpy.zeros(6*self.numNodes) + self.numQ_asym = numpy.zeros(6*self.numNodes) + self.numQ_R = numpy.zeros(6*self.numNodes) + self.N = numpy.zeros(6*self.numNodes) + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Initialize Timekeeping: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + self.t = 0 + self.tmax = 0 # will be set when run() is called + self.tidx = 0 + self.tseries[0] = 0 + + # Vectors holding the time that each node has been in a given state or in isolation: + self.timer_state = numpy.zeros((self.numNodes,1)) + self.timer_isolation = numpy.zeros(self.numNodes) + self.isolationTime = isolation_time + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Initialize Counts of inidividuals with each state: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + self.numE[0] = int(initE) + self.numI_pre[0] = int(initI_pre) + self.numI_sym[0] = int(initI_sym) + self.numI_asym[0] = int(initI_asym) + self.numH[0] = int(initH) + self.numR[0] = int(initR) + self.numF[0] = int(initF) + self.numQ_S[0] = int(initQ_S) + self.numQ_E[0] = int(initQ_E) + self.numQ_pre[0] = int(initQ_pre) + self.numQ_sym[0] = int(initQ_sym) + self.numQ_asym[0] = int(initQ_asym) + self.numQ_R[0] = int(initQ_R) + self.numS[0] = (self.numNodes - self.numE[0] - self.numI_pre[0] - self.numI_sym[0] - self.numI_asym[0] - self.numH[0] - self.numR[0] + - self.numQ_S[0] - self.numQ_E[0] - self.numQ_pre[0] - self.numQ_sym[0] - self.numQ_asym[0] - self.numQ_R[0] - self.numF[0]) + self.N[0] = self.numNodes - self.numF[0] + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Node states: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + self.S = 1 + self.E = 2 + self.I_pre = 3 + self.I_sym = 4 + self.I_asym = 5 + self.H = 6 + self.R = 7 + self.F = 8 + self.Q_S = 11 + self.Q_E = 12 + self.Q_pre = 13 + self.Q_sym = 14 + self.Q_asym = 15 + self.Q_R = 17 + + self.X = numpy.array( [self.S]*int(self.numS[0]) + [self.E]*int(self.numE[0]) + + [self.I_pre]*int(self.numI_pre[0]) + [self.I_sym]*int(self.numI_sym[0]) + [self.I_asym]*int(self.numI_asym[0]) + + [self.H]*int(self.numH[0]) + [self.R]*int(self.numR[0]) + [self.F]*int(self.numF[0]) + + [self.Q_S]*int(self.numQ_S[0]) + [self.Q_E]*int(self.numQ_E[0]) + + [self.Q_pre]*int(self.numQ_pre[0]) + [self.Q_sym]*int(self.numQ_sym[0]) + [self.Q_asym]*int(self.numQ_asym[0]) + + [self.Q_R]*int(self.numQ_R[0]) + ).reshape((self.numNodes,1)) + numpy.random.shuffle(self.X) + + self.store_Xseries = store_Xseries + if store_Xseries: + self.Xseries = numpy.zeros(shape=(6*self.numNodes, self.numNodes), dtype='uint8') + self.Xseries[0,:] = self.X.T + + self.transitions = { + 'StoE': {'currentState':self.S, 'newState':self.E}, + 'StoQS': {'currentState':self.S, 'newState':self.Q_S}, + 'EtoIPRE': {'currentState':self.E, 'newState':self.I_pre}, + 'EtoQE': {'currentState':self.E, 'newState':self.Q_E}, + 'IPREtoISYM': {'currentState':self.I_pre, 'newState':self.I_sym}, + 'IPREtoIASYM': {'currentState':self.I_pre, 'newState':self.I_asym}, + 'IPREtoQPRE': {'currentState':self.I_pre, 'newState':self.Q_pre}, + 'ISYMtoH': {'currentState':self.I_sym, 'newState':self.H}, + 'ISYMtoR': {'currentState':self.I_sym, 'newState':self.R}, + 'ISYMtoQSYM': {'currentState':self.I_sym, 'newState':self.Q_sym}, + 'IASYMtoR': {'currentState':self.I_asym, 'newState':self.R}, + 'IASYMtoQASYM': {'currentState':self.I_asym, 'newState':self.Q_asym}, + 'HtoR': {'currentState':self.H, 'newState':self.R}, + 'HtoF': {'currentState':self.H, 'newState':self.F}, + 'RtoS': {'currentState':self.R, 'newState':self.S}, + 'QStoQE': {'currentState':self.Q_S, 'newState':self.Q_E}, + 'QEtoQPRE': {'currentState':self.Q_E, 'newState':self.Q_pre}, + 'QPREtoQSYM': {'currentState':self.Q_pre, 'newState':self.Q_sym}, + 'QPREtoQASYM': {'currentState':self.Q_pre, 'newState':self.Q_asym}, + 'QSYMtoH': {'currentState':self.Q_sym, 'newState':self.H}, + 'QSYMtoQR': {'currentState':self.Q_sym, 'newState':self.Q_R}, + 'QASYMtoQR': {'currentState':self.Q_asym, 'newState':self.Q_R}, + '_toS': {'currentState':True, 'newState':self.S}, + } + + self.transition_mode = transition_mode + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Initialize other node metadata: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + self.tested = numpy.array([False]*self.numNodes).reshape((self.numNodes,1)) + self.positive = numpy.array([False]*self.numNodes).reshape((self.numNodes,1)) + self.numTested = numpy.zeros(6*self.numNodes) + self.numPositive = numpy.zeros(6*self.numNodes) + + self.testedInCurrentState = numpy.array([False]*self.numNodes).reshape((self.numNodes,1)) + + self.infectionsLog = [] + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Initialize node subgroup data series: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + self.nodeGroupData = None + if node_groups: + self.nodeGroupData = {} + for groupName, nodeList in node_groups.items(): + self.nodeGroupData[groupName] = {'nodes': numpy.array(nodeList), + 'mask': numpy.isin(range(self.numNodes), nodeList).reshape((self.numNodes,1))} + self.nodeGroupData[groupName]['numS'] = numpy.zeros(6*self.numNodes) + self.nodeGroupData[groupName]['numE'] = numpy.zeros(6*self.numNodes) + self.nodeGroupData[groupName]['numI_pre'] = numpy.zeros(6*self.numNodes) + self.nodeGroupData[groupName]['numI_sym'] = numpy.zeros(6*self.numNodes) + self.nodeGroupData[groupName]['numI_asym'] = numpy.zeros(6*self.numNodes) + self.nodeGroupData[groupName]['numH'] = numpy.zeros(6*self.numNodes) + self.nodeGroupData[groupName]['numR'] = numpy.zeros(6*self.numNodes) + self.nodeGroupData[groupName]['numF'] = numpy.zeros(6*self.numNodes) + self.nodeGroupData[groupName]['numQ_S'] = numpy.zeros(6*self.numNodes) + self.nodeGroupData[groupName]['numQ_E'] = numpy.zeros(6*self.numNodes) + self.nodeGroupData[groupName]['numQ_pre'] = numpy.zeros(6*self.numNodes) + self.nodeGroupData[groupName]['numQ_sym'] = numpy.zeros(6*self.numNodes) + self.nodeGroupData[groupName]['numQ_asym'] = numpy.zeros(6*self.numNodes) + self.nodeGroupData[groupName]['numQ_R'] = numpy.zeros(6*self.numNodes) + self.nodeGroupData[groupName]['N'] = numpy.zeros(6*self.numNodes) + self.nodeGroupData[groupName]['numPositive'] = numpy.zeros(6*self.numNodes) + self.nodeGroupData[groupName]['numTested'] = numpy.zeros(6*self.numNodes) + self.nodeGroupData[groupName]['numS'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.S) + self.nodeGroupData[groupName]['numE'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.E) + self.nodeGroupData[groupName]['numI_pre'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.I_pre) + self.nodeGroupData[groupName]['numI_sym'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.I_sym) + self.nodeGroupData[groupName]['numI_asym'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.I_asym) + self.nodeGroupData[groupName]['numH'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.H) + self.nodeGroupData[groupName]['numR'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.R) + self.nodeGroupData[groupName]['numF'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.F) + self.nodeGroupData[groupName]['numQ_S'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_E) + self.nodeGroupData[groupName]['numQ_E'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_E) + self.nodeGroupData[groupName]['numQ_pre'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_pre) + self.nodeGroupData[groupName]['numQ_I_sym'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_I_sym) + self.nodeGroupData[groupName]['numQ_I_asym'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_I_asym) + self.nodeGroupData[groupName]['numQ_R'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_E) + self.nodeGroupData[groupName]['N'][0] = self.numNodes - self.numF[0] + + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def update_parameters(self): + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Model graphs: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + self.G = self.parameters['G'] + # Adjacency matrix: + if type(self.G)==numpy.ndarray: + self.A = scipy.sparse.csr_matrix(self.G) + elif type(self.G)==networkx.classes.graph.Graph: + self.A = networkx.adj_matrix(self.G) # adj_matrix gives scipy.sparse csr_matrix + else: + raise BaseException("Input an adjacency matrix or networkx object only.") + self.numNodes = int(self.A.shape[1]) + self.degree = numpy.asarray(self.node_degrees(self.A)).astype(float) + #---------------------------------------- + if self.parameters['G_Q'] is None: + self.G_Q = self.G # If no Q graph is provided, use G in its place + else: + self.G_Q = self.parameters['G_Q'] + # Quarantine Adjacency matrix: + if type(self.G_Q)==numpy.ndarray: + self.A_Q = scipy.sparse.csr_matrix(self.G_Q) + elif type(self.G_Q)==networkx.classes.graph.Graph: + self.A_Q = networkx.adj_matrix(self.G_Q) # adj_matrix gives scipy.sparse csr_matrix + else: + raise BaseException("Input an adjacency matrix or networkx object only.") + self.numNodes_Q = int(self.A_Q.shape[1]) + self.degree_Q = numpy.asarray(self.node_degrees(self.A_Q)).astype(float) + #---------------------------------------- + assert(self.numNodes == self.numNodes_Q), "The normal and quarantine adjacency graphs must be of the same size." + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Model parameters: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + self.beta = numpy.array(self.parameters['beta']).reshape((self.numNodes, 1)) if isinstance(self.parameters['beta'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['beta'], shape=(self.numNodes,1)) + self.beta_asym = (numpy.array(self.parameters['beta_asym']).reshape((self.numNodes, 1)) if isinstance(self.parameters['beta_asym'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['beta_asym'], shape=(self.numNodes,1))) if self.parameters['beta_asym'] is not None else self.beta + self.sigma = numpy.array(self.parameters['sigma']).reshape((self.numNodes, 1)) if isinstance(self.parameters['sigma'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['sigma'], shape=(self.numNodes,1)) + self.lamda = numpy.array(self.parameters['lamda']).reshape((self.numNodes, 1)) if isinstance(self.parameters['lamda'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['lamda'], shape=(self.numNodes,1)) + self.gamma = numpy.array(self.parameters['gamma']).reshape((self.numNodes, 1)) if isinstance(self.parameters['gamma'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['gamma'], shape=(self.numNodes,1)) + self.eta = numpy.array(self.parameters['eta']).reshape((self.numNodes, 1)) if isinstance(self.parameters['eta'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['eta'], shape=(self.numNodes,1)) + self.gamma_asym = (numpy.array(self.parameters['gamma_asym']).reshape((self.numNodes, 1)) if isinstance(self.parameters['gamma_asym'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['gamma_asym'], shape=(self.numNodes,1))) if self.parameters['gamma_asym'] is not None else self.gamma + self.gamma_H = (numpy.array(self.parameters['gamma_H']).reshape((self.numNodes, 1)) if isinstance(self.parameters['gamma_H'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['gamma_H'], shape=(self.numNodes,1))) if self.parameters['gamma_H'] is not None else self.gamma + self.mu_H = numpy.array(self.parameters['mu_H']).reshape((self.numNodes, 1)) if isinstance(self.parameters['mu_H'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['mu_H'], shape=(self.numNodes,1)) + self.alpha = numpy.array(self.parameters['alpha']).reshape((self.numNodes, 1)) if isinstance(self.parameters['alpha'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['alpha'], shape=(self.numNodes,1)) + self.xi = numpy.array(self.parameters['xi']).reshape((self.numNodes, 1)) if isinstance(self.parameters['xi'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['xi'], shape=(self.numNodes,1)) + self.mu_0 = numpy.array(self.parameters['mu_0']).reshape((self.numNodes, 1)) if isinstance(self.parameters['mu_0'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['mu_0'], shape=(self.numNodes,1)) + self.nu = numpy.array(self.parameters['nu']).reshape((self.numNodes, 1)) if isinstance(self.parameters['nu'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['nu'], shape=(self.numNodes,1)) + self.a = numpy.array(self.parameters['a']).reshape((self.numNodes, 1)) if isinstance(self.parameters['a'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['a'], shape=(self.numNodes,1)) + self.h = numpy.array(self.parameters['h']).reshape((self.numNodes, 1)) if isinstance(self.parameters['h'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['h'], shape=(self.numNodes,1)) + self.f = numpy.array(self.parameters['f']).reshape((self.numNodes, 1)) if isinstance(self.parameters['f'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['f'], shape=(self.numNodes,1)) + self.p = numpy.array(self.parameters['p']).reshape((self.numNodes, 1)) if isinstance(self.parameters['p'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['p'], shape=(self.numNodes,1)) + self.o = numpy.array(self.parameters['o']).reshape((self.numNodes, 1)) if isinstance(self.parameters['o'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['o'], shape=(self.numNodes,1)) + + self.rand_a = numpy.random.rand(self.a.shape[0], self.a.shape[1]) + self.rand_h = numpy.random.rand(self.h.shape[0], self.h.shape[1]) + self.rand_f = numpy.random.rand(self.f.shape[0], self.f.shape[1]) + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # External infection introduction variables: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + self.prevalence_ext = numpy.array(self.parameters['prevalence_ext']).reshape((self.numNodes, 1)) if isinstance(self.parameters['prevalence_ext'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['prevalence_ext'], shape=(self.numNodes,1)) + + #---------------------------------------- + # Testing-related parameters: + #---------------------------------------- + self.beta_Q = (numpy.array(self.parameters['beta_Q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['beta_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['beta_Q'], shape=(self.numNodes,1))) if self.parameters['beta_Q'] is not None else self.beta + self.sigma_Q = (numpy.array(self.parameters['sigma_Q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['sigma_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['sigma_Q'], shape=(self.numNodes,1))) if self.parameters['sigma_Q'] is not None else self.sigma + self.lamda_Q = (numpy.array(self.parameters['lamda_Q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['lamda_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['lamda_Q'], shape=(self.numNodes,1))) if self.parameters['lamda_Q'] is not None else self.lamda + self.gamma_Q_sym = (numpy.array(self.parameters['gamma_Q_sym']).reshape((self.numNodes, 1)) if isinstance(self.parameters['gamma_Q_sym'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['gamma_Q_sym'], shape=(self.numNodes,1))) if self.parameters['gamma_Q_sym'] is not None else self.gamma + self.gamma_Q_asym = (numpy.array(self.parameters['gamma_Q_asym']).reshape((self.numNodes, 1)) if isinstance(self.parameters['gamma_Q_asym'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['gamma_Q_asym'], shape=(self.numNodes,1))) if self.parameters['gamma_Q_asym'] is not None else self.gamma + self.eta_Q = (numpy.array(self.parameters['eta_Q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['eta_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['eta_Q'], shape=(self.numNodes,1))) if self.parameters['eta_Q'] is not None else self.eta + self.alpha_Q = (numpy.array(self.parameters['alpha_Q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['alpha_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['alpha_Q'], shape=(self.numNodes,1))) if self.parameters['alpha_Q'] is not None else self.alpha + self.theta_S = numpy.array(self.parameters['theta_S']).reshape((self.numNodes, 1)) if isinstance(self.parameters['theta_S'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['theta_S'], shape=(self.numNodes,1)) + self.theta_E = numpy.array(self.parameters['theta_E']).reshape((self.numNodes, 1)) if isinstance(self.parameters['theta_E'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['theta_E'], shape=(self.numNodes,1)) + self.theta_pre = numpy.array(self.parameters['theta_pre']).reshape((self.numNodes, 1)) if isinstance(self.parameters['theta_pre'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['theta_pre'], shape=(self.numNodes,1)) + self.theta_sym = numpy.array(self.parameters['theta_sym']).reshape((self.numNodes, 1)) if isinstance(self.parameters['theta_sym'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['theta_sym'], shape=(self.numNodes,1)) + self.theta_asym = numpy.array(self.parameters['theta_asym']).reshape((self.numNodes, 1)) if isinstance(self.parameters['theta_asym'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['theta_asym'], shape=(self.numNodes,1)) + self.phi_S = numpy.array(self.parameters['phi_S']).reshape((self.numNodes, 1)) if isinstance(self.parameters['phi_S'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['phi_S'], shape=(self.numNodes,1)) + self.phi_E = numpy.array(self.parameters['phi_E']).reshape((self.numNodes, 1)) if isinstance(self.parameters['phi_E'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['phi_E'], shape=(self.numNodes,1)) + self.phi_pre = numpy.array(self.parameters['phi_pre']).reshape((self.numNodes, 1)) if isinstance(self.parameters['phi_pre'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['phi_pre'], shape=(self.numNodes,1)) + self.phi_sym = numpy.array(self.parameters['phi_sym']).reshape((self.numNodes, 1)) if isinstance(self.parameters['phi_sym'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['phi_sym'], shape=(self.numNodes,1)) + self.phi_asym = numpy.array(self.parameters['phi_asym']).reshape((self.numNodes, 1)) if isinstance(self.parameters['phi_asym'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['phi_asym'], shape=(self.numNodes,1)) + self.psi_S = numpy.array(self.parameters['psi_S']).reshape((self.numNodes, 1)) if isinstance(self.parameters['psi_S'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['psi_S'], shape=(self.numNodes,1)) + self.psi_E = numpy.array(self.parameters['psi_E']).reshape((self.numNodes, 1)) if isinstance(self.parameters['psi_E'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['psi_E'], shape=(self.numNodes,1)) + self.psi_pre = numpy.array(self.parameters['psi_pre']).reshape((self.numNodes, 1)) if isinstance(self.parameters['psi_pre'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['psi_pre'], shape=(self.numNodes,1)) + self.psi_sym = numpy.array(self.parameters['psi_sym']).reshape((self.numNodes, 1)) if isinstance(self.parameters['psi_sym'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['psi_sym'], shape=(self.numNodes,1)) + self.psi_asym = numpy.array(self.parameters['psi_asym']).reshape((self.numNodes, 1)) if isinstance(self.parameters['psi_asym'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['psi_asym'], shape=(self.numNodes,1)) + self.q = numpy.array(self.parameters['q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['q'], shape=(self.numNodes,1)) + + #---------------------------------------- + + self.beta_pairwise_mode = self.parameters['beta_pairwise_mode'] + + #---------------------------------------- + # Global transmission parameters: + #---------------------------------------- + if (self.beta_pairwise_mode == 'infected' or self.beta_pairwise_mode is None): + self.beta_global = numpy.full_like(self.beta, fill_value=numpy.mean(self.beta)) + self.beta_Q_global = numpy.full_like(self.beta_Q, fill_value=numpy.mean(self.beta_Q)) + self.beta_asym_global = numpy.full_like(self.beta_asym, fill_value=numpy.mean(self.beta_asym)) + elif self.beta_pairwise_mode == 'infectee': + self.beta_global = self.beta + self.beta_Q_global = self.beta_Q + self.beta_asym_global = self.beta_asym + elif self.beta_pairwise_mode == 'min': + self.beta_global = numpy.minimum(self.beta, numpy.mean(beta)) + self.beta_Q_global = numpy.minimum(self.beta_Q, numpy.mean(beta_Q)) + self.beta_asym_global = numpy.minimum(self.beta_asym, numpy.mean(beta_asym)) + elif self.beta_pairwise_mode == 'max': + self.beta_global = numpy.maximum(self.beta, numpy.mean(beta)) + self.beta_Q_global = numpy.maximum(self.beta_Q, numpy.mean(beta_Q)) + self.beta_asym_global = numpy.maximum(self.beta_asym, numpy.mean(beta_asym)) + elif self.beta_pairwise_mode == 'mean': + self.beta_global = (self.beta + numpy.full_like(self.beta, fill_value=numpy.mean(self.beta)))/2 + self.beta_Q_global = (self.beta_Q + numpy.full_like(self.beta_Q, fill_value=numpy.mean(self.beta_Q)))/2 + self.beta_asym_global = (self.beta_asym + numpy.full_like(self.beta_asym, fill_value=numpy.mean(self.beta_asym)))/2 + + #---------------------------------------- + # Local transmission parameters: + #---------------------------------------- + self.beta_local = self.beta if self.parameters['beta_local'] is None else numpy.array(self.parameters['beta_local']) if isinstance(self.parameters['beta_local'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['beta_local'], shape=(self.numNodes,1)) + self.beta_Q_local = self.beta_Q if self.parameters['beta_Q_local'] is None else numpy.array(self.parameters['beta_Q_local']) if isinstance(self.parameters['beta_Q_local'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['beta_Q_local'], shape=(self.numNodes,1)) + self.beta_asym_local = None if self.parameters['beta_asym_local'] is None else numpy.array(self.parameters['beta_asym_local']) if isinstance(self.parameters['beta_asym_local'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['beta_asym_local'], shape=(self.numNodes,1)) + #---------------------------------------- + if (self.beta_local.ndim == 2 and self.beta_local.shape[0] == self.numNodes and self.beta_local.shape[1] == self.numNodes): + self.A_beta_pairwise = self.beta_local + elif ((self.beta_local.ndim == 1 and self.beta_local.shape[0] == self.numNodes) or (self.beta_local.ndim == 2 and (self.beta_local.shape[0] == self.numNodes or self.beta_local.shape[1] == self.numNodes))): + self.beta_local = self.beta_local.reshape((self.numNodes,1)) + # Pre-multiply beta values by the adjacency matrix ("transmission weight connections") + A_beta_pairwise_byInfected = scipy.sparse.csr_matrix.multiply(self.A, self.beta_local.T).tocsr() + A_beta_pairwise_byInfectee = scipy.sparse.csr_matrix.multiply(self.A, self.beta_local).tocsr() + #------------------------------ + # Compute the effective pairwise beta values as a function of the infected/infectee pair: + if self.beta_pairwise_mode == 'infected': + self.A_beta_pairwise = A_beta_pairwise_byInfected + elif self.beta_pairwise_mode == 'infectee': + self.A_beta_pairwise = A_beta_pairwise_byInfectee + elif self.beta_pairwise_mode == 'min': + self.A_beta_pairwise = scipy.sparse.csr_matrix.minimum(A_beta_pairwise_byInfected, A_beta_pairwise_byInfectee) + elif self.beta_pairwise_mode == 'max': + self.A_beta_pairwise = scipy.sparse.csr_matrix.maximum(A_beta_pairwise_byInfected, A_beta_pairwise_byInfectee) + elif (self.beta_pairwise_mode == 'mean' or self.beta_pairwise_mode is None): + self.A_beta_pairwise = (A_beta_pairwise_byInfected + A_beta_pairwise_byInfectee)/2 + else: + print("Unrecognized beta_pairwise_mode value (support for 'infected', 'infectee', 'min', 'max', and 'mean').") + else: + print("Invalid values given for beta_local (expected 1xN list/array or NxN 2d array)") + #---------------------------------------- + if (self.beta_Q_local.ndim == 2 and self.beta_Q_local.shape[0] == self.numNodes and self.beta_Q_local.shape[1] == self.numNodes): + self.A_Q_beta_Q_pairwise = self.beta_Q_local + elif ((self.beta_Q_local.ndim == 1 and self.beta_Q_local.shape[0] == self.numNodes) or (self.beta_Q_local.ndim == 2 and (self.beta_Q_local.shape[0] == self.numNodes or self.beta_Q_local.shape[1] == self.numNodes))): + self.beta_Q_local = self.beta_Q_local.reshape((self.numNodes,1)) + # Pre-multiply beta_Q values by the isolation adjacency matrix ("transmission weight connections") + A_Q_beta_Q_pairwise_byInfected = scipy.sparse.csr_matrix.multiply(self.A_Q, self.beta_Q_local.T).tocsr() + A_Q_beta_Q_pairwise_byInfectee = scipy.sparse.csr_matrix.multiply(self.A_Q, self.beta_Q_local).tocsr() + #------------------------------ + # Compute the effective pairwise beta values as a function of the infected/infectee pair: + if self.beta_pairwise_mode == 'infected': + self.A_Q_beta_Q_pairwise = A_Q_beta_Q_pairwise_byInfected + elif self.beta_pairwise_mode == 'infectee': + self.A_Q_beta_Q_pairwise = A_Q_beta_Q_pairwise_byInfectee + elif self.beta_pairwise_mode == 'min': + self.A_Q_beta_Q_pairwise = scipy.sparse.csr_matrix.minimum(A_Q_beta_Q_pairwise_byInfected, A_Q_beta_Q_pairwise_byInfectee) + elif self.beta_pairwise_mode == 'max': + self.A_Q_beta_Q_pairwise = scipy.sparse.csr_matrix.maximum(A_Q_beta_Q_pairwise_byInfected, A_Q_beta_Q_pairwise_byInfectee) + elif (self.beta_pairwise_mode == 'mean' or self.beta_pairwise_mode is None): + self.A_Q_beta_Q_pairwise = (A_Q_beta_Q_pairwise_byInfected + A_Q_beta_Q_pairwise_byInfectee)/2 + else: + print("Unrecognized beta_pairwise_mode value (support for 'infected', 'infectee', 'min', 'max', and 'mean').") + else: + print("Invalid values given for beta_Q_local (expected 1xN list/array or NxN 2d array)") + #---------------------------------------- + if self.beta_asym_local is None: + self.A_beta_asym_pairwise = None + elif (self.beta_asym_local.ndim == 2 and self.beta_asym_local.shape[0] == self.numNodes and self.beta_asym_local.shape[1] == self.numNodes): + self.A_beta_asym_pairwise = self.beta_asym_local + elif ((self.beta_asym_local.ndim == 1 and self.beta_asym_local.shape[0] == self.numNodes) or (self.beta_asym_local.ndim == 2 and (self.beta_asym_local.shape[0] == self.numNodes or self.beta_asym_local.shape[1] == self.numNodes))): + self.beta_asym_local = self.beta_asym_local.reshape((self.numNodes,1)) + # Pre-multiply beta_asym values by the adjacency matrix ("transmission weight connections") + A_beta_asym_pairwise_byInfected = scipy.sparse.csr_matrix.multiply(self.A, self.beta_asym_local.T).tocsr() + A_beta_asym_pairwise_byInfectee = scipy.sparse.csr_matrix.multiply(self.A, self.beta_asym_local).tocsr() + #------------------------------ + # Compute the effective pairwise beta values as a function of the infected/infectee pair: + if self.beta_pairwise_mode == 'infected': + self.A_beta_asym_pairwise = A_beta_asym_pairwise_byInfected + elif self.beta_pairwise_mode == 'infectee': + self.A_beta_asym_pairwise = A_beta_asym_pairwise_byInfectee + elif self.beta_pairwise_mode == 'min': + self.A_beta_asym_pairwise = scipy.sparse.csr_matrix.minimum(A_beta_asym_pairwise_byInfected, A_beta_asym_pairwise_byInfectee) + elif self.beta_pairwise_mode == 'max': + self.A_beta_asym_pairwise = scipy.sparse.csr_matrix.maximum(A_beta_asym_pairwise_byInfected, A_beta_asym_pairwise_byInfectee) + elif (self.beta_pairwise_mode == 'mean' or self.beta_pairwise_mode is None): + self.A_beta_asym_pairwise = (A_beta_asym_pairwise_byInfected + A_beta_asym_pairwise_byInfectee)/2 + else: + print("Unrecognized beta_pairwise_mode value (support for 'infected', 'infectee', 'min', 'max', and 'mean').") + else: + print("Invalid values given for beta_asym_local (expected 1xN list/array or NxN 2d array)") + + #---------------------------------------- + # Degree-based transmission scaling parameters: + #---------------------------------------- + self.delta_pairwise_mode = self.parameters['delta_pairwise_mode'] + with numpy.errstate(divide='ignore'): # ignore log(0) warning, then convert log(0) = -inf -> 0.0 + self.delta = numpy.log(self.degree)/numpy.log(numpy.mean(self.degree)) if self.parameters['delta'] is None else numpy.array(self.parameters['delta']) if isinstance(self.parameters['delta'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['delta'], shape=(self.numNodes,1)) + self.delta_Q = numpy.log(self.degree_Q)/numpy.log(numpy.mean(self.degree_Q)) if self.parameters['delta_Q'] is None else numpy.array(self.parameters['delta_Q']) if isinstance(self.parameters['delta_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['delta_Q'], shape=(self.numNodes,1)) + self.delta[numpy.isneginf(self.delta)] = 0.0 + self.delta_Q[numpy.isneginf(self.delta_Q)] = 0.0 + #---------------------------------------- + if (self.delta.ndim == 2 and self.delta.shape[0] == self.numNodes and self.delta.shape[1] == self.numNodes): + self.A_delta_pairwise = self.delta + elif ((self.delta.ndim == 1 and self.delta.shape[0] == self.numNodes) or (self.delta.ndim == 2 and (self.delta.shape[0] == self.numNodes or self.delta.shape[1] == self.numNodes))): + self.delta = self.delta.reshape((self.numNodes,1)) + # Pre-multiply delta values by the adjacency matrix ("transmission weight connections") + A_delta_pairwise_byInfected = scipy.sparse.csr_matrix.multiply(self.A, self.delta.T).tocsr() + A_delta_pairwise_byInfectee = scipy.sparse.csr_matrix.multiply(self.A, self.delta).tocsr() + #------------------------------ + # Compute the effective pairwise delta values as a function of the infected/infectee pair: + if self.delta_pairwise_mode == 'infected': + self.A_delta_pairwise = A_delta_pairwise_byInfected + elif self.delta_pairwise_mode == 'infectee': + self.A_delta_pairwise = A_delta_pairwise_byInfectee + elif self.delta_pairwise_mode == 'min': + self.A_delta_pairwise = scipy.sparse.csr_matrix.minimum(A_delta_pairwise_byInfected, A_delta_pairwise_byInfectee) + elif self.delta_pairwise_mode == 'max': + self.A_delta_pairwise = scipy.sparse.csr_matrix.maximum(A_delta_pairwise_byInfected, A_delta_pairwise_byInfectee) + elif self.delta_pairwise_mode == 'mean': + self.A_delta_pairwise = (A_delta_pairwise_byInfected + A_delta_pairwise_byInfectee)/2 + elif self.delta_pairwise_mode is None: + self.A_delta_pairwise = self.A + else: + print("Unrecognized delta_pairwise_mode value (support for 'infected', 'infectee', 'min', 'max', and 'mean').") + else: + print("Invalid values given for delta (expected 1xN list/array or NxN 2d array)") + #---------------------------------------- + if (self.delta_Q.ndim == 2 and self.delta_Q.shape[0] == self.numNodes and self.delta_Q.shape[1] == self.numNodes): + self.A_Q_delta_Q_pairwise = self.delta_Q + elif ((self.delta_Q.ndim == 1 and self.delta_Q.shape[0] == self.numNodes) or (self.delta_Q.ndim == 2 and (self.delta_Q.shape[0] == self.numNodes or self.delta_Q.shape[1] == self.numNodes))): + self.delta_Q = self.delta_Q.reshape((self.numNodes,1)) + # Pre-multiply delta_Q values by the isolation adjacency matrix ("transmission weight connections") + A_Q_delta_Q_pairwise_byInfected = scipy.sparse.csr_matrix.multiply(self.A_Q, self.delta_Q).tocsr() + A_Q_delta_Q_pairwise_byInfectee = scipy.sparse.csr_matrix.multiply(self.A_Q, self.delta_Q.T).tocsr() + #------------------------------ + # Compute the effective pairwise delta values as a function of the infected/infectee pair: + if self.delta_pairwise_mode == 'infected': + self.A_Q_delta_Q_pairwise = A_Q_delta_Q_pairwise_byInfected + elif self.delta_pairwise_mode == 'infectee': + self.A_Q_delta_Q_pairwise = A_Q_delta_Q_pairwise_byInfectee + elif self.delta_pairwise_mode == 'min': + self.A_Q_delta_Q_pairwise = scipy.sparse.csr_matrix.minimum(A_Q_delta_Q_pairwise_byInfected, A_Q_delta_Q_pairwise_byInfectee) + elif self.delta_pairwise_mode == 'max': + self.A_Q_delta_Q_pairwise = scipy.sparse.csr_matrix.maximum(A_Q_delta_Q_pairwise_byInfected, A_Q_delta_Q_pairwise_byInfectee) + elif self.delta_pairwise_mode == 'mean': + self.A_Q_delta_Q_pairwise = (A_Q_delta_Q_pairwise_byInfected + A_Q_delta_Q_pairwise_byInfectee)/2 + elif self.delta_pairwise_mode is None: + self.A_Q_delta_Q_pairwise = self.A + else: + print("Unrecognized delta_pairwise_mode value (support for 'infected', 'infectee', 'min', 'max', and 'mean').") + else: + print("Invalid values given for delta_Q (expected 1xN list/array or NxN 2d array)") + + #---------------------------------------- + # Pre-calculate the pairwise delta*beta values: + #---------------------------------------- + self.A_deltabeta = scipy.sparse.csr_matrix.multiply(self.A_delta_pairwise, self.A_beta_pairwise) + self.A_Q_deltabeta_Q = scipy.sparse.csr_matrix.multiply(self.A_Q_delta_Q_pairwise, self.A_Q_beta_Q_pairwise) + if self.A_beta_asym_pairwise is not None: + self.A_deltabeta_asym = scipy.sparse.csr_matrix.multiply(self.A_delta_pairwise, self.A_beta_asym_pairwise) + else: + self.A_deltabeta_asym = None + + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def node_degrees(self, Amat): + return Amat.sum(axis=0).reshape(self.numNodes,1) # sums of adj matrix cols + + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def total_num_susceptible(self, t_idx=None): + if t_idx is None: + return self.numS[:] + self.numQ_S[:] + else: + return self.numS[t_idx] + self.numQ_S[t_idx] + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def total_num_infected(self, t_idx=None): + if t_idx is None: + return (self.numE[:] + self.numI_pre[:] + self.numI_sym[:] + self.numI_asym[:] + self.numH[:] + + self.numQ_E[:] + self.numQ_pre[:] + self.numQ_sym[:] + self.numQ_asym[:]) + else: + return (self.numE[t_idx] + self.numI_pre[t_idx] + self.numI_sym[t_idx] + self.numI_asym[t_idx] + self.numH[t_idx] + + self.numQ_E[t_idx] + self.numQ_pre[t_idx] + self.numQ_sym[t_idx] + self.numQ_asym[t_idx]) + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def total_num_isolated(self, t_idx=None): + if t_idx is None: + return self.numQ_S[:] + self.numQ_E[:] + self.numQ_pre[:] + self.numQ_sym[:] + self.numQ_asym[:] + self.numQ_R[:] + else: + return self.numQ_S[t_idx] + self.numQ_E[t_idx] + self.numQ_pre[t_idx] + self.numQ_sym[t_idx] + self.numQ_asym[t_idx] + self.numQ_R[t_idx] + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def total_num_tested(self, t_idx=None): + if t_idx is None: + return self.numTested[:] + else: + return self.numTested[t_idx] + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def total_num_positive(self, t_idx=None): + if t_idx is None: + return self.numPositive[:] + else: + return self.numPositive[t_idx] + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def total_num_recovered(self, t_idx=None): + if t_idx is None: + return self.numR[:] + self.numQ_R[:] + else: + return self.numR[t_idx] + self.numQ_R[t_idx] + + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def calc_propensities(self): + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Pre-calculate matrix multiplication terms that may be used in multiple propensity calculations, + # and check to see if their computation is necessary before doing the multiplication + #------------------------------------ + + self.transmissionTerms_I = numpy.zeros(shape=(self.numNodes,1)) + if (numpy.any(self.numI_sym[self.tidx]) or numpy.any(self.numI_asym[self.tidx]) or numpy.any(self.numI_pre[self.tidx])): + if self.A_deltabeta_asym is not None: + self.transmissionTerms_sym = numpy.asarray(scipy.sparse.csr_matrix.dot(self.A_deltabeta, self.X==self.I_sym)) + self.transmissionTerms_asym = numpy.asarray(scipy.sparse.csr_matrix.dot(self.A_deltabeta_asym, ((self.X==self.I_pre)|(self.X==self.I_asym)))) + self.transmissionTerms_I = self.transmissionTerms_sym+self.transmissionTerms_asym + else: + self.transmissionTerms_I = numpy.asarray(scipy.sparse.csr_matrix.dot(self.A_deltabeta, ((self.X==self.I_sym)|(self.X==self.I_pre)|(self.X==self.I_asym)))) + + #------------------------------------ + + self.transmissionTerms_Q = numpy.zeros(shape=(self.numNodes,1)) + if (numpy.any(self.numQ_pre[self.tidx]) or numpy.any(self.numQ_sym[self.tidx]) or numpy.any(self.numQ_asym[self.tidx])): + self.transmissionTerms_Q = numpy.asarray(scipy.sparse.csr_matrix.dot(self.A_Q_deltabeta_Q, ((self.X==self.Q_pre)|(self.X==self.Q_sym)|(self.X==self.Q_asym)))) + + #------------------------------------ + + self.transmissionTerms_IQ = numpy.zeros(shape=(self.numNodes,1)) + if (numpy.any(self.numQ_S[self.tidx]) and (numpy.any(self.numI_sym[self.tidx]) or numpy.any(self.numI_asym[self.tidx]) or numpy.any(self.numI_pre[self.tidx]))): + self.transmissionTerms_IQ = numpy.asarray(scipy.sparse.csr_matrix.dot(self.A_Q_deltabeta_Q, ((self.X==self.I_sym)|(self.X==self.I_pre)|(self.X==self.I_asym)))) + + #------------------------------------ + + numContacts_Q = numpy.zeros(shape=(self.numNodes,1)) + if (numpy.any(self.positive) and (numpy.any(self.phi_S) or numpy.any(self.phi_E) or numpy.any(self.phi_pre) or numpy.any(self.phi_sym) or numpy.any(self.phi_asym))): + numContacts_Q = numpy.asarray(scipy.sparse.csr_matrix.dot(self.A, ((self.positive)&(self.X!=self.R)&(self.X!=self.Q_R)&(self.X!=self.F)))) + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + propensities_StoE = ( self.alpha * + (self.o*(self.beta_global*self.prevalence_ext) + + (1-self.o)*( + self.p*((self.beta_global*self.numI_sym[self.tidx] + self.beta_asym_global*(self.numI_pre[self.tidx] + self.numI_asym[self.tidx]) + + self.q*self.beta_Q_global*(self.numQ_pre[self.tidx] + self.numQ_sym[self.tidx] + self.numQ_asym[self.tidx]))/self.N[self.tidx]) + + (1-self.p)*(numpy.divide(self.transmissionTerms_I, self.degree, out=numpy.zeros_like(self.degree), where=self.degree!=0) + + numpy.divide(self.transmissionTerms_Q, self.degree_Q, out=numpy.zeros_like(self.degree_Q), where=self.degree_Q!=0)))) + )*(self.X==self.S) + + propensities_QStoQE = numpy.zeros_like(propensities_StoE) + if numpy.any(self.X==self.Q_S): + propensities_QStoQE = ( self.alpha_Q * + (self.o*(self.q*self.beta_global*self.prevalence_ext) + + (1-self.o)*( + self.p*(self.q*(self.beta_global*self.numI_sym[self.tidx] + self.beta_asym_global*(self.numI_pre[self.tidx] + self.numI_asym[self.tidx]) + + self.beta_Q_global*(self.numQ_pre[self.tidx] + self.numQ_sym[self.tidx] + self.numQ_asym[self.tidx]))/self.N[self.tidx]) + + (1-self.p)*(numpy.divide(self.transmissionTerms_IQ+self.transmissionTerms_Q, self.degree_Q, out=numpy.zeros_like(self.degree_Q), where=self.degree_Q!=0)))) + )*(self.X==self.Q_S) + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + if self.transition_mode == 'time_in_state': + + propensities_EtoIPRE = 1e5 * ((self.X==self.E) & numpy.greater(self.timer_state, 1/self.sigma)) + + propensities_IPREtoISYM = 1e5 * ((self.X==self.I_pre) & numpy.greater(self.timer_state, 1/self.lamda) & numpy.greater_equal(self.rand_a, self.a)) + + propensities_IPREtoIASYM = 1e5 * ((self.X==self.I_pre) & numpy.greater(self.timer_state, 1/self.lamda) & numpy.less(self.rand_a, self.a)) + + propensities_ISYMtoR = 1e5 * ((self.X==self.I_sym) & numpy.greater(self.timer_state, 1/self.gamma) & numpy.greater_equal(self.rand_h, self.h)) + + propensities_ISYMtoH = 1e5 * ((self.X==self.I_sym) & numpy.greater(self.timer_state, 1/self.eta) & numpy.less(self.rand_h, self.h)) + + propensities_IASYMtoR = 1e5 * ((self.X==self.I_asym) & numpy.greater(self.timer_state, 1/self.gamma)) + + propensities_HtoR = 1e5 * ((self.X==self.H) & numpy.greater(self.timer_state, 1/self.gamma_H) & numpy.greater_equal(self.rand_f, self.f)) + + propensities_HtoF = 1e5 * ((self.X==self.H) & numpy.greater(self.timer_state, 1/self.mu_H) & numpy.less(self.rand_f, self.f)) + + propensities_StoQS = numpy.zeros_like(propensities_StoE) + + propensities_EtoQE = numpy.zeros_like(propensities_StoE) + + propensities_IPREtoQPRE = numpy.zeros_like(propensities_StoE) + + propensities_ISYMtoQSYM = numpy.zeros_like(propensities_StoE) + + propensities_IASYMtoQASYM = numpy.zeros_like(propensities_StoE) + + propensities_QEtoQPRE = 1e5 * ((self.X==self.Q_E) & numpy.greater(self.timer_state, 1/self.sigma_Q)) + + propensities_QPREtoQSYM = 1e5 * ((self.X==self.Q_pre) & numpy.greater(self.timer_state, 1/self.lamda_Q) & numpy.greater_equal(self.rand_a, self.a)) + + propensities_QPREtoQASYM = 1e5 * ((self.X==self.Q_pre) & numpy.greater(self.timer_state, 1/self.lamda_Q) & numpy.less(self.rand_a, self.a)) + + propensities_QSYMtoQR = 1e5 * ((self.X==self.Q_sym) & numpy.greater(self.timer_state, 1/self.gamma_Q_sym) & numpy.greater_equal(self.rand_h, self.h)) + + propensities_QSYMtoH = 1e5 * ((self.X==self.Q_sym) & numpy.greater(self.timer_state, 1/self.eta_Q) & numpy.less(self.rand_h, self.h)) + + propensities_QASYMtoQR = 1e5 * ((self.X==self.Q_asym) & numpy.greater(self.timer_state, 1/self.gamma_Q_asym)) + + propensities_RtoS = 1e5 * ((self.X==self.R) & numpy.greater(self.timer_state, 1/self.xi)) + + propensities__toS = 1e5 * ((self.X!=self.F) & numpy.greater(self.timer_state, 1/self.nu)) + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + else: # exponential_rates + + propensities_EtoIPRE = self.sigma * (self.X==self.E) + + propensities_IPREtoISYM = self.lamda * ((self.X==self.I_pre) & (numpy.greater_equal(self.rand_a, self.a))) + + propensities_IPREtoIASYM = self.lamda * ((self.X==self.I_pre) & (numpy.less(self.rand_a, self.a))) + + propensities_ISYMtoR = self.gamma * ((self.X==self.I_sym) & (numpy.greater_equal(self.rand_h, self.h))) + + propensities_ISYMtoH = self.eta * ((self.X==self.I_sym) & (numpy.less(self.rand_h, self.h))) + + propensities_IASYMtoR = self.gamma_asym * (self.X==self.I_asym) + + propensities_HtoR = self.gamma_H * ((self.X==self.H) & (numpy.greater_equal(self.rand_f, self.f))) + + propensities_HtoF = self.mu_H * ((self.X==self.H) & (numpy.less(self.rand_f, self.f))) + + propensities_StoQS = (self.theta_S + self.phi_S*numContacts_Q)*self.psi_S * (self.X==self.S) + + propensities_EtoQE = (self.theta_E + self.phi_E*numContacts_Q)*self.psi_E * (self.X==self.E) + + propensities_IPREtoQPRE = (self.theta_pre + self.phi_pre*numContacts_Q)*self.psi_pre * (self.X==self.I_pre) + + propensities_ISYMtoQSYM = (self.theta_sym + self.phi_sym*numContacts_Q)*self.psi_sym * (self.X==self.I_sym) + + propensities_IASYMtoQASYM = (self.theta_asym + self.phi_asym*numContacts_Q)*self.psi_asym * (self.X==self.I_asym) + + propensities_QEtoQPRE = self.sigma_Q * (self.X==self.Q_E) + + propensities_QPREtoQSYM = self.lamda_Q * ((self.X==self.Q_pre) & (numpy.greater_equal(self.rand_a, self.a))) + + propensities_QPREtoQASYM = self.lamda_Q * ((self.X==self.Q_pre) & (numpy.less(self.rand_a, self.a))) + + propensities_QSYMtoQR = self.gamma_Q_sym * ((self.X==self.Q_sym) & (numpy.greater_equal(self.rand_h, self.h))) + + propensities_QSYMtoH = self.eta_Q * ((self.X==self.Q_sym) & (numpy.less(self.rand_h, self.h))) + + propensities_QASYMtoQR = self.gamma_Q_asym * (self.X==self.Q_asym) + + propensities_RtoS = self.xi * (self.X==self.R) + + propensities__toS = self.nu * (self.X!=self.F) + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + propensities = numpy.hstack([propensities_StoE, propensities_EtoIPRE, propensities_IPREtoISYM, propensities_IPREtoIASYM, + propensities_ISYMtoR, propensities_ISYMtoH, propensities_IASYMtoR, propensities_HtoR, propensities_HtoF, + propensities_StoQS, propensities_EtoQE, propensities_IPREtoQPRE, propensities_ISYMtoQSYM, propensities_IASYMtoQASYM, + propensities_QStoQE, propensities_QEtoQPRE, propensities_QPREtoQSYM, propensities_QPREtoQASYM, + propensities_QSYMtoQR, propensities_QSYMtoH, propensities_QASYMtoQR, propensities_RtoS, propensities__toS]) + + columns = [ 'StoE', 'EtoIPRE', 'IPREtoISYM', 'IPREtoIASYM', + 'ISYMtoR', 'ISYMtoH', 'IASYMtoR', 'HtoR', 'HtoF', + 'StoQS', 'EtoQE', 'IPREtoQPRE', 'ISYMtoQSYM', 'IASYMtoQASYM', + 'QStoQE', 'QEtoQPRE', 'QPREtoQSYM', 'QPREtoQASYM', + 'QSYMtoQR', 'QSYMtoH', 'QASYMtoQR', 'RtoS', '_toS' ] + + return propensities, columns + + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def set_isolation(self, node, isolate): + # Move this node in/out of the appropriate isolation state: + if isolate == True: + if self.X[node] == self.S: + self.X[node] = self.Q_S + elif self.X[node] == self.E: + self.X[node] = self.Q_E + elif self.X[node] == self.I_pre: + self.X[node] = self.Q_pre + elif self.X[node] == self.I_sym: + self.X[node] = self.Q_sym + elif self.X[node] == self.I_asym: + self.X[node] = self.Q_asym + elif self.X[node] == self.R: + self.X[node] = self.Q_R + elif isolate == False: + if self.X[node] == self.Q_S: + self.X[node] = self.S + elif self.X[node] == self.Q_E: + self.X[node] = self.E + elif self.X[node] == self.Q_pre: + self.X[node] = self.I_pre + elif self.X[node] == self.Q_sym: + self.X[node] = self.I_sym + elif self.X[node] == self.Q_asym: + self.X[node] = self.I_asym + elif self.X[node] == self.Q_R: + self.X[node] = self.R + # Reset the isolation timer: + self.timer_isolation[node] = 0 + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def set_tested(self, node, tested): + self.tested[node] = tested + self.testedInCurrentState[node] = tested + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def set_positive(self, node, positive): + self.positive[node] = positive + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def introduce_exposures(self, num_new_exposures): + exposedNodes = numpy.random.choice(range(self.numNodes), size=num_new_exposures, replace=False) + for exposedNode in exposedNodes: + if self.X[exposedNode] == self.S: + self.X[exposedNode] = self.E + elif self.X[exposedNode] == self.Q_S: + self.X[exposedNode] = self.Q_E + + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def increase_data_series_length(self): + self.tseries = numpy.pad(self.tseries, [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.numS = numpy.pad(self.numS, [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.numE = numpy.pad(self.numE, [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.numI_pre = numpy.pad(self.numI_pre, [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.numI_sym = numpy.pad(self.numI_sym, [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.numI_asym = numpy.pad(self.numI_asym, [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.numH = numpy.pad(self.numH, [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.numR = numpy.pad(self.numR, [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.numF = numpy.pad(self.numF, [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.numQ_S = numpy.pad(self.numQ_S, [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.numQ_E = numpy.pad(self.numQ_E, [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.numQ_pre = numpy.pad(self.numQ_pre, [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.numQ_sym = numpy.pad(self.numQ_sym, [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.numQ_asym = numpy.pad(self.numQ_asym, [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.numQ_R = numpy.pad(self.numQ_R, [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.N = numpy.pad(self.N, [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.numTested = numpy.pad(self.numTested, [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.numPositive = numpy.pad(self.numPositive, [(0, 6*self.numNodes)], mode='constant', constant_values=0) + + if self.store_Xseries: + self.Xseries = numpy.pad(self.Xseries, [(0, 6*self.numNodes), (0,0)], mode='constant', constant_values=0) + + if self.nodeGroupData: + for groupName in self.nodeGroupData: + self.nodeGroupData[groupName]['numS'] = numpy.pad(self.nodeGroupData[groupName]['numS'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.nodeGroupData[groupName]['numE'] = numpy.pad(self.nodeGroupData[groupName]['numE'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.nodeGroupData[groupName]['numI_pre'] = numpy.pad(self.nodeGroupData[groupName]['numI_pre'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.nodeGroupData[groupName]['numI_sym'] = numpy.pad(self.nodeGroupData[groupName]['numI_sym'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.nodeGroupData[groupName]['numI_asym'] = numpy.pad(self.nodeGroupData[groupName]['numI_asym'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.nodeGroupData[groupName]['numH'] = numpy.pad(self.nodeGroupData[groupName]['numH'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.nodeGroupData[groupName]['numR'] = numpy.pad(self.nodeGroupData[groupName]['numR'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.nodeGroupData[groupName]['numF'] = numpy.pad(self.nodeGroupData[groupName]['numF'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.nodeGroupData[groupName]['numQ_S'] = numpy.pad(self.nodeGroupData[groupName]['numQ_S'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.nodeGroupData[groupName]['numQ_E'] = numpy.pad(self.nodeGroupData[groupName]['numQ_E'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.nodeGroupData[groupName]['numQ_pre'] = numpy.pad(self.nodeGroupData[groupName]['numQ_pre'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.nodeGroupData[groupName]['numQ_sym'] = numpy.pad(self.nodeGroupData[groupName]['numQ_sym'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.nodeGroupData[groupName]['numQ_asym'] = numpy.pad(self.nodeGroupData[groupName]['numQ_asym'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.nodeGroupData[groupName]['numQ_R'] = numpy.pad(self.nodeGroupData[groupName]['numQ_R'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.nodeGroupData[groupName]['N'] = numpy.pad(self.nodeGroupData[groupName]['N'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.nodeGroupData[groupName]['numTested'] = numpy.pad(self.nodeGroupData[groupName]['numTested'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.nodeGroupData[groupName]['numPositive'] = numpy.pad(self.nodeGroupData[groupName]['numPositive'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) + + return None + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def finalize_data_series(self): + self.tseries = numpy.array(self.tseries, dtype=float)[:self.tidx+1] + self.numS = numpy.array(self.numS, dtype=float)[:self.tidx+1] + self.numE = numpy.array(self.numE, dtype=float)[:self.tidx+1] + self.numI_pre = numpy.array(self.numI_pre, dtype=float)[:self.tidx+1] + self.numI_sym = numpy.array(self.numI_sym, dtype=float)[:self.tidx+1] + self.numI_asym = numpy.array(self.numI_asym, dtype=float)[:self.tidx+1] + self.numH = numpy.array(self.numH, dtype=float)[:self.tidx+1] + self.numR = numpy.array(self.numR, dtype=float)[:self.tidx+1] + self.numF = numpy.array(self.numF, dtype=float)[:self.tidx+1] + self.numQ_S = numpy.array(self.numQ_S, dtype=float)[:self.tidx+1] + self.numQ_E = numpy.array(self.numQ_E, dtype=float)[:self.tidx+1] + self.numQ_pre = numpy.array(self.numQ_pre, dtype=float)[:self.tidx+1] + self.numQ_sym = numpy.array(self.numQ_sym, dtype=float)[:self.tidx+1] + self.numQ_asym = numpy.array(self.numQ_asym, dtype=float)[:self.tidx+1] + self.numQ_R = numpy.array(self.numQ_R, dtype=float)[:self.tidx+1] + self.N = numpy.array(self.N, dtype=float)[:self.tidx+1] + self.numTested = numpy.array(self.numTested, dtype=float)[:self.tidx+1] + self.numPositive = numpy.array(self.numPositive, dtype=float)[:self.tidx+1] + + if self.store_Xseries: + self.Xseries = self.Xseries[:self.tidx+1, :] + + if self.nodeGroupData: + for groupName in self.nodeGroupData: + self.nodeGroupData[groupName]['numS'] = numpy.array(self.nodeGroupData[groupName]['numS'], dtype=float)[:self.tidx+1] + self.nodeGroupData[groupName]['numE'] = numpy.array(self.nodeGroupData[groupName]['numE'], dtype=float)[:self.tidx+1] + self.nodeGroupData[groupName]['numI_pre'] = numpy.array(self.nodeGroupData[groupName]['numI_pre'], dtype=float)[:self.tidx+1] + self.nodeGroupData[groupName]['numI_sym'] = numpy.array(self.nodeGroupData[groupName]['numI_sym'], dtype=float)[:self.tidx+1] + self.nodeGroupData[groupName]['numI_asym'] = numpy.array(self.nodeGroupData[groupName]['numI_asym'], dtype=float)[:self.tidx+1] + self.nodeGroupData[groupName]['numR'] = numpy.array(self.nodeGroupData[groupName]['numR'], dtype=float)[:self.tidx+1] + self.nodeGroupData[groupName]['numF'] = numpy.array(self.nodeGroupData[groupName]['numF'], dtype=float)[:self.tidx+1] + self.nodeGroupData[groupName]['numQ_S'] = numpy.array(self.nodeGroupData[groupName]['numQ_S'], dtype=float)[:self.tidx+1] + self.nodeGroupData[groupName]['numQ_E'] = numpy.array(self.nodeGroupData[groupName]['numQ_E'], dtype=float)[:self.tidx+1] + self.nodeGroupData[groupName]['numQ_pre'] = numpy.array(self.nodeGroupData[groupName]['numQ_pre'], dtype=float)[:self.tidx+1] + self.nodeGroupData[groupName]['numQ_sym'] = numpy.array(self.nodeGroupData[groupName]['numQ_sym'], dtype=float)[:self.tidx+1] + self.nodeGroupData[groupName]['numQ_asym'] = numpy.array(self.nodeGroupData[groupName]['numQ_asym'], dtype=float)[:self.tidx+1] + self.nodeGroupData[groupName]['numQ_R'] = numpy.array(self.nodeGroupData[groupName]['numQ_R'], dtype=float)[:self.tidx+1] + self.nodeGroupData[groupName]['N'] = numpy.array(self.nodeGroupData[groupName]['N'], dtype=float)[:self.tidx+1] + self.nodeGroupData[groupName]['numTested'] = numpy.array(self.nodeGroupData[groupName]['numTested'], dtype=float)[:self.tidx+1] + self.nodeGroupData[groupName]['numPositive'] = numpy.array(self.nodeGroupData[groupName]['numPositive'], dtype=float)[:self.tidx+1] + + return None + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def run_iteration(self, max_dt=None): + + max_dt = self.tmax if max_dt is None else max_dt + + if self.tidx >= len(self.tseries) - 1: + # Room has run out in the timeseries storage arrays; double the size of these arrays: + self.increase_data_series_length() + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Generate 2 random numbers uniformly distributed in (0,1) + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + r1 = numpy.random.rand() + r2 = numpy.random.rand() + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Calculate propensities + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + propensities, transitionTypes = self.calc_propensities() + + if propensities.sum() > 0: + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Calculate alpha + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + propensities_flat = propensities.ravel(order='F') + cumsum = propensities_flat.cumsum() + alpha = propensities_flat.sum() + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Compute the time until the next event takes place + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + tau = (1/alpha)*numpy.log(float(1/r1)) + + if tau > max_dt: + # If the time to next event exceeds the max allowed interval, + # advance the system time by the max allowed interval, + # but do not execute any events (recalculate Gillespie interval/event next iteration) + self.t += max_dt + self.timer_state += max_dt + # Update testing and isolation timers/statuses + isolatedNodes = numpy.argwhere((self.X==self.Q_S)|(self.X==self.Q_E)|(self.X==self.Q_pre)|(self.X==self.Q_sym)|(self.X==self.Q_asym)|(self.X==self.Q_R))[:,0].flatten() + self.timer_isolation[isolatedNodes] = self.timer_isolation[isolatedNodes] + max_dt + nodesExitingIsolation = numpy.argwhere(self.timer_isolation >= self.isolationTime) + for isoNode in nodesExitingIsolation: + self.set_isolation(node=isoNode, isolate=False) + # return without any further event execution + return True + else: + self.t += tau + self.timer_state += tau + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Compute which event takes place + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + transitionIdx = numpy.searchsorted(cumsum,r2*alpha) + transitionNode = transitionIdx % self.numNodes + transitionType = transitionTypes[ int(transitionIdx/self.numNodes) ] + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Perform updates triggered by rate propensities: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + assert(self.X[transitionNode] == self.transitions[transitionType]['currentState'] and self.X[transitionNode]!=self.F), "Assertion error: Node "+str(transitionNode)+" has unexpected current state "+str(self.X[transitionNode])+" given the intended transition of "+str(transitionType)+"." + self.X[transitionNode] = self.transitions[transitionType]['newState'] + + self.testedInCurrentState[transitionNode] = False + + self.timer_state[transitionNode] = 0.0 + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + # Save information about infection events when they occur: + if (transitionType == 'StoE' or transitionType == 'QStoQE'): + transitionNode_GNbrs = list(self.G[transitionNode].keys()) + transitionNode_GQNbrs = list(self.G_Q[transitionNode].keys()) + self.infectionsLog.append({ 't': self.t, + 'infected_node': transitionNode, + 'infection_type': transitionType, + 'infected_node_degree': self.degree[transitionNode], + 'local_contact_nodes': transitionNode_GNbrs, + 'local_contact_node_states': self.X[transitionNode_GNbrs].flatten(), + 'isolation_contact_nodes': transitionNode_GQNbrs, + 'isolation_contact_node_states':self.X[transitionNode_GQNbrs].flatten() }) + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + if (transitionType in ['EtoQE', 'IPREtoQPRE', 'ISYMtoQSYM', 'IASYMtoQASYM', 'ISYMtoH']): + self.set_positive(node=transitionNode, positive=True) + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + else: + + tau = 0.01 + self.t += tau + self.timer_state += tau + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + self.tidx += 1 + + self.tseries[self.tidx] = self.t + self.numS[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.S), a_min=0, a_max=self.numNodes) + self.numE[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.E), a_min=0, a_max=self.numNodes) + self.numI_pre[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.I_pre), a_min=0, a_max=self.numNodes) + self.numI_sym[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.I_sym), a_min=0, a_max=self.numNodes) + self.numI_asym[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.I_asym), a_min=0, a_max=self.numNodes) + self.numH[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.H), a_min=0, a_max=self.numNodes) + self.numR[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.R), a_min=0, a_max=self.numNodes) + self.numF[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.F), a_min=0, a_max=self.numNodes) + self.numQ_S[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.Q_S), a_min=0, a_max=self.numNodes) + self.numQ_E[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.Q_E), a_min=0, a_max=self.numNodes) + self.numQ_pre[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.Q_pre), a_min=0, a_max=self.numNodes) + self.numQ_sym[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.Q_sym), a_min=0, a_max=self.numNodes) + self.numQ_asym[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.Q_asym), a_min=0, a_max=self.numNodes) + self.numQ_R[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.Q_R), a_min=0, a_max=self.numNodes) + self.numTested[self.tidx] = numpy.clip(numpy.count_nonzero(self.tested), a_min=0, a_max=self.numNodes) + self.numPositive[self.tidx] = numpy.clip(numpy.count_nonzero(self.positive), a_min=0, a_max=self.numNodes) + + self.N[self.tidx] = numpy.clip((self.numNodes - self.numF[self.tidx]), a_min=0, a_max=self.numNodes) + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Update testing and isolation statuses + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + isolatedNodes = numpy.argwhere((self.X==self.Q_S)|(self.X==self.Q_E)|(self.X==self.Q_pre)|(self.X==self.Q_sym)|(self.X==self.Q_asym)|(self.X==self.Q_R))[:,0].flatten() + self.timer_isolation[isolatedNodes] = self.timer_isolation[isolatedNodes] + tau + + nodesExitingIsolation = numpy.argwhere(self.timer_isolation >= self.isolationTime) + for isoNode in nodesExitingIsolation: + self.set_isolation(node=isoNode, isolate=False) + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Store system states + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + if self.store_Xseries: + self.Xseries[self.tidx,:] = self.X.T + + if self.nodeGroupData: + for groupName in self.nodeGroupData: + self.nodeGroupData[groupName]['numS'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.S) + self.nodeGroupData[groupName]['numE'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.E) + self.nodeGroupData[groupName]['numI_pre'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.I_pre) + self.nodeGroupData[groupName]['numI_sym'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.I_sym) + self.nodeGroupData[groupName]['numI_asym'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.I_asym) + self.nodeGroupData[groupName]['numH'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.H) + self.nodeGroupData[groupName]['numR'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.R) + self.nodeGroupData[groupName]['numF'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.F) + self.nodeGroupData[groupName]['numQ_S'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_S) + self.nodeGroupData[groupName]['numQ_E'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_E) + self.nodeGroupData[groupName]['numQ_pre'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_pre) + self.nodeGroupData[groupName]['numQ_sym'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_sym) + self.nodeGroupData[groupName]['numQ_asym'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_asym) + self.nodeGroupData[groupName]['numQ_R'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_R) + self.nodeGroupData[groupName]['N'][self.tidx] = numpy.clip((self.nodeGroupData[groupName]['numS'][0] + self.nodeGroupData[groupName]['numE'][0] + self.nodeGroupData[groupName]['numI'][0] + self.nodeGroupData[groupName]['numQ_E'][0] + self.nodeGroupData[groupName]['numQ_I'][0] + self.nodeGroupData[groupName]['numR'][0]), a_min=0, a_max=self.numNodes) + self.nodeGroupData[groupName]['numTested'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.tested) + self.nodeGroupData[groupName]['numPositive'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.positive) + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Terminate if tmax reached or num infections is 0: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + if (self.t >= self.tmax or (self.total_num_infected(self.tidx) < 1 and self.total_num_isolated(self.tidx) < 1)): + self.finalize_data_series() + return False + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + return True + + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def run(self, T, checkpoints=None, max_dt=None, min_dt=None, print_interval=10, verbose='t'): + if T > 0: + self.tmax += T + else: + return False + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Pre-process checkpoint values: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + if checkpoints: + numCheckpoints = len(checkpoints['t']) + for chkpt_param, chkpt_values in checkpoints.items(): + assert(isinstance(chkpt_values, (list, numpy.ndarray)) and len(chkpt_values)==numCheckpoints), "Expecting a list of values with length equal to number of checkpoint times ("+str(numCheckpoints)+") for each checkpoint parameter." + checkpointIdx = numpy.searchsorted(checkpoints['t'], self.t) # Finds 1st index in list greater than given val + if checkpointIdx >= numCheckpoints: + # We are out of checkpoints, stop checking them: + checkpoints = None + else: + checkpointTime = checkpoints['t'][checkpointIdx] + + #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% + # Run the simulation loop: + #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% + print_reset = True + running = True + while running: + + running = self.run_iteration() + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Handle checkpoints if applicable: + if checkpoints: + if self.t >= checkpointTime: + if verbose is not False: + print("[Checkpoint: Updating parameters]") + # A checkpoint has been reached, update param values: + for param in list(self.parameters.keys()): + if param in list(checkpoints.keys()): + self.parameters.update({param: checkpoints[param][checkpointIdx]}) + # Update parameter data structures and scenario flags: + self.update_parameters() + # Update the next checkpoint time: + checkpointIdx = numpy.searchsorted(checkpoints['t'], self.t) # Finds 1st index in list greater than given val + if checkpointIdx >= numCheckpoints: + # We are out of checkpoints, stop checking them: + checkpoints = None + else: + checkpointTime = checkpoints['t'][checkpointIdx] + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + if print_interval: + if (print_reset and (int(self.t) % print_interval == 0)): + if verbose == "t": + print("t = %.2f" % self.t) + if verbose == True: + print("t = %.2f" % self.t) + print("\t S = " + str(self.numS[self.tidx])) + print("\t E = " + str(self.numE[self.tidx])) + print("\t I_pre = " + str(self.numI_pre[self.tidx])) + print("\t I_sym = " + str(self.numI_sym[self.tidx])) + print("\t I_asym = " + str(self.numI_asym[self.tidx])) + print("\t H = " + str(self.numH[self.tidx])) + print("\t R = " + str(self.numR[self.tidx])) + print("\t F = " + str(self.numF[self.tidx])) + print("\t Q_S = " + str(self.numQ_S[self.tidx])) + print("\t Q_E = " + str(self.numQ_E[self.tidx])) + print("\t Q_pre = " + str(self.numQ_pre[self.tidx])) + print("\t Q_sym = " + str(self.numQ_sym[self.tidx])) + print("\t Q_asym = " + str(self.numQ_asym[self.tidx])) + print("\t Q_R = " + str(self.numQ_R[self.tidx])) + + print_reset = False + elif (not print_reset and (int(self.t) % 10 != 0)): + print_reset = True + + return True + + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def plot(self, ax=None, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_sym='line', plot_I_asym='line', + plot_H='line', plot_R='line', plot_F='line', + plot_Q_E='line', plot_Q_pre='line', plot_Q_sym='line', plot_Q_asym='line', + plot_Q_S='line', plot_Q_R='line', combine_Q_infected=True, + color_S='tab:green', color_E='orange', color_I_pre='tomato', color_I_sym='crimson', color_I_asym='#F0909B', + color_H='violet', color_R='tab:blue', color_F='black', + color_Q_E='orange', color_Q_pre='tomato', color_Q_sym='crimson', color_Q_asym='#F0909B', + color_Q_S='tab:green', color_Q_R='tab:blue', color_Q_infected='tab:purple', + color_reference='#E0E0E0', + dashed_reference_results=None, dashed_reference_label='reference', + shaded_reference_results=None, shaded_reference_label='reference', + vlines=[], vline_colors=[], vline_styles=[], vline_labels=[], + ylim=None, xlim=None, legend=True, title=None, side_title=None, plot_percentages=True): + + import matplotlib.pyplot as pyplot + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Create an Axes object if None provided: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + if not ax: + fig, ax = pyplot.subplots() + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Prepare data series to be plotted: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + Sseries = self.numS/getattr(self, self.plotting_number_property) if plot_percentages else self.numS + Eseries = self.numE/getattr(self, self.plotting_number_property) if plot_percentages else self.numE + I_preseries = self.numI_pre/getattr(self, self.plotting_number_property) if plot_percentages else self.numI_pre + I_symseries = self.numI_sym/getattr(self, self.plotting_number_property) if plot_percentages else self.numI_sym + I_asymseries = self.numI_asym/getattr(self, self.plotting_number_property) if plot_percentages else self.numI_asym + Rseries = self.numR/getattr(self, self.plotting_number_property) if plot_percentages else self.numR + Hseries = self.numH/getattr(self, self.plotting_number_property) if plot_percentages else self.numH + Fseries = self.numF/getattr(self, self.plotting_number_property) if plot_percentages else self.numF + Q_Sseries = self.numQ_S/getattr(self, self.plotting_number_property) if plot_percentages else self.numQ_S + Q_Eseries = self.numQ_E/getattr(self, self.plotting_number_property) if plot_percentages else self.numQ_E + Q_preseries = self.numQ_pre/getattr(self, self.plotting_number_property) if plot_percentages else self.numQ_pre + Q_asymseries = self.numQ_asym/getattr(self, self.plotting_number_property) if plot_percentages else self.numQ_asym + Q_symseries = self.numQ_sym/getattr(self, self.plotting_number_property) if plot_percentages else self.numQ_sym + Q_Rseries = self.numQ_R/getattr(self, self.plotting_number_property) if plot_percentages else self.numQ_R + Q_infectedseries = (Q_Eseries + Q_preseries + Q_asymseries + Q_symseries) + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Draw the reference data: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + if dashed_reference_results: + dashedReference_tseries = dashed_reference_results.tseries[::int(getattr(self, self.plotting_number_property)/100)] + dashedReference_infectedStack = dashed_reference_results.total_num_infected()[::int(getattr(self, self.plotting_number_property)/100)] / (getattr(self, self.plotting_number_property) if plot_percentages else 1) + ax.plot(dashedReference_tseries, dashedReference_infectedStack, color='#E0E0E0', linestyle='--', label='Total infections ('+dashed_reference_label+')', zorder=0) + if shaded_reference_results: + shadedReference_tseries = shaded_reference_results.tseries + shadedReference_infectedStack = shaded_reference_results.total_num_infected() / (getattr(self, self.plotting_number_property) if plot_percentages else 1) + ax.fill_between(shaded_reference_results.tseries, shadedReference_infectedStack, 0, color='#EFEFEF', label='Total infections ('+shaded_reference_label+')', zorder=0) + ax.plot(shaded_reference_results.tseries, shadedReference_infectedStack, color='#E0E0E0', zorder=1) + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Draw the stacked variables: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + topstack = numpy.zeros_like(self.tseries) + if (any(Fseries) and plot_F=='stacked'): + ax.fill_between(numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, topstack+Fseries), topstack, color=color_F, alpha=0.75, label='$F$', zorder=2) + ax.plot( numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, topstack+Fseries), color=color_F, zorder=3) + topstack = topstack+Fseries + + if (any(Hseries) and plot_H=='stacked'): + ax.fill_between(numpy.ma.masked_where(Hseries<=0, self.tseries), numpy.ma.masked_where(Hseries<=0, topstack+Hseries), topstack, color=color_H, alpha=0.75, label='$H$', zorder=2) + ax.plot( numpy.ma.masked_where(Hseries<=0, self.tseries), numpy.ma.masked_where(Hseries<=0, topstack+Hseries), color=color_H, zorder=3) + topstack = topstack+Hseries + + if (combine_Q_infected and any(Q_infectedseries) and plot_Q_E=='stacked' and plot_Q_pre=='stacked' and plot_Q_sym=='stacked' and plot_Q_asym=='stacked'): + ax.fill_between(numpy.ma.masked_where(Q_infectedseries<=0, self.tseries), numpy.ma.masked_where(Q_infectedseries<=0, topstack+Q_infectedseries), topstack, facecolor=color_Q_infected, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_{infected}$', zorder=2) + ax.plot( numpy.ma.masked_where(Q_infectedseries<=0, self.tseries), numpy.ma.masked_where(Q_infectedseries<=0, topstack+Q_infectedseries), color=color_Q_infected, zorder=3) + topstack = topstack+Q_infectedseries + + if (not combine_Q_infected and any(Q_Eseries) and plot_Q_E=='stacked'): + ax.fill_between(numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, topstack+Q_Eseries), topstack, facecolor=color_Q_E, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_E$', zorder=2) + ax.plot( numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, topstack+Q_Eseries), color=color_Q_E, zorder=3) + topstack = topstack+Q_Eseries + + if (any(Eseries) and plot_E=='stacked'): + ax.fill_between(numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, topstack+Eseries), topstack, color=color_E, alpha=0.75, label='$E$', zorder=2) + ax.plot( numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, topstack+Eseries), color=color_E, zorder=3) + topstack = topstack+Eseries + + if (not combine_Q_infected and any(Q_preseries) and plot_Q_pre=='stacked'): + ax.fill_between(numpy.ma.masked_where(Q_preseries<=0, self.tseries), numpy.ma.masked_where(Q_preseries<=0, topstack+Q_preseries), topstack, facecolor=color_Q_pre, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_{pre}$', zorder=2) + ax.plot( numpy.ma.masked_where(Q_preseries<=0, self.tseries), numpy.ma.masked_where(Q_preseries<=0, topstack+Q_preseries), color=color_Q_pre, zorder=3) + topstack = topstack+Q_preseries + + if (any(I_preseries) and plot_I_pre=='stacked'): + ax.fill_between(numpy.ma.masked_where(I_preseries<=0, self.tseries), numpy.ma.masked_where(I_preseries<=0, topstack+I_preseries), topstack, color=color_I_pre, alpha=0.75, label='$I_{pre}$', zorder=2) + ax.plot( numpy.ma.masked_where(I_preseries<=0, self.tseries), numpy.ma.masked_where(I_preseries<=0, topstack+I_preseries), color=color_I_pre, zorder=3) + topstack = topstack+I_preseries + + if (not combine_Q_infected and any(Q_symseries) and plot_Q_sym=='stacked'): + ax.fill_between(numpy.ma.masked_where(Q_symseries<=0, self.tseries), numpy.ma.masked_where(Q_symseries<=0, topstack+Q_symseries), topstack, facecolor=color_Q_sym, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_{sym}$', zorder=2) + ax.plot( numpy.ma.masked_where(Q_symseries<=0, self.tseries), numpy.ma.masked_where(Q_symseries<=0, topstack+Q_symseries), color=color_Q_sym, zorder=3) + topstack = topstack+Q_symseries + + if (any(I_symseries) and plot_I_sym=='stacked'): + ax.fill_between(numpy.ma.masked_where(I_symseries<=0, self.tseries), numpy.ma.masked_where(I_symseries<=0, topstack+I_symseries), topstack, color=color_I_sym, alpha=0.75, label='$I_{sym}$', zorder=2) + ax.plot( numpy.ma.masked_where(I_symseries<=0, self.tseries), numpy.ma.masked_where(I_symseries<=0, topstack+I_symseries), color=color_I_sym, zorder=3) + topstack = topstack+I_symseries + + if (not combine_Q_infected and any(Q_asymseries) and plot_Q_asym=='stacked'): + ax.fill_between(numpy.ma.masked_where(Q_asymseries<=0, self.tseries), numpy.ma.masked_where(Q_asymseries<=0, topstack+Q_asymseries), topstack, facecolor=color_Q_asym, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_{asym}$', zorder=2) + ax.plot( numpy.ma.masked_where(Q_asymseries<=0, self.tseries), numpy.ma.masked_where(Q_asymseries<=0, topstack+Q_asymseries), color=color_Q_asym, zorder=3) + topstack = topstack+Q_asymseries + + if (any(I_asymseries) and plot_I_asym=='stacked'): + ax.fill_between(numpy.ma.masked_where(I_asymseries<=0, self.tseries), numpy.ma.masked_where(I_asymseries<=0, topstack+I_asymseries), topstack, color=color_I_asym, alpha=0.75, label='$I_{asym}$', zorder=2) + ax.plot( numpy.ma.masked_where(I_asymseries<=0, self.tseries), numpy.ma.masked_where(I_asymseries<=0, topstack+I_asymseries), color=color_I_asym, zorder=3) + topstack = topstack+I_asymseries + + if (any(Q_Rseries) and plot_Q_R=='stacked'): + ax.fill_between(numpy.ma.masked_where(Q_Rseries<=0, self.tseries), numpy.ma.masked_where(Q_Rseries<=0, topstack+Q_Rseries), topstack, facecolor=color_Q_R, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_R$', zorder=2) + ax.plot( numpy.ma.masked_where(Q_Rseries<=0, self.tseries), numpy.ma.masked_where(Q_Rseries<=0, topstack+Q_Rseries), color=color_Q_R, zorder=3) + topstack = topstack+Q_Rseries + + if (any(Rseries) and plot_R=='stacked'): + ax.fill_between(numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, topstack+Rseries), topstack, color=color_R, alpha=0.75, label='$R$', zorder=2) + ax.plot( numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, topstack+Rseries), color=color_R, zorder=3) + topstack = topstack+Rseries + + if (any(Q_Sseries) and plot_Q_S=='stacked'): + ax.fill_between(numpy.ma.masked_where(Q_Sseries<=0, self.tseries), numpy.ma.masked_where(Q_Sseries<=0, topstack+Q_Sseries), topstack, facecolor=color_Q_S, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_S$', zorder=2) + ax.plot( numpy.ma.masked_where(Q_Sseries<=0, self.tseries), numpy.ma.masked_where(Q_Sseries<=0, topstack+Q_Sseries), color=color_Q_S, zorder=3) + topstack = topstack+Q_Sseries + + if (any(Sseries) and plot_S=='stacked'): + ax.fill_between(numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, topstack+Sseries), topstack, color=color_S, alpha=0.75, label='$S$', zorder=2) + ax.plot( numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, topstack+Sseries), color=color_S, zorder=3) + topstack = topstack+Sseries + + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Draw the shaded variables: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + if (any(Fseries) and plot_F=='shaded'): + ax.fill_between(numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, Fseries), 0, color=color_F, alpha=0.75, label='$F$', zorder=4) + ax.plot( numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, Fseries), color=color_F, zorder=5) + + if (any(Hseries) and plot_H=='shaded'): + ax.fill_between(numpy.ma.masked_where(Hseries<=0, self.tseries), numpy.ma.masked_where(Hseries<=0, Hseries), 0, color=color_H, alpha=0.75, label='$H$', zorder=4) + ax.plot( numpy.ma.masked_where(Hseries<=0, self.tseries), numpy.ma.masked_where(Hseries<=0, Hseries), color=color_H, zorder=5) + + if (combine_Q_infected and any(Q_infectedseries) and plot_Q_E=='shaded' and plot_Q_pre=='shaded' and plot_Q_sym=='shaded' and plot_Q_asym=='shaded'): + ax.fill_between(numpy.ma.masked_where(Q_infectedseries<=0, self.tseries), numpy.ma.masked_where(Q_infectedseries<=0, Q_infectedseries), 0, color=color_Q_infected, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_{infected}$', zorder=4) + ax.plot( numpy.ma.masked_where(Q_infectedseries<=0, self.tseries), numpy.ma.masked_where(Q_infectedseries<=0, Q_infectedseries), color=color_Q_infected, zorder=5) + + if (not combine_Q_infected and any(Q_Eseries) and plot_Q_E=='shaded'): + ax.fill_between(numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, Q_Eseries), 0, facecolor=color_Q_E, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_E$', zorder=4) + ax.plot( numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, Q_Eseries), color=color_Q_E, zorder=5) + + if (any(Eseries) and plot_E=='shaded'): + ax.fill_between(numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, Eseries), 0, color=color_E, alpha=0.75, label='$E$', zorder=4) + ax.plot( numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, Eseries), color=color_E, zorder=5) + + if (not combine_Q_infected and any(Q_preseries) and plot_Q_pre=='shaded'): + ax.fill_between(numpy.ma.masked_where(Q_preseries<=0, self.tseries), numpy.ma.masked_where(Q_preseries<=0, Q_preseries), 0, facecolor=color_Q_pre, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_{pre}$', zorder=4) + ax.plot( numpy.ma.masked_where(Q_preseries<=0, self.tseries), numpy.ma.masked_where(Q_preseries<=0, Q_preseries), color=color_Q_pre, zorder=5) + + if (any(I_preseries) and plot_I_pre=='shaded'): + ax.fill_between(numpy.ma.masked_where(I_preseries<=0, self.tseries), numpy.ma.masked_where(I_preseries<=0, I_preseries), 0, color=color_I_pre, alpha=0.75, label='$I_{pre}$', zorder=4) + ax.plot( numpy.ma.masked_where(I_preseries<=0, self.tseries), numpy.ma.masked_where(I_preseries<=0, I_preseries), color=color_I_pre, zorder=5) + + if (not combine_Q_infected and any(Q_symseries) and plot_Q_sym=='shaded'): + ax.fill_between(numpy.ma.masked_where(Q_symseries<=0, self.tseries), numpy.ma.masked_where(Q_symseries<=0, Q_symseries), 0, facecolor=color_Q_sym, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_{sym}$', zorder=4) + ax.plot( numpy.ma.masked_where(Q_symseries<=0, self.tseries), numpy.ma.masked_where(Q_symseries<=0, Q_symseries), color=color_Q_sym, zorder=5) + + if (any(I_symseries) and plot_I_sym=='shaded'): + ax.fill_between(numpy.ma.masked_where(I_symseries<=0, self.tseries), numpy.ma.masked_where(I_symseries<=0, I_symseries), 0, color=color_I_sym, alpha=0.75, label='$I_{sym}$', zorder=4) + ax.plot( numpy.ma.masked_where(I_symseries<=0, self.tseries), numpy.ma.masked_where(I_symseries<=0, I_symseries), color=color_I_sym, zorder=5) + + if (not combine_Q_infected and any(Q_asymseries) and plot_Q_asym=='shaded'): + ax.fill_between(numpy.ma.masked_where(Q_asymseries<=0, self.tseries), numpy.ma.masked_where(Q_asymseries<=0, Q_asymseries), 0, facecolor=color_Q_asym, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_{asym}$', zorder=4) + ax.plot( numpy.ma.masked_where(Q_asymseries<=0, self.tseries), numpy.ma.masked_where(Q_asymseries<=0, Q_asymseries), color=color_Q_asym, zorder=5) + + if (any(I_asymseries) and plot_I_asym=='shaded'): + ax.fill_between(numpy.ma.masked_where(I_asymseries<=0, self.tseries), numpy.ma.masked_where(I_asymseries<=0, I_asymseries), 0, color=color_I_asym, alpha=0.75, label='$I_{asym}$', zorder=4) + ax.plot( numpy.ma.masked_where(I_asymseries<=0, self.tseries), numpy.ma.masked_where(I_asymseries<=0, I_asymseries), color=color_I_asym, zorder=5) + + if (any(Q_Rseries) and plot_Q_R=='shaded'): + ax.fill_between(numpy.ma.masked_where(Q_Rseries<=0, self.tseries), numpy.ma.masked_where(Q_Rseries<=0, Q_Rseries), 0, facecolor=color_Q_R, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_R$', zorder=4) + ax.plot( numpy.ma.masked_where(Q_Rseries<=0, self.tseries), numpy.ma.masked_where(Q_Rseries<=0, Q_Rseries), color=color_Q_R, zorder=5) + + if (any(Rseries) and plot_R=='shaded'): + ax.fill_between(numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, Rseries), 0, color=color_R, alpha=0.75, label='$R$', zorder=4) + ax.plot( numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, Rseries), color=color_R, zorder=5) + + if (any(Q_Sseries) and plot_Q_S=='shaded'): + ax.fill_between(numpy.ma.masked_where(Q_Sseries<=0, self.tseries), numpy.ma.masked_where(Q_Sseries<=0, Q_Sseries), 0, facecolor=color_Q_S, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_S$', zorder=4) + ax.plot( numpy.ma.masked_where(Q_Sseries<=0, self.tseries), numpy.ma.masked_where(Q_Sseries<=0, Q_Sseries), color=color_Q_S, zorder=5) + + if (any(Sseries) and plot_S=='shaded'): + ax.fill_between(numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, Sseries), 0, color=color_S, alpha=0.75, label='$S$', zorder=4) + ax.plot( numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, Sseries), color=color_S, zorder=5) + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Draw the line variables: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + if (any(Fseries) and plot_F=='line'): + ax.plot(numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, Fseries), color=color_F, label='$F$', zorder=6) + + if (any(Hseries) and plot_H=='line'): + ax.plot(numpy.ma.masked_where(Hseries<=0, self.tseries), numpy.ma.masked_where(Hseries<=0, Hseries), color=color_H, label='$H$', zorder=6) + + if (combine_Q_infected and any(Q_infectedseries) and plot_Q_E=='line' and plot_Q_pre=='line' and plot_Q_sym=='line' and plot_Q_asym=='line'): + ax.plot(numpy.ma.masked_where(Q_infectedseries<=0, self.tseries), numpy.ma.masked_where(Q_infectedseries<=0, Q_infectedseries), color=color_Q_infected, label='$Q_{infected}$', zorder=6) + + if (not combine_Q_infected and any(Q_Eseries) and plot_Q_E=='line'): + ax.plot(numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, Q_Eseries), color=color_Q_E, label='$Q_E$', zorder=6) + + if (any(Eseries) and plot_E=='line'): + ax.plot(numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, Eseries), color=color_E, label='$E$', zorder=6) + + if (not combine_Q_infected and any(Q_preseries) and plot_Q_pre=='line'): + ax.plot(numpy.ma.masked_where(Q_preseries<=0, self.tseries), numpy.ma.masked_where(Q_preseries<=0, Q_preseries), color=color_Q_pre, label='$Q_{pre}$', zorder=6) + + if (any(I_preseries) and plot_I_pre=='line'): + ax.plot(numpy.ma.masked_where(I_preseries<=0, self.tseries), numpy.ma.masked_where(I_preseries<=0, I_preseries), color=color_I_pre, label='$I_{pre}$', zorder=6) + + if (not combine_Q_infected and any(Q_symseries) and plot_Q_sym=='line'): + ax.plot(numpy.ma.masked_where(Q_symseries<=0, self.tseries), numpy.ma.masked_where(Q_symseries<=0, Q_symseries), color=color_Q_sym, label='$Q_{sym}$', zorder=6) + + if (any(I_symseries) and plot_I_sym=='line'): + ax.plot(numpy.ma.masked_where(I_symseries<=0, self.tseries), numpy.ma.masked_where(I_symseries<=0, I_symseries), color=color_I_sym, label='$I_{sym}$', zorder=6) + + if (not combine_Q_infected and any(Q_asymseries) and plot_Q_asym=='line'): + ax.plot(numpy.ma.masked_where(Q_asymseries<=0, self.tseries), numpy.ma.masked_where(Q_asymseries<=0, Q_asymseries), color=color_Q_asym, label='$Q_{asym}$', zorder=6) + + if (any(I_asymseries) and plot_I_asym=='line'): + ax.plot(numpy.ma.masked_where(I_asymseries<=0, self.tseries), numpy.ma.masked_where(I_asymseries<=0, I_asymseries), color=color_I_asym, label='$I_{asym}$', zorder=6) + + if (any(Q_Rseries) and plot_Q_R=='line'): + ax.plot(numpy.ma.masked_where(Q_Rseries<=0, self.tseries), numpy.ma.masked_where(Q_Rseries<=0, Q_Rseries), color=color_Q_R, linestyle='--', label='$Q_R$', zorder=6) + + if (any(Rseries) and plot_R=='line'): + ax.plot(numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, Rseries), color=color_R, label='$R$', zorder=6) + + if (any(Q_Sseries) and plot_Q_S=='line'): + ax.plot(numpy.ma.masked_where(Q_Sseries<=0, self.tseries), numpy.ma.masked_where(Q_Sseries<=0, Q_Sseries), color=color_Q_S, linestyle='--', label='$Q_S$', zorder=6) + + if (any(Sseries) and plot_S=='line'): + ax.plot(numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, Sseries), color=color_S, label='$S$', zorder=6) + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Draw the vertical line annotations: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + if (len(vlines)>0 and len(vline_colors)==0): + vline_colors = ['gray']*len(vlines) + if (len(vlines)>0 and len(vline_labels)==0): + vline_labels = [None]*len(vlines) + if (len(vlines)>0 and len(vline_styles)==0): + vline_styles = [':']*len(vlines) + for vline_x, vline_color, vline_style, vline_label in zip(vlines, vline_colors, vline_styles, vline_labels): + if vline_x is not None: + ax.axvline(x=vline_x, color=vline_color, linestyle=vline_style, alpha=1, label=vline_label) + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Draw the plot labels: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + ax.set_xlabel('days') + ax.set_ylabel('percent of population' if plot_percentages else 'number of individuals') + ax.set_xlim(0, (max(self.tseries) if not xlim else xlim)) + ax.set_ylim(0, ylim) + if plot_percentages: + ax.set_yticklabels(['{:,.0%}'.format(y) for y in ax.get_yticks()]) + if legend: + legend_handles, legend_labels = ax.get_legend_handles_labels() + ax.legend(legend_handles[::-1], legend_labels[::-1], loc='upper right', facecolor='white', edgecolor='none', framealpha=0.9, prop={'size': 8}) + if title: + ax.set_title(title, size=12) + if side_title: + ax.annotate(side_title, (0, 0.5), xytext=(-45, 0), ha='right', va='center', + size=12, rotation=90, xycoords='axes fraction', textcoords='offset points') + + return ax + + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def figure_basic(self, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_sym='line', plot_I_asym='line', + plot_H='line', plot_R='line', plot_F='line', + plot_Q_E='line', plot_Q_pre='line', plot_Q_sym='line', plot_Q_asym='line', + plot_Q_S=False, plot_Q_R=False, combine_Q_infected=True, + color_S='tab:green', color_E='orange', color_I_pre='tomato', color_I_sym='crimson', color_I_asym='#F0909B', + color_H='violet', color_R='tab:blue', color_F='black', + color_Q_E='orange', color_Q_pre='tomato', color_Q_sym='crimson', color_Q_asym='#F0909B', + color_Q_S='tab:green', color_Q_R='tab:blue', color_Q_infected='tab:purple', + color_reference='#E0E0E0', + dashed_reference_results=None, dashed_reference_label='reference', + shaded_reference_results=None, shaded_reference_label='reference', + vlines=[], vline_colors=[], vline_styles=[], vline_labels=[], + ylim=None, xlim=None, legend=True, title=None, side_title=None, plot_percentages=True, + figsize=(12,8), use_seaborn=True, show=True): + + import matplotlib.pyplot as pyplot + + fig, ax = pyplot.subplots(figsize=figsize) + + if use_seaborn: + import seaborn + seaborn.set_style('ticks') + seaborn.despine() + + self.plot(ax=ax, plot_S=plot_S, plot_E=plot_E, plot_I_pre=plot_I_pre, plot_I_sym=plot_I_sym, plot_I_asym=plot_I_asym, + plot_H=plot_H, plot_R=plot_R, plot_F=plot_F, + plot_Q_E=plot_Q_E, plot_Q_pre=plot_Q_pre, plot_Q_sym=plot_Q_sym, plot_Q_asym=plot_Q_asym, + plot_Q_S=plot_Q_S, plot_Q_R=plot_Q_R, combine_Q_infected=combine_Q_infected, + color_S=color_S, color_E=color_E, color_I_pre=color_I_pre, color_I_sym=color_I_sym, color_I_asym=color_I_asym, + color_H=color_H, color_R=color_R, color_F=color_F, + color_Q_E=color_Q_E, color_Q_pre=color_Q_pre, color_Q_sym=color_Q_sym, color_Q_asym=color_Q_asym, + color_Q_S=color_Q_S, color_Q_R=color_Q_R, color_Q_infected=color_Q_infected, + color_reference=color_reference, + dashed_reference_results=dashed_reference_results, dashed_reference_label=dashed_reference_label, + shaded_reference_results=shaded_reference_results, shaded_reference_label=shaded_reference_label, + vlines=vlines, vline_colors=vline_colors, vline_styles=vline_styles, vline_labels=vline_labels, + ylim=ylim, xlim=xlim, legend=legend, title=title, side_title=side_title, plot_percentages=plot_percentages) + + if show: + pyplot.show() + + return fig, ax + + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def figure_infections(self, plot_S=False, plot_E='stacked', plot_I_pre='stacked', plot_I_sym='stacked', plot_I_asym='stacked', + plot_H='stacked', plot_R=False, plot_F='stacked', + plot_Q_E='stacked', plot_Q_pre='stacked', plot_Q_sym='stacked', plot_Q_asym='stacked', + plot_Q_S=False, plot_Q_R=False, combine_Q_infected=True, + color_S='tab:green', color_E='orange', color_I_pre='tomato', color_I_sym='crimson', color_I_asym='#F0909B', + color_H='violet', color_R='tab:blue', color_F='black', + color_Q_E='orange', color_Q_pre='tomato', color_Q_sym='crimson', color_Q_asym='#F0909B', + color_Q_S='tab:green', color_Q_R='tab:blue', color_Q_infected='tab:purple', + color_reference='#E0E0E0', + dashed_reference_results=None, dashed_reference_label='reference', + shaded_reference_results=None, shaded_reference_label='reference', + vlines=[], vline_colors=[], vline_styles=[], vline_labels=[], + ylim=None, xlim=None, legend=True, title=None, side_title=None, plot_percentages=True, + figsize=(12,8), use_seaborn=True, show=True): + + import matplotlib.pyplot as pyplot + + fig, ax = pyplot.subplots(figsize=figsize) + + if use_seaborn: + import seaborn + seaborn.set_style('ticks') + seaborn.despine() + + self.plot(ax=ax, plot_S=plot_S, plot_E=plot_E, plot_I_pre=plot_I_pre, plot_I_sym=plot_I_sym, plot_I_asym=plot_I_asym, + plot_H=plot_H, plot_R=plot_R, plot_F=plot_F, + plot_Q_E=plot_Q_E, plot_Q_pre=plot_Q_pre, plot_Q_sym=plot_Q_sym, plot_Q_asym=plot_Q_asym, + plot_Q_S=plot_Q_S, plot_Q_R=plot_Q_R, combine_Q_infected=combine_Q_infected, + color_S=color_S, color_E=color_E, color_I_pre=color_I_pre, color_I_sym=color_I_sym, color_I_asym=color_I_asym, + color_H=color_H, color_R=color_R, color_F=color_F, + color_Q_E=color_Q_E, color_Q_pre=color_Q_pre, color_Q_sym=color_Q_sym, color_Q_asym=color_Q_asym, + color_Q_S=color_Q_S, color_Q_R=color_Q_R, color_Q_infected=color_Q_infected, + color_reference=color_reference, + dashed_reference_results=dashed_reference_results, dashed_reference_label=dashed_reference_label, + shaded_reference_results=shaded_reference_results, shaded_reference_label=shaded_reference_label, + vlines=vlines, vline_colors=vline_colors, vline_styles=vline_styles, vline_labels=vline_labels, + ylim=ylim, xlim=xlim, legend=legend, title=title, side_title=side_title, plot_percentages=plot_percentages) + + if show: + pyplot.show() + + return fig, ax diff --git a/seirsplus/models/seirs_model.py b/seirsplus/models/seirs_model.py new file mode 100644 index 0000000..baf6383 --- /dev/null +++ b/seirsplus/models/seirs_model.py @@ -0,0 +1,268 @@ +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy +import scipy.integrate + +from .base_plotable_model import BasePlotableModel + + +class SEIRSModel(BasePlotableModel): + """ + A class to simulate the Deterministic SEIRS Model + =================================================== + Params: beta Rate of transmission (exposure) + sigma Rate of infection (upon exposure) + gamma Rate of recovery (upon infection) + xi Rate of re-susceptibility (upon recovery) + mu_I Rate of infection-related death + mu_0 Rate of baseline death + nu Rate of baseline birth + + beta_Q Rate of transmission (exposure) for individuals with detected infections + sigma_Q Rate of infection (upon exposure) for individuals with detected infections + gamma_Q Rate of recovery (upon infection) for individuals with detected infections + mu_Q Rate of infection-related death for individuals with detected infections + theta_E Rate of baseline testing for exposed individuals + theta_I Rate of baseline testing for infectious individuals + psi_E Probability of positive test results for exposed individuals + psi_I Probability of positive test results for exposed individuals + q Probability of quarantined individuals interacting with others + + initE Init number of exposed individuals + initI Init number of infectious individuals + initQ_E Init number of detected infectious individuals + initQ_I Init number of detected infectious individuals + initR Init number of recovered individuals + initF Init number of infection-related fatalities + (all remaining nodes initialized susceptible) + """ + + plotting_number_property = 'N' + """Property to access the number to base plotting on.""" + + + def __init__(self, initN, beta, sigma, gamma, xi=0, mu_I=0, mu_0=0, nu=0, p=0, + beta_Q=None, sigma_Q=None, gamma_Q=None, mu_Q=None, + theta_E=0, theta_I=0, psi_E=0, psi_I=0, q=0, + initE=0, initI=10, initQ_E=0, initQ_I=0, initR=0, initF=0): + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Model Parameters: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + self.beta = beta + self.sigma = sigma + self.gamma = gamma + self.xi = xi + self.mu_I = mu_I + self.mu_0 = mu_0 + self.nu = nu + self.p = p + + # Testing-related parameters: + self.beta_Q = beta_Q if beta_Q is not None else self.beta + self.sigma_Q = sigma_Q if sigma_Q is not None else self.sigma + self.gamma_Q = gamma_Q if gamma_Q is not None else self.gamma + self.mu_Q = mu_Q if mu_Q is not None else self.mu_I + self.theta_E = theta_E if theta_E is not None else self.theta_E + self.theta_I = theta_I if theta_I is not None else self.theta_I + self.psi_E = psi_E if psi_E is not None else self.psi_E + self.psi_I = psi_I if psi_I is not None else self.psi_I + self.q = q if q is not None else self.q + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Initialize Timekeeping: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + self.t = 0 + self.tmax = 0 # will be set when run() is called + self.tseries = numpy.array([0]) + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Initialize Counts of inidividuals with each state: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + self.N = numpy.array([int(initN)]) + self.numE = numpy.array([int(initE)]) + self.numI = numpy.array([int(initI)]) + self.numQ_E = numpy.array([int(initQ_E)]) + self.numQ_I = numpy.array([int(initQ_I)]) + self.numR = numpy.array([int(initR)]) + self.numF = numpy.array([int(initF)]) + self.numS = numpy.array([self.N[-1] - self.numE[-1] - self.numI[-1] - self.numQ_E[-1] - self.numQ_I[-1] - self.numR[-1] - self.numF[-1]]) + assert(self.numS[0] >= 0), "The specified initial population size N must be greater than or equal to the initial compartment counts." + + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + @staticmethod + def system_dfes(t, variables, beta, sigma, gamma, xi, mu_I, mu_0, nu, + beta_Q, sigma_Q, gamma_Q, mu_Q, theta_E, theta_I, psi_E, psi_I, q): + + S, E, I, Q_E, Q_I, R, F = variables # variables is a list with compartment counts as elements + + N = S + E + I + Q_E + Q_I + R + + dS = - (beta*S*I)/N - q*(beta_Q*S*Q_I)/N + xi*R + nu*N - mu_0*S + + dE = (beta*S*I)/N + q*(beta_Q*S*Q_I)/N - sigma*E - theta_E*psi_E*E - mu_0*E + + dI = sigma*E - gamma*I - mu_I*I - theta_I*psi_I*I - mu_0*I + + dDE = theta_E*psi_E*E - sigma_Q*Q_E - mu_0*Q_E + + dDI = theta_I*psi_I*I + sigma_Q*Q_E - gamma_Q*Q_I - mu_Q*Q_I - mu_0*Q_I + + dR = gamma*I + gamma_Q*Q_I - xi*R - mu_0*R + + dF = mu_I*I + mu_Q*Q_I + + return [dS, dE, dI, dDE, dDI, dR, dF] + + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def run_epoch(self, runtime, dt=0.1): + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Create a list of times at which the ODE solver should output system values. + # Append this list of times as the model's time series + t_eval = numpy.arange(start=self.t, stop=self.t+runtime, step=dt) + + # Define the range of time values for the integration: + t_span = [self.t, self.t+runtime] + + # Define the initial conditions as the system's current state: + # (which will be the t=0 condition if this is the first run of this model, + # else where the last sim left off) + + init_cond = [self.numS[-1], self.numE[-1], self.numI[-1], self.numQ_E[-1], self.numQ_I[-1], self.numR[-1], self.numF[-1]] + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Solve the system of differential eqns: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + solution = scipy.integrate.solve_ivp(lambda t, X: SEIRSModel.system_dfes(t, X, self.beta, self.sigma, self.gamma, self.xi, self.mu_I, self.mu_0, self.nu, + self.beta_Q, self.sigma_Q, self.gamma_Q, self.mu_Q, self.theta_E, self.theta_I, self.psi_E, self.psi_I, self.q + ), + t_span=t_span, y0=init_cond, t_eval=t_eval + ) + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Store the solution output as the model's time series and data series: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + self.tseries = numpy.append(self.tseries, solution['t']) + self.numS = numpy.append(self.numS, solution['y'][0]) + self.numE = numpy.append(self.numE, solution['y'][1]) + self.numI = numpy.append(self.numI, solution['y'][2]) + self.numQ_E = numpy.append(self.numQ_E, solution['y'][3]) + self.numQ_I = numpy.append(self.numQ_I, solution['y'][4]) + self.numR = numpy.append(self.numR, solution['y'][5]) + self.numF = numpy.append(self.numF, solution['y'][6]) + + self.t = self.tseries[-1] + + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def run(self, T, dt=0.1, checkpoints=None, verbose=False): + + if T > 0: + self.tmax += T + else: + return False + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Pre-process checkpoint values: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + if checkpoints: + numCheckpoints = len(checkpoints['t']) + paramNames = ['beta', 'sigma', 'gamma', 'xi', 'mu_I', 'mu_0', 'nu', + 'beta_Q', 'sigma_Q', 'gamma_Q', 'mu_Q', + 'theta_E', 'theta_I', 'psi_E', 'psi_I', 'q'] + for param in paramNames: + # For params that don't have given checkpoint values (or bad value given), + # set their checkpoint values to the value they have now for all checkpoints. + if (param not in list(checkpoints.keys()) + or not isinstance(checkpoints[param], (list, numpy.ndarray)) + or len(checkpoints[param])!=numCheckpoints): + checkpoints[param] = [getattr(self, param)]*numCheckpoints + + #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% + # Run the simulation loop: + #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% + if not checkpoints: + self.run_epoch(runtime=self.tmax, dt=dt) + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + print("t = %.2f" % self.t) + if verbose: + print("\t S = " + str(self.numS[-1])) + print("\t E = " + str(self.numE[-1])) + print("\t I = " + str(self.numI[-1])) + print("\t Q_E = " + str(self.numQ_E[-1])) + print("\t Q_I = " + str(self.numQ_I[-1])) + print("\t R = " + str(self.numR[-1])) + print("\t F = " + str(self.numF[-1])) + + + else: # checkpoints provided + for checkpointIdx, checkpointTime in enumerate(checkpoints['t']): + # Run the sim until the next checkpoint time: + self.run_epoch(runtime=checkpointTime-self.t, dt=dt) + # Having reached the checkpoint, update applicable parameters: + print("[Checkpoint: Updating parameters]") + for param in paramNames: + setattr(self, param, checkpoints[param][checkpointIdx]) + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + print("t = %.2f" % self.t) + if verbose: + print("\t S = " + str(self.numS[-1])) + print("\t E = " + str(self.numE[-1])) + print("\t I = " + str(self.numI[-1])) + print("\t Q_E = " + str(self.numQ_E[-1])) + print("\t Q_I = " + str(self.numQ_I[-1])) + print("\t R = " + str(self.numR[-1])) + print("\t F = " + str(self.numF[-1])) + + if self.t < self.tmax: + self.run_epoch(runtime=self.tmax-self.t, dt=dt) + + return True + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def total_num_susceptible(self, t_idx=None): + if t_idx is None: + return self.numS[:] + else: + return self.numS[t_idx] + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def total_num_infected(self, t_idx=None): + if t_idx is None: + return self.numE[:] + self.numI[:] + self.numQ_E[:] + self.numQ_I[:] + else: + return self.numE[t_idx] + self.numI[t_idx] + self.numQ_E[t_idx] + self.numQ_I[t_idx] + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def total_num_isolated(self, t_idx=None): + if t_idx is None: + return self.numQ_E[:] + self.numQ_I[:] + else: + return self.numQ_E[t_idx] + self.numQ_I[t_idx] + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def total_num_recovered(self, t_idx=None): + if t_idx is None: + return self.numR[:] + else: + return self.numR[t_idx] diff --git a/seirsplus/models/seirs_network_model.py b/seirsplus/models/seirs_network_model.py new file mode 100644 index 0000000..90a6a9f --- /dev/null +++ b/seirsplus/models/seirs_network_model.py @@ -0,0 +1,900 @@ +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import networkx +import numpy +import scipy + +from .base_plotable_model import BasePlotableModel + + +class SEIRSNetworkModel(BasePlotableModel): + """ + A class to simulate the SEIRS Stochastic Network Model + ====================================================== + Params: + G Network adjacency matrix (numpy array) or Networkx graph object. + beta Rate of transmission (global interactions) + beta_local Rate(s) of transmission between adjacent individuals (optional) + sigma Rate of progression to infectious state (inverse of latent period) + gamma Rate of recovery (inverse of symptomatic infectious period) + mu_I Rate of infection-related death + xi Rate of re-susceptibility (upon recovery) + mu_0 Rate of baseline death + nu Rate of baseline birth + p Probability of individuals interacting with global population + + G_Q Quarantine adjacency matrix (numpy array) or Networkx graph object. + beta_Q Rate of transmission for isolated individuals (global interactions) + beta_Q_local Rate(s) of transmission (exposure) for adjacent isolated individuals (optional) + sigma_Q Rate of progression to infectious state for isolated individuals + gamma_Q Rate of recovery for isolated individuals + mu_Q Rate of infection-related death for isolated individuals + q Probability of isolated individuals interacting with global population + isolation_time Time to remain in isolation upon positive test, self-isolation, etc. + + theta_E Rate of random testing for exposed individuals + theta_I Rate of random testing for infectious individuals + phi_E Rate of testing when a close contact has tested positive for exposed individuals + phi_I Rate of testing when a close contact has tested positive for infectious individuals + psi_E Probability of positive test for exposed individuals + psi_I Probability of positive test for infectious individuals + + initE Initial number of exposed individuals + initI Initial number of infectious individuals + initR Initial number of recovered individuals + initF Initial number of infection-related fatalities + initQ_S Initial number of isolated susceptible individuals + initQ_E Initial number of isolated exposed individuals + initQ_I Initial number of isolated infectious individuals + initQ_R Initial number of isolated recovered individuals + (all remaining nodes initialized susceptible) + """ + + plotting_number_property = 'numNodes' + """Property to access the number to base plotting on.""" + + def __init__(self, G, beta, sigma, gamma, + mu_I=0, alpha=1.0, xi=0, mu_0=0, nu=0, f=0, p=0, + beta_local=None, beta_pairwise_mode='infected', delta=None, delta_pairwise_mode=None, + G_Q=None, beta_Q=None, beta_Q_local=None, sigma_Q=None, gamma_Q=None, mu_Q=None, alpha_Q=None, delta_Q=None, + theta_E=0, theta_I=0, phi_E=0, phi_I=0, psi_E=1, psi_I=1, q=0, isolation_time=14, + initE=0, initI=0, initR=0, initF=0, initQ_E=0, initQ_I=0, + transition_mode='exponential_rates', node_groups=None, store_Xseries=False, seed=None): + + if seed is not None: + numpy.random.seed(seed) + self.seed = seed + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Model Parameters: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + self.parameters = { 'G':G, 'G_Q':G_Q, + 'beta':beta, 'sigma':sigma, 'gamma':gamma, 'mu_I':mu_I, + 'xi':xi, 'mu_0':mu_0, 'nu':nu, 'f':f, 'p':p, + 'beta_local':beta_local, 'beta_pairwise_mode':beta_pairwise_mode, + 'alpha':alpha, 'delta':delta, 'delta_pairwise_mode':delta_pairwise_mode, + 'beta_Q':beta_Q, 'beta_Q_local':beta_Q_local, 'sigma_Q':sigma_Q, 'gamma_Q':gamma_Q, 'mu_Q':mu_Q, + 'alpha_Q':alpha_Q, 'delta_Q':delta_Q, + 'theta_E':theta_E, 'theta_I':theta_I, 'phi_E':phi_E, 'phi_I':phi_I, 'psi_E':psi_E, 'psi_I':psi_I, + 'q':q, 'isolation_time':isolation_time, + 'initE':initE, 'initI':initI, 'initR':initR, 'initF':initF, + 'initQ_E':initQ_E, 'initQ_I':initQ_I } + self.update_parameters() + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Each node can undergo 4-6 transitions (sans vitality/re-susceptibility returns to S state), + # so there are ~numNodes*6 events/timesteps expected; initialize numNodes*6 timestep slots to start + # (will be expanded during run if needed for some reason) + self.tseries = numpy.zeros(6*self.numNodes) + self.numS = numpy.zeros(6*self.numNodes) + self.numE = numpy.zeros(6*self.numNodes) + self.numI = numpy.zeros(6*self.numNodes) + self.numR = numpy.zeros(6*self.numNodes) + self.numF = numpy.zeros(6*self.numNodes) + self.numQ_E = numpy.zeros(6*self.numNodes) + self.numQ_I = numpy.zeros(6*self.numNodes) + self.N = numpy.zeros(6*self.numNodes) + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Initialize Timekeeping: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + self.t = 0 + self.tmax = 0 # will be set when run() is called + self.tidx = 0 + self.tseries[0] = 0 + + # Vectors holding the time that each node has been in a given state or in isolation: + self.timer_state = numpy.zeros((self.numNodes,1)) + self.timer_isolation = numpy.zeros(self.numNodes) + self.isolationTime = isolation_time + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Initialize Counts of individuals with each state: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + self.numE[0] = int(initE) + self.numI[0] = int(initI) + self.numR[0] = int(initR) + self.numF[0] = int(initF) + self.numQ_E[0] = int(initQ_E) + self.numQ_I[0] = int(initQ_I) + self.numS[0] = (self.numNodes - self.numE[0] - self.numI[0] - self.numR[0] - self.numQ_E[0] - self.numQ_I[0] - self.numF[0]) + self.N[0] = self.numNodes - self.numF[0] + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Node states: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + self.S = 1 + self.E = 2 + self.I = 3 + self.R = 4 + self.F = 5 + self.Q_E = 6 + self.Q_I = 7 + + self.X = numpy.array( [self.S]*int(self.numS[0]) + [self.E]*int(self.numE[0]) + [self.I]*int(self.numI[0]) + + [self.R]*int(self.numR[0]) + [self.F]*int(self.numF[0]) + + [self.Q_E]*int(self.numQ_E[0]) + [self.Q_I]*int(self.numQ_I[0]) + ).reshape((self.numNodes,1)) + numpy.random.shuffle(self.X) + + self.store_Xseries = store_Xseries + if store_Xseries: + self.Xseries = numpy.zeros(shape=(6*self.numNodes, self.numNodes), dtype='uint8') + self.Xseries[0,:] = self.X.T + + self.transitions = { + 'StoE': {'currentState':self.S, 'newState':self.E}, + 'EtoI': {'currentState':self.E, 'newState':self.I}, + 'ItoR': {'currentState':self.I, 'newState':self.R}, + 'ItoF': {'currentState':self.I, 'newState':self.F}, + 'RtoS': {'currentState':self.R, 'newState':self.S}, + 'EtoQE': {'currentState':self.E, 'newState':self.Q_E}, + 'ItoQI': {'currentState':self.I, 'newState':self.Q_I}, + 'QEtoQI': {'currentState':self.Q_E, 'newState':self.Q_I}, + 'QItoR': {'currentState':self.Q_I, 'newState':self.R}, + 'QItoF': {'currentState':self.Q_I, 'newState':self.F}, + '_toS': {'currentState':True, 'newState':self.S}, + } + + self.transition_mode = transition_mode + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Initialize other node metadata: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + self.tested = numpy.array([False]*self.numNodes).reshape((self.numNodes,1)) + self.positive = numpy.array([False]*self.numNodes).reshape((self.numNodes,1)) + self.numTested = numpy.zeros(6*self.numNodes) + self.numPositive = numpy.zeros(6*self.numNodes) + + self.testedInCurrentState = numpy.array([False]*self.numNodes).reshape((self.numNodes,1)) + + self.infectionsLog = [] + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Initialize node subgroup data series: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + self.nodeGroupData = None + if node_groups: + self.nodeGroupData = {} + for groupName, nodeList in node_groups.items(): + self.nodeGroupData[groupName] = {'nodes': numpy.array(nodeList), + 'mask': numpy.isin(range(self.numNodes), nodeList).reshape((self.numNodes,1))} + self.nodeGroupData[groupName]['numS'] = numpy.zeros(6*self.numNodes) + self.nodeGroupData[groupName]['numE'] = numpy.zeros(6*self.numNodes) + self.nodeGroupData[groupName]['numI'] = numpy.zeros(6*self.numNodes) + self.nodeGroupData[groupName]['numR'] = numpy.zeros(6*self.numNodes) + self.nodeGroupData[groupName]['numF'] = numpy.zeros(6*self.numNodes) + self.nodeGroupData[groupName]['numQ_E'] = numpy.zeros(6*self.numNodes) + self.nodeGroupData[groupName]['numQ_I'] = numpy.zeros(6*self.numNodes) + self.nodeGroupData[groupName]['N'] = numpy.zeros(6*self.numNodes) + self.nodeGroupData[groupName]['numPositive'] = numpy.zeros(6*self.numNodes) + self.nodeGroupData[groupName]['numTested'] = numpy.zeros(6*self.numNodes) + self.nodeGroupData[groupName]['numS'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.S) + self.nodeGroupData[groupName]['numE'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.E) + self.nodeGroupData[groupName]['numI'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.I) + self.nodeGroupData[groupName]['numR'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.R) + self.nodeGroupData[groupName]['numF'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.F) + self.nodeGroupData[groupName]['numQ_E'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_E) + self.nodeGroupData[groupName]['numQ_I'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_I) + self.nodeGroupData[groupName]['N'][0] = self.numNodes - self.numF[0] + + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def update_parameters(self): + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Model graphs: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + self.G = self.parameters['G'] + # Adjacency matrix: + if type(self.G)==numpy.ndarray: + self.A = scipy.sparse.csr_matrix(self.G) + elif type(self.G)==networkx.classes.graph.Graph: + self.A = networkx.adj_matrix(self.G) # adj_matrix gives scipy.sparse csr_matrix + else: + raise BaseException("Input an adjacency matrix or networkx object only.") + self.numNodes = int(self.A.shape[1]) + self.degree = numpy.asarray(self.node_degrees(self.A)).astype(float) + #---------------------------------------- + if self.parameters['G_Q'] is None: + self.G_Q = self.G # If no Q graph is provided, use G in its place + else: + self.G_Q = self.parameters['G_Q'] + # Quarantine Adjacency matrix: + if type(self.G_Q)==numpy.ndarray: + self.A_Q = scipy.sparse.csr_matrix(self.G_Q) + elif type(self.G_Q)==networkx.classes.graph.Graph: + self.A_Q = networkx.adj_matrix(self.G_Q) # adj_matrix gives scipy.sparse csr_matrix + else: + raise BaseException("Input an adjacency matrix or networkx object only.") + self.numNodes_Q = int(self.A_Q.shape[1]) + self.degree_Q = numpy.asarray(self.node_degrees(self.A_Q)).astype(float) + #---------------------------------------- + assert(self.numNodes == self.numNodes_Q), "The normal and quarantine adjacency graphs must be of the same size." + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Model parameters: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + self.beta = numpy.array(self.parameters['beta']).reshape((self.numNodes, 1)) if isinstance(self.parameters['beta'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['beta'], shape=(self.numNodes,1)) + self.sigma = numpy.array(self.parameters['sigma']).reshape((self.numNodes, 1)) if isinstance(self.parameters['sigma'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['sigma'], shape=(self.numNodes,1)) + self.gamma = numpy.array(self.parameters['gamma']).reshape((self.numNodes, 1)) if isinstance(self.parameters['gamma'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['gamma'], shape=(self.numNodes,1)) + self.mu_I = numpy.array(self.parameters['mu_I']).reshape((self.numNodes, 1)) if isinstance(self.parameters['mu_I'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['mu_I'], shape=(self.numNodes,1)) + self.alpha = numpy.array(self.parameters['alpha']).reshape((self.numNodes, 1)) if isinstance(self.parameters['alpha'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['alpha'], shape=(self.numNodes,1)) + self.xi = numpy.array(self.parameters['xi']).reshape((self.numNodes, 1)) if isinstance(self.parameters['xi'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['xi'], shape=(self.numNodes,1)) + self.mu_0 = numpy.array(self.parameters['mu_0']).reshape((self.numNodes, 1)) if isinstance(self.parameters['mu_0'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['mu_0'], shape=(self.numNodes,1)) + self.nu = numpy.array(self.parameters['nu']).reshape((self.numNodes, 1)) if isinstance(self.parameters['nu'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['nu'], shape=(self.numNodes,1)) + self.f = numpy.array(self.parameters['f']).reshape((self.numNodes, 1)) if isinstance(self.parameters['f'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['f'], shape=(self.numNodes,1)) + self.p = numpy.array(self.parameters['p']).reshape((self.numNodes, 1)) if isinstance(self.parameters['p'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['p'], shape=(self.numNodes,1)) + + self.rand_f = numpy.random.rand(self.f.shape[0], self.f.shape[1]) + + #---------------------------------------- + # Testing-related parameters: + #---------------------------------------- + self.beta_Q = (numpy.array(self.parameters['beta_Q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['beta_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['beta_Q'], shape=(self.numNodes,1))) if self.parameters['beta_Q'] is not None else self.beta + self.sigma_Q = (numpy.array(self.parameters['sigma_Q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['sigma_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['sigma_Q'], shape=(self.numNodes,1))) if self.parameters['sigma_Q'] is not None else self.sigma + self.gamma_Q = (numpy.array(self.parameters['gamma_Q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['gamma_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['gamma_Q'], shape=(self.numNodes,1))) if self.parameters['gamma_Q'] is not None else self.gamma + self.mu_Q = (numpy.array(self.parameters['mu_Q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['mu_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['mu_Q'], shape=(self.numNodes,1))) if self.parameters['mu_Q'] is not None else self.mu_I + self.alpha_Q = (numpy.array(self.parameters['alpha_Q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['alpha_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['alpha_Q'], shape=(self.numNodes,1))) if self.parameters['alpha_Q'] is not None else self.alpha + self.theta_E = numpy.array(self.parameters['theta_E']).reshape((self.numNodes, 1)) if isinstance(self.parameters['theta_E'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['theta_E'], shape=(self.numNodes,1)) + self.theta_I = numpy.array(self.parameters['theta_I']).reshape((self.numNodes, 1)) if isinstance(self.parameters['theta_I'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['theta_I'], shape=(self.numNodes,1)) + self.phi_E = numpy.array(self.parameters['phi_E']).reshape((self.numNodes, 1)) if isinstance(self.parameters['phi_E'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['phi_E'], shape=(self.numNodes,1)) + self.phi_I = numpy.array(self.parameters['phi_I']).reshape((self.numNodes, 1)) if isinstance(self.parameters['phi_I'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['phi_I'], shape=(self.numNodes,1)) + self.psi_E = numpy.array(self.parameters['psi_E']).reshape((self.numNodes, 1)) if isinstance(self.parameters['psi_E'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['psi_E'], shape=(self.numNodes,1)) + self.psi_I = numpy.array(self.parameters['psi_I']).reshape((self.numNodes, 1)) if isinstance(self.parameters['psi_I'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['psi_I'], shape=(self.numNodes,1)) + self.q = numpy.array(self.parameters['q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['q'], shape=(self.numNodes,1)) + + #---------------------------------------- + + self.beta_pairwise_mode = self.parameters['beta_pairwise_mode'] + + #---------------------------------------- + # Global transmission parameters: + #---------------------------------------- + if (self.beta_pairwise_mode == 'infected' or self.beta_pairwise_mode is None): + self.beta_global = numpy.full_like(self.beta, fill_value=numpy.mean(self.beta)) + self.beta_Q_global = numpy.full_like(self.beta_Q, fill_value=numpy.mean(self.beta_Q)) + elif self.beta_pairwise_mode == 'infectee': + self.beta_global = self.beta + self.beta_Q_global = self.beta_Q + elif self.beta_pairwise_mode == 'min': + self.beta_global = numpy.minimum(self.beta, numpy.mean(beta)) + self.beta_Q_global = numpy.minimum(self.beta_Q, numpy.mean(beta_Q)) + elif self.beta_pairwise_mode == 'max': + self.beta_global = numpy.maximum(self.beta, numpy.mean(beta)) + self.beta_Q_global = numpy.maximum(self.beta_Q, numpy.mean(beta_Q)) + elif self.beta_pairwise_mode == 'mean': + self.beta_global = (self.beta + numpy.full_like(self.beta, fill_value=numpy.mean(self.beta)))/2 + self.beta_Q_global = (self.beta_Q + numpy.full_like(self.beta_Q, fill_value=numpy.mean(self.beta_Q)))/2 + + #---------------------------------------- + # Local transmission parameters: + #---------------------------------------- + self.beta_local = self.beta if self.parameters['beta_local'] is None else numpy.array(self.parameters['beta_local']) if isinstance(self.parameters['beta_local'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['beta_local'], shape=(self.numNodes,1)) + self.beta_Q_local = self.beta_Q if self.parameters['beta_Q_local'] is None else numpy.array(self.parameters['beta_Q_local']) if isinstance(self.parameters['beta_Q_local'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['beta_Q_local'], shape=(self.numNodes,1)) + + #---------------------------------------- + if (self.beta_local.ndim == 2 and self.beta_local.shape[0] == self.numNodes and self.beta_local.shape[1] == self.numNodes): + self.A_beta_pairwise = self.beta_local + elif ((self.beta_local.ndim == 1 and self.beta_local.shape[0] == self.numNodes) or (self.beta_local.ndim == 2 and (self.beta_local.shape[0] == self.numNodes or self.beta_local.shape[1] == self.numNodes))): + self.beta_local = self.beta_local.reshape((self.numNodes,1)) + # Pre-multiply beta values by the adjacency matrix ("transmission weight connections") + A_beta_pairwise_byInfected = scipy.sparse.csr_matrix.multiply(self.A, self.beta_local.T).tocsr() + A_beta_pairwise_byInfectee = scipy.sparse.csr_matrix.multiply(self.A, self.beta_local).tocsr() + #------------------------------ + # Compute the effective pairwise beta values as a function of the infected/infectee pair: + if self.beta_pairwise_mode == 'infected': + self.A_beta_pairwise = A_beta_pairwise_byInfected + elif self.beta_pairwise_mode == 'infectee': + self.A_beta_pairwise = A_beta_pairwise_byInfectee + elif self.beta_pairwise_mode == 'min': + self.A_beta_pairwise = scipy.sparse.csr_matrix.minimum(A_beta_pairwise_byInfected, A_beta_pairwise_byInfectee) + elif self.beta_pairwise_mode == 'max': + self.A_beta_pairwise = scipy.sparse.csr_matrix.maximum(A_beta_pairwise_byInfected, A_beta_pairwise_byInfectee) + elif self.beta_pairwise_mode == 'mean' or self.beta_pairwise_mode is None: + self.A_beta_pairwise = (A_beta_pairwise_byInfected + A_beta_pairwise_byInfectee)/2 + else: + print("Unrecognized beta_pairwise_mode value (support for 'infected', 'infectee', 'min', 'max', and 'mean').") + else: + print("Invalid values given for beta_local (expected 1xN list/array or NxN 2d array)") + #---------------------------------------- + if (self.beta_Q_local.ndim == 2 and self.beta_Q_local.shape[0] == self.numNodes and self.beta_Q_local.shape[1] == self.numNodes): + self.A_Q_beta_Q_pairwise = self.beta_Q_local + elif ((self.beta_Q_local.ndim == 1 and self.beta_Q_local.shape[0] == self.numNodes) or (self.beta_Q_local.ndim == 2 and (self.beta_Q_local.shape[0] == self.numNodes or self.beta_Q_local.shape[1] == self.numNodes))): + self.beta_Q_local = self.beta_Q_local.reshape((self.numNodes,1)) + # Pre-multiply beta_Q values by the isolation adjacency matrix ("transmission weight connections") + A_Q_beta_Q_pairwise_byInfected = scipy.sparse.csr_matrix.multiply(self.A_Q, self.beta_Q_local.T).tocsr() + A_Q_beta_Q_pairwise_byInfectee = scipy.sparse.csr_matrix.multiply(self.A_Q, self.beta_Q_local).tocsr() + #------------------------------ + # Compute the effective pairwise beta values as a function of the infected/infectee pair: + if self.beta_pairwise_mode == 'infected': + self.A_Q_beta_Q_pairwise = A_Q_beta_Q_pairwise_byInfected + elif self.beta_pairwise_mode == 'infectee': + self.A_Q_beta_Q_pairwise = A_Q_beta_Q_pairwise_byInfectee + elif self.beta_pairwise_mode == 'min': + self.A_Q_beta_Q_pairwise = scipy.sparse.csr_matrix.minimum(A_Q_beta_Q_pairwise_byInfected, A_Q_beta_Q_pairwise_byInfectee) + elif self.beta_pairwise_mode == 'max': + self.A_Q_beta_Q_pairwise = scipy.sparse.csr_matrix.maximum(A_Q_beta_Q_pairwise_byInfected, A_Q_beta_Q_pairwise_byInfectee) + elif (self.beta_pairwise_mode == 'mean' or self.beta_pairwise_mode is None): + self.A_Q_beta_Q_pairwise = (A_Q_beta_Q_pairwise_byInfected + A_Q_beta_Q_pairwise_byInfectee)/2 + else: + print("Unrecognized beta_pairwise_mode value (support for 'infected', 'infectee', 'min', 'max', and 'mean').") + else: + print("Invalid values given for beta_Q_local (expected 1xN list/array or NxN 2d array)") + #---------------------------------------- + + #---------------------------------------- + # Degree-based transmission scaling parameters: + #---------------------------------------- + self.delta_pairwise_mode = self.parameters['delta_pairwise_mode'] + with numpy.errstate(divide='ignore'): # ignore log(0) warning, then convert log(0) = -inf -> 0.0 + self.delta = numpy.log(self.degree)/numpy.log(numpy.mean(self.degree)) if self.parameters['delta'] is None else numpy.array(self.parameters['delta']) if isinstance(self.parameters['delta'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['delta'], shape=(self.numNodes,1)) + self.delta_Q = numpy.log(self.degree_Q)/numpy.log(numpy.mean(self.degree_Q)) if self.parameters['delta_Q'] is None else numpy.array(self.parameters['delta_Q']) if isinstance(self.parameters['delta_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['delta_Q'], shape=(self.numNodes,1)) + self.delta[numpy.isneginf(self.delta)] = 0.0 + self.delta_Q[numpy.isneginf(self.delta_Q)] = 0.0 + #---------------------------------------- + if (self.delta.ndim == 2 and self.delta.shape[0] == self.numNodes and self.delta.shape[1] == self.numNodes): + self.A_delta_pairwise = self.delta + elif ((self.delta.ndim == 1 and self.delta.shape[0] == self.numNodes) or (self.delta.ndim == 2 and (self.delta.shape[0] == self.numNodes or self.delta.shape[1] == self.numNodes))): + self.delta = self.delta.reshape((self.numNodes,1)) + # Pre-multiply delta values by the adjacency matrix ("transmission weight connections") + A_delta_pairwise_byInfected = scipy.sparse.csr_matrix.multiply(self.A, self.delta.T).tocsr() + A_delta_pairwise_byInfectee = scipy.sparse.csr_matrix.multiply(self.A, self.delta).tocsr() + #------------------------------ + # Compute the effective pairwise delta values as a function of the infected/infectee pair: + if self.delta_pairwise_mode == 'infected': + self.A_delta_pairwise = A_delta_pairwise_byInfected + elif self.delta_pairwise_mode == 'infectee': + self.A_delta_pairwise = A_delta_pairwise_byInfectee + elif self.delta_pairwise_mode == 'min': + self.A_delta_pairwise = scipy.sparse.csr_matrix.minimum(A_delta_pairwise_byInfected, A_delta_pairwise_byInfectee) + elif self.delta_pairwise_mode == 'max': + self.A_delta_pairwise = scipy.sparse.csr_matrix.maximum(A_delta_pairwise_byInfected, A_delta_pairwise_byInfectee) + elif self.delta_pairwise_mode == 'mean': + self.A_delta_pairwise = (A_delta_pairwise_byInfected + A_delta_pairwise_byInfectee)/2 + elif self.delta_pairwise_mode is None: + self.A_delta_pairwise = self.A + else: + print("Unrecognized delta_pairwise_mode value (support for 'infected', 'infectee', 'min', 'max', and 'mean').") + else: + print("Invalid values given for delta (expected 1xN list/array or NxN 2d array)") + #---------------------------------------- + if (self.delta_Q.ndim == 2 and self.delta_Q.shape[0] == self.numNodes and self.delta_Q.shape[1] == self.numNodes): + self.A_Q_delta_Q_pairwise = self.delta_Q + elif ((self.delta_Q.ndim == 1 and self.delta_Q.shape[0] == self.numNodes) or (self.delta_Q.ndim == 2 and (self.delta_Q.shape[0] == self.numNodes or self.delta_Q.shape[1] == self.numNodes))): + self.delta_Q = self.delta_Q.reshape((self.numNodes,1)) + # Pre-multiply delta_Q values by the isolation adjacency matrix ("transmission weight connections") + A_Q_delta_Q_pairwise_byInfected = scipy.sparse.csr_matrix.multiply(self.A_Q, self.delta_Q).tocsr() + A_Q_delta_Q_pairwise_byInfectee = scipy.sparse.csr_matrix.multiply(self.A_Q, self.delta_Q.T).tocsr() + #------------------------------ + # Compute the effective pairwise delta values as a function of the infected/infectee pair: + if self.delta_pairwise_mode == 'infected': + self.A_Q_delta_Q_pairwise = A_Q_delta_Q_pairwise_byInfected + elif self.delta_pairwise_mode == 'infectee': + self.A_Q_delta_Q_pairwise = A_Q_delta_Q_pairwise_byInfectee + elif self.delta_pairwise_mode == 'min': + self.A_Q_delta_Q_pairwise = scipy.sparse.csr_matrix.minimum(A_Q_delta_Q_pairwise_byInfected, A_Q_delta_Q_pairwise_byInfectee) + elif self.delta_pairwise_mode == 'max': + self.A_Q_delta_Q_pairwise = scipy.sparse.csr_matrix.maximum(A_Q_delta_Q_pairwise_byInfected, A_Q_delta_Q_pairwise_byInfectee) + elif self.delta_pairwise_mode == 'mean': + self.A_Q_delta_Q_pairwise = (A_Q_delta_Q_pairwise_byInfected + A_Q_delta_Q_pairwise_byInfectee)/2 + elif (self.delta_pairwise_mode is None): + self.A_Q_delta_Q_pairwise = self.A + else: + print("Unrecognized delta_pairwise_mode value (support for 'infected', 'infectee', 'min', 'max', and 'mean').") + else: + print("Invalid values given for delta_Q (expected 1xN list/array or NxN 2d array)") + + #---------------------------------------- + # Pre-calculate the pairwise delta*beta values: + #---------------------------------------- + self.A_deltabeta = scipy.sparse.csr_matrix.multiply(self.A_delta_pairwise, self.A_beta_pairwise) + self.A_Q_deltabeta_Q = scipy.sparse.csr_matrix.multiply(self.A_Q_delta_Q_pairwise, self.A_Q_beta_Q_pairwise) + + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def node_degrees(self, Amat): + return Amat.sum(axis=0).reshape(self.numNodes,1) # sums of adj matrix cols + + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def total_num_susceptible(self, t_idx=None): + if t_idx is None: + return self.numS[:] + else: + return self.numS[t_idx] + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def total_num_infected(self, t_idx=None): + if t_idx is None: + return self.numE[:] + self.numI[:] + self.numQ_E[:] + self.numQ_I[:] + else: + return self.numE[t_idx] + self.numI[t_idx] + self.numQ_E[t_idx] + self.numQ_I[t_idx] + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def total_num_isolated(self, t_idx=None): + if t_idx is None: + return self.numQ_E[:] + self.numQ_I[:] + else: + return self.numQ_E[t_idx] + self.numQ_I[t_idx] + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def total_num_tested(self, t_idx=None): + if t_idx is None: + return self.numTested[:] + else: + return self.numTested[t_idx] + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def total_num_positive(self, t_idx=None): + if t_idx is None: + return self.numPositive[:] + else: + return self.numPositive[t_idx] + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def total_num_recovered(self, t_idx=None): + if t_idx is None: + return self.numR[:] + else: + return self.numR[t_idx] + + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def calc_propensities(self): + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Pre-calculate matrix multiplication terms that may be used in multiple propensity calculations, + # and check to see if their computation is necessary before doing the multiplication + #------------------------------------ + + self.transmissionTerms_I = numpy.zeros(shape=(self.numNodes,1)) + if numpy.any(self.numI[self.tidx]): + self.transmissionTerms_I = numpy.asarray(scipy.sparse.csr_matrix.dot(self.A_deltabeta, self.X==self.I)) + + #------------------------------------ + + self.transmissionTerms_Q = numpy.zeros(shape=(self.numNodes,1)) + if numpy.any(self.numQ_I[self.tidx]): + self.transmissionTerms_Q = numpy.asarray(scipy.sparse.csr_matrix.dot(self.A_Q_deltabeta_Q, self.X==self.Q_I)) + + #------------------------------------ + + numContacts_Q = numpy.zeros(shape=(self.numNodes,1)) + if numpy.any(self.positive) and (numpy.any(self.phi_E) or numpy.any(self.phi_I)): + numContacts_Q = numpy.asarray(scipy.sparse.csr_matrix.dot(self.A, ((self.positive)&(self.X!=self.R)&(self.X!=self.F)))) + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + propensities_StoE = (self.alpha * + (self.p*((self.beta_global*self.numI[self.tidx] + self.q*self.beta_Q_global*self.numQ_I[self.tidx])/self.N[self.tidx]) + + (1-self.p)*(numpy.divide(self.transmissionTerms_I, self.degree, out=numpy.zeros_like(self.degree), where=self.degree!=0) + +numpy.divide(self.transmissionTerms_Q, self.degree_Q, out=numpy.zeros_like(self.degree_Q), where=self.degree_Q!=0))) + )*(self.X==self.S) + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + if self.transition_mode == 'time_in_state': + + propensities_EtoI = 1e5 * ((self.X==self.E) & numpy.greater(self.timer_state, 1/self.sigma)) + + propensities_ItoR = 1e5 * ((self.X==self.I) & numpy.greater(self.timer_state, 1/self.gamma) & numpy.greater_equal(self.rand_f, self.f)) + + propensities_ItoF = 1e5 * ((self.X==self.I) & numpy.greater(self.timer_state, 1/self.mu_I) & numpy.less(self.rand_f, self.f)) + + propensities_EtoQE = numpy.zeros_like(propensities_StoE) + + propensities_ItoQI = numpy.zeros_like(propensities_StoE) + + propensities_QEtoQI = 1e5 * ((self.X==self.Q_E) & numpy.greater(self.timer_state, 1/self.sigma_Q)) + + propensities_QItoR = 1e5 * ((self.X==self.Q_I) & numpy.greater(self.timer_state, 1/self.gamma_Q) & numpy.greater_equal(self.rand_f, self.f)) + + propensities_QItoF = 1e5 * ((self.X==self.Q_I) & numpy.greater(self.timer_state, 1/self.mu_Q) & numpy.less(self.rand_f, self.f)) + + propensities_RtoS = 1e5 * ((self.X==self.R) & numpy.greater(self.timer_state, 1/self.xi)) + + propensities__toS = 1e5 * ((self.X!=self.F) & numpy.greater(self.timer_state, 1/self.nu)) + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + else: # exponential_rates + + propensities_EtoI = self.sigma * (self.X==self.E) + + propensities_ItoR = self.gamma * ((self.X==self.I) & (numpy.greater_equal(self.rand_f, self.f))) + + propensities_ItoF = self.mu_I * ((self.X==self.I) & (numpy.less(self.rand_f, self.f))) + + propensities_EtoQE = (self.theta_E + self.phi_E*numContacts_Q)*self.psi_E * (self.X==self.E) + + propensities_ItoQI = (self.theta_I + self.phi_I*numContacts_Q)*self.psi_I * (self.X==self.I) + + propensities_QEtoQI = self.sigma_Q * (self.X==self.Q_E) + + propensities_QItoR = self.gamma_Q * ((self.X==self.Q_I) & (numpy.greater_equal(self.rand_f, self.f))) + + propensities_QItoF = self.mu_Q * ((self.X==self.Q_I) & (numpy.less(self.rand_f, self.f))) + + propensities_RtoS = self.xi * (self.X==self.R) + + propensities__toS = self.nu * (self.X!=self.F) + + + + + + + + + + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + propensities = numpy.hstack([propensities_StoE, propensities_EtoI, + propensities_ItoR, propensities_ItoF, + propensities_EtoQE, propensities_ItoQI, propensities_QEtoQI, + propensities_QItoR, propensities_QItoF, + propensities_RtoS, propensities__toS]) + + columns = ['StoE', 'EtoI', 'ItoR', 'ItoF', 'EtoQE', 'ItoQI', 'QEtoQI', 'QItoR', 'QItoF', 'RtoS', '_toS'] + + return propensities, columns + + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def set_isolation(self, node, isolate): + # Move this node in/out of the appropriate isolation state: + if isolate == True: + if self.X[node] == self.E: + self.X[node] = self.Q_E + self.timer_state = 0 + elif self.X[node] == self.I: + self.X[node] = self.Q_I + self.timer_state = 0 + elif isolate == False: + if self.X[node] == self.Q_E: + self.X[node] = self.E + self.timer_state = 0 + elif self.X[node] == self.Q_I: + self.X[node] = self.I + self.timer_state = 0 + # Reset the isolation timer: + self.timer_isolation[node] = 0 + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def set_tested(self, node, tested): + self.tested[node] = tested + self.testedInCurrentState[node] = tested + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def set_positive(self, node, positive): + self.positive[node] = positive + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def introduce_exposures(self, num_new_exposures): + exposedNodes = numpy.random.choice(range(self.numNodes), size=num_new_exposures, replace=False) + for exposedNode in exposedNodes: + if self.X[exposedNode] == self.S: + self.X[exposedNode] = self.E + + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def increase_data_series_length(self): + self.tseries = numpy.pad(self.tseries, [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.numS = numpy.pad(self.numS, [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.numE = numpy.pad(self.numE, [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.numI = numpy.pad(self.numI, [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.numR = numpy.pad(self.numR, [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.numF = numpy.pad(self.numF, [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.numQ_E = numpy.pad(self.numQ_E, [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.numQ_I = numpy.pad(self.numQ_I, [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.N = numpy.pad(self.N, [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.numTested = numpy.pad(self.numTested, [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.numPositive = numpy.pad(self.numPositive, [(0, 6*self.numNodes)], mode='constant', constant_values=0) + + if self.store_Xseries: + self.Xseries = numpy.pad(self.Xseries, [(0, 6*self.numNodes), (0,0)], mode='constant', constant_values=0) + + if self.nodeGroupData: + for groupName in self.nodeGroupData: + self.nodeGroupData[groupName]['numS'] = numpy.pad(self.nodeGroupData[groupName]['numS'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.nodeGroupData[groupName]['numE'] = numpy.pad(self.nodeGroupData[groupName]['numE'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.nodeGroupData[groupName]['numI'] = numpy.pad(self.nodeGroupData[groupName]['numI'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.nodeGroupData[groupName]['numR'] = numpy.pad(self.nodeGroupData[groupName]['numR'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.nodeGroupData[groupName]['numF'] = numpy.pad(self.nodeGroupData[groupName]['numF'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.nodeGroupData[groupName]['numQ_E'] = numpy.pad(self.nodeGroupData[groupName]['numQ_E'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.nodeGroupData[groupName]['numQ_I'] = numpy.pad(self.nodeGroupData[groupName]['numQ_I'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.nodeGroupData[groupName]['N'] = numpy.pad(self.nodeGroupData[groupName]['N'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.nodeGroupData[groupName]['numTested'] = numpy.pad(self.nodeGroupData[groupName]['numTested'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) + self.nodeGroupData[groupName]['numPositive'] = numpy.pad(self.nodeGroupData[groupName]['numPositive'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) + + return None + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def finalize_data_series(self): + self.tseries = numpy.array(self.tseries, dtype=float)[:self.tidx+1] + self.numS = numpy.array(self.numS, dtype=float)[:self.tidx+1] + self.numE = numpy.array(self.numE, dtype=float)[:self.tidx+1] + self.numI = numpy.array(self.numI, dtype=float)[:self.tidx+1] + self.numR = numpy.array(self.numR, dtype=float)[:self.tidx+1] + self.numF = numpy.array(self.numF, dtype=float)[:self.tidx+1] + self.numQ_E = numpy.array(self.numQ_E, dtype=float)[:self.tidx+1] + self.numQ_I = numpy.array(self.numQ_I, dtype=float)[:self.tidx+1] + self.N = numpy.array(self.N, dtype=float)[:self.tidx+1] + self.numTested = numpy.array(self.numTested, dtype=float)[:self.tidx+1] + self.numPositive = numpy.array(self.numPositive, dtype=float)[:self.tidx+1] + + if self.store_Xseries: + self.Xseries = self.Xseries[:self.tidx+1, :] + + if self.nodeGroupData: + for groupName in self.nodeGroupData: + self.nodeGroupData[groupName]['numS'] = numpy.array(self.nodeGroupData[groupName]['numS'], dtype=float)[:self.tidx+1] + self.nodeGroupData[groupName]['numE'] = numpy.array(self.nodeGroupData[groupName]['numE'], dtype=float)[:self.tidx+1] + self.nodeGroupData[groupName]['numI'] = numpy.array(self.nodeGroupData[groupName]['numI'], dtype=float)[:self.tidx+1] + self.nodeGroupData[groupName]['numR'] = numpy.array(self.nodeGroupData[groupName]['numR'], dtype=float)[:self.tidx+1] + self.nodeGroupData[groupName]['numF'] = numpy.array(self.nodeGroupData[groupName]['numF'], dtype=float)[:self.tidx+1] + self.nodeGroupData[groupName]['numQ_E'] = numpy.array(self.nodeGroupData[groupName]['numQ_E'], dtype=float)[:self.tidx+1] + self.nodeGroupData[groupName]['numQ_I'] = numpy.array(self.nodeGroupData[groupName]['numQ_I'], dtype=float)[:self.tidx+1] + self.nodeGroupData[groupName]['N'] = numpy.array(self.nodeGroupData[groupName]['N'], dtype=float)[:self.tidx+1] + self.nodeGroupData[groupName]['numTested'] = numpy.array(self.nodeGroupData[groupName]['numTested'], dtype=float)[:self.tidx+1] + self.nodeGroupData[groupName]['numPositive'] = numpy.array(self.nodeGroupData[groupName]['numPositive'], dtype=float)[:self.tidx+1] + + return None + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def run_iteration(self): + + if self.tidx >= len(self.tseries) - 1: + # Room has run out in the timeseries storage arrays; double the size of these arrays: + self.increase_data_series_length() + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Generate 2 random numbers uniformly distributed in (0,1) + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + r1 = numpy.random.rand() + r2 = numpy.random.rand() + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Calculate propensities + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + propensities, transitionTypes = self.calc_propensities() + + if propensities.sum() > 0: + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Calculate alpha + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + propensities_flat = propensities.ravel(order='F') + cumsum = propensities_flat.cumsum() + alpha = propensities_flat.sum() + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Compute the time until the next event takes place + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + tau = (1/alpha)*numpy.log(float(1/r1)) + self.t += tau + self.timer_state += tau + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Compute which event takes place + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + transitionIdx = numpy.searchsorted(cumsum,r2*alpha) + transitionNode = transitionIdx % self.numNodes + transitionType = transitionTypes[ int(transitionIdx/self.numNodes) ] + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Perform updates triggered by rate propensities: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + assert(self.X[transitionNode] == self.transitions[transitionType]['currentState'] and self.X[transitionNode]!=self.F), "Assertion error: Node "+str(transitionNode)+" has unexpected current state "+str(self.X[transitionNode])+" given the intended transition of "+str(transitionType)+"." + self.X[transitionNode] = self.transitions[transitionType]['newState'] + + self.testedInCurrentState[transitionNode] = False + + self.timer_state = 0.0 + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + # Save information about infection events when they occur: + if transitionType == 'StoE': + transitionNode_GNbrs = list(self.G[transitionNode].keys()) + transitionNode_GQNbrs = list(self.G_Q[transitionNode].keys()) + self.infectionsLog.append({ 't': self.t, + 'infected_node': transitionNode, + 'infection_type': transitionType, + 'infected_node_degree': self.degree[transitionNode], + 'local_contact_nodes': transitionNode_GNbrs, + 'local_contact_node_states': self.X[transitionNode_GNbrs].flatten(), + 'isolation_contact_nodes': transitionNode_GQNbrs, + 'isolation_contact_node_states':self.X[transitionNode_GQNbrs].flatten() }) + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + if (transitionType in ['EtoQE', 'ItoQI']): + self.set_positive(node=transitionNode, positive=True) + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + else: + + tau = 0.01 + self.t += tau + self.timer_state += tau + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + self.tidx += 1 + + self.tseries[self.tidx] = self.t + self.numS[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.S), a_min=0, a_max=self.numNodes) + self.numE[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.E), a_min=0, a_max=self.numNodes) + self.numI[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.I), a_min=0, a_max=self.numNodes) + self.numF[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.F), a_min=0, a_max=self.numNodes) + self.numQ_E[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.Q_E), a_min=0, a_max=self.numNodes) + self.numQ_I[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.Q_I), a_min=0, a_max=self.numNodes) + self.numTested[self.tidx] = numpy.clip(numpy.count_nonzero(self.tested), a_min=0, a_max=self.numNodes) + self.numPositive[self.tidx] = numpy.clip(numpy.count_nonzero(self.positive), a_min=0, a_max=self.numNodes) + + self.N[self.tidx] = numpy.clip((self.numNodes - self.numF[self.tidx]), a_min=0, a_max=self.numNodes) + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Update testing and isolation statuses + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + isolatedNodes = numpy.argwhere((self.X==self.Q_E)|(self.X==self.Q_I))[:,0].flatten() + self.timer_isolation[isolatedNodes] = self.timer_isolation[isolatedNodes] + tau + + nodesExitingIsolation = numpy.argwhere(self.timer_isolation >= self.isolationTime) + for isoNode in nodesExitingIsolation: + self.set_isolation(node=isoNode, isolate=False) + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Store system states + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + if self.store_Xseries: + self.Xseries[self.tidx,:] = self.X.T + + if self.nodeGroupData: + for groupName in self.nodeGroupData: + self.nodeGroupData[groupName]['numS'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.S) + self.nodeGroupData[groupName]['numE'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.E) + self.nodeGroupData[groupName]['numI'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.I) + self.nodeGroupData[groupName]['numR'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.R) + self.nodeGroupData[groupName]['numF'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.F) + self.nodeGroupData[groupName]['numQ_E'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_E) + self.nodeGroupData[groupName]['numQ_I'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_I) + self.nodeGroupData[groupName]['N'][self.tidx] = numpy.clip((self.nodeGroupData[groupName]['numS'][0] + self.nodeGroupData[groupName]['numE'][0] + self.nodeGroupData[groupName]['numI'][0] + self.nodeGroupData[groupName]['numQ_E'][0] + self.nodeGroupData[groupName]['numQ_I'][0] + self.nodeGroupData[groupName]['numR'][0]), a_min=0, a_max=self.numNodes) + self.nodeGroupData[groupName]['numTested'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.tested) + self.nodeGroupData[groupName]['numPositive'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.positive) + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Terminate if tmax reached or num infections is 0: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + if (self.t >= self.tmax or (self.total_num_infected(self.tidx) < 1 and self.total_num_isolated(self.tidx) < 1)): + self.finalize_data_series() + return False + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + return True + + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + def run(self, T, checkpoints=None, print_interval=10, verbose='t'): + if T > 0: + self.tmax += T + else: + return False + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Pre-process checkpoint values: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + if checkpoints: + numCheckpoints = len(checkpoints['t']) + for chkpt_param, chkpt_values in checkpoints.items(): + assert(isinstance(chkpt_values, (list, numpy.ndarray)) and len(chkpt_values)==numCheckpoints), "Expecting a list of values with length equal to number of checkpoint times ("+str(numCheckpoints)+") for each checkpoint parameter." + checkpointIdx = numpy.searchsorted(checkpoints['t'], self.t) # Finds 1st index in list greater than given val + if checkpointIdx >= numCheckpoints: + # We are out of checkpoints, stop checking them: + checkpoints = None + else: + checkpointTime = checkpoints['t'][checkpointIdx] + + #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% + # Run the simulation loop: + #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% + print_reset = True + running = True + while running: + + running = self.run_iteration() + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Handle checkpoints if applicable: + if checkpoints: + if self.t >= checkpointTime: + if verbose is not False: + print("[Checkpoint: Updating parameters]") + # A checkpoint has been reached, update param values: + for param in list(self.parameters.keys()): + if param in list(checkpoints.keys()): + self.parameters.update({param: checkpoints[param][checkpointIdx]}) + # Update parameter data structures and scenario flags: + self.update_parameters() + # Update the next checkpoint time: + checkpointIdx = numpy.searchsorted(checkpoints['t'], self.t) # Finds 1st index in list greater than given val + if checkpointIdx >= numCheckpoints: + # We are out of checkpoints, stop checking them: + checkpoints = None + else: + checkpointTime = checkpoints['t'][checkpointIdx] + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + if print_interval: + if (print_reset and (int(self.t) % print_interval == 0)): + if verbose == "t": + print("t = %.2f" % self.t) + if verbose == True: + print("t = %.2f" % self.t) + print("\t S = " + str(self.numS[self.tidx])) + print("\t E = " + str(self.numE[self.tidx])) + print("\t I = " + str(self.numI[self.tidx])) + print("\t R = " + str(self.numR[self.tidx])) + print("\t F = " + str(self.numF[self.tidx])) + print("\t Q_E = " + str(self.numQ_E[self.tidx])) + print("\t Q_I = " + str(self.numQ_I[self.tidx])) + print_reset = False + elif (not print_reset and (int(self.t) % 10 != 0)): + print_reset = True + + return True diff --git a/seirsplus/networks.py b/seirsplus/networks.py index c014a6a..26aa421 100644 --- a/seirsplus/networks.py +++ b/seirsplus/networks.py @@ -12,7 +12,7 @@ def generate_workplace_contact_network(num_cohorts=1, num_nodes_per_cohort=100, num_teams_per_cohort=10, mean_intracohort_degree=6, pct_contacts_intercohort=0.2, - farz_params={'alpha':5.0, 'gamma':5.0, 'beta':0.5, 'r':1, 'q':0.0, 'phi':10, + farz_params={'alpha':5.0, 'gamma':5.0, 'beta':0.5, 'r':1, 'q':0.0, 'phi':10, 'b':0, 'epsilon':1e-6, 'directed': False, 'weighted': False}, distancing_scales=[]): @@ -41,7 +41,7 @@ def generate_workplace_contact_network(num_cohorts=1, num_nodes_per_cohort=100, try: teams_indices['c'+str(i)+'-t'+str(team)].append(node) except KeyError: - teams_indices['c'+str(i)+'-t'+str(team)] = [node] + teams_indices['c'+str(i)+'-t'+str(team)] = [node] #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Establish inter-cohort contacts: @@ -64,14 +64,14 @@ def generate_workplace_contact_network(num_cohorts=1, num_nodes_per_cohort=100, cohorts_indices['c'+str(c)] = list(range(cohortStartIdx, cohortFinalIdx)) for team, indices in teams_indices.items(): - if('c'+str(c) in team): + if 'c' + str(c) in team: teams_indices[team] = [idx+cohortStartIdx for idx in indices] for i in list(range(cohortNetwork.number_of_nodes())): i_intraCohortDegree = cohortNetwork.degree[i] i_interCohortDegree = int( ((1/(1-pct_contacts_intercohort))*i_intraCohortDegree)-i_intraCohortDegree ) # Add intercohort edges: - if(len(cohortNetworks) > 1): + if len(cohortNetworks) > 1: for d in list(range(i_interCohortDegree)): j = numpy.random.choice(list(range(0, cohortStartIdx))+list(range(cohortFinalIdx+1, N))) workplaceNetwork.add_edge(i, j) @@ -113,7 +113,7 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F meanNumU20PerHousehold_givenU20 = household_stats['mean_num_under20_givenAtLeastOneUnder20'] #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Define major age groups (under 20, between 20-60, over 60), + # Define major age groups (under 20, between 20-60, over 60), # and calculate age distributions conditional on belonging (or not) to one of these groups: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ ageBrackets_U20 = ['0-9', '10-19'] @@ -138,14 +138,14 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Calculate the probabilities of a household having members in the major age groups, + # Calculate the probabilities of a household having members in the major age groups, # conditional on single/multi-occupancy: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ prob_u20 = pctHouseholdsWithMember_U20 # probability of household having at least 1 member under 20 prob_o60 = pctHouseholdsWithMember_O60 # probability of household having at least 1 member over 60 prob_eq1 = household_size_distn[1] # probability of household having 1 member prob_gt1 = 1 - prob_eq1 # probability of household having greater than 1 member - householdSituations_prob = {} + householdSituations_prob = {} householdSituations_prob['u20_o60_eq1'] = 0 # can't have both someone under 20 and over 60 in a household with 1 member householdSituations_prob['u20_NOTo60_eq1'] = 0 # assume no one under 20 lives on their own (data suggests <1% actually do) householdSituations_prob['NOTu20_o60_eq1'] = pctHouseholdsWithMember_O60_givenEq1*prob_eq1 @@ -165,16 +165,16 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F households = [] # List of dicts storing household data structures and metadata homelessNodes = N # Number of individuals to place in households curMemberIndex = 0 - while(homelessNodes > 0): - + while homelessNodes > 0: + household = {} household['situation'] = numpy.random.choice(list(householdSituations_prob.keys()), p=list(householdSituations_prob.values())) household['ageBrackets'] = [] - if(household['situation'] == 'NOTu20_o60_eq1'): - + if household['situation'] == 'NOTu20_o60_eq1': + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Household size is definitely 1 household['size'] = 1 @@ -182,17 +182,17 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # There is only 1 member in this household, and they are OVER 60; add them: household['ageBrackets'].append( numpy.random.choice(list(age_distn_givenO60.keys()), p=list(age_distn_givenO60.values())) ) - - elif(household['situation'] == 'NOTu20_NOTo60_eq1'): - + + elif household['situation'] == 'NOTu20_NOTo60_eq1': + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Household size is definitely 1 household['size'] = 1 - + # There is only 1 member in this household, and they are BETWEEN 20-60; add them: household['ageBrackets'].append( numpy.random.choice(list(age_distn_given20to60.keys()), p=list(age_distn_given20to60.values())) ) - - elif(household['situation'] == 'u20_o60_gt1'): + + elif household['situation'] == 'u20_o60_gt1': #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Draw a household size (given the situation, there's at least 2 members): @@ -210,13 +210,13 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # There's definitely one OVER 60 in this household, add an appropriate age bracket: household['ageBrackets'].append( numpy.random.choice(list(age_distn_givenO60.keys()), p=list(age_distn_givenO60.values())) ) - + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Any remaining members can be any age EXCLUDING Under 20 (all U20s already added): for m in range(household['size'] - len(household['ageBrackets'])): household['ageBrackets'].append( numpy.random.choice(list(age_distn_givenNOTU20.keys()), p=list(age_distn_givenNOTU20.values())) ) - elif(household['situation'] == 'u20_NOTo60_gt1'): + elif household['situation'] == 'u20_NOTo60_gt1': #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Draw a household size (given the situation, there's at least 2 members): @@ -240,7 +240,7 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F for m in range(household['size'] - len(household['ageBrackets'])): household['ageBrackets'].append( numpy.random.choice(list(age_distn_given20to60.keys()), p=list(age_distn_given20to60.values())) ) - elif(household['situation'] == 'NOTu20_o60_gt1'): + elif household['situation'] == 'NOTu20_o60_gt1': #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Draw a household size (given the situation, there's at least 2 members): @@ -252,14 +252,14 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # There's definitely one OVER 60 in this household, add an appropriate age bracket: household['ageBrackets'].append( numpy.random.choice(list(age_distn_givenO60.keys()), p=list(age_distn_givenO60.values())) ) - + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Any remaining members can be any age EXCLUDING UNDER 20: for m in range(household['size'] - len(household['ageBrackets'])): household['ageBrackets'].append( numpy.random.choice(list(age_distn_givenNOTU20.keys()), p=list(age_distn_givenNOTU20.values())) ) - elif(household['situation'] == 'NOTu20_NOTo60_gt1'): - + elif household['situation'] == 'NOTu20_NOTo60_gt1': + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Draw a household size (given the situation, there's at least 2 members): household['size'] = min(homelessNodes, max(2, numpy.random.choice(list(household_size_distn_givenGT1), p=list(household_size_distn_givenGT1.values()))) ) @@ -275,12 +275,12 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F for m in range(household['size'] - len(household['ageBrackets'])): household['ageBrackets'].append( numpy.random.choice(list(age_distn_given20to60.keys()), p=list(age_distn_given20to60.values())) ) - # elif(household['situation'] == 'u20_NOTo60_eq1'): + # elif (household['situation'] == 'u20_NOTo60_eq1'): # impossible by assumption - # elif(household['situation'] == 'u20_o60_eq1'): + # elif (household['situation'] == 'u20_o60_eq1'): # impossible - if(len(household['ageBrackets']) == household['size']): + if len(household['ageBrackets']) == household['size']: homelessNodes -= household['size'] @@ -310,23 +310,23 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F print("Num households: " +str(numHouseholds)) print("mean household size: " + str(meanHouseholdSize)) print() - - if(verbose): + + if verbose: print("Generated percent households with at least one member Under 20:") checkval = len([household for household in households if not set(household['ageBrackets']).isdisjoint(ageBrackets_U20)])/numHouseholds target = pctHouseholdsWithMember_U20 print("%.4f\t\t(%.4f from target)" % (checkval, checkval - target)) - + print("Generated percent households with at least one Over 60") checkval = len([household for household in households if not set(household['ageBrackets']).isdisjoint(ageBrackets_O60)])/numHouseholds target = pctHouseholdsWithMember_O60 print("%.4f\t\t(%.4f from target)" % (checkval, checkval - target)) - + print("Generated percent households with at least one Under 20 AND Over 60") checkval = len([household for household in households if not set(household['ageBrackets']).isdisjoint(ageBrackets_O60) and not set(household['ageBrackets']).isdisjoint(ageBrackets_U20)])/numHouseholds target = pctHouseholdsWithMember_U20andO60 print("%.4f\t\t(%.4f from target)" % (checkval, checkval - target)) - + print("Generated percent households with 1 total member who is Over 60") checkval = numpy.sum([1 for household in households if household['size']==1 and not set(household['ageBrackets']).isdisjoint(ageBrackets_O60)])/numHouseholds target = pctHouseholdsWithMember_O60_givenEq1*prob_eq1 @@ -344,7 +344,7 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F # Generate Contact Networks ######################################### ######################################### - + ######################################### # Generate baseline (no intervention) contact network: ######################################### @@ -352,13 +352,13 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Define the age groups and desired mean degree for each graph layer: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(layer_info is None): + if layer_info is None: # Use the following default data if none is provided: # Data source: https://www.medrxiv.org/content/10.1101/2020.03.19.20039107v1 layer_info = { '0-9': {'ageBrackets': ['0-9'], 'meanDegree': 8.6, 'meanDegree_CI': (0.0, 17.7) }, '10-19': {'ageBrackets': ['10-19'], 'meanDegree': 16.2, 'meanDegree_CI': (12.5, 19.8) }, - '20-59': {'ageBrackets': ['20-29', '30-39', '40-49', '50-59'], - 'meanDegree': ((age_distn_given20to60['20-29']+age_distn_given20to60['30-39'])*15.3 + (age_distn_given20to60['40-49']+age_distn_given20to60['50-59'])*13.8), + '20-59': {'ageBrackets': ['20-29', '30-39', '40-49', '50-59'], + 'meanDegree': ((age_distn_given20to60['20-29']+age_distn_given20to60['30-39'])*15.3 + (age_distn_given20to60['40-49']+age_distn_given20to60['50-59'])*13.8), 'meanDegree_CI': ( ((age_distn_given20to60['20-29']+age_distn_given20to60['30-39'])*12.6 + (age_distn_given20to60['40-49']+age_distn_given20to60['50-59'])*11.0), ((age_distn_given20to60['20-29']+age_distn_given20to60['30-39'])*17.9 + (age_distn_given20to60['40-49']+age_distn_given20to60['50-59'])*16.6) ) }, # '20-39': {'ageBrackets': ['20-29', '30-39'], 'meanDegree': 15.3, 'meanDegree_CI': (12.6, 17.9) }, # '40-59': {'ageBrackets': ['40-49', '50-59'], 'meanDegree': 13.8, 'meanDegree_CI': (11.0, 16.6) }, @@ -379,7 +379,7 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F curidx = 0 for layerGroup, layerInfo in layer_info.items(): print("Generating graph for "+layerGroup+"...") - + layerInfo['numIndividuals'] = numpy.sum([ageBrackets_numInPop[ageBracket] for ageBracket in layerInfo['ageBrackets']]) layerInfo['indices'] = range(curidx, curidx+layerInfo['numIndividuals']) @@ -397,23 +397,23 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F targetMeanDegreeRange = (targetMeanDegree+meanHouseholdSize-0.75, targetMeanDegree+meanHouseholdSize+0.75) if layer_generator=='FARZ' else layerInfo['meanDegree_CI'] # targetMeanDegreeRange = (targetMeanDegree+meanHouseholdSize-1, targetMeanDegree+meanHouseholdSize+1) - while(not graph_generated): + while not graph_generated: try: - if(layer_generator == 'LFR'): + if layer_generator == 'LFR': # print "TARGET MEAN DEGREE = " + str(targetMeanDegree) - + layerInfo['graph'] = networkx.generators.community.LFR_benchmark_graph( - n=layerInfo['numIndividuals'], - tau1=3, tau2=2, mu=0.5, - average_degree=int(targetMeanDegree), + n=layerInfo['numIndividuals'], + tau1=3, tau2=2, mu=0.5, + average_degree=int(targetMeanDegree), tol=1e-01, max_iters=200, seed=(None if graph_gen_attempts<10 else int(numpy.random.rand()*1000))) - elif(layer_generator == 'FARZ'): + elif layer_generator == 'FARZ': # https://github.com/rabbanyk/FARZ - layerInfo['graph'], layerInfo['communities'] = FARZ.generate(farz_params={ - 'n': layerInfo['numIndividuals'], + layerInfo['graph'], layerInfo['communities'] = FARZ.generate(farz_params={ + 'n': layerInfo['numIndividuals'], 'm': int(targetMeanDegree/2), # mean degree / 2 'k': int(layerInfo['numIndividuals']/50), # num communities 'alpha': 2.0, # clustering param @@ -421,24 +421,24 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F 'beta': 0.6, # prob within community edges 'r': 1, # max num communities node can be part of 'q': 0.5, # probability of multi-community membership - 'phi': 1, 'b': 0.0, 'epsilon': 0.0000001, + 'phi': 1, 'b': 0.0, 'epsilon': 0.0000001, 'directed': False, 'weighted': False}) - elif(layer_generator == 'BA'): + elif layer_generator == 'BA': pass else: print("Layer generator \""+layer_generator+"\" is not recognized (support for 'LFR', 'FARZ', 'BA'") - + nodeDegrees = [d[1] for d in layerInfo['graph'].degree()] meanDegree = numpy.mean(nodeDegrees) maxDegree = numpy.max(nodeDegrees) # Enforce that the generated graph has mean degree within the 95% CI of the mean for this group in the data: - if(meanDegree+meanHouseholdSize >= targetMeanDegreeRange[0] and meanDegree+meanHouseholdSize <= targetMeanDegreeRange[1]): - # if(meanDegree+meanHouseholdSize >= targetMeanDegree+meanHouseholdSize-1 and meanDegree+meanHouseholdSize <= targetMeanDegree+meanHouseholdSize+1): - - if(verbose): + if (meanDegree+meanHouseholdSize >= targetMeanDegreeRange[0] and meanDegree+meanHouseholdSize <= targetMeanDegreeRange[1]): + # if (meanDegree+meanHouseholdSize >= targetMeanDegree+meanHouseholdSize-1 and meanDegree+meanHouseholdSize <= targetMeanDegree+meanHouseholdSize+1): + + if verbose: print(layerGroup+" public mean degree = "+str((meanDegree))) print(layerGroup+" public max degree = "+str((maxDegree))) @@ -446,13 +446,13 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F # Create an adjacency matrix mask that will zero out all public edges # for any isolation groups but allow all public edges for other groups: - if(layerGroup in isolation_groups): + if layerGroup in isolation_groups: adjMatrices_isolation_mask.append(numpy.zeros(shape=networkx.adj_matrix(layerInfo['graph']).shape)) else: # adjMatrices_isolation_mask.append(numpy.ones(shape=networkx.adj_matrix(layerInfo['graph']).shape)) # The graph layer we just created represents the baseline (no dist) public connections; # this should be the superset of all connections that exist in any modification of the network, - # therefore it should work to use this baseline adj matrix as the mask instead of a block of 1s + # therefore it should work to use this baseline adj matrix as the mask instead of a block of 1s # (which uses unnecessary memory to store a whole block of 1s, ie not sparse) adjMatrices_isolation_mask.append(networkx.adj_matrix(layerInfo['graph'])) @@ -460,24 +460,24 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F else: graph_gen_attempts += 1 - if(graph_gen_attempts >= 1):# and graph_gen_attempts % 2): - if(meanDegree+meanHouseholdSize < targetMeanDegreeRange[0]): + if graph_gen_attempts >= 1:# and graph_gen_attempts % 2): + if meanDegree+meanHouseholdSize < targetMeanDegreeRange[0]: targetMeanDegree += 1 if layer_generator=='FARZ' else 0.05 - elif(meanDegree+meanHouseholdSize > targetMeanDegreeRange[1]): + elif meanDegree+meanHouseholdSize > targetMeanDegreeRange[1]: targetMeanDegree -= 1 if layer_generator=='FARZ' else 0.05 # reload(networkx) - if(verbose): + if verbose: # print("Try again... (mean degree = "+str(meanDegree)+"+"+str(meanHouseholdSize)+" is outside the target range for mean degree "+str(targetMeanDegreeRange)+")") print("\tTry again... (mean degree = %.2f+%.2f=%.2f is outside the target range for mean degree (%.2f, %.2f)" % (meanDegree, meanHouseholdSize, meanDegree+meanHouseholdSize, targetMeanDegreeRange[0], targetMeanDegreeRange[1])) - + # The networks LFR graph generator function has unreliable convergence. # If it fails to converge in allotted iterations, try again to generate. # If it is stuck (for some reason) and failing many times, reload networkx. except networkx.exception.ExceededMaxIterations: graph_gen_attempts += 1 - # if(graph_gen_attempts >= 10 and graph_gen_attempts % 10): + # if (graph_gen_attempts >= 10 and graph_gen_attempts % 10): # reload(networkx) - if(verbose): + if verbose: print("\tTry again... (networkx failed to converge on a graph)") #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -492,11 +492,11 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F ######################################### # Generate social distancing modifications to the baseline *public* contact network: ######################################### - # In-household connections are assumed to be unaffected by social distancing, + # In-household connections are assumed to be unaffected by social distancing, # and edges will be added to strongly connect households below. - + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Social distancing graphs are generated by randomly drawing (from an exponential distribution) + # Social distancing graphs are generated by randomly drawing (from an exponential distribution) # a number of edges for each node to *keep*, and other edges are removed. #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ G_baseline_NODIST = graphs['baseline'].copy() @@ -504,7 +504,7 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F for dist_scale in distancing_scales: graphs['distancingScale'+str(dist_scale)] = custom_exponential_graph(G_baseline_NODIST, scale=dist_scale) - if(verbose): + if verbose: nodeDegrees_baseline_public_DIST = [d[1] for d in graphs['distancingScale'+str(dist_scale)].degree()] print("Distancing Public Degree Pcts:") (unique, counts) = numpy.unique(nodeDegrees_baseline_public_DIST, return_counts=True) @@ -521,16 +521,16 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F ######################################### # Generate modifications to the contact network representing isolation of individuals in specified groups: ######################################### - if(len(isolation_groups) > 0): - + if len(isolation_groups) > 0: + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Assemble an adjacency matrix mask (from layer generation step) that will zero out - # all public contact edges for the isolation groups but allow all public edges for other groups. + # Assemble an adjacency matrix mask (from layer generation step) that will zero out + # all public contact edges for the isolation groups but allow all public edges for other groups. #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ A_isolation_mask = scipy.sparse.lil_matrix(scipy.sparse.block_diag(adjMatrices_isolation_mask)) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Then multiply each distancing graph by this mask to generate the corresponding + # Then multiply each distancing graph by this mask to generate the corresponding # distancing adjacency matrices where the isolation groups are isolated (no public edges), # and create graphs corresponding to the isolation intervention for each distancing level: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -538,7 +538,7 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F A_withIsolation = scipy.sparse.csr_matrix.multiply( networkx.adj_matrix(graph), A_isolation_mask ) graphs[graphName+'_isolation'] = networkx.from_scipy_sparse_matrix(A_withIsolation) - + ######################################### ######################################### @@ -576,7 +576,7 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F ######################################### # Check the connectivity of the fully constructed contacts graphs for each age group's layer: ######################################### - if(verbose): + if verbose: for graphName, graph in graphs.items(): nodeDegrees = [d[1] for d in graph.degree()] meanDegree= numpy.mean(nodeDegrees) @@ -607,14 +607,14 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F def household_country_data(country): - if(country=='US'): + if country == 'US': household_data = { - 'household_size_distn':{ 1: 0.283708848, - 2: 0.345103011, + 'household_size_distn':{ 1: 0.283708848, + 2: 0.345103011, 3: 0.150677793, - 4: 0.127649150, - 5: 0.057777709, - 6: 0.022624223, + 4: 0.127649150, + 5: 0.057777709, + 6: 0.022624223, 7: 0.012459266 }, 'age_distn':{'0-9': 0.121, @@ -627,7 +627,7 @@ def household_country_data(country): '70-79': 0.070, '80+' : 0.038 }, - 'household_stats':{ 'pct_with_under20': 0.3368, + 'household_stats':{ 'pct_with_under20': 0.3368, 'pct_with_over60': 0.3801, 'pct_with_under20_over60': 0.0341, 'pct_with_over60_givenSingleOccupant': 0.110, @@ -641,27 +641,27 @@ def household_country_data(country): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -# Defines a random exponential edge pruning mechanism +# Defines a random exponential edge pruning mechanism # where the mean degree be easily down-shifted #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ def custom_exponential_graph(base_graph=None, scale=100, min_num_edges=0, m=9, n=None): # If no base graph is provided, generate a random preferential attachment power law graph as a starting point. - if(base_graph): + if base_graph: graph = base_graph.copy() else: assert(n is not None), "Argument n (number of nodes) must be provided when no base graph is given." graph = networkx.barabasi_albert_graph(n=n, m=m) - # We modify the graph by probabilistically dropping some edges from each node. + # We modify the graph by probabilistically dropping some edges from each node. for node in graph: neighbors = list(graph[node].keys()) - if(len(neighbors) > 0): + if len(neighbors) > 0: quarantineEdgeNum = int( max(min(numpy.random.exponential(scale=scale, size=1), len(neighbors)), min_num_edges) ) quarantineKeepNeighbors = numpy.random.choice(neighbors, size=quarantineEdgeNum, replace=False) for neighbor in neighbors: - if(neighbor not in quarantineKeepNeighbors): + if neighbor not in quarantineKeepNeighbors: graph.remove_edge(node, neighbor) - + return graph #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ @@ -669,7 +669,7 @@ def custom_exponential_graph(base_graph=None, scale=100, min_num_edges=0, m=9, n def plot_degree_distn(graph, max_degree=None, show=True, use_seaborn=True): import matplotlib.pyplot as pyplot - if(use_seaborn): + if use_seaborn: import seaborn seaborn.set_style('ticks') seaborn.despine() @@ -688,18 +688,5 @@ def plot_degree_distn(graph, max_degree=None, show=True, use_seaborn=True): pyplot.xlabel('degree') pyplot.ylabel('num nodes') pyplot.legend(loc='upper right') - if(show): + if show: pyplot.show() - - - - - - - - - - - - - diff --git a/seirsplus/sim_loops.py b/seirsplus/sim_loops.py index b855915..b2b45b9 100644 --- a/seirsplus/sim_loops.py +++ b/seirsplus/sim_loops.py @@ -8,12 +8,12 @@ def run_tti_sim(model, T, max_dt=None, intervention_start_pct_infected=0, average_introductions_per_day=0, - testing_cadence='everyday', pct_tested_per_day=1.0, test_falseneg_rate='temporal', + testing_cadence='everyday', pct_tested_per_day=1.0, test_falseneg_rate='temporal', testing_compliance_symptomatic=[None], max_pct_tests_for_symptomatics=1.0, testing_compliance_traced=[None], max_pct_tests_for_traces=1.0, testing_compliance_random=[None], random_testing_degree_bias=0, tracing_compliance=[None], num_contacts_to_trace=None, pct_contacts_to_trace=1.0, tracing_lag=1, - isolation_compliance_symptomatic_individual=[None], isolation_compliance_symptomatic_groupmate=[None], + isolation_compliance_symptomatic_individual=[None], isolation_compliance_symptomatic_groupmate=[None], isolation_compliance_positive_individual=[None], isolation_compliance_positive_groupmate=[None], isolation_compliance_positive_contact=[None], isolation_compliance_positive_contactgroupmate=[None], isolation_lag_symptomatic=1, isolation_lag_positive=1, isolation_lag_contact=0, isolation_groups=None, @@ -26,7 +26,7 @@ def run_tti_sim(model, T, max_dt=None, # (0:Mon, 1:Tue, 2:Wed, 3:Thu, 4:Fri, 5:Sat, 6:Sun, 7:Mon, 8:Tues, ...) # For each cadence, testing is done on the day numbers included in the associated list. - if(cadence_testing_days is None): + if cadence_testing_days is None: cadence_testing_days = { 'everyday': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27], 'workday': [0, 1, 2, 3, 4, 7, 8, 9, 10, 11, 14, 15, 16, 17, 18, 21, 22, 23, 24, 25], @@ -39,8 +39,8 @@ def run_tti_sim(model, T, max_dt=None, #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(temporal_falseneg_rates is None): - temporal_falseneg_rates = { + if temporal_falseneg_rates is None: + temporal_falseneg_rates = { model.E: {0: 1.00, 1: 1.00, 2: 1.00, 3: 1.00}, model.I_pre: {0: 0.25, 1: 0.25, 2: 0.22}, model.I_sym: {0: 0.19, 1: 0.16, 2: 0.16, 3: 0.17, 4: 0.19, 5: 0.22, 6: 0.26, 7: 0.29, 8: 0.34, 9: 0.38, 10: 0.43, 11: 0.48, 12: 0.52, 13: 0.57, 14: 0.62, 15: 0.66, 16: 0.70, 17: 0.76, 18: 0.79, 19: 0.82, 20: 0.85, 21: 0.88, 22: 0.90, 23: 0.92, 24: 0.93, 25: 0.95, 26: 0.96, 27: 0.97, 28: 0.97, 29: 0.98, 30: 0.98, 31: 0.99}, @@ -82,26 +82,26 @@ def run_tti_sim(model, T, max_dt=None, #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Introduce exogenous exposures randomly: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(int(model.t)!=int(timeOfLastIntroduction)): + if int(model.t) != int(timeOfLastIntroduction): timeOfLastIntroduction = model.t numNewExposures = numpy.random.poisson(lam=average_introductions_per_day) - + model.introduce_exposures(num_new_exposures=numNewExposures) - if(numNewExposures > 0): + if numNewExposures > 0: print("[NEW EXPOSURE @ t = %.2f (%d exposed)]" % (model.t, numNewExposures)) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Execute testing policy at designated intervals: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(int(model.t)!=int(timeOfLastIntervention)): - + if int(model.t) != int(timeOfLastIntervention): + cadenceDayNumbers = [int(model.t % cadence_cycle_length)] - if(backlog_skipped_intervals): + if backlog_skipped_intervals: cadenceDayNumbers = [int(i % cadence_cycle_length) for i in numpy.arange(start=timeOfLastIntervention, stop=int(model.t), step=1.0)[1:]] + cadenceDayNumbers timeOfLastIntervention = model.t @@ -111,14 +111,14 @@ def run_tti_sim(model, T, max_dt=None, currentNumInfected = model.total_num_infected()[model.tidx] currentPctInfected = model.total_num_infected()[model.tidx]/model.numNodes - if(currentPctInfected >= intervention_start_pct_infected and not interventionOn): + if (currentPctInfected >= intervention_start_pct_infected and not interventionOn): interventionOn = True interventionStartTime = model.t - - if(interventionOn): + + if interventionOn: print("[INTERVENTIONS @ t = %.2f (%d (%.2f%%) infected)]" % (model.t, currentNumInfected, currentPctInfected*100)) - + nodeStates = model.X.flatten() nodeTestedStatuses = model.tested.flatten() nodeTestedInCurrentStateStatuses = model.testedInCurrentState.flatten() @@ -129,7 +129,7 @@ def run_tti_sim(model, T, max_dt=None, # tracingPoolQueue[0] = tracingPoolQueue[0]Queue.pop(0) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - + newIsolationGroup_symptomatic = [] newIsolationGroup_contact = [] @@ -139,22 +139,22 @@ def run_tti_sim(model, T, max_dt=None, numSelfIsolated_symptoms = 0 numSelfIsolated_symptomaticGroupmate = 0 - if(any(isolation_compliance_symptomatic_individual)): + if any(isolation_compliance_symptomatic_individual): symptomaticNodes = numpy.argwhere((nodeStates==model.I_sym)).flatten() for symptomaticNode in symptomaticNodes: - if(isolation_compliance_symptomatic_individual[symptomaticNode]): - if(model.X[symptomaticNode] == model.I_sym): - numSelfIsolated_symptoms += 1 + if isolation_compliance_symptomatic_individual[symptomaticNode]: + if model.X[symptomaticNode] == model.I_sym: + numSelfIsolated_symptoms += 1 newIsolationGroup_symptomatic.append(symptomaticNode) #---------------------------------------- # Isolate the GROUPMATES of this SYMPTOMATIC node without a test: #---------------------------------------- - if(isolation_groups is not None and any(isolation_compliance_symptomatic_groupmate)): + if (isolation_groups is not None and any(isolation_compliance_symptomatic_groupmate)): isolationGroupmates = next((group for group in isolation_groups if symptomaticNode in group), None) for isolationGroupmate in isolationGroupmates: - if(isolationGroupmate != symptomaticNode): - if(isolation_compliance_symptomatic_groupmate[isolationGroupmate]): + if isolationGroupmate != symptomaticNode: + if isolation_compliance_symptomatic_groupmate[isolationGroupmate]: numSelfIsolated_symptomaticGroupmate += 1 newIsolationGroup_symptomatic.append(isolationGroupmate) @@ -165,23 +165,23 @@ def run_tti_sim(model, T, max_dt=None, numSelfIsolated_positiveContact = 0 numSelfIsolated_positiveContactGroupmate = 0 - if(any(isolation_compliance_positive_contact) or any(isolation_compliance_positive_contactgroupmate)): + if (any(isolation_compliance_positive_contact) or any(isolation_compliance_positive_contactgroupmate)): for contactNode in tracingPoolQueue[0]: - if(isolation_compliance_positive_contact[contactNode]): + if isolation_compliance_positive_contact[contactNode]: newIsolationGroup_contact.append(contactNode) - numSelfIsolated_positiveContact += 1 + numSelfIsolated_positiveContact += 1 #---------------------------------------- # Isolate the GROUPMATES of this self-isolating CONTACT without a test: #---------------------------------------- - if(isolation_groups is not None and any(isolation_compliance_positive_contactgroupmate)): + if (isolation_groups is not None and any(isolation_compliance_positive_contactgroupmate)): isolationGroupmates = next((group for group in isolation_groups if contactNode in group), None) for isolationGroupmate in isolationGroupmates: - # if(isolationGroupmate != contactNode): - if(isolation_compliance_positive_contactgroupmate[isolationGroupmate]): + # if (isolationGroupmate != contactNode): + if isolation_compliance_positive_contactgroupmate[isolationGroupmate]: newIsolationGroup_contact.append(isolationGroupmate) numSelfIsolated_positiveContactGroupmate += 1 - + #---------------------------------------- # Update the nodeStates list after self-isolation updates to model.X: @@ -198,8 +198,8 @@ def run_tti_sim(model, T, max_dt=None, #---------------------------------------- symptomaticSelection = [] - if(any(testing_compliance_symptomatic)): - + if any(testing_compliance_symptomatic): + symptomaticPool = numpy.argwhere((testing_compliance_symptomatic==True) & (nodeTestedInCurrentStateStatuses==False) & (nodePositiveStatuses==False) @@ -207,8 +207,8 @@ def run_tti_sim(model, T, max_dt=None, ).flatten() numSymptomaticTests = min(len(symptomaticPool), max_symptomatic_tests_per_day) - - if(len(symptomaticPool) > 0): + + if len(symptomaticPool) > 0: symptomaticSelection = symptomaticPool[numpy.random.choice(len(symptomaticPool), min(numSymptomaticTests, len(symptomaticPool)), replace=False)] @@ -220,25 +220,25 @@ def run_tti_sim(model, T, max_dt=None, tracingSelection = [] randomSelection = [] - if(cadenceDayNumber in testingDays): + if cadenceDayNumber in testingDays: #---------------------------------------- - # Apply a designated portion of this day's tests + # Apply a designated portion of this day's tests # to individuals identified by CONTACT TRACING: #---------------------------------------- tracingPool = tracingPoolQueue.pop(0) - if(any(testing_compliance_traced)): + if any(testing_compliance_traced): numTracingTests = min(len(tracingPool), min(tests_per_day-len(symptomaticSelection), max_tracing_tests_per_day)) for trace in range(numTracingTests): traceNode = tracingPool.pop() - if((nodePositiveStatuses[traceNode]==False) + if ((nodePositiveStatuses[traceNode]==False) and (testing_compliance_traced[traceNode]==True) and (model.X[traceNode] != model.R) - and (model.X[traceNode] != model.Q_R) + and (model.X[traceNode] != model.Q_R) and (model.X[traceNode] != model.H) and (model.X[traceNode] != model.F)): tracingSelection.append(traceNode) @@ -247,25 +247,25 @@ def run_tti_sim(model, T, max_dt=None, # Apply the remainder of this day's tests to random testing: #---------------------------------------- - if(any(testing_compliance_random)): - + if any(testing_compliance_random): + testingPool = numpy.argwhere((testing_compliance_random==True) & (nodePositiveStatuses==False) & (nodeStates != model.R) - & (nodeStates != model.Q_R) + & (nodeStates != model.Q_R) & (nodeStates != model.H) & (nodeStates != model.F) ).flatten() numRandomTests = max(min(tests_per_day-len(tracingSelection)-len(symptomaticSelection), len(testingPool)), 0) - + testingPool_degrees = model.degree.flatten()[testingPool] testingPool_degreeWeights = numpy.power(testingPool_degrees,random_testing_degree_bias)/numpy.sum(numpy.power(testingPool_degrees,random_testing_degree_bias)) - if(len(testingPool) > 0): + if len(testingPool) > 0: randomSelection = testingPool[numpy.random.choice(len(testingPool), numRandomTests, p=testingPool_degreeWeights, replace=False)] - + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -282,9 +282,9 @@ def run_tti_sim(model, T, max_dt=None, numPositive = 0 numPositive_random = 0 numPositive_tracing = 0 - numPositive_symptomatic = 0 + numPositive_symptomatic = 0 numIsolated_positiveGroupmate = 0 - + newTracingPool = [] newIsolationGroup_positive = [] @@ -294,80 +294,80 @@ def run_tti_sim(model, T, max_dt=None, model.set_tested(testNode, True) numTested += 1 - if(i < len(symptomaticSelection)): + if i < len(symptomaticSelection): numTested_symptomatic += 1 - elif(i < len(symptomaticSelection)+len(tracingSelection)): + elif i < len(symptomaticSelection)+len(tracingSelection): numTested_tracing += 1 else: - numTested_random += 1 + numTested_random += 1 - # If the node to be tested is not infected, then the test is guaranteed negative, + # If the node to be tested is not infected, then the test is guaranteed negative, # so don't bother going through with doing the test: - if(model.X[testNode] == model.S or model.X[testNode] == model.Q_S): + if (model.X[testNode] == model.S or model.X[testNode] == model.Q_S): pass # Also assume that latent infections are not picked up by tests: - elif(model.X[testNode] == model.E or model.X[testNode] == model.Q_E): + elif (model.X[testNode] == model.E or model.X[testNode] == model.Q_E): pass - elif(model.X[testNode] == model.I_pre or model.X[testNode] == model.Q_pre - or model.X[testNode] == model.I_sym or model.X[testNode] == model.Q_sym + elif (model.X[testNode] == model.I_pre or model.X[testNode] == model.Q_pre + or model.X[testNode] == model.I_sym or model.X[testNode] == model.Q_sym or model.X[testNode] == model.I_asym or model.X[testNode] == model.Q_asym): - - if(test_falseneg_rate == 'temporal'): + + if test_falseneg_rate == 'temporal': testNodeState = model.X[testNode][0] testNodeTimeInState = model.timer_state[testNode][0] - if(testNodeState in list(temporal_falseneg_rates.keys())): + if testNodeState in list(temporal_falseneg_rates.keys()): falseneg_prob = temporal_falseneg_rates[testNodeState][ int(min(testNodeTimeInState, max(list(temporal_falseneg_rates[testNodeState].keys())))) ] else: falseneg_prob = 1.00 else: falseneg_prob = test_falseneg_rate - if(numpy.random.rand() < (1-falseneg_prob)): + if numpy.random.rand() < (1-falseneg_prob): # +++++++++++++++++++++++++++++++++++++++++++++ # The tested node has returned a positive test # +++++++++++++++++++++++++++++++++++++++++++++ numPositive += 1 - if(i < len(symptomaticSelection)): + if i < len(symptomaticSelection): numPositive_symptomatic += 1 - elif(i < len(symptomaticSelection)+len(tracingSelection)): + elif i < len(symptomaticSelection)+len(tracingSelection): numPositive_tracing += 1 else: - numPositive_random += 1 - + numPositive_random += 1 + # Update the node's state to the appropriate detected case state: model.set_positive(testNode, True) #---------------------------------------- # Add this positive node to the isolation group: #---------------------------------------- - if(isolation_compliance_positive_individual[testNode]): + if isolation_compliance_positive_individual[testNode]: newIsolationGroup_positive.append(testNode) #---------------------------------------- # Add the groupmates of this positive node to the isolation group: - #---------------------------------------- - if(isolation_groups is not None and any(isolation_compliance_positive_groupmate)): + #---------------------------------------- + if (isolation_groups is not None and any(isolation_compliance_positive_groupmate)): isolationGroupmates = next((group for group in isolation_groups if testNode in group), None) for isolationGroupmate in isolationGroupmates: - if(isolationGroupmate != testNode): - if(isolation_compliance_positive_groupmate[isolationGroupmate]): + if isolationGroupmate != testNode: + if isolation_compliance_positive_groupmate[isolationGroupmate]: numIsolated_positiveGroupmate += 1 newIsolationGroup_positive.append(isolationGroupmate) - #---------------------------------------- + #---------------------------------------- # Add this node's neighbors to the contact tracing pool: - #---------------------------------------- - if(any(tracing_compliance) or any(isolation_compliance_positive_contact) or any(isolation_compliance_positive_contactgroupmate)): - if(tracing_compliance[testNode]): + #---------------------------------------- + if (any(tracing_compliance) or any(isolation_compliance_positive_contact) or any(isolation_compliance_positive_contactgroupmate)): + if tracing_compliance[testNode]: testNodeContacts = list(model.G[testNode].keys()) numpy.random.shuffle(testNodeContacts) - if(num_contacts_to_trace is None): + if num_contacts_to_trace is None: numContactsToTrace = int(pct_contacts_to_trace*len(testNodeContacts)) else: numContactsToTrace = num_contacts_to_trace newTracingPool.extend(testNodeContacts[0:numContactsToTrace]) - + # Add the nodes to be isolated to the isolation queue: isolationQueue_positive.append(newIsolationGroup_positive) isolationQueue_symptomatic.append(newIsolationGroup_symptomatic) @@ -378,9 +378,9 @@ def run_tti_sim(model, T, max_dt=None, print("\t"+str(numTested_symptomatic) +"\ttested due to symptoms [+ "+str(numPositive_symptomatic)+" positive (%.2f %%) +]" % (numPositive_symptomatic/numTested_symptomatic*100 if numTested_symptomatic>0 else 0)) - print("\t"+str(numTested_tracing) +"\ttested as traces [+ "+str(numPositive_tracing)+" positive (%.2f %%) +]" % (numPositive_tracing/numTested_tracing*100 if numTested_tracing>0 else 0)) - print("\t"+str(numTested_random) +"\ttested randomly [+ "+str(numPositive_random)+" positive (%.2f %%) +]" % (numPositive_random/numTested_random*100 if numTested_random>0 else 0)) - print("\t"+str(numTested) +"\ttested TOTAL [+ "+str(numPositive)+" positive (%.2f %%) +]" % (numPositive/numTested*100 if numTested>0 else 0)) + print("\t"+str(numTested_tracing) +"\ttested as traces [+ "+str(numPositive_tracing)+" positive (%.2f %%) +]" % (numPositive_tracing/numTested_tracing*100 if numTested_tracing>0 else 0)) + print("\t"+str(numTested_random) +"\ttested randomly [+ "+str(numPositive_random)+" positive (%.2f %%) +]" % (numPositive_random/numTested_random*100 if numTested_random>0 else 0)) + print("\t"+str(numTested) +"\ttested TOTAL [+ "+str(numPositive)+" positive (%.2f %%) +]" % (numPositive/numTested*100 if numTested>0 else 0)) print("\t"+str(numSelfIsolated_symptoms) +" will isolate due to symptoms ("+str(numSelfIsolated_symptomaticGroupmate)+" as groupmates of symptomatic)") print("\t"+str(numPositive) +" will isolate due to positive test ("+str(numIsolated_positiveGroupmate)+" as groupmates of positive)") @@ -408,7 +408,7 @@ def run_tti_sim(model, T, max_dt=None, numIsolated += 1 print("\t"+str(numIsolated)+" entered isolation") - + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/seirsplus/utilities.py b/seirsplus/utilities.py index b8e9720..398b586 100644 --- a/seirsplus/utilities.py +++ b/seirsplus/utilities.py @@ -9,17 +9,17 @@ def gamma_dist(mean, coeffvar, N): def dist_info(dists, names=None, plot=False, bin_size=1, colors=None, reverse_plot=False): dists = [dists] if not isinstance(dists, list) else dists - names = [names] if(names is not None and not isinstance(names, list)) else (names if names is not None else [None]*len(dists)) - colors = [colors] if(colors is not None and not isinstance(colors, list)) else (colors if colors is not None else pyplot.rcParams['axes.prop_cycle'].by_key()['color']) - + names = [names] if (names is not None and not isinstance(names, list)) else (names if names is not None else [None]*len(dists)) + colors = [colors] if (colors is not None and not isinstance(colors, list)) else (colors if colors is not None else pyplot.rcParams['axes.prop_cycle'].by_key()['color']) + for i, (dist, name) in enumerate(zip(dists, names)): print((name+": " if name else "")+" mean = %.2f, std = %.2f, 95%% CI = (%.2f, %.2f)" % (numpy.mean(dist), numpy.std(dist), numpy.percentile(dist, 2.5), numpy.percentile(dist, 97.5))) print() - - if(plot): + + if plot: pyplot.hist(dist, bins=numpy.arange(0, int(max(dist)+1), step=bin_size), label=(name if name else False), color=colors[i], edgecolor='white', alpha=0.6, zorder=(-1*i if reverse_plot else i)) - - if(plot): + + if plot: pyplot.ylabel('num nodes') pyplot.legend(loc='upper right') pyplot.show() @@ -29,27 +29,27 @@ def network_info(networks, names=None, plot=False, bin_size=1, colors=None, reve import networkx networks = [networks] if not isinstance(networks, list) else networks names = [names] if not isinstance(names, list) else names - colors = [colors] if(colors is not None and not isinstance(colors, list)) else (colors if colors is not None else pyplot.rcParams['axes.prop_cycle'].by_key()['color']) - + colors = [colors] if (colors is not None and not isinstance(colors, list)) else (colors if colors is not None else pyplot.rcParams['axes.prop_cycle'].by_key()['color']) + for i, (network, name) in enumerate(zip(networks, names)): - + degree = [d[1] for d in network.degree()] - if(name): + if name: print(name+":") - print("Degree: mean = %.2f, std = %.2f, 95%% CI = (%.2f, %.2f)\n coeff var = %.2f" - % (numpy.mean(degree), numpy.std(degree), numpy.percentile(degree, 2.5), numpy.percentile(degree, 97.5), + print("Degree: mean = %.2f, std = %.2f, 95%% CI = (%.2f, %.2f)\n coeff var = %.2f" + % (numpy.mean(degree), numpy.std(degree), numpy.percentile(degree, 2.5), numpy.percentile(degree, 97.5), numpy.std(degree)/numpy.mean(degree))) r = networkx.degree_assortativity_coefficient(network) print("Assortativity: %.2f" % (r)) c = networkx.average_clustering(network) print("Clustering coeff: %.2f" % (c)) print() - - if(plot): + + if plot: pyplot.hist(degree, bins=numpy.arange(0, int(max(degree)+1), step=bin_size), label=(name+" degree" if name else False), color=colors[i], edgecolor='white', alpha=0.6, zorder=(-1*i if reverse_plot else i)) - - if(plot): + + if plot: pyplot.ylabel('num nodes') pyplot.legend(loc='upper right') pyplot.show() @@ -59,17 +59,3 @@ def results_summary(model): print("total percent infected: %0.2f%%" % ((model.total_num_infected()[-1]+model.total_num_recovered()[-1])/model.numNodes * 100) ) print("total percent fatality: %0.2f%%" % (model.numF[-1]/model.numNodes * 100) ) print("peak pct hospitalized: %0.2f%%" % (numpy.max(model.numH)/model.numNodes * 100) ) - - - - - - - - - - - - - -