Adaptive sampling strategies from uncertainty-aware reduced-order models can improve the efficiency of inexpensive label generation in amortized optimization pipelines.

Uncertainty-aware reduced-order models (ROM), when used to guide adaptive sampling of candidate solutions, will reduce the total number of expensive label evaluations required to achieve equivalent or superior optimization performance compared to uniform/random sampling baselines in amortized optimization pipelines. Specifically: given a fixed label budget B, an adaptive ROM-guided sampler will produce a trained amortized optimizer with ≥10% better objective value or ≥20% reduction in labels needed to reach a target objective threshold, across at least two distinct optimization domains.

Minimum Viable Test (MVT): Compare ROM-guided adaptive sampling vs. random sampling vs. simple uncertainty sampling in two domains — (A) a physics-based PDE optimization task (e.g., airfoil shape optimization using a ROM of CFD solver) and (B) a molecular property optimization task. Measure label efficiency curves (objective quality vs. number of labels consumed) and final amortized optimizer performance on held-out instances. Use 5 independent random seeds per condition. Total conditions: 3 samplers × 2 domains × 5 seeds = 30 runs.

# === ROM-Guided Adaptive Sampling for Amortized Optimization ===
# Architecture Overview

# --- Core Components ---

class ReducedOrderModel:
    """POD-based or GP-based ROM with uncertainty quantification."""
    def __init__(self, model_type='gp_ensemble', n_ensemble=5):
        self.model_type = model_type
        self.models = [GaussianProcessRegressor(
            kernel=Matern(nu=2.5) + WhiteKernel(),
            normalize_y=True
        ) for _ in range(n_ensemble)]
        self.is_fitted = False

    def fit(self, X: np.ndarray, y: np.ndarray):
        """Fit ensemble of GPs on collected (X, y) pairs."""
        # X: [N, D] input parameters, y: [N] objective values
        for i, model in enumerate(self.models):
            # Bootstrap resample for ensemble diversity
            idx = np.random.choice(len(X), len(X), replace=True)
            model.fit(X[idx], y[idx])
        self.is_fitted = True

    def predict_with_uncertainty(self, X_query: np.ndarray):
        """Returns mean and std of predictions."""
        preds = np.array([m.predict(X_query) for m in self.models])
        mu = preds.mean(axis=0)      # [N_query]
        sigma = preds.std(axis=0)    # [N_query]
        return mu, sigma

    def calibration_error(self, X_val, y_val):
        """Expected Calibration Error on validation set."""
        mu, sigma = self.predict_with_uncertainty(X_val)
        # Compute ECE via reliability diagram binning
        z_scores = np.abs(y_val - mu) / (sigma + 1e-8)
        ece = compute_ece(z_scores)  # standard ECE computation
        return ece


class AdaptiveSampler:
    """ROM-guided acquisition function for label generation."""
    def __init__(self, rom: ReducedOrderModel, beta: float = 1.0,
                 strategy: str = 'ucb'):
        self.rom = rom
        self.beta = beta  # exploration-exploitation tradeoff
        self.strategy = strategy

    def acquisition(self, X_candidates: np.ndarray) -> np.ndarray:
        """Score candidates; higher = more valuable to label."""
        mu, sigma = self.rom.predict_with_uncertainty(X_candidates)
        if self.strategy == 'ucb':
            return mu + self.beta * sigma
        elif self.strategy == 'ei':
            # Expected Improvement over current best
            y_best = self.current_best
            z = (mu - y_best) / (sigma + 1e-8)
            return (mu - y_best) * norm.cdf(z) + sigma * norm.pdf(z)
        elif self.strategy == 'random':
            return np.random.rand(len(X_candidates))

    def select_batch(self, X_pool: np.ndarray,
                     batch_size: int = 10) -> np.ndarray:
        """Select top-k candidates by acquisition score."""
        scores = self.acquisition(X_pool)
        top_k_idx = np.argsort(scores)[-batch_size:]
        return X_pool[top_k_idx], top_k_idx


class LabelGenerator:
    """Wraps the inexpensive label generation process."""
    def __init__(self, domain: str = 'airfoil'):
        self.domain = domain

    def generate(self, X: np.ndarray) -> np.ndarray:
        """Generate labels for input parameters X."""
        if self.domain == 'airfoil':
            return xfoil_evaluate(X)      # ~1-5s per sample
        elif self.domain == 'molecular':
            return xtb_evaluate(X)        # ~10-30s per sample
        elif self.domain == 'synthetic':
            return hartmann6(X)           # <1ms per sample


class AmortizedOptimizer(nn.Module):
    """MLP-based amortized optimizer trained on labeled dataset."""
    def __init__(self, input_dim: int, output_dim: int,
                 hidden_dim: int = 256, n_layers: int = 3):
        super().__init__()
        layers = [nn.Linear(input_dim, hidden_dim), nn.ReLU()]
        for _ in range(n_layers - 1):
            layers += [nn.Linear(hidden_dim, hidden_dim), nn.ReLU()]
        layers.append(nn.Linear(hidden_dim, output_dim))
        self.net = nn.Sequential(*layers)

    def forward(self, problem_context: torch.Tensor) -> torch.Tensor:
        return self.net(problem_context)


# --- Main Experiment Loop ---

def run_experiment(domain: str, sampler_type: str,
                   total_budget: int = 1000,
                   batch_size: int = 10,
                   seed: int = 42) -> dict:

    np.random.seed(seed)
    torch.manual_seed(seed)

    label_gen = LabelGenerator(domain=domain)
    rom = ReducedOrderModel(model_type='gp_ensemble', n_ensemble=5)
    sampler = AdaptiveSampler(rom, beta=1.0, strategy=sampler_type)

    # Step 1: Initial random exploration (warm-start)
    n_initial = max(50, total_budget // 10)
    X_pool = generate_candidate_pool(domain, n=10000)
    X_init = X_pool[np.random.choice(len(X_pool), n_initial, replace=False)]
    y_init = label_gen.generate(X_init)

    X_collected = X_init.copy()
    y_collected = y_init.copy()
    labels_used = [n_initial]
    best_obj_curve = [y_init.max()]

    # Step 2: Adaptive sampling loop
    remaining_budget = total_budget - n_initial
    n_rounds = remaining_budget // batch_size

    for round_idx in range(n_rounds):
        # Fit ROM on collected data
        if sampler_type != 'random':
            rom.fit(X_collected, y_collected)
            # Validate ROM quality
            if round_idx == 0:
                ece = rom.calibration_error(X_val, y_val)
                assert ece < 0.25, f"ROM poorly calibrated: ECE={ece:.3f}"

        # Select next batch
        X_remaining = set_difference(X_pool, X_collected)
        X_batch, batch_idx = sampler.select_batch(X_remaining, batch_size)

        # Generate labels
        y_batch = label_gen.generate(X_batch)

        # Update collected dataset
        X_collected = np.vstack([X_collected, X_batch])
        y_collected = np.concatenate([y_collected, y_batch])
        labels_used.append(len(X_collected))
        best_obj_curve.append(y_collected.max())

    # Step 3: Train amortized optimizer on collected labels
    amortized_opt = AmortizedOptimizer(
        input_dim=X_collected.shape[1],
        output_dim=1
    )
    train_amortized_optimizer(amortized_opt, X_collected, y_collected,
                               epochs=200, patience=20)

    # Step 4: Evaluate on held-out test instances
    test_performance = evaluate_on_test_set(amortized_opt, X_test, y_test)

    return {
        'labels_used': labels_used,
        'best_obj_curve': best_obj_curve,
        'test_mean_obj': test_performance['mean'],
        'test_top10_obj': test_performance['top10'],
        'rom_ece': ece if sampler_type != 'random' else None,
        'total_wall_time': time.time() - t_start
    }


# --- Evaluation Metrics ---

def label_efficiency_ratio(curve_adaptive, curve_random,
                            threshold_pct=0.90) -> float:
    """
    Compute ratio of labels needed to reach threshold_pct of
    random sampling's final performance.
    Returns: labels_random_needed / labels_adaptive_needed
    Values > 1.0 indicate adaptive is more efficient.
    """
    final_random_perf = curve_random[-1]
    target = threshold_pct * final_random_perf

    labels_random = find_first_crossing(curve_random, target)
    labels_adaptive = find_first_crossing(curve_adaptive, target)

    return labels_random / (labels_adaptive + 1e-8)


# --- Statistical Testing ---

def statistical_analysis(results_adaptive, results_random):
    efficiency_ratios = [
        label_efficiency_ratio(r_a['best_obj_curve'], r_r['best_obj_curve'])
        for r_a, r_r in zip(results_adaptive, results_random)
    ]
    t_stat, p_value = scipy.stats.ttest_1samp(efficiency_ratios, popmean=1.0)
    cohens_d = (np.mean(efficiency_ratios) - 1.0) / np.std(efficiency_ratios)
    return {
        'mean_efficiency_ratio': np.mean(efficiency_ratios),
        'p_value': p_value,
        'cohens_d': cohens_d,
        'ci_95': scipy.stats.t.interval(0.95, len(efficiency_ratios)-1,
                                         loc=np.mean(efficiency_ratios),
                                         scale=scipy.stats.sem(efficiency_ratios))
    }