# Basic example

In [1]:
Copied!

    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 eventpropjax.evnn import FFEvNN

    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 eventpropjax.evnn import FFEvNN from felice.neuron_models import WereRabbit jax.config.update("jax_enable_x64", True)
In [2]:
Copied!

    key = jrand.key(0)
    max_time = 100

    init_weights = jnp.asarray([[0.0], [0.2], [-0.1]])
    snn = FFEvNN(
        layers=[1],
        in_size=2,
        neuron_model=WereRabbit,
        solver=dfx.Tsit5(),
        # stepsize_controller=PIDController(
        #     rtol=1e-6, atol=1e-3, pcoeff=0.1, icoeff=0.3, dcoeff=0.0
        # ),
        max_solver_time=max_time,
        key=key,
        max_event_steps=10000,
        solver_stepsize=0.001,
        init_weights=init_weights,
        dtype=jnp.float64,
    )

key = jrand.key(0) max_time = 100 init_weights = jnp.asarray([[0.0], [0.2], [-0.1]]) snn = FFEvNN( layers=[1], in_size=2, neuron_model=WereRabbit, solver=dfx.Tsit5(), # stepsize_controller=PIDController( # rtol=1e-6, atol=1e-3, pcoeff=0.1, icoeff=0.3, dcoeff=0.0 # ), max_solver_time=max_time, key=key, max_event_steps=10000, solver_stepsize=0.001, init_weights=init_weights, dtype=jnp.float64, )
In [3]:
Copied!

    in_spikes = jnp.asarray([[10, 30], [20, 21]])
    comp_times = jnp.linspace(0.0, max_time, 2000)
    state = snn.state_at_t(in_spikes, comp_times)
    spikes = snn.spikes_until_t(in_spikes, max_time)

in_spikes = jnp.asarray([[10, 30], [20, 21]]) comp_times = jnp.linspace(0.0, max_time, 2000) state = snn.state_at_t(in_spikes, comp_times) spikes = snn.spikes_until_t(in_spikes, max_time)
In [4]:
Copied!

    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.neuron_model.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.neuron_model.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.neuron_model.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.neuron_model.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 [5]:
Copied!

    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].axvline(s, alpha=0.2, color="g", linestyle="--") for s in jnp.unique(spikes)]
        ax[0].legend(["x1", "x2", "Spike"])

        plot_vf(ax[1], snn, [-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[0, -1, 0], state[0, -1, 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].axvline(s, alpha=0.2, color="g", linestyle="--") for s in jnp.unique(spikes)] ax[0].legend(["x1", "x2", "Spike"]) plot_vf(ax[1], snn, [-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[0, -1, 0], state[0, -1, 1], ".", label="end") ax[1].legend() plt.show()
In [6]:
Copied!

    import equinox as eqx

    def loss_fn(model, spike_times, comp_times):
        out_states = model.state_at_t(spike_times, comp_times)
        logits = out_states[:, :, 0]

        return jnp.sum(logits)

    in_spikes = jnp.asarray([[0.01], [0.157]])
    comp_times = jnp.linspace(0.0, max_time, 10)

    loss, gradients = eqx.filter_value_and_grad(loss_fn)(snn, in_spikes, comp_times)
    print("Loss ", loss)
    print("Gradients ", gradients.neuron_model.weight_u)

import equinox as eqx def loss_fn(model, spike_times, comp_times): out_states = model.state_at_t(spike_times, comp_times) logits = out_states[:, :, 0] return jnp.sum(logits) in_spikes = jnp.asarray([[0.01], [0.157]]) comp_times = jnp.linspace(0.0, max_time, 10) loss, gradients = eqx.filter_value_and_grad(loss_fn)(snn, in_spikes, comp_times) print("Loss ", loss) print("Gradients ", gradients.neuron_model.weight_u)

    Loss  2.78585239490824
    Gradients  [[ 0.        ]
     [-1.81404788]
     [-1.42144198]]

In [ ]:
Copied!