File size: 2,556 Bytes
e66b66f | 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 | """Write a lightweight lazy-data protocol and content-determining coefficients."""
import argparse
from pathlib import Path
import numpy as np
import yaml
ROOT = Path(__file__).resolve().parents[1]
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--force", action="store_true")
args = parser.parse_args()
config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
data, generator = config["data"], config["generator"]
output = ROOT / data["root"] / data["protocol_file"]
output.parent.mkdir(parents=True, exist_ok=True)
if output.exists() and not args.force:
print(f"exists={output.relative_to(ROOT)}")
return
rng = np.random.default_rng(int(config["seed"]))
experiments = len(data["experiments"])
growth = float(generator["growth_rate_per_hour"]) * np.asarray(generator["experiment_growth_multipliers"])
np.savez_compressed(
output,
format_version=np.asarray(data["format_version"]),
generation_method=np.asarray("deterministic_structured_lazy_low_rank_v1"),
experiments=np.asarray(data["experiments"]),
times_hours=np.arange(int(data["time_steps"]), dtype=np.int16) * int(data["step_hours"]),
pressure_levels_hpa=np.asarray(data["pressure_levels_hpa"], dtype=np.int16),
latitudes_degrees=np.linspace(90.0, -90.0, int(data["latitude_points"])),
longitudes_degrees=np.linspace(0.0, 360.0, int(data["longitude_points"]), endpoint=False),
field_shape=np.asarray(data["field_shape"], dtype=np.int64),
spectral_shape=np.asarray(data["spectral_shape"], dtype=np.int64),
ensemble_size=np.asarray(data["ensemble_size"]),
triangular_truncation=np.asarray(data["triangular_truncation"]),
nsp=np.asarray(data["nsp"]),
phase=rng.uniform(0, 2 * np.pi, experiments),
amplitude=float(generator["perturbation_ms"]) * np.asarray(generator["compensation_scales"]),
growth_rate=growth,
spectral_phase=rng.uniform(0, 2 * np.pi, experiments),
spectral_amplitude=rng.uniform(0.8, 1.2, experiments) * 1e-5,
base_wind_ms=np.asarray(generator["base_wind_ms"]),
spectral_slope=np.asarray(generator["spectral_slope"]),
coefficient_seed=np.asarray(config["seed"]),
raw_fields_materialized=np.asarray(False),
)
print(f"protocol={output.relative_to(ROOT)} bytes={output.stat().st_size} logical_field={data['field_shape']} logical_spectral={data['spectral_shape']}")
if __name__ == "__main__":
main()
|