From 554372d36901c078b53942d2d4c852dd05dd5595 Mon Sep 17 00:00:00 2001 From: Jiwoo Lee Date: Wed, 20 Dec 2023 17:16:45 -0800 Subject: [PATCH] let function takes lon and lat key for data array name in case not standard --- pcmdi_metrics/utils/create_land_sea_mask.py | 21 +++++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/pcmdi_metrics/utils/create_land_sea_mask.py b/pcmdi_metrics/utils/create_land_sea_mask.py index a6758da23..d40c9fd87 100644 --- a/pcmdi_metrics/utils/create_land_sea_mask.py +++ b/pcmdi_metrics/utils/create_land_sea_mask.py @@ -7,7 +7,10 @@ def create_land_sea_mask( - obj: Union[xr.Dataset, xr.DataArray], as_boolean: bool = False + obj: Union[xr.Dataset, xr.DataArray], + lon_key: str = None, + lat_key: str = None, + as_boolean: bool = False, ) -> xr.DataArray: """Generate a land-sea mask (1 for land, 0 for sea) for a given xarray Dataset or DataArray. @@ -15,6 +18,10 @@ def create_land_sea_mask( ---------- obj : Union[xr.Dataset, xr.DataArray] The Dataset or DataArray object. + lon_key : str, optional + Name of DataArray for longitude, by default None + lat_key : str, optional + Name of DataArray for latitude, by default None as_boolean : bool, optional Set mask value to True (land) or False (ocean), by default False, thus 1 (land) and 0 (ocean). @@ -41,14 +48,16 @@ def create_land_sea_mask( land_mask = regionmask.defined_regions.natural_earth_v5_0_0.land_110 # Get the longitude and latitude from the xarray dataset - key_lon = xc.axis.get_dim_keys(obj, axis="X") - key_lat = xc.axis.get_dim_keys(obj, axis="Y") + if lon_key is None: + lon_key = xc.axis.get_dim_keys(obj, axis="X") + if lat_key is None: + lat_key = xc.axis.get_dim_keys(obj, axis="Y") - lon = obj[key_lon] - lat = obj[key_lat] + lon = obj[lon_key] + lat = obj[lat_key] # Mask the land-sea mask to match the dataset's coordinates - land_sea_mask = land_mask.mask(lon, lat) + land_sea_mask = land_mask.mask(lon, lat=lat) if not as_boolean: # Convert the land-sea mask to a boolean mask