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

roms_tools/__init__.py

@@ -14,6 +14,7 @@ from roms_tools.setup.surface_forcing import SurfaceForcing  # noqa: F401
 from roms_tools.setup.initial_conditions import InitialConditions  # 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.nesting import ChildGrid  # noqa: F401
 from roms_tools.tiling.partition import partition_netcdf  # noqa: F401
 from roms_tools.analysis.roms_output import ROMSOutput  # noqa: F401

roms_tools/analysis/roms_output.py

@@ -2,7 +2,6 @@ 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.utils import _load_data
 from roms_tools.regrid import LateralRegridFromROMS, VerticalRegridFromROMS
 from dataclasses import dataclass, field
 from typing import Union, Optional
@@ -15,8 +14,14 @@ from roms_tools import Grid
 from roms_tools.vertical_coordinate import (
     compute_depth_coordinates,
 )
-from roms_tools.utils import interpolate_from_rho_to_u, interpolate_from_rho_to_v
-from roms_tools.analysis.utils import _validate_plot_inputs, _generate_coordinate_range
+from roms_tools.utils import (
+    _load_data,
+    interpolate_from_rho_to_u,
+    interpolate_from_rho_to_v,
+    _generate_coordinate_range,
+    _remove_edge_nans,
+)
+from roms_tools.analysis.utils import _validate_plot_inputs
 
 
 @dataclass(kw_only=True)
@@ -301,7 +306,7 @@ class ROMSOutput:
                 lats = [lat]
                 title = title + f", lat = {lat}°N"
             else:
-                resolution = self._infer_nominal_horizontal_resolution()
+                resolution = self.grid._infer_nominal_horizontal_resolution()
                 lats = _generate_coordinate_range(
                     field.lat.min().values, field.lat.max().values, resolution
                 )
@@ -311,7 +316,7 @@ class ROMSOutput:
                 lons = [lon]
                 title = title + f", lon = {lon}°E"
             else:
-                resolution = self._infer_nominal_horizontal_resolution(lat)
+                resolution = self.grid._infer_nominal_horizontal_resolution(lat)
                 lons = _generate_coordinate_range(
                     field.lon.min().values, field.lon.max().values, resolution
                 )
@@ -327,33 +332,6 @@ class ROMSOutput:
         if compute_layer_depth:
             field = field.assign_coords({"layer_depth": layer_depth})
 
-        def _remove_edge_nans(field, xdim, layer_depth=None):
-            """Removes NaNs from the edges along the specified dimension."""
-            if xdim in field.dims:
-                if layer_depth is not None:
-                    nan_mask = layer_depth.isnull().sum(
-                        dim=[dim for dim in layer_depth.dims if dim != xdim]
-                    )
-                else:
-                    nan_mask = field.isnull().sum(
-                        dim=[dim for dim in field.dims if dim != xdim]
-                    )
-
-                # Find the valid indices where the sum of the nans is 0
-                valid_indices = np.where(nan_mask.values == 0)[0]
-
-                if len(valid_indices) > 0:
-                    first_valid = valid_indices[0]
-                    last_valid = valid_indices[-1]
-
-                    field = field.isel({xdim: slice(first_valid, last_valid + 1)})
-                    if layer_depth is not None:
-                        layer_depth = layer_depth.isel(
-                            {xdim: slice(first_valid, last_valid + 1)}
-                        )
-
-            return field, layer_depth
-
         if lat is not None:
             field, layer_depth = _remove_edge_nans(
                 field, "lon", layer_depth if "layer_depth" in locals() else None
@@ -459,7 +437,7 @@ class ROMSOutput:
             lon_deg = xr.where(lon_deg > 180, lon_deg - 360, lon_deg)
 
         if horizontal_resolution is None:
-            horizontal_resolution = self._infer_nominal_horizontal_resolution()
+            horizontal_resolution = self.grid._infer_nominal_horizontal_resolution()
         lons = _generate_coordinate_range(
             lon_deg.min().values, lon_deg.max().values, horizontal_resolution
         )
@@ -809,50 +787,6 @@ class ROMSOutput:
         ds = ds.assign_coords(coords_to_add)
         return ds
 
-    def _infer_nominal_horizontal_resolution(self, lat=None):
-        """Estimate the nominal horizontal resolution of the grid in degrees at a
-        specified latitude.
-
-        This method calculates the nominal horizontal resolution of the grid by first
-        determining the average grid spacing in meters. The spacing is then converted
-        to degrees, accounting for the Earth's curvature, and the latitude where the
-        resolution is being computed.
-
-        Parameters
-        ----------
-        lat : float, optional
-            Latitude (in degrees) at which to estimate the horizontal resolution.
-            If not provided, the resolution is calculated at the average latitude of
-            the grid (`lat_rho`).
-
-        Returns
-        -------
-        float
-            The estimated horizontal resolution in degrees, adjusted for the Earth's curvature.
-        """
-        # Earth radius in meters
-        r_earth = 6371315.0
-
-        if lat is None:
-            # Center latitude in degrees
-            lat = (self.grid.ds.lat_rho.max() + self.grid.ds.lat_rho.min()) / 2
-
-        # Convert latitude to radians
-        lat_rad = np.deg2rad(lat)
-
-        # Mean resolution in meters
-        resolution_in_m = (
-            (1 / self.grid.ds.pm).mean() + (1 / self.grid.ds.pn).mean()
-        ) / 2
-
-        # Meters per degree at the equator
-        meters_per_degree = 2 * np.pi * r_earth / 360
-
-        # Correct for latitude by multiplying by cos(latitude) for longitude
-        resolution_in_degrees = resolution_in_m / (meters_per_degree * np.cos(lat_rad))
-
-        return resolution_in_degrees
-
     def _compute_exponential_depth_levels(self, Nz=None, depth=None, h=None):
         """Compute vertical grid center and face depths using an exponential profile.
 

roms_tools/analysis/utils.py

@@ -1,6 +1,3 @@
-import numpy as np
-
-
 def _validate_plot_inputs(field, s, eta, xi, depth, lat, lon, include_boundary):
     """Validate input parameters for the plot method.
 
@@ -104,66 +101,3 @@ def _validate_plot_inputs(field, s, eta, xi, depth, lat, lon, include_boundary):
                     f"Invalid xi index: {xi} lies on the boundary, which is excluded when `include_boundary = False`. "
                     "Either set `include_boundary = True`, or adjust eta to avoid boundary values."
                 )
-
-
-def _generate_coordinate_range(min, max, resolution):
-    """Generate an array of target coordinates (e.g., latitude or longitude) within a
-    specified range, with a resolution that is rounded to the nearest value of the form
-    `1/n` (or integer).
-
-    This method generates an array of target coordinates between the provided `min` and `max`
-    values, ensuring that both `min` and `max` are included in the resulting range. The resolution
-    is rounded to the nearest fraction of the form `1/n` or an integer, based on the input.
-
-    Parameters
-    ----------
-    min : float
-        The minimum value (in degrees) of the coordinate range (inclusive).
-
-    max : float
-        The maximum value (in degrees) of the coordinate range (inclusive).
-
-    resolution : float
-        The spacing (in degrees) between each coordinate in the array. The resolution will
-        be rounded to the nearest value of the form `1/n` or an integer, depending on the size
-        of the resolution value.
-
-    Returns
-    -------
-    numpy.ndarray
-        An array of target coordinates generated from the specified range, with the resolution
-        rounded to a suitable fraction (e.g., `1/n`) or integer, depending on the input resolution.
-    """
-
-    # Find the closest fraction of the form 1/n or integer to match the resolution
-    resolution_rounded = None
-    min_diff = float("inf")  # Initialize the minimum difference as infinity
-
-    # Search for the best fraction or integer approximation to the resolution
-    for n in range(1, 1000):  # Try fractions 1/n, where n ranges from 1 to 999
-        if resolution <= 1:
-            fraction = (
-                1 / n
-            )  # For small resolutions (<= 1), consider fractions of the form 1/n
-        else:
-            fraction = n  # For larger resolutions (>1), consider integers (n)
-
-        diff = abs(
-            fraction - resolution
-        )  # Calculate the difference between the fraction and the resolution
-
-        if diff < min_diff:  # If the current fraction is a better approximation
-            min_diff = diff
-            resolution_rounded = fraction  # Update the best fraction (or integer) found
-
-    # Adjust the start and end of the range to include integer values
-    start_int = np.floor(min)  # Round the minimum value down to the nearest integer
-    end_int = np.ceil(max)  # Round the maximum value up to the nearest integer
-
-    # Generate the array of target coordinates, including both the min and max values
-    target = np.arange(start_int, end_int + resolution_rounded, resolution_rounded)
-
-    # Truncate any values that exceed max (including small floating point errors)
-    target = target[target <= end_int + 1e-10]
-
-    return target.astype(np.float32)

roms_tools/constants.py

@@ -0,0 +1,2 @@
+R_EARTH = 6371315.0  # Earth radius in meters
+NUM_TRACERS = 34  # Number of tracers (temperature, salinity, BGC tracers)

roms_tools/download.py

@@ -22,7 +22,7 @@ correction_data = pooch.create(
     },
 )
 
-# Create a Pooch object to manage the global topography data
+# Create a Pooch object to manage the global river data
 river_data = pooch.create(
     # Use the default cache folder for the operating system
     path=pooch.os_cache("roms-tools"),
@@ -32,6 +32,18 @@ river_data = pooch.create(
         "dai_trenberth_may2019.nc": "sha256:793849e6aa60d1f6bdb480c345515fb2453d903c0a30599241b3d752f53715ab",
     },
 )
+
+# Create a Pooch object to manage the global SAL TPXO data
+sal_data = pooch.create(
+    # Use the default cache folder for the operating system
+    path=pooch.os_cache("roms-tools"),
+    base_url="https://github.com/CWorthy-ocean/roms-tools-data/raw/main/",
+    # The registry specifies the files that can be fetched
+    registry={
+        "sal_tpxo9.v2a.nc": "sha256:5343d745b4374170a069bfc2c67dcdd9f1dc4eb0df7c0e6de5c004432c903f40",
+    },
+)
+
 # Create a Pooch object to manage the test data
 pup_test_data = pooch.create(
     # Use the default cache folder for the operating system
@@ -46,8 +58,19 @@ pup_test_data = pooch.create(
         "GLORYS_NA_20121231.nc": "03c1155087195deff76ad3f136d6a7f35bc01ccae3402f3d95557a2886d39e71",
         "ERA5_regional_test_data.nc": "bd12ce3b562fbea2a80a3b79ba74c724294043c28dc98ae092ad816d74eac794",
         "ERA5_global_test_data.nc": "8ed177ab64c02caf509b9fb121cf6713f286cc603b1f302f15f3f4eb0c21dc4f",
-        "TPXO_global_test_data.nc": "457bfe87a7b247ec6e04e3c7d3e741ccf223020c41593f8ae33a14f2b5255e60",
-        "TPXO_regional_test_data.nc": "11739245e2286d9c9d342dce5221e6435d2072b50028bef2e86a30287b3b4032",
+        "global_grid_tpxo10.v2.nc": "26eb97cd135cd6f2b4e894c5f11bf7f860ff19cec8dbaa9190e37d30ee6e744e",
+        "global_h_tpxo10.v2.nc": "ef60fae6d52fa514dcc59a737435d74aa798dc114b57f01b123aa39dbaffc592",
+        "global_u_tpxo10.v2.nc": "022e57e6287e51f52eb1e5296614b1086e0e22ecd0bd57c9fd8d0e155babf5c3",
+        "regional_grid_tpxo10v2a.nc": "c5022bfe93ead7cd46e836578645bd87cb5be63c736e660937c7f5703c968cbc",
+        "regional_h_tpxo10.v2.nc": "202fd0c197490ac460af12cd9fa1156aa40023c0023c705f145c596de5b5ad3d",
+        "regional_grid_tpxo10v2.nc": "0789b6a24ecb2ced522481dfcfb7282e32f999984747b9b9f46f044a8898d0ac",
+        "regional_grid_tpxo9v5a.nc": "497a2ae9e6adc7e4b06408dadb57734e2ad24afaa3f0e2e4fd90ebc6eafc2557",
+        "regional_h_tpxo10v2a.nc": "2df2f181f748a960e4072f975226f6f98f6a6c4d5b77da23057946585152d59c",
+        "regional_h_tpxo10v2.nc": "202fd0c197490ac460af12cd9fa1156aa40023c0023c705f145c596de5b5ad3d",
+        "regional_h_tpxo9v5a.nc": "c7e4d9ab73bc11dcb415f88c48131531488e1aed5113df5797e80d3d374607fc",
+        "regional_u_tpxo10v2a.nc": "2d1680ecd53242e858281a762221d91827999967f8e1f3cb7de3d23b47efe8c8",
+        "regional_u_tpxo10v2.nc": "3b0849473cbb7f9076ca907e4fc39eceda3c7d64659c121fa0692024d59dcdb3",
+        "regional_u_tpxo9v5a.nc": "b0cc5f6934d2e212549c7120d458d61a4963ba73d17055e67cc9e4312901b041",
         "CESM_BGC_coarse_global_clim.nc": "20806e4e99285d6de168d3236e2d9245f4e9106474b1464beaa266a73e6ef79f",
         "CESM_BGC_2012.nc": "e374d5df3c1be742d564fd26fd861c2d40af73be50a432c51d258171d5638eb6",
         "CESM_regional_test_data_one_time_slice.nc": "43b578ecc067c85f95d6b97ed7b9dc8da7846f07c95331c6ba7f4a3161036a17",
@@ -127,6 +150,26 @@ def download_correction_data(filename: str) -> str:
     return fname
 
 
+def download_sal_data(filename: str) -> str:
+    """Download the SAL data file.
+
+    Parameters
+    ----------
+    filename : str
+        The name of the test data file to be downloaded. Available options:
+        - "sal_tpxo9.v2a.nc"
+
+    Returns
+    -------
+    str
+        The path to the downloaded test data file.
+    """
+    # Fetch the file using Pooch, downloading if necessary
+    fname = sal_data.fetch(filename)
+
+    return fname
+
+
 def download_test_data(filename: str) -> str:
     """Download the test data file.
 

roms_tools/plot.py

@@ -11,7 +11,9 @@ def _plot(
     title="",
     with_dim_names=False,
     plot_data=True,
+    add_colorbar=True,
     kwargs={},
+    ax=None,
 ):
     """Plots a grid or field on a map with optional depth contours.
 
@@ -30,13 +32,17 @@ def _plot(
     plot_data : bool, optional
         If True, plots the provided field data on the map. If False, only the grid
         boundaries and optional depth contours are plotted. Default is True.
+    add_colorbar : bool, optional
+        If True, add colobar.
     kwargs : dict, optional
         Additional keyword arguments to pass to `pcolormesh` (e.g., colormap or color limits).
+    ax : matplotlib.axes.Axes, optional
+        Pre-existing axes to draw the plot on. If None, a new figure and axes are created.
 
     Returns
     -------
     matplotlib.figure.Figure
-        The generated figure with the plotted data.
+        The generated figure with the plotted data, if a new figure was created.
 
     Raises
     ------
@@ -56,18 +62,27 @@ def _plot(
 
     trans = _get_projection(lon_deg, lat_deg)
 
+    if ax is None:
+        fig, ax = plt.subplots(1, 1, figsize=(13, 7), subplot_kw={"projection": trans})
+
     lon_deg = lon_deg.values
     lat_deg = lat_deg.values
 
-    fig, ax = plt.subplots(1, 1, figsize=(13, 7), subplot_kw={"projection": trans})
-
     if c is not None:
         _add_boundary_to_ax(
             ax, lon_deg, lat_deg, trans, c, with_dim_names=with_dim_names
         )
 
     if plot_data:
-        _add_field_to_ax(ax, lon_deg, lat_deg, field, depth_contours, kwargs=kwargs)
+        _add_field_to_ax(
+            ax,
+            lon_deg,
+            lat_deg,
+            field,
+            depth_contours,
+            add_colorbar=add_colorbar,
+            kwargs=kwargs,
+        )
 
     ax.coastlines(
         resolution="50m", linewidth=0.5, color="black"
@@ -90,7 +105,9 @@ def _plot(
 
     ax.set_title(title)
 
-    return fig
+    # Only return fig if it was created inside the function (i.e., ax was not provided)
+    if ax is None:
+        return fig
 
 
 def _plot_nesting(parent_grid_ds, child_grid_ds, parent_straddle, with_dim_names=False):