Source code for pycif.plugins.datastreams.fields.lmdz_inicond_reg.read

from __future__ import annotations

import datetime

import numpy as np
import xarray as xr




def read(
    self,
    name: str,
    varnames: str,
    dates: list[tuple[datetime.datetime, datetime.datetime]],
    files: list[str],
    tracer: object | None = None,
    **kwargs,
) -> xr.DataArray:
    varnames = varnames if varnames else name

    if len(files) > 1:
        raise ValueError("multiple files provided for initial conditions")
    (path,) = files

    with xr.open_dataset(path) as ds:
        data = ds[varnames].values

    # Backward compatibility with old LMDZ files
    if data.shape[0] == 1:
        data = data[0, ...]
    if data.shape[2] == self.domain.nlon + 1:
        data = data[:, :, :-1]

    if data.shape != (self.domain.nlev, self.domain.nlat, self.domain.nlon):
        raise ValueError(
            f"unexpected initial conditions data shape {data.shape}, "
            f"should be ({self.domain.nlev}, {self.domain.nlat}, {self.domain.nlon})"
        )

    xmod = xr.DataArray(
        data[np.newaxis, :, :, :],
        coords={"time": [dates[0][0]]},
        dims=("time", "lev", "lat", "lon"),
    )

    return xmod