SDE examples (#2030)

Addresses #1987

* reproducibility test for random number generator
* C++ example demonstrating Ornstein-Uhlenbeck process
* Python example featuring stochastic mechanism
* stochastic mechanism catalogue

The Python example is contributed by Jannik Luboeinski's (@jlubo): https://github.com/jlubo/arbor_ou_lif_example
and adapted to fit into Arbor.

arbor/include/arbor/mechcat.hpp

@@ -112,6 +112,7 @@ class ARB_ARBOR_API mechanism_catalogue {
 ARB_ARBOR_API const mechanism_catalogue& global_default_catalogue();
 ARB_ARBOR_API const mechanism_catalogue& global_allen_catalogue();
 ARB_ARBOR_API const mechanism_catalogue& global_bbp_catalogue();
+ARB_ARBOR_API const mechanism_catalogue& global_stochastic_catalogue();
 
 // Load catalogue from disk.
 ARB_ARBOR_API const mechanism_catalogue load_catalogue(const std::string&);

doc/concepts/mechanisms.rst

@@ -77,6 +77,14 @@ Two catalogues are provided that collect mechanisms associated with specific pro
 * ``bbp_catalogue`` For models published by the Blue Brain Project (BBP).
 * ``allen_catalogue`` For models published on the Allen Brain Atlas Database.
 
+A fourth catalogue ``stochastic_catalogue`` provides mechanisms expressed as stochastic differential
+equations:
+
+* *ou_input* Synapse mechanism that can stochastically account for a population of *ou_input*
+  synapses.  The mechanism is similar to *expsyn_curr* but with the exponential decay being subject
+  to noise due to a Ornstein-Uhlenbeck process.
+
+
 .. _mechanisms_dynamic:
 
 Adding Catalogues to Arbor
@@ -316,6 +324,50 @@ Euler-Maruyama method.
 For specifics about the notation to define stochastic processes, please
 consult the :ref:`Arbor-specific NMODL extension <format-sde>`.
 
+.. note::
+
+   While the units of :math:`\textbf{f}(t, \textbf{X}(t))` represent the deterministic rate of
+   change (per millisecond),
+
+   .. math::
+
+        \left[\textbf{f}(t, \textbf{X}(t))\right] = \frac{\left[\textbf{X}(t)\right]}{ms},
+
+   the stochastic terms scale with :math:`t^{-1/2}`,
+
+   .. math::
+
+        \left[\textbf{l}_i(t, \textbf{X}(t))\right] = \frac{\left[\textbf{X}(t)\right]}{\sqrt{ms}}.
+
+
+**Example:** The popular Ornstein-Uhlenbeck process is described by a scalar linear mean-reverting SDE
+and can be written as
+
+.. math::
+
+    X^\prime = -\frac{1}{\tau} (X - \mu) + \sqrt{\frac{2}{τ}}  \sigma W,
+
+with white noise :math:`W`, and constant model parameters :math:`\tau`, :math:`\mu` and
+:math:`\sigma`. The relaxation time :math:`\tau` determines how fast the process reverts back to its
+mean value :math:`\mu`, and :math:`\sigma` controls the volatility (:math:`\mu` and :math:`\sigma`
+have the same units as :math:`X`). The expected value and variance can be computed analytically and
+yield
+
+.. math::
+
+    \begin{align*}
+    E[X]   &= \mu - \left( \mu - X_0\right) e^{-t/\tau}, \\
+    Var[X] &= \sigma^2 \left( 1 - e^{-2 t/\tau} \right),
+    \end{align*}
+
+which in the limit :math:`t \rightarrow \infty` converge to
+
+.. math::
+
+    \begin{align*}
+    E[X]   &= \mu, \\
+    Var[X] &= \sigma^2.
+    \end{align*}
 
 API
 ---

example/CMakeLists.txt

@@ -14,3 +14,4 @@ add_subdirectory(probe-demo)
 add_subdirectory(plasticity)
 add_subdirectory(lfp)
 add_subdirectory(diffusion)
+add_subdirectory(ornstein_uhlenbeck)

example/ornstein_uhlenbeck/CMakeLists.txt

@@ -0,0 +1,23 @@
+make_catalogue_standalone(
+    NAME ou
+    SOURCES "${CMAKE_CURRENT_SOURCE_DIR}"
+    MOD ornstein_uhlenbeck
+    CXX
+    CXX_FLAGS_TARGET ${ARB_CXX_FLAGS_TARGET_FULL}
+    VERBOSE ON)
+
+add_executable(ou EXCLUDE_FROM_ALL ou.cpp)
+add_dependencies(ou ou-catalogue)
+target_compile_options(ou PRIVATE ${ARB_CXX_FLAGS_TARGET_FULL})
+target_include_directories(ou PRIVATE "${CMAKE_CURRENT_BINARY_DIR}/generated/ou")
+
+target_link_libraries(ou PRIVATE arbor arborio ou-catalogue)
+if (ARB_USE_BUNDLED_FMT)
+    target_include_directories(ou PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/../../ext/fmt/include")
+    target_compile_definitions(ou PRIVATE FMT_HEADER_ONLY)
+else()
+    find_package(fmt REQUIRED)
+    target_link_libraries(ou PRIVATE fmt::fmt-header-only)
+endif()
+
+add_dependencies(examples ou)

example/ornstein_uhlenbeck/ornstein_uhlenbeck.mod

@@ -0,0 +1,38 @@
+: Ornstein-Uhlenbeck: linear mean reverting process
+: =================================================
+: dS(t) = kappa [mu - S(t)] dt + sigma dW(t)
+:     E[S(t)] = mu - (mu-S_0)*e^(-kappa t)
+:   Var[S(t)] = sigma^2/(2 kappa) [1 - e^(-2 kappa t)]
+: in the limit:
+:     lim t->∞   E[S(t)] = mu
+:     lim t->∞ Var[S(t)] = sigma^2/(2 kappa)
+
+NEURON {
+    SUFFIX ornstein_uhlenbeck
+    GLOBAL kappa, mu, sigma
+}
+
+PARAMETER {
+    kappa = 0.1
+    mu = 1
+    sigma = 0.1
+}
+
+STATE {
+    S
+}
+
+INITIAL {
+    S=2
+}
+
+BREAKPOINT {
+    SOLVE state METHOD stochastic
+}
+
+WHITE_NOISE { W }
+
+DERIVATIVE state {
+    S' = kappa*(mu - S) + sigma*W
+}
+

example/ornstein_uhlenbeck/ou.cpp

@@ -0,0 +1,181 @@
+#include <fmt/format.h>
+
+#include <arborio/label_parse.hpp>
+
+#include <arbor/assert.hpp>
+#include <arbor/recipe.hpp>
+#include <arbor/cable_cell_param.hpp>
+#include <arbor/cable_cell.hpp>
+#include <arbor/simulation.hpp>
+
+// forward declaration
+arb::mechanism_catalogue build_catalogue();
+
+// a single-cell recipe with probes
+class recipe: public arb::recipe {
+public:
+    recipe(unsigned ncvs) {
+        using namespace arb;
+        using namespace arborio::literals;
+
+        // build catalogue with stochastic mechanism
+        cell_gprop_.catalogue = global_default_catalogue();
+        cell_gprop_.catalogue.import(build_catalogue(), "");
+        cell_gprop_.default_parameters = neuron_parameter_defaults;
+
+        // paint the process on the whole cell
+        decor dec;
+        double const cv_size = 1.0;
+        dec.set_default(cv_policy_max_extent(cv_size));
+        dec.paint("(all)"_reg , density("hh"));
+        dec.paint("(all)"_reg , density("ornstein_uhlenbeck"));
+
+        // single-cell tree with ncvs control volumes
+        segment_tree tree;
+        tree.append(mnpos, {0, 0, 0.0, 4.0}, {0, 0, ncvs*cv_size, 4.0}, 1);
+        cell_ = cable_cell(morphology(tree), dec);
+    }
+
+    arb::cell_size_type num_cells() const override { return 1; }
+
+    arb::util::unique_any get_cell_description(arb::cell_gid_type) const override { return cell_; }
+
+    arb::cell_kind get_cell_kind(arb::cell_gid_type) const override { return arb::cell_kind::cable; }
+
+    std::vector<arb::probe_info> get_probes(arb::cell_gid_type) const override { return probes_; }
+
+    std::any get_global_properties(arb::cell_kind) const override { return cell_gprop_; }
+
+    void add_probe(arb::probe_tag tag, std::any address) { probes_.emplace_back(std::move(address), tag); }
+
+protected:
+    std::vector<arb::probe_info> probes_;
+    arb::cable_cell_global_properties cell_gprop_;
+    arb::cable_cell cell_;
+};
+
+// sampler for vector probes
+struct sampler {
+    sampler(std::vector<arb_value_type>& data, std::size_t n_cvs, std::size_t n_steps):
+     n_cvs_{n_cvs} , n_steps_{n_steps}, data_{data} {
+        data_.resize(n_cvs*n_steps);
+    }
+
+    void operator()(arb::probe_metadata pm, std::size_t n, const arb::sample_record* samples) {
+        const auto* m = arb::util::any_cast<const arb::mcable_list*>(pm.meta);
+        arb_assert(n_cvs_ == m->size());
+        arb_assert(n_steps_ == n);
+
+        for (std::size_t i=0; i<n; ++i) {
+            const auto* data = arb::util::any_cast<const arb::cable_sample_range*>(samples[i].data);
+            auto [lo, hi] = *data;
+            arb_assert(static_cast<std::size_t>(hi-lo) == n_cvs_);
+            for (std::size_t j=0; j<n_cvs_; ++j) {
+                data_[i*n_cvs_ + j] = lo[j];
+            }
+        }
+    }
+
+    std::size_t n_cvs_;
+    std::size_t n_steps_;
+    std::vector<arb_value_type>& data_;
+};
+
+// compute mean and variance online
+// uses Welford-Knuth algorithm for the variance
+struct accumulator {
+    std::size_t n_ = 0;
+    double mean_ = 0;
+    double var_ = 0;
+
+    accumulator& operator()(double sample) {
+        double const delta = sample - mean_;
+        mean_ += delta / (++n_);
+        var_ += delta * (sample - mean_);
+        return *this;
+    }
+
+    std::size_t n() const noexcept { return n_; }
+    double mean() const noexcept { return mean_; }
+    double variance() const noexcept { return n_ > 1 ? var_/(n_-1) : 0; }
+};
+
+int main(int argc, char** argv) {
+
+    unsigned ncvs = 5000;        // number of control volumes
+    double const dt = 1.0/1024;   // time step
+    unsigned nsteps = 500;       // number of time steps
+
+    // create recipe and add probes
+    recipe rec(ncvs);
+    rec.add_probe(1, arb::cable_probe_density_state_cell{"ornstein_uhlenbeck", "S"});
+
+    // make context and simulation objects
+    auto context = arb::make_context({1, -1});
+    arb::simulation sim = arb::simulation::create(rec)
+        .set_context(context)
+        .set_seed(137);
+
+    // setup sampler and add it to the simulation with regular schedule
+    std::vector<arb_value_type> data;
+    sampler s{data, ncvs, nsteps};
+    auto all_probes = [](arb::cell_member_type) { return true; };
+    sim.add_sampler(all_probes, arb::regular_schedule(dt), s, arb::sampling_policy::lax);
+
+    // run the simulation
+    sim.run(nsteps*dt, dt);
+
+    // evaluate the mean for each time step across the ensembe of realizations
+    // (each control volume is a independent realization of the Ornstein-Uhlenbeck process)
+    std::vector<accumulator> acc(nsteps);
+    for (std::size_t t=0; t<nsteps; ++t) {
+        for (std::size_t i=0; i<ncvs; ++i) {
+            acc[t](s.data_[t*ncvs+ i]);
+        }
+    }
+
+    // analytical solutions
+    auto expected = [](double t) -> std::pair<double,double> {
+        const double mu = 1.0;
+        const double S_0 = 2.0;
+        const double kappa = 0.1;
+        const double sigma = 0.1;
+        return {
+            mu - (mu-S_0)*std::exp(-kappa*t),
+            (sigma*sigma/(2*kappa))*(1.0 - std::exp(-2*kappa*t))
+        };
+    };
+
+    // print mean and expectation
+    for (std::size_t t=0; t<nsteps; t+=10) {
+        fmt::print("time = {:.5f}: mean = {:.5f} expected = {:.5f}\n",
+            dt*t, acc[t].mean(), expected(dt*t).first);
+    }
+    return 0;
+}
+
+// load mechanisms from library and add to new catalogue
+// =====================================================
+
+extern "C" {
+const void* get_catalogue(int*);
+}
+
+arb::mechanism_catalogue build_catalogue() {
+    arb::mechanism_catalogue cat;
+    int n=0;
+    const void* ptr = get_catalogue(&n);
+    const auto* mechs = reinterpret_cast<const arb_mechanism*>(ptr);
+    for (int i=0; i<n; ++i) {
+        const auto& mech = mechs[i];
+        auto ty = mech.type();
+        auto nm = ty.name;
+        auto ig = mech.i_gpu();
+        auto ic = mech.i_cpu();
+        arb_assert(ic || ig);
+        cat.add(nm, ty);
+        if (ic) cat.register_implementation(nm, std::make_unique<arb::mechanism>(ty, *ic));
+        if (ig) cat.register_implementation(nm, std::make_unique<arb::mechanism>(ty, *ig));
+    }
+    return cat;
+}

example/ornstein_uhlenbeck/readme.md

@@ -0,0 +1,7 @@
+## Miniapp for demonstrating stochastic processes
+
+The simulation consists of a single cell with a stochastic process (Ornstein-Uhlenbeck) painted on
+its control volumes. The stochastic process is described by a linear mean-reverting stochastic
+differential equation which is specified in the accompanying NMODL file. All processes start from
+the same initial condition and are averaged over the control volumes at each time step to generate an
+ensemble statistic. These results are then compared to the analytical solution.

mechanisms/CMakeLists.txt

@@ -21,6 +21,12 @@ make_catalogue(
   VERBOSE ${ARB_CAT_VERBOSE}
   ADD_DEPS ON)
 
+make_catalogue(
+  NAME stochastic
+  MOD ou_input
+  VERBOSE ${ARB_CAT_VERBOSE}
+  ADD_DEPS ON)
+
 # This re-exports
 set(arbor-builtin-mechanisms ${arbor-builtin-mechanisms} PARENT_SCOPE)