scripts/result.py

"""Evaluate and visualize completed dCDB9 inference, including HLL comparison."""
import csv
import json
from pathlib import Path

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np

from common import arguments, device_for, errors, load_config, output_path, save_json, sha256
from hll_reference import solve


def physical_diagnostics(fields, x, y):
    dx, dy = float(x[1] - x[0]), float(y[1] - y[0])
    output = []
    mass0 = fields[0, :, :, 0].astype(np.float64).sum() * dx * dy
    for field in fields:
        mass = field[:, :, 0].astype(np.float64).sum() * dx * dy
        reflected_x = field[:, ::-1].copy()
        reflected_x[:, :, 1] *= -1
        reflected_y = field[::-1].copy()
        reflected_y[:, :, 2] *= -1
        output.append({'mass_m3': mass, 'relative_mass_change_from_predicted_t0': (mass - mass0) / mass0,
                       'min_depth_m': float(field[:, :, 0].min()),
                       'negative_depth_fraction': float(np.mean(field[:, :, 0] < 0)),
                       'reflection_x_rmse': np.sqrt(np.mean((field - reflected_x) ** 2, axis=(0, 1))).tolist(),
                       'reflection_y_rmse': np.sqrt(np.mean((field - reflected_y) ** 2, axis=(0, 1))).tolist()})
    return output


def figures(results, x, y, times, fields, reference, history):
    extent = [float(x[0]), float(x[-1]), float(y[0]), float(y[-1])]
    names, units = ['h', 'u', 'v'], ['m', 'm/s', 'm/s']
    fig, axes = plt.subplots(5, 3, figsize=(13, 17), constrained_layout=True)
    for row in range(5):
        for j in range(3):
            ax = axes[row, j]
            if row < 3:
                data = fields[row, :, :, j]
                low = float(fields[:, :, :, j].min())
                high = float(fields[:, :, :, j].max())
                if j:
                    high = max(abs(low), abs(high)); low = -high
                image = ax.imshow(data, origin='lower', extent=extent, cmap='viridis' if j == 0 else 'RdBu_r',
                                  vmin=low, vmax=high)
                title = f'PINN {names[j]}, t={times[row]:g} s'
            else:
                k = row - 2
                data = np.abs(fields[k, :, :, j] - reference[k, :, :, j])
                image = ax.imshow(data, origin='lower', extent=extent, cmap='magma', vmin=0)
                title = f'|PINN - HLL| {names[j]}, t={times[k]:g} s'
            ax.set(title=title, xlabel='x [m]', ylabel='y [m]')
            fig.colorbar(image, ax=ax, label=units[j], shrink=0.8)
    fig.suptitle('dCDB9: 5000 full Adam updates; 5 x 300 N-LAAF PINN', fontsize=15)
    fig.savefig(results / 'dCDB9_fields_and_errors.png', dpi=170)
    fig.savefig(results / 'dCDB9_fields_and_errors.pdf')
    plt.close(fig)

    fig, axes = plt.subplots(3, 3, figsize=(13, 12), constrained_layout=True)
    mid = len(y) // 2
    for k, t in enumerate(times):
        for j in range(3):
            ax = axes[k, j]
            ax.plot(x, reference[k, mid, :, j], 'k-', label='HLL (Eq.55 IC at t=0)')
            ax.plot(x, fields[k, mid, :, j], '--', color='#c0392b', label='PINN')
            ax.set(title=f'{names[j]}, t={t:g} s, y={y[mid]:.2f} m', xlabel='x [m]', ylabel=units[j])
            ax.grid(alpha=0.25)
            if k == j == 0:
                ax.legend()
    fig.savefig(results / 'dCDB9_centerline_comparison.png', dpi=170)
    plt.close(fig)

    fig = plt.figure(figsize=(15, 4.5), constrained_layout=True)
    xx, yy = np.meshgrid(x, y)
    stride = 5
    for k, t in enumerate(times):
        ax = fig.add_subplot(1, 3, k + 1, projection='3d')
        ax.plot_surface(xx[::stride, ::stride], yy[::stride, ::stride], fields[k, ::stride, ::stride, 0],
                        cmap='viridis', linewidth=0)
        ax.set(title=f'PINN water depth, t={t:g} s', xlabel='x [m]', ylabel='y [m]', zlabel='h [m]')
    fig.savefig(results / 'dCDB9_water_depth_surfaces.png', dpi=170)
    plt.close(fig)

    fig, ax = plt.subplots(figsize=(9, 5), constrained_layout=True)
    steps = np.array([float(row['iteration']) for row in history])
    for key in ['loss', 'mass_mse', 'momentum_x_mse', 'momentum_y_mse', 'initial_h_mse', 'slope_recovery']:
        ax.semilogy(steps, np.maximum([float(row[key]) for row in history], 1e-20), label=key)
    ax.set(xlabel='Adam iteration', ylabel='Loss / component MSE', title='dCDB9 formal training (all 5000 iterations)')
    ax.legend(fontsize=8)
    ax.grid(alpha=0.25)
    fig.savefig(results / 'training_loss.png', dpi=170)
    plt.close(fig)


def main():
    args = arguments(__doc__).parse_args()
    config, root = load_config(args.config)
    device = device_for(config)
    results = output_path(root, config['paths']['results'])
    training = json.loads((results / 'training_summary.json').read_text())
    inference = json.loads((results / 'inference_summary.json').read_text())
    assert training['completed_iterations'] == 5000 and inference['state_dict_strict_load']
    with np.load(output_path(root, config['paths']['predictions'])) as data:
        x, y, times, fields = [data[k] for k in ['x', 'y', 'times', 'fields']]
    reference_path = output_path(root, config['paths']['hll_reference'])
    reference, hll_info = solve(config, x, y, times, reference_path, device)
    save_json(results / 'hll_manifest.json', hll_info)
    evaluation = []
    for k, t in enumerate(times):
        metric = {'time_s': float(t), 'reference': 'exact Eq.55 initial conditions' if k == 0 else 'independent same-case HLL',
                  **errors(fields[k].reshape(-1, 3), reference[k].reshape(-1, 3))}
        evaluation.append(metric)
        print('METRICS', metric, flush=True)
    metrics = {'status': 'completed', 'variables': ['h', 'u', 'v'], 'units': ['m', 'm/s', 'm/s'],
               'grid_shape': list(fields.shape), 'times': times.tolist(), 'errors': evaluation,
               'pinn_diagnostics': physical_diagnostics(fields, x, y), 'hll': hll_info,
               'paper_table2': {'initial': {'MAE': [0.0059, 0.0083, 0.0064], 'RMSE': [0.032, 0.010, 0.0077]},
                                'final': {'MAE': [0.0056, 0.017, 0.013], 'RMSE': [0.010, 0.028, 0.027]}},
               'training_iterations': 5000, 'checkpoint_iteration': inference['selected_iteration'],
               'checkpoint_sha256': inference['checkpoint_sha256'], 'predictions_sha256': inference['predictions_sha256'],
               'reference_sha256': sha256(reference_path),
               'limitations': ['Single documented seed; paper does not specify seed.',
                    'Boundary type supplemented from official case; paper 0.8 s overrides repository 0.5 s.',
                    'Recovery coefficient 1 is an explicit assumption; not reported in SWE paper.',
                    'HLL is independently implemented; original author benchmark arrays are not used.',
                    'Actual metrics are reported without claiming numerical equality to Table 2.']}
    save_json(results / 'metrics.json', metrics)
    with (results / 'metrics.csv').open('w', newline='') as handle:
        writer = csv.writer(handle)
        writer.writerow(['time_s', 'variable', 'unit', 'MAE', 'RMSE', 'reference'])
        for row in evaluation:
            for j, name in enumerate(metrics['variables']):
                writer.writerow([row['time_s'], name, metrics['units'][j], row['MAE'][j], row['RMSE'][j], row['reference']])
    np.savez_compressed(results / 'absolute_errors.npz', times=times, errors=np.abs(fields - reference))
    with (results / 'training_history.csv').open() as handle:
        history = list(csv.DictReader(handle))
    assert len(history) == 5000 and [int(r['iteration']) for r in history] == list(range(1, 5001))
    figures(results, x, y, times, fields, reference, history)
    rows = ['# dCDB9 正式实验结果', '',
            f'已完成5000次完整Adam更新；best checkpoint选择于第{inference["selected_iteration"]}步。',
            f'运行节点：{training["runtime"]["hostname"]}；设备：{training["runtime"]["device_name"]}；SLURM作业：{training["runtime"]["slurm_job_id"]}。',
            '每次更新使用32768内部点、8192边界点和8192初始点；无外部训练数据。', '',
            '| t (s) | 变量 | MAE | RMSE | 参考 |', '|---|---|---|---|---|']
    for row in evaluation:
        for j, name in enumerate(metrics['variables']):
            rows.append(f'| {row["time_s"]:g} | {name} | {row["MAE"][j]:.6g} | {row["RMSE"][j]:.6g} | {row["reference"]} |')
    rows += ['', '初始时刻使用解析初值；动态参考是同一实验独立实现的一阶HLL，0.04m网格与CFL=0.25，仅用于评估。',
             '论文终时刻h/u/v RMSE分别为0.010/0.028/0.027；本次实际结果如表，训练完成不等于精度达到论文。', '',
             '![fields](dCDB9_fields_and_errors.png)', '![loss](training_loss.png)', '', '实现差异与假设：']
    rows += ['- ' + item for item in metrics['limitations']]
    (results / 'report.md').write_text('\n'.join(rows) + '\n')
    print('RESULT_COMPLETE', {'metrics': str(results / 'metrics.json'),
          'figures': [str(p) for p in results.glob('*.png')], 'diagnostics': metrics['pinn_diagnostics']}, flush=True)


if __name__ == '__main__':
    main()