diff --git a/src/_BirthDeath.pyx b/src/_BirthDeath.pyx index 028b7a9..f4590ef 100644 --- a/src/_BirthDeath.pyx +++ b/src/_BirthDeath.pyx @@ -32,7 +32,7 @@ cdef class BirthDeathModel: RndmWrapper seed bint first_simulation, sampling_probability, memory_optimization - Py_ssize_t user_seed, sites, hapNum, currentHapNum, maxHapNum, addMemoryNum, popNum, susNum, bCounter, dCounter, sCounter, mCounter, iCounter, swapLockdown, migPlus, migNonPlus, globalInfectious, countsPerStep, good_attempt + Py_ssize_t user_seed, sites, hapNum, currentHapNum, maxHapNum, addMemoryNum, popNum, susNum, bCounter, dCounter, sCounter, mCounter, iCounter, swapLockdown, migPlus, migNonPlus, globalInfectious, countsPerStep, good_attempt, cCounter, conditionNum double currentTime, totalRate, totalMigrationRate, rn, tau_l Events events @@ -42,11 +42,16 @@ cdef class BirthDeathModel: Lockdowns loc npy_int64[::1] suscType, sizes, totalSusceptible, totalInfectious, lockdownON, hapToNum, numToHap, tree, tree_pop - npy_int64[:,::1] susceptible, infectious, initial_susceptible, initial_infectious + npy_int64[:,::1] susceptible, initial_susceptible, infectious, + npy_int64[:,:,::1] infectByCondition, initial_infectious - double[::1] bRate, dRate, sRate, tmRate, maxEffectiveBirthMigration, suscepCumulTransition, immunePopRate, infectPopRate, popRate, migPopRate, actualSizes, contactDensity, contactDensityBeforeLockdown, contactDensityAfterLockdown, startLD, endLD, samplingMultiplier, times - double[:,::1] mRate, susceptibility, tEventHapPopRate, suscepTransition, immuneSourcePopRate, hapPopRate, migrationRates, effectiveMigration - double[:,:,::1] hapMutType, eventHapPopRate, susceptHapPopRate + double[::1] maxEffectiveBirthMigration, suscepCumulTransition, immunePopRate, infectPopRate, popRate, migPopRate, actualSizes, contactDensity, contactDensityBeforeLockdown, contactDensityAfterLockdown, startLD, endLD, samplingMultiplier, times + double[::1] conditionPopRate, cumuConditionRate + double[:,::1] susceptibility, suscepTransition, immuneSourcePopRate, hapPopRate, migrationRates, effectiveMigration + double[:,::1] conditionRate,massivForThree, conditionPopHapRate, bRate, dRate, sRate, tmRate + double[:,:,::1] hapMutType, susceptHapPopRate, tEventHapPopRate, mRate + double[:,:,::1] conditionPopHapConRate, tEventConHapPopRate + double[:,:,:,::1] eventHapPopRate double[:,:,:,::1] PropensitiesMigr, PropensitiesMutatations double[:,:,::1] PropensitiesSuscep, PropensitiesTransmission @@ -60,7 +65,7 @@ cdef class BirthDeathModel: npy_int64[:,::1] infectiousDelta, susceptibleDelta - def __init__(self, number_of_sites, populations_number, number_of_susceptible_groups, seed, sampling_probability, memory_optimization): + def __init__(self, number_of_sites, populations_number, number_of_susceptible_groups, seed, sampling_probability, memory_optimization, conditionNum): self.user_seed = seed self.seed = RndmWrapper(seed=(self.user_seed, 0)) @@ -101,6 +106,7 @@ cdef class BirthDeathModel: self.sCounter = 0 self.mCounter = 0 self.iCounter = 0 + self.cCounter = 0 self.swapLockdown = 0 self.migPlus = 0 self.migNonPlus = 0 @@ -119,28 +125,50 @@ cdef class BirthDeathModel: # memory_optimization self.infectious = np.zeros((self.popNum, self.maxHapNum), dtype=np.int64) - self.tmRate = np.zeros(self.maxHapNum, dtype=float) - self.tEventHapPopRate = np.zeros((self.popNum, self.maxHapNum), dtype=float) + self.tmRate = np.zeros((self.maxHapNum, self.conditionNum), dtype=float) + self.massivForThree = np.zeros((self.popNum, 3), dtype=float) + self.conditionPopRate = np.zeros(self.popNum, dtype=float) + self.tEventHapPopRate = np.zeros((self.popNum, self.maxHapNum, self.conditionNum), dtype=float) self.hapPopRate = np.zeros((self.popNum, self.maxHapNum), dtype=float) - self.eventHapPopRate = np.zeros((self.popNum, self.maxHapNum, 4), dtype=float) + self.eventHapPopRate = np.zeros((self.popNum, self.maxHapNum, self.conditionNum, 4), dtype=float) self.susceptHapPopRate = np.zeros((self.popNum, self.maxHapNum, self.susNum), dtype=float) self.suscType = np.zeros(self.hapNum, dtype=np.int64) - self.bRate = np.zeros(self.hapNum, dtype=float) - self.dRate = np.zeros(self.hapNum, dtype=float) - self.sRate = np.zeros(self.hapNum, dtype=float) - self.mRate = np.zeros((self.hapNum, self.sites), dtype=float) + self.bRate = np.zeros(self.hapNum, self.conditionNum, dtype=float) + self.dRate = np.zeros(self.hapNum, self.conditionNum, dtype=float) + self.sRate = np.zeros(self.hapNum, self.conditionNum, dtype=float) + self.mRate = np.zeros((self.hapNum, self.conditionNum, self.sites), dtype=float) + self.conditionRate = np.zeros((self.conditionNum, self.conditionNum), dtype=float) + self.cumuConditionRate = np.zeros(self.conditionNum, dtype=float) + self.conditionPopHapRate = np.zeros((self.popNum, self.maxHapNum), dtype=float) + self.conditionPopHapConRate = np.zeros((self.popNum, self.maxHapNum, self.conditionNum),dtype=float) + self.infectByCondition = np.zeros((self.popNum, self.maxHapNum, self.conditionNum),dtype=np.int64) + self.tEventConHapPopRate = np.zeros((self.popNum, self.maxHapNum, self.conditionNum), dtype=float) self.susceptibility = np.zeros((self.hapNum, self.susNum), dtype=float) self.hapMutType = np.ones((self.hapNum, self.sites, 3), dtype=float) for hn in range(self.hapNum): - self.bRate[hn] = 2.0 - self.dRate[hn] = 1.0 - self.sRate[hn] = 0.01 - for s in range(self.sites): - self.mRate[hn, s] = 0.01 + for cn in range(self.conditionNum): + self.bRate[hn, cn] = 2.0 + self.dRate[hn, cn] = 1.0 + self.sRate[hn, cn] = 0.01 + for s in range(self.sites): + self.mRate[hn, cn, s] = 0.01 self.susceptibility[hn, 0] = 1.0 + for cn1 in range(self.conditionNum): + for cn2 in range(self.conditionNum): + if cn1 != cn2: + self.conditionRate[cn1, cn2] = 2 + else: + self.conditionRate[cn1, cn2] = 0 + + for cn1 in range(self.conditionNum): + self.cumuConditionRate[cn1] = 0 + for cn2 in range(self.conditionNum): + self.cumuConditionRate[cn1] += self.conditionRate[cn1, cn2] + + self.sizes = np.zeros(self.popNum, dtype=np.int64) self.totalSusceptible = np.zeros(self.popNum, dtype=np.int64) self.totalInfectious = np.zeros(self.popNum, dtype=np.int64) @@ -148,7 +176,7 @@ cdef class BirthDeathModel: self.susceptible = np.zeros((self.popNum, self.susNum), dtype=np.int64) self.initial_susceptible = np.zeros((self.popNum, self.susNum), dtype=np.int64) - self.initial_infectious = np.zeros((self.popNum, self.hapNum), dtype=np.int64) + self.initial_infectious = np.zeros((self.popNum, self.hapNum, self.conditionNum), dtype=np.int64) self.maxEffectiveBirthMigration = np.zeros(self.popNum, dtype=float) self.suscepCumulTransition = np.zeros(self.susNum, dtype=float) @@ -208,24 +236,26 @@ cdef class BirthDeathModel: if self.susceptible[0, sn] != 0: if self.memory_optimization: self.AddHaplotype(0, 0) - self.NewInfections(0, sn, 0) + self.NewInfections(0, sn, 0, 0) return @cython.boundscheck(False) @cython.wraparound(False) - cdef inline void NewInfections(self, Py_ssize_t pi, Py_ssize_t si, Py_ssize_t hi, Py_ssize_t num=1): + cdef inline void NewInfections(self, Py_ssize_t pi, Py_ssize_t si, Py_ssize_t hi, Py_ssize_t ci, Py_ssize_t num=1): self.susceptible[pi, si] -= num self.totalSusceptible[pi] -= num self.infectious[pi, hi] += num + self.infectByCondition[pi, hi, ci] += num self.totalInfectious[pi] += num self.globalInfectious += num @cython.boundscheck(False) @cython.wraparound(False) - cdef inline void NewRecoveries(self, Py_ssize_t pi, Py_ssize_t si, Py_ssize_t hi, Py_ssize_t num=1): + cdef inline void NewRecoveries(self, Py_ssize_t pi, Py_ssize_t si, Py_ssize_t hi, Py_ssize_t ci, Py_ssize_t num=1): self.susceptible[pi, si] += num self.totalSusceptible[pi] += num self.infectious[pi, hi] -= num + self.infectByCondition[pi, hi, ci] -= num self.totalInfectious[pi] -= num self.globalInfectious -= num @@ -235,12 +265,16 @@ cdef class BirthDeathModel: if self.addMemoryNum + self.maxHapNum > self.hapNum: self.addMemoryNum = self.hapNum - self.maxHapNum self.numToHap = np.concatenate((self.numToHap, np.zeros(self.addMemoryNum, dtype=np.int64))) - self.infectious = np.concatenate((self.infectious, np.zeros((self.popNum, self.addMemoryNum), dtype=np.int64)), axis=1) - self.tmRate = np.concatenate((self.tmRate, np.zeros(self.addMemoryNum, dtype=float))) - self.tEventHapPopRate = np.concatenate((self.tEventHapPopRate, np.zeros((self.popNum, self.addMemoryNum), dtype=float)), axis=1) + self.infectious = np.concatenate((self.infectious, np.zeros((self.popNum, self.addMemoryNum, self.conditionNum), dtype=np.int64)), axis=1) + self.tmRate = np.concatenate((self.tmRate, np.zeros(self.addMemoryNum, self.conditionNum, dtype=float))) + self.tEventHapPopRate = np.concatenate((self.tEventHapPopRate, np.zeros((self.popNum, self.addMemoryNum, self.conditionNum), dtype=float)), axis=1) self.hapPopRate = np.concatenate((self.hapPopRate, np.zeros((self.popNum, self.addMemoryNum), dtype=float)), axis=1) - self.eventHapPopRate = np.concatenate((self.eventHapPopRate, np.zeros((self.popNum, self.addMemoryNum, 4), dtype=float)), axis=1) + self.eventHapPopRate = np.concatenate((self.eventHapPopRate, np.zeros((self.popNum, self.addMemoryNum, self.conditionNum ,4), dtype=float)), axis=1) self.susceptHapPopRate = np.concatenate((self.susceptHapPopRate, np.zeros((self.popNum, self.addMemoryNum, self.susNum), dtype=float)), axis=1) + self.conditionPopHapRate = np.concatenate((self.conditionPopRate, np.zeros((self.popNum, self.addMemoryNum), dtype=float)), axis=1) + self.conditionPopHapConRate = np.concatenate((self.conditionPopRate, np.zeros((self.popNum, self.addMemoryNum, self.conditionNum), dtype=float)), axis=1) + self.infectByCondition = np.concatenate((self.infectByCondition, np.zeros((self.popNum, self.addMemoryNum, self.conditionNum),dtype=np.int64)), axis=1) + self.tEventConHapPopRate = np.concatenate((self.tEventConHapPopRate, np.zeros((self.popNum, self.addMemoryNum, self.conditionNum), dtype=float)),axis=1) self.maxHapNum += self.addMemoryNum @cython.boundscheck(False) @@ -273,23 +307,37 @@ cdef class BirthDeathModel: self.popRate[pn] = 0. for pn in range(self.popNum): for hn in range(self.currentHapNum): # hn - program number - self.tmRate[hn] = 0 - for s in range(self.sites): - self.tmRate[hn] += self.mRate[self.numToHap[hn], s] + for cn in range(self.conditionNum): + self.tmRate[hn, cn] = 0 + - self.eventHapPopRate[pn, hn, 0] = self.BirthRate(pn, hn) - self.eventHapPopRate[pn, hn, 1] = self.dRate[self.numToHap[hn]] - self.eventHapPopRate[pn, hn, 2] = self.sRate[self.numToHap[hn]]*self.samplingMultiplier[pn] - self.eventHapPopRate[pn, hn, 3] = self.tmRate[hn] - self.tEventHapPopRate[pn, hn] = 0 - for i in range(4): - self.tEventHapPopRate[pn, hn] += self.eventHapPopRate[pn, hn, i] - self.hapPopRate[pn, hn] = self.tEventHapPopRate[pn, hn] * self.infectious[pn, hn] + for cn in range(self.conditionNum): + for s in range(self.sites): + self.tmRate[hn, cn] += self.mRate[self.numToHap[hn], cn, s] + self.eventHapPopRate[pn, hn, cn, 0] = self.BirthRate(pn, hn, cn) + self.eventHapPopRate[pn, hn, cn, 1] = self.dRate[self.numToHap[hn], cn] + self.eventHapPopRate[pn, hn, cn, 2] = self.sRate[self.numToHap[hn], cn]*self.samplingMultiplier[pn] + self.eventHapPopRate[pn, hn, cn, 3] = self.tmRate[hn, cn] + self.tEventHapPopRate[pn, hn, cn] = 0 + for i in range(4): + self.tEventHapPopRate[pn, hn, cn] += self.eventHapPopRate[pn, hn, cn, i] + self.tEventConHapPopRate[pn, hn, cn] = self.tEventHapPopRate[pn, hn, cn] * self.infectByCondition[pn, hn, cn] + self.hapPopRate[pn, hn] += self.tEventConHapPopRate[pn, hn, cn] self.infectPopRate[pn] += self.hapPopRate[pn, hn] for sn in range(self.susNum): self.immuneSourcePopRate[pn, sn] = self.suscepCumulTransition[sn]*self.susceptible[pn, sn] self.immunePopRate[pn] += self.immuneSourcePopRate[pn, sn] - self.popRate[pn] = self.infectPopRate[pn] + self.immunePopRate[pn] + for hn in range(self.currentHapNum): + for cn in range(self.conditionNum): + self.conditionPopHapConRate[pn, hn, cn] = self.cumuConditionRate[cn] * self.infectByCondition[pn, hn, cn] + self.conditionPopHapRate[pn, hn] += self.conditionPopHapConRate[pn, hn, cn] + self.conditionPopRate[pn] += self.conditionPopHapRate[pn, hn] + + + self.popRate[pn] = self.infectPopRate[pn] + self.immunePopRate[pn] + self.conditionPopRate[pn] + self.massivForThree[0] = self.infectPopRate[pn] + self.massivForThree[1] = self.immunePopRate[pn] + self.conditionPopRate[2] = self.conditionPopRate[pn] self.totalRate += self.popRate[pn] maxEffectiveMigration = np.zeros(self.popNum, dtype=float) @@ -305,9 +353,10 @@ cdef class BirthDeathModel: maxEffectiveBirth = 0.0 for hn in range(self.currentHapNum): - for sn in range(self.susNum): - if self.bRate[self.numToHap[hn]]*self.susceptibility[self.numToHap[hn], sn] > maxEffectiveBirth: - maxEffectiveBirth = self.bRate[self.numToHap[hn]]*self.susceptibility[self.numToHap[hn], sn] + for cn in range(self.conditionNum): + for sn in range(self.susNum): + if self.bRate[self.numToHap[hn], cn]*self.susceptibility[self.numToHap[hn], sn] > maxEffectiveBirth: + maxEffectiveBirth = self.bRate[self.numToHap[hn], cn]*self.susceptibility[self.numToHap[hn], sn] self.totalMigrationRate = 0.0 for pn in range(self.popNum): @@ -344,7 +393,7 @@ cdef class BirthDeathModel: @cython.boundscheck(False) @cython.wraparound(False) @cython.cdivision(True) - cdef inline double BirthRate(self, Py_ssize_t pi, Py_ssize_t hi): # hi - program number + cdef inline double BirthRate(self, Py_ssize_t pi, Py_ssize_t hi, Py_ssize_t ci): # hi - program number cdef double ps = 0.0 for sn in range(self.susNum): @@ -352,7 +401,12 @@ cdef class BirthDeathModel: for pn in range(self.popNum): ps += self.susceptHapPopRate[pi, hi, sn]*self.migrationRates[pi, pn]*self.migrationRates[pi, pn]*self.contactDensity[pn]/self.actualSizes[pn] - return self.bRate[self.numToHap[hi]]*ps + for cn in range(self.conditionNum): + self.bRate[self.numToHap[hi], cn] = self.bRate[self.numToHap[hi], cn] * ps + cdef double hi_ps = 0.0 + for cn in range(self.conditionNum): + hi_ps += self.bRate[self.numToHap[hi], cn] + return hi_ps @cython.boundscheck(False) @cython.wraparound(False) @@ -400,7 +454,8 @@ cdef class BirthDeathModel: for sn in range(self.susNum): self.initial_susceptible[pn, sn] = self.susceptible[pn, sn] for hn in range(self.hapNum): - self.initial_infectious[pn, hn] = self.infectious[pn, hn] + for cn in range(self.conditionNum): + self.initial_infectious[pn, hn, cn] = self.infectByCondition[pn, hn, cn] self.first_simulation = True for pn in range(self.popNum): self.CheckLockdown(pn) @@ -446,24 +501,27 @@ cdef class BirthDeathModel: if self.totalRate > choose: self.rn = choose / self.totalRate pi, self.rn = fastChoose1(self.popRate, self.totalRate, self.rn) - choose = self.rn * self.popRate[pi] - if self.immunePopRate[pi] > choose: - self.rn = choose / self.immunePopRate[pi] + #choose = self.rn * self.popRate[pi] + ei, self.rn = fastChoose1(self.massivForThree[pi], self.totalRate, self.rn) + if ei == 1: + #self.rn = choose / self.immunePopRate[pi] self.ImmunityTransition(pi) + elif ei == 2: + self.ConditionTransition(pi) else: self.rn = (choose - self.immunePopRate[pi]) / self.infectPopRate[pi] hi, self.rn = fastChoose1(self.hapPopRate[pi], self.infectPopRate[pi], self.rn) # hi - program number - ei, self.rn = fastChoose1(self.eventHapPopRate[pi, hi], self.tEventHapPopRate[pi, hi], self.rn) + ci, self.rn = fastChoose1(self.tEventConHapPopRate[pi, hi], self.hapPopRate[pi, hi], self.rn) + ei, self.rn = fastChoose1(self.tEventHapPopRate[pi, hi], self.tEvenContHapPopRate[pi, hi, ci], self.rn) if ei == BIRTH: - self.Birth(pi, hi) + self.Birth(pi, hi, ci) elif ei == DEATH: - self.Death(pi, hi) + self.Death(pi, hi, ci) elif ei == SAMPLING: - self.Sampling(pi, hi) + self.Sampling(pi, hi, ci) else: - self.Mutation(pi, hi) + self.Mutation(pi, hi, ci) else: - self.rn = (choose - self.totalRate) / self.totalMigrationRate pi = self.GenerateMigration() return pi @@ -474,13 +532,15 @@ cdef class BirthDeathModel: if infect: self.infectPopRate[pi] = 0.0 for hn in range(self.currentHapNum): - self.eventHapPopRate[pi, hn, 0] = self.BirthRate(pi, hn) - tmp = (self.eventHapPopRate[pi, hn, 0] + - self.eventHapPopRate[pi, hn, 1] + - self.eventHapPopRate[pi, hn, 2] + - self.eventHapPopRate[pi, hn, 3] ) - self.tEventHapPopRate[pi, hn] = tmp - self.hapPopRate[pi, hn] = self.tEventHapPopRate[pi, hn] * self.infectious[pi, hn] + for cn in range(self.conditionNum): + self.eventHapPopRate[pi, hn, cn, 0] = self.BirthRate(pi, hn, cn) + tmp = (self.eventHapPopRate[pi, hn, cn, 0] + + self.eventHapPopRate[pi, hn, cn, 1] + + self.eventHapPopRate[pi, hn, cn, 2] + + self.eventHapPopRate[pi, hn, cn, 3] ) + self.tEventHapPopRate[pi, hn, cn] = tmp + self.tEventConHapPopRate[pi, hn, cn] = self.tEventHapPopRate[pi, hn, cn] * self.infectByCondition[pi, hn, cn] + self.hapPopRate[pi, hn] += self.tEventConHapPopRate[pi, hn, cn] self.infectPopRate[pi] += self.hapPopRate[pi, hn] if immune: @@ -489,7 +549,11 @@ cdef class BirthDeathModel: self.immunePopRate[pi] += self.immuneSourcePopRate[pi, sn] if infect or immune: - self.popRate[pi] = self.infectPopRate[pi] + self.immunePopRate[pi] + self.popRate[pi] = 0 + self.popRate[pi] = self.infectPopRate[pi] + self.immunePopRate[pi]+ self.conditionPopRate[pi] + self.massivForThree[0] = self.infectPopRate[pi] + self.massivForThree[1] = self.immunePopRate[pi] + self.massivForThree[2] = self.conditionPopRate[pi] self.totalRate = 0.0 for pn in range(self.popNum): self.totalRate += self.popRate[pn] @@ -520,14 +584,33 @@ cdef class BirthDeathModel: @cython.boundscheck(False) @cython.wraparound(False) - cdef void Birth(self, Py_ssize_t pi, Py_ssize_t hi): # hi - program number + cdef void ConditionTransition(self, Py_ssize_t pi): + cdef: + Py_ssize_t ssi, tsi, hsi + + hsi, self.rn = fastChoose1(self.conditionPopHapRate[pi], self.conditionPopRate[pi], self.rn) + ssi, self.rn = fastChoose1(self.conditionPopHapConRate[pi, hsi], self.conditionPopHapRate[pi, hsi], self.rn) + tsi, self.rn = fastChoose1_skip(self.conditionRate[ssi], self.cumuConditionRate[ssi], self.rn, ssi) + + self.infectByCondition[pi, hsi, ssi] -= 1 + self.infectByCondition[pi, hsi, tsi] += 1 + self.conditionPopHapConRate[pi, hsi, ssi] = self.cumuConditionRate[ssi] * self.infectByCondition[pi, hsi, ssi] + self.conditionPopHapConRate[pi, hsi, tsi] = self.cumuConditionRate[tsi] * self.infectByCondition[pi, hsi, tsi] + self.UpdateRates(pi, False, True, False) + + self.cCounter += 1 + self.events.AddEvent(self.currentTime, CONDITION, ssi, pi, tsi, 0) #ssi, tsi -- condition, add haplotype hsi?? + + @cython.boundscheck(False) + @cython.wraparound(False) + cdef void Birth(self, Py_ssize_t pi, Py_ssize_t hi, Py_ssize_t ci): # hi - program number cdef double ws = 0.0 for sn in range(self.susNum): ws += self.susceptHapPopRate[pi, hi, sn] si, self.rn = fastChoose1(self.susceptHapPopRate[pi, hi], ws, self.rn) - self.NewInfections(pi, si, hi) + self.NewInfections(pi, si, hi, 0) self.immuneSourcePopRate[pi, si] = self.suscepCumulTransition[si]*self.susceptible[pi, si] self.UpdateRates(pi, True, True, True) @@ -536,8 +619,8 @@ cdef class BirthDeathModel: @cython.boundscheck(False) @cython.wraparound(False) - cdef void Death(self, Py_ssize_t pi, Py_ssize_t hi, bint add_event = True): # hi - program number - self.NewRecoveries(pi, self.suscType[self.numToHap[hi]], hi) + cdef void Death(self, Py_ssize_t pi, Py_ssize_t hi, Py_ssize_t ci, bint add_event = True): # hi - program number + self.NewRecoveries(pi, self.suscType[self.numToHap[hi]], hi, ci) self.immuneSourcePopRate[pi, self.suscType[self.numToHap[hi]]] = self.susceptible[pi, self.suscType[self.numToHap[hi]]]*self.suscepCumulTransition[self.suscType[self.numToHap[hi]]] self.UpdateRates(pi, True, True, True) @@ -547,7 +630,7 @@ cdef class BirthDeathModel: @cython.boundscheck(False) @cython.wraparound(False) - cdef void Sampling(self, Py_ssize_t pi, Py_ssize_t hi): # hi - program number + cdef void Sampling(self, Py_ssize_t pi, Py_ssize_t hi, Py_ssize_t ci): # hi - program number self.Death(pi, hi, False) self.sCounter += 1 @@ -556,13 +639,13 @@ cdef class BirthDeathModel: @cython.boundscheck(False) @cython.wraparound(False) @cython.cdivision(True) - cdef void Mutation(self, Py_ssize_t pi, Py_ssize_t hi): # hi - program number + cdef void Mutation(self, Py_ssize_t pi, Py_ssize_t hi, Py_ssize_t ci): # hi - program number cdef: bint check = True Py_ssize_t ohi, mi, digit4, AS, DS, nhi ohi = self.numToHap[hi] # ohi - haplotype - mi, self.rn = fastChoose1(self.mRate[ohi], self.tmRate[hi], self.rn) + mi, self.rn = fastChoose1(self.mRate[ohi, ci], self.tmRate[hi, ci], self.rn) DS, self.rn = fastChoose1(self.hapMutType[ohi, mi], self.hapMutType[ohi, mi, 0] \ + self.hapMutType[ohi, mi, 1] + self.hapMutType[ohi, mi, 2], self.rn) nhi = self.Mutate(ohi, mi, DS) @@ -596,11 +679,12 @@ cdef class BirthDeathModel: tpi, self.rn = fastChoose1(self.migPopRate, self.totalMigrationRate, self.rn) spi, self.rn = fastChoose2_skip(self.totalInfectious, self.globalInfectious-self.totalInfectious[tpi], self.rn, skip=tpi) hi, self.rn = fastChoose2(self.infectious[spi], self.totalInfectious[spi], self.rn) # hi - program number + ci, self.rn = fastChoose1(self.conditionRate[spi], self.cumuConditionRate[spi], self.rn) si, self.rn = fastChoose2(self.susceptible[tpi], self.totalSusceptible[tpi], self.rn) - p_accept = self.effectiveMigration[spi, tpi]*self.bRate[self.numToHap[hi]]*self.susceptibility[self.numToHap[hi], si]/self.maxEffectiveBirthMigration[tpi] + p_accept = self.effectiveMigration[spi, tpi]*self.bRate[self.numToHap[hi], ci]*self.susceptibility[self.numToHap[hi], si]/self.maxEffectiveBirthMigration[tpi] if self.rn < p_accept: - self.NewInfections(tpi, si, hi) + self.NewInfections(tpi, si, hi, ci) self.UpdateRates(tpi, True, True, True) self.migPlus += 1 @@ -647,9 +731,10 @@ cdef class BirthDeathModel: self.susceptible[pn, sn] = self.initial_susceptible[pn, sn] self.totalSusceptible[pn] += self.initial_susceptible[pn, sn] for hn in range(self.hapNum): - self.infectious[pn, hn] = self.initial_infectious[pn, hn] - self.totalInfectious[pn] += self.initial_infectious[pn, hn] - self.globalInfectious += self.initial_infectious[pn, hn] + for cn in range(self.conditionNum): + self.infectious[pn, hn] += self.initial_infectious[pn, hn, cn] + self.totalInfectious[pn] += self.initial_infectious[pn, hn, cn] + self.globalInfectious += self.initial_infectious[pn, hn, cn] for pn in range(self.popNum): self.CheckLockdown(pn) self.UpdateAllRates() @@ -1306,104 +1391,76 @@ cdef class BirthDeathModel: else: self.addMemoryNum = amount - def set_transmission_rate(self, rate, haplotype): - self.check_rate_1(rate, 'transmission') - self.check_haplotype(haplotype) + def set_transmission_rate(self, rate, haplotype, condition): + self.check_value(rate, 'transmission rate') + self.check_index(haplotype, self.hapNum, 'haplotype', True) + self.check_index(condition, self.conditionNum, 'number of states') + haplotypes = self.create_list_for_cycles(haplotype, self.hapNum) + conditions = self.create_list_for_cycles(condition, self.conditionNum) - # for hn in self.get_list_haplotype(haplotype): - # self.bRate[hn] = rate + for hn in haplotypes: + for cn in conditions: + self.bRate[hn, cn] = rate - if isinstance(haplotype, str): - haplotypes = self.create_list_haplotypes(haplotype) - for haplotype in haplotypes: - self.bRate[haplotype] = rate - elif isinstance(haplotype, int): - self.bRate[haplotype] = rate - elif haplotype == None: - for hn in range(self.hapNum): - self.bRate[hn] = rate - def set_recovery_rate(self, rate, haplotype): - self.check_rate_1(rate, 'recovery') - self.check_haplotype(haplotype) + def set_recovery_rate(self, rate, haplotype, condition): + self.check_value(rate, 'recovery rate') + self.check_index(haplotype, self.hapNum, 'haplotype', True) + self.check_index(condition. self.conditionNum, 'number of states') - if isinstance(haplotype, str): - haplotypes = self.create_list_haplotypes(haplotype) - for haplotype in haplotypes: - self.dRate[haplotype] = rate - elif isinstance(haplotype, int): - self.dRate[haplotype] = rate - elif haplotype == None: - for hn in range(self.hapNum): - self.dRate[hn] = rate + haplotypes = self.create_list_for_cycles(haplotype, self.hapNum) + conditions = self.create_list_for_cycles(condition, self.conditionNum) - def set_sampling_rate(self, rate, haplotype): - self.check_haplotype(haplotype) + for hn in haplotypes: + for cn in conditions: + self.dRate[hn, cn] = rate + + def set_sampling_rate(self, rate, haplotype, condition): + self.check_index(haplotype, self.hapNum, 'haplotype', True) + self.check_value(condition, self.conditionNum, 'number of states') + conditions = self.create_list_for_cycles(condition, self.conditionNum) if self.sampling_probability == True: - self.check_rate_2(rate, 'sampling probability') + self.check_value(rate, 'sampling probability', edge=1) - if isinstance(haplotype, str): - haplotypes = self.create_list_haplotypes(haplotype) - for haplotype in haplotypes: - deathRate = self.dRate[haplotype] + self.sRate[haplotype] - self.dRate[haplotype] = (1 - rate) * deathRate - self.sRate[haplotype] = rate * deathRate - elif isinstance(haplotype, int): - deathRate = self.dRate[haplotype] + self.sRate[haplotype] - self.dRate[haplotype] = (1 - rate) * deathRate - self.sRate[haplotype] = rate * deathRate - elif haplotype == None: - for hn in range(self.hapNum): - deathRate = self.dRate[hn] + self.sRate[hn] - self.dRate[hn] = (1 - rate) * deathRate - self.sRate[hn] = rate * deathRate + haplotypes = self.create_list_for_cycles(haplotype, self.hapNum) + for hn in haplotypes: + for cn in conditions: + deathRate = self.dRate[hn, cn] + self.sRate[hn, cn] + self.dRate[hn, cn] = (1 - rate) * deathRate + self.sRate[hn, cn] = rate * deathRate elif self.sampling_probability == False: - self.check_rate_1(rate, 'sampling') - - if isinstance(haplotype, str): - haplotypes = self.create_list_haplotypes(haplotype) - for haplotype in haplotypes: - self.sRate[haplotype] = rate - elif isinstance(haplotype, int): - self.sRate[haplotype] = rate - elif haplotype == None: - for hn in range(self.hapNum): - self.sRate[hn] = rate + self.check_value(rate, 'sampling rate') - def set_mutation_rate(self, rate, probabilities, haplotype, mutation): - self.check_rate_none_1(rate, 'mutation rate') - self.check_list(probabilities, 'probabilities list', 4) - if isinstance(probabilities, list): - for i in range(4): - self.check_rate_1(probabilities[i], 'mutation probabilities') - self.check_haplotype(haplotype) - self.check_index(mutation, self.sites, 'mutation site') + haplotypes = self.create_list_for_cycles(haplotype, self.hapNum) + for hn in haplotypes: + for cn in conditions: + self.sRate[hn, cn] = rate - if isinstance(haplotype, str): - haplotypes = self.create_list_haplotypes(haplotype) - if isinstance(mutation, int): - for hn in haplotypes: - self.set_probabilities(hn, mutation, rate, probabilities) - elif mutation == None: - for hn in haplotypes: - for s in range(self.sites): - self.set_probabilities(hn, s, rate, probabilities) - elif isinstance(haplotype, int): - if isinstance(mutation, int): - self.set_probabilities(haplotype, mutation, rate, probabilities) - elif mutation == None: - for s in range(self.sites): - self.set_probabilities(haplotype, s, rate, probabilities) - elif haplotype == None: - if isinstance(mutation, int): - for hn in range(self.hapNum): - self.set_probabilities(hn, mutation, rate, probabilities) - elif mutation == None: - for hn in range(self.hapNum): - for s in range(self.sites): - self.set_probabilities(hn, s, rate, probabilities) + def set_mutation_rate(self, rate, haplotype, mutation, condition): + self.check_value(rate, 'mutation rate') + self.check_index(haplotype, self.hapNum, 'haplotype', True) + self.check_index(mutation, self.sites, 'mutation site') + self.check_value(condition, self.conditionNum, 'number of states') + + haplotypes = self.create_list_for_cycles(haplotype, self.hapNum) + conditions = self.create_list_for_cycles(condition, self.conditionNum) + sites = self.create_list_for_cycles(mutation, self.sites) + for hn in haplotypes: + for cn in conditions: + for s in sites: + self.mRate[hn, cn, s] = rate + + @property + def condition_rate(self): + return self.conditionRate + + def set_condition_rate(self, rate, c1, c2): + self.check_value(rate, 'condition_rate') + self.check_index(c1, self.conditionNum, 'first_state') + self.check_index(c2, self.conditionNum, 'second_state') + self.conditionRate[c1, c2] = rate def set_probabilities(self, haplotype, mutation, rate, probabilities): if isinstance(rate, (int, float)): @@ -1526,7 +1583,7 @@ cdef class BirthDeathModel: self.initial_susceptible[pn, source_type] -= amount self.initial_susceptible[pn, target_type] += amount - def set_infectious(self, amount, source_type, target_haplotype, population): + def set_infectious(self, amount, source_type, target_haplotype, population, condition): if self.first_simulation: raise ValueError('This function is available only before first simulation!') self.check_amount(amount) @@ -1538,13 +1595,14 @@ cdef class BirthDeathModel: self.check_amount_inf(amount, source_type, target_haplotype, population) self.initial_susceptible[population, source_type] -= amount - self.initial_infectious[population, target_haplotype] += amount + + self.initial_infectious[population, target_haplotype, condition] += amount elif population == None: for pn in range(self.popNum): self.check_amount_inf(amount, source_type, target_haplotype, pn) self.initial_susceptible[pn, source_type] -= amount - self.initial_infectious[pn, target_haplotype] += amount + self.initial_infectious[pn, target_haplotype, condition] += amount def set_contact_density(self, value, population): self.check_rate_1(value, 'contact density') @@ -2264,11 +2322,12 @@ cdef class BirthDeathModel: continue for sn in range(self.susNum): for hn in range(self.hapNum): - self.PropensitiesMigr[spn, tpn, sn, hn] = self.effectiveMigration[tpn, spn]*self.susceptible[tpn, sn]*\ - self.infectious[spn, hn]*self.bRate[hn]*self.susceptibility[hn, sn]*self.migrationRates[spn, spn] + for cn in range(self.conditionNum): + self.PropensitiesMigr[spn, tpn, sn, hn] = self.effectiveMigration[tpn, spn]*self.susceptible[tpn, sn]*\ + self.infectious[spn, hn]*self.bRate[hn, cn]*self.susceptibility[hn, sn]*self.migrationRates[spn, spn] - self.infectiousAuxTau[tpn, hn] += self.PropensitiesMigr[spn, tpn, sn, hn] - self.susceptibleAuxTau[tpn, sn] -= self.PropensitiesMigr[spn, tpn, sn, hn] + self.infectiousAuxTau[tpn, hn] += self.PropensitiesMigr[spn, tpn, sn, hn] + self.susceptibleAuxTau[tpn, sn] -= self.PropensitiesMigr[spn, tpn, sn, hn] for pn in range(self.popNum): #Susceptibility transition @@ -2284,38 +2343,42 @@ cdef class BirthDeathModel: #Infectious-realted event for hn in range(self.hapNum): #Recovery - self.PropensitiesRecovery[pn, hn] = self.dRate[hn]*self.infectious[pn, hn] + for cn in range(self.conditionNum): + self.PropensitiesRecovery[pn, hn] = self.dRate[hn, cn]*self.infectious[pn, hn] self.susceptibleAuxTau[pn, self.suscType[hn]] += self.PropensitiesRecovery[pn, hn] self.infectiousAuxTau[pn, hn] -= self.PropensitiesRecovery[pn, hn] #Sampling - self.PropensitiesSampling[pn, hn] = self.sRate[hn]*self.infectious[pn, hn]*self.samplingMultiplier[pn] + for cn in range(self.conditionNum): + self.PropensitiesSampling[pn, hn] = self.sRate[hn, cn]*self.infectious[pn, hn]*self.samplingMultiplier[pn] self.susceptibleAuxTau[pn, self.suscType[hn]] += self.PropensitiesSampling[pn, hn] self.infectiousAuxTau[pn, hn] -= self.PropensitiesSampling[pn, hn] #Mutation for s in range(self.sites): - for i in range(3): - self.PropensitiesMutatations[pn, hn, s, i] = self.mRate[hn, s]*self.hapMutType[hn, s, i]/\ - (self.hapMutType[hn, s, 0] + self.hapMutType[hn, s, 1] + self.hapMutType[hn, s, 2])*self.infectious[pn, hn] + for cn in range(self.conditionNum): + for i in range(3): + self.PropensitiesMutatations[pn, hn, s, i] = self.mRate[hn, cn, s]*self.hapMutType[hn, s, i]/\ + (self.hapMutType[hn, s, 0] + self.hapMutType[hn, s, 1] + self.hapMutType[hn, s, 2])*self.infectious[pn, hn] - self.infectiousAuxTau[pn, self.Mutate(hn, s, i)] += self.PropensitiesMutatations[pn, hn, s, i] - self.infectiousAuxTau[pn, hn] -= self.PropensitiesMutatations[pn, hn, s, i] + self.infectiousAuxTau[pn, self.Mutate(hn, s, i)] += self.PropensitiesMutatations[pn, hn, s, i] + self.infectiousAuxTau[pn, hn] -= self.PropensitiesMutatations[pn, hn, s, i] #Transmission for tpn in range(self.popNum): for hn in range(self.hapNum): - for sn in range(self.susNum): - self.PropensitiesTransmission[tpn, hn, sn] = 0.0 - for spn in range(self.popNum): - self.PropensitiesTransmission[tpn, hn, sn] += self.bRate[hn]*self.susceptibility[hn, sn]*self.migrationRates[tpn, spn]*\ - self.migrationRates[tpn, spn]*self.contactDensity[spn]*self.susceptible[tpn, sn]*self.infectious[tpn, hn]/\ - self.actualSizes[spn] - - self.infectiousAuxTau[tpn, hn] += self.PropensitiesTransmission[tpn, hn, sn] - self.susceptibleAuxTau[tpn, sn] -= self.PropensitiesTransmission[tpn, hn, sn] + for cn in range(self.conditionNum): + for sn in range(self.susNum): + self.PropensitiesTransmission[tpn, hn, sn] = 0.0 + for spn in range(self.popNum): + self.PropensitiesTransmission[tpn, hn, sn] += self.bRate[hn, cn]*self.susceptibility[hn, sn]*self.migrationRates[tpn, spn]*\ + self.migrationRates[tpn, spn]*self.contactDensity[spn]*self.susceptible[tpn, sn]*self.infectious[tpn, hn]/\ + self.actualSizes[spn] + + self.infectiousAuxTau[tpn, hn] += self.PropensitiesTransmission[tpn, hn, sn] + self.susceptibleAuxTau[tpn, sn] -= self.PropensitiesTransmission[tpn, hn, sn] @cython.cdivision(True) cdef Py_ssize_t Mutate(self, Py_ssize_t hi, Py_ssize_t s, Py_ssize_t DS): diff --git a/src/_interface.py b/src/_interface.py index be6dba3..ca2eb3b 100644 --- a/src/_interface.py +++ b/src/_interface.py @@ -7,7 +7,7 @@ import time class Simulator: - def __init__(self, number_of_sites=0, populations_number=1, number_of_susceptible_groups=1, seed=None, sampling_probability=False, memory_optimization=False): + def __init__(self, number_of_sites=0, populations_number=1, number_of_susceptible_groups=1, seed=None, sampling_probability=False, memory_optimization=False, conditionNum=1): self.fig = None if seed == None: seed = int(randrange(sys.maxsize)) @@ -15,7 +15,7 @@ def __init__(self, number_of_sites=0, populations_number=1, number_of_susceptibl self.simulation = BirthDeathModel(number_of_sites=number_of_sites, populations_number=populations_number, \ number_of_susceptible_groups=number_of_susceptible_groups, seed=seed, sampling_probability=sampling_probability, \ - memory_optimization=memory_optimization) + memory_optimization=memory_optimization, conditionNum=conditionNum) def print_basic_parameters(self): @@ -43,17 +43,24 @@ def set_step_haplotype(self, amount): self.simulation.set_step_haplotype(amount) - def set_transmission_rate(self, rate, haplotype=None): - self.simulation.set_transmission_rate(rate, haplotype) + def set_transmission_rate(self, rate, haplotype=None, condition=None): + self.simulation.set_transmission_rate(rate, haplotype, condition) - def set_recovery_rate(self, rate, haplotype=None): - self.simulation.set_recovery_rate(rate, haplotype) + def set_recovery_rate(self, rate, haplotype=None, condition=None): + self.simulation.set_recovery_rate(rate, haplotype, condition) - def set_sampling_rate(self, rate, haplotype=None): - self.simulation.set_sampling_rate(rate, haplotype) + def set_sampling_rate(self, rate, haplotype=None, condition=None): + self.simulation.set_sampling_rate(rate, haplotype, condition) - def set_mutation_rate(self, rate=None, probabilities=None, haplotype=None, mutation=None): - self.simulation.set_mutation_rate(rate, probabilities, haplotype, mutation) + def set_mutation_rate(self, rate=None, probabilities=None, haplotype=None, mutation=None, condition=None): + self.simulation.set_mutation_rate(rate, probabilities, haplotype, mutation, condition) + + @property + def condition_rate(self): + return self.simulation.condition_rate + + def set_condition_rate(self, rate, c1=0, c2=0): + self.simulation.set_condition_rate(rate, c1, c2) def set_susceptibility_type(self, susceptibility_type, haplotype=None): diff --git a/src/events.pxi b/src/events.pxi index 5a8cd7e..015b0f5 100644 --- a/src/events.pxi +++ b/src/events.pxi @@ -6,6 +6,7 @@ DEF MUTATION = 3 DEF SUSCCHANGE = 4 DEF MIGRATION = 5 DEF MULTITYPE = 6 +DEF CONDITION = 7 cdef class Event: cdef: @@ -94,6 +95,8 @@ cdef class multiEvent: tn = "SUS" elif self.type_ == MIGRATION: tn = "MIG" + elif self.type_ == CONDITION: + tn = "CON" print("num=", self.num, " time=", self.time, " type=", tn,