Clean up the time encoding of a dataset to use the ilamb default

ilambify_time.py

from pathlib import Path
from typing import Union

import numpy as np
import xarray as xr
from cftime import DatetimeNoLeap

COMPRESS_FILES = True


def fix_simple(filename: Path) -> xr.Dataset:
    """Fix the time records of the given file and add bounds."""
    ds = xr.open_dataset(filename)
    ds["time"] = [DatetimeNoLeap(t.dt.year, t.dt.month, 15) for t in ds["time"]]
    ds["time_bnds"] = xr.DataArray(
        data=np.array(
            [
                [DatetimeNoLeap(t.dt.year, t.dt.month, 1) for t in ds["time"]],
                [
                    DatetimeNoLeap(
                        t.dt.year if t.dt.month < 12 else (t.dt.year + 1),
                        (t.dt.month + 1) if t.dt.month < 12 else 1,
                        1,
                    )
                    for t in ds["time"]
                ],
            ]
        ).T,
        dims=("time", "nb"),
    )
    return ds


def fix_unit(unit: str) -> str:
    unit = unit.replace("$^{", "")
    unit = unit.replace("}$", "")
    return unit


def ilambify_time(
    filename: Union[str, Path], destination_path: Union[str, Path]
) -> None:
    # convert paths to Path objects if needed
    if not isinstance(filename, Path):
        filename = Path(filename)
    if not isinstance(destination_path, Path):
        destination_path = Path(destination_path)

    # first we need to peek at what is happening in time to decide how to fix it
    ds = xr.open_dataset(filename, decode_times=False)
    t = ds["time"]
    if "months" in t.attrs["units"]:
        raise NotImplementedError
    else:
        ds = fix_simple(filename)

    # ok ok we will fix some other things too
    for var, da in ds.items():
        if da.ndim > 2 and "units" in da.attrs:
            ds[var].attrs["units"] = fix_unit(da.attrs["units"])

    # now write the new file
    if not destination_path.exists():
        destination_path.mkdir(parents=True, exist_ok=True)
    ds.to_netcdf(
        destination_path / filename,
        encoding={
            "time": {"units": "days since 1850-01-01", "bounds": "time_bnds"},
            "time_bnds": {"units": "days since 1850-01-01"},
            list(ds)[0]: {"zlib": COMPRESS_FILES},
        },
    )


if __name__ == "__main__":
    ilambify_time("CLASSIC_S3_nbp.nc", destination_path="./junk")