Basic example

In [10]:

import diffrax as dfx
import jax
import jax.numpy as jnp
import jax.random as jrand
import matplotlib as mpl
import matplotlib.pyplot as plt

from felice.neuron_models import WereRabbit

jax.config.update("jax_enable_x64", True)

import diffrax as dfx import jax import jax.numpy as jnp import jax.random as jrand import matplotlib as mpl import matplotlib.pyplot as plt from felice.neuron_models import WereRabbit jax.config.update("jax_enable_x64", True)

In [11]:

key = jrand.key(0)
max_time = 40

model = WereRabbit(dtype=jnp.float64)


def state_at_t(comp_times):
    sol = dfx.diffeqsolve(
        terms=dfx.ODETerm(model.dynamics),
        solver=dfx.Tsit5(),
        t0=0.0,
        t1=max_time,
        dt0=1e-3,
        y0=model.init_state(1)
        + jrand.uniform(key, shape=(1, 2), minval=0.1, maxval=0.5),
        saveat=dfx.SaveAt(ts=comp_times),
        max_steps=100000,
    )

    return sol.ts, sol.ys

key = jrand.key(0) max_time = 40 model = WereRabbit(dtype=jnp.float64) def state_at_t(comp_times): sol = dfx.diffeqsolve( terms=dfx.ODETerm(model.dynamics), solver=dfx.Tsit5(), t0=0.0, t1=max_time, dt0=1e-3, y0=model.init_state(1) + jrand.uniform(key, shape=(1, 2), minval=0.1, maxval=0.5), saveat=dfx.SaveAt(ts=comp_times), max_steps=100000, ) return sol.ts, sol.ys

In [12]:

comp_times = jnp.linspace(0.0, max_time, 2000)
_, state = state_at_t(comp_times)

comp_times = jnp.linspace(0.0, max_time, 2000) _, state = state_at_t(comp_times)

In [13]:

def compute_nullclines(snn, u_range, v_range, resolution=200):
    """
    Compute nullclines
    du/dt = 0 (u-nullcline)
    dv/dt = 0 (v-nullcline)
    """
    u_vals = jnp.linspace(u_range[0], u_range[1], resolution)
    v_vals = jnp.linspace(v_range[0], v_range[1], resolution)
    U, V = jnp.meshgrid(u_vals, v_vals)

    UV = jnp.stack(
        [U.reshape(-1), V.reshape(-1), jnp.ones((resolution * resolution,))], axis=1
    )
    dS = snn.vector_field(UV)
    dU = dS[:, 0].reshape(U.shape)
    dV = dS[:, 1].reshape(V.shape)

    return U, V, dU, dV


def plot_vf(ax, snn, u_range, v_range):
    import numpy as np

    u_sparse = jnp.linspace(u_range[0], u_range[1], 20)
    v_sparse = jnp.linspace(v_range[0], v_range[1], 20)

    Us, Vs = jnp.meshgrid(u_sparse, v_sparse)

    U, V, dU, dV = compute_nullclines(snn, u_range, v_range, 200)

    UVs = jnp.stack([Us.reshape(-1), Vs.reshape(-1), jnp.ones((20 * 20,))], axis=1)
    dS = snn.vector_field(UVs)
    dUs = dS[:, 0].reshape(Us.shape)
    dVs = dS[:, 1].reshape(Vs.shape)

    # Normalize for visualization
    magnitude = np.sqrt(dUs**2 + dVs**2)
    magnitude[magnitude == 0] = 1
    dUs_norm = dUs / magnitude
    dVs_norm = dVs / magnitude

    # Nullclines
    ax.contour(U, V, dU, levels=[0], colors="blue", linewidths=1, linestyles="-")
    ax.contour(U, V, dV, levels=[0], colors="red", linewidths=1, linestyles="-")

    # Vector field
    ax.quiver(Us, Vs, dUs_norm, dVs_norm, magnitude, cmap="viridis", alpha=0.6)

    ax.set_xlabel("u (Prey)")
    ax.set_ylabel("v (Predator)")
    ax.set_title("Wererabbit: Phase Portrait")
    ax.legend(["u-nullcline (du/dt=0)", "v-nullcline (dv/dt=0)"], loc="upper right")
    ax.set_xlim(u_range)
    ax.set_ylim(v_range)
    ax.axhline(y=0, color="gray", linestyle="--", alpha=0.3)
    ax.axvline(x=0, color="gray", linestyle="--", alpha=0.3)

def compute_nullclines(snn, u_range, v_range, resolution=200): """ Compute nullclines du/dt = 0 (u-nullcline) dv/dt = 0 (v-nullcline) """ u_vals = jnp.linspace(u_range[0], u_range[1], resolution) v_vals = jnp.linspace(v_range[0], v_range[1], resolution) U, V = jnp.meshgrid(u_vals, v_vals) UV = jnp.stack( [U.reshape(-1), V.reshape(-1), jnp.ones((resolution * resolution,))], axis=1 ) dS = snn.vector_field(UV) dU = dS[:, 0].reshape(U.shape) dV = dS[:, 1].reshape(V.shape) return U, V, dU, dV def plot_vf(ax, snn, u_range, v_range): import numpy as np u_sparse = jnp.linspace(u_range[0], u_range[1], 20) v_sparse = jnp.linspace(v_range[0], v_range[1], 20) Us, Vs = jnp.meshgrid(u_sparse, v_sparse) U, V, dU, dV = compute_nullclines(snn, u_range, v_range, 200) UVs = jnp.stack([Us.reshape(-1), Vs.reshape(-1), jnp.ones((20 * 20,))], axis=1) dS = snn.vector_field(UVs) dUs = dS[:, 0].reshape(Us.shape) dVs = dS[:, 1].reshape(Vs.shape) # Normalize for visualization magnitude = np.sqrt(dUs**2 + dVs**2) magnitude[magnitude == 0] = 1 dUs_norm = dUs / magnitude dVs_norm = dVs / magnitude # Nullclines ax.contour(U, V, dU, levels=[0], colors="blue", linewidths=1, linestyles="-") ax.contour(U, V, dV, levels=[0], colors="red", linewidths=1, linestyles="-") # Vector field ax.quiver(Us, Vs, dUs_norm, dVs_norm, magnitude, cmap="viridis", alpha=0.6) ax.set_xlabel("u (Prey)") ax.set_ylabel("v (Predator)") ax.set_title("Wererabbit: Phase Portrait") ax.legend(["u-nullcline (du/dt=0)", "v-nullcline (dv/dt=0)"], loc="upper right") ax.set_xlim(u_range) ax.set_ylim(v_range) ax.axhline(y=0, color="gray", linestyle="--", alpha=0.3) ax.axvline(x=0, color="gray", linestyle="--", alpha=0.3)

In [14]:

with mpl.style.context("boilerplot.ieeetran"):
    fig, ax = plt.subplots(1, 2, figsize=(6.9, 2.6), dpi=200)
    ax[0].plot(comp_times, state[:, 0, 0], label="x1")
    ax[0].plot(comp_times, state[:, 0, 1], label="x2")
    ax[0].legend(["x1", "x2"])

    plot_vf(ax[1], model, [-0.2, 0.5], [-0.2, 0.5])
    ax[1].plot(state[:, 0, 0], state[:, 0, 1])
    ax[1].plot(state[0, 0, 0], state[0, 0, 1], ".", label="start")
    ax[1].plot(state[-1, 0, 0], state[-1, 0, 1], ".", label="end")
    ax[1].legend()
    plt.show()

with mpl.style.context("boilerplot.ieeetran"): fig, ax = plt.subplots(1, 2, figsize=(6.9, 2.6), dpi=200) ax[0].plot(comp_times, state[:, 0, 0], label="x1") ax[0].plot(comp_times, state[:, 0, 1], label="x2") ax[0].legend(["x1", "x2"]) plot_vf(ax[1], model, [-0.2, 0.5], [-0.2, 0.5]) ax[1].plot(state[:, 0, 0], state[:, 0, 1]) ax[1].plot(state[0, 0, 0], state[0, 0, 1], ".", label="start") ax[1].plot(state[-1, 0, 0], state[-1, 0, 1], ".", label="end") ax[1].legend() plt.show()

In [ ]: