Spaces:
Running on CPU Upgrade
Running on CPU Upgrade
File size: 10,921 Bytes
4db0438 5c3cfae 4db0438 5c3cfae 4db0438 5c3cfae 4db0438 5c3cfae 4db0438 5c3cfae 4db0438 5c3cfae 4db0438 5c3cfae 4db0438 5c3cfae 4db0438 5c3cfae 4db0438 5c3cfae 4db0438 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 | """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,
TOOL_REGISTRY,
)
from .latent_state import FullLatentState
from .noise import NoiseModel
from .output_generator import OutputGenerator
# Fallback costs per ActionType when the agent doesn't specify a known tool.
_BASE_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: ( 100, 0.5),
ActionType.REQUEST_SUBAGENT_REVIEW: ( 50, 0.25),
ActionType.SYNTHESIZE_CONCLUSION: ( 0, 0.5),
}
# Kept as public alias so existing imports (e.g. hackathon_environment) still work.
ACTION_COSTS = _BASE_ACTION_COSTS
def compute_action_cost(action: ExperimentAction) -> Tuple[float, float]:
"""Return (budget_cost, time_cost_days) for an action.
If the action specifies a ``method`` that exists in ``TOOL_REGISTRY``,
the tool's ``typical_cost_usd`` and ``typical_runtime_hours`` are used
(converted to days). Otherwise we fall back to the per-ActionType base
cost table.
"""
tool_spec = TOOL_REGISTRY.get(action.method or "")
if tool_spec is not None:
budget = tool_spec.typical_cost_usd
time_days = tool_spec.typical_runtime_hours / 24.0
return (budget, time_days)
return _BASE_ACTION_COSTS.get(action.action_type, (0.0, 0.0))
@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,
)
# ── internals ───────────────────────────────────────────────────────
def _apply_resource_cost(
self, s: FullLatentState, action: ExperimentAction
) -> None:
budget_cost, time_cost = compute_action_cost(action)
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.DESIGN_FOLLOWUP: "followup_designed",
ActionType.REQUEST_SUBAGENT_REVIEW: "subagent_review_requested",
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
p.n_cells_sequenced = self.noise.sample_count(
s.biology.n_true_cells * s.technical.capture_efficiency
)
if at in {ActionType.PERTURB_GENE, ActionType.PERTURB_COMPOUND}:
self._apply_perturbation_effects(s, action)
if at == ActionType.FILTER_DATA:
retain = self.noise.sample_qc_metric(0.85, 0.05, 0.5, 1.0)
base = p.n_cells_sequenced or s.biology.n_true_cells
p.n_cells_after_filter = max(100, int(base * retain))
s.last_retain_frac = 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)
s.last_n_clusters = p.n_clusters_found
def _apply_perturbation_effects(
self, s: FullLatentState, action: ExperimentAction
) -> None:
"""Fold perturbation-specific gene effects into true_de_genes so
downstream DE analysis reflects the perturbed biology."""
target = action.parameters.get("target", "")
effects = s.biology.perturbation_effects.get(target, {})
if not effects:
return
# Efficiency drawn from the same distribution as the output handler
# so latent state and observable output are coherent.
if action.action_type == ActionType.PERTURB_GENE:
efficiency = self.noise.sample_qc_metric(0.80, 0.12, 0.0, 1.0)
else:
efficiency = self.noise.sample_qc_metric(0.70, 0.15, 0.0, 1.0)
s.last_perturbation_efficiency = efficiency
for gene_map in s.biology.true_de_genes.values():
for gene, delta in effects.items():
gene_map[gene] = gene_map.get(gene, 0.0) + delta * efficiency
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]]
s.progress.n_de_genes_found = output.data.get("n_significant", 0)
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)
|