File size: 11,679 Bytes
76d61a0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 | 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,
}
|