diff --git a/unseen/fileio.py b/unseen/fileio.py index cbf418c..6156ae5 100644 --- a/unseen/fileio.py +++ b/unseen/fileio.py @@ -27,6 +27,8 @@ def open_dataset( chunks=None, metadata_file=None, variables=[], + regrid_name=None, + regrid_method="conservative", point_selection=None, lat_bnds=None, lon_bnds=None, @@ -72,6 +74,12 @@ def open_dataset( YAML file path specifying required file metadata changes variables : list, optional Subset of variables of interest + regrid_name : str, optional + Name of grid for sptial regridding in AUSXXXi format. + e.g. AUS005i is a 0.05 x 0.05 lat/lon grid + e.g. AUS300i is a 3.00 x 3.00 lat/lon grid + regrid_method : {"conservative", "bilinear", "nearest_s2d", "nearest_d2s"}, default "conservative" + Name of grid for sptial regridding in AUSXXXi format. point_selection : list, optional Point coordinates: [lat, lon] lat_bnds : list, optional @@ -182,6 +190,11 @@ def open_dataset( for var, target_units in units.items(): ds[var] = general_utils.convert_units(ds[var], target_units) + # Spatial regridding + if regrid_name: + new_grid = general_utils.create_grid(regrid_name) + ds = general_utils.regrid(ds, new_grid, method=regrid_method) + # Spatial subsetting and aggregation if point_selection: ds = spatial_selection.select_point( @@ -709,6 +722,19 @@ def _parse_command_line(): choices=["month"], help="Anomaly can monthly (month) or all times (none)", ) + parser.add_argument( + "--regrid_name", + type=str, + default=None, + help="Name of grid for sptial regridding in AUSXXXi format", + ) + parser.add_argument( + "--regrid_method", + type=str, + default="conservative", + choices=["conservative", "bilinear", "nearest_s2d", "nearest_d2s"], + help="Regridding method", + ) parser.add_argument( "--point_selection", type=float, @@ -855,6 +881,8 @@ def _main(): "chunks": args.input_chunks, "metadata_file": args.metadata_file, "variables": args.variables, + "regrid_name": args.regrid_name, + "regrid_method": args.regrid_method, "point_selection": args.point_selection, "lat_bnds": args.lat_bnds, "lon_bnds": args.lon_bnds, diff --git a/unseen/general_utils.py b/unseen/general_utils.py index 6c84895..bb4b09a 100644 --- a/unseen/general_utils.py +++ b/unseen/general_utils.py @@ -4,6 +4,7 @@ from matplotlib.ticker import AutoMinorLocator import matplotlib.pyplot as plt import numpy as np +import xarray as xr from xarray import Dataset from xclim.core import units from xesmf import Regridder @@ -76,6 +77,67 @@ def convert_units(da, target_units): return da +def create_grid(grid_name): + """Create a regular lat/lon grid. + + Parameters + ---------- + grid_name : str + Name of the grid. + + Returns + ------- + ds_grid : xarray.Dataset + Dataset with desired lat/lon axes. + + Notes + ----- + The only valid grids are in the AUSXXi format. e.g: + - AUS005i is a 0.05 x 0.05 grid across Australia. + - AUS050i is a 0.50 x 0.50 grid across Australia. + - AUS300i is a 3.00 x 3.00 grid across Australia. + """ + + assert len(grid_name) == 7, "grid_name must be AUSXXXi format" + assert grid_name[0:3] == "AUS", "AUSXXXi grids only" + # AGCD bounds + south_lat = -44.5 + north_lat = -10 + west_lon = 112 + east_lon = 156.25 + + step_start = grid_name[3] + step_end = grid_name[4:6] + step = float(f"{step_start}.{step_end}") + + ds_grid = xr.Dataset( + { + "lat": ( + ["lat"], + np.round(np.arange(south_lat, north_lat + step, step), decimals=2), + ), + "lon": ( + ["lon"], + np.round(np.arange(west_lon, east_lon + step, step), decimals=2), + ), + } + ) + ds_grid["lat"].attrs = { + "standard_name": "latitude", + "long_name": "latitude", + "units": "degrees_north", + "axis": "Y", + } + ds_grid["lon"].attrs = { + "standard_name": "longitude", + "long_name": "longitude", + "units": "degrees_east", + "axis": "X", + } + + return ds_grid + + def regrid(ds, ds_grid, method="conservative", **kwargs): """Regrid `ds` to the grid of `ds_grid` using xESMF.