| from __future__ import annotations |
|
|
| from pathlib import Path |
| from typing import Any |
|
|
| import numpy as np |
| import xarray as xr |
| from scipy.ndimage import binary_dilation |
|
|
| from .input_data import ConcatInput, L2AIIInput |
| from .utils import future_offsets, inv_zscore, load_stats, open_memmap, parse_time, zscore |
|
|
|
|
| class CILabel: |
| source_name = "ci" |
|
|
| def __init__(self, config: dict[str, Any]): |
| self.config = dict(config) |
| self.root = Path(self.config["root"]) |
| self.expected_hw = self.config.get("shape_hw") |
| if self.expected_hw is not None: |
| self.expected_hw = tuple(int(v) for v in self.expected_hw) |
| self._xr_open_kwargs = dict(decode_cf=False, mask_and_scale=False, decode_times=False) |
|
|
| @property |
| def row_shape(self) -> tuple[int, int]: |
| if not self.expected_hw: |
| raise ValueError("ci.shape_hw must be configured or inferred before memmap use") |
| return tuple(self.expected_hw) |
|
|
| def path_for_time(self, timestamp: str): |
| ts = parse_time(timestamp).strftime("%Y%m%d%H%M") |
| day = ts[:8] |
| nc = self.root / day / f"{ts}_label.nc" |
| if nc.exists(): |
| return nc |
| npz = self.root / day / f"{ts}_layers.npz" |
| if npz.exists(): |
| return npz |
| return None |
|
|
| def load_label(self, timestamp: str) -> np.ndarray: |
| path = self.path_for_time(timestamp) |
| if path is None: |
| raise FileNotFoundError(f"CI label not found for {timestamp}") |
| if path.suffix == ".npz": |
| with np.load(path, allow_pickle=False) as data: |
| if "masks" not in data: |
| raise KeyError(f"'masks' key not found in {path}") |
| masks = data["masks"] |
| if masks.ndim != 3: |
| raise ValueError(f"npz masks must be (N,H,W), got {masks.shape}: {path}") |
| label = (masks > 0).any(axis=0).astype(np.uint8) |
| else: |
| with xr.open_dataset(path, **self._xr_open_kwargs) as ds: |
| var = list(ds.data_vars)[0] |
| arr = ds[var].values |
| if arr.ndim != 2: |
| raise ValueError(f"CI label must be 2D, got {arr.shape}: {path}") |
| label = (~np.isnan(arr)).astype(np.uint8) |
|
|
| if self.expected_hw is None: |
| self.expected_hw = tuple(int(v) for v in label.shape) |
| if tuple(label.shape) != tuple(self.expected_hw): |
| raise ValueError(f"CI shape mismatch: {label.shape} != {self.expected_hw}") |
| return label |
|
|
| @staticmethod |
| def smooth(label: np.ndarray, base: float = 0.5, radius: int = 3) -> np.ndarray: |
| base = float(base) |
| radius = int(radius) |
| if not (0.0 < base < 1.0): |
| raise ValueError(f"base must be between 0 and 1, got {base}") |
| if radius < 0: |
| raise ValueError(f"radius must be >= 0, got {radius}") |
|
|
| positive = np.asarray(label) > 0.5 |
| smoothed = np.zeros(positive.shape, dtype=np.float32) |
| smoothed[positive] = 1.0 |
| if radius == 0 or not positive.any(): |
| return smoothed |
|
|
| structure = np.ones((3, 3), dtype=bool) |
| prev = positive |
| for distance in range(1, radius + 1): |
| dilated = binary_dilation(positive, structure=structure, iterations=distance) |
| ring = dilated & ~prev |
| smoothed[ring] = base ** distance |
| prev = dilated |
| return smoothed |
|
|
| def open_memmap(self, dat_path: str | Path, n_rows: int, dtype: str = "uint8", mode: str = "r") -> np.memmap: |
| return open_memmap(dat_path, dtype, (int(n_rows), *self.row_shape), mode=mode) |
|
|
| def load_memmap_row(self, dat_path: str | Path, row_idx: int, n_rows: int, dtype: str = "uint8") -> np.ndarray: |
| mm = self.open_memmap(dat_path, n_rows=n_rows, dtype=dtype, mode="r") |
| return np.asarray(mm[int(row_idx)]) |
|
|
|
|
| class BTLabel: |
| source_name = "bt" |
| file_dtype = "float16" |
|
|
| DEFAULT_CONDITIONS = { |
| "KI": {"op": ">=", "value": 30.0}, |
| "LI": {"op": "<=", "value": -2.0}, |
| "SI": {"op": "<=", "value": 2.0}, |
| "CAPE": {"op": ">=", "value": 500.0}, |
| "TTI": {"op": ">=", "value": 42.0}, |
| } |
|
|
| def __init__(self, config: dict[str, Any], concat_input: ConcatInput, l2_input: L2AIIInput): |
| self.config = dict(config) |
| self.concat_input = concat_input |
| self.l2_input = l2_input |
| self.future_minutes = int(self.config.get("future_minutes", 60)) |
| self.interval_minutes = int(self.config.get("interval_minutes", 10)) |
| self.lead_minutes = list(self.config.get("lead_minutes") or future_offsets(self.future_minutes, self.interval_minutes)) |
| self.var_name = str(self.config.get("var_name", "ir105")) |
| self.mask_var_name = str(self.config.get("mask_var_name", self.var_name)) |
| self.bt_threshold_k = float(self.config.get("bt_threshold_k", 233.0)) |
| self.expansion_km = float(self.config.get("expansion_km", 50.0)) |
| self.pixel_size_km = float(self.config.get("pixel_size_km", 2.0)) |
| self.conditions = dict(self.DEFAULT_CONDITIONS) |
| self.conditions.update(self.config.get("aii_conditions", {})) |
| self.stats_path = self.config.get("stats_path") or self.concat_input.stats_path |
| self.stats = load_stats(self.stats_path) |
| self.normalization = str(self.config.get("normalization", "zscore")).lower() |
| self.eps = float(self.config.get("eps", 1e-6)) |
| self.expected_hw = self.config.get("shape_hw") or self.concat_input.expected_hw |
| if self.expected_hw is not None: |
| self.expected_hw = tuple(int(v) for v in self.expected_hw) |
|
|
| @property |
| def row_shape(self) -> tuple[int, int, int]: |
| if not self.expected_hw: |
| raise ValueError("BT shape_hw must be configured or inferred before memmap use") |
| h, w = self.expected_hw |
| return (1, h, w) |
|
|
| @property |
| def mask_row_shape(self) -> tuple[int, int]: |
| if not self.expected_hw: |
| raise ValueError("BT shape_hw must be configured or inferred before memmap use") |
| return tuple(self.expected_hw) |
|
|
| def _condition_mask(self, arr: np.ndarray, op: str, value: float) -> np.ndarray: |
| arr = np.asarray(arr, dtype=np.float32) |
| finite = np.isfinite(arr) |
| if op == ">=": |
| return finite & (arr >= float(value)) |
| if op == ">": |
| return finite & (arr > float(value)) |
| if op == "<=": |
| return finite & (arr <= float(value)) |
| if op == "<": |
| return finite & (arr < float(value)) |
| raise ValueError(f"unsupported condition op: {op}") |
|
|
| def circular_footprint(self) -> np.ndarray: |
| radius_pixels = int(np.ceil(self.expansion_km / self.pixel_size_km)) |
| yy, xx = np.ogrid[-radius_pixels : radius_pixels + 1, -radius_pixels : radius_pixels + 1] |
| distance_km = np.sqrt(xx**2 + yy**2) * self.pixel_size_km |
| return distance_km <= self.expansion_km |
|
|
| def build_mask(self, timestamp: str, shape_hw: tuple[int, int]) -> np.ndarray: |
| aii = self.l2_input.load_raw_dict(timestamp) |
| masks = [] |
| for name, rule in self.conditions.items(): |
| if name not in aii: |
| raise KeyError(f"BT mask requires L2 variable {name!r}") |
| arr = np.asarray(aii[name], dtype=np.float32) |
| if tuple(arr.shape) != tuple(shape_hw): |
| raise ValueError(f"AII/BT shape mismatch for {name}: {arr.shape} != {shape_hw}") |
| masks.append(self._condition_mask(arr, str(rule["op"]), float(rule["value"]))) |
| aii_mask = np.logical_and.reduce(masks) |
|
|
| start_data = self.concat_input.load_raw_dict(timestamp) |
| start_bt = self.concat_input.get_variable(start_data, self.mask_var_name) |
| if tuple(start_bt.shape) != tuple(shape_hw): |
| raise ValueError(f"start mask-variable shape mismatch: {start_bt.shape} != {shape_hw}") |
|
|
| cold = np.isfinite(start_bt) & (start_bt <= self.bt_threshold_k) |
| expanded = binary_dilation(cold, structure=self.circular_footprint()) if self.expansion_km > 0 else cold |
| return aii_mask & ~expanded |
|
|
| def normalize_bt(self, arr: np.ndarray) -> np.ndarray: |
| arr = np.asarray(arr, dtype=np.float32) |
| if self.normalization in {"none", "raw", "false"}: |
| return arr |
| if self.normalization == "zscore": |
| stat = self.stats[self.var_name] |
| return zscore(arr, stat["mean"], stat["std"], self.eps) |
| raise ValueError(f"unsupported BT normalization: {self.normalization}") |
|
|
| def denormalize(self, arr: np.ndarray) -> np.ndarray: |
| arr = np.asarray(arr, dtype=np.float32) |
| if self.normalization in {"none", "raw", "false"}: |
| return arr |
| stat = self.stats[self.var_name] |
| return inv_zscore(arr, stat["mean"], stat["std"], self.eps) |
|
|
| def load_label(self, timestamp: str) -> np.ndarray: |
| data = self.concat_input.load_raw_dict(timestamp) |
| bt = self.concat_input.get_variable(data, self.var_name) |
| if self.expected_hw is None: |
| self.expected_hw = tuple(int(v) for v in bt.shape) |
| if tuple(bt.shape) != tuple(self.expected_hw): |
| raise ValueError(f"BT shape mismatch: {timestamp} {bt.shape} != {self.expected_hw}") |
|
|
| normalized = self.normalize_bt(bt) |
| return normalized[np.newaxis, ...].astype(np.float32, copy=False) |
|
|
| def load_mask(self, timestamp: str) -> np.ndarray: |
| data = self.concat_input.load_raw_dict(timestamp) |
| bt = self.concat_input.get_variable(data, self.var_name) |
| if self.expected_hw is None: |
| self.expected_hw = tuple(int(v) for v in bt.shape) |
| if tuple(bt.shape) != tuple(self.expected_hw): |
| raise ValueError(f"BT mask shape mismatch: {timestamp} {bt.shape} != {self.expected_hw}") |
| return self.build_mask(timestamp, tuple(self.expected_hw)).astype(np.uint8, copy=False) |
|
|
| def open_memmap(self, dat_path: str | Path, n_rows: int, dtype: str | None = None, mode: str = "r") -> np.memmap: |
| return open_memmap(dat_path, dtype or self.file_dtype, (int(n_rows), *self.row_shape), mode=mode) |
|
|
| def load_memmap_row(self, dat_path: str | Path, row_idx: int, n_rows: int, dtype: str | None = None) -> np.ndarray: |
| mm = self.open_memmap(dat_path, n_rows=n_rows, dtype=dtype or self.file_dtype, mode="r") |
| return np.asarray(mm[int(row_idx)], dtype=np.float32) |
|
|
