AID_BC package

Submodules

AID_BC.bias_corrector module

Bias-correction interfaces and implementations.

This module defines a common interface for bias-correction methods and provides wrappers for Quantile Mapping and entropy-regularized Optimal Transport.

All correction methods operate on two-dimensional arrays with shape (n_samples, n_features). In the climate-data workflow, samples correspond to time steps and features correspond to flattened spatial grid points, possibly concatenated across several variables.

class AID_BC.bias_corrector.BiasCorrector

Bases: ABC

Common interface for bias-correction methods.

All arrays follow the convention (n_samples, n_features). In the climate-data workflow, samples correspond to time steps and features correspond to flattened spatial grid points, possibly concatenated across several variables.

is_fitted

Whether the corrector has been fitted.

Type:

bool

n_features

Number of expected input features after fitting.

Type:

int or None

method_name = 'base'

__init__()

abstractmethod fit(reference, biased)

Fit the bias-correction model.

Parameters:

• reference (array-like) – Reference observations with shape (n_reference_samples, n_features), for example ERA5 data.

• biased (array-like) – Biased training observations with shape (n_biased_samples, n_features), for example historical CMIP6 data.

Returns:

Fitted corrector instance.

Return type:

BiasCorrector

abstractmethod transform(data)

Correct new biased data.

Parameters:

data (array-like) – Biased application data with shape (n_samples, n_features).

Returns:

Corrected data with the same shape as data.

Return type:

numpy.ndarray

fit_transform(reference, biased, data)

Fit the corrector and transform application data.

Parameters:

• reference (array-like) – Reference observations with shape (n_reference_samples, n_features), for example ERA5 data.

• biased (array-like) – Biased training observations with shape (n_biased_samples, n_features).

• data (array-like) – Biased application data with shape (n_samples, n_features), for example historical CMIP6 data.

Returns:

Corrected application data.

Return type:

numpy.ndarray

property diagnostics

Return method diagnostics.

Returns:

Method name, fitting state, and number of input features.

Return type:

dict

static validate_fit_arrays(reference, biased)

Validate reference and biased training arrays.

Parameters:

• reference (numpy.ndarray) – Reference array with shape (n_samples, n_features).

• biased (numpy.ndarray) – Biased training array with shape (n_samples, n_features).

Raises:

ValueError – If an array is not two-dimensional, contains no samples, contains non-finite values, or if the feature dimensions differ.

validate_transform_array(data)

Validate application data before transformation.

Parameters:

data (numpy.ndarray) – Application array with shape (n_samples, n_features).

Raises:

• RuntimeError – If the corrector has not been fitted or its feature count is unavailable.

• ValueError – If data is not two-dimensional, has an unexpected number of features, or contains non-finite values.

class AID_BC.bias_corrector.QuantileMappingCorrector

Bases: BiasCorrector

Wrapper around the existing SBCK Quantile Mapping implementation.

method_name = 'qm'

__init__()

fit(reference, biased)

Fit empirical Quantile Mapping.

Parameters:

• reference (array-like) – Reference observations with shape (n_reference_samples, n_features).

• biased (array-like) – Biased training observations with shape (n_biased_samples, n_features).

Returns:

Fitted corrector instance.

Return type:

QuantileMappingCorrector

transform(data)

Apply the fitted Quantile Mapping model.

Parameters:

data (array-like) – Biased application data with shape (n_samples, n_features).

Returns:

Corrected data as float32.

Return type:

numpy.ndarray

Raises:

RuntimeError – If the underlying Quantile Mapping model is unavailable.

class AID_BC.bias_corrector.OptimalTransportCorrector(epsilon=1000, num_iterations=15000, threshold=0.01, batch_size=16, dtype='float64', normalize=True)

Bases: BiasCorrector

Optimal Transport bias corrector.

Parameters:

• epsilon (float, default=1000) – Entropic regularization strength.

• num_iterations (int, default=15000) – Maximum number of Sinkhorn iterations.

• threshold (float, default=1e-2) – Convergence threshold.

• batch_size (int, default=16) – Number of application samples transported together.

• dtype ({"float32", "float64"}, default="float64") – Floating-point precision used by NumPy and JAX.

• normalize (bool, default=True) – Whether to fit and apply a shared per-feature normalization.

center

Per-feature normalization center.

Type:

numpy.ndarray or None

scale

Per-feature normalization scale.

Type:

numpy.ndarray or None

target_points

Normalized ERA5 target samples.

Type:

jax.Array or None

target_potential

Target dual potential produced by Sinkhorn.

Type:

jax.Array or None

solver

Sinkhorn solver used during fitting.

Type:

AID_BC.optimal_transport.SinkhornSolver or None

transport_function

Vectorized transport function.

Type:

callable or None

converged

Whether Sinkhorn converged.

Type:

bool or None

fitted_iterations

Number of iterations performed.

Type:

int or None

regularized_ot_cost

Regularized transport cost.

Type:

float or None

method_name = 'ot'

__init__(epsilon=1000, num_iterations=15000, threshold=0.01, batch_size=16, dtype='float64', normalize=True)

fit(reference, biased)

Fit the Optimal Transport map from CMIP6 to ERA5.

Parameters:

• reference (array-like) – ERA5 reference observations with shape (n_reference_samples, n_features).

• biased (array-like) – Historical CMIP6 observations with shape (n_biased_samples, n_features).

Returns:

Fitted corrector instance.

Return type:

OptimalTransportCorrector

Raises:

RuntimeError – If the Sinkhorn solver does not converge.

transform(data)

Apply the fitted Optimal Transport map in batches.

Parameters:

data (array-like) – Biased CMIP6 application data with shape (n_samples, n_features).

Returns:

Corrected data as float32.

Return type:

numpy.ndarray

Raises:

RuntimeError – If the transport function is unavailable.

property diagnostics

Return Optimal Transport fitting diagnostics.

Returns:

Base diagnostics together with backend, devices, solver configuration, convergence information, cost, dtype, batch size, and normalization state.

Return type:

dict

AID_BC.bias_corrector.create_bias_corrector(method, *, ot_epsilon=1000, ot_num_iterations=15000, ot_threshold=0.01, ot_batch_size=16, ot_dtype='float64', ot_normalize=True)

Create a bias-correction method.

Parameters:

• method ({"qm", "ot"}) – Bias-correction method identifier.

• ot_epsilon (float, default=1000) – Entropic regularization used by OT.

• ot_num_iterations (int, default=15000) – Maximum number of Sinkhorn iterations.

• ot_threshold (float, default=1e-2) – Sinkhorn convergence threshold.

• ot_batch_size (int, default=16) – Number of application samples transported together.

• ot_dtype ({"float32", "float64"}, default="float64") – Floating-point precision used by OT.

• ot_normalize (bool, default=True) – Whether OT applies shared per-feature normalization.

Returns:

New, unfitted corrector instance.

Return type:

QuantileMappingCorrector or OptimalTransportCorrector

Raises:

ValueError – If method is neither “qm” nor “ot”.

AID_BC.dataset module

AID_BC.dataset.resize_with_torchvision(da, target_shape)

Resize a DataArray’s spatial dimensions using torchvision’s bilinear interpolation with antialiasing, matching the resizing method used in DataPreprocessor.coarse_down_up during model training.

Parameters:

• da (xarray.DataArray) – Input data array with dims (time, latitude, longitude).

• target_shape (tuple of int) – Target (H, W) shape for latitude and longitude.

Returns:

Resized array of shape (time, H, W). Coordinates are NOT attached; the caller must reassign latitude/longitude coordinates explicitly.

Return type:

np.ndarray

class AID_BC.dataset.ClimateDataset(cmip6_path, variable_name, era5_path=None, logger=None)

Bases: object

Climate dataset preprocessing for native CMIP6-grid bias correction.

This class: - loads ERA5 and CMIP6 datasets, - renames coordinates to a common ERA5-style convention, - normalizes longitudes to [0, 360), - sorts latitudes consistently, - interpolates ERA5 onto the native CMIP6 grid when ERA5 is provided, - keeps CMIP6 on its native grid.

The class supports two modes:

1. ERA5 and CMIP6 mode Used for training or evaluation. ERA5 is interpolated onto the native CMIP6 grid and returned together with CMIP6.

2. CMIP6-only mode Used for application periods where ERA5 is not available or not required. CMIP6 is standardized and returned on its native grid.

era5_data

ERA5 variable interpolated onto the native CMIP6 grid.

Type:

xarray.DataArray or None

cmip6_data

CMIP6 variable on its native grid.

Type:

xarray.DataArray or None

__init__(cmip6_path, variable_name, era5_path=None, logger=None)

Initialize the climate dataset handler.

Parameters:

• cmip6_path (str) – Path to the CMIP6 NetCDF file.

• variable_name (str) – Name of the climate variable to process.

• era5_path (str, optional) – Path to the ERA5 NetCDF file. If None, only CMIP6 is prepared.

• logger (AID_BC.logger.Logger, optional) – Custom logger instance.

load_cmip6()

Load the CMIP6 dataset and check that the target variable exists.

Raises:

ValueError – If the requested variable is not found in the CMIP6 dataset.

load_era5()

Load the ERA5 dataset and check that the target variable exists.

Raises:

ValueError – If no ERA5 path is provided or if the requested variable is not found in the ERA5 dataset.

static rename_coordinates(ds)

Rename common CMIP-style coordinates to ERA5-style coordinate names.

Parameters:

ds (xarray.Dataset or xarray.DataArray) – Input dataset or data array.

Returns:

ds – Dataset or data array with standardized coordinate names.

Return type:

xarray.Dataset or xarray.DataArray

Raises:

ValueError – If latitude or longitude coordinates are missing.

static is_latitude_descending(ds)

Check whether latitude is ordered from north to south.

Parameters:

ds (xarray.Dataset or xarray.DataArray) – Input dataset or data array with a latitude coordinate.

Returns:

True if latitude is descending, False otherwise.

Return type:

bool

static normalize_longitudes(ds)

Normalize longitudes to [0, 360), sort them, and remove duplicates.

Parameters:

ds (xarray.Dataset or xarray.DataArray) – Input dataset or data array with a longitude coordinate.

Returns:

ds – Dataset or data array with normalized and sorted longitudes.

Return type:

xarray.Dataset or xarray.DataArray

static sort_latitude(ds, descending=False)

Sort latitude either ascending or descending. ERA5 often uses descending latitude.

Parameters:

• ds (xarray.Dataset or xarray.DataArray) – Input dataset or data array with a latitude coordinate.

• descending (bool, default=False) – If True, sort latitude from north to south. If False, sort latitude from south to north.

Returns:

ds – Dataset or data array with sorted latitude.

Return type:

xarray.Dataset or xarray.DataArray

static standardize_coordinates(ds, latitude_descending=False)

Apply coordinate naming and ordering conventions.

Parameters:

• ds (xarray.Dataset or xarray.DataArray) – Input dataset or data array.

• latitude_descending (bool, default=False) – Whether the output latitude coordinate should be descending.

Returns:

ds – Dataset or data array with standardized coordinates.

Return type:

xarray.Dataset or xarray.DataArray

prepare_dataset(latitude_descending=None)

Run the dataset preparation pipeline.

If ERA5 is available, ERA5 is interpolated onto the native CMIP6 grid. This mode is used for training and evaluation.

If ERA5 is not available, only CMIP6 is prepared. This mode is used for application periods, including future CMIP6 simulations.

Parameters:

latitude_descending (bool or None) –

Latitude direction to use when preparing CMIP6 without ERA5. Required when era5_path is None.

If ERA5 is provided, this argument is ignored because the original ERA5 latitude convention is detected automatically and used for the final output ordering.

Returns:

If ERA5 is provided, returns:

era5_on_cmip6xarray.DataArray

ERA5 variable interpolated onto the native CMIP6 grid.

cmip6_nativexarray.DataArray

CMIP6 variable on its native grid.

era5_lat_descendingbool

Whether the original ERA5 latitude coordinate is descending.

If ERA5 is not provided, returns:

cmip6_nativexarray.DataArray

CMIP6 variable on its native grid.

Return type:

tuple or xarray.DataArray

Raises:

ValueError – If latitude_descending is not provided in CMIP6-only mode.

close()

Close opened xarray datasets.

This should be called after the prepared DataArrays have been loaded into memory or are no longer needed. Closing datasets prevents too many NetCDF files from remaining open during multi-year processing.

AID_BC.diagnostics module

class AID_BC.diagnostics.PlotConfig

Bases: object

Central configuration for all plotting functions.

DEFAULT_SAVE_DIR = './results'

DEFAULT_FIGSIZE_MULTIPLIER = 4

COLORMAPS = {'10u': 'BrBG_r', '10v': 'BrBG_r', '2t': 'rainbow', 'TP': 'Blues', 'curl': 'seismic', 'd2m': 'rainbow', 'default': 'viridis', 'dewpoint': 'rainbow', 'divergence': 'seismic', 'error': 'Reds', 'humid': 'Greens', 'humidity': 'Greens', 'mae': 'Reds', 'meridional': 'BrBG_r', 'precipitation': 'Blues', 'pres': 'viridis', 'pressure': 'viridis', 'speed': 'coolwarm', 'ssr': 'seismic', 'st': 'rainbow', 'surface temperature': 'rainbow', 'temp': 'rainbow', 'temperature': 'rainbow', 'tp': 'Blues', 'wind': 'coolwarm', 'zonal': 'BrBG_r'}

FIXED_DIFF_RANGES = {'10u': (-5.0, 5.0), '10v': (-5.0, 5.0), '2t': (-5.0, 5.0), 'T2M': (-5.0, 5.0), 'TP': (-0.5, 0.5), 'U10': (-5.0, 5.0), 'V10': (-5.0, 5.0), 'VAR_10U': (-5.0, 5.0), 'VAR_10V': (-5.0, 5.0), 'VAR_2T': (-5.0, 5.0), 'VAR_D2M': (-5.0, 5.0), 'VAR_ST': (-5.0, 5.0), 'VAR_TP': (-0.5, 0.5), 'meridional': (-5.0, 5.0), 'temperature': (-5.0, 5.0), 'tp': (-0.5, 0.5)}

COASTLINE_w = 0.5

BORDER_w = 0.5

LAKE_w = 0.5

BORDER_STYLE = '--'

COLORBAR_h = 0.02

COLORBAR_PAD = 0.05

classmethod get_colormap(variable_name)

Get appropriate colormap for a variable.

classmethod get_plot_name(variable_name)

Convert variable name to readable plot name.

classmethod convert_units(variable_name, data)

Safe unit conversion when required. - NEVER modifies input - Returns a new array only if conversion is needed

static get_fixed_diff_range(var_name)

Get fixed visualization range for signed differences (Prediction − Truth).

AID_BC.diagnostics.to_numpy_4d(data)

Convert input data to a four-dimensional NumPy array.

The returned array follows the dimension order “[time, variable, latitude, longitude]”.

Parameters:

data (torch.Tensor, numpy.ndarray, xarray.DataArray, or list of xarray.DataArray) – Input data to convert.

Returns:

Four-dimensional array with shape “[time, variable, latitude, longitude]”.

Return type:

numpy.ndarray

Raises:

ValueError – If the converted data does not have four dimensions.

AID_BC.diagnostics.plot_validation_pdfs(predictions, targets, coarse_inputs=None, variable_names=None, filename='validation_pdfs.png', save_dir='./results', figsize_multiplier=4, save_npz=False)

Create PDF (Probability Density Function) plots comparing distributions of model predictions vs ground truth for all variables.

Parameters:

• predictions (torch.Tensor, numpy.ndarray, xarray.DataArray, or list of xarray.DataArray) – Model predictions of shape [batch_size, num_variables, h, w]

• targets (torch.Tensor, numpy.ndarray, xarray.DataArray, or list of xarray.DataArray) – Ground truth of shape [batch_size, num_variables, h, w]

• coarse_inputs (torch.Tensor, numpy.ndarray, xarray.DataArray, or list of xarray.DataArray, optional) – Coarse inputs of shape [batch_size, num_variables, h, w]

• variable_names (list of str, optional) – Names of the variables for subplot titles

• filename (str, optional) – Output filename

• save_dir (str, optional) – Directory to save the plot

• figsize_multiplier (int, optional) – Base size multiplier for subplots

• save_npz (bool, optional) – If True, saves the PDF diagnostics to a compressed .npz file.

Returns:

The function saves the plot to disk and does not return any value.

Return type:

None

Notes

• Creates horizontal subplots (one per variable) showing PDFs

• Each subplot shows up to 3 lines: Predictions, Ground Truth, and Coarse Inputs

• Uses automatic color and linestyle cycling based on global matplotlib settings

• Calculates and displays key statistics for each distribution

• Handles both PyTorch tensors and numpy arrays

Examples

>>> predictions = np.random.randn(10, 3, 64, 64)  # 10 samples, 3 variables
>>> targets = np.random.randn(10, 3, 64, 64)
>>> plot_validation_pdfs(predictions, targets, variable_names=['Temp', 'Pres', 'Humid'])

AID_BC.diagnostics.plot_power_spectra(predictions, targets, dlat, dlon, coarse_inputs=None, variable_names=None, filename='power_spectra_physical.png', save_dir='./results', figsize_multiplier=4, save_npz=False)

Calculate and plot power spectra with proper physical wavenumbers.

Parameters:

• predictions (torch.Tensor, numpy.ndarray, xarray.DataArray, or list of xarray.DataArray) – Model predictions of shape [batch_size, num_variables, nh, nw]

• targets (torch.Tensor, numpy.ndarray, xarray.DataArray, or list of xarray.DataArray) – Ground truth of shape [batch_size, num_variables, nh, nw]

• dlat (float) – Grid spacing in latitude (degrees)

• dlon (float) – Grid spacing in longitude (degrees)

• coarse_inputs (torch.Tensor, numpy.ndarray, xarray.DataArray, or list of xarray.DataArray, optional) – Coarse inputs of shape [batch_size, num_variables, nh, nw]

• variable_names (list of str, optional) – Names of the variable names for subplot titles

• filename (str, optional) – Output filename

• save_dir (str, optional) – Directory to save the plot

• figsize_multiplier (int, optional) – Base size multiplier for subplots

• save_npz (bool, optional) – If True, saves the PDF diagnostics to a compressed .npz file.

Return type:

None

AID_BC.diagnostics.calculate_psd2d_simple(field)

Simple 2D PSD calculation without preprocessing.

AID_BC.diagnostics.radial_average_psd(psd2d, k_mag, k_bins)

Radially average 2D PSD using wavenumber magnitude.

AID_BC.diagnostics.plot_qq_quantiles(predictions, targets, coarse_inputs, variable_names=None, units=None, quantiles=[0.9, 0.95, 0.975, 0.99, 0.995], filename='qq_quantiles.png', save_dir='./results', figsize_multiplier=4, save_npz=False)

Create QQ-plats at different quantiles comparing model predictions and coarse inputs against ground truth.

For each variable, plots quantiles of predictions and coarse inputs against quantiles of ground truth with a 1:1 reference line.

Parameters:

• predictions (torch.Tensor, numpy.ndarray, xarray.DataArray, or list of xarray.DataArray) – Model predictions of shape [batch_size, num_variables, h, w]

• targets (torch.Tensor, numpy.ndarray, xarray.DataArray, or list of xarray.DataArray) – Ground truth of shape [batch_size, num_variables, h, w]

• coarse_inputs (torch.Tensor, numpy.ndarray, xarray.DataArray, or list of xarray.DataArray) – Coarse inputs of shape [batch_size, num_variables, h, w]

• variable_names (list of str, optional) – Names of the variables for subplot titles. If None, uses [“VAR_0”, “VAR_1”, …]

• units (list of str, optional) – Units for each variable for axis labels. If None, uses empty strings.

• quantiles (list of float, optional) – Quantile values to plot (e.g., [0.90, 0.95, 0.975, 0.99, 0.995])

• filename (str, optional) – Output filename

• save_dir (str, optional) – Directory to save the plot

• figsize_multiplier (int, optional) – Base size multiplier for subplots

• save_npz (bool, optional) – If True, saves the PDF diagnostics to a compressed .npz file.

Returns:

save_path – Path to the saved figure

Return type:

str

AID_BC.diagnostics.plot_metrics_heatmap(summaries, metric_names, filename='validation_metrics_heatmap.png', save_dir='./results', figsize_multiplier=4)

Plot a heatmap of summary metrics.

Parameters:

• summaries (list of dict) – Summary metrics returned by summary_metrics.

• metric_names (list of str) – Metrics to display.

• filename (str, optional) – Output filename.

• save_dir (str, optional) – Directory where the image is saved.

• figsize_multiplier (float, optional) – Controls overall figure size.

Returns:

Path to the saved figure.

Return type:

str

AID_BC.diagnostics.plot_surface(predictions, targets, coarse_inputs, lat_1d, lon_1d, timestamp=None, variable_names=None, filename='forecast_plot.png', save_dir=None, figsize_multiplier=None)

Plot side-by-side forecast maps (coarse_inputs input, true target, model prediction, and difference) for one or more meteorological variables over a geographic domain.

Parameters:

• coarse_inputs (torch.Tensor or np.ndarray) – coarse_inputs-resolution input data with shape [1, n_vars, H, W].

• targets (torch.Tensor or np.ndarray) – Ground-truth high-resolution data with shape [1, n_vars, H, W].

• predictions (torch.Tensor or np.ndarray) – Model predictions at targets resolution with shape [1, n_vars, H, W].

• lat_1d (array-like) – 1D array of latitude coordinates with shape [H].

• lon_1d (array-like) – 1D array of longitude coordinates with shape [W].

• timestamp (datetime.datetime) – Forecast timestamp to include in the plot title.

• variable_names (list of str, optional) – Variable names or identifiers.

• filename (str, optional) – Output filename for saving the plot.

• save_dir (str, optional) – Directory to save the plot.

• figsize_multiplier (int, optional) – Base size multiplier for subplots.

Return type:

None

AID_BC.evaluater module

class AID_BC.evaluater.MetricTracker

Bases: object

Track summary statistics incrementally.

__init__()

reset()

Reset tracked statistics.

update(values)

Update statistics with new values.

Parameters:

values (array-like) – Values to accumulate.

getmean()

Return the mean.

getstd()

Return the standard deviation.

getmin()

Return the minimum.

getmax()

Return the maximum.

getsqrtmean()

Return the square root of the mean.

AID_BC.evaluater.summary_metrics(data, name, variable_name, ref=None)

Compute summary statistics for one dataset.

Parameters:

• data (array-like) – Data to evaluate.

• name (str) – Dataset name.

• ref (array-like or None, optional) – Reference data used to compute MAE and RMSE.

Returns:

Summary statistics.

Return type:

dict

AID_BC.evaluater.print_metrics(variable_name, summaries, logger)

Print summary metrics for one variable.

Parameters:

• variable_name (str) – Climate variable name.

• summaries (list of dict) – Summary statistics.

• logger (AID_BC.logger.Logger) – Logger used to report metrics.

AID_BC.evaluater.evaluate(reference, raw, corrected, variable_names, method, year, logger, run_name, results_dir='./results')

Evaluate bias-corrected CMIP6 data against ERA5.

Metrics are computed separately for each climate variable, while multivariable diagnostic figures contain one subplot per variable.

Parameters:

• reference (xarray.Dataset) – ERA5 reference data on the CMIP6 grid.

• raw (xarray.Dataset) – Raw CMIP6 application data.

• corrected (xarray.Dataset) – Bias-corrected CMIP6 data.

• variable_names (sequence of str) – Ordered climate variables to evaluate.

• method (str) – Bias-correction method name.

• year (int) – Evaluated application year.

• logger (AID_BC.logger.Logger) – Logger used to report evaluation progress.

• run_name (str) – Name of the bias-correction run used to create the corresponding subdirectory inside results_dir.

• results_dir (str or pathlib.Path, default="./results") – Root directory for evaluation results.

AID_BC.logger module

class AID_BC.logger.Logger(console_output=True, file_output=False, log_file='module_log_file.log', pretty_print=True, record=False)

Bases: object

__init__(console_output=True, file_output=False, log_file='module_log_file.log', pretty_print=True, record=False)

clear_logs()

Clear the stored Rich logs if record=True.

show_header(module_name)

Display startup banner.

start_task(task_name: str, description: str = '', **meta)

Display a clearly formatted ‘task start’ message with good spacing.

log_metrics()

Log pipeline metrics

info(message)

Formatted info message

warning(message)

Formatted warning message

success(message)

Custom success level (not default logging level)

step(step_name, message)

Highlight pipeline step events

exception(message, exception=None)

Display a formatted exception message with visual stack trace.

error(message, exception=None)

Display a formatted error log, optionally including exception trace.

AID_BC.main module

AID_BC.main.parse_args()

Parse command-line arguments.

Returns:

Parsed arguments.

Return type:

argparse.Namespace

AID_BC.main.parse_variable_paths(items, argument_name)

Parse variable-specific paths from command-line values.

Parameters:

• items (sequence of str or None) – Values written as VARIABLE=PATH. If None, no mapping is created.

• argument_name (str) – Name of the command-line argument, used in validation messages.

Returns:

Mapping from variable names to directory paths, or None when items is None.

Return type:

dict of str to str or None

Raises:

ValueError – If an item does not use the VARIABLE=PATH format, contains an empty variable name or path, or defines the same variable more than once.

AID_BC.main.validate_variable_paths(variable_names, paths, argument_name)

Validate that paths are defined for exactly the requested variables.

Parameters:

• variable_names (sequence of str) – Climate variables requested for the correction.

• paths (dict of str to str or None) – Mapping from variable names to directory paths.

• argument_name (str) – Name of the command-line argument, used in validation messages.

Raises:

ValueError – If the mapping is missing, if a requested variable has no path, or if a path is provided for an unrequested variable.

AID_BC.main.validate_args(args)

Validate general command-line arguments.

Parameters:

args (argparse.Namespace) – Parsed command-line arguments.

Raises:

ValueError – If year ranges are invalid, variables are missing or duplicated, QM is requested with more than one variable, OT parameters are non-positive, or both/neither ERA5 input modes are supplied.

AID_BC.main.build_path(root, year)

Build the yearly NetCDF path for a data directory.

Parameters:

• root (str or pathlib.Path) – Directory containing yearly NetCDF files.

• year (int) – Year used in the samples_<year>.nc filename.

Returns:

Full path to the yearly NetCDF file.

Return type:

str

AID_BC.main.check_same_spatial_grid(reference, target)

Check that two arrays use the same latitude-longitude grid.

Parameters:

• reference (xarray.DataArray) – Reference array defining the expected spatial grid.

• target (xarray.DataArray) – Array whose spatial grid is compared with reference.

Raises:

ValueError – If latitude or longitude dimensions are missing, dimension sizes differ, or coordinate values are not equal within numerical tolerance.

AID_BC.main.check_dataset_consistency(dataset, variable_names, label)

Validate dimensions, coordinates, and values across dataset variables.

Parameters:

• dataset (xarray.Dataset) – Dataset containing the climate variables to validate.

• variable_names (sequence of str) – Variable names expected in dataset.

• label (str) – Readable dataset label used in error messages.

Raises:

ValueError – If no variables are provided, a variable is missing, required dimensions are absent, spatial or temporal coordinates differ, or non-finite values are present.

AID_BC.main.select_month(data, month)

Select all time steps belonging to one calendar month.

Parameters:

• data (xarray.Dataset or xarray.DataArray) – Input climate data with a datetime-like time coordinate.

• month (int) – Calendar month in the inclusive range 1 to 12.

Returns:

Subset containing all selected-month samples across available years.

Return type:

xarray.Dataset or xarray.DataArray

AID_BC.main.prepare_dataarray(data)

Convert and standardize a climate data array for downstream processing.

The array is converted to float32, loaded into memory, and reordered to use the canonical dimension order time, latitude, longitude.

Parameters:

data (xarray.DataArray) – Climate data array to prepare.

Returns:

Prepared array with float32 values loaded into memory and dimensions ordered as time, latitude, longitude.

Return type:

xarray.DataArray

AID_BC.main.resolve_precomputed_variable_root(era5_on_cmip6_root, variable_name)

Resolve the precomputed ERA5 directory for one variable.

The function supports either a common directory containing metadata.json and yearly files or a parent directory containing one subdirectory per variable.

Parameters:

• era5_on_cmip6_root (str or pathlib.Path) – Common root or parent directory of precomputed ERA5-on-CMIP6 data.

• variable_name (str) – Climate variable whose directory is required.

Returns:

Directory containing the variable metadata and yearly files.

Return type:

pathlib.Path

AID_BC.main.load_precomputed_metadata(era5_on_cmip6_root, variable_name, logger)

Load and validate metadata for precomputed ERA5-on-CMIP6 data.

Parameters:

• era5_on_cmip6_root (str or pathlib.Path) – Common root or parent directory of precomputed ERA5 data.

• variable_name (str) – Requested climate variable.

• logger (AID_BC.logger.Logger) – Logger used to report the loaded metadata.

Returns:

• metadata (dict) – Metadata read from metadata.json.

• variable_root (pathlib.Path) – Directory containing the metadata and yearly variable files.

Raises:

• FileNotFoundError – If metadata.json cannot be found.

• ValueError – If the metadata variable name conflicts with variable_name.

AID_BC.main.open_dataarray(path, variable_name)

Open one variable from a NetCDF file and load it into memory.

Parameters:

• path (str or pathlib.Path) – Path to the NetCDF file.

• variable_name (str) – Variable to extract from the dataset.

Returns:

Requested variable loaded into memory.

Return type:

xarray.DataArray

Raises:

ValueError – If variable_name is not present in the file.

AID_BC.main.load_training_variable(train_start, train_end, cmip6_train_root, variable_name, logger, era5_root=None, era5_on_cmip6_root=None)

Load one training variable from ERA5 and historical CMIP6 files.

ERA5 data can either be loaded from raw files and prepared onto the CMIP6 grid, or loaded from previously prepared ERA5-on-CMIP6 files.

Parameters:

• train_start (int) – First training year.

• train_end (int) – Last training year, inclusive.

• cmip6_train_root (str or pathlib.Path) – Directory containing historical CMIP6 yearly files.

• variable_name (str) – Climate variable to load.

• logger (AID_BC.logger.Logger) – Logger used to report progress.

• era5_root (str or pathlib.Path or None, optional) – Directory containing raw ERA5 yearly files.

• era5_on_cmip6_root (str or pathlib.Path or None, optional) – Directory containing precomputed ERA5-on-CMIP6 fields.

Returns:

• reference_train (xarray.DataArray) – Concatenated ERA5 training data on the CMIP6 grid.

• biased_train (xarray.DataArray) – Concatenated historical CMIP6 training data.

• latitude_descending (bool) – Whether latitude is ordered north to south.

Raises:

• ValueError – If both or neither ERA5 input modes are provided, or if spatial grids or latitude orientation are inconsistent between years.

• RuntimeError – If no training year is loaded.

AID_BC.main.assemble_variable_dataset(variable_arrays, variable_names, label)

Assemble aligned climate variables into one dataset.

Auxiliary scalar coordinates, such as variable-specific measurement heights, are moved to variable attributes before merging. Time, latitude, and longitude coordinates must match exactly across variables.

Parameters:

• variable_arrays (dict of str to xarray.DataArray) – Mapping from climate variable names to data arrays.

• variable_names (sequence of str) – Ordered variables to include in the assembled dataset.

• label (str) – Human-readable dataset label used in error messages.

Returns:

Dataset containing all variables with shared time and spatial coordinates.

Return type:

xarray.Dataset

Raises:

ValueError – If variables cannot be aligned exactly or fail consistency checks.

AID_BC.main.load_training_data(train_start, train_end, era5_roots, era5_on_cmip6_roots, cmip6_train_roots, variable_names, logger)

Load and assemble all variables required for training.

Parameters:

• train_start (int) – First training year.

• train_end (int) – Last training year, inclusive.

• era5_roots (dict of str to str or None) – Raw ERA5 directory for each variable, or None when precomputed ERA5 data are used.

• era5_on_cmip6_roots (dict of str to str or None) – Precomputed ERA5-on-CMIP6 directory for each variable, or None when raw ERA5 data are used.

• cmip6_train_roots (dict of str to str) – Historical CMIP6 directory for each variable.

• variable_names (sequence of str) – Ordered climate variables to load.

• logger (AID_BC.logger.Logger) – Logger used to report progress.

Returns:

• reference_train (xarray.Dataset) – ERA5 training variables on the CMIP6 grid.

• biased_train (xarray.Dataset) – Historical CMIP6 training variables.

• latitude_descending (bool) – Common latitude orientation used by all variables.

Raises:

ValueError – If variables have inconsistent latitude orientation, grids, times, or values.

AID_BC.main.load_application_data(year, era5_roots, era5_on_cmip6_roots, cmip6_apply_roots, variable_names, latitude_descending, logger)

Load and assemble ERA5 reference and CMIP6 data for one application year.

Parameters:

• year (int) – Application year to load.

• era5_roots (dict of str to str or None) – Raw ERA5 directory for each variable.

• era5_on_cmip6_roots (dict of str to str or None) – Precomputed ERA5-on-CMIP6 directory for each variable.

• cmip6_apply_roots (dict of str to str) – CMIP6 application directory for each variable.

• variable_names (sequence of str) – Ordered climate variables to load.

• latitude_descending (bool) – Latitude orientation established during training.

• logger (AID_BC.logger.Logger) – Logger used to report progress.

Returns:

• reference_apply (xarray.Dataset) – ERA5 reference data on the CMIP6 grid.

• biased_apply (xarray.Dataset) – Raw CMIP6 application data.

AID_BC.main.flatten_dataset(dataset, variable_names)

Flatten and concatenate climate variables into a feature matrix.

Variables are concatenated by blocks in the order given by variable_names. Each block contains all latitude-longitude grid points for one variable.

Parameters:

• dataset (xarray.Dataset) – Dataset with dimensions time, latitude, and longitude.

• variable_names (sequence of str) – Ordered variables to concatenate.

Returns:

• matrix (numpy.ndarray) – Array of shape (time, n_variables * latitude * longitude).

• metadata (dict) – Variable order and original dimensions required for reconstruction.

Raises:

ValueError – If the dataset fails consistency checks.

AID_BC.main.reconstruct_dataset(corrected_2d, template_dataset, metadata, method_name)

Reconstruct corrected climate fields from a feature matrix.

Parameters:

• corrected_2d (numpy.ndarray) – Corrected matrix with shape (time, n_variables * latitude * longitude).

• template_dataset (xarray.Dataset) – Dataset providing coordinates, variable attributes, and output structure.

• metadata (dict) – Flattening metadata produced by flatten_dataset.

• method_name (str) – Name of the applied bias-correction method.

Returns:

Corrected variables reconstructed on their original grid.

Return type:

xarray.Dataset

Raises:

RuntimeError – If corrected_2d does not have the expected shape.

AID_BC.main.fit_bias_corrector(reference_train, biased_train, variable_names, corrector, logger)

Fit one bias corrector on flattened climate fields.

For QM, command-line validation ensures that only one climate variable is provided. For OT, all requested variables and grid points are concatenated and corrected jointly.

Parameters:

• reference_train (xarray.Dataset) – ERA5 training data on the CMIP6 grid.

• biased_train (xarray.Dataset) – Historical CMIP6 training data.

• variable_names (sequence of str) – Ordered variables used to construct the feature matrices.

• corrector (BiasCorrector) – Initialized bias-correction object.

• logger (AID_BC.logger.Logger) – Logger used to report fitting diagnostics.

Returns:

Fitted bias corrector.

Return type:

BiasCorrector

Raises:

ValueError – If ERA5 and CMIP6 spatial or feature dimensions are inconsistent.

AID_BC.main.fit_bias_correctors_by_month(reference_train, biased_train, variable_names, method, ot_kwargs, logger)

Fit one bias corrector for each calendar month.

Parameters:

• reference_train (xarray.Dataset) – ERA5 training data covering all training years.

• biased_train (xarray.Dataset) – Historical CMIP6 training data covering all training years.

• variable_names (sequence of str) – Ordered variables used during fitting.

• method ({'qm', 'ot'}) – Bias-correction method.

• ot_kwargs (dict) – Keyword arguments passed to create_bias_corrector.

• logger (AID_BC.logger.Logger) – Logger used to report monthly fitting progress.

Returns:

Fitted corrector indexed by calendar month from 1 to 12.

Return type:

dict of int to BiasCorrector

Raises:

RuntimeError – If ERA5 or CMIP6 contains no training samples for a calendar month.

AID_BC.main.apply_bias_corrector(biased_apply, variable_names, corrector, logger)

Apply a fitted bias corrector to climate fields.

Parameters:

• biased_apply (xarray.Dataset) – CMIP6 data to correct.

• variable_names (sequence of str) – Ordered variables used during fitting and application.

• corrector (BiasCorrector) – Fitted bias corrector.

• logger (AID_BC.logger.Logger) – Logger used to report application diagnostics.

Returns:

Corrected climate variables on the original CMIP6 grid.

Return type:

xarray.Dataset

Raises:

RuntimeError – If the transformed matrix shape differs from the input matrix shape.

AID_BC.main.apply_bias_corrector_by_month(biased_apply, variable_names, correctors_by_month, logger)

Apply month-specific correctors and reassemble the time series.

Parameters:

• biased_apply (xarray.Dataset) – CMIP6 data for one application year.

• variable_names (sequence of str) – Ordered variables used during fitting and application.

• correctors_by_month (dict of int to BiasCorrector) – Fitted corrector for each calendar month.

• logger (AID_BC.logger.Logger) – Logger used to report application progress.

Returns:

Corrected data concatenated and sorted chronologically.

Return type:

xarray.Dataset

Raises:

RuntimeError – If a required monthly corrector is missing or no month is corrected.

AID_BC.main.save_corrected_year(corrected, output_dirs, year, variable_names, logger)

Save each corrected variable in its own yearly NetCDF file.

Parameters:

• corrected (xarray.Dataset) – Corrected climate variables.

• output_dirs (dict of str to str) – Output directory for each variable.

• year (int) – Year written to the samples_<year>.nc filename.

• variable_names (sequence of str) – Variables to save.

• logger (AID_BC.logger.Logger) – Logger used to report output paths.

Raises:

ValueError – If the corrected dataset has no time samples or lacks a requested variable.

AID_BC.main.main()

Run the complete CMIP6 bias-correction workflow.

The workflow validates arguments, loads and aligns training variables, fits one global or monthly QM/OT corrector, applies the fitted correction to each requested year, and saves each corrected variable in its configured output directory.

AID_BC.optimal_transport module

Solve entropy-regularized optimal transport using the Sinkhorn algorithm.

The solver computes a squared-Euclidean cost matrix, estimates dual potentials with Sinkhorn iterations, and constructs transport functions from the target dual potential.

class AID_BC.optimal_transport.TransportSolution(potentials: tuple[Array, Array], cost_matrix: Array, epsilon: float, reg_ot_cost: Array, threshold: float, converged: Array, num_iterations: Array)

Bases: NamedTuple

Store the output of the entropy-regularized transport solver.

The source and target dual potentials define the transport plan. The remaining fields describe the transport problem and the solver’s outcome.

potentials: tuple[jax.Array, jax.Array]

Alias for field number 0

cost_matrix: jax.Array

Alias for field number 1

epsilon: float

Alias for field number 2

reg_ot_cost: jax.Array

Alias for field number 3

threshold: float

Alias for field number 4

converged: jax.Array

Alias for field number 5

num_iterations: jax.Array

Alias for field number 6

property fu

Return the source dual potential.

property gv

Return the target dual potential.

property cost

Return the pairwise cost matrix.

property transport_plan

Compute the transport plan from the fitted dual potentials.

Returns:

Transport plan with shape (n_source, n_target).

Return type:

jax.Array

class AID_BC.optimal_transport.OptimalTransport(epsilon, sharding=None, num_iterations=100, threshold=0.001, eps_marginal=1e-12)

Bases: object

Solve entropy-regularized optimal transport using Sinkhorn iterations.

The solver uses the squared-Euclidean cost and uniform empirical marginals. The target dual potential can be used to construct a transport function for new source samples.

Parameters:

• epsilon (float) – Entropic regularization parameter. Must be strictly positive.

• sharding (jax.sharding.Sharding or None, optional) – Optional sharding constraint applied to the cost matrix.

• num_iterations (int, default=100) – Maximum number of Sinkhorn iterations.

• threshold (float, default=1e-3) – Convergence threshold based on changes in the dual potentials.

• eps_marginal (float, default=1e-12) – Small value added to marginal masses before taking their logarithms.

__init__(epsilon, sharding=None, num_iterations=100, threshold=0.001, eps_marginal=1e-12)

solve(source, target)

Solve the transport problem using named source and target arguments.

Parameters:

• source (jax.Array, shape (n_source, n_features)) – Source samples.

• target (jax.Array, shape (n_target, n_features)) – Target samples.

Returns:

Fitted transport solution.

Return type:

TransportSolution

transport_fn(potential, y, weights=None)

Build a vectorized transport function from the target dual potential.

The transport map is computed from the gradient of the entropic source potential: T(x) = x - 0.5 * grad(f_epsilon(x)). The factor 0.5 comes from using squared-Euclidean costs without a factor of one half.

Parameters:

• potential (jax.Array, shape (n_target,)) – Target dual potential obtained from Sinkhorn iterations.

• y (jax.Array, shape (n_target, n_features)) – Target samples associated with the dual potential.

• weights (jax.Array or None, optional) – Target marginal weights. Uniform weights are used when None.

Returns:

Function accepting an array of shape (n_samples, n_features) and returning the transported samples with the same shape.

Return type:

callable

transport_fn_direct(potential, y, weights=None)

Build a direct transport function using the target dual potential.

This implementation evaluates the transport weights using exponentials without log-sum-exp stabilization and may be numerically unstable.

Parameters:

• potential (jax.Array, shape (n_target,)) – Target dual potential obtained from Sinkhorn iterations.

• y (jax.Array, shape (n_target, n_features)) – Target samples associated with the dual potential.

• weights (jax.Array or None, optional) – Target marginal weights. Uniform weights are used when None.

Returns:

Function transporting one source sample of shape (n_features,) to a target-space sample with the same shape.

Return type:

callable

Raises:

ValueError – If the target dual potential is not one-dimensional or its length does not match the number of target samples.

AID_BC.preprocess module

Precompute ERA5 resized onto the native CMIP6 grid, once, and store it as yearly NetCDF files.

Run this once per (era5_root, cmip6_root, variable) combination, then point main.py at the output directory with –era5_on_cmip6_root.

Example

Run the preprocessing for 2m air temperature:

python preprocess.py –era5_root /data/kkingston/data/CMIP6/ERA5/data_6hourly_t2m –cmip6_root /data/kkingston/data/CMIP6/CMIP6_historical/data_6hourly_tas –variable VAR_2T –start_year 1980 –end_year 2014 –output_dir /data/kkingston/data/CMIP6/ERA5_on_CMIP6/data_6hourly_t2m

AID_BC.preprocess.parse_args()

Parse command-line arguments.

Returns:

Parsed arguments.

Return type:

argparse.Namespace

AID_BC.preprocess.build_path(root, year)

Build the yearly NetCDF path for a data directory.

Parameters:

• root (str or pathlib.Path) – Directory containing yearly NetCDF files.

• year (int) – Year used in the samples_<year>.nc filename.

Returns:

Full path to the yearly NetCDF file.

Return type:

str

AID_BC.preprocess.process_year(year, era5_root, cmip6_root, variable_name, output_dir, logger)

Resize one year of ERA5 data onto the native CMIP6 grid and save it.

The function loads one ERA5 file and one CMIP6 file, prepares the two datasets with AID_BC.dataset.ClimateDataset, resizes ERA5 onto the CMIP6 grid, and writes the resulting ERA5 field as a yearly NetCDF file.

Parameters:

• year (int) – Year to preprocess.

• era5_root (str or pathlib.Path) – Directory containing raw ERA5 yearly files named samples_<year>.nc.

• cmip6_root (str or pathlib.Path) – Directory containing CMIP6 yearly files that define the target native grid.

• variable_name (str) – Name of the climate variable to preprocess.

• output_dir (str or pathlib.Path) – Directory in which the precomputed ERA5-on-CMIP6 file is written.

• logger (AID_BC.logger.Logger) – Logger used to report progress, retries, and output paths.

Returns:

True if the original ERA5 latitude coordinate is descending, otherwise False.

Return type:

bool

AID_BC.preprocess.main()

Run ERA5-on-CMIP6 preprocessing for the requested year range.

The function parses command-line arguments, preprocesses every year, verifies that ERA5 latitude orientation remains consistent across years, and writes a metadata.json file describing the generated dataset.

Raises:

ValueError – If the ERA5 latitude ordering changes between processed years.

AID_BC.quantile_mapping module

Perform empirical Quantile Mapping for bias correction.

The method estimates the empirical distributions of reference and biased data independently for each feature. Biased values are corrected by mapping their empirical cumulative probabilities to the corresponding quantiles of the reference distribution using linear interpolation.

class AID_BC.quantile_mapping.QM

Bases: object

Basic empirical Quantile Mapping bias corrector.

Given a reference dataset Y0 and a biased dataset X0, this estimates the empirical cumulative distribution function (CDF) of each, feature by feature, and corrects biased values by matching their quantiles to those of the reference distribution.

__init__()

fit(Y0, X0)

Fit the empirical distributions of the reference and biased data.

Parameters:

• Y0 (numpy.ndarray) – Reference dataset with shape (n_samples, n_features), e.g. ERA5.

• X0 (numpy.ndarray) – Biased dataset with shape (n_samples, n_features), e.g. historical model data.

predict(X0)

Apply the fitted quantile mapping to new biased data.

For each feature, a value x is mapped to its empirical non-exceedance probability p under the fitted biased distribution, then p is mapped back to a value under the fitted reference distribution (linear interpolation, clipped to the reference sample’s range).

Parameters:

X0 (numpy.ndarray) – Data to correct with shape (n_samples, n_features).

Returns:

Corrected data with shape (n_samples, n_features).

Return type:

numpy.ndarray

AID_BC.version module

Version information for AID_BC.

AID_BC.version.get_version()

Return the version string.