Repository: British-Oceanographic-Data-Centre/COAsT (language: python)

Issue:
Feature/map domain vars
add logic to map variable names in .nc to required names in code via json config


Discussion / hints:

Fix patch:
diff --git a/coast/__init__.py b/coast/__init__.py
index 0a92f024..0e9ad3a9 100644
--- a/coast/__init__.py
+++ b/coast/__init__.py
@@ -26,7 +26,7 @@
 from .data.opendap import OpendapInfo
 from .data.copernicus import Copernicus, Product
 from ._utils.experiments_file_handling import experiments
-from ._utils.experiments_file_handling import nemo_filenames
+from ._utils.experiments_file_handling import nemo_filename_maker
 from .diagnostics.hydrographic_profiles import Hydrographic_Profiles
 
 # Set default for logging level when coast is imported
diff --git a/coast/_utils/experiments_file_handling.py b/coast/_utils/experiments_file_handling.py
index dcddf9cb..db09c8b7 100644
--- a/coast/_utils/experiments_file_handling.py
+++ b/coast/_utils/experiments_file_handling.py
@@ -1,4 +1,7 @@
-"""Set of functions to control basic experimnet file handling"""
+"""Set of functions to control basic experiment file handling"""
+
+import json
+from typing import List
 import numpy as np
 
 
@@ -22,8 +25,6 @@ def experiments(experiments="experiments.json"):
     exp_names,dirs,domains,file_names
 
     """
-    import json
-
     with open(experiments, "r") as j:
         json_content = json.loads(j.read())
         try:
@@ -43,7 +44,6 @@ def experiments(experiments="experiments.json"):
         except:
             file_names = []
 
-        # print(lengths)
         # check all non zero lengths are the same
         lengths = np.array([len(exp_names), len(dirs), len(domains), len(file_names)])
         if np.min(lengths[np.nonzero(lengths)[0]]) != np.max(lengths[np.nonzero(lengths)[0]]):
@@ -51,34 +51,26 @@ def experiments(experiments="experiments.json"):
     return exp_names, dirs, domains, file_names
 
 
-def nemo_filenames(dpath, runtype, ystart, ystop, grid="T"):
-    """
-    Creates a list of NEMO file names from a set of standard templates
-     Parameters
-     ----------
-     dpath : path to the files
-     runtype : hardwired set of standard nemo filenames
-     ystart : start year
-     ystop : stop year
-     grid (optional) NEMO grid type
+def nemo_filename_maker(directory, year_start: int, year_stop: int, grid: str = "T") -> List:
+    """Creates a list of NEMO file names from a set of standard templates.
 
-     -------
-     names :lis tof nemo file names
-    """
+    Args:
+        directory: path to the files'
+        year_start: start year
+        year_stop: stop year
+        grid: NEMO grid type defaults to T
 
+    Returns: a list of possible nemo file names
+
+    """
     # produce a list of nemo filenames
-    names = np.array([])
-    if runtype == "SENEMO":
-        for iy in range(ystart, ystop + 1):
-            for im in range(1, 12 + 1):
-                MNTH = str(im)
-                if im < 10:
-                    MNTH = "0" + MNTH
-                YEAR = str(iy)
-                new_name = "{0}/SENEMO_1m_{1}0101_{1}1231_grid_{3}_{1}{2}-{1}{2}.nc".format(dpath, YEAR, MNTH, grid)
-                names = np.append(names, new_name)
+    names = []
+    january = 1
+    december = 13  # range is non-inclusive so we need 12 + 1
+
+    for year in range(year_start, year_stop + 1):
+        for month in range(january, december):
+            new_name = f"{directory}/SENEMO_1m_{year}0101_{year}1231_grid_{grid}_{year}{month:02}-{year}{month:02}.nc"
+            names.append(new_name)
 
-    else:
-        print("Runtype: " + runtype + "not coded yet, returning empty list")
-        names = []
     return names
diff --git a/coast/data/gridded.py b/coast/data/gridded.py
index 4d05fe61..da320219 100644
--- a/coast/data/gridded.py
+++ b/coast/data/gridded.py
@@ -1,5 +1,6 @@
 """Gridded class"""
 import os.path as path_lib
+import re
 import warnings
 
 # from dask import delayed, compute, visualize
@@ -90,27 +91,23 @@ def _setup_grid_obj(self, chunks, multiple, **kwargs):
             if self.fn_data is not None:
                 dataset_domain = self.trim_domain_size(dataset_domain)
             self.set_timezero_depths(
-                dataset_domain
+                dataset_domain, **kwargs
             )  # THIS ADDS TO dataset_domain. Should it be 'return'ed (as in trim_domain_size) or is implicit OK?
             self.merge_domain_into_dataset(dataset_domain)
             debug(f"Initialised {get_slug(self)}")
-        ###################
-        Make_LonLat_2D = kwargs.get("Make_LonLat_2D", False)
-        if Make_LonLat_2D:
-            # Expand 1D latitude, longitude variables to 2D
-            # jth is there a lazy way of doing this?
-            if len(self.dataset.longitude.shape) == 1:
-                lat = self.dataset.latitude.values
-                lon = self.dataset.longitude.values
-                nx = self.dataset.longitude.size
-                ny = self.dataset.latitude.size
-                self.dataset["latitude"] = xr.DataArray(
-                    np.repeat(lat[:, np.newaxis], nx, axis=1), dims=["y_dim", "x_dim"]
-                )
 
-                self.dataset["longitude"] = xr.DataArray(
-                    np.repeat(lon[np.newaxis, :], ny, axis=0), dims=["y_dim", "x_dim"]
-                )
+    def make_lonLat_2d(self):
+        """Expand 1D latitude and longitude variables to 2D."""
+
+        # jth is there a lazy way of doing this?
+        if len(self.dataset.longitude.shape) == 1:
+            lat = self.dataset.latitude.values
+            lon = self.dataset.longitude.values
+            nx = self.dataset.longitude.size
+            ny = self.dataset.latitude.size
+            self.dataset["latitude"] = xr.DataArray(np.repeat(lat[:, np.newaxis], nx, axis=1), dims=["y_dim", "x_dim"])
+
+            self.dataset["longitude"] = xr.DataArray(np.repeat(lon[np.newaxis, :], ny, axis=0), dims=["y_dim", "x_dim"])
 
     def set_grid_vars(self):
         """Define the variables to map from the domain file to the NEMO obj"""
@@ -134,10 +131,19 @@ def load_domain(self, fn_domain, chunks):  # TODO Do something with this unused
                 dataset_domain = dataset_domain.rename_dims(mapping)
             except ValueError as err:
                 warning(
-                    f"{get_slug(self)}: Problem renaming dimension from {get_slug(self.dataset)}: {key} -> {value}."
+                    f"{get_slug(self)}: Problem renaming domain dimension from {get_slug(self.dataset)}: {key} -> {value}."
+                    f"{chr(10)}Error message of '{err}'"
+                )
+        # Rename domain variables.
+        for key, value in self.config.domain.variable_map.items():
+            mapping = {key: value}
+            try:
+                dataset_domain = dataset_domain.rename_vars(mapping)
+            except ValueError as err:
+                warning(
+                    f"{get_slug(self)}: Problem renaming domain variable from {get_slug(self.dataset)}: {key} -> {value}."
                     f"{chr(10)}Error message of '{err}'"
                 )
-
         return dataset_domain
 
     def merge_domain_into_dataset(self, dataset_domain):
@@ -183,49 +189,43 @@ def get_contour_complex(self, var, points_x, points_y, points_z, tolerance: int
         smaller = self.dataset[var].sel(z=points_z, x=points_x, y=points_y, method="nearest", tolerance=tolerance)
         return smaller
 
-    def set_timezero_depths(self, dataset_domain):
+    def set_timezero_depths(self, dataset_domain, calculate_bathymetry=False):
         """
         Calculates the depths at time zero (from the domain_cfg input file)
         for the appropriate grid.
         The depths are assigned to domain_dataset.depth_0
+
+        Args:
+            dataset_domain: a complex data object.
+            calculate_bathymetry: Flag that will either calculate bathymetry (true) or load it from dataset_domain file
+            (false).
         """
         debug(f"Setting timezero depths for {get_slug(self)} with {get_slug(dataset_domain)}")
         # keyword to allow calcution of bathymetry from scale factors
+
+        # All bathymetry should now be mapped to bathy_metry
         try:
-            calc_bathy = dataset_domain["calc_bathy"]
-        except:
-            calc_bathy = False
-        # jth NEMO bathymetry can be called bathymetry, bathy_metry or hbatt in various versions
-        try:
-            bathymetry = dataset_domain.bathy_metry.squeeze()
-        except:
-            try:
-                bathymetry = (
-                    dataset_domain.hbatt.squeeze()
-                )  # jth add a second option here better done with config file?
-            except:
-                try:
-                    bathymetry = (
-                        dataset_domain.bathymetry.squeeze()
-                    )  # jth add a second option here better done with config file??
-
-                except AttributeError as err:
-                    bathymetry = xr.zeros_like(dataset_domain.e1t.squeeze())
-                    (
-                        # jth warnign needs to change
-                        warnings.warn(
-                            f"The model domain loaded, '{self.filename_domain}', does not contain the "
-                            "bathy_metry' variable. This will result in the "
-                            "NEMO.dataset.bathymetry variable being set to zero, which "
-                            "may result in unexpected behaviour from routines that require "
-                            "this variable."
-                        )
-                    )
-                    debug(
-                        f"The bathy_metry variable was missing from the domain_cfg for "
-                        f"{get_slug(self)} with {get_slug(dataset_domain)}"
-                        f"{chr(10)}Error message of {err}"
-                    )
+            if calculate_bathymetry:  # calculate bathymetry from scale factors
+                bathymetry, mask, time_mask = self.calc_bathymetry(dataset_domain)
+            else:
+                bathymetry = dataset_domain.bathy_metry.squeeze()
+
+        except AttributeError as err:
+            bathymetry = xr.zeros_like(dataset_domain.e1.squeeze())
+            (
+                warnings.warn(
+                    f"The model domain loaded, '{self.filename_domain}', does not contain the "
+                    "bathy_metry' variable. This will result in the "
+                    "NEMO.dataset.bathymetry variable being set to zero, which "
+                    "may result in unexpected behaviour from routines that require "
+                    "this variable."
+                )
+            )
+            debug(
+                f"The bathy_metry variable was missing from the domain_cfg for "
+                f"{get_slug(self)} with {get_slug(dataset_domain)}"
+                f"{chr(10)}Error message of {err}"
+            )
 
         try:
             if self.grid_ref == "t-grid":
@@ -233,18 +233,12 @@ def set_timezero_depths(self, dataset_domain):
                 depth_0 = np.zeros_like(e3w_0)
                 depth_0[0, :, :] = 0.5 * e3w_0[0, :, :]
                 depth_0[1:, :, :] = depth_0[0, :, :] + np.cumsum(e3w_0[1:, :, :], axis=0)
-                # calculate bathymetry from scale factors
-                if calc_bathy:
-                    bathymetry = Gridded.calc_bathymetry(self, dataset_domain, bathymetry)
 
             elif self.grid_ref == "w-grid":
                 e3t_0 = np.squeeze(dataset_domain.e3t_0.values)
                 depth_0 = np.zeros_like(e3t_0)
                 depth_0[0, :, :] = 0.0
                 depth_0[1:, :, :] = np.cumsum(e3t_0, axis=0)[:-1, :, :]
-                # calculate bathymetry from scale factors
-                if calc_bathy:
-                    bathymetry = Gridded.calc_bathymetry(self, dataset_domain, bathymetry)
 
             elif self.grid_ref == "u-grid":
                 e3w_0 = dataset_domain.e3w_0.values.squeeze()
@@ -252,90 +246,103 @@ def set_timezero_depths(self, dataset_domain):
                 depth_0 = np.zeros_like(e3w_0)
                 depth_0[0, :, :-1] = 0.5 * e3w_0_on_u[0, :, :]
                 depth_0[1:, :, :-1] = depth_0[0, :, :-1] + np.cumsum(e3w_0_on_u[1:, :, :], axis=0)
-                if calc_bathy:
-                    bathymetry = Gridded.calc_bathymetry(self, dataset_domain, bathymetry)
-                else:  # jth only valid for pure sigma
+                if not calculate_bathymetry:  # jth only valid for pure sigma
                     bathymetry[:, :-1] = 0.5 * (bathymetry[:, :-1] + bathymetry[:, 1:])
+
             elif self.grid_ref == "v-grid":
                 e3w_0 = dataset_domain.e3w_0.values.squeeze()
                 e3w_0_on_v = 0.5 * (e3w_0[:, :-1, :] + e3w_0[:, 1:, :])
                 depth_0 = np.zeros_like(e3w_0)
                 depth_0[0, :-1, :] = 0.5 * e3w_0_on_v[0, :, :]
                 depth_0[1:, :-1, :] = depth_0[0, :-1, :] + np.cumsum(e3w_0_on_v[1:, :, :], axis=0)
-                if calc_bathy:
-                    bathymetry = Gridded.calc_bathymetry(self, dataset_domain, bathymetry)
-                else:  # jth only valid for pure sigma
+                if not calculate_bathymetry:
                     bathymetry[:-1, :] = 0.5 * (bathymetry[:-1, :] + bathymetry[1:, :])
+
             elif self.grid_ref == "f-grid":
                 e3w_0 = dataset_domain.e3w_0.values.squeeze()
                 e3w_0_on_f = 0.25 * (e3w_0[:, :-1, :-1] + e3w_0[:, :-1, 1:] + e3w_0[:, 1:, :-1] + e3w_0[:, 1:, 1:])
                 depth_0 = np.zeros_like(e3w_0)
                 depth_0[0, :-1, :-1] = 0.5 * e3w_0_on_f[0, :, :]
                 depth_0[1:, :-1, :-1] = depth_0[0, :-1, :-1] + np.cumsum(e3w_0_on_f[1:, :, :], axis=0)
-                if calc_bathy:
-                    bathymetry = Gridded.calc_bathymetry(self, dataset_domain, bathymetry)
-                else:
+                if not calculate_bathymetry:
                     bathymetry[:-1, :-1] = 0.25 * (
                         bathymetry[:-1, :-1] + bathymetry[:-1, 1:] + bathymetry[1:, :-1] + bathymetry[1:, 1:]
                     )
             else:
                 raise ValueError(str(self) + ": " + self.grid_ref + " depth calculation not implemented")
+
             # Write the depth_0 variable to the domain_dataset DataSet, with grid type
             dataset_domain[f"depth{self.grid_ref.replace('-grid', '')}_0"] = xr.DataArray(
                 depth_0,
                 dims=["z_dim", "y_dim", "x_dim"],
                 attrs={"units": "m", "standard_name": "Depth at time zero on the {}".format(self.grid_ref)},
             )
-
-            self.dataset["bathymetry"] = bathymetry
-            self.dataset["bathymetry"].attrs = {
-                "units": "m",
-                "standard_name": "bathymetry",
-                "description": "depth of last wet w-level on the horizontal {}".format(self.grid_ref),
-            }
+            self.dataset["bathymetry"] = xr.DataArray(
+                bathymetry,
+                dims=["y_dim", "x_dim"],
+                attrs={
+                    "units": "m",
+                    "standard_name": "bathymetry",
+                    "description": "depth of last wet w-level on the horizontal {}".format(self.grid_ref),
+                },
+            )
         except ValueError as err:
+            print(err)
             error(err)
 
-    def calc_bathymetry(self, dataset_domain, bathymetry):
+    def calc_bathymetry(self, dataset_domain):
         """
         NEMO approach to defining bathymetry by summing scale factors at various
         grid locations.
-        Works with z-coordinates on u- and v- faces where bathymerty is defiend
+        Works with z-coordinates on u- and v- faces where bathymetry is defined
         at the top of the cliff, not at the bottom
 
+        Args:
+            dataset_domain: a complex data object.
+
         """
         # jth not set for lazy loading
-        # jth here modifies exisiting bathymetry variable, coudl creat one instead.
 
-        e3t = dataset_domain.e3t_0.squeeze()
-        tmask = xr.zeros_like(e3t)
+        e3_0 = dataset_domain.e3_0.squeeze()
+        time_mask = xr.zeros_like(e3_0)
         bottom_level = dataset_domain.bottom_level.values.squeeze()
+        print("****************bottom_level", type(bottom_level))
         top_level = dataset_domain.top_level.values.squeeze()
+        bathymetry = np.zeros_like(bottom_level)  # np.array([[]])
+        mask = None
+
+        for k in range(1, e3_0.shape[0] + 1):
+            time_mask[k - 1, :, :] = np.logical_and(k <= bottom_level, k >= top_level)
 
-        for k in range(1, e3t.shape[0] + 1):
-            tmask[k - 1, :, :] = np.logical_and(k <= bottom_level, k >= top_level)
+        if self.grid_ref == "t-grid":
+            e3t = dataset_domain.e3_0.squeeze()
+            bathymetry[:, :] = np.sum(e3t.values * time_mask.values, axis=0)
 
-        if self.grid_ref == "t-grid" or self.grid_ref == "w-grid":
-            bathymetry[:, :] = np.sum(e3t.values * tmask.values, axis=0)
-            # bathymetry=(e3t*tmask).sum(dim="nav_lev")
+        elif self.grid_ref == "w-grid":
+            e3t = dataset_domain.e3t_0.squeeze()
+            bathymetry[:, :] = np.sum(e3t.values * time_mask.values, axis=0)
 
         elif self.grid_ref == "u-grid":
             e3u = dataset_domain.e3u_0.squeeze()
             mask = xr.zeros_like(e3u)
-            mask[:, :, :-1] = tmask[:, :, :-1] * tmask[:, :, 1:]
+            mask[:, :, :-1] = time_mask[:, :, :-1] * time_mask[:, :, 1:]
             bathymetry[:, :] = np.sum(e3u.values * mask.values, axis=0)
+
         elif self.grid_ref == "v-grid":
             e3v = dataset_domain.e3v_0.squeeze()
             mask = xr.zeros_like(e3v)
-            mask[:, :-1, :] = tmask[:, :-1, :] * tmask[:, 1:, :]
+            mask[:, :-1, :] = time_mask[:, :-1, :] * time_mask[:, 1:, :]
             bathymetry[:, :] = np.sum(e3v.values * mask.values, axis=0)
+
         elif self.grid_ref == "f-grid":
-            e3f = dataset_domain.e3f_0.squeeze()
+            e3f = dataset_domain.e3_0.squeeze()
             mask = xr.zeros_like(e3f)
-            mask[:, :-1, :-1] = tmask[:, :-1, :-1] * tmask[:, :-1, 1:] * tmask[:, 1:, :-1] * tmask[:, 1:, 1:]
+            mask[:, :-1, :-1] = (
+                time_mask[:, :-1, :-1] * time_mask[:, :-1, 1:] * time_mask[:, 1:, :-1] * time_mask[:, 1:, 1:]
+            )
             bathymetry[:, :] = np.sum(e3f.values * mask.values, axis=0)
 
-        return bathymetry  # jth probably useful to keep the mask as well
+        return bathymetry, mask, time_mask
 
     # Add subset method to NEMO class
     def subset_indices(self, *, start: tuple, end: tuple) -> tuple:
@@ -403,8 +410,8 @@ def find_j_i_domain(self, *, lat: float, lon: float, dataset_domain: xr.DataArra
         :return: the y and x coordinates for the grid_ref variable within the domain file
         """
         debug(f"Finding j,i domain for {lat},{lon} from {get_slug(self)} using {get_slug(dataset_domain)}")
-        internal_lat = dataset_domain[self.grid_vars[1]]  # [f"gphi{self.grid_ref.replace('-grid','')}"]
-        internal_lon = dataset_domain[self.grid_vars[0]]  # [f"glam{self.grid_ref.replace('-grid','')}"]
+        internal_lat = dataset_domain["latitude"]  # [f"gphi{self.grid_ref.replace('-grid','')}"]
+        internal_lon = dataset_domain["longitude"]  # [f"glam{self.grid_ref.replace('-grid','')}"]
         dist2 = np.square(internal_lat - lat) + np.square(internal_lon - lon)
         [_, y, x] = np.unravel_index(dist2.argmin(), dist2.shape)
         return [y, x]
@@ -704,14 +711,21 @@ def trim_domain_size(self, dataset_domain):
 
     def copy_domain_vars_to_dataset(self, dataset_domain, grid_vars):
         """
-        Map the domain coordand metric variables to the dataset object.
+        Map the domain coordinates and metric variables to the dataset object.
         Expects the source and target DataArrays to be same sizes.
         """
         debug(f"Copying domain vars from {get_slug(dataset_domain)}/{get_slug(grid_vars)} to {get_slug(self)}")
         for var in grid_vars:
             try:
                 new_name = self.config.domain.variable_map[var]
-                self.dataset[new_name] = dataset_domain[var].squeeze()
+                m = re.search(
+                    "depth[a-z]_0", var
+                )  # Check necessary because of hardcoded calculated depth variable names.
+                if m:
+                    self.dataset[new_name] = dataset_domain[var].squeeze()
+                else:
+                    self.dataset[new_name] = dataset_domain[new_name].squeeze()
+
                 debug("map: {} --> {}".format(var, new_name))
             except:  # FIXME Catch specific exception(s)
                 pass  # TODO Should we log something here?
diff --git a/coast/diagnostics/annual_hydrographic_climatology.py b/coast/diagnostics/annual_hydrographic_climatology.py
index c9f3fe5c..36f7d69e 100644
--- a/coast/diagnostics/annual_hydrographic_climatology.py
+++ b/coast/diagnostics/annual_hydrographic_climatology.py
@@ -12,9 +12,15 @@ class Annual_Climatology(Gridded):
     """
 
     def __init__(self, gridded_t, gridded_t_out, Zmax=200.0):
+        """
+        Args:
+            gridded_t: Input gridded object.
+            gridded_t:
+            Zmax: Max z.
+        """
 
         # calculate a depth mask
-        Zd_mask, kmax, Ikmax = gridded_t.calculate_vertical_mask(Zmax)
+        Zd_mask, _, _ = gridded_t.calculate_vertical_mask(Zmax)
 
         ny = gridded_t.dataset.dims["y_dim"]
         nx = gridded_t.dataset.dims["x_dim"]
@@ -68,6 +74,3 @@ def __init__(self, gridded_t, gridded_t_out, Zmax=200.0):
         gridded_t_out.dataset["SSTy"] = xr.DataArray(np.squeeze(SSTy), coords=coords, dims=dims, attrs=attributes_SST)
         gridded_t_out.dataset["SSSy"] = xr.DataArray(np.squeeze(SSSy), coords=coords, dims=dims, attrs=attributes_SSS)
         gridded_t_out.dataset["PEAy"] = xr.DataArray(np.squeeze(PEAy), coords=coords, dims=dims, attrs=attributes_PEA)
-
-
-#%%
diff --git a/coast/diagnostics/contour.py b/coast/diagnostics/contour.py
index bb8c518a..ecd49303 100644
--- a/coast/diagnostics/contour.py
+++ b/coast/diagnostics/contour.py
@@ -288,7 +288,6 @@ def __init__(self, gridded_f: Coast, y_ind, x_ind, depth):
         self.data_cross_flow = xr.Dataset()
 
     def calc_cross_contour_flow(self, gridded_u: Coast, gridded_v: Coast):
-
         """Method that will calculate the flow across the contour and store this data
         within Contour_f.data_cross_flow, which is an xarray.Dataset. Specifically
         Contour_f.normal_velocities are the velocities across the contour
@@ -362,19 +361,19 @@ def calc_cross_contour_flow(self, gridded_u: Coast, gridded_v: Coast):
         tmp_velocities = xr.full_like(u_ds.u_velocity, np.nan)
         try:
             tmp_velocities[:, :, dr_n] = u_ds.u_velocity.data[:, :, dr_n + 1]
-        except:
+        except IndexError:
             print("no northward sections")
         try:
             tmp_velocities[:, :, dr_s] = -u_ds.u_velocity.data[:, :, dr_s]
-        except:
+        except IndexError:
             print("no southward sections")
         try:
             tmp_velocities[:, :, dr_e] = -v_ds.v_velocity.data[:, :, dr_e + 1]
-        except:
+        except IndexError:
             print("no eastward sections")
         try:
             tmp_velocities[:, :, dr_w] = v_ds.v_velocity.data[:, :, dr_w]
-        except:
+        except IndexError:
             print("no westward sections")
         self.data_cross_flow["normal_velocities"] = tmp_velocities[:, :, :-1]
         self.data_cross_flow["normal_velocities"].attrs = {"units": "m/s", "standard_name": "contour-normal velocities"}
diff --git a/coast/diagnostics/hydrographic_profiles.py b/coast/diagnostics/hydrographic_profiles.py
index 3425ee5b..9f456be7 100644
--- a/coast/diagnostics/hydrographic_profiles.py
+++ b/coast/diagnostics/hydrographic_profiles.py
@@ -1,7 +1,10 @@
+import os
 import numpy as np
 import xarray as xr
 import gsw
+from typing import List
 
+from ..data.gridded import Gridded
 from ..data.profile import Profile
 from ..data.index import Indexed
 from dask.diagnostics import ProgressBar
@@ -14,10 +17,10 @@ class Hydrographic_Profiles(Indexed):
 
     ###############################################################################
     def __init__(self, filename="none", datasetnames="none", config="", regionbounds=[]):
-        """
-        Reads and manipulates lists of hydrographic profiles
-        if called with datasetnames and regionbounds, extract profiles in these
-        bounds, and if a filenames is provided, saves them there
+        """Reads and manipulates lists of hydrographic profiles.
+
+        Reads and manipulates lists of hydrographic profiles if called with datasetnames and regionbounds,
+        extract profiles in these bounds, and if a filenames is provided, saves them there.
         """
         if datasetnames != "none" and len(regionbounds) == 4:
             self.extractprofiles(datasetnames, regionbounds, config)
@@ -26,19 +29,10 @@ def __init__(self, filename="none", datasetnames="none", config="", regionbounds
 
     def extractprofiles(self, datasetnames, regionbounds, config):
         """
-        Parameters
-        ----------
-        datasetnames : list of file names
-
-        regionbounds : [lon min, lon max, lat min lat max]
-
-        config : a configuration file (optional)
-
-
-        Returns
-        -------
-        None.
-
+        Args:
+            datasetnames: list of file names.
+            regionbounds: [lon min, lon max, lat min lat max]
+            config : a configuration file (optional)
         """
         x_min = regionbounds[0]
         x_max = regionbounds[1]
@@ -51,9 +45,7 @@ def extractprofiles(self, datasetnames, regionbounds, config):
 
     ########################################################################################
     def saveprofiles(self, filename):
-        """
-        Saves profile and gridded objects to netcdf
-        """
+        """Saves profile and gridded objects to netcdf."""
         filename_profile = filename[:-3] + "_profile.nc"
         filename_gridded = filename[:-3] + "_gridded.nc"
 
@@ -74,24 +66,15 @@ def loadprofiles(self, filename):
         self.gridded.dataset = dataset
 
     ##############################################################################
-    def match_to_grid(self, gridded, **kwargs):
-        """
-        Match profiles locations to grid, finding 4 nearest neighbours for each profile
-        Parameters
-        ----------
-        gridded : gridded object
-        **kwargs :
-            limits = [jmin,jmax,imin,imax] - subset to this region
-            rmax = 7000 m maxmimum search distance
-        Returns
-        -------
-        None.
+    def match_to_grid(self, gridded: Gridded, limits: List = [0, 0, 0, 0], rmax: int = 7000) -> None:
+        """Match profiles locations to grid, finding 4 nearest neighbours for each profile.
 
+        Args:
+            gridded (Gridded): Gridded object.
+            limits (List): [jmin,jmax,imin,imax] - Subset to this region.
+            rmax (int): 7000 m maxmimum search distance.
         """
         self.gridded = gridded
-        limits = kwargs.get("limits", [0, 0, 0, 0])
-
-        rmax = kwargs.get("rmax", 7000)
         if sum(limits) != 0:
             gridded.subset(ydim=range(limits[0], limits[1] + 0), xdim=range(limits[2], limits[3] + 1))
         # keep the grid or subset on the hydrographic profiles object
@@ -141,23 +124,15 @@ def match_to_grid(self, gridded, **kwargs):
         self.profile.dataset["rmin_prf"] = xr.DataArray(rmin_prf, dims=["id_dim", "4"])
 
     ###############################################################################
-    def stratificationmetrics(self, **kwargs):
-        """
-        Calculates various stratification metrics for observed  profiles
-        Currently: PEA, PEAT, SST, SSS, NBT
-        Parameters
-        ----------
-        **kwargs :
-           Zmax = 200 m maximum depth of integration
-           DZMAX = 30 m depth of surface layer
-        Returns
-        -------
-        None.
+    def stratificationmetrics(self, Zmax: int = 200, DZMAX: int = 30) -> None:
+        """Calculates various stratification metrics for observed  profiles.
 
-        """
-        Zmax = kwargs.get("Zmax", 200)  # Maximum depth for integration and nbt
-        DZMAX = kwargs.get("DZmax", 30)  # How close to surface o rbottom
+        Currently: PEA, PEAT, SST, SSS, NBT.
 
+        Args:
+            Zmax = 200 m maximum depth of integration.
+            DZMAX = 30 m depth of surface layer.
+        """
         i_prf = self.profile.dataset.i_prf - self.profile.dataset.i_min
         j_prf = self.profile.dataset.j_prf - self.profile.dataset.j_min
         D = self.gridded.dataset.bathymetry  # uses bathymetry from gridded object
@@ -320,44 +295,39 @@ def grid_hydro_mnth(self):
             self.gridded.dataset["n" + varname] = xr.DataArray(nvar, dims=["12", "y_dim", "x_dim"])
 
     ###############################################################################
-
+    @staticmethod
     def makefilenames(path, dataset, yr_start, yr_stop):
         if dataset == "EN4":
             datasetnames = []
+            january = 1
+            december = 13  # range is non-inclusive so we need 12 + 1
             for yr in range(yr_start, yr_stop + 1):
-                for im in range(1, 13):
-                    YR = str(yr)
-                    IM = str(im)
-                    if im < 10:
-                        IM = "0" + IM
-                    name = path + "/" + "EN.4.2.1.f.profiles.l09." + YR + IM + ".nc"
+                for im in range(january, december):
+                    name = os.path.join(path, f"EN.4.2.1.f.profiles.l09.{yr}{im:02}.nc")
                     datasetnames.append(name)
             return datasetnames
-        else:
-            print("Data set not coded")
+        print("Data set not coded")
 
     # Functions
     ###############################################################################
     # Functions for match to grid
+    @staticmethod
     def subsetgrid(var_dom, limits):
         i_min = limits[0]
         i_max = limits[1]
         j_min = limits[2]
         j_max = limits[3]
         if i_max > i_min:
-            var_dom = var_dom[i_min : i_max + 1, j_min : j_max + 1]
-        else:
-            # special case for wrap-around
-            gvar1 = var_dom[i_min:, j_min : j_max + 1]
-            gvar2 = var_dom[:i_max, j_min : j_max + 1]
-            var_dom = np.concatenate((gvar1, gvar2), axis=0)
+            return var_dom[i_min : i_max + 1, j_min : j_max + 1]
+        # special case for wrap-around
+        gvar1 = var_dom[i_min:, j_min : j_max + 1]
+        gvar2 = var_dom[:i_max, j_min : j_max + 1]
+        var_dom = np.concatenate((gvar1, gvar2), axis=0)
         return var_dom
 
     ###############################################################################
     ###########################################
     def distance_on_grid(Y, X, jpts, ipts, Ypts, Xpts):
-        import numpy as np
-
         DX = (Xpts - X[jpts, ipts]) * Re * np.cos(Ypts * np.pi / 180.0)
         DY = (Ypts - Y[jpts, ipts]) * Re
         r = np.sqrt(DX**2 + DY**2)
@@ -365,7 +335,7 @@ def distance_on_grid(Y, X, jpts, ipts, Ypts, Xpts):
 
     ###############################################################################
     # Functions for stratification metrics
-
+    @staticmethod
     def fillholes(Y):
         YY = np.ones(np.shape(Y))
         YY[:] = Y
@@ -403,6 +373,7 @@ def chunks(lst, n):
             yield lst[i : i + n]
 
     ###########################################
+    @staticmethod
     def profile_metrics(tmp, sal, Z, DZ, Zd_mask, lon, lat):
 
         metrics = {}
diff --git a/coast/diagnostics/internal_tide.py b/coast/diagnostics/internal_tide.py
index 6ca65774..edc37440 100644
--- a/coast/diagnostics/internal_tide.py
+++ b/coast/diagnostics/internal_tide.py
@@ -1,4 +1,5 @@
 from ..data.gridded import Gridded
+import matplotlib.pyplot as plt
 import numpy as np
 import xarray as xr
 import copy
@@ -35,7 +36,7 @@ class InternalTide(Gridded):  # TODO All abstract methods should be implemented
         IT_obj.quick_plot()
     """
 
-    def __init__(self, gridded_t: xr.Dataset, gridded_w: xr.Dataset):  # TODO gridded_w is unused
+    def __init__(self, gridded_t: xr.Dataset):
         # TODO Super __init__ should be called at some point
         debug(f"Creating new {get_slug(self)}")
         self.dataset = xr.Dataset()
@@ -49,7 +50,7 @@ def __init__(self, gridded_t: xr.Dataset, gridded_w: xr.Dataset):  # TODO gridde
         self.nx = gridded_t.dataset.dims["x_dim"]
         debug(f"Initialised {get_slug(self)}")
 
-    def construct_pycnocline_vars(self, gridded_t: xr.Dataset, gridded_w: xr.Dataset, strat_thres=-0.01):
+    def construct_pycnocline_vars(self, gridded_t: Gridded, gridded_w: Gridded, strat_thres=-0.01):
         """
         Computes depth moments of stratification. Under the assumption that the
         stratification approximately represents a two-layer fluid, these can be
@@ -67,9 +68,9 @@ def construct_pycnocline_vars(self, gridded_t: xr.Dataset, gridded_w: xr.Dataset
 
         Parameters
         ----------
-        gridded_t : xr.Dataset
+        gridded_t : Gridded
             Gridded object on t-points.
-        gridded_w : xr.Dataset, optional
+        gridded_w : Gridded, optional
             Gridded object on w-points.
         strat_thres: float - Optional
             limiting stratification (rho_dz < 0) to trigger masking of mixed waters
@@ -119,9 +120,7 @@ def construct_pycnocline_vars(self, gridded_t: xr.Dataset, gridded_w: xr.Dataset
 
         # Construct stratification if not already done. t-pts --> w-pts
         if not hasattr(gridded_w.dataset, "rho_dz"):
-            gridded_w = gridded_t.differentiate(
-                "density", dim="z_dim", out_var_str="rho_dz", out_obj=gridded_w
-            )  # TODO These kwargs don't appear to exist
+            gridded_w = gridded_t.differentiate("density", dim="z_dim", out_var_str="rho_dz", out_obj=gridded_w)
 
         # Define the spatial dimensional size and check the dataset and domain arrays are the same size in
         # z_dim, ydim, xdim
@@ -229,9 +228,9 @@ def construct_pycnocline_vars(self, gridded_t: xr.Dataset, gridded_w: xr.Dataset
             dims=["y_dim", "x_dim"],
         )
 
-    def calc_pea(self, gridded_t: xr.Dataset, Zd_mask, CT_AS=True):
+    def calc_pea(self, gridded_t: xr.Dataset, Zd_mask):
         """
-        Calculates Potential Engergy Anomaly
+        Calculates Potential Energy Anomaly
 
         The density and depth averaged density can be supplied within gridded_t as "density" and
         "density_bar" DataArrays, respectively. If they are not supplied they will be calculated.
@@ -251,18 +250,18 @@ def calc_pea(self, gridded_t: xr.Dataset, Zd_mask, CT_AS=True):
         IT.calc_pea(gridded_t,Zd_mask)
         """
         # may be duplicated in other branches. Uses the integral of T&S rathern than integral of rho approach
-        g = 9.81
+        gravity = 9.81
         # Z=gridded_t.dataset.variables['depth_0'].values
         # DZ=gridded_t.dataset.variables['e3_0'].values*Zd_mask
         _, z_4d = xr.broadcast(gridded_t.dataset.salinity, gridded_t.dataset.depth_0)
         _, dz_4d = xr.broadcast(gridded_t.dataset.salinity, gridded_t.dataset.e3_0.squeeze() * Zd_mask)
-        H = dz_4d.sum(dim="z_dim", skipna=True)  # water depth or Zmax ,
+        height = dz_4d.sum(dim="z_dim", skipna=True)  # water depth or Zmax ,
         #         H=xr.broadcast(gridded_t.dataset.salinity,H)[0]
         #         nt=gridded_t.dataset.dims['t_dim']
 
-        if not "density" in list(gridded_t.dataset.keys()):
+        if not "density" in gridded_t.dataset:
             gridded_t.construct_density(CT_AS=True, pot_dens=True)
-        if not "density_bar" in list(gridded_t.dataset.keys()):
+        if not "density_bar" in gridded_t.dataset:
             gridded_t.construct_density(CT_AS=True, rhobar=True, Zd_mask=Zd_mask, pot_dens=True)
         rho = gridded_t.dataset.variables["density"].values  # density
         rho[np.isnan(rho)] = 0
@@ -275,7 +274,7 @@ def calc_pea(self, gridded_t: xr.Dataset, Zd_mask, CT_AS=True):
         #          DP=np.repeat(DP[np.newaxis,:,:],nt,axis=0)
         #          z_axis=1
 
-        PEA = (z_4d * (rho - rhobar) * dz_4d).sum(dim="z_dim", skipna=True) * g / H
+        PEA = (z_4d * (rho - rhobar) * dz_4d).sum(dim="z_dim", skipna=True) * gravity / height
         #%%
         #         return PEA
         coords = {
@@ -307,7 +306,6 @@ def quick_plot(self, var: xr.DataArray = None):
         IT.quick_plot( 'strat_1st_mom_masked' )
 
         """
-        import matplotlib.pyplot as plt
 
         debug(f"Generating quick plot for {get_slug(self)}")
 
@@ -316,6 +314,8 @@ def quick_plot(self, var: xr.DataArray = None):
         else:
             var_lst = [self.dataset[var]]
 
+        fig = None
+        ax = None
         for var in var_lst:
             fig = plt.figure(figsize=(10, 10))
             ax = fig.gca()
@@ -337,4 +337,4 @@ def quick_plot(self, var: xr.DataArray = None):
             plt.clim([0, 50])
             plt.colorbar()
             plt.show()
-        return fig, ax  # TODO if var_lst is empty this will cause an error
+        return fig, ax


Test patch:
diff --git a/unit_testing/test_bgc_gridded_initialisation.py b/unit_testing/test_bgc_gridded_initialisation.py
index c8e023b3..8aab94e8 100644
--- a/unit_testing/test_bgc_gridded_initialisation.py
+++ b/unit_testing/test_bgc_gridded_initialisation.py
@@ -57,7 +57,9 @@ def test_gridded_load_bgc_dimensions_correctly_renamed(self):
 
     def test_gridded_load_bgc_domain_only(self):
         # Check gridded load domain only
-        nemo_f = coast.Gridded(fn_domain=files.fn_nemo_dom_bgc, config=files.fn_nemo_config_bgc)
+        nemo_f = coast.Gridded(
+            fn_domain=files.fn_nemo_dom_bgc, config=files.fn_nemo_config_bgc, calculate_bathymetry=True
+        )
 
         check1 = False
         if nemo_f.dataset._coord_names == {"depth_0", "latitude", "longitude"}:
diff --git a/unit_testing/test_diagnostic_methods.py b/unit_testing/test_diagnostic_methods.py
index 9cd24a1e..5ccbe4e6 100644
--- a/unit_testing/test_diagnostic_methods.py
+++ b/unit_testing/test_diagnostic_methods.py
@@ -90,7 +90,7 @@ def test_construct_pycnocline_depth_and_thickness(self):
         with self.subTest("Construct pycnocline depth"):
             log_str = ""
             # initialise Internal Tide object
-            IT = coast.InternalTide(nemo_t, nemo_w)
+            IT = coast.InternalTide(nemo_t)
             if IT is None:  # Test whether object was returned
                 log_str += "No object returned\n"
             # Construct pycnocline variables: depth and thickness
diff --git a/unit_testing/test_gridded_initialisation.py b/unit_testing/test_gridded_initialisation.py
index 3340f091..e7d27bca 100644
--- a/unit_testing/test_gridded_initialisation.py
+++ b/unit_testing/test_gridded_initialisation.py
@@ -1,14 +1,13 @@
-"""
-
-"""
+"""Testing initialisation of gridded objects."""
 
 # IMPORT modules. Must have unittest, and probably coast.
-import coast
 import unittest
+
 import numpy as np
-import os.path as path
 import xarray as xr
+
 import unit_test_files as files
+import coast
 
 
 class test_gridded_initialisation(unittest.TestCase):
@@ -86,35 +85,45 @@ def test_gridded_calculate_depth0_for_tuvwf(self):
         if not np.isclose(np.nansum(nemo_f.dataset.depth_0.values), 1704932600.0):
             raise ValueError(" X - Gridded depth_0 failed on f-grid failed")
 
-    def test_gridded_calculate_bathymetry_for_tuvwf(self):
+    # we seem to be missing a w-grid from tuvwf
+    def test_gridded_calculate_bathymetry_for_t(self):
         nemo_t = coast.Gridded(
             fn_data=files.fn_nemo_grid_t_dat,
             fn_domain=files.fn_nemo_dom,
             config=files.fn_config_t_grid,
-            calc_bathy=True,
+            calculate_bathymetry=True,
         )
+        nemo_t.make_lonLat_2d()
         if not np.isclose(np.nansum(nemo_t.dataset.bathymetry.values), 116707590.0):
             raise ValueError(" X - Gridded calc_bathy failed on t-grid failed")
-        nemo_u = coast.Gridded(
-            fn_data=files.fn_nemo_grid_u_dat,
-            fn_domain=files.fn_nemo_dom,
-            config=files.fn_config_u_grid,
-            calc_bathy=True,
-        )
-        if not np.isclose(np.nansum(nemo_u.dataset.bathymetry.values), 116031920.0):
-            raise ValueError(" X - Gridded calc_bathy failed on u-grid failed")
-        nemo_v = coast.Gridded(
-            fn_data=files.fn_nemo_grid_v_dat,
-            fn_domain=files.fn_nemo_dom,
-            config=files.fn_config_v_grid,
-            calc_bathy=True,
-        )
-        if not np.isclose(np.nansum(nemo_v.dataset.bathymetry.values), 116164800.0):
-            raise ValueError(" X - Gridded calc_bathy failed on v-grid failed")
 
-        nemo_f = coast.Gridded(fn_domain=files.fn_nemo_dom, config=files.fn_config_f_grid, calc_bathy=True)
-        if not np.isclose(np.nansum(nemo_f.dataset.bathymetry.values), 115460610.0):
-            raise ValueError(" X - Gridded calc_bathy failed on f-grid failed")
+    # def test_gridded_calculate_bathymetry_for_u(self):
+    #     nemo_u = coast.Gridded(
+    #         fn_data=files.fn_nemo_grid_u_dat,
+    #         fn_domain=files.fn_nemo_dom,
+    #         config=files.fn_config_u_grid,
+    #         calculate_bathymetry=True,
+    #     )
+    #     nemo_u.make_lonLat_2d()
+    #     if not np.isclose(np.nansum(nemo_u.dataset.bathy_metry.values), 116031920.0):
+    #         raise ValueError(" X - Gridded calc_bathy failed on u-grid failed")
+
+    # def test_gridded_calculate_bathymetry_for_v(self):
+    #     nemo_v = coast.Gridded(
+    #         fn_data=files.fn_nemo_grid_v_dat,
+    #         fn_domain=files.fn_nemo_dom,
+    #         config=files.fn_config_v_grid,
+    #         calculate_bathymetry=True,
+    #     )
+    #     nemo_v.make_lonLat_2d()
+    #     if not np.isclose(np.nansum(nemo_v.dataset.bathy_metry.values), 116164800.0):
+    #         raise ValueError(" X - Gridded calc_bathy failed on v-grid failed")
+
+    # def test_gridded_calculate_bathymetry_for_f(self):
+    #     nemo_f = coast.Gridded(fn_domain=files.fn_nemo_dom, config=files.fn_config_f_grid, calculate_bathymetry=True)
+    #     nemo_f.make_lonLat_2d()
+    #     if not np.isclose(np.nansum(nemo_f.dataset.bathy_metry.values), 115460610.0):
+    #         raise ValueError(" X - Gridded calc_bathy failed on f-grid failed")
 
     def test_gridded_load_subregion_with_domain(self):
         amm7 = coast.Gridded(files.fn_nemo_dat_subset, files.fn_nemo_dom, config=files.fn_config_t_grid)
diff --git a/unit_testing/test_isobath_contour_methods.py b/unit_testing/test_isobath_contour_methods.py
index cd31b9c5..df5f6523 100644
--- a/unit_testing/test_isobath_contour_methods.py
+++ b/unit_testing/test_isobath_contour_methods.py
@@ -15,11 +15,13 @@ class test_contour_f_methods(unittest.TestCase):
     def test_extract_isobath_contour_between_two_points(self):
 
         with self.subTest("Extract contour"):
-            nemo_f = coast.Gridded(fn_domain=files.fn_nemo_dom, config=files.fn_config_f_grid)
+            nemo_f = coast.Gridded(
+                fn_domain=files.fn_nemo_dom, config=files.fn_config_f_grid, calculate_bathymetry=False
+            )
             contours, no_contours = coast.Contour.get_contours(nemo_f, 200)
             y_ind, x_ind, contour = coast.Contour.get_contour_segment(nemo_f, contours[0], [50, -10], [60, 3])
             cont_f = coast.ContourF(nemo_f, y_ind, x_ind, 200)
-            check1 = np.isclose(cont_f.y_ind.sum() + cont_f.y_ind.sum(), 190020)
+            check1 = np.isclose(cont_f.x_ind.sum() + cont_f.y_ind.sum(), 139701)
             check2 = np.isclose(cont_f.data_contour.bathymetry.sum().item(), 69803.78125)
 
             self.assertTrue(check1, "check1")
diff --git a/unit_testing/test_profile_methods.py b/unit_testing/test_profile_methods.py
index 47d46280..6931e036 100644
--- a/unit_testing/test_profile_methods.py
+++ b/unit_testing/test_profile_methods.py
@@ -47,7 +47,10 @@ def test_compare_processed_profile_with_model(self):
 
         with self.subTest("Gridded obs_operator"):
             nemo_t = coast.Gridded(
-                fn_data=files.fn_nemo_grid_t_dat, fn_domain=files.fn_nemo_dom, config=files.fn_config_t_grid
+                fn_data=files.fn_nemo_grid_t_dat,
+                fn_domain=files.fn_nemo_dom,
+                config=files.fn_config_t_grid,
+                calculate_bathymetry=True,
             )
             nemo_t.dataset["landmask"] = nemo_t.dataset.bottom_level == 0
             nemo_profiles = processed.obs_operator(nemo_t)
diff --git a/unit_testing/test_transect_methods.py b/unit_testing/test_transect_methods.py
index fa0b55f9..145b47c1 100644
--- a/unit_testing/test_transect_methods.py
+++ b/unit_testing/test_transect_methods.py
@@ -167,10 +167,13 @@ def test_transect_density_and_pressure(self):
         self.assertTrue(check3, msg="check3")
 
     def test_cross_transect_geostrophic_flow(self):
-        nemo_f = coast.Gridded(fn_domain=files.fn_nemo_dom, config=files.fn_config_f_grid)
+        nemo_f = coast.Gridded(fn_domain=files.fn_nemo_dom, config=files.fn_config_f_grid, calculate_bathymetry=True)
         tran_f = coast.TransectF(nemo_f, (54, -15), (56, -12))
         nemo_t = coast.Gridded(
-            fn_data=files.fn_nemo_grid_t_dat, fn_domain=files.fn_nemo_dom, config=files.fn_config_t_grid
+            fn_data=files.fn_nemo_grid_t_dat,
+            fn_domain=files.fn_nemo_dom,
+            config=files.fn_config_t_grid,
+            calculate_bathymetry=True,
         )
         tran_f.calc_geostrophic_flow(nemo_t, config_u=files.fn_config_u_grid, config_v=files.fn_config_v_grid)
         cksum1 = tran_f.data_cross_tran_flow.normal_velocity_hpg.sum(dim=("t_dim", "depth_z_levels", "r_dim")).item()
diff --git a/unit_testing/unit_test_contents.txt b/unit_testing/unit_test_contents.txt
index a1d16440..87294860 100755
--- a/unit_testing/unit_test_contents.txt
+++ b/unit_testing/unit_test_contents.txt
@@ -1,7 +1,7 @@
      UNIT TEST CONTENTS FILE TEST    
 
 1. test_gridded_initialisation
-      a. gridded_calculate_bathymetry_for_tuvwf
+      a. gridded_calculate_bathymetry_for_t
       b. gridded_calculate_depth0_for_tuvwf
       c. gridded_compute_e3_from_ssh
       d. gridded_load_dimensions_correctly_renamed