run fast eval without target

config/evaluate/config_zarr2cf.yaml

@@ -89,13 +89,24 @@ 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
     std: precipitation_amount
     wg_unit: m
     std_unit: m
     level_type: sfc
+    paramId: 228
+    accumulate: true
 
 
 coordinates:

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

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
 

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,15 +422,17 @@ 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).
                     if len(sample_results[sample]) == n_fsteps:
                         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,