from dataclasses import dataclass as _dataclass
from typing import List as _List
from typing import Union as _Union
from ._network import Network
from ._networks import Networks
__all__ = ["Workspace"]
[docs]@_dataclass
class Workspace:
"""This class provides a workspace for the running calculation.
This pre-allocates all of the memory into arrays, which
can then be used via cython memory views
"""
#: Number of disease classes (stages)
n_inf_classes: int = 0
#: Number of wards (nodes)
nnodes: int = 0
#: Size of population in each disease stage for work infections
inf_tot: _List[int] = None
#: Size of population in each disease stage for play infections
pinf_tot: _List[int] = None
#: Number of wards with at least one individual in this disease stage
n_inf_wards: _List[int] = None
#: Total number of infections in each ward over the last day
#: This is also equal to the prevalence
total_inf_ward: _List[int] = None
#: Number of new infections in each ward over the last day
total_new_inf_ward: _List[int] = None
#: The incidence of the infection (sum of infections up to
#: disease_class == 2)
incidence: _List[int] = None
#: The size of the S population in each ward
S_in_wards: _List[int] = None
#: The size of the E population in each ward
E_in_wards: _List[int] = None
#: The size of the I population in each ward
I_in_wards: _List[int] = None
#: The size of the R population in each ward
R_in_wards: _List[int] = None
#: The sub-workspaces used for the subnets of a
#: multi-demographic Networks (list[Workspace])
subspaces = None
[docs] @staticmethod
def build(network: _Union[Network, Networks]):
"""Create the workspace needed to run the model for the
passed network
"""
params = network.params
workspace = Workspace()
if isinstance(network, Network):
n_inf_classes = params.disease_params.N_INF_CLASSES()
workspace.n_inf_classes = params.disease_params.N_INF_CLASSES()
workspace.nnodes = network.nnodes
size = workspace.nnodes + 1 # 1-indexed
from .utils._array import create_int_array
workspace.inf_tot = create_int_array(n_inf_classes, 0)
workspace.pinf_tot = create_int_array(n_inf_classes, 0)
workspace.n_inf_wards = create_int_array(n_inf_classes, 0)
workspace.total_inf_ward = create_int_array(size, 0)
workspace.total_new_inf_ward = create_int_array(size, 0)
workspace.incidence = create_int_array(size, 0)
workspace.S_in_wards = create_int_array(size, 0)
workspace.E_in_wards = create_int_array(size, 0)
workspace.I_in_wards = create_int_array(size, 0)
workspace.R_in_wards = create_int_array(size, 0)
elif isinstance(network, Networks):
workspace = Workspace.build(network.overall)
subspaces = []
for subnet in network.subnets:
subspaces.append(Workspace.build(subnet))
workspace.subspaces = subspaces
return workspace
[docs] def zero_all(self, zero_subspaces=True):
"""Reset the values of all of the arrays to zero.
By default we zero the subspace networks
(change this by setting zero_subspaces to False)
"""
for i in range(0, self.n_inf_classes):
self.inf_tot[i] = 0
self.pinf_tot[i] = 0
self.n_inf_wards[i] = 0
for i in range(0, self.nnodes+1):
self.total_inf_ward[i] = 0
self.total_new_inf_ward[i] = 0
self.incidence[i] = 0
self.S_in_wards[i] = 0
self.E_in_wards[i] = 0
self.I_in_wards[i] = 0
self.R_in_wards[i] = 0
if zero_subspaces and self.subspaces is not None:
for subspace in self.subspaces:
subspace.zero_all()