| """Compute complete-grid DKE products in latitude chunks and compact T719 spectra.""" |
|
|
| import os |
| import sys |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import yaml |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT)) |
| from model.pangu_icon_dke import PanguIconDKEDiagnostics |
|
|
|
|
| def main(): |
| config = yaml.safe_load((ROOT / "conf/config.yaml").read_text()) |
| checkpoint = torch.load(ROOT / config["paths"]["checkpoint"], map_location="cpu", weights_only=False) |
| if checkpoint["format_version"] != config["data"]["format_version"]: |
| raise ValueError("checkpoint protocol mismatch") |
| model = PanguIconDKEDiagnostics(checkpoint["protocol"], checkpoint["coefficients"]) |
| world = int(os.environ.get("WORLD_SIZE", "1")) |
| distributed = world > 1 |
| if distributed: |
| torch.distributed.init_process_group("gloo") |
| rank = torch.distributed.get_rank() if distributed else 0 |
| experiments, times, _, nlat, nlon = model.fields.field_shape |
| chunk = int(config["runtime"]["latitude_chunk"]) |
| local = {} |
| for experiment in range(rank, experiments, world): |
| global_dke = np.zeros(times, dtype=np.float64) |
| dke_maps = np.empty((times, nlat, nlon), dtype=np.float32) |
| spectrum = np.empty((times, int(checkpoint["protocol"]["triangular_truncation"]) + 1), dtype=np.float64) |
| weights_total = 0.0 |
| weighted_sum = np.zeros(times, dtype=np.float64) |
| for start in range(0, nlat, chunk): |
| stop = min(start + chunk, nlat) |
| dke = model.fields.dke_chunk_all_times(experiment, slice(start, stop)) |
| weights = np.cos(np.deg2rad(model.fields.latitudes[start:stop]))[None, :, None] |
| weighted_sum += np.sum(dke * weights, axis=(1, 2)) |
| dke_maps[:, start:stop] = dke |
| weights_total += float(nlon * weights.sum()) |
| global_dke[:] = weighted_sum / weights_total |
| for time in range(times): |
| spectrum[time] = model.fields.spectrum(experiment, time) |
| local[experiment] = (global_dke, dke_maps, spectrum) |
| if distributed: |
| gathered = [None] * world if rank == 0 else None |
| torch.distributed.gather_object(local, gathered, dst=0) |
| if rank == 0: |
| local = {key: value for shard in gathered for key, value in shard.items()} |
| if rank == 0: |
| ordered = [local[index] for index in range(experiments)] |
| global_dke = np.stack([item[0] for item in ordered]) |
| maps = np.stack([item[1] for item in ordered]) |
| spectra = np.stack([item[2] for item in ordered]) |
| output = ROOT / config["paths"]["inference"] |
| output.parent.mkdir(parents=True, exist_ok=True) |
| np.savez_compressed(output, global_dke=global_dke, dke_maps=maps, spectra=spectra, experiments=np.asarray(checkpoint["protocol"]["experiments"]), times_hours=np.asarray(checkpoint["protocol"]["times_hours"]), latitudes_degrees=np.asarray(checkpoint["protocol"]["latitudes_degrees"]), longitudes_degrees=np.asarray(checkpoint["protocol"]["longitudes_degrees"]), total_wavenumber=np.arange(spectra.shape[-1]), format_version=np.asarray(checkpoint["format_version"]), logical_field_shape=np.asarray(checkpoint["protocol"]["field_shape"]), logical_spectral_shape=np.asarray(checkpoint["protocol"]["spectral_shape"])) |
| print(f"output={output.relative_to(ROOT)} global={global_dke.shape} maps={maps.shape} spectra={spectra.shape}") |
| if distributed: |
| torch.distributed.destroy_process_group() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
