Cache njite commands
@@ -68,7 +68,7 @@ def sample(dist=None, par1=None, par2=None, size=None):
def set_seed(seed=None):
''' Reset the random seed -- complicated because of Numba '''
- @nb.njit((nb.int32,))
+ @nb.njit((nb.int32,), cache=True)
def set_seed_numba(seed):
return np.random.seed(seed)
@@ -151,7 +151,7 @@ def multinomial(probs, repeats):
return np.searchsorted(np.cumsum(probs), np.random.random(repeats))
-@nb.njit((nb.int32,)) # This hugely increases performance
+@nb.njit((nb.int32,), cache=True) # This hugely increases performance
def poisson(rate):
''' A Poisson trial '''
return np.random.poisson(rate, 1)[0]
@@ -163,7 +163,7 @@ def binomial_filter(prob, arr):
return arr[(np.random.random(len(arr)) < prob).nonzero()[0]]
-@nb.njit((nb.int32, nb.int32)) # This hugely increases performance
+@nb.njit((nb.int32, nb.int32), cache=True) # This hugely increases performance
def choose(max_n, n):
'''
Choose a subset of items (e.g., people) without replacement.
@@ -178,7 +178,7 @@ def choose(max_n, n):
return np.random.choice(max_n, n, replace=False)
-@nb.njit((nb.int32, nb.int32)) # This hugely increases performance
+@nb.njit((nb.int32, nb.int32), cache=True) # This hugely increases performance
def choose_r(max_n, n):
'''
Choose a subset of items (e.g., people), with replacement.
@@ -218,7 +218,7 @@ def choose_w(probs, n, unique=True):
#%% The core Covasim functions -- compute the infections
-@nb.njit(( nb.float32[:], nb.float32[:], nb.bool_[:], nb.bool_[:], nb.bool_[:], nb.float32, nb.float32, nb.float32))
+@nb.njit(( nb.float32[:], nb.float32[:], nb.bool_[:], nb.bool_[:], nb.bool_[:], nb.float32, nb.float32, nb.float32), cache=True)
def compute_probs(rel_trans, rel_sus, symp, diag, quar, asymp_factor, diag_factor, quar_trans):
''' Calculate relative transmissibility and susceptibility '''
f_asymp = symp + ~symp * asymp_factor # Asymptomatic factor, changes e.g. [0,1] with a factor of 0.8 to [0.8,1.0]
@@ -229,7 +229,7 @@ def compute_probs(rel_trans, rel_sus, symp, diag, qua
return rel_trans, rel_sus
-@nb.njit(( nb.float32, nb.int32[:], nb.int32[:], nb.float32[:], nb.float32[:], nb.float32[:]))
+@nb.njit(( nb.float32, nb.int32[:], nb.int32[:], nb.float32[:], nb.float32[:], nb.float32[:]), cache=True)
def compute_targets(beta, sources, targets, layer_betas, rel_trans, rel_sus):
''' The heaviest step of the model -- figure out who gets infected on this timestep '''
betas = beta * layer_betas * rel_trans[sources] * rel_sus[targets] # Calculate the raw transmission probabilities
