diff --git a/python/ctsm/site_and_regional/regional_case.py b/python/ctsm/site_and_regional/regional_case.py index 57bb8474f4..92892b9340 100644 --- a/python/ctsm/site_and_regional/regional_case.py +++ b/python/ctsm/site_and_regional/regional_case.py @@ -5,9 +5,13 @@ # -- Import Python Standard Libraries import logging import os +import argparse # -- 3rd party libraries import numpy as np +import xarray as xr +from tqdm import tqdm +from datetime import datetime # -- import local classes for this script from ctsm.site_and_regional.base_case import BaseCase, USRDAT_DIR @@ -49,6 +53,15 @@ class RegionalCase(BaseCase): Methods ------- + check_region_bounds + Check for the regional bounds + + check_region_lons + Check for the regional lons + + check_region_lats + Check for the regional lats + create_tag Create a tag for this region which is either region's name or a combination of bounds of this @@ -77,6 +90,7 @@ def __init__( create_landuse, create_datm, create_user_mods, + create_mesh, out_dir, overwrite, ): @@ -96,7 +110,9 @@ def __init__( self.lon1 = lon1 self.lon2 = lon2 self.reg_name = reg_name + self.create_mesh = create_mesh self.out_dir = out_dir + self.check_region_bounds() self.create_tag() def create_tag(self): @@ -112,6 +128,50 @@ def create_tag(self): str(self.lon1), str(self.lon2), str(self.lat1), str(self.lat2) ) + def check_region_bounds (self): + """ + Check for the regional bounds + """ + self.check_region_lons() + self.check_region_lats() + + def check_region_lons (self): + """ + Check for the regional lon bounds + """ + if self.lon1 >= self.lon2: + err_msg = """ + \n + ERROR: lon1 is bigger than lon2. + lon1 points to the westernmost longitude of the region. {} + lon2 points to the easternmost longitude of the region. {} + Please make sure lon1 is smaller than lon2. + + Please note that if longitude in -180-0, the code automatically + convert it to 0-360. + """.format( + self.lon1, self.lon2 + ) + raise argparse.ArgumentTypeError(err_msg) + + def check_region_lats (self): + """ + Check for the regional lat bound + """ + if self.lat1 >= self.lat2: + err_msg = """ + \n + ERROR: lat1 is bigger than lat2. + lat1 points to the westernmost longitude of the region. {} + lat2 points to the easternmost longitude of the region. {} + Please make sure lat1 is smaller than lat2. + + """.format( + self.lat1, self.lat2 + ) + raise argparse.ArgumentTypeError(err_msg) + + def create_domain_at_reg(self, indir, file): """ Create domain file for this RegionalCase class. @@ -226,3 +286,191 @@ def create_landuse_at_reg(self, indir, file, user_mods_dir): os.path.join(USRDAT_DIR, fluse_out) ) self.write_to_file(line, nl_clm) + + + def create_mesh_at_reg(self, mesh_dir, mesh_surf): + """ + Create a mesh subsetted for the RegionalCase class. + """ + logger.info( + "----------------------------------------------------------------------" + ) + logger.info( + "Subsetting mesh file for region: %s", + self.tag + ) + + today = datetime.today() + today_string = today.strftime("%y%m%d") + + + mesh_in = os.path.join(mesh_dir, mesh_surf) + mesh_out = os.path.join(self.out_dir, os.path.splitext(mesh_surf)[0]+'_'+self.tag+'_c'+today_string+'.nc') + + logger.info("mesh_in : %s", mesh_in) + logger.info("mesh_out : %s", mesh_out) + + self.mesh = mesh_out + + node_coords, subset_element, subset_node, conn_dict = self.subset_mesh_at_reg(mesh_in) + + f_in = xr.open_dataset (mesh_in) + self.write_mesh (f_in, node_coords, subset_element, subset_node, conn_dict, mesh_out) + + + def subset_mesh_at_reg (self, mesh_in): + """ + This function subsets the mesh based on lat and lon bounds given by RegionalCase class. + """ + f_in = xr.open_dataset (mesh_in) + elem_count = len (f_in['elementCount']) + elem_conn = f_in['elementConn'] + num_elem_conn = f_in['numElementConn'] + center_coords = f_in['centerCoords'] + node_count = len (f_in['nodeCount']) + node_coords = f_in['nodeCoords'] + + subset_element = [] + cnt = 0 + + for n in tqdm(range(elem_count)): + endx = elem_conn[n,:num_elem_conn[n].values].values + endx[:,] -= 1# convert to zero based index + endx = [int(xi) for xi in endx] + + nlon = node_coords[endx,0].values + nlat = node_coords[endx,1].values + + l1 = np.logical_or(nlon <= self.lon1,nlon >= self.lon2) + l2 = np.logical_or(nlat <= self.lat1,nlat >= self.lat2) + if np.any(np.logical_or(l1,l2)): + pass + else: + subset_element.append(n) + cnt+=1 + + + subset_node = [] + conn_dict = {} + cnt = 1 + + for n in range(node_count): + nlon = node_coords[n,0].values + nlat = node_coords[n,1].values + + l1 = np.logical_or(nlon <= self.lon1,nlon >= self.lon2) + l2 = np.logical_or(nlat <= self.lat1,nlat >= self.lat2) + if np.logical_or(l1,l2): + conn_dict[n+1] = -9999 + else: + subset_node.append(n) + conn_dict[n+1] = cnt + cnt+=1 + return node_coords, subset_element, subset_node, conn_dict + + + + def write_mesh (self, f_in, node_coords, subset_element, subset_node, conn_dict, mesh_out): + """ + This function writes out the subsetted mesh file. + """ + corner_pairs = f_in.variables['nodeCoords'][subset_node,] + + dimensions = f_in.dims + variables = f_in.variables + global_attributes = f_in.attrs + + + max_node_dim = len(f_in['maxNodePElement']) + + elem_count = len(subset_element) + elem_conn_out = np.empty(shape=[elem_count, max_node_dim]) + elem_conn_index = f_in.variables['elementConn'][subset_element,] + + for ni in range(elem_count): + for mi in range(max_node_dim): + ndx = int (elem_conn_index[ni,mi]) + elem_conn_out[ni,mi] = conn_dict[ndx] + + + num_elem_conn_out = np.empty(shape=[elem_count,]) + num_elem_conn_out[:] = f_in.variables['numElementConn'][subset_element,] + + center_coords_out = np.empty(shape=[elem_count,2]) + center_coords_out[:,:]=f_in.variables['centerCoords'][subset_element,:] + + if 'elementMask' in variables: + elem_mask_out = np.empty(shape=[elem_count,]) + elem_mask_out[:]=f_in.variables['elementMask'][subset_element,] + + if 'elementArea' in variables: + elem_area_out = np.empty(shape=[elem_count,]) + elem_area_out[:]=f_in.variables['elementArea'][subset_element,] + + # -- create output dataset + f_out = xr.Dataset() + + f_out['nodeCoords'] = xr.DataArray(corner_pairs, + dims=('nodeCount', 'coordDim'), + attrs={'units': 'degrees'}) + + f_out['elementConn'] = xr.DataArray(elem_conn_out, + dims=('elementCount', 'maxNodePElement'), + attrs={'long_name': 'Node indices that define the element connectivity'}) + f_out.elementConn.encoding = {'dtype': np.int32} + + f_out['numElementConn'] = xr.DataArray(num_elem_conn_out, + dims=('elementCount'), + attrs={'long_name': 'Number of nodes per element'}) + + f_out['centerCoords'] = xr.DataArray(center_coords_out, + dims=('elementCount', 'coordDim'), + attrs={'units': 'degrees'}) + + + #-- add mask + if 'elementMask' in variables: + f_out['elementMask'] = xr.DataArray(elem_mask_out, + dims=('elementCount'), + attrs={'units': 'unitless'}) + f_out.elementMask.encoding = {'dtype': np.int32} + + if 'elementArea' in variables: + f_out['elementArea'] = xr.DataArray(elem_area_out, + dims=('elementCount'), + attrs={'units': 'unitless'}) + + #-- setting fill values + for var in variables: + if '_FillValue' in f_in[var].encoding: + f_out[var].encoding['_FillValue'] = f_in[var].encoding['_FillValue'] + else: + f_out[var].encoding['_FillValue'] = None + + #-- add global attributes + for attr in global_attributes: + if attr != 'timeGenerated': + f_out.attrs[attr] = global_attributes[attr] + + f_out.attrs = {'title': 'ESMF unstructured grid file for a region', + 'created_by': 'subset_data', + 'date_created': '{}'.format(datetime.now()), + } + + f_out.to_netcdf(mesh_out) + logger.info("Successfully created file (mesh_out) %s", mesh_out) + + def write_shell_commands(self, namelist): + """ + writes out xml commands commands to a file (i.e. shell_commands) for single-point runs + """ + # write_to_file surrounds text with newlines + with open(namelist, "w") as nl_file: + self.write_to_file( + "# Change below line if you move the subset data directory", nl_file + ) + self.write_to_file( + "./xmlchange {}={}".format(USRDAT_DIR, self.out_dir), nl_file + ) + self.write_to_file("./xmlchange ATM_DOMAIN_MESH={}".format(str(self.mesh)), nl_file) + self.write_to_file("./xmlchange LND_DOMAIN_MESH={}".format(str(self.mesh)), nl_file) diff --git a/python/ctsm/subset_data.py b/python/ctsm/subset_data.py index a1b8d21c15..7230a72a35 100644 --- a/python/ctsm/subset_data.py +++ b/python/ctsm/subset_data.py @@ -170,7 +170,7 @@ def get_parser(): # -- region-specific parser options rg_parser.add_argument( "--lat1", - help="Region start latitude. [default: %(default)s]", + help="Region southernmost latitude. [default: %(default)s]", action="store", dest="lat1", required=False, @@ -179,7 +179,7 @@ def get_parser(): ) rg_parser.add_argument( "--lat2", - help="Region end latitude. [default: %(default)s]", + help="Region northernmost latitude. [default: %(default)s]", action="store", dest="lat2", required=False, @@ -188,7 +188,7 @@ def get_parser(): ) rg_parser.add_argument( "--lon1", - help="Region start longitude. [default: %(default)s]", + help="Region westernmost longitude. [default: %(default)s]", action="store", dest="lon1", required=False, @@ -197,7 +197,7 @@ def get_parser(): ) rg_parser.add_argument( "--lon2", - help="Region end longitude. [default: %(default)s]", + help="Region easternmost longitude. [default: %(default)s]", action="store", dest="lon2", required=False, @@ -398,10 +398,13 @@ def setup_files(args, defaults, cesmroot): 'fdomain_in': defaults.get("domain", "file"), 'fsurf_dir': os.path.join(defaults.get("main", "clmforcingindir"), os.path.join(defaults.get("surfdat", "dir"))), + 'mesh_dir': os.path.join(defaults.get("main", "clmforcingindir"), + defaults.get("surfdat", "mesh_dir")), 'fluse_dir': os.path.join(defaults.get("main", "clmforcingindir"), os.path.join(defaults.get("landuse", "dir"))), 'fsurf_in': fsurf_in, 'fluse_in': fluse_in, + 'mesh_surf' : defaults.get("surfdat","mesh_surf"), 'datm_tuple': DatmFiles(dir_input_datm, dir_output_datm, defaults.get(datm_type, "domain"), @@ -415,7 +418,6 @@ def setup_files(args, defaults, cesmroot): defaults.get(datm_type, 'precname'), defaults.get(datm_type, 'tpqwname')) } - return file_dict @@ -502,6 +504,7 @@ def subset_region(args, file_dict: dict): create_landuse = args.create_landuse, create_datm = args.create_datm, create_user_mods = args.create_user_mods, + create_mesh = args.create_mesh, out_dir = args.out_dir, overwrite = args.overwrite, ) @@ -517,11 +520,29 @@ def subset_region(args, file_dict: dict): region.create_surfdata_at_reg(file_dict["fsurf_dir"], file_dict["fsurf_in"], args.user_mods_dir) + if region.create_mesh: + region.create_mesh_at_reg (file_dict["mesh_dir"], file_dict["mesh_surf"]) + # -- Create CTSM transient landuse data file if region.create_landuse: region.create_landuse_at_reg(file_dict["fluse_dir"], file_dict["fluse_in"], args.user_mods_dir) + + # -- Write shell commands + if region.create_user_mods: + if not region.create_mesh: + err_msg = """ + \n + ERROR: For regional cases, you can not create user_mods + without creating the mesh file. + + Please rerun the script adding --create-mesh to subset the mesh file. + """ + raise argparse.ArgumentTypeError(err_msg) + + region.write_shell_commands(os.path.join(args.user_mods_dir, "shell_commands")) + logger.info("Successfully ran script for a regional case.") diff --git a/tools/site_and_regional/default_data.cfg b/tools/site_and_regional/default_data.cfg index f689c99044..630013d502 100644 --- a/tools/site_and_regional/default_data.cfg +++ b/tools/site_and_regional/default_data.cfg @@ -18,6 +18,8 @@ tpqwname = CLMGSWP3v1.TPQW dir = lnd/clm2/surfdata_map/release-clm5.0.18 surfdat_16pft = surfdata_0.9x1.25_hist_16pfts_Irrig_CMIP6_simyr2000_c190214.nc surfdat_78pft = surfdata_0.9x1.25_hist_78pfts_CMIP6_simyr2000_c190214.nc +mesh_dir = share/meshes/ +mesh_surf = fv0.9x1.25_141008_ESMFmesh.nc [landuse] dir = lnd/clm2/surfdata_map/release-clm5.0.18