clean up + bug fix

pcmdi_metrics/io/xcdat_dataset_io.py

@@ -199,22 +199,24 @@ def da_to_ds(d: Union[xr.Dataset, xr.DataArray], var: str = "variable") -> xr.Da
 
 
 def get_grid(
- ds: xr.Dataset,
+ d: Union[xr.Dataset, xr.DataArray],
 ) -> xr.Dataset:
  """Get grid information
 
  Parameters
  ----------
- ds : xr.Dataset
+ d : Union[xr.Dataset, xr.DataArray]
  xarray dataset to extract grid information that has latitude, longitude, and their bounds included
 
  Returns
  -------
  xr.Dataset
  xarray dataset with grid information
  """
- lat_key = get_latitude_key(ds)
- lon_key = get_longitude_key(ds)
- lat_bnds_key = get_latitude_bounds_key(ds)
- lon_bnds_key = get_longitude_bounds_key(ds)
- return ds[[lat_key, lon_key, lat_bnds_key, lon_bnds_key]]
+ if isinstance(d, xr.DataArray):
+ d = da_to_ds(d, d.name)
+ lat_key = get_latitude_key(d)
+ lon_key = get_longitude_key(d)
+ lat_bnds_key = get_latitude_bounds_key(d)
+ lon_bnds_key = get_longitude_bounds_key(d)
+ return d[[lat_key, lon_key, lat_bnds_key, lon_bnds_key]]

pcmdi_metrics/mean_climate/lib/__init__.py

@@ -18,3 +18,4 @@
 from .create_mean_climate_parser import create_mean_climate_parser # noqa
 from .load_and_regrid import load_and_regrid # noqa
 from .mean_climate_metrics_to_json import mean_climate_metrics_to_json # noqa
+from .pmp_parser import PMPParser # noqa

pcmdi_metrics/utils/grid.py

@@ -3,6 +3,7 @@
 import xcdat as xc
 
 from pcmdi_metrics.io import (
+ get_grid,
  get_latitude_bounds_key,
  get_latitude_key,
  get_longitude_bounds_key,
@@ -143,8 +144,23 @@ def regrid(
  data_var: str,
  target_grid: xr.Dataset,
  regrid_tool: str = "regrid2",
- regrid_method: str = None,
+ regrid_method: str = "bilinear",
+ fill_zero: bool = False,
 ) -> xr.Dataset:
+ """regrid the dataset to a given grid
+
+ Args:
+ ds (xr.Dataset): dataset to regrid
+ data_var (str): variable in the dataset
+ target_grid (xr.Dataset): grid for interpolate to
+ regrid_tool (str, optional): Regrid option: "regrid2" or "xesmf". Defaults to "regrid2".
+ regrid_method (str, optional): Regrid method option that is required for xesmf regridder. Defaults to "bilinear".
+ fill_zero (bool, optional): Fill NaN value with zero if exists. Defaluts to False
+
+ Returns:
+ xr.Dataset: Regridded dataset
+ """
+ target_grid = get_grid(target_grid) # To remove time dimension if exist
  # regrid
  if regrid_tool == "regrid2":
  ds_regridded = ds.regridder.horizontal(data_var, target_grid, tool=regrid_tool)
@@ -154,4 +170,9 @@ def regrid(
  ds_regridded = ds.regridder.horizontal(
  data_var, target_grid, tool=regrid_tool, method=regrid_method
  )
+
+ if fill_zero:
+ ds_regridded = ds_regridded.fillna(0)
+
+ ds_regridded = ds_regridded.bounds.add_missing_bounds() # just in case
  return ds_regridded

pcmdi_metrics/utils/land_sea_mask.py

@@ -5,11 +5,13 @@
 from typing import Union
 
 import numpy as np
+import numpy.ma as ma
 import regionmask
 import xarray as xr
 import xcdat as xc
 
 from pcmdi_metrics import resources
+from pcmdi_metrics.io import get_grid
 
 
 def create_land_sea_mask(
@@ -183,7 +185,24 @@ def apply_landmask(
  data_array = obj[data_var].copy()
 
  # Validate landfrac
+ generateLandSeaMask = False
+ where_method = "xarray"
  if landfrac is None:
+ generateLandSeaMask = True
+ else:
+ data_grid = get_grid(data_array)
+ lf_grid = get_grid(landfrac)
+
+ if data_grid.identical(lf_grid):
+ pass
+ else:
+ if data_grid.equals(lf_grid):
+ pass
+ else:
+ if data_grid.sizes == lf_grid.sizes:
+ where_method = "numpy"
+
+ if generateLandSeaMask:
  landfrac = create_land_sea_mask(data_array)
  warnings.warn(
  "landfrac is not provided thus generated using the 'create_land_sea_mask' function"
@@ -199,7 +218,7 @@ def apply_landmask(
 
  # Convert landfrac to a fraction if it's in percentage form
  if percentage:
- landfrac /= 100.0
+ landfrac = landfrac.copy() / 100.0
 
  # Validate keep_over parameter
  if keep_over not in ["land", "ocean"]:
@@ -208,10 +227,28 @@ def apply_landmask(
  )
 
  # Apply land and ocean masks
- if keep_over == "land":
- data_array = data_array.where(landfrac >= land_criteria)
- elif keep_over == "ocean":
- data_array = data_array.where(landfrac <= ocean_criteria)
+ if where_method == "xarray":
+ if keep_over == "land":
+ data_array = data_array.where(landfrac >= land_criteria)
+ elif keep_over == "ocean":
+ data_array = data_array.where(landfrac <= ocean_criteria)
+
+ elif where_method == "numpy":
+ # Expand data1 to match the shape of landfrac along the time dimension
+ expanded_landfrac = np.expand_dims(landfrac.data, axis=0)
+ expanded_landfrac = np.repeat(
+ expanded_landfrac, data_array.shape[0], axis=0
+ ) # Repeat along the time dimension
+
+ # Mask data based on landfrac
+ if keep_over == "land":
+ data_array.data = ma.masked_where(
+ expanded_landfrac < land_criteria, data_array.data
+ )
+ elif keep_over == "ocean":
+ data_array.data = ma.masked_where(
+ expanded_landfrac > ocean_criteria, data_array.data
+ )
 
  return data_array
 
@@ -233,7 +270,7 @@ def apply_oceanmask(
  )
  return masked_data_array
 
- 
+
 def generate_land_sea_mask__pcmdi(
  target_grid,
  source=None,

pcmdi_metrics/variability_mode/lib/eof_analysis.py

@@ -194,7 +194,7 @@ def linear_regression_on_globe_for_teleconnection(
  return eof_lr, slope, intercept
 
 
-def linear_regression(x, y):
+def linear_regression(x, y, debug=False):
  """
  NOTE: Proceed linear regression
  Input
@@ -211,8 +211,9 @@ def linear_regression(x, y):
  im = y.shape[2]
  jm = y.shape[1]
  y_2d = y.data.reshape(y.shape[0], jm * im)
- print("x.shape:", x.shape)
- print("y_2d.shape:", y_2d.shape)
+ if debug:
+ print("x.shape:", x.shape)
+ print("y_2d.shape:", y_2d.shape)
  # Linear regression
  slope_1d, intercept_1d = np.polyfit(np.array(x), np.array(y_2d), 1)
  # Retreive to variabile from numpy array

pcmdi_metrics/variability_mode/lib/lib_variability_mode.py

@@ -100,11 +100,15 @@ def read_data_in(
  # landmask if required
  if LandMask:
  # Extract SST (land region mask out)
- landfrac = None
+ landfrac = None # by default, generate a land sea mask
  if lf_path is not None:
  if os.path.isfile(lf_path):
- landfrac_ds = xcdat_open(lf_path)
- landfrac = landfrac_ds[var_lf]
+ try:
+ landfrac_ds = xcdat_open(lf_path)
+ landfrac = landfrac_ds[var_lf]
+ except Exception:
+ landfrac = None # if unsuccessful, generate a land sea mask
+
  data_timeseries = apply_landmask(data_timeseries, landfrac=landfrac)
 
  ds_time_subsetted[var_in_data] = data_timeseries

pcmdi_metrics/variability_mode/scripts_pcmdi/post_process_merge_jsons.py

@@ -17,12 +17,13 @@ def main():
  # mips = ['cmip3']
 
  # exps = ['historical', 'amip']
- # exps = ['historical']
- exps = ["amip"]
+ exps = ["historical"]
+ # exps = ["amip"]
  # exps = ['20c3m', 'amip']
  # exps = ['20c3m']
 
- case_id = "v20210119"
+ # case_id = "v20210119"
+ case_id = "v20240227"
 
  syear = 1900
  eyear = 2005

pcmdi_metrics/variability_mode/variability_modes_driver.py

@@ -532,7 +532,8 @@
  # -------------------------------------------------
  for model_path in model_path_list:
  print("model_path:", model_path)
- try:
+ # try:
+ if 1:
  if realization == "*":
  run = re.split("[._]", (model_path.split("/")[-1]).split("."))[
  run_in_modpath
@@ -627,7 +628,11 @@
 
  # Regrid (interpolation, model grid to ref grid)
  model_timeseries_season_regrid = regrid(
- model_timeseries_season, var, ref_grid_global
+ model_timeseries_season,
+ var,
+ ref_grid_global,
+ regrid_tool="xesmf",
+ fill_zero=True,
  )
  # crop to subdomain
  model_timeseries_season_regrid_subdomain = region_subset(
@@ -988,14 +993,14 @@
  run=run,
  cmec_flag=cmec,
  )
-
+ """
  except Exception as err:
  if debug:
  raise
  else:
  print("warning: failed for ", model, run, err)
  pass
-
+ """
 # ========================================================================
 # Dictionary to JSON: collective JSON at the end of model_realization loop
 # ------------------------------------------------------------------------