././@PaxHeader0000000000000000000000000000003400000000000011452 xustar000000000000000028 mtime=1597645612.5713499 seirsplus-1.0.9/0000755000076500000240000000000000000000000013647 5ustar00ryanstaff00000000000000././@PaxHeader0000000000000000000000000000003400000000000011452 xustar000000000000000028 mtime=1597645612.5710325 seirsplus-1.0.9/PKG-INFO0000644000076500000240000000057200000000000014750 0ustar00ryanstaff00000000000000Metadata-Version: 1.0 Name: seirsplus Version: 1.0.9 Summary: Models of SEIRS epidemic dynamics with extensions, including network-structured populations, testing, contact tracing, and social distancing. Home-page: https://github.com/ryansmcgee/SEIRS-network-model Author: Ryan Seamus McGee Author-email: ryansmcgee@gmail.com License: MIT Description: UNKNOWN Platform: UNKNOWN ././@PaxHeader0000000000000000000000000000002600000000000011453 xustar000000000000000022 mtime=1597020232.0 seirsplus-1.0.9/README.md0000644000076500000240000002752100000000000015135 0ustar00ryanstaff00000000000000# SEIRS+ Model Framework This package implements models of generalized SEIRS infectious disease dynamics with extensions that allow us to study the effect of social contact network structures, heterogeneities, stochasticity, and interventions, such as social distancing, testing, contact tracing, and isolation. #### *Latest Release: v1.0 (9 Aug 2020)* * Full rollout of [Extended SEIRS Model](https://github.com/ryansmcgee/seirsplus/wiki/Extended-SEIRS-model-description), [TTI simulation loop](https://github.com/ryansmcgee/seirsplus/wiki/TTI-Simulation-loop), and [network generators](https://github.com/ryansmcgee/seirsplus/wiki/network-generation) for [community](https://github.com/ryansmcgee/seirsplus/wiki/network-generation#demographic-community-network) and [workplace](https://github.com/ryansmcgee/seirsplus/wiki/network-generation#workplace-network) populations. * Example notebooks with in-depth walkthroughs of sophisticated [Community TTI](https://github.com/ryansmcgee/seirsplus/blob/master/examples/Extended_SEIRS_Community_TTI_Demo.ipynb) and [Workplace TTI](https://github.com/ryansmcgee/seirsplus/blob/master/examples/Extended_SEIRS_Workplace_TTI_Demo.ipynb) simulations. * The [`SEIRSModel`](https://github.com/ryansmcgee/seirsplus/wiki/SEIRSModel-class) and [`SEIRSNetworkModel`](https://github.com/ryansmcgee/seirsplus/wiki/SEIRSNetworkModel-class) classes have been refactored to be in line with [Extended SEIRS Model](https://github.com/ryansmcgee/seirsplus/wiki/Extended-SEIRS-model-description) conventions. The dynamics and core implementations remain the same, but some variable names have changed, etc. * The versions of these classes from before 9 Aug 2020 are now located in the `legacy_models.py` module. * Launch of new [wiki site](https://github.com/ryansmcgee/seirsplus/wiki) for thorough documentation of models and package features ## Overview #### Full documentation of this package's models, code, use cases, examples, and more can be found on [the wiki](https://github.com/ryansmcgee/seirsplus/wiki/) [**Basic SEIRS Model**](https://github.com/ryansmcgee/seirsplus/wiki/SEIRS-Model-Description) | [**Extended SEIRS Model**](https://github.com/ryansmcgee/seirsplus/wiki/Extended-SEIRS-Model-Description) :-----:|:-----: | [Model Description](https://github.com/ryansmcgee/seirsplus/wiki/SEIRS-Model-Description) | [Model Description](https://github.com/ryansmcgee/seirsplus/wiki/Extended-SEIRS-Model-Description) [`SEIRSNetworkModel` docs](https://github.com/ryansmcgee/seirsplus/wiki/SEIRSModel-class)
[`SEIRSModel` docs](https://github.com/ryansmcgee/seirsplus/wiki/SEIRSNetworkModel-class) | [`ExtSEIRSNetworkModel` docs](https://github.com/ryansmcgee/seirsplus/wiki/ExtSEIRSNetworkModel-class) [Basic SEIRS Mean-field Model Demo](https://github.com/ryansmcgee/seirsplus/blob/master/examples/Basic_SEIRS_Meanfield_Model_Demo.ipynb)
[Basic SEIRS Network Model Demo](https://github.com/ryansmcgee/seirsplus/blob/master/examples/Basic_SEIRS_Network_Model_Demo.ipynb) | [Extended SEIRS Community TTI Demo](https://github.com/ryansmcgee/seirsplus/blob/master/examples/Extended_SEIRS_Community_TTI_Demo.ipynb)
[Extended SEIRS Workplace TTI Demo](https://github.com/ryansmcgee/seirsplus/blob/master/examples/Extended_SEIRS_Workplace_TTI_Demo.ipynb) # ### SEIRS Dynamics The foundation of the models in this package is the classic SEIR model of infectious disease. The SEIR model is a standard compartmental model in which the population is divided into **susceptible (S)**, **exposed (E)**, **infectious (I)**, and **recovered (R)** individuals. A susceptible member of the population becomes infected (exposed) when making a transmissive contact with an infectious individual and then progresses to the infectious and finally recovered states. In the SEIRS model, recovered individuals may become resusceptible some time after recovering (although re-susceptibility can be excluded if not applicable or desired).

### Extended SEIRS Model This model extends the classic SEIRS model of infectious disease to represent pre-symptomatic, asymptomatic, and severely symptomatic disease states, which are of particular **relevance to the SARS-CoV-2 pandemic**. In this extended model, the infectious subpopulation is subdivided into **pre-symptomatic (*Ipre*)**, **asymptomatic (*Iasym*)**, **symptomatic (*Isym*)**, and **hospitalized (severely symptomatic, *IH*)**. All of these *I* compartments represent contagious individuals, but transmissibility, rates of recovery, and other parameters may vary between these disease states.

### Testing, Tracing, & Isolation The effect of isolation-based interventions (e.g., isolating individuals in response to testing or contact tracing) are modeled by introducing compartments representing quarantined individuals. An individual may be quarantined in any disease state, and every disease state has a corresponding quarantine compartment (with the exception of the hospitalized state, which is considered a quarantine state for transmission and other purposes). Quarantined individuals follow the same progression through the disease states, but the rates of transition or other parameters may be different. There are multiple methods by which individuals can be moved into or out of a quarantine state in this framework.

### Network Models Standard compartment models capture important features of infectious disease dynamics, but they are deterministic mean-field models that assume uniform mixing of the population (i.e., every individual in the population is equally likely to interact with every other individual). However, it is often important to consider stochasticity, heterogeneity, and the structure of contact networks when studying disease transmission, and many strategies for mitigating spread can be thought of as perturbing the contact network (e.g., social distancing) or making use of it (e.g., contact tracing). This package includes implementation of SEIRS models on stochastic dynamical networks. Individuals are represented as nodes in a network, and parameters, contacts, and interventions can be specified on a targeted individual basis. The network model enables rigorous analysis of transmission patterns and network-based interventions with respect to the properties of realistic contact networks. These SEIRS models can be simulated on any network. Network generation is largely left to the user, but some tools for [Network Generation](https://github.com/ryansmcgee/seirsplus/wiki/network-generation) are included in this package. Unlike mean-field compartment models, which do not model individual members of the population, the network model explicitly represents individuals as discrete nodes. All model parameters can be assigned to each node on an individual basis. Therefore, the network models support arbitrary parameter heterogeneity at the user's discretion. In addition, the specification of the contact network allows for heterogeneity in interaction patterns to be explicitly modeled as well. ## Code Usage This package was designed with broad usability in mind. Complex scenarios can be simulated with very few lines of code or, in many cases, no new coding or knowledge of Python by simply modifying the parameter values in the [example notebooks](https://github.com/ryansmcgee/seirsplus/tree/master/examples) provided. See the Quick Start section and the rest of the wiki documentation for more details. Don't be intimidated by the length of the wiki pages, running these models is quick and easy. The package does all the hard work for you. As an example, here's a complete script that simulates the SEIRS dynamics on a network with forms of social distancing, testing, contact tracing, and quarantining in only 10 lines of code.: ```python from seirsplus.models import * import networkx numNodes = 10000 baseGraph = networkx.barabasi_albert_graph(n=numNodes, m=9) G_normal = custom_exponential_graph(baseGraph, scale=100) # Social distancing interactions: G_distancing = custom_exponential_graph(baseGraph, scale=10) # Quarantine interactions: G_quarantine = custom_exponential_graph(baseGraph, scale=5) model = SEIRSNetworkModel(G=G_normal, beta=0.155, sigma=1/5.2, gamma=1/12.39, mu_I=0.0004, p=0.5, Q=G_quarantine, beta_D=0.155, sigma_D=1/5.2, gamma_D=1/12.39, mu_D=0.0004, theta_E=0.02, theta_I=0.02, phi_E=0.2, phi_I=0.2, psi_E=1.0, psi_I=1.0, q=0.5, initI=10) checkpoints = {'t': [20, 100], 'G': [G_distancing, G_normal], 'p': [0.1, 0.5], 'theta_E': [0.02, 0.02], 'theta_I': [0.02, 0.02], 'phi_E': [0.2, 0.2], 'phi_I': [0.2, 0.2]} model.run(T=300, checkpoints=checkpoints) model.figure_infections() ``` ### Quick Start Perhaps the best way to get started with these models is to dive into the [examples](https://github.com/ryansmcgee/seirsplus/tree/master/examples). These example notebooks walk through full simulations using each of the models included in this package with description of the steps involved. **These notebooks can also serve as ready-to-run sandboxes for trying your own simulation scenarios simply by changing the parameter values in the notebook.** ### Installing and Importing the Package All of the code needed to run the models is imported from the ```models``` module of this package. Other features that may be of interest are implemented in the `networks`, `sim_loops`, and `utilities` modules. #### Install the package using ```pip``` The package can be installed on your machine by entering this in the command line: ```> pip install seirsplus``` Then, the ```models``` module (and other modules) can be imported into your scripts as shown here: ```python from seirsplus.models import * from seirsplus.networks import * from seirsplus.sim_loops import * from seirsplus.utilities import * import networkx ``` #### *Alternatively, manually copy the code to your machine* *You can use the model code without installing a package by copying the ```models.py``` module file to a directory on your machine. For some of the features included in this package you will also need the `networks`, `sim_loops`, and `utilities` modules. In this case, the easiest way to use the modules is to place your scripts in the same directory as the modules, and import the modules as shown here:* ```python from models import * from networks import * from sim_loops import * from utilities import * ``` ### Full Code Documentation Complete documentation for all package classes and functions can be found throughout this wiki, including in-depth descriptions of concept, parameters, and how to initialize, run, and interface with the models. Some pages of note: * [`SEIRSModel`](https://github.com/ryansmcgee/seirsplus/wiki/SEIRSModel-class) * [`SEIRSNetworkModel`](https://github.com/ryansmcgee/seirsplus/wiki/SEIRSNetworkModel-class) * [`ExtSEIRSNetworkModel`](https://github.com/ryansmcgee/seirsplus/wiki/ExtSEIRSNetworkModel-class) * [Network Generation](https://github.com/ryansmcgee/seirsplus/wiki/Network-Generation) * [TTI Simulation Loop](https://github.com/ryansmcgee/seirsplus/wiki/TTI-Simulation-Loop) ././@PaxHeader0000000000000000000000000000003400000000000011452 xustar000000000000000028 mtime=1597645612.5676172 seirsplus-1.0.9/seirsplus/0000755000076500000240000000000000000000000015700 5ustar00ryanstaff00000000000000././@PaxHeader0000000000000000000000000000002600000000000011453 xustar000000000000000022 mtime=1596753884.0 seirsplus-1.0.9/seirsplus/FARZ.py0000644000076500000240000005153000000000000017020 0ustar00ryanstaff00000000000000################################################# # Adapted from https://github.com/rabbanyk/FARZ ################################################# import random import bisect import math import os def random_choice(values, weights=None , size = 1, replace = True): if weights is None: i = int(random.random() * len(values)) else : total = 0 cum_weights = [] for w in weights: total += w cum_weights.append(total) x = random.random() * total i = bisect.bisect(cum_weights, x) if size <=1: if len(values)>i: return values[i] else: return None else: cval = [values[j] for j in range(len(values)) if replace or i!=j] 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 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)) if i in self.memberships: self.memberships[i].append((cluster_id,s)) else: self.memberships[i] =[(cluster_id,s)] def write_groups(self, path): with open(path, 'w') as f: for g in self.groups: for i,s in g: f.write(str(i) + ' ') f.write('\n') class Graph: def __init__(self,directed=False, weighted=False): self.n = 0 self.counter = 0 self.max_degree = 0 self.directed = directed self.weighted = weighted self.edge_list = [] self.edge_time = [] self.deg = [] self.neigh = [[]] 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 self.edge_list.append((u,v,w) if uself.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 def to_nx(self): import networkx as nx G=nx.Graph() for u,v, w in self.edge_list: 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() for i in range(self.n): # This line works for networkx 1.10: # G.add_node(i, {'c':str(sorted(C.memberships[i]))}) # This line works for networkx 2.2: G.add_node(i, c=str(sorted(C.memberships[i]))) # G.add_node(i, {'c':int(C.memberships[i][0][0])}) for i in range(len(self.edge_list)): # for u,v, w in self.edge_list: u,v, w = self.edge_list[i] 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 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) for c in C.groups: for i,_ in c: for j,_ in c: q+= G.weight(i,j) - (G.deg[i]*G.deg[j]/(2*m)) q /= 2*m return q 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) else: p[j]= (wik * wjk) if len(p)<=0 or not normalize: return p maxp = p[max(p, key = lambda i: p[i])] for j in p: p[j] = p[j]*1.0 / maxp return p 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: 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 def degree_similarity(i, ids, G, gamma, normalize = True): p = [0]*len(ids) for ij,j in enumerate(ids): p[ij] = (G.deg[j] -G.deg[i])**2 if len(p)<=0 or not normalize: return p maxp = max(p) if maxp==0: return p p = [pi*1.0/maxp if gamma<0 else 1-pi*1.0/maxp for pi in p] return p def combine (a,b,alpha,gamma): return (a**alpha) / ((b+1)**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) 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) if c is None: return #Choose node within community tmp = choose_node(i, c, G, C, alpha, beta, gamma, epsilon) if tmp is None: return 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': 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 elif method == 'younger_less_active' : p = [G.n-i for i in range(G.n)] # younger less active else: p = [1 for i in range(G.n)] # uniform return random_choice(range(len(p)), p ) #, size=1, replace = False)[0] def assign(i, C, e=1, r=1, q = 0.5): p = [e +len(c) for c in C.groups] 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 '') # 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 ) # print(numpy.asarray((unique, counts)).T) # 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 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 '') write_to_file(G,C,path,name,format,farz_params) import sys def main(argv): import getopt FARZsetting = default_FARZ_setting.copy() batch_setting= default_batch_setting.copy() 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: print('invalid command, try -h to see usage and options') sys.exit(2) for opt, arg in opts: if opt == '-h': print('*** examples:') print('+ example 1: generate a network with 1000 nodes and about 5x1000 edges (m=5), with 4 communities, where 90% of edges fall within communities (beta=0.9)') print('> python FARZ.py -n 1000 -m 5 -k 4 --beta 0.9\n') print('+ example 2: generate a network with properties of example 1, where alpha = 0.2 and gamma = -0.8') print('> python FARZ.py -n 1000 -m 5 -k 4 --beta 0.9 --alpha 0.2 --gamma -0.8 \n') print('+ example 3: generate 10 sample networks with properties of example 1 and save them into ./data') print('> python FARZ.py --path ./data -s 10 -n 1000 -m 5 -k 4 --beta 0.9\n') print('+ example 4: repeat example 2, for beta that varies from 0.5 to 1 with 0.05 increments') 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)') print('-k: number of communities, default (4)') print('-b [or --beta]: the strength of community structure, i.e. the probability of edges to be formed within communities, default (0.8)') print('-a [or --alpha]: the strength of common neighbor\'s effect on edge formation edges, default (0.5)') print('-g [or --gamma]: the strength of degree similarity effect on edge formation, default (0.5), can be negative for networks with negative degree correlation') print('-p [or --phi]: the constant added to all community sizes, higher number makes the communities more balanced in size, default (1), which results in power law distribution for community sizes') print('-r: the number of communities each node can belong to, default (1)' ) print('-q: the probability of a node belonging to the multiple communities, default (0.5)' ) print('-e [or --epsilon]: the probability of noisy/random edges, default (0.0000001)') print('-t: the probability of also connecting to the neighbors of a node each nodes connects to. The default value is (0), but could be increased to a small number to achieve higher clustering coefficient. \n') print('*** batch parameters:') print('-s: the number of networks to be sampled with the given properties, default (1)' ) print('-o: the name of the output network, default (network)') print('--path : the path to write the network(s) to, default (.)') print('-f [or --format]: the format of output, list or gml, default (gml)') print('-v: the parameter to vary and sample networks for, default (None)') 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"): batch_setting['path'] = arg elif opt in ("-f", "--format"): if arg in supported_formats: batch_setting['format'] = arg else: print('Format not supported , choose from ',supported_formats,' or try -h to see the usage and options') sys.exit(2) elif opt in ("-s","--repeat"): try: batch_setting['repeat'] = int(arg) except ValueError: print('Invalid Number , try -h to see the usage and options') sys.exit(2) elif opt in ("-v", "--vary"): if (arg in list(default_ranges.keys())): batch_setting['vari'] = arg else: print('Invalid variable, choose form :', list(default_ranges.keys()), ', try -h to see the usage and options') sys.exit(2) elif opt in ("-c", "--range"): try: arange = [float(s) for s in arg[1:-1].split(',')] batch_setting['arange'] = arange except Error: print('Invalid range, should have the following form : [start,end,incrementBy], try -h to see the usage and options ') sys.exit(2) elif opt in ("-n"): try: FARZsetting['n'] = int(arg) except ValueError: print('Invalid Number , try -h to see the usage and options') sys.exit(2) elif opt in ("-k"): try: FARZsetting['k'] = int(arg) except ValueError: print('Invalid Number , try -h to see the usage and options') sys.exit(2) elif opt in ("-m"): try: FARZsetting['m'] = int(arg) except ValueError: print('Invalid Number , try -h to see the usage and options') sys.exit(2) elif opt in ("-a","--alpha"): try: FARZsetting['alpha'] = float(arg) except ValueError: print('Invalid Number , try -h to see the usage and options') 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) 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) 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) 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) elif opt in ("-q","--oProb"): try: FARZsetting['q'] = float(arg) except ValueError: print('Invalid Number , try -h to see the usage and options') sys.exit(2) elif opt in ("-d","--directed"): FARZsetting['directed'] = True elif opt in ("-w","--wighted"): 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) elif opt in ("-e","--epsilon"): 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, # "alpha":0.5, # "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././@PaxHeader0000000000000000000000000000002600000000000011453 xustar000000000000000022 mtime=1584956330.0 seirsplus-1.0.9/seirsplus/__init__.py0000644000076500000240000000000000000000000017777 0ustar00ryanstaff00000000000000././@PaxHeader0000000000000000000000000000002600000000000011453 xustar000000000000000022 mtime=1597020335.0 seirsplus-1.0.9/seirsplus/legacy_models.py0000644000076500000240000024750100000000000021072 0ustar00ryanstaff00000000000000from __future__ import absolute_import from __future__ import division from __future__ import print_function import networkx as networkx import numpy as numpy import scipy as scipy import scipy.integrate ######################################################## #@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@# #@ @# #@ [LEGACY] 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_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 mu_D 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 initD_E Init number of detected infectious 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) """ 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, 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 # Testing-related parameters: self.beta_D = beta_D if beta_D is not None else self.beta self.sigma_D = sigma_D if sigma_D is not None else self.sigma self.gamma_D = gamma_D if gamma_D is not None else self.gamma self.mu_D = mu_D if mu_D 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.numD_E = numpy.array([int(initD_E)]) self.numD_I = numpy.array([int(initD_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.numD_E[-1] - self.numD_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_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 N = S + E + I + D_E + D_I + R dS = - (beta*S*I)/N - q*(beta_D*S*D_I)/N + xi*R + nu*N - mu_0*S dE = (beta*S*I)/N + q*(beta_D*S*D_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_D*D_E - mu_0*D_E dDI = theta_I*psi_I*I + sigma_D*D_E - gamma_D*D_I - mu_D*D_I - mu_0*D_I dR = gamma*I + gamma_D*D_I - xi*R - mu_0*R dF = mu_I*I + mu_D*D_I return [dS, dE, dI, dDE, dDI, dR, dF] #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def run_epoch(self, runtime, dt=0.1): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Create a list of times at which the ODE solver should output system values. # Append this list of times as the model's timeseries t_eval = numpy.arange(start=self.t, stop=self.t+runtime, step=dt) # Define the range of time values for the integration: t_span = (self.t, self.t+runtime) # Define the initial conditions as the system's current state: # (which will be the t=0 condition if this is the first run of this model, # else where the last sim left off) init_cond = [self.numS[-1], self.numE[-1], self.numI[-1], self.numD_E[-1], self.numD_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_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 ) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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.numD_E = numpy.append(self.numD_E, solution['y'][3]) self.numD_I = numpy.append(self.numD_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_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), # 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 D_E = " + str(self.numD_E[-1])) 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']): # 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 D_E = " + str(self.numD_E[-1])) print("\t D_I = " + str(self.numD_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_infections(self, t_idx=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]) #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='line', plot_F='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', 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.numD_E+self.numD_I)/self.N if plot_percentages else (self.numD_E+self.numD_I) D_Eseries = self.numD_E/self.N if plot_percentages else self.numD_E 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 #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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.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): 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) 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_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'): 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'): 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'): 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_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'): 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'): 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'): 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_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'): 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'): 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'): 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_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', 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_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, 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_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', 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_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, 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 (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 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) sigma_D Rate of infection (upon exposure) for individuals with detected infections gamma_D Rate of recovery (upon infection) for individuals with detected infections mu_D 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 phi_E Rate of contact tracing testing for exposed individuals phi_I Rate of contact tracing 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 interaction outside adjacent nodes 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 initF Init number of infection-related fatalities (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, Q=None, beta_D=None, sigma_D=None, gamma_D=None, mu_D=None, beta_D_local=None, theta_E=0, theta_I=0, phi_E=0, phi_I=0, psi_E=1, psi_I=1, q=0, initE=0, initI=10, initD_E=0, initD_I=0, initR=0, initF=0, node_groups=None, store_Xseries=False): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Setup Adjacency matrix: self.update_G(G) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Setup Quarantine Adjacency matrix: if(Q is None): Q = G # If no Q graph is provided, use G in its place self.update_Q(Q) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Model Parameters: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ self.parameters = { 'beta':beta, 'sigma':sigma, 'gamma':gamma, 'xi':xi, 'mu_I':mu_I, 'mu_0':mu_0, 'nu':nu, 'beta_D':beta_D, 'sigma_D':sigma_D, 'gamma_D':gamma_D, 'mu_D':mu_D, '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 # (will be expanded during run if needed) self.tseries = numpy.zeros(5*self.numNodes) self.numE = numpy.zeros(5*self.numNodes) self.numI = numpy.zeros(5*self.numNodes) self.numD_E = numpy.zeros(5*self.numNodes) self.numD_I = numpy.zeros(5*self.numNodes) self.numR = numpy.zeros(5*self.numNodes) self.numF = numpy.zeros(5*self.numNodes) self.numS = numpy.zeros(5*self.numNodes) self.N = numpy.zeros(5*self.numNodes) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Initialize Timekeeping: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ self.t = 0 self.tmax = 0 # will be set when run() is called self.tidx = 0 self.tseries[0] = 0 #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Initialize Counts of inidividuals with each state: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ self.numE[0] = int(initE) self.numI[0] = int(initI) self.numD_E[0] = int(initD_E) self.numD_I[0] = int(initD_I) self.numR[0] = int(initR) 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: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ self.S = 1 self.E = 2 self.I = 3 self.D_E = 4 self.D_I = 5 self.R = 6 self.F = 7 self.X = numpy.array([self.S]*int(self.numS[0]) + [self.E]*int(self.numE[0]) + [self.I]*int(self.numI[0]) + [self.D_E]*int(self.numD_E[0]) + [self.D_I]*int(self.numD_I[0]) + [self.R]*int(self.numR[0]) + [self.F]*int(self.numF[0])).reshape((self.numNodes,1)) numpy.random.shuffle(self.X) self.store_Xseries = store_Xseries if(store_Xseries): self.Xseries = numpy.zeros(shape=(5*self.numNodes, self.numNodes), dtype='uint8') self.Xseries[0,:] = self.X.T self.transitions = { '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}, 'EtoDE': {'currentState':self.E, 'newState':self.D_E}, 'ItoDI': {'currentState':self.I, 'newState':self.D_I}, 'DEtoDI': {'currentState':self.D_E, 'newState':self.D_I}, 'DItoR': {'currentState':self.D_I, 'newState':self.R}, 'DItoF': {'currentState':self.D_I, 'newState':self.F}, '_toS': {'currentState':True, 'newState':self.S}, } #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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(5*self.numNodes) self.nodeGroupData[groupName]['numE'] = numpy.zeros(5*self.numNodes) self.nodeGroupData[groupName]['numI'] = numpy.zeros(5*self.numNodes) self.nodeGroupData[groupName]['numD_E'] = numpy.zeros(5*self.numNodes) self.nodeGroupData[groupName]['numD_I'] = numpy.zeros(5*self.numNodes) self.nodeGroupData[groupName]['numR'] = numpy.zeros(5*self.numNodes) self.nodeGroupData[groupName]['numF'] = numpy.zeros(5*self.numNodes) self.nodeGroupData[groupName]['N'] = numpy.zeros(5*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]['numD_E'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.D_E) self.nodeGroupData[groupName]['numD_I'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.D_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]['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] #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def update_parameters(self): import time updatestart = time.time() #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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.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_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.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 self.gamma_D = (numpy.array(self.parameters['gamma_D']).reshape((self.numNodes, 1)) if isinstance(self.parameters['gamma_D'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['gamma_D'], shape=(self.numNodes,1))) if self.parameters['gamma_D'] is not None else self.gamma self.mu_D = (numpy.array(self.parameters['mu_D']).reshape((self.numNodes, 1)) if isinstance(self.parameters['mu_D'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['mu_D'], shape=(self.numNodes,1))) if self.parameters['mu_D'] is not None else self.mu_I 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)) #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)): 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): self.beta_local.reshape((self.numNodes, 1)) 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)): 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): self.beta_D_local.reshape((self.numNodes, 1)) elif(self.beta_D_local.ndim == 2): self.beta_D_local.reshape((self.numNodes, self.numNodes)) else: self.beta_D_local = numpy.full_like(self.beta_D, fill_value=self.parameters['beta_D_local']) else: self.beta_D_local = self.beta_D # Pre-multiply beta values by the adjacency matrix ("transmission weight connections") if(self.beta_local.ndim == 1): 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): 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): 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): self.A_Q_beta_D = scipy.sparse.csr_matrix.multiply(self.A_Q, self.beta_D_local).tocsr() #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Update scenario flags: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ self.update_scenario_flags() #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def node_degrees(self, Amat): return Amat.sum(axis=0).reshape(self.numNodes,1) # sums of adj matrix cols #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def update_G(self, new_G): self.G = new_G # Adjacency matrix: if type(new_G)==numpy.ndarray: self.A = scipy.sparse.csr_matrix(new_G) elif type(new_G)==networkx.classes.graph.Graph: self.A = networkx.adj_matrix(new_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) return #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def update_Q(self, new_Q): self.Q = new_Q # Quarantine Adjacency matrix: if type(new_Q)==numpy.ndarray: self.A_Q = scipy.sparse.csr_matrix(new_Q) elif type(new_Q)==networkx.classes.graph.Graph: self.A_Q = networkx.adj_matrix(new_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." return #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 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))) 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)) 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)) #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_infections(self, t_idx=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]) #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ 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]) 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 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 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)) ) ) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ propensities_StoE = ( self.p*((self.beta*self.numI[self.tidx] + self.q*self.beta_D*self.numD_I[self.tidx])/self.N[self.tidx]) + (1-self.p)*numpy.divide((transmissionTerms_I + transmissionTerms_DI), self.degree, out=numpy.zeros_like(self.degree), where=self.degree!=0) )*(self.X==self.S) propensities_EtoI = self.sigma*(self.X==self.E) propensities_ItoR = self.gamma*(self.X==self.I) propensities_ItoF = self.mu_I*(self.X==self.I) # propensities_EtoDE = ( self.theta_E + numpy.divide((self.phi_E*numContacts_D), self.degree, out=numpy.zeros_like(self.degree), where=self.degree!=0) )*self.psi_E*(self.X==self.E) propensities_EtoDE = (self.theta_E + self.phi_E*numContacts_D)*self.psi_E*(self.X==self.E) # propensities_ItoDI = ( self.theta_I + numpy.divide((self.phi_I*numContacts_D), self.degree, out=numpy.zeros_like(self.degree), where=self.degree!=0) )*self.psi_I*(self.X==self.I) propensities_ItoDI = (self.theta_I + self.phi_I*numContacts_D)*self.psi_I*(self.X==self.I) propensities_DEtoDI = self.sigma_D*(self.X==self.D_E) propensities_DItoR = self.gamma_D*(self.X==self.D_I) propensities_DItoF = self.mu_D*(self.X==self.D_I) 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_EtoDE, propensities_ItoDI, propensities_DEtoDI, propensities_DItoR, propensities_DItoF, propensities_RtoS, propensities__toS]) columns = ['StoE', 'EtoI', 'ItoR', 'ItoF', 'EtoDE', 'ItoDI', 'DEtoDI', 'DItoR', 'DItoF', 'RtoS', '_toS'] return propensities, columns #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def increase_data_series_length(self): self.tseries= numpy.pad(self.tseries, [(0, 5*self.numNodes)], mode='constant', constant_values=0) self.numS = numpy.pad(self.numS, [(0, 5*self.numNodes)], mode='constant', constant_values=0) self.numE = numpy.pad(self.numE, [(0, 5*self.numNodes)], mode='constant', constant_values=0) self.numI = numpy.pad(self.numI, [(0, 5*self.numNodes)], mode='constant', constant_values=0) self.numD_E = numpy.pad(self.numD_E, [(0, 5*self.numNodes)], mode='constant', constant_values=0) self.numD_I = numpy.pad(self.numD_I, [(0, 5*self.numNodes)], mode='constant', constant_values=0) self.numR = numpy.pad(self.numR, [(0, 5*self.numNodes)], mode='constant', constant_values=0) 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): self.Xseries = numpy.pad(self.Xseries, [(0, 5*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, 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) self.nodeGroupData[groupName]['numI'] = numpy.pad(self.nodeGroupData[groupName]['numI'], [(0, 5*self.numNodes)], mode='constant', constant_values=0) self.nodeGroupData[groupName]['numD_E'] = numpy.pad(self.nodeGroupData[groupName]['numD_E'], [(0, 5*self.numNodes)], mode='constant', constant_values=0) self.nodeGroupData[groupName]['numD_I'] = numpy.pad(self.nodeGroupData[groupName]['numD_I'], [(0, 5*self.numNodes)], mode='constant', constant_values=0) self.nodeGroupData[groupName]['numR'] = numpy.pad(self.nodeGroupData[groupName]['numR'], [(0, 5*self.numNodes)], mode='constant', constant_values=0) self.nodeGroupData[groupName]['numF'] = numpy.pad(self.nodeGroupData[groupName]['numF'], [(0, 5*self.numNodes)], mode='constant', constant_values=0) self.nodeGroupData[groupName]['N'] = numpy.pad(self.nodeGroupData[groupName]['N'], [(0, 5*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.numD_E = numpy.array(self.numD_E, dtype=float)[:self.tidx+1] self.numD_I = numpy.array(self.numD_I, 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.N = numpy.array(self.N, 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]['numD_E'] = numpy.array(self.nodeGroupData[groupName]['numD_E'], dtype=float)[:self.tidx+1] self.nodeGroupData[groupName]['numD_I'] = numpy.array(self.nodeGroupData[groupName]['numD_I'], 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]['N'] = numpy.array(self.nodeGroupData[groupName]['N'], 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() #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 1. Generate 2 random numbers uniformly distributed in (0,1) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ r1 = numpy.random.rand() r2 = numpy.random.rand() #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 2. Calculate propensities #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ propensities, transitionTypes = self.calc_propensities() # Terminate when probability of all events is 0: if(propensities.sum() <= 0.0): self.finalize_data_series() return False #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 3. Calculate alpha #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ propensities_flat = propensities.ravel(order='F') cumsum = propensities_flat.cumsum() alpha = propensities_flat.sum() #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 4. Compute the time until the next event takes place #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ tau = (1/alpha)*numpy.log(float(1/r1)) self.t += tau #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 5. Compute which event takes place #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ transitionIdx = numpy.searchsorted(cumsum,r2*alpha) transitionNode = transitionIdx % self.numNodes transitionType = transitionTypes[ int(transitionIdx/self.numNodes) ] #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 6. Update node states and data series #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ 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.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.numD_E[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.D_E), a_min=0, a_max=self.numNodes) self.numD_I[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.D_I), a_min=0, a_max=self.numNodes) self.numR[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.R), a_min=0, a_max=self.numNodes) self.numF[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.F), a_min=0, a_max=self.numNodes) self.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): 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]['numD_E'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.D_E) self.nodeGroupData[groupName]['numD_I'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.D_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]['N'][self.tidx] = numpy.clip((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]), a_min=0, a_max=self.numNodes) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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)): 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: if('G' in list(checkpoints.keys())): self.update_G(checkpoints['G'][checkpointIdx]) 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())): 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 D_E = " + str(self.numD_E[self.tidx])) print("\t D_I = " + str(self.numD_I[self.tidx])) 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)): 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_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', 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.numD_E+self.numD_I)/self.numNodes if plot_percentages else (self.numD_E+self.numD_I) D_Eseries = self.numD_E/self.numNodes if plot_percentages else self.numD_E 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 #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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.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): 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) 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_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'): 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'): 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'): 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_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'): 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'): 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'): 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_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'): 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'): 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'): 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_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', 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_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, 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_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', 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_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, ylim=ylim, xlim=xlim, legend=legend, title=title, side_title=side_title, plot_percentages=plot_percentages) if(show): pyplot.show() return fig, ax #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%././@PaxHeader0000000000000000000000000000002600000000000011453 xustar000000000000000022 mtime=1597626994.0 seirsplus-1.0.9/seirsplus/models.py0000644000076500000240000066774700000000000017570 0ustar00ryanstaff00000000000000from __future__ import absolute_import from __future__ import division from __future__ import print_function import networkx as networkx import numpy as numpy import scipy as scipy import scipy.integrate ######################################################## #@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@# #@ @# #@ BASIC SEIRS MODELS @# #@ @# #@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@# ######################################################## class SEIRSModel(): """ A class to simulate the Deterministic SEIRS Model =================================================== Params: beta Rate of transmission (exposure) sigma Rate of infection (upon exposure) gamma Rate of recovery (upon infection) xi Rate of re-susceptibility (upon recovery) mu_I Rate of infection-related death mu_0 Rate of baseline death nu Rate of baseline birth beta_Q Rate of transmission (exposure) for individuals with detected infections sigma_Q Rate of infection (upon exposure) for individuals with detected infections gamma_Q Rate of recovery (upon infection) for individuals with detected infections mu_Q Rate of infection-related death for individuals with detected infections theta_E Rate of baseline testing for exposed individuals theta_I Rate of baseline testing for infectious individuals psi_E Probability of positive test results for exposed individuals psi_I Probability of positive test results for exposed individuals q Probability of quarantined individuals interacting with others initE Init number of exposed individuals initI Init number of infectious individuals initQ_E Init number of detected infectious individuals initQ_I Init number of detected infectious individuals initR Init number of recovered individuals initF Init number of infection-related fatalities (all remaining nodes initialized susceptible) """ def __init__(self, initN, beta, sigma, gamma, xi=0, mu_I=0, mu_0=0, nu=0, p=0, beta_Q=None, sigma_Q=None, gamma_Q=None, mu_Q=None, theta_E=0, theta_I=0, psi_E=0, psi_I=0, q=0, initE=0, initI=10, initQ_E=0, initQ_I=0, initR=0, initF=0): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Model Parameters: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ self.beta = beta self.sigma = sigma self.gamma = gamma self.xi = xi self.mu_I = mu_I self.mu_0 = mu_0 self.nu = nu self.p = p # Testing-related parameters: self.beta_Q = beta_Q if beta_Q is not None else self.beta self.sigma_Q = sigma_Q if sigma_Q is not None else self.sigma self.gamma_Q = gamma_Q if gamma_Q is not None else self.gamma self.mu_Q = mu_Q if mu_Q is not None else self.mu_I self.theta_E = theta_E if theta_E is not None else self.theta_E self.theta_I = theta_I if theta_I is not None else self.theta_I self.psi_E = psi_E if psi_E is not None else self.psi_E self.psi_I = psi_I if psi_I is not None else self.psi_I self.q = q if q is not None else self.q #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Initialize Timekeeping: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ self.t = 0 self.tmax = 0 # will be set when run() is called self.tseries = numpy.array([0]) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Initialize Counts of inidividuals with each state: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ self.N = numpy.array([int(initN)]) self.numE = numpy.array([int(initE)]) self.numI = numpy.array([int(initI)]) self.numQ_E = numpy.array([int(initQ_E)]) self.numQ_I = numpy.array([int(initQ_I)]) self.numR = numpy.array([int(initR)]) self.numF = numpy.array([int(initF)]) self.numS = numpy.array([self.N[-1] - self.numE[-1] - self.numI[-1] - self.numQ_E[-1] - self.numQ_I[-1] - self.numR[-1] - self.numF[-1]]) assert(self.numS[0] >= 0), "The specified initial population size N must be greater than or equal to the initial compartment counts." #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ @staticmethod def system_dfes(t, variables, beta, sigma, gamma, xi, mu_I, mu_0, nu, beta_Q, sigma_Q, gamma_Q, mu_Q, theta_E, theta_I, psi_E, psi_I, q): S, E, I, Q_E, Q_I, R, F = variables # varibles is a list with compartment counts as elements N = S + E + I + Q_E + Q_I + R dS = - (beta*S*I)/N - q*(beta_Q*S*Q_I)/N + xi*R + nu*N - mu_0*S dE = (beta*S*I)/N + q*(beta_Q*S*Q_I)/N - sigma*E - theta_E*psi_E*E - mu_0*E dI = sigma*E - gamma*I - mu_I*I - theta_I*psi_I*I - mu_0*I dDE = theta_E*psi_E*E - sigma_Q*Q_E - mu_0*Q_E dDI = theta_I*psi_I*I + sigma_Q*Q_E - gamma_Q*Q_I - mu_Q*Q_I - mu_0*Q_I dR = gamma*I + gamma_Q*Q_I - xi*R - mu_0*R dF = mu_I*I + mu_Q*Q_I return [dS, dE, dI, dDE, dDI, dR, dF] #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def run_epoch(self, runtime, dt=0.1): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Create a list of times at which the ODE solver should output system values. # Append this list of times as the model's timeseries t_eval = numpy.arange(start=self.t, stop=self.t+runtime, step=dt) # Define the range of time values for the integration: t_span = [self.t, self.t+runtime] # Define the initial conditions as the system's current state: # (which will be the t=0 condition if this is the first run of this model, # else where the last sim left off) init_cond = [self.numS[-1], self.numE[-1], self.numI[-1], self.numQ_E[-1], self.numQ_I[-1], self.numR[-1], self.numF[-1]] #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Solve the system of differential eqns: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ solution = scipy.integrate.solve_ivp(lambda t, X: SEIRSModel.system_dfes(t, X, self.beta, self.sigma, self.gamma, self.xi, self.mu_I, self.mu_0, self.nu, self.beta_Q, self.sigma_Q, self.gamma_Q, self.mu_Q, self.theta_E, self.theta_I, self.psi_E, self.psi_I, self.q ), t_span=t_span, y0=init_cond, t_eval=t_eval ) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Store the solution output as the model's time series and data series: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ self.tseries = numpy.append(self.tseries, solution['t']) self.numS = numpy.append(self.numS, solution['y'][0]) self.numE = numpy.append(self.numE, solution['y'][1]) self.numI = numpy.append(self.numI, solution['y'][2]) self.numQ_E = numpy.append(self.numQ_E, solution['y'][3]) self.numQ_I = numpy.append(self.numQ_I, solution['y'][4]) self.numR = numpy.append(self.numR, solution['y'][5]) self.numF = numpy.append(self.numF, solution['y'][6]) self.t = self.tseries[-1] #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def run(self, T, dt=0.1, checkpoints=None, verbose=False): if(T>0): self.tmax += T else: return False #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Pre-process checkpoint values: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(checkpoints): numCheckpoints = len(checkpoints['t']) paramNames = ['beta', 'sigma', 'gamma', 'xi', 'mu_I', 'mu_0', 'nu', 'beta_Q', 'sigma_Q', 'gamma_Q', 'mu_Q', 'theta_E', 'theta_I', 'psi_E', 'psi_I', 'q'] for param in paramNames: # For params that don't have given checkpoint values (or bad value given), # set their checkpoint values to the value they have now for all checkpoints. if(param not in list(checkpoints.keys()) or not isinstance(checkpoints[param], (list, numpy.ndarray)) or len(checkpoints[param])!=numCheckpoints): checkpoints[param] = [getattr(self, param)]*numCheckpoints #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% # Run the simulation loop: #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% if(not checkpoints): self.run_epoch(runtime=self.tmax, dt=dt) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ print("t = %.2f" % self.t) if(verbose): print("\t S = " + str(self.numS[-1])) print("\t E = " + str(self.numE[-1])) print("\t I = " + str(self.numI[-1])) print("\t Q_E = " + str(self.numQ_E[-1])) print("\t Q_I = " + str(self.numQ_I[-1])) print("\t R = " + str(self.numR[-1])) print("\t F = " + str(self.numF[-1])) else: # checkpoints provided for checkpointIdx, checkpointTime in enumerate(checkpoints['t']): # Run the sim until the next checkpoint time: self.run_epoch(runtime=checkpointTime-self.t, dt=dt) # Having reached the checkpoint, update applicable parameters: print("[Checkpoint: Updating parameters]") for param in paramNames: setattr(self, param, checkpoints[param][checkpointIdx]) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ print("t = %.2f" % self.t) if(verbose): print("\t S = " + str(self.numS[-1])) print("\t E = " + str(self.numE[-1])) print("\t I = " + str(self.numI[-1])) print("\t Q_E = " + str(self.numQ_E[-1])) print("\t Q_I = " + str(self.numQ_I[-1])) print("\t R = " + str(self.numR[-1])) print("\t F = " + str(self.numF[-1])) if(self.t < self.tmax): self.run_epoch(runtime=self.tmax-self.t, dt=dt) return True #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_susceptible(self, t_idx=None): if(t_idx is None): return (self.numS[:]) else: return (self.numS[t_idx]) #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_infected(self, t_idx=None): if(t_idx is None): return (self.numE[:] + self.numI[:] + self.numQ_E[:] + self.numQ_I[:]) else: return (self.numE[t_idx] + self.numI[t_idx] + self.numQ_E[t_idx] + self.numQ_I[t_idx]) #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_isolated(self, t_idx=None): if(t_idx is None): return (self.numQ_E[:] + self.numQ_I[:]) else: return (self.numQ_E[t_idx] + self.numQ_I[t_idx]) #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_recovered(self, t_idx=None): if(t_idx is None): return (self.numR[:]) else: return (self.numR[t_idx]) #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='line', plot_F='line', plot_Q_E='line', plot_Q_I='line', combine_Q=True, color_S='tab:green', color_E='orange', color_I='crimson', color_R='tab:blue', color_F='black', color_Q_E='mediumorchid', color_Q_I='mediumorchid', color_reference='#E0E0E0', dashed_reference_results=None, dashed_reference_label='reference', shaded_reference_results=None, shaded_reference_label='reference', vlines=[], vline_colors=[], vline_styles=[], vline_labels=[], ylim=None, xlim=None, legend=True, title=None, side_title=None, plot_percentages=True): import matplotlib.pyplot as pyplot #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Create an Axes object if None provided: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(not ax): fig, ax = pyplot.subplots() #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Prepare data series to be plotted: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Fseries = self.numF/self.N if plot_percentages else self.numF Eseries = self.numE/self.N if plot_percentages else self.numE Dseries = (self.numQ_E+self.numQ_I)/self.N if plot_percentages else (self.numQ_E+self.numQ_I) Q_Eseries = self.numQ_E/self.N if plot_percentages else self.numQ_E Q_Iseries = self.numQ_I/self.N if plot_percentages else self.numQ_I Iseries = self.numI/self.N if plot_percentages else self.numI Rseries = self.numR/self.N if plot_percentages else self.numR Sseries = self.numS/self.N if plot_percentages else self.numS #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Draw the reference data: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(dashed_reference_results): dashedReference_tseries = dashed_reference_results.tseries[::int(self.N/100)] dashedReference_IDEstack = (dashed_reference_results.numI + dashed_reference_results.numQ_I + dashed_reference_results.numQ_E + dashed_reference_results.numE)[::int(self.N/100)] / (self.N if plot_percentages else 1) ax.plot(dashedReference_tseries, dashedReference_IDEstack, color='#E0E0E0', linestyle='--', label='$I+D+E$ ('+dashed_reference_label+')', zorder=0) if(shaded_reference_results): shadedReference_tseries = shaded_reference_results.tseries shadedReference_IDEstack = (shaded_reference_results.numI + shaded_reference_results.numQ_I + shaded_reference_results.numQ_E + shaded_reference_results.numE) / (self.N if plot_percentages else 1) ax.fill_between(shaded_reference_results.tseries, shadedReference_IDEstack, 0, color='#EFEFEF', label='$I+D+E$ ('+shaded_reference_label+')', zorder=0) ax.plot(shaded_reference_results.tseries, shadedReference_IDEstack, color='#E0E0E0', zorder=1) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Draw the stacked variables: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ topstack = numpy.zeros_like(self.tseries) if(any(Fseries) and plot_F=='stacked'): ax.fill_between(numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, topstack+Fseries), topstack, color=color_F, alpha=0.5, label='$F$', zorder=2) ax.plot( numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, topstack+Fseries), color=color_F, zorder=3) topstack = topstack+Fseries if(any(Eseries) and plot_E=='stacked'): ax.fill_between(numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, topstack+Eseries), topstack, color=color_E, alpha=0.5, label='$E$', zorder=2) ax.plot( numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, topstack+Eseries), color=color_E, zorder=3) topstack = topstack+Eseries if(combine_Q and plot_Q_E=='stacked' and plot_Q_I=='stacked'): ax.fill_between(numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, topstack+Dseries), topstack, color=color_Q_E, alpha=0.5, label='$Q_{all}$', zorder=2) ax.plot( numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, topstack+Dseries), color=color_Q_E, zorder=3) topstack = topstack+Dseries else: if(any(Q_Eseries) and plot_Q_E=='stacked'): ax.fill_between(numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, topstack+Q_Eseries), topstack, color=color_Q_E, alpha=0.5, label='$Q_E$', zorder=2) ax.plot( numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, topstack+Q_Eseries), color=color_Q_E, zorder=3) topstack = topstack+Q_Eseries if(any(Q_Iseries) and plot_Q_I=='stacked'): ax.fill_between(numpy.ma.masked_where(Q_Iseries<=0, self.tseries), numpy.ma.masked_where(Q_Iseries<=0, topstack+Q_Iseries), topstack, color=color_Q_I, alpha=0.5, label='$Q_I$', zorder=2) ax.plot( numpy.ma.masked_where(Q_Iseries<=0, self.tseries), numpy.ma.masked_where(Q_Iseries<=0, topstack+Q_Iseries), color=color_Q_I, zorder=3) topstack = topstack+Q_Iseries if(any(Iseries) and plot_I=='stacked'): ax.fill_between(numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, topstack+Iseries), topstack, color=color_I, alpha=0.5, label='$I$', zorder=2) ax.plot( numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, topstack+Iseries), color=color_I, zorder=3) topstack = topstack+Iseries if(any(Rseries) and plot_R=='stacked'): ax.fill_between(numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, topstack+Rseries), topstack, color=color_R, alpha=0.5, label='$R$', zorder=2) ax.plot( numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, topstack+Rseries), color=color_R, zorder=3) topstack = topstack+Rseries if(any(Sseries) and plot_S=='stacked'): ax.fill_between(numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, topstack+Sseries), topstack, color=color_S, alpha=0.5, label='$S$', zorder=2) ax.plot( numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, topstack+Sseries), color=color_S, zorder=3) topstack = topstack+Sseries #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Draw the shaded variables: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(any(Fseries) and plot_F=='shaded'): ax.fill_between(numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, Fseries), 0, color=color_F, alpha=0.5, label='$F$', zorder=4) ax.plot( numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, Fseries), color=color_F, zorder=5) if(any(Eseries) and plot_E=='shaded'): ax.fill_between(numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, Eseries), 0, color=color_E, alpha=0.5, label='$E$', zorder=4) ax.plot( numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, Eseries), color=color_E, zorder=5) if(combine_Q and (any(Dseries) and plot_Q_E=='shaded' and plot_Q_E=='shaded')): ax.fill_between(numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, Dseries), 0, color=color_Q_E, alpha=0.5, label='$Q_{all}$', zorder=4) ax.plot( numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, Dseries), color=color_Q_E, zorder=5) else: if(any(Q_Eseries) and plot_Q_E=='shaded'): ax.fill_between(numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, Q_Eseries), 0, color=color_Q_E, alpha=0.5, label='$Q_E$', zorder=4) ax.plot( numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, Q_Eseries), color=color_Q_E, zorder=5) if(any(Q_Iseries) and plot_Q_I=='shaded'): ax.fill_between(numpy.ma.masked_where(Q_Iseries<=0, self.tseries), numpy.ma.masked_where(Q_Iseries<=0, Q_Iseries), 0, color=color_Q_I, alpha=0.5, label='$Q_I$', zorder=4) ax.plot( numpy.ma.masked_where(Q_Iseries<=0, self.tseries), numpy.ma.masked_where(Q_Iseries<=0, Q_Iseries), color=color_Q_I, zorder=5) if(any(Iseries) and plot_I=='shaded'): ax.fill_between(numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, Iseries), 0, color=color_I, alpha=0.5, label='$I$', zorder=4) ax.plot( numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, Iseries), color=color_I, zorder=5) if(any(Sseries) and plot_S=='shaded'): ax.fill_between(numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, Sseries), 0, color=color_S, alpha=0.5, label='$S$', zorder=4) ax.plot( numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, Sseries), color=color_S, zorder=5) if(any(Rseries) and plot_R=='shaded'): ax.fill_between(numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, Rseries), 0, color=color_R, alpha=0.5, label='$R$', zorder=4) ax.plot( numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, Rseries), color=color_R, zorder=5) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Draw the line variables: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(any(Fseries) and plot_F=='line'): ax.plot(numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, Fseries), color=color_F, label='$F$', zorder=6) if(any(Eseries) and plot_E=='line'): ax.plot(numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, Eseries), color=color_E, label='$E$', zorder=6) if(combine_Q and (any(Dseries) and plot_Q_E=='line' and plot_Q_E=='line')): ax.plot(numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, Dseries), color=color_Q_E, label='$Q_{all}$', zorder=6) else: if(any(Q_Eseries) and plot_Q_E=='line'): ax.plot(numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, Q_Eseries), color=color_Q_E, label='$Q_E$', zorder=6) if(any(Q_Iseries) and plot_Q_I=='line'): ax.plot(numpy.ma.masked_where(Q_Iseries<=0, self.tseries), numpy.ma.masked_where(Q_Iseries<=0, Q_Iseries), color=color_Q_I, label='$Q_I$', zorder=6) if(any(Iseries) and plot_I=='line'): ax.plot(numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, Iseries), color=color_I, label='$I$', zorder=6) if(any(Sseries) and plot_S=='line'): ax.plot(numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, Sseries), color=color_S, label='$S$', zorder=6) if(any(Rseries) and plot_R=='line'): ax.plot(numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, Rseries), color=color_R, label='$R$', zorder=6) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Draw the vertical line annotations: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(len(vlines)>0 and len(vline_colors)==0): vline_colors = ['gray']*len(vlines) if(len(vlines)>0 and len(vline_labels)==0): vline_labels = [None]*len(vlines) if(len(vlines)>0 and len(vline_styles)==0): vline_styles = [':']*len(vlines) for vline_x, vline_color, vline_style, vline_label in zip(vlines, vline_colors, vline_styles, vline_labels): if(vline_x is not None): ax.axvline(x=vline_x, color=vline_color, linestyle=vline_style, alpha=1, label=vline_label) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Draw the plot labels: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ ax.set_xlabel('days') ax.set_ylabel('percent of population' if plot_percentages else 'number of individuals') ax.set_xlim(0, (max(self.tseries) if not xlim else xlim)) ax.set_ylim(0, ylim) if(plot_percentages): ax.set_yticklabels(['{:,.0%}'.format(y) for y in ax.get_yticks()]) if(legend): legend_handles, legend_labels = ax.get_legend_handles_labels() ax.legend(legend_handles[::-1], legend_labels[::-1], loc='upper right', facecolor='white', edgecolor='none', framealpha=0.9, prop={'size': 8}) if(title): ax.set_title(title, size=12) if(side_title): ax.annotate(side_title, (0, 0.5), xytext=(-45, 0), ha='right', va='center', size=12, rotation=90, xycoords='axes fraction', textcoords='offset points') return ax #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def figure_basic(self, plot_S='line', plot_E='line', plot_I='line',plot_R='line', plot_F='line', plot_Q_E='line', plot_Q_I='line', combine_Q=True, color_S='tab:green', color_E='orange', color_I='crimson', color_R='tab:blue', color_F='black', color_Q_E='mediumorchid', color_Q_I='mediumorchid', color_reference='#E0E0E0', dashed_reference_results=None, dashed_reference_label='reference', shaded_reference_results=None, shaded_reference_label='reference', vlines=[], vline_colors=[], vline_styles=[], vline_labels=[], ylim=None, xlim=None, legend=True, title=None, side_title=None, plot_percentages=True, figsize=(12,8), use_seaborn=True, show=True): import matplotlib.pyplot as pyplot fig, ax = pyplot.subplots(figsize=figsize) if(use_seaborn): import seaborn seaborn.set_style('ticks') seaborn.despine() self.plot(ax=ax, plot_S=plot_S, plot_E=plot_E, plot_I=plot_I,plot_R=plot_R, plot_F=plot_F, plot_Q_E=plot_Q_E, plot_Q_I=plot_Q_I, combine_Q=combine_Q, color_S=color_S, color_E=color_E, color_I=color_I, color_R=color_R, color_F=color_F, color_Q_E=color_Q_E, color_Q_I=color_Q_I, color_reference=color_reference, dashed_reference_results=dashed_reference_results, dashed_reference_label=dashed_reference_label, shaded_reference_results=shaded_reference_results, shaded_reference_label=shaded_reference_label, vlines=vlines, vline_colors=vline_colors, vline_styles=vline_styles, vline_labels=vline_labels, ylim=ylim, xlim=xlim, legend=legend, title=title, side_title=side_title, plot_percentages=plot_percentages) if(show): pyplot.show() return fig, ax #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def figure_infections(self, plot_S=False, plot_E='stacked', plot_I='stacked',plot_R=False, plot_F=False, plot_Q_E='stacked', plot_Q_I='stacked', combine_Q=True, color_S='tab:green', color_E='orange', color_I='crimson', color_R='tab:blue', color_F='black', color_Q_E='mediumorchid', color_Q_I='mediumorchid', color_reference='#E0E0E0', dashed_reference_results=None, dashed_reference_label='reference', shaded_reference_results=None, shaded_reference_label='reference', vlines=[], vline_colors=[], vline_styles=[], vline_labels=[], ylim=None, xlim=None, legend=True, title=None, side_title=None, plot_percentages=True, figsize=(12,8), use_seaborn=True, show=True): import matplotlib.pyplot as pyplot fig, ax = pyplot.subplots(figsize=figsize) if(use_seaborn): import seaborn seaborn.set_style('ticks') seaborn.despine() self.plot(ax=ax, plot_S=plot_S, plot_E=plot_E, plot_I=plot_I,plot_R=plot_R, plot_F=plot_F, plot_Q_E=plot_Q_E, plot_Q_I=plot_Q_I, combine_Q=combine_Q, color_S=color_S, color_E=color_E, color_I=color_I, color_R=color_R, color_F=color_F, color_Q_E=color_Q_E, color_Q_I=color_Q_I, color_reference=color_reference, dashed_reference_results=dashed_reference_results, dashed_reference_label=dashed_reference_label, shaded_reference_results=shaded_reference_results, shaded_reference_label=shaded_reference_label, vlines=vlines, vline_colors=vline_colors, vline_styles=vline_styles, vline_labels=vline_labels, ylim=ylim, xlim=xlim, legend=legend, title=title, side_title=side_title, plot_percentages=plot_percentages) if(show): pyplot.show() return fig, ax #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% class SEIRSNetworkModel(): """ A class to simulate the SEIRS Stochastic Network Model ====================================================== Params: G Network adjacency matrix (numpy array) or Networkx graph object. beta Rate of transmission (global interactions) beta_local Rate(s) of transmission between adjacent individuals (optional) sigma Rate of progression to infectious state (inverse of latent period) gamma Rate of recovery (inverse of symptomatic infectious period) mu_I Rate of infection-related death xi Rate of re-susceptibility (upon recovery) mu_0 Rate of baseline death nu Rate of baseline birth p Probability of individuals interacting with global population G_Q Quarantine adjacency matrix (numpy array) or Networkx graph object. beta_Q Rate of transmission for isolated individuals (global interactions) beta_Q_local Rate(s) of transmission (exposure) for adjacent isolated individuals (optional) sigma_Q Rate of progression to infectious state for isolated individuals gamma_Q Rate of recovery for isolated individuals mu_Q Rate of infection-related death for isolated individuals q Probability of isolated individuals interacting with global population isolation_time Time to remain in isolation upon positive test, self-isolation, etc. theta_E Rate of random testing for exposed individuals theta_I Rate of random testing for infectious individuals phi_E Rate of testing when a close contact has tested positive for exposed individuals phi_I Rate of testing when a close contact has tested positive for infectious individuals psi_E Probability of positive test for exposed individuals psi_I Probability of positive test for infectious individuals initE Initial number of exposed individuals initI Initial number of infectious individuals initR Initial number of recovered individuals initF Initial number of infection-related fatalities initQ_S Initial number of isolated susceptible individuals initQ_E Initial number of isolated exposed individuals initQ_I Initial number of isolated infectious individuals initQ_R Initial number of isolated recovered individuals (all remaining nodes initialized susceptible) """ def __init__(self, G, beta, sigma, gamma, mu_I=0, alpha=1.0, xi=0, mu_0=0, nu=0, f=0, p=0, beta_local=None, beta_pairwise_mode='infected', delta=None, delta_pairwise_mode=None, G_Q=None, beta_Q=None, beta_Q_local=None, sigma_Q=None, gamma_Q=None, mu_Q=None, alpha_Q=None, delta_Q=None, theta_E=0, theta_I=0, phi_E=0, phi_I=0, psi_E=1, psi_I=1, q=0, isolation_time=14, initE=0, initI=0, initR=0, initF=0, initQ_E=0, initQ_I=0, transition_mode='exponential_rates', node_groups=None, store_Xseries=False, seed=None): if(seed is not None): numpy.random.seed(seed) self.seed = seed #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Model Parameters: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ self.parameters = { 'G':G, 'G_Q':G_Q, 'beta':beta, 'sigma':sigma, 'gamma':gamma, 'mu_I':mu_I, 'xi':xi, 'mu_0':mu_0, 'nu':nu, 'f':f, 'p':p, 'beta_local':beta_local, 'beta_pairwise_mode':beta_pairwise_mode, 'alpha':alpha, 'delta':delta, 'delta_pairwise_mode':delta_pairwise_mode, 'beta_Q':beta_Q, 'beta_Q_local':beta_Q_local, 'sigma_Q':sigma_Q, 'gamma_Q':gamma_Q, 'mu_Q':mu_Q, 'alpha_Q':alpha_Q, 'delta_Q':delta_Q, 'theta_E':theta_E, 'theta_I':theta_I, 'phi_E':phi_E, 'phi_I':phi_I, 'psi_E':psi_E, 'psi_I':psi_I, 'q':q, 'isolation_time':isolation_time, 'initE':initE, 'initI':initI, 'initR':initR, 'initF':initF, 'initQ_E':initQ_E, 'initQ_I':initQ_I } self.update_parameters() #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Each node can undergo 4-6 transitions (sans vitality/re-susceptibility returns to S state), # so there are ~numNodes*6 events/timesteps expected; initialize numNodes*6 timestep slots to start # (will be expanded during run if needed for some reason) self.tseries = numpy.zeros(6*self.numNodes) self.numS = numpy.zeros(6*self.numNodes) self.numE = numpy.zeros(6*self.numNodes) self.numI = numpy.zeros(6*self.numNodes) self.numR = numpy.zeros(6*self.numNodes) self.numF = numpy.zeros(6*self.numNodes) self.numQ_E = numpy.zeros(6*self.numNodes) self.numQ_I = numpy.zeros(6*self.numNodes) self.N = numpy.zeros(6*self.numNodes) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Initialize Timekeeping: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ self.t = 0 self.tmax = 0 # will be set when run() is called self.tidx = 0 self.tseries[0] = 0 # Vectors holding the time that each node has been in a given state or in isolation: self.timer_state = numpy.zeros((self.numNodes,1)) self.timer_isolation = numpy.zeros(self.numNodes) self.isolationTime = isolation_time #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Initialize Counts of inidividuals with each state: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ self.numE[0] = int(initE) self.numI[0] = int(initI) self.numR[0] = int(initR) self.numF[0] = int(initF) self.numQ_E[0] = int(initQ_E) self.numQ_I[0] = int(initQ_I) self.numS[0] = (self.numNodes - self.numE[0] - self.numI[0] - self.numR[0] - self.numQ_E[0] - self.numQ_I[0] - self.numF[0]) self.N[0] = self.numNodes - self.numF[0] #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Node states: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ self.S = 1 self.E = 2 self.I = 3 self.R = 4 self.F = 5 self.Q_E = 6 self.Q_I = 7 self.X = numpy.array( [self.S]*int(self.numS[0]) + [self.E]*int(self.numE[0]) + [self.I]*int(self.numI[0]) + [self.R]*int(self.numR[0]) + [self.F]*int(self.numF[0]) + [self.Q_E]*int(self.numQ_E[0]) + [self.Q_I]*int(self.numQ_I[0]) ).reshape((self.numNodes,1)) numpy.random.shuffle(self.X) self.store_Xseries = store_Xseries if(store_Xseries): self.Xseries = numpy.zeros(shape=(6*self.numNodes, self.numNodes), dtype='uint8') self.Xseries[0,:] = self.X.T self.transitions = { 'StoE': {'currentState':self.S, 'newState':self.E}, 'EtoI': {'currentState':self.E, 'newState':self.I}, 'ItoR': {'currentState':self.I, 'newState':self.R}, 'ItoF': {'currentState':self.I, 'newState':self.F}, 'RtoS': {'currentState':self.R, 'newState':self.S}, 'EtoQE': {'currentState':self.E, 'newState':self.Q_E}, 'ItoQI': {'currentState':self.I, 'newState':self.Q_I}, 'QEtoQI': {'currentState':self.Q_E, 'newState':self.Q_I}, 'QItoR': {'currentState':self.Q_I, 'newState':self.R}, 'QItoF': {'currentState':self.Q_I, 'newState':self.F}, '_toS': {'currentState':True, 'newState':self.S}, } self.transition_mode = transition_mode #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Initialize other node metadata: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ self.tested = numpy.array([False]*self.numNodes).reshape((self.numNodes,1)) self.positive = numpy.array([False]*self.numNodes).reshape((self.numNodes,1)) self.numTested = numpy.zeros(6*self.numNodes) self.numPositive = numpy.zeros(6*self.numNodes) self.testedInCurrentState = numpy.array([False]*self.numNodes).reshape((self.numNodes,1)) self.infectionsLog = [] #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Initialize node subgroup data series: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ self.nodeGroupData = None if(node_groups): self.nodeGroupData = {} for groupName, nodeList in node_groups.items(): self.nodeGroupData[groupName] = {'nodes': numpy.array(nodeList), 'mask': numpy.isin(range(self.numNodes), nodeList).reshape((self.numNodes,1))} self.nodeGroupData[groupName]['numS'] = numpy.zeros(6*self.numNodes) self.nodeGroupData[groupName]['numE'] = numpy.zeros(6*self.numNodes) self.nodeGroupData[groupName]['numI'] = numpy.zeros(6*self.numNodes) self.nodeGroupData[groupName]['numR'] = numpy.zeros(6*self.numNodes) self.nodeGroupData[groupName]['numF'] = numpy.zeros(6*self.numNodes) self.nodeGroupData[groupName]['numQ_E'] = numpy.zeros(6*self.numNodes) self.nodeGroupData[groupName]['numQ_I'] = numpy.zeros(6*self.numNodes) self.nodeGroupData[groupName]['N'] = numpy.zeros(6*self.numNodes) self.nodeGroupData[groupName]['numPositive'] = numpy.zeros(6*self.numNodes) self.nodeGroupData[groupName]['numTested'] = numpy.zeros(6*self.numNodes) self.nodeGroupData[groupName]['numS'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.S) self.nodeGroupData[groupName]['numE'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.E) self.nodeGroupData[groupName]['numI'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.I) self.nodeGroupData[groupName]['numR'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.R) self.nodeGroupData[groupName]['numF'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.F) self.nodeGroupData[groupName]['numQ_E'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_E) self.nodeGroupData[groupName]['numQ_I'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_I) self.nodeGroupData[groupName]['N'][0] = self.numNodes - self.numF[0] #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def update_parameters(self): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Model graphs: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ self.G = self.parameters['G'] # Adjacency matrix: if type(self.G)==numpy.ndarray: self.A = scipy.sparse.csr_matrix(self.G) elif type(self.G)==networkx.classes.graph.Graph: self.A = networkx.adj_matrix(self.G) # adj_matrix gives scipy.sparse csr_matrix else: raise BaseException("Input an adjacency matrix or networkx object only.") self.numNodes = int(self.A.shape[1]) self.degree = numpy.asarray(self.node_degrees(self.A)).astype(float) #---------------------------------------- if(self.parameters['G_Q'] is None): self.G_Q = self.G # If no Q graph is provided, use G in its place else: self.G_Q = self.parameters['G_Q'] # Quarantine Adjacency matrix: if type(self.G_Q)==numpy.ndarray: self.A_Q = scipy.sparse.csr_matrix(self.G_Q) elif type(self.G_Q)==networkx.classes.graph.Graph: self.A_Q = networkx.adj_matrix(self.G_Q) # adj_matrix gives scipy.sparse csr_matrix else: raise BaseException("Input an adjacency matrix or networkx object only.") self.numNodes_Q = int(self.A_Q.shape[1]) self.degree_Q = numpy.asarray(self.node_degrees(self.A_Q)).astype(float) #---------------------------------------- assert(self.numNodes == self.numNodes_Q), "The normal and quarantine adjacency graphs must be of the same size." #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Model parameters: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ self.beta = numpy.array(self.parameters['beta']).reshape((self.numNodes, 1)) if isinstance(self.parameters['beta'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['beta'], shape=(self.numNodes,1)) self.sigma = numpy.array(self.parameters['sigma']).reshape((self.numNodes, 1)) if isinstance(self.parameters['sigma'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['sigma'], shape=(self.numNodes,1)) self.gamma = numpy.array(self.parameters['gamma']).reshape((self.numNodes, 1)) if isinstance(self.parameters['gamma'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['gamma'], shape=(self.numNodes,1)) self.mu_I = numpy.array(self.parameters['mu_I']).reshape((self.numNodes, 1)) if isinstance(self.parameters['mu_I'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['mu_I'], shape=(self.numNodes,1)) self.alpha = numpy.array(self.parameters['alpha']).reshape((self.numNodes, 1)) if isinstance(self.parameters['alpha'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['alpha'], shape=(self.numNodes,1)) self.xi = numpy.array(self.parameters['xi']).reshape((self.numNodes, 1)) if isinstance(self.parameters['xi'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['xi'], shape=(self.numNodes,1)) self.mu_0 = numpy.array(self.parameters['mu_0']).reshape((self.numNodes, 1)) if isinstance(self.parameters['mu_0'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['mu_0'], shape=(self.numNodes,1)) self.nu = numpy.array(self.parameters['nu']).reshape((self.numNodes, 1)) if isinstance(self.parameters['nu'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['nu'], shape=(self.numNodes,1)) self.f = numpy.array(self.parameters['f']).reshape((self.numNodes, 1)) if isinstance(self.parameters['f'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['f'], shape=(self.numNodes,1)) self.p = numpy.array(self.parameters['p']).reshape((self.numNodes, 1)) if isinstance(self.parameters['p'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['p'], shape=(self.numNodes,1)) self.rand_f = numpy.random.rand(self.f.shape[0], self.f.shape[1]) #---------------------------------------- # Testing-related parameters: #---------------------------------------- self.beta_Q = (numpy.array(self.parameters['beta_Q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['beta_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['beta_Q'], shape=(self.numNodes,1))) if self.parameters['beta_Q'] is not None else self.beta self.sigma_Q = (numpy.array(self.parameters['sigma_Q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['sigma_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['sigma_Q'], shape=(self.numNodes,1))) if self.parameters['sigma_Q'] is not None else self.sigma self.gamma_Q = (numpy.array(self.parameters['gamma_Q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['gamma_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['gamma_Q'], shape=(self.numNodes,1))) if self.parameters['gamma_Q'] is not None else self.gamma self.mu_Q = (numpy.array(self.parameters['mu_Q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['mu_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['mu_Q'], shape=(self.numNodes,1))) if self.parameters['mu_Q'] is not None else self.mu_I self.alpha_Q = (numpy.array(self.parameters['alpha_Q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['alpha_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['alpha_Q'], shape=(self.numNodes,1))) if self.parameters['alpha_Q'] is not None else self.alpha self.theta_E = numpy.array(self.parameters['theta_E']).reshape((self.numNodes, 1)) if isinstance(self.parameters['theta_E'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['theta_E'], shape=(self.numNodes,1)) self.theta_I = numpy.array(self.parameters['theta_I']).reshape((self.numNodes, 1)) if isinstance(self.parameters['theta_I'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['theta_I'], shape=(self.numNodes,1)) self.phi_E = numpy.array(self.parameters['phi_E']).reshape((self.numNodes, 1)) if isinstance(self.parameters['phi_E'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['phi_E'], shape=(self.numNodes,1)) self.phi_I = numpy.array(self.parameters['phi_I']).reshape((self.numNodes, 1)) if isinstance(self.parameters['phi_I'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['phi_I'], shape=(self.numNodes,1)) self.psi_E = numpy.array(self.parameters['psi_E']).reshape((self.numNodes, 1)) if isinstance(self.parameters['psi_E'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['psi_E'], shape=(self.numNodes,1)) self.psi_I = numpy.array(self.parameters['psi_I']).reshape((self.numNodes, 1)) if isinstance(self.parameters['psi_I'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['psi_I'], shape=(self.numNodes,1)) self.q = numpy.array(self.parameters['q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['q'], shape=(self.numNodes,1)) #---------------------------------------- self.beta_pairwise_mode = self.parameters['beta_pairwise_mode'] #---------------------------------------- # Global transmission parameters: #---------------------------------------- if(self.beta_pairwise_mode == 'infected' or self.beta_pairwise_mode is None): self.beta_global = numpy.full_like(self.beta, fill_value=numpy.mean(self.beta)) self.beta_Q_global = numpy.full_like(self.beta_Q, fill_value=numpy.mean(self.beta_Q)) elif(self.beta_pairwise_mode == 'infectee'): self.beta_global = self.beta self.beta_Q_global = self.beta_Q elif(self.beta_pairwise_mode == 'min'): self.beta_global = numpy.minimum(self.beta, numpy.mean(beta)) self.beta_Q_global = numpy.minimum(self.beta_Q, numpy.mean(beta_Q)) elif(self.beta_pairwise_mode == 'max'): self.beta_global = numpy.maximum(self.beta, numpy.mean(beta)) self.beta_Q_global = numpy.maximum(self.beta_Q, numpy.mean(beta_Q)) elif(self.beta_pairwise_mode == 'mean'): self.beta_global = (self.beta + numpy.full_like(self.beta, fill_value=numpy.mean(self.beta)))/2 self.beta_Q_global = (self.beta_Q + numpy.full_like(self.beta_Q, fill_value=numpy.mean(self.beta_Q)))/2 #---------------------------------------- # Local transmission parameters: #---------------------------------------- self.beta_local = self.beta if self.parameters['beta_local'] is None else numpy.array(self.parameters['beta_local']) if isinstance(self.parameters['beta_local'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['beta_local'], shape=(self.numNodes,1)) self.beta_Q_local = self.beta_Q if self.parameters['beta_Q_local'] is None else numpy.array(self.parameters['beta_Q_local']) if isinstance(self.parameters['beta_Q_local'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['beta_Q_local'], shape=(self.numNodes,1)) #---------------------------------------- if(self.beta_local.ndim == 2 and self.beta_local.shape[0] == self.numNodes and self.beta_local.shape[1] == self.numNodes): self.A_beta_pairwise = self.beta_local elif((self.beta_local.ndim == 1 and self.beta_local.shape[0] == self.numNodes) or (self.beta_local.ndim == 2 and (self.beta_local.shape[0] == self.numNodes or self.beta_local.shape[1] == self.numNodes))): self.beta_local = self.beta_local.reshape((self.numNodes,1)) # Pre-multiply beta values by the adjacency matrix ("transmission weight connections") A_beta_pairwise_byInfected = scipy.sparse.csr_matrix.multiply(self.A, self.beta_local.T).tocsr() A_beta_pairwise_byInfectee = scipy.sparse.csr_matrix.multiply(self.A, self.beta_local).tocsr() #------------------------------ # Compute the effective pairwise beta values as a function of the infected/infectee pair: if(self.beta_pairwise_mode == 'infected'): self.A_beta_pairwise = A_beta_pairwise_byInfected elif(self.beta_pairwise_mode == 'infectee'): self.A_beta_pairwise = A_beta_pairwise_byInfectee elif(self.beta_pairwise_mode == 'min'): self.A_beta_pairwise = scipy.sparse.csr_matrix.minimum(A_beta_pairwise_byInfected, A_beta_pairwise_byInfectee) elif(self.beta_pairwise_mode == 'max'): self.A_beta_pairwise = scipy.sparse.csr_matrix.maximum(A_beta_pairwise_byInfected, A_beta_pairwise_byInfectee) elif(self.beta_pairwise_mode == 'mean' or self.beta_pairwise_mode is None): self.A_beta_pairwise = (A_beta_pairwise_byInfected + A_beta_pairwise_byInfectee)/2 else: print("Unrecognized beta_pairwise_mode value (support for 'infected', 'infectee', 'min', 'max', and 'mean').") else: print("Invalid values given for beta_local (expected 1xN list/array or NxN 2d array)") #---------------------------------------- if(self.beta_Q_local.ndim == 2 and self.beta_Q_local.shape[0] == self.numNodes and self.beta_Q_local.shape[1] == self.numNodes): self.A_Q_beta_Q_pairwise = self.beta_Q_local elif((self.beta_Q_local.ndim == 1 and self.beta_Q_local.shape[0] == self.numNodes) or (self.beta_Q_local.ndim == 2 and (self.beta_Q_local.shape[0] == self.numNodes or self.beta_Q_local.shape[1] == self.numNodes))): self.beta_Q_local = self.beta_Q_local.reshape((self.numNodes,1)) # Pre-multiply beta_Q values by the isolation adjacency matrix ("transmission weight connections") A_Q_beta_Q_pairwise_byInfected = scipy.sparse.csr_matrix.multiply(self.A_Q, self.beta_Q_local.T).tocsr() A_Q_beta_Q_pairwise_byInfectee = scipy.sparse.csr_matrix.multiply(self.A_Q, self.beta_Q_local).tocsr() #------------------------------ # Compute the effective pairwise beta values as a function of the infected/infectee pair: if(self.beta_pairwise_mode == 'infected'): self.A_Q_beta_Q_pairwise = A_Q_beta_Q_pairwise_byInfected elif(self.beta_pairwise_mode == 'infectee'): self.A_Q_beta_Q_pairwise = A_Q_beta_Q_pairwise_byInfectee elif(self.beta_pairwise_mode == 'min'): self.A_Q_beta_Q_pairwise = scipy.sparse.csr_matrix.minimum(A_Q_beta_Q_pairwise_byInfected, A_Q_beta_Q_pairwise_byInfectee) elif(self.beta_pairwise_mode == 'max'): self.A_Q_beta_Q_pairwise = scipy.sparse.csr_matrix.maximum(A_Q_beta_Q_pairwise_byInfected, A_Q_beta_Q_pairwise_byInfectee) elif(self.beta_pairwise_mode == 'mean' or self.beta_pairwise_mode is None): self.A_Q_beta_Q_pairwise = (A_Q_beta_Q_pairwise_byInfected + A_Q_beta_Q_pairwise_byInfectee)/2 else: print("Unrecognized beta_pairwise_mode value (support for 'infected', 'infectee', 'min', 'max', and 'mean').") else: print("Invalid values given for beta_Q_local (expected 1xN list/array or NxN 2d array)") #---------------------------------------- #---------------------------------------- # Degree-based transmission scaling parameters: #---------------------------------------- self.delta_pairwise_mode = self.parameters['delta_pairwise_mode'] with numpy.errstate(divide='ignore'): # ignore log(0) warning, then convert log(0) = -inf -> 0.0 self.delta = numpy.log(self.degree)/numpy.log(numpy.mean(self.degree)) if self.parameters['delta'] is None else numpy.array(self.parameters['delta']) if isinstance(self.parameters['delta'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['delta'], shape=(self.numNodes,1)) self.delta_Q = numpy.log(self.degree_Q)/numpy.log(numpy.mean(self.degree_Q)) if self.parameters['delta_Q'] is None else numpy.array(self.parameters['delta_Q']) if isinstance(self.parameters['delta_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['delta_Q'], shape=(self.numNodes,1)) self.delta[numpy.isneginf(self.delta)] = 0.0 self.delta_Q[numpy.isneginf(self.delta_Q)] = 0.0 #---------------------------------------- if(self.delta.ndim == 2 and self.delta.shape[0] == self.numNodes and self.delta.shape[1] == self.numNodes): self.A_delta_pairwise = self.delta elif((self.delta.ndim == 1 and self.delta.shape[0] == self.numNodes) or (self.delta.ndim == 2 and (self.delta.shape[0] == self.numNodes or self.delta.shape[1] == self.numNodes))): self.delta = self.delta.reshape((self.numNodes,1)) # Pre-multiply delta values by the adjacency matrix ("transmission weight connections") A_delta_pairwise_byInfected = scipy.sparse.csr_matrix.multiply(self.A, self.delta.T).tocsr() A_delta_pairwise_byInfectee = scipy.sparse.csr_matrix.multiply(self.A, self.delta).tocsr() #------------------------------ # Compute the effective pairwise delta values as a function of the infected/infectee pair: if(self.delta_pairwise_mode == 'infected'): self.A_delta_pairwise = A_delta_pairwise_byInfected elif(self.delta_pairwise_mode == 'infectee'): self.A_delta_pairwise = A_delta_pairwise_byInfectee elif(self.delta_pairwise_mode == 'min'): self.A_delta_pairwise = scipy.sparse.csr_matrix.minimum(A_delta_pairwise_byInfected, A_delta_pairwise_byInfectee) elif(self.delta_pairwise_mode == 'max'): self.A_delta_pairwise = scipy.sparse.csr_matrix.maximum(A_delta_pairwise_byInfected, A_delta_pairwise_byInfectee) elif(self.delta_pairwise_mode == 'mean'): self.A_delta_pairwise = (A_delta_pairwise_byInfected + A_delta_pairwise_byInfectee)/2 elif(self.delta_pairwise_mode is None): self.A_delta_pairwise = self.A else: print("Unrecognized delta_pairwise_mode value (support for 'infected', 'infectee', 'min', 'max', and 'mean').") else: print("Invalid values given for delta (expected 1xN list/array or NxN 2d array)") #---------------------------------------- if(self.delta_Q.ndim == 2 and self.delta_Q.shape[0] == self.numNodes and self.delta_Q.shape[1] == self.numNodes): self.A_Q_delta_Q_pairwise = self.delta_Q elif((self.delta_Q.ndim == 1 and self.delta_Q.shape[0] == self.numNodes) or (self.delta_Q.ndim == 2 and (self.delta_Q.shape[0] == self.numNodes or self.delta_Q.shape[1] == self.numNodes))): self.delta_Q = self.delta_Q.reshape((self.numNodes,1)) # Pre-multiply delta_Q values by the isolation adjacency matrix ("transmission weight connections") A_Q_delta_Q_pairwise_byInfected = scipy.sparse.csr_matrix.multiply(self.A_Q, self.delta_Q).tocsr() A_Q_delta_Q_pairwise_byInfectee = scipy.sparse.csr_matrix.multiply(self.A_Q, self.delta_Q.T).tocsr() #------------------------------ # Compute the effective pairwise delta values as a function of the infected/infectee pair: if(self.delta_pairwise_mode == 'infected'): self.A_Q_delta_Q_pairwise = A_Q_delta_Q_pairwise_byInfected elif(self.delta_pairwise_mode == 'infectee'): self.A_Q_delta_Q_pairwise = A_Q_delta_Q_pairwise_byInfectee elif(self.delta_pairwise_mode == 'min'): self.A_Q_delta_Q_pairwise = scipy.sparse.csr_matrix.minimum(A_Q_delta_Q_pairwise_byInfected, A_Q_delta_Q_pairwise_byInfectee) elif(self.delta_pairwise_mode == 'max'): self.A_Q_delta_Q_pairwise = scipy.sparse.csr_matrix.maximum(A_Q_delta_Q_pairwise_byInfected, A_Q_delta_Q_pairwise_byInfectee) elif(self.delta_pairwise_mode == 'mean'): self.A_Q_delta_Q_pairwise = (A_Q_delta_Q_pairwise_byInfected + A_Q_delta_Q_pairwise_byInfectee)/2 elif(self.delta_pairwise_mode is None): self.A_Q_delta_Q_pairwise = self.A else: print("Unrecognized delta_pairwise_mode value (support for 'infected', 'infectee', 'min', 'max', and 'mean').") else: print("Invalid values given for delta_Q (expected 1xN list/array or NxN 2d array)") #---------------------------------------- # Pre-calculate the pairwise delta*beta values: #---------------------------------------- self.A_deltabeta = scipy.sparse.csr_matrix.multiply(self.A_delta_pairwise, self.A_beta_pairwise) self.A_Q_deltabeta_Q = scipy.sparse.csr_matrix.multiply(self.A_Q_delta_Q_pairwise, self.A_Q_beta_Q_pairwise) #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def node_degrees(self, Amat): return Amat.sum(axis=0).reshape(self.numNodes,1) # sums of adj matrix cols #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_susceptible(self, t_idx=None): if(t_idx is None): return (self.numS[:]) else: return (self.numS[t_idx]) #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_infected(self, t_idx=None): if(t_idx is None): return (self.numE[:] + self.numI[:] + self.numQ_E[:] + self.numQ_I[:]) else: return (self.numE[t_idx] + self.numI[t_idx] + self.numQ_E[t_idx] + self.numQ_I[t_idx]) #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_isolated(self, t_idx=None): if(t_idx is None): return (self.numQ_E[:] + self.numQ_I[:]) else: return (self.numQ_E[t_idx] + self.numQ_I[t_idx]) #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_tested(self, t_idx=None): if(t_idx is None): return (self.numTested[:]) else: return (self.numTested[t_idx]) #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_positive(self, t_idx=None): if(t_idx is None): return (self.numPositive[:]) else: return (self.numPositive[t_idx]) #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_recovered(self, t_idx=None): if(t_idx is None): return (self.numR[:]) else: return (self.numR[t_idx]) #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def calc_propensities(self): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Pre-calculate matrix multiplication terms that may be used in multiple propensity calculations, # and check to see if their computation is necessary before doing the multiplication #------------------------------------ self.transmissionTerms_I = numpy.zeros(shape=(self.numNodes,1)) if(numpy.any(self.numI[self.tidx])): self.transmissionTerms_I = numpy.asarray(scipy.sparse.csr_matrix.dot(self.A_deltabeta, self.X==self.I)) #------------------------------------ self.transmissionTerms_Q = numpy.zeros(shape=(self.numNodes,1)) if(numpy.any(self.numQ_I[self.tidx])): self.transmissionTerms_Q = numpy.asarray(scipy.sparse.csr_matrix.dot(self.A_Q_deltabeta_Q, self.X==self.Q_I)) #------------------------------------ numContacts_Q = numpy.zeros(shape=(self.numNodes,1)) if(numpy.any(self.positive) and (numpy.any(self.phi_E) or numpy.any(self.phi_I))): numContacts_Q = numpy.asarray(scipy.sparse.csr_matrix.dot(self.A, ((self.positive)&(self.X!=self.R)&(self.X!=self.F)))) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ propensities_StoE = (self.alpha * (self.p*((self.beta_global*self.numI[self.tidx] + self.q*self.beta_Q_global*self.numQ_I[self.tidx])/self.N[self.tidx]) + (1-self.p)*(numpy.divide(self.transmissionTerms_I, self.degree, out=numpy.zeros_like(self.degree), where=self.degree!=0) +numpy.divide(self.transmissionTerms_Q, self.degree_Q, out=numpy.zeros_like(self.degree_Q), where=self.degree_Q!=0))) )*(self.X==self.S) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(self.transition_mode == 'time_in_state'): propensities_EtoI = 1e5 * ((self.X==self.E) & numpy.greater(self.timer_state, 1/self.sigma)) propensities_ItoR = 1e5 * ((self.X==self.I) & numpy.greater(self.timer_state, 1/self.gamma) & numpy.greater_equal(self.rand_f, self.f)) propensities_ItoF = 1e5 * ((self.X==self.I) & numpy.greater(self.timer_state, 1/self.mu_I) & numpy.less(self.rand_f, self.f)) propensities_EtoQE = numpy.zeros_like(propensities_StoE) propensities_ItoQI = numpy.zeros_like(propensities_StoE) propensities_QEtoQI = 1e5 * ((self.X==self.Q_E) & numpy.greater(self.timer_state, 1/self.sigma_Q)) propensities_QItoR = 1e5 * ((self.X==self.Q_I) & numpy.greater(self.timer_state, 1/self.gamma_Q) & numpy.greater_equal(self.rand_f, self.f)) propensities_QItoF = 1e5 * ((self.X==self.Q_I) & numpy.greater(self.timer_state, 1/self.mu_Q) & numpy.less(self.rand_f, self.f)) propensities_RtoS = 1e5 * ((self.X==self.R) & numpy.greater(self.timer_state, 1/self.xi)) propensities__toS = 1e5 * ((self.X!=self.F) & numpy.greater(self.timer_state, 1/self.nu)) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ else: # exponential_rates propensities_EtoI = self.sigma * (self.X==self.E) propensities_ItoR = self.gamma * ((self.X==self.I) & (numpy.greater_equal(self.rand_f, self.f))) propensities_ItoF = self.mu_I * ((self.X==self.I) & (numpy.less(self.rand_f, self.f))) propensities_EtoQE = (self.theta_E + self.phi_E*numContacts_Q)*self.psi_E * (self.X==self.E) propensities_ItoQI = (self.theta_I + self.phi_I*numContacts_Q)*self.psi_I * (self.X==self.I) propensities_QEtoQI = self.sigma_Q * (self.X==self.Q_E) propensities_QItoR = self.gamma_Q * ((self.X==self.Q_I) & (numpy.greater_equal(self.rand_f, self.f))) propensities_QItoF = self.mu_Q * ((self.X==self.Q_I) & (numpy.less(self.rand_f, self.f))) propensities_RtoS = self.xi * (self.X==self.R) propensities__toS = self.nu * (self.X!=self.F) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ propensities = numpy.hstack([propensities_StoE, propensities_EtoI, propensities_ItoR, propensities_ItoF, propensities_EtoQE, propensities_ItoQI, propensities_QEtoQI, propensities_QItoR, propensities_QItoF, propensities_RtoS, propensities__toS]) columns = ['StoE', 'EtoI', 'ItoR', 'ItoF', 'EtoQE', 'ItoQI', 'QEtoQI', 'QItoR', 'QItoF', 'RtoS', '_toS'] return propensities, columns #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def set_isolation(self, node, isolate): # Move this node in/out of the appropriate isolation state: if(isolate == True): if(self.X[node] == self.E): self.X[node] = self.Q_E elif(self.X[node] == self.I): self.X[node] = self.Q_I elif(isolate == False): if(self.X[node] == self.Q_E): self.X[node] = self.E elif(self.X[node] == self.Q_I): self.X[node] = self.I # Reset the isolation timer: self.timer_isolation[node] = 0 #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def set_tested(self, node, tested): self.tested[node] = tested self.testedInCurrentState[node] = tested #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def set_positive(self, node, positive): self.positive[node] = positive #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def introduce_exposures(self, num_new_exposures): exposedNodes = numpy.random.choice(range(self.numNodes), size=num_new_exposures, replace=False) for exposedNode in exposedNodes: if(self.X[exposedNode]==self.S): self.X[exposedNode] = self.E #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def increase_data_series_length(self): self.tseries = numpy.pad(self.tseries, [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.numS = numpy.pad(self.numS, [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.numE = numpy.pad(self.numE, [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.numI = numpy.pad(self.numI, [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.numR = numpy.pad(self.numR, [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.numF = numpy.pad(self.numF, [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.numQ_E = numpy.pad(self.numQ_E, [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.numQ_I = numpy.pad(self.numQ_I, [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.N = numpy.pad(self.N, [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.numTested = numpy.pad(self.numTested, [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.numPositive = numpy.pad(self.numPositive, [(0, 6*self.numNodes)], mode='constant', constant_values=0) if(self.store_Xseries): self.Xseries = numpy.pad(self.Xseries, [(0, 6*self.numNodes), (0,0)], mode='constant', constant_values=0) if(self.nodeGroupData): for groupName in self.nodeGroupData: self.nodeGroupData[groupName]['numS'] = numpy.pad(self.nodeGroupData[groupName]['numS'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.nodeGroupData[groupName]['numE'] = numpy.pad(self.nodeGroupData[groupName]['numE'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.nodeGroupData[groupName]['numI'] = numpy.pad(self.nodeGroupData[groupName]['numI'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.nodeGroupData[groupName]['numR'] = numpy.pad(self.nodeGroupData[groupName]['numR'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.nodeGroupData[groupName]['numF'] = numpy.pad(self.nodeGroupData[groupName]['numF'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.nodeGroupData[groupName]['numQ_E'] = numpy.pad(self.nodeGroupData[groupName]['numQ_E'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.nodeGroupData[groupName]['numQ_I'] = numpy.pad(self.nodeGroupData[groupName]['numQ_I'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.nodeGroupData[groupName]['N'] = numpy.pad(self.nodeGroupData[groupName]['N'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.nodeGroupData[groupName]['numTested'] = numpy.pad(self.nodeGroupData[groupName]['numTested'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.nodeGroupData[groupName]['numPositive'] = numpy.pad(self.nodeGroupData[groupName]['numPositive'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) return None #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def finalize_data_series(self): self.tseries = numpy.array(self.tseries, dtype=float)[:self.tidx+1] self.numS = numpy.array(self.numS, dtype=float)[:self.tidx+1] self.numE = numpy.array(self.numE, dtype=float)[:self.tidx+1] self.numI = numpy.array(self.numI, dtype=float)[:self.tidx+1] self.numR = numpy.array(self.numR, dtype=float)[:self.tidx+1] self.numF = numpy.array(self.numF, dtype=float)[:self.tidx+1] self.numQ_E = numpy.array(self.numQ_E, dtype=float)[:self.tidx+1] self.numQ_I = numpy.array(self.numQ_I, dtype=float)[:self.tidx+1] self.N = numpy.array(self.N, dtype=float)[:self.tidx+1] self.numTested = numpy.array(self.numTested, dtype=float)[:self.tidx+1] self.numPositive = numpy.array(self.numPositive, dtype=float)[:self.tidx+1] if(self.store_Xseries): self.Xseries = self.Xseries[:self.tidx+1, :] if(self.nodeGroupData): for groupName in self.nodeGroupData: self.nodeGroupData[groupName]['numS'] = numpy.array(self.nodeGroupData[groupName]['numS'], dtype=float)[:self.tidx+1] self.nodeGroupData[groupName]['numE'] = numpy.array(self.nodeGroupData[groupName]['numE'], dtype=float)[:self.tidx+1] self.nodeGroupData[groupName]['numI'] = numpy.array(self.nodeGroupData[groupName]['numI'], dtype=float)[:self.tidx+1] self.nodeGroupData[groupName]['numR'] = numpy.array(self.nodeGroupData[groupName]['numR'], dtype=float)[:self.tidx+1] self.nodeGroupData[groupName]['numF'] = numpy.array(self.nodeGroupData[groupName]['numF'], dtype=float)[:self.tidx+1] self.nodeGroupData[groupName]['numQ_E'] = numpy.array(self.nodeGroupData[groupName]['numQ_E'], dtype=float)[:self.tidx+1] self.nodeGroupData[groupName]['numQ_I'] = numpy.array(self.nodeGroupData[groupName]['numQ_I'], dtype=float)[:self.tidx+1] self.nodeGroupData[groupName]['N'] = numpy.array(self.nodeGroupData[groupName]['N'], dtype=float)[:self.tidx+1] self.nodeGroupData[groupName]['numTested'] = numpy.array(self.nodeGroupData[groupName]['numTested'], dtype=float)[:self.tidx+1] self.nodeGroupData[groupName]['numPositive'] = numpy.array(self.nodeGroupData[groupName]['numPositive'], dtype=float)[:self.tidx+1] return None #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def run_iteration(self): if(self.tidx >= len(self.tseries)-1): # Room has run out in the timeseries storage arrays; double the size of these arrays: self.increase_data_series_length() #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Generate 2 random numbers uniformly distributed in (0,1) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ r1 = numpy.random.rand() r2 = numpy.random.rand() #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Calculate propensities #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ propensities, transitionTypes = self.calc_propensities() if(propensities.sum() > 0): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Calculate alpha #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ propensities_flat = propensities.ravel(order='F') cumsum = propensities_flat.cumsum() alpha = propensities_flat.sum() #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Compute the time until the next event takes place #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ tau = (1/alpha)*numpy.log(float(1/r1)) self.t += tau self.timer_state += tau #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Compute which event takes place #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ transitionIdx = numpy.searchsorted(cumsum,r2*alpha) transitionNode = transitionIdx % self.numNodes transitionType = transitionTypes[ int(transitionIdx/self.numNodes) ] #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Perform updates triggered by rate propensities: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ assert(self.X[transitionNode] == self.transitions[transitionType]['currentState'] and self.X[transitionNode]!=self.F), "Assertion error: Node "+str(transitionNode)+" has unexpected current state "+str(self.X[transitionNode])+" given the intended transition of "+str(transitionType)+"." self.X[transitionNode] = self.transitions[transitionType]['newState'] self.testedInCurrentState[transitionNode] = False self.timer_state[transitionNode] = 0.0 #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Save information about infection events when they occur: if(transitionType == 'StoE'): transitionNode_GNbrs = list(self.G[transitionNode].keys()) transitionNode_GQNbrs = list(self.G_Q[transitionNode].keys()) self.infectionsLog.append({ 't': self.t, 'infected_node': transitionNode, 'infection_type': transitionType, 'infected_node_degree': self.degree[transitionNode], 'local_contact_nodes': transitionNode_GNbrs, 'local_contact_node_states': self.X[transitionNode_GNbrs].flatten(), 'isolation_contact_nodes': transitionNode_GQNbrs, 'isolation_contact_node_states':self.X[transitionNode_GQNbrs].flatten() }) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(transitionType in ['EtoQE', 'ItoQI']): self.set_positive(node=transitionNode, positive=True) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ else: tau = 0.01 self.t += tau self.timer_state += tau #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ self.tidx += 1 self.tseries[self.tidx] = self.t self.numS[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.S), a_min=0, a_max=self.numNodes) self.numE[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.E), a_min=0, a_max=self.numNodes) self.numI[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.I), a_min=0, a_max=self.numNodes) self.numF[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.F), a_min=0, a_max=self.numNodes) self.numQ_E[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.Q_E), a_min=0, a_max=self.numNodes) self.numQ_I[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.Q_I), a_min=0, a_max=self.numNodes) self.numTested[self.tidx] = numpy.clip(numpy.count_nonzero(self.tested), a_min=0, a_max=self.numNodes) self.numPositive[self.tidx] = numpy.clip(numpy.count_nonzero(self.positive), a_min=0, a_max=self.numNodes) self.N[self.tidx] = numpy.clip((self.numNodes - self.numF[self.tidx]), a_min=0, a_max=self.numNodes) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Update testing and isolation statuses #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ isolatedNodes = numpy.argwhere((self.X==self.Q_E)|(self.X==self.Q_I))[:,0].flatten() self.timer_isolation[isolatedNodes] = self.timer_isolation[isolatedNodes] + tau nodesExitingIsolation = numpy.argwhere(self.timer_isolation >= self.isolationTime) for isoNode in nodesExitingIsolation: self.set_isolation(node=isoNode, isolate=False) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Store system states #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(self.store_Xseries): self.Xseries[self.tidx,:] = self.X.T if(self.nodeGroupData): for groupName in self.nodeGroupData: self.nodeGroupData[groupName]['numS'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.S) self.nodeGroupData[groupName]['numE'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.E) self.nodeGroupData[groupName]['numI'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.I) self.nodeGroupData[groupName]['numR'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.R) self.nodeGroupData[groupName]['numF'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.F) self.nodeGroupData[groupName]['numQ_E'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_E) self.nodeGroupData[groupName]['numQ_I'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_I) self.nodeGroupData[groupName]['N'][self.tidx] = numpy.clip((self.nodeGroupData[groupName]['numS'][0] + self.nodeGroupData[groupName]['numE'][0] + self.nodeGroupData[groupName]['numI'][0] + self.nodeGroupData[groupName]['numQ_E'][0] + self.nodeGroupData[groupName]['numQ_I'][0] + self.nodeGroupData[groupName]['numR'][0]), a_min=0, a_max=self.numNodes) self.nodeGroupData[groupName]['numTested'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.tested) self.nodeGroupData[groupName]['numPositive'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.positive) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Terminate if tmax reached or num infections is 0: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(self.t >= self.tmax or (self.total_num_infected(self.tidx) < 1 and self.total_num_isolated(self.tidx) < 1)): self.finalize_data_series() return False #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ return True #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def run(self, T, checkpoints=None, print_interval=10, verbose='t'): if(T>0): self.tmax += T else: return False #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Pre-process checkpoint values: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(checkpoints): numCheckpoints = len(checkpoints['t']) for chkpt_param, chkpt_values in checkpoints.items(): assert(isinstance(chkpt_values, (list, numpy.ndarray)) and len(chkpt_values)==numCheckpoints), "Expecting a list of values with length equal to number of checkpoint times ("+str(numCheckpoints)+") for each checkpoint parameter." checkpointIdx = numpy.searchsorted(checkpoints['t'], self.t) # Finds 1st index in list greater than given val if(checkpointIdx >= numCheckpoints): # We are out of checkpoints, stop checking them: checkpoints = None else: checkpointTime = checkpoints['t'][checkpointIdx] #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% # Run the simulation loop: #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% print_reset = True running = True while running: running = self.run_iteration() #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Handle checkpoints if applicable: if(checkpoints): if(self.t >= checkpointTime): if(verbose is not False): print("[Checkpoint: Updating parameters]") # A checkpoint has been reached, update param values: for param in list(self.parameters.keys()): if(param in list(checkpoints.keys())): self.parameters.update({param: checkpoints[param][checkpointIdx]}) # Update parameter data structures and scenario flags: self.update_parameters() # Update the next checkpoint time: checkpointIdx = numpy.searchsorted(checkpoints['t'], self.t) # Finds 1st index in list greater than given val if(checkpointIdx >= numCheckpoints): # We are out of checkpoints, stop checking them: checkpoints = None else: checkpointTime = checkpoints['t'][checkpointIdx] #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(print_interval): if(print_reset and (int(self.t) % print_interval == 0)): if(verbose=="t"): print("t = %.2f" % self.t) if(verbose==True): print("t = %.2f" % self.t) print("\t S = " + str(self.numS[self.tidx])) print("\t E = " + str(self.numE[self.tidx])) print("\t I = " + str(self.numI[self.tidx])) print("\t R = " + str(self.numR[self.tidx])) print("\t F = " + str(self.numF[self.tidx])) print("\t Q_E = " + str(self.numQ_E[self.tidx])) print("\t Q_I = " + str(self.numQ_I[self.tidx])) print_reset = False elif(not print_reset and (int(self.t) % 10 != 0)): print_reset = True return True #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I='line',plot_R='line', plot_F='line', plot_Q_E='line', plot_Q_I='line', combine_D=True, color_S='tab:green', color_E='orange', color_I='crimson', color_R='tab:blue', color_F='black', color_Q_E='mediumorchid', color_Q_I='mediumorchid', color_reference='#E0E0E0', dashed_reference_results=None, dashed_reference_label='reference', shaded_reference_results=None, shaded_reference_label='reference', vlines=[], vline_colors=[], vline_styles=[], vline_labels=[], ylim=None, xlim=None, legend=True, title=None, side_title=None, plot_percentages=True): import matplotlib.pyplot as pyplot #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Create an Axes object if None provided: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(not ax): fig, ax = pyplot.subplots() #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Prepare data series to be plotted: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Fseries = self.numF/self.numNodes if plot_percentages else self.numF Eseries = self.numE/self.numNodes if plot_percentages else self.numE Dseries = (self.numQ_E+self.numQ_I)/self.numNodes if plot_percentages else (self.numQ_E+self.numQ_I) Q_Eseries = self.numQ_E/self.numNodes if plot_percentages else self.numQ_E Q_Iseries = self.numQ_I/self.numNodes if plot_percentages else self.numQ_I Iseries = self.numI/self.numNodes if plot_percentages else self.numI Rseries = self.numR/self.numNodes if plot_percentages else self.numR Sseries = self.numS/self.numNodes if plot_percentages else self.numS #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Draw the reference data: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(dashed_reference_results): dashedReference_tseries = dashed_reference_results.tseries[::int(self.numNodes/100)] dashedReference_IDEstack = (dashed_reference_results.numI + dashed_reference_results.numQ_I + dashed_reference_results.numQ_E + dashed_reference_results.numE)[::int(self.numNodes/100)] / (self.numNodes if plot_percentages else 1) ax.plot(dashedReference_tseries, dashedReference_IDEstack, color='#E0E0E0', linestyle='--', label='$I+D+E$ ('+dashed_reference_label+')', zorder=0) if(shaded_reference_results): shadedReference_tseries = shaded_reference_results.tseries shadedReference_IDEstack = (shaded_reference_results.numI + shaded_reference_results.numQ_I + shaded_reference_results.numQ_E + shaded_reference_results.numE) / (self.numNodes if plot_percentages else 1) ax.fill_between(shaded_reference_results.tseries, shadedReference_IDEstack, 0, color='#EFEFEF', label='$I+D+E$ ('+shaded_reference_label+')', zorder=0) ax.plot(shaded_reference_results.tseries, shadedReference_IDEstack, color='#E0E0E0', zorder=1) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Draw the stacked variables: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ topstack = numpy.zeros_like(self.tseries) if(any(Fseries) and plot_F=='stacked'): ax.fill_between(numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, topstack+Fseries), topstack, color=color_F, alpha=0.5, label='$F$', zorder=2) ax.plot( numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, topstack+Fseries), color=color_F, zorder=3) topstack = topstack+Fseries if(any(Eseries) and plot_E=='stacked'): ax.fill_between(numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, topstack+Eseries), topstack, color=color_E, alpha=0.5, label='$E$', zorder=2) ax.plot( numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, topstack+Eseries), color=color_E, zorder=3) topstack = topstack+Eseries if(combine_D and plot_Q_E=='stacked' and plot_Q_I=='stacked'): ax.fill_between(numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, topstack+Dseries), topstack, color=color_Q_E, alpha=0.5, label='$Q_{all}$', zorder=2) ax.plot( numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, topstack+Dseries), color=color_Q_E, zorder=3) topstack = topstack+Dseries else: if(any(Q_Eseries) and plot_Q_E=='stacked'): ax.fill_between(numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, topstack+Q_Eseries), topstack, color=color_Q_E, alpha=0.5, label='$Q_E$', zorder=2) ax.plot( numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, topstack+Q_Eseries), color=color_Q_E, zorder=3) topstack = topstack+Q_Eseries if(any(Q_Iseries) and plot_Q_I=='stacked'): ax.fill_between(numpy.ma.masked_where(Q_Iseries<=0, self.tseries), numpy.ma.masked_where(Q_Iseries<=0, topstack+Q_Iseries), topstack, color=color_Q_I, alpha=0.5, label='$Q_I$', zorder=2) ax.plot( numpy.ma.masked_where(Q_Iseries<=0, self.tseries), numpy.ma.masked_where(Q_Iseries<=0, topstack+Q_Iseries), color=color_Q_I, zorder=3) topstack = topstack+Q_Iseries if(any(Iseries) and plot_I=='stacked'): ax.fill_between(numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, topstack+Iseries), topstack, color=color_I, alpha=0.5, label='$I$', zorder=2) ax.plot( numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, topstack+Iseries), color=color_I, zorder=3) topstack = topstack+Iseries if(any(Rseries) and plot_R=='stacked'): ax.fill_between(numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, topstack+Rseries), topstack, color=color_R, alpha=0.5, label='$R$', zorder=2) ax.plot( numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, topstack+Rseries), color=color_R, zorder=3) topstack = topstack+Rseries if(any(Sseries) and plot_S=='stacked'): ax.fill_between(numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, topstack+Sseries), topstack, color=color_S, alpha=0.5, label='$S$', zorder=2) ax.plot( numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, topstack+Sseries), color=color_S, zorder=3) topstack = topstack+Sseries #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Draw the shaded variables: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(any(Fseries) and plot_F=='shaded'): ax.fill_between(numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, Fseries), 0, color=color_F, alpha=0.5, label='$F$', zorder=4) ax.plot( numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, Fseries), color=color_F, zorder=5) if(any(Eseries) and plot_E=='shaded'): ax.fill_between(numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, Eseries), 0, color=color_E, alpha=0.5, label='$E$', zorder=4) ax.plot( numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, Eseries), color=color_E, zorder=5) if(combine_D and (any(Dseries) and plot_Q_E=='shaded' and plot_Q_I=='shaded')): ax.fill_between(numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, Dseries), 0, color=color_Q_E, alpha=0.5, label='$Q_{all}$', zorder=4) ax.plot( numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, Dseries), color=color_Q_E, zorder=5) else: if(any(Q_Eseries) and plot_Q_E=='shaded'): ax.fill_between(numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, Q_Eseries), 0, color=color_Q_E, alpha=0.5, label='$Q_E$', zorder=4) ax.plot( numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, Q_Eseries), color=color_Q_E, zorder=5) if(any(Q_Iseries) and plot_Q_I=='shaded'): ax.fill_between(numpy.ma.masked_where(Q_Iseries<=0, self.tseries), numpy.ma.masked_where(Q_Iseries<=0, Q_Iseries), 0, color=color_Q_I, alpha=0.5, label='$Q_I$', zorder=4) ax.plot( numpy.ma.masked_where(Q_Iseries<=0, self.tseries), numpy.ma.masked_where(Q_Iseries<=0, Q_Iseries), color=color_Q_I, zorder=5) if(any(Iseries) and plot_I=='shaded'): ax.fill_between(numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, Iseries), 0, color=color_I, alpha=0.5, label='$I$', zorder=4) ax.plot( numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, Iseries), color=color_I, zorder=5) if(any(Sseries) and plot_S=='shaded'): ax.fill_between(numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, Sseries), 0, color=color_S, alpha=0.5, label='$S$', zorder=4) ax.plot( numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, Sseries), color=color_S, zorder=5) if(any(Rseries) and plot_R=='shaded'): ax.fill_between(numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, Rseries), 0, color=color_R, alpha=0.5, label='$R$', zorder=4) ax.plot( numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, Rseries), color=color_R, zorder=5) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Draw the line variables: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(any(Fseries) and plot_F=='line'): ax.plot(numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, Fseries), color=color_F, label='$F$', zorder=6) if(any(Eseries) and plot_E=='line'): ax.plot(numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, Eseries), color=color_E, label='$E$', zorder=6) if(combine_D and (any(Dseries) and plot_Q_E=='line' and plot_Q_I=='line')): ax.plot(numpy.ma.masked_where(Dseries<=0, self.tseries), numpy.ma.masked_where(Dseries<=0, Dseries), color=color_Q_E, label='$Q_{all}$', zorder=6) else: if(any(Q_Eseries) and plot_Q_E=='line'): ax.plot(numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, Q_Eseries), color=color_Q_E, label='$Q_E$', zorder=6) if(any(Q_Iseries) and plot_Q_I=='line'): ax.plot(numpy.ma.masked_where(Q_Iseries<=0, self.tseries), numpy.ma.masked_where(Q_Iseries<=0, Q_Iseries), color=color_Q_I, label='$Q_I$', zorder=6) if(any(Iseries) and plot_I=='line'): ax.plot(numpy.ma.masked_where(Iseries<=0, self.tseries), numpy.ma.masked_where(Iseries<=0, Iseries), color=color_I, label='$I$', zorder=6) if(any(Sseries) and plot_S=='line'): ax.plot(numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, Sseries), color=color_S, label='$S$', zorder=6) if(any(Rseries) and plot_R=='line'): ax.plot(numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, Rseries), color=color_R, label='$R$', zorder=6) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Draw the vertical line annotations: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(len(vlines)>0 and len(vline_colors)==0): vline_colors = ['gray']*len(vlines) if(len(vlines)>0 and len(vline_labels)==0): vline_labels = [None]*len(vlines) if(len(vlines)>0 and len(vline_styles)==0): vline_styles = [':']*len(vlines) for vline_x, vline_color, vline_style, vline_label in zip(vlines, vline_colors, vline_styles, vline_labels): if(vline_x is not None): ax.axvline(x=vline_x, color=vline_color, linestyle=vline_style, alpha=1, label=vline_label) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Draw the plot labels: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ ax.set_xlabel('days') ax.set_ylabel('percent of population' if plot_percentages else 'number of individuals') ax.set_xlim(0, (max(self.tseries) if not xlim else xlim)) ax.set_ylim(0, ylim) if(plot_percentages): ax.set_yticklabels(['{:,.0%}'.format(y) for y in ax.get_yticks()]) if(legend): legend_handles, legend_labels = ax.get_legend_handles_labels() ax.legend(legend_handles[::-1], legend_labels[::-1], loc='upper right', facecolor='white', edgecolor='none', framealpha=0.9, prop={'size': 8}) if(title): ax.set_title(title, size=12) if(side_title): ax.annotate(side_title, (0, 0.5), xytext=(-45, 0), ha='right', va='center', size=12, rotation=90, xycoords='axes fraction', textcoords='offset points') return ax #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def figure_basic(self, plot_S='line', plot_E='line', plot_I='line',plot_R='line', plot_F='line', plot_Q_E='line', plot_Q_I='line', combine_D=True, color_S='tab:green', color_E='orange', color_I='crimson', color_R='tab:blue', color_F='black', color_Q_E='mediumorchid', color_Q_I='mediumorchid', color_reference='#E0E0E0', dashed_reference_results=None, dashed_reference_label='reference', shaded_reference_results=None, shaded_reference_label='reference', vlines=[], vline_colors=[], vline_styles=[], vline_labels=[], ylim=None, xlim=None, legend=True, title=None, side_title=None, plot_percentages=True, figsize=(12,8), use_seaborn=True, show=True): import matplotlib.pyplot as pyplot fig, ax = pyplot.subplots(figsize=figsize) if(use_seaborn): import seaborn seaborn.set_style('ticks') seaborn.despine() self.plot(ax=ax, plot_S=plot_S, plot_E=plot_E, plot_I=plot_I,plot_R=plot_R, plot_F=plot_F, plot_Q_E=plot_Q_E, plot_Q_I=plot_Q_I, combine_D=combine_D, color_S=color_S, color_E=color_E, color_I=color_I, color_R=color_R, color_F=color_F, color_Q_E=color_Q_E, color_Q_I=color_Q_I, color_reference=color_reference, dashed_reference_results=dashed_reference_results, dashed_reference_label=dashed_reference_label, shaded_reference_results=shaded_reference_results, shaded_reference_label=shaded_reference_label, vlines=vlines, vline_colors=vline_colors, vline_styles=vline_styles, vline_labels=vline_labels, ylim=ylim, xlim=xlim, legend=legend, title=title, side_title=side_title, plot_percentages=plot_percentages) if(show): pyplot.show() return fig, ax #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def figure_infections(self, plot_S=False, plot_E='stacked', plot_I='stacked',plot_R=False, plot_F=False, plot_Q_E='stacked', plot_Q_I='stacked', combine_D=True, color_S='tab:green', color_E='orange', color_I='crimson', color_R='tab:blue', color_F='black', color_Q_E='mediumorchid', color_Q_I='mediumorchid', color_reference='#E0E0E0', dashed_reference_results=None, dashed_reference_label='reference', shaded_reference_results=None, shaded_reference_label='reference', vlines=[], vline_colors=[], vline_styles=[], vline_labels=[], ylim=None, xlim=None, legend=True, title=None, side_title=None, plot_percentages=True, figsize=(12,8), use_seaborn=True, show=True): import matplotlib.pyplot as pyplot fig, ax = pyplot.subplots(figsize=figsize) if(use_seaborn): import seaborn seaborn.set_style('ticks') seaborn.despine() self.plot(ax=ax, plot_S=plot_S, plot_E=plot_E, plot_I=plot_I,plot_R=plot_R, plot_F=plot_F, plot_Q_E=plot_Q_E, plot_Q_I=plot_Q_I, combine_D=combine_D, color_S=color_S, color_E=color_E, color_I=color_I, color_R=color_R, color_F=color_F, color_Q_E=color_Q_E, color_Q_I=color_Q_I, color_reference=color_reference, dashed_reference_results=dashed_reference_results, dashed_reference_label=dashed_reference_label, shaded_reference_results=shaded_reference_results, shaded_reference_label=shaded_reference_label, vlines=vlines, vline_colors=vline_colors, vline_styles=vline_styles, vline_labels=vline_labels, ylim=ylim, xlim=xlim, legend=legend, title=title, side_title=side_title, plot_percentages=plot_percentages) if(show): pyplot.show() return fig, ax #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% ######################################################## #@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@# #@ @# #@ EXTENDED SEIRS MODELS @# #@ @# #@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@# ######################################################## class ExtSEIRSNetworkModel(): """ A class to simulate the Extended SEIRS Stochastic Network Model =================================================== Params: G Network adjacency matrix (numpy array) or Networkx graph object. beta Rate of transmission (global interactions) beta_local Rate(s) of transmission between adjacent individuals (optional) beta_asym Rate of transmission (global interactions) beta_asym_local Rate(s) of transmission between adjacent individuals (optional) sigma Rate of progression to infectious state (inverse of latent period) lamda Rate of progression to infectious (a)symptomatic state (inverse of prodromal period) eta Rate of progression to hospitalized state (inverse of onset-to-admission period) gamma Rate of recovery for non-hospitalized symptomatic individuals (inverse of symptomatic infectious period) gamma_asym Rate of recovery for asymptomatic individuals (inverse of asymptomatic infectious period) gamma_H Rate of recovery for hospitalized symptomatic individuals (inverse of hospitalized infectious period) mu_H Rate of death for hospitalized individuals (inverse of admission-to-death period) xi Rate of re-susceptibility (upon recovery) mu_0 Rate of baseline death nu Rate of baseline birth a Probability of an infected individual remaining asymptomatic h Probability of a symptomatic individual being hospitalized f Probability of death for hospitalized individuals (case fatality rate) p Probability of individuals interacting with global population G_Q Quarantine adjacency matrix (numpy array) or Networkx graph object. beta_Q Rate of transmission for isolated individuals (global interactions) beta_Q_local Rate(s) of transmission (exposure) for adjacent isolated individuals (optional) sigma_Q Rate of progression to infectious state for isolated individuals lamda_Q Rate of progression to infectious (a)symptomatic state for isolated individuals eta_Q Rate of progression to hospitalized state for isolated individuals gamma_Q_sym Rate of recovery for non-hospitalized symptomatic individuals for isolated individuals gamma_Q_asym Rate of recovery for asymptomatic individuals for isolated individuals theta_E Rate of random testing for exposed individuals theta_pre Rate of random testing for infectious pre-symptomatic individuals theta_sym Rate of random testing for infectious symptomatic individuals theta_asym Rate of random testing for infectious asymptomatic individuals phi_E Rate of testing when a close contact has tested positive for exposed individuals phi_pre Rate of testing when a close contact has tested positive for infectious pre-symptomatic individuals phi_sym Rate of testing when a close contact has tested positive for infectious symptomatic individuals phi_asym Rate of testing when a close contact has tested positive for infectious asymptomatic individuals psi_E Probability of positive test for exposed individuals psi_pre Probability of positive test for infectious pre-symptomatic individuals psi_sym Probability of positive test for infectious symptomatic individuals psi_asym Probability of positive test for infectious asymptomatic individuals q Probability of isolated individuals interacting with global population isolation_time Time to remain in isolation upon positive test, self-isolation, etc. initE Initial number of exposed individuals initI_pre Initial number of infectious pre-symptomatic individuals initI_sym Initial number of infectious symptomatic individuals initI_asym Initial number of infectious asymptomatic individuals initH Initial number of hospitalized individuals initR Initial number of recovered individuals initF Initial number of infection-related fatalities initQ_S Initial number of isolated susceptible individuals initQ_E Initial number of isolated exposed individuals initQ_pre Initial number of isolated infectious pre-symptomatic individuals initQ_sym Initial number of isolated infectious symptomatic individuals initQ_asym Initial number of isolated infectious asymptomatic individuals initQ_R Initial number of isolated recovered individuals (all remaining nodes initialized susceptible) """ def __init__(self, G, beta, sigma, lamda, gamma, gamma_asym=None, eta=0, gamma_H=None, mu_H=0, alpha=1.0, xi=0, mu_0=0, nu=0, a=0, h=0, f=0, p=0, beta_local=None, beta_asym=None, beta_asym_local=None, beta_pairwise_mode='infected', delta=None, delta_pairwise_mode=None, G_Q=None, beta_Q=None, beta_Q_local=None, sigma_Q=None, lamda_Q=None, eta_Q=None, gamma_Q_sym=None, gamma_Q_asym=None, alpha_Q=None, delta_Q=None, theta_S=0, theta_E=0, theta_pre=0, theta_sym=0, theta_asym=0, phi_S=0, phi_E=0, phi_pre=0, phi_sym=0, phi_asym=0, psi_S=0, psi_E=1, psi_pre=1, psi_sym=1, psi_asym=1, q=0, isolation_time=14, initE=0, initI_pre=0, initI_sym=0, initI_asym=0, initH=0, initR=0, initF=0, initQ_S=0, initQ_E=0, initQ_pre=0, initQ_sym=0, initQ_asym=0, initQ_R=0, o=0, prevalence_ext=0, transition_mode='exponential_rates', node_groups=None, store_Xseries=False, seed=None): if(seed is not None): numpy.random.seed(seed) self.seed = seed #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Model Parameters: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ self.parameters = { 'G':G, 'G_Q':G_Q, 'beta':beta, 'sigma':sigma, 'lamda':lamda, 'gamma':gamma, 'eta':eta, 'gamma_asym':gamma_asym, 'gamma_H':gamma_H, 'mu_H':mu_H, 'xi':xi, 'mu_0':mu_0, 'nu':nu, 'a':a, 'h':h, 'f':f, 'p':p, 'beta_local':beta_local, 'beta_asym':beta_asym, 'beta_asym_local':beta_asym_local, 'beta_pairwise_mode':beta_pairwise_mode, 'alpha':alpha, 'delta':delta, 'delta_pairwise_mode':delta_pairwise_mode, 'lamda_Q':lamda_Q, 'beta_Q':beta_Q, 'beta_Q_local':beta_Q_local, 'alpha_Q':alpha_Q, 'sigma_Q':sigma_Q, 'eta_Q':eta_Q, 'gamma_Q_sym':gamma_Q_sym, 'gamma_Q_asym':gamma_Q_asym, 'delta_Q':delta_Q, 'theta_S':theta_S, 'theta_E':theta_E, 'theta_pre':theta_pre, 'theta_sym':theta_sym, 'theta_asym':theta_asym, 'phi_S':phi_S, 'phi_E':phi_E, 'phi_pre':phi_pre, 'phi_sym':phi_sym, 'phi_asym':phi_asym, 'psi_S':psi_S, 'psi_E':psi_E, 'psi_pre':psi_pre, 'psi_sym':psi_sym, 'psi_asym':psi_asym, 'q':q, 'isolation_time':isolation_time, 'initE':initE, 'initI_pre':initI_pre, 'initI_sym':initI_sym, 'initI_asym':initI_asym, 'initH':initH, 'initR':initR, 'initF':initF, 'initQ_S':initQ_S, 'initQ_E':initQ_E, 'initQ_pre':initQ_pre, 'initQ_sym':initQ_sym, 'initQ_asym':initQ_asym, 'initQ_R':initQ_R, 'o':o, 'prevalence_ext':prevalence_ext} self.update_parameters() #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Each node can undergo 4-6 transitions (sans vitality/re-susceptibility returns to S state), # so there are ~numNodes*6 events/timesteps expected; initialize numNodes*6 timestep slots to start # (will be expanded during run if needed for some reason) self.tseries = numpy.zeros(6*self.numNodes) self.numS = numpy.zeros(6*self.numNodes) self.numE = numpy.zeros(6*self.numNodes) self.numI_pre = numpy.zeros(6*self.numNodes) self.numI_sym = numpy.zeros(6*self.numNodes) self.numI_asym = numpy.zeros(6*self.numNodes) self.numH = numpy.zeros(6*self.numNodes) self.numR = numpy.zeros(6*self.numNodes) self.numF = numpy.zeros(6*self.numNodes) self.numQ_S = numpy.zeros(6*self.numNodes) self.numQ_E = numpy.zeros(6*self.numNodes) self.numQ_pre = numpy.zeros(6*self.numNodes) self.numQ_sym = numpy.zeros(6*self.numNodes) self.numQ_asym = numpy.zeros(6*self.numNodes) self.numQ_R = numpy.zeros(6*self.numNodes) self.N = numpy.zeros(6*self.numNodes) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Initialize Timekeeping: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ self.t = 0 self.tmax = 0 # will be set when run() is called self.tidx = 0 self.tseries[0] = 0 # Vectors holding the time that each node has been in a given state or in isolation: self.timer_state = numpy.zeros((self.numNodes,1)) self.timer_isolation = numpy.zeros(self.numNodes) self.isolationTime = isolation_time #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Initialize Counts of inidividuals with each state: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ self.numE[0] = int(initE) self.numI_pre[0] = int(initI_pre) self.numI_sym[0] = int(initI_sym) self.numI_asym[0] = int(initI_asym) self.numH[0] = int(initH) self.numR[0] = int(initR) self.numF[0] = int(initF) self.numQ_S[0] = int(initQ_S) self.numQ_E[0] = int(initQ_E) self.numQ_pre[0] = int(initQ_pre) self.numQ_sym[0] = int(initQ_sym) self.numQ_asym[0] = int(initQ_asym) self.numQ_R[0] = int(initQ_R) self.numS[0] = (self.numNodes - self.numE[0] - self.numI_pre[0] - self.numI_sym[0] - self.numI_asym[0] - self.numH[0] - self.numR[0] - self.numQ_S[0] - self.numQ_E[0] - self.numQ_pre[0] - self.numQ_sym[0] - self.numQ_asym[0] - self.numQ_R[0] - self.numF[0]) self.N[0] = self.numNodes - self.numF[0] #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Node states: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ self.S = 1 self.E = 2 self.I_pre = 3 self.I_sym = 4 self.I_asym = 5 self.H = 6 self.R = 7 self.F = 8 self.Q_S = 11 self.Q_E = 12 self.Q_pre = 13 self.Q_sym = 14 self.Q_asym = 15 self.Q_R = 17 self.X = numpy.array( [self.S]*int(self.numS[0]) + [self.E]*int(self.numE[0]) + [self.I_pre]*int(self.numI_pre[0]) + [self.I_sym]*int(self.numI_sym[0]) + [self.I_asym]*int(self.numI_asym[0]) + [self.H]*int(self.numH[0]) + [self.R]*int(self.numR[0]) + [self.F]*int(self.numF[0]) + [self.Q_S]*int(self.numQ_S[0]) + [self.Q_E]*int(self.numQ_E[0]) + [self.Q_pre]*int(self.numQ_pre[0]) + [self.Q_sym]*int(self.numQ_sym[0]) + [self.Q_asym]*int(self.numQ_asym[0]) + [self.Q_R]*int(self.numQ_R[0]) ).reshape((self.numNodes,1)) numpy.random.shuffle(self.X) self.store_Xseries = store_Xseries if(store_Xseries): self.Xseries = numpy.zeros(shape=(6*self.numNodes, self.numNodes), dtype='uint8') self.Xseries[0,:] = self.X.T self.transitions = { 'StoE': {'currentState':self.S, 'newState':self.E}, 'StoQS': {'currentState':self.S, 'newState':self.Q_S}, 'EtoIPRE': {'currentState':self.E, 'newState':self.I_pre}, 'EtoQE': {'currentState':self.E, 'newState':self.Q_E}, 'IPREtoISYM': {'currentState':self.I_pre, 'newState':self.I_sym}, 'IPREtoIASYM': {'currentState':self.I_pre, 'newState':self.I_asym}, 'IPREtoQPRE': {'currentState':self.I_pre, 'newState':self.Q_pre}, 'ISYMtoH': {'currentState':self.I_sym, 'newState':self.H}, 'ISYMtoR': {'currentState':self.I_sym, 'newState':self.R}, 'ISYMtoQSYM': {'currentState':self.I_sym, 'newState':self.Q_sym}, 'IASYMtoR': {'currentState':self.I_asym, 'newState':self.R}, 'IASYMtoQASYM': {'currentState':self.I_asym, 'newState':self.Q_asym}, 'HtoR': {'currentState':self.H, 'newState':self.R}, 'HtoF': {'currentState':self.H, 'newState':self.F}, 'RtoS': {'currentState':self.R, 'newState':self.S}, 'QStoQE': {'currentState':self.Q_S, 'newState':self.Q_E}, 'QEtoQPRE': {'currentState':self.Q_E, 'newState':self.Q_pre}, 'QPREtoQSYM': {'currentState':self.Q_pre, 'newState':self.Q_sym}, 'QPREtoQASYM': {'currentState':self.Q_pre, 'newState':self.Q_asym}, 'QSYMtoH': {'currentState':self.Q_sym, 'newState':self.H}, 'QSYMtoQR': {'currentState':self.Q_sym, 'newState':self.Q_R}, 'QASYMtoQR': {'currentState':self.Q_asym, 'newState':self.Q_R}, '_toS': {'currentState':True, 'newState':self.S}, } self.transition_mode = transition_mode #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Initialize other node metadata: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ self.tested = numpy.array([False]*self.numNodes).reshape((self.numNodes,1)) self.positive = numpy.array([False]*self.numNodes).reshape((self.numNodes,1)) self.numTested = numpy.zeros(6*self.numNodes) self.numPositive = numpy.zeros(6*self.numNodes) self.testedInCurrentState = numpy.array([False]*self.numNodes).reshape((self.numNodes,1)) self.infectionsLog = [] #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Initialize node subgroup data series: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ self.nodeGroupData = None if(node_groups): self.nodeGroupData = {} for groupName, nodeList in node_groups.items(): self.nodeGroupData[groupName] = {'nodes': numpy.array(nodeList), 'mask': numpy.isin(range(self.numNodes), nodeList).reshape((self.numNodes,1))} self.nodeGroupData[groupName]['numS'] = numpy.zeros(6*self.numNodes) self.nodeGroupData[groupName]['numE'] = numpy.zeros(6*self.numNodes) self.nodeGroupData[groupName]['numI_pre'] = numpy.zeros(6*self.numNodes) self.nodeGroupData[groupName]['numI_sym'] = numpy.zeros(6*self.numNodes) self.nodeGroupData[groupName]['numI_asym'] = numpy.zeros(6*self.numNodes) self.nodeGroupData[groupName]['numH'] = numpy.zeros(6*self.numNodes) self.nodeGroupData[groupName]['numR'] = numpy.zeros(6*self.numNodes) self.nodeGroupData[groupName]['numF'] = numpy.zeros(6*self.numNodes) self.nodeGroupData[groupName]['numQ_S'] = numpy.zeros(6*self.numNodes) self.nodeGroupData[groupName]['numQ_E'] = numpy.zeros(6*self.numNodes) self.nodeGroupData[groupName]['numQ_pre'] = numpy.zeros(6*self.numNodes) self.nodeGroupData[groupName]['numQ_sym'] = numpy.zeros(6*self.numNodes) self.nodeGroupData[groupName]['numQ_asym'] = numpy.zeros(6*self.numNodes) self.nodeGroupData[groupName]['numQ_R'] = numpy.zeros(6*self.numNodes) self.nodeGroupData[groupName]['N'] = numpy.zeros(6*self.numNodes) self.nodeGroupData[groupName]['numPositive'] = numpy.zeros(6*self.numNodes) self.nodeGroupData[groupName]['numTested'] = numpy.zeros(6*self.numNodes) self.nodeGroupData[groupName]['numS'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.S) self.nodeGroupData[groupName]['numE'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.E) self.nodeGroupData[groupName]['numI_pre'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.I_pre) self.nodeGroupData[groupName]['numI_sym'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.I_sym) self.nodeGroupData[groupName]['numI_asym'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.I_asym) self.nodeGroupData[groupName]['numH'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.H) self.nodeGroupData[groupName]['numR'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.R) self.nodeGroupData[groupName]['numF'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.F) self.nodeGroupData[groupName]['numQ_S'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_E) self.nodeGroupData[groupName]['numQ_E'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_E) self.nodeGroupData[groupName]['numQ_pre'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_pre) self.nodeGroupData[groupName]['numQ_I_sym'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_I_sym) self.nodeGroupData[groupName]['numQ_I_asym'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_I_asym) self.nodeGroupData[groupName]['numQ_R'][0] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_E) self.nodeGroupData[groupName]['N'][0] = self.numNodes - self.numF[0] #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def update_parameters(self): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Model graphs: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ self.G = self.parameters['G'] # Adjacency matrix: if type(self.G)==numpy.ndarray: self.A = scipy.sparse.csr_matrix(self.G) elif type(self.G)==networkx.classes.graph.Graph: self.A = networkx.adj_matrix(self.G) # adj_matrix gives scipy.sparse csr_matrix else: raise BaseException("Input an adjacency matrix or networkx object only.") self.numNodes = int(self.A.shape[1]) self.degree = numpy.asarray(self.node_degrees(self.A)).astype(float) #---------------------------------------- if(self.parameters['G_Q'] is None): self.G_Q = self.G # If no Q graph is provided, use G in its place else: self.G_Q = self.parameters['G_Q'] # Quarantine Adjacency matrix: if type(self.G_Q)==numpy.ndarray: self.A_Q = scipy.sparse.csr_matrix(self.G_Q) elif type(self.G_Q)==networkx.classes.graph.Graph: self.A_Q = networkx.adj_matrix(self.G_Q) # adj_matrix gives scipy.sparse csr_matrix else: raise BaseException("Input an adjacency matrix or networkx object only.") self.numNodes_Q = int(self.A_Q.shape[1]) self.degree_Q = numpy.asarray(self.node_degrees(self.A_Q)).astype(float) #---------------------------------------- assert(self.numNodes == self.numNodes_Q), "The normal and quarantine adjacency graphs must be of the same size." #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Model parameters: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ self.beta = numpy.array(self.parameters['beta']).reshape((self.numNodes, 1)) if isinstance(self.parameters['beta'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['beta'], shape=(self.numNodes,1)) self.beta_asym = (numpy.array(self.parameters['beta_asym']).reshape((self.numNodes, 1)) if isinstance(self.parameters['beta_asym'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['beta_asym'], shape=(self.numNodes,1))) if self.parameters['beta_asym'] is not None else self.beta self.sigma = numpy.array(self.parameters['sigma']).reshape((self.numNodes, 1)) if isinstance(self.parameters['sigma'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['sigma'], shape=(self.numNodes,1)) self.lamda = numpy.array(self.parameters['lamda']).reshape((self.numNodes, 1)) if isinstance(self.parameters['lamda'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['lamda'], shape=(self.numNodes,1)) self.gamma = numpy.array(self.parameters['gamma']).reshape((self.numNodes, 1)) if isinstance(self.parameters['gamma'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['gamma'], shape=(self.numNodes,1)) self.eta = numpy.array(self.parameters['eta']).reshape((self.numNodes, 1)) if isinstance(self.parameters['eta'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['eta'], shape=(self.numNodes,1)) self.gamma_asym = (numpy.array(self.parameters['gamma_asym']).reshape((self.numNodes, 1)) if isinstance(self.parameters['gamma_asym'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['gamma_asym'], shape=(self.numNodes,1))) if self.parameters['gamma_asym'] is not None else self.gamma self.gamma_H = (numpy.array(self.parameters['gamma_H']).reshape((self.numNodes, 1)) if isinstance(self.parameters['gamma_H'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['gamma_H'], shape=(self.numNodes,1))) if self.parameters['gamma_H'] is not None else self.gamma self.mu_H = numpy.array(self.parameters['mu_H']).reshape((self.numNodes, 1)) if isinstance(self.parameters['mu_H'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['mu_H'], shape=(self.numNodes,1)) self.alpha = numpy.array(self.parameters['alpha']).reshape((self.numNodes, 1)) if isinstance(self.parameters['alpha'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['alpha'], shape=(self.numNodes,1)) self.xi = numpy.array(self.parameters['xi']).reshape((self.numNodes, 1)) if isinstance(self.parameters['xi'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['xi'], shape=(self.numNodes,1)) self.mu_0 = numpy.array(self.parameters['mu_0']).reshape((self.numNodes, 1)) if isinstance(self.parameters['mu_0'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['mu_0'], shape=(self.numNodes,1)) self.nu = numpy.array(self.parameters['nu']).reshape((self.numNodes, 1)) if isinstance(self.parameters['nu'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['nu'], shape=(self.numNodes,1)) self.a = numpy.array(self.parameters['a']).reshape((self.numNodes, 1)) if isinstance(self.parameters['a'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['a'], shape=(self.numNodes,1)) self.h = numpy.array(self.parameters['h']).reshape((self.numNodes, 1)) if isinstance(self.parameters['h'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['h'], shape=(self.numNodes,1)) self.f = numpy.array(self.parameters['f']).reshape((self.numNodes, 1)) if isinstance(self.parameters['f'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['f'], shape=(self.numNodes,1)) self.p = numpy.array(self.parameters['p']).reshape((self.numNodes, 1)) if isinstance(self.parameters['p'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['p'], shape=(self.numNodes,1)) self.o = numpy.array(self.parameters['o']).reshape((self.numNodes, 1)) if isinstance(self.parameters['o'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['o'], shape=(self.numNodes,1)) self.rand_a = numpy.random.rand(self.a.shape[0], self.a.shape[1]) self.rand_h = numpy.random.rand(self.h.shape[0], self.h.shape[1]) self.rand_f = numpy.random.rand(self.f.shape[0], self.f.shape[1]) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # External infection introduction variables: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ self.prevalence_ext = numpy.array(self.parameters['prevalence_ext']).reshape((self.numNodes, 1)) if isinstance(self.parameters['prevalence_ext'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['prevalence_ext'], shape=(self.numNodes,1)) #---------------------------------------- # Testing-related parameters: #---------------------------------------- self.beta_Q = (numpy.array(self.parameters['beta_Q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['beta_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['beta_Q'], shape=(self.numNodes,1))) if self.parameters['beta_Q'] is not None else self.beta self.sigma_Q = (numpy.array(self.parameters['sigma_Q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['sigma_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['sigma_Q'], shape=(self.numNodes,1))) if self.parameters['sigma_Q'] is not None else self.sigma self.lamda_Q = (numpy.array(self.parameters['lamda_Q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['lamda_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['lamda_Q'], shape=(self.numNodes,1))) if self.parameters['lamda_Q'] is not None else self.lamda self.gamma_Q_sym = (numpy.array(self.parameters['gamma_Q_sym']).reshape((self.numNodes, 1)) if isinstance(self.parameters['gamma_Q_sym'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['gamma_Q_sym'], shape=(self.numNodes,1))) if self.parameters['gamma_Q_sym'] is not None else self.gamma self.gamma_Q_asym = (numpy.array(self.parameters['gamma_Q_asym']).reshape((self.numNodes, 1)) if isinstance(self.parameters['gamma_Q_asym'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['gamma_Q_asym'], shape=(self.numNodes,1))) if self.parameters['gamma_Q_asym'] is not None else self.gamma self.eta_Q = (numpy.array(self.parameters['eta_Q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['eta_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['eta_Q'], shape=(self.numNodes,1))) if self.parameters['eta_Q'] is not None else self.eta self.alpha_Q = (numpy.array(self.parameters['alpha_Q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['alpha_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['alpha_Q'], shape=(self.numNodes,1))) if self.parameters['alpha_Q'] is not None else self.alpha self.theta_S = numpy.array(self.parameters['theta_S']).reshape((self.numNodes, 1)) if isinstance(self.parameters['theta_S'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['theta_S'], shape=(self.numNodes,1)) self.theta_E = numpy.array(self.parameters['theta_E']).reshape((self.numNodes, 1)) if isinstance(self.parameters['theta_E'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['theta_E'], shape=(self.numNodes,1)) self.theta_pre = numpy.array(self.parameters['theta_pre']).reshape((self.numNodes, 1)) if isinstance(self.parameters['theta_pre'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['theta_pre'], shape=(self.numNodes,1)) self.theta_sym = numpy.array(self.parameters['theta_sym']).reshape((self.numNodes, 1)) if isinstance(self.parameters['theta_sym'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['theta_sym'], shape=(self.numNodes,1)) self.theta_asym = numpy.array(self.parameters['theta_asym']).reshape((self.numNodes, 1)) if isinstance(self.parameters['theta_asym'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['theta_asym'], shape=(self.numNodes,1)) self.phi_S = numpy.array(self.parameters['phi_S']).reshape((self.numNodes, 1)) if isinstance(self.parameters['phi_S'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['phi_S'], shape=(self.numNodes,1)) self.phi_E = numpy.array(self.parameters['phi_E']).reshape((self.numNodes, 1)) if isinstance(self.parameters['phi_E'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['phi_E'], shape=(self.numNodes,1)) self.phi_pre = numpy.array(self.parameters['phi_pre']).reshape((self.numNodes, 1)) if isinstance(self.parameters['phi_pre'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['phi_pre'], shape=(self.numNodes,1)) self.phi_sym = numpy.array(self.parameters['phi_sym']).reshape((self.numNodes, 1)) if isinstance(self.parameters['phi_sym'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['phi_sym'], shape=(self.numNodes,1)) self.phi_asym = numpy.array(self.parameters['phi_asym']).reshape((self.numNodes, 1)) if isinstance(self.parameters['phi_asym'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['phi_asym'], shape=(self.numNodes,1)) self.psi_S = numpy.array(self.parameters['psi_S']).reshape((self.numNodes, 1)) if isinstance(self.parameters['psi_S'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['psi_S'], shape=(self.numNodes,1)) self.psi_E = numpy.array(self.parameters['psi_E']).reshape((self.numNodes, 1)) if isinstance(self.parameters['psi_E'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['psi_E'], shape=(self.numNodes,1)) self.psi_pre = numpy.array(self.parameters['psi_pre']).reshape((self.numNodes, 1)) if isinstance(self.parameters['psi_pre'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['psi_pre'], shape=(self.numNodes,1)) self.psi_sym = numpy.array(self.parameters['psi_sym']).reshape((self.numNodes, 1)) if isinstance(self.parameters['psi_sym'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['psi_sym'], shape=(self.numNodes,1)) self.psi_asym = numpy.array(self.parameters['psi_asym']).reshape((self.numNodes, 1)) if isinstance(self.parameters['psi_asym'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['psi_asym'], shape=(self.numNodes,1)) self.q = numpy.array(self.parameters['q']).reshape((self.numNodes, 1)) if isinstance(self.parameters['q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['q'], shape=(self.numNodes,1)) #---------------------------------------- self.beta_pairwise_mode = self.parameters['beta_pairwise_mode'] #---------------------------------------- # Global transmission parameters: #---------------------------------------- if(self.beta_pairwise_mode == 'infected' or self.beta_pairwise_mode is None): self.beta_global = numpy.full_like(self.beta, fill_value=numpy.mean(self.beta)) self.beta_Q_global = numpy.full_like(self.beta_Q, fill_value=numpy.mean(self.beta_Q)) self.beta_asym_global = numpy.full_like(self.beta_asym, fill_value=numpy.mean(self.beta_asym)) elif(self.beta_pairwise_mode == 'infectee'): self.beta_global = self.beta self.beta_Q_global = self.beta_Q self.beta_asym_global = self.beta_asym elif(self.beta_pairwise_mode == 'min'): self.beta_global = numpy.minimum(self.beta, numpy.mean(beta)) self.beta_Q_global = numpy.minimum(self.beta_Q, numpy.mean(beta_Q)) self.beta_asym_global = numpy.minimum(self.beta_asym, numpy.mean(beta_asym)) elif(self.beta_pairwise_mode == 'max'): self.beta_global = numpy.maximum(self.beta, numpy.mean(beta)) self.beta_Q_global = numpy.maximum(self.beta_Q, numpy.mean(beta_Q)) self.beta_asym_global = numpy.maximum(self.beta_asym, numpy.mean(beta_asym)) elif(self.beta_pairwise_mode == 'mean'): self.beta_global = (self.beta + numpy.full_like(self.beta, fill_value=numpy.mean(self.beta)))/2 self.beta_Q_global = (self.beta_Q + numpy.full_like(self.beta_Q, fill_value=numpy.mean(self.beta_Q)))/2 self.beta_asym_global = (self.beta_asym + numpy.full_like(self.beta_asym, fill_value=numpy.mean(self.beta_asym)))/2 #---------------------------------------- # Local transmission parameters: #---------------------------------------- self.beta_local = self.beta if self.parameters['beta_local'] is None else numpy.array(self.parameters['beta_local']) if isinstance(self.parameters['beta_local'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['beta_local'], shape=(self.numNodes,1)) self.beta_Q_local = self.beta_Q if self.parameters['beta_Q_local'] is None else numpy.array(self.parameters['beta_Q_local']) if isinstance(self.parameters['beta_Q_local'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['beta_Q_local'], shape=(self.numNodes,1)) self.beta_asym_local = None if self.parameters['beta_asym_local'] is None else numpy.array(self.parameters['beta_asym_local']) if isinstance(self.parameters['beta_asym_local'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['beta_asym_local'], shape=(self.numNodes,1)) #---------------------------------------- if(self.beta_local.ndim == 2 and self.beta_local.shape[0] == self.numNodes and self.beta_local.shape[1] == self.numNodes): self.A_beta_pairwise = self.beta_local elif((self.beta_local.ndim == 1 and self.beta_local.shape[0] == self.numNodes) or (self.beta_local.ndim == 2 and (self.beta_local.shape[0] == self.numNodes or self.beta_local.shape[1] == self.numNodes))): self.beta_local = self.beta_local.reshape((self.numNodes,1)) # Pre-multiply beta values by the adjacency matrix ("transmission weight connections") A_beta_pairwise_byInfected = scipy.sparse.csr_matrix.multiply(self.A, self.beta_local.T).tocsr() A_beta_pairwise_byInfectee = scipy.sparse.csr_matrix.multiply(self.A, self.beta_local).tocsr() #------------------------------ # Compute the effective pairwise beta values as a function of the infected/infectee pair: if(self.beta_pairwise_mode == 'infected'): self.A_beta_pairwise = A_beta_pairwise_byInfected elif(self.beta_pairwise_mode == 'infectee'): self.A_beta_pairwise = A_beta_pairwise_byInfectee elif(self.beta_pairwise_mode == 'min'): self.A_beta_pairwise = scipy.sparse.csr_matrix.minimum(A_beta_pairwise_byInfected, A_beta_pairwise_byInfectee) elif(self.beta_pairwise_mode == 'max'): self.A_beta_pairwise = scipy.sparse.csr_matrix.maximum(A_beta_pairwise_byInfected, A_beta_pairwise_byInfectee) elif(self.beta_pairwise_mode == 'mean' or self.beta_pairwise_mode is None): self.A_beta_pairwise = (A_beta_pairwise_byInfected + A_beta_pairwise_byInfectee)/2 else: print("Unrecognized beta_pairwise_mode value (support for 'infected', 'infectee', 'min', 'max', and 'mean').") else: print("Invalid values given for beta_local (expected 1xN list/array or NxN 2d array)") #---------------------------------------- if(self.beta_Q_local.ndim == 2 and self.beta_Q_local.shape[0] == self.numNodes and self.beta_Q_local.shape[1] == self.numNodes): self.A_Q_beta_Q_pairwise = self.beta_Q_local elif((self.beta_Q_local.ndim == 1 and self.beta_Q_local.shape[0] == self.numNodes) or (self.beta_Q_local.ndim == 2 and (self.beta_Q_local.shape[0] == self.numNodes or self.beta_Q_local.shape[1] == self.numNodes))): self.beta_Q_local = self.beta_Q_local.reshape((self.numNodes,1)) # Pre-multiply beta_Q values by the isolation adjacency matrix ("transmission weight connections") A_Q_beta_Q_pairwise_byInfected = scipy.sparse.csr_matrix.multiply(self.A_Q, self.beta_Q_local.T).tocsr() A_Q_beta_Q_pairwise_byInfectee = scipy.sparse.csr_matrix.multiply(self.A_Q, self.beta_Q_local).tocsr() #------------------------------ # Compute the effective pairwise beta values as a function of the infected/infectee pair: if(self.beta_pairwise_mode == 'infected'): self.A_Q_beta_Q_pairwise = A_Q_beta_Q_pairwise_byInfected elif(self.beta_pairwise_mode == 'infectee'): self.A_Q_beta_Q_pairwise = A_Q_beta_Q_pairwise_byInfectee elif(self.beta_pairwise_mode == 'min'): self.A_Q_beta_Q_pairwise = scipy.sparse.csr_matrix.minimum(A_Q_beta_Q_pairwise_byInfected, A_Q_beta_Q_pairwise_byInfectee) elif(self.beta_pairwise_mode == 'max'): self.A_Q_beta_Q_pairwise = scipy.sparse.csr_matrix.maximum(A_Q_beta_Q_pairwise_byInfected, A_Q_beta_Q_pairwise_byInfectee) elif(self.beta_pairwise_mode == 'mean' or self.beta_pairwise_mode is None): self.A_Q_beta_Q_pairwise = (A_Q_beta_Q_pairwise_byInfected + A_Q_beta_Q_pairwise_byInfectee)/2 else: print("Unrecognized beta_pairwise_mode value (support for 'infected', 'infectee', 'min', 'max', and 'mean').") else: print("Invalid values given for beta_Q_local (expected 1xN list/array or NxN 2d array)") #---------------------------------------- if(self.beta_asym_local is None): self.A_beta_asym_pairwise = None elif(self.beta_asym_local.ndim == 2 and self.beta_asym_local.shape[0] == self.numNodes and self.beta_asym_local.shape[1] == self.numNodes): self.A_beta_asym_pairwise = self.beta_asym_local elif((self.beta_asym_local.ndim == 1 and self.beta_asym_local.shape[0] == self.numNodes) or (self.beta_asym_local.ndim == 2 and (self.beta_asym_local.shape[0] == self.numNodes or self.beta_asym_local.shape[1] == self.numNodes))): self.beta_asym_local = self.beta_asym_local.reshape((self.numNodes,1)) # Pre-multiply beta_asym values by the adjacency matrix ("transmission weight connections") A_beta_asym_pairwise_byInfected = scipy.sparse.csr_matrix.multiply(self.A, self.beta_asym_local.T).tocsr() A_beta_asym_pairwise_byInfectee = scipy.sparse.csr_matrix.multiply(self.A, self.beta_asym_local).tocsr() #------------------------------ # Compute the effective pairwise beta values as a function of the infected/infectee pair: if(self.beta_pairwise_mode == 'infected'): self.A_beta_asym_pairwise = A_beta_asym_pairwise_byInfected elif(self.beta_pairwise_mode == 'infectee'): self.A_beta_asym_pairwise = A_beta_asym_pairwise_byInfectee elif(self.beta_pairwise_mode == 'min'): self.A_beta_asym_pairwise = scipy.sparse.csr_matrix.minimum(A_beta_asym_pairwise_byInfected, A_beta_asym_pairwise_byInfectee) elif(self.beta_pairwise_mode == 'max'): self.A_beta_asym_pairwise = scipy.sparse.csr_matrix.maximum(A_beta_asym_pairwise_byInfected, A_beta_asym_pairwise_byInfectee) elif(self.beta_pairwise_mode == 'mean' or self.beta_pairwise_mode is None): self.A_beta_asym_pairwise = (A_beta_asym_pairwise_byInfected + A_beta_asym_pairwise_byInfectee)/2 else: print("Unrecognized beta_pairwise_mode value (support for 'infected', 'infectee', 'min', 'max', and 'mean').") else: print("Invalid values given for beta_asym_local (expected 1xN list/array or NxN 2d array)") #---------------------------------------- # Degree-based transmission scaling parameters: #---------------------------------------- self.delta_pairwise_mode = self.parameters['delta_pairwise_mode'] with numpy.errstate(divide='ignore'): # ignore log(0) warning, then convert log(0) = -inf -> 0.0 self.delta = numpy.log(self.degree)/numpy.log(numpy.mean(self.degree)) if self.parameters['delta'] is None else numpy.array(self.parameters['delta']) if isinstance(self.parameters['delta'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['delta'], shape=(self.numNodes,1)) self.delta_Q = numpy.log(self.degree_Q)/numpy.log(numpy.mean(self.degree_Q)) if self.parameters['delta_Q'] is None else numpy.array(self.parameters['delta_Q']) if isinstance(self.parameters['delta_Q'], (list, numpy.ndarray)) else numpy.full(fill_value=self.parameters['delta_Q'], shape=(self.numNodes,1)) self.delta[numpy.isneginf(self.delta)] = 0.0 self.delta_Q[numpy.isneginf(self.delta_Q)] = 0.0 #---------------------------------------- if(self.delta.ndim == 2 and self.delta.shape[0] == self.numNodes and self.delta.shape[1] == self.numNodes): self.A_delta_pairwise = self.delta elif((self.delta.ndim == 1 and self.delta.shape[0] == self.numNodes) or (self.delta.ndim == 2 and (self.delta.shape[0] == self.numNodes or self.delta.shape[1] == self.numNodes))): self.delta = self.delta.reshape((self.numNodes,1)) # Pre-multiply delta values by the adjacency matrix ("transmission weight connections") A_delta_pairwise_byInfected = scipy.sparse.csr_matrix.multiply(self.A, self.delta.T).tocsr() A_delta_pairwise_byInfectee = scipy.sparse.csr_matrix.multiply(self.A, self.delta).tocsr() #------------------------------ # Compute the effective pairwise delta values as a function of the infected/infectee pair: if(self.delta_pairwise_mode == 'infected'): self.A_delta_pairwise = A_delta_pairwise_byInfected elif(self.delta_pairwise_mode == 'infectee'): self.A_delta_pairwise = A_delta_pairwise_byInfectee elif(self.delta_pairwise_mode == 'min'): self.A_delta_pairwise = scipy.sparse.csr_matrix.minimum(A_delta_pairwise_byInfected, A_delta_pairwise_byInfectee) elif(self.delta_pairwise_mode == 'max'): self.A_delta_pairwise = scipy.sparse.csr_matrix.maximum(A_delta_pairwise_byInfected, A_delta_pairwise_byInfectee) elif(self.delta_pairwise_mode == 'mean'): self.A_delta_pairwise = (A_delta_pairwise_byInfected + A_delta_pairwise_byInfectee)/2 elif(self.delta_pairwise_mode is None): self.A_delta_pairwise = self.A else: print("Unrecognized delta_pairwise_mode value (support for 'infected', 'infectee', 'min', 'max', and 'mean').") else: print("Invalid values given for delta (expected 1xN list/array or NxN 2d array)") #---------------------------------------- if(self.delta_Q.ndim == 2 and self.delta_Q.shape[0] == self.numNodes and self.delta_Q.shape[1] == self.numNodes): self.A_Q_delta_Q_pairwise = self.delta_Q elif((self.delta_Q.ndim == 1 and self.delta_Q.shape[0] == self.numNodes) or (self.delta_Q.ndim == 2 and (self.delta_Q.shape[0] == self.numNodes or self.delta_Q.shape[1] == self.numNodes))): self.delta_Q = self.delta_Q.reshape((self.numNodes,1)) # Pre-multiply delta_Q values by the isolation adjacency matrix ("transmission weight connections") A_Q_delta_Q_pairwise_byInfected = scipy.sparse.csr_matrix.multiply(self.A_Q, self.delta_Q).tocsr() A_Q_delta_Q_pairwise_byInfectee = scipy.sparse.csr_matrix.multiply(self.A_Q, self.delta_Q.T).tocsr() #------------------------------ # Compute the effective pairwise delta values as a function of the infected/infectee pair: if(self.delta_pairwise_mode == 'infected'): self.A_Q_delta_Q_pairwise = A_Q_delta_Q_pairwise_byInfected elif(self.delta_pairwise_mode == 'infectee'): self.A_Q_delta_Q_pairwise = A_Q_delta_Q_pairwise_byInfectee elif(self.delta_pairwise_mode == 'min'): self.A_Q_delta_Q_pairwise = scipy.sparse.csr_matrix.minimum(A_Q_delta_Q_pairwise_byInfected, A_Q_delta_Q_pairwise_byInfectee) elif(self.delta_pairwise_mode == 'max'): self.A_Q_delta_Q_pairwise = scipy.sparse.csr_matrix.maximum(A_Q_delta_Q_pairwise_byInfected, A_Q_delta_Q_pairwise_byInfectee) elif(self.delta_pairwise_mode == 'mean'): self.A_Q_delta_Q_pairwise = (A_Q_delta_Q_pairwise_byInfected + A_Q_delta_Q_pairwise_byInfectee)/2 elif(self.delta_pairwise_mode is None): self.A_Q_delta_Q_pairwise = self.A else: print("Unrecognized delta_pairwise_mode value (support for 'infected', 'infectee', 'min', 'max', and 'mean').") else: print("Invalid values given for delta_Q (expected 1xN list/array or NxN 2d array)") #---------------------------------------- # Pre-calculate the pairwise delta*beta values: #---------------------------------------- self.A_deltabeta = scipy.sparse.csr_matrix.multiply(self.A_delta_pairwise, self.A_beta_pairwise) self.A_Q_deltabeta_Q = scipy.sparse.csr_matrix.multiply(self.A_Q_delta_Q_pairwise, self.A_Q_beta_Q_pairwise) if(self.A_beta_asym_pairwise is not None): self.A_deltabeta_asym = scipy.sparse.csr_matrix.multiply(self.A_delta_pairwise, self.A_beta_asym_pairwise) else: self.A_deltabeta_asym = None #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def node_degrees(self, Amat): return Amat.sum(axis=0).reshape(self.numNodes,1) # sums of adj matrix cols #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_susceptible(self, t_idx=None): if(t_idx is None): return (self.numS[:] + self.numQ_S[:]) else: return (self.numS[t_idx] + self.numQ_S[t_idx]) #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_infected(self, t_idx=None): if(t_idx is None): return (self.numE[:] + self.numI_pre[:] + self.numI_sym[:] + self.numI_asym[:] + self.numH[:] + self.numQ_E[:] + self.numQ_pre[:] + self.numQ_sym[:] + self.numQ_asym[:]) else: return (self.numE[t_idx] + self.numI_pre[t_idx] + self.numI_sym[t_idx] + self.numI_asym[t_idx] + self.numH[t_idx] + self.numQ_E[t_idx] + self.numQ_pre[t_idx] + self.numQ_sym[t_idx] + self.numQ_asym[t_idx]) #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_isolated(self, t_idx=None): if(t_idx is None): return (self.numQ_S[:] + self.numQ_E[:] + self.numQ_pre[:] + self.numQ_sym[:] + self.numQ_asym[:] + self.numQ_R[:]) else: return (self.numQ_S[t_idx] + self.numQ_E[t_idx] + self.numQ_pre[t_idx] + self.numQ_sym[t_idx] + self.numQ_asym[t_idx] + self.numQ_R[t_idx]) #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_tested(self, t_idx=None): if(t_idx is None): return (self.numTested[:]) else: return (self.numTested[t_idx]) #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_positive(self, t_idx=None): if(t_idx is None): return (self.numPositive[:]) else: return (self.numPositive[t_idx]) #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def total_num_recovered(self, t_idx=None): if(t_idx is None): return (self.numR[:] + self.numQ_R[:]) else: return (self.numR[t_idx] + self.numQ_R[t_idx]) #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def calc_propensities(self): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Pre-calculate matrix multiplication terms that may be used in multiple propensity calculations, # and check to see if their computation is necessary before doing the multiplication #------------------------------------ self.transmissionTerms_I = numpy.zeros(shape=(self.numNodes,1)) if(numpy.any(self.numI_sym[self.tidx]) or numpy.any(self.numI_asym[self.tidx]) or numpy.any(self.numI_pre[self.tidx])): if(self.A_deltabeta_asym is not None): self.transmissionTerms_sym = numpy.asarray(scipy.sparse.csr_matrix.dot(self.A_deltabeta, self.X==self.I_sym)) self.transmissionTerms_asym = numpy.asarray(scipy.sparse.csr_matrix.dot(self.A_deltabeta_asym, ((self.X==self.I_pre)|(self.X==self.I_asym)))) self.transmissionTerms_I = self.transmissionTerms_sym+self.transmissionTerms_asym else: self.transmissionTerms_I = numpy.asarray(scipy.sparse.csr_matrix.dot(self.A_deltabeta, ((self.X==self.I_sym)|(self.X==self.I_pre)|(self.X==self.I_asym)))) #------------------------------------ self.transmissionTerms_Q = numpy.zeros(shape=(self.numNodes,1)) if(numpy.any(self.numQ_pre[self.tidx]) or numpy.any(self.numQ_sym[self.tidx]) or numpy.any(self.numQ_asym[self.tidx])): self.transmissionTerms_Q = numpy.asarray(scipy.sparse.csr_matrix.dot(self.A_Q_deltabeta_Q, ((self.X==self.Q_pre)|(self.X==self.Q_sym)|(self.X==self.Q_asym)))) #------------------------------------ self.transmissionTerms_IQ = numpy.zeros(shape=(self.numNodes,1)) if(numpy.any(self.numQ_S[self.tidx]) and (numpy.any(self.numI_sym[self.tidx]) or numpy.any(self.numI_asym[self.tidx]) or numpy.any(self.numI_pre[self.tidx]))): self.transmissionTerms_IQ = numpy.asarray(scipy.sparse.csr_matrix.dot(self.A_Q_deltabeta_Q, ((self.X==self.I_sym)|(self.X==self.I_pre)|(self.X==self.I_asym)))) #------------------------------------ numContacts_Q = numpy.zeros(shape=(self.numNodes,1)) if(numpy.any(self.positive) and (numpy.any(self.phi_S) or numpy.any(self.phi_E) or numpy.any(self.phi_pre) or numpy.any(self.phi_sym) or numpy.any(self.phi_asym))): numContacts_Q = numpy.asarray(scipy.sparse.csr_matrix.dot(self.A, ((self.positive)&(self.X!=self.R)&(self.X!=self.Q_R)&(self.X!=self.F)))) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ propensities_StoE = ( self.alpha * (self.o*(self.beta_global*self.prevalence_ext) + (1-self.o)*( self.p*((self.beta_global*self.numI_sym[self.tidx] + self.beta_asym_global*(self.numI_pre[self.tidx] + self.numI_asym[self.tidx]) + self.q*self.beta_Q_global*(self.numQ_pre[self.tidx] + self.numQ_sym[self.tidx] + self.numQ_asym[self.tidx]))/self.N[self.tidx]) + (1-self.p)*(numpy.divide(self.transmissionTerms_I, self.degree, out=numpy.zeros_like(self.degree), where=self.degree!=0) + numpy.divide(self.transmissionTerms_Q, self.degree_Q, out=numpy.zeros_like(self.degree_Q), where=self.degree_Q!=0)))) )*(self.X==self.S) propensities_QStoQE = numpy.zeros_like(propensities_StoE) if(numpy.any(self.X==self.Q_S)): propensities_QStoQE = ( self.alpha_Q * (self.o*(self.q*self.beta_global*self.prevalence_ext) + (1-self.o)*( self.p*(self.q*(self.beta_global*self.numI_sym[self.tidx] + self.beta_asym_global*(self.numI_pre[self.tidx] + self.numI_asym[self.tidx]) + self.beta_Q_global*(self.numQ_pre[self.tidx] + self.numQ_sym[self.tidx] + self.numQ_asym[self.tidx]))/self.N[self.tidx]) + (1-self.p)*(numpy.divide(self.transmissionTerms_IQ+self.transmissionTerms_Q, self.degree_Q, out=numpy.zeros_like(self.degree_Q), where=self.degree_Q!=0)))) )*(self.X==self.Q_S) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(self.transition_mode == 'time_in_state'): propensities_EtoIPRE = 1e5 * ((self.X==self.E) & numpy.greater(self.timer_state, 1/self.sigma)) propensities_IPREtoISYM = 1e5 * ((self.X==self.I_pre) & numpy.greater(self.timer_state, 1/self.lamda) & numpy.greater_equal(self.rand_a, self.a)) propensities_IPREtoIASYM = 1e5 * ((self.X==self.I_pre) & numpy.greater(self.timer_state, 1/self.lamda) & numpy.less(self.rand_a, self.a)) propensities_ISYMtoR = 1e5 * ((self.X==self.I_sym) & numpy.greater(self.timer_state, 1/self.gamma) & numpy.greater_equal(self.rand_h, self.h)) propensities_ISYMtoH = 1e5 * ((self.X==self.I_sym) & numpy.greater(self.timer_state, 1/self.eta) & numpy.less(self.rand_h, self.h)) propensities_IASYMtoR = 1e5 * ((self.X==self.I_asym) & numpy.greater(self.timer_state, 1/self.gamma)) propensities_HtoR = 1e5 * ((self.X==self.H) & numpy.greater(self.timer_state, 1/self.gamma_H) & numpy.greater_equal(self.rand_f, self.f)) propensities_HtoF = 1e5 * ((self.X==self.H) & numpy.greater(self.timer_state, 1/self.mu_H) & numpy.less(self.rand_f, self.f)) propensities_StoQS = numpy.zeros_like(propensities_StoE) propensities_EtoQE = numpy.zeros_like(propensities_StoE) propensities_IPREtoQPRE = numpy.zeros_like(propensities_StoE) propensities_ISYMtoQSYM = numpy.zeros_like(propensities_StoE) propensities_IASYMtoQASYM = numpy.zeros_like(propensities_StoE) propensities_QEtoQPRE = 1e5 * ((self.X==self.Q_E) & numpy.greater(self.timer_state, 1/self.sigma_Q)) propensities_QPREtoQSYM = 1e5 * ((self.X==self.Q_pre) & numpy.greater(self.timer_state, 1/self.lamda_Q) & numpy.greater_equal(self.rand_a, self.a)) propensities_QPREtoQASYM = 1e5 * ((self.X==self.Q_pre) & numpy.greater(self.timer_state, 1/self.lamda_Q) & numpy.less(self.rand_a, self.a)) propensities_QSYMtoQR = 1e5 * ((self.X==self.Q_sym) & numpy.greater(self.timer_state, 1/self.gamma_Q_sym) & numpy.greater_equal(self.rand_h, self.h)) propensities_QSYMtoH = 1e5 * ((self.X==self.Q_sym) & numpy.greater(self.timer_state, 1/self.eta_Q) & numpy.less(self.rand_h, self.h)) propensities_QASYMtoQR = 1e5 * ((self.X==self.Q_asym) & numpy.greater(self.timer_state, 1/self.gamma_Q_asym)) propensities_RtoS = 1e5 * ((self.X==self.R) & numpy.greater(self.timer_state, 1/self.xi)) propensities__toS = 1e5 * ((self.X!=self.F) & numpy.greater(self.timer_state, 1/self.nu)) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ else: # exponential_rates propensities_EtoIPRE = self.sigma * (self.X==self.E) propensities_IPREtoISYM = self.lamda * ((self.X==self.I_pre) & (numpy.greater_equal(self.rand_a, self.a))) propensities_IPREtoIASYM = self.lamda * ((self.X==self.I_pre) & (numpy.less(self.rand_a, self.a))) propensities_ISYMtoR = self.gamma * ((self.X==self.I_sym) & (numpy.greater_equal(self.rand_h, self.h))) propensities_ISYMtoH = self.eta * ((self.X==self.I_sym) & (numpy.less(self.rand_h, self.h))) propensities_IASYMtoR = self.gamma_asym * (self.X==self.I_asym) propensities_HtoR = self.gamma_H * ((self.X==self.H) & (numpy.greater_equal(self.rand_f, self.f))) propensities_HtoF = self.mu_H * ((self.X==self.H) & (numpy.less(self.rand_f, self.f))) propensities_StoQS = (self.theta_S + self.phi_S*numContacts_Q)*self.psi_S * (self.X==self.S) propensities_EtoQE = (self.theta_E + self.phi_E*numContacts_Q)*self.psi_E * (self.X==self.E) propensities_IPREtoQPRE = (self.theta_pre + self.phi_pre*numContacts_Q)*self.psi_pre * (self.X==self.I_pre) propensities_ISYMtoQSYM = (self.theta_sym + self.phi_sym*numContacts_Q)*self.psi_sym * (self.X==self.I_sym) propensities_IASYMtoQASYM = (self.theta_asym + self.phi_asym*numContacts_Q)*self.psi_asym * (self.X==self.I_asym) propensities_QEtoQPRE = self.sigma_Q * (self.X==self.Q_E) propensities_QPREtoQSYM = self.lamda_Q * ((self.X==self.Q_pre) & (numpy.greater_equal(self.rand_a, self.a))) propensities_QPREtoQASYM = self.lamda_Q * ((self.X==self.Q_pre) & (numpy.less(self.rand_a, self.a))) propensities_QSYMtoQR = self.gamma_Q_sym * ((self.X==self.Q_sym) & (numpy.greater_equal(self.rand_h, self.h))) propensities_QSYMtoH = self.eta_Q * ((self.X==self.Q_sym) & (numpy.less(self.rand_h, self.h))) propensities_QASYMtoQR = self.gamma_Q_asym * (self.X==self.Q_asym) propensities_RtoS = self.xi * (self.X==self.R) propensities__toS = self.nu * (self.X!=self.F) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ propensities = numpy.hstack([propensities_StoE, propensities_EtoIPRE, propensities_IPREtoISYM, propensities_IPREtoIASYM, propensities_ISYMtoR, propensities_ISYMtoH, propensities_IASYMtoR, propensities_HtoR, propensities_HtoF, propensities_StoQS, propensities_EtoQE, propensities_IPREtoQPRE, propensities_ISYMtoQSYM, propensities_IASYMtoQASYM, propensities_QStoQE, propensities_QEtoQPRE, propensities_QPREtoQSYM, propensities_QPREtoQASYM, propensities_QSYMtoQR, propensities_QSYMtoH, propensities_QASYMtoQR, propensities_RtoS, propensities__toS]) columns = [ 'StoE', 'EtoIPRE', 'IPREtoISYM', 'IPREtoIASYM', 'ISYMtoR', 'ISYMtoH', 'IASYMtoR', 'HtoR', 'HtoF', 'StoQS', 'EtoQE', 'IPREtoQPRE', 'ISYMtoQSYM', 'IASYMtoQASYM', 'QStoQE', 'QEtoQPRE', 'QPREtoQSYM', 'QPREtoQASYM', 'QSYMtoQR', 'QSYMtoH', 'QASYMtoQR', 'RtoS', '_toS' ] return propensities, columns #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def set_isolation(self, node, isolate): # Move this node in/out of the appropriate isolation state: if(isolate == True): if(self.X[node] == self.S): self.X[node] = self.Q_S elif(self.X[node] == self.E): self.X[node] = self.Q_E elif(self.X[node] == self.I_pre): self.X[node] = self.Q_pre elif(self.X[node] == self.I_sym): self.X[node] = self.Q_sym elif(self.X[node] == self.I_asym): self.X[node] = self.Q_asym elif(self.X[node] == self.R): self.X[node] = self.Q_R elif(isolate == False): if(self.X[node] == self.Q_S): self.X[node] = self.S elif(self.X[node] == self.Q_E): self.X[node] = self.E elif(self.X[node] == self.Q_pre): self.X[node] = self.I_pre elif(self.X[node] == self.Q_sym): self.X[node] = self.I_sym elif(self.X[node] == self.Q_asym): self.X[node] = self.I_asym elif(self.X[node] == self.Q_R): self.X[node] = self.R # Reset the isolation timer: self.timer_isolation[node] = 0 #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def set_tested(self, node, tested): self.tested[node] = tested self.testedInCurrentState[node] = tested #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def set_positive(self, node, positive): self.positive[node] = positive #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def introduce_exposures(self, num_new_exposures): exposedNodes = numpy.random.choice(range(self.numNodes), size=num_new_exposures, replace=False) for exposedNode in exposedNodes: if(self.X[exposedNode]==self.S): self.X[exposedNode] = self.E elif(self.X[exposedNode]==self.Q_S): self.X[exposedNode] = self.Q_E #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def increase_data_series_length(self): self.tseries = numpy.pad(self.tseries, [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.numS = numpy.pad(self.numS, [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.numE = numpy.pad(self.numE, [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.numI_pre = numpy.pad(self.numI_pre, [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.numI_sym = numpy.pad(self.numI_sym, [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.numI_asym = numpy.pad(self.numI_asym, [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.numH = numpy.pad(self.numH, [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.numR = numpy.pad(self.numR, [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.numF = numpy.pad(self.numF, [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.numQ_S = numpy.pad(self.numQ_S, [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.numQ_E = numpy.pad(self.numQ_E, [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.numQ_pre = numpy.pad(self.numQ_pre, [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.numQ_sym = numpy.pad(self.numQ_sym, [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.numQ_asym = numpy.pad(self.numQ_asym, [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.numQ_R = numpy.pad(self.numQ_R, [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.N = numpy.pad(self.N, [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.numTested = numpy.pad(self.numTested, [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.numPositive = numpy.pad(self.numPositive, [(0, 6*self.numNodes)], mode='constant', constant_values=0) if(self.store_Xseries): self.Xseries = numpy.pad(self.Xseries, [(0, 6*self.numNodes), (0,0)], mode='constant', constant_values=0) if(self.nodeGroupData): for groupName in self.nodeGroupData: self.nodeGroupData[groupName]['numS'] = numpy.pad(self.nodeGroupData[groupName]['numS'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.nodeGroupData[groupName]['numE'] = numpy.pad(self.nodeGroupData[groupName]['numE'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.nodeGroupData[groupName]['numI_pre'] = numpy.pad(self.nodeGroupData[groupName]['numI_pre'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.nodeGroupData[groupName]['numI_sym'] = numpy.pad(self.nodeGroupData[groupName]['numI_sym'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.nodeGroupData[groupName]['numI_asym'] = numpy.pad(self.nodeGroupData[groupName]['numI_asym'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.nodeGroupData[groupName]['numH'] = numpy.pad(self.nodeGroupData[groupName]['numH'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.nodeGroupData[groupName]['numR'] = numpy.pad(self.nodeGroupData[groupName]['numR'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.nodeGroupData[groupName]['numF'] = numpy.pad(self.nodeGroupData[groupName]['numF'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.nodeGroupData[groupName]['numQ_S'] = numpy.pad(self.nodeGroupData[groupName]['numQ_S'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.nodeGroupData[groupName]['numQ_E'] = numpy.pad(self.nodeGroupData[groupName]['numQ_E'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.nodeGroupData[groupName]['numQ_pre'] = numpy.pad(self.nodeGroupData[groupName]['numQ_pre'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.nodeGroupData[groupName]['numQ_sym'] = numpy.pad(self.nodeGroupData[groupName]['numQ_sym'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.nodeGroupData[groupName]['numQ_asym'] = numpy.pad(self.nodeGroupData[groupName]['numQ_asym'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.nodeGroupData[groupName]['numQ_R'] = numpy.pad(self.nodeGroupData[groupName]['numQ_R'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.nodeGroupData[groupName]['N'] = numpy.pad(self.nodeGroupData[groupName]['N'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.nodeGroupData[groupName]['numTested'] = numpy.pad(self.nodeGroupData[groupName]['numTested'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) self.nodeGroupData[groupName]['numPositive'] = numpy.pad(self.nodeGroupData[groupName]['numPositive'], [(0, 6*self.numNodes)], mode='constant', constant_values=0) return None #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def finalize_data_series(self): self.tseries = numpy.array(self.tseries, dtype=float)[:self.tidx+1] self.numS = numpy.array(self.numS, dtype=float)[:self.tidx+1] self.numE = numpy.array(self.numE, dtype=float)[:self.tidx+1] self.numI_pre = numpy.array(self.numI_pre, dtype=float)[:self.tidx+1] self.numI_sym = numpy.array(self.numI_sym, dtype=float)[:self.tidx+1] self.numI_asym = numpy.array(self.numI_asym, dtype=float)[:self.tidx+1] self.numH = numpy.array(self.numH, dtype=float)[:self.tidx+1] self.numR = numpy.array(self.numR, dtype=float)[:self.tidx+1] self.numF = numpy.array(self.numF, dtype=float)[:self.tidx+1] self.numQ_S = numpy.array(self.numQ_S, dtype=float)[:self.tidx+1] self.numQ_E = numpy.array(self.numQ_E, dtype=float)[:self.tidx+1] self.numQ_pre = numpy.array(self.numQ_pre, dtype=float)[:self.tidx+1] self.numQ_sym = numpy.array(self.numQ_sym, dtype=float)[:self.tidx+1] self.numQ_asym = numpy.array(self.numQ_asym, dtype=float)[:self.tidx+1] self.numQ_R = numpy.array(self.numQ_R, dtype=float)[:self.tidx+1] self.N = numpy.array(self.N, dtype=float)[:self.tidx+1] self.numTested = numpy.array(self.numTested, dtype=float)[:self.tidx+1] self.numPositive = numpy.array(self.numPositive, dtype=float)[:self.tidx+1] if(self.store_Xseries): self.Xseries = self.Xseries[:self.tidx+1, :] if(self.nodeGroupData): for groupName in self.nodeGroupData: self.nodeGroupData[groupName]['numS'] = numpy.array(self.nodeGroupData[groupName]['numS'], dtype=float)[:self.tidx+1] self.nodeGroupData[groupName]['numE'] = numpy.array(self.nodeGroupData[groupName]['numE'], dtype=float)[:self.tidx+1] self.nodeGroupData[groupName]['numI_pre'] = numpy.array(self.nodeGroupData[groupName]['numI_pre'], dtype=float)[:self.tidx+1] self.nodeGroupData[groupName]['numI_sym'] = numpy.array(self.nodeGroupData[groupName]['numI_sym'], dtype=float)[:self.tidx+1] self.nodeGroupData[groupName]['numI_asym'] = numpy.array(self.nodeGroupData[groupName]['numI_asym'], dtype=float)[:self.tidx+1] self.nodeGroupData[groupName]['numR'] = numpy.array(self.nodeGroupData[groupName]['numR'], dtype=float)[:self.tidx+1] self.nodeGroupData[groupName]['numF'] = numpy.array(self.nodeGroupData[groupName]['numF'], dtype=float)[:self.tidx+1] self.nodeGroupData[groupName]['numQ_S'] = numpy.array(self.nodeGroupData[groupName]['numQ_S'], dtype=float)[:self.tidx+1] self.nodeGroupData[groupName]['numQ_E'] = numpy.array(self.nodeGroupData[groupName]['numQ_E'], dtype=float)[:self.tidx+1] self.nodeGroupData[groupName]['numQ_pre'] = numpy.array(self.nodeGroupData[groupName]['numQ_pre'], dtype=float)[:self.tidx+1] self.nodeGroupData[groupName]['numQ_sym'] = numpy.array(self.nodeGroupData[groupName]['numQ_sym'], dtype=float)[:self.tidx+1] self.nodeGroupData[groupName]['numQ_asym'] = numpy.array(self.nodeGroupData[groupName]['numQ_asym'], dtype=float)[:self.tidx+1] self.nodeGroupData[groupName]['numQ_R'] = numpy.array(self.nodeGroupData[groupName]['numQ_R'], dtype=float)[:self.tidx+1] self.nodeGroupData[groupName]['N'] = numpy.array(self.nodeGroupData[groupName]['N'], dtype=float)[:self.tidx+1] self.nodeGroupData[groupName]['numTested'] = numpy.array(self.nodeGroupData[groupName]['numTested'], dtype=float)[:self.tidx+1] self.nodeGroupData[groupName]['numPositive'] = numpy.array(self.nodeGroupData[groupName]['numPositive'], dtype=float)[:self.tidx+1] return None #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def run_iteration(self): 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' or transitionType == 'QStoQE'): transitionNode_GNbrs = list(self.G[transitionNode].keys()) transitionNode_GQNbrs = list(self.G_Q[transitionNode].keys()) self.infectionsLog.append({ 't': self.t, 'infected_node': transitionNode, 'infection_type': transitionType, 'infected_node_degree': self.degree[transitionNode], 'local_contact_nodes': transitionNode_GNbrs, 'local_contact_node_states': self.X[transitionNode_GNbrs].flatten(), 'isolation_contact_nodes': transitionNode_GQNbrs, 'isolation_contact_node_states':self.X[transitionNode_GQNbrs].flatten() }) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(transitionType in ['EtoQE', 'IPREtoQPRE', 'ISYMtoQSYM', 'IASYMtoQASYM', 'ISYMtoH']): self.set_positive(node=transitionNode, positive=True) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ else: tau = 0.01 self.t += tau self.timer_state += tau #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ self.tidx += 1 self.tseries[self.tidx] = self.t self.numS[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.S), a_min=0, a_max=self.numNodes) self.numE[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.E), a_min=0, a_max=self.numNodes) self.numI_pre[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.I_pre), a_min=0, a_max=self.numNodes) self.numI_sym[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.I_sym), a_min=0, a_max=self.numNodes) self.numI_asym[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.I_asym), a_min=0, a_max=self.numNodes) self.numH[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.H), a_min=0, a_max=self.numNodes) self.numR[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.R), a_min=0, a_max=self.numNodes) self.numF[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.F), a_min=0, a_max=self.numNodes) self.numQ_S[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.Q_S), a_min=0, a_max=self.numNodes) self.numQ_E[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.Q_E), a_min=0, a_max=self.numNodes) self.numQ_pre[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.Q_pre), a_min=0, a_max=self.numNodes) self.numQ_sym[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.Q_sym), a_min=0, a_max=self.numNodes) self.numQ_asym[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.Q_asym), a_min=0, a_max=self.numNodes) self.numQ_R[self.tidx] = numpy.clip(numpy.count_nonzero(self.X==self.Q_R), a_min=0, a_max=self.numNodes) self.numTested[self.tidx] = numpy.clip(numpy.count_nonzero(self.tested), a_min=0, a_max=self.numNodes) self.numPositive[self.tidx] = numpy.clip(numpy.count_nonzero(self.positive), a_min=0, a_max=self.numNodes) self.N[self.tidx] = numpy.clip((self.numNodes - self.numF[self.tidx]), a_min=0, a_max=self.numNodes) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Update testing and isolation statuses #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ isolatedNodes = numpy.argwhere((self.X==self.Q_S)|(self.X==self.Q_E)|(self.X==self.Q_pre)|(self.X==self.Q_sym)|(self.X==self.Q_asym)|(self.X==self.Q_R))[:,0].flatten() self.timer_isolation[isolatedNodes] = self.timer_isolation[isolatedNodes] + tau nodesExitingIsolation = numpy.argwhere(self.timer_isolation >= self.isolationTime) for isoNode in nodesExitingIsolation: self.set_isolation(node=isoNode, isolate=False) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Store system states #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(self.store_Xseries): self.Xseries[self.tidx,:] = self.X.T if(self.nodeGroupData): for groupName in self.nodeGroupData: self.nodeGroupData[groupName]['numS'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.S) self.nodeGroupData[groupName]['numE'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.E) self.nodeGroupData[groupName]['numI_pre'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.I_pre) self.nodeGroupData[groupName]['numI_sym'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.I_sym) self.nodeGroupData[groupName]['numI_asym'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.I_asym) self.nodeGroupData[groupName]['numH'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.H) self.nodeGroupData[groupName]['numR'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.R) self.nodeGroupData[groupName]['numF'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.F) self.nodeGroupData[groupName]['numQ_S'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_S) self.nodeGroupData[groupName]['numQ_E'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_E) self.nodeGroupData[groupName]['numQ_pre'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_pre) self.nodeGroupData[groupName]['numQ_sym'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_sym) self.nodeGroupData[groupName]['numQ_asym'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_asym) self.nodeGroupData[groupName]['numQ_R'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.X==self.Q_R) self.nodeGroupData[groupName]['N'][self.tidx] = numpy.clip((self.nodeGroupData[groupName]['numS'][0] + self.nodeGroupData[groupName]['numE'][0] + self.nodeGroupData[groupName]['numI'][0] + self.nodeGroupData[groupName]['numQ_E'][0] + self.nodeGroupData[groupName]['numQ_I'][0] + self.nodeGroupData[groupName]['numR'][0]), a_min=0, a_max=self.numNodes) self.nodeGroupData[groupName]['numTested'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.tested) self.nodeGroupData[groupName]['numPositive'][self.tidx] = numpy.count_nonzero(self.nodeGroupData[groupName]['mask']*self.positive) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Terminate if tmax reached or num infections is 0: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(self.t >= self.tmax or (self.total_num_infected(self.tidx) < 1 and self.total_num_isolated(self.tidx) < 1)): self.finalize_data_series() return False #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ return True #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def run(self, T, checkpoints=None, print_interval=10, verbose='t'): if(T>0): self.tmax += T else: return False #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Pre-process checkpoint values: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(checkpoints): numCheckpoints = len(checkpoints['t']) for chkpt_param, chkpt_values in checkpoints.items(): assert(isinstance(chkpt_values, (list, numpy.ndarray)) and len(chkpt_values)==numCheckpoints), "Expecting a list of values with length equal to number of checkpoint times ("+str(numCheckpoints)+") for each checkpoint parameter." checkpointIdx = numpy.searchsorted(checkpoints['t'], self.t) # Finds 1st index in list greater than given val if(checkpointIdx >= numCheckpoints): # We are out of checkpoints, stop checking them: checkpoints = None else: checkpointTime = checkpoints['t'][checkpointIdx] #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% # Run the simulation loop: #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% print_reset = True running = True while running: running = self.run_iteration() #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Handle checkpoints if applicable: if(checkpoints): if(self.t >= checkpointTime): if(verbose is not False): print("[Checkpoint: Updating parameters]") # A checkpoint has been reached, update param values: for param in list(self.parameters.keys()): if(param in list(checkpoints.keys())): self.parameters.update({param: checkpoints[param][checkpointIdx]}) # Update parameter data structures and scenario flags: self.update_parameters() # Update the next checkpoint time: checkpointIdx = numpy.searchsorted(checkpoints['t'], self.t) # Finds 1st index in list greater than given val if(checkpointIdx >= numCheckpoints): # We are out of checkpoints, stop checking them: checkpoints = None else: checkpointTime = checkpoints['t'][checkpointIdx] #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(print_interval): if(print_reset and (int(self.t) % print_interval == 0)): if(verbose=="t"): print("t = %.2f" % self.t) if(verbose==True): print("t = %.2f" % self.t) print("\t S = " + str(self.numS[self.tidx])) print("\t E = " + str(self.numE[self.tidx])) print("\t I_pre = " + str(self.numI_pre[self.tidx])) print("\t I_sym = " + str(self.numI_sym[self.tidx])) print("\t I_asym = " + str(self.numI_asym[self.tidx])) print("\t H = " + str(self.numH[self.tidx])) print("\t R = " + str(self.numR[self.tidx])) print("\t F = " + str(self.numF[self.tidx])) print("\t Q_S = " + str(self.numQ_S[self.tidx])) print("\t Q_E = " + str(self.numQ_E[self.tidx])) print("\t Q_pre = " + str(self.numQ_pre[self.tidx])) print("\t Q_sym = " + str(self.numQ_sym[self.tidx])) print("\t Q_asym = " + str(self.numQ_asym[self.tidx])) print("\t Q_R = " + str(self.numQ_R[self.tidx])) print_reset = False elif(not print_reset and (int(self.t) % 10 != 0)): print_reset = True return True #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def plot(self, ax=None, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_sym='line', plot_I_asym='line', plot_H='line', plot_R='line', plot_F='line', plot_Q_E='line', plot_Q_pre='line', plot_Q_sym='line', plot_Q_asym='line', plot_Q_S='line', plot_Q_R='line', combine_Q_infected=True, color_S='tab:green', color_E='orange', color_I_pre='tomato', color_I_sym='crimson', color_I_asym='#F0909B', color_H='violet', color_R='tab:blue', color_F='black', color_Q_E='orange', color_Q_pre='tomato', color_Q_sym='crimson', color_Q_asym='#F0909B', color_Q_S='tab:green', color_Q_R='tab:blue', color_Q_infected='tab:purple', color_reference='#E0E0E0', dashed_reference_results=None, dashed_reference_label='reference', shaded_reference_results=None, shaded_reference_label='reference', vlines=[], vline_colors=[], vline_styles=[], vline_labels=[], ylim=None, xlim=None, legend=True, title=None, side_title=None, plot_percentages=True): import matplotlib.pyplot as pyplot #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Create an Axes object if None provided: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(not ax): fig, ax = pyplot.subplots() #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Prepare data series to be plotted: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Sseries = self.numS/self.numNodes if plot_percentages else self.numS Eseries = self.numE/self.numNodes if plot_percentages else self.numE I_preseries = self.numI_pre/self.numNodes if plot_percentages else self.numI_pre I_symseries = self.numI_sym/self.numNodes if plot_percentages else self.numI_sym I_asymseries = self.numI_asym/self.numNodes if plot_percentages else self.numI_asym Rseries = self.numR/self.numNodes if plot_percentages else self.numR Hseries = self.numH/self.numNodes if plot_percentages else self.numH Fseries = self.numF/self.numNodes if plot_percentages else self.numF Q_Sseries = self.numQ_S/self.numNodes if plot_percentages else self.numQ_S Q_Eseries = self.numQ_E/self.numNodes if plot_percentages else self.numQ_E Q_preseries = self.numQ_pre/self.numNodes if plot_percentages else self.numQ_pre Q_asymseries = self.numQ_asym/self.numNodes if plot_percentages else self.numQ_asym Q_symseries = self.numQ_sym/self.numNodes if plot_percentages else self.numQ_sym Q_Rseries = self.numQ_R/self.numNodes if plot_percentages else self.numQ_R Q_infectedseries = (Q_Eseries + Q_preseries + Q_asymseries + Q_symseries) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Draw the reference data: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(dashed_reference_results): dashedReference_tseries = dashed_reference_results.tseries[::int(self.numNodes/100)] dashedReference_infectedStack = dashed_reference_results.total_num_infected()[::int(self.numNodes/100)] / (self.numNodes if plot_percentages else 1) ax.plot(dashedReference_tseries, dashedReference_infectedStack, color='#E0E0E0', linestyle='--', label='Total infections ('+dashed_reference_label+')', zorder=0) if(shaded_reference_results): shadedReference_tseries = shaded_reference_results.tseries shadedReference_infectedStack = shaded_reference_results.total_num_infected() / (self.numNodes if plot_percentages else 1) ax.fill_between(shaded_reference_results.tseries, shadedReference_infectedStack, 0, color='#EFEFEF', label='Total infections ('+shaded_reference_label+')', zorder=0) ax.plot(shaded_reference_results.tseries, shadedReference_infectedStack, color='#E0E0E0', zorder=1) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Draw the stacked variables: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ topstack = numpy.zeros_like(self.tseries) if(any(Fseries) and plot_F=='stacked'): ax.fill_between(numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, topstack+Fseries), topstack, color=color_F, alpha=0.75, label='$F$', zorder=2) ax.plot( numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, topstack+Fseries), color=color_F, zorder=3) topstack = topstack+Fseries if(any(Hseries) and plot_H=='stacked'): ax.fill_between(numpy.ma.masked_where(Hseries<=0, self.tseries), numpy.ma.masked_where(Hseries<=0, topstack+Hseries), topstack, color=color_H, alpha=0.75, label='$H$', zorder=2) ax.plot( numpy.ma.masked_where(Hseries<=0, self.tseries), numpy.ma.masked_where(Hseries<=0, topstack+Hseries), color=color_H, zorder=3) topstack = topstack+Hseries if(combine_Q_infected and any(Q_infectedseries) and plot_Q_E=='stacked' and plot_Q_pre=='stacked' and plot_Q_sym=='stacked' and plot_Q_asym=='stacked'): ax.fill_between(numpy.ma.masked_where(Q_infectedseries<=0, self.tseries), numpy.ma.masked_where(Q_infectedseries<=0, topstack+Q_infectedseries), topstack, facecolor=color_Q_infected, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_{infected}$', zorder=2) ax.plot( numpy.ma.masked_where(Q_infectedseries<=0, self.tseries), numpy.ma.masked_where(Q_infectedseries<=0, topstack+Q_infectedseries), color=color_Q_infected, zorder=3) topstack = topstack+Q_infectedseries if(not combine_Q_infected and any(Q_Eseries) and plot_Q_E=='stacked'): ax.fill_between(numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, topstack+Q_Eseries), topstack, facecolor=color_Q_E, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_E$', zorder=2) ax.plot( numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, topstack+Q_Eseries), color=color_Q_E, zorder=3) topstack = topstack+Q_Eseries if(any(Eseries) and plot_E=='stacked'): ax.fill_between(numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, topstack+Eseries), topstack, color=color_E, alpha=0.75, label='$E$', zorder=2) ax.plot( numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, topstack+Eseries), color=color_E, zorder=3) topstack = topstack+Eseries if(not combine_Q_infected and any(Q_preseries) and plot_Q_pre=='stacked'): ax.fill_between(numpy.ma.masked_where(Q_preseries<=0, self.tseries), numpy.ma.masked_where(Q_preseries<=0, topstack+Q_preseries), topstack, facecolor=color_Q_pre, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_{pre}$', zorder=2) ax.plot( numpy.ma.masked_where(Q_preseries<=0, self.tseries), numpy.ma.masked_where(Q_preseries<=0, topstack+Q_preseries), color=color_Q_pre, zorder=3) topstack = topstack+Q_preseries if(any(I_preseries) and plot_I_pre=='stacked'): ax.fill_between(numpy.ma.masked_where(I_preseries<=0, self.tseries), numpy.ma.masked_where(I_preseries<=0, topstack+I_preseries), topstack, color=color_I_pre, alpha=0.75, label='$I_{pre}$', zorder=2) ax.plot( numpy.ma.masked_where(I_preseries<=0, self.tseries), numpy.ma.masked_where(I_preseries<=0, topstack+I_preseries), color=color_I_pre, zorder=3) topstack = topstack+I_preseries if(not combine_Q_infected and any(Q_symseries) and plot_Q_sym=='stacked'): ax.fill_between(numpy.ma.masked_where(Q_symseries<=0, self.tseries), numpy.ma.masked_where(Q_symseries<=0, topstack+Q_symseries), topstack, facecolor=color_Q_sym, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_{sym}$', zorder=2) ax.plot( numpy.ma.masked_where(Q_symseries<=0, self.tseries), numpy.ma.masked_where(Q_symseries<=0, topstack+Q_symseries), color=color_Q_sym, zorder=3) topstack = topstack+Q_symseries if(any(I_symseries) and plot_I_sym=='stacked'): ax.fill_between(numpy.ma.masked_where(I_symseries<=0, self.tseries), numpy.ma.masked_where(I_symseries<=0, topstack+I_symseries), topstack, color=color_I_sym, alpha=0.75, label='$I_{sym}$', zorder=2) ax.plot( numpy.ma.masked_where(I_symseries<=0, self.tseries), numpy.ma.masked_where(I_symseries<=0, topstack+I_symseries), color=color_I_sym, zorder=3) topstack = topstack+I_symseries if(not combine_Q_infected and any(Q_asymseries) and plot_Q_asym=='stacked'): ax.fill_between(numpy.ma.masked_where(Q_asymseries<=0, self.tseries), numpy.ma.masked_where(Q_asymseries<=0, topstack+Q_asymseries), topstack, facecolor=color_Q_asym, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_{asym}$', zorder=2) ax.plot( numpy.ma.masked_where(Q_asymseries<=0, self.tseries), numpy.ma.masked_where(Q_asymseries<=0, topstack+Q_asymseries), color=color_Q_asym, zorder=3) topstack = topstack+Q_asymseries if(any(I_asymseries) and plot_I_asym=='stacked'): ax.fill_between(numpy.ma.masked_where(I_asymseries<=0, self.tseries), numpy.ma.masked_where(I_asymseries<=0, topstack+I_asymseries), topstack, color=color_I_asym, alpha=0.75, label='$I_{asym}$', zorder=2) ax.plot( numpy.ma.masked_where(I_asymseries<=0, self.tseries), numpy.ma.masked_where(I_asymseries<=0, topstack+I_asymseries), color=color_I_asym, zorder=3) topstack = topstack+I_asymseries if(any(Q_Rseries) and plot_Q_R=='stacked'): ax.fill_between(numpy.ma.masked_where(Q_Rseries<=0, self.tseries), numpy.ma.masked_where(Q_Rseries<=0, topstack+Q_Rseries), topstack, facecolor=color_Q_R, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_R$', zorder=2) ax.plot( numpy.ma.masked_where(Q_Rseries<=0, self.tseries), numpy.ma.masked_where(Q_Rseries<=0, topstack+Q_Rseries), color=color_Q_R, zorder=3) topstack = topstack+Q_Rseries if(any(Rseries) and plot_R=='stacked'): ax.fill_between(numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, topstack+Rseries), topstack, color=color_R, alpha=0.75, label='$R$', zorder=2) ax.plot( numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, topstack+Rseries), color=color_R, zorder=3) topstack = topstack+Rseries if(any(Q_Sseries) and plot_Q_S=='stacked'): ax.fill_between(numpy.ma.masked_where(Q_Sseries<=0, self.tseries), numpy.ma.masked_where(Q_Sseries<=0, topstack+Q_Sseries), topstack, facecolor=color_Q_S, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_S$', zorder=2) ax.plot( numpy.ma.masked_where(Q_Sseries<=0, self.tseries), numpy.ma.masked_where(Q_Sseries<=0, topstack+Q_Sseries), color=color_Q_S, zorder=3) topstack = topstack+Q_Sseries if(any(Sseries) and plot_S=='stacked'): ax.fill_between(numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, topstack+Sseries), topstack, color=color_S, alpha=0.75, label='$S$', zorder=2) ax.plot( numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, topstack+Sseries), color=color_S, zorder=3) topstack = topstack+Sseries #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Draw the shaded variables: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(any(Fseries) and plot_F=='shaded'): ax.fill_between(numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, Fseries), 0, color=color_F, alpha=0.75, label='$F$', zorder=4) ax.plot( numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, Fseries), color=color_F, zorder=5) if(any(Hseries) and plot_H=='shaded'): ax.fill_between(numpy.ma.masked_where(Hseries<=0, self.tseries), numpy.ma.masked_where(Hseries<=0, Hseries), 0, color=color_H, alpha=0.75, label='$H$', zorder=4) ax.plot( numpy.ma.masked_where(Hseries<=0, self.tseries), numpy.ma.masked_where(Hseries<=0, Hseries), color=color_H, zorder=5) if(combine_Q_infected and any(Q_infectedseries) and plot_Q_E=='shaded' and plot_Q_pre=='shaded' and plot_Q_sym=='shaded' and plot_Q_asym=='shaded'): ax.fill_between(numpy.ma.masked_where(Q_infectedseries<=0, self.tseries), numpy.ma.masked_where(Q_infectedseries<=0, Q_infectedseries), 0, color=color_Q_infected, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_{infected}$', zorder=4) ax.plot( numpy.ma.masked_where(Q_infectedseries<=0, self.tseries), numpy.ma.masked_where(Q_infectedseries<=0, Q_infectedseries), color=color_Q_infected, zorder=5) if(not combine_Q_infected and any(Q_Eseries) and plot_Q_E=='shaded'): ax.fill_between(numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, Q_Eseries), 0, facecolor=color_Q_E, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_E$', zorder=4) ax.plot( numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, Q_Eseries), color=color_Q_E, zorder=5) if(any(Eseries) and plot_E=='shaded'): ax.fill_between(numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, Eseries), 0, color=color_E, alpha=0.75, label='$E$', zorder=4) ax.plot( numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, Eseries), color=color_E, zorder=5) if(not combine_Q_infected and any(Q_preseries) and plot_Q_pre=='shaded'): ax.fill_between(numpy.ma.masked_where(Q_preseries<=0, self.tseries), numpy.ma.masked_where(Q_preseries<=0, Q_preseries), 0, facecolor=color_Q_pre, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_{pre}$', zorder=4) ax.plot( numpy.ma.masked_where(Q_preseries<=0, self.tseries), numpy.ma.masked_where(Q_preseries<=0, Q_preseries), color=color_Q_pre, zorder=5) if(any(I_preseries) and plot_I_pre=='shaded'): ax.fill_between(numpy.ma.masked_where(I_preseries<=0, self.tseries), numpy.ma.masked_where(I_preseries<=0, I_preseries), 0, color=color_I_pre, alpha=0.75, label='$I_{pre}$', zorder=4) ax.plot( numpy.ma.masked_where(I_preseries<=0, self.tseries), numpy.ma.masked_where(I_preseries<=0, I_preseries), color=color_I_pre, zorder=5) if(not combine_Q_infected and any(Q_symseries) and plot_Q_sym=='shaded'): ax.fill_between(numpy.ma.masked_where(Q_symseries<=0, self.tseries), numpy.ma.masked_where(Q_symseries<=0, Q_symseries), 0, facecolor=color_Q_sym, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_{sym}$', zorder=4) ax.plot( numpy.ma.masked_where(Q_symseries<=0, self.tseries), numpy.ma.masked_where(Q_symseries<=0, Q_symseries), color=color_Q_sym, zorder=5) if(any(I_symseries) and plot_I_sym=='shaded'): ax.fill_between(numpy.ma.masked_where(I_symseries<=0, self.tseries), numpy.ma.masked_where(I_symseries<=0, I_symseries), 0, color=color_I_sym, alpha=0.75, label='$I_{sym}$', zorder=4) ax.plot( numpy.ma.masked_where(I_symseries<=0, self.tseries), numpy.ma.masked_where(I_symseries<=0, I_symseries), color=color_I_sym, zorder=5) if(not combine_Q_infected and any(Q_asymseries) and plot_Q_asym=='shaded'): ax.fill_between(numpy.ma.masked_where(Q_asymseries<=0, self.tseries), numpy.ma.masked_where(Q_asymseries<=0, Q_asymseries), 0, facecolor=color_Q_asym, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_{asym}$', zorder=4) ax.plot( numpy.ma.masked_where(Q_asymseries<=0, self.tseries), numpy.ma.masked_where(Q_asymseries<=0, Q_asymseries), color=color_Q_asym, zorder=5) if(any(I_asymseries) and plot_I_asym=='shaded'): ax.fill_between(numpy.ma.masked_where(I_asymseries<=0, self.tseries), numpy.ma.masked_where(I_asymseries<=0, I_asymseries), 0, color=color_I_asym, alpha=0.75, label='$I_{asym}$', zorder=4) ax.plot( numpy.ma.masked_where(I_asymseries<=0, self.tseries), numpy.ma.masked_where(I_asymseries<=0, I_asymseries), color=color_I_asym, zorder=5) if(any(Q_Rseries) and plot_Q_R=='shaded'): ax.fill_between(numpy.ma.masked_where(Q_Rseries<=0, self.tseries), numpy.ma.masked_where(Q_Rseries<=0, Q_Rseries), 0, facecolor=color_Q_R, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_R$', zorder=4) ax.plot( numpy.ma.masked_where(Q_Rseries<=0, self.tseries), numpy.ma.masked_where(Q_Rseries<=0, Q_Rseries), color=color_Q_R, zorder=5) if(any(Rseries) and plot_R=='shaded'): ax.fill_between(numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, Rseries), 0, color=color_R, alpha=0.75, label='$R$', zorder=4) ax.plot( numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, Rseries), color=color_R, zorder=5) if(any(Q_Sseries) and plot_Q_S=='shaded'): ax.fill_between(numpy.ma.masked_where(Q_Sseries<=0, self.tseries), numpy.ma.masked_where(Q_Sseries<=0, Q_Sseries), 0, facecolor=color_Q_S, alpha=0.75, hatch='//////', edgecolor='white', linewidth=0.0, label='$Q_S$', zorder=4) ax.plot( numpy.ma.masked_where(Q_Sseries<=0, self.tseries), numpy.ma.masked_where(Q_Sseries<=0, Q_Sseries), color=color_Q_S, zorder=5) if(any(Sseries) and plot_S=='shaded'): ax.fill_between(numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, Sseries), 0, color=color_S, alpha=0.75, label='$S$', zorder=4) ax.plot( numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, Sseries), color=color_S, zorder=5) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Draw the line variables: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(any(Fseries) and plot_F=='line'): ax.plot(numpy.ma.masked_where(Fseries<=0, self.tseries), numpy.ma.masked_where(Fseries<=0, Fseries), color=color_F, label='$F$', zorder=6) if(any(Hseries) and plot_H=='line'): ax.plot(numpy.ma.masked_where(Hseries<=0, self.tseries), numpy.ma.masked_where(Hseries<=0, Hseries), color=color_H, label='$H$', zorder=6) if(combine_Q_infected and any(Q_infectedseries) and plot_Q_E=='line' and plot_Q_pre=='line' and plot_Q_sym=='line' and plot_Q_asym=='line'): ax.plot(numpy.ma.masked_where(Q_infectedseries<=0, self.tseries), numpy.ma.masked_where(Q_infectedseries<=0, Q_infectedseries), color=color_Q_infected, label='$Q_{infected}$', zorder=6) if(not combine_Q_infected and any(Q_Eseries) and plot_Q_E=='line'): ax.plot(numpy.ma.masked_where(Q_Eseries<=0, self.tseries), numpy.ma.masked_where(Q_Eseries<=0, Q_Eseries), color=color_Q_E, label='$Q_E$', zorder=6) if(any(Eseries) and plot_E=='line'): ax.plot(numpy.ma.masked_where(Eseries<=0, self.tseries), numpy.ma.masked_where(Eseries<=0, Eseries), color=color_E, label='$E$', zorder=6) if(not combine_Q_infected and any(Q_preseries) and plot_Q_pre=='line'): ax.plot(numpy.ma.masked_where(Q_preseries<=0, self.tseries), numpy.ma.masked_where(Q_preseries<=0, Q_preseries), color=color_Q_pre, label='$Q_{pre}$', zorder=6) if(any(I_preseries) and plot_I_pre=='line'): ax.plot(numpy.ma.masked_where(I_preseries<=0, self.tseries), numpy.ma.masked_where(I_preseries<=0, I_preseries), color=color_I_pre, label='$I_{pre}$', zorder=6) if(not combine_Q_infected and any(Q_symseries) and plot_Q_sym=='line'): ax.plot(numpy.ma.masked_where(Q_symseries<=0, self.tseries), numpy.ma.masked_where(Q_symseries<=0, Q_symseries), color=color_Q_sym, label='$Q_{sym}$', zorder=6) if(any(I_symseries) and plot_I_sym=='line'): ax.plot(numpy.ma.masked_where(I_symseries<=0, self.tseries), numpy.ma.masked_where(I_symseries<=0, I_symseries), color=color_I_sym, label='$I_{sym}$', zorder=6) if(not combine_Q_infected and any(Q_asymseries) and plot_Q_asym=='line'): ax.plot(numpy.ma.masked_where(Q_asymseries<=0, self.tseries), numpy.ma.masked_where(Q_asymseries<=0, Q_asymseries), color=color_Q_asym, label='$Q_{asym}$', zorder=6) if(any(I_asymseries) and plot_I_asym=='line'): ax.plot(numpy.ma.masked_where(I_asymseries<=0, self.tseries), numpy.ma.masked_where(I_asymseries<=0, I_asymseries), color=color_I_asym, label='$I_{asym}$', zorder=6) if(any(Q_Rseries) and plot_Q_R=='line'): ax.plot(numpy.ma.masked_where(Q_Rseries<=0, self.tseries), numpy.ma.masked_where(Q_Rseries<=0, Q_Rseries), color=color_Q_R, linestyle='--', label='$Q_R$', zorder=6) if(any(Rseries) and plot_R=='line'): ax.plot(numpy.ma.masked_where(Rseries<=0, self.tseries), numpy.ma.masked_where(Rseries<=0, Rseries), color=color_R, label='$R$', zorder=6) if(any(Q_Sseries) and plot_Q_S=='line'): ax.plot(numpy.ma.masked_where(Q_Sseries<=0, self.tseries), numpy.ma.masked_where(Q_Sseries<=0, Q_Sseries), color=color_Q_S, linestyle='--', label='$Q_S$', zorder=6) if(any(Sseries) and plot_S=='line'): ax.plot(numpy.ma.masked_where(Sseries<=0, self.tseries), numpy.ma.masked_where(Sseries<=0, Sseries), color=color_S, label='$S$', zorder=6) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Draw the vertical line annotations: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(len(vlines)>0 and len(vline_colors)==0): vline_colors = ['gray']*len(vlines) if(len(vlines)>0 and len(vline_labels)==0): vline_labels = [None]*len(vlines) if(len(vlines)>0 and len(vline_styles)==0): vline_styles = [':']*len(vlines) for vline_x, vline_color, vline_style, vline_label in zip(vlines, vline_colors, vline_styles, vline_labels): if(vline_x is not None): ax.axvline(x=vline_x, color=vline_color, linestyle=vline_style, alpha=1, label=vline_label) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Draw the plot labels: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ ax.set_xlabel('days') ax.set_ylabel('percent of population' if plot_percentages else 'number of individuals') ax.set_xlim(0, (max(self.tseries) if not xlim else xlim)) ax.set_ylim(0, ylim) if(plot_percentages): ax.set_yticklabels(['{:,.0%}'.format(y) for y in ax.get_yticks()]) if(legend): legend_handles, legend_labels = ax.get_legend_handles_labels() ax.legend(legend_handles[::-1], legend_labels[::-1], loc='upper right', facecolor='white', edgecolor='none', framealpha=0.9, prop={'size': 8}) if(title): ax.set_title(title, size=12) if(side_title): ax.annotate(side_title, (0, 0.5), xytext=(-45, 0), ha='right', va='center', size=12, rotation=90, xycoords='axes fraction', textcoords='offset points') return ax #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def figure_basic(self, plot_S='line', plot_E='line', plot_I_pre='line', plot_I_sym='line', plot_I_asym='line', plot_H='line', plot_R='line', plot_F='line', plot_Q_E='line', plot_Q_pre='line', plot_Q_sym='line', plot_Q_asym='line', plot_Q_S=False, plot_Q_R=False, combine_Q_infected=True, color_S='tab:green', color_E='orange', color_I_pre='tomato', color_I_sym='crimson', color_I_asym='#F0909B', color_H='violet', color_R='tab:blue', color_F='black', color_Q_E='orange', color_Q_pre='tomato', color_Q_sym='crimson', color_Q_asym='#F0909B', color_Q_S='tab:green', color_Q_R='tab:blue', color_Q_infected='tab:purple', color_reference='#E0E0E0', dashed_reference_results=None, dashed_reference_label='reference', shaded_reference_results=None, shaded_reference_label='reference', vlines=[], vline_colors=[], vline_styles=[], vline_labels=[], ylim=None, xlim=None, legend=True, title=None, side_title=None, plot_percentages=True, figsize=(12,8), use_seaborn=True, show=True): import matplotlib.pyplot as pyplot fig, ax = pyplot.subplots(figsize=figsize) if(use_seaborn): import seaborn seaborn.set_style('ticks') seaborn.despine() self.plot(ax=ax, plot_S=plot_S, plot_E=plot_E, plot_I_pre=plot_I_pre, plot_I_sym=plot_I_sym, plot_I_asym=plot_I_asym, plot_H=plot_H, plot_R=plot_R, plot_F=plot_F, plot_Q_E=plot_Q_E, plot_Q_pre=plot_Q_pre, plot_Q_sym=plot_Q_sym, plot_Q_asym=plot_Q_asym, plot_Q_S=plot_Q_S, plot_Q_R=plot_Q_R, combine_Q_infected=combine_Q_infected, color_S=color_S, color_E=color_E, color_I_pre=color_I_pre, color_I_sym=color_I_sym, color_I_asym=color_I_asym, color_H=color_H, color_R=color_R, color_F=color_F, color_Q_E=color_Q_E, color_Q_pre=color_Q_pre, color_Q_sym=color_Q_sym, color_Q_asym=color_Q_asym, color_Q_S=color_Q_S, color_Q_R=color_Q_R, color_Q_infected=color_Q_infected, color_reference=color_reference, dashed_reference_results=dashed_reference_results, dashed_reference_label=dashed_reference_label, shaded_reference_results=shaded_reference_results, shaded_reference_label=shaded_reference_label, vlines=vlines, vline_colors=vline_colors, vline_styles=vline_styles, vline_labels=vline_labels, ylim=ylim, xlim=xlim, legend=legend, title=title, side_title=side_title, plot_percentages=plot_percentages) if(show): pyplot.show() return fig, ax #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def figure_infections(self, plot_S=False, plot_E='stacked', plot_I_pre='stacked', plot_I_sym='stacked', plot_I_asym='stacked', plot_H='stacked', plot_R=False, plot_F='stacked', plot_Q_E='stacked', plot_Q_pre='stacked', plot_Q_sym='stacked', plot_Q_asym='stacked', plot_Q_S=False, plot_Q_R=False, combine_Q_infected=True, color_S='tab:green', color_E='orange', color_I_pre='tomato', color_I_sym='crimson', color_I_asym='#F0909B', color_H='violet', color_R='tab:blue', color_F='black', color_Q_E='orange', color_Q_pre='tomato', color_Q_sym='crimson', color_Q_asym='#F0909B', color_Q_S='tab:green', color_Q_R='tab:blue', color_Q_infected='tab:purple', color_reference='#E0E0E0', dashed_reference_results=None, dashed_reference_label='reference', shaded_reference_results=None, shaded_reference_label='reference', vlines=[], vline_colors=[], vline_styles=[], vline_labels=[], ylim=None, xlim=None, legend=True, title=None, side_title=None, plot_percentages=True, figsize=(12,8), use_seaborn=True, show=True): import matplotlib.pyplot as pyplot fig, ax = pyplot.subplots(figsize=figsize) if(use_seaborn): import seaborn seaborn.set_style('ticks') seaborn.despine() self.plot(ax=ax, plot_S=plot_S, plot_E=plot_E, plot_I_pre=plot_I_pre, plot_I_sym=plot_I_sym, plot_I_asym=plot_I_asym, plot_H=plot_H, plot_R=plot_R, plot_F=plot_F, plot_Q_E=plot_Q_E, plot_Q_pre=plot_Q_pre, plot_Q_sym=plot_Q_sym, plot_Q_asym=plot_Q_asym, plot_Q_S=plot_Q_S, plot_Q_R=plot_Q_R, combine_Q_infected=combine_Q_infected, color_S=color_S, color_E=color_E, color_I_pre=color_I_pre, color_I_sym=color_I_sym, color_I_asym=color_I_asym, color_H=color_H, color_R=color_R, color_F=color_F, color_Q_E=color_Q_E, color_Q_pre=color_Q_pre, color_Q_sym=color_Q_sym, color_Q_asym=color_Q_asym, color_Q_S=color_Q_S, color_Q_R=color_Q_R, color_Q_infected=color_Q_infected, color_reference=color_reference, dashed_reference_results=dashed_reference_results, dashed_reference_label=dashed_reference_label, shaded_reference_results=shaded_reference_results, shaded_reference_label=shaded_reference_label, vlines=vlines, vline_colors=vline_colors, vline_styles=vline_styles, vline_labels=vline_labels, ylim=ylim, xlim=xlim, legend=legend, title=title, side_title=side_title, plot_percentages=plot_percentages) if(show): pyplot.show() return fig, ax #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% ././@PaxHeader0000000000000000000000000000002600000000000011453 xustar000000000000000022 mtime=1597124061.0 seirsplus-1.0.9/seirsplus/networks.py0000644000076500000240000011717600000000000020143 0ustar00ryanstaff00000000000000from __future__ import division import numpy import scipy import networkx from . import FARZ from .models import * import matplotlib.pyplot as pyplot 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, 'b':0, 'epsilon':1e-6, 'directed': False, 'weighted': False}, distancing_scales=[]): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Generate FARZ networks of intra-cohort contacts: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ cohortNetworks = [] teams_indices = {} for i in range(num_cohorts): numNodes = num_nodes_per_cohort[i] if isinstance(num_nodes_per_cohort, list) else num_nodes_per_cohort numTeams = num_teams_per_cohort[i] if isinstance(num_teams_per_cohort, list) else num_teams_per_cohort cohortMeanDegree = mean_intracohort_degree[i] if isinstance(mean_intracohort_degree, list) else mean_intracohort_degree farz_params.update({'n':numNodes, 'k':numTeams, 'm':cohortMeanDegree}) cohortNetwork, cohortTeamLabels = FARZ.generate(farz_params=farz_params) cohortNetworks.append(cohortNetwork) for node, teams in cohortTeamLabels.items(): for team in teams: try: teams_indices['c'+str(i)+'-t'+str(team)].append(node) except KeyError: teams_indices['c'+str(i)+'-t'+str(team)] = [node] #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Establish inter-cohort contacts: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ cohortsAdjMatrices = [networkx.adj_matrix(cohortNetwork) for cohortNetwork in cohortNetworks] workplaceAdjMatrix = scipy.sparse.block_diag(cohortsAdjMatrices) workplaceNetwork = networkx.from_scipy_sparse_matrix(workplaceAdjMatrix) N = workplaceNetwork.number_of_nodes() cohorts_indices = {} cohortStartIdx = -1 cohortFinalIdx = -1 for c, cohortNetwork in enumerate(cohortNetworks): cohortStartIdx = cohortFinalIdx + 1 cohortFinalIdx = cohortStartIdx + cohortNetwork.number_of_nodes() - 1 cohorts_indices['c'+str(c)] = list(range(cohortStartIdx, cohortFinalIdx)) for team, indices in teams_indices.items(): 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): 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) return workplaceNetwork, cohorts_indices, teams_indices # %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% # %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% def generate_demographic_contact_network(N, demographic_data, layer_generator='FARZ', layer_info=None, distancing_scales=[], isolation_groups=[], verbose=False): graphs = {} age_distn = demographic_data['age_distn'] household_size_distn = demographic_data['household_size_distn'] household_stats = demographic_data['household_stats'] ######################################### # Preprocess Demographic Statistics: ######################################### meanHouseholdSize = numpy.average(list(household_size_distn.keys()), weights=list(household_size_distn.values())) # print("mean household size: " + str(meanHouseholdSize)) # Calculate the distribution of household sizes given that the household has more than 1 member: household_size_distn_givenGT1 = {key: value/(1-household_size_distn[1]) for key, value in household_size_distn.items()} household_size_distn_givenGT1[1] = 0 # Percent of households with at least one member under 20: pctHouseholdsWithMember_U20 = household_stats['pct_with_under20'] # Percent of households with at least one member over 60: pctHouseholdsWithMember_O60 = household_stats['pct_with_over60'] # Percent of households with at least one member under 20 AND at least one over 60: pctHouseholdsWithMember_U20andO60 = household_stats['pct_with_under20_over60'] # Percent of SINGLE OCCUPANT households where the occupant is over 60: pctHouseholdsWithMember_O60_givenEq1 = household_stats['pct_with_over60_givenSingleOccupant'] # Average number of members Under 20 in households with at least one member Under 20: meanNumU20PerHousehold_givenU20 = household_stats['mean_num_under20_givenAtLeastOneUnder20'] #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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'] totalPctU20 = numpy.sum([age_distn[bracket] for bracket in ageBrackets_U20]) age_distn_givenU20 = {bracket: pct/totalPctU20 for bracket, pct in age_distn.items() if bracket in ageBrackets_U20} ageBrackets_20to60 = ['20-29', '30-39', '40-49', '50-59'] totalPct20to60 = numpy.sum([age_distn[bracket] for bracket in ageBrackets_20to60]) age_distn_given20to60 = {bracket: pct/totalPct20to60 for bracket, pct in age_distn.items() if bracket in ageBrackets_20to60} ageBrackets_O60 = ['60-69', '70-79', '80+'] totalPctO60 = numpy.sum([age_distn[bracket] for bracket in ageBrackets_O60]) age_distn_givenO60 = {bracket: pct/totalPctO60 for bracket, pct in age_distn.items() if bracket in ageBrackets_O60} ageBrackets_NOTU20 = ['20-29', '30-39', '40-49', '50-59', '60-69', '70-79', '80+'] totalPctNOTU20 = numpy.sum([age_distn[bracket] for bracket in ageBrackets_NOTU20]) age_distn_givenNOTU20 = {bracket: pct/totalPctNOTU20 for bracket, pct in age_distn.items() if bracket in ageBrackets_NOTU20} ageBrackets_NOTO60 = ['0-9', '10-19', '20-29', '30-39', '40-49', '50-59'] totalPctNOTO60 = numpy.sum([age_distn[bracket] for bracket in ageBrackets_NOTO60]) age_distn_givenNOTO60 = {bracket: pct/totalPctNOTO60 for bracket, pct in age_distn.items() if bracket in ageBrackets_NOTO60} #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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['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 householdSituations_prob['NOTu20_NOTo60_eq1'] = (1 - pctHouseholdsWithMember_O60_givenEq1)*prob_eq1 householdSituations_prob['u20_o60_gt1'] = pctHouseholdsWithMember_U20andO60 householdSituations_prob['u20_NOTo60_gt1'] = prob_u20 - householdSituations_prob['u20_o60_gt1'] - householdSituations_prob['u20_NOTo60_eq1'] - householdSituations_prob['u20_o60_eq1'] householdSituations_prob['NOTu20_o60_gt1'] = prob_o60 - householdSituations_prob['u20_o60_gt1'] - householdSituations_prob['NOTu20_o60_eq1'] - householdSituations_prob['u20_o60_eq1'] householdSituations_prob['NOTu20_NOTo60_gt1'] = prob_gt1 - householdSituations_prob['u20_o60_gt1'] - householdSituations_prob['NOTu20_o60_gt1'] - householdSituations_prob['u20_NOTo60_gt1'] assert(numpy.sum(list(householdSituations_prob.values())) == 1.0), "Household situation probabilities must do not sum to 1" ######################################### ######################################### # Randomly construct households following the size and age distributions defined above: ######################################### ######################################### households = [] # List of dicts storing household data structures and metadata 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'): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Household size is definitely 1 household['size'] = 1 #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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'): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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()))) ) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # There's definitely at least one UNDER 20 in this household, add an appropriate age bracket: household['ageBrackets'].append( numpy.random.choice(list(age_distn_givenU20.keys()), p=list(age_distn_givenU20.values())) ) # Figure out how many additional Under 20 to add given there is at least one U20; add them: # > Must leave room for at least one Over 60 (see minmax terms) numAdditionalU20_givenAtLeastOneU20 = min(max(0, numpy.random.poisson(meanNumU20PerHousehold_givenU20-1)), household['size']-len(household['ageBrackets'])-1) for k in range(numAdditionalU20_givenAtLeastOneU20): household['ageBrackets'].append( numpy.random.choice(list(age_distn_givenU20.keys()), p=list(age_distn_givenU20.values())) ) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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'): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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()))) ) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # There's definitely at least one UNDER 20 in this household, add an appropriate age bracket: household['ageBrackets'].append( numpy.random.choice(list(age_distn_givenU20.keys()), p=list(age_distn_givenU20.values())) ) # Figure out how many additional Under 20 to add given there is at least one U20; add them: # > NOT CURRENTLY ASSUMING that there must be at least one non-Under20 member in every household (doing so makes total % U20 in households too low) numAdditionalU20_givenAtLeastOneU20 = min(max(0, numpy.random.poisson(meanNumU20PerHousehold_givenU20-1)), household['size']-len(household['ageBrackets'])) for k in range(numAdditionalU20_givenAtLeastOneU20): household['ageBrackets'].append( numpy.random.choice(list(age_distn_givenU20.keys()), p=list(age_distn_givenU20.values())) ) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # There are no OVER 60 in this household. #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Remaining members can be any age EXCLUDING OVER 60 and 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_given20to60.keys()), p=list(age_distn_given20to60.values())) ) elif(household['situation'] == 'NOTu20_o60_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()))) ) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # There are no UNDER 20 in this household. #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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'): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # 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()))) ) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # There are no UNDER 20 in this household. #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # There are no OVER 60 in this household. #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Remaining household members can be any age BETWEEN 20 TO 60, add as many as needed to meet the household size: 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'): # impossible by assumption # elif(household['situation'] == 'u20_o60_eq1'): # impossible if(len(household['ageBrackets']) == household['size']): homelessNodes -= household['size'] households.append(household) else: print("Household size does not match number of age brackets assigned. "+household['situation']) numHouseholds = len(households) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Check the frequencies of constructed households against the target distributions: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ print("Generated overall age distribution:") for ageBracket in sorted(age_distn): age_freq = numpy.sum([len([age for age in household['ageBrackets'] if age==ageBracket]) for household in households])/N print(str(ageBracket)+": %.4f\t(%.4f from target)" % (age_freq, (age_freq - age_distn[ageBracket])) ) print() print("Generated household size distribution:") for size in sorted(household_size_distn): size_freq = numpy.sum([1 for household in households if household['size']==size])/numHouseholds print(str(size)+": %.4f\t(%.4f from target)" % (size_freq, (size_freq - household_size_distn[size])) ) 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 print("%.4f\t\t(%.4f from target)" % (checkval, checkval - target)) print("Generated mean num members Under 20 given at least one member is Under 20") checkval = numpy.mean([numpy.in1d(household['ageBrackets'], ageBrackets_U20).sum() for household in households if not set(household['ageBrackets']).isdisjoint(ageBrackets_U20)]) target = meanNumU20PerHousehold_givenU20 print("%.4f\t\t(%.4f from target)" % (checkval, checkval - target)) # ######################################### ######################################### # Generate Contact Networks ######################################### ######################################### ######################################### # Generate baseline (no intervention) contact network: ######################################### #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Define the age groups and desired mean degree for each graph layer: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(layer_info is None): # Use the following default data if none is provided: # Data source: https://www.medrxiv.org/content/10.1101/2020.03.19.20039107v1 layer_info = { '0-9': {'ageBrackets': ['0-9'], 'meanDegree': 8.6, 'meanDegree_CI': (0.0, 17.7) }, '10-19': {'ageBrackets': ['10-19'], 'meanDegree': 16.2, 'meanDegree_CI': (12.5, 19.8) }, '20-59': {'ageBrackets': ['20-29', '30-39', '40-49', '50-59'], 'meanDegree': ((age_distn_given20to60['20-29']+age_distn_given20to60['30-39'])*15.3 + (age_distn_given20to60['40-49']+age_distn_given20to60['50-59'])*13.8), '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) }, '60+': {'ageBrackets': ['60-69', '70-79', '80+'], 'meanDegree': 13.9, 'meanDegree_CI': (7.3, 20.5) } } # Count the number of individuals in each age bracket in the generated households: ageBrackets_numInPop = {ageBracket: numpy.sum([len([age for age in household['ageBrackets'] if age==ageBracket]) for household in households]) for ageBracket, __ in age_distn.items()} #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Generate a graph layer for each age group, representing the public contacts for each age group: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ adjMatrices = [] adjMatrices_isolation_mask = [] individualAgeGroupLabels = [] 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']) curidx += layerInfo['numIndividuals'] individualAgeGroupLabels[min(layerInfo['indices']):max(layerInfo['indices'])] = [layerGroup]*layerInfo['numIndividuals'] graph_generated = False graph_gen_attempts = 0 # Note, we generate a graph with average_degree parameter = target mean degree - meanHousehold size # so that when in-household edges are added each graph's mean degree will be close to the target mean targetMeanDegree = layerInfo['meanDegree']-int(meanHouseholdSize) 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): try: 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), 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'], 'm': int(targetMeanDegree/2), # mean degree / 2 'k': int(layerInfo['numIndividuals']/50), # num communities 'alpha': 2.0, # clustering param 'gamma': -0.6, # assortativity param '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, 'directed': False, 'weighted': False}) elif(layer_generator == 'BA'): pass else: print("Layer generator \""+layer_generator+"\" is not recognized (support for 'LFR', 'FARZ', 'BA'") nodeDegrees = [d[1] for d in layerInfo['graph'].degree()] meanDegree = numpy.mean(nodeDegrees) maxDegree = numpy.max(nodeDegrees) # Enforce that the generated graph has mean degree within the 95% CI of the mean for this group in the data: if(meanDegree+meanHouseholdSize >= targetMeanDegreeRange[0] and meanDegree+meanHouseholdSize <= targetMeanDegreeRange[1]): # if(meanDegree+meanHouseholdSize >= targetMeanDegree+meanHouseholdSize-1 and meanDegree+meanHouseholdSize <= targetMeanDegree+meanHouseholdSize+1): if(verbose): print(layerGroup+" public mean degree = "+str((meanDegree))) print(layerGroup+" public max degree = "+str((maxDegree))) adjMatrices.append(networkx.adj_matrix(layerInfo['graph'])) # 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): adjMatrices_isolation_mask.append(numpy.zeros(shape=networkx.adj_matrix(layerInfo['graph']).shape)) else: # adjMatrices_isolation_mask.append(numpy.ones(shape=networkx.adj_matrix(layerInfo['graph']).shape)) # The graph layer we just created represents the baseline (no dist) public connections; # this should be the superset of all connections that exist in any modification of the network, # therefore it should work to use this baseline adj matrix as the mask instead of a block of 1s # (which uses unnecessary memory to store a whole block of 1s, ie not sparse) adjMatrices_isolation_mask.append(networkx.adj_matrix(layerInfo['graph'])) graph_generated = True else: graph_gen_attempts += 1 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]): targetMeanDegree -= 1 if layer_generator=='FARZ' else 0.05 # reload(networkx) if(verbose): # print("Try again... (mean degree = "+str(meanDegree)+"+"+str(meanHouseholdSize)+" is outside the target range for mean degree "+str(targetMeanDegreeRange)+")") print("\tTry again... (mean degree = %.2f+%.2f=%.2f is outside the target range for mean degree (%.2f, %.2f)" % (meanDegree, meanHouseholdSize, meanDegree+meanHouseholdSize, targetMeanDegreeRange[0], targetMeanDegreeRange[1])) # The networks LFR graph generator function has unreliable convergence. # If it fails to converge in allotted iterations, try again to generate. # If it is stuck (for some reason) and failing many times, reload networkx. except networkx.exception.ExceededMaxIterations: graph_gen_attempts += 1 # if(graph_gen_attempts >= 10 and graph_gen_attempts % 10): # reload(networkx) if(verbose): print("\tTry again... (networkx failed to converge on a graph)") #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Assemble an graph for the full population out of the adjacencies generated for each layer: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ A_baseline = scipy.sparse.lil_matrix(scipy.sparse.block_diag(adjMatrices)) # Create a networkx Graph object from the adjacency matrix: G_baseline = networkx.from_scipy_sparse_matrix(A_baseline) graphs['baseline'] = G_baseline ######################################### # Generate social distancing modifications to the baseline *public* contact network: ######################################### # 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) # a number of edges for each node to *keep*, and other edges are removed. #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ G_baseline_NODIST = graphs['baseline'].copy() # Social distancing interactions: for dist_scale in distancing_scales: graphs['distancingScale'+str(dist_scale)] = custom_exponential_graph(G_baseline_NODIST, scale=dist_scale) 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) print([str(unique)+": "+str(count/N) for (unique, count) in zip(unique, counts)]) # pyplot.hist(nodeDegrees_baseline_public_NODIST, bins=range(int(max(nodeDegrees_baseline_public_NODIST))), alpha=0.5, color='tab:blue', label='Public Contacts (no dist)') pyplot.hist(nodeDegrees_baseline_public_DIST, bins=range(int(max(nodeDegrees_baseline_public_DIST))), alpha=0.5, color='tab:purple', label='Public Contacts (distancingScale'+str(dist_scale)+')') pyplot.xlim(0,40) pyplot.xlabel('degree') pyplot.ylabel('num nodes') pyplot.legend(loc='upper right') pyplot.show() ######################################### # 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. #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ 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 # distancing adjacency matrices where the isolation groups are isolated (no public edges), # and create graphs corresponding to the isolation intervention for each distancing level: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ for graphName, graph in graphs.items(): A_withIsolation = scipy.sparse.csr_matrix.multiply( networkx.adj_matrix(graph), A_isolation_mask ) graphs[graphName+'_isolation'] = networkx.from_scipy_sparse_matrix(A_withIsolation) ######################################### ######################################### # Add edges between housemates to strongly connect households: ######################################### ######################################### # Apply to all distancing graphs # Create a copy of the list of node indices for each age group (graph layer) to draw from: for layerGroup, layerInfo in layer_info.items(): layerInfo['selection_indices'] = list(layerInfo['indices']) individualAgeBracketLabels = [None]*N # Go through each household, look up what the age brackets of the members should be, # and randomly select nodes from corresponding age groups (graph layers) to place in the given household. # Strongly connect the nodes selected for each household by adding edges to the adjacency matrix. for household in households: household['indices'] = [] for ageBracket in household['ageBrackets']: ageGroupIndices = next(layer_info[item]['selection_indices'] for item in layer_info if ageBracket in layer_info[item]["ageBrackets"]) memberIndex = ageGroupIndices.pop() household['indices'].append(memberIndex) individualAgeBracketLabels[memberIndex] = ageBracket for memberIdx in household['indices']: nonselfIndices = [i for i in household['indices'] if memberIdx!=i] for housemateIdx in nonselfIndices: # Apply to all distancing graphs for graphName, graph in graphs.items(): graph.add_edge(memberIdx, housemateIdx) ######################################### # Check the connectivity of the fully constructed contacts graphs for each age group's layer: ######################################### if(verbose): for graphName, graph in graphs.items(): nodeDegrees = [d[1] for d in graph.degree()] meanDegree= numpy.mean(nodeDegrees) maxDegree= numpy.max(nodeDegrees) components = sorted(networkx.connected_components(graph), key=len, reverse=True) numConnectedComps = len(components) largestConnectedComp = graph.subgraph(components[0]) print(graphName+": Overall mean degree = "+str((meanDegree))) print(graphName+": Overall max degree = "+str((maxDegree))) print(graphName+": number of connected components = {0:d}".format(numConnectedComps)) print(graphName+": largest connected component = {0:d}".format(len(largestConnectedComp))) for layerGroup, layerInfo in layer_info.items(): nodeDegrees_group = networkx.adj_matrix(graph)[min(layerInfo['indices']):max(layerInfo['indices']), :].sum(axis=1) print("\t"+graphName+": "+layerGroup+" final graph mean degree = "+str(numpy.mean(nodeDegrees_group))) print("\t"+graphName+": "+layerGroup+" final graph max degree = "+str(numpy.max(nodeDegrees_group))) pyplot.hist(nodeDegrees_group, bins=range(int(max(nodeDegrees_group))), alpha=0.5, label=layerGroup) # pyplot.hist(nodeDegrees, bins=range(int(max(nodeDegrees))), alpha=0.5, color='black', label=graphName) pyplot.xlim(0,40) pyplot.xlabel('degree') pyplot.ylabel('num nodes') pyplot.legend(loc='upper right') pyplot.show() ######################################### return graphs, individualAgeBracketLabels, households def household_country_data(country): if(country=='US'): household_data = { 'household_size_distn':{ 1: 0.283708848, 2: 0.345103011, 3: 0.150677793, 4: 0.127649150, 5: 0.057777709, 6: 0.022624223, 7: 0.012459266 }, 'age_distn':{'0-9': 0.121, '10-19': 0.131, '20-29': 0.137, '30-39': 0.133, '40-49': 0.124, '50-59': 0.131, '60-69': 0.115, '70-79': 0.070, '80+' : 0.038 }, 'household_stats':{ 'pct_with_under20': 0.3368, 'pct_with_over60': 0.3801, 'pct_with_under20_over60': 0.0341, 'pct_with_over60_givenSingleOccupant': 0.110, 'mean_num_under20_givenAtLeastOneUnder20': 1.91 } } return household_data #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Defines a random exponential edge pruning mechanism # where the mean degree be easily down-shifted #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ def custom_exponential_graph(base_graph=None, scale=100, min_num_edges=0, m=9, n=None): # If no base graph is provided, generate a random preferential attachment power law graph as a starting point. if(base_graph): graph = base_graph.copy() else: assert(n is not None), "Argument n (number of nodes) must be provided when no base graph is given." graph = networkx.barabasi_albert_graph(n=n, m=m) # We modify the graph by probabilistically dropping some edges from each node. for node in graph: neighbors = list(graph[node].keys()) 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): graph.remove_edge(node, neighbor) return graph #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ #^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ def plot_degree_distn(graph, max_degree=None, show=True, use_seaborn=True): import matplotlib.pyplot as pyplot if(use_seaborn): import seaborn seaborn.set_style('ticks') seaborn.despine() # Get a list of the node degrees: if type(graph)==numpy.ndarray: nodeDegrees = graph.sum(axis=0).reshape((graph.shape[0],1)) # sums of adj matrix cols elif type(graph)==networkx.classes.graph.Graph: nodeDegrees = [d[1] for d in graph.degree()] else: raise BaseException("Input an adjacency matrix or networkx object only.") # Calculate the mean degree: meanDegree = numpy.mean(nodeDegrees) # Generate a histogram of the node degrees: pyplot.hist(nodeDegrees, bins=range(max(nodeDegrees)), alpha=0.75, color='tab:blue', label=('mean degree = %.1f' % meanDegree)) pyplot.xlim(0, max(nodeDegrees) if not max_degree else max_degree) pyplot.xlabel('degree') pyplot.ylabel('num nodes') pyplot.legend(loc='upper right') if(show): pyplot.show() ././@PaxHeader0000000000000000000000000000002600000000000011453 xustar000000000000000022 mtime=1597133814.0 seirsplus-1.0.9/seirsplus/sim_loops.py0000644000076500000240000005753700000000000020277 0ustar00ryanstaff00000000000000from __future__ import division import pickle import numpy import time def run_tti_sim(model, T, intervention_start_pct_infected=0, average_introductions_per_day=0, 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_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, cadence_testing_days=None, cadence_cycle_length=28, temporal_falseneg_rates=None ): #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Testing cadences involve a repeating 28 day cycle starting on a Monday # (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): 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], 'semiweekly': [0, 3, 7, 10, 14, 17, 21, 24], 'weekly': [0, 7, 14, 21], 'biweekly': [0, 14], 'monthly': [0], 'cycle_start': [0] } #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(temporal_falseneg_rates is None): temporal_falseneg_rates = { model.E: {0: 1.00, 1: 1.00, 2: 1.00, 3: 1.00}, model.I_pre: {0: 0.25, 1: 0.25, 2: 0.22}, model.I_sym: {0: 0.19, 1: 0.16, 2: 0.16, 3: 0.17, 4: 0.19, 5: 0.22, 6: 0.26, 7: 0.29, 8: 0.34, 9: 0.38, 10: 0.43, 11: 0.48, 12: 0.52, 13: 0.57, 14: 0.62, 15: 0.66, 16: 0.70, 17: 0.76, 18: 0.79, 19: 0.82, 20: 0.85, 21: 0.88, 22: 0.90, 23: 0.92, 24: 0.93, 25: 0.95, 26: 0.96, 27: 0.97, 28: 0.97, 29: 0.98, 30: 0.98, 31: 0.99}, model.I_asym: {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}, model.Q_E: {0: 1.00, 1: 1.00, 2: 1.00, 3: 1.00}, model.Q_pre: {0: 0.25, 1: 0.25, 2: 0.22}, model.Q_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}, model.Q_asym: {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}, } #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% # Custom simulation loop: #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% interventionOn = False interventionStartTime = None timeOfLastIntervention = -1 timeOfLastIntroduction = -1 testingDays = cadence_testing_days[testing_cadence] cadenceDayNumber = 0 tests_per_day = int(model.numNodes * pct_tested_per_day) max_tracing_tests_per_day = int(tests_per_day * max_pct_tests_for_traces) max_symptomatic_tests_per_day = int(tests_per_day * max_pct_tests_for_symptomatics) tracingPoolQueue = [[] for i in range(tracing_lag)] isolationQueue_symptomatic = [[] for i in range(isolation_lag_symptomatic)] isolationQueue_positive = [[] for i in range(isolation_lag_positive)] isolationQueue_contact = [[] for i in range(isolation_lag_contact)] model.tmax = T running = True while running: running = model.run_iteration() #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Introduce exogenous exposures randomly: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(int(model.t)!=int(timeOfLastIntroduction)): timeOfLastIntroduction = model.t numNewExposures = numpy.random.poisson(lam=average_introductions_per_day) model.introduce_exposures(num_new_exposures=numNewExposures) if(numNewExposures > 0): print("[NEW EXPOSURE @ t = %.2f (%d exposed)]" % (model.t, numNewExposures)) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # Execute testing policy at designated intervals: #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ if(int(model.t)!=int(timeOfLastIntervention)): cadenceDayNumber = int(model.t % cadence_cycle_length) timeOfLastIntervention = model.t 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): 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() nodePositiveStatuses = model.positive.flatten() #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # tracingPoolQueue[0] = tracingPoolQueue[0]Queue.pop(0) #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ newIsolationGroup_symptomatic = [] newIsolationGroup_contact = [] #---------------------------------------- # Isolate SYMPTOMATIC cases without a test: #---------------------------------------- numSelfIsolated_symptoms = 0 numSelfIsolated_symptomaticGroupmate = 0 if(any(isolation_compliance_symptomatic_individual)): symptomaticNodes = numpy.argwhere((nodeStates==model.I_sym)).flatten() for symptomaticNode in symptomaticNodes: if(isolation_compliance_symptomatic_individual[symptomaticNode]): if(model.X[symptomaticNode] == model.I_sym): numSelfIsolated_symptoms += 1 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)): 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]): numSelfIsolated_symptomaticGroupmate += 1 newIsolationGroup_symptomatic.append(isolationGroupmate) #---------------------------------------- # Isolate the CONTACTS of detected POSITIVE cases without a test: #---------------------------------------- numSelfIsolated_positiveContact = 0 numSelfIsolated_positiveContactGroupmate = 0 if(any(isolation_compliance_positive_contact) or any(isolation_compliance_positive_contactgroupmate)): for contactNode in tracingPoolQueue[0]: 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)): 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]): newIsolationGroup_contact.append(isolationGroupmate) numSelfIsolated_positiveContactGroupmate += 1 #---------------------------------------- # Update the nodeStates list after self-isolation updates to model.X: #---------------------------------------- nodeStates = model.X.flatten() #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #---------------------------------------- # Allow SYMPTOMATIC individuals to self-seek tests # regardless of cadence testing days #---------------------------------------- symptomaticSelection = [] if(any(testing_compliance_symptomatic)): symptomaticPool = numpy.argwhere((testing_compliance_symptomatic==True) & (nodeTestedInCurrentStateStatuses==False) & (nodePositiveStatuses==False) & ((nodeStates==model.I_sym)|(nodeStates==model.Q_sym)) ).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)] #---------------------------------------- # Test individuals randomly and via contact tracing # on cadence testing days: #---------------------------------------- tracingSelection = [] randomSelection = [] if(cadenceDayNumber in testingDays): #---------------------------------------- # Apply a designated portion of this day's tests # to individuals identified by CONTACT TRACING: #---------------------------------------- tracingPool = tracingPoolQueue.pop(0) 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) and (testing_compliance_traced[traceNode]==True) and (model.X[traceNode] != model.R) and (model.X[traceNode] != model.Q_R) and (model.X[traceNode] != model.H) and (model.X[traceNode] != model.F)): tracingSelection.append(traceNode) #---------------------------------------- # Apply the remainder of this day's tests to random testing: #---------------------------------------- if(any(testing_compliance_random)): testingPool = numpy.argwhere((testing_compliance_random==True) & (nodePositiveStatuses==False) & (nodeStates != model.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)] #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #---------------------------------------- # Perform the tests on the selected individuals: #---------------------------------------- selectedToTest = numpy.concatenate((symptomaticSelection, tracingSelection, randomSelection)).astype(int) numTested = 0 numTested_random = 0 numTested_tracing = 0 numTested_symptomatic = 0 numPositive = 0 numPositive_random = 0 numPositive_tracing = 0 numPositive_symptomatic = 0 numIsolated_positiveGroupmate = 0 newTracingPool = [] newIsolationGroup_positive = [] for i, testNode in enumerate(selectedToTest): model.set_tested(testNode, True) numTested += 1 if(i < len(symptomaticSelection)): numTested_symptomatic += 1 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): 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 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] 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)): # +++++++++++++++++++++++++++++++++++++++++++++ # The tested node has returned a positive test # +++++++++++++++++++++++++++++++++++++++++++++ numPositive += 1 if(i < len(symptomaticSelection)): numPositive_symptomatic += 1 elif(i < len(symptomaticSelection)+len(tracingSelection)): numPositive_tracing += 1 else: numPositive_random += 1 # Update the node's state to the appropriate detected case state: model.set_positive(testNode, True) #---------------------------------------- # Add this positive node to the isolation group: #---------------------------------------- if(isolation_compliance_positive_individual[testNode]): 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)): 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]): 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()) numpy.random.shuffle(testNodeContacts) if(num_contacts_to_trace is None): numContactsToTrace = int(pct_contacts_to_trace*len(testNodeContacts)) else: numContactsToTrace = num_contacts_to_trace newTracingPool.extend(testNodeContacts[0:numContactsToTrace]) # Add the nodes to be isolated to the isolation queue: isolationQueue_positive.append(newIsolationGroup_positive) isolationQueue_symptomatic.append(newIsolationGroup_symptomatic) isolationQueue_contact.append(newIsolationGroup_contact) # Add the nodes to be traced to the tracing queue: tracingPoolQueue.append(newTracingPool) 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(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)") print("\t"+str(numSelfIsolated_positiveContact) +" will isolate due to positive contact ("+str(numSelfIsolated_positiveContactGroupmate)+" as groupmates of contact)") #---------------------------------------- # Update the status of nodes who are to be isolated: #---------------------------------------- numIsolated = 0 isolationGroup_symptomatic = isolationQueue_symptomatic.pop(0) for isolationNode in isolationGroup_symptomatic: model.set_isolation(isolationNode, True) numIsolated += 1 isolationGroup_contact = isolationQueue_contact.pop(0) for isolationNode in isolationGroup_contact: model.set_isolation(isolationNode, True) numIsolated += 1 isolationGroup_positive = isolationQueue_positive.pop(0) for isolationNode in isolationGroup_positive: model.set_isolation(isolationNode, True) numIsolated += 1 print("\t"+str(numIsolated)+" entered isolation") #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ interventionInterval = (interventionStartTime, model.t) return interventionInterval #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% #%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% ././@PaxHeader0000000000000000000000000000002600000000000011453 xustar000000000000000022 mtime=1597565467.0 seirsplus-1.0.9/seirsplus/utilities.py0000644000076500000240000000602100000000000020264 0ustar00ryanstaff00000000000000import numpy import matplotlib.pyplot as pyplot def gamma_dist(mean, coeffvar, N): scale = mean*coeffvar**2 shape = mean/scale return numpy.random.gamma(scale=scale, shape=shape, size=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']) 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') pyplot.show() def network_info(networks, names=None, plot=False, bin_size=1, colors=None, reverse_plot=False): 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']) 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), 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') pyplot.show() 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) ) ././@PaxHeader0000000000000000000000000000003400000000000011452 xustar000000000000000028 mtime=1597645612.5705636 seirsplus-1.0.9/seirsplus.egg-info/0000755000076500000240000000000000000000000017372 5ustar00ryanstaff00000000000000././@PaxHeader0000000000000000000000000000002600000000000011453 xustar000000000000000022 mtime=1597645612.0 seirsplus-1.0.9/seirsplus.egg-info/PKG-INFO0000644000076500000240000000057200000000000020473 0ustar00ryanstaff00000000000000Metadata-Version: 1.0 Name: seirsplus Version: 1.0.9 Summary: Models of SEIRS epidemic dynamics with extensions, including network-structured populations, testing, contact tracing, and social distancing. Home-page: https://github.com/ryansmcgee/SEIRS-network-model Author: Ryan Seamus McGee Author-email: ryansmcgee@gmail.com License: MIT Description: UNKNOWN Platform: UNKNOWN ././@PaxHeader0000000000000000000000000000002600000000000011453 xustar000000000000000022 mtime=1597645612.0 seirsplus-1.0.9/seirsplus.egg-info/SOURCES.txt0000644000076500000240000000056100000000000021260 0ustar00ryanstaff00000000000000README.md setup.py seirsplus/FARZ.py seirsplus/__init__.py seirsplus/legacy_models.py seirsplus/models.py seirsplus/networks.py seirsplus/sim_loops.py seirsplus/utilities.py seirsplus.egg-info/PKG-INFO seirsplus.egg-info/SOURCES.txt seirsplus.egg-info/dependency_links.txt seirsplus.egg-info/not-zip-safe seirsplus.egg-info/requires.txt seirsplus.egg-info/top_level.txt././@PaxHeader0000000000000000000000000000002600000000000011453 xustar000000000000000022 mtime=1597645612.0 seirsplus-1.0.9/seirsplus.egg-info/dependency_links.txt0000644000076500000240000000000100000000000023440 0ustar00ryanstaff00000000000000 ././@PaxHeader0000000000000000000000000000002600000000000011453 xustar000000000000000022 mtime=1597020578.0 seirsplus-1.0.9/seirsplus.egg-info/not-zip-safe0000644000076500000240000000000100000000000021620 0ustar00ryanstaff00000000000000 ././@PaxHeader0000000000000000000000000000002600000000000011453 xustar000000000000000022 mtime=1597645612.0 seirsplus-1.0.9/seirsplus.egg-info/requires.txt0000644000076500000240000000002500000000000021767 0ustar00ryanstaff00000000000000numpy scipy networkx ././@PaxHeader0000000000000000000000000000002600000000000011453 xustar000000000000000022 mtime=1597645612.0 seirsplus-1.0.9/seirsplus.egg-info/top_level.txt0000644000076500000240000000001200000000000022115 0ustar00ryanstaff00000000000000seirsplus ././@PaxHeader0000000000000000000000000000003300000000000011451 xustar000000000000000027 mtime=1597645612.571457 seirsplus-1.0.9/setup.cfg0000644000076500000240000000004600000000000015470 0ustar00ryanstaff00000000000000[egg_info] tag_build = tag_date = 0 ././@PaxHeader0000000000000000000000000000002600000000000011453 xustar000000000000000022 mtime=1597645522.0 seirsplus-1.0.9/setup.py0000644000076500000240000000124200000000000015360 0ustar00ryanstaff00000000000000import setuptools # with open("README.md", "r") as fh: # long_description = fh.read() setuptools.setup( packages=setuptools.find_packages(), name="seirsplus", version='1.0.9', description='Models of SEIRS epidemic dynamics with extensions, including network-structured populations, testing, contact tracing, and social distancing.', # long_description=long_description, # long_description_content_type="text/markdown", url="https://github.com/ryansmcgee/SEIRS-network-model", author='Ryan Seamus McGee', author_email='ryansmcgee@gmail.com', license='MIT', install_requires=['numpy', 'scipy', 'networkx'], zip_safe=False)