Skip to content

Commit

Permalink
Diffusion: tests, fixed coefficients in the solver, & tutorial (#2226)
Browse files Browse the repository at this point in the history 
By careful re-analysis of the diffusion solver we found that the
coeffiecients of the LHS need an additional factor of the CV Volume.

Scopes be creeping. During the analysis it became apparent that the
during coupling to the cable equation proper raises at lot of potential
problems.

Conversion of ionic current densities $i_X$ to mass transfer is $\dot
c_X = \frac{i_XA}{q F V}$ with volume, area, Faraday's constant, and
charge. Now, people are interested in using neutral species $n$, ie
$q=0$ for which we also should expect $i_n=0$. Yet we need to special
case on this (and assert zero current) to avoid ill-defined division.
That's the smaller problem.

The more worrying issue is this:
We construct a coupling term from the cable equation to the ionic
diffusion via $\dot c_X = \frac{i_XA}{q F V}$. Yet, there's no
backreaction to the cable equation, unless the user explicitly
constructs it. How? Well, the way ions should feed back to the voltage
is this:
`iX = g(U-eX)`
this is the conductance model and should be found in an NMODL file and
`eX` $\sim \ln\frac{X_i}{X_o}$
which is the Nernst equation found, again, in a builtin NMODL file. But,
for technical reasons we had to invent a special diffusive concentration
`Xd` which is _not_ identical to `Xi`. For proper behaviour, we should
have used `Xd` instead of `Xi` in the Nernst term above. So, a
non-standard Nernst module needs to be used.

This brings me to the great change in this context: remove the offending
term and carefully document how to retrofit it in NMODL and add the
proper Nernst model, too.

fixes #2145 
requires #2209

---------

Co-authored-by: Jannik Luboeinski <[email protected]>
Co-authored-by: Jannik Luboeinski <[email protected]>
  • Loading branch information
@thorstenhater @jlubo
3 people authored Jan 22, 2025
1 parent 637712f commit 257e74f
Show file tree
Hide file tree
Showing 37 changed files with 1,137 additions and 418 deletions.
115 changes: 48 additions & 67 deletions arbor/backends/gpu/diffusion.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
}

Expand All @@ -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;

Expand Down Expand Up @@ -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:
Expand All @@ -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
Expand Down
44 changes: 19 additions & 25 deletions arbor/backends/gpu/diffusion.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
30 changes: 8 additions & 22 deletions arbor/backends/gpu/diffusion_state.hpp
View file
Original file line number Diff line number Diff line change
@@ -1,9 +1,7 @@
#pragma once

#include <cstring>

#include <vector>
#include <type_traits>

#include <arbor/common_types.hpp>

Expand All @@ -21,7 +19,6 @@ namespace gpu {

template <typename T, typename I>
struct diffusion_state {
public:
using value_type = T;
using size_type = I;

Expand All @@ -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]
Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -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);
Expand All @@ -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
Expand All @@ -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());
Expand All @@ -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);
}

Expand Down
2 changes: 0 additions & 2 deletions arbor/backends/gpu/shared_state.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,6 @@ void ion_state::init_concentration() {
}

void ion_state::zero_current() {
memory::fill(gX_, 0);
memory::fill(iX_, 0);
}

Expand Down Expand Up @@ -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.
Expand Down
5 changes: 1 addition & 4 deletions arbor/backends/gpu/shared_state.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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
Expand Down
6 changes: 2 additions & 4 deletions arbor/backends/multicore/cable_solver.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
Loading

0 comments on commit 257e74f

Please sign in to comment.