noaa-goes-visualization/suvi/corruptions.py

"""Catalogue of synthetic frame corruptions used by the test bench.

Each entry models a failure mode actually seen in (or plausible for) this archive.
Corruptions are seeded and reproducible: the same case seed and frame always yield
byte-identical output, so a bench run can be repeated exactly.

Two kinds, because they break a frame at different layers:

* **array** corruptions damage the pixels; the bench re-writes a valid FITS around
  the result.
* **file** corruptions damage the bytes on disk, producing files that are not valid
  FITS at all -- truncated downloads, missing HDUs, bit rot.

The ``recompute_stats`` flag is the subtle part.  A real eclipse frame has a header
whose ``IMG_MEAN`` agrees with its dim pixels, because NOAA computed it from them; a
frame damaged in transit keeps the *original* header over broken pixels.  Setting
this correctly per mode is what makes the bench's verdict on header-only detection
honest -- otherwise header_v1 would appear to catch corruptions it could never see.
"""

from dataclasses import dataclass

import numpy as np

from . import fitsio

#: Corruption strength in [0, 1].  0 is a barely-perceptible defect, 1 is total loss.
DEFAULT_SEVERITY = 1.0


@dataclass(frozen=True)
class Corruption:
    """One failure mode."""

    name: str
    #: Which of the four groups this belongs to, for per-group reporting.
    group: str
    #: 'array' (damages pixels) or 'file' (damages bytes on disk).
    kind: str
    apply: callable
    #: Whether the header's radiance statistics are recomputed from the damaged
    #: pixels.  True models a fault upstream of NOAA's header generation.
    recompute_stats: bool = False
    #: Whether this mode needs a second frame to draw from.
    needs_donor: bool = False


# ------------------------------------------------------------------ dropout / blackout


def _eclipse_dim(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    """The archive's most common real failure: Earth shadow drops radiance ~1000x."""
    factor = 10.0 ** (-4.0 * severity)
    dimmed = image * factor
    # Real eclipse frames keep sensor read noise, so they are not exactly zero.
    noise = rng.normal(0.0, float(np.abs(image).mean()) * 1e-4, image.shape)
    return (dimmed + noise).astype(np.float32), {"DEGRADED": True, "ECLIPSE": 2}


def _all_zero(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    return np.zeros_like(image), {"EMPTY": True}


def _nan_fill(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    """Undefined pixels over part or all of the frame."""
    out = image.copy()
    if severity >= 1.0:
        out[:] = np.nan
    else:
        mask = rng.random(image.shape) < severity
        out[mask] = np.nan
    return out, {}


def _zblank_fill(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    """Fill with the FITS blank sentinel rather than NaN."""
    out = image.copy()
    mask = rng.random(image.shape) < severity if severity < 1.0 else np.ones(image.shape, bool)
    out[mask] = fitsio.ZBLANK
    return out, {}


# -------------------------------------------------------------------------- structural


def _truncate(raw, rng, severity=DEFAULT_SEVERITY):
    """A download cut short.  Keeps at least the primary header."""
    keep = max(fitsio.BLOCK, int(len(raw) * (1.0 - 0.9 * severity)))
    return raw[:keep]


def _drop_image_hdu(raw, rng, severity=DEFAULT_SEVERITY):
    """Only the primary header survives -- the blank-HDU case the old filter hit."""
    return raw[: fitsio.BLOCK]


def _block_corruption(raw, rng, severity=DEFAULT_SEVERITY):
    """Random bytes overwritten inside the data unit, leaving the header intact."""
    out = bytearray(raw)
    start = fitsio.BLOCK * 8  # past the headers
    if len(out) <= start:
        return bytes(out)
    span = max(1, int((len(out) - start) * 0.02 * severity))
    offset = int(rng.integers(start, len(out) - span))
    out[offset : offset + span] = rng.integers(0, 256, span, dtype=np.uint8).tobytes()
    return bytes(out)


def _torn_frame(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    """Part of the frame comes from another observation -- a bad merge."""
    if donor is None:
        return image, {}
    out = image.copy()
    split = int(image.shape[0] * (1.0 - severity * 0.5))
    out[split:, :] = donor[split:, :]
    return out, {}


def _dropped_rows(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    """Whole scan lines lost, as from a dropped packet."""
    out = image.copy()
    count = max(1, int(image.shape[0] * 0.3 * severity))
    rows = rng.choice(image.shape[0], size=count, replace=False)
    out[rows, :] = 0.0
    return out, {}


# ------------------------------------------------------------------------ radiometric


def _gain_shift(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    """Calibration drift: everything scaled by a constant factor."""
    factor = 1.0 + 4.0 * severity * (1 if rng.random() < 0.5 else -0.2)
    return (image * factor).astype(np.float32), {}


def _offset_shift(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    """A constant added to every pixel -- a dark-current or bias fault."""
    return (image + float(np.abs(image).mean()) * 5.0 * severity).astype(np.float32), {}


def _saturate(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    """A blowout that drives a large fraction of the frame to the ceiling."""
    ceiling = float(np.nanmax(image)) or 1.0
    boosted = image * (1.0 + 50.0 * severity)
    return np.minimum(boosted, ceiling * 50.0).astype(np.float32), {}


def _gaussian_noise(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    scale = float(np.nanstd(image)) * severity
    return (image + rng.normal(0.0, scale, image.shape)).astype(np.float32), {}


def _salt_pepper(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    """Cosmic-ray hits and dead pixels."""
    out = image.copy()
    fraction = 0.05 * severity
    mask = rng.random(image.shape) < fraction
    extreme = float(np.nanmax(image)) or 1.0
    out[mask] = np.where(rng.random(int(mask.sum())) < 0.5, 0.0, extreme * 10.0)
    return out, {}


# ------------------------------------------------------------------ geometric/temporal


def _translate(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    """Mispointing: the solar disc sits off centre.

    The shift is a fraction of the frame, not a fixed pixel count.  200 px was right for
    the archive's native 1280, but training samples are half that, and a fixed count
    both doubles the apparent severity and -- below 200 px -- indexes past the frame and
    raises.  Scaling by ``min(h, w) / 1280`` leaves 1280 frames byte-identical.
    """
    h, w = image.shape
    shift = int(200 * severity * min(h, w) / 1280)
    dx = int(rng.integers(-shift, shift + 1)) if shift else 0
    dy = int(rng.integers(-shift, shift + 1)) if shift else 0
    out = np.zeros_like(image)
    xs, xd = (max(0, -dx), max(0, dx))
    ys, yd = (max(0, -dy), max(0, dy))
    height, width = h - abs(dy), w - abs(dx)
    out[yd : yd + height, xd : xd + width] = image[ys : ys + height, xs : xs + width]
    return out, {"CRPIX1": (w + 1) / 2.0 + dx, "CRPIX2": (h + 1) / 2.0 + dy}


def _rotate(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    """Wrong roll angle -- the disc is round, so only structure reveals this."""
    import cv2 as cv

    angle = 180.0 * severity
    centre = ((image.shape[1] - 1) / 2.0, (image.shape[0] - 1) / 2.0)
    matrix = cv.getRotationMatrix2D(centre, angle, 1.0)
    rotated = cv.warpAffine(
        np.nan_to_num(image), matrix, (image.shape[1], image.shape[0]), flags=cv.INTER_LINEAR
    )
    return rotated.astype(np.float32), {"CROTA": angle}


def _yaw_flip(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    """The spacecraft's twice-yearly yaw flip applied when it should not be."""
    return np.flip(np.flip(image, 0), 1).copy(), {"YAW_FLIP": 1}


def _header_only(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    """Pixels untouched; the metadata lies.

    The single largest real failure class (~48% of all BAD frames): for nine days
    in Feb-Mar 2023 every file declared WAVELNTH=171 and a displaced CRPIX while
    the images themselves were pristine.  A detector must flag these, even though
    the pixels would have been fine to use.  The wrong band is drawn at random; when it happens to match the
    frame's real band, the displaced CRPIX still trips the pointing rule, which
    is precisely how the real episode presented on the 171A band itself.
    """
    wrong_band = int(rng.choice([94, 131, 171, 195, 284, 304]))
    shift = 10.0 + 10.0 * severity          # comfortably past the 8 px rule
    return image.copy(), {"WAVELNTH": wrong_band,
                          "CRPIX1": 640.5 + shift, "CRPIX2": 640.5 + shift}


def _frozen(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    """The feed stalled: this frame is a byte-for-byte repeat of a neighbour."""
    return (donor.copy() if donor is not None else image), {}


def _wrong_time(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    """A frame from a different observation filed under this timestamp."""
    return (donor.copy() if donor is not None else image), {}


# --------------------------------------------- detector faults seen in the real archive
#
# Modelled on the hand-confirmed defects in labels/real_defects.json.  Bars and
# blocks lie along the detector's rows and columns, which in the delivered image
# sit at -REFCROTA, so each draws a rotation and records it as the frame's
# REFCROTA: a detector reading the header then finds the damage where it would
# find the real thing.

#: REFCROTA values seen across the archive span roughly this range (degrees).
_REFCROTA_RANGE = (5.0, 27.0)


def _detector_axes(shape, refcrota):
    """Pixel coordinates rotated onto the detector's row/column axes, centred."""
    height, width = shape
    yy, xx = np.ogrid[:height, :width]
    theta = np.radians(-refcrota)
    x, y = xx - (width - 1) / 2.0, yy - (height - 1) / 2.0
    along = x * np.cos(theta) - y * np.sin(theta)
    across = x * np.sin(theta) + y * np.cos(theta)
    return along, across


def _bright(image):
    """A brightness that reads as saturated in this frame's own band."""
    finite = image[np.isfinite(image)]
    return float(np.percentile(finite, 99.9)) if finite.size else 1.0


def _bar_mask(image, rng, severity):
    refcrota = float(rng.uniform(*_REFCROTA_RANGE))
    along, across = _detector_axes(image.shape, refcrota)
    line = across if rng.random() < 0.5 else along       # a detector row or column
    size = min(image.shape)
    width = size * (0.015 + 0.035 * severity)
    offset = float(rng.uniform(-0.45, 0.45)) * size
    return np.abs(line - offset) < width / 2.0, refcrota


def _detector_bar(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    """A saturated bar along a whole detector row or column (232 in the audit)."""
    mask, refcrota = _bar_mask(image, rng, severity)
    out = image.copy()
    out[mask] = _bright(image) * (0.5 + 4.5 * severity)
    return out, {"REFCROTA": refcrota}


def _noise_bar(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    """A detector-row bar filled with noise rather than light (the audit's #2140)."""
    mask, refcrota = _bar_mask(image, rng, severity)
    out = image.copy()
    sigma = _bright(image) * (0.1 + 0.9 * severity)
    out[mask] = np.abs(rng.normal(0.0, sigma, int(mask.sum()))).astype(np.float32)
    return out, {"REFCROTA": refcrota}


def _bright_block(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    """A saturated rectangle on the detector axes, reaching in from a corner."""
    refcrota = float(rng.uniform(*_REFCROTA_RANGE))
    along, across = _detector_axes(image.shape, refcrota)
    size = min(image.shape)
    half = size * (0.1 + 0.2 * severity)
    corner = size * 0.42
    u0 = corner * (1 if rng.random() < 0.5 else -1)
    v0 = corner * (1 if rng.random() < 0.5 else -1)
    mask = (np.abs(along - u0) < half) & (np.abs(across - v0) < half)
    out = image.copy()
    out[mask] = _bright(image) * (0.5 + 1.5 * severity)
    return out, {"REFCROTA": refcrota}


def _ghost_disc(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    """A displaced second copy of the Sun over the real one (63 in the audit)."""
    size = min(image.shape)
    distance = size * float(rng.uniform(0.15, 0.6))
    angle = float(rng.uniform(0.0, 2 * np.pi))
    dy, dx = int(round(distance * np.sin(angle))), int(round(distance * np.cos(angle)))
    shifted = np.zeros_like(image)
    h, w = image.shape
    ys, yd = max(0, -dy), max(0, dy)
    xs, xd = max(0, -dx), max(0, dx)
    shifted[yd:h - ys, xd:w - xs] = image[ys:h - yd, xs:w - xd]
    weight = 0.3 + 0.7 * severity
    return (np.nan_to_num(image) + weight * np.nan_to_num(shifted)).astype(np.float32), {}


def _short_exposure(image, rng, severity=DEFAULT_SEVERITY, donor=None):
    """No long exposure: scaled-up short exposures, noise everywhere (196 in the audit).

    NOAA flags these honestly -- N_LONG=0 -- and computes the statistics from the
    noisy pixels, so the header both tells and agrees.
    """
    finite = image[np.isfinite(image)]
    level = float(np.mean(np.abs(finite))) if finite.size else 1.0
    sigma = level * (0.2 + 1.8 * severity)
    noisy = image + rng.normal(0.0, sigma, image.shape)
    return noisy.astype(np.float32), {"N_LONG": 0, "NUM_IMGS": 1}


CATALOG = {
    corruption.name: corruption
    for corruption in (
        # Dropout / blackout -- the header follows the pixels, as NOAA computes it
        # from the image it actually produced.
        Corruption("eclipse_dim", "dropout", "array", _eclipse_dim, recompute_stats=True),
        Corruption("all_zero", "dropout", "array", _all_zero, recompute_stats=True),
        Corruption("nan_fill", "dropout", "array", _nan_fill, recompute_stats=False),
        Corruption("zblank_fill", "dropout", "array", _zblank_fill, recompute_stats=False),
        # Structural -- damage after the file was written, so the header is stale.
        Corruption("truncate", "structural", "file", _truncate),
        Corruption("drop_image_hdu", "structural", "file", _drop_image_hdu),
        Corruption("block_corruption", "structural", "file", _block_corruption),
        Corruption("torn_frame", "structural", "array", _torn_frame, needs_donor=True),
        Corruption("dropped_rows", "structural", "array", _dropped_rows),
        # Radiometric -- an instrument fault upstream of header generation.
        Corruption("gain_shift", "radiometric", "array", _gain_shift, recompute_stats=True),
        Corruption("offset_shift", "radiometric", "array", _offset_shift, recompute_stats=True),
        Corruption("saturate", "radiometric", "array", _saturate, recompute_stats=True),
        Corruption("gaussian_noise", "radiometric", "array", _gaussian_noise),
        Corruption("salt_pepper", "radiometric", "array", _salt_pepper),
        # Geometric / temporal -- the modes single-frame detectors are worst at.
        Corruption("translate", "geometric", "array", _translate, recompute_stats=True),
        Corruption("rotate", "geometric", "array", _rotate, recompute_stats=True),
        Corruption("yaw_flip", "geometric", "array", _yaw_flip),
        Corruption("frozen", "geometric", "array", _frozen, needs_donor=True),
        Corruption("wrong_time", "geometric", "array", _wrong_time, needs_donor=True),
        # Metadata -- the pixels are fine; the header is not.
        Corruption("header_only", "metadata", "array", _header_only),
        # Detector faults from the real-defect audit.  The header statistics follow
        # the pixels: NOAA computed them from the frame as delivered.
        Corruption("detector_bar", "detector", "array", _detector_bar, recompute_stats=True),
        Corruption("noise_bar", "detector", "array", _noise_bar, recompute_stats=True),
        Corruption("bright_block", "detector", "array", _bright_block, recompute_stats=True),
        Corruption("ghost_disc", "detector", "array", _ghost_disc, recompute_stats=True),
        Corruption("short_exposure", "detector", "array", _short_exposure,
                   recompute_stats=True),
    )
}

GROUPS = sorted({corruption.group for corruption in CATALOG.values()})

#: The measured real-world failure mixture: every BAD verdict of the production
#: judgment over 2023-01-01..2025-04-01 (judge run 46, 78,952 of 2,949,615 frames),
#: folded onto the catalogue by the rule that fired:
#:   ECLIPSE flag 33,061 + IMG_MEAN dark 2,264 + lunar transit 479  -> eclipse_dim
#:   header episode: wrong WAVELNTH 28,037 + sun-centre offset 5,485 -> header_only
#:   no long exposure (N_LONG=0) 4,568                               -> short_exposure
#:   disc_v1 radius/limb failures 1,433 (pointing drift)             -> translate
#:   field_v1 noise-only lines 133                                   -> noise_bar
#:   field_v1 bright lines / far field 867, split as the hand-confirmed
#:     defects field_v1 catches split (labels/real_defects.json:
#:     232 bars, 29 blocks, 57 ghosts)                               -> detector_bar,
#:                                                                      bright_block, ghost_disc
#:   EMPTY flag 75                                                   -> all_zero
#:   IMG_MEAN blowout 48                                             -> saturate
#:   temporal_v2 alone 2,502 -- frames that decoded but are out of step with
#:     their clean neighbours, of no single identified kind -- spread over the
#:     remaining pixel modes, so all stay reachable.
#:   unreadable containers: 122 frames in 3.04M (the earlier archive scan; the
#:     index leaves stubs out, so run 46 met none) -- spread over the file modes.
#: Used wherever damage is meant to look exactly like the archive: the bench's
#: trace-driven cases.  Training draws from :data:`TRAINING_MIX` instead, which
#: is this smoothed toward uniform so rare modes stay reachable.
_FIELD_SPLIT = {"detector_bar": 232, "bright_block": 29, "ghost_disc": 57}
_REAL_COUNTS = {
    "eclipse_dim": 33_061 + 2_264 + 479,
    "header_only": 28_037 + 5_485,
    "short_exposure": 4_568,
    "translate": 1_433,
    "noise_bar": 133,
    **{name: 867 * share / sum(_FIELD_SPLIT.values()) for name, share in _FIELD_SPLIT.items()},
    "all_zero": 75,
    "saturate": 48,
}
_REAL_TAIL = [name for name in sorted(CATALOG)
              if name not in _REAL_COUNTS and CATALOG[name].kind == "array"]
_FILE_MODES = [name for name in sorted(CATALOG) if CATALOG[name].kind == "file"]
_REAL_COUNTS.update({name: 2_502 / len(_REAL_TAIL) for name in _REAL_TAIL})
_REAL_COUNTS.update({name: 122 / len(_FILE_MODES) for name in _FILE_MODES})
REAL_MIX = {name: count / sum(_REAL_COUNTS.values()) for name, count in _REAL_COUNTS.items()}
_REAL_MODES = tuple(sorted(REAL_MIX))
_REAL_WEIGHTS = np.array([REAL_MIX[name] for name in _REAL_MODES], dtype=np.float64)
_REAL_WEIGHTS /= _REAL_WEIGHTS.sum()

#: Severity is drawn over (0.05, 1.0), not the old (0.5, 1.0): the floor made
#: every simulated defect blatant, while reality's dominant defects are subtle --
#: penumbral eclipses sit at 0.3-1.0x normal brightness (eclipse_dim's
#: 10**(-4*severity) only reaches that regime below severity ~0.13), and the
#: header episode is pixel-invisible entirely.
SEVERITY_RANGE = (0.05, 1.0)


#: How much of the *training* mixture is spent on uniform mode coverage rather
#: than on realism.  The sampler used to draw its synthetic damage uniformly over
#: the catalogue, on the argument that every mode should keep appearing; because
#: that path fires far more often than the trace-BAD path, the realized training
#: distribution came out ~70% modes that are ~0.1% of reality, and only 29% on
#: eclipse_dim plus header_only, which together are ~88% of real failures.  Two
#: of the over-represented modes (`yaw_flip`, `rotate`) have no archive instance
#: at all.
#:
#: Smoothing the archive mixture toward uniform keeps every mode reachable
#: without letting coverage become the dominant signal, and -- the point of
#: doing it here rather than per path -- makes the realized distribution
#: independent of how often each damage path fires, which is what went wrong.
COVERAGE_WEIGHT = 0.10
TRAINING_MIX = {name: (1.0 - COVERAGE_WEIGHT) * REAL_MIX[name]
                      + COVERAGE_WEIGHT / len(CATALOG)
                for name in sorted(CATALOG)}
_TRAINING_MODES = tuple(sorted(TRAINING_MIX))
_TRAINING_WEIGHTS = np.array([TRAINING_MIX[name] for name in _TRAINING_MODES],
                             dtype=np.float64)
_TRAINING_WEIGHTS /= _TRAINING_WEIGHTS.sum()


def draw_real(rng):
    """One mode, distributed as the real archive's failures are.

    For *evaluation*: the bench measures performance against reality, so it
    simulates reality.  Training uses :func:`draw_training`.
    """
    return str(rng.choice(_REAL_MODES, p=_REAL_WEIGHTS))


def draw_training(rng):
    """One mode from the archive mixture smoothed toward uniform coverage.

    Used for every damaged frame the sampler produces, whichever path chose to
    damage it, so the realized mixture is a property of this function rather
    than of the ratio between two call sites.
    """
    return str(rng.choice(_TRAINING_MODES, p=_TRAINING_WEIGHTS))


def draw_severity(rng):
    low, high = SEVERITY_RANGE
    return float(rng.uniform(low, high))


def by_group(group):
    return [name for name, c in CATALOG.items() if c.group == group]


def apply_array(name, image, seed, severity=DEFAULT_SEVERITY, donor=None):
    """Apply an array corruption.  Returns (image, header_overrides).

    Deterministic in `seed`, so a bench case is exactly reproducible.
    """
    corruption = CATALOG[name]
    if corruption.kind != "array":
        raise ValueError(f"{name} is a {corruption.kind} corruption, not an array one")
    if corruption.needs_donor and donor is None:
        raise ValueError(f"{name} requires a donor frame")
    rng = np.random.default_rng(seed)
    return corruption.apply(np.asarray(image, dtype=np.float32), rng,
                            severity=severity, donor=donor)


def apply_file(name, raw, seed, severity=DEFAULT_SEVERITY):
    """Apply a file-level corruption to raw FITS bytes."""
    corruption = CATALOG[name]
    if corruption.kind != "file":
        raise ValueError(f"{name} is a {corruption.kind} corruption, not a file one")
    rng = np.random.default_rng(seed)
    return corruption.apply(raw, rng, severity=severity)


def recomputed_stats(image):
    """Header statistics consistent with `image`, as NOAA would have written them."""
    finite = image[np.isfinite(image)]
    if finite.size == 0:
        return {"IMG_MIN": 0.0, "IMG_MAX": 0.0, "IMG_MEAN": 0.0, "IMG_SDEV": 0.0}
    return {
        "IMG_MIN": float(finite.min()),
        "IMG_MAX": float(finite.max()),
        "IMG_MEAN": float(finite.mean()),
        "IMG_SDEV": float(finite.std()),
    }