from __future__ import annotations
import xarray as xr
_VALID_NBP = (40, 80, 160, 320)
def _get_ncell(nbp: int) -> int:
# Base icosahedral grid parameters
n_grid_face = 10
n_grid_edge = 20
n_grid_vertex = 12
ncell = (
n_grid_face * (nbp - 1) * (nbp - 1) + n_grid_edge * (nbp - 1) + n_grid_vertex
)
return ncell
[docs]
def get_nbp(self) -> int:
"""Get nbp number from domain
Parameters
----------
domain : Plugin
LMDZ domain plugin
Returns
-------
int
nbp number
"""
ncell = self.nlon
for nbp in _VALID_NBP:
if ncell == _get_ncell(nbp):
return nbp
raise ValueError(f"Invalid number of cells in domain: {ncell}")
[docs]
def get_time_splits(self) -> tuple[int, int]:
"""Get time step splits for the domain
Returns
-------
tuple[int, int]
Dynamics time step split, physics time step split
"""
nbp = get_nbp(self)
if nbp == 80:
return 30, 15
else:
raise NotImplementedError(f"Time split not implemented for nbp{nbp} resolution")
[docs]
def get_domain_coords(self, vertical: bool = False) -> dict[str, xr.DataArray]:
"""Get the formatted domain coordinates for writing NetCDF files
Parameters
----------
vertical : bool, optional
Include vertical coordinates, by default False
Returns
-------
dict[str, DataArray]
Coordinates
"""
lat = self.zlat[0, :]
lon = self.zlon[0, :]
lat_vertices = self.zlatc.T
lon_vertices = self.zlonc.T
coords = {
# fmt: off
"lat": xr.DataArray(lat, dims=["cell"], attrs={
"standard_name": "latitude",
"long_name": "latitude",
"units": "degrees_north",
"bounds": "lat_bnds"
}),
"lon": xr.DataArray(lon, dims=["cell"], attrs={
"standard_name": "longitude",
"long_name": "longitude",
"units": "degrees_east",
"bounds": "lon_bnds"
}),
"lat_bnds": xr.DataArray(lat_vertices, dims=["cell", "vertex"], attrs={
"standard_name": "latitude_bounds",
"long_name": "latitude bounds",
"units": "degrees_north"
}),
"lon_bnds": xr.DataArray(lon_vertices, dims=["cell", "vertex"], attrs={
"standard_name": "longitude_bounds",
"long_name": "longitude bounds",
"units": "degrees_east"
}),
# fmt: on
}
if vertical:
coords.update(
{
# fmt: off
"ap": xr.DataArray(self.sigma_a, dims=["lev"], attrs={
"standard_name": "atmosphere_hybrid_sigma_pressure_coordinate_ap",
"long_name": "Ap coefficient at layer interface",
"units": "Pa"
}),
"bp": xr.DataArray(self.sigma_b, dims=["lev"], attrs={
"standard_name": "atmosphere_hybrid_sigma_pressure_coordinate_bp",
"long_name": "B coefficient at layer interface",
"units": "1"
}),
# fmt: on
}
)
return coords
