Adaptive sampling in parametrized dynamical systems can improve the efficiency of surrogate-based amortized optimization when problem parameters vary continuously.

Adaptive sampling strategies applied to the parameter space of parametrized dynamical systems will reduce the total number of high-fidelity simulator evaluations required to achieve a fixed surrogate model accuracy (measured by normalized RMSE ≤ 0.05) compared to uniform/random sampling baselines, when used within a surrogate-based amortized optimization pipeline and when problem parameters vary continuously over a compact domain. Specifically, adaptive sampling will achieve equivalent optimization quality (within 2% of optimal objective value) using ≤ 50% of the simulator calls required by uniform sampling.

Minimum Viable Test (MVT): Compare adaptive sampling (uncertainty-guided, e.g., active learning with GP or ensemble) vs. uniform random sampling vs. Latin Hypercube Sampling (LHS) as surrogate training data acquisition strategies on 3 parametrized dynamical systems of increasing complexity. Evaluate surrogate accuracy and downstream optimization quality as a function of simulator call budget. Run 10 independent trials per condition. Primary metric: simulator calls to reach RMSE ≤ 0.05 and optimization gap ≤ 2%.

Full Validation: Extend to 6 benchmark systems, 5 adaptive strategies, ablation over parameter space dimensionality (d = 2, 5, 10, 20), and real-world physics simulators (e.g., fluid dynamics, epidemiological models). Include wall-clock time profiling and scalability analysis.

# Adaptive Sampling for Surrogate-Based Amortized Optimization
# Architecture Overview

import numpy as np
from scipy.integrate import solve_ivp
import torch
import gpytorch
from botorch.models import SingleTaskGP
from botorch.acquisition import qExpectedImprovement, MaxVariance
from botorch.fit import fit_gpytorch_mll

# ── 1. PARAMETRIZED DYNAMICAL SYSTEM INTERFACE ──────────────────────────────
class ParametrizedDynamicalSystem:
    """Abstract interface for parametrized ODE/PDE systems."""
    def __init__(self, param_bounds: dict):
        self.param_bounds = param_bounds  # {name: (low, high)}
    
    def simulate(self, params: np.ndarray) -> np.ndarray:
        """Run high-fidelity simulator; returns QoI scalar or vector."""
        raise NotImplementedError
    
    def param_dim(self) -> int:
        return len(self.param_bounds)

class LorenzSystem(ParametrizedDynamicalSystem):
    def simulate(self, params):
        sigma, rho, beta = params
        def lorenz(t, y):
            return [sigma*(y[1]-y[0]), y[0]*(rho-y[2])-y[1], y[0]*y[1]-beta*y[2]]
        sol = solve_ivp(lorenz, [0, 50], [1,1,1], dense_output=False, 
                        t_eval=np.linspace(40,50,100))
        return np.array([np.mean(sol.y[0]**2)])  # QoI: mean energy

# ── 2. SAMPLING STRATEGIES ───────────────────────────────────────────────────
class SamplingStrategy:
    def select_next_query(self, X_observed, y_observed, surrogate, bounds):
        raise NotImplementedError

class UniformRandomSampling(SamplingStrategy):
    def select_next_query(self, X_observed, y_observed, surrogate, bounds):
        d = bounds.shape[1]
        return np.random.uniform(bounds[0], bounds[1], size=(1, d))

class AdaptiveGPSampling(SamplingStrategy):
    """Max-variance (uncertainty) based adaptive sampling."""
    def select_next_query(self, X_observed, y_observed, surrogate, bounds):
        # Optimize acquisition function (max posterior variance)
        from botorch.optim import optimize_acqf
        acq = MaxVariance(surrogate)
        candidate, _ = optimize_acqf(
            acq, bounds=torch.tensor(bounds, dtype=torch.float64),
            q=1, num_restarts=10, raw_samples=512
        )
        return candidate.numpy()

class EnsembleAdaptiveSampling(SamplingStrategy):
    """BALD criterion using deep ensemble disagreement."""
    def select_next_query(self, X_observed, y_observed, surrogate, bounds):
        candidates = np.random.uniform(bounds[0], bounds[1], size=(1000, bounds.shape[1]))
        X_cand = torch.tensor(candidates, dtype=torch.float32)
        preds = torch.stack([m(X_cand) for m in surrogate.members])  # (n_members, N, 1)
        variance = preds.var(dim=0).squeeze()
        best_idx = variance.argmax().item()
        return candidates[best_idx:best_idx+1]

# ── 3. SURROGATE MODEL ───────────────────────────────────────────────────────
class GPSurrogate:
    def __init__(self):
        self.model = None
        self.likelihood = None
    
    def fit(self, X: np.ndarray, y: np.ndarray):
        X_t = torch.tensor(X, dtype=torch.float64)
        y_t = torch.tensor(y.squeeze(), dtype=torch.float64)
        self.likelihood = gpytorch.likelihoods.GaussianLikelihood()
        self.model = SingleTaskGP(X_t, y_t.unsqueeze(-1))
        mll = gpytorch.mlls.ExactMarginalLogLikelihood(self.likelihood, self.model)
        fit_gpytorch_mll(mll)
    
    def predict(self, X: np.ndarray):
        X_t = torch.tensor(X, dtype=torch.float64)
        self.model.eval()
        with torch.no_grad(), gpytorch.settings.fast_pred_var():
            pred = self.model(X_t)
        return pred.mean.numpy(), pred.variance.numpy()
    
    def rmse(self, X_test, y_test):
        y_pred, _ = self.predict(X_test)
        return np.sqrt(np.mean((y_pred - y_test.squeeze())**2))

# ── 4. AMORTIZED OPTIMIZATION ────────────────────────────────────────────────
def amortized_optimize(surrogate, bounds, n_restarts=50):
    """Optimize objective over parameter space using surrogate."""
    from scipy.optimize import minimize
    best_val, best_x = np.inf, None
    for _ in range(n_restarts):
        x0 = np.random.uniform(bounds[0], bounds[1])
        def neg_surrogate(x):
            mu, _ = surrogate.predict(x.reshape(1,-1))
            return mu[0]
        res = minimize(neg_surrogate, x0, method='L-BFGS-B',
                      bounds=list(zip(bounds[0], bounds[1])))
        if res.fun < best_val:
            best_val, best_x = res.fun, res.x
    return best_x, best_val

# ── 5. MAIN EXPERIMENT LOOP ──────────────────────────────────────────────────
def run_experiment(system, strategy, budgets, n_trials=10, n0=100):
    results = {b: {'rmse': [], 'opt_gap': [], 'wall_time': []} for b in budgets}
    
    # Ground truth optimum (expensive, computed once)
    grid = np.random.uniform(system.param_bounds_array[0], 
                             system.param_bounds_array[1], size=(5000, system.param_dim()))
    y_grid = np.array([system.simulate(x) for x in grid])
    true_opt_val = y_grid.min()
    
    # Test set for RMSE evaluation
    X_test = np.random.uniform(system.param_bounds_array[0],
                               system.param_bounds_array[1], size=(500, system.param_dim()))
    y_test = np.array([system.simulate(x) for x in X_test])
    
    for trial in range(n_trials):
        np.random.seed(trial)
        # Warm-start with LHS (identical across strategies)
        from scipy.stats.qmc import LatinHypercube
        sampler = LatinHypercube(d=system.param_dim())
        X = sampler.random(n=n0)
        X = X * (system.param_bounds_array[1] - system.param_bounds_array[0]) + system.param_bounds_array[0]
        y = np.array([system.simulate(x) for x in X])
        
        surrogate = GPSurrogate()
        
        for budget in sorted(budgets):
            # Adaptive refinement to reach budget
            while len(X) < budget:
                surrogate.fit(X, y)
                x_next = strategy.select_next_query(X, y, surrogate.model, 
                                                    system.param_bounds_array)
                y_next = system.simulate(x_next.squeeze())
                X = np.vstack([X, x_next])
                y = np.append(y, y_next)
            
            surrogate.fit(X, y)
            rmse = surrogate.rmse(X_test, y_test)
            _, opt_val = amortized_optimize(surrogate, system.param_bounds_array)
            opt_gap = abs(opt_val - true_opt_val) / (abs(true_opt_val) + 1e-8) * 100
            
            results[budget]['rmse'].append(rmse)
            results[budget]['opt_gap'].append(opt_gap)
    
    return results

# ── 6. STATISTICAL ANALYSIS ──────────────────────────────────────────────────
def analyze_results(adaptive_results, uniform_results, budgets):
    from scipy.stats import wilcoxon
    for b in budgets:
        rmse_adaptive = adaptive_results[b]['rmse']
        rmse_uniform = uniform_results[b]['rmse']
        stat, p = wilcoxon(rmse_adaptive, rmse_uniform)
        savings = (1 - np.mean(rmse_adaptive)/np.mean(rmse_uniform)) * 100
        print(f"Budget {b}: RMSE savings={savings:.1f}%, p={p:.4f}")
    
    # Find simulator calls to threshold
    for strategy_name, res in [('adaptive', adaptive_results), ('uniform', uniform_results)]:
        for b in budgets:
            if np.mean(res[b]['rmse']) <= 0.05 and np.mean(res[b]['opt_gap']) <= 2.0:
                print(f"{strategy_name}: threshold reached at budget={b}")
                break