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

roms_tools/setup/grid.py

@@ -1,28 +1,31 @@
-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
+from typing import Dict, List, Union
 
+import matplotlib.pyplot as plt
 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 roms_tools.setup.mask import _add_mask, _add_velocity_masks
-from roms_tools.vertical_coordinate import compute_depth_coordinates, sigma_stretch
+
+from roms_tools.constants import MAXIMUM_GRID_SIZE, R_EARTH
 from roms_tools.plot import _plot, _section_plot
+from roms_tools.setup.mask import _add_mask, _add_velocity_masks
+from roms_tools.setup.topography import _add_topography
 from roms_tools.setup.utils import (
+    _pop_grid_data,
+    _write_to_yaml,
+    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,
 )
-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 +72,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,7 +81,7 @@ 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
@@ -390,7 +393,8 @@ 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
 
@@ -445,12 +449,7 @@ class Grid:
                 field=field, title=title, with_dim_names=with_dim_names, plot_data=False
             )
 
-    def plot_vertical_coordinate(
-        self,
-        s=None,
-        eta=None,
-        xi=None,
-    ) -> None:
+    def plot_vertical_coordinate(self, s=None, eta=None, xi=None, ax=None) -> None:
         """Plot the layer depth for a given eta-, xi-, or s-slice.
 
         Parameters
@@ -462,7 +461,7 @@ class Grid:
         xi : int, optional
             The xi-index to plot. Default is None.
         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
         -------
@@ -526,6 +525,7 @@ class Grid:
                 interface_depth=interface_depth,
                 title=title,
                 kwargs=kwargs,
+                ax=ax,
             )
 
     def save(self, filepath: Union[str, Path]) -> None:
@@ -557,13 +557,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 +588,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 +634,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,18 +753,28 @@ 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:
@@ -686,29 +786,11 @@ 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)
+        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(
@@ -894,13 +976,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."
             )
@@ -1136,6 +1217,51 @@ class Grid:
 
         return resolution_in_degrees
 
+    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
+        ----------
+        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
+        -------
+        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 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)
+
+        z_faces = np.exp(k / h)
+
+        # Normalize
+        z_faces -= z_faces[0]
+        z_faces *= max_depth / z_faces[-1]
+        z_faces[0] = 0.0
+
+        # Calculate center depths (average between adjacent face depths)
+        z_centers = (z_faces[:-1] + z_faces[1:]) / 2
+
+        # 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 z_centers, z_faces
+
 
 def _rotate(coords, rot):
     """Rotate grid counterclockwise relative to surface of Earth by rot degrees."""

roms_tools/setup/initial_conditions.py

@@ -1,37 +1,40 @@
-import xarray as xr
-import numpy as np
 import importlib.metadata
-from dataclasses import dataclass, field
-from typing import Dict, Union, List, Optional
-import matplotlib.pyplot as plt
-from pathlib import Path
 import logging
+from dataclasses import dataclass, field
 from datetime import datetime
+from pathlib import Path
+from typing import Dict, List, Optional, Union
+
+import matplotlib.pyplot as plt
+import numpy as np
+import xarray as xr
+
 from roms_tools import Grid
+from roms_tools.plot import _line_plot, _plot, _profile_plot, _section_plot
 from roms_tools.regrid import LateralRegridToROMS, VerticalRegridToROMS
-from roms_tools.plot import _plot, _section_plot, _profile_plot, _line_plot
+from roms_tools.setup.datasets import CESMBGCDataset, GLORYSDataset, UnifiedBGCDataset
+from roms_tools.setup.utils import (
+    _from_yaml,
+    _to_dict,
+    _write_to_yaml,
+    compute_barotropic_velocity,
+    compute_missing_bgc_variables,
+    get_target_coords,
+    get_variable_metadata,
+    nan_check,
+    rotate_velocities,
+    substitute_nans_by_fillvalue,
+)
 from roms_tools.utils import (
-    transpose_dimensions,
-    save_datasets,
     get_dask_chunks,
     interpolate_from_rho_to_u,
     interpolate_from_rho_to_v,
+    save_datasets,
+    transpose_dimensions,
 )
 from roms_tools.vertical_coordinate import (
-    compute_depth_coordinates,
     compute_depth,
-)
-from roms_tools.setup.datasets import GLORYSDataset, CESMBGCDataset, UnifiedBGCDataset
-from roms_tools.setup.utils import (
-    nan_check,
-    substitute_nans_by_fillvalue,
-    get_variable_metadata,
-    get_target_coords,
-    rotate_velocities,
-    compute_barotropic_velocity,
-    compute_missing_bgc_variables,
-    _to_yaml,
-    _from_yaml,
+    compute_depth_coordinates,
 )
 
 
@@ -683,7 +686,7 @@ class InitialConditions:
             visualize the layering of the water column. For clarity, the number of layer
             contours displayed is limited to a maximum of 10. Default is False.
         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. Default is None.
+            The axes to plot on. If None, a new figure is created. Note that this argument is ignored for 2D horizontal plots. Default is None.
 
         Returns
         -------
@@ -926,7 +929,8 @@ class InitialConditions:
             The path to the YAML file where the parameters will be saved.
         """
 
-        _to_yaml(self, filepath)
+        forcing_dict = _to_dict(self)
+        _write_to_yaml(forcing_dict, filepath)
 
     @classmethod
     def from_yaml(

roms_tools/setup/mask.py

@@ -1,12 +1,14 @@
-import xarray as xr
+import warnings
+
 import numpy as np
 import regionmask
-import warnings
+import xarray as xr
 from scipy.ndimage import label
+
 from roms_tools.setup.utils import (
+    handle_boundaries,
     interpolate_from_rho_to_u,
     interpolate_from_rho_to_v,
-    handle_boundaries,
 )
 
 

roms_tools/setup/nesting.py

@@ -1,23 +1,25 @@
-import numpy as np
-import xarray as xr
-from scipy.interpolate import griddata
+import logging
 from dataclasses import dataclass, field
-from typing import Dict, Union, Any, Tuple
 from pathlib import Path
-import logging
-from scipy.interpolate import interp1d
+from typing import Any, Dict, Tuple, Union
+
+import numpy as np
+import xarray as xr
+from scipy.interpolate import griddata, interp1d
+
 from roms_tools import Grid
 from roms_tools.plot import _plot_nesting
-from roms_tools.utils import save_datasets
 from roms_tools.setup.topography import _clip_depth
 from roms_tools.setup.utils import (
+    _from_yaml,
+    _to_dict,
+    _write_to_yaml,
+    get_boundary_coords,
     interpolate_from_rho_to_u,
     interpolate_from_rho_to_v,
-    get_boundary_coords,
     wrap_longitudes,
-    _to_yaml,
-    _from_yaml,
 )
+from roms_tools.utils import save_datasets
 
 
 @dataclass(kw_only=True)
@@ -209,7 +211,8 @@ class ChildGrid(Grid):
             The path to the YAML file where the parameters will be saved.
         """
 
-        _to_yaml(self, filepath)
+        forcing_dict = _to_dict(self)
+        _write_to_yaml(forcing_dict, filepath)
 
     @classmethod
     def from_yaml(cls, filepath: Union[str, Path]) -> "ChildGrid":