""" 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}