prax

JAX implementation of phase response curve

Installation

prax is created based on jax, and please install jax at first. See JAX page for installation.

After the installation of jax, prax can be installed with pip directly from GitHub, with the following command:

pip install git+https://github.com/yonesuke/prax.git

Quickstart

We give an example on how to use this package with Van der Pol oscillator.

First, import packages:

import jax.numpy as jnp
from prax import Oscillator
from jax.config import config; config.update("jax_enable_x64", True)

import matplotlib.pyplot as plt

Create an oscillator class by inheriting Oscillator class:

class VanderPol(Oscillator):
    def __init__(self, mu, dt=0.01, eps=10**-5):
        super().__init__(n_dim=2, dt=dt, eps=eps)
        self.mu = mu

    def forward(self, state):
        x, y = state
        vx = y
        vy = self.mu * (1.0 - x*x) * y - x
        return jnp.array([vx, vy])

model = VanderPol(mu=0.2)

Find periodic orbit (choose init_val nicely so that it goes to periodic orbit):

init_val = jnp.array([0.1, 0.2])
model.find_periodic_orbit(init_val)
print(model.period) # 6.3088767
plt.plot(model.ts, model.periodic_orbit)

Calculate phase response curve:

model.calc_phase_response()
plt.plot(model.ts, model.phase_response_curve)

Galleries

See examples directory!!

Van der Pol equation [code]

class VanderPol(Oscillator):
    def __init__(self, mu, dt=0.01, eps=10**-5):
        super().__init__(n_dim=2, dt=dt, eps=eps)
        self.mu = mu

    def forward(self, state):
        x, y = state
        vx = y
        vy = self.mu * (1.0 - x*x) * y - x
        return jnp.array([vx, vy])

model = VanderPol(mu=0.2)

Stuart Landau equation [code]

class StuartLandau(Oscillator):
    def __init__(self, dt=0.01, eps=10**-5):
        super().__init__(n_dim=2, dt=dt, eps=eps)

    def forward(self, state):
        x, y = state
        vx = x - y - x * (x * x + y * y)
        vy = x + y - y * (x * x + y * y)
        return jnp.array([vx, vy])

model = StuartLandau()

FitzHugh-Nagumo equation [code]

class FitzHughNagumo(Oscillator):
    def __init__(self, params, dt=0.01, eps=10**-5):
        super().__init__(n_dim=2, dt=dt, eps=eps)
        self.a, self.b, self.c = params

    def forward(self, state):
        x, y = state
        vx = self.c * (x - x ** 3 - y)
        vy = x - self.b * y + self.a
        return jnp.array([vx, vy])

model = FitzHughNagumo(params=[0.2, 0.5, 10.0])

Brusselator equation [code]

class Brusselator(Oscillator):
    def __init__(self, params, dt=0.01, eps=10**-5):
        super().__init__(n_dim=2, dt=dt, eps=eps)
        self.a, self.b = params

    def forward(self, state):
        x, y = state
        vx = self.a - (self.b + 1.0) * x + x * x * y
        vy = self.b * x - x * x * y
        return jnp.array([vx, vy])

model = Brusselator(params=[1.0, 3.0])

Hodgkin Huxley equation [code]

class HodgkinHuxley(Oscillator):
    def __init__(self, input_current, C=1.0, G_Na=120.0, G_K=36.0, G_L=0.3, E_Na=50.0, E_K=-77.0, E_L=-54.4, dt=0.01, eps=10**-5):
        super().__init__(n_dim=4, dt=dt, eps=eps)
        self.input_current = input_current
        self.C = C
        self.G_Na = G_Na
        self.G_K = G_K
        self.G_L = G_L
        self.E_Na = E_Na
        self.E_K = E_K
        self.E_L = E_L

    def alpha_m(self, V):
        return 0.1*(V+40.0)/(1.0 - jnp.exp(-(V+40.0) / 10.0))
    
    def beta_m(self, V):
        return 4.0*jnp.exp(-(V+65.0) / 18.0)
    
    def alpha_h(self, V):
        return 0.07*jnp.exp(-(V+65.0) / 20.0)
    
    def beta_h(self, V):
        return 1.0/(1.0 + jnp.exp(-(V+35.0) / 10.0))
    
    def alpha_n(self, V):
        return 0.01*(V+55.0)/(1.0 - jnp.exp(-(V+55.0) / 10.0))
    
    def beta_n(self, V):
        return 0.125*jnp.exp(-(V+65) / 80.0)

    def forward(self, state):
        V, m, h, n = state
        dVdt = self.G_Na * (m ** 3) * h * (self.E_Na - V) + self.G_K * (n ** 4) * (self.E_K - V) + self.G_L * (self.E_L - V) + self.input_current
        dVdt /= self.C
        dmdt = self.alpha_m(V) * (1.0 - m) - self.beta_m(V) * m
        dhdt = self.alpha_h(V) * (1.0 - h) - self.beta_h(V) * h
        dndt = self.alpha_n(V) * (1.0 - n) - self.beta_n(V) * n
        return jnp.array([dVdt, dmdt, dhdt, dndt])

model = HodgkinHuxley(input_current=30.0)