bics/felice-models
5
0 Fork 0
mirror of https://github.com/bics-rug/felice-models.git synced 2026-09-30 21:03:06 +02:00
Code Issues Packages Projects Releases Wiki Activity

Deployed 9fabbde with MkDocs version: 1.6.1

Browse Source
This commit is contained in:
F.M. Quintana Velazquez committed 2026-02-27 17:48:38 +01:00
1 parent 76339b2417
commit 55f1381105
16 files changed
+1529 -491
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

felice_models.pdf
BIN
View File
Binary file not shown.
img/lif_plot.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 98 KiB

index.html
+1 -54
View File
@@ -321,17 +321,6 @@
</ul>
</nav>
</li>
<li class="md-nav__item">
<a href="#quick-start" class="md-nav__link">
<span class="md-ellipsis">
Quick Start
</span>
</a>
</li>
</ul>
@@ -921,17 +910,6 @@
</ul>
</nav>
</li>
<li class="md-nav__item">
<a href="#quick-start" class="md-nav__link">
<span class="md-ellipsis">
Quick Start
</span>
</a>
</li>
</ul>
@@ -962,7 +940,7 @@
<h1 id="felice">Felice</h1>
<p>This project provides a <a href="https://github.com/google/jax">JAX</a> implementation of the different neuron models in felice</p>
<h2 id="overview">Overview</h2>
<p>The framework is built on top of <a href="">EventPropJax</a> and leverages JAX's automatic differentiation for efficient simulation and training of SNNs using event-based exact gradients.</p>
<p>The framework is built on top of diffrax and leverages JAX's automatic differentiation for efficient simulation and training of analogue models.</p>
<h3 id="key-features">Key Features</h3>
<ul>
<li><strong>Delay learning</strong></li>
@@ -981,37 +959,6 @@
<p>For GPU acceleration with CUDA 13:</p>
<div class="language-bash highlight"><pre><span></span><code><span id="__span-1-1"><a id="__codelineno-1-1" name="__codelineno-1-1" href="#__codelineno-1-1"></a>uv<span class="w"> </span>sync<span class="w"> </span>--extra<span class="w"> </span>cuda
</span></code></pre></div>
<h2 id="quick-start">Quick Start</h2>
<p>Here's a simple example using the WereRabbit neuron model:</p>
<div class="language-py highlight"><pre><span></span><code><span id="__span-2-1"><a id="__codelineno-2-1" name="__codelineno-2-1" href="#__codelineno-2-1"></a><span class="kn">import</span><span class="w"> </span><span class="nn">diffrax</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">dfx</span>
</span><span id="__span-2-2"><a id="__codelineno-2-2" name="__codelineno-2-2" href="#__codelineno-2-2"></a><span class="kn">import</span><span class="w"> </span><span class="nn">jax.numpy</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">jnp</span>
</span><span id="__span-2-3"><a id="__codelineno-2-3" name="__codelineno-2-3" href="#__codelineno-2-3"></a><span class="kn">import</span><span class="w"> </span><span class="nn">jax.random</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">jrand</span>
</span><span id="__span-2-4"><a id="__codelineno-2-4" name="__codelineno-2-4" href="#__codelineno-2-4"></a><span class="kn">from</span><span class="w"> </span><span class="nn">eventpropjax.evnn</span><span class="w"> </span><span class="kn">import</span> <span class="n">FFEvNN</span>
</span><span id="__span-2-5"><a id="__codelineno-2-5" name="__codelineno-2-5" href="#__codelineno-2-5"></a><span class="kn">from</span><span class="w"> </span><span class="nn">felice.neuron_models</span><span class="w"> </span><span class="kn">import</span> <span class="n">WereRabbit</span>
</span><span id="__span-2-6"><a id="__codelineno-2-6" name="__codelineno-2-6" href="#__codelineno-2-6"></a>
</span><span id="__span-2-7"><a id="__codelineno-2-7" name="__codelineno-2-7" href="#__codelineno-2-7"></a><span class="c1"># Initialize random key and parameters</span>
</span><span id="__span-2-8"><a id="__codelineno-2-8" name="__codelineno-2-8" href="#__codelineno-2-8"></a><span class="n">key</span> <span class="o">=</span> <span class="n">jrand</span><span class="o">.</span><span class="n">key</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
</span><span id="__span-2-9"><a id="__codelineno-2-9" name="__codelineno-2-9" href="#__codelineno-2-9"></a><span class="n">max_time</span> <span class="o">=</span> <span class="mf">300e-3</span>
</span><span id="__span-2-10"><a id="__codelineno-2-10" name="__codelineno-2-10" href="#__codelineno-2-10"></a>
</span><span id="__span-2-11"><a id="__codelineno-2-11" name="__codelineno-2-11" href="#__codelineno-2-11"></a><span class="c1"># Create a feedforward event-driven neural network</span>
</span><span id="__span-2-12"><a id="__codelineno-2-12" name="__codelineno-2-12" href="#__codelineno-2-12"></a><span class="n">snn</span> <span class="o">=</span> <span class="n">FFEvNN</span><span class="p">(</span>
</span><span id="__span-2-13"><a id="__codelineno-2-13" name="__codelineno-2-13" href="#__codelineno-2-13"></a> <span class="n">layers</span><span class="o">=</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
</span><span id="__span-2-14"><a id="__codelineno-2-14" name="__codelineno-2-14" href="#__codelineno-2-14"></a> <span class="n">in_size</span><span class="o">=</span><span class="mi">2</span><span class="p">,</span>
</span><span id="__span-2-15"><a id="__codelineno-2-15" name="__codelineno-2-15" href="#__codelineno-2-15"></a> <span class="n">neuron_model</span><span class="o">=</span><span class="n">WereRabbit</span><span class="p">,</span>
</span><span id="__span-2-16"><a id="__codelineno-2-16" name="__codelineno-2-16" href="#__codelineno-2-16"></a> <span class="n">solver</span><span class="o">=</span><span class="n">dfx</span><span class="o">.</span><span class="n">Tsit5</span><span class="p">(),</span>
</span><span id="__span-2-17"><a id="__codelineno-2-17" name="__codelineno-2-17" href="#__codelineno-2-17"></a> <span class="n">max_solver_time</span><span class="o">=</span><span class="n">max_time</span><span class="p">,</span>
</span><span id="__span-2-18"><a id="__codelineno-2-18" name="__codelineno-2-18" href="#__codelineno-2-18"></a> <span class="n">key</span><span class="o">=</span><span class="n">key</span><span class="p">,</span>
</span><span id="__span-2-19"><a id="__codelineno-2-19" name="__codelineno-2-19" href="#__codelineno-2-19"></a> <span class="n">max_event_steps</span><span class="o">=</span><span class="mi">1000000</span><span class="p">,</span>
</span><span id="__span-2-20"><a id="__codelineno-2-20" name="__codelineno-2-20" href="#__codelineno-2-20"></a> <span class="n">solver_stepsize</span><span class="o">=</span><span class="mf">1e-6</span><span class="p">,</span>
</span><span id="__span-2-21"><a id="__codelineno-2-21" name="__codelineno-2-21" href="#__codelineno-2-21"></a> <span class="n">rtol</span><span class="o">=</span><span class="mf">10.0</span><span class="p">,</span>
</span><span id="__span-2-22"><a id="__codelineno-2-22" name="__codelineno-2-22" href="#__codelineno-2-22"></a> <span class="n">atol</span><span class="o">=</span><span class="mf">0.0</span><span class="p">,</span>
</span><span id="__span-2-23"><a id="__codelineno-2-23" name="__codelineno-2-23" href="#__codelineno-2-23"></a> <span class="n">Ibias</span><span class="o">=</span><span class="mf">30e-12</span><span class="p">,</span>
</span><span id="__span-2-24"><a id="__codelineno-2-24" name="__codelineno-2-24" href="#__codelineno-2-24"></a><span class="p">)</span>
</span><span id="__span-2-25"><a id="__codelineno-2-25" name="__codelineno-2-25" href="#__codelineno-2-25"></a>
</span><span id="__span-2-26"><a id="__codelineno-2-26" name="__codelineno-2-26" href="#__codelineno-2-26"></a><span class="c1"># Simulate with input spikes</span>
</span><span id="__span-2-27"><a id="__codelineno-2-27" name="__codelineno-2-27" href="#__codelineno-2-27"></a><span class="n">in_spikes</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">asarray</span><span class="p">([[</span><span class="mf">0.0</span><span class="p">],</span> <span class="p">[</span><span class="mf">0.157</span><span class="p">]])</span>
</span><span id="__span-2-28"><a id="__codelineno-2-28" name="__codelineno-2-28" href="#__codelineno-2-28"></a><span class="n">spikes</span> <span class="o">=</span> <span class="n">snn</span><span class="o">.</span><span class="n">spikes_until_t</span><span class="p">(</span><span class="n">in_spikes</span><span class="p">,</span> <span class="n">max_time</span><span class="p">)</span>
</span></code></pre></div>
<p>See the <a href="scripts/examples/neuron_models/">examples</a> directory for more detailed usage examples.</p>
neuron_models/fhn/IMG_20260227_161313_688.jpg
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.2 MiB

neuron_models/fhn/IMG_20260227_161324_112.jpg
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.2 MiB

neuron_models/fhn/IMG_20260227_161333_222.jpg
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.3 MiB

neuron_models/fhn/IMG_20260227_161350_882.jpg
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.3 MiB

neuron_models/fhn/IMG_20260227_161401_705.jpg
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.2 MiB

neuron_models/fhn/IMG_20260227_161410_686.jpg
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.2 MiB

neuron_models/fhn/IMG_20260227_161423_789.jpg
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.4 MiB

neuron_models/fhn/fhn/index.html
+251 -61
View File
@@ -1564,7 +1564,7 @@ span.linenos.special { color: #000000; background-color: #ffffc0; padding-left:
<div class="jp-Collapser jp-InputCollapser jp-Cell-inputCollapser">
</div>
<div class="jp-InputArea jp-Cell-inputArea">
<div class="jp-InputPrompt jp-InputArea-prompt">In [ ]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="jp-InputPrompt jp-InputArea-prompt">In [2]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="CodeMirror cm-s-jupyter">
<div class="zeroclipboard-container">
<clipboard-copy ,="" aria-label="Copy to Clipboard" for="cell-1">
@@ -1583,9 +1583,8 @@ span.linenos.special { color: #000000; background-color: #ffffc0; padding-left:
<span class="kn">import</span><span class="w"> </span><span class="nn">jax.random</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">jrand</span>
<span class="kn">import</span><span class="w"> </span><span class="nn">matplotlib</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">mpl</span>
<span class="kn">import</span><span class="w"> </span><span class="nn">matplotlib.pyplot</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">plt</span>
<span class="kn">from</span><span class="w"> </span><span class="nn">eventpropjax.evnn</span><span class="w"> </span><span class="kn">import</span> <span class="n">FFEvNN</span>
<span class="kn">from</span><span class="w"> </span><span class="nn">felice.neuron_models</span><span class="w"> </span><span class="kn">import</span> <span class="n">FHN</span>
<span class="kn">from</span><span class="w"> </span><span class="nn">felice.neuron_models</span><span class="w"> </span><span class="kn">import</span> <span class="n">FHNRS</span>
</pre></div>
<div class="clipboard-copy-txt" id="cell-1">import diffrax as dfx
import jax
@@ -1593,9 +1592,8 @@ 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 FHN</div>
from felice.neuron_models import FHNRS</div>
</div>
</div>
</div>
@@ -1607,7 +1605,7 @@ from felice.neuron_models import FHN</div>
<div class="jp-Collapser jp-InputCollapser jp-Cell-inputCollapser">
</div>
<div class="jp-InputArea jp-Cell-inputArea">
<div class="jp-InputPrompt jp-InputArea-prompt">In [2]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="jp-InputPrompt jp-InputArea-prompt">In [23]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="CodeMirror cm-s-jupyter">
<div class="zeroclipboard-container">
<clipboard-copy ,="" aria-label="Copy to Clipboard" for="cell-2">
@@ -1623,32 +1621,64 @@ from felice.neuron_models import FHN</div>
<div class="highlight-ipynb hl-python"><pre><span></span><span class="n">key</span> <span class="o">=</span> <span class="n">jrand</span><span class="o">.</span><span class="n">key</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
<span class="n">max_time</span> <span class="o">=</span> <span class="mi">200</span>
<span class="n">snn</span> <span class="o">=</span> <span class="n">FFEvNN</span><span class="p">(</span>
<span class="n">layers</span><span class="o">=</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
<span class="n">in_size</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
<span class="n">neuron_model</span><span class="o">=</span><span class="n">FHN</span><span class="p">,</span>
<span class="n">solver</span><span class="o">=</span><span class="n">dfx</span><span class="o">.</span><span class="n">Dopri5</span><span class="p">(),</span>
<span class="n">max_solver_time</span><span class="o">=</span><span class="n">max_time</span><span class="p">,</span>
<span class="n">key</span><span class="o">=</span><span class="n">key</span><span class="p">,</span>
<span class="n">max_event_steps</span><span class="o">=</span><span class="mi">1000000</span><span class="p">,</span>
<span class="n">solver_stepsize</span><span class="o">=</span><span class="mf">0.1</span><span class="p">,</span>
<span class="n">init_weights</span><span class="o">=</span><span class="mf">2.0</span><span class="p">,</span>
<span class="n">neuron_model</span> <span class="o">=</span> <span class="n">FHNRS</span><span class="p">(</span>
<span class="n">gmax_pasive</span><span class="o">=</span><span class="mf">2.0</span><span class="p">,</span>
<span class="n">Erev_pasive</span><span class="o">=</span><span class="mf">0.0</span><span class="p">,</span>
<span class="n">a_fast</span><span class="o">=-</span><span class="mf">2.0</span><span class="p">,</span>
<span class="n">voff_fast</span><span class="o">=</span><span class="mf">0.0</span><span class="p">,</span>
<span class="n">tau_fast</span><span class="o">=</span><span class="mf">0.0</span><span class="p">,</span>
<span class="n">a_slow</span><span class="o">=</span><span class="mf">0.5</span><span class="p">,</span>
<span class="n">voff_slow</span><span class="o">=</span><span class="mf">1.0</span><span class="p">,</span>
<span class="n">tau_slow</span><span class="o">=</span><span class="mf">50.0</span><span class="p">,</span>
<span class="n">vthr</span><span class="o">=</span><span class="n">jnp</span><span class="o">.</span><span class="n">inf</span><span class="p">,</span>
<span class="p">)</span>
<span class="k">def</span><span class="w"> </span><span class="nf">state_at_t</span><span class="p">(</span><span class="n">comp_times</span><span class="p">):</span>
<span class="n">sol</span> <span class="o">=</span> <span class="n">dfx</span><span class="o">.</span><span class="n">diffeqsolve</span><span class="p">(</span>
<span class="n">terms</span><span class="o">=</span><span class="n">dfx</span><span class="o">.</span><span class="n">ODETerm</span><span class="p">(</span><span class="n">neuron_model</span><span class="o">.</span><span class="n">dynamics</span><span class="p">),</span>
<span class="n">solver</span><span class="o">=</span><span class="n">dfx</span><span class="o">.</span><span class="n">Tsit5</span><span class="p">(),</span>
<span class="n">t0</span><span class="o">=</span><span class="mf">0.0</span><span class="p">,</span>
<span class="n">t1</span><span class="o">=</span><span class="n">max_time</span><span class="p">,</span>
<span class="n">dt0</span><span class="o">=</span><span class="mf">1e-3</span><span class="p">,</span>
<span class="n">y0</span><span class="o">=</span><span class="n">neuron_model</span><span class="o">.</span><span class="n">init_state</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>
<span class="o">+</span> <span class="n">jrand</span><span class="o">.</span><span class="n">uniform</span><span class="p">(</span><span class="n">key</span><span class="p">,</span> <span class="n">shape</span><span class="o">=</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">3</span><span class="p">),</span> <span class="n">minval</span><span class="o">=</span><span class="mf">0.1</span><span class="p">,</span> <span class="n">maxval</span><span class="o">=</span><span class="mf">0.5</span><span class="p">),</span>
<span class="n">saveat</span><span class="o">=</span><span class="n">dfx</span><span class="o">.</span><span class="n">SaveAt</span><span class="p">(</span><span class="n">ts</span><span class="o">=</span><span class="n">comp_times</span><span class="p">),</span>
<span class="n">max_steps</span><span class="o">=</span><span class="mi">200000</span><span class="p">,</span>
<span class="p">)</span>
<span class="k">return</span> <span class="n">sol</span><span class="o">.</span><span class="n">ts</span><span class="p">,</span> <span class="n">sol</span><span class="o">.</span><span class="n">ys</span>
</pre></div>
<div class="clipboard-copy-txt" id="cell-2">key = jrand.key(0)
max_time = 200
snn = FFEvNN(
layers=[1],
in_size=1,
neuron_model=FHN,
solver=dfx.Dopri5(),
max_solver_time=max_time,
key=key,
max_event_steps=1000000,
solver_stepsize=0.1,
init_weights=2.0,
)</div>
neuron_model = FHNRS(
gmax_pasive=2.0,
Erev_pasive=0.0,
a_fast=-2.0,
voff_fast=0.0,
tau_fast=0.0,
a_slow=0.5,
voff_slow=1.0,
tau_slow=50.0,
vthr=jnp.inf,
)
def state_at_t(comp_times):
sol = dfx.diffeqsolve(
terms=dfx.ODETerm(neuron_model.dynamics),
solver=dfx.Tsit5(),
t0=0.0,
t1=max_time,
dt0=1e-3,
y0=neuron_model.init_state(1)
+ jrand.uniform(key, shape=(1, 3), minval=0.1, maxval=0.5),
saveat=dfx.SaveAt(ts=comp_times),
max_steps=200000,
)
return sol.ts, sol.ys</div>
</div>
</div>
</div>
@@ -1660,7 +1690,7 @@ snn = FFEvNN(
<div class="jp-Collapser jp-InputCollapser jp-Cell-inputCollapser">
</div>
<div class="jp-InputArea jp-Cell-inputArea">
<div class="jp-InputPrompt jp-InputArea-prompt">In [3]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="jp-InputPrompt jp-InputArea-prompt">In [24]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="CodeMirror cm-s-jupyter">
<div class="zeroclipboard-container">
<clipboard-copy ,="" aria-label="Copy to Clipboard" for="cell-3">
@@ -1674,9 +1704,9 @@ snn = FFEvNN(
</clipboard-copy>
</div>
<div class="highlight-ipynb hl-python"><pre><span></span><span class="n">v_range</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="o">-</span><span class="mf">3.1</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mf">0.1</span><span class="p">)</span>
<span class="n">VI_inst</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">vmap</span><span class="p">(</span><span class="n">snn</span><span class="o">.</span><span class="n">neuron_model</span><span class="o">.</span><span class="n">IV_inst</span><span class="p">)(</span><span class="n">v_range</span><span class="p">)</span>
<span class="n">VI_fast</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">vmap</span><span class="p">(</span><span class="n">snn</span><span class="o">.</span><span class="n">neuron_model</span><span class="o">.</span><span class="n">IV_fast</span><span class="p">)(</span><span class="n">v_range</span><span class="p">)</span>
<span class="n">VI_slow</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">vmap</span><span class="p">(</span><span class="n">snn</span><span class="o">.</span><span class="n">neuron_model</span><span class="o">.</span><span class="n">IV_slow</span><span class="p">)(</span><span class="n">v_range</span><span class="p">)</span>
<span class="n">VI_inst</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">vmap</span><span class="p">(</span><span class="n">neuron_model</span><span class="o">.</span><span class="n">IV_inst</span><span class="p">)(</span><span class="n">v_range</span><span class="p">)</span>
<span class="n">VI_fast</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">vmap</span><span class="p">(</span><span class="n">neuron_model</span><span class="o">.</span><span class="n">IV_fast</span><span class="p">)(</span><span class="n">v_range</span><span class="p">)</span>
<span class="n">VI_slow</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">vmap</span><span class="p">(</span><span class="n">neuron_model</span><span class="o">.</span><span class="n">IV_slow</span><span class="p">)(</span><span class="n">v_range</span><span class="p">)</span>
<span class="k">with</span> <span class="n">mpl</span><span class="o">.</span><span class="n">style</span><span class="o">.</span><span class="n">context</span><span class="p">(</span><span class="s2">"boilerplot.ieeetran"</span><span class="p">):</span>
<span class="n">fig</span><span class="p">,</span> <span class="n">ax</span> <span class="o">=</span> <span class="n">plt</span><span class="o">.</span><span class="n">subplots</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="n">figsize</span><span class="o">=</span><span class="p">(</span><span class="mf">6.9</span><span class="p">,</span> <span class="mf">2.3</span><span class="p">),</span> <span class="n">dpi</span><span class="o">=</span><span class="mf">200.0</span><span class="p">,</span> <span class="n">sharey</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
@@ -1686,9 +1716,9 @@ snn = FFEvNN(
<span class="n">plt</span><span class="o">.</span><span class="n">show</span><span class="p">()</span>
</pre></div>
<div class="clipboard-copy-txt" id="cell-3">v_range = jnp.arange(-3.1, 3, 0.1)
VI_inst = jax.vmap(snn.neuron_model.IV_inst)(v_range)
VI_fast = jax.vmap(snn.neuron_model.IV_fast)(v_range)
VI_slow = jax.vmap(snn.neuron_model.IV_slow)(v_range)
VI_inst = jax.vmap(neuron_model.IV_inst)(v_range)
VI_fast = jax.vmap(neuron_model.IV_fast)(v_range)
VI_slow = jax.vmap(neuron_model.IV_slow)(v_range)
with mpl.style.context("boilerplot.ieeetran"):
fig, ax = plt.subplots(1, 3, figsize=(6.9, 2.3), dpi=200.0, sharey=True)
@@ -1707,7 +1737,7 @@ with mpl.style.context("boilerplot.ieeetran"):
<div class="jp-OutputArea-child">
<div class="jp-OutputPrompt jp-OutputArea-prompt"></div>
<div class="jp-RenderedImage jp-OutputArea-output" tabindex="0">
<img alt="No description has been provided for this image" class="" src="data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAABXsAAAHjCAYAAABy0synAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjgsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvwVt1zgAAAAlwSFlzAAAewgAAHsIBbtB1PgAAmANJREFUeJzs3QV0XMfZxvFHzJZBtszMTAInDjMzmEmOG06aNs0XaANN26RJGibLTInDzAy2JNsyM7NkGQQrhu/ckbVx2I60urt7/79zdLzvdbqajmS92mfnzgRUVVVVCQAAAAAAAADg0wLtHgAAAAAAAAAAoPYIewEAAAAAAADADxD2AgAAAAAAAIAfIOwFAAAAAAAAAD9A2AsAAAAAAAAAfoCwFwAAAAAAAAD8AGEvAAAAAAAAAPgBwl4AAAAAAAAA8AOEvQAAAAAAAADgBwh7AQAAAAAAAMAPEPYCAAAAAAAAgB8g7AUAAAAAAAAAP0DYCwAAAAAAAAB+gLAXAAAAAAAAAPwAYS8AAAAAAAAA+AHCXgAAAAAAAADwA4S9AAAAAAAAAOAHguUwlZWVysnJMY8jIyMVEBBg95AAAPhVVVVVKiwsNI/j4uIUGHhs79PS9wAATul5FvoeAMBJfe+XOC7stYLe+Ph4u4cBAMAxy8rKUrNmzY7pf0PfAwA4pedZ6HsAACf1vV/CNg4AAAAAAAAA4Acct7LX2rrhyNQ8KirKI5/H5XK5VxB78vMwBuaA7wP+LfDzwL9/Jh75vzmyhx0t+l798dXvMX/6/N4wBrs/vzeMwe7P7w1jsPvze8MY7Oh5P/3f8XrPv7/HvGEMdn9+bxiD3Z/fG8Zg9+f3hjHY/fmd3Pd+iePC3iP36LUmvj6++PX1eRgDc8D3Af8W+Hng3z8T/8g+8/Q9e/jq95g/fX5vGIPdn98bxmD35/eGMdj9+b1hDPXV8376v+P1nnO+x7xhDHZ/fm8Yg92f3xvGYPfn94Yx2P35ndb3fgnbOAAAAAAAAACAH3Dcyt76YiX41ql6Th+D3bxhDuweg92f3xt4wxzYPQa7P7838IY58IYx+PP/N28Yg93sngO7P7+3jMFuds+B3Z/fW8ZgN7vnwO7P74T/f94wBrvZPQd2f35vGYPd7J4Duz+/t4zBbnbPQZQXfQ18YmXvnDlzNHr0aPXr18+cTBcSEqLY2FglJibq3//+twoKCuweIgAAAAAAAADYyidW9j733HP6/vvv1aNHDw0cOFCNGzc2mx0vWLBAGRkZmjp1qr766iu1bNnS7qECAAAAAAAAgC18Iux99NFH1aVLFxPyHmn//v26+OKL9e233+q2227TvHnzbBsjAAAAAAAAANjJJ8LepKSkX7zepEkT/etf/9KJJ56ojz/+uN7HBd/aswT24fsAfB/ACfhZB74PwM8DOAl9D3wfgJ8H3skn9uz9LcHB1Xl1WFiY3UMBAAAAAAAAANv4xMreX5Ofn697773XPL7wwguP+X/vcrl+811KAADqy6/1pN/qVXX1OSz0PQCAP/W833s++h4AwN/6nk+GvdZWDXPnzlVlZaX7gDYr8D377LP10EMPHfPzxcfH/+rfsfUAAKA+RUdHe/xz0PcAAE7peRb6HgDASX3PJ8Pe1atXa8aMGT+6Nnz4cD322GOKjY21bVwAAAAAAAAAYLeAKh9cwlpWVqbt27frrbfe0j//+U8FBATojTfeMAe1/R5riXRNom6tDv6123e4rQcA4C239tSsTCooKDjm/kTfAwA4pefVPAev9wAATul7fhP2HiktLU1DhgxR69attW7dOkVERBx186/LiQQAwBNq27foewAAX1EXPYu+BwDwFZ7qWYHycUlJSerZs6d27NihRYsW2T0cAAAAAAAAALCFz4e9lprkOzs72+6hAACgvOIyrdyV654JH7+JBgAAAABw2KrducotKpO38vmwNycnR8uWLTOPu3btavdwAAAOd6iwVCNT0zTsxYVasTOXoBcAAAAA/ETG1gO66oWFGjctXQUl5V75es/rw97Vq1drzpw5Ki4u/tnfrV+/XldccYVKSkqUnJysPn362DJGAAAsOQUluvrFhVq+M1f5JeW65eVMVVbJHCQKAAAAAPBd323M0egp1SHvku2H9NRnG+SNguXlrK0ZRo4cqUmTJmnAgAHmILbS0lJt375dS5YsUWVlpXr06KGXX37Z7qECABwsK69YI1LTtDG7wNRNY8L03MhBCgok6AUAAAAAX/bF2mxNmr1YpeWVpj6hS5xuOb2rVy7s8fqwt1evXnrwwQf1zTffaO3atcrMzFRZWZkaN26s0047TZdeeqnGjRunsLAwu4cKAHConQcLTdC7bX+hqVvEhmvuxGR1iKub01QBAAAAAPb4cOVe3ThvicoqqrdsOL1HvJ4ZMUBhwUFe+SXx+rC3adOmuvPOO+0eBgAAv2jbfpeGT07TrkNFpm7TOEJzU5LVpnEkMwYAAAAAPuytpbv05/nLVGHtzyfpvD4t9PjV/RUS5L0743p92AsAgLeytmwYkbpQWXklpu4YF6U5E5PUIjbC7qEBAAAAAGph/qId+ttry1VzBtulA1vp4cv6KtiLg145Pex1uVzux1FR3GoLADh6a/bkaWRqmva7Sk3dNT5as1OS1CwmvE571JG9qq6e00LfAwB4C0/0vJ8+H30PAHAsZi3YqnveWuWuhyW21YMX91ZgHZzJ4qm+V8PRYW98fLz7cVVNTA8AwO9YsTNXo6am6VBhmal7tWygWROS1DgqtM7mLjo6us6/DvQ9AIA38kTPs9D3AAB/ROo3m/XP99a467HHtdc/LuhZZ4exearv1XB02AsAwLFavO2Axk7NUH5Juan7t2moGeMTFRsRwmQCAAAAgA97+vMNeuTj9e762pM76fazutVZ0FsfHB32ZmVlcTsPAOCoLdi0XxNmZKiwtMLUiR0aa+rYBEWH1X07LSgocN/ac+TKpNqg7wEAvJEnep6FvgcAOFrWHf+PfrxeT3+x0X3t1tO76qbTOtd50OupvlfD0WGvtW8TezcBAI7GV+v36ZqZi1RSXmnqE7rE6cVRgxURGuSxHuWJ56TvAQC8jad6E30PAHC0Qe+D761R6rdb3Nf+75zumnRSJ3mCp1+TOTrsBQDgaHyyOkvXz1mi0orqoPe07s30zIiBCg/xTNALAAAAAPC8ysoq/f3tlZq9cLv72n0X9tKY49r77PQT9gIA8BveW75HN7+UqfLK6oM8z+ndXE9cPUChwYHMGwAAAAD4qIrKKv3tteV6dfFOU1u7Nfz7kj66OrGtfBlhLwAAv+L1JTv1l1eW6XDOq4v7t9QjV/RTcBBBLwAAAAD4qrKKSv15/jK9s2y3qYMCA/ToFf108YBW8nWEvQAA/IK5adt115srVHU46L06oY0evKSP+SUAAAAAAOCbSsordNO8TH20KsvUwYEBemrYAJ3Tp4X8AWEvAAA/Me27LbrvndXuesyQdvrHBb0USNALAAAAAD6ruKxC185erC/W7TN1aFCgnhs5UKf1iJe/IOwFAOAIz325SQ99uNZdTzqxo+44p7sCrA2cAAAAAAA+qbC0XCkzFun7TftNHR4SqMmjB+uELk3lTwh7AQCQVFVVpcc/3aAnPtvgno+bTuuiW0/vQtALAAAAAD4sv7hM46ZlaNG2g6aOCg3SlLEJSu7YRP7G0WGvy+VyP46KirJ1LAAAe4Pe/3y4Vi98tdl97fazu+m6kzvb3qOO7FV19ZwW+h4AwFt4ouf99PnoewDgXIcKSzVmarqW7cw1dUx4sGaMT9TAto38qu/VcHTYGx8f/6MX+gAA56msrNL9767W9O+3uq/9/fyeGj+0g63jio6OrvPnpO8BALyRJ3qehb4HANhfUKKRU9K1Zk+emYxGkSGaNSFJvVvF+l3fq+HosBcA4GwVlVW6640Veiljh6mtbXkfvLiPhie1tXtoAAAAAIBayM4r1ojUNG3ILjB1XHSY5qQkqVvzGL+eV0eHvVlZWdzOAwAOVV5Rqb++ulxvZO4ydWCA9N/L++myQa3tHppRUFDgvrXnyJVJtUHfAwB4I0/0PAt9DwCca9ehIo2YvFBb9xeaunmDcM2ZmKROTT27qtbOvlfD0WGvtW8TezcBgPOUllfq5pcy9cHKvaYODgzQ/67qrwv6tZS38ER/ou8BALyRp16T0fcAwJm27y/UsMkLTeBradUwQvMmJqttk0h5A09nkY4OewEAzlNcVqHr5yzRZ2uzTR0aFKinhw/Qmb2a2z00AAAAAEAtbNpXoBGT07Q3r9jU7ZtEau7EZLVsGOGYeSXsBQA4RlFpha6ZtUjfbMgxdVhwoF4cPVgndW1q99AAAAAAALWwdm+eRqamKaeg1NRdmkWbPXqbNQh31LwS9gIAHKGgpFzjp2cofcsBU0eGBil1zGAd1ynO7qEBAAAAAGph5a5cjZySpkOFZabu2aKBZk1IVJPoMMfNK2EvAMDv5RaVaczUdC3dccjUMWHBmj4+QYPaNbZ7aAAAAACAWliy/aB5vZdfXG7qfm0aaua4RMVGhjhyXgl7AQB+7YCrVKOmpGnV7jxTx0aEmHd4+7ZuaPfQAAAAAAC1sHDzfk2YniFXaYWpE9o30tSxCYoJd2bQayHsBQD4rez8Yo1KTde6rHxTN4kK1eyUJPVo0cDuoQEAAAAAauHr9fvMmSzFZZWmPr5zE00ePViRoc6OO539/x4A4Lf25BaZU1g357hM3SwmTHMnJqlzsxi7hwYAAAAAqIVPV2fpujlLVFpRHfSe0q2pnhs5SOEhQY6fV8JeAIDf2XGgUMNTF2rHgSJTt2oYYU5hbR8XZffQAAAAAAC18N7yPbr5pUyVV1aZ+uxezfXksAEKDQ5kXp2+stflql7tZYmKIgAALine truncated
<img alt="No description has been provided for this image" class="" src="data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAABXsAAAHjCAYAAABy0synAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjgsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvwVt1zgAAAAlwSFlzAAAewgAAHsIBbtB1PgAAkjRJREFUeJzs3Xd4VHXaxvF7Jj2TRggkQOi9E0jBrth7R5EOsa7d1XXXdV3L+uq69q6hB1TsvfdGCoTee2gJIaRN+sx5/wgZQEEpMzlTvp/r2suZ34Qzt5HNk/PMr1gMwzAEAAAAAAAAAPBpVrMDAAAAAAAAAACOHs1eAAAAAAAAAPADNHsBAAAAAAAAwA/Q7AUAAAAAAAAAP0CzFwAAAAAAAAD8AM1eAAAAAAAAAPADNHsBAAAAAAAAwA/Q7AUAAAAAAAAAP0CzFwAAAAAAAAD8AM1eAAAAAAAAAPADNHsBAAAAAAAAwA/Q7AUAAAAAAAAAP0CzFwAAAAAAAAD8AM1eAAAAAAAAAPADNHsBAAAAAAAAwA/Q7AUAAAAAAAAAP0CzFwAAAAAAAAD8AM1eAAAAAAAAAPADNHsBAAAAAAAAwA/4VLO3vr5ezzzzjI4//njFx8crPDxcycnJOvvss/XGG2+YHQ8AAAAAAAAATBNsdoBDtWXLFp155plavny5EhISdNxxx8lms6mwsFA//PCDbDabrrjiCrNjAgAAAAAAAIApfKLZW1NTo9NPP10rV67Uv//9b/3jH/9QSEiI6/Xq6mqtXr3axIQAAAAAAAAAYC6faPb+3//9n1auXKlrrrlG99133+9ej4yM1JAhQ1o+GAAAAAAAAAB4CYthGIbZIf5IQ0OD2rdvr5KSEq1Zs0Y9evQwOxIAAAAAAAAAeB2vn9m7YMEClZSUqH379urRo4eWLFmid955R9u2bVOrVq10wgkn6Oyzz5bVemhnzTmdTpWUlEhqmhFssVg8GR8AgKNiGIaqq6slSQkJCYdc75pR9wAAvuJoa55E3QMA+A531L0D8fpm7+LFiyVJycnJuvvuu/Xf//5X+05GfvTRR5WSkqL33ntPnTp1+tPrlZSUKDEx0WN5AQDwlKKiIrVt2/aw/gx1DwDgi46k5knUPQCAbzrSuncg7mkZe9CuXbskSQUFBXr00Ud1ww03aNWqVSovL9eXX36pXr16qaCgQOeee64aGhpMTgsAAAAAAAAA5vD6mb3Ns3gbGho0atQoPffcc67XTjvtNH355Zfq3bu3li5dqtdff11jx479w+tFRkYe0vtWVVUdeWhJdrvd9YlyUVGRbDbbUV2PDL75/t6Qwez394YMZr+/N2Qw+/29IYPZ7/9nGaKiov70zx9qDTuSP0Pd8/3394YMZr+/N2Qw+/29IYPZ7+8NGcx+f2/IYEbNO5w/R93z/ff3hgxmv783ZDD7/b0hg9nv7w0ZzH5/b8hgVt07EK9v9kZHR7seX3vttb97vVOnTjr33HP19ttv66uvvvrTZu++ezb90X98d/6lsNlspvxFJ4N3vb83ZDD7/b0hg9nv7w0ZzH5/b8hg9vsfKMPBbjr3/aXhSPYdpO4F5vt7Qwaz398bMpj9/t6Qwez394YMZr+/N2RoqZr32z9H3Quc9/eGDGa/vzdkMPv9vSGD2e/vDRnMfn9vyNCSde9AvL7Z261btwM+PtDXbN++/bCubfZ/fAAAmrVEPaLuAQC8QUvVIuoeAMAbtHQt8vpm79ChQ2WxWGQYhkpKStSxY8fffU3zaauHMi26pdhstv0OkgvUDGbzhu+B2RnMfn9v4A3fA7MzmP3+3sAbvgfekMFTvOHfzRsymM3s74HZ7+8tGcxm9vfA7Pf3lgxmM/t7YPb7e5o3/Pt5Qwazmf09MPv9vSWD2cz+Hpj9/t6SwWxmfw/Mfv99ef0BbUlJSTr++OMlSV999dXvXm9oaND3338vSUpPT2/RbAAAAAAAAADgLby+2StJ9913nyTp//7v/zRv3jzXeGNjo+644w6tX79e0dHRmjhxolkRAQAAAAAAAMBUXr+NgySdeuqpevDBB3XvvffqhBNOUHp6upKSkrRgwQJt3LhREREReu2111ybGgMAAAAAAABAoLEY3rKhxCH44osv9NRTTyknJ0eVlZVKSkrSqaeeqr/97W/q06fPIV3Dbre79vatqqpiw34AgFc72rpF3QMA+Ap31CzqHgDAV3iqZvnEzN5mZ5xxhs444wyzYwAAAAAAAACA1/GJPXsBAAAAAAAAAH+MZi8AAAAAAAAA+AGavQAAAAAAAADgB2j2AgAAAAAAAIAf8KkD2tzNbre7HnNKKwDAmzTXqH1rlbuuKVH3AADewxM177fXo+4BALyFp+pes4Bu9iYmJroeG4ZhYhIAAPYXFRXl9mtS9wAA3sgTNU+i7gEAvJOn6l4ztnEAAMDNahscWr+zyuwYAAAAAIAAE9Aze4uKiljOAwBwK3tdozJn5Gt1UaXmXneMurc5sk9tq6qamsV2u32/mUlHg7oHAPBGnqh5EnUPAOCdPFX3mgV0s9dms1H8AQBuU17doAnTc1WwuUySdM3MfH1x20kKsloO+1qeqE/UPQCAN/JUbaLuAQC8kadrU0A3ewEAcJedlXUaOyVHK3dUSpJiI0L0v8sHH1GjFwAAAACAI0GzFwCAo7StrEZjsnK0vqTpNNWEqFDNmpyhvu1iTE4GAAAAAAgkNHsBADgKG0vsGp2Vo61lNZKk9rHhys7MULcj3KsXAAAAAIAjRbMXAIAjtGpHpcZMydHOyjpJUpfWkcrOzFByq0iTkwEAAAAAAhHNXgAAjsDiLWUaNzVXZdUNkqTeidGalZmuttHhJicDAAAAAAQqmr0AABym3A2lmjQ9T1V1jZKkwcmxmjEpXXGRoSYnAwAAAAAEMpq9AAAchu9WFeu67PmqbXBKktK7xmvK+FRFh4eYnAwAAAAAEOho9gIAcIg+XbJdN79eoAaHIUk6uXcbvTh6mCJCg0xOBgAAAAAAzV4AAA7J2/O36M63FsnZ1OfV2QOS9PSVKQoNtpobDAAAAACAPWj2AgDwJ2b9ulH3vr/M9fyyYcl65JKBCg6i0QsAAAAA8B4B3ey12+2uxzabzcQkAABv9eJ36/ToZytdz8cf01n3nd9fVqvFo+/bXKP2rVXuuqZE3QMAeA9P1LzfXo+6BwDwFp6qe80CutmbmJjoemwYholJAADexjAMPfb5Kr3w3TrX2A0nd9edZ/aWxeLZRq8kRUVFuf2a1D0AgDfyRM2TqHsAAO/kqbrXLKCbvQAAHIjTaej+D5dpxq+bXGN3ntlbfzmlh4mpAAAAAAD4YwHd7C0qKmI5DwBgP40Op+5+Z4nemr/FNfbAhf017pguLZqjqqpKUtPSnn1nJh0N6h4AwBt5ouZJ1D0AgHfyVN1rFtDNXpvNRvEHALjUNzp16xsF+mTJDkmS1SL997LBumxYcotn8UR9ou4BALyRp2oTdQ8A4I08XZsCutkLAECzmnqHrsuer+9X75QkhQRZ9MyVKTp7YDuTkwEAAAAAcGho9gIAAl5lbYMmz8hX7oZSSVJ4iFUvjRmmk3u3NTkZAAAAAACHjmYvACCg7bbXa8K0XC3aUi5JigoL1tQJaUrvGm9yMgAAAAAADg/NXgBAwCqurNXYrFytKqqUJMVFhmjmpHQNSo4zNxgAAAAAAEeAZi8AICBt2V2tMVk52rirWpLUJjpM2ZMz1Dsp2uRkAAAAAAAcGZq9AICAs35nlcZk5Whbea0kqUNchGZnZqhLAid2AwAAAAB8F81eAEBAWbG9QmOn5Kikql6S1C3BpuzMDLWPizA5GQAAAAAAR4dmLwAgYBRs3q3xU3NVUdsoSerbLkazJqcrISrM5GQAAAAAABw9mr0AgIDw67pdypyRJ3u9Q5KU0ilO0yekKzYyxORkAAAAAAC4B81eAIDf+2Zlka7PXqC6Rqck6djurfXquFTZwiiDAAAAAAD/wV0uAMCvfbR4m259faEanYYk6dQ+bfX86KEKDwkyORkAAAAAAO4V0M1eu93uemyzcQI7APibuXmFuvudxdrT59V5g9rpySuGKCTIam6wQ9Bco/atVe66pkTdAwB4D0/UvN9ej7oHAPAWnqp7zQK62ZuYmOh6bBiGiUkAAO429acNeuCj5a7nV6R21MOXDFSQ1WJiqkMXFRXl9mtS9wAA3sgTNU+i7gEAvJOn6l6zgG72AgD8j2EYeu6btXr8y9WusUnHddW95/WVxeIbjV4AAAAAAI5EQDd7i4qKWM4DAH7EMAw98ulKvfzDetfYzaf21G2n9fS5Rm9VVZWkpqU9+85MOhrUPQCAN/JEzZOoewAA7+SputcsoJu9NpuN4g8AfsLpNHTv+0s1O2eza+wf5/TRNSd2NzHVkfNEfaLuAQC8kadqE3UPAOCNPF2bArrZCwDwD40Op+58a7HeLdgqSbJYpAcvHKAxwzubnAwAAAAAgJbj/ceRAwDwB+oaHfrLnAWuRm+Q1aInRw6h0QsAAAAAcLtPl2zXxhK72TEOimYvAMBnVdc3KnNGvj5fViRJCg2y6sXRQ3VRSgeTkwEAAAAA/M3cvELdMGeBRmflaMvuarPjHBDNXgCAT6qobdC4Kbn6cU2JJCkiJEhTJqTqjP5JJicDAAAAAPibGb9s1F1vL5ZhSFvLajQ3f4vZkQ6IPXsBAD6n1F6vcVNztHRrhSQpOixY0yamKbVLvMnJAAAAAAD+5sXv1unRz1a6nk88rotuO62niYkOzmdn9t51112yWCyyWCx66KGHzI4DAGghRRW1uuLlX12N3nhbqF67ZjiNXgAAAACAWxmGoSe+WLVfo/fGU3roX+f1k8ViMTHZwfnkzN5ffvlFjz/+uCwWiwzDMDsOAKCFFJZWa3RWLine truncated
</div>
</div>
</div>
@@ -1719,7 +1749,7 @@ with mpl.style.context("boilerplot.ieeetran"):
<div class="jp-Collapser jp-InputCollapser jp-Cell-inputCollapser">
</div>
<div class="jp-InputArea jp-Cell-inputArea">
<div class="jp-InputPrompt jp-InputArea-prompt">In [4]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="jp-InputPrompt jp-InputArea-prompt">In [25]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="CodeMirror cm-s-jupyter">
<div class="zeroclipboard-container">
<clipboard-copy ,="" aria-label="Copy to Clipboard" for="cell-4">
@@ -1732,27 +1762,23 @@ with mpl.style.context("boilerplot.ieeetran"):
</div>
</clipboard-copy>
</div>
<div class="highlight-ipynb hl-python"><pre><span></span><span class="n">in_spikes</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">asarray</span><span class="p">([[</span><span class="mf">0.00</span><span class="p">]])</span>
<span class="n">comp_times</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">linspace</span><span class="p">(</span><span class="mf">0.0</span><span class="p">,</span> <span class="n">max_time</span><span class="p">,</span> <span class="mi">500</span><span class="p">)</span>
<span class="n">state</span> <span class="o">=</span> <span class="n">snn</span><span class="o">.</span><span class="n">state_at_t</span><span class="p">(</span><span class="n">in_spikes</span><span class="p">,</span> <span class="n">comp_times</span><span class="p">)</span>
<span class="n">spikes</span> <span class="o">=</span> <span class="n">snn</span><span class="o">.</span><span class="n">spikes_until_t</span><span class="p">(</span><span class="n">in_spikes</span><span class="p">,</span> <span class="n">max_time</span><span class="p">)</span>
<div class="highlight-ipynb hl-python"><pre><span></span><span class="n">comp_times</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">linspace</span><span class="p">(</span><span class="mf">0.0</span><span class="p">,</span> <span class="n">max_time</span><span class="p">,</span> <span class="mi">500</span><span class="p">)</span>
<span class="n">_</span><span class="p">,</span> <span class="n">state</span> <span class="o">=</span> <span class="n">state_at_t</span><span class="p">(</span><span class="n">comp_times</span><span class="p">)</span>
</pre></div>
<div class="clipboard-copy-txt" id="cell-4">in_spikes = jnp.asarray([[0.00]])
comp_times = jnp.linspace(0.0, max_time, 500)
state = snn.state_at_t(in_spikes, comp_times)
spikes = snn.spikes_until_t(in_spikes, max_time)</div>
<div class="clipboard-copy-txt" id="cell-4">comp_times = jnp.linspace(0.0, max_time, 500)
_, state = state_at_t(comp_times)</div>
</div>
</div>
</div>
</div>
</div>
</div><div class="jp-Cell jp-CodeCell jp-Notebook-cell">
<div class="jp-Cell jp-CodeCell jp-Notebook-cell" id="cell-id=b73fb06d">
</div><div class="jp-Cell jp-CodeCell jp-Notebook-cell jp-mod-noOutputs">
<div class="jp-Cell jp-CodeCell jp-Notebook-cell jp-mod-noOutputs" id="cell-id=06ae6a94">
<div class="jp-Cell-inputWrapper" tabindex="0">
<div class="jp-Collapser jp-InputCollapser jp-Cell-inputCollapser">
</div>
<div class="jp-InputArea jp-Cell-inputArea">
<div class="jp-InputPrompt jp-InputArea-prompt">In [5]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="jp-InputPrompt jp-InputArea-prompt">In [26]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="CodeMirror cm-s-jupyter">
<div class="zeroclipboard-container">
<clipboard-copy ,="" aria-label="Copy to Clipboard" for="cell-5">
@@ -1765,16 +1791,152 @@ spikes = snn.spikes_until_t(in_spikes, max_time)</div>
</div>
</clipboard-copy>
</div>
<div class="highlight-ipynb hl-python"><pre><span></span><span class="k">def</span><span class="w"> </span><span class="nf">compute_nullclines</span><span class="p">(</span><span class="n">neuron_model</span><span class="p">,</span> <span class="n">u_range</span><span class="p">,</span> <span class="n">v_range</span><span class="p">,</span> <span class="n">resolution</span><span class="o">=</span><span class="mi">200</span><span class="p">):</span>
<span class="w"> </span><span class="sd">"""</span>
<span class="sd"> Compute nullclines</span>
<span class="sd"> du/dt = 0 (u-nullcline)</span>
<span class="sd"> dv/dt = 0 (v-nullcline)</span>
<span class="sd"> """</span>
<span class="n">u_vals</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">linspace</span><span class="p">(</span><span class="n">u_range</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">u_range</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">resolution</span><span class="p">)</span>
<span class="n">v_vals</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">linspace</span><span class="p">(</span><span class="n">v_range</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">v_range</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">resolution</span><span class="p">)</span>
<span class="n">U</span><span class="p">,</span> <span class="n">V</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">meshgrid</span><span class="p">(</span><span class="n">u_vals</span><span class="p">,</span> <span class="n">v_vals</span><span class="p">)</span>
<span class="n">UV</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">stack</span><span class="p">(</span>
<span class="p">[</span><span class="n">U</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">),</span> <span class="n">V</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">),</span> <span class="n">jnp</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">resolution</span> <span class="o">*</span> <span class="n">resolution</span><span class="p">,))],</span> <span class="n">axis</span><span class="o">=</span><span class="mi">1</span>
<span class="p">)</span>
<span class="n">dS</span> <span class="o">=</span> <span class="n">neuron_model</span><span class="o">.</span><span class="n">dynamics</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="n">UV</span><span class="p">,</span> <span class="p">{})</span>
<span class="n">dU</span> <span class="o">=</span> <span class="n">dS</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">U</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
<span class="n">dV</span> <span class="o">=</span> <span class="n">dS</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">V</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
<span class="k">return</span> <span class="n">U</span><span class="p">,</span> <span class="n">V</span><span class="p">,</span> <span class="n">dU</span><span class="p">,</span> <span class="n">dV</span>
<span class="k">def</span><span class="w"> </span><span class="nf">plot_vf</span><span class="p">(</span><span class="n">ax</span><span class="p">,</span> <span class="n">neuron_model</span><span class="p">,</span> <span class="n">u_range</span><span class="p">,</span> <span class="n">v_range</span><span class="p">):</span>
<span class="kn">import</span><span class="w"> </span><span class="nn">numpy</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">np</span>
<span class="n">u_sparse</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">linspace</span><span class="p">(</span><span class="n">u_range</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">u_range</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="mi">30</span><span class="p">)</span>
<span class="n">v_sparse</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">linspace</span><span class="p">(</span><span class="n">v_range</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">v_range</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="mi">30</span><span class="p">)</span>
<span class="n">Us</span><span class="p">,</span> <span class="n">Vs</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">meshgrid</span><span class="p">(</span><span class="n">u_sparse</span><span class="p">,</span> <span class="n">v_sparse</span><span class="p">)</span>
<span class="n">U</span><span class="p">,</span> <span class="n">V</span><span class="p">,</span> <span class="n">dU</span><span class="p">,</span> <span class="n">dV</span> <span class="o">=</span> <span class="n">compute_nullclines</span><span class="p">(</span><span class="n">neuron_model</span><span class="p">,</span> <span class="n">u_range</span><span class="p">,</span> <span class="n">v_range</span><span class="p">,</span> <span class="mi">200</span><span class="p">)</span>
<span class="n">UVs</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">stack</span><span class="p">([</span><span class="n">Us</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">),</span> <span class="n">Vs</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">),</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ones</span><span class="p">((</span><span class="mi">30</span> <span class="o">*</span> <span class="mi">30</span><span class="p">,))],</span> <span class="n">axis</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
<span class="n">dS</span> <span class="o">=</span> <span class="n">neuron_model</span><span class="o">.</span><span class="n">dynamics</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="n">UVs</span><span class="p">,</span> <span class="p">{})</span>
<span class="n">dUs</span> <span class="o">=</span> <span class="n">dS</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">Us</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
<span class="n">dVs</span> <span class="o">=</span> <span class="n">dS</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">Vs</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
<span class="c1"># Normalize for visualization</span>
<span class="n">magnitude</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">sqrt</span><span class="p">(</span><span class="n">dUs</span><span class="o">**</span><span class="mi">2</span> <span class="o">+</span> <span class="n">dVs</span><span class="o">**</span><span class="mi">2</span><span class="p">)</span>
<span class="n">magnitude</span><span class="p">[</span><span class="n">magnitude</span> <span class="o">==</span> <span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="mi">1</span>
<span class="n">dUs_norm</span> <span class="o">=</span> <span class="n">dUs</span> <span class="o">/</span> <span class="n">magnitude</span>
<span class="n">dVs_norm</span> <span class="o">=</span> <span class="n">dVs</span> <span class="o">/</span> <span class="n">magnitude</span>
<span class="c1"># Nullclines</span>
<span class="n">ax</span><span class="o">.</span><span class="n">contour</span><span class="p">(</span><span class="n">U</span><span class="p">,</span> <span class="n">V</span><span class="p">,</span> <span class="n">dU</span><span class="p">,</span> <span class="n">levels</span><span class="o">=</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">colors</span><span class="o">=</span><span class="s2">"blue"</span><span class="p">,</span> <span class="n">linewidths</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">linestyles</span><span class="o">=</span><span class="s2">"-"</span><span class="p">)</span>
<span class="n">ax</span><span class="o">.</span><span class="n">contour</span><span class="p">(</span><span class="n">U</span><span class="p">,</span> <span class="n">V</span><span class="p">,</span> <span class="n">dV</span><span class="p">,</span> <span class="n">levels</span><span class="o">=</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">colors</span><span class="o">=</span><span class="s2">"red"</span><span class="p">,</span> <span class="n">linewidths</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">linestyles</span><span class="o">=</span><span class="s2">"-"</span><span class="p">)</span>
<span class="c1"># Vector field</span>
<span class="n">ax</span><span class="o">.</span><span class="n">quiver</span><span class="p">(</span><span class="n">Us</span><span class="p">,</span> <span class="n">Vs</span><span class="p">,</span> <span class="n">dUs_norm</span><span class="p">,</span> <span class="n">dVs_norm</span><span class="p">,</span> <span class="n">magnitude</span><span class="p">,</span> <span class="n">cmap</span><span class="o">=</span><span class="s2">"viridis"</span><span class="p">,</span> <span class="n">alpha</span><span class="o">=</span><span class="mf">0.6</span><span class="p">)</span>
<span class="n">ax</span><span class="o">.</span><span class="n">set_xlabel</span><span class="p">(</span><span class="s2">"v"</span><span class="p">)</span>
<span class="n">ax</span><span class="o">.</span><span class="n">set_ylabel</span><span class="p">(</span><span class="s2">"w"</span><span class="p">)</span>
<span class="n">ax</span><span class="o">.</span><span class="n">legend</span><span class="p">([</span><span class="s2">"u-nullcline (du/dt=0)"</span><span class="p">,</span> <span class="s2">"v-nullcline (dv/dt=0)"</span><span class="p">],</span> <span class="n">loc</span><span class="o">=</span><span class="s2">"upper right"</span><span class="p">)</span>
<span class="n">ax</span><span class="o">.</span><span class="n">set_xlim</span><span class="p">(</span><span class="n">u_range</span><span class="p">)</span>
<span class="n">ax</span><span class="o">.</span><span class="n">set_ylim</span><span class="p">(</span><span class="n">v_range</span><span class="p">)</span>
<span class="n">ax</span><span class="o">.</span><span class="n">axhline</span><span class="p">(</span><span class="n">y</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span> <span class="n">color</span><span class="o">=</span><span class="s2">"gray"</span><span class="p">,</span> <span class="n">linestyle</span><span class="o">=</span><span class="s2">"--"</span><span class="p">,</span> <span class="n">alpha</span><span class="o">=</span><span class="mf">0.3</span><span class="p">)</span>
<span class="n">ax</span><span class="o">.</span><span class="n">axvline</span><span class="p">(</span><span class="n">x</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span> <span class="n">color</span><span class="o">=</span><span class="s2">"gray"</span><span class="p">,</span> <span class="n">linestyle</span><span class="o">=</span><span class="s2">"--"</span><span class="p">,</span> <span class="n">alpha</span><span class="o">=</span><span class="mf">0.3</span><span class="p">)</span>
</pre></div>
<div class="clipboard-copy-txt" id="cell-5">def compute_nullclines(neuron_model, 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.zeros((resolution * resolution,))], axis=1
)
dS = neuron_model.dynamics(0, UV, {})
dU = dS[:, 0].reshape(U.shape)
dV = dS[:, 1].reshape(V.shape)
return U, V, dU, dV
def plot_vf(ax, neuron_model, u_range, v_range):
import numpy as np
u_sparse = jnp.linspace(u_range[0], u_range[1], 30)
v_sparse = jnp.linspace(v_range[0], v_range[1], 30)
Us, Vs = jnp.meshgrid(u_sparse, v_sparse)
U, V, dU, dV = compute_nullclines(neuron_model, u_range, v_range, 200)
UVs = jnp.stack([Us.reshape(-1), Vs.reshape(-1), jnp.ones((30 * 30,))], axis=1)
dS = neuron_model.dynamics(0, 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("v")
ax.set_ylabel("w")
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)</div>
</div>
</div>
</div>
</div>
</div>
</div><div class="jp-Cell jp-CodeCell jp-Notebook-cell">
<div class="jp-Cell jp-CodeCell jp-Notebook-cell" id="cell-id=b73fb06d">
<div class="jp-Cell-inputWrapper" tabindex="0">
<div class="jp-Collapser jp-InputCollapser jp-Cell-inputCollapser">
</div>
<div class="jp-InputArea jp-Cell-inputArea">
<div class="jp-InputPrompt jp-InputArea-prompt">In [27]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="CodeMirror cm-s-jupyter">
<div class="zeroclipboard-container">
<clipboard-copy ,="" aria-label="Copy to Clipboard" for="cell-6">
<div>
<span class="notice" hidden="">Copied!</span>
<svg aria-hidden="true" class="clipboard-copy-icon" data-view-component="true" height="20" version="1.1" viewbox="0 0 16 16" width="20">
<path d="M0 6.75C0 5.784.784 5 1.75 5h1.5a.75.75 0 010 1.5h-1.5a.25.25 0 00-.25.25v7.5c0 .138.112.25.25.25h7.5a.25.25 0 00.25-.25v-1.5a.75.75 0 011.5 0v1.5A1.75 1.75 0 019.25 16h-7.5A1.75 1.75 0 010 14.25v-7.5z" fill="currentColor" fill-rule="evenodd"></path>
<path d="M5 1.75C5 .784 5.784 0 6.75 0h7.5C15.216 0 16 .784 16 1.75v7.5A1.75 1.75 0 0114.25 11h-7.5A1.75 1.75 0 015 9.25v-7.5zm1.75-.25a.25.25 0 00-.25.25v7.5c0 .138.112.25.25.25h7.5a.25.25 0 00.25-.25v-7.5a.25.25 0 00-.25-.25h-7.5z" fill="currentColor" fill-rule="evenodd"></path>
</svg>
</div>
</clipboard-copy>
</div>
<div class="highlight-ipynb hl-python"><pre><span></span><span class="k">with</span> <span class="n">mpl</span><span class="o">.</span><span class="n">style</span><span class="o">.</span><span class="n">context</span><span class="p">(</span><span class="s2">"boilerplot.ieeetran"</span><span class="p">):</span>
<span class="n">fig</span><span class="p">,</span> <span class="n">ax</span> <span class="o">=</span> <span class="n">plt</span><span class="o">.</span><span class="n">subplots</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="n">figsize</span><span class="o">=</span><span class="p">(</span><span class="mf">6.9</span><span class="p">,</span> <span class="mf">2.6</span><span class="p">),</span> <span class="n">dpi</span><span class="o">=</span><span class="mi">200</span><span class="p">)</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">comp_times</span><span class="p">,</span> <span class="n">state</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="p">:,</span> <span class="mi">0</span><span class="p">])</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">comp_times</span><span class="p">,</span> <span class="n">state</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="p">:,</span> <span class="mi">1</span><span class="p">],</span> <span class="s2">"--"</span><span class="p">)</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">comp_times</span><span class="p">,</span> <span class="n">state</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">])</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">comp_times</span><span class="p">,</span> <span class="n">state</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">],</span> <span class="s2">"--"</span><span class="p">)</span>
<span class="c1"># ax[0].plot(comp_times, state[0, :, 2], "-.")</span>
<span class="p">[</span><span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">axvline</span><span class="p">(</span><span class="n">s</span><span class="p">,</span> <span class="n">alpha</span><span class="o">=</span><span class="mf">0.2</span><span class="p">,</span> <span class="n">color</span><span class="o">=</span><span class="s2">"g"</span><span class="p">,</span> <span class="n">linestyle</span><span class="o">=</span><span class="s2">"--"</span><span class="p">)</span> <span class="k">for</span> <span class="n">s</span> <span class="ow">in</span> <span class="n">jnp</span><span class="o">.</span><span class="n">unique</span><span class="p">(</span><span class="n">spikes</span><span class="p">)]</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">set_xlabel</span><span class="p">(</span><span class="s2">"Time (ms)"</span><span class="p">)</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">legend</span><span class="p">([</span><span class="s2">"v"</span><span class="p">,</span> <span class="s2">"vslow"</span><span class="p">,</span> <span class="s2">"syn"</span><span class="p">,</span> <span class="s2">"Spike"</span><span class="p">])</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">legend</span><span class="p">([</span><span class="s2">"v"</span><span class="p">,</span> <span class="s2">"vslow"</span><span class="p">,</span> <span class="s2">"syn"</span><span class="p">])</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">state</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="p">:,</span> <span class="mi">0</span><span class="p">],</span> <span class="n">state</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="p">:,</span> <span class="mi">1</span><span class="p">])</span>
<span class="n">plot_vf</span><span class="p">(</span><span class="n">ax</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">neuron_model</span><span class="p">,</span> <span class="p">[</span><span class="o">-</span><span class="mi">2</span><span class="p">,</span> <span class="mi">2</span><span class="p">],</span> <span class="p">[</span><span class="o">-</span><span class="mi">2</span><span class="p">,</span> <span class="mi">2</span><span class="p">])</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">state</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span> <span class="n">state</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">])</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">state</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span> <span class="n">state</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">],</span> <span class="s2">"."</span><span class="p">,</span> <span class="n">label</span><span class="o">=</span><span class="s2">"start"</span><span class="p">)</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">state</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span> <span class="n">state</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">],</span> <span class="s2">"."</span><span class="p">,</span> <span class="n">label</span><span class="o">=</span><span class="s2">"end"</span><span class="p">)</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">set_xlabel</span><span class="p">(</span><span class="s2">"v"</span><span class="p">)</span>
@@ -1782,16 +1944,17 @@ spikes = snn.spikes_until_t(in_spikes, max_time)</div>
<span class="n">ax</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">legend</span><span class="p">()</span>
<span class="n">plt</span><span class="o">.</span><span class="n">show</span><span class="p">()</span>
</pre></div>
<div class="clipboard-copy-txt" id="cell-5">with mpl.style.context("boilerplot.ieeetran"):
<div class="clipboard-copy-txt" id="cell-6">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])
ax[0].plot(comp_times, state[0, :, 1], "--")
ax[0].plot(comp_times, state[:, 0, 0])
ax[0].plot(comp_times, state[:, 0, 1], "--")
# ax[0].plot(comp_times, state[0, :, 2], "-.")
[ax[0].axvline(s, alpha=0.2, color="g", linestyle="--") for s in jnp.unique(spikes)]
ax[0].set_xlabel("Time (ms)")
ax[0].legend(["v", "vslow", "syn", "Spike"])
ax[0].legend(["v", "vslow", "syn"])
ax[1].plot(state[0, :, 0], state[0, :, 1])
plot_vf(ax[1], neuron_model, [-2, 2], [-2, 2])
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].set_xlabel("v")
@@ -1809,7 +1972,7 @@ spikes = snn.spikes_until_t(in_spikes, max_time)</div>
<div class="jp-OutputArea-child">
<div class="jp-OutputPrompt jp-OutputArea-prompt"></div>
<div class="jp-RenderedImage jp-OutputArea-output" tabindex="0">
<img alt="No description has been provided for this image" class="" src="data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAABXsAAAIfCAYAAAAhag9YAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjgsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvwVt1zgAAAAlwSFlzAAAewgAAHsIBbtB1PgABAABJREFUeJzsnQd4HOX19Y96l1wky3LvveAGNraB2FRD6J0QCCHlI0AgyT90AiGUACnUEJIQSkIJvRkw2DRjAzY22Ni49yZbstV7+Z47s7P7zuystJK2zOyeH8+yZWZ3x7tXszPnPe+5Ca2tra0ghBBCCCGEEEIIIYQQ4moSo70BhBBCCCGEEEIIIYQQQroOxV5CCCGEEEIIIYQQQgiJASj2EkIIIYQQQgghhBBCSAxAsZcQQgghhBBCCCGEEEJiAIq9hBBCCCGEEEIIIYQQEgNQ7CWEEEIIIYQQQgghhJAYgGIvIYQQQgghhBBCCCGExAAUewkhhBBCCCGEEEIIISQGoNhLCCGEEEIIIYQQQgghMQDFXkIIIYQQQgghhBBCCIkBKPYSQgghhBBCCCGEEEJIDECxlxBCCCGEEEIIIYQQQmIAir2EEEIIIYQQQgghhBASA1DsJYQQQgghhBBCCCGEkBiAYi8hhBBCCCGEEEIIIYTEABR7CSGEEEIIIYQQQgghJAag2EsIIYQQQgghhBBCCCExQDLijJaWFpSUlGi3MzMzkZCQEO1NIoQQQgiJOVpbW1FTU6Pdzs/PR2IiPQZugsfMhBBCCCHuPF6OO7FXhN7CwsJobwYhhBBCSNxQXFyMXr16RXszSAfgMTMhhBBCiDuPl+POYlFdXR3tTSCEEEIIiSt4/EUIIYQQQkhkiDtnr0Q32FFVVdXuSYrhCBa1PSsrK6j3i+Tz3LCNnX2eG7axs89zwzZ29nlu2MbOPs8N29jZ57lhGzv7PG4jP/94rq3OPq+jz8nOzg76+Is4F/U760iNhbpmnfg+sfIekXqfWHmPSL0P34OfF2vYGX8n/Hvk5xWpfX04jpfjTuxVM3q3bNnitUh35MuTdTvzZUfyeW7Yxs4+zw3b2NnnuWEbO/s8N2xjZ5/nhm3s7PPcsI2dfR63kZ9/PNdWZ58XzHOMAfT9+/djyJAh2m32SHAfgb6zUJzwdLZmnfg+sfIekXqfWHmPSL0P34OfF+vLGX8nkXqfWHmPSL2Pm9+j2iZ1IJTHy3En9kajyAkhhBBC4gXj2IrHWLGD6jyRRiKEEEIIIaTzhLuXWNyJveqJR0fdvJE8uO3M+3V2GyP9vM7gln8bP5PofR5deT83/L11Fjf82/h3E73Poyvv54ba6ixu+Le55e/G7jaJTyJVs5HeX4SLWPq8YuU9IkWsfF6xVMORIJY+r1h5j0gRS59XLP1b3E7cib1OYeuhrWhqaUJyYjIGdx8c7c0hDvnut5VtQ2pdapt1wdohna2v9mqHtUXCVVusL0LcTTjzEAkhhBBC4o0qT+yZXXZvKKDYGyUqGyrR2NyIlKSUaG0CceB3X9VYheSE5DbrgrVDOltf7dUOa4uEq7ZYX4S4G0afEUIIIYSEjnAPoieG9dUJIYQQQgghhBBCCCGERAQ6e8MM80T4mXSkRlYVr9LccfEO/274mbBO+HfDfQn3r4Q4FR6nxC/87uMTfu/xC7/7+CQrBnKB6ewlhBBCCCGEEEIIIYSQGIBiLyGEEEIIIYQQQgghhMQAFHsJIYQQQgghhBBCCCEkBnC82NvY2IiFCxfi//7v/zBt2jR069YNKSkp6N27N0499VS8/fbb0d5EQgghhBBCCCGEEEIIiTqOb9D28ccf47jjjtNui8A7a9YsLSx57dq1ePPNN7XLT3/6Uzz22GNISEiI9uYSQgghhBASU1RXV3tvy3E4IYQQQgjp+rGVeowVV2JvYmIizjrrLPzyl7/E7NmzTcteeOEFXHTRRXj88ccxc+ZM/PCHP4zadhJCCCGEEBKLFBYWem+7vTs1IYQQQki0yc7ODuvrO17snTNnjnax47zzzsP777+Pf/3rX3j66addJfYWZBagubUZSQlJ0d4U4qDvPpi6YO2QcNUXa4tw30UIIYQQQggh7sbxYm97TJo0SbveuXMn3ERRTlG0N4E48LsPpi5YOyRc9cXaIuGqLdYXIe6muLiY8Q2EEEIIISGiqqrKG+OgzqAKFa4Xezdu3KhdFxVRPCWEEEIIISTUSE4vs3oJIYQQQkJDuI+rXC327tu3D08++aR2W3J9O0pbQcg8oCVOZOfBGry4fCdOGNcbY/vkRXtzCCGEkIDHU+FqOEEIIYS4Eck8b2huwaHqRhRX1OmXynrs99w+ILcr61FcUY+SqvqobmtuejKGF+ZgRGE2BudnYUh+NgYXZKFf9wykJTOKkhCnk9Dq0i4LTU1NOPHEE7Fw4UKMHz8ey5cvR2pqarvPkxOPYIKQXfqxkBjnqHs/xI6DNcjLSME3vzs+2ptDCCGEICEhIaipahxIdxfqMTO/P0JIvNDS0oqaxmbsKavFtpJqbCuVSw027a/C2j0VqKpvivYmOp7pQ3rgpHFFOHFcbxTmpkd7cwiJy+Mt1zp7f/7zn2tCb8+ePfHSSy8FJfQ6iVXFq9DY3IiUpBRMKJwQ7c0hDvjum5pb8NTyj1HULRX9uuXY1oUIvQ0JW7G/rhmrinuzdkjQ9dXesmCWk/imK7UlLN+9EgkJLawvQgghhIQcOZfaU1aHdfsqsKG4El/vLMfKHYdQWt3ATzvCfL7loHb53RtrbJefPL4IV88djpG9c/jdEBImXCn2/vKXv8S//vUvdO/eHe+//z5GjBjhqmYTj3y4CU8vX47LZw/AEYN7Rfz9iTN5aul23PL2GmSnJ+C/lx/pt7y+qTkq20UIIcLbq/egpKoG500d3KnfvT8uXIKzphThR0cO4wcaow0mrISr4QQhhJDYo6KuEev3VWLN7nKs3VuBFTvKNDctiT3eXr1Xu1j5vxNG4qdHDUFKUmJUtouQWMJ1Yu+vf/1rPPjgg+jWrRsWLFiASZMmuarZRE1DE+57bz0aEmpw6+tr8N41FHuJzh1vrQUSgMq6Ju0gZ1pf8yezvyK6uU2EkPjl040H8NDCTWhFM1ISU3B4v4799srvXlNCK57/cifF3hiE8QyEEEIk/qCkuh5rdlfg651lWL79ID7bVBo3H0yvnDTNqTq8Vw5G9pac22z0yEpFdloyMtOSkJmShOQYEDEbm1tQXd+knbNKvvBH6/fj1ZW7setQbZdfW44X5WKQn52Gpy87HGP65Hb5tQmJN1wl9v72t7/Fn//8Z+Tl5WlC79SpU+E2DnIaCQmCjFT/P03JjersgdeybQfRMzsVw3pxqgwJLXWNzdoBvTRr6Nc9kx9vjPLG13u8t1/+ahfuOTX45za3dD4DX6ZhyknFiMIcujwIIYSQMCN9aySTVhqFrdlTgcUbS/Dh+v2aqBdLiAA7vm8eJvTPw4S+3TCwZyZy0pORmZqMrLQkpCcnITGx/Tz6eERct90yU7VL/x6ZmDKwO359/MiA0RolVQ34dnc5/vz+Bs2x3RGkSd28Bz/13u+dm47Xr5zJHGBCYknsvf7663HfffdpQq9EN0ybNg1uREbACGmP9BT/Ud+95XUd/uDKaxpx1fMr8cmGA9r9I4f2xN8umoK8zBR+CcSEDAg8+PE6nDAuHxP65Qf16Xy2qQRX/HcFymsbkZqUiF8fP0KbehVMsybiLqRRiUFqSsc6MJd2spv0fe+twyMfbtZu98lLxz8vmUZnByGEENIB4bauUcS2eqzeXa4d6326scT10Qijeufg8ME9NJFxaEE2umWmIDcjBdmpyRRoHYY4mXvnpWuXY8eYY50amlqwpaQKf/toM15XTAVtsa+iDkfctdB7/6ELJuH7E/uEfLsJiQVcIfbefPPN+OMf/+iNbnCr0CtQ7CWdZU95rd8BXHvc8vq3XqFXWLK5FD99Zjme+fERSE12/zQiEjrOeWwpGhL244N1e/HeNd9rd/3KukZc88LXmtArNDS34O531iErLRk/mD6QX02MUdfQHHAwatehGjy4cCPG983G90YV+T23M26gl77a5RV6hT3ldTj/8aV486pZGNgz8ln7hMQ7kr9swNgOQqKHzNjbXVaLr7Yf0gbdP9l4AMUui3obkp+FGUN7aoLtkHyPWJueguz0ZCTRTRs3yLnoqN65eOD8SdpFPa6UeMP31hS3+xpXPbdSuwh/OmcizprSL6zbTEg4jq3UY6y4EnvfeOMN3HnnndrtYcOG4ZFHHrFdLz8/H/fffz+cToVHGCGkLex03H0WZ297Wq/8UL61yn+U9IutB/H4J5tx5Zzh/BKIp5Y6Ps3+0Y82a1P8rPz+rbWYPqQHI0NijFrF2ZuWZHb2/vTpr/D1Pmm00YypA/1d4fsr6zpcjyIeW6moa8L1L6/Gsz85gu5xQiKM2mivM78ZhBB/5G+ppqFZi/mT38rPt0jGbYlmznA6/XtkYPbwAswelo8RvXPQIzNVi0GIhUxaEl0kFu7vF/viOsVY8s9Pt+ChRZvafN6vX/xGuwhvXTUL4/rmhX1bCekK2dnZCCeOF3sPHjzovb18+XLtYsfAgQNdIfYaLjhC2qLF5kRqT5lF7G3nI3xm6XYYUZnXHjsCs0fka+5Nyc8Ux9yZk/uhT7cMfhFEy2br6MnJKyt2abclvmHhr4/GvxZvxZNLtmlTsu6evw7/utS9MzBI22JvqsXZq+WveZI7ymx+4zraXPKbXeXYcbDGGz3z+A+n4oS/fKI5mZZuKdWm+p0+ydLBkhBCCHEQcqwkU87FgfvmN3uCcilGi6NHFODY0b0waUB39O+eidyMZA6qEseQl5GiZQIbucCb9lfi2D9/0uZzTnlosXY9oV8eXrtiLine truncated
<img alt="No description has been provided for this image" class="" src="data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAABXsAAAIfCAYAAAAhag9YAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjgsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvwVt1zgAAAAlwSFlzAAAewgAAHsIBbtB1PgABAABJREFUeJzs3XeYG9W5+PHvzKittknaorK9N3cMGEwvoYYewiWEnoQkl0tIQi6kkh/pJCEBAgRCKKGXi+lgU1wBY1xwb2uv6/auLs3M7w/tatfYXmlt2Wub83keP1aZOXPmTNnRO++cI+m6riMIgiAIgiAIgiAIgiAIgiAc1uSxroAgCIIgCIIgCIIgCIIgCIKw/0SwVxAEQRAEQRAEQRAEQRAE4Qgggr2CIAiCIAiCIAiCIAiCIAhHABHsFQRBEARBEARBEARBEARBOAKIYK8gCIIgCIIgCIIgCIIgCMIRQAR7BUEQBEEQBEEQBEEQBEEQjgAi2CsIgiAIgiAIgiAIgiAIgnAEEMFeQRAEQRAEQRAEQRAEQRCEI4AI9gqCIAiCIAiCIAiCIAiCIBwBRLBXEARBEARBEARBEARBEAThCCCCvYIgCIIgCIIgCIIgCIIgCEcAEewVBEEQBEEQBEEQBEEQBEE4AohgryAIgiAIgiAIgiAIgiAIwhFABHsFQRAEQRAEQRAEQRAEQRCOACLYKwiCIAiCIAiCIAiCIAiCcAQQwV5BEARBEARBEARBEARBEIQjgAj2CoIgCIIgCIIgCIIgCIIgHAFEsFcQBEEQBEEQBEEQBEEQBOEIIIK9giAIgiAIgiAIgiAIgiAIRwAR7BUEQRAEQRCEgyQSifD+++9z2223cfTRR2Oz2TAajbhcLi644ALefPPNfS77vffe49xzzyU3N5e0tDRqa2v52c9+htfrTeEaCIIgCIIgCIcySdd1fawrIQiCIAiCIAhfBu+99x5nnnkmAC6Xi6OOOor09HRWr17NypUrAfj2t7/NQw89hCRJSZd7zz338MMf/hBJkjjxxBNxOp3MmzePlpYWampqmD9/Prm5uQdknQRBEARBEIRDhwj2CoIgCIIgCMJB8sEHH/DAAw9wyy23cOKJJ+7y3fPPP883vvENVFXliSee4Oqrr06qzKVLl3LUUUchyzKvv/4655xzDgB+v58LLriA999/n0svvZSXXnop5esjCIIgCIIgHFpEsFcQBEEQBEEQDhE33ngjjz76KKeffjrvvfdeUvNcfvnlvPjii9x444088sgju3y3ZcsWysvL0TSNNWvWUFtbeyCqLQiCIAiCIBwiRJ+9giAIgiAIgnCImDx5MgDbtm1LavpwOBzv5/fKK6/c7fuSkhKmT58OwCuvvJKiWgqCIAiCIAiHKhHsFQRBEARBEIRDxIYNGwBwu91JTb9+/Xr8fj8AU6dO3eM0g58vXbo0BTUUBEEQBEEQDmWGsa7A4UjTNDo6OgCwWq2jGjxDEARBEAThy0DX9XgQMjc3F1kWOQaJtLS08PjjjwNw6aWXJjXP5s2bAbDZbGRmZu5xmqKiol2mTWTwWtfn8414rZuenp5UeYIgCIIgCIcrn8+3x88Hr3XT09MPuWtdEezdBx0dHTidzrGuhiAIgiAIwmGhtbWV/Pz8sa7GIS0ajXLVVVfR29vL+PHj+c53vpPUfP39/cDIgdeMjAwA+vr6kipTXOsKgiAIgiAk71C71j10ws6Hkb1F9QVBEARBEITdiWunxG666Sbef/99cnJyeOmllzCZTGNWF7G9BEEQBEEQkneoXTuJzN59YLVa4683bdp0QKP3Pp8vnlnR2tp6QB+XOxKXdSSu08Fc1pG4TgdzWUfiOh3MZR2J63SkLutIXKeDuawjcZ0A2traKC8vB3a9dhJ2d8stt/Doo49it9uZNWsW1dXVSc872HXDSD8yvF4vAFlZWUmVOXx7rVy5kry8PABW960hrEZwWZzkmnPIzszebd7R7GObva00epspsuZSZM1FD6n7tH8OX+ZDC2dSk1dAWVYOOeb0hN2tjVRfTdd5ZeMqHBYr5Vl2CjKzMQx7RHNfjqc5mzfT2dvL1SecMOr13NMyt3v9rNzRSmmOnZIcGznpe+92I5n69vgCvL1kHYU52ZTk2XA7sjAqyj6fO3r7+inwxPqfnrtoBdWlHjLSzCPOE4lAcTHEktY7gDJaW1viy5zz2Qb6+oMUOG0UOLPJs2ciy7uu82jqu3ZjCyvW7MDjysaWaeTYqQ2jXs8vLvP+B16nqMiJ22XD5crGbhu567+R6quqGjNmLCbNYsLtycblyiYvLwuDYd+3y5wPV9Pc3MmN3z4HgKambeTm2vdpXVtbW9mxrY/Pl27B5bbh9thwuW1kZFpGtZ7DdXf7mPnW5zhd2bjcdtweG4pBY9q0aRgMBp555hnq6+uTrm9fbz+egth+OPuDxZRXFGB37N/5YcHcdfT1+vEU2nF77OTmZe2yH45222xc38KaldtxF9ixO8xMnFyb9Lx7WuaG9Zv5ZE4jLo8NT6EDV4Eds8W4T+sKoEZV3p2xhLR0M55CB+4iB1k26z6t66D3X/uUC688DYDPP1tNSUUBBqMy6nUdXOa6ZVtYPHsN+QX22L9CB7lu+25lJlvfjuYuXvrb26hRLf6ZbILb7vsuAL+94W8Y5GFtKsF5N5xGcV3Bbo/zJ7tMTdPobO1l5+Z2dja1UzW+GHe5Y5+v3YYv960ZC8jKzGDSsRWjmm9Py/T5QmzZ0klnZz8nnFC9y7G0L/tDVNXY0dLNmg3buez86aOad0/LtFqtdPT62Lizk4KcLIqdez+/JVNfn8+HNxii4423qPrFzzB1dQLgBVwD0xxy17q6MGqtra06oAN6a2vrAV2W1+uNL8vr9YplHaLLOVKXdSSu08Fc1pG4TgdzWUfiOh2pyzoS1+lgLutIXCddP7jXS4ezH/7whzqg22w2fdGiRaOef/ny5fF27uvr2+M0t956qw7ol112WVJl7s+22599bF/nHYtlfpnquz/z7ut8l1yi6zD475hd5tU0LaXLHV7eodi+mqbtdZ33dZmapqW0vslsk9HUd0/leb1evby8XK+urtZXrlw5qvr29/enfP9NtM6H0/E2VvPu63bZ2zKjkWhK67t9Y4u+fskmfUdji97b0adHI9Fd5t2+aYe+s7FVb1y+RV/18Xr9s1nL9Q1LN+/XMg+VeUV9vyAS0fWf/UzXJWnoj5Pbrbe//PIhe60rMnsFQRAEQRAEYQz85Cc/4a9//SvZ2dnMnDmTqVOnjrqMmpoarFYrfr+fzz77jFNPPXW3aT777DMApkyZst91FoSD4Zxz4P/+b/Ddebt8l+rBsQ/1wbYPRP0O9TY81Ms7UGV+2aS6DRVDclnBySqo2EPf9aGhl7b8bDFQ6ZfBtm1w5ZUwf/7QZ2edBU8+ibb3ucacCPbug+EH9JF0cKenp6Pr+hG1rCNxnQ6mI7X9jsT972A6EtvvSF3WwXKktt+RuK8fTEfq9VKq3H777dx9991kZ2cza9Ysjj766H0qx2Qycd555/Hiiy/yzDPP7Bbs3bJlCx999BEAF1988X7XO5H92Z/3dd6xWOb+GKv6Hk7tdM45Q6+nTPkloz2FiH3p0Jx3rP7eHW77w5dlXzrc2mh/liv2pQM77wFb5uuvw7XXQldX7L2iwO9+Bz/+McgytLXt0zIPBjFAmyAIgiAIgiAcRD//+c/54x//iM1mSzrQe//991NbW8vVV1+923e33347kiTx2GOP8c4778Q/9/v93HDDDaiqyqWXXkptbW1K10MQDpSCApg8OfZ6yRLYuXNs6yMIgiB8iYTDcOutcMEFQ4He4mKYNw9+8pNYoPcQJzJ7BUEQBEEQBOEgee211/jtb38LQGVlJf/4xz/2OF1ubi5//vOf4+87OjpYt24dLpdrt2mnTJnCX/7yF374wx9y7rnncvLJJ5Ofn8+8efNobm6mpqaGhx566MCskCAcIOefD0uXxl6/8QZ8+9tjWx9BEAThS6CxEa64Aga6wALg4ovh0UfBnvxAlmNNBHsFQRAEQRAE4SDpGswQIdaX7mfDf0wMU1JSskuwN5Fbb72V8ePH85e//IVPP/0Un89HcXExd9xxB3fccQeZmZn7XXdBOJguuADuuiv2+rXXRLBXgN7eXmRZJisra6yrIgjCkej55+Fb34L+/th7kwn+8hf4/vfhMOunW9KPxE7iDjCfz0dGRgYAXq9X9EMnCIIgCILwBeJ66fAltp1wKNA0KCqKdeFgNkN7O4h7FoIgCELKBQLwgx/Aww8PfVZZCS+8MNSn0B4cytdLh35HE4IgCIIgCIIgjAmfzxf/JwgHkyzDhRfGXodC8PbbY1sfQRAE4Qi0Zg0ce+yugd5vfCPWYfxeAr2Hw7WRCPYKgiAIgiAIgrBHTqeTjIyMeOaKIBxMl1469Pqll8auHoIgCMIRRtfh3/+GqVNhxYrYZ2lpsc/+858RHyUZvC5yOp0HqbKjJ/rs3U/DI/mHUsq2IAiCIAjCWBi8NjqUsx0EQTg8nHwy5ORAZye8+Sb4/WC1jnWtBEEQhMNafz9897vw9NNDn40bF+uzt75+7OqVQiKzdz+JbAdBEARBEIQhh0O2g5C81tZWvF4vXq+Xvkgfmq6lrOz+SICQGklZeWE1Sl84mLLyADqD/pSW1x0IkMohU7yhMBFVTVl5EVXFH0rdNgHwBkL7PK/BAJdcEnvt9w8EfIPhlLZhKBRB01JXXlTViEZTt00AwuFoassLpb68VG4TNartcZ0bGxtZt24djY2Noy4zEAinompxfn8opescCIRRo6k7v0ajKsFgao9lnze151e/L7VtmOr9cKypUZWAL3VtHo2odLT2pqw8VdXYvr0rZW2u6zrNbb34U3isdvb5Line truncated
</div>
</div>
</div>
@@ -1824,7 +1987,7 @@ spikes = snn.spikes_until_t(in_spikes, max_time)</div>
<div class="jp-InputPrompt jp-InputArea-prompt">In [ ]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="CodeMirror cm-s-jupyter">
<div class="zeroclipboard-container">
<clipboard-copy ,="" aria-label="Copy to Clipboard" for="cell-6">
<clipboard-copy ,="" aria-label="Copy to Clipboard" for="cell-7">
<div>
<span class="notice" hidden="">Copied!</span>
<svg aria-hidden="true" class="clipboard-copy-icon" data-view-component="true" height="20" version="1.1" viewbox="0 0 16 16" width="20">
@@ -1836,7 +1999,34 @@ spikes = snn.spikes_until_t(in_spikes, max_time)</div>
</div>
<div class="highlight-ipynb hl-python"><pre><span></span>
</pre></div>
<div class="clipboard-copy-txt" id="cell-6"></div>
<div class="clipboard-copy-txt" id="cell-7"></div>
</div>
</div>
</div>
</div>
</div>
</div><div class="jp-Cell jp-CodeCell jp-Notebook-cell jp-mod-noOutputs">
<div class="jp-Cell jp-CodeCell jp-Notebook-cell jp-mod-noOutputs" id="cell-id=daf213b5">
<div class="jp-Cell-inputWrapper" tabindex="0">
<div class="jp-Collapser jp-InputCollapser jp-Cell-inputCollapser">
</div>
<div class="jp-InputArea jp-Cell-inputArea">
<div class="jp-InputPrompt jp-InputArea-prompt">In [ ]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="CodeMirror cm-s-jupyter">
<div class="zeroclipboard-container">
<clipboard-copy ,="" aria-label="Copy to Clipboard" for="cell-8">
<div>
<span class="notice" hidden="">Copied!</span>
<svg aria-hidden="true" class="clipboard-copy-icon" data-view-component="true" height="20" version="1.1" viewbox="0 0 16 16" width="20">
<path d="M0 6.75C0 5.784.784 5 1.75 5h1.5a.75.75 0 010 1.5h-1.5a.25.25 0 00-.25.25v7.5c0 .138.112.25.25.25h7.5a.25.25 0 00.25-.25v-1.5a.75.75 0 011.5 0v1.5A1.75 1.75 0 019.25 16h-7.5A1.75 1.75 0 010 14.25v-7.5z" fill="currentColor" fill-rule="evenodd"></path>
<path d="M5 1.75C5 .784 5.784 0 6.75 0h7.5C15.216 0 16 .784 16 1.75v7.5A1.75 1.75 0 0114.25 11h-7.5A1.75 1.75 0 015 9.25v-7.5zm1.75-.25a.25.25 0 00-.25.25v7.5c0 .138.112.25.25.25h7.5a.25.25 0 00.25-.25v-7.5a.25.25 0 00-.25-.25h-7.5z" fill="currentColor" fill-rule="evenodd"></path>
</svg>
</div>
</clipboard-copy>
</div>
<div class="highlight-ipynb hl-python"><pre><span></span>
</pre></div>
<div class="clipboard-copy-txt" id="cell-8"></div>
</div>
</div>
</div>
neuron_models/index.html
+5
View File
@@ -860,6 +860,11 @@
<td>Exponential Integrate-and-Fire neuron model</td>
<td>...</td>
</tr>
<tr>
<td><a href="lif/">LIF</a></td>
<td>Leaky Integrate-and-Fire neuron model</td>
<td>...</td>
</tr>
</tbody>
</table>
neuron_models/lif/index.html
+890
View File
@@ -0,0 +1,890 @@
<!doctype html>
<html lang="en" class="no-js">
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width,initial-scale=1">
<link rel="icon" href="../../img/felice.png">
<meta name="generator" content="mkdocs-1.6.1, mkdocs-material-9.7.1">
<title>LIF - Felice models</title>
<link rel="stylesheet" href="../../assets/stylesheets/main.484c7ddc.min.css">
<link rel="stylesheet" href="../../assets/stylesheets/palette.ab4e12ef.min.css">
<link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
<link rel="stylesheet" href="https://fonts.googleapis.com/css?family=Roboto:300,300i,400,400i,700,700i%7CRoboto+Mono:400,400i,700,700i&display=fallback">
<style>:root{--md-text-font:"Roboto";--md-code-font:"Roboto Mono"}</style>
<link rel="stylesheet" href="../../css/print-site.css">
<link rel="stylesheet" href="../../css/print-site-material.css">
<link rel="stylesheet" href="../../assets/_mkdocstrings.css">
<link rel="stylesheet" href="../../stylesheets/extra.css">
<script>__md_scope=new URL("../..",location),__md_hash=e=>[...e].reduce(((e,_)=>(e<<5)-e+_.charCodeAt(0)),0),__md_get=(e,_=localStorage,t=__md_scope)=>JSON.parse(_.getItem(t.pathname+"."+e)),__md_set=(e,_,t=localStorage,a=__md_scope)=>{try{t.setItem(a.pathname+"."+e,JSON.stringify(_))}catch(e){}}</script>
</head>
<body dir="ltr" data-md-color-scheme="default" data-md-color-primary="teal" data-md-color-accent="indigo">
<input class="md-toggle" data-md-toggle="drawer" type="checkbox" id="__drawer" autocomplete="off">
<input class="md-toggle" data-md-toggle="search" type="checkbox" id="__search" autocomplete="off">
<label class="md-overlay" for="__drawer"></label>
<div data-md-component="skip">
<a href="#lif" class="md-skip">
Skip to content
</a>
</div>
<div data-md-component="announce">
</div>
<header class="md-header md-header--shadow" data-md-component="header">
<nav class="md-header__inner md-grid" aria-label="Header">
<a href="../.." title="Felice models" class="md-header__button md-logo" aria-label="Felice models" data-md-component="logo">
<img src="../../img/felice.png" alt="logo">
</a>
<label class="md-header__button md-icon" for="__drawer">
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24"><path d="M3 6h18v2H3zm0 5h18v2H3zm0 5h18v2H3z"/></svg>
</label>
<div class="md-header__title" data-md-component="header-title">
<div class="md-header__ellipsis">
<div class="md-header__topic">
<span class="md-ellipsis">
Felice models
</span>
</div>
<div class="md-header__topic" data-md-component="header-topic">
<span class="md-ellipsis">
LIF
</span>
</div>
</div>
</div>
<label class="md-header__button md-icon" for="__search">
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24"><path d="M9.5 3A6.5 6.5 0 0 1 16 9.5c0 1.61-.59 3.09-1.56 4.23l.27.27h.79l5 5-1.5 1.5-5-5v-.79l-.27-.27A6.52 6.52 0 0 1 9.5 16 6.5 6.5 0 0 1 3 9.5 6.5 6.5 0 0 1 9.5 3m0 2C7 5 5 7 5 9.5S7 14 9.5 14 14 12 14 9.5 12 5 9.5 5"/></svg>
</label>
<div class="md-search" data-md-component="search" role="dialog">
<label class="md-search__overlay" for="__search"></label>
<div class="md-search__inner" role="search">
<form class="md-search__form" name="search">
<input type="text" class="md-search__input" name="query" aria-label="Search" placeholder="Search" autocapitalize="off" autocorrect="off" autocomplete="off" spellcheck="false" data-md-component="search-query" required>
<label class="md-search__icon md-icon" for="__search">
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24"><path d="M9.5 3A6.5 6.5 0 0 1 16 9.5c0 1.61-.59 3.09-1.56 4.23l.27.27h.79l5 5-1.5 1.5-5-5v-.79l-.27-.27A6.52 6.52 0 0 1 9.5 16 6.5 6.5 0 0 1 3 9.5 6.5 6.5 0 0 1 9.5 3m0 2C7 5 5 7 5 9.5S7 14 9.5 14 14 12 14 9.5 12 5 9.5 5"/></svg>
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24"><path d="M20 11v2H8l5.5 5.5-1.42 1.42L4.16 12l7.92-7.92L13.5 5.5 8 11z"/></svg>
</label>
<nav class="md-search__options" aria-label="Search">
<button type="reset" class="md-search__icon md-icon" title="Clear" aria-label="Clear" tabindex="-1">
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24"><path d="M19 6.41 17.59 5 12 10.59 6.41 5 5 6.41 10.59 12 5 17.59 6.41 19 12 13.41 17.59 19 19 17.59 13.41 12z"/></svg>
</button>
</nav>
</form>
<div class="md-search__output">
<div class="md-search__scrollwrap" tabindex="0" data-md-scrollfix>
<div class="md-search-result" data-md-component="search-result">
<div class="md-search-result__meta">
Initializing search
</div>
<ol class="md-search-result__list" role="presentation"></ol>
</div>
</div>
</div>
</div>
</div>
</nav>
</header>
<div class="md-container" data-md-component="container">
<main class="md-main" data-md-component="main">
<div class="md-main__inner md-grid">
<div class="md-sidebar md-sidebar--primary" data-md-component="sidebar" data-md-type="navigation" >
<div class="md-sidebar__scrollwrap">
<div class="md-sidebar__inner">
<nav class="md-nav md-nav--primary" aria-label="Navigation" data-md-level="0">
<label class="md-nav__title" for="__drawer">
<a href="../.." title="Felice models" class="md-nav__button md-logo" aria-label="Felice models" data-md-component="logo">
<img src="../../img/felice.png" alt="logo">
</a>
Felice models
</label>
<ul class="md-nav__list" data-md-scrollfix>
<li class="md-nav__item">
<a href="../.." class="md-nav__link">
<span class="md-ellipsis">
Home
</span>
</a>
</li>
<li class="md-nav__item md-nav__item--nested">
<input class="md-nav__toggle md-toggle " type="checkbox" id="__nav_2" >
<div class="md-nav__link md-nav__container">
<a href="../" class="md-nav__link ">
<span class="md-ellipsis">
Neuron Models
</span>
</a>
<label class="md-nav__link " for="__nav_2" id="__nav_2_label" tabindex="0">
<span class="md-nav__icon md-icon"></span>
</label>
</div>
<nav class="md-nav" data-md-level="1" aria-labelledby="__nav_2_label" aria-expanded="false">
<label class="md-nav__title" for="__nav_2">
<span class="md-nav__icon md-icon"></span>
Neuron Models
</label>
<ul class="md-nav__list" data-md-scrollfix>
<li class="md-nav__item md-nav__item--nested">
<input class="md-nav__toggle md-toggle " type="checkbox" id="__nav_2_2" >
<div class="md-nav__link md-nav__container">
<a href="../wererabbit/" class="md-nav__link ">
<span class="md-ellipsis">
WereRabbit
</span>
</a>
<label class="md-nav__link " for="__nav_2_2" id="__nav_2_2_label" tabindex="0">
<span class="md-nav__icon md-icon"></span>
</label>
</div>
<nav class="md-nav" data-md-level="2" aria-labelledby="__nav_2_2_label" aria-expanded="false">
<label class="md-nav__title" for="__nav_2_2">
<span class="md-nav__icon md-icon"></span>
WereRabbit
</label>
<ul class="md-nav__list" data-md-scrollfix>
<li class="md-nav__item">
<a href="../wererabbit/wererabbit/" class="md-nav__link">
<span class="md-ellipsis">
Basic example
</span>
</a>
</li>
</ul>
</nav>
</li>
<li class="md-nav__item md-nav__item--nested">
<input class="md-nav__toggle md-toggle " type="checkbox" id="__nav_2_3" >
<div class="md-nav__link md-nav__container">
<a href="../fhn/" class="md-nav__link ">
<span class="md-ellipsis">
FitzHugh-Nagumo
</span>
</a>
<label class="md-nav__link " for="__nav_2_3" id="__nav_2_3_label" tabindex="0">
<span class="md-nav__icon md-icon"></span>
</label>
</div>
<nav class="md-nav" data-md-level="2" aria-labelledby="__nav_2_3_label" aria-expanded="false">
<label class="md-nav__title" for="__nav_2_3">
<span class="md-nav__icon md-icon"></span>
FitzHugh-Nagumo
</label>
<ul class="md-nav__list" data-md-scrollfix>
<li class="md-nav__item">
<a href="../fhn/fhn/" class="md-nav__link">
<span class="md-ellipsis">
Example
</span>
</a>
</li>
</ul>
</nav>
</li>
<li class="md-nav__item md-nav__item--nested">
<input class="md-nav__toggle md-toggle " type="checkbox" id="__nav_2_4" >
<div class="md-nav__link md-nav__container">
<a href="../snowball/" class="md-nav__link ">
<span class="md-ellipsis">
Snowball
</span>
</a>
</div>
<nav class="md-nav" data-md-level="2" aria-labelledby="__nav_2_4_label" aria-expanded="false">
<label class="md-nav__title" for="__nav_2_4">
<span class="md-nav__icon md-icon"></span>
Snowball
</label>
<ul class="md-nav__list" data-md-scrollfix>
</ul>
</nav>
</li>
</ul>
</nav>
</li>
<li class="md-nav__item md-nav__item--nested">
<input class="md-nav__toggle md-toggle " type="checkbox" id="__nav_3" >
<div class="md-nav__link md-nav__container">
<a href="../../api/" class="md-nav__link ">
<span class="md-ellipsis">
API Reference
</span>
</a>
<label class="md-nav__link " for="__nav_3" id="__nav_3_label" tabindex="0">
<span class="md-nav__icon md-icon"></span>
</label>
</div>
<nav class="md-nav" data-md-level="1" aria-labelledby="__nav_3_label" aria-expanded="false">
<label class="md-nav__title" for="__nav_3">
<span class="md-nav__icon md-icon"></span>
API Reference
</label>
<ul class="md-nav__list" data-md-scrollfix>
<li class="md-nav__item">
<a href="../../api/neuron_models/" class="md-nav__link">
<span class="md-ellipsis">
Neuron Models
</span>
</a>
</li>
<li class="md-nav__item">
<a href="../../api/solver/" class="md-nav__link">
<span class="md-ellipsis">
Solver
</span>
</a>
</li>
<li class="md-nav__item">
<a href="../../api/datasets/" class="md-nav__link">
<span class="md-ellipsis">
Datasets
</span>
</a>
</li>
</ul>
</nav>
</li>
</ul>
</nav>
</div>
</div>
</div>
<div class="md-sidebar md-sidebar--secondary" data-md-component="sidebar" data-md-type="toc" >
<div class="md-sidebar__scrollwrap">
<div class="md-sidebar__inner">
<nav class="md-nav md-nav--secondary" aria-label="Table of contents">
<label class="md-nav__title" for="__toc">
<span class="md-nav__icon md-icon"></span>
Table of contents
</label>
<ul class="md-nav__list" data-md-component="toc" data-md-scrollfix>
<li class="md-nav__item">
<a href="#circuit-design" class="md-nav__link">
<span class="md-ellipsis">
Circuit Design
</span>
</a>
</li>
<li class="md-nav__item">
<a href="#circuit-simulation" class="md-nav__link">
<span class="md-ellipsis">
Circuit Simulation
</span>
</a>
</li>
<li class="md-nav__item">
<a href="#referennces" class="md-nav__link">
<span class="md-ellipsis">
Referennces
</span>
</a>
</li>
</ul>
</nav>
</div>
</div>
</div>
<div class="md-content" data-md-component="content">
<article class="md-content__inner md-typeset">
<h1 id="lif">LIF</h1>
<h2 id="circuit-design">Circuit Design</h2>
<p>W/L = 4/3</p>
<h2 id="circuit-simulation">Circuit Simulation</h2>
<p><img alt="lif, output plot" src="/docs-site/docs/img/lif_plot.png" /> Fig.1 The dynamics of leaky integrate and fire neuron. The grey signal is the input spikes, the yellow signal is the membrane potential and the dark blue is the output spikes from the neuron.</p>
<h2 id="referennces">Referennces</h2>
<ol>
<li>Sourikopoulos I, Hedayat S, Loyez C, Danneville F, Hoel V, Mercier E and Cappy A (2017) A 4-fJ/Spike Artificial Neuron in 65 nm CMOS Technology. Front. Neurosci. 11:123. doi: 10.3389/fnins.2017.00123</li>
</ol>
</article>
</div>
<script>var target=document.getElementById(location.hash.slice(1));target&&target.name&&(target.checked=target.name.startsWith("__tabbed_"))</script>
</div>
</main>
<footer class="md-footer">
<div class="md-footer-meta md-typeset">
<div class="md-footer-meta__inner md-grid">
<div class="md-copyright">
<div class="md-copyright__highlight">
Felice models © 2026 by University of Groningen (Fernando M. Quintana) | Code is licensed under <a href="https://opensource.org/licenses/MIT">MIT License</a> | Docs is licensed under <a href="https://creativecommons.org/licenses/by/4.0/">CC BY 4.0</a><img src="https://mirrors.creativecommons.org/presskit/icons/cc.svg" alt="" style="max-width: 1em;max-height:1em;margin-left: .2em;"><img src="https://mirrors.creativecommons.org/presskit/icons/by.svg" alt="" style="max-width: 1em;max-height:1em;margin-left: .2em;">
</div>
Made with
<a href="https://squidfunk.github.io/mkdocs-material/" target="_blank" rel="noopener">
Material for MkDocs
</a>
</div>
</div>
</div>
</footer>
</div>
<div class="md-dialog" data-md-component="dialog">
<div class="md-dialog__inner md-typeset"></div>
</div>
<script id="__config" type="application/json">{"annotate": null, "base": "../..", "features": ["navigation.path", "navigation.indexes", "content.code.copy", "content.code.select", "content.code.annotate"], "search": "../../assets/javascripts/workers/search.2c215733.min.js", "tags": null, "translations": {"clipboard.copied": "Copied to clipboard", "clipboard.copy": "Copy to clipboard", "search.result.more.one": "1 more on this page", "search.result.more.other": "# more on this page", "search.result.none": "No matching documents", "search.result.one": "1 matching document", "search.result.other": "# matching documents", "search.result.placeholder": "Type to start searching", "search.result.term.missing": "Missing", "select.version": "Select version"}, "version": null}</script>
<script src="../../assets/javascripts/bundle.79ae519e.min.js"></script>
<script src="../../js/print-site.js"></script>
<script src="../../javascripts/mathjax.js"></script>
<script src="https://polyfill.io/v3/polyfill.min.js?features=es6"></script>
<script src="https://unpkg.com/mathjax@3/es5/tex-mml-chtml.js"></script>
</body>
</html>
neuron_models/wererabbit/wererabbit/index.html
+62 -141
View File
@@ -1564,7 +1564,7 @@ span.linenos.special { color: #000000; background-color: #ffffc0; padding-left:
<div class="jp-Collapser jp-InputCollapser jp-Cell-inputCollapser">
</div>
<div class="jp-InputArea jp-Cell-inputArea">
<div class="jp-InputPrompt jp-InputArea-prompt">In [1]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="jp-InputPrompt jp-InputArea-prompt">In [10]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="CodeMirror cm-s-jupyter">
<div class="zeroclipboard-container">
<clipboard-copy ,="" aria-label="Copy to Clipboard" for="cell-1">
@@ -1583,7 +1583,6 @@ span.linenos.special { color: #000000; background-color: #ffffc0; padding-left:
<span class="kn">import</span><span class="w"> </span><span class="nn">jax.random</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">jrand</span>
<span class="kn">import</span><span class="w"> </span><span class="nn">matplotlib</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">mpl</span>
<span class="kn">import</span><span class="w"> </span><span class="nn">matplotlib.pyplot</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">plt</span>
<span class="kn">from</span><span class="w"> </span><span class="nn">eventpropjax.evnn</span><span class="w"> </span><span class="kn">import</span> <span class="n">FFEvNN</span>
<span class="kn">from</span><span class="w"> </span><span class="nn">felice.neuron_models</span><span class="w"> </span><span class="kn">import</span> <span class="n">WereRabbit</span>
@@ -1595,7 +1594,6 @@ 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
@@ -1611,7 +1609,7 @@ jax.config.update("jax_enable_x64", True)</div>
<div class="jp-Collapser jp-InputCollapser jp-Cell-inputCollapser">
</div>
<div class="jp-InputArea jp-Cell-inputArea">
<div class="jp-InputPrompt jp-InputArea-prompt">In [2]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="jp-InputPrompt jp-InputArea-prompt">In [11]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="CodeMirror cm-s-jupyter">
<div class="zeroclipboard-container">
<clipboard-copy ,="" aria-label="Copy to Clipboard" for="cell-2">
@@ -1625,44 +1623,46 @@ jax.config.update("jax_enable_x64", True)</div>
</clipboard-copy>
</div>
<div class="highlight-ipynb hl-python"><pre><span></span><span class="n">key</span> <span class="o">=</span> <span class="n">jrand</span><span class="o">.</span><span class="n">key</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
<span class="n">max_time</span> <span class="o">=</span> <span class="mi">100</span>
<span class="n">max_time</span> <span class="o">=</span> <span class="mi">40</span>
<span class="n">init_weights</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">asarray</span><span class="p">([[</span><span class="mf">0.0</span><span class="p">],</span> <span class="p">[</span><span class="mf">0.2</span><span class="p">],</span> <span class="p">[</span><span class="o">-</span><span class="mf">0.1</span><span class="p">]])</span>
<span class="n">snn</span> <span class="o">=</span> <span class="n">FFEvNN</span><span class="p">(</span>
<span class="n">layers</span><span class="o">=</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
<span class="n">in_size</span><span class="o">=</span><span class="mi">2</span><span class="p">,</span>
<span class="n">neuron_model</span><span class="o">=</span><span class="n">WereRabbit</span><span class="p">,</span>
<span class="n">model</span> <span class="o">=</span> <span class="n">WereRabbit</span><span class="p">(</span><span class="n">dtype</span><span class="o">=</span><span class="n">jnp</span><span class="o">.</span><span class="n">float64</span><span class="p">)</span>
<span class="k">def</span><span class="w"> </span><span class="nf">state_at_t</span><span class="p">(</span><span class="n">comp_times</span><span class="p">):</span>
<span class="n">sol</span> <span class="o">=</span> <span class="n">dfx</span><span class="o">.</span><span class="n">diffeqsolve</span><span class="p">(</span>
<span class="n">terms</span><span class="o">=</span><span class="n">dfx</span><span class="o">.</span><span class="n">ODETerm</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">dynamics</span><span class="p">),</span>
<span class="n">solver</span><span class="o">=</span><span class="n">dfx</span><span class="o">.</span><span class="n">Tsit5</span><span class="p">(),</span>
<span class="c1"># stepsize_controller=PIDController(</span>
<span class="c1"># rtol=1e-6, atol=1e-3, pcoeff=0.1, icoeff=0.3, dcoeff=0.0</span>
<span class="c1"># ),</span>
<span class="n">max_solver_time</span><span class="o">=</span><span class="n">max_time</span><span class="p">,</span>
<span class="n">key</span><span class="o">=</span><span class="n">key</span><span class="p">,</span>
<span class="n">max_event_steps</span><span class="o">=</span><span class="mi">10000</span><span class="p">,</span>
<span class="n">solver_stepsize</span><span class="o">=</span><span class="mf">0.001</span><span class="p">,</span>
<span class="n">init_weights</span><span class="o">=</span><span class="n">init_weights</span><span class="p">,</span>
<span class="n">dtype</span><span class="o">=</span><span class="n">jnp</span><span class="o">.</span><span class="n">float64</span><span class="p">,</span>
<span class="p">)</span>
<span class="n">t0</span><span class="o">=</span><span class="mf">0.0</span><span class="p">,</span>
<span class="n">t1</span><span class="o">=</span><span class="n">max_time</span><span class="p">,</span>
<span class="n">dt0</span><span class="o">=</span><span class="mf">1e-3</span><span class="p">,</span>
<span class="n">y0</span><span class="o">=</span><span class="n">model</span><span class="o">.</span><span class="n">init_state</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>
<span class="o">+</span> <span class="n">jrand</span><span class="o">.</span><span class="n">uniform</span><span class="p">(</span><span class="n">key</span><span class="p">,</span> <span class="n">shape</span><span class="o">=</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">),</span> <span class="n">minval</span><span class="o">=</span><span class="mf">0.1</span><span class="p">,</span> <span class="n">maxval</span><span class="o">=</span><span class="mf">0.5</span><span class="p">),</span>
<span class="n">saveat</span><span class="o">=</span><span class="n">dfx</span><span class="o">.</span><span class="n">SaveAt</span><span class="p">(</span><span class="n">ts</span><span class="o">=</span><span class="n">comp_times</span><span class="p">),</span>
<span class="n">max_steps</span><span class="o">=</span><span class="mi">100000</span><span class="p">,</span>
<span class="p">)</span>
<span class="k">return</span> <span class="n">sol</span><span class="o">.</span><span class="n">ts</span><span class="p">,</span> <span class="n">sol</span><span class="o">.</span><span class="n">ys</span>
</pre></div>
<div class="clipboard-copy-txt" id="cell-2">key = jrand.key(0)
max_time = 100
max_time = 40
init_weights = jnp.asarray([[0.0], [0.2], [-0.1]])
snn = FFEvNN(
layers=[1],
in_size=2,
neuron_model=WereRabbit,
model = WereRabbit(dtype=jnp.float64)
def state_at_t(comp_times):
sol = dfx.diffeqsolve(
terms=dfx.ODETerm(model.dynamics),
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,
)</div>
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</div>
</div>
</div>
</div>
@@ -1674,7 +1674,7 @@ snn = FFEvNN(
<div class="jp-Collapser jp-InputCollapser jp-Cell-inputCollapser">
</div>
<div class="jp-InputArea jp-Cell-inputArea">
<div class="jp-InputPrompt jp-InputArea-prompt">In [3]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="jp-InputPrompt jp-InputArea-prompt">In [12]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="CodeMirror cm-s-jupyter">
<div class="zeroclipboard-container">
<clipboard-copy ,="" aria-label="Copy to Clipboard" for="cell-3">
@@ -1687,15 +1687,11 @@ snn = FFEvNN(
</div>
</clipboard-copy>
</div>
<div class="highlight-ipynb hl-python"><pre><span></span><span class="n">in_spikes</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">asarray</span><span class="p">([[</span><span class="mi">10</span><span class="p">,</span> <span class="mi">30</span><span class="p">],</span> <span class="p">[</span><span class="mi">20</span><span class="p">,</span> <span class="mi">21</span><span class="p">]])</span>
<span class="n">comp_times</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">linspace</span><span class="p">(</span><span class="mf">0.0</span><span class="p">,</span> <span class="n">max_time</span><span class="p">,</span> <span class="mi">2000</span><span class="p">)</span>
<span class="n">state</span> <span class="o">=</span> <span class="n">snn</span><span class="o">.</span><span class="n">state_at_t</span><span class="p">(</span><span class="n">in_spikes</span><span class="p">,</span> <span class="n">comp_times</span><span class="p">)</span>
<span class="n">spikes</span> <span class="o">=</span> <span class="n">snn</span><span class="o">.</span><span class="n">spikes_until_t</span><span class="p">(</span><span class="n">in_spikes</span><span class="p">,</span> <span class="n">max_time</span><span class="p">)</span>
<div class="highlight-ipynb hl-python"><pre><span></span><span class="n">comp_times</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">linspace</span><span class="p">(</span><span class="mf">0.0</span><span class="p">,</span> <span class="n">max_time</span><span class="p">,</span> <span class="mi">2000</span><span class="p">)</span>
<span class="n">_</span><span class="p">,</span> <span class="n">state</span> <span class="o">=</span> <span class="n">state_at_t</span><span class="p">(</span><span class="n">comp_times</span><span class="p">)</span>
</pre></div>
<div class="clipboard-copy-txt" id="cell-3">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)</div>
<div class="clipboard-copy-txt" id="cell-3">comp_times = jnp.linspace(0.0, max_time, 2000)
_, state = state_at_t(comp_times)</div>
</div>
</div>
</div>
@@ -1707,7 +1703,7 @@ spikes = snn.spikes_until_t(in_spikes, max_time)</div>
<div class="jp-Collapser jp-InputCollapser jp-Cell-inputCollapser">
</div>
<div class="jp-InputArea jp-Cell-inputArea">
<div class="jp-InputPrompt jp-InputArea-prompt">In [4]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="jp-InputPrompt jp-InputArea-prompt">In [13]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="CodeMirror cm-s-jupyter">
<div class="zeroclipboard-container">
<clipboard-copy ,="" aria-label="Copy to Clipboard" for="cell-4">
@@ -1733,7 +1729,7 @@ spikes = snn.spikes_until_t(in_spikes, max_time)</div>
<span class="n">UV</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">stack</span><span class="p">(</span>
<span class="p">[</span><span class="n">U</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">),</span> <span class="n">V</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">),</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ones</span><span class="p">((</span><span class="n">resolution</span> <span class="o">*</span> <span class="n">resolution</span><span class="p">,))],</span> <span class="n">axis</span><span class="o">=</span><span class="mi">1</span>
<span class="p">)</span>
<span class="n">dS</span> <span class="o">=</span> <span class="n">snn</span><span class="o">.</span><span class="n">neuron_model</span><span class="o">.</span><span class="n">vector_field</span><span class="p">(</span><span class="n">UV</span><span class="p">)</span>
<span class="n">dS</span> <span class="o">=</span> <span class="n">snn</span><span class="o">.</span><span class="n">vector_field</span><span class="p">(</span><span class="n">UV</span><span class="p">)</span>
<span class="n">dU</span> <span class="o">=</span> <span class="n">dS</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">U</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
<span class="n">dV</span> <span class="o">=</span> <span class="n">dS</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">V</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
@@ -1751,7 +1747,7 @@ spikes = snn.spikes_until_t(in_spikes, max_time)</div>
<span class="n">U</span><span class="p">,</span> <span class="n">V</span><span class="p">,</span> <span class="n">dU</span><span class="p">,</span> <span class="n">dV</span> <span class="o">=</span> <span class="n">compute_nullclines</span><span class="p">(</span><span class="n">snn</span><span class="p">,</span> <span class="n">u_range</span><span class="p">,</span> <span class="n">v_range</span><span class="p">,</span> <span class="mi">200</span><span class="p">)</span>
<span class="n">UVs</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">stack</span><span class="p">([</span><span class="n">Us</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">),</span> <span class="n">Vs</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">),</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ones</span><span class="p">((</span><span class="mi">20</span> <span class="o">*</span> <span class="mi">20</span><span class="p">,))],</span> <span class="n">axis</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
<span class="n">dS</span> <span class="o">=</span> <span class="n">snn</span><span class="o">.</span><span class="n">neuron_model</span><span class="o">.</span><span class="n">vector_field</span><span class="p">(</span><span class="n">UVs</span><span class="p">)</span>
<span class="n">dS</span> <span class="o">=</span> <span class="n">snn</span><span class="o">.</span><span class="n">vector_field</span><span class="p">(</span><span class="n">UVs</span><span class="p">)</span>
<span class="n">dUs</span> <span class="o">=</span> <span class="n">dS</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">Us</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
<span class="n">dVs</span> <span class="o">=</span> <span class="n">dS</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">Vs</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
@@ -1790,7 +1786,7 @@ spikes = snn.spikes_until_t(in_spikes, max_time)</div>
UV = jnp.stack(
[U.reshape(-1), V.reshape(-1), jnp.ones((resolution * resolution,))], axis=1
)
dS = snn.neuron_model.vector_field(UV)
dS = snn.vector_field(UV)
dU = dS[:, 0].reshape(U.shape)
dV = dS[:, 1].reshape(V.shape)
@@ -1808,7 +1804,7 @@ def plot_vf(ax, snn, u_range, v_range):
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)
dS = snn.vector_field(UVs)
dUs = dS[:, 0].reshape(Us.shape)
dVs = dS[:, 1].reshape(Vs.shape)
@@ -1844,7 +1840,7 @@ def plot_vf(ax, snn, u_range, v_range):
<div class="jp-Collapser jp-InputCollapser jp-Cell-inputCollapser">
</div>
<div class="jp-InputArea jp-Cell-inputArea">
<div class="jp-InputPrompt jp-InputArea-prompt">In [5]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="jp-InputPrompt jp-InputArea-prompt">In [14]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="CodeMirror cm-s-jupyter">
<div class="zeroclipboard-container">
<clipboard-copy ,="" aria-label="Copy to Clipboard" for="cell-5">
@@ -1859,29 +1855,27 @@ def plot_vf(ax, snn, u_range, v_range):
</div>
<div class="highlight-ipynb hl-python"><pre><span></span><span class="k">with</span> <span class="n">mpl</span><span class="o">.</span><span class="n">style</span><span class="o">.</span><span class="n">context</span><span class="p">(</span><span class="s2">"boilerplot.ieeetran"</span><span class="p">):</span>
<span class="n">fig</span><span class="p">,</span> <span class="n">ax</span> <span class="o">=</span> <span class="n">plt</span><span class="o">.</span><span class="n">subplots</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="n">figsize</span><span class="o">=</span><span class="p">(</span><span class="mf">6.9</span><span class="p">,</span> <span class="mf">2.6</span><span class="p">),</span> <span class="n">dpi</span><span class="o">=</span><span class="mi">200</span><span class="p">)</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">comp_times</span><span class="p">,</span> <span class="n">state</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="p">:,</span> <span class="mi">0</span><span class="p">],</span> <span class="n">label</span><span class="o">=</span><span class="s2">"x1"</span><span class="p">)</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">comp_times</span><span class="p">,</span> <span class="n">state</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="p">:,</span> <span class="mi">1</span><span class="p">],</span> <span class="n">label</span><span class="o">=</span><span class="s2">"x2"</span><span class="p">)</span>
<span class="p">[</span><span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">axvline</span><span class="p">(</span><span class="n">s</span><span class="p">,</span> <span class="n">alpha</span><span class="o">=</span><span class="mf">0.2</span><span class="p">,</span> <span class="n">color</span><span class="o">=</span><span class="s2">"g"</span><span class="p">,</span> <span class="n">linestyle</span><span class="o">=</span><span class="s2">"--"</span><span class="p">)</span> <span class="k">for</span> <span class="n">s</span> <span class="ow">in</span> <span class="n">jnp</span><span class="o">.</span><span class="n">unique</span><span class="p">(</span><span class="n">spikes</span><span class="p">)]</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">legend</span><span class="p">([</span><span class="s2">"x1"</span><span class="p">,</span> <span class="s2">"x2"</span><span class="p">,</span> <span class="s2">"Spike"</span><span class="p">])</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">comp_times</span><span class="p">,</span> <span class="n">state</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span> <span class="n">label</span><span class="o">=</span><span class="s2">"x1"</span><span class="p">)</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">comp_times</span><span class="p">,</span> <span class="n">state</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">],</span> <span class="n">label</span><span class="o">=</span><span class="s2">"x2"</span><span class="p">)</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">legend</span><span class="p">([</span><span class="s2">"x1"</span><span class="p">,</span> <span class="s2">"x2"</span><span class="p">])</span>
<span class="n">plot_vf</span><span class="p">(</span><span class="n">ax</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">snn</span><span class="p">,</span> <span class="p">[</span><span class="o">-</span><span class="mf">0.2</span><span class="p">,</span> <span class="mf">0.5</span><span class="p">],</span> <span class="p">[</span><span class="o">-</span><span class="mf">0.2</span><span class="p">,</span> <span class="mf">0.5</span><span class="p">])</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">state</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="p">:,</span> <span class="mi">0</span><span class="p">],</span> <span class="n">state</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="p">:,</span> <span class="mi">1</span><span class="p">])</span>
<span class="n">plot_vf</span><span class="p">(</span><span class="n">ax</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">model</span><span class="p">,</span> <span class="p">[</span><span class="o">-</span><span class="mf">0.2</span><span class="p">,</span> <span class="mf">0.5</span><span class="p">],</span> <span class="p">[</span><span class="o">-</span><span class="mf">0.2</span><span class="p">,</span> <span class="mf">0.5</span><span class="p">])</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">state</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span> <span class="n">state</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">])</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">state</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span> <span class="n">state</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">],</span> <span class="s2">"."</span><span class="p">,</span> <span class="n">label</span><span class="o">=</span><span class="s2">"start"</span><span class="p">)</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">state</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span> <span class="n">state</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">],</span> <span class="s2">"."</span><span class="p">,</span> <span class="n">label</span><span class="o">=</span><span class="s2">"end"</span><span class="p">)</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">state</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span> <span class="n">state</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">],</span> <span class="s2">"."</span><span class="p">,</span> <span class="n">label</span><span class="o">=</span><span class="s2">"end"</span><span class="p">)</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">legend</span><span class="p">()</span>
<span class="n">plt</span><span class="o">.</span><span class="n">show</span><span class="p">()</span>
</pre></div>
<div class="clipboard-copy-txt" id="cell-5">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"])
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], snn, [-0.2, 0.5], [-0.2, 0.5])
ax[1].plot(state[0, :, 0], state[0, :, 1])
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[0, -1, 0], state[0, -1, 1], ".", label="end")
ax[1].plot(state[-1, 0, 0], state[-1, 0, 1], ".", label="end")
ax[1].legend()
plt.show()</div>
</div>
@@ -1895,80 +1889,7 @@ def plot_vf(ax, snn, u_range, v_range):
<div class="jp-OutputArea-child">
<div class="jp-OutputPrompt jp-OutputArea-prompt"></div>
<div class="jp-RenderedImage jp-OutputArea-output" tabindex="0">
<img alt="No description has been provided for this image" class="" src="data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAABXsAAAIfCAYAAAAhag9YAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjgsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvwVt1zgAAAAlwSFlzAAAewgAAHsIBbtB1PgABAABJREFUeJzs3Xd4W/XZxvHv0bQlee9tx1nOJINASAIhAcre42VDB4WyS1sooxTaUrqBAoUuoAXKDjNAmJmQvfdwYifxdjwkW/u8f8ialke8kzyf6/Jl6egsSUe2dOs5z09RVVVFCCGEEEIIIYQQQgghxBFNM9g7IIQQQgghhBBCCCGEEKL3JOwVQgghhBBCCCGEEEKIo4CEvUIIIYQQQgghhBBCCHEUkLBXCCGEEEIIIYQQQgghjgIS9gohhBBCCCGEEEIIIcRRQMJeIYQQQgghhBBCCCGEOApI2CuEEEIIIYQQQgghhBBHAQl7hRBCCCGEEEIIIYQQ4iggYa8QQgghhBBCCCGEEEIcBSTsFUIIIYQQQgghhBBCiKOAhL1CCCGEEEIIIYQQQghxFJCwVwghhBBCCCGEEEIIIY4CEvYKIYQQQgghhBBCCCHEUUDCXiGEEEIIIYQQQgghhDgKSNgrhBBCCCGEEEIIIYQQRwEJe4UQQgghhBBCCCGEEOIoIGGvEEIIIYQQQgghhBBCHAUk7BVCCCGEEEIIIYQQQoijgIS9QgghhBBCCCGEEEIIcRSQsFcIIYQQQgghjjFOpxOTyYSiKFx11VWdzvvkk0+iKAqKonDyySd3Ou/bb78dmPett97qy10+6rz44ouBx6qnfvnLX6IoCoWFhYO6jr7kf0wif0wmE0VFRVx22WV88MEHUZf9+uuvA/Pv3bt3YHf8CBD6+IT+aDQa4uPjGTt2LDfffDPr1q0b7F0dkuT4EkcKCXuFEEIIIYQQ4hhjMBg44YQTAFi8eHGn8y5atChwecWKFTgcjm7N21UwLIa+vXv3BsKtr7/+elD3pbW1lb179/LWW29x/vnnc8EFF3R6LIruU1WV5uZmtmzZwvPPP8+UKVP4zW9+M6j75D/uXnzxxUHdj8NRWFiIoij88pe/HOxdEcc4CXuFEEIIIYQQ4hjkD2P379/Pnj17OpxvyZIlABiNRhwOB8uXL+9wXn/YO3r0aNLT0/twb8Wx5uqrr6a5uTnwU1FRwdtvv82YMWMAeP/997n77rsHeS+PXM8991zgsW1qamLPnj08/fTTJCYm4vV6efDBB3njjTcGezeFED0gYa8QQgghhBBCHINCK29DK3JDbdu2jerqakwmE1dccUWn8zY2NrJhw4Z26xZD2y9/+UtUVR1yp6XrdDosFkvgJzMzk4svvphFixaRkpICwD/+8Q8qKysHeU+PTEajMfDYxsXFUVRUxK233sq8efMC8zz66KODuIdDz+zZs1FVFVVVh0zbEyGikbBXCCGEEEIIIY5B06dPR6/XAx0HuP4WDyeeeCKnnnpq2LRIS5cuxev1AhL2iv6TkpLC7bffDoDb7earr74a5D06usyePZvZs2cDsHnzZioqKgZ3h4QQh03CXiGEEEIIIYQ4BplMJqZOnQp0HPb6p8+aNYuZM2cCsGzZMjweT4fzApxyyintbrfb7Tz99NPMnTuX9PR0DAYD6enpnHXWWbzxxhuoqhp1HyIHMtu/fz933303o0aNwmw2Rx0sqbKykgceeIApU6aQlJSE0WgkPz+fa665hhUrVnT4mNxwww0oihIIu5YvX87VV19Nfn4+BoMhrJpPVVVWrFjBgw8+yEknnURKSgp6vZ6kpCSOP/54Hn74Yerq6jrcVqSysjJuvfVWhg0bRkxMDOnp6VxyySV888033V7Hpk2buP7668nLy8NoNJKTk8N1113H1q1bO1ymowHaCgsLKSoqClw/9dRT2w3sNVh9fMeOHRu4vH///g7ns9lsPProo4wbNw6z2UxCQgKzZ8/ucvDAjRs38thjjzF79mwyMjLQ6/XEx8czYcIE7rnnHsrLyztd3maz8fvf/56TTjqJ5ORk9Ho9aWlpjBkzhssuu4x//OMf2Gy2Dpf/+uuvufbaaykqKiI2Npb4+HiOO+44fvGLX3Do0KFOt90XOnt8/fdt+vTpJCcnB46zyy67jE8//bTDdUb2f7bb7fz+97/n+OOPJykpKdCfd/bs2WGDFt54443tjrvQPr4D+ZrtaIA2/z7s27cPgEceeaTdPksfXzGgVCGEEEIIIYQQx6R7771XBVRAPXDgQLvbCwoKVED9/PPPVVVV1czMTBVQV6xY0W7e6dOnq4BaVFTU7rZNmzapRUVFgW1F+znvvPNUm83WbtkXXnghMM/y5cvV5OTkdsuWlpYG5n/jjTdUs9nc6bYeffTRqI/H9ddfrwLqKaecoj733HOqVqsNW66goCAw77vvvtvpNgA1MzNTXbt2bdRthd6vVatWRb1fgKrRaNSnnnoq6joefvjhwH599NFHamxsbNR1GAwG9e233+5yHaH8z31nP1999VXYMv7pp5xyStRtdYd/Hddff32H87z99tuB+X73u98Fpn/11VeB6d98841aUlLS4b7/6le/irrudevWdXm/4+Li1AULFkRd/uDBg+qIESO6XMfKlSvbLWu329Vrr7220+UyMjKiLtsdoY/PCy+80OF8t99+e9hrzm/z5s1qfn5+p/t37bXXqk6ns906S0tLA/O8/fbb6oQJE9ot+8ILL6innHJKl49d6L4P5Gs29PEL/Zvj34fOfh5++OFuPktC9J5U9gohhBBCCCHEMaqzvr3l5eXs27cPnU7HiSeeCBCo7o2ct7W1lVWrVgHtq3r379/P7NmzKS0tpaCggOeff54dO3ZQX1/P1q1beeSRRzAajXzwwQfccsstne7vJZdcgslk4t///jdlZWVUVVXx8ccfk5SUBMD8+fO54oorsNlszJw5k3feeYfy8nLq6upYsWIF11xzDQC/+MUveOGFFzrczrZt27jtttuYNm0aH3/8MVVVVZSVlfHXv/41MI9Op+P888/n+eefZ8mSJezevZva2lo2bdrE888/z8iRI6msrOTiiy/Gbrd3er8uu+wydDod//znPykvL6eyspLXX3+doqIivF4vd9xxB59//nmHyzc2NnLNNdeQlZXFG2+8QWVlJeXl5fzjH/8gLS0Np9PJlVdeyZYtWzrdj1Bbtmxh8+bNgevz588PGzCtubmZWbNmdXt9fSn0fuTk5ESd5+qrr6ampoann36aPXv2UFtby+eff864ceMAX0VztIpnRVGYM2cOTzzxBAsXLmTHjh3U1dWxbds2XnnlFaZOnUpzczNXXHFF1H7BP/vZz9i5cydarZaHHnqIdevWUV1dTXV1NWvXruXZZ59tV73qd/311/Pf//4XnU7HnXfeyfLly6mtreXgwYO88cYbjB49mqqqKs4999x+7VUc7fGtr6/njDPOoKysDKPRyCOPPML27dupra1l0aJFfOc73wHgv//9Lz/+8Y87Xf+dd97Jjh07eOSRR9i6dSu1tbWsWLEi8Hprbm4OzBs6kJz/x/86DjXQr9lQzz//PM3NzeTn5wPw85//vN0+33///d1enxC9NthpsxBCCCGEEEKIwdHY2KhqNBoVUG+55Zaw215++WUVUKdNmxaY9sQTT6iAesEFF4TN++WXXwYq2P7973+H3XbBBReogDp27Fi1vr4+6n7Mnz8/rMo1VGgFbEpKilpWVhZ1Ha2trWpGRoYKqFdffbXq9Xqjzvezn/1MBdT09HS1tbU17LbQCr2ZM2eqDocj6jq6o7m5WS0uLo76mETeL5PJpG7ZsqXdPAcOHAjcp7Fjx7a73V+VC6hZWVlqRUVFu3k2bdqkxsTEqIB69tlnd7iOyMpeVQ2vxoys4o3GP29/VvY2NDQEHhONRqOWl5cHbgutvLRYLOrWrVvbLb9///5ABfS999572PvncrnUGTNmqID6i1/8ot3t/grtu++++7DWO2/evMB9+vjjj6PO09DQEKgavu222w5737tT2bts2TJVURQVUEeOHBmYftdddwWWfffdd9st5/F41AsvvDAwz8aNG8NuDz2WAPWjjz7qdF+72k+/gXzNdlTZ6+evhpcqXjHYpLJXCCGEEEIIIY5R/l6g0L5aN7Rfr5//8pIlS8J67IYuG1otXFpayvvvvw/AE088EajAjXTWWWcFem6+8sorHe7vT3/6U/Ly8qLe9tprr1FVVYXJZOKZZ56JWjkJ8PDDD2M2m6murmbBggUdbuuPf/wjBoOhw9u7YrFYuPjiiwH47LPPOp331ltvpaSkpN307OxsHnjgAcA3WNby5cs7XMeDDz5IZmZmu+ljx47lRz/6EQCffPJJvw64paoqqqr2Sy9fm83GZ599xty5c6mqqgLg2muvJTc3N+r8t99+O6NHj243PScnh9NPPx2AlStXHvZ+6HQ6rrrqKiD68+p2uwHfc3c4nnjiCQCuvPJKzjzzzKjzJCQkBI6HV199tcM+1z1RW1vLf/7zHy688MLAen/+858D4PF4An1yzznnHC644IJ2y2s0Gp5++ml0Oh0A//jHPzrc1plnnsnZZ5/dZ/vuN5CvWSGGMgl7hRBCCCGEEOIY5g9nt2zZEjY40eLFi4HwsHfixInExcVRV1cXdoq/P+zNycmhuLg4MP2LL75AVVViYmKYNm0aVqu1wx9/6NxZAHfOOed0eJu/zcH06dPRarUdbsfr9QZCwI62lZKSwgknnNDhtvzcbjcvvfQS559/Pvn5+ZhMprBBmf7whz8AsH379k7X4w+YornkkksCl5cuXdqrdXi93sMa8G0wvfTSS2GPpcVi4YwzzmD16tUAzJ49m2eeeabD5c8666wObxs1ahRAp60Q5s2bLine truncated
</div>
</div>
</div>
</div>
</div>
</div><div class="jp-Cell jp-CodeCell jp-Notebook-cell">
<div class="jp-Cell jp-CodeCell jp-Notebook-cell" id="cell-id=44e22685">
<div class="jp-Cell-inputWrapper" tabindex="0">
<div class="jp-Collapser jp-InputCollapser jp-Cell-inputCollapser">
</div>
<div class="jp-InputArea jp-Cell-inputArea">
<div class="jp-InputPrompt jp-InputArea-prompt">In [6]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="CodeMirror cm-s-jupyter">
<div class="zeroclipboard-container">
<clipboard-copy ,="" aria-label="Copy to Clipboard" for="cell-6">
<div>
<span class="notice" hidden="">Copied!</span>
<svg aria-hidden="true" class="clipboard-copy-icon" data-view-component="true" height="20" version="1.1" viewbox="0 0 16 16" width="20">
<path d="M0 6.75C0 5.784.784 5 1.75 5h1.5a.75.75 0 010 1.5h-1.5a.25.25 0 00-.25.25v7.5c0 .138.112.25.25.25h7.5a.25.25 0 00.25-.25v-1.5a.75.75 0 011.5 0v1.5A1.75 1.75 0 019.25 16h-7.5A1.75 1.75 0 010 14.25v-7.5z" fill="currentColor" fill-rule="evenodd"></path>
<path d="M5 1.75C5 .784 5.784 0 6.75 0h7.5C15.216 0 16 .784 16 1.75v7.5A1.75 1.75 0 0114.25 11h-7.5A1.75 1.75 0 015 9.25v-7.5zm1.75-.25a.25.25 0 00-.25.25v7.5c0 .138.112.25.25.25h7.5a.25.25 0 00.25-.25v-7.5a.25.25 0 00-.25-.25h-7.5z" fill="currentColor" fill-rule="evenodd"></path>
</svg>
</div>
</clipboard-copy>
</div>
<div class="highlight-ipynb hl-python"><pre><span></span><span class="kn">import</span><span class="w"> </span><span class="nn">equinox</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">eqx</span>
<span class="k">def</span><span class="w"> </span><span class="nf">loss_fn</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">spike_times</span><span class="p">,</span> <span class="n">comp_times</span><span class="p">):</span>
<span class="n">out_states</span> <span class="o">=</span> <span class="n">model</span><span class="o">.</span><span class="n">state_at_t</span><span class="p">(</span><span class="n">spike_times</span><span class="p">,</span> <span class="n">comp_times</span><span class="p">)</span>
<span class="n">logits</span> <span class="o">=</span> <span class="n">out_states</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">0</span><span class="p">]</span>
<span class="k">return</span> <span class="n">jnp</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">logits</span><span class="p">)</span>
<span class="n">in_spikes</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">asarray</span><span class="p">([[</span><span class="mf">0.01</span><span class="p">],</span> <span class="p">[</span><span class="mf">0.157</span><span class="p">]])</span>
<span class="n">comp_times</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">linspace</span><span class="p">(</span><span class="mf">0.0</span><span class="p">,</span> <span class="n">max_time</span><span class="p">,</span> <span class="mi">10</span><span class="p">)</span>
<span class="n">loss</span><span class="p">,</span> <span class="n">gradients</span> <span class="o">=</span> <span class="n">eqx</span><span class="o">.</span><span class="n">filter_value_and_grad</span><span class="p">(</span><span class="n">loss_fn</span><span class="p">)(</span><span class="n">snn</span><span class="p">,</span> <span class="n">in_spikes</span><span class="p">,</span> <span class="n">comp_times</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="s2">"Loss "</span><span class="p">,</span> <span class="n">loss</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="s2">"Gradients "</span><span class="p">,</span> <span class="n">gradients</span><span class="o">.</span><span class="n">neuron_model</span><span class="o">.</span><span class="n">weight_u</span><span class="p">)</span>
</pre></div>
<div class="clipboard-copy-txt" id="cell-6">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)</div>
</div>
</div>
</div>
</div>
<div class="jp-Cell-outputWrapper">
<div class="jp-Collapser jp-OutputCollapser jp-Cell-outputCollapser">
</div>
<div class="jp-OutputArea jp-Cell-outputArea">
<div class="jp-OutputArea-child">
<div class="jp-OutputPrompt jp-OutputArea-prompt"></div>
<div class="jp-RenderedText jp-OutputArea-output" data-mime-type="text/plain" tabindex="0">
<pre>Loss 2.78585239490824
Gradients [[ 0. ]
[-1.81404788]
[-1.42144198]]
</pre>
<img alt="No description has been provided for this image" class="" src="data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAABXsAAAIfCAYAAAAhag9YAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjgsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvwVt1zgAAAAlwSFlzAAAewgAAHsIBbtB1PgABAABJREFUeJzs3Xd4W+X1wPGvti157x07znJ2QiZJIBDK3nsH+qODMgpNGWWUQkuhhbKhjLI3ISRAEjLJ3ntPj3jEe1sesqT7+0OWLFmS4ziO7YTzeR4/ka/ee9/3Svc6V0fnnlelKIqCEEIIIYQQQgghhBBCiFOauqcHIIQQQgghhBBCCCGEEOLESbBXCCGEEEIIIYQQQgghTgMS7BVCCCGEEEIIIYQQQojTgAR7hRBCCCGEEEIIIYQQ4jQgwV4hhBBCCCGEEEIIIYQ4DUiwVwghhBBCCCGEEEIIIU4DEuwVQgghhBBCCCGEEEKI04AEe4UQQgghhBBCCCGEEOI0IMFeIYQQQgghhBBCCCGEOA1IsFcIIYQQQgghhBBCCCFOAxLsFUIIIYQQQgghhBBCiNOABHuFEEIIIYQQQgghhBDiNCDBXiGEEEIIIYQQQgghhDgNSLBXCCGEEEIIIYQQQgghTgMS7BVCCCGEEEIIIYQQQojTgAR7hRBCCCGEEEIIIYQQ4jQgwV4hhBBCCCGEEEIIIYQ4DUiwVwghhBBCCCGEEEIIIU4DEuwVQgghhBBCCCGEEEKI04AEe4UQQgghhBDiF85isWA0GlGpVNx8883ttn311VdRqVSoVCrOOuusdtvOmjXL1fbbb7/tyiGfdj766CPXa9VZf/vb31CpVKSmpvboNrqS8zVp+2M0GklLS+O6667jxx9/9Lnu8uXLXe1zcnK6d+CnAPfXx/1HrVYTEhLCkCFD+P3vf8/27dt7eqi9khxforeSYK8QQgghhBBC/MLp9XrGjx8PwKpVq9ptu3LlStfjjRs30tTU1KG2xwoMi94vJyfHFdxavnx5j46loaGBnJwcvv32Wy6//HKuuOKKdo9F0XGKolBbW8vevXt55513OOOMM3j22Wd7dEzO4+6jjz7q0XEcj9TUVFQqFX/72996eijiF0aCvUIIIYQQQgghXMHY/Px8srKy/LZbvXo1AAaDgaamJjZs2OC3rTPYO2jQIGJiYrpwtOKX5pZbbqG2ttb1U1hYyKxZsxg8eDAAP/zwAw8++GAPj/LU9fbbb7te25qaGrKysnjjjTcICwvDbrfzxBNP8M033/T0MIUQHSDBXiGEEEIIIYQQHpm37hm57vbv309JSQlGo5Ebbrih3bbV1dXs3LnTa9uid/vb3/6Goii97rZ0rVZLUFCQ6ycuLo6rr76alStXEhkZCcB7771HUVFRD4/01GQwGFyvbXBwMGlpadxzzz3Mnj3b1eaZZ57pwRH2PlOnTkVRFBRF6TVlT4QACfYKIYQQQgghhAAmTpyITqcD/AdwnSUeJkyYwDnnnOOxrK01a9Zgt9sBCfaKkycyMpL77rsPAKvVyrJly3p4RKeXqVOnMnXqVAD27NlDYWFhzw5ICHFMEuwVQgghhBBCCIHRaGTMmDGA/2Cvc/mUKVOYPHkyAGvXrsVms/ltC3D22Wd7Pd/Y2Mgbb7zBtGnTiImJQa/XExMTw0UXXcQ333yDoig+x9B2IrP8/HwefPBBBg4ciMlk8jlZUlFREY8//jhnnHEG4eHhGAwGUlJSuPXWW9m4caPf1+SOO+5ApVK5gl0bNmzglltuISUlBb1e75HNpygKGzdu5IknnuDMM88kMjISnU5HeHg4Y8eO5amnnqK8vNxvX23l5uZyzz330LdvXwICAoiJieGaa65h3bp1Hd7G7t27mT59OsnJyRgMBhITE7n99tvZt2+f33X8TdCWmppKWlqa6/dzzjnHa2KvnqrjO2TIENfj/Px8v+3MZjPPPPMMQ4cOxWQyERoaytSpU485eeCuXbv45z//ydSpU4mNjUWn0xESEsLw4cOZMWMGeXl57a5vNpv597//zZlnnklERAQ6nY7o6GgGDx7Mddddx3vvvYfZbPa7/vLly7nttttIS0sjMDCQkJAQRo4cyV//+lcqKyvb7bsrtPf6Ovdt4sSJREREuI6z6667joULF/rdZtv6z42Njfz73/9m7NixhIeHu+rzTp061WPSwjvvvNPruHOv49ud56y/CdqcYzhy5AgATz/9tNeYpY6vOKkUIYQQQgghhBBCUZRHHnlEARRAKSgo8Hq+T58+CqAsWbJEURRFiYuLUwBl48aNXm0nTpyoAEpaWprXc7t371bS0tJcffn6ueyyyxSz2ey17ocffuhqs2HDBiUiIsJr3ezsbFf7b775RjGZTO329cwzz/h8PaZPn64Aytlnn628/fbbikaj8VivT58+rrZz5sxptw9AiYuLU7Zt2+azL/f92rx5s8/9AhS1Wq289tprPrfx1FNPucY1b948JTAw0Oc29Hq9MmvWrGNuw53zvW/vZ9myZR7rOJefffbZPvvqCOc2pk+f7rfNrFmzXO3+9a9/uZYvW7bMtXzdunVKRkaG37H//e9/97nt7du3H3O/g4ODlUWLFvlc/+jRo0r//v2PuY1NmzZ5rdvY2Kjcdttt7a4XGxvrc92OcH99PvzwQ7/t7rvvPo9zzmnPnj1KSkpKu+O77bbbFIvF4rXN7OxsV5tZs2Ypw4cP91r3ww8/VM4+++xjvnbuY+/Oc9b99XP/m+McQ3s/Tz31VAffJSGOn2T2CiGEEEIIIYQA2q/bm5eXx5EjR9BqtUyYMAHAld3btm1DQwObN28GvLN68/PzmTp1KtnZ2fTp04d33nmHgwcPUlFRwb59+3j66acxGAz8+OOP3H333e2O95prrsFoNPLBBx+Qm5tLcXExP/30E+Hh4QDMnz+fG264AbPZzOTJk/nuu+/Iy8ujvLycjRs3cuuttwLw17/+lQ8//NBvP/v37+fee+9l3Lhx/PTTTxQXF5Obm8vrr7/uaqPVarn88st55513WL16NZmZmZSVlbF7927eeecdBgwYQFFREVdffTWNjY3t7td1112HVqvlf//7H3l5eRQVFfH111+TlpaG3W7n/vvvZ8mSJX7Xr66u5tZbbyU+Pp5vvvmGoqIi8vLyeO+994iOjsZisXDTTTexd+/edsfhbu/evezZs8f1+/z58z0mTKutrWXKlCkd3l5Xct+PxMREn21uueUWSktLeeONN8jKyqKsrIwlS5YwdOhQwJHR7CvjWaVSce655/LKK6+wYsUKDh48SHl5Ofv37+fzzz9nzJgx1NbWcsMNN/isF/zwww9z6NAhNBoNTz75JNu3b6ekpISSkhK2bdvGW2+95ZW96jR9+nQ+/fRTtFotf/zjH9mwYQNlZWUcPXqUb775hkGDBlFcXMyll156UmsV+3p9KyoqOP/888nNzcVgMPD0009z4MABysrKWLlyJRdccAEAn376KX/605/a3f4f//hHDh48yNNPP82+ffsoKytj48aNrvOttrbW1dZ9Ijnnj/M8dtfd56y7d955h9raWlJSUgD4y1/+4jXmxx57rMPbE+K49XS0WQghhBBCCCFE71BdXa2o1WoFUO6++26P5z777DMFUMaNG+da9sorryiAcsUVV3i0/fnnn10ZbB988IHHc1dccYUCKEOGDFEqKip8jmP+/PkeWa7u3DNgIyMjldzcXJ/baGhoUGJjYxVAueWWWxS73e6z3cMPP6wASkxMjNLQ0ODxnHuG3uTJk5Wmpiaf2+iI2tpaJT093edr0na/jEajsnfvXq82BQUFrn0aMmSI1/POrFxAiY+PVwoLC73a7N69WwkICFAA5eKLL/a7jbaZvYrimY3ZNovXF2fbk5nZW1VV5XpN1Gq1kpeX53rOPfMyKChI2bdvn9f6+fn5rgzoRx555LjH19zcrEyaNEkBlL/+9a9ezzsztB988MHj2u7s2bNd+/TTTz/5bFNVVeXKGr733nuPe+wdyexdu3atolKpFEAZMGCAa/kDDzzgWnfOnDle69lsNuXKK690tdm1a5fH8+7HEqDMmzev3bEea5xO3XnO+svsdXJmw0sWr+huktkrhBBCCCGEEALAVQsUvLN13ev1Ojkfr1692qPGrvu67tnC2dnZ/PDDDwC88sorrgzcti666CJXzc3PP//c73gfeughkpOTfT731VdfUVxcjNFo5M033/SZOQnw1FNPYTKZKCkpYdGiRX77evHFF9Hr9X6fP5agoCCuvvpqABYvXtxu23vuuYeMjAyv5QkJCTz++OOAY7KsDRs2+N3GE088QVxcnNfyIUOG8Ic//AGABQsWnNQJtxRFQVGUk1LL12w2s3jxYqZNm0ZxcTEAt912G0lJST7b33fffQwaNMhreWJiIr/61a8A2LRp03GPQ6vVcvPNNwO+31er1Qo43rvj8corrwBw0003ceGFF/psExoa6joevvjiC791rjujrKyMTz75hCuvvNK13b/85S8A2Gw2V53cSy65hCuuuMJrfbVazRtvvIFWqwXgvffe89vXhRdeyMUXX9xlY3fqznNWiN5Egr1CCCGEEEIIIVycwdm9e/d6TE60atUqwDPYO2LECIKDgykvL/e4xd8Z7E1MTCQ9Pd21fOnSpSiKQkBAAOPGjaOurs7vjzPo3F4A7pJLLvH7nLPMwcSJE9FoNH77sdvtriCgv74iIyMZP368376crFYrH3/8MZdffjkpKSkYjUaPSZleeOEFAA4cONDudpwBJl+uueYa1+M1a9ac0DbsdvtxTfjWkz7++GOP1zIoKIjzzz+fLVu2ADB16lTefPNNv+tfdNFFfp8bOHAgQLulEGbPns31119Penq6ayJA588999wD+H5fncfxCy+8Line truncated
</div>
</div>
</div>
@@ -1983,7 +1904,7 @@ Gradients [[ 0. ]
<div class="jp-InputPrompt jp-InputArea-prompt">In [ ]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="CodeMirror cm-s-jupyter">
<div class="zeroclipboard-container">
<clipboard-copy ,="" aria-label="Copy to Clipboard" for="cell-7">
<clipboard-copy ,="" aria-label="Copy to Clipboard" for="cell-6">
<div>
<span class="notice" hidden="">Copied!</span>
<svg aria-hidden="true" class="clipboard-copy-icon" data-view-component="true" height="20" version="1.1" viewbox="0 0 16 16" width="20">
@@ -1995,7 +1916,7 @@ Gradients [[ 0. ]
</div>
<div class="highlight-ipynb hl-python"><pre><span></span>
</pre></div>
<div class="clipboard-copy-txt" id="cell-7"></div>
<div class="clipboard-copy-txt" id="cell-6"></div>
</div>
</div>
</div>
print_page/index.html
+319 -234
View File
@@ -997,7 +997,7 @@
<section class="print-page" id="index" heading-number="1"><h1 id="index-felice">Felice</h1>
<p>This project provides a <a href="https://github.com/google/jax">JAX</a> implementation of the different neuron models in felice</p>
<h2 id="index-overview">Overview</h2>
<p>The framework is built on top of <a href="#.">EventPropJax</a> and leverages JAX's automatic differentiation for efficient simulation and training of SNNs using event-based exact gradients.</p>
<p>The framework is built on top of diffrax and leverages JAX's automatic differentiation for efficient simulation and training of analogue models.</p>
<h3 id="index-key-features">Key Features</h3>
<ul>
<li><strong>Delay learning</strong></li>
@@ -1016,37 +1016,6 @@
<p>For GPU acceleration with CUDA 13:</p>
<div class="language-bash highlight"><pre><span></span><code><span id="__span-1-1"><a id="__codelineno-1-1" name="__codelineno-1-1" href="#index-__codelineno-1-1"></a>uv<span class="w"> </span>sync<span class="w"> </span>--extra<span class="w"> </span>cuda
</span></code></pre></div>
<h2 id="index-quick-start">Quick Start</h2>
<p>Here's a simple example using the WereRabbit neuron model:</p>
<div class="language-py highlight"><pre><span></span><code><span id="__span-2-1"><a id="__codelineno-2-1" name="__codelineno-2-1" href="#index-__codelineno-2-1"></a><span class="kn">import</span><span class="w"> </span><span class="nn">diffrax</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">dfx</span>
</span><span id="__span-2-2"><a id="__codelineno-2-2" name="__codelineno-2-2" href="#index-__codelineno-2-2"></a><span class="kn">import</span><span class="w"> </span><span class="nn">jax.numpy</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">jnp</span>
</span><span id="__span-2-3"><a id="__codelineno-2-3" name="__codelineno-2-3" href="#index-__codelineno-2-3"></a><span class="kn">import</span><span class="w"> </span><span class="nn">jax.random</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">jrand</span>
</span><span id="__span-2-4"><a id="__codelineno-2-4" name="__codelineno-2-4" href="#index-__codelineno-2-4"></a><span class="kn">from</span><span class="w"> </span><span class="nn">eventpropjax.evnn</span><span class="w"> </span><span class="kn">import</span> <span class="n">FFEvNN</span>
</span><span id="__span-2-5"><a id="__codelineno-2-5" name="__codelineno-2-5" href="#index-__codelineno-2-5"></a><span class="kn">from</span><span class="w"> </span><span class="nn">felice.neuron_models</span><span class="w"> </span><span class="kn">import</span> <span class="n">WereRabbit</span>
</span><span id="__span-2-6"><a id="__codelineno-2-6" name="__codelineno-2-6" href="#index-__codelineno-2-6"></a>
</span><span id="__span-2-7"><a id="__codelineno-2-7" name="__codelineno-2-7" href="#index-__codelineno-2-7"></a><span class="c1"># Initialize random key and parameters</span>
</span><span id="__span-2-8"><a id="__codelineno-2-8" name="__codelineno-2-8" href="#index-__codelineno-2-8"></a><span class="n">key</span> <span class="o">=</span> <span class="n">jrand</span><span class="o">.</span><span class="n">key</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
</span><span id="__span-2-9"><a id="__codelineno-2-9" name="__codelineno-2-9" href="#index-__codelineno-2-9"></a><span class="n">max_time</span> <span class="o">=</span> <span class="mf">300e-3</span>
</span><span id="__span-2-10"><a id="__codelineno-2-10" name="__codelineno-2-10" href="#index-__codelineno-2-10"></a>
</span><span id="__span-2-11"><a id="__codelineno-2-11" name="__codelineno-2-11" href="#index-__codelineno-2-11"></a><span class="c1"># Create a feedforward event-driven neural network</span>
</span><span id="__span-2-12"><a id="__codelineno-2-12" name="__codelineno-2-12" href="#index-__codelineno-2-12"></a><span class="n">snn</span> <span class="o">=</span> <span class="n">FFEvNN</span><span class="p">(</span>
</span><span id="__span-2-13"><a id="__codelineno-2-13" name="__codelineno-2-13" href="#index-__codelineno-2-13"></a> <span class="n">layers</span><span class="o">=</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
</span><span id="__span-2-14"><a id="__codelineno-2-14" name="__codelineno-2-14" href="#index-__codelineno-2-14"></a> <span class="n">in_size</span><span class="o">=</span><span class="mi">2</span><span class="p">,</span>
</span><span id="__span-2-15"><a id="__codelineno-2-15" name="__codelineno-2-15" href="#index-__codelineno-2-15"></a> <span class="n">neuron_model</span><span class="o">=</span><span class="n">WereRabbit</span><span class="p">,</span>
</span><span id="__span-2-16"><a id="__codelineno-2-16" name="__codelineno-2-16" href="#index-__codelineno-2-16"></a> <span class="n">solver</span><span class="o">=</span><span class="n">dfx</span><span class="o">.</span><span class="n">Tsit5</span><span class="p">(),</span>
</span><span id="__span-2-17"><a id="__codelineno-2-17" name="__codelineno-2-17" href="#index-__codelineno-2-17"></a> <span class="n">max_solver_time</span><span class="o">=</span><span class="n">max_time</span><span class="p">,</span>
</span><span id="__span-2-18"><a id="__codelineno-2-18" name="__codelineno-2-18" href="#index-__codelineno-2-18"></a> <span class="n">key</span><span class="o">=</span><span class="n">key</span><span class="p">,</span>
</span><span id="__span-2-19"><a id="__codelineno-2-19" name="__codelineno-2-19" href="#index-__codelineno-2-19"></a> <span class="n">max_event_steps</span><span class="o">=</span><span class="mi">1000000</span><span class="p">,</span>
</span><span id="__span-2-20"><a id="__codelineno-2-20" name="__codelineno-2-20" href="#index-__codelineno-2-20"></a> <span class="n">solver_stepsize</span><span class="o">=</span><span class="mf">1e-6</span><span class="p">,</span>
</span><span id="__span-2-21"><a id="__codelineno-2-21" name="__codelineno-2-21" href="#index-__codelineno-2-21"></a> <span class="n">rtol</span><span class="o">=</span><span class="mf">10.0</span><span class="p">,</span>
</span><span id="__span-2-22"><a id="__codelineno-2-22" name="__codelineno-2-22" href="#index-__codelineno-2-22"></a> <span class="n">atol</span><span class="o">=</span><span class="mf">0.0</span><span class="p">,</span>
</span><span id="__span-2-23"><a id="__codelineno-2-23" name="__codelineno-2-23" href="#index-__codelineno-2-23"></a> <span class="n">Ibias</span><span class="o">=</span><span class="mf">30e-12</span><span class="p">,</span>
</span><span id="__span-2-24"><a id="__codelineno-2-24" name="__codelineno-2-24" href="#index-__codelineno-2-24"></a><span class="p">)</span>
</span><span id="__span-2-25"><a id="__codelineno-2-25" name="__codelineno-2-25" href="#index-__codelineno-2-25"></a>
</span><span id="__span-2-26"><a id="__codelineno-2-26" name="__codelineno-2-26" href="#index-__codelineno-2-26"></a><span class="c1"># Simulate with input spikes</span>
</span><span id="__span-2-27"><a id="__codelineno-2-27" name="__codelineno-2-27" href="#index-__codelineno-2-27"></a><span class="n">in_spikes</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">asarray</span><span class="p">([[</span><span class="mf">0.0</span><span class="p">],</span> <span class="p">[</span><span class="mf">0.157</span><span class="p">]])</span>
</span><span id="__span-2-28"><a id="__codelineno-2-28" name="__codelineno-2-28" href="#index-__codelineno-2-28"></a><span class="n">spikes</span> <span class="o">=</span> <span class="n">snn</span><span class="o">.</span><span class="n">spikes_until_t</span><span class="p">(</span><span class="n">in_spikes</span><span class="p">,</span> <span class="n">max_time</span><span class="p">)</span>
</span></code></pre></div>
<p>See the <a href="#scripts-examples-neuron_models">examples</a> directory for more detailed usage examples.</p></section>
<section class='print-page md-section' id='section-2' heading-number='2'>
<h1>Neuron Models<a class='headerlink' href='#section-2' title='Permanent link'></a>
@@ -1078,6 +1047,11 @@
<td>Exponential Integrate-and-Fire neuron model</td>
<td>...</td>
</tr>
<tr>
<td><a href="#neuron_models-lif">LIF</a></td>
<td>Leaky Integrate-and-Fire neuron model</td>
<td>...</td>
</tr>
</tbody>
</table></section>
<section class='print-page md-section' id='section-2-2' heading-number='2.2'>
@@ -1941,7 +1915,7 @@ span.linenos.special { color: #000000; background-color: #ffffc0; padding-left:
<div class="jp-Collapser jp-InputCollapser jp-Cell-inputCollapser">
</div>
<div class="jp-InputArea jp-Cell-inputArea">
<div class="jp-InputPrompt jp-InputArea-prompt">In [1]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="jp-InputPrompt jp-InputArea-prompt">In [10]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="CodeMirror cm-s-jupyter">
<div class="zeroclipboard-container">
<clipboard-copy ,="" aria-label="Copy to Clipboard" for="cell-1">
@@ -1960,7 +1934,6 @@ span.linenos.special { color: #000000; background-color: #ffffc0; padding-left:
<span class="kn">import</span><span class="w"> </span><span class="nn">jax.random</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">jrand</span>
<span class="kn">import</span><span class="w"> </span><span class="nn">matplotlib</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">mpl</span>
<span class="kn">import</span><span class="w"> </span><span class="nn">matplotlib.pyplot</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">plt</span>
<span class="kn">from</span><span class="w"> </span><span class="nn">eventpropjax.evnn</span><span class="w"> </span><span class="kn">import</span> <span class="n">FFEvNN</span>
<span class="kn">from</span><span class="w"> </span><span class="nn">felice.neuron_models</span><span class="w"> </span><span class="kn">import</span> <span class="n">WereRabbit</span>
@@ -1972,7 +1945,6 @@ 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
@@ -1988,7 +1960,7 @@ jax.config.update("jax_enable_x64", True)</div>
<div class="jp-Collapser jp-InputCollapser jp-Cell-inputCollapser">
</div>
<div class="jp-InputArea jp-Cell-inputArea">
<div class="jp-InputPrompt jp-InputArea-prompt">In [2]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="jp-InputPrompt jp-InputArea-prompt">In [11]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="CodeMirror cm-s-jupyter">
<div class="zeroclipboard-container">
<clipboard-copy ,="" aria-label="Copy to Clipboard" for="cell-2">
@@ -2002,44 +1974,46 @@ jax.config.update("jax_enable_x64", True)</div>
</clipboard-copy>
</div>
<div class="highlight-ipynb hl-python"><pre><span></span><span class="n">key</span> <span class="o">=</span> <span class="n">jrand</span><span class="o">.</span><span class="n">key</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
<span class="n">max_time</span> <span class="o">=</span> <span class="mi">100</span>
<span class="n">max_time</span> <span class="o">=</span> <span class="mi">40</span>
<span class="n">init_weights</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">asarray</span><span class="p">([[</span><span class="mf">0.0</span><span class="p">],</span> <span class="p">[</span><span class="mf">0.2</span><span class="p">],</span> <span class="p">[</span><span class="o">-</span><span class="mf">0.1</span><span class="p">]])</span>
<span class="n">snn</span> <span class="o">=</span> <span class="n">FFEvNN</span><span class="p">(</span>
<span class="n">layers</span><span class="o">=</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
<span class="n">in_size</span><span class="o">=</span><span class="mi">2</span><span class="p">,</span>
<span class="n">neuron_model</span><span class="o">=</span><span class="n">WereRabbit</span><span class="p">,</span>
<span class="n">model</span> <span class="o">=</span> <span class="n">WereRabbit</span><span class="p">(</span><span class="n">dtype</span><span class="o">=</span><span class="n">jnp</span><span class="o">.</span><span class="n">float64</span><span class="p">)</span>
<span class="k">def</span><span class="w"> </span><span class="nf">state_at_t</span><span class="p">(</span><span class="n">comp_times</span><span class="p">):</span>
<span class="n">sol</span> <span class="o">=</span> <span class="n">dfx</span><span class="o">.</span><span class="n">diffeqsolve</span><span class="p">(</span>
<span class="n">terms</span><span class="o">=</span><span class="n">dfx</span><span class="o">.</span><span class="n">ODETerm</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">dynamics</span><span class="p">),</span>
<span class="n">solver</span><span class="o">=</span><span class="n">dfx</span><span class="o">.</span><span class="n">Tsit5</span><span class="p">(),</span>
<span class="c1"># stepsize_controller=PIDController(</span>
<span class="c1"># rtol=1e-6, atol=1e-3, pcoeff=0.1, icoeff=0.3, dcoeff=0.0</span>
<span class="c1"># ),</span>
<span class="n">max_solver_time</span><span class="o">=</span><span class="n">max_time</span><span class="p">,</span>
<span class="n">key</span><span class="o">=</span><span class="n">key</span><span class="p">,</span>
<span class="n">max_event_steps</span><span class="o">=</span><span class="mi">10000</span><span class="p">,</span>
<span class="n">solver_stepsize</span><span class="o">=</span><span class="mf">0.001</span><span class="p">,</span>
<span class="n">init_weights</span><span class="o">=</span><span class="n">init_weights</span><span class="p">,</span>
<span class="n">dtype</span><span class="o">=</span><span class="n">jnp</span><span class="o">.</span><span class="n">float64</span><span class="p">,</span>
<span class="p">)</span>
<span class="n">t0</span><span class="o">=</span><span class="mf">0.0</span><span class="p">,</span>
<span class="n">t1</span><span class="o">=</span><span class="n">max_time</span><span class="p">,</span>
<span class="n">dt0</span><span class="o">=</span><span class="mf">1e-3</span><span class="p">,</span>
<span class="n">y0</span><span class="o">=</span><span class="n">model</span><span class="o">.</span><span class="n">init_state</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>
<span class="o">+</span> <span class="n">jrand</span><span class="o">.</span><span class="n">uniform</span><span class="p">(</span><span class="n">key</span><span class="p">,</span> <span class="n">shape</span><span class="o">=</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">),</span> <span class="n">minval</span><span class="o">=</span><span class="mf">0.1</span><span class="p">,</span> <span class="n">maxval</span><span class="o">=</span><span class="mf">0.5</span><span class="p">),</span>
<span class="n">saveat</span><span class="o">=</span><span class="n">dfx</span><span class="o">.</span><span class="n">SaveAt</span><span class="p">(</span><span class="n">ts</span><span class="o">=</span><span class="n">comp_times</span><span class="p">),</span>
<span class="n">max_steps</span><span class="o">=</span><span class="mi">100000</span><span class="p">,</span>
<span class="p">)</span>
<span class="k">return</span> <span class="n">sol</span><span class="o">.</span><span class="n">ts</span><span class="p">,</span> <span class="n">sol</span><span class="o">.</span><span class="n">ys</span>
</pre></div>
<div class="clipboard-copy-txt" id="cell-2">key = jrand.key(0)
max_time = 100
max_time = 40
init_weights = jnp.asarray([[0.0], [0.2], [-0.1]])
snn = FFEvNN(
layers=[1],
in_size=2,
neuron_model=WereRabbit,
model = WereRabbit(dtype=jnp.float64)
def state_at_t(comp_times):
sol = dfx.diffeqsolve(
terms=dfx.ODETerm(model.dynamics),
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,
)</div>
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</div>
</div>
</div>
</div>
@@ -2051,7 +2025,7 @@ snn = FFEvNN(
<div class="jp-Collapser jp-InputCollapser jp-Cell-inputCollapser">
</div>
<div class="jp-InputArea jp-Cell-inputArea">
<div class="jp-InputPrompt jp-InputArea-prompt">In [3]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="jp-InputPrompt jp-InputArea-prompt">In [12]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="CodeMirror cm-s-jupyter">
<div class="zeroclipboard-container">
<clipboard-copy ,="" aria-label="Copy to Clipboard" for="cell-3">
@@ -2064,15 +2038,11 @@ snn = FFEvNN(
</div>
</clipboard-copy>
</div>
<div class="highlight-ipynb hl-python"><pre><span></span><span class="n">in_spikes</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">asarray</span><span class="p">([[</span><span class="mi">10</span><span class="p">,</span> <span class="mi">30</span><span class="p">],</span> <span class="p">[</span><span class="mi">20</span><span class="p">,</span> <span class="mi">21</span><span class="p">]])</span>
<span class="n">comp_times</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">linspace</span><span class="p">(</span><span class="mf">0.0</span><span class="p">,</span> <span class="n">max_time</span><span class="p">,</span> <span class="mi">2000</span><span class="p">)</span>
<span class="n">state</span> <span class="o">=</span> <span class="n">snn</span><span class="o">.</span><span class="n">state_at_t</span><span class="p">(</span><span class="n">in_spikes</span><span class="p">,</span> <span class="n">comp_times</span><span class="p">)</span>
<span class="n">spikes</span> <span class="o">=</span> <span class="n">snn</span><span class="o">.</span><span class="n">spikes_until_t</span><span class="p">(</span><span class="n">in_spikes</span><span class="p">,</span> <span class="n">max_time</span><span class="p">)</span>
<div class="highlight-ipynb hl-python"><pre><span></span><span class="n">comp_times</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">linspace</span><span class="p">(</span><span class="mf">0.0</span><span class="p">,</span> <span class="n">max_time</span><span class="p">,</span> <span class="mi">2000</span><span class="p">)</span>
<span class="n">_</span><span class="p">,</span> <span class="n">state</span> <span class="o">=</span> <span class="n">state_at_t</span><span class="p">(</span><span class="n">comp_times</span><span class="p">)</span>
</pre></div>
<div class="clipboard-copy-txt" id="cell-3">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)</div>
<div class="clipboard-copy-txt" id="cell-3">comp_times = jnp.linspace(0.0, max_time, 2000)
_, state = state_at_t(comp_times)</div>
</div>
</div>
</div>
@@ -2084,7 +2054,7 @@ spikes = snn.spikes_until_t(in_spikes, max_time)</div>
<div class="jp-Collapser jp-InputCollapser jp-Cell-inputCollapser">
</div>
<div class="jp-InputArea jp-Cell-inputArea">
<div class="jp-InputPrompt jp-InputArea-prompt">In [4]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="jp-InputPrompt jp-InputArea-prompt">In [13]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="CodeMirror cm-s-jupyter">
<div class="zeroclipboard-container">
<clipboard-copy ,="" aria-label="Copy to Clipboard" for="cell-4">
@@ -2110,7 +2080,7 @@ spikes = snn.spikes_until_t(in_spikes, max_time)</div>
<span class="n">UV</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">stack</span><span class="p">(</span>
<span class="p">[</span><span class="n">U</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">),</span> <span class="n">V</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">),</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ones</span><span class="p">((</span><span class="n">resolution</span> <span class="o">*</span> <span class="n">resolution</span><span class="p">,))],</span> <span class="n">axis</span><span class="o">=</span><span class="mi">1</span>
<span class="p">)</span>
<span class="n">dS</span> <span class="o">=</span> <span class="n">snn</span><span class="o">.</span><span class="n">neuron_model</span><span class="o">.</span><span class="n">vector_field</span><span class="p">(</span><span class="n">UV</span><span class="p">)</span>
<span class="n">dS</span> <span class="o">=</span> <span class="n">snn</span><span class="o">.</span><span class="n">vector_field</span><span class="p">(</span><span class="n">UV</span><span class="p">)</span>
<span class="n">dU</span> <span class="o">=</span> <span class="n">dS</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">U</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
<span class="n">dV</span> <span class="o">=</span> <span class="n">dS</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">V</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
@@ -2128,7 +2098,7 @@ spikes = snn.spikes_until_t(in_spikes, max_time)</div>
<span class="n">U</span><span class="p">,</span> <span class="n">V</span><span class="p">,</span> <span class="n">dU</span><span class="p">,</span> <span class="n">dV</span> <span class="o">=</span> <span class="n">compute_nullclines</span><span class="p">(</span><span class="n">snn</span><span class="p">,</span> <span class="n">u_range</span><span class="p">,</span> <span class="n">v_range</span><span class="p">,</span> <span class="mi">200</span><span class="p">)</span>
<span class="n">UVs</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">stack</span><span class="p">([</span><span class="n">Us</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">),</span> <span class="n">Vs</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">),</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ones</span><span class="p">((</span><span class="mi">20</span> <span class="o">*</span> <span class="mi">20</span><span class="p">,))],</span> <span class="n">axis</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
<span class="n">dS</span> <span class="o">=</span> <span class="n">snn</span><span class="o">.</span><span class="n">neuron_model</span><span class="o">.</span><span class="n">vector_field</span><span class="p">(</span><span class="n">UVs</span><span class="p">)</span>
<span class="n">dS</span> <span class="o">=</span> <span class="n">snn</span><span class="o">.</span><span class="n">vector_field</span><span class="p">(</span><span class="n">UVs</span><span class="p">)</span>
<span class="n">dUs</span> <span class="o">=</span> <span class="n">dS</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">Us</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
<span class="n">dVs</span> <span class="o">=</span> <span class="n">dS</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">Vs</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
@@ -2167,7 +2137,7 @@ spikes = snn.spikes_until_t(in_spikes, max_time)</div>
UV = jnp.stack(
[U.reshape(-1), V.reshape(-1), jnp.ones((resolution * resolution,))], axis=1
)
dS = snn.neuron_model.vector_field(UV)
dS = snn.vector_field(UV)
dU = dS[:, 0].reshape(U.shape)
dV = dS[:, 1].reshape(V.shape)
@@ -2185,7 +2155,7 @@ def plot_vf(ax, snn, u_range, v_range):
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)
dS = snn.vector_field(UVs)
dUs = dS[:, 0].reshape(Us.shape)
dVs = dS[:, 1].reshape(Vs.shape)
@@ -2221,7 +2191,7 @@ def plot_vf(ax, snn, u_range, v_range):
<div class="jp-Collapser jp-InputCollapser jp-Cell-inputCollapser">
</div>
<div class="jp-InputArea jp-Cell-inputArea">
<div class="jp-InputPrompt jp-InputArea-prompt">In [5]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="jp-InputPrompt jp-InputArea-prompt">In [14]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="CodeMirror cm-s-jupyter">
<div class="zeroclipboard-container">
<clipboard-copy ,="" aria-label="Copy to Clipboard" for="cell-5">
@@ -2236,29 +2206,27 @@ def plot_vf(ax, snn, u_range, v_range):
</div>
<div class="highlight-ipynb hl-python"><pre><span></span><span class="k">with</span> <span class="n">mpl</span><span class="o">.</span><span class="n">style</span><span class="o">.</span><span class="n">context</span><span class="p">(</span><span class="s2">"boilerplot.ieeetran"</span><span class="p">):</span>
<span class="n">fig</span><span class="p">,</span> <span class="n">ax</span> <span class="o">=</span> <span class="n">plt</span><span class="o">.</span><span class="n">subplots</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="n">figsize</span><span class="o">=</span><span class="p">(</span><span class="mf">6.9</span><span class="p">,</span> <span class="mf">2.6</span><span class="p">),</span> <span class="n">dpi</span><span class="o">=</span><span class="mi">200</span><span class="p">)</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">comp_times</span><span class="p">,</span> <span class="n">state</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="p">:,</span> <span class="mi">0</span><span class="p">],</span> <span class="n">label</span><span class="o">=</span><span class="s2">"x1"</span><span class="p">)</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">comp_times</span><span class="p">,</span> <span class="n">state</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="p">:,</span> <span class="mi">1</span><span class="p">],</span> <span class="n">label</span><span class="o">=</span><span class="s2">"x2"</span><span class="p">)</span>
<span class="p">[</span><span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">axvline</span><span class="p">(</span><span class="n">s</span><span class="p">,</span> <span class="n">alpha</span><span class="o">=</span><span class="mf">0.2</span><span class="p">,</span> <span class="n">color</span><span class="o">=</span><span class="s2">"g"</span><span class="p">,</span> <span class="n">linestyle</span><span class="o">=</span><span class="s2">"--"</span><span class="p">)</span> <span class="k">for</span> <span class="n">s</span> <span class="ow">in</span> <span class="n">jnp</span><span class="o">.</span><span class="n">unique</span><span class="p">(</span><span class="n">spikes</span><span class="p">)]</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">legend</span><span class="p">([</span><span class="s2">"x1"</span><span class="p">,</span> <span class="s2">"x2"</span><span class="p">,</span> <span class="s2">"Spike"</span><span class="p">])</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">comp_times</span><span class="p">,</span> <span class="n">state</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span> <span class="n">label</span><span class="o">=</span><span class="s2">"x1"</span><span class="p">)</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">comp_times</span><span class="p">,</span> <span class="n">state</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">],</span> <span class="n">label</span><span class="o">=</span><span class="s2">"x2"</span><span class="p">)</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">legend</span><span class="p">([</span><span class="s2">"x1"</span><span class="p">,</span> <span class="s2">"x2"</span><span class="p">])</span>
<span class="n">plot_vf</span><span class="p">(</span><span class="n">ax</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">snn</span><span class="p">,</span> <span class="p">[</span><span class="o">-</span><span class="mf">0.2</span><span class="p">,</span> <span class="mf">0.5</span><span class="p">],</span> <span class="p">[</span><span class="o">-</span><span class="mf">0.2</span><span class="p">,</span> <span class="mf">0.5</span><span class="p">])</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">state</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="p">:,</span> <span class="mi">0</span><span class="p">],</span> <span class="n">state</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="p">:,</span> <span class="mi">1</span><span class="p">])</span>
<span class="n">plot_vf</span><span class="p">(</span><span class="n">ax</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">model</span><span class="p">,</span> <span class="p">[</span><span class="o">-</span><span class="mf">0.2</span><span class="p">,</span> <span class="mf">0.5</span><span class="p">],</span> <span class="p">[</span><span class="o">-</span><span class="mf">0.2</span><span class="p">,</span> <span class="mf">0.5</span><span class="p">])</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">state</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span> <span class="n">state</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">])</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">state</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span> <span class="n">state</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">],</span> <span class="s2">"."</span><span class="p">,</span> <span class="n">label</span><span class="o">=</span><span class="s2">"start"</span><span class="p">)</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">state</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span> <span class="n">state</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">],</span> <span class="s2">"."</span><span class="p">,</span> <span class="n">label</span><span class="o">=</span><span class="s2">"end"</span><span class="p">)</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">state</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span> <span class="n">state</span><span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">],</span> <span class="s2">"."</span><span class="p">,</span> <span class="n">label</span><span class="o">=</span><span class="s2">"end"</span><span class="p">)</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">legend</span><span class="p">()</span>
<span class="n">plt</span><span class="o">.</span><span class="n">show</span><span class="p">()</span>
</pre></div>
<div class="clipboard-copy-txt" id="cell-5">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"])
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], snn, [-0.2, 0.5], [-0.2, 0.5])
ax[1].plot(state[0, :, 0], state[0, :, 1])
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[0, -1, 0], state[0, -1, 1], ".", label="end")
ax[1].plot(state[-1, 0, 0], state[-1, 0, 1], ".", label="end")
ax[1].legend()
plt.show()</div>
</div>
@@ -2272,80 +2240,7 @@ def plot_vf(ax, snn, u_range, v_range):
<div class="jp-OutputArea-child">
<div class="jp-OutputPrompt jp-OutputArea-prompt"></div>
<div class="jp-RenderedImage jp-OutputArea-output" tabindex="0">
<img alt="No description has been provided for this image" class="" src="data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAABXsAAAIfCAYAAAAhag9YAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjgsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvwVt1zgAAAAlwSFlzAAAewgAAHsIBbtB1PgABAABJREFUeJzs3Xd4W/XZxvHv0bQlee9tx1nOJINASAIhAcre42VDB4WyS1sooxTaUrqBAoUuoAXKDjNAmJmQvfdwYifxdjwkW/u8f8ialke8kzyf6/Jl6egsSUe2dOs5z09RVVVFCCGEEEIIIYQQQgghxBFNM9g7IIQQQgghhBBCCCGEEKL3JOwVQgghhBBCCCGEEEKIo4CEvUIIIYQQQgghhBBCCHEUkLBXCCGEEEIIIYQQQgghjgIS9gohhBBCCCGEEEIIIcRRQMJeIYQQQgghhBBCCCGEOApI2CuEEEIIIYQQQgghhBBHAQl7hRBCCCGEEEIIIYQQ4iggYa8QQgghhBBCCCGEEEIcBSTsFUIIIYQQQgghhBBCiKOAhL1CCCGEEEIIIYQQQghxFJCwVwghhBBCCCGEEEIIIY4CEvYKIYQQQgghhBBCCCHEUUDCXiGEEEIIIYQQQgghhDgKSNgrhBBCCCGEEEIIIYQQRwEJe4UQQgghhBBCCCGEEOIoIGGvEEIIIYQQQgghhBBCHAUk7BVCCCGEEEIIIYQQQoijgIS9QgghhBBCCCGEEEIIcRSQsFcIIYQQQgghjjFOpxOTyYSiKFx11VWdzvvkk0+iKAqKonDyySd3Ou/bb78dmPett97qy10+6rz44ouBx6qnfvnLX6IoCoWFhYO6jr7kf0wif0wmE0VFRVx22WV88MEHUZf9+uuvA/Pv3bt3YHf8CBD6+IT+aDQa4uPjGTt2LDfffDPr1q0b7F0dkuT4EkcKCXuFEEIIIYQQ4hhjMBg44YQTAFi8eHGn8y5atChwecWKFTgcjm7N21UwLIa+vXv3BsKtr7/+elD3pbW1lb179/LWW29x/vnnc8EFF3R6LIruU1WV5uZmtmzZwvPPP8+UKVP4zW9+M6j75D/uXnzxxUHdj8NRWFiIoij88pe/HOxdEcc4CXuFEEIIIYQQ4hjkD2P379/Pnj17OpxvyZIlABiNRhwOB8uXL+9wXn/YO3r0aNLT0/twb8Wx5uqrr6a5uTnwU1FRwdtvv82YMWMAeP/997n77rsHeS+PXM8991zgsW1qamLPnj08/fTTJCYm4vV6efDBB3njjTcGezeFED0gYa8QQgghhBBCHINCK29DK3JDbdu2jerqakwmE1dccUWn8zY2NrJhw4Z26xZD2y9/+UtUVR1yp6XrdDosFkvgJzMzk4svvphFixaRkpICwD/+8Q8qKysHeU+PTEajMfDYxsXFUVRUxK233sq8efMC8zz66KODuIdDz+zZs1FVFVVVh0zbEyGikbBXCCGEEEIIIY5B06dPR6/XAx0HuP4WDyeeeCKnnnpq2LRIS5cuxev1AhL2iv6TkpLC7bffDoDb7earr74a5D06usyePZvZs2cDsHnzZioqKgZ3h4QQh03CXiGEEEIIIYQ4BplMJqZOnQp0HPb6p8+aNYuZM2cCsGzZMjweT4fzApxyyintbrfb7Tz99NPMnTuX9PR0DAYD6enpnHXWWbzxxhuoqhp1HyIHMtu/fz933303o0aNwmw2Rx0sqbKykgceeIApU6aQlJSE0WgkPz+fa665hhUrVnT4mNxwww0oihIIu5YvX87VV19Nfn4+BoMhrJpPVVVWrFjBgw8+yEknnURKSgp6vZ6kpCSOP/54Hn74Yerq6jrcVqSysjJuvfVWhg0bRkxMDOnp6VxyySV888033V7Hpk2buP7668nLy8NoNJKTk8N1113H1q1bO1ymowHaCgsLKSoqClw/9dRT2w3sNVh9fMeOHRu4vH///g7ns9lsPProo4wbNw6z2UxCQgKzZ8/ucvDAjRs38thjjzF79mwyMjLQ6/XEx8czYcIE7rnnHsrLyztd3maz8fvf/56TTjqJ5ORk9Ho9aWlpjBkzhssuu4x//OMf2Gy2Dpf/+uuvufbaaykqKiI2Npb4+HiOO+44fvGLX3Do0KFOt90XOnt8/fdt+vTpJCcnB46zyy67jE8//bTDdUb2f7bb7fz+97/n+OOPJykpKdCfd/bs2WGDFt54443tjrvQPr4D+ZrtaIA2/z7s27cPgEceeaTdPksfXzGgVCGEEEIIIYQQx6R7771XBVRAPXDgQLvbCwoKVED9/PPPVVVV1czMTBVQV6xY0W7e6dOnq4BaVFTU7rZNmzapRUVFgW1F+znvvPNUm83WbtkXXnghMM/y5cvV5OTkdsuWlpYG5n/jjTdUs9nc6bYeffTRqI/H9ddfrwLqKaecoj733HOqVqsNW66goCAw77vvvtvpNgA1MzNTXbt2bdRthd6vVatWRb1fgKrRaNSnnnoq6joefvjhwH599NFHamxsbNR1GAwG9e233+5yHaH8z31nP1999VXYMv7pp5xyStRtdYd/Hddff32H87z99tuB+X73u98Fpn/11VeB6d98841aUlLS4b7/6le/irrudevWdXm/4+Li1AULFkRd/uDBg+qIESO6XMfKlSvbLWu329Vrr7220+UyMjKiLtsdoY/PCy+80OF8t99+e9hrzm/z5s1qfn5+p/t37bXXqk6ns906S0tLA/O8/fbb6oQJE9ot+8ILL6innHJKl49d6L4P5Gs29PEL/Zvj34fOfh5++OFuPktC9J5U9gohhBBCCCHEMaqzvr3l5eXs27cPnU7HiSeeCBCo7o2ct7W1lVWrVgHtq3r379/P7NmzKS0tpaCggOeff54dO3ZQX1/P1q1beeSRRzAajXzwwQfccsstne7vJZdcgslk4t///jdlZWVUVVXx8ccfk5SUBMD8+fO54oorsNlszJw5k3feeYfy8nLq6upYsWIF11xzDQC/+MUveOGFFzrczrZt27jtttuYNm0aH3/8MVVVVZSVlfHXv/41MI9Op+P888/n+eefZ8mSJezevZva2lo2bdrE888/z8iRI6msrOTiiy/Gbrd3er8uu+wydDod//znPykvL6eyspLXX3+doqIivF4vd9xxB59//nmHyzc2NnLNNdeQlZXFG2+8QWVlJeXl5fzjH/8gLS0Np9PJlVdeyZYtWzrdj1Bbtmxh8+bNgevz588PGzCtubmZWbNmdXt9fSn0fuTk5ESd5+qrr6ampoann36aPXv2UFtby+eff864ceMAX0VztIpnRVGYM2cOTzzxBAsXLmTHjh3U1dWxbds2XnnlFaZOnUpzczNXXHFF1H7BP/vZz9i5cydarZaHHnqIdevWUV1dTXV1NWvXruXZZ59tV73qd/311/Pf//4XnU7HnXfeyfLly6mtreXgwYO88cYbjB49mqqqKs4999x+7VUc7fGtr6/njDPOoKysDKPRyCOPPML27dupra1l0aJFfOc73wHgv//9Lz/+8Y87Xf+dd97Jjh07eOSRR9i6dSu1tbWsWLEi8Hprbm4OzBs6kJz/x/86DjXQr9lQzz//PM3NzeTn5wPw85//vN0+33///d1enxC9NthpsxBCCCGEEEKIwdHY2KhqNBoVUG+55Zaw215++WUVUKdNmxaY9sQTT6iAesEFF4TN++WXXwYq2P7973+H3XbBBReogDp27Fi1vr4+6n7Mnz8/rMo1VGgFbEpKilpWVhZ1Ha2trWpGRoYKqFdffbXq9Xqjzvezn/1MBdT09HS1tbU17LbQCr2ZM2eqDocj6jq6o7m5WS0uLo76mETeL5PJpG7ZsqXdPAcOHAjcp7Fjx7a73V+VC6hZWVlqRUVFu3k2bdqkxsTEqIB69tlnd7iOyMpeVQ2vxoys4o3GP29/VvY2NDQEHhONRqOWl5cHbgutvLRYLOrWrVvbLb9///5ABfS999572PvncrnUGTNmqID6i1/8ot3t/grtu++++7DWO2/evMB9+vjjj6PO09DQEKgavu222w5737tT2bts2TJVURQVUEeOHBmYftdddwWWfffdd9st5/F41AsvvDAwz8aNG8NuDz2WAPWjjz7qdF+72k+/gXzNdlTZ6+evhpcqXjHYpLJXCCGEEEIIIY5R/l6g0L5aN7Rfr5//8pIlS8J67IYuG1otXFpayvvvvw/AE088EajAjXTWWWcFem6+8sorHe7vT3/6U/Ly8qLe9tprr1FVVYXJZOKZZ56JWjkJ8PDDD2M2m6murmbBggUdbuuPf/wjBoOhw9u7YrFYuPjiiwH47LPPOp331ltvpaSkpN307OxsHnjgAcA3WNby5cs7XMeDDz5IZmZmu+ljx47lRz/6EQCffPJJvw64paoqqqr2Sy9fm83GZ599xty5c6mqqgLg2muvJTc3N+r8t99+O6NHj243PScnh9NPPx2AlStXHvZ+6HQ6rrrqKiD68+p2uwHfc3c4nnjiCQCuvPJKzjzzzKjzJCQkBI6HV199tcM+1z1RW1vLf/7zHy688MLAen/+858D4PF4An1yzznnHC644IJ2y2s0Gp5++ml0Oh0A//jHPzrc1plnnsnZZ5/dZ/vuN5CvWSGGMgl7hRBCCCGEEOIY5g9nt2zZEjY40eLFi4HwsHfixInExcVRV1cXdoq/P+zNycmhuLg4MP2LL75AVVViYmKYNm0aVqu1wx9/6NxZAHfOOed0eJu/zcH06dPRarUdbsfr9QZCwI62lZKSwgknnNDhtvzcbjcvvfQS559/Pvn5+ZhMprBBmf7whz8AsH379k7X4w+YornkkksCl5cuXdqrdXi93sMa8G0wvfTSS2GPpcVi4YwzzmD16tUAzJ49m2eeeabD5c8666wObxs1ahRAp60Q5s2bLine truncated
</div>
</div>
</div>
</div>
</div>
</div><div class="jp-Cell jp-CodeCell jp-Notebook-cell">
<div class="jp-Cell jp-CodeCell jp-Notebook-cell" id="cell-id=44e22685">
<div class="jp-Cell-inputWrapper" tabindex="0">
<div class="jp-Collapser jp-InputCollapser jp-Cell-inputCollapser">
</div>
<div class="jp-InputArea jp-Cell-inputArea">
<div class="jp-InputPrompt jp-InputArea-prompt">In [6]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="CodeMirror cm-s-jupyter">
<div class="zeroclipboard-container">
<clipboard-copy ,="" aria-label="Copy to Clipboard" for="cell-6">
<div>
<span class="notice" hidden="">Copied!</span>
<svg aria-hidden="true" class="clipboard-copy-icon" data-view-component="true" height="20" version="1.1" viewbox="0 0 16 16" width="20">
<path d="M0 6.75C0 5.784.784 5 1.75 5h1.5a.75.75 0 010 1.5h-1.5a.25.25 0 00-.25.25v7.5c0 .138.112.25.25.25h7.5a.25.25 0 00.25-.25v-1.5a.75.75 0 011.5 0v1.5A1.75 1.75 0 019.25 16h-7.5A1.75 1.75 0 010 14.25v-7.5z" fill="currentColor" fill-rule="evenodd"></path>
<path d="M5 1.75C5 .784 5.784 0 6.75 0h7.5C15.216 0 16 .784 16 1.75v7.5A1.75 1.75 0 0114.25 11h-7.5A1.75 1.75 0 015 9.25v-7.5zm1.75-.25a.25.25 0 00-.25.25v7.5c0 .138.112.25.25.25h7.5a.25.25 0 00.25-.25v-7.5a.25.25 0 00-.25-.25h-7.5z" fill="currentColor" fill-rule="evenodd"></path>
</svg>
</div>
</clipboard-copy>
</div>
<div class="highlight-ipynb hl-python"><pre><span></span><span class="kn">import</span><span class="w"> </span><span class="nn">equinox</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">eqx</span>
<span class="k">def</span><span class="w"> </span><span class="nf">loss_fn</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">spike_times</span><span class="p">,</span> <span class="n">comp_times</span><span class="p">):</span>
<span class="n">out_states</span> <span class="o">=</span> <span class="n">model</span><span class="o">.</span><span class="n">state_at_t</span><span class="p">(</span><span class="n">spike_times</span><span class="p">,</span> <span class="n">comp_times</span><span class="p">)</span>
<span class="n">logits</span> <span class="o">=</span> <span class="n">out_states</span><span class="p">[:,</span> <span class="p">:,</span> <span class="mi">0</span><span class="p">]</span>
<span class="k">return</span> <span class="n">jnp</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">logits</span><span class="p">)</span>
<span class="n">in_spikes</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">asarray</span><span class="p">([[</span><span class="mf">0.01</span><span class="p">],</span> <span class="p">[</span><span class="mf">0.157</span><span class="p">]])</span>
<span class="n">comp_times</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">linspace</span><span class="p">(</span><span class="mf">0.0</span><span class="p">,</span> <span class="n">max_time</span><span class="p">,</span> <span class="mi">10</span><span class="p">)</span>
<span class="n">loss</span><span class="p">,</span> <span class="n">gradients</span> <span class="o">=</span> <span class="n">eqx</span><span class="o">.</span><span class="n">filter_value_and_grad</span><span class="p">(</span><span class="n">loss_fn</span><span class="p">)(</span><span class="n">snn</span><span class="p">,</span> <span class="n">in_spikes</span><span class="p">,</span> <span class="n">comp_times</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="s2">"Loss "</span><span class="p">,</span> <span class="n">loss</span><span class="p">)</span>
<span class="nb">print</span><span class="p">(</span><span class="s2">"Gradients "</span><span class="p">,</span> <span class="n">gradients</span><span class="o">.</span><span class="n">neuron_model</span><span class="o">.</span><span class="n">weight_u</span><span class="p">)</span>
</pre></div>
<div class="clipboard-copy-txt" id="cell-6">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)</div>
</div>
</div>
</div>
</div>
<div class="jp-Cell-outputWrapper">
<div class="jp-Collapser jp-OutputCollapser jp-Cell-outputCollapser">
</div>
<div class="jp-OutputArea jp-Cell-outputArea">
<div class="jp-OutputArea-child">
<div class="jp-OutputPrompt jp-OutputArea-prompt"></div>
<div class="jp-RenderedText jp-OutputArea-output" data-mime-type="text/plain" tabindex="0">
<pre>Loss 2.78585239490824
Gradients [[ 0. ]
[-1.81404788]
[-1.42144198]]
</pre>
<img alt="No description has been provided for this image" class="" src="data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAABXsAAAIfCAYAAAAhag9YAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjgsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvwVt1zgAAAAlwSFlzAAAewgAAHsIBbtB1PgABAABJREFUeJzs3Xd4W+X1wPGvti157x07znJ2QiZJIBDK3nsH+qODMgpNGWWUQkuhhbKhjLI3ISRAEjLJ3ntPj3jEe1sesqT7+0OWLFmS4ziO7YTzeR4/ka/ee9/3Svc6V0fnnlelKIqCEEIIIYQQQgghhBBCiFOauqcHIIQQQgghhBBCCCGEEOLESbBXCCGEEEIIIYQQQgghTgMS7BVCCCGEEEIIIYQQQojTgAR7hRBCCCGEEEIIIYQQ4jQgwV4hhBBCCCGEEEIIIYQ4DUiwVwghhBBCCCGEEEIIIU4DEuwVQgghhBBCCCGEEEKI04AEe4UQQgghhBBCCCGEEOI0IMFeIYQQQgghhBBCCCGEOA1IsFcIIYQQQgghhBBCCCFOAxLsFUIIIYQQQgghhBBCiNOABHuFEEIIIYQQQgghhBDiNCDBXiGEEEIIIYQQQgghhDgNSLBXCCGEEEIIIYQQQgghTgMS7BVCCCGEEEIIIYQQQojTgAR7hRBCCCGEEEIIIYQQ4jQgwV4hhBBCCCGEEEIIIYQ4DUiwVwghhBBCCCGEEEIIIU4DEuwVQgghhBBCCCGEEEKI04AEe4UQQgghhBDiF85isWA0GlGpVNx8883ttn311VdRqVSoVCrOOuusdtvOmjXL1fbbb7/tyiGfdj766CPXa9VZf/vb31CpVKSmpvboNrqS8zVp+2M0GklLS+O6667jxx9/9Lnu8uXLXe1zcnK6d+CnAPfXx/1HrVYTEhLCkCFD+P3vf8/27dt7eqi9khxforeSYK8QQgghhBBC/MLp9XrGjx8PwKpVq9ptu3LlStfjjRs30tTU1KG2xwoMi94vJyfHFdxavnx5j46loaGBnJwcvv32Wy6//HKuuOKKdo9F0XGKolBbW8vevXt55513OOOMM3j22Wd7dEzO4+6jjz7q0XEcj9TUVFQqFX/72996eijiF0aCvUIIIYQQQgghXMHY/Px8srKy/LZbvXo1AAaDgaamJjZs2OC3rTPYO2jQIGJiYrpwtOKX5pZbbqG2ttb1U1hYyKxZsxg8eDAAP/zwAw8++GAPj/LU9fbbb7te25qaGrKysnjjjTcICwvDbrfzxBNP8M033/T0MIUQHSDBXiGEEEIIIYQQHpm37hm57vbv309JSQlGo5Ebbrih3bbV1dXs3LnTa9uid/vb3/6Goii97rZ0rVZLUFCQ6ycuLo6rr76alStXEhkZCcB7771HUVFRD4/01GQwGFyvbXBwMGlpadxzzz3Mnj3b1eaZZ57pwRH2PlOnTkVRFBRF6TVlT4QACfYKIYQQQgghhAAmTpyITqcD/AdwnSUeJkyYwDnnnOOxrK01a9Zgt9sBCfaKkycyMpL77rsPAKvVyrJly3p4RKeXqVOnMnXqVAD27NlDYWFhzw5ICHFMEuwVQgghhBBCCIHRaGTMmDGA/2Cvc/mUKVOYPHkyAGvXrsVms/ltC3D22Wd7Pd/Y2Mgbb7zBtGnTiImJQa/XExMTw0UXXcQ333yDoig+x9B2IrP8/HwefPBBBg4ciMlk8jlZUlFREY8//jhnnHEG4eHhGAwGUlJSuPXWW9m4caPf1+SOO+5ApVK5gl0bNmzglltuISUlBb1e75HNpygKGzdu5IknnuDMM88kMjISnU5HeHg4Y8eO5amnnqK8vNxvX23l5uZyzz330LdvXwICAoiJieGaa65h3bp1Hd7G7t27mT59OsnJyRgMBhITE7n99tvZt2+f33X8TdCWmppKWlqa6/dzzjnHa2KvnqrjO2TIENfj/Px8v+3MZjPPPPMMQ4cOxWQyERoaytSpU485eeCuXbv45z//ydSpU4mNjUWn0xESEsLw4cOZMWMGeXl57a5vNpv597//zZlnnklERAQ6nY7o6GgGDx7Mddddx3vvvYfZbPa7/vLly7nttttIS0sjMDCQkJAQRo4cyV//+lcqKyvb7bsrtPf6Ovdt4sSJREREuI6z6667joULF/rdZtv6z42Njfz73/9m7NixhIeHu+rzTp061WPSwjvvvNPruHOv49ud56y/CdqcYzhy5AgATz/9tNeYpY6vOKkUIYQQQgghhBBCUZRHHnlEARRAKSgo8Hq+T58+CqAsWbJEURRFiYuLUwBl48aNXm0nTpyoAEpaWprXc7t371bS0tJcffn6ueyyyxSz2ey17ocffuhqs2HDBiUiIsJr3ezsbFf7b775RjGZTO329cwzz/h8PaZPn64Aytlnn628/fbbikaj8VivT58+rrZz5sxptw9AiYuLU7Zt2+azL/f92rx5s8/9AhS1Wq289tprPrfx1FNPucY1b948JTAw0Oc29Hq9MmvWrGNuw53zvW/vZ9myZR7rOJefffbZPvvqCOc2pk+f7rfNrFmzXO3+9a9/uZYvW7bMtXzdunVKRkaG37H//e9/97nt7du3H3O/g4ODlUWLFvlc/+jRo0r//v2PuY1NmzZ5rdvY2Kjcdttt7a4XGxvrc92OcH99PvzwQ7/t7rvvPo9zzmnPnj1KSkpKu+O77bbbFIvF4rXN7OxsV5tZs2Ypw4cP91r3ww8/VM4+++xjvnbuY+/Oc9b99XP/m+McQ3s/Tz31VAffJSGOn2T2CiGEEEIIIYQA2q/bm5eXx5EjR9BqtUyYMAHAld3btm1DQwObN28GvLN68/PzmTp1KtnZ2fTp04d33nmHgwcPUlFRwb59+3j66acxGAz8+OOP3H333e2O95prrsFoNPLBBx+Qm5tLcXExP/30E+Hh4QDMnz+fG264AbPZzOTJk/nuu+/Iy8ujvLycjRs3cuuttwLw17/+lQ8//NBvP/v37+fee+9l3Lhx/PTTTxQXF5Obm8vrr7/uaqPVarn88st55513WL16NZmZmZSVlbF7927eeecdBgwYQFFREVdffTWNjY3t7td1112HVqvlf//7H3l5eRQVFfH111+TlpaG3W7n/vvvZ8mSJX7Xr66u5tZbbyU+Pp5vvvmGoqIi8vLyeO+994iOjsZisXDTTTexd+/edsfhbu/evezZs8f1+/z58z0mTKutrWXKlCkd3l5Xct+PxMREn21uueUWSktLeeONN8jKyqKsrIwlS5YwdOhQwJHR7CvjWaVSce655/LKK6+wYsUKDh48SHl5Ofv37+fzzz9nzJgx1NbWcsMNN/isF/zwww9z6NAhNBoNTz75JNu3b6ekpISSkhK2bdvGW2+95ZW96jR9+nQ+/fRTtFotf/zjH9mwYQNlZWUcPXqUb775hkGDBlFcXMyll156UmsV+3p9KyoqOP/888nNzcVgMPD0009z4MABysrKWLlyJRdccAEAn376KX/605/a3f4f//hHDh48yNNPP82+ffsoKytj48aNrvOttrbW1dZ9Ijnnj/M8dtfd56y7d955h9raWlJSUgD4y1/+4jXmxx57rMPbE+K49XS0WQghhBBCCCFE71BdXa2o1WoFUO6++26P5z777DMFUMaNG+da9sorryiAcsUVV3i0/fnnn10ZbB988IHHc1dccYUCKEOGDFEqKip8jmP+/PkeWa7u3DNgIyMjldzcXJ/baGhoUGJjYxVAueWWWxS73e6z3cMPP6wASkxMjNLQ0ODxnHuG3uTJk5Wmpiaf2+iI2tpaJT093edr0na/jEajsnfvXq82BQUFrn0aMmSI1/POrFxAiY+PVwoLC73a7N69WwkICFAA5eKLL/a7jbaZvYrimY3ZNovXF2fbk5nZW1VV5XpN1Gq1kpeX53rOPfMyKChI2bdvn9f6+fn5rgzoRx555LjH19zcrEyaNEkBlL/+9a9ezzsztB988MHj2u7s2bNd+/TTTz/5bFNVVeXKGr733nuPe+wdyexdu3atolKpFEAZMGCAa/kDDzzgWnfOnDle69lsNuXKK690tdm1a5fH8+7HEqDMmzev3bEea5xO3XnO+svsdXJmw0sWr+huktkrhBBCCCGEEALAVQsUvLN13ev1Ojkfr1692qPGrvu67tnC2dnZ/PDDDwC88sorrgzcti666CJXzc3PP//c73gfeughkpOTfT731VdfUVxcjNFo5M033/SZOQnw1FNPYTKZKCkpYdGiRX77evHFF9Hr9X6fP5agoCCuvvpqABYvXtxu23vuuYeMjAyv5QkJCTz++OOAY7KsDRs2+N3GE088QVxcnNfyIUOG8Ic//AGABQsWnNQJtxRFQVGUk1LL12w2s3jxYqZNm0ZxcTEAt912G0lJST7b33fffQwaNMhreWJiIr/61a8A2LRp03GPQ6vVcvPNNwO+31er1Qo43rvj8corrwBw0003ceGFF/psExoa6joevvjiC791rjujrKyMTz75hCuvvNK13b/85S8A2Gw2V53cSy65hCuuuMJrfbVazRtvvIFWqwXgvffe89vXhRdeyMUXX9xlY3fqznNWiN5Egr1CCCGEEEIIIVycwdm9e/d6TE60atUqwDPYO2LECIKDgykvL/e4xd8Z7E1MTCQ9Pd21fOnSpSiKQkBAAOPGjaOurs7vjzPo3F4A7pJLLvH7nLPMwcSJE9FoNH77sdvtriCgv74iIyMZP368376crFYrH3/8MZdffjkpKSkYjUaPSZleeOEFAA4cONDudpwBJl+uueYa1+M1a9ac0DbsdvtxTfjWkz7++GOP1zIoKIjzzz+fLVu2ADB16lTefPNNv+tfdNFFfp8bOHAgQLulEGbPns31119Penq6ayJA588999wD+H5fncfxCy+8Line truncated
</div>
</div>
</div>
@@ -2360,7 +2255,7 @@ Gradients [[ 0. ]
<div class="jp-InputPrompt jp-InputArea-prompt">In [ ]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="CodeMirror cm-s-jupyter">
<div class="zeroclipboard-container">
<clipboard-copy ,="" aria-label="Copy to Clipboard" for="cell-7">
<clipboard-copy ,="" aria-label="Copy to Clipboard" for="cell-6">
<div>
<span class="notice" hidden="">Copied!</span>
<svg aria-hidden="true" class="clipboard-copy-icon" data-view-component="true" height="20" version="1.1" viewbox="0 0 16 16" width="20">
@@ -2372,7 +2267,7 @@ Gradients [[ 0. ]
</div>
<div class="highlight-ipynb hl-python"><pre><span></span>
</pre></div>
<div class="clipboard-copy-txt" id="cell-7"></div>
<div class="clipboard-copy-txt" id="cell-6"></div>
</div>
</div>
</div>
@@ -3133,7 +3028,7 @@ span.linenos.special { color: #000000; background-color: #ffffc0; padding-left:
<div class="jp-Collapser jp-InputCollapser jp-Cell-inputCollapser">
</div>
<div class="jp-InputArea jp-Cell-inputArea">
<div class="jp-InputPrompt jp-InputArea-prompt">In [ ]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="jp-InputPrompt jp-InputArea-prompt">In [2]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="CodeMirror cm-s-jupyter">
<div class="zeroclipboard-container">
<clipboard-copy ,="" aria-label="Copy to Clipboard" for="cell-1">
@@ -3152,9 +3047,8 @@ span.linenos.special { color: #000000; background-color: #ffffc0; padding-left:
<span class="kn">import</span><span class="w"> </span><span class="nn">jax.random</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">jrand</span>
<span class="kn">import</span><span class="w"> </span><span class="nn">matplotlib</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">mpl</span>
<span class="kn">import</span><span class="w"> </span><span class="nn">matplotlib.pyplot</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">plt</span>
<span class="kn">from</span><span class="w"> </span><span class="nn">eventpropjax.evnn</span><span class="w"> </span><span class="kn">import</span> <span class="n">FFEvNN</span>
<span class="kn">from</span><span class="w"> </span><span class="nn">felice.neuron_models</span><span class="w"> </span><span class="kn">import</span> <span class="n">FHN</span>
<span class="kn">from</span><span class="w"> </span><span class="nn">felice.neuron_models</span><span class="w"> </span><span class="kn">import</span> <span class="n">FHNRS</span>
</pre></div>
<div class="clipboard-copy-txt" id="cell-1">import diffrax as dfx
import jax
@@ -3162,9 +3056,8 @@ 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 FHN</div>
from felice.neuron_models import FHNRS</div>
</div>
</div>
</div>
@@ -3176,7 +3069,7 @@ from felice.neuron_models import FHN</div>
<div class="jp-Collapser jp-InputCollapser jp-Cell-inputCollapser">
</div>
<div class="jp-InputArea jp-Cell-inputArea">
<div class="jp-InputPrompt jp-InputArea-prompt">In [2]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="jp-InputPrompt jp-InputArea-prompt">In [23]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="CodeMirror cm-s-jupyter">
<div class="zeroclipboard-container">
<clipboard-copy ,="" aria-label="Copy to Clipboard" for="cell-2">
@@ -3192,32 +3085,64 @@ from felice.neuron_models import FHN</div>
<div class="highlight-ipynb hl-python"><pre><span></span><span class="n">key</span> <span class="o">=</span> <span class="n">jrand</span><span class="o">.</span><span class="n">key</span><span class="p">(</span><span class="mi">0</span><span class="p">)</span>
<span class="n">max_time</span> <span class="o">=</span> <span class="mi">200</span>
<span class="n">snn</span> <span class="o">=</span> <span class="n">FFEvNN</span><span class="p">(</span>
<span class="n">layers</span><span class="o">=</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span>
<span class="n">in_size</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
<span class="n">neuron_model</span><span class="o">=</span><span class="n">FHN</span><span class="p">,</span>
<span class="n">solver</span><span class="o">=</span><span class="n">dfx</span><span class="o">.</span><span class="n">Dopri5</span><span class="p">(),</span>
<span class="n">max_solver_time</span><span class="o">=</span><span class="n">max_time</span><span class="p">,</span>
<span class="n">key</span><span class="o">=</span><span class="n">key</span><span class="p">,</span>
<span class="n">max_event_steps</span><span class="o">=</span><span class="mi">1000000</span><span class="p">,</span>
<span class="n">solver_stepsize</span><span class="o">=</span><span class="mf">0.1</span><span class="p">,</span>
<span class="n">init_weights</span><span class="o">=</span><span class="mf">2.0</span><span class="p">,</span>
<span class="n">neuron_model</span> <span class="o">=</span> <span class="n">FHNRS</span><span class="p">(</span>
<span class="n">gmax_pasive</span><span class="o">=</span><span class="mf">2.0</span><span class="p">,</span>
<span class="n">Erev_pasive</span><span class="o">=</span><span class="mf">0.0</span><span class="p">,</span>
<span class="n">a_fast</span><span class="o">=-</span><span class="mf">2.0</span><span class="p">,</span>
<span class="n">voff_fast</span><span class="o">=</span><span class="mf">0.0</span><span class="p">,</span>
<span class="n">tau_fast</span><span class="o">=</span><span class="mf">0.0</span><span class="p">,</span>
<span class="n">a_slow</span><span class="o">=</span><span class="mf">0.5</span><span class="p">,</span>
<span class="n">voff_slow</span><span class="o">=</span><span class="mf">1.0</span><span class="p">,</span>
<span class="n">tau_slow</span><span class="o">=</span><span class="mf">50.0</span><span class="p">,</span>
<span class="n">vthr</span><span class="o">=</span><span class="n">jnp</span><span class="o">.</span><span class="n">inf</span><span class="p">,</span>
<span class="p">)</span>
<span class="k">def</span><span class="w"> </span><span class="nf">state_at_t</span><span class="p">(</span><span class="n">comp_times</span><span class="p">):</span>
<span class="n">sol</span> <span class="o">=</span> <span class="n">dfx</span><span class="o">.</span><span class="n">diffeqsolve</span><span class="p">(</span>
<span class="n">terms</span><span class="o">=</span><span class="n">dfx</span><span class="o">.</span><span class="n">ODETerm</span><span class="p">(</span><span class="n">neuron_model</span><span class="o">.</span><span class="n">dynamics</span><span class="p">),</span>
<span class="n">solver</span><span class="o">=</span><span class="n">dfx</span><span class="o">.</span><span class="n">Tsit5</span><span class="p">(),</span>
<span class="n">t0</span><span class="o">=</span><span class="mf">0.0</span><span class="p">,</span>
<span class="n">t1</span><span class="o">=</span><span class="n">max_time</span><span class="p">,</span>
<span class="n">dt0</span><span class="o">=</span><span class="mf">1e-3</span><span class="p">,</span>
<span class="n">y0</span><span class="o">=</span><span class="n">neuron_model</span><span class="o">.</span><span class="n">init_state</span><span class="p">(</span><span class="mi">1</span><span class="p">)</span>
<span class="o">+</span> <span class="n">jrand</span><span class="o">.</span><span class="n">uniform</span><span class="p">(</span><span class="n">key</span><span class="p">,</span> <span class="n">shape</span><span class="o">=</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">3</span><span class="p">),</span> <span class="n">minval</span><span class="o">=</span><span class="mf">0.1</span><span class="p">,</span> <span class="n">maxval</span><span class="o">=</span><span class="mf">0.5</span><span class="p">),</span>
<span class="n">saveat</span><span class="o">=</span><span class="n">dfx</span><span class="o">.</span><span class="n">SaveAt</span><span class="p">(</span><span class="n">ts</span><span class="o">=</span><span class="n">comp_times</span><span class="p">),</span>
<span class="n">max_steps</span><span class="o">=</span><span class="mi">200000</span><span class="p">,</span>
<span class="p">)</span>
<span class="k">return</span> <span class="n">sol</span><span class="o">.</span><span class="n">ts</span><span class="p">,</span> <span class="n">sol</span><span class="o">.</span><span class="n">ys</span>
</pre></div>
<div class="clipboard-copy-txt" id="cell-2">key = jrand.key(0)
max_time = 200
snn = FFEvNN(
layers=[1],
in_size=1,
neuron_model=FHN,
solver=dfx.Dopri5(),
max_solver_time=max_time,
key=key,
max_event_steps=1000000,
solver_stepsize=0.1,
init_weights=2.0,
)</div>
neuron_model = FHNRS(
gmax_pasive=2.0,
Erev_pasive=0.0,
a_fast=-2.0,
voff_fast=0.0,
tau_fast=0.0,
a_slow=0.5,
voff_slow=1.0,
tau_slow=50.0,
vthr=jnp.inf,
)
def state_at_t(comp_times):
sol = dfx.diffeqsolve(
terms=dfx.ODETerm(neuron_model.dynamics),
solver=dfx.Tsit5(),
t0=0.0,
t1=max_time,
dt0=1e-3,
y0=neuron_model.init_state(1)
+ jrand.uniform(key, shape=(1, 3), minval=0.1, maxval=0.5),
saveat=dfx.SaveAt(ts=comp_times),
max_steps=200000,
)
return sol.ts, sol.ys</div>
</div>
</div>
</div>
@@ -3229,7 +3154,7 @@ snn = FFEvNN(
<div class="jp-Collapser jp-InputCollapser jp-Cell-inputCollapser">
</div>
<div class="jp-InputArea jp-Cell-inputArea">
<div class="jp-InputPrompt jp-InputArea-prompt">In [3]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="jp-InputPrompt jp-InputArea-prompt">In [24]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="CodeMirror cm-s-jupyter">
<div class="zeroclipboard-container">
<clipboard-copy ,="" aria-label="Copy to Clipboard" for="cell-3">
@@ -3243,9 +3168,9 @@ snn = FFEvNN(
</clipboard-copy>
</div>
<div class="highlight-ipynb hl-python"><pre><span></span><span class="n">v_range</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">arange</span><span class="p">(</span><span class="o">-</span><span class="mf">3.1</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mf">0.1</span><span class="p">)</span>
<span class="n">VI_inst</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">vmap</span><span class="p">(</span><span class="n">snn</span><span class="o">.</span><span class="n">neuron_model</span><span class="o">.</span><span class="n">IV_inst</span><span class="p">)(</span><span class="n">v_range</span><span class="p">)</span>
<span class="n">VI_fast</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">vmap</span><span class="p">(</span><span class="n">snn</span><span class="o">.</span><span class="n">neuron_model</span><span class="o">.</span><span class="n">IV_fast</span><span class="p">)(</span><span class="n">v_range</span><span class="p">)</span>
<span class="n">VI_slow</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">vmap</span><span class="p">(</span><span class="n">snn</span><span class="o">.</span><span class="n">neuron_model</span><span class="o">.</span><span class="n">IV_slow</span><span class="p">)(</span><span class="n">v_range</span><span class="p">)</span>
<span class="n">VI_inst</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">vmap</span><span class="p">(</span><span class="n">neuron_model</span><span class="o">.</span><span class="n">IV_inst</span><span class="p">)(</span><span class="n">v_range</span><span class="p">)</span>
<span class="n">VI_fast</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">vmap</span><span class="p">(</span><span class="n">neuron_model</span><span class="o">.</span><span class="n">IV_fast</span><span class="p">)(</span><span class="n">v_range</span><span class="p">)</span>
<span class="n">VI_slow</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">vmap</span><span class="p">(</span><span class="n">neuron_model</span><span class="o">.</span><span class="n">IV_slow</span><span class="p">)(</span><span class="n">v_range</span><span class="p">)</span>
<span class="k">with</span> <span class="n">mpl</span><span class="o">.</span><span class="n">style</span><span class="o">.</span><span class="n">context</span><span class="p">(</span><span class="s2">"boilerplot.ieeetran"</span><span class="p">):</span>
<span class="n">fig</span><span class="p">,</span> <span class="n">ax</span> <span class="o">=</span> <span class="n">plt</span><span class="o">.</span><span class="n">subplots</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="n">figsize</span><span class="o">=</span><span class="p">(</span><span class="mf">6.9</span><span class="p">,</span> <span class="mf">2.3</span><span class="p">),</span> <span class="n">dpi</span><span class="o">=</span><span class="mf">200.0</span><span class="p">,</span> <span class="n">sharey</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
@@ -3255,9 +3180,9 @@ snn = FFEvNN(
<span class="n">plt</span><span class="o">.</span><span class="n">show</span><span class="p">()</span>
</pre></div>
<div class="clipboard-copy-txt" id="cell-3">v_range = jnp.arange(-3.1, 3, 0.1)
VI_inst = jax.vmap(snn.neuron_model.IV_inst)(v_range)
VI_fast = jax.vmap(snn.neuron_model.IV_fast)(v_range)
VI_slow = jax.vmap(snn.neuron_model.IV_slow)(v_range)
VI_inst = jax.vmap(neuron_model.IV_inst)(v_range)
VI_fast = jax.vmap(neuron_model.IV_fast)(v_range)
VI_slow = jax.vmap(neuron_model.IV_slow)(v_range)
with mpl.style.context("boilerplot.ieeetran"):
fig, ax = plt.subplots(1, 3, figsize=(6.9, 2.3), dpi=200.0, sharey=True)
@@ -3276,7 +3201,7 @@ with mpl.style.context("boilerplot.ieeetran"):
<div class="jp-OutputArea-child">
<div class="jp-OutputPrompt jp-OutputArea-prompt"></div>
<div class="jp-RenderedImage jp-OutputArea-output" tabindex="0">
<img alt="No description has been provided for this image" class="" src="data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAABXsAAAHjCAYAAABy0synAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjgsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvwVt1zgAAAAlwSFlzAAAewgAAHsIBbtB1PgAAmANJREFUeJzs3QV0XMfZxvFHzJZBtszMTAInDjMzmEmOG06aNs0XaANN26RJGibLTInDzAy2JNsyM7NkGQQrhu/ckbVx2I60urt7/79zdLzvdbqajmS92mfnzgRUVVVVCQAAAAAAAADg0wLtHgAAAAAAAAAAoPYIewEAAAAAAADADxD2AgAAAAAAAIAfIOwFAAAAAAAAAD9A2AsAAAAAAAAAfoCwFwAAAAAAAAD8AGEvAAAAAAAAAPgBwl4AAAAAAAAA8AOEvQAAAAAAAADgBwh7AQAAAAAAAMAPEPYCAAAAAAAAgB8g7AUAAAAAAAAAP0DYCwAAAAAAAAB+gLAXAAAAAAAAAPwAYS8AAAAAAAAA+AHCXgAAAAAAAADwA4S9AAAAAAAAAOAHguUwlZWVysnJMY8jIyMVEBBg95AAAPhVVVVVKiwsNI/j4uIUGHhs79PS9wAATul5FvoeAMBJfe+XOC7stYLe+Ph4u4cBAMAxy8rKUrNmzY7pf0PfAwA4pedZ6HsAACf1vV/CNg4AAAAAAAAA4Acct7LX2rrhyNQ8KirKI5/H5XK5VxB78vMwBuaA7wP+LfDzwL9/Jh75vzmyhx0t+l798dXvMX/6/N4wBrs/vzeMwe7P7w1jsPvze8MY7Oh5P/3f8XrPv7/HvGEMdn9+bxiD3Z/fG8Zg9+f3hjHY/fmd3Pd+iePC3iP36LUmvj6++PX1eRgDc8D3Af8W+Hng3z8T/8g+8/Q9e/jq95g/fX5vGIPdn98bxmD35/eGMdj9+b1hDPXV8376v+P1nnO+x7xhDHZ/fm8Yg92f3xvGYPfn94Yx2P35ndb3fgnbOAAAAAAAAACAH3Dcyt76YiX41ql6Th+D3bxhDuweg92f3xt4wxzYPQa7P7838IY58IYx+PP/N28Yg93sngO7P7+3jMFuds+B3Z/fW8ZgN7vnwO7P74T/f94wBrvZPQd2f35vGYPd7J4Duz+/t4zBbnbPQZQXfQ18YmXvnDlzNHr0aPXr18+cTBcSEqLY2FglJibq3//+twoKCuweIgAAAAAAAADYyidW9j733HP6/vvv1aNHDw0cOFCNGzc2mx0vWLBAGRkZmjp1qr766iu1bNnS7qECAAAAAAAAgC18Iux99NFH1aVLFxPyHmn//v26+OKL9e233+q2227TvHnzbBsjAAAAAAAAANjJJ8LepKSkX7zepEkT/etf/9KJJ56ojz/+uN7HBd/aswT24fsAfB/ACfhZB74PwM8DOAl9D3wfgJ8H3skn9uz9LcHB1Xl1WFiY3UMBAAAAAAAAANv4xMreX5Ofn697773XPL7wwguP+X/vcrl+811KAADqy6/1pN/qVXX1OSz0PQCAP/W833s++h4AwN/6nk+GvdZWDXPnzlVlZaX7gDYr8D377LP10EMPHfPzxcfH/+rfsfUAAKA+RUdHe/xz0PcAAE7peRb6HgDASX3PJ8Pe1atXa8aMGT+6Nnz4cD322GOKjY21bVwAAAAAAAAAYLeAKh9cwlpWVqbt27frrbfe0j//+U8FBATojTfeMAe1/R5riXRNom6tDv6123e4rQcA4C239tSsTCooKDjm/kTfAwA4pefVPAev9wAATul7fhP2HiktLU1DhgxR69attW7dOkVERBx186/LiQQAwBNq27foewAAX1EXPYu+BwDwFZ7qWYHycUlJSerZs6d27NihRYsW2T0cAAAAAAAAALCFz4e9lprkOzs72+6hAACgvOIyrdyV654JH7+JBgAAAABw2KrducotKpO38vmwNycnR8uWLTOPu3btavdwAAAOd6iwVCNT0zTsxYVasTOXoBcAAAAA/ETG1gO66oWFGjctXQUl5V75es/rw97Vq1drzpw5Ki4u/tnfrV+/XldccYVKSkqUnJysPn362DJGAAAsOQUluvrFhVq+M1f5JeW65eVMVVbJHCQKAAAAAPBd323M0egp1SHvku2H9NRnG+SNguXlrK0ZRo4cqUmTJmnAgAHmILbS0lJt375dS5YsUWVlpXr06KGXX37Z7qECABwsK69YI1LTtDG7wNRNY8L03MhBCgok6AUAAAAAX/bF2mxNmr1YpeWVpj6hS5xuOb2rVy7s8fqwt1evXnrwwQf1zTffaO3atcrMzFRZWZkaN26s0047TZdeeqnGjRunsLAwu4cKAHConQcLTdC7bX+hqVvEhmvuxGR1iKub01QBAAAAAPb4cOVe3ThvicoqqrdsOL1HvJ4ZMUBhwUFe+SXx+rC3adOmuvPOO+0eBgAAv2jbfpeGT07TrkNFpm7TOEJzU5LVpnEkMwYAAAAAPuytpbv05/nLVGHtzyfpvD4t9PjV/RUS5L0743p92AsAgLeytmwYkbpQWXklpu4YF6U5E5PUIjbC7qEBAAAAAGph/qId+ttry1VzBtulA1vp4cv6KtiLg145Pex1uVzux1FR3GoLADh6a/bkaWRqmva7Sk3dNT5as1OS1CwmvE571JG9qq6e00LfAwB4C0/0vJ8+H30PAHAsZi3YqnveWuWuhyW21YMX91ZgHZzJ4qm+V8PRYW98fLz7cVVNTA8AwO9YsTNXo6am6VBhmal7tWygWROS1DgqtM7mLjo6us6/DvQ9AIA38kTPs9D3AAB/ROo3m/XP99a467HHtdc/LuhZZ4exearv1XB02AsAwLFavO2Axk7NUH5Juan7t2moGeMTFRsRwmQCAAAAgA97+vMNeuTj9e762pM76fazutVZ0FsfHB32ZmVlcTsPAOCoLdi0XxNmZKiwtMLUiR0aa+rYBEWH1X07LSgocN/ac+TKpNqg7wEAvJEnep6FvgcAOFrWHf+PfrxeT3+x0X3t1tO76qbTOtd50OupvlfD0WGvtW8TezcBAI7GV+v36ZqZi1RSXmnqE7rE6cVRgxURGuSxHuWJ56TvAQC8jad6E30PAHC0Qe+D761R6rdb3Nf+75zumnRSJ3mCp1+TOTrsBQDgaHyyOkvXz1mi0orqoPe07s30zIiBCg/xTNALAAAAAPC8ysoq/f3tlZq9cLv72n0X9tKY49r77PQT9gIA8BveW75HN7+UqfLK6oM8z+ndXE9cPUChwYHMGwAAAAD4qIrKKv3tteV6dfFOU1u7Nfz7kj66OrGtfBlhLwAAv+L1JTv1l1eW6XDOq4v7t9QjV/RTcBBBLwAAAAD4qrKKSv15/jK9s2y3qYMCA/ToFf108YBW8nWEvQAA/IK5adt115srVHU46L06oY0evKSP+SUAAAAAAOCbSsordNO8TH20KsvUwYEBemrYAJ3Tp4X8AWEvAAA/Me27LbrvndXuesyQdvrHBb0USNALAAAAAD6ruKxC185erC/W7TN1aFCgnhs5UKf1iJe/IOwFAOAIz325SQ99uNZdTzqxo+44p7sCrA2cAAAAAAA+qbC0XCkzFun7TftNHR4SqMmjB+uELk3lTwh7AQCQVFVVpcc/3aAnPtvgno+bTuuiW0/vQtALAAAAAD4sv7hM46ZlaNG2g6aOCg3SlLEJSu7YRP7G0WGvy+VyP46KirJ1LAAAe4Pe/3y4Vi98tdl97fazu+m6kzvb3qOO7FV19ZwW+h4AwFt4ouf99PnoewDgXIcKSzVmarqW7cw1dUx4sGaMT9TAto38qu/VcHTYGx8f/6MX+gAA56msrNL9767W9O+3uq/9/fyeGj+0g63jio6OrvPnpO8BALyRJ3qehb4HANhfUKKRU9K1Zk+emYxGkSGaNSFJvVvF+l3fq+HosBcA4GwVlVW6640Veiljh6mtbXkfvLiPhie1tXtoAAAAAIBayM4r1ojUNG3ILjB1XHSY5qQkqVvzGL+eV0eHvVlZWdzOAwAOVV5Rqb++ulxvZO4ydWCA9N/L++myQa3tHppRUFDgvrXnyJVJtUHfAwB4I0/0PAt9DwCca9ehIo2YvFBb9xeaunmDcM2ZmKROTT27qtbOvlfD0WGvtW8TezcBgPOUllfq5pcy9cHKvaYODgzQ/67qrwv6tZS38ER/ou8BALyRp16T0fcAwJm27y/UsMkLTeBradUwQvMmJqttk0h5A09nkY4OewEAzlNcVqHr5yzRZ2uzTR0aFKinhw/Qmb2a2z00AAAAAEAtbNpXoBGT07Q3r9jU7ZtEau7EZLVsGOGYeSXsBQA4RlFpha6ZtUjfbMgxdVhwoF4cPVgndW1q99AAAAAAALWwdm+eRqamKaeg1NRdmkWbPXqbNQh31LwS9gIAHKGgpFzjp2cofcsBU0eGBil1zGAd1ynO7qEBAAAAAGph5a5cjZySpkOFZabu2aKBZk1IVJPoMMfNK2EvAMDv5RaVaczUdC3dccjUMWHBmj4+QYPaNbZ7aAAAAACAWliy/aB5vZdfXG7qfm0aaua4RMVGhjhyXgl7AQB+7YCrVKOmpGnV7jxTx0aEmHd4+7ZuaPfQAAAAAAC1sHDzfk2YniFXaYWpE9o30tSxCYoJd2bQayHsBQD4rez8Yo1KTde6rHxTN4kK1eyUJPVo0cDuoQEAAAAAauHr9fvMmSzFZZWmPr5zE00ePViRoc6OO539/x4A4Lf25BaZU1g357hM3SwmTHMnJqlzsxi7hwYAAAAAqIVPV2fpujlLVFpRHfSe0q2pnhs5SOEhQY6fV8JeAIDf2XGgUMNTF2rHgSJTt2oYYU5hbR8XZffQAAAAAAC18N7yPbr5pUyVV1aZ+uxezfXksAEKDQ5kXp2+stflql7tZYmKIgAALine truncated
<img alt="No description has been provided for this image" class="" src="data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAABXsAAAHjCAYAAABy0synAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjgsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvwVt1zgAAAAlwSFlzAAAewgAAHsIBbtB1PgAAkjRJREFUeJzs3Xd4VHXaxvF7Jj2TRggkQOi9E0jBrth7R5EOsa7d1XXXdV3L+uq69q6hB1TsvfdGCoTee2gJIaRN+sx5/wgZQEEpMzlTvp/r2suZ34Qzt5HNk/PMr1gMwzAEAAAAAAAAAPBpVrMDAAAAAAAAAACOHs1eAAAAAAAAAPADNHsBAAAAAAAAwA/Q7AUAAAAAAAAAP0CzFwAAAAAAAAD8AM1eAAAAAAAAAPADNHsBAAAAAAAAwA/Q7AUAAAAAAAAAP0CzFwAAAAAAAAD8AM1eAAAAAAAAAPADNHsBAAAAAAAAwA/Q7AUAAAAAAAAAP0CzFwAAAAAAAAD8AM1eAAAAAAAAAPADNHsBAAAAAAAAwA/Q7AUAAAAAAAAAP0CzFwAAAAAAAAD8AM1eAAAAAAAAAPADNHsBAAAAAAAAwA/4VLO3vr5ezzzzjI4//njFx8crPDxcycnJOvvss/XGG2+YHQ8AAAAAAAAATBNsdoBDtWXLFp155plavny5EhISdNxxx8lms6mwsFA//PCDbDabrrjiCrNjAgAAAAAAAIApfKLZW1NTo9NPP10rV67Uv//9b/3jH/9QSEiI6/Xq6mqtXr3axIQAAAAAAAAAYC6faPb+3//9n1auXKlrrrlG99133+9ej4yM1JAhQ1o+GAAAAAAAAAB4CYthGIbZIf5IQ0OD2rdvr5KSEq1Zs0Y9evQwOxIAAAAAAAAAeB2vn9m7YMEClZSUqH379urRo4eWLFmid955R9u2bVOrVq10wgkn6Oyzz5bVemhnzTmdTpWUlEhqmhFssVg8GR8AgKNiGIaqq6slSQkJCYdc75pR9wAAvuJoa55E3QMA+A531L0D8fpm7+LFiyVJycnJuvvuu/Xf//5X+05GfvTRR5WSkqL33ntPnTp1+tPrlZSUKDEx0WN5AQDwlKKiIrVt2/aw/gx1DwDgi46k5knUPQCAbzrSuncg7mkZe9CuXbskSQUFBXr00Ud1ww03aNWqVSovL9eXX36pXr16qaCgQOeee64aGhpMTgsAAAAAAAAA5vD6mb3Ns3gbGho0atQoPffcc67XTjvtNH355Zfq3bu3li5dqtdff11jx479w+tFRkYe0vtWVVUdeWhJdrvd9YlyUVGRbDbbUV2PDL75/t6Qwez394YMZr+/N2Qw+/29IYPZ7/9nGaKiov70zx9qDTuSP0Pd8/3394YMZr+/N2Qw+/29IYPZ7+8NGcx+f2/IYEbNO5w/R93z/ff3hgxmv783ZDD7/b0hg9nv7w0ZzH5/b8hgVt07EK9v9kZHR7seX3vttb97vVOnTjr33HP19ttv66uvvvrTZu++ezb90X98d/6lsNlspvxFJ4N3vb83ZDD7/b0hg9nv7w0ZzH5/b8hg9vsfKMPBbjr3/aXhSPYdpO4F5vt7Qwaz398bMpj9/t6Qwez394YMZr+/N2RoqZr32z9H3Quc9/eGDGa/vzdkMPv9vSGD2e/vDRnMfn9vyNCSde9AvL7Z261btwM+PtDXbN++/bCubfZ/fAAAmrVEPaLuAQC8QUvVIuoeAMAbtHQt8vpm79ChQ2WxWGQYhkpKStSxY8fffU3zaauHMi26pdhstv0OkgvUDGbzhu+B2RnMfn9v4A3fA7MzmP3+3sAbvgfekMFTvOHfzRsymM3s74HZ7+8tGcxm9vfA7Pf3lgxmM/t7YPb7e5o3/Pt5Qwazmf09MPv9vSWD2cz+Hpj9/t6SwWxmfw/Mfv99ef0BbUlJSTr++OMlSV999dXvXm9oaND3338vSUpPT2/RbAAAAAAAAADgLby+2StJ9913nyTp//7v/zRv3jzXeGNjo+644w6tX79e0dHRmjhxolkRAQAAAAAAAMBUXr+NgySdeuqpevDBB3XvvffqhBNOUHp6upKSkrRgwQJt3LhREREReu2111ybGgMAAAAAAABAoLEY3rKhxCH44osv9NRTTyknJ0eVlZVKSkrSqaeeqr/97W/q06fPIV3Dbre79vatqqpiw34AgFc72rpF3QMA+Ap31CzqHgDAV3iqZvnEzN5mZ5xxhs444wyzYwAAAAAAAACA1/GJPXsBAAAAAAAAAH+MZi8AAAAAAAAA+AGavQAAAAAAAADgB2j2AgAAAAAAAIAf8KkD2tzNbre7HnNKKwDAmzTXqH1rlbuuKVH3AADewxM177fXo+4BALyFp+pes4Bu9iYmJroeG4ZhYhIAAPYXFRXl9mtS9wAA3sgTNU+i7gEAvJOn6l4ztnEAAMDNahscWr+zyuwYAAAAAIAAE9Aze4uKiljOAwBwK3tdozJn5Gt1UaXmXneMurc5sk9tq6qamsV2u32/mUlHg7oHAPBGnqh5EnUPAOCdPFX3mgV0s9dms1H8AQBuU17doAnTc1WwuUySdM3MfH1x20kKsloO+1qeqE/UPQCAN/JUbaLuAQC8kadrU0A3ewEAcJedlXUaOyVHK3dUSpJiI0L0v8sHH1GjFwAAAACAI0GzFwCAo7StrEZjsnK0vqTpNNWEqFDNmpyhvu1iTE4GAAAAAAgkNHsBADgKG0vsGp2Vo61lNZKk9rHhys7MULcj3KsXAAAAAIAjRbMXAIAjtGpHpcZMydHOyjpJUpfWkcrOzFByq0iTkwEAAAAAAhHNXgAAjsDiLWUaNzVXZdUNkqTeidGalZmuttHhJicDAAAAAAQqmr0AABym3A2lmjQ9T1V1jZKkwcmxmjEpXXGRoSYnAwAAAAAEMpq9AAAchu9WFeu67PmqbXBKktK7xmvK+FRFh4eYnAwAAAAAEOho9gIAcIg+XbJdN79eoAaHIUk6uXcbvTh6mCJCg0xOBgAAAAAAzV4AAA7J2/O36M63FsnZ1OfV2QOS9PSVKQoNtpobDAAAAACAPWj2AgDwJ2b9ulH3vr/M9fyyYcl65JKBCg6i0QsAAAAA8B4B3ey12+2uxzabzcQkAABv9eJ36/ToZytdz8cf01n3nd9fVqvFo+/bXKP2rVXuuqZE3QMAeA9P1LzfXo+6BwDwFp6qe80CutmbmJjoemwYholJAADexjAMPfb5Kr3w3TrX2A0nd9edZ/aWxeLZRq8kRUVFuf2a1D0AgDfyRM2TqHsAAO/kqbrXLKCbvQAAHIjTaej+D5dpxq+bXGN3ntlbfzmlh4mpAAAAAAD4YwHd7C0qKmI5DwBgP40Op+5+Z4nemr/FNfbAhf017pguLZqjqqpKUtPSnn1nJh0N6h4AwBt5ouZJ1D0AgHfyVN1rFtDNXpvNRvEHALjUNzp16xsF+mTJDkmS1SL997LBumxYcotn8UR9ou4BALyRp2oTdQ8A4I08XZsCutkLAECzmnqHrsuer+9X75QkhQRZ9MyVKTp7YDuTkwEAAAAAcGho9gIAAl5lbYMmz8hX7oZSSVJ4iFUvjRmmk3u3NTkZAAAAAACHjmYvACCg7bbXa8K0XC3aUi5JigoL1tQJaUrvGm9yMgAAAAAADg/NXgBAwCqurNXYrFytKqqUJMVFhmjmpHQNSo4zNxgAAAAAAEeAZi8AICBt2V2tMVk52rirWpLUJjpM2ZMz1Dsp2uRkAAAAAAAcGZq9AICAs35nlcZk5Whbea0kqUNchGZnZqhLAid2AwAAAAB8F81eAEBAWbG9QmOn5Kikql6S1C3BpuzMDLWPizA5GQAAAAAAR4dmLwAgYBRs3q3xU3NVUdsoSerbLkazJqcrISrM5GQAAAAAABw9mr0AgIDw67pdypyRJ3u9Q5KU0ilO0yekKzYyxORkAAAAAAC4B81eAIDf+2Zlka7PXqC6Rqck6djurfXquFTZwiiDAAAAAAD/wV0uAMCvfbR4m259faEanYYk6dQ+bfX86KEKDwkyORkAAAAAAO4V0M1eu93uemyzcQI7APibuXmFuvudxdrT59V5g9rpySuGKCTIam6wQ9Bco/atVe66pkTdAwB4D0/UvN9ej7oHAPAWnqp7zQK62ZuYmOh6bBiGiUkAAO429acNeuCj5a7nV6R21MOXDFSQ1WJiqkMXFRXl9mtS9wAA3sgTNU+i7gEAvJOn6l6zgG72AgD8j2EYeu6btXr8y9WusUnHddW95/WVxeIbjV4AAAAAAI5EQDd7i4qKWM4DAH7EMAw98ulKvfzDetfYzaf21G2n9fS5Rm9VVZWkpqU9+85MOhrUPQCAN/JEzZOoewAA7+SputcsoJu9NpuN4g8AfsLpNHTv+0s1O2eza+wf5/TRNSd2NzHVkfNEfaLuAQC8kadqE3UPAOCNPF2bArrZCwDwD40Op+58a7HeLdgqSbJYpAcvHKAxwzubnAwAAAAAgJbj/ceRAwDwB+oaHfrLnAWuRm+Q1aInRw6h0QsAAAAAcLtPl2zXxhK72TEOimYvAMBnVdc3KnNGvj5fViRJCg2y6sXRQ3VRSgeTkwEAAAAA/M3cvELdMGeBRmflaMvuarPjHBDNXgCAT6qobdC4Kbn6cU2JJCkiJEhTJqTqjP5JJicDAAAAAPibGb9s1F1vL5ZhSFvLajQ3f4vZkQ6IPXsBAD6n1F6vcVNztHRrhSQpOixY0yamKbVLvMnJAAAAAAD+5sXv1unRz1a6nk88rotuO62niYkOzmdn9t51112yWCyyWCx66KGHzI4DAGghRRW1uuLlX12N3nhbqF67ZjiNXgAAAACAWxmGoSe+WLVfo/fGU3roX+f1k8ViMTHZwfnkzN5ffvlFjz/+uCwWiwzDMDsOAKCFFJZWa3RWLine truncated
</div>
</div>
</div>
@@ -3288,7 +3213,7 @@ with mpl.style.context("boilerplot.ieeetran"):
<div class="jp-Collapser jp-InputCollapser jp-Cell-inputCollapser">
</div>
<div class="jp-InputArea jp-Cell-inputArea">
<div class="jp-InputPrompt jp-InputArea-prompt">In [4]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="jp-InputPrompt jp-InputArea-prompt">In [25]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="CodeMirror cm-s-jupyter">
<div class="zeroclipboard-container">
<clipboard-copy ,="" aria-label="Copy to Clipboard" for="cell-4">
@@ -3301,27 +3226,23 @@ with mpl.style.context("boilerplot.ieeetran"):
</div>
</clipboard-copy>
</div>
<div class="highlight-ipynb hl-python"><pre><span></span><span class="n">in_spikes</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">asarray</span><span class="p">([[</span><span class="mf">0.00</span><span class="p">]])</span>
<span class="n">comp_times</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">linspace</span><span class="p">(</span><span class="mf">0.0</span><span class="p">,</span> <span class="n">max_time</span><span class="p">,</span> <span class="mi">500</span><span class="p">)</span>
<span class="n">state</span> <span class="o">=</span> <span class="n">snn</span><span class="o">.</span><span class="n">state_at_t</span><span class="p">(</span><span class="n">in_spikes</span><span class="p">,</span> <span class="n">comp_times</span><span class="p">)</span>
<span class="n">spikes</span> <span class="o">=</span> <span class="n">snn</span><span class="o">.</span><span class="n">spikes_until_t</span><span class="p">(</span><span class="n">in_spikes</span><span class="p">,</span> <span class="n">max_time</span><span class="p">)</span>
<div class="highlight-ipynb hl-python"><pre><span></span><span class="n">comp_times</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">linspace</span><span class="p">(</span><span class="mf">0.0</span><span class="p">,</span> <span class="n">max_time</span><span class="p">,</span> <span class="mi">500</span><span class="p">)</span>
<span class="n">_</span><span class="p">,</span> <span class="n">state</span> <span class="o">=</span> <span class="n">state_at_t</span><span class="p">(</span><span class="n">comp_times</span><span class="p">)</span>
</pre></div>
<div class="clipboard-copy-txt" id="cell-4">in_spikes = jnp.asarray([[0.00]])
comp_times = jnp.linspace(0.0, max_time, 500)
state = snn.state_at_t(in_spikes, comp_times)
spikes = snn.spikes_until_t(in_spikes, max_time)</div>
<div class="clipboard-copy-txt" id="cell-4">comp_times = jnp.linspace(0.0, max_time, 500)
_, state = state_at_t(comp_times)</div>
</div>
</div>
</div>
</div>
</div>
</div><div class="jp-Cell jp-CodeCell jp-Notebook-cell">
<div class="jp-Cell jp-CodeCell jp-Notebook-cell" id="cell-id=b73fb06d">
</div><div class="jp-Cell jp-CodeCell jp-Notebook-cell jp-mod-noOutputs">
<div class="jp-Cell jp-CodeCell jp-Notebook-cell jp-mod-noOutputs" id="cell-id=06ae6a94">
<div class="jp-Cell-inputWrapper" tabindex="0">
<div class="jp-Collapser jp-InputCollapser jp-Cell-inputCollapser">
</div>
<div class="jp-InputArea jp-Cell-inputArea">
<div class="jp-InputPrompt jp-InputArea-prompt">In [5]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="jp-InputPrompt jp-InputArea-prompt">In [26]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="CodeMirror cm-s-jupyter">
<div class="zeroclipboard-container">
<clipboard-copy ,="" aria-label="Copy to Clipboard" for="cell-5">
@@ -3334,16 +3255,152 @@ spikes = snn.spikes_until_t(in_spikes, max_time)</div>
</div>
</clipboard-copy>
</div>
<div class="highlight-ipynb hl-python"><pre><span></span><span class="k">def</span><span class="w"> </span><span class="nf">compute_nullclines</span><span class="p">(</span><span class="n">neuron_model</span><span class="p">,</span> <span class="n">u_range</span><span class="p">,</span> <span class="n">v_range</span><span class="p">,</span> <span class="n">resolution</span><span class="o">=</span><span class="mi">200</span><span class="p">):</span>
<span class="w"> </span><span class="sd">"""</span>
<span class="sd"> Compute nullclines</span>
<span class="sd"> du/dt = 0 (u-nullcline)</span>
<span class="sd"> dv/dt = 0 (v-nullcline)</span>
<span class="sd"> """</span>
<span class="n">u_vals</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">linspace</span><span class="p">(</span><span class="n">u_range</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">u_range</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">resolution</span><span class="p">)</span>
<span class="n">v_vals</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">linspace</span><span class="p">(</span><span class="n">v_range</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">v_range</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">resolution</span><span class="p">)</span>
<span class="n">U</span><span class="p">,</span> <span class="n">V</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">meshgrid</span><span class="p">(</span><span class="n">u_vals</span><span class="p">,</span> <span class="n">v_vals</span><span class="p">)</span>
<span class="n">UV</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">stack</span><span class="p">(</span>
<span class="p">[</span><span class="n">U</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">),</span> <span class="n">V</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">),</span> <span class="n">jnp</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">resolution</span> <span class="o">*</span> <span class="n">resolution</span><span class="p">,))],</span> <span class="n">axis</span><span class="o">=</span><span class="mi">1</span>
<span class="p">)</span>
<span class="n">dS</span> <span class="o">=</span> <span class="n">neuron_model</span><span class="o">.</span><span class="n">dynamics</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="n">UV</span><span class="p">,</span> <span class="p">{})</span>
<span class="n">dU</span> <span class="o">=</span> <span class="n">dS</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">U</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
<span class="n">dV</span> <span class="o">=</span> <span class="n">dS</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">V</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
<span class="k">return</span> <span class="n">U</span><span class="p">,</span> <span class="n">V</span><span class="p">,</span> <span class="n">dU</span><span class="p">,</span> <span class="n">dV</span>
<span class="k">def</span><span class="w"> </span><span class="nf">plot_vf</span><span class="p">(</span><span class="n">ax</span><span class="p">,</span> <span class="n">neuron_model</span><span class="p">,</span> <span class="n">u_range</span><span class="p">,</span> <span class="n">v_range</span><span class="p">):</span>
<span class="kn">import</span><span class="w"> </span><span class="nn">numpy</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">np</span>
<span class="n">u_sparse</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">linspace</span><span class="p">(</span><span class="n">u_range</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">u_range</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="mi">30</span><span class="p">)</span>
<span class="n">v_sparse</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">linspace</span><span class="p">(</span><span class="n">v_range</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">v_range</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="mi">30</span><span class="p">)</span>
<span class="n">Us</span><span class="p">,</span> <span class="n">Vs</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">meshgrid</span><span class="p">(</span><span class="n">u_sparse</span><span class="p">,</span> <span class="n">v_sparse</span><span class="p">)</span>
<span class="n">U</span><span class="p">,</span> <span class="n">V</span><span class="p">,</span> <span class="n">dU</span><span class="p">,</span> <span class="n">dV</span> <span class="o">=</span> <span class="n">compute_nullclines</span><span class="p">(</span><span class="n">neuron_model</span><span class="p">,</span> <span class="n">u_range</span><span class="p">,</span> <span class="n">v_range</span><span class="p">,</span> <span class="mi">200</span><span class="p">)</span>
<span class="n">UVs</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">stack</span><span class="p">([</span><span class="n">Us</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">),</span> <span class="n">Vs</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="o">-</span><span class="mi">1</span><span class="p">),</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ones</span><span class="p">((</span><span class="mi">30</span> <span class="o">*</span> <span class="mi">30</span><span class="p">,))],</span> <span class="n">axis</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
<span class="n">dS</span> <span class="o">=</span> <span class="n">neuron_model</span><span class="o">.</span><span class="n">dynamics</span><span class="p">(</span><span class="mi">0</span><span class="p">,</span> <span class="n">UVs</span><span class="p">,</span> <span class="p">{})</span>
<span class="n">dUs</span> <span class="o">=</span> <span class="n">dS</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">Us</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
<span class="n">dVs</span> <span class="o">=</span> <span class="n">dS</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">reshape</span><span class="p">(</span><span class="n">Vs</span><span class="o">.</span><span class="n">shape</span><span class="p">)</span>
<span class="c1"># Normalize for visualization</span>
<span class="n">magnitude</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">sqrt</span><span class="p">(</span><span class="n">dUs</span><span class="o">**</span><span class="mi">2</span> <span class="o">+</span> <span class="n">dVs</span><span class="o">**</span><span class="mi">2</span><span class="p">)</span>
<span class="n">magnitude</span><span class="p">[</span><span class="n">magnitude</span> <span class="o">==</span> <span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="mi">1</span>
<span class="n">dUs_norm</span> <span class="o">=</span> <span class="n">dUs</span> <span class="o">/</span> <span class="n">magnitude</span>
<span class="n">dVs_norm</span> <span class="o">=</span> <span class="n">dVs</span> <span class="o">/</span> <span class="n">magnitude</span>
<span class="c1"># Nullclines</span>
<span class="n">ax</span><span class="o">.</span><span class="n">contour</span><span class="p">(</span><span class="n">U</span><span class="p">,</span> <span class="n">V</span><span class="p">,</span> <span class="n">dU</span><span class="p">,</span> <span class="n">levels</span><span class="o">=</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">colors</span><span class="o">=</span><span class="s2">"blue"</span><span class="p">,</span> <span class="n">linewidths</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">linestyles</span><span class="o">=</span><span class="s2">"-"</span><span class="p">)</span>
<span class="n">ax</span><span class="o">.</span><span class="n">contour</span><span class="p">(</span><span class="n">U</span><span class="p">,</span> <span class="n">V</span><span class="p">,</span> <span class="n">dV</span><span class="p">,</span> <span class="n">levels</span><span class="o">=</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">colors</span><span class="o">=</span><span class="s2">"red"</span><span class="p">,</span> <span class="n">linewidths</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span> <span class="n">linestyles</span><span class="o">=</span><span class="s2">"-"</span><span class="p">)</span>
<span class="c1"># Vector field</span>
<span class="n">ax</span><span class="o">.</span><span class="n">quiver</span><span class="p">(</span><span class="n">Us</span><span class="p">,</span> <span class="n">Vs</span><span class="p">,</span> <span class="n">dUs_norm</span><span class="p">,</span> <span class="n">dVs_norm</span><span class="p">,</span> <span class="n">magnitude</span><span class="p">,</span> <span class="n">cmap</span><span class="o">=</span><span class="s2">"viridis"</span><span class="p">,</span> <span class="n">alpha</span><span class="o">=</span><span class="mf">0.6</span><span class="p">)</span>
<span class="n">ax</span><span class="o">.</span><span class="n">set_xlabel</span><span class="p">(</span><span class="s2">"v"</span><span class="p">)</span>
<span class="n">ax</span><span class="o">.</span><span class="n">set_ylabel</span><span class="p">(</span><span class="s2">"w"</span><span class="p">)</span>
<span class="n">ax</span><span class="o">.</span><span class="n">legend</span><span class="p">([</span><span class="s2">"u-nullcline (du/dt=0)"</span><span class="p">,</span> <span class="s2">"v-nullcline (dv/dt=0)"</span><span class="p">],</span> <span class="n">loc</span><span class="o">=</span><span class="s2">"upper right"</span><span class="p">)</span>
<span class="n">ax</span><span class="o">.</span><span class="n">set_xlim</span><span class="p">(</span><span class="n">u_range</span><span class="p">)</span>
<span class="n">ax</span><span class="o">.</span><span class="n">set_ylim</span><span class="p">(</span><span class="n">v_range</span><span class="p">)</span>
<span class="n">ax</span><span class="o">.</span><span class="n">axhline</span><span class="p">(</span><span class="n">y</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span> <span class="n">color</span><span class="o">=</span><span class="s2">"gray"</span><span class="p">,</span> <span class="n">linestyle</span><span class="o">=</span><span class="s2">"--"</span><span class="p">,</span> <span class="n">alpha</span><span class="o">=</span><span class="mf">0.3</span><span class="p">)</span>
<span class="n">ax</span><span class="o">.</span><span class="n">axvline</span><span class="p">(</span><span class="n">x</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span> <span class="n">color</span><span class="o">=</span><span class="s2">"gray"</span><span class="p">,</span> <span class="n">linestyle</span><span class="o">=</span><span class="s2">"--"</span><span class="p">,</span> <span class="n">alpha</span><span class="o">=</span><span class="mf">0.3</span><span class="p">)</span>
</pre></div>
<div class="clipboard-copy-txt" id="cell-5">def compute_nullclines(neuron_model, 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.zeros((resolution * resolution,))], axis=1
)
dS = neuron_model.dynamics(0, UV, {})
dU = dS[:, 0].reshape(U.shape)
dV = dS[:, 1].reshape(V.shape)
return U, V, dU, dV
def plot_vf(ax, neuron_model, u_range, v_range):
import numpy as np
u_sparse = jnp.linspace(u_range[0], u_range[1], 30)
v_sparse = jnp.linspace(v_range[0], v_range[1], 30)
Us, Vs = jnp.meshgrid(u_sparse, v_sparse)
U, V, dU, dV = compute_nullclines(neuron_model, u_range, v_range, 200)
UVs = jnp.stack([Us.reshape(-1), Vs.reshape(-1), jnp.ones((30 * 30,))], axis=1)
dS = neuron_model.dynamics(0, 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("v")
ax.set_ylabel("w")
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)</div>
</div>
</div>
</div>
</div>
</div>
</div><div class="jp-Cell jp-CodeCell jp-Notebook-cell">
<div class="jp-Cell jp-CodeCell jp-Notebook-cell" id="cell-id=b73fb06d">
<div class="jp-Cell-inputWrapper" tabindex="0">
<div class="jp-Collapser jp-InputCollapser jp-Cell-inputCollapser">
</div>
<div class="jp-InputArea jp-Cell-inputArea">
<div class="jp-InputPrompt jp-InputArea-prompt">In [27]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="CodeMirror cm-s-jupyter">
<div class="zeroclipboard-container">
<clipboard-copy ,="" aria-label="Copy to Clipboard" for="cell-6">
<div>
<span class="notice" hidden="">Copied!</span>
<svg aria-hidden="true" class="clipboard-copy-icon" data-view-component="true" height="20" version="1.1" viewbox="0 0 16 16" width="20">
<path d="M0 6.75C0 5.784.784 5 1.75 5h1.5a.75.75 0 010 1.5h-1.5a.25.25 0 00-.25.25v7.5c0 .138.112.25.25.25h7.5a.25.25 0 00.25-.25v-1.5a.75.75 0 011.5 0v1.5A1.75 1.75 0 019.25 16h-7.5A1.75 1.75 0 010 14.25v-7.5z" fill="currentColor" fill-rule="evenodd"></path>
<path d="M5 1.75C5 .784 5.784 0 6.75 0h7.5C15.216 0 16 .784 16 1.75v7.5A1.75 1.75 0 0114.25 11h-7.5A1.75 1.75 0 015 9.25v-7.5zm1.75-.25a.25.25 0 00-.25.25v7.5c0 .138.112.25.25.25h7.5a.25.25 0 00.25-.25v-7.5a.25.25 0 00-.25-.25h-7.5z" fill="currentColor" fill-rule="evenodd"></path>
</svg>
</div>
</clipboard-copy>
</div>
<div class="highlight-ipynb hl-python"><pre><span></span><span class="k">with</span> <span class="n">mpl</span><span class="o">.</span><span class="n">style</span><span class="o">.</span><span class="n">context</span><span class="p">(</span><span class="s2">"boilerplot.ieeetran"</span><span class="p">):</span>
<span class="n">fig</span><span class="p">,</span> <span class="n">ax</span> <span class="o">=</span> <span class="n">plt</span><span class="o">.</span><span class="n">subplots</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="n">figsize</span><span class="o">=</span><span class="p">(</span><span class="mf">6.9</span><span class="p">,</span> <span class="mf">2.6</span><span class="p">),</span> <span class="n">dpi</span><span class="o">=</span><span class="mi">200</span><span class="p">)</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">comp_times</span><span class="p">,</span> <span class="n">state</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="p">:,</span> <span class="mi">0</span><span class="p">])</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">comp_times</span><span class="p">,</span> <span class="n">state</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="p">:,</span> <span class="mi">1</span><span class="p">],</span> <span class="s2">"--"</span><span class="p">)</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">comp_times</span><span class="p">,</span> <span class="n">state</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">])</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">comp_times</span><span class="p">,</span> <span class="n">state</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">],</span> <span class="s2">"--"</span><span class="p">)</span>
<span class="c1"># ax[0].plot(comp_times, state[0, :, 2], "-.")</span>
<span class="p">[</span><span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">axvline</span><span class="p">(</span><span class="n">s</span><span class="p">,</span> <span class="n">alpha</span><span class="o">=</span><span class="mf">0.2</span><span class="p">,</span> <span class="n">color</span><span class="o">=</span><span class="s2">"g"</span><span class="p">,</span> <span class="n">linestyle</span><span class="o">=</span><span class="s2">"--"</span><span class="p">)</span> <span class="k">for</span> <span class="n">s</span> <span class="ow">in</span> <span class="n">jnp</span><span class="o">.</span><span class="n">unique</span><span class="p">(</span><span class="n">spikes</span><span class="p">)]</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">set_xlabel</span><span class="p">(</span><span class="s2">"Time (ms)"</span><span class="p">)</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">legend</span><span class="p">([</span><span class="s2">"v"</span><span class="p">,</span> <span class="s2">"vslow"</span><span class="p">,</span> <span class="s2">"syn"</span><span class="p">,</span> <span class="s2">"Spike"</span><span class="p">])</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">legend</span><span class="p">([</span><span class="s2">"v"</span><span class="p">,</span> <span class="s2">"vslow"</span><span class="p">,</span> <span class="s2">"syn"</span><span class="p">])</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">state</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="p">:,</span> <span class="mi">0</span><span class="p">],</span> <span class="n">state</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="p">:,</span> <span class="mi">1</span><span class="p">])</span>
<span class="n">plot_vf</span><span class="p">(</span><span class="n">ax</span><span class="p">[</span><span class="mi">1</span><span class="p">],</span> <span class="n">neuron_model</span><span class="p">,</span> <span class="p">[</span><span class="o">-</span><span class="mi">2</span><span class="p">,</span> <span class="mi">2</span><span class="p">],</span> <span class="p">[</span><span class="o">-</span><span class="mi">2</span><span class="p">,</span> <span class="mi">2</span><span class="p">])</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">state</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span> <span class="n">state</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">])</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">state</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span> <span class="n">state</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">],</span> <span class="s2">"."</span><span class="p">,</span> <span class="n">label</span><span class="o">=</span><span class="s2">"start"</span><span class="p">)</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">plot</span><span class="p">(</span><span class="n">state</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span> <span class="n">state</span><span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">],</span> <span class="s2">"."</span><span class="p">,</span> <span class="n">label</span><span class="o">=</span><span class="s2">"end"</span><span class="p">)</span>
<span class="n">ax</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">set_xlabel</span><span class="p">(</span><span class="s2">"v"</span><span class="p">)</span>
@@ -3351,16 +3408,17 @@ spikes = snn.spikes_until_t(in_spikes, max_time)</div>
<span class="n">ax</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span><span class="o">.</span><span class="n">legend</span><span class="p">()</span>
<span class="n">plt</span><span class="o">.</span><span class="n">show</span><span class="p">()</span>
</pre></div>
<div class="clipboard-copy-txt" id="cell-5">with mpl.style.context("boilerplot.ieeetran"):
<div class="clipboard-copy-txt" id="cell-6">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])
ax[0].plot(comp_times, state[0, :, 1], "--")
ax[0].plot(comp_times, state[:, 0, 0])
ax[0].plot(comp_times, state[:, 0, 1], "--")
# ax[0].plot(comp_times, state[0, :, 2], "-.")
[ax[0].axvline(s, alpha=0.2, color="g", linestyle="--") for s in jnp.unique(spikes)]
ax[0].set_xlabel("Time (ms)")
ax[0].legend(["v", "vslow", "syn", "Spike"])
ax[0].legend(["v", "vslow", "syn"])
ax[1].plot(state[0, :, 0], state[0, :, 1])
plot_vf(ax[1], neuron_model, [-2, 2], [-2, 2])
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].set_xlabel("v")
@@ -3378,7 +3436,7 @@ spikes = snn.spikes_until_t(in_spikes, max_time)</div>
<div class="jp-OutputArea-child">
<div class="jp-OutputPrompt jp-OutputArea-prompt"></div>
<div class="jp-RenderedImage jp-OutputArea-output" tabindex="0">
<img alt="No description has been provided for this image" class="" src="data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAABXsAAAIfCAYAAAAhag9YAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjgsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvwVt1zgAAAAlwSFlzAAAewgAAHsIBbtB1PgABAABJREFUeJzsnQd4HOX19Y96l1wky3LvveAGNraB2FRD6J0QCCHlI0AgyT90AiGUACnUEJIQSkIJvRkw2DRjAzY22Ni49yZbstV7+Z47s7P7zuystJK2zOyeH8+yZWZ3x7tXszPnPe+5Ca2tra0ghBBCCCGEEEIIIYQQ4moSo70BhBBCCCGEEEIIIYQQQroOxV5CCCGEEEIIIYQQQgiJASj2EkIIIYQQQgghhBBCSAxAsZcQQgghhBBCCCGEEEJiAIq9hBBCCCGEEEIIIYQQEgNQ7CWEEEIIIYQQQgghhJAYgGIvIYQQQgghhBBCCCGExAAUewkhhBBCCCGEEEIIISQGoNhLCCGEEEIIIYQQQgghMQDFXkIIIYQQQgghhBBCCIkBKPYSQgghhBBCCCGEEEJIDECxlxBCCCGEEEIIIYQQQmIAir2EEEIIIYQQQgghhBASA1DsJYQQQgghhBBCCCGEkBiAYi8hhBBCCCGEEEIIIYTEABR7CSGEEEIIIYQQQgghJAag2EsIIYQQQgghhBBCCCExQDLijJaWFpSUlGi3MzMzkZCQEO1NIoQQQgiJOVpbW1FTU6Pdzs/PR2IiPQZugsfMhBBCCCHuPF6OO7FXhN7CwsJobwYhhBBCSNxQXFyMXr16RXszSAfgMTMhhBBCiDuPl+POYlFdXR3tTSCEEEIIiSt4/EUIIYQQQkhkiDtnr0Q32FFVVdXuSYrhCBa1PSsrK6j3i+Tz3LCNnX2eG7axs89zwzZ29nlu2MbOPs8N29jZ57lhGzv7PG4jP/94rq3OPq+jz8nOzg76+Is4F/U760iNhbpmnfg+sfIekXqfWHmPSL0P34OfF2vYGX8n/Hvk5xWpfX04jpfjTuxVM3q3bNnitUh35MuTdTvzZUfyeW7Yxs4+zw3b2NnnuWEbO/s8N2xjZ5/nhm3s7PPcsI2dfR63kZ9/PNdWZ58XzHOMAfT9+/djyJAh2m32SHAfgb6zUJzwdLZmnfg+sfIekXqfWHmPSL0P34OfF+vLGX8nkXqfWHmPSL2Pm9+j2iZ1IJTHy3En9kajyAkhhBBC4gXj2IrHWLGD6jyRRiKEEEIIIaTzhLuXWNyJveqJR0fdvJE8uO3M+3V2GyP9vM7gln8bP5PofR5deT83/L11Fjf82/h3E73Poyvv54ba6ixu+Le55e/G7jaJTyJVs5HeX4SLWPq8YuU9IkWsfF6xVMORIJY+r1h5j0gRS59XLP1b3E7cib1OYeuhrWhqaUJyYjIGdx8c7c0hDvnut5VtQ2pdapt1wdohna2v9mqHtUXCVVusL0LcTTjzEAkhhBBC4o0qT+yZXXZvKKDYGyUqGyrR2NyIlKSUaG0CceB3X9VYheSE5DbrgrVDOltf7dUOa4uEq7ZYX4S4G0afEUIIIYSEjnAPoieG9dUJIYQQQgghhBBCCCGERAQ6e8MM80T4mXSkRlYVr9LccfEO/274mbBO+HfDfQn3r4Q4FR6nxC/87uMTfu/xC7/7+CQrBnKB6ewlhBBCCCGEEEIIIYSQGIBiLyGEEEIIIYQQQgghhMQAFHsJIYQQQgghhBBCCCEkBnC82NvY2IiFCxfi//7v/zBt2jR069YNKSkp6N27N0499VS8/fbb0d5EQgghhBBCCCGEEEIIiTqOb9D28ccf47jjjtNui8A7a9YsLSx57dq1ePPNN7XLT3/6Uzz22GNISEiI9uYSQgghhBASU1RXV3tvy3E4IYQQQgjp+rGVeowVV2JvYmIizjrrLPzyl7/E7NmzTcteeOEFXHTRRXj88ccxc+ZM/PCHP4zadhJCCCGEEBKLFBYWem+7vTs1IYQQQki0yc7ODuvrO17snTNnjnax47zzzsP777+Pf/3rX3j66addJfYWZBagubUZSQlJ0d4U4qDvPpi6YO2QcNUXa4tw30UIIYQQQggh7sbxYm97TJo0SbveuXMn3ERRTlG0N4E48LsPpi5YOyRc9cXaIuGqLdYXIe6muLiY8Q2EEEIIISGiqqrKG+OgzqAKFa4Xezdu3KhdFxVRPCWEEEIIISTUSE4vs3oJIYQQQkJDuI+rXC327tu3D08++aR2W3J9O0pbQcg8oCVOZOfBGry4fCdOGNcbY/vkRXtzCCGEkIDHU+FqOEEIIYS4Eck8b2huwaHqRhRX1OmXynrs99w+ILcr61FcUY+SqvqobmtuejKGF+ZgRGE2BudnYUh+NgYXZKFf9wykJTOKkhCnk9Dq0i4LTU1NOPHEE7Fw4UKMHz8ey5cvR2pqarvPkxOPYIKQXfqxkBjnqHs/xI6DNcjLSME3vzs+2ptDCCGEICEhIaipahxIdxfqMTO/P0JIvNDS0oqaxmbsKavFtpJqbCuVSw027a/C2j0VqKpvivYmOp7pQ3rgpHFFOHFcbxTmpkd7cwiJy+Mt1zp7f/7zn2tCb8+ePfHSSy8FJfQ6iVXFq9DY3IiUpBRMKJwQ7c0hDvjum5pb8NTyj1HULRX9uuXY1oUIvQ0JW7G/rhmrinuzdkjQ9dXesmCWk/imK7UlLN+9EgkJLawvQgghhIQcOZfaU1aHdfsqsKG4El/vLMfKHYdQWt3ATzvCfL7loHb53RtrbJefPL4IV88djpG9c/jdEBImXCn2/vKXv8S//vUvdO/eHe+//z5GjBjhqmYTj3y4CU8vX47LZw/AEYN7Rfz9iTN5aul23PL2GmSnJ+C/lx/pt7y+qTkq20UIIcLbq/egpKoG500d3KnfvT8uXIKzphThR0cO4wcaow0mrISr4QQhhJDYo6KuEev3VWLN7nKs3VuBFTvKNDctiT3eXr1Xu1j5vxNG4qdHDUFKUmJUtouQWMJ1Yu+vf/1rPPjgg+jWrRsWLFiASZMmuarZRE1DE+57bz0aEmpw6+tr8N41FHuJzh1vrQUSgMq6Ju0gZ1pf8yezvyK6uU2EkPjl040H8NDCTWhFM1ISU3B4v4799srvXlNCK57/cifF3hiE8QyEEEIk/qCkuh5rdlfg651lWL79ID7bVBo3H0yvnDTNqTq8Vw5G9pac22z0yEpFdloyMtOSkJmShOQYEDEbm1tQXd+knbNKvvBH6/fj1ZW7setQbZdfW44X5WKQn52Gpy87HGP65Hb5tQmJN1wl9v72t7/Fn//8Z+Tl5WlC79SpU+E2DnIaCQmCjFT/P03JjersgdeybQfRMzsVw3pxqgwJLXWNzdoBvTRr6Nc9kx9vjPLG13u8t1/+ahfuOTX45za3dD4DX6ZhyknFiMIcujwIIYSQMCN9aySTVhqFrdlTgcUbS/Dh+v2aqBdLiAA7vm8eJvTPw4S+3TCwZyZy0pORmZqMrLQkpCcnITGx/Tz6eERct90yU7VL/x6ZmDKwO359/MiA0RolVQ34dnc5/vz+Bs2x3RGkSd28Bz/13u+dm47Xr5zJHGBCYknsvf7663HfffdpQq9EN0ybNg1uREbACGmP9BT/Ud+95XUd/uDKaxpx1fMr8cmGA9r9I4f2xN8umoK8zBR+CcSEDAg8+PE6nDAuHxP65Qf16Xy2qQRX/HcFymsbkZqUiF8fP0KbehVMsybiLqRRiUFqSsc6MJd2spv0fe+twyMfbtZu98lLxz8vmUZnByGEENIB4bauUcS2eqzeXa4d6326scT10Qijeufg8ME9NJFxaEE2umWmIDcjBdmpyRRoHYY4mXvnpWuXY8eYY50amlqwpaQKf/toM15XTAVtsa+iDkfctdB7/6ELJuH7E/uEfLsJiQVcIfbefPPN+OMf/+iNbnCr0CtQ7CWdZU95rd8BXHvc8vq3XqFXWLK5FD99Zjme+fERSE12/zQiEjrOeWwpGhL244N1e/HeNd9rd/3KukZc88LXmtArNDS34O531iErLRk/mD6QX02MUdfQHHAwatehGjy4cCPG983G90YV+T23M26gl77a5RV6hT3ldTj/8aV486pZGNgz8ln7hMQ7kr9swNgOQqKHzNjbXVaLr7Yf0gbdP9l4AMUui3obkp+FGUN7aoLtkHyPWJueguz0ZCTRTRs3yLnoqN65eOD8SdpFPa6UeMP31hS3+xpXPbdSuwh/OmcizprSL6zbTEg4jq3UY6y4EnvfeOMN3HnnndrtYcOG4ZFHHrFdLz8/H/fffz+cToVHGCGkLex03H0WZ297Wq/8UL61yn+U9IutB/H4J5tx5Zzh/BKIp5Y6Ps3+0Y82a1P8rPz+rbWYPqQHI0NijFrF2ZuWZHb2/vTpr/D1Pmm00YypA/1d4fsr6zpcjyIeW6moa8L1L6/Gsz85gu5xQiKM2mivM78ZhBB/5G+ppqFZi/mT38rPt0jGbYlmznA6/XtkYPbwAswelo8RvXPQIzNVi0GIhUxaEl0kFu7vF/viOsVY8s9Pt+ChRZvafN6vX/xGuwhvXTUL4/rmhX1bCekK2dnZCCeOF3sPHjzovb18+XLtYsfAgQNdIfYaLjhC2qLF5kRqT5lF7G3nI3xm6XYYUZnXHjsCs0fka+5Nyc8Ux9yZk/uhT7cMfhFEy2br6MnJKyt2abclvmHhr4/GvxZvxZNLtmlTsu6evw7/utS9MzBI22JvqsXZq+WveZI7ymx+4zraXPKbXeXYcbDGGz3z+A+n4oS/fKI5mZZuKdWm+p0+ydLBkhBCCHEQcqwkU87FgfvmN3uCcilGi6NHFODY0b0waUB39O+eidyMZA6qEseQl5GiZQIbucCb9lfi2D9/0uZzTnlosXY9oV8eXrtiLine truncated
<img alt="No description has been provided for this image" class="" src="data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAABXsAAAIfCAYAAAAhag9YAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjgsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvwVt1zgAAAAlwSFlzAAAewgAAHsIBbtB1PgABAABJREFUeJzs3XeYG9W5+PHvzKittknaorK9N3cMGEwvoYYewiWEnoQkl0tIQi6kkh/pJCEBAgRCKKGXi+lgU1wBY1xwb2uv6/auLs3M7w/tatfYXmlt2Wub83keP1aZOXPmTNnRO++cI+m6riMIgiAIgiAIgiAIgiAIgiAc1uSxroAgCIIgCIIgCIIgCIIgCIKw/0SwVxAEQRAEQRAEQRAEQRAE4Qgggr2CIAiCIAiCIAiCIAiCIAhHABHsFQRBEARBEARBEARBEARBOAKIYK8gCIIgCIIgCIIgCIIgCMIRQAR7BUEQBEEQBEEQBEEQBEEQjgAi2CsIgiAIgiAIgiAIgiAIgnAEEMFeQRAEQRAEQRAEQRAEQRCEI4AI9gqCIAiCIAiCIAiCIAiCIBwBRLBXEARBEARBEARBEARBEAThCCCCvYIgCIIgCIIgCIIgCIIgCEcAEewVBEEQBEEQBEEQBEEQBEE4AohgryAIgiAIgiAIgiAIgiAIwhFABHsFQRAEQRAEQRAEQRAEQRCOACLYKwiCIAiCIAiCIAiCIAiCcAQQwV5BEARBEARBEARBEARBEIQjgAj2CoIgCIIgCIIgCIIgCIIgHAFEsFcQBEEQBEEQBEEQBEEQBOEIIIK9giAIgiAIgiAIgiAIgiAIRwAR7BUEQRAEQRCEgyQSifD+++9z2223cfTRR2Oz2TAajbhcLi644ALefPPNfS77vffe49xzzyU3N5e0tDRqa2v52c9+htfrTeEaCIIgCIIgCIcySdd1fawrIQiCIAiCIAhfBu+99x5nnnkmAC6Xi6OOOor09HRWr17NypUrAfj2t7/NQw89hCRJSZd7zz338MMf/hBJkjjxxBNxOp3MmzePlpYWampqmD9/Prm5uQdknQRBEARBEIRDhwj2CoIgCIIgCMJB8sEHH/DAAw9wyy23cOKJJ+7y3fPPP883vvENVFXliSee4Oqrr06qzKVLl3LUUUchyzKvv/4655xzDgB+v58LLriA999/n0svvZSXXnop5esjCIIgCIIgHFpEsFcQBEEQBEEQDhE33ngjjz76KKeffjrvvfdeUvNcfvnlvPjii9x444088sgju3y3ZcsWysvL0TSNNWvWUFtbeyCqLQiCIAiCIBwiRJ+9giAIgiAIgnCImDx5MgDbtm1LavpwOBzv5/fKK6/c7fuSkhKmT58OwCuvvJKiWgqCIAiCIAiHKhHsFQRBEARBEIRDxIYNGwBwu91JTb9+/Xr8fj8AU6dO3eM0g58vXbo0BTUUBEEQBEEQDmWGsa7A4UjTNDo6OgCwWq2jGjxDEARBEAThy0DX9XgQMjc3F1kWOQaJtLS08PjjjwNw6aWXJjXP5s2bAbDZbGRmZu5xmqKiol2mTWTwWtfn8414rZuenp5UeYIgCIIgCIcrn8+3x88Hr3XT09MPuWtdEezdBx0dHTidzrGuhiAIgiAIwmGhtbWV/Pz8sa7GIS0ajXLVVVfR29vL+PHj+c53vpPUfP39/cDIgdeMjAwA+vr6kipTXOsKgiAIgiAk71C71j10ws6Hkb1F9QVBEARBEITdiWunxG666Sbef/99cnJyeOmllzCZTGNWF7G9BEEQBEEQkneoXTuJzN59YLVa4683bdp0QKP3Pp8vnlnR2tp6QB+XOxKXdSSu08Fc1pG4TgdzWUfiOh3MZR2J63SkLutIXKeDuawjcZ0A2traKC8vB3a9dhJ2d8stt/Doo49it9uZNWsW1dXVSc872HXDSD8yvF4vAFlZWUmVOXx7rVy5kry8PABW960hrEZwWZzkmnPIzszebd7R7GObva00epspsuZSZM1FD6n7tH8OX+ZDC2dSk1dAWVYOOeb0hN2tjVRfTdd5ZeMqHBYr5Vl2CjKzMQx7RHNfjqc5mzfT2dvL1SecMOr13NMyt3v9rNzRSmmOnZIcGznpe+92I5n69vgCvL1kHYU52ZTk2XA7sjAqyj6fO3r7+inwxPqfnrtoBdWlHjLSzCPOE4lAcTHEktY7gDJaW1viy5zz2Qb6+oMUOG0UOLPJs2ciy7uu82jqu3ZjCyvW7MDjysaWaeTYqQ2jXs8vLvP+B16nqMiJ22XD5crGbhu567+R6quqGjNmLCbNYsLtycblyiYvLwuDYd+3y5wPV9Pc3MmN3z4HgKambeTm2vdpXVtbW9mxrY/Pl27B5bbh9thwuW1kZFpGtZ7DdXf7mPnW5zhd2bjcdtweG4pBY9q0aRgMBp555hnq6+uTrm9fbz+egth+OPuDxZRXFGB37N/5YcHcdfT1+vEU2nF77OTmZe2yH45222xc38KaldtxF9ixO8xMnFyb9Lx7WuaG9Zv5ZE4jLo8NT6EDV4Eds8W4T+sKoEZV3p2xhLR0M55CB+4iB1k26z6t66D3X/uUC688DYDPP1tNSUUBBqMy6nUdXOa6ZVtYPHsN+QX22L9CB7lu+25lJlvfjuYuXvrb26hRLf6ZbILb7vsuAL+94W8Y5GFtKsF5N5xGcV3Bbo/zJ7tMTdPobO1l5+Z2dja1UzW+GHe5Y5+v3YYv960ZC8jKzGDSsRWjmm9Py/T5QmzZ0klnZz8nnFC9y7G0L/tDVNXY0dLNmg3buez86aOad0/LtFqtdPT62Lizk4KcLIqdez+/JVNfn8+HNxii4423qPrFzzB1dQLgBVwD0xxy17q6MGqtra06oAN6a2vrAV2W1+uNL8vr9YplHaLLOVKXdSSu08Fc1pG4TgdzWUfiOh2pyzoS1+lgLutIXCddP7jXS4ezH/7whzqg22w2fdGiRaOef/ny5fF27uvr2+M0t956qw7ol112WVJl7s+22599bF/nHYtlfpnquz/z7ut8l1yi6zD475hd5tU0LaXLHV7eodi+mqbtdZ33dZmapqW0vslsk9HUd0/leb1evby8XK+urtZXrlw5qvr29/enfP9NtM6H0/E2VvPu63bZ2zKjkWhK67t9Y4u+fskmfUdji97b0adHI9Fd5t2+aYe+s7FVb1y+RV/18Xr9s1nL9Q1LN+/XMg+VeUV9vyAS0fWf/UzXJWnoj5Pbrbe//PIhe60rMnsFQRAEQRAEYQz85Cc/4a9//SvZ2dnMnDmTqVOnjrqMmpoarFYrfr+fzz77jFNPPXW3aT777DMApkyZst91FoSD4Zxz4P/+b/Ddebt8l+rBsQ/1wbYPRP0O9TY81Ms7UGV+2aS6DRVDclnBySqo2EPf9aGhl7b8bDFQ6ZfBtm1w5ZUwf/7QZ2edBU8+ibb3ucacCPbug+EH9JF0cKenp6Pr+hG1rCNxnQ6mI7X9jsT972A6EtvvSF3WwXKktt+RuK8fTEfq9VKq3H777dx9991kZ2cza9Ysjj766H0qx2Qycd555/Hiiy/yzDPP7Bbs3bJlCx999BEAF1988X7XO5H92Z/3dd6xWOb+GKv6Hk7tdM45Q6+nTPkloz2FiH3p0Jx3rP7eHW77w5dlXzrc2mh/liv2pQM77wFb5uuvw7XXQldX7L2iwO9+Bz/+McgytLXt0zIPBjFAmyAIgiAIgiAcRD//+c/54x//iM1mSzrQe//991NbW8vVV1+923e33347kiTx2GOP8c4778Q/9/v93HDDDaiqyqWXXkptbW1K10MQDpSCApg8OfZ6yRLYuXNs6yMIgiB8iYTDcOutcMEFQ4He4mKYNw9+8pNYoPcQJzJ7BUEQBEEQBOEgee211/jtb38LQGVlJf/4xz/2OF1ubi5//vOf4+87OjpYt24dLpdrt2mnTJnCX/7yF374wx9y7rnncvLJJ5Ofn8+8efNobm6mpqaGhx566MCskCAcIOefD0uXxl6/8QZ8+9tjWx9BEAThS6CxEa64Aga6wALg4ovh0UfBnvxAlmNNBHsFQRAEQRAE4SDpGswQIdaX7mfDf0wMU1JSskuwN5Fbb72V8ePH85e//IVPP/0Un89HcXExd9xxB3fccQeZmZn7XXdBOJguuADuuiv2+rXXRLBXgN7eXmRZJisra6yrIgjCkej55+Fb34L+/th7kwn+8hf4/vfhMOunW9KPxE7iDjCfz0dGRgYAXq9X9EMnCIIgCILwBeJ66fAltp1wKNA0KCqKdeFgNkN7O4h7FoIgCELKBQLwgx/Aww8PfVZZCS+8MNSn0B4cytdLh35HE4IgCIIgCIIgjAmfzxf/JwgHkyzDhRfGXodC8PbbY1sfQRAE4Qi0Zg0ce+yugd5vfCPWYfxeAr2Hw7WRCPYKgiAIgiAIgrBHTqeTjIyMeOaKIBxMl1469Pqll8auHoIgCMIRRtfh3/+GqVNhxYrYZ2lpsc/+858RHyUZvC5yOp0HqbKjJ/rs3U/DI/mHUsq2IAiCIAjCWBi8NjqUsx0EQTg8nHwy5ORAZye8+Sb4/WC1jnWtBEEQhMNafz9897vw9NNDn40bF+uzt75+7OqVQiKzdz+JbAdBEARBEIQhh0O2g5C81tZWvF4vXq+Xvkgfmq6lrOz+SICQGklZeWE1Sl84mLLyADqD/pSW1x0IkMohU7yhMBFVTVl5EVXFH0rdNgHwBkL7PK/BAJdcEnvt9w8EfIPhlLZhKBRB01JXXlTViEZTt00AwuFoassLpb68VG4TNartcZ0bGxtZt24djY2Noy4zEAinompxfn8opescCIRRo6k7v0ajKsFgao9lnze151e/L7VtmOr9cKypUZWAL3VtHo2odLT2pqw8VdXYvr0rZW2u6zrNbb34U3isdvb5Line truncated
</div>
</div>
</div>
@@ -3393,7 +3451,7 @@ spikes = snn.spikes_until_t(in_spikes, max_time)</div>
<div class="jp-InputPrompt jp-InputArea-prompt">In [ ]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="CodeMirror cm-s-jupyter">
<div class="zeroclipboard-container">
<clipboard-copy ,="" aria-label="Copy to Clipboard" for="cell-6">
<clipboard-copy ,="" aria-label="Copy to Clipboard" for="cell-7">
<div>
<span class="notice" hidden="">Copied!</span>
<svg aria-hidden="true" class="clipboard-copy-icon" data-view-component="true" height="20" version="1.1" viewbox="0 0 16 16" width="20">
@@ -3405,7 +3463,34 @@ spikes = snn.spikes_until_t(in_spikes, max_time)</div>
</div>
<div class="highlight-ipynb hl-python"><pre><span></span>
</pre></div>
<div class="clipboard-copy-txt" id="cell-6"></div>
<div class="clipboard-copy-txt" id="cell-7"></div>
</div>
</div>
</div>
</div>
</div>
</div><div class="jp-Cell jp-CodeCell jp-Notebook-cell jp-mod-noOutputs">
<div class="jp-Cell jp-CodeCell jp-Notebook-cell jp-mod-noOutputs" id="cell-id=daf213b5">
<div class="jp-Cell-inputWrapper" tabindex="0">
<div class="jp-Collapser jp-InputCollapser jp-Cell-inputCollapser">
</div>
<div class="jp-InputArea jp-Cell-inputArea">
<div class="jp-InputPrompt jp-InputArea-prompt">In [ ]:</div><div class="jp-CodeMirrorEditor jp-Editor jp-InputArea-editor" data-type="inline">
<div class="CodeMirror cm-s-jupyter">
<div class="zeroclipboard-container">
<clipboard-copy ,="" aria-label="Copy to Clipboard" for="cell-8">
<div>
<span class="notice" hidden="">Copied!</span>
<svg aria-hidden="true" class="clipboard-copy-icon" data-view-component="true" height="20" version="1.1" viewbox="0 0 16 16" width="20">
<path d="M0 6.75C0 5.784.784 5 1.75 5h1.5a.75.75 0 010 1.5h-1.5a.25.25 0 00-.25.25v7.5c0 .138.112.25.25.25h7.5a.25.25 0 00.25-.25v-1.5a.75.75 0 011.5 0v1.5A1.75 1.75 0 019.25 16h-7.5A1.75 1.75 0 010 14.25v-7.5z" fill="currentColor" fill-rule="evenodd"></path>
<path d="M5 1.75C5 .784 5.784 0 6.75 0h7.5C15.216 0 16 .784 16 1.75v7.5A1.75 1.75 0 0114.25 11h-7.5A1.75 1.75 0 015 9.25v-7.5zm1.75-.25a.25.25 0 00-.25.25v7.5c0 .138.112.25.25.25h7.5a.25.25 0 00.25-.25v-7.5a.25.25 0 00-.25-.25h-7.5z" fill="currentColor" fill-rule="evenodd"></path>
</svg>
</div>
</clipboard-copy>
</div>
<div class="highlight-ipynb hl-python"><pre><span></span>
</pre></div>
<div class="clipboard-copy-txt" id="cell-8"></div>
</div>
</div>
</div>
search/search_index.json
+1 -1
View File
@@ -1 +1 @@
{"config":{"lang":["en"],"separator":"[\\s\\-]+","pipeline":["stopWordFilter"],"fields":{"title":{"boost":1000.0},"text":{"boost":1.0},"tags":{"boost":1000000.0}}},"docs":[{"location":"","title":"Felice","text":"<p>This project provides a JAX implementation of the different neuron models in felice</p>"},{"location":"#overview","title":"Overview","text":"<p>The framework is built on top of EventPropJax and leverages JAX's automatic differentiation for efficient simulation and training of SNNs using event-based exact gradients.</p>"},{"location":"#key-features","title":"Key Features","text":"<ul> <li>Delay learning</li> <li>Non-linear neuron models<ul> <li>WereRabbit Neuron Model: Implementation of a dual-state oscillatory neuron model with bistable dynamics</li> <li>FHN Neuron Model</li> <li>Snowball Neuron Model</li> </ul> </li> </ul>"},{"location":"#installation","title":"\ud83d\udce6 Installation","text":"<p>Felice uses uv for dependency management. To install:</p> <pre><code>uv sync\n</code></pre>"},{"location":"#cuda-support-optional","title":"CUDA Support (Optional)","text":"<p>For GPU acceleration with CUDA 13:</p> <pre><code>uv sync --extra cuda\n</code></pre>"},{"location":"#quick-start","title":"Quick Start","text":"<p>Here's a simple example using the WereRabbit neuron model:</p> <pre><code>import diffrax as dfx\nimport jax.numpy as jnp\nimport jax.random as jrand\nfrom eventpropjax.evnn import FFEvNN\nfrom felice.neuron_models import WereRabbit\n\n# Initialize random key and parameters\nkey = jrand.key(0)\nmax_time = 300e-3\n\n# Create a feedforward event-driven neural network\nsnn = FFEvNN(\n layers=[1],\n in_size=2,\n neuron_model=WereRabbit,\n solver=dfx.Tsit5(),\n max_solver_time=max_time,\n key=key,\n max_event_steps=1000000,\n solver_stepsize=1e-6,\n rtol=10.0,\n atol=0.0,\n Ibias=30e-12,\n)\n\n# Simulate with input spikes\nin_spikes = jnp.asarray([[0.0], [0.157]])\nspikes = snn.spikes_until_t(in_spikes, max_time)\n</code></pre> <p>See the examples directory for more detailed usage examples.</p>"},{"location":"api/","title":"API Reference","text":"<p>API documentation for Felice.</p>"},{"location":"api/#modules","title":"Modules","text":"<ul> <li>Neuron Models - Neuron model implementations</li> <li>Solver - Zero-clipping solver</li> <li>Datasets - Built-in datasets</li> </ul>"},{"location":"api/datasets/","title":"Datasets","text":""},{"location":"api/datasets/#felice.datasets","title":"<code>felice.datasets</code>","text":""},{"location":"api/neuron_models/","title":"Neuron Models","text":""},{"location":"api/neuron_models/#felice.neuron_models","title":"<code>felice.neuron_models</code>","text":""},{"location":"api/neuron_models/#felice.neuron_models-classes","title":"Classes","text":""},{"location":"api/neuron_models/#felice.neuron_models.Boomerang","title":"<code>Boomerang</code>","text":"<p> Bases: <code>Module</code></p> Source code in <code>felice/neuron_models/boomerang.py</code> <pre><code>class Boomerang(eqx.Module):\n rtol: float = eqx.field(static=True)\n atol: float = eqx.field(static=True)\n\n u0: float = eqx.field(static=True)\n v0: float = eqx.field(static=True)\n\n alpha: float = eqx.field(static=True) # I_n0 / I_bias ratio\n beta: float = eqx.field(static=True) # k / U_t (inverse thermal scale)\n gamma: float = eqx.field(static=True) # coupling coefficient\n rho: float = eqx.field(static=True) # tanh steepness\n sigma: float = eqx.field(static=True) # bias scaling (s * I_bias)\n\n dtype: DTypeLike = eqx.field(static=True)\n\n def __init__(\n self,\n *,\n atol: float = 1e-6,\n rtol: float = 1e-4,\n alpha: float = 0.0129,\n beta: float = 15.6,\n gamma: float = 0.26,\n rho: float = 30.0,\n sigma: float = 0.6,\n dtype: DTypeLike = jnp.float32,\n ):\n r\"\"\"Initialize the WereRabbit neuron model.\n\n Args:\n key: JAX random key for weight initialization.\n n_neurons: Number of neurons in this layer.\n in_size: Number of input connections (excluding recurrent connections).\n wmask: Binary mask defining connectivity pattern of shape (in_plus_neurons, neurons).\n rtol: Relative tolerance for the spiking fixpoint calculation.\n atol: Absolute tolerance for the spiking fixpoint calculation.\n alpha: Current scaling parameter $\\alpha = I_{n0}/I_{bias}$ (default: 0.0129)\n beta: Exponential slope $\\beta = \\kappa/U_t$ (default: 15.6)\n gamma: Coupling parameter $\\gamma = 26e^{-2}$\n rho: Steepness of the tanh function $\\rho$ (default: 5)\n sigma: Fixpoint distance scaling $\\sigma$ (default: 0.6)\n wlim: Limit for weight initialization. If None, uses init_weights.\n wmean: Mean value for weight initialization.\n init_weights: Optional initial weight values. Line truncated
{"config":{"lang":["en"],"separator":"[\\s\\-]+","pipeline":["stopWordFilter"],"fields":{"title":{"boost":1000.0},"text":{"boost":1.0},"tags":{"boost":1000000.0}}},"docs":[{"location":"","title":"Felice","text":"<p>This project provides a JAX implementation of the different neuron models in felice</p>"},{"location":"#overview","title":"Overview","text":"<p>The framework is built on top of diffrax and leverages JAX's automatic differentiation for efficient simulation and training of analogue models.</p>"},{"location":"#key-features","title":"Key Features","text":"<ul> <li>Delay learning</li> <li>Non-linear neuron models<ul> <li>WereRabbit Neuron Model: Implementation of a dual-state oscillatory neuron model with bistable dynamics</li> <li>FHN Neuron Model</li> <li>Snowball Neuron Model</li> </ul> </li> </ul>"},{"location":"#installation","title":"\ud83d\udce6 Installation","text":"<p>Felice uses uv for dependency management. To install:</p> <pre><code>uv sync\n</code></pre>"},{"location":"#cuda-support-optional","title":"CUDA Support (Optional)","text":"<p>For GPU acceleration with CUDA 13:</p> <pre><code>uv sync --extra cuda\n</code></pre> <p>See the examples directory for more detailed usage examples.</p>"},{"location":"api/","title":"API Reference","text":"<p>API documentation for Felice.</p>"},{"location":"api/#modules","title":"Modules","text":"<ul> <li>Neuron Models - Neuron model implementations</li> <li>Solver - Zero-clipping solver</li> <li>Datasets - Built-in datasets</li> </ul>"},{"location":"api/datasets/","title":"Datasets","text":""},{"location":"api/datasets/#felice.datasets","title":"<code>felice.datasets</code>","text":""},{"location":"api/neuron_models/","title":"Neuron Models","text":""},{"location":"api/neuron_models/#felice.neuron_models","title":"<code>felice.neuron_models</code>","text":""},{"location":"api/neuron_models/#felice.neuron_models-classes","title":"Classes","text":""},{"location":"api/neuron_models/#felice.neuron_models.Boomerang","title":"<code>Boomerang</code>","text":"<p> Bases: <code>Module</code></p> Source code in <code>felice/neuron_models/boomerang.py</code> <pre><code>class Boomerang(eqx.Module):\n rtol: float = eqx.field(static=True)\n atol: float = eqx.field(static=True)\n\n u0: float = eqx.field(static=True)\n v0: float = eqx.field(static=True)\n\n alpha: float = eqx.field(static=True) # I_n0 / I_bias ratio\n beta: float = eqx.field(static=True) # k / U_t (inverse thermal scale)\n gamma: float = eqx.field(static=True) # coupling coefficient\n rho: float = eqx.field(static=True) # tanh steepness\n sigma: float = eqx.field(static=True) # bias scaling (s * I_bias)\n\n dtype: DTypeLike = eqx.field(static=True)\n\n def __init__(\n self,\n *,\n atol: float = 1e-6,\n rtol: float = 1e-4,\n alpha: float = 0.0129,\n beta: float = 15.6,\n gamma: float = 0.26,\n rho: float = 30.0,\n sigma: float = 0.6,\n dtype: DTypeLike = jnp.float32,\n ):\n r\"\"\"Initialize the WereRabbit neuron model.\n\n Args:\n key: JAX random key for weight initialization.\n n_neurons: Number of neurons in this layer.\n in_size: Number of input connections (excluding recurrent connections).\n wmask: Binary mask defining connectivity pattern of shape (in_plus_neurons, neurons).\n rtol: Relative tolerance for the spiking fixpoint calculation.\n atol: Absolute tolerance for the spiking fixpoint calculation.\n alpha: Current scaling parameter $\\alpha = I_{n0}/I_{bias}$ (default: 0.0129)\n beta: Exponential slope $\\beta = \\kappa/U_t$ (default: 15.6)\n gamma: Coupling parameter $\\gamma = 26e^{-2}$\n rho: Steepness of the tanh function $\\rho$ (default: 5)\n sigma: Fixpoint distance scaling $\\sigma$ (default: 0.6)\n wlim: Limit for weight initialization. If None, uses init_weights.\n wmean: Mean value for weight initialization.\n init_weights: Optional initial weight values. If None, weights are randomly initialized.\n fan_in_mode: Mode for fan-in based weight initialization ('sqrt', 'linear').\n dtype: Data type for arrays (default: float32).\n \"\"\"\n self.dtype = dtype\n\n self.alpha = alpha\n self.beta = beta\n self.gamma = gamma\n self.rho = rho\n self.sigma = sigma\n\n self.rtol = rtol\n self.atol = atol\n\n def fn(y, _):\n return self.vector_field(y[0], y[1])\n\n solver: optx.AbstractRootFinder = optx.Newton(rtol=1e-8, atol=1e-8)\n y0 = (jnp.array(0.3), jnp.array(0.3))\n u0, v0 = optx.root_find(fn, solver, y0).value\n self.u0 = u0.item()\n self.v0 = v0.item()\n\n def init_state(self, n_neurons: int) -&gt; Float[Array, \"neurons 2\"]:\n \"\"\"Initialize the neurLine truncated