Fitzhugh-Nagumo dynamical system. With animation code.

fitzhugh-nagumo_model_ramp_input.py

"""Fitzhugh-Nagumo dynamical system. With animation code.

Author: Daniel Müller-Komorowska @scidanm
"""

import numpy as np
from scipy.integrate import solve_ivp
from matplotlib.animation import FuncAnimation
import matplotlib.pyplot as plt
import matplotlib as mpl

np.random.seed(1000) # Set random seed (for reproducibility)

"""Hyperparameters of the simulation"""
tmin = 0.0
tmax = 500
dt=0.1
t = np.arange(tmin,tmax,dt)
animate = True

"""Parameters of the FitzHugh-Nagumo Model"""
a = 0.7
b = 0.8
tau = 12.5


def dw(Vm, w):
    return (Vm + a - b * w) / tau


def dV(Vm, w, I):
    return Vm - ((Vm ** 3) / 3) - w + I

def It(t, slope, intercept):
    return t * slope + intercept

I = It(t, 0.001, 0.3)
I = I[np.newaxis,np.newaxis,:]

x = np.arange(-0.5,1.7, 0.1)
y = np.linspace(-2, 2, int(x.shape[0]/2))

x_mesh, y_mesh = np.meshgrid(x, y)
x_mesh_batch = x_mesh[:,:,np.newaxis]
y_mesh_batch = y_mesh[:,:,np.newaxis]

dw_result = dw(y_mesh_batch, x_mesh_batch)
dV_result = dV(y_mesh_batch, x_mesh_batch, I)

def compute_derivatives(t, y):

    dy = np.zeros((2,))
    
    Vm = y[0]
    w = y[1]
    
    # dVm/dt    
    dy[0] = dV(Vm, w, It(t, 0.001, 0.2))

    
    # dn/dt
    dy[1] = dw(Vm, w)
    
    return dy

# Solve ODE system
Y = np.array([-1.07, -0.462])
result = solve_ivp(compute_derivatives, (tmin,tmax), Y, t_eval=t)

"""Animation"""
if animate:
    mpl.rcParams['grid.linewidth'] = 3
    mpl.rcParams['font.size'] = 12
    fig, ax = plt.subplots(2, gridspec_kw={'height_ratios': [2,1]})
    fig.set_size_inches(16,9)
    results_v = result.y[0]
    results_u = result.y[1]

    def update_plot(i):
        ax[0].clear()
        ax[1].clear()
        ax[0].quiver(x, y, dw_result[:,:,0]/dw_result[:,:,0].max(),
                   dV_result[:,:,i]/dV_result[:,:,i].max(), pivot='mid',
                   headwidth=2.5, headlength=3, width=0.002)
        ax[0].set_xlabel("potassium activation n")
        ax[0].set_ylabel("voltage (mV)")
        ax[0].scatter(x_mesh, y_mesh, color='r', s=2)
        if i<100:
            ax[0].plot(results_u[0:i], results_v[0:i], color='k', alpha=0.9)
        else:
            ax[0].plot(results_u[0:i-100], results_v[0:i-100], color='r', alpha=0.5)
            ax[0].plot(results_u[0+i-100:i], results_v[0+i-100:i], color='k', alpha=0.9)
        ax[0].set_xlabel("w")
        ax[0].set_ylabel("V")
        ax[1].plot(t[0:i], results_v[0:i], color='k')
        ax[1].set_xlim((t[0], t[-1]))
        ax[1].set_ylim((results_v.min()-0.1, results_v.max()+0.1))
        ax[1].set_ylabel("V")
        ax[1].set_xlabel("time (ms)")
        txt = ax[1].text(5, 1.5, "I=" + "{0:.5f}".format(I[0,0,i]), size=12)
        
    anim = FuncAnimation(fig, update_plot, frames=np.arange(0, I.shape[2],12), interval=1)
    anim.save('output_animation.mp4', dpi=150, fps = 30, writer='ffmpeg')