pyfly-lightning 0.1.0__py3-none-any.whl

pyfly_lightning/__init__.py

@@ -0,0 +1,29 @@
+"""pyfly-lightning: a frozen-connectome substrate with Lightning-style training.
+
+The distribution is `pyfly-lightning` and the import is `pyfly_lightning`, following
+the `pytorch-lightning` / `pytorch_lightning` convention. A bare `pyfly` import name
+would collide with the unrelated `pyfly` package on PyPI (a load-testing framework),
+and two distributions installing the same top-level directory get merged in
+site-packages: one __init__.py wins and uninstalling either removes both.
+
+    import pyfly_lightning as pfl
+
+The substrate is a real Drosophila connectome (MaleCNS v1.0 by default).
+The only trainable parts are the encode/decode adapters around declared ports.
+See PLAN.md for the design, and docs/census-wm-ports.md for why CX is the
+default working-memory port.
+"""
+from pyfly_lightning.api import FitResult, Fly, Port, fit, load
+from pyfly_lightning.auto import Plan, plan
+from pyfly_lightning.data.download import cache_dir, fetch
+from pyfly_lightning.data.graph import SparseConnectome, load_malecns
+from pyfly_lightning.data.registry import DATASETS, describe
+from pyfly_lightning.data.schema import NeuronTable
+from pyfly_lightning.envs import CueDelayChoice
+from pyfly_lightning.freeze import FreezeState, FrozenError
+
+__version__ = "0.1.0"
+__all__ = ["DATASETS", "describe", "fetch", "cache_dir",
+           "NeuronTable", "SparseConnectome", "load_malecns",
+           "FreezeState", "FrozenError", "CueDelayChoice",
+           "Fly", "Port", "load", "fit", "FitResult", "Plan", "plan"]

pyfly_lightning/__main__.py

@@ -0,0 +1,5 @@
+"""`python -m pyfly_lightning` runs the same entry point as the `pyfly-lightning` command."""
+from pyfly_lightning.cli import main
+
+if __name__ == "__main__":
+    raise SystemExit(main())

pyfly_lightning/api.py

@@ -0,0 +1,322 @@
+"""The documented user-facing surface (PLAN.md section 4).
+
+    import pyfly_lightning as pfl
+
+    fly = pfl.load("malecns10")
+    fly.io(inputs=["cb_sensory"], outputs=["vnc_motor"])
+    sub = fly.slice(budget=4096, include_pools=("CX",))
+    fly.port(sub, "working_memory", selector={"class": "CX"}, read=True, write=True)
+    fly.gate(sub, plastic=True)
+    bus = fly.reward_bus(sub, selector={"class": "DAN"})
+    model = fly.model(sub, in_dim=784, n_classes=10)
+
+Thin wrappers over `Substrate`, `CreditChannel` and `FlyModule`. Kept deliberately
+thin: the point is that the four primitives the design calls for -- declare the I/O
+boundary, slice, port, gate -- are addressable by name, not that this layer does
+anything the core cannot.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import numpy as np
+
+from pyfly_lightning.data.graph import load_connectome
+from pyfly_lightning.data.schema import NeuronTable
+from pyfly_lightning.model.substrate import Substrate
+
+
+@dataclass
+class Port:
+    """A named read/write handle on a set of substrate nodes."""
+    name: str
+    positions: np.ndarray
+    body_ids: np.ndarray
+    read: bool = True
+    write: bool = True
+    protected: bool = True
+
+    def __len__(self) -> int:
+        return int(self.positions.size)
+
+    def __repr__(self) -> str:
+        mode = ("rw" if self.read and self.write else
+                "r" if self.read else "w" if self.write else "-")
+        return (f"<Port {self.name!r} n={len(self)} mode={mode} "
+                f"protected={self.protected}>")
+
+
+@dataclass
+class Fly:
+    """A dataset plus the primitives to build a model on it."""
+
+    dataset: str = "malecns10"
+    _table: object = None
+    _connectome: object = None
+    inputs: tuple = ()
+    outputs: tuple = ()
+    ports: dict = field(default_factory=dict)
+
+    # ---------- lazy data ----------
+    @property
+    def table(self) -> NeuronTable:
+        if self._table is None:
+            self._table = NeuronTable.load(dataset=self.dataset)
+        return self._table
+
+    @property
+    def connectome(self):
+        """Lazily load the graph for THIS dataset.
+
+        Hardcoding one loader here meant `load("flywire783")` silently returned the
+        MaleCNS graph, so a user could run a whole experiment on the wrong connectome
+        without a single warning.
+        """
+        if self._connectome is None:
+            self._connectome = load_connectome(self.dataset)
+        return self._connectome
+
+    def __repr__(self) -> str:
+        return (f"<Fly {self.dataset} ports={list(self.ports)} "
+                f"in={len(self.inputs)} out={len(self.outputs)}>")
+
+    # ---------- the four primitives ----------
+    def io(self, inputs=None, outputs=None) -> dict:
+        """Declare the I/O boundary from the release's own superclasses."""
+        if inputs is not None:
+            self.inputs = tuple(inputs)
+        if outputs is not None:
+            self.outputs = tuple(outputs)
+        in_super = {"ol_sensory", "cb_sensory", "vnc_sensory",
+                    "sensory_ascending", "sensory_descending"}
+        out_super = {"vnc_motor", "cb_motor", "descending_neuron",
+                     "vnc_efferent", "cb_endocrine"}
+        ask_in = set(self.inputs) | in_super
+        ask_out = set(self.outputs) | out_super
+        return {"inputs": self.table._mask_super(self.table.df, ask_in)["bodyId"].values,
+                "outputs": self.table._mask_super(self.table.df, ask_out)["bodyId"].values}
+
+    def slice(self, **kwargs) -> Substrate:
+        """Build a substrate slice. Extra inputs come from the declared boundary."""
+        kwargs.setdefault("n_inputs", 512)
+        kwargs.setdefault("n_outputs", 128)
+        kwargs.setdefault("budget", 4096)
+        return Substrate.build(self.connectome, self.table, **kwargs)
+
+    def _positions(self, substrate: Substrate, selector) -> tuple:
+        if isinstance(selector, str):
+            ids = self.table.pool(selector)["bodyId"].values
+        elif isinstance(selector, dict):
+            df = self.table.df
+            mask = np.ones(len(df), bool)
+            for k, v in selector.items():
+                if k == "cell_class":
+                    k = "class"
+                mask &= (df[k] == v).to_numpy()
+            ids = df.loc[mask, "bodyId"].values
+        else:
+            ids = np.asarray(selector, dtype=np.int64)
+        pos = substrate.index().get_indexer(np.asarray(ids))
+        pos = pos[pos >= 0]
+        return pos, substrate.node_ids[pos]
+
+    def port(self, substrate: Substrate, name: str, selector=None, read=True,
+             write=True, protected=True, allow_empty: bool = False) -> Port:
+        """Expose a named read/write handle on real neurons.
+
+        This is the answer to "reserve neurons for working memory": you do not
+        reserve neurons, you tap an existing pool. See section 2.
+
+        Raises if the selector matches nothing INSIDE THE SLICE. A slice is grown
+        from the declared inputs, so a pool that is not reachable from them will
+        simply be absent -- and a silently empty port would look like a working
+        handle that does nothing. Pass `include_pools=` to `slice()` to force the
+        pool in.
+        """
+        if selector is None:
+            pos, ids = substrate.input_pos, substrate.node_ids[substrate.input_pos]
+        else:
+            pos, ids = self._positions(substrate, selector)
+        if pos.size == 0 and not allow_empty:
+            raise ValueError(
+                f"selector {selector!r} matched no neurons inside this slice "
+                f"({substrate.n:,} nodes). Slices are grown from the declared sensory "
+                f"inputs, so a pool that is not reachable from them is absent. Pass "
+                f"include_pools=... to Substrate.build/Fly.slice, or allow_empty=True.")
+        p = Port(name=name, positions=pos, body_ids=ids, read=read, write=write,
+                 protected=protected)
+        self.ports[name] = p
+        return p
+
+    def gate(self, substrate: Substrate, positions=None, name=None, ablate=False,
+             plastic=False, module=None) -> dict:
+        """Change how a slice may be used. Two DISTINCT meanings, never conflated.
+
+        ablate    structural: zero the in/out edges. Changes what the slice can do.
+        plastic   parametric: permit this slice's weights to change. Changes how.
+
+        `ablate` on a substrate goes through `FlyModule.apply_topology_mask` when a
+        module is supplied, so the freeze precondition is enforced.
+        """
+        if positions is None and name is not None:
+            positions = self.ports[name].positions
+        out = {"ablate": bool(ablate), "plastic": bool(plastic)}
+        if ablate:
+            if module is None:
+                raise ValueError("ablate requires `module=` so the topology freeze "
+                                 "precondition is enforced")
+            out["dropped"] = module.apply_topology_mask(
+                np.asarray(positions, dtype=np.int64))
+        del substrate
+        return out
+
+    def reward_bus(self, substrate: Substrate, selector=None, policy=None,
+                   **kwargs):
+        """Build a CreditChannel onto neuromodulatory targets.
+
+        `policy` defaults to NoPlasticity: stimulating and measuring is allowed,
+        changing weights is opt-in and is an assumption, not a property of the
+        connectome (section 3).
+        """
+        from pyfly_lightning.reward import CreditChannel
+        if selector is None:
+            selector = {"class": "DAN"}
+        pos, _ = self._positions(substrate, selector)
+        return CreditChannel(targets=pos, **kwargs)
+
+    def model(self, substrate: Substrate, **kwargs):
+        from pyfly_lightning.train import FlyModule
+        return FlyModule(substrate, **kwargs)
+
+    def plasticity(self, module, positions, policy=None, **kwargs):
+        from pyfly_lightning.reward import NoPlasticity, PlasticPool
+        return PlasticPool(module, np.asarray(positions, dtype=np.int64),
+                           policy=policy or NoPlasticity(), **kwargs)
+
+
+def load(dataset: str = "malecns10") -> Fly:
+    """Load a dataset handle. Data is fetched lazily on first use."""
+    return Fly(dataset=dataset)
+
+
+@dataclass
+class FitResult:
+    """Everything a run must report, including the controls."""
+    accuracy: float
+    chance: float
+    plan: dict
+    controls: dict
+    verdict: str
+    model: object = None
+    substrate: dict = field(default_factory=dict)
+
+    def __str__(self) -> str:
+        return (f"<FitResult acc={self.accuracy:.4f} chance={self.chance:.4f} "
+                f"verdict={self.verdict}>")
+
+    def as_dict(self) -> dict:
+        return {"accuracy": self.accuracy, "chance": self.chance, "plan": self.plan,
+                "controls": self.controls, "verdict": self.verdict,
+                "substrate": self.substrate}
+
+
+def those_controls_are_required() -> str:
+    return ("RewiredNull, ShuffledWeightNull, PooledFeatureNull and NoPlasticityControl "
+            "are built in. A run whose claim does not beat all of them is reported "
+            "UNSUPPORTED.")
+
+
+def fit(X, y, *, dataset: str = "malecns10", temporal: bool = False,
+        modulation: bool = False, plan=None, device: str = "cpu", seed: int = 0,
+        val_frac: float = 0.25, verbose: bool = False) -> FitResult:
+    """Fit the connectome substrate to a task. One call, no connectome knowledge needed.
+
+    Runs: plan -> substrate slice -> ES routing -> trained decoder -> optional
+    modulation channel -> mandatory controls -> verdict.
+
+    The verdict is `SUPPORTED` only if the fit beats every null. On every task
+    measured so far it does not, and the reason is documented in
+    docs/architecture.md sections 11 and 22: the connectome's specific wiring is
+    interchangeable with its own rewiring. The framework's value is that the
+    INTERFACE is optimised for you, not that the substrate is better.
+    """
+    import torch
+
+    from pyfly_lightning.auto import plan as _plan
+    from pyfly_lightning.model.nulls import rewired_substrate, shuffled_weight_substrate
+    from pyfly_lightning.train.pipeline import PipelineConfig, StagedPipeline
+
+    X = np.asarray(X, np.float32)
+    y = np.asarray(y, np.int64)
+    n, d = X.shape
+    C = int(y.max()) + 1
+    pl = plan or _plan(n, d, C, temporal=temporal, modulation=modulation)
+
+    rng = np.random.default_rng(seed)
+    perm = rng.permutation(n)
+    n_val = max(1, int(round(val_frac * n)))
+    itr, iva = perm[:n - n_val], perm[n - n_val:]
+
+    fly = load(dataset)
+    sub = fly.slice(budget=pl.budget, include_pools=pl.pools, n_inputs=pl.n_inputs,
+                    n_outputs=pl.n_outputs, seed=seed)
+    dan_pos = None
+    if modulation:
+        pos, _ = fly._positions(sub, {"class": "DAN"})
+        dan_pos = pos if pos.size else None
+
+    Xt, yt = torch.as_tensor(X[itr], device=device), torch.as_tensor(y[itr], device=device)
+    Xv, yv = torch.as_tensor(X[iva], device=device), torch.as_tensor(y[iva], device=device)
+    chance = float(np.bincount(y[iva], minlength=C).max() / len(iva))
+
+    cfg = PipelineConfig(latent=pl.latent, T=pl.T, seed=seed,
+                         es_iters=pl.es_iters, es_popsize=pl.es_popsize,
+                         use_dan=modulation, train_dan_map=modulation)
+    model = StagedPipeline(sub, in_dim=d, n_classes=C, dan_pos=dan_pos,
+                           config=cfg, device=device).to(device)
+    hist = model.stage1_es(Xt, yt)
+    model.stage2_decode(Xt, yt, Xv, yv)
+    if modulation and dan_pos is not None:
+        model.stage3_modulate(Xt, yt, Xv, yv)
+    acc = model.accuracy(Xv, yv)
+    if verbose:
+        print(f"  plan: {pl}")
+        print(f"  ES {hist[0]:.4f} -> {max(hist):.4f}; final acc {acc:.4f} (chance {chance:.4f})")
+
+    def eval_substrate(s):
+        m = StagedPipeline(s, in_dim=d, n_classes=C, dan_pos=dan_pos,
+                           config=PipelineConfig(latent=pl.latent, T=pl.T, seed=seed,
+                                                 es_iters=pl.es_iters,
+                                                 es_popsize=pl.es_popsize,
+                                                 use_dan=modulation,
+                                                 train_dan_map=modulation),
+                           device=device).to(device)
+        m.stage1_es(Xt, yt)
+        m.stage2_decode(Xt, yt, Xv, yv)
+        return m.accuracy(Xv, yv)
+
+    rewired = eval_substrate(rewired_substrate(sub, seed=seed))
+    shuffled = eval_substrate(shuffled_weight_substrate(sub, seed=seed))
+
+    # the pooled-feature null: ridge on the RAW features, same readout class
+    Ftr = torch.as_tensor(X[itr], device=device)
+    Fva = torch.as_tensor(X[iva], device=device)
+    Y = torch.zeros((len(itr), C), device=device)
+    Y[torch.arange(len(itr), device=device), yt.long()] = 1.0
+    A = Ftr.T @ Ftr + torch.eye(d, device=device)
+    Wr = torch.linalg.solve(A, Ftr.T @ Y)
+    pooled = float(((Fva @ Wr).argmax(1) == yv).float().mean())
+
+    controls = {"fly": round(acc, 4), "rewired": round(rewired, 4),
+                "shuffled_weight": round(shuffled, 4), "pooled_features": round(pooled, 4),
+                "no_plasticity": True, "chance": round(chance, 4)}
+    beats = all(acc > controls[k] for k in ("rewired", "shuffled_weight",
+                                            "pooled_features"))
+    return FitResult(
+        accuracy=round(acc, 4), chance=round(chance, 4), plan=pl.as_dict(),
+        controls=controls,
+        verdict="SUPPORTED" if beats else "UNSUPPORTED",
+        model=model,
+        substrate={"n": sub.n, "edges": sub.n_edges,
+                   "n_groups": sub.input_surface.n_groups},
+    )

pyfly_lightning/auto.py

@@ -0,0 +1,128 @@
+"""Automatic configuration: the user brings a task, the system picks the substrate.
+
+This is the deliverable behind "let users apply the connectome to a task without
+researching its wiring, inputs or outputs". Every default here is derived from a
+measurement in docs/architecture.md, not from taste, and `Plan.reason` records
+which measurement drove which choice so a surprising result can be traced.
+
+What the defaults encode
+------------------------
+* CX is forced in (section 22). Without it, injected activity dies within 2 steps
+  (`late_echo` 0.0); with it, activity recirculates for 60+ steps. This is the one
+  measured contribution the connectome makes, and it is a property of the recurrent
+  pool, not of specific weights.
+* The substrate is kept small. Measured: 4096 nodes / 119k edges already recirculate
+  (late_echo 30) while 16384 nodes / 1.45M edges cost 12x more for the same memory.
+  Larger is not better here and is 12x slower.
+* Routing is grouped at `medium` granularity (~10 modality groups, ~330 parameters).
+  At `fine` granularity the router alone is 41,877 parameters, which no
+  full-covariance optimiser can fit.
+* The ES budget scales with the data, capped so a fit on a small dataset finishes in
+  minutes. Section 22 shows the optimiser is worth +7.6 to +15.3 points over naive
+  input, so it is not optional -- but neither is it free.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import numpy as np
+
+
+@dataclass
+class Plan:
+    budget: int = 4096
+    n_inputs: int = 256
+    n_outputs: int = 64
+    pools: tuple = ("CX",)
+    latent: int = 16
+    T: int = 4
+    es_iters: int = 15
+    es_popsize: int = 16
+    readout: str = "ridge+head"
+    reason: dict = field(default_factory=dict)
+
+    @property
+    def es_evaluations(self) -> int:
+        return self.es_iters * self.es_popsize
+
+    def as_dict(self) -> dict:
+        d = {"budget": self.budget, "n_inputs": self.n_inputs,
+             "n_outputs": self.n_outputs, "pools": list(self.pools),
+             "latent": self.latent, "T": self.T,
+             "es_iters": self.es_iters, "es_popsize": self.es_popsize,
+             "es_evaluations": self.es_evaluations, "readout": self.readout}
+        d["reason"] = self.reason
+        return d
+
+    def __str__(self) -> str:
+        return (f"<Plan {self.budget} nodes, pools={list(self.pools)}, latent={self.latent}, "
+                f"T={self.T}, ES={self.es_evaluations} evals>")
+
+
+# Measured anchors from docs/architecture.md.
+_MIN_BUDGET_FOR_MEMORY = 4096      # late_echo 30 at 4096; 0.0 without CX (section 22)
+_MAX_USEFUL_BUDGET = 8192          # 16384 costs 12x the edges for the same late_echo
+_DAN_COVERAGE_OK = 8192            # DAN reaches 71% at 4096, 83% at 8192 (section 22)
+
+
+def plan(n_samples: int, n_features: int, n_classes: int, *,
+         temporal: bool = False, modulation: bool = True,
+         es_budget_cap: int = 2000) -> Plan:
+    """Pick a substrate and an optimiser budget for a task of this shape."""
+    if n_samples < 2 or n_features < 1 or n_classes < 2:
+        raise ValueError("need at least 2 samples, 1 feature and 2 classes")
+
+    pools = ["CX"]
+    if modulation:
+        pools.append("DAN")
+
+    # Scale the substrate, and especially the READOUT dimension, with the data.
+    # Measured on iris: 64 readout features trained on 120 samples lost to 4 raw
+    # features (0.7133 vs 0.8267). The substrate was not the problem -- the feature
+    # count was. A small dataset needs a narrow interface, so the readout width is
+    # derived from n_samples first and the substrate is sized to match.
+    n_outputs = int(np.clip(n_samples // 16, 8, 64))
+    if n_samples < 1000:
+        budget = 2048               # 4096+ recirculates but is pointless on 150 rows
+    elif n_samples < 3000:
+        budget = _MIN_BUDGET_FOR_MEMORY
+    else:
+        budget = max(_MIN_BUDGET_FOR_MEMORY,
+                     _DAN_COVERAGE_OK if modulation else 0)
+    budget = int(min(budget, _MAX_USEFUL_BUDGET))
+    if modulation and budget < _DAN_COVERAGE_OK:
+        pools = ["CX"]              # no room for a channel with real targets
+
+    # Latent controls the ROUTER's dimensionality, which is what ES has to search.
+    latent = int(np.clip(2 ** int(np.ceil(np.log2(max(4, n_features)))), 8, 64))
+    # The sequence length must exceed the hop distance for the readout to see it.
+    T = 4 if not temporal else 8
+
+    # ES budget scales with the data but is capped: measured is better than assumed,
+    # and 600 evaluations already bought +7.6 to +15.3 points (section 22).
+    popsize = 16
+    iters = int(np.clip(round(n_samples / 40), 5, 60))
+    while iters * popsize > es_budget_cap and iters > 5:
+        iters -= 1
+
+    return Plan(
+        budget=budget, n_inputs=int(min(256, max(64, n_features * 8))),
+        n_outputs=n_outputs, pools=tuple(pools), latent=latent, T=T,
+        es_iters=iters, es_popsize=popsize,
+        reason={
+            "budget": (f"CX is forced in because without it late_echo = 0.0 and activity "
+                       f"dies in 2 steps; capped at {_MAX_USEFUL_BUDGET} because 16384 "
+                       f"costs 12x the edges for the same late_echo"),
+            "pools": ("DAN added because the modulation channel needs in-slice targets "
+                      "-- it reaches 83% of the slice at 8192 nodes"
+                      if modulation else
+                      "CX only; modulation channel disabled by request"),
+            "latent": f"scaled from n_features={n_features} to bound the ES search space",
+            "n_outputs": (f"{n_outputs} readout features = n_samples/16, because 64 "
+                          f"features on 120 samples lost to 4 raw features on iris"),
+            "readout_width": f"n_outputs={n_outputs} scaled from n_samples={n_samples}",
+            "T": ("8 for a temporal task so the sequence exceeds the hop distance"
+                  if temporal else "4; the feed-forward readout needs one pass"),
+            "es": (f"{iters} x {popsize} = {iters * popsize} evaluations, scaled from "
+                   f"n_samples={n_samples} and capped at {es_budget_cap}"),
+        })

pyfly_lightning/cli.py

@@ -0,0 +1,107 @@
+"""Console entry point.
+
+    pyfly describe                       list datasets and their caveats
+    pyfly get malecns10                  download (sha256-verified, cached)
+    pyfly verify malecns10               check the cache
+    pyfly plan --n-samples 150 --n-features 4 --n-classes 3
+    pyfly fit --csv iris.csv --target species --out report.json
+    pyfly fit --npy X.npy y.npy --modulation
+
+`plan` and `fit` are the point: a user brings a task and gets a fitted model with
+the mandatory controls, without studying the connectome's wiring, inputs or outputs.
+"""
+from __future__ import annotations
+
+import json
+import sys
+
+
+def _load_xy(args):
+    import numpy as np
+    if args.csv:
+        import pandas as pd
+        df = pd.read_csv(args.csv)
+        if args.target not in df.columns:
+            raise SystemExit(f"target column {args.target!r} not in {list(df.columns)}")
+        y = df[args.target]
+        X = df.drop(columns=[args.target]).select_dtypes("number")
+        if X.shape[1] == 0:
+            raise SystemExit("no numeric feature columns found")
+        codes, _ = pd.factorize(y)
+        return X.to_numpy(np.float32), codes.astype(np.int64)
+    if len(args.npy) != 2:
+        raise SystemExit("--npy takes exactly two paths: X.npy y.npy")
+    return (np.load(args.npy[0]).astype(np.float32),
+            np.load(args.npy[1]).astype(np.int64))
+
+
+def cmd_plan(args) -> int:
+    from pyfly_lightning.auto import plan
+    p = plan(args.n_samples, args.n_features, args.n_classes,
+             modulation=not args.no_modulation)
+    print(json.dumps(p.as_dict(), indent=2))
+    print(f"\n{p}")
+    return 0
+
+
+def cmd_fit(args) -> int:
+    import pyfly_lightning as pfl
+    X, y = _load_xy(args)
+    print(f"data: {X.shape[0]} samples, {X.shape[1]} features, "
+          f"{int(y.max()) + 1} classes")
+    res = pfl.fit(X, y, dataset=args.dataset, modulation=args.modulation,
+                    device=args.device, seed=args.seed, verbose=True)
+    print()
+    print(json.dumps(res.as_dict(), indent=2))
+    print(f"\n{res}")
+    if res.verdict == "UNSUPPORTED":
+        print("\nverdict=UNSUPPORTED: the fit did not beat every null. The four scores "
+              "above are in the JSON; see the README's measurement table and "
+              "docs/architecture.md for the protocols behind them.")
+    if args.out:
+        with open(args.out, "w") as f:
+            json.dump(res.as_dict(), f, indent=2)
+        print(f"wrote {args.out}")
+    return 0
+
+
+def main(argv: list[str] | None = None) -> int:
+    argv = list(sys.argv[1:] if argv is None else argv)
+    from pyfly_lightning.data.download import main as dl_main
+
+    if not argv or argv[0] in ("-h", "--help", "help"):
+        print(__doc__)
+        return 0
+
+    cmd = argv[0]
+    if cmd in ("describe", "get", "verify"):
+        return dl_main(argv)
+
+    import argparse
+    p = argparse.ArgumentParser(prog=f"pyfly {cmd}", description=__doc__)
+    if cmd == "plan":
+        p.add_argument("--n-samples", type=int, required=True)
+        p.add_argument("--n-features", type=int, required=True)
+        p.add_argument("--n-classes", type=int, default=2)
+        p.add_argument("--no-modulation", action="store_true")
+        a = p.parse_args(argv[1:])
+        return cmd_plan(a)
+    if cmd == "fit":
+        p.add_argument("--csv", default=None)
+        p.add_argument("--target", default=None)
+        p.add_argument("--npy", nargs="*", default=[])
+        p.add_argument("--dataset", default="malecns10")
+        p.add_argument("--modulation", action="store_true")
+        p.add_argument("--device", default="cpu")
+        p.add_argument("--seed", type=int, default=0)
+        p.add_argument("--out", default=None)
+        a = p.parse_args(argv[1:])
+        if not a.csv and len(a.npy) != 2:
+            raise SystemExit("pass either --csv FILE --target COL, or --npy X.npy y.npy")
+        return cmd_fit(a)
+    print(f"unknown command {cmd!r}", file=sys.stderr)
+    return 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

pyfly_lightning/data/__init__.py

@@ -0,0 +1,7 @@
+from pyfly_lightning.data.download import cache_dir, fetch, fetch_dataset, verify
+from pyfly_lightning.data.graph import SparseConnectome, load_malecns
+from pyfly_lightning.data.registry import DATASETS, describe, files_for
+from pyfly_lightning.data.schema import NeuronTable
+
+__all__ = ["DATASETS", "describe", "files_for", "fetch", "fetch_dataset",
+           "cache_dir", "verify", "NeuronTable", "SparseConnectome", "load_malecns"]

pyfly_lightning/data/_checksums.json

@@ -0,0 +1,14 @@
+{
+  "flywire783/Supplemental_file1_neuron_annotations.tsv": {
+    "bytes": 31718505,
+    "sha256": "9a4f8b2f843196074431ebd7cd883536afa1be86c8a4ce90970441e8be81d1be"
+  },
+  "malecns10/body-annotations.feather": {
+    "bytes": 14483314,
+    "sha256": "2177e246113e4cfbf1e7772ec37c6da1955ff22e8063d0b1f833101f99a9a3b2"
+  },
+  "malecns10/connectome-weights-significant-only.feather": {
+    "bytes": 502169298,
+    "sha256": "5c536423a62a688e59e7b441f9c04d6272c9a1f017e35814cf561f8c275d9e9e"
+  }
+}