"""Full 5000-update dCDB9 experiment; no reduced-training mode.""" import csv import os from pathlib import Path import random import sys import time import numpy as np import torch sys.path.insert(0, str(Path(__file__).resolve().parents[1])) from model import PINNShallowWater from common import (arguments, checkpoint_save, device_for, errors, load_config, output_path, runtime_info, save_json, sha256) from fake_data import generate def check_formal_config(config): required = [config['experiment'] == 'dCDB9', config['model']['hidden_layers'] == 5, config['model']['width'] == 300, config['training']['iterations'] == 5000, config['training']['learning_rate'] == 0.001, config['training']['scheduler'] is None, config['training']['early_stopping'] is False, [config['sampling'][k] for k in ('interior', 'boundary', 'initial')] == [32768, 8192, 8192], config['loss']['pde'] == [10, 1, 1], config['loss']['initial'] == [100, 100, 100], config['loss']['positive_depth'] == config['loss']['dry_velocity'] == 10, config['loss']['entropy'] == 0, config['physics']['t_range'] == [0, 0.8]] if not all(required): raise ValueError('Configuration violates the formal paper/user dCDB9 experiment.') def residual_terms(model, x, gravity, create_graph): state, residual = model.residuals(x, gravity, create_graph=create_graph) mask = (state[:, 0] < 0).to(state.dtype) return torch.cat((residual.square().mean(dim=0), (state.square() * mask[:, None]).mean(dim=0))) def boundary_initial_terms(model, data): ni, nb = len(data['initial']), len(data['boundary']) # One graph for the IC and the original + partner BC coordinates. values = model(torch.cat((data['initial'], data['boundary'], data['boundary_pair']))) initial = (values[:ni] - data['initial_target']).square().mean(dim=0) diff = values[ni:ni + nb] - values[ni + nb:] # Sum of 12 per-constraint MSEs (four sides, three variables). boundary = diff.reshape(4, nb // 4, 3).square().mean(dim=1).sum(dim=0) return initial, boundary def weights(config, device): loss = config['loss'] return (torch.tensor(loss['pde'] + [loss['positive_depth'], loss['dry_velocity'], loss['dry_velocity']], device=device), torch.tensor(loss['initial'], device=device), torch.tensor(loss['boundary'], device=device)) def validation(model, data, config, weight_vectors): wp, wi, wb = weight_vectors terms = torch.zeros(6, device=wp.device) n = len(data['interior']) chunk = min(4096, config['training']['pde_chunk_size']) with torch.enable_grad(): for start in range(0, n, chunk): batch = data['interior'][start:start + chunk].detach().requires_grad_(True) terms += residual_terms(model, batch, config['physics']['gravity'], False).detach() * (len(batch) / n) with torch.no_grad(): li, lb = boundary_initial_terms(model, data) recovery = model.slope_recovery() value = terms @ wp + li @ wi + lb @ wb + config['loss']['slope_recovery'] * recovery target = data['initial_target'] diff = model(data['initial']) - target ic_mae = diff.abs().mean(dim=0).cpu().tolist() ic_rmse = diff.square().mean(dim=0).sqrt().cpu().tolist() return {'objective': value.item(), 'pde_mse': terms[:3].cpu().tolist(), 'constraint_mse': terms[3:].cpu().tolist(), 'boundary_mse_sum_per_variable': lb.cpu().tolist(), 'ic_mae': ic_mae, 'ic_rmse': ic_rmse, 'slope_recovery': recovery.item()} def to_device(arrays, prefix, device): return {key: torch.as_tensor(arrays[prefix + key], device=device) for key in ['interior', 'initial', 'initial_target', 'boundary', 'boundary_pair']} def main(): parser = arguments(__doc__) parser.add_argument('--resume', type=Path) parser.add_argument('--check-config', action='store_true') args = parser.parse_args() config, root = load_config(args.config) check_formal_config(config) if args.check_config: print('FORMAL_CONFIG_OK: 5x300, 32768/8192/8192 points, 5000 Adam updates, no decay/early stopping.') return device = device_for(config) torch.set_num_threads(min(int(os.environ.get('SLURM_CPUS_PER_TASK', '8')), 8)) seed = config['seed'] random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) samples_path = output_path(root, config['paths']['samples']) if not samples_path.exists(): generate(config, root) sample_hash = sha256(samples_path) with np.load(samples_path) as arrays: train = to_device(arrays, '', device) val = to_device(arrays, 'val_', device) for key, expected in [('interior', 32768), ('initial', 8192), ('boundary', 8192)]: assert train[key].shape == (expected, 3), (key, train[key].shape) model = PINNShallowWater(**config['model']).to(device) optimizer = torch.optim.Adam(model.parameters(), lr=config['training']['learning_rate'], betas=tuple(config['training']['betas']), eps=config['training']['epsilon'], weight_decay=config['training']['weight_decay']) wp, wi, wb = weight_vectors = weights(config, device) best_path = output_path(root, config['paths']['best_checkpoint']) last_path = output_path(root, config['paths']['last_checkpoint']) result_dir = output_path(root, config['paths']['results']) result_dir.mkdir(parents=True, exist_ok=True) history_path = result_dir / 'training_history.csv' best, best_step, first, prior_seconds = float('inf'), 0, 1, 0.0 hardware = runtime_info(device) if args.resume: checkpoint = torch.load(args.resume, map_location=device, weights_only=False) if checkpoint['config'] != config or checkpoint['sampling_sha256'] != sample_hash: raise ValueError('Resume config or sampling data differs from the saved experiment.') model.load_state_dict(checkpoint['model_state_dict'], strict=True) optimizer.load_state_dict(checkpoint['optimizer_state_dict']) first = checkpoint['iteration'] + 1 best, best_step = checkpoint['best_validation_objective'], checkpoint['best_iteration'] prior_seconds = checkpoint['elapsed_seconds'] torch.set_rng_state(checkpoint['torch_rng'].cpu()) torch.cuda.set_rng_state_all([x.cpu() for x in checkpoint['cuda_rng']]) np.random.set_state(checkpoint['numpy_rng']) random.setstate(checkpoint['python_rng']) if history_path.exists(): with history_path.open() as source: rows = list(csv.DictReader(source)) rows = [r for r in rows if int(r['iteration']) < first] with history_path.open('w', newline='') as dest: writer = csv.DictWriter(dest, fieldnames=list(rows[0]) if rows else fieldnames()) writer.writeheader() writer.writerows(rows) elif best_path.exists() or last_path.exists(): raise FileExistsError('Existing experiment weights found. Use --resume explicitly.') print('FORMAL_TRAINING_START', {'runtime': hardware, 'parameters': sum(p.numel() for p in model.parameters()), 'sampling_sha256': sample_hash, 'iterations': 5000, 'starting_iteration': first, 'samples_every_update': [32768, 8192, 8192], 'external_training_data': False}, flush=True) save_json(result_dir / 'runtime_environment.json', hardware) start_time = time.monotonic() fields = fieldnames() mode = 'a' if args.resume and history_path.exists() else 'w' with history_path.open(mode, newline='', buffering=1) as handle: writer = csv.DictWriter(handle, fieldnames=fields) if mode == 'w': writer.writeheader() last_validation = None for iteration in range(first, config['training']['iterations'] + 1): model.train() optimizer.zero_grad(set_to_none=True) pde_terms = torch.zeros(6, device=device) n = len(train['interior']) chunk = config['training']['pde_chunk_size'] for begin in range(0, n, chunk): points = train['interior'][begin:begin + chunk].detach().requires_grad_(True) part = residual_terms(model, points, config['physics']['gravity'], True) fraction = len(points) / n (fraction * (part @ wp)).backward() pde_terms += fraction * part.detach() initial_loss, boundary_loss = boundary_initial_terms(model, train) recovery = model.slope_recovery() remaining = initial_loss @ wi + boundary_loss @ wb + config['loss']['slope_recovery'] * recovery remaining.backward() loss = pde_terms @ wp + remaining.detach() if not torch.isfinite(loss): raise FloatingPointError(f'Nonfinite loss at iteration {iteration}') optimizer.step() values = torch.cat((loss.reshape(1), pde_terms, initial_loss.detach(), boundary_loss.detach(), recovery.detach().reshape(1))).cpu().tolist() elapsed = time.monotonic() - start_time + prior_seconds row = dict(zip(fields[:15], [iteration] + values)) row.update(learning_rate=optimizer.param_groups[0]['lr'], elapsed_seconds=elapsed) writer.writerow(row) if iteration % config['training']['print_every'] == 0: print('TRAIN', row, flush=True) if iteration == 1 or iteration % config['training']['validation_every'] == 0 or iteration == 5000: model.eval() last_validation = validation(model, val, config, weight_vectors) print('EVAL', {'iteration': iteration, **last_validation}, flush=True) save_json(result_dir / f'validation_latest.json', {'iteration': iteration, **last_validation}) if last_validation['objective'] < best: best, best_step = last_validation['objective'], iteration checkpoint_save(best_path, {'model_class': 'PINNShallowWater', 'model_state_dict': {k: v.detach().cpu() for k, v in model.state_dict().items()}, 'iteration': iteration, 'completed_iterations': iteration, 'config': config, 'validation': last_validation, 'sampling_sha256': sample_hash, 'runtime': hardware}) if iteration % config['training']['checkpoint_every'] == 0 or iteration == 5000: checkpoint_save(last_path, {'model_class': 'PINNShallowWater', 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'iteration': iteration, 'completed_iterations': iteration, 'best_iteration': best_step, 'best_validation_objective': best, 'config': config, 'sampling_sha256': sample_hash, 'runtime': hardware, 'elapsed_seconds': elapsed, 'torch_rng': torch.get_rng_state(), 'cuda_rng': torch.cuda.get_rng_state_all(), 'numpy_rng': np.random.get_state(), 'python_rng': random.getstate()}) save_json(result_dir / 'training_progress.json', {'completed_iterations': iteration, 'required_iterations': 5000, 'best_iteration': best_step, 'best_validation_objective': best, 'elapsed_seconds': elapsed}) handle.flush() completed = torch.load(last_path, map_location='cpu', weights_only=False) assert completed['iteration'] == 5000 best_checkpoint = torch.load(best_path, map_location='cpu', weights_only=False) best_checkpoint['completed_iterations'] = 5000 best_checkpoint['formal_run_complete'] = True checkpoint_save(best_path, best_checkpoint) summary = {'status': 'completed', 'completed_iterations': 5000, 'optimizer_updates': 5000, 'stopped_early': False, 'best_iteration': best_step, 'best_validation_objective': best, 'last_validation': last_validation, 'elapsed_seconds': elapsed, 'runtime': hardware, 'best_checkpoint_sha256': sha256(best_path), 'last_checkpoint_sha256': sha256(last_path), 'sampling_sha256': sample_hash, 'history_rows': 5000, 'external_training_data': False, 'samples_every_update': [32768, 8192, 8192]} save_json(result_dir / 'training_summary.json', summary) print('FORMAL_TRAINING_COMPLETE', summary, flush=True) def fieldnames(): return ['iteration', 'loss', 'mass_mse', 'momentum_x_mse', 'momentum_y_mse', 'positive_depth_mse', 'dry_u_mse', 'dry_v_mse', 'initial_h_mse', 'initial_u_mse', 'initial_v_mse', 'boundary_h_mse', 'boundary_u_mse', 'boundary_v_mse', 'slope_recovery', 'learning_rate', 'elapsed_seconds'] if __name__ == '__main__': main()