code/causalcelljepa/dynamics.py

# Unpaired, batch-matched population sampling and the locked Stage 2 world model.
# Neither cell-line IDs nor pairwise control/outcome correspondences enter the model or loss.
import json
import math
import time
from collections import Counter
from copy import deepcopy
from hashlib import sha256
from importlib.metadata import version
from pathlib import Path

import h5py
import numpy as np
import torch
import yaml
from geomloss import SamplesLoss
from torch import nn
from torch.nn import functional as F
from torch.utils.data import Dataset

from causalcelljepa.resources import file_sha256
from causalcelljepa.training import (
    _data_loader,
    _git_state,
    _json_log,
    _runtime_environment,
    _runtime_source_hash,
    epoch_order,
    initial_train_state,
    learning_rate_schedule,
    load_checkpoint,
    save_checkpoint,
    seed_everything,
)


class LatentPopulationDataset(Dataset):
    """Sample independent outcome/control sets with identical batch multiplicities."""

    def __init__(
        self,
        cache_path,
        action_path,
        manifest_path,
        mode,
        population_size=32,
        seed=0,
        outcome_role=None,
        control_role=None,
        context="K562",
    ):
        self.cache_path, self.mode, self.context = Path(cache_path), mode, context
        self.population_size, self.seed, self.epoch = population_size, seed, 0
        self.manifest = json.loads(Path(manifest_path).read_text())
        with h5py.File(self.cache_path, "r") as cache:
            roles = cache["role"].asstr()[:]
            self.targets = cache["target"].asstr()[:]
            self.batches = cache["source_batch"].asstr()[:]
            contexts = cache["context"].asstr()[:]
        expected_targets = None
        if outcome_role is None:
            outcome_role = self.manifest["conditions"][mode]["role"]
            expected_targets = self.manifest["conditions"][mode]["targets"]
        control_role = control_role or self.manifest["conditions"]["controls"]["role"]
        outcome_indices = np.flatnonzero(roles == outcome_role)
        control_indices = np.flatnonzero(roles == control_role)
        assert set(contexts[outcome_indices]) == set(contexts[control_indices]) == {context}
        self.condition_targets = sorted(set(self.targets[outcome_indices]))
        assert expected_targets is None or len(self.condition_targets) == expected_targets
        self.outcomes = {
            target: outcome_indices[self.targets[outcome_indices] == target]
            for target in self.condition_targets
        }
        self.controls = {
            batch: control_indices[self.batches[control_indices] == batch]
            for batch in sorted(set(self.batches[control_indices]))
        }
        assert all(len(indices) >= population_size for indices in self.outcomes.values())
        action = torch.load(action_path, map_location="cpu", weights_only=True)
        action_index = {target: index for index, target in enumerate(action["targets"])}
        self.action = {
            target: (action["embedding"][action_index[target]], action["known"][action_index[target]])
            for target in self.condition_targets
        }
        normalization = self.manifest["normalization"]
        self.mean = np.asarray(normalization["latent_mean"], dtype=np.float32)
        self.scale = (
            np.asarray(normalization["latent_std"], dtype=np.float32)
            * normalization["dimension_scale"]
        )
        self._cache = None

    def __len__(self):
        return len(self.condition_targets)

    def set_epoch(self, epoch):
        self.epoch = epoch

    def sample_indices(self, index):
        """Expose deterministic indices for leakage and batch-matching audits."""
        target = self.condition_targets[index]
        generator = np.random.default_rng(
            int.from_bytes(
                sha256(f"{self.seed}\0{self.mode}\0{self.epoch}\0{target}".encode()).digest()[:8],
                "little",
            )
        )
        outcomes = generator.choice(
            self.outcomes[target], self.population_size, replace=False
        )
        controls = np.concatenate(
            [
                generator.choice(self.controls[batch], count, replace=False)
                for batch, count in sorted(Counter(self.batches[outcomes]).items())
            ]
        )
        generator.shuffle(outcomes)
        generator.shuffle(controls)
        return controls, outcomes, target

    def __getitem__(self, index):
        if self._cache is None:
            self._cache = h5py.File(self.cache_path, "r")
        controls, outcomes, target = self.sample_indices(index)
        populations = []
        for indices in (controls, outcomes):
            order = np.argsort(indices)
            latents = self._cache["latent"][indices[order]][np.argsort(order)]
            populations.append(torch.from_numpy((latents - self.mean) / self.scale))
        action, known = self.action[target]
        return {
            "control": populations[0],
            "perturbed": populations[1],
            "action": action,
            "action_known": known,
            "target": target,
        }

    def __getstate__(self):
        state = self.__dict__.copy()
        state["_cache"] = None
        return state


class ConditionalTransitionBlock(nn.Module):
    """Inject the context-action interaction into attention and feed-forward paths."""

    def __init__(self, dim=256, heads=8, ffn_dim=1024, dropout=0.1):
        super().__init__()
        self.condition = nn.Linear(dim, dim)
        self.attention_norm = nn.LayerNorm(dim)
        self.attention = nn.MultiheadAttention(dim, heads, dropout=dropout, batch_first=True)
        self.ffn_norm = nn.LayerNorm(dim)
        self.ffn = nn.Sequential(
            nn.Linear(dim, ffn_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(ffn_dim, dim),
            nn.Dropout(dropout),
        )

    def forward(self, states, interaction):
        condition = self.condition(interaction).unsqueeze(1)
        attended, _ = self.attention(
            self.attention_norm(states + condition),
            self.attention_norm(states + condition),
            self.attention_norm(states + condition),
            need_weights=False,
        )
        states = states + attended
        return states + self.ffn(self.ffn_norm(states + condition))


class ModalityAttentiveActionProjection(nn.Module):
    """Project frozen action teachers and fuse them with sample-specific attention."""

    def __init__(
        self,
        modality_dims,
        output_dim,
        modality_dropout=0.0,
        modality_availability=False,
    ):
        super().__init__()
        assert sum(modality_dims) > 0 and 0 <= modality_dropout < 1
        self.modality_dims = tuple(modality_dims)
        self.modality_dropout = modality_dropout
        self.modality_availability = modality_availability
        self.projectors = nn.ModuleList(
            nn.Sequential(nn.LayerNorm(width), nn.Linear(width, output_dim))
            for width in modality_dims
        )
        self.query = nn.Parameter(torch.randn(output_dim) * 0.02)
        self.output = nn.LayerNorm(output_dim)

    def projected_and_weights(self, action):
        feature_width = sum(self.modality_dims)
        expected_width = feature_width + (
            len(self.modality_dims) if self.modality_availability else 0
        )
        assert action.shape[1] == expected_width
        blocks = action[:, :feature_width].split(self.modality_dims, 1)
        projected = torch.stack(
            [
                projector(block)
                for projector, block in zip(self.projectors, blocks, strict=True)
            ],
            1,
        )
        scores = (torch.tanh(projected) * self.query).sum(2) / math.sqrt(projected.shape[2])
        visible = (
            action[:, feature_width:].bool()
            if self.modality_availability
            else torch.ones_like(scores, dtype=torch.bool)
        )
        if self.training and self.modality_dropout:
            visible &= torch.rand_like(scores) >= self.modality_dropout
        missing = ~visible.any(1)
        if missing.any():
            available = (
                action[:, feature_width:].bool()
                if self.modality_availability
                else torch.ones_like(visible)
            )
            fallback = available.float().argmax(1)
            visible[missing, fallback[missing]] = True
        weights = scores.masked_fill(~visible, -torch.inf).softmax(1)
        return projected, weights

    def forward(self, action):
        projected, weights = self.projected_and_weights(action)
        return self.output((weights.unsqueeze(2) * projected).sum(1))


class ContextConditionedModalityProjection(nn.Module):
    """Fuse available action teachers using an identity-free control-state query."""

    def __init__(self, modality_dims, output_dim, modality_dropout=0.0):
        super().__init__()
        assert sum(modality_dims) > 0 and 0 <= modality_dropout < 1
        self.modality_dims = tuple(modality_dims)
        self.modality_dropout = modality_dropout
        self.projectors = nn.ModuleList(
            nn.Sequential(nn.LayerNorm(width), nn.Linear(width, output_dim))
            for width in modality_dims
        )
        self.query = nn.Parameter(torch.randn(output_dim) * 0.02)
        self.context_query = nn.Sequential(nn.LayerNorm(output_dim), nn.Linear(output_dim, output_dim))
        nn.init.zeros_(self.context_query[-1].weight)
        nn.init.zeros_(self.context_query[-1].bias)
        self.output = nn.LayerNorm(output_dim)

    def projected_and_weights(self, action, context):
        feature_width = sum(self.modality_dims)
        assert action.shape[1] == feature_width + len(self.modality_dims)
        blocks = action[:, :feature_width].split(self.modality_dims, 1)
        availability = action[:, feature_width:].bool()
        projected = torch.stack(
            [
                projector(block)
                for projector, block in zip(self.projectors, blocks, strict=True)
            ],
            1,
        )
        query = self.query.unsqueeze(0) + self.context_query(context)
        scores = (torch.tanh(projected) * query.unsqueeze(1)).sum(2) / math.sqrt(
            projected.shape[2]
        )
        visible = availability.clone()
        if self.training and self.modality_dropout:
            visible &= torch.rand_like(scores) >= self.modality_dropout
        missing = ~visible.any(1)
        if missing.any():
            fallback = availability.float().argmax(1)
            visible[missing, fallback[missing]] = True
        weights = scores.masked_fill(~visible, -torch.inf).softmax(1)
        return projected, weights

    def forward(self, action, context):
        projected, weights = self.projected_and_weights(action, context)
        return self.output((weights.unsqueeze(2) * projected).sum(1))


class PopulationDynamics(nn.Module):
    """Permutation-equivariant residual dynamics conditioned on baseline context and action."""

    def __init__(
        self,
        cell_dim=256,
        action_input_dim=320,
        action_dim=256,
        context_blocks=2,
        transition_blocks=3,
        heads=8,
        ffn_dim=1024,
        dropout=0.1,
        context_mode="set_transformer",
        action_modalities=None,
        action_modality_dropout=0.0,
        action_context_conditioned=False,
        action_modality_availability=False,
    ):
        super().__init__()
        assert context_mode in {"set_transformer", "mean", "none"}
        availability_width = len(action_modalities or ()) if action_modality_availability else 0
        assert action_modalities is None or sum(action_modalities) + availability_width == action_input_dim
        assert not action_context_conditioned or action_modality_availability
        self.context_mode = context_mode
        self.action_context_conditioned = action_context_conditioned
        context_layer = nn.TransformerEncoderLayer(
            cell_dim,
            heads,
            ffn_dim,
            dropout,
            activation="gelu",
            batch_first=True,
            norm_first=True,
        )
        self.context_blocks = nn.TransformerEncoder(
            context_layer,
            context_blocks,
            nn.LayerNorm(cell_dim),
            enable_nested_tensor=False,
        )
        self.pool_query = nn.Parameter(torch.randn(1, cell_dim) * 0.02)
        self.pool = nn.MultiheadAttention(cell_dim, heads, dropout=dropout, batch_first=True)
        self.context_output = nn.LayerNorm(cell_dim)
        if action_context_conditioned:
            self.action_projection = ContextConditionedModalityProjection(
                action_modalities, action_dim, action_modality_dropout
            )
        elif action_modalities:
            self.action_projection = ModalityAttentiveActionProjection(
                action_modalities,
                action_dim,
                action_modality_dropout,
                action_modality_availability,
            )
        else:
            self.action_projection = nn.Sequential(
                nn.LayerNorm(action_input_dim), nn.Linear(action_input_dim, action_dim)
            )
        self.unknown_action = nn.Parameter(torch.randn(action_dim) * 0.02)
        self.interaction = nn.Sequential(
            nn.LayerNorm(4 * cell_dim),
            nn.Linear(4 * cell_dim, ffn_dim),
            nn.GELU(),
            nn.Linear(ffn_dim, cell_dim),
        )
        self.transition = nn.ModuleList(
            ConditionalTransitionBlock(cell_dim, heads, ffn_dim, dropout)
            for _ in range(transition_blocks)
        )
        self.delta = nn.Sequential(nn.LayerNorm(cell_dim), nn.Linear(cell_dim, cell_dim))
        nn.init.zeros_(self.delta[-1].weight)
        nn.init.zeros_(self.delta[-1].bias)

    def forward(self, control, action_embedding, action_known):
        if self.context_mode == "set_transformer":
            encoded = self.context_blocks(control)
            query = self.pool_query.unsqueeze(0).expand(control.shape[0], -1, -1)
            context, _ = self.pool(query, encoded, encoded, need_weights=False)
            context = self.context_output(context.squeeze(1))
        elif self.context_mode == "mean":
            context = self.context_output(control.mean(1))
        else:
            context = torch.zeros_like(control[:, 0])
        projected = (
            self.action_projection(action_embedding, context)
            if self.action_context_conditioned
            else self.action_projection(action_embedding)
        )
        action = torch.where(
            action_known.unsqueeze(1), projected, self.unknown_action.unsqueeze(0)
        )
        interaction = self.interaction(
            torch.cat((context, action, context * action, torch.abs(context - action)), dim=1)
        )
        states = control
        for block in self.transition:
            states = block(states, interaction)
        return control + self.delta(states)


class FrozenLowRankEffectAnchor(nn.Module):
    """Apply a frozen low-rank ESM-to-latent perturbation-effect map."""

    def __init__(self, checkpoint):
        super().__init__()
        assert checkpoint["format_version"] == 1
        assert checkpoint["architecture"] in {
            "esm2_low_rank_latent_effect_ridge",
            "multiteacher_low_rank_latent_effect_ridge",
        }
        for name in ("x_mean", "x_std", "y_mean", "components", "weights"):
            value = checkpoint[name].detach().float()
            assert torch.isfinite(value).all()
            self.register_buffer(name, value)
        input_indices = checkpoint.get("input_indices", torch.arange(len(self.x_mean)))
        input_indices = input_indices.detach().long()
        assert input_indices.ndim == 1 and len(input_indices) == len(self.x_mean)
        assert len(input_indices.unique()) == len(input_indices) and int(input_indices.min()) >= 0
        self.register_buffer("input_indices", input_indices)
        assert self.weights.shape == (len(self.x_mean), len(self.components))
        assert self.components.shape[1] == len(self.y_mean)

    def _load_from_state_dict(
        self,
        state_dict,
        prefix,
        local_metadata,
        strict,
        missing_keys,
        unexpected_keys,
        error_msgs,
    ):
        # Checkpoints frozen before modality-subset anchors did not serialize this buffer.
        key = f"{prefix}input_indices"
        if key not in state_dict:
            state_dict[key] = self.input_indices
        super()._load_from_state_dict(
            state_dict,
            prefix,
            local_metadata,
            strict,
            missing_keys,
            unexpected_keys,
            error_msgs,
        )

    def forward(self, action_embedding, action_known):
        selected = action_embedding.index_select(1, self.input_indices)
        standardized = (selected - self.x_mean) / self.x_std
        predicted = standardized @ self.weights @ self.components + self.y_mean
        return torch.where(action_known.unsqueeze(1), predicted, self.y_mean.unsqueeze(0))


class AnchoredPopulationDynamics(PopulationDynamics):
    """Separate a transferable mean action effect from population heterogeneity."""

    def __init__(
        self,
        *args,
        effect_anchor,
        anchor_gain_max=1.0,
        mean_residual_max_ratio=0.0,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        assert anchor_gain_max >= 1
        assert 0 <= mean_residual_max_ratio <= 1
        self.effect_anchor = FrozenLowRankEffectAnchor(effect_anchor)
        self.anchor_gain_max = float(anchor_gain_max)
        self.mean_residual_max_ratio = float(mean_residual_max_ratio)
        cell_dim = self.effect_anchor.y_mean.numel()
        ffn_dim = self.interaction[1].out_features
        dropout = self.transition[0].ffn[2].p
        self.mean_residual = nn.Sequential(
            nn.LayerNorm(cell_dim),
            nn.Linear(cell_dim, ffn_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(ffn_dim, cell_dim),
        )
        nn.init.zeros_(self.mean_residual[-1].weight)
        nn.init.zeros_(self.mean_residual[-1].bias)
        self.anchor_gain = nn.Sequential(nn.LayerNorm(cell_dim), nn.Linear(cell_dim, 1))
        nn.init.zeros_(self.anchor_gain[-1].weight)
        nn.init.zeros_(self.anchor_gain[-1].bias)
        self.minimum_anchor_norm = float(
            effect_anchor["report"]["training_null_effect_threshold"]
        )
        self.residual_gate_threshold = None

    def configure_residual_gate(self, checkpoint):
        """Attach a source-control-calibrated confidence gate without changing weights."""
        assert checkpoint["format_version"] == 1
        assert checkpoint["architecture"] == "control_population_residual_gate"
        for name in ("center", "scale"):
            value = checkpoint[name].detach().float()
            assert value.shape == self.effect_anchor.y_mean.shape
            self.register_buffer(f"residual_gate_{name}", value, persistent=False)
        self.residual_gate_threshold = float(checkpoint["threshold"])
        self.residual_gate_temperature = float(checkpoint["temperature"])
        assert self.residual_gate_temperature > 0

    def residual_gate_confidence(self, control):
        """Return one smooth residual-retention confidence per control population."""
        assert self.residual_gate_threshold is not None
        score = ((control.mean(1) - self.residual_gate_center) / self.residual_gate_scale).square().mean(1)
        excess = torch.relu(score - self.residual_gate_threshold)
        return torch.exp(-excess / self.residual_gate_temperature)

    def _scaled_anchor(self, action, anchor):
        if self.anchor_gain_max == 1:
            return anchor
        log_limit = math.log(self.anchor_gain_max)
        gain = torch.exp(torch.tanh(self.anchor_gain(action)) * log_limit)
        return anchor * gain

    def _bounded_mean_residual(self, action, anchor):
        if self.mean_residual_max_ratio == 0:
            return torch.zeros_like(anchor)
        residual = self.mean_residual(action)
        residual_norm = torch.linalg.vector_norm(residual, dim=1, keepdim=True)
        anchor_norm = torch.linalg.vector_norm(anchor, dim=1, keepdim=True).clamp_min(
            self.minimum_anchor_norm
        )
        maximum = self.mean_residual_max_ratio * anchor_norm
        scale = torch.minimum(torch.ones_like(residual_norm), maximum / residual_norm.clamp_min(1e-12))
        return residual * scale

    def forward(self, control, action_embedding, action_known):
        if self.context_mode == "set_transformer":
            encoded = self.context_blocks(control)
            query = self.pool_query.unsqueeze(0).expand(control.shape[0], -1, -1)
            context, _ = self.pool(query, encoded, encoded, need_weights=False)
            context = self.context_output(context.squeeze(1))
        elif self.context_mode == "mean":
            context = self.context_output(control.mean(1))
        else:
            context = torch.zeros_like(control[:, 0])
        projected = (
            self.action_projection(action_embedding, context)
            if self.action_context_conditioned
            else self.action_projection(action_embedding)
        )
        action = torch.where(
            action_known.unsqueeze(1), projected, self.unknown_action.unsqueeze(0)
        )
        interaction = self.interaction(
            torch.cat((context, action, context * action, torch.abs(context - action)), dim=1)
        )
        states = control
        for block in self.transition:
            states = block(states, interaction)
        population_residual = self.delta(states)
        population_residual = population_residual - population_residual.mean(1, keepdim=True)
        if self.residual_gate_threshold is not None:
            population_residual *= self.residual_gate_confidence(control)[:, None, None]
        anchor = self._scaled_anchor(
            action, self.effect_anchor(action_embedding, action_known)
        )
        mean_residual = self._bounded_mean_residual(action, anchor)
        return control + anchor.unsqueeze(1) + mean_residual.unsqueeze(1) + population_residual


def dynamics_loss(predicted, observed, control, config, median_distance, null_threshold):
    """Locked unpaired Sinkhorn + MMD + direction + magnitude population objective."""
    blur = config["sinkhorn_blur_ratio"] * median_distance
    sinkhorn = SamplesLoss(
        "sinkhorn",
        p=config["sinkhorn_p"],
        blur=blur,
        debias=config["sinkhorn_debias"],
        backend=config["sinkhorn_backend"],
    )(predicted, observed).mean()
    bandwidth = config["mmd_bandwidth_ratio"] * median_distance
    xx = torch.exp(-torch.cdist(predicted, predicted).square() / (2 * bandwidth**2)).mean((1, 2))
    yy = torch.exp(-torch.cdist(observed, observed).square() / (2 * bandwidth**2)).mean((1, 2))
    xy = torch.exp(-torch.cdist(predicted, observed).square() / (2 * bandwidth**2)).mean((1, 2))
    mmd = (xx + yy - 2 * xy).clamp_min(0).mean()
    true_effect = observed.mean(1) - control.mean(1)
    predicted_effect = predicted.mean(1) - control.mean(1)
    true_magnitude = torch.linalg.vector_norm(true_effect, dim=1)
    predicted_magnitude = torch.linalg.vector_norm(predicted_effect, dim=1)
    active = true_magnitude >= null_threshold
    direction = (
        (
            1
            - F.cosine_similarity(
                predicted_effect[active],
                true_effect[active],
                eps=max(null_threshold, 1e-8),
            )
        ).mean()
        if active.any()
        else predicted.sum() * 0
    )
    magnitude = F.smooth_l1_loss(predicted_magnitude, true_magnitude)
    weights = config["weights"]
    total = (
        weights["sinkhorn"] * sinkhorn
        + weights["mmd"] * mmd
        + weights["direction"] * direction
        + weights["magnitude"] * magnitude
    )
    return {
        "loss": total,
        "sinkhorn": sinkhorn,
        "mmd": mmd,
        "direction": direction,
        "magnitude": magnitude,
        "direction_active": active.float().mean(),
    }


def dynamics_objective(predicted, observed, control, config, statistics):
    """Switch only the training criterion for the locked pseudo-paired comparator."""
    objective = config.get("objective", "unpaired_distribution")
    if objective == "pseudo_paired_mse":
        pointwise = F.mse_loss(predicted, observed)
        return {"loss": pointwise, "pointwise_mse": pointwise}
    assert objective == "unpaired_distribution"
    return dynamics_loss(
        predicted,
        observed,
        control,
        config["loss"],
        statistics["normalization"]["median_training_latent_distance"],
        statistics["direction"]["null_effect_threshold"],
    )


def build_dynamics_model(config):
    """Build the primary model or an explicitly declared post-primary revision."""
    model = config["model"]
    assert model["cell_dim"] == model["action_dim"]
    assert model["predicted_population_size"] == config["data"]["population_size"]
    assert model["cell_line_id"] is False and model["stochastic_noise"] is False
    arguments = (
        model["cell_dim"],
        model["action_input_dim"],
        model["action_dim"],
        model["context_blocks"],
        model["transition_blocks"],
        model["heads"],
        model["ffn_dim"],
        model["dropout"],
        model.get("context_mode", "set_transformer"),
    )
    fusion = {
        "action_modalities": model.get("action_modalities"),
        "action_modality_dropout": model.get("action_modality_dropout", 0.0),
        "action_context_conditioned": model.get("action_context_conditioned", False),
        "action_modality_availability": model.get("action_modality_availability", False),
    }
    architecture = model.get("architecture", "population_dynamics")
    if architecture == "population_dynamics":
        return PopulationDynamics(*arguments, **fusion)
    assert architecture == "anchored_decomposed_population_dynamics"
    anchor = config["effect_anchor"]
    path = Path(anchor["output_path"])
    assert (path.stat().st_size, file_sha256(path)) == (
        anchor["output_bytes"],
        anchor["output_sha256"],
    )
    manifest = json.loads(Path(anchor["manifest_path"]).read_text())
    declared = manifest.pop("manifest_sha256")
    assert declared == anchor["manifest_sha256"] == sha256(
        json.dumps(manifest, sort_keys=True, separators=(",", ":")).encode()
    ).hexdigest()
    assert manifest["artifact"]["sha256"] == anchor["output_sha256"]
    checkpoint = torch.load(path, map_location="cpu", weights_only=True)
    return AnchoredPopulationDynamics(
        *arguments,
        effect_anchor=checkpoint,
        anchor_gain_max=model["anchor_gain_max"],
        mean_residual_max_ratio=model["mean_residual_max_ratio"],
        **fusion,
    )


def dynamics_ablation_configs(path="configs/ablations.yaml"):
    """Materialize the fixed matched-capacity dynamics ablations from one pinned base."""
    path = Path(path)
    specification = yaml.safe_load(path.read_text())
    base_path = Path(specification["base_config_path"])
    assert file_sha256(base_path) == specification["base_config_sha256"]
    base = yaml.safe_load(base_path.read_text())
    configs = {}
    for name, changes in specification["experiments"].items():
        config = deepcopy(base)
        config["model"]["context_mode"] = changes["context_mode"]
        config["loss"]["weights"]["direction"] = changes["direction_weight"]
        config["training"]["output_directory"] = changes["output_directory"]
        config["training"]["resume_from"] = None
        config["ablation"] = {
            "name": name,
            "hypothesis": changes["hypothesis"],
            "base_config_path": str(base_path),
            "base_config_sha256": specification["base_config_sha256"],
        }
        configs[name] = config
    return configs, specification


def dynamics_replication_configs(path="configs/stage2_replication.yaml"):
    """Vary only Stage 2 stochastic seeds while retaining the frozen primary protocol."""
    path = Path(path)
    specification = yaml.safe_load(path.read_text())
    base_path = Path(specification["base_config_path"])
    assert file_sha256(base_path) == specification["base_config_sha256"]
    base = yaml.safe_load(base_path.read_text())
    split = json.loads(Path(base["inputs"]["replogle_manifest_path"]).read_text())["targets"][
        "split"
    ]
    assert split["seed"] == base["seed"] == specification["target_split_seed"]
    configs = {}
    for entry in specification["model_seeds"]:
        seed = entry["seed"]
        assert seed != specification["target_split_seed"] and seed not in configs
        config = deepcopy(base)
        config["seed"] = seed
        config["training"]["output_directory"] = entry["output_directory"]
        config["training"]["resume_from"] = None
        config["replication"] = {
            "model_and_sampling_seed": seed,
            "target_split_seed": specification["target_split_seed"],
            "base_config_path": str(base_path),
            "base_config_sha256": specification["base_config_sha256"],
        }
        configs[seed] = config
    return configs, specification


def _normalized_role_effects(base, roles):
    """Compute batch-matched latent effects for explicitly named K562 roles."""
    inputs = base["inputs"]
    manifest = json.loads(Path(inputs["dynamics_manifest_path"]).read_text())
    normalization = manifest["normalization"]
    mean = np.asarray(normalization["latent_mean"], dtype=np.float32)
    scale = np.asarray(normalization["latent_std"], dtype=np.float32) * normalization[
        "dimension_scale"
    ]
    with h5py.File(inputs["latent_cache_path"], "r") as cache:
        cache_roles = cache["role"].asstr()[:]
        targets = cache["target"].asstr()[:]
        batches = cache["source_batch"].asstr()[:]
        contexts = cache["context"].asstr()[:]
        control_indices = np.flatnonzero(cache_roles == base["data"]["control_role"])
        assert set(contexts[control_indices]) == {base["data"]["context"]} == {"K562"}
        control_means = {
            batch: ((cache["latent"][control_indices[batches[control_indices] == batch]] - mean) / scale).mean(0)
            for batch in sorted(set(batches[control_indices]))
        }
        effects = {}
        target_ids = {}
        for role in roles:
            role_indices = np.flatnonzero(cache_roles == role)
            assert len(role_indices) and set(contexts[role_indices]) == {"K562"}
            target_ids[role] = sorted(set(targets[role_indices]))
            effects[role] = {}
            for target in target_ids[role]:
                selected = role_indices[targets[role_indices] == target]
                outcome = ((cache["latent"][selected] - mean) / scale).mean(0)
                counts = Counter(batches[selected])
                matched_control = sum(
                    count * control_means[batch] for batch, count in counts.items()
                ) / len(selected)
                effects[role][target] = (outcome - matched_control).astype(np.float32)
    return effects, target_ids, manifest


def _action_overridden_config(base, specification):
    """Replace only the frozen action input from a self-hashed cache manifest."""
    if "action_manifest_path" not in specification:
        return base
    manifest = json.loads(Path(specification["action_manifest_path"]).read_text())
    declared = manifest.pop("manifest_sha256")
    assert declared == specification["action_manifest_sha256"] == sha256(
        json.dumps(manifest, sort_keys=True, separators=(",", ":")).encode()
    ).hexdigest()
    artifact = manifest["artifact"]
    assert (Path(artifact["path"]).stat().st_size, file_sha256(artifact["path"])) == (
        artifact["bytes"], artifact["sha256"]
    )
    base["inputs"].update(
        {
            "action_cache_path": artifact["path"],
            "action_cache_bytes": artifact["bytes"],
            "action_cache_sha256": artifact["sha256"],
            "action_manifest_path": specification["action_manifest_path"],
            "action_manifest_sha256": declared,
        }
    )
    base["model"]["action_input_dim"] = artifact["input_dim"]
    base["model"]["action_modalities"] = artifact["modality_dims"]
    base["model"]["action_modality_availability"] = artifact.get(
        "modality_availability", False
    )
    return base


def effect_anchor_input_indices(action, anchor):
    """Select declared action modalities and their availability bits for an anchor."""
    modalities = action.get("modalities")
    modality_dims = action.get("modality_dims")
    selected_names = anchor.get("input_modalities")
    if selected_names is None:
        return torch.arange(action["embedding"].shape[1])
    assert modalities is not None and modality_dims is not None
    assert len(modalities) == len(modality_dims)
    assert len(selected_names) == len(set(selected_names))
    selected = set(selected_names)
    assert selected == selected.intersection(modalities)
    feature_width = sum(modality_dims)
    availability_width = len(modality_dims) if action.get("modality_availability", False) else 0
    assert action["embedding"].shape[1] == feature_width + availability_width
    indices, offset = [], 0
    selected_positions = []
    for position, (name, width) in enumerate(zip(modalities, modality_dims, strict=True)):
        if name in selected:
            indices.extend(range(offset, offset + width))
            selected_positions.append(position)
        offset += width
    if availability_width:
        indices.extend(feature_width + position for position in selected_positions)
    return torch.tensor(indices, dtype=torch.long)


def prepare_effect_anchor(path="configs/anchored_dynamics.yaml"):
    """Fit the frozen ESM-to-latent anchor without test or RPE1 outcomes."""
    path = Path(path)
    specification = yaml.safe_load(path.read_text())
    base_path = Path(specification["base_config_path"])
    assert file_sha256(base_path) == specification["base_config_sha256"]
    base = _action_overridden_config(yaml.safe_load(base_path.read_text()), specification)
    anchor = specification["effect_anchor"]
    roles = (anchor["fit_outcome_role"], anchor["selection_outcome_role"])
    assert roles == ("dynamics_train", "perturbation_ood_validation")
    effects, targets, dynamics_manifest = _normalized_role_effects(base, roles)
    action = torch.load(base["inputs"]["action_cache_path"], map_location="cpu", weights_only=True)
    input_indices = effect_anchor_input_indices(action, anchor)
    action_map = {
        target: (
            action["embedding"][index].index_select(0, input_indices).numpy(),
            bool(action["known"][index]),
        )
        for index, target in enumerate(action["targets"])
    }
    train = effects[roles[0]]
    known_targets = [target for target in targets[roles[0]] 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(anchor["rank"], len(known_targets))]
    standardized = (x - x_mean) / x_std
    scores = (y - y_mean) @ components.T
    gram, cross = standardized.T @ standardized, standardized.T @ scores
    validation = effects[roles[1]]
    validation_targets = targets[roles[1]]
    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 anchor["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
        correlations = [
            float(np.corrcoef(predicted, observed)[0, 1])
            for predicted, observed in zip(prediction, validation_y)
        ]
        magnitude_error = np.abs(
            np.linalg.norm(prediction, axis=1) - np.linalg.norm(validation_y, axis=1)
        )
        candidates.append(
            {
                "alpha": float(alpha),
                "mse": float(np.mean(np.square(prediction - validation_y))),
                "mean_effect_pearson": float(np.mean(correlations)),
                "mean_magnitude_absolute_error": float(np.mean(magnitude_error)),
                "weights": weights,
            }
        )
    selected = min(candidates, key=lambda item: (item["mse"], item["alpha"]))
    report = {
        "fit_outcome_role": roles[0],
        "selection_outcome_role": roles[1],
        "fit_targets": len(train),
        "fit_targets_with_known_action": len(known_targets),
        "selection_targets": len(validation_targets),
        "rank": len(components),
        "input_modalities": anchor.get("input_modalities", action.get("modalities")),
        "input_dimensions": len(input_indices),
        "input_indices_sha256": sha256(
            json.dumps(input_indices.tolist(), separators=(",", ":")).encode()
        ).hexdigest(),
        "selected_ridge": selected["alpha"],
        "selection_mse": selected["mse"],
        "selection_mean_effect_pearson": selected["mean_effect_pearson"],
        "selection_mean_magnitude_absolute_error": selected[
            "mean_magnitude_absolute_error"
        ],
        "ridge_candidates": [item["alpha"] for item in candidates],
        "ridge_validation_mse": [item["mse"] for item in candidates],
        "training_null_effect_threshold": dynamics_manifest["direction"][
            "null_effect_threshold"
        ],
        "target_sha256": {
            role: sha256("\n".join(targets[role]).encode()).hexdigest() for role in roles
        },
        "leakage": {
            "contexts": ["K562"],
            "sealed_test_outcomes_used": False,
            "rpe1_outcomes_used": False,
        },
    }
    checkpoint = {
        "format_version": 1,
        "architecture": anchor.get(
            "checkpoint_architecture", "esm2_low_rank_latent_effect_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(selected["weights"].astype(np.float32)),
        "input_indices": input_indices,
        "report": report,
    }
    output = Path(anchor["output_path"])
    output.parent.mkdir(parents=True, exist_ok=True)
    torch.save(checkpoint, output)
    artifact = {"path": str(output), "bytes": output.stat().st_size, "sha256": file_sha256(output)}
    manifest = {
        "format_version": 1,
        "architecture": checkpoint["architecture"],
        "artifact": artifact,
        "fit": report,
        "source": {
            "base_config_path": str(base_path),
            "base_config_sha256": specification["base_config_sha256"],
            "latent_cache_sha256": base["inputs"]["latent_cache_sha256"],
            "action_cache_sha256": base["inputs"]["action_cache_sha256"],
            "dynamics_manifest_sha256": dynamics_manifest["manifest_sha256"],
        },
    }
    manifest["manifest_sha256"] = sha256(
        json.dumps(manifest, sort_keys=True, separators=(",", ":")).encode()
    ).hexdigest()
    Path(anchor["manifest_path"]).write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
    return manifest


def anchored_dynamics_configs(path="configs/anchored_dynamics.yaml"):
    """Materialize the preregistered validation-only anchored candidates."""
    path = Path(path)
    specification = yaml.safe_load(path.read_text())
    base_path = Path(specification["base_config_path"])
    assert file_sha256(base_path) == specification["base_config_sha256"]
    base = _action_overridden_config(yaml.safe_load(base_path.read_text()), specification)
    anchor = specification["effect_anchor"]
    assert anchor["output_sha256"] != "PENDING" and anchor["manifest_sha256"] != "PENDING"
    artifact_path = Path(anchor["output_path"])
    assert (artifact_path.stat().st_size, file_sha256(artifact_path)) == (
        anchor["output_bytes"],
        anchor["output_sha256"],
    )
    manifest = json.loads(Path(anchor["manifest_path"]).read_text())
    declared = manifest.pop("manifest_sha256")
    assert declared == anchor["manifest_sha256"] == sha256(
        json.dumps(manifest, sort_keys=True, separators=(",", ":")).encode()
    ).hexdigest()
    configs = {}
    for name, experiment in specification["experiments"].items():
        config = deepcopy(base)
        config["seed"] = specification["seed"]
        config["effect_anchor"] = deepcopy(anchor)
        config["model"]["architecture"] = "anchored_decomposed_population_dynamics"
        config["model"]["anchor_gain_max"] = experiment["anchor_gain_max"]
        config["model"]["mean_residual_max_ratio"] = experiment[
            "mean_residual_max_ratio"
        ]
        config["model"]["action_modality_dropout"] = experiment.get(
            "action_modality_dropout", 0.0
        )
        config["model"]["action_context_conditioned"] = experiment.get(
            "action_context_conditioned", False
        )
        config["training"]["output_directory"] = experiment["output_directory"]
        config["training"]["resume_from"] = None
        config["revision"] = {
            **deepcopy(specification["revision"]),
            "candidate": name,
            "anchor_gain_max": experiment["anchor_gain_max"],
            "mean_residual_max_ratio": experiment["mean_residual_max_ratio"],
            "action_modality_dropout": experiment.get("action_modality_dropout", 0.0),
            "action_context_conditioned": experiment.get("action_context_conditioned", False),
            "target_split_seed": base["seed"],
        }
        configs[name] = config
    return configs, specification


def anchored_selected_entry(training_manifest, selection_manifest):
    """Resolve the frozen candidate while proving selection preceded sealed evaluation."""
    assert training_manifest["protocol"]["sealed_test_outcomes_used_for_fit_or_selection"] is False
    assert training_manifest["protocol"]["rpe1_perturbed_outcomes_used_for_fit_or_selection"] is False
    assert selection_manifest["leakage"] == {
        "context": "K562",
        "outcome_role": "perturbation_ood_validation",
        "rpe1_outcomes_used": False,
        "sealed_test_outcomes_used": False,
    }
    selected = selection_manifest["selected"]
    entry = training_manifest["artifacts"]["candidates"][selected["candidate"]]
    assert entry["best_checkpoint"] == {
        key: selected[key] for key in ("bytes", "path", "sha256")
    }
    assert (selected["best_validation_epoch"], selected["best_validation_loss"]) == (
        entry["full_run"]["best_validation_epoch"],
        entry["full_run"]["best_validation_loss"],
    )
    return selected["candidate"], entry


def select_multiteacher_candidate(training_manifest, specification):
    """Apply the frozen validation-loss margin without consulting test outcomes."""
    protocol = training_manifest["protocol"]
    assert protocol["sealed_test_outcomes_used_for_fit_or_selection"] is False
    assert protocol["rpe1_perturbed_outcomes_used_for_fit_or_selection"] is False
    rule = specification["selection"]
    assert rule["viewed_test_outcomes_used"] is False
    assert rule["checkpoint_rule"] == "minimum_original_latent_validation_loss"
    fallback = rule["fallback_candidate"]
    candidates = training_manifest["artifacts"]["candidates"]
    assert fallback in candidates
    alternatives = [
        name
        for name, experiment in specification["experiments"].items()
        if experiment.get("action_modality_dropout", 0.0) > 0
    ]
    assert len(alternatives) == 1 and alternatives[0] in candidates
    alternative = alternatives[0]
    fallback_loss = candidates[fallback]["full_run"]["best_validation_loss"]
    alternative_loss = candidates[alternative]["full_run"]["best_validation_loss"]
    improvement = fallback_loss - alternative_loss
    selected = (
        alternative
        if improvement >= rule["dropout_minimum_loss_improvement"]
        else fallback
    )
    return selected, improvement


def prepare_dynamics_manifest(config, config_path="configs/dynamics.yaml"):
    """Fit state-space statistics without reading validation, test, or RPE1 outcomes."""
    inputs, data, normalization = config["inputs"], config["data"], config["normalization"]
    for kind in ("latent", "action"):
        path = Path(inputs[f"{kind}_cache_path"])
        assert (path.stat().st_size, file_sha256(path)) == (
            inputs[f"{kind}_cache_bytes"],
            inputs[f"{kind}_cache_sha256"],
        )
    for kind in ("action", "replogle"):
        manifest = json.loads(Path(inputs[f"{kind}_manifest_path"]).read_text())
        assert manifest["manifest_sha256"] == inputs[f"{kind}_manifest_sha256"]

    with h5py.File(inputs["latent_cache_path"], "r") as cache:
        roles = cache["role"].asstr()[:]
        targets = cache["target"].asstr()[:]
        contexts = cache["context"].asstr()[:]
        batches = cache["source_batch"].asstr()[:]
        cell_ids = cache["cell_id"].asstr()[:]
        fit_indices = np.flatnonzero(np.isin(roles, normalization["fit_roles"]))
        fit_latents = cache["latent"][fit_indices]
        mean = fit_latents.mean(0, dtype=np.float64)
        variance = np.square(fit_latents.astype(np.float64) - mean).mean(0)
        std = np.sqrt(variance)
        assert np.isfinite(mean).all() and np.isfinite(std).all() and (std > 0).all()
        dimension_scale = math.sqrt(fit_latents.shape[1])
        rng = np.random.default_rng(config["seed"])
        distance_indices = np.sort(
            rng.choice(fit_indices, normalization["distance_sample_cells"], replace=False)
        )
        distance_latents = (cache["latent"][distance_indices] - mean) / std / dimension_scale
        distances = np.linalg.norm(
            distance_latents - distance_latents[rng.permutation(len(distance_latents))], axis=1
        )
        median_distance = float(np.median(distances))

        control_indices = np.flatnonzero(roles == data["control_role"])
        assert set(contexts[control_indices]) == {data["context"]}
        control_latents = (cache["latent"][control_indices] - mean) / std / dimension_scale
        control_batches = batches[control_indices]
        control_means = {
            batch: control_latents[control_batches == batch].mean(0)
            for batch in sorted(set(control_batches))
        }
        train_indices = np.flatnonzero(roles == data["train_outcome_role"])
        validation_indices = np.flatnonzero(roles == data["validation_outcome_role"])
        assert set(contexts[train_indices]) == set(contexts[validation_indices]) == {
            data["context"]
        }
        train_latents = (cache["latent"][train_indices] - mean) / std / dimension_scale
        train_targets, train_batches = targets[train_indices], batches[train_indices]
        effect_norms = {}
        for target in sorted(set(train_targets)):
            selected = train_targets == target
            batch_counts = Counter(train_batches[selected])
            matched_control = sum(
                count * control_means[batch] for batch, count in batch_counts.items()
            ) / selected.sum()
            effect_norms[target] = float(
                np.linalg.norm(train_latents[selected].mean(0) - matched_control)
            )

    threshold = float(
        np.quantile(
            list(effect_norms.values()), normalization["null_effect_quantile"], method="linear"
        )
    )
    replogle = json.loads(Path(inputs["replogle_manifest_path"]).read_text())
    split = replogle["targets"]["split"]["targets"]
    assert set(effect_norms) == set(split["train"])
    assert set(targets[validation_indices]) == set(split["validation"])
    report = {
        "format_version": 1,
        "config_sha256": file_sha256(config_path),
        "inputs": {
            key: value for key, value in inputs.items() if key.endswith(("_bytes", "_sha256"))
        },
        "normalization": {
            "fit_roles": normalization["fit_roles"],
            "fit_cells": len(fit_indices),
            "method": normalization["method"],
            "latent_mean": mean.tolist(),
            "latent_std": std.tolist(),
            "dimension_scale": dimension_scale,
            "distance_sample_cells": len(distance_indices),
            "distance_sample_cell_ids_sha256": sha256(
                "\n".join(sorted(cell_ids[distance_indices])).encode()
            ).hexdigest(),
            "median_training_latent_distance": median_distance,
        },
        "conditions": {
            "train": {
                "role": data["train_outcome_role"],
                "targets": len(set(targets[train_indices])),
                "cells": len(train_indices),
                "cells_per_target": dict(sorted(Counter(targets[train_indices]).items())),
            },
            "validation": {
                "role": data["validation_outcome_role"],
                "targets": len(set(targets[validation_indices])),
                "cells": len(validation_indices),
                "cells_per_target": dict(sorted(Counter(targets[validation_indices]).items())),
            },
            "controls": {
                "role": data["control_role"],
                "cells": len(control_indices),
                "batches": dict(sorted(Counter(control_batches).items())),
            },
        },
        "direction": {
            "null_effect_quantile": normalization["null_effect_quantile"],
            "null_effect_threshold": threshold,
            "training_effect_norms": effect_norms,
            "excluded_training_targets": sum(value < threshold for value in effect_norms.values()),
        },
        "leakage": {
            "normalization_roles": sorted(set(roles[fit_indices])),
            "outcome_context": sorted(set(contexts[train_indices])),
            "validation_used_for_statistics": False,
            "rpe1_outcomes_used_for_statistics": False,
            "sealed_test_outcomes_used_for_statistics": False,
        },
        "runtime": {"numpy": np.__version__, "h5py": version("h5py")},
    }
    report["manifest_sha256"] = sha256(
        json.dumps(report, sort_keys=True, separators=(",", ":")).encode()
    ).hexdigest()
    Path(inputs["dynamics_manifest_path"]).write_text(
        json.dumps(report, indent=2, sort_keys=True) + "\n"
    )
    return report


def state_ablation_config(path):
    """Materialize a frozen alternative state space into the locked Stage 2 model."""
    path = Path(path)
    specification = yaml.safe_load(path.read_text())
    base_path = Path(specification["base_config_path"])
    assert file_sha256(base_path) == specification["base_config_sha256"]
    protocol = json.loads(Path(specification["specification_manifest_path"]).read_text())
    protocol_sha256 = protocol.pop("manifest_sha256")
    assert protocol_sha256 == specification["specification_manifest_sha256"] == sha256(
        json.dumps(protocol, sort_keys=True, separators=(",", ":")).encode()
    ).hexdigest()
    state_manifest_path = Path(specification["cache_manifest_path"])
    state_manifest = json.loads(state_manifest_path.read_text())
    declared = state_manifest.pop("manifest_sha256")
    assert declared == sha256(
        json.dumps(state_manifest, sort_keys=True, separators=(",", ":")).encode()
    ).hexdigest()
    artifact = state_manifest["artifact"]
    assert (Path(artifact["path"]).stat().st_size, file_sha256(artifact["path"])) == (
        artifact["bytes"],
        artifact["sha256"],
    )
    assert artifact["cell_dim"] == yaml.safe_load(base_path.read_text())["model"]["cell_dim"]
    assert state_manifest["source"]["config_sha256"] == file_sha256(path)
    assert (
        state_manifest["source"]["specification_manifest_sha256"]
        == specification["specification_manifest_sha256"]
    )
    dynamics_manifest = json.loads(Path(specification["dynamics_manifest_path"]).read_text())
    dynamics_sha256 = dynamics_manifest.pop("manifest_sha256")
    assert dynamics_sha256 == state_manifest["source"]["dynamics_manifest_sha256"] == sha256(
        json.dumps(dynamics_manifest, sort_keys=True, separators=(",", ":")).encode()
    ).hexdigest()
    assert dynamics_manifest["config_sha256"] == file_sha256(path)
    assert dynamics_manifest["inputs"]["latent_cache_sha256"] == artifact["sha256"]
    config = yaml.safe_load(base_path.read_text())
    config["inputs"].update(
        {
            "latent_cache_path": artifact["path"],
            "latent_cache_bytes": artifact["bytes"],
            "latent_cache_sha256": artifact["sha256"],
            "latent_manifest_path": str(state_manifest_path),
            "latent_manifest_sha256": declared,
            "dynamics_manifest_path": specification["dynamics_manifest_path"],
        }
    )
    config["training"]["output_directory"] = specification["output_directory"]
    config["training"]["resume_from"] = None
    config["ablation"] = {
        "name": specification["name"],
        "hypothesis": specification["hypothesis"],
        "base_config_path": str(base_path),
        "base_config_sha256": specification["base_config_sha256"],
        "cache_manifest_sha256": declared,
    }
    return config, specification, {**state_manifest, "manifest_sha256": declared}


def learned_target_id_config(path="configs/learned_target_id.yaml"):
    """Materialize the categorical-action comparator from its pinned cache manifest."""
    path = Path(path)
    specification = yaml.safe_load(path.read_text())
    base_path = Path(specification["base_config_path"])
    assert file_sha256(base_path) == specification["base_config_sha256"]
    base = yaml.safe_load(base_path.read_text())
    cache_manifest_path = Path(specification["cache_manifest_path"])
    cache_manifest = json.loads(cache_manifest_path.read_text())
    declared = cache_manifest.pop("manifest_sha256")
    assert declared == sha256(
        json.dumps(cache_manifest, sort_keys=True, separators=(",", ":")).encode()
    ).hexdigest()
    artifact = cache_manifest["artifact"]
    assert artifact["path"] == specification["action_cache_path"]
    assert Path(artifact["path"]).stat().st_size == artifact["bytes"]
    assert file_sha256(artifact["path"]) == artifact["sha256"]
    assert cache_manifest["source"]["config_sha256"] == file_sha256(path)
    assert (
        cache_manifest["source"]["specification_manifest_sha256"]
        == specification["specification_manifest_sha256"]
    )
    config = deepcopy(base)
    config["inputs"].update(
        {
            "action_cache_path": artifact["path"],
            "action_cache_bytes": artifact["bytes"],
            "action_cache_sha256": artifact["sha256"],
            "action_manifest_path": str(cache_manifest_path),
            "action_manifest_sha256": declared,
        }
    )
    config["model"]["action_input_dim"] = artifact["input_dim"]
    config["training"]["output_directory"] = specification["output_directory"]
    config["training"]["resume_from"] = None
    config["ablation"] = {
        "name": "learned_target_id",
        "hypothesis": "structured ESM-2 actions enable unseen-target transfer",
        "base_config_path": str(base_path),
        "base_config_sha256": specification["base_config_sha256"],
        "specification_manifest_sha256": specification["specification_manifest_sha256"],
        "cache_manifest_sha256": declared,
    }
    return config, specification, {**cache_manifest, "manifest_sha256": declared}


def dynamics_provenance(config, config_path="configs/dynamics.yaml"):
    """Bind checkpoints to exact caches, manifests, configuration, code, and runtime."""
    inputs = config["inputs"]
    for kind in ("latent", "action"):
        assert file_sha256(inputs[f"{kind}_cache_path"]) == inputs[f"{kind}_cache_sha256"]
    manifests = {}
    for kind in ("replogle", "action", "dynamics"):
        path = inputs[f"{kind}_manifest_path"]
        payload = json.loads(Path(path).read_text())
        declared = payload.pop("manifest_sha256")
        assert declared == sha256(
            json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()
        ).hexdigest()
        if f"{kind}_manifest_sha256" in inputs:
            assert declared == inputs[f"{kind}_manifest_sha256"]
        manifests[kind] = declared
    if "latent_manifest_path" in inputs:
        payload = json.loads(Path(inputs["latent_manifest_path"]).read_text())
        declared = payload.pop("manifest_sha256")
        assert declared == inputs["latent_manifest_sha256"] == sha256(
            json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()
        ).hexdigest()
        manifests["latent"] = declared
    cache_sha256 = {
        "latent": inputs["latent_cache_sha256"],
        "action": inputs["action_cache_sha256"],
    }
    if "effect_anchor" in config:
        anchor = config["effect_anchor"]
        assert file_sha256(anchor["output_path"]) == anchor["output_sha256"]
        payload = json.loads(Path(anchor["manifest_path"]).read_text())
        declared = payload.pop("manifest_sha256")
        assert declared == anchor["manifest_sha256"] == sha256(
            json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()
        ).hexdigest()
        manifests["effect_anchor"] = declared
        cache_sha256["effect_anchor"] = anchor["output_sha256"]
    return {
        "config_sha256": file_sha256(config_path),
        "cache_sha256": cache_sha256,
        "manifest_sha256": manifests,
        "runtime_source_sha256": _runtime_source_hash(),
        "runtime_environment": _runtime_environment(),
        "geomloss": version("geomloss"),
        "git": _git_state(),
    }


def _device_batch(batch, device):
    non_blocking = device.type == "cuda"
    return {
        key: batch[key].to(device, non_blocking=non_blocking)
        for key in ("control", "perturbed", "action", "action_known")
    }


@torch.no_grad()
def validate_dynamics(model, dataset, config, device, max_batches=None):
    """Evaluate only frozen K562 perturbation-OOD validation conditions."""
    was_training = model.training
    model.eval()
    dataset.set_epoch(-1)
    training = config["training"]
    loader = _data_loader(
        dataset,
        range(len(dataset)),
        training["batch_size"],
        training["num_workers"],
        config["seed"] + 1,
        training["pin_memory"] and device.type == "cuda",
        training["persistent_workers"],
        training["prefetch_factor"],
    )
    statistics = dataset.manifest
    totals, conditions = {}, 0
    for index, raw_batch in enumerate(loader):
        if max_batches is not None and index >= max_batches:
            break
        batch = _device_batch(raw_batch, device)
        predicted = model(batch["control"], batch["action"], batch["action_known"])
        losses = dynamics_objective(
            predicted, batch["perturbed"], batch["control"], config, statistics
        )
        size = len(raw_batch["target"])
        conditions += size
        for name, value in losses.items():
            totals[name] = totals.get(name, 0.0) + float(value.detach()) * size
    model.train(was_training)
    return {name: value / conditions for name, value in totals.items()}


def train_dynamics(
    config,
    device,
    max_steps=None,
    validation_batches=None,
    config_path="configs/dynamics.yaml",
):
    """Train or exactly resume condition-weighted Stage 2 dynamics."""
    training, data = config["training"], config["data"]
    assert training["optimizer"] == "AdamW"
    seed_everything(config["seed"], training["deterministic_algorithms"])
    datasets = {
        mode: LatentPopulationDataset(
            config["inputs"]["latent_cache_path"],
            config["inputs"]["action_cache_path"],
            config["inputs"]["dynamics_manifest_path"],
            mode,
            data["population_size"],
            config["seed"],
        )
        for mode in ("train", "validation")
    }
    model = build_dynamics_model(config).to(device)
    optimizer = torch.optim.AdamW(
        model.parameters(), lr=training["learning_rate"], weight_decay=training["weight_decay"]
    )
    steps_per_epoch = math.ceil(len(datasets["train"]) / training["batch_size"])
    total_steps = training["epochs"] * steps_per_epoch
    scheduler = torch.optim.lr_scheduler.LambdaLR(
        optimizer,
        learning_rate_schedule(
            total_steps,
            training["warmup_fraction"],
            training["minimum_learning_rate_fraction"],
        ),
    )
    provenance = dynamics_provenance(config, config_path)
    configuration = deepcopy(config)
    output = Path(training["output_directory"])
    output.mkdir(parents=True, exist_ok=True)
    latest, best = output / "latest.pt", output / "best.pt"
    state = initial_train_state()
    if training["resume_from"]:
        state, saved_configuration = load_checkpoint(
            training["resume_from"], model, optimizer, scheduler, provenance, device
        )
        saved_configuration["training"]["resume_from"] = None
        current = deepcopy(configuration)
        current["training"]["resume_from"] = None
        assert saved_configuration == current
        if state["complete"]:
            model.load_state_dict(torch.load(best, map_location=device, weights_only=False)["model"])
            return model, state, {
                "train_conditions": len(datasets["train"]),
                "validation_conditions": len(datasets["validation"]),
                "elapsed_seconds": 0.0,
                "checkpoint": str(latest),
                "best_checkpoint": str(best),
                "already_complete": True,
            }
    statistics = datasets["train"].manifest
    started = time.perf_counter()
    model.train()
    for epoch in range(state["epoch"], training["epochs"]):
        datasets["train"].set_epoch(epoch)
        order = epoch_order(range(len(datasets["train"])), config["seed"], epoch)
        start = state["batch_in_epoch"] * training["batch_size"]
        loader = _data_loader(
            datasets["train"],
            order[start:],
            training["batch_size"],
            training["num_workers"],
            config["seed"] + 2 + epoch,
            training["pin_memory"] and device.type == "cuda",
            training["persistent_workers"],
            training["prefetch_factor"],
        )
        for offset, raw_batch in enumerate(loader, start=state["batch_in_epoch"]):
            batch = _device_batch(raw_batch, device)
            optimizer.zero_grad(set_to_none=True)
            predicted = model(batch["control"], batch["action"], batch["action_known"])
            losses = dynamics_objective(
                predicted, batch["perturbed"], batch["control"], config, statistics
            )
            if not all(torch.isfinite(value) for value in losses.values()):
                raise FloatingPointError(
                    f"Non-finite dynamics loss at optimizer step {state['global_step']}"
                )
            losses["loss"].backward()
            gradient_norm = torch.nn.utils.clip_grad_norm_(
                model.parameters(), training["gradient_clip_norm"], error_if_nonfinite=True
            )
            optimizer.step()
            scheduler.step()
            state["global_step"] += 1
            state["batch_in_epoch"] = offset + 1
            _json_log(
                output / "training.jsonl",
                {
                    "event": "train_step",
                    "epoch": epoch,
                    "global_step": state["global_step"],
                    "conditions": len(raw_batch["target"]),
                    **{name: float(value.detach()) for name, value in losses.items()},
                    "gradient_norm": float(gradient_norm),
                    "learning_rate": scheduler.get_last_lr()[0],
                },
            )
            if state["global_step"] % training["checkpoint_every_steps"] == 0:
                save_checkpoint(
                    latest, model, optimizer, scheduler, state, configuration, provenance
                )
            if max_steps is not None and state["global_step"] >= max_steps:
                save_checkpoint(
                    latest, model, optimizer, scheduler, state, configuration, provenance
                )
                return model, state, {
                    "train_conditions": len(datasets["train"]),
                    "validation_conditions": len(datasets["validation"]),
                    "elapsed_seconds": time.perf_counter() - started,
                    "checkpoint": str(latest),
                }
        validation = validate_dynamics(
            model, datasets["validation"], config, device, validation_batches
        )
        _json_log(
            output / "training.jsonl",
            {"event": "validation", "epoch": epoch, "global_step": state["global_step"], **validation},
        )
        improved = validation["loss"] < state["best_validation_loss"]
        state["best_validation_loss"] = min(state["best_validation_loss"], validation["loss"])
        if improved:
            state["best_validation_epoch"] = epoch
        state["epochs_without_improvement"] = 0 if improved else state["epochs_without_improvement"] + 1
        state["epoch"], state["batch_in_epoch"] = epoch + 1, 0
        save_checkpoint(latest, model, optimizer, scheduler, state, configuration, provenance)
        if improved:
            save_checkpoint(best, model, optimizer, scheduler, state, configuration, provenance)
        if state["epochs_without_improvement"] >= training["early_stopping_patience"]:
            state["complete"], state["completion_reason"] = True, "early_stopping"
            break
    else:
        state["complete"], state["completion_reason"] = True, "configured_epochs"
    save_checkpoint(latest, model, optimizer, scheduler, state, configuration, provenance)
    model.load_state_dict(torch.load(best, map_location=device, weights_only=False)["model"])
    return model, state, {
        "train_conditions": len(datasets["train"]),
        "validation_conditions": len(datasets["validation"]),
        "elapsed_seconds": time.perf_counter() - started,
        "checkpoint": str(latest),
        "best_checkpoint": str(best),
    }