Precommit changes

MPAS-Dev · Mar 13, 2023 · 87d24d8 · 87d24d8
1 parent 7e9f59d
commit 87d24d8
Show file tree

Hide file tree

Showing 4 changed files with 56 additions and 70 deletions.
diff --git a/compass/ocean/tests/turbulence_closure/__init__.py b/compass/ocean/tests/turbulence_closure/__init__.py
@@ -1,7 +1,7 @@
-from compass.testgroup import TestGroup
 from compass.ocean.tests.turbulence_closure.decomp_test import DecompTest
 from compass.ocean.tests.turbulence_closure.default import Default
 from compass.ocean.tests.turbulence_closure.restart_test import RestartTest
+from compass.testgroup import TestGroup
 
 
 class TurbulenceClosure(TestGroup):
@@ -23,12 +23,14 @@ def __init__(self, mpas_core):
         for resolution in [1, 2, 1e4]:
             for forcing in ['cooling', 'evaporation']:
                 self.add_test_case(
-                    Default(test_group=self, resolution=resolution, forcing=forcing))
+                    Default(test_group=self, resolution=resolution,
+                            forcing=forcing))
 
 
 def configure(resolution, forcing, config):
     """
-    Modify the configuration options for one of the turbulence closure test cases
+    Modify the configuration options for one of the turbulence closure test
+    cases
 
     Parameters
     ----------
@@ -61,7 +63,6 @@ def configure(resolution, forcing, config):
         vert_levels = 128
         bottom_depth = 128.0
 
-
     config.set('turbulence_closure', 'nx', f'{nx}')
     config.set('turbulence_closure', 'ny', f'{ny}')
     config.set('turbulence_closure', 'dc', f'{resolution}')
@@ -71,12 +72,14 @@ def configure(resolution, forcing, config):
     if forcing == 'cooling':
         config.set('turbulence_closure', 'surface_heat_flux', '-100')
         config.set('turbulence_closure', 'surface_freshwater_flux', '0')
-        config.set('turbulence_closure', 'interior_temperature_gradient', '0.1')
+        config.set('turbulence_closure', 'interior_temperature_gradient',
+                   '0.1')
         config.set('turbulence_closure', 'interior_salinity_gradient', '0')
         config.set('turbulence_closure', 'wind_stress_zonal', '0')
     if forcing == 'evaporation':
         config.set('turbulence_closure', 'surface_heat_flux', '0')
         config.set('turbulence_closure', 'surface_freshwater_flux', '0.429')
         config.set('turbulence_closure', 'interior_temperature_gradient', '0')
-        config.set('turbulence_closure', 'interior_salinity_gradient', '-0.025')
+        config.set('turbulence_closure', 'interior_salinity_gradient',
+                   '-0.025')
         config.set('turbulence_closure', 'wind_stress_zonal', '0')
diff --git a/compass/ocean/tests/turbulence_closure/default/__init__.py b/compass/ocean/tests/turbulence_closure/default/__init__.py
@@ -1,8 +1,8 @@
-from compass.testcase import TestCase
-from compass.ocean.tests.turbulence_closure.initial_state import InitialState
+from compass.ocean.tests import turbulence_closure
 from compass.ocean.tests.turbulence_closure.forward import Forward
+from compass.ocean.tests.turbulence_closure.initial_state import InitialState
 from compass.ocean.tests.turbulence_closure.viz import Viz
-from compass.ocean.tests import turbulence_closure
+from compass.testcase import TestCase
 
 
 class Default(TestCase):
@@ -45,21 +45,24 @@ def __init__(self, test_group, resolution, forcing='cooling'):
 
         if resolution <= 5:
             ntasks = 128
-            plot_comparison=True
+            plot_comparison = True
         else:
             ntasks = 4
-            plot_comparison=False
+            plot_comparison = False
 
         self.add_step(
             InitialState(test_case=self, resolution=resolution))
         self.add_step(
-            Forward(test_case=self, ntasks=ntasks, openmp_threads=1, resolution=resolution))
-        self.add_step(Viz(test_case=self, resolution=resolution, forcing=forcing, do_comparison=plot_comparison))
+            Forward(test_case=self, ntasks=ntasks, openmp_threads=1,
+                    resolution=resolution))
+        self.add_step(Viz(test_case=self, resolution=resolution,
+                          forcing=forcing, do_comparison=plot_comparison))
 
     def configure(self):
         """
         Modify the configuration options for this test case.
         """
-        turbulence_closure.configure(self.resolution, self.forcing, self.config)
+        turbulence_closure.configure(self.resolution, self.forcing,
+                                     self.config)
 
     # no run() is needed because we're doing the default: running all steps
diff --git a/compass/ocean/tests/turbulence_closure/initial_state.py b/compass/ocean/tests/turbulence_closure/initial_state.py
@@ -1,14 +1,13 @@
-import xarray as xr
-import numpy as np
 import matplotlib.pyplot as plt
-
-from compass.ocean.tests.turbulence_closure.boundary_values import add_boundary_arrays
-from random import random, seed
-
-from mpas_tools.planar_hex import make_planar_hex_mesh
+import numpy as np
+import xarray as xr
 from mpas_tools.io import write_netcdf
 from mpas_tools.mesh.conversion import convert, cull
+from mpas_tools.planar_hex import make_planar_hex_mesh
 
+from compass.ocean.tests.turbulence_closure.boundary_values import (
+    add_boundary_arrays,
+)
 from compass.ocean.vertical import init_vertical_coord
 from compass.step import Step
 
@@ -61,12 +60,18 @@ def run(self):
         interior_temperature_shape = section.get('interior_temperature_shape')
         interior_salinity_shape = section.get('interior_salinity_shape')
         surface_salinity = section.getfloat('surface_salinity')
-        mixed_layer_salinity_gradient = section.getfloat('mixed_layer_salinity_gradient')
-        mixed_layer_temperature_gradient = section.getfloat('mixed_layer_temperature_gradient')
-        mixed_layer_depth_salinity = section.getfloat('mixed_layer_depth_salinity')
-        mixed_layer_depth_temperature = section.getfloat('mixed_layer_depth_temperature')
-        interior_salinity_gradient = section.getfloat('interior_salinity_gradient')
-        interior_temperature_gradient = section.getfloat('interior_temperature_gradient')
+        mixed_layer_salinity_gradient = \
+            section.getfloat('mixed_layer_salinity_gradient')
+        mixed_layer_temperature_gradient = \
+            section.getfloat('mixed_layer_temperature_gradient')
+        mixed_layer_depth_salinity = \
+            section.getfloat('mixed_layer_depth_salinity')
+        mixed_layer_depth_temperature = \
+            section.getfloat('mixed_layer_depth_temperature')
+        interior_salinity_gradient = \
+            section.getfloat('interior_salinity_gradient')
+        interior_temperature_gradient = \
+            section.getfloat('interior_temperature_gradient')
         interior_temperature_c0 = section.getfloat('interior_temperature_c0')
         interior_temperature_c1 = section.getfloat('interior_temperature_c1')
         interior_temperature_c2 = section.getfloat('interior_temperature_c2')
@@ -86,7 +91,8 @@ def run(self):
         coriolis_parameter = section.getfloat('coriolis_parameter')
         bottom_depth = config.getfloat('vertical_grid', 'bottom_depth')
 
-        dsMesh = make_planar_hex_mesh(nx=nx, ny=ny, dc=dc, nonperiodic_x=x_nonperiodic,
+        dsMesh = make_planar_hex_mesh(nx=nx, ny=ny, dc=dc,
+                                      nonperiodic_x=x_nonperiodic,
                                       nonperiodic_y=y_nonperiodic)
 
         write_netcdf(dsMesh, 'base_mesh.nc')
@@ -99,15 +105,15 @@ def run(self):
         ds = dsMesh.copy()
         dsForcing = dsMesh.copy()
         xCell = ds.xCell
-        yCell = ds.yCell
 
         ds['bottomDepth'] = bottom_depth * xr.ones_like(xCell)
         ds['ssh'] = xr.zeros_like(xCell)
 
         init_vertical_coord(config, ds)
 
         dsForcing['windStressZonal'] = wind_stress_zonal * xr.ones_like(xCell)
-        dsForcing['windStressMeridional'] = wind_stress_meridional * xr.ones_like(xCell)
+        dsForcing['windStressMeridional'] = \
+            wind_stress_meridional * xr.ones_like(xCell)
         dsForcing['sensibleHeatFlux'] = surface_heat_flux * xr.ones_like(xCell)
         dsForcing['rainFlux'] = surface_freshwater_flux * xr.ones_like(xCell)
 
@@ -124,7 +130,8 @@ def run(self):
         zMid = ds.refZMid.values
 
         temperature_at_mld = (surface_temperature -
-            mixed_layer_temperature_gradient * mixed_layer_depth_temperature)
+                              mixed_layer_temperature_gradient *
+                              mixed_layer_depth_temperature)
         if interior_temperature_shape == 'linear':
             initial_temperature = np.where(
                 zMid >= -mixed_layer_depth_temperature,
@@ -144,7 +151,8 @@ def run(self):
             print('interior_temperature_shape is not supported')
 
         salinity_at_mld = (surface_salinity -
-            mixed_layer_salinity_gradient * mixed_layer_depth_salinity)
+                           mixed_layer_salinity_gradient *
+                           mixed_layer_depth_salinity)
         if interior_salinity_shape == 'linear':
             initial_salinity = np.where(
                 zMid >= -mixed_layer_depth_salinity,
@@ -162,38 +170,13 @@ def run(self):
                 interior_salinity_c5 * zMid**5.)
         else:
             print('interior_salinity_shape is not supported')
-#        surf_indices = np.where(ds.refZMid.values >= -mixed_layer_depth_temperature)[0]
-#
-#        if len(surf_indices) > 0:
-#            temperature[0,:,surf_indices] = surface_temperature + mixed_layer_temperature_gradient* \
-#			                                ds.zMid[0,:,surf_indices].values
-#
-#        int_indices = np.where(ds.refZMid.values < -mixed_layer_depth_temperature)[0]
-#
-#        if len(int_indices) > 0:
-#            temperature[0,:,int_indices] = surface_temperature - mixed_layer_temperature_gradient* \
-#													 mixed_layer_depth_temperature + interior_temperature_gradient* \
-#													 (ds.zMid[0,:,int_indices] + \
-#													  mixed_layer_depth_temperature)
-#
-#        temperature[0,:,0] += disturbance_amplitude*2*(np.random.rand(len(ds.xCell.values)) - 0.5)
-#
-#        surf_indices = np.where(ds.refZMid.values >= -mixed_layer_depth_salinity)[0]
-#        if len(surf_indices) > 0:
-#            salinity[0,:,surf_indices] = surface_salinity - mixed_layer_salinity_gradient * \
-#			                               	ds.zMid[0,:,surf_indices]
-#
-#        int_indices = np.where(ds.refZMid.values < -mixed_layer_depth_salinity)[0]
-#        if len(int_indices) > 0:
-#            salinity[0,:,int_indices] = surface_salinity - mixed_layer_salinity_gradient* \
-#									  				  mixed_layer_depth_salinity + interior_salinity_gradient* \
-#													  (ds.zMid[0,:,int_indices] + \
-#													   mixed_layer_depth_salinity)
 
         normalVelocity = normalVelocity.expand_dims(dim='Time', axis=0)
         ds['normalVelocity'] = normalVelocity
 
         temperature[0, :, :] = initial_temperature
+        temperature[0, :, 0] += disturbance_amplitude * 2. * \
+            (np.random.rand(len(ds.xCell.values)) - 0.5)
         salinity[0, :, :] = initial_salinity
         ds['temperature'] = temperature
         ds['salinity'] = salinity
@@ -209,14 +192,14 @@ def run(self):
         plt.plot(initial_temperature, zMid, 'k-')
         plt.xlabel('PT (C)')
         plt.ylabel('z (m)')
-        plt.savefig(f'pt_depth_t0h.png',
+        plt.savefig('pt_depth_t0h.png',
                     bbox_inches='tight', dpi=200)
         plt.close()
 
         plt.figure(dpi=100)
         plt.plot(initial_salinity, zMid, 'k-')
         plt.xlabel('SA (g/kg)')
         plt.ylabel('z (m)')
-        plt.savefig(f'sa_depth_t0h.png',
+        plt.savefig('sa_depth_t0h.png',
                     bbox_inches='tight', dpi=200)
         plt.close()
diff --git a/compass/ocean/tests/turbulence_closure/viz.py b/compass/ocean/tests/turbulence_closure/viz.py
@@ -1,6 +1,6 @@
-import xarray as xr
-import numpy as np
 import matplotlib.pyplot as plt
+import numpy as np
+import xarray as xr
 
 from compass.step import Step
 
@@ -25,7 +25,6 @@ def __init__(self, test_case, resolution, forcing, do_comparison=False):
         self.add_input_file(filename='mesh.nc',
                             target='../initial_state/culled_mesh.nc')
 
-
         if do_comparison:
             if forcing == 'cooling':
                 forcing_name = 'c02'
@@ -68,9 +67,9 @@ def run(self):
         else:
             # This routine is not generic but should not be used as
             # daysSinceStartOfSim is included in the streams file
-            time0 = 2.0 + (7.0/24.0)
-            dt = 0.5/24.0
-            time = np.linspace(time0 + dt, time0 + nt*dt, num=nt)
+            time0 = 2.0 + (7.0 / 24.0)
+            dt = 0.5 / 24.0
+            time = np.linspace(time0 + dt, time0 + nt * dt, num=nt)
 
         if self.do_comparison:
             time_palm = ds_palm.time
@@ -87,8 +86,6 @@ def run(self):
         ds = ds.sortby('yEdge')
 
         # Get mesh variables
-        xEdge = dsInit.xEdge
-        yEdge = dsInit.yEdge
         xCell = dsInit.xCell
         yCell = dsInit.yCell
         zMid = dsInit.refZMid
@@ -188,7 +185,7 @@ def run(self):
         plt.scatter(np.divide(xCell, 1e3),
                     np.divide(yCell, 1e3),
                     s=markersize, c=w_zTop.values,
-                    cmap='cmo.balance', vmin=-1.*cmax, vmax=cmax)
+                    cmap='cmo.balance', vmin=-1. * cmax, vmax=cmax)
         plt.xlabel('x (km)')
         plt.ylabel('y (km)')
         cbar = plt.colorbar()