Source code for openghg.transform._regrid
from typing import Optional, Tuple, Union, cast
import numpy as np
import xarray as xr
from numpy import ndarray
from openghg.types import construct_xesmf_import_error, ArrayLike, ArrayLikeMatch
def _getGridCC(lat: ndarray, lon: ndarray) -> Tuple[ndarray, ndarray]:
"""
Create a cell centered meshgrid for from 1D arrays of lon, lat
This meshgrid defines the bounds of each cell.
"""
# dx = lon[2]-lon[1]
# dy = lat[2]-lat[1]
dx = np.mean(lon[1:] - lon[:-1])
dy = np.mean(lat[1:] - lat[:-1])
lon = np.append(lon, lon[-1] + dx)
lat = np.append(lat, lat[-1] + dy)
lon -= dx / 2.0
lat -= dy / 2.0
LON, LAT = np.meshgrid(lon, lat)
return LON, LAT
# def _create_xesmf_grid_uniform_cc(lat: ndarray, lon: ndarray) -> xr.Dataset:
# """
# Creates a Dataset ready to be used by the xesmf regridder from 1D arrays
# of latitude and longitude values.
# This includes both the centre points and bounds of the latitude and longitude
# grids (needed for conservative regridding methods but not for bilinear etc.)
# """
# LON, LAT = np.meshgrid(lon, lat)
# LON_b, LAT_b = _getGridCC(lat, lon)
# grid = xr.Dataset(
# {
# "lon": (["x", "y"], LON),
# "lat": (["x", "y"], LAT),
# "lon_b": (["x_b", "y_b"], LON_b),
# "lat_b": (["x_b", "y_b"], LAT_b),
# }
# )
# return grid
def _create_uniform_coords(lat: ndarray, lon: ndarray) -> xr.Dataset:
"""
Creates a Dataset ready to be used by the xesmf regridder from 1D arrays
of latitude and longitude values.
"""
ds_out = xr.Dataset(
{
"lat": (["lat"], lat, {"units": "degrees_north"}),
"lon": (["lon"], lon, {"units": "degrees_east"}),
}
)
return ds_out
def convert_to_ndarray(array: Union[ndarray, xr.DataArray]) -> ndarray:
"""Check and extract underlying numpy array from DataArray as necessary"""
if isinstance(array, xr.DataArray):
values = array.values
else:
values = array
return values
[docs]
def regrid_uniform_cc(
data: ArrayLikeMatch,
lat_out: ArrayLike,
lon_out: ArrayLike,
lat_in: Optional[ArrayLike] = None,
lon_in: Optional[ArrayLike] = None,
latlon: Optional[list] = None,
method: str = "conservative",
) -> ArrayLikeMatch:
"""
Regrid data between two uniform, cell centered grids.
All coordinates (lat_out, lon_out, lat_in, lon_in) should be for the centre
of the representative cell and in degrees.
Adapted from code written by @DTHoare
Args:
data: Data to be regridded.
Data must have dimensions (lat, lon) if 2D or (time, lat, lon) if 3D.
lat_out: 1D array for output latitude grid
lon_out: 1D array for output longituide grid
lat_in: 1D array for input latitude grid.
Only used if data is a numpy array and not a DataArray
lon_in: 1D array for input longitude grid.
Only used if data is a numpy array and not a DataArray
latlon: Names for latitude and longitude coordinates within data.
Default = ["lat", "lon"]
method: Method to use for regridding. Mainly use:
- "conservative"
- "conservative_normed" (ignores NaN values)
See xesmf documentation for full list of options.
Returns:
ndarray / DataArray : Regridded data using specified method
"""
try:
import xesmf # type:ignore
except ImportError as e:
raise ImportError(construct_xesmf_import_error(e))
if latlon is None:
latlon = ["lat", "lon"]
if isinstance(data, xr.DataArray):
lat_in_extracted = data[latlon[0]].values
lon_in_extracted = data[latlon[1]].values
if lat_in is not None:
if not np.isclose(lat_in, lat_in_extracted):
raise ValueError(
"Input for 'lat_in' should not be supplied if 'data' is a DataArray object.\n"
"Please check 'lat_out' have been supplied correctly as well."
)
if lon_in is not None:
if not np.isclose(lon_in, lon_in_extracted):
raise ValueError(
"Input for 'lon_in' should not be supplied if 'data' is a DataArray object.\n"
"Please check 'lon_out' has been supplied correctly as well."
)
lat_in = lat_in_extracted
lon_in = lon_in_extracted
lat_in = cast(ArrayLike, lat_in)
lon_in = cast(ArrayLike, lon_in)
if data.shape != (lat_in.size, lon_in.size):
raise ValueError(
f"Shape of input 'data' {data.shape}"
f"does not match 'lat_in' and 'lon_in' lengths: {len(lat_in)}, {len(lon_in)}"
)
lat_in = convert_to_ndarray(lat_in)
lon_in = convert_to_ndarray(lon_in)
lat_out = convert_to_ndarray(lat_out)
lon_out = convert_to_ndarray(lon_out)
# 09/03/2023: Don't seem to need inputs with centre and bounding boxes for
# uniform grid and xesmf "conservative" method anymore.
input_grid = _create_uniform_coords(lat_in, lon_in)
output_grid = _create_uniform_coords(lat_out, lon_out)
regridder = xesmf.Regridder(input_grid, output_grid, method)
regridded: ArrayLikeMatch = regridder(data)
# # 09/03/2023: Don't need this for uniform grid anymore due to changes above
# # but a variant may be needed for non-uniform grids e.g. tropomi data.
# if isinstance(regridded, xr.DataArray):
# from scipy.sparse import coo_matrix # type:ignore
# # Checking dimensions and mapping back lat_out and lon_out
# # May be overkill but attempting to make sure this is done correctly.
# # The regridded DataArray will contain dimensions of ('x', 'y') for
# # data, 'lat' and 'lon' coordinates e.g. lon = [[-10,-10],[0,0]]
# # This is to allow for the generic case where ('lat', 'lon') is not
# # a rectilinear (uniform) grid.
# # Since we are creating a uniform grid we want to flatten this and
# # put data back on ('lat', 'lon') dimension.
# regridded_stack = regridded.stack(z=("x", "y"))
# lat_coord = regridded_stack["lat"]
# lon_coord = regridded_stack["lon"]
# # Find coords within our lat_out and lon_out grid for our regridded output
# # Digitize is technically a binning operation outputting the number of the
# # bin the value is within so to get indicies we can subtract 1.
# lat_index = np.digitize(lat_coord, lat_out) - 1
# lon_index = np.digitize(lon_coord, lon_out) - 1
# # Create a sparse matrix from the flattened data and the lat, lon index values.
# # Reshape to our required lat_out, lon_out and output as a numpy array
# shape = (len(lat_out), len(lon_out))
# regridded_grid = coo_matrix((regridded_stack.data, (lat_index, lon_index)), shape=shape).toarray()
# regridded = xr.DataArray(
# regridded_grid,
# dims=("lat", "lon"),
# coords={"lat": lat_out, "lon": lon_out},
# attrs=regridded.attrs,
# )
# # # Alternative, more simplistic approach. May not be generic
# # regridded = regridded.drop(labels=("lat","lon"))
# # regridded = regridded.assign_coords(**{"x":output_lat,"y":output_lon})
# # regridded = regridded.rename({"x":"lat","y":"lon"})
# # In later versions of xesmf (0.7?) this no longer seems to be needed
# # as weight files are not stored on disk.
# regridder.clean_weight_file()
return regridded
# def regrid_betweenGrids(data, input_grid, output_grid, method="conservative"):
# """
# Regrid data from predefined input_grid and output_grid
# Inputs
# data - numpy array
# input_grid, output_grid - Dataset of the form:
# xr.Dataset({'lat': (['x', 'y'], LAT),
# 'lon': (['x', 'y'], LON),
# 'lat_b': (['x_b', 'y_b'], LAT_b),
# 'lon_b': (['x_b', 'y_b'], LON_b)})
# where lat and lon give cell centre locations, and lat_b and lon_b give the cell bounds (corners)
# method - string describing method to use.
# Should be one of "conservative" or "conservative_normed".
# Note that to use the conservative_normed method you need to have installed the "masking" development
# branch of the xesmf package (contact Daniel for more details).
# With the 'masking' branch of xESMF you can include a mask in the input_grid to ignore nan values
# returns
# regridded numpy array
# """
# try:
# import xesmf
# except ImportError:
# raise ImportError(xesmf_error_message)
# #if 'mask' in input_grid:
# # # !!! Requires the 'masking' branch of xESMF to be manually installed !!!
# # method = 'conservative_normed'
# #else:
# #method = 'conservative'
# regridder = xesmf.Regridder(input_grid, output_grid, method)
# regridded = regridder( data )
# regridder.clean_weight_file()
# return regridded