data/data.py

"""nika-5 data layer — SDOML v2 (NASA SDO, ML-ready, public domain).

Source: s3://gov-nasa-hdrl-data1/contrib/fdl-sdoml/fdl-sdoml-v2/ (anonymous).
  sdomlv2_small.zarr       AIA 9 channels, 2010, [T≈6135, 512, 512] f32 (per-channel T!)
  sdomlv2_hmi_small.zarr   HMI Bx/By/Bz,   2010, [T≈25540, 512, 512] f32
  sdomlv2.zarr / _hmi.zarr full 2010-2020 stores, same layout by year

Design facts that shape everything downstream:
  * Each channel is a separate zarr array WITH ITS OWN time axis; timestamps
    live in attrs["T_OBS"]. Nothing is aligned for you — align() does it.
  * AIA channels are a temperature ladder (log T is the physical modality
    axis); HMI B is the causal substrate. See README.
"""
from dataclasses import dataclass
from functools import cached_property
from pathlib import Path

import numpy as np
import s3fs
import zarr

BASE = "gov-nasa-hdrl-data1/contrib/fdl-sdoml/fdl-sdoml-v2"
from lib import AIA_CHANNELS, HMI_CHANNELS, LOG_T  # noqa: F401


def _fs():
    return s3fs.S3FileSystem(anon=True, client_kwargs={"region_name": "us-west-2"})


@dataclass
class SDOML:
    year: str = "2010"
    small: bool = True
    cache_dir: str = "data/cache"

    @cached_property
    def _aia(self):
        name = "sdomlv2_small.zarr" if self.small else "sdomlv2.zarr"
        return zarr.open(s3fs.S3Map(f"{BASE}/{name}", s3=_fs(), check=False), mode="r")[self.year]

    @cached_property
    def _hmi(self):
        name = "sdomlv2_hmi_small.zarr" if self.small else "sdomlv2_hmi.zarr"
        return zarr.open(s3fs.S3Map(f"{BASE}/{name}", s3=_fs(), check=False), mode="r")[self.year]

    @staticmethod
    def _parse(ts: str) -> str:
        """T_OBS -> ISO. Handles TAI suffixes, dotted dates, and :60 seconds
        (TAI leap-second convention that datetime64 rejects)."""
        s = ts[:19].replace("_TAI", "").replace(".", "-", 2).replace("_", "T")
        if s[17:19] == "60":
            s = s[:17] + "59"
        return s

    @cached_property
    def times(self) -> dict:
        """Per-channel numpy datetime64 arrays parsed from T_OBS attrs."""
        out = {}
        for ch in AIA_CHANNELS:
            out[ch] = np.array([self._parse(t) for t in self._aia[ch].attrs["T_OBS"]], dtype="datetime64[s]")
        for ch in HMI_CHANNELS:
            out[ch] = np.array([self._parse(t) for t in self._hmi[ch].attrs["T_OBS"]], dtype="datetime64[s]")
        return out

    def align(self, when: str, channels=None, tol_minutes: int = 30) -> dict:
        """Nearest frame index per channel to `when` (ISO string), within tol."""
        channels = channels or (AIA_CHANNELS + HMI_CHANNELS)
        target = np.datetime64(when)
        picks = {}
        for ch in channels:
            ts = self.times[ch]
            i = int(np.argmin(np.abs(ts - target)))
            dt = abs((ts[i] - target) / np.timedelta64(1, "m"))
            if dt <= tol_minutes:
                picks[ch] = (i, str(ts[i]))
        return picks

    def frame(self, when: str, channels=None, tol_minutes: int = 30) -> dict:
        """Aligned multimodal snapshot: {channel: [512,512] float32}, cached."""
        channels = channels or (AIA_CHANNELS + HMI_CHANNELS)
        key = when.replace(":", "").replace("-", "") + "_" + "-".join(channels)
        cache = Path(self.cache_dir) / f"{key}.npz"
        if cache.exists():
            z = np.load(cache)
            return {ch: z[ch] for ch in z.files}
        picks = self.align(when, channels, tol_minutes)
        out = {}
        for ch, (i, _) in picks.items():
            src = self._aia if ch in AIA_CHANNELS else self._hmi
            out[ch] = np.asarray(src[ch][i], dtype=np.float32)
        cache.parent.mkdir(parents=True, exist_ok=True)
        np.savez_compressed(cache, **out)
        return out