diff --git a/compass/ocean/tests/internal_wave/forward.py b/compass/ocean/tests/internal_wave/forward.py index 16e71e4354..ff09bb94b3 100644 --- a/compass/ocean/tests/internal_wave/forward.py +++ b/compass/ocean/tests/internal_wave/forward.py @@ -52,9 +52,12 @@ def __init__(self, test_case, name='forward', subdir=None, cores=1, self.add_streams_file('compass.ocean.tests.internal_wave', 'streams.forward') - self.add_input_file(filename='init.nc', target='../initial_state/ocean.nc') - self.add_input_file(filename='mesh.nc', target='../initial_state/culled_mesh.nc') - self.add_input_file(filename='graph.info', target='../initial_state/culled_graph.info') + self.add_input_file(filename='init.nc', + target='../initial_state/ocean.nc') + self.add_input_file(filename='mesh.nc', + target='../initial_state/culled_mesh.nc') + self.add_input_file(filename='graph.info', + target='../initial_state/culled_graph.info') self.add_model_as_input() diff --git a/compass/ocean/tests/internal_wave/initial_state.py b/compass/ocean/tests/internal_wave/initial_state.py index e7a63e5a86..018756dd2d 100644 --- a/compass/ocean/tests/internal_wave/initial_state.py +++ b/compass/ocean/tests/internal_wave/initial_state.py @@ -96,24 +96,24 @@ def run(self): if use_distances: perturbation_width = amplitude_width_dist else: - perturbation_width = (yMax - yMin) * amplitude_width_frac + perturbation_width = (yMax - yMin) * amplitude_width_frac # Set stratified temperature - temp_vert = (bottom_temperature + - (surface_temperature - bottom_temperature) * - ((ds.refZMid + bottom_depth) / bottom_depth)) + temp_vert = (bottom_temperature + + (surface_temperature - bottom_temperature) * + ((ds.refZMid + bottom_depth) / bottom_depth)) depth_frac = xarray.zeros_like(temp_vert) refBottomDepth = ds['refBottomDepth'] for k in range(1, vert_levels): depth_frac[k] = refBottomDepth[k-1] / refBottomDepth[vert_levels-1] - # If cell is in the southern half, outside the sin width, subtract + # If cell is in the southern half, outside the sin width, subtract # temperature difference - frac = xarray.where( numpy.abs(yCell - yMid) < perturbation_width, - numpy.cos(0.5 * numpy.pi * (yCell - yMid) - / perturbation_width) - * numpy.sin(numpy.pi * depth_frac), + frac = xarray.where(numpy.abs(yCell - yMid) < perturbation_width, + numpy.cos(0.5 * numpy.pi * (yCell - yMid) / + perturbation_width) * + numpy.sin(numpy.pi * depth_frac), 0.) temperature = temp_vert - temperature_difference * frac diff --git a/compass/ocean/tests/internal_wave/rpe_test/analysis.py b/compass/ocean/tests/internal_wave/rpe_test/analysis.py index 5c38dae0b6..a41c82c912 100644 --- a/compass/ocean/tests/internal_wave/rpe_test/analysis.py +++ b/compass/ocean/tests/internal_wave/rpe_test/analysis.py @@ -78,10 +78,10 @@ def _plot(nx, ny, filename, nus): nCol = 2 iTime = [0, 1] time = ['1', '21'] - + fig, axs = plt.subplots(nRow, nCol, figsize=( 4.0 * nCol, 3.7 * nRow), constrained_layout=True) - + for iRow in range(nRow): ncfile = Dataset('output_' + str(iRow + 1) + '.nc', 'r') var = ncfile.variables['temperature'] @@ -89,11 +89,11 @@ def _plot(nx, ny, filename, nus): for iCol in range(nCol): ax = axs[iRow, iCol] dis = ax.imshow( - var[iTime[iCol], 0::4, :].T, - extent=[0, 250, 500, 0], - aspect='0.5', - cmap='jet', - vmin=10, + var[iTime[iCol], 0::4, :].T, + extent=[0, 250, 500, 0], + aspect='0.5', + cmap='jet', + vmin=10, vmax=20) if iRow == nRow - 1: ax.set_xlabel('x, km') diff --git a/compass/ocean/tests/internal_wave/ten_day_test/__init__.py b/compass/ocean/tests/internal_wave/ten_day_test/__init__.py index 976545fa6d..13848ba5ce 100644 --- a/compass/ocean/tests/internal_wave/ten_day_test/__init__.py +++ b/compass/ocean/tests/internal_wave/ten_day_test/__init__.py @@ -41,4 +41,3 @@ def validate(self): compare_variables(test_case=self, variables=['layerThickness', 'normalVelocity'], filename1='forward/output.nc') - diff --git a/compass/ocean/tests/internal_wave/viz.py b/compass/ocean/tests/internal_wave/viz.py index 49e57557e7..5e5544220f 100644 --- a/compass/ocean/tests/internal_wave/viz.py +++ b/compass/ocean/tests/internal_wave/viz.py @@ -33,17 +33,17 @@ def run(self): """ ds = xarray.open_dataset('output.nc') - figsize = [6.4,4.8] + figsize = [6.4, 4.8] markersize = 20 if 'Time' not in ds.dims: print('Dataset missing time dimension') return - nSteps = ds.sizes['Time'] # number of timesteps + nSteps = ds.sizes['Time'] # number of timesteps tend = nSteps - 1 - for j in [0,tend]: + for j in [0, tend]: ds1 = ds.isel(Time=j) - + # prep all variables for uNormal plot ds1 = ds1.sortby('yEdge') @@ -57,10 +57,10 @@ def run(self): xCell = ds1.xCell yCell = numpy.zeros((nCells)) yCell = ds1.yCell - + xEdge_mid = numpy.median(xEdge) - edgeMask_x = numpy.equal(xEdge,xEdge_mid) - + edgeMask_x = numpy.equal(xEdge, xEdge_mid) + zIndex = xarray.DataArray(data=numpy.arange(nVertLevels), dims='nVertLevels') @@ -74,122 +74,131 @@ def run(self): zMid = numpy.zeros((nCells, nVertLevels)) for zIndex in range(nVertLevels): - zMid[:, zIndex] = (zInterface[:, zIndex] + - numpy.divide(zInterface[:, zIndex + 1] - - zInterface[:, zIndex ],2.)) - - # Solve for lateral boundaries of uNormal at cell centers for x-section + zMid[:, zIndex] = (zInterface[:, zIndex] + + numpy.divide(zInterface[:, zIndex + 1] + - zInterface[:, zIndex], 2.)) + + # Solve for lateral boundaries of uNormal at cell centers for + # x-section cellsOnEdge = ds1.cellsOnEdge - cellsOnEdge_x = cellsOnEdge[edgeMask_x,:] + cellsOnEdge_x = cellsOnEdge[edgeMask_x, :] yEdges = numpy.zeros((len(cellsOnEdge_x)+1)) for i in range(len(cellsOnEdge_x)): - if cellsOnEdge[i,1] == 0: - yEdges[i] = yCell[cellsOnEdge_x[i,0]-1] - yEdges[i+1] = yCell[cellsOnEdge_x[i,0]-1] - elif cellsOnEdge[i,1]== 0: - yEdges[i] = yCell[cellsOnEdge_x[i,1]-1] - yEdges[i+1] = yCell[cellsOnEdge_x[i,1]-1] + if cellsOnEdge[i, 1] == 0: + yEdges[i] = yCell[cellsOnEdge_x[i, 0] - 1] + yEdges[i+1] = yCell[cellsOnEdge_x[i, 0] - 1] + elif cellsOnEdge[i, 1] == 0: + yEdges[i] = yCell[cellsOnEdge_x[i, 1] - 1] + yEdges[i+1] = yCell[cellsOnEdge_x[i, 1] - 1] else: - yEdges[i] = min(yCell[cellsOnEdge_x[i,0]-1], - yCell[cellsOnEdge_x[i,1]-1]) - yEdges[i+1] = max(yCell[cellsOnEdge_x[i,0]-1], - yCell[cellsOnEdge_x[i,1]-1]) - - zInterfaces_mesh,yEdges_mesh = numpy.meshgrid(zInterface[0,:],yEdges) - + yEdges[i] = min(yCell[cellsOnEdge_x[i, 0] - 1], + yCell[cellsOnEdge_x[i, 1] - 1]) + yEdges[i+1] = max(yCell[cellsOnEdge_x[i, 0] - 1], + yCell[cellsOnEdge_x[i, 1] - 1]) + + zInterfaces_mesh, yEdges_mesh = numpy.meshgrid(zInterface[0, :], + yEdges) + normalVelocity = numpy.zeros((nCells, nVertLevels)) normalVelocity = ds1.normalVelocity - normalVelocity_xmesh = normalVelocity[edgeMask_x,:] - + normalVelocity_xmesh = normalVelocity[edgeMask_x, :] + # Figures plt.figure(figsize=figsize, dpi=100) cmax = numpy.max(numpy.abs(normalVelocity_xmesh.values)) - plt.pcolormesh(numpy.divide(yEdges_mesh,1e3), - zInterfaces_mesh, + plt.pcolormesh(numpy.divide(yEdges_mesh, 1e3), + zInterfaces_mesh, normalVelocity_xmesh.values, cmap='RdBu', vmin=-1.*cmax, vmax=cmax) plt.xlabel('y (km)') plt.ylabel('z (m)') cbar = plt.colorbar() cbar.ax.set_title('uNormal (m/s)') - plt.savefig('uNormal_depth_section_t{}.png'.format(j), bbox_inches='tight', dpi=200) + plt.savefig('uNormal_depth_section_t{}.png'.format(j), + bbox_inches='tight', dpi=200) plt.close() - - # --------------------------------------------------------------------- + + # ------------------------------------------------------------------ # Plot cell-centered variables - # --------------------------------------------------------------------- + # ------------------------------------------------------------------ # Prep variables for cell quantities - cellIndex = numpy.subtract(cellsOnEdge_x[1:,0],1) + cellIndex = numpy.subtract(cellsOnEdge_x[1:, 0], 1) yEdge = numpy.zeros((nEdges)) yEdge = ds1.yEdge yEdge_x = yEdge[edgeMask_x] - - zInterfaces_mesh,yCells_mesh = numpy.meshgrid(zInterface[0,:],yEdge_x) - + + zInterfaces_mesh, yCells_mesh = numpy.meshgrid(zInterface[0, :], + yEdge_x) + # Import cell quantities layerThickness = numpy.zeros((nCells, nVertLevels)) layerThickness = ds1.layerThickness - layerThickness_x = layerThickness[cellIndex,:] + layerThickness_x = layerThickness[cellIndex, :] temperature = numpy.zeros((nCells, nVertLevels)) temperature = ds1.temperature temperature_z = temperature.mean(dim='nCells') zMid_z = zMid.mean(axis=0) - temperature_x = temperature[cellIndex,:] + temperature_x = temperature[cellIndex, :] salinity = numpy.zeros((nCells, nVertLevels)) salinity = ds1.salinity - salinity_x = salinity[cellIndex,:] + salinity_x = salinity[cellIndex, :] w = numpy.zeros((nCells, nVertLevels)) w = ds1.vertVelocityTop - w_x = w[cellIndex,:] - + w_x = w[cellIndex, :] + # Figures plt.figure(figsize=figsize, dpi=100) - plt.plot(temperature_z.values,zMid_z) + plt.plot(temperature_z.values, zMid_z) plt.xlabel('PT (C)') plt.ylabel('z (m)') - plt.savefig('pt_depth_t{}.png'.format(j), bbox_inches='tight', dpi=200) + plt.savefig('pt_depth_t{}.png'.format(j), + bbox_inches='tight', dpi=200) plt.close() - + plt.figure(figsize=figsize, dpi=100) - plt.pcolormesh(numpy.divide(yCells_mesh,1e3), - zInterfaces_mesh, - temperature_x.values,cmap='viridis') + plt.pcolormesh(numpy.divide(yCells_mesh, 1e3), + zInterfaces_mesh, + temperature_x.values, cmap='viridis') plt.xlabel('y (km)') plt.ylabel('z (m)') cbar = plt.colorbar() cbar.ax.set_title('PT (C)') - plt.savefig('pt_depth_section_t{}.png'.format(j), bbox_inches='tight', dpi=200) + plt.savefig('pt_depth_section_t{}.png'.format(j), + bbox_inches='tight', dpi=200) plt.close() - + plt.figure(figsize=figsize, dpi=100) - plt.pcolormesh(numpy.divide(yCells_mesh,1e3), - zInterfaces_mesh, - salinity_x.values,cmap='viridis') + plt.pcolormesh(numpy.divide(yCells_mesh, 1e3), + zInterfaces_mesh, + salinity_x.values, cmap='viridis') plt.xlabel('y (km)') plt.ylabel('z (m)') cbar = plt.colorbar() cbar.ax.set_title('SA (g/kg)') - plt.savefig('sa_depth_section_t{}.png'.format(j), bbox_inches='tight', dpi=200) + plt.savefig('sa_depth_section_t{}.png'.format(j), + bbox_inches='tight', dpi=200) plt.close() plt.figure(figsize=figsize, dpi=100) - plt.pcolormesh(numpy.divide(yCells_mesh,1e3), - zInterfaces_mesh, - w_x.values,cmap='viridis') + plt.pcolormesh(numpy.divide(yCells_mesh, 1e3), + zInterfaces_mesh, + w_x.values, cmap='viridis') plt.xlabel('y (km)') plt.ylabel('z (m)') cbar = plt.colorbar() cbar.ax.set_title('h (m)') - plt.savefig('w_depth_section_t{}.png'.format(j), bbox_inches='tight', dpi=200) + plt.savefig('w_depth_section_t{}.png'.format(j), + bbox_inches='tight', dpi=200) plt.close() plt.figure(figsize=figsize, dpi=100) - plt.pcolormesh(numpy.divide(yCells_mesh,1e3), - zInterfaces_mesh, - layerThickness_x.values,cmap='viridis') + plt.pcolormesh(numpy.divide(yCells_mesh, 1e3), + zInterfaces_mesh, + layerThickness_x.values, cmap='viridis') plt.xlabel('y (km)') plt.ylabel('z (m)') cbar = plt.colorbar() cbar.ax.set_title('h (m)') - plt.savefig('layerThickness_depth_section_t{}.png'.format(j), bbox_inches='tight', dpi=200) + plt.savefig('layerThickness_depth_section_t{}.png'.format(j), + bbox_inches='tight', dpi=200) plt.close()