#+TITLE: NeuroML2 and Arbor
#+AUTHOR: T. Hater
#+EMAIL: t.hater@fz-juelich.de

* Arbor 💘 NeuroML2

*HPCNS*
31.01.2022

Thorsten Hater (=t.hater@fz-juelich.de=)

* What is NeuroML2?
  #+begin_quote
  XML based format for the description of neuro-scientific simulations.
  #+end_quote
- Versions
  - v1 :: not used anymore, conductance models only
  - v2 :: built on top of LEMS, only relevant version as of now
  - v3 :: upcoming, iterates on v2 and prepares for ModeCI
- Coverage
  - Cell/Ion channel dynamics
  - Morphology and parametrisations
  - Connectivity/networks
  - Units and quantities
- NML2 is then _used_ in LEMS to add
  - Simulation
  - Visualisation/Data extraction

~LEMS~ = Low Entropy Model Specification (language)

* Example: Cell
#+begin_src xml
<cell id="cell">
  <morphology ="morphology" />
    <segment id="0" name="root">
      <proximal x="0" y="0" z="0" diameter="17.841242"/>
      <distal   x="0" y="0" z="0" diameter="17.841242"/>
    </segment>

    <segmentGroup id="soma">
        <member segment="0"/>
    </segmentGroup>
  </morphology>

  <biophysicalProperties id="properties">
    <membraneProperties>
      <channelDensity id="leak" ionChannel="pas" g="0.3mS_per_cm2" e="-54.387mV" ion="none"/>
      <specificCapacitance value="1.0 uF_per_cm2"/>
      <initMembPotential value="-65mV"/>
    </membraneProperties>
  </biophysicalProperties>
</cell>
#+end_src

* The 10'000ft View
#+begin_example
┌───────────────────────┐      ┌───────────────────────┐
│                       │      │                       │
│  LEMS Schema XSD      │      │  NML2 Schema XSD      │
│                       │      │                       │
└───────────┬───────────┘      └───────────┬───────────┘
            │                              │
            │                              │
┌───────────▼───────────┐      ┌───────────▼───────────┐
│                       │      │                       │
│  NeuroML2 Definitions │      │  NML2 Instances       │
│                       │      │                       │
│  <ComponentType>      ├─────►│  <Component>          │
│                       │      │                       │
│                       │      │  <ComponentType>      │
│                       │      │                       │
└───────────────────────┘      └───────────────────────┘
#+end_example


* =NeuroML2= and ~Arbor~
#+begin_quote
It's the worst system we have, except all others, which are even worse.
#+end_quote
- Status
  - _done (SY)_ load morphologies from =.nml= files.
  - _in progress_
    1. compiling the dynamics specification to Arbor's native format.
    2. extracting cell parametrisations
  - _planned_: Reading networks.
- Why?
  - only decently portable and somewhat modular format
  - interesting eco-system: OSB, NML-DB, ...
  - target a new demographic of scientists with existing models
  - =Arborio= inferior olive models in NML2
  - reference implementation has quite poor performance
  - personal feud with EDEN

* Ion Channel Dynamics in Arbor

- Arbor ingests =NMODL= files per default
  - Actually, we compile =NMODL= to a set of =C= callbacks (the ABI)
  - Decoupling Arbor from the mechanisms was one of my projects in '21
- =NMODL= is *not* anyones favourite language, not by quite some margin.
  - ...but, it _is_ kind of the default for MC neurons
  - (We are working on an alternative called =ArbLang=, also via the ABI)
- Status Quo: =NML2= channels in Arbor
  1. Install =jnml=
  2. Run =jnml -neuron my_channel.nml=
  3. Edit the resulting =.mod= to fit Arbor's dialect
  4. Manually copy out parametrisations, channel assignments, ...
  5. Run the simulation, complain that it's way too slow

* Example: Ion-Channel
#+begin_src xml
<?xml version="1.0" encoding="UTF-8"?>
<neuroml xmlns="http://www.neuroml.org/schema/neuroml2"
  xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
  xsi:schemaLocation="http://www.neuroml.org/schema/neuroml2 ../Schemas/NeuroML2/NeuroML_v2beta4.xsd"
         id="NML2_SimpleIonChannel">

    <ionChannelHH id="NaConductance" conductance="10pS" species="na">
        <gateHHrates id="m" instances="3">
            <forwardRate type="HHExpRate" rate="1per_ms" midpoint="-40mV" scale="10mV"/>
            <reverseRate type="HHExpRate" rate="4per_ms" midpoint="-65mV" scale="-18mV"/>
        </gateHHrates>
    </ionChannelHH>
</neuroml>
#+end_src

* Bringing =NML2= Channels to Arbor

- Need to AOT compile channels to Arbor's ABI, possibly via =NMODL= and =modcc=
- Why compilation?
  - Each =.nml= file may define new kinds of mechanisms.
  - Adding a single line in NML2 may change dynamics substantially
  - Channels might be deeply nested, so optimisation is crucial
- But, what do we actually compile?
  - NML2 is really a library of LEMS definitions
  - So, we compile LEMS =<Dynamics>= specifications
  - Here, the output is usually a conductance or a current
  - Types like =<cell>= then define ODEs ~U'(t) = g I(t)~

* =nmlcc:= a Compiler for NML2

- Written in Rust, because Rust is great for the job
  - pattern matching and ADTs _alone_ are amazing for compilers
  - good libraries: XML, parsing, ...
  - artifacts are easy to produce for consumers (=cargo run=)
- It's actually two compilers, because we need more of them
  1. Data model: LEMS schema =.xsd= to Rust =struct= and =enum=
  2. LEMS/NML2: XML to abstract syntax tree (AST)
- From the AST we can extract and export
  - NMODL       :: implemented, being tested
  - Cell        :: implemented, being tested
  - Simulations :: working on it
  - Networks    :: next step, somewhere in the future
- Covers only the multi-compartment aspects of =NML2=
- Find us here [[https://github.com/thorstenhater/nmlcc]]
  - =sloc= reports ~~9k~ lines of Rust, thereof ~~6k~ auto-generated
  - test-suite and CI enabled

* Unpacking =NML2=

** A Gate: =HHExp=

- Consider an excerpt from our example before
  #+begin_src xml
  <gateHHrates id="m" instances="3">
    <forwardRate type="HHExpRate" rate="1per_ms" midpoint="-40mV" scale="10mV"/>
    <reverseRate type="HHExpRate" rate="4per_ms" midpoint="-65mV" scale="-18mV"/>
  </gateHHrates>
  #+end_src
- Opening the lid on =HHExpRate=
  #+begin_src xml
  <ComponentType name="baseVoltageDepRate">
    <Requirement name="v" dimension="voltage"/>
  </ComponentType>

  <ComponentType name="baseHHRate" extends="baseVoltageDepRate">
    <Parameter name="rate" dimension="per_time"/>
    <Parameter name="midpoint" dimension="voltage"/>
    <Parameter name="scale" dimension="voltage"/>
  </ComponentType>

  <ComponentType name="HHExpRate" extends="baseHHRate">
    <Dynamics>
      <DerivedVariable name="r" value="rate * exp((v - midpoint)/scale)" />
    </Dynamics>
  </ComponentType>
  #+end_src
- So, =ComponentType= is essentially a _class_ with single inheritance.
  - Then, =forwardRate= is an _instance_ (in form of a member variable)
  - =gateHHrates= declares =forwardRate= to be of type =baseHHRate=
  - =HHExpRate= can fill that slot since it inherits =baseHHRate=
- =Dynamics= defines the time evolution of an instance
  - =Variables=, =DerivedVariable=
  - ODEs via =StateVariable=
  - =KineticScheme= for reactions

** But what _is_ a Gate Actually _Doing_?

#+begin_src xml
<ComponentType name="gateHHrates" extends="gate">
  <Child name="forwardRate" type="baseVoltageDepRate"/>
  <Child name="reverseRate" type="baseVoltageDepRate"/>

  <Dynamics>
    <StateVariable name="q" dimension="none" exposure="q"/>
    <DerivedVariable name="alpha" dimension="per_time" select="forwardRate/r"/>
    <DerivedVariable name="beta"  dimension="per_time" select="reverseRate/r"/>
    <DerivedVariable name="inf"   dimension="none"     value="alpha/(alpha+beta)"/>
    <DerivedVariable name="tau"   dimension="time"     value="1/(alpha+beta)"/>

    <TimeDerivative variable="q" value="(inf - q) / tau"/>

    <OnStart>
      <StateAssignment variable="q" value="inf"/>
    </OnStart>
  </Dynamics>
</ComponentType>
#+end_src

- Much simplified adaption.
- Here we define the dynamics of a =gate= in terms of an ODE.
- Note, that we can swap the instantiations of =forwardRate= and =reverseRate=.
- Thus we can compose a whole zoo of HH-like channels.
- In the higher layers, =q= will be used together with similar variables of
  other gates to determine conductance densities =g= and in turn current
  densities =i=, and so on, up to Ohm's law.

* Example: Running =nmlcc=
** Running =nmlcc=
#+begin_src bash
$> # Build data model (this is *not* needed in regular use)
$> cargo run --bin schema
[... download & compile dependencies ...]
[... compile ...]
   Compiling nml2 v0.1.0 (/Users/hater/src/nml2)
    Finished dev [unoptimized + debuginfo] target(s) in 1.51s
     Running `target/debug/schema`
$> # Compile nml to NMODL
$> cargo run -- nmodl --type ionChannelHH --parameter='-*' example/nml-minimal-channel.xml
[... download & compile dependencies ...]
[... compile ...]
    Finished dev [unoptimized + debuginfo] target(s) in 0.17s
     Running `target/debug/nmlcc nmodl --type ionChannelHH '--parameter=-*' example/nml-minimal-channel.xml`
[... logging ...]
$> # Show output
$> cat NaConductance.mod
#+end_src

** Output, Compressed to Fit
#+begin_src
NEURON { ... }
STATE { gates_m_q }
INITIAL { ... }

DERIVATIVE dstate {
  LOCAL gates_m_forwardRate_r, gates_m_reverseRate_r, gates_m_tau, gates_m_inf
  gates_m_forwardRate_r = exp(0.1 * (40 + v))
  gates_m_reverseRate_r = 4 * exp(-0.056 * (65 + v))
  gates_m_tau = (gates_m_forwardRate_r + gates_m_reverseRate_r)^-1
  gates_m_inf = gates_m_forwardRate_r * (gates_m_forwardRate_r + gates_m_reverseRate_r)^-1
  gates_m_q' = (gates_m_inf + -1 * gates_m_q) * gates_m_tau^-1
}

BREAKPOINT {
  SOLVE dstate METHOD cnexp
  ina = 0.00000001 * gates_m_q^3 * (v + -1 * ena)
}
#+end_src

* =NML2=: Prêt-à-simuler with ~nmlcc bundle~

This is still tedious and involves manually copying out parameters from XML.

So, let's do better:
#+begin_src bash
$> rm -rf hhcell                                       # no cheating
$> nmlcc bundle openworm/HHCellNetwork.net.nml hhcell  # all-in-one exporter
$> ls hhcell
Permissions Size User  Date Modified Name
drwxr-xr-x     - hater 29 Jan 14:01  acc               # cell model and region assignments
drwxr-xr-x     - hater 29 Jan 14:01  cat               # NMODL dynamics
drwxr-xr-x     - hater 29 Jan 14:01  mrf               # morphology files
.rw-r--r--  1.1k hater 29 Jan 14:01  main.hhcell.py    # template simulation for id=hhcell
$> python3 main.hhcell.py                              # runs, but without stimuli
#+end_src
From here, we need to tweak =main.<id>.py= according to our needs
- placement of stimuli    :: part of the network specification, but the list is given
- measurements            :: that you will have to decide yourself
- simulation =t= and =dt=     :: ditto; it is given but in LEMS not NML2

* Outlook

NML2 specificies the _whole_ simulation flow and we really, really want to
exploit this for performance. This is also nothing we would like to do manually,
since it involves lots of boilerplate and repetition.

** Specialised Mechanisms
Consider this channel assignment
#+begin_src xml
<channelDensity id="leak" ionChannel="pas" g="0.3 mS_per_cm2" e="-54.387 mV" ion="none"/>
#+end_src
in Arbor and NEURON =g= and =e= are runtime variables, but they never change
after the simulation has been started. (In Arbor, that might change sometime soon)

So, we can specialise, possibly per assigned region, and make those _constants_,
allowing even more speed by uncovering more optimisation potential and reducing
memory traffic.

** Merging Mechanisms
#+begin_src xml
<membraneProperties>
  <channelDensity id="pas" ionChannel="pas" condDensity="0.3 mS_per_cm2" erev="-54.387mV" ion="none"/>
  <channelDensity id="na"  ionChannel="na"  condDensity="120 mS_per_cm2" erev="50.0 mV" ion="na"/>
  <channelDensity id="k"   ionChannel="k"   condDensity=" 36 mS_per_cm2" erev="-77mV" ion="k"/>
</membraneProperties>
#+end_src
By essentially the same idea, we can merge all mechanisms on a region in to one.
We create one such _'super-mechanism'_ per assignment. This also gives a
significant speed-up, by reducing memory accesses, function call overheads,
sharing RO data etc.

* Performance
- HH tutorial cell from =https://github.com/openworm/hodgkin_huxley_tutorial=
  - simulation settings: ~t=1000 ms~ and ~dt=0.0025~
  - soma-only, ~d=17.8um~ discretized into ~0.1um~ segments
- Arbor 0.6 Release, single thread
  - times are measured across ~sim.run(...)~
  - thus model building is included
|-------------------+----------------+----------+---------------+----------|
| ARB_VECTORIZE=    | *OFF* t_wall/s | Speed-up | *ON* t_wall/s | Speed-up |
|-------------------+----------------+----------+---------------+----------|
| jnml              |         13.266 |      1.0 |         7.988 |      1.0 |
| nmlcc 0.2         |          5.934 |      2.2 |         2.827 |      2.8 |
| + CF + SM \dagger |          6.210 |      2.1 |         2.727 |      2.9 |
|-------------------+----------------+----------+---------------+----------|
| Arbor HH          |          6.376 |      2.1 |         2.960 |      2.7 |
|-------------------+----------------+----------+---------------+----------|
| hand-optimised    |          5.829 |      2.3 |         2.911 |      2.7 |
| + CF              |          5.731 |      2.3 |         2.782 |      2.9 |
| + SM              |          5.212 |      2.5 |         2.504 |      3.2 |
|-------------------+----------------+----------+---------------+----------|
#+TBLFM: $3=@2$2/$2;%.1f::$5=@2$4/$4;%.1f

\dagger Incomplete feature

* Recap

- Enabling a workflow to bring =NML2= projects to Arbor.
  - direct conversion of =NML2= dynamics into Arbor NMODL
  - bespoke representations optimised for Arbor
  - bio-physical parameters / assignments / simulation settings
  - merging and specialising mechanisms based on _'runtime'_ information
- Missing
  - extraction of networks
  - lots of testing
- Nice to have
  - micro-optimisations: CSE, some algebraic passes, ...
  - NML2-specific peephole optimisations
- Working with
  - ArborIO   :: inferior olive model in NML2, so far hand-ported
    - will drive features added to =nmlcc=
  - NML2 devs :: P. Gleeson et al
    - correctness testing
    - include Arbor in NML2 test framework

* Fin
#+begin_example

    ███        ▄█    █▄       ▄████████ ███▄▄▄▄      ▄█   ▄█▄    ▄████████
▀█████████▄   ███    ███     ███    ███ ███▀▀▀██▄   ███ ▄███▀   ███    ███
   ▀███▀▀██   ███    ███     ███    ███ ███   ███   ███▐██▀     ███    █▀
    ███   ▀  ▄███▄▄▄▄███▄▄   ███    ███ ███   ███  ▄█████▀      ███
    ███     ▀▀███▀▀▀▀███▀  ▀███████████ ███   ███ ▀▀█████▄    ▀███████████
    ███       ███    ███     ███    ███ ███   ███   ███▐██▄            ███
    ███       ███    ███     ███    ███ ███   ███   ███ ▀███▄    ▄█    ███
   ▄████▀     ███    █▀      ███    █▀   ▀█   █▀    ███   ▀█▀  ▄████████▀
                                                    ▀
#+end_example