bioai/simulation/wet_lab.py

"""bioai.simulation.wet_lab -- Virtual wet-lab validation simulator.

Simulates the cellular knockdown pipeline for dsRNA biopesticide candidates
using a 6-stage Monte Carlo model. This is NOT a replacement for real wet-lab
testing — it's a computational screening tool that estimates which candidates
are most likely to succeed before spending real lab time and money.

The 6 stages (each modeled with biological literature-informed kinetics):

  1. DELIVERY  — fraction of applied dsRNA that reaches target pest cells.
     Modeled as a function of the PINN-predicted environmental half-life
     (longer half-life = more dsRNA survives to reach the pest).

  2. UPTAKE    — cellular uptake efficiency. Modeled as a function of
     dsRNA length (21-nt siRNAs uptake more efficiently than 200-nt
     precursors) and GC content (moderate GC = better uptake).

  3. DICER     — probability of being correctly processed by Dicer into
     the active 21-nt siRNA. Modeled using sequence features (no internal
     repeats, moderate GC = better Dicer processing).

  4. RISC      — guide-strand loading efficiency. Modeled using
     thermodynamic asymmetry (Reynolds rule: lower 5' antisense binding
     energy = better guide loading).

  5. CLEAVAGE  — target mRNA cleavage rate. Uses the CNN-predicted
     efficacy score as the base rate, with noise to simulate biological
     variability.

  6. PHENOTYPE — phenotypic response (mortality / growth reduction).
     Modeled as a Hill dose-response curve: knockdown -> phenotype.

The final knockdown percentage is the product of all 6 stages, with noise
added at each step. We run N=1000 Monte Carlo trials per candidate and
report the mean, 95% confidence interval, and probability of achieving
>70% knockdown (the threshold for a "functional" siRNA per Reynolds 2004).

Scientific basis:
  - Reynolds et al. 2004 (Nature Biotechnology) — siRNA efficacy rules
  - Schwarz et al. 2003 (Cell) — thermodynamic asymmetry determines guide strand
  - Fire et al. 1998 (Nature) — RNAi discovery, dose-response kinetics
  - Bhatt et al. 2004 (NAR) — Dicer processing preferences
  - Parrish et al. 2000 (Molecular Cell) — dsRNA uptake in C. elegans

Limitations:
  - The noise model is Gaussian; real biological variability is heavier-tailed.
  - The dose-response curve is a simple Hill function; real pest response
    varies by species, life stage, and environmental conditions.
  - Off-target effects are not modeled here (handled separately by the
    k-mer index in the ranker).
  - This simulation CANNOT replace real wet-lab validation. It's a
    screening tool to prioritize candidates for lab testing.
"""

from __future__ import annotations

import math
from dataclasses import dataclass, field
from typing import Dict, List

import numpy as np


# ─────────────────────────────────────────────────────────────────────────────
# Constants (literature-informed defaults)
# ─────────────────────────────────────────────────────────────────────────────
N_TRIALS = 1000
KNOCKDOWN_THRESHOLD = 0.70  # Reynolds 2004 "functional" threshold


@dataclass
class SimulationResult:
    """Result of a single candidate's wet-lab simulation."""
    sirna_seq: str
    mean_knockdown: float           # mean fraction knocked down (0-1)
    ci_low: float                   # 95% CI lower bound
    ci_high: float                  # 95% CI upper bound
    prob_above_70: float            # P(knockdown > 70%)
    # Per-stage mean efficiencies (0-1)
    delivery_efficiency: float
    uptake_efficiency: float
    dicer_efficiency: float
    risc_loading: float
    cleavage_rate: float
    phenotypic_response: float
    n_trials: int = N_TRIALS


# ─────────────────────────────────────────────────────────────────────────────
# WetLabSimulator
# ─────────────────────────────────────────────────────────────────────────────
class WetLabSimulator:
    """Monte Carlo simulator for the dsRNA cellular knockdown pipeline.

    Parameters
    ----------
    n_trials:
        Number of Monte Carlo trials per candidate (default 1000).
    rng_seed:
        Random seed for reproducibility (default 42).
    """

    def __init__(self, n_trials: int = N_TRIALS, rng_seed: int = 42):
        self.n_trials = n_trials
        self.rng = np.random.default_rng(rng_seed)

    def simulate_candidate(
        self,
        sirna_seq: str,
        cnn_efficacy: float,
        half_life_hours: float,
        reynolds_score: int = 8,
    ) -> SimulationResult:
        """Run the 6-stage Monte Carlo simulation for one candidate.

        Models an in vitro cell-culture screening assay (the standard first
        step in dsRNA biopesticide validation). In this context:
        - Delivery is via lipofection (80-95% efficient, not field spray)
        - The cellular machinery (Dicer, RISC) operates at high efficiency
        - The CNN efficacy is the primary driver of knockdown
        - Other stages act as moderate modifiers (0.7-1.0x)

        Parameters
        ----------
        sirna_seq:
            The 21-nt siRNA sequence.
        cnn_efficacy:
            CNN-predicted efficacy (0-1) — primary driver of knockdown.
        half_life_hours:
            PINN-predicted environmental half-life (modifies delivery stage).
        reynolds_score:
            Reynolds 2004 rule score (0-8). Higher = better siRNA design.
        """
        gc = self._gc_content(sirna_seq)
        has_repeats = self._has_internal_repeats(sirna_seq)

        knockdowns = np.zeros(self.n_trials)
        delivery_arr = np.zeros(self.n_trials)
        uptake_arr = np.zeros(self.n_trials)
        dicer_arr = np.zeros(self.n_trials)
        risc_arr = np.zeros(self.n_trials)
        cleavage_arr = np.zeros(self.n_trials)
        phenotype_arr = np.zeros(self.n_trials)

        for i in range(self.n_trials):
            # Stage 1: Delivery (lipofection in vitro) — 80-95% efficient
            # Longer half-life slightly improves delivery stability
            hl_bonus = (half_life_hours - 24) / (168 - 24) * 0.08  # 0-8% bonus
            delivery = 0.85 + hl_bonus + self.rng.normal(0, 0.04)
            delivery = np.clip(delivery, 0.70, 0.98)
            delivery_arr[i] = delivery

            # Stage 2: Uptake — 75-95%, GC-dependent
            # Moderate GC (0.3-0.6) is optimal for uptake
            gc_factor = 1.0 - abs(gc - 0.45) * 0.8
            gc_factor = np.clip(gc_factor, 0.7, 1.0)
            uptake = (0.80 * gc_factor) + self.rng.normal(0, 0.05)
            uptake = np.clip(uptake, 0.60, 0.95)
            uptake_arr[i] = uptake

            # Stage 3: Dicer processing — 80-95%, Reynolds-dependent
            dicer_base = 0.75 + (reynolds_score / 8.0) * 0.20
            if has_repeats:
                dicer_base *= 0.85
            dicer = dicer_base + self.rng.normal(0, 0.04)
            dicer = np.clip(dicer, 0.65, 0.97)
            dicer_arr[i] = dicer

            # Stage 4: RISC loading — 70-95%, thermodynamic asymmetry dependent
            risc_base = 0.65 + (reynolds_score / 8.0) * 0.30
            risc = risc_base + self.rng.normal(0, 0.06)
            risc = np.clip(risc, 0.55, 0.95)
            risc_arr[i] = risc

            # Stage 5: Target cleavage — driven by CNN efficacy
            # The CNN predicts the intrinsic cleavage efficiency; add biological noise
            cleavage = cnn_efficacy + self.rng.normal(0, 0.08)
            cleavage = np.clip(cleavage, 0.20, 0.98)
            cleavage_arr[i] = cleavage

            # Stage 6: Phenotypic response — mRNA knockdown -> phenotype
            # In cell culture, phenotype tracks mRNA knockdown closely (R^2 ~ 0.85)
            # Use a soft saturating function rather than full Hill
            cumulative = delivery * uptake * dicer * risc * cleavage
            # Soft saturation: phenotype = cumulative^0.85 (sublinear, realistic)
            phenotype = cumulative ** 0.85 + self.rng.normal(0, 0.04)
            phenotype = np.clip(phenotype, 0.0, 0.95)
            phenotype_arr[i] = phenotype

            knockdowns[i] = phenotype

        mean_kd = float(np.mean(knockdowns))
        ci_low = float(np.percentile(knockdowns, 2.5))
        ci_high = float(np.percentile(knockdowns, 97.5))
        prob_above_70 = float(np.mean(knockdowns > KNOCKDOWN_THRESHOLD))

        return SimulationResult(
            sirna_seq=sirna_seq,
            mean_knockdown=mean_kd,
            ci_low=ci_low,
            ci_high=ci_high,
            prob_above_70=prob_above_70,
            delivery_efficiency=float(np.mean(delivery_arr)),
            uptake_efficiency=float(np.mean(uptake_arr)),
            dicer_efficiency=float(np.mean(dicer_arr)),
            risc_loading=float(np.mean(risc_arr)),
            cleavage_rate=float(np.mean(cleavage_arr)),
            phenotypic_response=float(np.mean(phenotype_arr)),
        )

    def simulate_batch(
        self,
        candidates: List[Dict],
        pest_species: str = "unknown",
    ) -> List[SimulationResult]:
        """Run the simulation for a batch of candidates.

        Each candidate dict should contain:
            - sirna_seq: the 21-nt sequence
            - efficacy: CNN-predicted efficacy (0-1)
            - half_life_hours: PINN-predicted half-life
            - reynolds_score (optional): 0-8, defaults to 8 if missing
        """
        results = []
        for cand in candidates:
            seq = cand.get("sirna_seq", "")
            eff = cand.get("efficacy", 0.5)
            hl = cand.get("half_life_hours", 48.0)
            rs = cand.get("reynolds_score", 8)
            result = self.simulate_candidate(seq, eff, hl, rs)
            results.append(result)
        return results

    # ─── Sequence helpers ────────────────────────────────────────────────
    @staticmethod
    def _gc_content(seq: str) -> float:
        seq = seq.upper().replace("U", "T")
        if not seq:
            return 0.0
        return (seq.count("G") + seq.count("C")) / len(seq)

    @staticmethod
    def _has_internal_repeats(seq: str) -> bool:
        """Check for internal repeats (>4 consecutive identical bases)."""
        seq = seq.upper()
        for i in range(len(seq) - 4):
            if seq[i] == seq[i + 1] == seq[i + 2] == seq[i + 3] == seq[i + 4]:
                return True
        return False


# ─────────────────────────────────────────────────────────────────────────────
# Result serialization
# ─────────────────────────────────────────────────────────────────────────────
def result_to_dict(result: SimulationResult) -> Dict:
    """Convert a SimulationResult to a JSON-serializable dict."""
    return {
        "sirna_seq": result.sirna_seq,
        "mean_knockdown": round(result.mean_knockdown, 4),
        "ci_low": round(result.ci_low, 4),
        "ci_high": round(result.ci_high, 4),
        "prob_above_70": round(result.prob_above_70, 4),
        "delivery_efficiency": round(result.delivery_efficiency, 4),
        "uptake_efficiency": round(result.uptake_efficiency, 4),
        "dicer_efficiency": round(result.dicer_efficiency, 4),
        "risc_loading": round(result.risc_loading, 4),
        "cleavage_rate": round(result.cleavage_rate, 4),
        "phenotypic_response": round(result.phenotypic_response, 4),
        "n_trials": result.n_trials,
    }