mlate/optimize.py

"""
Scaffold optimisation
=====================

Searches the formulation space for the composition and printing conditions that
maximise WSSQ, using the deployed classifiers as the objective.

Extracted from the Streamlit application deliberately. In the previous version
the objective function lived inside the interface, which meant it could not be
run headlessly, could not be tested, and could not produce a result the
manuscript was able to report. Here the interface calls this module rather than
containing it.

How a classifier becomes a continuous objective
-----------------------------------------------
Printability and Cell Response are ordinal classes, so a classifier returns a
probability vector rather than a number, and an optimiser handed the arg-max
would be climbing a step function with four or five levels. Instead each model's
full probability vector is collapsed to its EXPECTED CLASS VALUE,

    E[class] = sum_k  P(class = k) * label_k

which is continuous, respects the ordinal spacing of the labels, and carries the
model's uncertainty: a formulation the model is unsure about lands between
classes rather than being rounded confidently to one. Those two expected values
are then passed to WSSQ.

What is searched, and what is not
---------------------------------
Biomaterial concentrations, cell density and the printing parameters are
searched within user-supplied ranges. The CELL LINE IS NOT: it is chosen by the
user and held fixed for the run. This matters more than it looks. WSSQ falls
back to printability alone when Cell Response is 1, and an acellular
formulation therefore scores on printability only - so an optimiser allowed to
vary the cell line could raise its objective simply by removing the cells.
Fixing the line makes the acellular case a mode the user selects, not a
degenerate optimum the search can drift into.

Extrapolation is reported, not prevented
----------------------------------------
Nothing stops a user setting a range beyond anything in the corpus, and nothing
here silently clips it. `distance_report` instead states how far the returned
optimum sits outside the observed range of each variable and how close it is to
the nearest real formulation, so a candidate that the models are extrapolating
to is visible as such rather than presented with the same confidence as one
inside the training distribution.
"""

from __future__ import annotations

from dataclasses import dataclass, field

import numpy as np
import pandas as pd

from mlate import config as cfg
from mlate import wssq as wssq_mod

# Optuna's own logging is noisy inside a UI and uninformative in a script.
try:  # pragma: no cover - import guard only
    import optuna
    optuna.logging.set_verbosity(optuna.logging.WARNING)
except Exception:  # pragma: no cover
    optuna = None


@dataclass
class Variable:
    """One searchable input, with the range and grid the user specified."""
    name: str
    low: float
    high: float
    step: float | None = None

    def suggest(self, trial):
        if self.low >= self.high:
            return float(self.low)
        if self.step:
            return float(trial.suggest_float(self.name, self.low, self.high,
                                             step=self.step))
        return float(trial.suggest_float(self.name, self.low, self.high))


@dataclass
class SearchSpace:
    """Everything the optimiser may vary, plus the cell line it may not."""
    cell_line: str
    biomaterials: list[Variable] = field(default_factory=list)
    printing: list[Variable] = field(default_factory=list)
    cell_density: Variable | None = None

    @property
    def is_acellular(self) -> bool:
        return self.cell_line == cfg.ACELLULAR_TOKEN


def class_distribution(estimator, X: np.ndarray, classes: np.ndarray
                       ) -> tuple[np.ndarray, list[dict[float, float]]]:
    """
    Expected class values and the distributions they were taken over.

    Both come from a single `predict_proba` over the whole batch, because the
    two are wanted together and a second call would double the cost of a
    search. They carry different information and both are used. An expected
    value of 2.5 is what a confident prediction split evenly between classes 2
    and 3 produces, and also what a diffuse prediction over 1 and 4 produces;
    the optimiser needs only the number, while the protocol-generation prompt
    is given the distribution so that it can state which case it describes.

    Falls back to the predicted label when a model exposes no probabilities;
    five of the conventional classifiers do not, and for those the objective is
    a step function rather than a smooth one and the distributions are empty.
    """
    labels = np.asarray(classes, dtype=float)
    if not hasattr(estimator, "predict_proba"):
        idx = np.asarray(estimator.predict(X)).astype(int)
        return labels[idx], [{} for _ in idx]
    proba = np.asarray(estimator.predict_proba(X), dtype=float)
    return proba @ labels, [{float(l): float(pr) for l, pr in zip(labels, row)}
                            for row in proba]


def expected_class_value(estimator, X: np.ndarray,
                         classes: np.ndarray) -> np.ndarray:
    """
    Probability-weighted mean of the class labels.

    Falls back to the predicted label when a model exposes no probabilities;
    five of the conventional classifiers do not, and for those the objective is
    a step function rather than a smooth one. The application ranks models by
    benchmarked performance and the probability-capable models occupy the top of
    both leaderboards, so this path is a safeguard rather than the normal case.
    """
    labels = np.asarray(classes, dtype=float)
    if hasattr(estimator, "predict_proba"):
        return np.asarray(estimator.predict_proba(X), dtype=float) @ labels
    idx = np.asarray(estimator.predict(X)).astype(int)
    return labels[idx]


# Candidates scored per pass through the models. Thirty-two because the cost
# of an in-context model is the pass over its context and is flat in the batch
# size up to at least sixty-four, while a larger batch buys the sampler fewer
# updates for no further saving.
DEFAULT_BATCH_SIZE = 32


def _as_scorer(model):
    """
    Accept either a pickled classifier bundle or a `mlate.serving.Servable`.

    The three model families no longer share a storage format - a deep network
    is a state dictionary and a foundation model is a stored context - so the
    application passes the uniform wrapper from `mlate.serving`. A plain bundle
    is still accepted, because the analysis stages load their own artefacts
    directly and have no reason to go through the serving layer.
    """
    if isinstance(model, dict):
        return {"estimator": model["estimator"],
                "classes": np.asarray(model["classes"], dtype=float)}
    return {"estimator": model, "classes": np.asarray(model.classes,
                                                      dtype=float)}


class Objective:
    """
    WSSQ of one candidate formulation, as the optimiser sees it.

    Holds the preprocessor and both models so they are loaded once rather than
    per trial, and records every evaluated candidate so the search can be
    inspected afterwards instead of yielding only its arg-max.
    """

    def __init__(self, space: SearchSpace, preprocessor, feature_columns,
                 print_model, cell_model,
                 print_weight: float = wssq_mod.DEFAULT_PRINT_WEIGHT,
                 cell_weight: float = wssq_mod.DEFAULT_CELL_WEIGHT):
        self.space = space
        self.pre = preprocessor
        self.feature_columns = list(feature_columns)
        self.print_model = _as_scorer(print_model)
        self.cell_model = _as_scorer(cell_model)
        self.print_weight = print_weight
        self.cell_weight = cell_weight
        self.history: list[dict] = []

    def candidate_frame(self, values: dict | list[dict]) -> pd.DataFrame:
        """One row per candidate, in the raw feature layout the preprocessor
        expects."""
        batch = [values] if isinstance(values, dict) else list(values)
        rows = []
        for v in batch:
            row = {c: 0.0 for c in self.feature_columns}
            row.update({k: x for k, x in v.items() if k in row})
            row[cfg.CELL_COLS[0]] = self.space.cell_line
            rows.append(row)
        return pd.DataFrame(rows)[self.feature_columns]

    def evaluate_batch(self, batch: list[dict]) -> list[dict]:
        """
        Score many candidates in one pass through each model.

        This is the form the search uses, and the reason a foundation model can
        be offered at all. Their cost is dominated by the pass over the 2,646
        context rows, which is paid once per call however many candidates ride
        in it: on this corpus a batch of thirty-two costs what a single row
        costs. Scoring candidates one at a time, as the previous release did,
        charged that pass once per candidate and made an in-context model look
        two orders of magnitude more expensive than it is. The conventional
        classifiers gain from it too, for the ordinary reason that per-call
        Python overhead is then divided across the batch.
        """
        X = np.asarray(self.pre.transform(self.candidate_frame(batch)),
                       dtype=float)

        exp_p, proba_p = class_distribution(
            self.print_model["estimator"], X, self.print_model["classes"])

        # An acellular formulation has no cell response to predict. Asking the
        # model anyway would feed the optimiser a number with no referent; the
        # rating scale defines class 1 as "not applicable", so that is what is
        # used and WSSQ takes its printability-only branch.
        if self.space.is_acellular:
            exp_c = np.ones(len(batch))
            proba_c = [{}] * len(batch)
        else:
            exp_c, proba_c = class_distribution(
                self.cell_model["estimator"], X, self.cell_model["classes"])

        out = []
        for i in range(len(batch)):
            out.append({
                "wssq": float(wssq_mod.compute_wssq(
                    float(exp_p[i]), float(exp_c[i]),
                    self.print_weight, self.cell_weight)),
                "expected_printability": float(exp_p[i]),
                "expected_cell_response": float(exp_c[i]),
                "printability_proba": proba_p[i],
                "cell_response_proba": proba_c[i]})
        return out

    def evaluate(self, values: dict) -> dict:
        """Score one explicit formulation. Used to report the winner."""
        return self.evaluate_batch([values])[0]

    def suggest(self, trial) -> dict:
        """One candidate drawn from the search space, as parameter values."""
        values = {v.name: v.suggest(trial)
                  for v in self.space.biomaterials + self.space.printing}
        if self.space.cell_density is not None and not self.space.is_acellular:
            values[cfg.CELL_COLS[1]] = self.space.cell_density.suggest(trial)
        else:
            values[cfg.CELL_COLS[1]] = 0.0
        return values

    def record(self, values: dict, out: dict) -> None:
        """Append one evaluated candidate to the search history."""
        # history is exported to a spreadsheet; the two probability dictionaries
        # do not belong in it.
        self.history.append({**values, **{k: v for k, v in out.items()
                                          if not k.endswith("_proba")}})

    def __call__(self, trial) -> float:
        """Sequential form, kept for a caller that drives Optuna directly."""
        values = self.suggest(trial)
        out = self.evaluate(values)
        self.record(values, out)
        return out["wssq"]


def optimise(objective: Objective, n_trials: int = 100,
             seed: int = cfg.RANDOM_STATE, progress=None,
             batch_size: int = DEFAULT_BATCH_SIZE):
    """
    Run the search and return (best_values, best_score, study).

    Tree-structured Parzen estimation, which is the Bayesian method used in the
    published version: it models P(parameters | score) and samples where good
    scores are likely, rather than sampling the space uniformly. `multivariate`
    and `group` let it model interactions between parameters instead of treating
    each in isolation, which matters here because printability depends on
    combinations - a concentration that prints well at one nozzle diameter does
    not at another.

    Candidates are drawn and scored in batches through Optuna's ask-and-tell
    interface rather than one at a time. The trade is small and one-sided: the
    candidates within a batch are proposed from the same posterior, so the
    sampler sees `batch_size` fewer updates over a run, against a reduction in
    scoring cost of one to two orders of magnitude. It is what makes the
    in-context foundation models usable in a search at all - see
    `Objective.evaluate_batch` - and it leaves the conventional classifiers
    faster as well. Set `batch_size=1` to recover strictly sequential
    behaviour.
    """
    if optuna is None:  # pragma: no cover
        raise RuntimeError("optuna is required for scaffold optimisation")

    sampler = optuna.samplers.TPESampler(
        seed=seed, n_startup_trials=min(30, max(5, n_trials // 4)),
        multivariate=True, group=True, consider_prior=True)
    study = optuna.create_study(direction="maximize", sampler=sampler)

    done = 0
    while done < n_trials:
        size = min(max(1, batch_size), n_trials - done)
        trials = [study.ask() for _ in range(size)]
        batch = [objective.suggest(t) for t in trials]
        results = objective.evaluate_batch(batch)
        for trial, values, out in zip(trials, batch, results):
            objective.record(values, out)
            study.tell(trial, out["wssq"])
        done += size
        if progress is not None:
            progress(done, n_trials, study.best_value)

    return study.best_trial.params, float(study.best_value), study


def distance_report(values: dict, corpus: pd.DataFrame) -> dict:
    """
    How far outside the observed data the proposed formulation sits.

    Two numbers, because they answer different questions. `out_of_range` names
    the variables whose value exceeds anything recorded in the corpus - a
    concentration no one has published. `nearest_neighbour_distance` is the
    scaled Euclidean distance to the closest real formulation, which catches the
    subtler case where every individual value is ordinary but the combination
    has never been attempted.
    """
    out_of_range = {}
    for name, v in values.items():
        if name not in corpus.columns:
            continue
        col = pd.to_numeric(corpus[name], errors="coerce")
        col = col[col > 0]
        if len(col) == 0:
            continue
        lo, hi = float(col.min()), float(col.max())
        if v > 0 and (v < lo or v > hi):
            out_of_range[name] = {"value": float(v), "observed_min": lo,
                                  "observed_max": hi}

    shared = [c for c in values if c in corpus.columns]
    nn = float("nan")
    if shared:
        sub = corpus[shared].apply(pd.to_numeric, errors="coerce").fillna(0.0)
        span = (sub.max() - sub.min()).replace(0, 1.0)
        target = pd.Series({c: values[c] for c in shared})
        nn = float((((sub - target) / span) ** 2).sum(axis=1).pow(0.5).min())

    return {"n_out_of_range": len(out_of_range),
            "out_of_range": out_of_range,
            "nearest_neighbour_distance": nn}