From 7dbf3934fd36c25e879a2272ac5c034c6eef3d30 Mon Sep 17 00:00:00 2001 From: Carolyn Begeman Date: Tue, 18 Jul 2023 12:28:31 -0500 Subject: [PATCH] Assume mesh variables only in mesh file --- .../ocean/tests/baroclinic_channel/init.py | 6 ++--- .../tests/baroclinic_channel/rpe/analysis.py | 1 - polaris/ocean/tests/baroclinic_channel/viz.py | 6 ++--- .../cosine_bell/analysis.py | 18 ++++++++------- .../global_convergence/cosine_bell/init.py | 22 +++++++++---------- .../global_convergence/cosine_bell/viz.py | 8 +++++-- .../ocean/tests/inertial_gravity_wave/init.py | 10 ++++----- .../ocean/tests/inertial_gravity_wave/viz.py | 14 +++++++----- .../ocean/tests/manufactured_solution/init.py | 10 ++++----- .../ocean/tests/manufactured_solution/viz.py | 14 +++++++----- polaris/ocean/tests/single_column/init.py | 8 +++---- 11 files changed, 65 insertions(+), 52 deletions(-) diff --git a/polaris/ocean/tests/baroclinic_channel/init.py b/polaris/ocean/tests/baroclinic_channel/init.py index 9ff07ac32..9f28d3dad 100644 --- a/polaris/ocean/tests/baroclinic_channel/init.py +++ b/polaris/ocean/tests/baroclinic_channel/init.py @@ -140,7 +140,7 @@ def run(self): temperature = temperature.expand_dims(dim='Time', axis=0) - normal_velocity = xr.zeros_like(ds.xEdge) + normal_velocity = xr.zeros_like(ds_mesh.xEdge) normal_velocity, _ = xr.broadcast(normal_velocity, ds.refBottomDepth) normal_velocity = normal_velocity.transpose('nEdges', 'nVertLevels') normal_velocity = normal_velocity.expand_dims(dim='Time', axis=0) @@ -149,8 +149,8 @@ def run(self): ds['salinity'] = salinity * xr.ones_like(temperature) ds['normalVelocity'] = normal_velocity ds['fCell'] = coriolis_parameter * xr.ones_like(x_cell) - ds['fEdge'] = coriolis_parameter * xr.ones_like(ds.xEdge) - ds['fVertex'] = coriolis_parameter * xr.ones_like(ds.xVertex) + ds['fEdge'] = coriolis_parameter * xr.ones_like(ds_mesh.xEdge) + ds['fVertex'] = coriolis_parameter * xr.ones_like(ds_mesh.xVertex) ds.attrs['nx'] = nx ds.attrs['ny'] = ny diff --git a/polaris/ocean/tests/baroclinic_channel/rpe/analysis.py b/polaris/ocean/tests/baroclinic_channel/rpe/analysis.py index 206e7f6fa..f8078ad56 100644 --- a/polaris/ocean/tests/baroclinic_channel/rpe/analysis.py +++ b/polaris/ocean/tests/baroclinic_channel/rpe/analysis.py @@ -149,5 +149,4 @@ def _plot(nx, ny, lx, ly, filename, nus, rpe): ax.set_ylabel('y, km') if iCol == num_files - 1: fig.colorbar(dis, ax=axs[num_files - 1], aspect=40) - plt.savefig(filename) diff --git a/polaris/ocean/tests/baroclinic_channel/viz.py b/polaris/ocean/tests/baroclinic_channel/viz.py index 4f724b41f..60671dc79 100644 --- a/polaris/ocean/tests/baroclinic_channel/viz.py +++ b/polaris/ocean/tests/baroclinic_channel/viz.py @@ -22,8 +22,8 @@ def __init__(self, test_case): """ super().__init__(test_case=test_case, name='viz') self.add_input_file( - filename='initial_state.nc', - target='../init/initial_state.nc') + filename='mesh.nc', + target='../init/culled_mesh.nc') self.add_input_file( filename='output.nc', target='../forward/output.nc') @@ -32,7 +32,7 @@ def run(self): """ Run this step of the test case """ - ds_mesh = xr.load_dataset('initial_state.nc') + ds_mesh = xr.load_dataset('mesh.nc') ds = xr.load_dataset('output.nc') t_index = ds.sizes['Time'] - 1 plot_horiz_field(ds, ds_mesh, 'temperature', diff --git a/polaris/ocean/tests/global_convergence/cosine_bell/analysis.py b/polaris/ocean/tests/global_convergence/cosine_bell/analysis.py index efdb09abc..4de7889fa 100644 --- a/polaris/ocean/tests/global_convergence/cosine_bell/analysis.py +++ b/polaris/ocean/tests/global_convergence/cosine_bell/analysis.py @@ -45,6 +45,9 @@ def __init__(self, test_case, resolutions, icosahedral): mesh_name = f'Icos{resolution}' else: mesh_name = f'QU{resolution}' + self.add_input_file( + filename=f'{mesh_name}_mesh.nc', + target=f'../{mesh_name}/mesh/mesh.nc') self.add_input_file( filename=f'{mesh_name}_init.nc', target=f'../{mesh_name}/init/initial_state.nc') @@ -127,7 +130,8 @@ def rmse(self, mesh_name): psi0 = config.getfloat('cosine_bell', 'psi0') pd = config.getfloat('cosine_bell', 'vel_pd') - init = xr.open_dataset(f'{mesh_name}_init.nc') + ds_mesh = xr.open_dataset(f'{mesh_name}_mesh.nc') + ds_init = xr.open_dataset(f'{mesh_name}_init.nc') # find time since the beginning of run ds = xr.open_dataset(f'{mesh_name}_output.nc') for j in range(len(ds.xtime)): @@ -142,7 +146,7 @@ def rmse(self, mesh_name): t = 86400.0 * DY + HR * 3600. + MN # find new location of blob center # center is based on equatorial velocity - R = init.sphere_radius + R = ds_mesh.sphere_radius distTrav = 2.0 * 3.14159265 * R / (86400.0 * pd) * t # distance in radians is distRad = distTrav / R @@ -151,9 +155,9 @@ def rmse(self, mesh_name): newLon -= 2.0 * np.pi # construct analytic tracer - tracer = np.zeros_like(init.tracer1[0, :, 0].values) - latC = init.latCell.values - lonC = init.lonCell.values + tracer = np.zeros_like(ds_init.tracer1[0, :, 0].values) + latC = ds_mesh.latCell.values + lonC = ds_mesh.lonCell.values temp = R * np.arccos(np.sin(latCent) * np.sin(latC) + np.cos(latCent) * np.cos(latC) * np.cos( lonC - newLon)) @@ -165,6 +169,4 @@ def rmse(self, mesh_name): tracerF = ds.tracer1[sliceTime, :, 0].values rmseValue = np.sqrt(np.mean((tracerF - tracer)**2)) - init.close() - ds.close() - return rmseValue, init.dims['nCells'] + return rmseValue, ds_init.dims['nCells'] diff --git a/polaris/ocean/tests/global_convergence/cosine_bell/init.py b/polaris/ocean/tests/global_convergence/cosine_bell/init.py index a1e329bec..769769796 100644 --- a/polaris/ocean/tests/global_convergence/cosine_bell/init.py +++ b/polaris/ocean/tests/global_convergence/cosine_bell/init.py @@ -51,21 +51,21 @@ def run(self): section = config['vertical_grid'] bottom_depth = section.getfloat('bottom_depth') - dsMesh = xr.open_dataset('mesh.nc') - angleEdge = dsMesh.angleEdge - latCell = dsMesh.latCell - latEdge = dsMesh.latEdge - lonCell = dsMesh.lonCell - sphere_radius = dsMesh.sphere_radius + ds_mesh = xr.open_dataset('mesh.nc') + angleEdge = ds_mesh.angleEdge + latCell = ds_mesh.latCell + latEdge = ds_mesh.latEdge + lonCell = ds_mesh.lonCell + sphere_radius = ds_mesh.sphere_radius - ds = dsMesh.copy() + ds = ds_mesh.copy() ds['bottomDepth'] = bottom_depth * xr.ones_like(latCell) ds['ssh'] = xr.zeros_like(latCell) init_vertical_coord(config, ds) - temperature_array = temperature * xr.ones_like(ds.latCell) + temperature_array = temperature * xr.ones_like(ds_mesh.latCell) temperature_array, _ = xr.broadcast(temperature_array, ds.refZMid) ds['temperature'] = temperature_array.expand_dims(dim='Time', axis=0) ds['salinity'] = salinity * xr.ones_like(ds.temperature) @@ -92,8 +92,8 @@ def run(self): velocity_array, _ = xr.broadcast(velocity, ds.refZMid) ds['normalVelocity'] = velocity_array.expand_dims(dim='Time', axis=0) - ds['fCell'] = xr.zeros_like(ds.xCell) - ds['fEdge'] = xr.zeros_like(ds.xEdge) - ds['fVertex'] = xr.zeros_like(ds.xVertex) + ds['fCell'] = xr.zeros_like(ds_mesh.xCell) + ds['fEdge'] = xr.zeros_like(ds_mesh.xEdge) + ds['fVertex'] = xr.zeros_like(ds_mesh.xVertex) write_netcdf(ds, 'initial_state.nc') diff --git a/polaris/ocean/tests/global_convergence/cosine_bell/viz.py b/polaris/ocean/tests/global_convergence/cosine_bell/viz.py index b1a24e7b5..a5be36535 100644 --- a/polaris/ocean/tests/global_convergence/cosine_bell/viz.py +++ b/polaris/ocean/tests/global_convergence/cosine_bell/viz.py @@ -86,6 +86,9 @@ def __init__(self, test_case, name, subdir, viz_map, mesh_name): The name of the mesh """ super().__init__(test_case=test_case, name=name, subdir=subdir) + self.add_input_file( + filename='mesh.nc', + target='../init/mesh.nc') self.add_input_file( filename='initial_state.nc', target='../init/initial_state.nc') @@ -109,12 +112,13 @@ def run(self): remapper = viz_map.get_remapper() + ds_mesh = xr.open_dataset('mesh.nc') ds_init = xr.open_dataset('initial_state.nc') ds_init = ds_init[['tracer1', ]].isel(Time=0, nVertLevels=0) ds_init = remapper.remap(ds_init) ds_init.to_netcdf('remapped_init.nc') - plot_global(ds_init.lon.values, ds_init.lat.values, + plot_global(ds_mesh.lon.values, ds_mesh.lat.values, ds_init.tracer1.values, out_filename='init.png', config=config, colormap_section='cosine_bell_viz', @@ -125,7 +129,7 @@ def run(self): ds_out = remapper.remap(ds_out) ds_out.to_netcdf('remapped_final.nc') - plot_global(ds_out.lon.values, ds_out.lat.values, + plot_global(ds_mesh.lon.values, ds_mesh.lat.values, ds_out.tracer1.values, out_filename='final.png', config=config, colormap_section='cosine_bell_viz', diff --git a/polaris/ocean/tests/inertial_gravity_wave/init.py b/polaris/ocean/tests/inertial_gravity_wave/init.py index dabe9086b..39134b162 100644 --- a/polaris/ocean/tests/inertial_gravity_wave/init.py +++ b/polaris/ocean/tests/inertial_gravity_wave/init.py @@ -70,14 +70,14 @@ def run(self): ds = ds_mesh.copy() - ds['ssh'] = xr.zeros_like(ds.xCell) - ds['bottomDepth'] = bottom_depth * xr.ones_like(ds.xCell) + ds['ssh'] = xr.zeros_like(ds_mesh.xCell) + ds['bottomDepth'] = bottom_depth * xr.ones_like(ds_mesh.xCell) init_vertical_coord(config, ds) - ds['fCell'] = f0 * xr.ones_like(ds.xCell) - ds['fEdge'] = f0 * xr.ones_like(ds.xEdge) - ds['fVertex'] = f0 * xr.ones_like(ds.xVertex) + ds['fCell'] = f0 * xr.ones_like(ds_mesh.xCell) + ds['fEdge'] = f0 * xr.ones_like(ds_mesh.xEdge) + ds['fVertex'] = f0 * xr.ones_like(ds_mesh.xVertex) ds_mesh['maxLevelCell'] = ds.maxLevelCell exact_solution = ExactSolution(ds, config) diff --git a/polaris/ocean/tests/inertial_gravity_wave/viz.py b/polaris/ocean/tests/inertial_gravity_wave/viz.py index 025d5635b..5c1d47eb9 100644 --- a/polaris/ocean/tests/inertial_gravity_wave/viz.py +++ b/polaris/ocean/tests/inertial_gravity_wave/viz.py @@ -38,6 +38,9 @@ def __init__(self, test_case, resolutions): self.resolutions = resolutions for resolution in resolutions: + self.add_input_file( + filename=f'mesh_{resolution}km.nc', + target=f'../{resolution}km/initial_state/culled_mesh.nc') self.add_input_file( filename=f'init_{resolution}km.nc', target=f'../{resolution}km/init/initial_state.nc') @@ -63,9 +66,10 @@ def run(self): fig, axes = plt.subplots(nrows=nres, ncols=3, figsize=(12, 2 * nres)) rmse = [] for i, res in enumerate(resolutions): - init = xr.open_dataset(f'init_{res}km.nc') + ds_mesh = xr.open_dataset(f'mesh_{res}km.nc') + ds_init = xr.open_dataset(f'init_{res}km.nc') ds = xr.open_dataset(f'output_{res}km.nc') - exact = ExactSolution(init, config) + exact = ExactSolution(ds_init, config) t0 = datetime.datetime.strptime(ds.xtime.values[0].decode(), '%Y-%m-%d_%H:%M:%S') @@ -81,13 +85,13 @@ def run(self): if i == 0: error_range = np.max(np.abs(ds.ssh_error.values)) - plot_horiz_field(ds, init, 'ssh', ax=axes[i, 0], + plot_horiz_field(ds, ds_mesh, 'ssh', ax=axes[i, 0], cmap='cmo.balance', t_index=ds.sizes["Time"] - 1, vmin=-eta0, vmax=eta0, cmap_title="SSH (m)") - plot_horiz_field(ds, init, 'ssh_exact', ax=axes[i, 1], + plot_horiz_field(ds, ds_mesh, 'ssh_exact', ax=axes[i, 1], cmap='cmo.balance', vmin=-eta0, vmax=eta0, cmap_title="SSH (m)") - plot_horiz_field(ds, init, 'ssh_error', ax=axes[i, 2], + plot_horiz_field(ds, ds_mesh, 'ssh_error', ax=axes[i, 2], cmap='cmo.balance', cmap_title=r"$\Delta$ SSH (m)", vmin=-error_range, vmax=error_range) diff --git a/polaris/ocean/tests/manufactured_solution/init.py b/polaris/ocean/tests/manufactured_solution/init.py index 8a4c47587..04d77c27c 100644 --- a/polaris/ocean/tests/manufactured_solution/init.py +++ b/polaris/ocean/tests/manufactured_solution/init.py @@ -70,14 +70,14 @@ def run(self): ds = ds_mesh.copy() - ds['ssh'] = xr.zeros_like(ds.xCell) - ds['bottomDepth'] = bottom_depth * xr.ones_like(ds.xCell) + ds['ssh'] = xr.zeros_like(ds_mesh.xCell) + ds['bottomDepth'] = bottom_depth * xr.ones_like(ds_mesh.xCell) init_vertical_coord(config, ds) - ds['fCell'] = coriolis_parameter * xr.ones_like(ds.xCell) - ds['fEdge'] = coriolis_parameter * xr.ones_like(ds.xEdge) - ds['fVertex'] = coriolis_parameter * xr.ones_like(ds.xVertex) + ds['fCell'] = coriolis_parameter * xr.ones_like(ds_mesh.xCell) + ds['fEdge'] = coriolis_parameter * xr.ones_like(ds_mesh.xEdge) + ds['fVertex'] = coriolis_parameter * xr.ones_like(ds_mesh.xVertex) ds_mesh['maxLevelCell'] = ds.maxLevelCell diff --git a/polaris/ocean/tests/manufactured_solution/viz.py b/polaris/ocean/tests/manufactured_solution/viz.py index 5fdcfd5bb..4b01818ae 100644 --- a/polaris/ocean/tests/manufactured_solution/viz.py +++ b/polaris/ocean/tests/manufactured_solution/viz.py @@ -38,6 +38,9 @@ def __init__(self, test_case, resolutions): self.resolutions = resolutions for resolution in resolutions: + self.add_input_file( + filename=f'mesh_{resolution}km.nc', + target=f'../{resolution}km/initial_state/culled_mesh.nc') self.add_input_file( filename=f'init_{resolution}km.nc', target=f'../{resolution}km/init/initial_state.nc') @@ -62,9 +65,10 @@ def run(self): fig, axes = plt.subplots(nrows=nres, ncols=3, figsize=(12, 2 * nres)) rmse = [] for i, res in enumerate(resolutions): - init = xr.open_dataset(f'init_{res}km.nc') + ds_mesh = xr.open_dataset(f'mesh_{res}km.nc') + ds_init = xr.open_dataset(f'init_{res}km.nc') ds = xr.open_dataset(f'output_{res}km.nc') - exact = ExactSolution(config, init) + exact = ExactSolution(config, ds_init) t0 = datetime.datetime.strptime(ds.xtime.values[0].decode(), '%Y-%m-%d_%H:%M:%S') @@ -80,13 +84,13 @@ def run(self): if i == 0: error_range = np.max(np.abs(ds.ssh_error.values)) - plot_horiz_field(ds, init, 'ssh', ax=axes[i, 0], + plot_horiz_field(ds, ds_mesh, 'ssh', ax=axes[i, 0], cmap='cmo.balance', t_index=ds.sizes["Time"] - 1, vmin=-eta0, vmax=eta0, cmap_title="SSH") - plot_horiz_field(ds, init, 'ssh_exact', ax=axes[i, 1], + plot_horiz_field(ds, ds_mesh, 'ssh_exact', ax=axes[i, 1], cmap='cmo.balance', vmin=-eta0, vmax=eta0, cmap_title="SSH") - plot_horiz_field(ds, init, 'ssh_error', ax=axes[i, 2], + plot_horiz_field(ds, ds_mesh, 'ssh_error', ax=axes[i, 2], cmap='cmo.balance', cmap_title="dSSH", vmin=-error_range, vmax=error_range) diff --git a/polaris/ocean/tests/single_column/init.py b/polaris/ocean/tests/single_column/init.py index 3da037224..ce33a9822 100644 --- a/polaris/ocean/tests/single_column/init.py +++ b/polaris/ocean/tests/single_column/init.py @@ -57,7 +57,7 @@ def run(self): write_netcdf(ds_mesh, 'culled_mesh.nc') ds = ds_mesh.copy() - x_cell = ds.xCell + x_cell = ds_mesh.xCell bottom_depth = config.getfloat('vertical_grid', 'bottom_depth') ds['bottomDepth'] = bottom_depth * xr.ones_like(x_cell) ds['ssh'] = xr.zeros_like(x_cell) @@ -116,7 +116,7 @@ def run(self): salinity = salinity.expand_dims(dim='Time', axis=0) normal_velocity, _ = xr.broadcast( - xr.zeros_like(ds.xEdge), ds.refBottomDepth) + xr.zeros_like(ds_mesh.xEdge), ds.refBottomDepth) normal_velocity = normal_velocity.transpose('nEdges', 'nVertLevels') normal_velocity = normal_velocity.expand_dims(dim='Time', axis=0) @@ -127,8 +127,8 @@ def run(self): ds['salinity'] = salinity ds['normalVelocity'] = normal_velocity ds['fCell'] = coriolis_parameter * xr.ones_like(x_cell) - ds['fEdge'] = coriolis_parameter * xr.ones_like(ds.xEdge) - ds['fVertex'] = coriolis_parameter * xr.ones_like(ds.xVertex) + ds['fEdge'] = coriolis_parameter * xr.ones_like(ds_mesh.xEdge) + ds['fVertex'] = coriolis_parameter * xr.ones_like(ds_mesh.xVertex) ds.attrs['nx'] = nx ds.attrs['ny'] = ny