flycns 0.6.0__py3-none-any.whl

flycns/__init__.py

@@ -0,0 +1,7 @@
+"""flycns: fly connectome releases compiled for simulation, two compound eyes, and the whole CNS in motion.
+
+The display version (``X.XX.XXX``) lives in the ``VERSION`` file at the repository root; ``__version__`` is its
+semantic form, the one the package index and ``pyproject.toml`` carry.
+"""
+
+__version__ = "0.6.0"

flycns/compiled.py

@@ -0,0 +1,150 @@
+"""The compiled directory: the one thing both simulators read.
+
+A compiled connectome is a directory holding ``manifest.json`` and one little-endian binary file per array. The
+manifest lists every array with its dtype, shape and SHA-256, plus the release it came from, the hashes of the source
+tables, the string tables the index arrays point into, and the counts the compiler measured. The reader refuses a
+directory whose arrays do not match their recorded hashes: a simulation must never run on a graph that is not the one
+that was compiled.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+SCHEMA = "flycns.compiled/1"
+
+#: dtypes the format allows, by name; each is written little-endian.
+DTYPES = {
+    "int8": np.int8,
+    "uint8": np.uint8,
+    "int16": np.int16,
+    "uint16": np.uint16,
+    "int32": np.int32,
+    "uint32": np.uint32,
+    "int64": np.int64,
+    "float32": np.float32,
+}
+
+
+class CompiledError(ValueError):
+    """A compiled directory that cannot be trusted: a missing file, a wrong shape, or a hash that differs."""
+
+
+def sha256_bytes(data: bytes) -> str:
+    return hashlib.sha256(data).hexdigest()
+
+
+def sha256_file(path: Path, chunk: int = 1 << 24) -> str:
+    digest = hashlib.sha256()
+    with open(path, "rb") as handle:
+        while block := handle.read(chunk):
+            digest.update(block)
+    return digest.hexdigest()
+
+
+def _dtype_name(array: np.ndarray) -> str:
+    for name, dtype in DTYPES.items():
+        if array.dtype == np.dtype(dtype):
+            return name
+    raise CompiledError(f"dtype {array.dtype} is not allowed in a compiled directory")
+
+
+def write_compiled(directory: Path, arrays: dict[str, np.ndarray], meta: dict[str, Any],
+                   schema: str = SCHEMA) -> dict[str, Any]:
+    """Write ``arrays`` and a manifest into ``directory``; return the manifest.
+
+    ``meta`` carries the release, sources, counts and string tables; it is stored as given under the manifest's
+    ``release``, ``sources``, ``counts`` and ``strings`` keys. Array files are written in name order so two
+    compilations of the same inputs produce the same bytes.
+    """
+    directory = Path(directory)
+    directory.mkdir(parents=True, exist_ok=True)
+    entries = []
+    for name in sorted(arrays):
+        array = np.ascontiguousarray(arrays[name])
+        dtype = _dtype_name(array)
+        data = array.astype(array.dtype.newbyteorder("<"), copy=False).tobytes(order="C")
+        file_name = f"{name}.bin"
+        (directory / file_name).write_bytes(data)
+        entries.append(
+            {"name": name, "dtype": dtype, "shape": list(array.shape), "file": file_name, "bytes": len(data),
+             "sha256": sha256_bytes(data)}
+        )
+    manifest = {
+        "schema": schema,
+        "release": meta.get("release", {}),
+        "sources": meta.get("sources", []),
+        "counts": meta.get("counts", {}),
+        "strings": meta.get("strings", {}),
+        "arrays": entries,
+    }
+    # LF on every system, so the same inputs give the same bytes on Windows and Linux
+    (directory / "manifest.json").write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8",
+                                             newline="\n")
+    return manifest
+
+
+@dataclass
+class Compiled:
+    """A verified compiled connectome: its manifest, arrays and string tables."""
+
+    directory: Path
+    manifest: dict[str, Any]
+    arrays: dict[str, np.ndarray] = field(repr=False)
+
+    @property
+    def strings(self) -> dict[str, list[str]]:
+        return self.manifest["strings"]
+
+    @property
+    def counts(self) -> dict[str, Any]:
+        return self.manifest["counts"]
+
+    @property
+    def n_neurons(self) -> int:
+        return int(self.arrays["neuron_body_id"].shape[0])
+
+    @property
+    def n_edges(self) -> int:
+        return int(self.arrays["csr_indices"].shape[0])
+
+    def __getitem__(self, name: str) -> np.ndarray:
+        return self.arrays[name]
+
+
+def read_compiled(directory: Path, verify: bool = True, schema: str = SCHEMA) -> Compiled:
+    """Read a compiled directory, checking every array against the manifest.
+
+    With ``verify`` (the default) each file's SHA-256 is recomputed; any difference, a missing file or a size that
+    does not match the declared shape raises :class:`CompiledError`.
+    """
+    directory = Path(directory)
+    manifest_path = directory / "manifest.json"
+    if not manifest_path.is_file():
+        raise CompiledError(f"{directory} holds no manifest.json")
+    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
+    if manifest.get("schema") != schema:
+        raise CompiledError(f"schema {manifest.get('schema')!r} is not {schema!r}")
+    arrays: dict[str, np.ndarray] = {}
+    for entry in manifest["arrays"]:
+        path = directory / entry["file"]
+        if not path.is_file():
+            raise CompiledError(f"array {entry['name']} is missing its file {entry['file']}")
+        data = path.read_bytes()
+        if verify and sha256_bytes(data) != entry["sha256"]:
+            raise CompiledError(
+                f"array {entry['name']}: SHA-256 {sha256_bytes(data)} differs from the manifest's {entry['sha256']}"
+            )
+        dtype = np.dtype(DTYPES[entry["dtype"]]).newbyteorder("<")
+        shape = tuple(entry["shape"])
+        expected = int(np.prod(shape, dtype=np.int64)) * dtype.itemsize
+        if len(data) != expected:
+            raise CompiledError(f"array {entry['name']}: {len(data)} bytes where shape {shape} needs {expected}")
+        arrays[entry["name"]] = np.frombuffer(data, dtype=dtype).reshape(shape).astype(dtype.newbyteorder("="))
+    return Compiled(directory=directory, manifest=manifest, arrays=arrays)

flycns/data/flyvis-1.2.0-ensemble/LICENSE-flyvis.txt

@@ -0,0 +1,9 @@
+MIT License
+
+Copyright (c) 2023 Janne K. Lappalainen, Fabian D. Tschopp, Mason McGill, Jakob H. Macke, Srinivas C. Turaga
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

flycns/data/flyvis-1.2.0-ensemble/manifest.json

@@ -0,0 +1,248 @@
+{
+  "arrays": [
+    {
+      "bytes": 13000,
+      "dtype": "float32",
+      "file": "bias.bin",
+      "name": "bias",
+      "sha256": "aab655c3bb8801e3b9583b67a8299877a8278d0c23a03611f25ba46a927242ae",
+      "shape": [
+        50,
+        65
+      ]
+    },
+    {
+      "bytes": 4710,
+      "dtype": "int16",
+      "file": "group_du.bin",
+      "name": "group_du",
+      "sha256": "d317a8269e34bb7c24e5fe8b674203e5ef8821409f63842288527035d73c8083",
+      "shape": [
+        2355
+      ]
+    },
+    {
+      "bytes": 4710,
+      "dtype": "int16",
+      "file": "group_dv.bin",
+      "name": "group_dv",
+      "sha256": "fbb9e3bd3d5d77f0cc3942dac8cfdfb93eae90a9215a991046d0d38d4246accf",
+      "shape": [
+        2355
+      ]
+    },
+    {
+      "bytes": 9420,
+      "dtype": "float32",
+      "file": "group_n_syn.bin",
+      "name": "group_n_syn",
+      "sha256": "19ae8debe3345a1c545cabdbe8ea98e6687fc7723349d2a9f7beeb007f7ea746",
+      "shape": [
+        2355
+      ]
+    },
+    {
+      "bytes": 4710,
+      "dtype": "uint16",
+      "file": "group_pair.bin",
+      "name": "group_pair",
+      "sha256": "ca442065d4dd5bf0b15bd4e75576616e2de00dca0e702c1ca170263cdc77bc7f",
+      "shape": [
+        2355
+      ]
+    },
+    {
+      "bytes": 604,
+      "dtype": "int8",
+      "file": "pair_sign.bin",
+      "name": "pair_sign",
+      "sha256": "5725dc8fd49f11dfd62e51420feffb5eef4276c5bab8becc1d3fb476fa36c023",
+      "shape": [
+        604
+      ]
+    },
+    {
+      "bytes": 1208,
+      "dtype": "uint16",
+      "file": "pair_source.bin",
+      "name": "pair_source",
+      "sha256": "e4457b0d98807d82c371726f4d46aca1e0dab4d2d39c22f45563c04859dc0e59",
+      "shape": [
+        604
+      ]
+    },
+    {
+      "bytes": 1208,
+      "dtype": "uint16",
+      "file": "pair_target.bin",
+      "name": "pair_target",
+      "sha256": "205cf8d6bc7c7202e4bc173aeeeedb1a0d840eb751392d0a207df3a460bd64dd",
+      "shape": [
+        604
+      ]
+    },
+    {
+      "bytes": 120800,
+      "dtype": "float32",
+      "file": "strength.bin",
+      "name": "strength",
+      "sha256": "cd36b8a4982233534bae69084faf6719e1cd32fa5b3fa7dfdbfc31ee0d6972b8",
+      "shape": [
+        50,
+        604
+      ]
+    },
+    {
+      "bytes": 13000,
+      "dtype": "float32",
+      "file": "time_const_s.bin",
+      "name": "time_const_s",
+      "sha256": "9f35013972404b1a453a7621cb6609627c04f1a48cdc5b16621ac0c5a33486a1",
+      "shape": [
+        50,
+        65
+      ]
+    }
+  ],
+  "counts": {
+    "groups": 2355,
+    "models": 50,
+    "pairs": 604,
+    "types": 65
+  },
+  "release": {
+    "ensemble": "flow/0000",
+    "model": "flyvis",
+    "paper": "doi:10.1038/s41586-024-07939-3",
+    "version": "1.2.0"
+  },
+  "schema": "flycns.compiled/1",
+  "sources": [
+    {
+      "checkpoints_sha256": [
+        "d8a57a022aa18ca338599be713af9db29b89cde4b00098ffe42b78f0d583e46a",
+        "f351b025e25fc02e6ab9e2a80f6fd641fc62f5b3f33210e0a2318af5cc0614ae",
+        "34c70cdc875f9602053f31acf03443cc071982923a9407a367666086637ae866",
+        "b6764982671b5f0051e682360adc52f8ed6493e478fb9c0edbf66f1d256e1501",
+        "d3e5ffd015020dec45ccb87c92ceb8091a8f8e65a83a553a9dc3887dd1b6072d",
+        "237db8bd6b015aacd54a0719014ef5a146236a6b4e2ab769b81d314907753ffa",
+        "6fb45c35839d04f78073b5a8e67fabcf59131fbf340925459952a5ff1de3b4ea",
+        "cd1b53eef47db91d3529613a98d0879a739390dcfc53d6148a6953ff9fb13469",
+        "e54b1f0790f97ef791b7aeee056fb88a2bfc7ab4329dd9f7dca10300eb945897",
+        "37c2676a83587dd7642cbbd3a0d0be219cb999d587ac1830285bdc838f4891f6",
+        "ed8be3c396c7a549b2d8b85cd5b14527c43eec58b5ddaae054d18ab9a2b6a300",
+        "bad9dabe0e9432a1f305a4544a847a1fdf7da871021f6d69b1e52a3fa7fcda67",
+        "d0c816909660cae35811b1631b3ff8cbd40a3acc076b6b083f4c88abbfc19f60",
+        "bde19f62d1a9ce2fce20fce68b72cced1a082a5192c210f4a69c89708e1476d8",
+        "0bcac347539e36750a00adf71e7c1236a341c4b363cf6b84c0fb6fed349cf967",
+        "02d2cd384f7f8040582ab33bd52fa2c44edce50ca5be27897d0f105475d819cb",
+        "8fbfc7d39d48eca441e0e23a735dcbedb0a3e364b3a9de48cfa651d0e7c8e7bc",
+        "299ac9ab1a50e5dfc5f1716627a9b09f7c5abd8820436eda63f4a178c555e1e1",
+        "8c738269fdbecce5da3e4cc5d44718910be1adc61fc693959792eb1817960277",
+        "6fb18d56f4c014ccf7a53b6914df6fe93f8f594d28b056ebe056a919d70aa596",
+        "03061634c6c83a2b4560005d0e97f40103ef5767d0a831b5205256bc47919ef9",
+        "9285061e7e0af9b9e23e16a976f3ad158808a0a708c24590de1b85c23aea0727",
+        "dd864e23444915ad4bc4fab9d1c401567923c31388c174d3b3caf2b7874d4383",
+        "41a283e29e36ad7dbf2be0604138abd7887b87a1e7d168cf17d3f21103d5b88d",
+        "8400223db1a0702c2ec4eb863e47918fdcc5b998f6718c4f8b0f1a41264f386a",
+        "3c780db8bc7c1d5fa614e2475e4dc42b2ad15375b522f6e7ed8e153fe62847e4",
+        "6981fb1933705bfcb13b072f4c26520e3c3af5cb7e87652efbb6877ddf68b3e9",
+        "5f3723c0011f6d561ecb3861b2a2b8e13ea923011fd23a10f7c33ff0cbb6b6e7",
+        "b19c5c1e94760ec1479bea2adba122441c2c4c95e2ca174526a922d964395b64",
+        "2454e817485656d040fcee69eddc30e6ff7cd587ecd6a908ef3564e03ec95835",
+        "2b7a878d5f4328e2a26291522d2807e5b5a47410da16c15fa0fd55c195d464f7",
+        "1b6a145b53375ba5eea53c30aacd1ad916f756458acdce1c11dfbf733bb6c49a",
+        "e8b23729a4951ab3584d8d6fb7bf9b54e4a675ab05fa4afb3d49d9dde07ea858",
+        "4dead734972b429850aca796eb23874e16d449036c892a99cbae0b02856b1ef6",
+        "7145f3e4504557d0439ab95c79a1b300fcdef9b8889c65022bd5318db7097980",
+        "2c59918d9bd267330d9a8c2a9f42da1dc5335901a6c1c62fdcfdc10faa3ff737",
+        "052b7e31cebc6b60f73e37e166911dc6889795abbaa5f38717d5374ab1d0cc6d",
+        "93f25c9d4dc6a6f93d7a8d19fc372a9eeedc5455b388779de8152763337baeee",
+        "a336ca91868486678c4dfbf871ef8f1c009c12974df20289b8697a79f1db86e0",
+        "a280006d540ceb2d1e6e0cd25e912500235ea79e3c258cf3f93ea518459083d7",
+        "67a85cdcfb55d9e59fbe750a365c044d333d782ab89c0d44d75ffe9a736fac84",
+        "76759a8d7d714bf7886000e41eca56b8c02aecd2df7feb630f4aa98f645f6f32",
+        "d57aa0a374be4f63a2e8d955a73b15fbcc72b799fcb7906ff5cd6366b42a0f6d",
+        "9536964c530de66dd3356e04c47d4057a2764e25a40faf04a45c8e8083612d76",
+        "efc1850265ed305bdc2b6c645fe8a74faa2b2ea2a238d097148ca490b83fcd80",
+        "19323db328787e33b91ac0473046f0e361ef4dadf0394e3f389b5359b9fdedce",
+        "604b2cf4dddbd7fbaefad7b350903ace541acb1c412242705dccf246cd904750",
+        "f5f98dbc53a94606a07e21079268b57a26950be66fe2f879fb06647a35d870fb",
+        "781eae13aaef429af6b98f0e4a436974a41f399ed532aabfc2d7973d28545bfe",
+        "b0dd6e351eed6b3ac18941fb02760c97426fe71aba1c194b0bd24a09d8d5960f"
+      ],
+      "key": "flyvis-1.2.0-pretrained-flow-0000",
+      "license": "MIT (Copyright (c) 2023 Janne K. Lappalainen, Fabian D. Tschopp, Mason McGill, Jakob H. Macke, Srinivas C. Turaga)",
+      "url": "https://github.com/TuragaLab/flyvis (results_pretrained_models.zip, via flyvis download_pretrained)"
+    }
+  ],
+  "strings": {
+    "type": [
+      "R1",
+      "R2",
+      "R3",
+      "R4",
+      "R5",
+      "R6",
+      "R7",
+      "R8",
+      "L1",
+      "L2",
+      "L3",
+      "L4",
+      "L5",
+      "Lawf1",
+      "Lawf2",
+      "Am",
+      "C2",
+      "C3",
+      "CT1(Lo1)",
+      "CT1(M10)",
+      "Mi1",
+      "Mi2",
+      "Mi3",
+      "Mi4",
+      "Mi9",
+      "Mi10",
+      "Mi11",
+      "Mi12",
+      "Mi13",
+      "Mi14",
+      "Mi15",
+      "T1",
+      "T2",
+      "T2a",
+      "T3",
+      "T4a",
+      "T4b",
+      "T4c",
+      "T4d",
+      "T5a",
+      "T5b",
+      "T5c",
+      "T5d",
+      "Tm1",
+      "Tm2",
+      "Tm3",
+      "Tm4",
+      "Tm5Y",
+      "Tm5a",
+      "Tm5b",
+      "Tm5c",
+      "Tm9",
+      "Tm16",
+      "Tm20",
+      "Tm28",
+      "Tm30",
+      "TmY3",
+      "TmY4",
+      "TmY5a",
+      "TmY9",
+      "TmY10",
+      "TmY13",
+      "TmY14",
+      "TmY15",
+      "TmY18"
+    ]
+  }
+}

flycns/data/flyvis-1.2.0-ensemble/pair_sign.bin

@@ -0,0 +1 @@
+��������������������������������������



��������






































































����
�����


�������������

�������������������������������




























����������������������������������������������������������������������������������������



























































































































































































�����




















































����������

















��������������������������

flycns/dynamics/__init__.py

@@ -0,0 +1,18 @@
+"""Neuron dynamics: the published spiking model, the graded optic lobe, and their coupling."""
+
+from .graded import GradedNetwork, GradedReference, GradedTorch
+from .hybrid import HybridNetwork, HybridReference, HybridTorch, build_hybrid
+from .lif import (
+                  Drive,
+                  LIFParams,
+                  LIFReference,
+                  LIFTorch,
+                  Run,
+                  photoreceptor_drive,
+                  stabilised_weights,
+                  synaptic_weights,
+)
+
+__all__ = ["Drive", "GradedNetwork", "GradedReference", "GradedTorch", "HybridNetwork", "HybridReference",
+           "HybridTorch", "LIFParams", "LIFReference", "LIFTorch", "Run", "build_hybrid", "photoreceptor_drive",
+           "stabilised_weights", "synaptic_weights"]

flycns/dynamics/graded.py

@@ -0,0 +1,184 @@
+"""Graded (non-spiking) point neurons, as flyvis computes them (Lappalainen et al., Nature 2024).
+
+``PPNeuronIGRSynapses`` in flyvis 1.2.0: passive point neurons with instantaneous graded release,
+
+    tau_eff dV/dt = -V + b + sum_j w_ij max(V_j, 0) + x,    tau_eff = max(tau, dt),
+
+integrated by forward Euler. ``GradedReference`` (NumPy, float64) and ``GradedTorch`` (PyTorch, float32) take the
+same network and give the same run; both are held to flyvis's own recording of its network 000 (see
+``docs/design/features/graded/design.md``).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+
+
+@dataclass
+class GradedNetwork:
+    """A network of graded neurons.
+
+    Light enters through ``input_neuron``: each of those neurons receives the intensity of the column
+    ``input_column`` names (flyvis adds the same intensity to R1 to R8 of a column; on MaleCNS a column holds as many
+    photoreceptors as the release reconstructed there, plus stand-ins).
+    """
+
+    bias: np.ndarray            # (n,) resting potential of each neuron
+    time_const_s: np.ndarray    # (n,) time constant of each neuron, seconds
+    source: np.ndarray          # (e,) presynaptic neuron of each connection
+    target: np.ndarray          # (e,) postsynaptic neuron of each connection
+    weight: np.ndarray          # (e,) sign x synapse count x unitary strength
+    input_neuron: np.ndarray    # (k,) the neurons that receive light
+    input_column: np.ndarray    # (k,) the column each of them sees
+    n_columns: int
+
+    @classmethod
+    def from_input_index(cls, bias, time_const_s, source, target, weight, input_index) -> GradedNetwork:
+        """flyvis's layout: ``input_index`` has one row per input type and one column per eye column."""
+        input_index = np.asarray(input_index)
+        types, columns = input_index.shape
+        return cls(bias=bias, time_const_s=time_const_s, source=source, target=target, weight=weight,
+                   input_neuron=input_index.reshape(-1), input_column=np.tile(np.arange(columns), types),
+                   n_columns=columns)
+
+    @property
+    def n(self) -> int:
+        return len(self.bias)
+
+    def column_current(self, intensity: np.ndarray) -> np.ndarray:
+        """Per-neuron input for one frame of per-column intensities (zero off the input neurons)."""
+        current = np.zeros(self.n)
+        np.add.at(current, np.asarray(self.input_neuron, dtype=np.int64),
+                  np.asarray(intensity, dtype=np.float64)[np.asarray(self.input_column, dtype=np.int64)])
+        return current
+
+
+class GradedReference:
+    """The NumPy engine: float64 state; each neuron's input summed with ``bincount`` in connection order."""
+
+    def __init__(self, network: GradedNetwork, dt_s: float):
+        self.net = network
+        self.dt = float(dt_s)
+        self.bias = np.asarray(network.bias, dtype=np.float64)
+        self.rate = 1.0 / np.maximum(np.asarray(network.time_const_s, dtype=np.float64), self.dt)
+        self.source = np.asarray(network.source, dtype=np.int64)
+        self.target = np.asarray(network.target, dtype=np.int64)
+        self.weight = np.asarray(network.weight, dtype=np.float64)
+
+    def column_current(self, intensity: np.ndarray) -> np.ndarray:
+        return self.net.column_current(intensity)
+
+    def synaptic_input(self, v: np.ndarray) -> np.ndarray:
+        release = np.maximum(v[self.source], 0.0)
+        return np.bincount(self.target, weights=self.weight * release, minlength=self.net.n)
+
+    def step(self, v: np.ndarray, current: np.ndarray) -> np.ndarray:
+        """One Euler step with a per-neuron input ``current``."""
+        velocity = self.rate * (-v + self.bias + self.synaptic_input(v) + current)
+        return v + velocity * self.dt
+
+    def steady_state(self, t_pre_s: float, grey: float = 0.5, initial: np.ndarray | None = None) -> np.ndarray:
+        """The state after ``t_pre_s`` of uniform intensity ``grey``, from the resting potentials unless given."""
+        v = self.bias.copy() if initial is None else np.asarray(initial, dtype=np.float64).copy()
+        current = self.net.column_current(np.full(self.net.n_columns, grey))
+        for _ in range(int(t_pre_s / self.dt)):
+            v = self.step(v, current)
+        return v
+
+    def run(self, intensity: np.ndarray, initial: np.ndarray | None = None,
+            record: np.ndarray | None = None) -> tuple[np.ndarray, np.ndarray]:
+        """Step through frames of per-column intensity (frames x columns); returns (final state, activity), the
+        activity after each step for the neurons in ``record`` (all when omitted), as float32."""
+        v = self.bias.copy() if initial is None else np.asarray(initial, dtype=np.float64).copy()
+        record = np.arange(self.net.n) if record is None else np.asarray(record)
+        out = np.zeros((len(intensity), len(record)), dtype=np.float32)
+        for k, frame in enumerate(np.asarray(intensity)):
+            v = self.step(v, self.net.column_current(frame))
+            out[k] = v[record]
+        return v, out
+
+
+class GradedTorch:
+    """The same dynamics in PyTorch (float32), for the GPU."""
+
+    def __init__(self, network: GradedNetwork, dt_s: float, device: str | None = None):
+        import torch
+
+        self.torch = torch
+        self.net = network
+        self.dt = float(dt_s)
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        as_f32 = lambda a: torch.as_tensor(np.asarray(a, dtype=np.float32), device=self.device)  # noqa: E731
+        as_i64 = lambda a: torch.as_tensor(np.asarray(a, dtype=np.int64), device=self.device)    # noqa: E731
+        self.bias = as_f32(network.bias)
+        self.rate = 1.0 / torch.clamp(as_f32(network.time_const_s), min=self.dt)
+        self.source = as_i64(network.source)
+        self.target = as_i64(network.target)
+        self.weight = as_f32(network.weight)
+        self.input_neuron = as_i64(network.input_neuron)
+        self.input_column = as_i64(network.input_column)
+
+    def column_current(self, intensity):
+        torch = self.torch
+        current = torch.zeros(self.net.n, device=self.device)
+        frame = torch.as_tensor(np.asarray(intensity, dtype=np.float32), device=self.device)
+        return current.index_add_(0, self.input_neuron, frame[self.input_column])
+
+    def step(self, v, current):
+        synaptic = self.torch.zeros_like(v).index_add_(0, self.target, self.weight * self.torch.relu(v[self.source]))
+        return v + self.rate * (-v + self.bias + synaptic + current) * self.dt
+
+    def steady_state(self, t_pre_s: float, grey: float = 0.5, initial=None):
+        torch = self.torch
+        v = self.bias.clone() if initial is None else torch.as_tensor(np.asarray(initial, dtype=np.float32),
+                                                                       device=self.device).clone()
+        current = self.column_current(np.full(self.net.n_columns, grey))
+        for _ in range(int(t_pre_s / self.dt)):
+            v = self.step(v, current)
+        return v
+
+    def run(self, intensity: np.ndarray, initial=None, record: np.ndarray | None = None):
+        torch = self.torch
+        v = self.bias.clone() if initial is None else torch.as_tensor(np.asarray(initial, dtype=np.float32),
+                                                                       device=self.device).clone()
+        record_t = torch.as_tensor(np.arange(self.net.n) if record is None else np.asarray(record),
+                                   device=self.device)
+        out = torch.zeros((len(intensity), len(record_t)), device=self.device)
+        for k, frame in enumerate(np.asarray(intensity)):
+            v = self.step(v, self.column_current(frame))
+            out[k] = v[record_t]
+        return v.cpu().numpy(), out.cpu().numpy()
+
+    def run_batch(self, intensity: np.ndarray, initial, record: np.ndarray) -> np.ndarray:
+        """Several stimuli at once: ``intensity`` is (stimuli, frames, columns), every stimulus starts from
+        ``initial``; returns the recorded neurons after each step, (stimuli, frames, recorded), as float32. Each
+        stimulus is the same computation as ``run``; batching only shares the kernel launches."""
+        torch = self.torch
+        stimuli = torch.as_tensor(np.asarray(intensity, dtype=np.float32), device=self.device)
+        batch, frames, _ = stimuli.shape
+        matrix = self._matrix()
+        v = torch.as_tensor(np.asarray(initial, dtype=np.float32), device=self.device)[:, None].repeat(1, batch)
+        record_t = torch.as_tensor(np.asarray(record, dtype=np.int64), device=self.device)
+        out = torch.zeros((batch, frames, len(record_t)), device=self.device)
+        rate, bias = self.rate[:, None], self.bias[:, None]
+        for k in range(frames):                                   # state is (neurons, stimuli)
+            current = torch.zeros_like(v).index_add_(0, self.input_neuron, stimuli[:, k, self.input_column].T)
+            synaptic = matrix @ torch.relu(v)
+            v = v + rate * (-v + bias + synaptic + current) * self.dt
+            out[:, k] = v[record_t].T
+        return out.cpu().numpy()
+
+    def _matrix(self):
+        """The weights as a sparse (target x source) CSR matrix, built once: batches multiply by it instead of
+        materialising one release per connection and stimulus."""
+        if getattr(self, "_csr", None) is None:
+            torch = self.torch
+            order = torch.argsort(self.target * self.net.n + self.source)
+            rows = torch.bincount(self.target[order], minlength=self.net.n)
+            crow = torch.zeros(self.net.n + 1, dtype=torch.int64, device=self.device)
+            crow[1:] = torch.cumsum(rows, 0)
+            self._csr = torch.sparse_csr_tensor(crow, self.source[order], self.weight[order],
+                                                size=(self.net.n, self.net.n))
+        return self._csr