176 lines
7.7 KiB
Python
176 lines
7.7 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
from contextlib import contextmanager
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import re
|
|
import shutil
|
|
import tempfile
|
|
|
|
import numpy as np
|
|
|
|
|
|
NAME = re.compile(r"^[A-Za-z][A-Za-z0-9_]*$")
|
|
COORDS = {"lat", "lon", "lev", "levp", "time"}
|
|
|
|
|
|
def _atomic_json(path: Path, value: object) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
fd, temporary = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent)
|
|
os.fchmod(fd, 0o640)
|
|
try:
|
|
with os.fdopen(fd, "w", encoding="utf-8") as stream:
|
|
json.dump(value, stream, separators=(",", ":"), allow_nan=False)
|
|
stream.flush()
|
|
os.fsync(stream.fileno())
|
|
os.replace(temporary, path)
|
|
finally:
|
|
if os.path.exists(temporary):
|
|
os.unlink(temporary)
|
|
|
|
|
|
def _metadata(path: Path) -> dict:
|
|
result = {}
|
|
sidecar = path.with_name(path.stem + "_metadata.npz")
|
|
if not sidecar.is_file():
|
|
return result
|
|
with np.load(sidecar, allow_pickle=False) as archive:
|
|
for name in archive.files:
|
|
row = archive[name]
|
|
if row.ndim != 1 or row.size < 3:
|
|
continue
|
|
result[name] = (str(row[1]), str(row[2]))
|
|
return result
|
|
|
|
|
|
def latitude_weights(lat):
|
|
nodes, weights = np.polynomial.legendre.leggauss(len(lat))
|
|
order = np.argsort(lat)
|
|
sine = np.sin(np.deg2rad(lat[order]))
|
|
if np.allclose(sine, nodes, atol=1e-6, rtol=0):
|
|
result = np.empty_like(weights)
|
|
result[order] = weights
|
|
return result, 'Gaussian quadrature'
|
|
edges = np.concatenate(([-1.0], (sine[:-1] + sine[1:]) / 2, [1.0]))
|
|
result = np.empty_like(sine)
|
|
result[order] = np.diff(edges)
|
|
return result, 'Latitude cell area'
|
|
|
|
|
|
def _header(archive, name):
|
|
with archive.zip.open(name + '.npy') as stream:
|
|
version = np.lib.format.read_magic(stream)
|
|
reader = np.lib.format.read_array_header_1_0 if version == (1, 0) else np.lib.format.read_array_header_2_0
|
|
return reader(stream)
|
|
|
|
|
|
@contextmanager
|
|
def _field(archive, name, destination, memory_limit):
|
|
if archive.zip.getinfo(name + '.npy').file_size <= memory_limit:
|
|
yield archive[name]
|
|
return
|
|
fd, temporary = tempfile.mkstemp(prefix='.field-', suffix='.npy', dir=destination)
|
|
array = None
|
|
try:
|
|
with os.fdopen(fd, 'wb') as output, archive.zip.open(name + '.npy') as source:
|
|
shutil.copyfileobj(source, output, length=1024 * 1024)
|
|
array = np.load(temporary, mmap_mode='r', allow_pickle=False)
|
|
yield array
|
|
finally:
|
|
if array is not None:
|
|
del array
|
|
os.unlink(temporary)
|
|
|
|
|
|
def prepare_results(npz_path, destination, dataset_id, config=None):
|
|
if config is None:
|
|
if __package__:
|
|
from .config import load_config
|
|
else:
|
|
from config import load_config
|
|
config = load_config()
|
|
source, dest = Path(npz_path), Path(destination)
|
|
if not re.fullmatch(r"[a-f0-9]{32}", str(dataset_id)):
|
|
raise ValueError("dataset_id must contain 32 lowercase hexadecimal characters")
|
|
dest.mkdir(parents=True, exist_ok=True)
|
|
with np.load(source, allow_pickle=False) as archive:
|
|
required = ("lat", "lon", "time")
|
|
if any(key not in archive for key in required):
|
|
raise ValueError("NPZ must contain lat, lon, and time coordinates")
|
|
lat = np.asarray(archive["lat"], dtype=np.float64)
|
|
lon = np.asarray(archive["lon"], dtype=np.float64)
|
|
time = np.asarray(archive["time"], dtype=np.float64)
|
|
levels = np.asarray(archive["lev"], dtype=np.float64).tolist() if "lev" in archive else [0.0]
|
|
if any(a.ndim != 1 or not a.size or not np.isfinite(a).all() for a in (lat, lon, time)):
|
|
raise ValueError("coordinates must be non-empty finite vectors")
|
|
metadata = _metadata(source)
|
|
time_unit = metadata.get("time", ("", ""))[1] or ""
|
|
fields, total_output = [], 0
|
|
weights, weighting = latitude_weights(lat)
|
|
weights = weights[:, None]
|
|
for name in archive.files:
|
|
if name in COORDS or not NAME.fullmatch(name):
|
|
continue
|
|
shape, _, dtype = _header(archive, name)
|
|
if dtype.kind not in 'fiu':
|
|
continue
|
|
if len(shape) not in (3, 4) or shape[0] != len(time) or shape[-2:] != (len(lat), len(lon)):
|
|
continue
|
|
if len(shape) == 3:
|
|
field_shape = (shape[0], 1, shape[1], shape[2])
|
|
else:
|
|
field_shape = shape
|
|
if shape[1] != len(levels):
|
|
continue
|
|
byte_count = int(np.prod(field_shape, dtype=np.int64)) * 4
|
|
total_output += byte_count
|
|
if total_output > config['max_result_bytes']:
|
|
raise ValueError("Result output exceeds the configured storage limit.")
|
|
required_space = byte_count + archive.zip.getinfo(name + '.npy').file_size
|
|
if shutil.disk_usage(dest).free - required_space < config['minimum_free_bytes']:
|
|
raise ValueError("Not enough free space to prepare model output.")
|
|
label, unit = metadata.get(name, (name, ""))
|
|
minimum, maximum, means = [], [], []
|
|
fd, temporary = tempfile.mkstemp(prefix=f".{name}.", suffix=".tmp", dir=dest)
|
|
os.fchmod(fd, 0o640)
|
|
try:
|
|
with os.fdopen(fd, "wb") as stream, _field(archive, name, dest, config['max_field_memory_bytes']) as array:
|
|
for ti in range(shape[0]):
|
|
frame = array[ti] if len(shape) == 4 else array[ti][None, ...]
|
|
for li in range(field_shape[1]):
|
|
plane = np.array(frame[li], dtype=np.float32, copy=True)
|
|
np.nan_to_num(plane, copy=False, nan=np.nan, posinf=np.nan, neginf=np.nan)
|
|
finite = np.isfinite(plane)
|
|
minimum.append(float(np.min(plane[finite])) if finite.any() else None)
|
|
maximum.append(float(np.max(plane[finite])) if finite.any() else None)
|
|
weighted = np.where(finite, plane * weights, 0.0)
|
|
denominator = np.sum(np.where(finite, weights, 0.0), dtype=np.float64)
|
|
means.append(float(np.sum(weighted, dtype=np.float64) / denominator) if denominator else None)
|
|
stream.write(np.asarray(plane, dtype="<f4", order="C").tobytes(order="C"))
|
|
stream.flush()
|
|
os.fsync(stream.fileno())
|
|
os.replace(temporary, dest / f"{name}.f32")
|
|
finally:
|
|
if os.path.exists(temporary):
|
|
os.unlink(temporary)
|
|
_atomic_json(dest / f"{name}.json", {"mean": np.asarray(means, dtype=object).reshape(field_shape[:2]).tolist(), "min": np.asarray(minimum, dtype=object).reshape(field_shape[:2]).tolist(), "max": np.asarray(maximum, dtype=object).reshape(field_shape[:2]).tolist()})
|
|
fields.append({"name": name, "label": label, "unit": unit, "shape": list(field_shape), "min": min((x for x in minimum if x is not None), default=None), "max": max((x for x in maximum if x is not None), default=None)})
|
|
index = {"id": str(dataset_id), "lat": lat.tolist(), "lon": lon.tolist(), "time": time.tolist(), "time_unit": time_unit, "levels": levels, "fields": fields, "mean_weighting": weighting}
|
|
_atomic_json(dest / "index.json", index)
|
|
return index
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("npz_path")
|
|
parser.add_argument("destination")
|
|
parser.add_argument("dataset_id")
|
|
args = parser.parse_args()
|
|
prepare_results(args.npz_path, args.destination, args.dataset_id)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|