roms-tools 2.7.0__py3-none-any.whl → 3.1.0__py3-none-any.whl

roms_tools/setup/grid.py

@@ -1,28 +1,30 @@
-import time
+import importlib.metadata
 import logging
-from dataclasses import dataclass, field, asdict
+import re
+import time
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
 
 import numpy as np
 import xarray as xr
-import matplotlib.pyplot as plt
 import yaml
-import importlib.metadata
-from typing import Dict, Union, List
-from roms_tools.constants import R_EARTH
-from roms_tools.utils import save_datasets
-from roms_tools.setup.topography import _add_topography
+from matplotlib.axes import Axes
+
+from roms_tools.constants import MAXIMUM_GRID_SIZE, R_EARTH
+from roms_tools.plot import plot
 from roms_tools.setup.mask import _add_mask, _add_velocity_masks
-from roms_tools.vertical_coordinate import compute_depth_coordinates, sigma_stretch
-from roms_tools.plot import _plot, _section_plot
+from roms_tools.setup.topography import _add_topography
 from roms_tools.setup.utils import (
+    extract_single_value,
+    gc_dist,
+    get_target_coords,
     interpolate_from_rho_to_u,
     interpolate_from_rho_to_v,
-    get_target_coords,
-    gc_dist,
-    _pop_grid_data,
+    pop_grid_data,
+    write_to_yaml,
 )
-from roms_tools.setup.utils import extract_single_value
-from pathlib import Path
+from roms_tools.utils import save_datasets
+from roms_tools.vertical_coordinate import compute_depth_coordinates, sigma_stretch
 
 
 @dataclass(kw_only=True)
@@ -69,7 +71,7 @@ class Grid:
     theta_s : float, optional
         The surface control parameter. Must satisfy 0 < theta_s <= 10. The default is 5.0.
     theta_b : float, optional
-        The bottom control parameter. Must satisfy 0 < theta_b <= 4. The default is 2.0.
+        The bottom control parameter. Must satisfy 0 < theta_b <= 10. The default is 2.0.
     hc : float, optional
         The critical depth (in meters). The default is 300.0.
     verbose: bool, optional
@@ -78,28 +80,46 @@ class Grid:
     Raises
     ------
     ValueError
-        If you try to create a grid with domain size larger than 20000 km.
+        If you try to create a grid with domain size larger than 25000 km.
     """
 
     nx: int
+    """Number of grid points in the x-direction."""
     ny: int
+    """Number of grid points in the y-direction."""
     size_x: float
+    """Domain size in the x-direction (in kilometers)."""
     size_y: float
+    """Domain size in the y-direction (in kilometers)."""
     center_lon: float
+    """Longitude of grid center."""
     center_lat: float
+    """Latitude of grid center."""
     rot: float = 0
+    """Rotation of grid x-direction from lines of constant latitude, measured in
+    degrees."""
     N: int = 100
+    """The number of vertical levels."""
     theta_s: float = 5.0
+    """The surface control parameter."""
     theta_b: float = 2.0
+    """The bottom control parameter."""
     hc: float = 300.0
-    topography_source: Dict[str, Union[str, Path, List[Union[str, Path]]]] = None
+    """The critical depth (in meters)."""
+    topography_source: dict[str, str | Path | list[str | Path]] = None
+    """Dictionary specifying the source of the topography data."""
     hmin: float = 5.0
+    """The minimum ocean depth (in meters)."""
     verbose: bool = False
+    """Whether to print grid generation steps with timing."""
+
     ds: xr.Dataset = field(init=False, repr=False)
+    """An xarray Dataset containing post-processed variables ready for input into
+    ROMS."""
     straddle: bool = field(init=False, repr=False)
+    """Whether the grid straddles the dateline."""
 
     def __post_init__(self):
-
         self._input_checks()
 
         # Horizontal grid
@@ -146,7 +166,6 @@ class Grid:
                 )
 
     def _create_mask(self, verbose=False) -> None:
-
         if verbose:
             start_time = time.time()
             logging.info("=== Creating the mask ===")
@@ -193,7 +212,6 @@ class Grid:
             This method updates the internal dataset (`self.ds`) in place by adding or overwriting the
             topography variable. It does not return any value.
         """
-
         topography_source = topography_source or self.topography_source
         hmin = hmin or self.hmin
 
@@ -258,7 +276,6 @@ class Grid:
         None
             This method modifies the dataset in place by adding vertical coordinate variables.
         """
-
         N = N or self.N
         theta_s = theta_s or self.theta_s
         theta_b = theta_b or self.theta_b
@@ -293,9 +310,9 @@ class Grid:
         cs_w, sigma_w = sigma_stretch(theta_s, theta_b, N, "w")
 
         ds["sigma_r"] = sigma_r.astype(np.float32)
-        ds["sigma_r"].attrs[
-            "long_name"
-        ] = "Fractional vertical stretching coordinate at rho-points"
+        ds["sigma_r"].attrs["long_name"] = (
+            "Fractional vertical stretching coordinate at rho-points"
+        )
         ds["sigma_r"].attrs["units"] = "nondimensional"
 
         ds["Cs_r"] = cs_r.astype(np.float32)
@@ -303,9 +320,9 @@ class Grid:
         ds["Cs_r"].attrs["units"] = "nondimensional"
 
         ds["sigma_w"] = sigma_w.astype(np.float32)
-        ds["sigma_w"].attrs[
-            "long_name"
-        ] = "Fractional vertical stretching coordinate at w-points"
+        ds["sigma_w"].attrs["long_name"] = (
+            "Fractional vertical stretching coordinate at w-points"
+        )
         ds["sigma_w"].attrs["units"] = "nondimensional"
 
         ds["Cs_w"] = cs_w.astype(np.float32)
@@ -338,7 +355,6 @@ class Grid:
         The check is based on whether the longitudinal differences between adjacent
         points exceed 300 degrees, indicating a potential wraparound of longitude.
         """
-
         if (
             np.abs(self.ds.lon_rho.diff("xi_rho")).max() > 300
             or np.abs(self.ds.lon_rho.diff("eta_rho")).max() > 300
@@ -347,6 +363,8 @@ class Grid:
         else:
             self.straddle = False
 
+        self.ds.attrs["straddle"] = str(self.straddle)
+
     def _coarsen(self):
         """Update the grid by adding grid variables that are coarsened versions of the
         original fine-resoluion grid variables. The coarsening is by a factor of two.
@@ -390,66 +408,49 @@ class Grid:
                 "long_name", ds[fine_var].attrs.get("Long_name", "")
             )
             ds[coarse_var].attrs["long_name"] = f"{long_name} on coarsened grid"
-            ds[coarse_var].attrs["units"] = ds[fine_var].attrs["units"]
+            if "units" in ds[fine_var].attrs:
+                ds[coarse_var].attrs["units"] = ds[fine_var].attrs["units"]
 
         self.ds = ds
 
     def plot(
-        self, bathymetry: bool = True, title: str = None, with_dim_names: bool = False
+        self,
+        with_dim_names: bool = False,
+        save_path: str | None = None,
     ) -> None:
         """Plot the grid.
 
         Parameters
         ----------
-        bathymetry : bool, optional
-            Whether or not to plot the bathymetry. Default is True.
-        title : str, optional
-            The title of the plot. If not provided, it will be set to a default.
         with_dim_names : bool, optional
             Whether or not to plot the dimension names. Default is False.
 
+        save_path : str, optional
+            Path to save the generated plot. If None, the plot is shown interactively.
+            Default is None.
 
         Returns
         -------
         None
             This method does not return any value. It generates and displays a plot.
         """
-
-        field = self.ds.h.where(self.ds.mask_rho)
-        lat_deg = self.ds.lat_rho
-        lon_deg = self.ds.lon_rho
-        if self.straddle:
-            lon_deg = xr.where(lon_deg > 180, lon_deg - 360, lon_deg)
-        field = field.assign_coords({"lon": lon_deg, "lat": lat_deg})
-
-        if bathymetry:
-            if title is None:
-                title = "ROMS grid and bathymetry"
-
-            vmax = field.max().values
-            vmin = field.min().values
-            cmap = plt.colormaps.get_cmap("YlGnBu")
-            cmap.set_bad(color="gray")
-            kwargs = {"vmax": vmax, "vmin": vmin, "cmap": cmap}
-
-            _plot(
-                field=field,
-                title=title,
-                with_dim_names=with_dim_names,
-                kwargs=kwargs,
-            )
-        else:
-            if title is None:
-                title = "ROMS grid"
-            _plot(
-                field=field, title=title, with_dim_names=with_dim_names, plot_data=False
-            )
+        field = self.ds["h"]
+
+        plot(
+            field=field,
+            grid_ds=self.ds,
+            with_dim_names=with_dim_names,
+            save_path=save_path,
+            cmap_name="YlGnBu",
+        )
 
     def plot_vertical_coordinate(
         self,
-        s=None,
-        eta=None,
-        xi=None,
+        s: int | None = None,
+        eta: int | None = None,
+        xi: int | None = None,
+        max_nr_layer_contours: int | None = None,
+        ax: Axes | None = None,
     ) -> None:
         """Plot the layer depth for a given eta-, xi-, or s-slice.
 
@@ -461,8 +462,10 @@ class Grid:
             The eta-index to plot. Default is None.
         xi : int, optional
             The xi-index to plot. Default is None.
+        max_nr_layer_contours : int, optional
+            Number of layer contours displayed. Only relevant when plotting a vertical section. Default is all.
         ax : matplotlib.axes.Axes, optional
-            The axes to plot on. If None, a new figure is created. Note that this argument does not work for 2D horizontal plots.
+            The axes to plot on. If None, a new figure is created. Note that this argument is ignored for 2D horizontal plots.
 
         Returns
         -------
@@ -474,61 +477,34 @@ class Grid:
         ValueError
             If not exactly one of s, eta, xi is specified.
         """
-
-        if sum(s is not None for s in [s, eta, xi]) != 1:
+        if sum(i is not None for i in [s, eta, xi]) != 1:
             raise ValueError("Exactly one of s, eta, or xi must be specified.")
 
-        depth = compute_depth_coordinates(
-            self.ds, zeta=0, depth_type="layer", location="rho", eta=eta, xi=xi
-        )
-
-        title = "Layer depth at rho-points"
-        if eta is not None:
-            title = title + f", eta_rho = {self.ds.eta_rho[eta].item()}"
-        if xi is not None:
-            title = title + f", xi_rho = {self.ds.xi_rho[xi].item()}"
         if s is not None:
-            title = title + f", s_rho = {depth.s_rho[s].item()}"
-            depth = depth.isel(s_rho=s)
-
-            vmax = depth.max().values
-            vmin = depth.min().values
-            cmap = plt.colormaps.get_cmap("YlGnBu")
-            cmap.set_bad(color="gray")
-            kwargs = {"vmax": vmax, "vmin": vmin, "cmap": cmap}
-
-            lat_deg = self.ds.lat_rho
-            lon_deg = self.ds.lon_rho
-            if self.straddle:
-                lon_deg = xr.where(lon_deg > 180, lon_deg - 360, lon_deg)
-            depth = depth.assign_coords({"lon": lon_deg, "lat": lat_deg})
-
-            _plot(
-                field=depth.where(self.ds.mask_rho),
-                depth_contours=False,
-                title=title,
-                kwargs=kwargs,
+            field = compute_depth_coordinates(
+                self.ds, zeta=0, depth_type="layer", location="rho"
             )
+            layer_contours = False
+            add_colorbar = True
         else:
-            field = xr.zeros_like(depth)
-            field = field.assign_coords({"layer_depth": depth})
-
-            interface_depth = compute_depth_coordinates(
-                self.ds, zeta=0, depth_type="interface", eta=eta, xi=xi
-            )
-
-            cmap = plt.colormaps.get_cmap("YlGnBu")
-            cmap.set_bad(color="gray")
-            kwargs = {"vmax": 0.0, "vmin": 0.0, "cmap": cmap, "add_colorbar": False}
-
-            _section_plot(
-                field=field,
-                interface_depth=interface_depth,
-                title=title,
-                kwargs=kwargs,
-            )
+            field = xr.zeros_like(self.ds.s_rho * self.ds["h"])
+            field.attrs["long_name"] = "layer depth at rho-points"
+            layer_contours = True
+            add_colorbar = False
+
+        plot(
+            field=field,
+            grid_ds=self.ds,
+            s=s,
+            eta=eta,
+            xi=xi,
+            layer_contours=layer_contours,
+            max_nr_layer_contours=max_nr_layer_contours,
+            cmap_name="YlGnBu",
+            add_colorbar=add_colorbar,
+        )
 
-    def save(self, filepath: Union[str, Path]) -> None:
+    def save(self, filepath: str | Path) -> None:
         """Save the grid information to a netCDF4 file.
 
         Parameters
@@ -541,7 +517,6 @@ class Grid:
         List[Path]
             A list of Path objects for the filenames that were saved.
         """
-
         # Ensure filepath is a Path object
         filepath = Path(filepath)
 
@@ -557,13 +532,29 @@ class Grid:
         return saved_filenames
 
     @classmethod
-    def from_file(cls, filepath: Union[str, Path], verbose: bool = False) -> "Grid":
+    def from_file(
+        cls,
+        filepath: str | Path,
+        theta_s: float | None = None,
+        theta_b: float | None = None,
+        hc: float | None = None,
+        N: int | None = None,
+        verbose: bool = False,
+    ) -> "Grid":
         """Create a Grid instance from an existing file.
 
         Parameters
         ----------
         filepath : Union[str, Path]
             Path to the file containing the grid information.
+        theta_s : float, optional
+            Surface stretching parameter for vertical coordinate.
+        theta_b : float, optional
+            Bottom stretching parameter for vertical coordinate.
+        hc : float, optional
+            Critical depth for vertical coordinate.
+        N : int, optional
+            Number of vertical levels for vertical coordinate.
         verbose: bool, optional
             Indicates whether to print grid generation steps with timing. Defaults to False.
 
@@ -572,6 +563,15 @@ class Grid:
         Grid
             A new instance of Grid populated with data from the file.
         """
+        # Validate that either all or none of the vertical grid parameters are provided
+        user_params = [theta_s, theta_b, hc, N]
+        num_provided = sum(p is not None for p in user_params)
+        if 0 < num_provided < 4:
+            raise ValueError(
+                "If specifying vertical coordinate parameters, you must provide all of: "
+                "theta_s, theta_b, hc, and N."
+            )
+
         # Load the dataset from the file
         ds = xr.open_dataset(filepath)
 
@@ -609,53 +609,118 @@ class Grid:
                 ds = grid.ds.set_coords(var)
                 grid.ds = ds
 
-        # Update vertical coordinate if necessary
-        if not all(var in grid.ds for var in ["Cs_r", "Cs_w"]):
-            logging.warning("Vertical coordinates (Cs_r, Cs_w) not found in grid file.")
-            N = 100
-            theta_s = 5.0
-            theta_b = 2.0
-            hc = 300.0
+        # Check if vertical coordinates already exist
+        if all(var in grid.ds for var in ["Cs_r", "Cs_w"]):
+            prior_Cs_r = grid.ds.Cs_r.copy(deep=True)
+            prior_Cs_w = grid.ds.Cs_w.copy(deep=True)
+        else:
+            prior_Cs_r = None
+            prior_Cs_w = None
 
+        # Case: user provides all vertical coordinate parameters
+        if num_provided == 4:
             grid.update_vertical_coordinate(
                 N=N, theta_s=theta_s, theta_b=theta_b, hc=hc, verbose=True
             )
+
+            if prior_Cs_r is not None and prior_Cs_w is not None:
+                # Check for consistency between provided and file-stored vertical coordinates
+                if (grid.ds.Cs_r.shape != prior_Cs_r.shape) or (
+                    grid.ds.Cs_w.shape != prior_Cs_w.shape
+                ):
+                    raise ValueError(
+                        "The grid file contains vertical coordinates (Cs_r, Cs_w), but their shapes "
+                        "are inconsistent with the provided N."
+                    )
+                if not (
+                    np.allclose(grid.ds.Cs_r, prior_Cs_r)
+                    and np.allclose(grid.ds.Cs_w, prior_Cs_w)
+                ):
+                    raise ValueError(
+                        "The grid file contains vertical coordinates (Cs_r, Cs_w), "
+                        "but they are inconsistent with the provided theta_s, theta_b, hc, and N."
+                    )
+
+        # Case: user did not provide vertical coordinate parameters
+        elif prior_Cs_r is None or prior_Cs_w is None:
+            logging.warning("Vertical coordinates (Cs_r, Cs_w) not found in grid file.")
+
+            # Use fallback parameters
+            grid.update_vertical_coordinate(
+                N=100, theta_s=5.0, theta_b=2.0, hc=300.0, verbose=True
+            )
+
+        # Final fallback: get vertical coordinate parameters from attributes if available
         else:
-            grid.theta_s = ds.attrs["theta_s"].item()
-            grid.theta_b = ds.attrs["theta_b"].item()
-            grid.hc = ds.attrs["hc"].item()
-            grid.N = len(ds.s_rho)
+            if not all(attr in grid.ds.attrs for attr in ["theta_s", "theta_b", "hc"]):
+                raise ValueError(
+                    "Missing vertical coordinate attributes in grid file: "
+                    "'theta_s', 'theta_b', or 'hc'."
+                )
+
+        # Assign vertical coordinate metadata
+        grid.theta_s = grid.ds.attrs["theta_s"].item()
+        grid.theta_b = grid.ds.attrs["theta_b"].item()
+        grid.hc = grid.ds.attrs["hc"].item()
+        grid.N = len(grid.ds.s_rho)
 
         # Manually set the remaining attributes by extracting parameters from dataset
         grid.nx = ds.sizes["xi_rho"] - 2
         grid.ny = ds.sizes["eta_rho"] - 2
         if "center_lon" in ds.attrs:
-            center_lon = ds.attrs["center_lon"]
+            center_lon = float(ds.attrs["center_lon"])
         elif "tra_lon" in ds:
-            center_lon = extract_single_value(ds["tra_lon"])
+            center_lon = float(extract_single_value(ds["tra_lon"]))
+        elif "title" in ds.attrs:
+            match = re.search(r"Lon:\s*(-?\d+(?:\.\d+)?)", ds.attrs["title"])
+            if match:
+                center_lon = float(match.group(1))
+            else:
+                raise ValueError(
+                    "Could not extract 'center_lon' from title attribute. "
+                    "Expected format: '... Lon: <value> ...'"
+                )
         else:
             raise ValueError(
-                "Missing grid information: 'center_lon' attribute or 'tra_lon' variable "
+                "Missing grid information: 'center_lon' attribute, 'tra_lon' variable, or 'Lon:' in 'title' attribute "
                 "must be present in the dataset."
             )
         grid.center_lon = center_lon
         if "center_lat" in ds.attrs:
-            center_lat = ds.attrs["center_lat"]
+            center_lat = float(ds.attrs["center_lat"])
         elif "tra_lat" in ds:
-            center_lat = extract_single_value(ds["tra_lat"])
+            center_lat = float(extract_single_value(ds["tra_lat"]))
+        elif "title" in ds.attrs:
+            match = re.search(r"Lat:\s*(-?\d+(?:\.\d+)?)", ds.attrs["title"])
+            if match:
+                center_lat = float(match.group(1))
+            else:
+                raise ValueError(
+                    "Could not extract 'center_lat' from title attribute. "
+                    "Expected format: '... Lon: <value> ...'"
+                )
         else:
             raise ValueError(
-                "Missing grid information: 'center_lat' attribute or 'tra_lat' variable "
+                "Missing grid information: 'center_lat' attribute, 'tra_lat' variable, or 'Lat:' in 'title' attribute "
                 "must be present in the dataset."
             )
         grid.center_lat = center_lat
         if "rot" in ds.attrs:
-            rot = ds.attrs["rot"]
+            rot = float(ds.attrs["rot"])
         elif "rotate" in ds:
-            rot = extract_single_value(ds["rotate"])
+            rot = float(extract_single_value(ds["rotate"]))
+        elif "title" in ds.attrs:
+            match = re.search(r"rotate:\s*(-?\d+(?:\.\d+)?)", ds.attrs["title"])
+            if match:
+                rot = float(match.group(1))
+            else:
+                raise ValueError(
+                    "Could not extract 'rot' from title attribute. "
+                    "Expected format: '... rotate: <value> ...'"
+                )
         else:
             raise ValueError(
-                "Missing grid information: 'rot' attribute or 'rotate' variable "
+                "Missing grid information: 'rot' attribute, 'rotate' variable, or 'rotate:' in 'title' attribute "
                 "must be present in the dataset."
             )
         grid.rot = rot
@@ -663,21 +728,31 @@ class Grid:
         for attr in [
             "size_x",
             "size_y",
-            "topography_source",
             "hmin",
         ]:
             if attr in ds.attrs:
-                if attr == "topography_source":
-                    a = {"name": ds.attrs[attr]}
-                else:
-                    a = ds.attrs[attr]
+                a = float(ds.attrs[attr])
             else:
                 a = None
+
             object.__setattr__(grid, attr, a)
 
+        if "topography_source_name" in ds.attrs:
+            if "topography_source_path" in ds.attrs:
+                a = {
+                    "name": ds.attrs["topography_source_name"],
+                    "path": ds.attrs["topography_source_path"],
+                }
+            else:
+                a = {"name": ds.attrs["topography_source_name"]}
+        else:
+            a = None
+
+        object.__setattr__(grid, "topography_source", a)
+
         return grid
 
-    def to_yaml(self, filepath: Union[str, Path]) -> None:
+    def to_yaml(self, filepath: str | Path) -> None:
         """Export the parameters of the class to a YAML file, including the version of
         roms-tools.
 
@@ -686,34 +761,16 @@ class Grid:
         filepath : Union[str, Path]
             The path to the YAML file where the parameters will be saved.
         """
-
-        filepath = Path(filepath)
-
         data = asdict(self)
-        data = _pop_grid_data(data)
+        data = pop_grid_data(data)
+        forcing_dict = {self.__class__.__name__: data}
 
-        # Include the version of roms-tools
-        try:
-            roms_tools_version = importlib.metadata.version("roms-tools")
-        except importlib.metadata.PackageNotFoundError:
-            roms_tools_version = "unknown"
-
-        # Create header
-        header = f"---\nroms_tools_version: {roms_tools_version}\n---\n"
-
-        # Use the class name as the top-level key
-        yaml_data = {self.__class__.__name__: data}
-
-        with filepath.open("w") as file:
-            # Write header
-            file.write(header)
-            # Write YAML data
-            yaml.dump(yaml_data, file, default_flow_style=False, sort_keys=False)
+        write_to_yaml(forcing_dict, filepath)
 
     @classmethod
     def from_yaml(
         cls,
-        filepath: Union[str, Path],
+        filepath: str | Path,
         section_name: str = "Grid",
         verbose: bool = False,
     ) -> "Grid":
@@ -744,7 +801,6 @@ class Grid:
         Issues a warning if the ROMS-Tools version in the YAML header does not match the
         currently installed version.
         """
-
         filepath = Path(filepath)
         # Read the entire file content
         with filepath.open("r") as file:
@@ -787,7 +843,8 @@ class Grid:
 
     def __repr__(self) -> str:
         """Return a string representation of the object with non-None attributes,
-        excluding 'ds'."""
+        excluding 'ds'.
+        """
         cls = self.__class__
         cls_name = cls.__name__
         # Filter attributes to exclude 'ds' and those with None values
@@ -894,13 +951,12 @@ class Grid:
     def _raise_if_domain_size_too_large(self):
         """Raise a ValueError if the domain size exceeds the allowable threshold.
 
-        Checks if either the x or y domain size exceeds 20,000 km and raises an error
-        with appropriate details if the threshold is surpassed.
+        Checks if either the x or y domain size exceeds the threshold and raises an
+        error with appropriate details if the threshold is surpassed.
         """
-        threshold = 20000
-        if self.size_x > threshold or self.size_y > threshold:
+        if self.size_x > MAXIMUM_GRID_SIZE or self.size_y > MAXIMUM_GRID_SIZE:
             raise ValueError(
-                f"Domain size exceeds the allowable limit of {threshold} km. "
+                f"Domain size exceeds the allowable limit of {MAXIMUM_GRID_SIZE} km. "
                 f"Received dimensions: size_x = {self.size_x} km, size_y = {self.size_y} km. "
                 "Please reduce the domain size to meet the threshold."
             )
@@ -918,42 +974,58 @@ class Grid:
             - lonv, latv: 2D arrays of longitudes and latitudes at v-points.
             - lonq, latq: 2D arrays of longitudes and latitudes at cell corners.
         """
-
         # initially define the domain to be longer in x-direction (dimension "length")
         # than in y-direction (dimension "width") to keep grid distortion minimal
         if self.size_y > self.size_x:
-            domain_length, domain_width = self.size_y * 1e3, self.size_x * 1e3  # in m
+            domain_length = self.size_y * 1e3  # in meters
+            domain_width = self.size_x * 1e3
             nl, nw = self.ny, self.nx
         else:
-            domain_length, domain_width = self.size_x * 1e3, self.size_y * 1e3  # in m
+            domain_length = self.size_x * 1e3
+            domain_width = self.size_y * 1e3
             nl, nw = self.nx, self.ny
 
-        domain_length_in_degrees = domain_length / R_EARTH
-        domain_width_in_degrees = domain_width / R_EARTH
+        dlon = domain_length / R_EARTH
+        dlat = domain_width / R_EARTH
 
-        # Generate 1D longitude arrays at cell centers and corners
-        lon_array_1d_in_degrees = domain_length_in_degrees * (
-            np.arange(-0.5, nl + 1.5) / nl - 0.5
-        )
-        lonq_array_1d_in_degrees_q = domain_length_in_degrees * (
-            np.arange(-1, nl + 2) / nl - 0.5
-        )
+        # 1D longitude arrays at centers and corners (cell and cell+1)
+        lon_array_1d_in_degrees = dlon * np.arange(-0.5, nl + 1.5) / nl - dlon / 2
+        lonq_array_1d_in_degrees = dlon * np.arange(-1, nl + 2) / nl - dlon / 2
 
-        # Mercator projection for latitude
-        y1 = np.log(np.tan(np.pi / 4 - domain_width_in_degrees / 4))
-        y2 = np.log(np.tan(np.pi / 4 + domain_width_in_degrees / 4))
+        mul = 1.0
+        # Central longitude spacing (in radians)
+        dlon_cen = dlon / nl
 
-        # Generate 1D latitude arrays with inverse Mercator projection
-        lat_array_1d_in_degrees = np.arctan(
-            np.sinh((y2 - y1) * (np.arange(-0.5, nw + 1.5) / nw) + y1)
-        )
+        # Iteratively adjust latitude to make grid cells square using aspect ratio correction
+        for _ in range(100):
+            # Compute Mercator y-limits
+            y1 = np.log(np.tan(np.pi / 4 - dlat / 4))
+            y2 = np.log(np.tan(np.pi / 4 + dlat / 4))
+
+            # Mercator-transformed latitude coordinates
+            y = (y2 - y1) * np.arange(-0.5, nw + 1.5) / nw + y1
+            lat_array_1d_in_degrees = np.arctan(np.sinh(y))  # Inverse Mercator
+
+            # Central latitude spacing (in radians)
+            i_mid = round(nw / 2)
+            dlat_cen = 0.5 * (
+                lat_array_1d_in_degrees[i_mid + 1] - lat_array_1d_in_degrees[i_mid - 1]
+            )
+
+            # Update multiplicative scaling factor
+            mul = dlat_cen / dlon_cen * domain_length / domain_width * nw / nl
+
+            # Update dlat for next iteration
+            dlat /= mul
+
+        # Latitudes of corners
         latq_array_1d_in_degrees = np.arctan(
-            np.sinh((y2 - y1) * (np.arange(-1, nw + 2) / nw) + y1)
+            np.sinh((y2 - y1) * np.arange(-1, nw + 2) / nw + y1)
         )
 
         # 2D grids for cell centers and corners
         lon, lat = np.meshgrid(lon_array_1d_in_degrees, lat_array_1d_in_degrees)
-        lonq, latq = np.meshgrid(lonq_array_1d_in_degrees_q, latq_array_1d_in_degrees)
+        lonq, latq = np.meshgrid(lonq_array_1d_in_degrees, latq_array_1d_in_degrees)
 
         if self.size_y > self.size_x:
             # Rotate grid by 90 degrees because until here the grid has been defined
@@ -1003,7 +1075,6 @@ class Grid:
             Dataset with variables: lon_rho, lat_rho, lon_u, lat_u, lon_v, lat_v,
             angle, f (Coriolis parameter), pm, pn.
         """
-
         ds = xr.Dataset()
 
         lon_rho = xr.Variable(
@@ -1097,49 +1168,54 @@ class Grid:
 
         return ds
 
-    def _infer_nominal_horizontal_resolution(self, lat=None):
-        """Estimate the nominal horizontal resolution of the grid in degrees at a
-        specified latitude.
-
-        This method calculates the nominal horizontal resolution of the grid by first
-        determining the average grid spacing in meters. The spacing is then converted
-        to degrees, accounting for the Earth's curvature, and the latitude where the
-        resolution is being computed.
+    def _compute_exponential_depth_levels(self, Nz=None, max_depth=None, h=None):
+        """Compute vertical grid center and face depths using an exponential profile.
 
         Parameters
         ----------
-        lat : float, optional
-            Latitude (in degrees) at which to estimate the horizontal resolution.
-            If not provided, the resolution is calculated at the average latitude of
-            the grid (`lat_rho`).
+        Nz : int, optional
+            Number of vertical levels. Defaults to `self.N`.
+
+        max_depth : float, optional
+            Total depth of the domain. Defaults to `self.ds.h.max().values`.
+
+        h : float, optional
+            Scaling parameter for the exponential profile. Defaults to `Nz / 4.5`.
 
         Returns
         -------
-        float
-            The estimated horizontal resolution in degrees, adjusted for the Earth's curvature.
+        tuple of numpy.ndarray
+            Depth values at the vertical grid centers (`z_centers`) and grid faces (`z_faces`),
+            both rounded to two decimal places.
         """
-        if lat is None:
-            # Center latitude in degrees
-            lat = (self.ds.lat_rho.max() + self.ds.lat_rho.min()) / 2
+        if Nz is None:
+            Nz = self.N
+        if max_depth is None:
+            max_depth = self.ds.h.max().values
+        if h is None:
+            h = Nz / 4.5
+
+        k = np.arange(1, Nz + 2)
 
-        # Convert latitude to radians
-        lat_rad = np.deg2rad(lat)
+        z_faces = np.exp(k / h)
 
-        # Mean resolution in meters
-        resolution_in_m = ((1 / self.ds.pm).mean() + (1 / self.ds.pn).mean()) / 2
+        # Normalize
+        z_faces -= z_faces[0]
+        z_faces *= max_depth / z_faces[-1]
+        z_faces[0] = 0.0
 
-        # Meters per degree at the equator
-        meters_per_degree = 2 * np.pi * R_EARTH / 360
+        # Calculate center depths (average between adjacent face depths)
+        z_centers = (z_faces[:-1] + z_faces[1:]) / 2
 
-        # Correct for latitude by multiplying by cos(latitude) for longitude
-        resolution_in_degrees = resolution_in_m / (meters_per_degree * np.cos(lat_rad))
+        # Round both z_faces and z_centers to two decimal places
+        z_faces = np.round(z_faces, 2)
+        z_centers = np.round(z_centers, 2)
 
-        return resolution_in_degrees
+        return z_centers, z_faces
 
 
 def _rotate(coords, rot):
     """Rotate grid counterclockwise relative to surface of Earth by rot degrees."""
-
     (coords["lon"], coords["lat"]) = _rot_sphere(coords["lon"], coords["lat"], rot)
     (coords["lonu"], coords["latu"]) = _rot_sphere(coords["lonu"], coords["latu"], rot)
     (coords["lonv"], coords["latv"]) = _rot_sphere(coords["lonv"], coords["latv"], rot)
@@ -1149,8 +1225,7 @@ def _rotate(coords, rot):
 
 
 def _translate(coords, tra_lat, tra_lon):
-    """Translate grid so that the centre lies at the position (tra_lat, tra_lon)"""
-
+    """Translate grid so that the centre lies at the position (tra_lat, tra_lon)."""
     (lon, lat) = _tra_sphere(coords["lon"], coords["lat"], tra_lat)
     (lonu, latu) = _tra_sphere(coords["lonu"], coords["latu"], tra_lat)
     (lonv, latv) = _tra_sphere(coords["lonv"], coords["latv"], tra_lat)
@@ -1249,7 +1324,6 @@ def _tra_sphere(lon, lat, tra):
     tuple
         Translated longitude and latitude arrays (lon, lat) in radians.
     """
-
     # Convert translation angle from degrees to radians
     tra = tra * np.pi / 180
 
@@ -1304,7 +1378,6 @@ def _compute_coordinate_metrics(coords):
     -----
     Boundary values of `pn` and `pm` are copied from adjacent interior values.
     """
-
     # pm = 1/dx
     pmu = gc_dist(
         coords["lonu"][:, :-1],
@@ -1356,7 +1429,6 @@ def _compute_angle(coords):
         An array of angles (in radians) of the local grid's positive x-axis
         relative to east for each grid point.
     """
-
     # Compute differences in latitudes and longitudes
     dellat = coords["latu"][:, 1:] - coords["latu"][:, :-1]
     dellon = coords["lonu"][:, 1:] - coords["lonu"][:, :-1]
@@ -1395,7 +1467,6 @@ def _f2c(f):
     fc : xarray.DataArray
         Output DataArray with modified dimensions and values.
     """
-
     fc = _f2c_xdir(f)
     fc = fc.transpose()
     fc = _f2c_xdir(fc)