| """Plot completed training and held-out bubble radius predictions.""" |
| import csv |
| import json |
| import sys |
| from pathlib import Path |
| sys.path.insert(0,str(Path(__file__).resolve().parents[1])) |
| import numpy as np |
| import matplotlib |
| matplotlib.use('Agg') |
| import matplotlib.pyplot as plt |
| from common import arguments,config,path,write_json |
|
|
|
|
| def run(cfg): |
| root=path(cfg['paths']['results']) |
| a=np.load(root/'predictions.npz',allow_pickle=False) |
| with (root/'history.csv').open() as f:history=list(csv.DictReader(f)) |
| epochs=np.array([int(r['epoch']) for r in history]) |
| plt.rcParams.update({'font.size':10,'axes.spines.top':False,'axes.spines.right':False}) |
| fig,axes=plt.subplots(1,3,figsize=(14,3.8),constrained_layout=True) |
| for key,label in [('train_weighted_loss','vRBA weighted loss'),('train_sampled_mse_normalized','Sampled MSE'),('validation_mse_normalized','Validation MSE')]: |
| axes[0].semilogy(epochs,[float(r[key]) for r in history],label=label,lw=1) |
| axes[0].set(xlabel='Epoch',ylabel='Normalized loss');axes[0].legend() |
| axes[1].semilogy(epochs,[float(r['validation_relative_l2_normalized_batch_mean']) for r in history],label='Validation, normalized') |
| axes[1].semilogy(epochs,[float(r['validation_relative_l2_physical_global']) for r in history],label='Validation, physical') |
| axes[1].set(xlabel='Epoch',ylabel='Relative L2');axes[1].legend() |
| axes[2].plot(epochs,[float(r['learning_rate']) for r in history]);axes[2].set(xlabel='Epoch',ylabel='Learning rate') |
| fig.savefig(root/'training_curves.png',dpi=180);fig.savefig(root/'training_curves.pdf');plt.close(fig) |
| error=a['per_trajectory_relative_l2'];order=np.argsort(error) |
| chosen=[int(order[-1]),int(order[len(order)//2]),int(np.random.default_rng(23).integers(len(order)))] |
| fig,axes=plt.subplots(3,3,figsize=(14,10),constrained_layout=True) |
| for col,idx in enumerate(chosen): |
| tt=a['time']*1e6 |
| for row in [0,1]: |
| axes[row,col].plot(tt,a['target'][idx]*1e6,label='Reference',lw=1.8) |
| axes[row,col].plot(tt,a['prediction'][idx]*1e6,'--',label='vRBA DeepONet',lw=1.2) |
| axes[row,col].set(xlabel='Time (microseconds)',ylabel='Radius (micrometers)') |
| axes[0,col].set_title(f'{["Worst","Median","Seeded random"][col]}: trajectory {a["original_indices"][idx]}\nrelative L2 = {error[idx]:.3e}') |
| axes[1,col].set_xlim(0,80);axes[1,col].set_title('Initial oscillations') |
| axes[0,col].legend() |
| axes[2,col].plot(tt,np.abs(a['prediction'][idx]-a['target'][idx])*1e6,color='tab:red') |
| axes[2,col].set(xlabel='Time (microseconds)',ylabel='Absolute radius error (micrometers)') |
| fig.savefig(root/'bubble_predictions.png',dpi=180);fig.savefig(root/'bubble_predictions.pdf');plt.close(fig) |
| fig,axes=plt.subplots(1,2,figsize=(11,4),constrained_layout=True) |
| im=axes[0].imshow(np.abs(a['prediction']-a['target'])*1e6,aspect='auto',origin='lower',extent=[0,500,900,1000]) |
| axes[0].set(xlabel='Time (microseconds)',ylabel='Original test trajectory index',title='Absolute radius error') |
| fig.colorbar(im,ax=axes[0],label='Micrometers') |
| axes[1].hist(error,bins=20);axes[1].set(xlabel='Per-trajectory relative L2 (physical)',ylabel='Count') |
| fig.savefig(root/'test_error_distribution.png',dpi=180);plt.close(fig) |
| files=['training_curves.png','training_curves.pdf','bubble_predictions.png','bubble_predictions.pdf','test_error_distribution.png'] |
| assert all((root/f).stat().st_size>1000 for f in files) |
| write_json(root/'visualization_manifest.json',dict(files=files,selected_original_indices=a['original_indices'][chosen].tolist(), |
| completed_epochs=int(epochs[-1]),finite_predictions=bool(np.isfinite(a['prediction']).all()), |
| test_samples=int(len(error)),plot_units='microseconds, micrometers')) |
| print('VISUALIZATION_COMPLETE '+json.dumps(dict(directory=str(root),files=files)),flush=True) |
|
|
|
|
| if __name__=='__main__': |
| a=arguments(__doc__).parse_args();run(config(a.config)) |
|
|
