File size: 4,013 Bytes
3c8d50e | 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 | """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))
|