Add KineticIon support and family-phased cell updates#89
Open
TLDSZ wants to merge 7 commits into
Open
Conversation
2.add update_policy
2.solve the independent ion with non-indepent channel situation
Contributor
Reviewer's GuideIntroduces a declarative KineticIon template and associated runtime plumbing, adds an explicit Cell.update_policy with a new family_phased scheduling mode that orders voltage, channels, and ions, fixes mixed update-order for current-driven ions via cached step-start current, extends runtime ion-family detection and mechanism probing, and updates tests and documentation accordingly. Sequence diagram for Cell.update family_phased schedulingsequenceDiagram
participant Cell
participant DHS as dhs_voltage_step
participant Chan as RuntimeChannel
participant Ion as RuntimeIon
rect rgb(230,230,250)
Cell->>Cell: _cache_family_phased_currents(point_V_before)
Cell->>DHS: dhs_voltage_step(self,t,dt,I_ext)
DHS-->>Cell: V updated
end
rect rgb(220,255,220)
Cell->>Chan: _channel_phase_call(pre_integral,point_V_after,ion)
Cell->>Chan: _channel_phase_call(compute_derivative,point_V_after,ion)
Cell->>Chan: _channel_phase_call(post_integral,point_V_after,ion)
end
rect rgb(220,240,255)
Cell->>Ion: _run_ion_hook(_ion_compute_derivative_hook,point_V_after)
Ion-->>Ion: _ion_compute_derivative_hook(V)
Ion-->>Ion: derivative(Ci,V,total_current=_cached_total_current)
Cell->>Ion: _run_ion_hook(_ion_post_integral_hook,point_V_after)
end
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
Contributor
There was a problem hiding this comment.
Hey - I've found 5 issues, and left some high level feedback:
- The
KineticIon.uses_total_currentpath in_ion_compute_derivative_hookhard-requires_cached_total_currentand raises at runtime if it is missing; consider either guarding this path behind an explicitupdate_policy/owner check or providing a clearer public contract for how kinetic ions are expected to be driven outsideCell(update_policy='family_phased'). - The
Cell._update_family_phasedimplementation introduces a fairly involved parallel scheduling stack (_PhaseTarget,_collect_*,_run_independent_phase_solver); it may be worth isolating this into a dedicated module or reusing more of the existing solver abstraction to avoid duplication and make it easier to keep phase ordering consistent across future solvers.
Prompt for AI Agents
Please address the comments from this code review:
## Overall Comments
- The `KineticIon.uses_total_current` path in `_ion_compute_derivative_hook` hard-requires `_cached_total_current` and raises at runtime if it is missing; consider either guarding this path behind an explicit `update_policy`/owner check or providing a clearer public contract for how kinetic ions are expected to be driven outside `Cell(update_policy='family_phased')`.
- The `Cell._update_family_phased` implementation introduces a fairly involved parallel scheduling stack (`_PhaseTarget`, `_collect_*`, `_run_independent_phase_solver`); it may be worth isolating this into a dedicated module or reusing more of the existing solver abstraction to avoid duplication and make it easier to keep phase ordering consistent across future solvers.
## Individual Comments
### Comment 1
<location path="braincell/_base_ion.py" line_range="222-228" />
<code_context>
node.reset_state(V, ion_info, batch_size)
def update(self, V, *args, **kwargs):
+ if isinstance(self, IndependentIntegration):
+ self.make_integration(V)
+ ion_info = self.pack_info()
+ for key, node in brainstate.graph.nodes(self, Channel, allowed_hierarchy=(1, 1)).items():
+ if isinstance(node, IndependentIntegration):
+ node.update(V, ion_info)
+ return
ion_info = self.pack_info()
for key, node in brainstate.graph.nodes(self, Channel, allowed_hierarchy=(1, 1)).items():
</code_context>
<issue_to_address>
**issue (bug_risk):** IndependentIntegration ions skip non-independent child channels in update().
With this override, when `self` is an `IndependentIntegration`, only child channels that are also `IndependentIntegration` get their `update` called; any other `Channel` children are skipped entirely. If such mixed configurations are valid, those non-independent channels will stop updating once the parent becomes independent. Consider either calling `update` on all child channels:
```python
if isinstance(self, IndependentIntegration):
self.make_integration(V)
ion_info = self.pack_info()
for _, node in brainstate.graph.nodes(self, Channel, allowed_hierarchy=(1, 1)).items():
node.update(V, ion_info)
return
```
or clearly documenting/enforcing that independent ions cannot own non-independent channels.
</issue_to_address>
### Comment 2
<location path="braincell/_multi_compartment/cell.py" line_range="1814-1825" />
<code_context>
+ else:
+ ion._run_ion_hook("_ion_post_integral_hook", point_V_after)
+
+ def _cache_family_phased_currents(self, point_V):
+ runtime = self.runtime
+ for ion_name, ion in runtime.ions.items():
+ if getattr(type(ion), "uses_total_current", False):
+ # Cache a step-start current snapshot for ion dynamics only.
+ # The voltage phase still evaluates membrane current through the
+ # existing path; this cache merely prevents ion dynamics from
+ # re-reading a newer current after channel / voltage states have
+ # already advanced.
+ ion._cached_total_current = ion.current(point_V, include_external=True)
+ ion._cached_total_current_policy = self._update_policy
+ ion._cached_total_current_time = self._resolve_t()
+
def reset_state(self, batch_size=None) -> None:
</code_context>
<issue_to_address>
**suggestion:** Cached current metadata fields are written but never read, suggesting either dead state or a missing consumer.
The `_cache_family_phased_currents` helper writes three attributes on each ion:
```python
aio n._cached_total_current = ...
ion._cached_total_current_policy = self._update_policy
ion._cached_total_current_time = self._resolve_t()
```
Only `_cached_total_current` is read later (in `DynamicNernstIon` / `KineticIon` hooks); the policy/time fields are unused in the new code. If they’re meant for validation (e.g., checking the cache’s update policy or timestamp), consider adding those checks now. Otherwise, drop the unused attributes to avoid dead state and implied but unenforced invariants.
```suggestion
def _cache_family_phased_currents(self, point_V):
runtime = self.runtime
for ion_name, ion in runtime.ions.items():
if getattr(type(ion), "uses_total_current", False):
# Cache a step-start current snapshot for ion dynamics only.
# The voltage phase still evaluates membrane current through the
# existing path; this cache merely prevents ion dynamics from
# re-reading a newer current after channel / voltage states have
# already advanced.
ion._cached_total_current = ion.current(point_V, include_external=True)
```
</issue_to_address>
### Comment 3
<location path="braincell/ion/_base_test.py" line_range="113-122" />
<code_context>
+class _SimpleKineticIon(Calcium, KineticIon):
</code_context>
<issue_to_address>
**suggestion (testing):** Add tests for KineticIon/_Specs validation error paths (missing Ci, duplicate species/factors, invalid conserves/reactions).
The new KineticIon/_Specs logic has substantial validation (unique factor/species names, required 'Ci', valid Conserve membership and size, no algebraic 'Ci', reaction stoichiometry rules, and source/target constraints), but current tests only hit the happy path via _SimpleKineticIon. Please add table-driven tests with minimal dummy KineticIon subclasses that trigger these failures (e.g., missing 'Ci', duplicate species, Conserve with non-species/algebraic 'Ci', invalid stoichiometry, Source targeting algebraic species) and assert the expected ValueError/TypeError from init_state or species_values.
Suggested implementation:
```python
import pytest
from braincell.ion._base import Reaction
```
```python
from braincell.ion._base import Source
from braincell.ion._base import Species
from braincell.quad.protocol import DiffEqState
class _MissingCiIon(KineticIon):
"""KineticIon without a 'Ci' species – should fail Specs validation."""
default_Co = 2.0 * u.mM
default_valence = 2
factors = (
Factor("cyto", lambda self: self.cyt_volume),
)
# Intentionally omit 'Ci'
species = (
Species("B", init=1.0 * u.mM, factor="cyto"),
)
class _DuplicateSpeciesIon(KineticIon):
"""KineticIon with duplicate species names – should fail Specs validation."""
default_Co = 2.0 * u.mM
default_valence = 2
factors = (
Factor("cyto", lambda self: self.cyt_volume),
)
species = (
Species("Ci", init=0.1 * u.mM, factor="cyto"),
# Duplicate name 'Ci'
Species("Ci", init=0.2 * u.mM, factor="cyto"),
)
class _AlgebraicCiIon(KineticIon):
"""KineticIon with algebraic 'Ci' – Specs should forbid algebraic Ci."""
default_Co = 2.0 * u.mM
default_valence = 2
factors = (
Factor("cyto", lambda self: self.cyt_volume),
)
# Mark Ci as algebraic / non-differential according to Species API
species = (
Species("Ci", init=0.1 * u.mM, factor="cyto", algebraic=True),
)
class _InvalidConserveIon(KineticIon):
"""KineticIon with invalid Conserve definition (non-species/algebraic Ci)."""
default_Co = 2.0 * u.mM
default_valence = 2
factors = (
Factor("cyto", lambda self: self.cyt_volume),
)
species = (
Species("Ci", init=0.1 * u.mM, factor="cyto"),
Species("B", init=1.0 * u.mM, factor="cyto"),
)
# Conserve definition is intentionally invalid; exact API may differ.
conserves = (
Conserve("total_B", members=("Ci", "nonexistent_species")),
)
class _InvalidReactionStoichIon(KineticIon):
"""KineticIon with invalid reaction stoichiometry rules."""
default_Co = 2.0 * u.mM
default_valence = 2
factors = (
Factor("cyto", lambda self: self.cyt_volume),
)
species = (
Species("Ci", init=0.1 * u.mM, factor="cyto"),
Species("B", init=1.0 * u.mM, factor="cyto"),
)
# Reaction with illegal stoichiometry (e.g., Ci only on one side, or
# violating charge/particle conservation) – adjust to match actual rules.
reactions = (
Reaction("bad_reaction", reactants={"Ci": 1}, products={"B": 2}),
)
class _SourceToAlgebraicSpeciesIon(KineticIon):
"""KineticIon with a Source targeting an algebraic species."""
default_Co = 2.0 * u.mM
default_valence = 2
factors = (
Factor("cyto", lambda self: self.cyt_volume),
)
species = (
Species("Ci", init=0.1 * u.mM, factor="cyto"),
Species("A", init=0.0 * u.mM, factor="cyto", algebraic=True),
)
# Source is intentionally targeting algebraic species 'A'
sources = (
Source("bad_source", target="A", current=lambda self, t: 0 * u.mM / u.ms),
)
@pytest.mark.parametrize(
"ion_cls, error_substring",
(
(_MissingCiIon, "Ci"),
(_DuplicateSpeciesIon, "duplicate"),
(_AlgebraicCiIon, "algebraic"),
(_InvalidConserveIon, "Conserve"),
(_InvalidReactionStoichIon, "reaction"),
(_SourceToAlgebraicSpeciesIon, "Source"),
),
)
def test_kinetic_ion_specs_validation_errors(ion_cls, error_substring):
"""
Table-driven tests that exercise KineticIon/_Specs validation error paths.
Each dummy KineticIon subclass has a deliberately invalid Specs definition.
We assert that initializing state (or computing species values) raises the
appropriate validation error.
"""
with pytest.raises((ValueError, TypeError)) as excinfo:
ion = ion_cls() # __init__ signature may require args; see additional_changes.
# Trigger Specs validation. Depending on implementation, this may
# happen on __init__, init_state, or species_values.
#
# The following line assumes init_state exists and takes a DiffEqState
# (or similar) – adjust as needed.
ion.init_state(DiffEqState())
assert error_substring in str(excinfo.value)
class _RecorderChannel(Channel):
```
1. Ensure the following imports exist (either already present or added near the top of the file): `from braincell.ion._base import KineticIon, Factor, Conserve, Calcium` and the units alias `from braincell import u` (or whatever the project uses for units). Adjust the `Source`, `Reaction`, `Species`, `Factor`, `Conserve`, and `DiffEqState` imports to match the actual module layout.
2. The `Species` constructor in your codebase may not use an `algebraic` keyword; replace `algebraic=True` with the correct way to define algebraic/non-differential species (e.g., `ode=None`, `expr=...`, or a specific flag) and adjust the tests accordingly.
3. The `Conserve`, `Reaction`, and `Source` constructors in this patch are based on typical signatures. Update the arguments (`members=`, `reactants=`, `products=`, `target=`, etc.) to match your actual API and to create configurations that violate your specific validation rules (e.g., invalid membership, illegal stoichiometry, or Sources targeting algebraic species).
4. The `ion_cls()` call in `test_kinetic_ion_specs_validation_errors` assumes a no-arg constructor; if your `KineticIon` subclasses require parameters (e.g., morphology, channel density, or diffusion coefficients), update the test to construct them with minimal valid arguments.
5. The `ion.init_state(DiffEqState())` call is a placeholder for whatever method actually triggers `_Specs` validation in your implementation (`init_state`, `species_values`, or another initializer). Replace this with the correct method and a minimal `DiffEqState` instance (or other required arguments) to ensure each invalid Specs configuration raises the expected `ValueError`/`TypeError`.
6. If your validation errors use specific custom exception types or distinct messages, narrow the `pytest.raises` type and replace the generic `error_substring` checks with exact message matches or more precise assertions.
</issue_to_address>
### Comment 4
<location path="braincell/_multi_compartment/cell_test.py" line_range="296-305" />
<code_context>
+ with self.assertRaises(ValueError):
+ Cell(_simple_cell().morpho, cv_policy=CVPerBranch(), update_policy="bad")
+
+ def test_family_phased_channel_runs_before_ion(self):
+ cell = self._build_family_phased_cell()
+ cell.update_policy = "family_phased"
+ cell.init_state()
+ ion = cell.get_ion("ca_dyn")
+ channel = ion.channels["ca_dep"]
+ ion.Ci.value = jnp.asarray([1.0]) * u.mM
+ with brainstate.environ.context(dt=0.1 * u.ms):
+ cell.update()
+
+ self.assertEqual(_PHASE_LOG, ["channel_compute", "ion_compute"])
+
+ def test_family_phased_ion_uses_cached_current(self):
</code_context>
<issue_to_address>
**suggestion (testing):** Exercise family_phased behaviour with independently integrated channels/ions to cover the independent-path helpers.
The new family-phased code adds helpers and special handling for IndependentIntegration channels/ions, but the tests only cover the regular (non-independent) path. Please add or extend a test so that a channel, an ion, or both use IndependentIntegration with non-trivial substeps, and assert that (1) family order is preserved and (2) their make_integration/solver path is exercised (e.g., via call logging). This will ensure the independent path remains covered and robust to future refactors.
Suggested implementation:
```python
def test_family_phased_ion_uses_cached_current(self):
cell = self._build_family_phased_cell()
cell.update_policy = "family_phased"
cell.init_state()
ion = cell.get_ion("ca_dyn")
channel = ion.channels["ca_dep"]
ion.Ci.value = jnp.asarray([1.0]) * u.mM
channel.p.value = jnp.asarray([999.0])
with brainstate.environ.context(dt=0.1 * u.ms):
cell.update()
def test_family_phased_independent_integration_channel_and_ion(self):
"""Exercise IndependentIntegration helpers under the family_phased policy.
This ensures that:
1. The independent integration path (make_integration / solver) is exercised.
2. Family order is preserved (channel runs before ion) even for independent paths.
"""
# Build a cell where both the channel and ion are configured for independent integration
cell = self._build_family_phased_cell(
independent_channel=True,
independent_ion=True,
)
cell.update_policy = "family_phased"
cell.init_state()
ion = cell.get_ion("ca_dyn")
channel = ion.channels["ca_dep"]
# Initialise some non-trivial state so that the integrators actually run
ion.Ci.value = jnp.asarray([1.0]) * u.mM
channel.p.value = jnp.asarray([0.5])
# Clear phase log so we only see events from this test
_PHASE_LOG.clear()
# Use a non-trivial number of substeps to exercise the independent solver helpers
with brainstate.environ.context(dt=0.1 * u.ms, substeps=5):
cell.update()
# Assert that the independent integration helpers were invoked.
# These entries should be appended by the IndependentIntegration-aware
# make_integration / solver wrappers for the channel and ion.
self.assertIn("independent_channel_make_integration", _PHASE_LOG)
self.assertIn("independent_ion_make_integration", _PHASE_LOG)
# Also ensure that the compute phase order is preserved for the family:
# the channel's compute should precede the ion's compute in the phase log.
self.assertLess(
_PHASE_LOG.index("independent_channel_compute"),
_PHASE_LOG.index("independent_ion_compute"),
)
```
To fully implement this test and satisfy the comment:
1. **Extend `_build_family_phased_cell`** (in this file or the corresponding helper module) to accept `independent_channel` and `independent_ion` keyword arguments. When these are `True`, construct the `"ca_dep"` channel and `"ca_dyn"` ion using the IndependentIntegration path (e.g. by wrapping their integration in the IndependentIntegration helper or by selecting the IndependentIntegration-based implementation).
2. **Hook into the IndependentIntegration helpers** used by channels/ions so they append the following markers to `_PHASE_LOG`:
- `"independent_channel_make_integration"` when the independent channel's `make_integration`/solver wrapper is invoked.
- `"independent_ion_make_integration"` when the independent ion's `make_integration`/solver wrapper is invoked.
- `"independent_channel_compute"` and `"independent_ion_compute"` when their respective compute phases run under the independent path.
The exact place to log these should be inside the IndependentIntegration-specific `make_integration`/solver implementations for the channel and ion, consistent with how the existing `"channel_compute"` / `"ion_compute"` markers are emitted for the non-independent path.
3. Ensure that the IndependentIntegration helpers respect the `substeps` argument from `brainstate.environ.context`, so the `substeps=5` in the new test actually drives the independent solver path (rather than being ignored).
</issue_to_address>
### Comment 5
<location path="braincell/ion/_base.py" line_range="332" />
<code_context>
raise NotImplementedError
+
+
+class KineticIon(IndependentIntegration):
+ """Template for NMODL-style kinetic ion species.
+
</code_context>
<issue_to_address>
**issue (complexity):** Consider caching _Specs, _Species, _Conserve, and _Flux as per-instance adapters in _init_kinetic_ion and reusing them in hooks to simplify the runtime pipeline and avoid redundant constructions.
You can reduce both structural and cognitive complexity without changing behavior by instantiating the runtime helpers once per ion instance instead of rebuilding `_Specs`, `_Species`, `_Conserve`, and `_Flux` in every hook.
That flattens the helper stack and makes the data flow clearer (“this instance has its runtime adapter”) while also avoiding repeated allocations.
### 1. Cache runtime adapters on the instance
In `_init_kinetic_ion`, build and cache the spec/runtime helpers:
```python
class KineticIon(IndependentIntegration):
...
def _init_kinetic_ion(
self,
*,
Co=None,
temp=None,
valence=None,
solver: str = "rk4",
substeps: int = 1,
):
...
self.temp = braintools.init.param(temp, self.varshape, allow_none=False)
# NEW: cache specs + runtime helpers
self._specs = _Specs.for_type(type(self))
self._species = _Species(self, self._specs)
self._conserve = _Conserve(self, self._specs, self._species)
self._flux = _Flux(self, self._specs, self._species)
```
This keeps all existing validation/caching behavior but turns `_Specs` into a per-type “compile-time” artifact and `_Species`/`_Conserve`/`_Flux` into per-instance “runtime adapters” with explicit ownership.
### 2. Reuse cached adapters in hooks
Then simplify the hooks to use the cached adapters instead of re-creating them:
```python
class KineticIon(IndependentIntegration):
...
def species_values(self):
"""Return the current full visible species view."""
return self._conserve.resolve()
def _ion_init_state_hook(self, V, batch_size: int = None):
"""Initialize runtime species and project algebraic species."""
self._species.init(batch_size=batch_size)
self._conserve.writeback(V)
def _ion_reset_state_hook(self, V, batch_size: int = None):
"""Reset runtime species and project algebraic species."""
self._species.reset(batch_size=batch_size)
self._conserve.writeback(V)
def _ion_compute_derivative_hook(self, V):
"""Resolve algebraic species and write diffeq derivatives."""
total_current = None
if type(self).uses_total_current:
if not hasattr(self, "_cached_total_current"):
raise RuntimeError(
f"{type(self).__name__} requires a cached total current before compute_derivative()."
)
total_current = self._cached_total_current
species_values = self._conserve.resolve(V)
self._flux.compute(V, species_values, total_current=total_current)
def _ion_post_integral_hook(self, V):
"""Refresh cached algebraic species after one integration step."""
self._conserve.writeback(V)
```
Key benefits:
- Eliminates repeated `_Specs.for_type(type(self))` and `_Species/_Conserve/_Flux` constructions.
- Makes the runtime pipeline more discoverable (`self._species`, `self._conserve`, `self._flux` are clearly the adapters).
- Preserves all existing semantics and public APIs.
This is a small, targeted refactor that directly addresses point (1) from the review (“deep helper stack and indirection”) without removing features or changing the declarative DSL.
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
Comment on lines
+222
to
+228
| if isinstance(self, IndependentIntegration): | ||
| self.make_integration(V) | ||
| ion_info = self.pack_info() | ||
| for key, node in brainstate.graph.nodes(self, Channel, allowed_hierarchy=(1, 1)).items(): | ||
| if isinstance(node, IndependentIntegration): | ||
| node.update(V, ion_info) | ||
| return |
Contributor
There was a problem hiding this comment.
issue (bug_risk): IndependentIntegration ions skip non-independent child channels in update().
With this override, when self is an IndependentIntegration, only child channels that are also IndependentIntegration get their update called; any other Channel children are skipped entirely. If such mixed configurations are valid, those non-independent channels will stop updating once the parent becomes independent. Consider either calling update on all child channels:
if isinstance(self, IndependentIntegration):
self.make_integration(V)
ion_info = self.pack_info()
for _, node in brainstate.graph.nodes(self, Channel, allowed_hierarchy=(1, 1)).items():
node.update(V, ion_info)
returnor clearly documenting/enforcing that independent ions cannot own non-independent channels.
Comment on lines
+1814
to
+1825
| def _cache_family_phased_currents(self, point_V): | ||
| runtime = self.runtime | ||
| for ion_name, ion in runtime.ions.items(): | ||
| if getattr(type(ion), "uses_total_current", False): | ||
| # Cache a step-start current snapshot for ion dynamics only. | ||
| # The voltage phase still evaluates membrane current through the | ||
| # existing path; this cache merely prevents ion dynamics from | ||
| # re-reading a newer current after channel / voltage states have | ||
| # already advanced. | ||
| ion._cached_total_current = ion.current(point_V, include_external=True) | ||
| ion._cached_total_current_policy = self._update_policy | ||
| ion._cached_total_current_time = self._resolve_t() |
Contributor
There was a problem hiding this comment.
suggestion: Cached current metadata fields are written but never read, suggesting either dead state or a missing consumer.
The _cache_family_phased_currents helper writes three attributes on each ion:
aio n._cached_total_current = ...
ion._cached_total_current_policy = self._update_policy
ion._cached_total_current_time = self._resolve_t()Only _cached_total_current is read later (in DynamicNernstIon / KineticIon hooks); the policy/time fields are unused in the new code. If they’re meant for validation (e.g., checking the cache’s update policy or timestamp), consider adding those checks now. Otherwise, drop the unused attributes to avoid dead state and implied but unenforced invariants.
Suggested change
| def _cache_family_phased_currents(self, point_V): | |
| runtime = self.runtime | |
| for ion_name, ion in runtime.ions.items(): | |
| if getattr(type(ion), "uses_total_current", False): | |
| # Cache a step-start current snapshot for ion dynamics only. | |
| # The voltage phase still evaluates membrane current through the | |
| # existing path; this cache merely prevents ion dynamics from | |
| # re-reading a newer current after channel / voltage states have | |
| # already advanced. | |
| ion._cached_total_current = ion.current(point_V, include_external=True) | |
| ion._cached_total_current_policy = self._update_policy | |
| ion._cached_total_current_time = self._resolve_t() | |
| def _cache_family_phased_currents(self, point_V): | |
| runtime = self.runtime | |
| for ion_name, ion in runtime.ions.items(): | |
| if getattr(type(ion), "uses_total_current", False): | |
| # Cache a step-start current snapshot for ion dynamics only. | |
| # The voltage phase still evaluates membrane current through the | |
| # existing path; this cache merely prevents ion dynamics from | |
| # re-reading a newer current after channel / voltage states have | |
| # already advanced. | |
| ion._cached_total_current = ion.current(point_V, include_external=True) |
Comment on lines
+113
to
+122
| class _SimpleKineticIon(Calcium, KineticIon): | ||
| default_Co = 2.0 * u.mM | ||
| default_valence = 2 | ||
|
|
||
| factors = ( | ||
| Factor("cyto", lambda self: self.cyt_volume), | ||
| ) | ||
| species = ( | ||
| Species("Ci", init=0.1 * u.mM, factor="cyto"), | ||
| Species("B", init=1.0 * u.mM, factor="cyto"), |
Contributor
There was a problem hiding this comment.
suggestion (testing): Add tests for KineticIon/_Specs validation error paths (missing Ci, duplicate species/factors, invalid conserves/reactions).
The new KineticIon/_Specs logic has substantial validation (unique factor/species names, required 'Ci', valid Conserve membership and size, no algebraic 'Ci', reaction stoichiometry rules, and source/target constraints), but current tests only hit the happy path via _SimpleKineticIon. Please add table-driven tests with minimal dummy KineticIon subclasses that trigger these failures (e.g., missing 'Ci', duplicate species, Conserve with non-species/algebraic 'Ci', invalid stoichiometry, Source targeting algebraic species) and assert the expected ValueError/TypeError from init_state or species_values.
Suggested implementation:
import pytest
from braincell.ion._base import Reactionfrom braincell.ion._base import Source
from braincell.ion._base import Species
from braincell.quad.protocol import DiffEqState
class _MissingCiIon(KineticIon):
"""KineticIon without a 'Ci' species – should fail Specs validation."""
default_Co = 2.0 * u.mM
default_valence = 2
factors = (
Factor("cyto", lambda self: self.cyt_volume),
)
# Intentionally omit 'Ci'
species = (
Species("B", init=1.0 * u.mM, factor="cyto"),
)
class _DuplicateSpeciesIon(KineticIon):
"""KineticIon with duplicate species names – should fail Specs validation."""
default_Co = 2.0 * u.mM
default_valence = 2
factors = (
Factor("cyto", lambda self: self.cyt_volume),
)
species = (
Species("Ci", init=0.1 * u.mM, factor="cyto"),
# Duplicate name 'Ci'
Species("Ci", init=0.2 * u.mM, factor="cyto"),
)
class _AlgebraicCiIon(KineticIon):
"""KineticIon with algebraic 'Ci' – Specs should forbid algebraic Ci."""
default_Co = 2.0 * u.mM
default_valence = 2
factors = (
Factor("cyto", lambda self: self.cyt_volume),
)
# Mark Ci as algebraic / non-differential according to Species API
species = (
Species("Ci", init=0.1 * u.mM, factor="cyto", algebraic=True),
)
class _InvalidConserveIon(KineticIon):
"""KineticIon with invalid Conserve definition (non-species/algebraic Ci)."""
default_Co = 2.0 * u.mM
default_valence = 2
factors = (
Factor("cyto", lambda self: self.cyt_volume),
)
species = (
Species("Ci", init=0.1 * u.mM, factor="cyto"),
Species("B", init=1.0 * u.mM, factor="cyto"),
)
# Conserve definition is intentionally invalid; exact API may differ.
conserves = (
Conserve("total_B", members=("Ci", "nonexistent_species")),
)
class _InvalidReactionStoichIon(KineticIon):
"""KineticIon with invalid reaction stoichiometry rules."""
default_Co = 2.0 * u.mM
default_valence = 2
factors = (
Factor("cyto", lambda self: self.cyt_volume),
)
species = (
Species("Ci", init=0.1 * u.mM, factor="cyto"),
Species("B", init=1.0 * u.mM, factor="cyto"),
)
# Reaction with illegal stoichiometry (e.g., Ci only on one side, or
# violating charge/particle conservation) – adjust to match actual rules.
reactions = (
Reaction("bad_reaction", reactants={"Ci": 1}, products={"B": 2}),
)
class _SourceToAlgebraicSpeciesIon(KineticIon):
"""KineticIon with a Source targeting an algebraic species."""
default_Co = 2.0 * u.mM
default_valence = 2
factors = (
Factor("cyto", lambda self: self.cyt_volume),
)
species = (
Species("Ci", init=0.1 * u.mM, factor="cyto"),
Species("A", init=0.0 * u.mM, factor="cyto", algebraic=True),
)
# Source is intentionally targeting algebraic species 'A'
sources = (
Source("bad_source", target="A", current=lambda self, t: 0 * u.mM / u.ms),
)
@pytest.mark.parametrize(
"ion_cls, error_substring",
(
(_MissingCiIon, "Ci"),
(_DuplicateSpeciesIon, "duplicate"),
(_AlgebraicCiIon, "algebraic"),
(_InvalidConserveIon, "Conserve"),
(_InvalidReactionStoichIon, "reaction"),
(_SourceToAlgebraicSpeciesIon, "Source"),
),
)
def test_kinetic_ion_specs_validation_errors(ion_cls, error_substring):
"""
Table-driven tests that exercise KineticIon/_Specs validation error paths.
Each dummy KineticIon subclass has a deliberately invalid Specs definition.
We assert that initializing state (or computing species values) raises the
appropriate validation error.
"""
with pytest.raises((ValueError, TypeError)) as excinfo:
ion = ion_cls() # __init__ signature may require args; see additional_changes.
# Trigger Specs validation. Depending on implementation, this may
# happen on __init__, init_state, or species_values.
#
# The following line assumes init_state exists and takes a DiffEqState
# (or similar) – adjust as needed.
ion.init_state(DiffEqState())
assert error_substring in str(excinfo.value)
class _RecorderChannel(Channel):- Ensure the following imports exist (either already present or added near the top of the file):
from braincell.ion._base import KineticIon, Factor, Conserve, Calciumand the units aliasfrom braincell import u(or whatever the project uses for units). Adjust theSource,Reaction,Species,Factor,Conserve, andDiffEqStateimports to match the actual module layout. - The
Speciesconstructor in your codebase may not use analgebraickeyword; replacealgebraic=Truewith the correct way to define algebraic/non-differential species (e.g.,ode=None,expr=..., or a specific flag) and adjust the tests accordingly. - The
Conserve,Reaction, andSourceconstructors in this patch are based on typical signatures. Update the arguments (members=,reactants=,products=,target=, etc.) to match your actual API and to create configurations that violate your specific validation rules (e.g., invalid membership, illegal stoichiometry, or Sources targeting algebraic species). - The
ion_cls()call intest_kinetic_ion_specs_validation_errorsassumes a no-arg constructor; if yourKineticIonsubclasses require parameters (e.g., morphology, channel density, or diffusion coefficients), update the test to construct them with minimal valid arguments. - The
ion.init_state(DiffEqState())call is a placeholder for whatever method actually triggers_Specsvalidation in your implementation (init_state,species_values, or another initializer). Replace this with the correct method and a minimalDiffEqStateinstance (or other required arguments) to ensure each invalid Specs configuration raises the expectedValueError/TypeError. - If your validation errors use specific custom exception types or distinct messages, narrow the
pytest.raisestype and replace the genericerror_substringchecks with exact message matches or more precise assertions.
Comment on lines
+296
to
+305
| def test_family_phased_channel_runs_before_ion(self): | ||
| cell = self._build_family_phased_cell() | ||
| cell.update_policy = "family_phased" | ||
| cell.init_state() | ||
| ion = cell.get_ion("ca_dyn") | ||
| channel = ion.channels["ca_dep"] | ||
| ion.Ci.value = jnp.asarray([1.0]) * u.mM | ||
| with brainstate.environ.context(dt=0.1 * u.ms): | ||
| cell.update() | ||
|
|
Contributor
There was a problem hiding this comment.
suggestion (testing): Exercise family_phased behaviour with independently integrated channels/ions to cover the independent-path helpers.
The new family-phased code adds helpers and special handling for IndependentIntegration channels/ions, but the tests only cover the regular (non-independent) path. Please add or extend a test so that a channel, an ion, or both use IndependentIntegration with non-trivial substeps, and assert that (1) family order is preserved and (2) their make_integration/solver path is exercised (e.g., via call logging). This will ensure the independent path remains covered and robust to future refactors.
Suggested implementation:
def test_family_phased_ion_uses_cached_current(self):
cell = self._build_family_phased_cell()
cell.update_policy = "family_phased"
cell.init_state()
ion = cell.get_ion("ca_dyn")
channel = ion.channels["ca_dep"]
ion.Ci.value = jnp.asarray([1.0]) * u.mM
channel.p.value = jnp.asarray([999.0])
with brainstate.environ.context(dt=0.1 * u.ms):
cell.update()
def test_family_phased_independent_integration_channel_and_ion(self):
"""Exercise IndependentIntegration helpers under the family_phased policy.
This ensures that:
1. The independent integration path (make_integration / solver) is exercised.
2. Family order is preserved (channel runs before ion) even for independent paths.
"""
# Build a cell where both the channel and ion are configured for independent integration
cell = self._build_family_phased_cell(
independent_channel=True,
independent_ion=True,
)
cell.update_policy = "family_phased"
cell.init_state()
ion = cell.get_ion("ca_dyn")
channel = ion.channels["ca_dep"]
# Initialise some non-trivial state so that the integrators actually run
ion.Ci.value = jnp.asarray([1.0]) * u.mM
channel.p.value = jnp.asarray([0.5])
# Clear phase log so we only see events from this test
_PHASE_LOG.clear()
# Use a non-trivial number of substeps to exercise the independent solver helpers
with brainstate.environ.context(dt=0.1 * u.ms, substeps=5):
cell.update()
# Assert that the independent integration helpers were invoked.
# These entries should be appended by the IndependentIntegration-aware
# make_integration / solver wrappers for the channel and ion.
self.assertIn("independent_channel_make_integration", _PHASE_LOG)
self.assertIn("independent_ion_make_integration", _PHASE_LOG)
# Also ensure that the compute phase order is preserved for the family:
# the channel's compute should precede the ion's compute in the phase log.
self.assertLess(
_PHASE_LOG.index("independent_channel_compute"),
_PHASE_LOG.index("independent_ion_compute"),
)To fully implement this test and satisfy the comment:
-
Extend
_build_family_phased_cell(in this file or the corresponding helper module) to acceptindependent_channelandindependent_ionkeyword arguments. When these areTrue, construct the"ca_dep"channel and"ca_dyn"ion using the IndependentIntegration path (e.g. by wrapping their integration in the IndependentIntegration helper or by selecting the IndependentIntegration-based implementation). -
Hook into the IndependentIntegration helpers used by channels/ions so they append the following markers to
_PHASE_LOG:"independent_channel_make_integration"when the independent channel'smake_integration/solver wrapper is invoked."independent_ion_make_integration"when the independent ion'smake_integration/solver wrapper is invoked."independent_channel_compute"and"independent_ion_compute"when their respective compute phases run under the independent path.
The exact place to log these should be inside the IndependentIntegration-specific
make_integration/solver implementations for the channel and ion, consistent with how the existing"channel_compute"/"ion_compute"markers are emitted for the non-independent path. -
Ensure that the IndependentIntegration helpers respect the
substepsargument frombrainstate.environ.context, so thesubsteps=5in the new test actually drives the independent solver path (rather than being ignored).
| raise NotImplementedError | ||
|
|
||
|
|
||
| class KineticIon(IndependentIntegration): |
Contributor
There was a problem hiding this comment.
issue (complexity): Consider caching _Specs, _Species, _Conserve, and _Flux as per-instance adapters in _init_kinetic_ion and reusing them in hooks to simplify the runtime pipeline and avoid redundant constructions.
You can reduce both structural and cognitive complexity without changing behavior by instantiating the runtime helpers once per ion instance instead of rebuilding _Specs, _Species, _Conserve, and _Flux in every hook.
That flattens the helper stack and makes the data flow clearer (“this instance has its runtime adapter”) while also avoiding repeated allocations.
1. Cache runtime adapters on the instance
In _init_kinetic_ion, build and cache the spec/runtime helpers:
class KineticIon(IndependentIntegration):
...
def _init_kinetic_ion(
self,
*,
Co=None,
temp=None,
valence=None,
solver: str = "rk4",
substeps: int = 1,
):
...
self.temp = braintools.init.param(temp, self.varshape, allow_none=False)
# NEW: cache specs + runtime helpers
self._specs = _Specs.for_type(type(self))
self._species = _Species(self, self._specs)
self._conserve = _Conserve(self, self._specs, self._species)
self._flux = _Flux(self, self._specs, self._species)This keeps all existing validation/caching behavior but turns _Specs into a per-type “compile-time” artifact and _Species/_Conserve/_Flux into per-instance “runtime adapters” with explicit ownership.
2. Reuse cached adapters in hooks
Then simplify the hooks to use the cached adapters instead of re-creating them:
class KineticIon(IndependentIntegration):
...
def species_values(self):
"""Return the current full visible species view."""
return self._conserve.resolve()
def _ion_init_state_hook(self, V, batch_size: int = None):
"""Initialize runtime species and project algebraic species."""
self._species.init(batch_size=batch_size)
self._conserve.writeback(V)
def _ion_reset_state_hook(self, V, batch_size: int = None):
"""Reset runtime species and project algebraic species."""
self._species.reset(batch_size=batch_size)
self._conserve.writeback(V)
def _ion_compute_derivative_hook(self, V):
"""Resolve algebraic species and write diffeq derivatives."""
total_current = None
if type(self).uses_total_current:
if not hasattr(self, "_cached_total_current"):
raise RuntimeError(
f"{type(self).__name__} requires a cached total current before compute_derivative()."
)
total_current = self._cached_total_current
species_values = self._conserve.resolve(V)
self._flux.compute(V, species_values, total_current=total_current)
def _ion_post_integral_hook(self, V):
"""Refresh cached algebraic species after one integration step."""
self._conserve.writeback(V)Key benefits:
- Eliminates repeated
_Specs.for_type(type(self))and_Species/_Conserve/_Fluxconstructions. - Makes the runtime pipeline more discoverable (
self._species,self._conserve,self._fluxare clearly the adapters). - Preserves all existing semantics and public APIs.
This is a small, targeted refactor that directly addresses point (1) from the review (“deep helper stack and indirection”) without removing features or changing the declarative DSL.
one is graph node, another is nodetree
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
KineticIontemplate for NMODL-style ion species, includingFactor,Species,Reaction,Source, andConservebuilding blocks.Cell(update_policy=...)with a newfamily_phasedmode that advances voltage first, then channels, then ions, while keepinglegacyas the default behavior.start total current for current-driven ion dynamics.
braincell/ion/Cerebellum_ion_import_plan.mdto match the new kinetic-ion direction and the mixed ion/channel updatebehavior.
How Has This Been Tested
python3 -m pytest braincell/_base_ion_test.py braincell/_compute/runtime_test.py braincell/_multi_compartment/cell_test.py braincell/_multi_compartment/probes_test.py braincell/ion/_base_test.py104 passed in 9.08sKineticIoninit/reset behavior, conservation resolution, derivative computation, post-integral projection, andpack_info()behavior.Cell.update_policydefault/validation behavior plus family-phased channel-before-ion execution and cached-current usagefor current-driven ions.
Types of changes
Summary by Sourcery
Introduce declarative kinetic ion support and an explicit cell update policy for family-phased ion/channel scheduling.
New Features:
Factor,Species,Reaction,Source,Conserve) and aKineticIonbase class for NMODL-style reaction-network ions.Cellto select anupdate_policy, including a newfamily_phasedmode that orders voltage, channels, and ions explicitly while preserving the existinglegacybehavior.MechanismProbeand ion/mechanism probes to read plain attributes in addition to runtime state fields.Bug Fixes:
Enhancements:
Tests:
pack_info()behavior.Cell.update_policy, family-phased execution ordering, and cached-current usage for current-driven ions.