| from __future__ import annotations |
|
|
| from pathlib import Path |
| from typing import Any |
|
|
| import numpy as np |
| import pandas as pd |
| import torch |
| from torch.utils.data import Dataset |
|
|
| from .utils import ( |
| future_offsets, |
| load_yaml, |
| parse_time, |
| read_json, |
| source_existing_dat_path, |
| source_meta_path, |
| status_columns, |
| time_offsets, |
| ) |
|
|
|
|
| class SourceStore: |
| def __init__(self, output_root: str | Path, source: str): |
| self.output_root = Path(output_root) |
| self.source = source |
| self.dat_path = source_existing_dat_path(self.output_root, source) |
| self.meta_path = source_meta_path(self.dat_path) |
| self.meta = read_json(self.meta_path) |
| self.dtype = np.dtype(self.meta["dtype"]) |
| self.row_shape = tuple(int(v) for v in self.meta["row_shape"]) |
| self.row_count = int(self.meta["row_count"]) |
| self.shape = (self.row_count, *self.row_shape) |
| self._mm = np.memmap(self.dat_path, dtype=self.dtype, mode="r", shape=self.shape) |
|
|
| def get(self, idx: int) -> np.ndarray: |
| idx = int(idx) |
| if idx < 0 or idx >= self.row_count: |
| raise IndexError(f"{self.source} index out of range: {idx} / {self.row_count}") |
| return np.asarray(self._mm[idx], dtype=np.float32) |
|
|
|
|
| class FastDataset(Dataset): |
| def __init__( |
| self, |
| config: str | Path | dict[str, Any], |
| split: str | None = None, |
| required_inputs: list[str] | None = None, |
| required_labels: list[str] | None = None, |
| dtype: torch.dtype = torch.float32, |
| ): |
| self.config = load_yaml(config) if not isinstance(config, dict) else dict(config) |
| self.output_root = Path(self.config["output_root"]) |
| self.catalog_path = Path(self.config.get("catalog_path", self.output_root / "catalog.csv")) |
| self.catalog = pd.read_csv(self.catalog_path, dtype={"timestamp": str}) |
| self.catalog["timestamp"] = self.catalog["timestamp"].astype(str) |
|
|
| self.required_inputs = list(required_inputs if required_inputs is not None else self.config.get("required_inputs", ["concat"])) |
| self.required_labels = list(required_labels if required_labels is not None else self.config.get("required_labels", ["ci"])) |
| self.dtype = dtype |
|
|
| window = self.config.get("input_window", {}) |
| self.input_offsets = list(window.get("offset_minutes") or time_offsets( |
| int(window.get("past_minutes", 50)), |
| int(window.get("interval_minutes", 10)), |
| )) |
| bt_cfg = self.config.get("labels", {}).get("bt", {}) |
| self.bt_source = str(bt_cfg.get("source", bt_cfg.get("source_name", "bt"))) |
| self.use_bt_mask = bool(bt_cfg.get("use_mask", True)) |
| self.bt_mask_source = str(bt_cfg.get("mask_source", "bt_mask")) |
| self.bt_mask_regions = tuple(str(region).lower() for region in bt_cfg.get("mask_regions", ["object", "context"])) |
| valid_mask_regions = {"object", "context"} |
| unknown_mask_regions = sorted(set(self.bt_mask_regions) - valid_mask_regions) |
| if unknown_mask_regions: |
| raise ValueError(f"unsupported BT mask regions: {unknown_mask_regions}") |
| if self.use_bt_mask and not self.bt_mask_regions: |
| raise ValueError("labels.bt.mask_regions must contain at least one of: object, context") |
| self.bt_offsets = list(bt_cfg.get("lead_minutes") or future_offsets( |
| int(bt_cfg.get("future_minutes", 60)), |
| int(bt_cfg.get("interval_minutes", 10)), |
| )) |
|
|
| self.time_to_row = {str(row.timestamp): i for i, row in self.catalog.iterrows()} |
| store_sources = {self._catalog_source_name(source) for source in self.required_inputs + self.required_labels} |
| if "bt" in self.required_labels and self.use_bt_mask: |
| store_sources.add(self.bt_mask_source) |
| available_sources = self._available_sources() |
| self.stores = { |
| source: SourceStore(self.output_root, source) |
| for source in store_sources |
| if source in available_sources |
| } |
| self.samples = self._build_samples(split) |
|
|
| def _catalog_source_name(self, name: str) -> str: |
| if name == "ci": |
| return "ci_hard" |
| if name == "bt": |
| return self.bt_source |
| return name |
|
|
| def _available_sources(self) -> set[str]: |
| out = set() |
| for col in self.catalog.columns: |
| if col.endswith("_status"): |
| out.add(col[: -len("_status")]) |
| return out |
|
|
| def _filter_split(self, times: list[str], split: str | None) -> list[str]: |
| if not split: |
| return times |
| ranges = self.config.get("splits", {}).get(split) |
| if ranges is None: |
| raise KeyError(f"split {split!r} not found in config.splits") |
| parsed_ranges = [] |
| for item in ranges: |
| start, end = (item["start"], item["end"]) if isinstance(item, dict) else item |
| parsed_ranges.append((parse_time(start), parse_time(end))) |
| kept = [] |
| for ts in times: |
| t = parse_time(ts) |
| if any(start <= t <= end for start, end in parsed_ranges): |
| kept.append(ts) |
| return kept |
|
|
| def _is_ok(self, timestamp: str, source: str) -> bool: |
| row_idx = self.time_to_row.get(timestamp) |
| if row_idx is None: |
| return False |
| idx_col, status_col = status_columns(source) |
| row = self.catalog.iloc[row_idx] |
| return status_col in row and row[status_col] == "ok" and not pd.isna(row[idx_col]) |
|
|
| def _idx_for(self, timestamp: str, source: str) -> int: |
| row_idx = self.time_to_row[timestamp] |
| idx_col, status_col = status_columns(source) |
| row = self.catalog.iloc[row_idx] |
| if row[status_col] != "ok" or pd.isna(row[idx_col]): |
| raise RuntimeError(f"catalog says source not ok: {timestamp} {source} status={row[status_col]}") |
| return int(row[idx_col]) |
|
|
| def _input_times(self, sample_time: str) -> list[str]: |
| base = parse_time(sample_time) |
| return [(base + np.timedelta64(int(offset), "m")).strftime("%Y%m%d%H%M") for offset in self.input_offsets] |
|
|
| def _bt_times(self, sample_time: str) -> list[str]: |
| base = parse_time(sample_time) |
| return [(base + np.timedelta64(int(offset), "m")).strftime("%Y%m%d%H%M") for offset in self.bt_offsets] |
|
|
| def _valid_sample(self, sample_time: str) -> bool: |
| for source in self.required_inputs: |
| catalog_source = self._catalog_source_name(source) |
| for ts in self._input_times(sample_time): |
| if not self._is_ok(ts, catalog_source): |
| return False |
| for label in self.required_labels: |
| catalog_source = self._catalog_source_name(label) |
| if label == "bt": |
| if self.use_bt_mask and not self._is_ok(sample_time, self.bt_mask_source): |
| return False |
| for ts in self._bt_times(sample_time): |
| if not self._is_ok(ts, catalog_source): |
| return False |
| else: |
| if not self._is_ok(sample_time, catalog_source): |
| return False |
| return True |
|
|
| def _build_samples(self, split: str | None) -> list[str]: |
| times = sorted(self.catalog["timestamp"].astype(str).tolist()) |
| times = self._filter_split(times, split) |
| return [ts for ts in times if self._valid_sample(ts)] |
|
|
| def __len__(self) -> int: |
| return len(self.samples) |
|
|
| def _load_input(self, source: str, sample_time: str) -> torch.Tensor: |
| catalog_source = self._catalog_source_name(source) |
| store = self.stores[catalog_source] |
| rows = [store.get(self._idx_for(ts, catalog_source)) for ts in self._input_times(sample_time)] |
| return torch.from_numpy(np.stack(rows, axis=0).astype(np.float32, copy=False)).to(dtype=self.dtype) |
|
|
| def _load_label(self, label: str, sample_time: str) -> torch.Tensor: |
| catalog_source = self._catalog_source_name(label) |
| store = self.stores[catalog_source] |
| if label == "bt": |
| rows = [store.get(self._idx_for(ts, catalog_source)) for ts in self._bt_times(sample_time)] |
| bt = np.stack(rows, axis=0).astype(np.float32, copy=False) |
| if self.use_bt_mask: |
| packed_mask = self.stores[self.bt_mask_source].get( |
| self._idx_for(sample_time, self.bt_mask_source) |
| ) |
| bt = self.apply_bt_mask(bt, packed_mask) |
| return torch.from_numpy(bt.astype(np.float32, copy=False)).to(dtype=self.dtype) |
| if label == "bt_mask": |
| arr = store.get(self._idx_for(sample_time, catalog_source)) |
| return torch.from_numpy(arr.astype(np.float32, copy=False)).to(dtype=self.dtype) |
| arr = store.get(self._idx_for(sample_time, catalog_source)) |
| return torch.from_numpy(arr.astype(np.float32, copy=False)).to(dtype=self.dtype) |
|
|
| def decode_bt_mask(self, mask_row: np.ndarray) -> np.ndarray: |
| """Return one boolean BT-loss mask per configured forecast lead.""" |
| if self.bt_mask_source == "bt_mask": |
| static_mask = np.asarray(mask_row, dtype=bool) |
| if static_mask.ndim != 2: |
| raise ValueError(f"bt_mask row must be 2-D, got {static_mask.shape}") |
| return np.repeat(static_mask[np.newaxis, ...], len(self.bt_offsets), axis=0) |
|
|
| if self.bt_mask_source != "bt_lineage_mask": |
| raise ValueError(f"unsupported labels.bt.mask_source: {self.bt_mask_source!r}") |
|
|
| packed = np.asarray(mask_row, dtype=np.uint16) |
| if packed.ndim != 2: |
| raise ValueError(f"bt_lineage_mask row must be 2-D, got {packed.shape}") |
|
|
| meta = self.stores[self.bt_mask_source].meta |
| object_bits = {int(lead): int(bit) for lead, bit in dict(meta.get("object_bits", {})).items()} |
| context_bits = {int(lead): int(bit) for lead, bit in dict(meta.get("context_bits", {})).items()} |
| missing_leads = [int(lead) for lead in self.bt_offsets if int(lead) not in object_bits or int(lead) not in context_bits] |
| if missing_leads: |
| raise ValueError( |
| f"bt_lineage_mask metadata has no object/context bits for leads {missing_leads}" |
| ) |
|
|
| lead_masks = [] |
| for lead in self.bt_offsets: |
| valid = np.zeros(packed.shape, dtype=bool) |
| if "object" in self.bt_mask_regions: |
| valid |= (packed & np.uint16(1 << object_bits[int(lead)])) != 0 |
| if "context" in self.bt_mask_regions: |
| valid |= (packed & np.uint16(1 << context_bits[int(lead)])) != 0 |
| lead_masks.append(valid) |
| return np.stack(lead_masks, axis=0) |
|
|
| def apply_bt_mask(self, bt: np.ndarray, mask_row: np.ndarray) -> np.ndarray: |
| """Replace pixels outside the configured lead-specific BT mask with NaN.""" |
| bt = np.asarray(bt, dtype=np.float32) |
| valid = self.decode_bt_mask(mask_row) |
| expected_shape = (bt.shape[0], *bt.shape[-2:]) |
| if tuple(valid.shape) != expected_shape: |
| raise ValueError(f"BT/mask shape mismatch: bt={bt.shape}, mask={valid.shape}") |
| return np.where(valid[:, np.newaxis, ...], bt, np.nan) |
|
|
| def __getitem__(self, idx: int) -> dict[str, Any]: |
| sample_time = self.samples[int(idx)] |
| return { |
| "inputs": {name: self._load_input(name, sample_time) for name in self.required_inputs}, |
| "labels": {name: self._load_label(name, sample_time) for name in self.required_labels}, |
| "time": sample_time, |
| } |
|
|
