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

ci/environment-with-xesmf.yml

@@ -3,7 +3,8 @@ channels:
   - conda-forge
   - nodefaults
 dependencies:
-  - python>=3.10
+  - python>=3.11
+  - numpy>2,<2.3
   - xesmf
   # testing
   - zarr

ci/environment.yml

@@ -3,7 +3,8 @@ channels:
   - conda-forge
   - nodefaults
 dependencies:
-  - python>=3.10
+  - python>=3.11
+  - numpy>2,<2.3
   # testing
   - zarr
   - pytest

roms_tools/__init__.py

@@ -1,5 +1,5 @@
+import logging
 from importlib.metadata import version as _version
-import logging  # noqa: F811
 
 try:
     __version__ = _version("roms_tools")
@@ -7,17 +7,18 @@ except ImportError:  # pragma: no cover
     # Local copy or not installed with setuptools
     __version__ = "9999"
 
-
-from roms_tools.setup.grid import Grid  # noqa: F401
-from roms_tools.setup.tides import TidalForcing  # noqa: F401
-from roms_tools.setup.surface_forcing import SurfaceForcing  # noqa: F401
-from roms_tools.setup.initial_conditions import InitialConditions  # noqa: F401
+# grid must be imported first
+from roms_tools.setup.grid import Grid  # noqa: I001, F401
+from roms_tools.analysis.roms_output import ROMSOutput  # noqa: F401
 from roms_tools.setup.boundary_forcing import BoundaryForcing  # noqa: F401
-from roms_tools.setup.river_forcing import RiverForcing  # noqa: F401
-from roms_tools.setup.cdr_forcing import CDRVolumePointSource  # noqa: F401
+from roms_tools.setup.cdr_forcing import CDRForcing  # noqa: F401
+from roms_tools.setup.cdr_release import TracerPerturbation, VolumeRelease  # noqa: F401
+from roms_tools.setup.initial_conditions import InitialConditions  # noqa: F401
 from roms_tools.setup.nesting import ChildGrid  # noqa: F401
+from roms_tools.setup.river_forcing import RiverForcing  # noqa: F401
+from roms_tools.setup.surface_forcing import SurfaceForcing  # noqa: F401
+from roms_tools.setup.tides import TidalForcing  # noqa: F401
 from roms_tools.tiling.partition import partition_netcdf  # noqa: F401
-from roms_tools.analysis.roms_output import ROMSOutput  # noqa: F401
 
 # Configure logging when the package is imported
 logging.basicConfig(level=logging.INFO, format="%(levelname)s - %(message)s")

roms_tools/analysis/roms_output.py

@@ -1,27 +1,30 @@
-import xarray as xr
-import numpy as np
-import matplotlib.pyplot as plt
-from roms_tools.plot import _plot, _section_plot, _profile_plot, _line_plot
-from roms_tools.regrid import LateralRegridFromROMS, VerticalRegridFromROMS
-from dataclasses import dataclass, field
-from typing import Union, Optional
-from pathlib import Path
-import re
 import logging
+import re
 import warnings
+from dataclasses import dataclass, field
 from datetime import datetime, timedelta
+from pathlib import Path
+from typing import Optional, Union
+
+import matplotlib.pyplot as plt
+import numpy as np
+import xarray as xr
+
 from roms_tools import Grid
-from roms_tools.vertical_coordinate import (
-    compute_depth_coordinates,
-)
+from roms_tools.analysis.utils import _validate_plot_inputs
+from roms_tools.plot import _line_plot, _plot, _profile_plot, _section_plot
+from roms_tools.regrid import LateralRegridFromROMS, VerticalRegridFromROMS
 from roms_tools.utils import (
+    _generate_coordinate_range,
     _load_data,
+    _remove_edge_nans,
     interpolate_from_rho_to_u,
     interpolate_from_rho_to_v,
-    _generate_coordinate_range,
-    _remove_edge_nans,
+    normalize_longitude,
+)
+from roms_tools.vertical_coordinate import (
+    compute_depth_coordinates,
 )
-from roms_tools.analysis.utils import _validate_plot_inputs
 
 
 @dataclass(kw_only=True)
@@ -136,7 +139,7 @@ class ROMSOutput:
             By default, depth contours are not shown (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.
 
         save_path : str, optional
             Path to save the generated plot. If None, the plot is shown interactively.
@@ -212,6 +215,9 @@ class ROMSOutput:
         lon_deg = self.grid.ds[f"lon_{loc}"]
         if self.grid.straddle:
             lon_deg = xr.where(lon_deg > 180, lon_deg - 360, lon_deg)
+        if lon is not None:
+            lon = normalize_longitude(lon, self.grid.straddle)
+
         field = field.assign_coords({"lon": lon_deg, "lat": lat_deg})
 
         # Mask the field
@@ -450,7 +456,7 @@ class ROMSOutput:
 
         # Prepare vertical regrid
         if depth_levels is None:
-            depth_levels, _ = self._compute_exponential_depth_levels()
+            depth_levels, _ = self.grid._compute_exponential_depth_levels()
 
         # Ensure depth_levels is an xarray.DataArray
         if not isinstance(depth_levels, xr.DataArray):
@@ -786,52 +792,3 @@ class ROMSOutput:
         # Add all necessary coordinates in one go
         ds = ds.assign_coords(coords_to_add)
         return ds
-
-    def _compute_exponential_depth_levels(self, Nz=None, 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 `len(self.ds.s_rho)`.
-
-        depth : float, optional
-            Total depth of the domain. Defaults to `grid.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 = len(self.ds.s_rho)
-        if depth is None:
-            depth = self.grid.ds.h.max().values
-        if h is None:
-            h = Nz / 4.5
-
-        k = np.arange(1, Nz + 2)
-
-        # Define the exponential profile function
-        def exponential_profile(k, Nz, h):
-            return np.exp(k / h)
-
-        z_faces = np.vectorize(exponential_profile)(k, Nz, h)
-
-        # Normalize
-        z_faces -= z_faces[0]
-        z_faces *= 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

roms_tools/constants.py

@@ -1,2 +1,5 @@
 R_EARTH = 6371315.0  # Earth radius in meters
+MAXIMUM_GRID_SIZE = 25000  # in km
+UPPER_BOUND_THETA_S = 10  # upper bound for surface vertical stretching parameter
+UPPER_BOUND_THETA_B = 10  # upper bound for bottom vertical stretching parameter
 NUM_TRACERS = 34  # Number of tracers (temperature, salinity, BGC tracers)

roms_tools/download.py

@@ -79,7 +79,7 @@ pup_test_data = pooch.create(
         "CESM_BGC_SURFACE_2012.nc": "3c4d156adca97909d0fac36bf50b99583ab37d8020d7a3e8511e92abf2331b38",
         "CESM_surface_global_test_data_climatology.nc": "a072757110c6f7b716a98f867688ef4195a5966741d2f368201ac24617254e35",
         "CESM_surface_global_test_data.nc": "874106ffbc8b1b220db09df1551bbb89d22439d795b4d1e5a24ee775e9a7bf6e",
-        "coarsened_UNIFIED_bgc_dataset.nc": "58b510dfd9125e8bfd92342af8cbf7c2639ef9a42c21f3ee74fc31b361e3750c",
+        "coarsened_UNIFIED_bgc_dataset.nc": "c5eaeca6f7640f2358afe5c46ace99ff4879948df8a778a2f21967de034f8d0a",
         "grid_created_with_matlab.nc": "fd537ef8159fabb18e38495ec8d44e2fa1b7fb615fcb1417dd4c0e1bb5f4e41d",
         "etopo5_coarsened_and_shifted.nc": "9a5cb4b38c779d22ddb0ad069b298b9722db34ca85a89273eccca691e89e6f96",
         "srtm15_coarsened.nc": "48bc8f4beecfdca9c192b13f4cbeef1455f49d8261a82563aaec5757e100dff9",

roms_tools/plot.py

@@ -1,7 +1,7 @@
 import cartopy.crs as ccrs
 import matplotlib.pyplot as plt
-import xarray as xr
 import numpy as np
+import xarray as xr
 
 
 def _plot(
@@ -41,8 +41,8 @@ def _plot(
 
     Returns
     -------
-    matplotlib.figure.Figure
-        The generated figure with the plotted data, if a new figure was created.
+    matplotlib.figure.Figure, optional
+        The generated figure with the plotted data, only returned if `ax` is None.
 
     Raises
     ------
@@ -232,8 +232,8 @@ def _section_plot(field, interface_depth=None, title="", kwargs={}, ax=None):
 
     Returns
     -------
-    matplotlib.figure.Figure
-        The generated figure with the plotted section.
+    matplotlib.figure.Figure, optional
+        The generated figure with the plotted section, only returned if `ax` is None.
 
     Raises
     ------
@@ -304,7 +304,8 @@ def _section_plot(field, interface_depth=None, title="", kwargs={}, ax=None):
         xlabel = xdim
     ax.set_xlabel(xlabel)
 
-    return fig
+    if ax is None:
+        return fig
 
 
 def _profile_plot(field, title="", ax=None):
@@ -326,7 +327,7 @@ def _profile_plot(field, title="", ax=None):
     Returns
     -------
     matplotlib.figure.Figure
-        The generated figure with the plotted profile.
+        The generated figure with the plotted profile, only returned if `ax` is None.
 
     Raises
     ------
@@ -361,7 +362,8 @@ def _profile_plot(field, title="", ax=None):
     ax.set_ylabel("Depth [m]")
     ax.grid()
 
-    return fig
+    if ax is None:
+        return fig
 
 
 def _line_plot(field, title="", ax=None):
@@ -380,7 +382,8 @@ def _line_plot(field, title="", ax=None):
     Returns
     -------
     matplotlib.figure.Figure
-        The generated figure with the plotted data and highlighted NaN regions.
+        The generated figure with the plotted data and highlighted NaN regions, only
+        returned if `ax` is None.
 
     Raises
     ------
@@ -447,7 +450,8 @@ def _line_plot(field, title="", ax=None):
         xlabel = xdim
     ax.set_xlabel(xlabel)
 
-    return fig
+    if ax is None:
+        return fig
 
 
 def _add_boundary_to_ax(

roms_tools/regrid.py

@@ -1,7 +1,8 @@
-import xgcm
-import xarray as xr
 import warnings
 
+import xarray as xr
+import xgcm
+
 
 class LateralRegridToROMS:
     """Handles lateral regridding of data onto a new spatial grid."""

roms_tools/setup/boundary_forcing.py

@@ -1,39 +1,42 @@
-import xarray as xr
-import numpy as np
-from scipy.ndimage import label
-import logging
 import importlib.metadata
-from typing import Dict, Union, List, Optional
+import logging
 from dataclasses import dataclass, field
 from datetime import datetime
-import matplotlib.pyplot as plt
 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 scipy.ndimage import label
+
 from roms_tools import Grid
+from roms_tools.plot import _line_plot, _section_plot
 from roms_tools.regrid import LateralRegridToROMS, VerticalRegridToROMS
-from roms_tools.utils import save_datasets
-from roms_tools.vertical_coordinate import compute_depth
-from roms_tools.plot import _section_plot, _line_plot
-from roms_tools.utils import (
-    interpolate_from_rho_to_u,
-    interpolate_from_rho_to_v,
-    transpose_dimensions,
-)
-from roms_tools.setup.datasets import GLORYSDataset, CESMBGCDataset, UnifiedBGCDataset
+from roms_tools.setup.datasets import CESMBGCDataset, GLORYSDataset, UnifiedBGCDataset
 from roms_tools.setup.utils import (
-    get_variable_metadata,
-    group_dataset,
-    get_target_coords,
-    rotate_velocities,
+    _from_yaml,
+    _to_dict,
+    _write_to_yaml,
+    add_time_info_to_ds,
     compute_barotropic_velocity,
     compute_missing_bgc_variables,
-    one_dim_fill,
+    get_boundary_coords,
+    get_target_coords,
+    get_variable_metadata,
+    group_dataset,
     nan_check,
+    one_dim_fill,
+    rotate_velocities,
     substitute_nans_by_fillvalue,
-    convert_to_roms_time,
-    get_boundary_coords,
-    _to_yaml,
-    _from_yaml,
 )
+from roms_tools.utils import (
+    interpolate_from_rho_to_u,
+    interpolate_from_rho_to_v,
+    save_datasets,
+    transpose_dimensions,
+)
+from roms_tools.vertical_coordinate import compute_depth
 
 
 @dataclass(kw_only=True)
@@ -686,8 +689,7 @@ class BoundaryForcing:
         ds.attrs["theta_b"] = self.grid.ds.attrs["theta_b"]
         ds.attrs["hc"] = self.grid.ds.attrs["hc"]
 
-        # Convert the time coordinate to the format expected by ROMS
-        ds, bry_time = convert_to_roms_time(
+        ds, bry_time = add_time_info_to_ds(
             ds, self.model_reference_date, data.climatology
         )
 
@@ -1003,7 +1005,8 @@ class BoundaryForcing:
             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(