flyboost-connectome 0.1.0__tar.gz

flyboost_connectome-0.1.0/PKG-INFO

@@ -0,0 +1,145 @@
+Metadata-Version: 2.4
+Name: flyboost-connectome
+Version: 0.1.0
+Summary: Scikit-learn compatible boosting estimators built on a Drosophila connectome reservoir
+Author: FlyBoost contributors
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: numpy>=1.26
+Requires-Dist: scikit-learn>=1.6
+Requires-Dist: torch>=2.4
+Provides-Extra: dev
+Requires-Dist: pytest>=8; extra == "dev"
+Requires-Dist: build>=1.2; extra == "dev"
+
+# FlyBoost
+
+A small scikit-learn-compatible wrapper around the FlyBoost prototype from the
+research notebook. The connectome tensor file is **not** bundled; point
+`graph_path` at your existing `malecns_nomotor.pt`.
+
+## Install
+
+From the project directory:
+
+```bash
+pip install .
+```
+
+Editable development install:
+
+```bash
+pip install -e .
+```
+
+You can also install the source archive directly:
+
+```bash
+pip install flyboost-connectome-0.1.0.tar.gz
+```
+
+## Binary classification
+
+```python
+from flyboost import FlyBoostClassifier
+
+model = FlyBoostClassifier(
+    graph_path="data/malecns_nomotor.pt",
+    n_flies=8,
+    epochs_per_fly=8,
+    fly_steps=8,
+    batch_size=32,
+    fly_lr=2e-3,
+    fly_mode="real",      # "real" | "shuffled" | "frozen"
+    random_state=42,
+    verbose=1,
+)
+
+model.fit(
+    X_train,
+    y_train,
+    eval_set=[(X_val, y_val)],
+)
+
+pred = model.predict(X_val)
+proba = model.predict_proba(X_val)
+print(model.score(X_val, y_val))
+print(model.evals_result())
+```
+
+Switching controls requires only one parameter:
+
+```python
+real = FlyBoostClassifier(fly_mode="real", ...)
+shuffled = FlyBoostClassifier(fly_mode="shuffled", ...)
+frozen = FlyBoostClassifier(fly_mode="frozen", ...)
+```
+
+All boosted shuffled flies share one shuffled topology for the same graph and
+`shuffle_seed`; their sensory mappings still differ by stage seed.
+
+## Regression
+
+```python
+from flyboost import FlyBoostRegressor
+
+reg = FlyBoostRegressor(
+    graph_path="data/malecns_nomotor.pt",
+    n_flies=8,
+    epochs_per_fly=8,
+    fly_mode="real",
+    verbose=1,
+)
+reg.fit(X_train, y_train, eval_set=[(X_val, y_val)])
+yhat = reg.predict(X_val)
+print(reg.score(X_val, y_val))  # sklearn R^2
+```
+
+## sklearn ecosystem
+
+Because the estimators inherit from `BaseEstimator` and the appropriate mixin,
+normal sklearn parameter plumbing works:
+
+```python
+from sklearn.base import clone
+from sklearn.pipeline import make_pipeline
+from sklearn.preprocessing import StandardScaler
+from sklearn.model_selection import GridSearchCV
+
+base = FlyBoostClassifier(graph_path="data/malecns_nomotor.pt")
+clone(base)
+
+pipe = make_pipeline(StandardScaler(), base)
+
+search = GridSearchCV(
+    base,
+    {
+        "fly_mode": ["real", "shuffled", "frozen"],
+        "n_flies": [1, 8, 64],
+        "epochs_per_fly": [1, 8, 64],
+    },
+    cv=3,
+)
+```
+
+For the compute-matched control experiments, set only:
+
+```python
+FlyBoostClassifier(n_flies=1, epochs_per_fly=64, ...)
+FlyBoostClassifier(n_flies=8, epochs_per_fly=8, ...)
+FlyBoostClassifier(n_flies=64, epochs_per_fly=1, ...)
+```
+
+## Verbosity
+
+- `verbose=0`: no training prints
+- `verbose=1`: stage metrics
+- `verbose>=2`: per-fly epoch metrics plus control-construction diagnostics
+
+## Current scope
+
+- classifier: binary targets
+- regressor: single-output targets
+- `eval_set=[(X_val, y_val), ...]` supported
+- `evals_result()` stores stage-wise metrics
+- graph file stays external to keep the package lightweight

flyboost_connectome-0.1.0/README.md

@@ -0,0 +1,131 @@
+# FlyBoost
+
+A small scikit-learn-compatible wrapper around the FlyBoost prototype from the
+research notebook. The connectome tensor file is **not** bundled; point
+`graph_path` at your existing `malecns_nomotor.pt`.
+
+## Install
+
+From the project directory:
+
+```bash
+pip install .
+```
+
+Editable development install:
+
+```bash
+pip install -e .
+```
+
+You can also install the source archive directly:
+
+```bash
+pip install flyboost-connectome-0.1.0.tar.gz
+```
+
+## Binary classification
+
+```python
+from flyboost import FlyBoostClassifier
+
+model = FlyBoostClassifier(
+    graph_path="data/malecns_nomotor.pt",
+    n_flies=8,
+    epochs_per_fly=8,
+    fly_steps=8,
+    batch_size=32,
+    fly_lr=2e-3,
+    fly_mode="real",      # "real" | "shuffled" | "frozen"
+    random_state=42,
+    verbose=1,
+)
+
+model.fit(
+    X_train,
+    y_train,
+    eval_set=[(X_val, y_val)],
+)
+
+pred = model.predict(X_val)
+proba = model.predict_proba(X_val)
+print(model.score(X_val, y_val))
+print(model.evals_result())
+```
+
+Switching controls requires only one parameter:
+
+```python
+real = FlyBoostClassifier(fly_mode="real", ...)
+shuffled = FlyBoostClassifier(fly_mode="shuffled", ...)
+frozen = FlyBoostClassifier(fly_mode="frozen", ...)
+```
+
+All boosted shuffled flies share one shuffled topology for the same graph and
+`shuffle_seed`; their sensory mappings still differ by stage seed.
+
+## Regression
+
+```python
+from flyboost import FlyBoostRegressor
+
+reg = FlyBoostRegressor(
+    graph_path="data/malecns_nomotor.pt",
+    n_flies=8,
+    epochs_per_fly=8,
+    fly_mode="real",
+    verbose=1,
+)
+reg.fit(X_train, y_train, eval_set=[(X_val, y_val)])
+yhat = reg.predict(X_val)
+print(reg.score(X_val, y_val))  # sklearn R^2
+```
+
+## sklearn ecosystem
+
+Because the estimators inherit from `BaseEstimator` and the appropriate mixin,
+normal sklearn parameter plumbing works:
+
+```python
+from sklearn.base import clone
+from sklearn.pipeline import make_pipeline
+from sklearn.preprocessing import StandardScaler
+from sklearn.model_selection import GridSearchCV
+
+base = FlyBoostClassifier(graph_path="data/malecns_nomotor.pt")
+clone(base)
+
+pipe = make_pipeline(StandardScaler(), base)
+
+search = GridSearchCV(
+    base,
+    {
+        "fly_mode": ["real", "shuffled", "frozen"],
+        "n_flies": [1, 8, 64],
+        "epochs_per_fly": [1, 8, 64],
+    },
+    cv=3,
+)
+```
+
+For the compute-matched control experiments, set only:
+
+```python
+FlyBoostClassifier(n_flies=1, epochs_per_fly=64, ...)
+FlyBoostClassifier(n_flies=8, epochs_per_fly=8, ...)
+FlyBoostClassifier(n_flies=64, epochs_per_fly=1, ...)
+```
+
+## Verbosity
+
+- `verbose=0`: no training prints
+- `verbose=1`: stage metrics
+- `verbose>=2`: per-fly epoch metrics plus control-construction diagnostics
+
+## Current scope
+
+- classifier: binary targets
+- regressor: single-output targets
+- `eval_set=[(X_val, y_val), ...]` supported
+- `evals_result()` stores stage-wise metrics
+- graph file stays external to keep the package lightweight

flyboost_connectome-0.1.0/pyproject.toml

@@ -0,0 +1,25 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "flyboost-connectome"
+version = "0.1.0"
+description = "Scikit-learn compatible boosting estimators built on a Drosophila connectome reservoir"
+readme = "README.md"
+requires-python = ">=3.10"
+authors = [{name = "FlyBoost contributors"}]
+dependencies = [
+  "numpy>=1.26",
+  "scikit-learn>=1.6",
+  "torch>=2.4",
+]
+
+[project.optional-dependencies]
+dev = ["pytest>=8", "build>=1.2"]
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]

flyboost_connectome-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

flyboost_connectome-0.1.0/src/flyboost/__init__.py

@@ -0,0 +1,14 @@
+from .estimators import FlyBoostClassifier, FlyBoostRegressor
+from .fly import FrozenTrainableFly, ShuffledTrainableFly, TrainableFly
+from .graph import FlyGraph
+
+__all__ = [
+    "FlyBoostClassifier",
+    "FlyBoostRegressor",
+    "FlyGraph",
+    "TrainableFly",
+    "ShuffledTrainableFly",
+    "FrozenTrainableFly",
+]
+
+__version__ = "0.1.0"

flyboost_connectome-0.1.0/src/flyboost/estimators.py

@@ -0,0 +1,508 @@
+from __future__ import annotations
+
+import copy
+import math
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+import torch
+from sklearn.base import BaseEstimator, ClassifierMixin, RegressorMixin
+from sklearn.utils.multiclass import check_classification_targets, type_of_target
+from sklearn.utils.validation import check_is_fitted, validate_data
+
+from .fly import FLY_MODES
+from .graph import FlyGraph, load_graph_cached
+
+
+@dataclass
+class BoostStage:
+    state_dict: dict[str, torch.Tensor]
+
+
+def _as_float_tensor(x: np.ndarray) -> torch.Tensor:
+    return torch.tensor(np.asarray(x, dtype=np.float32), dtype=torch.float32)
+
+
+def _resolve_device(device: str) -> str:
+    if device != "auto":
+        return str(device)
+    return "cuda" if torch.cuda.is_available() else "cpu"
+
+
+def _check_sample_weight(sample_weight, n_samples: int) -> torch.Tensor:
+    if sample_weight is None:
+        return torch.ones(n_samples, dtype=torch.float32)
+    sw = np.asarray(sample_weight, dtype=np.float32)
+    if sw.ndim != 1 or len(sw) != n_samples:
+        raise ValueError(
+            f"sample_weight must have shape ({n_samples},), got {sw.shape}"
+        )
+    if np.any(sw < 0):
+        raise ValueError("sample_weight cannot contain negative values")
+    if not np.any(sw > 0):
+        raise ValueError("sample_weight must contain at least one positive value")
+    return torch.from_numpy(sw)
+
+
+class _FlyBoostBase(BaseEstimator):
+    """Shared sklearn-facing plumbing for FlyBoost estimators."""
+
+    def __init__(
+        self,
+        graph_path: str = "data/malecns_nomotor.pt",
+        n_flies: int = 8,
+        learning_rate: float = 0.2,
+        fly_steps: int = 8,
+        epochs_per_fly: int = 8,
+        batch_size: int = 32,
+        fly_lr: float = 2e-3,
+        fly_gain_lr: float = 0.0,
+        fly_mode: str = "real",
+        leak: float = 0.5,
+        max_log_gain: float = math.log(2.0),
+        shuffle_seed: int = 20260915,
+        device: str = "cpu",
+        random_state: int | None = 42,
+        verbose: int = 0,
+    ):
+        self.graph_path = graph_path
+        self.n_flies = n_flies
+        self.learning_rate = learning_rate
+        self.fly_steps = fly_steps
+        self.epochs_per_fly = epochs_per_fly
+        self.batch_size = batch_size
+        self.fly_lr = fly_lr
+        self.fly_gain_lr = fly_lr if fly_gain_lr == 0 else fly_gain_lr
+        self.fly_mode = fly_mode
+        self.leak = leak
+        self.max_log_gain = max_log_gain
+        self.shuffle_seed = shuffle_seed
+        self.device = device
+        self.random_state = random_state
+        self.verbose = verbose
+
+    def _validate_hyperparameters(self) -> None:
+        if self.fly_mode not in FLY_MODES:
+            raise ValueError(
+                f"fly_mode must be one of {sorted(FLY_MODES)}, got {self.fly_mode!r}"
+            )
+        if int(self.n_flies) < 1:
+            raise ValueError("n_flies must be >= 1")
+        if int(self.epochs_per_fly) < 1:
+            raise ValueError("epochs_per_fly must be >= 1")
+        if int(self.fly_steps) < 1:
+            raise ValueError("fly_steps must be >= 1")
+        if int(self.batch_size) < 1:
+            raise ValueError("batch_size must be >= 1")
+        if float(self.learning_rate) <= 0:
+            raise ValueError("learning_rate must be > 0")
+        if float(self.fly_lr) <= 0:
+            raise ValueError("fly_lr must be > 0")
+        if not (0 < float(self.leak) <= 1):
+            raise ValueError("leak must be in (0, 1]")
+
+    @property
+    def _seed(self) -> int:
+        return 0 if self.random_state is None else int(self.random_state)
+
+    def _graph(self) -> FlyGraph:
+        path = str(Path(self.graph_path).expanduser().resolve())
+        if not Path(path).exists():
+            raise FileNotFoundError(
+                f"FlyBoost graph not found: {path}. "
+                "Pass graph_path=... pointing to malecns_nomotor.pt."
+            )
+        return load_graph_cached(path)
+
+    def _new_fly(self, graph: FlyGraph, stage: int):
+        fly_cls = FLY_MODES[self.fly_mode]
+        kwargs = dict(
+            graph=graph,
+            n_features=self.n_features_in_,
+            steps=self.fly_steps,
+            leak=self.leak,
+            max_log_gain=self.max_log_gain,
+            seed=self._seed + int(stage),
+            device=self.device_,
+            verbose=self.verbose,
+        )
+        if self.fly_mode == "shuffled":
+            kwargs["shuffle_seed"] = self.shuffle_seed
+
+        # nn.Linear initializes from PyTorch's RNG.  Fork/restore the RNG so
+        # random_state makes repeated fits deterministic without perturbing
+        # the caller's global RNG state.
+        init_seed = self._seed + 20000 + int(stage)
+        cuda_devices = []
+        if str(self.device_).startswith("cuda") and torch.cuda.is_available():
+            dev = torch.device(self.device_)
+            cuda_devices = [dev.index if dev.index is not None else torch.cuda.current_device()]
+        with torch.random.fork_rng(devices=cuda_devices):
+            torch.manual_seed(init_seed)
+            if cuda_devices:
+                torch.cuda.manual_seed_all(init_seed)
+            return fly_cls(**kwargs)
+
+    def _make_optimizer(self, fly) -> torch.optim.Optimizer:
+        # Only genuinely trainable tensors enter the optimizer.  This makes
+        # frozen mode readout-only by construction.
+        params = [p for p in fly.parameters() if p.requires_grad]
+        return torch.optim.Adam([
+                {
+                    "params": [
+                        fly.pre_theta,
+                        fly.post_theta,
+                        fly.sensor_theta,
+                    ],
+                    "lr": self.fly_gain_lr,
+                },
+                {
+                    "params": fly.readout.parameters(),
+                    "lr": self.fly_lr,
+                },
+            ])
+
+    def _batched_predict(self, fly, x: torch.Tensor) -> torch.Tensor:
+        out = []
+        with torch.no_grad():
+            for start in range(0, len(x), int(self.batch_size)):
+                xb = x[start : start + int(self.batch_size)].to(self.device_)
+                out.append(fly(xb).cpu())
+        return torch.cat(out) if out else torch.empty(0, dtype=torch.float32)
+
+    def _fit_fly_to_target(
+        self,
+        fly,
+        x: torch.Tensor,
+        target: torch.Tensor,
+        weight: torch.Tensor,
+        stage: int,
+    ) -> None:
+        opt = self._make_optimizer(fly)
+        gen = torch.Generator(device="cpu")
+        gen.manual_seed(self._seed + 10000 + int(stage))
+
+        for epoch in range(int(self.epochs_per_fly)):
+            perm = torch.randperm(len(x), generator=gen)
+            weighted_loss_sum = 0.0
+            seen_weight = 0.0
+
+            for start in range(0, len(x), int(self.batch_size)):
+                idx = perm[start : start + int(self.batch_size)]
+                xb = x[idx].to(self.device_)
+                tb = target[idx].to(self.device_)
+                wb = weight[idx].to(self.device_)
+
+                opt.zero_grad(set_to_none=True)
+                pred = fly(xb)
+                denom = wb.sum().clamp_min(1e-8)
+                loss = (wb * (pred - tb).square()).sum() / denom
+                loss.backward()
+                torch.nn.utils.clip_grad_norm_(
+                    [p for p in fly.parameters() if p.requires_grad], 5.0
+                )
+                opt.step()
+
+                batch_weight = float(wb.sum().detach().cpu())
+                weighted_loss_sum += float(loss.detach().cpu()) * batch_weight
+                seen_weight += batch_weight
+
+            if int(self.verbose) >= 2:
+                mean_loss = weighted_loss_sum / max(seen_weight, 1e-12)
+                print(
+                    f"  fly {stage + 1}/{self.n_flies} "
+                    f"epoch {epoch + 1}/{self.epochs_per_fly} "
+                    f"weighted_mse={mean_loss:.5f}"
+                )
+
+    @staticmethod
+    def _freeze_and_save(fly) -> BoostStage:
+        fly.eval()
+        for p in fly.parameters():
+            p.requires_grad_(False)
+        state = {
+            k: v.detach().cpu().clone()
+            for k, v in fly.state_dict().items()
+        }
+        return BoostStage(state)
+
+    def _predict_raw_tensor(self, X) -> torch.Tensor:
+        check_is_fitted(self, attributes=["stages_", "n_features_in_"])
+        Xv = validate_data(
+            self, X, reset=False, dtype=np.float32, ensure_2d=True
+        )
+        x = _as_float_tensor(Xv)
+        graph = self._graph()
+        raw = torch.full(
+            (len(x),), float(self.base_score_), dtype=torch.float32
+        )
+        for stage, saved in enumerate(self.stages_):
+            fly = self._new_fly(graph, stage)
+            fly.load_state_dict(saved.state_dict, strict=True)
+            fly.eval()
+            raw += float(self.learning_rate) * self._batched_predict(fly, x)
+            del fly
+        return raw
+
+    def evals_result(self) -> dict:
+        check_is_fitted(self, attributes=["evals_result_"])
+        return copy.deepcopy(self.evals_result_)
+
+
+class FlyBoostClassifier(ClassifierMixin, _FlyBoostBase):
+    """
+    Binary classifier using stage-wise Fly connectome weak learners.
+
+    Parameters
+    ----------
+    fly_mode : {"real", "shuffled", "frozen"}, default="real"
+        Selects the biological-connectome model, degree-preserving shuffled
+        control, or frozen-connectome/readout-only control.
+
+    Notes
+    -----
+    This first packaged version intentionally supports binary classification.
+    """
+
+    def __sklearn_tags__(self):
+        tags = super().__sklearn_tags__()
+        if tags.classifier_tags is not None:
+            tags.classifier_tags.multi_class = False
+            tags.classifier_tags.poor_score = True
+        return tags
+
+    def fit(self, X, y, eval_set=None):
+        self._validate_hyperparameters()
+        Xv, yv = validate_data(
+            self,
+            X,
+            y,
+            reset=True,
+            dtype=np.float32,
+            ensure_2d=True,
+        )
+        check_classification_targets(yv)
+        y_type = type_of_target(yv, input_name="y", raise_unknown=True)
+        if y_type != "binary":
+            raise ValueError(
+                "Only binary classification is supported. "
+                f"The type of the target is {y_type}."
+            )
+        classes = np.unique(yv)
+        if len(classes) != 2:
+            raise ValueError(
+                "Only binary classification with exactly two classes is supported; "
+                f"got {len(classes)} class" + ("es" if len(classes) != 1 else "") + "."
+            )
+        self.classes_ = classes
+        y01 = (np.asarray(yv) == classes[1]).astype(np.float32)
+
+        self.device_ = _resolve_device(self.device)
+        x = _as_float_tensor(Xv)
+        y_t = torch.from_numpy(y01)
+        sw = torch.ones(len(x), dtype=torch.float32)
+
+        graph = self._graph()
+        self.stages_ = []
+        self.evals_result_ = {"train": {"logloss": [], "accuracy": []}}
+
+        weighted_pos = (sw * y_t).sum() / sw.sum().clamp_min(1e-8)
+        weighted_pos = weighted_pos.clamp(1e-4, 1 - 1e-4)
+        self.base_score_ = float(torch.log(weighted_pos / (1 - weighted_pos)))
+        ensemble = torch.full_like(y_t, self.base_score_)
+
+        eval_data = self._prepare_classifier_eval_set(eval_set)
+        eval_ensembles = [
+            torch.full((len(ex),), self.base_score_, dtype=torch.float32)
+            for ex, _ in eval_data
+        ]
+        for i in range(len(eval_data)):
+            self.evals_result_[f"validation_{i}"] = {
+                "logloss": [],
+                "accuracy": [],
+            }
+
+        for stage in range(int(self.n_flies)):
+            p = torch.sigmoid(ensemble)
+            g = p - y_t
+            h = (p * (1 - p)).clamp_min(1e-3)
+            z = (-g / h).clamp(-8.0, 8.0)
+            fit_weight = h * sw
+
+            fly = self._new_fly(graph, stage)
+            self._fit_fly_to_target(fly, x, z, fit_weight, stage)
+
+            ensemble += float(self.learning_rate) * self._batched_predict(fly, x)
+            train_logloss = float(
+                torch.nn.functional.binary_cross_entropy_with_logits(
+                    ensemble, y_t, weight=sw, reduction="sum"
+                )
+                / sw.sum().clamp_min(1e-8)
+            )
+            train_acc = float(((ensemble > 0).float() == y_t).float().mean())
+            self.evals_result_["train"]["logloss"].append(train_logloss)
+            self.evals_result_["train"]["accuracy"].append(train_acc)
+
+            msg = (
+                f"[stage {stage + 1}] train_logloss={train_logloss:.5f} "
+                f"train_acc={train_acc:.3f}"
+            )
+
+            for i, (ex, ey) in enumerate(eval_data):
+                eval_ensembles[i] += float(self.learning_rate) * self._batched_predict(
+                    fly, ex
+                )
+                logits = eval_ensembles[i]
+                val_logloss = float(
+                    torch.nn.functional.binary_cross_entropy_with_logits(logits, ey)
+                )
+                val_acc = float(((logits > 0).float() == ey).float().mean())
+                key = f"validation_{i}"
+                self.evals_result_[key]["logloss"].append(val_logloss)
+                self.evals_result_[key]["accuracy"].append(val_acc)
+                msg += (
+                    f" {key}_logloss={val_logloss:.5f} "
+                    f"{key}_acc={val_acc:.3f}"
+                )
+
+            if int(self.verbose) >= 1:
+                print(msg)
+
+            self.stages_.append(self._freeze_and_save(fly))
+            del fly
+
+        return self
+
+    def _prepare_classifier_eval_set(self, eval_set):
+        if eval_set is None:
+            return []
+        out = []
+        for X_eval, y_eval in eval_set:
+            Xe = validate_data(
+                self, X_eval, reset=False, dtype=np.float32, ensure_2d=True
+            )
+            ye = np.asarray(y_eval)
+            unknown = np.setdiff1d(np.unique(ye), self.classes_)
+            if len(unknown):
+                raise ValueError(
+                    f"eval_set contains labels not seen in fit(): {unknown.tolist()}"
+                )
+            y01 = (ye == self.classes_[1]).astype(np.float32)
+            if len(Xe) != len(y01):
+                raise ValueError("X_eval and y_eval have inconsistent lengths")
+            out.append((_as_float_tensor(Xe), torch.from_numpy(y01)))
+        return out
+
+    def decision_function(self, X) -> np.ndarray:
+        return self._predict_raw_tensor(X).numpy()
+
+    def predict_proba(self, X) -> np.ndarray:
+        logits = self._predict_raw_tensor(X)
+        p1 = torch.sigmoid(logits).numpy()
+        return np.column_stack([1.0 - p1, p1])
+
+    def predict(self, X) -> np.ndarray:
+        logits = self._predict_raw_tensor(X).numpy()
+        return self.classes_[(logits > 0).astype(np.int64)]
+
+
+class FlyBoostRegressor(RegressorMixin, _FlyBoostBase):
+    """Single-output squared-error regressor using Fly connectome weak learners."""
+
+    def __sklearn_tags__(self):
+        tags = super().__sklearn_tags__()
+        if tags.regressor_tags is not None:
+            tags.regressor_tags.poor_score = True
+        return tags
+
+    def fit(self, X, y, eval_set=None):
+        self._validate_hyperparameters()
+        Xv, yv = validate_data(
+            self,
+            X,
+            y,
+            reset=True,
+            dtype=np.float32,
+            ensure_2d=True,
+            y_numeric=True,
+        )
+        y_arr = np.asarray(yv, dtype=np.float32)
+        if y_arr.ndim != 1:
+            raise ValueError("FlyBoostRegressor supports single-output regression only")
+
+        self.device_ = _resolve_device(self.device)
+        x = _as_float_tensor(Xv)
+        y_t = torch.from_numpy(y_arr)
+        sw = torch.ones(len(x), dtype=torch.float32)
+
+        graph = self._graph()
+        self.stages_ = []
+        self.evals_result_ = {"train": {"rmse": [], "mae": []}}
+
+        self.base_score_ = float((sw * y_t).sum() / sw.sum().clamp_min(1e-8))
+        ensemble = torch.full_like(y_t, self.base_score_)
+
+        eval_data = self._prepare_regression_eval_set(eval_set)
+        eval_ensembles = [
+            torch.full((len(ex),), self.base_score_, dtype=torch.float32)
+            for ex, _ in eval_data
+        ]
+        for i in range(len(eval_data)):
+            self.evals_result_[f"validation_{i}"] = {"rmse": [], "mae": []}
+
+        for stage in range(int(self.n_flies)):
+            residual = y_t - ensemble
+            fly = self._new_fly(graph, stage)
+            self._fit_fly_to_target(fly, x, residual, sw, stage)
+
+            ensemble += float(self.learning_rate) * self._batched_predict(fly, x)
+            err = ensemble - y_t
+            train_rmse = float(torch.sqrt((sw * err.square()).sum() / sw.sum()))
+            train_mae = float((sw * err.abs()).sum() / sw.sum())
+            self.evals_result_["train"]["rmse"].append(train_rmse)
+            self.evals_result_["train"]["mae"].append(train_mae)
+
+            msg = (
+                f"[stage {stage + 1}] train_rmse={train_rmse:.5f} "
+                f"train_mae={train_mae:.5f}"
+            )
+
+            for i, (ex, ey) in enumerate(eval_data):
+                eval_ensembles[i] += float(self.learning_rate) * self._batched_predict(
+                    fly, ex
+                )
+                err_eval = eval_ensembles[i] - ey
+                rmse = float(torch.sqrt(torch.mean(err_eval.square())))
+                mae = float(torch.mean(err_eval.abs()))
+                key = f"validation_{i}"
+                self.evals_result_[key]["rmse"].append(rmse)
+                self.evals_result_[key]["mae"].append(mae)
+                msg += f" {key}_rmse={rmse:.5f} {key}_mae={mae:.5f}"
+
+            if int(self.verbose) >= 1:
+                print(msg)
+
+            self.stages_.append(self._freeze_and_save(fly))
+            del fly
+
+        return self
+
+    def _prepare_regression_eval_set(self, eval_set):
+        if eval_set is None:
+            return []
+        out = []
+        for X_eval, y_eval in eval_set:
+            Xe = validate_data(
+                self, X_eval, reset=False, dtype=np.float32, ensure_2d=True
+            )
+            ye = np.asarray(y_eval, dtype=np.float32)
+            if ye.ndim != 1:
+                raise ValueError("Each regression y_eval must be 1-dimensional")
+            if len(Xe) != len(ye):
+                raise ValueError("X_eval and y_eval have inconsistent lengths")
+            out.append((_as_float_tensor(Xe), torch.from_numpy(ye)))
+        return out
+
+    def predict(self, X) -> np.ndarray:
+        return self._predict_raw_tensor(X).numpy()

flyboost_connectome-0.1.0/src/flyboost/fly.py

@@ -0,0 +1,225 @@
+from __future__ import annotations
+
+import math
+
+import torch
+from torch import nn
+
+from .graph import FlyGraph
+
+
+class TrainableFly(nn.Module):
+    def __init__(
+        self,
+        graph: FlyGraph,
+        n_features: int,
+        steps: int = 4,
+        leak: float = 0.5,
+        max_log_gain: float = math.log(2.0),
+        seed: int = 0,
+        device: str = "cpu",
+        verbose: int = 0,
+    ):
+        super().__init__()
+        self.n = graph.num_nodes
+        self.n_features = int(n_features)
+        self.steps = int(steps)
+        self.leak = float(leak)
+        self.max_log_gain = float(max_log_gain)
+        self.device_name = device
+        self.verbose = int(verbose)
+
+        W = torch.sparse_coo_tensor(
+            graph.edge_index.to(device),
+            graph.base_weight.to(device),
+            size=(self.n, self.n),
+            dtype=torch.float32,
+            device=device,
+        ).coalesce()
+        W.requires_grad_(False)
+
+        # Large fixed graph tensors are deliberately excluded from state_dict.
+        self.register_buffer("W", W, persistent=False)
+        self.register_buffer("sensory_idx", graph.sensory_idx.to(device), persistent=False)
+        self.register_buffer("readout_idx", graph.readout_idx.to(device), persistent=False)
+
+        gen = torch.Generator(device="cpu")
+        gen.manual_seed(int(seed))
+        perm = graph.sensory_idx[
+            torch.randperm(len(graph.sensory_idx), generator=gen)
+        ]
+        feat = torch.arange(len(perm), dtype=torch.long) % self.n_features
+        inv = torch.full((graph.num_nodes,), -1, dtype=torch.long)
+        inv[perm] = feat
+        sensory_feature = inv[graph.sensory_idx]
+        self.register_buffer(
+            "sensory_feature", sensory_feature.to(device), persistent=False
+        )
+
+        self.pre_theta = nn.Parameter(torch.zeros(self.n, device=device))
+        self.post_theta = nn.Parameter(torch.zeros(self.n, device=device))
+        self.sensor_theta = nn.Parameter(
+            torch.zeros(len(graph.sensory_idx), device=device)
+        )
+        self.readout = nn.Linear(
+            len(graph.readout_idx), 1, bias=True, device=device
+        )
+
+    def _gain(self, theta: torch.Tensor) -> torch.Tensor:
+        return torch.exp(self.max_log_gain * torch.tanh(theta))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x.to(self.pre_theta.device, dtype=torch.float32)
+        b = x.shape[0]
+        h = torch.zeros((b, self.n), dtype=torch.float32, device=x.device)
+
+        drive = torch.zeros_like(h)
+        drive[:, self.sensory_idx] = (
+            x[:, self.sensory_feature] * self._gain(self.sensor_theta)
+        )
+
+        pre_gain = self._gain(self.pre_theta)
+        post_gain = self._gain(self.post_theta)
+
+        for _ in range(self.steps):
+            h_pre = h * pre_gain
+            recurrent = torch.sparse.mm(self.W, h_pre.T).T
+            recurrent = recurrent * post_gain
+            proposal = torch.tanh(recurrent + drive)
+            h = (1.0 - self.leak) * h + self.leak * proposal
+
+        return self.readout(h[:, self.readout_idx]).squeeze(-1)
+
+
+class ShuffledTrainableFly(TrainableFly):
+    """
+    Null-model fly.
+
+    Preserves the neuron populations, edge count, exact source out-degrees,
+    exact destination in-degrees, and source-associated signed weights while
+    destroying the biological pre -> post pairing.
+    """
+
+    _cached_graph: FlyGraph | None = None
+    _cached_source_id: int | None = None
+    _cached_shuffle_seed: int | None = None
+
+    def __init__(
+        self,
+        graph: FlyGraph,
+        n_features: int,
+        steps: int = 4,
+        leak: float = 0.5,
+        max_log_gain: float = math.log(2.0),
+        seed: int = 0,
+        device: str = "cpu",
+        verbose: int = 0,
+        shuffle_seed: int = 20260915,
+    ):
+        cls = self.__class__
+        if (
+            cls._cached_graph is None
+            or cls._cached_source_id != id(graph)
+            or cls._cached_shuffle_seed != int(shuffle_seed)
+        ):
+            if verbose >= 2:
+                print("[ShuffledTrainableFly] building shuffled connectome...")
+
+            edge_index = graph.edge_index.cpu()
+            base_weight = graph.base_weight.cpu()
+
+            # prepare_malecns.py convention:
+            # edge_index[0] = post, edge_index[1] = pre
+            post = edge_index[0]
+            pre = edge_index[1]
+            E = post.numel()
+
+            gen = torch.Generator(device="cpu")
+            gen.manual_seed(int(shuffle_seed))
+            perm = torch.randperm(E, generator=gen)
+            shuffled_post = post[perm]
+            shuffled_weight = base_weight.clone()
+
+            incoming_abs = torch.bincount(
+                shuffled_post,
+                weights=shuffled_weight.abs(),
+                minlength=graph.num_nodes,
+            )
+            shuffled_weight = (
+                shuffled_weight
+                / incoming_abs[shuffled_post].clamp_min(1e-12)
+            )
+
+            shuffled_edge_index = torch.stack([shuffled_post, pre], dim=0)
+            cls._cached_graph = FlyGraph(
+                num_nodes=graph.num_nodes,
+                edge_index=shuffled_edge_index,
+                base_weight=shuffled_weight,
+                sensory_idx=graph.sensory_idx,
+                readout_idx=graph.readout_idx,
+            )
+            cls._cached_source_id = id(graph)
+            cls._cached_shuffle_seed = int(shuffle_seed)
+
+            if verbose >= 2:
+                print(
+                    f"[ShuffledTrainableFly] done: "
+                    f"{graph.num_nodes:,} neurons, {E:,} shuffled edges"
+                )
+
+        super().__init__(
+            graph=cls._cached_graph,
+            n_features=n_features,
+            steps=steps,
+            leak=leak,
+            max_log_gain=max_log_gain,
+            seed=seed,
+            device=device,
+            verbose=verbose,
+        )
+
+
+class FrozenTrainableFly(TrainableFly):
+    """Real connectome with all neuron/sensory gains frozen; readout only trains."""
+
+    def __init__(
+        self,
+        graph: FlyGraph,
+        n_features: int,
+        steps: int = 4,
+        leak: float = 0.5,
+        max_log_gain: float = math.log(2.0),
+        seed: int = 0,
+        device: str = "cpu",
+        verbose: int = 0,
+    ):
+        super().__init__(
+            graph=graph,
+            n_features=n_features,
+            steps=steps,
+            leak=leak,
+            max_log_gain=max_log_gain,
+            seed=seed,
+            device=device,
+            verbose=verbose,
+        )
+
+        self.pre_theta.requires_grad_(False)
+        self.post_theta.requires_grad_(False)
+        self.sensor_theta.requires_grad_(False)
+
+        if verbose >= 2:
+            n_trainable = sum(
+                p.numel() for p in self.parameters() if p.requires_grad
+            )
+            print(
+                f"[FrozenTrainableFly] trainable params = "
+                f"{n_trainable:,} (readout only)"
+            )
+
+
+FLY_MODES = {
+    "real": TrainableFly,
+    "shuffled": ShuffledTrainableFly,
+    "frozen": FrozenTrainableFly,
+}

flyboost_connectome-0.1.0/src/flyboost/graph.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from functools import lru_cache
+from pathlib import Path
+
+import torch
+
+
+@dataclass
+class FlyGraph:
+    num_nodes: int
+    edge_index: torch.Tensor
+    base_weight: torch.Tensor
+    sensory_idx: torch.Tensor
+    readout_idx: torch.Tensor
+
+    @classmethod
+    def load(cls, path: str | Path) -> "FlyGraph":
+        d = torch.load(path, map_location="cpu", weights_only=False)
+        return cls(
+            int(d["num_nodes"]),
+            d["edge_index"].long(),
+            d["base_weight"].float(),
+            d["sensory_idx"].long(),
+            d["readout_idx"].long(),
+        )
+
+
+@lru_cache(maxsize=4)
+def load_graph_cached(path: str) -> FlyGraph:
+    """Load a graph once per absolute path inside the current Python process."""
+    resolved = str(Path(path).expanduser().resolve())
+    return FlyGraph.load(resolved)

flyboost_connectome-0.1.0/src/flyboost_connectome.egg-info/PKG-INFO

@@ -0,0 +1,145 @@
+Metadata-Version: 2.4
+Name: flyboost-connectome
+Version: 0.1.0
+Summary: Scikit-learn compatible boosting estimators built on a Drosophila connectome reservoir
+Author: FlyBoost contributors
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: numpy>=1.26
+Requires-Dist: scikit-learn>=1.6
+Requires-Dist: torch>=2.4
+Provides-Extra: dev
+Requires-Dist: pytest>=8; extra == "dev"
+Requires-Dist: build>=1.2; extra == "dev"
+
+# FlyBoost
+
+A small scikit-learn-compatible wrapper around the FlyBoost prototype from the
+research notebook. The connectome tensor file is **not** bundled; point
+`graph_path` at your existing `malecns_nomotor.pt`.
+
+## Install
+
+From the project directory:
+
+```bash
+pip install .
+```
+
+Editable development install:
+
+```bash
+pip install -e .
+```
+
+You can also install the source archive directly:
+
+```bash
+pip install flyboost-connectome-0.1.0.tar.gz
+```
+
+## Binary classification
+
+```python
+from flyboost import FlyBoostClassifier
+
+model = FlyBoostClassifier(
+    graph_path="data/malecns_nomotor.pt",
+    n_flies=8,
+    epochs_per_fly=8,
+    fly_steps=8,
+    batch_size=32,
+    fly_lr=2e-3,
+    fly_mode="real",      # "real" | "shuffled" | "frozen"
+    random_state=42,
+    verbose=1,
+)
+
+model.fit(
+    X_train,
+    y_train,
+    eval_set=[(X_val, y_val)],
+)
+
+pred = model.predict(X_val)
+proba = model.predict_proba(X_val)
+print(model.score(X_val, y_val))
+print(model.evals_result())
+```
+
+Switching controls requires only one parameter:
+
+```python
+real = FlyBoostClassifier(fly_mode="real", ...)
+shuffled = FlyBoostClassifier(fly_mode="shuffled", ...)
+frozen = FlyBoostClassifier(fly_mode="frozen", ...)
+```
+
+All boosted shuffled flies share one shuffled topology for the same graph and
+`shuffle_seed`; their sensory mappings still differ by stage seed.
+
+## Regression
+
+```python
+from flyboost import FlyBoostRegressor
+
+reg = FlyBoostRegressor(
+    graph_path="data/malecns_nomotor.pt",
+    n_flies=8,
+    epochs_per_fly=8,
+    fly_mode="real",
+    verbose=1,
+)
+reg.fit(X_train, y_train, eval_set=[(X_val, y_val)])
+yhat = reg.predict(X_val)
+print(reg.score(X_val, y_val))  # sklearn R^2
+```
+
+## sklearn ecosystem
+
+Because the estimators inherit from `BaseEstimator` and the appropriate mixin,
+normal sklearn parameter plumbing works:
+
+```python
+from sklearn.base import clone
+from sklearn.pipeline import make_pipeline
+from sklearn.preprocessing import StandardScaler
+from sklearn.model_selection import GridSearchCV
+
+base = FlyBoostClassifier(graph_path="data/malecns_nomotor.pt")
+clone(base)
+
+pipe = make_pipeline(StandardScaler(), base)
+
+search = GridSearchCV(
+    base,
+    {
+        "fly_mode": ["real", "shuffled", "frozen"],
+        "n_flies": [1, 8, 64],
+        "epochs_per_fly": [1, 8, 64],
+    },
+    cv=3,
+)
+```
+
+For the compute-matched control experiments, set only:
+
+```python
+FlyBoostClassifier(n_flies=1, epochs_per_fly=64, ...)
+FlyBoostClassifier(n_flies=8, epochs_per_fly=8, ...)
+FlyBoostClassifier(n_flies=64, epochs_per_fly=1, ...)
+```
+
+## Verbosity
+
+- `verbose=0`: no training prints
+- `verbose=1`: stage metrics
+- `verbose>=2`: per-fly epoch metrics plus control-construction diagnostics
+
+## Current scope
+
+- classifier: binary targets
+- regressor: single-output targets
+- `eval_set=[(X_val, y_val), ...]` supported
+- `evals_result()` stores stage-wise metrics
+- graph file stays external to keep the package lightweight

flyboost_connectome-0.1.0/src/flyboost_connectome.egg-info/SOURCES.txt

@@ -0,0 +1,12 @@
+README.md
+pyproject.toml
+src/flyboost/__init__.py
+src/flyboost/estimators.py
+src/flyboost/fly.py
+src/flyboost/graph.py
+src/flyboost_connectome.egg-info/PKG-INFO
+src/flyboost_connectome.egg-info/SOURCES.txt
+src/flyboost_connectome.egg-info/dependency_links.txt
+src/flyboost_connectome.egg-info/requires.txt
+src/flyboost_connectome.egg-info/top_level.txt
+tests/test_estimators.py

flyboost_connectome-0.1.0/src/flyboost_connectome.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

flyboost_connectome-0.1.0/src/flyboost_connectome.egg-info/requires.txt

@@ -0,0 +1,7 @@
+numpy>=1.26
+scikit-learn>=1.6
+torch>=2.4
+
+[dev]
+pytest>=8
+build>=1.2

flyboost_connectome-0.1.0/src/flyboost_connectome.egg-info/top_level.txt

@@ -0,0 +1 @@
+flyboost

flyboost_connectome-0.1.0/tests/test_estimators.py

@@ -0,0 +1,99 @@
+import numpy as np
+import torch
+from sklearn.base import clone
+from sklearn.pipeline import make_pipeline
+from sklearn.preprocessing import StandardScaler
+
+from flyboost import FlyBoostClassifier, FlyBoostRegressor
+
+
+def _write_toy_graph(path):
+    # Small directed graph in the same [post, pre] format as MaleCNS.
+    n = 12
+    pre = torch.tensor([0,1,2,3,4,5,6,7,8,9,10,11,0,2,4,6,8,10])
+    post = torch.tensor([1,2,3,4,5,6,7,8,9,10,11,0,6,7,8,9,10,11])
+    w = torch.tensor([
+        0.5,0.4,0.3,0.4,0.5,0.3,0.4,0.5,0.3,0.4,0.5,0.3,
+        -0.2,0.2,-0.2,0.2,-0.2,0.2,
+    ], dtype=torch.float32)
+    torch.save({
+        "num_nodes": n,
+        "edge_index": torch.stack([post, pre]),
+        "base_weight": w,
+        "sensory_idx": torch.tensor([0,1,2,3,4,5]),
+        "readout_idx": torch.tensor([8,9,10,11]),
+    }, path)
+
+
+def test_classifier_modes_and_sklearn_plumbing(tmp_path):
+    graph = tmp_path / "toy.pt"
+    _write_toy_graph(graph)
+    rng = np.random.default_rng(0)
+    X = rng.normal(size=(32, 4)).astype(np.float32)
+    y = (X[:, 0] + 0.5 * X[:, 1] > 0).astype(int)
+
+    for mode in ["real", "shuffled", "frozen"]:
+        clf = FlyBoostClassifier(
+            graph_path=str(graph),
+            n_flies=2,
+            fly_steps=2,
+            epochs_per_fly=1,
+            batch_size=8,
+            fly_lr=1e-2,
+            fly_mode=mode,
+            random_state=1,
+            verbose=0,
+        )
+        clone(clf)
+        clf.fit(X, y, eval_set=[(X, y)])
+        pred = clf.predict(X)
+        proba = clf.predict_proba(X)
+        assert pred.shape == (len(X),)
+        assert proba.shape == (len(X), 2)
+        assert np.allclose(proba.sum(axis=1), 1, atol=1e-5)
+        assert len(clf.evals_result()["validation_0"]["logloss"]) == 2
+
+
+def test_pipeline(tmp_path):
+    graph = tmp_path / "toy.pt"
+    _write_toy_graph(graph)
+    rng = np.random.default_rng(1)
+    X = rng.normal(size=(24, 3)).astype(np.float32)
+    y = (X[:, 0] > 0).astype(int)
+    pipe = make_pipeline(
+        StandardScaler(),
+        FlyBoostClassifier(
+            graph_path=str(graph),
+            n_flies=1,
+            fly_steps=1,
+            epochs_per_fly=1,
+            batch_size=8,
+            fly_lr=1e-2,
+            verbose=0,
+        ),
+    )
+    pipe.fit(X, y)
+    assert pipe.predict(X).shape == (len(X),)
+
+
+def test_regressor(tmp_path):
+    graph = tmp_path / "toy.pt"
+    _write_toy_graph(graph)
+    rng = np.random.default_rng(2)
+    X = rng.normal(size=(32, 4)).astype(np.float32)
+    y = (2 * X[:, 0] - X[:, 1] + 0.1 * rng.normal(size=32)).astype(np.float32)
+    reg = FlyBoostRegressor(
+        graph_path=str(graph),
+        n_flies=2,
+        fly_steps=2,
+        epochs_per_fly=1,
+        batch_size=8,
+        fly_lr=1e-2,
+        fly_mode="real",
+        verbose=0,
+    )
+    clone(reg)
+    reg.fit(X, y, eval_set=[(X, y)])
+    pred = reg.predict(X)
+    assert pred.shape == (len(X),)
+    assert len(reg.evals_result()["validation_0"]["rmse"]) == 2