liniting fixes

MPAS-Dev · Nov 1, 2024 · 92623c0 · 92623c0
1 parent b427336
commit 92623c0
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 31 deletions.
diff --git a/compass/ocean/tests/utility/create_salin_restoring/__init__.py b/compass/ocean/tests/utility/create_salin_restoring/__init__.py
@@ -1,5 +1,9 @@
-from compass.ocean.tests.utility.create_salin_restoring.salinity_restoring import Salinity
-from compass.ocean.tests.utility.create_salin_restoring.extrap_salin import ExtrapSalin
+from compass.ocean.tests.utility.create_salin_restoring.salinity_restoring import (
+    Salinity,
+)
+from compass.ocean.tests.utility.create_salin_restoring.extrap_salin import (
+    ExtrapSalin,
+)
 from compass.testcase import TestCase
 
 

diff --git a/compass/ocean/tests/utility/create_salin_restoring/extrap_salin.py b/compass/ocean/tests/utility/create_salin_restoring/extrap_salin.py
@@ -56,6 +56,7 @@ def run(self):
         # extrapolate horizontally using the ocean mask
         _extrap(self.woa_filename)
 
+
 def _extrap(out_filename):
 
     in_filename = 'woa_surface_salinity_monthly.nc'
@@ -75,12 +76,12 @@ def _extrap(out_filename):
     lon_no_halo = list(range(2, nlon + 2))
 
     for i in range(12):
-        valid = np.isfinite(field[i,:,:])
+        valid = np.isfinite(field[i, :, :])
         orig_mask = valid
         prev_fill_count = 0
         while True:
             valid_weight_sum = _extrap_with_halo(valid, kernel, valid,
-                                               lon_with_halo, lon_no_halo)
+                                                 lon_with_halo, lon_no_halo)
 
             new_valid = valid_weight_sum > threshold
 
@@ -93,9 +94,9 @@ def _extrap(out_filename):
                 # no change so we're done
                 break
 
-            field_extrap = _extrap_with_halo(field[i,:,:], kernel, valid,
-                                           lon_with_halo, lon_no_halo)
-            field[i,fill_mask] = field_extrap[fill_mask]/valid_weight_sum[fill_mask]
+            field_extrap = _extrap_with_halo(field[i, :, :], kernel, valid,
+                                             lon_with_halo, lon_no_halo)
+            field[i, fill_mask] = field_extrap[fill_mask] / valid_weight_sum[fill_mask]
 
             valid = new_valid
             prev_fill_count = fill_count
@@ -107,6 +108,7 @@ def _extrap(out_filename):
 
     ds.to_netcdf(out_filename)
 
+
 def _extrap_with_halo(field, kernel, valid, lon_with_halo, lon_no_halo):
     field = field.copy()
     field[np.logical_not(valid)] = 0.

diff --git a/compass/ocean/tests/utility/create_salin_restoring/salinity_restoring.py b/compass/ocean/tests/utility/create_salin_restoring/salinity_restoring.py
@@ -18,10 +18,12 @@ def __init__(self, test_case):
 
         Parameters
         ----------
-        test_case : compass.ocean.tests.utility.create_salin_restoring.CreateSalinRestoring
-            The test case this step belongs to
+        test_case : compass.ocean.tests.utility.create_salin_restoring.
+        CreateSalinRestoring
+        The test case this step belongs to
         """
-        super().__init__(test_case, name='salinity_restoring', ntasks=1, min_tasks=1)
+        super().__init__(test_case, name='salinity_restoring', ntasks=1,
+                         min_tasks=1)
         self.add_output_file(filename='woa_surface_salinity_monthly.nc')
 
     def setup(self):
@@ -36,22 +38,21 @@ def setup(self):
 
         woa_dir = 'salinity/netcdf/decav91C0/0.25'
 
-        woa_files = dict(
-                      jan='woa23_decav91C0_s01_04.nc',
-                      feb='woa23_decav91C0_s02_04.nc',
-                      mar='woa23_decav91C0_s03_04.nc',
-                      apr='woa23_decav91C0_s04_04.nc',
-                      may='woa23_decav91C0_s05_04.nc',
-                      jun='woa23_decav91C0_s06_04.nc',
-                      jul='woa23_decav91C0_s07_04.nc',
-                      aug='woa23_decav91C0_s08_04.nc',
-                      sep='woa23_decav91C0_s09_04.nc',
-                      octo='woa23_decav91C0_s10_04.nc',
-                      nov='woa23_decav91C0_s11_04.nc',
-                      dec='woa23_decav91C0_s12_04.nc')
-
-        for month in ['jan', 'feb', 'mar', 'apr', 'may', 'jun', 'jul', 'aug', 'sep',
-			          'octo', 'nov', 'dec']:
+        woa_files = dict(jan='woa23_decav91C0_s01_04.nc',
+                         feb='woa23_decav91C0_s02_04.nc',
+                         mar='woa23_decav91C0_s03_04.nc',
+                         apr='woa23_decav91C0_s04_04.nc',
+                         may='woa23_decav91C0_s05_04.nc',
+                         jun='woa23_decav91C0_s06_04.nc',
+                         jul='woa23_decav91C0_s07_04.nc',
+                         aug='woa23_decav91C0_s08_04.nc',
+                         sep='woa23_decav91C0_s09_04.nc',
+                         octo='woa23_decav91C0_s10_04.nc',
+                         nov='woa23_decav91C0_s11_04.nc',
+                         dec='woa23_decav91C0_s12_04.nc')
+
+        for month in ['jan', 'feb', 'mar', 'apr', 'may', 'jun', 'jul',
+                      'aug', 'sep', 'octo', 'nov', 'dec']:
             woa_filename = woa_files[month]
             woa_url = f'{base_url}/{woa_dir}/{woa_filename}'
 
@@ -74,8 +75,8 @@ def run(self):
             ds_out[f'{var}_bnds'] = ds_jan[f'{var}_bnds']
 
         slices = list()
-        for month in ['jan', 'feb', 'mar', 'apr', 'may', 'jun', 'jul', 'aug', 'sep',
-                      'octo', 'nov', 'dec']:
+        for month in ['jan', 'feb', 'mar', 'apr', 'may', 'jun', 'jul',
+                      'aug', 'sep', 'octo', 'nov', 'dec']:
             ds = xr.open_dataset(
                     f'woa_salin_{month}.nc',
                     decode_times=False).isel(depth=0).drop_vars('depth')
@@ -85,8 +86,8 @@ def run(self):
         ds_out['s_an'].attrs = ds_jan['s_an'].attrs
 
         # Change names of variables and alter attributes
-        ds_out = ds_out.rename_dims({'time':'Time'})
-        ds_out = ds_out.rename_vars({'s_an':'SALT'})
+        ds_out = ds_out.rename_dims({'time': 'Time'})
+        ds_out = ds_out.rename_vars({'s_an': 'SALT'})
         ds_out = ds_out.drop_vars('time')
 
         # Create a time index
@@ -98,7 +99,7 @@ def run(self):
         for i in range(1, 13):
             xtime_string = f"0000-{i:02d}-15_00:00:00"
             xtime_list.append(xtime_string)
-        xtime_out = np.array(xtime_list,dtype='S64')
+        xtime_out = np.array(xtime_list, dtype='S64')
 
         ds_out = ds_out.assign(xtime=xr.DataArray(xtime_out, dims=['Time']))