| """Data contract for trajectory-guided mutation and truncation design.""" |
| from __future__ import annotations |
|
|
| import json |
| from dataclasses import dataclass, field |
| from enum import Enum |
| from typing import Any, Mapping |
|
|
|
|
| class RegionObjective(str, Enum): |
| WEAKEN = "weaken" |
| PRESERVE = "preserve" |
| MONITOR = "monitor" |
|
|
|
|
| @dataclass(frozen=True, order=True) |
| class ResidueKey: |
| segid: str |
| resid: int |
| resname: str |
|
|
| def __post_init__(self) -> None: |
| if not self.resname.strip(): |
| raise ValueError("resname must not be empty") |
|
|
| @property |
| def token(self) -> str: |
| return f"{self.segid or '_'}:{self.resid}" |
|
|
| def to_dict(self) -> dict[str, Any]: |
| return {"segid": self.segid, "resid": self.resid, "resname": self.resname} |
|
|
| @classmethod |
| def from_dict(cls, value: Mapping[str, Any]) -> "ResidueKey": |
| return cls(str(value.get("segid", "")), int(value["resid"]), str(value["resname"])) |
|
|
|
|
| @dataclass(frozen=True) |
| class RegionSpec: |
| region_id: str |
| label: str |
| selection: str |
| objective: RegionObjective |
| cutoff_A: float = 4.0 |
|
|
| def __post_init__(self) -> None: |
| if not self.region_id.strip() or not self.label.strip() or not self.selection.strip(): |
| raise ValueError("region_id, label, and selection must not be empty") |
| if self.cutoff_A <= 0: |
| raise ValueError("cutoff_A must be positive") |
|
|
| def to_dict(self) -> dict[str, Any]: |
| return { |
| "region_id": self.region_id, |
| "label": self.label, |
| "selection": self.selection, |
| "objective": self.objective.value, |
| "cutoff_A": self.cutoff_A, |
| } |
|
|
| @classmethod |
| def from_dict(cls, value: Mapping[str, Any]) -> "RegionSpec": |
| return cls( |
| region_id=str(value["region_id"]), |
| label=str(value["label"]), |
| selection=str(value["selection"]), |
| objective=RegionObjective(value["objective"]), |
| cutoff_A=float(value.get("cutoff_A", 4.0)), |
| ) |
|
|
|
|
| @dataclass(frozen=True) |
| class TruncationSpec: |
| truncation_id: str |
| segid: str |
| start_resid: int |
| end_resid: int |
|
|
| def __post_init__(self) -> None: |
| if not self.truncation_id.strip(): |
| raise ValueError("truncation_id must not be empty") |
| if self.start_resid > self.end_resid: |
| raise ValueError("start_resid must be <= end_resid") |
|
|
| def contains(self, residue: ResidueKey) -> bool: |
| segid_matches = not self.segid or self.segid == "*" or residue.segid == self.segid |
| return segid_matches and self.start_resid <= residue.resid <= self.end_resid |
|
|
| def to_dict(self) -> dict[str, Any]: |
| return { |
| "truncation_id": self.truncation_id, |
| "segid": self.segid, |
| "start_resid": self.start_resid, |
| "end_resid": self.end_resid, |
| } |
|
|
| @classmethod |
| def from_dict(cls, value: Mapping[str, Any]) -> "TruncationSpec": |
| return cls( |
| truncation_id=str(value["truncation_id"]), |
| segid=str(value.get("segid", "")), |
| start_resid=int(value["start_resid"]), |
| end_resid=int(value["end_resid"]), |
| ) |
|
|
|
|
| @dataclass(frozen=True) |
| class EngineeringConfig: |
| protein_selection: str |
| regions: tuple[RegionSpec, ...] |
| protected_residues: tuple[str, ...] = field(default_factory=tuple) |
| mutation_budget: int = 5 |
| min_occupancy: float = 0.05 |
| preservation_weight: float = 1.0 |
| risk_weight: float = 1.0 |
| o2prime_bonus: float = 0.25 |
| max_preservation_burden: float = 0.40 |
| max_cumulative_risk: float = 1.50 |
| coverage_weight: float = 0.75 |
| truncation_coverage_gate: float = 0.50 |
| risk_by_residue: Mapping[str, float] = field(default_factory=dict) |
| truncations: tuple[TruncationSpec, ...] = field(default_factory=tuple) |
| schema_version: int = 1 |
|
|
| def __post_init__(self) -> None: |
| if not self.protein_selection.strip(): |
| raise ValueError("protein_selection must not be empty") |
| if not self.regions: |
| raise ValueError("at least one region is required") |
| region_ids = [item.region_id for item in self.regions] |
| if len(region_ids) != len(set(region_ids)): |
| raise ValueError("region_id values must be unique") |
| if not any(item.objective == RegionObjective.WEAKEN for item in self.regions): |
| raise ValueError("at least one weaken region is required") |
| if self.mutation_budget < 1: |
| raise ValueError("mutation_budget must be >= 1") |
| for label, value in ( |
| ("min_occupancy", self.min_occupancy), |
| ("max_preservation_burden", self.max_preservation_burden), |
| ("truncation_coverage_gate", self.truncation_coverage_gate), |
| ): |
| if not 0 <= value <= 1: |
| raise ValueError(f"{label} must be between 0 and 1") |
| if any(float(value) < 0 for value in self.risk_by_residue.values()): |
| raise ValueError("risk values must be non-negative") |
|
|
| def is_protected(self, residue: ResidueKey) -> bool: |
| return residue.token in self.protected_residues or f"*:{residue.resid}" in self.protected_residues |
|
|
| def risk_for(self, residue: ResidueKey) -> float: |
| return float( |
| self.risk_by_residue.get( |
| residue.token, self.risk_by_residue.get(f"*:{residue.resid}", 0.0) |
| ) |
| ) |
|
|
| def to_dict(self) -> dict[str, Any]: |
| return { |
| "schema_version": self.schema_version, |
| "protein_selection": self.protein_selection, |
| "regions": [item.to_dict() for item in self.regions], |
| "protected_residues": list(self.protected_residues), |
| "mutation_budget": self.mutation_budget, |
| "min_occupancy": self.min_occupancy, |
| "preservation_weight": self.preservation_weight, |
| "risk_weight": self.risk_weight, |
| "o2prime_bonus": self.o2prime_bonus, |
| "max_preservation_burden": self.max_preservation_burden, |
| "max_cumulative_risk": self.max_cumulative_risk, |
| "coverage_weight": self.coverage_weight, |
| "truncation_coverage_gate": self.truncation_coverage_gate, |
| "risk_by_residue": dict(self.risk_by_residue), |
| "truncations": [item.to_dict() for item in self.truncations], |
| } |
|
|
| @classmethod |
| def from_dict(cls, value: Mapping[str, Any]) -> "EngineeringConfig": |
| return cls( |
| protein_selection=str(value["protein_selection"]), |
| regions=tuple(RegionSpec.from_dict(item) for item in value["regions"]), |
| protected_residues=tuple(str(item) for item in value.get("protected_residues", [])), |
| mutation_budget=int(value.get("mutation_budget", 5)), |
| min_occupancy=float(value.get("min_occupancy", 0.05)), |
| preservation_weight=float(value.get("preservation_weight", 1.0)), |
| risk_weight=float(value.get("risk_weight", 1.0)), |
| o2prime_bonus=float(value.get("o2prime_bonus", 0.25)), |
| max_preservation_burden=float(value.get("max_preservation_burden", 0.40)), |
| max_cumulative_risk=float(value.get("max_cumulative_risk", 1.50)), |
| coverage_weight=float(value.get("coverage_weight", 0.75)), |
| truncation_coverage_gate=float(value.get("truncation_coverage_gate", 0.50)), |
| risk_by_residue={ |
| str(key): float(risk) for key, risk in value.get("risk_by_residue", {}).items() |
| }, |
| truncations=tuple( |
| TruncationSpec.from_dict(item) for item in value.get("truncations", []) |
| ), |
| schema_version=int(value.get("schema_version", 1)), |
| ) |
|
|
| def to_json(self) -> str: |
| return json.dumps(self.to_dict(), indent=2, allow_nan=False) |
|
|
|
|
| @dataclass(frozen=True) |
| class ContactObservation: |
| frame_index: int |
| protein: ResidueKey |
| region_id: str |
| nucleotide: ResidueKey |
| min_distance_A: float |
| o2prime_contact: bool = False |
|
|
|
|
| @dataclass(frozen=True) |
| class NucleotideContact: |
| nucleotide: ResidueKey |
| occupancy: float |
|
|
| def to_dict(self) -> dict[str, Any]: |
| return {"nucleotide": self.nucleotide.to_dict(), "occupancy": self.occupancy} |
|
|
| @classmethod |
| def from_dict(cls, value: Mapping[str, Any]) -> "NucleotideContact": |
| return cls(ResidueKey.from_dict(value["nucleotide"]), float(value["occupancy"])) |
|
|
|
|
| @dataclass(frozen=True) |
| class RegionFingerprint: |
| protein: ResidueKey |
| region_id: str |
| contact_occupancy: float |
| mean_contact_distance_A: float |
| p10_contact_distance_A: float |
| contact_events: int |
| longest_run_frames: int |
| o2prime_occupancy: float |
| nucleotide_contacts: tuple[NucleotideContact, ...] = field(default_factory=tuple) |
|
|
| @property |
| def nucleotide_coverage(self) -> int: |
| return len(self.nucleotide_contacts) |
|
|
| def to_dict(self) -> dict[str, Any]: |
| return { |
| "protein": self.protein.to_dict(), |
| "region_id": self.region_id, |
| "contact_occupancy": self.contact_occupancy, |
| "mean_contact_distance_A": self.mean_contact_distance_A, |
| "p10_contact_distance_A": self.p10_contact_distance_A, |
| "contact_events": self.contact_events, |
| "longest_run_frames": self.longest_run_frames, |
| "o2prime_occupancy": self.o2prime_occupancy, |
| "nucleotide_contacts": [item.to_dict() for item in self.nucleotide_contacts], |
| } |
|
|
| @classmethod |
| def from_dict(cls, value: Mapping[str, Any]) -> "RegionFingerprint": |
| return cls( |
| protein=ResidueKey.from_dict(value["protein"]), |
| region_id=str(value["region_id"]), |
| contact_occupancy=float(value["contact_occupancy"]), |
| mean_contact_distance_A=float(value["mean_contact_distance_A"]), |
| p10_contact_distance_A=float(value["p10_contact_distance_A"]), |
| contact_events=int(value["contact_events"]), |
| longest_run_frames=int(value["longest_run_frames"]), |
| o2prime_occupancy=float(value.get("o2prime_occupancy", 0.0)), |
| nucleotide_contacts=tuple( |
| NucleotideContact.from_dict(item) for item in value.get("nucleotide_contacts", []) |
| ), |
| ) |
|
|
|
|
| @dataclass(frozen=True) |
| class MutationCandidate: |
| protein: ResidueKey |
| suggested_mutation: str |
| target_engagement: float |
| preservation_burden: float |
| structural_risk: float |
| o2prime_signal: float |
| design_score: float |
| target_pairs: tuple[str, ...] |
| eligible: bool |
| exclusion_reason: str = "" |
| rationale: tuple[str, ...] = field(default_factory=tuple) |
|
|
| def to_dict(self) -> dict[str, Any]: |
| return { |
| "protein": self.protein.to_dict(), |
| "suggested_mutation": self.suggested_mutation, |
| "target_engagement": self.target_engagement, |
| "preservation_burden": self.preservation_burden, |
| "structural_risk": self.structural_risk, |
| "o2prime_signal": self.o2prime_signal, |
| "design_score": self.design_score, |
| "target_pairs": list(self.target_pairs), |
| "eligible": self.eligible, |
| "exclusion_reason": self.exclusion_reason, |
| "rationale": list(self.rationale), |
| } |
|
|
| @classmethod |
| def from_dict(cls, value: Mapping[str, Any]) -> "MutationCandidate": |
| return cls( |
| protein=ResidueKey.from_dict(value["protein"]), |
| suggested_mutation=str(value["suggested_mutation"]), |
| target_engagement=float(value["target_engagement"]), |
| preservation_burden=float(value["preservation_burden"]), |
| structural_risk=float(value["structural_risk"]), |
| o2prime_signal=float(value.get("o2prime_signal", 0.0)), |
| design_score=float(value["design_score"]), |
| target_pairs=tuple(str(item) for item in value.get("target_pairs", [])), |
| eligible=bool(value["eligible"]), |
| exclusion_reason=str(value.get("exclusion_reason", "")), |
| rationale=tuple(str(item) for item in value.get("rationale", [])), |
| ) |
|
|
|
|
| @dataclass(frozen=True) |
| class MutationSet: |
| selected: tuple[MutationCandidate, ...] |
| target_pair_coverage: float |
| preservation_burden: float |
| cumulative_risk: float |
| warnings: tuple[str, ...] = field(default_factory=tuple) |
|
|
| def to_dict(self) -> dict[str, Any]: |
| return { |
| "selected": [item.to_dict() for item in self.selected], |
| "target_pair_coverage": self.target_pair_coverage, |
| "preservation_burden": self.preservation_burden, |
| "cumulative_risk": self.cumulative_risk, |
| "warnings": list(self.warnings), |
| } |
|
|
| @classmethod |
| def from_dict(cls, value: Mapping[str, Any]) -> "MutationSet": |
| return cls( |
| selected=tuple(MutationCandidate.from_dict(item) for item in value["selected"]), |
| target_pair_coverage=float(value["target_pair_coverage"]), |
| preservation_burden=float(value["preservation_burden"]), |
| cumulative_risk=float(value["cumulative_risk"]), |
| warnings=tuple(str(item) for item in value.get("warnings", [])), |
| ) |
|
|
|
|
| @dataclass(frozen=True) |
| class TruncationAudit: |
| truncation: TruncationSpec |
| target_coverage_removed: float |
| preserve_coverage_removed: float |
| unwanted_interaction_remaining: float |
| distributed_interface_warning: bool |
| remaining_candidate_residues: tuple[str, ...] |
| flexibility_ratio: float | None |
| interpretation: str |
|
|
| def to_dict(self) -> dict[str, Any]: |
| return { |
| "truncation": self.truncation.to_dict(), |
| "target_coverage_removed": self.target_coverage_removed, |
| "preserve_coverage_removed": self.preserve_coverage_removed, |
| "unwanted_interaction_remaining": self.unwanted_interaction_remaining, |
| "distributed_interface_warning": self.distributed_interface_warning, |
| "remaining_candidate_residues": list(self.remaining_candidate_residues), |
| "flexibility_ratio": self.flexibility_ratio, |
| "interpretation": self.interpretation, |
| } |
|
|
| @classmethod |
| def from_dict(cls, value: Mapping[str, Any]) -> "TruncationAudit": |
| ratio = value.get("flexibility_ratio") |
| return cls( |
| truncation=TruncationSpec.from_dict(value["truncation"]), |
| target_coverage_removed=float(value["target_coverage_removed"]), |
| preserve_coverage_removed=float(value["preserve_coverage_removed"]), |
| unwanted_interaction_remaining=float(value["unwanted_interaction_remaining"]), |
| distributed_interface_warning=bool(value["distributed_interface_warning"]), |
| remaining_candidate_residues=tuple( |
| str(item) for item in value.get("remaining_candidate_residues", []) |
| ), |
| flexibility_ratio=None if ratio is None else float(ratio), |
| interpretation=str(value["interpretation"]), |
| ) |
|
|
|
|
| @dataclass(frozen=True) |
| class RegionComparison: |
| region_id: str |
| objective: RegionObjective |
| baseline_mass: float |
| variant_mass: float |
| retention_pct: float | None |
| objective_change: float | None |
|
|
| def to_dict(self) -> dict[str, Any]: |
| return { |
| "region_id": self.region_id, |
| "objective": self.objective.value, |
| "baseline_mass": self.baseline_mass, |
| "variant_mass": self.variant_mass, |
| "retention_pct": self.retention_pct, |
| "objective_change": self.objective_change, |
| } |
|
|
|
|
| @dataclass(frozen=True) |
| class VariantComparison: |
| baseline_label: str |
| variant_label: str |
| regions: tuple[RegionComparison, ...] |
|
|
| def to_dict(self) -> dict[str, Any]: |
| return { |
| "baseline_label": self.baseline_label, |
| "variant_label": self.variant_label, |
| "regions": [item.to_dict() for item in self.regions], |
| } |
|
|
|
|
| @dataclass(frozen=True) |
| class DesignReport: |
| config: EngineeringConfig |
| n_frames: int |
| frame_interval_ps: float | None |
| source_label: str |
| fingerprints: tuple[RegionFingerprint, ...] |
| candidates: tuple[MutationCandidate, ...] |
| mutation_set: MutationSet |
| truncation_audits: tuple[TruncationAudit, ...] |
| provenance: Mapping[str, Any] = field(default_factory=dict) |
|
|
| def to_dict(self) -> dict[str, Any]: |
| return { |
| "schema_version": 1, |
| "config": self.config.to_dict(), |
| "n_frames": self.n_frames, |
| "frame_interval_ps": self.frame_interval_ps, |
| "source_label": self.source_label, |
| "fingerprints": [item.to_dict() for item in self.fingerprints], |
| "candidates": [item.to_dict() for item in self.candidates], |
| "mutation_set": self.mutation_set.to_dict(), |
| "truncation_audits": [item.to_dict() for item in self.truncation_audits], |
| "provenance": dict(self.provenance), |
| } |
|
|
| def to_json(self) -> str: |
| return json.dumps(self.to_dict(), indent=2, allow_nan=False) |
|
|
|
|
