geno_lewm.surprise.score

score

Surprise scoring for action-conditioned DNA world-model predictions.

Aggregation module-attribute

Aggregation: TypeAlias = Literal['mean', 'max', 'median']

Supported aggregation modes for multi-step predictor outputs.

SurpriseResult dataclass

SurpriseResult(sigma_raw: float, sigma_calibrated: float, bucket_id: str, confidence: float, low_confidence: bool)

Calibrated surprise score for one edit.

to_dict

to_dict() -> dict[str, float | str | bool]

Return a JSON-native payload for CLI and JSONL outputs.

Source code in geno_lewm/surprise/score.py

def to_dict(self) -> dict[str, float | str | bool]:
    """Return a JSON-native payload for CLI and JSONL outputs."""
    return {
        "sigma_raw": self.sigma_raw,
        "sigma_calibrated": self.sigma_calibrated,
        "bucket_id": self.bucket_id,
        "confidence": self.confidence,
        "low_confidence": self.low_confidence,
    }

score_variant

score_variant(variant: EditSpec, encoder: object, action_encoder: object, predictor: object, calibration: CalibrationTable, *, reference_window: str, window_start_bp: int = 0, region: str | Sequence[str] | None = None, repeat: str | Sequence[str] | None = None, aggregation: str = 'mean', min_bucket_size: int = DEFAULT_MIN_BUCKET_SIZE) -> SurpriseResult

Score one edit against a caller-supplied reference window.

The scorer is intentionally model-object agnostic: callers can pass the concrete training-time modules or small deterministic fakes. FASTA-backed window extraction is available through :func:score_vcf; checkpoint loading is owned by higher runtime layers.

Source code in geno_lewm/surprise/score.py

def score_variant(
    variant: EditSpec,
    encoder: object,
    action_encoder: object,
    predictor: object,
    calibration: CalibrationTable,
    *,
    reference_window: str,
    window_start_bp: int = 0,
    region: str | Sequence[str] | None = None,
    repeat: str | Sequence[str] | None = None,
    aggregation: str = "mean",
    min_bucket_size: int = DEFAULT_MIN_BUCKET_SIZE,
) -> SurpriseResult:
    """Score one edit against a caller-supplied reference window.

    The scorer is intentionally model-object agnostic: callers can pass
    the concrete training-time modules or small deterministic fakes.
    FASTA-backed window extraction is available through :func:`score_vcf`;
    checkpoint loading is owned by higher runtime layers.
    """
    _require_calibration_table(calibration)
    min_size = _require_positive_int("min_bucket_size", min_bucket_size)
    bucket_id, sigma_raw = _raw_surprise(
        variant,
        encoder,
        action_encoder,
        predictor,
        reference_window=reference_window,
        window_start_bp=window_start_bp,
        region=region,
        repeat=repeat,
        aggregation=aggregation,
    )
    bucket = calibration.resolve(bucket_id, min_bucket_size=min_size)
    return SurpriseResult(
        sigma_raw=sigma_raw,
        sigma_calibrated=_cdf_percentile(bucket, sigma_raw),
        bucket_id=bucket.bucket_id,
        confidence=bucket.confidence,
        low_confidence=bucket.low_confidence,
    )

raw_surprise_example

raw_surprise_example(variant: EditSpec, encoder: object, action_encoder: object, predictor: object, *, reference_window: str, window_start_bp: int = 0, region: str | Sequence[str] | None = None, repeat: str | Sequence[str] | None = None, aggregation: str = 'mean') -> CalibrationExample

Score one reference variant into a :class:CalibrationExample.

Runs the model (no calibration table needed) and returns the variant's functional bucket_id paired with its raw surprise, ready to feed :func:geno_lewm.surprise.calibration.build_calibration_table.

Source code in geno_lewm/surprise/score.py

def raw_surprise_example(
    variant: EditSpec,
    encoder: object,
    action_encoder: object,
    predictor: object,
    *,
    reference_window: str,
    window_start_bp: int = 0,
    region: str | Sequence[str] | None = None,
    repeat: str | Sequence[str] | None = None,
    aggregation: str = "mean",
) -> CalibrationExample:
    """Score one reference variant into a :class:`CalibrationExample`.

    Runs the model (no calibration table needed) and returns the variant's
    functional ``bucket_id`` paired with its raw surprise, ready to feed
    :func:`geno_lewm.surprise.calibration.build_calibration_table`.
    """
    bucket_id, sigma_raw = _raw_surprise(
        variant,
        encoder,
        action_encoder,
        predictor,
        reference_window=reference_window,
        window_start_bp=window_start_bp,
        region=region,
        repeat=repeat,
        aggregation=aggregation,
    )
    return CalibrationExample(bucket_id=bucket_id, sigma_raw=sigma_raw)

score_vcf

score_vcf(vcf_path: str | Path, encoder: object, action_encoder: object, predictor: object, calibration: CalibrationTable, output_path: str | Path, *, reference_windows: Mapping[str, str] | None = None, reference_fasta: str | Path | None = None, window_bp: int = DEFAULT_WINDOW_BP, window_start_bp: int = 0, region: str | Sequence[str] | None = None, repeat: str | Sequence[str] | None = None, aggregation: str = 'mean', show_progress: bool = True, batch_size: int = 64, min_bucket_size: int = DEFAULT_MIN_BUCKET_SIZE) -> Path

Score VCF rows and write one JSON object per scored alternate.

Pass reference_fasta for local FASTA-backed window extraction. reference_windows remains useful for tests and already-extracted windows. Mapping keys are tried in this order: chrom:pos:ref:alt, chrom:pos, then chrom.

Source code in geno_lewm/surprise/score.py

def score_vcf(
    vcf_path: str | Path,
    encoder: object,
    action_encoder: object,
    predictor: object,
    calibration: CalibrationTable,
    output_path: str | Path,
    *,
    reference_windows: Mapping[str, str] | None = None,
    reference_fasta: str | Path | None = None,
    window_bp: int = DEFAULT_WINDOW_BP,
    window_start_bp: int = 0,
    region: str | Sequence[str] | None = None,
    repeat: str | Sequence[str] | None = None,
    aggregation: str = "mean",
    show_progress: bool = True,
    batch_size: int = 64,
    min_bucket_size: int = DEFAULT_MIN_BUCKET_SIZE,
) -> Path:
    """Score VCF rows and write one JSON object per scored alternate.

    Pass ``reference_fasta`` for local FASTA-backed window extraction.
    ``reference_windows`` remains useful for tests and already-extracted
    windows. Mapping keys are tried in this order:
    ``chrom:pos:ref:alt``, ``chrom:pos``, then ``chrom``.
    """
    if not isinstance(show_progress, bool):
        raise InputError(
            "show_progress must be a bool",
            details={"type": type(show_progress).__name__},
        )
    del show_progress

    output = Path(output_path)
    output.parent.mkdir(parents=True, exist_ok=True)
    with output.open("w", encoding="utf-8") as handle:
        for record in _iter_vcf_scores(
            vcf_path,
            encoder,
            action_encoder,
            predictor,
            calibration,
            reference_windows=reference_windows,
            reference_fasta=reference_fasta,
            window_bp=window_bp,
            window_start_bp=window_start_bp,
            region=region,
            repeat=repeat,
            aggregation=aggregation,
            batch_size=batch_size,
            min_bucket_size=min_bucket_size,
        ):
            variant = record.variant
            handle.write(
                json.dumps(
                    {
                        "schema_version": SCORE_JSONL_SCHEMA_VERSION,
                        "generated_by": SCORE_JSONL_GENERATED_BY,
                        "chrom": variant.chrom,
                        "pos": variant.pos,
                        "ref": variant.ref,
                        "alt": variant.alt,
                        **record.result.to_dict(),
                    },
                    sort_keys=True,
                )
                + "\n"
            )
    return output

build_calibration_examples_from_vcf

build_calibration_examples_from_vcf(vcf_path: str | Path, encoder: object, action_encoder: object, predictor: object, *, reference_windows: Mapping[str, str] | None = None, reference_fasta: str | Path | None = None, window_bp: int = DEFAULT_WINDOW_BP, window_start_bp: int = 0, region: str | Sequence[str] | None = None, repeat: str | Sequence[str] | None = None, aggregation: str = 'mean') -> list[CalibrationExample]

Score a background VCF into :class:CalibrationExample rows.

Runs the model over every scoreable VCF alternate (FASTA-backed window extraction) and records (bucket_id, sigma_raw) per variant. No calibration table is required, so the result can seed :func:geno_lewm.surprise.calibration.build_calibration_table.

Source code in geno_lewm/surprise/score.py

def build_calibration_examples_from_vcf(
    vcf_path: str | Path,
    encoder: object,
    action_encoder: object,
    predictor: object,
    *,
    reference_windows: Mapping[str, str] | None = None,
    reference_fasta: str | Path | None = None,
    window_bp: int = DEFAULT_WINDOW_BP,
    window_start_bp: int = 0,
    region: str | Sequence[str] | None = None,
    repeat: str | Sequence[str] | None = None,
    aggregation: str = "mean",
) -> list[CalibrationExample]:
    """Score a background VCF into :class:`CalibrationExample` rows.

    Runs the model over every scoreable VCF alternate (FASTA-backed window
    extraction) and records ``(bucket_id, sigma_raw)`` per variant. No
    calibration table is required, so the result can seed
    :func:`geno_lewm.surprise.calibration.build_calibration_table`.
    """
    if reference_windows is None and reference_fasta is None:
        raise InputError(
            "calibration requires reference_windows or reference_fasta",
            remediation="pass a local FASTA path or pre-extracted windows",
        )
    reference_sequences = (
        None if reference_fasta is None else _load_reference_fasta(reference_fasta)
    )
    examples: list[CalibrationExample] = []
    for variant in _iter_vcf_variants(vcf_path):
        ref_window = _reference_window_for_variant(
            variant,
            reference_windows,
            reference_sequences,
            window_start_bp=window_start_bp,
            window_bp=window_bp,
        )
        examples.append(
            raw_surprise_example(
                variant,
                encoder,
                action_encoder,
                predictor,
                reference_window=ref_window.sequence,
                window_start_bp=ref_window.start_bp,
                region=region,
                repeat=repeat,
                aggregation=aggregation,
            )
        )
    if not examples:
        raise VcfParseError(
            "VCF contains no scoreable variant rows", details={"path": str(vcf_path)}
        )
    return examples