bci: add preregistered signature metrics

2 files changed, 245 insertions, 1 deletions

diff --git a/sdil/bci.py b/sdil/bci.py
index 9847ffe..c8408e3 100644
--- a/sdil/bci.py
+++ b/sdil/bci.py
@@ -224,7 +224,8 @@ class BCISDIL:
             error = self.cfg.target - cursor
             loss = 0.5 * error.square()
 
-            did_perturb = global_step % self.cfg.perturb_every == 0
+            did_perturb = (learn_vectorizer
+                           and global_step % self.cfg.perturb_every == 0)
             causal_target = None
             if did_perturb:
                 causal_target = self.causal_targets(
@@ -270,6 +271,11 @@ def run_day(model, trajectories, day, *, control_gain=1.0,
     success = torch.zeros_like(active)
     events = []
     for step_index in range(cfg.steps_per_episode):
+        previous_soma_for_record = soma
+        previous_error_for_record = (torch.full(
+            (episodes,), abs(cfg.target), device=model.device, dtype=model.dtype)
+                                     if previous_abs_error is None
+                                     else previous_abs_error)
         record = model.step(
             trajectories.context[day, :, step_index],
             trajectories.process_noise[day, :, step_index],
@@ -288,6 +294,8 @@ def run_day(model, trajectories, day, *, control_gain=1.0,
             } | {
                 "active": active.detach().cpu().clone(),
                 "success": success.detach().cpu().clone(),
+                "previous_soma": previous_soma_for_record.detach().cpu().clone(),
+                "previous_abs_error": previous_error_for_record.detach().cpu().clone(),
                 "step": step_index,
             })
         soma = record["soma"]
diff --git a/sdil/bci_metrics.py b/sdil/bci_metrics.py
new file mode 100644
index 0000000..f786b8e
--- /dev/null
+++ b/sdil/bci_metrics.py
@@ -0,0 +1,236 @@
+"""Predetermined metrics for the preregistered continuous-BCI signatures."""
+import math
+
+import torch
+
+
+def pearson(x, y):
+    x = torch.as_tensor(x, dtype=torch.float64).flatten()
+    y = torch.as_tensor(y, dtype=torch.float64).flatten()
+    finite = torch.isfinite(x) & torch.isfinite(y)
+    x, y = x[finite], y[finite]
+    if x.numel() < 3:
+        return 0.0
+    x = x - x.mean()
+    y = y - y.mean()
+    denominator = x.square().sum().sqrt() * y.square().sum().sqrt()
+    if denominator <= 1e-15:
+        return 0.0
+    return (x @ y / denominator).item()
+
+
+def balanced_accuracy(labels, predictions):
+    labels = torch.as_tensor(labels, dtype=torch.bool)
+    predictions = torch.as_tensor(predictions, dtype=torch.bool)
+    scores = []
+    for value in (False, True):
+        mask = labels == value
+        if bool(mask.any()):
+            scores.append((predictions[mask] == labels[mask]).double().mean())
+    return torch.stack(scores).mean().item() if scores else 0.5
+
+
+def grouped_linear_predictions(features, target, groups, *, classification,
+                               n_folds=5, ridge=1e-3):
+    """Out-of-group-fold linear scores with train-fold standardization."""
+    x = torch.as_tensor(features, dtype=torch.float64)
+    y = torch.as_tensor(target, dtype=torch.float64).flatten()
+    groups = torch.as_tensor(groups, dtype=torch.long).flatten()
+    if x.ndim == 1:
+        x = x.unsqueeze(1)
+    if x.shape[0] != y.numel() or groups.numel() != y.numel():
+        raise ValueError("features, targets, and groups must have matching rows")
+    if x.shape[0] < n_folds * 2:
+        return torch.zeros_like(y)
+    prediction = torch.zeros_like(y)
+    for fold in range(n_folds):
+        test = groups.remainder(n_folds) == fold
+        train = ~test
+        if not bool(test.any()) or train.sum() <= x.shape[1] + 1:
+            continue
+        train_x, train_y = x[train], y[train]
+        mean = train_x.mean(0)
+        scale = train_x.std(0, unbiased=False).clamp_min(1e-8)
+        train_x = (train_x - mean) / scale
+        test_x = (x[test] - mean) / scale
+        train_x = torch.cat((train_x, torch.ones(train_x.shape[0], 1)), dim=1)
+        test_x = torch.cat((test_x, torch.ones(test_x.shape[0], 1)), dim=1)
+        if classification:
+            binary = train_y > 0.5
+            positives = binary.sum().item()
+            negatives = binary.numel() - positives
+            if positives == 0 or negatives == 0:
+                prediction[test] = 1.0 if positives else -1.0
+                continue
+            signed_y = binary.double().mul(2).sub(1)
+            weights = torch.where(
+                binary,
+                torch.full_like(train_y, binary.numel() / (2.0 * positives)),
+                torch.full_like(train_y, binary.numel() / (2.0 * negatives)))
+            root_weight = weights.sqrt()
+            fit_x = train_x * root_weight.unsqueeze(1)
+            fit_y = signed_y * root_weight
+        else:
+            fit_x, fit_y = train_x, train_y
+        penalty = torch.eye(fit_x.shape[1], dtype=torch.float64) * ridge
+        penalty[-1, -1] = 0.0
+        gram = fit_x.t() @ fit_x + penalty
+        rhs = fit_x.t() @ fit_y
+        try:
+            coefficient = torch.linalg.solve(gram, rhs)
+        except RuntimeError:
+            coefficient = torch.linalg.pinv(gram) @ rhs
+        prediction[test] = test_x @ coefficient
+    return prediction
+
+
+def grouped_classification_accuracy(features, labels, groups):
+    labels = torch.as_tensor(labels, dtype=torch.bool)
+    if labels.numel() < 10 or labels.all() or (~labels).all():
+        return 0.5, torch.zeros(labels.numel(), dtype=torch.float64)
+    score = grouped_linear_predictions(
+        features, labels.double(), groups, classification=True)
+    return balanced_accuracy(labels, score > 0), score
+
+
+def grouped_regression_correlation(features, target, groups):
+    prediction = grouped_linear_predictions(
+        features, target, groups, classification=False)
+    return pearson(prediction, target)
+
+
+def annotate_day_events(events, day, final_success, episode_offset):
+    """Attach stable episode groups and final outcomes to collected records."""
+    final_success = torch.as_tensor(final_success, dtype=torch.bool)
+    episode_ids = torch.arange(final_success.numel(), dtype=torch.long) + episode_offset
+    for event in events:
+        event["day"] = day
+        event["episode_ids"] = episode_ids.clone()
+        event["final_success"] = final_success.clone()
+    return events
+
+
+def flatten_active(events):
+    fields = {
+        "soma": [], "previous_soma": [], "raw": [], "innovation": [],
+        "abs_error": [], "previous_abs_error": [], "groups": [], "day": [],
+        "outcome": [],
+    }
+    for event in events:
+        active = event["active"].bool()
+        if not bool(active.any()):
+            continue
+        count = active.sum().item()
+        fields["soma"].append(event["base_soma"][active].double())
+        fields["previous_soma"].append(event["previous_soma"][active].double())
+        fields["raw"].append(event["raw_apical"][active].double())
+        fields["innovation"].append(event["innovation"][active].double())
+        fields["abs_error"].append(event["abs_error"][active].double())
+        fields["previous_abs_error"].append(
+            event["previous_abs_error"][active].double())
+        fields["groups"].append(event["episode_ids"][active].long())
+        fields["day"].append(torch.full((count,), int(event["day"]), dtype=torch.long))
+        fields["outcome"].append(event["final_success"][active].bool())
+    if not fields["soma"]:
+        raise ValueError("no active BCI events were collected")
+    return {key: torch.cat(value, dim=0) for key, value in fields.items()}
+
+
+def episode_population_features(events, field):
+    """Mean of the last four active events for every episode."""
+    values = {}
+    outcomes = {}
+    for event in events:
+        for row, group in enumerate(event["episode_ids"].tolist()):
+            if not bool(event["active"][row]):
+                continue
+            values.setdefault(group, []).append(event[field][row].double())
+            outcomes[group] = bool(event["final_success"][row])
+    groups = sorted(values)
+    features = torch.stack([
+        torch.stack(values[group][-4:]).mean(0) for group in groups])
+    labels = torch.tensor([outcomes[group] for group in groups], dtype=torch.bool)
+    return features, labels, torch.tensor(groups, dtype=torch.long)
+
+
+def population_correlations(flat, left, right):
+    return [pearson(flat[left][:, cell], flat[right][:, cell])
+            for cell in range(flat[left].shape[1])]
+
+
+def signature_metrics(train_events, evaluation_events, cfg, role):
+    train = flatten_active(train_events)
+    role = torch.as_tensor(role, dtype=torch.float64)
+    soma_corr = population_correlations(train, "innovation", "soma")
+    raw_corr = population_correlations(train, "raw", "soma")
+
+    # Surrounding-network prediction of each causal cell's amplification sign.
+    event_decoder, confidence_corr = [], []
+    causal_cells = range(cfg.n_plus + cfg.n_minus)
+    for cell in causal_cells:
+        keep = [index for index in range(cfg.n_neurons) if index != cell]
+        labels = train["innovation"][:, cell] > 0
+        score, distance = grouped_classification_accuracy(
+            train["previous_soma"][:, keep], labels, train["groups"])
+        event_decoder.append(score)
+        confidence_corr.append(pearson(distance, train["innovation"][:, cell]))
+
+    diff = (train["innovation"][:, :cfg.n_plus].mean(1)
+            - train["innovation"][:, cfg.n_plus:cfg.n_plus + cfg.n_minus].mean(1))
+    improvement = train["previous_abs_error"] - train["abs_error"]
+    improving = improvement > 0
+    worsening = improvement < 0
+    sign_index = 0.0
+    if bool(improving.any()) and bool(worsening.any()):
+        sign_index = 0.5 * (diff[improving].mean() - diff[worsening].mean()).item()
+
+    role_aligned = train["innovation"] @ role
+    velocity_corr = grouped_regression_correlation(
+        role_aligned, improvement, train["groups"])
+    error_corr = grouped_regression_correlation(
+        role_aligned, train["abs_error"], train["groups"])
+
+    residual_features, outcomes, outcome_groups = episode_population_features(
+        evaluation_events, "innovation")
+    soma_features, soma_outcomes, soma_groups = episode_population_features(
+        evaluation_events, "base_soma")
+    if not torch.equal(outcomes, soma_outcomes) or not torch.equal(outcome_groups, soma_groups):
+        raise ValueError("residual and soma outcome populations are not paired")
+    residual_outcome, _ = grouped_classification_accuracy(
+        residual_features, outcomes, outcome_groups)
+    soma_outcome, _ = grouped_classification_accuracy(
+        soma_features, outcomes, outcome_groups)
+
+    early = train["day"] <= 2
+    late = train["day"] >= cfg.days - 3
+    if not bool(early.any()) or not bool(late.any()):
+        longitudinal = 0.0
+    else:
+        early_residual = train["innovation"][early].mean(0)
+        activity_change = (train["soma"][late].mean(0)
+                           - train["soma"][early].mean(0))
+        longitudinal = pearson(early_residual, activity_change)
+
+    values = {
+        "mean_abs_residual_soma_corr": sum(abs(v) for v in soma_corr) / len(soma_corr),
+        "mean_abs_raw_soma_corr": sum(abs(v) for v in raw_corr) / len(raw_corr),
+        "raw_minus_residual_abs_soma_corr": (
+            sum(abs(v) for v in raw_corr) / len(raw_corr)
+            - sum(abs(v) for v in soma_corr) / len(soma_corr)),
+        "surrounding_event_decoder_balanced_acc": sum(event_decoder) / len(event_decoder),
+        "decoder_distance_residual_corr": sum(confidence_corr) / len(confidence_corr),
+        "residual_outcome_decoder_balanced_acc": residual_outcome,
+        "soma_outcome_decoder_balanced_acc": soma_outcome,
+        "residual_minus_soma_outcome_acc": residual_outcome - soma_outcome,
+        "causal_role_sign_inversion_index": sign_index,
+        "role_aligned_velocity_cv_corr": velocity_corr,
+        "role_aligned_error_cv_corr": error_corr,
+        "velocity_minus_error_abs_cv_corr": abs(velocity_corr) - abs(error_corr),
+        "early_residual_late_activity_change_corr": longitudinal,
+        "active_training_events": int(train["soma"].shape[0]),
+        "evaluation_episodes": int(outcomes.numel()),
+        "evaluation_success_fraction": outcomes.double().mean().item(),
+    }
+    if not all(math.isfinite(value) for value in values.values()):
+        raise ValueError("non-finite BCI signature metric")
+    return values