coevolving_gene_expression_dynamics.py

"""
Evolutionary simulation of gene regulatory networks.

Author: Dr. Tuan Pham

This code simulates the evolution of a set of gene regulatory networks over T generations.
The fitness of each network is determined by the average expression level of the genes. The networks are 
randomly mutated and the fitness is recalculated after each mutation. The results are plotted and saved in a 
directory called 'tutorial1'.

Usage:
    Run the 'simulate()' function.

Requirements:
    numpy, matplotlib, tqdm

"""
import itertools
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
from tqdm.auto import tqdm


def integrate_GRN(N, beta, sigma, J, dt, tmax):
    """
    Function to integrate the random trajectories of gene expression level.
    """
    x = np.random.rand(N)
    t = 0
    while t < tmax:
        noise = sigma * np.random.randn(N)
        x = x + (np.tanh(beta * J @ x) - x + noise) * dt
        t += dt
    return x


def evolve(M, Ns, L, N, c, beta, sigma, T, dt, tmax):
    """
    Evolve a set of M networks with N nodes over T generations.
    Parameters
    ----------
    M : Total number of networks.
    Ns : Number of selected networks.
    L : Average fitness repetition count.
    N : Number of nodes per network.
    c : Connection probability.
    beta : `integrate_GRN` function parameter.
    sigma : `integrate_GRN` function parameter.
    T : Number of generations.
    dt : Time step for `integrate_GRN` function.
    tmax : Maximum time for `integrate_GRN` function.
    Returns
    -------
    average_fitness : Array containing the average fitness of each network at each generation.
    """

    Js = np.random.normal(0, 1 / np.sqrt(N), (M, N, N))
    Js[:, np.arange(N), np.arange(N)] = 0
    Js[np.random.uniform(size=(M, N, N)) > c] = 0

    m = int(M / Ns)
    f = np.empty((T, M))
    Vg  = np.zeros(T)
    Vip  = np.zeros(T)
    for t in tqdm(range(T), leave=False):

        # Compute the average fitness of the networks after L repetitions.
        av_fit_net = np.empty(M)
        Vip_of_each_net = np.zeros(M)
        for k in range(M):
            fitness = np.zeros(L)
            for l in range(L):
                genes = integrate_GRN(N, beta, sigma, Js[k], dt, tmax)
                fitness[l] = np.sum(genes) / N
            av_fit_net[k] = np.sum(fitness) / L
            Vip_of_each_net[k] = np.var(fitness)

        # Select Ns networks with the best averages.
        best_fit = np.argsort(av_fit_net)[-Ns:]

        # From the best selected networks, generate M/Ns mutated networks.
        Jm = np.empty((Ns, m, N, N))
        for i, k in enumerate(best_fit):
            J_parent = Js[k]
            chosen_pairs = np.random.choice(N, size=(m, 2))
            new_pairs = np.random.choice(N, size=(m, 2))
            for j, (idel, inew) in enumerate(zip(chosen_pairs, new_pairs)):
                new_J = J_parent.copy()
                new_J[idel] = 0
                new_J[inew] = np.random.randn()
                Jm[i, j] = new_J

        # Flatten the generated networks to become the N parents at the next iter
        Js = Jm.reshape(-1, N, N)

        f[t] = av_fit_net
        Vg[t]  = np.var(av_fit_net)
        Vip[t] = np.sum(Vip_of_each_net)/M

    return f,  Vg, Vip


def simulate():
    # Define the parameters to run the simulation with
    N_list = [100]
    beta_list = [1]
    sigma_list = [1]

    # Create an ouput directory to store the figures
    out_dir = Path("tutorial1")
    out_dir.mkdir(exist_ok=True, parents=True)

    # Run the simulation for all the combinations of parameters
    combinations = list(itertools.product(N_list, beta_list, sigma_list))
    for N, beta, sigma in tqdm(combinations):
        fitness, Vg, Vip = evolve(M=20, Ns=2, L=10, N=N, c=0.1, beta=beta, sigma=sigma, T=50, dt=0.1, tmax=30)

        # Plot the resulting fitness and save it in the output directory.
        plt.figure(dpi=150, figsize=(7,4))
        title_str = rf"$\beta={beta:.0e}$  $\sigma={sigma:.0e}$"
        plt.suptitle(title_str)
        
        plt.subplot(1,2,1)
        plt.xlabel("Generations")
        plt.ylabel("Observables")
        plt.plot(Vg, label="mutational_variance")
        plt.plot(Vip, label="isogenic_variance")
        plt.legend()
        
        plt.subplot(1,2,2)
        plt.xlabel("Generations")
        plt.ylabel("fitness")
        plt.plot(fitness)
        
        plt.tight_layout()
        plt.savefig(out_dir / f"gene_dynamics_beta{beta:.0e}_sigma{sigma:.0e}.png", dpi=300)
        plt.close()


if __name__ == "__main__":
    simulate()