lib.py

"""Shared constants and utilities. Everything else imports from here."""
import math

import torch

# ── channels & physics ───────────────────────────────────────────────
AIA_CHANNELS = ["94A", "131A", "171A", "193A", "211A", "304A", "335A", "1600A", "1700A"]
HMI_CHANNELS = ["Bx", "By", "Bz"]

# log10 peak response temperature per AIA channel (AIA instrument papers)
LOG_T = {"94A": 6.8, "131A": 5.6, "171A": 5.8, "193A": 6.2, "211A": 6.3,
         "304A": 4.7, "335A": 6.4, "1600A": 5.0, "1700A": 3.7}
LOGT = torch.tensor([LOG_T[c] for c in AIA_CHANNELS])
LOGT_LO, LOGT_HI = 3.5, 7.0

# robust per-channel scales (99.5th pct, 2010 sample); I -> asinh(I/s)
# p99.5/scale ~= 0.5 for every channel — unequal scales silently reweight
# the loss toward whichever channel is normalized hottest
AIA_SCALE = torch.tensor([12.0, 60.0, 1400.0, 1800.0, 900.0, 350.0, 90.0, 280.0, 3440.0])
B_SCALE = 500.0

# ── data splits (AIA small store covers August 2010 only) ────────────
TRAIN_DAYS = ["2010-08-07", "2010-08-10", "2010-08-14"]
VAL_DAYS = ["2010-08-01", "2010-08-21"]          # eruptive, quieter
HOURS = range(0, 24, 2)


def timestamps(days):
    return [f"{d}T{h:02d}:00:00" for d in days for h in HOURS]


# ── tensor utilities ─────────────────────────────────────────────────

def norm_logt(t):
    return (t - LOGT_LO) / (LOGT_HI - LOGT_LO) * 2 - 1


def fourier(x, freqs=(1, 2, 4, 8)):
    return torch.cat([f(2 * math.pi * k * x) for k in freqs for f in (torch.sin, torch.cos)], -1)


def mlp(d_in, d_hid, d_out, depth=2):
    import torch.nn as nn
    layers = [nn.Linear(d_in, d_hid), nn.SiLU()]
    for _ in range(depth - 1):
        layers += [nn.Linear(d_hid, d_hid), nn.SiLU()]
    return nn.Sequential(*layers, nn.Linear(d_hid, d_out))


def scatter_mean(src, index, n):
    out = torch.zeros(*src.shape[:-2], n, src.shape[-1], device=src.device, dtype=src.dtype)
    out.index_add_(-2, index, src)
    deg = torch.zeros(n, 1, device=src.device, dtype=src.dtype)
    deg.index_add_(0, index, torch.ones(index.shape[0], 1, device=src.device, dtype=src.dtype))
    return out / deg.clamp(min=1)


def prep(frame, res):
    """Raw frame -> normalized pooled tensors (the space all scoring lives in)."""
    pool = lambda x: torch.as_tensor(x).view(res, 512 // res, res, 512 // res).mean((1, 3))
    out = {c: torch.asinh(pool(frame[c]) / s) for c, s in zip(AIA_CHANNELS, AIA_SCALE.tolist())}
    out |= {c: torch.tanh(pool(frame[c]) / B_SCALE) for c in ("Bx", "By", "Bz")}
    return out