utils.py

add_month_day_dims(daily_ts, monthly_ts, time_dim='time', spatial_dims=('lat', 'lon'))

Reshape daily and monthly data to have explicit month (M) and day (T) dimensions.

Here we assume maximum 31 days in a month, and invalid day entries will be padded with NaN.

Returns:

Name	Type	Description
daily_m	xr.DataArray - dims: (M, T, H, W)
monthly_m	xr.DataArray - dims: (M, H, W)
padded_days_mask	xr.DataArray - dims: (M, T=31), bool, True where day is padded
time_features	xr.DataArray - dims: (M, T, 2)

Source code in climanet/utils.py

def add_month_day_dims(
    daily_ts: xr.DataArray,  # (time, H, W) daily
    monthly_ts: xr.DataArray,  # (time, H, W) monthly
    time_dim: str = "time",
    spatial_dims: tuple[str, str] = ("lat", "lon"),
):
    """Reshape daily and monthly data to have explicit month (M) and day (T) dimensions.

    Here we assume maximum 31 days in a month, and invalid day entries will be
    padded with NaN.

    Returns
    -------
    daily_m : xr.DataArray - dims: (M, T, H, W)
    monthly_m : xr.DataArray - dims: (M, H, W)
    padded_days_mask : xr.DataArray - dims: (M, T=31), bool, True where day is padded
    time_features : xr.DataArray - dims: (M, T, 2)
    """
    # Month key as integer YYYYMM
    dkey = daily_ts[time_dim].dt.year * 100 + daily_ts[time_dim].dt.month
    mkey = monthly_ts[time_dim].dt.year * 100 + monthly_ts[time_dim].dt.month

    # Unique month keys preserving order
    _, idx = np.unique(dkey.values, return_index=True)
    month_keys = dkey.values[np.sort(idx)]

    # Add M (month key) and T (day of month) coordinates to daily data
    daily_indexed = (
        daily_ts.assign_coords(
            M=(time_dim, dkey.values), T=(time_dim, daily_ts[time_dim].dt.day.values)
        )
        .set_index({time_dim: ("M", "T")})
        .unstack(time_dim)
        .reindex(T=np.arange(1, 32), M=month_keys)
    )

    # fix chunks
    daily_indexed = daily_indexed.chunk(
        {
            "M": 1,
            "T": -1,
            "lat": 100,
            "lon": 100,
        }
    )

    # Force dim order: (M, T, H, W) (and keep any other non-time dims after M,T)
    other_dims = [d for d in daily_ts.dims if d != time_dim]  # e.g. ["H", "W"]
    daily_indexed = daily_indexed.transpose("M", "T", *other_dims)

    # Build padded days mask from daily_indexed (NaN locations)
    padded_days_mask = ~daily_indexed.notnull().any(dim=spatial_dims)

    # Preserve the original time coordinates for monthly data (M,)
    month_time = xr.DataArray(
        monthly_ts[time_dim].values,
        dims="M",
        coords={"M": mkey.values},
    )

    # Align monthly data to same month keys/order
    monthly_m = (
        monthly_ts.assign_coords(M=(time_dim, mkey.values))
        .swap_dims({time_dim: "M"})
        .drop_vars(time_dim)
        .sel(M=month_keys)
        .assign_coords(M=month_time.sel(M=month_keys).values)
    )

    # Build aligned datetime array (M,T)
    time_da = daily_ts[time_dim]

    # time_indexed is (M,T) with NaT for padded days
    time_indexed = (
        time_da.assign_coords(
            M=(time_dim, dkey.values), T=(time_dim, time_da.dt.day.values)
        )
        .set_index({time_dim: ("M", "T")})
        .unstack(time_dim)
        .reindex(T=np.arange(1, 32), M=month_keys)
    )

    # month-of_year (moy), day-of-year (doy) [and hour-of-day (hod) if applicable], fill NaT with 0 inplace
    # here we choose to use the tropical year length (365.2422 day, which we round to 365.24) as the
    # period to return to the position of the sun relative to the Earth
    moy_period = 12.0
    doy_period = 365.24
    hod_period = 24.0

    moy = time_indexed.dt.month.fillna(0)
    doy = time_indexed.dt.dayofyear.fillna(0)

    if "hour" in dir(time_indexed.dt):
        hod = time_indexed.dt.hour.fillna(0)
    else:
        hod = xr.zeros_like(doy)

    # create phase from day and hod
    moy_phase = 2 * np.pi * (moy - 1.0) / moy_period
    doy_phase = 2 * np.pi * doy / doy_period
    hod_phase = 2 * np.pi * hod / hod_period

    # Stack cyclic encodings into time_features (M,T,3)
    time_features = xr.concat(
        [moy_phase, doy_phase, hod_phase], dim="feature"
    ).transpose("M", "T", "feature")

    # fix chunks
    time_features = time_features.chunk(
        {
            "M": 1,
            "T": -1,
            "feature": -1,
        }
    )

    return daily_indexed, monthly_m, padded_days_mask, time_features

add_month_hour_dims(hourly_ts, monthly_ts, time_dim='time', spatial_dims=('lat', 'lon'))

Reshape hourly and monthly data to have explicit month (M) and hour (T) dimensions.

Here we assume maximum 31 days in a month with 24 hours per day = 744 hours maximum. Invalid hour entries will be padded with NaN.

Returns:

Name	Type	Description
hourly_m	xr.DataArray - dims: (M, T=744, H, W)
monthly_m	xr.DataArray - dims: (M, H, W)
padded_hours_mask	xr.DataArray - dims: (M, T=744), bool, True where hour is padded
time_features	xr.DataArray - dims: (M, T=744, 2)

Source code in climanet/utils.py

def add_month_hour_dims(
    hourly_ts: xr.DataArray,  # (time, H, W) hourly
    monthly_ts: xr.DataArray,  # (time, H, W) monthly
    time_dim: str = "time",
    spatial_dims: tuple[str, str] = ("lat", "lon"),
):
    """Reshape hourly and monthly data to have explicit month (M) and hour (T) dimensions.

    Here we assume maximum 31 days in a month with 24 hours per day = 744 hours maximum.
    Invalid hour entries will be padded with NaN.

    Returns
    -------
    hourly_m : xr.DataArray - dims: (M, T=744, H, W)
    monthly_m : xr.DataArray - dims: (M, H, W)
    padded_hours_mask : xr.DataArray - dims: (M, T=744), bool, True where hour is padded
    time_features : xr.DataArray - dims: (M, T=744, 2)
    """
    # Month key as integer YYYYMM
    hkey = hourly_ts[time_dim].dt.year * 100 + hourly_ts[time_dim].dt.month
    mkey = monthly_ts[time_dim].dt.year * 100 + monthly_ts[time_dim].dt.month

    # Unique month keys preserving order
    _, idx = np.unique(hkey.values, return_index=True)
    month_keys = hkey.values[np.sort(idx)]

    # Create hour-of-month coordinate (1-744)
    # hour_of_month = (day_of_month - 1) * 24 + hour_of_day + 1
    day_of_month = hourly_ts[time_dim].dt.day.values
    hour_of_day = hourly_ts[time_dim].dt.hour.values
    hour_of_month = (day_of_month - 1) * 24 + hour_of_day + 1

    # Add M (month key) and T (hour of month) coordinates to hourly data
    hourly_indexed = (
        hourly_ts.assign_coords(M=(time_dim, hkey.values), T=(time_dim, hour_of_month))
        .set_index({time_dim: ("M", "T")})
        .unstack(time_dim)
        .reindex(T=np.arange(1, 745), M=month_keys)  # 744 = 31 days * 24 hours
    )

    # fix chunks
    hourly_indexed = hourly_indexed.chunk(
        {
            "M": 1,
            "T": -1,
            "lat": 100,
            "lon": 100,
        }
    )

    # Force dim order: (M, T, H, W)
    other_dims = [d for d in hourly_ts.dims if d != time_dim]
    hourly_indexed = hourly_indexed.transpose("M", "T", *other_dims)

    # Build padded hours mask from hourly_indexed (NaN locations)
    padded_hours_mask = ~hourly_indexed.notnull().any(dim=spatial_dims)

    # Preserve the original time coordinates for monthly data (M,)
    month_time = xr.DataArray(
        monthly_ts[time_dim].values,
        dims="M",
        coords={"M": mkey.values},
    )

    # Align monthly data to same month keys/order
    monthly_m = (
        monthly_ts.assign_coords(M=(time_dim, mkey.values))
        .swap_dims({time_dim: "M"})
        .drop_vars(time_dim)
        .sel(M=month_keys)
        .assign_coords(M=month_time.sel(M=month_keys).values)
    )

    # Build aligned datetime array (M, T)
    time_da = hourly_ts[time_dim]

    # time_indexed is (M, T) with NaT for padded hours
    time_indexed = (
        time_da.assign_coords(M=(time_dim, hkey.values), T=(time_dim, hour_of_month))
        .set_index({time_dim: ("M", "T")})
        .unstack(time_dim)
        .reindex(T=np.arange(1, 745), M=month_keys)
    )

    # Determine month-of-year, day-of-year (doy) and hour-of-day (hod)
    moy_period = 12.0
    doy_period = 365.24
    hod_period = 24.0

    moy = time_indexed.dt.month.fillna(0)
    doy = time_indexed.dt.dayofyear.fillna(0)
    hod = time_indexed.dt.hour.fillna(0)

    # Create phase from month, day and hour
    moy_phase = 2 * np.pi * (moy - 1.0) / moy_period
    doy_phase = 2 * np.pi * doy / doy_period
    hod_phase = 2 * np.pi * hod / hod_period

    # Stack cyclic encodings into time_features (M, T, 3)
    time_features = xr.concat(
        [moy_phase, doy_phase, hod_phase], dim="feature"
    ).transpose("M", "T", "feature")

    # fix chunks
    time_features = time_features.chunk(
        {
            "M": 1,
            "T": -1,
            "feature": -1,
        }
    )

    return hourly_indexed, monthly_m, padded_hours_mask, time_features

calc_stats(arr, mean_axis=0)

Calculate mean and std along the specified axis, ignoring NaNs.

Args: arr: Input array containing NaNs to ignore. shape is (M, T, H, W) mean_axis: Axis along which to compute mean and std (default is 0 for month) Returns: mean: Mean values along the specified axis, shape (M,) std: Standard deviation along the specified axis, shape (M,)

Source code in climanet/utils.py

def calc_stats(arr: np.ndarray, mean_axis: int = 0) -> tuple[np.ndarray, np.ndarray]:
    """Calculate mean and std along the specified axis, ignoring NaNs.

    Args:
        arr: Input array containing NaNs to ignore. shape is (M, T, H, W)
        mean_axis: Axis along which to compute mean and std (default is 0 for month)
    Returns:
        mean: Mean values along the specified axis, shape (M,)
        std: Standard deviation along the specified axis, shape (M,)
    """
    axes_to_reduce = tuple(i for i in range(arr.ndim) if i != mean_axis)

    mean = np.nanmean(arr, axis=axes_to_reduce)  # shape: (M,)
    std = np.nanstd(arr, axis=axes_to_reduce)  # shape: (M,)
    return mean, std

compute_masked_loss(pred, target, land_mask)

Compute L1 loss masked to ocean pixels only.

Source code in climanet/utils.py

def compute_masked_loss(
    pred: torch.Tensor, target: torch.Tensor, land_mask: torch.Tensor
) -> torch.Tensor:
    """Compute L1 loss masked to ocean pixels only."""
    ocean = (~land_mask).to(pred.device).unsqueeze(1)

    # Mask for valid (non-NaN) target values
    valid = ~torch.isnan(target)
    target = torch.nan_to_num(target, nan=0.0)

    mask = ocean & valid
    loss = torch.nn.functional.l1_loss(pred, target, reduction="none")
    loss = loss * mask

    num = loss.sum(dim=(-2, -1))
    denom = mask.sum(dim=(-2, -1)).clamp_min(1)

    return (num / denom).mean()

configure_compute_resources(model, device, compute_threads, dataloader_num_workers)

Configure model for multi-GPU and set CPU thread usage for compute (training or prediction).

Args: model: the PyTorch model to configure device: device to run on ("cpu" or "cuda") compute_threads: number of threads to use for compute when device is CPU. If None, it will be set automatically. dataloader_num_workers: how many subprocesses to use for data loading. See torch DataLoader docs for details. Returns: The model, potentially wrapped in DataParallel if using multiple GPUs.

Source code in climanet/utils.py

def configure_compute_resources(
    model: torch.nn.Module,
    device: str,
    compute_threads: int,
    dataloader_num_workers: int,
) -> torch.nn.Module:
    """Configure model for multi-GPU and set CPU thread usage for compute (training or prediction).

    Args:
        model: the PyTorch model to configure
        device: device to run on ("cpu" or "cuda")
        compute_threads: number of threads to use for compute when device is CPU.
            If None, it will be set automatically.
        dataloader_num_workers: how many subprocesses to use for data loading.
            See torch DataLoader docs for details.
    Returns:
        The model, potentially wrapped in DataParallel if using multiple GPUs.
    """
    if device == "cpu":
        if compute_threads is None:
            total_cpus = psutil.cpu_count()
            # keep 1 for main thread
            compute_threads = max(1, total_cpus - dataloader_num_workers - 1)
        torch.set_num_threads(compute_threads)
    elif device == "cuda":
        num_gpus = torch.cuda.device_count()
        if num_gpus > 1:
            model = torch.nn.DataParallel(model)
    return model

data_preparation(input_data, monthly_data, time_dim='time', run_dir='.', calculate_residuals=True, is_hourly=False, save_to_zarr=False)

Prepare the data for training.

Args: input_data (xr.DataArray): The input data (daily or hourly). monthly_data (xr.DataArray): The monthly data. time_dim (str): The name of the time dimension in the data arrays. run_dir (str): Directory to save the preprocessed data. calculate_residuals (bool): Whether to calculate residuals between input and monthly data. is_hourly (bool): Whether the input data is hourly (True) or daily (False). save_to_zarr (bool): Whether to save the preprocessed data to zarr files. Returns: tuple: A tuple containing the following xarray.DataArray objects: - input_da: The reshaped input data with dimensions (M, T, H, W). - input_da_nan_mask: A boolean mask indicating NaN locations in the input data. - monthly_da: The reshaped monthly data with dimensions (M, H, W). - padded_days_mask: A boolean mask indicating padded days/hours in the input data. - time_features: A DataArray containing cyclic time features (month, day, hour).

Source code in climanet/utils.py

def data_preparation(
    input_data: xr.DataArray,
    monthly_data: xr.DataArray,
    time_dim="time",
    run_dir=".",
    calculate_residuals=True,
    is_hourly=False,
    save_to_zarr=False,
):
    """Prepare the data for training.

    Args:
        input_data (xr.DataArray): The input data (daily or hourly).
        monthly_data (xr.DataArray): The monthly data.
        time_dim (str): The name of the time dimension in the data arrays.
        run_dir (str): Directory to save the preprocessed data.
        calculate_residuals (bool): Whether to calculate residuals between input and monthly data.
        is_hourly (bool): Whether the input data is hourly (True) or daily (False).
        save_to_zarr (bool): Whether to save the preprocessed data to zarr files.
    Returns:
        tuple: A tuple containing the following xarray.DataArray objects:
            - input_da: The reshaped input data with dimensions (M, T, H, W).
            - input_da_nan_mask: A boolean mask indicating NaN locations in the input data.
            - monthly_da: The reshaped monthly data with dimensions (M, H, W).
            - padded_days_mask: A boolean mask indicating padded days/hours in the input data.
            - time_features: A DataArray containing cyclic time features (month, day, hour).

    """
    if time_dim not in input_data.dims or time_dim not in monthly_data.dims:
        raise ValueError(f"Time dimension '{time_dim}' not found in input data")

    var_name = input_data.name
    if not var_name:
        raise ValueError("Input data must have a name (variable name)")

    if calculate_residuals:
        input_data_averaged = input_data.resample({time_dim: "MS"}).mean(skipna=True)
        input_data_averaged[time_dim] = monthly_data[time_dim]
        monthly_data_res = monthly_data - input_data_averaged
    else:
        monthly_data_res = monthly_data

    if is_hourly:
        # hours_per_day == 24
        # Reshape daily → (M, T=31*24, H, W), monthly → (M, H, W),
        # and get padded_days_mask → (M, T=31*24)
        input_da, monthly_da, padded_days_mask, time_features = add_month_hour_dims(
            input_data, monthly_data_res, time_dim=time_dim
        )
    else:
        # Reshape daily → (M, T=31, H, W), monthly → (M, H, W),
        # and get padded_days_mask → (M, T=31)
        input_da, monthly_da, padded_days_mask, time_features = add_month_day_dims(
            input_data, monthly_data_res, time_dim=time_dim
        )

    # Precompute the NaN mask before filling NaNs
    # input_da_nan_mask: True where NaN (i.e. missing ocean data, not land)
    input_da_nan_mask = input_da.isnull()

    # NaNs will be filled with 0 in-place
    input_da = input_da.fillna(0.0).astype("float32")

    monthly_da = monthly_da.astype("float32")

    # rechunk data
    input_da = input_da.chunk({"M": 1, "T": -1, "lat": 100, "lon": 100})
    input_da_nan_mask = input_da_nan_mask.chunk({"M": 1, "T": -1, "lat": 100, "lon": 100})
    monthly_da = monthly_da.chunk({"M": 1, "lat": 100, "lon": 100})
    padded_days_mask = padded_days_mask.chunk({"M": 1})
    time_features = time_features.chunk({"M": 1})

    # set names
    input_da_nan_mask.name = var_name
    monthly_da.name = var_name
    time_features.name = var_name
    padded_days_mask.name = var_name

    # compression for boolean masks
    encoding_input_da_nan_mask = {
        input_da_nan_mask.name: {
            "dtype": "bool",
            "compressor": numcodecs.Blosc(
                cname="zstd",
                clevel=3,
                shuffle=numcodecs.Blosc.BITSHUFFLE,
            ),
        }
    }

    encoding_padded_days_mask = {
        padded_days_mask.name: {
            "dtype": "bool",
            "compressor": numcodecs.Blosc(
                cname="zstd",
                clevel=3,
                shuffle=numcodecs.Blosc.BITSHUFFLE,
            ),
        }
    }

    if save_to_zarr:
        # Create the run directory if it doesn't exist
        data_path = Path(run_dir).resolve()
        data_path.mkdir(parents=True, exist_ok=True)

        input_da_path = data_path / "input_da.zarr"
        input_da_nan_mask_path = data_path / "input_da_nan_mask.zarr"
        monthly_da_path = data_path / "monthly_da.zarr"
        padded_days_mask_path = data_path / "padded_days_mask.zarr"
        time_features_path = data_path / "time_features.zarr"

        # these will be saved as xr.Dataset
        input_da.to_zarr(input_da_path, mode="w", zarr_format=2, consolidated=True)
        input_da_nan_mask.to_zarr(
            input_da_nan_mask_path,
            mode="w",
            encoding=encoding_input_da_nan_mask,
            zarr_format=2,
            consolidated=True,
        )
        monthly_da.to_zarr(monthly_da_path, mode="w", zarr_format=2, consolidated=True)
        padded_days_mask.to_zarr(
            padded_days_mask_path,
            mode="w",
            encoding=encoding_padded_days_mask,
            zarr_format=2,
            consolidated=True,
        )
        time_features.to_zarr(
            time_features_path, mode="w", zarr_format=2, consolidated=True
        )

    return input_da, input_da_nan_mask, monthly_da, padded_days_mask, time_features

data_split(data_folder, filename_pattern='*_hr_ERA5dc_masked_tos.nc', train_range=(2018, 2020), validation_range=(2021, 2021), test_range=(2022, 2022))

Split the data into training, validation, and test sets based on the provided year ranges.

Source code in climanet/utils.py

def data_split(
    data_folder,
    filename_pattern="*_hr_ERA5dc_masked_tos.nc",
    train_range=(2018, 2020),
    validation_range=(2021, 2021),
    test_range=(2022, 2022),
):
    """
    Split the data into training, validation, and test sets based on the provided year ranges.
    """
    data_folder = Path(data_folder)

    splits = {
        "train": [],
        "validation": [],
        "test": [],
    }

    for file in data_folder.rglob(filename_pattern):
        year = int(file.stem[:4])

        if train_range[0] <= year <= train_range[1]:
            splits["train"].append(file)
        if validation_range[0] <= year <= validation_range[1]:
            splits["validation"].append(file)
        if test_range[0] <= year <= test_range[1]:
            splits["test"].append(file)

    for lst in splits.values():
        lst.sort()

    return splits

load_model(model_path, device)

Helper function to load a model from a checkpoint.

Source code in climanet/utils.py

def load_model(model_path: str, device: str):
    """Helper function to load a model from a checkpoint."""
    checkpoint = torch.load(model_path, map_location=device, weights_only=False)
    model = SpatioTemporalModel(**checkpoint["model_config"])
    model.load_state_dict(checkpoint["model_state_dict"])
    return model.to(device)

plot_histograms(target, predictions, label='SST K', legend_labels=('Target', 'Prediction'), bins=30)

Plot histograms of target and predictions in the same figure for comparison.

Source code in climanet/utils.py

def plot_histograms(
    target, predictions, label="SST K", legend_labels=("Target", "Prediction"), bins=30
):
    """Plot histograms of target and predictions in the same figure for comparison."""
    _, axs = plt.subplots(
        nrows=len(target.time),
        ncols=1,
        figsize=(8, 4 * len(target.time)),
        constrained_layout=True,
        squeeze=False,
    )

    for t in range(len(target.time)):
        target_t = target.isel(time=t)
        pred_t = predictions.isel(time=t)

        # Target histogram
        axs[t, 0].hist(
            target_t.values.flatten(), bins=bins, alpha=0.7, color="blue", density=True
        )
        axs[t, 0].set_xlabel(label)
        axs[t, 0].set_ylabel("Probability Density")
        axs[t, 0].grid(True, alpha=0.3)

        # Prediction histogram (overlaid)
        axs[t, 0].hist(
            pred_t.values.flatten(), bins=bins, alpha=0.7, color="orange", density=True
        )
        axs[t, 0].legend(legend_labels)
        axs[t, 0].set_title(
            f"Histogram {legend_labels[0]} vs {legend_labels[1]}, month={target.time.dt.strftime('%Y-%m-%d').values[t]}"
        )

    plt.show()

plot_loss(run_dir, list_loss_var, unit='K', figsize=(10, 5))

Plot training and validation loss from TensorBoard logs.

Args: run_dir (str | Path): Directory containing TensorBoard logs. list_loss_var (list[str]): List of loss variable names to plot. unit (str, optional): Unit of the loss values. Defaults to "K". figsize (Tuple[int, int], optional): Size of the figure. Defaults to (10, 5).

Source code in climanet/utils.py

def plot_loss(
    run_dir: str | Path,
    list_loss_var: list[str],
    unit: str = "K",
    figsize: tuple[int, int] = (10, 5),
):
    """Plot training and validation loss from TensorBoard logs.

    Args:
        run_dir (str | Path): Directory containing TensorBoard logs.
        list_loss_var (list[str]): List of loss variable names to plot.
        unit (str, optional): Unit of the loss values. Defaults to "K".
        figsize (Tuple[int, int], optional): Size of the figure. Defaults to (10, 5).
    """
    # Load saved training status with tbparse.SummaryReader
    reader = SummaryReader(run_dir)

    # plot training and validation loss
    plt.figure(figsize=figsize)
    for loss_var in list_loss_var:
        loss = reader.scalars[reader.scalars["tag"] == loss_var]
        plt.plot(loss["step"], loss["value"], label=loss_var)
    plt.xlabel("Step")
    plt.ylabel(f"Average loss per epoch ({unit})")
    plt.legend()

plot_nobs_vs_err(nobs, err_baseline, err_predictions)

Plot number of observations vs error for each month.

The three inputs are expected to be xarray DataArrays with dimensions (time, lat, lon). They should share the same spatial and temporal coordinates.

Args: nobs (xr.DataArray): Number of observations per grid cell per month. Dimensions: (time, lat, lon) err_baseline (xr.DataArray): Baseline error per grid cell per month. Dimensions: (time, lat, lon) err_predictions (xr.DataArray): Prediction error per grid cell per month. Dimensions: (time, lat, lon)

Source code in climanet/utils.py

def plot_nobs_vs_err(
    nobs: xr.DataArray, err_baseline: xr.DataArray, err_predictions: xr.DataArray
):
    """Plot number of observations vs error for each month.

    The three inputs are expected to be xarray DataArrays with dimensions (time, lat, lon).
    They should share the same spatial and temporal coordinates.

    Args:
        nobs (xr.DataArray): Number of observations per grid cell per month. Dimensions: (time, lat, lon)
        err_baseline (xr.DataArray): Baseline error per grid cell per month. Dimensions: (time, lat, lon)
        err_predictions (xr.DataArray): Prediction error per grid cell per month. Dimensions: (time, lat, lon)
    """
    _, axes = plt.subplots(nobs.sizes["time"], 1, figsize=(5 * nobs.sizes["time"], 8))
    if nobs.sizes["time"] == 1:
        axes = [axes]

    for i, ax in enumerate(axes):
        ax.set_title(f"Month = {err_baseline.time.dt.strftime('%Y-%m-%d').values[i]}")

        # Get unique number of observations for this month, ignoring NaNs and zeros
        n_obs_unique = np.unique(nobs.isel(time=i).values)
        n_obs_unique = n_obs_unique[(~np.isnan(n_obs_unique)) & (n_obs_unique > 0)]
        n_obs_unique = n_obs_unique.astype(int)

        err_by_n_obs_baseline = []
        err_by_n_obs_predictions = []

        for id_obs in n_obs_unique:
            # Baseline error
            err_arr = (
                err_baseline.isel(time=i)
                .where(nobs.isel(time=i) == id_obs)
                .values.flatten()
            )
            err_arr = err_arr[~np.isnan(err_arr)]
            if len(err_arr) == 0:
                err_arr = np.array([np.nan])
            err_by_n_obs_baseline.append(np.abs(err_arr))

            # Prediction error
            err_arr = (
                err_predictions.isel(time=i)
                .where(nobs.isel(time=i) == id_obs)
                .values.flatten()
            )
            err_arr = err_arr[~np.isnan(err_arr)]
            if len(err_arr) == 0:
                err_arr = np.array([np.nan])
            err_by_n_obs_predictions.append(np.abs(err_arr))

        h1 = ax.violinplot(
            err_by_n_obs_baseline,
            positions=n_obs_unique,
            showmedians=True,
            showextrema=True,
            points=500,
        )
        h2 = ax.violinplot(
            err_by_n_obs_predictions,
            positions=n_obs_unique,
            showmedians=True,
            showextrema=True,
            points=500,
        )

        # Style: thinner outlines + less prominent extrema
        for body in h1["bodies"]:
            body.set_facecolor("tab:blue")
            body.set_edgecolor("tab:blue")
            body.set_alpha(0.45)
            body.set_linewidth(0.5)

        for body in h2["bodies"]:
            body.set_facecolor("tab:orange")
            body.set_edgecolor("tab:orange")
            body.set_alpha(0.45)
            body.set_linewidth(0.5)

        for h in (h1, h2):
            h["cmedians"].set_linewidth(0.9)
            h["cmedians"].set_alpha(0.9)

            h["cbars"].set_linewidth(0.35)
            h["cbars"].set_alpha(0.2)
            h["cmins"].set_linewidth(0.35)
            h["cmins"].set_alpha(0.2)
            h["cmaxes"].set_linewidth(0.35)
            h["cmaxes"].set_alpha(0.2)

        ax.set_xlabel("Number of Observations")
        ax.set_ylabel("Symmetric log-scaled Absolute Error (K)")

        # Non-linear y-axis: keeps detail near 0 and compresses larger values.
        ax.set_yscale("symlog", linthresh=0.05, linscale=0.8, base=10)

        # Show major ticks as plain decimals instead of scientific/log notation.
        ax.yaxis.set_major_locator(
            mticker.SymmetricalLogLocator(base=10, linthresh=0.05)
        )
        ax.yaxis.set_major_formatter(
            mticker.FuncFormatter(lambda y, _: f"{y:.3f}".rstrip("0").rstrip("."))
        )
        ax.yaxis.set_minor_formatter(mticker.NullFormatter())

        ax.legend(
            [h1["bodies"][0], h2["bodies"][0]],
            ["Baseline", "Prediction"],
            loc="upper right",
        )

    plt.tight_layout()

pred_to_numpy(pred, orig_H=None, orig_W=None, land_mask=None)

pred: (B, M, H_pad,W_pad) or (B, H, W) torch tensor orig_H/W: original sizes before padding (optional) land_mask: (B, H_pad,W_pad) or (B, H,W) bool; if given, land will be set to NaN returns: (H,W) numpy array

Source code in climanet/utils.py

def pred_to_numpy(pred, orig_H=None, orig_W=None, land_mask=None):
    """
    pred: (B, M, H_pad,W_pad) or (B, H, W) torch tensor
    orig_H/W: original sizes before padding (optional)
    land_mask: (B, H_pad,W_pad) or (B, H,W) bool; if given, land will be set to NaN
    returns: (H,W) numpy array
    """
    # crop to original size if provided
    if orig_H is not None and orig_W is not None:
        pred = pred[..., :orig_H, :orig_W]
        if land_mask is not None:
            land_mask = land_mask[..., :orig_H, :orig_W]

    # set land to NaN (broadcast mask across batch)
    if land_mask is not None:
        pred = pred.clone().to(torch.float32)
        land_mask = land_mask.bool()
        land_mask = land_mask.unsqueeze(1)  # (B, H,W) -> (B, 1, H, W) for broadcasting
        pred = torch.where(land_mask, torch.full_like(pred, float("nan")), pred)

    return pred.detach().cpu().numpy()

read_st_data(data_path='.', var_name='tos')

Read preprocessed spatio-temporal data from zarr files. Args: data_path (str): Path to the directory containing the zarr files. var_name (str): Name of the variable to read from the zarr files.

Returns: tuple: A tuple containing the following xarray.DataArray objects: - input_da - input_da_nan_mask - monthly_da - padded_days_mask - time_features

Source code in climanet/utils.py

def read_st_data(data_path=".", var_name="tos"):
    """Read preprocessed spatio-temporal data from zarr files.
    Args:
        data_path (str): Path to the directory containing the zarr files.
        var_name (str): Name of the variable to read from the zarr files.

    Returns:
        tuple: A tuple containing the following xarray.DataArray objects:
            - input_da
            - input_da_nan_mask
            - monthly_da
            - padded_days_mask
            - time_features
    """

    # make filenames
    data_path = Path(data_path).resolve()

    input_da_path = data_path / "input_da.zarr"
    input_da_nan_mask_path = data_path / "input_da_nan_mask.zarr"
    monthly_da_path = data_path / "monthly_da.zarr"
    padded_days_mask_path = data_path / "padded_days_mask.zarr"
    time_features_path = data_path / "time_features.zarr"

    # Check if the zarr files already exist, if so, open them and return the datasets
    input_da = xr.open_zarr(input_da_path)[var_name]
    input_da_nan_mask = xr.open_zarr(input_da_nan_mask_path)[var_name]
    monthly_da = xr.open_zarr(monthly_da_path)[var_name]
    padded_days_mask = xr.open_zarr(padded_days_mask_path)[var_name]
    time_features = xr.open_zarr(time_features_path)[var_name]

    # if one of the datasets is None, we need to compute them
    return input_da, input_da_nan_mask, monthly_da, padded_days_mask, time_features

regrid_to_boundary_centered_grid(da, roll=False)

Interpolates a DataArray from its current center-based grid onto a new grid whose coordinates are derived from user-specified boundaries.

Includes robust handling for 0-360 vs -180-180 longitude domains.

Assumes dimensions are named 'lat' and 'lon'.

Source code in climanet/utils.py

def regrid_to_boundary_centered_grid(da: xr.DataArray, roll=False) -> xr.DataArray:
    """
    Interpolates a DataArray from its current center-based grid onto a new
    grid whose coordinates are derived from user-specified boundaries.

    Includes robust handling for 0-360 vs -180-180 longitude domains.

    Assumes dimensions are named 'lat' and 'lon'.
    """
    print("Starting regridding process...")

    # --- 0. Longitude Domain Check and Correction ---

    input_lon = da["longitude"]

    # Check if roll for 0-360 to -180-180 is requested
    if roll:
        print("Applying cyclic roll to -180 to 180...")

        # Calculate the index closest to 180 degrees
        lon_diff = np.abs(input_lon - 180.0)
        # We need to roll such that the 180-degree line is moved to the edge
        # and the new array starts near -180
        roll_amount = (
            int(lon_diff.argmin().item() + (input_lon.size / 2)) % input_lon.size
        )

        # Roll the DataArray and its coordinates
        da = da.roll(longitude=roll_amount, roll_coords=True)

        # Correct the longitude coordinate values: shift values > 180 down by 360
        new_lon_coords = da["longitude"].where(
            da["longitude"] <= 180, da["longitude"] - 360
        )

        # Assign the corrected and sorted coordinates
        da = da.assign_coords(longitude=new_lon_coords).sortby("longitude")
        print(
            f"Longitudes adjusted. New range: {da['longitude'].min().item():.2f} "
            f"to {da['longitude'].max().item():.2f}"
        )

    # --- 1. Define Target Grid Boundaries ---

    # Target latitude boundaries: -90.0 up to 90.0 in 0.25 degree steps
    # (721 points)
    lat_bnds = np.linspace(-90.0, 90.0, 721)

    # Target longitude boundaries: -180.0 up to 180.0 in 0.25 degree steps
    # (1441 points)
    lon_bnds = np.linspace(-180.0, 180.0, 1441)

    # --- 2. Calculate New Grid Centers (Coordinates) ---

    # New latitude centers are the average of consecutive boundaries
    # (720 points)
    new_lats = (lat_bnds[:-1] + lat_bnds[1:]) / 2.0

    # New longitude centers are the average of consecutive boundaries
    # (1440 points)
    new_lons = (lon_bnds[:-1] + lon_bnds[1:]) / 2.0

    # --- 3. Interpolate the Data ---

    # Use linear interpolation (suitable for gappy data) to map data onto the
    # new centers. xarray handles the NaNs automatically.
    da_regridded = da.interp(latitude=new_lats, longitude=new_lons, method="linear")

    print(f"Regridding complete. New dimensions: {da_regridded.dims}")
    return da_regridded

save_model(model, optimizer, run_dir, filename='best_model.pth', verbose=True)

Save model state and config to disk.

Source code in climanet/utils.py

def save_model(
    model: torch.nn.Module,
    optimizer: torch.optim.Optimizer,
    run_dir: str,
    filename="best_model.pth",
    verbose: bool = True,
) -> None:
    """Save model state and config to disk."""
    Path(run_dir).mkdir(parents=True, exist_ok=True)
    model_path = Path(run_dir) / filename
    model = model.module if hasattr(model, "module") else model
    torch.save(
        {
            "model_state_dict": model.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
            "model_config": model.config,
        },
        model_path,
    )
    if verbose:
        print(f"Model saved to {model_path}")

setup_logging(log_dir)

Set up TensorBoard logging directory and writer.

Source code in climanet/utils.py

def setup_logging(log_dir: str) -> SummaryWriter:
    """Set up TensorBoard logging directory and writer."""
    Path(log_dir).mkdir(parents=True, exist_ok=True)
    timestamp_utc = time.strftime("%Y%m%dT%H%M%S", time.gmtime())
    return SummaryWriter(log_dir, filename_suffix=f"_UTC{timestamp_utc}")