diff --git a/config/evaluate/config_zarr2cf.yaml b/config/evaluate/config_zarr2cf.yaml index 75677858b0..c686970a9c 100644 --- a/config/evaluate/config_zarr2cf.yaml +++ b/config/evaluate/config_zarr2cf.yaml @@ -89,6 +89,15 @@ variables: wg_unit: Pa std_unit: Pa level_type: sfc + tp: + var: tp + long: total_precipitation + std: precipitation_amount + wg_unit: m + std_unit: m + level_type: sfc + accumulate: true + paramId: 228 tp_imerg_0: var: tp_imerg_0 long: imerg_total_precipitation @@ -96,6 +105,8 @@ variables: wg_unit: m std_unit: m level_type: sfc + paramId: 228 + accumulate: true coordinates: diff --git a/link_results.sh b/link_results.sh new file mode 100644 index 0000000000..7ba034ffe8 --- /dev/null +++ b/link_results.sh @@ -0,0 +1,65 @@ +#!/usr/bin/env bash +# Create incrementally-numbered symlinks in results/xq6/ that point to the +# validation zip files inside all results/xq6XXXXX source directories. +# +# Usage: +# bash link_results.sh # actually create the symlinks +# bash link_results.sh --dry-run # only print what would be done +set -euo pipefail + +DRY_RUN=0 +if [[ "${1:-}" == "--dry-run" || "${1:-}" == "-n" ]]; then + DRY_RUN=1 + echo "== DRY RUN: no changes will be made ==" >&2 +fi + +#results/xq6 +BASE="results" +RUNID="xq6" +TARGET="${BASE}/${RUNID}" +CHKPT="validation_chkpt00000" + +if [[ "${DRY_RUN}" -eq 0 ]]; then + mkdir -p "${TARGET}" +fi + +# Global counter for the target rank numbering. +dest=0 + +# Discover all source run directories (e.g. results/xq650000 ... results/xq650015), +# sorted numerically, while skipping the aggregated target directory itself. +for src_dir in $(ls -d "${BASE}/${RUNID}"[0-9]*/ 2>/dev/null | sort); do + src_dir="${src_dir%/}" # strip trailing slash + src_name="$(basename "${src_dir}")" + + if [[ "${src_name}" == "${RUNID}" ]]; then + continue # never link the target into itself + fi + if [[ ! -d "${src_dir}" ]]; then + echo "Skipping missing directory: ${src_dir}" >&2 + continue + fi + + # Iterate source ranks in order (0,1,2,3,...). + for src_file in $(ls "${src_dir}/${CHKPT}_rank"*.zip 2>/dev/null | sort); do + dest_name=$(printf "%s_rank%04d.zip" "${CHKPT}" "${dest}") + dest_path="${TARGET}/${dest_name}" + + # Relative link target as seen from inside ${TARGET} (results/xq6/), + # hence the leading ../ to step out of xq6/ back into results/. + rel_src="../${src_name}/$(basename "${src_file}")" + + if [[ -e "${dest_path}" || -L "${dest_path}" ]]; then + echo "Exists, skipping: ${dest_path}" >&2 + else + if [[ "${DRY_RUN}" -eq 1 ]]; then + echo "[dry-run] ln -s ${rel_src} ${dest_path}" + else + ln -s "${rel_src}" "${dest_path}" + echo "ln -s ${rel_src} ${dest_path}" + fi + fi + + dest=$((dest + 1)) + done +done diff --git a/packages/common/src/weathergen/common/io.py b/packages/common/src/weathergen/common/io.py index 685cd5d8e7..718a0e1b6f 100644 --- a/packages/common/src/weathergen/common/io.py +++ b/packages/common/src/weathergen/common/io.py @@ -442,6 +442,14 @@ def _get_datasets(self, key: ItemKey): for name, dataset in group.groups() } + def _has_group(self, item: ItemKey) -> bool: + """Check if an output item exists in this store.""" + assert self.data_root is not None, "ZarrIO must be opened before accessing data." + try: + return self.data_root.get(item.path) is not None + except KeyError: + return False + def _get_group(self, item: ItemKey, create: bool) -> zarr.Array | zarr.Group: assert self.data_root is not None, "ZarrIO must be opened before accessing data." if create: @@ -492,16 +500,31 @@ def _create_dataset(self, group: zarr.Group, name: str, array: NDArray): @functools.cached_property def forecast_offset(self) -> int: - fstep0_datasets = self._get_datasets(self.example_key) + key = self.example_key + if not self._has_group(key): + # No fstep 0 group at all => no targets at fstep 0 => offset 1. + _logger.debug(f"No group at {key.path}, inferring forecast_offset=1.") + return 1 + fstep0_datasets = self._get_datasets(key) return ItemKey._infer_forecast_offset(fstep0_datasets) @functools.cached_property def example_key(self) -> ItemKey: try: sample, example_sample = next(self.data_root.groups()) + # Find the first stream that has prediction/target data (not just source) + for stream, example_stream in example_sample.groups(): + fstep_keys = sorted(example_stream.group_keys(), key=int) + for fk in fstep_keys: + fstep_group = example_stream[fk] + child_names = set(fstep_group.group_keys()) + if "prediction" in child_names or "target" in child_names: + # Return fstep 0 of this stream for correct forecast_offset detection + return ItemKey(sample, 0, stream) + # Fallback: use first stream / fstep 0 stream, example_stream = next(example_sample.groups()) fstep = 0 - except StopIteration as e: + except (StopIteration, IndexError) as e: msg = f"Data store at: {self._store_path} is empty." raise FileNotFoundError(msg) from e @@ -526,10 +549,11 @@ def forecast_steps(self) -> list[int]: _, example_sample = next(self.data_root.groups()) _, example_stream = next(example_sample.groups()) - all_steps = sorted(list(example_stream.group_keys())) + all_steps = sorted(example_stream.group_keys(), key=int) if self.forecast_offset == 1: - return all_steps[1:] # exclude fstep with no targets/preds + # exclude fstep with no targets/preds (may be absent from the store entirely) + return [step for step in all_steps if int(step) != 0] else: return all_steps diff --git a/packages/evaluate/src/weathergen/evaluate/export/export_core.py b/packages/evaluate/src/weathergen/evaluate/export/export_core.py index eeea8576b0..a43be509fb 100644 --- a/packages/evaluate/src/weathergen/evaluate/export/export_core.py +++ b/packages/evaluate/src/weathergen/evaluate/export/export_core.py @@ -184,9 +184,17 @@ def get_channels(channels, stream: str, fname_zarr: str) -> list[str]: with zarrio_reader(fname_zarr) as zio: zio_forecast_steps = sorted([int(step) for step in zio.forecast_steps]) dummy_out = zio.get_data(0, stream, zio_forecast_steps[0]) - all_channels = dummy_out.target.channels + all_channels = dummy_out.prediction.channels if channels is not None: + channels = list(channels) + # "10ff" (10m wind speed) is derived from its u/v components, so + # make sure both are exported whenever "10ff" is requested. + if "10ff" in channels: + for component in ("10u", "10v"): + if component not in channels: + channels.append(component) + existing_channels = set(all_channels) & set(channels) if existing_channels != set(channels): missing_channels = set(channels) - set(existing_channels) @@ -221,7 +229,9 @@ def get_grid_type(data_type, stream: str, fname_zarr: str) -> str: # TODO: this will change after restructuring the lead time. -def get_source_info(fname_zarr, stream, samples) -> tuple[list[np.datetime64], list[np.datetime64]]: +def get_source_info( + fname_zarr, stream, samples, fstep_hours: int = 6 +) -> tuple[list[np.datetime64], list[np.datetime64]]: """ Retrieve source interval boundaries from the source group at forecast step 0. @@ -232,6 +242,11 @@ def get_source_info(fname_zarr, stream, samples) -> tuple[list[np.datetime64], l The ``source_end`` also serves as the reference (initialisation) time. + If the store has no ``source`` group (e.g. prediction-only stores with + ``forecast_offset=1``), the reference time is derived from the earliest + available ``prediction`` group instead: its valid time minus the lead + time (``fstep * fstep_hours``) gives the initialisation time. + Parameters ---------- fname_zarr : str @@ -240,6 +255,10 @@ def get_source_info(fname_zarr, stream, samples) -> tuple[list[np.datetime64], l Stream name to retrieve data for (e.g., 'ERA5'). samples : list List of samples to process. + fstep_hours : int + Number of hours between consecutive forecast steps. Used to convert + a prediction's valid time back to the initialisation time when no + source group is present. Returns ------- @@ -252,16 +271,35 @@ def get_source_info(fname_zarr, stream, samples) -> tuple[list[np.datetime64], l source_starts = [] source_ends = [] with zarrio_reader(fname_zarr) as zio: + available_fsteps = sorted(int(step) for step in zio.forecast_steps) for sample in tqdm(samples, desc="Getting source info"): - group_path = f"{sample}/{stream}/0/source" - source_group = zio.data_root.get(group_path) - - if source_group is None: - raise FileNotFoundError(f"Zarr group '{group_path}' not found in {fname_zarr}") - - times_arr = np.asarray(source_group["times"]).astype("datetime64[ns]") - source_start = np.min(times_arr) - source_end = np.max(times_arr) + source_group = zio.data_root.get(f"{sample}/{stream}/0/source") + + if source_group is not None: + times_arr = np.asarray(source_group["times"]).astype("datetime64[ns]") + source_start = np.min(times_arr) + source_end = np.max(times_arr) + else: + # Prediction-only store (forecast_offset=1): no source group. + # Derive the init time from the earliest available prediction + # group. + pred_group = None + for fstep in available_fsteps: + candidate = zio.data_root.get(f"{sample}/{stream}/{fstep}/prediction") + if candidate is not None: + pred_group = candidate + break + if pred_group is None: + raise FileNotFoundError( + f"No 'source' group and no 'prediction' group found for " + f"sample {sample}, stream '{stream}' in {fname_zarr}" + ) + times_arr = np.asarray(pred_group["times"]).astype("datetime64[ns]") + # The init (reference) time is one forecast step before the + # earliest prediction valid time: init = min(times) - fstep_hours. + lead = np.timedelta64(fstep_hours, "h").astype("timedelta64[ns]") + source_start = np.min(times_arr) - lead + source_end = np.min(times_arr) - lead _logger.debug(f"Sample {sample}: source_interval=[{source_start} .. {source_end}]") source_starts.append(source_start) @@ -315,7 +353,9 @@ def export_model_outputs(data_type: str, config: OmegaConf, **kwargs) -> None: for stream in streams: grid_type = get_grid_type(data_type, stream, fname_zarr) channels = get_channels(channels, stream, fname_zarr) - source_starts, source_ends = get_source_info(fname_zarr, stream, samples) + source_starts, source_ends = get_source_info( + fname_zarr, stream, samples, fstep_hours=kwargs.get("fstep_hours", 6) + ) kwargs["grid_type"] = grid_type kwargs["channels"] = channels kwargs["data_type"] = data_type @@ -382,7 +422,7 @@ def export_model_outputs(data_type: str, config: OmegaConf, **kwargs) -> None: for sample, _fstep, data in pool.imap_unordered( get_data_worker, batch_tasks, chunksize=1 ): - sample_results[sample].append(data) + sample_results[sample].append((_fstep, data)) pbar.update(1) # Check if this sample is complete (all fsteps received). @@ -390,7 +430,9 @@ def export_model_outputs(data_type: str, config: OmegaConf, **kwargs) -> None: b_idx = sample_to_batch_idx[sample] source_start = batch_source_starts[b_idx] source_end = batch_source_ends[b_idx] - results_iter = iter(sample_results[sample]) + # Sort by forecast step so accumulation is in order. + sample_results[sample].sort(key=lambda x: x[0]) + results_iter = iter([d for _, d in sample_results[sample]]) processed = parser.process_sample( results_iter, ref_time=source_end, diff --git a/packages/evaluate/src/weathergen/evaluate/export/parsers/quaver_parser.py b/packages/evaluate/src/weathergen/evaluate/export/parsers/quaver_parser.py index a19077ccb4..93ed5de81c 100644 --- a/packages/evaluate/src/weathergen/evaluate/export/parsers/quaver_parser.py +++ b/packages/evaluate/src/weathergen/evaluate/export/parsers/quaver_parser.py @@ -87,6 +87,20 @@ def process_sample( ------- None """ + # Identify variables that need accumulation across forecast steps + accum_vars = { + var + for var in self.channels + if self.mapping.get(var, self.mapping.get(var.split("_")[0] if "_" in var else var, {})).get( + "accumulate", False + ) + } + # Running accumulator: {var_name: 1D numpy array} + accum_state: dict[str, np.ndarray] = {} + + if accum_vars: + _logger.info(f"Accumulating total precipitation for variables: {accum_vars}") + for result in fstep_iterator_results: if result is None: continue @@ -115,6 +129,25 @@ def process_sample( field_data = da_sub.sel(channel=var) field_data = self.scale_data(field_data, var) + field_values = field_data.values.copy() + + # Clamp negative precipitation to zero. + if var in accum_vars: + field_values = np.maximum(field_values, 0.0) + + # Accumulate precipitation: replace per-step values with + # running total (current step + all previous steps). + if var in accum_vars: + if var in accum_state: + accum_state[var] = accum_state[var] + field_values + else: + accum_state[var] = field_values.copy() + field_values = accum_state[var].copy() + _logger.info( + f"[Worker] Accumulated {var}: step sum={field_values.sum():.6g}, " + f"total sum={accum_state[var].sum():.6g}" + ) + template_field = self.template_cache.get((var, level), None) if template_field is None: _logger.error(f"Template for var={var}, level={level} not found. Skipping.") @@ -123,12 +156,14 @@ def process_sample( metadata = self.get_metadata( ref_time=ref_time, valid_time=vt, + source_interval_start=source_interval_start, source_interval_end=source_interval_end, level=level, + var=var, ) encoded = self.encoder.encode( - values=field_data.values, + values=field_values, template=template_field, metadata=metadata, ) @@ -153,10 +188,20 @@ def extract_var_info(self, var: str) -> tuple[str, str, str]: tuple[str, str, str] Variable short name, level, and level type. """ - var_short = var.split("_")[0] if "_" in var else var - level = int(var.split("_")[-1]) if "_" in var else "sfc" + # Try full variable name first, then fall back to first token before '_' + if var in self.mapping: + var_short = var + var_config = self.mapping[var] + level = "sfc" + elif "_" in var: + var_short = var.split("_")[0] + var_config = self.mapping.get(var_short, {}) + level = int(var.split("_")[-1]) + else: + var_short = var + var_config = self.mapping.get(var_short, {}) + level = "sfc" - var_config = self.mapping.get(var_short, {}) if not var_config: raise ValueError( f"Variable '{var} (using: {var_short})' not found in configuration mapping." @@ -209,7 +254,7 @@ def get_output_filename(self, level_type: str) -> Path: """ return ( Path(self.output_dir) - / f"{self.data_type}_{level_type}_{self.run_id}_{self.expver}.{self.file_extension}" + / f"{self.data_type}_{level_type}_{self.run_id}_{self.expver}_rank{int(self.rank):04d}.{self.file_extension}" ) def assign_coords(self, data: xr.DataArray) -> xr.DataArray: @@ -236,26 +281,36 @@ def get_metadata( self, ref_time: pd.Timestamp, valid_time: np.datetime64, + source_interval_start: np.datetime64, source_interval_end: np.datetime64, level: str, + var: str = None, ): """ Add metadata to the dataset attributes. - The GRIB ``step`` is computed as ``valid_time - source_interval_end`` - (in hours), i.e. the lead time relative to the end of the - conditioning window. + The GRIB ``date``/``time`` is set to ``source_interval_start`` + (the true initialisation time of the forecast). The GRIB ``step`` + is computed as ``valid_time - source_interval_start`` (in hours). """ - step_hours = int((valid_time - source_interval_end) / np.timedelta64(1, "h")) + step_hours = int((valid_time - source_interval_start) / np.timedelta64(1, "h")) metadata = { - "date": ref_time, + "date": pd.Timestamp(source_interval_start), "step": step_hours, "expver": self.expver, "marsClass": "rd", } if level != "sfc": metadata["level"] = level + + # Override paramId if specified in the variable config. + if var is not None: + var_config = self.mapping.get(var, self.mapping.get(var.split("_")[0] if "_" in var else var, {})) + param_id = var_config.get("paramId") + if param_id is not None: + metadata["paramId"] = param_id + return metadata def save(self, encoded_fields: list, level_type: str): diff --git a/packages/evaluate/src/weathergen/evaluate/io/data/io_orchestration.py b/packages/evaluate/src/weathergen/evaluate/io/data/io_orchestration.py index 4131bf20f0..bb20c391c0 100644 --- a/packages/evaluate/src/weathergen/evaluate/io/data/io_orchestration.py +++ b/packages/evaluate/src/weathergen/evaluate/io/data/io_orchestration.py @@ -83,6 +83,7 @@ class IOState: offset: np.timedelta64 | None = ( None # fallback offset in hours for init_time when source_interval is missing ) + anemoi_target_cfg: dict | None = None # if set, read targets from anemoi dataset # --------------------------------------------------------------------------- @@ -254,6 +255,7 @@ def _build_io_state( n_io_workers: int, ens_select: EnsembleSelect, rank: str = "", + inference_cfg: dict | None = None, ) -> IOState: """Resolve all I/O parameters that are shared between the two impl paths.""" zarr_path = str(fname_zarr) @@ -278,6 +280,19 @@ def _build_io_state( if isinstance(regrid_opts, bool) and regrid_opts: regrid_opts = {"target_grid": [1.5, 1.5]} + # ---- Resolve anemoi target config from inference config ---- + anemoi_target_cfg = None + if inference_cfg: + stream_info = inference_cfg.get("streams", {}).get(stream, {}) + if stream_info.get("type") in ("anemoi", "anemoi_operan") and stream_info.get("filenames"): + data_path = inference_cfg.get("data_path_anemoi", "") + filename = str(Path(data_path) / stream_info["filenames"][0]) + anemoi_target_cfg = { + "filename": filename, + "channels": stream_info.get("val_target_channels", []), + } + _logger.info(f"Anemoi target source: {filename}") + return IOState( run_id=run_id, zarr_path=zarr_path, @@ -295,10 +310,11 @@ def _build_io_state( coords=coords, lat=lat, lon=lon, - n_workers=n_io_workers, + n_workers=min(n_io_workers, 20) if anemoi_target_cfg is not None else n_io_workers, rank=rank, offset=offset, regrid_opts=regrid_opts, + anemoi_target_cfg=anemoi_target_cfg, ) @@ -315,6 +331,7 @@ def _parallel_read( backend: str, label: str, regrid_opts: dict, + anemoi_target_cfg: dict | None = None, ) -> tuple[list, bool]: """Dispatch _read_sample over samples, with parallel→sequential fallback. @@ -333,6 +350,7 @@ def _parallel_read( read_coords=need_coords, is_gridded=is_gridded, regrid_opts=regrid_opts, + anemoi_target_cfg=anemoi_target_cfg, ) calls = [delayed(_read_sample)(sample=s, **kwargs) for s in samples] @@ -550,6 +568,7 @@ def get_data_dirstore(state: IOState) -> ReaderOutput: backend=state.backend, label=f"RUN {state.run_id} [rank {state.rank}] - {state.stream} fstep {fs}", regrid_opts=state.regrid_opts, + anemoi_target_cfg=state.anemoi_target_cfg, ) # If _parallel_read fell back to sequential, honour that for the rest if fell_back: @@ -583,7 +602,7 @@ def get_data_dirstore(state: IOState) -> ReaderOutput: del results - if n_workers > 1: + if n_workers > 1 and state.backend == "loky": get_reusable_executor().shutdown(wait=True) _logger.info( @@ -623,7 +642,13 @@ def get_data_zipstore(state: IOState) -> ReaderOutput: read_coords=not state.is_gridded, is_gridded=state.is_gridded, regrid_opts=state.regrid_opts, + anemoi_target_cfg=state.anemoi_target_cfg, ) + if state.anemoi_target_cfg is not None: + _logger.info( + f"RUN {state.run_id} [rank {state.rank}] - {state.stream}: " + f"Target data will be read from anemoi dataset (not zarr)." + ) calls = [ delayed(_read_sample)(sample=s, fsteps=[fs], **kwargs) for s in state.samples @@ -637,6 +662,11 @@ def get_data_zipstore(state: IOState) -> ReaderOutput: verbose=5, ) + _logger.info( + f"RUN {state.run_id} [rank {state.rank}] - {state.stream}: " + f"dispatch_parallel returned {len(flat_results)} results. Assembling..." + ) + # --- Re-group: flat_results[sample_idx * n_fsteps + fstep_idx] -------- n_fsteps = len(state.fsteps) # Gather per-sample results in the same shape as get_data_dirstore expects @@ -694,7 +724,7 @@ def get_data_zipstore(state: IOState) -> ReaderOutput: del flat_results - if state.n_workers > 1: + if state.n_workers > 1 and state.backend == "loky": get_reusable_executor().shutdown(wait=True) _logger.info( diff --git a/packages/evaluate/src/weathergen/evaluate/io/data/io_workers.py b/packages/evaluate/src/weathergen/evaluate/io/data/io_workers.py index 643246d422..6acf9198ee 100644 --- a/packages/evaluate/src/weathergen/evaluate/io/data/io_workers.py +++ b/packages/evaluate/src/weathergen/evaluate/io/data/io_workers.py @@ -15,7 +15,9 @@ import contextlib import logging +import threading +import anemoi.datasets import numpy as np import zarr from numpy.typing import NDArray @@ -24,6 +26,108 @@ _logger = logging.getLogger(__name__) +# Module-level cache so that threading workers reuse the same anemoi dataset +# handle across tasks dispatched to the same process. +_anemoi_cache: dict[str, tuple] = {} +_anemoi_lock = threading.Lock() + + +def _open_anemoi_dataset(anemoi_cfg: dict) -> tuple: + """Open the anemoi dataset once and precompute target channel indices. + + Uses a module-level cache keyed by filename so that threading workers + reuse the same handle. A lock ensures only one thread opens the dataset. + + Returns ``(ds, target_idx, ds_dates)`` for reuse across forecast steps. + """ + filename = anemoi_cfg["filename"] + if filename in _anemoi_cache: + return _anemoi_cache[filename] + + with _anemoi_lock: + # Double-check after acquiring lock + if filename in _anemoi_cache: + return _anemoi_cache[filename] + + target_channels = anemoi_cfg["channels"] + + _logger.info(f"Opening anemoi dataset {filename} (lazy)") + ds = anemoi.datasets.open_dataset(filename) + target_idx = [ds.variables.index(ch) for ch in target_channels] + ds_dates = ds.dates.astype("datetime64[s]") + result = (ds, target_idx, ds_dates) + _anemoi_cache[filename] = result + return result + + +def _read_anemoi_target( + ds, + target_idx: list[int], + ds_dates: np.ndarray, + times: np.ndarray, + channel_idxs: list[int] | None, +) -> np.ndarray: + """Read target data from a pre-opened anemoi dataset for the given valid times. + + Parameters + ---------- + ds : anemoi dataset handle + Already-opened dataset (from ``_open_anemoi_dataset``). + target_idx : list[int] + Indices of val_target_channels within ``ds.variables``. + ds_dates : np.ndarray + ``ds.dates`` cast to ``datetime64[s]`` (precomputed). + times : np.ndarray + Valid times (datetime64) to extract from the dataset. + channel_idxs : list[int] | None + Channel indices to select (applied after reading). + + Returns + ------- + np.ndarray + Target data array, shape ``(n_points, n_channels)``, float32. + """ + unique_times = np.unique(times) + + # Handle empty times (some fsteps have no data) + if len(unique_times) == 0: + n_ch = len(target_idx) + if channel_idxs is not None: + n_ch = len(channel_idxs) + return np.empty((0, n_ch), dtype=np.float32) + + # Read one time-slice at a time (typically ≤6 per fstep) + all_data = [] + for ut in unique_times: + ut_s = np.datetime64(ut, "s") + matches = np.where(ds_dates == ut_s)[0] + + if len(matches) == 0: + raise ValueError( + f"Time {ut} not found in anemoi dataset. " + f"Available range: {ds_dates[0]} .. {ds_dates[-1]}" + ) + + idx = int(matches[0]) + # ds[idx] → shape (n_variables, n_ens, n_gridpoints) + data_slice = np.asarray(ds[idx]) + # Select target channels and squeeze ensemble dim → (n_gridpoints, n_target_channels) + data_slice = data_slice[target_idx, 0, :].T + all_data.append(data_slice) + + # For gridded data every grid point shares the same unique time, + # so each slice already has the right shape. + if len(all_data) == 1: + target_data = all_data[0] + else: + target_data = np.concatenate(all_data, axis=0) + + # Apply the same early channel selection as the zarr path + if channel_idxs is not None: + target_data = target_data[:, channel_idxs] + + return target_data.astype(np.float32) + def _compute_early_channel_selection( read_channels: list[str], @@ -92,6 +196,7 @@ def _read_sample( read_coords: bool = False, is_gridded: bool = True, regrid_opts: dict | None = None, + anemoi_target_cfg: dict | None = None, ) -> tuple[list[NDArray], list[NDArray], list[NDArray], dict]: """ Read all forecast steps for one sample via direct zarr array access. @@ -172,22 +277,40 @@ def _read_sample( if source_interval: break - for fs in fsteps: + # Open anemoi dataset once for all fsteps (if configured) + anemoi_handle = None + if anemoi_target_cfg is not None: + anemoi_handle = _open_anemoi_dataset(anemoi_target_cfg) + + for fi, fs in enumerate(fsteps): base = f"{sample}/{stream}/{fs}" # Direct array access — bypasses OutputDataset/as_xarray/dask entirely pred_data = np.asarray(ds[f"{base}/prediction/data"]) - target_data = np.asarray(ds[f"{base}/target/data"]) times_data = np.asarray(ds[f"{base}/prediction/times"]) - # Select channels by index + if anemoi_handle is not None: + # Read target from anemoi dataset + anemoi_ds, target_idx, ds_dates = anemoi_handle + if fi == 0: + _logger.info( + f"Sample {sample} fstep {fs}: reading target from anemoi " + f"(times shape={times_data.shape}, unique={len(np.unique(times_data))})" + ) + target_data = _read_anemoi_target( + anemoi_ds, target_idx, ds_dates, times_data, channel_idxs + ) + else: + target_data = np.asarray(ds[f"{base}/target/data"]) + if channel_idxs is not None: + target_data = target_data[:, channel_idxs] + + # Select channels by index for prediction if channel_idxs is not None: pred_data = ( pred_data[:, channel_idxs] if pred_data.ndim == 2 else pred_data[:, channel_idxs, :] ) - target_data = target_data[:, channel_idxs] - - # Handle sub-steps (gridded data with multiple valid_times per fstep). + # Handle sub-steps (gridded data with multiple valid_times per fstep). # For scatter/observation data each observation has its own timestamp, # so splitting by unique time would create one tiny array per obs — # thousands of them — causing the assembly code to hang. diff --git a/packages/evaluate/src/weathergen/evaluate/io/wegen_reader.py b/packages/evaluate/src/weathergen/evaluate/io/wegen_reader.py index 6589813305..ffb1ccd36e 100644 --- a/packages/evaluate/src/weathergen/evaluate/io/wegen_reader.py +++ b/packages/evaluate/src/weathergen/evaluate/io/wegen_reader.py @@ -616,6 +616,9 @@ def get_data( self._num_io_workers, ens_select, rank=rank_file.stem.split("rank")[-1], + inference_cfg=self.inference_cfg + if self.eval_cfg.get("type") == "anemoi-target" + else None, ) get_data_fn = get_data_zipstore if state.is_zip else get_data_dirstore result = get_data_fn(state) diff --git a/packages/evaluate/src/weathergen/evaluate/run_evaluation.py b/packages/evaluate/src/weathergen/evaluate/run_evaluation.py index e6af331012..a273124356 100755 --- a/packages/evaluate/src/weathergen/evaluate/run_evaluation.py +++ b/packages/evaluate/src/weathergen/evaluate/run_evaluation.py @@ -162,7 +162,7 @@ def get_reader( region: str | None = None, metric: dict[str, object] | None = None, ): - if reader_type == "zarr": + if reader_type == "zarr" or reader_type == "anemoi-target": reader = WeatherGenZarrReader(run, run_id, private_paths) elif reader_type == "csv": reader = CsvReader(run, run_id, private_paths) @@ -220,11 +220,12 @@ def _process_stream( _logger.info(f"Stream {stream} not found for run {run_id}. Skipping.") return run_id, stream, {}, {} - needs_plotting = stream_dict.get("plotting") and type_ == "zarr" + is_zarr_like = type_ in ("zarr", "anemoi-target") + needs_plotting = stream_dict.get("plotting") and is_zarr_like needs_scoring = stream_dict.get("evaluation", False) output_data = None - if (needs_plotting or needs_scoring) and type_ == "zarr": + if (needs_plotting or needs_scoring) and is_zarr_like: available_data = reader.check_availability(stream, mode="evaluation") output_data = None @@ -247,9 +248,9 @@ def _process_stream( if not needs_scoring: return run_id, stream, {}, {} - plot_score_maps = plot_score_options.get("plot_score_maps", False) and type_ == "zarr" + plot_score_maps = plot_score_options.get("plot_score_maps", False) and is_zarr_like plot_score_init_time_series = ( - plot_score_options.get("plot_score_init_time_series", False) and type_ == "zarr" + plot_score_options.get("plot_score_init_time_series", False) and is_zarr_like ) stream_loaded_scores, recomputable_metrics = reader.load_scores(stream, regions, metrics)