From 92623c0a9c59ddadad077df7da17c188934ba28d Mon Sep 17 00:00:00 2001 From: Luke Van Roekel Date: Thu, 31 Oct 2024 22:50:36 -0500 Subject: [PATCH] liniting fixes --- .../create_salin_restoring/__init__.py | 8 ++- .../create_salin_restoring/extrap_salin.py | 12 +++-- .../salinity_restoring.py | 49 ++++++++++--------- 3 files changed, 38 insertions(+), 31 deletions(-) diff --git a/compass/ocean/tests/utility/create_salin_restoring/__init__.py b/compass/ocean/tests/utility/create_salin_restoring/__init__.py index 20c243c38..5ab2e1009 100644 --- 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 index c7add1dd1..d9e358e57 100644 --- 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 index 38d44b011..06e59255b 100644 --- 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']))