from osgeo import ogr
import itertools
import pandas as pd
import copy
import numpy as np
from itertools import zip_longest
from logging import info, debug
from ......utils.check.errclass import CifError, CifValueError
# Ignore future warnings from pyproj
import warnings
new_pyproj = True
try:
from pyproj import Proj, Transformer
except ImportError:
from pyproj import Proj, transform
new_pyproj = False
[docs]
def find_gridcells(
domain_in, domain_out,
forward_direction=True,
grid_to_surface=False,
chunksize=2e6
):
# Reversing domains if reverse direction
if not forward_direction:
domain_tmp = domain_in
domain_in = domain_out
domain_out = domain_tmp
zlonc_in = domain_in.zlonc
zlatc_in = domain_in.zlatc
zlon_in = domain_in.zlon
zlat_in = domain_in.zlat
# Check for NaNs in domain_out
zlon_out = domain_out.zlon
zlat_out = domain_out.zlat
if np.any(~pd.notnull(zlon_out) | ~pd.notnull(zlat_out)):
raise CifError(
"There are NaNs in the coordinates to fit in the original grid. "
"This may come from NaNs or errors in the way you generated your "
"observations. Please check your yml and observations."
)
is_unstructured = getattr(domain_in, "unstructured_domain", False)
isregular = (
np.sum(zlonc_in[0, np.newaxis] - zlonc_in) == 0
and np.sum(zlatc_in[:, 0, np.newaxis] - zlatc_in) == 0
)
# Force ascending order if is regular
discont = (
180
if getattr(domain_in, "projection", "gps") == "gps"
else np.ptp(zlonc_in)
)
if isregular:
ordered_lon = np.all(
np.diff(np.unwrap(zlonc_in[0], discont=discont)) > 0
)
if not ordered_lon:
if not np.all(np.diff(zlonc_in[0]) < 0):
raise CifValueError(
"Longitudes are neither in ascending or descending order. "
"I can't apply `find_grid_cells`. Please check your domain"
)
zlonc_in = zlonc_in[:, ::-1]
ordered_lat = np.all(np.diff(zlatc_in[:, 0]) > 0)
if not ordered_lat:
if not np.all(np.diff(zlatc_in[:, 0]) < 0):
raise CifValueError(
"Latitudes are neither in ascending or descending order. "
"I can't apply `find_grid_cells`. Please check your domain"
)
zlatc_in = zlatc_in[::-1]
# Remove duplicates to reduce the number of stations to locate
ds = pd.DataFrame({
"lon": np.round(np.asarray(zlon_out).flatten(), 5),
"lat": np.round(np.asarray(zlat_out).flatten(), 5),
"i": np.nan,
"j": np.nan
})
locations = ds.groupby(["lon", "lat"]).first()
# Check that there is at least one station
if len(locations) == 0:
# If not data, just returns empty output
if domain_out.zlon.size == 0:
weights = {
"i": np.zeros((0, 0)),
"j": np.zeros((0, 0)),
"wgt": np.zeros((0, 1)),
}
if reversed:
weights = {
"i": np.zeros((zlon_in.size, 0)),
"j": np.zeros((zlon_in.size, 0)),
"wgt": np.zeros((zlon_in.size, 0)),
}
return weights
# raise Exception("Trying to find gridcells from an empty datastore. "
# "Please check your input observations in: $WORKDIR/datavect")
raise CifError("There is no valid location (lon/lat) in your observation. "
"This error can be caused by, e.g., all lon/lat being NaNs. "
"Please check the observation inputs in: $WORKDIR/datavect")
# Makes simplified operations if regular
lon = locations.index.get_level_values("lon").values
lat = locations.index.get_level_values("lat").values
if isregular or is_unstructured:
listi = []
listj = []
# Loops over chunks to make the code faster
for lon_chunk, lat_chunk in zip(
zip_longest(*[iter(lon)] * int(chunksize), fillvalue=-999),
zip_longest(*[iter(lat)] * int(chunksize), fillvalue=-999),
):
lon_chunk = np.array(lon_chunk)[np.array(lon_chunk) != -999]
lat_chunk = np.array(lat_chunk)[np.array(lat_chunk) != -999]
if isregular:
i, j = find_gridcell(
lon_chunk,
lat_chunk,
zlonc_in,
zlatc_in,
isregular=isregular,
discont=discont,
)
if not ordered_lon:
j = domain_in.nlon - 1 - j
if not ordered_lat:
i = domain_in.nlat - 1 - i
else:
# For unstructured domains, take closest point
# TODO: generalize
dist = (zlon_in.T
- lon_chunk) ** 2 \
+ (zlat_in.T
- lat_chunk) ** 2
i = np.argmin(dist, axis=0)
j = np.zeros(len(lon_chunk))
listi.extend(list(i))
listj.extend(list(j))
else:
# Loop over (lon, lat) tuples
k = 0
nlocs = np.floor(locations.size / 10) + 1
listi = []
listj = []
for lon_tmp, lat_tmp in zip(lon, lat):
i, j = find_gridcell(
lon_tmp, lat_tmp, zlonc_in, zlatc_in,
isregular=isregular,
is_unstructured=is_unstructured)
listi.append(i)
listj.append(j)
k += 1
if k % nlocs == 0:
info(f"{k * 10.0 / nlocs}%")
# Putting i, j to locations datastore
locations["i"] = listi
locations["j"] = listj
# Distribute values back to the full datastore with duplicates
ds.loc[:, ["i", "j"]] = \
locations.iloc[ds.groupby(
["lon", "lat"]).ngroup().astype(int), :].values
# Produce weights for later use for reprojection
weights = {
"i": ds["i"].values[:, np.newaxis],
"j": ds["j"].values[:, np.newaxis],
"wgt": ds["j"].values[:, np.newaxis] * 0 + 1}
# Filter out observations outside the domain
inside = ~np.isnan(ds["i"]) & ~np.isnan(ds["j"])
if forward_direction:
weights = {
"i": weights["i"][inside],
"j": weights["j"][inside],
"wgt": weights["wgt"][inside],
"filtered": np.where(inside)[0],
"non_filtered": np.where(~inside)[0]
}
return weights
# # Reversing weights if reverse direction
# nlat, nlon = zlon_in.shape
# ind_out = np.ravel_multi_index(
# np.concatenate([weights["i"][inside],
# weights["j"][inside]], axis=1).T.astype(int),
# (nlat, nlon), order="F")
# filtered = np.unique(ind_out)
# Reversing weights if reverse direction
nlat, nlon = zlon_in.shape
groups = ds.groupby(['i', 'j'])
filtered = np.ravel_multi_index(
groups.count().reset_index().loc[:, ["i", "j"]].values.astype(int).T,
(nlat, nlon), order="F")
# Initialize areas if needed
if grid_to_surface:
domain = domain_in
if not hasattr(domain, "areas"):
domain.calc_areas()
# Fill output groups with input coordinates
out_weights = {
"j": [list(groups.groups[k]) for k in groups.groups if not np.isnan(k[0])],
"i": [len(groups.groups[k]) * [0] for k in groups.groups if not np.isnan(k[0])],
"wgt": [
len(groups.groups[k]) * [1] if not grid_to_surface
else len(groups.groups[k]) * [1 / domain.areas[tuple(np.array(k).astype(int))]]
for k in groups.groups if not np.isnan(k[0])],
"filtered": filtered,
# "non_filtered": filtered
}
out_weights = {
"i": np.array(list(
itertools.zip_longest(*out_weights["i"], fillvalue=0))).T,
"j": np.array(list(
itertools.zip_longest(*out_weights["j"], fillvalue=0))).T,
"wgt": np.array(list(
itertools.zip_longest(*out_weights["wgt"], fillvalue=np.nan))).T,
"filtered": filtered,
# "non_filtered": filtered
}
return out_weights
[docs]
def find_gridcell(
lon,
lat,
zlonc,
zlatc,
orig_proj="epsg:4326" if new_pyproj else Proj(init="epsg:4326"),
isregular=False,
is_unstructured=False,
discont=180,
):
"""Finds the grid cell corresponding to a coordinate.
Args:
lon (np.array): longitude of the point to find
lat (np.array): latitude of the point to find
zlonc (np.array): longitudes of the corners of the grid
zlatc (np.array): latitudes of the corners of the grid
orig_proj (projection): the projection for reporting coordinates.
Default is WSG84
isregular (Boolean): if True, simplified operations are computed.
Default is False
Returns:
i, j the grid cell ID
Notes:
- For very regular grids, this script could be made more effileient
and shorter, but the objective is to be able to deal with any domain
- zlonc and zlatc must be in ascending order
"""
# If regular domain make simplified operations
if isregular and not is_unstructured:
try:
# Catch and ignore divide by zero
with np.errstate(divide='ignore', invalid='ignore'):
xlon = 1.0 / np.unwrap(
np.radians(zlonc[0, np.newaxis] - lon[:, np.newaxis]),
discont=np.radians(discont)
)
xlat = 1.0 / (zlatc[:, 0] - lat[:, np.newaxis])
xlon[xlon == np.inf] = -np.inf
xlat[xlat == np.inf] = -np.inf
i = xlat.argmin(axis=1)
j = xlon.argmin(axis=1)
nlat, nlon = zlonc[:-1, :-1].shape
# Put NaNs for points outside the domain
maski = (lat < zlatc.min()) | (lat > zlatc.max())
maskj = (lon < np.degrees(np.unwrap(np.radians(zlonc),
discont=np.radians(discont)).min())) \
| (lon > np.degrees(np.unwrap(np.radians(zlonc),
discont=np.radians(discont)).max()))
i = np.where(maski | maskj, np.nan, i)
j = np.where(maski | maskj, np.nan, j)
# Points in last column or row should be shifted
# as corners were used
maski = (i >= nlat)
maskj = (j >= nlon)
i = np.where(maski, nlat - 1, i)
j = np.where(maskj, nlon - 1, j)
return i, j
# If only one (lon, lat) tuple was provided
except TypeError as e:
info(e)
xlon = 1.0 / np.unwrap(np.radians(zlonc[0] - lon),
discont=np.radians(discont))
xlat = 1.0 / (zlatc[:, 0] - lat)
return xlat.argmin(), xlon.argmin()
# Grid shape
nlat, nlon = zlonc.shape
# Define measurement point geometry in local coordinates
point = ogr.Geometry(ogr.wkbPoint)
point.AddPoint(0, 0)
# Define local projection
if new_pyproj:
target_proj = \
f"+proj=laea +lat_0={lat} +lon_0={lon} +x_0=0 +y_0=0 "\
"+ellps=WGS84 +units=m +no_defs "
t = Transformer.from_crs(orig_proj, target_proj, always_xy=True)
zxc, zyc = t.transform((zlonc + 180) % 360 - 180, zlatc)
else:
target_proj = Proj(
f"+proj=laea +lat_0={lat} +lon_0={lon} +x_0=0 +y_0=0 "
"+ellps=WGS84 +units=m +no_defs "
)
zxc, zyc = transform(
orig_proj, target_proj, (zlonc + 180) % 360 - 180, zlatc
)
dist = zxc ** 2 + zyc ** 2
imin, jmin = np.unravel_index(dist.argmin(), dist.shape)
imin = imin % nlat
jmin = jmin % nlon
for i in range(imin - 3, imin + 2):
for j in range(jmin - 3, jmin + 2):
# Looping if the minimum is close to the domain side
i = max(min(i, nlat - 2), 0)
j = max(min(j, nlon - 2), 0)
# Create ring
ring = ogr.Geometry(ogr.wkbLinearRing)
_ = ring.AddPoint(zxc[i, j], zyc[i, j])
_ = ring.AddPoint(zxc[i + 1, j], zyc[i + 1, j])
_ = ring.AddPoint(zxc[i + 1, j + 1], zyc[i + 1, j + 1])
_ = ring.AddPoint(zxc[i, j + 1], zyc[i, j + 1])
_ = ring.AddPoint(zxc[i, j], zyc[i, j])
# Create polygon
poly = ogr.Geometry(ogr.wkbPolygon)
_ = poly.AddGeometry(ring)
# Buffering the polygon to include points on the edge
buffer = 1.0
poly = poly.Buffer(buffer)
if point.Within(poly):
return i, j
# If no grid cell was fund, raise exception
#
# raise IndexError(
# "No index was found for measurement at {}, {}".format(lon, lat))
# TODO: Decide whether we raise exception
# when the observation is outside the domain
return np.nan, np.nan
