Merge branch 'feature/535_stratification_diag' of https://github.com/…

…British-Oceanographic-Data-Centre/COAsT into feature/535_stratification_diag

coast/data/profile.py

@@ -156,16 +156,16 @@ def subset_indices_lonlat_box(self, lonbounds, latbounds):
         if lonbounds[0] < lonbounds[1]:
             ind = general_utils.subset_indices_lonlat_box(
                 self.dataset.longitude, self.dataset.latitude, lonbounds[0], lonbounds[1], latbounds[0], latbounds[1]
-        )
+            )
         else:
             ind1 = general_utils.subset_indices_lonlat_box(
-                self.dataset.longitude, self.dataset.latitude, lonbounds[0], 180.0 , latbounds[0], latbounds[1]
+                self.dataset.longitude, self.dataset.latitude, lonbounds[0], 180.0, latbounds[0], latbounds[1]
             )
             ind2 = general_utils.subset_indices_lonlat_box(
                 self.dataset.longitude, self.dataset.latitude, -180.0, lonbounds[1], latbounds[0], latbounds[1]
-        )
-            ind={}
-            ind[0] = np.concatenate((ind1[0],ind2[0]))
+            )
+            ind = {}
+            ind[0] = np.concatenate((ind1[0], ind2[0]))
         return Profile(dataset=self.dataset.isel(id_dim=ind[0]))
 
     def extract_en4_profiles(self, dataset_names, region_bounds, chunks: dict = {}):
@@ -538,7 +538,7 @@ def match_to_grid(self, gridded, limits=[0, 0, 0, 0], rmax=25.0) -> None:
         grd = gridded.dataset
         if "bottom_level" in grd:
             grd["landmask"] = grd.bottom_level == 0
-        else: #resort to using bathymetry
+        else:  # resort to using bathymetry
             grd["landmask"] = grd.bathymetry == 0
 
         lon_prf = prf["longitude"]
@@ -608,7 +608,7 @@ def match_to_grid(self, gridded, limits=[0, 0, 0, 0], rmax=25.0) -> None:
         self.dataset["rmin_prf"] = xr.DataArray(rmin_prf, dims=["id_dim", "NNs"])
         self.dataset["ind_good"] = xr.DataArray(ind_good, dims=["Ngood"])
 
-    def gridded_to_profile_2d(self, gridded, variable,limits=[0,0,0,0],rmax=25.0) -> None:
+    def gridded_to_profile_2d(self, gridded, variable, limits=[0, 0, 0, 0], rmax=25.0) -> None:
         """
         Evaluated a gridded data variable on each profile. Here just 2D, but could be extended to 3 or 4D
 
@@ -623,13 +623,13 @@ def gridded_to_profile_2d(self, gridded, variable,limits=[0,0,0,0],rmax=25.0) ->
         """
         # ensure there are indices in profile
         if not "ind_x" in self.dataset:
-            self.match_to_grid(gridded,limits=limits,rmax=rmax)
+            self.match_to_grid(gridded, limits=limits, rmax=rmax)
         #
         prf = self.dataset
         grd = gridded.dataset
         if "botton_level" in grd:
             grd["landmask"] = grd.bottom_level == 0
-        else: # resort to bathymetry for mask
+        else:  # resort to bathymetry for mask
             grd["landmask"] = grd.bathymetry == 0
 
         nprof = self.dataset.id_dim.shape[0]

coast/diagnostics/profile_stratification.py

@@ -39,7 +39,7 @@ def __init__(self, profile: xr.Dataset):
         self.nz = profile.dataset.dims["z_dim"]
         debug(f"Initialised {get_slug(self)}")
 
-    def clean_data(profile: xr.Dataset, gridded: xr.Dataset, Zmax,limits=[0,0,0,0],rmax=25.):
+    def clean_data(profile: xr.Dataset, gridded: xr.Dataset, Zmax, limits=[0, 0, 0, 0], rmax=25.0):
         """
         Cleaning data for stratification metric calculations
         Stage 1:...
@@ -68,7 +68,7 @@ def first_nonzero(arr, axis=0, invalid_val=np.nan):
             return np.where(mask.any(axis=axis), mask.argmax(axis=axis), invalid_val)
 
         if "bathymetry" in gridded.dataset:
-            profile.gridded_to_profile_2d(gridded, "bathymetry",limits=limits,rmax=rmax)
+            profile.gridded_to_profile_2d(gridded, "bathymetry", limits=limits, rmax=rmax)
             D_prf = profile.dataset.bathymetry.values
             z = profile.dataset.depth
             test_surface = z < np.minimum(dz_max, 0.25 * np.repeat(D_prf[:, np.newaxis], n_depth, axis=1))
@@ -143,7 +143,7 @@ def first_nonzero(arr, axis=0, invalid_val=np.nan):
         # %%
         return profile
 
-    def calc_pea(self, profile: xr.Dataset, gridded: xr.Dataset, Zmax, rmax=25.0, limits=[0,0,0,0]):
+    def calc_pea(self, profile: xr.Dataset, gridded: xr.Dataset, Zmax, rmax=25.0, limits=[0, 0, 0, 0]):
         """
         Calculates Potential Energy Anomaly