import copy
import itertools
from pathlib import Path
import numpy as np
from .....utils.classes.domains import Domain
from .utils import GribDataset, get_grid_type, get_jscan
from .....utils.check.errclass import CifFileNotFoundError, CifValueError
[docs]
def make_regular_grid(lat, lon, jscan):
"""Build center/corner lon-lat grids for a regular lat-lon GRIB grid.
Args:
lat: 1D array of latitude centers, as read from the GRIB file.
lon: 1D array of longitude centers, as read from the GRIB file.
jscan: Value of the GRIB ``jScansPositively`` key; if ``0``, the
latitude axis is flipped to increasing order.
Returns:
tuple: ``(zlonc, zlatc, zlon, zlat, nlon, nlat)`` where ``zlon``/
``zlat`` are the 2D meshgrid of cell centers, ``zlonc``/``zlatc``
the 2D meshgrid of cell corners, and ``nlon``/``nlat`` the grid
dimensions.
"""
dx = np.unique(np.round(np.diff(lon), 6))[0]
dy = np.abs(np.unique(np.round(np.diff(lat), 6))[0])
# Flip latitudes if jscan = 0
if jscan == 0:
lat2 = np.flip(lat)
lat = copy.deepcopy(lat2)
# Define lon/lat variables
nlon = lon.size
nlat = lat.size
zlon, zlat = np.meshgrid(lon, lat)
latc = np.append(lat - dy / 2, lat[-1] + dy / 2)
lonc = np.append(lon - dx / 2, lon[-1] + dx / 2)
zlonc, zlatc = np.meshgrid(lonc, latc)
return zlonc, zlatc, zlon, zlat, nlon, nlat
[docs]
def make_gaussian_grid(lat, lon):
"""Build an unstructured cell-vertex grid for a reduced-Gaussian GRIB grid.
Groups consecutive grid points into blocks of constant latitude (each
block corresponding to a Gaussian latitude ring), and computes the
4-vertex lon/lat corners of each cell from the spacing within its ring
and the mid-points to the neighboring rings (poles are extended to
+/-90 degrees).
Args:
lat: 1D array of latitude values, one per grid point (repeated
within each latitude ring), as read from the GRIB file.
lon: 1D array of longitude values, one per grid point, as read from
the GRIB file.
Returns:
tuple: ``(zlonc, zlatc, zlon, zlat, nlon, nlat)`` where ``zlonc``/
``zlatc`` are arrays of shape ``(4, ncell)`` giving the cell
vertices, ``zlon``/``zlat`` are the cell centers reshaped to
``(1, ncell)``, ``nlon`` is the number of cells and ``nlat`` is 1
(the grid is unstructured).
Raises:
CifValueError: If a latitude block is not constant, or if the
longitude spacing within a block is not regular.
"""
ncell = lat.shape[0]
lat_increasing = np.all(np.diff(lat) >= 0)
(indices,) = np.where(np.diff(lat) != 0)
indices = np.concatenate(([0], indices + 1, [ncell]))
lat_vertices = np.zeros((ncell, 4), dtype=lat.dtype)
lon_vertices = np.zeros((ncell, 4), dtype=lon.dtype)
for ind_start, ind_end in zip(indices[:-1], indices[1:]):
lat_block = lat[ind_start:ind_end]
lon_block = lon[ind_start:ind_end]
delta_lon = np.round(np.diff(lon_block), 12)
if len(np.unique(lat_block)) != 1:
raise CifValueError("Latitude values are not constant within the block")
if len(np.unique(delta_lon)) != 1:
raise CifValueError(
"Longitude values are not regularly spaced within the block"
)
lat_block = lat_block[0]
delta_lon = delta_lon[0]
if ind_start == 0:
if lat_increasing:
lat_bnds = (-90.0, 0.5 * (lat_block + lat[ind_end]))
else:
lat_bnds = (0.5 * (lat_block + lat[ind_end]), 90.0)
elif ind_end == ncell:
if lat_increasing:
lat_bnds = (0.5 * (lat[ind_start - 1] + lat_block), 90.0)
else:
lat_bnds = (-90.0, 0.5 * (lat[ind_start - 1] + lat_block))
else:
lat_bnds = (
0.5 * (lat[ind_start - 1] + lat_block),
0.5 * (lat_block + lat[ind_end]),
)
lon_bnds = np.concatenate(
(lon_block - 0.5 * delta_lon, [lon_block[-1] + 0.5 * delta_lon])
)
lat_vertices[ind_start:ind_end, 0] = lat_bnds[0]
lat_vertices[ind_start:ind_end, 1] = lat_bnds[1]
lat_vertices[ind_start:ind_end, 2] = lat_bnds[1]
lat_vertices[ind_start:ind_end, 3] = lat_bnds[0]
lon_vertices[ind_start:ind_end, 0] = lon_bnds[:-1]
lon_vertices[ind_start:ind_end, 1] = lon_bnds[:-1]
lon_vertices[ind_start:ind_end, 2] = lon_bnds[1:]
lon_vertices[ind_start:ind_end, 3] = lon_bnds[1:]
if np.any(lon_vertices > 180):
lon_vertices[lon_vertices > 180] -= 360
zlonc = lon_vertices.T
zlatc = lat_vertices.T
zlon = lon[np.newaxis, :]
zlat = lat[np.newaxis, :]
nlon = ncell
nlat = 1
return zlonc, zlatc, zlon, zlat, nlon, nlat
[docs]
def make_octahedral_grid(lat, lon, file_corner):
"""Build an unstructured cell-vertex grid for an octahedral GRIB grid.
Reads pre-computed corner coordinates from a csv file (since octahedral
grid geometry cannot be derived analytically from the center
coordinates alone).
Args:
lat: 1D array of latitude centers, as read from the GRIB file.
lon: 1D array of longitude centers, as read from the GRIB file.
file_corner: Path to a csv file with one line per grid cell, each
giving its 4 ``[lon, lat]`` corner pairs.
Returns:
tuple: ``(zlonc, zlatc, zlon, zlat, nlon, nlat)`` where ``zlonc``/
``zlatc`` are arrays of shape ``(4, ncell)`` giving the cell
vertices, ``zlon``/``zlat`` are the cell centers reshaped to
``(1, ncell)``, ``nlon`` is the number of cells (number of lines in
``file_corner``) and ``nlat`` is 1 (the grid is unstructured).
"""
# Read the csv coordinates
with open(file_corner, "r") as f:
lines = f.readlines()
data = np.zeros((len(lines), 4, 2))
for k, ln in enumerate(lines):
l = ln.strip().replace("[", "").replace("]", "").replace(",", "").split()
data[k] = np.array(l).astype(float).reshape(1, 4, 2)
zlonc = data[:, [0, 2, 3, 1], 0].T
zlatc = data[:, [0, 2, 3, 1], 1].T
zlon = lon[np.newaxis, :]
zlat = lat[np.newaxis, :]
nlon = len(lines)
nlat = 1
return zlonc, zlatc, zlon, zlat, nlon, nlat
[docs]
def get_domain(ref_dir, ref_file, input_interval, target_dir, tracer=None):
"""Build the pyCIF domain from an ECMWF GRIB reference file.
Opens the domain reference file (the first fetched input file, or
``tracer.domain_file`` if given), determines the grid type via
:func:`~pycif.plugins.datastreams.fields.grib2_ecmwf.utils.get_grid_type`
(``"octahedral"``, ``"reduced_gaussian"`` or ``"regular"``), builds the
corresponding center/corner grid, and reconstructs the hybrid vertical
coefficients ``sigma_a``/``sigma_b`` from the GRIB ``pv`` key (cropped
to the levels actually present via the ``hybrid`` key), or uses a
single surface level if ``tracer.surface`` is set.
Args:
ref_dir: Unused directly (kept for interface consistency);
reference files are taken from ``tracer.input_files``.
ref_file: Unused directly (kept for interface consistency).
input_interval: Unused directly (kept for interface consistency).
target_dir: Unused directly (kept for interface consistency).
tracer: Tracer/component configuration; ``tracer.input_files``,
``tracer.domain_file``, ``tracer.filter_by_keys_dict`` and
``tracer.surface`` are used.
Returns:
pycif.utils.classes.domains.Domain: The domain built from the
reference file's horizontal grid and vertical coefficients.
Raises:
CifFileNotFoundError: If the octahedral grid corner-coordinates
file cannot be found.
CifValueError: If the grid type is not one of the supported types.
"""
input_files = list(itertools.chain.from_iterable(tracer.input_files.values()))
domain_file = input_files[0][0]
if hasattr(tracer, "domain_file"):
domain_file = Path(domain_file).parent.joinpath(tracer.domain_file)
ds = GribDataset(
domain_file,
read_keys=["pv"],
filter_by_keys=tracer.filter_by_keys_dict,
)
lon = ds["longitude"]
lat = ds["latitude"]
grid_type = get_grid_type(domain_file, tracer.filter_by_keys_dict)
if grid_type == "octahedral":
jscan = 0
else:
jscan = get_jscan(domain_file, tracer.filter_by_keys_dict)
if grid_type == "octahedral":
path = Path(tracer.dir_hcoord_corner, tracer.file_hcoord_corner)
if not path.is_file():
raise CifFileNotFoundError(
f"Path to the corner coordinates does not exist, "
"please check that it is properly defined in your yml.\n"
f"Current definition is:\n"
f"\t- dir_hcoord_corner = {tracer.dir_hcoord_corner}\n"
f"\t- file_hcoord_corner = {tracer.file_hcoord_corner}\n"
)
zlonc, zlatc, zlon, zlat, nlon, nlat = make_octahedral_grid(lat, lon, path)
unstructured_domain = True
elif grid_type == "reduced_gaussian":
zlonc, zlatc, zlon, zlat, nlon, nlat = make_gaussian_grid(lat, lon)
unstructured_domain = True
elif grid_type == "regular":
zlonc, zlatc, zlon, zlat, nlon, nlat = make_regular_grid(lat, lon, jscan)
unstructured_domain = False
else:
raise CifValueError(f"Unsupported grid type {grid_type}")
# Reconstruct alpha and beta
if tracer.surface:
ecm_nlevs = 1
sigma_a = np.array([0])
sigma_b = np.array([1])
sigma_a_mid = np.array([0])
sigma_b_mid = np.array([1])
# Initializes domain
domain = Domain(
nlon=nlon,
nlat=nlat,
zlon=zlon,
zlat=zlat,
zlonc=zlonc,
zlatc=zlatc,
nlev=ecm_nlevs,
pressure_unit="Pa",
sigma_b=sigma_b,
sigma_a=sigma_a,
sigma_a_mid=sigma_a_mid,
sigma_b_mid=sigma_b_mid,
unstructured_domain=unstructured_domain,
)
else:
pv = ds.get_attr("pv")
hybrid = np.array(ds["hybrid"]).flatten().astype(int)
ecm_nlevs = int(len(pv) / 2 - 1)
sigma_a = pv[: ecm_nlevs + 1][::-1]
sigma_b = pv[ecm_nlevs + 1:][::-1]
# sigma_a = np.zeros(ecm_nlevs + 1)
# sigma_b = np.zeros(ecm_nlevs + 1)
# for ii in range(ecm_nlevs - 1):
# sigma_a[ii] = (pv[ecm_nlevs - ii] + pv[ecm_nlevs - ii - 1]) / 2
# sigma_b[ii] = \
# (pv[1 + 2 * ecm_nlevs - ii] + pv[1 + 2 * ecm_nlevs - ii - 1]) / 2
# Vertical crop if all hybrid levels are not available
ind_crop = np.concatenate(
[ecm_nlevs - hybrid[::-1], ecm_nlevs + 1 - hybrid[[0]]]
)
sigma_a = sigma_a[ind_crop]
sigma_b = sigma_b[ind_crop]
ecm_nlevs = hybrid.size
# Initializes domain
domain = Domain(
nlon=nlon,
nlat=nlat,
zlon=zlon,
zlat=zlat,
zlonc=zlonc,
zlatc=zlatc,
nlev=ecm_nlevs,
pressure_unit="Pa",
sigma_b=sigma_b,
sigma_a=sigma_a,
unstructured_domain=unstructured_domain,
)
return domain
