Ux/units (#2233)

The core issue here is to add units to the user facing API. I decided on using the LLNL/units library, which offers conversion and checking at runtime. Runtime is a requirement -- as much as I love static guarantees --, but keeping the interface uniform between Python and C++ is a must. While setting this up, I noticed the severe lack of IDE/LSP support for Arbor, so I added typing stubs using https://github.com/sizmailov/pybind11-stubgen. The conjunction of typing and units exposed misuse of pybind11 in several places, so next I had to massage the ordering of bindings, adjust the specification of default arguments, and add the odd missing binding. The schedule/event generator interface was tightened up, hiding the `*_impl` structs and exposing only the type erased `schedule` object. That in turn required de-generification of the Poisson schedule. Now, Mersenne twister is the only choice and I will remove that later on for the CBRNG we are already using elsewhere. Currently, units are used for: - [X] simulation - [X] schedule/generator - [x] paintables - [X] placeables - [X] iclamp - [X] threshold - [X] connections - [X] gap junctions Adding units to mechanism interfaces is _interesting_ but requires more work and thought, so I'll defer that to a later point in time. We'd need to adjust modcc to expose and **check** units and devise a scheme to handle missing units. Generic TODOs; some might spin off into separate issues. - [x] ~~rename py::iclamp OR cpp::i_clamp for consistency~~ covered by #2239 - [x] use scale/base for iexpr paintables for consistency with scaled_mech - [x] ~~Use CBRNG for Poisson schedule~~ covered by #2243 - [ ] Automate stub generation. A wishlist item, requires installing extra software. - [x] Properly integrate units w/ spack. NB. Units doesn't have a spackage. Closes #1983 Closes #2032 --------- Co-authored-by: boeschf <[email protected]>
arbor-sim · Jan 19, 2024 · 60445a4 · 60445a4
1 parent 9f18a50
commit 60445a4
Show file tree

Hide file tree

Showing 162 changed files with 28,834 additions and 24,765 deletions.
diff --git a/.github/workflows/test-spack.yml b/.github/workflows/test-spack.yml
@@ -35,6 +35,7 @@ jobs:
         uses: actions/checkout@v3
         with:
           path: arbor
+          submodules: recursive
 
       - name: clone spack develop
         if: ${{ matrix.spack-version == 'develop' }}

diff --git a/.gitmodules b/.gitmodules
@@ -26,3 +26,6 @@
 	path = ext/pugixml
 	url = https://github.com/zeux/pugixml.git
 	branch = master
+[submodule "ext/units"]
+	path = ext/units
+	url = https://github.com/LLNL/units.git
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -259,6 +259,9 @@ install(FILES mechanisms/BuildModules.cmake DESTINATION ${ARB_INSTALL_DATADIR})
 cmake_dependent_option(ARB_USE_BUNDLED_FMT "Use bundled FMT lib." ON "ARB_USE_BUNDLED_LIBS" OFF)
 cmake_dependent_option(ARB_USE_BUNDLED_PUGIXML "Use bundled XML lib." ON "ARB_USE_BUNDLED_LIBS" OFF)
 cmake_dependent_option(ARB_USE_BUNDLED_GTEST "Use bundled GoogleTest." ON "ARB_USE_BUNDLED_LIBS" OFF)
+# TODO When we get a units spack package...
+#cmake_dependent_option(ARB_USE_BUNDLED_UNITS "Use bundled LLNL units." ON "ARB_USE_BUNDLED_LIBS" OFF)
+set(ARB_USE_BUNDLED_UNITS ON CACHE STRING "Use bundled LLNL units.")
 
 cmake_dependent_option(ARB_USE_BUNDLED_JSON "Use bundled Niels Lohmann's json library." ON "ARB_USE_BUNDLED_LIBS" OFF)
 if(NOT ARB_USE_BUNDLED_JSON)
@@ -290,9 +293,20 @@ else()
     endif()
 endif()
 
+add_library(ext-units INTERFACE)
+if(ARB_USE_BUNDLED_UNITS)
+    target_link_libraries(ext-units INTERFACE units::units)
+else()
+    message(FATAL, "TODO: At the time of Arbor 0.10.0 there is no Spack package")
+endif()
+
+
 add_subdirectory(ext)
 install(TARGETS ext-hwloc EXPORT arbor-targets)
 install(TARGETS ext-random123 EXPORT arbor-targets)
+target_link_libraries(arbor-public-deps INTERFACE ext-units)
+install(TARGETS ext-units EXPORT arbor-targets)
+install(TARGETS units compile_flags_target EXPORT arbor-targets)
 
 # Keep track of packages we need to add to the generated CMake config
 # file for arbor.

diff --git a/arbor/backends/gpu/threshold_watcher.hpp b/arbor/backends/gpu/threshold_watcher.hpp
@@ -72,7 +72,7 @@ class threshold_watcher {
         v_prev_(num_cv),
         // TODO: allocates enough space for 10 spikes per watch.
         // A more robust approach might be needed to avoid overflows.
-        stack_(10*size(), context.gpu)
+        stack_(100*size(), context.gpu)
     {
         crossings_.reserve(stack_.capacity());
         // reset() needs to be called before this is ready for use

diff --git a/arbor/cable_cell.cpp b/arbor/cable_cell.cpp
@@ -2,17 +2,13 @@
 #include <sstream>
 #include <unordered_map>
 #include <variant>
-#include <vector>
 
 #include <arbor/cable_cell.hpp>
 #include <arbor/morph/label_dict.hpp>
 #include <arbor/morph/morphology.hpp>
 #include <arbor/morph/mprovider.hpp>
 #include <arbor/util/pp_util.hpp>
 
-#include "util/piecewise.hpp"
-#include "util/rangeutil.hpp"
-#include "util/span.hpp"
 #include "util/strprintf.hpp"
 
 namespace arb {
@@ -37,7 +33,7 @@ std::string show(const paintable& item) {
             else if constexpr (std::is_same_v<axial_resistivity, T>) {
                 os << "axial-resistivity";
             }
-            else if constexpr (std::is_same_v<temperature_K, T>) {
+            else if constexpr (std::is_same_v<temperature, T>) {
                 os << "temperature-kelvin";
             }
             else if constexpr (std::is_same_v<membrane_capacitance, T>) {

diff --git a/arbor/cable_cell_param.cpp b/arbor/cable_cell_param.cpp
@@ -75,30 +75,30 @@ cable_cell_parameter_set neuron_parameter_defaults = {
 std::vector<defaultable> cable_cell_parameter_set::serialize() const {
     std::vector<defaultable> D;
     if (init_membrane_potential) {
-        D.push_back(arb::init_membrane_potential{*this->init_membrane_potential});
+        D.push_back(arb::init_membrane_potential{*this->init_membrane_potential*units::mV});
     }
     if (temperature_K) {
-        D.push_back(arb::temperature_K{*this->temperature_K});
+        D.push_back(arb::temperature{*this->temperature_K*units::Kelvin});
     }
     if (axial_resistivity) {
-        D.push_back(arb::axial_resistivity{*this->axial_resistivity});
+        D.push_back(arb::axial_resistivity{*this->axial_resistivity*units::Ohm*units::cm});
     }
     if (membrane_capacitance) {
-        D.push_back(arb::membrane_capacitance{*this->membrane_capacitance});
+        D.push_back(arb::membrane_capacitance{*this->membrane_capacitance*units::F/units::m2});
     }
 
     for (const auto& [name, data]: ion_data) {
         if (data.init_int_concentration) {
-            D.push_back(init_int_concentration{name, *data.init_int_concentration});
+            D.push_back(init_int_concentration{name, *data.init_int_concentration*units::mM});
         }
         if (data.init_ext_concentration) {
-            D.push_back(init_ext_concentration{name, *data.init_ext_concentration});
+            D.push_back(init_ext_concentration{name, *data.init_ext_concentration*units::mM});
         }
         if (data.init_reversal_potential) {
-            D.push_back(init_reversal_potential{name, *data.init_reversal_potential});
+            D.push_back(init_reversal_potential{name, *data.init_reversal_potential*units::mV});
         }
         if (data.diffusivity) {
-            D.push_back(ion_diffusivity{name, *data.diffusivity});
+            D.push_back(ion_diffusivity{name, *data.diffusivity*units::m2/units::s});
         }
     }
 
@@ -133,32 +133,32 @@ decor& decor::set_default(defaultable what) {
             [this] (auto&& p) {
                 using T = std::decay_t<decltype(p)>;
                 if constexpr (std::is_same_v<init_membrane_potential, T>) {
-                    if (p.value.type() != iexpr_type::scalar) throw cable_cell_error{"Default values cannot have a scale."};
-                    defaults_.init_membrane_potential = *p.value.get_scalar();
+                    if (p.scale.type() != iexpr_type::scalar) throw cable_cell_error{"Default values cannot have a scale."};
+                    defaults_.init_membrane_potential = *p.scale.get_scalar()*p.value;
                 }
                 else if constexpr (std::is_same_v<axial_resistivity, T>) {
-                    if (p.value.type() != iexpr_type::scalar) throw cable_cell_error{"Default values cannot have a scale."};
-                    defaults_.axial_resistivity = *p.value.get_scalar();
+                    if (p.scale.type() != iexpr_type::scalar) throw cable_cell_error{"Default values cannot have a scale."};
+                    defaults_.axial_resistivity = *p.scale.get_scalar()*p.value;
                 }
-                else if constexpr (std::is_same_v<temperature_K, T>) {
-                    if (p.value.type() != iexpr_type::scalar) throw cable_cell_error{"Default values cannot have a scale."};
-                    defaults_.temperature_K = *p.value.get_scalar();
+                else if constexpr (std::is_same_v<temperature, T>) {
+                    if (p.scale.type() != iexpr_type::scalar) throw cable_cell_error{"Default values cannot have a scale."};
+                    defaults_.temperature_K = *p.scale.get_scalar()*p.value;
                 }
                 else if constexpr (std::is_same_v<membrane_capacitance, T>) {
-                    if (p.value.type() != iexpr_type::scalar) throw cable_cell_error{"Default values cannot have a scale."};
-                    defaults_.membrane_capacitance = *p.value.get_scalar();
+                    if (p.scale.type() != iexpr_type::scalar) throw cable_cell_error{"Default values cannot have a scale."};
+                    defaults_.membrane_capacitance = *p.scale.get_scalar()*p.value;
                 }
                 else if constexpr (std::is_same_v<init_int_concentration, T>) {
-                    if (p.value.type() != iexpr_type::scalar) throw cable_cell_error{"Default values cannot have a scale."};
-                    defaults_.ion_data[p.ion].init_int_concentration = *p.value.get_scalar();
+                    if (p.scale.type() != iexpr_type::scalar) throw cable_cell_error{"Default values cannot have a scale."};
+                    defaults_.ion_data[p.ion].init_int_concentration = *p.scale.get_scalar()*p.value;
                 }
                 else if constexpr (std::is_same_v<init_ext_concentration, T>) {
-                    if (p.value.type() != iexpr_type::scalar) throw cable_cell_error{"Default values cannot have a scale."};
-                    defaults_.ion_data[p.ion].init_ext_concentration = *p.value.get_scalar();
+                    if (p.scale.type() != iexpr_type::scalar) throw cable_cell_error{"Default values cannot have a scale."};
+                    defaults_.ion_data[p.ion].init_ext_concentration = *p.scale.get_scalar()*p.value;
                 }
                 else if constexpr (std::is_same_v<init_reversal_potential, T>) {
-                    if (p.value.type() != iexpr_type::scalar) throw cable_cell_error{"Default values cannot have a scale."};
-                    defaults_.ion_data[p.ion].init_reversal_potential = *p.value.get_scalar();
+                    if (p.scale.type() != iexpr_type::scalar) throw cable_cell_error{"Default values cannot have a scale."};
+                    defaults_.ion_data[p.ion].init_reversal_potential = *p.scale.get_scalar()*p.value;
                 }
                 else if constexpr (std::is_same_v<ion_reversal_potential_method, T>) {
                     defaults_.reversal_potential_method[p.ion] = p.method;
@@ -167,8 +167,9 @@ decor& decor::set_default(defaultable what) {
                     defaults_.discretization = std::forward<cv_policy>(p);
                 }
                 else if constexpr (std::is_same_v<ion_diffusivity, T>) {
-                    if (p.value.type() != iexpr_type::scalar) throw cable_cell_error{"Default values cannot have a scale."};
-                    defaults_.ion_data[p.ion].diffusivity = p.value.get_scalar();
+                    if (p.scale.type() != iexpr_type::scalar) throw cable_cell_error{"Default values cannot have a scale."};
+                    auto s = p.scale.get_scalar();
+                    defaults_.ion_data[p.ion].diffusivity = s ? std::optional{*s*p.value} : s;
                 }
             },
             what);

diff --git a/arbor/fvm_layout.cpp b/arbor/fvm_layout.cpp
@@ -1,7 +1,6 @@
 #include <algorithm>
 #include <optional>
 #include <set>
-#include <stdexcept>
 #include <unordered_set>
 #include <unordered_map>
 #include <vector>
@@ -20,17 +19,11 @@
 #include "fvm_layout.hpp"
 #include "threading/threading.hpp"
 #include "util/maputil.hpp"
-#include "util/meta.hpp"
-#include "util/partition.hpp"
 #include "util/piecewise.hpp"
 #include "util/pw_over_cable.hpp"
 #include "util/rangeutil.hpp"
-#include "util/transform.hpp"
-#include "util/unique.hpp"
 #include "util/strprintf.hpp"
 
-#include <iostream>
-
 namespace arb {
 
 using util::assign;
@@ -275,16 +268,16 @@ fvm_cv_discretize(const cable_cell& cell, const cable_cell_parameter_set& global
 
     double dflt_resistivity = *(dflt.axial_resistivity | global_dflt.axial_resistivity);
     double dflt_capacitance = *(dflt.membrane_capacitance | global_dflt.membrane_capacitance);
-    double dflt_potential =   *(dflt.init_membrane_potential | global_dflt.init_membrane_potential);
+    double dflt_potential   = *(dflt.init_membrane_potential | global_dflt.init_membrane_potential);
     double dflt_temperature = *(dflt.temperature_K | global_dflt.temperature_K);
 
-    const auto& assignments = cell.region_assignments();
-    const auto& resistivity = assignments.get<axial_resistivity>();
-    const auto& capacitance = assignments.get<membrane_capacitance>();
-    const auto& potential   = assignments.get<init_membrane_potential>();
-    const auto& temperature = assignments.get<temperature_K>();
-    const auto& diffusivity = assignments.get<ion_diffusivity>();
-    const auto& provider    = cell.provider();
+    const auto& assignments   = cell.region_assignments();
+    const auto& resistivity   = assignments.get<axial_resistivity>();
+    const auto& capacitance   = assignments.get<membrane_capacitance>();
+    const auto& potential     = assignments.get<init_membrane_potential>();
+    const auto& temperature_K = assignments.get<temperature>();
+    const auto& diffusivity   = assignments.get<ion_diffusivity>();
+    const auto& provider      = cell.provider();
 
     struct inv_diff {
         iexpr value;
@@ -310,8 +303,9 @@ fvm_cv_discretize(const cable_cell& cell, const cable_cell_parameter_set& global
         auto diffusive = std::any_of(data.begin(),
                                      data.end(),
                                      [](const auto& kv) {
-                                         const auto& v = kv.second.value.get_scalar();
-                                         return !v || *v != 0.0 || *v == *v;
+                                         const auto& [k, v] = kv;
+                                         auto s = v.scale.get_scalar();
+                                         return !s || *s*v.value != 0.0;
                                      });
         if (diffusive) {
             // Provide a (non-sensical) default.
@@ -337,7 +331,7 @@ fvm_cv_discretize(const cable_cell& cell, const cable_cell_parameter_set& global
         for (msize_t i = 0; i<n_branch; ++i) {
             auto cable = mcable{i, 0., 1.};
             auto scale_param = [&, ion=ion](const auto&,
-                                   const inv_diff& par) {
+                                   const inv_diff& par) -> double {
                 auto ie = thingify(par.value, provider);
                 auto sc = ie->eval(provider, cable);
                 if (def <= 0.0 || std::isnan(def)) {
@@ -360,9 +354,9 @@ fvm_cv_discretize(const cable_cell& cell, const cable_cell_parameter_set& global
     for (msize_t i = 0; i<n_branch; ++i) {
         auto cable = mcable{i, 0., 1.};
         auto scale_param = [&](const auto&,
-                               const axial_resistivity& par) {
-            auto ie = thingify(par.value, provider);
-            auto sc = ie->eval(provider, cable);
+                               const axial_resistivity& par) -> double {
+            auto ie = thingify(par.scale, provider);
+            auto sc = par.value*ie->eval(provider, cable);
             return sc;
         };
         ax_res_0.emplace_back(pw_over_cable(resistivity, cable, dflt_resistivity, scale_param));
@@ -419,15 +413,15 @@ fvm_cv_discretize(const cable_cell& cell, const cable_cell_parameter_set& global
         double cv_length = 0;
 
         for (mcable cable: cv_cables) {
-            auto scale_param = [&](const auto&, const auto& par) {
-                auto ie = thingify(par.value, provider);
-                auto sc = ie->eval(provider, cable);
+            auto scale_param = [&](const auto&, const auto& par) -> double {
+                auto ie = thingify(par.scale, provider);
+                auto sc = par.value*ie->eval(provider, cable);
                 return sc;
             };
 
-            auto pw_capacitance = pw_over_cable(capacitance, cable, dflt_capacitance, scale_param);
-            auto pw_potential   = pw_over_cable(potential,   cable, dflt_potential,   scale_param);
-            auto pw_temperature = pw_over_cable(temperature, cable, dflt_temperature, scale_param);
+            auto pw_capacitance = pw_over_cable(capacitance,   cable, dflt_capacitance, scale_param);
+            auto pw_potential   = pw_over_cable(potential,     cable, dflt_potential,   scale_param);
+            auto pw_temperature = pw_over_cable(temperature_K, cable, dflt_temperature, scale_param);
 
             D.cv_area[i]                 += embedding.integrate_area(cable);
             D.cv_capacitance[i]          += embedding.integrate_area(cable.branch, pw_capacitance);
@@ -544,7 +538,7 @@ bool cables_intersect_location(Seq&& cables, const mlocation& x) {
     auto eqr = std::equal_range(begin(cables), end(cables), x.branch, cmp_branch{});
 
     return util::any_of(util::make_range(eqr),
-        [&x](const mcable& c) { return c.prox_pos<=x.pos && x.pos<=c.dist_pos; });
+                        [&x](const mcable& c) { return c.prox_pos<=x.pos && x.pos<=c.dist_pos; });
 }
 
 voltage_reference_pair fvm_voltage_reference_points(const morphology& morph, const cv_geometry& geom, arb_size_type cell_idx, const mlocation& site) {
@@ -1312,9 +1306,9 @@ make_ion_config(fvm_ion_map build_data,
 
             for (const mcable& cable: data.D.geometry.cables(cv)) {
                 auto scale_param = [&](const auto&,
-                                   const auto& par) {
-                    auto ie = thingify(par.value, provider);
-                    auto sc = ie->eval(provider, cable);
+                                   const auto& par) -> double {
+                    auto ie = thingify(par.scale, provider);
+                    auto sc = par.value*ie->eval(provider, cable);
                     return sc;
                 };
 

diff --git a/arbor/include/arbor/cable_cell.hpp b/arbor/include/arbor/cable_cell.hpp
@@ -4,7 +4,6 @@
 #include <string>
 #include <unordered_map>
 #include <utility>
-#include <variant>
 #include <vector>
 
 #include <arbor/export.hpp>
@@ -239,7 +238,7 @@ using location_assignment =
 
 using cable_cell_region_map = static_typed_map<region_assignment,
     density, voltage_process, init_membrane_potential, axial_resistivity,
-    temperature_K, membrane_capacitance, init_int_concentration,
+    temperature, membrane_capacitance, init_int_concentration,
     ion_diffusivity, init_ext_concentration, init_reversal_potential>;
 
 using cable_cell_location_map = static_typed_map<location_assignment,