| """Transition dynamics engine — the heart of the biological simulator.
|
|
|
| Orchestrates latent-state updates, output generation, resource accounting,
|
| and constraint propagation for every agent action.
|
| """
|
|
|
| from __future__ import annotations
|
|
|
| from copy import deepcopy
|
| from dataclasses import dataclass, field
|
| from typing import Any, Dict, List, Optional, Tuple
|
|
|
| from models import (
|
| ActionType,
|
| ExperimentAction,
|
| IntermediateOutput,
|
| OutputType,
|
| )
|
|
|
| from .latent_state import FullLatentState
|
| from .noise import NoiseModel
|
| from .output_generator import OutputGenerator
|
|
|
|
|
| ACTION_COSTS: Dict[ActionType, Tuple[float, float]] = {
|
| ActionType.COLLECT_SAMPLE: (5_000, 7.0),
|
| ActionType.SELECT_COHORT: ( 500, 1.0),
|
| ActionType.PREPARE_LIBRARY: (8_000, 3.0),
|
| ActionType.CULTURE_CELLS: (3_000, 14.0),
|
| ActionType.PERTURB_GENE: (2_000, 3.0),
|
| ActionType.PERTURB_COMPOUND: (1_000, 2.0),
|
| ActionType.SEQUENCE_CELLS: (15_000, 5.0),
|
| ActionType.RUN_QC: ( 100, 0.5),
|
| ActionType.FILTER_DATA: ( 50, 0.25),
|
| ActionType.NORMALIZE_DATA: ( 50, 0.25),
|
| ActionType.INTEGRATE_BATCHES: ( 100, 0.5),
|
| ActionType.CLUSTER_CELLS: ( 100, 0.5),
|
| ActionType.DIFFERENTIAL_EXPRESSION: ( 100, 0.5),
|
| ActionType.TRAJECTORY_ANALYSIS: ( 200, 1.0),
|
| ActionType.PATHWAY_ENRICHMENT: ( 100, 0.5),
|
| ActionType.REGULATORY_NETWORK_INFERENCE: ( 300, 1.0),
|
| ActionType.MARKER_SELECTION: ( 100, 0.5),
|
| ActionType.VALIDATE_MARKER: (5_000, 14.0),
|
| ActionType.DESIGN_FOLLOWUP: ( 0, 0.5),
|
| ActionType.REQUEST_SUBAGENT_REVIEW: ( 0, 0.25),
|
| ActionType.SYNTHESIZE_CONCLUSION: ( 0, 0.5),
|
| }
|
|
|
|
|
| @dataclass
|
| class TransitionResult:
|
| """Bundle returned by the transition engine after one step."""
|
|
|
| next_state: FullLatentState
|
| output: IntermediateOutput
|
| reward_components: Dict[str, float] = field(default_factory=dict)
|
| hard_violations: List[str] = field(default_factory=list)
|
| soft_violations: List[str] = field(default_factory=list)
|
| done: bool = False
|
|
|
|
|
| class TransitionEngine:
|
| """Applies one action to the latent state, producing the next state
|
| and a simulated intermediate output.
|
|
|
| The engine delegates output generation to ``OutputGenerator`` and
|
| constraint checking to external rule engines (injected at call time).
|
| """
|
|
|
| def __init__(self, noise: NoiseModel):
|
| self.noise = noise
|
| self.output_gen = OutputGenerator(noise)
|
|
|
| def step(
|
| self,
|
| state: FullLatentState,
|
| action: ExperimentAction,
|
| *,
|
| hard_violations: Optional[List[str]] = None,
|
| soft_violations: Optional[List[str]] = None,
|
| ) -> TransitionResult:
|
| s = deepcopy(state)
|
| s.step_count += 1
|
| step_idx = s.step_count
|
|
|
| hard_v = hard_violations or []
|
| soft_v = soft_violations or []
|
|
|
| if hard_v:
|
| output = IntermediateOutput(
|
| output_type=OutputType.FAILURE_REPORT,
|
| step_index=step_idx,
|
| success=False,
|
| summary=f"Action blocked: {'; '.join(hard_v)}",
|
| )
|
| return TransitionResult(
|
| next_state=s,
|
| output=output,
|
| hard_violations=hard_v,
|
| soft_violations=soft_v,
|
| )
|
|
|
| self._apply_resource_cost(s, action)
|
|
|
| if s.resources.budget_exhausted or s.resources.time_exhausted:
|
| output = IntermediateOutput(
|
| output_type=OutputType.FAILURE_REPORT,
|
| step_index=step_idx,
|
| success=False,
|
| summary="Resources exhausted",
|
| )
|
| return TransitionResult(
|
| next_state=s, output=output, done=True,
|
| hard_violations=["resources_exhausted"],
|
| )
|
|
|
| self._update_progress(s, action)
|
|
|
| output = self.output_gen.generate(action, s, step_idx)
|
|
|
| if soft_v:
|
| output.quality_score *= 0.5
|
| output.warnings.extend(soft_v)
|
|
|
| self._propagate_artifacts(s, action, output)
|
|
|
| done = action.action_type == ActionType.SYNTHESIZE_CONCLUSION
|
|
|
| return TransitionResult(
|
| next_state=s,
|
| output=output,
|
| soft_violations=soft_v,
|
| done=done,
|
| )
|
|
|
|
|
|
|
| def _apply_resource_cost(
|
| self, s: FullLatentState, action: ExperimentAction
|
| ) -> None:
|
| budget_cost, time_cost = ACTION_COSTS.get(
|
| action.action_type, (0.0, 0.0)
|
| )
|
| s.resources.budget_used += budget_cost
|
| s.resources.time_used_days += time_cost
|
| if action.action_type in {
|
| ActionType.RUN_QC, ActionType.FILTER_DATA,
|
| ActionType.NORMALIZE_DATA, ActionType.INTEGRATE_BATCHES,
|
| ActionType.CLUSTER_CELLS, ActionType.DIFFERENTIAL_EXPRESSION,
|
| ActionType.TRAJECTORY_ANALYSIS, ActionType.PATHWAY_ENRICHMENT,
|
| ActionType.REGULATORY_NETWORK_INFERENCE, ActionType.MARKER_SELECTION,
|
| }:
|
| s.resources.compute_hours_used += time_cost * 8
|
|
|
| def _update_progress(
|
| self, s: FullLatentState, action: ExperimentAction
|
| ) -> None:
|
| at = action.action_type
|
| p = s.progress
|
| _MAP = {
|
| ActionType.COLLECT_SAMPLE: "samples_collected",
|
| ActionType.SELECT_COHORT: "cohort_selected",
|
| ActionType.PREPARE_LIBRARY: "library_prepared",
|
| ActionType.CULTURE_CELLS: "cells_cultured",
|
| ActionType.PERTURB_GENE: "perturbation_applied",
|
| ActionType.PERTURB_COMPOUND: "perturbation_applied",
|
| ActionType.SEQUENCE_CELLS: "cells_sequenced",
|
| ActionType.RUN_QC: "qc_performed",
|
| ActionType.FILTER_DATA: "data_filtered",
|
| ActionType.NORMALIZE_DATA: "data_normalized",
|
| ActionType.INTEGRATE_BATCHES: "batches_integrated",
|
| ActionType.CLUSTER_CELLS: "cells_clustered",
|
| ActionType.DIFFERENTIAL_EXPRESSION: "de_performed",
|
| ActionType.TRAJECTORY_ANALYSIS: "trajectories_inferred",
|
| ActionType.PATHWAY_ENRICHMENT: "pathways_analyzed",
|
| ActionType.REGULATORY_NETWORK_INFERENCE: "networks_inferred",
|
| ActionType.MARKER_SELECTION: "markers_discovered",
|
| ActionType.VALIDATE_MARKER: "markers_validated",
|
| ActionType.SYNTHESIZE_CONCLUSION: "conclusion_reached",
|
| }
|
| flag = _MAP.get(at)
|
| if flag:
|
| setattr(p, flag, True)
|
|
|
| if at == ActionType.COLLECT_SAMPLE:
|
| n = action.parameters.get("n_samples", 6)
|
| s.resources.samples_available += n
|
|
|
| if at == ActionType.SEQUENCE_CELLS:
|
| s.resources.sequencing_lanes_used += 1
|
|
|
| if at == ActionType.FILTER_DATA:
|
| retain = self.noise.sample_qc_metric(0.85, 0.05, 0.5, 1.0)
|
| p.n_cells_after_filter = max(
|
| 100, int(s.biology.n_true_cells * retain)
|
| )
|
|
|
| if at == ActionType.CLUSTER_CELLS:
|
| n_true = len(s.biology.cell_populations) or 5
|
| p.n_clusters_found = self.noise.sample_cluster_count(n_true, 0.8)
|
|
|
| def _propagate_artifacts(
|
| self,
|
| s: FullLatentState,
|
| action: ExperimentAction,
|
| output: IntermediateOutput,
|
| ) -> None:
|
| if action.action_type == ActionType.DIFFERENTIAL_EXPRESSION:
|
| top = output.data.get("top_genes", [])
|
| s.discovered_de_genes = [g["gene"] for g in top[:20]]
|
|
|
| if action.action_type == ActionType.CLUSTER_CELLS:
|
| s.discovered_clusters = output.data.get("cluster_names", [])
|
|
|
| if action.action_type == ActionType.MARKER_SELECTION:
|
| s.progress.n_markers_found = output.data.get("n_candidates", 0)
|
|
|
