| |
| |
| import json |
| from collections import defaultdict |
| from copy import deepcopy |
| from hashlib import sha256 |
| from pathlib import Path |
|
|
| import h5py |
| import numpy as np |
| import torch |
| from geomloss import SamplesLoss |
| from scipy.stats import wilcoxon |
| from torch.utils.data import DataLoader |
|
|
| from causalcelljepa.dynamics import ( |
| LatentPopulationDataset, |
| anchored_selected_entry, |
| build_dynamics_model, |
| ) |
| from causalcelljepa.resources import file_sha256 |
| from causalcelljepa.training import _git_state, _runtime_environment, _runtime_source_hash |
|
|
|
|
| def _role_effects(config, roles): |
| inputs = config["inputs"] |
| manifest = json.loads(Path(inputs["dynamics_manifest_path"]).read_text()) |
| normalization = manifest["normalization"] |
| mean = np.asarray(normalization["latent_mean"], dtype=np.float32) |
| scale = np.asarray(normalization["latent_std"], dtype=np.float32) * normalization[ |
| "dimension_scale" |
| ] |
| with h5py.File(inputs["latent_cache_path"], "r") as cache: |
| cache_roles = cache["role"].asstr()[:] |
| targets = cache["target"].asstr()[:] |
| batches = cache["source_batch"].asstr()[:] |
| controls = np.flatnonzero(cache_roles == "control_train") |
| control_means = { |
| batch: ((cache["latent"][controls[batches[controls] == batch]] - mean) / scale).mean(0) |
| for batch in sorted(set(batches[controls])) |
| } |
| effects = {} |
| for role in roles: |
| role_indices = np.flatnonzero(cache_roles == role) |
| effects[role] = {} |
| for target in sorted(set(targets[role_indices])): |
| indices = role_indices[targets[role_indices] == target] |
| outcome = ((cache["latent"][indices] - mean) / scale).mean(0) |
| counts = {batch: np.count_nonzero(batches[indices] == batch) for batch in set(batches[indices])} |
| control = sum(count * control_means[batch] for batch, count in counts.items()) / len(indices) |
| effects[role][target] = outcome - control |
| return effects, normalization |
|
|
|
|
| def fit_linear_baseline(config): |
| """Fit a low-rank ESM-to-effect ridge baseline and tune alpha on K562 validation.""" |
| effects, _ = _role_effects(config, ("dynamics_train", "perturbation_ood_validation")) |
| action = torch.load(config["inputs"]["action_cache_path"], map_location="cpu", weights_only=True) |
| action_map = { |
| target: (action["embedding"][index].numpy(), bool(action["known"][index])) |
| for index, target in enumerate(action["targets"]) |
| } |
| train = effects["dynamics_train"] |
| known_targets = [target for target in sorted(train) if action_map[target][1]] |
| x = np.stack([action_map[target][0] for target in known_targets]).astype(np.float64) |
| y = np.stack([train[target] for target in known_targets]).astype(np.float64) |
| x_mean, x_std = x.mean(0), x.std(0).clip(1e-8) |
| y_mean = y.mean(0) |
| _, _, components = np.linalg.svd(y - y_mean, full_matrices=False) |
| components = components[: config["baselines"]["linear_rank"]] |
| scores = (y - y_mean) @ components.T |
| standardized = (x - x_mean) / x_std |
| gram, cross = standardized.T @ standardized, standardized.T @ scores |
| validation = effects["perturbation_ood_validation"] |
| validation_targets = sorted(validation) |
| validation_x = np.stack([action_map[target][0] for target in validation_targets]) |
| validation_y = np.stack([validation[target] for target in validation_targets]) |
| candidates = [] |
| for alpha in config["baselines"]["ridge_candidates"]: |
| weights = np.linalg.solve(gram + alpha * np.eye(gram.shape[0]), cross) |
| prediction = ((validation_x - x_mean) / x_std) @ weights @ components + y_mean |
| candidates.append((float(np.mean((prediction - validation_y) ** 2)), alpha, weights)) |
| validation_mse, alpha, weights = min(candidates, key=lambda item: (item[0], item[1])) |
| predictions = { |
| target: ( |
| ((embedding - x_mean) / x_std) @ weights @ components + y_mean |
| if known |
| else y_mean |
| ).astype(np.float32) |
| for target, (embedding, known) in action_map.items() |
| } |
| return predictions, np.mean(np.stack(list(train.values())), axis=0).astype(np.float32), { |
| "fit_outcome_role": "dynamics_train", |
| "selection_outcome_role": "perturbation_ood_validation", |
| "fit_targets": len(train), |
| "fit_targets_with_known_action": len(known_targets), |
| "selection_targets": len(validation_targets), |
| "rank": len(components), |
| "selected_ridge": alpha, |
| "selection_mse": validation_mse, |
| "ridge_candidates": [float(item[1]) for item in candidates], |
| "ridge_validation_mse": [float(item[0]) for item in candidates], |
| } |
|
|
|
|
| def population_metrics(predicted, observed, control, median_distance, config): |
| """Return one value per perturbation-condition for all locked latent metrics.""" |
| sinkhorn = SamplesLoss( |
| "sinkhorn", |
| p=2, |
| blur=config["reference_sinkhorn_blur_ratio"] * median_distance, |
| debias=True, |
| backend="tensorized", |
| )(predicted, observed).reshape(-1) |
| bandwidth = config["mmd_bandwidth_ratio"] * median_distance |
| xx = torch.exp(-torch.cdist(predicted, predicted).square() / (2 * bandwidth**2)).mean((1, 2)) |
| yy = torch.exp(-torch.cdist(observed, observed).square() / (2 * bandwidth**2)).mean((1, 2)) |
| xy = torch.exp(-torch.cdist(predicted, observed).square() / (2 * bandwidth**2)).mean((1, 2)) |
| true_effect = observed.mean(1) - control.mean(1) |
| predicted_effect = predicted.mean(1) - control.mean(1) |
| true_magnitude = torch.linalg.vector_norm(true_effect, dim=1) |
| predicted_magnitude = torch.linalg.vector_norm(predicted_effect, dim=1) |
| centered_true = true_effect - true_effect.mean(1, keepdim=True) |
| centered_predicted = predicted_effect - predicted_effect.mean(1, keepdim=True) |
| denominator = torch.linalg.vector_norm(centered_true, dim=1) * torch.linalg.vector_norm( |
| centered_predicted, dim=1 |
| ) |
| pearson = (centered_true * centered_predicted).sum(1) / denominator |
| true_rank = true_effect.argsort(1).argsort(1).float() |
| predicted_rank = predicted_effect.argsort(1).argsort(1).float() |
| true_rank -= true_rank.mean(1, keepdim=True) |
| predicted_rank -= predicted_rank.mean(1, keepdim=True) |
| rank_denominator = torch.linalg.vector_norm(true_rank, dim=1) * torch.linalg.vector_norm( |
| predicted_rank, dim=1 |
| ) |
| spearman = (true_rank * predicted_rank).sum(1) / rank_denominator |
| invalid = (denominator <= 1e-12) | (predicted_effect.std(1) <= 1e-12) |
| pearson[invalid], spearman[invalid] = torch.nan, torch.nan |
| predicted_covariance = torch.bmm( |
| (predicted - predicted.mean(1, keepdim=True)).transpose(1, 2), |
| predicted - predicted.mean(1, keepdim=True), |
| ) / (predicted.shape[1] - 1) |
| observed_covariance = torch.bmm( |
| (observed - observed.mean(1, keepdim=True)).transpose(1, 2), |
| observed - observed.mean(1, keepdim=True), |
| ) / (observed.shape[1] - 1) |
| return { |
| "effect_pearson": pearson, |
| "effect_spearman": spearman, |
| "direction_cosine": torch.nn.functional.cosine_similarity(predicted_effect, true_effect), |
| "magnitude_ratio": predicted_magnitude / true_magnitude.clamp_min(1e-12), |
| "magnitude_absolute_error": torch.abs(predicted_magnitude - true_magnitude), |
| "sinkhorn": sinkhorn, |
| "mmd": (xx + yy - 2 * xy).clamp_min(0), |
| "energy_distance": ( |
| 2 * torch.cdist(predicted, observed).mean((1, 2)) |
| - torch.cdist(predicted, predicted).mean((1, 2)) |
| - torch.cdist(observed, observed).mean((1, 2)) |
| ).clamp_min(0), |
| "centroid_shift_error": torch.linalg.vector_norm(predicted_effect - true_effect, dim=1), |
| "covariance_shift_error": torch.linalg.matrix_norm( |
| predicted_covariance - observed_covariance, ord="fro" |
| ), |
| } |
|
|
|
|
| def paired_condition_comparisons(records, bootstrap_resamples, seed): |
| """Compare models on matched targets with paired bootstrap CIs and Wilcoxon tests.""" |
| excluded = {"regime", "context", "outcome_role", "target", "repeat", "model"} |
| metric_names = sorted(set(records[0]) - excluded) |
| higher = {"effect_pearson", "effect_spearman", "direction_cosine"} |
| closer = {"magnitude_ratio"} |
| target_values = defaultdict(list) |
| for record in records: |
| for metric in metric_names: |
| if record[metric] is not None: |
| target_values[record["regime"], record["model"], record["target"], metric].append( |
| record[metric] |
| ) |
| averaged = {key: float(np.mean(values)) for key, values in target_values.items()} |
| comparisons = [] |
| regimes = sorted({record["regime"] for record in records}) |
| for regime in regimes: |
| for baseline in ("no_change", "mean_effect", "linear_esm", "pseudo_paired"): |
| for metric in metric_names: |
| targets = sorted( |
| { |
| key[2] |
| for key in averaged |
| if key[0] == regime |
| and key[1] == "causalcelljepa" |
| and key[3] == metric |
| and (regime, baseline, key[2], metric) in averaged |
| } |
| ) |
| if not targets: |
| continue |
| model = np.asarray( |
| [averaged[regime, "causalcelljepa", target, metric] for target in targets] |
| ) |
| reference = np.asarray( |
| [averaged[regime, baseline, target, metric] for target in targets] |
| ) |
| if metric in higher: |
| improvement, direction = model - reference, "higher_is_better" |
| elif metric in closer: |
| improvement, direction = ( |
| np.abs(reference - 1) - np.abs(model - 1), |
| "closer_to_one_is_better", |
| ) |
| else: |
| improvement, direction = reference - model, "lower_is_better" |
| generator = np.random.default_rng( |
| (int.from_bytes(sha256(f"{regime}\0{baseline}\0{metric}".encode()).digest()[:8], "little") + seed) |
| % (1 << 64) |
| ) |
| bootstrap = improvement[ |
| generator.integers(0, len(improvement), (bootstrap_resamples, len(improvement))) |
| ].mean(1) |
| p_value = ( |
| 1.0 |
| if np.all(improvement == 0) |
| else float(wilcoxon(improvement, alternative="two-sided", method="approx").pvalue) |
| ) |
| comparisons.append( |
| { |
| "regime": regime, |
| "baseline": baseline, |
| "metric": metric, |
| "direction": direction, |
| "targets": len(targets), |
| "causalcelljepa_mean": float(model.mean()), |
| "baseline_mean": float(reference.mean()), |
| "mean_improvement": float(improvement.mean()), |
| "median_improvement": float(np.median(improvement)), |
| "mean_improvement_bootstrap_95ci": [ |
| float(value) for value in np.quantile(bootstrap, (0.025, 0.975)) |
| ], |
| "wilcoxon_two_sided_p": p_value, |
| } |
| ) |
| p_values = np.asarray([comparison["wilcoxon_two_sided_p"] for comparison in comparisons]) |
| order = np.argsort(p_values) |
| ranked = p_values[order] * len(p_values) / np.arange(1, len(p_values) + 1) |
| adjusted = np.minimum.accumulate(ranked[::-1])[::-1].clip(max=1) |
| q_values = np.empty_like(adjusted) |
| q_values[order] = adjusted |
| for comparison, q_value in zip(comparisons, q_values): |
| comparison["benjamini_hochberg_q"] = float(q_value) |
| return comparisons |
|
|
|
|
| def paired_model_comparisons(records, model_pairs, bootstrap_resamples, seed): |
| """Compare arbitrary learned models on matched perturbation-condition units.""" |
| excluded = {"regime", "context", "outcome_role", "target", "repeat", "model"} |
| metric_names = sorted(set(records[0]) - excluded) |
| higher = {"effect_pearson", "effect_spearman", "direction_cosine"} |
| closer = {"magnitude_ratio"} |
| target_values = defaultdict(list) |
| for record in records: |
| for metric in metric_names: |
| if record[metric] is not None: |
| target_values[record["regime"], record["model"], record["target"], metric].append( |
| record[metric] |
| ) |
| averaged = {key: float(np.mean(values)) for key, values in target_values.items()} |
| comparisons = [] |
| regimes = sorted({record["regime"] for record in records}) |
| for regime in regimes: |
| for pair in model_pairs: |
| candidate, reference = pair["candidate"], pair["reference"] |
| for metric in metric_names: |
| targets = sorted( |
| { |
| key[2] |
| for key in averaged |
| if key[0] == regime |
| and key[1] == candidate |
| and key[3] == metric |
| and (regime, reference, key[2], metric) in averaged |
| } |
| ) |
| if not targets: |
| continue |
| candidate_values = np.asarray( |
| [averaged[regime, candidate, target, metric] for target in targets] |
| ) |
| reference_values = np.asarray( |
| [averaged[regime, reference, target, metric] for target in targets] |
| ) |
| if metric in higher: |
| improvement = candidate_values - reference_values |
| direction = "higher_is_better" |
| elif metric in closer: |
| improvement = np.abs(reference_values - 1) - np.abs(candidate_values - 1) |
| direction = "closer_to_one_is_better" |
| else: |
| improvement = reference_values - candidate_values |
| direction = "lower_is_better" |
| key = f"{regime}\0{candidate}\0{reference}\0{metric}" |
| generator = np.random.default_rng( |
| (int.from_bytes(sha256(key.encode()).digest()[:8], "little") + seed) |
| % (1 << 64) |
| ) |
| bootstrap = improvement[ |
| generator.integers( |
| 0, len(improvement), (bootstrap_resamples, len(improvement)) |
| ) |
| ].mean(1) |
| p_value = ( |
| 1.0 |
| if np.all(improvement == 0) |
| else float( |
| wilcoxon(improvement, alternative="two-sided", method="approx").pvalue |
| ) |
| ) |
| comparisons.append( |
| { |
| "regime": regime, |
| "candidate": candidate, |
| "reference": reference, |
| "hypothesis": pair["hypothesis"], |
| "metric": metric, |
| "direction": direction, |
| "targets": len(targets), |
| "candidate_mean": float(candidate_values.mean()), |
| "reference_mean": float(reference_values.mean()), |
| "mean_improvement": float(improvement.mean()), |
| "median_improvement": float(np.median(improvement)), |
| "mean_improvement_bootstrap_95ci": [ |
| float(value) for value in np.quantile(bootstrap, (0.025, 0.975)) |
| ], |
| "wilcoxon_two_sided_p": p_value, |
| } |
| ) |
| if not comparisons: |
| return comparisons |
| p_values = np.asarray([comparison["wilcoxon_two_sided_p"] for comparison in comparisons]) |
| order = np.argsort(p_values) |
| ranked = p_values[order] * len(p_values) / np.arange(1, len(p_values) + 1) |
| adjusted = np.minimum.accumulate(ranked[::-1])[::-1].clip(max=1) |
| q_values = np.empty_like(adjusted) |
| q_values[order] = adjusted |
| for comparison, q_value in zip(comparisons, q_values): |
| comparison["benjamini_hochberg_q"] = float(q_value) |
| return comparisons |
|
|
|
|
| def _condition_metric_summaries(records, bootstrap_resamples, seed): |
| excluded = {"regime", "context", "outcome_role", "target", "repeat", "model"} |
| metric_names = sorted(set(records[0]) - excluded) |
| target_values = defaultdict(list) |
| for record in records: |
| for metric in metric_names: |
| if record[metric] is not None: |
| target_values[record["regime"], record["model"], record["target"], metric].append( |
| record[metric] |
| ) |
| grouped = defaultdict(list) |
| for (regime, model, _target, metric), values in target_values.items(): |
| grouped[regime, model, metric].append(float(np.mean(values))) |
| summaries = [] |
| for key, values in sorted(grouped.items()): |
| regime, model, metric = key |
| values = np.asarray(values) |
| generator = np.random.default_rng( |
| int.from_bytes(sha256("\0".join(key).encode()).digest()[:8], "little") + seed |
| ) |
| bootstrap = values[ |
| generator.integers(0, len(values), (bootstrap_resamples, len(values))) |
| ].mean(1) |
| summaries.append( |
| { |
| "regime": regime, |
| "model": model, |
| "metric": metric, |
| "targets": len(values), |
| "mean": float(values.mean()), |
| "median": float(np.median(values)), |
| "mean_bootstrap_95ci": [ |
| float(value) for value in np.quantile(bootstrap, (0.025, 0.975)) |
| ], |
| } |
| ) |
| return summaries |
|
|
|
|
| def _retrieval_summaries(signatures, regimes, models): |
| summaries = [] |
| for regime in regimes: |
| targets = sorted(key[2] for key in signatures if key[:2] == (regime, "true")) |
| truth = np.stack( |
| [signatures[regime, "true", target][1] / signatures[regime, "true", target][0] |
| for target in targets] |
| ) |
| truth /= np.linalg.norm(truth, axis=1, keepdims=True).clip(1e-12) |
| for model in models: |
| predicted = np.stack( |
| [signatures[regime, model, target][1] / signatures[regime, model, target][0] |
| for target in targets] |
| ) |
| predicted /= np.linalg.norm(predicted, axis=1, keepdims=True).clip(1e-12) |
| similarity = predicted @ truth.T |
| ranks = np.asarray( |
| [ |
| 1 |
| + np.count_nonzero(row > row[index]) |
| + 0.5 * (np.count_nonzero(row == row[index]) - 1) |
| for index, row in enumerate(similarity) |
| ] |
| ) |
| summaries.append( |
| { |
| "regime": regime, |
| "model": model, |
| "targets": len(targets), |
| "top_1": float(np.mean(ranks <= 1)), |
| "top_5": float(np.mean(ranks <= 5)), |
| "mean_reciprocal_rank": float(np.mean(1 / ranks)), |
| "median_rank": float(np.median(ranks)), |
| } |
| ) |
| return summaries |
|
|
|
|
| def _self_hashed_manifest(path, expected_sha256): |
| payload = json.loads(Path(path).read_text()) |
| declared = payload.pop("manifest_sha256") |
| computed = sha256( |
| json.dumps(payload, sort_keys=True, separators=(",", ":")).encode() |
| ).hexdigest() |
| assert declared == expected_sha256 == computed |
| return payload, declared |
|
|
|
|
| @torch.no_grad() |
| def run_evaluation(config, regimes=None, repeats=None, max_conditions=None, output_directory=None): |
| """Run deterministic population evaluation and summarize at the target-condition level.""" |
| inputs = config["inputs"] |
| for kind in ("latent_cache", "action_cache", "checkpoint", "pseudo_paired_checkpoint"): |
| assert file_sha256(inputs[f"{kind}_path"]) == inputs[f"{kind}_sha256"] |
| selection = json.loads(Path(inputs["selection_manifest_path"]).read_text()) |
| declared = selection.pop("manifest_sha256") |
| assert declared == inputs["selection_manifest_sha256"] == sha256( |
| json.dumps(selection, sort_keys=True, separators=(",", ":")).encode() |
| ).hexdigest() |
| checkpoint = torch.load(inputs["checkpoint_path"], map_location="cpu", weights_only=False) |
| assert checkpoint["configuration"]["loss"]["sinkhorn_blur_ratio"] == selection["selected"][ |
| "blur_ratio" |
| ] |
| model = build_dynamics_model(checkpoint["configuration"]).eval() |
| model.load_state_dict(checkpoint["model"]) |
| pseudo_manifest = json.loads(Path(inputs["pseudo_paired_manifest_path"]).read_text()) |
| pseudo_declared = pseudo_manifest.pop("manifest_sha256") |
| assert pseudo_declared == inputs["pseudo_paired_manifest_sha256"] == sha256( |
| json.dumps(pseudo_manifest, sort_keys=True, separators=(",", ":")).encode() |
| ).hexdigest() |
| assert ( |
| pseudo_manifest["artifacts"]["best_checkpoint"]["sha256"] |
| == inputs["pseudo_paired_checkpoint_sha256"] |
| ) |
| pseudo_checkpoint = torch.load( |
| inputs["pseudo_paired_checkpoint_path"], map_location="cpu", weights_only=False |
| ) |
| assert pseudo_checkpoint["configuration"]["objective"] == "pseudo_paired_mse" |
| pseudo_model = build_dynamics_model(pseudo_checkpoint["configuration"]).eval() |
| pseudo_model.load_state_dict(pseudo_checkpoint["model"]) |
| linear_effect, mean_effect, baseline_report = fit_linear_baseline(config) |
| dynamics_manifest = json.loads(Path(inputs["dynamics_manifest_path"]).read_text()) |
| median_distance = dynamics_manifest["normalization"]["median_training_latent_distance"] |
| regimes = regimes or config["regimes"] |
| repeats = repeats or config["sampling"]["repeats"] |
| output = Path(output_directory or config["output_directory"]) |
| assert not output.exists() |
| output.mkdir(parents=True) |
| records, signatures = [], defaultdict(lambda: [0, None]) |
| for regime, specification in regimes.items(): |
| dataset = LatentPopulationDataset( |
| inputs["latent_cache_path"], |
| inputs["action_cache_path"], |
| inputs["dynamics_manifest_path"], |
| regime, |
| config["sampling"]["population_size"], |
| config["seed"], |
| specification["outcome_role"], |
| specification["control_role"], |
| specification["context"], |
| ) |
| indices = range(min(len(dataset), max_conditions or len(dataset))) |
| for repeat in range(repeats): |
| dataset.set_epoch(repeat) |
| loader = DataLoader( |
| dataset, |
| batch_size=config["sampling"]["batch_size"], |
| sampler=list(indices), |
| num_workers=config["sampling"]["num_workers"], |
| ) |
| for batch in loader: |
| control, observed = batch["control"], batch["perturbed"] |
| predictions = { |
| "causalcelljepa": model(control, batch["action"], batch["action_known"]), |
| "pseudo_paired": pseudo_model( |
| control, batch["action"], batch["action_known"] |
| ), |
| "no_change": control, |
| "mean_effect": control + torch.from_numpy(mean_effect)[None, None], |
| "linear_esm": control |
| + torch.from_numpy(np.stack([linear_effect[target] for target in batch["target"]]))[ |
| :, None |
| ], |
| } |
| assert all(predicted.dtype == observed.dtype for predicted in predictions.values()) |
| true_effect = observed.mean(1) - control.mean(1) |
| for index, target in enumerate(batch["target"]): |
| key = (regime, "true", target) |
| signatures[key][0] += 1 |
| value = true_effect[index].numpy() |
| signatures[key][1] = ( |
| value if signatures[key][1] is None else signatures[key][1] + value |
| ) |
| for baseline, predicted in predictions.items(): |
| predicted_effect = (predicted.mean(1) - control.mean(1)).detach() |
| metrics = population_metrics( |
| predicted, observed, control, median_distance, config["metrics"] |
| ) |
| for index, target in enumerate(batch["target"]): |
| record = { |
| "regime": regime, |
| "context": specification["context"], |
| "outcome_role": specification["outcome_role"], |
| "target": target, |
| "repeat": repeat, |
| "model": baseline, |
| } |
| for name, values in metrics.items(): |
| value = float(values[index].detach()) |
| record[name] = value if np.isfinite(value) else None |
| records.append(record) |
| key = (regime, baseline, target) |
| signatures[key][0] += 1 |
| value = predicted_effect[index].numpy() |
| signatures[key][1] = ( |
| value if signatures[key][1] is None else signatures[key][1] + value |
| ) |
| with (output / "condition_metrics.jsonl").open("w") as handle: |
| for record in records: |
| handle.write(json.dumps(record, sort_keys=True) + "\n") |
| target_values = defaultdict(list) |
| metric_names = [key for key in records[0] if key not in {"regime", "context", "outcome_role", "target", "repeat", "model"}] |
| for record in records: |
| for metric in metric_names: |
| if record[metric] is not None: |
| target_values[record["regime"], record["model"], record["target"], metric].append( |
| record[metric] |
| ) |
| summary = {"condition_metrics": [], "retrieval": []} |
| grouped = defaultdict(list) |
| for (regime, baseline, target, metric), values in target_values.items(): |
| grouped[regime, baseline, metric].append(float(np.mean(values))) |
| for key, values in sorted(grouped.items()): |
| regime, baseline, metric = key |
| values = np.asarray(values) |
| generator = np.random.default_rng( |
| int.from_bytes(sha256("\0".join(key).encode()).digest()[:8], "little") + config["seed"] |
| ) |
| bootstrap = values[ |
| generator.integers(0, len(values), (config["metrics"]["bootstrap_resamples"], len(values))) |
| ].mean(1) |
| summary["condition_metrics"].append( |
| { |
| "regime": regime, |
| "model": baseline, |
| "metric": metric, |
| "targets": len(values), |
| "mean": float(values.mean()), |
| "median": float(np.median(values)), |
| "mean_bootstrap_95ci": [float(x) for x in np.quantile(bootstrap, (0.025, 0.975))], |
| } |
| ) |
| for regime in regimes: |
| targets = sorted({key[2] for key in signatures if key[0] == regime}) |
| truth = np.stack([signatures[regime, "true", target][1] / signatures[regime, "true", target][0] for target in targets]) |
| truth /= np.linalg.norm(truth, axis=1, keepdims=True).clip(1e-12) |
| for baseline in sorted({record["model"] for record in records}): |
| predicted = np.stack([signatures[regime, baseline, target][1] / signatures[regime, baseline, target][0] for target in targets]) |
| predicted /= np.linalg.norm(predicted, axis=1, keepdims=True).clip(1e-12) |
| similarity = predicted @ truth.T |
| ranks = np.asarray( |
| [ |
| 1 + np.count_nonzero(row > row[index]) + 0.5 * (np.count_nonzero(row == row[index]) - 1) |
| for index, row in enumerate(similarity) |
| ] |
| ) |
| summary["retrieval"].append( |
| { |
| "regime": regime, |
| "model": baseline, |
| "targets": len(targets), |
| "top_1": float(np.mean(ranks <= 1)), |
| "top_5": float(np.mean(ranks <= 5)), |
| "mean_reciprocal_rank": float(np.mean(1 / ranks)), |
| "median_rank": float(np.median(ranks)), |
| } |
| ) |
| provenance = { |
| "config": deepcopy(config), |
| "executed_regimes": regimes, |
| "executed_repeats": repeats, |
| "maximum_conditions_per_regime": max_conditions, |
| "checkpoint_provenance": checkpoint["provenance"], |
| "pseudo_paired_checkpoint_provenance": pseudo_checkpoint["provenance"], |
| "file_sha256": { |
| kind: inputs[f"{kind}_sha256"] |
| for kind in ("latent_cache", "action_cache", "checkpoint", "pseudo_paired_checkpoint") |
| }, |
| "selection_manifest_sha256": declared, |
| "pseudo_paired_manifest_sha256": pseudo_declared, |
| "runtime_source_sha256": _runtime_source_hash(), |
| "runtime_environment": _runtime_environment(), |
| "git": _git_state(), |
| } |
| (output / "summary.json").write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n") |
| (output / "baseline_report.json").write_text(json.dumps(baseline_report, indent=2, sort_keys=True) + "\n") |
| (output / "provenance.json").write_text(json.dumps(provenance, indent=2, sort_keys=True) + "\n") |
| return summary, baseline_report, provenance |
|
|
|
|
| @torch.no_grad() |
| def run_ablation_evaluation( |
| config, |
| base_config, |
| regimes=None, |
| repeats=None, |
| max_conditions=None, |
| output_directory=None, |
| ): |
| """Evaluate frozen mechanism ablations under the already locked four-regime protocol.""" |
| inputs = config["inputs"] |
| assert file_sha256(config["base_evaluation_config_path"]) == config[ |
| "base_evaluation_config_sha256" |
| ] |
| base_manifest, base_manifest_sha256 = _self_hashed_manifest( |
| inputs["base_evaluation_manifest_path"], |
| inputs["base_evaluation_manifest_sha256"], |
| ) |
| model_source = config.get("model_source", "ablation") |
| manifest_kinds = { |
| "ablation": "ablation", |
| "stage2_replication": "replication_training", |
| "anchored": "anchored_training", |
| "multiteacher": "multiteacher_training", |
| "contextual_multiteacher": "contextual_multiteacher_training", |
| } |
| assert model_source in manifest_kinds |
| manifest_kind = manifest_kinds[model_source] |
| model_manifest, model_manifest_sha256 = _self_hashed_manifest( |
| inputs[f"{manifest_kind}_manifest_path"], |
| inputs[f"{manifest_kind}_manifest_sha256"], |
| ) |
| gate_manifest_sha256 = None |
| if "residual_gate" in config: |
| gate = config["residual_gate"] |
| gate_manifest, gate_manifest_sha256 = _self_hashed_manifest( |
| gate["manifest_path"], gate["manifest_sha256"] |
| ) |
| assert gate_manifest["artifact"] == { |
| key: gate[key] for key in ("path", "bytes", "sha256") |
| } |
| assert gate_manifest["leakage"] == { |
| "roles_read": ["control_train"], "contexts_read": ["K562"], |
| "perturbed_outcomes_used": False, "rpe1_cells_used": False, |
| } |
| selected_candidate = selected_entry = selection_manifest_sha256 = None |
| frozen_selection_sources = {"anchored", "multiteacher", "contextual_multiteacher"} |
| if model_source in frozen_selection_sources: |
| selection_manifest, selection_manifest_sha256 = _self_hashed_manifest( |
| inputs[f"{model_source}_selection_manifest_path"], |
| inputs[f"{model_source}_selection_manifest_sha256"], |
| ) |
| selected_candidate, selected_entry = anchored_selected_entry( |
| model_manifest, selection_manifest |
| ) |
| expected = ( |
| "anchored_control_ood_gated" |
| if "residual_gate" in config |
| else f"{model_source}_selected" |
| ) |
| assert config["models"] == [expected] |
| for kind in ("base_condition_metrics", "base_summary", "base_provenance"): |
| path = Path(inputs[f"{kind}_path"]) |
| assert path.stat().st_size == inputs[f"{kind}_bytes"] |
| assert file_sha256(path) == inputs[f"{kind}_sha256"] |
| assert ( |
| base_manifest["artifacts"]["condition_metrics"]["sha256"] |
| == inputs["base_condition_metrics_sha256"] |
| ) |
| assert ( |
| base_manifest["artifacts"]["condition_metrics"]["records"] |
| == inputs["base_condition_metrics_records"] |
| ) |
| assert base_manifest["artifacts"]["summary"]["sha256"] == inputs["base_summary_sha256"] |
| assert ( |
| base_manifest["artifacts"]["provenance"]["sha256"] |
| == inputs["base_provenance_sha256"] |
| ) |
| base_provenance = json.loads(Path(inputs["base_provenance_path"]).read_text()) |
| assert base_provenance["config"] == base_config |
| assert base_provenance["git"]["dirty"] is False |
|
|
| models, model_configs, checkpoint_provenance = {}, {}, {} |
| for name in config["models"]: |
| if model_source == "ablation": |
| entry = model_manifest["experiments"][name] |
| elif model_source == "stage2_replication": |
| seed = str(config["model_seeds"][name]) |
| entry = model_manifest["artifacts"]["seeds"][seed] |
| else: |
| entry = selected_entry |
| artifact = entry["best_checkpoint"] |
| checkpoint_path = Path(artifact["path"]) |
| assert checkpoint_path.stat().st_size == artifact["bytes"] |
| assert file_sha256(checkpoint_path) == artifact["sha256"] |
| checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False) |
| if model_source == "ablation": |
| assert checkpoint["configuration"]["ablation"]["name"] == name |
| assert checkpoint["configuration"]["model"]["context_mode"] == entry[ |
| "mechanism" |
| ]["context_mode"] |
| assert ( |
| float(checkpoint["configuration"]["loss"]["weights"]["direction"]) |
| == entry["mechanism"]["direction_weight"] |
| ) |
| elif model_source == "stage2_replication": |
| assert checkpoint["configuration"]["seed"] == int(seed) |
| assert checkpoint["configuration"]["replication"] == { |
| "model_and_sampling_seed": int(seed), |
| "target_split_seed": base_config["seed"], |
| "base_config_path": "configs/dynamics.yaml", |
| "base_config_sha256": model_manifest["source"]["base_config_sha256"], |
| } |
| assert checkpoint["state"]["best_validation_epoch"] == entry["full_run"][ |
| "best_validation_epoch" |
| ] |
| assert checkpoint["state"]["best_validation_loss"] == entry["full_run"][ |
| "best_validation_loss" |
| ] |
| else: |
| assert checkpoint["configuration"]["revision"]["candidate"] == selected_candidate |
| assert checkpoint["state"]["best_validation_epoch"] == entry["full_run"][ |
| "best_validation_epoch" |
| ] |
| assert checkpoint["state"]["best_validation_loss"] == entry["full_run"][ |
| "best_validation_loss" |
| ] |
| assert checkpoint["provenance"]["git"]["dirty"] is False |
| model = build_dynamics_model(checkpoint["configuration"]).eval() |
| model.load_state_dict(checkpoint["model"]) |
| if "residual_gate" in config: |
| gate = config["residual_gate"] |
| gate_path = Path(gate["path"]) |
| assert (gate_path.stat().st_size, file_sha256(gate_path)) == ( |
| gate["bytes"], gate["sha256"] |
| ) |
| model.configure_residual_gate(torch.load(gate_path, map_location="cpu", weights_only=True)) |
| models[name] = model |
| model_configs[name] = checkpoint["configuration"] |
| checkpoint_provenance[name] = checkpoint["provenance"] |
|
|
| dynamics_manifest = json.loads(Path(base_config["inputs"]["dynamics_manifest_path"]).read_text()) |
| median_distance = dynamics_manifest["normalization"]["median_training_latent_distance"] |
| regimes = regimes or base_config["regimes"] |
| repeats = repeats or base_config["sampling"]["repeats"] |
| output = Path(output_directory or config["output_directory"]) |
| assert not output.exists() |
| output.mkdir(parents=True) |
| records, signatures = [], defaultdict(lambda: [0, None]) |
| model_input_paths = { |
| ( |
| item["inputs"]["latent_cache_path"], |
| item["inputs"]["action_cache_path"], |
| item["inputs"]["dynamics_manifest_path"], |
| ) |
| for item in model_configs.values() |
| } |
| assert len(model_input_paths) == 1 |
| latent_cache_path, action_cache_path, dynamics_manifest_path = model_input_paths.pop() |
| assert latent_cache_path == base_config["inputs"]["latent_cache_path"] |
| assert dynamics_manifest_path == base_config["inputs"]["dynamics_manifest_path"] |
| for regime, specification in regimes.items(): |
| dataset = LatentPopulationDataset( |
| latent_cache_path, |
| action_cache_path, |
| dynamics_manifest_path, |
| regime, |
| base_config["sampling"]["population_size"], |
| base_config["seed"], |
| specification["outcome_role"], |
| specification["control_role"], |
| specification["context"], |
| ) |
| indices = range(min(len(dataset), max_conditions or len(dataset))) |
| for repeat in range(repeats): |
| dataset.set_epoch(repeat) |
| loader = DataLoader( |
| dataset, |
| batch_size=base_config["sampling"]["batch_size"], |
| sampler=list(indices), |
| num_workers=base_config["sampling"]["num_workers"], |
| ) |
| for batch in loader: |
| control, observed = batch["control"], batch["perturbed"] |
| predictions = { |
| name: model(control, batch["action"], batch["action_known"]) |
| for name, model in models.items() |
| } |
| confidences = { |
| name: model.residual_gate_confidence(control) |
| for name, model in models.items() |
| if getattr(model, "residual_gate_threshold", None) is not None |
| } |
| assert all(predicted.dtype == observed.dtype for predicted in predictions.values()) |
| true_effect = observed.mean(1) - control.mean(1) |
| for index, target in enumerate(batch["target"]): |
| key = (regime, "true", target) |
| signatures[key][0] += 1 |
| value = true_effect[index].numpy() |
| signatures[key][1] = ( |
| value if signatures[key][1] is None else signatures[key][1] + value |
| ) |
| for name, predicted in predictions.items(): |
| predicted_effect = (predicted.mean(1) - control.mean(1)).detach() |
| metrics = population_metrics( |
| predicted, |
| observed, |
| control, |
| median_distance, |
| base_config["metrics"], |
| ) |
| for index, target in enumerate(batch["target"]): |
| record = { |
| "regime": regime, |
| "context": specification["context"], |
| "outcome_role": specification["outcome_role"], |
| "target": target, |
| "repeat": repeat, |
| "model": name, |
| } |
| for metric, values in metrics.items(): |
| value = float(values[index].detach()) |
| record[metric] = value if np.isfinite(value) else None |
| if name in confidences: |
| record["residual_gate_confidence"] = float(confidences[name][index]) |
| records.append(record) |
| key = (regime, name, target) |
| signatures[key][0] += 1 |
| value = predicted_effect[index].numpy() |
| signatures[key][1] = ( |
| value if signatures[key][1] is None else signatures[key][1] + value |
| ) |
|
|
| metrics_path = output / "condition_metrics.jsonl" |
| with metrics_path.open("w") as handle: |
| for record in records: |
| handle.write(json.dumps(record, sort_keys=True) + "\n") |
| base_records = [ |
| json.loads(line) |
| for line in Path(inputs["base_condition_metrics_path"]).read_text().splitlines() |
| ] |
| assert len(base_records) == inputs["base_condition_metrics_records"] |
| reference_records = [] |
| reference_summary = {"condition_metrics": [], "retrieval": []} |
| reference_models = set(config.get("reference_models", [])) |
| if "reference_condition_metrics_path" in inputs: |
| assert reference_models |
| reference_path = Path(inputs["reference_condition_metrics_path"]) |
| assert (reference_path.stat().st_size, file_sha256(reference_path)) == ( |
| inputs["reference_condition_metrics_bytes"], |
| inputs["reference_condition_metrics_sha256"], |
| ) |
| reference_records = [json.loads(line) for line in reference_path.read_text().splitlines()] |
| assert len(reference_records) == inputs["reference_condition_metrics_records"] |
| reference_summary_path = Path(inputs["reference_summary_path"]) |
| assert (reference_summary_path.stat().st_size, file_sha256(reference_summary_path)) == ( |
| inputs["reference_summary_bytes"], |
| inputs["reference_summary_sha256"], |
| ) |
| reference_summary = json.loads(reference_summary_path.read_text()) |
| base_summary = json.loads(Path(inputs["base_summary_path"]).read_text()) |
| summary = { |
| "condition_metrics": sorted( |
| base_summary["condition_metrics"] |
| + [ |
| item |
| for item in reference_summary["condition_metrics"] |
| if item["model"] in reference_models |
| ] |
| + _condition_metric_summaries( |
| records, |
| base_config["metrics"]["bootstrap_resamples"], |
| base_config["seed"], |
| ), |
| key=lambda item: (item["regime"], item["model"], item["metric"]), |
| ), |
| "retrieval": sorted( |
| base_summary["retrieval"] |
| + [ |
| item |
| for item in reference_summary["retrieval"] |
| if item["model"] in reference_models |
| ] |
| + _retrieval_summaries(signatures, regimes, models), |
| key=lambda item: (item["regime"], item["model"]), |
| ), |
| } |
| comparisons = paired_model_comparisons( |
| base_records + reference_records + records, |
| config["comparisons"], |
| base_config["metrics"]["bootstrap_resamples"], |
| base_config["seed"], |
| ) |
| provenance = { |
| "config": deepcopy(config), |
| "base_evaluation_config": deepcopy(base_config), |
| "executed_regimes": regimes, |
| "executed_repeats": repeats, |
| "maximum_conditions_per_regime": max_conditions, |
| "checkpoint_provenance": checkpoint_provenance, |
| "file_sha256": { |
| "ablation_condition_metrics": file_sha256(metrics_path), |
| "base_condition_metrics": inputs["base_condition_metrics_sha256"], |
| "base_summary": inputs["base_summary_sha256"], |
| "base_provenance": inputs["base_provenance_sha256"], |
| }, |
| "base_evaluation_manifest_sha256": base_manifest_sha256, |
| "runtime_source_sha256": _runtime_source_hash(), |
| "runtime_environment": _runtime_environment(), |
| "git": _git_state(), |
| } |
| provenance[f"{manifest_kind}_manifest_sha256"] = model_manifest_sha256 |
| if "residual_gate" in config: |
| provenance["file_sha256"]["residual_gate"] = config["residual_gate"]["sha256"] |
| provenance["residual_gate_manifest_sha256"] = gate_manifest_sha256 |
| if selection_manifest_sha256 is not None: |
| provenance[f"{model_source}_selection_manifest_sha256"] = selection_manifest_sha256 |
| if reference_records: |
| provenance["file_sha256"]["reference_condition_metrics"] = inputs[ |
| "reference_condition_metrics_sha256" |
| ] |
| provenance["file_sha256"]["reference_summary"] = inputs["reference_summary_sha256"] |
| (output / "summary.json").write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n") |
| (output / "paired_comparisons.json").write_text( |
| json.dumps(comparisons, indent=2, sort_keys=True) + "\n" |
| ) |
| (output / "provenance.json").write_text( |
| json.dumps(provenance, indent=2, sort_keys=True) + "\n" |
| ) |
| return summary, comparisons, provenance |
|
|
