roms-tools 1.7.0__py3-none-any.whl → 2.1.0__py3-none-any.whl

roms_tools/setup/grid.py

@@ -1,4 +1,5 @@
-import copy
+import time
+import logging
 from dataclasses import dataclass, field, asdict
 
 import numpy as np
@@ -6,24 +7,32 @@ import xarray as xr
 import matplotlib.pyplot as plt
 import yaml
 import importlib.metadata
-
-from typing import Union
-from roms_tools.setup.topography import _add_topography_and_mask, _add_velocity_masks
-from roms_tools.setup.plot import _plot, _section_plot, _profile_plot, _line_plot
-from roms_tools.setup.utils import interpolate_from_rho_to_u, interpolate_from_rho_to_v
+from typing import Dict, Union, List
+from roms_tools.setup.topography import _add_topography
+from roms_tools.setup.mask import _add_mask, _add_velocity_masks
+from roms_tools.setup.plot import _plot, _section_plot
+from roms_tools.setup.utils import (
+    interpolate_from_rho_to_u,
+    interpolate_from_rho_to_v,
+    get_target_coords,
+    gc_dist,
+)
 from roms_tools.setup.vertical_coordinate import sigma_stretch, compute_depth
 from roms_tools.setup.utils import extract_single_value, save_datasets
-import logging
 from pathlib import Path
 
-RADIUS_OF_EARTH = 6371315.0  # in m
-
 
 @dataclass(frozen=True, kw_only=True)
 class Grid:
-    """A single ROMS grid.
+    """A single ROMS grid, used for creating, plotting, and then saving a new ROMS
+    domain grid.
+
+    The grid generation consists of four steps:
 
-    Used for creating, plotting, and then saving a new ROMS domain grid.
+    1.  Creating the horizontal grid
+    2.  Creating the mask
+    3.  Generating the topography
+    4.  Preparing the vertical coordinate system
 
     Parameters
     ----------
@@ -43,6 +52,15 @@ class Grid:
         Rotation of grid x-direction from lines of constant latitude, measured in degrees.
         Positive values represent a counterclockwise rotation.
         The default is 0, which means that the x-direction of the grid is aligned with lines of constant latitude.
+    topography_source : Dict[str, Union[str, Path]], optional
+        Dictionary specifying the source of the topography data:
+
+        - "name" (str): The name of the topography data source (e.g., "SRTM15").
+        - "path" (Union[str, Path, List[Union[str, Path]]]): The path to the raw data file. Can be a string or a Path object.
+
+        The default is "ETOPO5", which does not require a path.
+    hmin : float, optional
+       The minimum ocean depth (in meters). The default is 5.0.
     N : int, optional
         The number of vertical levels. The default is 100.
     theta_s : float, optional
@@ -51,11 +69,8 @@ class Grid:
         The bottom control parameter. Must satisfy 0 < theta_b <= 4. The default is 2.0.
     hc : float, optional
         The critical depth (in meters). The default is 300.0.
-    topography_source : str, optional
-        Specifies the data source to use for the topography. Options are
-        "ETOPO5". The default is "ETOPO5".
-    hmin : float, optional
-        The minimum ocean depth (in meters). The default is 5.0.
+    verbose: bool, optional
+        Indicates whether to print grid generation steps with timing. Defaults to False.
 
     Raises
     ------
@@ -74,161 +89,166 @@ class Grid:
     theta_s: float = 5.0
     theta_b: float = 2.0
     hc: float = 300.0
-    topography_source: str = "ETOPO5"
+    topography_source: Dict[str, Union[str, Path, List[Union[str, Path]]]] = None
     hmin: float = 5.0
+    verbose: bool = False
     ds: xr.Dataset = field(init=False, repr=False)
     straddle: bool = field(init=False, repr=False)
 
     def __post_init__(self):
-        ds = _make_grid_ds(
-            nx=self.nx,
-            ny=self.ny,
-            size_x=self.size_x,
-            size_y=self.size_y,
-            center_lon=self.center_lon,
-            center_lat=self.center_lat,
-            rot=self.rot,
-        )
-        # Calling object.__setattr__ is ugly but apparently this really is the best (current) way to combine __post_init__ with a frozen dataclass
-        # see https://stackoverflow.com/questions/53756788/how-to-set-the-value-of-dataclass-field-in-post-init-when-frozen-true
-        object.__setattr__(self, "ds", ds)
 
-        # Update self.ds with topography and mask information
-        self.update_topography_and_mask(
-            topography_source=self.topography_source,
-            hmin=self.hmin,
-        )
+        self._input_checks()
+
+        # Horizontal grid
+        self._create_horizontal_grid()
 
         # Check if the Greenwich meridian goes through the domain.
         self._straddle()
 
-        object.__setattr__(self, "ds", ds)
+        # Mask
+        self._create_mask(verbose=self.verbose)
 
-        # Update the grid by adding grid variables that are coarsened versions of the original grid variables
+        # Coarsen the dataset if needed
         self._coarsen()
 
+        # Topography and mask
+        self.update_topography(
+            topography_source=self.topography_source,
+            hmin=self.hmin,
+            verbose=self.verbose,
+        )
+
+        # Vertical coordinate system
         self.update_vertical_coordinate(
-            N=self.N, theta_s=self.theta_s, theta_b=self.theta_b, hc=self.hc
+            N=self.N,
+            theta_s=self.theta_s,
+            theta_b=self.theta_b,
+            hc=self.hc,
+            verbose=self.verbose,
         )
 
-    def update_topography_and_mask(self, hmin, topography_source="ETOPO5") -> None:
-        """Update the grid dataset by adding or overwriting the topography and land/sea
-        mask.
+    def _input_checks(self):
+        if self.topography_source is None:
+            object.__setattr__(self, "topography_source", {"name": "ETOPO5"})
 
-        This method processes the topography data and generates a land/sea mask.
-        It applies several steps, including interpolating topography, smoothing
-        the topography over the entire domain and locally, and filling in enclosed basins. The
-        processed topography and mask are added to the grid's dataset as new variables.
+        if "name" not in self.topography_source:
+            raise ValueError(
+                "`topography_source` must include a 'name' key specifying the data source."
+            )
 
-        Parameters
-        ----------
-        hmin : float
-            The minimum ocean depth (in meters).
-        topography_source : str
-            Specifies the data source to use for the topography. Options are
-            "ETOPO5". Default is "ETOPO5".
+        if self.topography_source["name"] != "ETOPO5":
+            if "path" not in self.topography_source:
+                raise ValueError(
+                    "`topography_source` must include a 'path' key when the 'name' is not 'ETOPO5'."
+                )
 
-        Returns
-        -------
-        None
-            This method modifies the dataset in place and does not return a value.
-        """
+    def _create_mask(self, verbose=False) -> None:
+
+        if verbose:
+            start_time = time.time()
+            logging.info("=== Creating the mask ===")
+        ds = _add_mask(self.ds)
+
+        if verbose:
+            logging.info(f"Total time: {time.time() - start_time:.3f} seconds")
+            logging.info(
+                "========================================================================================================"
+            )
 
-        ds = _add_topography_and_mask(self.ds, topography_source, hmin)
-        # Assign the updated dataset back to the frozen dataclass
         object.__setattr__(self, "ds", ds)
-        object.__setattr__(self, "topography_source", topography_source)
-        object.__setattr__(self, "hmin", hmin)
 
-    def _straddle(self) -> None:
-        """Check if the Greenwich meridian goes through the domain.
+    def update_topography(
+        self, topography_source=None, hmin=None, verbose=False
+    ) -> None:
+        """Update the grid dataset with processed topography.
 
-        This method sets the `straddle` attribute to `True` if the Greenwich meridian
-        (0° longitude) intersects the domain defined by `lon_rho`. Otherwise, it sets
-        the `straddle` attribute to `False`.
+        This method performs several key operations, including regridding the topography, smoothing the
+        topography over the entire domain and locally.
+        The processed topography is then added to the grid's dataset.
 
-        The check is based on whether the longitudinal differences between adjacent
-        points exceed 300 degrees, indicating a potential wraparound of longitude.
-        """
+        Parameters
+        ----------
+        topography_source : dict, optional
+            A dictionary specifying the source of the topography data. The dictionary should
+            contain the following keys:
+            - "name" (str): The name of the topography data source (e.g., "SRTM15").
+            - "path" (Union[str, Path): The path to the raw data file.
 
-        if (
-            np.abs(self.ds.lon_rho.diff("xi_rho")).max() > 300
-            or np.abs(self.ds.lon_rho.diff("eta_rho")).max() > 300
-        ):
-            object.__setattr__(self, "straddle", True)
-        else:
-            object.__setattr__(self, "straddle", False)
+            If not provided, `topography_source` will remain unchanged (i.e., the existing value will not be overwritten).
 
-    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.
+        hmin : float, optional
+            The minimum ocean depth (in meters).
+            If not provided, `hmin` will remain unchanged (i.e., the existing value will not be overwritten).
 
-        The specific variables being coarsened are:
-        - `lon_rho` -> `lon_coarse`: Longitude at rho points.
-        - `lat_rho` -> `lat_coarse`: Latitude at rho points.
-        - `angle` -> `angle_coarse`: Angle between the xi axis and true east.
-        - `mask_rho` -> `mask_coarse`: Land/sea mask at rho points.
+        verbose : bool, optional
+            If True, the method will print detailed information about the grid generation process,
+            including the timing of each step. Defaults to False.
 
         Returns
         -------
         None
-
-        Modifies
-        --------
-        self.ds : xr.Dataset
-            The dataset attribute of the Grid instance is updated with the new coarser variables.
+            This method updates the internal dataset (`self.ds`) in place by adding or overwriting the
+            topography variable. It does not return any value.
         """
-        d = {
-            "angle": "angle_coarse",
-            "mask_rho": "mask_coarse",
-            "lat_rho": "lat_coarse",
-            "lon_rho": "lon_coarse",
-        }
 
-        for fine_var, coarse_var in d.items():
-            fine_field = self.ds[fine_var]
-            if self.straddle and fine_var == "lon_rho":
-                fine_field = xr.where(fine_field > 180, fine_field - 360, fine_field)
+        topography_source = topography_source or self.topography_source
+        hmin = hmin or self.hmin
 
-            coarse_field = _f2c(fine_field)
-            if fine_var == "lon_rho":
-                coarse_field = xr.where(
-                    coarse_field < 0, coarse_field + 360, coarse_field
-                )
-            if coarse_var in ["lon_coarse", "lat_coarse"]:
-                ds = self.ds.assign_coords({coarse_var: coarse_field})
-                object.__setattr__(self, "ds", ds)
-            else:
-                self.ds[coarse_var] = coarse_field
+        # Extract target coordinates for processing
+        target_coords = get_target_coords(self)
 
-        self.ds["mask_coarse"] = xr.where(self.ds["mask_coarse"] > 0.5, 1, 0).astype(
-            np.int32
+        # If verbose is enabled, start the timer and print the start message
+        if verbose:
+            start_time = time.time()
+            logging.info(
+                f"=== Generating the topography using {topography_source['name']} data and hmin = {hmin} meters ==="
+            )
+
+        # Add topography and mask to the dataset
+        ds = _add_topography(
+            ds=self.ds,
+            target_coords=target_coords,
+            topography_source=topography_source,
+            hmin=hmin,
+            verbose=verbose,
         )
 
-        for fine_var, coarse_var in d.items():
-            self.ds[coarse_var].attrs[
-                "long_name"
-            ] = f"{self.ds[fine_var].attrs['long_name']} on coarsened grid"
-            self.ds[coarse_var].attrs["units"] = self.ds[fine_var].attrs["units"]
+        # If verbose is enabled, print elapsed time and a separator
+        if verbose:
+            logging.info(f"Total time: {time.time() - start_time:.3f} seconds")
+            logging.info(
+                "========================================================================================================"
+            )
 
-    def update_vertical_coordinate(self, N, theta_s, theta_b, hc) -> None:
+        # Update the grid's dataset and related attributes
+        object.__setattr__(self, "ds", ds)
+        object.__setattr__(self, "topography_source", topography_source)
+        object.__setattr__(self, "hmin", hmin)
+
+    def update_vertical_coordinate(
+        self, N=None, theta_s=None, theta_b=None, hc=None, verbose=False
+    ) -> None:
         """Create vertical coordinate variables for the ROMS grid.
 
-        This method computes the S-coordinate stretching curves and depths
-        at various grid points (rho, u, v) using the specified parameters.
-        The computed depths and stretching curves are added to the dataset
-        as new coordinates, along with their corresponding attributes.
+        This method computes the S-coordinate stretching curves at rho- and
+        w-points.
 
         Parameters
         ----------
         N : int
             Number of vertical levels.
+            If not provided, `N` will remain unchanged (i.e., the existing value will not be overwritten).
         theta_s : float
             S-coordinate surface control parameter.
+            If not provided, `theta_s` will remain unchanged (i.e., the existing value will not be overwritten).
         theta_b : float
             S-coordinate bottom control parameter.
+            If not provided, `theta_b` will remain unchanged (i.e., the existing value will not be overwritten).
         hc : float
             Critical depth (m) used in ROMS vertical coordinate stretching.
+            If not provided, `hc` will remain unchanged (i.e., the existing value will not be overwritten).
+        verbose: bool, optional
+            Indicates whether to print vertical coordinate generation steps with timing. Defaults to False.
 
         Returns
         -------
@@ -236,6 +256,17 @@ class Grid:
             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
+        hc = hc or self.hc
+
+        if verbose:
+            start_time = time.time()
+            logging.info(
+                f"=== Preparing the vertical coordinate system using N = {N}, theta_s = {theta_s}, theta_b = {theta_b}, hc = {hc} ==="
+            )
+
         ds = self.ds
         # need to drop vertical coordinates because they could cause conflict if N changed
         vars_to_drop = [
@@ -245,6 +276,8 @@ class Grid:
             "interface_depth_rho",
             "interface_depth_u",
             "interface_depth_v",
+            "sigma_r",
+            "sigma_w",
             "Cs_w",
             "Cs_r",
         ]
@@ -253,74 +286,125 @@ class Grid:
             if var in ds.variables:
                 ds = ds.drop_vars(var)
 
-        h = ds.h
-
         cs_r, sigma_r = sigma_stretch(theta_s, theta_b, N, "r")
-        zr = compute_depth(h * 0, h, hc, cs_r, sigma_r)
         cs_w, sigma_w = sigma_stretch(theta_s, theta_b, N, "w")
-        zw = compute_depth(h * 0, h, hc, cs_w, sigma_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["units"] = "nondimensional"
 
         ds["Cs_r"] = cs_r.astype(np.float32)
-        ds["Cs_r"].attrs["long_name"] = "S-coordinate stretching curves at rho-points"
+        ds["Cs_r"].attrs["long_name"] = "Vertical stretching function at rho-points"
         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["units"] = "nondimensional"
+
         ds["Cs_w"] = cs_w.astype(np.float32)
-        ds["Cs_w"].attrs["long_name"] = "S-coordinate stretching curves at w-points"
+        ds["Cs_w"].attrs["long_name"] = "Vertical stretching function at w-points"
         ds["Cs_w"].attrs["units"] = "nondimensional"
 
         ds.attrs["theta_s"] = np.float32(theta_s)
         ds.attrs["theta_b"] = np.float32(theta_b)
         ds.attrs["hc"] = np.float32(hc)
 
-        depth = -zr
-        depth.attrs["long_name"] = "Layer depth at rho-points"
-        depth.attrs["units"] = "m"
+        if verbose:
+            logging.info(f"Total time: {time.time() - start_time:.3f} seconds")
+            logging.info(
+                "========================================================================================================"
+            )
 
-        depth_u = interpolate_from_rho_to_u(depth)
-        depth_u.attrs["long_name"] = "Layer depth at u-points"
-        depth_u.attrs["units"] = "m"
+        object.__setattr__(self, "ds", ds)
+        object.__setattr__(self, "theta_s", theta_s)
+        object.__setattr__(self, "theta_b", theta_b)
+        object.__setattr__(self, "hc", hc)
+        object.__setattr__(self, "N", N)
 
-        depth_v = interpolate_from_rho_to_v(depth)
-        depth_v.attrs["long_name"] = "Layer depth at v-points"
-        depth_v.attrs["units"] = "m"
+    def _straddle(self) -> None:
+        """Check if the Greenwich meridian goes through the domain.
 
-        interface_depth = -zw
-        interface_depth.attrs["long_name"] = "Interface depth at rho-points"
-        interface_depth.attrs["units"] = "m"
+        This method sets the `straddle` attribute to `True` if the Greenwich meridian
+        (0° longitude) intersects the domain defined by `lon_rho`. Otherwise, it sets
+        the `straddle` attribute to `False`.
+
+        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
+        ):
+            object.__setattr__(self, "straddle", True)
+        else:
+            object.__setattr__(self, "straddle", False)
 
-        interface_depth_u = interpolate_from_rho_to_u(interface_depth)
-        interface_depth_u.attrs["long_name"] = "Interface depth at u-points"
-        interface_depth_u.attrs["units"] = "m"
+    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.
 
-        interface_depth_v = interpolate_from_rho_to_v(interface_depth)
-        interface_depth_v.attrs["long_name"] = "Interface depth at v-points"
-        interface_depth_v.attrs["units"] = "m"
+        The specific variables being coarsened are:
+        - `lon_rho` -> `lon_coarse`: Longitude at rho points.
+        - `lat_rho` -> `lat_coarse`: Latitude at rho points.
+        - `angle` -> `angle_coarse`: Angle between the xi axis and true east.
+        - `mask_rho` -> `mask_coarse`: Land/sea mask at rho points.
+        """
+        d = {
+            "angle": "angle_coarse",
+            "mask_rho": "mask_coarse",
+            "lat_rho": "lat_coarse",
+            "lon_rho": "lon_coarse",
+        }
 
-        ds = ds.assign_coords(
-            {
-                "layer_depth_rho": depth.astype(np.float32),
-                "layer_depth_u": depth_u.astype(np.float32),
-                "layer_depth_v": depth_v.astype(np.float32),
-                "interface_depth_rho": interface_depth.astype(np.float32),
-                "interface_depth_u": interface_depth_u.astype(np.float32),
-                "interface_depth_v": interface_depth_v.astype(np.float32),
-            }
-        )
-        ds = ds.drop_vars(["eta_rho", "xi_rho"])
+        ds = self.ds
+
+        for fine_var, coarse_var in d.items():
+            fine_field = ds[fine_var]
+            if self.straddle and fine_var == "lon_rho":
+                fine_field = xr.where(fine_field > 180, fine_field - 360, fine_field)
+
+            coarse_field = _f2c(fine_field)
+            if fine_var == "lon_rho":
+                coarse_field = xr.where(
+                    coarse_field < 0, coarse_field + 360, coarse_field
+                )
+            if coarse_var in ["lon_coarse", "lat_coarse"]:
+                ds = ds.assign_coords({coarse_var: coarse_field})
+            else:
+                ds[coarse_var] = coarse_field
+
+            del fine_field, coarse_field
+
+        ds["mask_coarse"] = xr.where(ds["mask_coarse"] > 0.5, 1, 0).astype(np.int32)
+
+        for fine_var, coarse_var in d.items():
+            long_name = ds[fine_var].attrs.get(
+                "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"]
 
         object.__setattr__(self, "ds", ds)
-        object.__setattr__(self, "theta_s", theta_s)
-        object.__setattr__(self, "theta_b", theta_b)
-        object.__setattr__(self, "hc", hc)
-        object.__setattr__(self, "N", N)
 
-    def plot(self, bathymetry: bool = False) -> None:
+    def plot(
+        self, bathymetry: bool = False, title: str = None, with_dim_names: bool = False
+    ) -> None:
         """Plot the grid.
 
         Parameters
         ----------
-        bathymetry : bool
+        bathymetry : bool, optional
             Whether or not to plot the bathymetry. Default is False.
+        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.
+
 
         Returns
         -------
@@ -329,6 +413,8 @@ class Grid:
         """
 
         if bathymetry:
+            if title is None:
+                title = "ROMS grid and bathymetry"
             field = self.ds.h.where(self.ds.mask_rho)
             field = field.assign_coords(
                 {"lon": self.ds.lon_rho, "lat": self.ds.lat_rho}
@@ -344,28 +430,30 @@ class Grid:
                 self.ds,
                 field=field,
                 straddle=self.straddle,
+                title=title,
+                with_dim_names=with_dim_names,
                 kwargs=kwargs,
             )
         else:
-            _plot(self.ds, straddle=self.straddle)
+            if title is None:
+                title = "ROMS grid"
+            _plot(
+                self.ds,
+                straddle=self.straddle,
+                title=title,
+                with_dim_names=with_dim_names,
+            )
 
     def plot_vertical_coordinate(
-        self, varname="layer_depth_rho", s=None, eta=None, xi=None, ax=None
+        self,
+        s=None,
+        eta=None,
+        xi=None,
     ) -> None:
-        """Plot the vertical coordinate system for a given eta-, xi-, or s-slice.
+        """Plot the layer depth for a given eta-, xi-, or s-slice.
 
         Parameters
         ----------
-        varname : str, optional
-            The vertical coordinate field to plot. Options include:
-
-            - "layer_depth_rho": Layer depth at rho-points.
-            - "layer_depth_u": Layer depth at u-points.
-            - "layer_depth_v": Layer depth at v-points.
-            - "interface_depth_rho": Interface depth at rho-points.
-            - "interface_depth_u": Interface depth at u-points.
-            - "interface_depth_v": Interface depth at v-points.
-
         s: int, optional
             The s-index to plot. Default is None.
         eta : int, optional
@@ -383,105 +471,67 @@ class Grid:
         Raises
         ------
         ValueError
-            If the specified varname is not one of the valid options.
-            If none of s, eta, xi are specified.
+            If not exactly one of s, eta, xi is specified.
         """
 
-        if not any([s is not None, eta is not None, xi is not None]):
-            raise ValueError("At least one of s, eta, or xi must be specified.")
+        title = "Layer depth at rho-points"
 
-        self.ds[varname].load()
-        field = self.ds[varname].squeeze()
+        if sum(s is not None for s in [s, eta, xi]) != 1:
+            raise ValueError("Exactly one of s, eta, or xi must be specified.")
 
-        if all(dim in field.dims for dim in ["eta_rho", "xi_rho"]):
-            interface_depth = self.ds.interface_depth_rho
-            field = field.where(self.ds.mask_rho)
-            field = field.assign_coords(
-                {"lon": self.ds.lon_rho, "lat": self.ds.lat_rho}
-            )
-        elif all(dim in field.dims for dim in ["eta_rho", "xi_u"]):
-            interface_depth = self.ds.interface_depth_u
-            field = field.where(self.ds.mask_u)
-            field = field.assign_coords({"lon": self.ds.lon_u, "lat": self.ds.lat_u})
-        elif all(dim in field.dims for dim in ["eta_v", "xi_rho"]):
-            interface_depth = self.ds.interface_depth_v
-            field = field.where(self.ds.mask_v)
-            field = field.assign_coords({"lon": self.ds.lon_v, "lat": self.ds.lat_v})
-
-        # slice the field as desired
-        title = field.long_name
-        if s is not None:
-            if "s_rho" in field.dims:
-                title = title + f", s_rho = {field.s_rho[s].item()}"
-                field = field.isel(s_rho=s)
-            elif "s_w" in field.dims:
-                title = title + f", s_w = {field.s_w[s].item()}"
-                field = field.isel(s_w=s)
-            else:
-                raise ValueError(
-                    f"None of the expected dimensions (s_rho, s_w) found in ds[{varname}]."
-                )
+        h = self.ds["h"]
+        h = h.assign_coords({"lon": self.ds.lon_rho, "lat": self.ds.lat_rho})
 
+        # slice the bathymetry as desired
         if eta is not None:
-            if "eta_rho" in field.dims:
-                title = title + f", eta_rho = {field.eta_rho[eta].item()}"
-                field = field.isel(eta_rho=eta)
-                interface_depth = interface_depth.isel(eta_rho=eta)
-            elif "eta_v" in field.dims:
-                title = title + f", eta_v = {field.eta_v[eta].item()}"
-                field = field.isel(eta_v=eta)
-                interface_depth = interface_depth.isel(eta_v=eta)
-            else:
-                raise ValueError(
-                    f"None of the expected dimensions (eta_rho, eta_v) found in ds[{varname}]."
-                )
+            title = title + f", eta_rho = {h.eta_rho[eta].item()}"
+            h = h.isel(eta_rho=eta)
         if xi is not None:
-            if "xi_rho" in field.dims:
-                title = title + f", xi_rho = {field.xi_rho[xi].item()}"
-                field = field.isel(xi_rho=xi)
-                interface_depth = interface_depth.isel(xi_rho=xi)
-            elif "xi_u" in field.dims:
-                title = title + f", xi_u = {field.xi_u[xi].item()}"
-                field = field.isel(xi_u=xi)
-                interface_depth = interface_depth.isel(xi_u=xi)
-            else:
-                raise ValueError(
-                    f"None of the expected dimensions (xi_rho, xi_u) found in ds[{varname}]."
-                )
+            title = title + f", xi_rho = {h.xi_rho[xi].item()}"
+            h = h.isel(xi_rho=xi)
 
         if eta is None and xi is None:
-            vmax = field.max().values
-            vmin = field.min().values
+            layer_depth = compute_depth(0, h, self.hc, self.ds.Cs_r, self.ds.sigma_r)
+            title = title + f", s_rho = {layer_depth.s_rho[s].item()}"
+            layer_depth = layer_depth.isel(s_rho=s)
+
+            layer_depth.attrs["long_name"] = "Layer depth"
+            layer_depth.attrs["units"] = "m"
+
+            vmax = layer_depth.max().values
+            vmin = layer_depth.min().values
             cmap = plt.colormaps.get_cmap("YlGnBu")
             cmap.set_bad(color="gray")
             kwargs = {"vmax": vmax, "vmin": vmin, "cmap": cmap}
 
             _plot(
                 self.ds,
-                field=field,
+                field=layer_depth.where(self.ds.mask_rho),
                 straddle=self.straddle,
                 depth_contours=False,
                 title=title,
                 kwargs=kwargs,
             )
         else:
-            if len(field.dims) == 2:
-                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(
-                    xr.zeros_like(field),
-                    interface_depth=interface_depth,
-                    title=title,
-                    kwargs=kwargs,
-                    ax=ax,
-                )
-            else:
-                if "s_rho" in field.dims or "s_w" in field.dims:
-                    _profile_plot(field, title=title, ax=ax)
-                else:
-                    _line_plot(field, title=title, ax=ax)
+            layer_depth = compute_depth(0, h, self.hc, self.ds.Cs_r, self.ds.sigma_r)
+            layer_depth.attrs["long_name"] = "Layer depth"
+            layer_depth.attrs["units"] = "m"
+            field = xr.zeros_like(layer_depth)
+            field = field.assign_coords({"layer_depth": layer_depth})
+
+            interface_depth = compute_depth(
+                0, h, self.hc, self.ds.Cs_w, self.ds.sigma_w
+            )
+            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,
+            )
 
     def save(
         self, filepath: Union[str, Path], np_eta: int = None, np_xi: int = None
@@ -532,13 +582,15 @@ class Grid:
         return saved_filenames
 
     @classmethod
-    def from_file(cls, filepath: Union[str, Path]) -> "Grid":
+    def from_file(cls, filepath: Union[str, Path], 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.
+        verbose: bool, optional
+            Indicates whether to print grid generation steps with timing. Defaults to False.
 
         Returns
         -------
@@ -584,13 +636,14 @@ class Grid:
 
         # 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
 
             grid.update_vertical_coordinate(
-                N=N, theta_s=theta_s, theta_b=theta_b, hc=hc
+                N=N, theta_s=theta_s, theta_b=theta_b, hc=hc, verbose=True
             )
         else:
             object.__setattr__(grid, "theta_s", ds.attrs["theta_s"].item())
@@ -639,7 +692,10 @@ class Grid:
             "hmin",
         ]:
             if attr in ds.attrs:
-                a = ds.attrs[attr]
+                if attr == "topography_source":
+                    a = {"name": ds.attrs[attr]}
+                else:
+                    a = ds.attrs[attr]
             else:
                 a = None
             object.__setattr__(grid, attr, a)
@@ -661,6 +717,7 @@ class Grid:
         data = asdict(self)
         data.pop("ds", None)
         data.pop("straddle", None)
+        data.pop("verbose", None)
 
         # Include the version of roms-tools
         try:
@@ -681,18 +738,38 @@ class Grid:
             yaml.dump(yaml_data, file, default_flow_style=False, sort_keys=False)
 
     @classmethod
-    def from_yaml(cls, filepath: Union[str, Path]) -> "Grid":
+    def from_yaml(
+        cls,
+        filepath: Union[str, Path],
+        section_name: str = "Grid",
+        verbose: bool = False,
+    ) -> "Grid":
         """Create an instance of the class from a YAML file.
 
         Parameters
         ----------
         filepath : Union[str, Path]
             The path to the YAML file from which the parameters will be read.
+        section_name : str, optional
+            The name of the YAML section containing the grid configuration. Defaults to "Grid".
+        verbose : bool, optional
+            Indicates whether to print grid generation steps with timing. Defaults to False.
 
         Returns
         -------
         Grid
-            An instance of the Grid class.
+            An instance of the Grid class initialized with the parameters from the YAML file.
+
+        Raises
+        ------
+        ValueError
+            If the ROMS-Tools version is not found in the YAML file or if the specified section
+            does not exist in the file.
+
+        Warnings
+        --------
+        Issues a warning if the ROMS-Tools version in the YAML header does not match the
+        currently installed version.
         """
 
         filepath = Path(filepath)
@@ -712,8 +789,8 @@ class Grid:
                 continue
             if "roms_tools_version" in doc:
                 header_data = doc
-            elif "Grid" in doc:
-                grid_data = doc["Grid"]
+            elif section_name in doc:
+                grid_data = doc[section_name]
 
         if header_data is None:
             raise ValueError("Version of ROMS-Tools not found in the YAML file.")
@@ -733,164 +810,341 @@ class Grid:
 
         if grid_data is None:
             raise ValueError("No Grid configuration found in the YAML file.")
+        return cls(**grid_data, verbose=verbose)
 
-        return cls(**grid_data)
-
-    # override __repr__ method to only print attributes that are actually set
     def __repr__(self) -> str:
+        """Return a string representation of the object with non-None attributes,
+        excluding 'ds'."""
         cls = self.__class__
         cls_name = cls.__name__
-        # Create a dictionary of attribute names and values, filtering out those that are not set and 'ds'
+        # Filter attributes to exclude 'ds' and those with None values
         attr_dict = {
             k: v for k, v in self.__dict__.items() if k != "ds" and v is not None
         }
         attr_str = ", ".join(f"{k}={v!r}" for k, v in attr_dict.items())
         return f"{cls_name}({attr_str})"
 
+    def _create_horizontal_grid(self) -> xr.Dataset():
+        """Create the horizontal grid based on a Mercator projection and store it in the
+        'ds' attribute.
 
-def _make_grid_ds(
-    nx: int,
-    ny: int,
-    size_x: float,
-    size_y: float,
-    center_lon: float,
-    center_lat: float,
-    rot: float,
-) -> xr.Dataset:
-    _raise_if_domain_size_too_large(size_x, size_y)
-
-    initial_lon_lat_vars = _make_initial_lon_lat_ds(size_x, size_y, nx, ny)
-
-    # rotate coordinate system
-    rotated_lon_lat_vars = _rotate(*initial_lon_lat_vars, rot)
-
-    # translate coordinate system
-    translated_lon_lat_vars = _translate(*rotated_lon_lat_vars, center_lat, center_lon)
-    lon, lat, lonu, latu, lonv, latv, lonq, latq = translated_lon_lat_vars
-
-    # compute 1/dx and 1/dy
-    pm, pn = _compute_coordinate_metrics(lon, lonu, latu, lonv, latv)
-
-    # compute angle of local grid positive x-axis relative to east
-    ang = _compute_angle(lon, lonu, latu, lonq)
-
-    # make sure lons are in [0, 360] range
-    lon[lon < 0] = lon[lon < 0] + 2 * np.pi
-    lonu[lonu < 0] = lonu[lonu < 0] + 2 * np.pi
-    lonv[lonv < 0] = lonv[lonv < 0] + 2 * np.pi
-    lonq[lonq < 0] = lonq[lonq < 0] + 2 * np.pi
-
-    ds = _create_grid_ds(
-        lon,
-        lat,
-        lonu,
-        latu,
-        lonv,
-        latv,
-        lonq,
-        latq,
-        pm,
-        pn,
-        ang,
-        rot,
-        center_lon,
-        center_lat,
-    )
+        Parameters
+        ----------
+        None
 
-    ds = _add_global_metadata(ds, size_x, size_y, center_lon, center_lat, rot)
+        Returns
+        -------
+        xr.Dataset
+            The created horizontal grid dataset, including coordinates, grid metrics, angles, and metadata.
 
-    return ds
+        Notes
+        -----
+        - Longitude values are adjusted to fall within the range [0, 360].
+        - Grid rotation and translation are applied based on the specified parameters.
+        """
+        if self.verbose:
+            start_time = time.time()
+            logging.info("=== Creating the horizontal grid ===")
 
+        self._raise_if_domain_size_too_large()
 
-def _raise_if_domain_size_too_large(size_x, size_y):
-    threshold = 20000
-    if size_x > threshold or size_y > threshold:
-        raise ValueError("Domain size has to be smaller than %g km" % threshold)
+        coords = self._make_initial_lon_lat_ds()
 
+        # rotate coordinate system
+        coords = _rotate(coords, self.rot)
 
-def _make_initial_lon_lat_ds(size_x, size_y, nx, ny):
-    # Mercator projection around the equator
+        # translate coordinate system
+        coords = _translate(coords, self.center_lat, self.center_lon)
 
-    # initially define the domain to be longer in x-direction (dimension "length")
-    # than in y-direction (dimension "width") to keep grid distortion minimal
-    if size_y > size_x:
-        domain_length, domain_width = size_y * 1e3, size_x * 1e3  # in m
-        nl, nw = ny, nx
-    else:
-        domain_length, domain_width = size_x * 1e3, size_y * 1e3  # in m
-        nl, nw = nx, ny
+        # compute 1/dx and 1/dy
+        coords["pm"], coords["pn"] = _compute_coordinate_metrics(coords)
 
-    domain_length_in_degrees = domain_length / RADIUS_OF_EARTH
-    domain_width_in_degrees = domain_width / RADIUS_OF_EARTH
+        # compute angle of local grid positive x-axis relative to east
+        coords["angle"] = _compute_angle(coords)
 
-    # 1d array describing the longitudes at cell centers
-    x = np.arange(-0.5, nl + 1.5, 1)
-    lon_array_1d_in_degrees = (
-        domain_length_in_degrees * x / nl - domain_length_in_degrees / 2
-    )
-    # 1d array describing the longitudes at cell corners (or vorticity points "q")
-    xq = np.arange(-1, nl + 2, 1)
-    lonq_array_1d_in_degrees_q = (
-        domain_length_in_degrees * xq / nl - domain_length_in_degrees / 2
-    )
+        # make sure lons are in [0, 360] range
+        for lon in ["lon", "lonu", "lonv", "lonq"]:
+            coords[lon][coords[lon] < 0] = coords[lon][coords[lon] < 0] + 2 * np.pi
+
+        ds = self._create_grid_ds(coords)
+
+        ds = self._add_global_metadata(ds)
+
+        if self.verbose:
+            logging.info(f"Total time: {time.time() - start_time:.3f} seconds")
+            logging.info(
+                "========================================================================================================"
+            )
+
+        object.__setattr__(self, "ds", ds)
+
+    def _add_global_metadata(self, ds):
+        """Add global metadata and attributes to the dataset.
+
+        Parameters
+        ----------
+        ds : xr.Dataset
+            Dataset to which global metadata and attributes will be added.
+
+        Returns
+        -------
+        xr.Dataset
+            The dataset with added global metadata, including grid type, tool version,
+            grid dimensions, center coordinates, and rotation.
+
+        Notes
+        -----
+        - The "spherical" attribute indicates the grid type and is set to "T" (spherical).
+        - The ROMS-Tools version is included as "roms_tools_version". If unavailable, it defaults to "unknown".
+        """
+        ds["spherical"] = xr.DataArray(np.array("T", dtype="S1"))
+        ds["spherical"].attrs["Long_name"] = "Grid type logical switch"
+        ds["spherical"].attrs["option_T"] = "spherical"
 
-    # convert degrees latitude to y-coordinate using Mercator projection
-    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))
+        ds.attrs["title"] = "ROMS grid created by ROMS-Tools"
 
-    # linearly space points in y-space
-    y = (y2 - y1) * np.arange(-0.5, nw + 1.5, 1) / nw + y1
-    yq = (y2 - y1) * np.arange(-1, nw + 2) / nw + y1
+        # Include the version of roms-tools
+        try:
+            roms_tools_version = importlib.metadata.version("roms-tools")
+        except importlib.metadata.PackageNotFoundError:
+            roms_tools_version = "unknown"
+
+        ds.attrs["roms_tools_version"] = roms_tools_version
+        ds.attrs["size_x"] = self.size_x
+        ds.attrs["size_y"] = self.size_y
+        ds.attrs["center_lon"] = self.center_lon
+        ds.attrs["center_lat"] = self.center_lat
+        ds.attrs["rot"] = self.rot
+
+        return ds
+
+    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.
+        """
+        threshold = 20000
+        if self.size_x > threshold or self.size_y > threshold:
+            raise ValueError(
+                f"Domain size exceeds the allowable limit of {threshold} 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."
+            )
+
+    def _make_initial_lon_lat_ds(self):
+        """Generate initial longitude and latitude arrays with Mercator projection
+        around the equator.
+
+        Returns
+        -------
+        dict
+            A dictionary containing the following arrays:
+            - lon, lat: 2D arrays of longitudes and latitudes at cell centers.
+            - lonu, latu: 2D arrays of longitudes and latitudes at u-points.
+            - lonv, latv: 2D arrays of longitudes and latitudes at v-points.
+            - lonq, latq: 2D arrays of longitudes and latitudes at cell corners.
+        """
+
+        r_earth = 6371315.0
+
+        # 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
+            nl, nw = self.ny, self.nx
+        else:
+            domain_length, domain_width = self.size_x * 1e3, self.size_y * 1e3  # in m
+            nl, nw = self.nx, self.ny
+
+        domain_length_in_degrees = domain_length / r_earth
+        domain_width_in_degrees = 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
+        )
+
+        # 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))
+
+        # 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)
+        )
+        latq_array_1d_in_degrees = np.arctan(
+            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)
+
+        if self.size_y > self.size_x:
+            # Rotate grid by 90 degrees because until here the grid has been defined
+            # to be longer in x-direction than in y-direction
+
+            lon, lat = _rot_sphere(lon, lat, 90)
+            lonq, latq = _rot_sphere(lonq, latq, 90)
+
+            lon = np.transpose(np.flip(lon, 0))
+            lat = np.transpose(np.flip(lat, 0))
+            lonq = np.transpose(np.flip(lonq, 0))
+            latq = np.transpose(np.flip(latq, 0))
+
+        # Inference for u- and v-point coordinates
+        lonu = 0.5 * (lon[:, :-1] + lon[:, 1:])
+        latu = 0.5 * (lat[:, :-1] + lat[:, 1:])
+        lonv = 0.5 * (lon[:-1, :] + lon[1:, :])
+        latv = 0.5 * (lat[:-1, :] + lat[1:, :])
+
+        coords = {
+            "lon": lon,
+            "lat": lat,
+            "lonu": lonu,
+            "latu": latu,
+            "lonv": lonv,
+            "latv": latv,
+            "lonq": lonq,
+            "latq": latq,
+        }
+
+        return coords
+
+    def _create_grid_ds(self, coords):
+        """Create an xarray Dataset with grid coordinates and metrics.
+
+        Parameters
+        ----------
+        coords : dict
+            Dictionary containing:
+            - lon, lat, lonu, latu, lonv, latv : 1d arrays of coordinates (degrees)
+            - angle : 2d array (radians)
+            - pm, pn : 2d arrays (meter^-1)
 
-    # inverse Mercator projections
-    lat_array_1d_in_degrees = np.arctan(np.sinh(y))
-    latq_array_1d_in_degrees = np.arctan(np.sinh(yq))
+        Returns
+        -------
+        xarray.Dataset
+            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(
+            data=coords["lon"] * 180 / np.pi,
+            dims=["eta_rho", "xi_rho"],
+            attrs={"long_name": "longitude of rho-points", "units": "degrees East"},
+        )
+        lat_rho = xr.Variable(
+            data=coords["lat"] * 180 / np.pi,
+            dims=["eta_rho", "xi_rho"],
+            attrs={"long_name": "latitude of rho-points", "units": "degrees North"},
+        )
+        lon_u = xr.Variable(
+            data=coords["lonu"] * 180 / np.pi,
+            dims=["eta_rho", "xi_u"],
+            attrs={"long_name": "longitude of u-points", "units": "degrees East"},
+        )
+        lat_u = xr.Variable(
+            data=coords["latu"] * 180 / np.pi,
+            dims=["eta_rho", "xi_u"],
+            attrs={"long_name": "latitude of u-points", "units": "degrees North"},
+        )
+        lon_v = xr.Variable(
+            data=coords["lonv"] * 180 / np.pi,
+            dims=["eta_v", "xi_rho"],
+            attrs={"long_name": "longitude of v-points", "units": "degrees East"},
+        )
+        lat_v = xr.Variable(
+            data=coords["latv"] * 180 / np.pi,
+            dims=["eta_v", "xi_rho"],
+            attrs={"long_name": "latitude of v-points", "units": "degrees North"},
+        )
+        # lon_q = xr.Variable(
+        #    data=coords["lonq"] * 180 / np.pi,
+        #    dims=["eta_psi", "xi_psi"],
+        #    attrs={"long_name": "longitude of psi-points", "units": "degrees East"},
+        # )
+        # lat_q = xr.Variable(
+        #    data=coords["latq"] * 180 / np.pi,
+        #    dims=["eta_psi", "xi_psi"],
+        #    attrs={"long_name": "latitude of psi-points", "units": "degrees North"},
+        # )
 
-    # 2d grid at cell centers
-    lon, lat = np.meshgrid(lon_array_1d_in_degrees, lat_array_1d_in_degrees)
-    # 2d grid at cell corners
-    lonq, latq = np.meshgrid(lonq_array_1d_in_degrees_q, latq_array_1d_in_degrees)
+        ds = ds.assign_coords(
+            {
+                "lat_rho": lat_rho,
+                "lon_rho": lon_rho,
+                "lat_u": lat_u,
+                "lon_u": lon_u,
+                "lat_v": lat_v,
+                "lon_v": lon_v,
+                # "lat_psi": lat_q,
+                # "lon_psi": lon_q,
+            }
+        )
 
-    if size_y > size_x:
-        # Rotate grid by 90 degrees because until here the grid has been defined
-        # to be longer in x-direction than in y-direction
+        ds["angle"] = xr.Variable(
+            data=coords["angle"],
+            dims=["eta_rho", "xi_rho"],
+            attrs={"long_name": "Angle between xi axis and east", "units": "radians"},
+        )
 
-        lon, lat = _rot_sphere(lon, lat, 90)
-        lonq, latq = _rot_sphere(lonq, latq, 90)
+        # Coriolis frequency
+        f0 = 4 * np.pi * np.sin(coords["lat"]) / (24 * 3600)
 
-        lon = np.transpose(np.flip(lon, 0))
-        lat = np.transpose(np.flip(lat, 0))
-        lonq = np.transpose(np.flip(lonq, 0))
-        latq = np.transpose(np.flip(latq, 0))
+        ds["f"] = xr.Variable(
+            data=f0,
+            dims=["eta_rho", "xi_rho"],
+            attrs={
+                "long_name": "Coriolis parameter at rho-points",
+                "units": "second-1",
+            },
+        )
 
-    # infer longitudes and latitudes at u- and v-points
-    lonu = 0.5 * (lon[:, :-1] + lon[:, 1:])
-    latu = 0.5 * (lat[:, :-1] + lat[:, 1:])
-    lonv = 0.5 * (lon[:-1, :] + lon[1:, :])
-    latv = 0.5 * (lat[:-1, :] + lat[1:, :])
+        ds["pm"] = xr.Variable(
+            data=coords["pm"],
+            dims=["eta_rho", "xi_rho"],
+            attrs={
+                "long_name": "Curvilinear coordinate metric in xi-direction",
+                "units": "meter-1",
+            },
+        )
+        ds["pn"] = xr.Variable(
+            data=coords["pn"],
+            dims=["eta_rho", "xi_rho"],
+            attrs={
+                "long_name": "Curvilinear coordinate metric in eta-direction",
+                "units": "meter-1",
+            },
+        )
 
-    # TODO wrap up into temporary container Dataset object?
-    return lon, lat, lonu, latu, lonv, latv, lonq, latq
+        return ds
 
 
-def _rotate(lon, lat, lonu, latu, lonv, latv, lonq, latq, rot):
+def _rotate(coords, rot):
     """Rotate grid counterclockwise relative to surface of Earth by rot degrees."""
 
-    (lon, lat) = _rot_sphere(lon, lat, rot)
-    (lonu, latu) = _rot_sphere(lonu, latu, rot)
-    (lonv, latv) = _rot_sphere(lonv, latv, rot)
-    (lonq, latq) = _rot_sphere(lonq, latq, rot)
+    (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)
+    (coords["lonq"], coords["latq"]) = _rot_sphere(coords["lonq"], coords["latq"], rot)
 
-    return lon, lat, lonu, latu, lonv, latv, lonq, latq
+    return coords
 
 
-def _translate(lon, lat, lonu, latu, lonv, latv, lonq, latq, tra_lat, tra_lon):
+def _translate(coords, tra_lat, tra_lon):
     """Translate grid so that the centre lies at the position (tra_lat, tra_lon)"""
 
-    (lon, lat) = _tra_sphere(lon, lat, tra_lat)
-    (lonu, latu) = _tra_sphere(lonu, latu, tra_lat)
-    (lonv, latv) = _tra_sphere(lonv, latv, tra_lat)
-    (lonq, latq) = _tra_sphere(lonq, latq, tra_lat)
+    (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)
+    (lonq, latq) = _tra_sphere(coords["lonq"], coords["latq"], tra_lat)
 
     lon = lon + tra_lon * np.pi / 180
     lonu = lonu + tra_lon * np.pi / 180
@@ -902,133 +1156,171 @@ def _translate(lon, lat, lonu, latu, lonv, latv, lonq, latq, tra_lat, tra_lon):
     lonv[lonv < -np.pi] = lonv[lonv < -np.pi] + 2 * np.pi
     lonq[lonq < -np.pi] = lonq[lonq < -np.pi] + 2 * np.pi
 
-    return lon, lat, lonu, latu, lonv, latv, lonq, latq
+    coords = {
+        "lon": lon,
+        "lat": lat,
+        "lonu": lonu,
+        "latu": latu,
+        "lonv": lonv,
+        "latv": latv,
+        "lonq": lonq,
+        "latq": latq,
+    }
+
+    return coords
 
 
 def _rot_sphere(lon, lat, rot):
-    (n, m) = np.shape(lon)
-    # convert rotation angle from degrees to radians
+    """Rotate longitude and latitude coordinates on a sphere.
+
+    Parameters
+    ----------
+    lon : ndarray
+        2D array of longitudes in radians.
+    lat : ndarray
+        2D array of latitudes in radians.
+    rot : float
+        Rotation angle in degrees.
+
+    Returns
+    -------
+    tuple
+        Rotated longitude and latitude arrays (lon, lat) in radians.
+    """
+    # Convert rotation angle from degrees to radians
     rot = rot * np.pi / 180
 
-    # translate into Cartesian coordinates x,y,z
-    # conventions:  (lon,lat) = (0,0)  corresponds to (x,y,z) = ( 0,-r, 0)
-    #               (lon,lat) = (0,90) corresponds to (x,y,z) = ( 0, 0, r)
+    # Convert spherical coordinates to Cartesian coordinates (x, y, z)
     x1 = np.sin(lon) * np.cos(lat)
     y1 = np.cos(lon) * np.cos(lat)
     z1 = np.sin(lat)
 
-    # We will rotate these points around the small circle defined by
-    # the intersection of the sphere and the plane that
-    # is orthogonal to the line through (lon,lat) (0,0) and (180,0)
-
-    # The rotation is in that plane around its intersection with
-    # aforementioned line.
-
-    # Since the plane is orthogonal to the y-axis (in my definition at least),
-    # Rotations in the plane of the small circle maintain constant y and are around
-    # (x,y,z) = (0,y1,0)
-
+    # Calculate the radial distance in the x-z plane
     rp1 = np.sqrt(x1**2 + z1**2)
 
-    ap1 = np.pi / 2 * np.ones((n, m))
-    ap1[np.abs(x1) > 1e-7] = np.arctan(
-        np.abs(z1[np.abs(x1) > 1e-7] / x1[np.abs(x1) > 1e-7])
-    )
+    # Compute azimuthal angle
+    ap1 = np.arctan2(np.abs(z1), np.abs(x1))
     ap1[x1 < 0] = np.pi - ap1[x1 < 0]
     ap1[z1 < 0] = -ap1[z1 < 0]
 
+    # Apply rotation to the azimuthal angle
     ap2 = ap1 + rot
     x2 = rp1 * np.cos(ap2)
     y2 = y1
     z2 = rp1 * np.sin(ap2)
 
-    lon = np.pi / 2 * np.ones((n, m))
-    lon[abs(y2) > 1e-7] = np.arctan(
-        np.abs(x2[np.abs(y2) > 1e-7] / y2[np.abs(y2) > 1e-7])
-    )
-    lon[y2 < 0] = np.pi - lon[y2 < 0]
-    lon[x2 < 0] = -lon[x2 < 0]
+    # Recompute longitude and latitude
+    lon_rot = np.arctan2(np.abs(x2), np.abs(y2))
+    lon_rot[y2 < 0] = np.pi - lon_rot[y2 < 0]
+    lon_rot[x2 < 0] = -lon_rot[x2 < 0]
 
     pr2 = np.sqrt(x2**2 + y2**2)
-    lat = np.pi / 2 * np.ones((n, m))
-    lat[np.abs(pr2) > 1e-7] = np.arctan(
-        np.abs(z2[np.abs(pr2) > 1e-7] / pr2[np.abs(pr2) > 1e-7])
-    )
-    lat[z2 < 0] = -lat[z2 < 0]
+    lat_rot = np.arctan2(np.abs(z2), pr2)
+    lat_rot[z2 < 0] = -lat_rot[z2 < 0]
 
-    return (lon, lat)
+    return lon_rot, lat_rot
 
 
 def _tra_sphere(lon, lat, tra):
-    (n, m) = np.shape(lon)
-    tra = tra * np.pi / 180  # translation in latitude direction
+    """Translate longitude and latitude coordinates on a sphere in the latitude
+    direction.
 
-    # translate into x,y,z
-    # conventions:  (lon,lat) = (0,0)  corresponds to (x,y,z) = ( 0,-r, 0)
-    #               (lon,lat) = (0,90) corresponds to (x,y,z) = ( 0, 0, r)
-    x1 = np.sin(lon) * np.cos(lat)
-    y1 = np.cos(lon) * np.cos(lat)
-    z1 = np.sin(lat)
+    Parameters
+    ----------
+    lon : ndarray
+        2D array of longitudes in radians.
+    lat : ndarray
+        2D array of latitudes in radians.
+    tra : float
+        Translation angle in degrees.
 
-    # We will rotate these points around the small circle defined by
-    # the intersection of the sphere and the plane that
-    # is orthogonal to the line through (lon,lat) (90,0) and (-90,0)
+    Returns
+    -------
+    tuple
+        Translated longitude and latitude arrays (lon, lat) in radians.
+    """
 
-    # The rotation is in that plane around its intersection with
-    # aforementioned line.
+    # Convert translation angle from degrees to radians
+    tra = tra * np.pi / 180
 
-    # Since the plane is orthogonal to the x-axis (in my definition at least),
-    # Rotations in the plane of the small circle maintain constant x and are around
-    # (x,y,z) = (x1,0,0)
+    # Convert spherical coordinates to Cartesian coordinates (x, y, z)
+    x1 = np.sin(lon) * np.cos(lat)
+    y1 = np.cos(lon) * np.cos(lat)
+    z1 = np.sin(lat)
 
+    # Radial distance in the y-z plane
     rp1 = np.sqrt(y1**2 + z1**2)
 
-    ap1 = np.pi / 2 * np.ones((n, m))
-    ap1[np.abs(y1) > 1e-7] = np.arctan(
-        np.abs(z1[np.abs(y1) > 1e-7] / y1[np.abs(y1) > 1e-7])
-    )
+    # Compute azimuthal angle in the y-z plane
+    ap1 = np.arctan2(np.abs(z1), np.abs(y1))
     ap1[y1 < 0] = np.pi - ap1[y1 < 0]
     ap1[z1 < 0] = -ap1[z1 < 0]
 
+    # Apply translation in the azimuthal angle
     ap2 = ap1 + tra
-    x2 = x1
     y2 = rp1 * np.cos(ap2)
     z2 = rp1 * np.sin(ap2)
 
-    ## transformation from (x,y,z) to (lat,lon)
-    lon = np.pi / 2 * np.ones((n, m))
-    lon[np.abs(y2) > 1e-7] = np.arctan(
-        np.abs(x2[np.abs(y2) > 1e-7] / y2[np.abs(y2) > 1e-7])
-    )
-    lon[y2 < 0] = np.pi - lon[y2 < 0]
-    lon[x2 < 0] = -lon[x2 < 0]
+    # Convert back to spherical coordinates
+    lon_rot = np.arctan2(np.abs(x1), np.abs(y2))
+    lon_rot[y2 < 0] = np.pi - lon_rot[y2 < 0]
+    lon_rot[x1 < 0] = -lon_rot[x1 < 0]
 
-    pr2 = np.sqrt(x2**2 + y2**2)
-    lat = np.pi / (2 * np.ones((n, m)))
-    lat[np.abs(pr2) > 1e-7] = np.arctan(
-        np.abs(z2[np.abs(pr2) > 1e-7] / pr2[np.abs(pr2) > 1e-7])
-    )
-    lat[z2 < 0] = -lat[z2 < 0]
+    pr2 = np.sqrt(x1**2 + y2**2)
+    lat_rot = np.arctan2(np.abs(z2), pr2)
+    lat_rot[z2 < 0] = -lat_rot[z2 < 0]
+
+    return lon_rot, lat_rot
 
-    return (lon, lat)
 
+def _compute_coordinate_metrics(coords):
+    """Compute the reciprocal of grid spacing (`pn` and `pm`) in the latitude and
+    longitude directions.
+
+    Parameters
+    ----------
+    coords : dict
+        A dictionary containing coordinate arrays 'lonu', 'latu', 'lonv', and 'latv' for the u- and v-velocity points.
 
-def _compute_coordinate_metrics(lon, lonu, latu, lonv, latv):
-    """Compute the curvilinear coordinate metrics pn and pm, defined as 1/grid
-    spacing."""
+    Returns
+    -------
+    pn : ndarray
+        The metric for the latitude direction (1/dy).
+
+    pm : ndarray
+        The metric for the longitude direction (1/dx).
+
+    Notes
+    -----
+    Boundary values of `pn` and `pm` are copied from adjacent interior values.
+    """
 
     # pm = 1/dx
-    pmu = gc_dist(lonu[:, :-1], latu[:, :-1], lonu[:, 1:], latu[:, 1:])
-    pm = 0 * lon
+    pmu = gc_dist(
+        coords["lonu"][:, :-1],
+        coords["latu"][:, :-1],
+        coords["lonu"][:, 1:],
+        coords["latu"][:, 1:],
+        input_in_degrees=False,
+    )
+    pm = np.zeros_like(coords["lon"])
     pm[:, 1:-1] = pmu
+    # Handle boundary conditions
     pm[:, 0] = pm[:, 1]
     pm[:, -1] = pm[:, -2]
     pm = 1 / pm
 
     # pn = 1/dy
-    pnv = gc_dist(lonv[:-1, :], latv[:-1, :], lonv[1:, :], latv[1:, :])
-    pn = 0 * lon
+    pnv = gc_dist(
+        coords["lonv"][:-1, :],
+        coords["latv"][:-1, :],
+        coords["lonv"][1:, :],
+        coords["latv"][1:, :],
+        input_in_degrees=False,
+    )
+    pn = np.zeros_like(coords["lon"])
     pn[1:-1, :] = pnv
+    # Handle boundary conditions
     pn[0, :] = pn[1, :]
     pn[-1, :] = pn[-2, :]
     pn = 1 / pn
@@ -1036,179 +1328,50 @@ def _compute_coordinate_metrics(lon, lonu, latu, lonv, latv):
     return pn, pm
 
 
-def gc_dist(lon1, lat1, lon2, lat2):
-    # Distance between 2 points along a great circle
-    # lat and lon in radians!!
-    # 2008, Jeroen Molemaker, UCLA
-
-    dlat = lat2 - lat1
-    dlon = lon2 - lon1
+def _compute_angle(coords):
+    """Compute angles of the local grid's positive x-axis relative to east.
 
-    dang = 2 * np.arcsin(
-        np.sqrt(
-            np.sin(dlat / 2) ** 2 + np.cos(lat2) * np.cos(lat1) * np.sin(dlon / 2) ** 2
-        )
-    )  # haversine function
+    The angle is computed for each grid cell using the latitude and longitude
+    differences between neighboring grid points. The result is wrapped to
+    the range [-π, π] and adjusted based on longitude and latitude conditions.
 
-    dis = RADIUS_OF_EARTH * dang
+    Parameters
+    ----------
+    coords : dict
+        A dictionary containing 'latu' (latitudes) and 'lonu' (longitudes) arrays.
 
-    return dis
+    Returns
+    -------
+    ang : ndarray
+        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]
 
-def _compute_angle(lon, lonu, latu, lonq):
-    """Compute angles of local grid positive x-axis relative to east."""
+    # Normalize longitude differences to the range [-π, π]
+    dellon = (dellon + np.pi) % (2 * np.pi) - np.pi
+    dellon *= np.cos(0.5 * (coords["latu"][:, 1:] + coords["latu"][:, :-1]))
 
-    dellat = latu[:, 1:] - latu[:, :-1]
-    dellon = lonu[:, 1:] - lonu[:, :-1]
-    dellon[dellon > np.pi] = dellon[dellon > np.pi] - 2 * np.pi
-    dellon[dellon < -np.pi] = dellon[dellon < -np.pi] + 2 * np.pi
-    dellon = dellon * np.cos(0.5 * (latu[:, 1:] + latu[:, :-1]))
+    # Compute the angle in radians
+    ang_s = np.arctan2(dellat, dellon)
 
-    ang = copy.copy(lon)
-    ang_s = np.arctan(dellat / (dellon + 1e-16))
-    ang_s[(dellon < 0) & (dellat < 0)] = ang_s[(dellon < 0) & (dellat < 0)] - np.pi
-    ang_s[(dellon < 0) & (dellat >= 0)] = ang_s[(dellon < 0) & (dellat >= 0)] + np.pi
-    ang_s[ang_s > np.pi] = ang_s[ang_s > np.pi] - np.pi
-    ang_s[ang_s < -np.pi] = ang_s[ang_s < -np.pi] + np.pi
+    # Adjust angles based on longitude and latitude conditions
+    ang_s[(dellon < 0) & (dellat < 0)] -= np.pi
+    ang_s[(dellon < 0) & (dellat >= 0)] += np.pi
+    ang_s = np.mod(ang_s + np.pi, 2 * np.pi) - np.pi  # Ensure angles are in [-π, π]
 
+    # Create output array and set angles
+    ang = np.zeros_like(coords["lon"])
     ang[:, 1:-1] = ang_s
-    ang[:, 0] = ang[:, 1]
-    ang[:, -1] = ang[:, -2]
+    ang[:, 0] = ang[:, 1]  # Set first column to the second column
+    ang[:, -1] = ang[:, -2]  # Set last column to the second-to-last column
 
     return ang
 
 
-def _create_grid_ds(
-    lon,
-    lat,
-    lonu,
-    latu,
-    lonv,
-    latv,
-    lonq,
-    latq,
-    pm,
-    pn,
-    angle,
-    rot,
-    center_lon,
-    center_lat,
-):
-    ds = xr.Dataset()
-
-    lon_rho = xr.Variable(
-        data=lon * 180 / np.pi,
-        dims=["eta_rho", "xi_rho"],
-        attrs={"long_name": "longitude of rho-points", "units": "degrees East"},
-    )
-    lat_rho = xr.Variable(
-        data=lat * 180 / np.pi,
-        dims=["eta_rho", "xi_rho"],
-        attrs={"long_name": "latitude of rho-points", "units": "degrees North"},
-    )
-    lon_u = xr.Variable(
-        data=lonu * 180 / np.pi,
-        dims=["eta_rho", "xi_u"],
-        attrs={"long_name": "longitude of u-points", "units": "degrees East"},
-    )
-    lat_u = xr.Variable(
-        data=latu * 180 / np.pi,
-        dims=["eta_rho", "xi_u"],
-        attrs={"long_name": "latitude of u-points", "units": "degrees North"},
-    )
-    lon_v = xr.Variable(
-        data=lonv * 180 / np.pi,
-        dims=["eta_v", "xi_rho"],
-        attrs={"long_name": "longitude of v-points", "units": "degrees East"},
-    )
-    lat_v = xr.Variable(
-        data=latv * 180 / np.pi,
-        dims=["eta_v", "xi_rho"],
-        attrs={"long_name": "latitude of v-points", "units": "degrees North"},
-    )
-    # lon_q = xr.Variable(
-    #    data=lonq * 180 / np.pi,
-    #    dims=["eta_psi", "xi_psi"],
-    #    attrs={"long_name": "longitude of psi-points", "units": "degrees East"},
-    # )
-    # lat_q = xr.Variable(
-    #    data=latq * 180 / np.pi,
-    #    dims=["eta_psi", "xi_psi"],
-    #    attrs={"long_name": "latitude of psi-points", "units": "degrees North"},
-    # )
-
-    ds = ds.assign_coords(
-        {
-            "lat_rho": lat_rho,
-            "lon_rho": lon_rho,
-            "lat_u": lat_u,
-            "lon_u": lon_u,
-            "lat_v": lat_v,
-            "lon_v": lon_v,
-            # "lat_psi": lat_q,
-            # "lon_psi": lon_q,
-        }
-    )
-
-    ds["angle"] = xr.Variable(
-        data=angle,
-        dims=["eta_rho", "xi_rho"],
-        attrs={"long_name": "Angle between xi axis and east", "units": "radians"},
-    )
-
-    # Coriolis frequency
-    f0 = 4 * np.pi * np.sin(lat) / (24 * 3600)
-
-    ds["f"] = xr.Variable(
-        data=f0,
-        dims=["eta_rho", "xi_rho"],
-        attrs={"long_name": "Coriolis parameter at rho-points", "units": "second-1"},
-    )
-
-    ds["pm"] = xr.Variable(
-        data=pm,
-        dims=["eta_rho", "xi_rho"],
-        attrs={
-            "long_name": "Curvilinear coordinate metric in xi-direction",
-            "units": "meter-1",
-        },
-    )
-    ds["pn"] = xr.Variable(
-        data=pn,
-        dims=["eta_rho", "xi_rho"],
-        attrs={
-            "long_name": "Curvilinear coordinate metric in eta-direction",
-            "units": "meter-1",
-        },
-    )
-
-    return ds
-
-
-def _add_global_metadata(ds, size_x, size_y, center_lon, center_lat, rot):
-
-    ds["spherical"] = xr.DataArray(np.array("T", dtype="S1"))
-    ds["spherical"].attrs["Long_name"] = "Grid type logical switch"
-    ds["spherical"].attrs["option_T"] = "spherical"
-
-    ds.attrs["title"] = "ROMS grid created by ROMS-Tools"
-
-    # Include the version of roms-tools
-    try:
-        roms_tools_version = importlib.metadata.version("roms-tools")
-    except importlib.metadata.PackageNotFoundError:
-        roms_tools_version = "unknown"
-
-    ds.attrs["roms_tools_version"] = roms_tools_version
-    ds.attrs["size_x"] = size_x
-    ds.attrs["size_y"] = size_y
-    ds.attrs["center_lon"] = center_lon
-    ds.attrs["center_lat"] = center_lat
-    ds.attrs["rot"] = rot
-
-    return ds
-
-
 def _f2c(f):
     """Coarsen input xarray DataArray f in both x- and y-direction.