| """Data, functional, provenance and computation boundary shared by entrypoints.""" |
| import argparse |
| import csv |
| import hashlib |
| import json |
| import os |
| import random |
| import socket |
| import sys |
| from pathlib import Path |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| if (ROOT / '.deps').exists(): |
| sys.path.insert(0, str(ROOT / '.deps')) |
| sys.path.insert(0, str(ROOT)) |
|
|
| import numpy as np |
| import torch |
| import yaml |
|
|
|
|
| def parse_args(description): |
| p = argparse.ArgumentParser(description=description) |
| p.add_argument('--config', type=Path, default=ROOT / 'config/config.yaml') |
| return p |
|
|
|
|
| def load_config(path): |
| with open(path) as f: |
| cfg = yaml.safe_load(f) |
| return cfg |
|
|
|
|
| def save_json(path, value): |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| temporary = path.with_suffix(path.suffix + '.tmp') |
| temporary.write_text(json.dumps(value, indent=2, ensure_ascii=False, allow_nan=False)) |
| temporary.replace(path) |
|
|
|
|
| def save_torch(path, value): |
| path = Path(path) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| temporary = path.with_suffix(path.suffix + '.tmp') |
| torch.save(value, temporary) |
| temporary.replace(path) |
|
|
|
|
|
|
| def model_state_on_cpu(model): |
| """Export PyTorch state while handling neuraloperator 1.0.2 metadata. |
| |
| BaseModel.state_dict inserts an unprefixed dictionary key '_metadata' even |
| when nested in NOTS. That entry is not a parameter/buffer and must not be |
| passed to the enclosing nn.Module's strict loader. The model architecture |
| is saved separately as checkpoint.model_config. OrderedDict._metadata (the |
| PyTorch version information attribute) is preserved. |
| """ |
| state = model.state_dict() |
| metadata = state.pop('_metadata', None) |
| if metadata is not None and not isinstance(metadata, dict): |
| raise TypeError('Unexpected neuraloperator initialization metadata') |
| for name, value in state.items(): |
| if not torch.is_tensor(value): |
| raise TypeError(f'Unexpected non-tensor model state: {name}') |
| state[name] = value.detach().cpu() |
| return state |
|
|
|
|
| def write_csv(path, rows): |
| path = Path(path) |
| if not rows: |
| return |
| temporary = path.with_suffix('.csv.tmp') |
| with temporary.open('w', newline='') as f: |
| w = csv.DictWriter(f, fieldnames=list(rows[0])) |
| w.writeheader() |
| w.writerows(rows) |
| temporary.replace(path) |
|
|
|
|
| def sha256(path): |
| h = hashlib.sha256() |
| with open(path, 'rb') as f: |
| for chunk in iter(lambda: f.read(1024 * 1024), b''): |
| h.update(chunk) |
| return h.hexdigest() |
|
|
|
|
| def config_hash(cfg): |
| return hashlib.sha256(json.dumps(cfg, sort_keys=True).encode()).hexdigest() |
|
|
|
|
| def case_guard(cfg): |
| """Refuse login-node execution and attachment to any pre-existing job.""" |
| if cfg['runtime']['require_case_slurm_job']: |
| actual = os.environ.get('SLURM_JOB_ID', '') |
| registry = Path(cfg['paths']['output_dir']) / '.onescience/jobs.json' |
| if not actual or not registry.is_file(): |
| raise RuntimeError('Compute requires a newly submitted case-owned SLURM job') |
| jobs = json.loads(registry.read_text())['jobs'] |
| if actual not in [str(j['job_id']) for j in jobs]: |
| raise RuntimeError(f'Refusing unregistered job {actual}') |
| if not os.environ.get('SLURM_JOB_NODELIST'): |
| raise RuntimeError('Missing allocated compute node list') |
| torch.set_num_threads(cfg['runtime']['num_threads']) |
| device = torch.device(cfg['runtime']['device']) |
| if device.type == 'cuda' and not torch.cuda.is_available(): |
| raise RuntimeError('Allocated accelerator is unavailable; CPU fallback is disabled') |
| return device |
|
|
|
|
| def seed_all(seed): |
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
| torch.cuda.manual_seed_all(seed) |
|
|
|
|
| def rng_state(): |
| return {'python': random.getstate(), 'numpy': np.random.get_state(), |
| 'torch': torch.get_rng_state(), 'cuda': torch.cuda.get_rng_state_all()} |
|
|
|
|
| def restore_rng(state): |
| random.setstate(state['python']) |
| np.random.set_state(state['numpy']) |
| torch.set_rng_state(state['torch']) |
| torch.cuda.set_rng_state_all(state['cuda']) |
|
|
|
|
| def load_data(cfg, split='train'): |
| path = Path(cfg['paths']['data_dir']) / cfg['paths'][split + '_file'] |
| raw = torch.load(path, map_location='cpu', weights_only=True) |
| out = [] |
| for key in ('x', 'y'): |
| a = raw[key].float() |
| if a.ndim == 3: |
| a = a.unsqueeze(1) |
| if a.ndim != 4 or tuple(a.shape[1:]) != (1, 16, 16): |
| raise ValueError(f'{path}:{key} invalid shape {a.shape}') |
| if not torch.isfinite(a).all(): |
| raise ValueError(f'{path}:{key} contains nonfinite data') |
| out.append(a.contiguous()) |
| if out[0].shape != out[1].shape: |
| raise ValueError('Input and pressure pairing mismatch') |
| if split == 'train' and len(out[0]) != cfg['experiment']['pool_size']: |
| raise ValueError('The required 1000-candidate official pool is unavailable') |
| meta = {'path': str(path), 'sha256': sha256(path), 'n': len(out[0]), |
| 'shape': list(out[0].shape), 'x_range': [out[0].min().item(), out[0].max().item()], |
| 'y_range': [out[1].min().item(), out[1].max().item()], 'dtype': 'float32'} |
| return out[0], out[1], meta |
|
|
|
|
| def negative_total_flow_rates(u, a): |
| """Appendix D.1: integral(a grad(u).n), maximize negative physical outflow. |
| |
| First-order one-sided normal differences and rectangular edge quadrature. |
| Equal 16x16 spacings cancel analytically; corners count once per edge. |
| This deliberately uses original permeability values, not a display mask. |
| """ |
| if u.shape != a.shape or u.ndim != 4: |
| raise ValueError('Functional requires matching [B,1,H,W] tensors') |
| h, w = u.shape[-2:] |
| left = -(a[..., :, 0] * (u[..., :, 1] - u[..., :, 0])).sum(-1) * (w / h) |
| right = (a[..., :, -1] * (u[..., :, -1] - u[..., :, -2])).sum(-1) * (w / h) |
| bottom = -(a[..., 0, :] * (u[..., 1, :] - u[..., 0, :])).sum(-1) * (h / w) |
| top = (a[..., -1, :] * (u[..., -1, :] - u[..., -2, :])).sum(-1) * (h / w) |
| return (left + right + bottom + top).mean(-1) |
|
|
|
|
| @torch.no_grad() |
| def predict(model, xs, batch_size): |
| model.eval() |
| return torch.cat([model(xs[i:i + batch_size]) for i in range(0, len(xs), batch_size)]) |
|
|
|
|
| def hardware_record(): |
| return {'hostname': socket.gethostname(), 'job_id': os.environ.get('SLURM_JOB_ID'), |
| 'node_list': os.environ.get('SLURM_JOB_NODELIST'), 'torch': torch.__version__, |
| 'accelerator': torch.cuda.get_device_name(), 'device_count': torch.cuda.device_count(), |
| 'device_memory_bytes': torch.cuda.get_device_properties(0).total_memory} |
|
|
