diff --git a/modcc/kineticrewriter.cpp b/modcc/kineticrewriter.cpp
index 32b6d52a6424aca58c5bda5628db027272c00480..bb6ef0b77f3a6857d95521f4b91bed6bdfb21538 100644
--- a/modcc/kineticrewriter.cpp
+++ b/modcc/kineticrewriter.cpp
@@ -38,8 +38,8 @@ expression_ptr kinetic_rewrite(BlockExpression* block) {
// KineticRewriter implementation follows.
-void KineticRewriter::visit(ConserveExpression*) {
- // Deliberately ignoring these for now!
+void KineticRewriter::visit(ConserveExpression* e) {
+ statements_.push_back(e->clone());
}
void KineticRewriter::visit(ReactionExpression* e) {
diff --git a/modcc/msparse.hpp b/modcc/msparse.hpp
index 2af48cd53bb471b60bb80074f09c1b5dacdd203d..b6e2292cf9b87bf7fa7425fad110b9ea093dd9ac 100644
--- a/modcc/msparse.hpp
+++ b/modcc/msparse.hpp
@@ -96,6 +96,10 @@ public:
data_.push_back(e);
}
+ void clear() {
+ data_.clear();
+ }
+
// Return index into entry list which has column `c`.
unsigned index(unsigned c) const {
auto i = std::lower_bound(data_.begin(), data_.end(), c,
diff --git a/modcc/solvers.cpp b/modcc/solvers.cpp
index 25e4df5b7f45ecb99ab8ea51db0881b555f56113..96872a51e4def816d828094a6835d0618c14d93f 100644
--- a/modcc/solvers.cpp
+++ b/modcc/solvers.cpp
@@ -210,6 +210,11 @@ void SparseSolverVisitor::visit(AssignmentExpression *e) {
return;
}
+ if (conserve_ && !A_[deq_index_].empty()) {
+ deq_index_++;
+ return;
+ }
+
auto s = deriv->name();
auto expanded_rhs = substitute(rhs, local_expr_);
linear_test_result r = linear_test(expanded_rhs, dvars_);
@@ -221,6 +226,7 @@ void SparseSolverVisitor::visit(AssignmentExpression *e) {
// Populate sparse symbolic matrix for GE.
if (s!=dvars_[deq_index_]) {
error({"ICE: inconsistent ordering of derivative assignments", loc});
+ return;
}
auto dt_expr = make_expression<IdentifierExpression>(loc, "dt");
@@ -265,11 +271,81 @@ void SparseSolverVisitor::visit(AssignmentExpression *e) {
++deq_index_;
}
+void SparseSolverVisitor::visit(ConserveExpression *e) {
+ if (A_.empty()) {
+ unsigned n = dvars_.size();
+ A_ = symge::sym_matrix(n, n);
+ }
+ conserve_ = true;
+
+ auto loc = e->location();
+ scope_ptr scope = e->scope();
+
+ int row_idx;
+
+ // Find a row that contains one of the state variables in the conserve statement
+ auto& l = e->lhs()->is_stoich()->terms().front();
+ auto ident = l->is_stoich_term()->ident()->is_identifier();
+ if (ident) {
+ auto it = std::find(dvars_.begin(), dvars_.end(), ident->name());
+ if (it!=dvars_.end()) {
+ row_idx = it - dvars_.begin();
+ } else {
+ error({"CONSERVE statement unknown is not a state variable", loc});
+ return;
+ }
+ }
+ else {
+ error({"ICE: coefficient in state variable is not an identifier", loc});
+ return;
+ }
+
+ // Replace that row with the conserve statement
+ A_[row_idx].clear();
+
+ for (unsigned j = 0; j < dvars_.size(); ++j) {
+ auto state = dvars_[j];
+
+ auto& terms = e->lhs()->is_stoich()->terms();
+ auto it = std::find_if(terms.begin(), terms.end(), [&state](expression_ptr& p)
+ { return p->is_stoich_term()->ident()->is_identifier()->name() == state;});
+
+ if (it != terms.end()) {
+ auto expr = (*it)->is_stoich_term()->coeff()->clone();
+
+ auto local_a_term = make_unique_local_assign(scope, expr.get(), "a_");
+ auto a_ = local_a_term.id->is_identifier()->spelling();
+
+ statements_.push_back(std::move(local_a_term.local_decl));
+ statements_.push_back(std::move(local_a_term.assignment));
+
+ A_[row_idx].push_back({j, symtbl_.define(a_)});
+ }
+ }
+
+
+ expression_ptr expr = e->rhs()->clone();
+ auto local_a_term = make_unique_local_assign(scope, expr.get(), "a_");
+ auto a_ = local_a_term.id->is_identifier()->spelling();
+
+ statements_.push_back(std::move(local_a_term.local_decl));
+ statements_.push_back(std::move(local_a_term.assignment));
+
+ conserve_rhs_.push_back(a_);
+ conserve_idx_.push_back(row_idx);
+
+}
+
void SparseSolverVisitor::finalize() {
std::vector<symge::symbol> rhs;
for (const auto& var: dvars_) {
rhs.push_back(symtbl_.define(var));
}
+ if (conserve_) {
+ for (unsigned i = 0; i < conserve_idx_.size(); ++i) {
+ rhs[conserve_idx_[i]] = symtbl_.define(conserve_rhs_[i]);
+ }
+ }
A_.augment(rhs);
symge::gj_reduce(A_, symtbl_);
diff --git a/modcc/solvers.hpp b/modcc/solvers.hpp
index 33f5b00bb50fe9b94e0ff17651151d5e6d145fbe..e31084dfc0a87137ca0d46c276543d076aa08087 100644
--- a/modcc/solvers.hpp
+++ b/modcc/solvers.hpp
@@ -82,6 +82,12 @@ protected:
// 'Symbol table' for symbolic manipulation.
symge::symbol_table symtbl_;
+ // Flag to indicate whether conserve statements are part of the system
+ bool conserve_ = false;
+
+ // rhs of conserve statement
+ std::vector<std::string> conserve_rhs_;
+ std::vector<unsigned> conserve_idx_;
public:
using SolverVisitorBase::visit;
@@ -90,11 +96,15 @@ public:
virtual void visit(BlockExpression* e) override;
virtual void visit(AssignmentExpression *e) override;
+ virtual void visit(ConserveExpression *e) override;
virtual void finalize() override;
virtual void reset() override {
deq_index_ = 0;
local_expr_.clear();
symtbl_.clear();
+ conserve_rhs_.clear();
+ conserve_idx_.clear();
+ conserve_ = false;
SolverVisitorBase::reset();
}
};
diff --git a/test/unit/CMakeLists.txt b/test/unit/CMakeLists.txt
index 45d8b30b01036e3690f5f8b69d476d68d9f8783f..2fb9912b8db2bb04f4eb0af11e7ebb4271e390a1 100644
--- a/test/unit/CMakeLists.txt
+++ b/test/unit/CMakeLists.txt
@@ -2,6 +2,10 @@
set(test_mechanisms
celsius_test
+ test0_kin_diff
+ test0_kin_conserve
+ test1_kin_diff
+ test1_kin_conserve
fixed_ica_current
point_ica_current
linear_ca_conc
@@ -79,6 +83,7 @@ set(unit_sources
test_fvm_lowered.cpp
test_glob_basic.cpp
test_mc_cell_group.cpp
+ test_kinetic.cpp
test_lexcmp.cpp
test_lif_cell_group.cpp
test_maputil.cpp
diff --git a/test/unit/mod/test0_kin_conserve.mod b/test/unit/mod/test0_kin_conserve.mod
new file mode 100644
index 0000000000000000000000000000000000000000..f6a2b8ad894bf64ca62bb31e7ece53461468406b
--- /dev/null
+++ b/test/unit/mod/test0_kin_conserve.mod
@@ -0,0 +1,30 @@
+NEURON {
+ SUFFIX test0_kin_conserve
+}
+
+STATE {
+ s d h
+}
+
+BREAKPOINT {
+ SOLVE state METHOD sparse
+}
+
+KINETIC state {
+ LOCAL alpha1, beta1, alpha2, beta2
+ alpha1 = 2
+ beta1 = 0.6
+ alpha2 = 3
+ beta2 = 0.7
+
+ ~ s <-> h (alpha1, beta1)
+ ~ d <-> s (alpha2, beta2)
+
+ CONSERVE s + d + h = 1
+}
+
+INITIAL {
+ h = 0.2
+ d = 0.3
+ s = 1-d-h
+}
diff --git a/test/unit/mod/test0_kin_diff.mod b/test/unit/mod/test0_kin_diff.mod
new file mode 100644
index 0000000000000000000000000000000000000000..1de564f84d167bb1725e8002af8973e6dc12607e
--- /dev/null
+++ b/test/unit/mod/test0_kin_diff.mod
@@ -0,0 +1,28 @@
+NEURON {
+ SUFFIX test0_kin_diff
+}
+
+STATE {
+ s d h
+}
+
+BREAKPOINT {
+ SOLVE state METHOD sparse
+}
+
+KINETIC state {
+ LOCAL alpha1, beta1, alpha2, beta2
+ alpha1 = 2
+ beta1 = 0.6
+ alpha2 = 3
+ beta2 = 0.7
+
+ ~ s <-> h (alpha1, beta1)
+ ~ d <-> s (alpha2, beta2)
+}
+
+INITIAL {
+ h = 0.2
+ d = 0.3
+ s = 1-d-h
+}
diff --git a/test/unit/mod/test1_kin_conserve.mod b/test/unit/mod/test1_kin_conserve.mod
new file mode 100644
index 0000000000000000000000000000000000000000..b754f39256823f08bccfa394da0b34977359b09f
--- /dev/null
+++ b/test/unit/mod/test1_kin_conserve.mod
@@ -0,0 +1,32 @@
+NEURON {
+ SUFFIX test1_kin_conserve
+}
+
+STATE {
+ a b x y
+}
+
+BREAKPOINT {
+ SOLVE state METHOD sparse
+}
+
+KINETIC state {
+ LOCAL alpha1, beta1, alpha2, beta2
+ alpha1 = 2
+ beta1 = 0.6
+ alpha2 = 3
+ beta2 = 0.7
+
+ ~ a <-> b (alpha1, beta1)
+ ~ x <-> y (alpha2, beta2)
+
+ CONSERVE a + b = 1
+ CONSERVE x + y = 1
+}
+
+INITIAL {
+ a = 0.2
+ b = 1 - a
+ x = 0.6
+ y = 1 - x
+}
diff --git a/test/unit/mod/test1_kin_diff.mod b/test/unit/mod/test1_kin_diff.mod
new file mode 100644
index 0000000000000000000000000000000000000000..e27b80b56bf38c99b32cfc3e795da5cde1794143
--- /dev/null
+++ b/test/unit/mod/test1_kin_diff.mod
@@ -0,0 +1,29 @@
+NEURON {
+ SUFFIX test1_kin_diff
+}
+
+STATE {
+ a b x y
+}
+
+BREAKPOINT {
+ SOLVE state METHOD sparse
+}
+
+KINETIC state {
+ LOCAL alpha, beta, gamma, delta
+ alpha = 2
+ beta = 0.6
+ gamma = 3
+ delta = 0.7
+
+ ~ a <-> b (alpha, beta)
+ ~ x <-> y (gamma, delta)
+}
+
+INITIAL {
+ a = 0.2
+ b = 1 - a
+ x = 0.6
+ y = 1 - x
+}
diff --git a/test/unit/test_kinetic.cpp b/test/unit/test_kinetic.cpp
new file mode 100644
index 0000000000000000000000000000000000000000..71f1a636f03cea4998d895aa4d460e58e2f99479
--- /dev/null
+++ b/test/unit/test_kinetic.cpp
@@ -0,0 +1,117 @@
+#include <vector>
+
+#include <arbor/mechanism.hpp>
+#include <arbor/version.hpp>
+
+#include "backends/multicore/fvm.hpp"
+
+#ifdef ARB_GPU_ENABLED
+#include "backends/gpu/fvm.hpp"
+#endif
+
+#include "common.hpp"
+#include "mech_private_field_access.hpp"
+#include "fvm_lowered_cell.hpp"
+#include "fvm_lowered_cell_impl.hpp"
+#include "sampler_map.hpp"
+#include "simple_recipes.hpp"
+#include "unit_test_catalogue.hpp"
+
+using namespace arb;
+
+using backend = arb::multicore::backend;
+using fvm_cell = arb::fvm_lowered_cell_impl<backend>;
+
+using shared_state = backend::shared_state;
+ACCESS_BIND(std::unique_ptr<shared_state> fvm_cell::*, private_state_ptr, &fvm_cell::state_)
+
+template <typename backend>
+void run_kinetic_test(std::string mech_name,
+ std::vector<std::string> variables,
+ std::vector<fvm_value_type> t0_values,
+ std::vector<fvm_value_type> t1_values) {
+
+ auto cat = make_unit_test_catalogue();
+
+ fvm_size_type ncell = 1;
+ fvm_size_type ncv = 1;
+ std::vector<fvm_index_type> cv_to_intdom(ncv, 0);
+
+ std::vector<fvm_gap_junction> gj = {};
+ auto instance = cat.instance<backend>(mech_name);
+ auto& kinetic_test = instance.mech;
+
+ std::vector<fvm_value_type> temp(ncv, 300.);
+ std::vector<fvm_value_type> vinit(ncv, -65);
+
+ auto shared_state = std::make_unique<typename backend::shared_state>(
+ ncell, cv_to_intdom, gj, vinit, temp, kinetic_test->data_alignment());
+
+ 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);
+ }
+
+ kinetic_test->instantiate(0, *shared_state, overrides, layout);
+ shared_state->reset();
+
+ kinetic_test->initialize();
+
+ for (unsigned i = 0; i < variables.size(); i++) {
+ for (unsigned j = 0; j < ncv; j++) {
+ EXPECT_NEAR(t0_values[i], mechanism_field(kinetic_test.get(), variables[i]).at(j), 1e-6);
+ }
+ }
+
+ shared_state->update_time_to(0.5, 0.5);
+ shared_state->set_dt();
+
+ kinetic_test->nrn_state();
+
+ for (unsigned i = 0; i < variables.size(); i++) {
+ for (unsigned j = 0; j < ncv; j++) {
+ EXPECT_NEAR(t1_values[i], mechanism_field(kinetic_test.get(), variables[i]).at(j), 1e-6);
+ }
+ }
+}
+
+TEST(mech_kinetic, kinetic_1_conserve) {
+ std::vector<std::string> variables = {"s", "h", "d"};
+ std::vector<fvm_value_type> t0_values = {0.5, 0.2, 0.3};
+ std::vector<fvm_value_type> t1_values = {0.380338, 0.446414, 0.173247};
+
+ run_kinetic_test<multicore::backend>("test0_kin_diff", variables, t0_values, t1_values);
+ run_kinetic_test<multicore::backend>("test0_kin_conserve", variables, t0_values, t1_values);
+}
+
+TEST(mech_kinetic, kinetic_2_conserve) {
+ std::vector<std::string> variables = {"a", "b", "x", "y"};
+ std::vector<fvm_value_type> t0_values = {0.2, 0.8, 0.6, 0.4};
+ std::vector<fvm_value_type> t1_values = {0.217391304, 0.782608696, 0.33333333, 0.66666666};
+
+ run_kinetic_test<multicore::backend>("test1_kin_diff", variables, t0_values, t1_values);
+ run_kinetic_test<multicore::backend>("test1_kin_conserve", variables, t0_values, t1_values);
+}
+
+#ifdef ARB_GPU_ENABLED
+TEST(mech_kinetic_gpu, kinetic_1_conserve) {
+ std::vector<std::string> variables = {"s", "h", "d"};
+ std::vector<fvm_value_type> t0_values = {0.5, 0.2, 0.3};
+ std::vector<fvm_value_type> t1_values = {0.380338, 0.446414, 0.173247};
+
+ run_kinetic_test<gpu::backend>("test0_kin_diff", variables, t0_values, t1_values);
+ run_kinetic_test<gpu::backend>("test0_kin_conserve", variables, t0_values, t1_values);
+}
+
+TEST(mech_kinetic_gpu, kinetic_2_conserve) {
+ std::vector<std::string> variables = {"a", "b", "x", "y"};
+ std::vector<fvm_value_type> t0_values = {0.2, 0.8, 0.6, 0.4};
+ std::vector<fvm_value_type> t1_values = {0.217391304, 0.782608696, 0.33333333, 0.66666666};
+
+ run_kinetic_test<gpu::backend>("test1_kin_diff", variables, t0_values, t1_values);
+ run_kinetic_test<gpu::backend>("test1_kin_conserve", variables, t0_values, t1_values);
+}
+#endif
diff --git a/test/unit/unit_test_catalogue.cpp b/test/unit/unit_test_catalogue.cpp
index fc4987662ab6ead860b7af4044fac0072851613d..ccd007e3591192edad4139097df90dc8cfd18dd9 100644
--- a/test/unit/unit_test_catalogue.cpp
+++ b/test/unit/unit_test_catalogue.cpp
@@ -8,6 +8,10 @@
#include "unit_test_catalogue.hpp"
#include "mechanisms/celsius_test.hpp"
+#include "mechanisms/test0_kin_diff.hpp"
+#include "mechanisms/test0_kin_conserve.hpp"
+#include "mechanisms/test1_kin_diff.hpp"
+#include "mechanisms/test1_kin_conserve.hpp"
#include "mechanisms/fixed_ica_current.hpp"
#include "mechanisms/point_ica_current.hpp"
#include "mechanisms/linear_ca_conc.hpp"
@@ -36,6 +40,10 @@ mechanism_catalogue make_unit_test_catalogue() {
mechanism_catalogue cat;
ADD_MECH(cat, celsius_test)
+ ADD_MECH(cat, test0_kin_diff)
+ ADD_MECH(cat, test0_kin_conserve)
+ ADD_MECH(cat, test1_kin_diff)
+ ADD_MECH(cat, test1_kin_conserve)
ADD_MECH(cat, fixed_ica_current)
ADD_MECH(cat, point_ica_current)
ADD_MECH(cat, linear_ca_conc)