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