diff --git a/arbor/backends/gpu/diffusion.cu b/arbor/backends/gpu/diffusion.cu index 697dc44686..692a4c3e27 100644 --- a/arbor/backends/gpu/diffusion.cu +++ b/arbor/backends/gpu/diffusion.cu @@ -16,33 +16,22 @@ namespace kernels { /// - compute the RHS of the linear system to solve. template __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 __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, - 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, + 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, - 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, + rhs, d, u, level_meta, level_lengths, level_parents, block_index, num_cells); } } // namespace gpu diff --git a/arbor/backends/gpu/diffusion.hpp b/arbor/backends/gpu/diffusion.hpp index cf63890bfc..19e97c7a66 100644 --- a/arbor/backends/gpu/diffusion.hpp +++ b/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 diff --git a/arbor/backends/gpu/diffusion_state.hpp b/arbor/backends/gpu/diffusion_state.hpp index aae9c23503..e2a68a5008 100644 --- a/arbor/backends/gpu/diffusion_state.hpp +++ b/arbor/backends/gpu/diffusion_state.hpp @@ -1,9 +1,7 @@ #pragma once #include - #include -#include #include @@ -21,7 +19,6 @@ namespace gpu { template 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& p, const std::vector& cell_cv_divs, const std::vector& face_diffusivity, - const std::vector& area) { + const std::vector& 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 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); } diff --git a/arbor/backends/gpu/shared_state.cpp b/arbor/backends/gpu/shared_state.cpp index b7f4277109..30252eb7f7 100644 --- a/arbor/backends/gpu/shared_state.cpp +++ b/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. diff --git a/arbor/backends/gpu/shared_state.hpp b/arbor/backends/gpu/shared_state.hpp index 39f3302021..0abbf49c69 100644 --- a/arbor/backends/gpu/shared_state.hpp +++ b/arbor/backends/gpu/shared_state.hpp @@ -9,10 +9,7 @@ #include #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 diff --git a/arbor/backends/multicore/cable_solver.hpp b/arbor/backends/multicore/cable_solver.hpp index bc748486aa..113510838e 100644 --- a/arbor/backends/multicore/cable_solver.hpp +++ b/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 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(); diff --git a/arbor/backends/multicore/diffusion_solver.hpp b/arbor/backends/multicore/diffusion_solver.hpp index 7cee6e4bfb..567a38f56a 100644 --- a/arbor/backends/multicore/diffusion_solver.hpp +++ b/arbor/backends/multicore/diffusion_solver.hpp @@ -22,7 +22,7 @@ struct diffusion_solver { array d; // [μS] array u; // [μS] - array cv_area; // [μm^2] + array cv_volume; // [μm^3] array invariant_d; // [μS] invariant part of matrix diagonal diffusion_solver() = default; @@ -35,13 +35,12 @@ struct diffusion_solver { diffusion_solver(const std::vector& p, const std::vector& cell_cv_divs, const std::vector& diff, - const std::vector& area): + const std::vector& volume): parent_index(p.begin(), p.end()), cell_cv_divs(cell_cv_divs.begin(), cell_cv_divs.end()), d(size(), 0), u(size(), 0), - cv_area(area.begin(), area.end()), - invariant_d(size(), 0) - { + cv_volume(volume.begin(), volume.end()), + invariant_d(size(), 0) { // Sanity check arb_assert(diff.size() == size()); arb_assert(cell_cv_divs.back() == (index_type)size()); @@ -62,37 +61,22 @@ struct diffusion_solver { // Assemble and solve the matrix // Assemble the matrix + // concentration [mM] (per control volume) // dt [ms] (scalar) - // internal conc [mM] (per control volume) - // voltage [mV] (per control volume) - // current density [A.m^-2] (per control volume and species) - // diffusivity [m^2/s] (per control volume) - // charge [e] - // diff. concentration - // * will be overwritten by the solution template void solve(T& concentration, - value_type dt, - const_view voltage, - const_view current, - const_view conductivity, - arb_value_type q) { + value_type dt) { auto cell_cv_part = util::partition_view(cell_cv_divs); index_type ncells = cell_cv_part.size(); + + value_type _1_dt = 1e-3/dt; // 1/µs // loop over submatrices for (auto m: util::make_span(0, ncells)) { - value_type _1_dt = 1e-3/dt; // 1/µs for (auto i: util::make_span(cell_cv_part[m])) { - auto u = voltage[i]; // mV - auto g = conductivity[i]; // µS - auto J = current[i]; // A/m^2 - auto A = 1e-3*cv_area[i]; // 1e-9·m² auto X = concentration[i]; // mM - // conversion from current density to concentration change - // using Faraday's constant - auto F = A/(q*96.485332); - d[i] = _1_dt + F*g + invariant_d[i]; - concentration[i] = _1_dt*X + F*(u*g - J); + auto V = cv_volume[i]; // m^3 + d[i] = _1_dt*V + invariant_d[i]; + concentration[i] = _1_dt*V*X; } } solve(concentration); @@ -103,22 +87,23 @@ struct diffusion_solver { void solve(T& rhs) { // loop over submatrices for (const auto& [first, last]: util::partition_view(cell_cv_divs)) { - if (first >= last) continue; // skip cell with no CVs + if (first >= last || d[first] == 0) continue; // skip cell with no CVs // backward sweep - for(auto i=last-1; i>first; --i) { - auto pi = parent_index[i]; - auto fac = -u[i] / d[i]; - d[pi] = std::fma(fac, u[i], d[pi]); - rhs[pi] = std::fma(fac, rhs[i], rhs[pi]); + for(int i = last - 1; i > first; --i) { + const auto pi = parent_index[i]; + const auto factor = u[i] / d[i]; + d[pi] -= factor * u[i]; + rhs[pi] -= factor * rhs[i]; } // solve root rhs[first] /= d[first]; // forward sweep - for(auto i=first+1; iXo_.data(); ion_state.diffusive_concentration = oion->Xd_.data(); ion_state.ionic_charge = oion->charge.data(); - ion_state.conductivity = oion->gX_.data(); } // Initialize state and parameter vectors with default values. diff --git a/arbor/backends/multicore/shared_state.hpp b/arbor/backends/multicore/shared_state.hpp index 3674d635d7..5353ecd8da 100644 --- a/arbor/backends/multicore/shared_state.hpp +++ b/arbor/backends/multicore/shared_state.hpp @@ -14,16 +14,12 @@ #include #include "fvm_layout.hpp" - #include "util/padded_alloc.hpp" #include "util/rangeutil.hpp" - #include "threading/threading.hpp" - #include "backends/common_types.hpp" #include "backends/rand_fwd.hpp" #include "backends/shared_state_base.hpp" - #include "backends/multicore/threshold_watcher.hpp" #include "backends/multicore/multicore_common.hpp" #include "backends/multicore/partition_by_constraint.hpp" @@ -250,7 +246,7 @@ ARB_ARBOR_API std::ostream& operator<<(std::ostream& o, const shared_state& s); } // namespace multicore // Xd and gX are the only things that persist -ARB_SERDES_ENABLE_EXT(multicore::ion_state, Xd_, gX_); +ARB_SERDES_ENABLE_EXT(multicore::ion_state, Xd_); ARB_SERDES_ENABLE_EXT(multicore::mech_storage, data_, // NOTE(serdes) ion_states_, this is just a bunch of pointers diff --git a/arbor/backends/shared_state_base.hpp b/arbor/backends/shared_state_base.hpp index 0c1742a0d5..51bd2afd63 100644 --- a/arbor/backends/shared_state_base.hpp +++ b/arbor/backends/shared_state_base.hpp @@ -7,6 +7,7 @@ #include "backends/common_types.hpp" #include "fvm_layout.hpp" +#include "timestep_range.hpp" #include "util/rangeutil.hpp" #include "timestep_range.hpp" #include "event_lane.hpp" @@ -17,6 +18,8 @@ namespace arb { template struct shared_state_base { + using diff_solver = typename ion_state::solver_type; + void update_time_to(const timestep_range::timestep& ts) { auto d = static_cast(this); d->time_to = ts.t_end(); @@ -49,7 +52,10 @@ struct shared_state_base { void configure_solver(const fvm_cv_discretization& disc) { auto d = static_cast(this); - d->solver = {disc.geometry.cv_parent, disc.geometry.cell_cv_divs, disc.cv_capacitance, disc.face_conductance}; + d->solver = {disc.geometry.cv_parent, + disc.geometry.cell_cv_divs, + disc.cv_capacitance, + disc.face_conductance}; } void add_ion(const std::string& ion_name, @@ -65,11 +71,11 @@ struct shared_state_base { const std::unordered_map& ions) { auto d = static_cast(this); for (const auto& [ion, data]: ions) { - std::unique_ptr solver = nullptr; - if (data.is_diffusive) solver = std::make_unique(disc.geometry.cv_parent, - disc.geometry.cell_cv_divs, - data.face_diffusivity, - disc.cv_area); + std::unique_ptr solver = nullptr; + if (data.is_diffusive) solver = std::make_unique(disc.geometry.cv_parent, + disc.geometry.cell_cv_divs, + data.face_diffusivity, + disc.cv_volume); d->add_ion(ion, data, std::move(solver)); } } @@ -139,11 +145,7 @@ struct shared_state_base { for (auto& [ion, data]: d->ion_data) { if (data.solver) { data.solver->solve(data.Xd_, - d->dt, - d->voltage, - data.iX_, - data.gX_, - data.charge[0]); + d->dt); } } } diff --git a/arbor/fvm_layout.cpp b/arbor/fvm_layout.cpp index 158edc4fb6..2adf810489 100644 --- a/arbor/fvm_layout.cpp +++ b/arbor/fvm_layout.cpp @@ -212,33 +212,32 @@ ARB_ARBOR_API cv_geometry& append(cv_geometry& geom, const cv_geometry& right) { ARB_ARBOR_API fvm_cv_discretization& append(fvm_cv_discretization& dczn, const fvm_cv_discretization& right) { using util::append; - append(dczn.geometry, right.geometry); - - // Those in L and R: merge - for (auto& [ion, data]: dczn.diffusive_ions) { - const auto& rhs = right.diffusive_ions.find(ion); - if (rhs != right.diffusive_ions.end()) { - append(data.axial_inv_diffusivity, rhs->second.axial_inv_diffusivity); - append(data.face_diffusivity, rhs->second.face_diffusivity); + // Merge diffusive ion data, scan ions in L and R, then... + // ... those in L and R: append R's data to that of L + for (auto& [ion, lhs]: dczn.diffusive_ions) { + if (auto rhs = right.diffusive_ions.find(ion); rhs != right.diffusive_ions.end()) { + append(lhs.axial_resistivity, rhs->second.axial_resistivity); + append(lhs.face_diffusivity, rhs->second.face_diffusivity); } } - // Those only in R: add to L - for (auto& [ion, data]: right.diffusive_ions) { - const auto& lhs = dczn.diffusive_ions.find(ion); - if (lhs == dczn.diffusive_ions.end()) { - dczn.diffusive_ions[ion].axial_inv_diffusivity = data.axial_inv_diffusivity; - dczn.diffusive_ions[ion].face_diffusivity = data.face_diffusivity; + + // ... those only in R: add to L + for (auto& [ion, rhs]: right.diffusive_ions) { + if (0 == dczn.diffusive_ions.count(ion)) { + dczn.diffusive_ions[ion].axial_resistivity = rhs.axial_resistivity; + dczn.diffusive_ions[ion].face_diffusivity = rhs.face_diffusivity; } } - append(dczn.face_conductance, right.face_conductance); - append(dczn.cv_area, right.cv_area); - append(dczn.cv_capacitance, right.cv_capacitance); + append(dczn.geometry, right.geometry); + append(dczn.face_conductance, right.face_conductance); + append(dczn.cv_area, right.cv_area); + append(dczn.cv_volume, right.cv_volume); + append(dczn.cv_capacitance, right.cv_capacitance); append(dczn.init_membrane_potential, right.init_membrane_potential); - append(dczn.temperature_K, right.temperature_K); - append(dczn.diam_um, right.diam_um); - - append(dczn.axial_resistivity, right.axial_resistivity); + append(dczn.temperature_K, right.temperature_K); + append(dczn.diam_um, right.diam_um); + append(dczn.axial_resistivity, right.axial_resistivity); return dczn; } @@ -255,8 +254,9 @@ fvm_cv_discretize(const cable_cell& cell, if (D.geometry.empty()) return D; auto n_cv = D.geometry.size(); - D.face_conductance.resize(n_cv); + D.face_conductance.resize(n_cv, 0.0); D.cv_area.resize(n_cv); + D.cv_volume.resize(n_cv); D.cv_capacitance.resize(n_cv); D.init_membrane_potential.resize(n_cv); D.temperature_K.resize(n_cv); @@ -275,13 +275,9 @@ fvm_cv_discretize(const cable_cell& cell, const auto& diffusivity = assignments.get(); const auto& provider = cell.provider(); - struct inv_diff { - iexpr value; - }; - // Set up for ion diffusivity std::unordered_map diffusive_ions; - std::unordered_map> inverse_diffusivity; + std::unordered_map> ion_diffusivity; // Collect all eglible ions: those where any cable has finite diffusivity for (const auto& [ion, data]: global_dflt.ion_data) { @@ -305,30 +301,30 @@ fvm_cv_discretize(const cable_cell& cell, }); if (diffusive) { // Provide a (non-sensical) default. - if (!diffusive_ions.count(ion)) diffusive_ions[ion] = {NAN}; - auto& inv = inverse_diffusivity[ion]; - for (const auto& [k, v]: data) inv.insert(k, {1.0/v.value}); + if (!diffusive_ions.count(ion)) diffusive_ions[ion] = {}; + auto& diff = ion_diffusivity[ion]; + for (const auto& [k, v]: data) diff.insert(k, v.value); } } // Remap diffusivity to resistivity for (auto& [ion, data]: diffusive_ions) { - auto& id_map = inverse_diffusivity[ion]; + auto& id_map = ion_diffusivity[ion]; arb_value_type def = data.default_value; if (def <= 0.0 || std::isnan(def)) { throw make_cc_error("Illegal global diffusivity '{}' for ion '{}'; possibly unset." " Please define a positive global or cell default.", def, ion); } // Write inverse diffusivity / diffuse resistivity map - auto& id = data.axial_inv_diffusivity; + auto& id = data.axial_resistivity; id.resize(1); msize_t n_branch = D.geometry.n_branch(0); id.reserve(n_branch); for (msize_t i = 0; i double { - auto ie = thingify(par.value, provider); + auto scale_param = [&, ion=ion](const auto&, const iexpr& par) { + auto ii = 1.0/par; + auto ie = thingify(ii, provider); auto sc = ie->eval(provider, cable); if (def <= 0.0 || std::isnan(def)) { throw make_cc_error("Illegal diffusivity '{}' for ion '{}' at cable {}." @@ -340,7 +336,7 @@ fvm_cv_discretize(const cable_cell& cell, id[0].push_back(pw); } // Prepare conductivity map - data.face_diffusivity.resize(n_cv); + data.face_diffusivity.resize(n_cv, 0.0); } D.axial_resistivity.resize(1); @@ -374,7 +370,9 @@ fvm_cv_discretize(const cable_cell& cell, } arb_index_type p = D.geometry.cv_parent[i]; - if (p!=-1) { + bool has_parent = p != -1; + + if (has_parent) { auto parent_cables = D.geometry.cables(p); msize_t bid = cv_cables.front().branch; double parent_refpt = 0; @@ -396,9 +394,12 @@ fvm_cv_discretize(const cable_cell& cell, mcable span{bid, parent_refpt, cv_refpt}; double resistance = embedding.integrate_ixa(span, D.axial_resistivity[0].at(bid)); D.face_conductance[i] = 100/resistance; // 100 scales to µS. + // Compute + auto len = embedding.integrate_length(span); for (auto& [ion, info]: diffusive_ions) { - double resistance = embedding.integrate_ixa(span, info.axial_inv_diffusivity[0].at(bid)); - info.face_diffusivity[i] = 1.0/resistance; // scale to m^2/s + // TODO scale to m/s^2 + auto sigma = len/embedding.integrate_ixa(span, info.axial_resistivity[0].at(bid)); + info.face_diffusivity[i] = sigma; } } @@ -408,7 +409,7 @@ fvm_cv_discretize(const cable_cell& cell, D.diam_um[i] = 0; double cv_length = 0; - for (mcable cable: cv_cables) { + for (const mcable& cable: cv_cables) { auto scale_param = [&](const auto&, const auto& par) -> double { auto ie = thingify(par.scale, provider); auto sc = par.value*ie->eval(provider, cable); @@ -426,30 +427,27 @@ fvm_cv_discretize(const cable_cell& cell, cv_length += embedding.integrate_length(cable); } - if (D.cv_area[i]>0) { - auto A = D.cv_area[i]; - D.init_membrane_potential[i] /= A; - D.temperature_K[i] /= A; + double area = D.cv_area[i]; + if (cv_length > 0) { + D.diam_um[i] = area/(cv_length*math::pi); + } + D.cv_volume[i] = 0.25*area*D.diam_um[i]; - for (auto& [ion, info]: diffusive_ions) { - info.face_diffusivity[i] /= A; - } - // If parent is trivial, and there is no grandparent, then we can use values from this CV - // to get initial values for the parent. (The other case, when there is a grandparent, is - // caught below.) - if (p!=-1 && D.geometry.cv_parent[p]==-1 && D.cv_area[p]==0) { + if (area > 0) { + D.init_membrane_potential[i] /= area; + D.temperature_K[i] /= area; + // If parent is trivial, and there is no grandparent, use values from this CV + // as initial values for the parent. + // The other case, when there is a grandparent, is caught below when i == p. + if (has_parent && D.geometry.cv_parent[p] == -1 && D.cv_area[p] == 0) { D.init_membrane_potential[p] = D.init_membrane_potential[i]; - D.temperature_K[p] = D.temperature_K[i]; + D.temperature_K[p] = D.temperature_K[i]; } } - else if (p!=-1) { + else if (has_parent) { // Use parent CV to get a sensible initial value for voltage and temp on zero-size CVs. D.init_membrane_potential[i] = D.init_membrane_potential[p]; - D.temperature_K[i] = D.temperature_K[p]; - } - - if (cv_length>0) { - D.diam_um[i] = D.cv_area[i]/(cv_length*math::pi); + D.temperature_K[i] = D.temperature_K[p]; } } @@ -1219,8 +1217,6 @@ make_density_mechanism_config(const region_assignment& assignments, for (const auto& [name, cables]: assignments) { const auto& info = data.catalogue[name]; auto config = make_mechanism_config(info, arb_mechanism_kind_density); - - auto parameters = ordered_parameters(info); auto n_param = parameters.size(); diff --git a/arbor/fvm_layout.hpp b/arbor/fvm_layout.hpp index 9ba9a9d615..9b8ec3bd79 100644 --- a/arbor/fvm_layout.hpp +++ b/arbor/fvm_layout.hpp @@ -142,7 +142,7 @@ struct fvm_diffusion_info { using value_type = arb_value_type; value_type default_value; std::vector face_diffusivity; - std::vector> axial_inv_diffusivity; + std::vector> axial_resistivity; fvm_diffusion_info(value_type d): default_value(d) {} fvm_diffusion_info(): fvm_diffusion_info{0.0} {} @@ -162,6 +162,7 @@ struct fvm_cv_discretization { // Following members have one element per CV. std::vector face_conductance; // [µS] std::vector cv_area; // [µm²] + std::vector cv_volume; // [µm3]// std::vector cv_capacitance; // [pF] std::vector init_membrane_potential; // [mV] std::vector temperature_K; // [K] diff --git a/arbor/include/arbor/gpu/cuda_api.hpp b/arbor/include/arbor/gpu/cuda_api.hpp index 532b9ec8f9..193796e80b 100644 --- a/arbor/include/arbor/gpu/cuda_api.hpp +++ b/arbor/include/arbor/gpu/cuda_api.hpp @@ -100,27 +100,11 @@ inline api_error_type device_mem_get_info(ARGS &&... args) { // Wrappers around CUDA addition functions. // CUDA 8 introduced support for atomicAdd with double precision, but only for -// Pascal GPUs (__CUDA_ARCH__ >= 600). These wrappers provide a portable -// atomic addition interface that chooses the appropriate implementation. - -#if __CUDA_ARCH__ < 600 // Maxwell or older (no native double precision atomic addition) -__device__ -inline double gpu_atomic_add(double* address, double val) { - using I = unsigned long long int; - I* address_as_ull = (I*)address; - I old = *address_as_ull, assumed; - do { - assumed = old; - old = atomicCAS(address_as_ull, assumed, __double_as_longlong(val+__longlong_as_double(assumed))); - } while (assumed != old); - return __longlong_as_double(old); -} -#else // use build in atomicAdd for double precision from Pascal onwards +// Pascal GPUs (__CUDA_ARCH__ >= 600). We assume no one is running anything older. __device__ inline double gpu_atomic_add(double* address, double val) { return atomicAdd(address, val); } -#endif __device__ inline double gpu_atomic_sub(double* address, double val) { @@ -138,11 +122,14 @@ inline float gpu_atomic_sub(float* address, float val) { } /// Warp-Level Primitives - __device__ __inline__ unsigned ballot(unsigned mask, unsigned is_root) { return __ballot_sync(mask, is_root); } +__device__ __inline__ unsigned active_mask() { + return __activemask(); +} + __device__ __inline__ unsigned any(unsigned mask, unsigned width) { return __any_sync(mask, width); } diff --git a/arbor/include/arbor/gpu/hip_api.hpp b/arbor/include/arbor/gpu/hip_api.hpp index 019f235a8e..ec1f5ea453 100644 --- a/arbor/include/arbor/gpu/hip_api.hpp +++ b/arbor/include/arbor/gpu/hip_api.hpp @@ -122,8 +122,7 @@ inline float gpu_atomic_sub(float* address, float val) { template __device__ __inline__ std::enable_if_t< !std::is_same_v, double>, std::decay_t> -shfl(T x, int lane) -{ +shfl(T x, int lane) { return __shfl(x, lane); } @@ -142,6 +141,10 @@ __device__ __inline__ unsigned ballot(unsigned mask, unsigned is_root) { return __ballot(is_root); } +__device__ __inline__ unsigned active_mask() { + return __activemask(); +} + __device__ __inline__ unsigned any(unsigned mask, unsigned width) { return __any(width); } diff --git a/arbor/include/arbor/mechanism_abi.h b/arbor/include/arbor/mechanism_abi.h index b5bca6f24b..91f67a9f9e 100644 --- a/arbor/include/arbor/mechanism_abi.h +++ b/arbor/include/arbor/mechanism_abi.h @@ -54,7 +54,6 @@ inline const char* arb_mechanism_kind_str(const arb_mechanism_kind& mech) { // Ion state variables; view into shared_state typedef struct arb_ion_state { arb_value_type* current_density; - arb_value_type* conductivity; arb_value_type* reversal_potential; arb_value_type* internal_concentration; arb_value_type* external_concentration; diff --git a/doc/tutorial/full-feature-diffusion.rst b/doc/tutorial/full-feature-diffusion.rst new file mode 100644 index 0000000000..57b83c393f --- /dev/null +++ b/doc/tutorial/full-feature-diffusion.rst @@ -0,0 +1,160 @@ +.. _tutorialfullfeaturediffusion: + +A Diffusion Process with Full Feedback +====================================== + +Originally we intended for the axial diffusion process to be fully integrated +into the cable model (linked to the internal concentration, see :ref:`here `). However, for various technical reasons, this has proven +infeasible and was replaced by introduction of an extra variable ``Xd`` which +is subject to diffusion as opposed to the internal concentration. + +This tutorial will show you how to build a simulation that fully integrates +diffusive particle dynamics with a cable cell neuron. While we focus on a single +neuron, this will naturally work as well in a network of cells. + +Theory +------ + +In order to build a complete model, we will need to review the foundational +equations as implemented in Arbor. We assume a single ion (in general: particle) species :math:`s` with +valence :math:`q_s` and no leak currents such that :math:`I_m = I_s`. More ion +species can be added, but do not change the fundamental mechanisms. First, the +cable equation + +.. math:: + + \frac{1}{C}\dot U_m = \sigma \Delta U_m + I_m + +where + +.. math:: + + I_s = g_s (U_m - E_s) + +with the reversal potential :math:`E_s`. The current :math:`I_s` -- or actually +the conductance :math:`g_s` -- is usually computed by a density mechanism. + +Note that we use a slight simplifications here; :math:`\sigma` and the cable +radius :math:`r` are taken as constant along the dendrite. Second, the diffusion +equation under similar assumptions + +.. math:: + + \dot X_s = \beta \Delta X_s + +where :math:`X_s` is the diffusive concentration of species :math:`s`. Finally, +the Nernst equation + +.. math:: + + E_s = \frac{R\cdot T}{q_s\cdot F}\log\left(\frac{X_o}{X_i}\right) + +where :math:`X_{i,o}` are the internal/external concentrations of species +:math:`s` If your model does not use the Nernst equation to compute :math:`E_s`, +you can ignore this part, but be aware that this approach will not have a closed +feedback loop. + +Incorporating Diffusion +----------------------- + +We will leave the cable equation and the current computation unaltered. However, +from the intuition that a current is caused by movement of ions across the membrane, +one would expect :math:`I_s` to exact a change in :math:`X_s`. Thus, we will add an +extra term to the diffusion equation + +.. math:: + + \dot X_s = \beta \Delta X_s + \frac{I_s}{q_s\cdot F\cdot V} + +where Faraday's constant :math:`F` scales from current to change in molar +amount, in conjunction with the charge per ion :math:`q_s`, which we scale again +by the volume :math:`V` to arrive at a change in concentration. + +We implement this as an NMODL file intended to be added to the full cell that +should look like this (cf. :ref:`here `): + +.. code-block:: + + NEURON { + SUFFIX current_to_delta_x + USEION x READ ix WRITE xd, xi VALENCE zx + GLOBAL F + RANGE xi0 + } + + PARAMETER { + F = 96485.3321233100184 (coulomb/mole) : Faraday's constant + diam : compartment diameter + xi0 : initial concentration + } + + STATE { xi } + + INITIAL { + xi = xi0 + xd = xi0 + } + + BREAKPOINT { + : area = pi * length * radius + : volume = pi * length * radius^2 + : thus area/volume = radius + xd = xd + 2*ix/(q*F*diam) + xi = xd + } + +where we exploit our knowledge about the cylindrical geometry of the CVs to +compute :math:`\frac{A}{V} = \frac{1}{r}`. This mechanism also turns the +internal concentration ``xi`` into a ``STATE`` variable, which is the standard +way of handling concentration mechanisms, and couples it to ``xd`` directly. +Note that this requires explicit intialisation of both ``xi`` and ``xd``. + +Setting up a Simulation +----------------------- + +Having laid the foundation, adding this to a simulation is pretty simple. Save +the NMODL file, add it to your local catalogue, and compile & load that via the +usual method (cf. :ref:`here `). + +First, declare the ion -- we'll use a new species ``X`` here, but any name will +do -- by calling ``set_ion("X", valence=1)`` on your global properties object. +This object is part of the ``recipe`` or ``single_cell_model``, the use of both +is explained in other tutorials. When using an existing ion, this upfront +declaration is not needed and you'll get the default values for this ion. + +.. code-block:: python + + dec = (A.decor() + # Add our new ion to the cell; the `int_con` value has no effect. + .set_ion("X", int_con=0.0, ext_con=42.0, diff=0.005, method="nernst/X") ) + # Place a synapse that _directly_ adds to the diffusive concentration + .place("(location 0 0.5)", A.synapse("inject/x=X", {"alpha": 200.0}), "Zap") + # also add an exponential decay to Xd + .paint("(all)", A.density("decay/x=X")) + # turn iX into a change in Xd and bind Xi to Xd + .paint("(all)", A.density("current_to_delta_x/x=X", {"xi0": 10.0})) + # ... + ) + +While simple, note some subtleties around our custom concentration mechanism: +- The mechanism ``current_to_delta_x`` uses ``xi`` as a ``STATE`` and is thus + solely responsible for managing its value. This makes adding an explicit + initialisation via ``xi0`` necessary. Only one mechanism with this property + should exist. See above for an alternative. +- The change in ``xd`` due to events arriving at the synapse ``Zap`` will be + synchronised with ``xi`` in our custom mechanism. If no concentration + mechanism is used, the synapse needs to be modified to write to ``xi`` as well. +- By using ``xi=xd``, the Nernst mechanism will pick up the correct value for + ``xi``. If that is not your intention, you will have to provide a modified + version of ``nernst`` in which ``xi`` is replaced with ``xd``. + +Final notes +----------- + +Although the related theory is somewhat intricate, as this tutorial has shown, adding an ion (or another particle) with diffusion and full +transmembrane current dynamics to a simulation is quite straightforward. You +might also want to consider changing the external concentration ``Xo`` according to the +ion current ``iX``. This was not shown here for two reasons. First, Arbor does +not handle extra-cellular dynamics and thus has no extra-cellular diffusion. +Second, the method for handling this is identical to what we have done for +``xi``, so including it doesn't add any insight. diff --git a/doc/tutorial/index.rst b/doc/tutorial/index.rst index be8790a0cc..3ce7f9d8d9 100644 --- a/doc/tutorial/index.rst +++ b/doc/tutorial/index.rst @@ -1,3 +1,4 @@ + .. _tutorial: Tutorials @@ -81,6 +82,7 @@ Advanced nmodl plasticity connectivity + full-feature-diffusion Demonstrations -------------- diff --git a/modcc/identifier.hpp b/modcc/identifier.hpp index 08a3f40fdb..648891c24f 100644 --- a/modcc/identifier.hpp +++ b/modcc/identifier.hpp @@ -49,7 +49,6 @@ enum class sourceKind { dt, time, ion_current, - ion_conductivity, ion_current_density, ion_diffusive, ion_revpot, diff --git a/modcc/module.cpp b/modcc/module.cpp index 8289f5e74f..acd28f1c01 100644 --- a/modcc/module.cpp +++ b/modcc/module.cpp @@ -37,7 +37,6 @@ class NrnCurrentRewriter: public BlockRewriterBase { if (src==sourceKind::current_density || src==sourceKind::current || src==sourceKind::ion_current_density || - src==sourceKind::ion_conductivity || src==sourceKind::ion_current) { return src; @@ -59,8 +58,6 @@ class NrnCurrentRewriter: public BlockRewriterBase { using BlockRewriterBase::visit; virtual void finalize() override { - // Take current name 'iX' and strip off leading 'i' to get ion name. - auto i2ion = [this](const auto& name) { return id("conductivity_" + name.substr(1) + "_"); }; if (has_current_update_) { expression_ptr current_sum, conductivity_sum; for (auto& curr: current_vars_) { @@ -79,11 +76,6 @@ class NrnCurrentRewriter: public BlockRewriterBase { conductivity_sum = make_expression( Location{}, std::move(conductivity_sum), cond->clone()); } - if (name != non_specific_current) { - statements_.push_back(make_expression(loc_, - i2ion(name), - cond->clone())); - } } if (current_sum) { statements_.push_back(make_expression(loc_, @@ -825,11 +817,6 @@ bool Module::add_variables_to_symbols() { for(auto const& ion: neuron_block_.ions) { for(auto const& var: ion.write) { update_ion_symbols(var, accessKind::write, ion.name); - auto name = "conductivity_" + ion.name + "_"; - if (cond.find(name) == cond.end()) { - create_indexed_variable(name, sourceKind::ion_conductivity, accessKind::write, ion.name, var.location); - cond.insert(name); - } } for(auto const& var: ion.read) { diff --git a/modcc/printer/cprinter.cpp b/modcc/printer/cprinter.cpp index 36ec914e6d..b3b87649ee 100644 --- a/modcc/printer/cprinter.cpp +++ b/modcc/printer/cprinter.cpp @@ -1,4 +1,3 @@ -#include #include #include #include @@ -19,7 +18,6 @@ using io::indent; using io::popindent; -using io::quote; constexpr bool with_profiling() { #ifdef ARB_HAVE_PROFILING @@ -129,6 +127,7 @@ ARB_LIBMODCC_API std::string emit_cpp_source(const Module& module_, const printe APIMethod* write_ions_api = find_api_method(module_, "write_ions"); bool with_simd = opt.simd.abi!=simd_spec::none; + bool is_point_proc = (module_.kind() == moduleKind::point) || (module_.kind() == moduleKind::junction); options_trace_codegen = opt.trace_codegen; @@ -320,6 +319,8 @@ ARB_LIBMODCC_API std::string emit_cpp_source(const Module& module_, const printe out << popindent << "}\n\n"; if (net_receive_api) { + auto flags = net_recv_flags; + flags.point(is_point_proc); out << fmt::format(FMT_COMPILE("static void apply_events(arb_mechanism_ppack* pp, arb_deliverable_event_stream* stream_ptr) {{\n" " PPACK_IFACE_BLOCK;\n" " auto [begin_, end_] = *stream_ptr;\n" @@ -327,13 +328,15 @@ ARB_LIBMODCC_API std::string emit_cpp_source(const Module& module_, const printe " [[maybe_unused]] auto [i_, {0}] = *begin_;\n"), net_receive_api->args().empty() ? "weight" : net_receive_api->args().front()->is_argument()->name()); out << indent << indent; - emit_api_body(out, net_receive_api, net_recv_flags); + emit_api_body(out, net_receive_api, flags); out << popindent << "}\n" << popindent << "}\n\n"; } else { out << "static void apply_events(arb_mechanism_ppack*, arb_deliverable_event_stream*) {}\n\n"; } if(post_event_api) { + auto flags = post_evt_flags; + flags.point(is_point_proc); const std::string time_arg = post_event_api->args().empty() ? "time" : post_event_api->args().front()->is_argument()->name(); out << fmt::format(FMT_COMPILE("static void post_event(arb_mechanism_ppack* pp) {{\n" " PPACK_IFACE_BLOCK;\n" @@ -347,7 +350,7 @@ ARB_LIBMODCC_API std::string emit_cpp_source(const Module& module_, const printe pp_var_pfx, time_arg); out << indent << indent << indent << indent; - emit_api_body(out, post_event_api, post_evt_flags); + emit_api_body(out, post_event_api, flags); out << popindent << "}\n" << popindent << "}\n" << popindent << "}\n" << popindent << "}\n"; } else { out << "static void post_event(arb_mechanism_ppack*) {}\n"; diff --git a/modcc/printer/gpuprinter.cpp b/modcc/printer/gpuprinter.cpp index 00e040549e..03adf44b64 100644 --- a/modcc/printer/gpuprinter.cpp +++ b/modcc/printer/gpuprinter.cpp @@ -1,7 +1,5 @@ -#include #include #include -#include #include #include @@ -10,14 +8,12 @@ #include "gpuprinter.hpp" #include "expression.hpp" -#include "io/ostream_wrappers.hpp" #include "io/prefixbuf.hpp" #include "printer/cexpr_emit.hpp" #include "printer/printerutil.hpp" using io::indent; using io::popindent; -using io::quote; static std::string scaled(double coeff) { std::stringstream ss; @@ -221,18 +217,22 @@ ARB_LIBMODCC_API std::string emit_gpu_cu_source(const Module& module_, const pri "}}\n\n"), pp_var_pfx); } - emit_api_kernel(state_api); + // TODO This needs additive? + emit_api_kernel(state_api, true); emit_api_kernel(current_api, true); emit_api_kernel(write_ions_api); // event delivery if (net_receive_api) { + ApiFlags flags = {false, false, true}; // No CV loop, no PPACK, use additive + flags.point(is_point_proc); out << fmt::format(FMT_COMPILE("__global__\n" "void apply_events(arb_mechanism_ppack params_, const arb_deliverable_event_data* begin_,\n" " const arb_deliverable_event_data* end_) {{\n" " PPACK_IFACE_BLOCK;\n" " const auto ii_ = threadIdx.x + blockDim.x*blockIdx.x;\n" - " if (ii_ < (end_ - begin_)) {{\n" + " auto n_ = end_ - begin_;\n" + " if (ii_ < n_) {{\n" " const auto tid_ = (begin_ + ii_)->mech_index;\n" " if ((ii_ > 0) && ((begin_ + (ii_ - 1))->mech_index == tid_)) return;\n" " for (auto i_ = begin_ + ii_; i_ < end_; ++i_) {{\n" @@ -240,7 +240,7 @@ ARB_LIBMODCC_API std::string emit_gpu_cu_source(const Module& module_, const pri " [[maybe_unused]] auto {0} = i_->weight;\n"), net_receive_api->args().empty() ? "weight" : net_receive_api->args().front()->is_argument()->name()); out << indent << indent << indent; - emit_api_body_cu(out, net_receive_api, ApiFlags{}.point(is_point_proc).loop(false).iface(false)); + emit_api_body_cu(out, net_receive_api, flags); out << popindent << "}\n" << popindent << "}\n" << popindent << "}\n\n"; } @@ -313,20 +313,15 @@ ARB_LIBMODCC_API std::string emit_gpu_cu_source(const Module& module_, const pri } else { emit_empty_wrapper("post_event"); } - if (net_receive_api) { - auto api_name = "apply_events"; - out << fmt::format(FMT_COMPILE("void {}_{}_(arb_mechanism_ppack* p, arb_deliverable_event_stream* stream_ptr) {{"), class_name, api_name); - if(!net_receive_api->body()->statements().empty()) { - out << fmt::format(FMT_COMPILE("\n" - " const arb_deliverable_event_data* const begin = stream_ptr->begin;\n" - " const arb_deliverable_event_data* const end = stream_ptr->end;\n" - " ::arb::gpu::launch_1d(end - begin, 128, {}, *p, begin, end);\n"), - api_name); - } - out << "}\n\n"; + if (net_receive_api && !net_receive_api->body()->statements().empty()) { + out << fmt::format(FMT_COMPILE("void {}_apply_events_(arb_mechanism_ppack* p, arb_deliverable_event_stream* stream_ptr) {{\n" + " const arb_deliverable_event_data* const begin = stream_ptr->begin;\n" + " const arb_deliverable_event_data* const end = stream_ptr->end;\n" + " ::arb::gpu::launch_1d(end - begin, 128, apply_events, *p, begin, end);\n" + "}}\n\n"), class_name); } else { - auto api_name = "apply_events"; - out << fmt::format(FMT_COMPILE("void {}_{}_(arb_mechanism_ppack* p, arb_deliverable_event_stream* events) {{}}\n\n"), class_name, api_name); + out << fmt::format(FMT_COMPILE("void {}_apply_events_(arb_mechanism_ppack* p, arb_deliverable_event_stream* events) {{}}\n\n"), + class_name); } out << namespace_declaration_close(ns_components); return out.str(); @@ -394,8 +389,8 @@ void emit_api_body_cu(std::ostream& out, APIMethod* e, const ApiFlags& flags) { // update an indexed variable, like current or conductance. // This is the case if one of the external variables "is_write". auto it = std::find_if(indexed_vars.begin(), indexed_vars.end(), - [](auto& sym){return sym->external_variable()->is_write();}); - if (it!=indexed_vars.end()) { + [](auto& sym){return sym->external_variable()->is_write();}); + if (it != indexed_vars.end()) { out << "unsigned lane_mask_ = arb::gpu::ballot(0xffffffff, tid_ valence, optional int_con, - optional ext_con, optional rev_pot, - py::object method, optional diff) { + optional valence, + optional int_con, + optional ext_con, + optional rev_pot, + py::object method, + optional diff) { if (!props.ion_species.count(ion) && !valence) { throw std::runtime_error(util::pprintf("New ion species: '{}', missing valence", ion)); } @@ -672,7 +675,7 @@ void register_cells(py::module& m) { props.default_parameters.reversal_potential_method[ion] = *m; } }, - "ion"_a, "valence"_a=py::none(), "int_con"_a=py::none(), "ext_con"_a=py::none(), "rev_pot"_a=py::none(), "method"_a =py::none(), "diff"_a=py::none(), + "ion"_a, py::kw_only(), "valence"_a=py::none(), "int_con"_a=py::none(), "ext_con"_a=py::none(), "rev_pot"_a=py::none(), "method"_a =py::none(), "diff"_a=py::none(), "Set the global default properties of ion species named 'ion'.\n" " * valence: valence of the ion species [e].\n" " * int_con: initial internal concentration [mM].\n" @@ -742,7 +745,7 @@ void register_cells(py::module& m) { if (tempK) d.set_default(arb::temperature{*tempK}); return d; }, - "Vm"_a=py::none(), "cm"_a=py::none(), "rL"_a=py::none(), "tempK"_a=py::none(), + py::kw_only(), "Vm"_a=py::none(), "cm"_a=py::none(), "rL"_a=py::none(), "tempK"_a=py::none(), "Set default values for cable and cell properties:\n" " * Vm: initial membrane voltage [mV].\n" " * cm: membrane capacitance [F/m²].\n" @@ -762,7 +765,7 @@ void register_cells(py::module& m) { if (auto m = maybe_method(method)) d.set_default(arb::ion_reversal_potential_method{ion, *m}); return d; }, - "ion"_a, "int_con"_a=py::none(), "ext_con"_a=py::none(), "rev_pot"_a=py::none(), "method"_a =py::none(), "diff"_a=py::none(), + "ion"_a, py::kw_only(), "int_con"_a=py::none(), "ext_con"_a=py::none(), "rev_pot"_a=py::none(), "method"_a =py::none(), "diff"_a=py::none(), "Set the cell-level properties of ion species named 'ion'.\n" " * int_con: initial internal concentration [mM].\n" " * ext_con: initial external concentration [mM].\n" @@ -844,7 +847,7 @@ void register_cells(py::module& m) { } return dec; }, - "region"_a, "Vm"_a=py::none(), "cm"_a=py::none(), "rL"_a=py::none(), "tempK"_a=py::none(), + "region"_a, py::kw_only(), "Vm"_a=py::none(), "cm"_a=py::none(), "rL"_a=py::none(), "tempK"_a=py::none(), "Set cable properties on a region.\n" "Set global default values for cable and cell properties.\n" " * Vm: initial membrane voltage [mV].\n" diff --git a/python/context.cpp b/python/context.cpp index 23f128c2d0..7b9f2b4025 100644 --- a/python/context.cpp +++ b/python/context.cpp @@ -105,7 +105,6 @@ context_shim make_context_shim(proc_allocation_shim alloc, pybind11::object mpi, context_shim make_context_shim() { return context_shim{arb::make_context(arbenv::default_allocation())}; - }; // pybind diff --git a/python/env.cpp b/python/env.cpp index d10a9bf5c0..f9343a9c4e 100644 --- a/python/env.cpp +++ b/python/env.cpp @@ -1,6 +1,8 @@ #include #include "conversion.hpp" +#include + #include #include #include @@ -23,7 +25,7 @@ namespace pyarb { #else auto err = "Private GPU: Invalid MPI Communicator."; if (can_convert_to_mpi_comm(mpi)) { - return arbenv::find_private_gpu(can_convert_to_mpi_comm(mpi)); + return arbenv::find_private_gpu(convert_to_mpi_comm(mpi)); } else if (auto c = py2optional(mpi, err)) { return arbenv::find_private_gpu(c->comm); diff --git a/python/example/connectivity/brunel.py b/python/example/connectivity/brunel.py index 1426880994..784a61598f 100644 --- a/python/example/connectivity/brunel.py +++ b/python/example/connectivity/brunel.py @@ -36,7 +36,7 @@ def network_description(self): (random {self.seed} {self.p}))""" inh = f"(gid-range {self.n_exc} {self.N})" weight = f"""(if-else (source-cell {inh}) - (scalar {self.g*self.weight}) + (scalar {self.g * self.weight}) (scalar {self.weight}))""" delay = "(scalar 0.5)" return A.network_description(rand, weight, delay, {}) diff --git a/python/test/fixtures.py b/python/test/fixtures.py index 1071653855..e9e5b15e1f 100644 --- a/python/test/fixtures.py +++ b/python/test/fixtures.py @@ -4,12 +4,8 @@ from functools import lru_cache as cache from pathlib import Path import subprocess -import atexit import inspect -_mpi_enabled = A.__config__["mpi"] -_mpi4py_enabled = A.__config__["mpi4py"] - # The API of `functools`'s caches went through a bunch of breaking changes from # 3.6 to 3.9. Patch them up in a local `cache` function. try: @@ -72,14 +68,23 @@ def repo_path(): return Path(__file__).parent.parent.parent -def _finalize_mpi(): - print("Context fixture finalizing mpi") - if _mpi4py_enabled: - from mpi4py import MPI +def get_mpi_comm_world(): + """ + Obtain MPI_COMM_WORLD as --- in order --- + 1. MPI4PY.MPI.COMM_WORLD + 2. Arbor MPI + 3. None + """ + if A.config()["mpi"]: + if A.config()["mpi4py"]: + from mpi4py import MPI - MPI.Finalize() - else: - A.mpi_finalize() + return MPI.COMM_WORLD + else: + if not A.mpi_is_initialized(): + A.mpi_init() + return A.mpi_comm() + return None @_fixture @@ -87,21 +92,20 @@ def context(): """ Fixture that produces an MPI sensitive `A.context` """ - if _mpi_enabled: - if _mpi4py_enabled: - from mpi4py import MPI + return A.context(mpi=get_mpi_comm_world()) - if not MPI.Is_initialized(): - print("Context fixture initializing mpi4py", flush=True) - MPI.Initialize() - atexit.register(_finalize_mpi) - return A.context(A.proc_allocation(), mpi=MPI.COMM_WORLD) - elif not A.mpi_is_initialized(): - print("Context fixture initializing mpi", flush=True) - A.mpi_init() - atexit.register(_finalize_mpi) - return A.context(A.proc_allocation(), mpi=A.mpi_comm()) - return A.context(A.proc_allocation()) + +@_fixture +def single_context(): + """ + Fixture that produces an non-MPI `A.context`, which includes + a GPU, if enabled. + """ + mpi = None + gpu = None + if A.config()["gpu"]: + gpu = 0 + return A.context(mpi=mpi, gpu_id=gpu) class _BuildCatError(Exception): @@ -174,6 +178,19 @@ def dummy_catalogue(repo_path): return A.load_catalogue(str(cat_path)) +@_singleton_fixture +@repo_path() +def diffusion_catalogue(repo_path): + """ + Fixture that returns an `arbor.catalogue` + which contains mechanisms `neuron_with_diffusion` + and `synapse_with_diffusion`. + """ + path = repo_path / "test" / "unit" / "diffusion" + cat_path = _build_cat("diffusion", path) + return A.load_catalogue(str(cat_path)) + + @_fixture class empty_recipe(A.recipe): """ diff --git a/python/test/readme.md b/python/test/readme.md index 515366df87..ac4cc673e0 100644 --- a/python/test/readme.md +++ b/python/test/readme.md @@ -39,6 +39,6 @@ In subfolders `unit`/`unit_distributed`: ## Naming convention -- modules: `test_xxxs` (ending with `s` since module can consist of multiple classes) -- class(es): `TestXxxs` (ending with `s` since class can consist of multiple test functions) +- modules: `test_xxxs` (if applicable use the plural with `s` as modules can comprise multiple classes) +- class(es): `TestXxxs` (if applicable use the plural with `s` as classes can comprise multiple test functions) - functions: `test_yyy` diff --git a/python/test/unit/test_diffusion.py b/python/test/unit/test_diffusion.py new file mode 100644 index 0000000000..946c9c894a --- /dev/null +++ b/python/test/unit/test_diffusion.py @@ -0,0 +1,430 @@ +# -*- coding: utf-8 -*- + +import unittest +import arbor as A +from arbor import units as U +import numpy as np +from .. import fixtures + +""" +Tests for the concentration and amount of diffusive particles across time and morphology. +Three different morphological structures are considered: 1 segment ("soma only"), 2 segments +("soma with dendrite"), and 3 segments ("soma with two dendrites"). + +NOTE: Internally, Arbor only knows concentrations. Thus, particle amounts have to be computed + from concentrations by integrating over the volume of the morphology. The total amount + of particles should be conserved unless there is deliberate injection or removal of + particles. +""" + + +class recipe(A.recipe): + # Constructor + # - cat: catalogue of custom mechanisms + # - cell: cell description + # - probes: list of probes + # - inject_remove: list of dictionaries of the form [ {"time" :