import numpy as np
import xarray as xr
from ......utils.check.errclass import CifError
[docs]
def scale2map(x, tracer, dates, dom, ensemble=False):
"""Unpack a control-vector slice to a spatial map at the model grid resolution.
Projects the 1-D (or 2-D ensemble) control-vector slice onto the full
model domain according to the tracer's horizontal resolution type:
* ``hpixels`` — each element maps to one grid cell (direct reshape).
* ``bands`` — each element fills a lat/lon band region.
* ``ibands`` — each element fills a rectangular index-space band.
* ``regions`` — each element fills a pre-defined geographic region
(NaN outside any region).
* ``global`` — a single element broadcasts over the entire domain.
Args:
x (np.ndarray): control-vector slice of shape ``(ndates, vresoldim,
nhresol)`` or ``(nsamples, ndates, vresoldim, nhresol)`` when
*ensemble* is ``True``.
tracer: tracer plugin instance with attributes ``hresol``,
``vresoldim``, ``levels``, ``domain``, and band/region
definitions.
dates (array-like): datetime array of length ``ndates`` used as the
``time`` coordinate of the output DataArray.
dom: domain plugin providing grid coordinates (``zlon``, ``zlat``).
ensemble (bool): if ``True``, *x* carries a leading ensemble
dimension and the output retains it.
Returns:
xarray.DataArray: shape ``(time, lev, lat, lon)`` or
``(ens, time, lev, lat, lon)`` when *ensemble* is ``True``.
Raises:
Exception: if ``tracer.hresol`` is not one of the recognised types.
"""
if ensemble:
nsamples = len(x)
ndates = len(x[0])
else:
nsamples = 1
ndates = len(x)
x = x[np.newaxis]
if not getattr(tracer, "is_lbc", False):
nlat, nlon = dom.zlat.shape
zlon = dom.zlon
zlat = dom.zlat
else:
nlon = dom.nlon_side
nlat = dom.nlat_side
zlon = dom.zlon_side
zlat = dom.zlat_side
if tracer.hresol == "hpixels":
xmap = x.reshape(nsamples, -1, tracer.vresoldim, nlat, nlon)
elif tracer.hresol == "bands":
xmap = np.zeros((nsamples, ndates, tracer.vresoldim, nlat, nlon))
iband = 0
for lat1, lat2 in zip(tracer.bands_lat[:-1], tracer.bands_lat[1:]):
for lon1, lon2 in zip(tracer.bands_lon[:-1], tracer.bands_lon[1:]):
reg = (
(zlon >= lon1)
& (zlon < lon2)
& (zlat >= lat1)
& (zlat < lat2)
)
xmap[..., reg] = x[..., iband, np.newaxis]
iband += 1
elif tracer.hresol == "ibands":
xmap = np.zeros((nsamples, ndates, tracer.vresoldim, nlat, nlon))
iband = 0
for i1, i2 in zip(tracer.bands_i[:-1], tracer.bands_i[1:]):
for j1, j2 in zip(tracer.bands_j[:-1], tracer.bands_j[1:]):
xmap[..., i1:i2, j1:j2] = \
x[..., iband, np.newaxis, np.newaxis]
iband += 1
elif tracer.hresol == "regions":
# As of today, pixels outside any region of the control vector
# are set to one
xmap = np.zeros(
(nsamples, ndates, tracer.vresoldim, nlat, nlon)) + np.nan
for regID, reg in enumerate(tracer.region_ids):
xmap[..., tracer.regions == reg] = x[..., regID, np.newaxis]
elif tracer.hresol == "global":
xmap = x[..., np.newaxis] \
* np.ones((nsamples, ndates, tracer.vresoldim, nlat, nlon))
else:
raise CifError(f"{tracer.hresol} is not recognized")
# Putting in xarray dataset
xmap = xr.DataArray(
xmap,
coords={"time": dates, "lev": tracer.levels},
dims=("ens", "time", "lev", "lat", "lon"),
)
if not ensemble and nsamples == 1:
return xmap[0]
return xmap
[docs]
def map2scale(xmap, tracer, dom,
region_scale_area=False,
region_max_val=False):
"""Project a spatial sensitivity map onto the control-vector horizontal grid.
The adjoint of :func:`scale2map`: aggregates a model-space 4-D field
``(time, lev, lat, lon)`` to the control-vector horizontal resolution:
* ``hpixels`` — reshape: each grid cell maps to one control element.
* ``bands`` — sum over all cells in each lat/lon band.
* ``ibands`` — sum over all cells in each index-space band.
* ``regions`` — sum (or area-weighted mean, or max) over each region.
* ``global`` — sum over the entire domain.
Args:
xmap (np.ndarray): sensitivity field, shape
``(time, lev, lat, lon)``.
tracer: tracer plugin instance (same as passed to :func:`scale2map`).
dom: domain plugin providing grid coordinates.
region_scale_area (bool): for ``regions`` resolution, weight each
cell by its area and normalise by the total region area instead
of summing.
region_max_val (bool): for ``regions`` resolution, take the
spatial maximum instead of the sum.
Returns:
np.ndarray: aggregated slice of shape
``(ndates, vresoldim, nhresol)``.
Raises:
Exception: if the input does not have exactly 4 dimensions, or if
``tracer.hresol`` is not recognised.
"""
# Checking that xmap has the correct dimension
if not len(xmap.shape) == 4:
raise CifError(
f"Warning! map2scale expects inputs data of dimension:(time, levels, lat, lon). Got only {len(xmap.shape)} dimensions instead"
)
ndates = xmap.shape[0]
if not getattr(tracer, "is_lbc", False):
nlon = dom.nlon
nlat = dom.nlat
zlon = dom.zlon
zlat = dom.zlat
else:
nlon = dom.nlon_side
nlat = dom.nlat_side
zlon = dom.zlon_side
zlat = dom.zlat_side
# Dealing different resolution types
if tracer.hresol == "hpixels":
x = xmap.reshape((ndates, tracer.vresoldim, -1))
elif tracer.hresol == "bands":
x = np.zeros((ndates, tracer.vresoldim, tracer.nbands))
iband = 0
for lat1, lat2 in zip(tracer.bands_lat[:-1], tracer.bands_lat[1:]):
for lon1, lon2 in zip(tracer.bands_lon[:-1], tracer.bands_lon[1:]):
reg = (
(zlon >= lon1)
& (zlon < lon2)
& (zlat >= lat1)
& (zlat < lat2)
)
x[..., iband] = np.sum(xmap[..., reg], axis=2)
iband += 1
elif tracer.hresol == "ibands":
x = np.zeros((ndates, tracer.vresoldim, tracer.nbands))
iband = 0
for i1, i2 in zip(tracer.bands_i[:-1], tracer.bands_i[1:]):
for j1, j2 in zip(tracer.bands_j[:-1], tracer.bands_j[1:]):
x[..., iband] = np.sum(xmap[..., i1:i2, j1:j2], axis=(2, 3))
iband += 1
elif tracer.hresol == "regions":
x = np.zeros((ndates, tracer.vresoldim, tracer.nregions))
if not region_max_val and not region_scale_area:
mask_regions = ~np.isnan(tracer.regions)
_, idx, _ = np.unique(
tracer.regions[mask_regions],
return_counts=True, return_inverse=True
)
for idate in range(ndates):
for ilev in range(tracer.vresoldim):
x[idate, ilev] = np.bincount(
idx, xmap[idate, ilev, mask_regions].flatten())
else:
for regID, reg in enumerate(tracer.region_ids):
mask = tracer.regions == reg
if region_scale_area:
x[..., regID] = np.sum(
xmap[..., mask] * tracer.domain.areas[..., mask],
axis=2)
x[..., regID] /= tracer.region_areas[regID]
elif region_max_val:
x[..., regID] = np.max(xmap[..., mask])
elif tracer.hresol == "global":
x = xmap.sum(axis=(2, 3))[..., np.newaxis]
else:
raise CifError(f"{tracer.hresol} is not recognized")
return x
[docs]
def vmap2vaggreg(data, tracer, dom, tracer_id):
"""Aggregate a full 3-D sensitivity field to the control-vector vertical resolution.
Three vertical resolutions are supported:
* ``column`` — sum over all levels, returning a single vertical layer.
* ``vpixels`` — keep each level individually (no aggregation); validates
that the number of model levels matches ``tracer.vresoldim``.
* ``kbands`` — sum within each contiguous vertical band defined by
``tracer.kbands``.
Args:
data (np.ndarray): sensitivity field, shape ``(time, lev, lat, lon)``.
tracer: tracer plugin instance with ``vresol``, ``vresoldim``,
``nlev``, and (for ``kbands``) ``nvbands`` / ``kbands``.
dom: domain plugin (unused; kept for API consistency).
tracer_id (tuple): ``(component, parameter)`` for error messages.
Returns:
np.ndarray: aggregated field of shape
``(time, vresoldim, lat, lon)``.
Raises:
Exception: if the input does not have 4 dimensions, if ``vpixels``
level count mismatches, or if ``vresol`` is not recognised.
"""
# Checking that xmap has the correct dimension
if not len(data.shape) == 4:
raise CifError(
f"Warning! vmap2vaggreg expects inputs data of dimension: (time, levels, lat, lon). Got only {len(data.shape)} dimensions instead."
)
# Flat fields are directly returned
if tracer.vresol == "column":
return data.sum(axis=1)[:, np.newaxis, ...]
elif tracer.vresol == "vpixels":
if data.shape[1] != tracer.vresoldim:
raise CifError(f"Adjoint sensitivity have {data.shape[1]} levels, while the corresponding control vector component has {tracer.vresoldim} levels. Please check the corresponding paragraph of your Yaml ({tracer_id[0]}/{tracer_id[1]}): the nlev ({tracer.nlev}) should be consistent with the number of levels in the sensitivity ({data.shape[1]})")
return data
elif tracer.vresol == "kbands":
outshape = list(data.shape)
outshape[1] = tracer.nvbands
outdata = np.zeros(outshape)
iband = 0
for k1, k2 in zip(tracer.kbands[:-1], tracer.kbands[1:]):
outdata[:, iband, ...] = data[:, k1:k2, ...].sum(axis=1)
iband += 1
return outdata
else:
raise CifError(f"{tracer.vresol} is not recognized")
