diff --git a/compass/landice/tests/ismip6_GrIS_forcing/file_finders.py b/compass/landice/tests/ismip6_GrIS_forcing/file_finders.py index 780287471..9f41ca0eb 100644 --- a/compass/landice/tests/ismip6_GrIS_forcing/file_finders.py +++ b/compass/landice/tests/ismip6_GrIS_forcing/file_finders.py @@ -99,7 +99,7 @@ def get_filename(self, GCM, scenario, variable, start=2015, end=2100): # with output filename, combine the files to a single self.__combine_files(yearly_files, out_fp) - return out_fp, yearly_files + return out_fp def __find_yearly_files(self, GCM, scenario, variable, start, end): """ diff --git a/compass/landice/tests/ismip6_GrIS_forcing/process_forcing.py b/compass/landice/tests/ismip6_GrIS_forcing/process_forcing.py index b1b4ad346..d8f8b0cf4 100644 --- a/compass/landice/tests/ismip6_GrIS_forcing/process_forcing.py +++ b/compass/landice/tests/ismip6_GrIS_forcing/process_forcing.py @@ -1,16 +1,14 @@ +import os + import xarray as xr from mpas_tools.io import write_netcdf -from mpas_tools.logging import check_call +from compass.landice.tests.ismip6_GrIS_forcing.utilities import ( + add_xtime, + remap_variables, +) from compass.step import Step - -def datetime_2_xtime(ds, var="Time", xtime_fmt="%Y-%m-%d_%H:%M:%S"): - - return xr.apply_ufunc(lambda t: t.strftime(xtime_fmt).ljust(64), - ds[var], vectorize=True, output_dtypes=["S"]) - - renaming_dict = {"thermal_forcing": "ismip6_2dThermalForcing", "basin_runoff": "ismip6Runoff", "aSMB": "sfcMassBal", @@ -55,102 +53,93 @@ def run(self): if expr_name in ["ctrl", "hist"]: continue - ocean_forcing_datastes = [] - # loop over ocean variables seperately - for var in ocean_vars: - var_fp = self.test_case.findForcingFiles(GCM, scenario, var) + atm_fn = f"gis_atm_forcing_{GCM}_{scenario}_{start}--{end}.nc" + self.process_variables(GCM, scenario, start, end, atmosphere_vars, + atm_fn) + + ocn_fn = f"gis_ocn_forcing_{GCM}_{scenario}_{start}--{end}.nc" + self.process_variables(GCM, scenario, start, end, ocean_vars, + ocn_fn) + + def process_variables(self, GCM, scenario, start, end, + variables, output_fn): - ocean_forcing_datastes.append( - self.process_variable(GCM, scenario, var, - var_fp, start, end)) - ocean_forcing_ds = xr.combined_by_coords( - ocean_forcing_datastes, combine_attrs="drop_conflicts") + forcing_datastes = [] - write_netcdf(ocean_forcing_ds, "ocean_forcing_test.nc") + for var in variables: + var_fp = self.test_case.findForcingFiles(GCM, scenario, var) - # loop over atmosphere variables - for var in atmosphere_vars: - # this pretty hacky with the tuple, fix the file finder - var_fp, _ = self.test_case.findForcingFiles(GCM, scenario, var) + ds = self.process_variable(GCM, scenario, var, var_fp, start, end) - self.process_variable(GCM, scenario, var, var_fp, start, end) + forcing_datastes.append(ds) + + forcing_ds = xr.merge(forcing_datastes) + + write_netcdf(forcing_ds, output_fn) def process_variable(self, GCM, scenario, var, forcing_fp, start, end): - # create the remapped file name - remapped_fp = f"gis_{var}_{GCM}_{scenario}_{start}-{end}.nc" + # + config = self.config + # + renamed_var = renaming_dict[var] + + # create the remapped file name, using original variable name + remapped_fn = f"remapped_{GCM}_{scenario}_{var}_{start}-{end}.nc" + # follow the directory structure of the concatenated source files + remapped_fp = os.path.join(f"{GCM}-{scenario}/{var}", remapped_fn) # remap the forcing file - self.remap_variable(forcing_fp, - remapped_fp, - self.test_case.ismip6_2_mali_weights) + remap_variables(forcing_fp, + remapped_fp, + self.test_case.ismip6_2_mali_weights, + variables=[var], + logger=self.logger) # Now that forcing has been regridded onto the MALI grid, we can drop # ISMIP6 grid variables. Special note: the char field `mapping` # particularily causes problem with `xtime` (40byte vs. 64byte chars) vars_2_drop = ["time_bounds", "lat_vertices", "lon_vertices", "area", - "mapping", "lat", "lon"] + "mapping", "lat", "lon", "polar_stereographic"] # open the remapped file for post processing remapped_ds = xr.open_dataset(remapped_fp, drop_variables=vars_2_drop, use_cftime=True) - # convert time to xtime - remapped_ds["xtime"] = datetime_2_xtime(remapped_ds, var="time") + # mask of desired time indices + mask = remapped_ds.time.dt.year.isin(range(start, end)) + # drop the non-desired time indices from remapped dataset + remapped_ds = remapped_ds.where(mask, drop=True) + # add xtime variable based on `time` + remapped_ds["xtime"] = add_xtime(remapped_ds, var="time") # rename the variable/dimensions to match MPAS/MALI conventions remapped_ds = remapped_ds.rename({"ncol": "nCells", "time": "Time", - var: renaming_dict[var]}) + var: renamed_var}) + + # drop the unneeded attributes from the dataset + for _var in remapped_ds: + remapped_ds[_var].attrs.pop("grid_mapping", None) + remapped_ds[_var].attrs.pop("cell_measures", None) # SMB is not a var in ISMIP6 forcing files, need to use `SMB_ref` # and the processed `aSMB` to produce a `SMB` forcing - if var == "aSMB": + if renamed_var == "sfcMassBal": # open the reference climatology file smb_ref = xr.open_dataset(self.test_case.smb_ref_climatology) # squeeze the empty time dimension so that broadcasting works smb_ref = smb_ref.squeeze("Time") # add climatology to anomoly to make full forcing field - remapped_ds[renaming_dict[var]] += smb_ref["sfcMassBal"] - - # need to process surfaceAirTemperature - if var == "surfaceAirTemperature": - pass - - # write_netcdf(remapped_ds, remapped_fp) - return remapped_ds - - def remap_variable(self, input_file, output_file, weights_file): - """ - """ - - # remap the forcing file onto the MALI mesh - args = ["ncremap", - "-i", input_file, - "-o", output_file, - "-m", weights_file] - - check_call(args, logger=self.logger) - - def rename_variable_and_trim_dataset(self, ds, var_src, var_dest): + remapped_ds[renamed_var] += smb_ref[renamed_var] - # drop unnecessary variables - ds = ds.drop_vars(["lat_vertices", "area", "lon_vertices", - "lat", "lon"]) - - ds = ds.rename({"ncol": "nCells", - "time": "Time", - var_src: var_dest}) - - # need to make `Time` the unlimited dimension, which also prevents - # `time` dimension for being added back to the dataset - # ds.encoding["unlimited_dims"] = {"Time"} - - return ds - - def create_xtime(self, ds): + if renamed_var == "surfaceAirTemperature": + # read the mesh path/name from the config file + mali_mesh_fp = config.get("ISMIP6_GrIS_Forcing", "MALI_mesh_fp") + # squeeze the empty time dimension so that broadcasting works + mali_ds = xr.open_dataset(mali_mesh_fp).squeeze("Time") - ds["xtime"] = datetime_2_xtime(ds, var="Time") + remapped_ds[renamed_var] += mali_ds[renamed_var] - return ds + return remapped_ds diff --git a/compass/landice/tests/ismip6_GrIS_forcing/utilities.py b/compass/landice/tests/ismip6_GrIS_forcing/utilities.py new file mode 100644 index 000000000..b553ef8eb --- /dev/null +++ b/compass/landice/tests/ismip6_GrIS_forcing/utilities.py @@ -0,0 +1,50 @@ +import xarray as xr +from mpas_tools.logging import check_call + + +def _datetime_2_xtime(ds, var="Time", xtime_fmt="%Y-%m-%d_%H:%M:%S"): + + return xr.apply_ufunc(lambda t: t.strftime(xtime_fmt).ljust(64), + ds[var], vectorize=True, output_dtypes=["S"]) + + +def add_xtime(ds, var="Time"): + """ + ds["xtime"] = add_xtime(ds) + """ + + # ensure time variable has been properly parsed as a datetime object + if not hasattr(ds[var], "dt"): + raise TypeError(f"The {var} variable passed has not been parsed as" + " a datetime object, so conversion to xtime string" + " will not work.\n\nTry using ds = xr.open_dataset" + "(..., use_cftime=True).") + + # xtime related attributes + attrs = {"units": "unitless", + "long_name": "model time, with format \'YYYY-MM-DD_HH:MM:SS\'"} + + # compute xtime dataarray from time variable passed + xtime = _datetime_2_xtime(ds, var=var) + + # update attributes of dataarray + xtime.attrs = attrs + + return xtime + + +def remap_variables(in_fp, out_fp, weights_fp, variables=None, logger=None): + """ + """ + + args = ["ncremap", + "-i", in_fp, + "-o", out_fp, + "-m", weights_fp] + + if variables and not isinstance(variables, list): + raise TypeError("`variables` kwarg must be a list of strings.") + elif variables: + args += ["-v"] + variables + + check_call(args, logger=logger)