diff --git a/compass/ocean/tests/global_ocean/__init__.py b/compass/ocean/tests/global_ocean/__init__.py index 3d266a9cc9..f07959f5fa 100644 --- a/compass/ocean/tests/global_ocean/__init__.py +++ b/compass/ocean/tests/global_ocean/__init__.py @@ -13,6 +13,7 @@ from compass.ocean.tests.global_ocean.performance_test import PerformanceTest from compass.ocean.tests.global_ocean.restart_test import RestartTest from compass.ocean.tests.global_ocean.threads_test import ThreadsTest +from compass.ocean.tests.global_ocean.wave_mesh import WaveMesh from compass.testgroup import TestGroup @@ -62,10 +63,15 @@ def __init__(self, mpas_core): # A test case for making E3SM support files from an existing mesh self.add_test_case(FilesForE3SM(test_group=self)) + # Test cases for creating waves mesh + self.add_test_case(WaveMesh(test_group=self)) + self._add_tests(mesh_names=['Icos'], include_wave_mesh=True) + def _add_tests(self, mesh_names, high_res_topography=True, include_rk4=False, include_regression=False, include_phc=False, - include_en4_1900=False): + include_en4_1900=False, + include_wave_mesh=False): """ Add test cases for the given mesh(es) """ default_ic = 'WOA23' @@ -116,10 +122,10 @@ def _add_tests(self, mesh_names, high_res_topography=True, time_integrator=time_integrator) self.add_test_case(dynamic_adjustment_test) - self.add_test_case( - FilesForE3SM( - test_group=self, mesh=mesh_test, init=init_test, - dynamic_adjustment=dynamic_adjustment_test)) + e3sm_files_test = FilesForE3SM( + test_group=self, mesh=mesh_test, init=init_test, + dynamic_adjustment=dynamic_adjustment_test) + self.add_test_case(e3sm_files_test) if include_rk4: time_integrator = 'RK4' @@ -167,3 +173,8 @@ def _add_tests(self, mesh_names, high_res_topography=True, FilesForE3SM( test_group=self, mesh=mesh_test, init=init_test, dynamic_adjustment=dynamic_adjustment_test)) + + if include_wave_mesh: + self.add_test_case( + WaveMesh(test_group=self, ocean_mesh=mesh_test, + ocean_init=init_test)) diff --git a/compass/ocean/tests/global_ocean/wave_mesh/__init__.py b/compass/ocean/tests/global_ocean/wave_mesh/__init__.py new file mode 100644 index 0000000000..037e08b99a --- /dev/null +++ b/compass/ocean/tests/global_ocean/wave_mesh/__init__.py @@ -0,0 +1,99 @@ +# from compass.ocean.tests.global_ocean.metadata import ( +# get_author_and_email_from_git, +# ) +from compass.ocean.tests.global_ocean.wave_mesh.base_mesh import WavesBaseMesh +from compass.ocean.tests.global_ocean.wave_mesh.cull_mesh import WavesCullMesh +from compass.ocean.tests.global_ocean.wave_mesh.remap_files import ( + WavesRemapFiles, +) +from compass.ocean.tests.global_ocean.wave_mesh.rotate_mesh import ( + WavesRotateMesh, +) +from compass.ocean.tests.global_ocean.wave_mesh.scrip_file import ( + WavesScripFile, +) +from compass.ocean.tests.global_ocean.wave_mesh.uost_files import ( + WavesUostFiles, +) +from compass.testcase import TestCase + +# from compass.validate import compare_variables + + +class WaveMesh(TestCase): + """ + A test case for creating a global MPAS-Ocean mesh + + Attributes + ---------- + """ + def __init__(self, test_group, ocean_mesh=None, ocean_init=None): + """ + Create test case for creating a global MPAS-Ocean mesh + + Parameters + ---------- + test_group : compass.ocean.tests.global_ocean.GlobalOcean + The global ocean test group that this test case belongs to + + """ + name = 'wave_mesh' + if ocean_mesh is not None: + subdir = f'{ocean_mesh.mesh_name}/{name}' + else: + subdir = name + + super().__init__(test_group=test_group, name=name, subdir=subdir) + + self.package = 'compass.ocean.tests.global_ocean.wave_mesh' + self.mesh_config_filename = 'wave_mesh.cfg' + + base_mesh_step = WavesBaseMesh(test_case=self, + ocean_base_mesh=ocean_mesh, + ocean_culled_mesh=ocean_init) + self.add_step(base_mesh_step) + + cull_mesh_step = WavesCullMesh(test_case=self, + ocean_mesh=ocean_init, + wave_base_mesh=base_mesh_step) + self.add_step(cull_mesh_step) + + scrip_file_step = WavesScripFile(test_case=self, + wave_culled_mesh=cull_mesh_step, + ocean_mesh=ocean_init) + self.add_step(scrip_file_step) + + remap_file_step = WavesRemapFiles(test_case=self, + wave_scrip=scrip_file_step) + self.add_step(remap_file_step) + + rotate_mesh_step = WavesRotateMesh(test_case=self, + wave_culled_mesh=cull_mesh_step) + self.add_step(rotate_mesh_step) + + uost_file_step = WavesUostFiles(test_case=self, + wave_culled_mesh=cull_mesh_step, + wave_rotate_mesh=rotate_mesh_step) + self.add_step(uost_file_step) + + def configure(self, config=None): + """ + Modify the configuration options for this test case + + config : compass.config.CompassConfigParser, optional + Configuration options to update if not those for this test case + """ + if config is None: + config = self.config + config.add_from_package('compass.mesh', 'mesh.cfg', exception=True) + + # get_author_and_email_from_git(config) + +# def validate(self): +# """ +# Test cases can override this method to perform validation of variables +# and timers +# """ +# variables = ['xCell', 'yCell', 'zCell'] +# compare_variables(test_case=self, variables=variables, +# filename1='cull_mesh/culled_mesh.nc') diff --git a/compass/ocean/tests/global_ocean/wave_mesh/base_mesh.py b/compass/ocean/tests/global_ocean/wave_mesh/base_mesh.py new file mode 100644 index 0000000000..b6d9ec76e5 --- /dev/null +++ b/compass/ocean/tests/global_ocean/wave_mesh/base_mesh.py @@ -0,0 +1,390 @@ +import timeit + +import cartopy.crs as ccrs +import cartopy.feature as cfeature +import jigsawpy +import matplotlib.pyplot as plt +import numpy as np +from mpas_tools.cime.constants import constants +from mpas_tools.mesh.cull import write_map_culled_to_base +from netCDF4 import Dataset +from scipy import interpolate, spatial + +from compass.mesh import QuasiUniformSphericalMeshStep + + +class WavesBaseMesh(QuasiUniformSphericalMeshStep): + """ + A step for creating wave mesh based on an ocean mesh + """ + def __init__(self, test_case, ocean_base_mesh, ocean_culled_mesh, + name='base_mesh', subdir=None): + + super().__init__(test_case=test_case, name=name, subdir=subdir, + cell_width=None) + + self.ocean_base_mesh = ocean_base_mesh + self.ocean_culled_mesh = ocean_culled_mesh + + def setup(self): + + super().setup() + self.opts.init_file = 'init.msh' + + # Set links to base mesh + if self.ocean_base_mesh is not None: + mesh_path = self.ocean_base_mesh.steps['base_mesh'].path + ocean_base_mesh_path = f'{mesh_path}/base_mesh.nc' + else: + if self.config.has_option('wave_mesh', 'ocean_base_mesh'): + ocean_base_mesh_path = self.config.get('wave_mesh', + 'ocean_base_mesh') + else: + raise ValueError('ocean_base_mesh option required ' + 'in wave_mesh section in cfg file') + + self.add_input_file( + filename='ocean_base_mesh.nc', + work_dir_target=ocean_base_mesh_path) + + # Set links to culled mesh + if self.ocean_culled_mesh is not None: + mesh_path = self.ocean_culled_mesh.steps['initial_state'].path + ocean_culled_mesh_path = f'{mesh_path}/initial_state.nc' + else: + if self.config.has_option('wave_mesh', 'ocean_culled_mesh'): + ocean_culled_mesh_path = self.config.get('wave_mesh', + 'ocean_culled_mesh') + else: + raise ValueError('ocean_culled_mesh option required ' + 'in wave_mesh section in cfg file') + + self.add_input_file( + filename='ocean_culled_mesh.nc', + work_dir_target=ocean_culled_mesh_path) + + def build_cell_width_lat_lon(self): + """ + Create cell width array for this mesh on a regular latitude-longitude + grid + + Returns + ------- + cellWidth : numpy.array + m x n array of cell width in km + + lon : numpy.array + longitude in degrees (length n and between -180 and 180) + + lat : numpy.array + longitude in degrees (length m and between -90 and 90) + """ + km = 1000.0 + + lon_min = -180.0 + lon_max = 180.0 + dlon = self.config.getfloat('wave_mesh', 'hfun_grid_spacing') + nlon = int((lon_max - lon_min) / dlon) + 1 + lat_min = -90.0 + lat_max = 90.0 + nlat = int((lat_max - lat_min) / dlon) + 1 + + xlon = np.linspace(lon_min, lon_max, nlon) + ylat = np.linspace(lat_min, lat_max, nlat) + + earth_radius = constants['SHR_CONST_REARTH'] / km + cell_width = self.cell_widthVsLatLon(xlon, ylat, + earth_radius, + 'ocean_culled_mesh.nc') + cell_width = cell_width / km + + self.create_initial_points('ocean_base_mesh.nc', + 'ocean_culled_mesh.nc', + xlon, ylat, cell_width, + earth_radius, self.opts.init_file) + + hfun_slope_lim = self.config.getfloat('wave_mesh', 'hfun_slope_lim') + cell_width = self.limit_spacing_gradient(xlon, ylat, cell_width, + earth_radius, hfun_slope_lim) + + return cell_width, xlon, ylat + + def cell_widthVsLatLon(self, lon, lat, sphere_radius, ocean_mesh): + + config = self.config + + depth_threshold_refined = config.getfloat('wave_mesh', + 'depth_threshold_refined') + dist_threshold_refined = config.getfloat('wave_mesh', + 'distance_threshold_refined') + depth_threshold_global = config.getfloat('wave_mesh', + 'depth_threshold_global') + dist_threshold_global = config.getfloat('wave_mesh', + 'distance_threshold_global') + refined_res = config.getfloat('wave_mesh', 'refined_res') + maxres = config.getfloat('wave_mesh', 'maxres') + + # Create structrued grid points + nlon = lon.size + nlat = lat.size + Lon_grd, Lat_grd = np.meshgrid(np.radians(lon), np.radians(lat)) + xy_pts = np.vstack((Lon_grd.ravel(), Lat_grd.ravel())).T + + # Create structured grid points and background cell_with array + cell_width = np.zeros(Lon_grd.shape) + maxres + + # Get ocean mesh variables + nc_file = Dataset(ocean_mesh, 'r') + areaCell = nc_file.variables['areaCell'][:] + lonCell = nc_file.variables['lonCell'][:] + latCell = nc_file.variables['latCell'][:] + bottomDepth = nc_file.variables['bottomDepth'][:] + + # Transform 0,360 range to -180,180 + idx, = np.where(lonCell > np.pi) + lonCell[idx] = lonCell[idx] - 2.0 * np.pi + idx, = np.where(lonCell < -np.pi) + lonCell[idx] = lonCell[idx] + 2.0 * np.pi + + # Interpolate cellWidth onto background grid + cellWidth = 2.0 * np.sqrt(areaCell / np.pi) + hfun = interpolate.NearestNDInterpolator((lonCell, latCell), cellWidth) + hfun_interp = hfun(xy_pts) + hfun_grd = np.reshape(hfun_interp, (nlat, nlon)) + + # Interpolate bathymetry onto background grid + bathy = interpolate.NearestNDInterpolator((lonCell, latCell), + bottomDepth) + bathy_interp = bathy(xy_pts) + bathy_grd = np.reshape(bathy_interp, (nlat, nlon)) + + # Get distance to coasts + D = self.distance_to_shapefile_points(lon, lat, + sphere_radius, reggrid=True) + + # Apply refined region criteria + idxx, idxy = np.where((bathy_grd < depth_threshold_refined) & + (bathy_grd > 0.0) & + (D < dist_threshold_refined) & + (hfun_grd < refined_res)) + cell_width[idxx, idxy] = hfun_grd[idxx, idxy] + + # Apply global region criteria + idxx, idxy = np.where((bathy_grd < depth_threshold_global) & + (bathy_grd > 0.0) & + (D < dist_threshold_global)) + cell_width[idxx, idxy] = hfun_grd[idxx, idxy] + + # Plot + fig = plt.figure(figsize=(16, 8)) + + ax1 = fig.add_subplot(4, 1, 1, projection=ccrs.PlateCarree()) + plt1 = ax1.contourf(lon, lat, cell_width, transform=ccrs.PlateCarree()) + ax1.set_title('waves mesh cell width') + fig.colorbar(plt1, ax=ax1) + + ax2 = fig.add_subplot(4, 1, 2, projection=ccrs.PlateCarree()) + plt2 = ax2.contourf(lon, lat, hfun_grd, transform=ccrs.PlateCarree()) + ax2.set_title('ocean cell width') + fig.colorbar(plt2, ax=ax2) + + ax3 = fig.add_subplot(4, 1, 3, projection=ccrs.PlateCarree()) + plt3 = ax3.contourf(lon, lat, bathy_grd, transform=ccrs.PlateCarree()) + ax3.set_title('bathymetry') + fig.colorbar(plt3, ax=ax3) + + ax4 = fig.add_subplot(4, 1, 4, projection=ccrs.PlateCarree()) + plt4 = ax4.contourf(lon, lat, D, transform=ccrs.PlateCarree()) + ax4.set_title('distance to coast') + fig.colorbar(plt4, ax=ax4) + + fig.savefig('cell_width.png', bbox_inches='tight', dpi=400) + plt.close() + + return cell_width + + def limit_spacing_gradient(self, lon, lat, cell_width, + sphere_radius, dhdx): + + print("Smoothing h(x) via |dh/dx| limits...") + + opts = jigsawpy.jigsaw_jig_t() + opts.hfun_file = "spac_pre_smooth.msh" + opts.jcfg_file = "opts_pre_smooth.jig" + opts.verbosity = +1 + + spac = jigsawpy.jigsaw_msh_t() + spac.mshID = "ellipsoid-grid" + spac.radii = np.full(3, sphere_radius, dtype=spac.REALS_t) + spac.xgrid = np.radians(lon) + spac.ygrid = np.radians(lat) + spac.value = cell_width.astype(spac.REALS_t) + spac.slope = np.full(spac.value.shape, dhdx, dtype=spac.REALS_t) + jigsawpy.savemsh(opts.hfun_file, spac) + + jigsawpy.cmd.marche(opts, spac) + + return spac.value + + def distance_to_shapefile_points(self, lon, lat, + sphere_radius, reggrid=False): + + # Get coastline coordinates from shapefile + features = cfeature.GSHHSFeature(scale='l', levels=[1, 5]) + + pt_list = [] + for feature in features.geometries(): + + if feature.length < 20: + continue + + for coord in feature.exterior.coords: + pt_list.append([coord[0], coord[1]]) + + coast_pts = np.radians(np.array(pt_list)) + + # Convert coastline points to x,y,z and create kd-tree + npts = coast_pts.shape[0] + coast_pts_xyz = np.zeros((npts, 3)) + x, y, z = self.lonlat2xyz(coast_pts[:, 0], coast_pts[:, 1], + sphere_radius) + coast_pts_xyz[:, 0] = x + coast_pts_xyz[:, 1] = y + coast_pts_xyz[:, 2] = z + tree = spatial.KDTree(coast_pts_xyz) + + # Make sure longitude and latitude are in radians + if np.amax(lon) < 2.1 * np.pi: + lonr = lon + latr = lat + else: + lonr = np.radians(lon) + latr = np.radians(lat) + + # Convert backgound grid coordinates to x,y,z + # and put in a nx x 3 array for kd-tree query + if reggrid: + Lon, Lat = np.meshgrid(lonr, latr) + else: + Lon = lonr + Lat = latr + X, Y, Z = self.lonlat2xyz(Lon, Lat, sphere_radius) + pts = np.vstack([X.ravel(), Y.ravel(), Z.ravel()]).T + + # Find distances of background grid coordinates to the coast + print(" Finding distance") + start = timeit.default_timer() + d, idx = tree.query(pts) + end = timeit.default_timer() + print(" Done") + print(" " + str(end - start) + " seconds") + + D = np.reshape(d, Lon.shape) + + return D + + def create_initial_points(self, base_mesh, culled_mesh, + lon, lat, hfunction, + sphere_radius, outfile): + + # Open MPAS culled mesh and get cell variables + nc_file_culled = Dataset(culled_mesh, 'r') + lonCell_culled = nc_file_culled.variables['lonCell'][:] + nCells_culled = lonCell_culled.shape[0] + + # Open MPAS base mesh and get cell variables + nc_file_base = Dataset(base_mesh, 'r') + lonCell_base = nc_file_base.variables['lonCell'][:] + latCell_base = nc_file_base.variables['latCell'][:] + + # Transform 0,360 range to -180,180 + idx, = np.where(lonCell_base > np.pi) + lonCell_base[idx] = lonCell_base[idx] - 2.0 * np.pi + + # Interpolate hfunction onto mesh cell centers + hfun = interpolate.RegularGridInterpolator( + (np.radians(lon), np.radians(lat)), + hfunction.T) + mesh_pts = np.vstack((lonCell_base, latCell_base)).T + hfun_interp = hfun(mesh_pts) + + # Find cells in refined region of waves mesh + max_res = np.amax(hfunction) + idx, = np.where(hfun_interp < 0.5 * max_res) + + # Find boundary cells + # some meshes have non-zero values + # in the extra columns of cellsOnCell + # in this case, these must be zeroed out + # to correctly identify boundary cells + nEdgesOnCell_culled = nc_file_culled.variables['nEdgesOnCell'][:] + cellsOnCell_culled = nc_file_culled.variables['cellsOnCell'][:] + nz = np.zeros(cellsOnCell_culled.shape, dtype=bool) + for i in range(nCells_culled): + nz[i, 0:nEdgesOnCell_culled[i]] = True + cellsOnCell_culled[~nz] = 0 + nCellsOnCell = np.count_nonzero(cellsOnCell_culled, axis=1) + is_boundary_cell = np.equal(nCellsOnCell, nEdgesOnCell_culled) + idx_bnd, = np.where(is_boundary_cell == False) # noqa: E712 + + # Map from culled mesh cells to base mesh cells + write_map_culled_to_base(base_mesh, culled_mesh, + 'base_to_culled_map.nc') + map_file = Dataset('base_to_culled_map.nc', 'r') + culled_to_base = map_file.variables['mapCulledToBaseCell'][:] + base_idx_bnd = culled_to_base[idx_bnd] + + # Find cells in base mesh connected to culled mesh boundary cells + cellsOnCell_base = nc_file_base.variables['cellsOnCell'][:] + base_idx_xbnd = np.concatenate(cellsOnCell_base[base_idx_bnd]) + base_idx_xbnd = base_idx_xbnd - 1 + + # Force inclusion of all boundary cells, eliminating any duplicates + idx = np.union1d(idx, base_idx_bnd) + idx = np.union1d(idx, base_idx_xbnd) + + # Get coordinates of points + lon = lonCell_base[idx] + lat = latCell_base[idx] + + # Include north pole point + lon = np.append(lon, 0.0) + lat = np.append(lat, 0.5 * np.pi) + + # Change to Cartesian coordinates + x, y, z = self.lonlat2xyz(lon, lat, sphere_radius) + + # Get coordinates and ID into structured array + # (for use with np.savetxt) + pt_list = [] + npt = lon.size + for i in range(npt): + # ID of -1 specifies that node is fixed + pt_list.append((x[i], y[i], z[i], -1)) + pt_type = np.dtype({'names': ['x', 'y', 'z', 'id'], + 'formats': [np.float64, np.float64, + np.float64, np.int32]}) + pts = np.array(pt_list, dtype=pt_type) + + # Write initial conditions file + f = open(outfile, 'w') + f.write('# Initial coordinates \n') + f.write('MSHID=3;EUCLIDEAN-MESH\n') + f.write('NDIMS=3\n') + f.write(f'POINT={npt}\n') + np.savetxt(f, pts, fmt='%.12e;%.12e;%.12e;%2i') + f.close() + + init = jigsawpy.jigsaw_msh_t() + jigsawpy.loadmsh(self.opts.init_file, init) + jigsawpy.savevtk("init.vtk", init) + + return + + def lonlat2xyz(self, lon, lat, R): + + x = R * np.multiply(np.cos(lon), np.cos(lat)) + y = R * np.multiply(np.sin(lon), np.cos(lat)) + z = R * np.sin(lat) + + return x, y, z diff --git a/compass/ocean/tests/global_ocean/wave_mesh/cull_mesh.py b/compass/ocean/tests/global_ocean/wave_mesh/cull_mesh.py new file mode 100644 index 0000000000..4598be3681 --- /dev/null +++ b/compass/ocean/tests/global_ocean/wave_mesh/cull_mesh.py @@ -0,0 +1,56 @@ +from importlib.resources import read_text + +from jinja2 import Template +from mpas_tools.logging import check_call + +from compass import Step + + +class WavesCullMesh(Step): + """ + A step for creating wave mesh based on an ocean mesh + """ + def __init__(self, test_case, ocean_mesh, wave_base_mesh, + name='cull_mesh', subdir=None): + + super().__init__(test_case=test_case, name=name, subdir=subdir) + + self.ocean_mesh = ocean_mesh + + wave_base_mesh_path = wave_base_mesh.path + self.add_input_file( + filename='wave_base_mesh.nc', + work_dir_target=f'{wave_base_mesh_path}/base_mesh.nc') + + def setup(self): + + super().setup() + + if self.ocean_mesh is not None: + mesh_path = self.ocean_mesh.steps['initial_state'].path + culled_mesh_path = f'{mesh_path}/initial_state.nc' + else: + if self.config.has_option('wave_mesh', 'ocean_culled_mesh'): + culled_mesh_path = self.config.get('wave_mesh', + 'ocean_culled_mesh') + else: + raise ValueError('ocean_culled_mesh option required ' + 'in wave_mesh section in cfg file') + + self.add_input_file( + filename='ocean_culled_mesh.nc', + work_dir_target=culled_mesh_path) + + template = Template(read_text( + 'compass.ocean.tests.global_ocean.wave_mesh', + 'cull_mesh.template')) + + text = template.render() + text = f'{text}\n' + + with open(f'{self.work_dir}/cull_waves_mesh.nml', 'w') as file: + file.write(text) + + def run(self): + + check_call('ocean_cull_wave_mesh', logger=self.logger) diff --git a/compass/ocean/tests/global_ocean/wave_mesh/cull_mesh.template b/compass/ocean/tests/global_ocean/wave_mesh/cull_mesh.template new file mode 100644 index 0000000000..0c5dfe2b6b --- /dev/null +++ b/compass/ocean/tests/global_ocean/wave_mesh/cull_mesh.template @@ -0,0 +1,8 @@ +&inputs + waves_mesh_file = 'wave_base_mesh.nc' + ocean_mesh_file = 'ocean_culled_mesh.nc' +/ +&output + waves_mesh_culled_vtk = 'wave_mesh_culled.vtk' + waves_mesh_culled_gmsh = 'wave_mesh_culled.msh' +/ diff --git a/compass/ocean/tests/global_ocean/wave_mesh/remap_files.py b/compass/ocean/tests/global_ocean/wave_mesh/remap_files.py new file mode 100644 index 0000000000..79f5e42a1d --- /dev/null +++ b/compass/ocean/tests/global_ocean/wave_mesh/remap_files.py @@ -0,0 +1,104 @@ +from datetime import date +from distutils.spawn import find_executable + +from mpas_tools.logging import check_call + +from compass import Step + + +class WavesRemapFiles(Step): + """ + A step for creating remapping files for wave mesh + """ + def __init__(self, test_case, wave_scrip, + name='remap_files', subdir=None): + + super().__init__(test_case=test_case, name=name, subdir=subdir) + + wave_scrip_file_path = wave_scrip.path + self.add_input_file( + filename='wave_scrip.nc', + work_dir_target=f'{wave_scrip_file_path}/wave_mesh_scrip.nc') + + ocean_scrip_file_path = wave_scrip.path + self.add_input_file( + filename='ocean_scrip.nc', + work_dir_target=f'{ocean_scrip_file_path}/ocean_mesh_scrip.nc') + + def make_remapping_files(self, grid1, grid2, + grid1_shortname, grid2_shortname, + datestamp, reg1, reg2, map_type, nprocs=1): + + if map_type == 'conserve': + map_abbrev = 'aave' + elif map_type == 'bilinear': + map_abbrev = 'blin' + elif map_type == 'patch': + map_abbrev = 'patc' + elif map_type == 'neareststod': + map_abbrev = 'nstod' + elif map_type == 'nearsetdtos': + map_abbrev = 'ndtos' + else: + print('map type not recognized') + raise SystemExit(0) + + exe = find_executable('ESMF_RegridWeightGen') + parallel_executable = self.config.get('parallel', + 'parallel_executable') + ntasks = self.ntasks + parallel_args = parallel_executable.split(' ') + if 'srun' in parallel_args: + parallel_args.extend(['-n', f'{ntasks}']) + else: # presume mpirun syntax + parallel_args.extend(['-np', f'{ntasks}']) + + flags = ['--ignore_unmapped', '--ignore_degenerate'] + + # wave to ocean remap + map_name = f'map_{grid1_shortname}_TO_{grid2_shortname}'\ + f'_{map_abbrev}.{datestamp}.nc' + + args = [exe, + '--source', grid1, + '--destination', grid2, + '--method', map_type, + '--weight', map_name] + args.extend(flags) + if reg1: + args.append('--src_regional') + if reg2: + args.append('--dst_regional') + + cmd = parallel_args + args + check_call(cmd, logger=self.logger) + + # ocean to wave remap + map_name = f'map_{grid2_shortname}_TO_{grid1_shortname}'\ + f'_{map_abbrev}.{datestamp}.nc' + + args = [exe, + '--source', grid2, + '--destination', grid1, + '--method', map_type, + '--weight', map_name] + args.extend(flags) + if reg1: + args.append('--dst_regional') + if reg2: + args.append('--src_regional') + + cmd = parallel_args + args + check_call(cmd, logger=self.logger) + + def run(self): + today = date.today() + creation_date = today.strftime("%Y%m%d") + for map_type in ['conserve', 'bilinear', 'neareststod']: + self.make_remapping_files('wave_scrip.nc', + 'ocean_scrip.nc', + 'wave', # This should be more specific + 'ocean', # This should be more spedific + creation_date, + True, True, + map_type) diff --git a/compass/ocean/tests/global_ocean/wave_mesh/rotate_mesh.py b/compass/ocean/tests/global_ocean/wave_mesh/rotate_mesh.py new file mode 100644 index 0000000000..483805377a --- /dev/null +++ b/compass/ocean/tests/global_ocean/wave_mesh/rotate_mesh.py @@ -0,0 +1,39 @@ +from importlib.resources import read_text + +from jinja2 import Template +from mpas_tools.logging import check_call + +from compass import Step + + +class WavesRotateMesh(Step): + """ + A step for creating wave mesh based on an ocean mesh + """ + def __init__(self, test_case, wave_culled_mesh, + name='rotate_mesh', subdir=None): + + super().__init__(test_case=test_case, name=name, subdir=subdir) + + wave_culled_mesh_path = wave_culled_mesh.path + self.add_input_file( + filename='wave_mesh_culled.msh', + work_dir_target=f'{wave_culled_mesh_path}/wave_mesh_culled.msh') + + def setup(self): + + super().setup() + + template = Template(read_text( + 'compass.ocean.tests.global_ocean.wave_mesh', + 'rotate_mesh.template')) + + text = template.render() + text = f'{text}\n' + + with open(f'{self.work_dir}/rotate.nml', 'w') as file: + file.write(text) + + def run(self): + + check_call('ocean_rotate_wave_mesh', logger=self.logger) diff --git a/compass/ocean/tests/global_ocean/wave_mesh/rotate_mesh.template b/compass/ocean/tests/global_ocean/wave_mesh/rotate_mesh.template new file mode 100644 index 0000000000..e7cd97d008 --- /dev/null +++ b/compass/ocean/tests/global_ocean/wave_mesh/rotate_mesh.template @@ -0,0 +1,7 @@ +&inputs + LON_POLE = -42.8906d0 + LAT_POLE = 72.3200d0 + wind_file = 'null' + mesh_file = 'wave_mesh_culled.msh' + mesh_file_out = 'wave_culled_mesh_RTD.msh' +/ diff --git a/compass/ocean/tests/global_ocean/wave_mesh/scrip_file.py b/compass/ocean/tests/global_ocean/wave_mesh/scrip_file.py new file mode 100644 index 0000000000..73e23f88f7 --- /dev/null +++ b/compass/ocean/tests/global_ocean/wave_mesh/scrip_file.py @@ -0,0 +1,62 @@ +from importlib.resources import read_text + +from jinja2 import Template +from mpas_tools.logging import check_call +from mpas_tools.scrip.from_mpas import scrip_from_mpas + +from compass import Step + + +class WavesScripFile(Step): + """ + A step for creating the scrip file for the wave mesh + """ + def __init__(self, test_case, wave_culled_mesh, ocean_mesh, + name='scrip_file', subdir=None): + + super().__init__(test_case=test_case, name=name, subdir=subdir) + + self.ocean_mesh = ocean_mesh + + wave_culled_mesh_path = wave_culled_mesh.path + self.add_input_file( + filename='wave_mesh_culled.msh', + work_dir_target=f'{wave_culled_mesh_path}/wave_mesh_culled.msh') + + def setup(self): + + super().setup() + + if self.ocean_mesh is not None: + mesh_path = self.ocean_mesh.steps['initial_state'].path + ocean_mesh_path = f'{mesh_path}/initial_state.nc' + else: + if self.config.has_option('wave_mesh', 'ocean_culled_mesh'): + ocean_mesh_path = self.config.get('wave_mesh', + 'ocean_culled_mesh') + else: + raise ValueError('ocean_culled_mesh option required ' + 'in wave_mesh section in cfg file') + + self.add_input_file( + filename='ocean_mesh.nc', + work_dir_target=ocean_mesh_path) + + template = Template(read_text( + 'compass.ocean.tests.global_ocean.wave_mesh', + 'scrip_file.template')) + + text = template.render() + text = f'{text}\n' + + with open(f'{self.work_dir}/scrip.nml', 'w') as file: + file.write(text) + + def run(self): + """ + Create scrip files for wave mesh + """ + local_filename = 'ocean_mesh_scrip.nc' + scrip_from_mpas('ocean_mesh.nc', local_filename) + + check_call('ocean_scrip_wave_mesh', logger=self.logger) diff --git a/compass/ocean/tests/global_ocean/wave_mesh/scrip_file.template b/compass/ocean/tests/global_ocean/wave_mesh/scrip_file.template new file mode 100644 index 0000000000..c6d3582026 --- /dev/null +++ b/compass/ocean/tests/global_ocean/wave_mesh/scrip_file.template @@ -0,0 +1,6 @@ +&inputs + waves_mesh_file = 'wave_mesh_culled.msh' +/ +&outputs + waves_scrip_file = 'wave_mesh_scrip.nc' +/ diff --git a/compass/ocean/tests/global_ocean/wave_mesh/uost_files.py b/compass/ocean/tests/global_ocean/wave_mesh/uost_files.py new file mode 100644 index 0000000000..7c1623650a --- /dev/null +++ b/compass/ocean/tests/global_ocean/wave_mesh/uost_files.py @@ -0,0 +1,271 @@ +import matplotlib.pyplot as plt +import numpy as np + +from alphaBetaLab.alphaBetaLab.abEstimateAndSave import ( + abEstimateAndSaveTriangularEtopo1, + triMeshSpecFromMshFile, +) +# importing from alphaBetaLab the needed components +from alphaBetaLab.alphaBetaLab.abOptionManager import abOptions +from compass import Step + + +class WavesUostFiles(Step): + """ + A step for creating the unresolved obstacles file for wave mesh + """ + def __init__(self, test_case, wave_culled_mesh, wave_rotate_mesh, + name='uost_files', subdir=None): + + super().__init__(test_case=test_case, name=name, subdir=subdir) + + # other things INIT should do? + culled_mesh_path = wave_culled_mesh.path + self.add_input_file( + filename='wave_mesh_culled.msh', + work_dir_target=f'{culled_mesh_path}/wave_mesh_culled.msh') + + angled_path = wave_rotate_mesh.path + self.add_input_file( + filename='angled.d', + work_dir_target=f'{angled_path}/angled.d') + + self.add_input_file( + filename='etopo1_180.nc', + target='etopo1_180.nc', + database='bathymetry_database') + + def run(self): + """ + Create unresolved obstacles for wave mesh and spectral resolution + """ + + ndir = self.config.getint('wave_mesh', 'ndir') + nfreq = self.config.getint('wave_mesh', 'nfreq') + + dirs = np.linspace(0, 2 * np.pi, ndir) + minfrq = .035 + if (nfreq == 50): + frqfactor = 1.07 + elif (nfreq == 36): + frqfactor = 1.10 + elif (nfreq == 25): + frqfactor = 1.147 + else: + print("ERROR: Spectral resolution not supported.") + print("Number of wave freqencies must be 25, 36, or 50.") + + freqs = [minfrq * (frqfactor ** i) for i in range(1, nfreq + 1)] + + # definition of the spatial mesh + gridname = 'waves_mesh_culled' # SB NOTE: should be flexible + mshfile = 'wave_mesh_culled.msh' + triMeshSpec = triMeshSpecFromMshFile(mshfile) + + # path of the etopo1 bathymetry + etopoFilePath = 'etopo1_180.nc' + + # output directory + outputDestDir = 'output/' + + # number of cores for parallel computing + nParWorker = 1 + #nParWorker = self.config.getint('parallel', 'cores_per_node') + + # this option indicates that the computation + # should be skipped for cells smaller than 3 km + minSizeKm = 3 + opt = abOptions(minSizeKm=minSizeKm) + + # instruction to do the computation and save the output + # abEstimateAndSaveTriangularGebco( + # dirs, freqs, gridname, triMeshSpec, + # etopoFilePath, outputDestDir, nParWorker, abOptions=opt) + abEstimateAndSaveTriangularEtopo1( + dirs, freqs, gridname, triMeshSpec, etopoFilePath, + outputDestDir, nParWorker, abOptions=opt) + + theta = np.radians(np.linspace(0.0, 360.0, ndir, endpoint=False)) + freq = np.linspace(0.0, 1.0, nfreq) + Theta, Freq = np.meshgrid(theta, freq) + + data = np.loadtxt('angled.d') + angled = data[:, 2] + + filename_local_in = 'obstructions_local.waves_mesh_culled.in' + filename_local_out = 'obstructions_local.waves_mesh_culled.rtd.in' + lon_local, lat_local, nodes, sizes, alpha_local_avg, beta_local_avg, alpha_spec, beta_spec = self.read_alpha_beta(filename_local_in, nfreq) + alpha_interp, beta_interp = self.rotate_and_interpolate(Theta, nodes, angled, alpha_spec, beta_spec) + header = '$WAVEWATCH III LOCAL OBSTRUCTIONS' + self.write_alpha_beta(filename_local_out, header, nodes, lon_local, lat_local, sizes, alpha_interp, beta_interp) + self.plot_alpha_beta_spectra(Theta, Freq, alpha_spec, beta_spec, alpha_interp, beta_interp, angled, nodes, 'local') + + filename_shadow_in = 'obstructions_shadow.waves_mesh_culled.in' + filename_shadow_out = 'obstructions_shadow.waves_mesh_culled.rtd.in' + lon_shadow, lat_shadow, nodes, sizes, alpha_shadow_avg, beta_shadow_avg, alpha_spec, beta_spec = self.read_alpha_beta(filename_shadow_in, nfreq) + alpha_interp, beta_interp = self.rotate_and_interpolate(Theta, nodes, angled, alpha_spec, beta_spec) + header = '$WAVEWATCH III SHADOW OBSTRUCTIONS' + self.write_alpha_beta(filename_shadow_out, header, nodes, lon_shadow, lat_shadow, sizes, alpha_interp, beta_interp) + self.plot_alpha_beta_spectra(Theta, Freq, alpha_spec, beta_spec, alpha_interp, beta_interp, angled, nodes, 'shadow') + + def write_alpha_beta(self, filename, header, nodes, lon, lat, sizes, alpha_spec, beta_spec): + + n = alpha_spec.shape[0] + nfreq = alpha_spec.shape[1] + ndir = alpha_spec.shape[2] + + lines = [] + lines.append(header) + lines.append(str(n)) + + for i in range(n): + lines.append('$ ilon ilat of the cell. lon: {:.8f}, lat: {:.8f}'.format(lon[i], lat[i])) + lines.append(str(nodes[i]) + ' 1') + lines.append(sizes[i]) + lines.append('$ mean alpha: {:.16}'.format(np.mean(alpha_spec[i, :, :]))) + lines.append('$ mean beta: {:.16}'.format(np.mean(beta_spec[i, :, :]))) + lines.append('$alpha by ik, ith') + for j in range(nfreq): + line = '' + for k in range(ndir): + line = line + '{:.2f} '.format(alpha_spec[i, j, k]) + lines.append(line) + lines.append('$beta by ik, ith') + for j in range(nfreq): + line = '' + for k in range(ndir): + line = line + '{:.2f} '.format(beta_spec[i, j, k]) + lines.append(line) + + f = open(filename, 'w') + for line in lines: + f.write(line + '\n') + f.close() + + def rotate_and_interpolate(self, Theta, nodes, angled, + alpha_spec, beta_spec): + + n = alpha_spec.shape[0] + nfreq = alpha_spec.shape[1] + ndir = alpha_spec.shape[2] + + alpha_interp = np.zeros((n, nfreq, ndir)) + beta_interp = np.zeros((n, nfreq, ndir)) + for i in range(n): + nd = nodes[i] - 1 + Theta2 = Theta - angled[nd] * np.pi / 180.0 + for j in range(nfreq): + alpha_interp[i, j, :] = np.interp(Theta2[j, :], + Theta[j, :], + alpha_spec[i, j, :], + period=2.0 * np.pi) + beta_interp[i, j, :] = np.interp(Theta2[j, :], + Theta[j, :], + beta_spec[i, j, :], + period=2.0 * np.pi) + + return [alpha_interp, beta_interp] + + def plot_alpha_beta_spectra(self, Theta, Freq, alpha_spec, beta_spec, + alpha_interp, beta_interp, angled, nodes, kind): + + for i in range(10): + print(i) + + fig = plt.figure(figsize=[8, 4]) + ax = fig.add_subplot(2, 2, 1, polar=True) + ax.set_theta_zero_location("N") + ax.set_theta_direction(-1) + ax.contourf(Theta, Freq, alpha_spec[i, :, :], 30) + + ax = fig.add_subplot(2, 2, 2, polar=True) + ax.set_theta_zero_location("N") + ax.set_theta_direction(-1) + ax.contourf(Theta, Freq, beta_spec[i, :, :], 30) + + ax = fig.add_subplot(2, 2, 3, polar=True) + ax.set_theta_zero_location("N") + ax.set_theta_direction(-1) + ax.contourf(Theta, Freq, alpha_interp[i, :, :], 30) + ax.set_title('AnglD = ' + str(angled[nodes[i] - 1])) + + ax = fig.add_subplot(2, 2, 4, polar=True) + ax.set_theta_zero_location("N") + ax.set_theta_direction(-1) + ax.contourf(Theta, Freq, beta_interp[i, :, :], 30) + ax.set_title('AnglD = ' + str(angled[nodes[i] - 1])) + + plt.savefig(kind + '_spec_' + str(i) + '.png', bbox_inches='tight') + + def read_alpha_beta(self, filename, nfreq): + f = open(filename, 'r') + lines = f.read().splitlines() + + nodes = [] + lon = [] + lat = [] + sizes = [] + alpha_avg = [] + beta_avg = [] + alpha_spec = [] + beta_spec = [] + + line = 1 # header comment + n = int(lines[line]) + for i in range(n): + + line = line + 1 # lon lat comment + + text = lines[line] + text_sp = text.split() + x = float(text_sp[7].replace(',', '')) + y = float(text_sp[9]) + + line = line + 1 # node number + nodes.append(int(lines[line].split()[0])) + + line = line + 1 # sizes comment + line = line + 1 # sizes + sizes.append(lines[line]) + + line = line + 1 # mean alpha + text = lines[line] + text_sp = text.split() + a = float(text_sp[-1]) + + line = line + 1 # mean beta + text = lines[line] + text_sp = text.split() + b = float(text_sp[-1]) + + line = line + 1 # alpha comment + spectrum = [] + for i in range(nfreq): + line = line + 1 + spectrum.append(lines[line].split()) + alpha_spec.append(spectrum) + del spectrum + + line = line + 1 # beta comment + spectrum = [] + for i in range(nfreq): + line = line + 1 + spectrum.append(lines[line].split()) + beta_spec.append(spectrum) + del spectrum + + lon.append(x) + lat.append(y) + alpha_avg.append(a) + beta_avg.append(b) + + lon = np.array(lon) + lat = np.array(lat) + nodes = np.array(nodes) + alpha_avg = np.array(alpha_avg) + beta_avg = np.array(beta_avg) + alpha_spec = np.array(alpha_spec) + beta_spec = np.array(beta_spec) + + return [lon, lat, nodes, sizes, alpha_avg, beta_avg, + alpha_spec, beta_spec] diff --git a/compass/ocean/tests/global_ocean/wave_mesh/wave_mesh.cfg b/compass/ocean/tests/global_ocean/wave_mesh/wave_mesh.cfg new file mode 100644 index 0000000000..bec4b05128 --- /dev/null +++ b/compass/ocean/tests/global_ocean/wave_mesh/wave_mesh.cfg @@ -0,0 +1,42 @@ +[wave_mesh] + +# Spacing of grid used to specify wave mesh resolution, units of deg +hfun_grid_spacing = 0.5 + +# Resolution gradient limit +hfun_slope_lim = 0.15 + +# Ocean mesh resolution threshold for where depth_threshold_refined +# and distance_threshold_refined, criteria apply. This can be used +# to set more flexible threshold values for regionally refined ocean +# meshes, units of m. +refined_res = 20000.0 + +# Ocean resolution gets applied where depth is less than this threshold +# and where ocean resolution is less than refined_res, units of m +depth_threshold_refined = 1000.0 + +# Maximum distance from coast where ocean resolution gets applied +# where ocean resolutoin is less than refined_res, units of km +distance_threshold_refined = 300.0 + +# Ocean resolution gets applied where depth is less than this threshold, units of m +depth_threshold_global = 1000.0 + +# Maximum distance from ocast where ocean resolution gets applied, units of km +distance_threshold_global = 300.0 + +# Maximum resolution of waves mesh, units of m +maxres = 225000.0 + +# Number of spectral directions bins +ndir = 36 + +# Number of spectral frequency bins +nfreq = 36 + +# Path to existing base ocean mesh (optional) +# ocean_base_mesh = + +# Path to existing culled ocean mesh (optional) +# ocean_culled_mesh =