Reinitialize mechanisms after updating the ion state (#897)

If ion concentrations are read and written in the INITIAL blocks of mechanisms, this ensures that the reading mechanisms will receive the correct values.
Fixes #896

arbor/backends/gpu/shared_state.cpp

@@ -50,20 +50,19 @@ std::pair<fvm_value_type, fvm_value_type> minmax_value_impl(fvm_size_type n, con
 
 ion_state::ion_state(
     int charge,
-    const std::vector<fvm_index_type>& cv,
-    const std::vector<fvm_value_type>& init_iconc,
-    const std::vector<fvm_value_type>& init_econc,
-    const std::vector<fvm_value_type>& init_erev,
+    const fvm_ion_config& ion_data,
     unsigned // alignment/padding ignored.
 ):
-    node_index_(make_const_view(cv)),
-    iX_(cv.size(), NAN),
-    eX_(cv.size(), NAN),
-    Xi_(cv.size(), NAN),
-    Xo_(cv.size(), NAN),
-    init_Xi_(make_const_view(init_iconc)),
-    init_Xo_(make_const_view(init_econc)),
-    init_eX_(make_const_view(init_erev)),
+    node_index_(make_const_view(ion_data.cv)),
+    iX_(ion_data.cv.size(), NAN),
+    eX_(ion_data.cv.size(), NAN),
+    Xi_(ion_data.cv.size(), NAN),
+    Xo_(ion_data.cv.size(), NAN),
+    init_Xi_(make_const_view(ion_data.init_iconc)),
+    init_Xo_(make_const_view(ion_data.init_econc)),
+    reset_Xi_(make_const_view(ion_data.reset_iconc)),
+    reset_Xo_(make_const_view(ion_data.reset_econc)),
+    init_eX_(make_const_view(ion_data.init_revpot)),
     charge(1u, charge)
 {
     arb_assert(node_index_.size()==init_Xi_.size());
@@ -82,7 +81,8 @@ void ion_state::zero_current() {
 
 void ion_state::reset() {
     zero_current();
-    init_concentration();
+    memory::copy(reset_Xi_, Xi_);
+    memory::copy(reset_Xo_, Xo_);
     memory::copy(init_eX_, eX_);
 }
 
@@ -120,14 +120,11 @@ shared_state::shared_state(
 void shared_state::add_ion(
     const std::string& ion_name,
     int charge,
-    const std::vector<fvm_index_type>& cv,
-    const std::vector<fvm_value_type>& init_iconc,
-    const std::vector<fvm_value_type>& init_econc,
-    const std::vector<fvm_value_type>& init_erev)
+    const fvm_ion_config& ion_info)
 {
     ion_data.emplace(std::piecewise_construct,
         std::forward_as_tuple(ion_name),
-        std::forward_as_tuple(charge, cv, init_iconc, init_econc, init_erev, 1u));
+        std::forward_as_tuple(charge, ion_info, 1u));
 }
 
 void shared_state::reset() {

arbor/backends/gpu/shared_state.hpp

@@ -7,6 +7,8 @@
 
 #include <arbor/fvm_types.hpp>
 
+#include "fvm_layout.hpp"
+
 #include "backends/gpu/gpu_store_types.hpp"
 
 namespace arb {
@@ -33,6 +35,8 @@ struct ion_state {
 
     array init_Xi_;     // (mM) area-weighted initial internal concentration
     array init_Xo_;     // (mM) area-weighted initial external concentration
+    array reset_Xi_;    // (mM) area-weighted user-set internal concentration
+    array reset_Xo_;    // (mM) area-weighted user-set internal concentration
     array init_eX_;     // (mM) initial reversal potential
 
     array charge;       // charge of ionic species (global, length 1)
@@ -41,10 +45,7 @@ struct ion_state {
 
     ion_state(
         int charge,
-        const std::vector<fvm_index_type>& cv,
-        const std::vector<fvm_value_type>& init_Xi,
-        const std::vector<fvm_value_type>& init_Xo,
-        const std::vector<fvm_value_type>& init_eX,
+        const fvm_ion_config& ion_data,
         unsigned align
     );
 
@@ -97,10 +98,7 @@ struct shared_state {
     void add_ion(
         const std::string& ion_name,
         int charge,
-        const std::vector<fvm_index_type>& cv,
-        const std::vector<fvm_value_type>& init_iconc,
-        const std::vector<fvm_value_type>& init_econc,
-        const std::vector<fvm_value_type>& init_erev);
+        const fvm_ion_config& ion_data);
 
     void zero_currents();
 

arbor/backends/multicore/shared_state.cpp

@@ -48,21 +48,20 @@ using pad = util::padded_allocator<>;
 
 ion_state::ion_state(
     int charge,
-    const std::vector<fvm_index_type>& cv,
-    const std::vector<fvm_value_type>& init_Xi,
-    const std::vector<fvm_value_type>& init_Xo,
-    const std::vector<fvm_value_type>& init_eX,
+    const fvm_ion_config& ion_data,
     unsigned align
 ):
     alignment(min_alignment(align)),
-    node_index_(cv.begin(), cv.end(), pad(alignment)),
-    iX_(cv.size(), NAN, pad(alignment)),
-    eX_(init_eX.begin(), init_eX.end(), pad(alignment)),
-    Xi_(cv.size(), NAN, pad(alignment)),
-    Xo_(cv.size(), NAN, pad(alignment)),
-    init_Xi_(init_Xi.begin(), init_Xi.end(), pad(alignment)),
-    init_Xo_(init_Xo.begin(), init_Xo.end(), pad(alignment)),
-    init_eX_(init_eX.begin(), init_eX.end(), pad(alignment)),
+    node_index_(ion_data.cv.begin(), ion_data.cv.end(), pad(alignment)),
+    iX_(ion_data.cv.size(), NAN, pad(alignment)),
+    eX_(ion_data.init_revpot.begin(), ion_data.init_revpot.end(), pad(alignment)),
+    Xi_(ion_data.cv.size(), NAN, pad(alignment)),
+    Xo_(ion_data.cv.size(), NAN, pad(alignment)),
+    init_Xi_(ion_data.init_iconc.begin(), ion_data.init_iconc.end(), pad(alignment)),
+    init_Xo_(ion_data.init_econc.begin(), ion_data.init_econc.end(), pad(alignment)),
+    reset_Xi_(ion_data.reset_iconc.begin(), ion_data.reset_iconc.end(), pad(alignment)),
+    reset_Xo_(ion_data.reset_econc.begin(), ion_data.reset_econc.end(), pad(alignment)),
+    init_eX_(ion_data.init_revpot.begin(), ion_data.init_revpot.end(), pad(alignment)),
     charge(1u, charge, pad(alignment))
 {
     arb_assert(node_index_.size()==init_Xi_.size());
@@ -82,7 +81,8 @@ void ion_state::zero_current() {
 
 void ion_state::reset() {
     zero_current();
-    init_concentration();
+    std::copy(reset_Xi_.begin(), reset_Xi_.end(), Xi_.begin());
+    std::copy(reset_Xo_.begin(), reset_Xo_.end(), Xo_.begin());
     std::copy(init_eX_.begin(), init_eX_.end(), eX_.begin());
 }
 
@@ -134,14 +134,11 @@ shared_state::shared_state(
 void shared_state::add_ion(
     const std::string& ion_name,
     int charge,
-    const std::vector<fvm_index_type>& cv,
-    const std::vector<fvm_value_type>& init_iconc,
-    const std::vector<fvm_value_type>& init_econc,
-    const std::vector<fvm_value_type>& init_erev)
+    const fvm_ion_config& ion_info)
 {
     ion_data.emplace(std::piecewise_construct,
         std::forward_as_tuple(ion_name),
-        std::forward_as_tuple(charge, cv, init_iconc, init_econc, init_erev, alignment));
+        std::forward_as_tuple(charge, ion_info, alignment));
 }
 
 void shared_state::reset() {

arbor/backends/multicore/shared_state.hpp

@@ -20,6 +20,7 @@
 #include "multi_event_stream.hpp"
 #include "threshold_watcher.hpp"
 
+#include "fvm_layout.hpp"
 #include "multicore_common.hpp"
 
 namespace arb {
@@ -48,6 +49,8 @@ struct ion_state {
 
     array init_Xi_;         // (mM) area-weighted initial internal concentration
     array init_Xo_;         // (mM) area-weighted initial external concentration
+    array reset_Xi_;        // (mM) area-weighted user-set internal concentration
+    array reset_Xo_;        // (mM) area-weighted user-set internal concentration
     array init_eX_;         // (mV) initial reversal potential
 
     array charge;           // charge of ionic species (global value, length 1)
@@ -56,10 +59,7 @@ struct ion_state {
 
     ion_state(
         int charge,
-        const std::vector<fvm_index_type>& cv,
-        const std::vector<fvm_value_type>& init_Xi,
-        const std::vector<fvm_value_type>& init_Xo,
-        const std::vector<fvm_value_type>& init_eX,
+        const fvm_ion_config& ion_data,
         unsigned align
     );
 
@@ -115,10 +115,7 @@ struct shared_state {
     void add_ion(
         const std::string& ion_name,
         int charge,
-        const std::vector<fvm_index_type>& cv,
-        const std::vector<fvm_value_type>& init_iconc,
-        const std::vector<fvm_value_type>& init_econc,
-        const std::vector<fvm_value_type>& init_erev);
+        const fvm_ion_config& ion_data);
 
     void zero_currents();
 

arbor/fvm_layout.cpp

@@ -756,6 +756,8 @@ fvm_mechanism_data fvm_build_mechanism_data(const cable_cell_global_properties&
 
         ion_config.init_iconc.resize(ion_config.cv.size());
         ion_config.init_econc.resize(ion_config.cv.size());
+        ion_config.reset_iconc.resize(ion_config.cv.size());
+        ion_config.reset_econc.resize(ion_config.cv.size());
         ion_config.init_revpot.resize(ion_config.cv.size());
 
         for_each_cv_in_segments(D, keys(seg_ion_data),
@@ -767,6 +769,9 @@ fvm_mechanism_data fvm_build_mechanism_data(const cable_cell_global_properties&
                 auto& seg_ion_entry = seg_ion_map[seg];
 
                 value_type weight = area/D.cv_area[cv];
+                ion_config.reset_iconc[i] += weight*seg_ion_entry.ion_data.init_int_concentration;
+                ion_config.reset_econc[i] += weight*seg_ion_entry.ion_data.init_ext_concentration;
+
                 if (!seg_ion_entry.mech_writes_iconc) {
                     ion_config.init_iconc[i] += weight*seg_ion_entry.ion_data.init_int_concentration;
                 }

arbor/fvm_layout.hpp

@@ -136,8 +136,13 @@ struct fvm_ion_config {
     std::vector<value_type> init_iconc;
     std::vector<value_type> init_econc;
 
+    // Normalized area contribution of default concentration contribution in corresponding CV set by users
+    std::vector<value_type> reset_iconc;
+    std::vector<value_type> reset_econc;
+
     // Ion-specific (initial) reversal potential per CV.
     std::vector<value_type> init_revpot;
+
 };
 
 struct fvm_mechanism_data {

arbor/fvm_lowered_cell_impl.hpp

@@ -164,6 +164,18 @@ void fvm_lowered_cell_impl<Backend>::reset() {
 
     update_ion_state();
 
+    state_->zero_currents();
+
+    // Note: mechanisms must be initialized again after the ion state is updated,
+    // as mechanisms can read/write the ion_state within the initialize block
+    for (auto& m: revpot_mechanisms_) {
+        m->initialize();
+    }
+
+    for (auto& m: mechanisms_) {
+        m->initialize();
+    }
+
     // NOTE: Threshold watcher reset must come after the voltage values are set,
     // as voltage is implicitly read by watcher to set initial state.
     threshold_watcher_.reset();
@@ -422,7 +434,7 @@ void fvm_lowered_cell_impl<B>::initialize(
         const std::string& ion_name = i.first;
 
         if (auto charge = value_by_key(global_props.ion_species, ion_name)) {
-            state_->add_ion(ion_name, *charge, i.second.cv, i.second.init_iconc, i.second.init_econc, i.second.init_revpot);
+            state_->add_ion(ion_name, *charge, i.second);
         }
         else {
             throw cable_cell_error("unrecognized ion '"+ion_name+"' in mechanism");

test/unit/CMakeLists.txt

@@ -25,6 +25,8 @@ set(test_mechanisms
     read_eX
     write_multiple_eX
     write_eX
+    read_cai_init
+    write_cai_breakpoint
 )
 
 include(${PROJECT_SOURCE_DIR}/mechanisms/BuildModules.cmake)

test/unit/mod/read_cai_init.mod

+: Test mechanism with linear response to ica.
+
+NEURON {
+    SUFFIX read_cai_init
+    USEION ca READ cai VALENCE 2
+}
+
+PARAMETER {}
+
+ASSIGNED {}
+
+STATE {
+    s
+}
+
+INITIAL {
+    s = cai
+}
+
+BREAKPOINT {
+    SOLVE states
+}
+
+DERIVATIVE states {
+    s = cai
+}
+

test/unit/mod/write_cai_breakpoint.mod

+: Test mechanism with linear response to ica.
+
+NEURON {
+    SUFFIX write_cai_breakpoint
+    USEION ca WRITE cai READ ica VALENCE 2
+}
+
+PARAMETER {}
+
+ASSIGNED {}
+
+STATE {
+    cai
+}
+
+INITIAL {}
+
+BREAKPOINT {
+    SOLVE states
+}
+
+DERIVATIVE states {
+    cai = 5.2e-4
+}
\ No newline at end of file

test/unit/test_fvm_lowered.cpp

@@ -21,6 +21,7 @@
 #include "execution_context.hpp"
 #include "fvm_lowered_cell.hpp"
 #include "fvm_lowered_cell_impl.hpp"
+#include "mech_private_field_access.hpp"
 #include "sampler_map.hpp"
 #include "util/meta.hpp"
 #include "util/maputil.hpp"
@@ -543,6 +544,68 @@ TEST(fvm_lowered, read_valence) {
 }
 
 // Test correct scaling of ionic currents in reading and writing
+TEST(fvm_lowered, ionic_concentrations) {
+    auto cat = make_unit_test_catalogue();
+
+    // one cell, one CV:
+    fvm_size_type ncell = 1;
+    fvm_size_type ncv = 1;
+    std::vector<fvm_index_type> cv_to_intdom(ncv, 0);
+    std::vector<fvm_value_type> temp(ncv, 23);
+    std::vector<fvm_value_type> diam(ncv, 1.);
+    std::vector<fvm_value_type> vinit(ncv, -65);
+    std::vector<fvm_gap_junction> gj = {};
+
+    fvm_ion_config ion_config;
+    mechanism_layout layout;
+    mechanism_overrides overrides;
+
+    layout.weight.assign(ncv, 1.);
+    for (fvm_size_type i = 0; i<ncv; ++i) {
+        layout.cv.push_back(i);
+        ion_config.cv.push_back(i);
+    }
+    ion_config.init_revpot.assign(ncv, 0.);
+    ion_config.init_econc.assign(ncv, 0.);
+    ion_config.init_iconc.assign(ncv, 0.);
+    ion_config.reset_econc.assign(ncv, 0.);
+    ion_config.reset_iconc.assign(ncv, 2.3e-4);
+
+    auto read_cai  = cat.instance<backend>("read_cai_init");
+    auto write_cai = cat.instance<backend>("write_cai_breakpoint");
+
+    auto& read_cai_mech  = read_cai.mech;
+    auto& write_cai_mech = write_cai.mech;
+
+    auto shared_state = std::make_unique<typename backend::shared_state>(
+            ncell, cv_to_intdom, gj, vinit, temp, diam, read_cai_mech->data_alignment());
+    shared_state->add_ion("ca", 2, ion_config);
+
+    read_cai_mech->instantiate(0, *shared_state, overrides, layout);
+    write_cai_mech->instantiate(1, *shared_state, overrides, layout);
+
+    shared_state->reset();
+
+    // expect 2.3 value in state 's' in read_cai_init after init:
+    read_cai_mech->initialize();
+    write_cai_mech->initialize();
+
+    std::vector<fvm_value_type> expected_s_values(ncv, 2.3e-4);
+
+    EXPECT_EQ(expected_s_values, mechanism_field(read_cai_mech.get(), "s"));
+
+    // expect 5.2 + 2.3 value in state 's' in read_cai_init after state update:
+    read_cai_mech->nrn_state();
+    write_cai_mech->nrn_state();
+
+    read_cai_mech->write_ions();
+    write_cai_mech->write_ions();
+
+    read_cai_mech->nrn_state();
+
+    expected_s_values.assign(ncv, 7.5e-4);
+    EXPECT_EQ(expected_s_values, mechanism_field(read_cai_mech.get(), "s"));
+}
 
 TEST(fvm_lowered, ionic_currents) {
     soma_cell_builder b(6);

test/unit/unit_test_catalogue.cpp

@@ -31,6 +31,8 @@
 #include "mechanisms/read_eX.hpp"
 #include "mechanisms/write_multiple_eX.hpp"
 #include "mechanisms/write_eX.hpp"
+#include "mechanisms/read_cai_init.hpp"
+#include "mechanisms/write_cai_breakpoint.hpp"
 
 #include "../gtest.h"
 
@@ -74,6 +76,8 @@ mechanism_catalogue make_unit_test_catalogue() {
     ADD_MECH(cat, read_eX)
     ADD_MECH(cat, write_multiple_eX)
     ADD_MECH(cat, write_eX)
+    ADD_MECH(cat, read_cai_init)
+    ADD_MECH(cat, write_cai_breakpoint)
 
     return cat;
 }