Waveform Relaxation for Gap Junctions

arbor/backends/multicore/shared_state.cpp

@@ -477,6 +477,7 @@ void shared_state::instantiate(arb::mechanism& m, unsigned id, const mechanism_o
     m.ppack_.vec_di           = cv_to_intdom.data();
     m.ppack_.vec_dt           = dt_cv.data();
     m.ppack_.vec_v            = voltage.data();
+    m.ppack_.vec_v_peer       = voltage.data();
     m.ppack_.vec_i            = current_density.data();
     m.ppack_.vec_g            = conductivity.data();
     m.ppack_.temperature_degC = temperature_degC.data();

arbor/fvm_lowered_cell_impl.hpp

@@ -7,6 +7,7 @@
 // implementation details may be tested in the unit tests.
 // It should otherwise only be used in `fvm_lowered_cell.cpp`.
 
+#include <array>
 #include <cmath>
 #include <iterator>
 #include <memory>
@@ -188,6 +189,7 @@ void fvm_lowered_cell_impl<Backend>::reset() {
     threshold_watcher_.reset();
 }
 
+
 template <typename Backend>
 fvm_integration_result fvm_lowered_cell_impl<Backend>::integrate(
     value_type tfinal,
@@ -218,116 +220,199 @@ fvm_integration_result fvm_lowered_cell_impl<Backend>::integrate(
     // per-compartment dt probably not a win on GPU), possibly rumbling
     // complete fvm state into shared state object.
 
-    while (remaining_steps) {
-        // Update any required reversal potentials based on ionic concs.
+    // Assume that: 
+    // - all gap junctions are subject to WR
+    // - we only have a single cell group
 
-        for (auto& m: revpot_mechanisms_) {
-            m->update_current();
-        }
+    //traces:            gj mech id                                values               
+    std::unordered_map<arb_index_type, std::vector<std::vector<arb_value_type>>> traces_v, traces_v_prev;
+
+    //Map m->ppack_.peer_index to corresponding peer voltage in traces_v_prev
+    std::vector<arb_index_type> peer_ix;
 
-        // Deliver events and accumulate mechanism current contributions.
-
-        PE(advance_integrate_events);
-        state_->deliverable_events.mark_until_after(state_->time);
-        PL();
-
-        PE(advance_integrate_current_zero);
-        state_->zero_currents();
-        PL();
-        for (auto& m: mechanisms_) {
-            auto state = state_->deliverable_events.marked_events();
-            arb_deliverable_event_stream events;
-            events.n_streams = state.n;
-            events.begin     = state.begin_offset;
-            events.end       = state.end_offset;
-            events.events    = (arb_deliverable_event_data*) state.ev_data; // FIXME(TH): This relies on bit-castability
-            m->deliver_events(events);
-            m->update_current();
-        }
+    auto max_steps = remaining_steps;
+    value_type tmin_reset = tmin_;
 
-        PE(advance_integrate_events);
-        state_->deliverable_events.drop_marked_events();
+    //WR iterations
+    int wr_max = 1;
+    auto eps = 1e-15;
+    for (int wr_it = 0; wr_it < wr_max; wr_it++){
 
-        // Update event list and integration step times.
+        //Reset remaining_steps
+        if (wr_it > 0) {
+            remaining_steps = dt_steps(tmin_reset, tfinal, dt_max);
+        }
+        //Reset error variables for WR break condition
+        std::vector<arb_value_type> err(peer_ix.size(), 0.);
+        auto b = 0;
+
+        while (remaining_steps) {
+            auto step = max_steps-remaining_steps;
+
+            // Update any required reversal potentials based on ionic concs.
+            for (auto& m: revpot_mechanisms_) {
+                m->update_current();
+            }
 
-        state_->update_time_to(dt_max, tfinal);
-        state_->deliverable_events.event_time_if_before(state_->time_to);
-        state_->set_dt();
-        PL();
+            // Deliver events and accumulate mechanism current contributions.
 
-        // Add stimulus current contributions.
-        // (Note: performed after dt, time_to calculation, in case we
-        // want to use mean current contributions as opposed to point
-        // sample.)
+            PE(advance_integrate_events);
+            state_->deliverable_events.mark_until_after(state_->time);
+            PL();
 
-        PE(advance_integrate_stimuli)
-        state_->add_stimulus_current();
-        PL();
+            PE(advance_integrate_current_zero);
+            state_->zero_currents();
+            PL();
+            for (auto& m: mechanisms_) {
+                auto state = state_->deliverable_events.marked_events();
+                arb_deliverable_event_stream events;
+                events.n_streams = state.n;
+                events.begin     = state.begin_offset;
+                events.end       = state.end_offset;
+                events.events    = (arb_deliverable_event_data*) state.ev_data; // FIXME(TH): This relies on bit-castability
+                m->deliver_events(events);
+
+                //Feed back previous traces to vec_v_peer
+                //vec_v_peer defaults to vec_v
+                if (m->kind() == arb_mechanism_kind_gap_junction) {
+                    if (wr_it == 0 && step == 0) {
+                        std::unordered_map<arb_index_type, std::deque<arb_index_type>> peer_pos;
+                        for (auto ix = 0; ix < m->ppack_.width; ++ix) {
+                            peer_pos[m->ppack_.node_index[ix]].push_back(ix);
+                        }
+                        for (auto it = 0; it < m->ppack_.width; ++it) {
+                            auto p = m->ppack_.peer_index[it];
+                            peer_ix.push_back(peer_pos[p][0]);
+                            peer_pos[p].pop_front();
+                        }
+                    }
+                    if (wr_it > 0) {
+                        auto gj = m->mechanism_id();
+                        for (auto i = 0; i < m->ppack_.width; ++i){ 
+                            auto peer = peer_ix[i];
+                            auto v_prev = traces_v_prev[gj][step][peer];
+                            m->ppack_.vec_v_peer[i] = v_prev;
+                        }
+                    }
+                }
+
+                m->update_current();
+
+                //update traces
+                if (m->kind() == arb_mechanism_kind_gap_junction) {
+                    auto gj = m->mechanism_id();
+                    std::vector<value_type> v_step = {};
+                    std::vector<value_type> t_step = {};
+                    for(int ix = 0; ix < m->ppack_.width; ++ix) {
+                        auto node_cv = m->ppack_.node_index[ix];
+                        v_step.push_back(state_->voltage[node_cv]);
+                        t_step.push_back(state_->dt_cv[node_cv]);
+
+                        auto err_cv = state_->voltage[node_cv] - traces_v_prev[gj][step][node_cv];
+                        err[node_cv] += (err_cv*err_cv);
+                        if (err[node_cv] > eps) {
+                            b += 1;
+                        }
+                    }
+                    traces_v[gj].push_back(v_step);
+                }
+            }
+
 
-        // Take samples at cell time if sample time in this step interval.
+            PE(advance_integrate_events);
+            state_->deliverable_events.drop_marked_events();
 
-        PE(advance_integrate_samples);
-        sample_events_.mark_until(state_->time_to);
-        state_->take_samples(sample_events_.marked_events(), sample_time_, sample_value_);
-        sample_events_.drop_marked_events();
-        PL();
+            // Update event list and integration step times.
 
-        // Integrate voltage by matrix solve.
+            state_->update_time_to(dt_max, tfinal);
+            state_->deliverable_events.event_time_if_before(state_->time_to);
+            state_->set_dt();
+            PL();
 
-        PE(advance_integrate_matrix_build);
-        matrix_.assemble(state_->dt_intdom, state_->voltage, state_->current_density, state_->conductivity);
-        PL();
-        PE(advance_integrate_matrix_solve);
-        matrix_.solve(state_->voltage);
-        PL();
+            // Add stimulus current contributions.
+            // (Note: performed after dt, time_to calculation, in case we
+            // want to use mean current contributions as opposed to point
+            // sample.)
 
-        // Integrate mechanism state.
+            PE(advance_integrate_stimuli)
+            state_->add_stimulus_current();
+            PL();
 
-        for (auto& m: mechanisms_) {
-            m->update_state();
-        }
+            // Take samples at cell time if sample time in this step interval.
 
-        // Update ion concentrations.
+            PE(advance_integrate_samples);
+            sample_events_.mark_until(state_->time_to);
+            state_->take_samples(sample_events_.marked_events(), sample_time_, sample_value_);
+            sample_events_.drop_marked_events();
+            PL();
 
-        PE(advance_integrate_ionupdate);
-        update_ion_state();
-        PL();
+            // Integrate voltage by matrix solve.
 
-        // Update time and test for spike threshold crossings.
+            PE(advance_integrate_matrix_build);
+            matrix_.assemble(state_->dt_intdom, state_->voltage, state_->current_density, state_->conductivity);
+            PL();
+            PE(advance_integrate_matrix_solve);
+            matrix_.solve(state_->voltage);
+            PL();
 
-        PE(advance_integrate_threshold);
-        threshold_watcher_.test(&state_->time_since_spike);
-        PL();
+            // Integrate mechanism state.
 
-        PE(advance_integrate_post)
-        if (post_events_) {
             for (auto& m: mechanisms_) {
-                m->post_event();
+                m->update_state();
             }
-        }
-        PL();
 
-        std::swap(state_->time_to, state_->time);
-        state_->time_ptr = state_->time.data();
+            // Update ion concentrations.
+
+            PE(advance_integrate_ionupdate);
+            update_ion_state();
+            PL();
 
-        // Check for non-physical solutions:
+            // Update time and test for spike threshold crossings.
 
-        if (check_voltage_mV_>0) {
-            PE(advance_integrate_physicalcheck);
-            assert_voltage_bounded(check_voltage_mV_);
+            PE(advance_integrate_threshold);
+            threshold_watcher_.test(&state_->time_since_spike);
             PL();
-        }
 
-        // Check for end of integration.
+            PE(advance_integrate_post)
+            if (post_events_) {
+                for (auto& m: mechanisms_) {
+                    m->post_event();
+                }
+            }
+            PL();
+
+            std::swap(state_->time_to, state_->time);
+            state_->time_ptr = state_->time.data();
+
+            // Check for non-physical solutions:
+
+            if (check_voltage_mV_>0) {
+                PE(advance_integrate_physicalcheck);
+                assert_voltage_bounded(check_voltage_mV_);
+                PL();
+            }
+
+            // Check for end of integration.
+
+            PE(advance_integrate_stepsupdate);
+            if (!--remaining_steps) {
+                tmin_ = state_->time_bounds().first;
+                remaining_steps = dt_steps(tmin_, tfinal, dt_max);
+            }
+            PL();
 
-        PE(advance_integrate_stepsupdate);
-        if (!--remaining_steps) {
-            tmin_ = state_->time_bounds().first;
-            remaining_steps = dt_steps(tmin_, tfinal, dt_max);
         }
-        PL();
-    }
 
+        //break WR if difference between previous and current trace small enough
+        if (wr_it > 0 && b == 0) {
+            break;
+        }
+
+        //reset traces
+        traces_v_prev = traces_v;
+        traces_v = {};
+    }
+
     set_tmin(tfinal);
 
     const auto& crossings = threshold_watcher_.crossings();

arbor/include/arbor/mechanism_abi.h

@@ -86,6 +86,7 @@ typedef struct arb_mechanism_ppack {
     const arb_value_type* vec_t;
     arb_value_type* vec_dt;
     arb_value_type* vec_v;
+    arb_value_type* vec_v_peer;                  // Other side of gap junction, default to vec_v.
     arb_value_type* vec_i;
     arb_value_type* vec_g;
     arb_value_type* temperature_degC;

modcc/printer/cprinter.cpp

@@ -273,6 +273,7 @@ std::string emit_cpp_source(const Module& module_, const printer_options& opt) {
                                    "[[maybe_unused]] auto* {0}vec_t             = pp->vec_t;\\\n"
                                    "[[maybe_unused]] auto* {0}vec_dt            = pp->vec_dt;\\\n"
                                    "[[maybe_unused]] auto* {0}vec_v             = pp->vec_v;\\\n"
+                                   "[[maybe_unused]] auto* {0}vec_v_peer        = pp->vec_v_peer;\\\n"
                                    "[[maybe_unused]] auto* {0}vec_i             = pp->vec_i;\\\n"
                                    "[[maybe_unused]] auto* {0}vec_g             = pp->vec_g;\\\n"
                                    "[[maybe_unused]] auto* {0}temperature_degC  = pp->temperature_degC;\\\n"

modcc/printer/printerutil.cpp

@@ -155,7 +155,7 @@ indexed_variable_info decode_indexed_variable(IndexedVariable* sym) {
         v.readonly = true;
         break;
     case sourceKind::peer_voltage:
-        v.data_var="vec_v";
+        v.data_var="vec_v_peer";
         v.other_index_var = "peer_index";
         v.node_index_var = "";
         v.index_var_kind = index_kind::other;