| |
| |
| import json |
| from collections import Counter, defaultdict |
| from copy import deepcopy |
| from hashlib import sha256 |
| from itertools import product |
| from pathlib import Path |
|
|
| import anndata as ad |
| import h5py |
| import numpy as np |
| import torch |
| import yaml |
| from scipy import sparse |
| from scipy.stats import rankdata, t, wilcoxon |
| from torch.utils.data import DataLoader |
|
|
| from causalcelljepa.dynamics import ( |
| LatentPopulationDataset, |
| anchored_selected_entry, |
| build_dynamics_model, |
| learned_target_id_config, |
| state_ablation_config, |
| ) |
| from causalcelljepa.evaluation import fit_linear_baseline |
| from causalcelljepa.representations import build_autoencoder |
| from causalcelljepa.resources import file_sha256 |
| from causalcelljepa.training import _git_state, _runtime_environment, _runtime_source_hash |
|
|
|
|
| def normalized_hvg_expression(counts, hvg_columns, library_size=10_000): |
| """Normalize by each full measured library before selecting the frozen HVGs.""" |
| counts = np.asarray(counts, dtype=np.float32) |
| totals = counts.sum(1) |
| assert np.isfinite(counts).all() and (counts >= 0).all() and (totals > 0).all() |
| return np.log1p( |
| counts[:, hvg_columns] * (np.float32(library_size) / totals)[:, None] |
| ).astype(np.float32, copy=False) |
|
|
|
|
| def write_expression_cache(config, output_path=None, maximum_cells=None, verify_raw=True): |
| """Stream raw H5AD rows into a latent-aligned normalized 3,000-HVG cache.""" |
| inputs, cache_config = config["inputs"], config["expression_cache"] |
| replogle = json.loads(Path(inputs["replogle_manifest_path"]).read_text()) |
| assert replogle["manifest_sha256"] == inputs["replogle_manifest_sha256"] |
| replogle_config = yaml.safe_load(Path(inputs["replogle_config_path"]).read_text()) |
| assert file_sha256(inputs["replogle_config_path"]) == replogle["runtime"]["config_sha256"] |
| latent_path = Path(inputs["latent_cache_path"]) |
| assert (latent_path.stat().st_size, file_sha256(latent_path)) == ( |
| inputs["latent_cache_bytes"], |
| inputs["latent_cache_sha256"], |
| ) |
| output = Path(output_path or cache_config["output_path"]) |
| assert cache_config["dtype"] == "float32" and not output.exists() |
| output.parent.mkdir(parents=True, exist_ok=True) |
| temporary = output.with_suffix(output.suffix + ".tmp") |
| assert not temporary.exists() |
| with h5py.File(latent_path, "r") as latent: |
| cells = min(int(latent.attrs["cells"]), maximum_cells or int(latent.attrs["cells"])) |
| assert maximum_cells is not None or cells == cache_config["expected_cells"] |
| contexts = latent["context"].asstr()[:cells] |
| source_rows = latent["source_row"][:cells] |
| roles = latent["role"].asstr()[:cells] |
| with h5py.File(temporary, "w") as destination: |
| destination.attrs.update( |
| { |
| "format_version": 1, |
| "cells": cells, |
| "hvg_count": cache_config["hvg_count"], |
| "dtype": cache_config["dtype"], |
| "library_size": replogle_config["data"]["library_size"], |
| "latent_cache_sha256": inputs["latent_cache_sha256"], |
| "replogle_manifest_sha256": inputs["replogle_manifest_sha256"], |
| "hvg_sha256": replogle["genes"]["hvg_sha256"], |
| "role_counts_json": json.dumps(dict(sorted(Counter(roles).items()))), |
| "provenance_json": json.dumps( |
| { |
| "config_sha256": file_sha256("configs/readout.yaml"), |
| "runtime_source_sha256": _runtime_source_hash(), |
| "runtime_environment": _runtime_environment(), |
| "git": _git_state(), |
| }, |
| sort_keys=True, |
| ), |
| } |
| ) |
| expression = destination.create_dataset( |
| "expression", |
| (cells, cache_config["hvg_count"]), |
| dtype="f4", |
| chunks=(cache_config["block_size"], cache_config["hvg_count"]), |
| ) |
| hvg_ids = replogle["genes"]["hvg_gene_ids"] |
| for context in replogle_config["data"]["contexts"]: |
| source = replogle_config["data"]["files"][context] |
| path = Path(inputs["raw_directory"]) / source["filename"] |
| assert path.stat().st_size == source["bytes"] |
| if verify_raw: |
| assert file_sha256(path) == source["sha256"] |
| positions = np.flatnonzero(contexts == context) |
| if not len(positions): |
| continue |
| assert np.array_equal(positions, np.arange(positions[0], positions[-1] + 1)) |
| data = ad.read_h5ad(path, backed="r") |
| columns = np.asarray([data.var_names.get_loc(gene) for gene in hvg_ids]) |
| rows = source_rows[positions] |
| assert np.all(rows[1:] > rows[:-1]) |
| for start in range(0, len(positions), cache_config["block_size"]): |
| selected = positions[start : start + cache_config["block_size"]] |
| counts = np.asarray(data.X[source_rows[selected]]) |
| expression[selected] = normalized_hvg_expression( |
| counts, columns, replogle_config["data"]["library_size"] |
| ) |
| data.file.close() |
| destination.flush() |
| temporary.replace(output) |
| return { |
| "output": str(output), |
| "bytes": output.stat().st_size, |
| "sha256": file_sha256(output), |
| "cells": cells, |
| "hvg_count": cache_config["hvg_count"], |
| "role_counts": dict(sorted(Counter(roles).items())), |
| } |
|
|
|
|
| def decoder_split(cell_ids, roles, fit_roles, seed, validation_fraction, maximum_cells=None): |
| """Create an exact deterministic cell split restricted to representation-visible roles.""" |
| eligible = np.flatnonzero(np.isin(roles, fit_roles)) |
| ranked = sorted( |
| eligible, |
| key=lambda index: sha256(f"{seed}\0readout\0{cell_ids[index]}".encode()).digest(), |
| ) |
| if maximum_cells is not None: |
| ranked = ranked[:maximum_cells] |
| validation_cells = max(1, round(validation_fraction * len(ranked))) |
| return np.sort(ranked[validation_cells:]), np.sort(ranked[:validation_cells]) |
|
|
|
|
| def sufficient_statistics(latent, expression, indices, mean, scale, block_size): |
| """Accumulate the exact multivariate linear-regression sufficient statistics.""" |
| width, genes = latent.shape[1] + 1, expression.shape[1] |
| gram = np.zeros((width, width), dtype=np.float64) |
| cross = np.zeros((width, genes), dtype=np.float64) |
| response_square = 0.0 |
| for start in range(0, len(indices), block_size): |
| selected = indices[start : start + block_size] |
| design = np.empty((len(selected), width), dtype=np.float32) |
| design[:, :-1] = (latent[selected] - mean) / scale |
| design[:, -1] = 1 |
| response = expression[selected].astype(np.float32) |
| gram += design.T @ design |
| cross += design.T @ response |
| response_square += float(np.square(response.astype(np.float64)).sum()) |
| return {"cells": len(indices), "gram": gram, "cross": cross, "response_square": response_square} |
|
|
|
|
| def ridge_solution(statistics, alpha): |
| """Solve the linear decoder while leaving the intercept unpenalized.""" |
| penalty = np.eye(statistics["gram"].shape[0], dtype=np.float64) * alpha |
| penalty[-1, -1] = 0 |
| return np.linalg.solve(statistics["gram"] + penalty, statistics["cross"]) |
|
|
|
|
| def regression_mse(statistics, solution): |
| """Evaluate an un-clipped linear decoder exactly from sufficient statistics.""" |
| residual = ( |
| statistics["response_square"] |
| - 2 * np.sum(solution * statistics["cross"]) |
| + np.sum(solution * (statistics["gram"] @ solution)) |
| ) |
| return max(0.0, float(residual)) / (statistics["cells"] * statistics["cross"].shape[1]) |
|
|
|
|
| def fit_readout(config, maximum_cells=None): |
| """Select and refit the decoder using only explicitly permitted cell roles.""" |
| inputs, decoder = config["inputs"], config["decoder"] |
| assert decoder["architecture"] == "linear" |
| dynamics = json.loads(Path(inputs["dynamics_manifest_path"]).read_text()) |
| assert dynamics["manifest_sha256"] == inputs["dynamics_manifest_sha256"] |
| latent_mean = np.asarray(dynamics["normalization"]["latent_mean"], dtype=np.float32) |
| latent_scale = ( |
| np.asarray(dynamics["normalization"]["latent_std"], dtype=np.float32) |
| * dynamics["normalization"]["dimension_scale"] |
| ) |
| with h5py.File(inputs["latent_cache_path"], "r") as latent, h5py.File( |
| config["expression_cache"]["output_path"], "r" |
| ) as expression_cache: |
| cells = int(expression_cache.attrs["cells"]) |
| assert expression_cache.attrs["latent_cache_sha256"] == inputs["latent_cache_sha256"] |
| cell_ids, roles = latent["cell_id"].asstr()[:cells], latent["role"].asstr()[:cells] |
| train, validation = decoder_split( |
| cell_ids, |
| roles, |
| decoder["fit_roles"], |
| config["seed"], |
| decoder["validation_fraction"], |
| maximum_cells, |
| ) |
| train_stats = sufficient_statistics( |
| latent["latent"], |
| expression_cache["expression"], |
| train, |
| latent_mean, |
| latent_scale, |
| decoder["block_size"], |
| ) |
| validation_stats = sufficient_statistics( |
| latent["latent"], |
| expression_cache["expression"], |
| validation, |
| latent_mean, |
| latent_scale, |
| decoder["block_size"], |
| ) |
| candidates = [] |
| for alpha in decoder["ridge_candidates"]: |
| solution = ridge_solution(train_stats, alpha) |
| candidates.append((regression_mse(validation_stats, solution), alpha)) |
| validation_mse, selected_alpha = min(candidates) |
| combined = { |
| key: train_stats[key] + validation_stats[key] |
| for key in ("cells", "gram", "cross", "response_square") |
| } |
| solution = ridge_solution(combined, selected_alpha) |
| baseline_mean = train_stats["cross"][-1] / train_stats["cells"] |
| baseline = np.vstack( |
| (np.zeros((solution.shape[0] - 1, solution.shape[1])), baseline_mean) |
| ) |
| baseline_mse = regression_mse(validation_stats, baseline) |
| used_roles = sorted(set(roles[np.concatenate((train, validation))])) |
| assert used_roles == sorted(decoder["fit_roles"]) |
| report = { |
| "architecture": decoder["architecture"], |
| "fit_roles": used_roles, |
| "fit_cells": combined["cells"], |
| "train_cells": train_stats["cells"], |
| "validation_cells": validation_stats["cells"], |
| "validation_fraction": decoder["validation_fraction"], |
| "ridge_candidates": decoder["ridge_candidates"], |
| "ridge_validation_mse": [value for value, _ in candidates], |
| "selected_ridge": selected_alpha, |
| "selected_validation_mse": validation_mse, |
| "gene_mean_validation_mse": baseline_mse, |
| "validation_explained_fraction": 1 - validation_mse / baseline_mse, |
| "split_cell_ids_sha256": { |
| "train": sha256("\n".join(sorted(cell_ids[train])).encode()).hexdigest(), |
| "validation": sha256("\n".join(sorted(cell_ids[validation])).encode()).hexdigest(), |
| }, |
| } |
| checkpoint = { |
| "format_version": 1, |
| "weights": torch.from_numpy(solution[:-1].astype(np.float32)), |
| "bias": torch.from_numpy(solution[-1].astype(np.float32)), |
| "latent_mean": torch.from_numpy(latent_mean), |
| "latent_scale": torch.from_numpy(latent_scale), |
| "output_clamp_min": decoder["output_clamp_min"], |
| "report": report, |
| "provenance": { |
| "config_sha256": file_sha256("configs/readout.yaml"), |
| "latent_cache_sha256": inputs["latent_cache_sha256"], |
| "expression_cache_sha256": file_sha256(config["expression_cache"]["output_path"]), |
| "replogle_manifest_sha256": inputs["replogle_manifest_sha256"], |
| "dynamics_manifest_sha256": inputs["dynamics_manifest_sha256"], |
| "runtime_source_sha256": _runtime_source_hash(), |
| "runtime_environment": _runtime_environment(), |
| "git": _git_state(), |
| }, |
| } |
| assert torch.isfinite(checkpoint["weights"]).all() and torch.isfinite(checkpoint["bias"]).all() |
| return checkpoint |
|
|
|
|
| def decode_normalized_latents(latents, checkpoint): |
| """Decode dynamics-space latents to nonnegative normalized log expression.""" |
| return (latents @ checkpoint["weights"] + checkpoint["bias"]).clamp_min( |
| checkpoint["output_clamp_min"] |
| ) |
|
|
|
|
| def decode_representation_centroids(latents, decoder): |
| """Map normalized state centroids through one frozen representation-specific readout.""" |
| if decoder["kind"] == "jepa_linear": |
| return decode_normalized_latents(latents, decoder["checkpoint"]) |
| raw = latents * decoder["latent_scale"] + decoder["latent_mean"] |
| if decoder["kind"] == "pca_inverse": |
| expression = raw @ decoder["components"] + decoder["expression_mean"] |
| else: |
| assert decoder["kind"] == "autoencoder" |
| expression = decoder["model"].decoder(raw) |
| return expression.clamp_min(0) |
|
|
|
|
| def _bh_adjust(p_values): |
| 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) |
| result = np.empty_like(adjusted) |
| result[order] = adjusted |
| return result |
|
|
|
|
| def expression_truth(config, regimes=None, maximum_conditions=None): |
| """Compute batch-matched true effects and batch-level DE calls once per condition.""" |
| inputs, metric_config = config["inputs"], config["metrics"] |
| regimes = regimes or config["regimes"] |
| truth, report = {}, {} |
| with h5py.File(inputs["latent_cache_path"], "r") as latent, h5py.File( |
| inputs["expression_cache_path"], "r" |
| ) as expression_cache: |
| roles = latent["role"].asstr()[:] |
| targets = latent["target"].asstr()[:] |
| batches = latent["source_batch"].asstr()[:] |
| contexts = latent["context"].asstr()[:] |
| expression = expression_cache["expression"] |
| for regime, specification in regimes.items(): |
| control_indices = np.flatnonzero(roles == specification["control_role"]) |
| assert set(contexts[control_indices]) == {specification["context"]} |
| control_means = { |
| batch: expression[np.sort(control_indices[batches[control_indices] == batch])].mean( |
| 0, dtype=np.float64 |
| ) |
| for batch in sorted(set(batches[control_indices])) |
| } |
| outcome_indices = np.flatnonzero(roles == specification["outcome_role"]) |
| assert set(contexts[outcome_indices]) == {specification["context"]} |
| condition_targets = sorted(set(targets[outcome_indices])) |
| condition_targets = condition_targets[: maximum_conditions or len(condition_targets)] |
| report[regime] = { |
| "context": specification["context"], |
| "outcome_role": specification["outcome_role"], |
| "targets": len(condition_targets), |
| "conditions_with_no_degs": 0, |
| } |
| for target in condition_targets: |
| indices = np.sort(outcome_indices[targets[outcome_indices] == target]) |
| observed = expression[indices].astype(np.float64) |
| target_batches = batches[indices] |
| batch_effects, matched_control = [], np.zeros(observed.shape[1]) |
| for batch in sorted(set(target_batches)): |
| selected = target_batches == batch |
| batch_effects.append(observed[selected].mean(0) - control_means[batch]) |
| matched_control += selected.sum() * control_means[batch] |
| batch_effects = np.stack(batch_effects) |
| effect = observed.mean(0) - matched_control / len(indices) |
| if len(batch_effects) > 1: |
| batch_mean = batch_effects.mean(0) |
| standard_error = batch_effects.std(0, ddof=1) / np.sqrt(len(batch_effects)) |
| statistic = np.divide( |
| np.abs(batch_mean), |
| standard_error, |
| out=np.where(batch_mean == 0, 0.0, np.inf), |
| where=standard_error > 0, |
| ) |
| q_value = _bh_adjust(2 * t.sf(statistic, len(batch_effects) - 1)) |
| else: |
| q_value = np.ones(observed.shape[1]) |
| deg = (q_value <= metric_config["deg_batch_fdr"]) & ( |
| np.abs(effect) >= metric_config["deg_min_abs_effect"] |
| ) |
| report[regime]["conditions_with_no_degs"] += int(not deg.any()) |
| truth[regime, target] = { |
| "effect": effect.astype(np.float32), |
| "deg": deg, |
| "q_value": q_value.astype(np.float32), |
| "batches": len(batch_effects), |
| "cells": len(indices), |
| } |
| return truth, report |
|
|
|
|
| def _correlation(predicted, observed, rank=False): |
| if np.std(predicted) <= 1e-12 or np.std(observed) <= 1e-12: |
| return None |
| if rank: |
| predicted, observed = rankdata(predicted), rankdata(observed) |
| predicted, observed = predicted - predicted.mean(), observed - observed.mean() |
| return float(predicted @ observed / (np.linalg.norm(predicted) * np.linalg.norm(observed))) |
|
|
|
|
| def _binary_ranking_metrics(scores, labels): |
| positives, negatives = int(labels.sum()), int((~labels).sum()) |
| if not positives or not negatives: |
| return None, None |
| ranks = rankdata(scores) |
| auroc = (ranks[labels].sum() - positives * (positives + 1) / 2) / ( |
| positives * negatives |
| ) |
| order = np.argsort(-scores, kind="stable") |
| sorted_scores, sorted_labels = scores[order], labels[order] |
| cumulative = np.cumsum(sorted_labels) |
| group_ends = np.r_[np.flatnonzero(sorted_scores[1:] != sorted_scores[:-1]), len(scores) - 1] |
| true_positives = cumulative[group_ends] |
| precision = true_positives / (group_ends + 1) |
| average_precision = np.sum(np.diff(np.r_[0, true_positives]) * precision) / positives |
| return float(average_precision), float(auroc) |
|
|
|
|
| def gene_effect_metrics(predicted, observed, target_index, deg, top_genes=(20, 50, 100)): |
| """Evaluate one perturbation-condition, explicitly labeling outcome-selected scopes.""" |
| result = {} |
| scopes = {"all": np.arange(len(observed))} |
| excluded = np.arange(len(observed)) if target_index is None else np.delete( |
| np.arange(len(observed)), target_index |
| ) |
| scopes["target_excluded"] = excluded |
| for scope, indices in scopes.items(): |
| predicted_scope, observed_scope = predicted[indices], observed[indices] |
| predicted_magnitude, observed_magnitude = ( |
| np.linalg.norm(predicted_scope), |
| np.linalg.norm(observed_scope), |
| ) |
| result[f"{scope}_effect_pearson"] = _correlation(predicted_scope, observed_scope) |
| result[f"{scope}_effect_spearman"] = _correlation( |
| predicted_scope, observed_scope, rank=True |
| ) |
| result[f"{scope}_direction_cosine"] = ( |
| float(predicted_scope @ observed_scope / (predicted_magnitude * observed_magnitude)) |
| if predicted_magnitude > 1e-12 and observed_magnitude > 1e-12 |
| else None |
| ) |
| result[f"{scope}_magnitude_ratio"] = float( |
| predicted_magnitude / max(observed_magnitude, 1e-12) |
| ) |
| result[f"{scope}_magnitude_absolute_error"] = float( |
| abs(predicted_magnitude - observed_magnitude) |
| ) |
| ranked_true = np.argsort(-np.abs(observed), kind="stable") |
| ranked_predicted = np.argsort(-np.abs(predicted), kind="stable") |
| for top in top_genes: |
| indices = ranked_true[:top] |
| result[f"retrospective_top{top}_effect_pearson"] = _correlation( |
| predicted[indices], observed[indices] |
| ) |
| result[f"retrospective_top{top}_effect_spearman"] = _correlation( |
| predicted[indices], observed[indices], rank=True |
| ) |
| result[f"retrospective_top{top}_overlap"] = float( |
| len(set(indices) & set(ranked_predicted[:top])) / top |
| ) |
| auprc, auroc = _binary_ranking_metrics(np.abs(predicted), deg) |
| result["deg_auprc"], result["deg_auroc"] = auprc, auroc |
| result["deg_sign_accuracy"] = ( |
| float(np.mean(np.sign(predicted[deg]) == np.sign(observed[deg]))) if deg.any() else None |
| ) |
| return result |
|
|
|
|
| def _load_pathway_matrix(path, hvg_genes): |
| gene_index = {gene: index for index, gene in enumerate(hvg_genes)} |
| rows, columns, labels = [], [], [] |
| with Path(path).open() as handle: |
| for row, line in enumerate(handle): |
| fields = line.rstrip("\n").split("\t") |
| indices = sorted({gene_index[gene] for gene in fields[2:] if gene in gene_index}) |
| labels.append(fields[0]) |
| rows.extend([row] * len(indices)) |
| columns.extend(indices) |
| matrix = sparse.csr_matrix( |
| (np.ones(len(rows), dtype=np.float32), (rows, columns)), |
| shape=(len(labels), len(hvg_genes)), |
| ) |
| return matrix, labels |
|
|
|
|
| def pathway_scores(effect, matrix): |
| """Return an analytic size-normalized signed rank-enrichment score.""" |
| ranks = rankdata(effect).astype(np.float32) |
| if ranks.std() == 0: |
| return np.zeros(matrix.shape[0], dtype=np.float32) |
| ranks = (ranks - ranks.mean()) / ranks.std() |
| sizes = np.asarray(matrix.sum(1)).ravel() |
| denominator = np.sqrt(sizes * (len(effect) - sizes) / (len(effect) - 1)) |
| return np.asarray(matrix @ ranks).ravel() / denominator |
|
|
|
|
| def pathway_agreement(predicted, observed, matrix, top_k=(10, 25, 50)): |
| predicted_score, observed_score = pathway_scores(predicted, matrix), pathway_scores( |
| observed, matrix |
| ) |
| result = { |
| "pathway_nes_pearson": _correlation(predicted_score, observed_score), |
| "pathway_rank_spearman": _correlation(predicted_score, observed_score, rank=True), |
| "pathway_nes_rmse": float(np.sqrt(np.mean(np.square(predicted_score - observed_score)))), |
| } |
| predicted_order, observed_order = ( |
| np.argsort(-np.abs(predicted_score)), |
| np.argsort(-np.abs(observed_score)), |
| ) |
| for top in top_k: |
| if np.std(predicted) <= 1e-12: |
| result[f"pathway_top{top}_jaccard"] = None |
| else: |
| intersection = len(set(predicted_order[:top]) & set(observed_order[:top])) |
| result[f"pathway_top{top}_jaccard"] = float( |
| intersection / (2 * top - intersection) |
| ) |
| return result |
|
|
|
|
| _TRANSCRIPTOMIC_METADATA = { |
| "regime", |
| "context", |
| "outcome_role", |
| "target", |
| "repeat", |
| "model", |
| "truth_batches", |
| "truth_cells", |
| "true_deg_count", |
| "target_in_hvg", |
| } |
|
|
|
|
| def _condition_summary(records, resamples, seed): |
| target_values = defaultdict(list) |
| for record in records: |
| for metric in set(record) - _TRANSCRIPTOMIC_METADATA: |
| 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))) |
| summary = [] |
| for key, values in sorted(grouped.items()): |
| 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), (resamples, len(values))) |
| ].mean(1) |
| summary.append( |
| { |
| "regime": key[0], |
| "model": key[1], |
| "metric": key[2], |
| "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 summary |
|
|
|
|
| def _paired_comparisons( |
| records, |
| resamples, |
| seed, |
| baselines=("no_change", "mean_effect", "linear_esm", "pseudo_paired"), |
| ): |
| target_values = defaultdict(list) |
| for record in records: |
| for metric in set(record) - _TRANSCRIPTOMIC_METADATA: |
| 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 = [] |
| for regime in sorted({record["regime"] for record in records}): |
| metrics = sorted({key[3] for key in averaged if key[:2] == (regime, "causalcelljepa")}) |
| for baseline in baselines: |
| for metric in metrics: |
| targets = sorted( |
| key[2] |
| for key in averaged |
| if key[:2] == (regime, "causalcelljepa") |
| and key[3] == metric |
| and (regime, baseline, key[2], metric) in averaged |
| ) |
| if not targets: |
| continue |
| model = np.asarray([averaged[regime, "causalcelljepa", x, metric] for x in targets]) |
| reference = np.asarray([averaged[regime, baseline, x, metric] for x in targets]) |
| if metric.endswith("magnitude_ratio"): |
| improvement, direction = np.abs(reference - 1) - np.abs(model - 1), "closer_to_one_is_better" |
| elif metric.endswith(("absolute_error", "rmse")): |
| improvement, direction = reference - model, "lower_is_better" |
| else: |
| improvement, direction = model - reference, "higher_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), (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()), |
| "mean_improvement_bootstrap_95ci": [ |
| float(value) for value in np.quantile(bootstrap, (0.025, 0.975)) |
| ], |
| "wilcoxon_two_sided_p": p_value, |
| } |
| ) |
| q_values = _bh_adjust(np.asarray([item["wilcoxon_two_sided_p"] for item in comparisons])) |
| for comparison, q_value in zip(comparisons, q_values): |
| comparison["benjamini_hochberg_q"] = float(q_value) |
| return comparisons |
|
|
|
|
| def _paired_transcriptomic_models(records, pairs, resamples, seed): |
| """Compare arbitrary models on matched perturbations using transcriptomic metric directions.""" |
| target_values = defaultdict(list) |
| for record in records: |
| for metric in set(record) - _TRANSCRIPTOMIC_METADATA: |
| 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 = [] |
| for regime in sorted({record["regime"] for record in records}): |
| for pair in pairs: |
| candidate, reference = pair["candidate"], pair["reference"] |
| metrics = sorted({ |
| key[3] |
| for key in averaged |
| if key[:2] == (regime, candidate) |
| and (regime, reference, key[2], key[3]) in averaged |
| }) |
| for metric in metrics: |
| targets = sorted( |
| key[2] |
| for key in averaged |
| if key[:2] == (regime, candidate) |
| and key[3] == metric |
| and (regime, reference, key[2], metric) in averaged |
| ) |
| 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.endswith("magnitude_ratio"): |
| improvement = np.abs(reference_values - 1) - np.abs(candidate_values - 1) |
| direction = "closer_to_one_is_better" |
| elif metric.endswith(("absolute_error", "rmse")): |
| improvement, direction = reference_values - candidate_values, "lower_is_better" |
| else: |
| improvement, direction = candidate_values - reference_values, "higher_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), (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, |
| } |
| ) |
| q_values = _bh_adjust(np.asarray([item["wilcoxon_two_sided_p"] for item in comparisons])) |
| for comparison, q_value in zip(comparisons, q_values): |
| comparison["benjamini_hochberg_q"] = float(q_value) |
| return comparisons |
|
|
|
|
| @torch.no_grad() |
| def run_transcriptomic_evaluation( |
| config, regimes=None, repeats=None, maximum_conditions=None, output_directory=None |
| ): |
| """Decode all frozen models and evaluate gene effects without fitting on outcomes.""" |
| inputs = config["inputs"] |
| for kind in ( |
| "latent_cache", |
| "expression_cache", |
| "action_cache", |
| "checkpoint", |
| "pseudo_paired_checkpoint", |
| "readout_checkpoint", |
| "go_gmt", |
| ): |
| assert file_sha256(inputs[f"{kind}_path"]) == inputs[f"{kind}_sha256"] |
| for kind in ("readout", "go"): |
| manifest = json.loads(Path(inputs[f"{kind}_manifest_path"]).read_text()) |
| declared = manifest.pop("manifest_sha256") |
| assert declared == inputs[f"{kind}_manifest_sha256"] == sha256( |
| json.dumps(manifest, sort_keys=True, separators=(",", ":")).encode() |
| ).hexdigest() |
| primary = torch.load(inputs["checkpoint_path"], map_location="cpu", weights_only=False) |
| pseudo = torch.load( |
| inputs["pseudo_paired_checkpoint_path"], map_location="cpu", weights_only=False |
| ) |
| readout = torch.load(inputs["readout_checkpoint_path"], map_location="cpu", weights_only=False) |
| model = build_dynamics_model(primary["configuration"]).eval() |
| model.load_state_dict(primary["model"]) |
| pseudo_model = build_dynamics_model(pseudo["configuration"]).eval() |
| pseudo_model.load_state_dict(pseudo["model"]) |
| assert pseudo["configuration"]["objective"] == "pseudo_paired_mse" |
| linear_effect, mean_effect, baseline_report = fit_linear_baseline(config) |
| replogle = json.loads(Path(inputs["replogle_manifest_path"]).read_text()) |
| hvg_genes = replogle["genes"]["hvg_gene_names"] |
| hvg_index = {gene: index for index, gene in enumerate(hvg_genes)} |
| pathway_matrix, pathway_labels = _load_pathway_matrix(inputs["go_gmt_path"], hvg_genes) |
| assert len(pathway_labels) == 4328 |
| regimes = regimes or config["regimes"] |
| repeats = repeats or config["sampling"]["repeats"] |
| truth, truth_report = expression_truth(config, regimes, maximum_conditions) |
| 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), maximum_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 = batch["control"] |
| 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[x] for x in batch["target"]]))[ |
| :, None |
| ], |
| } |
| control_expression = decode_normalized_latents(control.mean(1), readout) |
| for baseline, predicted in predictions.items(): |
| predicted_effect = ( |
| decode_normalized_latents(predicted.mean(1), readout) - control_expression |
| ).numpy() |
| for index, target in enumerate(batch["target"]): |
| observed = truth[regime, target] |
| metrics = gene_effect_metrics( |
| predicted_effect[index], |
| observed["effect"], |
| hvg_index.get(target), |
| observed["deg"], |
| config["metrics"]["retrospective_top_genes"], |
| ) |
| record = { |
| "regime": regime, |
| "context": specification["context"], |
| "outcome_role": specification["outcome_role"], |
| "target": target, |
| "repeat": repeat, |
| "model": baseline, |
| "truth_batches": observed["batches"], |
| "truth_cells": observed["cells"], |
| "true_deg_count": int(observed["deg"].sum()), |
| "target_in_hvg": target in hvg_index, |
| **metrics, |
| } |
| records.append(record) |
| key = (regime, baseline, target) |
| signatures[key][0] += 1 |
| signatures[key][1] = ( |
| predicted_effect[index] |
| if signatures[key][1] is None |
| else signatures[key][1] + predicted_effect[index] |
| ) |
| pathway_records, retrieval = [], [] |
| for regime in regimes: |
| targets = sorted(target for key, target in truth if key == regime) |
| true_signatures = np.stack([truth[regime, target]["effect"] for target in targets]) |
| normalized_truth = true_signatures / np.linalg.norm( |
| true_signatures, axis=1, keepdims=True |
| ).clip(1e-12) |
| for baseline in sorted({record["model"] for record in records}): |
| predicted_signatures = np.stack( |
| [ |
| signatures[regime, baseline, target][1] |
| / signatures[regime, baseline, target][0] |
| for target in targets |
| ] |
| ) |
| normalized_prediction = predicted_signatures / np.linalg.norm( |
| predicted_signatures, axis=1, keepdims=True |
| ).clip(1e-12) |
| similarity = normalized_prediction @ normalized_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) |
| ] |
| ) |
| 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)), |
| } |
| ) |
| for index, target in enumerate(targets): |
| pathway_records.append( |
| { |
| "regime": regime, |
| "context": regimes[regime]["context"], |
| "outcome_role": regimes[regime]["outcome_role"], |
| "target": target, |
| "repeat": 0, |
| "model": baseline, |
| "truth_batches": truth[regime, target]["batches"], |
| "truth_cells": truth[regime, target]["cells"], |
| "true_deg_count": int(truth[regime, target]["deg"].sum()), |
| "target_in_hvg": target in hvg_index, |
| **pathway_agreement( |
| predicted_signatures[index], |
| true_signatures[index], |
| pathway_matrix, |
| config["metrics"]["pathway_top_k"], |
| ), |
| } |
| ) |
| resamples = config["metrics"]["bootstrap_resamples"] |
| summary = { |
| "condition_metrics": _condition_summary(records, resamples, config["seed"]), |
| "pathway_metrics": _condition_summary(pathway_records, resamples, config["seed"]), |
| "retrieval": retrieval, |
| } |
| paired = { |
| "condition_comparisons": _paired_comparisons(records, resamples, config["seed"]), |
| "pathway_comparisons": _paired_comparisons( |
| pathway_records, resamples, config["seed"] |
| ), |
| } |
| provenance = { |
| "config": deepcopy(config), |
| "executed_regimes": regimes, |
| "executed_repeats": repeats, |
| "maximum_conditions_per_regime": maximum_conditions, |
| "checkpoint_provenance": primary["provenance"], |
| "pseudo_paired_checkpoint_provenance": pseudo["provenance"], |
| "readout_provenance": readout["provenance"], |
| "file_sha256": { |
| kind: inputs[f"{kind}_sha256"] |
| for kind in ( |
| "latent_cache", |
| "expression_cache", |
| "action_cache", |
| "checkpoint", |
| "pseudo_paired_checkpoint", |
| "readout_checkpoint", |
| "go_gmt", |
| ) |
| }, |
| "runtime_source_sha256": _runtime_source_hash(), |
| "runtime_environment": _runtime_environment(), |
| "git": _git_state(), |
| } |
| for filename, payload in ( |
| ("summary.json", summary), |
| ("paired_comparisons.json", paired), |
| ("truth_report.json", truth_report), |
| ("baseline_report.json", baseline_report), |
| ("provenance.json", provenance), |
| ): |
| (output / filename).write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") |
| for filename, values in ( |
| ("condition_metrics.jsonl", records), |
| ("pathway_metrics.jsonl", pathway_records), |
| ): |
| with (output / filename).open("w") as handle: |
| for value in values: |
| handle.write(json.dumps(value, sort_keys=True) + "\n") |
| return summary, paired, truth_report, provenance |
|
|
|
|
| @torch.no_grad() |
| def run_remaining_comparator_evaluation( |
| config, base_config, regimes=None, repeats=None, maximum_conditions=None, output_directory=None |
| ): |
| """Evaluate frozen action/state comparators through leakage-safe representation readouts.""" |
| inputs = config["inputs"] |
| assert file_sha256(config["base_transcriptomics_config_path"]) == config[ |
| "base_transcriptomics_config_sha256" |
| ] |
| model_source = config.get("model_source", "comparator") |
| manifest_kinds = { |
| "comparator": "comparator", |
| "stage2_replication": "replication_training", |
| "anchored": "anchored_training", |
| "multiteacher": "multiteacher_training", |
| "contextual_multiteacher": "contextual_multiteacher_training", |
| } |
| assert model_source in manifest_kinds |
| model_manifest_kind = manifest_kinds[model_source] |
| manifests = {} |
| for name in ("base", model_manifest_kind): |
| payload = json.loads(Path(inputs[f"{name}_manifest_path"]).read_text()) |
| declared = payload.pop("manifest_sha256") |
| assert declared == inputs[f"{name}_manifest_sha256"] == sha256( |
| json.dumps(payload, sort_keys=True, separators=(",", ":")).encode() |
| ).hexdigest() |
| manifests[name] = payload |
| selected_candidate = selected_entry = selection_manifest_sha256 = None |
| frozen_selection_sources = {"anchored", "multiteacher", "contextual_multiteacher"} |
| if model_source in frozen_selection_sources: |
| selection = json.loads( |
| Path(inputs[f"{model_source}_selection_manifest_path"]).read_text() |
| ) |
| selection_manifest_sha256 = selection.pop("manifest_sha256") |
| assert ( |
| selection_manifest_sha256 |
| == inputs[f"{model_source}_selection_manifest_sha256"] |
| == sha256( |
| json.dumps(selection, sort_keys=True, separators=(",", ":")).encode() |
| ).hexdigest() |
| ) |
| selected_candidate, selected_entry = anchored_selected_entry( |
| manifests[model_manifest_kind], selection |
| ) |
| assert config["models"] == {f"{model_source}_selected": selected_candidate} |
| base_artifacts = manifests["base"]["artifacts"]["predictive_evaluation"] |
| for name in ("condition_metrics", "pathway_metrics", "summary", "provenance"): |
| path = Path(inputs[f"base_{name}_path"]) |
| assert (path.stat().st_size, file_sha256(path)) == ( |
| inputs[f"base_{name}_bytes"], |
| inputs[f"base_{name}_sha256"], |
| ) |
| assert base_artifacts[name]["sha256"] == inputs[f"base_{name}_sha256"] |
| assert base_artifacts["condition_metrics"]["records"] == inputs[ |
| "base_condition_metrics_records" |
| ] |
| assert base_artifacts["pathway_metrics"]["records"] == inputs[ |
| "base_pathway_metrics_records" |
| ] |
| base_provenance = json.loads(Path(inputs["base_provenance_path"]).read_text()) |
| assert base_provenance["config"] == base_config and base_provenance["git"]["dirty"] is False |
| reference_models = set(config.get("reference_models", [])) |
| reference_summary = {"condition_metrics": [], "pathway_metrics": [], "retrieval": []} |
| if reference_models: |
| for name in ("condition_metrics", "pathway_metrics", "summary"): |
| path = Path(inputs[f"reference_{name}_path"]) |
| assert (path.stat().st_size, file_sha256(path)) == ( |
| inputs[f"reference_{name}_bytes"], |
| inputs[f"reference_{name}_sha256"], |
| ) |
| assert len( |
| Path(inputs["reference_condition_metrics_path"]).read_text().splitlines() |
| ) == inputs["reference_condition_metrics_records"] |
| assert len( |
| Path(inputs["reference_pathway_metrics_path"]).read_text().splitlines() |
| ) == inputs["reference_pathway_metrics_records"] |
| reference_summary = json.loads(Path(inputs["reference_summary_path"]).read_text()) |
| for kind in ("latent_cache", "expression_cache", "readout_checkpoint", "go_gmt"): |
| assert file_sha256(base_config["inputs"][f"{kind}_path"]) == base_config["inputs"][ |
| f"{kind}_sha256" |
| ] |
|
|
| readout = torch.load( |
| base_config["inputs"]["readout_checkpoint_path"], map_location="cpu", weights_only=False |
| ) |
| assert readout["report"]["fit_roles"] == [ |
| "control_inference", |
| "control_train", |
| "dynamics_train", |
| ] |
| models, decoders, model_configs, checkpoint_provenance = {}, {}, {}, {} |
| representation_provenance = {"jepa_linear": readout["provenance"]} |
| model_manifest = manifests[model_manifest_kind] |
| assert model_manifest["protocol"]["rpe1_perturbed_outcomes_used_for_fit_or_selection"] is False |
| assert model_manifest["protocol"]["sealed_test_outcomes_used_for_fit_or_selection"] is False |
| for name, path in config["models"].items(): |
| if model_source == "stage2_replication": |
| entry = model_manifest["artifacts"]["seeds"][str(path)]["best_checkpoint"] |
| decoder = {"kind": "jepa_linear", "checkpoint": readout} |
| elif model_source in frozen_selection_sources: |
| entry = selected_entry["best_checkpoint"] |
| decoder = {"kind": "jepa_linear", "checkpoint": readout} |
| elif name == "learned_target_id": |
| model_config, _, representation_manifest = learned_target_id_config(path) |
| policy = representation_manifest["policy"] |
| assert policy["heldout_target_identity_used_for_fit"] is False |
| assert policy["validation_target_identity_used_for_fit"] is False |
| decoder = {"kind": "jepa_linear", "checkpoint": readout} |
| else: |
| model_config, specification, representation_manifest = state_ablation_config(path) |
| policy = representation_manifest["policy"] |
| assert policy["rpe1_perturbed_outcomes_used_for_fit"] is False |
| assert policy["sealed_test_outcomes_used_for_fit"] is False |
| dynamics_manifest = json.loads( |
| Path(model_config["inputs"]["dynamics_manifest_path"]).read_text() |
| ) |
| normalization = dynamics_manifest["normalization"] |
| mean = torch.tensor(normalization["latent_mean"], dtype=torch.float32) |
| scale = torch.tensor(normalization["latent_std"], dtype=torch.float32) * normalization[ |
| "dimension_scale" |
| ] |
| if name == "pca_state": |
| with h5py.File(model_config["inputs"]["latent_cache_path"], "r") as cache: |
| components = np.ascontiguousarray(cache["component"][:]) |
| expression_mean = np.ascontiguousarray(cache["expression_mean"][:]) |
| artifact = representation_manifest["artifact"] |
| assert sha256(components.tobytes()).hexdigest() == artifact["component_sha256"] |
| assert sha256(expression_mean.tobytes()).hexdigest() == artifact[ |
| "expression_mean_sha256" |
| ] |
| decoder = { |
| "kind": "pca_inverse", |
| "latent_mean": mean, |
| "latent_scale": scale, |
| "components": torch.from_numpy(components), |
| "expression_mean": torch.from_numpy(expression_mean), |
| } |
| else: |
| assert name == "autoencoder_state" |
| artifact = representation_manifest["training"]["best_checkpoint"] |
| assert (Path(artifact["path"]).stat().st_size, file_sha256(artifact["path"])) == ( |
| artifact["bytes"], |
| artifact["sha256"], |
| ) |
| representation = torch.load( |
| artifact["path"], map_location="cpu", weights_only=False |
| ) |
| assert representation["configuration"] == specification |
| autoencoder = build_autoencoder(specification).eval() |
| autoencoder.load_state_dict(representation["model"]) |
| decoder = { |
| "kind": "autoencoder", |
| "latent_mean": mean, |
| "latent_scale": scale, |
| "model": autoencoder, |
| } |
| representation_provenance[name] = representation["provenance"] |
| if model_source == "comparator": |
| entry = model_manifest["experiments"][name]["best_checkpoint"] |
| assert (Path(entry["path"]).stat().st_size, file_sha256(entry["path"])) == ( |
| entry["bytes"], |
| entry["sha256"], |
| ) |
| checkpoint = torch.load(entry["path"], map_location="cpu", weights_only=False) |
| if model_source == "stage2_replication": |
| model_config = checkpoint["configuration"] |
| assert model_config["seed"] == path |
| assert model_config["replication"]["model_and_sampling_seed"] == path |
| assert model_config["replication"]["target_split_seed"] == base_config["seed"] |
| assert checkpoint["state"]["best_validation_epoch"] == model_manifest["artifacts"][ |
| "seeds" |
| ][str(path)]["full_run"]["best_validation_epoch"] |
| elif model_source in frozen_selection_sources: |
| model_config = checkpoint["configuration"] |
| assert model_config["revision"]["candidate"] == path == selected_candidate |
| assert checkpoint["state"]["best_validation_epoch"] == selected_entry["full_run"][ |
| "best_validation_epoch" |
| ] |
| else: |
| assert checkpoint["configuration"] == model_config |
| assert checkpoint["provenance"]["git"]["dirty"] is False |
| model = build_dynamics_model(model_config).eval() |
| model.load_state_dict(checkpoint["model"]) |
| models[name], decoders[name], model_configs[name] = model, decoder, model_config |
| checkpoint_provenance[name] = checkpoint["provenance"] |
| if model_source == "comparator": |
| representation_provenance[f"{name}_manifest_sha256"] = representation_manifest[ |
| "manifest_sha256" |
| ] |
|
|
| replogle = json.loads(Path(base_config["inputs"]["replogle_manifest_path"]).read_text()) |
| hvg_genes = replogle["genes"]["hvg_gene_names"] |
| hvg_index = {gene: index for index, gene in enumerate(hvg_genes)} |
| pathway_matrix, pathway_labels = _load_pathway_matrix( |
| base_config["inputs"]["go_gmt_path"], hvg_genes |
| ) |
| assert len(pathway_labels) == 4328 |
| regimes = regimes or base_config["regimes"] |
| repeats = repeats or base_config["sampling"]["repeats"] |
| truth, truth_report = expression_truth(base_config, regimes, maximum_conditions) |
| output = Path(output_directory or config["output_directory"]) |
| assert not output.exists() |
| output.mkdir(parents=True) |
| records, signatures = [], defaultdict(lambda: [0, None]) |
| for regime, regime_config in regimes.items(): |
| for name, model in models.items(): |
| model_config = model_configs[name] |
| dataset = LatentPopulationDataset( |
| model_config["inputs"]["latent_cache_path"], |
| model_config["inputs"]["action_cache_path"], |
| model_config["inputs"]["dynamics_manifest_path"], |
| regime, |
| base_config["sampling"]["population_size"], |
| base_config["seed"], |
| regime_config["outcome_role"], |
| regime_config["control_role"], |
| regime_config["context"], |
| ) |
| indices = range(min(len(dataset), maximum_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 = batch["control"] |
| predicted = model(control, batch["action"], batch["action_known"]) |
| predicted_effect = ( |
| decode_representation_centroids(predicted.mean(1), decoders[name]) |
| - decode_representation_centroids(control.mean(1), decoders[name]) |
| ).numpy() |
| for index, target in enumerate(batch["target"]): |
| observed = truth[regime, target] |
| records.append( |
| { |
| "regime": regime, |
| "context": regime_config["context"], |
| "outcome_role": regime_config["outcome_role"], |
| "target": target, |
| "repeat": repeat, |
| "model": name, |
| "truth_batches": observed["batches"], |
| "truth_cells": observed["cells"], |
| "true_deg_count": int(observed["deg"].sum()), |
| "target_in_hvg": target in hvg_index, |
| **gene_effect_metrics( |
| predicted_effect[index], |
| observed["effect"], |
| hvg_index.get(target), |
| observed["deg"], |
| base_config["metrics"]["retrospective_top_genes"], |
| ), |
| } |
| ) |
| key = (regime, name, target) |
| signatures[key][0] += 1 |
| signatures[key][1] = ( |
| predicted_effect[index] |
| if signatures[key][1] is None |
| else signatures[key][1] + predicted_effect[index] |
| ) |
|
|
| pathway_records, retrieval = [], [] |
| for regime, regime_config in regimes.items(): |
| targets = sorted(target for key, target in truth if key == regime) |
| true_signatures = np.stack([truth[regime, target]["effect"] for target in targets]) |
| normalized_truth = true_signatures / np.linalg.norm( |
| true_signatures, axis=1, keepdims=True |
| ).clip(1e-12) |
| for name in models: |
| predicted_signatures = np.stack( |
| [signatures[regime, name, target][1] / signatures[regime, name, target][0] for target in targets] |
| ) |
| normalized_prediction = predicted_signatures / np.linalg.norm( |
| predicted_signatures, axis=1, keepdims=True |
| ).clip(1e-12) |
| similarity = normalized_prediction @ normalized_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) |
| ] |
| ) |
| retrieval.append( |
| { |
| "regime": regime, |
| "model": name, |
| "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)), |
| } |
| ) |
| for index, target in enumerate(targets): |
| observed = truth[regime, target] |
| pathway_records.append( |
| { |
| "regime": regime, |
| "context": regime_config["context"], |
| "outcome_role": regime_config["outcome_role"], |
| "target": target, |
| "repeat": 0, |
| "model": name, |
| "truth_batches": observed["batches"], |
| "truth_cells": observed["cells"], |
| "true_deg_count": int(observed["deg"].sum()), |
| "target_in_hvg": target in hvg_index, |
| **pathway_agreement( |
| predicted_signatures[index], |
| true_signatures[index], |
| pathway_matrix, |
| base_config["metrics"]["pathway_top_k"], |
| ), |
| } |
| ) |
|
|
| base_summary = json.loads(Path(inputs["base_summary_path"]).read_text()) |
| resamples = base_config["metrics"]["bootstrap_resamples"] |
| summary = { |
| "condition_metrics": sorted( |
| base_summary["condition_metrics"] |
| + [ |
| item |
| for item in reference_summary["condition_metrics"] |
| if item["model"] in reference_models |
| ] |
| + _condition_summary(records, resamples, base_config["seed"]), |
| key=lambda item: (item["regime"], item["model"], item["metric"]), |
| ), |
| "pathway_metrics": sorted( |
| base_summary["pathway_metrics"] |
| + [ |
| item |
| for item in reference_summary["pathway_metrics"] |
| if item["model"] in reference_models |
| ] |
| + _condition_summary(pathway_records, 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, |
| key=lambda item: (item["regime"], item["model"]), |
| ), |
| } |
| base_models = { |
| model |
| for pair in config["comparisons"] |
| for model in (pair["candidate"], pair["reference"]) |
| if model not in models |
| } |
| base_condition_lines = Path(inputs["base_condition_metrics_path"]).read_text().splitlines() |
| assert len(base_condition_lines) == inputs["base_condition_metrics_records"] |
| if reference_models: |
| base_condition_lines += Path(inputs["reference_condition_metrics_path"]).read_text().splitlines() |
| base_records = [] |
| for line in base_condition_lines: |
| record = json.loads(line) |
| if record["model"] in base_models: |
| base_records.append(record) |
| base_pathway_lines = Path(inputs["base_pathway_metrics_path"]).read_text().splitlines() |
| assert len(base_pathway_lines) == inputs["base_pathway_metrics_records"] |
| if reference_models: |
| base_pathway_lines += Path(inputs["reference_pathway_metrics_path"]).read_text().splitlines() |
| base_pathway_records = [] |
| for line in base_pathway_lines: |
| record = json.loads(line) |
| if record["model"] in base_models: |
| base_pathway_records.append(record) |
| paired = { |
| "condition_comparisons": _paired_transcriptomic_models( |
| base_records + records, config["comparisons"], resamples, base_config["seed"] |
| ), |
| "pathway_comparisons": _paired_transcriptomic_models( |
| base_pathway_records + pathway_records, |
| config["comparisons"], |
| resamples, |
| base_config["seed"], |
| ), |
| } |
| provenance = { |
| "config": deepcopy(config), |
| "base_transcriptomics_config": deepcopy(base_config), |
| "executed_regimes": regimes, |
| "executed_repeats": repeats, |
| "maximum_conditions_per_regime": maximum_conditions, |
| "checkpoint_provenance": checkpoint_provenance, |
| "representation_provenance": representation_provenance, |
| "base_artifact_sha256": { |
| name: inputs[f"base_{name}_sha256"] |
| for name in ("condition_metrics", "pathway_metrics", "summary", "provenance") |
| }, |
| "base_manifest_sha256": inputs["base_manifest_sha256"], |
| "runtime_source_sha256": _runtime_source_hash(), |
| "runtime_environment": _runtime_environment(), |
| "git": _git_state(), |
| } |
| provenance[f"{model_manifest_kind}_manifest_sha256"] = inputs[ |
| f"{model_manifest_kind}_manifest_sha256" |
| ] |
| if selection_manifest_sha256 is not None: |
| provenance[f"{model_source}_selection_manifest_sha256"] = selection_manifest_sha256 |
| if reference_models: |
| provenance["reference_artifact_sha256"] = { |
| name: inputs[f"reference_{name}_sha256"] |
| for name in ("condition_metrics", "pathway_metrics", "summary") |
| } |
| expected_records = repeats * len(models) * sum(item["targets"] for item in truth_report.values()) |
| assert len(records) == expected_records |
| assert all( |
| value is None or not isinstance(value, float) or np.isfinite(value) |
| for record in records + pathway_records |
| for value in record.values() |
| ) |
| for filename, payload in ( |
| ("summary.json", summary), |
| ("paired_comparisons.json", paired), |
| ("truth_report.json", truth_report), |
| ("provenance.json", provenance), |
| ): |
| (output / filename).write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") |
| for filename, values in ( |
| ("condition_metrics.jsonl", records), |
| ("pathway_metrics.jsonl", pathway_records), |
| ): |
| with (output / filename).open("w") as handle: |
| for value in values: |
| handle.write(json.dumps(value, sort_keys=True) + "\n") |
| return summary, paired, truth_report, provenance |
|
|
|
|
| @torch.no_grad() |
| def run_readout_oracle(config, regimes=None, maximum_conditions=None, output_directory=None): |
| """Audit the decoder ceiling using observed outcome latents, never as a predictor.""" |
| inputs = config["inputs"] |
| for kind in ("latent_cache", "expression_cache", "readout_checkpoint", "go_gmt"): |
| assert file_sha256(inputs[f"{kind}_path"]) == inputs[f"{kind}_sha256"] |
| readout = torch.load(inputs["readout_checkpoint_path"], map_location="cpu", weights_only=False) |
| replogle = json.loads(Path(inputs["replogle_manifest_path"]).read_text()) |
| hvg_genes = replogle["genes"]["hvg_gene_names"] |
| hvg_index = {gene: index for index, gene in enumerate(hvg_genes)} |
| pathway_matrix, pathway_labels = _load_pathway_matrix(inputs["go_gmt_path"], hvg_genes) |
| assert len(pathway_labels) == 4328 |
| regimes = regimes or config["regimes"] |
| truth, truth_report = expression_truth(config, regimes, maximum_conditions) |
| records, pathway_records = [], [] |
| with h5py.File(inputs["latent_cache_path"], "r") as latent: |
| roles = latent["role"].asstr()[:] |
| targets = latent["target"].asstr()[:] |
| batches = latent["source_batch"].asstr()[:] |
| contexts = latent["context"].asstr()[:] |
| mean, scale = readout["latent_mean"].numpy(), readout["latent_scale"].numpy() |
| for regime, specification in regimes.items(): |
| control_indices = np.flatnonzero(roles == specification["control_role"]) |
| assert set(contexts[control_indices]) == {specification["context"]} |
| control_means = { |
| batch: (latent["latent"][np.sort(control_indices[batches[control_indices] == batch])].mean(0) - mean) |
| / scale |
| for batch in sorted(set(batches[control_indices])) |
| } |
| outcome_indices = np.flatnonzero(roles == specification["outcome_role"]) |
| condition_targets = sorted(set(targets[outcome_indices])) |
| condition_targets = condition_targets[: maximum_conditions or len(condition_targets)] |
| for target in condition_targets: |
| indices = np.sort(outcome_indices[targets[outcome_indices] == target]) |
| target_batches = batches[indices] |
| observed_latent = (latent["latent"][indices].mean(0) - mean) / scale |
| matched_control = sum( |
| np.count_nonzero(target_batches == batch) * control_means[batch] |
| for batch in set(target_batches) |
| ).astype(np.float32) / len(indices) |
| decoded_effect = ( |
| decode_normalized_latents(torch.from_numpy(observed_latent[None]), readout) |
| - decode_normalized_latents(torch.from_numpy(matched_control[None]), readout) |
| )[0].numpy() |
| observed = truth[regime, target] |
| metadata = { |
| "regime": regime, |
| "context": specification["context"], |
| "outcome_role": specification["outcome_role"], |
| "target": target, |
| "repeat": 0, |
| "model": "observed_latent_readout", |
| "truth_batches": observed["batches"], |
| "truth_cells": observed["cells"], |
| "true_deg_count": int(observed["deg"].sum()), |
| "target_in_hvg": target in hvg_index, |
| } |
| records.append( |
| { |
| **metadata, |
| **gene_effect_metrics( |
| decoded_effect, |
| observed["effect"], |
| hvg_index.get(target), |
| observed["deg"], |
| config["metrics"]["retrospective_top_genes"], |
| ), |
| } |
| ) |
| pathway_records.append( |
| { |
| **metadata, |
| **pathway_agreement( |
| decoded_effect, |
| observed["effect"], |
| pathway_matrix, |
| config["metrics"]["pathway_top_k"], |
| ), |
| } |
| ) |
| summary = { |
| "condition_metrics": _condition_summary( |
| records, config["metrics"]["bootstrap_resamples"], config["seed"] |
| ), |
| "pathway_metrics": _condition_summary( |
| pathway_records, config["metrics"]["bootstrap_resamples"], config["seed"] |
| ), |
| } |
| provenance = { |
| "diagnostic_only": True, |
| "observed_outcomes_used_as_model_inputs": True, |
| "valid_predictive_baseline": False, |
| "executed_regimes": regimes, |
| "maximum_conditions_per_regime": maximum_conditions, |
| "readout_provenance": readout["provenance"], |
| "runtime_source_sha256": _runtime_source_hash(), |
| "runtime_environment": _runtime_environment(), |
| "git": _git_state(), |
| } |
| output = Path(output_directory or "artifacts/readout_oracle") |
| assert not output.exists() |
| output.mkdir(parents=True) |
| for filename, payload in ( |
| ("summary.json", summary), |
| ("truth_report.json", truth_report), |
| ("provenance.json", provenance), |
| ): |
| (output / filename).write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") |
| for filename, values in ( |
| ("condition_metrics.jsonl", records), |
| ("pathway_metrics.jsonl", pathway_records), |
| ): |
| with (output / filename).open("w") as handle: |
| for value in values: |
| handle.write(json.dumps(value, sort_keys=True) + "\n") |
| return summary, truth_report, provenance |
|
|
|
|
| def _expression_role_effects(config, roles, maximum_targets=None): |
| inputs = config["inputs"] |
| with h5py.File(inputs["latent_cache_path"], "r") as latent, h5py.File( |
| inputs["expression_cache_path"], "r" |
| ) as expression_cache: |
| cache_roles = latent["role"].asstr()[:] |
| targets = latent["target"].asstr()[:] |
| batches = latent["source_batch"].asstr()[:] |
| contexts = latent["context"].asstr()[:] |
| control_indices = np.flatnonzero(cache_roles == "control_train") |
| assert set(contexts[control_indices]) == {"K562"} |
| expression = expression_cache["expression"] |
| control_means = { |
| batch: expression[np.sort(control_indices[batches[control_indices] == batch])].mean( |
| 0, dtype=np.float64 |
| ) |
| for batch in sorted(set(batches[control_indices])) |
| } |
| effects = {} |
| for role in roles: |
| outcome_indices = np.flatnonzero(cache_roles == role) |
| assert set(contexts[outcome_indices]) == {"K562"} |
| role_targets = sorted(set(targets[outcome_indices])) |
| role_targets = role_targets[: maximum_targets or len(role_targets)] |
| effects[role] = {} |
| for target in role_targets: |
| indices = np.sort(outcome_indices[targets[outcome_indices] == target]) |
| target_batches = batches[indices] |
| matched_control = sum( |
| np.count_nonzero(target_batches == batch) * control_means[batch] |
| for batch in set(target_batches) |
| ) / len(indices) |
| effects[role][target] = ( |
| expression[indices].mean(0, dtype=np.float64) - matched_control |
| ).astype(np.float32) |
| return effects |
|
|
|
|
| def fit_direct_gene_baseline(config, maximum_targets=None): |
| """Fit ESM-to-expression effects with K562 train/validation outcomes only.""" |
| direct, base = config["direct_gene"], config["transcriptomics"] |
| effects = _expression_role_effects( |
| base, |
| (direct["fit_outcome_role"], direct["selection_outcome_role"]), |
| maximum_targets, |
| ) |
| action = torch.load(base["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[direct["fit_outcome_role"]] |
| 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[: min(direct["rank"], len(known_targets))] |
| scores = (y - y_mean) @ components.T |
| standardized = (x - x_mean) / x_std |
| gram, cross = standardized.T @ standardized, standardized.T @ scores |
| validation = effects[direct["selection_outcome_role"]] |
| 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 direct["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])) |
| checkpoint = { |
| "format_version": 1, |
| "architecture": "direct_gene_esm_low_rank_ridge", |
| "x_mean": torch.from_numpy(x_mean.astype(np.float32)), |
| "x_std": torch.from_numpy(x_std.astype(np.float32)), |
| "y_mean": torch.from_numpy(y_mean.astype(np.float32)), |
| "components": torch.from_numpy(components.astype(np.float32)), |
| "weights": torch.from_numpy(weights.astype(np.float32)), |
| "report": { |
| "fit_outcome_role": direct["fit_outcome_role"], |
| "selection_outcome_role": direct["selection_outcome_role"], |
| "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_validation_mse": [value for value, _, _ in candidates], |
| }, |
| "provenance": { |
| "config_sha256": file_sha256("configs/direct_gene.yaml"), |
| "latent_cache_sha256": base["inputs"]["latent_cache_sha256"], |
| "expression_cache_sha256": base["inputs"]["expression_cache_sha256"], |
| "action_cache_sha256": base["inputs"]["action_cache_sha256"], |
| "runtime_source_sha256": _runtime_source_hash(), |
| "runtime_environment": _runtime_environment(), |
| "git": _git_state(), |
| }, |
| } |
| return checkpoint |
|
|
|
|
| def direct_gene_predictions(checkpoint, action_path): |
| action = torch.load(action_path, map_location="cpu", weights_only=True) |
| standardized = (action["embedding"] - checkpoint["x_mean"]) / checkpoint["x_std"] |
| predicted = ( |
| standardized @ checkpoint["weights"] @ checkpoint["components"] |
| + checkpoint["y_mean"] |
| ) |
| predicted[~action["known"]] = checkpoint["y_mean"] |
| return {target: predicted[index].numpy() for index, target in enumerate(action["targets"])} |
|
|
|
|
| def fit_kernel_gene_student(config, maximum_targets=None): |
| """Select a low-rank action kernel using only K562 train/validation effects.""" |
| kernel, base = config["kernel_gene"], config["transcriptomics"] |
| assert kernel["experimental_status"] == "exploratory_post_test_requires_external_confirmation" |
| for prefix in ("direct_checkpoint", "action"): |
| path = Path(kernel[f"{prefix}_path"]) |
| assert (path.stat().st_size, file_sha256(path)) == ( |
| kernel[f"{prefix}_bytes"], |
| kernel[f"{prefix}_sha256"], |
| ) |
| if "action_manifest_path" in kernel: |
| action_manifest = json.loads(Path(kernel["action_manifest_path"]).read_text()) |
| declared = action_manifest.pop("manifest_sha256") |
| assert declared == kernel["action_manifest_sha256"] == sha256( |
| json.dumps(action_manifest, sort_keys=True, separators=(",", ":")).encode() |
| ).hexdigest() |
| assert action_manifest["artifact"]["sha256"] == kernel["action_sha256"] |
| direct = torch.load(kernel["direct_checkpoint_path"], map_location="cpu", weights_only=True) |
| action = torch.load(kernel["action_path"], map_location="cpu", weights_only=True) |
| assert direct["architecture"] == "direct_gene_esm_low_rank_ridge" |
| effects = _expression_role_effects( |
| base, |
| (kernel["fit_outcome_role"], kernel["selection_outcome_role"]), |
| maximum_targets, |
| ) |
| action_map = { |
| target: (action["embedding"][index].numpy(), bool(action["known"][index])) |
| for index, target in enumerate(action["targets"]) |
| } |
| train = effects[kernel["fit_outcome_role"]] |
| train_targets = [target for target in sorted(train) if action_map[target][1]] |
| validation = effects[kernel["selection_outcome_role"]] |
| validation_targets = sorted(validation) |
| assert all(action_map[target][1] for target in validation_targets) |
| train_y = np.stack([train[target] for target in train_targets]).astype(np.float64) |
| validation_y = np.stack([validation[target] for target in validation_targets]) |
| components = direct["components"].numpy().astype(np.float64) |
| y_mean = direct["y_mean"].numpy().astype(np.float64) |
| train_scores = (train_y - y_mean) @ components.T |
| candidates, solutions = [], [] |
| for block, (start, end) in kernel["feature_blocks"].items(): |
| train_x = np.stack( |
| [action_map[target][0][start:end] for target in train_targets] |
| ).astype(np.float64) |
| validation_x = np.stack( |
| [action_map[target][0][start:end] for target in validation_targets] |
| ).astype(np.float64) |
| mean, std = train_x.mean(0), train_x.std(0).clip(1e-8) |
| train_x, validation_x = (train_x - mean) / std, (validation_x - mean) / std |
| for kind in kernel["kernels"]: |
| if kind == "linear": |
| problems = [ |
| ( |
| None, |
| None, |
| validation_x, |
| train_x.T @ train_x, |
| train_x.T @ train_scores, |
| kernel["linear_ridge_candidates"], |
| ) |
| ] |
| else: |
| assert kind == "rbf" |
| train_square = np.square(train_x).sum(1) |
| train_distance = np.maximum( |
| train_square[:, None] + train_square[None, :] - 2 * train_x @ train_x.T, |
| 0, |
| ) |
| median_distance = float(np.median(train_distance[train_distance > 0])) |
| validation_distance = np.maximum( |
| np.square(validation_x).sum(1)[:, None] |
| + train_square[None, :] |
| - 2 * validation_x @ train_x.T, |
| 0, |
| ) |
| problems = [ |
| ( |
| float(multiplier), |
| median_distance, |
| np.exp(-multiplier * validation_distance / median_distance), |
| np.exp(-multiplier * train_distance / median_distance), |
| train_scores, |
| kernel["rbf_ridge_candidates"], |
| ) |
| for multiplier in kernel["rbf_bandwidth_multipliers"] |
| ] |
| for multiplier, median, design, gram, cross, ridges in problems: |
| values, vectors = np.linalg.eigh(gram) |
| projected = vectors.T @ cross |
| for alpha in ridges: |
| coefficient = vectors @ (projected / (values[:, None] + alpha)) |
| prediction = design @ coefficient @ components + y_mean |
| candidates.append( |
| { |
| "feature_block": block, |
| "kernel": kind, |
| "bandwidth_multiplier": multiplier, |
| "ridge": float(alpha), |
| "selection_mse": float(np.mean((prediction - validation_y) ** 2)), |
| "selection_mean_effect_pearson": float( |
| np.mean( |
| [ |
| np.corrcoef(predicted, observed)[0, 1] |
| for predicted, observed in zip(prediction, validation_y) |
| ] |
| ) |
| ), |
| "median_squared_distance": median, |
| } |
| ) |
| solutions.append( |
| {"weights": coefficient} |
| if kind == "linear" |
| else {"dual": coefficient, "train_x": train_x} |
| ) |
| selected_index = min( |
| range(len(candidates)), |
| key=lambda index: (candidates[index]["selection_mse"], index), |
| ) |
| selected, solution = candidates[selected_index], solutions[selected_index] |
| start, end = kernel["feature_blocks"][selected["feature_block"]] |
| selected_x = np.stack( |
| [action_map[target][0][start:end] for target in train_targets] |
| ).astype(np.float64) |
| mean, std = selected_x.mean(0), selected_x.std(0).clip(1e-8) |
| return { |
| "format_version": 1, |
| "architecture": "kernel_gene_low_rank_student", |
| "feature_block": selected["feature_block"], |
| "feature_slice": [start, end], |
| "kernel": selected["kernel"], |
| "x_mean": torch.from_numpy(mean.astype(np.float32)), |
| "x_std": torch.from_numpy(std.astype(np.float32)), |
| "y_mean": direct["y_mean"], |
| "components": direct["components"], |
| **{ |
| name: torch.from_numpy(value.astype(np.float32)) |
| for name, value in solution.items() |
| }, |
| "report": { |
| "experimental_status": kernel["experimental_status"], |
| "fit_outcome_role": kernel["fit_outcome_role"], |
| "selection_outcome_role": kernel["selection_outcome_role"], |
| "fit_targets": len(train), |
| "fit_targets_with_known_action": len(train_targets), |
| "selection_targets": len(validation_targets), |
| "rank": len(components), |
| "selection_rule": "minimum validation MSE; config order breaks exact ties", |
| "selected": selected, |
| "candidates": candidates, |
| "sealed_test_outcomes_used_for_fit_or_selection": False, |
| "rpe1_perturbed_outcomes_used_for_fit_or_selection": False, |
| }, |
| "provenance": { |
| "config_sha256": file_sha256(kernel.get("config_path", "configs/kernel_gene.yaml")), |
| "direct_checkpoint_sha256": kernel["direct_checkpoint_sha256"], |
| "action_sha256": kernel["action_sha256"], |
| "latent_cache_sha256": base["inputs"]["latent_cache_sha256"], |
| "expression_cache_sha256": base["inputs"]["expression_cache_sha256"], |
| "runtime_source_sha256": _runtime_source_hash(), |
| "runtime_environment": _runtime_environment(), |
| "git": _git_state(), |
| }, |
| } |
|
|
|
|
| def fit_external_response_student(config, maximum_targets=None): |
| """Fit a multi-view action student without validation or sealed-target outcome leakage.""" |
| from causalcelljepa.external_evaluation import grouped_expression_moments |
|
|
| specification, base = config["external_response"], config["transcriptomics"] |
| assert specification["experimental_status"] == "exploratory_post_test_not_confirmatory" |
| assert specification["fit_outcome_role"] == "dynamics_train" |
| assert specification["selection_outcome_role"] == "perturbation_ood_validation" |
| for name in ("action", "reference_checkpoint"): |
| path = Path(specification[f"{name}_path"]) |
| assert (path.stat().st_size, file_sha256(path)) == ( |
| specification[f"{name}_bytes"], |
| specification[f"{name}_sha256"], |
| ) |
| external_config_path = Path(specification["external_config_path"]) |
| assert file_sha256(external_config_path) == specification["external_config_sha256"] |
| external_config = yaml.safe_load(external_config_path.read_text()) |
| external_manifest = json.loads(Path(specification["external_manifest_path"]).read_text()) |
| declared = external_manifest.pop("manifest_sha256") |
| assert declared == specification["external_manifest_sha256"] == sha256( |
| json.dumps(external_manifest, sort_keys=True, separators=(",", ":")).encode() |
| ).hexdigest() |
| assert external_manifest["protocol"]["role"] == "external_test_only" |
|
|
| effects = _expression_role_effects( |
| base, |
| (specification["fit_outcome_role"], specification["selection_outcome_role"]), |
| maximum_targets, |
| ) |
| train, validation = ( |
| effects[specification["fit_outcome_role"]], |
| effects[specification["selection_outcome_role"]], |
| ) |
| action = torch.load(specification["action_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_targets = [target for target in sorted(train) if action_map[target][1]] |
| validation_targets = sorted(validation) |
| assert all(action_map[target][1] for target in validation_targets) |
| with h5py.File(base["inputs"]["latent_cache_path"], "r") as cache: |
| roles, cache_targets = cache["role"].asstr()[:], cache["target"].asstr()[:] |
| excluded_targets = set().union( |
| *(set(cache_targets[roles == role]) for role in specification["excluded_outcome_roles"]) |
| ) |
| assert not set(train_targets).intersection(excluded_targets) |
|
|
| replogle = json.loads(Path(base["inputs"]["replogle_manifest_path"]).read_text()) |
| hvg_genes = replogle["genes"]["hvg_gene_names"] |
| hvg_index = {gene: index for index, gene in enumerate(hvg_genes)} |
| train_y = np.stack([train[target] for target in train_targets]).astype(np.float64) |
| validation_y = np.stack([validation[target] for target in validation_targets]) |
| y_mean, centered_y = train_y.mean(0), train_y - train_y.mean(0) |
| target_index = {target: index for index, target in enumerate(train_targets)} |
| external_blocks, external_report = {}, {} |
| for context in specification["contexts"]: |
| source = external_config["source"]["files"][context] |
| path = Path(external_config["source"]["raw_directory"]) / source["filename"] |
| assert path.stat().st_size == source["bytes"] |
| assert file_sha256(path) == external_manifest["contexts"][context]["source_sha256"] |
| data = ad.read_h5ad(path, backed="r") |
| labels = data.obs[external_config["source"]["perturbation_column"]].astype(str).to_numpy() |
| eligible = set(external_manifest["contexts"][context]["eligible_targets"]) |
| admitted = sorted(set(train_targets).intersection(eligible)) |
| assert not set(admitted).intersection(excluded_targets) |
| present_genes = [gene for gene in hvg_genes if gene in data.var_names] |
| columns = np.asarray([data.var_names.get_loc(gene) for gene in present_genes]) |
| groups = [external_config["source"]["control_label"], *admitted] |
| means, _, counts = grouped_expression_moments( |
| data.X, labels, groups, columns, specification["source_block_size"] |
| ) |
| data.file.close() |
| assert counts[0] == source["expected_controls"] |
| assert counts[1:].min() >= external_config["protocol"]["population_size"] |
| observed = means[1:] - means[0] |
| observed -= observed.mean(0) |
| rows = np.zeros((len(train_targets), len(present_genes)), dtype=np.float64) |
| positions = np.asarray([target_index[target] for target in admitted]) |
| rows[positions] = observed |
| gene_positions = np.asarray([hvg_index[gene] for gene in present_genes]) |
| reference_rms = np.sqrt(np.mean(centered_y[np.ix_(positions, gene_positions)] ** 2)) |
| source_rms = np.sqrt(np.mean(observed**2)) |
| scale = float(reference_rms / source_rms) |
| external_blocks[context] = rows * scale |
| external_report[context] = { |
| "fit_targets": len(admitted), |
| "fit_target_sha256": sha256("\n".join(admitted).encode()).hexdigest(), |
| "genes": len(present_genes), |
| "cells": int(counts.sum()), |
| "effect_rms_scale_to_k562": scale, |
| "source_sha256": external_manifest["contexts"][context]["source_sha256"], |
| } |
|
|
| start, end = specification["feature_slice"] |
| train_x = np.stack( |
| [action_map[target][0][start:end] for target in train_targets] |
| ).astype(np.float64) |
| validation_x = np.stack( |
| [action_map[target][0][start:end] for target in validation_targets] |
| ).astype(np.float64) |
| x_mean, x_std = train_x.mean(0), train_x.std(0).clip(1e-8) |
| train_x, validation_x = (train_x - x_mean) / x_std, (validation_x - x_mean) / x_std |
| train_square = np.square(train_x).sum(1) |
| train_distance = np.maximum( |
| train_square[:, None] + train_square[None, :] - 2 * train_x @ train_x.T, 0 |
| ) |
| median_distance = float(np.median(train_distance[train_distance > 0])) |
| validation_distance = np.maximum( |
| np.square(validation_x).sum(1)[:, None] |
| + train_square[None, :] |
| - 2 * validation_x @ train_x.T, |
| 0, |
| ) |
| kernels = {} |
| for bandwidth in specification["rbf_bandwidth_multipliers"]: |
| gram = np.exp(-bandwidth * train_distance / median_distance) |
| values, vectors = np.linalg.eigh(gram) |
| kernels[bandwidth] = ( |
| values, |
| vectors, |
| np.exp(-bandwidth * validation_distance / median_distance), |
| ) |
|
|
| reference = torch.load( |
| specification["reference_checkpoint_path"], map_location="cpu", weights_only=True |
| ) |
| reference_prediction = kernel_gene_predictions(reference, specification["action_path"]) |
| reference_y = np.stack([reference_prediction[target] for target in validation_targets]) |
| reference_metrics = { |
| "mse": float(np.mean((reference_y - validation_y) ** 2)), |
| "mean_effect_pearson": float( |
| np.mean( |
| [ |
| np.corrcoef(predicted, observed)[0, 1] |
| for predicted, observed in zip(reference_y, validation_y) |
| ] |
| ) |
| ), |
| "mean_magnitude_absolute_error": float( |
| np.mean( |
| np.abs( |
| np.linalg.norm(reference_y, axis=1) |
| - np.linalg.norm(validation_y, axis=1) |
| ) |
| ) |
| ), |
| } |
| candidates, solutions = [], [] |
| weight_grid = specification["external_view_weights"] |
| for weights in ( |
| dict(zip(specification["contexts"], values, strict=True)) |
| for values in product(weight_grid, repeat=len(specification["contexts"])) |
| ): |
| joined = np.concatenate( |
| [centered_y] |
| + [weights[context] * external_blocks[context] for context in specification["contexts"]], |
| axis=1, |
| ) |
| _, _, components = np.linalg.svd(joined, full_matrices=False) |
| for rank in specification["ranks"]: |
| rank = min(rank, len(train_targets)) |
| selected_components = components[:rank] |
| scores = joined @ selected_components.T |
| gene_head = selected_components[:, : train_y.shape[1]] |
| for bandwidth, (values, vectors, design) in kernels.items(): |
| projected = vectors.T @ scores |
| for ridge in specification["rbf_ridge_candidates"]: |
| dual = vectors @ (projected / (values[:, None] + ridge)) |
| unscaled = design @ dual @ gene_head + y_mean |
| correlation = float( |
| np.mean( |
| [ |
| np.corrcoef(predicted, observed)[0, 1] |
| for predicted, observed in zip(unscaled, validation_y) |
| ] |
| ) |
| ) |
| for output_scale in specification["output_scale_candidates"]: |
| prediction = output_scale * unscaled |
| candidate = { |
| "external_view_weights": weights, |
| "rank": rank, |
| "bandwidth_multiplier": float(bandwidth), |
| "ridge": float(ridge), |
| "output_scale": float(output_scale), |
| "median_squared_distance": median_distance, |
| "selection_mse": float(np.mean((prediction - validation_y) ** 2)), |
| "selection_mean_effect_pearson": correlation, |
| "selection_mean_magnitude_absolute_error": float( |
| np.mean( |
| np.abs( |
| np.linalg.norm(prediction, axis=1) |
| - np.linalg.norm(validation_y, axis=1) |
| ) |
| ) |
| ), |
| } |
| candidates.append(candidate) |
| solutions.append((dual, gene_head)) |
| gates = specification["validation_gates"] |
| admissible = [ |
| index |
| for index, item in enumerate(candidates) |
| if sum(item["external_view_weights"].values()) > 0 |
| and item["selection_mse"] |
| <= reference_metrics["mse"] * (1 - gates["minimum_relative_mse_improvement"]) |
| and item["selection_mean_effect_pearson"] |
| >= reference_metrics["mean_effect_pearson"] + gates["minimum_pearson_improvement"] |
| and item["selection_mean_magnitude_absolute_error"] |
| <= reference_metrics["mean_magnitude_absolute_error"] |
| + gates["maximum_magnitude_error_degradation"] |
| ] |
| selected_index = min( |
| admissible or range(len(candidates)), |
| key=lambda index: (candidates[index]["selection_mse"], index), |
| ) |
| selected, (dual, gene_head) = candidates[selected_index], solutions[selected_index] |
| selected["passes_validation_gate"] = bool(admissible) |
| internal_control = min( |
| (item for item in candidates if sum(item["external_view_weights"].values()) == 0), |
| key=lambda item: item["selection_mse"], |
| ) |
| report = { |
| "experimental_status": specification["experimental_status"], |
| "protocol_amendment": specification["protocol_amendment"], |
| "fit_outcome_role": specification["fit_outcome_role"], |
| "selection_outcome_role": specification["selection_outcome_role"], |
| "fit_targets": len(train_targets), |
| "selection_targets": len(validation_targets), |
| "external": external_report, |
| "excluded_outcome_roles": specification["excluded_outcome_roles"], |
| "excluded_target_count": len(excluded_targets), |
| "excluded_target_sha256": sha256("\n".join(sorted(excluded_targets)).encode()).hexdigest(), |
| "reference": reference_metrics, |
| "internal_no_external_control": internal_control, |
| "selected": selected, |
| "candidate_count": len(candidates), |
| "validation_gates": gates, |
| "sealed_test_outcomes_used_for_fit_or_selection": False, |
| "rpe1_perturbed_outcomes_used_for_fit_or_selection": False, |
| "prior_external_results_make_this_confirmatory": False, |
| } |
| return { |
| "format_version": 1, |
| "architecture": "external_response_multiview_kernel_gene_student", |
| "feature_slice": specification["feature_slice"], |
| "kernel": "rbf", |
| "x_mean": torch.from_numpy(x_mean.astype(np.float32)), |
| "x_std": torch.from_numpy(x_std.astype(np.float32)), |
| "y_mean": torch.from_numpy(y_mean.astype(np.float32)), |
| "components": torch.from_numpy(gene_head.astype(np.float32)), |
| "train_x": torch.from_numpy(train_x.astype(np.float32)), |
| "dual": torch.from_numpy(dual.astype(np.float32)), |
| "output_scale": selected["output_scale"], |
| "report": report, |
| "provenance": { |
| "config_sha256": file_sha256(specification["config_path"]), |
| "external_config_sha256": specification["external_config_sha256"], |
| "external_manifest_sha256": declared, |
| "reference_checkpoint_sha256": specification["reference_checkpoint_sha256"], |
| "action_sha256": specification["action_sha256"], |
| "latent_cache_sha256": base["inputs"]["latent_cache_sha256"], |
| "expression_cache_sha256": base["inputs"]["expression_cache_sha256"], |
| "runtime_source_sha256": _runtime_source_hash(), |
| "runtime_environment": _runtime_environment(), |
| "git": _git_state(), |
| }, |
| } |
|
|
|
|
| def kernel_gene_predictions(checkpoint, action_path): |
| """Apply a frozen linear or RBF action student without outcome access.""" |
| action = torch.load(action_path, map_location="cpu", weights_only=True) |
| start, end = checkpoint["feature_slice"] |
| standardized = ( |
| action["embedding"][:, start:end] - checkpoint["x_mean"] |
| ) / checkpoint["x_std"] |
| if checkpoint["kernel"] == "linear": |
| scores = standardized @ checkpoint["weights"] |
| else: |
| distance = torch.cdist(standardized, checkpoint["train_x"]).square() |
| selected = checkpoint["report"]["selected"] |
| similarity = torch.exp( |
| -selected["bandwidth_multiplier"] |
| * distance |
| / selected["median_squared_distance"] |
| ) |
| scores = similarity @ checkpoint["dual"] |
| predicted = checkpoint.get("output_scale", 1.0) * ( |
| scores @ checkpoint["components"] + checkpoint["y_mean"] |
| ) |
| predicted[~action["known"]] = checkpoint.get("output_scale", 1.0) * checkpoint[ |
| "y_mean" |
| ] |
| return {target: predicted[index].numpy() for index, target in enumerate(action["targets"])} |
|
|
|
|
| def score_effect_predictions(base, model_name, predictions, maximum_conditions=None): |
| """Score frozen gene-effect predictions through the common reporting-only endpoint.""" |
| truth, truth_report = expression_truth(base, maximum_conditions=maximum_conditions) |
| replogle = json.loads(Path(base["inputs"]["replogle_manifest_path"]).read_text()) |
| hvg_genes = replogle["genes"]["hvg_gene_names"] |
| hvg_index = {gene: index for index, gene in enumerate(hvg_genes)} |
| pathway_matrix, pathway_labels = _load_pathway_matrix( |
| base["inputs"]["go_gmt_path"], hvg_genes |
| ) |
| assert len(pathway_labels) == 4328 |
| records, signatures, seen = [], defaultdict(lambda: [0, None]), Counter() |
| for regime, target, repeat, predicted in predictions: |
| assert (regime, target) in truth and repeat >= 0 |
| predicted = np.asarray(predicted, dtype=np.float32) |
| assert predicted.shape == (len(hvg_genes),) and np.isfinite(predicted).all() |
| observed = truth[regime, target] |
| specification = base["regimes"][regime] |
| metadata = { |
| "regime": regime, |
| "context": specification["context"], |
| "outcome_role": specification["outcome_role"], |
| "target": target, |
| "repeat": int(repeat), |
| "model": model_name, |
| "truth_batches": observed["batches"], |
| "truth_cells": observed["cells"], |
| "true_deg_count": int(observed["deg"].sum()), |
| "target_in_hvg": target in hvg_index, |
| } |
| records.append( |
| { |
| **metadata, |
| **gene_effect_metrics( |
| predicted, |
| observed["effect"], |
| hvg_index.get(target), |
| observed["deg"], |
| base["metrics"]["retrospective_top_genes"], |
| ), |
| } |
| ) |
| key = regime, model_name, target |
| signatures[key][0] += 1 |
| signatures[key][1] = predicted if signatures[key][1] is None else signatures[key][1] + predicted |
| seen[regime, target, int(repeat)] += 1 |
| assert {(regime, target) for regime, target, _ in seen} == set(truth) |
| assert set(seen.values()) == {1} |
|
|
| pathway_records, retrieval = [], [] |
| for regime, specification in base["regimes"].items(): |
| targets = sorted(target for key, target in truth if key == regime) |
| predicted_signatures = np.stack( |
| [signatures[regime, model_name, target][1] / signatures[regime, model_name, target][0] for target in targets] |
| ) |
| true_signatures = np.stack([truth[regime, target]["effect"] for target in targets]) |
| normalized_prediction = predicted_signatures / np.linalg.norm( |
| predicted_signatures, axis=1, keepdims=True |
| ).clip(1e-12) |
| normalized_truth = true_signatures / np.linalg.norm( |
| true_signatures, axis=1, keepdims=True |
| ).clip(1e-12) |
| similarity = normalized_prediction @ normalized_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) |
| ] |
| ) |
| retrieval.append( |
| { |
| "regime": regime, |
| "model": model_name, |
| "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)), |
| } |
| ) |
| for index, target in enumerate(targets): |
| observed = truth[regime, target] |
| pathway_records.append( |
| { |
| "regime": regime, |
| "context": specification["context"], |
| "outcome_role": specification["outcome_role"], |
| "target": target, |
| "repeat": 0, |
| "model": model_name, |
| "truth_batches": observed["batches"], |
| "truth_cells": observed["cells"], |
| "true_deg_count": int(observed["deg"].sum()), |
| "target_in_hvg": target in hvg_index, |
| **pathway_agreement( |
| predicted_signatures[index], |
| true_signatures[index], |
| pathway_matrix, |
| base["metrics"]["pathway_top_k"], |
| ), |
| } |
| ) |
| return records, pathway_records, retrieval, truth_report |
|
|
|
|
| def run_state_baseline_evaluation(config, base): |
| """Evaluate validation-frozen State predictions and pair them with frozen references.""" |
| evaluation, inputs = config["evaluation"], config["inputs"] |
| prediction_report = json.loads(Path(config["prediction"]["report_path"]).read_text()) |
| prediction_path = Path(prediction_report["path"]) |
| assert (prediction_path.stat().st_size, file_sha256(prediction_path)) == ( |
| prediction_report["bytes"], |
| prediction_report["sha256"], |
| ) |
| assert (Path(inputs["checkpoint_path"]).stat().st_size, file_sha256(inputs["checkpoint_path"])) == ( |
| inputs["checkpoint_bytes"], |
| inputs["checkpoint_sha256"], |
| ) |
| with h5py.File(prediction_path, "r") as prediction: |
| assert prediction.attrs["checkpoint_sha256"] == inputs["checkpoint_sha256"] |
| assert not prediction.attrs["test_outcome_expression_read"] |
| assert len(prediction["predicted_effect"]) == prediction_report["records"] |
| predictions = ( |
| ( |
| prediction["regime"].asstr()[index], |
| prediction["target"].asstr()[index], |
| int(prediction["repeat"][index]), |
| prediction["predicted_effect"][index], |
| ) |
| for index in range(prediction_report["records"]) |
| ) |
| records, pathways, retrieval, truth_report = score_effect_predictions( |
| base, config["prediction"]["model_name"], predictions |
| ) |
| reference_records, reference_pathways = [], [] |
| for reference in evaluation["references"]: |
| for name, destination in ( |
| ("condition_metrics", reference_records), |
| ("pathway_metrics", reference_pathways), |
| ): |
| path = Path(reference[f"{name}_path"]) |
| assert (path.stat().st_size, file_sha256(path)) == ( |
| reference[f"{name}_bytes"], |
| reference[f"{name}_sha256"], |
| ) |
| lines = path.read_text().splitlines() |
| assert len(lines) == reference[f"{name}_records"] |
| destination.extend( |
| value |
| for value in map(json.loads, lines) |
| if value["model"] in reference["models"] |
| ) |
| resamples = base["metrics"]["bootstrap_resamples"] |
| summary = { |
| "condition_metrics": _condition_summary(records, resamples, base["seed"]), |
| "pathway_metrics": _condition_summary(pathways, resamples, base["seed"]), |
| "retrieval": retrieval, |
| } |
| paired = { |
| "condition_comparisons": _paired_transcriptomic_models( |
| reference_records + records, evaluation["comparisons"], resamples, base["seed"] |
| ), |
| "pathway_comparisons": _paired_transcriptomic_models( |
| reference_pathways + pathways, evaluation["comparisons"], resamples, base["seed"] |
| ), |
| } |
| provenance = { |
| "config": deepcopy(config), |
| "base_transcriptomics_config": deepcopy(base), |
| "reporting_only": True, |
| "checkpoint_selected_on": "K562 perturbation_ood_validation only", |
| "rpe1_perturbed_outcomes_used_for_fit_or_selection": False, |
| "sealed_test_outcomes_used_for_fit_or_selection": False, |
| "prediction_report": prediction_report, |
| "checkpoint_sha256": inputs["checkpoint_sha256"], |
| "runtime_source_sha256": _runtime_source_hash(), |
| "runtime_environment": _runtime_environment(), |
| "git": _git_state(), |
| } |
| output = Path(evaluation["output_directory"]) |
| assert not output.exists() |
| output.mkdir(parents=True) |
| for filename, payload in ( |
| ("summary.json", summary), |
| ("paired_comparisons.json", paired), |
| ("truth_report.json", truth_report), |
| ("provenance.json", provenance), |
| ): |
| (output / filename).write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") |
| for filename, values in (("condition_metrics.jsonl", records), ("pathway_metrics.jsonl", pathways)): |
| with (output / filename).open("w") as handle: |
| for value in values: |
| handle.write(json.dumps(value, sort_keys=True) + "\n") |
| return summary, paired, truth_report, provenance |
|
|
|
|
| def run_kernel_gene_evaluation(config, maximum_conditions=None, output_directory=None): |
| """Report a selection-frozen kernel student against immutable references.""" |
| kernel, evaluation, base = config["kernel_gene"], config["evaluation"], config["transcriptomics"] |
| references, reference_evaluation = [], evaluation |
| while "reference_config_path" in reference_evaluation: |
| assert file_sha256(reference_evaluation["reference_config_path"]) == reference_evaluation[ |
| "reference_config_sha256" |
| ] |
| parent = yaml.safe_load(Path(reference_evaluation["reference_config_path"]).read_text())[ |
| "evaluation" |
| ] |
| manifest = json.loads(Path(reference_evaluation["reference_manifest_path"]).read_text()) |
| declared_reference = manifest.pop("manifest_sha256") |
| assert declared_reference == reference_evaluation["reference_manifest_sha256"] == sha256( |
| json.dumps(manifest, sort_keys=True, separators=(",", ":")).encode() |
| ).hexdigest() |
| references.append( |
| { |
| "models": [parent["model_name"]], |
| **{ |
| f"{name}_{field}": manifest["artifacts"][name][field] |
| for name in ("condition_metrics", "pathway_metrics") |
| for field in ("path", "bytes", "records", "sha256") |
| }, |
| } |
| ) |
| reference_evaluation = parent |
| references = reference_evaluation["references"] + references |
| comparisons = [ |
| {**pair, "candidate": evaluation["model_name"]} for pair in evaluation["comparisons"] |
| ] |
| selection = json.loads(Path(kernel["selection_manifest_path"]).read_text()) |
| declared = selection.pop("manifest_sha256") |
| assert declared == kernel["selection_manifest_sha256"] == sha256( |
| json.dumps(selection, sort_keys=True, separators=(",", ":")).encode() |
| ).hexdigest() |
| checkpoint_artifact = selection["artifacts"]["checkpoint"] |
| checkpoint_path = Path(checkpoint_artifact["path"]) |
| assert (checkpoint_path.stat().st_size, file_sha256(checkpoint_path)) == ( |
| checkpoint_artifact["bytes"], |
| checkpoint_artifact["sha256"], |
| ) |
| action_path = Path(kernel["action_path"]) |
| assert (action_path.stat().st_size, file_sha256(action_path)) == ( |
| kernel["action_bytes"], |
| kernel["action_sha256"], |
| ) |
| checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True) |
| selected = checkpoint["report"]["selected"] |
| frozen = selection["selection"]["selected"] |
| if checkpoint["architecture"] == "external_response_multiview_kernel_gene_student": |
| assert { |
| "external_view_weights": selected["external_view_weights"], |
| "rank": selected["rank"], |
| "bandwidth_multiplier": selected["bandwidth_multiplier"], |
| "ridge": selected["ridge"], |
| "output_scale": selected["output_scale"], |
| "mse": selected["selection_mse"], |
| } == { |
| key: frozen[key] |
| for key in ( |
| "external_view_weights", |
| "rank", |
| "bandwidth_multiplier", |
| "ridge", |
| "output_scale", |
| "mse", |
| ) |
| } |
| else: |
| assert ( |
| selected["feature_block"], |
| selected["kernel"], |
| selected["bandwidth_multiplier"], |
| selected["ridge"], |
| selected["selection_mse"], |
| ) == ( |
| frozen["feature_block"], |
| frozen["kernel"], |
| frozen["bandwidth_multiplier"], |
| frozen["ridge"], |
| frozen["mse"], |
| ) |
| predicted = kernel_gene_predictions(checkpoint, action_path) |
| gate_report = None |
| if "control_gate" in kernel: |
| gate_config = kernel["control_gate"] |
| reference_selection = json.loads(Path(gate_config["reference_manifest_path"]).read_text()) |
| reference_declared = reference_selection.pop("manifest_sha256") |
| assert reference_declared == gate_config["reference_manifest_sha256"] == sha256( |
| json.dumps(reference_selection, sort_keys=True, separators=(",", ":")).encode() |
| ).hexdigest() |
| reference_artifact = reference_selection["artifacts"]["checkpoint"] |
| reference_path = Path(reference_artifact["path"]) |
| assert (reference_path.stat().st_size, file_sha256(reference_path)) == ( |
| reference_artifact["bytes"], reference_artifact["sha256"] |
| ) |
| reference_checkpoint = torch.load(reference_path, map_location="cpu", weights_only=True) |
| reference_prediction = kernel_gene_predictions(reference_checkpoint, action_path) |
| gate_manifest = json.loads(Path(gate_config["gate_manifest_path"]).read_text()) |
| gate_declared = gate_manifest.pop("manifest_sha256") |
| assert gate_declared == gate_config["gate_manifest_sha256"] == sha256( |
| json.dumps(gate_manifest, sort_keys=True, separators=(",", ":")).encode() |
| ).hexdigest() |
| gate_artifact = gate_manifest["artifact"] |
| gate_path = Path(gate_artifact["path"]) |
| assert (gate_path.stat().st_size, file_sha256(gate_path)) == ( |
| gate_artifact["bytes"], gate_artifact["sha256"] |
| ) |
| gate = torch.load(gate_path, map_location="cpu", weights_only=True) |
| dynamics_manifest = json.loads(Path(base["inputs"]["dynamics_manifest_path"]).read_text()) |
| normalization = dynamics_manifest["normalization"] |
| latent_mean = np.asarray(normalization["latent_mean"]) |
| latent_scale = np.asarray(normalization["latent_std"]) * normalization["dimension_scale"] |
| with h5py.File(base["inputs"]["latent_cache_path"], "r") as latent: |
| roles, targets = latent["role"].asstr()[:], latent["target"].asstr()[:] |
| if "control_gate" in kernel: |
| confidences = {} |
| for role in sorted({value["control_role"] for value in base["regimes"].values()}): |
| control = ((latent["latent"][roles == role] - latent_mean) / latent_scale).mean(0) |
| score = float(np.square((control - gate["center"].numpy()) / gate["scale"].numpy()).mean()) |
| confidence = float( |
| np.exp(-max(0, score - gate["threshold"]) / gate["temperature"]) |
| ) |
| confidences[role] = {"score": score, "reference_confidence": confidence} |
| assert confidences["control_train"]["reference_confidence"] >= gate_config[ |
| "minimum_in_domain_confidence" |
| ] |
| assert confidences["control_inference"]["reference_confidence"] <= gate_config[ |
| "maximum_shifted_context_confidence" |
| ] |
| gate_report = { |
| "formula": "reference_confidence * STRING_GO + (1-reference_confidence) * external_response", |
| "controls": confidences, |
| "reference_manifest_sha256": reference_declared, |
| "gate_manifest_sha256": gate_declared, |
| } |
| rows = [] |
| for regime, specification in base["regimes"].items(): |
| regime_targets = sorted(set(targets[roles == specification["outcome_role"]])) |
| regime_targets = regime_targets[: maximum_conditions or len(regime_targets)] |
| confidence = ( |
| gate_report["controls"][specification["control_role"]]["reference_confidence"] |
| if gate_report |
| else 0.0 |
| ) |
| rows.extend( |
| ( |
| regime, |
| target, |
| 0, |
| confidence * reference_prediction[target] + (1 - confidence) * predicted[target] |
| if gate_report |
| else predicted[target], |
| ) |
| for target in regime_targets |
| ) |
| records, pathways, retrieval, truth_report = score_effect_predictions( |
| base, evaluation["model_name"], rows, maximum_conditions |
| ) |
| reference_records, reference_pathways = [], [] |
| for reference in references: |
| for name, destination in ( |
| ("condition_metrics", reference_records), |
| ("pathway_metrics", reference_pathways), |
| ): |
| path = Path(reference[f"{name}_path"]) |
| assert (path.stat().st_size, file_sha256(path)) == ( |
| reference[f"{name}_bytes"], |
| reference[f"{name}_sha256"], |
| ) |
| lines = path.read_text().splitlines() |
| assert len(lines) == reference[f"{name}_records"] |
| destination.extend( |
| value for value in map(json.loads, lines) if value["model"] in reference["models"] |
| ) |
| resamples = base["metrics"]["bootstrap_resamples"] |
| summary = { |
| "condition_metrics": _condition_summary(records, resamples, base["seed"]), |
| "pathway_metrics": _condition_summary(pathways, resamples, base["seed"]), |
| "retrieval": retrieval, |
| } |
| paired = { |
| "condition_comparisons": _paired_transcriptomic_models( |
| reference_records + records, comparisons, resamples, base["seed"] |
| ), |
| "pathway_comparisons": _paired_transcriptomic_models( |
| reference_pathways + pathways, comparisons, resamples, base["seed"] |
| ), |
| } |
| provenance = { |
| "config": deepcopy(config), |
| "experimental_status": selection["experimental_status"], |
| "reporting_only": True, |
| "selection_manifest_sha256": declared, |
| "selection_checkpoint_frozen_before_test_evaluation": True, |
| "mixture_defined_after_component_test_evaluation": gate_report is not None, |
| "rpe1_perturbed_outcomes_used_for_fit_or_selection": False, |
| "sealed_test_outcomes_used_for_fit_or_selection": False, |
| "checkpoint_provenance": checkpoint["provenance"], |
| "control_gate": gate_report, |
| "maximum_conditions_per_regime": maximum_conditions, |
| "runtime_source_sha256": _runtime_source_hash(), |
| "runtime_environment": _runtime_environment(), |
| "git": _git_state(), |
| } |
| output = Path(output_directory or evaluation["output_directory"]) |
| assert not output.exists() |
| output.mkdir(parents=True) |
| for filename, payload in ( |
| ("summary.json", summary), |
| ("paired_comparisons.json", paired), |
| ("truth_report.json", truth_report), |
| ("provenance.json", provenance), |
| ): |
| (output / filename).write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") |
| for filename, values in (("condition_metrics.jsonl", records), ("pathway_metrics.jsonl", pathways)): |
| with (output / filename).open("w") as handle: |
| for value in values: |
| handle.write(json.dumps(value, sort_keys=True) + "\n") |
| return summary, paired, truth_report, provenance |
|
|
|
|
| def run_direct_gene_evaluation(config, checkpoint): |
| """Evaluate the frozen direct gene-space baseline against the frozen five-model result.""" |
| base, direct = config["transcriptomics"], config["direct_gene"] |
| manifest = json.loads(Path(config["transcriptomics_manifest_path"]).read_text()) |
| declared = manifest.pop("manifest_sha256") |
| assert declared == config["transcriptomics_manifest_sha256"] == sha256( |
| json.dumps(manifest, sort_keys=True, separators=(",", ":")).encode() |
| ).hexdigest() |
| prior_directory = Path(base["output_directory"]) |
| prior_hash = manifest["artifacts"]["predictive_evaluation"]["condition_metrics"][ |
| "sha256" |
| ] |
| assert file_sha256(prior_directory / "condition_metrics.jsonl") == prior_hash |
| predictions = direct_gene_predictions(checkpoint, base["inputs"]["action_cache_path"]) |
| truth, truth_report = expression_truth(base) |
| replogle = json.loads(Path(base["inputs"]["replogle_manifest_path"]).read_text()) |
| hvg_genes = replogle["genes"]["hvg_gene_names"] |
| hvg_index = {gene: index for index, gene in enumerate(hvg_genes)} |
| pathway_matrix, pathway_labels = _load_pathway_matrix( |
| base["inputs"]["go_gmt_path"], hvg_genes |
| ) |
| assert len(pathway_labels) == 4328 |
| records, pathway_records, retrieval = [], [], [] |
| for regime, specification in base["regimes"].items(): |
| targets = sorted(target for key, target in truth if key == regime) |
| predicted_signatures = np.stack([predictions[target] for target in targets]) |
| true_signatures = np.stack([truth[regime, target]["effect"] for target in targets]) |
| normalized_prediction = predicted_signatures / np.linalg.norm( |
| predicted_signatures, axis=1, keepdims=True |
| ).clip(1e-12) |
| normalized_truth = true_signatures / np.linalg.norm( |
| true_signatures, axis=1, keepdims=True |
| ).clip(1e-12) |
| similarity = normalized_prediction @ normalized_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) |
| ] |
| ) |
| retrieval.append( |
| { |
| "regime": regime, |
| "model": "direct_gene_esm", |
| "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)), |
| } |
| ) |
| for index, target in enumerate(targets): |
| observed = truth[regime, target] |
| metadata = { |
| "regime": regime, |
| "context": specification["context"], |
| "outcome_role": specification["outcome_role"], |
| "target": target, |
| "repeat": 0, |
| "model": "direct_gene_esm", |
| "truth_batches": observed["batches"], |
| "truth_cells": observed["cells"], |
| "true_deg_count": int(observed["deg"].sum()), |
| "target_in_hvg": target in hvg_index, |
| } |
| records.append( |
| { |
| **metadata, |
| **gene_effect_metrics( |
| predicted_signatures[index], |
| observed["effect"], |
| hvg_index.get(target), |
| observed["deg"], |
| base["metrics"]["retrospective_top_genes"], |
| ), |
| } |
| ) |
| pathway_records.append( |
| { |
| **metadata, |
| **pathway_agreement( |
| predicted_signatures[index], |
| observed["effect"], |
| pathway_matrix, |
| base["metrics"]["pathway_top_k"], |
| ), |
| } |
| ) |
| prior_records = [ |
| json.loads(line) |
| for line in (prior_directory / "condition_metrics.jsonl").read_text().splitlines() |
| if '"model": "causalcelljepa"' in line |
| ] |
| prior_pathways = [ |
| json.loads(line) |
| for line in (prior_directory / "pathway_metrics.jsonl").read_text().splitlines() |
| if '"model": "causalcelljepa"' in line |
| ] |
| resamples = base["metrics"]["bootstrap_resamples"] |
| summary = { |
| "condition_metrics": _condition_summary(records, resamples, base["seed"]), |
| "pathway_metrics": _condition_summary(pathway_records, resamples, base["seed"]), |
| "retrieval": retrieval, |
| } |
| paired = { |
| "condition_comparisons": _paired_comparisons( |
| prior_records + records, |
| resamples, |
| base["seed"], |
| baselines=("direct_gene_esm",), |
| ), |
| "pathway_comparisons": _paired_comparisons( |
| prior_pathways + pathway_records, |
| resamples, |
| base["seed"], |
| baselines=("direct_gene_esm",), |
| ), |
| } |
| provenance = { |
| "config": deepcopy(config), |
| "checkpoint_provenance": checkpoint["provenance"], |
| "prior_transcriptomics_manifest_sha256": declared, |
| "runtime_source_sha256": _runtime_source_hash(), |
| "runtime_environment": _runtime_environment(), |
| "git": _git_state(), |
| } |
| output = Path(direct["output_directory"]) |
| assert not output.exists() |
| output.mkdir(parents=True) |
| for filename, payload in ( |
| ("summary.json", summary), |
| ("paired_comparisons.json", paired), |
| ("truth_report.json", truth_report), |
| ("fit_report.json", checkpoint["report"]), |
| ("provenance.json", provenance), |
| ): |
| (output / filename).write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") |
| for filename, values in ( |
| ("condition_metrics.jsonl", records), |
| ("pathway_metrics.jsonl", pathway_records), |
| ): |
| with (output / filename).open("w") as handle: |
| for value in values: |
| handle.write(json.dumps(value, sort_keys=True) + "\n") |
| return summary, paired, truth_report, provenance |
|
|
