| """Aggregate actual trials into Figure 1 style plots and the flow-rate table.""" |
| import csv |
| import json |
| from pathlib import Path |
|
|
| from common import case_guard, load_config, load_data, parse_args, save_json, write_csv |
| import numpy as np |
| import torch |
| import matplotlib |
| matplotlib.use('Agg') |
| import matplotlib.pyplot as plt |
|
|
|
|
| def read_trajectory(path, budget): |
| with open(path) as f: |
| rows = list(csv.DictReader(f)) |
| if len(rows) != budget or [int(r['iteration']) for r in rows] != list(range(1, budget + 1)): |
| raise RuntimeError(f'Incomplete or nonsequential trial: {path}') |
| regrets = np.array([float(r['regret']) for r in rows]) |
| if not np.isfinite(regrets).all() or (regrets < -1e-7).any(): |
| raise ValueError(f'Invalid regret: {path}') |
| cumulative = regrets.cumsum() |
| if not np.allclose(cumulative, [float(r['cumulative_regret']) for r in rows]): |
| raise ValueError('Saved cumulative regret does not match sum of instant regrets') |
| return regrets |
|
|
|
|
| def save_fig(fig, out, name): |
| for ext in ('png', 'pdf'): |
| fig.savefig(out / f'{name}.{ext}', dpi=300, bbox_inches='tight') |
| plt.close(fig) |
|
|
|
|
| def main(): |
| args = parse_args(__doc__).parse_args() |
| cfg = load_config(args.config) |
| case_guard(cfg) |
| out = Path(cfg['paths']['output_dir']) / 'results' |
| n = cfg['experiment']['iterations'] |
| seeds = cfg['experiment']['seeds'] |
| times = np.arange(1, n + 1) |
| actual = {'NOTS': np.stack([read_trajectory(out / f'trial_{s}/trajectory.csv', n) for s in seeds])} |
| if cfg['experiment']['run_random_baseline']: |
| actual['RS'] = np.stack([read_trajectory(out / f'rs_{s}.csv', n) for s in seeds]) |
| statistics = {} |
| for method, regrets in actual.items(): |
| cum = regrets.cumsum(axis=1) |
| avg = cum / times |
| finals = avg[:, -1] |
| statistics[method] = {'final_average_regret_mean': float(finals.mean()), |
| 'final_average_regret_std': float(finals.std(ddof=0)), |
| 'final_average_regret_sample_std': float(finals.std(ddof=1)), |
| 'per_trial_final_average_regret': finals.tolist(), |
| 'trials': len(seeds), 'iterations': n, 'provenance': 'measured'} |
| write_csv(out / f'{method.lower()}_curves.csv', [ |
| {'iteration': int(t), 'cumulative_mean': float(cum[:, i].mean()), |
| 'cumulative_std': float(cum[:, i].std()), 'average_mean': float(avg[:, i].mean()), |
| 'average_std': float(avg[:, i].std())} for i, t in enumerate(times)]) |
| save_json(out / 'metrics.json', {'metric_definition': 'mean_trials(sum_t(true_max - true_selected_t)/300)', |
| 'std_ddof': 0, 'seeds': seeds, 'methods': statistics, |
| 'published_references': cfg['reference'], |
| 'unrerun_baselines': ['BO', 'BFO', 'GP-TS', 'STO-NTS']}) |
| rows = [] |
| for method, ref in cfg['reference'].items(): |
| value = statistics.get(method) |
| rows.append({'method': method, 'measured_mean': value['final_average_regret_mean'] if value else '', |
| 'measured_std': value['final_average_regret_std'] if value else '', |
| 'paper_mean': ref['mean'], 'paper_std': ref['std'], |
| 'status': 'measured' if value else 'published_reference_only'}) |
| write_csv(out / 'table1_flow_rates.csv', rows) |
| plt.rcParams.update({'font.size': 10, 'axes.spines.top': False, 'axes.spines.right': False}) |
| fig = plt.figure(figsize=(12, 7.2), constrained_layout=True) |
| grid = fig.add_gridspec(2, 4, height_ratios=[1.1, 1]) |
| ax0, ax1 = fig.add_subplot(grid[0, :2]), fig.add_subplot(grid[0, 2:]) |
| for method, regrets in actual.items(): |
| cumulative = regrets.cumsum(axis=1) |
| for ax, matrix in [(ax0, cumulative), (ax1, cumulative / times)]: |
| mean, std = matrix.mean(0), matrix.std(0) |
| line, = ax.plot(times, mean, label=f'{method} (measured)', lw=2) |
| ax.fill_between(times, mean - std, mean + std, color=line.get_color(), alpha=0.18) |
| ax0.set_ylabel('Cumulative regret $R_t$') |
| ax1.set_ylabel('Average regret $R_t/t$') |
| for ax in (ax0, ax1): |
| ax.set_xlabel('Oracle query') |
| ax.legend() |
| ax.grid(alpha=0.2) |
| x, y, meta = load_data(cfg) |
| oracle = torch.load(out / 'oracle.pt', map_location='cpu', weights_only=True) |
| best, worst = oracle['best_index'], oracle['worst_index'] |
| |
| panels = [(x[best, 0], 'Best permeability a', 'gray', x.min().item(), x.max().item()), |
| (y[best, 0], 'Best pressure u', 'viridis', y.min().item(), y.max().item()), |
| (x[worst, 0], 'Worst permeability a', 'gray', x.min().item(), x.max().item()), |
| (y[worst, 0], 'Worst pressure u', 'viridis', y.min().item(), y.max().item())] |
| for i, (a, title, cmap, vmin, vmax) in enumerate(panels): |
| ax = fig.add_subplot(grid[1, i]) |
| im = ax.imshow(a.numpy(), origin='lower', extent=[0, 1, 0, 1], cmap=cmap, |
| interpolation='nearest', vmin=vmin, vmax=vmax) |
| ax.set_title(title) |
| ax.set_xlabel('x') |
| ax.set_ylabel('y') |
| fig.colorbar(im, ax=ax, shrink=0.7) |
| fig.suptitle('Darcy negative total flow rates: 10 independent trials, 300 queries\nShading: ±1 trial standard deviation; reference-only methods are omitted') |
| save_fig(fig, out, 'figure1') |
| inf = torch.load(out / 'inference_test.pt', map_location='cpu', weights_only=True) |
| target, pred = inf['target'][0, 0].numpy(), inf['prediction'][0, 0].numpy() |
| fig, axes = plt.subplots(1, 4, figsize=(12, 3), constrained_layout=True) |
| for ax, field, title, cmap in zip(axes, |
| [inf['x'][0, 0].numpy(), target, pred, np.abs(pred - target)], |
| ['Test permeability', 'Oracle pressure', 'Loaded checkpoint pressure', 'Absolute error'], |
| ['gray', 'viridis', 'viridis', 'magma']): |
| kwargs = {'vmin': float(min(target.min(), pred.min())), 'vmax': float(max(target.max(), pred.max()))} if 'pressure' in title else {} |
| im = ax.imshow(field, origin='lower', extent=[0, 1, 0, 1], cmap=cmap, **kwargs) |
| ax.set_title(title) |
| fig.colorbar(im, ax=ax, shrink=0.7) |
| save_fig(fig, out, 'checkpoint_inference') |
| fig, ax = plt.subplots(figsize=(8, 4), constrained_layout=True) |
| for seed in seeds: |
| with (out / f'trial_{seed}/training_loss.jsonl').open() as f: |
| losses = [json.loads(line) for line in f] |
| final_epochs = {r['iteration']: r['mse'] for r in losses if r['epoch'] == cfg['training']['epochs']} |
| ax.plot(list(final_epochs), list(final_epochs.values()), alpha=0.6, label=f'seed {seed}') |
| ax.set(xlabel='Oracle query', ylabel='Final epoch training pressure MSE', yscale='log') |
| ax.legend(ncol=2, fontsize=8) |
| save_fig(fig, out, 'training_loss') |
| measured = statistics['NOTS'] |
| report = ('# Darcy flow rates reproduction results\n\n' |
| f'NOTS measured R300/300 = {measured["final_average_regret_mean"]:.8g} ± ' |
| f'{measured["final_average_regret_std"]:.8g} (population std, ten trials).\n\n' |
| 'Paper Table 1 reference: 0.012 ± 0.001. This is a literature value, not measured data.\n\n' |
| 'All measured trajectories contain 300 queries including one random initial query. ' |
| 'Training and acquisition use the same physical pressure scale. ' |
| 'The paper main-text flow sign conflicts with Appendix D.1; the latter is used. ' |
| 'Published BO/BFO/GP-TS/STO-NTS entries are reference-only, not reruns.\n') |
| (out / 'report.md').write_text(report) |
| print('RESULT ' + json.dumps(statistics), flush=True) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
