noaa-goes-visualization/tests/test_corruptions.py

import numpy as np
import pytest

from conftest import solar_disc, write_fits
from suvi import corruptions, detectors, fitsio

ARRAY_MODES = [n for n, c in corruptions.CATALOG.items() if c.kind == "array"]
FILE_MODES = [n for n, c in corruptions.CATALOG.items() if c.kind == "file"]


@pytest.fixture
def image():
    return solar_disc(peak=3.0)


@pytest.fixture
def donor():
    return solar_disc(peak=3.0, centre=(700, 640))


def apply(name, image, donor=None, seed=7, severity=1.0):
    return corruptions.apply_array(name, image, seed, severity, donor)


# ------------------------------------------------------------------------ catalogue


def test_catalogue_covers_all_six_groups():
    assert set(corruptions.GROUPS) == {"dropout", "structural", "radiometric",
                                       "geometric", "metadata", "detector"}
    for group in corruptions.GROUPS:
        assert corruptions.by_group(group), f"{group} has no modes"


def test_every_entry_is_self_consistent():
    for name, corruption in corruptions.CATALOG.items():
        assert corruption.name == name
        assert corruption.kind in ("array", "file")
        assert callable(corruption.apply)


# -------------------------------------------------------------------- determinism


@pytest.mark.parametrize("name", ARRAY_MODES)
def test_array_corruptions_are_deterministic(name, image, donor):
    first, header_a = apply(name, image, donor)
    second, header_b = apply(name, image, donor)
    np.testing.assert_array_equal(np.nan_to_num(first), np.nan_to_num(second))
    assert header_a == header_b


@pytest.mark.parametrize("name", [n for n in ARRAY_MODES if n != "header_only"])
def test_array_corruptions_change_the_data(name, image, donor):
    damaged, _ = apply(name, image, donor)
    assert damaged.shape == image.shape
    assert not np.array_equal(np.nan_to_num(damaged), np.nan_to_num(image))


def test_header_only_lies_in_the_header_and_nowhere_else(image, donor):
    """The largest real failure class: pristine pixels under falsified metadata.
    The detector must see it (wrong band, or a CRPIX shift past the 8 px rule)
    while the pixels stay byte-identical."""
    damaged, overrides = apply("header_only", image, donor)
    np.testing.assert_array_equal(damaged, image)
    assert overrides["WAVELNTH"] in (94, 131, 171, 195, 284, 304)
    assert abs(overrides["CRPIX1"] - 640.5) > 8.0
    assert abs(overrides["CRPIX2"] - 640.5) > 8.0


def test_the_real_mixture_is_normalised_and_covers_every_mode():
    assert set(corruptions.REAL_MIX) == set(corruptions.CATALOG)
    assert sum(corruptions.REAL_MIX.values()) == pytest.approx(1.0, abs=1e-6)
    # The two dominant real classes dominate the mixture too.
    assert corruptions.REAL_MIX["eclipse_dim"] > 0.4
    assert corruptions.REAL_MIX["header_only"] > 0.4
    rng = np.random.default_rng(0)
    draws = [corruptions.draw_real(rng) for _ in range(2000)]
    assert draws.count("eclipse_dim") > 800
    assert draws.count("header_only") > 700
    for severity in (corruptions.draw_severity(rng) for _ in range(100)):
        low, high = corruptions.SEVERITY_RANGE
        assert low <= severity <= high


def test_the_training_mixture_is_the_archive_smoothed_not_uniform():
    """Training damage must look mostly like reality, while still reaching every
    mode.  The sampler previously drew its synthetic damage uniformly, which put
    ~70% of training damage on modes that are ~0.1% of the archive."""
    assert set(corruptions.TRAINING_MIX) == set(corruptions.CATALOG)
    assert sum(corruptions.TRAINING_MIX.values()) == pytest.approx(1.0, abs=1e-6)
    # Reality still dominates: the two big classes are ~88% of archive failures,
    # and training keeps 90% of that weight on them.
    real = corruptions.TRAINING_MIX["eclipse_dim"] + corruptions.TRAINING_MIX["header_only"]
    assert real > 0.75, real
    for name in corruptions.CATALOG:
        assert corruptions.TRAINING_MIX[name] == pytest.approx(
            (1 - corruptions.COVERAGE_WEIGHT) * corruptions.REAL_MIX[name]
            + corruptions.COVERAGE_WEIGHT / len(corruptions.CATALOG))
    # But nothing is unreachable, so an unseen fault is not out of distribution.
    assert all(weight > 0.0 for weight in corruptions.TRAINING_MIX.values())
    rare = [name for name in corruptions.CATALOG
            if corruptions.REAL_MIX[name] < 0.01]
    assert sum(corruptions.TRAINING_MIX[name] for name in rare) < 0.15


def test_training_draws_follow_the_training_mixture():
    rng = np.random.default_rng(0)
    draws = [corruptions.draw_training(rng) for _ in range(4000)]
    assert draws.count("eclipse_dim") / 4000 == pytest.approx(
        corruptions.TRAINING_MIX["eclipse_dim"], abs=0.03)
    assert draws.count("header_only") / 4000 == pytest.approx(
        corruptions.TRAINING_MIX["header_only"], abs=0.03)
    assert len(set(draws)) == len(corruptions.CATALOG)      # every mode appeared


def test_evaluation_still_draws_the_unsmoothed_archive_mixture():
    """The bench measures against reality, so it must not inherit the training
    mixture's coverage smoothing."""
    rng = np.random.default_rng(1)
    draws = [corruptions.draw_real(rng) for _ in range(4000)]
    assert draws.count("eclipse_dim") / 4000 == pytest.approx(
        corruptions.REAL_MIX["eclipse_dim"], abs=0.03)


@pytest.mark.parametrize("name", ARRAY_MODES)
def test_different_seeds_stay_reproducible(name, image, donor):
    first, _ = apply(name, image, donor, seed=1)
    again, _ = apply(name, image, donor, seed=1)
    np.testing.assert_array_equal(np.nan_to_num(first), np.nan_to_num(again))


@pytest.mark.parametrize("name", FILE_MODES)
def test_file_corruptions_are_deterministic(name, tmp_path, image):
    path = write_fits(str(tmp_path / "src.fits"), image)
    raw = open(path, "rb").read()
    first = corruptions.apply_file(name, raw, 7, 1.0)
    second = corruptions.apply_file(name, raw, 7, 1.0)
    assert first == second
    assert first != raw


# ------------------------------------------------------------------ specific modes


def test_eclipse_dim_reproduces_the_real_failure(image):
    """The archive's commonest fault: radiance collapses, flags are set."""
    damaged, header = apply("eclipse_dim", image)
    assert header["DEGRADED"] is True and header["ECLIPSE"] == 2
    assert abs(float(damaged.mean())) < float(image.mean()) / 1000


def test_all_zero_sets_the_empty_flag(image):
    damaged, header = apply("all_zero", image)
    assert not damaged.any() and header["EMPTY"] is True


def test_nan_fill_at_full_severity_covers_everything(image):
    damaged, _ = apply("nan_fill", image, severity=1.0)
    assert np.isnan(damaged).all()


def test_nan_fill_at_partial_severity_is_partial(image):
    damaged, _ = apply("nan_fill", image, severity=0.3)
    fraction = np.isnan(damaged).mean()
    assert 0.1 < fraction < 0.5


def test_zblank_fill_uses_the_fits_sentinel(image):
    damaged, _ = apply("zblank_fill", image, severity=1.0)
    assert (damaged == fitsio.ZBLANK).all()


def test_frozen_returns_the_donor_exactly(image, donor):
    damaged, _ = apply("frozen", image, donor)
    np.testing.assert_array_equal(damaged, donor)


def test_torn_frame_mixes_two_observations(image, donor):
    damaged, _ = apply("torn_frame", image, donor, severity=1.0)
    assert np.array_equal(damaged[-1], donor[-1])
    assert np.array_equal(damaged[0], image[0])


def test_dropped_rows_blanks_whole_scan_lines(image):
    damaged, _ = apply("dropped_rows", image, severity=1.0)
    blank = [row for row in range(damaged.shape[0]) if not damaged[row].any()]
    assert len(blank) >= 1


def test_yaw_flip_is_a_180_degree_rotation(image):
    damaged, header = apply("yaw_flip", image)
    assert header["YAW_FLIP"] == 1
    np.testing.assert_array_equal(damaged, np.flip(np.flip(image, 0), 1))


def test_translate_moves_the_disc_and_updates_the_wcs(image):
    damaged, header = apply("translate", image, severity=1.0)
    assert "CRPIX1" in header and "CRPIX2" in header
    assert not np.array_equal(damaged, image)


def test_saturate_raises_the_ceiling(image):
    damaged, _ = apply("saturate", image, severity=1.0)
    assert float(damaged.max()) > float(image.max())


def test_modes_requiring_a_donor_say_so(image):
    for name, corruption in corruptions.CATALOG.items():
        if corruption.needs_donor:
            with pytest.raises(ValueError, match="donor"):
                corruptions.apply_array(name, image, 1, 1.0, None)


def test_truncate_keeps_at_least_one_block(tmp_path, image):
    path = write_fits(str(tmp_path / "src.fits"), image)
    raw = open(path, "rb").read()
    damaged = corruptions.apply_file("truncate", raw, 1, 1.0)
    assert fitsio.BLOCK <= len(damaged) < len(raw)


def test_drop_image_hdu_leaves_only_the_primary_header(tmp_path, image):
    path = write_fits(str(tmp_path / "src.fits"), image)
    raw = open(path, "rb").read()
    damaged = corruptions.apply_file("drop_image_hdu", raw, 1, 1.0)
    assert len(damaged) == fitsio.BLOCK
    out = str(tmp_path / "damaged.fits")
    open(out, "wb").write(damaged)
    data, error = fitsio.read_image(out)
    assert data is None and error


def test_block_corruption_leaves_the_header_readable(tmp_path, image):
    """Simulates bit rot: the header still parses, the data no longer matches it."""
    path = write_fits(str(tmp_path / "src.fits"), image)
    raw = open(path, "rb").read()
    damaged = corruptions.apply_file("block_corruption", raw, 3, 1.0)
    out = str(tmp_path / "damaged.fits")
    open(out, "wb").write(damaged)
    cards, error = fitsio.read_header(out)
    assert cards.get("WAVELNTH") == 171
    assert "datasum mismatch" in fitsio.verify_datasums(out)


# ---------------------------------------------------------------------- interface


def test_wrong_kind_is_rejected(image, tmp_path):
    with pytest.raises(ValueError, match="not an array one"):
        corruptions.apply_array("truncate", image, 1)
    with pytest.raises(ValueError, match="not a file one"):
        corruptions.apply_file("eclipse_dim", b"x" * 3000, 1)


def test_unknown_mode_raises(image):
    with pytest.raises(KeyError):
        corruptions.apply_array("does_not_exist", image, 1)


def test_recomputed_stats_match_the_damaged_pixels(image):
    stats = corruptions.recomputed_stats(image)
    assert stats["IMG_MEAN"] == pytest.approx(float(image.mean()), rel=1e-5)
    assert stats["IMG_MAX"] == pytest.approx(float(image.max()), rel=1e-5)


def test_recomputed_stats_survive_an_all_nan_frame():
    stats = corruptions.recomputed_stats(np.full((4, 4), np.nan, np.float32))
    assert stats == {"IMG_MIN": 0.0, "IMG_MAX": 0.0, "IMG_MEAN": 0.0, "IMG_SDEV": 0.0}


def test_stats_are_recomputed_only_where_the_header_would_follow():
    """A real eclipse has a matching header; bit rot does not.  See the module docs."""
    assert corruptions.CATALOG["eclipse_dim"].recompute_stats is True
    assert corruptions.CATALOG["gain_shift"].recompute_stats is True
    assert corruptions.CATALOG["block_corruption"].recompute_stats is False
    assert corruptions.CATALOG["salt_pepper"].recompute_stats is False



# ------------------------------------------------- detector faults from the real audit

DETECTOR_MODES = ("detector_bar", "noise_bar", "bright_block", "ghost_disc", "short_exposure")


@pytest.mark.parametrize("mode", DETECTOR_MODES)
def test_detector_modes_are_deterministic_and_change_the_pixels(mode):
    image = solar_disc(peak=1.0)
    a, header_a = corruptions.apply_array(mode, image, seed=5, severity=0.7)
    b, header_b = corruptions.apply_array(mode, image, seed=5, severity=0.7)
    assert np.array_equal(a, b, equal_nan=True) and header_a == header_b
    assert not np.array_equal(a, image)
    assert a.shape == image.shape and a.dtype == np.float32


@pytest.mark.parametrize("mode", DETECTOR_MODES)
def test_detector_modes_work_at_the_training_resolution(mode):
    """The sampler damages 640px shard frames, not 1280px files."""
    small = solar_disc(size=640, radius=193, peak=1.0)
    out, _ = corruptions.apply_array(mode, small, seed=1, severity=0.5)
    assert out.shape == (640, 640) and np.isfinite(out).all()


@pytest.mark.parametrize("mode,source", [
    ("detector_bar", "field_v1"), ("noise_bar", "field_v1"), ("bright_block", "field_v1"),
    ("short_exposure", "header_v1"),
])
def test_the_production_judgment_catches_each_injected_detector_fault(tmp_path, mode, source):
    """Injected the way the bench writes them: damaged pixels, header overrides,
    statistics recomputed from the damage."""
    from conftest import write_fits

    image = solar_disc(peak=1.0)
    damaged, overrides = corruptions.apply_array(mode, image, seed=3, severity=0.8)
    headers = {"REFCROTA": 18.0, **corruptions.recomputed_stats(damaged), **overrides}
    path = str(tmp_path / "d" / f"{mode}.fits")
    write_fits(path, damaged, headers=headers)
    cards, error = fitsio.read_header(path)
    got, verdict, _ = detectors.judge_frame((16, 171, 0), cards, error,
                                            lambda: (damaged, None))
    assert verdict.verdict == detectors.BAD, (mode, verdict)
    assert got == source


def test_bars_and_blocks_record_the_rotation_they_were_drawn_at():
    image = solar_disc(peak=1.0)
    for mode in ("detector_bar", "noise_bar", "bright_block"):
        _, header = corruptions.apply_array(mode, image, seed=2, severity=0.5)
        low, high = corruptions._REFCROTA_RANGE
        assert low <= header["REFCROTA"] <= high


def test_short_exposure_declares_itself_in_the_header():
    _, header = corruptions.apply_array("short_exposure", solar_disc(peak=1.0), seed=0,
                                        severity=0.3)
    assert header == {"N_LONG": 0, "NUM_IMGS": 1}



class _FixedDraw:
    """An rng whose uniform() returns queued values: pins a corruption's geometry."""

    def __init__(self, *values):
        self.values = list(values)

    def uniform(self, low, high):
        return self.values.pop(0)


def _judge(image):
    return detectors.field_v1(image, 171, {"diam_sun": 771.98}, refcrota=18.0)


def test_a_ghost_reaching_the_far_field_is_caught_by_field_v1():
    image = solar_disc(peak=1.0)
    far, _ = corruptions._ghost_disc(image, _FixedDraw(0.55, 0.7), severity=0.8)
    assert _judge(far).verdict == detectors.BAD


def test_a_ghost_overlapping_the_disc_is_invisible_to_single_frame_checks():
    """Documents the gap the clean-neighbour series check (temporal_v2) exists for:
    four real ghost limbs escaped every single-frame detector in the audit."""
    image = solar_disc(peak=1.0)
    near, _ = corruptions._ghost_disc(image, _FixedDraw(0.15, 0.7), severity=0.8)
    assert _judge(near).verdict == detectors.GOOD



def test_the_real_mixture_carries_the_measured_detector_faults():
    """Weights come from the production judgment of the whole archive (run 46):
    short-exposure static is the third-largest real failure, and the bar, block
    and ghost modes keep the proportions of the hand-confirmed defects."""
    mix = corruptions.REAL_MIX
    assert 0.05 < mix["short_exposure"] < 0.07
    ranked = sorted(mix, key=mix.get, reverse=True)
    assert ranked[:3] == ["eclipse_dim", "header_only", "short_exposure"]
    assert mix["detector_bar"] > mix["ghost_disc"] > mix["bright_block"] > 0
    assert mix["detector_bar"] / mix["bright_block"] == pytest.approx(232 / 29)


def test_container_damage_is_weighted_by_real_unreadable_frames_only():
    """Frames the series check rejects decoded fine, so none of that residue may
    land on the file modes; those carry only the measured unreadable count."""
    file_modes = [n for n, c in corruptions.CATALOG.items() if c.kind == "file"]
    assert file_modes
    assert sum(corruptions.REAL_MIX[n] for n in file_modes) < 0.002
    pixel_tail = corruptions.REAL_MIX["salt_pepper"]
    assert all(corruptions.REAL_MIX[n] < pixel_tail for n in file_modes)