Fold the dataset sampler into mlcast (mlcast.sampling)

[franchg]

Summary

Folds the previously separate mlcast-dataset-sampler tool into the mlcast
package as a new mlcast.sampling subpackage. This removes the duplication
between the two repos and gives training a single in-package path from a source
radar Zarr to a sampled training set.

The design is a clean producer → consumer seam:

• Producer (mlcast stats) — scans a source Zarr and writes a stats
  parquet: an index of candidate datacubes, each tagged with nan_count /
  sum / mean / frac_wet. The parquet schema + metadata footer is the
  contract (mlcast.sampling.stats_spec).
• Consumer (training) — SourceDataIndexedDataset reads that index, and a
  pluggable Sampler reshapes the candidate pool at dataset init.

What's new

Data-prep CLI (folded into the existing mlcast argparse CLI as subcommands):

• mlcast stats <zarr> — cumsum sliding-window scan → zstd stats parquet (CPU
  via bottleneck, or --device cuda for the GPU backend).
• mlcast validate-stats <parquet> — validates a parquet against the contract;
  prints the schema, metadata, and a row preview.

Sampler registry (mlcast.sampling.samplers):

• Sampler ABC + @register_sampler / get_sampler (mirrors the existing
  normalization registry).
• UniformSampler(keep_fraction) and
  ImportanceSampler(column, q_min, scale, mean_weight) — per-row accept/reject,
  run once at init so the dataset length is fixed and reused every epoch.
• Applied per split by SourceDataDataModule (train_sampler / eval_sampler):
  importance on train, uniform on val/test, so validation/test stay representative.

Training dataset:

• SourceDataIndexedDataset reads the parquet index_path (and still accepts a
  legacy CSV index, used without a sampler).

Notable changes

• bottleneck added to core dependencies; the unused fire dependency removed.
• mlcast CLI startup made lazy — mlcast -h ~4.2s → ~0.08s by deferring the
  training stack to the train path.
• Example updated: csv_path → index_path.
• README gained a "Preparing training data" section documenting the workflow.

Testing

• uv run pytest → 114 passed, 1 skipped (the skip is the CUDA-only GPU
  backend parity test).
• ruff check clean.
• New tests cover sampler semantics + registry, the stats _process_chunk
  against two independent reference oracles, GPU/CPU parity, and the parquet
  contract.

Out of scope / follow-ups

• Decommissioning the standalone mlcast-dataset-sampler repo (archive, or a
  thin shim pointing here).

[leifdenby]

This is looking really nice! I do think we should maybe name things slightly differently here, maybe we can make the step that computes the stats be called "time-sequence stats" instead? I.e. the CLI call would be uv run mlcast timeseq-stats instead?

I think I have also discovered a bug when trying to apply this to the DMI dataset, but maybe I am doing something incorrect. Steps to reproduce:

1. download to sample data from the DMI dataset where we know it is raining, e.g. from 2024-12-31T00:00 is the jupyter notebook

uv run examples/scripts/download_mlcast_dataset_sample.py precipitation.dmi_10_minutes --duration PT4H --start-time 2024-12-3

2. run time-sequence stats collection:

uv run mlcast stats data/mlcast-source-datasets/dmi/v0.1.1/dmi_10_minutes.zarr --data-var db

this reports 0 datacube candidates. Even with just looking at two successive timesteps there are no valid datacubes found:

uv run mlcast stats data/mlcast-source-datasets/dmi/v0.1.1/dmi_10_minutes.zarr --data-var dbz --time-depth 2

However, the data contains more than 7dBZ rainrate values:

import xarray as xr

ds_dmi = xr.open_zarr("../data/mlcast-source-datasets/dmi/v0.1.1/dmi_10_minutes.zarr")
ds_dmi

da = ds_dmi.dbz.isel(time=slice(0, 10))

da.plot.imshow(x="x", y="y", col="time", col_wrap=5)

Any ideas what could be going on @franchg?

[franchg]

Thanks for reviewing this @leifdenby!
The root cause is the --time-step-minutes default of 5: the DMI dataset is 10-minute data, so the time-continuity check saw a "gap" between every pair of consecutive frames and rejected every candidate window. Your exact runs work once the cadence is passed:
uv run mlcast stats data/mlcast-source-datasets/dmi/v0.1.1/dmi_10_minutes.zarr --data-var dbz --time-step-minutes 10
I verified against the full DMI archive: one day of data with default parameters gives 182,778 datacube candidates (and the data-kind auto-detection correctly picks reflectivity with the 7 dBZ wet threshold from the dbz attrs).

The real bug was that this failure mode was silent, a misconfigured run happily wrote an empty parquet. I've pushed 76c2724 so stats now exits non-zero with a diagnostic instead. Your original command now fails with:

ERROR 0 gap-free time windows: expecting a 5-min step between frames (--time-step-minutes)
but the dominant observed step is 10 min (100.0% of 144 steps). If that is the dataset's
real cadence, rerun with --time-step-minutes 10.

Could you re-test on with the change I made?

As a follow-up (separate change) I'd like to auto-detect the cadence from the time coordinate, so the flag becomes an override rather than something you have to know.

One thought on timeseq-stats rename: it names only the time axis, but these are stats over sliding (time × y × x) windows, which the tool, README, and output already call datacube candidates. So my counter-proposal would be mlcast datacube-stats which says exactly what the stats are over, in vocabulary the project already uses, and validate-stats still reads as its companion. Another alternative woul be to use a verb/action: mlcast scan (scan the archive for candidates, short and clear) or mlcast build-index (the training dataset then consumes the parquet as index_path).
Any preference?