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

roms_tools/setup/download.py

@@ -24,6 +24,16 @@ correction_data = pooch.create(
     },
 )
 
+# Create a Pooch object to manage the global topography data
+river_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={
+        "dai_trenberth_may2019.nc": "sha256:793849e6aa60d1f6bdb480c345515fb2453d903c0a30599241b3d752f53715ab",
+    },
+)
 # Create a Pooch object to manage the test data
 pup_test_data = pooch.create(
     # Use the default cache folder for the operating system
@@ -77,6 +87,26 @@ def fetch_topo(topography_source: str) -> xr.Dataset:
     return ds
 
 
+def download_river_data(filename: str) -> xr.Dataset:
+    """Download river data file.
+
+    Parameters
+    ----------
+    filename : str
+        The name of the test data file to be downloaded. Available options:
+        - "dai_trenberth_may2019.nc"
+    Returns
+    -------
+    str
+        The path to the downloaded test data file.
+    """
+
+    # Fetch the file using Pooch, downloading if necessary
+    fname = river_data.fetch(filename)
+
+    return fname
+
+
 def download_correction_data(filename: str) -> str:
     """Download the correction data file.
 

roms_tools/setup/initial_conditions.py

@@ -1,8 +1,7 @@
 import xarray as xr
 import numpy as np
-import yaml
 import importlib.metadata
-from dataclasses import dataclass, field, asdict
+from dataclasses import dataclass, field
 from typing import Dict, Union, List, Optional
 from roms_tools.setup.grid import Grid
 from datetime import datetime
@@ -16,6 +15,8 @@ from roms_tools.setup.utils import (
     rotate_velocities,
     compute_barotropic_velocity,
     transpose_dimensions,
+    _to_yaml,
+    _from_yaml,
 )
 from roms_tools.setup.regrid import LateralRegrid, VerticalRegrid
 from roms_tools.setup.plot import _plot, _section_plot, _profile_plot, _line_plot
@@ -387,17 +388,22 @@ class InitialConditions:
         ds["Cs_w"] = self.grid.ds["Cs_w"]
 
         # Preserve absolute time coordinate for readability
-        ds = ds.assign_coords({"abs_time": ds["time"]})
+        abs_time = ds["time"]
+        attrs = [key for key in abs_time.attrs]
+        for attr in attrs:
+            del abs_time.attrs[attr]
+        abs_time.attrs["long_name"] = "absolute time"
+        ds = ds.assign_coords({"abs_time": abs_time})
 
         # Translate the time coordinate to days since the model reference date
         model_reference_date = np.datetime64(self.model_reference_date)
 
-        # Convert the time coordinate to the format expected by ROMS (days since model reference date)
+        # Convert the time coordinate to the format expected by ROMS (seconds since model reference date)
         ocean_time = (ds["time"] - model_reference_date).astype("float64") * 1e-9
         ds = ds.assign_coords(ocean_time=("time", ocean_time.data.astype("float64")))
         ds["ocean_time"].attrs[
             "long_name"
-        ] = f"seconds since {str(self.model_reference_date)}"
+        ] = f"relative time: seconds since {str(self.model_reference_date)}"
         ds["ocean_time"].attrs["units"] = "seconds"
         ds = ds.swap_dims({"time": "ocean_time"})
         ds = ds.drop_vars("time")
@@ -753,48 +759,8 @@ class InitialConditions:
         filepath : Union[str, Path]
             The path to the YAML file where the parameters will be saved.
         """
-        filepath = Path(filepath)
-
-        # Serialize Grid data
-        grid_data = asdict(self.grid)
-        grid_data.pop("ds", None)  # Exclude non-serializable fields
-        grid_data.pop("straddle", None)
-
-        # 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"
-
-        grid_yaml_data = {"Grid": grid_data}
-
-        initial_conditions_data = {
-            "InitialConditions": {
-                "ini_time": self.ini_time.isoformat(),
-                "source": self.source,
-            }
-        }
-        # Include bgc_source if it's not None
-        if self.bgc_source is not None:
-            initial_conditions_data["InitialConditions"]["bgc_source"] = self.bgc_source
-
-        initial_conditions_data["InitialConditions"][
-            "model_reference_date"
-        ] = self.model_reference_date.isoformat()
-
-        yaml_data = {
-            **grid_yaml_data,
-            **initial_conditions_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)
+        _to_yaml(self, filepath)
 
     @classmethod
     def from_yaml(
@@ -815,35 +781,7 @@ class InitialConditions:
             An instance of the InitialConditions class.
         """
         filepath = Path(filepath)
-        # Read the entire file content
-        with filepath.open("r") as file:
-            file_content = file.read()
-
-        # Split the content into YAML documents
-        documents = list(yaml.safe_load_all(file_content))
-
-        initial_conditions_data = None
-
-        # Process the YAML documents
-        for doc in documents:
-            if doc is None:
-                continue
-            if "InitialConditions" in doc:
-                initial_conditions_data = doc["InitialConditions"]
-                break
-
-        if initial_conditions_data is None:
-            raise ValueError(
-                "No InitialConditions configuration found in the YAML file."
-            )
-
-        # Convert from string to datetime
-        for date_string in ["model_reference_date", "ini_time"]:
-            initial_conditions_data[date_string] = datetime.fromisoformat(
-                initial_conditions_data[date_string]
-            )
 
         grid = Grid.from_yaml(filepath)
-
-        # Create and return an instance of InitialConditions
-        return cls(grid=grid, **initial_conditions_data, use_dask=use_dask)
+        initial_conditions_params = _from_yaml(cls, filepath)
+        return cls(grid=grid, **initial_conditions_params, use_dask=use_dask)

roms_tools/setup/plot.py

@@ -58,16 +58,71 @@ def _plot(
     if straddle:
         lon_deg = xr.where(lon_deg > 180, lon_deg - 360, lon_deg)
 
-    # Define projections
-    proj = ccrs.PlateCarree()
-
-    trans = ccrs.NearsidePerspective(
-        central_longitude=lon_deg.mean().values, central_latitude=lat_deg.mean().values
-    )
+    trans = _get_projection(lon_deg, lat_deg)
 
     lon_deg = lon_deg.values
     lat_deg = lat_deg.values
 
+    fig, ax = plt.subplots(1, 1, figsize=(13, 7), subplot_kw={"projection": trans})
+
+    _add_plot_to_ax(
+        ax, lon_deg, lat_deg, trans, field, depth_contours, c, title, kwargs=kwargs
+    )
+
+
+def _add_plot_to_ax(
+    ax,
+    lon_deg,
+    lat_deg,
+    trans,
+    field=None,
+    depth_contours=False,
+    c="red",
+    title="",
+    add_colorbar=True,
+    kwargs=None,
+):
+    """Plots a grid or field on a map with optional depth contours.
+
+    Parameters
+    ----------
+    ax : matplotlib.axes._axes.Axes
+        The axes on which to plot the data (Cartopy axis with projection).
+
+    lon_deg : np.ndarray
+        Longitude values in degrees.
+
+    lat_deg : np.ndarray
+        Latitude values in degrees.
+
+    trans : cartopy.crs.Projection
+        The projection for transforming coordinates.
+
+    field : xarray.DataArray, optional
+        Field data to plot (e.g., temperature, salinity). If None, only the grid is plotted.
+
+    depth_contours : bool, optional
+        If True, adds depth contours to the plot.
+
+    c : str, optional
+        Color of the grid boundary (default is 'red').
+
+    title : str, optional
+        Title of the plot.
+
+    add_colorbar : bool, optional
+        If True, add colobar.
+
+    kwargs : dict, optional
+        Additional keyword arguments passed to `pcolormesh` (e.g., colormap, limits).
+
+    Notes
+    -----
+    - If `field` is provided, a colorbar is added.
+    - If `depth_contours` is True, the field’s `layer_depth` is used to add contours.
+    """
+    proj = ccrs.PlateCarree()
+
     # find corners
     corners = [
         (lon_deg[0, 0], lat_deg[0, 0]),
@@ -80,14 +135,13 @@ def _plot(
     transformed_corners = [trans.transform_point(lo, la, proj) for lo, la in corners]
     transformed_lons, transformed_lats = zip(*transformed_corners)
 
-    fig, ax = plt.subplots(1, 1, figsize=(13, 7), subplot_kw={"projection": trans})
-
-    ax.plot(
-        list(transformed_lons) + [transformed_lons[0]],
-        list(transformed_lats) + [transformed_lats[0]],
-        "o-",
-        c=c,
-    )
+    if c is not None:
+        ax.plot(
+            list(transformed_lons) + [transformed_lons[0]],
+            list(transformed_lats) + [transformed_lats[0]],
+            "o-",
+            c=c,
+        )
 
     ax.coastlines(
         resolution="50m", linewidth=0.5, color="black"
@@ -104,13 +158,21 @@ def _plot(
             label = f"{field.Long_name} [{field.units}]"
         else:
             label = ""
-        plt.colorbar(p, label=label)
+        if add_colorbar:
+            plt.colorbar(p, label=label)
 
     if depth_contours:
         cs = ax.contour(lon_deg, lat_deg, field.layer_depth, transform=proj, colors="k")
         ax.clabel(cs, inline=True, fontsize=10)
 
 
+def _get_projection(lon, lat):
+
+    return ccrs.NearsidePerspective(
+        central_longitude=lon.mean().values, central_latitude=lat.mean().values
+    )
+
+
 def _section_plot(field, interface_depth=None, title="", kwargs={}, ax=None):
 
     if ax is None: