""" This module includes the definition for a RegionalCase classs. """ # -- Import libraries # -- Import Python Standard Libraries import logging import os import argparse from datetime import datetime # -- 3rd party libraries import numpy as np import xarray as xr # -- import local classes for this script from ctsm.site_and_regional.base_case import BaseCase, USRDAT_DIR from ctsm.site_and_regional.mesh_type import MeshType from ctsm.utils import add_tag_to_filename from ctsm.utils import abort logger = logging.getLogger(__name__) class RegionalCase(BaseCase): """ A class to encapsulate regional cases. ... Attributes ---------- lat1 : float first (left) latitude of a region. lat1 : float second (right) latitude of a region. lon1 : float first (bottom) longitude of a region. lon2 : float second (top) longitude of a region. reg_name: str -- default = None Region's name create_domain : bool flag for creating domain file create_mesh : bool flag for creating mesh file create_surfdata : bool flag for creating surface dataset create_landuse : bool flag for creating landuse file create_datm : bool flag for creating DATM files create_user_mods : bool flag for creating user mods files and folders overwrite : bool flag for over-writing files if they already exist Methods ------- create_tag Create a tag for this region which is either region's name or a combination of bounds of this region lat1-lat2_lon1-lon2 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_domain_at_reg Create domain file at this region create_surfdata_at_reg Create surface dataset at this region extract_mesh_at_reg Extract mesh from the domain dataset created by create_domain_at_reg create_landuse_at_reg Create landuse file at this region write_shell_commands(namelist) write out xml commands to a file for usermods (i.e. shell_commands) for regional settings. """ # pylint: disable=too-many-instance-attributes # the ones we have are useful def __init__( self, lat1, lat2, lon1, lon2, reg_name, create_domain, create_surfdata, create_landuse, create_datm, create_user_mods, create_mesh, out_dir, overwrite, ): """ Initializes RegionalCase with the given arguments. """ super().__init__( create_domain, create_surfdata, create_landuse, create_datm, create_user_mods, overwrite, ) self.lat1 = lat1 self.lat2 = lat2 self.lon1 = lon1 self.lon2 = lon2 self.reg_name = reg_name self.create_mesh = create_mesh self.mesh = None self.out_dir = out_dir self.check_region_bounds() self.create_tag() self.ni = None self.nj = None def create_tag(self): """ Create a tag for a region which is either the region name or the lat1-lat2_lon1-lon2 if the region name does not exist. """ if self.reg_name: self.tag = self.reg_name else: self.tag = "{}-{}_{}-{}".format( 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. """ # specify files fdomain_in = os.path.join(indir, file) fdomain_out = add_tag_to_filename(fdomain_in, self.tag) logger.info("fdomain_in: %s", fdomain_in) logger.info("fdomain_out: %s", os.path.join(self.out_dir, fdomain_out)) logger.info("Creating domain file at region: %s", self.tag) # create 1d coordinate variables to enable sel() method f_in = self.create_1d_coord(fdomain_in, "xc", "yc", "ni", "nj") lat = f_in["lat"] lon = f_in["lon"] # subset longitude and latitude arrays xind = np.where((lon >= self.lon1) & (lon <= self.lon2))[0] yind = np.where((lat >= self.lat1) & (lat <= self.lat2))[0] f_out = f_in.isel(nj=yind, ni=xind) # update attributes self.update_metadata(f_out) f_out.attrs["Created_from"] = fdomain_in # mode 'w' overwrites file wfile = os.path.join(self.out_dir, fdomain_out) self.write_to_netcdf(f_out, wfile) logger.info("Successfully created file (fdomain_out) %s", wfile) f_in.close() f_out.close() if self.create_mesh: mesh_out = os.path.join( self.out_dir, os.path.splitext(fdomain_out)[0] + "_ESMF_UNSTRUCTURED_MESH.nc", ) self.mesh = mesh_out logger.info("creating mesh file from domain file: %s", wfile) ds = xr.open_dataset(wfile, mask_and_scale=False, decode_times=False).transpose() self.extract_mesh_at_reg(ds) def create_surfdata_at_reg(self, indir, file, user_mods_dir, specify_fsurf_out): """ Create surface data file for this RegionalCase class. """ logger.info("Creating surface dataset file at region: %s", self.tag) # specify files fsurf_in = os.path.join(indir, file) if specify_fsurf_out is None: fsurf_out = add_tag_to_filename(fsurf_in, self.tag, replace_res=True) else: fsurf_out = specify_fsurf_out logger.info("fsurf_in: %s", fsurf_in) logger.info("fsurf_out: %s", os.path.join(self.out_dir, fsurf_out)) # create 1d coordinate variables to enable sel() method f_in = self.create_1d_coord(fsurf_in, "LONGXY", "LATIXY", "lsmlon", "lsmlat") lat = f_in["lat"] lon = f_in["lon"] # subset longitude and latitude arrays xind = np.where((lon >= self.lon1) & (lon <= self.lon2))[0] yind = np.where((lat >= self.lat1) & (lat <= self.lat2))[0] f_out = f_in.isel(lsmlat=yind, lsmlon=xind) # update attributes self.update_metadata(f_out) f_out.attrs["Created_from"] = fsurf_in # mode 'w' overwrites file wfile = os.path.join(self.out_dir, fsurf_out) self.write_to_netcdf(f_out, wfile) logger.info("created file (fsurf_out) %s", wfile) f_in.close() f_out.close() # write to user_nl_clm if specified if self.create_user_mods: with open(os.path.join(user_mods_dir, "user_nl_clm"), "a") as nl_clm: line = "fsurdat = '${}'".format(os.path.join(USRDAT_DIR, fsurf_out)) self.write_to_file(line, nl_clm) def extract_mesh_at_reg(self, ds_in): """ Create Mesh from Surface dataset netcdf file. """ logger.info("Creating meshfile for at region: %s", self.tag) lat_name = "yc" lon_name = "xc" lats = ds_in[lat_name].astype(np.float32) lons = ds_in[lon_name].astype(np.float32) self.ni = len(lats.ni) self.nj = len(lats.nj) mask = ds_in["mask"].astype(np.float32) if mask.max() > 1.0 or mask.min() < 0.0: abort("Mask variable is not within 0 to 1") this_mesh = MeshType(lats, lons, mask=mask) this_mesh.calculate_corners() this_mesh.calculate_nodes() this_mesh.create_esmf(self.mesh) def create_landuse_at_reg(self, indir, file, user_mods_dir): """ Create land use data file for this RegionalCase class. """ logger.info("Creating landuse file at region: %s", self.tag) # specify files fluse_in = os.path.join(indir, file) fluse_out = add_tag_to_filename(fluse_in, self.tag, replace_res=True) logger.info("fluse_in: %s", fluse_in) logger.info("fluse_out: %s", os.path.join(self.out_dir, fluse_out)) # create 1d coordinate variables to enable sel() method f_in = self.create_1d_coord(fluse_in, "LONGXY", "LATIXY", "lsmlon", "lsmlat") lat = f_in["lat"] lon = f_in["lon"] # subset longitude and latitude arrays xind = np.where((lon >= self.lon1) & (lon <= self.lon2))[0] yind = np.where((lat >= self.lat1) & (lat <= self.lat2))[0] f_out = f_in.isel(lsmlat=yind, lsmlon=xind) # update attributes self.update_metadata(f_out) f_out.attrs["Created_from"] = fluse_in # mode 'w' overwrites file wfile = os.path.join(self.out_dir, fluse_out) self.write_to_netcdf(f_out, wfile) logger.info("Successfully created file (fluse_out) %s", wfile) f_in.close() f_out.close() if self.create_user_mods: with open(os.path.join(user_mods_dir, "user_nl_clm"), "a") as nl_clm: # line = "landuse = '${}'".format(os.path.join(USRDAT_DIR, fluse_out)) line = "flanduse_timeseries = '${}'".format(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"] node_count = len(f_in["nodeCount"]) node_coords = f_in["nodeCoords"] subset_element = [] cnt = 0 for n in 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 # -- reverse logic # l1 = np.logical_and(nlon >= self.lon1,nlon <= self.lon2) # l2 = np.logical_and(nlat >= self.lat1,nlat <= self.lat2) # if np.any(l1) and np.any(l2): # subset_node.append(n) # conn_dict[n+1] = cnt # cnt+=1 # else: # conn_dict[n+1] = -9999 return node_coords, subset_element, subset_node, conn_dict @staticmethod def write_mesh(f_in, node_coords, subset_element, subset_node, conn_dict, mesh_out): # pylint: disable=unused-argument """ This function writes out the subsetted mesh file. """ corner_pairs = f_in.variables["nodeCoords"][ subset_node, ] 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 n in range(elem_count): for m in range(max_node_dim): ndx = int(elem_conn_index[n, m]) elem_conn_out[n, m] = 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.numElementConn.encoding = {"dtype": np.int32} 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) self.write_to_file("./xmlchange MASK_MESH={}".format(str(str(self.mesh))), nl_file) self.write_to_file("./xmlchange ATM_NX={}".format(str(str(self.ni))), nl_file) self.write_to_file("./xmlchange LND_NX={}".format(str(str(self.ni))), nl_file) self.write_to_file("./xmlchange ATM_NY={}".format(str(str(self.nj))), nl_file) self.write_to_file("./xmlchange LND_NY={}".format(str(str(self.nj))), nl_file)