jeremy/noaa-goes-visualization
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Rewrote processing pipeline including bad-image detection with multiple overlapping methods (header + disc based detection seems best for now).

Browse source
This commit is contained in:
Jeremy Karst 2026-09-01 22:36:13 -04:00
parent 830a7eb469
commit 22ac9c68b1
39 changed files with 11131 additions and 1631 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
.gitignore vendored
View file

@ -1 +1,4 @@
.vscode .vscode
.venv
__pycache__/
*.pyc

1298
bench.py Normal file
View file

File diff suppressed because it is too large Load diff

73
ffmpeg_video.py
View file

@ -1,8 +1,9 @@
import argparse
import subprocess import subprocess
import os import os
import shutil import sys
import datetime import datetime
import time import shutil as _shutil
import calendar import calendar
import numpy as np import numpy as np
@ -10,15 +11,44 @@ from sortedcontainers import SortedDict
import tqdm import tqdm
from PIL import Image from PIL import Image
years = [2024] sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
sources = ["goes16", "goes18"]
ffmpeg_path = r"..\ffmpeg.exe"
for year in years: 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: for source in sources:
path_to_images = f"..\\composite\\{source}" path_to_images = args.images or os.path.join(
output_file = f"..\\{source}_{year}_nofilt.mp4" os.path.dirname(paths.data_root()), "composite", source)
interp_file = f"..\\{source}_{year}_interp_nofilt.mp4" 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()) 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()) stoptime = calendar.timegm(datetime.datetime(year+1, 1, 1, tzinfo=datetime.timezone.utc).timetuple())
@ -51,7 +81,9 @@ for year in years:
assert np.sum((unique % interval) > 0) == 0 # Ensure all our timestamps align perfectly with our interval assert np.sum((unique % interval) > 0) == 0 # Ensure all our timestamps align perfectly with our interval
prevtime = files_by_time.peekitem(0)[0] - 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)): 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 framejump = (t - prevtime) // interval
if framejump < max_frame_interp: if framejump < max_frame_interp:
for i in range(framejump - 1): # Fill black frames to fill gaps, these will later be interpolated with ffmpeg for i in range(framejump - 1): # Fill black frames to fill gaps, these will later be interpolated with ffmpeg
@ -69,5 +101,26 @@ for year in years:
print(command_line) print(command_line)
pipe = subprocess.Popen(command_line, shell=True, stdout=subprocess.PIPE).stdout pipe = subprocess.Popen(command_line, shell=True, stdout=subprocess.PIPE).stdout
output = pipe.read().decode() pipe.read()
pipe.close() 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: <root>/composite/<source>)")
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)

942
filter_FITS.py
View file

@ -1,281 +1,733 @@
#!/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 os
import sys
import time
from collections import defaultdict from collections import defaultdict
from multiprocessing import Queue, Process from multiprocessing import Pool
import re
import traceback
import datetime
import calendar
import tqdm
import numpy as np 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(): sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
""" Set the priority of the process to lowest possible."""
import sys from suvi import db, detectors, fitsio, index, paths, vfs
try:
sys.getwindowsversion()
except AttributeError: def parse_int_list(text):
isWindows = False 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: else:
isWindows = True verdict = detectors.disc_v1(image, slot[1], values)
verdicts[name] = (verdict, (time.perf_counter_ns() - started) // 1000)
if isWindows: return frame_id, verdicts
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) 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
measurement_names = ["094", "131", "171", "195", "284", "304"] frame_detectors = [n for n in wanted if n in detectors.FRAME_DETECTORS]
thresholds = [0.050, 0.10 , 1.00, 1.40, 1.0, 2.50] series_detectors = [n for n in wanted if n in detectors.SERIES_DETECTORS]
max_center_skew = 7 if series_detectors:
expected_dims = 1280 print(f"NOTE: {series_detectors} need a whole series in memory; use bench.py "
half_dims = expected_dims // 2 "for those, or run them here one band at a time via --wavelengths.")
valid_radii = [383, 394]
avg_radius = (valid_radii[0] + valid_radii[1]) // 2
max_radius_error = 80
ratio_above_thresh_max = 0.4 where, params = ["1=1"], []
ratio_above_thresh_min = 0.07 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
def filter_fits(work_queue): runs = {
lowpriority() name: db.create_detector_run(conn, name, {"driver": "filter_FITS"}, notes=args.notes)
idealcircle = np.zeros((expected_dims // 2, expected_dims // 2)) for name in frame_detectors
cv.circle(idealcircle, (expected_dims // 4, expected_dims // 4), avg_radius // 2, 1, -1) }
idealcircle_axis = np.average(idealcircle, 0) print(f"Running {frame_detectors} over {len(rows)} frames with {args.workers} workers")
while True:
try: jobs = [
job = work_queue.get() (
if job is None: row["id"],
return (row["satellite"], row["wavelength"], row["t_start"]),
measurement = job.split("dr_suvi-l2-")[1].split("_")[0][2:] paths.abspath(row["path"], root),
idx = measurement_names.index(measurement) frozenset(frame_detectors),
try: )
data = fits.getdata(job) for row in rows
except IndexError: # This can happen with a blank or corrupt HDU in the .fits file ]
new_name = job.split(".fits")[0] + "_e.fits" started = time.time()
os.rename(job, new_name) 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 continue
assert data.shape[0] == expected_dims results = []
assert data.shape[1] == expected_dims reliever = vfs.Reliever(label=f"calibrate {wavelength}A")
data = cv.resize(data, dsize=(expected_dims // 2, expected_dims // 2), interpolation=cv.INTER_LINEAR) with Pool(args.workers, initializer=fitsio.quiet_astropy) as pool:
filtered_data = np.copy(data) for measured in pool.imap_unordered(_calibration_measure, candidates, 8):
above_thresh_indexes = data > thresholds[idx] reliever.tick()
filtered_data[above_thresh_indexes] = thresholds[idx] if measured is None:
filtered_data /= thresholds[idx] continue
ratio_above_thresh = np.count_nonzero(above_thresh_indexes) / data.shape[0] / data.shape[1] if any(v is None for v in measured.values()):
results.append(measured) # keep, so unmeasurable rate is visible
if ratio_above_thresh > ratio_above_thresh_max: continue
print(f"Exceeded ratio_above_thresh_max, possible data corruption in file: {job}") results.append(measured)
new_name = job.split(".fits")[0] + "_e.fits" if len(results) >= args.samples_per_band:
os.rename(job, new_name) break
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')
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 continue
if ratio_above_thresh < ratio_above_thresh_min: bounds = {}
print(f"Below ratio_above_thresh_min, possible bad data in file: {job}") print(f"{wavelength:>5}A n={len(usable)} unmeasurable={unmeasurable}")
new_name = job.split(".fits")[0] + "_e.fits" for name in ("limb_contrast", "radius_ratio", "limb_width"):
os.rename(job, new_name) 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]
plt.figure(figsize=[10.24, 7.68]) lower = low - half if side in ("lower", "both") else floor
plt.title(f"Below ratio_above_thresh_min [{ratio_above_thresh:0.2f}]") upper = high + half if side in ("upper", "both") else ceiling
plt.imshow(filtered_data, cmap='jet') # Keep each limit on the physical side of zero. A negative floor on a
plt.savefig(job.split(".fits")[0] + "_e.jpg") # non-negative quantity is not a loose rule, it is no rule at all.
plt.close('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
continue if not table:
print("\nNothing calibrated.")
return 1
# Decide whether the data is valid based on solar disc placement and levels. print("\n" + "=" * 78)
xavg = np.average(filtered_data, 0) print("Paste into DISC_BOUNDS in suvi/detectors.py:\n")
yavg = np.average(filtered_data, 1) print("DISC_BOUNDS = {")
axis_indices = list(range(half_dims)) for wavelength, bounds in sorted(table.items()):
print(f" {wavelength}: {{")
# Center Estimation for name, (lower, upper) in bounds.items():
weighted_avg_center_x = np.average(axis_indices, 0, xavg) print(f' "{name}": ({lower}, {upper}),')
weighted_avg_center_y = np.average(axis_indices, 0, yavg) print(" },")
print("}")
good_center = True if args.json:
if (half_dims // 2 - weighted_avg_center_x > max_center_skew) or (half_dims // 2 - weighted_avg_center_y > max_center_skew): with open(args.json, "w") as handle:
print(f"Could not find centered solar disc in file: {job}") json.dump(table, handle, indent=2)
good_center = False print(f"\nAlso written to {args.json}")
return 0
# 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: # --------------------------------------------------------------------------- sample
# We have validated this file, rename it appropriately
new_name = job.split(".fits")[0] + "_f.fits" #: The production combination: a frame is bad if either detector says so.
os.rename(job, new_name) 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: else:
new_name = job.split(".fits")[0] + "_e.fits" disc = detectors.disc_v1(image, band, values)
os.rename(job, new_name)
# Plot Results bad = [v for v in (header, disc) if v.verdict == "bad"]
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]) 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
axes[0][0].axis('off')
axes[0][1].plot(axis_indices, xavg) def _render_sample(job):
axes[0][1].plot(axis_indices, idealcircle_axis) """Render one annotated composite for review."""
# axes[0][1].fill_between(axis_indices, xavg, idealcircle_axis, hatch="//", edgecolor="red", facecolor="none") import merger_FITS
from PIL import ImageDraw, ImageFont
axes[1][0].plot(yavg, axis_indices) satellite, timestamp, band_paths, caption, out_path = job
axes[1][0].plot(idealcircle_axis, axis_indices) 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_xaxis()
axes[1][0].invert_yaxis() axes[1][0].invert_yaxis()
axes[1][1].imshow(shown, aspect="auto", vmin=0, vmax=1)
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[0][1].sharex(axes[1][1])
axes[1][0].sharey(axes[1][1]) axes[1][0].sharey(axes[1][1])
axes[1][1].yaxis.set_ticks_position("right") figure.suptitle(f"{os.path.basename(row['path'])}\n{row['reason']}", fontsize=8)
figure.savefig(
plt.savefig(job.split(".fits")[0] + "_e.jpg") os.path.join(args.out, os.path.basename(row["path"]).replace(".fits", ".jpg")),
plt.close('all') dpi=110,
)
plt.close(figure)
return 0
except KeyboardInterrupt: # ------------------------------------------------------------------------------ cli
return
except Exception as e:
print(f"Error on file: {job} - {e}") def build_parser():
traceback.print_exception(e) 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__": if __name__ == "__main__":
stored_fits_dirs = [r"..\Data\goes16\l2\data\suvi-l2-ci094\2024", sys.exit(main())
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()

168
make_comparison_video.py Normal file
View file

@ -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())

270
merger.py
View file

@ -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")

346
merger_FITS.py
View file

@ -1,10 +1,9 @@
import os import os
import time import sys
import calendar import calendar
import datetime import datetime
from collections import defaultdict from collections import defaultdict
from multiprocessing import Queue, Process from multiprocessing import Queue, Process
import re
import warnings import warnings
import queue import queue
import traceback import traceback
@ -16,6 +15,8 @@ from matplotlib import pyplot as plt
from astropy.io import fits from astropy.io import fits
import palettable import palettable
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
def lowpriority(): def lowpriority():
""" Set the priority of the process to below-normal.""" """ Set the priority of the process to below-normal."""
@ -268,55 +269,70 @@ def hsl_to_rgb(img):
else: else:
return np.stack([r, g, b], -1) return np.stack([r, g, b], -1)
def generate_composite(work_queue, result_queue): #: Band labels drawn on the composite, and the order `composite_from_arrays`
lowpriority() #: expects its six arrays in.
cmaps =[palettable.cmocean.sequential.Ice_5.mpl_colormap, 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.Ice_20.mpl_colormap,
palettable.cmocean.sequential.Turbid_5_r.mpl_colormap, palettable.cmocean.sequential.Turbid_5_r.mpl_colormap,
palettable.cmocean.sequential.Turbid_20_r.mpl_colormap, palettable.cmocean.sequential.Turbid_20_r.mpl_colormap,
plt.colormaps.get_cmap('gist_heat'), plt.colormaps.get_cmap('gist_heat'),
plt.colormaps.get_cmap('afmhot')] plt.colormaps.get_cmap('afmhot')]
# Modify colormaps to start at perfect black (when they otherwise start at very dark colors) # Start these maps at perfect black rather than a very dark colour.
for cmi in range(0,4): for cmi in range(0, 4):
for c in ['red','green','blue']: for c in ['red', 'green', 'blue']:
for i in range(3): for i in range(3):
cmaps[cmi]._segmentdata[c][0][i] = 0.0 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. font_path = os.path.join(os.path.dirname(os.path.abspath(__file__)),
vmins = [0.050, 0.05, 00.100, 00.10, 00.100, 00.1] "OpenSans-Regular.ttf")
vmaxs = [8.000, 8.00, 20.000, 30.00, 40.000, 90.0] _render_cache.update(
gammas = [0.375, 0.40, 00.425, 00.45, 00.475, 00.5] cmaps=cmaps,
trimx = 64 fnt1=ImageFont.truetype(font_path, size=24),
trimy = 100 fnt2=ImageFont.truetype(font_path, size=16),
fnt1 = ImageFont.truetype("OpenSans-Regular.ttf", size = 24) )
fnt2 = ImageFont.truetype("OpenSans-Regular.ttf", size = 16) return _render_cache
while True:
try:
job = work_queue.get() def composite_from_arrays(arrays, timestamp):
if job is None: """Render one composite from six in-memory bands. Returns a PIL Image.
result_queue.cancel_join_thread()
return `arrays` are the six 1280x1280 radiance arrays in `image_names` order.
files_this_timestamp, timestamp, processed_images_dir, synthetic_data = job """
synthetic_filepath = os.path.join(processed_images_dir, f"Composite-{int(timestamp)}_s.jpg") assets = _render_assets()
normal_filepath = os.path.join(processed_images_dir, f"Composite-{int(timestamp)}.jpg") cmaps, fnt1, fnt2 = assets["cmaps"], assets["fnt1"], assets["fnt2"]
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 = [] base_imgs = []
for i in range(6): 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 = 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) 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]) 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.figure(image_names[i])
# plt.imshow(base_imgs[-1]) # plt.imshow(base_imgs[-1])
@ -420,6 +436,47 @@ def generate_composite(work_queue, result_queue):
xdimtxtoff = composite_image_data.shape[1] - 44 xdimtxtoff = composite_image_data.shape[1] - 44
ydimtxtoff = i//2*new_dimy + new_dimy / 2.0 - 14 ydimtxtoff = i//2*new_dimy + new_dimy / 2.0 - 14
ImageDraw.Draw(img).text((xdimtxtoff, ydimtxtoff), image_names[i], font = fnt2) 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) img.save(filepath, quality = 95)
result_queue.put(("Created", timestamp)) result_queue.put(("Created", timestamp))
@ -430,34 +487,70 @@ def generate_composite(work_queue, result_queue):
result_queue.put((e, timestamp)) result_queue.put((e, timestamp))
image_names = ["094A", "131A", "171A", "195A", "284A", "304A"]
if __name__ == "__main__": if __name__ == "__main__":
stored_fits_dirs = [r"..\Data\goes16\l2\data", r"..\Data\goes18\l2\data"] import argparse
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$" from suvi import db, paths
nworkers = 16 parser = argparse.ArgumentParser(
max_time_gap = 3 description="Render SUVI composite images from indexed FITS frames.",
fill_missing_data = False )
ignore_errors = True 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()
if ignore_errors: def parse_day(text):
regex_filename = r"dr_suvi-l2-ci\d{3}_g(16|18)_s\S*.fits$" 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] file_prefixes = ["dr_suvi-l2-ci" + n[:-1] for n in image_names]
# Testing work_queue = Queue(maxsize=nworkers)
# stored_fits_dirs = [r"..\fits_test_2024"]
work_queue = Queue(maxsize = nworkers)
result_queue = Queue() result_queue = Queue()
workers = [] workers = []
for i in range(nworkers): for i in range(nworkers):
p = Process(target = generate_composite, args = (work_queue, result_queue), daemon=True) p = Process(target=generate_composite, args=(work_queue, result_queue), daemon=True)
p.start() p.start()
workers.append(p) workers.append(p)
@ -465,61 +558,54 @@ if __name__ == "__main__":
ncreated = 0 ncreated = 0
nexists = 0 nexists = 0
nfailed = 0 nfailed = 0
black_image = Image.fromarray(np.zeros((1080, 1920, 3), dtype='uint8'))
try: try:
for stored_fits_dir, processed_images_dir in zip(stored_fits_dirs, processed_images_dirs): for satellite in satellites:
files_by_timestamp = defaultdict(list) processed_images_dir = os.path.join(composite_root, f"goes{satellite}")
num_found_files = 0
os.makedirs(processed_images_dir, exist_ok=True) 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.") # 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: 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 continue
# Collect completed jobs and record completion status
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: while True:
try: try:
result = result_queue.get_nowait() result = result_queue.get_nowait()
@ -528,40 +614,32 @@ if __name__ == "__main__":
elif result[0] == "Created": elif result[0] == "Created":
ncreated += 1 ncreated += 1
else: else:
print(f"A worker encountered an exception on job {result[1]}: {result[0]}") print(f"A worker failed on job {result[1]}: {result[0]}")
nfailed += 1 nfailed += 1
except queue.Empty: except queue.Empty:
break break
# Submit new jobs files_this_timestamp = sorted(files_by_timestamp[timestamp])
files_this_timestamp = files_by_timestamp[timestamp] if len(files_this_timestamp) == 6:
files_this_timestamp = sorted(files_this_timestamp) last_good_files = files_this_timestamp
if (not len(files_this_timestamp) == 6): last_good_file_times = [timestamp] * 6
if fill_missing_data: work_queue.put((files_this_timestamp, timestamp, processed_images_dir, False))
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.") elif fill_missing_data:
files_for_job = [] files_for_job = []
for i, prefix in enumerate(file_prefixes): for i, prefix in enumerate(file_prefixes):
found = False found = next(
for f in files_this_timestamp: (f for f in files_this_timestamp
filename = os.path.split(f)[-1] if os.path.basename(f).startswith(prefix)),
if filename.startswith(prefix): None,
files_for_job.append(f) )
last_good_files[i] = f if found is not None:
files_for_job.append(found)
last_good_files[i] = found
last_good_file_times[i] = timestamp last_good_file_times[i] = timestamp
found = True elif (timestamp - last_good_file_times[i]) // interval <= max_time_gap:
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]) 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: if len(files_for_job) == 6:
work_queue.put((files_for_job, timestamp, processed_images_dir, True)) 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: except KeyboardInterrupt:
print("Finishing current jobs and exiting") print("Finishing current jobs and exiting")
@ -569,11 +647,11 @@ if __name__ == "__main__":
for _ in range(nworkers): for _ in range(nworkers):
try: try:
work_queue.put(None, timeout=10.0) work_queue.put(None, timeout=10.0)
except: except Exception:
break break
for w in workers: for w in workers:
# w.join(10.0)
w.join() w.join()
conn.close()
print(f"Created {ncreated} | Already had {nexists} | Failed {nfailed}") print(f"Created {ncreated} | Already had {nexists} | Failed {nfailed}")

366
migrate_unrename.py Normal file
View file

@ -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: <root>/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())

125
migrate_urlcache.py Normal file
View file

@ -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())

148
puller.py
View file

@ -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'(?<=<a href=")([^ ?]*)(?=">)' # 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}")

170
puller_fits.py
View file

@ -1,11 +1,11 @@
import os import os
import sys
import urllib.request import urllib.request
import urllib.parse import urllib.parse
import re import re
import time import time
import random import random
import datetime import datetime
import json
from threading import Thread from threading import Thread
import queue import queue
import math import math
@ -14,12 +14,13 @@ from queue import Empty
import tqdm 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/" directory_url = r"https://data.ngdc.noaa.gov/platforms/solar-space-observing-satellites/goes/"
stored_images_dir = os.path.abspath(os.path.join("..", "Data")) stored_images_dir = paths.data_root()
# 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"] 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 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 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 ndownloadworkers = 3 # Be nice to the servers, this value is how many threads will be downloading files at the same time
@ -138,15 +139,28 @@ def file_download_worker(download_work_queue, download_result_queue, attempt_cou
if __name__ == "__main__": if __name__ == "__main__":
file_info_cache = {} import argparse
try:
print(f"Attempting to load file records from cache: {file_database_path}") parser = argparse.ArgumentParser(
with open(file_database_path, 'r') as f: description="Mirror NOAA GOES SUVI L2 FITS into the local archive.",
file_info_cache = json.loads(f.read()) )
print(f"File records loaded from cache: {len(file_info_cache)} records found.") parser.add_argument("--db", default=None, help="SQLite index (default: $SUVI_DB)")
except Exception as e: parser.add_argument("--root", default=None, help="archive root (default: $SUVI_DATA_ROOT)")
print(f"Load failed, starting with empty cache") parser.add_argument("--interval", type=int, default=fetch_interval,
file_info_cache = {} 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_work_queue = queue.Queue()
query_result_queue = queue.Queue() query_result_queue = queue.Queue()
@ -162,6 +176,33 @@ if __name__ == "__main__":
t.start() t.start()
workers.append(t) 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: try:
while True: while True:
fetched_image_count = 0 fetched_image_count = 0
@ -177,78 +218,62 @@ if __name__ == "__main__":
r_success, r_url, r_t = download_result_queue.get_nowait() r_success, r_url, r_t = download_result_queue.get_nowait()
if r_success: if r_success:
fetched_image_count += 1 fetched_image_count += 1
file_info_cache[r_url] = r_t 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: else:
failed_image_count += 1 failed_image_count += 1
except queue.Empty: except queue.Empty:
break break
flush_records()
file_portion_of_link = l[urllen:] file_portion_of_link = l[urllen:]
filepath = os.path.join(stored_images_dir, file_portion_of_link.replace("/", os.sep)) 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 recorded_mtime = db.get_remote_mtime(conn, l)
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: # The archive no longer carries _f/_e suffixes, so a frame is either
f.write(json.dumps(file_info_cache)) # 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}") print(f"Downloaded {fetched_image_count} | Already had {already_had_image_count} | Failed {failed_image_count}")
if fetch_interval > 0: if fetch_interval > 0:
print(f"Run complete!, sleeping for {fetch_interval} seconds.") 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) time.sleep(fetch_interval)
else: else:
print("Run complete!, exiting...") print("Run complete!, exiting...")
break break
except KeyboardInterrupt: except KeyboardInterrupt:
print("Saving file database and shutting down.") print("Saving download records and shutting down.")
except Empty: except Empty:
print("Work Complete, Shutting down...") print("Work Complete, Shutting down...")
except Exception as e:
print(f"Unhandled Exception during run: {e}")
print("Shutting down")
for _ in range(nfetchworkers): for _ in range(nfetchworkers):
try: try:
query_work_queue.put(None) query_work_queue.put(None)
except: except Exception:
break break
time.sleep(1) time.sleep(1)
@ -261,13 +286,22 @@ if __name__ == "__main__":
for _ in range(ndownloadworkers): for _ in range(ndownloadworkers):
try: try:
download_work_queue.put(None, timeout=5.0) download_work_queue.put(None, timeout=5.0)
except: except Exception:
break break
with open(file_database_path, 'w') as f: # Drain any results the workers finished after the main loop ended.
f.write(json.dumps(file_info_cache)) 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...") print("Waiting for workers to shutdown...")
for w in workers: for w in workers:
w.join(5.0) w.join(5.0)

5
pytest.ini Normal file
View file

@ -0,0 +1,5 @@
[pytest]
testpaths = tests
filterwarnings =
# astropy warns about the deliberately damaged files these tests construct.
ignore::astropy.utils.exceptions.AstropyUserWarning

69
reclaim.py Normal file
View file

@ -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())

14
requirements.txt
View file

@ -1,8 +1,12 @@
astropy astropy
scikit-image
tqdm
numpy
opencv-python
palettable
matplotlib matplotlib
numpy
opencv-python-headless
palettable
pillow pillow
scikit-image
sortedcontainers
tqdm
# Development
pytest

6
suvi/__init__.py Normal file
View file

@ -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.
"""

397
suvi/cases.py Normal file
View file

@ -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
)

280
suvi/corruptions.py Normal file
View file

@ -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()),
}

615
suvi/db.py Normal file
View file

@ -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)

742
suvi/detectors.py Normal file
View file

@ -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)

349
suvi/fillers.py Normal file
View file

@ -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.

294
suvi/fitsio.py Normal file
View file

@ -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"

253
suvi/index.py Normal file
View file

@ -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)

354
suvi/metrics.py Normal file
View file

@ -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())
},
}

185
suvi/paths.py Normal file
View file

@ -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<wavelength>\d{3})"
r"_g(?P<satellite>\d{2})"
r"_s(?P<t_start>\d{8}T\d{6})Z"
r"_e(?P<t_end>\d{8}T\d{6})Z"
r"_v(?P<version>[0-9]+(?:-[0-9]+)*)"
r"(?P<labels>(?:_[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
]

187
suvi/vfs.py Normal file
View file

@ -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

115
tests/conftest.py Normal file
View file

@ -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())

381
tests/test_cases.py Normal file
View file

@ -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

165
tests/test_comparison_video.py Normal file
View file

@ -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

90
tests/test_composite.py Normal file
View file

@ -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

209
tests/test_corruptions.py Normal file
View file

@ -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

285
tests/test_db.py Normal file
View file

@ -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()

476
tests/test_detectors.py Normal file
View file

@ -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

215
tests/test_fillers.py Normal file
View file

@ -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

229
tests/test_fitsio.py Normal file
View file

@ -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()

401
tests/test_index.py Normal file
View file

@ -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()

350
tests/test_metrics.py Normal file
View file

@ -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

355
tests/test_migrations.py Normal file
View file

@ -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"

115
tests/test_paths.py Normal file
View file

@ -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)

215
tests/test_vfs.py Normal file
View file

@ -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