| """Visualize independent Burgers predictions, errors and the actual training log.""" |
| import csv |
| import json |
|
|
| import matplotlib |
| matplotlib.use('Agg') |
| import matplotlib.pyplot as plt |
| import numpy as np |
| from common import parser, load_config, resolve, save_json, update_run_manifest |
|
|
|
|
| def main(): |
| args = parser(__doc__).parse_args() |
| cfg = load_config(args.config) |
| folder = resolve(cfg['paths']['results']) |
| data = np.load(folder / 'test_predictions.npz') |
| pred, truth, times, x = (data[k] for k in ['prediction', 'reference', 'times', 'x']) |
| if pred.shape != (200, 11, 33) or not np.isfinite(pred).all(): |
| raise ValueError('Invalid independent test predictions') |
| summary = json.loads((folder / 'test_metrics.json').read_text()) |
| rmse = np.sqrt(np.mean((pred - truth)**2, axis=(1, 2))) |
| plot_dir = folder / 'figures' |
| plot_dir.mkdir(exist_ok=True) |
| dpi = cfg['visualization']['dpi'] |
| files = [] |
|
|
| def save(fig, name): |
| for suffix in ['png', 'pdf']: |
| destination = plot_dir / f'{name}.{suffix}' |
| fig.savefig(destination, dpi=dpi, bbox_inches='tight') |
| files.append(str(destination.relative_to(folder))) |
| plt.close(fig) |
|
|
| selected = list(dict.fromkeys(cfg['visualization']['sample_indices'] + [int(rmse.argmax())])) |
| for i in selected: |
| uid = int(data['sample_ids'][i]) |
| fig, axes = plt.subplots(1, 3, figsize=(12, 3.5), constrained_layout=True) |
| vlim = max(float(np.max(np.abs(truth[i]))), float(np.max(np.abs(pred[i])))) |
| for ax, field, title in zip(axes[:2], [truth[i], pred[i]], ['Reference', 'RONOM']): |
| im = ax.pcolormesh(times, x, field.T, shading='nearest', cmap='RdBu_r', |
| vmin=-vlim, vmax=vlim) |
| ax.set(title=title, xlabel='t', ylabel='x') |
| fig.colorbar(im, ax=axes[:2], label='u(x,t)') |
| im = axes[2].pcolormesh(times, x, np.abs(pred[i] - truth[i]).T, |
| shading='nearest', cmap='magma') |
| axes[2].set(title='Absolute error', xlabel='t', ylabel='x') |
| fig.colorbar(im, ax=axes[2], label='|prediction - reference|') |
| fig.suptitle(f'Independent test case {uid}, trajectory RMSE = {rmse[i]:.3e}') |
| save(fig, f'test_case_{uid}_fields') |
| fig, axes = plt.subplots(1, 3, figsize=(12, 3.2), constrained_layout=True) |
| for ax, j in zip(axes, [0, 5, 10]): |
| ax.plot(x, truth[i, j], color='black', label='Reference') |
| ax.plot(x, pred[i, j], '--o', ms=3, color='tab:blue', label='RONOM') |
| ax.set(title=f't = {times[j]:.1f}', xlabel='x', ylabel='u(x,t)') |
| ax.grid(alpha=.2) |
| axes[0].legend() |
| fig.suptitle(f'Independent test case {uid}') |
| save(fig, f'test_case_{uid}_profiles') |
|
|
| epochs = np.array([], dtype=int) |
| if (folder / 'training_history.csv').exists(): |
| with (folder / 'training_history.csv').open() as f: |
| history = list(csv.DictReader(f)) |
| epochs = np.array([int(r['epoch']) for r in history]) |
| if len(epochs) != 5000 or not np.array_equal(epochs, np.arange(1, 5001)): |
| raise ValueError('Expected one complete formal training history, epochs 1 through 5000') |
| val_rows = [r for r in history if r['validation_mse']] |
| fig, ax = plt.subplots(figsize=(7, 4)) |
| ax.semilogy(epochs, [float(r['train_mse']) for r in history], label='Training MSE') |
| ax.semilogy([int(r['epoch']) for r in val_rows], |
| [float(r['validation_mse']) for r in val_rows], 'o-', label='Validation MSE') |
| ax.axvline(summary['checkpoint_epoch'], color='grey', ls=':', label='Selected epoch') |
| ax.set(xlabel='Epoch', ylabel='MSE', title='Formal 5000-epoch training') |
| ax.legend(); ax.grid(alpha=.2) |
| save(fig, 'training_curve') |
| else: |
| print('Training history is absent in this results directory; generating independent test figures.', flush=True) |
| fig, axes = plt.subplots(1, 2, figsize=(10, 3.6), constrained_layout=True) |
| axes[0].plot(times, summary['rmse_by_time'], 'o-') |
| axes[0].set(xlabel='t', ylabel='RMSE across independent test cases and x', title='Error by time') |
| axes[1].hist(rmse * 1000, bins=25, edgecolor='white') |
| axes[1].set(xlabel='Trajectory RMSE (× 10⁻³)', ylabel='Test cases', title='200 independent trajectories') |
| save(fig, 'test_error_summary') |
| comparison = folder / 'table1a_comparison.csv' |
| with comparison.open('w', newline='') as f: |
| writer = csv.writer(f) |
| writer.writerow(['source', 'split', 'samples', 'rmse_mean', 'rmse_std_ddof1', 'units']) |
| writer.writerow(['paper_Table1a_RONOM', 'paper_test', 200, .00134, .00162, 'raw']) |
| writer.writerow(['this_run', 'independent_test_800_999', 200, |
| summary['trajectory_rmse_mean'], summary['trajectory_rmse_std_ddof1'], 'raw']) |
| save_json(folder / 'visualization_manifest.json', { |
| 'figures': files, 'sample_ids': [int(data['sample_ids'][i]) for i in selected], |
| 'dpi': dpi, 'test_shape': list(pred.shape), 'training_epochs': len(epochs) if len(epochs) else None, |
| 'colormaps': {'field': 'RdBu_r', 'absolute_error': 'magma'}, |
| 'note': 'Validation selected weights. Test figures and metrics are post-training only.'}) |
| update_run_manifest(cfg, status='completed', figure_files=len(files), |
| visualization_manifest='visualization_manifest.json') |
| print(f'VISUALIZATION COMPLETE: {len(files)} figures in {plot_dir}', flush=True) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
