From fe78dbfe85467d21035b1cd7724c53fd4cb9d53e Mon Sep 17 00:00:00 2001 From: Noel Schutt Date: Sat, 13 Mar 2021 10:10:15 -0500 Subject: [PATCH 01/12] Cleaned up trailing whitespace. Ran `s/\s\+$//e`. --- seirsplus/FARZ.py | 178 +++++------ seirsplus/legacy_models.py | 210 ++++++------- seirsplus/models.py | 588 ++++++++++++++++++------------------- seirsplus/networks.py | 125 ++++---- seirsplus/sim_loops.py | 72 ++--- seirsplus/utilities.py | 32 +- 6 files changed, 586 insertions(+), 619 deletions(-) diff --git a/seirsplus/FARZ.py b/seirsplus/FARZ.py index f81acc4..309f657 100644 --- a/seirsplus/FARZ.py +++ b/seirsplus/FARZ.py @@ -5,7 +5,7 @@ import random import bisect import math -import os +import os def random_choice(values, weights=None , size = 1, replace = True): if weights is None: @@ -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,9 +398,9 @@ 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"): @@ -410,12 +410,12 @@ def main(argv): 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())): batch_setting['vari'] = arg @@ -428,9 +428,9 @@ 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) + sys.exit(2) elif opt in ("-n"): - try: FARZsetting['n'] = int(arg) + try: FARZsetting['n'] = int(arg) except ValueError: print('Invalid Number , try -h to see the usage and options') sys.exit(2) @@ -445,30 +445,30 @@ def main(argv): 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..333af64 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,7 +82,7 @@ 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: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -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 # varibles is a list with compartment counts as elements N = S + E + I + D_E + D_I + R @@ -139,7 +139,7 @@ def run_epoch(self, runtime, dt=0.1): 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 ) @@ -177,7 +177,7 @@ def run(self, T, dt=0.1, checkpoints=None, verbose=False): self.tmax += T else: return False - + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Pre-process checkpoint values: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -187,10 +187,10 @@ def run(self, T, dt=0.1, checkpoints=None, verbose=False): '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)) + or not isinstance(checkpoints[param], (list, numpy.ndarray)) or len(checkpoints[param])!=numCheckpoints): checkpoints[param] = [getattr(self, param)]*numCheckpoints @@ -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']): @@ -244,9 +244,9 @@ 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[:]) + 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): @@ -279,10 +279,10 @@ 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): dashedReference_tseries = dashed_reference_results.tseries[::int(self.N/100)] @@ -331,7 +331,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin 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: @@ -364,7 +364,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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'): @@ -413,7 +413,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin 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): @@ -443,8 +443,8 @@ 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) @@ -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): @@ -480,9 +480,9 @@ 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): @@ -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, @@ -554,15 +554,15 @@ def __init__(self, G, beta, sigma, gamma, xi=0, mu_I=0, mu_0=0, nu=0, beta_local #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -613,7 +613,7 @@ def __init__(self, G, beta, sigma, gamma, xi=0, mu_I=0, mu_0=0, nu=0, beta_local 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}, @@ -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 @@ -717,7 +717,7 @@ def update_parameters(self): 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): self.A_beta = scipy.sparse.csr_matrix.multiply(self.A, numpy.tile(self.beta_local, (1,self.numNodes))).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)) @@ -791,35 +791,35 @@ 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[:]) + 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) @@ -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] @@ -920,7 +920,7 @@ def finalize_data_series(self): return None #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def run_iteration(self): @@ -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) @@ -1027,7 +1027,7 @@ def run(self, T, checkpoints=None, print_interval=10, verbose='t'): 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 + checkpoints = None else: checkpointTime = checkpoints['t'][checkpointIdx] @@ -1060,7 +1060,7 @@ def run(self, T, checkpoints=None, print_interval=10, verbose='t'): 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 + checkpoints = None else: checkpointTime = checkpoints['t'][checkpointIdx] #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -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): @@ -1117,10 +1117,10 @@ 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): dashedReference_tseries = dashed_reference_results.tseries[::int(self.numNodes/100)] @@ -1169,7 +1169,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin 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: @@ -1202,7 +1202,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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'): @@ -1251,7 +1251,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin 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): @@ -1281,8 +1281,8 @@ 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) @@ -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): @@ -1318,9 +1318,9 @@ 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): @@ -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 index 2b03171..339feee 100644 --- a/seirsplus/models.py +++ b/seirsplus/models.py @@ -20,14 +20,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_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 @@ -37,32 +37,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 initQ_E Init number of detected infectious individuals - initQ_I Init number of detected infectious individuals - initR Init number of recovered 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) + (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, + 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 + 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 @@ -81,7 +81,7 @@ 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: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -103,7 +103,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_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 + 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 @@ -138,7 +138,7 @@ def run_epoch(self, runtime, dt=0.1): 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.numQ_E[-1], self.numQ_I[-1], self.numR[-1], self.numF[-1]] @@ -148,7 +148,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_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 ) @@ -176,7 +176,7 @@ def run(self, T, dt=0.1, checkpoints=None, verbose=False): self.tmax += T else: return False - + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Pre-process checkpoint values: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -186,10 +186,10 @@ def run(self, T, dt=0.1, checkpoints=None, verbose=False): '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), + # 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 not isinstance(checkpoints[param], (list, numpy.ndarray)) or len(checkpoints[param])!=numCheckpoints): checkpoints[param] = [getattr(self, param)]*numCheckpoints @@ -210,7 +210,7 @@ def run(self, T, dt=0.1, checkpoints=None, verbose=False): 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']): @@ -243,7 +243,7 @@ def run(self, T, dt=0.1, checkpoints=None, verbose=False): def total_num_susceptible(self, t_idx=None): if(t_idx is None): - return (self.numS[:]) + return (self.numS[:]) else: return (self.numS[t_idx]) @@ -251,9 +251,9 @@ def total_num_susceptible(self, t_idx=None): def total_num_infected(self, t_idx=None): if(t_idx is None): - return (self.numE[:] + self.numI[:] + self.numQ_E[:] + self.numQ_I[:]) + 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]) + return (self.numE[t_idx] + self.numI[t_idx] + self.numQ_E[t_idx] + self.numQ_I[t_idx]) #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ @@ -267,9 +267,9 @@ def total_num_isolated(self, t_idx=None): def total_num_recovered(self, t_idx=None): if(t_idx is None): - return (self.numR[:]) + return (self.numR[:]) else: - return (self.numR[t_idx]) + return (self.numR[t_idx]) #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ @@ -279,8 +279,8 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin 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', + 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): @@ -302,10 +302,10 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin 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 + Sseries = self.numS/self.N if plot_percentages else self.numS #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Draw the reference data: + # Draw the reference data: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(dashed_reference_results): dashedReference_tseries = dashed_reference_results.tseries[::int(self.N/100)] @@ -354,7 +354,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin 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: @@ -387,7 +387,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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'): @@ -436,7 +436,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin 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 @@ -447,8 +447,8 @@ def figure_basic(self, plot_S='line', plot_E='line', plot_I='line',plot_R='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', + 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): @@ -466,8 +466,8 @@ def figure_basic(self, plot_S='line', plot_E='line', plot_I='line',plot_R='line' 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, + 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) @@ -484,8 +484,8 @@ def figure_infections(self, plot_S=False, plot_E='stacked', plot_I='stacked',plo 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', + 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): @@ -503,9 +503,9 @@ def figure_infections(self, plot_S=False, plot_E='stacked', plot_I='stacked',plo 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, + 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): @@ -525,50 +525,50 @@ class SEIRSNetworkModel(): """ A class to simulate the SEIRS Stochastic Network Model ====================================================== - Params: + 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) + 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 - + 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 + 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 - + 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 + 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) + (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, + 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, + 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): @@ -579,21 +579,21 @@ def __init__(self, G, beta, sigma, gamma, # 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':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, + '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, + '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 + # 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) @@ -617,7 +617,7 @@ def __init__(self, G, beta, sigma, gamma, 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: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -640,8 +640,8 @@ def __init__(self, G, beta, sigma, gamma, 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.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)) @@ -652,7 +652,7 @@ def __init__(self, G, beta, sigma, gamma, self.Xseries = numpy.zeros(shape=(6*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}, @@ -673,8 +673,8 @@ def __init__(self, G, beta, sigma, gamma, #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ 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.numTested = numpy.zeros(6*self.numNodes) + self.numPositive = numpy.zeros(6*self.numNodes) self.testedInCurrentState = numpy.array([False]*self.numNodes).reshape((self.numNodes,1)) @@ -708,7 +708,7 @@ def __init__(self, G, beta, sigma, gamma, 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] - + #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ @@ -757,9 +757,9 @@ def update_parameters(self): 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: #---------------------------------------- @@ -787,18 +787,18 @@ def update_parameters(self): 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 + 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)) + 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)) + 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: #---------------------------------------- @@ -812,7 +812,7 @@ def update_parameters(self): 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() + 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'): @@ -871,7 +871,7 @@ def update_parameters(self): 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() + 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'): @@ -922,7 +922,7 @@ def update_parameters(self): #---------------------------------------- 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) - + #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ @@ -936,7 +936,7 @@ def node_degrees(self, Amat): def total_num_susceptible(self, t_idx=None): if(t_idx is None): - return (self.numS[:]) + return (self.numS[:]) else: return (self.numS[t_idx]) @@ -944,7 +944,7 @@ def total_num_susceptible(self, t_idx=None): def total_num_infected(self, t_idx=None): if(t_idx is None): - return (self.numE[:] + self.numI[:] + self.numQ_E[:] + self.numQ_I[:]) + 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]) @@ -952,15 +952,15 @@ def total_num_infected(self, t_idx=None): def total_num_isolated(self, t_idx=None): if(t_idx is None): - return (self.numQ_E[:] + self.numQ_I[:]) + return (self.numQ_E[:] + self.numQ_I[:]) else: - return (self.numQ_E[t_idx] + self.numQ_I[t_idx]) + 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[:]) + return (self.numTested[:]) else: return (self.numTested[t_idx]) @@ -968,7 +968,7 @@ def total_num_tested(self, t_idx=None): def total_num_positive(self, t_idx=None): if(t_idx is None): - return (self.numPositive[:]) + return (self.numPositive[:]) else: return (self.numPositive[t_idx]) @@ -976,14 +976,14 @@ def total_num_positive(self, t_idx=None): def total_num_recovered(self, t_idx=None): if(t_idx is None): - return (self.numR[:]) + return (self.numR[:]) else: return (self.numR[t_idx]) #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - + def calc_propensities(self): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -991,7 +991,7 @@ def calc_propensities(self): # and check to see if their computation is necessary before doing the multiplication #------------------------------------ - self.transmissionTerms_I = numpy.zeros(shape=(self.numNodes,1)) + 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)) @@ -1037,7 +1037,7 @@ def calc_propensities(self): 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)) + propensities__toS = 1e5 * ((self.X!=self.F) & numpy.greater(self.timer_state, 1/self.nu)) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -1063,20 +1063,20 @@ def calc_propensities(self): 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 = numpy.hstack([propensities_StoE, propensities_EtoI, + propensities_ItoR, propensities_ItoF, + propensities_EtoQE, propensities_ItoQI, propensities_QEtoQI, propensities_QItoR, propensities_QItoF, propensities_RtoS, propensities__toS]) @@ -1128,7 +1128,7 @@ def introduce_exposures(self, num_new_exposures): #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def increase_data_series_length(self): self.tseries = numpy.pad(self.tseries, [(0, 6*self.numNodes)], mode='constant', constant_values=0) @@ -1161,7 +1161,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] @@ -1195,7 +1195,7 @@ def finalize_data_series(self): return None #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def run_iteration(self): @@ -1278,7 +1278,7 @@ def run_iteration(self): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ 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) @@ -1288,7 +1288,7 @@ def run_iteration(self): 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) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -1352,7 +1352,7 @@ def run(self, T, checkpoints=None, print_interval=10, verbose='t'): 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 + checkpoints = None else: checkpointTime = checkpoints['t'][checkpointIdx] @@ -1381,7 +1381,7 @@ def run(self, T, checkpoints=None, print_interval=10, verbose='t'): 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 + checkpoints = None else: checkpointTime = checkpoints['t'][checkpointIdx] #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -1415,8 +1415,8 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin 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', + 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): @@ -1438,10 +1438,10 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin 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 + Sseries = self.numS/self.numNodes if plot_percentages else self.numS #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Draw the reference data: + # Draw the reference data: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(dashed_reference_results): dashedReference_tseries = dashed_reference_results.tseries[::int(self.numNodes/100)] @@ -1490,7 +1490,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin 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: @@ -1523,7 +1523,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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'): @@ -1572,7 +1572,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin 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 @@ -1583,8 +1583,8 @@ def figure_basic(self, plot_S='line', plot_E='line', plot_I='line',plot_R='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', + 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): @@ -1602,8 +1602,8 @@ def figure_basic(self, plot_S='line', plot_E='line', plot_I='line',plot_R='line' 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, + 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) @@ -1620,8 +1620,8 @@ def figure_infections(self, plot_S=False, plot_E='stacked', plot_I='stacked',plo 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', + 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): @@ -1639,9 +1639,9 @@ def figure_infections(self, plot_S=False, plot_E='stacked', plot_I='stacked',plo 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, + 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): @@ -1671,56 +1671,56 @@ class ExtSEIRSNetworkModel(): """ A class to simulate the Extended SEIRS Stochastic Network Model =================================================== - Params: + 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) + 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 - + 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 + 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. - + 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 + 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 @@ -1728,15 +1728,15 @@ class ExtSEIRSNetworkModel(): 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) + (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, + 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, + 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, + 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): @@ -1749,26 +1749,26 @@ def __init__(self, G, beta, sigma, lamda, gamma, # 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':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, + '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, + '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, + '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 + # 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) @@ -1799,7 +1799,7 @@ def __init__(self, G, beta, sigma, lamda, gamma, 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: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -1816,7 +1816,7 @@ def __init__(self, G, beta, sigma, lamda, gamma, 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.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] @@ -1837,12 +1837,12 @@ def __init__(self, G, beta, sigma, lamda, gamma, 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.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_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) @@ -1852,7 +1852,7 @@ def __init__(self, G, beta, sigma, lamda, gamma, self.Xseries = numpy.zeros(shape=(6*self.numNodes, self.numNodes), dtype='uint8') self.Xseries[0,:] = self.X.T - self.transitions = { + 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}, @@ -1885,8 +1885,8 @@ def __init__(self, G, beta, sigma, lamda, gamma, #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ 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.numTested = numpy.zeros(6*self.numNodes) + self.numPositive = numpy.zeros(6*self.numNodes) self.testedInCurrentState = numpy.array([False]*self.numNodes).reshape((self.numNodes,1)) @@ -1934,7 +1934,7 @@ def __init__(self, G, beta, sigma, lamda, gamma, 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] - + #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ @@ -2000,7 +2000,7 @@ def update_parameters(self): # 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: #---------------------------------------- @@ -2040,22 +2040,22 @@ def update_parameters(self): 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_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_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_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: #---------------------------------------- @@ -2069,7 +2069,7 @@ def update_parameters(self): 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() + 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'): @@ -2153,7 +2153,7 @@ def update_parameters(self): 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() + 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'): @@ -2208,7 +2208,7 @@ def update_parameters(self): self.A_deltabeta_asym = scipy.sparse.csr_matrix.multiply(self.A_delta_pairwise, self.A_beta_asym_pairwise) else: self.A_deltabeta_asym = None - + #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ @@ -2222,7 +2222,7 @@ def node_degrees(self, Amat): def total_num_susceptible(self, t_idx=None): if(t_idx is None): - return (self.numS[:] + self.numQ_S[:]) + return (self.numS[:] + self.numQ_S[:]) else: return (self.numS[t_idx] + self.numQ_S[t_idx]) @@ -2231,7 +2231,7 @@ def total_num_susceptible(self, t_idx=None): 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[:]) + + 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]) @@ -2240,15 +2240,15 @@ def total_num_infected(self, t_idx=None): 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[:]) + 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]) + 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[:]) + return (self.numTested[:]) else: return (self.numTested[t_idx]) @@ -2256,7 +2256,7 @@ def total_num_tested(self, t_idx=None): def total_num_positive(self, t_idx=None): if(t_idx is None): - return (self.numPositive[:]) + return (self.numPositive[:]) else: return (self.numPositive[t_idx]) @@ -2264,14 +2264,14 @@ def total_num_positive(self, t_idx=None): def total_num_recovered(self, t_idx=None): if(t_idx is None): - return (self.numR[:] + self.numQ_R[:]) + return (self.numR[:] + self.numQ_R[:]) else: return (self.numR[t_idx] + self.numQ_R[t_idx]) #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - + def calc_propensities(self): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -2279,7 +2279,7 @@ def calc_propensities(self): # and check to see if their computation is necessary before doing the multiplication #------------------------------------ - self.transmissionTerms_I = numpy.zeros(shape=(self.numNodes,1)) + 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)) @@ -2296,7 +2296,7 @@ def calc_propensities(self): #------------------------------------ - self.transmissionTerms_IQ = numpy.zeros(shape=(self.numNodes,1)) + 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)))) @@ -2319,7 +2319,7 @@ def calc_propensities(self): propensities_QStoQE = numpy.zeros_like(propensities_StoE) if(numpy.any(self.X==self.Q_S)): - propensities_QStoQE = ( self.alpha_Q * + 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]) @@ -2422,15 +2422,15 @@ def calc_propensities(self): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ 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_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', + 'ISYMtoR', 'ISYMtoH', 'IASYMtoR', 'HtoR', 'HtoF', + 'StoQS', 'EtoQE', 'IPREtoQPRE', 'ISYMtoQSYM', 'IASYMtoQASYM', + 'QStoQE', 'QEtoQPRE', 'QPREtoQSYM', 'QPREtoQASYM', 'QSYMtoQR', 'QSYMtoH', 'QASYMtoQR', 'RtoS', '_toS' ] return propensities, columns @@ -2493,7 +2493,7 @@ def introduce_exposures(self, num_new_exposures): #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def increase_data_series_length(self): self.tseries = numpy.pad(self.tseries, [(0, 6*self.numNodes)], mode='constant', constant_values=0) @@ -2540,7 +2540,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] @@ -2587,7 +2587,7 @@ def finalize_data_series(self): return None #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def run_iteration(self, max_dt=None): @@ -2688,7 +2688,7 @@ def run_iteration(self, max_dt=None): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ 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) @@ -2706,7 +2706,7 @@ def run_iteration(self, max_dt=None): 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) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -2777,7 +2777,7 @@ def run(self, T, checkpoints=None, max_dt=None, min_dt=None, print_interval=10, 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 + checkpoints = None else: checkpointTime = checkpoints['t'][checkpointIdx] @@ -2806,7 +2806,7 @@ def run(self, T, checkpoints=None, max_dt=None, min_dt=None, print_interval=10, 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 + checkpoints = None else: checkpointTime = checkpoints['t'][checkpointIdx] #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -2833,7 +2833,7 @@ def run(self, T, checkpoints=None, max_dt=None, min_dt=None, print_interval=10, 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 @@ -2846,15 +2846,15 @@ def run(self, T, checkpoints=None, max_dt=None, min_dt=None, print_interval=10, 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_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_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_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', + 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): @@ -2871,22 +2871,22 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_ #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ 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_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) + 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: + # Draw the reference data: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(dashed_reference_results): dashedReference_tseries = dashed_reference_results.tseries[::int(self.numNodes/100)] @@ -2906,12 +2906,12 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_ 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) @@ -2921,12 +2921,12 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_ 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) @@ -2936,7 +2936,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_ 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) @@ -2946,7 +2946,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_ 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) @@ -2956,7 +2956,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_ 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) @@ -2966,7 +2966,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_ 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) @@ -2976,19 +2976,19 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_ 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) @@ -2996,11 +2996,11 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_ 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) @@ -3008,7 +3008,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_ 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) @@ -3016,7 +3016,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_ 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) @@ -3024,7 +3024,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_ 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) @@ -3036,7 +3036,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_ 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) @@ -3046,37 +3046,37 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_ #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ 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) @@ -3085,7 +3085,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_ 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) @@ -3119,7 +3119,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_ 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 @@ -3128,15 +3128,15 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_ 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_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_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_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', + 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): @@ -3152,15 +3152,15 @@ def figure_basic(self, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_ 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_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_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_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, + 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) @@ -3175,15 +3175,15 @@ def figure_basic(self, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_ 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_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_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_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', + 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): @@ -3199,29 +3199,23 @@ def figure_infections(self, plot_S=False, plot_E='stacked', plot_I_pre='stacked' 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_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_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_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, + 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/networks.py b/seirsplus/networks.py index c014a6a..a123617 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: @@ -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 @@ -166,15 +166,15 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F homelessNodes = N # Number of individuals to place in households curMemberIndex = 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,16 +182,16 @@ 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'): - + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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'): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -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,9 +275,9 @@ 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']): @@ -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): 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: ######################################### @@ -357,8 +357,8 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F # 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']) @@ -402,18 +402,18 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F 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'): # 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,7 +421,7 @@ 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'): @@ -429,7 +429,7 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F 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) @@ -437,7 +437,7 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F # 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): print(layerGroup+" public mean degree = "+str((meanDegree))) print(layerGroup+" public max degree = "+str((maxDegree))) @@ -452,7 +452,7 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F # 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'])) @@ -469,7 +469,7 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F 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. @@ -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() @@ -522,15 +522,15 @@ 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): - + #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # 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) - + ######################################### ######################################### @@ -609,12 +609,12 @@ def household_country_data(country): 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,7 +641,7 @@ 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): @@ -652,7 +652,7 @@ def custom_exponential_graph(base_graph=None, scale=100, min_num_edges=0, m=9, n 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): @@ -661,7 +661,7 @@ def custom_exponential_graph(base_graph=None, scale=100, min_num_edges=0, m=9, n for neighbor in neighbors: if(neighbor not in quarantineKeepNeighbors): graph.remove_edge(node, neighbor) - + return graph #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ @@ -690,16 +690,3 @@ def plot_degree_distn(graph, max_degree=None, show=True, use_seaborn=True): pyplot.legend(loc='upper right') if(show): pyplot.show() - - - - - - - - - - - - - diff --git a/seirsplus/sim_loops.py b/seirsplus/sim_loops.py index b855915..3a085da 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, @@ -40,7 +40,7 @@ def run_tti_sim(model, T, max_dt=None, #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(temporal_falseneg_rates is None): - temporal_falseneg_rates = { + 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}, @@ -87,7 +87,7 @@ def run_tti_sim(model, T, max_dt=None, timeOfLastIntroduction = model.t numNewExposures = numpy.random.poisson(lam=average_introductions_per_day) - + model.introduce_exposures(num_new_exposures=numNewExposures) if(numNewExposures > 0): @@ -98,7 +98,7 @@ def run_tti_sim(model, T, max_dt=None, #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(int(model.t)!=int(timeOfLastIntervention)): - + cadenceDayNumbers = [int(model.t % cadence_cycle_length)] if(backlog_skipped_intervals): @@ -114,11 +114,11 @@ def run_tti_sim(model, T, max_dt=None, if(currentPctInfected >= intervention_start_pct_infected and not interventionOn): interventionOn = True interventionStartTime = model.t - + 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 = [] @@ -144,7 +144,7 @@ def run_tti_sim(model, T, max_dt=None, for symptomaticNode in symptomaticNodes: if(isolation_compliance_symptomatic_individual[symptomaticNode]): if(model.X[symptomaticNode] == model.I_sym): - numSelfIsolated_symptoms += 1 + numSelfIsolated_symptoms += 1 newIsolationGroup_symptomatic.append(symptomaticNode) #---------------------------------------- @@ -169,7 +169,7 @@ def run_tti_sim(model, T, max_dt=None, for contactNode in tracingPoolQueue[0]: 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: @@ -181,7 +181,7 @@ def run_tti_sim(model, T, max_dt=None, if(isolation_compliance_positive_contactgroupmate[isolationGroupmate]): newIsolationGroup_contact.append(isolationGroupmate) numSelfIsolated_positiveContactGroupmate += 1 - + #---------------------------------------- # Update the nodeStates list after self-isolation updates to model.X: @@ -199,7 +199,7 @@ def run_tti_sim(model, T, max_dt=None, symptomaticSelection = [] if(any(testing_compliance_symptomatic)): - + symptomaticPool = numpy.argwhere((testing_compliance_symptomatic==True) & (nodeTestedInCurrentStateStatuses==False) & (nodePositiveStatuses==False) @@ -207,7 +207,7 @@ def run_tti_sim(model, T, max_dt=None, ).flatten() numSymptomaticTests = min(len(symptomaticPool), max_symptomatic_tests_per_day) - + if(len(symptomaticPool) > 0): symptomaticSelection = symptomaticPool[numpy.random.choice(len(symptomaticPool), min(numSymptomaticTests, len(symptomaticPool)), replace=False)] @@ -223,7 +223,7 @@ def run_tti_sim(model, T, max_dt=None, 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: #---------------------------------------- @@ -238,7 +238,7 @@ def run_tti_sim(model, T, max_dt=None, 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) @@ -248,24 +248,24 @@ def run_tti_sim(model, T, max_dt=None, #---------------------------------------- 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): 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 = [] @@ -299,19 +299,19 @@ def run_tti_sim(model, T, max_dt=None, 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): 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): 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'): testNodeState = model.X[testNode][0] testNodeTimeInState = model.timer_state[testNode][0] @@ -332,8 +332,8 @@ def run_tti_sim(model, T, max_dt=None, 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) @@ -345,7 +345,7 @@ def run_tti_sim(model, T, max_dt=None, #---------------------------------------- # Add the groupmates of this positive node to the isolation group: - #---------------------------------------- + #---------------------------------------- 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: @@ -354,9 +354,9 @@ def run_tti_sim(model, T, max_dt=None, 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]): testNodeContacts = list(model.G[testNode].keys()) @@ -367,7 +367,7 @@ def run_tti_sim(model, T, max_dt=None, 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..9f189be 100644 --- a/seirsplus/utilities.py +++ b/seirsplus/utilities.py @@ -11,14 +11,14 @@ def dist_info(dists, names=None, plot=False, bin_size=1, colors=None, reverse_pl 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']) - + 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): 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): pyplot.ylabel('num nodes') pyplot.legend(loc='upper right') @@ -30,25 +30,25 @@ def network_info(networks, names=None, plot=False, bin_size=1, colors=None, reve 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']) - + for i, (network, name) in enumerate(zip(networks, names)): - + degree = [d[1] for d in network.degree()] 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): 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): pyplot.ylabel('num nodes') pyplot.legend(loc='upper right') @@ -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) ) - - - - - - - - - - - - - - From 9468fcd5745e44cd46aa04720d1af5a63c9f892d Mon Sep 17 00:00:00 2001 From: Noel Schutt Date: Sat, 13 Mar 2021 10:15:54 -0500 Subject: [PATCH 02/12] Added missing space after `if` statements. --- seirsplus/FARZ.py | 2 +- seirsplus/legacy_models.py | 262 +++++++-------- seirsplus/models.py | 662 ++++++++++++++++++------------------- seirsplus/networks.py | 70 ++-- seirsplus/sim_loops.py | 86 ++--- seirsplus/utilities.py | 16 +- 6 files changed, 549 insertions(+), 549 deletions(-) diff --git a/seirsplus/FARZ.py b/seirsplus/FARZ.py index 309f657..b103ecd 100644 --- a/seirsplus/FARZ.py +++ b/seirsplus/FARZ.py @@ -198,7 +198,7 @@ def choose_node(i,c, G, C, alpha, beta, gamma, epsilon): j = trim_ids[ind] p[ind] = (cn[j]**alpha )/ ((dd[ind]+1)** gamma) - if(sum(p)==0): return None + if (sum(p)==0): return None tmp = random_choice(range(len(p)), p ) #, size=1, replace = False) # TODO add weights /direction/attributes if tmp is None: return None diff --git a/seirsplus/legacy_models.py b/seirsplus/legacy_models.py index 333af64..66fcb71 100644 --- a/seirsplus/legacy_models.py +++ b/seirsplus/legacy_models.py @@ -173,7 +173,7 @@ 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 @@ -181,7 +181,7 @@ def run(self, T, dt=0.1, checkpoints=None, verbose=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', @@ -189,7 +189,7 @@ def run(self, T, dt=0.1, checkpoints=None, verbose=False): 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()) + 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 @@ -197,13 +197,13 @@ def run(self, T, dt=0.1, checkpoints=None, verbose=False): #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% # 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])) @@ -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,7 +243,7 @@ def run(self, T, dt=0.1, checkpoints=None, verbose=False): #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_infections(self, t_idx=None): - if(t_idx is None): + 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]) @@ -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() #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -284,11 +284,11 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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,36 +298,36 @@ 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 @@ -336,64 +336,64 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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,14 +403,14 @@ 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') @@ -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() @@ -448,7 +448,7 @@ def figure_basic(self, plot_S='line', plot_E='line', plot_I='line',plot_R='line' 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 @@ -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() @@ -485,7 +485,7 @@ def figure_infections(self, plot_S=False, plot_E='stacked', plot_I='stacked',plo 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 @@ -547,7 +547,7 @@ 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) @@ -609,7 +609,7 @@ 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 @@ -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), @@ -688,30 +688,30 @@ 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']) @@ -719,14 +719,14 @@ def update_parameters(self): 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() #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -790,7 +790,7 @@ def update_scenario_flags(self): #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_infections(self, t_idx=None): - if(t_idx is None): + 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]) @@ -805,18 +805,18 @@ 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)) ) ) @@ -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) @@ -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] @@ -924,7 +924,7 @@ def finalize_data_series(self): 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 @@ -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,12 +1020,12 @@ 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 else: @@ -1042,23 +1042,23 @@ 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 else: @@ -1067,11 +1067,11 @@ def run(self, T, checkpoints=None, print_interval=10, verbose='t'): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - 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 @@ -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() #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -1122,11 +1122,11 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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,36 +1136,36 @@ 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 @@ -1174,64 +1174,64 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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,14 +1241,14 @@ 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') @@ -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() @@ -1286,7 +1286,7 @@ def figure_basic(self, plot_S='line', plot_E='line', plot_I='line',plot_R='line' 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 @@ -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() @@ -1323,7 +1323,7 @@ def figure_infections(self, plot_S=False, plot_E='stacked', plot_I='stacked',plo 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 diff --git a/seirsplus/models.py b/seirsplus/models.py index 339feee..104f1e7 100644 --- a/seirsplus/models.py +++ b/seirsplus/models.py @@ -172,7 +172,7 @@ 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 @@ -180,7 +180,7 @@ def run(self, T, dt=0.1, checkpoints=None, verbose=False): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Pre-process checkpoint values: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(checkpoints): + if (checkpoints): numCheckpoints = len(checkpoints['t']) paramNames = ['beta', 'sigma', 'gamma', 'xi', 'mu_I', 'mu_0', 'nu', 'beta_Q', 'sigma_Q', 'gamma_Q', 'mu_Q', @@ -188,7 +188,7 @@ def run(self, T, dt=0.1, checkpoints=None, verbose=False): 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()) + 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 @@ -196,13 +196,13 @@ def run(self, T, dt=0.1, checkpoints=None, verbose=False): #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% # 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])) @@ -224,7 +224,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])) @@ -233,7 +233,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 @@ -242,7 +242,7 @@ def run(self, T, dt=0.1, checkpoints=None, verbose=False): #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_susceptible(self, t_idx=None): - if(t_idx is None): + if (t_idx is None): return (self.numS[:]) else: return (self.numS[t_idx]) @@ -250,7 +250,7 @@ def total_num_susceptible(self, t_idx=None): #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_infected(self, t_idx=None): - if(t_idx is 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]) @@ -258,7 +258,7 @@ def total_num_infected(self, t_idx=None): #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_isolated(self, t_idx=None): - if(t_idx is 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]) @@ -266,7 +266,7 @@ def total_num_isolated(self, t_idx=None): #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_recovered(self, t_idx=None): - if(t_idx is None): + if (t_idx is None): return (self.numR[:]) else: return (self.numR[t_idx]) @@ -289,7 +289,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() #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -307,11 +307,11 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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.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): + 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) @@ -321,36 +321,36 @@ 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_Q and plot_Q_E=='stacked' and plot_Q_I=='stacked'): + 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'): + 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'): + 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'): + 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 @@ -359,64 +359,64 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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_Q and (any(Dseries) and plot_Q_E=='shaded' and plot_Q_E=='shaded')): + 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'): + 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'): + 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'): + 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_Q and (any(Dseries) and plot_Q_E=='line' and plot_Q_E=='line')): + 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'): + 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'): + 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'): + 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) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -426,14 +426,14 @@ 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') @@ -457,7 +457,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() @@ -471,7 +471,7 @@ def figure_basic(self, plot_S='line', plot_E='line', plot_I='line',plot_R='line' 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 @@ -494,7 +494,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() @@ -508,7 +508,7 @@ def figure_infections(self, plot_S=False, plot_E='stacked', plot_I='stacked',plo 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 @@ -571,7 +571,7 @@ def __init__(self, G, beta, sigma, gamma, 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): + if (seed is not None): numpy.random.seed(seed) self.seed = seed @@ -648,7 +648,7 @@ def __init__(self, G, beta, sigma, gamma, numpy.random.shuffle(self.X) self.store_Xseries = store_Xseries - if(store_Xseries): + if (store_Xseries): self.Xseries = numpy.zeros(shape=(6*self.numNodes, self.numNodes), dtype='uint8') self.Xseries[0,:] = self.X.T @@ -684,7 +684,7 @@ def __init__(self, G, beta, sigma, gamma, # 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), @@ -728,7 +728,7 @@ def update_parameters(self): 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): + 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'] @@ -783,19 +783,19 @@ def update_parameters(self): #---------------------------------------- # Global transmission parameters: #---------------------------------------- - if(self.beta_pairwise_mode == 'infected' or self.beta_pairwise_mode is None): + 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'): + elif (self.beta_pairwise_mode == 'infectee'): self.beta_global = self.beta self.beta_Q_global = self.beta_Q - elif(self.beta_pairwise_mode == 'min'): + 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'): + 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'): + 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 @@ -806,48 +806,48 @@ def update_parameters(self): 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): + 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))): + 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'): + if (self.beta_pairwise_mode == 'infected'): self.A_beta_pairwise = A_beta_pairwise_byInfected - elif(self.beta_pairwise_mode == 'infectee'): + elif (self.beta_pairwise_mode == 'infectee'): self.A_beta_pairwise = A_beta_pairwise_byInfectee - elif(self.beta_pairwise_mode == 'min'): + 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'): + 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): + 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): + 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))): + 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'): + if (self.beta_pairwise_mode == 'infected'): self.A_Q_beta_Q_pairwise = A_Q_beta_Q_pairwise_byInfected - elif(self.beta_pairwise_mode == 'infectee'): + elif (self.beta_pairwise_mode == 'infectee'): self.A_Q_beta_Q_pairwise = A_Q_beta_Q_pairwise_byInfectee - elif(self.beta_pairwise_mode == 'min'): + 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'): + 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): + 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').") @@ -865,52 +865,52 @@ def update_parameters(self): 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): + 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))): + 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'): + if (self.delta_pairwise_mode == 'infected'): self.A_delta_pairwise = A_delta_pairwise_byInfected - elif(self.delta_pairwise_mode == 'infectee'): + elif (self.delta_pairwise_mode == 'infectee'): self.A_delta_pairwise = A_delta_pairwise_byInfectee - elif(self.delta_pairwise_mode == 'min'): + 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'): + 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'): + 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): + 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): + 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))): + 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'): + if (self.delta_pairwise_mode == 'infected'): self.A_Q_delta_Q_pairwise = A_Q_delta_Q_pairwise_byInfected - elif(self.delta_pairwise_mode == 'infectee'): + elif (self.delta_pairwise_mode == 'infectee'): self.A_Q_delta_Q_pairwise = A_Q_delta_Q_pairwise_byInfectee - elif(self.delta_pairwise_mode == 'min'): + 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'): + 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'): + 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): + 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').") @@ -935,7 +935,7 @@ def node_degrees(self, Amat): #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_susceptible(self, t_idx=None): - if(t_idx is None): + if (t_idx is None): return (self.numS[:]) else: return (self.numS[t_idx]) @@ -943,7 +943,7 @@ def total_num_susceptible(self, t_idx=None): #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_infected(self, t_idx=None): - if(t_idx is 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]) @@ -951,7 +951,7 @@ def total_num_infected(self, t_idx=None): #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_isolated(self, t_idx=None): - if(t_idx is 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]) @@ -959,7 +959,7 @@ def total_num_isolated(self, t_idx=None): #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_tested(self, t_idx=None): - if(t_idx is None): + if (t_idx is None): return (self.numTested[:]) else: return (self.numTested[t_idx]) @@ -967,7 +967,7 @@ def total_num_tested(self, t_idx=None): #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_positive(self, t_idx=None): - if(t_idx is None): + if (t_idx is None): return (self.numPositive[:]) else: return (self.numPositive[t_idx]) @@ -975,7 +975,7 @@ def total_num_positive(self, t_idx=None): #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_recovered(self, t_idx=None): - if(t_idx is None): + if (t_idx is None): return (self.numR[:]) else: return (self.numR[t_idx]) @@ -992,19 +992,19 @@ def calc_propensities(self): #------------------------------------ self.transmissionTerms_I = numpy.zeros(shape=(self.numNodes,1)) - if(numpy.any(self.numI[self.tidx])): + 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])): + 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))): + 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)))) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -1017,7 +1017,7 @@ def calc_propensities(self): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(self.transition_mode == 'time_in_state'): + if (self.transition_mode == 'time_in_state'): propensities_EtoI = 1e5 * ((self.X==self.E) & numpy.greater(self.timer_state, 1/self.sigma)) @@ -1090,18 +1090,18 @@ def calc_propensities(self): 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): + 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): + 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): + 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): + elif (self.X[node] == self.Q_I): self.X[node] = self.I self.timer_state = 0 # Reset the isolation timer: @@ -1123,7 +1123,7 @@ def set_positive(self, 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): + if (self.X[exposedNode]==self.S): self.X[exposedNode] = self.E @@ -1143,10 +1143,10 @@ def increase_data_series_length(self): 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): + if (self.store_Xseries): self.Xseries = numpy.pad(self.Xseries, [(0, 6*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, 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) @@ -1176,10 +1176,10 @@ def finalize_data_series(self): 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): + 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] @@ -1199,7 +1199,7 @@ def finalize_data_series(self): 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() @@ -1214,7 +1214,7 @@ def run_iteration(self): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ propensities, transitionTypes = self.calc_propensities() - if(propensities.sum() > 0): + if (propensities.sum() > 0): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Calculate alpha @@ -1250,7 +1250,7 @@ def run_iteration(self): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Save information about infection events when they occur: - if(transitionType == 'StoE'): + if (transitionType == 'StoE'): transitionNode_GNbrs = list(self.G[transitionNode].keys()) transitionNode_GQNbrs = list(self.G_Q[transitionNode].keys()) self.infectionsLog.append({ 't': self.t, @@ -1264,7 +1264,7 @@ def run_iteration(self): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(transitionType in ['EtoQE', 'ItoQI']): + if (transitionType in ['EtoQE', 'ItoQI']): self.set_positive(node=transitionNode, positive=True) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -1305,10 +1305,10 @@ def run_iteration(self): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Store system states #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - 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) @@ -1324,7 +1324,7 @@ def run_iteration(self): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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)): + 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 @@ -1337,7 +1337,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 @@ -1345,12 +1345,12 @@ 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 else: @@ -1367,19 +1367,19 @@ 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: 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 else: @@ -1388,11 +1388,11 @@ def run(self, T, checkpoints=None, print_interval=10, verbose='t'): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - 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])) @@ -1402,7 +1402,7 @@ def run(self, T, checkpoints=None, print_interval=10, verbose='t'): 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)): + elif (not print_reset and (int(self.t) % 10 != 0)): print_reset = True return True @@ -1425,7 +1425,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() #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -1443,11 +1443,11 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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.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): + 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) @@ -1457,36 +1457,36 @@ 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_Q_E=='stacked' and plot_Q_I=='stacked'): + 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'): + 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'): + 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'): + 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 @@ -1495,64 +1495,64 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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_Q_E=='shaded' and plot_Q_I=='shaded')): + 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'): + 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'): + 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'): + 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_Q_E=='line' and plot_Q_I=='line')): + 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'): + 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'): + 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'): + 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) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -1562,14 +1562,14 @@ 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') @@ -1593,7 +1593,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() @@ -1607,7 +1607,7 @@ def figure_basic(self, plot_S='line', plot_E='line', plot_I='line',plot_R='line' 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 @@ -1630,7 +1630,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() @@ -1644,7 +1644,7 @@ def figure_infections(self, plot_S=False, plot_E='stacked', plot_I='stacked',plo 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 @@ -1741,7 +1741,7 @@ def __init__(self, G, beta, sigma, lamda, gamma, o=0, prevalence_ext=0, transition_mode='exponential_rates', node_groups=None, store_Xseries=False, seed=None): - if(seed is not None): + if (seed is not None): numpy.random.seed(seed) self.seed = seed @@ -1848,7 +1848,7 @@ def __init__(self, G, beta, sigma, lamda, gamma, numpy.random.shuffle(self.X) self.store_Xseries = store_Xseries - if(store_Xseries): + if (store_Xseries): self.Xseries = numpy.zeros(shape=(6*self.numNodes, self.numNodes), dtype='uint8') self.Xseries[0,:] = self.X.T @@ -1896,7 +1896,7 @@ def __init__(self, G, beta, sigma, lamda, gamma, # 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), @@ -1954,7 +1954,7 @@ def update_parameters(self): 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): + 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'] @@ -2035,23 +2035,23 @@ def update_parameters(self): #---------------------------------------- # Global transmission parameters: #---------------------------------------- - if(self.beta_pairwise_mode == 'infected' or self.beta_pairwise_mode is None): + 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'): + 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'): + 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'): + 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'): + 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 @@ -2063,74 +2063,74 @@ def update_parameters(self): 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): + 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))): + 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'): + if (self.beta_pairwise_mode == 'infected'): self.A_beta_pairwise = A_beta_pairwise_byInfected - elif(self.beta_pairwise_mode == 'infectee'): + elif (self.beta_pairwise_mode == 'infectee'): self.A_beta_pairwise = A_beta_pairwise_byInfectee - elif(self.beta_pairwise_mode == 'min'): + 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'): + 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): + 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): + 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))): + 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'): + if (self.beta_pairwise_mode == 'infected'): self.A_Q_beta_Q_pairwise = A_Q_beta_Q_pairwise_byInfected - elif(self.beta_pairwise_mode == 'infectee'): + elif (self.beta_pairwise_mode == 'infectee'): self.A_Q_beta_Q_pairwise = A_Q_beta_Q_pairwise_byInfectee - elif(self.beta_pairwise_mode == 'min'): + 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'): + 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): + 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): + 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): + 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))): + 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'): + if (self.beta_pairwise_mode == 'infected'): self.A_beta_asym_pairwise = A_beta_asym_pairwise_byInfected - elif(self.beta_pairwise_mode == 'infectee'): + elif (self.beta_pairwise_mode == 'infectee'): self.A_beta_asym_pairwise = A_beta_asym_pairwise_byInfectee - elif(self.beta_pairwise_mode == 'min'): + 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'): + 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): + 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').") @@ -2147,52 +2147,52 @@ def update_parameters(self): 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): + 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))): + 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'): + if (self.delta_pairwise_mode == 'infected'): self.A_delta_pairwise = A_delta_pairwise_byInfected - elif(self.delta_pairwise_mode == 'infectee'): + elif (self.delta_pairwise_mode == 'infectee'): self.A_delta_pairwise = A_delta_pairwise_byInfectee - elif(self.delta_pairwise_mode == 'min'): + 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'): + 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'): + 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): + 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): + 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))): + 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'): + if (self.delta_pairwise_mode == 'infected'): self.A_Q_delta_Q_pairwise = A_Q_delta_Q_pairwise_byInfected - elif(self.delta_pairwise_mode == 'infectee'): + elif (self.delta_pairwise_mode == 'infectee'): self.A_Q_delta_Q_pairwise = A_Q_delta_Q_pairwise_byInfectee - elif(self.delta_pairwise_mode == 'min'): + 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'): + 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'): + 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): + 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').") @@ -2204,7 +2204,7 @@ def update_parameters(self): #---------------------------------------- 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): + 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 @@ -2221,7 +2221,7 @@ def node_degrees(self, Amat): #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_susceptible(self, t_idx=None): - if(t_idx is None): + if (t_idx is None): return (self.numS[:] + self.numQ_S[:]) else: return (self.numS[t_idx] + self.numQ_S[t_idx]) @@ -2229,7 +2229,7 @@ def total_num_susceptible(self, t_idx=None): #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_infected(self, t_idx=None): - if(t_idx is 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: @@ -2239,7 +2239,7 @@ def total_num_infected(self, t_idx=None): #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_isolated(self, t_idx=None): - if(t_idx is 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]) @@ -2247,7 +2247,7 @@ def total_num_isolated(self, t_idx=None): #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_tested(self, t_idx=None): - if(t_idx is None): + if (t_idx is None): return (self.numTested[:]) else: return (self.numTested[t_idx]) @@ -2255,7 +2255,7 @@ def total_num_tested(self, t_idx=None): #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_positive(self, t_idx=None): - if(t_idx is None): + if (t_idx is None): return (self.numPositive[:]) else: return (self.numPositive[t_idx]) @@ -2263,7 +2263,7 @@ def total_num_positive(self, t_idx=None): #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_recovered(self, t_idx=None): - if(t_idx is None): + if (t_idx is None): return (self.numR[:] + self.numQ_R[:]) else: return (self.numR[t_idx] + self.numQ_R[t_idx]) @@ -2280,8 +2280,8 @@ def calc_propensities(self): #------------------------------------ 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): + 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 @@ -2291,19 +2291,19 @@ def calc_propensities(self): #------------------------------------ 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])): + 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]))): + 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))): + 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)))) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -2318,7 +2318,7 @@ def calc_propensities(self): )*(self.X==self.S) propensities_QStoQE = numpy.zeros_like(propensities_StoE) - if(numpy.any(self.X==self.Q_S)): + 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)*( @@ -2329,7 +2329,7 @@ def calc_propensities(self): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(self.transition_mode == 'time_in_state'): + if (self.transition_mode == 'time_in_state'): propensities_EtoIPRE = 1e5 * ((self.X==self.E) & numpy.greater(self.timer_state, 1/self.sigma)) @@ -2441,31 +2441,31 @@ def calc_propensities(self): 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): + if (isolate == True): + if (self.X[node] == self.S): self.X[node] = self.Q_S - elif(self.X[node] == self.E): + elif (self.X[node] == self.E): self.X[node] = self.Q_E - elif(self.X[node] == self.I_pre): + elif (self.X[node] == self.I_pre): self.X[node] = self.Q_pre - elif(self.X[node] == self.I_sym): + elif (self.X[node] == self.I_sym): self.X[node] = self.Q_sym - elif(self.X[node] == self.I_asym): + elif (self.X[node] == self.I_asym): self.X[node] = self.Q_asym - elif(self.X[node] == self.R): + elif (self.X[node] == self.R): self.X[node] = self.Q_R - elif(isolate == False): - if(self.X[node] == self.Q_S): + elif (isolate == False): + if (self.X[node] == self.Q_S): self.X[node] = self.S - elif(self.X[node] == self.Q_E): + elif (self.X[node] == self.Q_E): self.X[node] = self.E - elif(self.X[node] == self.Q_pre): + elif (self.X[node] == self.Q_pre): self.X[node] = self.I_pre - elif(self.X[node] == self.Q_sym): + elif (self.X[node] == self.Q_sym): self.X[node] = self.I_sym - elif(self.X[node] == self.Q_asym): + elif (self.X[node] == self.Q_asym): self.X[node] = self.I_asym - elif(self.X[node] == self.Q_R): + elif (self.X[node] == self.Q_R): self.X[node] = self.R # Reset the isolation timer: self.timer_isolation[node] = 0 @@ -2486,9 +2486,9 @@ def set_positive(self, 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): + if (self.X[exposedNode]==self.S): self.X[exposedNode] = self.E - elif(self.X[exposedNode]==self.Q_S): + elif (self.X[exposedNode]==self.Q_S): self.X[exposedNode] = self.Q_E @@ -2515,10 +2515,10 @@ def increase_data_series_length(self): 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): + if (self.store_Xseries): self.Xseries = numpy.pad(self.Xseries, [(0, 6*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, 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) @@ -2562,10 +2562,10 @@ def finalize_data_series(self): 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): + 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] @@ -2593,7 +2593,7 @@ 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): + 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() @@ -2608,7 +2608,7 @@ def run_iteration(self, max_dt=None): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ propensities, transitionTypes = self.calc_propensities() - if(propensities.sum() > 0): + if (propensities.sum() > 0): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Calculate alpha @@ -2622,7 +2622,7 @@ def run_iteration(self, max_dt=None): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ tau = (1/alpha)*numpy.log(float(1/r1)) - if(tau > max_dt): + 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) @@ -2660,7 +2660,7 @@ def run_iteration(self, max_dt=None): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Save information about infection events when they occur: - if(transitionType == 'StoE' or transitionType == 'QStoQE'): + 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, @@ -2674,7 +2674,7 @@ def run_iteration(self, max_dt=None): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(transitionType in ['EtoQE', 'IPREtoQPRE', 'ISYMtoQSYM', 'IASYMtoQASYM', 'ISYMtoH']): + if (transitionType in ['EtoQE', 'IPREtoQPRE', 'ISYMtoQSYM', 'IASYMtoQASYM', 'ISYMtoH']): self.set_positive(node=transitionNode, positive=True) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -2723,10 +2723,10 @@ def run_iteration(self, max_dt=None): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Store system states #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - 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) @@ -2749,7 +2749,7 @@ def run_iteration(self, max_dt=None): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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)): + 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 @@ -2762,7 +2762,7 @@ def run_iteration(self, max_dt=None): #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def run(self, T, checkpoints=None, max_dt=None, min_dt=None, print_interval=10, verbose='t'): - if(T>0): + if (T>0): self.tmax += T else: return False @@ -2770,12 +2770,12 @@ def run(self, T, checkpoints=None, max_dt=None, min_dt=None, print_interval=10, #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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 else: @@ -2792,19 +2792,19 @@ def run(self, T, checkpoints=None, max_dt=None, min_dt=None, print_interval=10, #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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: 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 else: @@ -2813,11 +2813,11 @@ def run(self, T, checkpoints=None, max_dt=None, min_dt=None, print_interval=10, #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - 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])) @@ -2835,7 +2835,7 @@ def run(self, T, checkpoints=None, max_dt=None, min_dt=None, print_interval=10, print("\t Q_R = " + str(self.numQ_R[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 @@ -2863,7 +2863,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_ #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Create an Axes object if None provided: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(not ax): + if (not ax): fig, ax = pyplot.subplots() #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -2888,11 +2888,11 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_ #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Draw the reference data: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(dashed_reference_results): + 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): + 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) @@ -2902,77 +2902,77 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_ # 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.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'): + 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'): + 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'): + 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'): + 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'): + 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'): + 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'): + 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'): + 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'): + 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'): + 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'): + 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'): + 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'): + 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'): + 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 @@ -2981,125 +2981,125 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_ #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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.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'): + 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'): + 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'): + 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'): + 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'): + 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'): + 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'): + 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'): + 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'): + 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'): + 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'): + 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'): + 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'): + 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'): + 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'): + 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'): + 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'): + 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'): + 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'): + 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'): + 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'): + 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'): + 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'): + 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'): + 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'): + 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'): + 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'): + 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'): + 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'): + 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): + 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) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -3109,14 +3109,14 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_ 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') @@ -3145,7 +3145,7 @@ def figure_basic(self, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_ fig, ax = pyplot.subplots(figsize=figsize) - if(use_seaborn): + if (use_seaborn): import seaborn seaborn.set_style('ticks') seaborn.despine() @@ -3164,7 +3164,7 @@ def figure_basic(self, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_ 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 @@ -3192,7 +3192,7 @@ def figure_infections(self, plot_S=False, plot_E='stacked', plot_I_pre='stacked' fig, ax = pyplot.subplots(figsize=figsize) - if(use_seaborn): + if (use_seaborn): import seaborn seaborn.set_style('ticks') seaborn.despine() @@ -3211,7 +3211,7 @@ def figure_infections(self, plot_S=False, plot_E='stacked', plot_I_pre='stacked' 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 diff --git a/seirsplus/networks.py b/seirsplus/networks.py index a123617..b57ec05 100644 --- a/seirsplus/networks.py +++ b/seirsplus/networks.py @@ -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) @@ -173,7 +173,7 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F household['ageBrackets'] = [] - if(household['situation'] == 'NOTu20_o60_eq1'): + if (household['situation'] == 'NOTu20_o60_eq1'): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Household size is definitely 1 @@ -183,7 +183,7 @@ 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 @@ -192,7 +192,7 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F # 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): @@ -216,7 +216,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_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): @@ -258,7 +258,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_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): @@ -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'] @@ -311,7 +311,7 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F 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 @@ -352,7 +352,7 @@ 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) }, @@ -399,7 +399,7 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F while(not graph_generated): try: - if(layer_generator == 'LFR'): + if (layer_generator == 'LFR'): # print "TARGET MEAN DEGREE = " + str(targetMeanDegree) @@ -409,7 +409,7 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F 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={ @@ -424,7 +424,7 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F 'phi': 1, 'b': 0.0, 'epsilon': 0.0000001, 'directed': False, 'weighted': False}) - elif(layer_generator == 'BA'): + elif (layer_generator == 'BA'): pass else: @@ -435,10 +435,10 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F 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 (meanDegree+meanHouseholdSize >= targetMeanDegreeRange[0] and meanDegree+meanHouseholdSize <= targetMeanDegreeRange[1]): + # if (meanDegree+meanHouseholdSize >= targetMeanDegree+meanHouseholdSize-1 and meanDegree+meanHouseholdSize <= targetMeanDegree+meanHouseholdSize+1): - if(verbose): + if (verbose): print(layerGroup+" public mean degree = "+str((meanDegree))) print(layerGroup+" public max degree = "+str((maxDegree))) @@ -446,7 +446,7 @@ 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)) @@ -460,13 +460,13 @@ 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])) @@ -475,9 +475,9 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F # 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)") #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -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,7 +521,7 @@ 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 @@ -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,7 +607,7 @@ 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, @@ -646,7 +646,7 @@ def household_country_data(country): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ 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." @@ -655,11 +655,11 @@ def custom_exponential_graph(base_graph=None, scale=100, min_num_edges=0, m=9, n # 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,5 +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 3a085da..8ac8ca3 100644 --- a/seirsplus/sim_loops.py +++ b/seirsplus/sim_loops.py @@ -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,7 +39,7 @@ def run_tti_sim(model, T, max_dt=None, #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if(temporal_falseneg_rates is None): + 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}, @@ -82,7 +82,7 @@ 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 @@ -90,18 +90,18 @@ def run_tti_sim(model, T, max_dt=None, 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,11 +111,11 @@ 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)) @@ -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): + 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,20 +165,20 @@ 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 #---------------------------------------- # 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 @@ -198,7 +198,7 @@ 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) @@ -208,7 +208,7 @@ def run_tti_sim(model, T, max_dt=None, 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,7 +220,7 @@ 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 @@ -229,13 +229,13 @@ def run_tti_sim(model, T, max_dt=None, 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) @@ -247,7 +247,7 @@ 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) @@ -262,7 +262,7 @@ def run_tti_sim(model, T, max_dt=None, 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)] @@ -294,42 +294,42 @@ 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 # 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 + 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 @@ -340,28 +340,28 @@ def run_tti_sim(model, T, max_dt=None, #---------------------------------------- # 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 diff --git a/seirsplus/utilities.py b/seirsplus/utilities.py index 9f189be..70d5e9a 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,13 +29,13 @@ 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), @@ -46,10 +46,10 @@ def network_info(networks, names=None, plot=False, bin_size=1, colors=None, reve 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() From 31eac1f39164cf2b867e6e3f0450ca5b0c8824c8 Mon Sep 17 00:00:00 2001 From: Noel Schutt Date: Sat, 13 Mar 2021 10:42:51 -0500 Subject: [PATCH 03/12] Cleaned up superfluous parens warnings. --- seirsplus/FARZ.py | 16 ++-- seirsplus/legacy_models.py | 148 ++++++++++++++++++------------------- seirsplus/models.py | 110 +++++++++++++-------------- seirsplus/networks.py | 64 ++++++++-------- seirsplus/sim_loops.py | 64 ++++++++-------- seirsplus/utilities.py | 10 +-- 6 files changed, 206 insertions(+), 206 deletions(-) diff --git a/seirsplus/FARZ.py b/seirsplus/FARZ.py index b103ecd..4b3444f 100644 --- a/seirsplus/FARZ.py +++ b/seirsplus/FARZ.py @@ -10,7 +10,7 @@ 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: @@ -198,7 +198,7 @@ def choose_node(i,c, G, C, alpha, beta, gamma, epsilon): j = trim_ids[ind] p[ind] = (cn[j]**alpha )/ ((dd[ind]+1)** gamma) - if (sum(p)==0): return None + if sum(p) == 0: return None tmp = random_choice(range(len(p)), p ) #, size=1, replace = False) # TODO add weights /direction/attributes if tmp is None: return None @@ -237,7 +237,7 @@ 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()0): + if T > 0: self.tmax += T else: return False @@ -181,7 +181,7 @@ def run(self, T, dt=0.1, checkpoints=None, verbose=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', @@ -197,13 +197,13 @@ def run(self, T, dt=0.1, checkpoints=None, verbose=False): #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% # 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])) @@ -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] #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ @@ -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() #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -284,11 +284,11 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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) @@ -393,7 +393,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin 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,14 +403,14 @@ 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') @@ -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() @@ -448,7 +448,7 @@ def figure_basic(self, plot_S='line', plot_E='line', plot_I='line',plot_R='line' 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 @@ -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() @@ -485,7 +485,7 @@ def figure_infections(self, plot_S=False, plot_E='stacked', plot_I='stacked',plo 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 @@ -547,7 +547,7 @@ 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) @@ -609,7 +609,7 @@ 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 @@ -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), @@ -688,30 +688,30 @@ 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']) @@ -719,14 +719,14 @@ def update_parameters(self): 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() #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -790,10 +790,10 @@ 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] #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ @@ -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) @@ -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] @@ -924,7 +924,7 @@ def finalize_data_series(self): 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 @@ -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) @@ -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,12 +1020,12 @@ 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 else: @@ -1042,23 +1042,23 @@ 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 else: @@ -1067,11 +1067,11 @@ def run(self, T, checkpoints=None, print_interval=10, verbose='t'): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if (print_interval): + if print_interval: if (print_reset and (int(self.t) % print_interval == 0)): - if (verbose=="t"): + 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])) @@ -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() #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -1122,11 +1122,11 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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) @@ -1231,7 +1231,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin 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,14 +1241,14 @@ 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') @@ -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() @@ -1286,7 +1286,7 @@ def figure_basic(self, plot_S='line', plot_E='line', plot_I='line',plot_R='line' 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 @@ -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() @@ -1323,7 +1323,7 @@ def figure_infections(self, plot_S=False, plot_E='stacked', plot_I='stacked',plo 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 diff --git a/seirsplus/models.py b/seirsplus/models.py index 104f1e7..0e0a1bd 100644 --- a/seirsplus/models.py +++ b/seirsplus/models.py @@ -172,7 +172,7 @@ 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 @@ -180,7 +180,7 @@ def run(self, T, dt=0.1, checkpoints=None, verbose=False): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Pre-process checkpoint values: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if (checkpoints): + if checkpoints: numCheckpoints = len(checkpoints['t']) paramNames = ['beta', 'sigma', 'gamma', 'xi', 'mu_I', 'mu_0', 'nu', 'beta_Q', 'sigma_Q', 'gamma_Q', 'mu_Q', @@ -196,13 +196,13 @@ def run(self, T, dt=0.1, checkpoints=None, verbose=False): #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% # 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])) @@ -224,7 +224,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])) @@ -233,7 +233,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 @@ -242,7 +242,7 @@ def run(self, T, dt=0.1, checkpoints=None, verbose=False): #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_susceptible(self, t_idx=None): - if (t_idx is None): + if t_idx is None: return (self.numS[:]) else: return (self.numS[t_idx]) @@ -250,26 +250,26 @@ def total_num_susceptible(self, t_idx=None): #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_infected(self, t_idx=None): - if (t_idx is None): - return (self.numE[:] + self.numI[:] + self.numQ_E[:] + self.numQ_I[:]) + 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]) + 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[:]) + if t_idx is None: + return self.numQ_E[:] + self.numQ_I[:] else: - return (self.numQ_E[t_idx] + self.numQ_I[t_idx]) + 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[:]) + if t_idx is None: + return self.numR[:] else: - return (self.numR[t_idx]) + return self.numR[t_idx] #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ @@ -289,7 +289,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() #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -307,11 +307,11 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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.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): + 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) @@ -416,7 +416,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin 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) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -426,14 +426,14 @@ 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') @@ -457,7 +457,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() @@ -471,7 +471,7 @@ def figure_basic(self, plot_S='line', plot_E='line', plot_I='line',plot_R='line' 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 @@ -508,7 +508,7 @@ def figure_infections(self, plot_S=False, plot_E='stacked', plot_I='stacked',plo 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 @@ -571,7 +571,7 @@ def __init__(self, G, beta, sigma, gamma, 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): + if seed is not None: numpy.random.seed(seed) self.seed = seed @@ -648,7 +648,7 @@ def __init__(self, G, beta, sigma, gamma, numpy.random.shuffle(self.X) self.store_Xseries = store_Xseries - if (store_Xseries): + if store_Xseries: self.Xseries = numpy.zeros(shape=(6*self.numNodes, self.numNodes), dtype='uint8') self.Xseries[0,:] = self.X.T @@ -684,7 +684,7 @@ def __init__(self, G, beta, sigma, gamma, # 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), @@ -728,7 +728,7 @@ def update_parameters(self): 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): + 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'] @@ -786,16 +786,16 @@ def update_parameters(self): 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'): + elif self.beta_pairwise_mode == 'infectee': self.beta_global = self.beta self.beta_Q_global = self.beta_Q - elif (self.beta_pairwise_mode == 'min'): + 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'): + 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'): + 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 @@ -815,15 +815,15 @@ def update_parameters(self): 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'): + if self.beta_pairwise_mode == 'infected': self.A_beta_pairwise = A_beta_pairwise_byInfected - elif (self.beta_pairwise_mode == 'infectee'): + elif self.beta_pairwise_mode == 'infectee': self.A_beta_pairwise = A_beta_pairwise_byInfectee - elif (self.beta_pairwise_mode == 'min'): + 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'): + 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): + 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').") @@ -839,13 +839,13 @@ def update_parameters(self): 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'): + if self.beta_pairwise_mode == 'infected': self.A_Q_beta_Q_pairwise = A_Q_beta_Q_pairwise_byInfected - elif (self.beta_pairwise_mode == 'infectee'): + elif self.beta_pairwise_mode == 'infectee': self.A_Q_beta_Q_pairwise = A_Q_beta_Q_pairwise_byInfectee - elif (self.beta_pairwise_mode == 'min'): + 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'): + 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 @@ -874,17 +874,17 @@ def update_parameters(self): 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'): + if self.delta_pairwise_mode == 'infected': self.A_delta_pairwise = A_delta_pairwise_byInfected - elif (self.delta_pairwise_mode == 'infectee'): + elif self.delta_pairwise_mode == 'infectee': self.A_delta_pairwise = A_delta_pairwise_byInfectee - elif (self.delta_pairwise_mode == 'min'): + 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'): + 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'): + 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): + 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').") @@ -900,15 +900,15 @@ def update_parameters(self): 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'): + if self.delta_pairwise_mode == 'infected': self.A_Q_delta_Q_pairwise = A_Q_delta_Q_pairwise_byInfected - elif (self.delta_pairwise_mode == 'infectee'): + elif self.delta_pairwise_mode == 'infectee': self.A_Q_delta_Q_pairwise = A_Q_delta_Q_pairwise_byInfectee - elif (self.delta_pairwise_mode == 'min'): + 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'): + 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'): + 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 diff --git a/seirsplus/networks.py b/seirsplus/networks.py index b57ec05..26aa421 100644 --- a/seirsplus/networks.py +++ b/seirsplus/networks.py @@ -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) @@ -165,7 +165,7 @@ 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 = {} @@ -173,7 +173,7 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F household['ageBrackets'] = [] - if (household['situation'] == 'NOTu20_o60_eq1'): + if household['situation'] == 'NOTu20_o60_eq1': #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Household size is definitely 1 @@ -183,7 +183,7 @@ 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 @@ -192,7 +192,7 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F # 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): @@ -216,7 +216,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_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): @@ -258,7 +258,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_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): @@ -280,7 +280,7 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F # elif (household['situation'] == 'u20_o60_eq1'): # impossible - if (len(household['ageBrackets']) == household['size']): + if len(household['ageBrackets']) == household['size']: homelessNodes -= household['size'] @@ -311,7 +311,7 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F 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 @@ -352,7 +352,7 @@ 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) }, @@ -397,9 +397,9 @@ 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) @@ -409,7 +409,7 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F 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={ @@ -424,7 +424,7 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F 'phi': 1, 'b': 0.0, 'epsilon': 0.0000001, 'directed': False, 'weighted': False}) - elif (layer_generator == 'BA'): + elif layer_generator == 'BA': pass else: @@ -438,7 +438,7 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F 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 verbose: print(layerGroup+" public mean degree = "+str((meanDegree))) print(layerGroup+" public max degree = "+str((maxDegree))) @@ -446,7 +446,7 @@ 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)) @@ -460,13 +460,13 @@ 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])) @@ -477,7 +477,7 @@ def generate_demographic_contact_network(N, demographic_data, layer_generator='F graph_gen_attempts += 1 # 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)") #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -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,7 +521,7 @@ 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 @@ -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,7 +607,7 @@ 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, @@ -646,7 +646,7 @@ def household_country_data(country): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ 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." @@ -655,11 +655,11 @@ def custom_exponential_graph(base_graph=None, scale=100, min_num_edges=0, m=9, n # 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,5 +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 8ac8ca3..b2b45b9 100644 --- a/seirsplus/sim_loops.py +++ b/seirsplus/sim_loops.py @@ -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,7 +39,7 @@ def run_tti_sim(model, T, max_dt=None, #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if (temporal_falseneg_rates is None): + 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}, @@ -82,7 +82,7 @@ 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 @@ -90,18 +90,18 @@ def run_tti_sim(model, T, max_dt=None, 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 @@ -115,7 +115,7 @@ def run_tti_sim(model, T, max_dt=None, interventionOn = True interventionStartTime = model.t - if (interventionOn): + if interventionOn: print("[INTERVENTIONS @ t = %.2f (%d (%.2f%%) infected)]" % (model.t, currentNumInfected, currentPctInfected*100)) @@ -139,11 +139,11 @@ 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): + if isolation_compliance_symptomatic_individual[symptomaticNode]: + if model.X[symptomaticNode] == model.I_sym: numSelfIsolated_symptoms += 1 newIsolationGroup_symptomatic.append(symptomaticNode) @@ -153,8 +153,8 @@ def run_tti_sim(model, T, max_dt=None, 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) @@ -167,7 +167,7 @@ def run_tti_sim(model, T, max_dt=None, 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 @@ -178,7 +178,7 @@ def run_tti_sim(model, T, max_dt=None, 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 isolation_compliance_positive_contactgroupmate[isolationGroupmate]: newIsolationGroup_contact.append(isolationGroupmate) numSelfIsolated_positiveContactGroupmate += 1 @@ -198,7 +198,7 @@ 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) @@ -208,7 +208,7 @@ def run_tti_sim(model, T, max_dt=None, 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,7 +220,7 @@ 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 @@ -229,7 +229,7 @@ def run_tti_sim(model, T, max_dt=None, 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)) @@ -247,7 +247,7 @@ 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) @@ -262,7 +262,7 @@ def run_tti_sim(model, T, max_dt=None, 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)] @@ -294,9 +294,9 @@ 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 @@ -312,24 +312,24 @@ def run_tti_sim(model, T, max_dt=None, 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 @@ -340,7 +340,7 @@ def run_tti_sim(model, T, max_dt=None, #---------------------------------------- # Add this positive node to the isolation group: #---------------------------------------- - if (isolation_compliance_positive_individual[testNode]): + if isolation_compliance_positive_individual[testNode]: newIsolationGroup_positive.append(testNode) #---------------------------------------- @@ -349,8 +349,8 @@ def run_tti_sim(model, T, max_dt=None, 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) @@ -358,10 +358,10 @@ def run_tti_sim(model, T, max_dt=None, # 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 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 diff --git a/seirsplus/utilities.py b/seirsplus/utilities.py index 70d5e9a..398b586 100644 --- a/seirsplus/utilities.py +++ b/seirsplus/utilities.py @@ -16,10 +16,10 @@ def dist_info(dists, names=None, plot=False, bin_size=1, colors=None, reverse_pl 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() @@ -35,7 +35,7 @@ def network_info(networks, names=None, plot=False, bin_size=1, colors=None, reve 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), @@ -46,10 +46,10 @@ def network_info(networks, names=None, plot=False, bin_size=1, colors=None, reve 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() From 1267f375f70b7ad29c9a515b71f5698bdd4bdde0 Mon Sep 17 00:00:00 2001 From: Noel Schutt Date: Sat, 13 Mar 2021 11:11:46 -0500 Subject: [PATCH 04/12] Spellcheck. --- seirsplus/legacy_models.py | 6 +++--- seirsplus/models.py | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/seirsplus/legacy_models.py b/seirsplus/legacy_models.py index e93c61f..d6d4e55 100644 --- a/seirsplus/legacy_models.py +++ b/seirsplus/legacy_models.py @@ -84,7 +84,7 @@ def __init__(self, initN, beta, sigma, gamma, xi=0, mu_I=0, mu_0=0, nu=0, p=0, 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,7 +132,7 @@ 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: diff --git a/seirsplus/models.py b/seirsplus/models.py index 0e0a1bd..c2f500e 100644 --- a/seirsplus/models.py +++ b/seirsplus/models.py @@ -103,7 +103,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_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 + 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 @@ -131,7 +131,7 @@ 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: @@ -619,7 +619,7 @@ def __init__(self, G, beta, sigma, gamma, self.isolationTime = isolation_time #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - # Initialize Counts of inidividuals with each state: + # Initialize Counts of individuals with each state: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ self.numE[0] = int(initE) self.numI[0] = int(initI) From f99218cf340d841d80413bb5677ae5507b93c4bc Mon Sep 17 00:00:00 2001 From: Noel Schutt Date: Sat, 13 Mar 2021 12:34:21 -0500 Subject: [PATCH 05/12] Split the models into separate files for ease of extension. --- seirsplus/models/__init__.py | 5 + .../extended_seirs_network_model.py} | 1670 +---------------- seirsplus/models/seirs_model.py | 504 +++++ seirsplus/models/seirs_network_model.py | 1136 +++++++++++ 4 files changed, 1648 insertions(+), 1667 deletions(-) create mode 100644 seirsplus/models/__init__.py rename seirsplus/{models.py => models/extended_seirs_network_model.py} (52%) create mode 100644 seirsplus/models/seirs_model.py create mode 100644 seirsplus/models/seirs_network_model.py 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.py b/seirsplus/models/extended_seirs_network_model.py similarity index 52% rename from seirsplus/models.py rename to seirsplus/models/extended_seirs_network_model.py index c2f500e..143f9eb 100644 --- a/seirsplus/models.py +++ b/seirsplus/models/extended_seirs_network_model.py @@ -2,1669 +2,9 @@ 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 # 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] - - -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ - - 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 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[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 @# -#@ @# -#@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@# -######################################################## +import networkx +import numpy +import scipy class ExtSEIRSNetworkModel(): @@ -3215,7 +1555,3 @@ def figure_infections(self, plot_S=False, plot_E='stacked', plot_I_pre='stacked' 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..115b8d5 --- /dev/null +++ b/seirsplus/models/seirs_model.py @@ -0,0 +1,504 @@ +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import numpy +import scipy.integrate + + +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 # 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] + + +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ +#^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + + 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 diff --git a/seirsplus/models/seirs_network_model.py b/seirsplus/models/seirs_network_model.py new file mode 100644 index 0000000..0266b4f --- /dev/null +++ b/seirsplus/models/seirs_network_model.py @@ -0,0 +1,1136 @@ +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import networkx +import numpy +import scipy + + +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 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[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 From 0ac31e06f08b553218f5797743e3e6fc31e5a176 Mon Sep 17 00:00:00 2001 From: Noel Schutt <1065410+schutt@users.noreply.github.com> Date: Sat, 13 Mar 2021 13:37:54 -0500 Subject: [PATCH 06/12] Reapplied superfluous parens fixes. --- .../models/extended_seirs_network_model.py | 200 +++++++++--------- seirsplus/models/seirs_model.py | 6 +- seirsplus/models/seirs_network_model.py | 122 +++++------ 3 files changed, 164 insertions(+), 164 deletions(-) diff --git a/seirsplus/models/extended_seirs_network_model.py b/seirsplus/models/extended_seirs_network_model.py index 143f9eb..b804912 100644 --- a/seirsplus/models/extended_seirs_network_model.py +++ b/seirsplus/models/extended_seirs_network_model.py @@ -188,7 +188,7 @@ def __init__(self, G, beta, sigma, lamda, gamma, numpy.random.shuffle(self.X) self.store_Xseries = store_Xseries - if (store_Xseries): + if store_Xseries: self.Xseries = numpy.zeros(shape=(6*self.numNodes, self.numNodes), dtype='uint8') self.Xseries[0,:] = self.X.T @@ -236,7 +236,7 @@ def __init__(self, G, beta, sigma, lamda, gamma, # 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), @@ -294,7 +294,7 @@ def update_parameters(self): 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): + 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'] @@ -379,19 +379,19 @@ def update_parameters(self): 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'): + 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'): + 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'): + 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'): + 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 @@ -412,13 +412,13 @@ def update_parameters(self): 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'): + if self.beta_pairwise_mode == 'infected': self.A_beta_pairwise = A_beta_pairwise_byInfected - elif (self.beta_pairwise_mode == 'infectee'): + elif self.beta_pairwise_mode == 'infectee': self.A_beta_pairwise = A_beta_pairwise_byInfectee - elif (self.beta_pairwise_mode == 'min'): + 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'): + 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 @@ -436,13 +436,13 @@ def update_parameters(self): 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'): + if self.beta_pairwise_mode == 'infected': self.A_Q_beta_Q_pairwise = A_Q_beta_Q_pairwise_byInfected - elif (self.beta_pairwise_mode == 'infectee'): + elif self.beta_pairwise_mode == 'infectee': self.A_Q_beta_Q_pairwise = A_Q_beta_Q_pairwise_byInfectee - elif (self.beta_pairwise_mode == 'min'): + 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'): + 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 @@ -451,7 +451,7 @@ def update_parameters(self): else: print("Invalid values given for beta_Q_local (expected 1xN list/array or NxN 2d array)") #---------------------------------------- - if (self.beta_asym_local is None): + 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 @@ -462,13 +462,13 @@ def update_parameters(self): 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'): + if self.beta_pairwise_mode == 'infected': self.A_beta_asym_pairwise = A_beta_asym_pairwise_byInfected - elif (self.beta_pairwise_mode == 'infectee'): + elif self.beta_pairwise_mode == 'infectee': self.A_beta_asym_pairwise = A_beta_asym_pairwise_byInfectee - elif (self.beta_pairwise_mode == 'min'): + 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'): + 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 @@ -496,17 +496,17 @@ def update_parameters(self): 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'): + if self.delta_pairwise_mode == 'infected': self.A_delta_pairwise = A_delta_pairwise_byInfected - elif (self.delta_pairwise_mode == 'infectee'): + elif self.delta_pairwise_mode == 'infectee': self.A_delta_pairwise = A_delta_pairwise_byInfectee - elif (self.delta_pairwise_mode == 'min'): + 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'): + 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'): + 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): + 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').") @@ -522,17 +522,17 @@ def update_parameters(self): 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'): + if self.delta_pairwise_mode == 'infected': self.A_Q_delta_Q_pairwise = A_Q_delta_Q_pairwise_byInfected - elif (self.delta_pairwise_mode == 'infectee'): + elif self.delta_pairwise_mode == 'infectee': self.A_Q_delta_Q_pairwise = A_Q_delta_Q_pairwise_byInfectee - elif (self.delta_pairwise_mode == 'min'): + 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'): + 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'): + 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): + 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').") @@ -544,7 +544,7 @@ def update_parameters(self): #---------------------------------------- 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): + 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 @@ -561,15 +561,15 @@ def node_degrees(self, Amat): #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_susceptible(self, t_idx=None): - if (t_idx is None): - return (self.numS[:] + self.numQ_S[:]) + if t_idx is None: + return self.numS[:] + self.numQ_S[:] else: - return (self.numS[t_idx] + self.numQ_S[t_idx]) + return self.numS[t_idx] + self.numQ_S[t_idx] #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_infected(self, t_idx=None): - if (t_idx is 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: @@ -579,34 +579,34 @@ def total_num_infected(self, t_idx=None): #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 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[:]) + 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]) + 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[:]) + if t_idx is None: + return self.numTested[:] else: - return (self.numTested[t_idx]) + return self.numTested[t_idx] #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_positive(self, t_idx=None): - if (t_idx is None): - return (self.numPositive[:]) + if t_idx is None: + return self.numPositive[:] else: - return (self.numPositive[t_idx]) + return self.numPositive[t_idx] #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_recovered(self, t_idx=None): - if (t_idx is None): - return (self.numR[:] + self.numQ_R[:]) + if t_idx is None: + return self.numR[:] + self.numQ_R[:] else: - return (self.numR[t_idx] + self.numQ_R[t_idx]) + return self.numR[t_idx] + self.numQ_R[t_idx] #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ @@ -621,7 +621,7 @@ def calc_propensities(self): 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): + 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 @@ -658,7 +658,7 @@ def calc_propensities(self): )*(self.X==self.S) propensities_QStoQE = numpy.zeros_like(propensities_StoE) - if (numpy.any(self.X==self.Q_S)): + 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)*( @@ -669,7 +669,7 @@ def calc_propensities(self): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if (self.transition_mode == 'time_in_state'): + if self.transition_mode == 'time_in_state': propensities_EtoIPRE = 1e5 * ((self.X==self.E) & numpy.greater(self.timer_state, 1/self.sigma)) @@ -781,31 +781,31 @@ def calc_propensities(self): 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): + if isolate == True: + if self.X[node] == self.S: self.X[node] = self.Q_S - elif (self.X[node] == self.E): + elif self.X[node] == self.E: self.X[node] = self.Q_E - elif (self.X[node] == self.I_pre): + elif self.X[node] == self.I_pre: self.X[node] = self.Q_pre - elif (self.X[node] == self.I_sym): + elif self.X[node] == self.I_sym: self.X[node] = self.Q_sym - elif (self.X[node] == self.I_asym): + elif self.X[node] == self.I_asym: self.X[node] = self.Q_asym - elif (self.X[node] == self.R): + elif self.X[node] == self.R: self.X[node] = self.Q_R - elif (isolate == False): - if (self.X[node] == self.Q_S): + elif isolate == False: + if self.X[node] == self.Q_S: self.X[node] = self.S - elif (self.X[node] == self.Q_E): + elif self.X[node] == self.Q_E: self.X[node] = self.E - elif (self.X[node] == self.Q_pre): + elif self.X[node] == self.Q_pre: self.X[node] = self.I_pre - elif (self.X[node] == self.Q_sym): + elif self.X[node] == self.Q_sym: self.X[node] = self.I_sym - elif (self.X[node] == self.Q_asym): + elif self.X[node] == self.Q_asym: self.X[node] = self.I_asym - elif (self.X[node] == self.Q_R): + elif self.X[node] == self.Q_R: self.X[node] = self.R # Reset the isolation timer: self.timer_isolation[node] = 0 @@ -826,9 +826,9 @@ def set_positive(self, 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): + if self.X[exposedNode] == self.S: self.X[exposedNode] = self.E - elif (self.X[exposedNode]==self.Q_S): + elif self.X[exposedNode] == self.Q_S: self.X[exposedNode] = self.Q_E @@ -855,10 +855,10 @@ def increase_data_series_length(self): 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): + if self.store_Xseries: self.Xseries = numpy.pad(self.Xseries, [(0, 6*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, 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) @@ -902,10 +902,10 @@ def finalize_data_series(self): 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): + 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] @@ -933,7 +933,7 @@ 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): + 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() @@ -948,7 +948,7 @@ def run_iteration(self, max_dt=None): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ propensities, transitionTypes = self.calc_propensities() - if (propensities.sum() > 0): + if propensities.sum() > 0: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Calculate alpha @@ -962,7 +962,7 @@ def run_iteration(self, max_dt=None): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ tau = (1/alpha)*numpy.log(float(1/r1)) - if (tau > max_dt): + 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) @@ -1063,10 +1063,10 @@ def run_iteration(self, max_dt=None): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Store system states #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - 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) @@ -1102,7 +1102,7 @@ def run_iteration(self, max_dt=None): #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def run(self, T, checkpoints=None, max_dt=None, min_dt=None, print_interval=10, verbose='t'): - if (T>0): + if T > 0: self.tmax += T else: return False @@ -1110,12 +1110,12 @@ def run(self, T, checkpoints=None, max_dt=None, min_dt=None, print_interval=10, #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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 else: @@ -1132,19 +1132,19 @@ def run(self, T, checkpoints=None, max_dt=None, min_dt=None, print_interval=10, #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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: 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 else: @@ -1153,11 +1153,11 @@ def run(self, T, checkpoints=None, max_dt=None, min_dt=None, print_interval=10, #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if (print_interval): + if print_interval: if (print_reset and (int(self.t) % print_interval == 0)): - if (verbose=="t"): + 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])) @@ -1203,7 +1203,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_ #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Create an Axes object if None provided: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if (not ax): + if not ax: fig, ax = pyplot.subplots() #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -1228,11 +1228,11 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_ #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Draw the reference data: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if (dashed_reference_results): + 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): + 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) @@ -1439,7 +1439,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_ 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) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -1449,14 +1449,14 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_ 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') @@ -1485,7 +1485,7 @@ def figure_basic(self, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_ fig, ax = pyplot.subplots(figsize=figsize) - if (use_seaborn): + if use_seaborn: import seaborn seaborn.set_style('ticks') seaborn.despine() @@ -1504,7 +1504,7 @@ def figure_basic(self, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_ 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 @@ -1532,7 +1532,7 @@ def figure_infections(self, plot_S=False, plot_E='stacked', plot_I_pre='stacked' fig, ax = pyplot.subplots(figsize=figsize) - if (use_seaborn): + if use_seaborn: import seaborn seaborn.set_style('ticks') seaborn.despine() @@ -1551,7 +1551,7 @@ def figure_infections(self, plot_S=False, plot_E='stacked', plot_I_pre='stacked' 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 diff --git a/seirsplus/models/seirs_model.py b/seirsplus/models/seirs_model.py index 115b8d5..3483ebc 100644 --- a/seirsplus/models/seirs_model.py +++ b/seirsplus/models/seirs_model.py @@ -233,9 +233,9 @@ def run(self, T, dt=0.1, checkpoints=None, verbose=False): def total_num_susceptible(self, t_idx=None): if t_idx is None: - return (self.numS[:]) + return self.numS[:] else: - return (self.numS[t_idx]) + return self.numS[t_idx] #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ @@ -243,7 +243,7 @@ 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] + return self.numE[t_idx] + self.numI[t_idx] + self.numQ_E[t_idx] + self.numQ_I[t_idx] #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/seirsplus/models/seirs_network_model.py b/seirsplus/models/seirs_network_model.py index 0266b4f..b7e37b9 100644 --- a/seirsplus/models/seirs_network_model.py +++ b/seirsplus/models/seirs_network_model.py @@ -421,50 +421,50 @@ def node_degrees(self, Amat): #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_susceptible(self, t_idx=None): - if (t_idx is None): - return (self.numS[:]) + if t_idx is None: + return self.numS[:] else: - return (self.numS[t_idx]) + 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[:]) + 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]) + 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[:]) + if t_idx is None: + return self.numQ_E[:] + self.numQ_I[:] else: - return (self.numQ_E[t_idx] + self.numQ_I[t_idx]) + 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[:]) + if t_idx is None: + return self.numTested[:] else: - return (self.numTested[t_idx]) + return self.numTested[t_idx] #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_positive(self, t_idx=None): - if (t_idx is None): - return (self.numPositive[:]) + if t_idx is None: + return self.numPositive[:] else: - return (self.numPositive[t_idx]) + return self.numPositive[t_idx] #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_recovered(self, t_idx=None): - if (t_idx is None): - return (self.numR[:]) + if t_idx is None: + return self.numR[:] else: - return (self.numR[t_idx]) + return self.numR[t_idx] #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ @@ -478,19 +478,19 @@ def calc_propensities(self): #------------------------------------ self.transmissionTerms_I = numpy.zeros(shape=(self.numNodes,1)) - if (numpy.any(self.numI[self.tidx])): + 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])): + 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))): + 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)))) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -503,7 +503,7 @@ def calc_propensities(self): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if (self.transition_mode == 'time_in_state'): + if self.transition_mode == 'time_in_state': propensities_EtoI = 1e5 * ((self.X==self.E) & numpy.greater(self.timer_state, 1/self.sigma)) @@ -576,18 +576,18 @@ def calc_propensities(self): 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): + 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): + 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): + 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): + elif self.X[node] == self.Q_I: self.X[node] = self.I self.timer_state = 0 # Reset the isolation timer: @@ -609,7 +609,7 @@ def set_positive(self, 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): + if self.X[exposedNode] == self.S: self.X[exposedNode] = self.E @@ -629,10 +629,10 @@ def increase_data_series_length(self): 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): + if self.store_Xseries: self.Xseries = numpy.pad(self.Xseries, [(0, 6*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, 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) @@ -662,10 +662,10 @@ def finalize_data_series(self): 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): + 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] @@ -685,7 +685,7 @@ def finalize_data_series(self): 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() @@ -700,7 +700,7 @@ def run_iteration(self): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ propensities, transitionTypes = self.calc_propensities() - if (propensities.sum() > 0): + if propensities.sum() > 0: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Calculate alpha @@ -736,7 +736,7 @@ def run_iteration(self): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Save information about infection events when they occur: - if (transitionType == 'StoE'): + if transitionType == 'StoE': transitionNode_GNbrs = list(self.G[transitionNode].keys()) transitionNode_GQNbrs = list(self.G_Q[transitionNode].keys()) self.infectionsLog.append({ 't': self.t, @@ -791,10 +791,10 @@ def run_iteration(self): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Store system states #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - 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) @@ -823,7 +823,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 @@ -831,12 +831,12 @@ 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 else: @@ -853,19 +853,19 @@ 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: 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 else: @@ -874,11 +874,11 @@ def run(self, T, checkpoints=None, print_interval=10, verbose='t'): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ - if (print_interval): + if print_interval: if (print_reset and (int(self.t) % print_interval == 0)): - if (verbose=="t"): + 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])) @@ -911,7 +911,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() #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -929,11 +929,11 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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.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): + 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) @@ -1038,7 +1038,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin 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) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -1048,14 +1048,14 @@ 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') @@ -1079,7 +1079,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() @@ -1093,7 +1093,7 @@ def figure_basic(self, plot_S='line', plot_E='line', plot_I='line',plot_R='line' 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 @@ -1116,7 +1116,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() @@ -1130,7 +1130,7 @@ def figure_infections(self, plot_S=False, plot_E='stacked', plot_I='stacked',plo 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 From e07b868c8683dbdab7fbc069b0e9f9f3cd4795a8 Mon Sep 17 00:00:00 2001 From: Noel Schutt <1065410+schutt@users.noreply.github.com> Date: Sat, 13 Mar 2021 13:42:08 -0500 Subject: [PATCH 07/12] Removed superfluous parens. --- seirsplus/models/seirs_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/seirsplus/models/seirs_model.py b/seirsplus/models/seirs_model.py index 3483ebc..fc2c9ad 100644 --- a/seirsplus/models/seirs_model.py +++ b/seirsplus/models/seirs_model.py @@ -484,7 +484,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() From 4eaf10081344042cf83c2c838d4391b8d0fae9fc Mon Sep 17 00:00:00 2001 From: Noel Schutt <1065410+schutt@users.noreply.github.com> Date: Sat, 13 Mar 2021 13:53:37 -0500 Subject: [PATCH 08/12] Use common name between basic and network models. --- seirsplus/models/seirs_network_model.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/seirsplus/models/seirs_network_model.py b/seirsplus/models/seirs_network_model.py index b7e37b9..9495daa 100644 --- a/seirsplus/models/seirs_network_model.py +++ b/seirsplus/models/seirs_network_model.py @@ -898,7 +898,7 @@ def run(self, T, checkpoints=None, print_interval=10, verbose='t'): #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 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, + 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', @@ -951,7 +951,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin 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'): + 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 @@ -987,7 +987,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin 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')): + 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: @@ -1014,7 +1014,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin 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')): + 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'): @@ -1066,7 +1066,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 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, + 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', @@ -1085,7 +1085,7 @@ def figure_basic(self, plot_S='line', plot_E='line', plot_I='line',plot_R='line' 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, + 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, @@ -1103,7 +1103,7 @@ def figure_basic(self, plot_S='line', plot_E='line', plot_I='line',plot_R='line' #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 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, + 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', @@ -1122,7 +1122,7 @@ def figure_infections(self, plot_S=False, plot_E='stacked', plot_I='stacked',plo 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, + 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, From d2e51728b7f41a5c15297b01abd92c0f2bcffe43 Mon Sep 17 00:00:00 2001 From: Noel Schutt <1065410+schutt@users.noreply.github.com> Date: Sat, 13 Mar 2021 13:55:05 -0500 Subject: [PATCH 09/12] Fix doubled parameter. --- seirsplus/models/seirs_model.py | 2 +- seirsplus/models/seirs_network_model.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/seirsplus/models/seirs_model.py b/seirsplus/models/seirs_model.py index fc2c9ad..5395e6c 100644 --- a/seirsplus/models/seirs_model.py +++ b/seirsplus/models/seirs_model.py @@ -382,7 +382,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin 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')): + 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'): diff --git a/seirsplus/models/seirs_network_model.py b/seirsplus/models/seirs_network_model.py index 9495daa..4c1a03d 100644 --- a/seirsplus/models/seirs_network_model.py +++ b/seirsplus/models/seirs_network_model.py @@ -1014,7 +1014,7 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='lin 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')): + 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'): From 4e5032814605c936263fe9338607062fe4842bc0 Mon Sep 17 00:00:00 2001 From: Noel Schutt <1065410+schutt@users.noreply.github.com> Date: Sat, 13 Mar 2021 18:43:45 -0500 Subject: [PATCH 10/12] Split the plotting out of the basic SEIRS and SEIRS network models into a separate base class. The plotting was slightly altered by adding a property to select between N and number of nodes as the basis of the plotting. This makes the plotting methods identical. --- seirsplus/models/base_plotable_model.py | 262 ++++++++++++++++++++++++ seirsplus/models/seirs_model.py | 250 +--------------------- seirsplus/models/seirs_network_model.py | 250 +--------------------- 3 files changed, 276 insertions(+), 486 deletions(-) create mode 100644 seirsplus/models/base_plotable_model.py 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/seirs_model.py b/seirsplus/models/seirs_model.py index 5395e6c..baf6383 100644 --- a/seirsplus/models/seirs_model.py +++ b/seirsplus/models/seirs_model.py @@ -5,8 +5,10 @@ import numpy import scipy.integrate +from .base_plotable_model import BasePlotableModel -class SEIRSModel(): + +class SEIRSModel(BasePlotableModel): """ A class to simulate the Deterministic SEIRS Model =================================================== @@ -37,6 +39,10 @@ class SEIRSModel(): (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, @@ -260,245 +266,3 @@ def total_num_recovered(self, t_idx=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_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/seirs_network_model.py b/seirsplus/models/seirs_network_model.py index 4c1a03d..9297362 100644 --- a/seirsplus/models/seirs_network_model.py +++ b/seirsplus/models/seirs_network_model.py @@ -6,8 +6,10 @@ import numpy import scipy +from .base_plotable_model import BasePlotableModel -class SEIRSNetworkModel(): + +class SEIRSNetworkModel(BasePlotableModel): """ A class to simulate the SEIRS Stochastic Network Model ====================================================== @@ -49,6 +51,10 @@ class SEIRSNetworkModel(): 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, @@ -892,245 +898,3 @@ def run(self, T, checkpoints=None, print_interval=10, verbose='t'): 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_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.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_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 From 31f8574e92482a4a4974e05dbec23e3e0f7e58d0 Mon Sep 17 00:00:00 2001 From: Noel Schutt <1065410+schutt@users.noreply.github.com> Date: Sat, 13 Mar 2021 19:50:23 -0500 Subject: [PATCH 11/12] Corrected timer state in the network model. --- seirsplus/models/seirs_network_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/seirsplus/models/seirs_network_model.py b/seirsplus/models/seirs_network_model.py index 9297362..90a6a9f 100644 --- a/seirsplus/models/seirs_network_model.py +++ b/seirsplus/models/seirs_network_model.py @@ -737,7 +737,7 @@ def run_iteration(self): self.testedInCurrentState[transitionNode] = False - self.timer_state[transitionNode] = 0.0 + self.timer_state = 0.0 #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ From 1deb63b9a8e52733c7dba5803088a62bbee55de3 Mon Sep 17 00:00:00 2001 From: Noel Schutt <1065410+schutt@users.noreply.github.com> Date: Sat, 13 Mar 2021 20:15:33 -0500 Subject: [PATCH 12/12] Update the extended SEIRS network model to extend the base model. Note that because of how the extra parameters are handled, this doesn't actually abstract anything. It is done for consistency and to simplify future unification efforts. --- .../models/extended_seirs_network_model.py | 42 +++++++++++-------- 1 file changed, 24 insertions(+), 18 deletions(-) diff --git a/seirsplus/models/extended_seirs_network_model.py b/seirsplus/models/extended_seirs_network_model.py index b804912..26e1a36 100644 --- a/seirsplus/models/extended_seirs_network_model.py +++ b/seirsplus/models/extended_seirs_network_model.py @@ -6,8 +6,10 @@ import numpy import scipy +from .base_plotable_model import BasePlotableModel -class ExtSEIRSNetworkModel(): + +class ExtSEIRSNetworkModel(BasePlotableModel): """ A class to simulate the Extended SEIRS Stochastic Network Model =================================================== @@ -70,6 +72,10 @@ class ExtSEIRSNetworkModel(): 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, @@ -1209,32 +1215,32 @@ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_ #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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 + 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(self.numNodes/100)] - dashedReference_infectedStack = dashed_reference_results.total_num_infected()[::int(self.numNodes/100)] / (self.numNodes if plot_percentages else 1) + 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() / (self.numNodes if plot_percentages else 1) + 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)