feat: integrate CV-space runtime, trainable channels, and optimization workflows

.gitignore

@@ -160,6 +160,7 @@ target/
 # Jupyter Notebook
 .ipynb_checkpoints
 examples/neuron_compare/channel_no_conc/artifacts/
+examples/experimental/**/artifacts/
 
 # IPython
 profile_default/

AGENTS.md

@@ -11,11 +11,12 @@ Biologically detailed brain cell modeling in BrainX.
 5. Every correction: reflect on the mistake, plan to avoid repeating it.
 6. All updates must be happened on the worktree branch, not main.
 7. Use `brainstate.random` instead of `jax.random` directly for all random number generation.
-8. **All prose lives under `docs/`; never leave a stray `.md` inside `braincell/`.** Two homes, each with a filename rule:
+8. **All durable prose lives under `docs/`; never leave a stray `.md` inside `braincell/`.** Two homes, each with a filename rule:
     - `docs/specs/YYYY-MM-DD-<slug>.md` — the spec and plan for one change, written *before* implementation. The date prefix is the creation date, so the directory reads chronologically.
     - `docs/design/<topic>.md` — durable design notes, invariants, and architecture maps that outlive any single change. Group a multi-document topic in its own subdirectory (`docs/design/network/`).
 
     Name a file for what it documents, not where the code happens to sit: `io-swc-reader-invariants.md`, never `README.md` or `notes.md`. Give it an `# H1` that matches.
+    - Explicit exception: an unstable experiment directory under `examples/experimental/` may contain one local `README.md` that maps the files, status, commands, and historical context for that directory. It must not define public API or duplicate durable design/results prose; link to `docs/` and generated artifacts for those.
 9. Tests should >90% coverage, but focus on meaningful tests that cover edge cases and critical paths, not just trivial lines.
 10. Co-locate tests with the code under test: each module `foo.py` has its tests in a sibling `foo_test.py` (suffix style — never a separate `tests/` directory, never the `test_*.py` prefix). See [Testing](#testing) for the full rule.
 11. **Never drive a model with a bare Python `for`/`while` loop when it runs repeatedly.** Python loops execute op-by-op (dispatch overhead, no fusion) and trace fresh each step; the `brainstate.transform` primitives lower the whole loop into one compiled XLA program, tracing the body only once. Pick by shape of the work:

CONTRIBUTING.md

@@ -31,7 +31,7 @@ You can help by:
 
 If you are planning a larger change, open an issue first so the scope and API impact can be discussed before implementation.
 
-`TODO.md` is the project design document. It tracks the architectural intent and the current implementation state of every subsystem, using `[x]` shipped / `[~]` partial / `[ ]` planned markers. Read the relevant section before starting substantial work, and update it when your change moves a subsystem forward.
+`docs/design/TODO.md` is the project design document. It tracks the architectural intent and the current implementation state of every subsystem, using `[x]` shipped / `[~]` partial / `[ ]` planned markers. Read the relevant section before starting substantial work, and update it when your change moves a subsystem forward.
 
 ## Development setup
 
@@ -176,7 +176,7 @@ Before opening a pull request:
 1. make sure your branch is based on the latest target branch state
 2. run `pre-commit run --all-files`
 3. run the relevant tests locally
-4. update documentation, examples, `TODO.md`, or `changelog.md` if your change is user-facing
+4. update documentation, examples, `docs/design/TODO.md`, or `changelog.md` if your change is user-facing
 5. review the pull request template in `.github/PULL_REQUEST_TEMPLATE.md`
 
 When opening a pull request, include:

MANIFEST.in

@@ -29,7 +29,6 @@ exclude CLAUDE.md
 exclude CODE_OF_CONDUCT.md
 exclude CONTRIBUTING.md
 exclude SECURITY.md
-exclude TODO.md
 exclude changelog.md
 exclude conftest.py
 exclude requirements-dev.txt

SECURITY.md

@@ -82,7 +82,7 @@ documented API allows are in scope. In particular:
 - **Numerical disagreement is not a vulnerability.** A simulation that diverges from
   NEURON, an integrator that loses accuracy on a stiff channel, or a solver that fails
   to converge is a correctness bug — please open a normal GitHub issue.
-- **Non-public API surface.** Per section 6 of `TODO.md`, only the documented public
+- **Non-public API surface.** Per section 6 of `docs/design/TODO.md`, only the documented public
   API is a stable surface; anything else is internal and may change without
   deprecation. Reports that depend on reaching into private modules will generally be
   treated as ordinary bugs.

braincell/__init__.py

@@ -27,7 +27,7 @@
     hidden_state,
     state_grouping,
 )
-from . import quad, mech, channel, synapse, ion, filter, morph
+from . import quad, mech, channel, synapse, ion, filter, morph, trainable
 from ._base_channel import (
     Channel,
     IonInfo,
@@ -58,6 +58,7 @@
     Cell,
     CellView,
     ChannelView,
+    ClampView,
     IonView,
     MultiCompartment,
     RunResult,
@@ -121,6 +122,7 @@
     "CellView",
     "Channel",
     "ChannelView",
+    "ClampView",
     "CompositeByTypePolicy",
     "ConnectionView",
     "CurrentClamp",
@@ -181,5 +183,6 @@
     "state",
     "state_grouping",
     "synapse",
+    "trainable",
     "vis",
 ]

braincell/__init___test.py

@@ -42,6 +42,7 @@
     "network",
     "quad",
     "synapse",
+    "trainable",
     "vis",
 )
 

braincell/_base_channel.py

@@ -33,6 +33,7 @@
 import brainunit as u
 
 from braincell._typing import ArrayLike, Size
+from braincell._parameter_schema import RuntimeParameterState
 from ._misc import TreeNode
 from .mech import NoEventInput, ParameterSpec, StateSpec
 from .quad.protocol import DiffEqModule, DiffEqSingleState, IndependentIntegration
@@ -336,6 +337,20 @@ def current(self, V, *args):
 
     __module__ = 'braincell'
 
+    parameters: Mapping[str, ParameterSpec] = {}
+    states: Mapping[str, StateSpec] = {}
+
+    def __getattribute__(self, name: str):
+        value = super().__getattribute__(name)
+        return value.dense_value(masked=True) if isinstance(value, RuntimeParameterState) else value
+
+    def __setattr__(self, name: str, value) -> None:
+        current = vars(self).get(name)
+        if isinstance(current, RuntimeParameterState) and not isinstance(value, RuntimeParameterState):
+            current.value = value
+            return
+        super().__setattr__(name, value)
+
 
 class Synapse(IonChannel):
     """Base class for vectorized runtime point-synapse mechanisms.

braincell/_base_ion.py

@@ -181,7 +181,7 @@ def _mask_inactive_current(current, point_mask):
         Current density returned by a channel. May be a
         :class:`brainunit.Quantity`.
     point_mask : array-like of bool
-        Boolean mask whose ``True`` entries mark active runtime points.
+        Legacy-named boolean mask whose ``True`` entries mark active density CV rows.
 
     Returns
     -------
@@ -191,9 +191,9 @@ def _mask_inactive_current(current, point_mask):
 
     Notes
     -----
-    Dense runtime layouts store full point-shaped state and use masks to
-    disable points outside the painted region. This helper keeps inactive
-    points from contributing to ion-current totals.
+    Dense density layouts store full CV-shaped state and use masks to disable
+    CVs outside the painted region. This helper keeps inactive rows from
+    contributing to ion-current totals.
     """
     return _where_active(current, point_mask, 0.0)
 
@@ -207,7 +207,7 @@ def _safe_inactive_voltage(V, point_mask):
         Membrane potential passed to a channel. May be a
         :class:`brainunit.Quantity`.
     point_mask : array-like of bool
-        Boolean mask whose ``True`` entries mark active runtime points.
+        Legacy-named boolean mask whose ``True`` entries mark active density CV rows.
 
     Returns
     -------

braincell/_base_neuron_test.py

@@ -18,7 +18,9 @@
 import unittest
 
 import brainunit as u
+import jax
 import jax.numpy as jnp
+import numpy as np
 
 from braincell._base_neuron import HHTypedNeuron
 
@@ -73,6 +75,25 @@ def test_single_compartment_inherits_get_spike(self) -> None:
         spk = sc.get_spike(jnp.array([-10.0]) * u.mV, jnp.array([10.0]) * u.mV)
         self.assertGreater(float(spk[0]), 0.0)
 
+    def test_default_crossing_surrogate_gradient_and_finite_support(self) -> None:
+        from braincell._single_compartment.base import SingleCompartment
+
+        sc = SingleCompartment(size=1, V_th=0.0 * u.mV)
+
+        def crossing(last_v_mantissa, next_v_mantissa):
+            last_v = jnp.asarray([last_v_mantissa]) * u.mV
+            next_v = jnp.asarray([next_v_mantissa]) * u.mV
+            return sc.get_spike(last_v, next_v)[0]
+
+        self.assertEqual(float(crossing(-10.0, 10.0)), 1.0)
+        grad_last, grad_next = jax.grad(crossing, argnums=(0, 1))(-10.0, 10.0)
+        np.testing.assert_allclose(grad_last, -0.0075, rtol=1e-6, atol=1e-8)
+        np.testing.assert_allclose(grad_next, 0.0075, rtol=1e-6, atol=1e-8)
+
+        far_grad_last, far_grad_next = jax.grad(crossing, argnums=(0, 1))(-40.0, 40.0)
+        np.testing.assert_allclose(far_grad_last, 0.0, atol=0.0)
+        np.testing.assert_allclose(far_grad_next, 0.0, atol=0.0)
+
 
 if __name__ == "__main__":
     unittest.main()

braincell/_compute/__init___test.py

@@ -470,12 +470,13 @@ def test_graph_matches_the_declared_layering(self) -> None:
         # edge — not merely an illegal one — surfaces in review.
         expected = {
             "__init__": set(),
-            "bindings": {"ions", "layouts"},
+            "bindings": {"ions", "layouts", "parameters"},
             "bridge": set(),
             "ions": {"layouts"},
-            "layouts": set(),
+            "layouts": {"parameters"},
+            "parameters": set(),
             "scheduling": set(),
-            "state": {"bindings", "bridge", "layouts"},
+            "state": {"bindings", "bridge", "layouts", "parameters"},
             # ``table`` builds its rows by matching mechanisms against
             # layout signatures, so it reads ``layouts`` directly rather
             # than through ``state``. ``layouts`` is a leaf, so the edge