diff --git a/src/neat/actions/install.py b/src/neat/actions/install.py index a9d788f9..945746bf 100644 --- a/src/neat/actions/install.py +++ b/src/neat/actions/install.py @@ -447,8 +447,6 @@ def _compile_nest( ): from pynestml.frontend.pynestml_frontend import generate_nest_compartmental_target - print("!!! codegen_opts in _compile_nest:", codegen_opts) - # assert that `model_name` is a pure name assert not "/" in model_name assert not "." in model_name @@ -538,6 +536,21 @@ def _install_models( code uses a fast polynomial approximation only for dynamic propagator ``exp()`` terms in hot loops; all other exponentials use ``std::exp``/``std::expf``. + - ``single_precision_propagator_exp_mode``: ``"bounded"`` or + ``"plain"`` (default: ``"bounded"``). Only used when + ``fp_precision="single"`` and ``use_fastexp=False``. Selects + bounded or raw ``std::expf`` evaluation for propagator exponentials. + - ``with_profiling``: bool (default: ``False``). If ``True``, generated + models expose cumulative profiling recordables for matrix assembly, + Hines solves, and current evaluation. + - ``with_detailed_recordables``: bool (default: ``False``). If + ``True``, generated models expose additional multimeter recordables + for runtime propagators and pure helper functions such as ``*_inf_*`` + and ``tau_*`` values. + - ``freeze_exp_mode``: ``"none"`` or ``"freeze_init"`` (default: + ``"none"``). ``"freeze_init"`` replaces runtime ``exp`` evaluations + in ``f_numstep()`` with values computed once in ``pre_run_hook()`` + from the initialized model state and timestep. """ model_name = _resolve_model_name(model_name, channel_path_arg) diff --git a/src/neat/factorydefaults.py b/src/neat/factorydefaults.py index 793d04d2..fa93269c 100644 --- a/src/neat/factorydefaults.py +++ b/src/neat/factorydefaults.py @@ -23,7 +23,7 @@ import numpy.typing as npt from dataclasses import dataclass, field -from typing import Dict, List, Optional +from typing import Dict, List, Optional, Tuple @dataclass @@ -133,6 +133,13 @@ class FitParams: # ]) # ms # ) + # admittance-kernel correction: minimal degree of the rational fit + min_degree: int = 4 + # admittance-kernel correction: maximal degree of the rational fit + max_degree: int = 40 + # admittance-kernel correction: maximum relative error of the fit + max_rel_error: float = 1e-4 + @dataclass class MechParams: diff --git a/src/neat/modelreduction/compartmentfitter.py b/src/neat/modelreduction/compartmentfitter.py index eeb24615..09a5c891 100755 --- a/src/neat/modelreduction/compartmentfitter.py +++ b/src/neat/modelreduction/compartmentfitter.py @@ -28,7 +28,7 @@ from ..trees.stree import STree from ..trees.phystree import PhysTree from ..trees.compartmenttree import CompartmentTree -from ..tools.kernelextraction import Kernel +from ..tools.kernelextraction import Kernel, fExpFitter, FourierTools from ..channels.ionchannels import SPDict from ..factorydefaults import FitParams, MechParams from .cachetrees import CachedGreensTree, CachedSOVTree, EquilibriumTree @@ -595,6 +595,47 @@ def fit_concentration(self, fit_arg, ion): return ctree, locs + def _get_passified_greenstree(self, suffix="_passified_"): + """ + Construct a `CachedGreensTree` whose membrane has been linearized + ("passified") around the equilibrium potentials of the full tree. + + This helper is used by `fit_passive` (when + `use_all_channels=True`). The + resulting tree is cache-backed using `self.cache_path` / + `self.cache_name + "_eq" + suffix` and `... + "_gf" + suffix`, so + repeated calls reuse cached evaluations. + + Parameters + ---------- + suffix: str + Cache name suffix. Defaults to ``"_passified_"``, matching the + existing convention in `fit_passive`. + + Returns + ------- + `neat.CachedGreensTree` + """ + fit_tree = EquilibriumTree(self) + fit_tree.set_cache_params( + cache_path=self.cache_path, + cache_name=self.cache_name + "_eq" + suffix, + save_cache=self.save_cache, + recompute_cache=self.recompute_cache, + ) + # set the channels to passive + fit_tree.as_passive_membrane() + # convert to a greens tree for further evaluation + fit_tree = CachedGreensTree( + fit_tree, + cache_path=self.cache_path, + cache_name=self.cache_name + "_gf" + suffix, + save_cache=self.save_cache, + recompute_cache=self.recompute_cache, + ) + fit_tree.set_comp_tree(eps=self.fit_cfg.fit_comptree_eps) + return fit_tree + def fit_passive(self, fit_arg, use_all_channels=True, pprint=False): """ Fit the steady state passive model, consisting only of leak and coupling @@ -626,24 +667,7 @@ def fit_passive(self, fit_arg, use_all_channels=True, pprint=False): suffix = f"_passified_" if use_all_channels: - fit_tree = EquilibriumTree(self) - fit_tree.set_cache_params( - cache_path=self.cache_path, - cache_name=self.cache_name + "_eq" + suffix, - save_cache=self.save_cache, - recompute_cache=self.recompute_cache, - ) - # set the channels to passive - fit_tree.as_passive_membrane() - # convert to a greens tree for further evaluation - fit_tree = CachedGreensTree( - fit_tree, - cache_path=self.cache_path, - cache_name=self.cache_name + "_gf" + suffix, - save_cache=self.save_cache, - recompute_cache=self.recompute_cache, - ) - fit_tree.set_comp_tree(eps=self.fit_cfg.fit_comptree_eps) + fit_tree = self._get_passified_greenstree(suffix=suffix) else: fit_tree = self.create_tree_gf( [], # empty list of channel to include @@ -1210,12 +1234,214 @@ def fit_e_eq(self, fit_arg): return ctree, locs + def compute_admittance_correction(self, fit_arg, kernel_correction, pprint=False): + """ + Add admittance-kernel-correcting dummy compartments to a set of host + compartments in the reduced model. + + For each host compartment ``p`` specified by ``kernel_correction``, + the missing admittance + :math:`\Delta Y_p(s) = Y_p^{\\mathrm{full}}(s) - Y_p^{\\mathrm{red}}(s)` + is computed by comparing the driving-point admittance of the + passified full model against the leak-only driving-point admittance + of the reduced model on a log-spaced frequency grid spanning + ``self.fit_cfg.freq_band`` Hz. The residual is then fit by a sum of + zero-DC high-pass rational terms + + .. math:: + + \\Delta Y_p(s) \\approx \\sum_q \\alpha_q \\frac{s}{s + b_q} + + with ``alpha_q > 0`` and ``b_q > 0`` (see ``min_degree``, + ``max_degree`` and ``max_rel_error`` on ``FitParams``). Each + accepted term yields one passive dummy compartment with + ``g_c = alpha_q``, ``c = alpha_q / b_q``, ``g_l = 0``, attached + as a leaf to the host node. ``locs`` is extended with ``None`` + entries to keep its length in sync with the number of tree nodes. + + If for a given host the residual is negligible relative to + ``Y_p^full``, or if the fit is non-physical and degrades accuracy, + no dummy compartments are added. If the fit does not reach the + configured tolerance but improves accuracy over the uncorrected + reduction, it is kept and a warning is issued. + + Parameters + ---------- + fit_arg: see docstring of `CompartmentFitter.convert_fit_args()` + Specifying the fit that is being performed. + kernel_correction: list of int + Indices into the list of fit locations selecting the host + compartments to which an admittance correction is applied. + pprint: bool + + Returns + ------- + `neat.CompartmentTree` + The compartmenttree, with dummy compartments attached. + list of + The corresponding list of fit locations, extended with + ``None`` for each added dummy compartment. + """ + ctree, locs = self.convert_fit_arg(fit_arg) + + if kernel_correction is None or len(kernel_correction) == 0: + return ctree, locs + + # validate host indices + n_locs = len(locs) + for p in kernel_correction: + if not (0 <= p < n_locs): + raise IndexError( + f"`kernel_correction` index {p} is out of bounds for " + f"a fit with {n_locs} locations." + ) + + # frequency grid for admittance evaluation and fitting + ft = FourierTools( + np.array([0.0, 0.1]) + ) # dummy time array, we don't need it here + s_arr = ft.freqs_vfit + + # input impedances full model + gtree = self.create_tree_gf( + [], # empty list of channel to include + cache_name_suffix="_pas_", + ) + gtree.set_impedances_in_tree(freqs=s_arr, pprint=pprint) + zs_full = [gtree.calc_zf(loc, loc) for loc in locs] + + # to be corrected input impedances reduced model + z_mat = ctree.calc_impedance_matrix( + freqs=s_arr, channel_names=["L"], indexing="locs" + ) + zs_red = [z_mat[:, p, p] for p in kernel_correction] + + # negligibility threshold: an order of magnitude tighter than + # `max_rel_error`, so we only skip when the residual is truly tiny. + eps_skip = 0.1 * self.fit_cfg.max_rel_error + + # use the existing rational-fit engine + fef = fExpFitter() + + for ii, loc_idx in enumerate(kernel_correction): + host = ctree.get_nodes_from_loc_idxs(loc_idx) + y_full_p = 1.0 / zs_full[ii] + y_red_p = 1.0 / zs_red[ii] + dy_p = y_full_p - y_red_p + + # skip if the residual is negligible relative to the full admittance + if np.max(np.abs(dy_p) / np.abs(y_full_p)) < eps_skip: + continue + + # somatic capacitance correction + idx_bool = np.abs(s_arr) > 1e3 + (dca,), _, _, _ = np.linalg.lstsq( + s_arr[idx_bool][:, None].imag, dy_p[idx_bool].imag, rcond=None + ) + if pprint: + print(f"somatic capcitance correction dca = {dca} uF") + # apply capacitance correction to the host compartment and the residual + host.ca += dca + dy_p -= dca * s_arr + + # pl.figure("impedance full/red") + # ax1, ax2, ax3 = pl.subplot(311), pl.subplot(312), pl.subplot(313) + # ax1.plot(s_arr.imag, np.abs(zs_full[ii]), 'b', label="full") + # ax1.plot(s_arr.imag, np.abs(zs_red[ii]), 'r--', label="red") + # ax1.legend(loc=0) + # ax2.plot(s_arr.imag, (dy_p).real, 'g', label="full real") + # ax2.plot(s_arr.imag, (dy_p).imag, 'g--', label="full imag") + # ax2.plot(s_arr.imag, np.abs(dy_p), 'g:', label="full abs") + # ax2.plot(s_arr.imag, (dca * s_arr).real, 'y', label="fit real") + # ax2.plot(s_arr.imag, (dca * s_arr).imag, 'y--', label="fit imag") + # ax2.plot(s_arr.imag, np.abs(dca * s_arr), 'y:', label="fit abs") + + # ax3.plot(s_arr.imag, (dy_p - (dca * s_arr)).real, 'm', label="residual real") + # ax3.plot(s_arr.imag, (dy_p - (dca * s_arr)).imag, 'm--', label="residual imag") + # ax3.plot(s_arr.imag, np.abs(dy_p - (dca * s_arr)), 'm:', label="residual abs") + # pl.show() + + for Q in range(self.fit_cfg.min_degree, self.fit_cfg.max_degree + 1, 4): + alphas, gammas, _, rms = fef.fitFExp( + s_arr, + dy_p, + deg=Q, + rtol=1e-4, + realpoles=True, + initpoles="log10", + zerostart=False, + constrained=True, + reduce_numexp=False, + return_real=True, + ) + if pprint: + print( + f">>> admittance correction for host loc {loc_idx} and degree {Q}: rms = {rms}" + ) + + if rms < self.fit_cfg.max_rel_error: + print( + f"Error criterion satisfied for Q = {Q} | {rms} < {self.fit_cfg.max_rel_error}, stopping fit." + ) + print("alpha:\n", alphas, "\ngamma:\n", gammas) + break + + # f_corr = Kernel((alphas*1e-3, gammas*1e-3)) + # pl.figure(f"fit Q={Q}") + # pl.plot(s_arr.imag, dy_p.real, 'g', label="target real") + # pl.plot(s_arr.imag, dy_p.imag, 'g--', label="target real") + # pl.plot(s_arr.imag, np.abs(dy_p), 'g:', label="target real") + # pl.plot(s_arr.imag, f_corr.ft(s_arr).real + bias, 'y', label="fit real") + # pl.plot(s_arr.imag, f_corr.ft(s_arr).imag + bias, 'y--', label="fit imag") + # pl.plot(s_arr.imag, np.abs(f_corr.ft(s_arr)) + bias, 'y:', label="fit abs") + # pl.legend(loc=0) + # pl.show() + + if pprint: + print("Original compartment tree:\n", ctree) + + # attach dummy compartments to the host nodes according to fit results + for ar, gr in zip(alphas, gammas): + b_q = float(ar.real) + gc_q = float((-gr / ar).real) + ca_q = gc_q / b_q + + dummy_idx = max(n.index for n in ctree) + 1 + dummy = ctree.create_corresponding_node( + dummy_idx, + ca=ca_q, + g_c=gc_q, + g_l=0.0, + ) + # dummy has no associated MorphLoc + dummy.loc_idx = None + dummy.e_eq = host.e_eq + # leak reversal of the dummy follows the host leak + # reversal. The dummy has + # `g_l = 0`, so this is dynamically inert but keeps the + # convention explicit. + if "L" in host.currents: + dummy.currents["L"] = (0.0, host.currents["L"][1]) + else: + dummy.currents["L"] = (0.0, host.e_eq) + ctree.add_node_with_parent(dummy, host) + locs.append(None) + + if pprint: + print( + f"\nNew compartment tree after correcting compartment {host.index}:\n", + ctree, + ) + + return ctree, locs + def fit_model( self, loc_arg, fit_name="", alpha_inds=[0], use_all_channels_for_passive=True, + kernel_correction=None, pprint=False, ): """ @@ -1234,6 +1460,15 @@ def fit_model( Indices of all mode time-scales to be included in the fit use_all_channels_for_passive: bool (optional, default ``True``) Uses all channels in the tree to compute coupling conductances + kernel_correction: list of int or ``None`` (optional, default ``None``) + Indices into ``loc_arg`` (after bifurcation extension) selecting + host compartments to which an admittance-kernel correction will be + applied via additional passive dummy compartments. ``None`` or an + empty list disables the correction (default behavior). Note that this + correction can result in negative couplings / capacitances of dummy + compartments, which NEURON / Brian 2 exports do not support. + `NeuronCompartmentTree` and `Brian2CompartmentTree` will raise a warning + and ignore the correction. pprint: bool whether to print information @@ -1273,6 +1508,14 @@ def fit_model( # fit the resting potentials fit_arg = self.fit_e_eq(fit_arg) + # admittance-kernel correction by dummy compartments + if kernel_correction: + fit_arg = self.compute_admittance_correction( + fit_arg, + kernel_correction, + pprint=pprint, + ) + if fit_name == "temp": self.remove_fit(fit_name) else: diff --git a/src/neat/simulations/brian2/brian2model.py b/src/neat/simulations/brian2/brian2model.py index 2165dc4a..c25a702e 100644 --- a/src/neat/simulations/brian2/brian2model.py +++ b/src/neat/simulations/brian2/brian2model.py @@ -238,6 +238,24 @@ def __init__( self, arg=None, channel_storage: Optional[Dict[str, "IonChannel"]] = None ): super().__init__(arg) + + # Admittance-kernel correction dummy compartments (loc_idx is None) + # cannot be realized as Brian2 ``Cylinder`` sections: they have + # `g_l = 0` and the fake-geometry solve can produce non-physical + # (negative) radii / lengths for them. The correction is only + # meaningful for the NEST / analytic models, so we drop the dummy + # compartments here (``self`` is a copy of the input tree, so the + # caller's tree is left untouched). + if self.has_correction_compartments(): + warnings.warn( + "The compartment tree contains admittance-kernel correction " + "dummy compartments (loc_idx is None), which cannot be " + "represented in Brian2. They are ignored when building the " + "`Brian2CompartmentTree`.", + UserWarning, + ) + self.remove_correction_compartments() + self._channel_storage = channel_storage or getattr(self, "channel_storage", {}) # auxiliary constants for fake geometry calculation diff --git a/src/neat/simulations/neuron/neuronmodel.py b/src/neat/simulations/neuron/neuronmodel.py index 34e2e078..d6d5f815 100755 --- a/src/neat/simulations/neuron/neuronmodel.py +++ b/src/neat/simulations/neuron/neuronmodel.py @@ -290,7 +290,7 @@ def _make_section(self, factorlambda=1.0, pprint=False): compartment.diam = ( 2.0 * self.R ) # section radius [um] (NEURON takes diam = 2*r) - compartment.L = self.L # section length [um] + compartment.L = max(1e-9, self.L) # section length [um] # set number of segments if type(factorlambda) == float: # nseg according to NEURON book @@ -1535,6 +1535,22 @@ def __init__(self, ctree, fake_c_m=1.0, fake_r_a=100.0 * 1e-6, method=2): "`neat.NeuronCompartmentTree` can only be instantiated " "from a `neat.CompartmentTree` or derived class" ) + # Admittance-kernel correction dummy compartments (loc_idx is None) + # cannot be realized as NEURON sections: they have `g_l = 0` and the + # fake-geometry solve can produce non-physical (negative) radii / + # lengths for them. The correction is only meaningful for the + # NEST / analytic models, so we drop the dummy compartments here + # (working on a copy so the caller's tree is left untouched). + if ctree.has_correction_compartments(): + warnings.warn( + "The compartment tree contains admittance-kernel correction " + "dummy compartments (loc_idx is None), which cannot be " + "represented in NEURON. They are ignored when building the " + "`NeuronCompartmentTree`.", + UserWarning, + ) + ctree = ctree.__copy__() + ctree.remove_correction_compartments() super().__init__(ctree, types=[1, 3, 4]) self.equivalent_locs = ctree.get_equivalent_locs() self._create_reduced_neuron_model( diff --git a/src/neat/tools/kernelextraction.py b/src/neat/tools/kernelextraction.py index d746e6ee..5046e795 100755 --- a/src/neat/tools/kernelextraction.py +++ b/src/neat/tools/kernelextraction.py @@ -1959,6 +1959,7 @@ def inverse_fourier( constrained=True, reduce_numexp=False, ) + zk = Kernel({"a": alpha * 1e-3, "c": gamma * 1e-3}) if compute_time_derivative: dzk_dt = zk.diff() diff --git a/src/neat/trees/compartmenttree.py b/src/neat/trees/compartmenttree.py index c5b70fa7..d49e0bf4 100755 --- a/src/neat/trees/compartmenttree.py +++ b/src/neat/trees/compartmenttree.py @@ -87,13 +87,12 @@ def set_loc_idx(self, loc_idx): self._loc_idx = loc_idx def get_loc_idx(self): - if self._loc_idx is None: - raise AttributeError( - "`self.loc_idx` is undefined, this node has " - + "not been associated with a location" - ) - else: - return self._loc_idx + # `None` is a sentinel used by the admittance-kernel correction + # to mark dummy compartments that have no associated MorphLoc. + # Returning `None` here (instead of raising) lets the standard tree + # traversal helpers filter dummy nodes without needing exception + # handling, while still being explicit about "no location". + return self._loc_idx loc_idx = property(get_loc_idx, set_loc_idx) @@ -692,7 +691,10 @@ def set_e_eq(self, e_eq, indexing="locs"): ---------- e_eq: float or np.array of floats The equilibrium potential(s). If a float, the same potential is set - at every node. If a numpy array, must have the same length as `self` + at every node. If a numpy array with `indexing='tree'`, must have the + same length as `self`. If a numpy array with `indexing='locs'`, must + have the same length as the number of location-bearing nodes (dummy + compartments added by `compute_admittance_correction` are skipped). indexing: 'locs' or 'tree' The ordering of the equilibrium potentials. If 'locs', assumes the equilibrium potentials are in the order of the list of locations @@ -700,12 +702,17 @@ def set_e_eq(self, e_eq, indexing="locs"): of which nodes appear during iteration """ if isinstance(e_eq, float) or isinstance(e_eq, int): - e_eq = e_eq * np.ones(len(self), dtype=float) - elif indexing == "locs": - e_eq = self._permute_to_tree(np.array(e_eq)) + for node in self: + node.e_eq = float(e_eq) + return - for ii, node in enumerate(self): - node.e_eq = e_eq[ii] + if indexing == "locs": + e_eq_arr = self._permute_to_tree(np.array(e_eq)) + for node, val in zip(self._iter_loc_nodes(), e_eq_arr): + node.e_eq = val + else: + for ii, node in enumerate(self): + node.e_eq = e_eq[ii] def get_e_eq(self, indexing="locs"): """ @@ -736,15 +743,22 @@ def set_conc_eq(self, ion, conc_eq, indexing="locs"): Parameters ---------- conc_eq: `np.array` or float - The equilibrium concentrations [mM] + The equilibrium concentrations [mM]. If a numpy array with + `indexing='locs'`, dummy compartments (those with `loc_idx is None`) + are skipped. """ if isinstance(conc_eq, float) or isinstance(conc_eq, int): - conc_eq = conc_eq * np.ones(len(self), dtype=float) - elif indexing == "locs": - conc_eq = self._permute_to_tree(np.array(conc_eq)) + for node in self: + node.set_conc_eq(ion, float(conc_eq)) + return - for ii, node in enumerate(self): - node.set_conc_eq(ion, conc_eq[ii]) + if indexing == "locs": + conc_eq_arr = self._permute_to_tree(np.array(conc_eq)) + for node, val in zip(self._iter_loc_nodes(), conc_eq_arr): + node.set_conc_eq(ion, val) + else: + for ii, node in enumerate(self): + node.set_conc_eq(ion, conc_eq[ii]) def get_conc_eq(self, ion, indexing="locs"): """ @@ -879,16 +893,67 @@ def add_conc_mech(self, ion, params={}): for node in self: node.add_conc_mech(ion, params=params) + def _iter_loc_nodes(self): + """ + Iterate over tree nodes that correspond to an actual fit location, + i.e. nodes whose `loc_idx is not None`. Admittance-kernel correction + dummy compartments have `loc_idx is None` and are skipped. + """ + for node in self: + if node.loc_idx is not None: + yield node + + def has_correction_compartments(self): + """ + Whether the tree contains admittance-kernel correction dummy + compartments. These are leaf nodes attached by + `CompartmentFitter.compute_admittance_correction` and are marked + by ``loc_idx is None`` (they have no associated location). + + Returns + ------- + bool + """ + return any(node.loc_idx is None for node in self) + + def remove_correction_compartments(self): + """ + Remove all admittance-kernel correction dummy compartments + (those with ``loc_idx is None``) from the tree, in place. + + Dummy compartments are always attached as leaves to a host + compartment (see + `CompartmentFitter.compute_admittance_correction`), so removing + them leaves the structure of the location-bearing compartments + intact. Because the dummies are always assigned the highest node + indices, the remaining (location-bearing) compartments keep a + consecutive ``0 .. N-1`` indexing. + + Returns + ------- + int + The number of dummy compartments that were removed. + """ + dummy_nodes = [node for node in self if node.loc_idx is None] + for node in dummy_nodes: + # dummies are leaves, so this removes only the dummy itself + self.remove_node(node) + return len(dummy_nodes) + def permute_to_tree_idxs(self): """ - Returns index array for the permutation of location indices to tree indices + Returns index array for the permutation of location indices to tree + indices. Skips dummy nodes (those with `loc_idx is None`); the returned + array has length equal to the number of location-bearing nodes. """ - return np.array([node.loc_idx for node in self]) + return np.array([node.loc_idx for node in self._iter_loc_nodes()], dtype=int) def _permute_to_tree(self, mat): """ - Permutes the input array, which is assumed to be ordered according to the - location list, to the tree order + Permutes the input array, which is assumed to be ordered according to + the location list, to the tree order (location-bearing nodes only). + Dummy nodes (`loc_idx is None`) are skipped: the returned array has + length equal to the number of location-bearing nodes. """ index_arr = self.permute_to_tree_idxs() if mat.ndim == 1: @@ -898,14 +963,36 @@ def _permute_to_tree(self, mat): def permute_to_locs_idxs(self): """ - Return an index array that can be used to permute matrices that follow to tree - order to the location list order + Return an index array that can be used to permute tree-iteration-order + arrays (restricted to location-bearing nodes) to location list order. + Dummy nodes (`loc_idx is None`) are skipped. """ - loc_idxs = np.array([node.loc_idx for node in self]) + loc_idxs = np.array( + [node.loc_idx for node in self._iter_loc_nodes()], dtype=int + ) return np.argsort(loc_idxs) + def _loc_node_positions(self): + """ + Return an array of positions (in tree iteration order) of nodes whose + `loc_idx is not None`. Used to extract the location-bearing rows/cols + from a full tree-ordered array. + """ + return np.array( + [ii for ii, node in enumerate(self) if node.loc_idx is not None], + dtype=int, + ) + def _permuteToLocs(self, mat): - index_arr = self.permute_to_locs_idxs() + """ + Takes a tree-iteration-ordered array (potentially including entries for + admittance-correction dummy nodes), extracts the location-bearing + entries, and reorders them to location list order. + """ + sort_idxs = self.permute_to_locs_idxs() + # positions of location-bearing nodes in tree iteration order + loc_positions = self._loc_node_positions() + index_arr = loc_positions[sort_idxs] if mat.ndim == 1: return mat[index_arr] else: @@ -914,16 +1001,18 @@ def _permuteToLocs(self, mat): def get_equivalent_locs(self): """ Get list of fake locations in the same order as original list of locations - to which the compartment tree was fitted. + to which the compartment tree was fitted. Admittance-correction dummy + nodes (those with `loc_idx is None`) are skipped. Returns ------- list of tuple Tuple has the form `(node.index, .5)` """ - loc_idxs = [node.loc_idx for node in self] + real_nodes = list(self._iter_loc_nodes()) + loc_idxs = [node.loc_idx for node in real_nodes] index_arr = np.argsort(loc_idxs) - locs_unordered = [(node.index, 0.5) for node in self] + locs_unordered = [(node.index, 0.5) for node in real_nodes] return [locs_unordered[ind] for ind in index_arr] def calc_impedance_matrix( @@ -955,14 +1044,28 @@ def calc_impedance_matrix( frequency, the second and third dimension contain the impedance matrix for that frequency """ - return np.linalg.inv( - self.calc_system_matrix( - freqs=freqs, - channel_names=channel_names, - indexing=indexing, - use_conc=use_conc, - ) + # Build the full tree-order system matrix and invert it. When dummy + # compartments (loc_idx is None) are present in the tree, inverting + # the loc-only submatrix of S would give the impedance of a model + # without those dummies (the Schur complement that eliminates them + # would not be applied). Inverting the full tree-order S first and + # then extracting the loc-bearing rows/cols of Z yields the correct + # driving-point and transfer impedances among the location-bearing + # compartments with the dummy compartments still attached as load. + s_mat = self.calc_system_matrix( + freqs=freqs, + channel_names=channel_names, + indexing="tree", + use_conc=use_conc, ) + z_mat = np.linalg.inv(s_mat) + if indexing == "locs": + z_mat = self._permuteToLocs(z_mat) + elif indexing != "tree": + raise ValueError( + "invalid argument for `indexing`, has to be 'tree' or 'locs'" + ) + return z_mat def calc_impulse_response_matrix( self, @@ -1013,7 +1116,7 @@ def calc_impulse_response_matrix( zf_mat = self.calc_impedance_matrix( freqs=ft.freqs, channel_names=channel_names, - indexing=indexing, + indexing="tree", use_conc=use_conc, ) @@ -1045,6 +1148,15 @@ def calc_impulse_response_matrix( ) zt_mat[:, jj, ii] = zt_mat[:, ii, jj] + if indexing == "locs": + zt_mat = self._permuteToLocs(zt_mat) + if compute_time_derivative: + dzt_dt_mat = self._permuteToLocs(dzt_dt_mat) + elif indexing != "tree": + raise ValueError( + "invalid argument for `indexing`, has to be 'tree' or 'locs'" + ) + if compute_time_derivative: return zt_mat, dzt_dt_mat else: diff --git a/tests/test_brian2tree.py b/tests/test_brian2tree.py index 14e85451..1ff0ecee 100644 --- a/tests/test_brian2tree.py +++ b/tests/test_brian2tree.py @@ -322,7 +322,7 @@ def test_axon_brian2_neuron_comparison(self, pplot=False): ax.plot(res_brian2["t"][:idx], res_brian2["v_m"][ii][:idx], "bo--") pl.show() - def load_T_tree(self): + def load_T_tree(self, fit_locs=None, kernel_correction=None): tree = PhysTree( os.path.join(MORPHOLOGIES_PATH_PREFIX, "Ttree_segments.swc"), types=[1, 2, 3, 4], @@ -337,10 +337,54 @@ def load_T_tree(self): ca_chan, 0.005 * 1e6, 132.4579341637009, node_arg=[tree[1]] ) tree.fit_leak_current(-70.0, 15.0) + self.tree = tree - locs = [(n.index, 0.5) for n in tree] + if fit_locs is None: + fit_locs = [(n.index, 0.5) for n in tree] cfit = CompartmentFitter(tree, save_cache=False, recompute_cache=True) - self.ctree, _ = cfit.fit_model(locs) + self.ctree, _ = cfit.fit_model(fit_locs, kernel_correction=kernel_correction) + + def test_ignore_correction_compartments(self): + """ + A somatic fit with an admittance-kernel correction attaches passive + dummy compartments (``loc_idx is None``) to the soma. These cannot be + represented in Brian2, so `Brian2CompartmentTree` must drop them (with + a warning) while leaving the location-bearing compartments intact. + """ + self.load_T_tree(fit_locs=[(1, 0.5)], kernel_correction=[0]) + ctree = self.ctree + + # the correction should have induced dummy compartments + dummy_idxs = [n.index for n in ctree if n.loc_idx is None] + real_idxs = [n.index for n in ctree if n.loc_idx is not None] + assert len(dummy_idxs) > 0, ( + "the somatic admittance correction did not induce any dummy " + "compartments; the test cannot verify they are ignored" + ) + assert ctree.has_correction_compartments() + + # building the Brian2 model should warn about (and ignore) the dummies + with pytest.warns(UserWarning): + bct = Brian2CompartmentTree(ctree, channel_storage=ctree.channel_storage) + + # the dummy compartments are dropped from the Brian2 tree ... + assert not bct.has_correction_compartments() + for di in dummy_idxs: + assert bct[di] is None + # ... while the location-bearing compartments remain + for ri in real_idxs: + assert bct[ri] is not None + assert len([n for n in bct]) == len(real_idxs) + + # the caller's tree must be left untouched (dummies still present) + assert ctree.has_correction_compartments() + assert len([n for n in ctree if n.loc_idx is None]) == len(dummy_idxs) + + # the pruned model must be buildable (no non-physical geometry); the + # Brian2 model has one compartment per location-bearing node + brian2.start_scope() + bmodel = bct.init_model() + assert len(bmodel.v) == len(real_idxs) def test_dend_brian2_neuron_comparison(self, pplot=False): dt = 0.01 diff --git a/tests/test_compartmentfitter.py b/tests/test_compartmentfitter.py index ab52d28c..d15c82f3 100755 --- a/tests/test_compartmentfitter.py +++ b/tests/test_compartmentfitter.py @@ -29,11 +29,15 @@ MorphTree, PhysTree, GreensTree, + GreensTreeTime, SOVTree, NeuronCompartmentTree, + CompartmentFitter, + CompartmentTree, + CachedGreensTree, check_for_coreneuron, + FourierTools, ) -from neat import CompartmentFitter, CachedGreensTree import neat.modelreduction.compartmentfitter as compartmentfitter import channelcollection_for_tests as channelcollection @@ -832,7 +836,257 @@ def test_e_eq_fit(self): assert v_eq_sim == pytest.approx(v_eq_target) -@SKIP_ACTIVE_MECHS_WITH_CORENEURON +class TestAdmittanceCorrection: + """ + Tests for the admittance-kernel correction feature of `CompartmentFitter` + (see `CompartmentFitter.compute_admittance_correction`). + """ + + # --- fixture loaders (mirrored from TestCompartmentFitter) ----------- + + def load_T_tree(self): + fname = os.path.join(MORPHOLOGIES_PATH_PREFIX, "Tsovtree.swc") + self.tree = PhysTree(fname, types=[1, 3, 4]) + self.tree.set_physiology(0.8, 100.0 / 1e6) + self.tree.fit_leak_current(-75.0, 10.0) + self.tree.set_comp_tree() + + def load_ball_and_stick(self): + self.tree = PhysTree( + os.path.join(MORPHOLOGIES_PATH_PREFIX, "ball_and_stick.swc") + ) + self.tree.set_physiology(0.8, 100.0 / 1e6) + self.tree.set_leak_current(100.0, -75.0) + self.tree.set_comp_tree() + + def load_ball(self): + self.tree = PhysTree(os.path.join(MORPHOLOGIES_PATH_PREFIX, "ball.swc")) + self.tree.set_physiology(0.8, 100.0 / 1e6) + k_chan = channelcollection.Kv3_1() + self.tree.add_channel_current(k_chan, 0.766 * 1e6, -85.0) + na_chan = channelcollection.Na_Ta() + self.tree.add_channel_current(na_chan, 1.71 * 1e6, 50.0) + self.tree.fit_leak_current(-75.0, 10.0) + self.tree.set_v_ep(-75.0) + self.tree.set_comp_tree() + + def test_point_neuron_no_correction(self): + """ + For a point neuron the reduction is exact: `Delta Y(s) == 0`, + so `compute_admittance_correction` should attach no dummy + compartments. + """ + self.load_ball() + cm = CompartmentFitter(self.tree, save_cache=False, recompute_cache=True) + ctree, locs = cm.fit_model( + [(1, 0.5)], + kernel_correction=[0], + pprint=False, + ) + # no dummies attached + assert len(ctree) == 1 + # no None entries appended to locs + assert all(loc is not None for loc in locs) + + def test_equilibrium(self): + """ + With no input current, a corrected reduced model with zero-leak + dummies should remain at equilibrium: the dummy compartments add + no DC load, and their leak reversal equals the host's. + """ + self.load_ball_and_stick() + cm = CompartmentFitter( + self.tree, + save_cache=False, + recompute_cache=True, + ) + fit_locs = [(1, 0.5), (4, 0.5)] + ctree, locs = cm.fit_model( + fit_locs, + kernel_correction=[0], + pprint=False, + ) + # every dummy compartment has g_l = 0 and e_eq matching its host + n_dummies = 0 + for node in ctree: + if node.loc_idx is None: + n_dummies += 1 + # zero leak conductance + assert node.currents["L"][0] == pytest.approx(0.0) + # leak reversal equals e_eq (which equals host e_eq) + parent = node.parent_node + assert node.e_eq == pytest.approx(parent.e_eq) + # dummy zero-DC admittance: Y_dum(0) == 0 in isolation + # check via the steady-state system matrix on a tree + # containing just this dummy + host + # short-circuit: if no dummies were attached, the test has nothing + # to verify beyond the assertion that no extra nodes were added + if n_dummies == 0: + return + + # DC admittance of the full reduced+dummy tree at the host loc + # must equal that of the leak-only reduced tree (no dummies): + # Y_dum(0) == 0 by construction. + # Build a copy of ctree without dummies for comparison via the + # tree-indexed system matrix at freq=0. + z_hyb = ctree.calc_impedance_matrix( + freqs=0.0, channel_names=["L"], indexing="locs" + ) + # compare against the same model rebuilt without kernel correction + cm2 = CompartmentFitter( + self.tree, + save_cache=False, + recompute_cache=True, + ) + ctree_ref, _ = cm2.fit_model(fit_locs, pprint=False) + z_red = ctree_ref.calc_impedance_matrix( + freqs=0.0, channel_names=["L"], indexing="locs" + ) + # entries at fit-location indices should match at DC + assert np.allclose(z_hyb, z_red, rtol=5e-3, atol=1e-6) + + def test_discrete_update_T_tree(self, pplot=False): + """ + With kernel correction, the somatic driving-point admittance of + the reduced model should be closer to the full model than without + correction, on the configured frequency band. + """ + self.load_T_tree() + fit_locs = [(1, 0.5)] # , (4, 1.0), (5, 0.5), (8, 0.5)] + + cm = CompartmentFitter( + self.tree, + save_cache=False, + recompute_cache=True, + ) + ctree_red, _ = cm.fit_model(fit_locs, pprint=False) + + cm2 = CompartmentFitter( + self.tree, + save_cache=False, + recompute_cache=True, + ) + ctree_hyb, locs_hyb = cm2.fit_model( + fit_locs, + kernel_correction=[0], + pprint=False, + ) + + t_arr = np.linspace(0.0, 50.0, 1000) + gtt = GreensTreeTime(self.tree) + gtt.set_impedance(t_arr) + zt0 = gtt.calc_zt((1, 0.5), (1, 0.5), compute_time_derivative=False) + + z_red = ctree_red.calc_impulse_response_matrix(t_arr) + z_hyb = ctree_hyb.calc_impulse_response_matrix(t_arr) + + errt_red = np.linalg.norm(zt0 - z_red[:, 0, 0]) + errt_hyb = np.linalg.norm(zt0 - z_hyb[:, 0, 0]) + + # test whetehrcorrection significantlys reduce time-domain error + assert errt_hyb < errt_red - np.abs(errt_red - 1e1) + + # from neat import NestCompartmentTree + # # t1 = NeuronCompartmentTree(ctree_hyb) + # # t1.init_model() + + # channel_installer.load_or_install_nest_test_channels() + # t2 = NestCompartmentTree(ctree_hyb) + # t2.init_model("multichannel_test",1) + + if pplot: + import matplotlib.pyplot as plt + + plt.plot(t_arr, zt0, "g", label="full") + plt.plot(t_arr, z_red[:, 0, 0], "b", label="default") + plt.plot(t_arr, z_hyb[:, 0, 0], "r--", label="corrected") + plt.show() + + # admittance comparison on the same grid the correction uses + ft = FourierTools(t_arr) + s_arr = ft.freqs_vfit + + gtree = cm._get_passified_greenstree(suffix="_test_passified_") + gtree.set_impedances_in_tree(freqs=s_arr, pprint=False) + z_full = gtree.calc_impedance_matrix(fit_locs) + y_full_0 = 1.0 / z_full[:, 0, 0] + + z_red = ctree_red.calc_impedance_matrix( + freqs=s_arr, + channel_names=["L"], + indexing="locs", + ) + y_red_0 = 1.0 / z_red[:, 0, 0] + + z_hyb = ctree_hyb.calc_impedance_matrix( + freqs=s_arr, + channel_names=["L"], + indexing="locs", + ) + # extract the host's diagonal in the location-indexed matrix + y_hyb_0 = 1.0 / z_hyb[:, 0, 0] + + err_red = np.linalg.norm(y_full_0 - y_red_0) + err_hyb = np.linalg.norm(y_full_0 - y_hyb_0) + + # test whether correction significantly reduces error on the frequency grid + assert err_hyb < err_red - np.abs(err_red - 5e-3) + + def test_discrete_update_ball_and_stick(self): + """ + Same as `test_discrete_update_T_tree`, but on the ball-and-stick + model. + """ + self.load_ball_and_stick() + fit_locs = [(1, 0.5)] + + cm = CompartmentFitter( + self.tree, + save_cache=False, + recompute_cache=True, + ) + ctree_red, _ = cm.fit_model(fit_locs, pprint=False) + + cm2 = CompartmentFitter( + self.tree, + save_cache=False, + recompute_cache=True, + ) + ctree_hyb, locs_hyb = cm2.fit_model( + fit_locs, + kernel_correction=[0], + pprint=False, + ) + + # admittance comparison on the same grid the correction uses + ft = FourierTools(np.array([0.0, 50.0])) + s_arr = ft.freqs_vfit + + gtree = cm._get_passified_greenstree(suffix="_test_passified_") + gtree.set_impedances_in_tree(freqs=s_arr, pprint=False) + z_full = gtree.calc_impedance_matrix(fit_locs) + y_full_0 = 1.0 / z_full[:, 0, 0] + + z_red = ctree_red.calc_impedance_matrix( + freqs=s_arr, + channel_names=["L"], + indexing="locs", + ) + y_red_0 = 1.0 / z_red[:, 0, 0] + + z_hyb = ctree_hyb.calc_impedance_matrix( + freqs=s_arr, + channel_names=["L"], + indexing="locs", + ) + y_hyb_0 = 1.0 / z_hyb[:, 0, 0] + + err_red = np.linalg.norm(y_full_0 - y_red_0) + err_hyb = np.linalg.norm(y_full_0 - y_hyb_0) + + assert err_hyb < err_red - np.abs(err_red - 5e-3) + + def test_expansion_points(): kv3_1 = channelcollection.Kv3_1() na_ta = channelcollection.Na_Ta() @@ -872,4 +1126,10 @@ def test_expansion_points(): tcf.test_fit_storage() tcf.test_e_eq_fit() + tac = TestAdmittanceCorrection() + tac.test_point_neuron_no_correction() + tac.test_equilibrium() + tac.test_discrete_update_T_tree(pplot=True) + tac.test_discrete_update_ball_and_stick() + test_expansion_points() diff --git a/tests/test_nesttree.py b/tests/test_nesttree.py index cff9a767..a00fc963 100644 --- a/tests/test_nesttree.py +++ b/tests/test_nesttree.py @@ -34,7 +34,7 @@ from neat import PhysTree from neat import CompartmentNode, CompartmentTree -from neat import CompartmentFitter, NeuronCompartmentTree +from neat import CompartmentFitter, NeuronCompartmentTree, NeuronSimTree from neat import NestCompartmentNode, NestCompartmentTree, load_nest_model import channelcollection_for_tests as channelcollection @@ -308,9 +308,12 @@ def test_axon_nest_neuron_comparison(self, pplot=False): ax.plot(res_nest["times"], res_nest["v_comp2"], "bo--") pl.show() - def load_T_tree(self): + def _build_T_phystree(self): """ - Parameters taken from a BBP SST model for a subset of ion channels + Build the T-tree `neat.PhysTree` morphology with a subset of ion + channels at the soma (parameters taken from a BBP SST model). A fresh + tree is returned on every call so that independent fits do not share + (and mutate) the same underlying morphology. """ tree = PhysTree( os.path.join(MORPHOLOGIES_PATH_PREFIX, "Ttree_segments.swc"), @@ -332,10 +335,20 @@ def load_T_tree(self): # passive leak current tree.fit_leak_current(-70.0, 15.0) + return tree + + def load_T_tree(self, fit_locs=None, kernel_correction=None): + """ + Parameters taken from a BBP SST model for a subset of ion channels + """ + tree = self._build_T_phystree() + self.tree = tree + # simplify - locs = [(n.index, 0.5) for n in tree] + if fit_locs is None: + fit_locs = [(n.index, 0.5) for n in tree] cfit = CompartmentFitter(tree, save_cache=False, recompute_cache=True) - self.ctree, _ = cfit.fit_model(locs) + self.ctree, _ = cfit.fit_model(fit_locs, kernel_correction=kernel_correction) def test_dend_nest_neuron_comparison(self, pplot=False): dt = 0.01 @@ -493,11 +506,150 @@ def test_dend_nest_neuron_comparison(self, pplot=False): ax.plot(res_nest["times"], res_nest[f"h_NaTa_t{dend_idx}"], "g--", lw=2) pl.show() + def _simulate_nest_dc_step(self, ctree, amp, delay, dur, cal, dt, tmax): + """ + Inject a somatic DC current step into a NEST reduction of `ctree` + and return the somatic voltage trace (time aligned so that ``t = 0`` + is the end of the calibration window). + """ + idx0 = int(cal / dt) + + nest.ResetKernel() + channel_installer.load_or_install_nest_test_channels() + nest.SetKernelStatus(dict(resolution=dt)) + + csimtree_nest = NestCompartmentTree(ctree) + nestmodel = csimtree_nest.init_model("multichannel_test", 1) + # somatic current input port + nestmodel.receptors = [ + {"comp_idx": 0, "receptor_type": "curr_in"}, + ] + # somatic DC current step + dcg = nest.Create( + "step_current_generator", + { + "amplitude_times": [cal + dt, cal + delay, cal + delay + dur], + "amplitude_values": [0.0, amp, 0.0], + }, + ) + nest.Connect( + dcg, + nestmodel, + syn_spec={ + "synapse_model": "static_synapse", + "weight": 1.0, + "delay": dt, + "receptor_type": 0, + }, + ) + # somatic voltage recording + mm = nest.Create("multimeter", 1, {"record_from": ["v_comp0"], "interval": dt}) + nest.Connect(mm, nestmodel) + + nest.Simulate(cal + tmax) + res_nest = nest.GetStatus(mm, "events")[0] + + times = res_nest["times"][idx0:] - res_nest["times"][idx0] + v_comp0 = res_nest["v_comp0"][idx0:] + + return times, v_comp0 + + def test_admittance_correction_dc_step(self, pplot=False): + """ + Compare the somatic voltage response to a DC current step between + + (i) the full T morphology simulated in NEURON, + (ii) a NEST reduction with only the somatic compartment, and + (iii) a NEST reduction with only the somatic compartment plus an + admittance-kernel correction. + + The admittance correction adds passive dummy compartments that + reproduce the (frequency-dependent) load of the missing dendrites, + so the corrected reduction (iii) should track the full model (i) + more closely than the uncorrected reduction (ii). + """ + dt = 0.025 + cal = 200.0 # calibration / settling time [ms] + delay = 50.0 # step onset after calibration [ms] + dur = 20.0 # step duration [ms] + tmax = delay + dur + 100.0 # total recorded time [ms] + amp = 0.5 # hyperpolarizing step [nA], kept subthreshold + + # --- (i) full morphology in NEURON ------------------------------ + tree_full = self._build_T_phystree() + tree_full.set_comp_tree() + sim_full = NeuronSimTree(tree_full) + sim_full.set_default_tree("computational") + sim_full.init_model(dt=dt, t_calibrate=cal) + sim_full.store_locs([(1, 0.5)], name="rec locs") + sim_full.add_i_clamp((1, 0.5), amp, delay, dur) + res_full = sim_full.run(tmax) + sim_full.delete_model() + t_full, v_full = res_full["t"], res_full["v_m"][0] + + # --- (ii) uncorrected somatic NEST reduction -------------------- + cfit_red = CompartmentFitter( + self._build_T_phystree(), save_cache=False, recompute_cache=True + ) + ctree_red, _ = cfit_red.fit_model([(1, 0.5)]) + assert not ctree_red.has_correction_compartments() + t_red, v_red = self._simulate_nest_dc_step( + ctree_red, amp, delay, dur, cal, dt, tmax + ) + + # --- (iii) corrected somatic NEST reduction --------------------- + cfit_hyb = CompartmentFitter( + self._build_T_phystree(), save_cache=False, recompute_cache=True + ) + ctree_hyb, _ = cfit_hyb.fit_model([(1, 0.5)], kernel_correction=[0]) + # the correction must have induced dummy compartments, otherwise the + # corrected and uncorrected reductions are identical + assert ctree_hyb.has_correction_compartments() + t_hyb, v_hyb = self._simulate_nest_dc_step( + ctree_hyb, amp, delay, dur, cal, dt, tmax + ) + + # --- compare the step responses --------------------------------- + # work on a common length and compare the deflection from the + # pre-step baseline, which removes any small DC offset between the + # models and isolates the (transient) response shape that the + # correction acts on. + imax = min(len(v_full), len(v_red), len(v_hyb)) + i_base = int(delay / dt) - 1 + + def deflection(v): + return v[:imax] - v[i_base] + + d_full = deflection(v_full) + d_red = deflection(v_red) + d_hyb = deflection(v_hyb) + + rmse_red = np.sqrt(np.mean((d_full - d_red) ** 2)) + rmse_hyb = np.sqrt(np.mean((d_full - d_hyb) ** 2)) + + print(f"RMSE uncorrected reduction: {rmse_red:.4f} mV") + print(f"RMSE corrected reduction: {rmse_hyb:.4f} mV") + + if pplot: + pl.figure("admittance correction DC step") + pl.plot(t_full[:imax], d_full, "g-", label="full (NEURON)") + pl.plot(t_red[:imax], d_red, "b--", label="reduction (NEST)") + pl.plot(t_hyb[:imax], d_hyb, "r-.", label="reduction + correction (NEST)") + pl.xlabel("t [ms]") + pl.ylabel(r"$\Delta v$ [mV]") + pl.legend(loc=0) + pl.show() + + # the corrected reduction must track the full model much better + # as the uncorrected one + assert rmse_hyb < 0.1 * rmse_red + if __name__ == "__main__": tn = TestNest() # tn.test_model_construction() # tn.test_initialization() - tn.test_single_comp_nest_neuron_comparison(pplot=True) + # tn.test_single_comp_nest_neuron_comparison(pplot=True) # tn.test_axon_nest_neuron_comparison(pplot=True) # tn.test_dend_nest_neuron_comparison(pplot=True) + tn.test_admittance_correction_dc_step(pplot=True) diff --git a/tests/test_neurontree.py b/tests/test_neurontree.py index c2c4cd4f..54706572 100755 --- a/tests/test_neurontree.py +++ b/tests/test_neurontree.py @@ -29,9 +29,10 @@ import pytest import itertools -from neat import GreensTree +from neat import GreensTree, PhysTree from neat import CompartmentNode, CompartmentTree from neat import NeuronSimTree, NeuronCompartmentTree +from neat import CompartmentFitter from neat import check_for_coreneuron import neat.tools.kernelextraction as ke @@ -442,6 +443,36 @@ def load_multi_dend_model(self, w_locinds=True): if w_locinds: self.add_locinds() + def load_T_tree(self, kernel_correction=None): + """ + Load the T-tree morphology and fit a reduced model at the soma. + + When ``kernel_correction=[0]`` is passed, an admittance-kernel + correction is applied to the (single) somatic compartment, which + attaches passive dummy compartments (``loc_idx is None``) to the + soma. These dummy compartments cannot be represented in NEURON and + should be ignored by `NeuronCompartmentTree`. + """ + tree = PhysTree( + os.path.join(MORPHOLOGIES_PATH_PREFIX, "Ttree_segments.swc"), + types=[1, 2, 3, 4], + ) + tree.set_physiology(1.0, 100.0 / 1e6) + k_chan = channelcollection.SKv3_1() + tree.add_channel_current(k_chan, 0.653374 * 1e6, -85.0, node_arg=[tree[1]]) + na_chan = channelcollection.NaTa_t() + tree.add_channel_current(na_chan, 0.15 * 1e6, 50.0, node_arg=[tree[1]]) + ca_chan = channelcollection.Ca_HVA() + tree.add_channel_current( + ca_chan, 0.005 * 1e6, 132.4579341637009, node_arg=[tree[1]] + ) + tree.fit_leak_current(-70.0, 15.0) + self.tree = tree + + # fit a single somatic compartment + cfit = CompartmentFitter(tree, save_cache=False, recompute_cache=True) + self.ctree, _ = cfit.fit_model([(1, 0.5)], kernel_correction=kernel_correction) + def test_neuroncompartmentree_instantiation(self): self.load_multi_dend_model() # correct initialization @@ -451,6 +482,47 @@ def test_neuroncompartmentree_instantiation(self): with pytest.raises(ValueError): neuron_sim_tree = NeuronCompartmentTree(GreensTree()) + def test_ignore_correction_compartments(self): + """ + A somatic fit with an admittance-kernel correction attaches passive + dummy compartments (``loc_idx is None``) to the soma. These cannot be + represented in NEURON, so `NeuronCompartmentTree` must drop them (with + a warning) while leaving the location-bearing compartments intact. + """ + self.load_T_tree(kernel_correction=[0]) + ctree = self.ctree + + # the correction should have induced dummy compartments + dummy_idxs = [n.index for n in ctree if n.loc_idx is None] + real_idxs = [n.index for n in ctree if n.loc_idx is not None] + assert len(dummy_idxs) > 0, ( + "the somatic admittance correction did not induce any dummy " + "compartments; the test cannot verify they are ignored" + ) + assert ctree.has_correction_compartments() + + # building the NEURON model should warn about (and ignore) the dummies + with pytest.warns(UserWarning): + sim_tree = NeuronCompartmentTree(ctree) + + # the location-bearing compartments are present in the NEURON model ... + for ri in real_idxs: + assert sim_tree[ri] is not None + # ... while the dummy compartments are absent + for di in dummy_idxs: + assert sim_tree[di] is None + + # the number of NEURON sections matches the number of real compartments + assert len([n for n in sim_tree]) == len(real_idxs) + + # the caller's tree must be left untouched (dummies still present) + assert ctree.has_correction_compartments() + assert len([n for n in ctree if n.loc_idx is None]) == len(dummy_idxs) + + # the pruned model must be buildable (no non-physical geometry) + sim_tree.init_model(dt=0.1, t_calibrate=0.0) + sim_tree.delete_model() + def test_geometry1(self): fake_c_m = 1.0 fake_r_a = 100.0 * 1e-6 @@ -616,9 +688,11 @@ def test_impedance_properties_1(self): # create the two compartment model without locinds self.load_two_compartment_model(w_locinds=False) ctree = self.ctree - # check if error is raised if loc_idxs have not been set - with pytest.raises(AttributeError): - ctree.calc_impedance_matrix() + # check if matrix is 2x2 if indexing=='tree' and 0x0 otherwise + zmat_tree = ctree.calc_impedance_matrix(indexing="tree") + zmat_locs = ctree.calc_impedance_matrix(indexing="locs") + assert zmat_tree.shape == (2, 2) + assert zmat_locs.shape == (0, 0) # create the two compartment model with locinds self.load_two_compartment_model() @@ -1160,13 +1234,13 @@ def debug_print(pstr): # tn.test_channel_recording() # tn.test_recording_timestep() - # trn = TestReducedNeuron() + trn = TestReducedNeuron() # trn.test_geometry1() - # trn.test_impedance_properties_1() + trn.test_impedance_properties_1() # trn.test_geometry2() # trn.test_impedance_properties_2() - ts = TestStimuli() + # ts = TestStimuli() # ts.test_i_clamp(pplot=True) # ts.test_ou_processes() - ts.test_input_spiketrain(pplot=True) + # ts.test_input_spiketrain(pplot=True)