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, cv_group_by=None, cv_groups=None)

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 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
cv_group_by	str | Sequence[str] | None	"design" groups identical factor configurations; a sequence of factor names groups by those regime descriptors. Grouped folds keep all rows of a group together. When selected, cv_r2/cv_rmse report grouped validation; row_cv_r2 and row_cv_rmse retain shuffled row validation for comparison. Both group options unset preserve shuffled row validation.	None
cv_groups	Sequence[object] | NDArray[Any] | None	Explicit group labels aligned with all results rows, mutually exclusive with cv_group_by. Useful for shared dataset/replicate IDs supplied from metadata. Labels are finite numbers, strings, bytes or tuples thereof; missing labels and other objects are refused. One partition tests one caller-declared generalization target.	None

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(  # ruff: ignore[too-many-arguments]
    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,
    cv_group_by: str | Sequence[str] | None = None,
    cv_groups: Sequence[object] | NDArray[Any] | None = None,
) -> 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 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.
        cv_group_by: ``"design"`` groups identical factor configurations;
            a sequence of factor names groups by those regime descriptors.
            Grouped folds keep all rows of a group together. When selected,
            ``cv_r2``/``cv_rmse`` report grouped validation; ``row_cv_r2`` and
            ``row_cv_rmse`` retain shuffled row validation for comparison.
            Both group options unset preserve shuffled row validation.
        cv_groups: Explicit group labels aligned with all results rows, mutually
            exclusive with cv_group_by. Useful for shared dataset/replicate IDs
            supplied from metadata. Labels are finite numbers, strings, bytes or
            tuples thereof; missing labels and other objects are refused. One
            partition tests one caller-declared generalization target.

    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)
    external = (
        None
        if cv_groups is None
        else _external_group_labels(cv_groups, len(results.configs))
    )
    group_names, groups = _validation_groups(
        results.configs, factors, cv_group_by, external
    )

    surrogate = SurrogateModel(
        method=method,
        encoder=encoder,
        observable_names=[],
        models=[],
        cv_group_by=group_names,
        cv_group_source="external"
        if external is not None
        else "configuration"
        if group_names is not None
        else "rows",
        cv_group_labels=external,
        validation_groups=[] if groups is None else groups.tolist(),
    )
    low_accuracy: list[str] = []
    for j, name in enumerate(results.observable_names):
        y = results.scores[:, j]
        mask = np.isfinite(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])
        surrogate.models.append(model)
        surrogate.observable_names.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),
        )
        surrogate.row_cv_r2[name], surrogate.row_cv_rmse[name] = r2, rmse
        if groups is not None:
            r2, rmse = _cross_val_accuracy(
                x_full[mask],
                y[mask],
                method=method,
                seed=seed,
                n_estimators=n_estimators,
                n_folds=cv_folds,
                groups=groups[mask],
            )
        configs = [cfg for cfg, keep in zip(results.configs, mask, strict=True) if keep]
        surrogate.support_ranges[name], surrogate.support_levels[name] = (
            _training_support(configs, factors)
        )
        surrogate.cv_r2[name] = r2
        surrogate.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 surrogate.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 surrogate

Validation and support

fit_surrogate(..., cv_group_by="design") keeps identical factor configurations in one validation fold, preventing replicates of a design from appearing in both training and validation. Pass a sequence such as cv_group_by=["noise"] to hold out whole regimes. cv_r2 and cv_rmse then describe held-out groups; row_cv_r2 and row_cv_rmse retain shuffled row validation for comparison. With both group options unset, the existing row validation behavior is retained. Grouping is explicit: different design-point ids with identical factor values belong to the same design group. Fewer than two groups produce nan metrics; constant targets have undefined R². Non-finite scores are excluded separately for each observable. validation_groups exposes row-aligned group ids.

Use cv_groups=[meta["dataset_id"] for meta in results.metadata] for an explicit shared-dataset/replicate partition across configurations. It is mutually exclusive with cv_group_by and must align with all original rows. Finite scalar/tuple labels are validated without splitting equal numeric identities; missing labels are refused. cv_group_source and cv_group_labels preserve the source and external labels. One partition does not automatically block crossed dependencies. See the complete conditional workflow for the distinction between held-out designs, regimes and datasets.

model.support(configs) reports continuous factors outside the observed training range and categorical/discrete levels absent from finite training rows for each observable. Predictions warn by default and still extrapolate without clipping. Use warn_support=False to disable these warnings explicitly. Levels outside the factor's declared domain remain invalid. Bounds use raw factor units, including for log-scaled factors. Marginal ranges do not establish joint support, adequate sampling density, or calibrated prediction uncertainty.

trade_study.SurrogateModel(method, encoder, observable_names, models, cv_r2=dict(), cv_rmse=dict(), row_cv_r2=dict(), row_cv_rmse=dict(), cv_group_by=None, validation_groups=list(), support_ranges=dict(), support_levels=dict(), cv_group_source='rows', cv_group_labels=None) 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. 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.
row_cv_r2	dict[str, float]	Shuffled row-validation R², even when grouping is selected.
row_cv_rmse	dict[str, float]	Shuffled row-validation RMSE in observable units.
cv_group_by	tuple[str, ...] | None	Configuration grouping names, or None for rows/external groups.
cv_group_source	str	"rows", "configuration" or explicit "external" labels.
cv_group_labels	tuple[Hashable, ...] | None	Validated external labels aligned with results rows, or None. Equal finite numeric labels share one group.
validation_groups	list[int]	Group ids aligned with the original results rows.
support_ranges	dict[str, dict[str, tuple[float, float]]]	Observed continuous bounds per fitted observable.
support_levels	dict[str, dict[str, list[Any]]]	Observed levels per fitted observable.

support(configs)

Inspect observed support without predicting or emitting warnings.

Parameters:

Name	Type	Description	Default
configs	Sequence[dict[str, Any]]	Factor configurations to inspect.	required

Returns:

Type	Description
list[PredictionSupport]	One diagnostic per configuration and fitted observable. Support
list[PredictionSupport]	uses finite training rows for that observable, in raw factor units.

Source code in src/trade_study/surrogate.py

def support(self, configs: Sequence[dict[str, Any]]) -> list[PredictionSupport]:
    """Inspect observed support without predicting or emitting warnings.

    Args:
        configs: Factor configurations to inspect.

    Returns:
        One diagnostic per configuration and fitted observable. Support
        uses finite training rows for that observable, in raw factor units.

    """
    reports = []
    for index, config in enumerate(configs):
        for observable in self.observable_names:
            outside = {
                name: bounds
                for name, bounds in self.support_ranges.get(observable, {}).items()
                if not bounds[0] <= float(config[name]) <= bounds[1]
            }
            unseen = {
                name: config[name]
                for name, levels in self.support_levels.get(observable, {}).items()
                if config[name] not in levels
            }
            reports.append(PredictionSupport(index, observable, outside, unseen))
    return reports

predict(config, *, warn_support=True)

Predict observables for a single config.

Parameters:

Name	Type	Description	Default
config	dict[str, Any]	Factor-keyed config dict.	required
warn_support	bool	Warn on values outside observed training support.	True

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], *, warn_support: bool = True
) -> dict[str, float]:
    """Predict observables for a single config.

    Args:
        config: Factor-keyed config dict.
        warn_support: Warn on values outside observed training support.

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

predict_batch(configs, *, warn_support=True)

Predict observables for a batch of configs.

Parameters:

Name	Type	Description	Default
configs	Sequence[dict[str, Any]]	Sequence of factor-keyed config dicts.	required
warn_support	bool	Warn on values outside observed training support.	True

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]],
    *,
    warn_support: bool = True,
) -> dict[str, NDArray[np.float64]]:
    """Predict observables for a batch of configs.

    Args:
        configs: Sequence of factor-keyed config dicts.
        warn_support: Warn on values outside observed training support.

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

spread_batch(configs, *, warn_support=True)

Predictive spread per observable for a batch of configs.

For method="gp" this is the predictive standard deviation; for method="rf" it is the standard deviation of the individual trees' predictions, a relative (not calibrated) measure of disagreement.

Parameters:

Name	Type	Description	Default
configs	Sequence[dict[str, Any]]	Sequence of factor-keyed config dicts.	required
warn_support	bool	Warn on values outside observed training support.	True

Returns:

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

Source code in src/trade_study/surrogate.py

def spread_batch(
    self,
    configs: Sequence[dict[str, Any]],
    *,
    warn_support: bool = True,
) -> dict[str, NDArray[np.float64]]:
    """Predictive spread per observable for a batch of configs.

    For ``method="gp"`` this is the predictive standard deviation; for
    ``method="rf"`` it is the standard deviation of the individual trees'
    predictions, a relative (not calibrated) measure of disagreement.

    Args:
        configs: Sequence of factor-keyed config dicts.
        warn_support: Warn on values outside observed training support.

    Returns:
        Mapping from observable name to a length-``len(configs)`` array.
    """
    x = self._encode(configs, warn_support=warn_support)
    out: dict[str, NDArray[np.float64]] = {}
    for name, model in zip(self.observable_names, self.models, strict=True):
        if self.method == "gp":
            _, std = model.predict(x, return_std=True)
        else:
            std = np.std([tree.predict(x) for tree in model.estimators_], axis=0)
        out[name] = np.asarray(std, dtype=np.float64)
    return out

uncertainty(config, *, warn_support=True)

Predictive standard deviation per observable (GP only).

Parameters:

Name	Type	Description	Default
config	dict[str, Any]	Factor-keyed config dict.	required
warn_support	bool	Warn on values outside observed training support.	True

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], *, warn_support: bool = True
) -> dict[str, float]:
    """Predictive standard deviation per observable (GP only).

    Args:
        config: Factor-keyed config dict.
        warn_support: Warn on values outside observed training support.

    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._encode([config], warn_support=warn_support)
    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

trade_study.PredictionSupport(config_index, observable, outside_ranges, unseen_levels) dataclass

Marginal training-support diagnostics for one query and observable.

These checks flag extrapolation; they do not establish joint support or provide a calibrated uncertainty guarantee.

supported property

Whether every factor lies within its observed marginal support.

Returns:

Type	Description
bool	True when no range or level check failed.