398 lines
16 KiB
Python
398 lines
16 KiB
Python
|
"""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
|
||
|
|
)
|