examples/multi_fidelity_example.py

"""
Example: Multi-Fidelity Bayesian Optimization

Uses the physics model as a cheap low-fidelity source and experimental
measurements as the expensive high-fidelity source. The multi-fidelity GP
learns the correlation between fidelities to transfer knowledge.
"""

import torch
from torch import Tensor

from physics_informed_bo.experiment.parameter_space import ParameterSpace
from physics_informed_bo.models.multi_fidelity import MultiFidelitySurrogate


def physics_model(X: Tensor) -> Tensor:
    """Low-fidelity physics model."""
    x1, x2 = X[:, 0], X[:, 1]
    return torch.sin(x1) * x2 + x1 * 0.5


def true_function(X: Tensor) -> Tensor:
    """High-fidelity ground truth (simulating experiments)."""
    x1, x2 = X[:, 0], X[:, 1]
    return torch.sin(x1) * x2 + x1 * 0.5 + 0.3 * torch.cos(3 * x1 * x2)


def main():
    torch.manual_seed(42)

    # Define space
    space = ParameterSpace()
    space.add_continuous("x1", 0.0, 6.28)
    space.add_continuous("x2", 0.0, 5.0)

    # Multi-fidelity surrogate
    mf_model = MultiFidelitySurrogate(
        physics_fn=physics_model,
        device="cpu",
    )

    # Small set of expensive experiments
    X_exp = torch.tensor([
        [1.0, 2.0], [3.0, 1.0], [5.0, 4.0], [2.0, 3.0], [4.0, 2.5]
    ], dtype=torch.float64)
    y_exp = true_function(X_exp).unsqueeze(-1) + 0.05 * torch.randn(5, 1, dtype=torch.float64)

    # Build multi-fidelity dataset (physics=low, experiments=high)
    X_mf, y_mf = mf_model.build_multi_fidelity_data(
        X_experiment=X_exp,
        y_experiment=y_exp,
        n_physics_points=50,
    )

    print(f"Multi-fidelity dataset: {len(X_mf)} points "
          f"({len(X_mf) - len(X_exp)} physics + {len(X_exp)} experimental)")

    # Fit the model
    mf_model.fit(X_mf, y_mf)

    # Predict at new points (always at high fidelity)
    X_test = torch.tensor([[2.5, 2.5], [4.0, 3.0]], dtype=torch.float64)
    mean, var = mf_model.predict(X_test)

    print("\nPredictions (high fidelity):")
    for i, (m, v) in enumerate(zip(mean, var)):
        true_val = true_function(X_test[i:i+1]).item()
        print(f"  x={X_test[i].tolist()} -> pred={m.item():.3f} ± {v.sqrt().item():.3f} "
              f"(true={true_val:.3f})")


if __name__ == "__main__":
    main()