Diffusion: tests, fixed coefficients in the solver, & tutorial

arbor/backends/gpu/diffusion.cu

@@ -16,33 +16,22 @@ namespace kernels {
 ///     - compute the RHS of the linear system to solve.
 template <typename T, typename I>
 __global__
-void assemble_diffusion(
-        T* __restrict__ const d,
-        T* __restrict__ const rhs,
-        const T* __restrict__ const invariant_d,
-        const T* __restrict__ const concentration,
-        const T* __restrict__ const voltage,
-        const T* __restrict__ const current,
-        const T q,
-        const T* __restrict__ const conductivity,
-        const T* __restrict__ const area,
-        const T dt,
-        const I* __restrict__ const perm,
-        unsigned n) {
+void assemble_diffusion(T* __restrict__ const d,
+                        T* __restrict__ const rhs,
+                        const T* __restrict__ const invariant_d,
+                        const T* __restrict__ const concentration,
+                        const T* __restrict__ const volume,
+                        const T dt,
+                        const I* __restrict__ const perm,
+                        unsigned n) {
     const unsigned tid = threadIdx.x + blockDim.x*blockIdx.x;
     if (tid < n) {
-        const auto pid = perm[tid];
-        auto u = voltage[tid];        // mV
-        auto g = conductivity[tid];   // µS
-        auto J = current[tid];        // A/m^2
-        auto A = 1e-3*area[tid];      // 1e-9·m²
+        auto _1_dt = 1e-3/dt;         // 1/µs
+        auto pid = perm[tid];
+        auto V = volume[tid];         // um^3
         auto X = concentration[tid];  // mM
-        // conversion from current density to concentration change
-        // using Faraday's constant
-        auto F = A/(q*96.485332);
-
-        d[pid]   = 1e-3/dt   + F*g + invariant_d[tid];
-        rhs[pid] = 1e-3/dt*X + F*(u*g - J);
+        d[pid]   = _1_dt*V + invariant_d[tid];
+        rhs[pid] = _1_dt*V*X;
     }
 }
 
@@ -57,16 +46,14 @@ void assemble_diffusion(
 /// number of branches.
 template <typename T>
 __global__
-void solve_diffusion(
-    T* __restrict__ const rhs,
-    T* __restrict__ const d,
-    const T* __restrict__ const u,
-    const level_metadata* __restrict__ const level_meta,
-    const arb_index_type* __restrict__ const level_lengths,
-    const arb_index_type* __restrict__ const level_parents,
-    const arb_index_type* __restrict__ const block_index,
-    const arb_index_type* __restrict__ const num_matrix) // number of packed matrices = number of cells
-{
+void solve_diffusion(T* __restrict__ const rhs,
+                     T* __restrict__ const d,
+                     const T* __restrict__ const u,
+                     const level_metadata* __restrict__ const level_meta,
+                     const arb_index_type* __restrict__ const level_lengths,
+                     const arb_index_type* __restrict__ const level_parents,
+                     const arb_index_type* __restrict__ const block_index,
+                     const arb_index_type* __restrict__ const num_matrix) {// number of packed matrices = number of cells
     const auto tid = threadIdx.x;
     const auto bid = blockIdx.x;
 
@@ -195,23 +182,18 @@ void solve_diffusion(
 
 } // namespace kernels
 
-ARB_ARBOR_API void assemble_diffusion(
-    arb_value_type* d,
-    arb_value_type* rhs,
-    const arb_value_type* invariant_d,
-    const arb_value_type* concentration,
-    const arb_value_type* voltage,
-    const arb_value_type* current,
-    arb_value_type q,
-    const arb_value_type* conductivity,
-    const arb_value_type* area,
-    const arb_value_type dt,
-    const arb_index_type* perm,
-    unsigned n)
-{
-    launch_1d(n, 128, kernels::assemble_diffusion<arb_value_type, arb_index_type>,
-        d, rhs, invariant_d, concentration, voltage, current, q, conductivity, area,
-        dt, perm, n);
+ARB_ARBOR_API void assemble_diffusion(arb_value_type* d,
+                                      arb_value_type* rhs,
+                                      const arb_value_type* invariant_d,
+                                      const arb_value_type* concentration,
+                                      const arb_value_type* volume,
+                                      const arb_value_type dt,
+                                      const arb_index_type* perm,
+                                      unsigned n) {
+    launch_1d(n,
+              128,
+              kernels::assemble_diffusion<arb_value_type, arb_index_type>,
+              d, rhs, invariant_d, concentration, volume, dt, perm, n);
 }
 
 // Example:
@@ -231,22 +213,21 @@ ARB_ARBOR_API void assemble_diffusion(
 // num_levels   = [3, 2, 3, ...]
 // num_cells    = [2, 3, ...]
 // num_blocks   = level_start.size() - 1 = num_levels.size() = num_cells.size()
-ARB_ARBOR_API void solve_diffusion(
-    arb_value_type* rhs,
-    arb_value_type* d,                     // diagonal values
-    const arb_value_type* u,               // upper diagonal (and lower diagonal as the matrix is SPD)
-    const level_metadata* level_meta,      // information pertaining to each level
-    const arb_index_type* level_lengths,   // lengths of branches of every level concatenated
-    const arb_index_type* level_parents,   // parents of branches of every level concatenated
-    const arb_index_type* block_index,     // start index into levels for each gpu block
-    arb_index_type* num_cells,             // the number of cells packed into this single matrix
-    arb_index_type* padded_size,           // length of rhs, d, u, including padding
-    unsigned num_blocks,                   // number of blocks
-    unsigned blocksize)                    // size of each block
-{
-    launch(num_blocks, blocksize, kernels::solve_diffusion<arb_value_type>,
-        rhs, d, u, level_meta, level_lengths, level_parents, block_index,
-        num_cells);
+ARB_ARBOR_API void solve_diffusion(arb_value_type* rhs,
+                                   arb_value_type* d,                     // diagonal values
+                                   const arb_value_type* u,               // upper diagonal (and lower diagonal as the matrix is SPD)
+                                   const level_metadata* level_meta,      // information pertaining to each level
+                                   const arb_index_type* level_lengths,   // lengths of branches of every level concatenated
+                                   const arb_index_type* level_parents,   // parents of branches of every level concatenated
+                                   const arb_index_type* block_index,     // start index into levels for each gpu block
+                                   arb_index_type* num_cells,             // the number of cells packed into this single matrix
+                                   arb_index_type* padded_size,           // length of rhs, d, u, including padding
+                                   unsigned num_blocks,                   // number of blocks
+                                   unsigned blocksize) {                  // size of each block
+    launch(num_blocks,
+           blocksize,
+           kernels::solve_diffusion<arb_value_type>,
+           rhs, d, u, level_meta, level_lengths, level_parents, block_index, num_cells);
 }
 
 } // namespace gpu

arbor/backends/gpu/diffusion.hpp

@@ -8,32 +8,26 @@
 namespace arb {
 namespace gpu {
 
-ARB_ARBOR_API void assemble_diffusion(
-    arb_value_type* d,
-    arb_value_type* rhs,
-    const arb_value_type* invariant_d,
-    const arb_value_type* concentration,
-    const arb_value_type* voltage,
-    const arb_value_type* current,
-    const arb_value_type q,
-    const arb_value_type* conductivity,
-    const arb_value_type* area,
-    const arb_value_type dt,
-    const arb_index_type* perm,
-    unsigned n);
+ARB_ARBOR_API void assemble_diffusion(arb_value_type* d,
+                                      arb_value_type* rhs,
+                                      const arb_value_type* invariant_d,
+                                      const arb_value_type* concentration,
+                                      const arb_value_type* volume,
+                                      const arb_value_type dt,
+                                      const arb_index_type* perm,
+                                      unsigned n);
 
-ARB_ARBOR_API void solve_diffusion(
-    arb_value_type* rhs,
-    arb_value_type* d,                     // diagonal values
-    const arb_value_type* u,               // upper diagonal (and lower diagonal as the matrix is SPD)
-    const level_metadata* level_meta,      // information pertaining to each level
-    const arb_index_type* level_lengths,   // lengths of branches of every level concatenated
-    const arb_index_type* level_parents,   // parents of branches of every level concatenated
-    const arb_index_type* block_index,     // start index (exclusive) into levels for each gpu block
-    arb_index_type* num_cells,             // the number of cells packed into this single matrix
-    arb_index_type* padded_size,           // length of rhs, d, u, including padding
-    unsigned num_blocks,                   // number of blocks
-    unsigned blocksize);                   // size of each block
+ARB_ARBOR_API void solve_diffusion(arb_value_type* rhs,
+                                   arb_value_type* d,                     // diagonal values
+                                   const arb_value_type* u,               // upper diagonal (and lower diagonal as the matrix is SPD)
+                                   const level_metadata* level_meta,      // information pertaining to each level
+                                   const arb_index_type* level_lengths,   // lengths of branches of every level concatenated
+                                   const arb_index_type* level_parents,   // parents of branches of every level concatenated
+                                   const arb_index_type* block_index,     // start index (exclusive) into levels for each gpu block
+                                   arb_index_type* num_cells,             // the number of cells packed into this single matrix
+                                   arb_index_type* padded_size,           // length of rhs, d, u, including padding
+                                   unsigned num_blocks,                   // number of blocks
+                                   unsigned blocksize);                   // size of each block
 
 } // namespace gpu
 } // namespace arb

arbor/backends/gpu/diffusion_state.hpp

@@ -1,9 +1,7 @@
 #pragma once
 
 #include <cstring>
-
 #include <vector>
-#include <type_traits>
 
 #include <arbor/common_types.hpp>
 
@@ -21,7 +19,6 @@ namespace gpu {
 
 template <typename T, typename I>
 struct diffusion_state {
-public:
     using value_type = T;
     using size_type = I;
 
@@ -37,7 +34,7 @@ struct diffusion_state {
     array rhs;   // [nA]
 
     // Required for matrix assembly
-    array cv_area;             // [μm^2]
+    array cv_volume;           // [μm^3]
 
     // Invariant part of the matrix diagonal
     array invariant_d;         // [μS]
@@ -88,7 +85,7 @@ struct diffusion_state {
     diffusion_state(const std::vector<size_type>& p,
                     const std::vector<size_type>& cell_cv_divs,
                     const std::vector<value_type>& face_diffusivity,
-                    const std::vector<value_type>& area) {
+                    const std::vector<value_type>& volume) {
         using util::make_span;
         constexpr unsigned npos = unsigned(-1);
 
@@ -356,7 +353,6 @@ struct diffusion_state {
         // d, u, rhs        : packed
         // invariant_d      : flat
         // cv_to_cell       : flat
-        // area             : flat
 
         // the invariant part of d is stored in in flat form
         std::vector<value_type> invariant_d_tmp(matrix_size, 0);
@@ -382,8 +378,8 @@ struct diffusion_state {
         // transform u_shuffled values into packed u vector.
         flat_to_packed(u_shuffled, u);
 
-        // the invariant part of d and cv_area are in flat form
-        cv_area = memory::make_const_view(area);
+        // data in flat form
+        cv_volume = memory::make_const_view(volume);
         invariant_d = memory::make_const_view(invariant_d_tmp);
 
         // calculate the cv -> cell mappings
@@ -399,18 +395,12 @@ struct diffusion_state {
     // Afterwards the diagonal and RHS will have been set given dt, voltage, current, and conductivity.
     //   dt [ms] (scalar)
     //   voltage [mV]
-    //   current density [A/m²]
-    //   conductivity [kS/m²]
-    void assemble(const value_type dt, const_view concentration, const_view voltage, const_view current, const_view conductivity, arb_value_type q) {
+    void assemble(const value_type dt, const_view concentration) {
         assemble_diffusion(d.data(),
                            rhs.data(),
                            invariant_d.data(),
                            concentration.data(),
-                           voltage.data(),
-                           current.data(),
-                           q,
-                           conductivity.data(),
-                           cv_area.data(),
+                           cv_volume.data(),
                            dt,
                            perm.data(),
                            size());
@@ -433,12 +423,8 @@ struct diffusion_state {
     }
 
     void solve(array& concentration,
-               const value_type dt,
-               const_view voltage,
-               const_view current,
-               const_view conductivity,
-               arb_value_type q) {
-        assemble(dt, concentration, voltage, current, conductivity, q);
+               const value_type dt) {
+        assemble(dt, concentration);
         solve(concentration);
     }
 

arbor/backends/gpu/shared_state.cpp

@@ -74,7 +74,6 @@ void ion_state::init_concentration() {
 }
 
 void ion_state::zero_current() {
-    memory::fill(gX_, 0);
     memory::fill(iX_, 0);
 }
 
@@ -261,7 +260,6 @@ void shared_state::instantiate(mechanism& m,
         ion_state.external_concentration  = oion->Xo_.data();
         ion_state.diffusive_concentration = oion->Xd_.data();
         ion_state.ionic_charge            = oion->charge.data();
-        ion_state.conductivity            = oion->gX_.data();
     }
 
     // If there are no sites (is this ever meaningful?) there is nothing more to do.

arbor/backends/gpu/shared_state.hpp

@@ -9,10 +9,7 @@
 #include <arbor/fvm_types.hpp>
 
 #include "fvm_layout.hpp"
-#include "timestep_range.hpp"
-
 #include "threading/threading.hpp"
-
 #include "backends/common_types.hpp"
 #include "backends/shared_state_base.hpp"
 #include "backends/gpu/rand.hpp"
@@ -244,7 +241,7 @@ ARB_ARBOR_API std::ostream& operator<<(std::ostream& o, shared_state& s);
 
 } // namespace gpu
 
-ARB_SERDES_ENABLE_EXT(gpu::ion_state, Xd_, gX_);
+ARB_SERDES_ENABLE_EXT(gpu::ion_state, Xd_);
 ARB_SERDES_ENABLE_EXT(gpu::mech_storage,
                       data_,
                       // NOTE(serdes) ion_states_, this is just a bunch of pointers

arbor/backends/multicore/cable_solver.hpp

@@ -53,16 +53,14 @@ struct cable_solver {
                 const auto gij = cond[i];
                 u[i] = -gij;
                 invariant_d[i] += gij;
-                if (p[i]!=-1) { // root
-                    invariant_d[p[i]] += gij;
-                }
+                if (auto pi= p[i]; pi != -1) invariant_d[pi] += gij;
             }
         }
     }
 
     // Setup and solve the cable equation
     // * expects the voltage from its first argument
-    // * will likewise overwrite the first argument with the solction
+    // * will likewise overwrite the first argument with the solution
     template<typename T>
     void solve(T& rhs, const value_type dt, const_view current, const_view conductivity, const_view cv_area) {
         value_type * const ARB_NO_ALIAS d_ = d.data();