Beam pipeline

weatherbenchX.beam_pipeline.define_pipeline(root: Pipeline, times: TimeChunks, predictions_loader: DataLoader, targets_loader: DataLoader, metrics: Mapping[str, Metric], aggregator: Aggregator, out_path: str, max_chunks_per_aggregation_stage: int | None = 10, setup_fn: Callable[[], None] | None = None)

Defines the beam pipeline.

Parameters:

• root – Pipeline root.

• times – TimeChunks instance.

• predictions_loader – DataLoader instance.

• targets_loader – DataLoader instance.

• metrics – A dictionary of metrics to compute.

• aggregator – Aggregation instance.

• out_path – The full path to write the metrics to.

• max_chunks_per_aggregation_stage – The maximum number of chunks to aggregate in a single worker. If None, does aggregation in a single step. Default: 10

• setup_fn – (Optional) A function to call once per worker in LoadPredictionsAndTargets.