File size: 7,877 Bytes
5978445 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | """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']
# Common scales for best/worst pressure, original permeability for display.
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()
|