File size: 6,751 Bytes
5978445 | 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 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 | """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}
|