coop_merging.py

import numpy as np
import scipy
import matplotlib.pyplot as plt
import math 
import random
import copy
import statistics
from scipy.optimize import curve_fit
from scipy.integrate import odeint
from numba import jit
import numba as nb
from numbalsoda import lsoda_sig, lsoda
from numba import njit, cfunc

#Deterministic Growth
@cfunc(lsoda_sig)
def rhs(t, N, dN, par):
    dN[0]=N[0]*(1 - par[0] + par[0]*(par[1]*(N[0])/(N[0]+N[1])-par[2]))*(1-(N[0]+N[1])/par[3])
    dN[1]=N[1]*(1 - par[0] + par[0]*(par[1]*N[0]/(N[0]+N[1])))*(1-(N[0]+N[1])/par[3])

funcptr = rhs.address

#Deterministic Growth
@njit
def growth_event(in_numbers,w,b,c,t,K):   
    N_demes,N_types=np.shape(in_numbers)
    numbers=np.zeros((N_demes,N_types),dtype=np.float64)
    par=np.array([w,b,c,K])
    t_values = np.linspace(0,t,1000)
    for i in range(N_demes):
        usol = lsoda(funcptr, (in_numbers[i,:]).astype(np.float64), t_values, data = par)[0]
        numbers[i,:]=(usol[-1])
    return numbers

#Stochastic Growth

@njit
def count_nb_ele_nonzero(arr): 
    compt = 0
    for i in range(len(arr)):
        if arr[i] != 0:
            compt += 1
    return compt
    
@njit
def put_to0_negative_ele(arr):
    for i in range(len(arr)):
        if arr[i] < 0:
            arr[i] = 0
    return arr

@njit
def initialization10(Nm0, Nw0, D): #We start from one fully mutant deme in a wild-type structure; the mutant deme corresponds to the element of index 0 in the arrays Nm and Nw
    Nm    = np.zeros(D)
    Nm[0] = Nm0
    Nw    = np.zeros(D)
    for i in range(D):
        Nw[i]=(Nm0+Nw0)-Nm[i]
    #print('Initial state  ', 'Nm0: ', Nm, ' ', 'Nw0: ', Nw)
    return Nm, Nw
    
@njit
def fitnessM(b, c, w, Nm, Nw, D): #We create a np array containing the mutant fitnesses for each deme 
    fm = np.zeros(D)
    for i in range(D): 
        fm[i] = 1 - w + w * ((b*(Nm[i])/(Nm[i]+Nw[i])-c))
    return fm

@njit
def fitnessW(b, c, w, Nm, Nw, D): #We create a np array containing the wild-type fitnesses for each deme 
    fw = np.zeros(D)
    for i in range(D):
        fw[i] = 1 - w + w * b*Nm[i]/(Nm[i]+Nw[i])
    return fw
    
    
@njit
def compute_division_reactionrate(Nm, Nw, f, K, D):
    k = np.zeros(D)
    for i in range(D):
        k[i] = f[i]*(1-(Nw[i]+Nm[i])/K) 
    return k
    
@njit
def compute_total_rate(Nm, Nw, kw_p, km_p, D):
    #kw_p: division reaction rate for the wild-types
    #km_p: division reaction rate for the mutants
    
    ktot = 0
    for i in range(D):
        ktot += kw_p[i]*Nw[i] + km_p[i]*Nm[i]
    return ktot

@njit
def build_changetower(D): #We create an array change_tower in which we store all modifications which can happen at one Gillespie step 
    nb_reac = 2*D
    ind_arr = np.array([D, 2*D], dtype=np.int64)
    change_tower = np.empty((nb_reac,2), dtype=np.int64)
    change_tower[0,:] = [0,+1] 
    
    for i in range(1,ind_arr[0]):
        change_tower[i,:] = [0,+1] #Means here that we have 1 more wild-type in the deme i as we are looking at wild-type division reactions 
    
    for i in range(ind_arr[0], ind_arr[1]):
        change_tower[i,:] = [+1,0] #Means here that we have 1 more mutant in the deme i as we are looking at mutant division reactions

    return change_tower
    
@njit
def build_samplingtower(Nm, Nw, kw_p, km_p, D): #We create an array reac_tower in which we store all the possible reactions and from which we will randomly choose an event to happen in the simulation 
    nb_reac = 2*D  #Total number of possible reactions in Gillespie algorithm
    ind_arr = np.array([D, 2*D], dtype=np.int64) #Reactions are grouped by package, first the mutants divisions in the D demes, second the wild-types divisions in the D demes   
    reac_tower = np.zeros(nb_reac, dtype=np.float64)  
    
    ktot = compute_total_rate(Nm, Nw, kw_p, km_p, D)

    reac_tower[0] = kw_p[0]*Nw[0]/ktot

    for i in range(1,ind_arr[0]):
        reac_tower[i] = reac_tower[i-1] + kw_p[i]*Nw[i]/ktot
    
    for i in range(ind_arr[0], ind_arr[1]):
        reac_tower[i] = reac_tower[i-1] + km_p[i-ind_arr[0]]*Nm[i-ind_arr[0]]/ktot

    return reac_tower
    
#Gillespie Algorithm

@njit
def finalpopulation(b, c, w, Nm, Nw, K, D, t): 
    nb_loop = 0
    time = 0

    # Initialization
    N_C = np.zeros((D, 0), dtype=np.float64)
    N_F = np.zeros((D, 0), dtype=np.float64)
    times = [0.0] 

    ind_arr = np.array([D, 2*D])
    change_tower = build_changetower(D)
    
    while time < t:
        nb_loop += 1

        fm = fitnessM(b, c, w, Nm, Nw, D)
        fw = fitnessW(b, c, w, Nm, Nw, D)

        kw_p = compute_division_reactionrate(Nm, Nw, fw, K, D)
        kw_p = put_to0_negative_ele(kw_p)

        km_p = compute_division_reactionrate(Nm, Nw, fm, K, D)
        km_p = put_to0_negative_ele(km_p)

        ktot = compute_total_rate(Nm, Nw, kw_p, km_p, D)
        reac_tower = build_samplingtower(Nm, Nw, kw_p, km_p, D)

        r = np.random.uniform(0, 1) 
        dt = (1 / ktot) * np.log(1 / r)
        time += dt

        r_ = np.random.uniform(0, 1)
        ir = 0
        while reac_tower[ir] <= r_:
            ir += 1  
            
        if ir < ind_arr[1]:
            addNm = change_tower[ir, 0]
            addNw = change_tower[ir, 1]

            if ir < D:
                Nm[ir] += addNm
                Nw[ir] += addNw
            elif ir < 2 * D:
                Nm[ir - D] += addNm
                Nw[ir - D] += addNw

        times.append(time)

        Nm_reshaped = np.copy(Nm).reshape(D, 1)
        Nw_reshaped = np.copy(Nw).reshape(D, 1)

        N_C = np.hstack((N_C, Nm_reshaped))
        N_F = np.hstack((N_F, Nw_reshaped))

    return N_C, N_F, np.array(times) 


@njit
def growth_stochastique(in_numbers, b, c, w, K, t):
    in_numbers = in_numbers.astype(np.float64)

    D, N_types = in_numbers.shape
    Nm0 = in_numbers[:, 0]
    Nw0 = in_numbers[:, 1]

    mutants, cheaters, temps = finalpopulation(b, c, w, Nm0, Nw0, K, D, t)      
    return mutants, cheaters, temps

#End Stochastic Growth

#Migration and dilution event, simpson=True will sample from each deme proprotionally to its size after growth. 
#Here we implement the binomial version, as described in the paper.
@njit
def dilution_migration_event(in_numbers,migration_matrix,Nmin,simpson=False):
    N_demes, N_types= np.shape(in_numbers)
    #Final numbers that will be kept at the end of the migration/dilution phase
    new_numbers=np.zeros((N_demes, N_types), dtype=np.int64)
     
    for j in np.arange(N_demes):
        Nj=Nmin
        probabilities = np.zeros((N_demes,2),dtype=np.float64)
        for k in range(N_demes):
            probabilities[k,0]= max(min((in_numbers[k,0]/np.sum(in_numbers[k,:]))*migration_matrix[k,j],1),0)
            probabilities[k,1]= max(min((1-in_numbers[k,0]/np.sum(in_numbers[k,:]))*migration_matrix[k,j],1),0)
        proba = np.reshape(probabilities,2*N_demes)
        #while np.sum(proba)>1:
            #print(np.sum(proba))
        #proba = proba/np.sum(proba)
        for l in range(len(proba)):
                if proba[l]>1e-10:
                    proba[l]-= 1e-15
                    break
        samples = np.zeros(2*N_demes)
        samples = np.random.multinomial(Nj, proba)
        samples_reshape = samples.reshape(-1, 2)
        new_numbers[j,0]= np.sum(samples_reshape[:,0])
        new_numbers[j,1]= np.sum(samples_reshape[:,1])
        #Updating the numbers in j with the migrants that arrived from i
    return new_numbers
#Checking if the mutant is completely extinct
@njit
def extinct_mutant(numbers):
    N_demes = numbers.shape[0]
    for i in range(N_demes):
        if numbers[i,0]>0:
            return False
    return True

#Checking if the wild-type is completely extinct
@njit
def extinct_wild(numbers):
    N_demes, N_types= numbers.shape
    for i in range(N_demes):
        for j in range(N_types -1):
            if numbers[i,j+1]>0:
                return False
    return True


#We do a number of cycles starting from inital numbers in_numbers ; until the mutant is fixed or dies out
#nb_cycles : a priori number of how many cycles we do before observing fixation or extinction
#dilution_number = bottleneck
#growth_factor = t the growth parameter
@njit
def cycle(in_numbers, fitness_parameters, nb_cycles, growth_factor, dilution_number, determinism, start_follow_numbers, size_follow_numbers, start_cycle, print_frequency, save_dynamics=False, simpson=False):
    N_demes, N_types= in_numbers.shape
    if start_follow_numbers is None:
        follow_numbers=np.zeros((size_follow_numbers,N_demes,N_types), dtype=np.int64)
    else:
        follow_numbers=start_follow_numbers.copy()
    fixation=True
    numbers=in_numbers.copy()
    w,b,c,K=fitness_parameters
    #Booleans that check if mutants or wild-types are extinct
    keep_going=True
    for i in range(nb_cycles):
        end_cycle = nb_cycles
        if determinism == True: #Deterministic growth
            numbers=growth_event(numbers,w,b,c,growth_factor,K)
        else: #Stochastic growth
            mutants,cheaters,temps = growth_stochastique(numbers, b, c, w, K, growth_factor)
            temps=temps[1:]
            numbers=np.zeros((N_demes,N_types),dtype=np.float64)
            numbers[:,0]=mutants[:,-1]
            numbers[:,1]=cheaters[:,-1]
            #print(numbers)
        migration_matrix=np.ones((N_demes,N_demes),dtype=np.float64)
        for j in range(N_demes):
            migration_matrix[j,:]= np.ones(N_demes)*(numbers[j,0]+numbers[j,1])/np.sum(numbers[:,0]+numbers[:,1])
        #print(migration_matrix)
        numbers=dilution_migration_event(numbers,migration_matrix,dilution_number,simpson=False)
        if (start_cycle+i)%print_frequency==0 and ((i+start_cycle)/print_frequency)<size_follow_numbers:
            follow_numbers[int(i+start_cycle), :, :]=numbers
        if extinct_mutant(numbers):
            keep_going=False
            fixation=False
            end_cycle=i+start_cycle
            follow_numbers[int((i+start_cycle))+1:]=numbers
            break
        if extinct_wild(numbers):
            keep_going=False
            fixation=True
            #print("Fixation occurred at cycle ", i+start_cycle)
            end_cycle=i+start_cycle
            follow_numbers[int((i+start_cycle))+1:]=numbers
            break
       

    if keep_going:
        follow_numbers, end_cycle, fixation= cycle(numbers, fitness_parameters, nb_cycles, growth_factor, dilution_number, determinism, follow_numbers, size_follow_numbers, start_cycle+end_cycle, print_frequency, save_dynamics, simpson=False)
    return follow_numbers, end_cycle, fixation

#We compute the fixation probability starting from in_numbers, on nb_sim simulations.
#For each simulation we use function cycle, and see if the mutant is fixed or not
@njit
def fixation_probability(in_numbers, fitness_parameters, nb_sim, growth_factor, dilution_number, determinism, simpson,nb_cycles=1000,equal_contribution=False,size_follow_numbers=10000, print_frequency=1, save_dynamics=False,folder=None):
    fix_count=0
    average_fix_cycle=np.zeros(nb_sim)
    average_ex_cycle=np.zeros(nb_sim)
    for i in range(nb_sim):
        #if i % 10000 == 0:
        #print("Simulation", i)
        start_cycle=0
        start_follow_numbers=None
        follow_numbers, end_cycle, fixation = cycle(in_numbers, fitness_parameters, nb_cycles, growth_factor, dilution_number, determinism, start_follow_numbers, size_follow_numbers, start_cycle, print_frequency, save_dynamics=False, simpson=False)
        if fixation :
            fix_count+=1
            average_fix_cycle[i] = end_cycle
            #if save_dynamics:
                #if fix_count<10000:
                    #np.savez(folder+"/fix_"+str(fix_count),follow_numbers)
        else:
            average_ex_cycle[i] = end_cycle
            #if save_dynamics:
                #if i-fix_count<10000:
                    #np.savez(folder+"/ex_"+str(i-fix_count),follow_numbers)
    ex_count=nb_sim-fix_count
    if fix_count>0:
        fixation_cycle = np.sum(average_fix_cycle)/fix_count
    else :
        fixation_cycle = 0
    if ex_count>0:
        extinction_cycle = np.sum(average_ex_cycle)/ex_count
    else:
        extinction_cycle=0
    proba = fix_count/nb_sim
    return extinction_cycle,fixation_cycle, proba


#Computing the fixation probability
def calculate_fixation_probability(bvalue):
    w = 0.001
    c = 1
    nb_sim = 1000
    t = 3
    K=1000
    bottleneck = 10
    in_num_graph = np.array([[1.,9.],[0.,10.],[0.,10.],[0.,10.],[0.,10.]]).astype(np.int64) #Clique-structured population of D demes, bottleneck size B=10, starting with one mutant 
    #in_num_mixed=np.array([[1.,49.]]).astype(np.int64) #Well-mixed population of size D*B, starting with one mutant
    parameters = w, bvalue, c,K
    _, _, fixation_value = fixation_probability(in_num_graph, parameters, nb_sim, t, bottleneck,False,False,1000,False,10000, 1,False,None) #Deterministic or stochastic growth according to a boolean variable
    return fixation_value

def execution():
    c = 1
    bvalues = np.linspace(c,50,20) #Testing different values of benefit-to-cost ratio b/c
    bvalues=np.sort(bvalues)  
    fixation_values_list=np.zeros((20))
    for i in range(20):
        fixation_values_list[i]=calculate_fixation_probability(bvalues[i])
    return fixation_values_list

fixation_values = execution() #Fixation probabilities

print(fixation_values)