Surrogate

Cheap regression surrogates fit over a ResultsTable for predicting observables at untested configurations.

Install via the optional extra:

uv pip install 'trade-study[surrogate]'

trade_study.fit_surrogate(results, factors, *, method='gp', seed=0, n_estimators=200, cv_folds=5, warn_below_r2=0.0)

Fit a per-observable surrogate over a :class:ResultsTable.

Rows whose score column contains NaN are dropped on a per-observable basis (so a partially-evaluated trial still contributes to the observables it does have).

After fitting each observable's model on all its available rows, also computes a held-out cross-validated R^2/RMSE (#114) so callers can tell a well-fit surrogate from one that's effectively guessing in a sparse or noisy region of the design space -- neither backend reports this on its own (RF's cheap OOB score isn't available for GP, so CV is used uniformly for both, at the cost of cv_folds extra fits per observable).

Parameters:

Name	Type	Description	Default
results	ResultsTable	A :class:ResultsTable from a previous study run.	required
factors	list[Factor]	Factor definitions used to encode results.configs. Must cover every key in the configs.	required
method	str	"gp" for Gaussian process (Matern 1.5 + WhiteKernel) or "rf" for a random forest.	'gp'
seed	int	Random seed forwarded to the backend estimators.	0
n_estimators	int	Number of trees for the "rf" backend; ignored for "gp".	200
cv_folds	int	Number of cross-validation folds used to compute cv_r2/cv_rmse. Clamped down to the observable's available row count when fewer than cv_folds rows exist; skipped (nan) for an observable with fewer than 2 rows.	5
warn_below_r2	float | None	If not None, emit a UserWarning naming any observable whose cv_r2 falls below this threshold -- 0.0 (the default) flags a surrogate that predicts no better than the training mean. Pass None to disable.	0.0

Returns:

Type	Description
SurrogateModel	A fitted :class:SurrogateModel.

Raises:

Type	Description
ValueError	If method is unknown, results is empty, or no observable has at least two non-NaN training rows.

Source code in src/trade_study/surrogate.py

def fit_surrogate(
    results: ResultsTable,
    factors: list[Factor],
    *,
    method: str = "gp",
    seed: int = 0,
    n_estimators: int = 200,
    cv_folds: int = 5,
    warn_below_r2: float | None = 0.0,
) -> SurrogateModel:
    """Fit a per-observable surrogate over a :class:`ResultsTable`.

    Rows whose score column contains ``NaN`` are dropped on a
    per-observable basis (so a partially-evaluated trial still
    contributes to the observables it does have).

    After fitting each observable's model on all its available rows,
    also computes a held-out cross-validated R^2/RMSE (#114) so callers
    can tell a well-fit surrogate from one that's effectively guessing
    in a sparse or noisy region of the design space -- neither backend
    reports this on its own (RF's cheap OOB score isn't available for
    GP, so CV is used uniformly for both, at the cost of ``cv_folds``
    extra fits per observable).

    Args:
        results: A :class:`ResultsTable` from a previous study run.
        factors: Factor definitions used to encode ``results.configs``.
            Must cover every key in the configs.
        method: ``"gp"`` for Gaussian process (Matern 1.5 + WhiteKernel)
            or ``"rf"`` for a random forest.
        seed: Random seed forwarded to the backend estimators.
        n_estimators: Number of trees for the ``"rf"`` backend; ignored
            for ``"gp"``.
        cv_folds: Number of cross-validation folds used to compute
            ``cv_r2``/``cv_rmse``. Clamped down to the observable's
            available row count when fewer than ``cv_folds`` rows exist;
            skipped (``nan``) for an observable with fewer than 2 rows.
        warn_below_r2: If not ``None``, emit a ``UserWarning`` naming any
            observable whose ``cv_r2`` falls below this threshold --
            0.0 (the default) flags a surrogate that predicts no better
            than the training mean. Pass ``None`` to disable.

    Returns:
        A fitted :class:`SurrogateModel`.

    Raises:
        ValueError: If ``method`` is unknown, ``results`` is empty, or
            no observable has at least two non-NaN training rows.
    """
    if method not in _SUPPORTED_METHODS:
        msg = (
            f"Unknown surrogate method {method!r}. "
            f"Supported: {sorted(_SUPPORTED_METHODS)}"
        )
        raise ValueError(msg)
    if not results.configs:
        msg = "fit_surrogate: results table is empty"
        raise ValueError(msg)

    encoder = _FactorEncoder.from_factors(factors)
    x_full = encoder.transform(results.configs)

    models: list[Any] = []
    fitted_obs: list[str] = []
    cv_r2: dict[str, float] = {}
    cv_rmse: dict[str, float] = {}
    low_accuracy: list[str] = []
    for j, name in enumerate(results.observable_names):
        y = results.scores[:, j]
        mask = ~np.isnan(y)
        n_rows = int(mask.sum())
        if n_rows < 2:
            continue
        model = _make_estimator(method, seed=seed, n_estimators=n_estimators)
        model.fit(x_full[mask], y[mask])
        models.append(model)
        fitted_obs.append(name)

        r2, rmse = _cross_val_accuracy(
            x_full[mask],
            y[mask],
            method=method,
            seed=seed,
            n_estimators=n_estimators,
            n_folds=min(cv_folds, n_rows),
        )
        cv_r2[name] = r2
        cv_rmse[name] = rmse
        if warn_below_r2 is not None and np.isfinite(r2) and r2 < warn_below_r2:
            low_accuracy.append(name)

    if not models:
        msg = (
            "fit_surrogate: no observable has at least 2 non-NaN training "
            "rows; nothing to fit"
        )
        raise ValueError(msg)

    if low_accuracy:
        warnings.warn(
            f"fit_surrogate: cross-validated R^2 below {warn_below_r2} for "
            f"{low_accuracy} -- predictions for these observables may not "
            f"be trustworthy. See SurrogateModel.cv_r2 for exact values.",
            UserWarning,
            stacklevel=2,
        )

    return SurrogateModel(
        method=method,
        encoder=encoder,
        observable_names=fitted_obs,
        models=models,
        cv_r2=cv_r2,
        cv_rmse=cv_rmse,
    )

trade_study.SurrogateModel(method, encoder, observable_names, models, cv_r2=dict(), cv_rmse=dict()) dataclass

Fitted surrogate over a :class:ResultsTable.

Use :func:fit_surrogate to construct one. Per-observable backend estimators are stored in models; encoding is shared across them.

Attributes:

Name	Type	Description
method	str	"gp" or "rf".
encoder	_FactorEncoder	Factor encoder used at fit time.
observable_names	list[str]	Column names of the predicted observables.
models	list[Any]	One fitted scikit-learn estimator per observable.
cv_r2	dict[str, float]	Per-observable held-out cross-validated R^2 (#114). Missing for an observable if too few rows were available to run CV (fewer than 2 folds); float("nan") in that case.
cv_rmse	dict[str, float]	Per-observable held-out cross-validated RMSE, companion to cv_r2 in the observable's original units.

predict(config)

Predict observables for a single config.

Parameters:

Name	Type	Description	Default
config	dict[str, Any]	Factor-keyed config dict.	required

Returns:

Type	Description
dict[str, float]	Mapping from observable name to predicted scalar.

Source code in src/trade_study/surrogate.py

def predict(self, config: dict[str, Any]) -> dict[str, float]:
    """Predict observables for a single config.

    Args:
        config: Factor-keyed config dict.

    Returns:
        Mapping from observable name to predicted scalar.
    """
    x = self.encoder.transform([config])
    return {
        name: float(model.predict(x)[0])
        for name, model in zip(self.observable_names, self.models, strict=True)
    }

predict_batch(configs)

Predict observables for a batch of configs.

Parameters:

Name	Type	Description	Default
configs	Sequence[dict[str, Any]]	Sequence of factor-keyed config dicts.	required

Returns:

Type	Description
dict[str, NDArray[float64]]	Mapping from observable name to a length-len(configs)
dict[str, NDArray[float64]]	array of predictions.

Source code in src/trade_study/surrogate.py

def predict_batch(
    self,
    configs: Sequence[dict[str, Any]],
) -> dict[str, NDArray[np.float64]]:
    """Predict observables for a batch of configs.

    Args:
        configs: Sequence of factor-keyed config dicts.

    Returns:
        Mapping from observable name to a length-``len(configs)``
        array of predictions.
    """
    x = self.encoder.transform(configs)
    return {
        name: np.asarray(model.predict(x), dtype=np.float64)
        for name, model in zip(self.observable_names, self.models, strict=True)
    }

uncertainty(config)

Predictive standard deviation per observable (GP only).

Parameters:

Name	Type	Description	Default
config	dict[str, Any]	Factor-keyed config dict.	required

Returns:

Type	Description
dict[str, float]	Mapping from observable name to predictive standard deviation.

Raises:

Type	Description
NotImplementedError	If the backend does not expose calibrated uncertainties (currently anything other than "gp").

Source code in src/trade_study/surrogate.py

def uncertainty(self, config: dict[str, Any]) -> dict[str, float]:
    """Predictive standard deviation per observable (GP only).

    Args:
        config: Factor-keyed config dict.

    Returns:
        Mapping from observable name to predictive standard deviation.

    Raises:
        NotImplementedError: If the backend does not expose calibrated
            uncertainties (currently anything other than ``"gp"``).
    """
    if self.method != "gp":
        msg = (
            f"uncertainty() is only supported for method='gp'; "
            f"this surrogate uses method={self.method!r}"
        )
        raise NotImplementedError(msg)
    x = self.encoder.transform([config])
    out: dict[str, float] = {}
    for name, model in zip(self.observable_names, self.models, strict=True):
        _, std = model.predict(x, return_std=True)
        out[name] = float(std[0])
    return out