Utilities API Reference

The utilities module provides helper functions for climate data processing.

Overview

The utilities include:

  • Coordinate name detection and validation

  • Data selection and processing

  • Spatial averaging functions

  • Dask client management

  • Advanced chunking strategies and optimization

Available Functions

Coordinate Utilities

climate_diagnostics.utils.coord_utils.get_coord_name(xarray_like_obj: DataArray | Dataset, possible_names: List[str]) → str | None[source]

Find the name of a coordinate in an xarray object from a list of possible names.

This function checks for coordinate names in a case-sensitive manner first, then falls back to a case-insensitive check.

Parameters:

  • xarray_like_obj (xr.DataArray or xr.Dataset) – The xarray object to search for coordinates.

  • possible_names (list of str) – A list of possible coordinate names to look for.

Returns:

The found coordinate name, or None if no matching coordinate is found.

Return type:

str or None

climate_diagnostics.utils.coord_utils.filter_by_season(data_subset: DataArray | Dataset, season: str = 'annual') → DataArray | Dataset[source]

Filter climate data for a specific season using xarray’s time accessors.

This function implements robust seasonal filtering that handles various time coordinate formats, including standard datetime64 and cftime objects, in a performant, Dask-aware manner.

Parameters:

  • data_subset (xr.DataArray or xr.Dataset) – The climate data to filter by season. Must have a recognizable time dimension.

  • season (str, optional) – The season to filter by. Defaults to ‘annual’. Supported: ‘annual’, ‘jjas’, ‘djf’, ‘mam’, ‘son’, ‘jja’.

Returns:

The filtered data containing only the specified season.

Return type:

xr.DataArray or xr.Dataset

Raises:

ValueError – If a usable time coordinate cannot be found or processed, or if the season is invalid.

Notes

  • This function relies on xarray’s .dt accessor.

  • Note on ‘DJF’ (Winter): This function filters for Dec, Jan, and Feb based on the month index. It does NOT automatically shift December to align with the Jan/Feb of the following year.

Data Utilities

climate_diagnostics.utils.data_utils.validate_and_get_sel_slice(coord_val_param: float | int | slice | List | ndarray, data_coord: DataArray, coord_name_str: str) → Tuple[float | int | slice | List | ndarray, bool][source]

Validate a coordinate selection parameter against the data’s coordinate range.

Parameters:

  • coord_val_param (Union[float, int, slice, List, np.ndarray]) – The value(s) to select.

  • data_coord (xr.DataArray) – The coordinate array from the dataset (e.g., ds[‘lat’]).

  • coord_name_str (str) – Name of the coordinate for logging/error messages.

Returns:

Returns the sanitized selection value and a boolean indicating if method=’nearest’ is required.

Return type:

Tuple[SelectionValue, bool]

climate_diagnostics.utils.data_utils.select_process_data(xarray_obj: Dataset, variable: str, latitude: float | slice | List | None = None, longitude: float | slice | List | None = None, level: float | slice | List | None = None, time_range: slice | None = None, season: str = 'annual', year: int | None = None) → DataArray[source]

Select, filter, and process a data variable from the dataset.

Parameters:

  • xarray_obj (xr.Dataset) – The input dataset.

  • variable (str) – Name of the variable to select.

  • latitude (Optional) – Selection parameters (slice, list, or single value).

  • longitude (Optional) – Selection parameters (slice, list, or single value).

  • level (Optional) – Selection parameters (slice, list, or single value).

  • time_range (slice, optional) – Range of times to select.

  • season (str, optional) – Season to filter (e.g., ‘annual’, ‘djf’, ‘jjas’).

  • year (int, optional) – Specific year to filter.

Returns:

The processed DataArray.

Return type:

xr.DataArray

Raises:

ValueError – If the variable is missing or selection results in empty data.

climate_diagnostics.utils.data_utils.get_spatial_mean(data_var: DataArray, area_weighted: bool = True) → DataArray[source]

Calculate the spatial mean of a DataArray.

Parameters:

  • data_var (xr.DataArray) – Input data.

  • area_weighted (bool, optional) – If True (default), weights the mean by cos(latitude).

Returns:

The spatially averaged data.

Return type:

xr.DataArray

Dask Utilities

climate_diagnostics.utils.dask_utils.managed_dask_client(**kwargs: Any) → Iterator[Client][source]

A context manager to get an existing Dask client or create/manage a new one.

This function provides robust resource management for Dask clients in library code. If a client already exists, it yields that client without closing it on exit. If no client exists, it creates a new one, yields it, and then cleanly shuts it down upon exiting the context.

Parameters:

**kwargs (Any) – Keyword arguments for the dask.distributed.Client constructor.

Yields:

dask.distributed.Client – The active Dask client.

climate_diagnostics.utils.dask_utils.get_or_create_dask_client(**kwargs: Any) → Client[source]

Get an active Dask client or create a new one with specified settings.

This function provides a centralized way to manage Dask client connections. It reuses an existing client if available or creates a new one.

Parameters:

**kwargs (Any) – Keyword arguments to be passed to the dask.distributed.Client constructor.

Returns:

The active Dask client.

Return type:

dask.distributed.Client

Chunking Utilities

Basic Examples

Working with Coordinates

from climate_diagnostics.utils import get_coord_name

# Find coordinate names automatically
time_coord = get_coord_name(ds, ['time', 't'])
lat_coord = get_coord_name(ds, ['lat', 'latitude'])

print(f"Time coordinate: {time_coord}")
print(f"Latitude coordinate: {lat_coord}")

Spatial Averaging

from climate_diagnostics.utils import get_spatial_mean

# Calculate area-weighted spatial mean
global_mean = get_spatial_mean(ds.air, area_weighted=True)

# Simple spatial mean (no area weighting)
simple_mean = get_spatial_mean(ds.air, area_weighted=False)

Data Selection

from climate_diagnostics.utils import select_process_data

# Select and process data with automatic coordinate handling
processed_data = select_process_data(
    ds,
    variable="air",
    latitude=slice(30, 60),
    longitude=slice(-10, 40),
    season="jja"  # Summer season
)

Advanced Chunking

The chunking_utils module provides sophisticated tools for optimizing Dask chunking strategies, which is critical for performance when working with large climate datasets. The dynamic_chunk_calculator function automatically determines optimal chunk sizes based on the operation being performed and the desired performance characteristics.

from climate_diagnostics.utils.chunking_utils import (
    dynamic_chunk_calculator,
    suggest_chunking_strategy,
    print_chunking_info
)
import xarray as xr

# Load a sample dataset
ds = xr.tutorial.load_dataset("air_temperature")

# Calculate optimal chunks for a time-series analysis that is memory-intensive
optimal_chunks = dynamic_chunk_calculator(
    ds,
    operation_type='time-series',
    performance_priority='memory'
)
print("Optimal chunks for memory-optimized time-series analysis:", optimal_chunks)

# Rechunk the dataset with the optimal chunking scheme
ds_optimized = ds.chunk(optimal_chunks)

# Print chunking information to verify
print_chunking_info(ds_optimized)

Seasonal Filtering

from climate_diagnostics.utils import filter_by_season

# Filter data by season
summer_data = filter_by_season(ds, season="jja")
winter_data = filter_by_season(ds, season="djf")

Practical Usage

Complete Analysis Workflow

import xarray as xr
from climate_diagnostics.utils import (
    get_coord_name,
    select_process_data,
    get_spatial_mean
)

# Load data
ds = xr.open_dataset("temperature_data.nc")

# Check coordinates
time_coord = get_coord_name(ds, ['time', 't'])
print(f"Time coordinate found: {time_coord}")

# Select and process regional data
arctic_data = select_process_data(
    ds,
    variable="air",
    latitude=slice(60, 90),
    season="annual"
)

# Calculate regional mean
arctic_mean = get_spatial_mean(arctic_data, area_weighted=True)

# Plot results
import matplotlib.pyplot as plt
arctic_mean.plot()
plt.title("Arctic Mean Temperature")
plt.show()

Memory-Efficient Processing

from climate_diagnostics.utils import get_or_create_dask_client

# Ensure Dask client is available for large datasets
client = get_or_create_dask_client()

# Process large dataset
large_ds = xr.open_dataset("large_file.nc", chunks={'time': 100})
result = get_spatial_mean(large_ds.air, area_weighted=True)

# Compute result
computed_result = result.compute()

See Also