Skip to content

Commit

Permalink
Build out the regridding and processing step.
Browse files Browse the repository at this point in the history
- Now packing atm/ocn forcing vars into their own respective forcing
  files

- moved shared utilites to common file

- fixed bug in returning conctenated atm files from file finder
  • Loading branch information
andrewdnolan committed Jun 12, 2024
1 parent 5b96321 commit aee45e9
Show file tree
Hide file tree
Showing 3 changed files with 112 additions and 73 deletions.
2 changes: 1 addition & 1 deletion compass/landice/tests/ismip6_GrIS_forcing/file_finders.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
"""
Expand Down
133 changes: 61 additions & 72 deletions compass/landice/tests/ismip6_GrIS_forcing/process_forcing.py
Original file line number Diff line number Diff line change
@@ -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",
Expand Down Expand Up @@ -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
50 changes: 50 additions & 0 deletions compass/landice/tests/ismip6_GrIS_forcing/utilities.py
Original file line number Diff line number Diff line change
@@ -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)

0 comments on commit aee45e9

Please sign in to comment.