noaa-goes-visualization/tests/test_render_qa.py

"""render_qa: the composite layout, the three panel measures, and the scan.

Composites here are drawn, not rendered: a black canvas with a grey disc in each
panel where merger_FITS puts the Sun, which is all the measures look at.
"""

import json
import os

import numpy as np
import pytest
from PIL import Image, ImageDraw

import render_qa as qa
from suvi import paths

T0 = 1727913600                       # 2024-10-03 00:00 UTC
STEP = paths.CADENCE


def draw_composite(path, grey=100, bar=None, static=None, title="2024-10-03 00:00"):
    """A composite with a disc per panel.  `bar` is a panel name to draw a bright
    stripe across (into its far field); `static` a panel to fill with noise."""
    image = Image.new("RGB", qa.CANVAS)
    draw = ImageDraw.Draw(image)
    rng = np.random.default_rng(0)
    for name, (x0, y0, x1, y1) in zip(qa.PANELS, qa.panel_boxes()):
        cx, cy = (x0 + x1 - 1) / 2, (y0 + y1 - 1) / 2
        radius = 384 * (x1 - x0) / qa.FRAME_W
        draw.ellipse((cx - radius, cy - radius, cx + radius, cy + radius),
                     fill=(grey, grey // 2, 0))
        if name == bar:
            draw.rectangle((x0, y0 + 4, x1 - 1, y0 + 12), fill=(250, 250, 250))
        if name == static:
            noise = rng.integers(0, 120, (y1 - y0, x1 - x0, 3), dtype=np.uint8)
            image.paste(Image.fromarray(noise), (x0, y0))
    draw.text((602, 15), title, fill=(255, 255, 255))
    image.save(path, quality=95)
    return path


# ----------------------------------------------------------------- layout


def test_the_panels_tile_the_canvas_exactly():
    covered = np.zeros(qa.CANVAS[::-1], dtype=int)
    for x0, y0, x1, y1 in qa.panel_boxes():
        covered[y0:y1, x0:x1] += 1
    assert (covered == 1).all()
    assert qa.CANVAS == (1920, 1080)


def test_bands_alternate_left_and_right_in_merger_order():
    boxes = dict(zip(qa.PANELS, qa.panel_boxes()))
    assert qa.PANELS == ("94", "131", "171", "195", "284", "304", "composite")
    assert boxes["94"][:2] == (0, 0) and boxes["131"][0] == 1920 - qa.PANEL_W
    assert boxes["304"][1] == 2 * qa.PANEL_H
    assert boxes["composite"] == (qa.PANEL_W, 0, qa.PANEL_W + qa.FRAME_W, qa.FRAME_H)


def test_the_far_field_starts_beyond_the_limb_and_skips_the_title():
    masks = dict(zip(qa.PANELS, qa.far_field_masks()))
    side = masks["94"]
    assert not side[180, 192] and not side[180, 192 + 180]       # disc and inner corona
    assert side[0, 0] and side[-1, -1]                              # corners
    composite = masks["composite"]
    assert not composite[:qa.TITLE_ROWS].any()
    assert composite[qa.TITLE_ROWS, 0] and composite[-1, -1]


# --------------------------------------------------------------- measures


def test_a_clean_composite_measures_dark_and_quiet(tmp_path):
    levels, binned = qa.measure(draw_composite(str(tmp_path / "c.jpg")))
    assert len(levels) == len(binned) == 7
    for far, noise in levels:
        assert far < 0.02 and noise < 0.002


def test_a_bar_lights_only_its_own_panels_far_field(tmp_path):
    levels, _ = qa.measure(draw_composite(str(tmp_path / "c.jpg"), bar="195"))
    far = dict(zip(qa.PANELS, (level[0] for level in levels)))
    assert far["195"] > qa.LIMITS["far_field"]
    assert all(value < 0.02 for name, value in far.items() if name != "195")


def test_static_shows_as_noise(tmp_path):
    levels, _ = qa.measure(draw_composite(str(tmp_path / "c.jpg"), static="131"))
    noise = dict(zip(qa.PANELS, (level[1] for level in levels)))
    assert noise["131"] > qa.LIMITS["noise"]
    assert noise["94"] < 0.002


def test_the_title_never_counts(tmp_path):
    """The timestamp changes every frame; it must not read as far field or spike."""
    a = qa.measure(draw_composite(str(tmp_path / "a.jpg"), title="0000-00-00 00:00"))
    b = qa.measure(draw_composite(str(tmp_path / "b.jpg"), title="9999-99-99 99:99 XXXXXX"))
    assert a[0][6][0] == b[0][6][0] < 0.02
    assert np.array_equal(a[1][6], b[1][6])


def test_unreadable_or_foreign_images_measure_as_none(tmp_path):
    broken = tmp_path / "broken.jpg"
    broken.write_bytes(b"\xff\xd8 not a jpeg")
    small = tmp_path / "small.jpg"
    Image.new("RGB", (960, 540)).save(small)
    assert qa.measure(str(broken)) is None
    assert qa.measure(str(small)) is None
    assert qa.measure(str(tmp_path / "absent.jpg")) is None


@pytest.mark.parametrize("before, now, after, expected", [
    (0.2, 0.2, 0.2, 0.0),        # still
    (0.2, 0.3, 0.4, 0.0),        # ramp: a flare brightening
    (0.2, 0.6, 0.6, 0.0),        # step: an onset
    (0.2, 0.2, 0.6, 0.0),        # step the other side
    (0.2, 0.7, 0.2, 0.5),        # one frame out of line
    (0.7, 0.2, 0.7, 0.5),        # one frame dark
    (0.2, 0.8, 0.4, 0.4),        # out of line on a ramp: |0.8-0.3| - 0.1
])
def test_spike_is_zero_for_ramps_and_steps_and_measures_outliers(before, now, after,
                                                                 expected):
    frames = [np.full((4, 4), value, dtype=np.float32) for value in (before, now, after)]
    assert qa.spike(*frames) == pytest.approx(expected, abs=1e-6)


# ------------------------------------------------------------------- scan


@pytest.fixture
def stream(tmp_path):
    """Composites at T0.. with optional oddities; returns (directory, frames)."""
    directory = tmp_path / "render" / "learned" / "goes16"
    directory.mkdir(parents=True)

    def make(count, odd=(), skip=(), broken=()):
        for index in range(count):
            if index in skip:
                continue
            path = str(directory / f"Composite-{T0 + index * STEP}.jpg")
            if index in broken:
                with open(path, "wb") as handle:
                    handle.write(b"broken")
                continue
            draw_composite(path, grey=180 if index in odd else 100)
        return str(directory), qa.stream_frames(str(directory))

    return make


def test_scan_measures_spikes_against_both_neighbours(stream):
    directory, frames = stream(6, odd={3})
    results = qa.scan(frames, sorted(frames), workers=1)
    assert results[T0]["spike"] is None and results[T0 + 5 * STEP]["spike"] is None
    spikes = {i: max(results[T0 + i * STEP]["spike"]) for i in (1, 2, 3, 4)}
    assert spikes[3] > qa.LIMITS["spike"]
    assert spikes[2] < 1e-3 and spikes[4] < 1e-3 and spikes[1] < 1e-3


def test_a_gap_leaves_its_neighbours_without_a_spike_measure(stream):
    directory, frames = stream(7, skip={3})
    results = qa.scan(frames, list(range(T0, T0 + 7 * STEP, STEP)), workers=1)
    assert T0 + 3 * STEP not in results
    assert results[T0 + 2 * STEP]["spike"] is None
    assert results[T0 + 4 * STEP]["spike"] is None
    assert results[T0 + 1 * STEP]["spike"] is not None
    assert results[T0 + 5 * STEP]["spike"] is not None


def test_an_unreadable_composite_is_skipped_and_breaks_the_window(stream):
    directory, frames = stream(5, broken={2})
    results = qa.scan(frames, sorted(frames), workers=1)
    assert T0 + 2 * STEP not in results
    assert results[T0 + 1 * STEP]["spike"] is None
    assert results[T0 + 3 * STEP]["spike"] is None


def test_scan_reads_only_the_timeline(stream):
    directory, frames = stream(6)
    results = qa.scan(frames, [T0 + STEP, T0 + 2 * STEP, T0 + 99 * STEP], workers=1)
    assert sorted(results) == [T0 + STEP, T0 + 2 * STEP]


# ------------------------------------------------------------------- flag


def entry(far=0.0, noise=0.0, spike=0.0):
    return {"far_field": [far] + [0.0] * 6, "noise": [0.0] * 6 + [noise],
            "spike": None if spike is None else [0.0, spike] + [0.0] * 5}


def test_flag_lists_only_frames_over_a_limit_worst_first():
    results = {T0: entry(), T0 + STEP: entry(far=0.3), T0 + 2 * STEP: entry(spike=0.15),
               T0 + 3 * STEP: entry(noise=0.006), T0 + 4 * STEP: entry(far=0.15)}
    flagged = qa.flag(results)
    assert [when for when, _, _ in flagged] == [T0 + 2 * STEP, T0 + STEP, T0 + 3 * STEP]
    assert flagged[0][2] == [("131", "spike", 0.15)]
    assert flagged[1][1] == pytest.approx(2.0)
    assert flagged[2][2] == [("composite", "noise", 0.006)]


def test_a_frame_without_a_spike_measure_is_judged_on_the_rest():
    flagged = qa.flag({T0: entry(spike=None), T0 + STEP: entry(far=0.5, spike=None)})
    assert [when for when, _, _ in flagged] == [T0 + STEP]


def test_every_hit_is_kept_and_the_worst_leads():
    results = {T0: {"far_field": [0.3, 0.9] + [0.0] * 5, "noise": [0.0] * 7,
                    "spike": [0.03] + [0.0] * 6}}
    (_, severity, hits), = qa.flag(results)
    assert [h[:2] for h in hits] == [("131", "far_field"), ("94", "far_field"),
                                     ("94", "spike")]
    assert severity == pytest.approx(6.0)


# ----------------------------------------------------------- output files


def test_stream_names_come_from_the_last_three_path_parts():
    assert qa.stream_name("/r/compare_v8/new-detection/learned/goes16/") == \
        "new-detection_learned_goes16"


def test_contact_sheets_hold_the_worst_frames_up_to_the_limit(stream, tmp_path):
    directory, frames = stream(5)
    flagged = [(when, 2.0, [("195", "far_field", 0.3)]) for when in sorted(frames)]
    written = qa.contact_sheets(flagged, frames, str(tmp_path / "s"), columns=2, rows=1,
                                limit=3)
    assert [os.path.basename(p) for p in written] == ["s_01.jpg", "s_02.jpg"]
    with Image.open(written[0]) as sheet:
        assert sheet.size == (2 * 480, 270)
    with Image.open(written[1]) as sheet:
        assert sheet.size == (2 * 480, 270)
    assert qa.contact_sheets([], frames, str(tmp_path / "none")) == []


def test_main_writes_a_report_per_stream_and_flags_the_bad_frame(stream, tmp_path, capsys):
    directory, frames = stream(5, odd={2})
    out = tmp_path / "qa"
    assert qa.main([directory, "--out", str(out), "--workers", "1"]) == 0
    report = json.loads((out / "render_learned_goes16.json").read_text())
    assert report["frames"] == 5 and report["limits"] == qa.LIMITS
    assert [f["time"] for f in report["flagged"]] == [T0 + 2 * STEP]
    assert report["flagged"][0]["hits"][0]["metric"] == "spike"
    assert (out / "render_learned_goes16_01.jpg").exists()
    assert "1 flagged" in capsys.readouterr().out


def test_main_honours_periods(stream, tmp_path):
    directory, frames = stream(5, odd={2})
    out = tmp_path / "qa"
    qa.main([directory, "--out", str(out), "--workers", "1",
             "--periods", "2024-10-04..2024-10-05"])
    report = json.loads((out / "render_learned_goes16.json").read_text())
    assert report["frames"] == 0 and report["flagged"] == []


def test_streams_whose_reports_would_collide_are_refused(stream, tmp_path):
    directory, _ = stream(3)
    with pytest.raises(SystemExit, match="must differ"):
        qa.main([directory, directory + "/", "--out", str(tmp_path / "qa")])


def test_an_empty_stream_is_refused(tmp_path):
    empty = tmp_path / "a" / "b" / "c"
    empty.mkdir(parents=True)
    with pytest.raises(SystemExit, match="no composites"):
        qa.main([str(empty), "--out", str(tmp_path / "qa")])