rustwx_stormscope/rust_datasources.py

#!/usr/bin/env python
"""Fast GOES/MRMS inputs for StormScope, backed by the rust `stormscope_obs` extension.

- GOES: rust parallel-downloads the MCMIPC files; we decode the 8 StormScope channels with
  h5netcdf (fast, proven). lat/lon come from earth2studio's GOES.grid() (same ABI fixed grid).
- MRMS: rust parallel-downloads + decodes the GRIB2 composite (pure rust).

`goes_input` / `mrms_input` return (tensor, coords) shaped exactly like earth2studio's fetch_data
output, so StormScope's interpolator/sampler behave identically. The download - the operational
bottleneck - is parallelized; the model runs unchanged.
"""
import os
import time
from collections import OrderedDict
from pathlib import Path

import numpy as np
import torch

_CMI = {"abi01c": "CMI_C01", "abi02c": "CMI_C02", "abi03c": "CMI_C03", "abi07c": "CMI_C07",
        "abi08c": "CMI_C08", "abi09c": "CMI_C09", "abi10c": "CMI_C10", "abi13c": "CMI_C13"}


def _env(name, legacy_name=None, default=None):
    if name in os.environ:
        return os.environ[name]
    if legacy_name and legacy_name in os.environ:
        return os.environ[legacy_name]
    return default


def _cache_dir(kind):
    root = _env("RUSTWX_STORMSCOPE_CACHE_DIR", "SSFAST_CACHE_DIR")
    if root:
        return os.path.join(root, kind)
    return os.path.join(".", "cache", kind)


def _iso_times(start_date, lead_times):
    base = np.datetime64(start_date[0])
    return [str((base + np.timedelta64(lt)).astype("datetime64[s]")).replace(" ", "T")
            for lt in lead_times]


def _goes_decode_cache_path(path, variables):
    cache = Path(_env("RUSTWX_STORMSCOPE_GOES_DECODE_CACHE", "SSFAST_GOES_DECODE_CACHE", _cache_dir("goes_decoded")))
    var_key = "-".join(str(v) for v in variables)
    return cache / f"{Path(path).name}.{var_key}.f32.npy"


def _prune_decode_cache(cache_dir: Path):
    max_files = int(_env("RUSTWX_STORMSCOPE_GOES_DECODE_CACHE_MAX_FILES", "SSFAST_GOES_DECODE_CACHE_MAX_FILES", "96"))
    if max_files <= 0 or not cache_dir.exists():
        return
    now = time.time()
    # Keep the common case cheap; pruning every call is fine because the cache is small.
    files = sorted(
        (p for p in cache_dir.glob("*.npy") if p.is_file()),
        key=lambda p: p.stat().st_mtime,
        reverse=True,
    )
    for p in files[max_files:]:
        try:
            p.unlink()
        except OSError:
            pass
    max_age_hours = float(_env("RUSTWX_STORMSCOPE_GOES_DECODE_CACHE_MAX_AGE_HOURS", "SSFAST_GOES_DECODE_CACHE_MAX_AGE_HOURS", "18"))
    if max_age_hours > 0:
        cutoff = now - max_age_hours * 3600.0
        for p in files[:max_files]:
            try:
                if p.stat().st_mtime < cutoff:
                    p.unlink()
            except OSError:
                pass


def _read_goes_frame(path, variables):
    import xarray as xr

    cache_path = _goes_decode_cache_path(path, variables)
    try:
        if cache_path.exists() and cache_path.stat().st_mtime >= Path(path).stat().st_mtime:
            return np.load(cache_path)
    except Exception:
        pass

    ds = xr.open_dataset(path, engine="h5netcdf")  # mask_and_scale applies scale/offset -> physical
    try:
        frame = np.stack(
            [np.asarray(ds[_CMI[str(v)]].values, dtype=np.float32) for v in variables],
            axis=0,
        )
    finally:
        ds.close()

    try:
        cache_path.parent.mkdir(parents=True, exist_ok=True)
        tmp = cache_path.with_suffix(cache_path.suffix + ".tmp")
        with open(tmp, "wb") as f:
            np.save(f, frame)
        os.replace(tmp, cache_path)
        _prune_decode_cache(cache_path.parent)
    except Exception:
        pass
    return frame


def goes_input(satellite, scan_mode, start_date, variables, lead_times, device, hw):
    """Returns (x[1,L,C,H,W] float32 tensor, coords OrderedDict time/lead_time/variable/y/x)."""
    import stormscope_obs
    cache = _env("RUSTWX_STORMSCOPE_GOES_CACHE", "SSFAST_GOES_CACHE", _cache_dir("goes_nc"))
    times = _iso_times(start_date, lead_times)
    paths = stormscope_obs.download_goes_files(satellite, times, cache)
    frames = [_read_goes_frame(p, variables) for p in paths]  # [C, H, W]
    data = np.stack(frames, axis=0)[None]                  # [1, L, C, H, W]
    H, W = hw
    coords = OrderedDict([
        ("time", np.asarray(start_date)),
        ("lead_time", np.asarray(lead_times)),
        ("variable", np.asarray(list(variables))),
        ("y", np.arange(H)),
        ("x", np.arange(W)),
    ])
    return torch.from_numpy(np.ascontiguousarray(data)).to(device), coords


def mrms_input(start_date, lead_times, device):
    """Returns (x[1,L,1,H,W] tensor, coords time/lead_time/variable/lat/lon, lat1d, lon1d)."""
    import stormscope_obs
    times = _iso_times(start_date, lead_times)
    cache = _env("RUSTWX_STORMSCOPE_MRMS_CACHE", "SSFAST_MRMS_CACHE", _cache_dir("mrms_grib2"))
    res = stormscope_obs.fetch_mrms_sequence(times, cache)
    T, H, W = res["shape"]
    data = np.asarray(res["data"], dtype=np.float32).reshape(T, 1, H, W)[None]   # [1, L, 1, H, W]
    # Keep MRMS missing/no-coverage flags finite (-999/-99), matching earth2studio.
    # StormScopeMRMS.prep_input imputes values <= -20 dBZ to -10; NaNs poison the denoiser.
    north, south, west, east = res["geo"]   # row 0 = north, col 0 = west (MRMS N->S, W->E)
    lats = np.linspace(north, south, H).astype(np.float32)
    lons = np.linspace(west, east, W).astype(np.float32)
    lons = np.where(lons < 0.0, lons + 360.0, lons).astype(np.float32)
    coords = OrderedDict([
        ("time", np.asarray(start_date)),
        ("lead_time", np.asarray(lead_times)),
        ("variable", np.asarray(["refc"])),
        ("lat", lats),
        ("lon", lons),
    ])
    return torch.from_numpy(np.ascontiguousarray(data)).to(device), coords, lats, lons