Merge branch 'master' into jelic/cvode_visitors

docs/contents/longitudinal_diffusion.rst

@@ -0,0 +1,33 @@
+Longitudinal Diffusion
+======================
+
+The idea behind ``LONGITUDINAL_DIFFUSION`` is to allow a ``STATE`` variable to
+diffuse along a section, i.e. from one segment into a neighbouring segment.
+
+This problem is solved by registering callbacks. In particular, NEURON needs to
+be informed of the volume and diffusion rate. Additionally, the implicit
+time-stepping requires information about certain derivatives.
+
+Implementation in NMODL
+-----------------------
+
+The following ``KINETIC`` block
+
+  .. code-block::
+
+    KINETIC state {
+        COMPARTMENT vol {X}
+        LONGITUDINAL_DIFFUSION mu {X}
+
+        ~ X << (ica)
+    }
+
+Will undergo two transformations. The first is to create a system of ODEs that
+can be solved. This consumed the AST node. However, to print the code for
+longitudinal diffusion we require information from the ``COMPARTMENT`` and
+``LONGITUDINAL_DIFFUSION`` statements. This is why there's a second
+transformation, that runs before the other transformation, to extract the
+required information and store it a AST node called
+``LONGITUDINAL_DIFFUSION_BLOCK``. This block can then be converted into an
+"info" object, which is then used to print the callbacks.
+

docs/index.rst

@@ -24,6 +24,7 @@ About NMODL
    contents/pointers
    contents/cable_equations
    contents/globals
+   contents/longitudinal_diffusion
    contents/cvode
 
 .. toctree::

src/codegen/codegen_helper_visitor.cpp

@@ -9,8 +9,10 @@
 
 #include <algorithm>
 #include <cmath>
+#include <memory>
 
 #include "ast/all.hpp"
+#include "ast/constant_var.hpp"
 #include "codegen/codegen_naming.hpp"
 #include "parser/c11_driver.hpp"
 #include "visitors/visitor_utils.hpp"
@@ -853,5 +855,63 @@ void CodegenHelperVisitor::visit_after_block(const ast::AfterBlock& node) {
     info.before_after_blocks.push_back(&node);
 }
 
+static std::shared_ptr<ast::Compartment> find_compartment(
+    const ast::LongitudinalDiffusionBlock& node,
+    const std::string& var_name) {
+    const auto& compartment_block = node.get_compartment_statements();
+    for (const auto& stmt: compartment_block->get_statements()) {
+        auto comp = std::dynamic_pointer_cast<ast::Compartment>(stmt);
+
+        auto species = comp->get_species();
+        auto it = std::find_if(species.begin(), species.end(), [&var_name](auto var) {
+            return var->get_node_name() == var_name;
+        });
+
+        if (it != species.end()) {
+            return comp;
+        }
+    }
+
+    return nullptr;
+}
+
+void CodegenHelperVisitor::visit_longitudinal_diffusion_block(
+    const ast::LongitudinalDiffusionBlock& node) {
+    auto longitudinal_diffusion_block = node.get_longitudinal_diffusion_statements();
+    for (auto stmt: longitudinal_diffusion_block->get_statements()) {
+        auto diffusion = std::dynamic_pointer_cast<ast::LonDiffuse>(stmt);
+        auto rate_index_name = diffusion->get_index_name();
+        auto rate_expr = diffusion->get_rate();
+        auto species = diffusion->get_species();
+
+        auto process_compartment = [](const std::shared_ptr<ast::Compartment>& compartment)
+            -> std::pair<std::shared_ptr<ast::Name>, std::shared_ptr<ast::Expression>> {
+            std::shared_ptr<ast::Expression> volume_expr;
+            std::shared_ptr<ast::Name> volume_index_name;
+            if (!compartment) {
+                volume_index_name = nullptr;
+                volume_expr = std::make_shared<ast::Double>("1.0");
+            } else {
+                volume_index_name = compartment->get_index_name();
+                volume_expr = std::shared_ptr<ast::Expression>(compartment->get_volume()->clone());
+            }
+            return {std::move(volume_index_name), std::move(volume_expr)};
+        };
+
+        for (auto var: species) {
+            std::string state_name = var->get_value()->get_value();
+            auto compartment = find_compartment(node, state_name);
+            auto [volume_index_name, volume_expr] = process_compartment(compartment);
+
+            info.longitudinal_diffusion_info.insert(
+                {state_name,
+                 LongitudinalDiffusionInfo(volume_index_name,
+                                           std::shared_ptr<ast::Expression>(volume_expr),
+                                           rate_index_name,
+                                           std::shared_ptr<ast::Expression>(rate_expr->clone()))});
+        }
+    }
+}
+
 }  // namespace codegen
 }  // namespace nmodl

src/codegen/codegen_helper_visitor.hpp

@@ -115,6 +115,7 @@ class CodegenHelperVisitor: public visitor::ConstAstVisitor {
     void visit_verbatim(const ast::Verbatim& node) override;
     void visit_before_block(const ast::BeforeBlock& node) override;
     void visit_after_block(const ast::AfterBlock& node) override;
+    void visit_longitudinal_diffusion_block(const ast::LongitudinalDiffusionBlock& node) override;
 };
 
 /** @} */  // end of codegen_details

src/codegen/codegen_info.cpp

@@ -8,6 +8,8 @@
 #include "codegen/codegen_info.hpp"
 
 #include "ast/all.hpp"
+#include "ast/longitudinal_diffusion_block.hpp"
+#include "visitors/rename_visitor.hpp"
 #include "visitors/var_usage_visitor.hpp"
 #include "visitors/visitor_utils.hpp"
 
@@ -17,6 +19,48 @@ namespace codegen {
 
 using visitor::VarUsageVisitor;
 
+LongitudinalDiffusionInfo::LongitudinalDiffusionInfo(
+    const std::shared_ptr<ast::Name>& volume_index_name,
+    std::shared_ptr<ast::Expression> volume_expr,
+    const std::shared_ptr<ast::Name>& rate_index_name,
+    std::shared_ptr<ast::Expression> rate_expr)
+    : volume_index_name(volume_index_name ? volume_index_name->get_node_name() : std::string{})
+    , volume_expr(std::move(volume_expr))
+    , rate_index_name(rate_index_name ? rate_index_name->get_node_name() : std::string{})
+    , rate_expr(std::move(rate_expr)) {}
+
+std::shared_ptr<ast::Expression> LongitudinalDiffusionInfo::volume(
+    const std::string& index_name) const {
+    return substitute_index(index_name, volume_index_name, volume_expr);
+}
+std::shared_ptr<ast::Expression> LongitudinalDiffusionInfo::diffusion_rate(
+    const std::string& index_name) const {
+    return substitute_index(index_name, rate_index_name, rate_expr);
+}
+
+double LongitudinalDiffusionInfo::dfcdc(const std::string& /* index_name */) const {
+    // Needed as part of the Jacobian to stabalize
+    // the implicit time-integration. However,
+    // currently, it's set to `0.0` for simplicity.
+    return 0.0;
+}
+
+std::shared_ptr<ast::Expression> LongitudinalDiffusionInfo::substitute_index(
+    const std::string& index_name,
+    const std::string& old_index_name,
+    const std::shared_ptr<ast::Expression>& old_expr) const {
+    if (old_index_name == "") {
+        // The expression doesn't contain an index that needs substituting.
+        return old_expr;
+    }
+    auto new_expr = old_expr->clone();
+
+    auto v = visitor::RenameVisitor(old_index_name, index_name);
+    new_expr->accept(v);
+
+    return std::shared_ptr<ast::Expression>(dynamic_cast<ast::Expression*>(new_expr));
+}
+
 /// if any ion has write variable
 bool CodegenInfo::ion_has_write_variable() const noexcept {
     return std::any_of(ions.begin(), ions.end(), [](auto const& ion) {

src/codegen/codegen_info.hpp

@@ -284,6 +284,44 @@ struct IndexSemantics {
         , size(size) {}
 };
 
+/**
+ * \brief Information required to print LONGITUDINAL_DIFFUSION callbacks.
+ */
+class LongitudinalDiffusionInfo {
+  public:
+    LongitudinalDiffusionInfo(const std::shared_ptr<ast::Name>& index_name,
+                              std::shared_ptr<ast::Expression> volume_expr,
+                              const std::shared_ptr<ast::Name>& rate_index_name,
+                              std::shared_ptr<ast::Expression> rate_expr);
+    /// Volume of this species.
+    ///
+    /// If the volume expression is an indexed expression, the index in the
+    /// expression is substituted with `index_name`.
+    std::shared_ptr<ast::Expression> volume(const std::string& index_name) const;
+
+    /// Difusion rate of this species.
+    ///
+    /// If the diffusion expression is an indexed expression, the index in the
+    /// expression is substituted with `index_name`.
+    std::shared_ptr<ast::Expression> diffusion_rate(const std::string& index_name) const;
+
+    /// The value of what NEURON calls `dfcdc`.
+    double dfcdc(const std::string& /* index_name */) const;
+
+  protected:
+    std::shared_ptr<ast::Expression> substitute_index(
+        const std::string& index_name,
+        const std::string& old_index_name,
+        const std::shared_ptr<ast::Expression>& old_expr) const;
+
+  private:
+    std::string volume_index_name;
+    std::shared_ptr<ast::Expression> volume_expr;
+
+    std::string rate_index_name;
+    std::shared_ptr<ast::Expression> rate_expr;
+};
+
 
 /**
  * \class CodegenInfo
@@ -447,6 +485,9 @@ struct CodegenInfo {
     /// all factors defined in the mod file
     std::vector<const ast::FactorDef*> factor_definitions;
 
+    /// for each state, the information needed to print the callbacks.
+    std::map<std::string, LongitudinalDiffusionInfo> longitudinal_diffusion_info;
+
     /// ions used in the mod file
     std::vector<Ion> ions;