Implement LONGITUDINAL_DIFFUSION.

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
 
 .. toctree::
    :maxdepth: 3

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,62 @@ void CodegenHelperVisitor::visit_after_block(const ast::AfterBlock& node) {
     info.before_after_blocks.push_back(&node);
 }
 
+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,31 @@ struct IndexSemantics {
         , size(size) {}
 };
 
+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);
+    std::shared_ptr<ast::Expression> volume(const std::string& index_name) const;
+    std::shared_ptr<ast::Expression> diffusion_rate(const std::string& index_name) const;
+
+    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 +472,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;
 

src/codegen/codegen_neuron_cpp_visitor.cpp

@@ -36,7 +36,6 @@ using visitor::VarUsageVisitor;
 
 using symtab::syminfo::NmodlType;
 
-
 /****************************************************************************************/
 /*                              Generic information getters                             */
 /****************************************************************************************/
@@ -436,6 +435,79 @@ void CodegenNeuronCppVisitor::print_hoc_py_wrapper_function_definitions() {
     }
 }
 
+CodegenNeuronCppVisitor::ParamVector CodegenNeuronCppVisitor::ldifusfunc1_parameters() const {
+    return ParamVector{{"", "ldifusfunc2_t", "", "_f"},
+                       {"", "const _nrn_model_sorted_token&", "", "_sorted_token"},
+                       {"", "NrnThread&", "", "_nt"}};
+}
+
+
+CodegenNeuronCppVisitor::ParamVector CodegenNeuronCppVisitor::ldifusfunc3_parameters() const {
+    return ParamVector{{"", "int", "", "_i"},
+                       {"", "Memb_list*", "", "_ml_arg"},
+                       {"", "size_t", "", "id"},
+                       {"", "Datum*", "", "_ppvar"},
+                       {"", "double*", "", "_pdvol"},
+                       {"", "double*", "", "_pdfcdc"},
+                       {"", "Datum*", "", "/* _thread */"},
+                       {"", "NrnThread*", "", "nt"},
+                       {"", "const _nrn_model_sorted_token&", "", "_sorted_token"}};
+}
+
+void CodegenNeuronCppVisitor::print_longitudinal_diffusion_callbacks() {
+    auto coeff_callback_name = [](const std::string& var_name) {
+        return fmt::format("_diffusion_coefficient_{}", var_name);
+    };
+
+    auto space_name = [](const std::string& var_name) {
+        return fmt::format("_diffusion_space_{}", var_name);
+    };
+
+    for (auto [var_name, values]: info.longitudinal_diffusion_info) {
+        printer->fmt_line("static void* {};", space_name(var_name));
+        printer->fmt_push_block("static double {}({})",
+                                coeff_callback_name(var_name),
+                                get_parameter_str(ldifusfunc3_parameters()));
+
+        print_entrypoint_setup_code_from_memb_list();
+
+        auto volume_expr = values.volume("_i");
+        auto mu_expr = values.diffusion_rate("_i");
+
+        printer->add_indent();
+        printer->add_text("*_pdvol= ");
+        volume_expr->accept(*this);
+        printer->add_text(";");
+        printer->add_newline();
+
+        printer->add_line("*_pdfcdc = 0.0;");
+        printer->add_indent();
+        printer->add_text("return ");
+        mu_expr->accept(*this);
+        printer->add_text(";");
+        printer->add_newline();
+
+        printer->pop_block();
+    }
+
+    printer->fmt_push_block("static void _apply_diffusion_function({})",
+                            get_parameter_str(ldifusfunc1_parameters()));
+    for (auto [var_name, values]: info.longitudinal_diffusion_info) {
+        auto var = program_symtab->lookup(var_name);
+        size_t array_size = var->get_length();
+        printer->fmt_push_block("for(size_t _i = 0; _i < {}; ++_i)", array_size);
+        printer->fmt_line(
+            "(*_f)(mech_type, {}, &{}, _i, /* x pos */ {}, /* Dx pos */ {}, _sorted_token, _nt);",
+            coeff_callback_name(var_name),
+            space_name(var_name),
+            position_of_float_var(var_name),
+            position_of_float_var("D" + var_name));
+        printer->pop_block();
+    }
+    printer->pop_block();
+    printer->add_newline();
+}
+
 /****************************************************************************************/
 /*                           Code-specific helper routines                              */
 /****************************************************************************************/
@@ -1301,6 +1373,11 @@ void CodegenNeuronCppVisitor::print_mechanism_register() {
                           info.semantics[i].name);
     }
 
+    if (!info.longitudinal_diffusion_info.empty()) {
+        printer->fmt_line("hoc_register_ldifus1(_apply_diffusion_function);");
+    }
+
+
     if (info.write_concentration) {
         printer->fmt_line("nrn_writes_conc(mech_type, 0);");
     }
@@ -2275,6 +2352,7 @@ void CodegenNeuronCppVisitor::print_codegen_routines() {
     print_nrn_destructor_declaration();
     print_nrn_alloc();
     print_function_prototypes();
+    print_longitudinal_diffusion_callbacks();
     print_point_process_function_definitions();
     print_setdata_functions();
     print_check_table_entrypoint();
@@ -2457,6 +2535,15 @@ void CodegenNeuronCppVisitor::visit_watch_statement(const ast::WatchStatement& /
     return;
 }
 
+void CodegenNeuronCppVisitor::visit_longitudinal_diffusion_block(
+    const ast::LongitudinalDiffusionBlock& /* node */) {
+    // These are handled via `print_longitdudinal_*`.
+}
+
+void CodegenNeuronCppVisitor::visit_lon_diffuse(const ast::LonDiffuse& /* node */) {
+    // These are handled via `print_longitdudinal_*`.
+}
+
 void CodegenNeuronCppVisitor::visit_for_netcon(const ast::ForNetcon& node) {
     // The setup for enabling this loop is:
     //   double ** _fornetcon_data = ...;

src/codegen/codegen_neuron_cpp_visitor.hpp

@@ -250,6 +250,21 @@
 
     void print_hoc_py_wrapper_function_definitions();
 
+    /**
+     * Prints the callbacks required for LONGITUDINAL_DIFFUSION.
+     */
+    void print_longitudinal_diffusion_callbacks();
+
+    /**
+     * Parameters for what NEURON calls `ldifusfunc1_t`.
+     */
+    ParamVector ldifusfunc1_parameters() const;
+
+    /**
+     * Parameters for what NEURON calls `ldifusfunc3_t`.
+     */
+    ParamVector ldifusfunc3_parameters() const;
+
 
     /****************************************************************************************/
     /*                             Code-specific helper routines                            */
@@ -708,8 +723,8 @@
 
     void visit_watch_statement(const ast::WatchStatement& node) override;
     void visit_for_netcon(const ast::ForNetcon& node) override;
-
-
+    void visit_longitudinal_diffusion_block(const ast::LongitudinalDiffusionBlock& node);
+    void visit_lon_diffuse(const ast::LonDiffuse& node);
 
 
   public:

src/language/code_generator.cmake

@@ -115,6 +115,7 @@ set(AST_GENERATED_SOURCES
     ${PROJECT_BINARY_DIR}/src/ast/local_list_statement.hpp
     ${PROJECT_BINARY_DIR}/src/ast/local_var.hpp
     ${PROJECT_BINARY_DIR}/src/ast/lon_diffuse.hpp
+    ${PROJECT_BINARY_DIR}/src/ast/longitudinal_diffusion_block.hpp
     ${PROJECT_BINARY_DIR}/src/ast/model.hpp
     ${PROJECT_BINARY_DIR}/src/ast/mutex_lock.hpp
     ${PROJECT_BINARY_DIR}/src/ast/mutex_unlock.hpp