abbot 0.1.0__py3-none-any.whl

abbot/__init__.py

@@ -0,0 +1,159 @@
+"""abbot — an active-learning-for-labs layer on top of BoTorch.
+
+From A to B: a Bayesian guide for lab experiments. abbot is NOT a competing
+optimization framework; it is the composite-SBI + lab-experiment-orchestration
+layer on top of botorch (which already provides the surrogate-agnostic
+Model/Posterior layer, the BO/US loop primitives, async/batch, and
+serialization). If it drifts into "framework", stop.
+"""
+
+from abbot.composite import (
+    ObjectiveReducer,
+    Rescore,
+    composite_evaluate,
+)
+from abbot.cost import (
+    CallableCost,
+    CostModel,
+    FittedCost,
+    FixedCost,
+    PowerLawCost,
+    SourceCost,
+)
+from abbot.experiment import (
+    BatchExperiment,
+    Experiment,
+    ManualExperiment,
+    MockAnalyticExperiment,
+    MockDatasetExperiment,
+    TensorExperiment,
+    Verdict,
+)
+from abbot.goal import (
+    BandIntersection,
+    FunctionGoal,
+    Goal,
+    LevelSet,
+    PercentileSet,
+    TopK,
+    Wishlist,
+    goals_from_spec,
+)
+from abbot.history import DataStore, History, Record, ResumeMismatchError
+from abbot.interop.sbi import PosteriorSampler, TrainedPosterior, sbi_evaluate
+from abbot.loop import drive, run_campaign
+from abbot.model import (
+    EnsembleConstructor,
+    ModelConstructor,
+    ModelListConstructor,
+    MultiFidelityGPConstructor,
+    MultiSourceGPConstructor,
+    MultiTaskGPConstructor,
+    SingleTaskGPConstructor,
+    default_constructor,
+)
+from abbot.policy import (
+    BAXPolicy,
+    ParetoBO,
+    Policy,
+    SimpleBO,
+    UncertaintySampling,
+    default_policy,
+)
+from abbot.problem import (
+    Box,
+    ContinuousDomain,
+    EnumeratedDomain,
+    InequalitySpec,
+    IntegerDomain,
+    LinearConstraint,
+    ObjectiveSpec,
+    OutputSpec,
+    ProblemSpec,
+    Source,
+    StructuredDomain,
+)
+from abbot.session import (
+    Campaign,
+    CampaignResult,
+    CampaignState,
+    Candidate,
+    PersistentFaultError,
+    Targets,
+    apply_targets,
+    default_initial_designs,
+)
+from abbot.supervision import Supervisor
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "BAXPolicy",
+    "BandIntersection",
+    "BatchExperiment",
+    "Box",
+    "CallableCost",
+    "Campaign",
+    "CampaignResult",
+    "CampaignState",
+    "Candidate",
+    "ContinuousDomain",
+    "CostModel",
+    "DataStore",
+    "EnsembleConstructor",
+    "EnumeratedDomain",
+    "Experiment",
+    "FittedCost",
+    "FixedCost",
+    "FunctionGoal",
+    "Goal",
+    "History",
+    "InequalitySpec",
+    "IntegerDomain",
+    "LevelSet",
+    "LinearConstraint",
+    "ManualExperiment",
+    "MockAnalyticExperiment",
+    "MockDatasetExperiment",
+    "ModelConstructor",
+    "ModelListConstructor",
+    "MultiFidelityGPConstructor",
+    "MultiSourceGPConstructor",
+    "MultiTaskGPConstructor",
+    "ObjectiveReducer",
+    "ObjectiveSpec",
+    "OutputSpec",
+    "ParetoBO",
+    "PercentileSet",
+    "PersistentFaultError",
+    "Policy",
+    "PosteriorSampler",
+    "PowerLawCost",
+    "ProblemSpec",
+    "Record",
+    "Rescore",
+    "ResumeMismatchError",
+    "SimpleBO",
+    "SingleTaskGPConstructor",
+    "Source",
+    "SourceCost",
+    "StructuredDomain",
+    "Supervisor",
+    "Targets",
+    "TensorExperiment",
+    "TopK",
+    "TrainedPosterior",
+    "UncertaintySampling",
+    "Verdict",
+    "Wishlist",
+    "__version__",
+    "apply_targets",
+    "composite_evaluate",
+    "default_constructor",
+    "default_initial_designs",
+    "default_policy",
+    "drive",
+    "goals_from_spec",
+    "run_campaign",
+    "sbi_evaluate",
+]

abbot/acq_opt.py

@@ -0,0 +1,206 @@
+"""The CandidateOptimizer seam: candidate proposal decoupled from the surrogate.
+
+``Continuous`` covers any space with a continuous or
+integer axis (``optimize_acqf``; integer axes relaxed then rounded; enumerated axes
+alongside them enumerated by ``optimize_acqf_mixed``). ``DiscreteArgmax`` covers
+all-enumerated spaces (``optimize_acqf_discrete``). True combinatorial
+(non-enumerable) search is an extension point, not solved speculatively.
+
+``fixed_features={position: value}`` pins inputs, which is how a policy that doesn't
+choose sources proposes at their target (``spec.target_features()``).
+
+Decoupling proposal from the surrogate is what lets a discrete materials problem
+and a continuous accelerator problem share one policy.
+"""
+
+from __future__ import annotations
+
+from typing import Protocol, runtime_checkable
+
+import torch
+from botorch.acquisition import AcquisitionFunction
+from botorch.optim import optimize_acqf, optimize_acqf_discrete, optimize_acqf_mixed
+
+from abbot.problem import ContinuousDomain, EnumeratedDomain, IntegerDomain, ProblemSpec
+
+
+@runtime_checkable
+class CandidateOptimizer(Protocol):
+    """Maximize an acquisition function over the inputs, return (n, d).
+
+    ``fixed_features`` pins inputs ``{position: value}``; optimizers that predate it
+    may not accept it, and policies only pass it when the spec has sources.
+    """
+
+    def optimize(
+        self,
+        acq: AcquisitionFunction,
+        spec: ProblemSpec,
+        n: int,
+        *,
+        fixed_features: dict[int, float] | None = None,
+    ) -> torch.Tensor: ...
+
+
+# Above this many combinations of enumerated values, a mixed space is too big to
+# enumerate per proposal; pass the combinations you care about to DiscreteArgmax.
+MAX_MIXED_COMBINATIONS = 256
+
+
+class Continuous:
+    """Proposal via botorch ``optimize_acqf`` for any space with a continuous or
+    integer input. Integer inputs are optimized as continuous and rounded;
+    enumerated inputs alongside them are enumerated (``optimize_acqf_mixed``, one
+    sub-problem per combination of their values, up to
+    ``MAX_MIXED_COMBINATIONS``). With ``log=True`` inputs (and no enumerated ones
+    to enumerate), the restarts start from the best of ``raw_samples`` log-uniform
+    designs, so the search covers every decade rather than mostly the top one."""
+
+    def __init__(self, num_restarts: int = 10, raw_samples: int = 256) -> None:
+        self.num_restarts = num_restarts
+        self.raw_samples = raw_samples
+
+    def optimize(
+        self,
+        acq: AcquisitionFunction,
+        spec: ProblemSpec,
+        n: int,
+        *,
+        fixed_features: dict[int, float] | None = None,
+    ) -> torch.Tensor:
+        from itertools import product
+
+        inputs = spec.inputs
+        unsupported = [
+            d.name for d in inputs
+            if not isinstance(d, (ContinuousDomain, IntegerDomain, EnumeratedDomain))
+        ]
+        if unsupported:
+            raise NotImplementedError(
+                f"Continuous optimizer can't propose structured inputs {unsupported}."
+            )
+        fixed = dict(fixed_features or {})
+        integer = [j for j, d in enumerate(inputs) if isinstance(d, IntegerDomain)]
+
+        def round_integers(X: torch.Tensor) -> torch.Tensor:
+            if integer:
+                X = X.clone()
+                X[..., integer] = X[..., integer].round()
+            return X
+
+        enum = [
+            j for j, d in enumerate(inputs)
+            if isinstance(d, EnumeratedDomain) and j not in fixed
+        ]
+        kwargs = {
+            "acq_function": acq,
+            "bounds": spec.bounds_tensor(),
+            "q": n,
+            "num_restarts": self.num_restarts,
+            "raw_samples": self.raw_samples,
+            "inequality_constraints": spec.inequality_constraints(),
+            "post_processing_func": round_integers,
+        }
+        if not enum:
+            if set(spec.log_indices()) - set(fixed):
+                kwargs["batch_initial_conditions"] = self._log_uniform_starts(acq, spec, n, fixed)
+                kwargs.pop("raw_samples")
+            candidates, _ = optimize_acqf(fixed_features=fixed or None, **kwargs)
+            return candidates.detach()
+        combos = list(product(*(inputs[j].codes for j in enum)))
+        if len(combos) > MAX_MIXED_COMBINATIONS:
+            raise ValueError(
+                f"{len(combos)} combinations of enumerated values "
+                f"({[inputs[j].name for j in enum]}) is too many to enumerate per "
+                f"proposal (limit {MAX_MIXED_COMBINATIONS}); make some of them "
+                "continuous or integer, or use DiscreteArgmax(choices=...)."
+            )
+        fixed_list = [{**fixed, **dict(zip(enum, map(float, c)))} for c in combos]
+        candidates, _ = optimize_acqf_mixed(fixed_features_list=fixed_list, **kwargs)
+        return candidates.detach()
+
+
+    def _log_uniform_starts(self, acq, spec: ProblemSpec, n: int, fixed: dict[int, float]):
+        """``(num_restarts, n, d)`` starting batches: the best of ``raw_samples``
+        batches drawn by ``spec.sample_designs`` (log-uniform on log axes)."""
+        X = spec._sample(self.raw_samples * n, None, fixed).reshape(self.raw_samples, n, -1)
+        with torch.no_grad():
+            value = acq(X).reshape(-1)
+        best = value.nan_to_num(float("-inf")).topk(min(self.num_restarts, len(value))).indices
+        return X[best]
+
+
+class DiscreteArgmax:
+    """Enumerated proposal via botorch ``optimize_acqf_discrete``.
+
+    ``choices`` is the (N, d) tensor of enumerable candidates (every input). When
+    omitted it is built as the Cartesian product of each enumerated or integer
+    input's values (small spaces only — large products should pass ``choices``
+    explicitly). ``fixed_features`` keeps only the choices that match.
+    """
+
+    def __init__(self, choices: torch.Tensor | None = None) -> None:
+        self.choices = choices
+
+    def _enumerate(self, spec: ProblemSpec) -> torch.Tensor:
+        from itertools import product
+
+        axes = []
+        for d in spec.inputs:
+            if isinstance(d, EnumeratedDomain):
+                axes.append(d.codes)
+            elif isinstance(d, IntegerDomain):
+                axes.append(list(range(d.lower, d.upper + 1)))
+            else:
+                raise TypeError(
+                    "DiscreteArgmax without explicit `choices` requires every "
+                    f"input to be enumerated or integer; {d.name!r} is not."
+                )
+        return torch.tensor(list(product(*axes)), dtype=torch.double)
+
+    def optimize(
+        self,
+        acq: AcquisitionFunction,
+        spec: ProblemSpec,
+        n: int,
+        *,
+        fixed_features: dict[int, float] | None = None,
+    ) -> torch.Tensor:
+        if self.choices is None:
+            choices = self._enumerate(spec)
+        else:  # only those the spec allows (with travel limits, the ones in reach)
+            choices = self.choices.double()
+            bounds = spec.bounds_tensor()
+            choices = choices[((choices >= bounds[0] - 1e-9) & (choices <= bounds[1] + 1e-9)).all(-1)]
+            if not len(choices):
+                raise ValueError(
+                    "no choices lie within the bounds (with travel limits: within "
+                    "max_step of the last design)"
+                )
+        choices = choices.double()
+        for j, value in (fixed_features or {}).items():
+            choices = choices[torch.isclose(choices[:, j], torch.tensor(value).double())]
+        if not len(choices):
+            raise ValueError(f"no choices match the fixed inputs {fixed_features}")
+        candidates, _ = optimize_acqf_discrete(acq, q=n, choices=choices)
+        return candidates.detach()
+
+
+def optimize(optimizer, acq, spec: ProblemSpec, n: int, fixed_features=None) -> torch.Tensor:
+    """``optimizer.optimize``, passing ``fixed_features`` only when there are some,
+    so custom optimizers written before it existed keep working without sources."""
+    if fixed_features:
+        return optimizer.optimize(acq, spec, n, fixed_features=fixed_features)
+    return optimizer.optimize(acq, spec, n)
+
+
+def default_optimizer(spec: ProblemSpec) -> CandidateOptimizer:
+    """Pick the CandidateOptimizer the spec supports.
+
+    - all-enumerated inputs -> ``DiscreteArgmax`` (``optimize_acqf_discrete``)
+    - otherwise             -> ``Continuous`` (``optimize_acqf``; integers rounded,
+      enumerated inputs alongside continuous ones enumerated)
+    """
+    if all(isinstance(d, EnumeratedDomain) for d in spec.inputs):
+        return DiscreteArgmax()
+    return Continuous()

abbot/analysis.py

@@ -0,0 +1,256 @@
+"""Live-analysis built-ins, computed from a campaign's history + fitted model.
+
+These are pure functions of ``(History, spec[, model])`` so they are unit-testable
+without a UI and reusable by a notebook or an agent. The UI registers them through
+``Campaign``'s ``on_step`` hook and renders the results; callers can register their
+own analyses the same way. Kept deliberately small — the general-purpose "any
+analysis" surface is the hook itself; these are just the common basics.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from abbot.history import History
+from abbot.problem import ProblemSpec
+from abbot.session import canonical_Y, observed_front
+
+if TYPE_CHECKING:
+    from abbot.policy import BAXPolicy
+
+
+def _counts(r, spec: ProblemSpec) -> bool:
+    """A record counts toward the best so far: clean, measured, and satisfying
+    every output constraint."""
+    if not r.usable or r.outputs is None:
+        return False
+    if not spec.has_output_constraints:
+        return True
+    from abbot.feasibility import feasible
+
+    return bool(feasible(spec, r.outputs.double().reshape(1, -1))[0])
+
+
+def convergence(history: History, spec: ProblemSpec) -> list[float]:
+    """Running best objective value vs. step, in RAW units (single-objective only).
+
+    Element ``i`` is the best clean observed objective among the first ``i+1``
+    records, respecting direction (max of raw for MAXIMIZE, min for MINIMIZE).
+    Faulted records carry forward the previous best, and so do records that break
+    an output constraint (NaN until one satisfies them).
+    """
+    if spec.n_objectives != 1:
+        raise ValueError("convergence is single-objective; use hypervolume for MO.")
+    col = spec.objective_indices()[0]  # objective may not be column 0 (OBSERVE outputs)
+    maximize = spec.objectives[0].direction == "MAXIMIZE"
+    best = None
+    out: list[float] = []
+    for r in history.records:
+        if _counts(r, spec):
+            v = float(r.outputs.reshape(-1)[col])
+            if best is None:
+                best = v
+            else:
+                best = max(best, v) if maximize else min(best, v)
+        out.append(best if best is not None else float("nan"))
+    return out
+
+
+def hypervolume(history: History, spec: ProblemSpec) -> list[float]:
+    """Dominated hypervolume of the observed canonical outputs vs. step (MO).
+
+    Uses botorch's exact ``Hypervolume`` against an inferred reference point (from
+    the final canonical set, so the series is comparable across steps). Returns a
+    per-step list; steps before any clean point are 0.0. With output
+    constraints, only records that satisfy them count.
+    """
+    if spec.n_objectives < 2:
+        raise ValueError("hypervolume needs >=2 objectives; use convergence for 1.")
+    from botorch.utils.multi_objective.hypervolume import (
+        Hypervolume,
+        infer_reference_point,
+    )
+    from botorch.utils.multi_objective.pareto import is_non_dominated
+
+    clean = [r for r in history.records if _counts(r, spec)]
+    if not clean:
+        return [0.0] * len(history.records)
+    # Objective columns in canonical space, one row per clean record (same rule
+    # everywhere via spec.objective_view).
+    all_Y = spec.objective_view(torch.stack([r.outputs.reshape(-1) for r in clean]))
+    ref = infer_reference_point(all_Y[is_non_dominated(all_Y)])
+    hv = Hypervolume(ref_point=ref)
+
+    out: list[float] = []
+    acc: list[torch.Tensor] = []
+    for r in history.records:
+        if _counts(r, spec):
+            acc.append(spec.objective_view(r.outputs.reshape(1, -1))[0])
+        if acc:
+            Y = torch.stack(acc)
+            front = Y[is_non_dominated(Y)]
+            try:
+                out.append(float(hv.compute(front)))
+            except Exception:  # noqa: BLE001 - a degenerate front -> HV 0, keep the series going
+                out.append(0.0)
+        else:
+            out.append(0.0)
+    return out
+
+
+def recommended_set(history: History, spec: ProblemSpec, model=None) -> list[int]:
+    """Current recommended handover indices.
+
+    If a fitted ``model`` is given, uses the denoised posterior-mean recommendation
+    (via the policy path is the caller's job); here we expose the model-free
+    observed non-dominated set, which is also the loop's fallback.
+    """
+    return observed_front(history, spec)
+
+
+# --- model visualization ------------------------------------------------------
+# The surrogate is fit in maximize-canonical space (see session.canonical_Y), so
+# these un-apply the objective signs to report posterior values in RAW units.
+
+
+def _posterior_raw(model, X: torch.Tensor, spec: ProblemSpec):
+    """Posterior (mean, std) at X over ALL outputs, converted back to RAW units.
+
+    Returns tensors shaped (n, n_outputs). The sign flip maps canonical -> raw for
+    the mean (OBSERVE columns pass through); std is sign-invariant.
+    """
+    with torch.no_grad():
+        post = model.posterior(X.double())
+        mean = spec.to_raw(post.mean.reshape(X.shape[0], -1))  # canonical -> raw
+        std = post.variance.clamp_min(1e-12).sqrt().reshape(X.shape[0], -1)
+    return mean, std
+
+
+def parity(history: History, spec: ProblemSpec, model) -> dict:
+    """Predicted (posterior mean ± std) vs. observed, per OUTPUT — dimension-agnostic.
+
+    Returns ``{output_name: {"observed": [...], "predicted": [...], "std": [...]}}``
+    for every modeled output (incl. OBSERVE). The go-to "is the model any good?"
+    view when the design space is too high-D to plot.
+    """
+    clean = history.clean
+    if not len(clean):
+        return {}
+    X = clean.designs.double()
+    Y = clean.outputs.double()  # raw observed, full width
+    mean, std = _posterior_raw(model, X, spec)
+    out = {}
+    for j, o in enumerate(spec.outputs):
+        out[o.name] = {
+            "observed": Y[:, j].tolist(),
+            "predicted": mean[:, j].tolist(),
+            "std": std[:, j].tolist(),
+        }
+    return out
+
+
+def _anchor_point(history: History, spec: ProblemSpec) -> torch.Tensor:
+    """The design to hold the non-varied dimensions at: the recommended/best point."""
+    rec = observed_front(history, spec)
+    idx = rec[0] if rec else 0
+    return history.records[idx].candidate.reshape(-1).double()
+
+
+def slice_1d(
+    history: History,
+    spec: ProblemSpec,
+    model,
+    var: int = 0,
+    output: int = 0,
+    n: int = 80,
+) -> dict:
+    """1-D posterior slice: vary one variable over its range, others fixed at the
+    recommended point. ``output`` indexes ``spec.outputs``. Returns grid ``x``, raw
+    ``mean``, and ``std`` for a band.
+    """
+    bounds = spec.bounds_tensor()
+    lo, hi = float(bounds[0, var]), float(bounds[1, var])
+    anchor = _anchor_point(history, spec)
+    xs = torch.linspace(lo, hi, n, dtype=torch.double)
+    grid = anchor.repeat(n, 1)
+    grid[:, var] = xs
+    mean, std = _posterior_raw(model, grid, spec)
+    return {
+        "x": xs.tolist(),
+        "mean": mean[:, output].tolist(),
+        "std": std[:, output].tolist(),
+        "var_name": spec.variable_names[var],
+        "objective_name": spec.outputs[output].name,
+    }
+
+
+def slice_2d(
+    history: History,
+    spec: ProblemSpec,
+    model,
+    var_x: int = 0,
+    var_y: int = 1,
+    output: int = 0,
+    n: int = 40,
+) -> dict:
+    """2-D posterior-mean slice: vary two variables over a grid, others fixed at the
+    recommended point. ``output`` indexes ``spec.outputs``. Returns ``x``, ``y``
+    axes and a ``z`` matrix (raw mean).
+    """
+    bounds = spec.bounds_tensor()
+    xs = torch.linspace(float(bounds[0, var_x]), float(bounds[1, var_x]), n, dtype=torch.double)
+    ys = torch.linspace(float(bounds[0, var_y]), float(bounds[1, var_y]), n, dtype=torch.double)
+    anchor = _anchor_point(history, spec)
+    gx, gy = torch.meshgrid(xs, ys, indexing="xy")
+    grid = anchor.repeat(n * n, 1)
+    grid[:, var_x] = gx.reshape(-1)
+    grid[:, var_y] = gy.reshape(-1)
+    mean, _ = _posterior_raw(model, grid, spec)
+    z = mean[:, output].reshape(n, n)
+    return {
+        "x": xs.tolist(),
+        "y": ys.tolist(),
+        "z": z.tolist(),
+        "var_x_name": spec.variable_names[var_x],
+        "var_y_name": spec.variable_names[var_y],
+        "objective_name": spec.outputs[output].name,
+    }
+
+
+def bax_target_region(
+    policy: BAXPolicy, model, spec: ProblemSpec
+) -> dict[str, int]:
+    """Snapshot of the BAX discovered-target region on the current posterior mean.
+
+    Returns ``{"identified": int, "grid_size": int}``: how many candidate-grid
+    points the policy's goal selects from the posterior-MEAN prediction (the size
+    of the region found so far — should grow then stabilize), and the grid size.
+    When the policy applies output constraints, ``identified`` counts what the goal
+    selects among points predicted to satisfy them, and ``"before_constraints"``
+    what it selects from every point, so the narrowing is visible.
+    Pure function of ``(policy, model, spec)``; encapsulates grid + raw-posterior +
+    ``goal.identify_subspace`` so no BAX internals leak to the caller.
+    """
+    grid = policy.candidate_grid(spec)
+    mean, _ = _posterior_raw(model, grid, spec)
+    out = {
+        "identified": len(policy._identify(spec, mean, grid)),
+        "grid_size": int(grid.shape[0]),
+    }
+    if policy._constrains(spec):
+        out["before_constraints"] = len(policy._goal_for(spec).identify_subspace(mean, grid))
+    return out
+
+
+__all__ = [
+    "bax_target_region",
+    "canonical_Y",
+    "convergence",
+    "hypervolume",
+    "parity",
+    "recommended_set",
+    "slice_1d",
+    "slice_2d",
+]