engine_v2/types.py

"""Typed grammar for engine_v2.

The grammar covers the FULL DSL:

- ``Matrix``        — a (sub)matrix of opaque-ID columns.
- ``Vector``        — one score per patient.
- ``Scalar``        — a single number (e.g. an Associate / Effect output).
- ``Model``         — a fitted classifier (Fit) we can Apply.

``FeatureSet`` is a leaf payload (a list of opaque IDs carried by Select);
``Agg`` / ``Op`` / ``Kind`` are enum-like string payloads. ``Predicate``
is carried by ``Split`` as a small dataclass.
"""

from __future__ import annotations

from enum import Enum


class TType(str, Enum):
    """Return type of an AST node — what a parent slot demands."""

    MATRIX = "Matrix"
    VECTOR = "Vector"
    SCALAR = "Scalar"
    MODEL = "Model"


# Reduce aggregators. Must mirror dsl.operators._REDUCE_AGGS.
AGGS = ("mean", "median", "max", "min", "var")

# Pairwise vector combinators. ``protected_div`` guards against zero.
OPS = ("add", "sub", "mul", "protected_div", "mean")

# Association kinds for the ``Associate`` and ``Effect`` operators.
ASSOC_KINDS = ("pearson", "spearman")

# Predicate kinds — the variable a ``Split`` partitions on:
# - "score"        : threshold the input Vector at its median (NAME-BLIND).
# - "stage_late"   : clinical predicate, True for AJCC stage III / IV.
PREDICATE_KINDS = ("score", "stage_late")

# Adjustment columns the ``Effect`` operator residualises on. Clinical
# only — never gene names.
EFFECT_ADJUST = ("stage", "age")

# Targets the engine may Associate / Effect / Fit against. Each is a
# named column on the cohort labels — NOT a gene.
ASSOC_TARGETS = ("msi", "tmb")

MIN_GENES_PER_SET = 1
DEFAULT_MAX_GENES_PER_SET = 8
DEFAULT_MAX_DEPTH = 4
DEFAULT_MAX_NODES = 24

# Hard caps for the recursive ``Search`` operator. Off by default; if the
# nested cost is acceptable, enable via the search-rate knob — never
# silently drop.
SEARCH_MAX_K = 4
SEARCH_MAX_COLS = 200