Skip to content

climate_data

Climate Data

This package contains modules for extracting, processing, harmonizing, and downscaling climate data. It sources historical climate data from the European Centre for Medium-Range Weather Forecasts (ECMWF) ERA5 dataset and future climate data from the Coupled Model Intercomparison Project Phase 6 (CMIP6).

aggregate

hierarchy

Runner for the collate stage of the climate aggregates pipeline.

This module provides functions to collate block-level datasets into measure-level datasets. The compilation process: 1. Loads raw results from all blocks for a given hierarchy, measure, and scenario 2. Aggregates the data up the location hierarchy 3. Produces views for subset hierarchies 4. Saves the results as measure-level datasets

hierarchy(agg_version: str, hierarchy: list[str], agg_measure: list[str], agg_scenario: list[str], population_model_dir: str, output_dir: str, queue: str, dry_run: bool) -> None

Collate block-level datasets into measure-level datasets.

This command: 1. Identifies which combinations of hierarchy, measure, and scenario need compilation 2. Creates parallel jobs to collate each combination 3. Runs the jobs using jobmon

Source code in src/climate_data/aggregate/hierarchy.py
@click.command()
@clio.with_agg_version()
@clio.with_hierarchy(allow_all=True)
@clio.with_agg_measure(allow_all=True)
@clio.with_agg_scenario(allow_all=True)
@clio.with_input_directory("population-model", cdc.POPULATION_MODEL_ROOT)
@clio.with_output_directory(cdc.AGGREGATE_ROOT)
@clio.with_queue()
@clio.with_dry_run()
def hierarchy(
    agg_version: str,
    hierarchy: list[str],
    agg_measure: list[str],
    agg_scenario: list[str],
    population_model_dir: str,
    output_dir: str,
    queue: str,
    dry_run: bool,
) -> None:
    """Collate block-level datasets into measure-level datasets.

    This command:
    1. Identifies which combinations of hierarchy, measure, and scenario need compilation
    2. Creates parallel jobs to collate each combination
    3. Runs the jobs using jobmon
    """
    ca_data = ClimateAggregateData(output_dir)

    n_jobs = len(hierarchy) * len(agg_measure) * len(agg_scenario)
    print(f"Running {n_jobs} jobs")

    run_parallel_maybe_dry_run(
        runner="cdtask aggregate",
        task_name="hierarchy",
        node_args={
            "hierarchy": hierarchy,
            "agg-measure": agg_measure,
            "agg-scenario": agg_scenario,
        },
        task_args={
            "agg-version": agg_version,
            "population-model-dir": population_model_dir,
            "output-dir": output_dir,
        },
        task_resources={
            "queue": queue,
            "cores": 1,
            "memory": "50G",
            "runtime": "240m",
            "project": "proj_rapidresponse",
        },
        log_root=ca_data.log_dir("aggregate_hierarchy"),
        max_attempts=3,
        dry_run=dry_run,
    )

hierarchy_main(agg_version: str, hierarchy: str, measure: str, scenario: str, population_model_dir: str, output_dir: str, *, progress_bar: bool = False) -> None

Collate block-level datasets into measure/scenario-level datasets.

This function: 1. Loads all block-level datasets for a given hierarchy, measure, and scenario 2. Combines them into a single dataset 3. Aggregates the data up the location hierarchy 4. Produces views for subset hierarchies 5. Saves the results as measure-level datasets

Parameters

agg_version The version identifier hierarchy The full aggregation hierarchy to process measure The climate measure to process scenario The climate scenario to process population_model_dir Path to the population model directory output_dir Path to save results progress_bar Whether to show a progress bar

Source code in src/climate_data/aggregate/hierarchy.py
def hierarchy_main(
    agg_version: str,
    hierarchy: str,
    measure: str,
    scenario: str,
    population_model_dir: str,
    output_dir: str,
    *,
    progress_bar: bool = False,
) -> None:
    """Collate block-level datasets into measure/scenario-level datasets.

    This function:
    1. Loads all block-level datasets for a given hierarchy, measure, and scenario
    2. Combines them into a single dataset
    3. Aggregates the data up the location hierarchy
    4. Produces views for subset hierarchies
    5. Saves the results as measure-level datasets

    Parameters
    ----------
    agg_version
        The version identifier
    hierarchy
        The full aggregation hierarchy to process
    measure
        The climate measure to process
    scenario
        The climate scenario to process
    population_model_dir
        Path to the population model directory
    output_dir
        Path to save results
    progress_bar
        Whether to show a progress bar
    """
    print(f"Compiling {measure} for {scenario} in {hierarchy}")
    ca_data = ClimateAggregateData(output_dir)
    pm_data = PopulationModelData(population_model_dir)
    draws = cdc.DRAWS

    # Load hierarchy data for aggregation
    hierarchy_df = pm_data.load_subset_hierarchy(hierarchy)

    # Get all block keys
    modeling_frame = pm_data.load_modeling_frame()
    block_keys = modeling_frame.block_key.unique()
    intersecting = utils.blocks_with_shapefile_intersections(
        hierarchy, pm_data, modeling_frame
    )
    block_keys = [b for b in block_keys if b in intersecting]

    # Load and combine all block-level results
    desc_template = "{draw:5} {block_key:15}"
    pbar = tqdm.tqdm(
        total=len(block_keys) * len(draws),
        desc=desc_template.format(draw="DRAW", block_key="BLOCK KEY"),
        disable=not progress_bar,
    )

    all_results = []
    pop_df: pd.Series[Any] | None = None

    for draw in draws:
        save_population = (
            measure == "mean_temperature" and scenario == "ssp245" and draw == "000"
        )
        draw_results = []
        for block_key in block_keys:
            pbar.set_description(desc_template.format(draw=draw, block_key=block_key))

            draw_df = ca_data.load_raw_results(
                agg_version,
                hierarchy,
                block_key,
                draw,
                measure=measure,
                scenario=scenario,
            ).drop(columns=["scenario", "measure"])
            draw_results.append(draw_df)

            pbar.update()
        draw_df = (
            pd.concat(draw_results, ignore_index=True)
            .groupby(["location_id", "year_id"])
            .sum()
            .reset_index()
        )

        agg_df = utils.aggregate_climate_to_hierarchy(
            draw_df,
            hierarchy_df,
        ).set_index(["location_id", "year_id"])
        all_results.append(agg_df["value"].rename(draw))

        if save_population:
            pop_df = agg_df["population"]

    pbar.close()

    combined_results = pd.concat(all_results, axis=1)

    # Produce views for subset hierarchies
    subset_hierarchies = cdc.HIERARCHY_MAP[hierarchy]
    for subset_hierarchy in subset_hierarchies:
        # Load the subset hierarchy
        subset_hierarchy_df = pm_data.load_subset_hierarchy(subset_hierarchy)

        # Filter results to only include locations in the subset hierarchy
        subset_location_ids = subset_hierarchy_df["location_id"].tolist()
        subset_results = combined_results.loc[subset_location_ids]

        # Save results for the subset hierarchy
        ca_data.save_results(
            subset_results,
            agg_version,
            subset_hierarchy,
            scenario,
            measure,
        )

        if pop_df is not None:
            subset_pop = pop_df.loc[subset_location_ids].reset_index()
            ca_data.save_population(subset_pop, agg_version, subset_hierarchy)

hierarchy_task(agg_version: str, hierarchy: str, agg_measure: str, agg_scenario: str, population_model_dir: str, output_dir: str, *, progress_bar: bool) -> None

Collate block-level datasets into measure-level datasets.

This command collates results for a specific hierarchy, measure, and scenario.

Source code in src/climate_data/aggregate/hierarchy.py
@click.command()
@clio.with_agg_version()
@clio.with_hierarchy()
@clio.with_agg_measure()
@clio.with_agg_scenario()
@clio.with_input_directory("population-model", cdc.POPULATION_MODEL_ROOT)
@clio.with_output_directory(cdc.AGGREGATE_ROOT)
@clio.with_progress_bar()
def hierarchy_task(
    agg_version: str,
    hierarchy: str,
    agg_measure: str,
    agg_scenario: str,
    population_model_dir: str,
    output_dir: str,
    *,
    progress_bar: bool,
) -> None:
    """Collate block-level datasets into measure-level datasets.

    This command collates results for a specific hierarchy, measure, and scenario.
    """
    hierarchy_main(
        agg_version,
        hierarchy,
        agg_measure,
        agg_scenario,
        population_model_dir,
        output_dir,
        progress_bar=progress_bar,
    )

pixel

utils

aggregate_climate_to_hierarchy(data: pd.DataFrame, hierarchy: pd.DataFrame) -> pd.DataFrame

Create all aggregate climate values for a given hierarchy from most-detailed data.

Parameters

data The most-detailed climate data to aggregate. hierarchy The hierarchy to aggregate the data to.

Returns

pd.DataFrame The climate data with values for all levels of the hierarchy.

Source code in src/climate_data/aggregate/utils.py
def aggregate_climate_to_hierarchy(
    data: pd.DataFrame, hierarchy: pd.DataFrame
) -> pd.DataFrame:
    """Create all aggregate climate values for a given hierarchy from most-detailed data.

    Parameters
    ----------
    data
        The most-detailed climate data to aggregate.
    hierarchy
        The hierarchy to aggregate the data to.

    Returns
    -------
    pd.DataFrame
        The climate data with values for all levels of the hierarchy.
    """
    results = data.set_index("location_id").copy()

    # Most detailed locations can be at multiple levels of the hierarchy,
    # so we loop over all levels from most detailed to global, aggregating
    # level by level and appending the results to the data.
    for level in reversed(list(range(1, hierarchy.level.max() + 1))):
        level_mask = hierarchy.level == level
        parent_map = hierarchy.loc[level_mask].set_index("location_id").parent_id

        subset = results.loc[parent_map.index]
        subset["parent_id"] = parent_map

        parent_values = (
            subset.groupby(["year_id", "parent_id"])[["weighted_climate", "population"]]
            .sum()
            .reset_index()
            .rename(columns={"parent_id": "location_id"})
            .set_index("location_id")
        )
        results = pd.concat([results, parent_values])
    results = (
        results.reset_index()
        .sort_values(["location_id", "year_id"])
        .reset_index(drop=True)
    )

    # The block-level results arrive with object dtype, so a zero denominator
    # raises ZeroDivisionError from Python scalar arithmetic rather than giving
    # the inf that float division would. Coerce first, then divide only where
    # there are people: a population-weighted climate value is undefined for a
    # location with no population, and NaN says that honestly.
    #
    # Zero-population locations are real and expected. Some are uninhabited
    # (Antarctica, Bouvet Island, the Spratlys). Others are inhabited places the
    # population model carries no estimate for (Aland, Svalbard, Norfolk
    # Island), which is a coverage gap, not a modelling result -- hence the
    # count is reported rather than silently swallowed.
    results = results.astype({"weighted_climate": "float64", "population": "float64"})
    population = results["population"]
    no_population = ~(population > 0)
    results["value"] = results["weighted_climate"] / population.where(~no_population)

    if no_population.any():
        n_loc = results.loc[no_population, "location_id"].nunique()
        print(
            f"No population for {n_loc} location(s) covering "
            f"{int(no_population.sum())} location-year(s); "
            "their value is NaN."
        )

    return results

blocks_with_shapefile_intersections(hierarchy: str, pm_data: PopulationModelData, modeling_frame: gpd.GeoDataFrame | None = None) -> set[str]

Return the block_keys whose footprint intersects the hierarchy's raking shapes.

Blocks whose (dissolved) modeling-frame geometry intersects no raking polygon contribute only empty rows to the sum/sum aggregation, so the pipeline can skip them. The shapefile lookup is delegated to PopulationModelData.load_raking_shapes, which routes to the right file for whichever hierarchy is in use.

.. note::

This uses the block *geometry* vs the full shape set, whereas the production
emptiness gate in ``build_location_masks`` uses the raster *bounding box* vs
bbox-filtered shapes. Equivalence was validated empirically for ``lsae_1285``
(bit-for-bit identical output) but is frame/hierarchy-dependent — re-validate
if the modeling frame or a hierarchy's shapes change.
Parameters

hierarchy The full aggregation hierarchy whose raking shapes gate the blocks. pm_data PopulationModelData used to load the raking shapes (and the modeling frame when one is not supplied). modeling_frame Optional pre-loaded modeling frame. Loaded via pm_data.load_modeling_frame() when None; pass the already-loaded frame to avoid a redundant re-read.

Returns

set[str] The block_keys with at least one intersecting raking polygon.

Raises

ValueError If no block intersects any raking shape while blocks exist — a likely CRS/shapefile misconfiguration that would otherwise silently skip everything.

Source code in src/climate_data/aggregate/utils.py
def blocks_with_shapefile_intersections(
    hierarchy: str,
    pm_data: PopulationModelData,
    modeling_frame: gpd.GeoDataFrame | None = None,
) -> set[str]:
    """Return the block_keys whose footprint intersects the hierarchy's raking shapes.

    Blocks whose (dissolved) modeling-frame geometry intersects no raking polygon
    contribute only empty rows to the sum/sum aggregation, so the pipeline can skip
    them. The shapefile lookup is delegated to
    ``PopulationModelData.load_raking_shapes``, which routes to the right file for
    whichever hierarchy is in use.

    .. note::

        This uses the block *geometry* vs the full shape set, whereas the production
        emptiness gate in ``build_location_masks`` uses the raster *bounding box* vs
        bbox-filtered shapes. Equivalence was validated empirically for ``lsae_1285``
        (bit-for-bit identical output) but is frame/hierarchy-dependent — re-validate
        if the modeling frame or a hierarchy's shapes change.

    Parameters
    ----------
    hierarchy
        The full aggregation hierarchy whose raking shapes gate the blocks.
    pm_data
        PopulationModelData used to load the raking shapes (and the modeling frame
        when one is not supplied).
    modeling_frame
        Optional pre-loaded modeling frame. Loaded via
        ``pm_data.load_modeling_frame()`` when ``None``; pass the already-loaded
        frame to avoid a redundant re-read.

    Returns
    -------
    set[str]
        The block_keys with at least one intersecting raking polygon.

    Raises
    ------
    ValueError
        If no block intersects any raking shape while blocks exist — a likely
        CRS/shapefile misconfiguration that would otherwise silently skip
        everything.
    """
    if modeling_frame is None:
        modeling_frame = pm_data.load_modeling_frame()
    blocks_gdf = (
        modeling_frame[["block_key", "geometry"]].dissolve(by="block_key").reset_index()
    )
    shapes = pm_data.load_raking_shapes(hierarchy)
    if shapes.crs != blocks_gdf.crs:
        shapes = shapes.to_crs(blocks_gdf.crs)
    joined = gpd.sjoin(blocks_gdf, shapes, how="inner", predicate="intersects")
    intersecting = set(joined["block_key"].unique())
    if not intersecting and len(blocks_gdf):
        msg = (
            f"No blocks intersect the raking shapes for hierarchy '{hierarchy}' "
            f"({len(blocks_gdf)} blocks checked) — refusing to silently skip every "
            "block. Check the CRS, the shapefile path, and the sjoin."
        )
        raise ValueError(msg)
    return intersecting

build_bounds_map(raster_template: rt.RasterArray, shape_values: list[tuple[Polygon | MultiPolygon, int]]) -> dict[int, tuple[slice, slice]]

Build a map of location IDs to buffered slices of the raster template.

Parameters

raster_template The raster template to build the bounds map for. shape_values A list of tuples where the first element is a shapely Polygon or MultiPolygon in the CRS of the raster template and the second element is the location ID of the shape.

Returns

dict[int, tuple[slice, slice]] A dictionary mapping location IDs to a tuple of slices representing the bounds of the location in the raster template. The slices are buffered by 10 pixels to ensure that the entire shape is included in the mask.

Source code in src/climate_data/aggregate/utils.py
def build_bounds_map(
    raster_template: rt.RasterArray,
    shape_values: list[tuple[Polygon | MultiPolygon, int]],
) -> dict[int, tuple[slice, slice]]:
    """Build a map of location IDs to buffered slices of the raster template.

    Parameters
    ----------
    raster_template
        The raster template to build the bounds map for.
    shape_values
        A list of tuples where the first element is a shapely Polygon or MultiPolygon
        in the CRS of the raster template and the second element is the location ID
        of the shape.

    Returns
    -------
    dict[int, tuple[slice, slice]]
        A dictionary mapping location IDs to a tuple of slices representing the bounds
        of the location in the raster template. The slices are buffered by 10 pixels
        to ensure that the entire shape is included in the mask.
    """
    # The tranform maps pixel coordinates to the CRS coordinates.
    # This mask is the inverse of that transform.
    to_pixel = ~raster_template.transform

    bounds_map = {}
    for shp, loc_id in shape_values:
        xmin, ymin, xmax, ymax = shp.bounds
        pxmin, pymin = to_pixel * (xmin, ymax)
        pixel_buffer = 10
        pxmin = max(0, int(pxmin) - pixel_buffer)
        pymin = max(0, int(pymin) - pixel_buffer)
        pxmax, pymax = to_pixel * (xmax, ymin)
        pxmax = min(raster_template.width, int(pxmax) + pixel_buffer)
        pymax = min(raster_template.height, int(pymax) + pixel_buffer)
        bounds_map[loc_id] = (slice(pymin, pymax), slice(pxmin, pxmax))

    return bounds_map

build_location_masks(hierarchy: str, block_key: str, pm_data: PopulationModelData) -> tuple[dict[str, slice], dict[int, tuple[slice, slice, npt.NDArray[np.bool_]]], npt.NDArray[np.uint32]]

Build location masks for each location in the hierarchy.

Parameters

hierarchy The name of the hierarchy to build location masks for. Must be one of of the keys of the HIERARCHY_MAP constant. pm_data PopulationModelData object to load the population model data.

Returns

tuple[dict[str, slice], dict[int, tuple[slice, slice, npt.NDArray[np.bool_]]], npt.NDArray[np.uint32]] A three-tuple of: - climate_slice: a dict mapping "longitude"/"latitude" to slices bounding the location block, for subsetting climate rasters before processing (downstream operations scale with the number of pixels in the mask). - bounds_map: a dict mapping each location ID to (row_slice, col_slice, mask), where mask is a boolean array selecting that location's pixels within the sliced window. - location_mask: a 2D uint32 array where each location ID is written as a unique integer value.

Source code in src/climate_data/aggregate/utils.py
def build_location_masks(
    hierarchy: str,
    block_key: str,
    pm_data: PopulationModelData,
) -> tuple[
    dict[str, slice],
    dict[int, tuple[slice, slice, npt.NDArray[np.bool_]]],
    npt.NDArray[np.uint32],
]:
    """Build location masks for each location in the hierarchy.

    Parameters
    ----------
    hierarchy
        The name of the hierarchy to build location masks for. Must be one of
        of the keys of the HIERARCHY_MAP constant.
    pm_data
        PopulationModelData object to load the population model data.

    Returns
    -------
    tuple[dict[str, slice], dict[int, tuple[slice, slice, npt.NDArray[np.bool_]]], npt.NDArray[np.uint32]]
        A three-tuple of:
        - climate_slice: a dict mapping "longitude"/"latitude" to slices bounding
          the location block, for subsetting climate rasters before processing
          (downstream operations scale with the number of pixels in the mask).
        - bounds_map: a dict mapping each location ID to
          (row_slice, col_slice, mask), where mask is a boolean array selecting
          that location's pixels within the sliced window.
        - location_mask: a 2D uint32 array where each location ID is written as a
          unique integer value.
    """
    template = pm_data.load_results("2020q1", block_key)
    template_bbox = get_bbox(template, "EPSG:4326")
    bounds = template_bbox.bounds
    raking_shapes = pm_data.load_raking_shapes(hierarchy, bounds=bounds)
    raking_shapes = raking_shapes[raking_shapes.intersects(template_bbox)].to_crs(
        template.crs
    )

    # Get some bounds to subset the climate rasters with.
    lon_min, lat_min, lon_max, lat_max = bounds
    buffer = 0.1  # Degrees, this is the resolution of the climate raster
    lon_min, lon_max = max(-180, lon_min - buffer), min(180, lon_max + buffer)
    lat_min, lat_max = max(-90, lat_min - buffer), min(90, lat_max + buffer)
    climate_slice = {
        "longitude": slice(lon_min, lon_max),
        "latitude": slice(lat_max, lat_min),
    }

    shape_values = [
        (shape, loc_id)
        for loc_id, shape in raking_shapes.set_index("location_id")
        .geometry.to_dict()
        .items()
    ]
    bounds_map = build_bounds_map(template, shape_values)

    location_mask = np.zeros_like(template, dtype=np.uint32)
    location_mask = rasterize(
        shape_values,
        out=location_mask,
        transform=template.transform,
        merge_alg=MergeAlg.replace,
    )
    final_bounds_map = {
        location_id: (rows, cols, location_mask[rows, cols] == location_id)
        for location_id, (rows, cols) in bounds_map.items()
    }
    return climate_slice, final_bounds_map, location_mask

get_bbox(raster: rt.RasterArray, crs: str | None = None) -> shapely.Polygon

Get the bounding box of a raster array.

Parameters

raster The raster array to get the bounding box of. crs The CRS to return the bounding box in. If None, the bounding box is returned in the CRS of the raster.

Returns

shapely.Polybon The bounding box of the raster in the CRS specified by the crs parameter.

Source code in src/climate_data/aggregate/utils.py
def get_bbox(raster: rt.RasterArray, crs: str | None = None) -> shapely.Polygon:
    """Get the bounding box of a raster array.

    Parameters
    ----------
    raster
        The raster array to get the bounding box of.
    crs
        The CRS to return the bounding box in. If None, the bounding box
        is returned in the CRS of the raster.

    Returns
    -------
    shapely.Polybon
        The bounding box of the raster in the CRS specified by the crs parameter.
    """
    if raster.crs not in MAX_BOUNDS:
        msg = f"Unsupported CRS: {raster.crs}"
        raise ValueError(msg)

    xmin_clip, xmax_clip, ymin_clip, ymax_clip = MAX_BOUNDS[raster.crs]
    xmin, xmax, ymin, ymax = raster.bounds

    xmin = np.clip(xmin, xmin_clip, xmax_clip)
    xmax = np.clip(xmax, xmin_clip, xmax_clip)
    ymin = np.clip(ymin, ymin_clip, ymax_clip)
    ymax = np.clip(ymax, ymin_clip, ymax_clip)

    bbox = gpd.GeoSeries([shapely.box(xmin, ymin, xmax, ymax)], crs=raster.crs)
    out_bbox = bbox.to_crs(crs) if crs is not None else bbox.copy()

    # Check that our transformation didn't do something weird
    # (e.g. artificially clip the bounds or have the bounds extend over the
    # antimeridian)
    check_bbox = out_bbox.to_crs(raster.crs)
    area_change = (np.abs(bbox.area - check_bbox.area) / bbox.area).iloc[0]
    tolerance = 1e-6
    if area_change > tolerance:
        msg = f"Area change: {area_change}"
        raise ValueError(msg)

    return cast(shapely.Polygon, out_bbox.iloc[0])

cli

cdrun() -> None

Entry point for running climate downscale workflows.

Source code in src/climate_data/cli.py
6
7
8
@click.group()
def cdrun() -> None:
    """Entry point for running climate downscale workflows."""

cdtask() -> None

Entry point for running climate downscale tasks.

Source code in src/climate_data/cli.py
@click.group()
def cdtask() -> None:
    """Entry point for running climate downscale tasks."""

cli_options

Climate Data CLI Options

This module provides a set of CLI options for extracting climate data from the ERA5 and CMIP6 datasets. These options are used to specify the data to extract, such as the year, month, variable, and dataset. It also provides global variables representing the full space of valid values for these options.

resolve_run_mode_root(param_name: str, value: str, run_mode: str, *, aggregate: bool = False) -> str

Resolve a directory option to the run-mode root unless the user overrode it.

--run-mode selects the storage-root profile; an explicitly-passed --<param> still wins (detected via click's parameter source).

Source code in src/climate_data/cli_options.py
def resolve_run_mode_root(
    param_name: str, value: str, run_mode: str, *, aggregate: bool = False
) -> str:
    """Resolve a directory option to the run-mode root unless the user overrode it.

    ``--run-mode`` selects the storage-root profile; an explicitly-passed
    ``--<param>`` still wins (detected via click's parameter source).
    """
    ctx = click.get_current_context(silent=True)
    overridden = (
        ctx is not None
        and ctx.get_parameter_source(param_name)
        is not click.core.ParameterSource.DEFAULT
    )
    if overridden:
        return value
    root = cdc.aggregate_root(run_mode) if aggregate else cdc.model_root(run_mode)
    return str(root)

with_agg_measure(*, allow_all: bool = False) -> Callable[[Callable[P, T]], Callable[P, T]]

Add aggregation measure option to a command.

Source code in src/climate_data/cli_options.py
def with_agg_measure[**P, T](
    *,
    allow_all: bool = False,
) -> Callable[[Callable[P, T]], Callable[P, T]]:
    """Add aggregation measure option to a command."""
    return with_choice(
        "agg-measure",
        allow_all=allow_all,
        choices=cdc.AGGREGATION_MEASURES,
        help="Climate measure to process.",
    )

with_agg_scenario(*, allow_all: bool = False) -> Callable[[Callable[P, T]], Callable[P, T]]

Add aggregation scenario option to a command.

Source code in src/climate_data/cli_options.py
def with_agg_scenario[**P, T](
    *,
    allow_all: bool = False,
) -> Callable[[Callable[P, T]], Callable[P, T]]:
    """Add aggregation scenario option to a command."""
    return with_choice(
        "agg-scenario",
        allow_all=allow_all,
        choices=cdc.AGGREGATION_SCENARIOS,
        help="Climate scenario to process.",
    )

with_agg_version() -> Callable[[Callable[P, T]], Callable[P, T]]

Add aggregation version option to a command.

Source code in src/climate_data/cli_options.py
def with_agg_version[**P, T]() -> Callable[[Callable[P, T]], Callable[P, T]]:
    """Add aggregation version option to a command."""
    return click.option(
        "--agg-version",
        help="Aggregation version to process.",
        required=True,
    )

with_block_key(*, allow_all: bool = False) -> Callable[[Callable[P, T]], Callable[P, T]]

Add block key option to a command.

Source code in src/climate_data/cli_options.py
def with_block_key[**P, T](
    *,
    allow_all: bool = False,
) -> Callable[[Callable[P, T]], Callable[P, T]]:
    """Add block key option to a command."""
    return with_choice(
        "block-key",
        allow_all=allow_all,
        choices=None,  # Will be populated at runtime
        help="Block key to process.",
    )

with_concurrency_limit(*, default: int | None = None) -> Callable[[Callable[P, T]], Callable[P, T]]

Add the jobmon concurrency-limit option to a command.

Pass default to throttle a runner out of the box; leaving it unset means the option defaults to None, which run_parallel_maybe_dry_run drops so jobmon applies its own default (10000, effectively unthrottled).

Source code in src/climate_data/cli_options.py
def with_concurrency_limit[**P, T](
    *,
    default: int | None = None,
) -> Callable[[Callable[P, T]], Callable[P, T]]:
    """Add the jobmon concurrency-limit option to a command.

    Pass ``default`` to throttle a runner out of the box; leaving it unset means
    the option defaults to ``None``, which ``run_parallel_maybe_dry_run`` drops
    so jobmon applies its own default (10000, effectively unthrottled).
    """
    return click.option(
        "--concurrency-limit",
        type=click.INT,
        default=default,
        show_default=True,
        help=(
            "Cap on tasks running at once, to keep write latency on shared "
            "storage manageable. jobmon's own default is 10000, which is "
            "effectively unthrottled."
        ),
    )

with_debias_method() -> Callable[[Callable[P, T]], Callable[P, T]]

Add the option selecting the Jensen de-bias applied to a multiplicative anomaly.

Source code in src/climate_data/cli_options.py
def with_debias_method[**P, T]() -> Callable[[Callable[P, T]], Callable[P, T]]:
    """Add the option selecting the Jensen de-bias applied to a multiplicative anomaly."""
    return click.option(
        "--debias-method",
        type=click.Choice(list(cdc.DEBIAS_METHODS)),
        default="none",
        show_default=True,
        help=(
            "Correct the ratio-estimator (Jensen) bias in the multiplicative anomaly: "
            "'none' (ship as-is), 'loo' (leave-one-out, a direct out-of-sample estimate), "
            "or 'analytic' (second-order expansion). Only valid for "
            f"{', '.join(cdc.DEBIAS_VARIABLES)}."
        ),
    )

with_dry_day_rule() -> Callable[[Callable[P, T]], Callable[P, T]]

Add the option selecting how days the driving model reports as dry are treated.

Source code in src/climate_data/cli_options.py
def with_dry_day_rule[**P, T]() -> Callable[[Callable[P, T]], Callable[P, T]]:
    """Add the option selecting how days the driving model reports as dry are treated."""
    return click.option(
        "--dry-day-rule",
        type=click.Choice(list(cdc.DRY_DAY_RULES)),
        default="none",
        show_default=True,
        help=(
            "Treatment of dry model days in the multiplicative anomaly: 'none' (ship as-is) "
            "or 'preserve' (zero the anomaly on dry model days and renormalise the "
            "cell-month, leaving the monthly total unchanged). Only valid for "
            f"{', '.join(cdc.DRY_DAY_VARIABLES)}."
        ),
    )

with_dry_run() -> Callable[[Callable[P, T]], Callable[P, T]]

Add dry-run flag to a command.

Source code in src/climate_data/cli_options.py
def with_dry_run[**P, T]() -> Callable[[Callable[P, T]], Callable[P, T]]:
    """Add dry-run flag to a command."""
    return click.option(
        "--dry-run/--no-dry-run",
        "-n",
        default=False,
        show_default=True,
        help="Print sbatch-like previews instead of submitting jobs.",
    )

with_hierarchy(choices: Collection[str] = cdc.HIERARCHY_MAP, *, allow_all: bool = False, default: str | None = None) -> Callable[[Callable[P, T]], Callable[P, T]]

Add hierarchy option to a command.

Pass default to give the (single-value) option a default; this builds the option directly since with_choice always supplies its own default.

Source code in src/climate_data/cli_options.py
def with_hierarchy[**P, T](
    choices: Collection[str] = cdc.HIERARCHY_MAP,
    *,
    allow_all: bool = False,
    default: str | None = None,
) -> Callable[[Callable[P, T]], Callable[P, T]]:
    """Add hierarchy option to a command.

    Pass ``default`` to give the (single-value) option a default; this builds
    the option directly since ``with_choice`` always supplies its own default.
    """
    if default is not None:
        return click.option(
            "--hierarchy",
            type=click.Choice(list(choices)),
            default=default,
            show_default=True,
            help="Hierarchy to process.",
        )
    return with_choice(
        "hierarchy",
        allow_all=allow_all,
        choices=choices,
        help="Hierarchy to process.",
        convert=allow_all,
    )

with_location_id() -> Callable[[Callable[P, T]], Callable[P, T]]

Add location ID option to a command.

Source code in src/climate_data/cli_options.py
def with_location_id[**P, T]() -> Callable[[Callable[P, T]], Callable[P, T]]:
    """Add location ID option to a command."""
    return click.option(
        "--location-id",
        "-l",
        type=click.INT,
        help="Location ID to process.",
    )

with_run_mode() -> Callable[[Callable[P, T]], Callable[P, T]]

Add the run-mode option selecting the storage-root profile.

Source code in src/climate_data/cli_options.py
def with_run_mode[**P, T]() -> Callable[[Callable[P, T]], Callable[P, T]]:
    """Add the run-mode option selecting the storage-root profile."""
    return click.option(
        "--run-mode",
        type=click.Choice(list(cdc.RUN_MODES)),
        default="forecast",
        show_default=True,
        help=(
            "Storage-root profile: 'forecast' (production roots) or 'historical' "
            "(geospatial roots for the GBD historical exposure product)."
        ),
    )

with_year(years: Collection[str], *, allow_all: bool = False) -> Callable[[Callable[P, T]], Callable[P, T]]

Create a CLI option for selecting a year.

Source code in src/climate_data/cli_options.py
def with_year[**P, T](
    years: Collection[str],
    *,
    allow_all: bool = False,
) -> Callable[[Callable[P, T]], Callable[P, T]]:
    """Create a CLI option for selecting a year."""
    return with_choice(
        "year",
        "y",
        allow_all=allow_all,
        choices=years,
        help="Year to extract data for.",
        convert=allow_all,
    )

constants

aggregate_root(run_mode: str) -> Path

Aggregation working directory for the given run mode.

Source code in src/climate_data/constants.py
def aggregate_root(run_mode: str) -> Path:
    """Aggregation working directory for the given run mode."""
    if run_mode == "forecast":
        return RRA_ROOT / "climate-aggregates"
    if run_mode == "historical":
        return model_root("historical") / "aggregates"
    msg = f"Unknown run_mode {run_mode!r}; expected one of {list(RUN_MODES)}"
    raise ValueError(msg)

model_root(run_mode: str) -> Path

Downscaling working directory for the given run mode.

Source code in src/climate_data/constants.py
def model_root(run_mode: str) -> Path:
    """Downscaling working directory for the given run mode."""
    try:
        return _MODEL_ROOTS[run_mode]
    except KeyError:
        msg = f"Unknown run_mode {run_mode!r}; expected one of {list(RUN_MODES)}"
        raise ValueError(msg) from None

data

Climate Data Management

This module provides a class for managing the climate data used in the project. It includes methods for loading and saving data, as well as for accessing the various directories where data is stored. This abstraction allows for easy access to the data and ensures that all data is stored in a consistent and organized manner. It also provides a central location for managing the data, which makes it easier to update and maintain the path structure of the data as needed.

This module generally does not load or process data itself, though some exceptions are made for metadata which is generally loaded and cached on disk.

The main classes are: - PopulationModelData: Handles data from the gridded population modeling pipeline. This includes population estimates and projections as well as the location hierarchies for the population data. This class provides read-only access to the data. - ClimateData: Handles gridded climate data from the climate downscaling pipeline. This includes climate data for different scenarios and measures. This class provides read and write access to the data. - ClimateAggregateData: Handles the output data structure for climate aggregates. This includes raw results at the block level, final results at the measure level, and versioned results for different pipeline versions. This class provides both read and write access to the data.

ClimateAggregateData

Manages the output data structure for climate aggregates.

This class manages the file organization and paths for: 1. Reading and writing raw results at block level 2. Reading and writing final results at measure and scenario level 3. Versioning of results

Source code in src/climate_data/data.py
 797
 798
 799
 800
 801
 802
 803
 804
 805
 806
 807
 808
 809
 810
 811
 812
 813
 814
 815
 816
 817
 818
 819
 820
 821
 822
 823
 824
 825
 826
 827
 828
 829
 830
 831
 832
 833
 834
 835
 836
 837
 838
 839
 840
 841
 842
 843
 844
 845
 846
 847
 848
 849
 850
 851
 852
 853
 854
 855
 856
 857
 858
 859
 860
 861
 862
 863
 864
 865
 866
 867
 868
 869
 870
 871
 872
 873
 874
 875
 876
 877
 878
 879
 880
 881
 882
 883
 884
 885
 886
 887
 888
 889
 890
 891
 892
 893
 894
 895
 896
 897
 898
 899
 900
 901
 902
 903
 904
 905
 906
 907
 908
 909
 910
 911
 912
 913
 914
 915
 916
 917
 918
 919
 920
 921
 922
 923
 924
 925
 926
 927
 928
 929
 930
 931
 932
 933
 934
 935
 936
 937
 938
 939
 940
 941
 942
 943
 944
 945
 946
 947
 948
 949
 950
 951
 952
 953
 954
 955
 956
 957
 958
 959
 960
 961
 962
 963
 964
 965
 966
 967
 968
 969
 970
 971
 972
 973
 974
 975
 976
 977
 978
 979
 980
 981
 982
 983
 984
 985
 986
 987
 988
 989
 990
 991
 992
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
class ClimateAggregateData:
    """Manages the output data structure for climate aggregates.

    This class manages the file organization and paths for:
    1. Reading and writing raw results at block level
    2. Reading and writing final results at measure and scenario level
    3. Versioning of results
    """

    def __init__(
        self,
        root: str | Path = cdc.AGGREGATE_ROOT,
        *,
        read_only: bool = False,
    ) -> None:
        """Initialize the climate aggregate data manager.

        Parameters
        ----------
        root
            Path to the model root directory
        read_only
            If ``True``, do not create the root or logs directories. Use this for
            path construction and reads (including dry runs), so merely building a
            manager never writes to shared storage.
        """
        self._root = Path(root)
        self._read_only = read_only
        if not read_only:
            self._create_model_root()

    def _create_model_root(self) -> None:
        """Create the model root directory and logs directory."""
        mkdir(self.root, exist_ok=True)
        mkdir(self.logs, exist_ok=True)

    @property
    def root(self) -> Path:
        """Get the root directory for model data."""
        return self._root

    @property
    def logs(self) -> Path:
        """Get the directory for log files."""
        return self.root / "logs"

    def log_dir(self, step_name: str) -> Path:
        """Get the directory for logs from a specific pipeline step.

        Parameters
        ----------
        step_name
            The name of the pipeline step

        Returns
        -------
        Path
            The directory for step-specific logs
        """
        return self.logs / step_name

    def person_days_path(self, block_key: str, scenario: str, gcm_member: str) -> Path:
        """Path to a raw per-block person-days file.

        Different GBD vintages (gbd_2021/2023/2025, which have different location
        sets) are kept apart by the versioned root (``<output_dir>/<hierarchy>``,
        e.g. ``.../aggregates/gbd_2023``), mirroring the aggregate stage's
        ``version_root`` convention -- not by a segment in this path. The special
        runners build that root by appending ``hierarchy`` to ``output_dir``.
        """
        return (
            self.root
            / "erf-scratch"
            / "person-days"
            / block_key
            / f"{scenario}_{gcm_member}.parquet"
        )

    def load_person_days(
        self, block_key: str, scenario: str, gcm_member: str
    ) -> pd.DataFrame:
        """Load a raw per-block person-days file."""
        return pd.read_parquet(self.person_days_path(block_key, scenario, gcm_member))

    def compiled_person_days_path(
        self, subset_hierarchy: str, scenario: str, gcm_member: str
    ) -> Path:
        """Path to a compiled (hierarchy-aggregated) person-days file."""
        return (
            self.root
            / "erf-scratch"
            / "compiled-person-days"
            / subset_hierarchy
            / f"{scenario}_{gcm_member}.parquet"
        )

    def version_root(self, version: str) -> Path:
        """Get the root directory for a specific version.

        Parameters
        ----------
        version
            The version identifier

        Returns
        -------
        Path
            The directory for version-specific data
        """
        return self.root / version

    def raw_results_root(self, version: str) -> Path:
        """Get the directory for raw results (block-level).

        Parameters
        ----------
        version
            The version identifier

        Returns
        -------
        Path
            The directory for raw results
        """
        return self.version_root(version) / "raw-results"

    def raw_results_path(
        self, version: str, hierarchy: str, block_key: str, draw: str
    ) -> Path:
        """Get the path to raw results for a specific hierarchy, block, and draw.

        Parameters
        ----------
        version
            The version identifier
        hierarchy
            The location hierarchy
        block_key
            The block key
        draw
            The draw of the climate data (e.g. "000")

        Returns
        -------
        Path
            The path to the raw results file
        """
        root = self.raw_results_root(version)
        return root / hierarchy / block_key / f"{draw}.parquet"

    def save_raw_results(
        self,
        df: pd.DataFrame,
        version: str,
        hierarchy: str,
        block_key: str,
        draw: str,
    ) -> None:
        """Save raw results for a specific hierarchy, block, and draw.

        Parameters
        ----------
        df
            The results to save
        version
            The version identifier
        hierarchy
            The location hierarchy
        block_key
            The block key
        draw
            The draw of the climate data to save (e.g. "000")
        """
        path = self.raw_results_path(version, hierarchy, block_key, draw)
        mkdir(path.parent, exist_ok=True, parents=True)
        save_parquet(df, path)

    def load_raw_results(
        self,
        version: str,
        hierarchy: str,
        block_key: str,
        draw: str,
        measure: str | None = None,
        scenario: str | None = None,
    ) -> pd.DataFrame:
        """Load raw results for a specific hierarchy, block, and draw.

        Parameters
        ----------
        version
            The version identifier
        hierarchy
            The location hierarchy
        block_key
            The block key
        draw
            The draw of the climate data to load (e.g. "000")
        measure
            If provided, filter results to only include this measure
        scenario
            If provided, filter results to only include this scenario

        Returns
        -------
        pd.DataFrame
            The raw results
        """
        path = self.raw_results_path(version, hierarchy, block_key, draw)

        # Build filters for parquet's read_parquet function
        filters = []
        if measure is not None:
            filters.append(("measure", "==", measure))
        if scenario is not None:
            filters.append(("scenario", "==", scenario))

        return pd.read_parquet(path, filters=filters)

    def results_root(self, version: str) -> Path:
        """Get the directory for final results (measure-level).

        Parameters
        ----------
        version
            The version identifier

        Returns
        -------
        Path
            The directory for final results
        """
        return self.version_root(version) / "results"

    def population_path(self, version: str, hierarchy: str) -> Path:
        """Get the path to population data for a specific hierarchy.

        Parameters
        ----------
        version
            The version identifier
        hierarchy
            The location hierarchy

        Returns
        -------
        Path
            The path to the population data file
        """
        return self.results_root(version) / hierarchy / "population.parquet"

    def save_population(self, df: pd.DataFrame, version: str, hierarchy: str) -> None:
        """Save population data for a specific hierarchy.

        Parameters
        ----------
        df
            The population data to save
        version
            The version identifier
        hierarchy
            The location hierarchy
        """
        path = self.population_path(version, hierarchy)
        mkdir(path.parent, exist_ok=True, parents=True)
        save_parquet(df, path)

    def load_population(
        self, version: str, hierarchy: str, location_id: int | None = None
    ) -> pd.DataFrame:
        """Load population data for a specific hierarchy and optionally location.

        Parameters
        ----------
        version
            The version identifier
        hierarchy
            The location hierarchy
        location_id
            If provided, load only data for this location

        Returns
        -------
        pd.DataFrame
            The population data
        """
        path = self.population_path(version, hierarchy)
        if location_id is not None:
            filters = [("location_id", "==", location_id)]
            return pd.read_parquet(path, filters=filters)
        return pd.read_parquet(path)

    def results_path(
        self, version: str, hierarchy: str, scenario: str, measure: str
    ) -> Path:
        """Get the path to final results for a specific scenario and measure.

        Parameters
        ----------
        version
            The version identifier
        hierarchy
            The location hierarchy
        scenario
            The climate scenario
        measure
            The climate measure

        Returns
        -------
        Path
            The path to the results file
        """
        return self.results_root(version) / hierarchy / f"{measure}_{scenario}.parquet"

    def save_results(
        self,
        df: pd.DataFrame,
        version: str,
        hierarchy: str,
        scenario: str,
        measure: str,
    ) -> None:
        """Save final results for a specific scenario and measure.

        Parameters
        ----------
        df
            The results to save
        version
            The version identifier
        hierarchy
            The location hierarchy
        scenario
            The climate scenario
        measure
            The climate measure
        """
        path = self.results_path(version, hierarchy, scenario, measure)
        mkdir(path.parent, exist_ok=True, parents=True)
        save_parquet(df, path)

    def load_results(
        self,
        version: str,
        hierarchy: str,
        scenario: str,
        measure: str,
        location_id: int | None = None,
    ) -> pd.DataFrame:
        """Load final results for a specific scenario and measure.

        Parameters
        ----------
        version
            The version identifier
        hierarchy
            The location hierarchy
        scenario
            The climate scenario
        measure
            The climate measure
        location_id
            If provided, load only data for this location

        Returns
        -------
        pd.DataFrame
            The results
        """
        path = self.results_path(version, hierarchy, scenario, measure)
        if location_id is not None:
            filters = [("location_id", "==", location_id)]
            return pd.read_parquet(path, filters=filters)
        return pd.read_parquet(path)

    def diagnostics_root(self, version: str, hierarchy: str) -> Path:
        """Get the path to the diagnostics directory.

        Parameters
        ----------
        version
            The version identifier

        Returns
        -------
        Path
            The path to the diagnostics directory
        """
        return self.version_root(version) / "diagnostics" / hierarchy

    def grid_plots_pages_root(self, version: str, hierarchy: str) -> Path:
        return self.diagnostics_root(version, hierarchy) / "grid_plots_pages"

    def grid_plots_page_path(
        self, version: str, hierarchy: str, location_id: int
    ) -> Path:
        path = self.grid_plots_pages_root(version, hierarchy) / f"{location_id}.pdf"
        mkdir(path.parent, exist_ok=True, parents=True)
        return path

    def grid_plots_path(self, version: str, hierarchy: str) -> Path:
        path = self.diagnostics_root(version, hierarchy) / f"grid_plots_{hierarchy}.pdf"
        mkdir(path.parent, exist_ok=True, parents=True)
        return path

logs: Path property

Get the directory for log files.

root: Path property

Get the root directory for model data.

__init__(root: str | Path = cdc.AGGREGATE_ROOT, *, read_only: bool = False) -> None

Initialize the climate aggregate data manager.

Parameters

root Path to the model root directory read_only If True, do not create the root or logs directories. Use this for path construction and reads (including dry runs), so merely building a manager never writes to shared storage.

Source code in src/climate_data/data.py
def __init__(
    self,
    root: str | Path = cdc.AGGREGATE_ROOT,
    *,
    read_only: bool = False,
) -> None:
    """Initialize the climate aggregate data manager.

    Parameters
    ----------
    root
        Path to the model root directory
    read_only
        If ``True``, do not create the root or logs directories. Use this for
        path construction and reads (including dry runs), so merely building a
        manager never writes to shared storage.
    """
    self._root = Path(root)
    self._read_only = read_only
    if not read_only:
        self._create_model_root()

compiled_person_days_path(subset_hierarchy: str, scenario: str, gcm_member: str) -> Path

Path to a compiled (hierarchy-aggregated) person-days file.

Source code in src/climate_data/data.py
def compiled_person_days_path(
    self, subset_hierarchy: str, scenario: str, gcm_member: str
) -> Path:
    """Path to a compiled (hierarchy-aggregated) person-days file."""
    return (
        self.root
        / "erf-scratch"
        / "compiled-person-days"
        / subset_hierarchy
        / f"{scenario}_{gcm_member}.parquet"
    )

diagnostics_root(version: str, hierarchy: str) -> Path

Get the path to the diagnostics directory.

Parameters

version The version identifier

Returns

Path The path to the diagnostics directory

Source code in src/climate_data/data.py
def diagnostics_root(self, version: str, hierarchy: str) -> Path:
    """Get the path to the diagnostics directory.

    Parameters
    ----------
    version
        The version identifier

    Returns
    -------
    Path
        The path to the diagnostics directory
    """
    return self.version_root(version) / "diagnostics" / hierarchy

load_person_days(block_key: str, scenario: str, gcm_member: str) -> pd.DataFrame

Load a raw per-block person-days file.

Source code in src/climate_data/data.py
def load_person_days(
    self, block_key: str, scenario: str, gcm_member: str
) -> pd.DataFrame:
    """Load a raw per-block person-days file."""
    return pd.read_parquet(self.person_days_path(block_key, scenario, gcm_member))

load_population(version: str, hierarchy: str, location_id: int | None = None) -> pd.DataFrame

Load population data for a specific hierarchy and optionally location.

Parameters

version The version identifier hierarchy The location hierarchy location_id If provided, load only data for this location

Returns

pd.DataFrame The population data

Source code in src/climate_data/data.py
def load_population(
    self, version: str, hierarchy: str, location_id: int | None = None
) -> pd.DataFrame:
    """Load population data for a specific hierarchy and optionally location.

    Parameters
    ----------
    version
        The version identifier
    hierarchy
        The location hierarchy
    location_id
        If provided, load only data for this location

    Returns
    -------
    pd.DataFrame
        The population data
    """
    path = self.population_path(version, hierarchy)
    if location_id is not None:
        filters = [("location_id", "==", location_id)]
        return pd.read_parquet(path, filters=filters)
    return pd.read_parquet(path)

load_raw_results(version: str, hierarchy: str, block_key: str, draw: str, measure: str | None = None, scenario: str | None = None) -> pd.DataFrame

Load raw results for a specific hierarchy, block, and draw.

Parameters

version The version identifier hierarchy The location hierarchy block_key The block key draw The draw of the climate data to load (e.g. "000") measure If provided, filter results to only include this measure scenario If provided, filter results to only include this scenario

Returns

pd.DataFrame The raw results

Source code in src/climate_data/data.py
def load_raw_results(
    self,
    version: str,
    hierarchy: str,
    block_key: str,
    draw: str,
    measure: str | None = None,
    scenario: str | None = None,
) -> pd.DataFrame:
    """Load raw results for a specific hierarchy, block, and draw.

    Parameters
    ----------
    version
        The version identifier
    hierarchy
        The location hierarchy
    block_key
        The block key
    draw
        The draw of the climate data to load (e.g. "000")
    measure
        If provided, filter results to only include this measure
    scenario
        If provided, filter results to only include this scenario

    Returns
    -------
    pd.DataFrame
        The raw results
    """
    path = self.raw_results_path(version, hierarchy, block_key, draw)

    # Build filters for parquet's read_parquet function
    filters = []
    if measure is not None:
        filters.append(("measure", "==", measure))
    if scenario is not None:
        filters.append(("scenario", "==", scenario))

    return pd.read_parquet(path, filters=filters)

load_results(version: str, hierarchy: str, scenario: str, measure: str, location_id: int | None = None) -> pd.DataFrame

Load final results for a specific scenario and measure.

Parameters

version The version identifier hierarchy The location hierarchy scenario The climate scenario measure The climate measure location_id If provided, load only data for this location

Returns

pd.DataFrame The results

Source code in src/climate_data/data.py
def load_results(
    self,
    version: str,
    hierarchy: str,
    scenario: str,
    measure: str,
    location_id: int | None = None,
) -> pd.DataFrame:
    """Load final results for a specific scenario and measure.

    Parameters
    ----------
    version
        The version identifier
    hierarchy
        The location hierarchy
    scenario
        The climate scenario
    measure
        The climate measure
    location_id
        If provided, load only data for this location

    Returns
    -------
    pd.DataFrame
        The results
    """
    path = self.results_path(version, hierarchy, scenario, measure)
    if location_id is not None:
        filters = [("location_id", "==", location_id)]
        return pd.read_parquet(path, filters=filters)
    return pd.read_parquet(path)

log_dir(step_name: str) -> Path

Get the directory for logs from a specific pipeline step.

Parameters

step_name The name of the pipeline step

Returns

Path The directory for step-specific logs

Source code in src/climate_data/data.py
def log_dir(self, step_name: str) -> Path:
    """Get the directory for logs from a specific pipeline step.

    Parameters
    ----------
    step_name
        The name of the pipeline step

    Returns
    -------
    Path
        The directory for step-specific logs
    """
    return self.logs / step_name

person_days_path(block_key: str, scenario: str, gcm_member: str) -> Path

Path to a raw per-block person-days file.

Different GBD vintages (gbd_2021/2023/2025, which have different location sets) are kept apart by the versioned root (<output_dir>/<hierarchy>, e.g. .../aggregates/gbd_2023), mirroring the aggregate stage's version_root convention -- not by a segment in this path. The special runners build that root by appending hierarchy to output_dir.

Source code in src/climate_data/data.py
def person_days_path(self, block_key: str, scenario: str, gcm_member: str) -> Path:
    """Path to a raw per-block person-days file.

    Different GBD vintages (gbd_2021/2023/2025, which have different location
    sets) are kept apart by the versioned root (``<output_dir>/<hierarchy>``,
    e.g. ``.../aggregates/gbd_2023``), mirroring the aggregate stage's
    ``version_root`` convention -- not by a segment in this path. The special
    runners build that root by appending ``hierarchy`` to ``output_dir``.
    """
    return (
        self.root
        / "erf-scratch"
        / "person-days"
        / block_key
        / f"{scenario}_{gcm_member}.parquet"
    )

population_path(version: str, hierarchy: str) -> Path

Get the path to population data for a specific hierarchy.

Parameters

version The version identifier hierarchy The location hierarchy

Returns

Path The path to the population data file

Source code in src/climate_data/data.py
def population_path(self, version: str, hierarchy: str) -> Path:
    """Get the path to population data for a specific hierarchy.

    Parameters
    ----------
    version
        The version identifier
    hierarchy
        The location hierarchy

    Returns
    -------
    Path
        The path to the population data file
    """
    return self.results_root(version) / hierarchy / "population.parquet"

raw_results_path(version: str, hierarchy: str, block_key: str, draw: str) -> Path

Get the path to raw results for a specific hierarchy, block, and draw.

Parameters

version The version identifier hierarchy The location hierarchy block_key The block key draw The draw of the climate data (e.g. "000")

Returns

Path The path to the raw results file

Source code in src/climate_data/data.py
def raw_results_path(
    self, version: str, hierarchy: str, block_key: str, draw: str
) -> Path:
    """Get the path to raw results for a specific hierarchy, block, and draw.

    Parameters
    ----------
    version
        The version identifier
    hierarchy
        The location hierarchy
    block_key
        The block key
    draw
        The draw of the climate data (e.g. "000")

    Returns
    -------
    Path
        The path to the raw results file
    """
    root = self.raw_results_root(version)
    return root / hierarchy / block_key / f"{draw}.parquet"

raw_results_root(version: str) -> Path

Get the directory for raw results (block-level).

Parameters

version The version identifier

Returns

Path The directory for raw results

Source code in src/climate_data/data.py
def raw_results_root(self, version: str) -> Path:
    """Get the directory for raw results (block-level).

    Parameters
    ----------
    version
        The version identifier

    Returns
    -------
    Path
        The directory for raw results
    """
    return self.version_root(version) / "raw-results"

results_path(version: str, hierarchy: str, scenario: str, measure: str) -> Path

Get the path to final results for a specific scenario and measure.

Parameters

version The version identifier hierarchy The location hierarchy scenario The climate scenario measure The climate measure

Returns

Path The path to the results file

Source code in src/climate_data/data.py
def results_path(
    self, version: str, hierarchy: str, scenario: str, measure: str
) -> Path:
    """Get the path to final results for a specific scenario and measure.

    Parameters
    ----------
    version
        The version identifier
    hierarchy
        The location hierarchy
    scenario
        The climate scenario
    measure
        The climate measure

    Returns
    -------
    Path
        The path to the results file
    """
    return self.results_root(version) / hierarchy / f"{measure}_{scenario}.parquet"

results_root(version: str) -> Path

Get the directory for final results (measure-level).

Parameters

version The version identifier

Returns

Path The directory for final results

Source code in src/climate_data/data.py
def results_root(self, version: str) -> Path:
    """Get the directory for final results (measure-level).

    Parameters
    ----------
    version
        The version identifier

    Returns
    -------
    Path
        The directory for final results
    """
    return self.version_root(version) / "results"

save_population(df: pd.DataFrame, version: str, hierarchy: str) -> None

Save population data for a specific hierarchy.

Parameters

df The population data to save version The version identifier hierarchy The location hierarchy

Source code in src/climate_data/data.py
def save_population(self, df: pd.DataFrame, version: str, hierarchy: str) -> None:
    """Save population data for a specific hierarchy.

    Parameters
    ----------
    df
        The population data to save
    version
        The version identifier
    hierarchy
        The location hierarchy
    """
    path = self.population_path(version, hierarchy)
    mkdir(path.parent, exist_ok=True, parents=True)
    save_parquet(df, path)

save_raw_results(df: pd.DataFrame, version: str, hierarchy: str, block_key: str, draw: str) -> None

Save raw results for a specific hierarchy, block, and draw.

Parameters

df The results to save version The version identifier hierarchy The location hierarchy block_key The block key draw The draw of the climate data to save (e.g. "000")

Source code in src/climate_data/data.py
def save_raw_results(
    self,
    df: pd.DataFrame,
    version: str,
    hierarchy: str,
    block_key: str,
    draw: str,
) -> None:
    """Save raw results for a specific hierarchy, block, and draw.

    Parameters
    ----------
    df
        The results to save
    version
        The version identifier
    hierarchy
        The location hierarchy
    block_key
        The block key
    draw
        The draw of the climate data to save (e.g. "000")
    """
    path = self.raw_results_path(version, hierarchy, block_key, draw)
    mkdir(path.parent, exist_ok=True, parents=True)
    save_parquet(df, path)

save_results(df: pd.DataFrame, version: str, hierarchy: str, scenario: str, measure: str) -> None

Save final results for a specific scenario and measure.

Parameters

df The results to save version The version identifier hierarchy The location hierarchy scenario The climate scenario measure The climate measure

Source code in src/climate_data/data.py
def save_results(
    self,
    df: pd.DataFrame,
    version: str,
    hierarchy: str,
    scenario: str,
    measure: str,
) -> None:
    """Save final results for a specific scenario and measure.

    Parameters
    ----------
    df
        The results to save
    version
        The version identifier
    hierarchy
        The location hierarchy
    scenario
        The climate scenario
    measure
        The climate measure
    """
    path = self.results_path(version, hierarchy, scenario, measure)
    mkdir(path.parent, exist_ok=True, parents=True)
    save_parquet(df, path)

version_root(version: str) -> Path

Get the root directory for a specific version.

Parameters

version The version identifier

Returns

Path The directory for version-specific data

Source code in src/climate_data/data.py
def version_root(self, version: str) -> Path:
    """Get the root directory for a specific version.

    Parameters
    ----------
    version
        The version identifier

    Returns
    -------
    Path
        The directory for version-specific data
    """
    return self.root / version

ClimateData

Class for managing the climate data used in the project.

Source code in src/climate_data/data.py
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
class ClimateData:
    """Class for managing the climate data used in the project."""

    def __init__(
        self,
        root: str | Path = cdc.MODEL_ROOT,
        *,
        read_only: bool = False,
    ) -> None:
        self._root = Path(root)
        self._credentials_root = self._root / "credentials"
        self._read_only = read_only
        if not read_only:
            self._create_model_root()

    def _create_model_root(self) -> None:
        mkdir(self.root, exist_ok=True)
        mkdir(self.credentials_root, exist_ok=True)

        mkdir(self.extracted_data, exist_ok=True)
        mkdir(self.extracted_era5, exist_ok=True)
        mkdir(self.extracted_cmip6, exist_ok=True)
        mkdir(self.ncei_climate_stations, exist_ok=True)
        mkdir(self.open_topography_elevation, exist_ok=True)
        mkdir(self.rub_local_climate_zones, exist_ok=True)

        mkdir(self.downscale_model, exist_ok=True)
        mkdir(self.predictors, exist_ok=True)
        mkdir(self.training_data, exist_ok=True)

        mkdir(self.results, exist_ok=True)
        mkdir(self.results_metadata, exist_ok=True)
        mkdir(self.daily_results, exist_ok=True)
        mkdir(self.raw_daily_results, exist_ok=True)
        mkdir(self.annual_results, exist_ok=True)
        mkdir(self.raw_annual_results, exist_ok=True)

    @property
    def root(self) -> Path:
        return self._root

    @property
    def credentials_root(self) -> Path:
        return self._credentials_root

    ##################
    # Extracted data #
    ##################

    @property
    def extracted_data(self) -> Path:
        return self.root / "extracted_data"

    @property
    def extracted_era5(self) -> Path:
        return self.extracted_data / "era5"

    def extracted_era5_path(
        self, dataset: str, variable: str, year: int | str, month: str
    ) -> Path:
        return self.extracted_era5 / f"{dataset}_{variable}_{year}_{month}.nc"

    @property
    def extracted_cmip6(self) -> Path:
        return self.extracted_data / "cmip6"

    def load_koppen_geiger_model_inclusion(
        self, *, return_full_criteria: bool = False
    ) -> pd.DataFrame:
        meta_path = self.extracted_cmip6 / "koppen_geiger_model_inclusion.parquet"

        if not meta_path.exists():
            df = pd.read_html(
                "https://www.nature.com/articles/s41597-023-02549-6/tables/3"
            )[0]
            df.columns = [
                "source_id",
                "member_count",
                "mean_trend",
                "std_dev_trend",
                "transient_climate_response",
                "equilibrium_climate_sensitivity",
                "included_raw",
            ]
            df["included"] = df["included_raw"].apply({"Yes": True, "No": False}.get)
            save_parquet(df, meta_path)

        df = pd.read_parquet(meta_path)
        if return_full_criteria:
            return df
        return df[["source_id", "included"]]

    def load_cmip6_metadata(self) -> pd.DataFrame:
        meta_path = self.extracted_cmip6 / "cmip6-metadata.parquet"

        if not meta_path.exists():
            external_path = "https://storage.googleapis.com/cmip6/cmip6-zarr-consolidated-stores.csv"
            meta = pd.read_csv(external_path)
            save_parquet(meta, meta_path)

        return pd.read_parquet(meta_path)

    def extracted_cmip6_path(
        self,
        variable: str,
        experiment: str,
        gcm_member: str,
    ) -> Path:
        return self.extracted_cmip6 / f"{variable}_{experiment}_{gcm_member}.nc"

    def get_gcms(
        self,
        source_variables: Collection[str],
    ) -> list[str]:
        inclusion_meta = self.load_scenario_inclusion_metadata()[source_variables]
        inclusion_meta = inclusion_meta[inclusion_meta.all(axis=1)]
        gcms = []
        for model, variant in inclusion_meta.index.tolist():
            gcms.append(gcm_member_id(model, variant))
        return gcms

    @property
    def ncei_climate_stations(self) -> Path:
        return self.extracted_data / "ncei_climate_stations"

    def save_ncei_climate_stations(self, df: pd.DataFrame, year: int | str) -> None:
        if self._read_only:
            msg = "Cannot save NCEI climate stations to read-only data"
            raise ValueError(msg)
        path = self.ncei_climate_stations / f"{year}.parquet"
        save_parquet(df, path)

    def load_ncei_climate_stations(self, year: int | str) -> pd.DataFrame:
        return pd.read_parquet(self.ncei_climate_stations / f"{year}.parquet")

    @property
    def open_topography_elevation(self) -> Path:
        return self.extracted_data / "open_topography_elevation"

    @property
    def rub_local_climate_zones(self) -> Path:
        return self.extracted_data / "rub_local_climate_zones"

    ###################
    # Downscale model #
    ###################

    @property
    def downscale_model(self) -> Path:
        return self.root / "downscale_model"

    @property
    def predictors(self) -> Path:
        return self.downscale_model / "predictors"

    def save_predictor(
        self,
        predictor: rt.RasterArray,
        name: str,
        lat_start: int,
        lon_start: int,
    ) -> None:
        if self._read_only:
            msg = "Cannot save predictors to read-only data"
            raise ValueError(msg)
        path = self.predictors / f"{name}_{lat_start}_{lon_start}.tif"
        save_raster(predictor, path)

    def load_predictor(self, name: str) -> rt.RasterArray:
        paths = list(self.predictors.glob(f"{name}_*.tif"))
        return rt.load_mf_raster(paths)

    @property
    def training_data(self) -> Path:
        return self.downscale_model / "training_data"

    def save_training_data(self, df: pd.DataFrame, year: int | str) -> None:
        if self._read_only:
            msg = "Cannot save training data to read-only data"
            raise ValueError(msg)
        path = self.training_data / f"{year}.parquet"
        save_parquet(df, path)

    def load_training_data(self, year: int | str) -> pd.DataFrame:
        return pd.read_parquet(self.training_data / f"{year}.parquet")

    ###########
    # Results #
    ###########

    @property
    def results(self) -> Path:
        return self.root / "results"

    @property
    def results_metadata(self) -> Path:
        return self.results / "metadata"

    def save_scenario_metadata(self, df: pd.DataFrame) -> None:
        if self._read_only:
            msg = "Cannot save scenario metadata to read-only data"
            raise ValueError(msg)
        path = self.results_metadata / "scenario_metadata.parquet"
        save_parquet(df, path)

    def load_scenario_metadata(self) -> pd.DataFrame:
        path = self.results_metadata / "scenario_metadata.parquet"
        return pd.read_parquet(path)

    def save_scenario_inclusion_metadata(self, df: pd.DataFrame) -> None:
        if self._read_only:
            msg = "Cannot save scenario inclusion metadata to read-only data"
            raise ValueError(msg)
        # Need to save to our scripts directory for doc building
        scripts_root = Path(__file__).parent.parent.parent / "scripts"
        for root_dir in [self.results_metadata, scripts_root]:
            path = root_dir / "scenario_inclusion_metadata.parquet"
            save_parquet(df, path)

    def load_scenario_inclusion_metadata(self) -> pd.DataFrame:
        path = self.results_metadata / "scenario_inclusion_metadata.parquet"
        return pd.read_parquet(path)

    @property
    def daily_results(self) -> Path:
        return self.results / "daily"

    @property
    def raw_daily_results(self) -> Path:
        # NOTE: temporarily points at erf-scratch instead of self.daily_results / "raw".
        # This ignores self.root, so *any* ClimateData -- including one built on a tmp or
        # scratch root -- creates and writes here, because _create_model_root mkdirs it
        # unless read_only=True. That is why tests/conftest.py has to patch
        # cdc.AGGREGATE_ROOT to keep the suite hermetic; without it six tests fail on CI
        # runners while passing on the cluster, where the directory already exists.
        # If the redirect is still wanted it should be an explicit constructor argument
        # rather than an unconditional override.
        return cdc.AGGREGATE_ROOT / "erf-scratch"

    def raw_daily_results_path(
        self,
        scenario: str,
        variable: str,
        year: int | str,
        gcm_member: str,
    ) -> Path:
        return self.raw_daily_results / scenario / variable / f"{year}_{gcm_member}.nc"

    def save_raw_daily_results(
        self,
        results_ds: xr.Dataset,
        scenario: str,
        variable: str,
        year: int | str,
        gcm_member: str,
        encoding_kwargs: dict[str, Any],
    ) -> None:
        if self._read_only:
            msg = "Cannot save raw daily results to read-only data"
            raise ValueError(msg)
        path = self.raw_daily_results_path(scenario, variable, year, gcm_member)
        mkdir(path.parent, exist_ok=True, parents=True)
        save_xarray(results_ds, path, encoding_kwargs)

    def load_raw_daily_results(
        self,
        scenario: str,
        variable: str,
        year: int | str,
        gcm_member: str,
    ) -> xr.Dataset:
        path = self.raw_daily_results_path(scenario, variable, year, gcm_member)
        return xr.open_dataset(path)

    def daily_results_path(
        self,
        scenario: str,
        variable: str,
        year: int | str,
    ) -> Path:
        return self.daily_results / scenario / variable / f"{year}.nc"

    def save_daily_results(
        self,
        results_ds: xr.Dataset,
        scenario: str,
        variable: str,
        year: int | str,
        encoding_kwargs: dict[str, Any],
    ) -> None:
        if self._read_only:
            msg = "Cannot save daily results to read-only data"
            raise ValueError(msg)
        path = self.daily_results_path(scenario, variable, year)
        mkdir(path.parent, exist_ok=True, parents=True)
        save_xarray(results_ds, path, encoding_kwargs)

    def load_daily_results(
        self,
        scenario: str,
        variable: str,
        year: int | str,
    ) -> xr.Dataset:
        results_path = self.daily_results_path(scenario, variable, year)
        return xr.open_dataset(results_path)

    @property
    def annual_results(self) -> Path:
        return self.results / "annual"

    @property
    def raw_annual_results(self) -> Path:
        return self.annual_results / "raw"

    def raw_annual_results_path(
        self,
        scenario: str,
        variable: str,
        year: int | str,
        gcm_member: str,
    ) -> Path:
        return self.raw_annual_results / scenario / variable / f"{year}_{gcm_member}.nc"

    @staticmethod
    def combine_raw_annual(paths: list[Path]) -> xr.Dataset:
        """Open and combine raw annual ``.nc`` files into one dataset, sorted by year."""
        return xr.open_mfdataset(paths, combine="by_coords").sortby("year").compute()

    def load_raw_annual_mfdataset(
        self,
        scenario: str,
        variable: str,
        gcm_member: str | None = None,
    ) -> xr.Dataset:
        """Glob and combine all raw annual results for a scenario/variable.

        Pass ``gcm_member`` to restrict to a single member's files; otherwise every
        ``.nc`` under the scenario/variable directory is combined.
        """
        pattern = f"*{gcm_member}.nc" if gcm_member is not None else "*.nc"
        paths = sorted((self.raw_annual_results / scenario / variable).glob(pattern))
        return self.combine_raw_annual(paths)

    def save_raw_annual_results(
        self,
        results_ds: xr.Dataset,
        scenario: str,
        variable: str,
        year: int | str,
        gcm_member: str,
        encoding_kwargs: dict[str, Any],
    ) -> None:
        if self._read_only:
            msg = "Cannot save raw annual results to read-only data"
            raise ValueError(msg)
        path = self.raw_annual_results_path(scenario, variable, year, gcm_member)
        mkdir(path.parent, exist_ok=True, parents=True)
        save_xarray(results_ds, path, encoding_kwargs)

    @property
    def compiled_annual_results(self) -> Path:
        return self.raw_annual_results / "compiled"

    def compiled_annual_results_path(
        self,
        scenario: str,
        variable: str,
        gcm_member: str,
    ) -> Path:
        return self.compiled_annual_results / scenario / variable / f"{gcm_member}.nc"

    def list_gcm_members(self, scenario: str, variable: str) -> list[str]:
        return [
            p.stem
            for p in self.compiled_annual_results_path(
                scenario, variable, ""
            ).parent.glob("*.nc")
        ]

    def save_compiled_annual_results(
        self,
        results_ds: xr.Dataset,
        scenario: str,
        variable: str,
        gcm_member: str,
        encoding_kwargs: dict[str, Any],
    ) -> None:
        if self._read_only:
            msg = "Cannot save compiled annual results to read-only data"
            raise ValueError(msg)
        path = self.compiled_annual_results_path(scenario, variable, gcm_member)
        mkdir(path.parent, exist_ok=True, parents=True)
        save_xarray(results_ds, path, encoding_kwargs)

    def load_compiled_annual_results(
        self,
        scenario: str,
        variable: str,
        gcm_member: str,
    ) -> xr.Dataset:
        path = self.compiled_annual_results_path(scenario, variable, gcm_member)
        return xr.open_dataset(path)

    def annual_results_path(
        self,
        scenario: str,
        variable: str,
        draw: int | str,
    ) -> Path:
        return self.annual_results / scenario / variable / f"{draw:0>3}.nc"

    def link_annual_draw(
        self,
        draw: int | str,
        scenario: str,
        variable: str,
        gcm_member: str,
    ) -> None:
        if self._read_only:
            msg = "Cannot link annual draw to read-only data"
            raise ValueError(msg)
        source_path = self.compiled_annual_results_path(scenario, variable, gcm_member)
        dest_path = self.annual_results_path(scenario, variable, draw)
        mkdir(dest_path.parent, exist_ok=True, parents=True)
        if dest_path.exists():
            dest_path.unlink()
        dest_path.symlink_to(source_path)

    def draw_results_path(self, scenario: str, measure: str, draw: str) -> Path:
        """Get the path to annual results for a specific scenario, measure, and draw.

        Parameters
        ----------
        scenario
            The climate scenario (e.g. "ssp126")
        measure
            The climate measure (e.g. "mean_temperature")
        draw
            The draw of the climate data to load (e.g. "000")

        Returns
        -------
        Path
            The path to the results file
        """
        return self.annual_results / scenario / measure / f"{draw}.nc"

    def load_draw_results(self, scenario: str, measure: str, draw: str) -> xr.Dataset:
        """Load annual climate results for a specific scenario, measure, and draw.

        Parameters
        ----------
        scenario
            The climate scenario (e.g. "ssp126")
        measure
            The climate measure (e.g. "mean_temperature")
        draw
            The draw of the climate data to load (e.g. "000")

        Returns
        -------
        xr.Dataset
            The climate data in xarray format
        """
        path = self.annual_results_path(scenario, measure, draw)
        ds = xr.open_dataset(path, decode_coords="all")
        # The rioxarray accessor is untyped, so write_crs erases the Dataset type.
        return cast(xr.Dataset, ds.rio.write_crs("EPSG:4326"))

combine_raw_annual(paths: list[Path]) -> xr.Dataset staticmethod

Open and combine raw annual .nc files into one dataset, sorted by year.

Source code in src/climate_data/data.py
@staticmethod
def combine_raw_annual(paths: list[Path]) -> xr.Dataset:
    """Open and combine raw annual ``.nc`` files into one dataset, sorted by year."""
    return xr.open_mfdataset(paths, combine="by_coords").sortby("year").compute()

draw_results_path(scenario: str, measure: str, draw: str) -> Path

Get the path to annual results for a specific scenario, measure, and draw.

Parameters

scenario The climate scenario (e.g. "ssp126") measure The climate measure (e.g. "mean_temperature") draw The draw of the climate data to load (e.g. "000")

Returns

Path The path to the results file

Source code in src/climate_data/data.py
def draw_results_path(self, scenario: str, measure: str, draw: str) -> Path:
    """Get the path to annual results for a specific scenario, measure, and draw.

    Parameters
    ----------
    scenario
        The climate scenario (e.g. "ssp126")
    measure
        The climate measure (e.g. "mean_temperature")
    draw
        The draw of the climate data to load (e.g. "000")

    Returns
    -------
    Path
        The path to the results file
    """
    return self.annual_results / scenario / measure / f"{draw}.nc"

load_draw_results(scenario: str, measure: str, draw: str) -> xr.Dataset

Load annual climate results for a specific scenario, measure, and draw.

Parameters

scenario The climate scenario (e.g. "ssp126") measure The climate measure (e.g. "mean_temperature") draw The draw of the climate data to load (e.g. "000")

Returns

xr.Dataset The climate data in xarray format

Source code in src/climate_data/data.py
def load_draw_results(self, scenario: str, measure: str, draw: str) -> xr.Dataset:
    """Load annual climate results for a specific scenario, measure, and draw.

    Parameters
    ----------
    scenario
        The climate scenario (e.g. "ssp126")
    measure
        The climate measure (e.g. "mean_temperature")
    draw
        The draw of the climate data to load (e.g. "000")

    Returns
    -------
    xr.Dataset
        The climate data in xarray format
    """
    path = self.annual_results_path(scenario, measure, draw)
    ds = xr.open_dataset(path, decode_coords="all")
    # The rioxarray accessor is untyped, so write_crs erases the Dataset type.
    return cast(xr.Dataset, ds.rio.write_crs("EPSG:4326"))

load_raw_annual_mfdataset(scenario: str, variable: str, gcm_member: str | None = None) -> xr.Dataset

Glob and combine all raw annual results for a scenario/variable.

Pass gcm_member to restrict to a single member's files; otherwise every .nc under the scenario/variable directory is combined.

Source code in src/climate_data/data.py
def load_raw_annual_mfdataset(
    self,
    scenario: str,
    variable: str,
    gcm_member: str | None = None,
) -> xr.Dataset:
    """Glob and combine all raw annual results for a scenario/variable.

    Pass ``gcm_member`` to restrict to a single member's files; otherwise every
    ``.nc`` under the scenario/variable directory is combined.
    """
    pattern = f"*{gcm_member}.nc" if gcm_member is not None else "*.nc"
    paths = sorted((self.raw_annual_results / scenario / variable).glob(pattern))
    return self.combine_raw_annual(paths)

PopulationModelData

Handles population data and location hierarchies.

This class manages: 1. Population projections at different time points 2. Location hierarchies (GBD, LSAE, etc.) 3. Spatial data for aggregation

The population data is used as weights when aggregating climate data to different location hierarchies.

Source code in src/climate_data/data.py
class PopulationModelData:
    """Handles population data and location hierarchies.

    This class manages:
    1. Population projections at different time points
    2. Location hierarchies (GBD, LSAE, etc.)
    3. Spatial data for aggregation

    The population data is used as weights when aggregating climate data
    to different location hierarchies.
    """

    def __init__(
        self,
        root: str | Path = cdc.POPULATION_MODEL_ROOT,
    ) -> None:
        """Initialize the population model data manager.

        Parameters
        ----------
        root : str | Path
            Path to the population model root directory
        """
        self._root = Path(root)

    @property
    def root(self) -> Path:
        """Get the root directory for population model data."""
        return self._root

    @property
    def results(self) -> Path:
        """Get the directory containing current model results."""
        return Path(self.root, "results") / "current"

    @property
    def model_spec_path(self) -> Path:
        """Get the path to the model specification file."""
        return self.results / "specification.yaml"

    def load_model_spec(self) -> dict[str, Any]:
        """Load the model specification file.

        Returns
        -------
        dict
            The model specification containing paths and parameters
        """
        return cast(dict[str, Any], yaml.safe_load(self.model_spec_path.read_text()))

    def load_modeling_frame(self) -> gpd.GeoDataFrame:
        """Load the modeling frame containing spatial information.

        The modeling frame is a subdivision of the world into equal-area blocks.
        Each block is assigned a unique key that is used to parallelize
        pipeline steps in both population modeling and in this pipeline's
        aggregation step.

        Returns
        -------
        gpd.GeoDataFrame
            The modeling frame with spatial information and block keys
        """
        model_spec = self.load_model_spec()
        raw_root = Path(model_spec["output_root"])
        model_frame_path = raw_root.parent.parent / "modeling_frame.parquet"
        return gpd.read_parquet(model_frame_path)

    def load_results(self, time_point: str, block_key: str) -> rt.RasterArray:
        """Load population results for a specific time point and block.

        Parameters
        ----------
        time_point
            The time point to load (e.g. "2020q1")
        block_key
            The block key to load (e.g. "B-0021X-0003Y")

        Returns
        -------
        rt.RasterArray
            The population raster data
        """
        model_spec = self.load_model_spec()
        raw_root = Path(model_spec["output_root"])
        path = raw_root / "raked_predictions" / time_point / f"{block_key}.tif"
        return rt.load_raster(path)

    @property
    def raking_data(self) -> Path:
        """Get the directory containing data used to rake the population estimates.

        Raking enforces admin-level consistency between gridded population data
        and GBD/FHS population estimates. We'll use these same hierarchies to
        aggregate the climate data.

        """
        return self.root / "admin-inputs" / "raking"

    def load_raking_shapes(
        self,
        full_aggregation_hierarchy: str,
        bounds: tuple[float, float, float, float] | None = None,
    ) -> gpd.GeoDataFrame:
        """Load shapes for a full aggregation hierarchy within given bounds.

        Parameters
        ----------
        full_aggregation_hierarchy
            The full aggregation hierarchy to load (e.g. "gbd_2021")
        bounds
            The bounds to load (xmin, ymin, xmax, ymax). ``None`` (the default)
            loads all shapes for the hierarchy.

        Returns
        -------
        gpd.GeoDataFrame
            The shapes for the given hierarchy and bounds
        """
        if full_aggregation_hierarchy in cdc.GBD_HIERARCHIES:
            shape_path = (
                self.raking_data / f"shapes_{full_aggregation_hierarchy}.parquet"
            )
            gdf = gpd.read_parquet(shape_path, bbox=bounds)

            # We're using population data here instead of a hierarchy because
            # The populations include extra locations we've supplemented that aren't
            # modeled in GBD (e.g. locations with zero population or places that
            # GBD uses population scalars from WPP to model)
            pop_path = (
                self.raking_data / f"population_{full_aggregation_hierarchy}.parquet"
            )
            pop = pd.read_parquet(pop_path)

            keep_cols = ["location_id", "location_name", "most_detailed", "parent_id"]
            keep_mask = (
                (pop.year_id == pop.year_id.max())  # Year doesn't matter
                & (pop.most_detailed == 1)
            )
            out = gdf.merge(pop.loc[keep_mask, keep_cols], on="location_id", how="left")
        elif full_aggregation_hierarchy in ["lsae_1209", "lsae_1285"]:
            # This is only a2 geoms, so already most detailed
            shape_path = (
                self.raking_data
                / "gbd-inputs"
                / f"shapes_{full_aggregation_hierarchy}_a2.parquet"
            )
            out = gpd.read_parquet(shape_path, bbox=bounds)
        else:
            msg = f"Unknown pixel hierarchy: {full_aggregation_hierarchy}"
            raise ValueError(msg)
        return out

    def load_raking_populations(self, hierarchy: str) -> pd.DataFrame:
        path = self.raking_data / f"population_{hierarchy}.parquet"
        return pd.read_parquet(path)

    def load_lsae_mapping_shapes(self, admin_level: int) -> gpd.GeoDataFrame:
        """Load the LSAE mapping shapes for a given admin level.

        Parameters
        ----------
        admin_level
            The admin level to load (0, 1, or 2)

        Returns
        -------
        gpd.GeoDataFrame
            The LSAE mapping shapes for the given admin level
        """
        assert admin_level in [0, 1, 2]
        path = f"/home/j/WORK/11_geospatial/admin_shapefiles/current/lbd_standard_admin_{admin_level}_simplified.shp"
        gdf = (
            gpd.read_file(path)
            .rename(columns={"loc_id": "location_id"})
            .loc[:, ["location_id", "geometry"]]
        )
        return gdf

    def load_subset_hierarchy(self, subset_hierarchy: str) -> pd.DataFrame:
        """Load a subset location hierarchy.

        The subset hierarchy might be equal to the full aggregation hierarchy,
        but it might also be a subset of the full aggregation hierarchy.
        These hierarchies are used to provide different views of aggregated
        climate data.

        Parameters
        ----------
        subset_hierarchy
            The administrative hierarchy to load (e.g. "gbd_2021")

        Returns
        -------
        pd.DataFrame
            The hierarchy data with parent-child relationships
        """
        allowed_hierarchies = [
            *cdc.GBD_HIERARCHIES,
            "fhs_2021",
            "fhs_2023",
            "lsae_1209",
            "lsae_1285",
        ]
        if subset_hierarchy not in allowed_hierarchies:
            msg = f"Unknown admin hierarchy: {subset_hierarchy}"
            raise ValueError(msg)
        path = self.raking_data / "gbd-inputs" / f"hierarchy_{subset_hierarchy}.parquet"
        hierarchy_df = pd.read_parquet(path)
        if subset_hierarchy in cdc.GBD_HIERARCHIES:
            # NOTE: parent-drop list authored for gbd_2021/2023, verified against the
            # gbd_2025 hierarchy (2026-07-14): 39/41 ids are still present, so it is
            # applied to gbd_2025 as well. `4854` (J&K/Ladakh) was reorganized out of
            # gbd_2025 and simply no-ops there; `4919` was a typo for `4619` ("North
            # West England", present in all rounds, whose UTLA children were never
            # being dropped) and is corrected below.
            to_drop_parents = [
                ## FROM POPULATION MODEL RAKING DATA PREP
                # Drop UK UTLAs from these regions
                4618,
                4619,
                4620,
                4621,
                4622,
                4623,
                4624,
                4625,
                4626,
                # Drop the India urban/rural splits from these states
                4841,
                4842,
                4843,
                4844,
                4846,
                4849,
                4850,
                4851,
                4852,
                4853,
                4854,
                4855,
                4856,
                4857,
                4859,
                4860,
                4861,
                4862,
                4863,
                4864,
                4865,
                4867,
                4868,
                4869,
                4870,
                4871,
                4872,
                4873,
                4874,
                4875,
                44538,
                # Drop the Maori/non-Maori split from New Zealand
                72,
            ]
            hierarchy_df = hierarchy_df.loc[
                ~hierarchy_df["parent_id"].isin(to_drop_parents)
            ]
            hierarchy_df.loc[
                hierarchy_df["location_id"].isin(to_drop_parents), "most_detailed"
            ] = 1

        return hierarchy_df

model_spec_path: Path property

Get the path to the model specification file.

raking_data: Path property

Get the directory containing data used to rake the population estimates.

Raking enforces admin-level consistency between gridded population data and GBD/FHS population estimates. We'll use these same hierarchies to aggregate the climate data.

results: Path property

Get the directory containing current model results.

root: Path property

Get the root directory for population model data.

__init__(root: str | Path = cdc.POPULATION_MODEL_ROOT) -> None

Initialize the population model data manager.

Parameters

root : str | Path Path to the population model root directory

Source code in src/climate_data/data.py
def __init__(
    self,
    root: str | Path = cdc.POPULATION_MODEL_ROOT,
) -> None:
    """Initialize the population model data manager.

    Parameters
    ----------
    root : str | Path
        Path to the population model root directory
    """
    self._root = Path(root)

load_lsae_mapping_shapes(admin_level: int) -> gpd.GeoDataFrame

Load the LSAE mapping shapes for a given admin level.

Parameters

admin_level The admin level to load (0, 1, or 2)

Returns

gpd.GeoDataFrame The LSAE mapping shapes for the given admin level

Source code in src/climate_data/data.py
def load_lsae_mapping_shapes(self, admin_level: int) -> gpd.GeoDataFrame:
    """Load the LSAE mapping shapes for a given admin level.

    Parameters
    ----------
    admin_level
        The admin level to load (0, 1, or 2)

    Returns
    -------
    gpd.GeoDataFrame
        The LSAE mapping shapes for the given admin level
    """
    assert admin_level in [0, 1, 2]
    path = f"/home/j/WORK/11_geospatial/admin_shapefiles/current/lbd_standard_admin_{admin_level}_simplified.shp"
    gdf = (
        gpd.read_file(path)
        .rename(columns={"loc_id": "location_id"})
        .loc[:, ["location_id", "geometry"]]
    )
    return gdf

load_model_spec() -> dict[str, Any]

Load the model specification file.

Returns

dict The model specification containing paths and parameters

Source code in src/climate_data/data.py
def load_model_spec(self) -> dict[str, Any]:
    """Load the model specification file.

    Returns
    -------
    dict
        The model specification containing paths and parameters
    """
    return cast(dict[str, Any], yaml.safe_load(self.model_spec_path.read_text()))

load_modeling_frame() -> gpd.GeoDataFrame

Load the modeling frame containing spatial information.

The modeling frame is a subdivision of the world into equal-area blocks. Each block is assigned a unique key that is used to parallelize pipeline steps in both population modeling and in this pipeline's aggregation step.

Returns

gpd.GeoDataFrame The modeling frame with spatial information and block keys

Source code in src/climate_data/data.py
def load_modeling_frame(self) -> gpd.GeoDataFrame:
    """Load the modeling frame containing spatial information.

    The modeling frame is a subdivision of the world into equal-area blocks.
    Each block is assigned a unique key that is used to parallelize
    pipeline steps in both population modeling and in this pipeline's
    aggregation step.

    Returns
    -------
    gpd.GeoDataFrame
        The modeling frame with spatial information and block keys
    """
    model_spec = self.load_model_spec()
    raw_root = Path(model_spec["output_root"])
    model_frame_path = raw_root.parent.parent / "modeling_frame.parquet"
    return gpd.read_parquet(model_frame_path)

load_raking_shapes(full_aggregation_hierarchy: str, bounds: tuple[float, float, float, float] | None = None) -> gpd.GeoDataFrame

Load shapes for a full aggregation hierarchy within given bounds.

Parameters

full_aggregation_hierarchy The full aggregation hierarchy to load (e.g. "gbd_2021") bounds The bounds to load (xmin, ymin, xmax, ymax). None (the default) loads all shapes for the hierarchy.

Returns

gpd.GeoDataFrame The shapes for the given hierarchy and bounds

Source code in src/climate_data/data.py
def load_raking_shapes(
    self,
    full_aggregation_hierarchy: str,
    bounds: tuple[float, float, float, float] | None = None,
) -> gpd.GeoDataFrame:
    """Load shapes for a full aggregation hierarchy within given bounds.

    Parameters
    ----------
    full_aggregation_hierarchy
        The full aggregation hierarchy to load (e.g. "gbd_2021")
    bounds
        The bounds to load (xmin, ymin, xmax, ymax). ``None`` (the default)
        loads all shapes for the hierarchy.

    Returns
    -------
    gpd.GeoDataFrame
        The shapes for the given hierarchy and bounds
    """
    if full_aggregation_hierarchy in cdc.GBD_HIERARCHIES:
        shape_path = (
            self.raking_data / f"shapes_{full_aggregation_hierarchy}.parquet"
        )
        gdf = gpd.read_parquet(shape_path, bbox=bounds)

        # We're using population data here instead of a hierarchy because
        # The populations include extra locations we've supplemented that aren't
        # modeled in GBD (e.g. locations with zero population or places that
        # GBD uses population scalars from WPP to model)
        pop_path = (
            self.raking_data / f"population_{full_aggregation_hierarchy}.parquet"
        )
        pop = pd.read_parquet(pop_path)

        keep_cols = ["location_id", "location_name", "most_detailed", "parent_id"]
        keep_mask = (
            (pop.year_id == pop.year_id.max())  # Year doesn't matter
            & (pop.most_detailed == 1)
        )
        out = gdf.merge(pop.loc[keep_mask, keep_cols], on="location_id", how="left")
    elif full_aggregation_hierarchy in ["lsae_1209", "lsae_1285"]:
        # This is only a2 geoms, so already most detailed
        shape_path = (
            self.raking_data
            / "gbd-inputs"
            / f"shapes_{full_aggregation_hierarchy}_a2.parquet"
        )
        out = gpd.read_parquet(shape_path, bbox=bounds)
    else:
        msg = f"Unknown pixel hierarchy: {full_aggregation_hierarchy}"
        raise ValueError(msg)
    return out

load_results(time_point: str, block_key: str) -> rt.RasterArray

Load population results for a specific time point and block.

Parameters

time_point The time point to load (e.g. "2020q1") block_key The block key to load (e.g. "B-0021X-0003Y")

Returns

rt.RasterArray The population raster data

Source code in src/climate_data/data.py
def load_results(self, time_point: str, block_key: str) -> rt.RasterArray:
    """Load population results for a specific time point and block.

    Parameters
    ----------
    time_point
        The time point to load (e.g. "2020q1")
    block_key
        The block key to load (e.g. "B-0021X-0003Y")

    Returns
    -------
    rt.RasterArray
        The population raster data
    """
    model_spec = self.load_model_spec()
    raw_root = Path(model_spec["output_root"])
    path = raw_root / "raked_predictions" / time_point / f"{block_key}.tif"
    return rt.load_raster(path)

load_subset_hierarchy(subset_hierarchy: str) -> pd.DataFrame

Load a subset location hierarchy.

The subset hierarchy might be equal to the full aggregation hierarchy, but it might also be a subset of the full aggregation hierarchy. These hierarchies are used to provide different views of aggregated climate data.

Parameters

subset_hierarchy The administrative hierarchy to load (e.g. "gbd_2021")

Returns

pd.DataFrame The hierarchy data with parent-child relationships

Source code in src/climate_data/data.py
def load_subset_hierarchy(self, subset_hierarchy: str) -> pd.DataFrame:
    """Load a subset location hierarchy.

    The subset hierarchy might be equal to the full aggregation hierarchy,
    but it might also be a subset of the full aggregation hierarchy.
    These hierarchies are used to provide different views of aggregated
    climate data.

    Parameters
    ----------
    subset_hierarchy
        The administrative hierarchy to load (e.g. "gbd_2021")

    Returns
    -------
    pd.DataFrame
        The hierarchy data with parent-child relationships
    """
    allowed_hierarchies = [
        *cdc.GBD_HIERARCHIES,
        "fhs_2021",
        "fhs_2023",
        "lsae_1209",
        "lsae_1285",
    ]
    if subset_hierarchy not in allowed_hierarchies:
        msg = f"Unknown admin hierarchy: {subset_hierarchy}"
        raise ValueError(msg)
    path = self.raking_data / "gbd-inputs" / f"hierarchy_{subset_hierarchy}.parquet"
    hierarchy_df = pd.read_parquet(path)
    if subset_hierarchy in cdc.GBD_HIERARCHIES:
        # NOTE: parent-drop list authored for gbd_2021/2023, verified against the
        # gbd_2025 hierarchy (2026-07-14): 39/41 ids are still present, so it is
        # applied to gbd_2025 as well. `4854` (J&K/Ladakh) was reorganized out of
        # gbd_2025 and simply no-ops there; `4919` was a typo for `4619` ("North
        # West England", present in all rounds, whose UTLA children were never
        # being dropped) and is corrected below.
        to_drop_parents = [
            ## FROM POPULATION MODEL RAKING DATA PREP
            # Drop UK UTLAs from these regions
            4618,
            4619,
            4620,
            4621,
            4622,
            4623,
            4624,
            4625,
            4626,
            # Drop the India urban/rural splits from these states
            4841,
            4842,
            4843,
            4844,
            4846,
            4849,
            4850,
            4851,
            4852,
            4853,
            4854,
            4855,
            4856,
            4857,
            4859,
            4860,
            4861,
            4862,
            4863,
            4864,
            4865,
            4867,
            4868,
            4869,
            4870,
            4871,
            4872,
            4873,
            4874,
            4875,
            44538,
            # Drop the Maori/non-Maori split from New Zealand
            72,
        ]
        hierarchy_df = hierarchy_df.loc[
            ~hierarchy_df["parent_id"].isin(to_drop_parents)
        ]
        hierarchy_df.loc[
            hierarchy_df["location_id"].isin(to_drop_parents), "most_detailed"
        ] = 1

    return hierarchy_df

gcm_member_id(source: str, variant: str) -> str

The <source>_<variant> key identifying one CMIP6 ensemble member.

Both halves are load-bearing. A CMIP6 member_id is only unique within a source -- r1i1p1f1 is shared by 24 of the 22 extracted sources in ssp126 and 30 in ssp585 -- so a filename keyed on the variant alone collides across models. Defined once here because extract_cmip6_main writes these paths and ClimateData.get_gcms reads them; when they disagreed, the extract wrote pr_ssp126_r1i1p1f1.nc for every source, each overwriting the last, and the generate stage looked for a name nothing had written.

Source code in src/climate_data/data.py
def gcm_member_id(source: str, variant: str) -> str:
    """The `<source>_<variant>` key identifying one CMIP6 ensemble member.

    Both halves are load-bearing. A CMIP6 `member_id` is only unique *within* a source --
    `r1i1p1f1` is shared by 24 of the 22 extracted sources in ssp126 and 30 in ssp585 --
    so a filename keyed on the variant alone collides across models. Defined once here
    because `extract_cmip6_main` writes these paths and `ClimateData.get_gcms` reads them;
    when they disagreed, the extract wrote `pr_ssp126_r1i1p1f1.nc` for every source, each
    overwriting the last, and the generate stage looked for a name nothing had written.
    """
    return f"{source}_{variant}"

save_parquet(df: pd.DataFrame, output_path: str | Path) -> None

Save a pandas DataFrame to a file with standard parameters.

Parameters

df The DataFrame to save. output_path The path to save the DataFrame to.

Source code in src/climate_data/data.py
def save_parquet(
    df: pd.DataFrame,
    output_path: str | Path,
) -> None:
    """Save a pandas DataFrame to a file with standard parameters.

    Parameters
    ----------
    df
        The DataFrame to save.
    output_path
        The path to save the DataFrame to.
    """
    touch(output_path, clobber=True)
    df.to_parquet(output_path)

save_raster(raster: rt.RasterArray, output_path: str | Path, num_cores: int = 1, **kwargs: Any) -> None

Save a raster to a file with standard parameters.

Parameters

raster The raster to save. output_path The path to save the raster to. num_cores The number of cores to use for compression.

Source code in src/climate_data/data.py
def save_raster(
    raster: rt.RasterArray,
    output_path: str | Path,
    num_cores: int = 1,
    **kwargs: Any,
) -> None:
    """Save a raster to a file with standard parameters.

    Parameters
    ----------
    raster
        The raster to save.
    output_path
        The path to save the raster to.
    num_cores
        The number of cores to use for compression.
    """
    save_params = {
        "tiled": True,
        "blockxsize": 512,
        "blockysize": 512,
        "compress": "ZSTD",
        "predictor": 2,  # horizontal differencing
        "num_threads": num_cores,
        "bigtiff": "yes",
        **kwargs,
    }
    touch(output_path, clobber=True)
    raster.to_file(output_path, **save_params)

save_raster_to_cog(raster: rt.RasterArray, output_path: str | Path, num_cores: int = 1, resampling: str = 'nearest') -> None

Save a raster to a COG file.

A COG file is a cloud-optimized GeoTIFF that is optimized for use in cloud storage systems. This function saves the raster to a COG file with the specified resampling method.

Parameters

raster The raster to save. output_path The path to save the raster to. num_cores The number of cores to use for compression. resampling The resampling method to use when building the overviews.

Source code in src/climate_data/data.py
def save_raster_to_cog(
    raster: rt.RasterArray,
    output_path: str | Path,
    num_cores: int = 1,
    resampling: str = "nearest",
) -> None:
    """Save a raster to a COG file.

    A COG file is a cloud-optimized GeoTIFF that is optimized for use in cloud storage
    systems. This function saves the raster to a COG file with the specified resampling
    method.

    Parameters
    ----------
    raster
        The raster to save.
    output_path
        The path to save the raster to.
    num_cores
        The number of cores to use for compression.
    resampling
        The resampling method to use when building the overviews.
    """
    cog_save_params = {
        "driver": "COG",
        "overview_resampling": resampling,
    }
    save_raster(raster, output_path, num_cores, **cog_save_params)

save_xarray(ds: xr.Dataset, output_path: str | Path, encoding_kwargs: dict[str, Any]) -> None

Save an xarray dataset to a file with standard parameters.

Parameters

ds The dataset to save. output_path The path to save the dataset to. encoding_kwargs The encoding parameters to use when saving the dataset.

Source code in src/climate_data/data.py
def save_xarray(
    ds: xr.Dataset,
    output_path: str | Path,
    encoding_kwargs: dict[str, Any],
) -> None:
    """Save an xarray dataset to a file with standard parameters.

    Parameters
    ----------
    ds
        The dataset to save.
    output_path
        The path to save the dataset to.
    encoding_kwargs
        The encoding parameters to use when saving the dataset.
    """
    touch(output_path, clobber=True)
    encoding = {
        "dtype": "int16",
        "_FillValue": -32767,
        "zlib": True,
        "complevel": 1,
    }
    encoding.update(encoding_kwargs)
    ds.to_netcdf(output_path, encoding={"value": encoding})

diagnostics

utils

get_locations_depth_first(hierarchy: pd.DataFrame) -> list[int]

Return location ids sorted by a depth first search of the hierarchy.

Locations at the same level are sorted alphabetically by name.

Source code in src/climate_data/diagnostics/utils.py
def get_locations_depth_first(hierarchy: pd.DataFrame) -> list[int]:
    """Return location ids sorted by a depth first search of the hierarchy.

    Locations at the same level are sorted alphabetically by name.
    """

    def _get_locations(location: Any) -> list[int]:
        locs = [location.location_id]

        children = hierarchy[
            (hierarchy.parent_id == location.location_id)
            & (hierarchy.location_id != location.location_id)
        ]
        for child in children.sort_values("location_ascii_name").itertuples():
            locs.extend(_get_locations(child))
        return locs

    top_locs = hierarchy[hierarchy.location_id == hierarchy.parent_id]
    locations = []
    for top_loc in top_locs.sort_values("location_ascii_name").itertuples():
        locations.extend(_get_locations(top_loc))

    return locations

extract

Climate Data Extraction

This module contains pipelines for extracting climate data from various sources.

cmip6

CMIP6 Data Extraction

check_encoding_covers(data_min: float, data_max: float, offset: float, scale: float, variable: str, dtype: str = 'int16') -> None

Refuse to write values the declared encoding cannot represent.

Packing to a 16-bit integer rounds (value - offset) / scale, and anything outside the type's range wraps modulo 65536. to_netcdf does this silently, so the corruption is invisible until someone plots the result and finds negative rainfall. pr shipped with scale_factor=1e-9 for two years -- a 2.83 mm/day ceiling -- and produced 295 files in which 26.4% of sampled cells were wrong and 12.4% were negative.

Raising here makes the next such mistake a failed extract rather than a corrupt archive. An unsigned dtype also makes a negative value an error rather than a wrap, which for a flux like precipitation is the honest outcome.

Source code in src/climate_data/extract/cmip6.py
def check_encoding_covers(
    data_min: float,
    data_max: float,
    offset: float,
    scale: float,
    variable: str,
    dtype: str = "int16",
) -> None:
    """Refuse to write values the declared encoding cannot represent.

    Packing to a 16-bit integer rounds `(value - offset) / scale`, and anything outside
    the type's range wraps modulo 65536. `to_netcdf` does this silently, so the corruption
    is invisible until someone plots the result and finds negative rainfall. `pr` shipped
    with `scale_factor=1e-9` for two years -- a 2.83 mm/day ceiling -- and produced 295
    files in which 26.4% of sampled cells were wrong and 12.4% were negative.

    Raising here makes the next such mistake a failed extract rather than a corrupt
    archive. An unsigned dtype also makes a negative value an error rather than a wrap,
    which for a flux like precipitation is the honest outcome.
    """
    low, high = STORED_LIMITS[dtype]
    for label, value in (("minimum", data_min), ("maximum", data_max)):
        quotient = (value - offset) / scale
        # Round, because rounding is what the writer does -- `to_netcdf` packs with
        # round-half-to-even, not truncation. Comparing the unrounded quotient rejects
        # values the writer would have stored correctly: floating-point noise of
        # -6.19e-17 kg m-2 s-1 in a source GCM field is -6.19e-11 stored, which rounds
        # to 0, and cost eight CMIP6 `pr` members a re-extract. `round` raises on
        # nan/inf, so screen those out and let them fall through to the same message
        # rather than an opaque conversion error.
        stored = round(quotient) if math.isfinite(quotient) else quotient
        if not low <= stored <= high:
            msg = (
                f"The {label} value of {variable} ({value:g}) cannot be represented by"
                f" its {dtype} encoding (offset={offset:g}, scale={scale:g}): it would be"
                f" stored as {stored:.0f}, outside [{low}, {high}], and would wrap modulo"
                f" 65536. Widen encoding_scale for {variable} in"
                f" constants.CMIP6_VARIABLES."
            )
            raise ValueError(msg)

extract_cmip6(cmip6_source: list[str], cmip6_experiment: list[str], cmip6_variable: list[str], output_dir: str, queue: str, overwrite: bool, dry_run: bool) -> None

Extract CMIP6 data.

Extracts CMIP6 data for the given source, experiment, and variable. We use the the table at https://www.nature.com/articles/s41597-023-02549-6/tables/3 to determine which CMIP6 source_ids to include. See ClimateData.load_koppen_geiger_model_inclusion to load and examine this table. The extraction criteria does not completely capture model inclusion criteria as it does not account for the year range avaialable in the data. This determiniation is made when we proccess the data in later steps.

Fans out one job per ensemble member rather than one per (source, experiment). The member counts are wildly uneven -- MIROC6 has 50 pr members where most sources have one -- so grouping them made three jobs carry fifty times the work of a typical one. More importantly, extract_cmip6_main re-raises on failure, so a member the encoding guard rejects used to abandon every member behind it in the same job. One job per member contains that to the member that failed, and makes a resumed run skip the members already written instead of redoing whole groups.

The member space is not a cartesian product -- not every source publishes every variant for every experiment -- so it is enumerated from the metadata and passed as flat_node_args.

Source code in src/climate_data/extract/cmip6.py
@click.command()
@clio.with_cmip6_source(allow_all=True)
@clio.with_cmip6_experiment(allow_all=True)
@clio.with_cmip6_variable(allow_all=True)
@clio.with_output_directory(cdc.MODEL_ROOT)
@clio.with_queue()
@clio.with_overwrite()
@clio.with_dry_run()
def extract_cmip6(
    cmip6_source: list[str],
    cmip6_experiment: list[str],
    cmip6_variable: list[str],
    output_dir: str,
    queue: str,
    overwrite: bool,
    dry_run: bool,
) -> None:
    """Extract CMIP6 data.

    Extracts CMIP6 data for the given source, experiment, and variable. We use the
    the table at https://www.nature.com/articles/s41597-023-02549-6/tables/3 to determine
    which CMIP6 source_ids to include. See `ClimateData.load_koppen_geiger_model_inclusion`
    to load and examine this table. The extraction criteria does not completely
    capture model inclusion criteria as it does not account for the year range avaialable
    in the data. This determiniation is made when we proccess the data in later steps.

    Fans out one job per ensemble member rather than one per (source, experiment). The
    member counts are wildly uneven -- MIROC6 has 50 `pr` members where most sources have
    one -- so grouping them made three jobs carry fifty times the work of a typical one.
    More importantly, `extract_cmip6_main` re-raises on failure, so a member the encoding
    guard rejects used to abandon every member behind it in the same job. One job per
    member contains that to the member that failed, and makes a resumed run skip the
    members already written instead of redoing whole groups.

    The member space is not a cartesian product -- not every source publishes every
    variant for every experiment -- so it is enumerated from the metadata and passed as
    `flat_node_args`.
    """
    overwrite_arg = {"overwrite": None} if overwrite else {}

    cdata = ClimateData(output_dir)
    meta = cdata.load_cmip6_metadata()

    to_run = []
    complete = []
    for variable in cmip6_variable:
        for experiment in cmip6_experiment:
            for source in cmip6_source:
                members = select_members(meta, source, experiment, variable)
                for variant in members:
                    member = gcm_member_id(source, variant)
                    path = cdata.extracted_cmip6_path(variable, experiment, member)
                    if path.exists() and not overwrite:
                        complete.append(member)
                    else:
                        to_run.append((source, experiment, variable, member))

    if not to_run:
        print("All tasks already done.")
        return

    print(f"{len(complete)} tasks already done. Launching {len(to_run)} tasks")
    run_parallel_maybe_dry_run(
        runner="cdtask",
        task_name="extract cmip6",
        flat_node_args=(
            ("cmip6-source", "cmip6-experiment", "cmip6-variable", "gcm-member"),
            to_run,
        ),
        task_args={
            "output-dir": output_dir,
            **overwrite_arg,
        },
        task_resources={
            "queue": queue,
            "cores": 1,
            "memory": "10G",
            "runtime": "3000m",
            "project": "proj_rapidresponse",
        },
        max_attempts=1,
        concurrency_limit=50,
        dry_run=dry_run,
    )

load_cmip_data(zarr_path: str) -> xr.Dataset

Loads a CMIP6 dataset from a zarr path.

Source code in src/climate_data/extract/cmip6.py
def load_cmip_data(zarr_path: str) -> xr.Dataset:
    """Loads a CMIP6 dataset from a zarr path."""
    gcs = gcsfs.GCSFileSystem(token="anon")  # noqa: S106
    mapper = gcs.get_mapper(zarr_path)
    ds = xr.open_zarr(mapper, consolidated=True)
    ds = ds.drop_vars(
        ["lat_bnds", "lon_bnds", "time_bnds", "height", "time_bounds", "bnds"],
        errors="ignore",
    )
    return ds  # type: ignore[no-any-return]

select_members(meta: pd.DataFrame, cmip6_source: str, cmip6_experiment: str, cmip6_variable: str, gcm_member: str | None = None) -> dict[str, str]

The {member_id: zstore} this job should extract.

With gcm_member given, narrows to that one ensemble member so the runner can put each member in its own job. Shared with the runner, which enumerates the same space to build its task list -- if these two disagreed, the runner would submit jobs whose member does not exist and they would silently extract nothing.

Source code in src/climate_data/extract/cmip6.py
def select_members(
    meta: pd.DataFrame,
    cmip6_source: str,
    cmip6_experiment: str,
    cmip6_variable: str,
    gcm_member: str | None = None,
) -> dict[str, str]:
    """The `{member_id: zstore}` this job should extract.

    With `gcm_member` given, narrows to that one ensemble member so the runner can put
    each member in its own job. Shared with the runner, which enumerates the same space
    to build its task list -- if these two disagreed, the runner would submit jobs whose
    member does not exist and they would silently extract nothing.
    """
    table_id = cdc.CMIP6_VARIABLES.get(cmip6_variable).table_id
    mask = (
        (meta.source_id == cmip6_source)
        & (meta.experiment_id == cmip6_experiment)
        & (meta.variable_id == cmip6_variable)
        & (meta.table_id == table_id)
    )
    # `to_dict` on a pandas index is `dict[Hashable, Any]`; coerce once here so callers
    # get plain strings and do not each have to re-cast the variant.
    members: dict[str, str] = {}
    for variant, zstore in meta[mask].set_index("member_id").zstore.to_dict().items():
        members[str(variant)] = str(zstore)
    if gcm_member is None:
        return members

    selected = {}
    for variant, zstore in members.items():
        if gcm_member_id(cmip6_source, variant) == gcm_member:
            selected[variant] = zstore
    if not selected:
        available = []
        for variant in members:
            available.append(gcm_member_id(cmip6_source, variant))
        msg = (
            f"No CMIP6 member {gcm_member!r} for {cmip6_source} {cmip6_experiment}"
            f" {cmip6_variable}. Available: {sorted(available)}"
        )
        raise ValueError(msg)
    return selected

elevation

extract_elevation(model_name: str, output_dir: str, queue: str, dry_run: bool) -> None

Download elevation data from Open Topography.

Source code in src/climate_data/extract/elevation.py
@click.command()
@click.option(
    "--generate-name",
    required=True,
    type=click.Choice(ELEVATION_MODELS),
    help="Name of the elevation model to download.",
)
@clio.with_output_directory(cdc.MODEL_ROOT)
@clio.with_queue()
@clio.with_dry_run()
def extract_elevation(
    model_name: str,
    output_dir: str,
    queue: str,
    dry_run: bool,
) -> None:
    """Download elevation data from Open Topography."""
    invalid = True
    if invalid:
        msg = "Downloaded using aws cli, this implementation is not valid"
        raise NotImplementedError(msg)

    lat_starts = list(range(-90, 90, FETCH_SIZE))
    lon_starts = list(range(-180, 180, FETCH_SIZE))

    run_parallel_maybe_dry_run(
        runner="cdtask",
        task_name="extract elevation",
        node_args={
            "model-name": [model_name],
            "lat-start": lat_starts,
            "lon-start": lon_starts,
        },
        task_args={
            "output-dir": output_dir,
        },
        task_resources={
            "queue": queue,
            "cores": 1,
            "memory": "10G",
            "runtime": "240m",
            "project": "proj_rapidresponse",
        },
        dry_run=dry_run,
    )

extract_elevation_task(model_name: str, lat_start: int, lon_start: int, output_dir: str) -> None

Download elevation data from Open Topography.

Source code in src/climate_data/extract/elevation.py
@click.command()
@click.option(
    "--model-name",
    required=True,
    type=click.Choice(ELEVATION_MODELS),
    help="Name of the elevation model to download.",
)
@click.option(
    "--lat-start",
    required=True,
    type=int,
    help="Latitude of the top-left corner of the tile.",
)
@click.option(
    "--lon-start",
    required=True,
    type=int,
    help="Longitude of the top-left corner of the tile.",
)
@clio.with_output_directory(cdc.MODEL_ROOT)
def extract_elevation_task(
    model_name: str,
    lat_start: int,
    lon_start: int,
    output_dir: str,
) -> None:
    """Download elevation data from Open Topography."""
    invalid = True
    if invalid:
        msg = "Downloaded using aws cli, this implementation is not valid"
        raise NotImplementedError(msg)

    extract_elevation_main(model_name, lat_start, lon_start, output_dir)

era5

ERA5 Data Extraction

check_extract_year_floor(years: Sequence[str], *, allow_pre_floor: bool) -> None

Refuse a run that reaches below cdc.EXTRACT_YEAR_FLOOR unless asked to.

build_task_lists treats a missing output file as work to do, so after ERF's Sep2026 deletion every pre-1980 extract looks like a gap. The runner's --year defaults to ALL, which means the bare invocation would refill 3,238 files and silently undo the reclamation -- days of Copernicus queue to recover from. Guarding the outcome rather than the default also catches an explicitly typed --year ALL.

Source code in src/climate_data/extract/era5.py
def check_extract_year_floor(years: Sequence[str], *, allow_pre_floor: bool) -> None:
    """Refuse a run that reaches below `cdc.EXTRACT_YEAR_FLOOR` unless asked to.

    `build_task_lists` treats a missing output file as work to do, so after ERF's Sep2026
    deletion every pre-1980 extract looks like a gap. The runner's `--year` defaults to
    ALL, which means the bare invocation would refill 3,238 files and silently undo the
    reclamation -- days of Copernicus queue to recover from. Guarding the outcome rather
    than the default also catches an explicitly typed `--year ALL`.
    """
    below = [year for year in years if int(year) < int(cdc.EXTRACT_YEAR_FLOOR)]
    if not below or allow_pre_floor:
        return
    msg = (
        f"{len(below)} requested years fall below {cdc.EXTRACT_YEAR_FLOOR}"
        f" ({below[0]}-{below[-1]}). ERF deleted those extracts in Sep2026 and this step"
        f" would re-download them. The historical daily layer they feed is already built."
        f" Pass --allow-pre-1980 if a backfill is what you want."
    )
    raise click.UsageError(msg)

variables_for_full_expansion(era5_variables: Sequence[str]) -> list[str]

Drop never-read variables when the whole variable set was requested.

--era5-variable ALL resolves to every declared variable, so the default invocation downloads and stores surface_pressure for every month of every year although no stage opens it. Naming a variable explicitly still extracts it; this only narrows the meaning of "all" to "all the ones we use".

Source code in src/climate_data/extract/era5.py
def variables_for_full_expansion(era5_variables: Sequence[str]) -> list[str]:
    """Drop never-read variables when the whole variable set was requested.

    `--era5-variable ALL` resolves to every declared variable, so the default invocation
    downloads and stores `surface_pressure` for every month of every year although no
    stage opens it. Naming a variable explicitly still extracts it; this only narrows the
    meaning of "all" to "all the ones we use".
    """
    if set(era5_variables) != set(cdc.ERA5_VARIABLES):
        return list(era5_variables)

    keep = []
    skipped = []
    for variable in era5_variables:
        if variable in cdc.EXTRACT_UNUSED_VARIABLES:
            skipped.append(variable)
        else:
            keep.append(variable)

    if skipped:
        print(
            f"Skipping {', '.join(skipped)}: no pipeline stage reads"
            f" {'them' if len(skipped) > 1 else 'it'}."
            f" Name the variable explicitly to extract it anyway."
        )
    return keep

generate

historical_daily

drop_noncore_coords(ds: xr.Dataset) -> xr.Dataset

Drop coordinates the older and newer CDS extract formats disagree about.

Extracts pulled through the newer CDS API carry number (ensemble member) and expver ('0001' final ERA5, '0005' preliminary ERA5T) alongside the grid coords; older extracts carry neither, and are packed int16 rather than float32. xr.concat refuses to join datasets whose coordinates differ, so a look-ahead that crosses from an old extract into a new one has to be normalised first.

Which era a file belongs to tracks its download date, not its data year: surveying the archive, valid_time covers 1950-1989 and 2024 while time covers 1990-2023. So there are two seams, not the single 2023/2024 one this used to describe -- 1989->1990 crosses new format into old, 2023->2024 crosses old into new -- and the 1950-2023 regeneration passed through both.

Source code in src/climate_data/generate/historical_daily.py
def drop_noncore_coords(ds: xr.Dataset) -> xr.Dataset:
    """Drop coordinates the older and newer CDS extract formats disagree about.

    Extracts pulled through the newer CDS API carry `number` (ensemble member) and
    `expver` ('0001' final ERA5, '0005' preliminary ERA5T) alongside the grid coords;
    older extracts carry neither, and are packed int16 rather than float32. `xr.concat`
    refuses to join datasets whose coordinates differ, so a look-ahead that crosses from
    an old extract into a new one has to be normalised first.

    Which era a file belongs to tracks its *download* date, not its data year: surveying
    the archive, `valid_time` covers 1950-1989 and 2024 while `time` covers 1990-2023. So
    there are two seams, not the single 2023/2024 one this used to describe -- 1989->1990
    crosses new format into old, 2023->2024 crosses old into new -- and the 1950-2023
    regeneration passed through both.
    """
    extra = []
    for coord in ds.coords:
        if coord not in CORE_COORDS:
            extra.append(coord)
    return ds.drop_vars(extra)

load_variable_with_lookahead(cdata: ClimateData, variable: str, year: str, month: str, dataset: str) -> xr.Dataset

Load a month's hourly data plus the one sample that closes its final day.

Accumulation windows are stamped by their end, so the last day of a month is closed by the next month's first sample -- for December, by the next year's January. That single extra sample completes the final bin without opening a new one.

Raises if the look-ahead file is absent rather than degrading to the incomplete 23-hour window, which should stop the run rather than quietly shorten one day. Checked before loading the target month so the run fails fast. Note this reaches one year past the history range -- regenerating the last history year needs the following January -- which is why the extract tasks span cdc.EXTRACT_YEARS rather than cdc.HISTORY_YEARS.

Source code in src/climate_data/generate/historical_daily.py
def load_variable_with_lookahead(
    cdata: ClimateData,
    variable: str,
    year: str,
    month: str,
    dataset: str,
) -> xr.Dataset:
    """Load a month's hourly data plus the one sample that closes its final day.

    Accumulation windows are stamped by their end, so the last day of a month is closed
    by the next month's first sample -- for December, by the next *year's* January. That
    single extra sample completes the final bin without opening a new one.

    Raises if the look-ahead file is absent rather than degrading to the incomplete
    23-hour window, which should stop the run rather than quietly shorten one day.
    Checked before loading the target month so the run fails fast. Note this reaches one
    year past the history range -- regenerating the last history year needs the following
    January -- which is why the extract tasks span `cdc.EXTRACT_YEARS` rather than
    `cdc.HISTORY_YEARS`.
    """
    next_year, next_month = _next_month(year, int(month))
    lookahead_path = cdata.extracted_era5_path(dataset, variable, next_year, next_month)
    if not lookahead_path.exists():
        msg = (
            f"Cannot close {year}-{month} for {variable}: the accumulation window of its"
            f" final day ends in the following month, whose extract is missing at"
            f" {lookahead_path}. Extract it before generating {year}:\n"
            f"  cdtask extract era5_download -d {dataset} -x {variable}"
            f" -y {next_year} -m {next_month}\n"
            f"  cdtask extract era5_compress -d {dataset} -x {variable}"
            f" -y {next_year} -m {next_month}\n"
            f"Use the single-job tasks, not `cdrun extract era5`: the runner's --year"
            f" stops at the last history year, and it decides what to fetch by file"
            f" existence."
        )
        raise FileNotFoundError(msg)

    ds = drop_noncore_coords(load_variable(cdata, variable, year, month, dataset))
    lookahead = drop_noncore_coords(
        load_variable(cdata, variable, next_year, next_month, dataset).isel(time=[0])
    )

    # Taking sample zero assumes the next month opens on midnight, which is what closes
    # the target month's final day. ERA5-Land 1950_01 opens at 01:00 instead, so the
    # assumption does fail somewhere in the archive -- silently, at one day per month.
    stamp = pd.Timestamp(lookahead.time.to_index()[0])
    if (stamp.hour, stamp.minute) != (0, 0):
        msg = (
            f"Cannot close {year}-{month} for {variable}: the look-ahead at"
            f" {lookahead_path} opens at {stamp}, not midnight, so its first sample does"
            f" not close the final day of {year}-{month}."
        )
        raise ValueError(msg)

    # `join="exact"` rather than the default `"outer"`: a look-ahead on a different grid
    # should fail here, not silently NaN-pad both months onto a union grid. Note this
    # guards the single-level datasets only -- `load_variable` overwrites ERA5-Land's
    # coordinates from `cdc.ERA5_LAND_*`, so a genuine land-grid change would be relabelled
    # before reaching this point. The real land files differ across the CDS format change
    # by ~1e-5 degrees, which is why that overwrite exists.
    return xr.concat([ds, lookahead], dim="time", join="exact")

trim_to_month(ds: xr.Dataset, year: str, month: int) -> xr.Dataset

Drop collapsed bins that fall outside the target month.

The interval-aware collapse labels its first bin with the previous month's last day, since that bin holds that day's closing sample. Left in place, every month contributes a duplicate date at its seam.

Source code in src/climate_data/generate/historical_daily.py
def trim_to_month(ds: xr.Dataset, year: str, month: int) -> xr.Dataset:
    """Drop collapsed bins that fall outside the target month.

    The interval-aware collapse labels its first bin with the *previous* month's last
    day, since that bin holds that day's closing sample. Left in place, every month
    contributes a duplicate date at its seam.
    """
    dates = pd.to_datetime(ds.date.to_index())
    keep = np.flatnonzero((dates.year == int(year)) & (dates.month == month))
    return ds.isel(date=keep)

scenario_annual

bounded_source_variables(target_variable: str, anomaly_scheme: str) -> list[str]

Source variables of target_variable that anomaly_scheme cannot bias-correct.

Value-bounded variables are floored and clipped to cdc.VALUE_BOUNDS instead of being stabilised, so they run under cdc.BOUNDED_VARIABLE_SCHEMES only; see check_bounded_variable_scheme for why the other schemes are refused. The bounds are keyed on the daily variable, so the annual target has to be resolved through its sources the same way the de-bias check is.

Source code in src/climate_data/generate/scenario_annual.py
def bounded_source_variables(target_variable: str, anomaly_scheme: str) -> list[str]:
    """Source variables of `target_variable` that `anomaly_scheme` cannot bias-correct.

    Value-bounded variables are floored and clipped to `cdc.VALUE_BOUNDS` instead of being
    stabilised, so they run under `cdc.BOUNDED_VARIABLE_SCHEMES` only; see
    `check_bounded_variable_scheme` for why the other schemes are refused. The bounds are
    keyed on the daily variable, so the annual target has to be resolved through its
    sources the same way the de-bias check is.
    """
    if anomaly_scheme in cdc.BOUNDED_VARIABLE_SCHEMES:
        return []
    bounded = []
    for source_variable in TRANSFORM_MAP[target_variable].source_variables:
        if source_variable in cdc.VALUE_BOUNDS:
            bounded.append(source_variable)
    return bounded

forecast_jobs_for_anomaly_scheme(to_run: list[tuple[str, str, str, str]], anomaly_scheme: str) -> list[tuple[str, str, str, str]]

Drop forecast jobs whose variable the anomaly scheme cannot be applied to.

Filtered per job rather than per variable because the historical scenario never reaches generate_scenario_daily_main -- it reads the daily results off disk -- so the scheme does not constrain it, and an additive variable is perfectly runnable there. Only the forecast scenarios pass through compute_anomaly.

Two classes are dropped, and they are not mirror images. Additive variables run under the stabilised monthly scheme and no other. Value-bounded variables are the reverse: they run under cdc.BOUNDED_VARIABLE_SCHEMES only, which means they have to be dropped under the default scheme too -- this stage calls generate_scenario_daily_main in memory, so a bounded forecast job reaches check_bounded_variable_scheme and raises in the worker after being scheduled. That is the failure the daily launcher's filter already prevents.

A filter that removed every job is a usage error rather than an empty run, because the caller asked for work that cannot be done and an empty fan-out looks like success.

Source code in src/climate_data/generate/scenario_annual.py
def forecast_jobs_for_anomaly_scheme(
    to_run: list[tuple[str, str, str, str]],
    anomaly_scheme: str,
) -> list[tuple[str, str, str, str]]:
    """Drop forecast jobs whose variable the anomaly scheme cannot be applied to.

    Filtered per job rather than per variable because the `historical` scenario never
    reaches `generate_scenario_daily_main` -- it reads the daily results off disk -- so
    the scheme does not constrain it, and an additive variable is perfectly runnable
    there. Only the forecast scenarios pass through `compute_anomaly`.

    Two classes are dropped, and they are not mirror images. Additive variables run under
    the stabilised `monthly` scheme and no other. Value-bounded variables are the reverse:
    they run under `cdc.BOUNDED_VARIABLE_SCHEMES` only, which means they have to be dropped
    under the default scheme too -- this stage calls `generate_scenario_daily_main` in
    memory, so a bounded forecast job reaches `check_bounded_variable_scheme` and raises in
    the worker after being scheduled. That is the failure the daily launcher's filter
    already prevents.

    A filter that removed every job is a usage error rather than an empty run, because the
    caller asked for work that cannot be done and an empty fan-out looks like success.
    """
    keep = []
    skipped_additive = set()
    skipped_bounded = set()
    dropped_additive = 0
    dropped_bounded = 0
    for job in to_run:
        variable, scenario = job[0], job[1]
        forecast = scenario != "historical"
        if forecast and bounded_source_variables(variable, anomaly_scheme):
            skipped_bounded.add(variable)
            dropped_bounded += 1
        elif (
            forecast
            and anomaly_scheme != cdc.ANOMALY_SCHEME_MONTHLY
            and ANOMALY_TYPES[variable] != "multiplicative"
        ):
            skipped_additive.add(variable)
            dropped_additive += 1
        else:
            keep.append(job)

    if skipped_additive:
        print(
            f"Anomaly scheme '{anomaly_scheme}' applies to multiplicative variables only;"
            f" skipping {dropped_additive} forecast tasks for:"
            f" {', '.join(sorted(skipped_additive))}."
        )
    if skipped_bounded:
        print(
            f"Anomaly scheme '{anomaly_scheme}' does not apply to value-bounded variables,"
            f" which run under {', '.join(cdc.BOUNDED_VARIABLE_SCHEMES)} only;"
            f" skipping {dropped_bounded} forecast tasks for:"
            f" {', '.join(sorted(skipped_bounded))}."
        )
    if to_run and not keep:
        msg = (
            f"Anomaly scheme '{anomaly_scheme}' can run none of the selected variables, so"
            f" there is nothing to submit. Skipped:"
            f" {', '.join(sorted(skipped_additive | skipped_bounded))}."
        )
        if skipped_bounded:
            msg += (
                f" Value-bounded variables need"
                f" --anomaly-scheme {cdc.BOUNDED_VARIABLE_SCHEMES[0]}."
            )
        raise click.UsageError(msg)
    return keep

scenario_daily

apply_dry_day_rule(anomaly: xr.Dataset, target: xr.Dataset, dry_day_rule: str) -> xr.Dataset

Stop the eps offset from manufacturing rain on days the model reports as dry.

The multiplicative anomaly (T + eps)/(R + eps) is strictly positive even when the model reports no rain at all, so a rainless model day still receives E_ref(month) * a(d) > 0 of the ERA5 climatology. Worse, eps dominates the numerator for such a day, so every rainless day in a cell-month gets the identical anomaly 1/(R_m + eps) -- a flat positive floor rather than a dry spell. Wherever that floor clears the 0.1 mm/day wet-day threshold the pipeline reports a wet day that neither ERA5 nor the GCM has, which is what makes precipitation_days step at the boundary.

preserve zeroes the anomaly on those days and rescales the surviving days of the same cell-month so the month's summed anomaly is unchanged. Because that sum is preserved per GCM cell and interpolate_to_target_latlon is linear, the monthly -- and therefore the annual -- total on the target grid is untouched to floating point. Only the distribution across days moves. That is deliberate: this is a shape fix, exactly orthogonal to the Jensen de-bias, which is a level fix that leaves the shape alone. The two commute, because the de-bias scales a whole month uniformly and this rescale is scale-invariant.

A cell-month the model reports dry on every day has nothing to renormalise onto. Those are left exactly as they are rather than zeroed. Zeroing them is the variant that was measured and rejected: it loses up to 1.2% of the population-weighted annual total across the 1-3% of cell-months that are all-dry. Keeping them is what makes this rule total-preserving, and it is the whole difference between the two.

Source code in src/climate_data/generate/scenario_daily.py
def apply_dry_day_rule(
    anomaly: xr.Dataset,
    target: xr.Dataset,
    dry_day_rule: str,
) -> xr.Dataset:
    """Stop the ``eps`` offset from manufacturing rain on days the model reports as dry.

    The multiplicative anomaly ``(T + eps)/(R + eps)`` is strictly positive even when the
    model reports no rain at all, so a rainless model day still receives
    ``E_ref(month) * a(d) > 0`` of the ERA5 climatology. Worse, ``eps`` dominates the
    numerator for such a day, so *every* rainless day in a cell-month gets the identical
    anomaly ``1/(R_m + eps)`` -- a flat positive floor rather than a dry spell. Wherever that
    floor clears the 0.1 mm/day wet-day threshold the pipeline reports a wet day that neither
    ERA5 nor the GCM has, which is what makes ``precipitation_days`` step at the boundary.

    ``preserve`` zeroes the anomaly on those days and rescales the surviving days of the same
    cell-month so the month's *summed* anomaly is unchanged. Because that sum is preserved per
    GCM cell and ``interpolate_to_target_latlon`` is linear, the monthly -- and therefore the
    annual -- total on the target grid is untouched to floating point. Only the distribution
    across days moves. That is deliberate: this is a shape fix, exactly orthogonal to the
    Jensen de-bias, which is a level fix that leaves the shape alone. The two commute, because
    the de-bias scales a whole month uniformly and this rescale is scale-invariant.

    A cell-month the model reports dry on *every* day has nothing to renormalise onto. Those
    are left exactly as they are rather than zeroed. Zeroing them is the variant that was
    measured and rejected: it loses up to 1.2% of the population-weighted annual total across
    the 1-3% of cell-months that are all-dry. Keeping them is what makes this rule
    total-preserving, and it is the whole difference between the two.
    """
    if dry_day_rule == "none":
        return anomaly
    if dry_day_rule != "preserve":
        msg = f"Unknown dry-day rule: {dry_day_rule}"
        raise ValueError(msg)

    wet = target > cdc.DRY_DAY_THRESHOLD_MM
    kept = anomaly.where(wet, 0.0)

    month_total = anomaly.groupby("date.month").sum("date")
    kept_total = kept.groupby("date.month").sum("date")
    has_wet_day = kept_total > 0
    # 1.0 on an all-dry cell-month, so the restore below survives the rescale untouched.
    rescale = (month_total / kept_total.where(has_wet_day)).fillna(1.0)

    has_wet_day_daily = has_wet_day.sel(month=anomaly["date"].dt.month).drop_vars(
        "month"
    )
    kept = kept.where(has_wet_day_daily, anomaly)

    rescaled = kept.groupby("date.month") * rescale
    return rescaled.drop_vars("month")

apply_value_bounds(ds: xr.Dataset, bounds: tuple[float, float]) -> xr.Dataset

Clip every data variable to bounds (inclusive); NaN stays NaN.

Source code in src/climate_data/generate/scenario_daily.py
def apply_value_bounds(ds: xr.Dataset, bounds: tuple[float, float]) -> xr.Dataset:
    """Clip every data variable to ``bounds`` (inclusive); NaN stays NaN."""
    lo, hi = bounds
    if not lo < hi:
        msg = f"value bounds must satisfy lo < hi, got {bounds!r}."
        raise ValueError(msg)
    return ds.clip(lo, hi)

check_bounded_variable_scheme(target_variable: str, anomaly_scheme: str) -> None

Refuse a stabilised or yearly anomaly scheme for a value-bounded variable.

A bounded variable (relative humidity) has its raw inputs clipped to cdc.VALUE_BOUNDS before the ratio is formed, so the floor already does the job of the +1 / tapered eps in the eps-bearing schemes -- applying both would stabilise twice and shift the ratio away from the construction that was evaluated. The yearly family drops the monthly anchor that evaluation was made on. Called from the worker and, via variables_for_anomaly_scheme, from the launcher, so --target-variable all under the default scheme skips the variable with a message instead of queueing doomed jobs.

Source code in src/climate_data/generate/scenario_daily.py
def check_bounded_variable_scheme(target_variable: str, anomaly_scheme: str) -> None:
    """Refuse a stabilised or yearly anomaly scheme for a value-bounded variable.

    A bounded variable (relative humidity) has its raw inputs clipped to
    ``cdc.VALUE_BOUNDS`` before the ratio is formed, so the floor already does the job of
    the +1 / tapered eps in the eps-bearing schemes -- applying both would stabilise twice
    and shift the ratio away from the construction that was evaluated. The yearly family
    drops the monthly anchor that evaluation was made on. Called from the worker and, via
    ``variables_for_anomaly_scheme``, from the launcher, so ``--target-variable all`` under
    the default scheme skips the variable with a message instead of queueing doomed jobs.
    """
    bounds = cdc.VALUE_BOUNDS.get(target_variable)
    if bounds is None or anomaly_scheme in cdc.BOUNDED_VARIABLE_SCHEMES:
        return
    msg = (
        f"{target_variable!r} is value-bounded (inputs clipped to {bounds['input']}, output "
        f"to {bounds['output']}) and runs under {list(cdc.BOUNDED_VARIABLE_SCHEMES)} only; "
        f"got anomaly_scheme={anomaly_scheme!r}. Pass --anomaly-scheme "
        f"{cdc.BOUNDED_VARIABLE_SCHEMES[0]} for this variable."
    )
    raise ValueError(msg)

check_debias_variable(target_variable: str, debias_method: str, dry_day_rule: str = 'none') -> None

Refuse a correction for a variable it has not been validated against.

Called from the launchers as well as from the worker, so that --target-variable all fails in a second rather than after submitting thousands of doomed jobs.

Source code in src/climate_data/generate/scenario_daily.py
def check_debias_variable(
    target_variable: str, debias_method: str, dry_day_rule: str = "none"
) -> None:
    """Refuse a correction for a variable it has not been validated against.

    Called from the launchers as well as from the worker, so that ``--target-variable all``
    fails in a second rather than after submitting thousands of doomed jobs.
    """
    if debias_method != "none" and target_variable not in cdc.DEBIAS_VARIABLES:
        msg = (
            f"debias_method={debias_method!r} is not validated for {target_variable!r}. "
            f"Allowed: {list(cdc.DEBIAS_VARIABLES)}. Name the variable explicitly rather "
            "than using 'all'."
        )
        raise ValueError(msg)
    if dry_day_rule != "none" and target_variable not in cdc.DRY_DAY_VARIABLES:
        msg = (
            f"dry_day_rule={dry_day_rule!r} is not validated for {target_variable!r}. "
            f"Allowed: {list(cdc.DRY_DAY_VARIABLES)}. Name the variable explicitly rather "
            "than using 'all'."
        )
        raise ValueError(msg)

check_scheme_compatibility(anomaly_scheme: str, anomaly_type: str, debias_method: str, dry_day_rule: str) -> None

Reject combinations of the two correction axes that cannot both apply.

debias_method and dry_day_rule correct the (T + eps) / (R + eps) construction. Both monthly (constant eps) and monthly-taper (tapered eps) use it, so both accept them. The yearly and monthly-ratio families have no eps for those corrections to act on, so asking for either against them is a mistake rather than a no-op -- and silently ignoring the request would produce a file whose attrs claim a correction that was never applied.

Source code in src/climate_data/generate/scenario_daily.py
def check_scheme_compatibility(
    anomaly_scheme: str,
    anomaly_type: str,
    debias_method: str,
    dry_day_rule: str,
) -> None:
    """Reject combinations of the two correction axes that cannot both apply.

    `debias_method` and `dry_day_rule` correct the `(T + eps) / (R + eps)` construction.
    Both `monthly` (constant eps) and `monthly-taper` (tapered eps) use it, so both accept
    them. The yearly and monthly-ratio families have no eps for those corrections to act
    on, so asking for either against them is a mistake rather than a no-op -- and silently
    ignoring the request would produce a file whose attrs claim a correction that was
    never applied.
    """
    has_eps = anomaly_scheme in cdc.EPS_BEARING_SCHEMES
    if not has_eps and (debias_method != "none" or dry_day_rule != "none"):
        msg = (
            f"debias_method={debias_method!r} and dry_day_rule={dry_day_rule!r} cannot "
            f"be combined with anomaly_scheme={anomaly_scheme!r}: they correct the eps "
            f"stabiliser, which this scheme does not use. Schemes that carry an eps: "
            f"{list(cdc.EPS_BEARING_SCHEMES)}."
        )
        raise ValueError(msg)
    if anomaly_type != "multiplicative":
        msg = f"Anomaly scheme '{anomaly_scheme}' only applies to multiplicative variables."
        raise ValueError(msg)

compute_anomaly(reference: xr.Dataset, target: xr.Dataset, anomaly_type: str, *, debias_method: str, dry_day_rule: str, anomaly_scheme: str = cdc.ANOMALY_SCHEME_MONTHLY, eps_floor: float = cdc.DEFAULT_EPS_FLOOR, anomaly_cap: float | None = cdc.DEFAULT_ANOMALY_CAP) -> xr.Dataset

The forecast anomaly, optionally capped.

A thin wrapper so the ceiling applies to every scheme without threading it through each one: the scheme dispatch below has several return points, and duplicating the clip at each is how one of them ends up missing it.

Applied on the GCM grid, BEFORE interpolate_to_target_latlon. Capping afterwards would leave a blown-up cell already smeared across its neighbours by the regrid.

The cap bounds each cell's multiplier at the granularity the scheme anchors at -- annual for the yearly family, per calendar month for the monthly family -- and rescales the daily series to meet it.

It cannot be an elementwise clip. anomaly carries the target's daily date dimension, because the denominator is a reference-window mean with no date dim and the division broadcasts. Clipping elementwise bounds every day at anomaly_cap times the cell's reference mean, and against an annual denominator that ceiling sits inside the ordinary distribution of daily rainfall: a cell averaging 2 mm/day is cut at 40 mm/day, an unremarkable tropical wet day. The daily-over-annual ratio also carries the seasonal cycle, so such a clip bites hardest where seasonality is strongest rather than where the pathology is.

Measured on yearly-delta ssp126 before this change, an elementwise clip at 20 altered 20.5% of land pixels in 2083 and 15.5% in 2050; 82% of those had an annual anomaly at or below 2 -- cells projecting essentially no change -- and the global total moved -19.1% and -1.1% respectively.

Rescaling instead leaves any cell at or below the ceiling bit-identical to uncapped, brings a cell above it exactly to the ceiling, and preserves within-period shape so a wet day stays proportionally a wet day.

The cap still only removes precipitation -- it cannot conserve, and whatever it removes is gone from the total. That is intended: the values it removes are ones no cell's observed climatology supports. But it means a cap moves the level, so it must be chosen on evidence rather than set defensively.

Source code in src/climate_data/generate/scenario_daily.py
def compute_anomaly(
    reference: xr.Dataset,
    target: xr.Dataset,
    anomaly_type: str,
    *,
    debias_method: str,
    dry_day_rule: str,
    anomaly_scheme: str = cdc.ANOMALY_SCHEME_MONTHLY,
    eps_floor: float = cdc.DEFAULT_EPS_FLOOR,
    anomaly_cap: float | None = cdc.DEFAULT_ANOMALY_CAP,
) -> xr.Dataset:
    """The forecast anomaly, optionally capped.

    A thin wrapper so the ceiling applies to every scheme without threading it through
    each one: the scheme dispatch below has several return points, and duplicating the
    clip at each is how one of them ends up missing it.

    Applied on the GCM grid, BEFORE `interpolate_to_target_latlon`. Capping afterwards
    would leave a blown-up cell already smeared across its neighbours by the regrid.

    The cap bounds each cell's multiplier at the granularity the scheme anchors at --
    annual for the yearly family, per calendar month for the monthly family -- and
    rescales the daily series to meet it.

    It cannot be an elementwise clip. `anomaly` carries the target's daily date
    dimension, because the denominator is a reference-window mean with no date dim and
    the division broadcasts. Clipping elementwise bounds every *day* at `anomaly_cap`
    times the cell's reference mean, and against an annual denominator that ceiling sits
    inside the ordinary distribution of daily rainfall: a cell averaging 2 mm/day is cut
    at 40 mm/day, an unremarkable tropical wet day. The daily-over-annual ratio also
    carries the seasonal cycle, so such a clip bites hardest where seasonality is
    strongest rather than where the pathology is.

    Measured on `yearly-delta` ssp126 before this change, an elementwise clip at 20
    altered 20.5% of land pixels in 2083 and 15.5% in 2050; 82% of those had an annual
    anomaly at or below 2 -- cells projecting essentially no change -- and the global
    total moved -19.1% and -1.1% respectively.

    Rescaling instead leaves any cell at or below the ceiling bit-identical to uncapped,
    brings a cell above it exactly to the ceiling, and preserves within-period shape so
    a wet day stays proportionally a wet day.

    The cap still only removes precipitation -- it cannot conserve, and whatever it
    removes is gone from the total. That is intended: the values it removes are ones no
    cell's observed climatology supports. But it means a cap moves the level, so it must
    be chosen on evidence rather than set defensively.
    """
    anomaly = _compute_anomaly_uncapped(
        reference,
        target,
        anomaly_type,
        debias_method=debias_method,
        dry_day_rule=dry_day_rule,
        anomaly_scheme=anomaly_scheme,
        eps_floor=eps_floor,
    )
    if anomaly_cap is None:
        return anomaly
    if anomaly_type != "multiplicative":
        msg = (
            f"anomaly_cap={anomaly_cap!r} was requested for an additive anomaly. The cap "
            "bounds a multiplier; an additive anomaly is a difference in the variable's "
            "own units, where a ceiling of this kind is meaningless."
        )
        raise ValueError(msg)
    if anomaly_cap <= 0:
        msg = f"anomaly_cap must be positive, got {anomaly_cap!r}."
        raise ValueError(msg)
    if anomaly_scheme in cdc.YEARLY_ANOMALY_SCHEMES:
        # anchored to the annual level, so bound the annual multiplier
        over = anomaly.mean("date")
        factor = (anomaly_cap / over).where(over > anomaly_cap, 1.0)
        return anomaly * factor

    # the monthly family anchors each calendar month to ERA5 separately, so the ceiling
    # belongs per month. Bounding the annual mean instead would let one blown-up month
    # drag the rescale factor down and crush the other eleven, which are fine.
    over = anomaly.groupby("date.month").mean("date")
    factor = (anomaly_cap / over).where(over > anomaly_cap, 1.0)
    return anomaly.groupby("date.month") * factor

jensen_debias_factor(reference: xr.Dataset, reference_monthly: xr.Dataset, debias_method: str, eps: xr.Dataset | float = 1.0) -> xr.Dataset

The factor to divide a multiplicative anomaly by, per month and per GCM cell.

The anomaly is (T + 1) / (R + 1) with R a monthly mean over only five reference years. 1/(R + 1) is convex, so by Jensen's inequality the anomaly averages above 1 even when the target year is drawn from the same distribution as the reference period -- a level bias on every forecast year. This returns an estimate of that inflation.

loo -- leave-one-out. For each held-out reference year, form the multiplier of that year against the mean of the other years and average over the folds. Between reference years there is no climate signal, so an unbiased estimator would return 1; the excess is the bias, measured from the data with no series expansion. The held-out denominator averages n-1 years while the pipeline averages n, and the bias goes as 1/n, so the excess is rescaled by (n-1)/n.

This is provably >= 1: with u_y = T_y + 1 and S = sum_y u_y the held-out denominator is (S - u_y)/(n-1), so each fold is (n-1)*u_y/(S - u_y), convex in u_y, and Jensen gives mean_y >= f(S/n) = 1 with equality iff every reference year is identical. So dividing by it can only shrink the anomaly, never inflate it -- which is what makes the effect on a threshold count such as precipitation_days sign-definite.

analytic -- the second-order expansion 1 + Var(Rbar)/(R + 1)^2. Cheaper to reason about but it is a truncated series, and the neglected terms matter exactly where the correction is largest (near-zero R, where eps dominates the denominator). Kept for comparison; loo is the estimator this was built for.

reference_monthly is the pipeline's own denominator, passed in rather than recomputed so the analytic form squares precisely the value the anomaly divides by.

Source code in src/climate_data/generate/scenario_daily.py
def jensen_debias_factor(
    reference: xr.Dataset,
    reference_monthly: xr.Dataset,
    debias_method: str,
    eps: xr.Dataset | float = 1.0,
) -> xr.Dataset:
    """The factor to divide a multiplicative anomaly by, per month and per GCM cell.

    The anomaly is ``(T + 1) / (R + 1)`` with ``R`` a monthly mean over only five reference
    years. ``1/(R + 1)`` is convex, so by Jensen's inequality the anomaly averages above 1 even
    when the target year is drawn from the same distribution as the reference period -- a level
    bias on every forecast year. This returns an estimate of that inflation.

    ``loo`` -- leave-one-out. For each held-out reference year, form the multiplier of that
    year against the mean of the *other* years and average over the folds. Between reference
    years there is no climate signal, so an unbiased estimator would return 1; the excess is
    the bias, measured from the data with no series expansion. The held-out denominator
    averages ``n-1`` years while the pipeline averages ``n``, and the bias goes as ``1/n``, so
    the excess is rescaled by ``(n-1)/n``.

    This is provably ``>= 1``: with ``u_y = T_y + 1`` and ``S = sum_y u_y`` the held-out
    denominator is ``(S - u_y)/(n-1)``, so each fold is ``(n-1)*u_y/(S - u_y)``, convex in
    ``u_y``, and Jensen gives ``mean_y >= f(S/n) = 1`` with equality iff every reference year
    is identical. So dividing by it can only shrink the anomaly, never inflate it -- which is
    what makes the effect on a threshold count such as ``precipitation_days`` sign-definite.

    ``analytic`` -- the second-order expansion ``1 + Var(Rbar)/(R + 1)^2``. Cheaper to reason
    about but it is a truncated series, and the neglected terms matter exactly where the
    correction is largest (near-zero ``R``, where ``eps`` dominates the denominator). Kept for
    comparison; ``loo`` is the estimator this was built for.

    ``reference_monthly`` is the pipeline's own denominator, passed in rather than recomputed so
    the analytic form squares precisely the value the anomaly divides by.
    """
    by_year = _monthly_means_by_reference_year(reference)
    n_years = by_year.sizes["reference_year"]

    if debias_method == "loo":
        mean_year = by_year.mean("reference_year")
        folds = []
        valid = []
        for i in range(n_years):
            held_out = by_year.isel(reference_year=i)
            others = (n_years * mean_year - held_out) / (n_years - 1)
            fold_denominator = others + eps
            folds.append(
                (held_out + eps) / fold_denominator.where(fold_denominator > 0)
            )
            valid.append((fold_denominator > 0).astype("int8"))
        raw = xr.concat(folds, dim="reference_year").mean("reference_year")
        factor = 1.0 + ((n_years - 1) / n_years) * (raw - 1.0)
        # Where any fold is undefined the estimate is undefined, so apply no correction.
        # Averaging the surviving folds is NOT a repair: the undefined one is precisely
        # the large fold that lifts the mean above 1, so dropping it collapses the factor
        # (0.2 for a one-wet-four-dry cell-month) and the de-bias INFLATES the anomaly --
        # the opposite of its purpose. Only reachable when eps can be zero, i.e. under the
        # taper; the constant eps = 1 floors every denominator at 1.
        factor = factor.where(sum(valid) == n_years, 1.0).fillna(1.0)
    elif debias_method == "analytic":
        variance = by_year.var("reference_year", ddof=1)
        factor = 1.0 + (variance / n_years) / (reference_monthly + eps) ** 2
    else:
        msg = f"Unknown debias method: {debias_method}"
        raise ValueError(msg)

    factor = factor.drop_vars("reference_year", errors="ignore")
    values = factor.to_dataarray()
    if not bool(np.isfinite(values).all()):
        msg = (
            f"Non-finite value in the {debias_method} de-bias factor. Interpolation would "
            "silently fill it from a neighbour rather than surface it, so refusing to proceed."
        )
        raise ValueError(msg)
    minimum = float(values.min())
    if minimum < 1.0 - _DEBIAS_FACTOR_TOLERANCE:
        # Both estimators are >= 1 by construction -- loo by Jensen on a convex fold,
        # analytic because it adds a non-negative term. A value below 1 means the
        # construction's precondition failed, and dividing by it would inflate rather
        # than shrink the anomaly. Refuse rather than ship an anti-correction.
        msg = (
            f"The {debias_method} de-bias factor reached {minimum!r}, below its "
            "guaranteed lower bound of 1. Dividing by it would inflate the anomaly."
        )
        raise ValueError(msg)
    return factor

load_and_shift_longitude_and_correct_time(member_path: str | Path, year: str) -> xr.Dataset

Put a member's year onto the real Gregorian calendar, day by day.

The conversion is by DATE, never by value. interp_calendar used to resample onto the target axis by linear interpolation, so whenever the source calendar's year length differed from the target's -- a noleap member in a leap year -- every target day became a blend of two source days. convert_calendar maps each source day onto its own date instead and leaves 29 February missing; the reindex holds the output axis to exactly this year's days whatever the source calendar, and interpolate_na fills the gap from the nearest real day.

Every variable comes through here, not just precipitation, but what the blending cost depends on the field. Temperature is smooth, so blending barely moves it and the threshold measures it feeds -- days_over_30C, the suitability maps -- were knocked either way and largely cancelled in the spatial mean. Precipitation is spiky and mostly zero against a 0.1 mm cut sitting on the floor, so smearing a wet day onto its dry neighbours could only push them UP over the line, never below it. That one-way ratchet is why precipitation_days took a coherent ~13.5 d per noleap member in all 19 leap years 2024-2096 while everything else came out as cancelling noise. (CLIMATE-35)

No align_on is passed: it only takes effect when a 360_day calendar is involved, and no member in the extracts uses one. If that changes, align_on needs a deliberate choice rather than xarray's default.

Source code in src/climate_data/generate/scenario_daily.py
def load_and_shift_longitude_and_correct_time(
    member_path: str | Path,
    year: str,
) -> xr.Dataset:
    """Put a member's year onto the real Gregorian calendar, day by day.

    The conversion is by DATE, never by value. `interp_calendar` used to resample onto the
    target axis by linear interpolation, so whenever the source calendar's year length
    differed from the target's -- a `noleap` member in a leap year -- every target day
    became a blend of two source days. `convert_calendar` maps each source day onto its own
    date instead and leaves 29 February missing; the `reindex` holds the output axis to
    exactly this year's days whatever the source calendar, and `interpolate_na` fills the
    gap from the nearest real day.

    Every variable comes through here, not just precipitation, but what the blending cost
    depends on the field. Temperature is smooth, so blending barely moves it and the
    threshold measures it feeds -- `days_over_30C`, the suitability maps -- were knocked
    either way and largely cancelled in the spatial mean. Precipitation is spiky and mostly
    zero against a 0.1 mm cut sitting on the floor, so smearing a wet day onto its dry
    neighbours could only push them UP over the line, never below it. That one-way ratchet
    is why `precipitation_days` took a coherent ~13.5 d per noleap member in all 19 leap
    years 2024-2096 while everything else came out as cancelling noise. (CLIMATE-35)

    No `align_on` is passed: it only takes effect when a `360_day` calendar is involved,
    and no member in the extracts uses one. If that changes, `align_on` needs a deliberate
    choice rather than xarray's default.
    """
    time_slice = slice(f"{year}-01-01", f"{year}-12-31")
    time_range = pd.date_range(f"{year}-01-01", f"{year}-12-31")
    ds = load_and_shift_longitude(member_path, time_slice)
    ds = (
        ds.assign_coords(time=ds.time.dt.floor("D"))
        .convert_calendar("standard")
        .reindex(time=time_range)
        .interpolate_na(dim="time", method="nearest", fill_value="extrapolate")
        .rename({"time": "date"})
    )
    return ds

variables_for_anomaly_scheme(target_variables: list[str], anomaly_scheme: str, anomaly_types: dict[str, str]) -> list[str]

Drop target variables the anomaly scheme cannot be applied to.

The yearly schemes are defined for multiplicative variables only -- compute_anomaly raises for anything else -- while --target-variable defaults to ALL. Without this filter the default invocation of a yearly scheme submits additive jobs that are certain to fail, and they fail only after being scheduled and retried.

Skipped variables are named rather than dropped quietly, and selecting nothing but additive variables is an error rather than an empty run.

Source code in src/climate_data/generate/scenario_daily.py
def variables_for_anomaly_scheme(
    target_variables: list[str],
    anomaly_scheme: str,
    anomaly_types: dict[str, str],
) -> list[str]:
    """Drop target variables the anomaly scheme cannot be applied to.

    The yearly schemes are defined for multiplicative variables only -- `compute_anomaly`
    raises for anything else -- while `--target-variable` defaults to ALL. Without this
    filter the default invocation of a yearly scheme submits additive jobs that are
    certain to fail, and they fail only after being scheduled and retried.

    Skipped variables are named rather than dropped quietly, and selecting nothing but
    additive variables is an error rather than an empty run.
    """
    keep = []
    skip = []
    skip_bounded = []
    for variable in target_variables:
        if (
            variable in cdc.VALUE_BOUNDS
            and anomaly_scheme not in cdc.BOUNDED_VARIABLE_SCHEMES
        ):
            # Value-bounded variables are floored explicitly; see
            # check_bounded_variable_scheme for why the other schemes are refused.
            skip_bounded.append(variable)
        elif (
            anomaly_scheme != cdc.ANOMALY_SCHEME_MONTHLY
            and anomaly_types[variable] != "multiplicative"
        ):
            skip.append(variable)
        else:
            keep.append(variable)

    if skip:
        print(
            f"Anomaly scheme '{anomaly_scheme}' applies to multiplicative variables only;"
            f" skipping {len(skip)}: {', '.join(skip)}."
        )
    if skip_bounded:
        print(
            f"Anomaly scheme '{anomaly_scheme}' does not apply to value-bounded variables,"
            f" which run under {', '.join(cdc.BOUNDED_VARIABLE_SCHEMES)} only;"
            f" skipping {len(skip_bounded)}: {', '.join(skip_bounded)}."
        )
    if not keep:
        msg = (
            f"No multiplicative variables selected, so anomaly scheme '{anomaly_scheme}'"
            f" has nothing to run. Selected: {', '.join(target_variables)}."
        )
        raise click.UsageError(msg)
    return keep

utils

annual_mean_from_monthly(ds: xr.Dataset) -> xr.Dataset

Day-weighted annual mean of a 12-month climatology.

Weights are bound by month label, not position, so a permuted month coordinate still weights correctly; anything other than a full 1..12 coordinate is rejected rather than misweighted.

Source code in src/climate_data/generate/utils.py
def annual_mean_from_monthly(ds: xr.Dataset) -> xr.Dataset:
    """Day-weighted annual mean of a 12-month climatology.

    Weights are bound by month label, not position, so a permuted month
    coordinate still weights correctly; anything other than a full 1..12
    coordinate is rejected rather than misweighted.
    """
    months = np.arange(1, 13)
    if not np.array_equal(np.sort(ds["month"].to_numpy()), months):
        msg = f"expected a full 1..12 month coordinate, got {ds['month'].to_numpy()}"
        raise ValueError(msg)
    weights = xr.DataArray(cdc.DAYS_IN_MONTH, dims=["month"], coords={"month": months})
    return ds.weighted(weights).mean("month")

buck_vapor_pressure(temperature_c: xr.Dataset) -> xr.Dataset

Approximate vapor pressure of water.

https://en.wikipedia.org/wiki/Arden_Buck_equation https://journals.ametsoc.org/view/journals/apme/20/12/1520-0450_1981_020_1527_nefcvp_2_0_co_2.xml

Parameters

temperature_c Temperature in Celsius

Returns

xr.Dataset Vapor pressure in hPa

Source code in src/climate_data/generate/utils.py
def buck_vapor_pressure(temperature_c: xr.Dataset) -> xr.Dataset:
    """Approximate vapor pressure of water.

    https://en.wikipedia.org/wiki/Arden_Buck_equation
    https://journals.ametsoc.org/view/journals/apme/20/12/1520-0450_1981_020_1527_nefcvp_2_0_co_2.xml

    Parameters
    ----------
    temperature_c
        Temperature in Celsius

    Returns
    -------
    xr.Dataset
        Vapor pressure in hPa
    """
    over_water = 6.1121 * np.exp(
        (18.678 - temperature_c / 234.5) * (temperature_c / (257.14 + temperature_c))
    )
    over_ice = 6.1115 * np.exp(
        (23.036 - temperature_c / 333.7) * (temperature_c / (279.82 + temperature_c))
    )
    vp = xr.where(temperature_c > 0, over_water, over_ice)  # type: ignore[no-untyped-call]
    return vp  # type: ignore[no-any-return]

daily_accumulation_last(ds: xr.Dataset) -> xr.Dataset

Collapse a within-day cumulative variable to its daily total.

For ERA5-Land total_precipitation, which accumulates since 00Z, the day's total is the window's closing sample, so last reads it directly. Prefer this to a maximum: the two agree only while the window rises monotonically, and int16 packing can make it tick down in its final step, in which case the maximum is an earlier, quantisation-inflated sample. Measured over 741,270 land day-pixels, last matched the true close for 100.0000% against 99.9864% for max.

Note resample emits a regular axis with NaN for empty bins, unlike groupby, which emits observed groups only. Callers must trim the partial bins at each end.

skipna=False because the default steps back past an absent closing sample to hour 23, reintroducing the incomplete window this function exists to eliminate. NaN is the honest answer for a window that never closed. It is not, however, a detectable one: generate_historical_daily_main fills ERA5-Land NaNs from the interpolated single-level field -- that is how ocean pixels are supplied -- so such a pixel carries a complete 0.25 degree value rather than a truncated 0.1 degree one, and never reaches validate_output. Preferring the coarse-but-whole value is the point; detecting the substitution would need a separate check against the sea mask.

Source code in src/climate_data/generate/utils.py
def daily_accumulation_last(ds: xr.Dataset) -> xr.Dataset:
    """Collapse a within-day cumulative variable to its daily total.

    For ERA5-Land `total_precipitation`, which accumulates since 00Z, the day's total is
    the window's closing sample, so `last` reads it directly. Prefer this to a maximum:
    the two agree only while the window rises monotonically, and int16 packing can make
    it tick down in its final step, in which case the maximum is an earlier,
    quantisation-inflated sample. Measured over 741,270 land day-pixels, `last` matched
    the true close for 100.0000% against 99.9864% for `max`.

    Note `resample` emits a regular axis with NaN for empty bins, unlike `groupby`, which
    emits observed groups only. Callers must trim the partial bins at each end.

    `skipna=False` because the default steps back past an absent closing sample to hour
    23, reintroducing the incomplete window this function exists to eliminate. NaN is the
    honest answer for a window that never closed. It is not, however, a detectable one:
    `generate_historical_daily_main` fills ERA5-Land NaNs from the interpolated
    single-level field -- that is how ocean pixels are supplied -- so such a pixel carries
    a complete 0.25 degree value rather than a truncated 0.1 degree one, and never reaches
    `validate_output`. Preferring the coarse-but-whole value is the point; detecting the
    substitution would need a separate check against the sea mask.
    """
    resampled = ds.resample(
        time=ACCUMULATION_FREQ,
        closed=ACCUMULATION_CLOSED,
        label=ACCUMULATION_LABEL,
    ).last(skipna=False)
    return resampled.rename({"time": "date"})

daily_accumulation_sum(ds: xr.Dataset) -> xr.Dataset

Collapse a variable of per-hour increments to its daily total.

For ERA5 single-levels, "accumulations are over the hour ending at the validity date/time", so the increments sum. The same interval convention applies as for daily_accumulation_last, so the two stay on a common date axis -- they are merged downstream in generate_historical_daily_main.

Source code in src/climate_data/generate/utils.py
def daily_accumulation_sum(ds: xr.Dataset) -> xr.Dataset:
    """Collapse a variable of per-hour increments to its daily total.

    For ERA5 single-levels, "accumulations are over the hour ending at the validity
    date/time", so the increments sum. The same interval convention applies as for
    `daily_accumulation_last`, so the two stay on a common `date` axis -- they are merged
    downstream in `generate_historical_daily_main`.
    """
    resampled = ds.resample(
        time=ACCUMULATION_FREQ,
        closed=ACCUMULATION_CLOSED,
        label=ACCUMULATION_LABEL,
    ).sum()
    return resampled.rename({"time": "date"})

identity(ds: xr.Dataset) -> xr.Dataset

Identity transformation

Source code in src/climate_data/generate/utils.py
def identity(ds: xr.Dataset) -> xr.Dataset:
    """Identity transformation"""
    return ds

interpolate_to_target_latlon(ds: xr.Dataset, method: str = 'nearest', target_lon: xr.DataArray = cdc.TARGET_LONGITUDE, target_lat: xr.DataArray = cdc.TARGET_LATITUDE) -> xr.Dataset

Interpolate a dataset to a target latitude and longitude grid.

Parameters

ds Dataset to interpolate method Interpolation method target_lon Target longitude grid target_lat Target latitude grid

Returns

xr.Dataset Interpolated dataset

Source code in src/climate_data/generate/utils.py
def interpolate_to_target_latlon(
    ds: xr.Dataset,
    method: str = "nearest",
    target_lon: xr.DataArray = cdc.TARGET_LONGITUDE,
    target_lat: xr.DataArray = cdc.TARGET_LATITUDE,
) -> xr.Dataset:
    """Interpolate a dataset to a target latitude and longitude grid.

    Parameters
    ----------
    ds
        Dataset to interpolate
    method
        Interpolation method
    target_lon
        Target longitude grid
    target_lat
        Target latitude grid

    Returns
    -------
    xr.Dataset
        Interpolated dataset
    """
    return (
        ds.interp(longitude=target_lon, latitude=target_lat, method=method)  # type: ignore[arg-type]
        .interpolate_na(dim="longitude", method="nearest", fill_value="extrapolate")
        .sortby("latitude")
        .interpolate_na(dim="latitude", method="nearest", fill_value="extrapolate")
        .sortby("latitude", ascending=False)
    )

kelvin_to_celsius(temperature_k: xr.Dataset) -> xr.Dataset

Convert temperature from Kelvin to Celsius

Parameters

temperature_k Temperature in Kelvin

Returns

xr.Dataset Temperature in Celsius

Source code in src/climate_data/generate/utils.py
def kelvin_to_celsius(temperature_k: xr.Dataset) -> xr.Dataset:
    """Convert temperature from Kelvin to Celsius

    Parameters
    ----------
    temperature_k
        Temperature in Kelvin

    Returns
    -------
    xr.Dataset
        Temperature in Celsius
    """
    return temperature_k - 273.15

meter_to_millimeter(rainfall_m: xr.Dataset) -> xr.Dataset

Convert rainfall from meters to millimeters

Parameters

rainfall_m Rainfall in meters

Returns

xr.Dataset Rainfall in millimeters

Source code in src/climate_data/generate/utils.py
def meter_to_millimeter(rainfall_m: xr.Dataset) -> xr.Dataset:
    """Convert rainfall from meters to millimeters

    Parameters
    ----------
    rainfall_m
        Rainfall in meters

    Returns
    -------
    xr.Dataset
        Rainfall in millimeters
    """
    return 1000 * rainfall_m

parse_reference_years(reference_years: str) -> slice

Turn a START-END year string (inclusive, four-digit years) into a date slice.

Source code in src/climate_data/generate/utils.py
def parse_reference_years(reference_years: str) -> slice:
    """Turn a START-END year string (inclusive, four-digit years) into a date slice."""
    start, sep, end = reference_years.partition("-")
    valid = (
        bool(sep)
        and len(start) == 4  # noqa: PLR2004
        and len(end) == 4  # noqa: PLR2004
        and start.isdigit()
        and end.isdigit()
        and int(start) <= int(end)
    )
    if not valid:
        msg = (
            "reference-years must be 'START-END' with four-digit years and "
            f"START <= END, got {reference_years!r}"
        )
        raise ValueError(msg)
    return slice(f"{start}-01-01", f"{end}-12-31")

precipitation_flux_to_rainfall(precipitation_flux: xr.Dataset) -> xr.Dataset

Convert precipitation flux to rainfall

Parameters

precipitation_flux Precipitation flux in kg m-2 s-1

Returns

xr.Dataset Rainfall in mm/day

Source code in src/climate_data/generate/utils.py
def precipitation_flux_to_rainfall(precipitation_flux: xr.Dataset) -> xr.Dataset:
    """Convert precipitation flux to rainfall

    Parameters
    ----------
    precipitation_flux
        Precipitation flux in kg m-2 s-1

    Returns
    -------
    xr.Dataset
        Rainfall in mm/day
    """
    seconds_per_day = 86400
    mm_per_kg_m2 = 1
    return seconds_per_day * mm_per_kg_m2 * precipitation_flux

rh_percent(temperature_c: xr.Dataset, dewpoint_temperature_c: xr.Dataset) -> xr.Dataset

Calculate relative humidity from temperature and dewpoint temperature.

Parameters

temperature_c Temperature in Celsius dewpoint_temperature_c Dewpoint temperature in Celsius

Returns

xr.Dataset Relative humidity as a percentage

Source code in src/climate_data/generate/utils.py
def rh_percent(
    temperature_c: xr.Dataset, dewpoint_temperature_c: xr.Dataset
) -> xr.Dataset:
    """Calculate relative humidity from temperature and dewpoint temperature.

    Parameters
    ----------
    temperature_c
        Temperature in Celsius
    dewpoint_temperature_c
        Dewpoint temperature in Celsius

    Returns
    -------
    xr.Dataset
        Relative humidity as a percentage
    """
    # saturation vapour pressure
    svp = buck_vapor_pressure(temperature_c)
    # actual vapour pressure
    vp = buck_vapor_pressure(dewpoint_temperature_c)
    return 100 * vp / svp

scale_wind_speed_height(wind_speed_10m: xr.Dataset) -> xr.Dataset

Scaling wind speed from a height of 10 meters to a height of 2 meters

Reference: Bröde et al. (2012) https://doi.org/10.1007/s00484-011-0454-1

Parameters

wind_speed_10m The 10m wind speed [m/s]. May be signed (ie a velocity component)

Returns

xr.DataSet The 2m wind speed [m/s]. May be signed (ie a velocity component)

Source code in src/climate_data/generate/utils.py
def scale_wind_speed_height(wind_speed_10m: xr.Dataset) -> xr.Dataset:
    """Scaling wind speed from a height of 10 meters to a height of 2 meters

    Reference: Bröde et al. (2012)
    https://doi.org/10.1007/s00484-011-0454-1

    Parameters
    ----------
    wind_speed_10m
        The 10m wind speed [m/s]. May be signed (ie a velocity component)

    Returns
    -------
    xr.DataSet
        The 2m wind speed [m/s]. May be signed (ie a velocity component)
    """
    scale_factor = np.log10(2 / 0.01) / np.log10(10 / 0.01)
    return scale_factor * wind_speed_10m  # type: ignore[no-any-return]

vector_magnitude(x: xr.Dataset, y: xr.Dataset) -> xr.Dataset

Calculate the magnitude of a vector.

Source code in src/climate_data/generate/utils.py
def vector_magnitude(x: xr.Dataset, y: xr.Dataset) -> xr.Dataset:
    """Calculate the magnitude of a vector."""
    return np.sqrt(x**2 + y**2)  # type: ignore[return-value]

jobmon_utils

Helpers for working with jobmon in the Climate Data CLI.

The main entrypoint is run_parallel_maybe_dry_run which mirrors rra_tools.jobmon.run_parallel but adds a dry_run flag.

When dry_run is True, this helper: * Expands flat_node_args or node_args into per-job CLI argument sets. * Builds sbatch-like preview commands using the supplied task resources. * Prints a short summary plus representative example commands. * Returns a dummy success status without touching jobmon or the scheduler.

When dry_run is False, it simply delegates to jobmon.run_parallel with identical semantics.

run_parallel_maybe_dry_run(*, runner: str, task_name: str, task_resources: Mapping[str, Any], flat_node_args: tuple[Sequence[str], Sequence[Sequence[Any]]] | None = None, node_args: Mapping[str, Sequence[Any]] | None = None, task_args: Mapping[str, Any] | None = None, log_root: str | Path | None = None, max_attempts: int | None = None, concurrency_limit: int | None = None, dry_run: bool) -> Any

Wrapper around jobmon.run_parallel with optional dry-run behavior.

Parameters

runner The executable used for tasks, e.g. "cdtask" or "cdtask aggregate". task_name A short, human-readable name for the job group. task_resources Resource profile for the tasks (queue, cores, memory, runtime, project, ...). flat_node_args Tuple of (argument names, rows) describing per-node CLI arguments. node_args Mapping from argument name to a sequence of values; all combinations are used. task_args Shared CLI arguments that are the same for all nodes. log_root Optional log directory passed through to jobmon. max_attempts Maximum number of attempts; passed through to jobmon. Defaults to None to match jobmon.run_parallel (jobmon then applies its own default). dry_run If True, only print sbatch-like previews instead of submitting.

Source code in src/climate_data/jobmon_utils.py
def run_parallel_maybe_dry_run(
    *,
    runner: str,
    task_name: str,
    task_resources: Mapping[str, Any],
    flat_node_args: tuple[Sequence[str], Sequence[Sequence[Any]]] | None = None,
    node_args: Mapping[str, Sequence[Any]] | None = None,
    task_args: Mapping[str, Any] | None = None,
    log_root: str | Path | None = None,
    max_attempts: int | None = None,
    concurrency_limit: int | None = None,
    dry_run: bool,
) -> Any:
    """Wrapper around ``jobmon.run_parallel`` with optional dry-run behavior.

    Parameters
    ----------
    runner
        The executable used for tasks, e.g. ``"cdtask"`` or ``"cdtask aggregate"``.
    task_name
        A short, human-readable name for the job group.
    task_resources
        Resource profile for the tasks (queue, cores, memory, runtime, project, ...).
    flat_node_args
        Tuple of (argument names, rows) describing per-node CLI arguments.
    node_args
        Mapping from argument name to a sequence of values; all combinations are used.
    task_args
        Shared CLI arguments that are the same for all nodes.
    log_root
        Optional log directory passed through to jobmon.
    max_attempts
        Maximum number of attempts; passed through to jobmon. Defaults to ``None``
        to match ``jobmon.run_parallel`` (jobmon then applies its own default).
    dry_run
        If ``True``, only print sbatch-like previews instead of submitting.
    """
    if not dry_run:
        kwargs: dict[str, Any] = {
            "runner": runner,
            "task_name": task_name,
            "flat_node_args": flat_node_args,
            "node_args": node_args,
            "task_args": task_args,
            "task_resources": task_resources,
            "log_root": log_root,
            "max_attempts": max_attempts,
        }
        if concurrency_limit is not None:
            kwargs["concurrency_limit"] = concurrency_limit
        return jobmon.run_parallel(**kwargs)

    normalized_task_args = _normalize_task_args(task_args)

    jobs: list[dict[str, Any]] = []
    if flat_node_args is not None:
        jobs = _iter_jobs_from_flat_node_args(flat_node_args)
    elif node_args is not None:
        jobs = _iter_jobs_from_node_args(node_args)
    else:
        jobs = [{}]

    n_jobs = len(jobs)
    queue = task_resources.get("queue", "")
    cores = task_resources.get("cores", "")
    memory = task_resources.get("memory", "")
    runtime = task_resources.get("runtime", "")
    project = task_resources.get("project", "")

    print(
        "[DRY-RUN] "
        f"Would submit {n_jobs} job{'s' if n_jobs != 1 else ''} "
        f"for task '{task_name}' via runner '{runner}'."
    )
    print("  Resources:")
    print(f"    queue   = {queue}")
    print(f"    cores   = {cores}")
    print(f"    memory  = {memory}")
    print(f"    runtime = {runtime}")
    print(f"    project = {project}")

    n_examples = min(n_jobs, _MAX_EXAMPLE_JOBS)
    if n_examples > 0:
        print(f"  Example sbatch-like commands (showing {n_examples} of {n_jobs}):")
        for job in jobs[:n_examples]:
            line = _format_sbatch_like_line(
                runner=runner,
                task_name=task_name,
                job_args=job,
                task_args=normalized_task_args,
                task_resources=task_resources,
            )
            print("   ", line)
        if n_jobs > n_examples:
            print(
                f"  ... ({n_jobs - n_examples} more job"
                f"{'s' if n_jobs - n_examples != 1 else ''} not shown)"
            )

    # Return a dummy "success" status so callers that expect a jobmon status
    # code continue without raising errors.
    return "D"

special

download_era5_uncertainty

ERA5 Uncertainty Download (CLIMATE-22 gap-fill)

Downloads whole-year ERA5 reanalysis and ensemble_spread 2m temperature files from the Copernicus CDS, matching the layout Katrin Burkart produced by hand in /ihme/erf/ERA5/katrin_download_20240208/ for 2022/2023:

era5_{product_type}_{variable}_{year}.nc   (one file per whole year)

This fills the 2024/2025 gap that GBD temperature PAF work depends on. It is a standalone gap-fill: the ensemble_spread product is not (yet) consumed by the rest of this repo's pipeline, which only uses reanalysis mean temperature.

Notes / gotchas baked in below: * The ensemble products only exist on reanalysis-era5-single-levels (0.25° HRES / ~0.5° EDA), never on reanalysis-era5-land. * Both products are pulled 3-hourly to match the canonical GBD download and Katrin's 2022/2023 files. ensemble_spread (the EDA) is only 3-hourly; reanalysis (HRES) is subsampled from hourly to the same 3-hourly axis. * Credentials come from the caller's ~/.cdsapirc (cdsapi.Client() with no args), so no dependency on the shared per-user copernicus.yaml keyring (which does not carry every user).

download_era5_uncertainty(era5_variable: str, output_dir: str, queue: str, dry_run: bool) -> None

Download 2024/2025 ERA5 reanalysis + ensemble_spread (CLIMATE-22 gap-fill).

Source code in src/climate_data/special/download_era5_uncertainty.py
@click.command()
@clio.with_era5_variable()
@clio.with_output_directory(DEFAULT_OUTPUT_DIR)
@clio.with_queue()
@clio.with_dry_run()
def download_era5_uncertainty(
    era5_variable: str,
    output_dir: str,
    queue: str,
    dry_run: bool,
) -> None:
    """Download 2024/2025 ERA5 reanalysis + ensemble_spread (CLIMATE-22 gap-fill)."""
    node_args = {
        "year": ERA5_UNCERTAINTY_YEARS,
        "product-type": ERA5_PRODUCT_TYPES,
    }

    run_parallel_maybe_dry_run(
        runner="cdtask special",
        task_name="download_era5_uncertainty",
        node_args=node_args,
        task_args={
            "era5-variable": era5_variable,
            "output-dir": output_dir,
        },
        task_resources={
            "queue": queue,
            "cores": 1,
            "memory": "20G",
            "runtime": "720m",
            "project": "proj_rapidresponse",
        },
        max_attempts=1,
        dry_run=dry_run,
    )

download_era5_uncertainty_task(year: str, product_type: str, era5_variable: str, output_dir: str) -> None

Download one whole-year ERA5 file (single year, single product type).

Source code in src/climate_data/special/download_era5_uncertainty.py
@click.command()
@clio.with_year(years=cdc.HISTORY_YEARS)
@click.option(
    "--product-type",
    required=True,
    type=click.Choice(ERA5_PRODUCT_TYPES),
    help="ERA5 product type to download.",
)
@clio.with_era5_variable()
@clio.with_output_directory(DEFAULT_OUTPUT_DIR)
def download_era5_uncertainty_task(
    year: str,
    product_type: str,
    era5_variable: str,
    output_dir: str,
) -> None:
    """Download one whole-year ERA5 file (single year, single product type)."""
    download_era5_uncertainty_main(year, product_type, era5_variable, output_dir)

temperature_person_days

temperature_person_days_main(block_key: str, gcm_member: str, scenario: str, hierarchy: str, population_model_root: str, climate_data_root: str, output_dir: str, *, progress_bar: bool = False) -> None

Bin population into (temperature x temperature-zone) person-days per year.

The year span is derived from the temperature_zone actually on disk (not a hardcoded range): for the historical scenario this is the ERA5-only product spanning EXPOSURE_START_YEAR through the last year present (1990-2025 as delivered); for a forecast scenario, daily temperature is ERA5 before FORECAST_START_YEAR and the GCM scenario from then on, binned against that scenario's own zone. A year whose daily or population input is missing fails loudly rather than being skipped, so a gap in a product that must be square cannot slip through silently.

Source code in src/climate_data/special/temperature_person_days.py
def temperature_person_days_main(
    block_key: str,
    gcm_member: str,
    scenario: str,
    hierarchy: str,
    population_model_root: str,
    climate_data_root: str,
    output_dir: str,
    *,
    progress_bar: bool = False,
) -> None:
    """Bin population into (temperature x temperature-zone) person-days per year.

    The year span is derived from the ``temperature_zone`` actually on disk (not a
    hardcoded range): for the ``historical`` scenario this is the ERA5-only product
    spanning ``EXPOSURE_START_YEAR`` through the last year present (1990-2025 as
    delivered); for a forecast scenario, daily temperature is ERA5 before
    ``FORECAST_START_YEAR`` and the GCM scenario from then on, binned against that
    scenario's own zone. A year whose daily or population input is missing fails loudly
    rather than being skipped, so a gap in a product that must be square cannot slip
    through silently.
    """
    print(f"Aggregating {gcm_member} for {block_key}")
    pm_data = PopulationModelData(population_model_root)
    cd_data = ClimateData(climate_data_root, read_only=True)
    ca_data = ClimateAggregateData(Path(output_dir) / hierarchy)

    print("Building location masks")
    climate_slice, location_ids, location_idx = utils.build_location_index(
        hierarchy, block_key, pm_data
    )

    print("Building data index")
    temperature_bins = np.arange(-35, 45, 0.1)
    temperature_zone_bins = np.arange(-25, 35, 1)
    data_idx = pd.MultiIndex.from_product(
        [location_ids, temperature_bins, temperature_zone_bins]
    )
    out_template = np.zeros(
        (len(location_ids), len(temperature_bins), len(temperature_zone_bins)),
        dtype=np.float64,
    )

    print("Building historical temperature zone index")
    temperature_zone = cd_data.load_compiled_annual_results(
        scenario, "temperature_zone", gcm_member
    ).sel(**climate_slice)  # type: ignore[arg-type]
    historical_temperature_zone_idx = utils.to_idx(
        temperature_zone, temperature_zone_bins
    )

    print("Building temperature coordinates")
    temperature_coordinates = utils.get_temperature_coordinates(
        block_key, pm_data, temperature_zone
    )

    print("Aggregating temperature person days")
    # Drive the span from the zone actually on disk (historical: 1990..last present;
    # forecast: the compiled historical+forecast series). A year whose daily or
    # population input is missing now fails loudly rather than being skipped, so a gap
    # in a product that must be square can't slip through silently.
    zone_years = [int(y) for y in temperature_zone["year"].to_numpy()]
    zone_row_for_year = {y: i for i, y in enumerate(zone_years)}
    years = [y for y in zone_years if y >= cdc.EXPOSURE_START_YEAR]
    dfs = []
    for year in tqdm.tqdm(years, disable=not progress_bar):
        if scenario == "historical" or year < FORECAST_START_YEAR:
            temperature = cd_data.load_daily_results(
                "historical", "mean_temperature", year
            ).sel(**climate_slice)  # type: ignore[arg-type]
        else:
            temperature = cd_data.load_raw_daily_results(
                scenario, "mean_temperature", year, gcm_member
            ).sel(**climate_slice)  # type: ignore[arg-type]
        # Population nodata is stored as NaN. The aggregate stage tolerates this via
        # np.nansum, but here we accumulate into `out_arr` with `+=`, so a NaN pixel
        # would poison every output cell it touches (silently read back as 0, which
        # zeroed small locations like American Samoa). Zero-fill nodata first: no
        # modeled population contributes no person-days.
        pop_arr = pm_data.load_results(f"{year}q1", block_key)._ndarray.flatten()  # noqa: SLF001
        pop_arr = np.nan_to_num(pop_arr, nan=0.0)
        temperature_idx = utils.to_idx(temperature, temperature_bins)

        out_arr = out_template.copy()
        utils.compute_person_days(
            location_idx,
            temperature_idx,
            historical_temperature_zone_idx[zone_row_for_year[year]],
            pop_arr,
            temperature_coordinates,
            out_arr,
        )

        df = (
            pd.DataFrame({"person_days": out_arr.reshape(-1)}, index=data_idx)
            .assign(year=year)
            .set_index("year", append=True)
        )
        df.index.names = ["location_id", "temperature", "temperature_zone", "year"]
        df = df.reset_index()
        df["temperature"] = df["temperature"].round(1)
        df = df.set_index(["location_id", "year", "temperature_zone", "temperature"])[
            "person_days"
        ].unstack()
        df.columns.name = None
        dfs.append(df)

    if not dfs:
        msg = (
            f"No person-days produced for {scenario} {gcm_member} {block_key}: the "
            f"temperature zone has no year >= {cdc.EXPOSURE_START_YEAR}."
        )
        raise ValueError(msg)
    df = pd.concat(dfs)
    out_path = ca_data.person_days_path(block_key, scenario, gcm_member)
    mkdir(out_path.parent, parents=True, exist_ok=True)
    save_parquet(df, out_path)

temperature_zone

generate_temperature_zone_main(gcm_member: str, scenario: str, output_dir: str | Path) -> None

Generate the temperature zone for a given scenario and gcm member.

Parameters

gcm_member The gcm member to generate the temperature zone for. scenario The scenario to generate the temperature zone for. Pass historical (with gcm_member="era5") to build a pure-ERA5 zone from the raw historical annual mean temperature rather than the compiled historical+forecast series; the output then spans EXPOSURE_START_YEAR through the last historical year present on disk. output_dir The directory to save the temperature zone to (root for this run mode).

Source code in src/climate_data/special/temperature_zone.py
def generate_temperature_zone_main(
    gcm_member: str,
    scenario: str,
    output_dir: str | Path,
) -> None:
    """Generate the temperature zone for a given scenario and gcm member.

    Parameters
    ----------
    gcm_member
        The gcm member to generate the temperature zone for.
    scenario
        The scenario to generate the temperature zone for.  Pass ``historical``
        (with ``gcm_member="era5"``) to build a pure-ERA5 zone from the raw
        historical annual mean temperature rather than the compiled
        historical+forecast series; the output then spans ``EXPOSURE_START_YEAR``
        through the last historical year present on disk.
    output_dir
        The directory to save the temperature zone to (root for this run mode).
    """
    print(f"Generating temperature zone for {scenario} {gcm_member}")
    cdata = ClimateData(output_dir)
    if scenario == "historical":
        ds = cdata.load_raw_annual_mfdataset("historical", "mean_temperature")
    else:
        ds = cdata.load_compiled_annual_results(
            scenario, "mean_temperature", gcm_member
        )
    temperature_zone = (
        ds.rolling(year=10)
        .mean()
        .sel(year=slice(cdc.EXPOSURE_START_YEAR, int(cdc.FORECAST_YEARS[-1])))
    )
    print(f"Saving temperature zone for {scenario} {gcm_member}")
    cdata.save_compiled_annual_results(
        temperature_zone,
        scenario=scenario,
        variable="temperature_zone",
        gcm_member=gcm_member,
        encoding_kwargs={"scale_factor": 0.01, "add_offset": 0.0},
    )

utils

aggregate_to_hierarchy(data: pd.DataFrame, hierarchy: pd.DataFrame) -> pd.DataFrame

Create all aggregate climate values for a given hierarchy from most-detailed data.

Parameters

data The most-detailed climate data to aggregate. hierarchy The hierarchy to aggregate the data to.

Returns

pd.DataFrame The climate data with values for all levels of the hierarchy.

Source code in src/climate_data/special/utils.py
def aggregate_to_hierarchy(data: pd.DataFrame, hierarchy: pd.DataFrame) -> pd.DataFrame:
    """Create all aggregate climate values for a given hierarchy from most-detailed data.

    Parameters
    ----------
    data
        The most-detailed climate data to aggregate.
    hierarchy
        The hierarchy to aggregate the data to.

    Returns
    -------
    pd.DataFrame
        The climate data with values for all levels of the hierarchy.
    """
    agg_cols = sorted(
        set(data.columns) - {"location_id", "year_id", "temperature_zone"}
    )

    results = data.set_index("location_id")

    # Most detailed locations can be at multiple levels of the hierarchy,
    # so we loop over all levels from most detailed to global, aggregating
    # level by level and appending the results to the data.

    for level in reversed(list(range(1, hierarchy.level.max() + 1))):
        level_mask = hierarchy.level == level
        parent_map = hierarchy.loc[level_mask].set_index("location_id").parent_id

        subset = results.loc[results.index.intersection(parent_map.index)]
        subset["parent_id"] = parent_map

        parent_values = (
            subset.groupby(["year_id", "parent_id", "temperature_zone"])[agg_cols]
            .sum()
            .reset_index()
            .rename(columns={"parent_id": "location_id"})
            .set_index("location_id")
        )
        results = pd.concat([results, parent_values])
    results = (
        results.reset_index()
        .sort_values(["location_id", "year_id", "temperature_zone"])
        .reset_index(drop=True)
    )

    return results

utils

Climate Data Utilities

Utility functions for working with climate data.

make_raster_template(x_min: int | float, y_min: int | float, stride: int | float, resolution: int | float, crs: str = 'EPSG:4326') -> rt.RasterArray

Create a raster template with the specified dimensions and resolution.

A raster template is a RasterArray with a specified extent, resolution, and CRS. The data values are initialized to zero. This function is useful for creating a template to use when resampling another raster to a common grid.

Parameters

x_min The minimum x-coordinate of the raster. y_min The minimum y-coordinate of the raster. stride The length of one side of the raster in the x and y directions measured in the units of the provided coordinate reference system. resolution The resolution of the raster in the units of the provided coordinate reference system. crs The coordinate reference system of the generated raster.

Returns

rt.RasterArray A raster template with the specified dimensions and resolution.

Source code in src/climate_data/utils.py
def make_raster_template(
    x_min: int | float,
    y_min: int | float,
    stride: int | float,
    resolution: int | float,
    crs: str = "EPSG:4326",
) -> rt.RasterArray:
    """Create a raster template with the specified dimensions and resolution.

    A raster template is a RasterArray with a specified extent, resolution, and CRS. The data
    values are initialized to zero. This function is useful for creating a template to use
    when resampling another raster to a common grid.

    Parameters
    ----------
    x_min
        The minimum x-coordinate of the raster.
    y_min
        The minimum y-coordinate of the raster.
    stride
        The length of one side of the raster in the x and y directions measured in the units
        of the provided coordinate reference system.
    resolution
        The resolution of the raster in the units of the provided coordinate reference system.
    crs
        The coordinate reference system of the generated raster.

    Returns
    -------
    rt.RasterArray
        A raster template with the specified dimensions and resolution.
    """
    tolerance = 1e-12
    evenly_divides = (stride % resolution < tolerance) or (
        resolution - stride % resolution < tolerance
    )
    if not evenly_divides:
        msg = "Stride must be a multiple of resolution"
        raise ValueError(msg)

    transform = Affine(
        a=resolution,
        b=0,
        c=x_min,
        d=0,
        e=-resolution,
        f=y_min + stride,
    )

    n_pix = int(stride / resolution)

    data = np.zeros((n_pix, n_pix), dtype=np.int8)
    return rt.RasterArray(
        data,
        transform,
        crs=crs,
        no_data_value=-1,
    )

to_raster(ds: xr.DataArray, no_data_value: float | int, lat_col: str = 'lat', lon_col: str = 'lon', crs: str = 'EPSG:4326') -> rt.RasterArray

Convert an xarray DataArray to a RasterArray.

Parameters

ds The xarray DataArray to convert. no_data_value The value to use for missing data. This should be consistent with the dtype of the data. lat_col The name of the latitude coordinate in the dataset. lon_col The name of the longitude coordinate in the dataset. crs The coordinate reference system of the data.

Returns

rt.RasterArray The RasterArray representation of the input data.

Source code in src/climate_data/utils.py
def to_raster(
    ds: xr.DataArray,
    no_data_value: float | int,
    lat_col: str = "lat",
    lon_col: str = "lon",
    crs: str = "EPSG:4326",
) -> rt.RasterArray:
    """Convert an xarray DataArray to a RasterArray.

    Parameters
    ----------
    ds
        The xarray DataArray to convert.
    no_data_value
        The value to use for missing data. This should be consistent with the dtype of the data.
    lat_col
        The name of the latitude coordinate in the dataset.
    lon_col
        The name of the longitude coordinate in the dataset.
    crs
        The coordinate reference system of the data.

    Returns
    -------
    rt.RasterArray
        The RasterArray representation of the input data.
    """
    lat, lon = ds[lat_col].data, ds[lon_col].data

    dlat = (lat[1:] - lat[:-1]).mean()
    dlon = (lon[1:] - lon[:-1]).mean()

    transform = Affine(
        a=dlon,
        b=0.0,
        c=lon[0],
        d=0.0,
        e=-dlat,
        f=lat[-1],
    )
    return rt.RasterArray(
        data=ds.data[::-1],
        transform=transform,
        crs=crs,
        no_data_value=no_data_value,
    )