Files
felice-models/api/neuron_models/index.html

6099 lines
403 KiB
HTML

<!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="prev" href="../">
<link rel="next" href="../solver/">
<link rel="icon" href="../../img/felice.png">
<meta name="generator" content="mkdocs-1.6.1, mkdocs-material-9.7.1">
<title>Neuron Models - 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="#neuron-models" 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">
Neuron Models
</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="../../neuron_models/" 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="../../neuron_models/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="../../neuron_models/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="../../neuron_models/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="../../neuron_models/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="../../neuron_models/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>
<li class="md-nav__item md-nav__item--nested">
<input class="md-nav__toggle md-toggle " type="checkbox" id="__nav_2_5" >
<div class="md-nav__link md-nav__container">
<a href="../../neuron_models/lif/" class="md-nav__link ">
<span class="md-ellipsis">
LIF
</span>
</a>
</div>
<nav class="md-nav" data-md-level="2" aria-labelledby="__nav_2_5_label" aria-expanded="false">
<label class="md-nav__title" for="__nav_2_5">
<span class="md-nav__icon md-icon"></span>
LIF
</label>
<ul class="md-nav__list" data-md-scrollfix>
</ul>
</nav>
</li>
</ul>
</nav>
</li>
<li class="md-nav__item md-nav__item--active md-nav__item--nested">
<input class="md-nav__toggle md-toggle " type="checkbox" id="__nav_3" checked>
<div class="md-nav__link md-nav__container">
<a href="../" 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="true">
<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 md-nav__item--active">
<input class="md-nav__toggle md-toggle" type="checkbox" id="__toc">
<label class="md-nav__link md-nav__link--active" for="__toc">
<span class="md-ellipsis">
Neuron Models
</span>
<span class="md-nav__icon md-icon"></span>
</label>
<a href="./" class="md-nav__link md-nav__link--active">
<span class="md-ellipsis">
Neuron Models
</span>
</a>
<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="#felice.neuron_models" class="md-nav__link">
<span class="md-ellipsis">
neuron_models
</span>
</a>
<nav class="md-nav" aria-label="neuron_models">
<ul class="md-nav__list">
<li class="md-nav__item">
<a href="#felice.neuron_models-classes" class="md-nav__link">
<span class="md-ellipsis">
Classes
</span>
</a>
<nav class="md-nav" aria-label="Classes">
<ul class="md-nav__list">
<li class="md-nav__item">
<a href="#felice.neuron_models.Boomerang" class="md-nav__link">
<span class="md-ellipsis">
Boomerang
</span>
</a>
<nav class="md-nav" aria-label="Boomerang">
<ul class="md-nav__list">
<li class="md-nav__item">
<a href="#felice.neuron_models.Boomerang-functions" class="md-nav__link">
<span class="md-ellipsis">
Functions
</span>
</a>
<nav class="md-nav" aria-label="Functions">
<ul class="md-nav__list">
<li class="md-nav__item">
<a href="#felice.neuron_models.Boomerang.__init__" class="md-nav__link">
<span class="md-ellipsis">
__init__
</span>
</a>
</li>
<li class="md-nav__item">
<a href="#felice.neuron_models.Boomerang.init_state" class="md-nav__link">
<span class="md-ellipsis">
init_state
</span>
</a>
</li>
<li class="md-nav__item">
<a href="#felice.neuron_models.Boomerang.dynamics" class="md-nav__link">
<span class="md-ellipsis">
dynamics
</span>
</a>
</li>
<li class="md-nav__item">
<a href="#felice.neuron_models.Boomerang.spike_condition" class="md-nav__link">
<span class="md-ellipsis">
spike_condition
</span>
</a>
</li>
</ul>
</nav>
</li>
</ul>
</nav>
</li>
<li class="md-nav__item">
<a href="#felice.neuron_models.FHNRS" class="md-nav__link">
<span class="md-ellipsis">
FHNRS
</span>
</a>
<nav class="md-nav" aria-label="FHNRS">
<ul class="md-nav__list">
<li class="md-nav__item">
<a href="#felice.neuron_models.FHNRS-functions" class="md-nav__link">
<span class="md-ellipsis">
Functions
</span>
</a>
<nav class="md-nav" aria-label="Functions">
<ul class="md-nav__list">
<li class="md-nav__item">
<a href="#felice.neuron_models.FHNRS.__init__" class="md-nav__link">
<span class="md-ellipsis">
__init__
</span>
</a>
</li>
<li class="md-nav__item">
<a href="#felice.neuron_models.FHNRS.init_state" class="md-nav__link">
<span class="md-ellipsis">
init_state
</span>
</a>
</li>
<li class="md-nav__item">
<a href="#felice.neuron_models.FHNRS.IV_inst" class="md-nav__link">
<span class="md-ellipsis">
IV_inst
</span>
</a>
</li>
<li class="md-nav__item">
<a href="#felice.neuron_models.FHNRS.IV_fast" class="md-nav__link">
<span class="md-ellipsis">
IV_fast
</span>
</a>
</li>
<li class="md-nav__item">
<a href="#felice.neuron_models.FHNRS.IV_slow" class="md-nav__link">
<span class="md-ellipsis">
IV_slow
</span>
</a>
</li>
<li class="md-nav__item">
<a href="#felice.neuron_models.FHNRS.dynamics" class="md-nav__link">
<span class="md-ellipsis">
dynamics
</span>
</a>
</li>
<li class="md-nav__item">
<a href="#felice.neuron_models.FHNRS.spike_condition" class="md-nav__link">
<span class="md-ellipsis">
spike_condition
</span>
</a>
</li>
</ul>
</nav>
</li>
</ul>
</nav>
</li>
<li class="md-nav__item">
<a href="#felice.neuron_models.WereRabbit" class="md-nav__link">
<span class="md-ellipsis">
WereRabbit
</span>
</a>
<nav class="md-nav" aria-label="WereRabbit">
<ul class="md-nav__list">
<li class="md-nav__item">
<a href="#felice.neuron_models.WereRabbit-functions" class="md-nav__link">
<span class="md-ellipsis">
Functions
</span>
</a>
<nav class="md-nav" aria-label="Functions">
<ul class="md-nav__list">
<li class="md-nav__item">
<a href="#felice.neuron_models.WereRabbit.__init__" class="md-nav__link">
<span class="md-ellipsis">
__init__
</span>
</a>
</li>
<li class="md-nav__item">
<a href="#felice.neuron_models.WereRabbit.init_state" class="md-nav__link">
<span class="md-ellipsis">
init_state
</span>
</a>
</li>
<li class="md-nav__item">
<a href="#felice.neuron_models.WereRabbit.vector_field" class="md-nav__link">
<span class="md-ellipsis">
vector_field
</span>
</a>
</li>
<li class="md-nav__item">
<a href="#felice.neuron_models.WereRabbit.dynamics" class="md-nav__link">
<span class="md-ellipsis">
dynamics
</span>
</a>
</li>
<li class="md-nav__item">
<a href="#felice.neuron_models.WereRabbit.spike_condition" class="md-nav__link">
<span class="md-ellipsis">
spike_condition
</span>
</a>
</li>
</ul>
</nav>
</li>
</ul>
</nav>
</li>
</ul>
</nav>
</li>
</ul>
</nav>
</li>
</ul>
</nav>
</li>
<li class="md-nav__item">
<a href="../solver/" class="md-nav__link">
<span class="md-ellipsis">
Solver
</span>
</a>
</li>
<li class="md-nav__item">
<a href="../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="#felice.neuron_models" class="md-nav__link">
<span class="md-ellipsis">
neuron_models
</span>
</a>
<nav class="md-nav" aria-label="neuron_models">
<ul class="md-nav__list">
<li class="md-nav__item">
<a href="#felice.neuron_models-classes" class="md-nav__link">
<span class="md-ellipsis">
Classes
</span>
</a>
<nav class="md-nav" aria-label="Classes">
<ul class="md-nav__list">
<li class="md-nav__item">
<a href="#felice.neuron_models.Boomerang" class="md-nav__link">
<span class="md-ellipsis">
Boomerang
</span>
</a>
<nav class="md-nav" aria-label="Boomerang">
<ul class="md-nav__list">
<li class="md-nav__item">
<a href="#felice.neuron_models.Boomerang-functions" class="md-nav__link">
<span class="md-ellipsis">
Functions
</span>
</a>
<nav class="md-nav" aria-label="Functions">
<ul class="md-nav__list">
<li class="md-nav__item">
<a href="#felice.neuron_models.Boomerang.__init__" class="md-nav__link">
<span class="md-ellipsis">
__init__
</span>
</a>
</li>
<li class="md-nav__item">
<a href="#felice.neuron_models.Boomerang.init_state" class="md-nav__link">
<span class="md-ellipsis">
init_state
</span>
</a>
</li>
<li class="md-nav__item">
<a href="#felice.neuron_models.Boomerang.dynamics" class="md-nav__link">
<span class="md-ellipsis">
dynamics
</span>
</a>
</li>
<li class="md-nav__item">
<a href="#felice.neuron_models.Boomerang.spike_condition" class="md-nav__link">
<span class="md-ellipsis">
spike_condition
</span>
</a>
</li>
</ul>
</nav>
</li>
</ul>
</nav>
</li>
<li class="md-nav__item">
<a href="#felice.neuron_models.FHNRS" class="md-nav__link">
<span class="md-ellipsis">
FHNRS
</span>
</a>
<nav class="md-nav" aria-label="FHNRS">
<ul class="md-nav__list">
<li class="md-nav__item">
<a href="#felice.neuron_models.FHNRS-functions" class="md-nav__link">
<span class="md-ellipsis">
Functions
</span>
</a>
<nav class="md-nav" aria-label="Functions">
<ul class="md-nav__list">
<li class="md-nav__item">
<a href="#felice.neuron_models.FHNRS.__init__" class="md-nav__link">
<span class="md-ellipsis">
__init__
</span>
</a>
</li>
<li class="md-nav__item">
<a href="#felice.neuron_models.FHNRS.init_state" class="md-nav__link">
<span class="md-ellipsis">
init_state
</span>
</a>
</li>
<li class="md-nav__item">
<a href="#felice.neuron_models.FHNRS.IV_inst" class="md-nav__link">
<span class="md-ellipsis">
IV_inst
</span>
</a>
</li>
<li class="md-nav__item">
<a href="#felice.neuron_models.FHNRS.IV_fast" class="md-nav__link">
<span class="md-ellipsis">
IV_fast
</span>
</a>
</li>
<li class="md-nav__item">
<a href="#felice.neuron_models.FHNRS.IV_slow" class="md-nav__link">
<span class="md-ellipsis">
IV_slow
</span>
</a>
</li>
<li class="md-nav__item">
<a href="#felice.neuron_models.FHNRS.dynamics" class="md-nav__link">
<span class="md-ellipsis">
dynamics
</span>
</a>
</li>
<li class="md-nav__item">
<a href="#felice.neuron_models.FHNRS.spike_condition" class="md-nav__link">
<span class="md-ellipsis">
spike_condition
</span>
</a>
</li>
</ul>
</nav>
</li>
</ul>
</nav>
</li>
<li class="md-nav__item">
<a href="#felice.neuron_models.WereRabbit" class="md-nav__link">
<span class="md-ellipsis">
WereRabbit
</span>
</a>
<nav class="md-nav" aria-label="WereRabbit">
<ul class="md-nav__list">
<li class="md-nav__item">
<a href="#felice.neuron_models.WereRabbit-functions" class="md-nav__link">
<span class="md-ellipsis">
Functions
</span>
</a>
<nav class="md-nav" aria-label="Functions">
<ul class="md-nav__list">
<li class="md-nav__item">
<a href="#felice.neuron_models.WereRabbit.__init__" class="md-nav__link">
<span class="md-ellipsis">
__init__
</span>
</a>
</li>
<li class="md-nav__item">
<a href="#felice.neuron_models.WereRabbit.init_state" class="md-nav__link">
<span class="md-ellipsis">
init_state
</span>
</a>
</li>
<li class="md-nav__item">
<a href="#felice.neuron_models.WereRabbit.vector_field" class="md-nav__link">
<span class="md-ellipsis">
vector_field
</span>
</a>
</li>
<li class="md-nav__item">
<a href="#felice.neuron_models.WereRabbit.dynamics" class="md-nav__link">
<span class="md-ellipsis">
dynamics
</span>
</a>
</li>
<li class="md-nav__item">
<a href="#felice.neuron_models.WereRabbit.spike_condition" class="md-nav__link">
<span class="md-ellipsis">
spike_condition
</span>
</a>
</li>
</ul>
</nav>
</li>
</ul>
</nav>
</li>
</ul>
</nav>
</li>
</ul>
</nav>
</li>
</ul>
</nav>
</div>
</div>
</div>
<div class="md-content" data-md-component="content">
<nav class="md-path" aria-label="Navigation" >
<ol class="md-path__list">
<li class="md-path__item">
<a href="../.." class="md-path__link">
<span class="md-ellipsis">
Home
</span>
</a>
</li>
<li class="md-path__item">
<a href="../" class="md-path__link">
<span class="md-ellipsis">
API Reference
</span>
</a>
</li>
</ol>
</nav>
<article class="md-content__inner md-typeset">
<h1 id="neuron-models">Neuron Models</h1>
<div class="doc doc-object doc-module">
<h2 id="felice.neuron_models" class="doc doc-heading">
<code class="doc-symbol doc-symbol-heading doc-symbol-module"></code> <code>felice.neuron_models</code>
</h2>
<div class="doc doc-contents first">
<div class="doc doc-children">
<h3 id="felice.neuron_models-classes">Classes</h3>
<div class="doc doc-object doc-class">
<h4 id="felice.neuron_models.Boomerang" class="doc doc-heading">
<code class="doc-symbol doc-symbol-heading doc-symbol-class"></code> <code>Boomerang</code>
</h4>
<div class="doc doc-contents ">
<p class="doc doc-class-bases">
Bases: <code><span title="equinox.Module">Module</span></code></p>
<details class="mkdocstrings-source">
<summary>Source code in <code>felice/neuron_models/boomerang.py</code></summary>
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal"><a href="#__codelineno-0-10"> 10</a></span>
<span class="normal"><a href="#__codelineno-0-11"> 11</a></span>
<span class="normal"><a href="#__codelineno-0-12"> 12</a></span>
<span class="normal"><a href="#__codelineno-0-13"> 13</a></span>
<span class="normal"><a href="#__codelineno-0-14"> 14</a></span>
<span class="normal"><a href="#__codelineno-0-15"> 15</a></span>
<span class="normal"><a href="#__codelineno-0-16"> 16</a></span>
<span class="normal"><a href="#__codelineno-0-17"> 17</a></span>
<span class="normal"><a href="#__codelineno-0-18"> 18</a></span>
<span class="normal"><a href="#__codelineno-0-19"> 19</a></span>
<span class="normal"><a href="#__codelineno-0-20"> 20</a></span>
<span class="normal"><a href="#__codelineno-0-21"> 21</a></span>
<span class="normal"><a href="#__codelineno-0-22"> 22</a></span>
<span class="normal"><a href="#__codelineno-0-23"> 23</a></span>
<span class="normal"><a href="#__codelineno-0-24"> 24</a></span>
<span class="normal"><a href="#__codelineno-0-25"> 25</a></span>
<span class="normal"><a href="#__codelineno-0-26"> 26</a></span>
<span class="normal"><a href="#__codelineno-0-27"> 27</a></span>
<span class="normal"><a href="#__codelineno-0-28"> 28</a></span>
<span class="normal"><a href="#__codelineno-0-29"> 29</a></span>
<span class="normal"><a href="#__codelineno-0-30"> 30</a></span>
<span class="normal"><a href="#__codelineno-0-31"> 31</a></span>
<span class="normal"><a href="#__codelineno-0-32"> 32</a></span>
<span class="normal"><a href="#__codelineno-0-33"> 33</a></span>
<span class="normal"><a href="#__codelineno-0-34"> 34</a></span>
<span class="normal"><a href="#__codelineno-0-35"> 35</a></span>
<span class="normal"><a href="#__codelineno-0-36"> 36</a></span>
<span class="normal"><a href="#__codelineno-0-37"> 37</a></span>
<span class="normal"><a href="#__codelineno-0-38"> 38</a></span>
<span class="normal"><a href="#__codelineno-0-39"> 39</a></span>
<span class="normal"><a href="#__codelineno-0-40"> 40</a></span>
<span class="normal"><a href="#__codelineno-0-41"> 41</a></span>
<span class="normal"><a href="#__codelineno-0-42"> 42</a></span>
<span class="normal"><a href="#__codelineno-0-43"> 43</a></span>
<span class="normal"><a href="#__codelineno-0-44"> 44</a></span>
<span class="normal"><a href="#__codelineno-0-45"> 45</a></span>
<span class="normal"><a href="#__codelineno-0-46"> 46</a></span>
<span class="normal"><a href="#__codelineno-0-47"> 47</a></span>
<span class="normal"><a href="#__codelineno-0-48"> 48</a></span>
<span class="normal"><a href="#__codelineno-0-49"> 49</a></span>
<span class="normal"><a href="#__codelineno-0-50"> 50</a></span>
<span class="normal"><a href="#__codelineno-0-51"> 51</a></span>
<span class="normal"><a href="#__codelineno-0-52"> 52</a></span>
<span class="normal"><a href="#__codelineno-0-53"> 53</a></span>
<span class="normal"><a href="#__codelineno-0-54"> 54</a></span>
<span class="normal"><a href="#__codelineno-0-55"> 55</a></span>
<span class="normal"><a href="#__codelineno-0-56"> 56</a></span>
<span class="normal"><a href="#__codelineno-0-57"> 57</a></span>
<span class="normal"><a href="#__codelineno-0-58"> 58</a></span>
<span class="normal"><a href="#__codelineno-0-59"> 59</a></span>
<span class="normal"><a href="#__codelineno-0-60"> 60</a></span>
<span class="normal"><a href="#__codelineno-0-61"> 61</a></span>
<span class="normal"><a href="#__codelineno-0-62"> 62</a></span>
<span class="normal"><a href="#__codelineno-0-63"> 63</a></span>
<span class="normal"><a href="#__codelineno-0-64"> 64</a></span>
<span class="normal"><a href="#__codelineno-0-65"> 65</a></span>
<span class="normal"><a href="#__codelineno-0-66"> 66</a></span>
<span class="normal"><a href="#__codelineno-0-67"> 67</a></span>
<span class="normal"><a href="#__codelineno-0-68"> 68</a></span>
<span class="normal"><a href="#__codelineno-0-69"> 69</a></span>
<span class="normal"><a href="#__codelineno-0-70"> 70</a></span>
<span class="normal"><a href="#__codelineno-0-71"> 71</a></span>
<span class="normal"><a href="#__codelineno-0-72"> 72</a></span>
<span class="normal"><a href="#__codelineno-0-73"> 73</a></span>
<span class="normal"><a href="#__codelineno-0-74"> 74</a></span>
<span class="normal"><a href="#__codelineno-0-75"> 75</a></span>
<span class="normal"><a href="#__codelineno-0-76"> 76</a></span>
<span class="normal"><a href="#__codelineno-0-77"> 77</a></span>
<span class="normal"><a href="#__codelineno-0-78"> 78</a></span>
<span class="normal"><a href="#__codelineno-0-79"> 79</a></span>
<span class="normal"><a href="#__codelineno-0-80"> 80</a></span>
<span class="normal"><a href="#__codelineno-0-81"> 81</a></span>
<span class="normal"><a href="#__codelineno-0-82"> 82</a></span>
<span class="normal"><a href="#__codelineno-0-83"> 83</a></span>
<span class="normal"><a href="#__codelineno-0-84"> 84</a></span>
<span class="normal"><a href="#__codelineno-0-85"> 85</a></span>
<span class="normal"><a href="#__codelineno-0-86"> 86</a></span>
<span class="normal"><a href="#__codelineno-0-87"> 87</a></span>
<span class="normal"><a href="#__codelineno-0-88"> 88</a></span>
<span class="normal"><a href="#__codelineno-0-89"> 89</a></span>
<span class="normal"><a href="#__codelineno-0-90"> 90</a></span>
<span class="normal"><a href="#__codelineno-0-91"> 91</a></span>
<span class="normal"><a href="#__codelineno-0-92"> 92</a></span>
<span class="normal"><a href="#__codelineno-0-93"> 93</a></span>
<span class="normal"><a href="#__codelineno-0-94"> 94</a></span>
<span class="normal"><a href="#__codelineno-0-95"> 95</a></span>
<span class="normal"><a href="#__codelineno-0-96"> 96</a></span>
<span class="normal"><a href="#__codelineno-0-97"> 97</a></span>
<span class="normal"><a href="#__codelineno-0-98"> 98</a></span>
<span class="normal"><a href="#__codelineno-0-99"> 99</a></span>
<span class="normal"><a href="#__codelineno-0-100">100</a></span>
<span class="normal"><a href="#__codelineno-0-101">101</a></span>
<span class="normal"><a href="#__codelineno-0-102">102</a></span>
<span class="normal"><a href="#__codelineno-0-103">103</a></span>
<span class="normal"><a href="#__codelineno-0-104">104</a></span>
<span class="normal"><a href="#__codelineno-0-105">105</a></span>
<span class="normal"><a href="#__codelineno-0-106">106</a></span>
<span class="normal"><a href="#__codelineno-0-107">107</a></span>
<span class="normal"><a href="#__codelineno-0-108">108</a></span>
<span class="normal"><a href="#__codelineno-0-109">109</a></span>
<span class="normal"><a href="#__codelineno-0-110">110</a></span>
<span class="normal"><a href="#__codelineno-0-111">111</a></span>
<span class="normal"><a href="#__codelineno-0-112">112</a></span>
<span class="normal"><a href="#__codelineno-0-113">113</a></span>
<span class="normal"><a href="#__codelineno-0-114">114</a></span>
<span class="normal"><a href="#__codelineno-0-115">115</a></span>
<span class="normal"><a href="#__codelineno-0-116">116</a></span>
<span class="normal"><a href="#__codelineno-0-117">117</a></span>
<span class="normal"><a href="#__codelineno-0-118">118</a></span>
<span class="normal"><a href="#__codelineno-0-119">119</a></span>
<span class="normal"><a href="#__codelineno-0-120">120</a></span>
<span class="normal"><a href="#__codelineno-0-121">121</a></span>
<span class="normal"><a href="#__codelineno-0-122">122</a></span>
<span class="normal"><a href="#__codelineno-0-123">123</a></span>
<span class="normal"><a href="#__codelineno-0-124">124</a></span>
<span class="normal"><a href="#__codelineno-0-125">125</a></span>
<span class="normal"><a href="#__codelineno-0-126">126</a></span>
<span class="normal"><a href="#__codelineno-0-127">127</a></span>
<span class="normal"><a href="#__codelineno-0-128">128</a></span>
<span class="normal"><a href="#__codelineno-0-129">129</a></span>
<span class="normal"><a href="#__codelineno-0-130">130</a></span>
<span class="normal"><a href="#__codelineno-0-131">131</a></span>
<span class="normal"><a href="#__codelineno-0-132">132</a></span>
<span class="normal"><a href="#__codelineno-0-133">133</a></span>
<span class="normal"><a href="#__codelineno-0-134">134</a></span>
<span class="normal"><a href="#__codelineno-0-135">135</a></span>
<span class="normal"><a href="#__codelineno-0-136">136</a></span>
<span class="normal"><a href="#__codelineno-0-137">137</a></span>
<span class="normal"><a href="#__codelineno-0-138">138</a></span>
<span class="normal"><a href="#__codelineno-0-139">139</a></span>
<span class="normal"><a href="#__codelineno-0-140">140</a></span>
<span class="normal"><a href="#__codelineno-0-141">141</a></span>
<span class="normal"><a href="#__codelineno-0-142">142</a></span>
<span class="normal"><a href="#__codelineno-0-143">143</a></span>
<span class="normal"><a href="#__codelineno-0-144">144</a></span>
<span class="normal"><a href="#__codelineno-0-145">145</a></span>
<span class="normal"><a href="#__codelineno-0-146">146</a></span>
<span class="normal"><a href="#__codelineno-0-147">147</a></span>
<span class="normal"><a href="#__codelineno-0-148">148</a></span>
<span class="normal"><a href="#__codelineno-0-149">149</a></span>
<span class="normal"><a href="#__codelineno-0-150">150</a></span>
<span class="normal"><a href="#__codelineno-0-151">151</a></span>
<span class="normal"><a href="#__codelineno-0-152">152</a></span>
<span class="normal"><a href="#__codelineno-0-153">153</a></span>
<span class="normal"><a href="#__codelineno-0-154">154</a></span>
<span class="normal"><a href="#__codelineno-0-155">155</a></span>
<span class="normal"><a href="#__codelineno-0-156">156</a></span>
<span class="normal"><a href="#__codelineno-0-157">157</a></span>
<span class="normal"><a href="#__codelineno-0-158">158</a></span>
<span class="normal"><a href="#__codelineno-0-159">159</a></span>
<span class="normal"><a href="#__codelineno-0-160">160</a></span>
<span class="normal"><a href="#__codelineno-0-161">161</a></span>
<span class="normal"><a href="#__codelineno-0-162">162</a></span>
<span class="normal"><a href="#__codelineno-0-163">163</a></span>
<span class="normal"><a href="#__codelineno-0-164">164</a></span>
<span class="normal"><a href="#__codelineno-0-165">165</a></span>
<span class="normal"><a href="#__codelineno-0-166">166</a></span>
<span class="normal"><a href="#__codelineno-0-167">167</a></span>
<span class="normal"><a href="#__codelineno-0-168">168</a></span>
<span class="normal"><a href="#__codelineno-0-169">169</a></span>
<span class="normal"><a href="#__codelineno-0-170">170</a></span>
<span class="normal"><a href="#__codelineno-0-171">171</a></span></pre></div></td><td class="code"><div><pre><span></span><code><span id="__span-0-10"><a id="__codelineno-0-10" name="__codelineno-0-10"></a><span class="k">class</span><span class="w"> </span><span class="nc">Boomerang</span><span class="p">(</span><span class="n">eqx</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
</span><span id="__span-0-11"><a id="__codelineno-0-11" name="__codelineno-0-11"></a> <span class="n">rtol</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="n">eqx</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="n">static</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
</span><span id="__span-0-12"><a id="__codelineno-0-12" name="__codelineno-0-12"></a> <span class="n">atol</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="n">eqx</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="n">static</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
</span><span id="__span-0-13"><a id="__codelineno-0-13" name="__codelineno-0-13"></a>
</span><span id="__span-0-14"><a id="__codelineno-0-14" name="__codelineno-0-14"></a> <span class="n">u0</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="n">eqx</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="n">static</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
</span><span id="__span-0-15"><a id="__codelineno-0-15" name="__codelineno-0-15"></a> <span class="n">v0</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="n">eqx</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="n">static</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
</span><span id="__span-0-16"><a id="__codelineno-0-16" name="__codelineno-0-16"></a>
</span><span id="__span-0-17"><a id="__codelineno-0-17" name="__codelineno-0-17"></a> <span class="n">alpha</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="n">eqx</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="n">static</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span> <span class="c1"># I_n0 / I_bias ratio</span>
</span><span id="__span-0-18"><a id="__codelineno-0-18" name="__codelineno-0-18"></a> <span class="n">beta</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="n">eqx</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="n">static</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span> <span class="c1"># k / U_t (inverse thermal scale)</span>
</span><span id="__span-0-19"><a id="__codelineno-0-19" name="__codelineno-0-19"></a> <span class="n">gamma</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="n">eqx</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="n">static</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span> <span class="c1"># coupling coefficient</span>
</span><span id="__span-0-20"><a id="__codelineno-0-20" name="__codelineno-0-20"></a> <span class="n">rho</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="n">eqx</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="n">static</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span> <span class="c1"># tanh steepness</span>
</span><span id="__span-0-21"><a id="__codelineno-0-21" name="__codelineno-0-21"></a> <span class="n">sigma</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="n">eqx</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="n">static</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span> <span class="c1"># bias scaling (s * I_bias)</span>
</span><span id="__span-0-22"><a id="__codelineno-0-22" name="__codelineno-0-22"></a>
</span><span id="__span-0-23"><a id="__codelineno-0-23" name="__codelineno-0-23"></a> <span class="n">dtype</span><span class="p">:</span> <span class="n">DTypeLike</span> <span class="o">=</span> <span class="n">eqx</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="n">static</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
</span><span id="__span-0-24"><a id="__codelineno-0-24" name="__codelineno-0-24"></a>
</span><span id="__span-0-25"><a id="__codelineno-0-25" name="__codelineno-0-25"></a> <span class="k">def</span><span class="w"> </span><span class="fm">__init__</span><span class="p">(</span>
</span><span id="__span-0-26"><a id="__codelineno-0-26" name="__codelineno-0-26"></a> <span class="bp">self</span><span class="p">,</span>
</span><span id="__span-0-27"><a id="__codelineno-0-27" name="__codelineno-0-27"></a> <span class="o">*</span><span class="p">,</span>
</span><span id="__span-0-28"><a id="__codelineno-0-28" name="__codelineno-0-28"></a> <span class="n">atol</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1e-6</span><span class="p">,</span>
</span><span id="__span-0-29"><a id="__codelineno-0-29" name="__codelineno-0-29"></a> <span class="n">rtol</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1e-4</span><span class="p">,</span>
</span><span id="__span-0-30"><a id="__codelineno-0-30" name="__codelineno-0-30"></a> <span class="n">alpha</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.0129</span><span class="p">,</span>
</span><span id="__span-0-31"><a id="__codelineno-0-31" name="__codelineno-0-31"></a> <span class="n">beta</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">15.6</span><span class="p">,</span>
</span><span id="__span-0-32"><a id="__codelineno-0-32" name="__codelineno-0-32"></a> <span class="n">gamma</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.26</span><span class="p">,</span>
</span><span id="__span-0-33"><a id="__codelineno-0-33" name="__codelineno-0-33"></a> <span class="n">rho</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">30.0</span><span class="p">,</span>
</span><span id="__span-0-34"><a id="__codelineno-0-34" name="__codelineno-0-34"></a> <span class="n">sigma</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.6</span><span class="p">,</span>
</span><span id="__span-0-35"><a id="__codelineno-0-35" name="__codelineno-0-35"></a> <span class="n">dtype</span><span class="p">:</span> <span class="n">DTypeLike</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">float32</span><span class="p">,</span>
</span><span id="__span-0-36"><a id="__codelineno-0-36" name="__codelineno-0-36"></a> <span class="p">):</span>
</span><span id="__span-0-37"><a id="__codelineno-0-37" name="__codelineno-0-37"></a><span class="w"> </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;Initialize the WereRabbit neuron model.</span>
</span><span id="__span-0-38"><a id="__codelineno-0-38" name="__codelineno-0-38"></a>
</span><span id="__span-0-39"><a id="__codelineno-0-39" name="__codelineno-0-39"></a><span class="sd"> Args:</span>
</span><span id="__span-0-40"><a id="__codelineno-0-40" name="__codelineno-0-40"></a><span class="sd"> key: JAX random key for weight initialization.</span>
</span><span id="__span-0-41"><a id="__codelineno-0-41" name="__codelineno-0-41"></a><span class="sd"> n_neurons: Number of neurons in this layer.</span>
</span><span id="__span-0-42"><a id="__codelineno-0-42" name="__codelineno-0-42"></a><span class="sd"> in_size: Number of input connections (excluding recurrent connections).</span>
</span><span id="__span-0-43"><a id="__codelineno-0-43" name="__codelineno-0-43"></a><span class="sd"> wmask: Binary mask defining connectivity pattern of shape (in_plus_neurons, neurons).</span>
</span><span id="__span-0-44"><a id="__codelineno-0-44" name="__codelineno-0-44"></a><span class="sd"> rtol: Relative tolerance for the spiking fixpoint calculation.</span>
</span><span id="__span-0-45"><a id="__codelineno-0-45" name="__codelineno-0-45"></a><span class="sd"> atol: Absolute tolerance for the spiking fixpoint calculation.</span>
</span><span id="__span-0-46"><a id="__codelineno-0-46" name="__codelineno-0-46"></a><span class="sd"> alpha: Current scaling parameter $\alpha = I_{n0}/I_{bias}$ (default: 0.0129)</span>
</span><span id="__span-0-47"><a id="__codelineno-0-47" name="__codelineno-0-47"></a><span class="sd"> beta: Exponential slope $\beta = \kappa/U_t$ (default: 15.6)</span>
</span><span id="__span-0-48"><a id="__codelineno-0-48" name="__codelineno-0-48"></a><span class="sd"> gamma: Coupling parameter $\gamma = 26e^{-2}$</span>
</span><span id="__span-0-49"><a id="__codelineno-0-49" name="__codelineno-0-49"></a><span class="sd"> rho: Steepness of the tanh function $\rho$ (default: 5)</span>
</span><span id="__span-0-50"><a id="__codelineno-0-50" name="__codelineno-0-50"></a><span class="sd"> sigma: Fixpoint distance scaling $\sigma$ (default: 0.6)</span>
</span><span id="__span-0-51"><a id="__codelineno-0-51" name="__codelineno-0-51"></a><span class="sd"> wlim: Limit for weight initialization. If None, uses init_weights.</span>
</span><span id="__span-0-52"><a id="__codelineno-0-52" name="__codelineno-0-52"></a><span class="sd"> wmean: Mean value for weight initialization.</span>
</span><span id="__span-0-53"><a id="__codelineno-0-53" name="__codelineno-0-53"></a><span class="sd"> init_weights: Optional initial weight values. If None, weights are randomly initialized.</span>
</span><span id="__span-0-54"><a id="__codelineno-0-54" name="__codelineno-0-54"></a><span class="sd"> fan_in_mode: Mode for fan-in based weight initialization (&#39;sqrt&#39;, &#39;linear&#39;).</span>
</span><span id="__span-0-55"><a id="__codelineno-0-55" name="__codelineno-0-55"></a><span class="sd"> dtype: Data type for arrays (default: float32).</span>
</span><span id="__span-0-56"><a id="__codelineno-0-56" name="__codelineno-0-56"></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="__span-0-57"><a id="__codelineno-0-57" name="__codelineno-0-57"></a> <span class="bp">self</span><span class="o">.</span><span class="n">dtype</span> <span class="o">=</span> <span class="n">dtype</span>
</span><span id="__span-0-58"><a id="__codelineno-0-58" name="__codelineno-0-58"></a>
</span><span id="__span-0-59"><a id="__codelineno-0-59" name="__codelineno-0-59"></a> <span class="bp">self</span><span class="o">.</span><span class="n">alpha</span> <span class="o">=</span> <span class="n">alpha</span>
</span><span id="__span-0-60"><a id="__codelineno-0-60" name="__codelineno-0-60"></a> <span class="bp">self</span><span class="o">.</span><span class="n">beta</span> <span class="o">=</span> <span class="n">beta</span>
</span><span id="__span-0-61"><a id="__codelineno-0-61" name="__codelineno-0-61"></a> <span class="bp">self</span><span class="o">.</span><span class="n">gamma</span> <span class="o">=</span> <span class="n">gamma</span>
</span><span id="__span-0-62"><a id="__codelineno-0-62" name="__codelineno-0-62"></a> <span class="bp">self</span><span class="o">.</span><span class="n">rho</span> <span class="o">=</span> <span class="n">rho</span>
</span><span id="__span-0-63"><a id="__codelineno-0-63" name="__codelineno-0-63"></a> <span class="bp">self</span><span class="o">.</span><span class="n">sigma</span> <span class="o">=</span> <span class="n">sigma</span>
</span><span id="__span-0-64"><a id="__codelineno-0-64" name="__codelineno-0-64"></a>
</span><span id="__span-0-65"><a id="__codelineno-0-65" name="__codelineno-0-65"></a> <span class="bp">self</span><span class="o">.</span><span class="n">rtol</span> <span class="o">=</span> <span class="n">rtol</span>
</span><span id="__span-0-66"><a id="__codelineno-0-66" name="__codelineno-0-66"></a> <span class="bp">self</span><span class="o">.</span><span class="n">atol</span> <span class="o">=</span> <span class="n">atol</span>
</span><span id="__span-0-67"><a id="__codelineno-0-67" name="__codelineno-0-67"></a>
</span><span id="__span-0-68"><a id="__codelineno-0-68" name="__codelineno-0-68"></a> <span class="k">def</span><span class="w"> </span><span class="nf">fn</span><span class="p">(</span><span class="n">y</span><span class="p">,</span> <span class="n">_</span><span class="p">):</span>
</span><span id="__span-0-69"><a id="__codelineno-0-69" name="__codelineno-0-69"></a> <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">vector_field</span><span class="p">(</span><span class="n">y</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">y</span><span class="p">[</span><span class="mi">1</span><span class="p">])</span>
</span><span id="__span-0-70"><a id="__codelineno-0-70" name="__codelineno-0-70"></a>
</span><span id="__span-0-71"><a id="__codelineno-0-71" name="__codelineno-0-71"></a> <span class="n">solver</span><span class="p">:</span> <span class="n">optx</span><span class="o">.</span><span class="n">AbstractRootFinder</span> <span class="o">=</span> <span class="n">optx</span><span class="o">.</span><span class="n">Newton</span><span class="p">(</span><span class="n">rtol</span><span class="o">=</span><span class="mf">1e-8</span><span class="p">,</span> <span class="n">atol</span><span class="o">=</span><span class="mf">1e-8</span><span class="p">)</span>
</span><span id="__span-0-72"><a id="__codelineno-0-72" name="__codelineno-0-72"></a> <span class="n">y0</span> <span class="o">=</span> <span class="p">(</span><span class="n">jnp</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="mf">0.3</span><span class="p">),</span> <span class="n">jnp</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="mf">0.3</span><span class="p">))</span>
</span><span id="__span-0-73"><a id="__codelineno-0-73" name="__codelineno-0-73"></a> <span class="n">u0</span><span class="p">,</span> <span class="n">v0</span> <span class="o">=</span> <span class="n">optx</span><span class="o">.</span><span class="n">root_find</span><span class="p">(</span><span class="n">fn</span><span class="p">,</span> <span class="n">solver</span><span class="p">,</span> <span class="n">y0</span><span class="p">)</span><span class="o">.</span><span class="n">value</span>
</span><span id="__span-0-74"><a id="__codelineno-0-74" name="__codelineno-0-74"></a> <span class="bp">self</span><span class="o">.</span><span class="n">u0</span> <span class="o">=</span> <span class="n">u0</span><span class="o">.</span><span class="n">item</span><span class="p">()</span>
</span><span id="__span-0-75"><a id="__codelineno-0-75" name="__codelineno-0-75"></a> <span class="bp">self</span><span class="o">.</span><span class="n">v0</span> <span class="o">=</span> <span class="n">v0</span><span class="o">.</span><span class="n">item</span><span class="p">()</span>
</span><span id="__span-0-76"><a id="__codelineno-0-76" name="__codelineno-0-76"></a>
</span><span id="__span-0-77"><a id="__codelineno-0-77" name="__codelineno-0-77"></a> <span class="k">def</span><span class="w"> </span><span class="nf">init_state</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">n_neurons</span><span class="p">:</span> <span class="nb">int</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;neurons 2&quot;</span><span class="p">]:</span>
</span><span id="__span-0-78"><a id="__codelineno-0-78" name="__codelineno-0-78"></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;Initialize the neuron state variables.</span>
</span><span id="__span-0-79"><a id="__codelineno-0-79" name="__codelineno-0-79"></a>
</span><span id="__span-0-80"><a id="__codelineno-0-80" name="__codelineno-0-80"></a><span class="sd"> Args:</span>
</span><span id="__span-0-81"><a id="__codelineno-0-81" name="__codelineno-0-81"></a><span class="sd"> n_neurons: Number of neurons to initialize.</span>
</span><span id="__span-0-82"><a id="__codelineno-0-82" name="__codelineno-0-82"></a>
</span><span id="__span-0-83"><a id="__codelineno-0-83" name="__codelineno-0-83"></a><span class="sd"> Returns:</span>
</span><span id="__span-0-84"><a id="__codelineno-0-84" name="__codelineno-0-84"></a><span class="sd"> Initial state array of shape (neurons, 3) containing [u, v],</span>
</span><span id="__span-0-85"><a id="__codelineno-0-85" name="__codelineno-0-85"></a><span class="sd"> where u and v are the predator/prey membrane voltages.</span>
</span><span id="__span-0-86"><a id="__codelineno-0-86" name="__codelineno-0-86"></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="__span-0-87"><a id="__codelineno-0-87" name="__codelineno-0-87"></a>
</span><span id="__span-0-88"><a id="__codelineno-0-88" name="__codelineno-0-88"></a> <span class="n">u</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">full</span><span class="p">((</span><span class="n">n_neurons</span><span class="p">,),</span> <span class="bp">self</span><span class="o">.</span><span class="n">u0</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dtype</span><span class="p">)</span>
</span><span id="__span-0-89"><a id="__codelineno-0-89" name="__codelineno-0-89"></a> <span class="n">v</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">full</span><span class="p">((</span><span class="n">n_neurons</span><span class="p">,),</span> <span class="bp">self</span><span class="o">.</span><span class="n">v0</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dtype</span><span class="p">)</span>
</span><span id="__span-0-90"><a id="__codelineno-0-90" name="__codelineno-0-90"></a> <span class="n">x</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">u</span><span class="p">,</span> <span class="n">v</span><span class="p">],</span> <span class="n">axis</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
</span><span id="__span-0-91"><a id="__codelineno-0-91" name="__codelineno-0-91"></a> <span class="k">return</span> <span class="n">x</span>
</span><span id="__span-0-92"><a id="__codelineno-0-92" name="__codelineno-0-92"></a>
</span><span id="__span-0-93"><a id="__codelineno-0-93" name="__codelineno-0-93"></a> <span class="k">def</span><span class="w"> </span><span class="nf">vector_field</span><span class="p">(</span>
</span><span id="__span-0-94"><a id="__codelineno-0-94" name="__codelineno-0-94"></a> <span class="bp">self</span><span class="p">,</span> <span class="n">u</span><span class="p">:</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;...&quot;</span><span class="p">],</span> <span class="n">v</span><span class="p">:</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;...&quot;</span><span class="p">]</span>
</span><span id="__span-0-95"><a id="__codelineno-0-95" name="__codelineno-0-95"></a> <span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Tuple</span><span class="p">[</span><span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;...&quot;</span><span class="p">],</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;...&quot;</span><span class="p">]]:</span>
</span><span id="__span-0-96"><a id="__codelineno-0-96" name="__codelineno-0-96"></a> <span class="n">alpha</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">alpha</span>
</span><span id="__span-0-97"><a id="__codelineno-0-97" name="__codelineno-0-97"></a> <span class="n">beta</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">beta</span>
</span><span id="__span-0-98"><a id="__codelineno-0-98" name="__codelineno-0-98"></a> <span class="n">gamma</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">gamma</span>
</span><span id="__span-0-99"><a id="__codelineno-0-99" name="__codelineno-0-99"></a> <span class="n">sigma</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">sigma</span>
</span><span id="__span-0-100"><a id="__codelineno-0-100" name="__codelineno-0-100"></a> <span class="n">rho</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">rho</span>
</span><span id="__span-0-101"><a id="__codelineno-0-101" name="__codelineno-0-101"></a>
</span><span id="__span-0-102"><a id="__codelineno-0-102" name="__codelineno-0-102"></a> <span class="n">z</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">tanh</span><span class="p">(</span><span class="n">rho</span> <span class="o">*</span> <span class="p">(</span><span class="n">v</span> <span class="o">-</span> <span class="n">u</span><span class="p">))</span>
</span><span id="__span-0-103"><a id="__codelineno-0-103" name="__codelineno-0-103"></a> <span class="n">du</span> <span class="o">=</span> <span class="p">(</span><span class="mi">1</span> <span class="o">-</span> <span class="n">alpha</span> <span class="o">*</span> <span class="n">jnp</span><span class="o">.</span><span class="n">exp</span><span class="p">(</span><span class="n">beta</span> <span class="o">*</span> <span class="n">v</span><span class="p">)</span> <span class="o">*</span> <span class="p">(</span><span class="mi">1</span> <span class="o">-</span> <span class="n">gamma</span> <span class="o">*</span> <span class="p">(</span><span class="mf">0.3</span> <span class="o">-</span> <span class="n">u</span><span class="p">)))</span> <span class="o">+</span> <span class="n">sigma</span> <span class="o">*</span> <span class="n">z</span>
</span><span id="__span-0-104"><a id="__codelineno-0-104" name="__codelineno-0-104"></a> <span class="n">dv</span> <span class="o">=</span> <span class="p">(</span><span class="o">-</span><span class="mi">1</span> <span class="o">+</span> <span class="n">alpha</span> <span class="o">*</span> <span class="n">jnp</span><span class="o">.</span><span class="n">exp</span><span class="p">(</span><span class="n">beta</span> <span class="o">*</span> <span class="n">u</span><span class="p">)</span> <span class="o">*</span> <span class="p">(</span><span class="mi">1</span> <span class="o">+</span> <span class="n">gamma</span> <span class="o">*</span> <span class="p">(</span><span class="mf">0.3</span> <span class="o">-</span> <span class="n">v</span><span class="p">)))</span> <span class="o">+</span> <span class="n">sigma</span> <span class="o">*</span> <span class="n">z</span>
</span><span id="__span-0-105"><a id="__codelineno-0-105" name="__codelineno-0-105"></a>
</span><span id="__span-0-106"><a id="__codelineno-0-106" name="__codelineno-0-106"></a> <span class="k">return</span> <span class="n">du</span><span class="p">,</span> <span class="n">dv</span>
</span><span id="__span-0-107"><a id="__codelineno-0-107" name="__codelineno-0-107"></a>
</span><span id="__span-0-108"><a id="__codelineno-0-108" name="__codelineno-0-108"></a> <span class="k">def</span><span class="w"> </span><span class="nf">dynamics</span><span class="p">(</span>
</span><span id="__span-0-109"><a id="__codelineno-0-109" name="__codelineno-0-109"></a> <span class="bp">self</span><span class="p">,</span>
</span><span id="__span-0-110"><a id="__codelineno-0-110" name="__codelineno-0-110"></a> <span class="n">t</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
</span><span id="__span-0-111"><a id="__codelineno-0-111" name="__codelineno-0-111"></a> <span class="n">y</span><span class="p">:</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;neurons 2&quot;</span><span class="p">],</span>
</span><span id="__span-0-112"><a id="__codelineno-0-112" name="__codelineno-0-112"></a> <span class="n">args</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">],</span>
</span><span id="__span-0-113"><a id="__codelineno-0-113" name="__codelineno-0-113"></a> <span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;neurons 2&quot;</span><span class="p">]:</span>
</span><span id="__span-0-114"><a id="__codelineno-0-114" name="__codelineno-0-114"></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;Compute time derivatives of the neuron state variables.</span>
</span><span id="__span-0-115"><a id="__codelineno-0-115" name="__codelineno-0-115"></a>
</span><span id="__span-0-116"><a id="__codelineno-0-116" name="__codelineno-0-116"></a><span class="sd"> This implements the WereRabbit dynamics</span>
</span><span id="__span-0-117"><a id="__codelineno-0-117" name="__codelineno-0-117"></a>
</span><span id="__span-0-118"><a id="__codelineno-0-118" name="__codelineno-0-118"></a><span class="sd"> - du/dt: Predator dynamics</span>
</span><span id="__span-0-119"><a id="__codelineno-0-119" name="__codelineno-0-119"></a><span class="sd"> - dv/dt: WerePrey dynamics</span>
</span><span id="__span-0-120"><a id="__codelineno-0-120" name="__codelineno-0-120"></a>
</span><span id="__span-0-121"><a id="__codelineno-0-121" name="__codelineno-0-121"></a><span class="sd"> Args:</span>
</span><span id="__span-0-122"><a id="__codelineno-0-122" name="__codelineno-0-122"></a><span class="sd"> t: Current simulation time (unused but required by framework).</span>
</span><span id="__span-0-123"><a id="__codelineno-0-123" name="__codelineno-0-123"></a><span class="sd"> y: State array of shape (neurons, 2) containing [u, v].</span>
</span><span id="__span-0-124"><a id="__codelineno-0-124" name="__codelineno-0-124"></a><span class="sd"> args: Additional arguments (unused but required by framework).</span>
</span><span id="__span-0-125"><a id="__codelineno-0-125" name="__codelineno-0-125"></a>
</span><span id="__span-0-126"><a id="__codelineno-0-126" name="__codelineno-0-126"></a><span class="sd"> Returns:</span>
</span><span id="__span-0-127"><a id="__codelineno-0-127" name="__codelineno-0-127"></a><span class="sd"> Time derivatives of shape (neurons, 2) containing [du/dt, dv/dt].</span>
</span><span id="__span-0-128"><a id="__codelineno-0-128" name="__codelineno-0-128"></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="__span-0-129"><a id="__codelineno-0-129" name="__codelineno-0-129"></a> <span class="n">u</span> <span class="o">=</span> <span class="n">y</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span>
</span><span id="__span-0-130"><a id="__codelineno-0-130" name="__codelineno-0-130"></a> <span class="n">v</span> <span class="o">=</span> <span class="n">y</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">]</span>
</span><span id="__span-0-131"><a id="__codelineno-0-131" name="__codelineno-0-131"></a>
</span><span id="__span-0-132"><a id="__codelineno-0-132" name="__codelineno-0-132"></a> <span class="n">du</span><span class="p">,</span> <span class="n">dv</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">vector_field</span><span class="p">(</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">)</span>
</span><span id="__span-0-133"><a id="__codelineno-0-133" name="__codelineno-0-133"></a> <span class="n">dxdt</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">du</span><span class="p">,</span> <span class="n">dv</span><span class="p">],</span> <span class="n">axis</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
</span><span id="__span-0-134"><a id="__codelineno-0-134" name="__codelineno-0-134"></a>
</span><span id="__span-0-135"><a id="__codelineno-0-135" name="__codelineno-0-135"></a> <span class="k">return</span> <span class="n">dxdt</span>
</span><span id="__span-0-136"><a id="__codelineno-0-136" name="__codelineno-0-136"></a>
</span><span id="__span-0-137"><a id="__codelineno-0-137" name="__codelineno-0-137"></a> <span class="k">def</span><span class="w"> </span><span class="nf">spike_condition</span><span class="p">(</span>
</span><span id="__span-0-138"><a id="__codelineno-0-138" name="__codelineno-0-138"></a> <span class="bp">self</span><span class="p">,</span>
</span><span id="__span-0-139"><a id="__codelineno-0-139" name="__codelineno-0-139"></a> <span class="n">t</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
</span><span id="__span-0-140"><a id="__codelineno-0-140" name="__codelineno-0-140"></a> <span class="n">y</span><span class="p">:</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;neurons 2&quot;</span><span class="p">],</span>
</span><span id="__span-0-141"><a id="__codelineno-0-141" name="__codelineno-0-141"></a> <span class="o">**</span><span class="n">kwargs</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">],</span>
</span><span id="__span-0-142"><a id="__codelineno-0-142" name="__codelineno-0-142"></a> <span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot; neurons&quot;</span><span class="p">]:</span>
</span><span id="__span-0-143"><a id="__codelineno-0-143" name="__codelineno-0-143"></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;Compute spike condition for event detection.</span>
</span><span id="__span-0-144"><a id="__codelineno-0-144" name="__codelineno-0-144"></a>
</span><span id="__span-0-145"><a id="__codelineno-0-145" name="__codelineno-0-145"></a><span class="sd"> A spike is triggered when the system reach to a fixpoint.</span>
</span><span id="__span-0-146"><a id="__codelineno-0-146" name="__codelineno-0-146"></a>
</span><span id="__span-0-147"><a id="__codelineno-0-147" name="__codelineno-0-147"></a><span class="sd"> INFO:</span>
</span><span id="__span-0-148"><a id="__codelineno-0-148" name="__codelineno-0-148"></a><span class="sd"> `has_spiked` is use to the system don&#39;t detect a continuos</span>
</span><span id="__span-0-149"><a id="__codelineno-0-149" name="__codelineno-0-149"></a><span class="sd"> spike when reach a fixpoint.</span>
</span><span id="__span-0-150"><a id="__codelineno-0-150" name="__codelineno-0-150"></a>
</span><span id="__span-0-151"><a id="__codelineno-0-151" name="__codelineno-0-151"></a><span class="sd"> Args:</span>
</span><span id="__span-0-152"><a id="__codelineno-0-152" name="__codelineno-0-152"></a><span class="sd"> t: Current simulation time (unused but required by the framework).</span>
</span><span id="__span-0-153"><a id="__codelineno-0-153" name="__codelineno-0-153"></a><span class="sd"> y: State array of shape (neurons, 3) containing [u, v, has_spiked].</span>
</span><span id="__span-0-154"><a id="__codelineno-0-154" name="__codelineno-0-154"></a><span class="sd"> **kwargs: Additional keyword arguments (unused).</span>
</span><span id="__span-0-155"><a id="__codelineno-0-155" name="__codelineno-0-155"></a>
</span><span id="__span-0-156"><a id="__codelineno-0-156" name="__codelineno-0-156"></a><span class="sd"> Returns:</span>
</span><span id="__span-0-157"><a id="__codelineno-0-157" name="__codelineno-0-157"></a><span class="sd"> Spike condition array of shape (neurons,). Positive values indicate spike.</span>
</span><span id="__span-0-158"><a id="__codelineno-0-158" name="__codelineno-0-158"></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="__span-0-159"><a id="__codelineno-0-159" name="__codelineno-0-159"></a> <span class="n">_atol</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">atol</span>
</span><span id="__span-0-160"><a id="__codelineno-0-160" name="__codelineno-0-160"></a> <span class="n">_rtol</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">rtol</span>
</span><span id="__span-0-161"><a id="__codelineno-0-161" name="__codelineno-0-161"></a> <span class="n">_norm</span> <span class="o">=</span> <span class="n">optx</span><span class="o">.</span><span class="n">rms_norm</span>
</span><span id="__span-0-162"><a id="__codelineno-0-162" name="__codelineno-0-162"></a>
</span><span id="__span-0-163"><a id="__codelineno-0-163" name="__codelineno-0-163"></a> <span class="n">vf</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">dynamics</span><span class="p">(</span><span class="n">t</span><span class="p">,</span> <span class="n">y</span><span class="p">,</span> <span class="p">{})</span>
</span><span id="__span-0-164"><a id="__codelineno-0-164" name="__codelineno-0-164"></a>
</span><span id="__span-0-165"><a id="__codelineno-0-165" name="__codelineno-0-165"></a> <span class="nd">@jax</span><span class="o">.</span><span class="n">vmap</span>
</span><span id="__span-0-166"><a id="__codelineno-0-166" name="__codelineno-0-166"></a> <span class="k">def</span><span class="w"> </span><span class="nf">calculate_norm</span><span class="p">(</span><span class="n">vf</span><span class="p">,</span> <span class="n">y</span><span class="p">):</span>
</span><span id="__span-0-167"><a id="__codelineno-0-167" name="__codelineno-0-167"></a> <span class="k">return</span> <span class="n">_atol</span> <span class="o">+</span> <span class="n">_rtol</span> <span class="o">*</span> <span class="n">_norm</span><span class="p">(</span><span class="n">y</span><span class="p">)</span> <span class="o">-</span> <span class="n">_norm</span><span class="p">(</span><span class="n">vf</span><span class="p">)</span>
</span><span id="__span-0-168"><a id="__codelineno-0-168" name="__codelineno-0-168"></a>
</span><span id="__span-0-169"><a id="__codelineno-0-169" name="__codelineno-0-169"></a> <span class="n">base_cond</span> <span class="o">=</span> <span class="n">calculate_norm</span><span class="p">(</span><span class="n">vf</span><span class="p">,</span> <span class="n">y</span><span class="p">)</span><span class="o">.</span><span class="n">repeat</span><span class="p">(</span><span class="mi">2</span><span class="p">)</span>
</span><span id="__span-0-170"><a id="__codelineno-0-170" name="__codelineno-0-170"></a>
</span><span id="__span-0-171"><a id="__codelineno-0-171" name="__codelineno-0-171"></a> <span class="k">return</span> <span class="n">base_cond</span>
</span></code></pre></div></td></tr></table></div>
</details>
<div class="doc doc-children">
<h5 id="felice.neuron_models.Boomerang-functions">Functions</h5>
<div class="doc doc-object doc-function">
<h6 id="felice.neuron_models.Boomerang.__init__" class="doc doc-heading">
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <code class="highlight language-python"><span class="fm">__init__</span><span class="p">(</span><span class="o">*</span><span class="p">,</span> <span class="n">atol</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1e-06</span><span class="p">,</span> <span class="n">rtol</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.0001</span><span class="p">,</span> <span class="n">alpha</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.0129</span><span class="p">,</span> <span class="n">beta</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">15.6</span><span class="p">,</span> <span class="n">gamma</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.26</span><span class="p">,</span> <span class="n">rho</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">30.0</span><span class="p">,</span> <span class="n">sigma</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.6</span><span class="p">,</span> <span class="n">dtype</span><span class="p">:</span> <span class="n">DTypeLike</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">float32</span><span class="p">)</span></code>
</h6>
<div class="doc doc-contents ">
<p>Initialize the WereRabbit neuron model.</p>
<p><span class="doc-section-title">Parameters:</span></p>
<table>
<thead>
<tr>
<th>Name</th>
<th>Type</th>
<th>Description</th>
<th>Default</th>
</tr>
</thead>
<tbody>
<tr class="doc-section-item">
<td>
<code>key</code>
</td>
<td>
</td>
<td>
<div class="doc-md-description">
<p>JAX random key for weight initialization.</p>
</div>
</td>
<td>
<em>required</em>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>n_neurons</code>
</td>
<td>
</td>
<td>
<div class="doc-md-description">
<p>Number of neurons in this layer.</p>
</div>
</td>
<td>
<em>required</em>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>in_size</code>
</td>
<td>
</td>
<td>
<div class="doc-md-description">
<p>Number of input connections (excluding recurrent connections).</p>
</div>
</td>
<td>
<em>required</em>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>wmask</code>
</td>
<td>
</td>
<td>
<div class="doc-md-description">
<p>Binary mask defining connectivity pattern of shape (in_plus_neurons, neurons).</p>
</div>
</td>
<td>
<em>required</em>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>rtol</code>
</td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Relative tolerance for the spiking fixpoint calculation.</p>
</div>
</td>
<td>
<code>0.0001</code>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>atol</code>
</td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Absolute tolerance for the spiking fixpoint calculation.</p>
</div>
</td>
<td>
<code>1e-06</code>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>alpha</code>
</td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Current scaling parameter <span class="arithmatex">\(\alpha = I_{n0}/I_{bias}\)</span> (default: 0.0129)</p>
</div>
</td>
<td>
<code>0.0129</code>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>beta</code>
</td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Exponential slope <span class="arithmatex">\(\beta = \kappa/U_t\)</span> (default: 15.6)</p>
</div>
</td>
<td>
<code>15.6</code>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>gamma</code>
</td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Coupling parameter <span class="arithmatex">\(\gamma = 26e^{-2}\)</span></p>
</div>
</td>
<td>
<code>0.26</code>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>rho</code>
</td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Steepness of the tanh function <span class="arithmatex">\(\rho\)</span> (default: 5)</p>
</div>
</td>
<td>
<code>30.0</code>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>sigma</code>
</td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Fixpoint distance scaling <span class="arithmatex">\(\sigma\)</span> (default: 0.6)</p>
</div>
</td>
<td>
<code>0.6</code>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>wlim</code>
</td>
<td>
</td>
<td>
<div class="doc-md-description">
<p>Limit for weight initialization. If None, uses init_weights.</p>
</div>
</td>
<td>
<em>required</em>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>wmean</code>
</td>
<td>
</td>
<td>
<div class="doc-md-description">
<p>Mean value for weight initialization.</p>
</div>
</td>
<td>
<em>required</em>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>init_weights</code>
</td>
<td>
</td>
<td>
<div class="doc-md-description">
<p>Optional initial weight values. If None, weights are randomly initialized.</p>
</div>
</td>
<td>
<em>required</em>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>fan_in_mode</code>
</td>
<td>
</td>
<td>
<div class="doc-md-description">
<p>Mode for fan-in based weight initialization ('sqrt', 'linear').</p>
</div>
</td>
<td>
<em>required</em>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>dtype</code>
</td>
<td>
<code><span title="jaxtyping.DTypeLike">DTypeLike</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Data type for arrays (default: float32).</p>
</div>
</td>
<td>
<code><span title="jax.numpy.float32">float32</span></code>
</td>
</tr>
</tbody>
</table>
<details class="mkdocstrings-source">
<summary>Source code in <code>felice/neuron_models/boomerang.py</code></summary>
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal"><a href="#__codelineno-0-25">25</a></span>
<span class="normal"><a href="#__codelineno-0-26">26</a></span>
<span class="normal"><a href="#__codelineno-0-27">27</a></span>
<span class="normal"><a href="#__codelineno-0-28">28</a></span>
<span class="normal"><a href="#__codelineno-0-29">29</a></span>
<span class="normal"><a href="#__codelineno-0-30">30</a></span>
<span class="normal"><a href="#__codelineno-0-31">31</a></span>
<span class="normal"><a href="#__codelineno-0-32">32</a></span>
<span class="normal"><a href="#__codelineno-0-33">33</a></span>
<span class="normal"><a href="#__codelineno-0-34">34</a></span>
<span class="normal"><a href="#__codelineno-0-35">35</a></span>
<span class="normal"><a href="#__codelineno-0-36">36</a></span>
<span class="normal"><a href="#__codelineno-0-37">37</a></span>
<span class="normal"><a href="#__codelineno-0-38">38</a></span>
<span class="normal"><a href="#__codelineno-0-39">39</a></span>
<span class="normal"><a href="#__codelineno-0-40">40</a></span>
<span class="normal"><a href="#__codelineno-0-41">41</a></span>
<span class="normal"><a href="#__codelineno-0-42">42</a></span>
<span class="normal"><a href="#__codelineno-0-43">43</a></span>
<span class="normal"><a href="#__codelineno-0-44">44</a></span>
<span class="normal"><a href="#__codelineno-0-45">45</a></span>
<span class="normal"><a href="#__codelineno-0-46">46</a></span>
<span class="normal"><a href="#__codelineno-0-47">47</a></span>
<span class="normal"><a href="#__codelineno-0-48">48</a></span>
<span class="normal"><a href="#__codelineno-0-49">49</a></span>
<span class="normal"><a href="#__codelineno-0-50">50</a></span>
<span class="normal"><a href="#__codelineno-0-51">51</a></span>
<span class="normal"><a href="#__codelineno-0-52">52</a></span>
<span class="normal"><a href="#__codelineno-0-53">53</a></span>
<span class="normal"><a href="#__codelineno-0-54">54</a></span>
<span class="normal"><a href="#__codelineno-0-55">55</a></span>
<span class="normal"><a href="#__codelineno-0-56">56</a></span>
<span class="normal"><a href="#__codelineno-0-57">57</a></span>
<span class="normal"><a href="#__codelineno-0-58">58</a></span>
<span class="normal"><a href="#__codelineno-0-59">59</a></span>
<span class="normal"><a href="#__codelineno-0-60">60</a></span>
<span class="normal"><a href="#__codelineno-0-61">61</a></span>
<span class="normal"><a href="#__codelineno-0-62">62</a></span>
<span class="normal"><a href="#__codelineno-0-63">63</a></span>
<span class="normal"><a href="#__codelineno-0-64">64</a></span>
<span class="normal"><a href="#__codelineno-0-65">65</a></span>
<span class="normal"><a href="#__codelineno-0-66">66</a></span>
<span class="normal"><a href="#__codelineno-0-67">67</a></span>
<span class="normal"><a href="#__codelineno-0-68">68</a></span>
<span class="normal"><a href="#__codelineno-0-69">69</a></span>
<span class="normal"><a href="#__codelineno-0-70">70</a></span>
<span class="normal"><a href="#__codelineno-0-71">71</a></span>
<span class="normal"><a href="#__codelineno-0-72">72</a></span>
<span class="normal"><a href="#__codelineno-0-73">73</a></span>
<span class="normal"><a href="#__codelineno-0-74">74</a></span>
<span class="normal"><a href="#__codelineno-0-75">75</a></span></pre></div></td><td class="code"><div><pre><span></span><code><span id="__span-0-25"><a id="__codelineno-0-25" name="__codelineno-0-25"></a><span class="k">def</span><span class="w"> </span><span class="fm">__init__</span><span class="p">(</span>
</span><span id="__span-0-26"><a id="__codelineno-0-26" name="__codelineno-0-26"></a> <span class="bp">self</span><span class="p">,</span>
</span><span id="__span-0-27"><a id="__codelineno-0-27" name="__codelineno-0-27"></a> <span class="o">*</span><span class="p">,</span>
</span><span id="__span-0-28"><a id="__codelineno-0-28" name="__codelineno-0-28"></a> <span class="n">atol</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1e-6</span><span class="p">,</span>
</span><span id="__span-0-29"><a id="__codelineno-0-29" name="__codelineno-0-29"></a> <span class="n">rtol</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1e-4</span><span class="p">,</span>
</span><span id="__span-0-30"><a id="__codelineno-0-30" name="__codelineno-0-30"></a> <span class="n">alpha</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.0129</span><span class="p">,</span>
</span><span id="__span-0-31"><a id="__codelineno-0-31" name="__codelineno-0-31"></a> <span class="n">beta</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">15.6</span><span class="p">,</span>
</span><span id="__span-0-32"><a id="__codelineno-0-32" name="__codelineno-0-32"></a> <span class="n">gamma</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.26</span><span class="p">,</span>
</span><span id="__span-0-33"><a id="__codelineno-0-33" name="__codelineno-0-33"></a> <span class="n">rho</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">30.0</span><span class="p">,</span>
</span><span id="__span-0-34"><a id="__codelineno-0-34" name="__codelineno-0-34"></a> <span class="n">sigma</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.6</span><span class="p">,</span>
</span><span id="__span-0-35"><a id="__codelineno-0-35" name="__codelineno-0-35"></a> <span class="n">dtype</span><span class="p">:</span> <span class="n">DTypeLike</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">float32</span><span class="p">,</span>
</span><span id="__span-0-36"><a id="__codelineno-0-36" name="__codelineno-0-36"></a><span class="p">):</span>
</span><span id="__span-0-37"><a id="__codelineno-0-37" name="__codelineno-0-37"></a><span class="w"> </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;Initialize the WereRabbit neuron model.</span>
</span><span id="__span-0-38"><a id="__codelineno-0-38" name="__codelineno-0-38"></a>
</span><span id="__span-0-39"><a id="__codelineno-0-39" name="__codelineno-0-39"></a><span class="sd"> Args:</span>
</span><span id="__span-0-40"><a id="__codelineno-0-40" name="__codelineno-0-40"></a><span class="sd"> key: JAX random key for weight initialization.</span>
</span><span id="__span-0-41"><a id="__codelineno-0-41" name="__codelineno-0-41"></a><span class="sd"> n_neurons: Number of neurons in this layer.</span>
</span><span id="__span-0-42"><a id="__codelineno-0-42" name="__codelineno-0-42"></a><span class="sd"> in_size: Number of input connections (excluding recurrent connections).</span>
</span><span id="__span-0-43"><a id="__codelineno-0-43" name="__codelineno-0-43"></a><span class="sd"> wmask: Binary mask defining connectivity pattern of shape (in_plus_neurons, neurons).</span>
</span><span id="__span-0-44"><a id="__codelineno-0-44" name="__codelineno-0-44"></a><span class="sd"> rtol: Relative tolerance for the spiking fixpoint calculation.</span>
</span><span id="__span-0-45"><a id="__codelineno-0-45" name="__codelineno-0-45"></a><span class="sd"> atol: Absolute tolerance for the spiking fixpoint calculation.</span>
</span><span id="__span-0-46"><a id="__codelineno-0-46" name="__codelineno-0-46"></a><span class="sd"> alpha: Current scaling parameter $\alpha = I_{n0}/I_{bias}$ (default: 0.0129)</span>
</span><span id="__span-0-47"><a id="__codelineno-0-47" name="__codelineno-0-47"></a><span class="sd"> beta: Exponential slope $\beta = \kappa/U_t$ (default: 15.6)</span>
</span><span id="__span-0-48"><a id="__codelineno-0-48" name="__codelineno-0-48"></a><span class="sd"> gamma: Coupling parameter $\gamma = 26e^{-2}$</span>
</span><span id="__span-0-49"><a id="__codelineno-0-49" name="__codelineno-0-49"></a><span class="sd"> rho: Steepness of the tanh function $\rho$ (default: 5)</span>
</span><span id="__span-0-50"><a id="__codelineno-0-50" name="__codelineno-0-50"></a><span class="sd"> sigma: Fixpoint distance scaling $\sigma$ (default: 0.6)</span>
</span><span id="__span-0-51"><a id="__codelineno-0-51" name="__codelineno-0-51"></a><span class="sd"> wlim: Limit for weight initialization. If None, uses init_weights.</span>
</span><span id="__span-0-52"><a id="__codelineno-0-52" name="__codelineno-0-52"></a><span class="sd"> wmean: Mean value for weight initialization.</span>
</span><span id="__span-0-53"><a id="__codelineno-0-53" name="__codelineno-0-53"></a><span class="sd"> init_weights: Optional initial weight values. If None, weights are randomly initialized.</span>
</span><span id="__span-0-54"><a id="__codelineno-0-54" name="__codelineno-0-54"></a><span class="sd"> fan_in_mode: Mode for fan-in based weight initialization (&#39;sqrt&#39;, &#39;linear&#39;).</span>
</span><span id="__span-0-55"><a id="__codelineno-0-55" name="__codelineno-0-55"></a><span class="sd"> dtype: Data type for arrays (default: float32).</span>
</span><span id="__span-0-56"><a id="__codelineno-0-56" name="__codelineno-0-56"></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="__span-0-57"><a id="__codelineno-0-57" name="__codelineno-0-57"></a> <span class="bp">self</span><span class="o">.</span><span class="n">dtype</span> <span class="o">=</span> <span class="n">dtype</span>
</span><span id="__span-0-58"><a id="__codelineno-0-58" name="__codelineno-0-58"></a>
</span><span id="__span-0-59"><a id="__codelineno-0-59" name="__codelineno-0-59"></a> <span class="bp">self</span><span class="o">.</span><span class="n">alpha</span> <span class="o">=</span> <span class="n">alpha</span>
</span><span id="__span-0-60"><a id="__codelineno-0-60" name="__codelineno-0-60"></a> <span class="bp">self</span><span class="o">.</span><span class="n">beta</span> <span class="o">=</span> <span class="n">beta</span>
</span><span id="__span-0-61"><a id="__codelineno-0-61" name="__codelineno-0-61"></a> <span class="bp">self</span><span class="o">.</span><span class="n">gamma</span> <span class="o">=</span> <span class="n">gamma</span>
</span><span id="__span-0-62"><a id="__codelineno-0-62" name="__codelineno-0-62"></a> <span class="bp">self</span><span class="o">.</span><span class="n">rho</span> <span class="o">=</span> <span class="n">rho</span>
</span><span id="__span-0-63"><a id="__codelineno-0-63" name="__codelineno-0-63"></a> <span class="bp">self</span><span class="o">.</span><span class="n">sigma</span> <span class="o">=</span> <span class="n">sigma</span>
</span><span id="__span-0-64"><a id="__codelineno-0-64" name="__codelineno-0-64"></a>
</span><span id="__span-0-65"><a id="__codelineno-0-65" name="__codelineno-0-65"></a> <span class="bp">self</span><span class="o">.</span><span class="n">rtol</span> <span class="o">=</span> <span class="n">rtol</span>
</span><span id="__span-0-66"><a id="__codelineno-0-66" name="__codelineno-0-66"></a> <span class="bp">self</span><span class="o">.</span><span class="n">atol</span> <span class="o">=</span> <span class="n">atol</span>
</span><span id="__span-0-67"><a id="__codelineno-0-67" name="__codelineno-0-67"></a>
</span><span id="__span-0-68"><a id="__codelineno-0-68" name="__codelineno-0-68"></a> <span class="k">def</span><span class="w"> </span><span class="nf">fn</span><span class="p">(</span><span class="n">y</span><span class="p">,</span> <span class="n">_</span><span class="p">):</span>
</span><span id="__span-0-69"><a id="__codelineno-0-69" name="__codelineno-0-69"></a> <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">vector_field</span><span class="p">(</span><span class="n">y</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">y</span><span class="p">[</span><span class="mi">1</span><span class="p">])</span>
</span><span id="__span-0-70"><a id="__codelineno-0-70" name="__codelineno-0-70"></a>
</span><span id="__span-0-71"><a id="__codelineno-0-71" name="__codelineno-0-71"></a> <span class="n">solver</span><span class="p">:</span> <span class="n">optx</span><span class="o">.</span><span class="n">AbstractRootFinder</span> <span class="o">=</span> <span class="n">optx</span><span class="o">.</span><span class="n">Newton</span><span class="p">(</span><span class="n">rtol</span><span class="o">=</span><span class="mf">1e-8</span><span class="p">,</span> <span class="n">atol</span><span class="o">=</span><span class="mf">1e-8</span><span class="p">)</span>
</span><span id="__span-0-72"><a id="__codelineno-0-72" name="__codelineno-0-72"></a> <span class="n">y0</span> <span class="o">=</span> <span class="p">(</span><span class="n">jnp</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="mf">0.3</span><span class="p">),</span> <span class="n">jnp</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="mf">0.3</span><span class="p">))</span>
</span><span id="__span-0-73"><a id="__codelineno-0-73" name="__codelineno-0-73"></a> <span class="n">u0</span><span class="p">,</span> <span class="n">v0</span> <span class="o">=</span> <span class="n">optx</span><span class="o">.</span><span class="n">root_find</span><span class="p">(</span><span class="n">fn</span><span class="p">,</span> <span class="n">solver</span><span class="p">,</span> <span class="n">y0</span><span class="p">)</span><span class="o">.</span><span class="n">value</span>
</span><span id="__span-0-74"><a id="__codelineno-0-74" name="__codelineno-0-74"></a> <span class="bp">self</span><span class="o">.</span><span class="n">u0</span> <span class="o">=</span> <span class="n">u0</span><span class="o">.</span><span class="n">item</span><span class="p">()</span>
</span><span id="__span-0-75"><a id="__codelineno-0-75" name="__codelineno-0-75"></a> <span class="bp">self</span><span class="o">.</span><span class="n">v0</span> <span class="o">=</span> <span class="n">v0</span><span class="o">.</span><span class="n">item</span><span class="p">()</span>
</span></code></pre></div></td></tr></table></div>
</details>
</div>
</div>
<div class="doc doc-object doc-function">
<h6 id="felice.neuron_models.Boomerang.init_state" class="doc doc-heading">
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <code class="highlight language-python"><span class="n">init_state</span><span class="p">(</span><span class="n">n_neurons</span><span class="p">:</span> <span class="nb">int</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s1">&#39;neurons 2&#39;</span><span class="p">]</span></code>
</h6>
<div class="doc doc-contents ">
<p>Initialize the neuron state variables.</p>
<p><span class="doc-section-title">Parameters:</span></p>
<table>
<thead>
<tr>
<th>Name</th>
<th>Type</th>
<th>Description</th>
<th>Default</th>
</tr>
</thead>
<tbody>
<tr class="doc-section-item">
<td>
<code>n_neurons</code>
</td>
<td>
<code><span title="int">int</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Number of neurons to initialize.</p>
</div>
</td>
<td>
<em>required</em>
</td>
</tr>
</tbody>
</table>
<p><span class="doc-section-title">Returns:</span></p>
<table>
<thead>
<tr>
<th>Type</th>
<th>Description</th>
</tr>
</thead>
<tbody>
<tr class="doc-section-item">
<td>
<code><span title="jaxtyping.Float">Float</span>[<span title="jaxtyping.Array">Array</span>, &#39;neurons 2&#39;]</code>
</td>
<td>
<div class="doc-md-description">
<p>Initial state array of shape (neurons, 3) containing [u, v],</p>
</div>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code><span title="jaxtyping.Float">Float</span>[<span title="jaxtyping.Array">Array</span>, &#39;neurons 2&#39;]</code>
</td>
<td>
<div class="doc-md-description">
<p>where u and v are the predator/prey membrane voltages.</p>
</div>
</td>
</tr>
</tbody>
</table>
<details class="mkdocstrings-source">
<summary>Source code in <code>felice/neuron_models/boomerang.py</code></summary>
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal"><a href="#__codelineno-0-77">77</a></span>
<span class="normal"><a href="#__codelineno-0-78">78</a></span>
<span class="normal"><a href="#__codelineno-0-79">79</a></span>
<span class="normal"><a href="#__codelineno-0-80">80</a></span>
<span class="normal"><a href="#__codelineno-0-81">81</a></span>
<span class="normal"><a href="#__codelineno-0-82">82</a></span>
<span class="normal"><a href="#__codelineno-0-83">83</a></span>
<span class="normal"><a href="#__codelineno-0-84">84</a></span>
<span class="normal"><a href="#__codelineno-0-85">85</a></span>
<span class="normal"><a href="#__codelineno-0-86">86</a></span>
<span class="normal"><a href="#__codelineno-0-87">87</a></span>
<span class="normal"><a href="#__codelineno-0-88">88</a></span>
<span class="normal"><a href="#__codelineno-0-89">89</a></span>
<span class="normal"><a href="#__codelineno-0-90">90</a></span>
<span class="normal"><a href="#__codelineno-0-91">91</a></span></pre></div></td><td class="code"><div><pre><span></span><code><span id="__span-0-77"><a id="__codelineno-0-77" name="__codelineno-0-77"></a><span class="k">def</span><span class="w"> </span><span class="nf">init_state</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">n_neurons</span><span class="p">:</span> <span class="nb">int</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;neurons 2&quot;</span><span class="p">]:</span>
</span><span id="__span-0-78"><a id="__codelineno-0-78" name="__codelineno-0-78"></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;Initialize the neuron state variables.</span>
</span><span id="__span-0-79"><a id="__codelineno-0-79" name="__codelineno-0-79"></a>
</span><span id="__span-0-80"><a id="__codelineno-0-80" name="__codelineno-0-80"></a><span class="sd"> Args:</span>
</span><span id="__span-0-81"><a id="__codelineno-0-81" name="__codelineno-0-81"></a><span class="sd"> n_neurons: Number of neurons to initialize.</span>
</span><span id="__span-0-82"><a id="__codelineno-0-82" name="__codelineno-0-82"></a>
</span><span id="__span-0-83"><a id="__codelineno-0-83" name="__codelineno-0-83"></a><span class="sd"> Returns:</span>
</span><span id="__span-0-84"><a id="__codelineno-0-84" name="__codelineno-0-84"></a><span class="sd"> Initial state array of shape (neurons, 3) containing [u, v],</span>
</span><span id="__span-0-85"><a id="__codelineno-0-85" name="__codelineno-0-85"></a><span class="sd"> where u and v are the predator/prey membrane voltages.</span>
</span><span id="__span-0-86"><a id="__codelineno-0-86" name="__codelineno-0-86"></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="__span-0-87"><a id="__codelineno-0-87" name="__codelineno-0-87"></a>
</span><span id="__span-0-88"><a id="__codelineno-0-88" name="__codelineno-0-88"></a> <span class="n">u</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">full</span><span class="p">((</span><span class="n">n_neurons</span><span class="p">,),</span> <span class="bp">self</span><span class="o">.</span><span class="n">u0</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dtype</span><span class="p">)</span>
</span><span id="__span-0-89"><a id="__codelineno-0-89" name="__codelineno-0-89"></a> <span class="n">v</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">full</span><span class="p">((</span><span class="n">n_neurons</span><span class="p">,),</span> <span class="bp">self</span><span class="o">.</span><span class="n">v0</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dtype</span><span class="p">)</span>
</span><span id="__span-0-90"><a id="__codelineno-0-90" name="__codelineno-0-90"></a> <span class="n">x</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">u</span><span class="p">,</span> <span class="n">v</span><span class="p">],</span> <span class="n">axis</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
</span><span id="__span-0-91"><a id="__codelineno-0-91" name="__codelineno-0-91"></a> <span class="k">return</span> <span class="n">x</span>
</span></code></pre></div></td></tr></table></div>
</details>
</div>
</div>
<div class="doc doc-object doc-function">
<h6 id="felice.neuron_models.Boomerang.dynamics" class="doc doc-heading">
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <code class="highlight language-python"><span class="n">dynamics</span><span class="p">(</span><span class="n">t</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span> <span class="n">y</span><span class="p">:</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s1">&#39;neurons 2&#39;</span><span class="p">],</span> <span class="n">args</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">])</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s1">&#39;neurons 2&#39;</span><span class="p">]</span></code>
</h6>
<div class="doc doc-contents ">
<p>Compute time derivatives of the neuron state variables.</p>
<p>This implements the WereRabbit dynamics</p>
<div class="language-text highlight"><pre><span></span><code>- du/dt: Predator dynamics
- dv/dt: WerePrey dynamics
</code></pre></div>
<p><span class="doc-section-title">Parameters:</span></p>
<table>
<thead>
<tr>
<th>Name</th>
<th>Type</th>
<th>Description</th>
<th>Default</th>
</tr>
</thead>
<tbody>
<tr class="doc-section-item">
<td>
<code>t</code>
</td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Current simulation time (unused but required by framework).</p>
</div>
</td>
<td>
<em>required</em>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>y</code>
</td>
<td>
<code><span title="jaxtyping.Float">Float</span>[<span title="jaxtyping.Array">Array</span>, &#39;neurons 2&#39;]</code>
</td>
<td>
<div class="doc-md-description">
<p>State array of shape (neurons, 2) containing [u, v].</p>
</div>
</td>
<td>
<em>required</em>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>args</code>
</td>
<td>
<code><span title="typing.Dict">Dict</span>[<span title="str">str</span>, <span title="typing.Any">Any</span>]</code>
</td>
<td>
<div class="doc-md-description">
<p>Additional arguments (unused but required by framework).</p>
</div>
</td>
<td>
<em>required</em>
</td>
</tr>
</tbody>
</table>
<p><span class="doc-section-title">Returns:</span></p>
<table>
<thead>
<tr>
<th>Type</th>
<th>Description</th>
</tr>
</thead>
<tbody>
<tr class="doc-section-item">
<td>
<code><span title="jaxtyping.Float">Float</span>[<span title="jaxtyping.Array">Array</span>, &#39;neurons 2&#39;]</code>
</td>
<td>
<div class="doc-md-description">
<p>Time derivatives of shape (neurons, 2) containing [du/dt, dv/dt].</p>
</div>
</td>
</tr>
</tbody>
</table>
<details class="mkdocstrings-source">
<summary>Source code in <code>felice/neuron_models/boomerang.py</code></summary>
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal"><a href="#__codelineno-0-108">108</a></span>
<span class="normal"><a href="#__codelineno-0-109">109</a></span>
<span class="normal"><a href="#__codelineno-0-110">110</a></span>
<span class="normal"><a href="#__codelineno-0-111">111</a></span>
<span class="normal"><a href="#__codelineno-0-112">112</a></span>
<span class="normal"><a href="#__codelineno-0-113">113</a></span>
<span class="normal"><a href="#__codelineno-0-114">114</a></span>
<span class="normal"><a href="#__codelineno-0-115">115</a></span>
<span class="normal"><a href="#__codelineno-0-116">116</a></span>
<span class="normal"><a href="#__codelineno-0-117">117</a></span>
<span class="normal"><a href="#__codelineno-0-118">118</a></span>
<span class="normal"><a href="#__codelineno-0-119">119</a></span>
<span class="normal"><a href="#__codelineno-0-120">120</a></span>
<span class="normal"><a href="#__codelineno-0-121">121</a></span>
<span class="normal"><a href="#__codelineno-0-122">122</a></span>
<span class="normal"><a href="#__codelineno-0-123">123</a></span>
<span class="normal"><a href="#__codelineno-0-124">124</a></span>
<span class="normal"><a href="#__codelineno-0-125">125</a></span>
<span class="normal"><a href="#__codelineno-0-126">126</a></span>
<span class="normal"><a href="#__codelineno-0-127">127</a></span>
<span class="normal"><a href="#__codelineno-0-128">128</a></span>
<span class="normal"><a href="#__codelineno-0-129">129</a></span>
<span class="normal"><a href="#__codelineno-0-130">130</a></span>
<span class="normal"><a href="#__codelineno-0-131">131</a></span>
<span class="normal"><a href="#__codelineno-0-132">132</a></span>
<span class="normal"><a href="#__codelineno-0-133">133</a></span>
<span class="normal"><a href="#__codelineno-0-134">134</a></span>
<span class="normal"><a href="#__codelineno-0-135">135</a></span></pre></div></td><td class="code"><div><pre><span></span><code><span id="__span-0-108"><a id="__codelineno-0-108" name="__codelineno-0-108"></a><span class="k">def</span><span class="w"> </span><span class="nf">dynamics</span><span class="p">(</span>
</span><span id="__span-0-109"><a id="__codelineno-0-109" name="__codelineno-0-109"></a> <span class="bp">self</span><span class="p">,</span>
</span><span id="__span-0-110"><a id="__codelineno-0-110" name="__codelineno-0-110"></a> <span class="n">t</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
</span><span id="__span-0-111"><a id="__codelineno-0-111" name="__codelineno-0-111"></a> <span class="n">y</span><span class="p">:</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;neurons 2&quot;</span><span class="p">],</span>
</span><span id="__span-0-112"><a id="__codelineno-0-112" name="__codelineno-0-112"></a> <span class="n">args</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">],</span>
</span><span id="__span-0-113"><a id="__codelineno-0-113" name="__codelineno-0-113"></a><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;neurons 2&quot;</span><span class="p">]:</span>
</span><span id="__span-0-114"><a id="__codelineno-0-114" name="__codelineno-0-114"></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;Compute time derivatives of the neuron state variables.</span>
</span><span id="__span-0-115"><a id="__codelineno-0-115" name="__codelineno-0-115"></a>
</span><span id="__span-0-116"><a id="__codelineno-0-116" name="__codelineno-0-116"></a><span class="sd"> This implements the WereRabbit dynamics</span>
</span><span id="__span-0-117"><a id="__codelineno-0-117" name="__codelineno-0-117"></a>
</span><span id="__span-0-118"><a id="__codelineno-0-118" name="__codelineno-0-118"></a><span class="sd"> - du/dt: Predator dynamics</span>
</span><span id="__span-0-119"><a id="__codelineno-0-119" name="__codelineno-0-119"></a><span class="sd"> - dv/dt: WerePrey dynamics</span>
</span><span id="__span-0-120"><a id="__codelineno-0-120" name="__codelineno-0-120"></a>
</span><span id="__span-0-121"><a id="__codelineno-0-121" name="__codelineno-0-121"></a><span class="sd"> Args:</span>
</span><span id="__span-0-122"><a id="__codelineno-0-122" name="__codelineno-0-122"></a><span class="sd"> t: Current simulation time (unused but required by framework).</span>
</span><span id="__span-0-123"><a id="__codelineno-0-123" name="__codelineno-0-123"></a><span class="sd"> y: State array of shape (neurons, 2) containing [u, v].</span>
</span><span id="__span-0-124"><a id="__codelineno-0-124" name="__codelineno-0-124"></a><span class="sd"> args: Additional arguments (unused but required by framework).</span>
</span><span id="__span-0-125"><a id="__codelineno-0-125" name="__codelineno-0-125"></a>
</span><span id="__span-0-126"><a id="__codelineno-0-126" name="__codelineno-0-126"></a><span class="sd"> Returns:</span>
</span><span id="__span-0-127"><a id="__codelineno-0-127" name="__codelineno-0-127"></a><span class="sd"> Time derivatives of shape (neurons, 2) containing [du/dt, dv/dt].</span>
</span><span id="__span-0-128"><a id="__codelineno-0-128" name="__codelineno-0-128"></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="__span-0-129"><a id="__codelineno-0-129" name="__codelineno-0-129"></a> <span class="n">u</span> <span class="o">=</span> <span class="n">y</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span>
</span><span id="__span-0-130"><a id="__codelineno-0-130" name="__codelineno-0-130"></a> <span class="n">v</span> <span class="o">=</span> <span class="n">y</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">]</span>
</span><span id="__span-0-131"><a id="__codelineno-0-131" name="__codelineno-0-131"></a>
</span><span id="__span-0-132"><a id="__codelineno-0-132" name="__codelineno-0-132"></a> <span class="n">du</span><span class="p">,</span> <span class="n">dv</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">vector_field</span><span class="p">(</span><span class="n">u</span><span class="p">,</span> <span class="n">v</span><span class="p">)</span>
</span><span id="__span-0-133"><a id="__codelineno-0-133" name="__codelineno-0-133"></a> <span class="n">dxdt</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">du</span><span class="p">,</span> <span class="n">dv</span><span class="p">],</span> <span class="n">axis</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
</span><span id="__span-0-134"><a id="__codelineno-0-134" name="__codelineno-0-134"></a>
</span><span id="__span-0-135"><a id="__codelineno-0-135" name="__codelineno-0-135"></a> <span class="k">return</span> <span class="n">dxdt</span>
</span></code></pre></div></td></tr></table></div>
</details>
</div>
</div>
<div class="doc doc-object doc-function">
<h6 id="felice.neuron_models.Boomerang.spike_condition" class="doc doc-heading">
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <code class="highlight language-python"><span class="n">spike_condition</span><span class="p">(</span><span class="n">t</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span> <span class="n">y</span><span class="p">:</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s1">&#39;neurons 2&#39;</span><span class="p">],</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">])</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s1">&#39; neurons&#39;</span><span class="p">]</span></code>
</h6>
<div class="doc doc-contents ">
<p>Compute spike condition for event detection.</p>
<p>A spike is triggered when the system reach to a fixpoint.</p>
<details class="info" open>
<summary>INFO</summary>
<p><code>has_spiked</code> is use to the system don't detect a continuos
spike when reach a fixpoint.</p>
</details>
<p><span class="doc-section-title">Parameters:</span></p>
<table>
<thead>
<tr>
<th>Name</th>
<th>Type</th>
<th>Description</th>
<th>Default</th>
</tr>
</thead>
<tbody>
<tr class="doc-section-item">
<td>
<code>t</code>
</td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Current simulation time (unused but required by the framework).</p>
</div>
</td>
<td>
<em>required</em>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>y</code>
</td>
<td>
<code><span title="jaxtyping.Float">Float</span>[<span title="jaxtyping.Array">Array</span>, &#39;neurons 2&#39;]</code>
</td>
<td>
<div class="doc-md-description">
<p>State array of shape (neurons, 3) containing [u, v, has_spiked].</p>
</div>
</td>
<td>
<em>required</em>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>**kwargs</code>
</td>
<td>
<code><span title="typing.Dict">Dict</span>[<span title="str">str</span>, <span title="typing.Any">Any</span>]</code>
</td>
<td>
<div class="doc-md-description">
<p>Additional keyword arguments (unused).</p>
</div>
</td>
<td>
<code>{}</code>
</td>
</tr>
</tbody>
</table>
<p><span class="doc-section-title">Returns:</span></p>
<table>
<thead>
<tr>
<th>Type</th>
<th>Description</th>
</tr>
</thead>
<tbody>
<tr class="doc-section-item">
<td>
<code><span title="jaxtyping.Float">Float</span>[<span title="jaxtyping.Array">Array</span>, &#39; neurons&#39;]</code>
</td>
<td>
<div class="doc-md-description">
<p>Spike condition array of shape (neurons,). Positive values indicate spike.</p>
</div>
</td>
</tr>
</tbody>
</table>
<details class="mkdocstrings-source">
<summary>Source code in <code>felice/neuron_models/boomerang.py</code></summary>
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal"><a href="#__codelineno-0-137">137</a></span>
<span class="normal"><a href="#__codelineno-0-138">138</a></span>
<span class="normal"><a href="#__codelineno-0-139">139</a></span>
<span class="normal"><a href="#__codelineno-0-140">140</a></span>
<span class="normal"><a href="#__codelineno-0-141">141</a></span>
<span class="normal"><a href="#__codelineno-0-142">142</a></span>
<span class="normal"><a href="#__codelineno-0-143">143</a></span>
<span class="normal"><a href="#__codelineno-0-144">144</a></span>
<span class="normal"><a href="#__codelineno-0-145">145</a></span>
<span class="normal"><a href="#__codelineno-0-146">146</a></span>
<span class="normal"><a href="#__codelineno-0-147">147</a></span>
<span class="normal"><a href="#__codelineno-0-148">148</a></span>
<span class="normal"><a href="#__codelineno-0-149">149</a></span>
<span class="normal"><a href="#__codelineno-0-150">150</a></span>
<span class="normal"><a href="#__codelineno-0-151">151</a></span>
<span class="normal"><a href="#__codelineno-0-152">152</a></span>
<span class="normal"><a href="#__codelineno-0-153">153</a></span>
<span class="normal"><a href="#__codelineno-0-154">154</a></span>
<span class="normal"><a href="#__codelineno-0-155">155</a></span>
<span class="normal"><a href="#__codelineno-0-156">156</a></span>
<span class="normal"><a href="#__codelineno-0-157">157</a></span>
<span class="normal"><a href="#__codelineno-0-158">158</a></span>
<span class="normal"><a href="#__codelineno-0-159">159</a></span>
<span class="normal"><a href="#__codelineno-0-160">160</a></span>
<span class="normal"><a href="#__codelineno-0-161">161</a></span>
<span class="normal"><a href="#__codelineno-0-162">162</a></span>
<span class="normal"><a href="#__codelineno-0-163">163</a></span>
<span class="normal"><a href="#__codelineno-0-164">164</a></span>
<span class="normal"><a href="#__codelineno-0-165">165</a></span>
<span class="normal"><a href="#__codelineno-0-166">166</a></span>
<span class="normal"><a href="#__codelineno-0-167">167</a></span>
<span class="normal"><a href="#__codelineno-0-168">168</a></span>
<span class="normal"><a href="#__codelineno-0-169">169</a></span>
<span class="normal"><a href="#__codelineno-0-170">170</a></span>
<span class="normal"><a href="#__codelineno-0-171">171</a></span></pre></div></td><td class="code"><div><pre><span></span><code><span id="__span-0-137"><a id="__codelineno-0-137" name="__codelineno-0-137"></a><span class="k">def</span><span class="w"> </span><span class="nf">spike_condition</span><span class="p">(</span>
</span><span id="__span-0-138"><a id="__codelineno-0-138" name="__codelineno-0-138"></a> <span class="bp">self</span><span class="p">,</span>
</span><span id="__span-0-139"><a id="__codelineno-0-139" name="__codelineno-0-139"></a> <span class="n">t</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
</span><span id="__span-0-140"><a id="__codelineno-0-140" name="__codelineno-0-140"></a> <span class="n">y</span><span class="p">:</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;neurons 2&quot;</span><span class="p">],</span>
</span><span id="__span-0-141"><a id="__codelineno-0-141" name="__codelineno-0-141"></a> <span class="o">**</span><span class="n">kwargs</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">],</span>
</span><span id="__span-0-142"><a id="__codelineno-0-142" name="__codelineno-0-142"></a><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot; neurons&quot;</span><span class="p">]:</span>
</span><span id="__span-0-143"><a id="__codelineno-0-143" name="__codelineno-0-143"></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;Compute spike condition for event detection.</span>
</span><span id="__span-0-144"><a id="__codelineno-0-144" name="__codelineno-0-144"></a>
</span><span id="__span-0-145"><a id="__codelineno-0-145" name="__codelineno-0-145"></a><span class="sd"> A spike is triggered when the system reach to a fixpoint.</span>
</span><span id="__span-0-146"><a id="__codelineno-0-146" name="__codelineno-0-146"></a>
</span><span id="__span-0-147"><a id="__codelineno-0-147" name="__codelineno-0-147"></a><span class="sd"> INFO:</span>
</span><span id="__span-0-148"><a id="__codelineno-0-148" name="__codelineno-0-148"></a><span class="sd"> `has_spiked` is use to the system don&#39;t detect a continuos</span>
</span><span id="__span-0-149"><a id="__codelineno-0-149" name="__codelineno-0-149"></a><span class="sd"> spike when reach a fixpoint.</span>
</span><span id="__span-0-150"><a id="__codelineno-0-150" name="__codelineno-0-150"></a>
</span><span id="__span-0-151"><a id="__codelineno-0-151" name="__codelineno-0-151"></a><span class="sd"> Args:</span>
</span><span id="__span-0-152"><a id="__codelineno-0-152" name="__codelineno-0-152"></a><span class="sd"> t: Current simulation time (unused but required by the framework).</span>
</span><span id="__span-0-153"><a id="__codelineno-0-153" name="__codelineno-0-153"></a><span class="sd"> y: State array of shape (neurons, 3) containing [u, v, has_spiked].</span>
</span><span id="__span-0-154"><a id="__codelineno-0-154" name="__codelineno-0-154"></a><span class="sd"> **kwargs: Additional keyword arguments (unused).</span>
</span><span id="__span-0-155"><a id="__codelineno-0-155" name="__codelineno-0-155"></a>
</span><span id="__span-0-156"><a id="__codelineno-0-156" name="__codelineno-0-156"></a><span class="sd"> Returns:</span>
</span><span id="__span-0-157"><a id="__codelineno-0-157" name="__codelineno-0-157"></a><span class="sd"> Spike condition array of shape (neurons,). Positive values indicate spike.</span>
</span><span id="__span-0-158"><a id="__codelineno-0-158" name="__codelineno-0-158"></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="__span-0-159"><a id="__codelineno-0-159" name="__codelineno-0-159"></a> <span class="n">_atol</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">atol</span>
</span><span id="__span-0-160"><a id="__codelineno-0-160" name="__codelineno-0-160"></a> <span class="n">_rtol</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">rtol</span>
</span><span id="__span-0-161"><a id="__codelineno-0-161" name="__codelineno-0-161"></a> <span class="n">_norm</span> <span class="o">=</span> <span class="n">optx</span><span class="o">.</span><span class="n">rms_norm</span>
</span><span id="__span-0-162"><a id="__codelineno-0-162" name="__codelineno-0-162"></a>
</span><span id="__span-0-163"><a id="__codelineno-0-163" name="__codelineno-0-163"></a> <span class="n">vf</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">dynamics</span><span class="p">(</span><span class="n">t</span><span class="p">,</span> <span class="n">y</span><span class="p">,</span> <span class="p">{})</span>
</span><span id="__span-0-164"><a id="__codelineno-0-164" name="__codelineno-0-164"></a>
</span><span id="__span-0-165"><a id="__codelineno-0-165" name="__codelineno-0-165"></a> <span class="nd">@jax</span><span class="o">.</span><span class="n">vmap</span>
</span><span id="__span-0-166"><a id="__codelineno-0-166" name="__codelineno-0-166"></a> <span class="k">def</span><span class="w"> </span><span class="nf">calculate_norm</span><span class="p">(</span><span class="n">vf</span><span class="p">,</span> <span class="n">y</span><span class="p">):</span>
</span><span id="__span-0-167"><a id="__codelineno-0-167" name="__codelineno-0-167"></a> <span class="k">return</span> <span class="n">_atol</span> <span class="o">+</span> <span class="n">_rtol</span> <span class="o">*</span> <span class="n">_norm</span><span class="p">(</span><span class="n">y</span><span class="p">)</span> <span class="o">-</span> <span class="n">_norm</span><span class="p">(</span><span class="n">vf</span><span class="p">)</span>
</span><span id="__span-0-168"><a id="__codelineno-0-168" name="__codelineno-0-168"></a>
</span><span id="__span-0-169"><a id="__codelineno-0-169" name="__codelineno-0-169"></a> <span class="n">base_cond</span> <span class="o">=</span> <span class="n">calculate_norm</span><span class="p">(</span><span class="n">vf</span><span class="p">,</span> <span class="n">y</span><span class="p">)</span><span class="o">.</span><span class="n">repeat</span><span class="p">(</span><span class="mi">2</span><span class="p">)</span>
</span><span id="__span-0-170"><a id="__codelineno-0-170" name="__codelineno-0-170"></a>
</span><span id="__span-0-171"><a id="__codelineno-0-171" name="__codelineno-0-171"></a> <span class="k">return</span> <span class="n">base_cond</span>
</span></code></pre></div></td></tr></table></div>
</details>
</div>
</div>
</div>
</div>
</div>
<div class="doc doc-object doc-class">
<h4 id="felice.neuron_models.FHNRS" class="doc doc-heading">
<code class="doc-symbol doc-symbol-heading doc-symbol-class"></code> <code>FHNRS</code>
</h4>
<div class="doc doc-contents ">
<p class="doc doc-class-bases">
Bases: <code><span title="equinox.Module">Module</span></code></p>
<p>FitzHugh-Nagumo neuron model</p>
<p>Model for FitzHugh-Nagumo neuron, with a hardware implementation proposed by
Ribar-Sepulchre. This implementation uses a dual-timescale dynamics with fast
and slow currents to produce oscillatory spiking behavior.</p>
<p>The dynamics are governed by:</p>
<div class="arithmatex">\[
\begin{align}
C\frac{dv}{dt} &amp;= I_{app} - I_{passive} - I_{fast} - I_{slow} \\
\frac{dv_{slow}}{dt} &amp;= \frac{v - v_{slow}}{\tau_{slow}} \\
\frac{dI_{app}}{dt} &amp;= -\frac{I_{app}}{\tau_{syn}}
\end{align}
\]</div>
<p>where the currents are:</p>
<ul>
<li><span class="arithmatex">\(I_{passive} = g_{max}(v - E_{rev})\)</span></li>
<li><span class="arithmatex">\(I_{fast} = a_{fast} \tanh(v - v_{off,fast})\)</span></li>
<li><span class="arithmatex">\(I_{slow} = a_{slow} \tanh(v_{slow} - v_{off,slow})\)</span></li>
</ul>
<details class="references" open>
<summary>References</summary>
<ul>
<li>Ribar, L., &amp; Sepulchre, R. (2019). Neuromodulation of neuromorphic circuits. IEEE Transactions on Circuits and Systems I: Regular Papers, 66(8), 3028-3040.</li>
</ul>
</details>
<p><span class="doc-section-title">Attributes:</span></p>
<table>
<thead>
<tr>
<th>Name</th>
<th>Type</th>
<th>Description</th>
</tr>
</thead>
<tbody>
<tr class="doc-section-item">
<td><code><span title="felice.neuron_models.FHNRS.reset_grad_preserve">reset_grad_preserve</span></code></td>
<td>
</td>
<td>
<div class="doc-md-description">
<p>Preserve the gradient when the neuron spikes by doing a soft reset.</p>
</div>
</td>
</tr>
<tr class="doc-section-item">
<td><code><span title="felice.neuron_models.FHNRS.gmax_pasive">gmax_pasive</span></code></td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Maximal conductance of the passive current.</p>
</div>
</td>
</tr>
<tr class="doc-section-item">
<td><code><span title="felice.neuron_models.FHNRS.Erev_pasive">Erev_pasive</span></code></td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Reversal potential for the passive current.</p>
</div>
</td>
</tr>
<tr class="doc-section-item">
<td><code><span title="felice.neuron_models.FHNRS.a_fast">a_fast</span></code></td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Amplitude parameter for the fast current dynamics.</p>
</div>
</td>
</tr>
<tr class="doc-section-item">
<td><code><span title="felice.neuron_models.FHNRS.voff_fast">voff_fast</span></code></td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Voltage offset for the fast current activation.</p>
</div>
</td>
</tr>
<tr class="doc-section-item">
<td><code><span title="felice.neuron_models.FHNRS.tau_fast">tau_fast</span></code></td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Time constant for the fast current (typically zero for instantaneous).</p>
</div>
</td>
</tr>
<tr class="doc-section-item">
<td><code><span title="felice.neuron_models.FHNRS.a_slow">a_slow</span></code></td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Amplitude parameter for the slow current dynamics.</p>
</div>
</td>
</tr>
<tr class="doc-section-item">
<td><code><span title="felice.neuron_models.FHNRS.voff_slow">voff_slow</span></code></td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Voltage offset for the slow current activation.</p>
</div>
</td>
</tr>
<tr class="doc-section-item">
<td><code><span title="felice.neuron_models.FHNRS.tau_slow">tau_slow</span></code></td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Time constant for the slow recovery variable.</p>
</div>
</td>
</tr>
<tr class="doc-section-item">
<td><code><span title="felice.neuron_models.FHNRS.vthr">vthr</span></code></td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Voltage threshold for spike generation.</p>
</div>
</td>
</tr>
<tr class="doc-section-item">
<td><code><span title="felice.neuron_models.FHNRS.C">C</span></code></td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Membrane capacitance.</p>
</div>
</td>
</tr>
<tr class="doc-section-item">
<td><code><span title="felice.neuron_models.FHNRS.tsyn">tsyn</span></code></td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Synaptic time constant for input current decay.</p>
</div>
</td>
</tr>
<tr class="doc-section-item">
<td><code><span title="felice.neuron_models.FHNRS.weights">weights</span></code></td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Synaptic weight matrix of shape (in_plus_neurons, neurons).</p>
</div>
</td>
</tr>
</tbody>
</table>
<details class="mkdocstrings-source">
<summary>Source code in <code>felice/neuron_models/fhn.py</code></summary>
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal"><a href="#__codelineno-0-8"> 8</a></span>
<span class="normal"><a href="#__codelineno-0-9"> 9</a></span>
<span class="normal"><a href="#__codelineno-0-10"> 10</a></span>
<span class="normal"><a href="#__codelineno-0-11"> 11</a></span>
<span class="normal"><a href="#__codelineno-0-12"> 12</a></span>
<span class="normal"><a href="#__codelineno-0-13"> 13</a></span>
<span class="normal"><a href="#__codelineno-0-14"> 14</a></span>
<span class="normal"><a href="#__codelineno-0-15"> 15</a></span>
<span class="normal"><a href="#__codelineno-0-16"> 16</a></span>
<span class="normal"><a href="#__codelineno-0-17"> 17</a></span>
<span class="normal"><a href="#__codelineno-0-18"> 18</a></span>
<span class="normal"><a href="#__codelineno-0-19"> 19</a></span>
<span class="normal"><a href="#__codelineno-0-20"> 20</a></span>
<span class="normal"><a href="#__codelineno-0-21"> 21</a></span>
<span class="normal"><a href="#__codelineno-0-22"> 22</a></span>
<span class="normal"><a href="#__codelineno-0-23"> 23</a></span>
<span class="normal"><a href="#__codelineno-0-24"> 24</a></span>
<span class="normal"><a href="#__codelineno-0-25"> 25</a></span>
<span class="normal"><a href="#__codelineno-0-26"> 26</a></span>
<span class="normal"><a href="#__codelineno-0-27"> 27</a></span>
<span class="normal"><a href="#__codelineno-0-28"> 28</a></span>
<span class="normal"><a href="#__codelineno-0-29"> 29</a></span>
<span class="normal"><a href="#__codelineno-0-30"> 30</a></span>
<span class="normal"><a href="#__codelineno-0-31"> 31</a></span>
<span class="normal"><a href="#__codelineno-0-32"> 32</a></span>
<span class="normal"><a href="#__codelineno-0-33"> 33</a></span>
<span class="normal"><a href="#__codelineno-0-34"> 34</a></span>
<span class="normal"><a href="#__codelineno-0-35"> 35</a></span>
<span class="normal"><a href="#__codelineno-0-36"> 36</a></span>
<span class="normal"><a href="#__codelineno-0-37"> 37</a></span>
<span class="normal"><a href="#__codelineno-0-38"> 38</a></span>
<span class="normal"><a href="#__codelineno-0-39"> 39</a></span>
<span class="normal"><a href="#__codelineno-0-40"> 40</a></span>
<span class="normal"><a href="#__codelineno-0-41"> 41</a></span>
<span class="normal"><a href="#__codelineno-0-42"> 42</a></span>
<span class="normal"><a href="#__codelineno-0-43"> 43</a></span>
<span class="normal"><a href="#__codelineno-0-44"> 44</a></span>
<span class="normal"><a href="#__codelineno-0-45"> 45</a></span>
<span class="normal"><a href="#__codelineno-0-46"> 46</a></span>
<span class="normal"><a href="#__codelineno-0-47"> 47</a></span>
<span class="normal"><a href="#__codelineno-0-48"> 48</a></span>
<span class="normal"><a href="#__codelineno-0-49"> 49</a></span>
<span class="normal"><a href="#__codelineno-0-50"> 50</a></span>
<span class="normal"><a href="#__codelineno-0-51"> 51</a></span>
<span class="normal"><a href="#__codelineno-0-52"> 52</a></span>
<span class="normal"><a href="#__codelineno-0-53"> 53</a></span>
<span class="normal"><a href="#__codelineno-0-54"> 54</a></span>
<span class="normal"><a href="#__codelineno-0-55"> 55</a></span>
<span class="normal"><a href="#__codelineno-0-56"> 56</a></span>
<span class="normal"><a href="#__codelineno-0-57"> 57</a></span>
<span class="normal"><a href="#__codelineno-0-58"> 58</a></span>
<span class="normal"><a href="#__codelineno-0-59"> 59</a></span>
<span class="normal"><a href="#__codelineno-0-60"> 60</a></span>
<span class="normal"><a href="#__codelineno-0-61"> 61</a></span>
<span class="normal"><a href="#__codelineno-0-62"> 62</a></span>
<span class="normal"><a href="#__codelineno-0-63"> 63</a></span>
<span class="normal"><a href="#__codelineno-0-64"> 64</a></span>
<span class="normal"><a href="#__codelineno-0-65"> 65</a></span>
<span class="normal"><a href="#__codelineno-0-66"> 66</a></span>
<span class="normal"><a href="#__codelineno-0-67"> 67</a></span>
<span class="normal"><a href="#__codelineno-0-68"> 68</a></span>
<span class="normal"><a href="#__codelineno-0-69"> 69</a></span>
<span class="normal"><a href="#__codelineno-0-70"> 70</a></span>
<span class="normal"><a href="#__codelineno-0-71"> 71</a></span>
<span class="normal"><a href="#__codelineno-0-72"> 72</a></span>
<span class="normal"><a href="#__codelineno-0-73"> 73</a></span>
<span class="normal"><a href="#__codelineno-0-74"> 74</a></span>
<span class="normal"><a href="#__codelineno-0-75"> 75</a></span>
<span class="normal"><a href="#__codelineno-0-76"> 76</a></span>
<span class="normal"><a href="#__codelineno-0-77"> 77</a></span>
<span class="normal"><a href="#__codelineno-0-78"> 78</a></span>
<span class="normal"><a href="#__codelineno-0-79"> 79</a></span>
<span class="normal"><a href="#__codelineno-0-80"> 80</a></span>
<span class="normal"><a href="#__codelineno-0-81"> 81</a></span>
<span class="normal"><a href="#__codelineno-0-82"> 82</a></span>
<span class="normal"><a href="#__codelineno-0-83"> 83</a></span>
<span class="normal"><a href="#__codelineno-0-84"> 84</a></span>
<span class="normal"><a href="#__codelineno-0-85"> 85</a></span>
<span class="normal"><a href="#__codelineno-0-86"> 86</a></span>
<span class="normal"><a href="#__codelineno-0-87"> 87</a></span>
<span class="normal"><a href="#__codelineno-0-88"> 88</a></span>
<span class="normal"><a href="#__codelineno-0-89"> 89</a></span>
<span class="normal"><a href="#__codelineno-0-90"> 90</a></span>
<span class="normal"><a href="#__codelineno-0-91"> 91</a></span>
<span class="normal"><a href="#__codelineno-0-92"> 92</a></span>
<span class="normal"><a href="#__codelineno-0-93"> 93</a></span>
<span class="normal"><a href="#__codelineno-0-94"> 94</a></span>
<span class="normal"><a href="#__codelineno-0-95"> 95</a></span>
<span class="normal"><a href="#__codelineno-0-96"> 96</a></span>
<span class="normal"><a href="#__codelineno-0-97"> 97</a></span>
<span class="normal"><a href="#__codelineno-0-98"> 98</a></span>
<span class="normal"><a href="#__codelineno-0-99"> 99</a></span>
<span class="normal"><a href="#__codelineno-0-100">100</a></span>
<span class="normal"><a href="#__codelineno-0-101">101</a></span>
<span class="normal"><a href="#__codelineno-0-102">102</a></span>
<span class="normal"><a href="#__codelineno-0-103">103</a></span>
<span class="normal"><a href="#__codelineno-0-104">104</a></span>
<span class="normal"><a href="#__codelineno-0-105">105</a></span>
<span class="normal"><a href="#__codelineno-0-106">106</a></span>
<span class="normal"><a href="#__codelineno-0-107">107</a></span>
<span class="normal"><a href="#__codelineno-0-108">108</a></span>
<span class="normal"><a href="#__codelineno-0-109">109</a></span>
<span class="normal"><a href="#__codelineno-0-110">110</a></span>
<span class="normal"><a href="#__codelineno-0-111">111</a></span>
<span class="normal"><a href="#__codelineno-0-112">112</a></span>
<span class="normal"><a href="#__codelineno-0-113">113</a></span>
<span class="normal"><a href="#__codelineno-0-114">114</a></span>
<span class="normal"><a href="#__codelineno-0-115">115</a></span>
<span class="normal"><a href="#__codelineno-0-116">116</a></span>
<span class="normal"><a href="#__codelineno-0-117">117</a></span>
<span class="normal"><a href="#__codelineno-0-118">118</a></span>
<span class="normal"><a href="#__codelineno-0-119">119</a></span>
<span class="normal"><a href="#__codelineno-0-120">120</a></span>
<span class="normal"><a href="#__codelineno-0-121">121</a></span>
<span class="normal"><a href="#__codelineno-0-122">122</a></span>
<span class="normal"><a href="#__codelineno-0-123">123</a></span>
<span class="normal"><a href="#__codelineno-0-124">124</a></span>
<span class="normal"><a href="#__codelineno-0-125">125</a></span>
<span class="normal"><a href="#__codelineno-0-126">126</a></span>
<span class="normal"><a href="#__codelineno-0-127">127</a></span>
<span class="normal"><a href="#__codelineno-0-128">128</a></span>
<span class="normal"><a href="#__codelineno-0-129">129</a></span>
<span class="normal"><a href="#__codelineno-0-130">130</a></span>
<span class="normal"><a href="#__codelineno-0-131">131</a></span>
<span class="normal"><a href="#__codelineno-0-132">132</a></span>
<span class="normal"><a href="#__codelineno-0-133">133</a></span>
<span class="normal"><a href="#__codelineno-0-134">134</a></span>
<span class="normal"><a href="#__codelineno-0-135">135</a></span>
<span class="normal"><a href="#__codelineno-0-136">136</a></span>
<span class="normal"><a href="#__codelineno-0-137">137</a></span>
<span class="normal"><a href="#__codelineno-0-138">138</a></span>
<span class="normal"><a href="#__codelineno-0-139">139</a></span>
<span class="normal"><a href="#__codelineno-0-140">140</a></span>
<span class="normal"><a href="#__codelineno-0-141">141</a></span>
<span class="normal"><a href="#__codelineno-0-142">142</a></span>
<span class="normal"><a href="#__codelineno-0-143">143</a></span>
<span class="normal"><a href="#__codelineno-0-144">144</a></span>
<span class="normal"><a href="#__codelineno-0-145">145</a></span>
<span class="normal"><a href="#__codelineno-0-146">146</a></span>
<span class="normal"><a href="#__codelineno-0-147">147</a></span>
<span class="normal"><a href="#__codelineno-0-148">148</a></span>
<span class="normal"><a href="#__codelineno-0-149">149</a></span>
<span class="normal"><a href="#__codelineno-0-150">150</a></span>
<span class="normal"><a href="#__codelineno-0-151">151</a></span>
<span class="normal"><a href="#__codelineno-0-152">152</a></span>
<span class="normal"><a href="#__codelineno-0-153">153</a></span>
<span class="normal"><a href="#__codelineno-0-154">154</a></span>
<span class="normal"><a href="#__codelineno-0-155">155</a></span>
<span class="normal"><a href="#__codelineno-0-156">156</a></span>
<span class="normal"><a href="#__codelineno-0-157">157</a></span>
<span class="normal"><a href="#__codelineno-0-158">158</a></span>
<span class="normal"><a href="#__codelineno-0-159">159</a></span>
<span class="normal"><a href="#__codelineno-0-160">160</a></span>
<span class="normal"><a href="#__codelineno-0-161">161</a></span>
<span class="normal"><a href="#__codelineno-0-162">162</a></span>
<span class="normal"><a href="#__codelineno-0-163">163</a></span>
<span class="normal"><a href="#__codelineno-0-164">164</a></span>
<span class="normal"><a href="#__codelineno-0-165">165</a></span>
<span class="normal"><a href="#__codelineno-0-166">166</a></span>
<span class="normal"><a href="#__codelineno-0-167">167</a></span>
<span class="normal"><a href="#__codelineno-0-168">168</a></span>
<span class="normal"><a href="#__codelineno-0-169">169</a></span>
<span class="normal"><a href="#__codelineno-0-170">170</a></span>
<span class="normal"><a href="#__codelineno-0-171">171</a></span>
<span class="normal"><a href="#__codelineno-0-172">172</a></span>
<span class="normal"><a href="#__codelineno-0-173">173</a></span>
<span class="normal"><a href="#__codelineno-0-174">174</a></span>
<span class="normal"><a href="#__codelineno-0-175">175</a></span>
<span class="normal"><a href="#__codelineno-0-176">176</a></span>
<span class="normal"><a href="#__codelineno-0-177">177</a></span>
<span class="normal"><a href="#__codelineno-0-178">178</a></span>
<span class="normal"><a href="#__codelineno-0-179">179</a></span>
<span class="normal"><a href="#__codelineno-0-180">180</a></span>
<span class="normal"><a href="#__codelineno-0-181">181</a></span>
<span class="normal"><a href="#__codelineno-0-182">182</a></span>
<span class="normal"><a href="#__codelineno-0-183">183</a></span>
<span class="normal"><a href="#__codelineno-0-184">184</a></span>
<span class="normal"><a href="#__codelineno-0-185">185</a></span>
<span class="normal"><a href="#__codelineno-0-186">186</a></span>
<span class="normal"><a href="#__codelineno-0-187">187</a></span>
<span class="normal"><a href="#__codelineno-0-188">188</a></span>
<span class="normal"><a href="#__codelineno-0-189">189</a></span>
<span class="normal"><a href="#__codelineno-0-190">190</a></span>
<span class="normal"><a href="#__codelineno-0-191">191</a></span>
<span class="normal"><a href="#__codelineno-0-192">192</a></span>
<span class="normal"><a href="#__codelineno-0-193">193</a></span>
<span class="normal"><a href="#__codelineno-0-194">194</a></span>
<span class="normal"><a href="#__codelineno-0-195">195</a></span>
<span class="normal"><a href="#__codelineno-0-196">196</a></span>
<span class="normal"><a href="#__codelineno-0-197">197</a></span>
<span class="normal"><a href="#__codelineno-0-198">198</a></span>
<span class="normal"><a href="#__codelineno-0-199">199</a></span>
<span class="normal"><a href="#__codelineno-0-200">200</a></span>
<span class="normal"><a href="#__codelineno-0-201">201</a></span>
<span class="normal"><a href="#__codelineno-0-202">202</a></span>
<span class="normal"><a href="#__codelineno-0-203">203</a></span>
<span class="normal"><a href="#__codelineno-0-204">204</a></span>
<span class="normal"><a href="#__codelineno-0-205">205</a></span>
<span class="normal"><a href="#__codelineno-0-206">206</a></span>
<span class="normal"><a href="#__codelineno-0-207">207</a></span>
<span class="normal"><a href="#__codelineno-0-208">208</a></span>
<span class="normal"><a href="#__codelineno-0-209">209</a></span>
<span class="normal"><a href="#__codelineno-0-210">210</a></span>
<span class="normal"><a href="#__codelineno-0-211">211</a></span>
<span class="normal"><a href="#__codelineno-0-212">212</a></span>
<span class="normal"><a href="#__codelineno-0-213">213</a></span>
<span class="normal"><a href="#__codelineno-0-214">214</a></span>
<span class="normal"><a href="#__codelineno-0-215">215</a></span>
<span class="normal"><a href="#__codelineno-0-216">216</a></span>
<span class="normal"><a href="#__codelineno-0-217">217</a></span>
<span class="normal"><a href="#__codelineno-0-218">218</a></span>
<span class="normal"><a href="#__codelineno-0-219">219</a></span>
<span class="normal"><a href="#__codelineno-0-220">220</a></span>
<span class="normal"><a href="#__codelineno-0-221">221</a></span>
<span class="normal"><a href="#__codelineno-0-222">222</a></span>
<span class="normal"><a href="#__codelineno-0-223">223</a></span>
<span class="normal"><a href="#__codelineno-0-224">224</a></span>
<span class="normal"><a href="#__codelineno-0-225">225</a></span>
<span class="normal"><a href="#__codelineno-0-226">226</a></span>
<span class="normal"><a href="#__codelineno-0-227">227</a></span>
<span class="normal"><a href="#__codelineno-0-228">228</a></span>
<span class="normal"><a href="#__codelineno-0-229">229</a></span>
<span class="normal"><a href="#__codelineno-0-230">230</a></span>
<span class="normal"><a href="#__codelineno-0-231">231</a></span>
<span class="normal"><a href="#__codelineno-0-232">232</a></span>
<span class="normal"><a href="#__codelineno-0-233">233</a></span>
<span class="normal"><a href="#__codelineno-0-234">234</a></span>
<span class="normal"><a href="#__codelineno-0-235">235</a></span></pre></div></td><td class="code"><div><pre><span></span><code><span id="__span-0-8"><a id="__codelineno-0-8" name="__codelineno-0-8"></a><span class="k">class</span><span class="w"> </span><span class="nc">FHNRS</span><span class="p">(</span><span class="n">eqx</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
</span><span id="__span-0-9"><a id="__codelineno-0-9" name="__codelineno-0-9"></a><span class="w"> </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;FitzHugh-Nagumo neuron model</span>
</span><span id="__span-0-10"><a id="__codelineno-0-10" name="__codelineno-0-10"></a>
</span><span id="__span-0-11"><a id="__codelineno-0-11" name="__codelineno-0-11"></a><span class="sd"> Model for FitzHugh-Nagumo neuron, with a hardware implementation proposed by</span>
</span><span id="__span-0-12"><a id="__codelineno-0-12" name="__codelineno-0-12"></a><span class="sd"> Ribar-Sepulchre. This implementation uses a dual-timescale dynamics with fast</span>
</span><span id="__span-0-13"><a id="__codelineno-0-13" name="__codelineno-0-13"></a><span class="sd"> and slow currents to produce oscillatory spiking behavior.</span>
</span><span id="__span-0-14"><a id="__codelineno-0-14" name="__codelineno-0-14"></a>
</span><span id="__span-0-15"><a id="__codelineno-0-15" name="__codelineno-0-15"></a><span class="sd"> The dynamics are governed by:</span>
</span><span id="__span-0-16"><a id="__codelineno-0-16" name="__codelineno-0-16"></a>
</span><span id="__span-0-17"><a id="__codelineno-0-17" name="__codelineno-0-17"></a><span class="sd"> $$</span>
</span><span id="__span-0-18"><a id="__codelineno-0-18" name="__codelineno-0-18"></a><span class="sd"> \begin{align}</span>
</span><span id="__span-0-19"><a id="__codelineno-0-19" name="__codelineno-0-19"></a><span class="sd"> C\frac{dv}{dt} &amp;= I_{app} - I_{passive} - I_{fast} - I_{slow} \\</span>
</span><span id="__span-0-20"><a id="__codelineno-0-20" name="__codelineno-0-20"></a><span class="sd"> \frac{dv_{slow}}{dt} &amp;= \frac{v - v_{slow}}{\tau_{slow}} \\</span>
</span><span id="__span-0-21"><a id="__codelineno-0-21" name="__codelineno-0-21"></a><span class="sd"> \frac{dI_{app}}{dt} &amp;= -\frac{I_{app}}{\tau_{syn}}</span>
</span><span id="__span-0-22"><a id="__codelineno-0-22" name="__codelineno-0-22"></a><span class="sd"> \end{align}</span>
</span><span id="__span-0-23"><a id="__codelineno-0-23" name="__codelineno-0-23"></a><span class="sd"> $$</span>
</span><span id="__span-0-24"><a id="__codelineno-0-24" name="__codelineno-0-24"></a>
</span><span id="__span-0-25"><a id="__codelineno-0-25" name="__codelineno-0-25"></a><span class="sd"> where the currents are:</span>
</span><span id="__span-0-26"><a id="__codelineno-0-26" name="__codelineno-0-26"></a>
</span><span id="__span-0-27"><a id="__codelineno-0-27" name="__codelineno-0-27"></a><span class="sd"> - $I_{passive} = g_{max}(v - E_{rev})$</span>
</span><span id="__span-0-28"><a id="__codelineno-0-28" name="__codelineno-0-28"></a><span class="sd"> - $I_{fast} = a_{fast} \tanh(v - v_{off,fast})$</span>
</span><span id="__span-0-29"><a id="__codelineno-0-29" name="__codelineno-0-29"></a><span class="sd"> - $I_{slow} = a_{slow} \tanh(v_{slow} - v_{off,slow})$</span>
</span><span id="__span-0-30"><a id="__codelineno-0-30" name="__codelineno-0-30"></a>
</span><span id="__span-0-31"><a id="__codelineno-0-31" name="__codelineno-0-31"></a><span class="sd"> References:</span>
</span><span id="__span-0-32"><a id="__codelineno-0-32" name="__codelineno-0-32"></a><span class="sd"> - Ribar, L., &amp; Sepulchre, R. (2019). Neuromodulation of neuromorphic circuits. IEEE Transactions on Circuits and Systems I: Regular Papers, 66(8), 3028-3040.</span>
</span><span id="__span-0-33"><a id="__codelineno-0-33" name="__codelineno-0-33"></a>
</span><span id="__span-0-34"><a id="__codelineno-0-34" name="__codelineno-0-34"></a><span class="sd"> Attributes:</span>
</span><span id="__span-0-35"><a id="__codelineno-0-35" name="__codelineno-0-35"></a><span class="sd"> reset_grad_preserve: Preserve the gradient when the neuron spikes by doing a soft reset.</span>
</span><span id="__span-0-36"><a id="__codelineno-0-36" name="__codelineno-0-36"></a><span class="sd"> gmax_pasive: Maximal conductance of the passive current.</span>
</span><span id="__span-0-37"><a id="__codelineno-0-37" name="__codelineno-0-37"></a><span class="sd"> Erev_pasive: Reversal potential for the passive current.</span>
</span><span id="__span-0-38"><a id="__codelineno-0-38" name="__codelineno-0-38"></a><span class="sd"> a_fast: Amplitude parameter for the fast current dynamics.</span>
</span><span id="__span-0-39"><a id="__codelineno-0-39" name="__codelineno-0-39"></a><span class="sd"> voff_fast: Voltage offset for the fast current activation.</span>
</span><span id="__span-0-40"><a id="__codelineno-0-40" name="__codelineno-0-40"></a><span class="sd"> tau_fast: Time constant for the fast current (typically zero for instantaneous).</span>
</span><span id="__span-0-41"><a id="__codelineno-0-41" name="__codelineno-0-41"></a><span class="sd"> a_slow: Amplitude parameter for the slow current dynamics.</span>
</span><span id="__span-0-42"><a id="__codelineno-0-42" name="__codelineno-0-42"></a><span class="sd"> voff_slow: Voltage offset for the slow current activation.</span>
</span><span id="__span-0-43"><a id="__codelineno-0-43" name="__codelineno-0-43"></a><span class="sd"> tau_slow: Time constant for the slow recovery variable.</span>
</span><span id="__span-0-44"><a id="__codelineno-0-44" name="__codelineno-0-44"></a><span class="sd"> vthr: Voltage threshold for spike generation.</span>
</span><span id="__span-0-45"><a id="__codelineno-0-45" name="__codelineno-0-45"></a><span class="sd"> C: Membrane capacitance.</span>
</span><span id="__span-0-46"><a id="__codelineno-0-46" name="__codelineno-0-46"></a><span class="sd"> tsyn: Synaptic time constant for input current decay.</span>
</span><span id="__span-0-47"><a id="__codelineno-0-47" name="__codelineno-0-47"></a><span class="sd"> weights: Synaptic weight matrix of shape (in_plus_neurons, neurons).</span>
</span><span id="__span-0-48"><a id="__codelineno-0-48" name="__codelineno-0-48"></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="__span-0-49"><a id="__codelineno-0-49" name="__codelineno-0-49"></a>
</span><span id="__span-0-50"><a id="__codelineno-0-50" name="__codelineno-0-50"></a> <span class="c1"># Pasive parameters</span>
</span><span id="__span-0-51"><a id="__codelineno-0-51" name="__codelineno-0-51"></a> <span class="n">gmax_pasive</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="n">eqx</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="n">static</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
</span><span id="__span-0-52"><a id="__codelineno-0-52" name="__codelineno-0-52"></a> <span class="n">Erev_pasive</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="n">eqx</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="n">static</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
</span><span id="__span-0-53"><a id="__codelineno-0-53" name="__codelineno-0-53"></a>
</span><span id="__span-0-54"><a id="__codelineno-0-54" name="__codelineno-0-54"></a> <span class="c1"># Fast current</span>
</span><span id="__span-0-55"><a id="__codelineno-0-55" name="__codelineno-0-55"></a> <span class="n">a_fast</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="n">eqx</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="n">static</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
</span><span id="__span-0-56"><a id="__codelineno-0-56" name="__codelineno-0-56"></a> <span class="n">voff_fast</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="n">eqx</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="n">static</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
</span><span id="__span-0-57"><a id="__codelineno-0-57" name="__codelineno-0-57"></a> <span class="n">tau_fast</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="n">eqx</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="n">static</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
</span><span id="__span-0-58"><a id="__codelineno-0-58" name="__codelineno-0-58"></a>
</span><span id="__span-0-59"><a id="__codelineno-0-59" name="__codelineno-0-59"></a> <span class="c1"># Slow current</span>
</span><span id="__span-0-60"><a id="__codelineno-0-60" name="__codelineno-0-60"></a> <span class="n">a_slow</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="n">eqx</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="n">static</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
</span><span id="__span-0-61"><a id="__codelineno-0-61" name="__codelineno-0-61"></a> <span class="n">voff_slow</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="n">eqx</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="n">static</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
</span><span id="__span-0-62"><a id="__codelineno-0-62" name="__codelineno-0-62"></a> <span class="n">tau_slow</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="n">eqx</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="n">static</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
</span><span id="__span-0-63"><a id="__codelineno-0-63" name="__codelineno-0-63"></a>
</span><span id="__span-0-64"><a id="__codelineno-0-64" name="__codelineno-0-64"></a> <span class="c1"># Neuron threshold</span>
</span><span id="__span-0-65"><a id="__codelineno-0-65" name="__codelineno-0-65"></a> <span class="n">vthr</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="n">eqx</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="n">static</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
</span><span id="__span-0-66"><a id="__codelineno-0-66" name="__codelineno-0-66"></a> <span class="n">C</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="n">eqx</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="n">static</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">default</span><span class="o">=</span><span class="mf">1.0</span><span class="p">)</span>
</span><span id="__span-0-67"><a id="__codelineno-0-67" name="__codelineno-0-67"></a>
</span><span id="__span-0-68"><a id="__codelineno-0-68" name="__codelineno-0-68"></a> <span class="c1"># Input synaptic time constant</span>
</span><span id="__span-0-69"><a id="__codelineno-0-69" name="__codelineno-0-69"></a> <span class="n">tsyn</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="n">eqx</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="n">static</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
</span><span id="__span-0-70"><a id="__codelineno-0-70" name="__codelineno-0-70"></a>
</span><span id="__span-0-71"><a id="__codelineno-0-71" name="__codelineno-0-71"></a> <span class="n">dtype</span><span class="p">:</span> <span class="n">DTypeLike</span> <span class="o">=</span> <span class="n">eqx</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="n">static</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
</span><span id="__span-0-72"><a id="__codelineno-0-72" name="__codelineno-0-72"></a>
</span><span id="__span-0-73"><a id="__codelineno-0-73" name="__codelineno-0-73"></a> <span class="k">def</span><span class="w"> </span><span class="fm">__init__</span><span class="p">(</span>
</span><span id="__span-0-74"><a id="__codelineno-0-74" name="__codelineno-0-74"></a> <span class="bp">self</span><span class="p">,</span>
</span><span id="__span-0-75"><a id="__codelineno-0-75" name="__codelineno-0-75"></a> <span class="o">*</span><span class="p">,</span>
</span><span id="__span-0-76"><a id="__codelineno-0-76" name="__codelineno-0-76"></a> <span class="n">tsyn</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="mf">1.0</span><span class="p">,</span>
</span><span id="__span-0-77"><a id="__codelineno-0-77" name="__codelineno-0-77"></a> <span class="n">C</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="mf">1.0</span><span class="p">,</span>
</span><span id="__span-0-78"><a id="__codelineno-0-78" name="__codelineno-0-78"></a> <span class="n">gmax_pasive</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="mf">1.0</span><span class="p">,</span>
</span><span id="__span-0-79"><a id="__codelineno-0-79" name="__codelineno-0-79"></a> <span class="n">Erev_pasive</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="mf">0.0</span><span class="p">,</span>
</span><span id="__span-0-80"><a id="__codelineno-0-80" name="__codelineno-0-80"></a> <span class="n">a_fast</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="o">-</span><span class="mf">2.0</span><span class="p">,</span>
</span><span id="__span-0-81"><a id="__codelineno-0-81" name="__codelineno-0-81"></a> <span class="n">voff_fast</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="mf">0.0</span><span class="p">,</span>
</span><span id="__span-0-82"><a id="__codelineno-0-82" name="__codelineno-0-82"></a> <span class="n">tau_fast</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="mf">0.0</span><span class="p">,</span>
</span><span id="__span-0-83"><a id="__codelineno-0-83" name="__codelineno-0-83"></a> <span class="n">a_slow</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="mf">2.0</span><span class="p">,</span>
</span><span id="__span-0-84"><a id="__codelineno-0-84" name="__codelineno-0-84"></a> <span class="n">voff_slow</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="mf">0.0</span><span class="p">,</span>
</span><span id="__span-0-85"><a id="__codelineno-0-85" name="__codelineno-0-85"></a> <span class="n">tau_slow</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="mf">50.0</span><span class="p">,</span>
</span><span id="__span-0-86"><a id="__codelineno-0-86" name="__codelineno-0-86"></a> <span class="n">vthr</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="mf">2.0</span><span class="p">,</span>
</span><span id="__span-0-87"><a id="__codelineno-0-87" name="__codelineno-0-87"></a> <span class="n">dtype</span><span class="p">:</span> <span class="n">DTypeLike</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">float32</span><span class="p">,</span>
</span><span id="__span-0-88"><a id="__codelineno-0-88" name="__codelineno-0-88"></a> <span class="p">):</span>
</span><span id="__span-0-89"><a id="__codelineno-0-89" name="__codelineno-0-89"></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;Initialize the FitzHugh-Nagumo neuron model.</span>
</span><span id="__span-0-90"><a id="__codelineno-0-90" name="__codelineno-0-90"></a>
</span><span id="__span-0-91"><a id="__codelineno-0-91" name="__codelineno-0-91"></a><span class="sd"> Args:</span>
</span><span id="__span-0-92"><a id="__codelineno-0-92" name="__codelineno-0-92"></a><span class="sd"> tsyn: Synaptic time constant for input current decay. Can be scalar or per-neuron array.</span>
</span><span id="__span-0-93"><a id="__codelineno-0-93" name="__codelineno-0-93"></a><span class="sd"> C: Membrane capacitance. Can be scalar or per-neuron array.</span>
</span><span id="__span-0-94"><a id="__codelineno-0-94" name="__codelineno-0-94"></a><span class="sd"> gmax_pasive: Maximal conductance of passive current. Can be scalar or per-neuron array.</span>
</span><span id="__span-0-95"><a id="__codelineno-0-95" name="__codelineno-0-95"></a><span class="sd"> Erev_pasive: Reversal potential for passive current. Can be scalar or per-neuron array.</span>
</span><span id="__span-0-96"><a id="__codelineno-0-96" name="__codelineno-0-96"></a><span class="sd"> a_fast: Amplitude of fast current. Can be scalar or per-neuron array.</span>
</span><span id="__span-0-97"><a id="__codelineno-0-97" name="__codelineno-0-97"></a><span class="sd"> voff_fast: Voltage offset for fast current activation. Can be scalar or per-neuron array.</span>
</span><span id="__span-0-98"><a id="__codelineno-0-98" name="__codelineno-0-98"></a><span class="sd"> tau_fast: Time constant for fast current (typically 0 for instantaneous). Can be scalar or per-neuron array.</span>
</span><span id="__span-0-99"><a id="__codelineno-0-99" name="__codelineno-0-99"></a><span class="sd"> a_slow: Amplitude of slow current. Can be scalar or per-neuron array.</span>
</span><span id="__span-0-100"><a id="__codelineno-0-100" name="__codelineno-0-100"></a><span class="sd"> voff_slow: Voltage offset for slow current activation. Can be scalar or per-neuron array.</span>
</span><span id="__span-0-101"><a id="__codelineno-0-101" name="__codelineno-0-101"></a><span class="sd"> tau_slow: Time constant for slow recovery variable. Can be scalar or per-neuron array.</span>
</span><span id="__span-0-102"><a id="__codelineno-0-102" name="__codelineno-0-102"></a><span class="sd"> vthr: Voltage threshold for spike generation. Can be scalar or per-neuron array.</span>
</span><span id="__span-0-103"><a id="__codelineno-0-103" name="__codelineno-0-103"></a><span class="sd"> dtype: Data type for arrays (default: float32).</span>
</span><span id="__span-0-104"><a id="__codelineno-0-104" name="__codelineno-0-104"></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="__span-0-105"><a id="__codelineno-0-105" name="__codelineno-0-105"></a> <span class="bp">self</span><span class="o">.</span><span class="n">dtype</span> <span class="o">=</span> <span class="n">dtype</span>
</span><span id="__span-0-106"><a id="__codelineno-0-106" name="__codelineno-0-106"></a>
</span><span id="__span-0-107"><a id="__codelineno-0-107" name="__codelineno-0-107"></a> <span class="bp">self</span><span class="o">.</span><span class="n">tsyn</span> <span class="o">=</span> <span class="n">tsyn</span>
</span><span id="__span-0-108"><a id="__codelineno-0-108" name="__codelineno-0-108"></a> <span class="bp">self</span><span class="o">.</span><span class="n">C</span> <span class="o">=</span> <span class="n">C</span>
</span><span id="__span-0-109"><a id="__codelineno-0-109" name="__codelineno-0-109"></a> <span class="bp">self</span><span class="o">.</span><span class="n">gmax_pasive</span> <span class="o">=</span> <span class="n">gmax_pasive</span>
</span><span id="__span-0-110"><a id="__codelineno-0-110" name="__codelineno-0-110"></a> <span class="bp">self</span><span class="o">.</span><span class="n">Erev_pasive</span> <span class="o">=</span> <span class="n">Erev_pasive</span>
</span><span id="__span-0-111"><a id="__codelineno-0-111" name="__codelineno-0-111"></a> <span class="bp">self</span><span class="o">.</span><span class="n">a_fast</span> <span class="o">=</span> <span class="n">a_fast</span>
</span><span id="__span-0-112"><a id="__codelineno-0-112" name="__codelineno-0-112"></a> <span class="bp">self</span><span class="o">.</span><span class="n">voff_fast</span> <span class="o">=</span> <span class="n">voff_fast</span>
</span><span id="__span-0-113"><a id="__codelineno-0-113" name="__codelineno-0-113"></a> <span class="bp">self</span><span class="o">.</span><span class="n">tau_fast</span> <span class="o">=</span> <span class="n">tau_fast</span>
</span><span id="__span-0-114"><a id="__codelineno-0-114" name="__codelineno-0-114"></a> <span class="bp">self</span><span class="o">.</span><span class="n">a_slow</span> <span class="o">=</span> <span class="n">a_slow</span>
</span><span id="__span-0-115"><a id="__codelineno-0-115" name="__codelineno-0-115"></a> <span class="bp">self</span><span class="o">.</span><span class="n">voff_slow</span> <span class="o">=</span> <span class="n">voff_slow</span>
</span><span id="__span-0-116"><a id="__codelineno-0-116" name="__codelineno-0-116"></a> <span class="bp">self</span><span class="o">.</span><span class="n">tau_slow</span> <span class="o">=</span> <span class="n">tau_slow</span>
</span><span id="__span-0-117"><a id="__codelineno-0-117" name="__codelineno-0-117"></a> <span class="bp">self</span><span class="o">.</span><span class="n">vthr</span> <span class="o">=</span> <span class="n">vthr</span>
</span><span id="__span-0-118"><a id="__codelineno-0-118" name="__codelineno-0-118"></a>
</span><span id="__span-0-119"><a id="__codelineno-0-119" name="__codelineno-0-119"></a> <span class="k">def</span><span class="w"> </span><span class="nf">init_state</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">n_neurons</span><span class="p">:</span> <span class="nb">int</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;neurons 3&quot;</span><span class="p">]:</span>
</span><span id="__span-0-120"><a id="__codelineno-0-120" name="__codelineno-0-120"></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;Initialize the neuron state variables.</span>
</span><span id="__span-0-121"><a id="__codelineno-0-121" name="__codelineno-0-121"></a>
</span><span id="__span-0-122"><a id="__codelineno-0-122" name="__codelineno-0-122"></a><span class="sd"> Args:</span>
</span><span id="__span-0-123"><a id="__codelineno-0-123" name="__codelineno-0-123"></a><span class="sd"> n_neurons: Number of neurons to initialize.</span>
</span><span id="__span-0-124"><a id="__codelineno-0-124" name="__codelineno-0-124"></a>
</span><span id="__span-0-125"><a id="__codelineno-0-125" name="__codelineno-0-125"></a><span class="sd"> Returns:</span>
</span><span id="__span-0-126"><a id="__codelineno-0-126" name="__codelineno-0-126"></a><span class="sd"> Initial state array of shape (neurons, 3) containing [v, v_slow, i_app],</span>
</span><span id="__span-0-127"><a id="__codelineno-0-127" name="__codelineno-0-127"></a><span class="sd"> where v is membrane voltage, v_slow is the slow recovery variable,</span>
</span><span id="__span-0-128"><a id="__codelineno-0-128" name="__codelineno-0-128"></a><span class="sd"> and i_app is the applied synaptic current.</span>
</span><span id="__span-0-129"><a id="__codelineno-0-129" name="__codelineno-0-129"></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="__span-0-130"><a id="__codelineno-0-130" name="__codelineno-0-130"></a> <span class="k">return</span> <span class="n">jnp</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">n_neurons</span><span class="p">,</span> <span class="mi">3</span><span class="p">),</span> <span class="n">dtype</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dtype</span><span class="p">)</span>
</span><span id="__span-0-131"><a id="__codelineno-0-131" name="__codelineno-0-131"></a>
</span><span id="__span-0-132"><a id="__codelineno-0-132" name="__codelineno-0-132"></a> <span class="k">def</span><span class="w"> </span><span class="nf">IV_inst</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">v</span><span class="p">:</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;...&quot;</span><span class="p">],</span> <span class="n">Vrest</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mi">0</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;...&quot;</span><span class="p">]:</span>
</span><span id="__span-0-133"><a id="__codelineno-0-133" name="__codelineno-0-133"></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;Compute instantaneous I-V relationship with fast and slow currents at rest.</span>
</span><span id="__span-0-134"><a id="__codelineno-0-134" name="__codelineno-0-134"></a>
</span><span id="__span-0-135"><a id="__codelineno-0-135" name="__codelineno-0-135"></a><span class="sd"> Args:</span>
</span><span id="__span-0-136"><a id="__codelineno-0-136" name="__codelineno-0-136"></a><span class="sd"> v: Membrane voltage.</span>
</span><span id="__span-0-137"><a id="__codelineno-0-137" name="__codelineno-0-137"></a><span class="sd"> Vrest: Resting voltage for both fast and slow currents (default: 0).</span>
</span><span id="__span-0-138"><a id="__codelineno-0-138" name="__codelineno-0-138"></a>
</span><span id="__span-0-139"><a id="__codelineno-0-139" name="__codelineno-0-139"></a><span class="sd"> Returns:</span>
</span><span id="__span-0-140"><a id="__codelineno-0-140" name="__codelineno-0-140"></a><span class="sd"> Total current at voltage v with both fast and slow currents evaluated at Vrest.</span>
</span><span id="__span-0-141"><a id="__codelineno-0-141" name="__codelineno-0-141"></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="__span-0-142"><a id="__codelineno-0-142" name="__codelineno-0-142"></a> <span class="n">I_pasive</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">gmax_pasive</span> <span class="o">*</span> <span class="p">(</span><span class="n">v</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">Erev_pasive</span><span class="p">)</span>
</span><span id="__span-0-143"><a id="__codelineno-0-143" name="__codelineno-0-143"></a> <span class="n">I_fast</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">a_fast</span> <span class="o">*</span> <span class="n">jnp</span><span class="o">.</span><span class="n">tanh</span><span class="p">(</span><span class="n">Vrest</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">voff_fast</span><span class="p">)</span>
</span><span id="__span-0-144"><a id="__codelineno-0-144" name="__codelineno-0-144"></a> <span class="n">I_slow</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">a_slow</span> <span class="o">*</span> <span class="n">jnp</span><span class="o">.</span><span class="n">tanh</span><span class="p">(</span><span class="n">Vrest</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">voff_slow</span><span class="p">)</span>
</span><span id="__span-0-145"><a id="__codelineno-0-145" name="__codelineno-0-145"></a>
</span><span id="__span-0-146"><a id="__codelineno-0-146" name="__codelineno-0-146"></a> <span class="k">return</span> <span class="n">I_pasive</span> <span class="o">+</span> <span class="n">I_fast</span> <span class="o">+</span> <span class="n">I_slow</span>
</span><span id="__span-0-147"><a id="__codelineno-0-147" name="__codelineno-0-147"></a>
</span><span id="__span-0-148"><a id="__codelineno-0-148" name="__codelineno-0-148"></a> <span class="k">def</span><span class="w"> </span><span class="nf">IV_fast</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">v</span><span class="p">:</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;...&quot;</span><span class="p">],</span> <span class="n">Vrest</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mi">0</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;...&quot;</span><span class="p">]:</span>
</span><span id="__span-0-149"><a id="__codelineno-0-149" name="__codelineno-0-149"></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;Compute I-V relationship with fast current at voltage v and slow current at rest.</span>
</span><span id="__span-0-150"><a id="__codelineno-0-150" name="__codelineno-0-150"></a>
</span><span id="__span-0-151"><a id="__codelineno-0-151" name="__codelineno-0-151"></a><span class="sd"> Args:</span>
</span><span id="__span-0-152"><a id="__codelineno-0-152" name="__codelineno-0-152"></a><span class="sd"> v: Membrane voltage for passive and fast currents.</span>
</span><span id="__span-0-153"><a id="__codelineno-0-153" name="__codelineno-0-153"></a><span class="sd"> Vrest: Resting voltage for slow current (default: 0).</span>
</span><span id="__span-0-154"><a id="__codelineno-0-154" name="__codelineno-0-154"></a>
</span><span id="__span-0-155"><a id="__codelineno-0-155" name="__codelineno-0-155"></a><span class="sd"> Returns:</span>
</span><span id="__span-0-156"><a id="__codelineno-0-156" name="__codelineno-0-156"></a><span class="sd"> Total current with fast dynamics responding to v and slow current at Vrest.</span>
</span><span id="__span-0-157"><a id="__codelineno-0-157" name="__codelineno-0-157"></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="__span-0-158"><a id="__codelineno-0-158" name="__codelineno-0-158"></a> <span class="n">I_pasive</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">gmax_pasive</span> <span class="o">*</span> <span class="p">(</span><span class="n">v</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">Erev_pasive</span><span class="p">)</span>
</span><span id="__span-0-159"><a id="__codelineno-0-159" name="__codelineno-0-159"></a> <span class="n">I_fast</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">a_fast</span> <span class="o">*</span> <span class="n">jnp</span><span class="o">.</span><span class="n">tanh</span><span class="p">(</span><span class="n">v</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">voff_fast</span><span class="p">)</span>
</span><span id="__span-0-160"><a id="__codelineno-0-160" name="__codelineno-0-160"></a> <span class="n">I_slow</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">a_slow</span> <span class="o">*</span> <span class="n">jnp</span><span class="o">.</span><span class="n">tanh</span><span class="p">(</span><span class="n">Vrest</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">voff_slow</span><span class="p">)</span>
</span><span id="__span-0-161"><a id="__codelineno-0-161" name="__codelineno-0-161"></a>
</span><span id="__span-0-162"><a id="__codelineno-0-162" name="__codelineno-0-162"></a> <span class="k">return</span> <span class="n">I_pasive</span> <span class="o">+</span> <span class="n">I_fast</span> <span class="o">+</span> <span class="n">I_slow</span>
</span><span id="__span-0-163"><a id="__codelineno-0-163" name="__codelineno-0-163"></a>
</span><span id="__span-0-164"><a id="__codelineno-0-164" name="__codelineno-0-164"></a> <span class="k">def</span><span class="w"> </span><span class="nf">IV_slow</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">v</span><span class="p">:</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;...&quot;</span><span class="p">],</span> <span class="n">Vrest</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mi">0</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;...&quot;</span><span class="p">]:</span>
</span><span id="__span-0-165"><a id="__codelineno-0-165" name="__codelineno-0-165"></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;Compute steady-state I-V relationship with all currents at voltage v.</span>
</span><span id="__span-0-166"><a id="__codelineno-0-166" name="__codelineno-0-166"></a>
</span><span id="__span-0-167"><a id="__codelineno-0-167" name="__codelineno-0-167"></a><span class="sd"> Args:</span>
</span><span id="__span-0-168"><a id="__codelineno-0-168" name="__codelineno-0-168"></a><span class="sd"> v: Membrane voltage for all currents.</span>
</span><span id="__span-0-169"><a id="__codelineno-0-169" name="__codelineno-0-169"></a><span class="sd"> Vrest: Unused parameter for API consistency (default: 0).</span>
</span><span id="__span-0-170"><a id="__codelineno-0-170" name="__codelineno-0-170"></a>
</span><span id="__span-0-171"><a id="__codelineno-0-171" name="__codelineno-0-171"></a><span class="sd"> Returns:</span>
</span><span id="__span-0-172"><a id="__codelineno-0-172" name="__codelineno-0-172"></a><span class="sd"> Total steady-state current with all currents responding to v.</span>
</span><span id="__span-0-173"><a id="__codelineno-0-173" name="__codelineno-0-173"></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="__span-0-174"><a id="__codelineno-0-174" name="__codelineno-0-174"></a> <span class="n">I_pasive</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">gmax_pasive</span> <span class="o">*</span> <span class="p">(</span><span class="n">v</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">Erev_pasive</span><span class="p">)</span>
</span><span id="__span-0-175"><a id="__codelineno-0-175" name="__codelineno-0-175"></a> <span class="n">I_fast</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">a_fast</span> <span class="o">*</span> <span class="n">jnp</span><span class="o">.</span><span class="n">tanh</span><span class="p">(</span><span class="n">v</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">voff_fast</span><span class="p">)</span>
</span><span id="__span-0-176"><a id="__codelineno-0-176" name="__codelineno-0-176"></a> <span class="n">I_slow</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">a_slow</span> <span class="o">*</span> <span class="n">jnp</span><span class="o">.</span><span class="n">tanh</span><span class="p">(</span><span class="n">v</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">voff_slow</span><span class="p">)</span>
</span><span id="__span-0-177"><a id="__codelineno-0-177" name="__codelineno-0-177"></a>
</span><span id="__span-0-178"><a id="__codelineno-0-178" name="__codelineno-0-178"></a> <span class="k">return</span> <span class="n">I_pasive</span> <span class="o">+</span> <span class="n">I_fast</span> <span class="o">+</span> <span class="n">I_slow</span>
</span><span id="__span-0-179"><a id="__codelineno-0-179" name="__codelineno-0-179"></a>
</span><span id="__span-0-180"><a id="__codelineno-0-180" name="__codelineno-0-180"></a> <span class="k">def</span><span class="w"> </span><span class="nf">dynamics</span><span class="p">(</span>
</span><span id="__span-0-181"><a id="__codelineno-0-181" name="__codelineno-0-181"></a> <span class="bp">self</span><span class="p">,</span>
</span><span id="__span-0-182"><a id="__codelineno-0-182" name="__codelineno-0-182"></a> <span class="n">t</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
</span><span id="__span-0-183"><a id="__codelineno-0-183" name="__codelineno-0-183"></a> <span class="n">y</span><span class="p">:</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;neurons 3&quot;</span><span class="p">],</span>
</span><span id="__span-0-184"><a id="__codelineno-0-184" name="__codelineno-0-184"></a> <span class="n">args</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">],</span>
</span><span id="__span-0-185"><a id="__codelineno-0-185" name="__codelineno-0-185"></a> <span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;neurons 3&quot;</span><span class="p">]:</span>
</span><span id="__span-0-186"><a id="__codelineno-0-186" name="__codelineno-0-186"></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;Compute time derivatives of the neuron state variables.</span>
</span><span id="__span-0-187"><a id="__codelineno-0-187" name="__codelineno-0-187"></a>
</span><span id="__span-0-188"><a id="__codelineno-0-188" name="__codelineno-0-188"></a><span class="sd"> This implements the FitzHugh-Nagumo dynamics with passive, fast, and slow currents:</span>
</span><span id="__span-0-189"><a id="__codelineno-0-189" name="__codelineno-0-189"></a><span class="sd"> - dv/dt: Fast membrane voltage dynamics</span>
</span><span id="__span-0-190"><a id="__codelineno-0-190" name="__codelineno-0-190"></a><span class="sd"> - dv_slow/dt: Slow recovery variable dynamics</span>
</span><span id="__span-0-191"><a id="__codelineno-0-191" name="__codelineno-0-191"></a><span class="sd"> - di_app/dt: Synaptic current decay</span>
</span><span id="__span-0-192"><a id="__codelineno-0-192" name="__codelineno-0-192"></a>
</span><span id="__span-0-193"><a id="__codelineno-0-193" name="__codelineno-0-193"></a><span class="sd"> Args:</span>
</span><span id="__span-0-194"><a id="__codelineno-0-194" name="__codelineno-0-194"></a><span class="sd"> t: Current simulation time (unused but required by framework).</span>
</span><span id="__span-0-195"><a id="__codelineno-0-195" name="__codelineno-0-195"></a><span class="sd"> y: State array of shape (neurons, 3) containing [v, v_slow, i_app].</span>
</span><span id="__span-0-196"><a id="__codelineno-0-196" name="__codelineno-0-196"></a><span class="sd"> args: Additional arguments (unused but required by framework).</span>
</span><span id="__span-0-197"><a id="__codelineno-0-197" name="__codelineno-0-197"></a>
</span><span id="__span-0-198"><a id="__codelineno-0-198" name="__codelineno-0-198"></a><span class="sd"> Returns:</span>
</span><span id="__span-0-199"><a id="__codelineno-0-199" name="__codelineno-0-199"></a><span class="sd"> Time derivatives of shape (neurons, 3) containing [dv/dt, dv_slow/dt, di_app/dt].</span>
</span><span id="__span-0-200"><a id="__codelineno-0-200" name="__codelineno-0-200"></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="__span-0-201"><a id="__codelineno-0-201" name="__codelineno-0-201"></a> <span class="n">v</span> <span class="o">=</span> <span class="n">y</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span>
</span><span id="__span-0-202"><a id="__codelineno-0-202" name="__codelineno-0-202"></a> <span class="n">v_slow</span> <span class="o">=</span> <span class="n">y</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">]</span>
</span><span id="__span-0-203"><a id="__codelineno-0-203" name="__codelineno-0-203"></a> <span class="n">i_app</span> <span class="o">=</span> <span class="n">y</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">]</span>
</span><span id="__span-0-204"><a id="__codelineno-0-204" name="__codelineno-0-204"></a>
</span><span id="__span-0-205"><a id="__codelineno-0-205" name="__codelineno-0-205"></a> <span class="n">I_pasive</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">gmax_pasive</span> <span class="o">*</span> <span class="p">(</span><span class="n">v</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">Erev_pasive</span><span class="p">)</span>
</span><span id="__span-0-206"><a id="__codelineno-0-206" name="__codelineno-0-206"></a> <span class="n">I_fast</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">a_fast</span> <span class="o">*</span> <span class="n">jnp</span><span class="o">.</span><span class="n">tanh</span><span class="p">(</span><span class="n">v</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">voff_fast</span><span class="p">)</span>
</span><span id="__span-0-207"><a id="__codelineno-0-207" name="__codelineno-0-207"></a> <span class="n">I_slow</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">a_slow</span> <span class="o">*</span> <span class="n">jnp</span><span class="o">.</span><span class="n">tanh</span><span class="p">(</span><span class="n">v_slow</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">voff_slow</span><span class="p">)</span>
</span><span id="__span-0-208"><a id="__codelineno-0-208" name="__codelineno-0-208"></a>
</span><span id="__span-0-209"><a id="__codelineno-0-209" name="__codelineno-0-209"></a> <span class="n">i_sum</span> <span class="o">=</span> <span class="n">I_pasive</span> <span class="o">+</span> <span class="n">I_fast</span> <span class="o">+</span> <span class="n">I_slow</span>
</span><span id="__span-0-210"><a id="__codelineno-0-210" name="__codelineno-0-210"></a>
</span><span id="__span-0-211"><a id="__codelineno-0-211" name="__codelineno-0-211"></a> <span class="n">dv_dt</span> <span class="o">=</span> <span class="p">(</span><span class="n">i_app</span> <span class="o">-</span> <span class="n">i_sum</span><span class="p">)</span> <span class="o">/</span> <span class="bp">self</span><span class="o">.</span><span class="n">C</span>
</span><span id="__span-0-212"><a id="__codelineno-0-212" name="__codelineno-0-212"></a> <span class="n">dvslow_dt</span> <span class="o">=</span> <span class="p">(</span><span class="n">v</span> <span class="o">-</span> <span class="n">v_slow</span><span class="p">)</span> <span class="o">/</span> <span class="bp">self</span><span class="o">.</span><span class="n">tau_slow</span>
</span><span id="__span-0-213"><a id="__codelineno-0-213" name="__codelineno-0-213"></a> <span class="n">di_dt</span> <span class="o">=</span> <span class="o">-</span><span class="n">i_app</span> <span class="o">/</span> <span class="bp">self</span><span class="o">.</span><span class="n">tsyn</span>
</span><span id="__span-0-214"><a id="__codelineno-0-214" name="__codelineno-0-214"></a>
</span><span id="__span-0-215"><a id="__codelineno-0-215" name="__codelineno-0-215"></a> <span class="k">return</span> <span class="n">jnp</span><span class="o">.</span><span class="n">stack</span><span class="p">([</span><span class="n">dv_dt</span><span class="p">,</span> <span class="n">dvslow_dt</span><span class="p">,</span> <span class="n">di_dt</span><span class="p">],</span> <span class="n">axis</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
</span><span id="__span-0-216"><a id="__codelineno-0-216" name="__codelineno-0-216"></a>
</span><span id="__span-0-217"><a id="__codelineno-0-217" name="__codelineno-0-217"></a> <span class="k">def</span><span class="w"> </span><span class="nf">spike_condition</span><span class="p">(</span>
</span><span id="__span-0-218"><a id="__codelineno-0-218" name="__codelineno-0-218"></a> <span class="bp">self</span><span class="p">,</span>
</span><span id="__span-0-219"><a id="__codelineno-0-219" name="__codelineno-0-219"></a> <span class="n">t</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
</span><span id="__span-0-220"><a id="__codelineno-0-220" name="__codelineno-0-220"></a> <span class="n">y</span><span class="p">:</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;neurons 3&quot;</span><span class="p">],</span>
</span><span id="__span-0-221"><a id="__codelineno-0-221" name="__codelineno-0-221"></a> <span class="o">**</span><span class="n">kwargs</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">],</span>
</span><span id="__span-0-222"><a id="__codelineno-0-222" name="__codelineno-0-222"></a> <span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot; neurons&quot;</span><span class="p">]:</span>
</span><span id="__span-0-223"><a id="__codelineno-0-223" name="__codelineno-0-223"></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;Compute spike condition for event detection.</span>
</span><span id="__span-0-224"><a id="__codelineno-0-224" name="__codelineno-0-224"></a>
</span><span id="__span-0-225"><a id="__codelineno-0-225" name="__codelineno-0-225"></a><span class="sd"> A spike is triggered when this function crosses zero (v &gt;= vthr).</span>
</span><span id="__span-0-226"><a id="__codelineno-0-226" name="__codelineno-0-226"></a>
</span><span id="__span-0-227"><a id="__codelineno-0-227" name="__codelineno-0-227"></a><span class="sd"> Args:</span>
</span><span id="__span-0-228"><a id="__codelineno-0-228" name="__codelineno-0-228"></a><span class="sd"> t: Current simulation time (unused but required by event detection).</span>
</span><span id="__span-0-229"><a id="__codelineno-0-229" name="__codelineno-0-229"></a><span class="sd"> y: State array of shape (neurons, 3) containing [v, v_slow, i_app].</span>
</span><span id="__span-0-230"><a id="__codelineno-0-230" name="__codelineno-0-230"></a><span class="sd"> **kwargs: Additional keyword arguments (unused).</span>
</span><span id="__span-0-231"><a id="__codelineno-0-231" name="__codelineno-0-231"></a>
</span><span id="__span-0-232"><a id="__codelineno-0-232" name="__codelineno-0-232"></a><span class="sd"> Returns:</span>
</span><span id="__span-0-233"><a id="__codelineno-0-233" name="__codelineno-0-233"></a><span class="sd"> Spike condition array of shape (neurons,). Positive values indicate v &gt; vthr.</span>
</span><span id="__span-0-234"><a id="__codelineno-0-234" name="__codelineno-0-234"></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="__span-0-235"><a id="__codelineno-0-235" name="__codelineno-0-235"></a> <span class="k">return</span> <span class="n">y</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">vthr</span>
</span></code></pre></div></td></tr></table></div>
</details>
<div class="doc doc-children">
<h5 id="felice.neuron_models.FHNRS-functions">Functions</h5>
<div class="doc doc-object doc-function">
<h6 id="felice.neuron_models.FHNRS.__init__" class="doc doc-heading">
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <code class="highlight language-python"><span class="fm">__init__</span><span class="p">(</span><span class="o">*</span><span class="p">,</span> <span class="n">tsyn</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="mf">1.0</span><span class="p">,</span> <span class="n">C</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="mf">1.0</span><span class="p">,</span> <span class="n">gmax_pasive</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="mf">1.0</span><span class="p">,</span> <span class="n">Erev_pasive</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="mf">0.0</span><span class="p">,</span> <span class="n">a_fast</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="o">-</span><span class="mf">2.0</span><span class="p">,</span> <span class="n">voff_fast</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="mf">0.0</span><span class="p">,</span> <span class="n">tau_fast</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="mf">0.0</span><span class="p">,</span> <span class="n">a_slow</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="mf">2.0</span><span class="p">,</span> <span class="n">voff_slow</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="mf">0.0</span><span class="p">,</span> <span class="n">tau_slow</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="mf">50.0</span><span class="p">,</span> <span class="n">vthr</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="mf">2.0</span><span class="p">,</span> <span class="n">dtype</span><span class="p">:</span> <span class="n">DTypeLike</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">float32</span><span class="p">)</span></code>
</h6>
<div class="doc doc-contents ">
<p>Initialize the FitzHugh-Nagumo neuron model.</p>
<p><span class="doc-section-title">Parameters:</span></p>
<table>
<thead>
<tr>
<th>Name</th>
<th>Type</th>
<th>Description</th>
<th>Default</th>
</tr>
</thead>
<tbody>
<tr class="doc-section-item">
<td>
<code>tsyn</code>
</td>
<td>
<code><span title="typing.Union">Union</span>[<span title="int">int</span>, <span title="float">float</span>, <span title="jax.numpy.ndarray">ndarray</span>]</code>
</td>
<td>
<div class="doc-md-description">
<p>Synaptic time constant for input current decay. Can be scalar or per-neuron array.</p>
</div>
</td>
<td>
<code>1.0</code>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>C</code>
</td>
<td>
<code><span title="typing.Union">Union</span>[<span title="int">int</span>, <span title="float">float</span>, <span title="jax.numpy.ndarray">ndarray</span>]</code>
</td>
<td>
<div class="doc-md-description">
<p>Membrane capacitance. Can be scalar or per-neuron array.</p>
</div>
</td>
<td>
<code>1.0</code>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>gmax_pasive</code>
</td>
<td>
<code><span title="typing.Union">Union</span>[<span title="int">int</span>, <span title="float">float</span>, <span title="jax.numpy.ndarray">ndarray</span>]</code>
</td>
<td>
<div class="doc-md-description">
<p>Maximal conductance of passive current. Can be scalar or per-neuron array.</p>
</div>
</td>
<td>
<code>1.0</code>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>Erev_pasive</code>
</td>
<td>
<code><span title="typing.Union">Union</span>[<span title="int">int</span>, <span title="float">float</span>, <span title="jax.numpy.ndarray">ndarray</span>]</code>
</td>
<td>
<div class="doc-md-description">
<p>Reversal potential for passive current. Can be scalar or per-neuron array.</p>
</div>
</td>
<td>
<code>0.0</code>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>a_fast</code>
</td>
<td>
<code><span title="typing.Union">Union</span>[<span title="int">int</span>, <span title="float">float</span>, <span title="jax.numpy.ndarray">ndarray</span>]</code>
</td>
<td>
<div class="doc-md-description">
<p>Amplitude of fast current. Can be scalar or per-neuron array.</p>
</div>
</td>
<td>
<code>-2.0</code>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>voff_fast</code>
</td>
<td>
<code><span title="typing.Union">Union</span>[<span title="int">int</span>, <span title="float">float</span>, <span title="jax.numpy.ndarray">ndarray</span>]</code>
</td>
<td>
<div class="doc-md-description">
<p>Voltage offset for fast current activation. Can be scalar or per-neuron array.</p>
</div>
</td>
<td>
<code>0.0</code>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>tau_fast</code>
</td>
<td>
<code><span title="typing.Union">Union</span>[<span title="int">int</span>, <span title="float">float</span>, <span title="jax.numpy.ndarray">ndarray</span>]</code>
</td>
<td>
<div class="doc-md-description">
<p>Time constant for fast current (typically 0 for instantaneous). Can be scalar or per-neuron array.</p>
</div>
</td>
<td>
<code>0.0</code>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>a_slow</code>
</td>
<td>
<code><span title="typing.Union">Union</span>[<span title="int">int</span>, <span title="float">float</span>, <span title="jax.numpy.ndarray">ndarray</span>]</code>
</td>
<td>
<div class="doc-md-description">
<p>Amplitude of slow current. Can be scalar or per-neuron array.</p>
</div>
</td>
<td>
<code>2.0</code>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>voff_slow</code>
</td>
<td>
<code><span title="typing.Union">Union</span>[<span title="int">int</span>, <span title="float">float</span>, <span title="jax.numpy.ndarray">ndarray</span>]</code>
</td>
<td>
<div class="doc-md-description">
<p>Voltage offset for slow current activation. Can be scalar or per-neuron array.</p>
</div>
</td>
<td>
<code>0.0</code>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>tau_slow</code>
</td>
<td>
<code><span title="typing.Union">Union</span>[<span title="int">int</span>, <span title="float">float</span>, <span title="jax.numpy.ndarray">ndarray</span>]</code>
</td>
<td>
<div class="doc-md-description">
<p>Time constant for slow recovery variable. Can be scalar or per-neuron array.</p>
</div>
</td>
<td>
<code>50.0</code>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>vthr</code>
</td>
<td>
<code><span title="typing.Union">Union</span>[<span title="int">int</span>, <span title="float">float</span>, <span title="jax.numpy.ndarray">ndarray</span>]</code>
</td>
<td>
<div class="doc-md-description">
<p>Voltage threshold for spike generation. Can be scalar or per-neuron array.</p>
</div>
</td>
<td>
<code>2.0</code>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>dtype</code>
</td>
<td>
<code><span title="jaxtyping.DTypeLike">DTypeLike</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Data type for arrays (default: float32).</p>
</div>
</td>
<td>
<code><span title="jax.numpy.float32">float32</span></code>
</td>
</tr>
</tbody>
</table>
<details class="mkdocstrings-source">
<summary>Source code in <code>felice/neuron_models/fhn.py</code></summary>
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal"><a href="#__codelineno-0-73"> 73</a></span>
<span class="normal"><a href="#__codelineno-0-74"> 74</a></span>
<span class="normal"><a href="#__codelineno-0-75"> 75</a></span>
<span class="normal"><a href="#__codelineno-0-76"> 76</a></span>
<span class="normal"><a href="#__codelineno-0-77"> 77</a></span>
<span class="normal"><a href="#__codelineno-0-78"> 78</a></span>
<span class="normal"><a href="#__codelineno-0-79"> 79</a></span>
<span class="normal"><a href="#__codelineno-0-80"> 80</a></span>
<span class="normal"><a href="#__codelineno-0-81"> 81</a></span>
<span class="normal"><a href="#__codelineno-0-82"> 82</a></span>
<span class="normal"><a href="#__codelineno-0-83"> 83</a></span>
<span class="normal"><a href="#__codelineno-0-84"> 84</a></span>
<span class="normal"><a href="#__codelineno-0-85"> 85</a></span>
<span class="normal"><a href="#__codelineno-0-86"> 86</a></span>
<span class="normal"><a href="#__codelineno-0-87"> 87</a></span>
<span class="normal"><a href="#__codelineno-0-88"> 88</a></span>
<span class="normal"><a href="#__codelineno-0-89"> 89</a></span>
<span class="normal"><a href="#__codelineno-0-90"> 90</a></span>
<span class="normal"><a href="#__codelineno-0-91"> 91</a></span>
<span class="normal"><a href="#__codelineno-0-92"> 92</a></span>
<span class="normal"><a href="#__codelineno-0-93"> 93</a></span>
<span class="normal"><a href="#__codelineno-0-94"> 94</a></span>
<span class="normal"><a href="#__codelineno-0-95"> 95</a></span>
<span class="normal"><a href="#__codelineno-0-96"> 96</a></span>
<span class="normal"><a href="#__codelineno-0-97"> 97</a></span>
<span class="normal"><a href="#__codelineno-0-98"> 98</a></span>
<span class="normal"><a href="#__codelineno-0-99"> 99</a></span>
<span class="normal"><a href="#__codelineno-0-100">100</a></span>
<span class="normal"><a href="#__codelineno-0-101">101</a></span>
<span class="normal"><a href="#__codelineno-0-102">102</a></span>
<span class="normal"><a href="#__codelineno-0-103">103</a></span>
<span class="normal"><a href="#__codelineno-0-104">104</a></span>
<span class="normal"><a href="#__codelineno-0-105">105</a></span>
<span class="normal"><a href="#__codelineno-0-106">106</a></span>
<span class="normal"><a href="#__codelineno-0-107">107</a></span>
<span class="normal"><a href="#__codelineno-0-108">108</a></span>
<span class="normal"><a href="#__codelineno-0-109">109</a></span>
<span class="normal"><a href="#__codelineno-0-110">110</a></span>
<span class="normal"><a href="#__codelineno-0-111">111</a></span>
<span class="normal"><a href="#__codelineno-0-112">112</a></span>
<span class="normal"><a href="#__codelineno-0-113">113</a></span>
<span class="normal"><a href="#__codelineno-0-114">114</a></span>
<span class="normal"><a href="#__codelineno-0-115">115</a></span>
<span class="normal"><a href="#__codelineno-0-116">116</a></span>
<span class="normal"><a href="#__codelineno-0-117">117</a></span></pre></div></td><td class="code"><div><pre><span></span><code><span id="__span-0-73"><a id="__codelineno-0-73" name="__codelineno-0-73"></a><span class="k">def</span><span class="w"> </span><span class="fm">__init__</span><span class="p">(</span>
</span><span id="__span-0-74"><a id="__codelineno-0-74" name="__codelineno-0-74"></a> <span class="bp">self</span><span class="p">,</span>
</span><span id="__span-0-75"><a id="__codelineno-0-75" name="__codelineno-0-75"></a> <span class="o">*</span><span class="p">,</span>
</span><span id="__span-0-76"><a id="__codelineno-0-76" name="__codelineno-0-76"></a> <span class="n">tsyn</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="mf">1.0</span><span class="p">,</span>
</span><span id="__span-0-77"><a id="__codelineno-0-77" name="__codelineno-0-77"></a> <span class="n">C</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="mf">1.0</span><span class="p">,</span>
</span><span id="__span-0-78"><a id="__codelineno-0-78" name="__codelineno-0-78"></a> <span class="n">gmax_pasive</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="mf">1.0</span><span class="p">,</span>
</span><span id="__span-0-79"><a id="__codelineno-0-79" name="__codelineno-0-79"></a> <span class="n">Erev_pasive</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="mf">0.0</span><span class="p">,</span>
</span><span id="__span-0-80"><a id="__codelineno-0-80" name="__codelineno-0-80"></a> <span class="n">a_fast</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="o">-</span><span class="mf">2.0</span><span class="p">,</span>
</span><span id="__span-0-81"><a id="__codelineno-0-81" name="__codelineno-0-81"></a> <span class="n">voff_fast</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="mf">0.0</span><span class="p">,</span>
</span><span id="__span-0-82"><a id="__codelineno-0-82" name="__codelineno-0-82"></a> <span class="n">tau_fast</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="mf">0.0</span><span class="p">,</span>
</span><span id="__span-0-83"><a id="__codelineno-0-83" name="__codelineno-0-83"></a> <span class="n">a_slow</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="mf">2.0</span><span class="p">,</span>
</span><span id="__span-0-84"><a id="__codelineno-0-84" name="__codelineno-0-84"></a> <span class="n">voff_slow</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="mf">0.0</span><span class="p">,</span>
</span><span id="__span-0-85"><a id="__codelineno-0-85" name="__codelineno-0-85"></a> <span class="n">tau_slow</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="mf">50.0</span><span class="p">,</span>
</span><span id="__span-0-86"><a id="__codelineno-0-86" name="__codelineno-0-86"></a> <span class="n">vthr</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="nb">int</span><span class="p">,</span> <span class="nb">float</span><span class="p">,</span> <span class="n">jnp</span><span class="o">.</span><span class="n">ndarray</span><span class="p">]</span> <span class="o">=</span> <span class="mf">2.0</span><span class="p">,</span>
</span><span id="__span-0-87"><a id="__codelineno-0-87" name="__codelineno-0-87"></a> <span class="n">dtype</span><span class="p">:</span> <span class="n">DTypeLike</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">float32</span><span class="p">,</span>
</span><span id="__span-0-88"><a id="__codelineno-0-88" name="__codelineno-0-88"></a><span class="p">):</span>
</span><span id="__span-0-89"><a id="__codelineno-0-89" name="__codelineno-0-89"></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;Initialize the FitzHugh-Nagumo neuron model.</span>
</span><span id="__span-0-90"><a id="__codelineno-0-90" name="__codelineno-0-90"></a>
</span><span id="__span-0-91"><a id="__codelineno-0-91" name="__codelineno-0-91"></a><span class="sd"> Args:</span>
</span><span id="__span-0-92"><a id="__codelineno-0-92" name="__codelineno-0-92"></a><span class="sd"> tsyn: Synaptic time constant for input current decay. Can be scalar or per-neuron array.</span>
</span><span id="__span-0-93"><a id="__codelineno-0-93" name="__codelineno-0-93"></a><span class="sd"> C: Membrane capacitance. Can be scalar or per-neuron array.</span>
</span><span id="__span-0-94"><a id="__codelineno-0-94" name="__codelineno-0-94"></a><span class="sd"> gmax_pasive: Maximal conductance of passive current. Can be scalar or per-neuron array.</span>
</span><span id="__span-0-95"><a id="__codelineno-0-95" name="__codelineno-0-95"></a><span class="sd"> Erev_pasive: Reversal potential for passive current. Can be scalar or per-neuron array.</span>
</span><span id="__span-0-96"><a id="__codelineno-0-96" name="__codelineno-0-96"></a><span class="sd"> a_fast: Amplitude of fast current. Can be scalar or per-neuron array.</span>
</span><span id="__span-0-97"><a id="__codelineno-0-97" name="__codelineno-0-97"></a><span class="sd"> voff_fast: Voltage offset for fast current activation. Can be scalar or per-neuron array.</span>
</span><span id="__span-0-98"><a id="__codelineno-0-98" name="__codelineno-0-98"></a><span class="sd"> tau_fast: Time constant for fast current (typically 0 for instantaneous). Can be scalar or per-neuron array.</span>
</span><span id="__span-0-99"><a id="__codelineno-0-99" name="__codelineno-0-99"></a><span class="sd"> a_slow: Amplitude of slow current. Can be scalar or per-neuron array.</span>
</span><span id="__span-0-100"><a id="__codelineno-0-100" name="__codelineno-0-100"></a><span class="sd"> voff_slow: Voltage offset for slow current activation. Can be scalar or per-neuron array.</span>
</span><span id="__span-0-101"><a id="__codelineno-0-101" name="__codelineno-0-101"></a><span class="sd"> tau_slow: Time constant for slow recovery variable. Can be scalar or per-neuron array.</span>
</span><span id="__span-0-102"><a id="__codelineno-0-102" name="__codelineno-0-102"></a><span class="sd"> vthr: Voltage threshold for spike generation. Can be scalar or per-neuron array.</span>
</span><span id="__span-0-103"><a id="__codelineno-0-103" name="__codelineno-0-103"></a><span class="sd"> dtype: Data type for arrays (default: float32).</span>
</span><span id="__span-0-104"><a id="__codelineno-0-104" name="__codelineno-0-104"></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="__span-0-105"><a id="__codelineno-0-105" name="__codelineno-0-105"></a> <span class="bp">self</span><span class="o">.</span><span class="n">dtype</span> <span class="o">=</span> <span class="n">dtype</span>
</span><span id="__span-0-106"><a id="__codelineno-0-106" name="__codelineno-0-106"></a>
</span><span id="__span-0-107"><a id="__codelineno-0-107" name="__codelineno-0-107"></a> <span class="bp">self</span><span class="o">.</span><span class="n">tsyn</span> <span class="o">=</span> <span class="n">tsyn</span>
</span><span id="__span-0-108"><a id="__codelineno-0-108" name="__codelineno-0-108"></a> <span class="bp">self</span><span class="o">.</span><span class="n">C</span> <span class="o">=</span> <span class="n">C</span>
</span><span id="__span-0-109"><a id="__codelineno-0-109" name="__codelineno-0-109"></a> <span class="bp">self</span><span class="o">.</span><span class="n">gmax_pasive</span> <span class="o">=</span> <span class="n">gmax_pasive</span>
</span><span id="__span-0-110"><a id="__codelineno-0-110" name="__codelineno-0-110"></a> <span class="bp">self</span><span class="o">.</span><span class="n">Erev_pasive</span> <span class="o">=</span> <span class="n">Erev_pasive</span>
</span><span id="__span-0-111"><a id="__codelineno-0-111" name="__codelineno-0-111"></a> <span class="bp">self</span><span class="o">.</span><span class="n">a_fast</span> <span class="o">=</span> <span class="n">a_fast</span>
</span><span id="__span-0-112"><a id="__codelineno-0-112" name="__codelineno-0-112"></a> <span class="bp">self</span><span class="o">.</span><span class="n">voff_fast</span> <span class="o">=</span> <span class="n">voff_fast</span>
</span><span id="__span-0-113"><a id="__codelineno-0-113" name="__codelineno-0-113"></a> <span class="bp">self</span><span class="o">.</span><span class="n">tau_fast</span> <span class="o">=</span> <span class="n">tau_fast</span>
</span><span id="__span-0-114"><a id="__codelineno-0-114" name="__codelineno-0-114"></a> <span class="bp">self</span><span class="o">.</span><span class="n">a_slow</span> <span class="o">=</span> <span class="n">a_slow</span>
</span><span id="__span-0-115"><a id="__codelineno-0-115" name="__codelineno-0-115"></a> <span class="bp">self</span><span class="o">.</span><span class="n">voff_slow</span> <span class="o">=</span> <span class="n">voff_slow</span>
</span><span id="__span-0-116"><a id="__codelineno-0-116" name="__codelineno-0-116"></a> <span class="bp">self</span><span class="o">.</span><span class="n">tau_slow</span> <span class="o">=</span> <span class="n">tau_slow</span>
</span><span id="__span-0-117"><a id="__codelineno-0-117" name="__codelineno-0-117"></a> <span class="bp">self</span><span class="o">.</span><span class="n">vthr</span> <span class="o">=</span> <span class="n">vthr</span>
</span></code></pre></div></td></tr></table></div>
</details>
</div>
</div>
<div class="doc doc-object doc-function">
<h6 id="felice.neuron_models.FHNRS.init_state" class="doc doc-heading">
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <code class="highlight language-python"><span class="n">init_state</span><span class="p">(</span><span class="n">n_neurons</span><span class="p">:</span> <span class="nb">int</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s1">&#39;neurons 3&#39;</span><span class="p">]</span></code>
</h6>
<div class="doc doc-contents ">
<p>Initialize the neuron state variables.</p>
<p><span class="doc-section-title">Parameters:</span></p>
<table>
<thead>
<tr>
<th>Name</th>
<th>Type</th>
<th>Description</th>
<th>Default</th>
</tr>
</thead>
<tbody>
<tr class="doc-section-item">
<td>
<code>n_neurons</code>
</td>
<td>
<code><span title="int">int</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Number of neurons to initialize.</p>
</div>
</td>
<td>
<em>required</em>
</td>
</tr>
</tbody>
</table>
<p><span class="doc-section-title">Returns:</span></p>
<table>
<thead>
<tr>
<th>Type</th>
<th>Description</th>
</tr>
</thead>
<tbody>
<tr class="doc-section-item">
<td>
<code><span title="jaxtyping.Float">Float</span>[<span title="jaxtyping.Array">Array</span>, &#39;neurons 3&#39;]</code>
</td>
<td>
<div class="doc-md-description">
<p>Initial state array of shape (neurons, 3) containing [v, v_slow, i_app],</p>
</div>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code><span title="jaxtyping.Float">Float</span>[<span title="jaxtyping.Array">Array</span>, &#39;neurons 3&#39;]</code>
</td>
<td>
<div class="doc-md-description">
<p>where v is membrane voltage, v_slow is the slow recovery variable,</p>
</div>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code><span title="jaxtyping.Float">Float</span>[<span title="jaxtyping.Array">Array</span>, &#39;neurons 3&#39;]</code>
</td>
<td>
<div class="doc-md-description">
<p>and i_app is the applied synaptic current.</p>
</div>
</td>
</tr>
</tbody>
</table>
<details class="mkdocstrings-source">
<summary>Source code in <code>felice/neuron_models/fhn.py</code></summary>
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal"><a href="#__codelineno-0-119">119</a></span>
<span class="normal"><a href="#__codelineno-0-120">120</a></span>
<span class="normal"><a href="#__codelineno-0-121">121</a></span>
<span class="normal"><a href="#__codelineno-0-122">122</a></span>
<span class="normal"><a href="#__codelineno-0-123">123</a></span>
<span class="normal"><a href="#__codelineno-0-124">124</a></span>
<span class="normal"><a href="#__codelineno-0-125">125</a></span>
<span class="normal"><a href="#__codelineno-0-126">126</a></span>
<span class="normal"><a href="#__codelineno-0-127">127</a></span>
<span class="normal"><a href="#__codelineno-0-128">128</a></span>
<span class="normal"><a href="#__codelineno-0-129">129</a></span>
<span class="normal"><a href="#__codelineno-0-130">130</a></span></pre></div></td><td class="code"><div><pre><span></span><code><span id="__span-0-119"><a id="__codelineno-0-119" name="__codelineno-0-119"></a><span class="k">def</span><span class="w"> </span><span class="nf">init_state</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">n_neurons</span><span class="p">:</span> <span class="nb">int</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;neurons 3&quot;</span><span class="p">]:</span>
</span><span id="__span-0-120"><a id="__codelineno-0-120" name="__codelineno-0-120"></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;Initialize the neuron state variables.</span>
</span><span id="__span-0-121"><a id="__codelineno-0-121" name="__codelineno-0-121"></a>
</span><span id="__span-0-122"><a id="__codelineno-0-122" name="__codelineno-0-122"></a><span class="sd"> Args:</span>
</span><span id="__span-0-123"><a id="__codelineno-0-123" name="__codelineno-0-123"></a><span class="sd"> n_neurons: Number of neurons to initialize.</span>
</span><span id="__span-0-124"><a id="__codelineno-0-124" name="__codelineno-0-124"></a>
</span><span id="__span-0-125"><a id="__codelineno-0-125" name="__codelineno-0-125"></a><span class="sd"> Returns:</span>
</span><span id="__span-0-126"><a id="__codelineno-0-126" name="__codelineno-0-126"></a><span class="sd"> Initial state array of shape (neurons, 3) containing [v, v_slow, i_app],</span>
</span><span id="__span-0-127"><a id="__codelineno-0-127" name="__codelineno-0-127"></a><span class="sd"> where v is membrane voltage, v_slow is the slow recovery variable,</span>
</span><span id="__span-0-128"><a id="__codelineno-0-128" name="__codelineno-0-128"></a><span class="sd"> and i_app is the applied synaptic current.</span>
</span><span id="__span-0-129"><a id="__codelineno-0-129" name="__codelineno-0-129"></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="__span-0-130"><a id="__codelineno-0-130" name="__codelineno-0-130"></a> <span class="k">return</span> <span class="n">jnp</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">n_neurons</span><span class="p">,</span> <span class="mi">3</span><span class="p">),</span> <span class="n">dtype</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dtype</span><span class="p">)</span>
</span></code></pre></div></td></tr></table></div>
</details>
</div>
</div>
<div class="doc doc-object doc-function">
<h6 id="felice.neuron_models.FHNRS.IV_inst" class="doc doc-heading">
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <code class="highlight language-python"><span class="n">IV_inst</span><span class="p">(</span><span class="n">v</span><span class="p">:</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="o">...</span><span class="p">],</span> <span class="n">Vrest</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mi">0</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="o">...</span><span class="p">]</span></code>
</h6>
<div class="doc doc-contents ">
<p>Compute instantaneous I-V relationship with fast and slow currents at rest.</p>
<p><span class="doc-section-title">Parameters:</span></p>
<table>
<thead>
<tr>
<th>Name</th>
<th>Type</th>
<th>Description</th>
<th>Default</th>
</tr>
</thead>
<tbody>
<tr class="doc-section-item">
<td>
<code>v</code>
</td>
<td>
<code><span title="jaxtyping.Float">Float</span>[<span title="jaxtyping.Array">Array</span>, ...]</code>
</td>
<td>
<div class="doc-md-description">
<p>Membrane voltage.</p>
</div>
</td>
<td>
<em>required</em>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>Vrest</code>
</td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Resting voltage for both fast and slow currents (default: 0).</p>
</div>
</td>
<td>
<code>0</code>
</td>
</tr>
</tbody>
</table>
<p><span class="doc-section-title">Returns:</span></p>
<table>
<thead>
<tr>
<th>Type</th>
<th>Description</th>
</tr>
</thead>
<tbody>
<tr class="doc-section-item">
<td>
<code><span title="jaxtyping.Float">Float</span>[<span title="jaxtyping.Array">Array</span>, ...]</code>
</td>
<td>
<div class="doc-md-description">
<p>Total current at voltage v with both fast and slow currents evaluated at Vrest.</p>
</div>
</td>
</tr>
</tbody>
</table>
<details class="mkdocstrings-source">
<summary>Source code in <code>felice/neuron_models/fhn.py</code></summary>
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal"><a href="#__codelineno-0-132">132</a></span>
<span class="normal"><a href="#__codelineno-0-133">133</a></span>
<span class="normal"><a href="#__codelineno-0-134">134</a></span>
<span class="normal"><a href="#__codelineno-0-135">135</a></span>
<span class="normal"><a href="#__codelineno-0-136">136</a></span>
<span class="normal"><a href="#__codelineno-0-137">137</a></span>
<span class="normal"><a href="#__codelineno-0-138">138</a></span>
<span class="normal"><a href="#__codelineno-0-139">139</a></span>
<span class="normal"><a href="#__codelineno-0-140">140</a></span>
<span class="normal"><a href="#__codelineno-0-141">141</a></span>
<span class="normal"><a href="#__codelineno-0-142">142</a></span>
<span class="normal"><a href="#__codelineno-0-143">143</a></span>
<span class="normal"><a href="#__codelineno-0-144">144</a></span>
<span class="normal"><a href="#__codelineno-0-145">145</a></span>
<span class="normal"><a href="#__codelineno-0-146">146</a></span></pre></div></td><td class="code"><div><pre><span></span><code><span id="__span-0-132"><a id="__codelineno-0-132" name="__codelineno-0-132"></a><span class="k">def</span><span class="w"> </span><span class="nf">IV_inst</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">v</span><span class="p">:</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;...&quot;</span><span class="p">],</span> <span class="n">Vrest</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mi">0</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;...&quot;</span><span class="p">]:</span>
</span><span id="__span-0-133"><a id="__codelineno-0-133" name="__codelineno-0-133"></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;Compute instantaneous I-V relationship with fast and slow currents at rest.</span>
</span><span id="__span-0-134"><a id="__codelineno-0-134" name="__codelineno-0-134"></a>
</span><span id="__span-0-135"><a id="__codelineno-0-135" name="__codelineno-0-135"></a><span class="sd"> Args:</span>
</span><span id="__span-0-136"><a id="__codelineno-0-136" name="__codelineno-0-136"></a><span class="sd"> v: Membrane voltage.</span>
</span><span id="__span-0-137"><a id="__codelineno-0-137" name="__codelineno-0-137"></a><span class="sd"> Vrest: Resting voltage for both fast and slow currents (default: 0).</span>
</span><span id="__span-0-138"><a id="__codelineno-0-138" name="__codelineno-0-138"></a>
</span><span id="__span-0-139"><a id="__codelineno-0-139" name="__codelineno-0-139"></a><span class="sd"> Returns:</span>
</span><span id="__span-0-140"><a id="__codelineno-0-140" name="__codelineno-0-140"></a><span class="sd"> Total current at voltage v with both fast and slow currents evaluated at Vrest.</span>
</span><span id="__span-0-141"><a id="__codelineno-0-141" name="__codelineno-0-141"></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="__span-0-142"><a id="__codelineno-0-142" name="__codelineno-0-142"></a> <span class="n">I_pasive</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">gmax_pasive</span> <span class="o">*</span> <span class="p">(</span><span class="n">v</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">Erev_pasive</span><span class="p">)</span>
</span><span id="__span-0-143"><a id="__codelineno-0-143" name="__codelineno-0-143"></a> <span class="n">I_fast</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">a_fast</span> <span class="o">*</span> <span class="n">jnp</span><span class="o">.</span><span class="n">tanh</span><span class="p">(</span><span class="n">Vrest</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">voff_fast</span><span class="p">)</span>
</span><span id="__span-0-144"><a id="__codelineno-0-144" name="__codelineno-0-144"></a> <span class="n">I_slow</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">a_slow</span> <span class="o">*</span> <span class="n">jnp</span><span class="o">.</span><span class="n">tanh</span><span class="p">(</span><span class="n">Vrest</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">voff_slow</span><span class="p">)</span>
</span><span id="__span-0-145"><a id="__codelineno-0-145" name="__codelineno-0-145"></a>
</span><span id="__span-0-146"><a id="__codelineno-0-146" name="__codelineno-0-146"></a> <span class="k">return</span> <span class="n">I_pasive</span> <span class="o">+</span> <span class="n">I_fast</span> <span class="o">+</span> <span class="n">I_slow</span>
</span></code></pre></div></td></tr></table></div>
</details>
</div>
</div>
<div class="doc doc-object doc-function">
<h6 id="felice.neuron_models.FHNRS.IV_fast" class="doc doc-heading">
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <code class="highlight language-python"><span class="n">IV_fast</span><span class="p">(</span><span class="n">v</span><span class="p">:</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="o">...</span><span class="p">],</span> <span class="n">Vrest</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mi">0</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="o">...</span><span class="p">]</span></code>
</h6>
<div class="doc doc-contents ">
<p>Compute I-V relationship with fast current at voltage v and slow current at rest.</p>
<p><span class="doc-section-title">Parameters:</span></p>
<table>
<thead>
<tr>
<th>Name</th>
<th>Type</th>
<th>Description</th>
<th>Default</th>
</tr>
</thead>
<tbody>
<tr class="doc-section-item">
<td>
<code>v</code>
</td>
<td>
<code><span title="jaxtyping.Float">Float</span>[<span title="jaxtyping.Array">Array</span>, ...]</code>
</td>
<td>
<div class="doc-md-description">
<p>Membrane voltage for passive and fast currents.</p>
</div>
</td>
<td>
<em>required</em>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>Vrest</code>
</td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Resting voltage for slow current (default: 0).</p>
</div>
</td>
<td>
<code>0</code>
</td>
</tr>
</tbody>
</table>
<p><span class="doc-section-title">Returns:</span></p>
<table>
<thead>
<tr>
<th>Type</th>
<th>Description</th>
</tr>
</thead>
<tbody>
<tr class="doc-section-item">
<td>
<code><span title="jaxtyping.Float">Float</span>[<span title="jaxtyping.Array">Array</span>, ...]</code>
</td>
<td>
<div class="doc-md-description">
<p>Total current with fast dynamics responding to v and slow current at Vrest.</p>
</div>
</td>
</tr>
</tbody>
</table>
<details class="mkdocstrings-source">
<summary>Source code in <code>felice/neuron_models/fhn.py</code></summary>
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal"><a href="#__codelineno-0-148">148</a></span>
<span class="normal"><a href="#__codelineno-0-149">149</a></span>
<span class="normal"><a href="#__codelineno-0-150">150</a></span>
<span class="normal"><a href="#__codelineno-0-151">151</a></span>
<span class="normal"><a href="#__codelineno-0-152">152</a></span>
<span class="normal"><a href="#__codelineno-0-153">153</a></span>
<span class="normal"><a href="#__codelineno-0-154">154</a></span>
<span class="normal"><a href="#__codelineno-0-155">155</a></span>
<span class="normal"><a href="#__codelineno-0-156">156</a></span>
<span class="normal"><a href="#__codelineno-0-157">157</a></span>
<span class="normal"><a href="#__codelineno-0-158">158</a></span>
<span class="normal"><a href="#__codelineno-0-159">159</a></span>
<span class="normal"><a href="#__codelineno-0-160">160</a></span>
<span class="normal"><a href="#__codelineno-0-161">161</a></span>
<span class="normal"><a href="#__codelineno-0-162">162</a></span></pre></div></td><td class="code"><div><pre><span></span><code><span id="__span-0-148"><a id="__codelineno-0-148" name="__codelineno-0-148"></a><span class="k">def</span><span class="w"> </span><span class="nf">IV_fast</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">v</span><span class="p">:</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;...&quot;</span><span class="p">],</span> <span class="n">Vrest</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mi">0</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;...&quot;</span><span class="p">]:</span>
</span><span id="__span-0-149"><a id="__codelineno-0-149" name="__codelineno-0-149"></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;Compute I-V relationship with fast current at voltage v and slow current at rest.</span>
</span><span id="__span-0-150"><a id="__codelineno-0-150" name="__codelineno-0-150"></a>
</span><span id="__span-0-151"><a id="__codelineno-0-151" name="__codelineno-0-151"></a><span class="sd"> Args:</span>
</span><span id="__span-0-152"><a id="__codelineno-0-152" name="__codelineno-0-152"></a><span class="sd"> v: Membrane voltage for passive and fast currents.</span>
</span><span id="__span-0-153"><a id="__codelineno-0-153" name="__codelineno-0-153"></a><span class="sd"> Vrest: Resting voltage for slow current (default: 0).</span>
</span><span id="__span-0-154"><a id="__codelineno-0-154" name="__codelineno-0-154"></a>
</span><span id="__span-0-155"><a id="__codelineno-0-155" name="__codelineno-0-155"></a><span class="sd"> Returns:</span>
</span><span id="__span-0-156"><a id="__codelineno-0-156" name="__codelineno-0-156"></a><span class="sd"> Total current with fast dynamics responding to v and slow current at Vrest.</span>
</span><span id="__span-0-157"><a id="__codelineno-0-157" name="__codelineno-0-157"></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="__span-0-158"><a id="__codelineno-0-158" name="__codelineno-0-158"></a> <span class="n">I_pasive</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">gmax_pasive</span> <span class="o">*</span> <span class="p">(</span><span class="n">v</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">Erev_pasive</span><span class="p">)</span>
</span><span id="__span-0-159"><a id="__codelineno-0-159" name="__codelineno-0-159"></a> <span class="n">I_fast</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">a_fast</span> <span class="o">*</span> <span class="n">jnp</span><span class="o">.</span><span class="n">tanh</span><span class="p">(</span><span class="n">v</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">voff_fast</span><span class="p">)</span>
</span><span id="__span-0-160"><a id="__codelineno-0-160" name="__codelineno-0-160"></a> <span class="n">I_slow</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">a_slow</span> <span class="o">*</span> <span class="n">jnp</span><span class="o">.</span><span class="n">tanh</span><span class="p">(</span><span class="n">Vrest</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">voff_slow</span><span class="p">)</span>
</span><span id="__span-0-161"><a id="__codelineno-0-161" name="__codelineno-0-161"></a>
</span><span id="__span-0-162"><a id="__codelineno-0-162" name="__codelineno-0-162"></a> <span class="k">return</span> <span class="n">I_pasive</span> <span class="o">+</span> <span class="n">I_fast</span> <span class="o">+</span> <span class="n">I_slow</span>
</span></code></pre></div></td></tr></table></div>
</details>
</div>
</div>
<div class="doc doc-object doc-function">
<h6 id="felice.neuron_models.FHNRS.IV_slow" class="doc doc-heading">
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <code class="highlight language-python"><span class="n">IV_slow</span><span class="p">(</span><span class="n">v</span><span class="p">:</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="o">...</span><span class="p">],</span> <span class="n">Vrest</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mi">0</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="o">...</span><span class="p">]</span></code>
</h6>
<div class="doc doc-contents ">
<p>Compute steady-state I-V relationship with all currents at voltage v.</p>
<p><span class="doc-section-title">Parameters:</span></p>
<table>
<thead>
<tr>
<th>Name</th>
<th>Type</th>
<th>Description</th>
<th>Default</th>
</tr>
</thead>
<tbody>
<tr class="doc-section-item">
<td>
<code>v</code>
</td>
<td>
<code><span title="jaxtyping.Float">Float</span>[<span title="jaxtyping.Array">Array</span>, ...]</code>
</td>
<td>
<div class="doc-md-description">
<p>Membrane voltage for all currents.</p>
</div>
</td>
<td>
<em>required</em>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>Vrest</code>
</td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Unused parameter for API consistency (default: 0).</p>
</div>
</td>
<td>
<code>0</code>
</td>
</tr>
</tbody>
</table>
<p><span class="doc-section-title">Returns:</span></p>
<table>
<thead>
<tr>
<th>Type</th>
<th>Description</th>
</tr>
</thead>
<tbody>
<tr class="doc-section-item">
<td>
<code><span title="jaxtyping.Float">Float</span>[<span title="jaxtyping.Array">Array</span>, ...]</code>
</td>
<td>
<div class="doc-md-description">
<p>Total steady-state current with all currents responding to v.</p>
</div>
</td>
</tr>
</tbody>
</table>
<details class="mkdocstrings-source">
<summary>Source code in <code>felice/neuron_models/fhn.py</code></summary>
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal"><a href="#__codelineno-0-164">164</a></span>
<span class="normal"><a href="#__codelineno-0-165">165</a></span>
<span class="normal"><a href="#__codelineno-0-166">166</a></span>
<span class="normal"><a href="#__codelineno-0-167">167</a></span>
<span class="normal"><a href="#__codelineno-0-168">168</a></span>
<span class="normal"><a href="#__codelineno-0-169">169</a></span>
<span class="normal"><a href="#__codelineno-0-170">170</a></span>
<span class="normal"><a href="#__codelineno-0-171">171</a></span>
<span class="normal"><a href="#__codelineno-0-172">172</a></span>
<span class="normal"><a href="#__codelineno-0-173">173</a></span>
<span class="normal"><a href="#__codelineno-0-174">174</a></span>
<span class="normal"><a href="#__codelineno-0-175">175</a></span>
<span class="normal"><a href="#__codelineno-0-176">176</a></span>
<span class="normal"><a href="#__codelineno-0-177">177</a></span>
<span class="normal"><a href="#__codelineno-0-178">178</a></span></pre></div></td><td class="code"><div><pre><span></span><code><span id="__span-0-164"><a id="__codelineno-0-164" name="__codelineno-0-164"></a><span class="k">def</span><span class="w"> </span><span class="nf">IV_slow</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">v</span><span class="p">:</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;...&quot;</span><span class="p">],</span> <span class="n">Vrest</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mi">0</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;...&quot;</span><span class="p">]:</span>
</span><span id="__span-0-165"><a id="__codelineno-0-165" name="__codelineno-0-165"></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;Compute steady-state I-V relationship with all currents at voltage v.</span>
</span><span id="__span-0-166"><a id="__codelineno-0-166" name="__codelineno-0-166"></a>
</span><span id="__span-0-167"><a id="__codelineno-0-167" name="__codelineno-0-167"></a><span class="sd"> Args:</span>
</span><span id="__span-0-168"><a id="__codelineno-0-168" name="__codelineno-0-168"></a><span class="sd"> v: Membrane voltage for all currents.</span>
</span><span id="__span-0-169"><a id="__codelineno-0-169" name="__codelineno-0-169"></a><span class="sd"> Vrest: Unused parameter for API consistency (default: 0).</span>
</span><span id="__span-0-170"><a id="__codelineno-0-170" name="__codelineno-0-170"></a>
</span><span id="__span-0-171"><a id="__codelineno-0-171" name="__codelineno-0-171"></a><span class="sd"> Returns:</span>
</span><span id="__span-0-172"><a id="__codelineno-0-172" name="__codelineno-0-172"></a><span class="sd"> Total steady-state current with all currents responding to v.</span>
</span><span id="__span-0-173"><a id="__codelineno-0-173" name="__codelineno-0-173"></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="__span-0-174"><a id="__codelineno-0-174" name="__codelineno-0-174"></a> <span class="n">I_pasive</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">gmax_pasive</span> <span class="o">*</span> <span class="p">(</span><span class="n">v</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">Erev_pasive</span><span class="p">)</span>
</span><span id="__span-0-175"><a id="__codelineno-0-175" name="__codelineno-0-175"></a> <span class="n">I_fast</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">a_fast</span> <span class="o">*</span> <span class="n">jnp</span><span class="o">.</span><span class="n">tanh</span><span class="p">(</span><span class="n">v</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">voff_fast</span><span class="p">)</span>
</span><span id="__span-0-176"><a id="__codelineno-0-176" name="__codelineno-0-176"></a> <span class="n">I_slow</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">a_slow</span> <span class="o">*</span> <span class="n">jnp</span><span class="o">.</span><span class="n">tanh</span><span class="p">(</span><span class="n">v</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">voff_slow</span><span class="p">)</span>
</span><span id="__span-0-177"><a id="__codelineno-0-177" name="__codelineno-0-177"></a>
</span><span id="__span-0-178"><a id="__codelineno-0-178" name="__codelineno-0-178"></a> <span class="k">return</span> <span class="n">I_pasive</span> <span class="o">+</span> <span class="n">I_fast</span> <span class="o">+</span> <span class="n">I_slow</span>
</span></code></pre></div></td></tr></table></div>
</details>
</div>
</div>
<div class="doc doc-object doc-function">
<h6 id="felice.neuron_models.FHNRS.dynamics" class="doc doc-heading">
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <code class="highlight language-python"><span class="n">dynamics</span><span class="p">(</span><span class="n">t</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span> <span class="n">y</span><span class="p">:</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s1">&#39;neurons 3&#39;</span><span class="p">],</span> <span class="n">args</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">])</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s1">&#39;neurons 3&#39;</span><span class="p">]</span></code>
</h6>
<div class="doc doc-contents ">
<p>Compute time derivatives of the neuron state variables.</p>
<p>This implements the FitzHugh-Nagumo dynamics with passive, fast, and slow currents:
- dv/dt: Fast membrane voltage dynamics
- dv_slow/dt: Slow recovery variable dynamics
- di_app/dt: Synaptic current decay</p>
<p><span class="doc-section-title">Parameters:</span></p>
<table>
<thead>
<tr>
<th>Name</th>
<th>Type</th>
<th>Description</th>
<th>Default</th>
</tr>
</thead>
<tbody>
<tr class="doc-section-item">
<td>
<code>t</code>
</td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Current simulation time (unused but required by framework).</p>
</div>
</td>
<td>
<em>required</em>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>y</code>
</td>
<td>
<code><span title="jaxtyping.Float">Float</span>[<span title="jaxtyping.Array">Array</span>, &#39;neurons 3&#39;]</code>
</td>
<td>
<div class="doc-md-description">
<p>State array of shape (neurons, 3) containing [v, v_slow, i_app].</p>
</div>
</td>
<td>
<em>required</em>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>args</code>
</td>
<td>
<code><span title="typing.Dict">Dict</span>[<span title="str">str</span>, <span title="typing.Any">Any</span>]</code>
</td>
<td>
<div class="doc-md-description">
<p>Additional arguments (unused but required by framework).</p>
</div>
</td>
<td>
<em>required</em>
</td>
</tr>
</tbody>
</table>
<p><span class="doc-section-title">Returns:</span></p>
<table>
<thead>
<tr>
<th>Type</th>
<th>Description</th>
</tr>
</thead>
<tbody>
<tr class="doc-section-item">
<td>
<code><span title="jaxtyping.Float">Float</span>[<span title="jaxtyping.Array">Array</span>, &#39;neurons 3&#39;]</code>
</td>
<td>
<div class="doc-md-description">
<p>Time derivatives of shape (neurons, 3) containing [dv/dt, dv_slow/dt, di_app/dt].</p>
</div>
</td>
</tr>
</tbody>
</table>
<details class="mkdocstrings-source">
<summary>Source code in <code>felice/neuron_models/fhn.py</code></summary>
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal"><a href="#__codelineno-0-180">180</a></span>
<span class="normal"><a href="#__codelineno-0-181">181</a></span>
<span class="normal"><a href="#__codelineno-0-182">182</a></span>
<span class="normal"><a href="#__codelineno-0-183">183</a></span>
<span class="normal"><a href="#__codelineno-0-184">184</a></span>
<span class="normal"><a href="#__codelineno-0-185">185</a></span>
<span class="normal"><a href="#__codelineno-0-186">186</a></span>
<span class="normal"><a href="#__codelineno-0-187">187</a></span>
<span class="normal"><a href="#__codelineno-0-188">188</a></span>
<span class="normal"><a href="#__codelineno-0-189">189</a></span>
<span class="normal"><a href="#__codelineno-0-190">190</a></span>
<span class="normal"><a href="#__codelineno-0-191">191</a></span>
<span class="normal"><a href="#__codelineno-0-192">192</a></span>
<span class="normal"><a href="#__codelineno-0-193">193</a></span>
<span class="normal"><a href="#__codelineno-0-194">194</a></span>
<span class="normal"><a href="#__codelineno-0-195">195</a></span>
<span class="normal"><a href="#__codelineno-0-196">196</a></span>
<span class="normal"><a href="#__codelineno-0-197">197</a></span>
<span class="normal"><a href="#__codelineno-0-198">198</a></span>
<span class="normal"><a href="#__codelineno-0-199">199</a></span>
<span class="normal"><a href="#__codelineno-0-200">200</a></span>
<span class="normal"><a href="#__codelineno-0-201">201</a></span>
<span class="normal"><a href="#__codelineno-0-202">202</a></span>
<span class="normal"><a href="#__codelineno-0-203">203</a></span>
<span class="normal"><a href="#__codelineno-0-204">204</a></span>
<span class="normal"><a href="#__codelineno-0-205">205</a></span>
<span class="normal"><a href="#__codelineno-0-206">206</a></span>
<span class="normal"><a href="#__codelineno-0-207">207</a></span>
<span class="normal"><a href="#__codelineno-0-208">208</a></span>
<span class="normal"><a href="#__codelineno-0-209">209</a></span>
<span class="normal"><a href="#__codelineno-0-210">210</a></span>
<span class="normal"><a href="#__codelineno-0-211">211</a></span>
<span class="normal"><a href="#__codelineno-0-212">212</a></span>
<span class="normal"><a href="#__codelineno-0-213">213</a></span>
<span class="normal"><a href="#__codelineno-0-214">214</a></span>
<span class="normal"><a href="#__codelineno-0-215">215</a></span></pre></div></td><td class="code"><div><pre><span></span><code><span id="__span-0-180"><a id="__codelineno-0-180" name="__codelineno-0-180"></a><span class="k">def</span><span class="w"> </span><span class="nf">dynamics</span><span class="p">(</span>
</span><span id="__span-0-181"><a id="__codelineno-0-181" name="__codelineno-0-181"></a> <span class="bp">self</span><span class="p">,</span>
</span><span id="__span-0-182"><a id="__codelineno-0-182" name="__codelineno-0-182"></a> <span class="n">t</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
</span><span id="__span-0-183"><a id="__codelineno-0-183" name="__codelineno-0-183"></a> <span class="n">y</span><span class="p">:</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;neurons 3&quot;</span><span class="p">],</span>
</span><span id="__span-0-184"><a id="__codelineno-0-184" name="__codelineno-0-184"></a> <span class="n">args</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">],</span>
</span><span id="__span-0-185"><a id="__codelineno-0-185" name="__codelineno-0-185"></a><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;neurons 3&quot;</span><span class="p">]:</span>
</span><span id="__span-0-186"><a id="__codelineno-0-186" name="__codelineno-0-186"></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;Compute time derivatives of the neuron state variables.</span>
</span><span id="__span-0-187"><a id="__codelineno-0-187" name="__codelineno-0-187"></a>
</span><span id="__span-0-188"><a id="__codelineno-0-188" name="__codelineno-0-188"></a><span class="sd"> This implements the FitzHugh-Nagumo dynamics with passive, fast, and slow currents:</span>
</span><span id="__span-0-189"><a id="__codelineno-0-189" name="__codelineno-0-189"></a><span class="sd"> - dv/dt: Fast membrane voltage dynamics</span>
</span><span id="__span-0-190"><a id="__codelineno-0-190" name="__codelineno-0-190"></a><span class="sd"> - dv_slow/dt: Slow recovery variable dynamics</span>
</span><span id="__span-0-191"><a id="__codelineno-0-191" name="__codelineno-0-191"></a><span class="sd"> - di_app/dt: Synaptic current decay</span>
</span><span id="__span-0-192"><a id="__codelineno-0-192" name="__codelineno-0-192"></a>
</span><span id="__span-0-193"><a id="__codelineno-0-193" name="__codelineno-0-193"></a><span class="sd"> Args:</span>
</span><span id="__span-0-194"><a id="__codelineno-0-194" name="__codelineno-0-194"></a><span class="sd"> t: Current simulation time (unused but required by framework).</span>
</span><span id="__span-0-195"><a id="__codelineno-0-195" name="__codelineno-0-195"></a><span class="sd"> y: State array of shape (neurons, 3) containing [v, v_slow, i_app].</span>
</span><span id="__span-0-196"><a id="__codelineno-0-196" name="__codelineno-0-196"></a><span class="sd"> args: Additional arguments (unused but required by framework).</span>
</span><span id="__span-0-197"><a id="__codelineno-0-197" name="__codelineno-0-197"></a>
</span><span id="__span-0-198"><a id="__codelineno-0-198" name="__codelineno-0-198"></a><span class="sd"> Returns:</span>
</span><span id="__span-0-199"><a id="__codelineno-0-199" name="__codelineno-0-199"></a><span class="sd"> Time derivatives of shape (neurons, 3) containing [dv/dt, dv_slow/dt, di_app/dt].</span>
</span><span id="__span-0-200"><a id="__codelineno-0-200" name="__codelineno-0-200"></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="__span-0-201"><a id="__codelineno-0-201" name="__codelineno-0-201"></a> <span class="n">v</span> <span class="o">=</span> <span class="n">y</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span>
</span><span id="__span-0-202"><a id="__codelineno-0-202" name="__codelineno-0-202"></a> <span class="n">v_slow</span> <span class="o">=</span> <span class="n">y</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">]</span>
</span><span id="__span-0-203"><a id="__codelineno-0-203" name="__codelineno-0-203"></a> <span class="n">i_app</span> <span class="o">=</span> <span class="n">y</span><span class="p">[:,</span> <span class="mi">2</span><span class="p">]</span>
</span><span id="__span-0-204"><a id="__codelineno-0-204" name="__codelineno-0-204"></a>
</span><span id="__span-0-205"><a id="__codelineno-0-205" name="__codelineno-0-205"></a> <span class="n">I_pasive</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">gmax_pasive</span> <span class="o">*</span> <span class="p">(</span><span class="n">v</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">Erev_pasive</span><span class="p">)</span>
</span><span id="__span-0-206"><a id="__codelineno-0-206" name="__codelineno-0-206"></a> <span class="n">I_fast</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">a_fast</span> <span class="o">*</span> <span class="n">jnp</span><span class="o">.</span><span class="n">tanh</span><span class="p">(</span><span class="n">v</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">voff_fast</span><span class="p">)</span>
</span><span id="__span-0-207"><a id="__codelineno-0-207" name="__codelineno-0-207"></a> <span class="n">I_slow</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">a_slow</span> <span class="o">*</span> <span class="n">jnp</span><span class="o">.</span><span class="n">tanh</span><span class="p">(</span><span class="n">v_slow</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">voff_slow</span><span class="p">)</span>
</span><span id="__span-0-208"><a id="__codelineno-0-208" name="__codelineno-0-208"></a>
</span><span id="__span-0-209"><a id="__codelineno-0-209" name="__codelineno-0-209"></a> <span class="n">i_sum</span> <span class="o">=</span> <span class="n">I_pasive</span> <span class="o">+</span> <span class="n">I_fast</span> <span class="o">+</span> <span class="n">I_slow</span>
</span><span id="__span-0-210"><a id="__codelineno-0-210" name="__codelineno-0-210"></a>
</span><span id="__span-0-211"><a id="__codelineno-0-211" name="__codelineno-0-211"></a> <span class="n">dv_dt</span> <span class="o">=</span> <span class="p">(</span><span class="n">i_app</span> <span class="o">-</span> <span class="n">i_sum</span><span class="p">)</span> <span class="o">/</span> <span class="bp">self</span><span class="o">.</span><span class="n">C</span>
</span><span id="__span-0-212"><a id="__codelineno-0-212" name="__codelineno-0-212"></a> <span class="n">dvslow_dt</span> <span class="o">=</span> <span class="p">(</span><span class="n">v</span> <span class="o">-</span> <span class="n">v_slow</span><span class="p">)</span> <span class="o">/</span> <span class="bp">self</span><span class="o">.</span><span class="n">tau_slow</span>
</span><span id="__span-0-213"><a id="__codelineno-0-213" name="__codelineno-0-213"></a> <span class="n">di_dt</span> <span class="o">=</span> <span class="o">-</span><span class="n">i_app</span> <span class="o">/</span> <span class="bp">self</span><span class="o">.</span><span class="n">tsyn</span>
</span><span id="__span-0-214"><a id="__codelineno-0-214" name="__codelineno-0-214"></a>
</span><span id="__span-0-215"><a id="__codelineno-0-215" name="__codelineno-0-215"></a> <span class="k">return</span> <span class="n">jnp</span><span class="o">.</span><span class="n">stack</span><span class="p">([</span><span class="n">dv_dt</span><span class="p">,</span> <span class="n">dvslow_dt</span><span class="p">,</span> <span class="n">di_dt</span><span class="p">],</span> <span class="n">axis</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
</span></code></pre></div></td></tr></table></div>
</details>
</div>
</div>
<div class="doc doc-object doc-function">
<h6 id="felice.neuron_models.FHNRS.spike_condition" class="doc doc-heading">
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <code class="highlight language-python"><span class="n">spike_condition</span><span class="p">(</span><span class="n">t</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span> <span class="n">y</span><span class="p">:</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s1">&#39;neurons 3&#39;</span><span class="p">],</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">])</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s1">&#39; neurons&#39;</span><span class="p">]</span></code>
</h6>
<div class="doc doc-contents ">
<p>Compute spike condition for event detection.</p>
<p>A spike is triggered when this function crosses zero (v &gt;= vthr).</p>
<p><span class="doc-section-title">Parameters:</span></p>
<table>
<thead>
<tr>
<th>Name</th>
<th>Type</th>
<th>Description</th>
<th>Default</th>
</tr>
</thead>
<tbody>
<tr class="doc-section-item">
<td>
<code>t</code>
</td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Current simulation time (unused but required by event detection).</p>
</div>
</td>
<td>
<em>required</em>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>y</code>
</td>
<td>
<code><span title="jaxtyping.Float">Float</span>[<span title="jaxtyping.Array">Array</span>, &#39;neurons 3&#39;]</code>
</td>
<td>
<div class="doc-md-description">
<p>State array of shape (neurons, 3) containing [v, v_slow, i_app].</p>
</div>
</td>
<td>
<em>required</em>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>**kwargs</code>
</td>
<td>
<code><span title="typing.Dict">Dict</span>[<span title="str">str</span>, <span title="typing.Any">Any</span>]</code>
</td>
<td>
<div class="doc-md-description">
<p>Additional keyword arguments (unused).</p>
</div>
</td>
<td>
<code>{}</code>
</td>
</tr>
</tbody>
</table>
<p><span class="doc-section-title">Returns:</span></p>
<table>
<thead>
<tr>
<th>Type</th>
<th>Description</th>
</tr>
</thead>
<tbody>
<tr class="doc-section-item">
<td>
<code><span title="jaxtyping.Float">Float</span>[<span title="jaxtyping.Array">Array</span>, &#39; neurons&#39;]</code>
</td>
<td>
<div class="doc-md-description">
<p>Spike condition array of shape (neurons,). Positive values indicate v &gt; vthr.</p>
</div>
</td>
</tr>
</tbody>
</table>
<details class="mkdocstrings-source">
<summary>Source code in <code>felice/neuron_models/fhn.py</code></summary>
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal"><a href="#__codelineno-0-217">217</a></span>
<span class="normal"><a href="#__codelineno-0-218">218</a></span>
<span class="normal"><a href="#__codelineno-0-219">219</a></span>
<span class="normal"><a href="#__codelineno-0-220">220</a></span>
<span class="normal"><a href="#__codelineno-0-221">221</a></span>
<span class="normal"><a href="#__codelineno-0-222">222</a></span>
<span class="normal"><a href="#__codelineno-0-223">223</a></span>
<span class="normal"><a href="#__codelineno-0-224">224</a></span>
<span class="normal"><a href="#__codelineno-0-225">225</a></span>
<span class="normal"><a href="#__codelineno-0-226">226</a></span>
<span class="normal"><a href="#__codelineno-0-227">227</a></span>
<span class="normal"><a href="#__codelineno-0-228">228</a></span>
<span class="normal"><a href="#__codelineno-0-229">229</a></span>
<span class="normal"><a href="#__codelineno-0-230">230</a></span>
<span class="normal"><a href="#__codelineno-0-231">231</a></span>
<span class="normal"><a href="#__codelineno-0-232">232</a></span>
<span class="normal"><a href="#__codelineno-0-233">233</a></span>
<span class="normal"><a href="#__codelineno-0-234">234</a></span>
<span class="normal"><a href="#__codelineno-0-235">235</a></span></pre></div></td><td class="code"><div><pre><span></span><code><span id="__span-0-217"><a id="__codelineno-0-217" name="__codelineno-0-217"></a><span class="k">def</span><span class="w"> </span><span class="nf">spike_condition</span><span class="p">(</span>
</span><span id="__span-0-218"><a id="__codelineno-0-218" name="__codelineno-0-218"></a> <span class="bp">self</span><span class="p">,</span>
</span><span id="__span-0-219"><a id="__codelineno-0-219" name="__codelineno-0-219"></a> <span class="n">t</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
</span><span id="__span-0-220"><a id="__codelineno-0-220" name="__codelineno-0-220"></a> <span class="n">y</span><span class="p">:</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;neurons 3&quot;</span><span class="p">],</span>
</span><span id="__span-0-221"><a id="__codelineno-0-221" name="__codelineno-0-221"></a> <span class="o">**</span><span class="n">kwargs</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">],</span>
</span><span id="__span-0-222"><a id="__codelineno-0-222" name="__codelineno-0-222"></a><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot; neurons&quot;</span><span class="p">]:</span>
</span><span id="__span-0-223"><a id="__codelineno-0-223" name="__codelineno-0-223"></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;Compute spike condition for event detection.</span>
</span><span id="__span-0-224"><a id="__codelineno-0-224" name="__codelineno-0-224"></a>
</span><span id="__span-0-225"><a id="__codelineno-0-225" name="__codelineno-0-225"></a><span class="sd"> A spike is triggered when this function crosses zero (v &gt;= vthr).</span>
</span><span id="__span-0-226"><a id="__codelineno-0-226" name="__codelineno-0-226"></a>
</span><span id="__span-0-227"><a id="__codelineno-0-227" name="__codelineno-0-227"></a><span class="sd"> Args:</span>
</span><span id="__span-0-228"><a id="__codelineno-0-228" name="__codelineno-0-228"></a><span class="sd"> t: Current simulation time (unused but required by event detection).</span>
</span><span id="__span-0-229"><a id="__codelineno-0-229" name="__codelineno-0-229"></a><span class="sd"> y: State array of shape (neurons, 3) containing [v, v_slow, i_app].</span>
</span><span id="__span-0-230"><a id="__codelineno-0-230" name="__codelineno-0-230"></a><span class="sd"> **kwargs: Additional keyword arguments (unused).</span>
</span><span id="__span-0-231"><a id="__codelineno-0-231" name="__codelineno-0-231"></a>
</span><span id="__span-0-232"><a id="__codelineno-0-232" name="__codelineno-0-232"></a><span class="sd"> Returns:</span>
</span><span id="__span-0-233"><a id="__codelineno-0-233" name="__codelineno-0-233"></a><span class="sd"> Spike condition array of shape (neurons,). Positive values indicate v &gt; vthr.</span>
</span><span id="__span-0-234"><a id="__codelineno-0-234" name="__codelineno-0-234"></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="__span-0-235"><a id="__codelineno-0-235" name="__codelineno-0-235"></a> <span class="k">return</span> <span class="n">y</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">vthr</span>
</span></code></pre></div></td></tr></table></div>
</details>
</div>
</div>
</div>
</div>
</div>
<div class="doc doc-object doc-class">
<h4 id="felice.neuron_models.WereRabbit" class="doc doc-heading">
<code class="doc-symbol doc-symbol-heading doc-symbol-class"></code> <code>WereRabbit</code>
</h4>
<div class="doc doc-contents ">
<p class="doc doc-class-bases">
Bases: <code><span title="equinox.Module">Module</span></code></p>
<p>WereRabbit Neuron Model</p>
<p>The WereRabbit model implements a predator-prey dynamic with bistable
switching behavior controlled by a "moon phase" parameter <span class="arithmatex">\(z\)</span>.</p>
<p>The dynamics are governed by:</p>
<div class="arithmatex">\[
\begin{align}
z &amp;= tanh(\rho (u-v)) \\
\frac{du}{dt} &amp;= z - z \alpha e^{\beta v} [1 + \gamma (0.5 - u)] - \sigma \\
\frac{dv}{dt} &amp;= -z - z \alpha e^{\beta u} [1 + \gamma (0.5 - v)] - \sigma
\end{align}
\]</div>
<p>where <span class="arithmatex">\(z\)</span> represents the "moon phase" that switches the predator-prey roles.</p>
<p><span class="doc-section-title">Attributes:</span></p>
<table>
<thead>
<tr>
<th>Name</th>
<th>Type</th>
<th>Description</th>
</tr>
</thead>
<tbody>
<tr class="doc-section-item">
<td><code><span title="felice.neuron_models.WereRabbit.alpha">alpha</span></code></td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Current scaling parameter <span class="arithmatex">\(\alpha = I_{n0}/I_{bias}\)</span> (default: 0.0129)</p>
</div>
</td>
</tr>
<tr class="doc-section-item">
<td><code><span title="felice.neuron_models.WereRabbit.beta">beta</span></code></td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Exponential slope <span class="arithmatex">\(\beta = \kappa/U_t\)</span> (default: 15.6)</p>
</div>
</td>
</tr>
<tr class="doc-section-item">
<td><code><span title="felice.neuron_models.WereRabbit.gamma">gamma</span></code></td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Coupling parameter <span class="arithmatex">\(\gamma = 26e^{-2}\)</span></p>
</div>
</td>
</tr>
<tr class="doc-section-item">
<td><code><span title="felice.neuron_models.WereRabbit.rho">rho</span></code></td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Steepness of the tanh function <span class="arithmatex">\(\rho\)</span> (default: 5)</p>
</div>
</td>
</tr>
<tr class="doc-section-item">
<td><code><span title="felice.neuron_models.WereRabbit.sigma">sigma</span></code></td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Fixpoint distance scaling <span class="arithmatex">\(\sigma\)</span> (default: 0.6)</p>
</div>
</td>
</tr>
<tr class="doc-section-item">
<td><code><span title="felice.neuron_models.WereRabbit.rtol">rtol</span></code></td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Relative tolerance for the spiking fixpoint calculation.</p>
</div>
</td>
</tr>
<tr class="doc-section-item">
<td><code><span title="felice.neuron_models.WereRabbit.atol">atol</span></code></td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Absolute tolerance for the spiking fixpoint calculation.</p>
</div>
</td>
</tr>
<tr class="doc-section-item">
<td><code><span title="felice.neuron_models.WereRabbit.weight_u">weight_u</span></code></td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Input weight for the predator.</p>
</div>
</td>
</tr>
<tr class="doc-section-item">
<td><code><span title="felice.neuron_models.WereRabbit.weight_v">weight_v</span></code></td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Input weight for the prey.</p>
</div>
</td>
</tr>
</tbody>
</table>
<details class="mkdocstrings-source">
<summary>Source code in <code>felice/neuron_models/wererabbit.py</code></summary>
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal"><a href="#__codelineno-0-10"> 10</a></span>
<span class="normal"><a href="#__codelineno-0-11"> 11</a></span>
<span class="normal"><a href="#__codelineno-0-12"> 12</a></span>
<span class="normal"><a href="#__codelineno-0-13"> 13</a></span>
<span class="normal"><a href="#__codelineno-0-14"> 14</a></span>
<span class="normal"><a href="#__codelineno-0-15"> 15</a></span>
<span class="normal"><a href="#__codelineno-0-16"> 16</a></span>
<span class="normal"><a href="#__codelineno-0-17"> 17</a></span>
<span class="normal"><a href="#__codelineno-0-18"> 18</a></span>
<span class="normal"><a href="#__codelineno-0-19"> 19</a></span>
<span class="normal"><a href="#__codelineno-0-20"> 20</a></span>
<span class="normal"><a href="#__codelineno-0-21"> 21</a></span>
<span class="normal"><a href="#__codelineno-0-22"> 22</a></span>
<span class="normal"><a href="#__codelineno-0-23"> 23</a></span>
<span class="normal"><a href="#__codelineno-0-24"> 24</a></span>
<span class="normal"><a href="#__codelineno-0-25"> 25</a></span>
<span class="normal"><a href="#__codelineno-0-26"> 26</a></span>
<span class="normal"><a href="#__codelineno-0-27"> 27</a></span>
<span class="normal"><a href="#__codelineno-0-28"> 28</a></span>
<span class="normal"><a href="#__codelineno-0-29"> 29</a></span>
<span class="normal"><a href="#__codelineno-0-30"> 30</a></span>
<span class="normal"><a href="#__codelineno-0-31"> 31</a></span>
<span class="normal"><a href="#__codelineno-0-32"> 32</a></span>
<span class="normal"><a href="#__codelineno-0-33"> 33</a></span>
<span class="normal"><a href="#__codelineno-0-34"> 34</a></span>
<span class="normal"><a href="#__codelineno-0-35"> 35</a></span>
<span class="normal"><a href="#__codelineno-0-36"> 36</a></span>
<span class="normal"><a href="#__codelineno-0-37"> 37</a></span>
<span class="normal"><a href="#__codelineno-0-38"> 38</a></span>
<span class="normal"><a href="#__codelineno-0-39"> 39</a></span>
<span class="normal"><a href="#__codelineno-0-40"> 40</a></span>
<span class="normal"><a href="#__codelineno-0-41"> 41</a></span>
<span class="normal"><a href="#__codelineno-0-42"> 42</a></span>
<span class="normal"><a href="#__codelineno-0-43"> 43</a></span>
<span class="normal"><a href="#__codelineno-0-44"> 44</a></span>
<span class="normal"><a href="#__codelineno-0-45"> 45</a></span>
<span class="normal"><a href="#__codelineno-0-46"> 46</a></span>
<span class="normal"><a href="#__codelineno-0-47"> 47</a></span>
<span class="normal"><a href="#__codelineno-0-48"> 48</a></span>
<span class="normal"><a href="#__codelineno-0-49"> 49</a></span>
<span class="normal"><a href="#__codelineno-0-50"> 50</a></span>
<span class="normal"><a href="#__codelineno-0-51"> 51</a></span>
<span class="normal"><a href="#__codelineno-0-52"> 52</a></span>
<span class="normal"><a href="#__codelineno-0-53"> 53</a></span>
<span class="normal"><a href="#__codelineno-0-54"> 54</a></span>
<span class="normal"><a href="#__codelineno-0-55"> 55</a></span>
<span class="normal"><a href="#__codelineno-0-56"> 56</a></span>
<span class="normal"><a href="#__codelineno-0-57"> 57</a></span>
<span class="normal"><a href="#__codelineno-0-58"> 58</a></span>
<span class="normal"><a href="#__codelineno-0-59"> 59</a></span>
<span class="normal"><a href="#__codelineno-0-60"> 60</a></span>
<span class="normal"><a href="#__codelineno-0-61"> 61</a></span>
<span class="normal"><a href="#__codelineno-0-62"> 62</a></span>
<span class="normal"><a href="#__codelineno-0-63"> 63</a></span>
<span class="normal"><a href="#__codelineno-0-64"> 64</a></span>
<span class="normal"><a href="#__codelineno-0-65"> 65</a></span>
<span class="normal"><a href="#__codelineno-0-66"> 66</a></span>
<span class="normal"><a href="#__codelineno-0-67"> 67</a></span>
<span class="normal"><a href="#__codelineno-0-68"> 68</a></span>
<span class="normal"><a href="#__codelineno-0-69"> 69</a></span>
<span class="normal"><a href="#__codelineno-0-70"> 70</a></span>
<span class="normal"><a href="#__codelineno-0-71"> 71</a></span>
<span class="normal"><a href="#__codelineno-0-72"> 72</a></span>
<span class="normal"><a href="#__codelineno-0-73"> 73</a></span>
<span class="normal"><a href="#__codelineno-0-74"> 74</a></span>
<span class="normal"><a href="#__codelineno-0-75"> 75</a></span>
<span class="normal"><a href="#__codelineno-0-76"> 76</a></span>
<span class="normal"><a href="#__codelineno-0-77"> 77</a></span>
<span class="normal"><a href="#__codelineno-0-78"> 78</a></span>
<span class="normal"><a href="#__codelineno-0-79"> 79</a></span>
<span class="normal"><a href="#__codelineno-0-80"> 80</a></span>
<span class="normal"><a href="#__codelineno-0-81"> 81</a></span>
<span class="normal"><a href="#__codelineno-0-82"> 82</a></span>
<span class="normal"><a href="#__codelineno-0-83"> 83</a></span>
<span class="normal"><a href="#__codelineno-0-84"> 84</a></span>
<span class="normal"><a href="#__codelineno-0-85"> 85</a></span>
<span class="normal"><a href="#__codelineno-0-86"> 86</a></span>
<span class="normal"><a href="#__codelineno-0-87"> 87</a></span>
<span class="normal"><a href="#__codelineno-0-88"> 88</a></span>
<span class="normal"><a href="#__codelineno-0-89"> 89</a></span>
<span class="normal"><a href="#__codelineno-0-90"> 90</a></span>
<span class="normal"><a href="#__codelineno-0-91"> 91</a></span>
<span class="normal"><a href="#__codelineno-0-92"> 92</a></span>
<span class="normal"><a href="#__codelineno-0-93"> 93</a></span>
<span class="normal"><a href="#__codelineno-0-94"> 94</a></span>
<span class="normal"><a href="#__codelineno-0-95"> 95</a></span>
<span class="normal"><a href="#__codelineno-0-96"> 96</a></span>
<span class="normal"><a href="#__codelineno-0-97"> 97</a></span>
<span class="normal"><a href="#__codelineno-0-98"> 98</a></span>
<span class="normal"><a href="#__codelineno-0-99"> 99</a></span>
<span class="normal"><a href="#__codelineno-0-100">100</a></span>
<span class="normal"><a href="#__codelineno-0-101">101</a></span>
<span class="normal"><a href="#__codelineno-0-102">102</a></span>
<span class="normal"><a href="#__codelineno-0-103">103</a></span>
<span class="normal"><a href="#__codelineno-0-104">104</a></span>
<span class="normal"><a href="#__codelineno-0-105">105</a></span>
<span class="normal"><a href="#__codelineno-0-106">106</a></span>
<span class="normal"><a href="#__codelineno-0-107">107</a></span>
<span class="normal"><a href="#__codelineno-0-108">108</a></span>
<span class="normal"><a href="#__codelineno-0-109">109</a></span>
<span class="normal"><a href="#__codelineno-0-110">110</a></span>
<span class="normal"><a href="#__codelineno-0-111">111</a></span>
<span class="normal"><a href="#__codelineno-0-112">112</a></span>
<span class="normal"><a href="#__codelineno-0-113">113</a></span>
<span class="normal"><a href="#__codelineno-0-114">114</a></span>
<span class="normal"><a href="#__codelineno-0-115">115</a></span>
<span class="normal"><a href="#__codelineno-0-116">116</a></span>
<span class="normal"><a href="#__codelineno-0-117">117</a></span>
<span class="normal"><a href="#__codelineno-0-118">118</a></span>
<span class="normal"><a href="#__codelineno-0-119">119</a></span>
<span class="normal"><a href="#__codelineno-0-120">120</a></span>
<span class="normal"><a href="#__codelineno-0-121">121</a></span>
<span class="normal"><a href="#__codelineno-0-122">122</a></span>
<span class="normal"><a href="#__codelineno-0-123">123</a></span>
<span class="normal"><a href="#__codelineno-0-124">124</a></span>
<span class="normal"><a href="#__codelineno-0-125">125</a></span>
<span class="normal"><a href="#__codelineno-0-126">126</a></span>
<span class="normal"><a href="#__codelineno-0-127">127</a></span>
<span class="normal"><a href="#__codelineno-0-128">128</a></span>
<span class="normal"><a href="#__codelineno-0-129">129</a></span>
<span class="normal"><a href="#__codelineno-0-130">130</a></span>
<span class="normal"><a href="#__codelineno-0-131">131</a></span>
<span class="normal"><a href="#__codelineno-0-132">132</a></span>
<span class="normal"><a href="#__codelineno-0-133">133</a></span>
<span class="normal"><a href="#__codelineno-0-134">134</a></span>
<span class="normal"><a href="#__codelineno-0-135">135</a></span>
<span class="normal"><a href="#__codelineno-0-136">136</a></span>
<span class="normal"><a href="#__codelineno-0-137">137</a></span>
<span class="normal"><a href="#__codelineno-0-138">138</a></span>
<span class="normal"><a href="#__codelineno-0-139">139</a></span>
<span class="normal"><a href="#__codelineno-0-140">140</a></span>
<span class="normal"><a href="#__codelineno-0-141">141</a></span>
<span class="normal"><a href="#__codelineno-0-142">142</a></span>
<span class="normal"><a href="#__codelineno-0-143">143</a></span>
<span class="normal"><a href="#__codelineno-0-144">144</a></span>
<span class="normal"><a href="#__codelineno-0-145">145</a></span>
<span class="normal"><a href="#__codelineno-0-146">146</a></span>
<span class="normal"><a href="#__codelineno-0-147">147</a></span>
<span class="normal"><a href="#__codelineno-0-148">148</a></span>
<span class="normal"><a href="#__codelineno-0-149">149</a></span>
<span class="normal"><a href="#__codelineno-0-150">150</a></span>
<span class="normal"><a href="#__codelineno-0-151">151</a></span>
<span class="normal"><a href="#__codelineno-0-152">152</a></span>
<span class="normal"><a href="#__codelineno-0-153">153</a></span>
<span class="normal"><a href="#__codelineno-0-154">154</a></span>
<span class="normal"><a href="#__codelineno-0-155">155</a></span>
<span class="normal"><a href="#__codelineno-0-156">156</a></span>
<span class="normal"><a href="#__codelineno-0-157">157</a></span>
<span class="normal"><a href="#__codelineno-0-158">158</a></span>
<span class="normal"><a href="#__codelineno-0-159">159</a></span>
<span class="normal"><a href="#__codelineno-0-160">160</a></span>
<span class="normal"><a href="#__codelineno-0-161">161</a></span>
<span class="normal"><a href="#__codelineno-0-162">162</a></span>
<span class="normal"><a href="#__codelineno-0-163">163</a></span>
<span class="normal"><a href="#__codelineno-0-164">164</a></span>
<span class="normal"><a href="#__codelineno-0-165">165</a></span>
<span class="normal"><a href="#__codelineno-0-166">166</a></span>
<span class="normal"><a href="#__codelineno-0-167">167</a></span>
<span class="normal"><a href="#__codelineno-0-168">168</a></span>
<span class="normal"><a href="#__codelineno-0-169">169</a></span>
<span class="normal"><a href="#__codelineno-0-170">170</a></span>
<span class="normal"><a href="#__codelineno-0-171">171</a></span>
<span class="normal"><a href="#__codelineno-0-172">172</a></span>
<span class="normal"><a href="#__codelineno-0-173">173</a></span>
<span class="normal"><a href="#__codelineno-0-174">174</a></span>
<span class="normal"><a href="#__codelineno-0-175">175</a></span>
<span class="normal"><a href="#__codelineno-0-176">176</a></span>
<span class="normal"><a href="#__codelineno-0-177">177</a></span>
<span class="normal"><a href="#__codelineno-0-178">178</a></span>
<span class="normal"><a href="#__codelineno-0-179">179</a></span>
<span class="normal"><a href="#__codelineno-0-180">180</a></span>
<span class="normal"><a href="#__codelineno-0-181">181</a></span>
<span class="normal"><a href="#__codelineno-0-182">182</a></span>
<span class="normal"><a href="#__codelineno-0-183">183</a></span>
<span class="normal"><a href="#__codelineno-0-184">184</a></span>
<span class="normal"><a href="#__codelineno-0-185">185</a></span>
<span class="normal"><a href="#__codelineno-0-186">186</a></span>
<span class="normal"><a href="#__codelineno-0-187">187</a></span>
<span class="normal"><a href="#__codelineno-0-188">188</a></span>
<span class="normal"><a href="#__codelineno-0-189">189</a></span>
<span class="normal"><a href="#__codelineno-0-190">190</a></span>
<span class="normal"><a href="#__codelineno-0-191">191</a></span>
<span class="normal"><a href="#__codelineno-0-192">192</a></span>
<span class="normal"><a href="#__codelineno-0-193">193</a></span>
<span class="normal"><a href="#__codelineno-0-194">194</a></span></pre></div></td><td class="code"><div><pre><span></span><code><span id="__span-0-10"><a id="__codelineno-0-10" name="__codelineno-0-10"></a><span class="k">class</span><span class="w"> </span><span class="nc">WereRabbit</span><span class="p">(</span><span class="n">eqx</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
</span><span id="__span-0-11"><a id="__codelineno-0-11" name="__codelineno-0-11"></a><span class="w"> </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;</span>
</span><span id="__span-0-12"><a id="__codelineno-0-12" name="__codelineno-0-12"></a><span class="sd"> WereRabbit Neuron Model</span>
</span><span id="__span-0-13"><a id="__codelineno-0-13" name="__codelineno-0-13"></a>
</span><span id="__span-0-14"><a id="__codelineno-0-14" name="__codelineno-0-14"></a><span class="sd"> The WereRabbit model implements a predator-prey dynamic with bistable </span>
</span><span id="__span-0-15"><a id="__codelineno-0-15" name="__codelineno-0-15"></a><span class="sd"> switching behavior controlled by a &quot;moon phase&quot; parameter $z$.</span>
</span><span id="__span-0-16"><a id="__codelineno-0-16" name="__codelineno-0-16"></a>
</span><span id="__span-0-17"><a id="__codelineno-0-17" name="__codelineno-0-17"></a><span class="sd"> The dynamics are governed by:</span>
</span><span id="__span-0-18"><a id="__codelineno-0-18" name="__codelineno-0-18"></a>
</span><span id="__span-0-19"><a id="__codelineno-0-19" name="__codelineno-0-19"></a><span class="sd"> $$</span>
</span><span id="__span-0-20"><a id="__codelineno-0-20" name="__codelineno-0-20"></a><span class="sd"> \begin{align}</span>
</span><span id="__span-0-21"><a id="__codelineno-0-21" name="__codelineno-0-21"></a><span class="sd"> z &amp;= tanh(\rho (u-v)) \\</span>
</span><span id="__span-0-22"><a id="__codelineno-0-22" name="__codelineno-0-22"></a><span class="sd"> \frac{du}{dt} &amp;= z - z \alpha e^{\beta v} [1 + \gamma (0.5 - u)] - \sigma \\</span>
</span><span id="__span-0-23"><a id="__codelineno-0-23" name="__codelineno-0-23"></a><span class="sd"> \frac{dv}{dt} &amp;= -z - z \alpha e^{\beta u} [1 + \gamma (0.5 - v)] - \sigma</span>
</span><span id="__span-0-24"><a id="__codelineno-0-24" name="__codelineno-0-24"></a><span class="sd"> \end{align}</span>
</span><span id="__span-0-25"><a id="__codelineno-0-25" name="__codelineno-0-25"></a><span class="sd"> $$</span>
</span><span id="__span-0-26"><a id="__codelineno-0-26" name="__codelineno-0-26"></a>
</span><span id="__span-0-27"><a id="__codelineno-0-27" name="__codelineno-0-27"></a><span class="sd"> where $z$ represents the &quot;moon phase&quot; that switches the predator-prey roles.</span>
</span><span id="__span-0-28"><a id="__codelineno-0-28" name="__codelineno-0-28"></a>
</span><span id="__span-0-29"><a id="__codelineno-0-29" name="__codelineno-0-29"></a><span class="sd"> Attributes:</span>
</span><span id="__span-0-30"><a id="__codelineno-0-30" name="__codelineno-0-30"></a><span class="sd"> alpha: Current scaling parameter $\alpha = I_{n0}/I_{bias}$ (default: 0.0129)</span>
</span><span id="__span-0-31"><a id="__codelineno-0-31" name="__codelineno-0-31"></a><span class="sd"> beta: Exponential slope $\beta = \kappa/U_t$ (default: 15.6)</span>
</span><span id="__span-0-32"><a id="__codelineno-0-32" name="__codelineno-0-32"></a><span class="sd"> gamma: Coupling parameter $\gamma = 26e^{-2}$</span>
</span><span id="__span-0-33"><a id="__codelineno-0-33" name="__codelineno-0-33"></a><span class="sd"> rho: Steepness of the tanh function $\rho$ (default: 5)</span>
</span><span id="__span-0-34"><a id="__codelineno-0-34" name="__codelineno-0-34"></a><span class="sd"> sigma: Fixpoint distance scaling $\sigma$ (default: 0.6)</span>
</span><span id="__span-0-35"><a id="__codelineno-0-35" name="__codelineno-0-35"></a>
</span><span id="__span-0-36"><a id="__codelineno-0-36" name="__codelineno-0-36"></a><span class="sd"> rtol: Relative tolerance for the spiking fixpoint calculation.</span>
</span><span id="__span-0-37"><a id="__codelineno-0-37" name="__codelineno-0-37"></a><span class="sd"> atol: Absolute tolerance for the spiking fixpoint calculation.</span>
</span><span id="__span-0-38"><a id="__codelineno-0-38" name="__codelineno-0-38"></a>
</span><span id="__span-0-39"><a id="__codelineno-0-39" name="__codelineno-0-39"></a><span class="sd"> weight_u: Input weight for the predator.</span>
</span><span id="__span-0-40"><a id="__codelineno-0-40" name="__codelineno-0-40"></a><span class="sd"> weight_v: Input weight for the prey.</span>
</span><span id="__span-0-41"><a id="__codelineno-0-41" name="__codelineno-0-41"></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="__span-0-42"><a id="__codelineno-0-42" name="__codelineno-0-42"></a>
</span><span id="__span-0-43"><a id="__codelineno-0-43" name="__codelineno-0-43"></a> <span class="n">dtype</span><span class="p">:</span> <span class="n">DTypeLike</span> <span class="o">=</span> <span class="n">eqx</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="n">static</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
</span><span id="__span-0-44"><a id="__codelineno-0-44" name="__codelineno-0-44"></a> <span class="n">rtol</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="n">eqx</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="n">static</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
</span><span id="__span-0-45"><a id="__codelineno-0-45" name="__codelineno-0-45"></a> <span class="n">atol</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="n">eqx</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="n">static</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
</span><span id="__span-0-46"><a id="__codelineno-0-46" name="__codelineno-0-46"></a>
</span><span id="__span-0-47"><a id="__codelineno-0-47" name="__codelineno-0-47"></a> <span class="n">alpha</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="n">eqx</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="n">static</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span> <span class="c1"># I_n0 / I_bias ratio</span>
</span><span id="__span-0-48"><a id="__codelineno-0-48" name="__codelineno-0-48"></a> <span class="n">beta</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="n">eqx</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="n">static</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span> <span class="c1"># k / U_t (inverse thermal scale)</span>
</span><span id="__span-0-49"><a id="__codelineno-0-49" name="__codelineno-0-49"></a> <span class="n">gamma</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="n">eqx</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="n">static</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span> <span class="c1"># coupling coefficient</span>
</span><span id="__span-0-50"><a id="__codelineno-0-50" name="__codelineno-0-50"></a> <span class="n">rho</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="n">eqx</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="n">static</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span> <span class="c1"># tanh steepness</span>
</span><span id="__span-0-51"><a id="__codelineno-0-51" name="__codelineno-0-51"></a> <span class="n">sigma</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="n">eqx</span><span class="o">.</span><span class="n">field</span><span class="p">(</span><span class="n">static</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span> <span class="c1"># bias scaling (s * I_bias)</span>
</span><span id="__span-0-52"><a id="__codelineno-0-52" name="__codelineno-0-52"></a>
</span><span id="__span-0-53"><a id="__codelineno-0-53" name="__codelineno-0-53"></a> <span class="k">def</span><span class="w"> </span><span class="fm">__init__</span><span class="p">(</span>
</span><span id="__span-0-54"><a id="__codelineno-0-54" name="__codelineno-0-54"></a> <span class="bp">self</span><span class="p">,</span>
</span><span id="__span-0-55"><a id="__codelineno-0-55" name="__codelineno-0-55"></a> <span class="o">*</span><span class="p">,</span>
</span><span id="__span-0-56"><a id="__codelineno-0-56" name="__codelineno-0-56"></a> <span class="n">atol</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1e-3</span><span class="p">,</span>
</span><span id="__span-0-57"><a id="__codelineno-0-57" name="__codelineno-0-57"></a> <span class="n">rtol</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1e-3</span><span class="p">,</span>
</span><span id="__span-0-58"><a id="__codelineno-0-58" name="__codelineno-0-58"></a> <span class="n">alpha</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.0129</span><span class="p">,</span>
</span><span id="__span-0-59"><a id="__codelineno-0-59" name="__codelineno-0-59"></a> <span class="n">beta</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">15.6</span><span class="p">,</span>
</span><span id="__span-0-60"><a id="__codelineno-0-60" name="__codelineno-0-60"></a> <span class="n">gamma</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.26</span><span class="p">,</span>
</span><span id="__span-0-61"><a id="__codelineno-0-61" name="__codelineno-0-61"></a> <span class="n">rho</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">5.0</span><span class="p">,</span>
</span><span id="__span-0-62"><a id="__codelineno-0-62" name="__codelineno-0-62"></a> <span class="n">sigma</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.6</span><span class="p">,</span>
</span><span id="__span-0-63"><a id="__codelineno-0-63" name="__codelineno-0-63"></a> <span class="n">dtype</span><span class="p">:</span> <span class="n">DTypeLike</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">float32</span><span class="p">,</span>
</span><span id="__span-0-64"><a id="__codelineno-0-64" name="__codelineno-0-64"></a> <span class="p">):</span>
</span><span id="__span-0-65"><a id="__codelineno-0-65" name="__codelineno-0-65"></a><span class="w"> </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;Initialize the WereRabbit neuron model.</span>
</span><span id="__span-0-66"><a id="__codelineno-0-66" name="__codelineno-0-66"></a>
</span><span id="__span-0-67"><a id="__codelineno-0-67" name="__codelineno-0-67"></a><span class="sd"> Args:</span>
</span><span id="__span-0-68"><a id="__codelineno-0-68" name="__codelineno-0-68"></a><span class="sd"> rtol: Relative tolerance for the spiking fixpoint calculation.</span>
</span><span id="__span-0-69"><a id="__codelineno-0-69" name="__codelineno-0-69"></a><span class="sd"> atol: Absolute tolerance for the spiking fixpoint calculation.</span>
</span><span id="__span-0-70"><a id="__codelineno-0-70" name="__codelineno-0-70"></a><span class="sd"> alpha: Current scaling parameter $\alpha = I_{n0}/I_{bias}$ (default: 0.0129)</span>
</span><span id="__span-0-71"><a id="__codelineno-0-71" name="__codelineno-0-71"></a><span class="sd"> beta: Exponential slope $\beta = \kappa/U_t$ (default: 15.6)</span>
</span><span id="__span-0-72"><a id="__codelineno-0-72" name="__codelineno-0-72"></a><span class="sd"> gamma: Coupling parameter $\gamma = 26e^{-2}$</span>
</span><span id="__span-0-73"><a id="__codelineno-0-73" name="__codelineno-0-73"></a><span class="sd"> rho: Steepness of the tanh function $\rho$ (default: 5)</span>
</span><span id="__span-0-74"><a id="__codelineno-0-74" name="__codelineno-0-74"></a><span class="sd"> sigma: Fixpoint distance scaling $\sigma$ (default: 0.6)</span>
</span><span id="__span-0-75"><a id="__codelineno-0-75" name="__codelineno-0-75"></a><span class="sd"> dtype: Data type for arrays (default: float32).</span>
</span><span id="__span-0-76"><a id="__codelineno-0-76" name="__codelineno-0-76"></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="__span-0-77"><a id="__codelineno-0-77" name="__codelineno-0-77"></a> <span class="bp">self</span><span class="o">.</span><span class="n">dtype</span> <span class="o">=</span> <span class="n">dtype</span>
</span><span id="__span-0-78"><a id="__codelineno-0-78" name="__codelineno-0-78"></a> <span class="bp">self</span><span class="o">.</span><span class="n">alpha</span> <span class="o">=</span> <span class="n">alpha</span>
</span><span id="__span-0-79"><a id="__codelineno-0-79" name="__codelineno-0-79"></a> <span class="bp">self</span><span class="o">.</span><span class="n">beta</span> <span class="o">=</span> <span class="n">beta</span>
</span><span id="__span-0-80"><a id="__codelineno-0-80" name="__codelineno-0-80"></a> <span class="bp">self</span><span class="o">.</span><span class="n">gamma</span> <span class="o">=</span> <span class="n">gamma</span>
</span><span id="__span-0-81"><a id="__codelineno-0-81" name="__codelineno-0-81"></a> <span class="bp">self</span><span class="o">.</span><span class="n">rho</span> <span class="o">=</span> <span class="n">rho</span>
</span><span id="__span-0-82"><a id="__codelineno-0-82" name="__codelineno-0-82"></a> <span class="bp">self</span><span class="o">.</span><span class="n">sigma</span> <span class="o">=</span> <span class="n">sigma</span>
</span><span id="__span-0-83"><a id="__codelineno-0-83" name="__codelineno-0-83"></a>
</span><span id="__span-0-84"><a id="__codelineno-0-84" name="__codelineno-0-84"></a> <span class="bp">self</span><span class="o">.</span><span class="n">rtol</span> <span class="o">=</span> <span class="n">rtol</span>
</span><span id="__span-0-85"><a id="__codelineno-0-85" name="__codelineno-0-85"></a> <span class="bp">self</span><span class="o">.</span><span class="n">atol</span> <span class="o">=</span> <span class="n">atol</span>
</span><span id="__span-0-86"><a id="__codelineno-0-86" name="__codelineno-0-86"></a>
</span><span id="__span-0-87"><a id="__codelineno-0-87" name="__codelineno-0-87"></a> <span class="k">def</span><span class="w"> </span><span class="nf">init_state</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">n_neurons</span><span class="p">:</span> <span class="nb">int</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;neurons 2&quot;</span><span class="p">]:</span>
</span><span id="__span-0-88"><a id="__codelineno-0-88" name="__codelineno-0-88"></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;Initialize the neuron state variables.</span>
</span><span id="__span-0-89"><a id="__codelineno-0-89" name="__codelineno-0-89"></a>
</span><span id="__span-0-90"><a id="__codelineno-0-90" name="__codelineno-0-90"></a><span class="sd"> Args:</span>
</span><span id="__span-0-91"><a id="__codelineno-0-91" name="__codelineno-0-91"></a><span class="sd"> n_neurons: Number of neurons to initialize.</span>
</span><span id="__span-0-92"><a id="__codelineno-0-92" name="__codelineno-0-92"></a>
</span><span id="__span-0-93"><a id="__codelineno-0-93" name="__codelineno-0-93"></a><span class="sd"> Returns:</span>
</span><span id="__span-0-94"><a id="__codelineno-0-94" name="__codelineno-0-94"></a><span class="sd"> Initial state array of shape (neurons, 3) containing [u, v, has_spiked],</span>
</span><span id="__span-0-95"><a id="__codelineno-0-95" name="__codelineno-0-95"></a><span class="sd"> where u and v are the predator/prey membrane voltages, has_spiked is a</span>
</span><span id="__span-0-96"><a id="__codelineno-0-96" name="__codelineno-0-96"></a><span class="sd"> variable that is 1 whenever the neuron spike and 0 otherwise .</span>
</span><span id="__span-0-97"><a id="__codelineno-0-97" name="__codelineno-0-97"></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="__span-0-98"><a id="__codelineno-0-98" name="__codelineno-0-98"></a> <span class="n">x1</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">n_neurons</span><span class="p">,),</span> <span class="n">dtype</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dtype</span><span class="p">)</span>
</span><span id="__span-0-99"><a id="__codelineno-0-99" name="__codelineno-0-99"></a> <span class="n">x2</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">n_neurons</span><span class="p">,),</span> <span class="n">dtype</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dtype</span><span class="p">)</span>
</span><span id="__span-0-100"><a id="__codelineno-0-100" name="__codelineno-0-100"></a> <span class="k">return</span> <span class="n">jnp</span><span class="o">.</span><span class="n">stack</span><span class="p">([</span><span class="n">x1</span><span class="p">,</span> <span class="n">x2</span><span class="p">],</span> <span class="n">axis</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
</span><span id="__span-0-101"><a id="__codelineno-0-101" name="__codelineno-0-101"></a>
</span><span id="__span-0-102"><a id="__codelineno-0-102" name="__codelineno-0-102"></a> <span class="k">def</span><span class="w"> </span><span class="nf">vector_field</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">y</span><span class="p">:</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;neurons 2&quot;</span><span class="p">])</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;neurons 2&quot;</span><span class="p">]:</span>
</span><span id="__span-0-103"><a id="__codelineno-0-103" name="__codelineno-0-103"></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;Compute vector field of the neuron state variables.</span>
</span><span id="__span-0-104"><a id="__codelineno-0-104" name="__codelineno-0-104"></a>
</span><span id="__span-0-105"><a id="__codelineno-0-105" name="__codelineno-0-105"></a><span class="sd"> This implements the WereRabbit dynamics</span>
</span><span id="__span-0-106"><a id="__codelineno-0-106" name="__codelineno-0-106"></a>
</span><span id="__span-0-107"><a id="__codelineno-0-107" name="__codelineno-0-107"></a><span class="sd"> - du/dt: Predator dynamics</span>
</span><span id="__span-0-108"><a id="__codelineno-0-108" name="__codelineno-0-108"></a><span class="sd"> - dv/dt: WerePrey dynamics</span>
</span><span id="__span-0-109"><a id="__codelineno-0-109" name="__codelineno-0-109"></a>
</span><span id="__span-0-110"><a id="__codelineno-0-110" name="__codelineno-0-110"></a><span class="sd"> Args:</span>
</span><span id="__span-0-111"><a id="__codelineno-0-111" name="__codelineno-0-111"></a><span class="sd"> y: State array of shape (neurons, 2) containing [u, v].</span>
</span><span id="__span-0-112"><a id="__codelineno-0-112" name="__codelineno-0-112"></a>
</span><span id="__span-0-113"><a id="__codelineno-0-113" name="__codelineno-0-113"></a><span class="sd"> Returns:</span>
</span><span id="__span-0-114"><a id="__codelineno-0-114" name="__codelineno-0-114"></a><span class="sd"> Time derivatives of shape (neurons, 2) containing [du/dt, dv/dt].</span>
</span><span id="__span-0-115"><a id="__codelineno-0-115" name="__codelineno-0-115"></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="__span-0-116"><a id="__codelineno-0-116" name="__codelineno-0-116"></a> <span class="n">u</span> <span class="o">=</span> <span class="n">y</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span>
</span><span id="__span-0-117"><a id="__codelineno-0-117" name="__codelineno-0-117"></a> <span class="n">v</span> <span class="o">=</span> <span class="n">y</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">]</span>
</span><span id="__span-0-118"><a id="__codelineno-0-118" name="__codelineno-0-118"></a>
</span><span id="__span-0-119"><a id="__codelineno-0-119" name="__codelineno-0-119"></a> <span class="n">z</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">tanh</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">rho</span> <span class="o">*</span> <span class="p">(</span><span class="n">u</span> <span class="o">-</span> <span class="n">v</span><span class="p">))</span>
</span><span id="__span-0-120"><a id="__codelineno-0-120" name="__codelineno-0-120"></a> <span class="n">du</span> <span class="o">=</span> <span class="p">(</span>
</span><span id="__span-0-121"><a id="__codelineno-0-121" name="__codelineno-0-121"></a> <span class="n">z</span> <span class="o">*</span> <span class="p">(</span><span class="mi">1</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">alpha</span> <span class="o">*</span> <span class="n">jnp</span><span class="o">.</span><span class="n">exp</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">beta</span> <span class="o">*</span> <span class="n">v</span><span class="p">)</span> <span class="o">*</span> <span class="p">(</span><span class="mi">1</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">gamma</span> <span class="o">*</span> <span class="p">(</span><span class="mf">0.5</span> <span class="o">-</span> <span class="n">u</span><span class="p">)))</span>
</span><span id="__span-0-122"><a id="__codelineno-0-122" name="__codelineno-0-122"></a> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">sigma</span>
</span><span id="__span-0-123"><a id="__codelineno-0-123" name="__codelineno-0-123"></a> <span class="p">)</span>
</span><span id="__span-0-124"><a id="__codelineno-0-124" name="__codelineno-0-124"></a> <span class="n">dv</span> <span class="o">=</span> <span class="p">(</span>
</span><span id="__span-0-125"><a id="__codelineno-0-125" name="__codelineno-0-125"></a> <span class="n">z</span>
</span><span id="__span-0-126"><a id="__codelineno-0-126" name="__codelineno-0-126"></a> <span class="o">*</span> <span class="p">(</span><span class="o">-</span><span class="mi">1</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">alpha</span> <span class="o">*</span> <span class="n">jnp</span><span class="o">.</span><span class="n">exp</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">beta</span> <span class="o">*</span> <span class="n">u</span><span class="p">)</span> <span class="o">*</span> <span class="p">(</span><span class="mi">1</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">gamma</span> <span class="o">*</span> <span class="p">(</span><span class="mf">0.5</span> <span class="o">-</span> <span class="n">v</span><span class="p">)))</span>
</span><span id="__span-0-127"><a id="__codelineno-0-127" name="__codelineno-0-127"></a> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">sigma</span>
</span><span id="__span-0-128"><a id="__codelineno-0-128" name="__codelineno-0-128"></a> <span class="p">)</span>
</span><span id="__span-0-129"><a id="__codelineno-0-129" name="__codelineno-0-129"></a>
</span><span id="__span-0-130"><a id="__codelineno-0-130" name="__codelineno-0-130"></a> <span class="n">dv</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">where</span><span class="p">(</span><span class="n">jnp</span><span class="o">.</span><span class="n">allclose</span><span class="p">(</span><span class="n">z</span><span class="p">,</span> <span class="mf">0.0</span><span class="p">),</span> <span class="n">dv</span> <span class="o">*</span> <span class="n">jnp</span><span class="o">.</span><span class="n">sign</span><span class="p">(</span><span class="n">v</span><span class="p">),</span> <span class="n">dv</span><span class="p">)</span>
</span><span id="__span-0-131"><a id="__codelineno-0-131" name="__codelineno-0-131"></a> <span class="n">du</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">where</span><span class="p">(</span><span class="n">jnp</span><span class="o">.</span><span class="n">allclose</span><span class="p">(</span><span class="n">z</span><span class="p">,</span> <span class="mf">0.0</span><span class="p">),</span> <span class="n">du</span> <span class="o">*</span> <span class="n">jnp</span><span class="o">.</span><span class="n">sign</span><span class="p">(</span><span class="n">u</span><span class="p">),</span> <span class="n">du</span><span class="p">)</span>
</span><span id="__span-0-132"><a id="__codelineno-0-132" name="__codelineno-0-132"></a>
</span><span id="__span-0-133"><a id="__codelineno-0-133" name="__codelineno-0-133"></a> <span class="k">return</span> <span class="n">jnp</span><span class="o">.</span><span class="n">stack</span><span class="p">([</span><span class="n">du</span><span class="p">,</span> <span class="n">dv</span><span class="p">],</span> <span class="n">axis</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
</span><span id="__span-0-134"><a id="__codelineno-0-134" name="__codelineno-0-134"></a>
</span><span id="__span-0-135"><a id="__codelineno-0-135" name="__codelineno-0-135"></a> <span class="k">def</span><span class="w"> </span><span class="nf">dynamics</span><span class="p">(</span>
</span><span id="__span-0-136"><a id="__codelineno-0-136" name="__codelineno-0-136"></a> <span class="bp">self</span><span class="p">,</span>
</span><span id="__span-0-137"><a id="__codelineno-0-137" name="__codelineno-0-137"></a> <span class="n">t</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
</span><span id="__span-0-138"><a id="__codelineno-0-138" name="__codelineno-0-138"></a> <span class="n">y</span><span class="p">:</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;neurons 2&quot;</span><span class="p">],</span>
</span><span id="__span-0-139"><a id="__codelineno-0-139" name="__codelineno-0-139"></a> <span class="n">args</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">],</span>
</span><span id="__span-0-140"><a id="__codelineno-0-140" name="__codelineno-0-140"></a> <span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;neurons 2&quot;</span><span class="p">]:</span>
</span><span id="__span-0-141"><a id="__codelineno-0-141" name="__codelineno-0-141"></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;Compute time derivatives of the neuron state variables.</span>
</span><span id="__span-0-142"><a id="__codelineno-0-142" name="__codelineno-0-142"></a>
</span><span id="__span-0-143"><a id="__codelineno-0-143" name="__codelineno-0-143"></a><span class="sd"> This implements the WereRabbit dynamics</span>
</span><span id="__span-0-144"><a id="__codelineno-0-144" name="__codelineno-0-144"></a>
</span><span id="__span-0-145"><a id="__codelineno-0-145" name="__codelineno-0-145"></a><span class="sd"> - du/dt: Predator dynamics</span>
</span><span id="__span-0-146"><a id="__codelineno-0-146" name="__codelineno-0-146"></a><span class="sd"> - dv/dt: WerePrey dynamics</span>
</span><span id="__span-0-147"><a id="__codelineno-0-147" name="__codelineno-0-147"></a>
</span><span id="__span-0-148"><a id="__codelineno-0-148" name="__codelineno-0-148"></a><span class="sd"> Args:</span>
</span><span id="__span-0-149"><a id="__codelineno-0-149" name="__codelineno-0-149"></a><span class="sd"> t: Current simulation time (unused but required by framework).</span>
</span><span id="__span-0-150"><a id="__codelineno-0-150" name="__codelineno-0-150"></a><span class="sd"> y: State array of shape (neurons, 3) containing [u, v, has_spiked].</span>
</span><span id="__span-0-151"><a id="__codelineno-0-151" name="__codelineno-0-151"></a><span class="sd"> args: Additional arguments (unused but required by framework).</span>
</span><span id="__span-0-152"><a id="__codelineno-0-152" name="__codelineno-0-152"></a>
</span><span id="__span-0-153"><a id="__codelineno-0-153" name="__codelineno-0-153"></a><span class="sd"> Returns:</span>
</span><span id="__span-0-154"><a id="__codelineno-0-154" name="__codelineno-0-154"></a><span class="sd"> Time derivatives of shape (neurons, 3) containing [du/dt, dv/dt, 0].</span>
</span><span id="__span-0-155"><a id="__codelineno-0-155" name="__codelineno-0-155"></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="__span-0-156"><a id="__codelineno-0-156" name="__codelineno-0-156"></a> <span class="n">dxdt</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">vector_field</span><span class="p">(</span><span class="n">y</span><span class="p">)</span>
</span><span id="__span-0-157"><a id="__codelineno-0-157" name="__codelineno-0-157"></a>
</span><span id="__span-0-158"><a id="__codelineno-0-158" name="__codelineno-0-158"></a> <span class="k">return</span> <span class="n">dxdt</span>
</span><span id="__span-0-159"><a id="__codelineno-0-159" name="__codelineno-0-159"></a>
</span><span id="__span-0-160"><a id="__codelineno-0-160" name="__codelineno-0-160"></a> <span class="k">def</span><span class="w"> </span><span class="nf">spike_condition</span><span class="p">(</span>
</span><span id="__span-0-161"><a id="__codelineno-0-161" name="__codelineno-0-161"></a> <span class="bp">self</span><span class="p">,</span>
</span><span id="__span-0-162"><a id="__codelineno-0-162" name="__codelineno-0-162"></a> <span class="n">t</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
</span><span id="__span-0-163"><a id="__codelineno-0-163" name="__codelineno-0-163"></a> <span class="n">y</span><span class="p">:</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;neurons 2&quot;</span><span class="p">],</span>
</span><span id="__span-0-164"><a id="__codelineno-0-164" name="__codelineno-0-164"></a> <span class="o">**</span><span class="n">kwargs</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">],</span>
</span><span id="__span-0-165"><a id="__codelineno-0-165" name="__codelineno-0-165"></a> <span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot; neurons&quot;</span><span class="p">]:</span>
</span><span id="__span-0-166"><a id="__codelineno-0-166" name="__codelineno-0-166"></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;Compute spike condition for event detection.</span>
</span><span id="__span-0-167"><a id="__codelineno-0-167" name="__codelineno-0-167"></a>
</span><span id="__span-0-168"><a id="__codelineno-0-168" name="__codelineno-0-168"></a><span class="sd"> A spike is triggered when the system reach to a fixpoint.</span>
</span><span id="__span-0-169"><a id="__codelineno-0-169" name="__codelineno-0-169"></a>
</span><span id="__span-0-170"><a id="__codelineno-0-170" name="__codelineno-0-170"></a><span class="sd"> INFO:</span>
</span><span id="__span-0-171"><a id="__codelineno-0-171" name="__codelineno-0-171"></a><span class="sd"> `has_spiked` is use to the system don&#39;t detect a continuos</span>
</span><span id="__span-0-172"><a id="__codelineno-0-172" name="__codelineno-0-172"></a><span class="sd"> spike when reach a fixpoint.</span>
</span><span id="__span-0-173"><a id="__codelineno-0-173" name="__codelineno-0-173"></a>
</span><span id="__span-0-174"><a id="__codelineno-0-174" name="__codelineno-0-174"></a><span class="sd"> Args:</span>
</span><span id="__span-0-175"><a id="__codelineno-0-175" name="__codelineno-0-175"></a><span class="sd"> t: Current simulation time (unused but required by the framework).</span>
</span><span id="__span-0-176"><a id="__codelineno-0-176" name="__codelineno-0-176"></a><span class="sd"> y: State array of shape (neurons, 3) containing [u, v, has_spiked].</span>
</span><span id="__span-0-177"><a id="__codelineno-0-177" name="__codelineno-0-177"></a><span class="sd"> **kwargs: Additional keyword arguments (unused).</span>
</span><span id="__span-0-178"><a id="__codelineno-0-178" name="__codelineno-0-178"></a>
</span><span id="__span-0-179"><a id="__codelineno-0-179" name="__codelineno-0-179"></a><span class="sd"> Returns:</span>
</span><span id="__span-0-180"><a id="__codelineno-0-180" name="__codelineno-0-180"></a><span class="sd"> Spike condition array of shape (neurons,). Positive values indicate spike.</span>
</span><span id="__span-0-181"><a id="__codelineno-0-181" name="__codelineno-0-181"></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="__span-0-182"><a id="__codelineno-0-182" name="__codelineno-0-182"></a> <span class="n">_atol</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">atol</span>
</span><span id="__span-0-183"><a id="__codelineno-0-183" name="__codelineno-0-183"></a> <span class="n">_rtol</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">rtol</span>
</span><span id="__span-0-184"><a id="__codelineno-0-184" name="__codelineno-0-184"></a> <span class="n">_norm</span> <span class="o">=</span> <span class="n">optx</span><span class="o">.</span><span class="n">rms_norm</span>
</span><span id="__span-0-185"><a id="__codelineno-0-185" name="__codelineno-0-185"></a>
</span><span id="__span-0-186"><a id="__codelineno-0-186" name="__codelineno-0-186"></a> <span class="n">vf</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">dynamics</span><span class="p">(</span><span class="n">t</span><span class="p">,</span> <span class="n">y</span><span class="p">,</span> <span class="p">{})</span>
</span><span id="__span-0-187"><a id="__codelineno-0-187" name="__codelineno-0-187"></a>
</span><span id="__span-0-188"><a id="__codelineno-0-188" name="__codelineno-0-188"></a> <span class="nd">@jax</span><span class="o">.</span><span class="n">vmap</span>
</span><span id="__span-0-189"><a id="__codelineno-0-189" name="__codelineno-0-189"></a> <span class="k">def</span><span class="w"> </span><span class="nf">calculate_norm</span><span class="p">(</span><span class="n">vf</span><span class="p">,</span> <span class="n">y</span><span class="p">):</span>
</span><span id="__span-0-190"><a id="__codelineno-0-190" name="__codelineno-0-190"></a> <span class="k">return</span> <span class="n">_atol</span> <span class="o">+</span> <span class="n">_rtol</span> <span class="o">*</span> <span class="n">_norm</span><span class="p">(</span><span class="n">y</span><span class="p">[:</span><span class="o">-</span><span class="mi">1</span><span class="p">])</span> <span class="o">-</span> <span class="n">_norm</span><span class="p">(</span><span class="n">vf</span><span class="p">[:</span><span class="o">-</span><span class="mi">1</span><span class="p">])</span>
</span><span id="__span-0-191"><a id="__codelineno-0-191" name="__codelineno-0-191"></a>
</span><span id="__span-0-192"><a id="__codelineno-0-192" name="__codelineno-0-192"></a> <span class="n">base_cond</span> <span class="o">=</span> <span class="n">calculate_norm</span><span class="p">(</span><span class="n">vf</span><span class="p">,</span> <span class="n">y</span><span class="p">)</span>
</span><span id="__span-0-193"><a id="__codelineno-0-193" name="__codelineno-0-193"></a>
</span><span id="__span-0-194"><a id="__codelineno-0-194" name="__codelineno-0-194"></a> <span class="k">return</span> <span class="n">base_cond</span>
</span></code></pre></div></td></tr></table></div>
</details>
<div class="doc doc-children">
<h5 id="felice.neuron_models.WereRabbit-functions">Functions</h5>
<div class="doc doc-object doc-function">
<h6 id="felice.neuron_models.WereRabbit.__init__" class="doc doc-heading">
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <code class="highlight language-python"><span class="fm">__init__</span><span class="p">(</span><span class="o">*</span><span class="p">,</span> <span class="n">atol</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.001</span><span class="p">,</span> <span class="n">rtol</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.001</span><span class="p">,</span> <span class="n">alpha</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.0129</span><span class="p">,</span> <span class="n">beta</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">15.6</span><span class="p">,</span> <span class="n">gamma</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.26</span><span class="p">,</span> <span class="n">rho</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">5.0</span><span class="p">,</span> <span class="n">sigma</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.6</span><span class="p">,</span> <span class="n">dtype</span><span class="p">:</span> <span class="n">DTypeLike</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">float32</span><span class="p">)</span></code>
</h6>
<div class="doc doc-contents ">
<p>Initialize the WereRabbit neuron model.</p>
<p><span class="doc-section-title">Parameters:</span></p>
<table>
<thead>
<tr>
<th>Name</th>
<th>Type</th>
<th>Description</th>
<th>Default</th>
</tr>
</thead>
<tbody>
<tr class="doc-section-item">
<td>
<code>rtol</code>
</td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Relative tolerance for the spiking fixpoint calculation.</p>
</div>
</td>
<td>
<code>0.001</code>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>atol</code>
</td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Absolute tolerance for the spiking fixpoint calculation.</p>
</div>
</td>
<td>
<code>0.001</code>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>alpha</code>
</td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Current scaling parameter <span class="arithmatex">\(\alpha = I_{n0}/I_{bias}\)</span> (default: 0.0129)</p>
</div>
</td>
<td>
<code>0.0129</code>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>beta</code>
</td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Exponential slope <span class="arithmatex">\(\beta = \kappa/U_t\)</span> (default: 15.6)</p>
</div>
</td>
<td>
<code>15.6</code>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>gamma</code>
</td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Coupling parameter <span class="arithmatex">\(\gamma = 26e^{-2}\)</span></p>
</div>
</td>
<td>
<code>0.26</code>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>rho</code>
</td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Steepness of the tanh function <span class="arithmatex">\(\rho\)</span> (default: 5)</p>
</div>
</td>
<td>
<code>5.0</code>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>sigma</code>
</td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Fixpoint distance scaling <span class="arithmatex">\(\sigma\)</span> (default: 0.6)</p>
</div>
</td>
<td>
<code>0.6</code>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>dtype</code>
</td>
<td>
<code><span title="jaxtyping.DTypeLike">DTypeLike</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Data type for arrays (default: float32).</p>
</div>
</td>
<td>
<code><span title="jax.numpy.float32">float32</span></code>
</td>
</tr>
</tbody>
</table>
<details class="mkdocstrings-source">
<summary>Source code in <code>felice/neuron_models/wererabbit.py</code></summary>
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal"><a href="#__codelineno-0-53">53</a></span>
<span class="normal"><a href="#__codelineno-0-54">54</a></span>
<span class="normal"><a href="#__codelineno-0-55">55</a></span>
<span class="normal"><a href="#__codelineno-0-56">56</a></span>
<span class="normal"><a href="#__codelineno-0-57">57</a></span>
<span class="normal"><a href="#__codelineno-0-58">58</a></span>
<span class="normal"><a href="#__codelineno-0-59">59</a></span>
<span class="normal"><a href="#__codelineno-0-60">60</a></span>
<span class="normal"><a href="#__codelineno-0-61">61</a></span>
<span class="normal"><a href="#__codelineno-0-62">62</a></span>
<span class="normal"><a href="#__codelineno-0-63">63</a></span>
<span class="normal"><a href="#__codelineno-0-64">64</a></span>
<span class="normal"><a href="#__codelineno-0-65">65</a></span>
<span class="normal"><a href="#__codelineno-0-66">66</a></span>
<span class="normal"><a href="#__codelineno-0-67">67</a></span>
<span class="normal"><a href="#__codelineno-0-68">68</a></span>
<span class="normal"><a href="#__codelineno-0-69">69</a></span>
<span class="normal"><a href="#__codelineno-0-70">70</a></span>
<span class="normal"><a href="#__codelineno-0-71">71</a></span>
<span class="normal"><a href="#__codelineno-0-72">72</a></span>
<span class="normal"><a href="#__codelineno-0-73">73</a></span>
<span class="normal"><a href="#__codelineno-0-74">74</a></span>
<span class="normal"><a href="#__codelineno-0-75">75</a></span>
<span class="normal"><a href="#__codelineno-0-76">76</a></span>
<span class="normal"><a href="#__codelineno-0-77">77</a></span>
<span class="normal"><a href="#__codelineno-0-78">78</a></span>
<span class="normal"><a href="#__codelineno-0-79">79</a></span>
<span class="normal"><a href="#__codelineno-0-80">80</a></span>
<span class="normal"><a href="#__codelineno-0-81">81</a></span>
<span class="normal"><a href="#__codelineno-0-82">82</a></span>
<span class="normal"><a href="#__codelineno-0-83">83</a></span>
<span class="normal"><a href="#__codelineno-0-84">84</a></span>
<span class="normal"><a href="#__codelineno-0-85">85</a></span></pre></div></td><td class="code"><div><pre><span></span><code><span id="__span-0-53"><a id="__codelineno-0-53" name="__codelineno-0-53"></a><span class="k">def</span><span class="w"> </span><span class="fm">__init__</span><span class="p">(</span>
</span><span id="__span-0-54"><a id="__codelineno-0-54" name="__codelineno-0-54"></a> <span class="bp">self</span><span class="p">,</span>
</span><span id="__span-0-55"><a id="__codelineno-0-55" name="__codelineno-0-55"></a> <span class="o">*</span><span class="p">,</span>
</span><span id="__span-0-56"><a id="__codelineno-0-56" name="__codelineno-0-56"></a> <span class="n">atol</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1e-3</span><span class="p">,</span>
</span><span id="__span-0-57"><a id="__codelineno-0-57" name="__codelineno-0-57"></a> <span class="n">rtol</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1e-3</span><span class="p">,</span>
</span><span id="__span-0-58"><a id="__codelineno-0-58" name="__codelineno-0-58"></a> <span class="n">alpha</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.0129</span><span class="p">,</span>
</span><span id="__span-0-59"><a id="__codelineno-0-59" name="__codelineno-0-59"></a> <span class="n">beta</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">15.6</span><span class="p">,</span>
</span><span id="__span-0-60"><a id="__codelineno-0-60" name="__codelineno-0-60"></a> <span class="n">gamma</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.26</span><span class="p">,</span>
</span><span id="__span-0-61"><a id="__codelineno-0-61" name="__codelineno-0-61"></a> <span class="n">rho</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">5.0</span><span class="p">,</span>
</span><span id="__span-0-62"><a id="__codelineno-0-62" name="__codelineno-0-62"></a> <span class="n">sigma</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">0.6</span><span class="p">,</span>
</span><span id="__span-0-63"><a id="__codelineno-0-63" name="__codelineno-0-63"></a> <span class="n">dtype</span><span class="p">:</span> <span class="n">DTypeLike</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">float32</span><span class="p">,</span>
</span><span id="__span-0-64"><a id="__codelineno-0-64" name="__codelineno-0-64"></a><span class="p">):</span>
</span><span id="__span-0-65"><a id="__codelineno-0-65" name="__codelineno-0-65"></a><span class="w"> </span><span class="sa">r</span><span class="sd">&quot;&quot;&quot;Initialize the WereRabbit neuron model.</span>
</span><span id="__span-0-66"><a id="__codelineno-0-66" name="__codelineno-0-66"></a>
</span><span id="__span-0-67"><a id="__codelineno-0-67" name="__codelineno-0-67"></a><span class="sd"> Args:</span>
</span><span id="__span-0-68"><a id="__codelineno-0-68" name="__codelineno-0-68"></a><span class="sd"> rtol: Relative tolerance for the spiking fixpoint calculation.</span>
</span><span id="__span-0-69"><a id="__codelineno-0-69" name="__codelineno-0-69"></a><span class="sd"> atol: Absolute tolerance for the spiking fixpoint calculation.</span>
</span><span id="__span-0-70"><a id="__codelineno-0-70" name="__codelineno-0-70"></a><span class="sd"> alpha: Current scaling parameter $\alpha = I_{n0}/I_{bias}$ (default: 0.0129)</span>
</span><span id="__span-0-71"><a id="__codelineno-0-71" name="__codelineno-0-71"></a><span class="sd"> beta: Exponential slope $\beta = \kappa/U_t$ (default: 15.6)</span>
</span><span id="__span-0-72"><a id="__codelineno-0-72" name="__codelineno-0-72"></a><span class="sd"> gamma: Coupling parameter $\gamma = 26e^{-2}$</span>
</span><span id="__span-0-73"><a id="__codelineno-0-73" name="__codelineno-0-73"></a><span class="sd"> rho: Steepness of the tanh function $\rho$ (default: 5)</span>
</span><span id="__span-0-74"><a id="__codelineno-0-74" name="__codelineno-0-74"></a><span class="sd"> sigma: Fixpoint distance scaling $\sigma$ (default: 0.6)</span>
</span><span id="__span-0-75"><a id="__codelineno-0-75" name="__codelineno-0-75"></a><span class="sd"> dtype: Data type for arrays (default: float32).</span>
</span><span id="__span-0-76"><a id="__codelineno-0-76" name="__codelineno-0-76"></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="__span-0-77"><a id="__codelineno-0-77" name="__codelineno-0-77"></a> <span class="bp">self</span><span class="o">.</span><span class="n">dtype</span> <span class="o">=</span> <span class="n">dtype</span>
</span><span id="__span-0-78"><a id="__codelineno-0-78" name="__codelineno-0-78"></a> <span class="bp">self</span><span class="o">.</span><span class="n">alpha</span> <span class="o">=</span> <span class="n">alpha</span>
</span><span id="__span-0-79"><a id="__codelineno-0-79" name="__codelineno-0-79"></a> <span class="bp">self</span><span class="o">.</span><span class="n">beta</span> <span class="o">=</span> <span class="n">beta</span>
</span><span id="__span-0-80"><a id="__codelineno-0-80" name="__codelineno-0-80"></a> <span class="bp">self</span><span class="o">.</span><span class="n">gamma</span> <span class="o">=</span> <span class="n">gamma</span>
</span><span id="__span-0-81"><a id="__codelineno-0-81" name="__codelineno-0-81"></a> <span class="bp">self</span><span class="o">.</span><span class="n">rho</span> <span class="o">=</span> <span class="n">rho</span>
</span><span id="__span-0-82"><a id="__codelineno-0-82" name="__codelineno-0-82"></a> <span class="bp">self</span><span class="o">.</span><span class="n">sigma</span> <span class="o">=</span> <span class="n">sigma</span>
</span><span id="__span-0-83"><a id="__codelineno-0-83" name="__codelineno-0-83"></a>
</span><span id="__span-0-84"><a id="__codelineno-0-84" name="__codelineno-0-84"></a> <span class="bp">self</span><span class="o">.</span><span class="n">rtol</span> <span class="o">=</span> <span class="n">rtol</span>
</span><span id="__span-0-85"><a id="__codelineno-0-85" name="__codelineno-0-85"></a> <span class="bp">self</span><span class="o">.</span><span class="n">atol</span> <span class="o">=</span> <span class="n">atol</span>
</span></code></pre></div></td></tr></table></div>
</details>
</div>
</div>
<div class="doc doc-object doc-function">
<h6 id="felice.neuron_models.WereRabbit.init_state" class="doc doc-heading">
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <code class="highlight language-python"><span class="n">init_state</span><span class="p">(</span><span class="n">n_neurons</span><span class="p">:</span> <span class="nb">int</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s1">&#39;neurons 2&#39;</span><span class="p">]</span></code>
</h6>
<div class="doc doc-contents ">
<p>Initialize the neuron state variables.</p>
<p><span class="doc-section-title">Parameters:</span></p>
<table>
<thead>
<tr>
<th>Name</th>
<th>Type</th>
<th>Description</th>
<th>Default</th>
</tr>
</thead>
<tbody>
<tr class="doc-section-item">
<td>
<code>n_neurons</code>
</td>
<td>
<code><span title="int">int</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Number of neurons to initialize.</p>
</div>
</td>
<td>
<em>required</em>
</td>
</tr>
</tbody>
</table>
<p><span class="doc-section-title">Returns:</span></p>
<table>
<thead>
<tr>
<th>Type</th>
<th>Description</th>
</tr>
</thead>
<tbody>
<tr class="doc-section-item">
<td>
<code><span title="jaxtyping.Float">Float</span>[<span title="jaxtyping.Array">Array</span>, &#39;neurons 2&#39;]</code>
</td>
<td>
<div class="doc-md-description">
<p>Initial state array of shape (neurons, 3) containing [u, v, has_spiked],</p>
</div>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code><span title="jaxtyping.Float">Float</span>[<span title="jaxtyping.Array">Array</span>, &#39;neurons 2&#39;]</code>
</td>
<td>
<div class="doc-md-description">
<p>where u and v are the predator/prey membrane voltages, has_spiked is a</p>
</div>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code><span title="jaxtyping.Float">Float</span>[<span title="jaxtyping.Array">Array</span>, &#39;neurons 2&#39;]</code>
</td>
<td>
<div class="doc-md-description">
<p>variable that is 1 whenever the neuron spike and 0 otherwise .</p>
</div>
</td>
</tr>
</tbody>
</table>
<details class="mkdocstrings-source">
<summary>Source code in <code>felice/neuron_models/wererabbit.py</code></summary>
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal"><a href="#__codelineno-0-87"> 87</a></span>
<span class="normal"><a href="#__codelineno-0-88"> 88</a></span>
<span class="normal"><a href="#__codelineno-0-89"> 89</a></span>
<span class="normal"><a href="#__codelineno-0-90"> 90</a></span>
<span class="normal"><a href="#__codelineno-0-91"> 91</a></span>
<span class="normal"><a href="#__codelineno-0-92"> 92</a></span>
<span class="normal"><a href="#__codelineno-0-93"> 93</a></span>
<span class="normal"><a href="#__codelineno-0-94"> 94</a></span>
<span class="normal"><a href="#__codelineno-0-95"> 95</a></span>
<span class="normal"><a href="#__codelineno-0-96"> 96</a></span>
<span class="normal"><a href="#__codelineno-0-97"> 97</a></span>
<span class="normal"><a href="#__codelineno-0-98"> 98</a></span>
<span class="normal"><a href="#__codelineno-0-99"> 99</a></span>
<span class="normal"><a href="#__codelineno-0-100">100</a></span></pre></div></td><td class="code"><div><pre><span></span><code><span id="__span-0-87"><a id="__codelineno-0-87" name="__codelineno-0-87"></a><span class="k">def</span><span class="w"> </span><span class="nf">init_state</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">n_neurons</span><span class="p">:</span> <span class="nb">int</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;neurons 2&quot;</span><span class="p">]:</span>
</span><span id="__span-0-88"><a id="__codelineno-0-88" name="__codelineno-0-88"></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;Initialize the neuron state variables.</span>
</span><span id="__span-0-89"><a id="__codelineno-0-89" name="__codelineno-0-89"></a>
</span><span id="__span-0-90"><a id="__codelineno-0-90" name="__codelineno-0-90"></a><span class="sd"> Args:</span>
</span><span id="__span-0-91"><a id="__codelineno-0-91" name="__codelineno-0-91"></a><span class="sd"> n_neurons: Number of neurons to initialize.</span>
</span><span id="__span-0-92"><a id="__codelineno-0-92" name="__codelineno-0-92"></a>
</span><span id="__span-0-93"><a id="__codelineno-0-93" name="__codelineno-0-93"></a><span class="sd"> Returns:</span>
</span><span id="__span-0-94"><a id="__codelineno-0-94" name="__codelineno-0-94"></a><span class="sd"> Initial state array of shape (neurons, 3) containing [u, v, has_spiked],</span>
</span><span id="__span-0-95"><a id="__codelineno-0-95" name="__codelineno-0-95"></a><span class="sd"> where u and v are the predator/prey membrane voltages, has_spiked is a</span>
</span><span id="__span-0-96"><a id="__codelineno-0-96" name="__codelineno-0-96"></a><span class="sd"> variable that is 1 whenever the neuron spike and 0 otherwise .</span>
</span><span id="__span-0-97"><a id="__codelineno-0-97" name="__codelineno-0-97"></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="__span-0-98"><a id="__codelineno-0-98" name="__codelineno-0-98"></a> <span class="n">x1</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">n_neurons</span><span class="p">,),</span> <span class="n">dtype</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dtype</span><span class="p">)</span>
</span><span id="__span-0-99"><a id="__codelineno-0-99" name="__codelineno-0-99"></a> <span class="n">x2</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="n">n_neurons</span><span class="p">,),</span> <span class="n">dtype</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">dtype</span><span class="p">)</span>
</span><span id="__span-0-100"><a id="__codelineno-0-100" name="__codelineno-0-100"></a> <span class="k">return</span> <span class="n">jnp</span><span class="o">.</span><span class="n">stack</span><span class="p">([</span><span class="n">x1</span><span class="p">,</span> <span class="n">x2</span><span class="p">],</span> <span class="n">axis</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
</span></code></pre></div></td></tr></table></div>
</details>
</div>
</div>
<div class="doc doc-object doc-function">
<h6 id="felice.neuron_models.WereRabbit.vector_field" class="doc doc-heading">
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <code class="highlight language-python"><span class="n">vector_field</span><span class="p">(</span><span class="n">y</span><span class="p">:</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s1">&#39;neurons 2&#39;</span><span class="p">])</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s1">&#39;neurons 2&#39;</span><span class="p">]</span></code>
</h6>
<div class="doc doc-contents ">
<p>Compute vector field of the neuron state variables.</p>
<p>This implements the WereRabbit dynamics</p>
<div class="language-text highlight"><pre><span></span><code>- du/dt: Predator dynamics
- dv/dt: WerePrey dynamics
</code></pre></div>
<p><span class="doc-section-title">Parameters:</span></p>
<table>
<thead>
<tr>
<th>Name</th>
<th>Type</th>
<th>Description</th>
<th>Default</th>
</tr>
</thead>
<tbody>
<tr class="doc-section-item">
<td>
<code>y</code>
</td>
<td>
<code><span title="jaxtyping.Float">Float</span>[<span title="jaxtyping.Array">Array</span>, &#39;neurons 2&#39;]</code>
</td>
<td>
<div class="doc-md-description">
<p>State array of shape (neurons, 2) containing [u, v].</p>
</div>
</td>
<td>
<em>required</em>
</td>
</tr>
</tbody>
</table>
<p><span class="doc-section-title">Returns:</span></p>
<table>
<thead>
<tr>
<th>Type</th>
<th>Description</th>
</tr>
</thead>
<tbody>
<tr class="doc-section-item">
<td>
<code><span title="jaxtyping.Float">Float</span>[<span title="jaxtyping.Array">Array</span>, &#39;neurons 2&#39;]</code>
</td>
<td>
<div class="doc-md-description">
<p>Time derivatives of shape (neurons, 2) containing [du/dt, dv/dt].</p>
</div>
</td>
</tr>
</tbody>
</table>
<details class="mkdocstrings-source">
<summary>Source code in <code>felice/neuron_models/wererabbit.py</code></summary>
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal"><a href="#__codelineno-0-102">102</a></span>
<span class="normal"><a href="#__codelineno-0-103">103</a></span>
<span class="normal"><a href="#__codelineno-0-104">104</a></span>
<span class="normal"><a href="#__codelineno-0-105">105</a></span>
<span class="normal"><a href="#__codelineno-0-106">106</a></span>
<span class="normal"><a href="#__codelineno-0-107">107</a></span>
<span class="normal"><a href="#__codelineno-0-108">108</a></span>
<span class="normal"><a href="#__codelineno-0-109">109</a></span>
<span class="normal"><a href="#__codelineno-0-110">110</a></span>
<span class="normal"><a href="#__codelineno-0-111">111</a></span>
<span class="normal"><a href="#__codelineno-0-112">112</a></span>
<span class="normal"><a href="#__codelineno-0-113">113</a></span>
<span class="normal"><a href="#__codelineno-0-114">114</a></span>
<span class="normal"><a href="#__codelineno-0-115">115</a></span>
<span class="normal"><a href="#__codelineno-0-116">116</a></span>
<span class="normal"><a href="#__codelineno-0-117">117</a></span>
<span class="normal"><a href="#__codelineno-0-118">118</a></span>
<span class="normal"><a href="#__codelineno-0-119">119</a></span>
<span class="normal"><a href="#__codelineno-0-120">120</a></span>
<span class="normal"><a href="#__codelineno-0-121">121</a></span>
<span class="normal"><a href="#__codelineno-0-122">122</a></span>
<span class="normal"><a href="#__codelineno-0-123">123</a></span>
<span class="normal"><a href="#__codelineno-0-124">124</a></span>
<span class="normal"><a href="#__codelineno-0-125">125</a></span>
<span class="normal"><a href="#__codelineno-0-126">126</a></span>
<span class="normal"><a href="#__codelineno-0-127">127</a></span>
<span class="normal"><a href="#__codelineno-0-128">128</a></span>
<span class="normal"><a href="#__codelineno-0-129">129</a></span>
<span class="normal"><a href="#__codelineno-0-130">130</a></span>
<span class="normal"><a href="#__codelineno-0-131">131</a></span>
<span class="normal"><a href="#__codelineno-0-132">132</a></span>
<span class="normal"><a href="#__codelineno-0-133">133</a></span></pre></div></td><td class="code"><div><pre><span></span><code><span id="__span-0-102"><a id="__codelineno-0-102" name="__codelineno-0-102"></a><span class="k">def</span><span class="w"> </span><span class="nf">vector_field</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">y</span><span class="p">:</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;neurons 2&quot;</span><span class="p">])</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;neurons 2&quot;</span><span class="p">]:</span>
</span><span id="__span-0-103"><a id="__codelineno-0-103" name="__codelineno-0-103"></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;Compute vector field of the neuron state variables.</span>
</span><span id="__span-0-104"><a id="__codelineno-0-104" name="__codelineno-0-104"></a>
</span><span id="__span-0-105"><a id="__codelineno-0-105" name="__codelineno-0-105"></a><span class="sd"> This implements the WereRabbit dynamics</span>
</span><span id="__span-0-106"><a id="__codelineno-0-106" name="__codelineno-0-106"></a>
</span><span id="__span-0-107"><a id="__codelineno-0-107" name="__codelineno-0-107"></a><span class="sd"> - du/dt: Predator dynamics</span>
</span><span id="__span-0-108"><a id="__codelineno-0-108" name="__codelineno-0-108"></a><span class="sd"> - dv/dt: WerePrey dynamics</span>
</span><span id="__span-0-109"><a id="__codelineno-0-109" name="__codelineno-0-109"></a>
</span><span id="__span-0-110"><a id="__codelineno-0-110" name="__codelineno-0-110"></a><span class="sd"> Args:</span>
</span><span id="__span-0-111"><a id="__codelineno-0-111" name="__codelineno-0-111"></a><span class="sd"> y: State array of shape (neurons, 2) containing [u, v].</span>
</span><span id="__span-0-112"><a id="__codelineno-0-112" name="__codelineno-0-112"></a>
</span><span id="__span-0-113"><a id="__codelineno-0-113" name="__codelineno-0-113"></a><span class="sd"> Returns:</span>
</span><span id="__span-0-114"><a id="__codelineno-0-114" name="__codelineno-0-114"></a><span class="sd"> Time derivatives of shape (neurons, 2) containing [du/dt, dv/dt].</span>
</span><span id="__span-0-115"><a id="__codelineno-0-115" name="__codelineno-0-115"></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="__span-0-116"><a id="__codelineno-0-116" name="__codelineno-0-116"></a> <span class="n">u</span> <span class="o">=</span> <span class="n">y</span><span class="p">[:,</span> <span class="mi">0</span><span class="p">]</span>
</span><span id="__span-0-117"><a id="__codelineno-0-117" name="__codelineno-0-117"></a> <span class="n">v</span> <span class="o">=</span> <span class="n">y</span><span class="p">[:,</span> <span class="mi">1</span><span class="p">]</span>
</span><span id="__span-0-118"><a id="__codelineno-0-118" name="__codelineno-0-118"></a>
</span><span id="__span-0-119"><a id="__codelineno-0-119" name="__codelineno-0-119"></a> <span class="n">z</span> <span class="o">=</span> <span class="n">jax</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">tanh</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">rho</span> <span class="o">*</span> <span class="p">(</span><span class="n">u</span> <span class="o">-</span> <span class="n">v</span><span class="p">))</span>
</span><span id="__span-0-120"><a id="__codelineno-0-120" name="__codelineno-0-120"></a> <span class="n">du</span> <span class="o">=</span> <span class="p">(</span>
</span><span id="__span-0-121"><a id="__codelineno-0-121" name="__codelineno-0-121"></a> <span class="n">z</span> <span class="o">*</span> <span class="p">(</span><span class="mi">1</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">alpha</span> <span class="o">*</span> <span class="n">jnp</span><span class="o">.</span><span class="n">exp</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">beta</span> <span class="o">*</span> <span class="n">v</span><span class="p">)</span> <span class="o">*</span> <span class="p">(</span><span class="mi">1</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">gamma</span> <span class="o">*</span> <span class="p">(</span><span class="mf">0.5</span> <span class="o">-</span> <span class="n">u</span><span class="p">)))</span>
</span><span id="__span-0-122"><a id="__codelineno-0-122" name="__codelineno-0-122"></a> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">sigma</span>
</span><span id="__span-0-123"><a id="__codelineno-0-123" name="__codelineno-0-123"></a> <span class="p">)</span>
</span><span id="__span-0-124"><a id="__codelineno-0-124" name="__codelineno-0-124"></a> <span class="n">dv</span> <span class="o">=</span> <span class="p">(</span>
</span><span id="__span-0-125"><a id="__codelineno-0-125" name="__codelineno-0-125"></a> <span class="n">z</span>
</span><span id="__span-0-126"><a id="__codelineno-0-126" name="__codelineno-0-126"></a> <span class="o">*</span> <span class="p">(</span><span class="o">-</span><span class="mi">1</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">alpha</span> <span class="o">*</span> <span class="n">jnp</span><span class="o">.</span><span class="n">exp</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">beta</span> <span class="o">*</span> <span class="n">u</span><span class="p">)</span> <span class="o">*</span> <span class="p">(</span><span class="mi">1</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">gamma</span> <span class="o">*</span> <span class="p">(</span><span class="mf">0.5</span> <span class="o">-</span> <span class="n">v</span><span class="p">)))</span>
</span><span id="__span-0-127"><a id="__codelineno-0-127" name="__codelineno-0-127"></a> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">sigma</span>
</span><span id="__span-0-128"><a id="__codelineno-0-128" name="__codelineno-0-128"></a> <span class="p">)</span>
</span><span id="__span-0-129"><a id="__codelineno-0-129" name="__codelineno-0-129"></a>
</span><span id="__span-0-130"><a id="__codelineno-0-130" name="__codelineno-0-130"></a> <span class="n">dv</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">where</span><span class="p">(</span><span class="n">jnp</span><span class="o">.</span><span class="n">allclose</span><span class="p">(</span><span class="n">z</span><span class="p">,</span> <span class="mf">0.0</span><span class="p">),</span> <span class="n">dv</span> <span class="o">*</span> <span class="n">jnp</span><span class="o">.</span><span class="n">sign</span><span class="p">(</span><span class="n">v</span><span class="p">),</span> <span class="n">dv</span><span class="p">)</span>
</span><span id="__span-0-131"><a id="__codelineno-0-131" name="__codelineno-0-131"></a> <span class="n">du</span> <span class="o">=</span> <span class="n">jnp</span><span class="o">.</span><span class="n">where</span><span class="p">(</span><span class="n">jnp</span><span class="o">.</span><span class="n">allclose</span><span class="p">(</span><span class="n">z</span><span class="p">,</span> <span class="mf">0.0</span><span class="p">),</span> <span class="n">du</span> <span class="o">*</span> <span class="n">jnp</span><span class="o">.</span><span class="n">sign</span><span class="p">(</span><span class="n">u</span><span class="p">),</span> <span class="n">du</span><span class="p">)</span>
</span><span id="__span-0-132"><a id="__codelineno-0-132" name="__codelineno-0-132"></a>
</span><span id="__span-0-133"><a id="__codelineno-0-133" name="__codelineno-0-133"></a> <span class="k">return</span> <span class="n">jnp</span><span class="o">.</span><span class="n">stack</span><span class="p">([</span><span class="n">du</span><span class="p">,</span> <span class="n">dv</span><span class="p">],</span> <span class="n">axis</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
</span></code></pre></div></td></tr></table></div>
</details>
</div>
</div>
<div class="doc doc-object doc-function">
<h6 id="felice.neuron_models.WereRabbit.dynamics" class="doc doc-heading">
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <code class="highlight language-python"><span class="n">dynamics</span><span class="p">(</span><span class="n">t</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span> <span class="n">y</span><span class="p">:</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s1">&#39;neurons 2&#39;</span><span class="p">],</span> <span class="n">args</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">])</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s1">&#39;neurons 2&#39;</span><span class="p">]</span></code>
</h6>
<div class="doc doc-contents ">
<p>Compute time derivatives of the neuron state variables.</p>
<p>This implements the WereRabbit dynamics</p>
<div class="language-text highlight"><pre><span></span><code>- du/dt: Predator dynamics
- dv/dt: WerePrey dynamics
</code></pre></div>
<p><span class="doc-section-title">Parameters:</span></p>
<table>
<thead>
<tr>
<th>Name</th>
<th>Type</th>
<th>Description</th>
<th>Default</th>
</tr>
</thead>
<tbody>
<tr class="doc-section-item">
<td>
<code>t</code>
</td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Current simulation time (unused but required by framework).</p>
</div>
</td>
<td>
<em>required</em>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>y</code>
</td>
<td>
<code><span title="jaxtyping.Float">Float</span>[<span title="jaxtyping.Array">Array</span>, &#39;neurons 2&#39;]</code>
</td>
<td>
<div class="doc-md-description">
<p>State array of shape (neurons, 3) containing [u, v, has_spiked].</p>
</div>
</td>
<td>
<em>required</em>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>args</code>
</td>
<td>
<code><span title="typing.Dict">Dict</span>[<span title="str">str</span>, <span title="typing.Any">Any</span>]</code>
</td>
<td>
<div class="doc-md-description">
<p>Additional arguments (unused but required by framework).</p>
</div>
</td>
<td>
<em>required</em>
</td>
</tr>
</tbody>
</table>
<p><span class="doc-section-title">Returns:</span></p>
<table>
<thead>
<tr>
<th>Type</th>
<th>Description</th>
</tr>
</thead>
<tbody>
<tr class="doc-section-item">
<td>
<code><span title="jaxtyping.Float">Float</span>[<span title="jaxtyping.Array">Array</span>, &#39;neurons 2&#39;]</code>
</td>
<td>
<div class="doc-md-description">
<p>Time derivatives of shape (neurons, 3) containing [du/dt, dv/dt, 0].</p>
</div>
</td>
</tr>
</tbody>
</table>
<details class="mkdocstrings-source">
<summary>Source code in <code>felice/neuron_models/wererabbit.py</code></summary>
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal"><a href="#__codelineno-0-135">135</a></span>
<span class="normal"><a href="#__codelineno-0-136">136</a></span>
<span class="normal"><a href="#__codelineno-0-137">137</a></span>
<span class="normal"><a href="#__codelineno-0-138">138</a></span>
<span class="normal"><a href="#__codelineno-0-139">139</a></span>
<span class="normal"><a href="#__codelineno-0-140">140</a></span>
<span class="normal"><a href="#__codelineno-0-141">141</a></span>
<span class="normal"><a href="#__codelineno-0-142">142</a></span>
<span class="normal"><a href="#__codelineno-0-143">143</a></span>
<span class="normal"><a href="#__codelineno-0-144">144</a></span>
<span class="normal"><a href="#__codelineno-0-145">145</a></span>
<span class="normal"><a href="#__codelineno-0-146">146</a></span>
<span class="normal"><a href="#__codelineno-0-147">147</a></span>
<span class="normal"><a href="#__codelineno-0-148">148</a></span>
<span class="normal"><a href="#__codelineno-0-149">149</a></span>
<span class="normal"><a href="#__codelineno-0-150">150</a></span>
<span class="normal"><a href="#__codelineno-0-151">151</a></span>
<span class="normal"><a href="#__codelineno-0-152">152</a></span>
<span class="normal"><a href="#__codelineno-0-153">153</a></span>
<span class="normal"><a href="#__codelineno-0-154">154</a></span>
<span class="normal"><a href="#__codelineno-0-155">155</a></span>
<span class="normal"><a href="#__codelineno-0-156">156</a></span>
<span class="normal"><a href="#__codelineno-0-157">157</a></span>
<span class="normal"><a href="#__codelineno-0-158">158</a></span></pre></div></td><td class="code"><div><pre><span></span><code><span id="__span-0-135"><a id="__codelineno-0-135" name="__codelineno-0-135"></a><span class="k">def</span><span class="w"> </span><span class="nf">dynamics</span><span class="p">(</span>
</span><span id="__span-0-136"><a id="__codelineno-0-136" name="__codelineno-0-136"></a> <span class="bp">self</span><span class="p">,</span>
</span><span id="__span-0-137"><a id="__codelineno-0-137" name="__codelineno-0-137"></a> <span class="n">t</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
</span><span id="__span-0-138"><a id="__codelineno-0-138" name="__codelineno-0-138"></a> <span class="n">y</span><span class="p">:</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;neurons 2&quot;</span><span class="p">],</span>
</span><span id="__span-0-139"><a id="__codelineno-0-139" name="__codelineno-0-139"></a> <span class="n">args</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">],</span>
</span><span id="__span-0-140"><a id="__codelineno-0-140" name="__codelineno-0-140"></a><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;neurons 2&quot;</span><span class="p">]:</span>
</span><span id="__span-0-141"><a id="__codelineno-0-141" name="__codelineno-0-141"></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;Compute time derivatives of the neuron state variables.</span>
</span><span id="__span-0-142"><a id="__codelineno-0-142" name="__codelineno-0-142"></a>
</span><span id="__span-0-143"><a id="__codelineno-0-143" name="__codelineno-0-143"></a><span class="sd"> This implements the WereRabbit dynamics</span>
</span><span id="__span-0-144"><a id="__codelineno-0-144" name="__codelineno-0-144"></a>
</span><span id="__span-0-145"><a id="__codelineno-0-145" name="__codelineno-0-145"></a><span class="sd"> - du/dt: Predator dynamics</span>
</span><span id="__span-0-146"><a id="__codelineno-0-146" name="__codelineno-0-146"></a><span class="sd"> - dv/dt: WerePrey dynamics</span>
</span><span id="__span-0-147"><a id="__codelineno-0-147" name="__codelineno-0-147"></a>
</span><span id="__span-0-148"><a id="__codelineno-0-148" name="__codelineno-0-148"></a><span class="sd"> Args:</span>
</span><span id="__span-0-149"><a id="__codelineno-0-149" name="__codelineno-0-149"></a><span class="sd"> t: Current simulation time (unused but required by framework).</span>
</span><span id="__span-0-150"><a id="__codelineno-0-150" name="__codelineno-0-150"></a><span class="sd"> y: State array of shape (neurons, 3) containing [u, v, has_spiked].</span>
</span><span id="__span-0-151"><a id="__codelineno-0-151" name="__codelineno-0-151"></a><span class="sd"> args: Additional arguments (unused but required by framework).</span>
</span><span id="__span-0-152"><a id="__codelineno-0-152" name="__codelineno-0-152"></a>
</span><span id="__span-0-153"><a id="__codelineno-0-153" name="__codelineno-0-153"></a><span class="sd"> Returns:</span>
</span><span id="__span-0-154"><a id="__codelineno-0-154" name="__codelineno-0-154"></a><span class="sd"> Time derivatives of shape (neurons, 3) containing [du/dt, dv/dt, 0].</span>
</span><span id="__span-0-155"><a id="__codelineno-0-155" name="__codelineno-0-155"></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="__span-0-156"><a id="__codelineno-0-156" name="__codelineno-0-156"></a> <span class="n">dxdt</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">vector_field</span><span class="p">(</span><span class="n">y</span><span class="p">)</span>
</span><span id="__span-0-157"><a id="__codelineno-0-157" name="__codelineno-0-157"></a>
</span><span id="__span-0-158"><a id="__codelineno-0-158" name="__codelineno-0-158"></a> <span class="k">return</span> <span class="n">dxdt</span>
</span></code></pre></div></td></tr></table></div>
</details>
</div>
</div>
<div class="doc doc-object doc-function">
<h6 id="felice.neuron_models.WereRabbit.spike_condition" class="doc doc-heading">
<code class="doc-symbol doc-symbol-heading doc-symbol-method"></code> <code class="highlight language-python"><span class="n">spike_condition</span><span class="p">(</span><span class="n">t</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span> <span class="n">y</span><span class="p">:</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s1">&#39;neurons 2&#39;</span><span class="p">],</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">])</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s1">&#39; neurons&#39;</span><span class="p">]</span></code>
</h6>
<div class="doc doc-contents ">
<p>Compute spike condition for event detection.</p>
<p>A spike is triggered when the system reach to a fixpoint.</p>
<details class="info" open>
<summary>INFO</summary>
<p><code>has_spiked</code> is use to the system don't detect a continuos
spike when reach a fixpoint.</p>
</details>
<p><span class="doc-section-title">Parameters:</span></p>
<table>
<thead>
<tr>
<th>Name</th>
<th>Type</th>
<th>Description</th>
<th>Default</th>
</tr>
</thead>
<tbody>
<tr class="doc-section-item">
<td>
<code>t</code>
</td>
<td>
<code><span title="float">float</span></code>
</td>
<td>
<div class="doc-md-description">
<p>Current simulation time (unused but required by the framework).</p>
</div>
</td>
<td>
<em>required</em>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>y</code>
</td>
<td>
<code><span title="jaxtyping.Float">Float</span>[<span title="jaxtyping.Array">Array</span>, &#39;neurons 2&#39;]</code>
</td>
<td>
<div class="doc-md-description">
<p>State array of shape (neurons, 3) containing [u, v, has_spiked].</p>
</div>
</td>
<td>
<em>required</em>
</td>
</tr>
<tr class="doc-section-item">
<td>
<code>**kwargs</code>
</td>
<td>
<code><span title="typing.Dict">Dict</span>[<span title="str">str</span>, <span title="typing.Any">Any</span>]</code>
</td>
<td>
<div class="doc-md-description">
<p>Additional keyword arguments (unused).</p>
</div>
</td>
<td>
<code>{}</code>
</td>
</tr>
</tbody>
</table>
<p><span class="doc-section-title">Returns:</span></p>
<table>
<thead>
<tr>
<th>Type</th>
<th>Description</th>
</tr>
</thead>
<tbody>
<tr class="doc-section-item">
<td>
<code><span title="jaxtyping.Float">Float</span>[<span title="jaxtyping.Array">Array</span>, &#39; neurons&#39;]</code>
</td>
<td>
<div class="doc-md-description">
<p>Spike condition array of shape (neurons,). Positive values indicate spike.</p>
</div>
</td>
</tr>
</tbody>
</table>
<details class="mkdocstrings-source">
<summary>Source code in <code>felice/neuron_models/wererabbit.py</code></summary>
<div class="language-python highlight"><table class="highlighttable"><tr><td class="linenos"><div class="linenodiv"><pre><span></span><span class="normal"><a href="#__codelineno-0-160">160</a></span>
<span class="normal"><a href="#__codelineno-0-161">161</a></span>
<span class="normal"><a href="#__codelineno-0-162">162</a></span>
<span class="normal"><a href="#__codelineno-0-163">163</a></span>
<span class="normal"><a href="#__codelineno-0-164">164</a></span>
<span class="normal"><a href="#__codelineno-0-165">165</a></span>
<span class="normal"><a href="#__codelineno-0-166">166</a></span>
<span class="normal"><a href="#__codelineno-0-167">167</a></span>
<span class="normal"><a href="#__codelineno-0-168">168</a></span>
<span class="normal"><a href="#__codelineno-0-169">169</a></span>
<span class="normal"><a href="#__codelineno-0-170">170</a></span>
<span class="normal"><a href="#__codelineno-0-171">171</a></span>
<span class="normal"><a href="#__codelineno-0-172">172</a></span>
<span class="normal"><a href="#__codelineno-0-173">173</a></span>
<span class="normal"><a href="#__codelineno-0-174">174</a></span>
<span class="normal"><a href="#__codelineno-0-175">175</a></span>
<span class="normal"><a href="#__codelineno-0-176">176</a></span>
<span class="normal"><a href="#__codelineno-0-177">177</a></span>
<span class="normal"><a href="#__codelineno-0-178">178</a></span>
<span class="normal"><a href="#__codelineno-0-179">179</a></span>
<span class="normal"><a href="#__codelineno-0-180">180</a></span>
<span class="normal"><a href="#__codelineno-0-181">181</a></span>
<span class="normal"><a href="#__codelineno-0-182">182</a></span>
<span class="normal"><a href="#__codelineno-0-183">183</a></span>
<span class="normal"><a href="#__codelineno-0-184">184</a></span>
<span class="normal"><a href="#__codelineno-0-185">185</a></span>
<span class="normal"><a href="#__codelineno-0-186">186</a></span>
<span class="normal"><a href="#__codelineno-0-187">187</a></span>
<span class="normal"><a href="#__codelineno-0-188">188</a></span>
<span class="normal"><a href="#__codelineno-0-189">189</a></span>
<span class="normal"><a href="#__codelineno-0-190">190</a></span>
<span class="normal"><a href="#__codelineno-0-191">191</a></span>
<span class="normal"><a href="#__codelineno-0-192">192</a></span>
<span class="normal"><a href="#__codelineno-0-193">193</a></span>
<span class="normal"><a href="#__codelineno-0-194">194</a></span></pre></div></td><td class="code"><div><pre><span></span><code><span id="__span-0-160"><a id="__codelineno-0-160" name="__codelineno-0-160"></a><span class="k">def</span><span class="w"> </span><span class="nf">spike_condition</span><span class="p">(</span>
</span><span id="__span-0-161"><a id="__codelineno-0-161" name="__codelineno-0-161"></a> <span class="bp">self</span><span class="p">,</span>
</span><span id="__span-0-162"><a id="__codelineno-0-162" name="__codelineno-0-162"></a> <span class="n">t</span><span class="p">:</span> <span class="nb">float</span><span class="p">,</span>
</span><span id="__span-0-163"><a id="__codelineno-0-163" name="__codelineno-0-163"></a> <span class="n">y</span><span class="p">:</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot;neurons 2&quot;</span><span class="p">],</span>
</span><span id="__span-0-164"><a id="__codelineno-0-164" name="__codelineno-0-164"></a> <span class="o">**</span><span class="n">kwargs</span><span class="p">:</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">],</span>
</span><span id="__span-0-165"><a id="__codelineno-0-165" name="__codelineno-0-165"></a><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Float</span><span class="p">[</span><span class="n">Array</span><span class="p">,</span> <span class="s2">&quot; neurons&quot;</span><span class="p">]:</span>
</span><span id="__span-0-166"><a id="__codelineno-0-166" name="__codelineno-0-166"></a><span class="w"> </span><span class="sd">&quot;&quot;&quot;Compute spike condition for event detection.</span>
</span><span id="__span-0-167"><a id="__codelineno-0-167" name="__codelineno-0-167"></a>
</span><span id="__span-0-168"><a id="__codelineno-0-168" name="__codelineno-0-168"></a><span class="sd"> A spike is triggered when the system reach to a fixpoint.</span>
</span><span id="__span-0-169"><a id="__codelineno-0-169" name="__codelineno-0-169"></a>
</span><span id="__span-0-170"><a id="__codelineno-0-170" name="__codelineno-0-170"></a><span class="sd"> INFO:</span>
</span><span id="__span-0-171"><a id="__codelineno-0-171" name="__codelineno-0-171"></a><span class="sd"> `has_spiked` is use to the system don&#39;t detect a continuos</span>
</span><span id="__span-0-172"><a id="__codelineno-0-172" name="__codelineno-0-172"></a><span class="sd"> spike when reach a fixpoint.</span>
</span><span id="__span-0-173"><a id="__codelineno-0-173" name="__codelineno-0-173"></a>
</span><span id="__span-0-174"><a id="__codelineno-0-174" name="__codelineno-0-174"></a><span class="sd"> Args:</span>
</span><span id="__span-0-175"><a id="__codelineno-0-175" name="__codelineno-0-175"></a><span class="sd"> t: Current simulation time (unused but required by the framework).</span>
</span><span id="__span-0-176"><a id="__codelineno-0-176" name="__codelineno-0-176"></a><span class="sd"> y: State array of shape (neurons, 3) containing [u, v, has_spiked].</span>
</span><span id="__span-0-177"><a id="__codelineno-0-177" name="__codelineno-0-177"></a><span class="sd"> **kwargs: Additional keyword arguments (unused).</span>
</span><span id="__span-0-178"><a id="__codelineno-0-178" name="__codelineno-0-178"></a>
</span><span id="__span-0-179"><a id="__codelineno-0-179" name="__codelineno-0-179"></a><span class="sd"> Returns:</span>
</span><span id="__span-0-180"><a id="__codelineno-0-180" name="__codelineno-0-180"></a><span class="sd"> Spike condition array of shape (neurons,). Positive values indicate spike.</span>
</span><span id="__span-0-181"><a id="__codelineno-0-181" name="__codelineno-0-181"></a><span class="sd"> &quot;&quot;&quot;</span>
</span><span id="__span-0-182"><a id="__codelineno-0-182" name="__codelineno-0-182"></a> <span class="n">_atol</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">atol</span>
</span><span id="__span-0-183"><a id="__codelineno-0-183" name="__codelineno-0-183"></a> <span class="n">_rtol</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">rtol</span>
</span><span id="__span-0-184"><a id="__codelineno-0-184" name="__codelineno-0-184"></a> <span class="n">_norm</span> <span class="o">=</span> <span class="n">optx</span><span class="o">.</span><span class="n">rms_norm</span>
</span><span id="__span-0-185"><a id="__codelineno-0-185" name="__codelineno-0-185"></a>
</span><span id="__span-0-186"><a id="__codelineno-0-186" name="__codelineno-0-186"></a> <span class="n">vf</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">dynamics</span><span class="p">(</span><span class="n">t</span><span class="p">,</span> <span class="n">y</span><span class="p">,</span> <span class="p">{})</span>
</span><span id="__span-0-187"><a id="__codelineno-0-187" name="__codelineno-0-187"></a>
</span><span id="__span-0-188"><a id="__codelineno-0-188" name="__codelineno-0-188"></a> <span class="nd">@jax</span><span class="o">.</span><span class="n">vmap</span>
</span><span id="__span-0-189"><a id="__codelineno-0-189" name="__codelineno-0-189"></a> <span class="k">def</span><span class="w"> </span><span class="nf">calculate_norm</span><span class="p">(</span><span class="n">vf</span><span class="p">,</span> <span class="n">y</span><span class="p">):</span>
</span><span id="__span-0-190"><a id="__codelineno-0-190" name="__codelineno-0-190"></a> <span class="k">return</span> <span class="n">_atol</span> <span class="o">+</span> <span class="n">_rtol</span> <span class="o">*</span> <span class="n">_norm</span><span class="p">(</span><span class="n">y</span><span class="p">[:</span><span class="o">-</span><span class="mi">1</span><span class="p">])</span> <span class="o">-</span> <span class="n">_norm</span><span class="p">(</span><span class="n">vf</span><span class="p">[:</span><span class="o">-</span><span class="mi">1</span><span class="p">])</span>
</span><span id="__span-0-191"><a id="__codelineno-0-191" name="__codelineno-0-191"></a>
</span><span id="__span-0-192"><a id="__codelineno-0-192" name="__codelineno-0-192"></a> <span class="n">base_cond</span> <span class="o">=</span> <span class="n">calculate_norm</span><span class="p">(</span><span class="n">vf</span><span class="p">,</span> <span class="n">y</span><span class="p">)</span>
</span><span id="__span-0-193"><a id="__codelineno-0-193" name="__codelineno-0-193"></a>
</span><span id="__span-0-194"><a id="__codelineno-0-194" name="__codelineno-0-194"></a> <span class="k">return</span> <span class="n">base_cond</span>
</span></code></pre></div></td></tr></table></div>
</details>
</div>
</div>
</div>
</div>
</div>
</div>
</div>
</div>
</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>