Spaces:
Paused
Paused
| """Experimental multi-compartment neuron variant with separable NMDA plateau axis.""" | |
| from __future__ import annotations | |
| from dataclasses import asdict, dataclass, field | |
| from typing import Any | |
| from .biophysical_contracts import ( | |
| DendriticNonlinearity, | |
| ModelProvenance, | |
| ModelTimescale, | |
| NeuronStructure, | |
| ) | |
| class CompartmentConfig: | |
| enabled: bool = False | |
| structure: NeuronStructure = NeuronStructure.MULTI_COMPARTMENT | |
| compartments: int = 3 | |
| dt_ms: float = 1.0 | |
| soma_tau_ms: float = 20.0 | |
| dendrite_tau_ms: float = 40.0 | |
| coupling: float = 0.08 | |
| resting_potential: float = -65.0 | |
| spike_threshold: float = -50.0 | |
| reset_potential: float = -65.0 | |
| dendritic_nonlinearity: DendriticNonlinearity = DendriticNonlinearity.LINEAR | |
| nmda_plateau_threshold: float = -35.0 | |
| nmda_plateau_gain: float = 4.0 | |
| enable_nmda_plateau: bool = False | |
| def __post_init__(self) -> None: | |
| if self.structure is not NeuronStructure.MULTI_COMPARTMENT: | |
| raise ValueError("CompartmentNeuron requires MULTI_COMPARTMENT structure") | |
| if self.compartments < 2: | |
| raise ValueError("compartments must be >= 2") | |
| if self.dt_ms <= 0.0: | |
| raise ValueError("dt_ms must be > 0") | |
| class CompartmentNeuron: | |
| neuron_id: int | |
| config: CompartmentConfig = CompartmentConfig() | |
| soma_v: float = -65.0 | |
| spike_counter: int = 0 | |
| last_spike_tick: int = -1 | |
| dendrite_v: list[float] = field(default_factory=list[float]) | |
| def __post_init__(self) -> None: | |
| expected = self.config.compartments - 1 | |
| if not self.dendrite_v: | |
| self.dendrite_v = [self.config.resting_potential] * expected | |
| elif len(self.dendrite_v) != expected: | |
| raise ValueError("dendrite_v must match configured compartments") | |
| def provenance_tag(self) -> ModelProvenance: | |
| payload = asdict(self.config) | |
| payload["structure"] = self.config.structure.value | |
| payload["dendritic_nonlinearity"] = self.config.dendritic_nonlinearity.value | |
| return ModelProvenance( | |
| model_id="multi-compartment-v1", | |
| version="mhrn-experimental-1", | |
| timescale=ModelTimescale.TICK, | |
| enabled=self.config.enabled, | |
| parameters=payload, | |
| ) | |
| def step( | |
| self, | |
| soma_current: float, | |
| dendritic_currents: tuple[float, ...], | |
| tick: int, | |
| ) -> bool: | |
| if not self.config.enabled: | |
| return False | |
| if len(dendritic_currents) != len(self.dendrite_v): | |
| raise ValueError("dendritic_currents must match configured dendrites") | |
| dt = self.config.dt_ms | |
| old = list(self.dendrite_v) | |
| for index, current in enumerate(dendritic_currents): | |
| v = old[index] | |
| neighbour = self.soma_v if index == 0 else old[index - 1] | |
| coupling = self.config.coupling * (neighbour - v) | |
| plateau = 0.0 | |
| if ( | |
| self.config.enable_nmda_plateau | |
| and self.config.dendritic_nonlinearity | |
| is DendriticNonlinearity.NMDA_PLATEAU | |
| and v >= self.config.nmda_plateau_threshold | |
| ): | |
| plateau = self.config.nmda_plateau_gain | |
| dv = ( | |
| (self.config.resting_potential - v) + current + coupling + plateau | |
| ) * (dt / self.config.dendrite_tau_ms) | |
| self.dendrite_v[index] = v + dv | |
| dendritic_drive = sum( | |
| self.config.coupling * (value - self.soma_v) for value in self.dendrite_v | |
| ) | |
| self.soma_v += ( | |
| (self.config.resting_potential - self.soma_v) | |
| + soma_current | |
| + dendritic_drive | |
| ) * (dt / self.config.soma_tau_ms) | |
| spiked = self.soma_v >= self.config.spike_threshold | |
| if spiked: | |
| self.soma_v = self.config.reset_potential | |
| self.spike_counter += 1 | |
| self.last_spike_tick = tick | |
| return spiked | |
| def to_dict(self) -> dict[str, Any]: | |
| return { | |
| "schema_version": 1, | |
| "neuron_id": self.neuron_id, | |
| "soma_v": self.soma_v, | |
| "dendrite_v": list(self.dendrite_v), | |
| "spike_counter": self.spike_counter, | |
| "last_spike_tick": self.last_spike_tick, | |
| "provenance": self.provenance_tag.to_dict(), | |
| } | |
| __all__ = ["CompartmentConfig", "CompartmentNeuron"] | |