import numpy as np
import os
from netCDF4 import Dataset
import datetime
import itertools
from osgeo import ogr
from logging import info, debug
from .....utils.check.errclass import CifError
[docs]
def hresol2dim(tracer, dom, **kwargs):
"""Computes the horizontal size of a control vector from its resolution
Args:
tracer (Plugin): definition of the tracer, including the resolution and
additional information on the resolution
domain (dict): the domain grid
Returns
int: the size of the control vector for this component
"""
if tracer.hresol == "hpixels":
out_dimensions = (
dom.zlon.size
if not getattr(tracer, "is_lbc", False)
else dom.zlon_side.size
)
elif tracer.hresol == "bands":
# Check that the lists are sorted
if (
not sorted(tracer.bands_lat) == tracer.bands_lat
or not sorted(tracer.bands_lon) == tracer.bands_lon
):
raise CifError(
"Bands in the control vector are not sorted!\n"
"Please check your configuration file."
)
# Compute the total number of dimensions in the grid
if not hasattr(tracer, "nbands"):
tracer.nbands = (len(tracer.bands_lat) - 1) * (
len(tracer.bands_lon) - 1
)
out_dimensions = tracer.nbands
elif tracer.hresol == "ibands":
# Check that the lists are sorted
if (
not sorted(tracer.bands_i) == tracer.bands_i
or not sorted(tracer.bands_j) == tracer.bands_j
):
raise CifError(
"Bands in the control vector are not sorted!\n"
"Please check your configuration file."
)
# Compute the total number of dimensions in the grid
if not hasattr(tracer, "nbands"):
tracer.nbands = (len(tracer.bands_i) - 1) * (
len(tracer.bands_j) - 1
)
out_dimensions = tracer.nbands
elif tracer.hresol == "regions":
if not hasattr(tracer, "regions"):
region_infos = tracer.regions_infos
if hasattr(region_infos, "read") \
and not getattr(region_infos, "default_read", False):
tracer.regions = region_infos.read(
"regions",
"regions",
[[datetime.datetime(1970, 1, 1),
datetime.datetime(1970, 1, 1)]],
[os.path.join(region_infos.dir, region_infos.file)],
tracer=region_infos
)[0, 0].values
file_regions = f"{region_infos.dir} / {region_infos.file}"
else:
file_regions = os.path.join(
region_infos.dir, region_infos.file)
with Dataset(file_regions, "r") as f:
tracer.regions = f.variables["regions"][:]
# Unmask arrays
if np.ma.isMaskedArray(tracer.regions):
tracer.regions = tracer.regions.data
# Check that regions have the correct dimensions
if tracer.regions.shape != dom.zlat.shape:
raise CifError(
f"Regions were not correctly defined in {file_regions}"
)
# Number of regions
tracer.region_ids = np.unique(
tracer.regions[~np.isnan(tracer.regions)])
tracer.nregions = len(tracer.region_ids)
# Deal with land sea mask
regions_lsm = getattr(tracer, 'regions_lsm', False)
if regions_lsm:
# Set ocean and land mask according to the region definition file
# Negative numbers are for ocean, positive numbers are for land
tracer.regions_lsmask = np.ndarray(tracer.regions.shape)
# Set ocean and land mask
tracer.regions_lsmask[tracer.regions < 0] = 0
tracer.regions_lsmask[tracer.regions > 0] = 1
# Default behaviour: optimize ocean boxes
tracer.inc_ocean = getattr(tracer, 'inc_ocean', True)
out_dimensions = tracer.nregions
elif tracer.hresol == "global":
out_dimensions = 1
else:
raise CifError(f"This resolution is not processed by the CIF: {tracer.hresol}")
# Pre-compute areas if not already done
if not hasattr(tracer.domain, "areas"):
if hasattr(tracer, "glob_err") \
or tracer.hresol not in ["global", "hpixels"]:
info("Computing areas")
tracer.domain.calc_areas(**kwargs)
# Stop here if not regions nor bands
if tracer.hresol in ["global", "hpixels"]:
return out_dimensions
# Compute region areas
debug("Computing region total areas. This can take a while...")
zlon = tracer.domain.zlon
zlat = tracer.domain.zlat
if tracer.hresol == "regions":
mask_regions = ~np.isnan(tracer.regions)
_, idx, _ = np.unique(
tracer.regions[mask_regions],
return_counts=True, return_inverse=True
)
tracer.region_areas = np.bincount(
idx, tracer.domain.areas[mask_regions].flatten())
elif tracer.hresol == "bands":
dlat = zlat[..., np.newaxis] - \
np.array(tracer.bands_lat)[np.newaxis, np.newaxis, 1:]
ilat = np.argmax(dlat < 0, axis=2).astype(float)
ilat[(zlat < np.min(tracer.bands_lat))
| (zlat > np.max(tracer.bands_lat))] = np.nan
dlon = zlon[..., np.newaxis] - \
np.array(tracer.bands_lon)[np.newaxis, np.newaxis, 1:]
ilon = np.argmax(dlon < 0, axis=2).astype(float)
ilon[(zlon < np.min(tracer.bands_lon))
| (zlon > np.max(tracer.bands_lon))] = np.nan
bands_idx = ilat * (len(tracer.bands_lon) - 1) + ilon
mask_bands = ~np.isnan(bands_idx)
_, idx, _ = np.unique(
bands_idx[mask_bands],
return_counts=True, return_inverse=True
)
tracer.region_areas = np.bincount(
idx, tracer.domain.areas[mask_bands].flatten())
elif tracer.hresol == "ibands":
ilon = -1 * np.ones_like(zlon)
jbands = np.maximum(0, tracer.bands_j[:-1]).astype(int)
jbands_index = np.arange(len(jbands) - 1)
ilon[:, np.array(tracer.bands_j[:-1]).astype(int)
] = np.arange(len(tracer.bands_j) - 1)
ilon = np.maximum.accumulate(ilon, axis=1)
ilat = -1 * np.ones_like(zlat)
ilat[np.array(tracer.bands_i[:-1]).astype(int)
] = np.arange(len(tracer.bands_i) - 1)[:, np.newaxis]
ilat = np.maximum.accumulate(ilat, axis=0)
# Crop outside of bands
meshj, meshi = np.meshgrid(
np.arange(zlat.shape[1]), np.arange(zlat.shape[0])
)
ilon[(meshj < np.min(tracer.bands_j))
| (meshj > np.max(tracer.bands_j))] = np.nan
ilat[(meshi < np.min(tracer.bands_i))
| (meshi > np.max(tracer.bands_i))] = np.nan
# Compute bands areas
bands_idx = ilat * (len(tracer.bands_j) - 1) + ilon
mask_bands = ~np.isnan(bands_idx)
_, idx, _ = np.unique(
bands_idx[mask_bands],
return_counts=True, return_inverse=True
)
tracer.region_areas = np.bincount(
idx, tracer.domain.areas[mask_bands].flatten())
else:
raise CifError(
f"This resolution is not processed by the CIF: {tracer.hresol}"
)
# Compute centroids of regions
if hasattr(tracer, "hcorrelations"):
debug("Computing centroids. This can take very long...")
# TODO: switch to geopandas to optimze the computation
tracer.centroids_lons = []
tracer.centroids_lats = []
# List of regions depending on type of resolution
reg_list = []
if tracer.hresol == "regions":
reg_list = tracer.region_ids
elif tracer.hresol in ["ibands", "bands"]:
reg_list = range(tracer.nbands)
tracer.region_areas2 = []
for regID in reg_list:
# Get list of lon/lat of the sub-region
zlon = tracer.domain.zlon
zlat = tracer.domain.zlat
if tracer.hresol == "regions":
mask = tracer.regions == regID
lons = zlon[mask]
lats = zlat[mask]
# tracer.region_areas.append(
# tracer.domain.areas[mask].sum())
elif tracer.hresol == "bands":
lat_bounds, lon_bounds = list(itertools.product(
zip(tracer.bands_lat[:-1], tracer.bands_lat[1:]),
zip(tracer.bands_lon[:-1], tracer.bands_lon[1:])
))[regID]
mask = (
(zlon >= lon_bounds[0])
& (zlon < lon_bounds[1])
& (zlat >= lat_bounds[0])
& (zlat < lat_bounds[1])
)
lons = zlon[mask]
lats = zlat[mask]
elif tracer.hresol == "ibands":
i_bounds, j_bounds = list(itertools.product(
zip(tracer.bands_i[:-1], tracer.bands_i[1:]),
zip(tracer.bands_j[:-1], tracer.bands_j[1:])
))[regID]
lons = zlon[i_bounds[0]:i_bounds[1],
j_bounds[0]:j_bounds[1]].flatten()
lats = zlat[i_bounds[0]:i_bounds[1],
j_bounds[0]:j_bounds[1]].flatten()
tracer.region_areas2.append(
tracer.domain.areas[i_bounds[0]:i_bounds[1],
j_bounds[0]:j_bounds[1]].sum())
else:
raise CifError(
f"This resolution is not processed by the CIF: {tracer.hresol}"
)
multipoint = ogr.Geometry(ogr.wkbMultiPoint)
for lon, lat in zip(lons, lats):
pt = ogr.Geometry(ogr.wkbPoint)
pt.AddPoint(float(lon), float(lat))
multipoint.AddGeometry(pt)
centroid = multipoint.Centroid()
tracer.centroids_lons.append(centroid.GetX())
tracer.centroids_lats.append(centroid.GetY())
tracer.region_areas = np.array(tracer.region_areas)[:, np.newaxis]
tracer.centroids_lons = np.array(
tracer.centroids_lons)[:, np.newaxis]
tracer.centroids_lats = np.array(
tracer.centroids_lats)[:, np.newaxis]
return out_dimensions
[docs]
def vresol2dim(tracer, dom, **kwargs):
"""Computes the horizontal size of a control vector from its resolution
Args:
tracer (Plugin): definition of the tracer, including the resolution and
additional information on the resolution
domain (dict): the domain grid
Returns
int: the size of the control vector for this component
"""
# Default vertical resolution is integrated columns
tracer.vresol = getattr(tracer, "vresol", "column")
# Getting the number of vertical levels form the tracer domain if not
# already defined (defaults to nlev = 1)
if not hasattr(tracer, "nlev"):
if hasattr(tracer, "domain"):
tracer.nlev = tracer.domain.nlev
else:
tracer.nlev = 1
# Loop over possible vertical resolution
if tracer.vresol == "vpixels":
tracer.levels = np.arange(tracer.nlev)
return tracer.nlev
elif tracer.vresol == "kbands":
# Compute the total number of dimensions in the grid
if not hasattr(tracer, "nvbands"):
tracer.nvbands = len(tracer.kbands) - 1
tracer.levels = np.array(tracer.kbands[:-1])
return tracer.nvbands
elif tracer.vresol == "column":
tracer.levels = np.array([0])
tracer.nlev = 1
return tracer.nlev
