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