runtime/src/resynthesis/molecular_geometry.py

"""All-atom pair, coordinate diffusion, affinity, and validity science heads.

Additive tensor-native modules inspired by Chai / AF3 / Boltz-2 / BioMatrix
patterns — not a copy of those systems. Heads attach beside the existing
text-chemistry path; zero-init residuals preserve Q/K/V identity until trained.

Hot-path contracts stay tensor-native (no Dict[str, Any] payloads).
Boundary receipts are host JSON only.
"""

from __future__ import annotations

import hashlib
from dataclasses import dataclass
from typing import cast

import torch
import torch.nn as nn
import torch.nn.functional as F

from resynthesis.config import RESYNTHESIS_HIDDEN_SIZE

MOLECULAR_INITIALIZATION_SCHEME = (
    "sha256_role_seeded_xavier_identity_residual_v1"
)


def _role_seeded_xavier_uniform_(
    tensor: torch.Tensor,
    role: str,
) -> None:
    """Initialize one growth tensor reproducibly without changing global RNG."""

    if tensor.device.type == "meta":
        return
    seed = int.from_bytes(
        hashlib.sha256(
            f"resynthesis.molecular.v1:{role}".encode("utf-8")
        ).digest()[:8],
        byteorder="little",
        signed=False,
    )
    generator = torch.Generator(device=tensor.device)
    generator.manual_seed(seed)
    nn.init.xavier_uniform_(tensor, generator=generator)


@dataclass(frozen=True)
class MolecularGeometryConfig:
    """Seed geometry — not caps (uncapped-policy: intentional)."""

    hidden_size: int = RESYNTHESIS_HIDDEN_SIZE
    pair_dim: int = 128
    atom_dim: int = 64
    diffusion_steps_seed: int = 8
    affinity_hidden: int = 256


@dataclass(frozen=True)
class MolecularInputPacket:
    """Target-free tensor input for native molecular geometry participation.

    ``noisy_coordinates`` are model inputs.  The clean coordinates and sampled
    noise target are deliberately absent: those remain at the trainer's loss
    boundary and cannot influence RBO, Fabric, attention, or expert routing.
    """

    atomic_numbers: torch.Tensor
    noisy_coordinates: torch.Tensor
    atom_mask: torch.Tensor
    diffusion_time: torch.Tensor

    def validated(self) -> "MolecularInputPacket":
        if self.atomic_numbers.ndim != 2:
            raise ValueError("molecular atomic numbers expect [batch, atoms]")
        if self.noisy_coordinates.shape != (*self.atomic_numbers.shape, 3):
            raise ValueError("molecular noisy coordinates expect [batch, atoms, 3]")
        if self.atom_mask.shape != self.atomic_numbers.shape:
            raise ValueError("molecular atom mask geometry differs")
        if self.diffusion_time.shape != (self.atomic_numbers.shape[0],):
            raise ValueError("molecular diffusion time expects [batch]")
        if self.atomic_numbers.dtype != torch.long:
            raise ValueError("molecular atomic numbers must be int64")
        if self.atom_mask.dtype != torch.bool:
            raise ValueError("molecular atom mask must be boolean")
        if not self.diffusion_time.is_floating_point():
            raise ValueError("molecular diffusion time must be floating point")
        torch._assert_async(
            ((self.atomic_numbers >= 0) & (self.atomic_numbers <= 118)).all(),
            "molecular atomic number is outside the periodic table",
        )
        torch._assert_async(
            self.atom_mask.any(dim=-1).all(),
            "molecular input contains no active atoms",
        )
        torch._assert_async(
            torch.isfinite(self.diffusion_time).all()
            & (self.diffusion_time >= 0).all()
            & (self.diffusion_time <= 1).all(),
            "molecular diffusion time is outside [0, 1]",
        )
        return self


@dataclass(frozen=True)
class MolecularSciencePacket:
    """Tensor-native molecular science outputs (no string-key hot-path dict)."""

    pair: torch.Tensor
    binder_logits: torch.Tensor
    potency: torch.Tensor
    developability: torch.Tensor
    validity: torch.Tensor
    physical_scores: torch.Tensor
    clash: torch.Tensor
    coord_noise: torch.Tensor
    diffusion_loss: torch.Tensor
    vibrational_spectrum: torch.Tensor

    def as_boundary_dict(self) -> dict[str, torch.Tensor]:
        """Explicit serialization adapter for receipts / logs only."""

        return {
            "pair": self.pair,
            "binder_logits": self.binder_logits,
            "potency": self.potency,
            "developability": self.developability,
            "validity": self.validity,
            "physical_scores": self.physical_scores,
            "clash": self.clash,
            "coord_noise": self.coord_noise,
            "diffusion_loss": self.diffusion_loss,
            "vibrational_spectrum": self.vibrational_spectrum,
        }


class AtomFeatureEncoder(nn.Module):
    """Encode per-atom/residue tokens into an atom latent."""

    def __init__(self, cfg: MolecularGeometryConfig | None = None) -> None:
        super().__init__()
        self.cfg = cfg or MolecularGeometryConfig()
        self.proj = nn.Linear(self.cfg.hidden_size, self.cfg.atom_dim, bias=False)
        _role_seeded_xavier_uniform_(self.proj.weight, "atom_encoder.proj.weight")

    def forward(self, hidden_t: torch.Tensor) -> torch.Tensor:
        if hidden_t.ndim != 3:
            raise ValueError("atom encoder expects [batch, tokens, hidden]")
        return cast(torch.Tensor, self.proj(hidden_t))


class AllAtomPairRepresentation(nn.Module):
    """Unified residue/atom pair features (AF3-style pair plane, additive).

    Builds ``pair = φ(a_i) + ψ(a_j) + outer`` without truncating the token
    graph. Pair dim is a seed; sequence extent is uncapped.
    """

    def __init__(self, cfg: MolecularGeometryConfig | None = None) -> None:
        super().__init__()
        self.cfg = cfg or MolecularGeometryConfig()
        self.left = nn.Linear(self.cfg.atom_dim, self.cfg.pair_dim, bias=False)
        self.right = nn.Linear(self.cfg.atom_dim, self.cfg.pair_dim, bias=False)
        self.outer = nn.Linear(self.cfg.atom_dim, self.cfg.pair_dim, bias=False)
        self.pair_update = nn.Linear(self.cfg.pair_dim, self.cfg.pair_dim, bias=False)
        _role_seeded_xavier_uniform_(self.left.weight, "pair_rep.left.weight")
        _role_seeded_xavier_uniform_(self.right.weight, "pair_rep.right.weight")
        _role_seeded_xavier_uniform_(self.outer.weight, "pair_rep.outer.weight")
        _role_seeded_xavier_uniform_(
            self.pair_update.weight,
            "pair_rep.pair_update.weight",
        )

    def forward(self, atom_t: torch.Tensor) -> torch.Tensor:
        if atom_t.ndim != 3:
            raise ValueError("pair representation expects [batch, atoms, atom_dim]")
        left_t = self.left(atom_t).unsqueeze(2)
        right_t = self.right(atom_t).unsqueeze(1)
        outer_scalar_t = torch.einsum("bid,bjd->bij", atom_t, atom_t).unsqueeze(-1)
        pair_t = left_t + right_t + outer_scalar_t
        return cast(torch.Tensor, self.pair_update(pair_t))


class CoordinateDiffusionHead(nn.Module):
    """Predict coordinate noise / score for all-atom generation (diffusion).

    Input coords ``[B, N, 3]`` + pair context → noise residual. Identity at
    zero residual scale.
    """

    def __init__(self, cfg: MolecularGeometryConfig | None = None) -> None:
        super().__init__()
        self.cfg = cfg or MolecularGeometryConfig()
        self.coord_in = nn.Linear(3, self.cfg.pair_dim, bias=False)
        self.time_in = nn.Linear(1, self.cfg.pair_dim, bias=False)
        self.pair_pool = nn.Linear(self.cfg.pair_dim, self.cfg.pair_dim, bias=False)
        self.noise_out = nn.Linear(self.cfg.pair_dim, 3, bias=False)
        self.residual_scale = nn.Parameter(torch.zeros(()))
        _role_seeded_xavier_uniform_(
            self.coord_in.weight,
            "diffusion.coord_in.weight",
        )
        _role_seeded_xavier_uniform_(
            self.time_in.weight,
            "diffusion.time_in.weight",
        )
        _role_seeded_xavier_uniform_(
            self.pair_pool.weight,
            "diffusion.pair_pool.weight",
        )
        nn.init.zeros_(self.noise_out.weight)

    def forward(
        self,
        coords_t: torch.Tensor,
        pair_t: torch.Tensor,
        *,
        diffusion_time_t: torch.Tensor | None = None,
        noise_t: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if coords_t.ndim != 3 or coords_t.shape[-1] != 3:
            raise ValueError("coords expect [batch, atoms, 3]")
        if pair_t.ndim != 4:
            raise ValueError("pair expect [batch, atoms, atoms, pair_dim]")
        active_time_t = (
            coords_t.new_zeros((coords_t.shape[0],))
            if diffusion_time_t is None
            else diffusion_time_t.to(device=coords_t.device, dtype=coords_t.dtype)
        )
        if active_time_t.shape != (coords_t.shape[0],):
            raise ValueError("diffusion time expects [batch]")
        atom_context_t = pair_t.mean(dim=2)
        fused_t = (
            self.coord_in(coords_t)
            + self.pair_pool(atom_context_t)
            + self.time_in(active_time_t[:, None, None])
        )
        # The learned projection and residual scale are both zero in inherited
        # pre-geometry checkpoints.  Multiplying those two zero surfaces would
        # create a permanently dead branch.  A deterministic, parameter-free
        # three-channel seed keeps the initial prediction exactly zero while
        # giving ``residual_scale`` a first-update gradient; after that update,
        # gradients also reach ``noise_out``.
        seed_noise_t = torch.tanh(fused_t[..., :3])
        pred_noise_t = torch.tanh(self.residual_scale) * (
            self.noise_out(fused_t) + seed_noise_t
        )
        if noise_t is None:
            target_t = coords_t.new_zeros(coords_t.shape)
        else:
            target_t = noise_t
        loss_t = F.mse_loss(pred_noise_t, target_t, reduction="none").mean(dim=-1)
        return pred_noise_t, loss_t


class VibrationalSpectrumHead(nn.Module):
    """Predict three normal modes per atom as wave-number/intensity pairs.

    QMugs stores ``3 * atom_count`` modes for each conformer.  The head keeps
    that native geometry as ``[batch, atoms, 3, 2]`` instead of flattening the
    spectrum into text.  Its zero residual preserves inherited behavior while
    a deterministic seed opens the first scale gradient.
    """

    def __init__(self, cfg: MolecularGeometryConfig | None = None) -> None:
        super().__init__()
        self.cfg = cfg or MolecularGeometryConfig()
        self.atom_context = nn.Linear(
            self.cfg.atom_dim,
            self.cfg.pair_dim,
            bias=False,
        )
        self.pair_context = nn.Linear(
            self.cfg.pair_dim,
            self.cfg.pair_dim,
            bias=False,
        )
        self.mode_out = nn.Linear(self.cfg.pair_dim, 6, bias=False)
        self.residual_scale = nn.Parameter(torch.zeros(()))
        _role_seeded_xavier_uniform_(
            self.atom_context.weight,
            "vibrational.atom_context.weight",
        )
        _role_seeded_xavier_uniform_(
            self.pair_context.weight,
            "vibrational.pair_context.weight",
        )
        nn.init.zeros_(self.mode_out.weight)

    def forward(
        self,
        atom_t: torch.Tensor,
        pair_t: torch.Tensor,
    ) -> torch.Tensor:
        if atom_t.ndim != 3:
            raise ValueError("vibrational atom context expects [batch, atoms, dim]")
        if pair_t.ndim != 4 or pair_t.shape[:2] != atom_t.shape[:2]:
            raise ValueError("vibrational pair context geometry differs")
        fused_t = self.atom_context(atom_t) + self.pair_context(pair_t.mean(dim=2))
        seed_modes_t = torch.tanh(fused_t[..., :6])
        modes_t = torch.tanh(self.residual_scale) * (
            self.mode_out(fused_t) + seed_modes_t
        )
        return cast(torch.Tensor, modes_t.reshape(*atom_t.shape[:2], 3, 2))


class AffinityHead(nn.Module):
    """Joint binder classification + potency regression (Boltz-2-style split)."""

    def __init__(self, cfg: MolecularGeometryConfig | None = None) -> None:
        super().__init__()
        self.cfg = cfg or MolecularGeometryConfig()
        self.pool = nn.Linear(self.cfg.hidden_size, self.cfg.affinity_hidden, bias=False)
        self.binder_logits = nn.Linear(self.cfg.affinity_hidden, 2, bias=True)
        self.potency = nn.Linear(self.cfg.affinity_hidden, 1, bias=True)
        _role_seeded_xavier_uniform_(self.pool.weight, "affinity.pool.weight")
        _role_seeded_xavier_uniform_(
            self.binder_logits.weight,
            "affinity.binder_logits.weight",
        )
        nn.init.zeros_(self.binder_logits.bias)
        nn.init.zeros_(self.potency.weight)
        nn.init.zeros_(self.potency.bias)

    def forward(self, hidden_t: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        if hidden_t.ndim != 3:
            raise ValueError("affinity head expects [batch, tokens, hidden]")
        pooled_t = self.pool(hidden_t.mean(dim=1))
        return self.binder_logits(pooled_t), self.potency(pooled_t).squeeze(-1)


class DevelopabilityHead(nn.Module):
    """Antibody/protein developability scores (aggregation, stability proxies)."""

    def __init__(self, cfg: MolecularGeometryConfig | None = None) -> None:
        super().__init__()
        self.cfg = cfg or MolecularGeometryConfig()
        self.proj = nn.Linear(self.cfg.hidden_size, 8, bias=True)
        _role_seeded_xavier_uniform_(
            self.proj.weight,
            "developability.proj.weight",
        )
        nn.init.zeros_(self.proj.bias)

    def forward(self, hidden_t: torch.Tensor) -> torch.Tensor:
        return torch.sigmoid(self.proj(hidden_t.mean(dim=1)))


class MolecularValidityHead(nn.Module):
    """Validity-preserving generation gate (Fragment-SELFIES / Molexar style)."""

    def __init__(self, cfg: MolecularGeometryConfig | None = None) -> None:
        super().__init__()
        self.cfg = cfg or MolecularGeometryConfig()
        self.proj = nn.Linear(self.cfg.hidden_size, 1, bias=True)
        nn.init.zeros_(self.proj.weight)
        nn.init.zeros_(self.proj.bias)

    def forward(self, hidden_t: torch.Tensor) -> torch.Tensor:
        return torch.sigmoid(self.proj(hidden_t)).squeeze(-1)


class PhysicalValidationHead(nn.Module):
    """Clash / stereochem / physical plausibility proxies from coords + hidden."""

    def __init__(self, cfg: MolecularGeometryConfig | None = None) -> None:
        super().__init__()
        self.cfg = cfg or MolecularGeometryConfig()
        self.hidden_proj = nn.Linear(self.cfg.hidden_size, 4, bias=True)
        _role_seeded_xavier_uniform_(
            self.hidden_proj.weight,
            "physical.hidden_proj.weight",
        )
        nn.init.zeros_(self.hidden_proj.bias)

    def forward(
        self,
        hidden_t: torch.Tensor,
        coords_t: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        scores_t = torch.sigmoid(self.hidden_proj(hidden_t.mean(dim=1)))
        clash_t = scores_t.new_zeros(scores_t.shape[0])
        if coords_t is not None:
            # Soft clash proxy: fraction of atom pairs under 1.0 Å (diagnostic).
            delta_t = coords_t.unsqueeze(2) - coords_t.unsqueeze(1)
            dist_t = delta_t.norm(dim=-1)
            eye = torch.eye(dist_t.shape[-1], device=dist_t.device, dtype=torch.bool)
            close_t = (dist_t < 1.0) & (~eye.unsqueeze(0))
            clash_t = close_t.float().mean(dim=(1, 2))
        return scores_t, clash_t


class MolecularScienceBank(nn.Module):
    """Composable bank of molecular science heads for the Resynthesis stack."""

    def __init__(self, cfg: MolecularGeometryConfig | None = None) -> None:
        super().__init__()
        self.cfg = cfg or MolecularGeometryConfig()
        self.atom_encoder = AtomFeatureEncoder(self.cfg)
        self.pair_rep = AllAtomPairRepresentation(self.cfg)
        self.diffusion = CoordinateDiffusionHead(self.cfg)
        self.vibrational = VibrationalSpectrumHead(self.cfg)
        self.affinity = AffinityHead(self.cfg)
        self.developability = DevelopabilityHead(self.cfg)
        self.validity = MolecularValidityHead(self.cfg)
        self.physical = PhysicalValidationHead(self.cfg)
        self.hidden_lift = nn.Linear(self.cfg.atom_dim, self.cfg.hidden_size, bias=False)
        self.blend = nn.Parameter(torch.zeros(()))
        _role_seeded_xavier_uniform_(
            self.hidden_lift.weight,
            "hidden_lift.weight",
        )

    def forward_hidden(self, hidden_t: torch.Tensor) -> torch.Tensor:
        """Apply the molecular residual without materializing auxiliary heads.

        Molecular coordinates and atom identities feed the auxiliary packet,
        but the residual returned to the reasoning graph is intentionally
        derived only from the model-produced hidden state. Earlier recursive
        attempts need that exact residual while only the final attempt's packet
        can contribute to the molecular loss.
        """

        text_atom_t = self.atom_encoder(hidden_t)
        residual_t = self.hidden_lift(text_atom_t)
        return hidden_t + torch.tanh(self.blend) * residual_t

    def forward(
        self,
        hidden_t: torch.Tensor,
        *,
        molecular_input: MolecularInputPacket | None = None,
        coords_t: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, MolecularSciencePacket]:
        """Return residual-blended hidden + tensor packet of molecular outputs."""

        if molecular_input is not None and coords_t is not None:
            raise ValueError("molecular input packet and legacy coordinates are exclusive")
        text_atom_t = self.atom_encoder(hidden_t)
        active_coords_t = coords_t
        atom_mask_t: torch.Tensor | None = None
        diffusion_time_t: torch.Tensor | None = None
        if molecular_input is None:
            atom_t = text_atom_t
        else:
            active_input = molecular_input.validated()
            atomic_numbers_t = active_input.atomic_numbers.to(device=hidden_t.device)
            active_coords_t = active_input.noisy_coordinates.to(
                device=hidden_t.device,
                dtype=hidden_t.dtype,
            )
            atom_mask_t = active_input.atom_mask.to(device=hidden_t.device)
            diffusion_time_t = active_input.diffusion_time.to(
                device=hidden_t.device,
                dtype=hidden_t.dtype,
            )
            feature_index_t = torch.arange(
                1,
                self.cfg.hidden_size + 1,
                device=hidden_t.device,
                dtype=hidden_t.dtype,
            )
            atomic_phase_t = atomic_numbers_t.to(dtype=hidden_t.dtype).unsqueeze(-1)
            atomic_hidden_t = (
                torch.sin(atomic_phase_t * feature_index_t * 0.017)
                + torch.cos(atomic_phase_t * feature_index_t * 0.031)
                + hidden_t.mean(dim=1, keepdim=True)
            )
            atomic_hidden_t = atomic_hidden_t * atom_mask_t.unsqueeze(-1).to(
                dtype=hidden_t.dtype
            )
            atom_t = self.atom_encoder(atomic_hidden_t)
        pair_t = self.pair_rep(atom_t)
        if atom_mask_t is not None:
            pair_mask_t = atom_mask_t.unsqueeze(2) & atom_mask_t.unsqueeze(1)
            pair_t = pair_t * pair_mask_t.unsqueeze(-1).to(dtype=pair_t.dtype)
        binder_logits_t, potency_t = self.affinity(hidden_t)
        develop_t = self.developability(hidden_t)
        validity_t = self.validity(hidden_t)
        physical_t, clash_t = self.physical(hidden_t, active_coords_t)
        if active_coords_t is None:
            active_coords_t = hidden_t.new_zeros(
                hidden_t.shape[0],
                hidden_t.shape[1],
                3,
            )
        noise_t, diffusion_loss_t = self.diffusion(
            active_coords_t,
            pair_t,
            diffusion_time_t=diffusion_time_t,
        )
        vibrational_spectrum_t = self.vibrational(atom_t, pair_t)
        if atom_mask_t is not None:
            active_mask_t = atom_mask_t.to(dtype=noise_t.dtype)
            noise_t = noise_t * active_mask_t.unsqueeze(-1)
            diffusion_loss_t = diffusion_loss_t * active_mask_t
            vibrational_spectrum_t = (
                vibrational_spectrum_t * active_mask_t[:, :, None, None]
            )
        residual_t = self.hidden_lift(text_atom_t)
        blended_t = hidden_t + torch.tanh(self.blend) * residual_t
        packet = MolecularSciencePacket(
            pair=pair_t,
            binder_logits=binder_logits_t,
            potency=potency_t,
            developability=develop_t,
            validity=validity_t,
            physical_scores=physical_t,
            clash=clash_t,
            coord_noise=noise_t,
            diffusion_loss=diffusion_loss_t,
            vibrational_spectrum=vibrational_spectrum_t,
        )
        return blended_t, packet