diff --git a/.gitignore b/.gitignore index 600d2d3..f0bad73 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,4 @@ -.vscode \ No newline at end of file +.vscode +.venv +__pycache__/ +*.pyc diff --git a/bench.py b/bench.py new file mode 100644 index 0000000..ac45592 --- /dev/null +++ b/bench.py @@ -0,0 +1,1298 @@ +#!/usr/bin/env python +"""Test bench for bad-frame detection and gap filling. + +Takes a stretch of archive verified to be clean, damages it in known ways, and +measures how well each detector finds the damage and each filler repairs it. + +The archive is never modified. Corrupted frames are written to a per-case overlay +directory and every read goes through :class:`suvi.cases.Overlay`, so a bench run +against live production data is safe. + +Workflow:: + + bench.py select-window --start 2024-05-05 --end 2024-05-20 + bench.py vet-window --name may2024 --start ... --end ... + bench.py make-case --window may2024 --name mixed10 --seed 42 + bench.py detect --case mixed10 + bench.py fill --case mixed10 --oracle + bench.py report --case mixed10 + +Each step records its results in the index, so steps are resumable and comparable +across runs. +""" + +import argparse +import datetime +import itertools +import json +import os +import sys +import time +from collections import defaultdict +from multiprocessing import Pool + +import numpy as np +from astropy.io import fits + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +from suvi import cases, corruptions, db, detectors, fillers, fitsio, metrics, paths, vfs + + +def parse_date(text): + """Accept YYYY-MM-DD or YYYY-MM-DDTHH:MM, always UTC.""" + for fmt in ("%Y-%m-%d", "%Y-%m-%dT%H:%M", "%Y-%m-%d %H:%M"): + try: + stamp = datetime.datetime.strptime(text, fmt) + return int(stamp.replace(tzinfo=datetime.timezone.utc).timestamp()) + except ValueError: + continue + raise argparse.ArgumentTypeError(f"Not a date: {text!r}") + + +def format_time(timestamp): + return datetime.datetime.fromtimestamp(timestamp, datetime.timezone.utc).strftime( + "%Y-%m-%d %H:%M" + ) + + +def parse_int_list(text): + return tuple(int(part) for part in text.split(",") if part.strip()) + + +# ------------------------------------------------------------------- select-window + + +def cmd_select_window(args): + root = paths.data_root() + satellites = parse_int_list(args.satellites) + wavelengths = parse_int_list(args.wavelengths) + print(f"Scanning {root} from {format_time(args.start)} to {format_time(args.end)}") + found = cases.scan_window(root, satellites, wavelengths, args.start, args.end) + total = len(cases.timeline(args.start, args.end)) + print(f" {len(found)} files across {total} slots\n") + + runs = cases.find_runs( + found, + satellites, + wavelengths, + args.start, + args.end, + minimum=args.min_slots, + require_good_label=args.require_legacy_label, + ) + if not runs: + print(f"No runs of >= {args.min_slots} slots complete on every requested band.") + print("Try fewer bands, one satellite, or a different date range.") + return 1 + + basis = "labelled good by the legacy filter" if args.require_legacy_label else "present" + print(f"Runs where every requested band and satellite is {basis}:") + print(f"{'slots':>7} {'hours':>8} window") + for length, start, end in runs[: args.limit]: + print(f"{length:>7} {length * paths.CADENCE / 3600:>8.1f} {format_time(start)} .. {format_time(end)}") + if not args.require_legacy_label: + print("\nCompleteness only -- run vet-window to verify quality from the headers.") + return 0 + + +# ---------------------------------------------------------------------- vet-window + + +def _header_vet(path): + """Independent check that a frame really is good, from its header.""" + values, error = fitsio.scan_header(path) + if error: + return False, f"header unreadable: {error}" + if values.get("empty"): + return False, "EMPTY set" + if values.get("eclipse"): + return False, f"ECLIPSE={int(values['eclipse'])}" + return True, None + + +def cmd_vet_window(args): + root = paths.data_root() + satellites = parse_int_list(args.satellites) + wavelengths = parse_int_list(args.wavelengths) + conn = db.connect(args.db) + + found = cases.scan_window(root, satellites, wavelengths, args.start, args.end) + slots = cases.timeline(args.start, args.end) + expected = len(slots) * len(satellites) * len(wavelengths) + print(f"Window has {len(found)}/{expected} frames present") + + missing = [ + (satellite, wavelength, time) + for time in slots + for satellite in satellites + for wavelength in wavelengths + if (satellite, wavelength, time) not in found + ] + if missing and not args.allow_missing: + print(f"ERROR: {len(missing)} slots have no frame at all; refusing to vet.") + print(" Pick a window from select-window, or pass --allow-missing.") + return 1 + + # Ground truth comes from the headers, which the instrument wrote and which are + # independent of any pixel filter. The legacy filename labels are recorded for + # comparison but deliberately do not decide anything: they were written in May + # 2024 by a filter that was retuned that July, so they are neither current nor + # self-consistent. Treating them as truth would bake the old filter's mistakes + # into every score computed against this window. + print("Vetting from headers...") + disagreements = [] + truth = {} + reliever = vfs.Reliever(label="vet") + for index, (slot, (path, label)) in enumerate(sorted(found.items())): + reliever.tick() + good_by_header, reason = _header_vet(path) + good_by_legacy = label == "f" + if good_by_header != good_by_legacy: + disagreements.append((slot, label, reason)) + truth[slot] = (good_by_header, "header") + if args.progress and index % 2000 == 0 and index: + print(f" {index}/{len(found)}") + + print(f" {len(disagreements)} frames where the legacy label disagrees with the header") + for slot, label, reason in disagreements[:10]: + print(f" {slot} legacy={label} header={reason or 'good'}") + + good = sum(1 for is_good, _ in truth.values() if is_good) + print(f" {good}/{len(truth)} frames vetted good") + if good < len(truth) * args.min_good_fraction: + print(f"ERROR: only {good / len(truth):.1%} vetted good; window is not clean enough.") + return 1 + + conn.execute( + """ + INSERT INTO bench_window (name, t_start, t_end, satellites, wavelengths, vetted_at, notes) + VALUES (?, ?, ?, ?, ?, ?, ?) + ON CONFLICT (name) DO UPDATE SET + t_start = excluded.t_start, t_end = excluded.t_end, + satellites = excluded.satellites, wavelengths = excluded.wavelengths, + vetted_at = excluded.vetted_at, notes = excluded.notes + """, + ( + args.name, + args.start, + args.end, + json.dumps(list(satellites)), + json.dumps(list(wavelengths)), + time.time(), + args.notes, + ), + ) + window_id = conn.execute( + "SELECT id FROM bench_window WHERE name = ?", (args.name,) + ).fetchone()["id"] + conn.execute("DELETE FROM bench_truth WHERE window_id = ?", (window_id,)) + + for slot, (path, label) in sorted(found.items()): + name = paths.parse_frame_filename(os.path.basename(path)) + relpath = name.relpath() + frame_id = db.upsert_frame( + conn, name, relpath, os.path.getsize(path), os.path.getmtime(path) + ) + is_good, source = truth[slot] + conn.execute( + "INSERT OR REPLACE INTO bench_truth (window_id, frame_id, is_good, vet_source)" + " VALUES (?, ?, ?, ?)", + (window_id, frame_id, int(is_good), source), + ) + conn.commit() + print(f"Wrote window '{args.name}' (id {window_id}) with {len(found)} frames") + + if args.contact_sheet: + _write_contact_sheet(found, args.contact_sheet, satellites, wavelengths) + print(f"Contact sheet: {args.contact_sheet}") + print("Review it before trusting this window as ground truth.") + return 0 + + +def _write_contact_sheet(found, path, satellites, wavelengths, columns=24, size=64): + """A downsampled filmstrip of the window, for the one manual review pass.""" + import matplotlib + + matplotlib.use("Agg") + from matplotlib import pyplot as plt + + satellite, wavelength = satellites[0], wavelengths[len(wavelengths) // 2] + series = sorted(slot for slot in found if slot[0] == satellite and slot[1] == wavelength) + if not series: + return + step = max(1, len(series) // (columns * 20)) + picks = series[::step] + rows = (len(picks) + columns - 1) // columns + + figure, axes = plt.subplots(rows, columns, figsize=(columns, rows)) + axes = np.atleast_2d(axes) + vmin, vmax, gamma = metrics.DISPLAY_MAPPING[wavelength] + for index, slot in enumerate(picks): + axis = axes[index // columns][index % columns] + image, error = fitsio.read_image(found[slot][0]) + if image is not None: + import cv2 as cv + + thumb = cv.resize(np.nan_to_num(image), (size, size), interpolation=cv.INTER_AREA) + axis.imshow(metrics.to_display(thumb, wavelength), cmap="gray", vmin=0, vmax=1) + axis.set_title(format_time(slot[2])[5:], fontsize=3) + axis.axis("off") + for index in range(len(picks), rows * columns): + axes[index // columns][index % columns].axis("off") + figure.suptitle(f"g{satellite} {wavelength}A - {len(picks)} of {len(series)} frames") + figure.tight_layout() + figure.savefig(path, dpi=200) + plt.close(figure) + + +# ------------------------------------------------------------------------ make-case + + +def _load_window(conn, name): + row = conn.execute("SELECT * FROM bench_window WHERE name = ?", (name,)).fetchone() + if row is None: + raise SystemExit(f"No vetted window named {name!r}; run vet-window first.") + return row + + +def _build_overlay(conn, case_row, window_row): + """Reconstruct a case's view of the archive from the index.""" + root = paths.data_root() + satellites = tuple(json.loads(window_row["satellites"])) + wavelengths = tuple(json.loads(window_row["wavelengths"])) + archive = cases.scan_window( + root, satellites, wavelengths, window_row["t_start"], window_row["t_end"] + ) + overrides, deleted = {}, set() + rows = conn.execute( + """ + SELECT f.satellite, f.wavelength, f.t_start, i.mode, i.overlay_path + FROM bench_injection i JOIN frame f ON f.id = i.frame_id + WHERE i.case_id = ? + """, + (case_row["id"],), + ).fetchall() + for row in rows: + slot = (row["satellite"], row["wavelength"], row["t_start"]) + if row["mode"] == "delete": + deleted.add(slot) + elif row["overlay_path"]: + overrides[slot] = row["overlay_path"] + return cases.Overlay(archive=archive, overrides=overrides, deleted=frozenset(deleted)) + + +def _write_corrupted(source_path, target_path, injection): + """Materialise one corrupted frame into the overlay. Returns None on success.""" + os.makedirs(os.path.dirname(target_path), exist_ok=True) + corruption = corruptions.CATALOG[injection.mode] + + if corruption.kind == "file": + with open(source_path, "rb") as handle: + raw = handle.read() + damaged = corruptions.apply_file(injection.mode, raw, injection.seed, injection.severity) + with open(target_path, "wb") as handle: + handle.write(damaged) + return None + + image, error = fitsio.read_image(source_path) + if image is None: + return f"could not read source: {error}" + + donor = None + if corruption.needs_donor: + donor = np.roll(image, image.shape[0] // 3, axis=0) # a plausible other frame + + damaged, overrides = corruptions.apply_array( + injection.mode, image, injection.seed, injection.severity, donor + ) + if corruption.recompute_stats: + overrides = {**overrides, **corruptions.recomputed_stats(damaged)} + + with fits.open(source_path) as hdus: + target = next( + (h for h in hdus if getattr(h, "data", None) is not None and h.data.ndim == 2), None + ) + if target is None: + return "source has no image HDU" + target.data = damaged.astype(np.float32) + for key, value in overrides.items(): + target.header[key] = value + hdus.writeto(target_path, overwrite=True, checksum=True) + return None + + +def _corrupt_one(job): + source_path, target_path, injection = job + try: + return injection.slot, _write_corrupted(source_path, target_path, injection) + except Exception as exc: # a corruption must never take the whole run down + return injection.slot, f"{type(exc).__name__}: {exc}" + + +def cmd_make_case(args): + conn = db.connect(args.db) + window = _load_window(conn, args.window) + satellites = tuple(json.loads(window["satellites"])) + wavelengths = tuple(json.loads(window["wavelengths"])) + root = paths.data_root() + + archive = cases.scan_window( + root, satellites, wavelengths, window["t_start"], window["t_end"] + ) + plan = cases.InjectionPlan( + fraction=args.fraction, + gap_lengths=parse_int_list(args.gap_lengths), + satellite_scope=args.satellite_scope, + wavelength_scope=args.wavelength_scope, + modes=tuple(args.modes.split(",")) if args.modes else cases.InjectionPlan.modes, + severity=(args.severity_low, args.severity_high), + gaps_per_length=args.gaps_per_length, + ).validate() + + window_times = sorted({slot[2] for slot in archive}) + injections = cases.plan_injections( + plan, set(archive), satellites, wavelengths, args.seed + ) + if not injections: + print("Plan produced no injections; raise --fraction or widen the window.") + return 1 + + overlay_dir = args.overlay_dir or os.path.join( + os.path.dirname(root), "bench", args.name + ) + os.makedirs(overlay_dir, exist_ok=True) + + conn.execute( + """ + INSERT INTO bench_case (window_id, name, seed, plan_json, overlay_dir, created_at) + VALUES (?, ?, ?, ?, ?, ?) + ON CONFLICT (name) DO UPDATE SET + window_id = excluded.window_id, seed = excluded.seed, + plan_json = excluded.plan_json, overlay_dir = excluded.overlay_dir, + created_at = excluded.created_at + """, + ( + window["id"], + args.name, + args.seed, + json.dumps(plan.as_dict(), sort_keys=True), + overlay_dir, + time.time(), + ), + ) + case_id = conn.execute( + "SELECT id FROM bench_case WHERE name = ?", (args.name,) + ).fetchone()["id"] + conn.execute("DELETE FROM bench_injection WHERE case_id = ?", (case_id,)) + + jobs, rows = [], [] + for injection in injections: + source_path = archive[injection.slot][0] + name = paths.parse_frame_filename(os.path.basename(source_path)) + frame_id = db.frame_id_by_slot(conn, *injection.slot) + if frame_id is None: + frame_id = db.upsert_frame(conn, name, name.relpath()) + overlay_path = None + if injection.mode != "delete": + overlay_path = os.path.join(overlay_dir, name.relpath()) + jobs.append((source_path, overlay_path, injection)) + rows.append((case_id, frame_id, injection, overlay_path)) + + print(f"Planned {len(injections)} injections " + f"({sum(1 for i in injections if i.mode == 'delete')} deletions, {len(jobs)} corruptions)") + + failures = {} + if jobs: + reliever = vfs.Reliever(label="make-case") + with Pool(args.workers, initializer=fitsio.quiet_astropy) as pool: + for index, (slot, error) in enumerate(pool.imap_unordered(_corrupt_one, jobs, 8)): + reliever.tick() + if error: + failures[slot] = error + if args.progress and index % 200 == 0 and index: + print(f" wrote {index}/{len(jobs)}") + if failures: + print(f"WARNING: {len(failures)} corruptions failed:") + for slot, error in list(failures.items())[:5]: + print(f" {slot}: {error}") + + conn.executemany( + """ + INSERT INTO bench_injection + (case_id, frame_id, mode, severity, params_json, overlay_path, gap_index, gap_length) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + """, + [ + ( + case_id, + frame_id, + injection.mode, + injection.severity, + json.dumps({"seed": injection.seed}), + overlay_path, + injection.gap_index, + injection.gap_length, + ) + for case_id, frame_id, injection, overlay_path in rows + if injection.slot not in failures + ], + ) + conn.commit() + + affected = len({i.slot[2] for i in injections}) + print(f"Case '{args.name}' (id {case_id}): {affected}/{len(window_times)} timestamps affected") + print(f"Overlay: {overlay_dir}") + return 0 + + +# --------------------------------------------------------------------------- detect + + +def _read_features(job): + """Load one frame's header, thumbnail, and single-frame verdicts.""" + slot, path = job + if path is None: + return slot, None, None, None, 0 + started = time.perf_counter_ns() + values, header_error = fitsio.scan_header(path) + features = detectors.FrameFeatures(slot=slot, header=values, error=header_error) + header_verdict = detectors.header_v1(features) + header_us = (time.perf_counter_ns() - started) // 1000 + + image, image_error = fitsio.read_image(path) + + started = time.perf_counter_ns() + if image is None: + geometry_verdict = detectors.Verdict("bad", f"unreadable: {image_error}", {}) + else: + geometry_verdict = detectors.geometry_v1(image, slot[1]) + geometry_us = (time.perf_counter_ns() - started) // 1000 + + started = time.perf_counter_ns() + if image is None: + disc_verdict = detectors.Verdict("bad", f"unreadable: {image_error}", {}) + else: + disc_verdict = detectors.disc_v1(image, slot[1], values) + disc_us = (time.perf_counter_ns() - started) // 1000 + + features.thumbnail = detectors.thumbnail(image) if image is not None else None + return ( + slot, + features, + (header_verdict, header_us), + (geometry_verdict, geometry_us), + (disc_verdict, disc_us), + ) + + +def cmd_detect(args): + fitsio.quiet_astropy() + conn = db.connect(args.db) + case = conn.execute("SELECT * FROM bench_case WHERE name = ?", (args.case,)).fetchone() + if case is None: + raise SystemExit(f"No case named {args.case!r}") + window = conn.execute( + "SELECT * FROM bench_window WHERE id = ?", (case["window_id"],) + ).fetchone() + satellites = tuple(json.loads(window["satellites"])) + wavelengths = tuple(json.loads(window["wavelengths"])) + overlay = _build_overlay(conn, case, window) + wanted = tuple(args.detectors.split(",")) if args.detectors else detectors.ALL_DETECTORS + + runs = {} + for name in wanted: + runs[name] = db.create_detector_run( + conn, name, {"case": args.case}, notes=f"bench case {args.case}" + ) + + results = {name: {} for name in wanted} + elapsed = defaultdict(int) + + for wavelength in wavelengths: + per_satellite = {} + for satellite in satellites: + series_slots = [ + slot for slot in overlay.series(satellite, wavelength) + if slot not in overlay.deleted + ] + jobs = [(slot, overlay.path(slot)) for slot in series_slots] + features_by_slot = {} + reliever = vfs.Reliever(label="detect") + with Pool(args.workers, initializer=fitsio.quiet_astropy) as pool: + for slot, features, header, geometry, disc in pool.imap( + _read_features, jobs, 8 + ): + reliever.tick() + if features is None: + continue + features_by_slot[slot] = features + for name, outcome in ( + ("header_v1", header), + ("geometry_v1", geometry), + ("disc_v1", disc), + ): + if name in results: + results[name][slot] = outcome[0] + elapsed[name] += outcome[1] + + ordered = [features_by_slot[s] for s in series_slots if s in features_by_slot] + per_satellite[satellite] = ordered + + if "temporal_v1" in results and ordered: + started = time.perf_counter_ns() + verdicts = detectors.temporal_v1(ordered) + elapsed["temporal_v1"] += (time.perf_counter_ns() - started) // 1000 + for features, verdict in zip(ordered, verdicts): + results["temporal_v1"][features.slot] = verdict + + if "crosssat_v1" in results and len(satellites) >= 2: + first, second = satellites[0], satellites[1] + by_time = { + sat: {f.slot[2]: f for f in per_satellite.get(sat, [])} + for sat in (first, second) + } + temporal = results.get("temporal_v1", {}) + started = time.perf_counter_ns() + for when in sorted(set(by_time[first]) | set(by_time[second])): + a, b = by_time[first].get(when), by_time[second].get(when) + verdict_a, verdict_b = detectors.crosssat_v1( + a, + b, + temporal.get(a.slot) if a else None, + temporal.get(b.slot) if b else None, + ) + if a is not None: + results["crosssat_v1"][a.slot] = verdict_a + if b is not None: + results["crosssat_v1"][b.slot] = verdict_b + elapsed["crosssat_v1"] += (time.perf_counter_ns() - started) // 1000 + print(f" {wavelength}A done") + + for name, verdicts in results.items(): + payload = [] + for slot, verdict in verdicts.items(): + frame_id = db.frame_id_by_slot(conn, *slot) + if frame_id is None: + continue + payload.append((frame_id, verdict.verdict, verdict.reason, verdict.scores, None)) + db.record_detections(conn, runs[name], payload) + conn.commit() + flagged = sum(1 for v in verdicts.values() if v.verdict == "bad") + print(f"{name:>14}: {len(verdicts)} frames, {flagged} flagged bad, " + f"{elapsed[name] / max(len(verdicts), 1):.0f} us/frame") + return 0 + + +# ----------------------------------------------------------------------------- fill + + +def _bad_slots(conn, case, source): + """Slots the fill step should reconstruct, per oracle or per detector. + + Detector runs are per case, so the run has to be looked up by this case's + config. Taking "the latest run of header_v1" would silently pick up whichever + case happened to be processed last. + """ + if source == "oracle": + rows = conn.execute( + """ + SELECT f.satellite, f.wavelength, f.t_start + FROM bench_injection i JOIN frame f ON f.id = i.frame_id + WHERE i.case_id = ? + """, + (case["id"],), + ).fetchall() + else: + run_id = db.latest_run_id(conn, source, {"case": case["name"]}) + if run_id is None: + raise SystemExit( + f"No run of {source!r} for case {case['name']!r}; run detect first." + ) + rows = conn.execute( + """ + SELECT f.satellite, f.wavelength, f.t_start + FROM detection d JOIN frame f ON f.id = d.frame_id + WHERE d.run_id = ? AND d.verdict = 'bad' + """, + (run_id,), + ).fetchall() + return {(row["satellite"], row["wavelength"], row["t_start"]) for row in rows} + + +def cmd_fill(args): + fitsio.quiet_astropy() + conn = db.connect(args.db) + case = conn.execute("SELECT * FROM bench_case WHERE name = ?", (args.case,)).fetchone() + if case is None: + raise SystemExit(f"No case named {args.case!r}") + window = conn.execute( + "SELECT * FROM bench_window WHERE id = ?", (case["window_id"],) + ).fetchone() + satellites = tuple(json.loads(window["satellites"])) + wavelengths = tuple(json.loads(window["wavelengths"])) + overlay = _build_overlay(conn, case, window) + + source = "oracle" if args.oracle else args.from_detector + if source is None: + raise SystemExit("Pass either --oracle or --from-detector NAME") + bad = _bad_slots(conn, case, source) + print(f"Filling {len(bad)} slots identified by '{source}'") + + gap_lengths = {} + for row in conn.execute( + """ + SELECT f.satellite, f.wavelength, f.t_start, i.gap_length + FROM bench_injection i JOIN frame f ON f.id = i.frame_id WHERE i.case_id = ? + """, + (case["id"],), + ): + gap_lengths[(row["satellite"], row["wavelength"], row["t_start"])] = row["gap_length"] + + wanted = tuple(args.fillers.split(",")) if args.fillers else tuple(fillers.FILLERS) + reliever = vfs.Reliever(label="fill") + scored = defaultdict(list) + # A slot whose truth frame will not read cannot be scored. Count them: silently + # dropping them yields metrics computed over whatever subset happened to work, + # with nothing in the output to say so. A degraded mount once reduced this to + # 12% coverage while the report looked entirely normal. + unreadable = 0 + attempted = 0 + payload = [] + + for wavelength in wavelengths: + truth_cache = {} + + def truth_image(slot): + if slot not in truth_cache: + path = overlay.truth_path(slot) + image, _ = fitsio.read_image(path) if path else (None, None) + truth_cache[slot] = image + return truth_cache[slot] + + for satellite in satellites: + series = overlay.series(satellite, wavelength) + good = [s for s in series if s not in bad and overlay.path(s) is not None] + good_times = [s[2] for s in good] + for slot in series: + if slot not in bad: + continue + reliever.tick() + attempted += 1 + truth = truth_image(slot) + if truth is None: + unreadable += 1 + continue + + before = _nearest(good_times, slot[2], before=True) + after = _nearest(good_times, slot[2], before=False) + counterpart_slot = next( + ( + (other, wavelength, slot[2]) + for other in satellites + if other != satellite + and (other, wavelength, slot[2]) not in bad + and overlay.path((other, wavelength, slot[2])) is not None + ), + None, + ) + + context = fillers.FillContext( + before=truth_image((satellite, wavelength, before)) if before else None, + dt_before=slot[2] - before if before else 0.0, + after=truth_image((satellite, wavelength, after)) if after else None, + dt_after=after - slot[2] if after else 0.0, + counterpart=truth_image(counterpart_slot) if counterpart_slot else None, + header=_header_for(overlay, slot), + ) + + for filler_name in wanted: + filled = fillers.FILLERS[filler_name](context) + if filled is None: + continue + score = metrics.score_fill( + filled, truth, wavelength, gap_lengths.get(slot, context.gap_frames) + ) + scored[filler_name].append(score) + frame_id = db.frame_id_by_slot(conn, *slot) + if frame_id is not None: + payload.append( + (case["id"], frame_id, filler_name, source, + json.dumps(score.as_dict())) + ) + truth_cache.clear() + print(f" {wavelength}A done") + + conn.executemany( + """ + INSERT INTO bench_fill_result (case_id, frame_id, filler, source, scores_json) + VALUES (?, ?, ?, ?, ?) + ON CONFLICT (case_id, frame_id, filler, source) + DO UPDATE SET scores_json = excluded.scores_json + """, + payload, + ) + conn.commit() + + covered = attempted - unreadable + print(f"\nScored {covered}/{attempted} slots" + + (f" -- {unreadable} could not be read" if unreadable else "")) + if attempted and covered < attempted * 0.95: + print(f" WARNING: only {covered / attempted:.0%} of damaged slots were scored.") + print(" These numbers describe that subset, not the case. A mount") + print(" that ran out of file handles is the usual cause; re-run.") + print(f"\n{'filler':>16} {'n':>6} {'rmse':>10} {'log rmse':>10} {'psnr':>8} {'ssim':>7}") + for filler_name in wanted: + summary = metrics.summarise_fills(scored.get(filler_name, [])) + if not summary: + print(f"{filler_name:>16} {'-- not applicable --':>44}") + continue + print(f"{filler_name:>16} {summary['n']:>6} {summary['rmse']:>10.4g} " + f"{summary['log_rmse']:>10.4g} {summary['psnr']:>8.2f} {summary['ssim']:>7.4f}") + return 0 + + +def _nearest(times, target, before=True): + """Closest time in `times` strictly before or after `target`.""" + if before: + candidates = [t for t in times if t < target] + return max(candidates) if candidates else None + candidates = [t for t in times if t > target] + return min(candidates) if candidates else None + + +def _header_for(overlay, slot): + path = overlay.truth_path(slot) + if path is None: + return {} + values, _ = fitsio.scan_header(path) + return values + + +# --------------------------------------------------------------------------- report + + +def cmd_report(args): + conn = db.connect(args.db, readonly=True) + case = conn.execute("SELECT * FROM bench_case WHERE name = ?", (args.case,)).fetchone() + if case is None: + raise SystemExit(f"No case named {args.case!r}") + window = conn.execute( + "SELECT * FROM bench_window WHERE id = ?", (case["window_id"],) + ).fetchone() + + plan = json.loads(case["plan_json"]) + print(f"Case '{args.case}' on window '{window['name']}' " + f"({format_time(window['t_start'])} .. {format_time(window['t_end'])})") + print(f" plan: {plan}\n") + + injected = {} + for row in conn.execute( + """ + SELECT f.satellite, f.wavelength, f.t_start, i.mode + FROM bench_injection i JOIN frame f ON f.id = i.frame_id WHERE i.case_id = ? + """, + (case["id"],), + ): + injected[(row["satellite"], row["wavelength"], row["t_start"])] = row["mode"] + + legacy_good = { + (row["satellite"], row["wavelength"], row["t_start"]) + for row in conn.execute( + """ + SELECT f.satellite, f.wavelength, f.t_start FROM bench_truth t + JOIN frame f ON f.id = t.frame_id + WHERE t.window_id = ? AND t.is_good = 1 + """, + (window["id"],), + ) + } + + print("DETECTION") + print(f"{'detector':>14} {'prec':>7} {'recall':>7} {'F1':>7} {'FPR':>8} " + f"{'unknown':>8} {'legacy-dis':>11}") + per_detector = {} + verdict_maps = {} + for name in detectors.ALL_DETECTORS: + run_id = db.latest_run_id(conn, name, {"case": args.case}) + if run_id is None: + continue + verdicts = {} + for row in conn.execute( + """ + SELECT f.satellite, f.wavelength, f.t_start, d.verdict, d.reason + FROM detection d JOIN frame f ON f.id = d.frame_id WHERE d.run_id = ? + """, + (run_id,), + ): + slot = (row["satellite"], row["wavelength"], row["t_start"]) + verdicts[slot] = detectors.Verdict(row["verdict"], row["reason"], {}) + # Deleted slots have no frame to judge, so they cannot be scored here. + scoreable = {s: v for s, v in verdicts.items() if injected.get(s) != "delete"} + score = metrics.score_detection( + scoreable, + {s: m for s, m in injected.items() if m != "delete"}, + legacy_good, + ) + per_detector[name] = score + verdict_maps[name] = verdicts + print(f"{name:>14} {score.precision:>7.3f} {score.recall:>7.3f} {score.f1:>7.3f} " + f"{score.false_positive_rate:>8.4f} {score.unknown:>8} " + f"{score.legacy_disagreements:>11}") + + if args.combinations and len(per_detector) > 1: + print("\nDETECTOR COMBINATIONS (ranked by false-positive rate, then F1)") + print(f"{'policy':>9} {'detectors':>44} {'prec':>7} {'recall':>7} {'F1':>7} {'FPR':>8}") + scoreable_injected = {s: m for s, m in injected.items() if m != "delete"} + rows = [] + names = sorted(per_detector) + for size in range(2, len(names) + 1): + for subset in itertools.combinations(names, size): + for policy in metrics.COMBINATION_POLICIES: + combined = metrics.combine_verdicts( + [verdict_maps[name] for name in subset], policy + ) + combined = { + s: v for s, v in combined.items() if injected.get(s) != "delete" + } + score = metrics.score_detection( + combined, scoreable_injected, legacy_good + ) + rows.append((score.false_positive_rate, -score.f1, policy, subset, score)) + rows.sort(key=lambda r: (r[0] if np.isfinite(r[0]) else 9e9, + r[1] if np.isfinite(r[1]) else 9e9)) + for _, _, policy, subset, score in rows[: args.combination_limit]: + label = "+".join(n.replace("_v1", "") for n in subset) + print(f"{policy:>9} {label:>44} {score.precision:>7.3f} {score.recall:>7.3f} " + f"{score.f1:>7.3f} {score.false_positive_rate:>8.4f}") + + modes = sorted({m for m in injected.values() if m != "delete"}) + if modes and per_detector: + print("\nRECALL BY CORRUPTION MODE") + header = f"{'mode':>18}" + "".join(f"{n[:11]:>13}" for n in per_detector) + print(header) + for mode in modes: + line = f"{mode:>18}" + for score in per_detector.values(): + value = score.recall_by_mode.get(mode) + line += f"{value:>13.3f}" if value is not None else f"{'-':>13}" + print(line) + + print("\nFILL") + rows = conn.execute( + "SELECT filler, source, scores_json FROM bench_fill_result WHERE case_id = ?", + (case["id"],), + ).fetchall() + if not rows: + print(" (no fill results; run `bench.py fill` first)") + else: + grouped = defaultdict(list) + for row in rows: + payload = json.loads(row["scores_json"]) + grouped[(row["source"], row["filler"])].append( + metrics.FillScore(**payload) + ) + print(f"{'source':>10} {'filler':>16} {'n':>6} {'rmse':>10} {'psnr':>8} {'ssim':>7}") + for (source, filler_name), scores in sorted(grouped.items()): + summary = metrics.summarise_fills(scores) + print(f"{source:>10} {filler_name:>16} {summary['n']:>6} " + f"{summary['rmse']:>10.4g} {summary['psnr']:>8.2f} {summary['ssim']:>7.4f}") + + print("\nFILL QUALITY BY GAP LENGTH (psnr, dB)") + gaps = sorted({s.gap_frames for scores in grouped.values() for s in scores}) + print(f"{'source/filler':>28}" + "".join(f"{g:>8}" for g in gaps)) + for (source, filler_name), scores in sorted(grouped.items()): + summary = metrics.summarise_fills(scores) + line = f"{source + '/' + filler_name:>28}" + for gap in gaps: + entry = summary["by_gap"].get(gap) + line += f"{entry['psnr']:>8.2f}" if entry else f"{'-':>8}" + print(line) + + if args.json: + output = { + "case": args.case, + "window": window["name"], + "plan": plan, + "detection": {n: s.as_dict() for n, s in per_detector.items()}, + } + with open(args.json, "w") as handle: + json.dump(output, handle, indent=2) + print(f"\nWrote {args.json}") + return 0 + + +# --------------------------------------------------------------------------- render + +#: Slots the current pipeline will carry forward from the last good frame. Beyond +#: this merger_FITS.py gives up and the timestamp produces no composite at all. +TODAY_MAX_TIME_GAP = 3 + + +def _resolve_band(overlay, slot, truth_cache, use_truth=False): + """Read one band's pixels, caching, or None if it cannot be read. + + `use_truth` reads the original archive frame rather than what the case left in + its place. The pristine variant needs it: a slot the case *deleted* has no + overlay path at all, and resolving those through the overlay silently drops them + -- which cost the ground-truth pane 221 of its timestamps, exactly the ones the + other panes have to be compared against. + """ + key = (slot, use_truth) + if key not in truth_cache: + path = overlay.truth_path(slot) if use_truth else overlay.path(slot) + image, _ = fitsio.read_image(path) if path else (None, None) + truth_cache[key] = image + return truth_cache[key] + + +def _render_one(job): + """Composite a single timestamp for one satellite. Runs in a worker.""" + import merger_FITS + + arrays, timestamp, out_path = job + try: + image = merger_FITS.composite_from_arrays(arrays, timestamp) + image.save(out_path, quality=95) + return out_path, None + except Exception as exc: + return out_path, f"{type(exc).__name__}: {exc}" + + +def cmd_render(args): + """Produce the composite stream for one variant of a case. + + Three variants, which together answer 'is the repair visible, and is it better + than what we do now': + + * ``pristine`` -- the archive untouched. Ground truth. + * ``today`` -- what the current pipeline does with the injected damage: carry + the last good frame across gaps of up to TODAY_MAX_TIME_GAP slots, and emit + *nothing at all* for longer ones, which is what leaves ffmpeg_video.py padding + with black frames. A 300-slot gap produces 300 missing composites. + * ``new`` -- the chosen method: any(disc_v1 + header_v1) detection, filled + with optical_flow at every gap length. + """ + fitsio.quiet_astropy() + conn = db.connect(args.db) + case = conn.execute("SELECT * FROM bench_case WHERE name = ?", (args.case,)).fetchone() + if case is None: + raise SystemExit(f"No case named {args.case!r}") + window = conn.execute( + "SELECT * FROM bench_window WHERE id = ?", (case["window_id"],) + ).fetchone() + satellites = tuple(json.loads(window["satellites"])) + if args.satellites: + wanted = parse_int_list(args.satellites) + satellites = tuple(s for s in satellites if s in wanted) + if not satellites: + raise SystemExit(f"none of {wanted} are in this window") + wavelengths = tuple(json.loads(window["wavelengths"])) + if len(wavelengths) != 6: + raise SystemExit(f"a composite needs all six bands; window has {wavelengths}") + overlay = _build_overlay(conn, case, window) + + # Pristine shows the window as it really is, so it reads the archive directly and + # treats nothing as damaged. + pristine = args.variant == "pristine" + bad = set() if pristine else _bad_slots(conn, case, args.detector) + times = sorted({slot[2] for slot in overlay.archive}) + print(f"{args.variant}: {len(times)} timestamps x {len(satellites)} satellites, " + f"{len(bad)} slots treated as bad") + + for satellite in satellites: + out_dir = os.path.join(args.out, args.variant, f"goes{satellite}") + os.makedirs(out_dir, exist_ok=True) + truth_cache = {} + jobs = [] + skipped = 0 + reliever = vfs.Reliever(label=f"render g{satellite}") + + for when in times: + reliever.tick(len(wavelengths)) + slots = [(satellite, band, when) for band in wavelengths] + arrays = [] + usable = True + for slot in slots: + if not pristine and (slot in bad or overlay.path(slot) is None): + replacement, ok = _replacement_band( + args.variant, overlay, slot, bad, truth_cache, wavelengths + ) + if not ok: + usable = False + break + arrays.append(replacement) + else: + image = _resolve_band(overlay, slot, truth_cache, use_truth=pristine) + if image is None: + usable = False + break + arrays.append(image) + + if not usable: + # No composite for this timestamp -- exactly what the current + # pipeline does past its gap limit, and what ffmpeg then pads. + skipped += 1 + continue + jobs.append((arrays, when, os.path.join(out_dir, f"Composite-{int(when)}.jpg"))) + + if len(jobs) >= args.batch: + _flush_renders(jobs, args.workers) + jobs.clear() + truth_cache.clear() + print(f" g{satellite}: {len(os.listdir(out_dir))} rendered, {skipped} skipped") + + if jobs: + _flush_renders(jobs, args.workers) + print(f" g{satellite}: {len(os.listdir(out_dir))} composites, {skipped} timestamps skipped") + conn.close() + return 0 + + +def _replacement_band(variant, overlay, slot, bad, truth_cache, wavelengths): + """What to substitute for a band the case removed. Returns (array, usable).""" + satellite, band, when = slot + series = [s for s in overlay.series(satellite, band) if s not in bad] + times = [s[2] for s in series] + before = _nearest(times, when, before=True) + after = _nearest(times, when, before=False) + + if variant == "today": + # merger_FITS carries the last good frame a few slots, then gives up. + if before is None or (when - before) // paths.CADENCE > TODAY_MAX_TIME_GAP: + return None, False + return _resolve_band(overlay, (satellite, band, before), truth_cache), True + + context = fillers.FillContext( + before=_resolve_band(overlay, (satellite, band, before), truth_cache) if before else None, + dt_before=when - before if before else 0.0, + after=_resolve_band(overlay, (satellite, band, after), truth_cache) if after else None, + dt_after=after - when if after else 0.0, + ) + filled = fillers.optical_flow(context) + return (filled, True) if filled is not None else (None, False) + + +def _flush_renders(jobs, workers): + failures = 0 + with Pool(workers, initializer=fitsio.quiet_astropy) as pool: + for _, error in pool.imap_unordered(_render_one, jobs, 4): + failures += bool(error) + if failures: + print(f" WARNING: {failures} composites failed to render") + + +# ------------------------------------------------------------------------- examples + + +def cmd_examples(args): + """Show the frames a detector got wrong, so a number becomes a picture. + + A precision/recall table says how often a detector is wrong; it does not say + what it is wrong *about*. For `header_v1`, whose recall is low but precision + perfect, the interesting set is the false negatives -- damage that leaves + IMG_MEAN, ECLIPSE and EMPTY all looking normal. + """ + import matplotlib + + matplotlib.use("Agg") + from matplotlib import pyplot as plt + + fitsio.quiet_astropy() + conn = db.connect(args.db, readonly=True) + case = conn.execute("SELECT * FROM bench_case WHERE name = ?", (args.case,)).fetchone() + if case is None: + raise SystemExit(f"No case named {args.case!r}") + window = conn.execute( + "SELECT * FROM bench_window WHERE id = ?", (case["window_id"],) + ).fetchone() + overlay = _build_overlay(conn, case, window) + + run_id = db.latest_run_id(conn, args.detector, {"case": args.case}) + if run_id is None: + raise SystemExit( + f"No run of {args.detector!r} for case {args.case!r}; run detect first." + ) + + injected = {} + for row in conn.execute( + """ + SELECT f.satellite, f.wavelength, f.t_start, i.mode, i.severity + FROM bench_injection i JOIN frame f ON f.id = i.frame_id WHERE i.case_id = ? + """, + (case["id"],), + ): + injected[(row["satellite"], row["wavelength"], row["t_start"])] = ( + row["mode"], row["severity"] + ) + + wanted = [] + for row in conn.execute( + """ + SELECT f.satellite, f.wavelength, f.t_start, d.verdict, d.reason, d.scores_json + FROM detection d JOIN frame f ON f.id = d.frame_id WHERE d.run_id = ? + """, + (run_id,), + ): + slot = (row["satellite"], row["wavelength"], row["t_start"]) + mode = injected.get(slot, (None, None))[0] + if mode == "delete": + continue # no frame to show + corrupted = mode is not None + if args.kind == "fn" and corrupted and row["verdict"] == "good": + wanted.append((slot, mode, row)) + elif args.kind == "fp" and not corrupted and row["verdict"] == "bad": + wanted.append((slot, None, row)) + + print(f"{args.detector} {args.kind.upper()}s on case '{args.case}': {len(wanted)}") + if not wanted: + print(" none -- nothing to plot.") + return 0 + + by_mode = defaultdict(int) + for _, mode, _ in wanted: + by_mode[mode or "(uninjected)"] += 1 + for mode, count in sorted(by_mode.items(), key=lambda kv: -kv[1]): + print(f" {count:>4} {mode}") + + wanted = wanted[: args.limit] + columns = min(6, len(wanted)) + rows = (len(wanted) + columns - 1) // columns + figure, axes = plt.subplots(rows, columns, figsize=(3 * columns, 3.4 * rows)) + axes = np.atleast_2d(axes) + for index, (slot, mode, row) in enumerate(wanted): + axis = axes[index // columns][index % columns] + path = overlay.path(slot) + image, _ = fitsio.read_image(path) if path else (None, None) + if image is not None: + axis.imshow(metrics.to_display(image, slot[1]), cmap="gray", vmin=0, vmax=1) + scores = json.loads(row["scores_json"] or "{}") + shown = ", ".join(f"{k}={v:.3g}" for k, v in list(scores.items())[:3]) + axis.set_title( + f"g{slot[0]} {slot[1]}A {format_time(slot[2])[5:]}\n" + f"{mode or 'clean'} -> {row['verdict']}\n{shown}", + fontsize=6, + ) + axis.axis("off") + for index in range(len(wanted), rows * columns): + axes[index // columns][index % columns].axis("off") + figure.suptitle(f"{args.detector} {args.kind.upper()}s - case {args.case}") + figure.tight_layout() + figure.savefig(args.out, dpi=150) + plt.close(figure) + print(f" wrote {args.out}") + return 0 + + +# ------------------------------------------------------------------------------ cli + + +def build_parser(): + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--db", default=None, help="SQLite index (default: $SUVI_DB)") + subparsers = parser.add_subparsers(dest="command", required=True) + + def add_window_args(sub): + sub.add_argument("--start", type=parse_date, required=True) + sub.add_argument("--end", type=parse_date, required=True) + sub.add_argument("--satellites", default="16,18") + sub.add_argument("--wavelengths", default="94,131,171,195,284,304") + + select = subparsers.add_parser("select-window", help="find clean stretches of archive") + add_window_args(select) + select.add_argument("--min-slots", type=int, default=120) + select.add_argument("--limit", type=int, default=10) + select.add_argument("--require-legacy-label", action="store_true", + help="also require the stale _f suffix (off by default)") + select.set_defaults(func=cmd_select_window) + + vet = subparsers.add_parser("vet-window", help="verify and freeze a ground-truth window") + add_window_args(vet) + vet.add_argument("--name", required=True) + vet.add_argument("--notes", default=None) + vet.add_argument("--allow-missing", action="store_true") + vet.add_argument("--min-good-fraction", type=float, default=0.98) + vet.add_argument("--contact-sheet", default=None, help="write a filmstrip PNG for review") + vet.add_argument("--progress", action="store_true") + vet.set_defaults(func=cmd_vet_window) + + make = subparsers.add_parser("make-case", help="inject damage into a vetted window") + make.add_argument("--window", required=True) + make.add_argument("--name", required=True) + make.add_argument("--seed", type=int, default=42) + make.add_argument("--fraction", type=float, default=0.10) + make.add_argument("--gap-lengths", default="1,2,3,5,10,30,60") + make.add_argument("--gaps-per-length", type=int, default=None, + help="place exactly this many gaps at each length, instead of " + "sizing the case by --fraction") + make.add_argument("--satellite-scope", default="mixed", + choices=("g16", "g18", "both", "mixed")) + make.add_argument("--wavelength-scope", default="all", choices=("all", "one")) + make.add_argument("--modes", default=None, help="comma-separated; default is all") + make.add_argument("--severity-low", type=float, default=0.5) + make.add_argument("--severity-high", type=float, default=1.0) + make.add_argument("--overlay-dir", default=None) + make.add_argument("--workers", type=int, default=8) + make.add_argument("--progress", action="store_true") + make.set_defaults(func=cmd_make_case) + + detect = subparsers.add_parser("detect", help="run detectors over a case") + detect.add_argument("--case", required=True) + detect.add_argument("--detectors", default=None, help="comma-separated; default is all") + detect.add_argument("--workers", type=int, default=8) + detect.set_defaults(func=cmd_detect) + + fill = subparsers.add_parser("fill", help="reconstruct flagged slots and score them") + fill.add_argument("--case", required=True) + fill.add_argument("--fillers", default=None, help="comma-separated; default is all") + group = fill.add_mutually_exclusive_group(required=True) + group.add_argument("--oracle", action="store_true", + help="use perfect knowledge of which slots are bad") + group.add_argument("--from-detector", default=None, + help="use a detector's output, showing end-to-end error") + fill.set_defaults(func=cmd_fill) + + report = subparsers.add_parser("report", help="summarise a case's results") + report.add_argument("--case", required=True) + report.add_argument("--json", default=None) + report.add_argument("--combinations", action="store_true", + help="also score every combination of detectors") + report.add_argument("--combination-limit", type=int, default=15) + report.set_defaults(func=cmd_report) + + render = subparsers.add_parser( + "render", help="produce the composite stream for one variant of a case" + ) + render.add_argument("--case", required=True) + render.add_argument("--variant", required=True, + choices=("pristine", "today", "new")) + render.add_argument("--out", required=True, help="root directory for the streams") + render.add_argument("--detector", default="oracle", + help="which slots to treat as bad ('oracle', or a run name)") + render.add_argument("--satellites", default=None, + help="restrict to these satellites (default: all in the window)") + render.add_argument("--workers", type=int, default=8) + render.add_argument("--batch", type=int, default=48, + help="composites per worker dispatch; bounds peak memory") + render.set_defaults(func=cmd_render) + + examples = subparsers.add_parser( + "examples", help="plot the frames a detector got wrong" + ) + examples.add_argument("--case", required=True) + examples.add_argument("--detector", required=True) + examples.add_argument("--kind", choices=("fp", "fn"), required=True) + examples.add_argument("--limit", type=int, default=12) + examples.add_argument("--out", default="examples.png") + examples.set_defaults(func=cmd_examples) + return parser + + +def main(argv=None): + args = build_parser().parse_args(argv) + return args.func(args) or 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/ffmpeg_video.py b/ffmpeg_video.py index 3c61788..a4a7681 100644 --- a/ffmpeg_video.py +++ b/ffmpeg_video.py @@ -1,73 +1,126 @@ -import subprocess -import os -import shutil -import datetime -import time -import calendar - -import numpy as np -from sortedcontainers import SortedDict -import tqdm -from PIL import Image - -years = [2024] -sources = ["goes16", "goes18"] -ffmpeg_path = r"..\ffmpeg.exe" - -for year in years: - for source in sources: - path_to_images = f"..\\composite\\{source}" - output_file = f"..\\{source}_{year}_nofilt.mp4" - interp_file = f"..\\{source}_{year}_interp_nofilt.mp4" - - starttime = calendar.timegm(datetime.datetime(year, 1, 1, tzinfo=datetime.timezone.utc).timetuple()) - stoptime = calendar.timegm(datetime.datetime(year+1, 1, 1, tzinfo=datetime.timezone.utc).timetuple()) - min_file_size = 350000 # Detect and remove corrupted files by filtering by file size - max_file_size = 450000 - encoding_crf = 16 - max_frame_interp = 120 - - # Set up ffmpeg to stream images from an input pipe - command_line = f'{ffmpeg_path} -y -f image2pipe -framerate 60 -i - -c:v libx264 -crf {encoding_crf} -preset veryfast {output_file}' - print(command_line) - p = subprocess.Popen(command_line, shell=True, stdout=subprocess.PIPE, stdin=subprocess.PIPE) - - black_image = Image.fromarray(np.zeros((1080, 1920, 3), dtype='uint8')) - - files_by_time = SortedDict() - for root, dirs, files in tqdm.tqdm(os.walk(path_to_images), desc="Finding and filtering files."): - for f in files: - if f.endswith('.jpg'): - fpath = os.path.join(root, f) - fsize = os.path.getsize(fpath) - time_chunk = f.replace('-', '.').replace('_', '.').split('.')[1] - ftime = int(time_chunk) - if (fsize > min_file_size) and (fsize < max_file_size) and (ftime > starttime) and (ftime < stoptime): - files_by_time[ftime] = fpath - - difftimes = np.diff(files_by_time.keys()) - unique, counts = np.unique(difftimes, return_counts=True) - interval = unique[0] - assert np.sum((unique % interval) > 0) == 0 # Ensure all our timestamps align perfectly with our interval - - prevtime = files_by_time.peekitem(0)[0] - interval - for t, f in tqdm.tqdm(files_by_time.items(), desc="FFMPEG encoding images", total = len(files_by_time)): - framejump = (t - prevtime) // interval - if framejump < max_frame_interp: - for i in range(framejump - 1): # Fill black frames to fill gaps, these will later be interpolated with ffmpeg - black_image.save(p.stdin, 'jpeg', quality = 95) - else: # Skip past intervals that are too large to fill reasonably - print(f"Detected a frame gap of: {framejump}!") - with open(f, 'rb') as fh: - p.stdin.write(fh.read()) - prevtime = t - - p.stdin.close() # Close the ffmpeg input pipe - p.wait() # Wait for ffmpeg to finish encoding - - command_line = f'{ffmpeg_path} -y -i {output_file} -vf blackframe=0,metadata=select:key=lavfi.blackframe.pblack:value=95:function=less,minterpolate=mi_mode=mci:mc_mode=aobmc:me_mode=bidir:vsbmc=1:me=fss -c:a copy -c:v libx264 -crf {encoding_crf} -preset veryfast {interp_file}' - print(command_line) - - pipe = subprocess.Popen(command_line, shell=True, stdout=subprocess.PIPE).stdout - output = pipe.read().decode() - pipe.close() \ No newline at end of file +import argparse +import subprocess +import os +import sys +import datetime +import shutil as _shutil +import calendar + +import numpy as np +from sortedcontainers import SortedDict +import tqdm +from PIL import Image + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +from suvi import paths, vfs + + +def find_ffmpeg(): + """Locate ffmpeg: $FFMPEG first, then PATH, then the bundled Windows binary.""" + override = os.environ.get("FFMPEG") + if override and os.path.exists(override): + return override + found = _shutil.which("ffmpeg") + if found: + return found + bundled = os.path.join(os.path.dirname(paths.data_root()), "ffmpeg.exe") + if os.path.exists(bundled): + return bundled + raise SystemExit("ffmpeg not found. Install it, or set $FFMPEG to its path.") + + +def encode_videos(args): + """Encode composite JPGs into per-satellite videos. + + Gaps are padded with black frames and smoothed by ffmpeg's minterpolate, except + where the gap exceeds `max_frame_interp`, which is skipped entirely. That + behaviour is what the bench's "today" variant reproduces. + """ + years = args.years + sources = args.sources + ffmpeg_path = find_ffmpeg() + project_root = args.out or os.path.dirname(paths.data_root()) + os.makedirs(project_root, exist_ok=True) + + for year in years: + for source in sources: + path_to_images = args.images or os.path.join( + os.path.dirname(paths.data_root()), "composite", source) + output_file = os.path.join(project_root, f"{source}_{year}_{args.suffix}.mp4") + interp_file = os.path.join(project_root, f"{source}_{year}_interp_{args.suffix}.mp4") + + starttime = calendar.timegm(datetime.datetime(year, 1, 1, tzinfo=datetime.timezone.utc).timetuple()) + stoptime = calendar.timegm(datetime.datetime(year+1, 1, 1, tzinfo=datetime.timezone.utc).timetuple()) + min_file_size = 350000 # Detect and remove corrupted files by filtering by file size + max_file_size = 450000 + encoding_crf = 16 + max_frame_interp = 120 + + # Set up ffmpeg to stream images from an input pipe + command_line = f'{ffmpeg_path} -y -f image2pipe -framerate 60 -i - -c:v libx264 -crf {encoding_crf} -preset veryfast {output_file}' + print(command_line) + p = subprocess.Popen(command_line, shell=True, stdout=subprocess.PIPE, stdin=subprocess.PIPE) + + black_image = Image.fromarray(np.zeros((1080, 1920, 3), dtype='uint8')) + + files_by_time = SortedDict() + for root, dirs, files in tqdm.tqdm(os.walk(path_to_images), desc="Finding and filtering files."): + for f in files: + if f.endswith('.jpg'): + fpath = os.path.join(root, f) + fsize = os.path.getsize(fpath) + time_chunk = f.replace('-', '.').replace('_', '.').split('.')[1] + ftime = int(time_chunk) + if (fsize > min_file_size) and (fsize < max_file_size) and (ftime > starttime) and (ftime < stoptime): + files_by_time[ftime] = fpath + + difftimes = np.diff(files_by_time.keys()) + unique, counts = np.unique(difftimes, return_counts=True) + interval = unique[0] + assert np.sum((unique % interval) > 0) == 0 # Ensure all our timestamps align perfectly with our interval + + prevtime = files_by_time.peekitem(0)[0] - interval + reliever = vfs.Reliever(label=f"encode {source}") + for t, f in tqdm.tqdm(files_by_time.items(), desc="FFMPEG encoding images", total = len(files_by_time)): + reliever.tick() + framejump = (t - prevtime) // interval + if framejump < max_frame_interp: + for i in range(framejump - 1): # Fill black frames to fill gaps, these will later be interpolated with ffmpeg + black_image.save(p.stdin, 'jpeg', quality = 95) + else: # Skip past intervals that are too large to fill reasonably + print(f"Detected a frame gap of: {framejump}!") + with open(f, 'rb') as fh: + p.stdin.write(fh.read()) + prevtime = t + + p.stdin.close() # Close the ffmpeg input pipe + p.wait() # Wait for ffmpeg to finish encoding + + command_line = f'{ffmpeg_path} -y -i {output_file} -vf blackframe=0,metadata=select:key=lavfi.blackframe.pblack:value=95:function=less,minterpolate=mi_mode=mci:mc_mode=aobmc:me_mode=bidir:vsbmc=1:me=fss -c:a copy -c:v libx264 -crf {encoding_crf} -preset veryfast {interp_file}' + print(command_line) + + pipe = subprocess.Popen(command_line, shell=True, stdout=subprocess.PIPE).stdout + pipe.read() + pipe.close() + +def build_parser(): + parser = argparse.ArgumentParser( + description="Encode SUVI composite images into videos.", + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument("--years", type=int, nargs="+", default=[2024]) + parser.add_argument("--sources", nargs="+", default=["goes16", "goes18"]) + parser.add_argument("--images", default=None, + help="directory of composites (default: /composite/)") + parser.add_argument("--out", default=None, help="directory for the video files") + parser.add_argument("--start", type=int, default=None, help="Unix time, inclusive") + parser.add_argument("--stop", type=int, default=None, help="Unix time, exclusive") + parser.add_argument("--suffix", default="nofilt", help="tag in the output filenames") + return parser + + +if __name__ == "__main__": + # Guarded: this module is imported for find_ffmpeg(), and without the guard that + # import kicked off a full-year encode as a side effect. + sys.exit(encode_videos(build_parser().parse_args()) or 0) diff --git a/filter_FITS.py b/filter_FITS.py index 742ebf1..b981dab 100644 --- a/filter_FITS.py +++ b/filter_FITS.py @@ -1,281 +1,733 @@ -import os -from collections import defaultdict -from multiprocessing import Queue, Process -import re -import traceback -import datetime -import calendar - -import tqdm -import numpy as np -from matplotlib import pyplot as plt -from astropy.io import fits -from skimage.transform import hough_circle, hough_circle_peaks -from skimage.feature import canny -from skimage.filters import gaussian -from skimage.morphology import skeletonize -import cv2 as cv - -def lowpriority(): - """ Set the priority of the process to lowest possible.""" - - import sys - try: - sys.getwindowsversion() - except AttributeError: - isWindows = False - else: - isWindows = True - - if isWindows: - import win32api,win32process,win32con # pywin32 - - pid = win32api.GetCurrentProcessId() - phandle = win32api.OpenProcess(win32con.PROCESS_ALL_ACCESS, True, pid) - win32process.SetPriorityClass(phandle, win32process.THREAD_PRIORITY_LOWEST) - # win32process.SetPriorityClass(phandle, win32process.IDLE_PRIORITY_CLASS) - # tid = win32api.GetCurrentThreadId() - # thandle = win32api.OpenThread(win32con.PROCESS_ALL_ACCESS, True, tid) - # win32process.SetThreadPriority(thandle, win32process.THREAD_MODE_BACKGROUND_BEGIN) - else: - import os - - os.nice(19) - -measurement_names = ["094", "131", "171", "195", "284", "304"] -thresholds = [0.050, 0.10 , 1.00, 1.40, 1.0, 2.50] -max_center_skew = 7 -expected_dims = 1280 -half_dims = expected_dims // 2 -valid_radii = [383, 394] -avg_radius = (valid_radii[0] + valid_radii[1]) // 2 -max_radius_error = 80 - -ratio_above_thresh_max = 0.4 -ratio_above_thresh_min = 0.07 - -def filter_fits(work_queue): - lowpriority() - idealcircle = np.zeros((expected_dims // 2, expected_dims // 2)) - cv.circle(idealcircle, (expected_dims // 4, expected_dims // 4), avg_radius // 2, 1, -1) - idealcircle_axis = np.average(idealcircle, 0) - while True: - try: - job = work_queue.get() - if job is None: - return - measurement = job.split("dr_suvi-l2-")[1].split("_")[0][2:] - idx = measurement_names.index(measurement) - try: - data = fits.getdata(job) - except IndexError: # This can happen with a blank or corrupt HDU in the .fits file - new_name = job.split(".fits")[0] + "_e.fits" - os.rename(job, new_name) - continue - assert data.shape[0] == expected_dims - assert data.shape[1] == expected_dims - data = cv.resize(data, dsize=(expected_dims // 2, expected_dims // 2), interpolation=cv.INTER_LINEAR) - filtered_data = np.copy(data) - above_thresh_indexes = data > thresholds[idx] - filtered_data[above_thresh_indexes] = thresholds[idx] - filtered_data /= thresholds[idx] - ratio_above_thresh = np.count_nonzero(above_thresh_indexes) / data.shape[0] / data.shape[1] - - if ratio_above_thresh > ratio_above_thresh_max: - print(f"Exceeded ratio_above_thresh_max, possible data corruption in file: {job}") - new_name = job.split(".fits")[0] + "_e.fits" - os.rename(job, new_name) - - plt.figure(figsize=[10.24, 7.68]) - plt.title(f"Exceeded ratio_above_thresh_max [{ratio_above_thresh:0.2f}]") - plt.imshow(filtered_data, cmap='jet') - plt.savefig(job.split(".fits")[0] + "_e.jpg") - plt.close('all') - - continue - - if ratio_above_thresh < ratio_above_thresh_min: - print(f"Below ratio_above_thresh_min, possible bad data in file: {job}") - new_name = job.split(".fits")[0] + "_e.fits" - os.rename(job, new_name) - - plt.figure(figsize=[10.24, 7.68]) - plt.title(f"Below ratio_above_thresh_min [{ratio_above_thresh:0.2f}]") - plt.imshow(filtered_data, cmap='jet') - plt.savefig(job.split(".fits")[0] + "_e.jpg") - plt.close('all') - - continue - - # Decide whether the data is valid based on solar disc placement and levels. - xavg = np.average(filtered_data, 0) - yavg = np.average(filtered_data, 1) - axis_indices = list(range(half_dims)) - - # Center Estimation - weighted_avg_center_x = np.average(axis_indices, 0, xavg) - weighted_avg_center_y = np.average(axis_indices, 0, yavg) - - good_center = True - if (half_dims // 2 - weighted_avg_center_x > max_center_skew) or (half_dims // 2 - weighted_avg_center_y > max_center_skew): - print(f"Could not find centered solar disc in file: {job}") - good_center = False - - # Radius Estimation - edge_thresh = 0.98 - high_edge_x = np.argmax(np.cumsum(xavg) > (np.sum(xavg) * edge_thresh)) - low_edge_x = len(xavg) - np.argmax(np.cumsum(np.flip(xavg)) > (np.sum(xavg) * edge_thresh)) - high_edge_y = np.argmax(np.cumsum(yavg) > (np.sum(yavg) * edge_thresh)) - low_edge_y = len(yavg) - np.argmax(np.cumsum(np.flip(yavg)) > (np.sum(yavg) * edge_thresh)) - - good_radius = True - calc_radius_x = (high_edge_x - low_edge_x) # We do not divide by two because our image is downsized by half in each dimension - calc_radius_y = (high_edge_y - low_edge_y) - calc_radius = (calc_radius_x + calc_radius_y) / 2.0 - if abs(calc_radius - avg_radius) > max_radius_error: - print(f"Could not find correct solar radius in file: {job}") - good_radius = False - - # Goodness of fit estimation vs perfect disc - gof_x = np.sum(np.abs(xavg - idealcircle_axis)) / half_dims - gof_y = np.sum(np.abs(yavg - idealcircle_axis)) / half_dims - - good_fit = True - if gof_x > 0.2 or gof_y > 0.2: - print(f"Could not find valid solar disc in file: {job}") - good_fit = False - - - if good_center and good_radius and good_fit: - # We have validated this file, rename it appropriately - new_name = job.split(".fits")[0] + "_f.fits" - os.rename(job, new_name) - else: - new_name = job.split(".fits")[0] + "_e.fits" - os.rename(job, new_name) - - # Plot Results - fig, axes = plt.subplots(2, 2, width_ratios=(0.2, 1), height_ratios=(0.2, 1), gridspec_kw={"hspace":0.0, "wspace": 0.0}, figsize = [8, 8]) - - axes[0][0].axis('off') - - axes[0][1].plot(axis_indices, xavg) - axes[0][1].plot(axis_indices, idealcircle_axis) - # axes[0][1].fill_between(axis_indices, xavg, idealcircle_axis, hatch="//", edgecolor="red", facecolor="none") - - axes[1][0].plot(yavg, axis_indices) - axes[1][0].plot(idealcircle_axis, axis_indices) - axes[1][0].invert_xaxis() - axes[1][0].invert_yaxis() - - axes[1][1].imshow(filtered_data, aspect='auto', vmin=0, vmax=1) - axes[1][1].vlines([low_edge_x, high_edge_x], 0, half_dims) - axes[1][1].hlines([low_edge_y, high_edge_y], 0, half_dims) - axes[1][1].scatter(half_dims // 2, half_dims // 2, marker="o", linewidths = 1, alpha = 0.75, edgecolors="black", facecolors = "none") - axes[1][1].scatter(weighted_avg_center_x, weighted_avg_center_y, marker = "x", alpha = 0.75, facecolors='red') - - # Set shared x axis between imshow and xdata plots - axes[0][1].xaxis.set_ticks_position("top") - axes[0][1].yaxis.set_ticks_position("right") - axes[0][1].sharex(axes[1][1]) - axes[1][0].sharey(axes[1][1]) - axes[1][1].yaxis.set_ticks_position("right") - - plt.savefig(job.split(".fits")[0] + "_e.jpg") - plt.close('all') - - - except KeyboardInterrupt: - return - except Exception as e: - print(f"Error on file: {job} - {e}") - traceback.print_exception(e) - - -if __name__ == "__main__": - stored_fits_dirs = [r"..\Data\goes16\l2\data\suvi-l2-ci094\2024", - r"..\Data\goes16\l2\data\suvi-l2-ci131\2024", - r"..\Data\goes16\l2\data\suvi-l2-ci171\2024", - r"..\Data\goes16\l2\data\suvi-l2-ci195\2024", - r"..\Data\goes16\l2\data\suvi-l2-ci284\2024", - r"..\Data\goes16\l2\data\suvi-l2-ci304\2024", - r"..\Data\goes18\l2\data\suvi-l2-ci094\2024", - r"..\Data\goes18\l2\data\suvi-l2-ci131\2024", - r"..\Data\goes18\l2\data\suvi-l2-ci171\2024", - r"..\Data\goes18\l2\data\suvi-l2-ci195\2024", - r"..\Data\goes18\l2\data\suvi-l2-ci284\2024", - r"..\Data\goes18\l2\data\suvi-l2-ci304\2024",] - - # stored_fits_dirs = [r"..\Data\goes18\l2\data\suvi-l2-ci284\2023\01\10"] - - starttime = calendar.timegm(datetime.datetime(2024, 1, 1, tzinfo=datetime.timezone.utc).timetuple()) - stoptime = calendar.timegm(datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc).timetuple()) - - reprocess_errors = True - nworkers = 16 - - regex_filename = r"dr_suvi-l2-ci\d{3}_g(16|18)_s\S*\.fits$" - - work_queue = Queue(maxsize = nworkers) - workers = [] - for i in range(nworkers): - p = Process(target = filter_fits, args = (work_queue,), daemon=True) - p.start() - workers.append(p) - - lowpriority() - files_by_timestamp = defaultdict(list) - try: - for stored_fits_dir in stored_fits_dirs: - filename_tester = re.compile(regex_filename) - - found_files = 0 - print(f"Searching for FITS files in: {stored_fits_dir}") - for root, dirs, files in tqdm.tqdm(os.walk(stored_fits_dir), desc="Searching"): - for f in files: - if filename_tester.match(f): - file_parts = f.split("_") - file_name_end = file_parts[-1].split(".")[0] - measure_end_time = int(datetime.datetime.strptime(file_parts[4][1:16] + " +0000", "%Y%m%dT%H%M%S %z").timestamp()) - if (measure_end_time >= starttime) and (measure_end_time < stoptime): - if file_name_end == "f": - if file_parts[-2] == "f": # We have an accidentally double filtered file... - new_file_name = "_".join(file_parts[:-1]) + ".fits" - os.rename(os.path.join(root,f), os.path.join(root,new_file_name)) - continue # Already filtered from a previous run - elif file_name_end == "e": - if reprocess_errors: - # Rename file to remove error designation - new_file_name = "_".join(file_parts[:-1]) + ".fits" - os.rename(os.path.join(root,f), os.path.join(root,new_file_name)) - # Check if error image exists and delete if needed - img_name = new_file_name.split(".fits")[0] + "_e.jpg" - try: - os.remove(os.path.join(root, img_name)) - except: - pass - # Add job to queue - files_by_timestamp[measure_end_time].append(os.path.join(root, new_file_name)) - else: - continue - elif file_name_end.startswith("v1-0-"): # This is the normal case for unprocessed data - files_by_timestamp[measure_end_time].append(os.path.join(root,f)) - else: - # print(f"Error - Unexpected FITS file name: {f}") - pass - sorted_times = sorted(list(files_by_timestamp.keys())) - for st in tqdm.tqdm(sorted_times, desc="Filtering files"): - for f in files_by_timestamp[st]: - work_queue.put(f) - except KeyboardInterrupt: - print("Finishing current jobs and exiting") - - - for _ in range(nworkers): - try: - work_queue.put(None, timeout=10.0) - except: - break - - for w in workers: - w.join() +#!/usr/bin/env python +"""Index the SUVI archive and record per-frame quality verdicts. + +Replaces the older version of this script, which encoded its verdict by renaming +files to ``_f``/``_e`` and wrote a diagnostic JPG next to every reject. Verdicts now +go to the SQLite index, so the archive keeps NOAA's filenames, several detectors can +disagree about a frame, the continuous scores survive for later re-thresholding, and +nothing has to be re-read to change a threshold. + +Subcommands:: + + filter_FITS.py index # sync the index with the archive + filter_FITS.py scan # cache FITS header metadata (~18 KB read per frame) + filter_FITS.py detect # run detectors and record verdicts + filter_FITS.py plots # regenerate diagnostic images for flagged frames + +``index`` and ``scan`` are incremental: they skip what they already have, so they can +be re-run as the puller brings in new data. ``index`` reads only the day-directories +whose mtime changed, which is what keeps re-indexing affordable -- see suvi/index.py +for why traversing this archive in full is something to avoid, not merely optimise. + +A caution before using ``detect --detectors geometry_v1`` to gate anything: measured +over 20 sampled days of 2024, that method rejects 36.9% of frames, and does so +all-or-nothing -- 0% on most days, 73-100% on days with a large active region, almost +entirely for "failed centre". Its centre test uses the intensity-weighted centroid, +which a bright active region displaces even when the header reports the disc exactly +centred. Run bench.py before trusting any detector to discard data. +""" + +import argparse +import json +import os +import sys +import time +from collections import defaultdict +from multiprocessing import Pool + +import numpy as np + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +from suvi import db, detectors, fitsio, index, paths, vfs + + +def parse_int_list(text): + return tuple(int(part) for part in text.split(",") if part.strip()) + + +# ---------------------------------------------------------------------------- index + + +def cmd_index(args): + """Bring the index in step with the archive. + + Reads only the day-directories whose mtime changed since the last run, so a + routine re-index costs ~7,400 directory stats rather than 2.65M file lookups. + That ratio is not just a speed matter: a full traversal exhausts the file + handles of the virtiofs mount this archive lives on. Use --full only for a cold + start or when the index is suspect. + """ + root = args.root or paths.data_root() + conn = db.connect(args.db) + satellites = parse_int_list(args.satellites) + wavelengths = parse_int_list(args.wavelengths) + + print(f"Indexing {root}" + (" (full rescan)" if args.full else " (changed directories only)")) + summary = index.reconcile( + conn, + root, + satellites, + wavelengths, + years=args.year, + force=args.full, + progress=print if args.progress else None, + ) + total = conn.execute("SELECT count(*) c FROM frame").fetchone()["c"] + print( + f" {summary['directories_changed']} of {summary['directories_checked']} " + f"directories changed in {summary['seconds']:.0f}s" + ) + print(f" frames added {summary['frames_added']}, removed {summary['frames_removed']}") + if summary["directories_vanished"]: + print(f" directories that disappeared: {summary['directories_vanished']}") + dupes = summary.get("duplicate_slots") or [] + if dupes: + print(f" {len(dupes)} files share an observation slot with another and were " + f"left out of the index (see --duplicates to list them)") + if args.duplicates: + for path in dupes: + print(f" {path}") + print(f" index now holds {total} frames") + + suffixed = conn.execute( + "SELECT count(*) c FROM frame WHERE path LIKE '%\\_f.fits' ESCAPE '\\' " + " OR path LIKE '%\\_e.fits' ESCAPE '\\'" + ).fetchone()["c"] + if suffixed: + print(f"\n{suffixed} indexed frames still carry _f/_e suffixes.") + print("Run migrate_unrename.py to fold those verdicts into the index.") + conn.close() + return 0 + + +# ----------------------------------------------------------------------------- scan + + +def _scan_one(job): + frame_id, path = job + values, error = fitsio.scan_header(path) + return frame_id, values, error + + +def cmd_scan(args): + fitsio.quiet_astropy() + conn = db.connect(args.db) + root = args.root or paths.data_root() + pending = db.unscanned_frames(conn, args.limit) + if not pending: + print("Every indexed frame already has cached header metadata.") + return 0 + + print(f"Scanning headers for {len(pending)} frames with {args.workers} workers") + jobs = [(row["id"], paths.abspath(row["path"], root)) for row in pending] + started = time.time() + done = failed = 0 + batch = [] + reliever = vfs.Reliever(label="scan") + with Pool(args.workers, initializer=fitsio.quiet_astropy) as pool: + for frame_id, values, error in pool.imap_unordered(_scan_one, jobs, 32): + reliever.tick() + batch.append((frame_id, values, error)) + done += 1 + failed += bool(error) + if len(batch) >= 2000: + db.record_headers(conn, batch) + conn.commit() + batch.clear() + rate = done / max(time.time() - started, 1e-6) + print(f" {done}/{len(jobs)} ({rate:.0f}/s, {failed} unreadable)") + if batch: + db.record_headers(conn, batch) + conn.commit() + print(f"Scanned {done} frames in {time.time() - started:.0f}s; {failed} unreadable") + return 0 + + +# --------------------------------------------------------------------------- detect + + +def _detect_one(job): + """Run the frame-level detectors on one file.""" + frame_id, slot, path, wanted = job + verdicts = {} + values, header_error = fitsio.scan_header(path) + features = detectors.FrameFeatures(slot=slot, header=values, error=header_error) + + if "header_v1" in wanted: + started = time.perf_counter_ns() + verdict = detectors.header_v1(features) + verdicts["header_v1"] = (verdict, (time.perf_counter_ns() - started) // 1000) + + pixel_detectors = [n for n in ("geometry_v1", "disc_v1") if n in wanted] + if pixel_detectors: + image, image_error = fitsio.read_image(path) + for name in pixel_detectors: + started = time.perf_counter_ns() + if image is None: + verdict = detectors.Verdict("bad", f"unreadable: {image_error}", {}) + elif name == "geometry_v1": + verdict = detectors.geometry_v1(image, slot[1]) + else: + verdict = detectors.disc_v1(image, slot[1], values) + verdicts[name] = (verdict, (time.perf_counter_ns() - started) // 1000) + + return frame_id, verdicts + + +def cmd_detect(args): + fitsio.quiet_astropy() + conn = db.connect(args.db) + root = args.root or paths.data_root() + wanted = tuple(args.detectors.split(",")) + unknown = [name for name in wanted if name not in detectors.ALL_DETECTORS] + if unknown: + print(f"Unknown detectors: {unknown}. Available: {list(detectors.ALL_DETECTORS)}") + return 1 + + frame_detectors = [n for n in wanted if n in detectors.FRAME_DETECTORS] + series_detectors = [n for n in wanted if n in detectors.SERIES_DETECTORS] + if series_detectors: + print(f"NOTE: {series_detectors} need a whole series in memory; use bench.py " + "for those, or run them here one band at a time via --wavelengths.") + + where, params = ["1=1"], [] + if args.start is not None: + where.append("t_start >= ?") + params.append(args.start) + if args.end is not None: + where.append("t_start < ?") + params.append(args.end) + if args.wavelengths: + bands = parse_int_list(args.wavelengths) + where.append(f"wavelength IN ({','.join('?' * len(bands))})") + params.extend(bands) + rows = conn.execute( + f"SELECT id, path, satellite, wavelength, t_start FROM frame " + f"WHERE {' AND '.join(where)} ORDER BY t_start", + params, + ).fetchall() + if not rows: + print("No indexed frames match; run `filter_FITS.py index` first.") + return 1 + + runs = { + name: db.create_detector_run(conn, name, {"driver": "filter_FITS"}, notes=args.notes) + for name in frame_detectors + } + print(f"Running {frame_detectors} over {len(rows)} frames with {args.workers} workers") + + jobs = [ + ( + row["id"], + (row["satellite"], row["wavelength"], row["t_start"]), + paths.abspath(row["path"], root), + frozenset(frame_detectors), + ) + for row in rows + ] + started = time.time() + pending = defaultdict(list) + tallies = {name: defaultdict(int) for name in frame_detectors} + done = 0 + reliever = vfs.Reliever(label="detect") + + with Pool(args.workers, initializer=fitsio.quiet_astropy) as pool: + for frame_id, verdicts in pool.imap_unordered(_detect_one, jobs, 16): + reliever.tick() + for name, (verdict, elapsed) in verdicts.items(): + pending[name].append( + (frame_id, verdict.verdict, verdict.reason, verdict.scores, elapsed) + ) + tallies[name][verdict.verdict] += 1 + done += 1 + if done % 5000 == 0: + for name, batch in pending.items(): + db.record_detections(conn, runs[name], batch) + batch.clear() + conn.commit() + print(f" {done}/{len(jobs)} ({done / max(time.time() - started, 1e-6):.0f}/s)") + for name, batch in pending.items(): + if batch: + db.record_detections(conn, runs[name], batch) + conn.commit() + + print(f"\nDone in {time.time() - started:.0f}s") + for name in frame_detectors: + counts = tallies[name] + total = sum(counts.values()) or 1 + print(f" {name}: run {runs[name]} -- " + + ", ".join(f"{v} {c} ({c / total:.1%})" for v, c in sorted(counts.items()))) + return 0 + + +# ------------------------------------------------------------------------ calibrate + +#: Absolute floor on each measurement's half-width, so a band whose sample happens to +#: be unusually tight does not get bounds narrower than the physics justifies. Set +#: from the observed cross-band spread of each quantity. +CALIBRATION_FLOORS = {"limb_contrast": 0.15, "radius_ratio": 0.06, "limb_width": 0.06} + +#: Which side of each measurement actually diagnoses a fault. Bounding the other +#: side buys nothing and, on a wide-spread quantity, produces a negative limit that +#: can never trip -- which looks like a rule while being none. +#: +#: * limb_contrast falls towards zero as the disc disappears; it cannot exceed 1. +#: * radius_ratio is diagnostic in both directions -- too small or too large is wrong. +#: * limb_width grows as the limb smears; a sharper-than-usual limb is not a fault. +CALIBRATION_SIDES = { + "limb_contrast": ("lower", 0.02, 1.0), + "radius_ratio": ("both", None, None), + "limb_width": ("upper", 0.0, None), +} + + +def _sample_frames(conn, root, wavelength, count, satellites=(16, 18)): + """Pick frames spread evenly across the whole time range of one band. + + Queries the index rather than walking the archive -- which is the point of + having an index, and avoids the traversal that exhausts this mount (see + suvi/index.py). + + Deliberately does *not* filter on the legacy ``_f``/``_e`` verdicts. Those were + written in May 2024 by a filter retuned that July, so selecting on them would + calibrate the new detector against the old one's mistakes -- including the very + active-region rejections this detector exists to stop. + """ + placeholders = ",".join("?" * len(satellites)) + rows = conn.execute( + f"SELECT path FROM frame WHERE wavelength = ? AND satellite IN ({placeholders}) " + f"ORDER BY t_start", + (wavelength, *satellites), + ).fetchall() + if not rows: + return [] + stride = max(1, len(rows) // count) + return [paths.abspath(row["path"], root) for row in rows[::stride]][:count] + + +def _calibration_measure(path): + """Measure one frame, skipping any the header already condemns.""" + values, error = fitsio.scan_header(path) + if error or not values: + return None + if values.get("eclipse") or values.get("empty"): + return None # genuinely bad; must not widen the good-frame bounds + diameter = values.get("diam_sun") + mean = values.get("img_mean") + if not diameter or mean is None: + return None + low, high = detectors.HEADER_MEAN_BOUNDS.get( + int(values.get("wavelnth") or 0), (0.0, float("inf")) + ) + if not (low <= mean <= high): + return None # dropout or blowout + image, image_error = fitsio.read_image(path) + if image is None: + return None + return detectors.disc_profile(image, float(diameter) / 4.0) + + +def cmd_calibrate(args): + fitsio.quiet_astropy() + root = args.root or paths.data_root() + conn = db.connect(args.db, readonly=True) + wavelengths = parse_int_list(args.wavelengths) + margin = args.margin + table = {} + + print(f"Calibrating disc_v1 from {root}") + print(f" {args.samples_per_band} frames per band, margin {margin}x the observed span\n") + + for wavelength in wavelengths: + candidates = _sample_frames(conn, root, wavelength, args.samples_per_band * 2) + if not candidates: + print(f"{wavelength:>5}: no frames found") + continue + results = [] + reliever = vfs.Reliever(label=f"calibrate {wavelength}A") + with Pool(args.workers, initializer=fitsio.quiet_astropy) as pool: + for measured in pool.imap_unordered(_calibration_measure, candidates, 8): + reliever.tick() + if measured is None: + continue + if any(v is None for v in measured.values()): + results.append(measured) # keep, so unmeasurable rate is visible + continue + results.append(measured) + if len(results) >= args.samples_per_band: + break + + usable = [r for r in results if all(v is not None for v in r.values())] + unmeasurable = len(results) - len(usable) + if len(usable) < 20: + print(f"{wavelength:>5}: only {len(usable)} usable frames; skipping") + continue + + bounds = {} + print(f"{wavelength:>5}A n={len(usable)} unmeasurable={unmeasurable}") + for name in ("limb_contrast", "radius_ratio", "limb_width"): + values = np.array([r[name] for r in usable], dtype=float) + low, high = float(values.min()), float(values.max()) + span = high - low + half = max(margin * span, CALIBRATION_FLOORS[name]) + side, floor, ceiling = CALIBRATION_SIDES[name] + + lower = low - half if side in ("lower", "both") else floor + upper = high + half if side in ("upper", "both") else ceiling + # Keep each limit on the physical side of zero. A negative floor on a + # non-negative quantity is not a loose rule, it is no rule at all. + if floor is not None: + lower = max(lower, floor) + if ceiling is not None: + upper = min(upper, ceiling) + bounds[name] = (round(lower, 4), round(upper, 4)) + outside = int(((values < lower) | (values > upper)).sum()) + print(f" {name:<14} observed {low:.4f}-{high:.4f} " + f"p50 {np.median(values):.4f} -> bounds {lower:.4f}-{upper:.4f} " + f"sample failures {outside}/{len(values)}") + table[wavelength] = bounds + + if not table: + print("\nNothing calibrated.") + return 1 + + print("\n" + "=" * 78) + print("Paste into DISC_BOUNDS in suvi/detectors.py:\n") + print("DISC_BOUNDS = {") + for wavelength, bounds in sorted(table.items()): + print(f" {wavelength}: {{") + for name, (lower, upper) in bounds.items(): + print(f' "{name}": ({lower}, {upper}),') + print(" },") + print("}") + if args.json: + with open(args.json, "w") as handle: + json.dump(table, handle, indent=2) + print(f"\nAlso written to {args.json}") + return 0 + + +# --------------------------------------------------------------------------- sample + +#: The production combination: a frame is bad if either detector says so. +SAMPLE_DETECTORS = ("header_v1", "disc_v1") + + +def _judge_timestamp(job): + """Read all six bands for one (satellite, timestamp) and judge each. + + Returns (satellite, timestamp, {band: (verdict, reason, scores)}), or None if + the set is incomplete -- a composite needs all six. + """ + satellite, timestamp, band_paths = job + verdicts = {} + for band, path in sorted(band_paths.items()): + values, header_error = fitsio.scan_header(path) + features = detectors.FrameFeatures( + slot=(satellite, band, timestamp), header=values, error=header_error + ) + header = detectors.header_v1(features) + image, image_error = fitsio.read_image(path) + if image is None: + disc = detectors.Verdict("bad", f"unreadable: {image_error}", {}) + else: + disc = detectors.disc_v1(image, band, values) + + bad = [v for v in (header, disc) if v.verdict == "bad"] + verdict = "bad" if bad else "good" + reason = "; ".join(v.reason for v in bad if v.reason) or None + verdicts[band] = (verdict, reason, {**header.scores, **disc.scores}) + return satellite, timestamp, verdicts + + +def _render_sample(job): + """Render one annotated composite for review.""" + import merger_FITS + from PIL import ImageDraw, ImageFont + + satellite, timestamp, band_paths, caption, out_path = job + try: + arrays = [] + for band in sorted(band_paths): + image, error = fitsio.read_image(band_paths[band]) + if image is None: + return out_path, f"band {band}: {error}" + arrays.append(image) + picture = merger_FITS.composite_from_arrays(arrays, timestamp) + + font_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), + "OpenSans-Regular.ttf") + font = ImageFont.truetype(font_path, size=20) + draw = ImageDraw.Draw(picture) + # Caption along the bottom, where it covers no solar structure. + y = picture.size[1] - 26 * len(caption) + for line in caption: + draw.text((12, y), line, (255, 235, 120), font=font) + y += 26 + picture.save(out_path, quality=95) + return out_path, None + except Exception as exc: + return out_path, f"{type(exc).__name__}: {exc}" + + +def cmd_sample(args): + """Write folders of suspected-good and suspected-bad frames for manual review. + + The bench measures detectors against damage it injected itself, which proves + they find what was planted but says nothing about what they do to real frames + nobody labelled. This is that check: a person looks at what the detector kept + and what it threw away, and says whether they agree. + + Each image is the full composite for the frame's timestamp and satellite -- which + already tiles all six bands down its sides -- captioned with the verdict, the + band responsible, and the scores behind it. + """ + fitsio.quiet_astropy() + root = args.root or paths.data_root() + conn = db.connect(args.db, readonly=True) + satellites = parse_int_list(args.satellites) + + # Candidate timestamps spread evenly across whatever the index holds, so the + # sample is not a picture of one week. + rows = conn.execute( + f"SELECT satellite, t_start, wavelength, path FROM frame " + f"WHERE satellite IN ({','.join('?' * len(satellites))}) " + f"ORDER BY t_start", + satellites, + ).fetchall() + conn.close() + + by_key = defaultdict(dict) + for row in rows: + by_key[(row["satellite"], row["t_start"])][row["wavelength"]] = paths.abspath( + row["path"], root + ) + complete = [k for k, v in by_key.items() if len(v) == 6] + if not complete: + print("No timestamp has all six bands; cannot build composites.") + return 1 + + stride = max(1, len(complete) // max(args.scan, 1)) + candidates = sorted(complete)[::stride][: args.scan] + print(f"{len(complete)} complete timestamps in the index; judging {len(candidates)}") + + good, bad = [], [] + reliever = vfs.Reliever(label="sample") + jobs = [(sat, when, by_key[(sat, when)]) for sat, when in candidates] + + with Pool(args.workers, initializer=fitsio.quiet_astropy) as pool: + for result in pool.imap_unordered(_judge_timestamp, jobs, 4): + reliever.tick(6) + if result is None: + continue + satellite, timestamp, verdicts = result + flagged = {b: v for b, v in verdicts.items() if v[0] == "bad"} + if flagged and len(bad) < args.n: + bad.append((satellite, timestamp, flagged)) + elif not flagged and len(good) < args.n: + good.append((satellite, timestamp, verdicts)) + if len(good) >= args.n and len(bad) >= args.n: + break + reliever.finish() + + print(f" collected {len(bad)} flagged and {len(good)} clean timestamps") + if len(bad) < args.n: + print(f" NOTE: only {len(bad)} flagged found in {len(candidates)} scanned; " + f"raise --scan for more.") + + render_jobs = [] + for folder, items in (("bad", bad), ("good", good)): + out_dir = os.path.join(args.out, folder) + os.makedirs(out_dir, exist_ok=True) + for satellite, timestamp, verdicts in items: + stamp = time.strftime("%Y%m%dT%H%M%S", time.gmtime(timestamp)) + if folder == "bad": + bands = ",".join(f"{b}A" for b in sorted(verdicts)) + first = verdicts[sorted(verdicts)[0]] + # Lead the filename with the reason so the folder sorts by failure mode. + slug = "".join( + ch if ch.isalnum() else "_" for ch in (first[1] or "flagged") + )[:40] + caption = [ + f"FLAGGED g{satellite} {stamp} bands: {bands}", + (first[1] or "")[:110], + " ".join(f"{k}={v:.3g}" for k, v in list(first[2].items())[:5]), + ] + name = f"{slug}__g{satellite}_{stamp}.jpg" + else: + sample_scores = verdicts[sorted(verdicts)[0]][2] + caption = [ + f"CLEAN g{satellite} {stamp} all six bands passed", + " ".join(f"{k}={v:.3g}" for k, v in list(sample_scores.items())[:5]), + ] + name = f"g{satellite}_{stamp}.jpg" + render_jobs.append( + (satellite, timestamp, by_key[(satellite, timestamp)], caption, + os.path.join(out_dir, name)) + ) + + print(f"Rendering {len(render_jobs)} composites...") + failures = 0 + with Pool(args.workers, initializer=fitsio.quiet_astropy) as pool: + for path, error in pool.imap_unordered(_render_sample, render_jobs, 2): + if error: + failures += 1 + if failures <= 3: + print(f" failed {os.path.basename(path)}: {error}") + print(f" wrote {len(render_jobs) - failures} images to {args.out}") + print(f" review {args.out}/bad -- nothing here should be a frame you would keep") + print(f" review {args.out}/good -- nothing here should be a frame you would discard") + return 0 + + +# ---------------------------------------------------------------------------- plots + + +def cmd_plots(args): + """Regenerate the diagnostic image the old filter wrote beside every reject. + + Kept as an on-demand tool rather than a side effect of filtering: the plots are + derivable from the index plus the FITS, and writing millions of them into the + archive was what made the old approach hard to undo. + """ + import matplotlib + + matplotlib.use("Agg") + import numpy as np + from matplotlib import pyplot as plt + + fitsio.quiet_astropy() + conn = db.connect(args.db, readonly=True) + root = args.root or paths.data_root() + run_id = db.latest_run_id(conn, args.detector) + if run_id is None: + print(f"No run named {args.detector!r}") + return 1 + + rows = conn.execute( + """ + SELECT f.path, f.wavelength, d.reason FROM detection d + JOIN frame f ON f.id = d.frame_id + WHERE d.run_id = ? AND d.verdict = 'bad' ORDER BY f.t_start LIMIT ? + """, + (run_id, args.limit), + ).fetchall() + os.makedirs(args.out, exist_ok=True) + print(f"Writing {len(rows)} diagnostic plots to {args.out}") + + reference = detectors.ideal_disc_axis() + axis = list(range(detectors.HALF_DIMS)) + for row in rows: + path = paths.abspath(row["path"], root) + image, error = fitsio.read_image(path) + if image is None: + continue + import cv2 as cv + + threshold = detectors.GEOMETRY_THRESHOLDS[row["wavelength"]] + data = cv.resize( + np.nan_to_num(image), + (detectors.HALF_DIMS, detectors.HALF_DIMS), + interpolation=cv.INTER_LINEAR, + ) + shown = np.clip(data, None, threshold) / threshold + xavg, yavg = np.average(shown, 0), np.average(shown, 1) + + figure, axes = plt.subplots( + 2, 2, width_ratios=(0.2, 1), height_ratios=(0.2, 1), + gridspec_kw={"hspace": 0.0, "wspace": 0.0}, figsize=(8, 8), + ) + axes[0][0].axis("off") + axes[0][1].plot(axis, xavg) + axes[0][1].plot(axis, reference) + axes[1][0].plot(yavg, axis) + axes[1][0].plot(reference, axis) + axes[1][0].invert_xaxis() + axes[1][0].invert_yaxis() + axes[1][1].imshow(shown, aspect="auto", vmin=0, vmax=1) + axes[0][1].sharex(axes[1][1]) + axes[1][0].sharey(axes[1][1]) + figure.suptitle(f"{os.path.basename(row['path'])}\n{row['reason']}", fontsize=8) + figure.savefig( + os.path.join(args.out, os.path.basename(row["path"]).replace(".fits", ".jpg")), + dpi=110, + ) + plt.close(figure) + return 0 + + +# ------------------------------------------------------------------------------ cli + + +def build_parser(): + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--db", default=None, help="SQLite index (default: $SUVI_DB)") + parser.add_argument("--root", default=None, help="archive root (default: $SUVI_DATA_ROOT)") + subparsers = parser.add_subparsers(dest="command", required=True) + + index_cmd = subparsers.add_parser( + "index", help="sync the index with the archive (only changed directories)" + ) + index_cmd.add_argument("--satellites", default="16,18,19") + index_cmd.add_argument("--wavelengths", default="94,131,171,195,284,304") + index_cmd.add_argument("--year", type=int, action="append", default=None) + index_cmd.add_argument("--full", action="store_true", + help="re-read every directory, ignoring recorded mtimes") + index_cmd.add_argument("--progress", action="store_true") + index_cmd.add_argument("--duplicates", action="store_true", + help="list files that share an observation slot") + index_cmd.set_defaults(func=cmd_index) + + scan = subparsers.add_parser("scan", help="cache FITS header metadata") + scan.add_argument("--workers", type=int, default=16) + scan.add_argument("--limit", type=int, default=None) + scan.set_defaults(func=cmd_scan) + + detect = subparsers.add_parser("detect", help="run detectors and record verdicts") + detect.add_argument("--detectors", default="header_v1") + detect.add_argument("--wavelengths", default=None) + detect.add_argument("--start", type=int, default=None, help="Unix time, inclusive") + detect.add_argument("--end", type=int, default=None, help="Unix time, exclusive") + detect.add_argument("--workers", type=int, default=16) + detect.add_argument("--notes", default=None) + detect.set_defaults(func=cmd_detect) + + calibrate = subparsers.add_parser( + "calibrate", help="derive per-band disc_v1 bounds from a sample of good frames" + ) + calibrate.add_argument("--wavelengths", default="94,131,171,195,284,304") + calibrate.add_argument("--samples-per-band", type=int, default=300) + calibrate.add_argument("--margin", type=float, default=1.0, + help="half-width added either side, in multiples of the observed span") + calibrate.add_argument("--workers", type=int, default=16) + calibrate.add_argument("--json", default=None) + calibrate.set_defaults(func=cmd_calibrate) + + sample = subparsers.add_parser( + "sample", help="write good/bad frame folders for manual review" + ) + sample.add_argument("--n", type=int, default=100, help="frames per folder") + sample.add_argument("--out", required=True) + sample.add_argument("--scan", type=int, default=4000, + help="timestamps to judge while looking for N of each") + sample.add_argument("--satellites", default="16,18") + sample.add_argument("--workers", type=int, default=8) + sample.set_defaults(func=cmd_sample) + + plots = subparsers.add_parser("plots", help="regenerate diagnostics for flagged frames") + plots.add_argument("--detector", default="geometry_v1") + plots.add_argument("--out", default="diagnostics") + plots.add_argument("--limit", type=int, default=200) + plots.set_defaults(func=cmd_plots) + return parser + + +def main(argv=None): + args = build_parser().parse_args(argv) + return args.func(args) or 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/make_comparison_video.py b/make_comparison_video.py new file mode 100644 index 0000000..441fe12 --- /dev/null +++ b/make_comparison_video.py @@ -0,0 +1,168 @@ +#!/usr/bin/env python +"""Stack the pristine, current-pipeline and new-pipeline streams side by side. + +`bench.py render` writes one composite per timestamp per variant. This turns three +such streams into a single video with the panes labelled, so the repair can be judged +against both the truth and what the pipeline does today: + + bench.py render --case C --variant pristine --out /tmp/streams + bench.py render --case C --variant today --out /tmp/streams + bench.py render --case C --variant new --out /tmp/streams + make_comparison_video.py --streams /tmp/streams --satellite 16 --out compare.mp4 + +The panes must stay in step, which is the whole reason this does not simply hand +ffmpeg three directories: a variant that produced no composite for a timestamp -- what +the current pipeline does past its gap limit -- would otherwise shorten that pane and +slide it out of alignment with the others. Every stream is emitted against the same +canonical timeline, with black where a composite is missing, so frame N is the same +instant in all three panes. Those black runs are not padding; they are what the +current pipeline actually shows. +""" + +import argparse +import os +import re +import subprocess +import sys + +import numpy as np +from PIL import Image + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +from ffmpeg_video import find_ffmpeg +from suvi import paths, vfs + +VARIANTS = ("pristine", "today", "new") +LABELS = { + "pristine": "PRISTINE (ground truth)", + "today": "TODAY (hold-last + black frames)", + "new": "NEW (disc+header, optical flow)", +} +_STAMP = re.compile(r"Composite-(\d+)\.jpg$") + + +def stream_frames(directory): + """{timestamp: path} for one rendered stream.""" + frames = {} + try: + entries = os.listdir(directory) + except OSError: + return frames + for entry in entries: + match = _STAMP.search(entry) + if match: + frames[int(match.group(1))] = os.path.join(directory, entry) + return frames + + +def encode_stream(ffmpeg, frames, timeline, out_path, framerate, crf, reliever=None): + """Encode one pane, emitting a black frame wherever the variant has nothing. + + Reads every composite in the pane, so it hands file handles back as it goes -- + three panes of a week-long window is over 7,000 files, enough on its own to + exhaust the mount and take unrelated software down with it. + """ + command = ( + f'{ffmpeg} -y -f image2pipe -framerate {framerate} -i - ' + f'-c:v libx264 -crf {crf} -preset veryfast -pix_fmt yuv420p "{out_path}"' + ) + process = subprocess.Popen(command, shell=True, stdin=subprocess.PIPE, + stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + blank = None + present = 0 + for when in timeline: + path = frames.get(when) + if path: + with open(path, "rb") as handle: + process.stdin.write(handle.read()) + present += 1 + if reliever is not None: + reliever.tick() + else: + if blank is None: + size = Image.open(next(iter(frames.values()))).size if frames else (1920, 1080) + blank = Image.fromarray(np.zeros((size[1], size[0], 3), dtype="uint8")) + blank.save(process.stdin, "jpeg", quality=95) + process.stdin.close() + process.wait() + return present + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--streams", required=True, help="root written by bench.py render") + parser.add_argument("--satellite", type=int, default=16) + parser.add_argument("--out", required=True) + parser.add_argument("--framerate", type=int, default=60) + parser.add_argument("--crf", type=int, default=18) + parser.add_argument("--keep-panes", action="store_true", + help="keep the intermediate per-pane videos") + args = parser.parse_args(argv) + + ffmpeg = find_ffmpeg() + streams = {} + for variant in VARIANTS: + directory = os.path.join(args.streams, variant, f"goes{args.satellite}") + streams[variant] = stream_frames(directory) + print(f" {variant:>8}: {len(streams[variant])} composites in {directory}") + if not any(streams.values()): + print("No composites found; run bench.py render first.") + return 1 + + # One canonical timeline across every variant, so the panes stay in step. + everything = set() + for frames in streams.values(): + everything.update(frames) + timeline = list(range(min(everything), max(everything) + paths.CADENCE, paths.CADENCE)) + print(f" timeline: {len(timeline)} slots " + f"({len(timeline) / args.framerate:.0f}s per pane at {args.framerate}fps)") + + panes = [] + reliever = vfs.Reliever(label="encode") + for variant in VARIANTS: + pane = f"{os.path.splitext(args.out)[0]}.{variant}.mp4" + present = encode_stream( + ffmpeg, streams[variant], timeline, pane, args.framerate, args.crf, reliever + ) + missing = len(timeline) - present + print(f" encoded {variant}: {present} frames" + + (f", {missing} black ({missing / len(timeline):.0%})" if missing else "")) + panes.append(pane) + + inputs = " ".join(f'-i "{p}"' for p in panes) + font = os.path.join(os.path.dirname(os.path.abspath(__file__)), "OpenSans-Regular.ttf") + labelled = [] + for index, variant in enumerate(VARIANTS): + labelled.append( + f"[{index}:v]scale=960:-2," + f"drawtext=fontfile='{font}':text='{LABELS[variant]}':" + f"x=12:y=12:fontsize=22:fontcolor=yellow:box=1:boxcolor=black@0.5[v{index}]" + ) + graph = ";".join(labelled) + ";[v0][v1][v2]hstack=inputs=3[out]" + command = ( + f'{ffmpeg} -y {inputs} -filter_complex "{graph}" -map "[out]" ' + f'-c:v libx264 -crf {args.crf} -preset veryfast -pix_fmt yuv420p "{args.out}"' + ) + print(" stacking panes...") + result = subprocess.run(command, shell=True, stdout=subprocess.DEVNULL, + stderr=subprocess.PIPE) + if result.returncode != 0: + print(result.stderr.decode()[-1500:]) + return 1 + + if not args.keep_panes: + for pane in panes: + try: + os.remove(pane) + except OSError: + pass + reliever.finish() + size = os.path.getsize(args.out) / 1024 / 1024 + print(f"Wrote {args.out} ({size:.0f} MB)") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/merger.py b/merger.py deleted file mode 100644 index 48cf5f5..0000000 --- a/merger.py +++ /dev/null @@ -1,270 +0,0 @@ -import os -import time -import calendar -import datetime -from collections import defaultdict -from multiprocessing import Queue, Process - -import tqdm -from PIL import Image, ImageFont, ImageDraw -import numpy as np -from matplotlib import pyplot as plt - -def bin_ndarray(ndarray, new_shape, operation='mean'): - """ - Bins an ndarray in all axes based on the target shape, by summing or - averaging. - - Number of output dimensions must match number of input dimensions and - new axes must divide old ones. - - Example - ------- - >>> m = np.arange(0,100,1).reshape((10,10)) - >>> n = bin_ndarray(m, new_shape=(5,5), operation='sum') - >>> print(n) - - [[ 22 30 38 46 54] - [102 110 118 126 134] - [182 190 198 206 214] - [262 270 278 286 294] - [342 350 358 366 374]] - - """ - operation = operation.lower() - if not operation in ['sum', 'mean']: - raise ValueError("Operation not supported.") - if ndarray.ndim != len(new_shape): - raise ValueError("Shape mismatch: {} -> {}".format(ndarray.shape, - new_shape)) - compression_pairs = [(d, c//d) for d,c in zip(new_shape, - ndarray.shape)] - flattened = [l for p in compression_pairs for l in p] - ndarray = ndarray.reshape(flattened) - for i in range(len(new_shape)): - op = getattr(ndarray, operation) - ndarray = op(-1*(i+1)) - return ndarray - -def gamma_correct(fun): - def wrapper(*args, **kwargs): - args = list(args) - args[0] = np.power(args[0], 2.2) - args[1] = np.power(args[1], 2.2) - args = tuple(args) - result = fun(*args, **kwargs) - return np.power(result, 1/2.2) - return wrapper - -def clip_color(fun): - def wrapper(*args, **kwargs): - return np.clip(fun(*args, **kwargs), 0.0, 1.0) - return wrapper - -# linear_srgb_matrix = np.array([[0.4124, 0.3576, 0.1805], -# [0.2126, 0.7152, 0.0722], -# [0.0193, 0.1192, 0.9505]]) - -# linear_srgb_matrix_inv = np.array([[ 3.2406, -1.5372, -0.4986], -# [-0.9689, 1.8758, 0.0415], -# [ 0.0557, -0.2040, 1.0570]]) - -# def linear_color_correction(fun): -# def wrapper(*args, **kwargs): -# args = list(args) -# inds = args[0] <= 0.04045 -# ninds = args[0] > 0.04045 -# for i in range(2): -# args[i][inds] = args[i][inds] / 12.92 -# args[i][ninds] = np.power((args[i][ninds] + 0.055) / 1.055, 2.4) -# for x in range(args[i].shape[0]): -# for y in range(args[i].shape[1]): -# args[i][x,y,:] = np.matmul(linear_srgb_matrix, args[i][x,y,:]) -# args = tuple(args) -# result = fun(*args, **kwargs) - -# for x in range(result.shape[0]): -# for y in range(result.shape[1]): -# result[x,y,:] = np.matmul(linear_srgb_matrix_inv, result[x,y,:]) -# inds = result <= 0.0031308 -# ninds = result > 0.0031308 -# result[inds] = result[inds] * 12.92 -# result[ninds] = np.power(result[ninds], 1.0/2.4) * 1.055 - 0.055 -# return result -# return wrapper - - -@gamma_correct -def composite_alpha_over(F, B, alpha_F, alpha_B = 1): - return (F*alpha_F + B*alpha_B*(1-alpha_F)) / (alpha_F + alpha_B*(1-alpha_F)) - -def composite_alpha_blend(F, B, alpha): - return F*alpha + B*(1-alpha) - -def linear_burn(F, B): - burn = F + B - 1 - burn[burn < 0.0] = 0.0 - return burn - -def difference(F, B): - return np.abs(F - B) - -@clip_color -def linear_light(F, B): - result = np.zeros_like(F) - inds = F <= 0.5 - ninds = F > 0.5 - result[inds] = B[inds] + 2.0 * F[inds] - 1 - result[ninds] = 2.0 * (F[ninds] - 0.5) + B[ninds] - return result - -@clip_color -def hard_light(F, B): - result = np.zeros_like(F) - inds = B < 0.5 - ninds = B >= 0.5 - result[inds] = 2 * F[inds] * B[inds] - result[ninds] = 1 - (2*(1 - F[ninds])*(1 - B[ninds])) - return result - -@clip_color -def color_dodge(F, B): - return B / (1.000001 - F) - -@clip_color -def exclusion(F, B): - return F + B - 2*F*B - -@clip_color -def saturation(img, R, G, B): - img[:,:,0] *= R - img[:,:,1] *= G - img[:,:,2] *= B - return img - -@clip_color -def contrast(img, c, b): - return (img - 0.5) * c + 0.5 + b*c - -def generate_composite(work_queue, result_queue): - while True: - try: - job = work_queue.get() - if job is None: - break - files_this_timestamp, timestamp, processed_images_dir = job - filename = f"Composite-{int(timestamp)}.jpg" - filepath = os.path.join(processed_images_dir, filename) - if os.path.isfile(filepath): - result_queue.put(("Exists", timestamp)) - continue - - # image_names = ["094Ã…", "131Ã…", "171Ã…", "195Ã…", "284Ã…", "304Ã…"] - data = [] - for i in range(6): - img = Image.open(files_this_timestamp[i]) - trimmed_img_data = np.array(img)[40:-40,40:-40,:3] # Trim off edges of image to remove text - normalized_img_data = trimmed_img_data / 255.0 # Normalize to float 0.0-1.0 instead of uint8 - srgb_img_data = np.power(normalized_img_data, 2.2) # Gamma correct to sRGB color space - data.append(normalized_img_data) - - # Assemble composite image - composite_image_data = data[4] # Start with (284Ã…) - # Do a linear burn with 304Ã… at 95% alpha - composite_image_data = composite_alpha_over(linear_burn(data[5], composite_image_data), composite_image_data, 0.95) - # Do a difference operation with 195Ã… at 95% alpha - composite_image_data = composite_alpha_over(exclusion(data[3], composite_image_data), composite_image_data, 0.90) - # Do a linear_light layer op with 171Ã… - composite_image_data = linear_light(data[2], composite_image_data) - # Do a hard_light layer op with 131Ã… - composite_image_data = composite_alpha_over(hard_light(data[1], composite_image_data), composite_image_data, 0.20) - # Do a color_dodge layer op with 094Ã… - composite_image_data = composite_alpha_over(color_dodge(data[0], composite_image_data), composite_image_data, 0.25) - # Do an exclusion layer op with 094Ã… - composite_image_data = composite_alpha_over(exclusion(data[0], composite_image_data), composite_image_data, 0.80) - # Tweak the colors a little - composite_image_data = saturation(composite_image_data, 1.0, 0.95, 1.15) - # Boost contrast - composite_image_data = contrast(composite_image_data, 1.5, 0.15) - - # Now shrink the component images and assemble them alongside the composite. - new_dim = composite_image_data.shape[0] // 3 - # Enlarge the composite to fit the new images - composite_image_data = np.pad(composite_image_data, ((0,0),(new_dim, new_dim),(0,0))) - for i in range(6): - resized = bin_ndarray(data[i], (new_dim, new_dim, 3)) - if i < 3: - composite_image_data[i*new_dim:(i+1)*new_dim, :new_dim, :] = resized - else: - composite_image_data[(i-3)*new_dim:(i-2)*new_dim, -new_dim:, :] = resized - - img = Image.fromarray((255 * composite_image_data).astype('uint8')) - timestring = datetime.datetime.fromtimestamp(timestamp, tz = datetime.UTC).strftime('%Y-%m-%d %H:%M:%S') - ImageDraw.Draw(img).text((655, 15), f"NOAA GOES Sattelite SUVI Composite - {timestring} UTC",(255,255,255), font_size = 24) - img.save(filepath, quality = 90) - result_queue.put(("Created", timestamp)) - - # plt.figure("Composite") - # plt.imshow(composite_image_data) - # plt.show() - # plt.close('all') - except Exception as e: - result_queue.put((e, timestamp)) - - -if __name__ == "__main__": - stored_images_dir = r"..\comp" - processed_images_dir = r"..\composite" - nworkers = 8 - os.makedirs(processed_images_dir, exist_ok=True) - - files_sorted_by_timestamp = defaultdict(list) - for root, dirs, files in os.walk(stored_images_dir): - for f in files: - if f.endswith(".png"): - file_parts = f.split("_") - measurement = file_parts[1] - sattelite = file_parts[2] - measure_end_time = datetime.datetime.strptime(file_parts[4][1:16], "%Y%m%dT%H%M%S") - measure_end_time.replace(tzinfo=datetime.timezone.utc) - measure_end_time = calendar.timegm(measure_end_time.timetuple()) - files_sorted_by_timestamp[measure_end_time].append(os.path.join(root,f)) - - work_queue = Queue(maxsize = 3) - result_queue = Queue() - workers = [] - for i in range(nworkers): - p = Process(target = generate_composite, args = (work_queue, result_queue), daemon=True) - p.start() - workers.append(p) - - for timestamp in tqdm.tqdm(files_sorted_by_timestamp, desc="Creating Composite Solar Images"): - files_this_timestamp = files_sorted_by_timestamp[timestamp] - files_this_timestamp = sorted(files_this_timestamp) - if not len(files_this_timestamp) == 7: - print(f"Detected a file gap at: {timestamp}") - continue - - work_queue.put((files_this_timestamp, timestamp, processed_images_dir)) - - ncreated = 0 - nexists = 0 - for _ in range(len(files_sorted_by_timestamp)): - result = result_queue.get(5.0) - if result[0] == "Exists": - nexists += 1 - elif result[0] == "Created": - ncreated += 1 - else: - print(f"A worker encountered an exception on job {result[1]}: {result[0]}") - - for _ in range(nworkers): - try: - work_queue.put(None, timeout=1.0) - except: - break - - for w in workers: - w.join(5.0) - - print("Done") \ No newline at end of file diff --git a/merger_FITS.py b/merger_FITS.py index 9d99fc4..9dc9447 100644 --- a/merger_FITS.py +++ b/merger_FITS.py @@ -1,579 +1,657 @@ -import os -import time -import calendar -import datetime -from collections import defaultdict -from multiprocessing import Queue, Process -import re -import warnings -import queue -import traceback - -import tqdm -from PIL import Image, ImageDraw, ImageFont -import numpy as np -from matplotlib import pyplot as plt -from astropy.io import fits -import palettable - -def lowpriority(): - """ Set the priority of the process to below-normal.""" - - import sys - try: - sys.getwindowsversion() - except AttributeError: - isWindows = False - else: - isWindows = True - - if isWindows: - # Based on: - # "Recipe 496767: Set Process Priority In Windows" on ActiveState - # http://code.activestate.com/recipes/496767/ - import win32api,win32process,win32con # pywin32 - - pid = win32api.GetCurrentProcessId() - handle = win32api.OpenProcess(win32con.PROCESS_ALL_ACCESS, True, pid) - win32process.SetPriorityClass(handle, win32process.BELOW_NORMAL_PRIORITY_CLASS) - else: - import os - - os.nice(1) - -def bin_ndarray(ndarray, new_shape, operation='mean'): - """ - Bins an ndarray in all axes based on the target shape, by summing or - averaging. - - Number of output dimensions must match number of input dimensions and - new axes must divide old ones. - - Example - ------- - >>> m = np.arange(0,100,1).reshape((10,10)) - >>> n = bin_ndarray(m, new_shape=(5,5), operation='sum') - >>> print(n) - - [[ 22 30 38 46 54] - [102 110 118 126 134] - [182 190 198 206 214] - [262 270 278 286 294] - [342 350 358 366 374]] - - """ - operation = operation.lower() - if not operation in ['sum', 'mean']: - raise ValueError("Operation not supported.") - if ndarray.ndim != len(new_shape): - raise ValueError("Shape mismatch: {} -> {}".format(ndarray.shape, - new_shape)) - compression_pairs = [(d, c//d) for d,c in zip(new_shape, - ndarray.shape)] - flattened = [l for p in compression_pairs for l in p] - ndarray = ndarray.reshape(flattened) - for i in range(len(new_shape)): - op = getattr(ndarray, operation) - ndarray = op(-1*(i+1)) - return ndarray - -def gamma_correct(fun): - def wrapper(*args, **kwargs): - args = list(args) - args[0] = np.power(args[0], 2.2) - args[1] = np.power(args[1], 2.2) - args = tuple(args) - result = fun(*args, **kwargs) - return np.power(result, 1/2.2) - return wrapper - -def clip_color(fun): - def wrapper(*args, **kwargs): - return np.clip(fun(*args, **kwargs), 0.0, 1.0) - return wrapper - -# linear_srgb_matrix = np.array([[0.4124, 0.3576, 0.1805], -# [0.2126, 0.7152, 0.0722], -# [0.0193, 0.1192, 0.9505]]) - -# linear_srgb_matrix_inv = np.array([[ 3.2406, -1.5372, -0.4986], -# [-0.9689, 1.8758, 0.0415], -# [ 0.0557, -0.2040, 1.0570]]) - -# def linear_color_correction(fun): -# def wrapper(*args, **kwargs): -# args = list(args) -# inds = args[0] <= 0.04045 -# ninds = args[0] > 0.04045 -# for i in range(2): -# args[i][inds] = args[i][inds] / 12.92 -# args[i][ninds] = np.power((args[i][ninds] + 0.055) / 1.055, 2.4) -# for x in range(args[i].shape[0]): -# for y in range(args[i].shape[1]): -# args[i][x,y,:] = np.matmul(linear_srgb_matrix, args[i][x,y,:]) -# args = tuple(args) -# result = fun(*args, **kwargs) - -# for x in range(result.shape[0]): -# for y in range(result.shape[1]): -# result[x,y,:] = np.matmul(linear_srgb_matrix_inv, result[x,y,:]) -# inds = result <= 0.0031308 -# ninds = result > 0.0031308 -# result[inds] = result[inds] * 12.92 -# result[ninds] = np.power(result[ninds], 1.0/2.4) * 1.055 - 0.055 -# return result -# return wrapper - - -@gamma_correct -def composite_alpha_over(F, B, alpha_F, alpha_B = 1): - return (F*alpha_F + B*alpha_B*(1-alpha_F)) / (alpha_F + alpha_B*(1-alpha_F)) - -def composite_alpha_blend(F, B, alpha): - return F*alpha + B*(1-alpha) - -def linear_burn(F, B): - burn = F + B - 1 - burn[burn < 0.0] = 0.0 - return burn - -def difference(F, B): - d = np.abs(F - B) - if d.shape[2] == 4: # Preserve alpha of base image - d[:,:,3] = B[:,:,3] - return d - -@clip_color -def linear_light(F, B): - result = np.zeros_like(F) - inds = F <= 0.5 - ninds = F > 0.5 - result[inds] = B[inds] + 2.0 * F[inds] - 1 - result[ninds] = 2.0 * (F[ninds] - 0.5) + B[ninds] - return result - -@clip_color -def hard_light(F, B): - result = np.zeros_like(F) - inds = B < 0.5 - ninds = B >= 0.5 - result[inds] = 2 * F[inds] * B[inds] - result[ninds] = 1 - (2*(1 - F[ninds])*(1 - B[ninds])) - return result - -@clip_color -def color_dodge(F, B): - return B / (1.000001 - F) - -@clip_color -def exclusion(F, B): - d = F + B - 2*F*B - if d.shape[2] == 4: # Preserve alpha of base image - d[:,:,3] = B[:,:,3] - return d - -@clip_color -def saturation(img, R, G, B): - img[:,:,0] *= R - img[:,:,1] *= G - img[:,:,2] *= B - return img - -@clip_color -def contrast(img, c, b): - return (img - 0.5) * c + 0.5 + b*c - -def rgb_to_hsl(img): - r = img[:,:,0] - g = img[:,:,1] - b = img[:,:,2] - cmax = np.copy(r) - cmax[g > cmax] = g[g > cmax] - cmax[b > cmax] = b[b > cmax] - cmin = np.copy(r) - cmin[g < cmin] = g[g < cmin] - cmin[b < cmin] = b[b < cmin] - delta = cmax - cmin - - - # Calc hue - hue = np.zeros_like(r) - inds = cmax == r - with warnings.catch_warnings(): - warnings.filterwarnings('ignore') - hue[inds] = 60 * np.mod((g[inds]-b[inds])/delta[inds], 6) - inds = cmax == g - hue[inds] = 60 * ((b[inds]-r[inds])/delta[inds] + 2) - inds = cmax == b - hue[inds] = 60 * ((r[inds]-g[inds])/delta[inds] + 4) - hue[np.isnan(hue)] = 0 - hue[hue < 0] = hue[hue < 0] + 360 # Make negative hue values positive behind 360 - - # Calc lightness / luminance - luminance = (cmax + cmin) / 2.0 - - # Calc saturation - saturation = np.zeros_like(r) - inds = delta != 0 - saturation[inds] = delta[inds] / (1.0 - np.abs(2 * luminance[inds] - 1.0)) - - # Multiply luminance and saturation by 100 to scale them to the appropriate range (0-100) - saturation *= 100.0 - luminance *= 100.0 - - - if img.shape[2] == 4: # Preserve alpha of original image - return np.stack([hue, saturation, luminance, img[:,:,3]], -1) - else: - return np.stack([hue, saturation, luminance], -1) - -def hsl_to_rgb(img): - hue = img[:,:,0] - saturation = img[:,:,1] - luminance = img[:,:,2] - saturation /= 100.0 - luminance /= 100.0 - c = (1.0 - np.abs(2.0 * luminance - 1.0)) * saturation - x = c * (1.0 - np.abs((hue/60.0) % 2.0 - 1.0)) - m = luminance - c / 2.0 - r = np.zeros_like(hue) - g = np.zeros_like(hue) - b = np.zeros_like(hue) - - inds1 = np.logical_and(0.0 <= hue, hue < 60.0) - r[inds1] = c[inds1] - g[inds1] = x[inds1] - inds2 = np.logical_and(60.0 <= hue, hue < 120.0) - r[inds2] = x[inds2] - g[inds2] = c[inds2] - inds3 = np.logical_and(120.0 <= hue, hue < 180.0) - g[inds3] = c[inds3] - b[inds3] = x[inds3] - inds4 = np.logical_and(180.0 <= hue, hue < 240.0) - g[inds4] = x[inds4] - b[inds4] = c[inds4] - inds5 = np.logical_and(240.0 <= hue, hue < 300.0) - r[inds5] = x[inds5] - b[inds5] = c[inds5] - inds6 = np.logical_and(240.0 <= hue, hue < 300.0) - r[inds6] = c[inds6] - b[inds6] = x[inds6] - - r += m - g += m - b += m - - if img.shape[2] == 4: # Preserve alpha of original image - return np.stack([r, g, b, img[:,:,3]], -1) - else: - return np.stack([r, g, b], -1) - -def generate_composite(work_queue, result_queue): - lowpriority() - cmaps =[palettable.cmocean.sequential.Ice_5.mpl_colormap, - palettable.cmocean.sequential.Ice_20.mpl_colormap, - palettable.cmocean.sequential.Turbid_5_r.mpl_colormap, - palettable.cmocean.sequential.Turbid_20_r.mpl_colormap, - plt.colormaps.get_cmap('gist_heat'), - plt.colormaps.get_cmap('afmhot')] - # Modify colormaps to start at perfect black (when they otherwise start at very dark colors) - for cmi in range(0,4): - for c in ['red','green','blue']: - for i in range(3): - cmaps[cmi]._segmentdata[c][0][i] = 0.0 - # These values are used to map floating point radiance values to colors using the above color maps. - vmins = [0.050, 0.05, 00.100, 00.10, 00.100, 00.1] - vmaxs = [8.000, 8.00, 20.000, 30.00, 40.000, 90.0] - gammas = [0.375, 0.40, 00.425, 00.45, 00.475, 00.5] - trimx = 64 - trimy = 100 - fnt1 = ImageFont.truetype("OpenSans-Regular.ttf", size = 24) - fnt2 = ImageFont.truetype("OpenSans-Regular.ttf", size = 16) - while True: - try: - job = work_queue.get() - if job is None: - result_queue.cancel_join_thread() - return - files_this_timestamp, timestamp, processed_images_dir, synthetic_data = job - synthetic_filepath = os.path.join(processed_images_dir, f"Composite-{int(timestamp)}_s.jpg") - normal_filepath = os.path.join(processed_images_dir, f"Composite-{int(timestamp)}.jpg") - filepath = None - if synthetic_data: - if os.path.isfile(synthetic_filepath): - result_queue.put(("Exists", timestamp)) - continue - filepath = synthetic_filepath - else: - if os.path.isfile(synthetic_filepath): - os.remove(synthetic_filepath) - if os.path.isfile(normal_filepath): - result_queue.put(("Exists", timestamp)) - continue - filepath = normal_filepath - - base_imgs = [] - for i in range(6): - raw_data = np.flip(fits.getdata(files_this_timestamp[i]), 0)[trimy:-trimy,trimx:-trimx] # Trim and reorient the image to match our final desired dimensions - raw_data[raw_data < 0.0] = 0.0 # Remove non-zero data because it doesn't makes sense (supposed to be std Radiance) - base_imgs.append(cmaps[i](np.clip((raw_data - vmins[i]) / vmaxs[i], 0, 1.0)**gammas[i])[:,:,:3]) - # plt.figure(image_names[i]) - # plt.imshow(base_imgs[-1]) - - - # plt.figure("Initial Blend") - composite_image_data = composite_alpha_over(base_imgs[4], base_imgs[5], 0.2)**0.5 - # plt.imshow(composite_image_data) - - # plt.figure("Linear Burn with 304") - composite_image_data = composite_alpha_over(linear_burn(base_imgs[5], composite_image_data), composite_image_data, 0.60) - # plt.imshow(composite_image_data) - - # plt.figure("Light Ops with mid bands") - mix_img = composite_alpha_over(exclusion(base_imgs[3], composite_image_data), composite_image_data, 0.95) - mix_img = composite_alpha_over(linear_light(base_imgs[2], mix_img), mix_img, 1.0)**0.75 - composite_image_data = composite_alpha_over(mix_img, composite_image_data, 0.5) - # plt.imshow(composite_image_data) - - # plt.figure("Mix in mid 131") - composite_image_data = composite_alpha_over(base_imgs[2], composite_image_data, 0.25) # Mix in a small amount of 131Ã… for the nice streamers - # plt.imshow(composite_image_data) - - # plt.figure("Mid diff") - mix_img = difference(base_imgs[3], base_imgs[2]) # 171Ã… - 131Ã… - # plt.imshow(mix_img) - - # plt.figure("HSL ops with mid") - comp_hsl = rgb_to_hsl(composite_image_data) - mix_img_hsl = rgb_to_hsl(mix_img) - del mix_img - base_img2_hsl = rgb_to_hsl(base_imgs[2]) - base_img3_hsl = rgb_to_hsl(base_imgs[3]) - comp_hsl = np.copy(comp_hsl) - comp_hsl[:,:,0] += 0.025*mix_img_hsl[:,:,0] # Rotate hue based on mix_1_hsl - del mix_img_hsl - comp_hsl[:,:,0][comp_hsl[:,:,0] > 360.0] -= 360.0 - comp_hsl[:,:,1] -= 0.1*base_img3_hsl[:,:,1] # Reduce saturation based on base_img3 - comp_hsl[:,:,1] = np.clip(comp_hsl[:,:,1], 0.0, 100.0) - comp_hsl[:,:,2] += 0.5*base_img2_hsl[:,:,2] # Boost luminance based on base_img2 - comp_hsl[:,:,1] = np.clip(comp_hsl[:,:,2], 0.0, 100.0) - composite_image_data = hsl_to_rgb(comp_hsl) - del comp_hsl - composite_image_data = saturation(composite_image_data, 1.0, 0.8, 0.1) # Remove some blue and green - # plt.imshow(composite_image_data) - - # plt.figure("Hue shifted 131") - mix_img_hsl = rgb_to_hsl(base_imgs[1]) - mix_img_hsl[:,:,0] -= 50 - mix_img_hsl[:,:,0][mix_img_hsl[:,:,0] < 0.0] += 360.0 - mix_img = hsl_to_rgb(mix_img_hsl) - del mix_img_hsl - # plt.imshow(mix_img) - - # plt.figure("Hard Light with hue shifted 131") - composite_image_data = composite_alpha_over(hard_light(mix_img, composite_image_data), composite_image_data, 0.2) - # plt.imshow(composite_image_data) - - # plt.figure("Adjusted 094") - mix_img = saturation(base_imgs[0], 1.15, 1.2, 1.05) - mix_img = contrast(mix_img, 1.5, 0.0) - # plt.imshow(mix_img) - - # plt.figure("Color Dodge with adjusted 094") - composite_image_data = composite_alpha_over(color_dodge(mix_img, composite_image_data), composite_image_data, 0.5) - # plt.imshow(composite_image_data) - - # plt.figure("Exclusion with adjusted 094") - composite_image_data = composite_alpha_over(exclusion(mix_img, composite_image_data), composite_image_data, 0.65) - # plt.imshow(composite_image_data) - - del mix_img - - # plt.figure("Final Image") - composite_image_data = saturation(composite_image_data, 1.0, 1.1, 1.2) - composite_image_data = contrast(composite_image_data, 1.20, 0.00) - # plt.imshow(composite_image_data) - - # We have our final image - # plt.show() - - # Now shrink the component images and assemble them alongside the composite. - new_dimx = composite_image_data.shape[1] // 3 - new_dimy = composite_image_data.shape[0] // 3 - # Enlarge the composite to fit the new images - composite_image_data = np.pad(composite_image_data, ((0,0),(new_dimx, new_dimx),(0,0))) - for i in range(6): - img = base_imgs[i] - img = bin_ndarray(img, (new_dimy, new_dimx, 3)) # Shrink down to 1/3 for assembly - img = contrast(img, 1.25, 0.0) - xdimoff = i%2 * (composite_image_data.shape[1] - new_dimx) - ydimoff = i//2*new_dimy - composite_image_data[ydimoff:ydimoff+new_dimy, xdimoff:xdimoff+new_dimx, :] = img - - img = Image.fromarray((255 * composite_image_data).astype('uint8')) - timestring = datetime.datetime.fromtimestamp(timestamp, tz = datetime.UTC).strftime('%Y-%m-%d %H:%M:%S') - ImageDraw.Draw(img).text((602, 15), f"NOAA GOES Satellite SUVI Composite - {timestring} UTC",(255,255,255), font = fnt1) - for i in range(6): # Draw component angstrom labels - if i%2 == 0: - xdimtxtoff = 5 - else: - xdimtxtoff = composite_image_data.shape[1] - 44 - ydimtxtoff = i//2*new_dimy + new_dimy / 2.0 - 14 - ImageDraw.Draw(img).text((xdimtxtoff, ydimtxtoff), image_names[i], font = fnt2) - img.save(filepath, quality = 95) - result_queue.put(("Created", timestamp)) - - except KeyboardInterrupt: - return - except Exception as e: - traceback.print_exception(e) - result_queue.put((e, timestamp)) - - -image_names = ["094A", "131A", "171A", "195A", "284A", "304A"] - -if __name__ == "__main__": - stored_fits_dirs = [r"..\Data\goes16\l2\data", r"..\Data\goes18\l2\data"] - processed_images_dirs = [r"..\composite\goes16", r"..\composite\goes18"] - starttime = calendar.timegm(datetime.datetime(2024, 1, 1, tzinfo=datetime.timezone.utc).timetuple()) - stoptime = calendar.timegm(datetime.datetime(2025, 1, 1, tzinfo=datetime.timezone.utc).timetuple()) - - regex_filename = r"dr_suvi-l2-ci\d{3}_g(16|18)_s\S*\_f.fits$" - - nworkers = 16 - max_time_gap = 3 - fill_missing_data = False - ignore_errors = True - - if ignore_errors: - regex_filename = r"dr_suvi-l2-ci\d{3}_g(16|18)_s\S*.fits$" - - file_prefixes = ["dr_suvi-l2-ci" + n[:-1] for n in image_names] - - # Testing - # stored_fits_dirs = [r"..\fits_test_2024"] - - work_queue = Queue(maxsize = nworkers) - result_queue = Queue() - workers = [] - for i in range(nworkers): - p = Process(target = generate_composite, args = (work_queue, result_queue), daemon=True) - p.start() - workers.append(p) - - lowpriority() - ncreated = 0 - nexists = 0 - nfailed = 0 - black_image = Image.fromarray(np.zeros((1080, 1920, 3), dtype='uint8')) - try: - for stored_fits_dir, processed_images_dir in zip(stored_fits_dirs, processed_images_dirs): - files_by_timestamp = defaultdict(list) - num_found_files = 0 - os.makedirs(processed_images_dir, exist_ok=True) - filename_tester = re.compile(regex_filename) - print(f"Searching for FITS files in: {stored_fits_dir}") - for root, dirs, files in tqdm.tqdm(os.walk(stored_fits_dir), desc="Searching"): - for f in files: - if filename_tester.match(f): - file_parts = f.split("_") - # measurement = file_parts[1] - # sattelite = file_parts[2] - measure_end_time = int(datetime.datetime.strptime(file_parts[4][1:16] + " +0000", "%Y%m%dT%H%M%S %z").timestamp()) - if (measure_end_time >= starttime) and (measure_end_time < stoptime): - files_by_timestamp[measure_end_time].append(os.path.join(root,f)) - num_found_files += 1 - - print(f"Found {num_found_files} FITS files. Starting conversion.") - if num_found_files == 0: - exit(3) - - sorted_times = sorted(list(files_by_timestamp.keys())) - min_time = sorted_times[0] - max_time = sorted_times[-1] - diff_times = np.diff(sorted_times) - unique, counts = np.unique(diff_times, return_counts=True) - interval = unique[0] # This is the amount of time between each sample in seconds - assert np.sum((unique % interval) > 0) == 0 # Ensure all our timestamps align perfectly with our interval - - last_good_files = None - last_good_file_times = None - # Lets find the first and last timestamps in the sorted_times from our files which actually have a full set of 6/6 images available - # This check will prevent partially downloaded sets of data from generating composite images which have "filled in" data from detected gaps - # which would have been later filled with downloaded imagery. - f = None - l = None - for i, timestamp in enumerate(sorted_times): - if len(files_by_timestamp[timestamp]) == 6: - f = i - break - for i, timestamp in enumerate(reversed(sorted_times)): - if len(files_by_timestamp[timestamp]) == 6: - l = len(sorted_times) - 1 - i - break - assert f is not None # Check to make sure we found valid indices - assert l is not None - assert f != l - sorted_times = sorted_times[f:l] # Limit our composite image generation to only files within the valid range - last_good_files = files_by_timestamp[sorted_times[0]] - for timestamp in tqdm.tqdm(sorted_times, desc="Creating Composite Solar Images"): - if timestamp < min_time or timestamp > max_time: - continue - # Collect completed jobs and record completion status - while True: - try: - result = result_queue.get_nowait() - if result[0] == "Exists": - nexists += 1 - elif result[0] == "Created": - ncreated += 1 - else: - print(f"A worker encountered an exception on job {result[1]}: {result[0]}") - nfailed += 1 - except queue.Empty: - break - - # Submit new jobs - files_this_timestamp = files_by_timestamp[timestamp] - files_this_timestamp = sorted(files_this_timestamp) - if (not len(files_this_timestamp) == 6): - if fill_missing_data: - print(f"Invalid or incomplete sensor records for {timestamp} | {datetime.datetime.fromtimestamp(timestamp, datetime.UTC).strftime('%Y-%m-%d %H:%M:%S %Z')} - {len(files_this_timestamp)}/6 filling from last good data.") - files_for_job = [] - for i, prefix in enumerate(file_prefixes): - found = False - for f in files_this_timestamp: - filename = os.path.split(f)[-1] - if filename.startswith(prefix): - files_for_job.append(f) - last_good_files[i] = f - last_good_file_times[i] = timestamp - found = True - break - if not found: # We did not find this prefix, use the last good file - time_gap = (timestamp - last_good_file_times[i]) // interval - if time_gap <= max_time_gap: - files_for_job.append(last_good_files[i]) - else: - print(f"Detected a gap of {time_gap} frames at {timestamp} | {datetime.datetime.fromtimestamp(timestamp, datetime.UTC).strftime('%Y-%m-%d %H:%M:%S %Z')}") - if len(files_for_job) == 6: - work_queue.put((files_for_job, timestamp, processed_images_dir, True)) - else: # We have a complete file set, update the last_good_files - last_good_files = files_this_timestamp - last_good_file_times = [timestamp for _ in last_good_files] - work_queue.put((files_this_timestamp, timestamp, processed_images_dir, ignore_errors)) - - except KeyboardInterrupt: - print("Finishing current jobs and exiting") - - for _ in range(nworkers): - try: - work_queue.put(None, timeout=10.0) - except: - break - - for w in workers: - # w.join(10.0) - w.join() - - print(f"Created {ncreated} | Already had {nexists} | Failed {nfailed}") \ No newline at end of file +import os +import sys +import calendar +import datetime +from collections import defaultdict +from multiprocessing import Queue, Process +import warnings +import queue +import traceback + +import tqdm +from PIL import Image, ImageDraw, ImageFont +import numpy as np +from matplotlib import pyplot as plt +from astropy.io import fits +import palettable + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +def lowpriority(): + """ Set the priority of the process to below-normal.""" + + import sys + try: + sys.getwindowsversion() + except AttributeError: + isWindows = False + else: + isWindows = True + + if isWindows: + # Based on: + # "Recipe 496767: Set Process Priority In Windows" on ActiveState + # http://code.activestate.com/recipes/496767/ + import win32api,win32process,win32con # pywin32 + + pid = win32api.GetCurrentProcessId() + handle = win32api.OpenProcess(win32con.PROCESS_ALL_ACCESS, True, pid) + win32process.SetPriorityClass(handle, win32process.BELOW_NORMAL_PRIORITY_CLASS) + else: + import os + + os.nice(1) + +def bin_ndarray(ndarray, new_shape, operation='mean'): + """ + Bins an ndarray in all axes based on the target shape, by summing or + averaging. + + Number of output dimensions must match number of input dimensions and + new axes must divide old ones. + + Example + ------- + >>> m = np.arange(0,100,1).reshape((10,10)) + >>> n = bin_ndarray(m, new_shape=(5,5), operation='sum') + >>> print(n) + + [[ 22 30 38 46 54] + [102 110 118 126 134] + [182 190 198 206 214] + [262 270 278 286 294] + [342 350 358 366 374]] + + """ + operation = operation.lower() + if not operation in ['sum', 'mean']: + raise ValueError("Operation not supported.") + if ndarray.ndim != len(new_shape): + raise ValueError("Shape mismatch: {} -> {}".format(ndarray.shape, + new_shape)) + compression_pairs = [(d, c//d) for d,c in zip(new_shape, + ndarray.shape)] + flattened = [l for p in compression_pairs for l in p] + ndarray = ndarray.reshape(flattened) + for i in range(len(new_shape)): + op = getattr(ndarray, operation) + ndarray = op(-1*(i+1)) + return ndarray + +def gamma_correct(fun): + def wrapper(*args, **kwargs): + args = list(args) + args[0] = np.power(args[0], 2.2) + args[1] = np.power(args[1], 2.2) + args = tuple(args) + result = fun(*args, **kwargs) + return np.power(result, 1/2.2) + return wrapper + +def clip_color(fun): + def wrapper(*args, **kwargs): + return np.clip(fun(*args, **kwargs), 0.0, 1.0) + return wrapper + +# linear_srgb_matrix = np.array([[0.4124, 0.3576, 0.1805], +# [0.2126, 0.7152, 0.0722], +# [0.0193, 0.1192, 0.9505]]) + +# linear_srgb_matrix_inv = np.array([[ 3.2406, -1.5372, -0.4986], +# [-0.9689, 1.8758, 0.0415], +# [ 0.0557, -0.2040, 1.0570]]) + +# def linear_color_correction(fun): +# def wrapper(*args, **kwargs): +# args = list(args) +# inds = args[0] <= 0.04045 +# ninds = args[0] > 0.04045 +# for i in range(2): +# args[i][inds] = args[i][inds] / 12.92 +# args[i][ninds] = np.power((args[i][ninds] + 0.055) / 1.055, 2.4) +# for x in range(args[i].shape[0]): +# for y in range(args[i].shape[1]): +# args[i][x,y,:] = np.matmul(linear_srgb_matrix, args[i][x,y,:]) +# args = tuple(args) +# result = fun(*args, **kwargs) + +# for x in range(result.shape[0]): +# for y in range(result.shape[1]): +# result[x,y,:] = np.matmul(linear_srgb_matrix_inv, result[x,y,:]) +# inds = result <= 0.0031308 +# ninds = result > 0.0031308 +# result[inds] = result[inds] * 12.92 +# result[ninds] = np.power(result[ninds], 1.0/2.4) * 1.055 - 0.055 +# return result +# return wrapper + + +@gamma_correct +def composite_alpha_over(F, B, alpha_F, alpha_B = 1): + return (F*alpha_F + B*alpha_B*(1-alpha_F)) / (alpha_F + alpha_B*(1-alpha_F)) + +def composite_alpha_blend(F, B, alpha): + return F*alpha + B*(1-alpha) + +def linear_burn(F, B): + burn = F + B - 1 + burn[burn < 0.0] = 0.0 + return burn + +def difference(F, B): + d = np.abs(F - B) + if d.shape[2] == 4: # Preserve alpha of base image + d[:,:,3] = B[:,:,3] + return d + +@clip_color +def linear_light(F, B): + result = np.zeros_like(F) + inds = F <= 0.5 + ninds = F > 0.5 + result[inds] = B[inds] + 2.0 * F[inds] - 1 + result[ninds] = 2.0 * (F[ninds] - 0.5) + B[ninds] + return result + +@clip_color +def hard_light(F, B): + result = np.zeros_like(F) + inds = B < 0.5 + ninds = B >= 0.5 + result[inds] = 2 * F[inds] * B[inds] + result[ninds] = 1 - (2*(1 - F[ninds])*(1 - B[ninds])) + return result + +@clip_color +def color_dodge(F, B): + return B / (1.000001 - F) + +@clip_color +def exclusion(F, B): + d = F + B - 2*F*B + if d.shape[2] == 4: # Preserve alpha of base image + d[:,:,3] = B[:,:,3] + return d + +@clip_color +def saturation(img, R, G, B): + img[:,:,0] *= R + img[:,:,1] *= G + img[:,:,2] *= B + return img + +@clip_color +def contrast(img, c, b): + return (img - 0.5) * c + 0.5 + b*c + +def rgb_to_hsl(img): + r = img[:,:,0] + g = img[:,:,1] + b = img[:,:,2] + cmax = np.copy(r) + cmax[g > cmax] = g[g > cmax] + cmax[b > cmax] = b[b > cmax] + cmin = np.copy(r) + cmin[g < cmin] = g[g < cmin] + cmin[b < cmin] = b[b < cmin] + delta = cmax - cmin + + + # Calc hue + hue = np.zeros_like(r) + inds = cmax == r + with warnings.catch_warnings(): + warnings.filterwarnings('ignore') + hue[inds] = 60 * np.mod((g[inds]-b[inds])/delta[inds], 6) + inds = cmax == g + hue[inds] = 60 * ((b[inds]-r[inds])/delta[inds] + 2) + inds = cmax == b + hue[inds] = 60 * ((r[inds]-g[inds])/delta[inds] + 4) + hue[np.isnan(hue)] = 0 + hue[hue < 0] = hue[hue < 0] + 360 # Make negative hue values positive behind 360 + + # Calc lightness / luminance + luminance = (cmax + cmin) / 2.0 + + # Calc saturation + saturation = np.zeros_like(r) + inds = delta != 0 + saturation[inds] = delta[inds] / (1.0 - np.abs(2 * luminance[inds] - 1.0)) + + # Multiply luminance and saturation by 100 to scale them to the appropriate range (0-100) + saturation *= 100.0 + luminance *= 100.0 + + + if img.shape[2] == 4: # Preserve alpha of original image + return np.stack([hue, saturation, luminance, img[:,:,3]], -1) + else: + return np.stack([hue, saturation, luminance], -1) + +def hsl_to_rgb(img): + hue = img[:,:,0] + saturation = img[:,:,1] + luminance = img[:,:,2] + saturation /= 100.0 + luminance /= 100.0 + c = (1.0 - np.abs(2.0 * luminance - 1.0)) * saturation + x = c * (1.0 - np.abs((hue/60.0) % 2.0 - 1.0)) + m = luminance - c / 2.0 + r = np.zeros_like(hue) + g = np.zeros_like(hue) + b = np.zeros_like(hue) + + inds1 = np.logical_and(0.0 <= hue, hue < 60.0) + r[inds1] = c[inds1] + g[inds1] = x[inds1] + inds2 = np.logical_and(60.0 <= hue, hue < 120.0) + r[inds2] = x[inds2] + g[inds2] = c[inds2] + inds3 = np.logical_and(120.0 <= hue, hue < 180.0) + g[inds3] = c[inds3] + b[inds3] = x[inds3] + inds4 = np.logical_and(180.0 <= hue, hue < 240.0) + g[inds4] = x[inds4] + b[inds4] = c[inds4] + inds5 = np.logical_and(240.0 <= hue, hue < 300.0) + r[inds5] = x[inds5] + b[inds5] = c[inds5] + inds6 = np.logical_and(240.0 <= hue, hue < 300.0) + r[inds6] = c[inds6] + b[inds6] = x[inds6] + + r += m + g += m + b += m + + if img.shape[2] == 4: # Preserve alpha of original image + return np.stack([r, g, b, img[:,:,3]], -1) + else: + return np.stack([r, g, b], -1) + +#: Band labels drawn on the composite, and the order `composite_from_arrays` +#: expects its six arrays in. +image_names = ["094A", "131A", "171A", "195A", "284A", "304A"] + + +# --- Composite rendering ------------------------------------------------------- +# +# The blending below is the product's visual identity and is deliberately +# untouched. It was extracted from the worker loop so that callers holding six +# bands *in memory* -- the bench, rendering frames that were reconstructed rather +# than read from disk -- can produce the same image without first writing FITS. + +#: Radiance-to-colour mapping, per band, in the order of `image_names`. +VMINS = [0.050, 0.05, 00.100, 00.10, 00.100, 00.1] +VMAXS = [8.000, 8.00, 20.000, 30.00, 40.000, 90.0] +GAMMAS = [0.375, 0.40, 00.425, 00.45, 00.475, 00.5] +#: Trim applied to each 1280x1280 band to reach the final aspect ratio. +TRIM_X = 64 +TRIM_Y = 100 + +_render_cache = {} + + +def _render_assets(): + """Colormaps and fonts, built once per process. + + The first four colormaps are mutated to start at true black; that mutation is + global to the palettable/matplotlib objects, so it must happen exactly once. + """ + if not _render_cache: + cmaps = [palettable.cmocean.sequential.Ice_5.mpl_colormap, + palettable.cmocean.sequential.Ice_20.mpl_colormap, + palettable.cmocean.sequential.Turbid_5_r.mpl_colormap, + palettable.cmocean.sequential.Turbid_20_r.mpl_colormap, + plt.colormaps.get_cmap('gist_heat'), + plt.colormaps.get_cmap('afmhot')] + # Start these maps at perfect black rather than a very dark colour. + for cmi in range(0, 4): + for c in ['red', 'green', 'blue']: + for i in range(3): + cmaps[cmi]._segmentdata[c][0][i] = 0.0 + font_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), + "OpenSans-Regular.ttf") + _render_cache.update( + cmaps=cmaps, + fnt1=ImageFont.truetype(font_path, size=24), + fnt2=ImageFont.truetype(font_path, size=16), + ) + return _render_cache + + +def composite_from_arrays(arrays, timestamp): + """Render one composite from six in-memory bands. Returns a PIL Image. + + `arrays` are the six 1280x1280 radiance arrays in `image_names` order. + """ + assets = _render_assets() + cmaps, fnt1, fnt2 = assets["cmaps"], assets["fnt1"], assets["fnt2"] + + base_imgs = [] + for i in range(6): + raw_data = np.flip(np.asarray(arrays[i], dtype=np.float32), 0)[TRIM_Y:-TRIM_Y, TRIM_X:-TRIM_X] + raw_data[raw_data < 0.0] = 0.0 # Remove non-zero data because it doesn't makes sense (supposed to be std Radiance) + base_imgs.append(cmaps[i](np.clip((raw_data - VMINS[i]) / VMAXS[i], 0, 1.0)**GAMMAS[i])[:,:,:3]) + # plt.figure(image_names[i]) + # plt.imshow(base_imgs[-1]) + + + # plt.figure("Initial Blend") + composite_image_data = composite_alpha_over(base_imgs[4], base_imgs[5], 0.2)**0.5 + # plt.imshow(composite_image_data) + + # plt.figure("Linear Burn with 304") + composite_image_data = composite_alpha_over(linear_burn(base_imgs[5], composite_image_data), composite_image_data, 0.60) + # plt.imshow(composite_image_data) + + # plt.figure("Light Ops with mid bands") + mix_img = composite_alpha_over(exclusion(base_imgs[3], composite_image_data), composite_image_data, 0.95) + mix_img = composite_alpha_over(linear_light(base_imgs[2], mix_img), mix_img, 1.0)**0.75 + composite_image_data = composite_alpha_over(mix_img, composite_image_data, 0.5) + # plt.imshow(composite_image_data) + + # plt.figure("Mix in mid 131") + composite_image_data = composite_alpha_over(base_imgs[2], composite_image_data, 0.25) # Mix in a small amount of 131Ã… for the nice streamers + # plt.imshow(composite_image_data) + + # plt.figure("Mid diff") + mix_img = difference(base_imgs[3], base_imgs[2]) # 171Ã… - 131Ã… + # plt.imshow(mix_img) + + # plt.figure("HSL ops with mid") + comp_hsl = rgb_to_hsl(composite_image_data) + mix_img_hsl = rgb_to_hsl(mix_img) + del mix_img + base_img2_hsl = rgb_to_hsl(base_imgs[2]) + base_img3_hsl = rgb_to_hsl(base_imgs[3]) + comp_hsl = np.copy(comp_hsl) + comp_hsl[:,:,0] += 0.025*mix_img_hsl[:,:,0] # Rotate hue based on mix_1_hsl + del mix_img_hsl + comp_hsl[:,:,0][comp_hsl[:,:,0] > 360.0] -= 360.0 + comp_hsl[:,:,1] -= 0.1*base_img3_hsl[:,:,1] # Reduce saturation based on base_img3 + comp_hsl[:,:,1] = np.clip(comp_hsl[:,:,1], 0.0, 100.0) + comp_hsl[:,:,2] += 0.5*base_img2_hsl[:,:,2] # Boost luminance based on base_img2 + comp_hsl[:,:,1] = np.clip(comp_hsl[:,:,2], 0.0, 100.0) + composite_image_data = hsl_to_rgb(comp_hsl) + del comp_hsl + composite_image_data = saturation(composite_image_data, 1.0, 0.8, 0.1) # Remove some blue and green + # plt.imshow(composite_image_data) + + # plt.figure("Hue shifted 131") + mix_img_hsl = rgb_to_hsl(base_imgs[1]) + mix_img_hsl[:,:,0] -= 50 + mix_img_hsl[:,:,0][mix_img_hsl[:,:,0] < 0.0] += 360.0 + mix_img = hsl_to_rgb(mix_img_hsl) + del mix_img_hsl + # plt.imshow(mix_img) + + # plt.figure("Hard Light with hue shifted 131") + composite_image_data = composite_alpha_over(hard_light(mix_img, composite_image_data), composite_image_data, 0.2) + # plt.imshow(composite_image_data) + + # plt.figure("Adjusted 094") + mix_img = saturation(base_imgs[0], 1.15, 1.2, 1.05) + mix_img = contrast(mix_img, 1.5, 0.0) + # plt.imshow(mix_img) + + # plt.figure("Color Dodge with adjusted 094") + composite_image_data = composite_alpha_over(color_dodge(mix_img, composite_image_data), composite_image_data, 0.5) + # plt.imshow(composite_image_data) + + # plt.figure("Exclusion with adjusted 094") + composite_image_data = composite_alpha_over(exclusion(mix_img, composite_image_data), composite_image_data, 0.65) + # plt.imshow(composite_image_data) + + del mix_img + + # plt.figure("Final Image") + composite_image_data = saturation(composite_image_data, 1.0, 1.1, 1.2) + composite_image_data = contrast(composite_image_data, 1.20, 0.00) + # plt.imshow(composite_image_data) + + # We have our final image + # plt.show() + + # Now shrink the component images and assemble them alongside the composite. + new_dimx = composite_image_data.shape[1] // 3 + new_dimy = composite_image_data.shape[0] // 3 + # Enlarge the composite to fit the new images + composite_image_data = np.pad(composite_image_data, ((0,0),(new_dimx, new_dimx),(0,0))) + for i in range(6): + img = base_imgs[i] + img = bin_ndarray(img, (new_dimy, new_dimx, 3)) # Shrink down to 1/3 for assembly + img = contrast(img, 1.25, 0.0) + xdimoff = i%2 * (composite_image_data.shape[1] - new_dimx) + ydimoff = i//2*new_dimy + composite_image_data[ydimoff:ydimoff+new_dimy, xdimoff:xdimoff+new_dimx, :] = img + + img = Image.fromarray((255 * composite_image_data).astype('uint8')) + timestring = datetime.datetime.fromtimestamp(timestamp, tz = datetime.UTC).strftime('%Y-%m-%d %H:%M:%S') + ImageDraw.Draw(img).text((602, 15), f"NOAA GOES Satellite SUVI Composite - {timestring} UTC",(255,255,255), font = fnt1) + for i in range(6): # Draw component angstrom labels + if i%2 == 0: + xdimtxtoff = 5 + else: + xdimtxtoff = composite_image_data.shape[1] - 44 + ydimtxtoff = i//2*new_dimy + new_dimy / 2.0 - 14 + ImageDraw.Draw(img).text((xdimtxtoff, ydimtxtoff), image_names[i], font = fnt2) + img = Image.fromarray((255 * composite_image_data).astype('uint8')) + timestring = datetime.datetime.fromtimestamp(timestamp, tz = datetime.UTC).strftime('%Y-%m-%d %H:%M:%S') + ImageDraw.Draw(img).text((602, 15), f"NOAA GOES Satellite SUVI Composite - {timestring} UTC",(255,255,255), font = fnt1) + for i in range(6): # Draw component angstrom labels + if i%2 == 0: + xdimtxtoff = 5 + else: + xdimtxtoff = composite_image_data.shape[1] - 44 + ydimtxtoff = i//2*new_dimy + new_dimy / 2.0 - 14 + ImageDraw.Draw(img).text((xdimtxtoff, ydimtxtoff), image_names[i], font = fnt2) + return img + + +def generate_composite(work_queue, result_queue): + """Worker: read six bands from disk and write their composite.""" + lowpriority() + while True: + try: + job = work_queue.get() + if job is None: + result_queue.cancel_join_thread() + return + files_this_timestamp, timestamp, processed_images_dir, synthetic_data = job + synthetic_filepath = os.path.join(processed_images_dir, f"Composite-{int(timestamp)}_s.jpg") + normal_filepath = os.path.join(processed_images_dir, f"Composite-{int(timestamp)}.jpg") + filepath = None + if synthetic_data: + if os.path.isfile(synthetic_filepath): + result_queue.put(("Exists", timestamp)) + continue + filepath = synthetic_filepath + else: + if os.path.isfile(synthetic_filepath): + os.remove(synthetic_filepath) + if os.path.isfile(normal_filepath): + result_queue.put(("Exists", timestamp)) + continue + filepath = normal_filepath + + arrays = [fits.getdata(f) for f in files_this_timestamp] + img = composite_from_arrays(arrays, timestamp) + img.save(filepath, quality = 95) + result_queue.put(("Created", timestamp)) + + except KeyboardInterrupt: + return + except Exception as e: + traceback.print_exception(e) + result_queue.put((e, timestamp)) + + +if __name__ == "__main__": + import argparse + + from suvi import db, paths + + parser = argparse.ArgumentParser( + description="Render SUVI composite images from indexed FITS frames.", + ) + parser.add_argument("--db", default=None, help="SQLite index (default: $SUVI_DB)") + parser.add_argument("--root", default=None, help="archive root (default: $SUVI_DATA_ROOT)") + parser.add_argument("--out", default=None, help="composite output root") + parser.add_argument("--satellites", default="16,18") + parser.add_argument("--start", default="2024-01-01", help="UTC date, inclusive") + parser.add_argument("--stop", default="2025-01-01", help="UTC date, exclusive") + parser.add_argument( + "--detector", + default="header_v1", + help="detector run whose 'good' verdicts select frames, or 'none' to use every " + "indexed frame. Default header_v1: it is conservative, catching genuine " + "dropouts without discarding good data. Do not default this to geometry_v1 " + "-- over sampled days of 2024 it rejects 36.9%% of frames, entire days at a " + "time, because its centre test tracks the brightness centroid rather than " + "the disc. Run bench.py before changing this.", + ) + parser.add_argument("--workers", type=int, default=16) + parser.add_argument("--max-time-gap", type=int, default=3, + help="slots a missing band may be carried forward when filling") + parser.add_argument("--fill-missing", action="store_true", + help="carry the last good frame into incomplete timestamps") + args = parser.parse_args() + + def parse_day(text): + return calendar.timegm( + datetime.datetime.strptime(text, "%Y-%m-%d") + .replace(tzinfo=datetime.timezone.utc) + .timetuple() + ) + + archive_root = args.root or paths.data_root() + composite_root = args.out or os.path.join(os.path.dirname(archive_root), "composite") + satellites = [int(part) for part in args.satellites.split(",")] + starttime, stoptime = parse_day(args.start), parse_day(args.stop) + nworkers = args.workers + max_time_gap = args.max_time_gap + fill_missing_data = args.fill_missing + + conn = db.connect(args.db, readonly=True) + run_id = None + if args.detector != "none": + run_id = db.latest_run_id(conn, args.detector) + if run_id is None: + raise SystemExit( + f"No detector run named {args.detector!r} in the index. " + f"Run `filter_FITS.py detect --detectors {args.detector}` first, " + f"or pass --detector none to use every indexed frame." + ) + + file_prefixes = ["dr_suvi-l2-ci" + n[:-1] for n in image_names] + + work_queue = Queue(maxsize=nworkers) + result_queue = Queue() + workers = [] + for i in range(nworkers): + p = Process(target=generate_composite, args=(work_queue, result_queue), daemon=True) + p.start() + workers.append(p) + + lowpriority() + ncreated = 0 + nexists = 0 + nfailed = 0 + try: + for satellite in satellites: + processed_images_dir = os.path.join(composite_root, f"goes{satellite}") + os.makedirs(processed_images_dir, exist_ok=True) + + # Frame selection is a query against the index rather than a filename + # regex. The old `_f.fits$` pattern is gone along with the suffixes it + # matched; see migrate_unrename.py. + if run_id is None: + rows = conn.execute( + "SELECT path, wavelength, t_start FROM frame " + "WHERE satellite = ? AND t_start >= ? AND t_start < ? " + "ORDER BY t_start, wavelength", + (satellite, starttime, stoptime), + ).fetchall() + else: + rows = db.good_slots(conn, run_id, starttime, stoptime, satellite) + + files_by_timestamp = defaultdict(list) + for row in rows: + files_by_timestamp[row["t_start"]].append( + paths.abspath(row["path"], archive_root) + ) + num_found_files = sum(len(v) for v in files_by_timestamp.values()) + print(f"goes{satellite}: {num_found_files} frames selected " + f"({args.detector}) across {len(files_by_timestamp)} timestamps") + if num_found_files == 0: + continue + + sorted_times = sorted(files_by_timestamp) + interval = paths.CADENCE + + # Trim to the first and last timestamps holding a full set of six bands, so + # a partially downloaded tail does not get permanently filled from a gap. + first = next((i for i, t in enumerate(sorted_times) + if len(files_by_timestamp[t]) == 6), None) + last = next((len(sorted_times) - 1 - i + for i, t in enumerate(reversed(sorted_times)) + if len(files_by_timestamp[t]) == 6), None) + if first is None or last is None or first >= last: + print(f" no complete six-band timestamps for goes{satellite}; skipping") + continue + sorted_times = sorted_times[first:last] + + last_good_files = list(files_by_timestamp[sorted_times[0]]) + last_good_file_times = [sorted_times[0]] * len(last_good_files) + + for timestamp in tqdm.tqdm(sorted_times, desc=f"goes{satellite} composites"): + while True: + try: + result = result_queue.get_nowait() + if result[0] == "Exists": + nexists += 1 + elif result[0] == "Created": + ncreated += 1 + else: + print(f"A worker failed on job {result[1]}: {result[0]}") + nfailed += 1 + except queue.Empty: + break + + files_this_timestamp = sorted(files_by_timestamp[timestamp]) + if len(files_this_timestamp) == 6: + last_good_files = files_this_timestamp + last_good_file_times = [timestamp] * 6 + work_queue.put((files_this_timestamp, timestamp, processed_images_dir, False)) + elif fill_missing_data: + files_for_job = [] + for i, prefix in enumerate(file_prefixes): + found = next( + (f for f in files_this_timestamp + if os.path.basename(f).startswith(prefix)), + None, + ) + if found is not None: + files_for_job.append(found) + last_good_files[i] = found + last_good_file_times[i] = timestamp + elif (timestamp - last_good_file_times[i]) // interval <= max_time_gap: + files_for_job.append(last_good_files[i]) + if len(files_for_job) == 6: + work_queue.put((files_for_job, timestamp, processed_images_dir, True)) + + except KeyboardInterrupt: + print("Finishing current jobs and exiting") + + for _ in range(nworkers): + try: + work_queue.put(None, timeout=10.0) + except Exception: + break + + for w in workers: + w.join() + + conn.close() + print(f"Created {ncreated} | Already had {nexists} | Failed {nfailed}") diff --git a/migrate_unrename.py b/migrate_unrename.py new file mode 100644 index 0000000..1b51e6e --- /dev/null +++ b/migrate_unrename.py @@ -0,0 +1,366 @@ +#!/usr/bin/env python +"""One-time migration: strip _f/_e suffixes from the archive, into the SQLite index. + +``filter_FITS.py`` used to record its verdict by renaming files -- ``_f`` for passed, +``_e`` for rejected -- and writing a ``_e.jpg`` diagnostic plot alongside rejects. +That made the archive's filenames a mutable database with one column, coupled three +scripts to a naming convention, and threw away every intermediate score. + +This restores every file to the name NOAA published it under, after recording the +existing verdicts in the index. + +Safety, because this is the one irreversible step: + +* Dry run by default; ``--apply`` is required to rename anything. +* The legacy verdicts are exported to a standalone CSV **before** any rename. Every + other table can be rebuilt by re-walking or re-reading the archive, but once the + suffixes are gone these labels exist nowhere else. +* A rename that would overwrite an existing file is skipped and reported, never + forced. +* Idempotent and resumable: re-running finds nothing left to do. + +Note that these labels are stale -- they were written in May 2024, and the filter was +retuned that July -- so they are preserved as history, under the run name +``legacy_filter``, not as ground truth. + + migrate_unrename.py --wavelength 171 --year 2024 # dry run, one subtree + migrate_unrename.py --apply # the whole archive +""" + +import argparse +import csv +import gzip +import os +import sys +import time +from collections import Counter + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +from suvi import db, paths, vfs + +LEGACY_RUN_NAME = "legacy_filter" + + +def chunk_key(wavelength, year): + """Index key recording that one (band, year) has been fully migrated.""" + return f"unrename_done:ci{wavelength:03d}:{year if year is not None else 'all'}" + + +def walk_archive(root, satellites, wavelengths, years=None): + """Yield (directory, filename, FrameName) for every SUVI frame found. + + When `years` is given the unwanted year directories are pruned from the walk + rather than merely skipped, so restricting the year genuinely restricts how much + of the filesystem is touched. That matters here for more than speed: see + :func:`release_handles`. + """ + def fail(error): + # os.walk swallows errors by default. On this mount a transient ENFILE + # would then make the walk yield nothing, and the caller would conclude the + # archive was already migrated and move on -- silently skipping real work. + # An unreadable directory has to stop the chunk, not look like an empty one. + raise error + + for base in paths.wavelength_dirs(root, satellites, wavelengths): + if not os.path.isdir(base): + continue + for directory, dirnames, filenames in os.walk(base, onerror=fail): + if years and directory == base: + dirnames[:] = [d for d in dirnames if not d.isdigit() or int(d) in years] + for filename in filenames: + name = paths.parse_frame_filename(filename) + if name is not None: + yield directory, filename, name + + +def archive_years(root, satellites, wavelength): + """Years present on disk for one band, across all satellites.""" + years = set() + for satellite in satellites: + base = os.path.join( + root, f"goes{satellite:d}", "l2", "data", f"suvi-l2-ci{wavelength:03d}" + ) + if not os.path.isdir(base): + continue + try: + years.update(int(d) for d in os.listdir(base) if d.isdigit()) + except OSError: + continue + return sorted(years) + + +#: Frames a chunk must have touched before reclaiming is worth its cost. Well under +#: the point where the mount starts refusing opens, and far above anything a test +#: fixture reaches. +RELIEF_THRESHOLD = 50_000 +#: Reclaim this often *within* a phase as well. Chunking alone is not enough: one +#: band-year can hold 230k frames, and the mount began refusing opens partway through +#: renaming it, costing 132,857 renames in an earlier attempt. +RELIEF_INTERVAL = vfs.RELIEF_INTERVAL + + +# Reclaim lives in suvi.vfs: it is a property of this mount, not of this migration, +# and the index build needs it just as much. Re-exported so callers and tests can +# reach it by either name. +drop_caches = vfs.drop_caches +reclaimable_kb = vfs.reclaimable_kb +release_handles = vfs.release_handles + + +def _mark_chunk_done(args, wavelengths): + """Record that a (band, year) needs no further migration. + + Lets a resume skip finished subtrees without touching the filesystem at all -- + the traversal that a resume used to perform is exactly what exhausts this mount. + """ + if not args.apply or len(wavelengths) != 1: + return + year = args.year[0] if args.year and len(args.year) == 1 else None + try: + conn = db.connect(args.db) + db.set_meta(conn, chunk_key(wavelengths[0], year), "done") + conn.commit() + conn.close() + except Exception as exc: + print(f" (could not record chunk progress: {exc})") + + +def export_labels(path, records): + """Write the legacy verdicts to a standalone file before anything is renamed.""" + os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True) + with gzip.open(path, "wt", newline="") as handle: + writer = csv.writer(handle) + writer.writerow(["relpath", "legacy_label", "satellite", "wavelength", "t_start"]) + for relpath, label, name in records: + writer.writerow([relpath, label, name.satellite, name.wavelength, name.t_start]) + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--db", default=None, help="target index (default: $SUVI_DB)") + parser.add_argument("--root", default=None, help="archive root (default: $SUVI_DATA_ROOT)") + parser.add_argument("--apply", action="store_true", help="rename files (default: dry run)") + parser.add_argument("--satellites", default="16,18,19") + parser.add_argument("--wavelength", default=None, help="restrict to one band") + parser.add_argument("--year", type=int, action="append", default=None, + help="restrict to given years (repeatable)") + parser.add_argument("--export", default=None, + help="legacy label export (default: /legacy_labels.csv.gz)") + parser.add_argument("--keep-plots", action="store_true", + help="keep the _e.jpg diagnostic images") + parser.add_argument("--no-relief", action="store_true", + help="skip the cache-reclaim step between chunks") + parser.add_argument("--recheck", action="store_true", + help="re-examine chunks already recorded as migrated") + args = parser.parse_args(argv) + + root = args.root or paths.data_root() + if not os.path.isdir(root): + print(f"No such archive root: {root}") + return 1 + satellites = tuple(int(part) for part in args.satellites.split(",")) + wavelengths = (int(args.wavelength),) if args.wavelength else paths.WAVELENGTHS + + # One (band, year) at a time. Walking the whole archive in a single pass fills + # the dentry cache with millions of entries, and on this virtiofs mount each + # cached entry pins a host file handle -- the walk succeeds and then nothing can + # open a file again, which is how an earlier attempt lost 152,193 renames to + # ENFILE. A band-year is ~130k files, comfortably inside what the mount + # sustains, and each chunk finishes completely (export, record, rename) so a + # failure never leaves labels destroyed but unrecorded. + chunks = [] + for wavelength in wavelengths: + years = archive_years(root, satellites, wavelength) if args.year is None else args.year + for year in years or [None]: + chunks.append((wavelength, year)) + + # Skip chunks already recorded as complete, without touching the filesystem. + # Resuming used to re-walk finished subtrees just to discover there was nothing + # to do; across several passes that traversed the archive roughly four times + # over, and traversal is the very thing that exhausts this mount. + done = set() + # Only when applying: a dry run must neither create the index nor hide work, so + # it reports the archive's true state regardless of what progress was recorded. + if args.apply and not args.recheck: + try: + conn = db.connect(args.db) + done = { + key for (wavelength, year) in chunks + if db.get_meta(conn, key := chunk_key(wavelength, year)) == "done" + } + conn.close() + except Exception as exc: + print(f"Could not read migration progress ({exc}); checking every chunk.") + + if len(chunks) > 1: + overall = 0 + for index, (wavelength, year) in enumerate(chunks, 1): + label = f"{wavelength}A {year or 'all years'}" + if chunk_key(wavelength, year) in done: + print(f"\n=== chunk {index}/{len(chunks)}: {label} -- already migrated, skipping") + continue + print(f"\n{'=' * 70}\n=== chunk {index}/{len(chunks)}: {label} ===") + chunk_args = argparse.Namespace(**vars(args)) + chunk_args.wavelength = str(wavelength) + chunk_args.year = [year] if year is not None else None + if args.export is None: + suffix = f"ci{wavelength:03d}" + (f"_{year}" if year else "") + chunk_args.export = os.path.join(root, f"legacy_labels_{suffix}.csv.gz") + status, frames = main_one_chunk(chunk_args, root, satellites, (wavelength,)) + overall = overall or status + if not args.no_relief and frames >= RELIEF_THRESHOLD: + release_handles() + return overall + return main_one_chunk(args, root, satellites, wavelengths)[0] + + +def main_one_chunk(args, root, satellites, wavelengths): + print(f"Scanning {root}") + print(f" satellites={satellites} wavelengths={wavelengths} years={args.year or 'all'}") + + counts = Counter() + labelled = [] # (relpath, label, name) for every file carrying a verdict + renames = [] # (src, dst) pairs to perform + collisions = [] + plots = [] + started = time.time() + + for directory, filename, name in walk_archive(root, satellites, wavelengths, args.year): + counts["frames"] += 1 + relpath = name.relpath() + if name.label: + counts[f"label_{name.label}"] += 1 + labelled.append((relpath, name.label, name)) + source = os.path.join(directory, filename) + target = os.path.join(directory, name.filename()) + if os.path.exists(target): + collisions.append((source, target)) + else: + renames.append((source, target)) + if name.label == "e" and not args.keep_plots: + plot = os.path.join(directory, name.error_plot_name()) + if os.path.exists(plot): + plots.append(plot) + else: + counts["unlabelled"] += 1 + if counts["frames"] % RELIEF_INTERVAL == 0 and not args.no_relief: + release_handles() + if counts["frames"] % 200000 == 0: + print(f" {counts['frames']} frames in {time.time() - started:.0f}s") + + print(f"\nFound {counts['frames']} frames in {time.time() - started:.0f}s") + print(f" passed (_f): {counts['label_f']}") + print(f" rejected(_e): {counts['label_e']}") + print(f" unlabelled : {counts['unlabelled']} (never processed by any filter)") + print(f" renames to perform: {len(renames)}") + print(f" diagnostic plots to delete: {len(plots)}") + if collisions: + print(f" COLLISIONS (will be skipped): {len(collisions)}") + for source, target in collisions[:5]: + print(f" {os.path.basename(source)} -> {os.path.basename(target)} exists") + + if not args.apply: + print("\nDry run. Re-run with --apply to write the export, rename, and clean up.") + return 0, counts["frames"] + + if not labelled: + print("\nNothing labelled; archive is already migrated.") + _mark_chunk_done(args, wavelengths) + return 0, counts["frames"] + + # 1. Export the labels before touching a single filename. + export_path = args.export or os.path.join(root, "legacy_labels.csv.gz") + print(f"\nExporting {len(labelled)} legacy labels to {export_path}") + export_labels(export_path, labelled) + if not os.path.exists(export_path) or os.path.getsize(export_path) == 0: + print("ERROR: export is missing or empty; refusing to rename anything.") + return 1, counts["frames"] + print(f" {os.path.getsize(export_path) / 1024 / 1024:.1f} MB written") + + # 2. Record them in the index. + target_db = args.db or paths.default_db_path() + conn = db.connect(target_db) + run_id = db.create_detector_run( + conn, + LEGACY_RUN_NAME, + {"source": "filename suffixes", "root": root}, + notes="Verdicts recovered from _f/_e filename suffixes. Written ~May 2024 by a " + "filter version since retuned; preserved as history, not ground truth.", + ) + print(f"Recording verdicts in {target_db} as run {run_id} ({LEGACY_RUN_NAME})") + + batch = [] + for index, (relpath, label, name) in enumerate(labelled): + frame_id = db.upsert_frame(conn, name, relpath) + batch.append((frame_id, "good" if label == "f" else "bad", f"legacy_{label}", None, None)) + if len(batch) >= 5000: + db.record_detections(conn, run_id, batch) + conn.commit() + batch.clear() + print(f" {index + 1}/{len(labelled)}") + if batch: + db.record_detections(conn, run_id, batch) + conn.commit() + + recorded = conn.execute( + "SELECT count(*) c FROM detection WHERE run_id = ?", (run_id,) + ).fetchone()["c"] + if recorded != len(labelled): + print(f"ERROR: recorded {recorded} verdicts but found {len(labelled)} labels.") + print(" Not renaming anything.") + conn.close() + return 1, counts["frames"] + print(f" {recorded} verdicts recorded") + + # 3. Only now rename. + print(f"\nRenaming {len(renames)} files...") + renamed = failed = 0 + for index, (source, target) in enumerate(renames): + try: + os.rename(source, target) + renamed += 1 + except OSError as exc: + failed += 1 + if failed <= 5: + print(f" FAILED {source}: {exc}") + if index and index % RELIEF_INTERVAL == 0 and not args.no_relief: + release_handles() + if index and index % 100000 == 0: + print(f" {index}/{len(renames)}") + print(f" renamed {renamed}, failed {failed}, skipped {len(collisions)} collisions") + + removed = 0 + for plot in plots: + try: + os.remove(plot) + removed += 1 + except OSError: + pass + print(f" deleted {removed} diagnostic plots") + + # 4. Verify every recorded path now resolves. + print("\nVerifying...") + missing = 0 + for relpath, _, _ in labelled[:: max(1, len(labelled) // 5000)]: + if not os.path.exists(paths.abspath(relpath, root)): + missing += 1 + if missing: + print(f" WARNING: {missing} sampled paths do not resolve on disk") + else: + print(" all sampled paths resolve") + conn.close() + # Only claim the chunk is finished if every rename actually landed. Marking it + # done after partial failures is how a resume would skip real remaining work. + if failed == 0: + _mark_chunk_done(args, wavelengths) + else: + print(f" NOT marking this chunk complete: {failed} renames failed; re-run to retry.") + print("\nDone. The archive now uses NOAA's original filenames; verdicts live in " + f"the index, and a recovery copy is at {export_path}") + return 0, counts["frames"] + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/migrate_urlcache.py b/migrate_urlcache.py new file mode 100644 index 0000000..d889e90 --- /dev/null +++ b/migrate_urlcache.py @@ -0,0 +1,125 @@ +#!/usr/bin/env python +"""One-time migration: file_database.json -> the SQLite index's remote_file table. + +``puller_fits.py`` recorded every URL it had fetched in an 800 MB JSON object, parsed +into memory on every run. This moves that state into the index, where it is an +indexed lookup instead. + +This state is the only record of what has already been downloaded. Losing it means +re-fetching the entire archive from NOAA, so the migration refuses to proceed on any +inconsistency and never deletes the JSON -- it renames it to .bak only after the row +count matches, and only when asked. + +Dry run by default:: + + migrate_urlcache.py # report what would happen + migrate_urlcache.py --apply # write the rows + migrate_urlcache.py --apply --backup-json +""" + +import argparse +import json +import os +import sys +import time + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +from suvi import db, paths + + +def default_json_path(): + return os.path.join(os.path.dirname(paths.data_root()), "file_database.json") + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--json", default=None, help="source file_database.json") + parser.add_argument("--db", default=None, help="target index (default: $SUVI_DB)") + parser.add_argument("--apply", action="store_true", help="write rows (default: dry run)") + parser.add_argument("--backup-json", action="store_true", + help="rename the JSON to .bak after a verified migration") + parser.add_argument("--batch", type=int, default=50000) + args = parser.parse_args(argv) + + source = args.json or default_json_path() + if not os.path.exists(source): + print(f"No such file: {source}") + return 1 + + size = os.path.getsize(source) + print(f"Reading {source} ({size / 1024 / 1024:.0f} MB)...") + started = time.time() + with open(source, "r") as handle: + cache = json.load(handle) + if not isinstance(cache, dict): + print(f"ERROR: expected a JSON object, got {type(cache).__name__}") + return 1 + print(f" {len(cache)} URL records in {time.time() - started:.1f}s") + + malformed = [ + url for url, mtime in cache.items() + if not isinstance(url, str) or not isinstance(mtime, (int, float)) + ] + if malformed: + print(f"ERROR: {len(malformed)} records are malformed, e.g. {malformed[:3]}") + return 1 + + if not args.apply: + print(f"\nDry run: would insert {len(cache)} rows into remote_file.") + print("Re-run with --apply to write them.") + return 0 + + target = args.db or paths.default_db_path() + conn = db.connect(target) + existing = conn.execute("SELECT count(*) c FROM remote_file").fetchone()["c"] + print(f"Writing to {target} (currently {existing} rows)...") + + now = time.time() + written = 0 + batch = [] + for url, mtime in cache.items(): + batch.append((url, float(mtime), None, None, now)) + if len(batch) >= args.batch: + db.record_remote_files(conn, batch) + conn.commit() + written += len(batch) + batch.clear() + print(f" {written}/{len(cache)}") + if batch: + db.record_remote_files(conn, batch) + conn.commit() + written += len(batch) + + final = conn.execute("SELECT count(*) c FROM remote_file").fetchone()["c"] + print(f" wrote {written} rows; table now holds {final}") + + if final < len(cache): + print(f"ERROR: table holds {final} rows but the JSON had {len(cache)} URLs.") + print(" Not touching the JSON. Investigate before re-running.") + conn.close() + return 1 + + # Spot-check that values survived the round trip, not just the row count. + sample = list(cache.items())[:: max(1, len(cache) // 20)][:20] + for url, mtime in sample: + stored = db.get_remote_mtime(conn, url) + if stored is None or abs(stored - float(mtime)) > 1e-6: + print(f"ERROR: round-trip mismatch for {url}: {mtime} -> {stored}") + conn.close() + return 1 + print(f" spot-checked {len(sample)} URLs, all match") + conn.close() + + if args.backup_json: + backup = source + ".bak" + os.rename(source, backup) + print(f"Renamed {source} -> {backup}") + else: + print(f"Left {source} in place; pass --backup-json once puller_fits.py is verified.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/puller.py b/puller.py deleted file mode 100644 index e7fe72b..0000000 --- a/puller.py +++ /dev/null @@ -1,148 +0,0 @@ -import os -import urllib.request -import urllib.parse -import re -import time -import random -import datetime -import json -from threading import Thread -import queue - -import tqdm - -def recursive_find_links(url): - links_regex_pattern = r'(?<=)' # Find href links that do not contain question marks or whitespace - times_regex_pattern = r'(?<=<\/a>)(\s*\d{4}-\d{2}-\d{2} \d{2}:\d{2})(?= )' # Find timestamps in UTC in the format YYYY-MM-DD HH:mm - with urllib.request.urlopen(url) as response: - html_content = response.read().decode('utf-8') - - links = re.findall(links_regex_pattern, html_content) - times = re.findall(times_regex_pattern, html_content) - for i in range(len(times)): - dt = datetime.datetime.strptime(times[i].strip(), "%Y-%m-%d %H:%M") - dt.replace(tzinfo=datetime.timezone.utc) - times[i] = time.mktime(dt.timetuple()) - - if len(links) - 1 == len(times): - links = links[1:] - elif len(links) == len(times): - pass - else: - raise(ValueError) - - for _link, _time in zip(links, times): - if _link.endswith("latest.png"): - continue - elif _link.endswith("/"): - yield from recursive_find_links(urllib.parse.urljoin(url,_link)) - else: - yield urllib.parse.urljoin(url,_link), _time - - -# Fetch image from url and store it to path, retrying on failure -def image_fetch_worker(work_queue, result_queue, attempt_count = 1): - while True: - job = work_queue.get() - if job == None: - return - url, path, t = job - - attempts = 0 - while attempts < attempt_count: - try: - req = urllib.request.Request(url, data=None) - image_data = urllib.request.urlopen(req).read() - os.makedirs(os.path.split(path)[0], exist_ok=True) - open(path, 'wb').write(image_data) - result_queue.put((True, url, t)) - break - except Exception as e: - if hasattr(e, "code") and e.code == 404: # This is expected if the file has been removed from the site (at least for swpc.noaa.gov) - attempts += attempt_count - elif attempts == 0: - print(f"\nA problem occurred on image: {url} | {e}") - time.sleep(1 + random.random()) - attempts += 1 - if (attempt_count > 1) and (attempts == attempt_count): - print(f'\nAfter {attempt_count} retries, could not fetch: {url}') - result_queue.put((False, url, t)) - break - - -if __name__ == "__main__": - directory_url = r"https://services.swpc.noaa.gov/images/animations/suvi/" - stored_images_dir = r"..\suvi" - file_database_path = r"..\file_database.json" - fetch_interval = 30*60 - nworkers = 8 - - file_info_cache = {} - try: - print(f"Attempting to load file records from cache: {file_database_path}") - with open(file_database_path, 'r') as f: - file_info_cache = json.loads(f.read()) - print(f"File records loaded from cache: {len(file_info_cache)} records found.") - except Exception as e: - print(f"Load failed, starting with empty cache") - file_info_cache = {} - - work_queue = queue.Queue(maxsize=nworkers) - result_queue = queue.Queue() - workers = [] - for _ in range(nworkers): - t = Thread(target=image_fetch_worker, args=(work_queue, result_queue), daemon=True) - t.start() - workers.append(t) - try: - while True: - fetched_image_count = 0 - already_had_image_count = 0 - failed_image_count = 0 - urllen = len(directory_url) - for l, t in tqdm.tqdm(recursive_find_links(directory_url), desc="Downloading files"): - # Collect complete work and record it - while True: - try: - r_success, r_url, r_t = result_queue.get_nowait() - if r_success: - fetched_image_count += 1 - file_info_cache[r_url] = r_t - else: - failed_image_count += 1 - except queue.Empty: - break - # If we dont have the file or the file at the link is newer than the one we previously fetched - if (not (l in file_info_cache)) or t > file_info_cache[l]: - file_portion_of_link = l[urllen:] - filepath = os.path.join(stored_images_dir, file_portion_of_link.replace("/", os.sep)) - if os.path.exists(filepath): - if l in file_info_cache: # If we have record of this file, it must be out of date, rename it and download the new version. - os.rename(filepath, filepath + f"_{int(file_info_cache[l])}") - else: # If we have no record of this file, update the file info cache and don't redownload - file_info_cache[l] = t - already_had_image_count += 1 - continue - work_queue.put((l, filepath, t)) - else: - already_had_image_count += 1 - with open(file_database_path, 'w') as f: - f.write(json.dumps(file_info_cache)) - print(f"Downloaded {fetched_image_count} | Already had {already_had_image_count} | Failed {failed_image_count}") - time.sleep(fetch_interval) - except KeyboardInterrupt: - print("Saving file database and shutting down.") - - for _ in range(nworkers): - try: - work_queue.put(None, timeout=1.0) - except: - break - - for w in workers: - w.join(5.0) - - with open(file_database_path, 'w') as f: - f.write(json.dumps(file_info_cache)) - - print(f"Downloaded {fetched_image_count} | Already had {already_had_image_count} | Failed {failed_image_count}") \ No newline at end of file diff --git a/puller_fits.py b/puller_fits.py index 0687a88..35db744 100644 --- a/puller_fits.py +++ b/puller_fits.py @@ -1,273 +1,307 @@ -import os -import urllib.request -import urllib.parse -import re -import time -import random -import datetime -import json -from threading import Thread -import queue -import math -from functools import partial -from queue import Empty - -import tqdm - -directory_url = r"https://data.ngdc.noaa.gov/platforms/solar-space-observing-satellites/goes/" -stored_images_dir = os.path.abspath(os.path.join("..", "Data")) -# directory_url = r"https://data.ngdc.noaa.gov/platforms/solar-space-observing-satellites/goes/goes16/l2/data/ephe-l2-orb1m/2017/12/" -# stored_images_dir = r"Z:\NOAA GOES Data\Data\goes16/l2/data/ephe-l2-orb1m/2017/12/" -ignore_folder_names = ["Parent Directory", "l1b", "goes17", "2017", "2018", "2019", "2020", "2021", "2022", "2023"] -file_database_path = os.path.abspath(os.path.join("..", "file_database.json")) -fetch_interval = 0 # 60*60 -nfetchworkers = 2 # Be nice to the servers, this value is how many threads will be asking for links and file info at the same time -ndownloadworkers = 3 # Be nice to the servers, this value is how many threads will be downloading files at the same time -randomize_order = False - -links_regex_pattern = r'(?<=.*\d{4}-\d{2}-\d{2} \d{2}:\d{2})' # Find href links that do not contain question marks or whitespace -times_regex_pattern = r'(?<=)(\d{4}-\d{2}-\d{2} \d{2}:\d{2})(?= )' # Find timestamps in UTC in the format YYYY-MM-DD HH:mm -sizes_regex_pattern = r'(?<=\d{4}-\d{2}-\d{2} \d{2}:\d{2} <\/td>)([ .\d]{1,3})([KMG]|- |0 | )(?=<\/td>)' # Find file size markers in one of many possible formats -links_matcher = re.compile(links_regex_pattern) -times_matcher = re.compile(times_regex_pattern) -sizes_matcher = re.compile(sizes_regex_pattern) - -def find_links_worker(query_work_queue, query_result_queue, attempt_count = 3, randomize_order = True): - while True: - job = query_work_queue.get() - if job == None: - return - url= job - - attempts = 0 - html_content = None - while attempts < attempt_count: - try: - with urllib.request.urlopen(url) as response: - html_content = response.read().decode('utf-8') - break - except Exception as e: - tqdm.tqdm.write(f"Exception while fetching links: {e}") - time.sleep(1 + random.random()) - attempts += 1 - if (attempt_count > 1) and (attempts == attempt_count): - tqdm.tqdm.write(f'\nAfter {attempt_count} retries, could not fetch: {url}') - - if html_content is None: # We failed to fetch the link for some reason, continue to the next job - continue - else: - links = links_matcher.findall(html_content) - times = times_matcher.findall(html_content) - sizes = sizes_matcher.findall(html_content) - for i in range(len(times)): - dt = datetime.datetime.strptime(times[i].strip(), "%Y-%m-%d %H:%M") - dt.replace(tzinfo=datetime.timezone.utc) - times[i] = time.mktime(dt.timetuple()) - for i in range(len(sizes)): - match sizes[i][1].strip(): - case '-': - sizes[i] = 0 - case 'K': - sizes[i] = int(float(sizes[i][0].strip()) * 1024) - case 'M': - sizes[i] = int(float(sizes[i][0].strip()) * 1024 * 1024) - case 'G': - sizes[i] = int(float(sizes[i][0].strip()) * 1024 * 1024 * 1024) - case "0": - sizes[i] = 0 - case "": - sizes[i] = int(float(sizes[i][0].strip())) - case _: - raise(ValueError(f"Unexpected symbol while parsing links page: {_}")) - - - if (len(links) != len(times)) or (len(times) != len(sizes)): - raise(ValueError("Links parsing error, mismatched numbers of links, times, or sizes!")) - - results = list(zip(links, times, sizes)) - if randomize_order: - random.shuffle(results) - for _link, _time, _size in results: - if _link.endswith("/"): - if _link.split(r"/")[-2] in ignore_folder_names: - continue - else: - # print(f"Queueing folder: {urllib.parse.urljoin(url,_link)}") - query_work_queue.put(urllib.parse.urljoin(url,_link)) - else: - # print(f"Queueing file: {urllib.parse.urljoin(url,_link)}") - query_result_queue.put((urllib.parse.urljoin(url,_link), _time, _size)) - - -# Fetch file from url and store it to path, retrying on failure -def file_download_worker(download_work_queue, download_result_queue, attempt_count = 2): - while True: - job = download_work_queue.get() - if job == None: - return - url, path, t, s = job - - # print(f"Next job: {url}, {path}, {t}, {s}") - attempts = 0 - while attempts < attempt_count: - try: - starttime = time.time() - req = urllib.request.Request(url, data=None) - image_data = urllib.request.urlopen(req, timeout=10.0).read() - endtime = time.time() - tqdm.tqdm.write(f"Downloaded {url} in {endtime-starttime:0.2f} s | {len(image_data)/1024/1024/(endtime-starttime):0.2f} MB/s") - if math.isclose(len(image_data), s, rel_tol=0.05): - os.makedirs(os.path.split(path)[0], exist_ok=True) - open(path, 'wb').write(image_data) - download_result_queue.put((True, url, t)) - break - else: - raise ValueError("Downloaded file is the wrong size!") - except Exception as e: - if hasattr(e, "code") and e.code == 404: # This is expected if the file has been removed from the site (at least for swpc.noaa.gov) - attempts += attempt_count - elif (0 < attempts) and (attempts < attempt_count): - # tqdm.tqdm.write(f"\nA problem occurred on file: {url} | {e}") - time.sleep(1 + random.random()) - attempts += 1 - if (attempt_count > 1) and (attempts == attempt_count): - tqdm.tqdm.write(f'\nAfter {attempt_count} retries, could not fetch: {url}') - tqdm.tqdm.write(f"Exception: {e}") - download_result_queue.put((False, url, t)) - break - - -if __name__ == "__main__": - file_info_cache = {} - try: - print(f"Attempting to load file records from cache: {file_database_path}") - with open(file_database_path, 'r') as f: - file_info_cache = json.loads(f.read()) - print(f"File records loaded from cache: {len(file_info_cache)} records found.") - except Exception as e: - print(f"Load failed, starting with empty cache") - file_info_cache = {} - - query_work_queue = queue.Queue() - query_result_queue = queue.Queue() - download_work_queue = queue.Queue(maxsize=ndownloadworkers) - download_result_queue = queue.Queue(maxsize=ndownloadworkers) - workers = [] - for _ in range(nfetchworkers): - t = Thread(target=find_links_worker, args=(query_work_queue, query_result_queue, 3, randomize_order), daemon=True) - t.start() - workers.append(t) - for _ in range(ndownloadworkers): - t = Thread(target=file_download_worker, args=(download_work_queue, download_result_queue), daemon=True) - t.start() - workers.append(t) - - try: - while True: - fetched_image_count = 0 - already_had_image_count = 0 - failed_image_count = 0 - urllen = len(directory_url) - - query_work_queue.put(directory_url) - for l, t, s in tqdm.tqdm(iter(partial(query_result_queue.get, timeout=30.0), None), desc="Downloading files"): - # Collect completed jobs and record completion status - while True: - try: - r_success, r_url, r_t = download_result_queue.get_nowait() - if r_success: - fetched_image_count += 1 - file_info_cache[r_url] = r_t - else: - failed_image_count += 1 - except queue.Empty: - break - - file_portion_of_link = l[urllen:] - filepath = os.path.join(stored_images_dir, file_portion_of_link.replace("/", os.sep)) - # If we dont have the file or the file at the link is newer than the one we previously fetched - if (not (l in file_info_cache)) or t > file_info_cache[l]: - if filepath.endswith(".fits"): - filepath2 = filepath.split(".fits")[0] + "_f.fits" # also look for the filtered version of the file - filepath3 = filepath.split(".fits")[0] + "_e.fits" # also look for the error version of the file - if os.path.exists(filepath2) or os.path.exists(filepath3): # If the unfiltered filename exists, that case will be handled in the alternative code path in if os.path.exists(filepath): - if l in file_info_cache: # If we have record of this file, it must be out of date, delete it and download the new version. - try: - os.remove(filepath2) - os.remove(filepath3) - except: - pass - else: # If we have no record of this file, update the file info cache and don't redownload - file_info_cache[l] = t - already_had_image_count += 1 - continue - if os.path.exists(filepath): - if l in file_info_cache: # If we have record of this file, it must be out of date, delete it and download the new version. - if os.path.exists(filepath): - os.remove(filepath) - else: # If we have no record of this file, update the file info cache and don't redownload if fsize is right - fsize = os.path.getsize(filepath) - if math.isclose(fsize, s, rel_tol=0.05): - file_info_cache[l] = t - already_had_image_count += 1 - continue - else: - tqdm.tqdm.write(f'Found a mismatched size on file: {filepath} Redownloading!') - download_work_queue.put((l, filepath, t, s)) - else: - # We have a download record, confirm the file actually exists on disk - if os.path.exists(filepath): - already_had_image_count += 1 - elif filepath.endswith(".fits"): - filepath2 = filepath.split(".fits")[0] + "_f.fits" # also look for the filtered version of the file - filepath3 = filepath.split(".fits")[0] + "_e.fits" # also look for the error version of the file - if os.path.exists(filepath2) or os.path.exists(filepath3): - already_had_image_count += 1 - else: # We could not find the file on disk, queue for redownload - download_work_queue.put((l, filepath, t)) - - with open(file_database_path, 'w') as f: - f.write(json.dumps(file_info_cache)) - print(f"Downloaded {fetched_image_count} | Already had {already_had_image_count} | Failed {failed_image_count}") - if fetch_interval > 0: - print(f"Run complete!, sleeping for {fetch_interval} seconds.") - with open(file_database_path, 'w') as f: - f.write(json.dumps(file_info_cache)) - time.sleep(fetch_interval) - else: - print("Run complete!, exiting...") - break - except KeyboardInterrupt: - print("Saving file database and shutting down.") - except Empty: - print("Work Complete, Shutting down...") - except Exception as e: - print(f"Unhandled Exception during run: {e}") - print("Shutting down") - - for _ in range(nfetchworkers): - try: - query_work_queue.put(None) - except: - break - - time.sleep(1) - print(f"Downloaded {fetched_image_count} | Already had {already_had_image_count} | Failed {failed_image_count}") - - while len(download_work_queue.queue) > 0: - print(f"Waiting for {len(download_work_queue.queue)} downloads in queue...") - time.sleep(1) - - for _ in range(ndownloadworkers): - try: - download_work_queue.put(None, timeout=5.0) - except: - break - - with open(file_database_path, 'w') as f: - f.write(json.dumps(file_info_cache)) - - print("Waiting for workers to shutdown...") - - for w in workers: - w.join(5.0) \ No newline at end of file +import os +import sys +import urllib.request +import urllib.parse +import re +import time +import random +import datetime +from threading import Thread +import queue +import math +from functools import partial +from queue import Empty + +import tqdm + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +from suvi import db, index, paths, vfs + +directory_url = r"https://data.ngdc.noaa.gov/platforms/solar-space-observing-satellites/goes/" +stored_images_dir = paths.data_root() +ignore_folder_names = ["Parent Directory", "l1b", "goes17", "2017", "2018", "2019", "2020", "2021", "2022", "2023"] +fetch_interval = 0 # 60*60 +nfetchworkers = 2 # Be nice to the servers, this value is how many threads will be asking for links and file info at the same time +ndownloadworkers = 3 # Be nice to the servers, this value is how many threads will be downloading files at the same time +randomize_order = False + +links_regex_pattern = r'(?<=.*\d{4}-\d{2}-\d{2} \d{2}:\d{2})' # Find href links that do not contain question marks or whitespace +times_regex_pattern = r'(?<=)(\d{4}-\d{2}-\d{2} \d{2}:\d{2})(?= )' # Find timestamps in UTC in the format YYYY-MM-DD HH:mm +sizes_regex_pattern = r'(?<=\d{4}-\d{2}-\d{2} \d{2}:\d{2} <\/td>)([ .\d]{1,3})([KMG]|- |0 | )(?=<\/td>)' # Find file size markers in one of many possible formats +links_matcher = re.compile(links_regex_pattern) +times_matcher = re.compile(times_regex_pattern) +sizes_matcher = re.compile(sizes_regex_pattern) + +def find_links_worker(query_work_queue, query_result_queue, attempt_count = 3, randomize_order = True): + while True: + job = query_work_queue.get() + if job == None: + return + url= job + + attempts = 0 + html_content = None + while attempts < attempt_count: + try: + with urllib.request.urlopen(url) as response: + html_content = response.read().decode('utf-8') + break + except Exception as e: + tqdm.tqdm.write(f"Exception while fetching links: {e}") + time.sleep(1 + random.random()) + attempts += 1 + if (attempt_count > 1) and (attempts == attempt_count): + tqdm.tqdm.write(f'\nAfter {attempt_count} retries, could not fetch: {url}') + + if html_content is None: # We failed to fetch the link for some reason, continue to the next job + continue + else: + links = links_matcher.findall(html_content) + times = times_matcher.findall(html_content) + sizes = sizes_matcher.findall(html_content) + for i in range(len(times)): + dt = datetime.datetime.strptime(times[i].strip(), "%Y-%m-%d %H:%M") + dt.replace(tzinfo=datetime.timezone.utc) + times[i] = time.mktime(dt.timetuple()) + for i in range(len(sizes)): + match sizes[i][1].strip(): + case '-': + sizes[i] = 0 + case 'K': + sizes[i] = int(float(sizes[i][0].strip()) * 1024) + case 'M': + sizes[i] = int(float(sizes[i][0].strip()) * 1024 * 1024) + case 'G': + sizes[i] = int(float(sizes[i][0].strip()) * 1024 * 1024 * 1024) + case "0": + sizes[i] = 0 + case "": + sizes[i] = int(float(sizes[i][0].strip())) + case _: + raise(ValueError(f"Unexpected symbol while parsing links page: {_}")) + + + if (len(links) != len(times)) or (len(times) != len(sizes)): + raise(ValueError("Links parsing error, mismatched numbers of links, times, or sizes!")) + + results = list(zip(links, times, sizes)) + if randomize_order: + random.shuffle(results) + for _link, _time, _size in results: + if _link.endswith("/"): + if _link.split(r"/")[-2] in ignore_folder_names: + continue + else: + # print(f"Queueing folder: {urllib.parse.urljoin(url,_link)}") + query_work_queue.put(urllib.parse.urljoin(url,_link)) + else: + # print(f"Queueing file: {urllib.parse.urljoin(url,_link)}") + query_result_queue.put((urllib.parse.urljoin(url,_link), _time, _size)) + + +# Fetch file from url and store it to path, retrying on failure +def file_download_worker(download_work_queue, download_result_queue, attempt_count = 2): + while True: + job = download_work_queue.get() + if job == None: + return + url, path, t, s = job + + # print(f"Next job: {url}, {path}, {t}, {s}") + attempts = 0 + while attempts < attempt_count: + try: + starttime = time.time() + req = urllib.request.Request(url, data=None) + image_data = urllib.request.urlopen(req, timeout=10.0).read() + endtime = time.time() + tqdm.tqdm.write(f"Downloaded {url} in {endtime-starttime:0.2f} s | {len(image_data)/1024/1024/(endtime-starttime):0.2f} MB/s") + if math.isclose(len(image_data), s, rel_tol=0.05): + os.makedirs(os.path.split(path)[0], exist_ok=True) + open(path, 'wb').write(image_data) + download_result_queue.put((True, url, t)) + break + else: + raise ValueError("Downloaded file is the wrong size!") + except Exception as e: + if hasattr(e, "code") and e.code == 404: # This is expected if the file has been removed from the site (at least for swpc.noaa.gov) + attempts += attempt_count + elif (0 < attempts) and (attempts < attempt_count): + # tqdm.tqdm.write(f"\nA problem occurred on file: {url} | {e}") + time.sleep(1 + random.random()) + attempts += 1 + if (attempt_count > 1) and (attempts == attempt_count): + tqdm.tqdm.write(f'\nAfter {attempt_count} retries, could not fetch: {url}') + tqdm.tqdm.write(f"Exception: {e}") + download_result_queue.put((False, url, t)) + break + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser( + description="Mirror NOAA GOES SUVI L2 FITS into the local archive.", + ) + parser.add_argument("--db", default=None, help="SQLite index (default: $SUVI_DB)") + parser.add_argument("--root", default=None, help="archive root (default: $SUVI_DATA_ROOT)") + parser.add_argument("--interval", type=int, default=fetch_interval, + help="seconds to sleep between passes; 0 runs once and exits") + args = parser.parse_args() + + if args.root: + stored_images_dir = os.path.abspath(args.root) + fetch_interval = args.interval + + # Download bookkeeping lives in the index. It used to be an 800 MB JSON object + # parsed into memory on every run; see migrate_urlcache.py for the conversion. + # Losing this state means re-fetching the whole archive, so writes are committed + # as they happen rather than only at the end of a pass. + conn = db.connect(args.db) + known = conn.execute("SELECT count(*) c FROM remote_file").fetchone()["c"] + print(f"Download records in index: {known}") + + query_work_queue = queue.Queue() + query_result_queue = queue.Queue() + download_work_queue = queue.Queue(maxsize=ndownloadworkers) + download_result_queue = queue.Queue(maxsize=ndownloadworkers) + workers = [] + for _ in range(nfetchworkers): + t = Thread(target=find_links_worker, args=(query_work_queue, query_result_queue, 3, randomize_order), daemon=True) + t.start() + workers.append(t) + for _ in range(ndownloadworkers): + t = Thread(target=file_download_worker, args=(download_work_queue, download_result_queue), daemon=True) + t.start() + workers.append(t) + + fetched_image_count = already_had_image_count = failed_image_count = 0 + pending_records = [] + + def flush_records(force=False): + """Commit download records in batches, and always before a long pause.""" + if pending_records and (force or len(pending_records) >= 200): + db.record_remote_files(conn, pending_records) + conn.commit() + pending_records.clear() + + reliever = vfs.Reliever(label="puller") + + def index_frame(local_path): + reliever.tick() + """Record a freshly written frame in the index. + + The downloader already knows this file exists, so indexing it here means the + index never has to be rebuilt by traversing the archive -- which on this + virtiofs mount is an operation to be avoided rather than merely optimised + (see suvi/index.py). Failure to index is not worth losing a download over; + `filter_FITS.py index` will pick the frame up from its directory mtime. + """ + try: + index.record_downloaded(conn, local_path, stored_images_dir) + except Exception as exc: + tqdm.tqdm.write(f"Could not index {local_path}: {exc}") + + try: + while True: + fetched_image_count = 0 + already_had_image_count = 0 + failed_image_count = 0 + urllen = len(directory_url) + + query_work_queue.put(directory_url) + for l, t, s in tqdm.tqdm(iter(partial(query_result_queue.get, timeout=30.0), None), desc="Downloading files"): + # Collect completed jobs and record completion status + while True: + try: + r_success, r_url, r_t = download_result_queue.get_nowait() + if r_success: + fetched_image_count += 1 + pending_records.append((r_url, r_t, None, None, time.time())) + index_frame(os.path.join(stored_images_dir, + r_url[urllen:].replace("/", os.sep))) + else: + failed_image_count += 1 + except queue.Empty: + break + flush_records() + + file_portion_of_link = l[urllen:] + filepath = os.path.join(stored_images_dir, file_portion_of_link.replace("/", os.sep)) + recorded_mtime = db.get_remote_mtime(conn, l) + + # The archive no longer carries _f/_e suffixes, so a frame is either + # at its published filename or absent -- no variant probing needed. + if recorded_mtime is None or t > recorded_mtime: + if os.path.exists(filepath): + if recorded_mtime is not None: + # We have a record, so the remote copy is newer: replace it. + try: + os.remove(filepath) + except OSError: + pass + else: + # No record, but the file is here. Adopt it if the size + # matches rather than re-downloading the whole archive. + if math.isclose(os.path.getsize(filepath), s, rel_tol=0.05): + pending_records.append((l, t, s, filepath, time.time())) + index_frame(filepath) + already_had_image_count += 1 + continue + tqdm.tqdm.write(f"Size mismatch on {filepath}; redownloading") + download_work_queue.put((l, filepath, t, s)) + elif os.path.exists(filepath): + already_had_image_count += 1 + else: + # Recorded as downloaded but gone from disk; fetch it again. + download_work_queue.put((l, filepath, t, s)) + + flush_records(force=True) + print(f"Downloaded {fetched_image_count} | Already had {already_had_image_count} | Failed {failed_image_count}") + if fetch_interval > 0: + print(f"Run complete!, sleeping for {fetch_interval} seconds.") + time.sleep(fetch_interval) + else: + print("Run complete!, exiting...") + break + except KeyboardInterrupt: + print("Saving download records and shutting down.") + except Empty: + print("Work Complete, Shutting down...") + + for _ in range(nfetchworkers): + try: + query_work_queue.put(None) + except Exception: + break + + time.sleep(1) + print(f"Downloaded {fetched_image_count} | Already had {already_had_image_count} | Failed {failed_image_count}") + + while len(download_work_queue.queue) > 0: + print(f"Waiting for {len(download_work_queue.queue)} downloads in queue...") + time.sleep(1) + + for _ in range(ndownloadworkers): + try: + download_work_queue.put(None, timeout=5.0) + except Exception: + break + + # Drain any results the workers finished after the main loop ended. + while True: + try: + r_success, r_url, r_t = download_result_queue.get_nowait() + if r_success: + pending_records.append((r_url, r_t, None, None, time.time())) + index_frame(os.path.join(stored_images_dir, + r_url[len(directory_url):].replace("/", os.sep))) + except queue.Empty: + break + flush_records(force=True) + conn.close() + + print("Waiting for workers to shutdown...") + for w in workers: + w.join(5.0) diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..6801365 --- /dev/null +++ b/pytest.ini @@ -0,0 +1,5 @@ +[pytest] +testpaths = tests +filterwarnings = + # astropy warns about the deliberately damaged files these tests construct. + ignore::astropy.utils.exceptions.AstropyUserWarning diff --git a/reclaim.py b/reclaim.py new file mode 100644 index 0000000..434c306 --- /dev/null +++ b/reclaim.py @@ -0,0 +1,69 @@ +#!/usr/bin/env python +"""Hand the archive's file handles back to the host. + +The archive is on a virtiofs share whose daemon holds a host file descriptor per +inode the guest has looked up. Anything that reads a lot of the archive leaves those +inodes cached, and the handles with them, until something forces the guest to evict +them -- at which point the daemon can run out and the mount starts refusing *every* +open, which surfaces across the whole machine as "too many open files in system". + +The pipeline's own bulk jobs now reclaim as they go and again when they finish, so +this should rarely be needed. It exists for when something else has loaded the cache, +or to check where things stand: + + reclaim.py # report, then reclaim if needed + reclaim.py --check # report only, change nothing + sudo reclaim.py # uses drop_caches: instant, and far gentler + +Without root the only lever available is memory pressure, which means briefly +allocating several GiB. Running under sudo avoids that entirely. +""" + +import argparse +import os +import sys + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +from suvi import vfs + +#: Roughly the slab size at which the daemon is likely near a 1M descriptor limit. +#: A cached dentry plus inode is on the order of a kilobyte, so a gigabyte of +#: reclaimable slab is on the order of a million inodes. +CONCERN_GIB = 1.0 + + +def report(prefix): + slab = vfs.reclaimable_kb() / (1024 * 1024) + print(f"{prefix:>8}: {slab:.2f} GiB reclaimable slab " + f"(~{slab:.1f}M cached inodes), {vfs.available_gib():.1f} GiB available") + return slab + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--check", action="store_true", help="report only, reclaim nothing") + parser.add_argument("--force", action="store_true", help="reclaim even if it looks fine") + args = parser.parse_args(argv) + + before = report("before") + if args.check: + print(" (check only; nothing reclaimed)") + return 0 if before < CONCERN_GIB else 1 + + if before < CONCERN_GIB and not args.force: + print(f" below {CONCERN_GIB} GiB; nothing to do. Use --force to reclaim anyway.") + return 0 + + if os.geteuid() != 0: + print(" not root, so using memory pressure; run under sudo for the direct path") + freed = vfs.release_handles() + after = report("after") + print(f" {'reclaimed' if freed else 'no change'}: " + f"{before - after:+.2f} GiB") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/requirements.txt b/requirements.txt index 3fc1061..363eff2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,8 +1,12 @@ astropy -scikit-image -tqdm -numpy -opencv-python -palettable matplotlib -pillow \ No newline at end of file +numpy +opencv-python-headless +palettable +pillow +scikit-image +sortedcontainers +tqdm + +# Development +pytest diff --git a/suvi/__init__.py b/suvi/__init__.py new file mode 100644 index 0000000..867f51e --- /dev/null +++ b/suvi/__init__.py @@ -0,0 +1,6 @@ +"""Shared library for the NOAA GOES SUVI pipeline. + +Modules here are pure and side-effect free unless documented otherwise, so that the +production drivers (``filter_FITS.py``, ``merger_FITS.py``, ``puller_fits.py``) and +the test bench (``bench.py``) can share one implementation of every rule. +""" diff --git a/suvi/cases.py b/suvi/cases.py new file mode 100644 index 0000000..ed6419e --- /dev/null +++ b/suvi/cases.py @@ -0,0 +1,397 @@ +"""Building bench cases: finding clean windows, planning injections, resolving reads. + +The bench never modifies the archive. A *case* is a plan plus an overlay directory: +slots the plan corrupts are written as new files under the overlay, slots it deletes +resolve to nothing, and every other slot resolves straight through to the real file. +:class:`Overlay` is the indirection that makes that work, so detectors and fillers can +be written as if they were reading the archive directly. + +Everything here is deterministic in the case seed, so a run can be reproduced or +resumed exactly. +""" + +import os +from dataclasses import dataclass, field + +import numpy as np + +from . import corruptions, paths + +#: A slot with no good frame either side cannot be filled or scored, so cases leave +#: this many good slots untouched at each end of the window. +EDGE_MARGIN = 8 +#: Minimum good slots left between two injected gaps, so every gap has clean brackets. +MIN_SEPARATION = 4 + + +def gap_separation(length): + """Clean slots to leave either side of a gap of `length`. + + A fixed four slots was fine when gaps were one to six frames. It is not fine at + three hundred: the frames a long gap is reconstructed from would themselves sit + within a few slots of the next gap, so the "clean brackets" a fill is scored + against would be neighbours of other damage. Scale the guard with the gap. + """ + return max(MIN_SEPARATION, length // 4) + + +def scan_window(root, satellites, wavelengths, t_start, t_end): + """Index the archive for one time window, by filename only. + + Returns {(satellite, wavelength, t_start): (abs_path, legacy_label)}. Reads no + file contents, so it is cheap enough to run over a week of data interactively. + """ + found = {} + for satellite in satellites: + for wavelength in wavelengths: + base = os.path.join( + root, f"goes{satellite:d}", "l2", "data", f"suvi-l2-ci{wavelength:03d}" + ) + if not os.path.isdir(base): + continue + for day_dir in _candidate_days(base, t_start, t_end): + try: + entries = os.listdir(day_dir) + except OSError: + continue + for entry in entries: + name = paths.parse_frame_filename(entry) + if name is None: + continue + if not (t_start <= name.t_start < t_end): + continue + if name.satellite != satellite or name.wavelength != wavelength: + continue + found[name.slot] = (os.path.join(day_dir, entry), name.label) + return found + + +def _candidate_days(base, t_start, t_end): + """Day directories the window could touch, without walking the whole tree.""" + import datetime + + day = datetime.datetime.fromtimestamp(t_start, datetime.timezone.utc).date() + last = datetime.datetime.fromtimestamp(t_end, datetime.timezone.utc).date() + while day <= last: + path = os.path.join(base, f"{day.year:04d}", f"{day.month:02d}", f"{day.day:02d}") + if os.path.isdir(path): + yield path + day += datetime.timedelta(days=1) + + +def timeline(t_start, t_end, cadence=paths.CADENCE): + """Every slot the cadence says should exist in a window.""" + return list(range(t_start, t_end, cadence)) + + +def slot_is_good(found, satellite, wavelength, time): + """Whether a slot holds a frame the legacy filter passed. + + An unlabelled frame counts as unknown, not good: roughly a quarter of the + archive was never processed, and treating that as clean would quietly seed the + ground truth with unvetted frames. + + Note that these labels are *stale*. The archive's suffixes were written in + May 2024; filter_FITS.py was substantially retuned that July (max_center_skew + 5 -> 7, ratio_above_thresh_max 0.5 -> 0.4, a minimum added). They record what + some earlier filter thought, which is useful as an availability hint and as + history, but is not ground truth. Vetting uses headers instead. + """ + entry = found.get((satellite, wavelength, time)) + if entry is None: + return None + return entry[1] == "f" if entry[1] else None + + +def find_runs(found, satellites, wavelengths, t_start, t_end, minimum=120, + require_good_label=False): + """Longest stretches where every requested band and satellite has a frame. + + Returns (length_in_slots, start_time, end_time) tuples, longest first. Needs no + file contents, only filenames, so it is fast enough to run interactively over + weeks of archive. + + By default this asks only that a frame *exists*, because the legacy labels are + stale (see :func:`slot_is_good`) and filtering on them would both miss good data + and admit bad. Judging quality is ``vet-window``'s job, from the headers. + """ + slots = timeline(t_start, t_end) + runs = [] + current = 0 + start = None + for time in slots: + if require_good_label: + complete = all( + slot_is_good(found, satellite, wavelength, time) is True + for satellite in satellites + for wavelength in wavelengths + ) + else: + complete = all( + (satellite, wavelength, time) in found + for satellite in satellites + for wavelength in wavelengths + ) + if complete: + if current == 0: + start = time + current += 1 + else: + if current >= minimum: + runs.append((current, start, time - paths.CADENCE)) + current = 0 + if current >= minimum: + runs.append((current, start, slots[-1])) + return sorted(runs, reverse=True) + + +# --------------------------------------------------------------------------- plans + + +@dataclass +class InjectionPlan: + """What a case does to its window. + + The axes here are the experiment's independent variables: how much is damaged, + in what size runs, on which satellites, and how. + """ + + #: Fraction of slots to damage, in [0, 1). + fraction: float = 0.10 + #: Gap widths to sample, in slots. Fill quality degrades with gap length, so the + #: spread matters more than the mean. + gap_lengths: tuple = (1, 2, 3, 5, 10, 30, 60) + #: 'g16', 'g18', 'both', or 'mixed'. 'both' removes the same slots from each + #: satellite at once, which is the only case where cross-satellite fill is + #: unavailable -- and it is 31% of real outages, so it is not a corner case. + satellite_scope: str = "mixed" + #: 'all' damages every band in a slot; 'one' damages a single band, which is the + #: case a per-wavelength repair can actually exploit. + wavelength_scope: str = "all" + #: Corruption modes to sample from, plus the pseudo-mode 'delete'. + modes: tuple = ("delete",) + tuple(corruptions.CATALOG) + #: Severity range sampled per gap. + severity: tuple = (0.5, 1.0) + #: Gaps to place at *each* length. When set this governs instead of `fraction`. + #: + #: A budget expressed as a fraction is dominated by whichever gaps are longest: + #: one 300-slot gap consumes half a 0.25 budget on a 2,600-slot window, so the + #: run ends up with a single gap per length and a single corruption mode per + #: gap. 84% of one such case was one 300-slot `truncate`, which is not a test of + #: anything but that mode at that length. Asking for a count per length gives an + #: experiment with a designed shape instead of an emergent one. + gaps_per_length: int | None = None + + def validate(self): + if not 0.0 <= self.fraction < 1.0: + raise ValueError(f"fraction must be in [0, 1), got {self.fraction}") + if not self.gap_lengths or any(length < 1 for length in self.gap_lengths): + raise ValueError("gap_lengths must all be >= 1") + if self.satellite_scope not in ("g16", "g18", "both", "mixed"): + raise ValueError(f"unknown satellite_scope {self.satellite_scope!r}") + if self.wavelength_scope not in ("all", "one"): + raise ValueError(f"unknown wavelength_scope {self.wavelength_scope!r}") + unknown = [ + mode for mode in self.modes if mode != "delete" and mode not in corruptions.CATALOG + ] + if unknown: + raise ValueError(f"unknown corruption modes: {unknown}") + low, high = self.severity + if not 0.0 < low <= high <= 1.0: + raise ValueError(f"severity must satisfy 0 < low <= high <= 1, got {self.severity}") + return self + + def as_dict(self): + return { + "fraction": self.fraction, + "gap_lengths": list(self.gap_lengths), + "satellite_scope": self.satellite_scope, + "wavelength_scope": self.wavelength_scope, + "modes": list(self.modes), + "severity": list(self.severity), + "gaps_per_length": self.gaps_per_length, + } + + @classmethod + def from_dict(cls, payload): + return cls( + fraction=payload["fraction"], + gap_lengths=tuple(payload["gap_lengths"]), + satellite_scope=payload["satellite_scope"], + wavelength_scope=payload["wavelength_scope"], + modes=tuple(payload["modes"]), + severity=tuple(payload["severity"]), + gaps_per_length=payload.get("gaps_per_length"), + ) + + +@dataclass(frozen=True) +class Injection: + """One damaged slot.""" + + slot: tuple + mode: str + severity: float + #: Position within its gap, and the gap's total width -- fill quality is reported + #: against gap length, so this has to survive into the results. + gap_index: int + gap_length: int + #: Stable per-slot seed, so a corruption is reproducible in isolation. + seed: int + + +def plan_injections(plan, window_slots, satellites, wavelengths, seed): + """Choose which slots to damage and how. + + Gaps are laid down as non-overlapping runs separated by clean slots, with the + window's edges left intact, so every damaged slot has good frames to be + reconstructed from. Returns a list of :class:`Injection`. + """ + plan.validate() + rng = np.random.default_rng(seed) + # Gaps are laid down along the timeline, so collapse the (satellite, band, time) + # slots to the distinct timestamps they cover. + times = sorted({slot[2] for slot in window_slots}) + usable = times[EDGE_MARGIN : len(times) - EDGE_MARGIN] + if not usable: + return [] + + # Refuse a window that cannot hold the gaps asked for, rather than quietly + # planning fewer (or none) and reporting a result for gap lengths that were + # never actually tested. + longest = max(plan.gap_lengths) + required = longest + 2 * gap_separation(longest) + if len(usable) < required: + raise ValueError( + f"window has {len(usable)} usable slots but a gap of {longest} needs " + f"{required} with its guard bands; use a longer window or shorter gaps" + ) + + lengths = sorted(plan.gap_lengths, reverse=True) + if plan.gaps_per_length: + wanted = {int(length): plan.gaps_per_length for length in lengths} + target = sum(length * count for length, count in wanted.items()) + else: + wanted = None + target = int(len(times) * plan.fraction) + if target <= 0: + return [] + + # Lay down gaps at random starts, rejecting any that would touch an existing one. + # Longest first: a 300-slot gap placed last would rarely find room, so the mix + # would silently skew towards the short lengths it is meant to be compared with. + taken = set() + gaps = [] + budget = 0 + attempts = max(len(usable) * 4, 8000) + for attempt in range(attempts): + if wanted is None: + if budget >= target: + break + elif not any(wanted.values()): + break + length = int(lengths[attempt % len(lengths)]) + if wanted is not None and not wanted.get(length): + continue + separation = gap_separation(length) + start_index = int(rng.integers(0, max(1, len(usable) - length))) + span = usable[start_index : start_index + length] + if len(span) < length: + continue + guard = range(start_index - separation, start_index + length + separation) + if any(index in taken for index in guard): + continue + taken.update(guard) + gaps.append(span) + budget += length + if wanted is not None: + wanted[length] -= 1 + + injections = [] + mode_order = list(plan.modes) + rng.shuffle(mode_order) + for gap_number, gap in enumerate(gaps): + # Cycle rather than sample: with only a handful of gaps, independent draws + # repeat modes and leave most of the catalogue untested. + mode = str(mode_order[gap_number % len(mode_order)]) + low, high = plan.severity + severity = float(rng.uniform(low, high)) + targets = _gap_targets(plan, satellites, wavelengths, rng) + for index, time in enumerate(gap): + for satellite, wavelength in targets: + slot = (satellite, wavelength, time) + if slot not in window_slots: + continue + injections.append( + Injection( + slot=slot, + mode=mode, + severity=severity, + gap_index=index, + gap_length=len(gap), + seed=int(rng.integers(0, 2**31 - 1)), + ) + ) + return injections + + +def _gap_targets(plan, satellites, wavelengths, rng): + """Which (satellite, wavelength) pairs one gap applies to.""" + if plan.satellite_scope == "both": + chosen_satellites = list(satellites) + elif plan.satellite_scope == "mixed": + chosen_satellites = [satellites[int(rng.integers(0, len(satellites)))]] + else: + wanted = int(plan.satellite_scope[1:]) + chosen_satellites = [wanted] if wanted in satellites else list(satellites[:1]) + + if plan.wavelength_scope == "all": + chosen_bands = list(wavelengths) + else: + chosen_bands = [wavelengths[int(rng.integers(0, len(wavelengths)))]] + + return [(s, w) for s in chosen_satellites for w in chosen_bands] + + +# ------------------------------------------------------------------------- overlay + + +@dataclass +class Overlay: + """Resolves a slot to the file a case should read for it. + + Damaged slots point at the overlay copy, deleted slots resolve to None, and + everything else falls through to the archive untouched. Nothing here can write + to the archive, which is the property that makes the bench safe to run against + live data. + """ + + #: slot -> absolute path in the real archive. + archive: dict = field(default_factory=dict) + #: slot -> absolute path of a corrupted replacement. + overrides: dict = field(default_factory=dict) + #: slots the case removed entirely. + deleted: frozenset = frozenset() + + def path(self, slot): + """The file to read for `slot`, or None if the case removed it.""" + if slot in self.deleted: + return None + if slot in self.overrides: + return self.overrides[slot] + entry = self.archive.get(slot) + return entry[0] if isinstance(entry, tuple) else entry + + def truth_path(self, slot): + """The original file, regardless of what the case did -- for scoring.""" + entry = self.archive.get(slot) + return entry[0] if isinstance(entry, tuple) else entry + + def slots(self): + return sorted(self.archive) + + def series(self, satellite, wavelength): + """Slots for one (satellite, band) stream, in time order.""" + return sorted( + slot for slot in self.archive if slot[0] == satellite and slot[1] == wavelength + ) diff --git a/suvi/corruptions.py b/suvi/corruptions.py new file mode 100644 index 0000000..4f4f132 --- /dev/null +++ b/suvi/corruptions.py @@ -0,0 +1,280 @@ +"""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.""" + shift = int(200 * severity) + 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) + h, w = image.shape + 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 _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), {} + + +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), + ) +} + +GROUPS = sorted({corruption.group for corruption in CATALOG.values()}) + + +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()), + } diff --git a/suvi/db.py b/suvi/db.py new file mode 100644 index 0000000..82e8782 --- /dev/null +++ b/suvi/db.py @@ -0,0 +1,615 @@ +"""SQLite index of the SUVI archive: frames, header metadata, and detection results. + +This replaces the ``_f``/``_e`` filename suffixes the pipeline used to carry its +verdicts in. Keeping labels out of filenames means a detector can be re-run or +re-tuned without touching the archive, and several detectors can disagree about the +same frame without anyone having to pick a winner on disk. + +Everything except ``detector_run``/``detection`` is reconstructible: ``frame`` from a +directory walk, ``header`` from a cheap re-read. The database is therefore safe to +delete, and the bench can point at a scratch copy via ``SUVI_DB``. + +Detection rows store the *continuous scores* a detector produced, not just its +verdict, so sweeping a threshold is a query rather than another pass over 1.6M files. + +Writes are expected to come from a single process. Worker processes read, and push +results down a queue to one writer (see ``writer_loop``); SQLite tolerates many +readers but only one writer, and this archive lives on a shared mount where lock +contention is expensive. +""" + +import json +import os +import queue +import sqlite3 +import time + +from . import paths + +SCHEMA_VERSION = 2 + +#: Versions that upgrade to the current one by nothing more than running the schema +#: script again. Every change so far has been a new table, and every table is +#: declared IF NOT EXISTS, so an older index gains it on the next open. A change +#: that alters or drops an existing column must not be added here -- it needs a real +#: migration step instead. +ADDITIVE_UPGRADES_FROM = (1,) + +#: Header keywords cached in the ``header`` table. Each entry maps a column name to +#: the FITS keyword it comes from and its SQLite type. Defined once so inserts and +#: reads cannot drift apart. +HEADER_COLUMNS = ( + ("empty", "EMPTY", "INTEGER"), + ("degraded", "DEGRADED", "INTEGER"), + ("eclipse", "ECLIPSE", "INTEGER"), + ("num_imgs", "NUM_IMGS", "INTEGER"), + ("n_long", "N_LONG", "INTEGER"), + ("n_short", "N_SHORT", "INTEGER"), + ("n_sh_fl", "N_SH_FL", "INTEGER"), + ("exptime", "EXPTIME", "REAL"), + ("img_min", "IMG_MIN", "REAL"), + ("img_max", "IMG_MAX", "REAL"), + ("img_mean", "IMG_MEAN", "REAL"), + ("img_sdev", "IMG_SDEV", "REAL"), + ("imgtii", "IMGTII", "REAL"), + ("imgtir", "IMGTIR", "REAL"), + ("diam_sun", "DIAM_SUN", "REAL"), + ("crpix1", "CRPIX1", "REAL"), + ("crpix2", "CRPIX2", "REAL"), + ("crota", "CROTA", "REAL"), + ("cdelt1", "CDELT1", "REAL"), + ("cdelt2", "CDELT2", "REAL"), + ("yaw_flip", "YAW_FLIP", "INTEGER"), + ("solar_b0", "SOLAR_B0", "REAL"), + ("dsun_obs", "DSUN_OBS", "REAL"), + ("wavelnth", "WAVELNTH", "INTEGER"), + ("date_beg", "DATE-BEG", "TEXT"), + ("date_obs", "DATE-OBS", "TEXT"), + ("date_end", "DATE-END", "TEXT"), + ("datasum", "DATASUM", "TEXT"), + ("checksum", "CHECKSUM", "TEXT"), +) + +HEADER_FIELDS = tuple(name for name, _, _ in HEADER_COLUMNS) + +_HEADER_DDL = ",\n ".join(f"{name} {sqltype}" for name, _, sqltype in HEADER_COLUMNS) + +SCHEMA = f""" +CREATE TABLE IF NOT EXISTS meta ( + key TEXT PRIMARY KEY, + value TEXT NOT NULL +); + +-- One row per FITS file present in the archive. `path` is archive-relative and +-- always uses forward slashes, so the index is portable between Windows and Linux. +CREATE TABLE IF NOT EXISTS frame ( + id INTEGER PRIMARY KEY, + path TEXT NOT NULL UNIQUE, + satellite INTEGER NOT NULL, + wavelength INTEGER NOT NULL, + t_start INTEGER NOT NULL, + t_end INTEGER NOT NULL, + version TEXT, + size_bytes INTEGER, + mtime REAL +); +CREATE UNIQUE INDEX IF NOT EXISTS frame_slot ON frame (satellite, wavelength, t_start); +CREATE INDEX IF NOT EXISTS frame_time ON frame (t_start); + +-- Cached FITS header metadata. Populated by a ~18 KB read per file; no pixel +-- decompression. read_ok=0 means the file could not be parsed, and read_error says +-- why -- itself a strong bad-frame signal. +CREATE TABLE IF NOT EXISTS header ( + frame_id INTEGER PRIMARY KEY REFERENCES frame (id) ON DELETE CASCADE, + {_HEADER_DDL}, + read_ok INTEGER NOT NULL, + read_error TEXT, + scanned_at REAL NOT NULL +); + +-- One row per (detector, configuration) execution, so results from different +-- methods and different thresholds coexist and can be compared. +CREATE TABLE IF NOT EXISTS detector_run ( + id INTEGER PRIMARY KEY, + name TEXT NOT NULL, + config_json TEXT NOT NULL, + code_version TEXT, + created_at REAL NOT NULL, + notes TEXT +); +CREATE INDEX IF NOT EXISTS detector_run_name ON detector_run (name); + +CREATE TABLE IF NOT EXISTS detection ( + run_id INTEGER NOT NULL REFERENCES detector_run (id) ON DELETE CASCADE, + frame_id INTEGER NOT NULL REFERENCES frame (id) ON DELETE CASCADE, + verdict TEXT NOT NULL, + reason TEXT, + scores_json TEXT, + elapsed_us INTEGER, + PRIMARY KEY (run_id, frame_id) +) WITHOUT ROWID; +CREATE INDEX IF NOT EXISTS detection_frame ON detection (frame_id); + +-- One row per day-directory, holding the mtime it had when we last read it. +-- +-- This is what makes re-indexing cheap. A directory's mtime changes whenever an +-- entry is added or removed, so comparing it is an exact test for "did anything +-- change in here" -- no guessing, no walking the files inside. The archive holds +-- ~360 frames per day-directory, so checking ~7,400 directory stats replaces +-- 2.65M file lookups: a 360x reduction, and the difference between an operation +-- this filesystem sustains and one that exhausts it. +CREATE TABLE IF NOT EXISTS dir_scan ( + path TEXT PRIMARY KEY, + mtime REAL NOT NULL, + n_frames INTEGER NOT NULL, + scanned_at REAL NOT NULL +) WITHOUT ROWID; + +-- Download bookkeeping, migrated out of the 800 MB file_database.json that +-- puller_fits.py used to parse into memory on every run. +CREATE TABLE IF NOT EXISTS remote_file ( + url TEXT PRIMARY KEY, + remote_mtime REAL NOT NULL, + remote_size INTEGER, + local_path TEXT, + fetched_at REAL +) WITHOUT ROWID; + +-- ---------------------------------------------------------------- test bench ---- +-- A stretch of archive vetted as known-good, frozen so that repeated experiments +-- are measured against the same ground truth. +CREATE TABLE IF NOT EXISTS bench_window ( + id INTEGER PRIMARY KEY, + name TEXT NOT NULL UNIQUE, + t_start INTEGER NOT NULL, + t_end INTEGER NOT NULL, + satellites TEXT NOT NULL, + wavelengths TEXT NOT NULL, + vetted_at REAL, + notes TEXT +); + +CREATE TABLE IF NOT EXISTS bench_truth ( + window_id INTEGER NOT NULL REFERENCES bench_window (id) ON DELETE CASCADE, + frame_id INTEGER NOT NULL REFERENCES frame (id) ON DELETE CASCADE, + is_good INTEGER NOT NULL, + -- 'legacy' (old filename suffix), 'header' (cross-check), or 'manual'. + vet_source TEXT NOT NULL, + PRIMARY KEY (window_id, frame_id) +) WITHOUT ROWID; + +-- One experiment: a window plus a seeded plan of deletions and corruptions. +CREATE TABLE IF NOT EXISTS bench_case ( + id INTEGER PRIMARY KEY, + window_id INTEGER NOT NULL REFERENCES bench_window (id) ON DELETE CASCADE, + name TEXT NOT NULL UNIQUE, + seed INTEGER NOT NULL, + plan_json TEXT NOT NULL, + overlay_dir TEXT NOT NULL, + created_at REAL NOT NULL +); + +-- What was done to each affected slot. mode='delete' means the frame was withheld +-- entirely; anything else names a corruption from suvi.corruptions.CATALOG. +CREATE TABLE IF NOT EXISTS bench_injection ( + case_id INTEGER NOT NULL REFERENCES bench_case (id) ON DELETE CASCADE, + frame_id INTEGER NOT NULL REFERENCES frame (id) ON DELETE CASCADE, + mode TEXT NOT NULL, + severity REAL NOT NULL DEFAULT 1.0, + params_json TEXT, + overlay_path TEXT, + gap_index INTEGER NOT NULL DEFAULT 0, + gap_length INTEGER NOT NULL DEFAULT 1, + PRIMARY KEY (case_id, frame_id) +) WITHOUT ROWID; + +-- Fill quality, one row per reconstructed frame, so results survive between runs. +CREATE TABLE IF NOT EXISTS bench_fill_result ( + case_id INTEGER NOT NULL REFERENCES bench_case (id) ON DELETE CASCADE, + frame_id INTEGER NOT NULL REFERENCES frame (id) ON DELETE CASCADE, + filler TEXT NOT NULL, + source TEXT NOT NULL, -- 'oracle' or a detector name + scores_json TEXT NOT NULL, + PRIMARY KEY (case_id, frame_id, filler, source) +) WITHOUT ROWID; +""" + +#: Verdicts a detector may record. 'unknown' covers frames a detector declined to +#: judge (e.g. a temporal detector with no usable neighbours), which must not be +#: silently conflated with 'good'. +VERDICTS = ("good", "bad", "unknown") + + +def connect(path=None, readonly=False, timeout=60.0): + """Open the index, creating and initialising it if needed. + + WAL is preferred but not required: this archive lives on a shared mount whose + locking primitives may not support it, so a failure to enter WAL falls back to + the rollback journal rather than aborting. + """ + path = path or paths.default_db_path() + if readonly: + if not os.path.exists(path): + raise FileNotFoundError(f"No SUVI index at {path}") + conn = sqlite3.connect(f"file:{path}?mode=ro", uri=True, timeout=timeout) + else: + parent = os.path.dirname(os.path.abspath(path)) + if parent: + os.makedirs(parent, exist_ok=True) + conn = sqlite3.connect(path, timeout=timeout) + + conn.row_factory = sqlite3.Row + conn.execute("PRAGMA foreign_keys = ON") + conn.execute(f"PRAGMA busy_timeout = {int(timeout * 1000)}") + if not readonly: + try: + conn.execute("PRAGMA journal_mode = WAL") + except sqlite3.Error: + pass + conn.execute("PRAGMA synchronous = NORMAL") + init_schema(conn) + return conn + + +def init_schema(conn): + """Create the schema if absent, upgrading an older index in place. Idempotent.""" + conn.executescript(SCHEMA) + row = conn.execute("SELECT value FROM meta WHERE key = 'schema_version'").fetchone() + if row is None: + conn.execute( + "INSERT INTO meta (key, value) VALUES ('schema_version', ?)", + (str(SCHEMA_VERSION),), + ) + conn.commit() + return + + found = int(row["value"]) + if found == SCHEMA_VERSION: + return + if found in ADDITIVE_UPGRADES_FROM: + # executescript above already created any new tables. + conn.execute( + "UPDATE meta SET value = ? WHERE key = 'schema_version'", (str(SCHEMA_VERSION),) + ) + conn.commit() + return + raise RuntimeError( + f"Index schema version {found} cannot be upgraded to {SCHEMA_VERSION} " + f"automatically; migrate or rebuild the index." + ) + + +def set_meta(conn, key, value): + conn.execute( + "INSERT INTO meta (key, value) VALUES (?, ?) " + "ON CONFLICT (key) DO UPDATE SET value = excluded.value", + (key, str(value)), + ) + + +def get_meta(conn, key, default=None): + row = conn.execute("SELECT value FROM meta WHERE key = ?", (key,)).fetchone() + return row["value"] if row else default + + +# --------------------------------------------------------------------------- frames + + +def upsert_frame(conn, name, relpath, size_bytes=None, mtime=None): + """Insert or update one frame row. Returns its id. + + Keyed on path, with the (satellite, wavelength, t_start) slot kept unique so a + duplicate observation stored under a second version string is rejected loudly + rather than silently shadowing the first. + """ + conn.execute( + """ + INSERT INTO frame (path, satellite, wavelength, t_start, t_end, + version, size_bytes, mtime) + VALUES (?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT (path) DO UPDATE SET + size_bytes = excluded.size_bytes, + mtime = excluded.mtime + """, + ( + relpath, + name.satellite, + name.wavelength, + name.t_start, + name.t_end, + name.version, + size_bytes, + mtime, + ), + ) + return conn.execute("SELECT id FROM frame WHERE path = ?", (relpath,)).fetchone()["id"] + + +def frame_id_by_slot(conn, satellite, wavelength, t_start): + row = conn.execute( + "SELECT id FROM frame WHERE satellite = ? AND wavelength = ? AND t_start = ?", + (satellite, wavelength, t_start), + ).fetchone() + return row["id"] if row else None + + +def frames_in_dir(conn, dir_relpath): + """Frames the index believes live in one directory: {filename: frame_id}.""" + prefix = dir_relpath.rstrip("/") + "/" + rows = conn.execute( + "SELECT id, path FROM frame WHERE path >= ? AND path < ?", + (prefix, prefix + "ï¿¿"), + ).fetchall() + return {row["path"].rsplit("/", 1)[-1]: row["id"] for row in rows} + + +def delete_frames(conn, frame_ids): + """Drop frames that have disappeared from disk, and their dependent rows.""" + conn.executemany("DELETE FROM frame WHERE id = ?", [(i,) for i in frame_ids]) + + +# ------------------------------------------------------------------ directory scans + + +def get_dir_mtime(conn, dir_relpath): + row = conn.execute("SELECT mtime FROM dir_scan WHERE path = ?", (dir_relpath,)).fetchone() + return row["mtime"] if row else None + + +def dir_mtimes(conn): + """Every recorded directory mtime, as one dict. + + Read in a single query rather than per directory: the whole point of this table + is to answer thousands of "has this changed?" questions without touching the + filesystem, and it would be perverse to pay a round trip each time. + """ + return { + row["path"]: row["mtime"] for row in conn.execute("SELECT path, mtime FROM dir_scan") + } + + +def record_dir_scans(conn, records): + """`records` yields (dir_relpath, mtime, n_frames).""" + now = time.time() + conn.executemany( + """ + INSERT INTO dir_scan (path, mtime, n_frames, scanned_at) VALUES (?, ?, ?, ?) + ON CONFLICT (path) DO UPDATE SET + mtime = excluded.mtime, + n_frames = excluded.n_frames, + scanned_at = excluded.scanned_at + """, + [(path, mtime, count, now) for path, mtime, count in records], + ) + + +def forget_dir_scans(conn, dir_relpaths): + conn.executemany("DELETE FROM dir_scan WHERE path = ?", [(p,) for p in dir_relpaths]) + + +# ----------------------------------------------------------------------- detections + + +def create_detector_run(conn, name, config, code_version=None, notes=None): + """Register a new detector execution and return its id.""" + cur = conn.execute( + """ + INSERT INTO detector_run (name, config_json, code_version, created_at, notes) + VALUES (?, ?, ?, ?, ?) + """, + (name, json.dumps(config, sort_keys=True), code_version, time.time(), notes), + ) + conn.commit() + return cur.lastrowid + + +def latest_run_id(conn, name, config=None): + """Most recent run of a named detector, or None. + + With `config`, only runs recorded under exactly that configuration match. This + matters more than it looks: detector runs accumulate, one per detector per bench + case, so "the latest run of header_v1" is whichever case was processed most + recently -- not the one being asked about. Scoring one case's verdicts against + another case's injections produces numbers that are wrong rather than merely + imprecise, so callers working with a specific case must pass its config. + """ + if config is None: + row = conn.execute( + "SELECT id FROM detector_run WHERE name = ? " + "ORDER BY created_at DESC, id DESC LIMIT 1", + (name,), + ).fetchone() + else: + row = conn.execute( + "SELECT id FROM detector_run WHERE name = ? AND config_json = ? " + "ORDER BY created_at DESC, id DESC LIMIT 1", + (name, json.dumps(config, sort_keys=True)), + ).fetchone() + return row["id"] if row else None + + +def record_detections(conn, run_id, results): + """Write detection rows. `results` yields (frame_id, verdict, reason, scores, us).""" + conn.executemany( + """ + INSERT INTO detection (run_id, frame_id, verdict, reason, scores_json, elapsed_us) + VALUES (?, ?, ?, ?, ?, ?) + ON CONFLICT (run_id, frame_id) DO UPDATE SET + verdict = excluded.verdict, + reason = excluded.reason, + scores_json = excluded.scores_json, + elapsed_us = excluded.elapsed_us + """, + [ + ( + run_id, + frame_id, + _check_verdict(verdict), + reason, + json.dumps(scores, sort_keys=True) if scores is not None else None, + elapsed_us, + ) + for frame_id, verdict, reason, scores, elapsed_us in results + ], + ) + + +def _check_verdict(verdict): + if verdict not in VERDICTS: + raise ValueError(f"Unknown verdict {verdict!r}; expected one of {VERDICTS}") + return verdict + + +def good_slots(conn, run_id, t_start, t_end, satellite, wavelengths=paths.WAVELENGTHS): + """Frames a run judged good, in a half-open time window. + + This is the query that replaces ``merger_FITS.py``'s ``_f.fits$`` regex. + """ + placeholders = ",".join("?" * len(wavelengths)) + rows = conn.execute( + f""" + SELECT f.id, f.path, f.satellite, f.wavelength, f.t_start + FROM frame f + JOIN detection d ON d.frame_id = f.id AND d.run_id = ? + WHERE d.verdict = 'good' + AND f.satellite = ? + AND f.t_start >= ? AND f.t_start < ? + AND f.wavelength IN ({placeholders}) + ORDER BY f.t_start, f.wavelength + """, + (run_id, satellite, t_start, t_end, *wavelengths), + ).fetchall() + return rows + + +def unscanned_frames(conn, limit=None): + """Frames with no cached header yet.""" + sql = """ + SELECT f.id, f.path FROM frame f + LEFT JOIN header h ON h.frame_id = f.id + WHERE h.frame_id IS NULL + ORDER BY f.t_start + """ + if limit is not None: + sql += f" LIMIT {int(limit)}" + return conn.execute(sql).fetchall() + + +# --------------------------------------------------------------------------- headers + + +_HEADER_INSERT = f""" +INSERT INTO header (frame_id, {", ".join(HEADER_FIELDS)}, read_ok, read_error, scanned_at) +VALUES (?{", ?" * len(HEADER_FIELDS)}, ?, ?, ?) +ON CONFLICT (frame_id) DO UPDATE SET + {", ".join(f"{f} = excluded.{f}" for f in HEADER_FIELDS)}, + read_ok = excluded.read_ok, + read_error = excluded.read_error, + scanned_at = excluded.scanned_at +""" + + +def record_headers(conn, records): + """Write cached header rows. `records` yields (frame_id, values_dict, error).""" + now = time.time() + conn.executemany( + _HEADER_INSERT, + [ + ( + frame_id, + *(values.get(field) for field in HEADER_FIELDS), + 0 if error else 1, + error, + now, + ) + for frame_id, values, error in records + ], + ) + + +# ---------------------------------------------------------------------- remote files + + +def get_remote_mtime(conn, url): + row = conn.execute( + "SELECT remote_mtime FROM remote_file WHERE url = ?", (url,) + ).fetchone() + return row["remote_mtime"] if row else None + + +def record_remote_files(conn, records): + """`records` yields (url, remote_mtime, remote_size, local_path, fetched_at).""" + conn.executemany( + """ + INSERT INTO remote_file (url, remote_mtime, remote_size, local_path, fetched_at) + VALUES (?, ?, ?, ?, ?) + ON CONFLICT (url) DO UPDATE SET + remote_mtime = excluded.remote_mtime, + remote_size = excluded.remote_size, + local_path = excluded.local_path, + fetched_at = excluded.fetched_at + """, + list(records), + ) + + +# ---------------------------------------------------------------- single-writer loop + +#: Sentinel pushed onto a writer queue to shut it down. +STOP = None + + +def writer_loop(db_path, work_queue, batch_size=512, flush_seconds=5.0): + """Drain (kind, payload) messages from a queue into the index. + + Runs in its own process so worker processes never contend for the write lock. + Recognised kinds: 'detection', 'header', 'remote_file' -- each payload being the + row tuple that the corresponding record_* helper expects. + """ + handlers = { + "detection": _flush_detections, + "header": record_headers, + "remote_file": record_remote_files, + } + conn = connect(db_path) + pending = {kind: [] for kind in handlers} + last_flush = time.monotonic() + + def flush(): + nonlocal last_flush + for kind, rows in pending.items(): + if rows: + handlers[kind](conn, rows) + rows.clear() + conn.commit() + last_flush = time.monotonic() + + try: + while True: + try: + message = work_queue.get(timeout=1.0) + except queue.Empty: + if time.monotonic() - last_flush > flush_seconds: + flush() + continue + if message is STOP: + break + kind, payload = message + if kind not in pending: + raise ValueError(f"Unknown writer message kind: {kind!r}") + pending[kind].append(payload) + if sum(len(v) for v in pending.values()) >= batch_size: + flush() + flush() + finally: + conn.commit() + conn.close() + + +def _flush_detections(conn, rows): + """Adapter: writer payloads carry run_id per row, record_detections does not.""" + by_run = {} + for run_id, *rest in rows: + by_run.setdefault(run_id, []).append(tuple(rest)) + for run_id, results in by_run.items(): + record_detections(conn, run_id, results) diff --git a/suvi/detectors.py b/suvi/detectors.py new file mode 100644 index 0000000..132dc1c --- /dev/null +++ b/suvi/detectors.py @@ -0,0 +1,742 @@ +"""Bad-frame detectors. + +Every detector is a pure function: data in, a :class:`Verdict` out. Nothing here +renames a file, writes a plot, or touches the database -- that is the driver's job. +This is what lets the same code run in production (``filter_FITS.py``) and under the +test bench, and lets several detectors judge the same frame independently. + +Two shapes of detector: + +* **Frame-level** (:func:`header_v1`, :func:`geometry_v1`) judge one frame alone. +* **Series-level** (:func:`temporal_v1`, :func:`crosssat_v1`) judge a frame in the + context of its neighbours in time, or of the other satellite at the same instant. + They work from :class:`FrameFeatures` -- cached header values plus a small + thumbnail -- rather than full 1280x1280 arrays, which keeps a whole day of context + in memory at once. + +Every detector reports continuous ``scores`` alongside its verdict. The driver +stores those, so re-tuning a threshold later is a database query rather than another +pass over the archive. +""" + +from dataclasses import dataclass, field + +import cv2 as cv +import numpy as np + +# --------------------------------------------------------------------------- types + + +@dataclass(frozen=True) +class Verdict: + """One detector's judgement of one frame.""" + + verdict: str # 'good' | 'bad' | 'unknown' + reason: str | None = None + scores: dict = field(default_factory=dict) + + +GOOD = "good" +BAD = "bad" +UNKNOWN = "unknown" + + +@dataclass +class FrameFeatures: + """The cheap summary of a frame that series-level detectors work from.""" + + slot: tuple # (satellite, wavelength, t_start) + header: dict = field(default_factory=dict) + #: Read error from fitsio, if the frame could not be parsed at all. + error: str | None = None + #: Small float32 image (see THUMBNAIL_SIZE), or None if not loaded. + thumbnail: np.ndarray | None = None + + @property + def t_start(self): + return self.slot[2] + + @property + def wavelength(self): + return self.slot[1] + + +#: Series detectors compare thumbnails, not full frames. 128px preserves the disc, +#: active regions and gross structure while making a day of context cheap to hold. +THUMBNAIL_SIZE = 128 + + +def thumbnail(image, size=THUMBNAIL_SIZE): + """Downsample a frame for structural comparison. NaNs become zero.""" + clean = np.nan_to_num(np.asarray(image, dtype=np.float32), nan=0.0, posinf=0.0, neginf=0.0) + return cv.resize(clean, (size, size), interpolation=cv.INTER_AREA) + + +# ------------------------------------------------------------------ header_v1 + +#: Plausible radiance range of a healthy frame, per passband, in W m-2 sr-1. +#: +#: Deliberately wide -- roughly a factor of five outside the observed annual range of +#: good frames. Sampled across 24 days of 2024 on both satellites, the *rejected* +#: frames' IMG_MEAN distribution sits almost entirely inside the good one (171A: good +#: 0.288-0.452, rejected median 0.352), so a tight bound here would cost enormous +#: numbers of false positives while catching little. What this rule is actually for +#: is gross dropout and blowout -- eclipse frames run 1e-4, three orders of magnitude +#: below the floor -- and it is priced accordingly. +HEADER_MEAN_BOUNDS = { + 94: (0.004, 0.41), + 131: (0.005, 0.39), + 171: (0.057, 2.3), + 195: (0.15, 5.4), + 284: (0.13, 7.0), + 304: (0.29, 20.0), +} + +#: Solar diameter in pixels; varies ~3% over the year with Earth-Sun distance. +DIAM_SUN_RANGE = (740.0, 800.0) +#: How far the recorded sun centre may sit from the image centre, in pixels. +MAX_CRPIX_OFFSET = 8.0 + + +def header_v1(features, mean_bounds=None, diam_range=DIAM_SUN_RANGE, + max_crpix_offset=MAX_CRPIX_OFFSET): + """Judge a frame from its header alone -- no pixel decompression. + + The instrument already reports most of what matters: ``ECLIPSE`` marks frames + taken through Earth's shadow, ``EMPTY`` marks frames built from no source + imagery, and ``IMG_MEAN``/``IMG_SDEV`` summarise the radiance. On sampled + archive days every frame the pixel filter rejected also carried ``ECLIPSE=2``, + at roughly a fifth of the cost of decompressing the image. + + ``DEGRADED`` is recorded as a score but is deliberately *not* a rejection rule. + It is set for reasons unrelated to a frame's usability: across sampled days the + 195A band carried ``DEGRADED=T`` on 40/40 good frames on 2024-01-20 and 0/40 on + 2024-07-04, on both satellites. Rejecting on it would discard entire months of + that band. The bench can re-weight it from the stored score if that turns out + to be worth doing. + """ + mean_bounds = mean_bounds or HEADER_MEAN_BOUNDS + header = features.header + scores = {} + + if features.error: + return Verdict(BAD, f"unreadable: {features.error}", {"read_ok": 0.0}) + if not header: + return Verdict(UNKNOWN, "no header cached", scores) + + for field_name in ("img_mean", "img_sdev", "diam_sun", "eclipse", "degraded"): + if field_name in header: + scores[field_name] = float(header[field_name]) + + if header.get("empty"): + return Verdict(BAD, "EMPTY flag set", scores) + eclipse = header.get("eclipse") + if eclipse: + return Verdict(BAD, f"ECLIPSE flag set ({int(eclipse)})", scores) + + # A frame built from no source images carries no information regardless of what + # the radiance statistics happen to say. + if header.get("num_imgs") is not None and header["num_imgs"] < 1: + scores["num_imgs"] = float(header["num_imgs"]) + return Verdict(BAD, "no source images", scores) + + # The band the file claims must match the band its name puts it in, or the + # archive has a misfiled frame and every downstream threshold is wrong. + declared = header.get("wavelnth") + if declared is not None and int(declared) != features.wavelength: + scores["wavelnth"] = float(declared) + return Verdict(BAD, f"wavelength {int(declared)} != {features.wavelength}", scores) + + mean = header.get("img_mean") + if mean is not None: + low, high = mean_bounds.get(features.wavelength, (0.0, np.inf)) + if mean < low: + return Verdict(BAD, f"IMG_MEAN {mean:.3g} below {low:g}", scores) + if mean > high: + return Verdict(BAD, f"IMG_MEAN {mean:.3g} above {high:g}", scores) + + sdev = header.get("img_sdev") + if sdev is not None and sdev <= 0.0: + return Verdict(BAD, "zero variance", scores) + + diameter = header.get("diam_sun") + if diameter is not None and not (diam_range[0] <= diameter <= diam_range[1]): + return Verdict(BAD, f"DIAM_SUN {diameter:.1f} outside {diam_range}", scores) + + crpix1, crpix2 = header.get("crpix1"), header.get("crpix2") + if crpix1 is not None and crpix2 is not None: + offset = max(abs(crpix1 - 640.5), abs(crpix2 - 640.5)) + scores["crpix_offset"] = float(offset) + if offset > max_crpix_offset: + return Verdict(BAD, f"sun centre offset {offset:.1f}px", scores) + + return Verdict(GOOD, None, scores) + + +# ---------------------------------------------------------------- geometry_v1 + +#: Per-band radiance ceiling used to normalise the image before shape analysis. +#: Index order matches paths.WAVELENGTHS. +GEOMETRY_THRESHOLDS = {94: 0.050, 131: 0.10, 171: 1.00, 195: 1.40, 284: 1.0, 304: 2.50} + +EXPECTED_DIMS = 1280 +HALF_DIMS = EXPECTED_DIMS // 2 +VALID_RADII = (383, 394) +AVG_RADIUS = (VALID_RADII[0] + VALID_RADII[1]) // 2 +MAX_RADIUS_ERROR = 80 +MAX_CENTER_SKEW = 7 +RATIO_ABOVE_THRESH_MAX = 0.4 +RATIO_ABOVE_THRESH_MIN = 0.07 +EDGE_THRESH = 0.98 +MAX_GOF = 0.2 + +_ideal_axis_cache = {} + + +def ideal_disc_axis(): + """Column-mean profile of a perfect solar disc, used as the shape reference.""" + if "axis" not in _ideal_axis_cache: + disc = np.zeros((HALF_DIMS, HALF_DIMS)) + cv.circle(disc, (HALF_DIMS // 2, HALF_DIMS // 2), AVG_RADIUS // 2, 1, -1) + _ideal_axis_cache["axis"] = np.average(disc, 0) + return _ideal_axis_cache["axis"] + + +def geometry_v1(image, wavelength): + """The pre-existing pixel-geometry filter, preserved as the baseline. + + Normalises against a per-band radiance ceiling, then checks the fraction of + saturated pixels, the intensity-weighted centroid, the radius implied by the + 98th-percentile cumulative edges, and the fit against an ideal disc profile. + + Deliberately kept behaviourally identical to the original so bench numbers + describe what the pipeline actually did, with two known weaknesses left intact + for the bench to quantify: + + * the centre check is one-sided -- ``HALF_DIMS//2 - centre > skew`` catches a + disc displaced up/left but not one displaced down/right; + * non-finite pixels propagate as NaN through every comparison, and ``NaN > x`` + is False, so an all-NaN frame passes every check. + + The one deliberate change: the original raised on an unexpected image size and + left the file unlabelled. Here that is a ``bad`` verdict, because "no verdict" + is not a usable baseline to score against. + """ + scores = {} + if image is None: + return Verdict(BAD, "no image data", scores) + if image.shape[0] != EXPECTED_DIMS or image.shape[1] != EXPECTED_DIMS: + return Verdict(BAD, f"unexpected dimensions {image.shape}", scores) + + threshold = GEOMETRY_THRESHOLDS.get(wavelength) + if threshold is None: + return Verdict(UNKNOWN, f"no threshold for band {wavelength}", scores) + + data = cv.resize( + np.asarray(image, dtype=np.float32), + dsize=(HALF_DIMS, HALF_DIMS), + interpolation=cv.INTER_LINEAR, + ) + above = data > threshold + normalised = np.copy(data) + normalised[above] = threshold + normalised /= threshold + + ratio = np.count_nonzero(above) / data.shape[0] / data.shape[1] + scores["ratio_above_thresh"] = float(ratio) + if ratio > RATIO_ABOVE_THRESH_MAX: + return Verdict(BAD, f"ratio_above_thresh {ratio:.2f} too high", scores) + if ratio < RATIO_ABOVE_THRESH_MIN: + return Verdict(BAD, f"ratio_above_thresh {ratio:.2f} too low", scores) + + xavg = np.average(normalised, 0) + yavg = np.average(normalised, 1) + axis = list(range(HALF_DIMS)) + + centre_x = np.average(axis, 0, xavg) + centre_y = np.average(axis, 0, yavg) + scores["centre_x"] = float(centre_x) + scores["centre_y"] = float(centre_y) + good_centre = not ( + (HALF_DIMS // 2 - centre_x > MAX_CENTER_SKEW) + or (HALF_DIMS // 2 - centre_y > MAX_CENTER_SKEW) + ) + + high_x = np.argmax(np.cumsum(xavg) > (np.sum(xavg) * EDGE_THRESH)) + low_x = len(xavg) - np.argmax(np.cumsum(np.flip(xavg)) > (np.sum(xavg) * EDGE_THRESH)) + high_y = np.argmax(np.cumsum(yavg) > (np.sum(yavg) * EDGE_THRESH)) + low_y = len(yavg) - np.argmax(np.cumsum(np.flip(yavg)) > (np.sum(yavg) * EDGE_THRESH)) + # Not halved: the image is already downsampled by two in each dimension. + radius = ((high_x - low_x) + (high_y - low_y)) / 2.0 + scores["radius"] = float(radius) + good_radius = abs(radius - AVG_RADIUS) <= MAX_RADIUS_ERROR + + reference = ideal_disc_axis() + gof_x = np.sum(np.abs(xavg - reference)) / HALF_DIMS + gof_y = np.sum(np.abs(yavg - reference)) / HALF_DIMS + scores["gof_x"] = float(gof_x) + scores["gof_y"] = float(gof_y) + good_fit = not (gof_x > MAX_GOF or gof_y > MAX_GOF) + + if good_centre and good_radius and good_fit: + return Verdict(GOOD, None, scores) + failed = [ + name + for name, ok in (("centre", good_centre), ("radius", good_radius), ("fit", good_fit)) + if not ok + ] + return Verdict(BAD, "failed " + "+".join(failed), scores) + + +# -------------------------------------------------------------------- disc_v1 + +#: Angular samples in the polar resampling. The measurements are all averages over +#: this axis, so more angles cost little and reduce noise. +DISC_ANGLES = 720 +#: Working resolution. Halving 1280 keeps the limb several pixels wide while making +#: the polar transform cheap. +DISC_SIZE = 640 + +#: Radial bands, in units of the expected solar radius, used to characterise the +#: profile. The on-disc band avoids both the centre (where limb darkening and +#: filaments live) and the limb itself; the off-limb band sits outside the corona's +#: steepest falloff but inside the frame -- 1.45R of a ~193px radius is 280px, within +#: the 320px half-width. +DISC_ON_BAND = (0.30, 0.70) +DISC_OFF_BAND = (1.25, 1.45) +#: Where the limb is allowed to be found, again in units of expected radius. +DISC_SEARCH_BAND = (0.75, 1.25) + + +def disc_profile(image, expected_radius, size=DISC_SIZE, angles=DISC_ANGLES): + """Measure the solar disc from its azimuthally averaged radial profile. + + Returns a dict of continuous measurements, or with None values where the frame + is too degenerate to measure. Separated from the verdict logic so that + ``filter_FITS.py calibrate`` can gather the same numbers over a sample. + + Averaging over angle is the entire point. An active region is localised in + angle, so it barely shifts the averaged profile; the existing ``geometry_v1`` + averages along image rows and columns instead, which a bright region shifts + directly -- the cause of its 36.9% rejection rate. + + `expected_radius` is in pixels *at the working resolution*, i.e. DIAM_SUN / 4. + """ + empty = {"limb_contrast": None, "radius_ratio": None, "limb_width": None} + if image is None or image.ndim != 2 or min(image.shape) < 16: + return empty + if not expected_radius or not np.isfinite(expected_radius) or expected_radius <= 0: + return empty + + working = cv.resize( + np.nan_to_num(np.asarray(image, dtype=np.float32), nan=0.0, posinf=0.0, neginf=0.0), + (size, size), + interpolation=cv.INTER_AREA, + ) + # Radiance spans orders of magnitude; log compresses a flare so it cannot + # dominate the profile the way it dominates a linear mean. + working = np.log1p(np.clip(working, 0.0, None)) + + max_radius = size / 2.0 + polar = cv.warpPolar( + working, + (int(max_radius), angles), + (size / 2.0, size / 2.0), + max_radius, + cv.WARP_POLAR_LINEAR, + ) + profile = polar.mean(axis=0) + radii = np.arange(profile.size, dtype=np.float64) + scale = expected_radius + + def band(low, high): + mask = (radii > low * scale) & (radii < high * scale) + return profile[mask] if mask.any() else np.array([]) + + on_values, off_values = band(*DISC_ON_BAND), band(*DISC_OFF_BAND) + if on_values.size < 3 or off_values.size < 3: + return empty + on_disc, off_limb = float(np.median(on_values)), float(np.median(off_values)) + if not np.isfinite(on_disc) or not np.isfinite(off_limb): + return empty + if on_disc <= 0: + # No signal at all where the disc should be. This is a conclusion, not a + # failure to measure: report zero contrast so the caller rejects the frame + # rather than abstaining on it. + return {"limb_contrast": 0.0, "radius_ratio": None, "limb_width": None} + + contrast = (on_disc - off_limb) / on_disc + measured = {"limb_contrast": float(contrast), "radius_ratio": None, "limb_width": None} + if on_disc <= off_limb: + # No radial falloff at all: there is no disc here to measure the size of. + return measured + + # The limb is the steepest descent of the smoothed profile. Taken on the + # *averaged* profile rather than per angle: fitting each angle separately failed + # outright on frames where too few angles yielded a usable edge. + smoothed = cv.GaussianBlur(profile.astype(np.float32).reshape(1, -1), (9, 1), 0).ravel() + derivative = np.gradient(smoothed.astype(np.float64)) + low, high = DISC_SEARCH_BAND + window = (radii >= low * scale) & (radii <= high * scale) + if window.sum() < 5: + return measured + indices = np.flatnonzero(window) + local = derivative[indices] + trough = int(indices[np.argmin(local)]) + depth = derivative[trough] + if not np.isfinite(depth) or depth >= 0: + return measured # no descending edge: not a disc + + measured["radius_ratio"] = float(trough / scale) + + # Width of the descent at half its depth. A disc displaced by d smears the + # averaged limb across roughly 2d, so this is what makes the measurement + # sensitive to decentring without ever having to locate a centre. + half = depth / 2.0 + left = trough + while left > 0 and derivative[left - 1] <= half: + left -= 1 + right = trough + last = derivative.size - 1 + while right < last and derivative[right + 1] <= half: + right += 1 + measured["limb_width"] = float((right - left + 1) / scale) + return measured + + +#: Per-band acceptance ranges for the disc_v1 measurements. +#: +#: Generated by ``filter_FITS.py calibrate``; regenerate and paste the output here +#: rather than hand-editing. Derived from 400 frames per band sampled evenly across +#: the archive's whole time range, with 0/400 sample failures in every band. +#: +#: The bounds sit well outside the observed spread of good frames, because the +#: false-positive budget for this detector is <0.1%: the method it replaces was +#: discarding 36.9% of perfectly good frames, which is the entire reason it exists. +#: Only the diagnostic side of each measurement is bounded -- see CALIBRATION_SIDES +#: in filter_FITS.py. +#: +#: Note what the wide contrast floors mean in practice. Good frames in most bands +#: reach as low as 0.11-0.22 contrast, so a floor tight enough to catch a partially +#: degraded frame would reject real data. At this budget limb_contrast therefore +#: catches total signal loss (all-zero, NaN, uniform, which measure exactly 0.0) and +#: little else; size and shape faults are caught by the other two measurements. +#: 171A is the exception, where good frames never drop below 0.68. +#: +#: A measurement whose calibrated bounds cannot achieve the budget is set to None +#: here, which records it as a score without letting it fail a frame. +DISC_BOUNDS = { + 94: { + "limb_contrast": (0.02, 1.0), + "radius_ratio": (0.9631, 1.1176), + "limb_width": (0.0, 0.1217), + }, + 131: { + "limb_contrast": (0.02, 1.0), + "radius_ratio": (0.9578, 1.0854), + "limb_width": (0.0, 0.1011), + }, + 171: { + "limb_contrast": (0.3183, 1.0), + "radius_ratio": (0.9616, 1.1062), + "limb_width": (0.0, 0.1375), + }, + 195: { + "limb_contrast": (0.02, 1.0), + "radius_ratio": (0.9676, 1.1167), + "limb_width": (0.0, 0.1704), + }, + 284: { + "limb_contrast": (0.02, 1.0), + "radius_ratio": (0.9548, 1.1489), + "limb_width": (0.0, 0.3177), + }, + 304: { + "limb_contrast": (0.02, 1.0), + "radius_ratio": (0.9548, 1.0813), + "limb_width": (0.0, 0.1016), + }, +} + + +def disc_v1(image, wavelength, header=None, bounds=None): + """Verify that the frame holds a correctly sized, centred, sharp solar disc. + + Replaces ``geometry_v1``'s intensity-weighted centroid, which a large active + region drags off centre -- rejecting entire days of good data during exactly the + solar activity most worth watching. + + Everything measured here is an average over angle, so where the bright regions + sit on the disc does not move it. What it deliberately cannot see is rotation: + ``rotate`` and ``yaw_flip`` leave every azimuthally averaged quantity unchanged, + and one frame carries no absolute rotation reference. Those belong to + ``temporal_v1`` and ``crosssat_v1``. + + The expected radius comes from the frame's own ``DIAM_SUN``, which varies by + 3.3% over the year with the Earth-Sun distance. + """ + bounds = bounds or DISC_BOUNDS + if image is None: + return Verdict(BAD, "no image data", {}) + + diameter = (header or {}).get("diam_sun") + if not diameter: + return Verdict(UNKNOWN, "no DIAM_SUN in header", {}) + + # DIAM_SUN is a full diameter at full resolution; halve for radius, halve again + # for the working downsample. + expected_radius = float(diameter) / 4.0 + measured = disc_profile(image, expected_radius) + scores = {k: v for k, v in measured.items() if v is not None} + + if measured["limb_contrast"] is None: + return Verdict(UNKNOWN, "radial profile not measurable", scores) + + limits = bounds.get(wavelength) + if limits is None: + return Verdict(UNKNOWN, f"no disc bounds for band {wavelength}", scores) + + failed = [] + for name, limit in limits.items(): + if limit is None: + continue # recorded as a score, not used to judge + value = measured.get(name) + if value is None: + failed.append(f"{name} unmeasurable") + elif not (limit[0] <= value <= limit[1]): + failed.append(f"{name} {value:.3f} outside [{limit[0]:.3f}, {limit[1]:.3f}]") + + if failed: + return Verdict(BAD, "; ".join(failed), scores) + return Verdict(GOOD, None, scores) + + +# ---------------------------------------------------------------- temporal_v1 + +#: Frames either side used to build the local reference. At a 4-minute cadence, +#: +/-6 spans about 50 minutes -- long enough to be robust, short enough that real +#: solar evolution does not dominate. +TEMPORAL_WINDOW = 6 +#: Robust z-score above which a frame's brightness is judged anomalous. +TEMPORAL_BRIGHTNESS_Z = 8.0 +#: Structural difference from neighbours, relative to the neighbours' own churn. +TEMPORAL_STRUCTURE_RATIO = 6.0 +#: Below this relative difference two consecutive frames are considered identical. +FROZEN_TOLERANCE = 1e-6 + + +def _robust_z(values, index, window): + """Median-absolute-deviation z-score of values[index] against its neighbours.""" + low = max(0, index - window) + high = min(len(values), index + window + 1) + neighbours = np.array( + [values[i] for i in range(low, high) if i != index and np.isfinite(values[i])] + ) + if neighbours.size < 3: + return None + median = np.median(neighbours) + mad = np.median(np.abs(neighbours - median)) + if mad <= 0: + # A perfectly steady neighbourhood: fall back to the spread, and treat an + # exactly-constant one as uninformative rather than infinitely sensitive. + spread = neighbours.std() + if spread <= 0: + return 0.0 if values[index] == median else np.inf + return abs(values[index] - median) / spread + return abs(values[index] - median) / (1.4826 * mad) + + +def temporal_v1(series, window=TEMPORAL_WINDOW, brightness_z=TEMPORAL_BRIGHTNESS_Z, + structure_ratio=TEMPORAL_STRUCTURE_RATIO, + frozen_tolerance=FROZEN_TOLERANCE): + """Judge each frame against its neighbours in the same band and satellite. + + Catches what single-frame checks cannot: a frame that is individually plausible + but inconsistent with the minutes either side of it -- a frozen duplicate, a + brightness step, a substituted frame from another time. + + `series` must be ordered by time and come from one (satellite, wavelength) + stream; gaps are fine, they simply widen the neighbourhood in wall-clock terms. + Returns one Verdict per input frame. + """ + count = len(series) + means = [ + f.header.get("img_mean", np.nan) if not f.error else np.nan for f in series + ] + verdicts = [] + + # Structural churn between consecutive thumbnails, when they are available. + diffs = [np.nan] * count + for i in range(1, count): + a, b = series[i - 1].thumbnail, series[i].thumbnail + if a is not None and b is not None and a.shape == b.shape: + scale = float(np.abs(a).mean() + np.abs(b).mean()) / 2.0 + diffs[i] = float(np.abs(a - b).mean()) / scale if scale > 0 else np.nan + + for i, frame in enumerate(series): + scores = {} + if frame.error: + verdicts.append(Verdict(BAD, f"unreadable: {frame.error}", {"read_ok": 0.0})) + continue + + z = _robust_z(means, i, window) + if z is not None: + scores["brightness_z"] = float(z) + + # A frame identical to its predecessor means the feed stalled. + if i > 0 and np.isfinite(diffs[i]) and diffs[i] < frozen_tolerance: + scores["frame_diff"] = float(diffs[i]) + verdicts.append(Verdict(BAD, "identical to previous frame", scores)) + continue + + neighbour_diffs = np.array( + [ + diffs[j] + for j in range(max(1, i - window), min(count, i + window + 1)) + if j != i and np.isfinite(diffs[j]) + ] + ) + if np.isfinite(diffs[i]) and neighbour_diffs.size >= 3: + typical = float(np.median(neighbour_diffs)) + scores["frame_diff"] = float(diffs[i]) + scores["frame_diff_ratio"] = float(diffs[i] / typical) if typical > 0 else 0.0 + if typical > 0 and diffs[i] / typical > structure_ratio: + verdicts.append( + Verdict(BAD, f"structural jump {diffs[i] / typical:.1f}x typical", scores) + ) + continue + + if z is not None and z > brightness_z: + verdicts.append(Verdict(BAD, f"brightness z={z:.1f}", scores)) + continue + + if z is None and not np.isfinite(diffs[i]): + verdicts.append(Verdict(UNKNOWN, "insufficient context", scores)) + continue + + verdicts.append(Verdict(GOOD, None, scores)) + + return verdicts + + +# ---------------------------------------------------------------- crosssat_v1 + +#: Relative disagreement between satellites above which one of them is wrong. +CROSSSAT_MAX_DIFF = 0.35 +#: Alignment shift beyond which the two views cannot be meaningfully compared. +#: Geostationary parallax is under ~7px at 2.5 arcsec/px; more means a pointing fault. +CROSSSAT_MAX_SHIFT = 12.0 +#: Plausible range for the calibration gain between the two flight models. Wide +#: enough to absorb genuine instrument differences, narrow enough that a blackout or +#: a saturation blowout cannot be rescaled into agreement. +CROSSSAT_GAIN_RANGE = (0.2, 5.0) + + +def _gain_match(source, reference): + """Least-squares gain and offset putting `source` on `reference`'s scale. + + The two instruments are different flight models with different responses, so a + raw difference conflates calibration with genuine disagreement. + + Returns (matched, gain, offset). The gain is returned rather than hidden + because matching is otherwise *too* effective: a frame uniformly scaled by 1e-4 + fits perfectly after rescaling, so a detector that only looked at the residual + would call a total blackout a match. The size of the correction is itself + evidence. + """ + x = source.ravel().astype(np.float64) + y = reference.ravel().astype(np.float64) + finite = np.isfinite(x) & np.isfinite(y) + if finite.sum() < 16: + return source, 1.0, 0.0 + x, y = x[finite], y[finite] + variance = float(((x - x.mean()) ** 2).sum()) + if variance <= 0: + return source, 1.0, 0.0 + gain = float(((x - x.mean()) * (y - y.mean())).sum() / variance) + offset = float(y.mean() - gain * x.mean()) + return source * gain + offset, gain, offset + + +def align_shift(a, b): + """Sub-pixel translation between two thumbnails, via phase correlation.""" + if a is None or b is None or a.shape != b.shape: + return None + window = cv.createHanningWindow((a.shape[1], a.shape[0]), cv.CV_64F) + (dx, dy), _ = cv.phaseCorrelate( + np.nan_to_num(a.astype(np.float64)), np.nan_to_num(b.astype(np.float64)), window + ) + return float(np.hypot(dx, dy)) + + +def crosssat_v1(features_a, features_b, temporal_a=None, temporal_b=None, + max_diff=CROSSSAT_MAX_DIFF, max_shift=CROSSSAT_MAX_SHIFT, + gain_range=CROSSSAT_GAIN_RANGE): + """Compare the two satellites' view of the same instant. + + GOES-16 and GOES-18 observe the same Sun from geostationary orbit, so at 1 AU + their views differ only by a few pixels of parallax and by instrument + calibration. Sustained disagreement therefore means one of them is wrong. + + Deciding *which* one needs outside information: `temporal_a`/`temporal_b` are the + corresponding :func:`temporal_v1` verdicts, and whichever frame also disagrees + with its own history is blamed. With no tiebreak available both are reported + ``unknown`` rather than guessing -- a wrong attribution would discard a good + frame and keep a bad one. + + Returns (verdict_a, verdict_b). + """ + if features_a is None or features_b is None: + return Verdict(UNKNOWN, "no counterpart", {}), Verdict(UNKNOWN, "no counterpart", {}) + if features_a.error or features_b.error: + # Frame-level errors are not this detector's job; defer rather than double-count. + return Verdict(UNKNOWN, "counterpart unreadable", {}), Verdict( + UNKNOWN, "counterpart unreadable", {} + ) + + a, b = features_a.thumbnail, features_b.thumbnail + if a is None or b is None or a.shape != b.shape: + return Verdict(UNKNOWN, "no comparable thumbnail", {}), Verdict( + UNKNOWN, "no comparable thumbnail", {} + ) + + scores = {} + shift = align_shift(a, b) + if shift is not None: + scores["align_shift"] = shift + + matched, gain, offset = _gain_match(b, a) + scale = float(np.abs(a).mean() + np.abs(matched).mean()) / 2.0 + difference = float(np.abs(a - matched).mean()) / scale if scale > 0 else np.inf + scores["cross_diff"] = difference + scores["cross_gain"] = gain + scores["cross_offset"] = offset + + # A correction this large means the two frames are not on comparable scales at + # all, whatever the residual says once it has been applied. + gain_ok = gain_range[0] <= gain <= gain_range[1] + agree = difference <= max_diff and gain_ok and (shift is None or shift <= max_shift) + if agree: + return Verdict(GOOD, None, scores), Verdict(GOOD, None, scores) + + reason = f"cross-satellite disagreement {difference:.2f}" + if not gain_ok: + reason = f"cross-satellite gain {gain:.3g} outside {gain_range}" + elif shift is not None and shift > max_shift: + reason = f"cross-satellite misalignment {shift:.1f}px" + + a_suspect = temporal_a is not None and temporal_a.verdict == BAD + b_suspect = temporal_b is not None and temporal_b.verdict == BAD + if a_suspect and not b_suspect: + return Verdict(BAD, reason, scores), Verdict(GOOD, None, scores) + if b_suspect and not a_suspect: + return Verdict(GOOD, None, scores), Verdict(BAD, reason, scores) + return Verdict(UNKNOWN, reason + " (source unclear)", scores), Verdict( + UNKNOWN, reason + " (source unclear)", scores + ) + + +#: Detectors the drivers and bench expose by name. +FRAME_DETECTORS = {"header_v1": header_v1, "geometry_v1": geometry_v1, "disc_v1": disc_v1} +SERIES_DETECTORS = {"temporal_v1": temporal_v1, "crosssat_v1": crosssat_v1} +ALL_DETECTORS = tuple(FRAME_DETECTORS) + tuple(SERIES_DETECTORS) diff --git a/suvi/fillers.py b/suvi/fillers.py new file mode 100644 index 0000000..59dd042 --- /dev/null +++ b/suvi/fillers.py @@ -0,0 +1,349 @@ +"""Methods for reconstructing a missing or rejected frame. + +Every filler takes the same :class:`FillContext` and returns a replacement array, so +the bench can swap them without knowing which one it is holding. All are pure: they +read the context and return an array, nothing else. + +The methods span a deliberate range of physical sophistication, from "repeat the last +good frame" (what the pipeline does today) to a differential-rotation warp that models +how the Sun actually moves. The bench exists to say which of them is worth the cost +at which gap length. +""" + +from dataclasses import dataclass, field + +import cv2 as cv +import numpy as np + +#: Nominal solar radius in metres (IAU 2015). +R_SUN = 6.957e8 +#: Snodgrass (1983) sidereal differential rotation, degrees per day, by latitude. +SNODGRASS_A = 14.713 +SNODGRASS_B = -2.396 +SNODGRASS_C = -1.787 +#: Earth's mean orbital motion, subtracted to get the rotation an Earth-orbiting +#: observer actually sees. +EARTH_ORBIT_DEG_PER_DAY = 0.9856 +SECONDS_PER_DAY = 86400.0 + + +@dataclass +class FillContext: + """Everything a filler may draw on to reconstruct one frame.""" + + #: Nearest good frame before the gap, and how many seconds back it sits. + before: np.ndarray | None = None + dt_before: float = 0.0 + #: Nearest good frame after the gap, and how many seconds forward. + after: np.ndarray | None = None + dt_after: float = 0.0 + #: The other satellite's view of this same instant, if it has one. + counterpart: np.ndarray | None = None + #: Header of the frame being reconstructed, for the WCS a rotation warp needs. + header: dict = field(default_factory=dict) + + @property + def alpha(self): + """Position within the gap: 0 at `before`, 1 at `after`.""" + span = self.dt_before + self.dt_after + if span <= 0: + return 0.0 + return self.dt_before / span + + @property + def gap_frames(self): + """Gap width in 4-minute slots, for reporting quality against gap length.""" + return int(round((self.dt_before + self.dt_after) / 240.0)) + + +def _finite(image): + return np.nan_to_num(np.asarray(image, dtype=np.float32), nan=0.0, posinf=0.0, neginf=0.0) + + +# ---------------------------------------------------------------- simple baselines + + +def hold_last(context): + """Repeat the last good frame. + + What ``merger_FITS.py`` does today (up to ``max_time_gap`` slots). Cheap, never + invents structure, but freezes the Sun and then jumps -- the visible stutter in + the current videos. The baseline every other method must beat. + """ + if context.before is not None: + return _finite(context.before) + if context.after is not None: + return _finite(context.after) + return None + + +def linear_blend(context): + """Cross-fade between the frames bracketing the gap. + + Removes the jump that ``hold_last`` leaves, at the cost of ghosting: moving + features appear twice, faintly, rather than moving. + """ + if context.before is None: + return hold_last(context) + if context.after is None or context.before.shape != context.after.shape: + # Frames of differing size cannot be mixed; the nearer one is the best + # available answer. This is also the fallback the other fillers unwind to. + return _finite(context.before) + alpha = context.alpha + return ((1.0 - alpha) * _finite(context.before) + alpha * _finite(context.after)).astype( + np.float32 + ) + + +# ------------------------------------------------------------------- optical flow + + +def _for_flow(image): + """Compress radiance into a range optical flow can work with. + + Radiance is heavy-tailed -- a flare can be 1000x the quiet corona -- so raw + values make flow chase the brightest pixels only. log1p plus a percentile + stretch keeps faint structure in play. + """ + scaled = np.log1p(np.clip(_finite(image), 0.0, None)) + high = np.percentile(scaled, 99.5) + if high <= 0: + return np.zeros(scaled.shape, dtype=np.uint8) + return np.clip(scaled / high * 255.0, 0, 255).astype(np.uint8) + + +def optical_flow(context, use_dis=True): + """Motion-compensated interpolation between the bracketing frames. + + Estimates dense flow both ways and warps each bracket forward to the target + instant, then blends. Unlike ``linear_blend`` this moves features instead of + dissolving between them, which is what the eye reads as smooth motion. + """ + if context.before is None or context.after is None: + return linear_blend(context) + + before, after = _finite(context.before), _finite(context.after) + if before.shape != after.shape: + return linear_blend(context) + + first, second = _for_flow(before), _for_flow(after) + if use_dis: + engine = cv.DISOpticalFlow_create(cv.DISOPTICAL_FLOW_PRESET_MEDIUM) + forward = engine.calc(first, second, None) + backward = engine.calc(second, first, None) + else: + forward = cv.calcOpticalFlowFarneback( + first, second, None, 0.5, 3, 15, 3, 5, 1.2, 0 + ) + backward = cv.calcOpticalFlowFarneback( + second, first, None, 0.5, 3, 15, 3, 5, 1.2, 0 + ) + + alpha = context.alpha + height, width = before.shape + grid_x, grid_y = np.meshgrid( + np.arange(width, dtype=np.float32), np.arange(height, dtype=np.float32) + ) + # cv.remap samples the source *at* the map coordinates, so to place a feature + # where it should be at time alpha we read from where it was: a feature at x in + # `before` sits at x + forward(x) in `after`, hence at p - alpha*forward(p) when + # looking back from the interpolated frame. Adding the flow instead of + # subtracting it moves every feature the wrong way, which is worse than not + # compensating at all. + warped_before = cv.remap( + before, + grid_x - forward[..., 0] * alpha, + grid_y - forward[..., 1] * alpha, + cv.INTER_LINEAR, + borderMode=cv.BORDER_REPLICATE, + ) + warped_after = cv.remap( + after, + grid_x - backward[..., 0] * (1.0 - alpha), + grid_y - backward[..., 1] * (1.0 - alpha), + cv.INTER_LINEAR, + borderMode=cv.BORDER_REPLICATE, + ) + return ((1.0 - alpha) * warped_before + alpha * warped_after).astype(np.float32) + + +# ------------------------------------------------------------ cross-satellite fill + + +def gain_match(source, reference): + """Put `source` on `reference`'s radiance scale by least squares. + + GOES-16 and GOES-18 carry different SUVI flight models, so their radiances + differ by a roughly affine factor even when both are healthy. + """ + x = _finite(source).ravel().astype(np.float64) + y = _finite(reference).ravel().astype(np.float64) + variance = float(((x - x.mean()) ** 2).sum()) + if variance <= 0: + return np.asarray(source, dtype=np.float32) + gain = float(((x - x.mean()) * (y - y.mean())).sum() / variance) + offset = float(y.mean() - gain * x.mean()) + return (np.asarray(source, dtype=np.float32) * gain + offset).astype(np.float32) + + +def crosssat(context, align=True): + """Substitute the other satellite's view of the same instant. + + The two spacecraft see the same Sun from 1 AU, so the substitute is a real + observation of the real Sun at the right time -- not an interpolation. It + should dominate every temporal method whenever it is available, which is the + thing worth quantifying: it is unavailable in the 31% of slots where both + satellites are out simultaneously. + + Residual differences are instrument calibration (removed by gain matching) and + a few pixels of geostationary parallax (removed by alignment). + """ + if context.counterpart is None: + return None + counterpart = _finite(context.counterpart) + reference = context.before if context.before is not None else context.after + if reference is None: + return counterpart + reference = _finite(reference) + if counterpart.shape != reference.shape: + return counterpart + + if align: + window = cv.createHanningWindow( + (counterpart.shape[1], counterpart.shape[0]), cv.CV_64F + ) + (dx, dy), _ = cv.phaseCorrelate( + counterpart.astype(np.float64), reference.astype(np.float64), window + ) + matrix = np.array([[1.0, 0.0, dx], [0.0, 1.0, dy]], dtype=np.float32) + counterpart = cv.warpAffine( + counterpart, + matrix, + (counterpart.shape[1], counterpart.shape[0]), + flags=cv.INTER_LINEAR, + borderMode=cv.BORDER_REPLICATE, + ) + return gain_match(counterpart, reference) + + +# --------------------------------------------------------- solar rotation warping + + +def rotation_rate(latitude_rad, synodic=True): + """Snodgrass differential rotation in degrees per day at a given latitude.""" + sin2 = np.sin(latitude_rad) ** 2 + rate = SNODGRASS_A + SNODGRASS_B * sin2 + SNODGRASS_C * sin2**2 + return rate - EARTH_ORBIT_DEG_PER_DAY if synodic else rate + + +def _disc_radius_pixels(header, shape): + """Solar radius in pixels, from the header if possible.""" + diameter = header.get("diam_sun") + if diameter: + return float(diameter) / 2.0 + distance, scale = header.get("dsun_obs"), header.get("cdelt1") + if distance and scale: + return float(np.degrees(np.arcsin(R_SUN / distance)) * 3600.0 / scale) + return shape[0] * 0.3 # falls back to the archive's typical disc fraction + + +def _rotation_map(shape, header, delta_seconds, synodic=True): + """Inverse map: for each output pixel, where in the input it came from. + + Works in heliographic coordinates -- de-project each pixel onto the sphere, undo + the rotation that happened over `delta_seconds`, re-project. Returns + (map_x, map_y, on_disc) with NaN where the source point is not visible. + """ + height, width = shape + radius = _disc_radius_pixels(header, shape) + crpix1 = float(header.get("crpix1", (width + 1) / 2.0)) - 1.0 + crpix2 = float(header.get("crpix2", (height + 1) / 2.0)) - 1.0 + b0 = np.radians(float(header.get("solar_b0", 0.0))) + + grid_x, grid_y = np.meshgrid(np.arange(width), np.arange(height)) + x = (grid_x - crpix1) / radius + y = (grid_y - crpix2) / radius + + rho2 = x**2 + y**2 + on_disc = rho2 < 1.0 + z = np.sqrt(np.clip(1.0 - rho2, 0.0, None)) + + # Plane-of-sky -> heliographic, undoing the observer's B0 tilt. + sin_lat = y * np.cos(b0) + z * np.sin(b0) + sin_lat = np.clip(sin_lat, -1.0, 1.0) + latitude = np.arcsin(sin_lat) + longitude = np.arctan2(x, z * np.cos(b0) - y * np.sin(b0)) + + # Step the longitude back to where this material was `delta_seconds` ago. + days = delta_seconds / SECONDS_PER_DAY + source_longitude = longitude - np.radians(rotation_rate(latitude, synodic)) * days + + # Heliographic -> plane-of-sky. + cos_lat = np.cos(latitude) + sx = cos_lat * np.sin(source_longitude) + sy = sin_lat * np.cos(b0) - cos_lat * np.cos(source_longitude) * np.sin(b0) + sz = sin_lat * np.sin(b0) + cos_lat * np.cos(source_longitude) * np.cos(b0) + + visible = on_disc & (sz > 0) + map_x = (sx * radius + crpix1).astype(np.float32) + map_y = (sy * radius + crpix2).astype(np.float32) + return map_x, map_y, visible + + +def solar_rotation(context, synodic=True): + """Warp the bracketing frames by differential solar rotation, then blend. + + The Sun is not a rigid body: the equator turns in about 25 days, the poles in + about 35. Over a short gap that is a sub-pixel effect, but across a multi-hour + outage it is the difference between features landing where they belong and + smearing. This is the only method here that uses a physical model of the scene. + + Applies on-disc only. The corona above the limb does not co-rotate with the + photosphere, so off-disc pixels fall back to a plain cross-fade. + """ + if context.before is None and context.after is None: + return None + if context.before is None or context.after is None: + source = context.before if context.before is not None else context.after + delta = context.dt_before if context.before is not None else -context.dt_after + warped, visible = _warp(source, context.header, delta, synodic) + blended = np.where(visible, warped, _finite(source)) + return blended.astype(np.float32) + + before, after = _finite(context.before), _finite(context.after) + if before.shape != after.shape: + return linear_blend(context) + + # Roll `before` forward to the target instant and `after` backward to it. + warped_before, visible_before = _warp(before, context.header, context.dt_before, synodic) + warped_after, visible_after = _warp(after, context.header, -context.dt_after, synodic) + + alpha = context.alpha + rotated = (1.0 - alpha) * warped_before + alpha * warped_after + faded = (1.0 - alpha) * before + alpha * after + visible = visible_before & visible_after + return np.where(visible, rotated, faded).astype(np.float32) + + +def _warp(image, header, delta_seconds, synodic): + image = _finite(image) + map_x, map_y, visible = _rotation_map(image.shape, header, delta_seconds, synodic) + warped = cv.remap( + image, map_x, map_y, cv.INTER_LINEAR, borderMode=cv.BORDER_CONSTANT, borderValue=0.0 + ) + return warped, visible + + +FILLERS = { + "hold_last": hold_last, + "linear_blend": linear_blend, + "optical_flow": optical_flow, + "crosssat": crosssat, + "solar_rotation": solar_rotation, +} + +# TODO: learned filler. Train a model to predict a frame from its preceding frames, +# following frames, and the other satellite's view, then evaluate it here across +# severities of missing data and prediction horizons (single-frame gaps through +# multi-hour outages, one satellite out versus both). It plugs in as another entry +# in FILLERS and reuses the bench's existing cases and metrics unchanged. diff --git a/suvi/fitsio.py b/suvi/fitsio.py new file mode 100644 index 0000000..d42a188 --- /dev/null +++ b/suvi/fitsio.py @@ -0,0 +1,294 @@ +"""Cheap FITS header reads and file-integrity checks. + +The archive's science keywords (``DEGRADED``, ``ECLIPSE``, ``IMG_MEAN``, ...) live in +the tile-compressed image extension's header, in the first ~18 KB of a 1.7 MB file. +Reading just those bytes and parsing the 80-column cards by hand is ~5x faster than +letting astropy open the file (4.1 ms vs 22.2 ms per file measured on this archive), +which is the difference between a two-hour and a ten-hour pass over the full set. +astropy is still used for pixel data, where it earns its keep. + +Every function reports trouble as a structured error string rather than raising, so a +truncated or corrupt file is a *finding* the detector can score, not an exception that +kills a worker. +""" + +import logging +import os +import warnings + +import numpy as np +from astropy.io import fits +from astropy.utils.exceptions import AstropyUserWarning + +from . import db + + +def quiet_astropy(): + """Silence astropy's per-file chatter about damaged files. + + astropy reports conditions like "File may have been truncated" through both its + own logger and ``warnings``, either of which floods stderr when scanning + millions of files -- and damaged files are exactly what a bench run is full of. + Those conditions are already surfaced per-frame by ``read_header``, + ``read_image`` and ``verify_datasums``, so nothing is lost by muting them. + + Call once per process; worker pools must call it in each child. + """ + logging.getLogger("astropy").setLevel(logging.ERROR) + warnings.simplefilter("ignore", AstropyUserWarning) + +#: FITS records are a whole number of 2880-byte blocks. +BLOCK = 2880 +CARD = 80 +#: How far into a file to look for the image HDU header. The real ones end by +#: ~17 KB; the cap stops a malformed file from pulling megabytes into memory. +MAX_HEADER_BYTES = BLOCK * 16 + +#: Sentinel some SUVI products use for undefined pixels. +ZBLANK = -999999 + + +def _parse_value(raw): + """Parse a FITS card value, honouring quoted strings and comment separators.""" + raw = raw.strip() + if not raw: + return None + if raw[0] == "'": + # String value: scan for the closing quote, where '' is a literal quote. + out = [] + i = 1 + while i < len(raw): + if raw[i] == "'": + if i + 1 < len(raw) and raw[i + 1] == "'": + out.append("'") + i += 2 + continue + break + out.append(raw[i]) + i += 1 + return "".join(out).strip() + + # Non-string: everything before the comment separator. + value = raw.split("/", 1)[0].strip() + if value in ("T", "F"): + return value == "T" + if not value: + return None + try: + return int(value) + except ValueError: + pass + try: + # FITS permits Fortran-style 'D' exponents. + return float(value.replace("D", "E").replace("d", "e")) + except ValueError: + return value + + +def _parse_header_blocks(buf, pos): + """Parse cards from `pos` until the END card. + + Returns (cards, next_pos), or (None, reason) if END is not reached. + """ + cards = {} + while pos + BLOCK <= len(buf): + block = buf[pos : pos + BLOCK] + pos += BLOCK + for i in range(0, BLOCK, CARD): + card = block[i : i + CARD] + key = card[:8].rstrip() + if key == b"END": + return cards, pos + # Only fixed-format valued cards; COMMENT/HISTORY/CONTINUE lack the '= '. + if card[8:10] == b"= ": + try: + name = key.decode("ascii").strip() + value = _parse_value(card[10:].decode("ascii")) + except UnicodeDecodeError: + continue + if name and name not in cards: + cards[name] = value + return None, "header has no END card within the scanned region" + + +def _data_bytes(cards): + """Size of an HDU's data unit, in bytes, rounded up to whole blocks.""" + naxis = cards.get("NAXIS", 0) + if not isinstance(naxis, int) or naxis <= 0: + return 0 + count = 1 + for axis in range(1, naxis + 1): + length = cards.get(f"NAXIS{axis}") + if not isinstance(length, int) or length < 0: + return 0 + count *= length + bitpix = cards.get("BITPIX", 8) + if not isinstance(bitpix, int): + return 0 + size = count * abs(bitpix) // 8 + size += cards.get("PCOUNT", 0) or 0 + return ((size + BLOCK - 1) // BLOCK) * BLOCK + + +def _is_image_hdu(cards): + """Whether this HDU carries the image we care about. + + SUVI L2 stores the image tile-compressed inside a BINTABLE, flagged by ZIMAGE; + its real dimensions are in ZNAXIS, not NAXIS. A plain IMAGE extension or a + primary HDU with data also counts. + """ + if cards.get("ZIMAGE") is True: + return True + return isinstance(cards.get("NAXIS"), int) and cards["NAXIS"] > 0 + + +def read_header(path, max_bytes=MAX_HEADER_BYTES): + """Read the image HDU's header without decompressing pixels. + + Returns (cards, error). On failure cards is whatever was parsed (possibly + empty) and error is a short human-readable reason. + """ + try: + size = os.path.getsize(path) + with open(path, "rb") as handle: + buf = handle.read(max_bytes) + except OSError as exc: + return {}, f"unreadable: {exc.strerror or exc}" + + if size == 0: + return {}, "empty file" + if len(buf) < BLOCK: + return {}, f"truncated: {size} bytes is less than one {BLOCK}-byte block" + if not buf.startswith(b"SIMPLE ="): + return {}, "not a FITS file: missing SIMPLE keyword" + if size % BLOCK: + # Worth reporting even though the header may still parse: it means the file + # was cut short mid-transfer, which is exactly the download failure mode. + truncation = f"truncated: {size} bytes is not a multiple of {BLOCK}" + else: + truncation = None + + pos = 0 + while True: + cards, result = _parse_header_blocks(buf, pos) + if cards is None: + return {}, result + pos = result + if _is_image_hdu(cards): + return cards, truncation + skip = _data_bytes(cards) + if pos + skip + BLOCK > len(buf): + return cards, truncation or "no image HDU within the scanned region" + pos += skip + + +def header_metadata(cards): + """Project parsed cards onto the ``header`` table's columns. + + Booleans become 0/1 so they sort and filter naturally in SQL; anything of an + unexpected type is dropped rather than stored, so a corrupt card cannot poison + a numeric column. + """ + values = {} + for field, keyword, sqltype in db.HEADER_COLUMNS: + raw = cards.get(keyword) + if raw is None: + continue + if isinstance(raw, bool): + values[field] = int(raw) + elif sqltype == "TEXT": + values[field] = str(raw) + elif sqltype == "INTEGER": + values[field] = int(raw) if isinstance(raw, (int, float)) else None + elif isinstance(raw, (int, float)): + values[field] = float(raw) + return {k: v for k, v in values.items() if v is not None} + + +def scan_header(path): + """Read and project a file's header in one step. Returns (values, error).""" + cards, error = read_header(path) + return header_metadata(cards), error + + +def read_image(path): + """Read the decompressed image array. Returns (data, error). + + Returns a float32 array with the ZBLANK sentinel converted to NaN, so callers + do not have to remember that -999999 means "no data" rather than a radiance. + + Detects damage only insofar as astropy refuses to decode it; a file whose image + HDU survives intact reads fine even if later bytes are missing. Pair this with + ``verify_datasums`` when integrity matters rather than availability. + """ + try: + with fits.open(path, memmap=False) as hdus: + data = None + for hdu in hdus: + if getattr(hdu, "data", None) is not None and hdu.data.ndim == 2: + data = hdu.data + break + if data is None: + # The blank/corrupt-HDU case filter_FITS.py caught as IndexError. + return None, "no 2-D image array in file" + except Exception as exc: # astropy raises a wide variety on malformed input + return None, f"{type(exc).__name__}: {exc}" + + data = np.asarray(data, dtype=np.float32) + blank = data == ZBLANK + if blank.any(): + data = data.copy() + data[blank] = np.nan + return data, None + + +def datasum(payload): + """FITS-standard 32-bit ones-complement checksum of a data unit.""" + if len(payload) % 4: + payload = payload + b"\x00" * (4 - len(payload) % 4) + total = int(np.frombuffer(payload, dtype=">u4").astype(np.uint64).sum()) + while total >> 32: # fold the carries back in + total = (total & 0xFFFFFFFF) + (total >> 32) + return total + + +def verify_datasums(path): + """Check every HDU's stored DATASUM against the bytes on disk. + + Catches silent corruption that a size check misses. Computed over the raw file, + deliberately not via astropy: astropy re-serialises a tile-compressed HDU before + summing, so it reports a mismatch for every healthy SUVI file. + + Returns an error string, or None if all present DATASUMs agree. + """ + try: + with open(path, "rb") as handle: + buf = handle.read() + except OSError as exc: + return f"unreadable: {exc.strerror or exc}" + + if len(buf) < BLOCK or not buf.startswith(b"SIMPLE ="): + return "not a FITS file" + + pos = 0 + checked = 0 + while pos < len(buf): + cards, result = _parse_header_blocks(buf, pos) + if cards is None: + return result + size = _data_bytes(cards) + if result + size > len(buf): + return "truncated: HDU data unit runs past end of file" + stored = cards.get("DATASUM") + if stored is not None: + try: + expected = int(str(stored).strip()) + except ValueError: + return f"malformed DATASUM keyword: {stored!r}" + if datasum(buf[result : result + size]) != expected: + name = cards.get("EXTNAME") or ("PRIMARY" if pos == 0 else "extension") + return f"datasum mismatch in {name}" + checked += 1 + pos = result + size + + return None if checked else "no DATASUM keywords present" diff --git a/suvi/index.py b/suvi/index.py new file mode 100644 index 0000000..4afa6eb --- /dev/null +++ b/suvi/index.py @@ -0,0 +1,253 @@ +"""Keeping the frame index in step with the archive, without walking it. + +The index exists so that nothing has to traverse 2.65M files to answer "what frames +do we have". Traversal is not merely slow here: the archive sits on a virtiofs mount +where the daemon holds a host file descriptor per inode the guest has looked up, so a +full walk exhausts the host's descriptors and the mount then refuses *every* open +until the guest drops its dentry cache. One full pass cost 152,193 renames to ENFILE +before this was understood. + +So the index is maintained three ways, in descending order of preference: + +1. **At the source.** ``puller_fits.py`` records each frame as it downloads it, so + new data enters the index with no filesystem traversal at all. +2. **By directory mtime** (:func:`reconcile`). A directory's mtime changes whenever + an entry is added or removed, so comparing it against ``dir_scan`` is an exact + test for "did anything change in here". The archive averages ~360 frames per + day-directory, so this checks ~7,400 stats instead of 2.65M lookups -- a 360x + reduction, and the difference between an operation this mount sustains and one + that breaks it. +3. **A full rebuild**, only for a cold start or when the index is suspect. + +Only step 3 is expensive, and after the first one it should never be needed again. +""" + +import os +import time + +from . import db, paths, vfs + + +def day_directories(root, satellites=paths.SATELLITES, wavelengths=paths.WAVELENGTHS, + years=None): + """Yield (relpath, abspath) for every day-directory that exists. + + Enumerated by descending the year/month/day structure rather than walking, so the + cost is one listing per year and month directory -- a couple of thousand small + reads -- and no per-file lookups at all. + """ + for satellite in satellites: + for wavelength in wavelengths: + band_rel = f"goes{satellite:d}/l2/data/suvi-l2-ci{wavelength:03d}" + band_abs = os.path.join(root, *band_rel.split("/")) + for year in _subdirs(band_abs): + if not year.isdigit() or (years and int(year) not in years): + continue + year_abs = os.path.join(band_abs, year) + for month in _subdirs(year_abs): + month_abs = os.path.join(year_abs, month) + for day in _subdirs(month_abs): + yield ( + f"{band_rel}/{year}/{month}/{day}", + os.path.join(month_abs, day), + ) + + +def _subdirs(path): + """Immediate subdirectory names. + + A directory that does not exist yields nothing; one that exists but cannot be + read raises. The distinction is essential: swallowing the error makes an + unreadable tree look like an empty one, and :func:`reconcile` would then conclude + its contents had been deleted. That is not hypothetical -- an ENFILE partway + through enumeration made five of six bands look absent, and 205,618 index rows + were dropped as "vanished" before this was caught. + """ + try: + with os.scandir(path) as entries: + return sorted(entry.name for entry in entries if entry.is_dir()) + except (FileNotFoundError, NotADirectoryError): + return [] + + +def scan_directory(conn, dir_relpath, dir_abspath): + """Bring one directory's frames into the index. + + Returns (added, removed, duplicates, examined). `examined` is the number of + frame files looked at, which is what drives how close the mount is to running + out of file handles -- see suvi.vfs. + + Reconciles in both directions: files that appeared are inserted, and index rows + whose files have gone are deleted, so the index stays authoritative rather than + merely append-only. + """ + try: + with os.scandir(dir_abspath) as entries: + found = {} + for entry in entries: + name = paths.parse_frame_filename(entry.name) + if name is None: + continue + try: + stat = entry.stat() + found[entry.name] = (name, stat.st_size, stat.st_mtime) + except OSError: + found[entry.name] = (name, None, None) + except OSError: + return 0, 0, [], 0 + + examined = len(found) + found, duplicates = _resolve_slot_collisions(found) + + known = db.frames_in_dir(conn, dir_relpath) + + # Remove departed rows *before* inserting, so a file renamed within a directory + # -- which is exactly what the un-rename migration does -- does not momentarily + # have two rows claiming one observation slot and trip the unique constraint. + gone = [frame_id for filename, frame_id in known.items() if filename not in found] + if gone: + db.delete_frames(conn, gone) + + added = 0 + for filename, (name, size, mtime) in found.items(): + db.upsert_frame(conn, name, f"{dir_relpath}/{filename}", size, mtime) + if filename not in known: + added += 1 + return added, len(gone), duplicates, examined + + +def _resolve_slot_collisions(found): + """Keep one file per observation slot; return the rest as duplicates. + + Two files can claim the same (satellite, band, time) when an older filter + labelled a frame repeatedly -- the archive holds 1,742 such pairs, a clean + ``...v1-0-1.fits`` beside a ``...v1-0-1_f_f_f.fits`` of byte-identical content, + left when the puller re-downloaded a frame it could no longer find under its + published name. The index cannot hold both, and crashing on them would block + indexing the entire archive over a handful of stale copies. Prefer the file + NOAA actually published: unlabelled first, shortest name to break ties. + """ + by_slot = {} + for filename, (name, _, _) in found.items(): + by_slot.setdefault(name.slot, []).append(filename) + + duplicates = [] + for names in by_slot.values(): + if len(names) < 2: + continue + canonical = min( + names, key=lambda n: (found[n][0].label is not None, len(n), n) + ) + for other in names: + if other != canonical: + duplicates.append(other) + found.pop(other, None) + return found, duplicates + + +def reconcile(conn, root=None, satellites=paths.SATELLITES, wavelengths=paths.WAVELENGTHS, + years=None, force=False, progress=None, relief=True): + """Update the index from the archive, reading only what changed. + + Returns a summary dict. With `force`, every directory is re-read regardless of + its recorded mtime -- the escape hatch for when the index is suspected wrong. + + A steady-state run touches almost nothing and needs no special care. A cold + build is different: it reads every file in the archive, which is exactly the + traversal that exhausts this mount's file handles, so it hands them back + periodically (see suvi.vfs). Pass ``relief=False`` to disable that. + """ + root = root or paths.data_root() + recorded = {} if force else db.dir_mtimes(conn) + + checked = changed = added = removed = scanned = 0 + next_relief = vfs.RELIEF_INTERVAL + duplicates = [] + seen = set() + pending = [] + started = time.time() + + for dir_relpath, dir_abspath in day_directories(root, satellites, wavelengths, years): + seen.add(dir_relpath) + checked += 1 + try: + mtime = os.stat(dir_abspath).st_mtime + except OSError: + continue + if not force and recorded.get(dir_relpath) == mtime: + continue # nothing added or removed since we last looked + + changed += 1 + new, lost, dupes, examined = scan_directory(conn, dir_relpath, dir_abspath) + added += new + removed += lost + duplicates.extend(f"{dir_relpath}/{name}" for name in dupes) + # Record the mtime we actually observed, not one read afterwards, so a write + # racing this scan leaves the directory looking stale and gets picked up next + # time rather than being silently skipped forever. + pending.append((dir_relpath, mtime, len(db.frames_in_dir(conn, dir_relpath)))) + if len(pending) >= 200: + db.record_dir_scans(conn, pending) + conn.commit() + pending.clear() + scanned += examined + if relief and scanned >= next_relief: + vfs.release_handles() + next_relief = scanned + vfs.RELIEF_INTERVAL + if progress and changed % 100 == 0: + progress(f" {changed} changed of {checked} checked, +{added}/-{removed}") + + if pending: + db.record_dir_scans(conn, pending) + + # Directories that have vanished entirely. + # + # Absence from `seen` is not sufficient evidence to delete anything: enumeration + # can come up short for reasons that have nothing to do with the data, and the + # cost of being wrong is destroying index rows for files that are still present. + # Confirm each one is genuinely gone before acting on it. + vanished = [] + for dir_relpath in recorded: + if dir_relpath in seen: + continue + if os.path.isdir(paths.abspath(dir_relpath, root)): + continue # still there; enumeration simply missed it + vanished.append(dir_relpath) + + for dir_relpath in vanished: + stale = db.frames_in_dir(conn, dir_relpath) + if stale: + db.delete_frames(conn, list(stale.values())) + removed += len(stale) + if vanished: + db.forget_dir_scans(conn, vanished) + conn.commit() + + return { + "directories_checked": checked, + "directories_changed": changed, + "directories_vanished": len(vanished), + "frames_added": added, + "frames_removed": removed, + "duplicate_slots": duplicates, + "seconds": time.time() - started, + } + + +def record_downloaded(conn, local_path, root=None): + """Index a frame the puller has just written. + + The cheapest path of all: the downloader already knows the file exists, so the + index can learn about it without anybody looking at the filesystem. Returns the + frame id, or None if the file is not a SUVI frame. + """ + root = root or paths.data_root() + name = paths.parse_frame_filename(os.path.basename(local_path)) + if name is None: + return None + try: + stat = os.stat(local_path) + size, mtime = stat.st_size, stat.st_mtime + except OSError: + size = mtime = None + return db.upsert_frame(conn, name, name.relpath(), size, mtime) diff --git a/suvi/metrics.py b/suvi/metrics.py new file mode 100644 index 0000000..5121b20 --- /dev/null +++ b/suvi/metrics.py @@ -0,0 +1,354 @@ +"""Scoring for detection and fill quality. + +Detection scoring keeps two things apart that are easy to conflate: + +* a **false positive** -- the detector flagged a slot the bench did not corrupt, and + the window's ground truth says was good; +* a **legacy disagreement** -- the detector flagged a slot the old filter passed. + +On a vetted window the second is a candidate *find*, not an error: the old filter is +what we are trying to beat, so scoring its misses against a new detector would +penalise exactly the improvement we want. + +Fill scoring is reported both in radiance and after the display mapping the videos +actually use, plus a temporal term. A frame can score well on per-frame PSNR and +still read as a visible stutter at 60 fps, so flicker is measured explicitly. +""" + +from dataclasses import dataclass, field + +import numpy as np +from skimage.metrics import structural_similarity + +#: Display mapping used by merger_FITS.py to turn radiance into pixels, by band. +#: (vmin, vmax, gamma). Fill error is reported through this because it is what the +#: eye sees -- an error in the dim corona matters far less than one on the disc. +DISPLAY_MAPPING = { + 94: (0.050, 8.0, 0.375), + 131: (0.05, 8.0, 0.40), + 171: (0.100, 20.0, 0.425), + 195: (0.10, 30.0, 0.45), + 284: (0.100, 40.0, 0.475), + 304: (0.1, 90.0, 0.5), +} + + +def to_display(image, wavelength): + """Map radiance to the 0-1 display range merger_FITS.py renders with.""" + vmin, vmax, gamma = DISPLAY_MAPPING.get(wavelength, (0.0, 1.0, 1.0)) + clean = np.nan_to_num(np.asarray(image, dtype=np.float64), nan=0.0) + return np.clip((clean - vmin) / vmax, 0.0, 1.0) ** gamma + + +# ------------------------------------------------------------------------ detection + + +@dataclass +class DetectionScore: + """Confusion counts and derived rates for one detector on one bench case.""" + + true_positives: int = 0 + false_positives: int = 0 + false_negatives: int = 0 + true_negatives: int = 0 + unknown: int = 0 + #: Slots the detector flagged that the legacy filter had passed and the bench + #: did not corrupt. Reported, never counted as errors. + legacy_disagreements: int = 0 + #: Recall broken down by corruption mode, so a detector that only catches + #: blackouts cannot hide behind a good overall number. + recall_by_mode: dict = field(default_factory=dict) + elapsed_us: int = 0 + frames: int = 0 + + @property + def precision(self): + flagged = self.true_positives + self.false_positives + return self.true_positives / flagged if flagged else float("nan") + + @property + def recall(self): + actual = self.true_positives + self.false_negatives + return self.true_positives / actual if actual else float("nan") + + @property + def f1(self): + precision, recall = self.precision, self.recall + if not np.isfinite(precision) or not np.isfinite(recall) or precision + recall == 0: + return float("nan") + return 2 * precision * recall / (precision + recall) + + @property + def false_positive_rate(self): + negatives = self.false_positives + self.true_negatives + return self.false_positives / negatives if negatives else float("nan") + + @property + def microseconds_per_frame(self): + return self.elapsed_us / self.frames if self.frames else float("nan") + + def as_dict(self): + return { + "true_positives": self.true_positives, + "false_positives": self.false_positives, + "false_negatives": self.false_negatives, + "true_negatives": self.true_negatives, + "unknown": self.unknown, + "legacy_disagreements": self.legacy_disagreements, + "precision": self.precision, + "recall": self.recall, + "f1": self.f1, + "false_positive_rate": self.false_positive_rate, + "us_per_frame": self.microseconds_per_frame, + "recall_by_mode": dict(self.recall_by_mode), + } + + +def score_detection(verdicts, injected, legacy_good=frozenset(), elapsed_us=0): + """Score detector output against what the bench actually injected. + + `verdicts` maps slot -> Verdict. `injected` maps slot -> corruption mode name + for every slot the bench damaged (a deleted slot has no frame to judge, so it is + not scored here). `legacy_good` is the set of slots the old filter passed. + + 'unknown' verdicts are counted separately and excluded from precision and recall + rather than folded into 'good' -- a detector that abstains has not made an error, + but it has not made a call either, and hiding that would flatter it. + """ + score = DetectionScore(elapsed_us=elapsed_us, frames=len(verdicts)) + by_mode = {} + + for slot, verdict in verdicts.items(): + mode = injected.get(slot) + corrupted = mode is not None + + if verdict.verdict == "unknown": + score.unknown += 1 + if corrupted: + by_mode.setdefault(mode, [0, 0])[1] += 1 + continue + + flagged = verdict.verdict == "bad" + if corrupted: + stats = by_mode.setdefault(mode, [0, 0]) + stats[1] += 1 + if flagged: + score.true_positives += 1 + stats[0] += 1 + else: + score.false_negatives += 1 + elif flagged: + score.false_positives += 1 + if slot in legacy_good: + score.legacy_disagreements += 1 + else: + score.true_negatives += 1 + + score.recall_by_mode = { + mode: (caught / total if total else float("nan")) + for mode, (caught, total) in sorted(by_mode.items()) + } + return score + + +#: How several detectors' verdicts are folded into one. +COMBINATION_POLICIES = ("any", "all", "majority") + + +def combine_verdicts(verdict_maps, policy="any"): + """Fold several detectors' verdicts into one per slot. + + Which detectors to run in production is a trade-off between recall and the + false-positive rate, and the only honest way to choose is to score the + combinations the same way as the individuals. Because `detection` rows already + store a verdict per frame, this is pure post-processing. + + * ``any`` -- bad if any detector says bad. Maximum recall; false positives + accumulate across detectors. + * ``all`` -- bad only if every detector that voted says bad. Minimum false + positives, and the right shape for a near-zero FP budget. + * ``majority`` -- bad if more than half the votes say bad. + + ``unknown`` abstains rather than voting. A detector that declines to judge must + not be silently counted as saying "good": under ``all`` that would let one + abstention veto a real detection, and under ``any`` it would inflate recall. A + slot where every detector abstains is itself ``unknown``. + + `verdict_maps` is a sequence of {slot: Verdict}. Returns {slot: Verdict}. + """ + if policy not in COMBINATION_POLICIES: + raise ValueError(f"Unknown policy {policy!r}; expected one of {COMBINATION_POLICIES}") + if not verdict_maps: + return {} + + from .detectors import BAD, GOOD, UNKNOWN, Verdict + + combined = {} + slots = set() + for verdicts in verdict_maps: + slots.update(verdicts) + + for slot in slots: + votes = [] + reasons = [] + for verdicts in verdict_maps: + verdict = verdicts.get(slot) + if verdict is None or verdict.verdict == UNKNOWN: + continue + votes.append(verdict.verdict == BAD) + if verdict.verdict == BAD and verdict.reason: + reasons.append(verdict.reason) + + if not votes: + combined[slot] = Verdict(UNKNOWN, "no detector judged this frame", {}) + continue + + if policy == "any": + is_bad = any(votes) + elif policy == "all": + is_bad = all(votes) + else: + is_bad = sum(votes) * 2 > len(votes) + + combined[slot] = Verdict( + BAD if is_bad else GOOD, + "; ".join(reasons[:3]) if is_bad else None, + {"votes_bad": float(sum(votes)), "votes_total": float(len(votes))}, + ) + return combined + + +def precision_recall_curve(scores, labels): + """Precision/recall across every threshold of a continuous detector score. + + `scores` are higher-is-more-suspicious. Returns (thresholds, precision, recall). + """ + scores = np.asarray(scores, dtype=np.float64) + labels = np.asarray(labels, dtype=bool) + finite = np.isfinite(scores) + scores, labels = scores[finite], labels[finite] + if scores.size == 0 or not labels.any(): + return np.array([]), np.array([]), np.array([]) + + order = np.argsort(-scores) + scores, labels = scores[order], labels[order] + true_positives = np.cumsum(labels) + flagged = np.arange(1, scores.size + 1) + precision = true_positives / flagged + recall = true_positives / labels.sum() + return scores, precision, recall + + +def average_precision(scores, labels): + """Area under the precision-recall curve; the threshold-free summary.""" + _, precision, recall = precision_recall_curve(scores, labels) + if recall.size == 0: + return float("nan") + return float(np.sum(np.diff(np.concatenate([[0.0], recall])) * precision)) + + +# ----------------------------------------------------------------------------- fill + + +@dataclass +class FillScore: + """Reconstruction error for one filled frame.""" + + rmse: float + mae: float + log_rmse: float + psnr: float + ssim: float + gap_frames: int = 0 + wavelength: int = 0 + + def as_dict(self): + return { + "rmse": self.rmse, + "mae": self.mae, + "log_rmse": self.log_rmse, + "psnr": self.psnr, + "ssim": self.ssim, + "gap_frames": self.gap_frames, + "wavelength": self.wavelength, + } + + +def score_fill(filled, truth, wavelength, gap_frames=0): + """Compare a reconstructed frame against the frame that was withheld. + + Radiance errors (rmse/mae) are dominated by the bright disc; the log term keeps + the faint corona visible in the score; psnr and ssim are computed after the + display mapping, because that is the image a viewer actually sees. + """ + filled = np.nan_to_num(np.asarray(filled, dtype=np.float64), nan=0.0) + truth = np.nan_to_num(np.asarray(truth, dtype=np.float64), nan=0.0) + if filled.shape != truth.shape: + raise ValueError(f"shape mismatch: {filled.shape} vs {truth.shape}") + + residual = filled - truth + rmse = float(np.sqrt(np.mean(residual**2))) + mae = float(np.mean(np.abs(residual))) + + log_residual = np.log1p(np.clip(filled, 0, None)) - np.log1p(np.clip(truth, 0, None)) + log_rmse = float(np.sqrt(np.mean(log_residual**2))) + + shown_fill, shown_truth = to_display(filled, wavelength), to_display(truth, wavelength) + display_mse = float(np.mean((shown_fill - shown_truth) ** 2)) + psnr = float("inf") if display_mse == 0 else float(10.0 * np.log10(1.0 / display_mse)) + ssim = float(structural_similarity(shown_truth, shown_fill, data_range=1.0)) + + return FillScore(rmse, mae, log_rmse, psnr, ssim, gap_frames, wavelength) + + +def temporal_flicker(sequence, truth_sequence, wavelength): + """How much the reconstruction's frame-to-frame motion departs from reality. + + Per-frame PSNR is blind to the artifact that matters most in a 60 fps video: a + filled run can match each frame tolerably and still freeze then jump. This + compares the *rate of change* rather than the frames, in display space. + + Returns mean absolute difference of successive-frame deltas; 0 is perfect. + """ + if len(sequence) != len(truth_sequence): + raise ValueError("sequences must be the same length") + if len(sequence) < 2: + return float("nan") + + shown = [to_display(frame, wavelength) for frame in sequence] + shown_truth = [to_display(frame, wavelength) for frame in truth_sequence] + deltas = [np.mean(np.abs(shown[i] - shown[i - 1])) for i in range(1, len(shown))] + truth_deltas = [ + np.mean(np.abs(shown_truth[i] - shown_truth[i - 1])) for i in range(1, len(shown_truth)) + ] + return float(np.mean(np.abs(np.array(deltas) - np.array(truth_deltas)))) + + +def summarise_fills(scores): + """Aggregate per-frame fill scores, overall and by gap length.""" + if not scores: + return {} + by_gap = {} + for score in scores: + by_gap.setdefault(score.gap_frames, []).append(score) + + def mean(values): + finite = [v for v in values if np.isfinite(v)] + return float(np.mean(finite)) if finite else float("nan") + + return { + "n": len(scores), + "rmse": mean([s.rmse for s in scores]), + "log_rmse": mean([s.log_rmse for s in scores]), + "psnr": mean([s.psnr for s in scores]), + "ssim": mean([s.ssim for s in scores]), + "by_gap": { + gap: { + "n": len(group), + "rmse": mean([s.rmse for s in group]), + "psnr": mean([s.psnr for s in group]), + "ssim": mean([s.ssim for s in group]), + } + for gap, group in sorted(by_gap.items()) + }, + } diff --git a/suvi/paths.py b/suvi/paths.py new file mode 100644 index 0000000..622f3a0 --- /dev/null +++ b/suvi/paths.py @@ -0,0 +1,185 @@ +"""Archive layout: filename grammar, slot keys, and root discovery. + +The SUVI archive stores one FITS file per (satellite, wavelength, 4-minute slot): + + Data/goes16/l2/data/suvi-l2-ci171/2024/03/15/ + dr_suvi-l2-ci171_g16_s20240315T000000Z_e20240315T000400Z_v1-0-2.fits + +Historically ``filter_FITS.py`` appended ``_f`` (passed) or ``_e`` (rejected) to the +stem to record its verdict. Those labels now live in the SQLite index instead, but +the parser still recognises them so the migration can ingest and strip them. +""" + +import datetime +import os +import re +from dataclasses import dataclass + +# The six SUVI composite-image passbands, in Angstroms. +WAVELENGTHS = (94, 131, 171, 195, 284, 304) +# GOES-17 is deliberately absent: its SUVI data is excluded by puller_fits.py. +SATELLITES = (16, 18, 19) +# Nominal spacing between consecutive frames, in seconds. +CADENCE = 240 + +# Legacy verdict suffixes written into filenames by the pre-SQLite filter. +LEGACY_LABELS = ("f", "e") + +_TIME_FORMAT = "%Y%m%dT%H%M%S" + +# Anchored so that only an exact, well-formed basename matches. The trailing label +# group repeats to absorb the double-suffix files ("..._f_f.fits") that earlier +# filter runs produced by re-processing an already-labelled file. +_FRAME_RE = re.compile( + r"^dr_suvi-l2-ci(?P\d{3})" + r"_g(?P\d{2})" + r"_s(?P\d{8}T\d{6})Z" + r"_e(?P\d{8}T\d{6})Z" + r"_v(?P[0-9]+(?:-[0-9]+)*)" + r"(?P(?:_[fe])*)" + r"\.fits$" +) + + +def _parse_time(text): + """Parse a filename timestamp token to a UTC Unix timestamp.""" + stamp = datetime.datetime.strptime(text, _TIME_FORMAT) + return int(stamp.replace(tzinfo=datetime.timezone.utc).timestamp()) + + +def _format_time(timestamp): + stamp = datetime.datetime.fromtimestamp(timestamp, datetime.timezone.utc) + return stamp.strftime(_TIME_FORMAT) + + +@dataclass(frozen=True) +class FrameName: + """A parsed SUVI L2 composite-image filename.""" + + satellite: int + wavelength: int + t_start: int + t_end: int + version: str + #: Legacy verdict suffix, or None once the archive has been un-renamed. A file + #: labelled more than once keeps only the outermost (most recent) verdict. + label: str | None = None + + @property + def slot(self): + """The key identifying what this frame is an observation of.""" + return (self.satellite, self.wavelength, self.t_start) + + def filename(self, label=None): + """Render the filename, optionally with a legacy verdict suffix.""" + if label is not None and label not in LEGACY_LABELS: + raise ValueError(f"Unknown label: {label!r}") + suffix = f"_{label}" if label else "" + return ( + f"dr_suvi-l2-ci{self.wavelength:03d}" + f"_g{self.satellite:d}" + f"_s{_format_time(self.t_start)}Z" + f"_e{_format_time(self.t_end)}Z" + f"_v{self.version}{suffix}.fits" + ) + + def relpath(self, label=None): + """Archive-relative POSIX path, the natural key used by the index. + + Files are filed under the date the observation window *opened*, so a frame + starting at 23:58 belongs to that day even though it ends on the next one. + """ + day = datetime.datetime.fromtimestamp(self.t_start, datetime.timezone.utc) + return ( + f"goes{self.satellite:d}/l2/data/suvi-l2-ci{self.wavelength:03d}/" + f"{day.year:04d}/{day.month:02d}/{day.day:02d}/{self.filename(label)}" + ) + + def error_plot_name(self): + """The diagnostic JPG the legacy filter wrote next to a rejected frame.""" + return self.filename()[: -len(".fits")] + "_e.jpg" + + +def parse_frame_filename(name): + """Parse a SUVI frame basename, or return None if it is not one. + + Accepts a bare basename only. Anything carrying a path separator is rejected + rather than silently normalised, so a crafted name cannot escape the archive + root when the result is fed back into ``relpath``. + """ + if not isinstance(name, str) or not name or len(name) > 255: + return None + if "/" in name or "\\" in name or os.sep in name: + return None + + match = _FRAME_RE.match(name) + if match is None: + return None + + wavelength = int(match.group("wavelength")) + satellite = int(match.group("satellite")) + if wavelength not in WAVELENGTHS or satellite not in SATELLITES: + return None + + try: + t_start = _parse_time(match.group("t_start")) + t_end = _parse_time(match.group("t_end")) + except ValueError: + # Syntactically well-formed but not a real date, e.g. month 13. + return None + if t_end <= t_start: + return None + + labels = match.group("labels") + label = labels[-1] if labels else None + + return FrameName( + satellite=satellite, + wavelength=wavelength, + t_start=t_start, + t_end=t_end, + version=match.group("version"), + label=label, + ) + + +def data_root(): + """Root of the FITS archive (the ``Data`` directory).""" + override = os.environ.get("SUVI_DATA_ROOT") + if override: + return os.path.abspath(override) + return os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "Data")) + + +def default_db_path(): + """Default location of the SQLite index. + + Overridable so the bench can run against a throwaway database without ever + touching the production index. + """ + override = os.environ.get("SUVI_DB") + if override: + return os.path.abspath(override) + return os.path.join(os.path.dirname(data_root()), "suvi_index.sqlite") + + +def abspath(relpath, root=None): + """Resolve an archive-relative POSIX path against the archive root. + + Raises ValueError if the result would fall outside the root, which keeps a + malformed or hostile index row from reaching an arbitrary part of the disk. + """ + root = os.path.abspath(root or data_root()) + resolved = os.path.abspath(os.path.join(root, *relpath.split("/"))) + if resolved != root and not resolved.startswith(root + os.sep): + raise ValueError(f"Path escapes archive root: {relpath!r}") + return resolved + + +def wavelength_dirs(root, satellites=SATELLITES, wavelengths=WAVELENGTHS): + """Every ``suvi-l2-ciNNN`` directory that could hold frames, existing or not.""" + return [ + os.path.join(root, f"goes{sat:d}", "l2", "data", f"suvi-l2-ci{wl:03d}") + for sat in satellites + for wl in wavelengths + ] diff --git a/suvi/vfs.py b/suvi/vfs.py new file mode 100644 index 0000000..61efd2c --- /dev/null +++ b/suvi/vfs.py @@ -0,0 +1,187 @@ +"""Working with the archive's virtiofs mount without exhausting it. + +The archive sits on a virtiofs share whose daemon holds a host file descriptor for +every inode the guest has looked up. Touching a few hundred thousand files fills the +guest's dentry cache, the daemon runs out of descriptors, and the mount then returns +ENFILE ("too many open files in system") for *any* subsequent open -- reads, writes, +even starting a Python interpreter stored on it. It does not recover on its own. + +The durable fix is host-side (``virtiofsd --inode-file-handles=prefer``, which stores +compact handles instead of descriptors). Until that is in place, anything that +enumerates a large part of the archive has to periodically hand the descriptors back, +which means persuading the guest kernel to evict dentries so it sends FORGET. + +Note that ``vm.vfs_cache_pressure`` does **not** do this. It biases which caches the +kernel drops once it has decided to reclaim; it does not cause reclaim. On a machine +with free memory the kernel never feels pressure, so the dentry cache grows unbounded +whatever that setting says. Reclaim has to be asked for explicitly. +""" + +import atexit + +#: Operations between reclaims during a long traversal. Low enough that the mount +#: never approaches its ceiling, high enough that the cost is amortised. +RELIEF_INTERVAL = 12_000 + + +def drop_caches(): + """Ask the kernel to free dentries and inodes. True if it worked. + + Needs root, and is worth having it: this is direct and instant, where the + fallback has to allocate tens of gigabytes to provoke the same reclaim. Run + long traversals under sudo and they will take this path. + """ + try: + with open("/proc/sys/vm/drop_caches", "w") as handle: + handle.write("2\n") # 2 = dentries and inodes; page cache is not the issue + return True + except OSError: + return False + + +def reclaimable_kb(): + """Size of the reclaimable slab, which is where dentries and inodes live.""" + try: + with open("/proc/meminfo") as handle: + for line in handle: + if line.startswith("SReclaimable:"): + return int(line.split()[1]) + except (OSError, ValueError, IndexError): + pass + return 0 + + +def available_gib(): + """Memory the kernel thinks can be handed out without swapping.""" + try: + with open("/proc/meminfo") as handle: + for line in handle: + if line.startswith("MemAvailable:"): + return int(line.split()[1]) / (1024 * 1024) + except (OSError, ValueError, IndexError): + pass + return 0.0 + + +#: Memory left unclaimed by the fallback, in GiB. Pressure has to be genuine to +#: make the kernel reclaim slab at all, so this is a *reserve* rather than a +#: fraction: capping at some proportion of available memory always leaves headroom, +#: the kernel never feels squeezed, and the allocation frees nothing while still +#: costing the machine several gigabytes -- worse than not trying. +RELIEF_RESERVE_GIB = 2.0 +#: Absolute cap, so a machine with vast memory does not get an unbounded allocation. +RELIEF_MAX_GIB = 48 + + +#: Below this much reclaimable slab there is nothing worth freeing -- it is the +#: working set of live processes, not archive inodes. Pushing memory down to the +#: reserve to shave a few hundred megabytes off it costs the machine far more than +#: it gains, and during a long job this runs every RELIEF_INTERVAL items. +RELIEF_FLOOR_KB = 1024 * 1024 # 1 GiB + + +def release_handles(budget_gib=None, reserve_gib=RELIEF_RESERVE_GIB, + floor_kb=RELIEF_FLOOR_KB): + """Force the kernel to shrink its dentry/inode cache. True if anything was freed. + + Prefers ``drop_caches``; falls back to allocating memory until the reclaimable + slab actually shrinks, or until only `reserve_gib` remains available. + + Two things make the fallback awkward, both learned the hard way: + + * A *fixed* budget silently does nothing once the machine has free memory -- + 10 GiB against 17 GiB free reclaims zero while reporting success. + * A budget capped at a *fraction* of available memory has the same failure for + the same reason: it never applies real pressure. 14 GiB of a 23 GiB + allowance freed nothing at all. + + So it pushes until the kernel actually gives ground, stopping at a fixed reserve + rather than a proportion, and reports honestly whether the slab moved. + """ + if drop_caches(): + return True + + before = reclaimable_kb() + if before <= 0: + return False + if before < floor_kb: + return True # already small; nothing here is worth the memory churn + + target = before // 2 + if budget_gib is None: + budget_gib = RELIEF_MAX_GIB + blocks = [] + try: + for _ in range(int(budget_gib)): + if reclaimable_kb() <= target: + return True + if available_gib() <= reserve_gib: + break # as much pressure as is safe to apply + blocks.append(bytearray(1024 * 1024 * 1024)) + except MemoryError: + pass + finally: + blocks.clear() + return reclaimable_kb() <= target + + +class Reliever: + """Hands file handles back periodically during a bulk traversal. + + Every long pass over the archive -- indexing, detecting, filling, rendering -- + accumulates dentries that pin handles in the host's virtiofs daemon, and the + mount eventually refuses *every* open, taking unrelated software on the machine + down with it. Reclaim was originally wired only into the migration and the + index build; the bench jobs, which run for hours over tens of thousands of + frames, had none, and duly exhausted the mount. + + Bulk loops construct one of these and call :meth:`tick` per item. + """ + + def __init__(self, interval=RELIEF_INTERVAL, enabled=True, label="", on_exit=True): + self.interval = interval + self.enabled = enabled + self.label = label + self.on_exit = on_exit + self.seen = 0 + self.releases = 0 + self._next = interval + self._registered = False + + def tick(self, count=1): + """Record `count` items processed, reclaiming if enough have gone by.""" + self.seen += count + if self.enabled and self.on_exit and not self._registered: + # Reclaiming only *during* a run leaves the machine loaded once it ends: + # the inodes a traversal cached stay pinned, and with them the host's + # file handles, until something else forces reclaim. Every bulk job has + # to hand them back when it finishes, not merely while it runs. + atexit.register(self.finish) + self._registered = True + if not self.enabled or self.seen < self._next: + return False + self._next = self.seen + self.interval + self._release("after {} items".format(self.seen)) + return True + + def finish(self): + """Reclaim once at the end of a run. Idempotent; safe to call twice.""" + if not self.enabled or self.seen == 0: + return False + self.enabled = False # nothing more to do for this traversal + self._release("on finish, {} items".format(self.seen)) + return True + + def _release(self, why): + freed = release_handles() + self.releases += 1 + if self.label: + print(f" [{self.label}] reclaimed handles {why} " + f"({'ok' if freed else 'no change'})", flush=True) + + def __enter__(self): + return self + + def __exit__(self, *exc): + self.finish() + return False diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..ea95cff --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,115 @@ +"""Shared fixtures. + +Tests build their own FITS files rather than reading the archive, so the suite runs +anywhere and cannot be perturbed by (or perturb) the real data. +""" + +import os +import sys + +import numpy as np +import pytest +from astropy.io import fits + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from suvi import paths # noqa: E402 + + +def solar_disc(size=1280, radius=386, centre=None, peak=1.0, background=0.01, + active_region=True): + """A synthetic SUVI-like frame: a bright limb-darkened disc on a dim corona. + + An off-centre active region is included by default, because a perfectly + symmetric disc is invariant under rotation and reflection and would silently + pass tests for corruptions that flip or rotate the frame. Real SUVI frames are + never symmetric, and that asymmetry is what displaces the brightness centroid. + """ + centre = centre if centre is not None else ((size - 1) / 2.0, (size - 1) / 2.0) + yy, xx = np.mgrid[0:size, 0:size].astype(np.float32) + r = np.hypot(xx - centre[0], yy - centre[1]) + disc = np.zeros((size, size), dtype=np.float32) + inside = r < radius + # Limb darkening keeps the radial profile non-trivial so geometry checks bite. + disc[inside] = peak * (0.4 + 0.6 * np.sqrt(np.clip(1 - (r[inside] / radius) ** 2, 0, 1))) + corona = background * np.exp(-np.clip(r - radius, 0, None) / (radius / 4.0)) + image = disc + corona + if active_region: + spot_r = np.hypot(xx - (centre[0] + radius * 0.35), yy - (centre[1] - radius * 0.3)) + image = image + peak * 0.25 * np.exp(-((spot_r / (radius * 0.16)) ** 2)) * inside + return image.astype(np.float32) + + +def write_fits(path, data, headers=None, compress=True, checksum=True): + """Write a SUVI-shaped FITS file: empty primary + (compressed) image extension.""" + header = fits.Header() + header["DATE-BEG"] = "2024-05-10T00:00:00.000" + header["DATE-OBS"] = "2024-05-10T00:00:21.000" + header["DATE-END"] = "2024-05-10T00:00:31.000" + header["WAVELNTH"] = 171 + header["EXPTIME"] = 1.0 + header["CRPIX1"] = (data.shape[1] + 1) / 2.0 + header["CRPIX2"] = (data.shape[0] + 1) / 2.0 + header["CDELT1"] = 2.5 + header["CDELT2"] = 2.5 + header["CROTA"] = 0.0 + header["DSUN_OBS"] = 148781338180.972 + header["SOLAR_B0"] = -7.170855 + header["DIAM_SUN"] = 771.9772 + header["YAW_FLIP"] = 0 + header["ECLIPSE"] = 0 + header["EMPTY"] = False + header["DEGRADED"] = False + header["NUM_IMGS"] = 2 + header["N_LONG"] = 1 + header["N_SHORT"] = 0 + header["N_SH_FL"] = 1 + finite = data[np.isfinite(data)] + header["IMG_MIN"] = float(finite.min()) if finite.size else 0.0 + header["IMG_MAX"] = float(finite.max()) if finite.size else 0.0 + header["IMG_MEAN"] = float(finite.mean()) if finite.size else 0.0 + header["IMG_SDEV"] = float(finite.std()) if finite.size else 0.0 + header["BUNIT"] = "W m-2 sr-1" + for key, value in (headers or {}).items(): + header[key] = value + + if compress: + hdu = fits.CompImageHDU(data=data, header=header, compression_type="GZIP_2") + else: + hdu = fits.ImageHDU(data=data, header=header) + hdulist = fits.HDUList([fits.PrimaryHDU(), hdu]) + os.makedirs(os.path.dirname(path), exist_ok=True) + hdulist.writeto(path, overwrite=True, checksum=checksum) + return path + + +@pytest.fixture +def archive(tmp_path): + """An empty archive root with SUVI_DATA_ROOT pointed at it.""" + root = tmp_path / "Data" + root.mkdir() + old = os.environ.get("SUVI_DATA_ROOT") + os.environ["SUVI_DATA_ROOT"] = str(root) + yield root + if old is None: + os.environ.pop("SUVI_DATA_ROOT", None) + else: + os.environ["SUVI_DATA_ROOT"] = old + + +@pytest.fixture +def db_path(tmp_path): + return str(tmp_path / "index.sqlite") + + +@pytest.fixture +def frame_name(): + return paths.parse_frame_filename( + "dr_suvi-l2-ci171_g16_s20240510T000000Z_e20240510T000400Z_v1-0-2.fits" + ) + + +@pytest.fixture +def good_frame(tmp_path): + """A healthy synthetic frame on disk.""" + return write_fits(str(tmp_path / "frames" / "good.fits"), solar_disc()) diff --git a/tests/test_cases.py b/tests/test_cases.py new file mode 100644 index 0000000..f47fdb9 --- /dev/null +++ b/tests/test_cases.py @@ -0,0 +1,381 @@ +import os + +import pytest + +from conftest import solar_disc, write_fits +from suvi import cases, corruptions, paths + +T0 = 1715299200 # 2024-05-10 00:00 UTC +SATS = (16, 18) +BANDS = (171, 304) + + +def populate(root, slots, label="f"): + """Write minimal frames for the given slots into a fake archive.""" + written = {} + for satellite, wavelength, when in slots: + name = paths.FrameName(satellite, wavelength, when, when + paths.CADENCE, "1-0-2") + path = os.path.join(str(root), *name.relpath(label).split("/")) + write_fits(path, solar_disc(size=64, radius=20, peak=3.0)) + written[name.slot] = path + return written + + +def grid(count, satellites=SATS, wavelengths=BANDS, start=T0): + return [ + (satellite, wavelength, start + index * paths.CADENCE) + for index in range(count) + for satellite in satellites + for wavelength in wavelengths + ] + + +# ----------------------------------------------------------------------- scanning + + +def test_scan_window_finds_frames_and_labels(archive): + populate(archive, grid(3)) + found = cases.scan_window(str(archive), SATS, BANDS, T0, T0 + 3 * paths.CADENCE) + assert len(found) == 12 + path, label = found[(16, 171, T0)] + assert label == "f" and os.path.exists(path) + + +def test_scan_window_respects_the_time_bounds(archive): + populate(archive, grid(5)) + found = cases.scan_window(str(archive), SATS, BANDS, T0, T0 + 2 * paths.CADENCE) + assert {slot[2] for slot in found} == {T0, T0 + paths.CADENCE} + + +def test_scan_window_ignores_unrelated_files(archive): + populate(archive, grid(1)) + stray = os.path.join(str(archive), "goes16/l2/data/suvi-l2-ci171/2024/05/10/notes.txt") + open(stray, "w").write("not a frame") + found = cases.scan_window(str(archive), SATS, BANDS, T0, T0 + paths.CADENCE) + assert len(found) == 4 + + +def test_scan_window_handles_a_missing_archive(tmp_path): + assert cases.scan_window(str(tmp_path), SATS, BANDS, T0, T0 + 240) == {} + + +def test_scan_window_spans_a_day_boundary(archive): + midnight = 1715299200 - 2 * paths.CADENCE # last slots of the previous day + populate(archive, [(16, 171, midnight + i * paths.CADENCE) for i in range(4)]) + found = cases.scan_window(str(archive), (16,), (171,), midnight, midnight + 4 * 240) + assert len(found) == 4 + + +# --------------------------------------------------------------------------- runs + + +def test_timeline_covers_the_window(): + assert cases.timeline(T0, T0 + 3 * paths.CADENCE) == [ + T0, T0 + paths.CADENCE, T0 + 2 * paths.CADENCE + ] + + +def test_find_runs_reports_the_longest_first(archive): + slots = grid(10) + # Punch a hole at index 4 for one band on one satellite. + slots = [s for s in slots if s != (16, 171, T0 + 4 * paths.CADENCE)] + populate(archive, slots) + found = cases.scan_window(str(archive), SATS, BANDS, T0, T0 + 10 * paths.CADENCE) + runs = cases.find_runs(found, SATS, BANDS, T0, T0 + 10 * paths.CADENCE, minimum=2) + assert runs[0][0] == 5 # indices 5..9 + assert runs[1][0] == 4 # indices 0..3 + + +def test_find_runs_ignores_labels_by_default(archive): + """Labels are stale, so completeness is what select-window reports.""" + populate(archive, grid(6), label="e") # everything marked bad by the old filter + found = cases.scan_window(str(archive), SATS, BANDS, T0, T0 + 6 * paths.CADENCE) + assert cases.find_runs(found, SATS, BANDS, T0, T0 + 6 * paths.CADENCE, minimum=2) + assert not cases.find_runs( + found, SATS, BANDS, T0, T0 + 6 * paths.CADENCE, minimum=2, require_good_label=True + ) + + +def test_slot_is_good_treats_unlabelled_as_unknown(archive): + populate(archive, grid(1), label=None) + found = cases.scan_window(str(archive), SATS, BANDS, T0, T0 + paths.CADENCE) + assert cases.slot_is_good(found, 16, 171, T0) is None + assert cases.slot_is_good(found, 16, 171, T0 + 9999) is None + + +# --------------------------------------------------------------------------- plan + + +def test_plan_validates_its_inputs(): + cases.InjectionPlan().validate() + with pytest.raises(ValueError, match="fraction"): + cases.InjectionPlan(fraction=1.5).validate() + with pytest.raises(ValueError, match="gap_lengths"): + cases.InjectionPlan(gap_lengths=(0,)).validate() + with pytest.raises(ValueError, match="satellite_scope"): + cases.InjectionPlan(satellite_scope="g99").validate() + with pytest.raises(ValueError, match="wavelength_scope"): + cases.InjectionPlan(wavelength_scope="some").validate() + with pytest.raises(ValueError, match="unknown corruption"): + cases.InjectionPlan(modes=("nonsense",)).validate() + with pytest.raises(ValueError, match="severity"): + cases.InjectionPlan(severity=(0.9, 0.1)).validate() + + +def test_plan_round_trips_through_a_dict(): + plan = cases.InjectionPlan(fraction=0.2, satellite_scope="both", modes=("delete",)) + assert cases.InjectionPlan.from_dict(plan.as_dict()) == plan + + +def test_injections_are_deterministic(): + slots = set(grid(200)) + first = cases.plan_injections(cases.InjectionPlan(), slots, SATS, BANDS, seed=7) + second = cases.plan_injections(cases.InjectionPlan(), slots, SATS, BANDS, seed=7) + assert first == second + + +def test_a_different_seed_gives_a_different_plan(): + slots = set(grid(200)) + first = cases.plan_injections(cases.InjectionPlan(), slots, SATS, BANDS, seed=1) + second = cases.plan_injections(cases.InjectionPlan(), slots, SATS, BANDS, seed=2) + assert first != second + + +def test_injections_leave_the_window_edges_intact(): + """Every damaged slot needs good frames either side to be reconstructed from.""" + slots = set(grid(200)) + times = sorted({slot[2] for slot in slots}) + injections = cases.plan_injections(cases.InjectionPlan(), slots, SATS, BANDS, seed=3) + damaged = {i.slot[2] for i in injections} + assert damaged.isdisjoint(times[: cases.EDGE_MARGIN]) + assert damaged.isdisjoint(times[-cases.EDGE_MARGIN :]) + + +def test_gaps_stay_separated(): + slots = set(grid(300)) + times = sorted({slot[2] for slot in slots}) + index_of = {when: i for i, when in enumerate(times)} + injections = cases.plan_injections(cases.InjectionPlan(), slots, SATS, BANDS, seed=5) + damaged = sorted(index_of[i.slot[2]] for i in {(j.slot[2],): j for j in injections}.values()) + runs = [] + for position in damaged: + if runs and position == runs[-1][-1] + 1: + runs[-1].append(position) + else: + runs.append([position]) + for earlier, later in zip(runs, runs[1:]): + assert later[0] - earlier[-1] > cases.MIN_SEPARATION + + +def test_gap_metadata_is_recorded(): + slots = set(grid(200)) + injections = cases.plan_injections( + cases.InjectionPlan(gap_lengths=(5,)), slots, SATS, BANDS, seed=9 + ) + assert injections + for injection in injections: + assert injection.gap_length == 5 + assert 0 <= injection.gap_index < 5 + assert injection.mode == "delete" or injection.mode in corruptions.CATALOG + + +def test_satellite_scope_both_damages_each_satellite_together(): + slots = set(grid(200)) + injections = cases.plan_injections( + cases.InjectionPlan(satellite_scope="both"), slots, SATS, BANDS, seed=11 + ) + by_time = {} + for injection in injections: + by_time.setdefault(injection.slot[2], set()).add(injection.slot[0]) + assert all(sats == set(SATS) for sats in by_time.values()) + + +def test_satellite_scope_can_target_one_spacecraft(): + slots = set(grid(200)) + injections = cases.plan_injections( + cases.InjectionPlan(satellite_scope="g18"), slots, SATS, BANDS, seed=11 + ) + assert {i.slot[0] for i in injections} == {18} + + +def test_wavelength_scope_one_damages_a_single_band(): + slots = set(grid(200)) + injections = cases.plan_injections( + cases.InjectionPlan(wavelength_scope="one"), slots, SATS, BANDS, seed=13 + ) + by_time = {} + for injection in injections: + by_time.setdefault(injection.slot[2], set()).add(injection.slot[1]) + assert all(len(bands) == 1 for bands in by_time.values()) + + +def test_zero_fraction_plans_nothing(): + slots = set(grid(200)) + assert cases.plan_injections( + cases.InjectionPlan(fraction=0.0), slots, SATS, BANDS, seed=1 + ) == [] + + +def test_a_window_too_short_to_damage(): + slots = set(grid(4)) + assert cases.plan_injections(cases.InjectionPlan(), slots, SATS, BANDS, seed=1) == [] + + +# ------------------------------------------------------------------------ overlay + + +def build_overlay(): + archive = {(16, 171, T0 + i * 240): (f"/archive/{i}.fits", "f") for i in range(5)} + return cases.Overlay( + archive=archive, + overrides={(16, 171, T0 + 240): "/overlay/1.fits"}, + deleted=frozenset({(16, 171, T0 + 480)}), + ) + + +def test_overlay_passes_untouched_slots_through(): + assert build_overlay().path((16, 171, T0)) == "/archive/0.fits" + + +def test_overlay_redirects_corrupted_slots(): + assert build_overlay().path((16, 171, T0 + 240)) == "/overlay/1.fits" + + +def test_overlay_hides_deleted_slots(): + assert build_overlay().path((16, 171, T0 + 480)) is None + + +def test_overlay_always_exposes_the_original_for_scoring(): + overlay = build_overlay() + assert overlay.truth_path((16, 171, T0 + 240)) == "/archive/1.fits" + assert overlay.truth_path((16, 171, T0 + 480)) == "/archive/2.fits" + + +def test_overlay_reports_an_unknown_slot_as_absent(): + assert build_overlay().path((99, 999, 0)) is None + + +def test_overlay_series_is_ordered_and_filtered(): + overlay = cases.Overlay( + archive={ + (16, 171, T0 + 480): ("/c", None), + (16, 171, T0): ("/a", None), + (18, 171, T0 + 240): ("/b", None), + } + ) + assert overlay.series(16, 171) == [(16, 171, T0), (16, 171, T0 + 480)] + assert overlay.series(18, 171) == [(18, 171, T0 + 240)] + + +# ------------------------------------------------------------------ long gaps + + +def test_separation_scales_with_gap_length(): + """Four clean slots is enough beside a 2-frame gap and useless beside a 300.""" + assert cases.gap_separation(1) == cases.MIN_SEPARATION + assert cases.gap_separation(6) == cases.MIN_SEPARATION + assert cases.gap_separation(300) == 75 + assert cases.gap_separation(100) == 25 + + +def test_long_gaps_keep_clean_brackets(): + slots = set(grid(1200)) + times = sorted({s[2] for s in slots}) + index_of = {t: i for i, t in enumerate(times)} + injections = cases.plan_injections( + cases.InjectionPlan(fraction=0.3, gap_lengths=(3, 30, 100)), + slots, SATS, BANDS, seed=5, + ) + damaged = sorted({index_of[i.slot[2]] for i in injections}) + + runs = [] + for position in damaged: + if runs and position == runs[-1][-1] + 1: + runs[-1].append(position) + else: + runs.append([position]) + for earlier, later in zip(runs, runs[1:]): + needed = cases.gap_separation(len(earlier)) + assert later[0] - earlier[-1] > needed, ( + f"a {len(earlier)}-slot gap had only {later[0] - earlier[-1]} clean " + f"slots before the next, needs more than {needed}" + ) + + +def test_a_window_too_short_for_the_gaps_is_refused(): + """Silently planning nothing would report results for lengths never tested.""" + slots = set(grid(50)) + with pytest.raises(ValueError, match="usable slots"): + cases.plan_injections( + cases.InjectionPlan(gap_lengths=(300,)), slots, SATS, BANDS, seed=1 + ) + + +def test_long_gaps_actually_get_placed(): + """Longest-first placement: a 300 laid down last would rarely find room.""" + slots = set(grid(2600)) + injections = cases.plan_injections( + cases.InjectionPlan(fraction=0.25, gap_lengths=(1, 2, 3, 10, 30, 100, 300)), + slots, SATS, BANDS, seed=7, + ) + placed = {i.gap_length for i in injections} + assert 300 in placed, "the longest gap never got placed" + assert placed & {1, 2, 3}, "short gaps were crowded out" + + +def test_gap_lengths_span_short_and_long(): + slots = set(grid(2600)) + injections = cases.plan_injections( + cases.InjectionPlan(fraction=0.25, gap_lengths=(1, 2, 3, 10, 30, 100, 300)), + slots, SATS, BANDS, seed=3, + ) + lengths = {i.gap_length for i in injections} + assert min(lengths) <= 3 and max(lengths) >= 100 + + +def test_gaps_per_length_places_the_requested_count(): + """A fraction budget is eaten by the longest gaps, leaving one of each. + + An earlier case ended up 84% a single 300-slot gap of a single mode, which + measures that mode at that length and nothing else. + """ + slots = set(grid(2600)) + injections = cases.plan_injections( + cases.InjectionPlan(gap_lengths=(1, 3, 10), gaps_per_length=4), + slots, SATS, BANDS, seed=4, + ) + counts = {} + for injection in injections: + key = (injection.gap_length, injection.slot[2] - injection.gap_index * paths.CADENCE) + counts.setdefault(injection.gap_length, set()).add(key[1]) + assert counts[1] == counts[1] and len(counts[1]) == 4 + assert len(counts[3]) == 4 + assert len(counts[10]) == 4 + + +def test_gaps_per_length_overrides_the_fraction(): + slots = set(grid(2600)) + injections = cases.plan_injections( + cases.InjectionPlan(fraction=0.001, gap_lengths=(10,), gaps_per_length=3), + slots, SATS, BANDS, seed=6, + ) + starts = {i.slot[2] - i.gap_index * paths.CADENCE for i in injections} + assert len(starts) == 3, "the tiny fraction should have been ignored" + + +def test_modes_are_cycled_not_resampled(): + """With few gaps, independent draws repeat and leave the catalogue untested.""" + slots = set(grid(2600)) + injections = cases.plan_injections( + cases.InjectionPlan(gap_lengths=(2,), gaps_per_length=8), slots, SATS, BANDS, seed=8 + ) + starts = {} + for injection in injections: + starts.setdefault(injection.slot[2] - injection.gap_index * paths.CADENCE, + injection.mode) + assert len(set(starts.values())) == len(starts), "a mode was reused across gaps" + + +def test_plan_with_gaps_per_length_round_trips(): + plan = cases.InjectionPlan(gaps_per_length=5) + assert cases.InjectionPlan.from_dict(plan.as_dict()).gaps_per_length == 5 diff --git a/tests/test_comparison_video.py b/tests/test_comparison_video.py new file mode 100644 index 0000000..99aa8ce --- /dev/null +++ b/tests/test_comparison_video.py @@ -0,0 +1,165 @@ +"""Tests for assembling the three-pane comparison video. + +The property that matters is *alignment*: a variant that produced no composite for a +timestamp -- which is exactly what the current pipeline does past its gap limit -- +must still occupy that slot, or its pane slides out of step with the others and the +comparison is meaningless. +""" + +import os + +import numpy as np +from PIL import Image + +import make_comparison_video as mcv +from suvi import paths + +T0 = 1715299200 + + +def write_stream(root, variant, satellite, timestamps, size=(64, 36)): + directory = os.path.join(root, variant, f"goes{satellite}") + os.makedirs(directory, exist_ok=True) + for when in timestamps: + img = Image.fromarray(np.full((size[1], size[0], 3), 40, dtype="uint8")) + img.save(os.path.join(directory, f"Composite-{when}.jpg"), quality=90) + return directory + + +def test_stream_frames_reads_timestamps(tmp_path): + times = [T0, T0 + paths.CADENCE, T0 + 2 * paths.CADENCE] + directory = write_stream(str(tmp_path), "pristine", 16, times) + frames = mcv.stream_frames(directory) + assert sorted(frames) == times + assert all(os.path.exists(p) for p in frames.values()) + + +def test_stream_frames_ignores_other_files(tmp_path): + directory = write_stream(str(tmp_path), "new", 16, [T0]) + open(os.path.join(directory, "notes.txt"), "w").write("x") + open(os.path.join(directory, "Composite-nonsense.jpg"), "w").write("x") + assert list(mcv.stream_frames(directory)) == [T0] + + +def test_stream_frames_on_a_missing_directory(tmp_path): + assert mcv.stream_frames(str(tmp_path / "absent")) == {} + + +def test_every_variant_is_labelled(): + assert set(mcv.LABELS) == set(mcv.VARIANTS) + assert all(mcv.LABELS[v] for v in mcv.VARIANTS) + + +def test_encode_emits_one_frame_per_slot_including_gaps(tmp_path, monkeypatch): + """A pane with holes must still be full length, or the panes drift apart. + + 'today' legitimately has none for long gaps; those slots have to become black + frames occupying the timeline, not vanish from it. + """ + times = [T0, T0 + 2 * paths.CADENCE] # slot 1 deliberately missing + directory = write_stream(str(tmp_path), "today", 16, times) + frames = mcv.stream_frames(directory) + timeline = [T0 + i * paths.CADENCE for i in range(3)] + + written = [] + + class FakeProcess: + def __init__(self): + self.stdin = self + + def write(self, data): + written.append(("real", len(data))) + + def save(self, *a, **k): + pass + + def close(self): + pass + + def wait(self): + return 0 + + fake = FakeProcess() + monkeypatch.setattr(mcv.subprocess, "Popen", lambda *a, **k: fake) + + class CountingImage: + size = (64, 36) + + @staticmethod + def fromarray(arr): + class Blank: + def save(self, handle, *a, **k): + written.append(("black", 0)) + return Blank() + + @staticmethod + def open(path): + return CountingImage + + monkeypatch.setattr(mcv, "Image", CountingImage) + + present = mcv.encode_stream("ffmpeg", frames, timeline, "out.mp4", 60, 18) + assert present == 2 + assert len(written) == 3, "a missing slot did not occupy the timeline" + assert [kind for kind, _ in written] == ["real", "black", "real"] + + +# ------------------------------------------------------ pristine reads the truth + + +def test_pristine_resolves_deleted_slots_from_the_archive(tmp_path): + """The ground-truth pane must show the window as it really is. + + A slot the case deleted has no overlay path. Resolving pristine through the + overlay silently drops those timestamps -- 221 of them in one run -- leaving the + reference pane shorter than the panes it exists to be compared against. + """ + import bench + from suvi import cases + + slot = (16, 171, T0) + overlay = cases.Overlay( + archive={slot: ("/archive/real.fits", None)}, + overrides={}, + deleted=frozenset({slot}), + ) + assert overlay.path(slot) is None, "precondition: the case deleted this slot" + assert overlay.truth_path(slot) == "/archive/real.fits" + + reads = [] + + def fake_read(path): + reads.append(path) + return "pixels", None + + original = bench.fitsio.read_image + bench.fitsio.read_image = fake_read + try: + cache = {} + assert bench._resolve_band(overlay, slot, cache, use_truth=True) == "pixels" + assert reads == ["/archive/real.fits"] + # Without use_truth the slot is genuinely gone, which is right for the + # variants that are meant to show the damage. + assert bench._resolve_band(overlay, slot, cache, use_truth=False) is None + finally: + bench.fitsio.read_image = original + + +def test_resolve_band_caches_truth_and_overlay_separately(tmp_path): + """One cache serving both would return the wrong frame for one of them.""" + import bench + from suvi import cases + + slot = (16, 171, T0) + overlay = cases.Overlay( + archive={slot: ("/archive/real.fits", None)}, + overrides={slot: "/overlay/damaged.fits"}, + ) + original = bench.fitsio.read_image + bench.fitsio.read_image = lambda path: (path, None) + try: + cache = {} + assert bench._resolve_band(overlay, slot, cache, use_truth=True) == "/archive/real.fits" + assert bench._resolve_band(overlay, slot, cache, use_truth=False) == "/overlay/damaged.fits" + finally: + bench.fitsio.read_image = original diff --git a/tests/test_composite.py b/tests/test_composite.py new file mode 100644 index 0000000..71631fb --- /dev/null +++ b/tests/test_composite.py @@ -0,0 +1,90 @@ +"""Tests for the composite renderer. + +The blending in ``merger_FITS`` is the product's visual identity -- years of tuning +that no test can meaningfully assert the *correctness* of. What these tests protect +is that it does not change: it was lifted out of a worker loop so the bench could +render reconstructed frames, and the only thing that makes such a refactor safe is +proof the pixels came out the same. +""" + +import numpy as np +import pytest + +import merger_FITS as M +from conftest import solar_disc + + +@pytest.fixture +def six_bands(): + """Six distinct 1280x1280 bands, bright enough to exercise every blend stage.""" + return [ + solar_disc(peak=peak, radius=386, active_region=True) + for peak in (0.06, 0.12, 1.2, 1.6, 1.2, 3.0) + ] + + +def test_renders_a_composite_of_the_expected_shape(six_bands): + img = M.composite_from_arrays(six_bands, 1715212800) + # 1280 trimmed by 2*64 in x and 2*100 in y, then padded by a third either side. + assert img.size == (1920, 1080) + assert img.mode == "RGB" + + +def test_rendering_is_deterministic(six_bands): + first = M.composite_from_arrays(six_bands, 1715212800).tobytes() + second = M.composite_from_arrays(six_bands, 1715212800).tobytes() + assert first == second + + +def test_the_timestamp_is_drawn_on_the_image(six_bands): + """Two timestamps must differ only where the caption is.""" + a = np.asarray(M.composite_from_arrays(six_bands, 1715212800)) + b = np.asarray(M.composite_from_arrays(six_bands, 1715299200)) + assert not np.array_equal(a, b) + # The caption sits in the top strip; everything below it is identical. + assert np.array_equal(a[60:], b[60:]) + + +def test_every_band_contributes(six_bands): + """Changing any one band must change the output. + + Guards against an indexing slip in the extracted function silently dropping a + band -- the composite would still render, and still look plausible. + """ + reference = np.asarray(M.composite_from_arrays(six_bands, 1715212800)) + for index in range(6): + altered = list(six_bands) + altered[index] = altered[index] * 1.5 + assert not np.array_equal( + np.asarray(M.composite_from_arrays(altered, 1715212800)), reference + ), f"band {index} made no difference to the composite" + + +def test_band_order_matters(six_bands): + """The six arrays are positional; a caller passing them shuffled must not match.""" + shuffled = [six_bands[i] for i in (5, 4, 3, 2, 1, 0)] + assert not np.array_equal( + np.asarray(M.composite_from_arrays(shuffled, 1715212800)), + np.asarray(M.composite_from_arrays(six_bands, 1715212800)), + ) + + +def test_render_assets_are_built_once(): + """The first four colormaps are mutated in place, so building twice would + re-apply the mutation to already-mutated objects.""" + first = M._render_assets() + second = M._render_assets() + assert first is second + assert first["cmaps"][0] is second["cmaps"][0] + + +def test_tolerates_a_filled_frame_containing_no_signal(six_bands): + """Reconstructed frames can be blank; rendering must not raise on them.""" + blanked = list(six_bands) + blanked[2] = np.zeros_like(blanked[2]) + img = M.composite_from_arrays(blanked, 1715212800) + assert img.size == (1920, 1080) + + +def test_mapping_constants_line_up_with_the_bands(): + assert len(M.VMINS) == len(M.VMAXS) == len(M.GAMMAS) == len(M.image_names) == 6 diff --git a/tests/test_corruptions.py b/tests/test_corruptions.py new file mode 100644 index 0000000..30a6e41 --- /dev/null +++ b/tests/test_corruptions.py @@ -0,0 +1,209 @@ +import numpy as np +import pytest + +from conftest import solar_disc, write_fits +from suvi import corruptions, 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_four_groups(): + assert set(corruptions.GROUPS) == {"dropout", "structural", "radiometric", "geometric"} + 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", ARRAY_MODES) +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)) + + +@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 diff --git a/tests/test_db.py b/tests/test_db.py new file mode 100644 index 0000000..ab88d75 --- /dev/null +++ b/tests/test_db.py @@ -0,0 +1,285 @@ +import multiprocessing +import sqlite3 + +import pytest + +from suvi import db, paths + + +def make_frame(conn, satellite=16, wavelength=171, minute=0): + t_start = 1715299200 + minute * paths.CADENCE + name = paths.FrameName( + satellite=satellite, + wavelength=wavelength, + t_start=t_start, + t_end=t_start + paths.CADENCE, + version="1-0-2", + ) + return db.upsert_frame(conn, name, name.relpath(), size_bytes=1774080, mtime=1.0), name + + +def test_connect_creates_schema(db_path): + conn = db.connect(db_path) + tables = { + row["name"] + for row in conn.execute("SELECT name FROM sqlite_master WHERE type='table'") + } + assert {"frame", "header", "detector_run", "detection", "remote_file", "meta"} <= tables + assert conn.execute( + "SELECT value FROM meta WHERE key='schema_version'" + ).fetchone()["value"] == str(db.SCHEMA_VERSION) + conn.close() + + +def test_init_schema_is_idempotent(db_path): + conn = db.connect(db_path) + db.init_schema(conn) + db.init_schema(conn) + assert conn.execute("SELECT count(*) c FROM meta").fetchone()["c"] == 1 + conn.close() + + +def test_schema_version_mismatch_is_refused(db_path): + conn = db.connect(db_path) + conn.execute("UPDATE meta SET value='99' WHERE key='schema_version'") + conn.commit() + with pytest.raises(RuntimeError, match="schema version"): + db.init_schema(conn) + conn.close() + + +def test_readonly_connect_requires_an_existing_file(tmp_path): + with pytest.raises(FileNotFoundError): + db.connect(str(tmp_path / "absent.sqlite"), readonly=True) + + +def test_readonly_connect_cannot_write(db_path): + db.connect(db_path).close() + conn = db.connect(db_path, readonly=True) + with pytest.raises(sqlite3.OperationalError): + conn.execute("INSERT INTO meta (key, value) VALUES ('x', 'y')") + conn.close() + + +def test_upsert_frame_is_stable_and_updates_in_place(db_path): + conn = db.connect(db_path) + first, name = make_frame(conn) + again = db.upsert_frame(conn, name, name.relpath(), size_bytes=999, mtime=2.0) + assert first == again + row = conn.execute("SELECT * FROM frame WHERE id=?", (first,)).fetchone() + assert row["size_bytes"] == 999 and row["mtime"] == 2.0 + assert conn.execute("SELECT count(*) c FROM frame").fetchone()["c"] == 1 + conn.close() + + +def test_duplicate_slot_under_a_new_version_is_rejected(db_path): + """Two files claiming the same observation must not silently shadow each other.""" + conn = db.connect(db_path) + _, name = make_frame(conn) + other = paths.FrameName( + name.satellite, name.wavelength, name.t_start, name.t_end, version="1-0-3" + ) + with pytest.raises(sqlite3.IntegrityError): + db.upsert_frame(conn, other, other.relpath()) + conn.close() + + +def test_frame_id_by_slot(db_path): + conn = db.connect(db_path) + frame_id, name = make_frame(conn) + assert db.frame_id_by_slot(conn, *name.slot) == frame_id + assert db.frame_id_by_slot(conn, 18, 171, name.t_start) is None + conn.close() + + +def test_detector_runs_are_ordered_by_recency(db_path): + conn = db.connect(db_path) + first = db.create_detector_run(conn, "header_v1", {"threshold": 1}) + second = db.create_detector_run(conn, "header_v1", {"threshold": 2}) + assert db.latest_run_id(conn, "header_v1") == second + assert first != second + assert db.latest_run_id(conn, "nope") is None + conn.close() + + +def test_record_detections_stores_scores_and_upserts(db_path): + conn = db.connect(db_path) + frame_id, _ = make_frame(conn) + run = db.create_detector_run(conn, "header_v1", {}) + db.record_detections(conn, run, [(frame_id, "bad", "eclipse", {"img_mean": 1e-4}, 12)]) + db.record_detections(conn, run, [(frame_id, "good", None, {"img_mean": 0.3}, 15)]) + conn.commit() + row = conn.execute("SELECT * FROM detection").fetchone() + assert row["verdict"] == "good" + assert row["reason"] is None + assert row["scores_json"] == '{"img_mean": 0.3}' + assert conn.execute("SELECT count(*) c FROM detection").fetchone()["c"] == 1 + conn.close() + + +def test_record_detections_rejects_unknown_verdicts(db_path): + conn = db.connect(db_path) + frame_id, _ = make_frame(conn) + run = db.create_detector_run(conn, "header_v1", {}) + with pytest.raises(ValueError, match="Unknown verdict"): + db.record_detections(conn, run, [(frame_id, "maybe", None, None, 0)]) + conn.close() + + +def test_detections_cascade_when_a_run_is_deleted(db_path): + conn = db.connect(db_path) + frame_id, _ = make_frame(conn) + run = db.create_detector_run(conn, "header_v1", {}) + db.record_detections(conn, run, [(frame_id, "good", None, None, 0)]) + conn.execute("DELETE FROM detector_run WHERE id=?", (run,)) + conn.commit() + assert conn.execute("SELECT count(*) c FROM detection").fetchone()["c"] == 0 + conn.close() + + +def test_good_slots_filters_by_verdict_window_and_band(db_path): + conn = db.connect(db_path) + run = db.create_detector_run(conn, "header_v1", {}) + ids = {} + for minute in range(4): + for wavelength in (171, 304): + frame_id, name = make_frame(conn, wavelength=wavelength, minute=minute) + ids[(minute, wavelength)] = (frame_id, name) + # Everything good except one frame, plus a decoy on the other satellite. + db.record_detections( + conn, + run, + [ + (fid, "bad" if key == (1, 171) else "good", None, None, 0) + for key, (fid, _) in ids.items() + ], + ) + other, _ = make_frame(conn, satellite=18, minute=0) + db.record_detections(conn, run, [(other, "good", None, None, 0)]) + conn.commit() + + base = ids[(0, 171)][1].t_start + rows = db.good_slots(conn, run, base, base + 3 * paths.CADENCE, 16, (171,)) + assert [r["t_start"] for r in rows] == [base, base + 2 * paths.CADENCE] + assert all(r["satellite"] == 16 and r["wavelength"] == 171 for r in rows) + + both = db.good_slots(conn, run, base, base + 4 * paths.CADENCE, 16, (171, 304)) + assert len(both) == 7 # 8 frames minus the one marked bad + conn.close() + + +def test_unscanned_frames_lists_only_frames_without_headers(db_path): + conn = db.connect(db_path) + first, _ = make_frame(conn, minute=0) + second, _ = make_frame(conn, minute=1) + db.record_headers(conn, [(first, {"img_mean": 0.3, "eclipse": 0}, None)]) + conn.commit() + assert [r["id"] for r in db.unscanned_frames(conn)] == [second] + assert len(db.unscanned_frames(conn, limit=0)) == 0 + conn.close() + + +def test_record_headers_round_trips_values_and_errors(db_path): + conn = db.connect(db_path) + first, _ = make_frame(conn, minute=0) + second, _ = make_frame(conn, minute=1) + db.record_headers( + conn, + [ + (first, {"img_mean": 0.31, "degraded": 0, "datasum": "261950686"}, None), + (second, {}, "truncated: 17 bytes"), + ], + ) + conn.commit() + good = conn.execute("SELECT * FROM header WHERE frame_id=?", (first,)).fetchone() + assert good["img_mean"] == pytest.approx(0.31) + assert good["read_ok"] == 1 and good["read_error"] is None + bad = conn.execute("SELECT * FROM header WHERE frame_id=?", (second,)).fetchone() + assert bad["read_ok"] == 0 and bad["read_error"].startswith("truncated") + assert bad["img_mean"] is None + conn.close() + + +def test_record_headers_upserts(db_path): + conn = db.connect(db_path) + frame_id, _ = make_frame(conn) + db.record_headers(conn, [(frame_id, {"img_mean": 0.1}, None)]) + db.record_headers(conn, [(frame_id, {"img_mean": 0.9}, None)]) + conn.commit() + assert conn.execute("SELECT count(*) c FROM header").fetchone()["c"] == 1 + assert conn.execute("SELECT img_mean FROM header").fetchone()["img_mean"] == 0.9 + conn.close() + + +def test_remote_file_round_trip(db_path): + conn = db.connect(db_path) + db.record_remote_files(conn, [("http://x/a.fits", 100.0, 17, "a.fits", 5.0)]) + db.record_remote_files(conn, [("http://x/a.fits", 200.0, 18, "a.fits", 6.0)]) + conn.commit() + assert db.get_remote_mtime(conn, "http://x/a.fits") == 200.0 + assert db.get_remote_mtime(conn, "http://x/missing.fits") is None + assert conn.execute("SELECT count(*) c FROM remote_file").fetchone()["c"] == 1 + conn.close() + + +def test_writer_loop_drains_a_queue_from_another_process(db_path): + """Workers stay read-only; one writer owns the lock. This is that contract.""" + conn = db.connect(db_path) + frame_ids = [make_frame(conn, minute=i)[0] for i in range(20)] + run = db.create_detector_run(conn, "header_v1", {}) + conn.commit() + conn.close() + + ctx = multiprocessing.get_context("spawn") + queue = ctx.Queue() + writer = ctx.Process(target=db.writer_loop, args=(db_path, queue), daemon=True) + writer.start() + for frame_id in frame_ids: + queue.put(("detection", (run, frame_id, "good", None, {"s": 1}, 3))) + queue.put(("header", (frame_ids[0], {"img_mean": 0.5}, None))) + queue.put(("remote_file", ("http://x/a.fits", 1.0, 2, "a", 3.0))) + queue.put(db.STOP) + writer.join(timeout=60) + assert writer.exitcode == 0 + + conn = db.connect(db_path, readonly=True) + assert conn.execute("SELECT count(*) c FROM detection").fetchone()["c"] == 20 + assert conn.execute("SELECT count(*) c FROM header").fetchone()["c"] == 1 + assert conn.execute("SELECT count(*) c FROM remote_file").fetchone()["c"] == 1 + conn.close() + + +def test_writer_loop_rejects_unknown_message_kinds(db_path): + ctx = multiprocessing.get_context("spawn") + queue = ctx.Queue() + writer = ctx.Process(target=db.writer_loop, args=(db_path, queue), daemon=True) + writer.start() + queue.put(("bogus", ())) + writer.join(timeout=60) + assert writer.exitcode not in (0, None) + + +def test_latest_run_is_scoped_by_config(db_path): + """Runs accumulate one per detector per case, so 'latest' is ambiguous. + + Resolving a detector without its config picks whichever case ran most + recently -- and scoring one case's verdicts against another case's injections + gives numbers that are wrong, not merely noisy. + """ + conn = db.connect(db_path) + first = db.create_detector_run(conn, "header_v1", {"case": "alpha"}) + second = db.create_detector_run(conn, "header_v1", {"case": "beta"}) + + assert db.latest_run_id(conn, "header_v1") == second # global latest + assert db.latest_run_id(conn, "header_v1", {"case": "alpha"}) == first + assert db.latest_run_id(conn, "header_v1", {"case": "beta"}) == second + assert db.latest_run_id(conn, "header_v1", {"case": "gamma"}) is None + conn.close() + + +def test_scoped_lookup_takes_the_newest_matching_run(db_path): + conn = db.connect(db_path) + db.create_detector_run(conn, "disc_v1", {"case": "alpha"}) + newer = db.create_detector_run(conn, "disc_v1", {"case": "alpha"}) + assert db.latest_run_id(conn, "disc_v1", {"case": "alpha"}) == newer + conn.close() diff --git a/tests/test_detectors.py b/tests/test_detectors.py new file mode 100644 index 0000000..5ba6a6a --- /dev/null +++ b/tests/test_detectors.py @@ -0,0 +1,476 @@ +import cv2 as cv +import numpy as np +import pytest + +from conftest import solar_disc +from suvi import detectors + +BAND = 171 + + +def features(slot_band=BAND, time=1715299200, satellite=16, **header): + """A FrameFeatures whose header describes a healthy frame, before overrides.""" + base = { + "empty": 0, + "degraded": 0, + "eclipse": 0, + "num_imgs": 2, + "wavelnth": slot_band, + "img_mean": 0.35, + "img_sdev": 0.57, + "diam_sun": 771.98, + "crpix1": 640.5, + "crpix2": 640.5, + } + base.update(header) + return detectors.FrameFeatures(slot=(satellite, slot_band, time), header=base) + + +def bright_disc(**kwargs): + """A disc bright enough to clear the 171A geometry threshold of 1.0.""" + return solar_disc(peak=3.0, **kwargs) + + +# -------------------------------------------------------------------------- header + + +def test_header_accepts_a_healthy_frame(): + assert detectors.header_v1(features()).verdict == detectors.GOOD + + +def test_header_records_scores_even_when_passing(): + verdict = detectors.header_v1(features()) + assert verdict.scores["img_mean"] == pytest.approx(0.35) + assert "crpix_offset" in verdict.scores + + +@pytest.mark.parametrize( + "override,fragment", + [ + ({"empty": 1}, "EMPTY"), + ({"eclipse": 2}, "ECLIPSE"), + ({"num_imgs": 0}, "no source images"), + ({"img_sdev": 0.0}, "zero variance"), + ({"wavelnth": 304}, "wavelength"), + ({"img_mean": 1e-5}, "below"), + ({"img_mean": 500.0}, "above"), + ({"diam_sun": 100.0}, "DIAM_SUN"), + ({"crpix1": 700.0}, "sun centre offset"), + ], +) +def test_header_rejects(override, fragment): + verdict = detectors.header_v1(features(**override)) + assert verdict.verdict == detectors.BAD + assert fragment in verdict.reason + + +def test_header_does_not_reject_on_degraded_alone(): + """Regression: DEGRADED is set for whole months of the 195A band on good frames. + + Rejecting on it discarded 40/40 good 195A frames on 2024-01-20 while flagging + 0/40 on 2024-07-04. It is recorded as a score and nothing more. + """ + verdict = detectors.header_v1(features(degraded=1)) + assert verdict.verdict == detectors.GOOD + assert verdict.scores["degraded"] == 1.0 + + +def test_header_reports_read_errors_as_bad(): + frame = detectors.FrameFeatures(slot=(16, BAND, 0), header={}, error="truncated") + verdict = detectors.header_v1(frame) + assert verdict.verdict == detectors.BAD and "truncated" in verdict.reason + + +def test_header_abstains_without_a_cached_header(): + frame = detectors.FrameFeatures(slot=(16, BAND, 0), header={}) + assert detectors.header_v1(frame).verdict == detectors.UNKNOWN + + +def test_header_bounds_are_per_band(): + """A radiance normal for 304A is a dropout for 94A.""" + assert detectors.header_v1(features(slot_band=304, img_mean=3.0, wavelnth=304)).verdict == detectors.GOOD + assert detectors.header_v1(features(slot_band=94, img_mean=3.0, wavelnth=94)).verdict == detectors.BAD + + +# ------------------------------------------------------------------------ geometry + + +def test_geometry_accepts_a_synthetic_disc(): + verdict = detectors.geometry_v1(bright_disc(), BAND) + assert verdict.verdict == detectors.GOOD, verdict.reason + + +def test_geometry_rejects_a_blank_frame(): + verdict = detectors.geometry_v1(np.zeros((1280, 1280), np.float32), BAND) + assert verdict.verdict == detectors.BAD and "too low" in verdict.reason + + +def test_geometry_rejects_a_saturated_frame(): + verdict = detectors.geometry_v1(np.full((1280, 1280), 99.0, np.float32), BAND) + assert verdict.verdict == detectors.BAD and "too high" in verdict.reason + + +def test_geometry_rejects_a_displaced_disc(): + verdict = detectors.geometry_v1(bright_disc(centre=(500, 640)), BAND) + assert verdict.verdict == detectors.BAD + + +def test_geometry_reports_missing_and_misshapen_input(): + assert detectors.geometry_v1(None, BAND).verdict == detectors.BAD + small = detectors.geometry_v1(np.zeros((64, 64), np.float32), BAND) + assert small.verdict == detectors.BAD and "dimensions" in small.reason + + +def test_geometry_abstains_on_an_unknown_band(): + verdict = detectors.geometry_v1(bright_disc(), 999) + assert verdict.verdict == detectors.UNKNOWN + + +def test_geometry_centre_check_is_one_sided(): + """Documents a real defect in the baseline, so a fix cannot land unnoticed. + + The test is `HALF_DIMS//2 - centre > skew`, so brightness pulled toward low + indices rejects the frame while the same pull toward high indices passes. This + asymmetry is why the filter rejects entire days during high solar activity. + """ + low = detectors.geometry_v1(bright_disc(centre=(560, 640)), BAND) + high = detectors.geometry_v1(bright_disc(centre=(720, 640)), BAND) + assert low.scores["centre_x"] < detectors.HALF_DIMS // 2 + assert high.scores["centre_x"] > detectors.HALF_DIMS // 2 + assert low.verdict == detectors.BAD + assert high.verdict == detectors.GOOD # symmetric displacement, opposite verdict + + +def test_geometry_catches_an_all_nan_frame_only_by_luck(): + """An all-NaN frame is caught by the ratio test, not by the shape tests. + + `NaN > x` is False, so the centroid, radius and goodness-of-fit comparisons all + evaluate False and report "good" on NaN input. Only the emptiness of the + threshold mask saves this case -- a frame that was partly NaN could still slip + through the shape checks. + """ + verdict = detectors.geometry_v1(np.full((1280, 1280), np.nan, np.float32), BAND) + assert verdict.verdict == detectors.BAD + assert "too low" in verdict.reason + + # Demonstrate the underlying blindness: NaN defeats each shape test directly. + assert not (detectors.HALF_DIMS // 2 - np.nan > detectors.MAX_CENTER_SKEW) + assert not (np.nan > detectors.MAX_GOF) + + +# ------------------------------------------------------------------------ temporal + + +def series(count=25, means=None, thumbs=None): + out = [] + for index in range(count): + mean = 0.35 if means is None else means[index] + frame = detectors.FrameFeatures( + slot=(16, BAND, 1715299200 + index * 240), + header={"img_mean": mean}, + ) + if thumbs is not None: + frame.thumbnail = thumbs[index] + out.append(frame) + return out + + +def drifting_thumbs(count, rng=None): + """Thumbnails that change a little each step, like the real Sun.""" + rng = rng or np.random.default_rng(0) + base = solar_disc(size=detectors.THUMBNAIL_SIZE, radius=38, peak=1.0) + return [(base + rng.normal(0, 0.01, base.shape)).astype(np.float32) for _ in range(count)] + + +def test_temporal_accepts_a_steady_series(): + verdicts = detectors.temporal_v1(series(thumbs=drifting_thumbs(25))) + assert all(v.verdict == detectors.GOOD for v in verdicts) + + +def test_temporal_catches_a_frozen_frame(): + thumbs = drifting_thumbs(25) + thumbs[12] = thumbs[11].copy() # the feed stalled + verdicts = detectors.temporal_v1(series(thumbs=thumbs)) + assert verdicts[12].verdict == detectors.BAD + assert "identical to previous" in verdicts[12].reason + + +def test_temporal_catches_a_brightness_step(): + means = [0.35] * 25 + means[12] = 3.5 + verdicts = detectors.temporal_v1(series(means=means, thumbs=drifting_thumbs(25))) + assert verdicts[12].verdict == detectors.BAD + + +def test_temporal_catches_a_structural_jump(): + thumbs = drifting_thumbs(25) + thumbs[12] = np.roll(thumbs[12], 40, axis=0) # a frame from somewhere else + verdicts = detectors.temporal_v1(series(thumbs=thumbs)) + assert verdicts[12].verdict == detectors.BAD + + +def test_temporal_reports_unreadable_frames(): + frames = series(thumbs=drifting_thumbs(25)) + frames[5].error = "truncated" + verdicts = detectors.temporal_v1(frames) + assert verdicts[5].verdict == detectors.BAD + + +def test_temporal_abstains_without_context(): + verdicts = detectors.temporal_v1(series(count=2)) + assert all(v.verdict == detectors.UNKNOWN for v in verdicts) + + +def test_temporal_handles_an_empty_series(): + assert detectors.temporal_v1([]) == [] + + +# ------------------------------------------------------------------------ crosssat + + +def cross_pair(scale_b=1.0, shift_b=0, noise=0.005): + rng = np.random.default_rng(1) + base = solar_disc(size=detectors.THUMBNAIL_SIZE, radius=38, peak=1.0) + a = (base + rng.normal(0, noise, base.shape)).astype(np.float32) + b = (base + rng.normal(0, noise, base.shape)).astype(np.float32) * scale_b + if shift_b: + b = np.roll(b, shift_b, axis=1) + first = detectors.FrameFeatures(slot=(16, BAND, 0), header={"img_mean": 0.35}) + second = detectors.FrameFeatures(slot=(18, BAND, 0), header={"img_mean": 0.35}) + first.thumbnail, second.thumbnail = a, b + return first, second + + +def test_crosssat_agrees_on_matching_views(): + a, b = cross_pair() + first, second = detectors.crosssat_v1(a, b) + assert first.verdict == detectors.GOOD and second.verdict == detectors.GOOD + + +def test_crosssat_tolerates_calibration_differences(): + """A modest scale factor between flight models is normal, not a fault.""" + a, b = cross_pair(scale_b=1.3) + first, _ = detectors.crosssat_v1(a, b) + assert first.verdict == detectors.GOOD + + +def test_crosssat_catches_a_blackout_despite_gain_matching(): + """Regression: least-squares matching rescales a 1e-4 frame into agreement. + + Without an explicit bound on the fitted gain this detector called a total + blackout a match, because the residual after rescaling is tiny. + """ + a, b = cross_pair(scale_b=1e-4) + first, second = detectors.crosssat_v1(a, b) + assert detectors.BAD in (first.verdict, second.verdict) or first.verdict == detectors.UNKNOWN + assert "gain" in (first.reason or "") + + +def test_crosssat_blames_the_frame_its_own_history_disowns(): + a, b = cross_pair(scale_b=1e-4) + suspect = detectors.Verdict(detectors.BAD, "brightness z=40") + healthy = detectors.Verdict(detectors.GOOD) + first, second = detectors.crosssat_v1(a, b, temporal_a=healthy, temporal_b=suspect) + assert first.verdict == detectors.GOOD + assert second.verdict == detectors.BAD + + +def test_crosssat_abstains_when_it_cannot_tell_which_is_wrong(): + a, b = cross_pair(scale_b=1e-4) + first, second = detectors.crosssat_v1(a, b) + assert first.verdict == detectors.UNKNOWN and second.verdict == detectors.UNKNOWN + + +def test_crosssat_abstains_without_a_counterpart(): + a, _ = cross_pair() + first, second = detectors.crosssat_v1(a, None) + assert first.verdict == detectors.UNKNOWN and second.verdict == detectors.UNKNOWN + + +def test_crosssat_abstains_without_thumbnails(): + a = detectors.FrameFeatures(slot=(16, BAND, 0), header={}) + b = detectors.FrameFeatures(slot=(18, BAND, 0), header={}) + first, _ = detectors.crosssat_v1(a, b) + assert first.verdict == detectors.UNKNOWN + + +# --------------------------------------------------------------------------- utils + + +def test_thumbnail_shrinks_and_removes_nans(): + image = solar_disc() + image[0:10, 0:10] = np.nan + thumb = detectors.thumbnail(image) + assert thumb.shape == (detectors.THUMBNAIL_SIZE, detectors.THUMBNAIL_SIZE) + assert np.isfinite(thumb).all() + + +def test_align_shift_measures_a_known_translation(): + base = solar_disc(size=128, radius=38) + shifted = np.roll(base, 5, axis=1) + assert detectors.align_shift(base, shifted) == pytest.approx(5.0, abs=1.0) + assert detectors.align_shift(base, base) == pytest.approx(0.0, abs=0.5) + + +def test_align_shift_returns_none_on_mismatched_shapes(): + assert detectors.align_shift(np.zeros((8, 8)), np.zeros((4, 4))) is None + assert detectors.align_shift(None, np.zeros((4, 4))) is None + + +def test_robust_z_handles_a_constant_neighbourhood(): + values = [1.0] * 10 + assert detectors._robust_z(values, 5, 4) == 0.0 + values[5] = 9.0 + assert detectors._robust_z(values, 5, 4) == np.inf + + +def test_robust_z_needs_enough_neighbours(): + assert detectors._robust_z([1.0, 2.0], 0, 4) is None + + +# ---------------------------------------------------------------------- disc_v1 + +DISC_HEADER = {"diam_sun": 772.0} + + +def quiet_disc(peak=1.1): + """A disc whose quiet regions sit near the 171A threshold, as real frames do.""" + return solar_disc(peak=peak, active_region=False) + + +def add_active_region(image, offset, strength=6.0, width=0.30, radius=386): + """Add a bright compact region at `offset` solar radii along +x from centre.""" + yy, xx = np.mgrid[0 : image.shape[0], 0 : image.shape[1]].astype(np.float32) + centre = (image.shape[1] - 1) / 2.0 + on_disc = np.hypot(xx - centre, yy - centre) < radius + spot = np.hypot(xx - (centre + radius * offset), yy - centre) + return image + strength * np.exp(-((spot / (radius * width)) ** 2)) * on_disc + + +def test_disc_accepts_a_synthetic_disc(): + assert detectors.disc_v1(bright_disc(), BAND, DISC_HEADER).verdict == detectors.GOOD + + +def test_disc_measurements_are_unmoved_by_an_active_region(): + """The whole point of the detector, and the fix for the 36.9% rejection rate. + + Every disc_v1 measurement is an average over angle, so where the bright regions + sit does not move it. geometry_v1 averages along image rows and columns + instead, so its centroid swings with the active region -- and because its centre + test is one-sided, an equal displacement is fatal on one limb and harmless on + the other. + """ + base = quiet_disc() + left = add_active_region(base, -0.45) + right = add_active_region(base, +0.45) + + measurements = [detectors.disc_profile(img, 772.0 / 4.0) for img in (base, left, right)] + for name in ("radius_ratio", "limb_width"): + values = [m[name] for m in measurements] + assert max(values) - min(values) < 1e-6, f"{name} moved: {values}" + contrasts = [m["limb_contrast"] for m in measurements] + assert max(contrasts) - min(contrasts) < 0.01 + + # Meanwhile the baseline's centroid swings, in opposite directions. + centroids = [detectors.geometry_v1(img, BAND).scores["centre_x"] for img in (left, base, right)] + assert centroids[0] < centroids[1] < centroids[2] + assert centroids[2] - centroids[0] > 5.0 + + for img in (base, left, right): + assert detectors.disc_v1(img, BAND, DISC_HEADER).verdict == detectors.GOOD + + +@pytest.mark.parametrize( + "image,label", + [ + (np.zeros((1280, 1280), np.float32), "all zero"), + (np.full((1280, 1280), np.nan, np.float32), "all NaN"), + (np.full((1280, 1280), 1.0, np.float32), "uniform field"), + ], +) +def test_disc_rejects_frames_with_no_disc(image, label): + """A frame with no radial structure is a conclusion, not an abstention.""" + verdict = detectors.disc_v1(image, BAND, DISC_HEADER) + assert verdict.verdict == detectors.BAD, label + assert verdict.scores["limb_contrast"] == 0.0 + + +def test_disc_detects_a_wrong_sized_disc(): + small = solar_disc(radius=300, peak=3.0, active_region=False) + measured = detectors.disc_profile(small, 772.0 / 4.0) + assert measured["radius_ratio"] < 0.85 + + +def test_disc_detects_a_displaced_disc_through_limb_smearing(): + """A decentred disc smears the azimuthally averaged limb; that is the signal.""" + base = solar_disc(peak=3.0, active_region=False) + sharp = detectors.disc_profile(base, 772.0 / 4.0) + for shift in (20, 40): + moved = detectors.disc_profile(np.roll(base, shift, axis=1), 772.0 / 4.0) + assert moved["limb_width"] > sharp["limb_width"] * 2, f"shift {shift}" + + +def test_disc_detects_a_blurred_limb(): + base = solar_disc(peak=3.0, active_region=False) + blurred = cv.GaussianBlur(base, (81, 81), 25) + assert ( + detectors.disc_profile(blurred, 772.0 / 4.0)["limb_width"] + > detectors.disc_profile(base, 772.0 / 4.0)["limb_width"] * 2 + ) + + +def test_disc_is_blind_to_rotation_by_construction(): + """Documents a deliberate limit: rotate/yaw_flip belong to the other detectors. + + Any azimuthally averaged quantity is rotation invariant, and a single frame + carries no absolute rotation reference beyond the CROTA header. + """ + image = solar_disc(peak=3.0) + flipped = np.flip(np.flip(image, 0), 1).copy() + original = detectors.disc_profile(image, 772.0 / 4.0) + rotated = detectors.disc_profile(flipped, 772.0 / 4.0) + for name, value in original.items(): + assert rotated[name] == pytest.approx(value, abs=1e-6), name + + +def test_disc_abstains_without_the_expected_radius(): + verdict = detectors.disc_v1(bright_disc(), BAND, {}) + assert verdict.verdict == detectors.UNKNOWN and "DIAM_SUN" in verdict.reason + + +def test_disc_abstains_on_an_uncalibrated_band(): + verdict = detectors.disc_v1(bright_disc(), 999, DISC_HEADER) + assert verdict.verdict == detectors.UNKNOWN + + +def test_disc_reports_a_missing_image(): + assert detectors.disc_v1(None, BAND, DISC_HEADER).verdict == detectors.BAD + + +def test_disc_profile_rejects_nonsense_input(): + for image in (None, np.zeros((4, 4), np.float32), np.zeros((8, 8, 3), np.float32)): + assert detectors.disc_profile(image, 193.0)["limb_contrast"] is None + assert detectors.disc_profile(bright_disc(), 0)["limb_contrast"] is None + assert detectors.disc_profile(bright_disc(), np.nan)["limb_contrast"] is None + + +def test_disc_honours_supplied_bounds(): + """Bounds are per band and injectable, so calibration is testable in isolation.""" + image = bright_disc() + permissive = {BAND: {"limb_contrast": (0.0, 1.0), "radius_ratio": None, "limb_width": None}} + strict = {BAND: {"limb_contrast": (0.999, 1.0), "radius_ratio": None, "limb_width": None}} + assert detectors.disc_v1(image, BAND, DISC_HEADER, permissive).verdict == detectors.GOOD + assert detectors.disc_v1(image, BAND, DISC_HEADER, strict).verdict == detectors.BAD + + +def test_disc_bounds_set_to_none_record_a_score_without_judging(): + image = bright_disc() + bounds = {BAND: {"limb_contrast": None, "radius_ratio": None, "limb_width": None}} + verdict = detectors.disc_v1(image, BAND, DISC_HEADER, bounds) + assert verdict.verdict == detectors.GOOD + assert "radius_ratio" in verdict.scores # still measured and reported + + +def test_disc_is_registered_as_a_frame_detector(): + assert "disc_v1" in detectors.FRAME_DETECTORS + assert "disc_v1" in detectors.ALL_DETECTORS diff --git a/tests/test_fillers.py b/tests/test_fillers.py new file mode 100644 index 0000000..da02589 --- /dev/null +++ b/tests/test_fillers.py @@ -0,0 +1,215 @@ +import numpy as np +import pytest + +from conftest import solar_disc +from suvi import fillers + +HEADER = { + "crpix1": 640.5, + "crpix2": 640.5, + "cdelt1": 2.5, + "diam_sun": 771.98, + "dsun_obs": 148781338180.972, + "solar_b0": -7.170855, +} + + +def context(**kwargs): + base = dict(dt_before=240.0, dt_after=240.0, header=HEADER) + base.update(kwargs) + return fillers.FillContext(**base) + + +# -------------------------------------------------------------------------- context + + +def test_alpha_locates_the_frame_within_its_gap(): + assert context(dt_before=240, dt_after=240).alpha == pytest.approx(0.5) + assert context(dt_before=240, dt_after=720).alpha == pytest.approx(0.25) + assert context(dt_before=0, dt_after=0).alpha == 0.0 + + +def test_gap_frames_counts_slots(): + assert context(dt_before=240, dt_after=240).gap_frames == 2 + assert context(dt_before=240, dt_after=2160).gap_frames == 10 + + +# --------------------------------------------------------------------- baselines + + +def test_hold_last_repeats_the_preceding_frame(): + before, after = np.ones((8, 8), np.float32), np.zeros((8, 8), np.float32) + np.testing.assert_array_equal(fillers.hold_last(context(before=before, after=after)), before) + + +def test_hold_last_falls_back_to_the_following_frame(): + after = np.full((8, 8), 3.0, np.float32) + np.testing.assert_array_equal(fillers.hold_last(context(after=after)), after) + + +def test_hold_last_gives_up_with_nothing_to_hold(): + assert fillers.hold_last(context()) is None + + +def test_linear_blend_is_exact_on_a_linear_ramp(): + """A quantity changing linearly in time must be recovered exactly.""" + before = np.full((8, 8), 10.0, np.float32) + after = np.full((8, 8), 20.0, np.float32) + filled = fillers.linear_blend(context(before=before, after=after)) + np.testing.assert_allclose(filled, 15.0, rtol=1e-6) + + skewed = fillers.linear_blend( + context(before=before, after=after, dt_before=240, dt_after=720) + ) + np.testing.assert_allclose(skewed, 12.5, rtol=1e-6) + + +def test_linear_blend_degrades_to_hold_last_at_a_boundary(): + before = np.full((8, 8), 7.0, np.float32) + np.testing.assert_array_equal(fillers.linear_blend(context(before=before)), before) + + +def test_fillers_tolerate_non_finite_input(): + before = np.full((8, 8), np.nan, np.float32) + after = np.ones((8, 8), np.float32) + filled = fillers.linear_blend(context(before=before, after=after)) + assert np.isfinite(filled).all() + + +# ------------------------------------------------------------------ optical flow + + +def test_optical_flow_tracks_a_translation(): + """A feature moving at constant speed should land mid-way, not appear twice.""" + base = solar_disc(size=256, radius=70, peak=3.0) + before = np.roll(base, -8, axis=1) + after = np.roll(base, 8, axis=1) + filled = fillers.optical_flow(context(before=before, after=after)) + blended = fillers.linear_blend(context(before=before, after=after)) + # The truth is the untranslated frame; flow should beat a plain cross-fade. + assert np.abs(filled - base).mean() < np.abs(blended - base).mean() + + +def test_optical_flow_falls_back_without_two_brackets(): + before = np.ones((32, 32), np.float32) + np.testing.assert_array_equal(fillers.optical_flow(context(before=before)), before) + + +def test_optical_flow_falls_back_on_mismatched_shapes(): + before = np.ones((32, 32), np.float32) + after = np.ones((16, 16), np.float32) + result = fillers.optical_flow(context(before=before, after=after)) + np.testing.assert_array_equal(result, before) + + +# ---------------------------------------------------------------------- crosssat + + +def test_crosssat_returns_nothing_without_a_counterpart(): + assert fillers.crosssat(context(before=np.ones((8, 8), np.float32))) is None + + +def test_crosssat_corrects_the_calibration_difference(): + """The other satellite's radiance scale must be matched, not copied blindly.""" + truth = solar_disc(size=128, radius=38, peak=2.0) + counterpart = truth * 0.5 + 0.3 # a different flight model's response + filled = fillers.crosssat(context(before=truth, counterpart=counterpart), align=False) + np.testing.assert_allclose(filled, truth, atol=1e-3) + + +def test_crosssat_aligns_a_parallax_shift(): + truth = solar_disc(size=128, radius=38, peak=2.0) + counterpart = np.roll(truth, 4, axis=1) + filled = fillers.crosssat(context(before=truth, counterpart=counterpart), align=True) + unaligned = fillers.crosssat(context(before=truth, counterpart=counterpart), align=False) + assert np.abs(filled - truth).mean() < np.abs(unaligned - truth).mean() + + +def test_gain_match_recovers_an_affine_transform(): + source = solar_disc(size=64, radius=20, peak=1.0) + reference = source * 3.0 - 0.5 + np.testing.assert_allclose(fillers.gain_match(source, reference), reference, atol=1e-4) + + +def test_gain_match_survives_a_constant_source(): + flat = np.ones((8, 8), np.float32) + result = fillers.gain_match(flat, np.arange(64, dtype=np.float32).reshape(8, 8)) + assert np.isfinite(result).all() + + +# --------------------------------------------------------------- solar rotation + + +def test_rotation_rate_is_fastest_at_the_equator(): + equator = fillers.rotation_rate(0.0) + mid = fillers.rotation_rate(np.radians(45)) + pole = fillers.rotation_rate(np.radians(80)) + assert equator > mid > pole + + +def test_synodic_rate_is_slower_than_sidereal(): + """An Earth-orbiting observer sees the Sun turn more slowly than the stars do.""" + assert fillers.rotation_rate(0.0, synodic=True) < fillers.rotation_rate(0.0, synodic=False) + difference = fillers.rotation_rate(0.0, synodic=False) - fillers.rotation_rate(0.0, True) + assert difference == pytest.approx(fillers.EARTH_ORBIT_DEG_PER_DAY) + + +def test_rotation_map_is_the_identity_at_zero_lag(): + shape = (128, 128) + header = dict(HEADER, crpix1=64.5, crpix2=64.5, diam_sun=76.0) + map_x, map_y, visible = fillers._rotation_map(shape, header, 0.0) + grid_x, grid_y = np.meshgrid(np.arange(shape[1]), np.arange(shape[0])) + np.testing.assert_allclose(map_x[visible], grid_x[visible], atol=0.01) + np.testing.assert_allclose(map_y[visible], grid_y[visible], atol=0.01) + + +def test_rotation_map_marks_the_off_disc_region_invisible(): + header = dict(HEADER, crpix1=64.5, crpix2=64.5, diam_sun=76.0) + _, _, visible = fillers._rotation_map((128, 128), header, 0.0) + assert not visible[0, 0] # a corner is outside the disc + assert visible[64, 64] # disc centre + + +def test_rotation_warp_round_trips(): + """Warping forward by dt then back by dt must return the original on-disc.""" + header = dict(HEADER, crpix1=64.5, crpix2=64.5, diam_sun=76.0) + image = solar_disc(size=128, radius=38, peak=2.0) + forward, visible = fillers._warp(image, header, 3600.0, True) + back, visible_back = fillers._warp(forward, header, -3600.0, True) + core = np.zeros_like(visible) + core[50:78, 50:78] = True # well inside the disc, away from limb foreshortening + np.testing.assert_allclose(back[core], image[core], atol=0.05) + + +def test_solar_rotation_reduces_to_a_blend_off_disc(): + header = dict(HEADER, crpix1=64.5, crpix2=64.5, diam_sun=76.0) + before = np.full((128, 128), 1.0, np.float32) + after = np.full((128, 128), 3.0, np.float32) + filled = fillers.solar_rotation(context(before=before, after=after, header=header)) + assert filled[0, 0] == pytest.approx(2.0, abs=1e-4) # corner: pure cross-fade + + +def test_solar_rotation_handles_a_single_bracket(): + header = dict(HEADER, crpix1=64.5, crpix2=64.5, diam_sun=76.0) + image = solar_disc(size=128, radius=38, peak=2.0) + filled = fillers.solar_rotation(context(before=image, header=header)) + assert filled is not None and filled.shape == image.shape + + +def test_solar_rotation_gives_up_with_no_brackets(): + assert fillers.solar_rotation(context()) is None + + +# ---------------------------------------------------------------------- registry + + +def test_every_registered_filler_is_callable_and_shape_preserving(): + before = solar_disc(size=64, radius=20, peak=2.0) + after = solar_disc(size=64, radius=20, peak=2.2) + header = dict(HEADER, crpix1=32.5, crpix2=32.5, diam_sun=40.0) + ctx = context(before=before, after=after, counterpart=after, header=header) + for name, filler in fillers.FILLERS.items(): + result = filler(ctx) + assert result is not None, name + assert result.shape == before.shape, name + assert np.isfinite(result).all(), name diff --git a/tests/test_fitsio.py b/tests/test_fitsio.py new file mode 100644 index 0000000..3107eb3 --- /dev/null +++ b/tests/test_fitsio.py @@ -0,0 +1,229 @@ +import numpy as np +import pytest +from astropy.io import fits + +from conftest import solar_disc, write_fits +from suvi import fitsio + + +@pytest.fixture +def disc(): + return solar_disc() + + +def corrupt(path, tmp_path, name, transform): + raw = open(path, "rb").read() + out = str(tmp_path / name) + with open(out, "wb") as handle: + handle.write(transform(raw)) + return out + + +# ------------------------------------------------------------------ value parsing + + +@pytest.mark.parametrize( + "raw,expected", + [ + (" T / comment", True), + (" F", False), + (" 171 / [angstrom]", 171), + (" 2.5 / arcsec", 2.5), + ("7.40304005902818E-05 / [W m-2]", pytest.approx(7.40304005902818e-05)), + ("1.5D3 / fortran exponent", 1500.0), + ("'UTC ' / principal time system", "UTC"), + ("'W m-2 sr-1'", "W m-2 sr-1"), + ("'0 ' / data unit checksum", "0"), + ("'it''s quoted' / embedded quote", "it's quoted"), + ("'a / b' / slash inside string", "a / b"), + ("", None), + (" -999999", -999999), + ], +) +def test_parse_value(raw, expected): + assert fitsio._parse_value(raw) == expected + + +# ------------------------------------------------------------------ header reading + + +def test_reads_header_of_a_compressed_image(good_frame): + cards, error = fitsio.read_header(good_frame) + assert error is None + assert cards["WAVELNTH"] == 171 + assert cards["ZIMAGE"] is True + assert cards["ZNAXIS1"] == 1280 + assert cards["DEGRADED"] is False + assert cards["IMG_MEAN"] > 0 + + +def test_reads_header_of_an_uncompressed_image(tmp_path, disc): + path = write_fits(str(tmp_path / "plain.fits"), disc, compress=False) + cards, error = fitsio.read_header(path) + assert error is None + assert cards["NAXIS"] == 2 and cards["NAXIS1"] == 1280 + assert cards["WAVELNTH"] == 171 + + +def test_header_metadata_projects_onto_db_columns(good_frame): + values, error = fitsio.scan_header(good_frame) + assert error is None + # Booleans become 0/1 so SQL can filter on them. + assert values["degraded"] == 0 and values["eclipse"] == 0 and values["empty"] == 0 + assert values["wavelnth"] == 171 + assert isinstance(values["img_mean"], float) + assert isinstance(values["datasum"], str) + assert set(values) <= set(fitsio.db.HEADER_FIELDS) + + +def test_header_metadata_drops_wrongly_typed_cards(): + values = fitsio.header_metadata({"IMG_MEAN": "not a number", "WAVELNTH": 171}) + assert "img_mean" not in values + assert values["wavelnth"] == 171 + + +def test_flags_are_read_from_a_degraded_frame(tmp_path, disc): + """The real eclipse failure mode: near-zero radiance with DEGRADED/ECLIPSE set.""" + path = write_fits( + str(tmp_path / "eclipse.fits"), + disc * 1e-4, + headers={"DEGRADED": True, "ECLIPSE": 2}, + ) + values, error = fitsio.scan_header(path) + assert error is None + assert values["degraded"] == 1 and values["eclipse"] == 2 + assert values["img_mean"] < 1e-3 + + +# ----------------------------------------------------------------------- integrity + + +def test_empty_file(tmp_path): + path = str(tmp_path / "empty.fits") + open(path, "wb").close() + cards, error = fitsio.read_header(path) + assert cards == {} and error == "empty file" + + +def test_missing_file(tmp_path): + cards, error = fitsio.read_header(str(tmp_path / "absent.fits")) + assert cards == {} and "unreadable" in error + + +def test_not_a_fits_file(tmp_path): + path = str(tmp_path / "junk.fits") + open(path, "wb").write(b"X" * 5000) + cards, error = fitsio.read_header(path) + assert cards == {} and "missing SIMPLE" in error + + +def test_shorter_than_one_block(tmp_path, good_frame): + path = corrupt(good_frame, tmp_path, "tiny.fits", lambda raw: raw[:100]) + _, error = fitsio.read_header(path) + assert "truncated" in error + + +def test_truncation_not_on_a_block_boundary_is_reported(tmp_path, good_frame): + path = corrupt(good_frame, tmp_path, "odd.fits", lambda raw: raw[:-100]) + _, error = fitsio.read_header(path) + assert "not a multiple of 2880" in error + + +def test_header_without_end_card(tmp_path, good_frame): + """A header whose END is gone must not be parsed as if it were complete.""" + path = corrupt( + good_frame, tmp_path, "noend.fits", lambda raw: raw.replace(b"END ", b"XXX ") + ) + cards, error = fitsio.read_header(path) + assert cards == {} and "END" in error + + +def test_primary_only_file_has_no_image_hdu(tmp_path): + path = str(tmp_path / "primary.fits") + fits.HDUList([fits.PrimaryHDU()]).writeto(path, overwrite=True) + cards, error = fitsio.read_header(path) + assert "no image HDU" in error + + +def test_max_header_bytes_is_respected(good_frame): + """A malformed file must not pull unbounded data into memory.""" + cards, error = fitsio.read_header(good_frame, max_bytes=fitsio.BLOCK) + assert cards == {} or error is not None + + +# ------------------------------------------------------------------------ datasums + + +def test_datasum_matches_the_stored_keyword(good_frame): + assert fitsio.verify_datasums(good_frame) is None + + +def test_datasum_detects_a_flipped_bit(tmp_path, good_frame): + def flip(raw): + offset = len(raw) - fitsio.BLOCK # inside the final data unit + return raw[:offset] + bytes([raw[offset] ^ 0xFF]) + raw[offset + 1 :] + + path = corrupt(good_frame, tmp_path, "flipped.fits", flip) + assert "datasum mismatch" in fitsio.verify_datasums(path) + + +def test_datasum_detects_truncation(tmp_path, good_frame): + path = corrupt(good_frame, tmp_path, "half.fits", lambda raw: raw[: len(raw) // 2]) + assert "truncated" in fitsio.verify_datasums(path) + + +def test_datasum_reports_when_no_keywords_are_present(tmp_path, disc): + path = write_fits(str(tmp_path / "nosum.fits"), disc, checksum=False) + assert fitsio.verify_datasums(path) == "no DATASUM keywords present" + + +def test_datasum_on_a_non_fits_file(tmp_path): + path = str(tmp_path / "junk.fits") + open(path, "wb").write(b"nope" * 1000) + assert fitsio.verify_datasums(path) == "not a FITS file" + + +def test_datasum_is_the_fits_ones_complement_sum(): + # Two words that overflow 32 bits, forcing an end-around carry. + payload = (0xFFFFFFFF).to_bytes(4, "big") + (0x00000002).to_bytes(4, "big") + assert fitsio.datasum(payload) == 2 + assert fitsio.datasum(b"") == 0 + # A trailing partial word is zero-padded on the right: 0x00000001 + 0x01000000. + assert fitsio.datasum(b"\x00\x00\x00\x01\x01") == 0x01000001 + + +# --------------------------------------------------------------------- image reads + + +def test_reads_pixel_data(good_frame, disc): + data, error = fitsio.read_image(good_frame) + assert error is None + assert data.shape == (1280, 1280) and data.dtype == np.float32 + assert np.allclose(data, disc, atol=1e-3) + + +def test_zblank_becomes_nan(tmp_path, disc): + holed = disc.copy() + holed[10:20, 10:20] = fitsio.ZBLANK + path = write_fits(str(tmp_path / "blank.fits"), holed) + data, error = fitsio.read_image(path) + assert error is None + assert np.isnan(data[10:20, 10:20]).all() + assert np.isfinite(data[500:600, 500:600]).all() + + +def test_image_read_reports_a_missing_array(tmp_path): + path = str(tmp_path / "primary.fits") + fits.HDUList([fits.PrimaryHDU()]).writeto(path, overwrite=True) + data, error = fitsio.read_image(path) + assert data is None and "no 2-D image array" in error + + +def test_image_read_reports_corruption_instead_of_raising(tmp_path, good_frame): + path = corrupt(good_frame, tmp_path, "broken.fits", lambda raw: raw[: len(raw) // 2]) + data, error = fitsio.read_image(path) + assert data is None and error + + +def test_quiet_astropy_is_callable(): + fitsio.quiet_astropy() diff --git a/tests/test_index.py b/tests/test_index.py new file mode 100644 index 0000000..6e26104 --- /dev/null +++ b/tests/test_index.py @@ -0,0 +1,401 @@ +"""Tests for keeping the index in step with the archive without traversing it. + +The property under test throughout is that a *re-run reads nothing it does not have +to*. That is not an optimisation here: a full traversal of this archive exhausts the +file handles of the virtiofs mount it lives on, and the mount then refuses every open +until the guest drops its dentry cache. So "how many directories did we open" is a +correctness-adjacent measurement, and several tests assert on it directly. +""" + +import os +import re + +import pytest + +from conftest import solar_disc, write_fits +from suvi import db, index, paths + +T0 = 1715299200 + + +def add_frame(root, satellite=16, wavelength=171, index_=0, day=None, label=None): + when = T0 + index_ * paths.CADENCE + name = paths.FrameName(satellite, wavelength, when, when + paths.CADENCE, "1-0-2") + relpath = name.relpath(label) + if day: + relpath = relpath.replace("/2024/05/10/", f"/{day}/") + path = os.path.join(str(root), *relpath.split("/")) + write_fits(path, solar_disc(size=32, radius=10, peak=3.0)) + return name, path, relpath + + +@pytest.fixture +def small_archive(archive): + for i in range(5): + add_frame(archive, index_=i) + for i in range(3): + add_frame(archive, wavelength=304, index_=i) + return archive + + +class CountingScandir: + """Wraps os.scandir so a test can assert how much of the tree was opened.""" + + def __init__(self, monkeypatch, only_under=None): + self.paths = [] + self._real = os.scandir + self._only = str(only_under) if only_under else None + monkeypatch.setattr(os, "scandir", self) + + def __call__(self, path=".", *args, **kwargs): + text = str(path) + if self._only is None or text.startswith(self._only): + self.paths.append(text) + return self._real(path, *args, **kwargs) + + #: .../YYYY/MM/DD -- the directories holding ~360 frames each. Month + #: directories must still be listed to discover day directories at all; it is + #: only the day directories, and the per-file lookups inside them, that a + #: no-change reconcile has to avoid. + _DAY_DIR = re.compile(r"/\d{4}/\d{2}/\d{2}/?$") + + def opened_day_dirs(self): + return [p for p in self.paths if self._DAY_DIR.search(p)] + + +# ------------------------------------------------------------- day_directories + + +def test_day_directories_finds_every_day(small_archive): + found = list(index.day_directories(str(small_archive), (16,), (171, 304))) + assert len(found) == 2 + for relpath, abspath in found: + assert relpath.endswith("2024/05/10") + assert os.path.isdir(abspath) + + +def test_day_directories_filters_by_year(small_archive): + add_frame(small_archive, day="2023/07/04") + assert len(list(index.day_directories(str(small_archive), (16,), (171,)))) == 2 + only = list(index.day_directories(str(small_archive), (16,), (171,), years=[2023])) + assert len(only) == 1 and only[0][0].endswith("2023/07/04") + + +def test_day_directories_tolerates_a_missing_archive(tmp_path): + assert list(index.day_directories(str(tmp_path), (16,), (171,))) == [] + + +# ------------------------------------------------------------------- reconcile + + +def test_first_reconcile_indexes_everything(small_archive, db_path): + conn = db.connect(db_path) + summary = index.reconcile(conn, str(small_archive), (16,), (171, 304)) + assert summary["frames_added"] == 8 + assert summary["directories_changed"] == 2 + assert conn.execute("SELECT count(*) c FROM frame").fetchone()["c"] == 8 + conn.close() + + +def test_second_reconcile_opens_no_day_directories(small_archive, db_path, monkeypatch): + """The heart of it: an unchanged archive must cost directory stats, nothing more. + + Re-reading directories that have not changed is what made a resume traverse the + whole archive, and traversal is what breaks the mount. + """ + conn = db.connect(db_path) + index.reconcile(conn, str(small_archive), (16,), (171, 304)) + + spy = CountingScandir(monkeypatch, only_under=small_archive) + summary = index.reconcile(conn, str(small_archive), (16,), (171, 304)) + assert summary["directories_changed"] == 0 + assert summary["frames_added"] == 0 + assert spy.opened_day_dirs() == [], "re-read a directory that had not changed" + conn.close() + + +def test_reconcile_picks_up_a_new_frame(small_archive, db_path): + conn = db.connect(db_path) + index.reconcile(conn, str(small_archive), (16,), (171,)) + add_frame(small_archive, index_=99) + summary = index.reconcile(conn, str(small_archive), (16,), (171,)) + assert summary["frames_added"] == 1 + assert summary["directories_changed"] == 1 + conn.close() + + +def test_reconcile_notices_a_deleted_frame(small_archive, db_path): + conn = db.connect(db_path) + index.reconcile(conn, str(small_archive), (16,), (171,)) + before = conn.execute("SELECT count(*) c FROM frame").fetchone()["c"] + + _, path, _ = add_frame(small_archive, index_=0) # existing file + os.remove(path) + summary = index.reconcile(conn, str(small_archive), (16,), (171,)) + assert summary["frames_removed"] == 1 + assert conn.execute("SELECT count(*) c FROM frame").fetchone()["c"] == before - 1 + conn.close() + + +def test_reconcile_handles_a_vanished_directory(small_archive, db_path): + import shutil + + conn = db.connect(db_path) + index.reconcile(conn, str(small_archive), (16,), (171, 304)) + shutil.rmtree(os.path.join(str(small_archive), "goes16/l2/data/suvi-l2-ci304")) + summary = index.reconcile(conn, str(small_archive), (16,), (171, 304)) + assert summary["directories_vanished"] == 1 + assert summary["frames_removed"] == 3 + assert conn.execute("SELECT count(*) c FROM frame").fetchone()["c"] == 5 + conn.close() + + +def test_force_rereads_everything(small_archive, db_path, monkeypatch): + conn = db.connect(db_path) + index.reconcile(conn, str(small_archive), (16,), (171,)) + spy = CountingScandir(monkeypatch, only_under=small_archive) + summary = index.reconcile(conn, str(small_archive), (16,), (171,), force=True) + assert summary["directories_changed"] == 1 + assert spy.opened_day_dirs(), "--full should re-read directories" + conn.close() + + +def test_reconcile_is_idempotent(small_archive, db_path): + conn = db.connect(db_path) + first = index.reconcile(conn, str(small_archive), (16,), (171,)) + second = index.reconcile(conn, str(small_archive), (16,), (171,)) + third = index.reconcile(conn, str(small_archive), (16,), (171,)) + assert first["frames_added"] == 5 + assert second["frames_added"] == third["frames_added"] == 0 + assert conn.execute("SELECT count(*) c FROM frame").fetchone()["c"] == 5 + conn.close() + + +def test_reconcile_ignores_non_frame_files(small_archive, db_path): + conn = db.connect(db_path) + day = os.path.join(str(small_archive), "goes16/l2/data/suvi-l2-ci171/2024/05/10") + open(os.path.join(day, "notes.txt"), "w").write("not a frame") + summary = index.reconcile(conn, str(small_archive), (16,), (171,)) + assert summary["frames_added"] == 5 + conn.close() + + +def test_reconcile_records_the_observed_mtime(small_archive, db_path): + conn = db.connect(db_path) + index.reconcile(conn, str(small_archive), (16,), (171,)) + relpath = "goes16/l2/data/suvi-l2-ci171/2024/05/10" + abspath = os.path.join(str(small_archive), *relpath.split("/")) + assert db.get_dir_mtime(conn, relpath) == pytest.approx(os.stat(abspath).st_mtime) + conn.close() + + +# ------------------------------------------------------------ record_downloaded + + +def test_record_downloaded_indexes_without_any_traversal(archive, db_path, monkeypatch): + """The cheapest path: the downloader already knows the file is there.""" + conn = db.connect(db_path) + _, path, relpath = add_frame(archive) + + spy = CountingScandir(monkeypatch, only_under=archive) + frame_id = index.record_downloaded(conn, path, str(archive)) + assert frame_id is not None + assert spy.paths == [], "indexing a download must not scan the archive" + row = conn.execute("SELECT path FROM frame WHERE id = ?", (frame_id,)).fetchone() + assert row["path"] == relpath + conn.close() + + +def test_record_downloaded_ignores_other_files(archive, db_path): + conn = db.connect(db_path) + other = os.path.join(str(archive), "readme.txt") + open(other, "w").write("x") + assert index.record_downloaded(conn, other, str(archive)) is None + conn.close() + + +def test_record_downloaded_survives_a_missing_file(archive, db_path): + conn = db.connect(db_path) + _, path, _ = add_frame(archive) + os.remove(path) + assert index.record_downloaded(conn, path, str(archive)) is not None + conn.close() + + +# ------------------------------------------------------------------ db helpers + + +def test_frames_in_dir_scopes_to_one_directory(small_archive, db_path): + conn = db.connect(db_path) + index.reconcile(conn, str(small_archive), (16,), (171, 304)) + in_171 = db.frames_in_dir(conn, "goes16/l2/data/suvi-l2-ci171/2024/05/10") + in_304 = db.frames_in_dir(conn, "goes16/l2/data/suvi-l2-ci304/2024/05/10") + assert len(in_171) == 5 and len(in_304) == 3 + assert all("ci171" in name for name in in_171) + conn.close() + + +def test_schema_upgrades_from_version_one(tmp_path): + """An index created before dir_scan existed must gain it, not be rejected.""" + path = str(tmp_path / "old.sqlite") + conn = db.connect(path) + conn.execute("DROP TABLE dir_scan") + conn.execute("UPDATE meta SET value = '1' WHERE key = 'schema_version'") + conn.commit() + conn.close() + + conn = db.connect(path) # re-open triggers the upgrade + assert db.get_meta(conn, "schema_version") == str(db.SCHEMA_VERSION) + conn.execute("SELECT count(*) FROM dir_scan") # table now exists + conn.close() + + +def test_meta_round_trip(db_path): + conn = db.connect(db_path) + assert db.get_meta(conn, "absent") is None + assert db.get_meta(conn, "absent", "fallback") == "fallback" + db.set_meta(conn, "unrename_done:ci171:2024", "done") + db.set_meta(conn, "unrename_done:ci171:2024", "done") # upsert, not duplicate + assert db.get_meta(conn, "unrename_done:ci171:2024") == "done" + conn.close() + + +# ------------------------------------------------------------ slot collisions + + +def test_two_files_claiming_one_slot_do_not_break_indexing(archive, db_path): + """The archive holds 1,742 such pairs; they must not block the whole index. + + An older filter labelled some frames repeatedly, and the puller later + re-downloaded a clean copy it could no longer find under the published name -- + leaving `X_v1-0-1.fits` beside a byte-identical `X_v1-0-1_f_f_f.fits`. + """ + _, path, _ = add_frame(archive) # canonical + add_frame(archive, label="f") # same slot, labelled + conn = db.connect(db_path) + summary = index.reconcile(conn, str(archive), (16,), (171,)) + assert summary["frames_added"] == 1 + assert len(summary["duplicate_slots"]) == 1 + assert summary["duplicate_slots"][0].endswith("_f.fits") + row = conn.execute("SELECT path FROM frame").fetchone() + assert row["path"].endswith("v1-0-2.fits"), "kept the labelled copy over the published one" + conn.close() + + +def test_the_published_name_wins_over_a_multiply_labelled_one(archive, db_path): + _, path, _ = add_frame(archive) + name = paths.parse_frame_filename(os.path.basename(path)) + triple = os.path.join(os.path.dirname(path), name.filename().replace(".fits", "_f_f_f.fits")) + write_fits(triple, solar_disc(size=32, radius=10, peak=3.0)) + + conn = db.connect(db_path) + summary = index.reconcile(conn, str(archive), (16,), (171,)) + assert summary["frames_added"] == 1 + assert conn.execute("SELECT path FROM frame").fetchone()["path"].endswith("v1-0-2.fits") + assert summary["duplicate_slots"][0].endswith("_f_f_f.fits") + conn.close() + + +def test_a_renamed_file_does_not_collide_with_its_old_row(archive, db_path): + """Un-renaming moves a file within its directory; reconcile must cope. + + The removal has to be applied before the insertion, or both names briefly claim + the same observation slot and the unique constraint fires. + """ + _, old_path, _ = add_frame(archive, label="f") + conn = db.connect(db_path) + index.reconcile(conn, str(archive), (16,), (171,)) + assert conn.execute("SELECT count(*) c FROM frame").fetchone()["c"] == 1 + + name = paths.parse_frame_filename(os.path.basename(old_path)) + os.rename(old_path, os.path.join(os.path.dirname(old_path), name.filename())) + + summary = index.reconcile(conn, str(archive), (16,), (171,)) + assert summary["frames_added"] == 1 and summary["frames_removed"] == 1 + assert conn.execute("SELECT count(*) c FROM frame").fetchone()["c"] == 1 + assert conn.execute("SELECT path FROM frame").fetchone()["path"].endswith("v1-0-2.fits") + conn.close() + + +# ------------------------------------------------- unreadable != empty + + +def test_unreadable_directory_raises_rather_than_looking_empty(small_archive, monkeypatch): + """Regression: swallowing OSError made five of six bands look absent. + + reconcile then treated their contents as deleted and dropped 205,618 index rows + for files that were still on disk. An unreadable directory has to stop the run. + """ + real = os.scandir + + def refuse(path=".", *args, **kwargs): + if "suvi-l2-ci304" in str(path): + raise OSError(23, "Too many open files in system") + return real(path, *args, **kwargs) + + monkeypatch.setattr(os, "scandir", refuse) + with pytest.raises(OSError): + list(index.day_directories(str(small_archive), (16,), (171, 304))) + + +def test_a_missing_directory_is_still_treated_as_empty(tmp_path): + """Only genuine absence may be silent.""" + assert index._subdirs(str(tmp_path / "does-not-exist")) == [] + + +def test_frames_are_not_purged_when_enumeration_comes_up_short(small_archive, db_path, monkeypatch): + """Defence in depth: absence from the scan is not proof a directory is gone. + + Even if enumeration misses a directory, its rows must survive as long as the + directory is still on disk. + """ + conn = db.connect(db_path) + index.reconcile(conn, str(small_archive), (16,), (171, 304)) + before = conn.execute("SELECT count(*) c FROM frame").fetchone()["c"] + assert before == 8 + + # Enumerate only one band, as though the other had been missed entirely. + real = index.day_directories + partial = list(real(str(small_archive), (16,), (171,))) + monkeypatch.setattr(index, "day_directories", lambda *a, **k: iter(partial)) + summary = index.reconcile(conn, str(small_archive), (16,), (171, 304)) + assert summary["directories_vanished"] == 0 + assert conn.execute("SELECT count(*) c FROM frame").fetchone()["c"] == before + conn.close() + + +def test_a_cold_build_reclaims_periodically(small_archive, db_path, monkeypatch): + """A cold build reads every file, which is what exhausts the mount's handles. + + It must hand them back as it goes, or the build cannot finish. + """ + calls = [] + monkeypatch.setattr(index.vfs, "release_handles", lambda *a, **k: calls.append(1)) + monkeypatch.setattr(index.vfs, "RELIEF_INTERVAL", 3) + conn = db.connect(db_path) + index.reconcile(conn, str(small_archive), (16,), (171, 304)) + assert calls, "cold build never reclaimed" + conn.close() + + +def test_reclaim_can_be_disabled(small_archive, db_path, monkeypatch): + def explode(*a, **k): + raise AssertionError("reclaimed despite relief=False") + + monkeypatch.setattr(index.vfs, "release_handles", explode) + monkeypatch.setattr(index.vfs, "RELIEF_INTERVAL", 1) + conn = db.connect(db_path) + index.reconcile(conn, str(small_archive), (16,), (171,), relief=False) + conn.close() + + +def test_a_steady_state_run_does_not_reclaim(small_archive, db_path, monkeypatch): + """Nothing changed means nothing was read, so there is nothing to hand back.""" + conn = db.connect(db_path) + index.reconcile(conn, str(small_archive), (16,), (171,)) + calls = [] + monkeypatch.setattr(index.vfs, "release_handles", lambda *a, **k: calls.append(1)) + monkeypatch.setattr(index.vfs, "RELIEF_INTERVAL", 1) + index.reconcile(conn, str(small_archive), (16,), (171,)) + assert calls == [] + conn.close() diff --git a/tests/test_metrics.py b/tests/test_metrics.py new file mode 100644 index 0000000..1a477bd --- /dev/null +++ b/tests/test_metrics.py @@ -0,0 +1,350 @@ +import numpy as np +import pytest + +from conftest import solar_disc +from suvi import detectors, metrics + +BAND = 171 + + +def verdict(kind, reason=None): + return detectors.Verdict(kind, reason, {}) + + +# ------------------------------------------------------------------------ display + + +def test_display_mapping_is_bounded(): + image = solar_disc(size=64, radius=20, peak=50.0) + shown = metrics.to_display(image, BAND) + assert shown.min() >= 0.0 and shown.max() <= 1.0 + + +def test_display_mapping_is_monotonic(): + ramp = np.linspace(0, 20, 100).reshape(10, 10) + shown = metrics.to_display(ramp, BAND) + assert np.all(np.diff(shown.ravel()) >= -1e-12) + + +def test_display_mapping_differs_per_band(): + ramp = np.full((4, 4), 5.0) + assert not np.allclose(metrics.to_display(ramp, 94), metrics.to_display(ramp, 304)) + + +def test_display_mapping_handles_nan(): + assert np.isfinite(metrics.to_display(np.full((4, 4), np.nan), BAND)).all() + + +# ---------------------------------------------------------------------- detection + + +def test_scores_a_perfect_detector(): + injected = {(16, BAND, 0): "eclipse_dim", (16, BAND, 240): "truncate"} + verdicts = { + (16, BAND, 0): verdict("bad"), + (16, BAND, 240): verdict("bad"), + (16, BAND, 480): verdict("good"), + (16, BAND, 720): verdict("good"), + } + score = metrics.score_detection(verdicts, injected) + assert (score.true_positives, score.false_negatives) == (2, 0) + assert (score.false_positives, score.true_negatives) == (0, 2) + assert score.precision == 1.0 and score.recall == 1.0 and score.f1 == 1.0 + assert score.false_positive_rate == 0.0 + + +def test_scores_a_detector_that_flags_everything(): + injected = {(16, BAND, 0): "eclipse_dim"} + verdicts = {(16, BAND, t): verdict("bad") for t in (0, 240, 480, 720)} + score = metrics.score_detection(verdicts, injected) + assert score.recall == 1.0 + assert score.precision == pytest.approx(0.25) + assert score.false_positive_rate == 1.0 + + +def test_unknown_verdicts_are_excluded_not_counted_as_good(): + """Abstaining is neither a hit nor a miss, and must not flatter the score.""" + injected = {(16, BAND, 0): "eclipse_dim"} + verdicts = {(16, BAND, 0): verdict("unknown"), (16, BAND, 240): verdict("unknown")} + score = metrics.score_detection(verdicts, injected) + assert score.unknown == 2 + assert score.true_positives == score.false_positives == 0 + assert score.false_negatives == score.true_negatives == 0 + assert np.isnan(score.precision) and np.isnan(score.recall) + + +def test_legacy_disagreements_are_reported_separately(): + """Flagging a frame the old filter passed may be a find, not an error.""" + injected = {} + verdicts = {(16, BAND, 0): verdict("bad"), (16, BAND, 240): verdict("bad")} + score = metrics.score_detection(verdicts, injected, legacy_good={(16, BAND, 0)}) + assert score.false_positives == 2 + assert score.legacy_disagreements == 1 + + +def test_recall_is_broken_down_by_mode(): + injected = { + (16, BAND, 0): "eclipse_dim", + (16, BAND, 240): "eclipse_dim", + (16, BAND, 480): "gain_shift", + (16, BAND, 720): "gain_shift", + } + verdicts = { + (16, BAND, 0): verdict("bad"), + (16, BAND, 240): verdict("bad"), + (16, BAND, 480): verdict("bad"), + (16, BAND, 720): verdict("good"), + } + score = metrics.score_detection(verdicts, injected) + assert score.recall_by_mode == {"eclipse_dim": 1.0, "gain_shift": 0.5} + + +def test_unknown_counts_against_a_modes_recall(): + injected = {(16, BAND, 0): "rotate", (16, BAND, 240): "rotate"} + verdicts = {(16, BAND, 0): verdict("bad"), (16, BAND, 240): verdict("unknown")} + score = metrics.score_detection(verdicts, injected) + assert score.recall_by_mode["rotate"] == 0.5 + + +def test_scoring_an_empty_case(): + score = metrics.score_detection({}, {}) + assert np.isnan(score.precision) and np.isnan(score.recall) and np.isnan(score.f1) + assert np.isnan(score.microseconds_per_frame) + + +def test_throughput_is_per_frame(): + verdicts = {(16, BAND, t): verdict("good") for t in range(0, 1000, 240)} + score = metrics.score_detection(verdicts, {}, elapsed_us=1000) + assert score.microseconds_per_frame == pytest.approx(1000 / len(verdicts)) + + +def test_as_dict_is_serialisable(): + score = metrics.score_detection({(16, BAND, 0): verdict("bad")}, {(16, BAND, 0): "x"}) + import json + + assert json.loads(json.dumps(score.as_dict()))["recall"] == 1.0 + + +# ---------------------------------------------------------------------- pr curves + + +def test_precision_recall_curve_on_a_separable_score(): + scores = [0.9, 0.8, 0.2, 0.1] + labels = [True, True, False, False] + thresholds, precision, recall = metrics.precision_recall_curve(scores, labels) + assert list(thresholds) == [0.9, 0.8, 0.2, 0.1] + assert precision[1] == 1.0 and recall[1] == 1.0 + assert metrics.average_precision(scores, labels) == pytest.approx(1.0) + + +def test_average_precision_of_an_uninformative_score(): + scores = [0.5, 0.5, 0.5, 0.5] + labels = [True, False, True, False] + assert 0.0 < metrics.average_precision(scores, labels) < 1.0 + + +def test_curves_handle_degenerate_input(): + assert metrics.average_precision([], []) != metrics.average_precision([], []) or True + assert np.isnan(metrics.average_precision([1.0], [False])) + thresholds, _, _ = metrics.precision_recall_curve([np.nan, np.inf], [True, False]) + assert len(thresholds) <= 1 + + +# --------------------------------------------------------------------------- fill + + +def test_identical_frames_score_perfectly(): + image = solar_disc(size=64, radius=20, peak=3.0) + score = metrics.score_fill(image, image, BAND) + assert score.rmse == 0.0 and score.mae == 0.0 and score.log_rmse == 0.0 + assert score.psnr == float("inf") + assert score.ssim == pytest.approx(1.0) + + +def test_worse_reconstructions_score_worse(): + truth = solar_disc(size=64, radius=20, peak=3.0) + close = truth + 0.01 + far = truth + 0.5 + assert metrics.score_fill(close, truth, BAND).rmse < metrics.score_fill(far, truth, BAND).rmse + assert metrics.score_fill(close, truth, BAND).psnr > metrics.score_fill(far, truth, BAND).psnr + assert metrics.score_fill(close, truth, BAND).ssim > metrics.score_fill(far, truth, BAND).ssim + + +def test_log_error_keeps_the_faint_corona_visible(): + """The log term makes an equal *relative* error count comparably everywhere. + + Radiance spans orders of magnitude, so a plain rmse is dominated by the bright + disc: tripling the faint corona and tripling the disc are the same mistake, but + rmse rates one ~1000x worse than the other. log_rmse compresses that gap so a + filler cannot look good by getting only the bright pixels right. + """ + faint = np.full((32, 32), 0.01) + bright = np.full((32, 32), 10.0) + faint_score = metrics.score_fill(faint * 3.0, faint, BAND) + bright_score = metrics.score_fill(bright * 3.0, bright, BAND) + + assert bright_score.rmse / faint_score.rmse == pytest.approx(1000, rel=0.01) + assert bright_score.log_rmse / faint_score.log_rmse < 100 + + +def test_fill_scoring_rejects_mismatched_shapes(): + with pytest.raises(ValueError, match="shape mismatch"): + metrics.score_fill(np.zeros((4, 4)), np.zeros((8, 8)), BAND) + + +def test_fill_scoring_tolerates_nan(): + truth = solar_disc(size=32, radius=10, peak=2.0) + filled = truth.copy() + filled[0:4, 0:4] = np.nan + score = metrics.score_fill(filled, truth, BAND) + assert np.isfinite(score.rmse) and np.isfinite(score.ssim) + + +# ------------------------------------------------------------------------ flicker + + +def test_flicker_is_zero_for_a_perfect_reconstruction(): + frames = [solar_disc(size=32, radius=10, peak=1.0 + i * 0.1) for i in range(5)] + assert metrics.temporal_flicker(frames, frames, BAND) == pytest.approx(0.0) + + +def test_flicker_penalises_a_freeze_that_psnr_forgives(): + """The artifact per-frame metrics miss: hold-then-jump instead of smooth motion.""" + truth = [solar_disc(size=32, radius=10, peak=1.0 + i * 0.4) for i in range(4)] + frozen = [truth[0], truth[0], truth[0], truth[3]] # hold, hold, then snap + smooth = truth + assert metrics.temporal_flicker(frozen, truth, BAND) > metrics.temporal_flicker( + smooth, truth, BAND + ) + + +def test_flicker_requires_matching_lengths(): + frames = [np.zeros((8, 8))] * 3 + with pytest.raises(ValueError, match="same length"): + metrics.temporal_flicker(frames, frames[:2], BAND) + + +def test_flicker_of_a_single_frame_is_undefined(): + assert np.isnan(metrics.temporal_flicker([np.zeros((8, 8))], [np.zeros((8, 8))], BAND)) + + +# ------------------------------------------------------------------------ summary + + +def test_summary_groups_by_gap_length(): + scores = [ + metrics.FillScore(1.0, 1.0, 0.1, 30.0, 0.9, gap_frames=1, wavelength=BAND), + metrics.FillScore(3.0, 3.0, 0.3, 20.0, 0.7, gap_frames=1, wavelength=BAND), + metrics.FillScore(5.0, 5.0, 0.5, 10.0, 0.5, gap_frames=10, wavelength=BAND), + ] + summary = metrics.summarise_fills(scores) + assert summary["n"] == 3 + assert summary["rmse"] == pytest.approx(3.0) + assert summary["by_gap"][1]["n"] == 2 + assert summary["by_gap"][1]["psnr"] == pytest.approx(25.0) + assert summary["by_gap"][10]["ssim"] == pytest.approx(0.5) + + +def test_summary_ignores_infinite_psnr_from_perfect_frames(): + scores = [ + metrics.FillScore(0.0, 0.0, 0.0, float("inf"), 1.0, 1, BAND), + metrics.FillScore(1.0, 1.0, 0.1, 20.0, 0.8, 1, BAND), + ] + assert metrics.summarise_fills(scores)["psnr"] == pytest.approx(20.0) + + +def test_summary_of_nothing(): + assert metrics.summarise_fills([]) == {} + + +# -------------------------------------------------------------------- combinations + + +def vmap(**pairs): + """Build a {slot: Verdict} from slot-suffix -> verdict-string pairs.""" + return {(16, BAND, int(k[1:])): verdict(v) for k, v in pairs.items()} + + +def test_any_policy_flags_if_one_detector_does(): + combined = metrics.combine_verdicts( + [vmap(t0="bad", t1="good"), vmap(t0="good", t1="good")], "any" + ) + assert combined[(16, BAND, 0)].verdict == "bad" + assert combined[(16, BAND, 1)].verdict == "good" + + +def test_all_policy_requires_unanimity(): + combined = metrics.combine_verdicts( + [vmap(t0="bad", t1="bad"), vmap(t0="good", t1="bad")], "all" + ) + assert combined[(16, BAND, 0)].verdict == "good" + assert combined[(16, BAND, 1)].verdict == "bad" + + +def test_majority_policy(): + maps = [vmap(t0="bad"), vmap(t0="bad"), vmap(t0="good")] + assert metrics.combine_verdicts(maps, "majority")[(16, BAND, 0)].verdict == "bad" + maps = [vmap(t0="bad"), vmap(t0="good"), vmap(t0="good")] + assert metrics.combine_verdicts(maps, "majority")[(16, BAND, 0)].verdict == "good" + + +def test_unknown_abstains_rather_than_voting(): + """An abstention must not act as a 'good' vote. + + Under 'all' that would let one abstaining detector veto a real detection; under + 'any' it would quietly inflate recall. + """ + combined = metrics.combine_verdicts( + [vmap(t0="bad"), vmap(t0="unknown")], "all" + ) + assert combined[(16, BAND, 0)].verdict == "bad" # the abstention is ignored + assert combined[(16, BAND, 0)].scores["votes_total"] == 1.0 + + +def test_all_detectors_abstaining_yields_unknown(): + combined = metrics.combine_verdicts([vmap(t0="unknown"), vmap(t0="unknown")], "any") + assert combined[(16, BAND, 0)].verdict == "unknown" + + +def test_slots_missing_from_one_map_are_still_judged(): + combined = metrics.combine_verdicts([vmap(t0="bad"), vmap(t1="good")], "any") + assert set(combined) == {(16, BAND, 0), (16, BAND, 1)} + assert combined[(16, BAND, 0)].verdict == "bad" + + +def test_single_detector_reduces_to_itself(): + single = vmap(t0="bad", t1="good", t2="unknown") + for policy in metrics.COMBINATION_POLICIES: + combined = metrics.combine_verdicts([single], policy) + for slot, original in single.items(): + assert combined[slot].verdict == original.verdict, policy + + +def test_combining_nothing(): + assert metrics.combine_verdicts([], "any") == {} + + +def test_unknown_policy_is_rejected(): + with pytest.raises(ValueError, match="Unknown policy"): + metrics.combine_verdicts([vmap(t0="bad")], "consensus") + + +def test_combined_verdicts_carry_the_reasons(): + maps = [vmap(t0="bad"), vmap(t0="bad")] + maps[0][(16, BAND, 0)] = detectors.Verdict("bad", "eclipse", {}) + maps[1][(16, BAND, 0)] = detectors.Verdict("bad", "limb_contrast 0.0", {}) + combined = metrics.combine_verdicts(maps, "any") + assert "eclipse" in combined[(16, BAND, 0)].reason + assert "limb_contrast" in combined[(16, BAND, 0)].reason + + +def test_a_combination_can_be_scored_like_any_detector(): + """Combinations feed straight into score_detection -- that is the point.""" + injected = {(16, BAND, 0): "eclipse_dim"} + narrow = vmap(t0="bad", t1="good") # perfect + noisy = vmap(t0="bad", t1="bad") # one false positive + strict = metrics.combine_verdicts([narrow, noisy], "all") + loose = metrics.combine_verdicts([narrow, noisy], "any") + assert metrics.score_detection(strict, injected).false_positive_rate == 0.0 + assert metrics.score_detection(loose, injected).false_positive_rate == 1.0 + assert metrics.score_detection(strict, injected).recall == 1.0 diff --git a/tests/test_migrations.py b/tests/test_migrations.py new file mode 100644 index 0000000..5da1672 --- /dev/null +++ b/tests/test_migrations.py @@ -0,0 +1,355 @@ +"""Tests for the two one-time migrations. + +These are the only steps in the project that mutate the archive or destroy state, so +the properties tested here are the ones that make them safe to run: dry runs change +nothing, re-runs are no-ops, collisions are skipped rather than forced, and nothing +is renamed until the labels it would destroy have been exported and recorded. +""" + +import csv +import gzip +import json +import os + +import pytest + +import migrate_unrename +import migrate_urlcache +from conftest import solar_disc, write_fits +from suvi import db, paths + +T0 = 1715299200 + + +def make_frame_file(root, satellite=16, wavelength=171, index=0, label=None): + when = T0 + index * paths.CADENCE + name = paths.FrameName(satellite, wavelength, when, when + paths.CADENCE, "1-0-2") + path = os.path.join(str(root), *name.relpath(label).split("/")) + write_fits(path, solar_disc(size=32, radius=10, peak=3.0)) + return name, path + + +# ------------------------------------------------------------------- url cache + + +@pytest.fixture +def url_json(tmp_path): + payload = {f"https://example.test/file{i}.fits": 1700000000.0 + i for i in range(50)} + path = tmp_path / "file_database.json" + path.write_text(json.dumps(payload)) + return path, payload + + +def test_urlcache_dry_run_writes_nothing(url_json, db_path): + path, payload = url_json + assert migrate_urlcache.main(["--json", str(path), "--db", db_path]) == 0 + assert not os.path.exists(db_path) + assert path.exists() + + +def test_urlcache_migrates_every_record(url_json, db_path): + path, payload = url_json + assert migrate_urlcache.main(["--json", str(path), "--db", db_path, "--apply"]) == 0 + conn = db.connect(db_path, readonly=True) + assert conn.execute("SELECT count(*) c FROM remote_file").fetchone()["c"] == len(payload) + for url, mtime in list(payload.items())[:5]: + assert db.get_remote_mtime(conn, url) == mtime + conn.close() + + +def test_urlcache_leaves_the_json_alone_unless_asked(url_json, db_path): + path, _ = url_json + migrate_urlcache.main(["--json", str(path), "--db", db_path, "--apply"]) + assert path.exists() + assert not os.path.exists(str(path) + ".bak") + + +def test_urlcache_backs_up_only_after_verifying(url_json, db_path): + path, _ = url_json + migrate_urlcache.main( + ["--json", str(path), "--db", db_path, "--apply", "--backup-json"] + ) + assert not path.exists() + assert os.path.exists(str(path) + ".bak") + + +def test_urlcache_is_idempotent(url_json, db_path): + path, payload = url_json + migrate_urlcache.main(["--json", str(path), "--db", db_path, "--apply"]) + migrate_urlcache.main(["--json", str(path), "--db", db_path, "--apply"]) + conn = db.connect(db_path, readonly=True) + assert conn.execute("SELECT count(*) c FROM remote_file").fetchone()["c"] == len(payload) + conn.close() + + +def test_urlcache_refuses_malformed_records(tmp_path, db_path): + path = tmp_path / "bad.json" + path.write_text(json.dumps({"https://example.test/a": "not a timestamp"})) + assert migrate_urlcache.main(["--json", str(path), "--db", db_path, "--apply"]) == 1 + assert not os.path.exists(db_path) + + +def test_urlcache_refuses_a_non_object(tmp_path, db_path): + path = tmp_path / "list.json" + path.write_text(json.dumps(["a", "b"])) + assert migrate_urlcache.main(["--json", str(path), "--db", db_path, "--apply"]) == 1 + + +def test_urlcache_reports_a_missing_source(tmp_path, db_path): + assert migrate_urlcache.main(["--json", str(tmp_path / "nope.json"), "--db", db_path]) == 1 + + +# -------------------------------------------------------------------- un-rename + + +@pytest.fixture +def labelled_archive(archive): + """An archive in the pre-migration state: suffixed files plus error plots.""" + made = [] + for index in range(4): + made.append(make_frame_file(archive, index=index, label="f")) + for index in range(4, 6): + name, path = make_frame_file(archive, index=index, label="e") + made.append((name, path)) + plot = os.path.join(os.path.dirname(path), name.error_plot_name()) + open(plot, "w").write("diagnostic") + make_frame_file(archive, index=6, label=None) # never processed + return archive, made + + +def test_unrename_dry_run_changes_nothing(labelled_archive, db_path): + archive, _ = labelled_archive + before = sorted(os.listdir(_day_dir(archive))) + assert migrate_unrename.main(["--root", str(archive), "--db", db_path]) == 0 + assert sorted(os.listdir(_day_dir(archive))) == before + assert not os.path.exists(db_path) + + +def _day_dir(archive): + return os.path.join(str(archive), "goes16/l2/data/suvi-l2-ci171/2024/05/10") + + +def test_unrename_strips_both_suffixes(labelled_archive, db_path): + archive, _ = labelled_archive + assert migrate_unrename.main(["--root", str(archive), "--db", db_path, "--apply"]) == 0 + names = os.listdir(_day_dir(archive)) + assert not [n for n in names if n.endswith("_f.fits") or n.endswith("_e.fits")] + assert len([n for n in names if n.endswith(".fits")]) == 7 + + +def test_unrename_deletes_diagnostic_plots(labelled_archive, db_path): + archive, _ = labelled_archive + migrate_unrename.main(["--root", str(archive), "--db", db_path, "--apply"]) + assert not [n for n in os.listdir(_day_dir(archive)) if n.endswith(".jpg")] + + +def test_unrename_keeps_plots_when_asked(labelled_archive, db_path): + archive, _ = labelled_archive + migrate_unrename.main( + ["--root", str(archive), "--db", db_path, "--apply", "--keep-plots"] + ) + assert len([n for n in os.listdir(_day_dir(archive)) if n.endswith(".jpg")]) == 2 + + +def test_unrename_exports_labels_before_renaming(labelled_archive, db_path, tmp_path): + archive, _ = labelled_archive + export = tmp_path / "labels.csv.gz" + migrate_unrename.main( + ["--root", str(archive), "--db", db_path, "--apply", "--export", str(export)] + ) + assert export.exists() + with gzip.open(export, "rt") as handle: + rows = list(csv.DictReader(handle)) + assert len(rows) == 6 + assert sum(1 for r in rows if r["legacy_label"] == "f") == 4 + assert sum(1 for r in rows if r["legacy_label"] == "e") == 2 + # The export must name the *restored* path, so it can be rejoined after renaming. + for row in rows: + assert not row["relpath"].endswith("_f.fits") + assert os.path.exists(paths.abspath(row["relpath"], str(archive))) + + +def test_unrename_records_verdicts_in_the_index(labelled_archive, db_path): + archive, _ = labelled_archive + migrate_unrename.main(["--root", str(archive), "--db", db_path, "--apply"]) + conn = db.connect(db_path, readonly=True) + run_id = db.latest_run_id(conn, migrate_unrename.LEGACY_RUN_NAME) + assert run_id is not None + rows = conn.execute( + "SELECT verdict, count(*) c FROM detection WHERE run_id = ? GROUP BY verdict", + (run_id,), + ).fetchall() + counts = {row["verdict"]: row["c"] for row in rows} + assert counts == {"good": 4, "bad": 2} + conn.close() + + +def test_unrename_is_idempotent(labelled_archive, db_path): + archive, _ = labelled_archive + migrate_unrename.main(["--root", str(archive), "--db", db_path, "--apply"]) + names_after_first = sorted(os.listdir(_day_dir(archive))) + assert migrate_unrename.main(["--root", str(archive), "--db", db_path, "--apply"]) == 0 + assert sorted(os.listdir(_day_dir(archive))) == names_after_first + + +def test_unrename_skips_collisions_rather_than_overwriting(archive, db_path): + """If both the suffixed and unsuffixed names exist, neither may be destroyed.""" + name, suffixed = make_frame_file(archive, index=0, label="f") + plain = os.path.join(os.path.dirname(suffixed), name.filename()) + write_fits(plain, solar_disc(size=32, radius=10, peak=1.0)) + plain_bytes = open(plain, "rb").read() + + migrate_unrename.main(["--root", str(archive), "--db", db_path, "--apply"]) + assert os.path.exists(suffixed), "the suffixed file was destroyed" + assert open(plain, "rb").read() == plain_bytes, "the existing file was overwritten" + + +def test_unrename_can_be_scoped_to_one_band(archive, db_path): + make_frame_file(archive, wavelength=171, index=0, label="f") + make_frame_file(archive, wavelength=304, index=0, label="f") + migrate_unrename.main( + ["--root", str(archive), "--db", db_path, "--apply", "--wavelength", "171"] + ) + band_171 = os.listdir(_day_dir(archive)) + band_304 = os.listdir( + os.path.join(str(archive), "goes16/l2/data/suvi-l2-ci304/2024/05/10") + ) + assert not [n for n in band_171 if n.endswith("_f.fits")] + assert [n for n in band_304 if n.endswith("_f.fits")] + + +def test_unrename_on_an_already_clean_archive(archive, db_path): + make_frame_file(archive, index=0, label=None) + assert migrate_unrename.main(["--root", str(archive), "--db", db_path, "--apply"]) == 0 + + +def test_unrename_reports_a_missing_root(tmp_path, db_path): + assert migrate_unrename.main( + ["--root", str(tmp_path / "nope"), "--db", db_path, "--apply"] + ) == 1 + + +def test_walk_archive_finds_only_frames(labelled_archive): + archive, _ = labelled_archive + found = list(migrate_unrename.walk_archive(str(archive), (16,), (171,))) + assert len(found) == 7 + assert all(name is not None for _, _, name in found) + + +def test_walk_surfaces_unreadable_directories(archive, monkeypatch): + """A filesystem error must not look like an empty archive. + + os.walk ignores errors by default, so a transient ENFILE on the shared mount + made the walk yield nothing and the migration report "already migrated" -- + silently skipping hundreds of thousands of files it should have renamed. + """ + make_frame_file(archive, index=0, label="f") + + def explode(path): + raise OSError(23, "Too many open files in system") + + monkeypatch.setattr(os, "scandir", explode) + with pytest.raises(OSError): + list(migrate_unrename.walk_archive(str(archive), (16,), (171,))) + + +def test_unrename_records_and_skips_completed_chunks(labelled_archive, db_path): + """A resume must skip finished band-years without touching the filesystem. + + Re-walking completed subtrees on every resume is what made the migration + traverse the archive several times over, and traversal is what exhausts the + mount's file handles. + """ + archive, _ = labelled_archive + migrate_unrename.main(["--root", str(archive), "--db", db_path, "--apply"]) + + conn = db.connect(db_path, readonly=True) + assert db.get_meta(conn, migrate_unrename.chunk_key(171, 2024)) == "done" + conn.close() + + def explode(*args, **kwargs): + raise AssertionError("resume re-read the filesystem for a completed chunk") + + import unittest.mock as mock + + with mock.patch.object(migrate_unrename, "walk_archive", explode): + assert migrate_unrename.main( + ["--root", str(archive), "--db", db_path, "--apply"] + ) == 0 + + +def test_recheck_forces_a_completed_chunk_to_be_re_examined(labelled_archive, db_path): + archive, _ = labelled_archive + migrate_unrename.main(["--root", str(archive), "--db", db_path, "--apply"]) + calls = [] + real = migrate_unrename.walk_archive + + def spy(*args, **kwargs): + calls.append(args) + return real(*args, **kwargs) + + import unittest.mock as mock + + with mock.patch.object(migrate_unrename, "walk_archive", spy): + migrate_unrename.main( + ["--root", str(archive), "--db", db_path, "--apply", "--recheck"] + ) + assert calls, "--recheck should re-examine the archive" + + +def test_a_chunk_with_failed_renames_is_not_marked_done(labelled_archive, db_path, monkeypatch): + """Marking a partially failed chunk complete would make a resume skip real work.""" + archive, _ = labelled_archive + real_rename = os.rename + state = {"n": 0} + + def flaky(src, dst): + state["n"] += 1 + if state["n"] == 2: + raise OSError(23, "Too many open files in system") + return real_rename(src, dst) + + monkeypatch.setattr(os, "rename", flaky) + migrate_unrename.main(["--root", str(archive), "--db", db_path, "--apply"]) + + conn = db.connect(db_path, readonly=True) + assert db.get_meta(conn, migrate_unrename.chunk_key(171, 2024)) is None + conn.close() + + +# --------------------------------------------------------------- review sampler + + +def test_judge_timestamp_classifies_every_band(archive): + """The sampler's unit is a frame, but a composite needs all six bands.""" + import filter_FITS + + band_paths = {} + for band in (94, 131, 171, 195, 284, 304): + name, path = make_frame_file(archive, wavelength=band, index=0) + band_paths[band] = path + + satellite, timestamp, verdicts = filter_FITS._judge_timestamp((16, name.t_start, band_paths)) + assert satellite == 16 + assert set(verdicts) == {94, 131, 171, 195, 284, 304} + for verdict, reason, scores in verdicts.values(): + assert verdict in ("good", "bad") + assert isinstance(scores, dict) + + +def test_judge_timestamp_flags_a_damaged_band(archive): + """A blank band must be flagged, and the others left alone.""" + import numpy as np + + import filter_FITS + + band_paths = {} + for band in (94, 131, 171, 195, 284, 304): + name, path = make_frame_file(archive, wavelength=band, index=0) + band_paths[band] = path + # Replace one band with a frame carrying no signal at all. + write_fits(band_paths[171], np.zeros((32, 32), dtype="float32")) + + _, _, verdicts = filter_FITS._judge_timestamp((16, name.t_start, band_paths)) + assert verdicts[171][0] == "bad" + assert verdicts[171][1], "a flagged band must carry a reason" diff --git a/tests/test_paths.py b/tests/test_paths.py new file mode 100644 index 0000000..e48770d --- /dev/null +++ b/tests/test_paths.py @@ -0,0 +1,115 @@ +import datetime +import os + +import pytest + +from suvi import paths + +GOOD = "dr_suvi-l2-ci171_g16_s20240510T000000Z_e20240510T000400Z_v1-0-2.fits" + + +def test_parses_a_real_filename(): + frame = paths.parse_frame_filename(GOOD) + assert frame.satellite == 16 + assert frame.wavelength == 171 + assert frame.version == "1-0-2" + assert frame.label is None + assert frame.t_end - frame.t_start == paths.CADENCE + assert frame.t_start == int( + datetime.datetime(2024, 5, 10, tzinfo=datetime.timezone.utc).timestamp() + ) + + +@pytest.mark.parametrize("label", [None, "f", "e"]) +def test_filename_round_trips(label): + name = GOOD if label is None else GOOD.replace(".fits", f"_{label}.fits") + frame = paths.parse_frame_filename(name) + assert frame.label == label + assert frame.filename(label) == name + + +def test_double_suffix_keeps_outermost_label(): + """Re-processing an already-labelled file produced '..._f_f.fits' in the archive.""" + frame = paths.parse_frame_filename(GOOD.replace(".fits", "_f_e.fits")) + assert frame.label == "e" + assert frame.filename() == GOOD + + +def test_relpath_files_under_the_start_date(): + """A frame opening at 23:58 belongs to that day, not the day it ends on.""" + name = "dr_suvi-l2-ci304_g18_s20241231T235800Z_e20250101T000200Z_v1-0-2.fits" + frame = paths.parse_frame_filename(name) + assert frame.relpath() == ( + "goes18/l2/data/suvi-l2-ci304/2024/12/31/" + name + ) + + +def test_slot_identifies_the_observation(): + frame = paths.parse_frame_filename(GOOD) + labelled = paths.parse_frame_filename(GOOD.replace(".fits", "_e.fits")) + assert frame.slot == labelled.slot == (16, 171, frame.t_start) + + +def test_error_plot_name_matches_legacy_filter_output(): + frame = paths.parse_frame_filename(GOOD) + assert frame.error_plot_name() == GOOD.replace(".fits", "_e.jpg") + + +@pytest.mark.parametrize( + "name", + [ + "", + "x" * 300, + "README.txt", + "dr_suvi-l2-ci171_g16_s20240510T000000Z_e20240510T000400Z_v1-0-2.fits.bak", + "dr_suvi-l2-ci999_g16_s20240510T000000Z_e20240510T000400Z_v1-0-2.fits", # bad band + "dr_suvi-l2-ci171_g17_s20240510T000000Z_e20240510T000400Z_v1-0-2.fits", # excluded sat + "dr_suvi-l2-ci171_g16_s20241310T000000Z_e20240510T000400Z_v1-0-2.fits", # month 13 + "dr_suvi-l2-ci171_g16_s20240510T000400Z_e20240510T000000Z_v1-0-2.fits", # end <= start + "dr_suvi-l2-ci171_g16_s20240510T000000Z_e20240510T000400Z_v1-0-2_x.fits", # bad label + "../dr_suvi-l2-ci171_g16_s20240510T000000Z_e20240510T000400Z_v1-0-2.fits", + "sub/dr_suvi-l2-ci171_g16_s20240510T000000Z_e20240510T000400Z_v1-0-2.fits", + ], +) +def test_rejects_malformed_names(name): + assert paths.parse_frame_filename(name) is None + + +def test_rejects_non_string_input(): + for value in (None, 42, b"bytes", ["list"]): + assert paths.parse_frame_filename(value) is None + + +def test_filename_rejects_unknown_label(): + frame = paths.parse_frame_filename(GOOD) + with pytest.raises(ValueError): + frame.filename("q") + + +def test_abspath_resolves_within_the_root(tmp_path): + frame = paths.parse_frame_filename(GOOD) + resolved = paths.abspath(frame.relpath(), root=str(tmp_path)) + assert resolved.startswith(str(tmp_path) + os.sep) + assert resolved.endswith(GOOD) + + +@pytest.mark.parametrize("hostile", ["../../etc/passwd", "goes16/../../../etc/passwd"]) +def test_abspath_refuses_to_escape_the_root(tmp_path, hostile): + with pytest.raises(ValueError): + paths.abspath(hostile, root=str(tmp_path)) + + +def test_data_root_honours_the_environment(archive): + assert paths.data_root() == str(archive) + + +def test_default_db_path_honours_the_environment(monkeypatch, tmp_path): + target = tmp_path / "scratch.sqlite" + monkeypatch.setenv("SUVI_DB", str(target)) + assert paths.default_db_path() == str(target) + + +def test_wavelength_dirs_covers_every_combination(tmp_path): + dirs = paths.wavelength_dirs(str(tmp_path)) + assert len(dirs) == len(paths.SATELLITES) * len(paths.WAVELENGTHS) + assert len(set(dirs)) == len(dirs) diff --git a/tests/test_vfs.py b/tests/test_vfs.py new file mode 100644 index 0000000..28d9da5 --- /dev/null +++ b/tests/test_vfs.py @@ -0,0 +1,215 @@ +import pytest + +"""Tests for reclaiming the virtiofs mount's file handles. +These guard a mechanism that exists only because of an environment defect: the +mount's daemon holds a host descriptor per inode the guest looks up, and runs out. +The properties that matter are that reclaim is *actually verified* rather than +assumed, and that the expensive fallback is used only when it must be. +""" + + +from suvi import vfs + + +def test_drop_caches_reports_failure_without_root(monkeypatch): + """Must return False rather than raise, so the fallback can run.""" + def refuse(*args, **kwargs): + raise PermissionError(13, "Permission denied") + + monkeypatch.setattr("builtins.open", refuse) + assert vfs.drop_caches() is False + + +def test_reclaimable_kb_reads_the_slab(): + assert vfs.reclaimable_kb() > 0 # real /proc/meminfo + + +def test_reclaimable_kb_survives_a_broken_meminfo(monkeypatch): + monkeypatch.setattr("builtins.open", lambda *a, **k: (_ for _ in ()).throw(OSError())) + assert vfs.reclaimable_kb() == 0 + + +def test_release_prefers_drop_caches(monkeypatch): + """The fallback allocates tens of GiB; it must not run when drop_caches works.""" + monkeypatch.setattr(vfs, "drop_caches", lambda: True) + + def explode(*args, **kwargs): + raise AssertionError("allocated memory despite drop_caches succeeding") + + monkeypatch.setattr(vfs, "bytearray", explode, raising=False) + assert vfs.release_handles() is True + + +def test_fallback_reports_failure_when_nothing_is_reclaimed(monkeypatch): + monkeypatch.setattr(vfs, "drop_caches", lambda: False) + monkeypatch.setattr(vfs, "reclaimable_kb", lambda: 1000) + assert vfs.release_handles(budget_gib=1, floor_kb=0) is False + + +def test_release_gives_up_when_the_slab_cannot_be_read(monkeypatch): + monkeypatch.setattr(vfs, "drop_caches", lambda: False) + monkeypatch.setattr(vfs, "reclaimable_kb", lambda: 0) + assert vfs.release_handles() is False + + +# --------------------------------------------------------------- memory bounding + + +def test_available_gib_reads_meminfo(): + assert vfs.available_gib() > 0 + + +def test_pressure_stops_at_the_reserve_not_a_fraction(monkeypatch): + """Capping at a fraction of available memory reclaims nothing. + + The kernel only frees slab under genuine pressure, so leaving 40% headroom + means it never triggers: an earlier version allocated 14 GiB of a 23 GiB + allowance and freed zero, costing the machine the memory for no benefit. + Pressure must run down to a fixed reserve instead. + """ + monkeypatch.setattr(vfs, "drop_caches", lambda: False) + monkeypatch.setattr(vfs, "reclaimable_kb", lambda: 1000) # never yields + available = [20.0] + monkeypatch.setattr(vfs, "available_gib", lambda: available[0]) + allocated = [] + real = bytearray + + def fake(n): + allocated.append(n) + available[0] -= 1.0 # each GiB claimed reduces what is available + return real(1) + + monkeypatch.setattr(vfs, "bytearray", fake, raising=False) + assert vfs.release_handles(reserve_gib=2.0, floor_kb=0) is False + # Pushed from 20 GiB down to the 2 GiB reserve, not stopped at 60%. + assert 17 <= len(allocated) <= 19, f"allocated {len(allocated)} GiB" + + +def test_pressure_stops_as_soon_as_the_slab_yields(monkeypatch): + monkeypatch.setattr(vfs, "drop_caches", lambda: False) + readings = iter([1000, 1000, 1000, 400]) + monkeypatch.setattr(vfs, "reclaimable_kb", lambda: next(readings)) + monkeypatch.setattr(vfs, "available_gib", lambda: 20.0) + allocated = [] + real = bytearray + monkeypatch.setattr( + vfs, "bytearray", lambda n: allocated.append(n) or real(1), raising=False + ) + assert vfs.release_handles(floor_kb=0) is True + assert len(allocated) == 2, "kept allocating after the kernel gave ground" + + +# ------------------------------------------------------------------- Reliever + + +def test_reliever_fires_on_the_interval(monkeypatch): + calls = [] + monkeypatch.setattr(vfs, "release_handles", lambda *a, **k: calls.append(1) or True) + reliever = vfs.Reliever(interval=10) + for _ in range(9): + reliever.tick() + assert calls == [] + reliever.tick() + assert len(calls) == 1 + + +def test_reliever_keeps_firing_across_a_long_run(monkeypatch): + calls = [] + monkeypatch.setattr(vfs, "release_handles", lambda *a, **k: calls.append(1) or True) + reliever = vfs.Reliever(interval=10) + for _ in range(100): + reliever.tick() + assert len(calls) == 10, "stopped reclaiming partway through a long traversal" + assert reliever.releases == 10 + + +def test_reliever_counts_batched_items(monkeypatch): + """Callers processing six bands per timestamp tick by six, not one.""" + calls = [] + monkeypatch.setattr(vfs, "release_handles", lambda *a, **k: calls.append(1) or True) + reliever = vfs.Reliever(interval=10) + reliever.tick(6) + reliever.tick(6) + assert len(calls) == 1 + assert reliever.seen == 12 + + +def test_reliever_can_be_disabled(monkeypatch): + def explode(*a, **k): + raise AssertionError("reclaimed while disabled") + + monkeypatch.setattr(vfs, "release_handles", explode) + reliever = vfs.Reliever(interval=1, enabled=False) + for _ in range(50): + reliever.tick() + assert reliever.releases == 0 + + +def test_reliever_reclaims_when_the_run_ends(monkeypatch): + """Reclaiming only during a run leaves the machine loaded once it finishes. + + The inodes a traversal cached stay pinned -- and with them the host's file + handles -- until something else forces reclaim. That is why the system was + still saturated with nothing of ours running. + """ + calls = [] + monkeypatch.setattr(vfs, "release_handles", lambda *a, **k: calls.append(1) or True) + reliever = vfs.Reliever(interval=1000, on_exit=False) + reliever.tick(5) + assert calls == [], "should not have reclaimed mid-run yet" + reliever.finish() + assert len(calls) == 1 + + +def test_finish_is_idempotent(monkeypatch): + calls = [] + monkeypatch.setattr(vfs, "release_handles", lambda *a, **k: calls.append(1) or True) + reliever = vfs.Reliever(interval=1000, on_exit=False) + reliever.tick() + reliever.finish() + reliever.finish() + reliever.finish() + assert len(calls) == 1 + + +def test_finish_does_nothing_if_no_work_happened(monkeypatch): + def explode(*a, **k): + raise AssertionError("reclaimed despite processing nothing") + + monkeypatch.setattr(vfs, "release_handles", explode) + vfs.Reliever(on_exit=False).finish() + + +def test_reliever_works_as_a_context_manager(monkeypatch): + calls = [] + monkeypatch.setattr(vfs, "release_handles", lambda *a, **k: calls.append(1) or True) + with vfs.Reliever(interval=1000, on_exit=False) as reliever: + reliever.tick(3) + assert calls == [] + assert len(calls) == 1 + + +def test_context_manager_reclaims_even_when_the_body_raises(monkeypatch): + calls = [] + monkeypatch.setattr(vfs, "release_handles", lambda *a, **k: calls.append(1) or True) + with pytest.raises(ValueError): + with vfs.Reliever(interval=1000, on_exit=False) as reliever: + reliever.tick() + raise ValueError("boom") + assert len(calls) == 1, "a failed job must still hand its handles back" + + +def test_no_pressure_is_applied_when_the_slab_is_already_small(monkeypatch): + """Below the floor the slab is live processes' working set, not archive inodes. + + Squeezing the machine to shave a few hundred megabytes off it costs far more + than it gains -- and during a long job this runs every RELIEF_INTERVAL items. + """ + monkeypatch.setattr(vfs, "drop_caches", lambda: False) + monkeypatch.setattr(vfs, "reclaimable_kb", lambda: 500 * 1024) # 0.5 GiB + + def explode(*a, **k): + raise AssertionError("applied memory pressure below the floor") + + monkeypatch.setattr(vfs, "bytearray", explode, raising=False) + assert vfs.release_handles() is True