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

roms_tools/setup/initial_conditions.py

@@ -91,14 +91,13 @@ class InitialConditions:
         self._input_checks()
 
         processed_fields = {}
-        processed_fields = self._process_data(processed_fields, type="physics")
+        processed_fields, variable_info = self._process_data(
+            processed_fields, type="physics"
+        )
 
         if self.bgc_source is not None:
-            processed_fields = self._process_data(processed_fields, type="bgc")
-
-        for var_name in processed_fields.keys():
-            processed_fields[var_name] = transpose_dimensions(
-                processed_fields[var_name]
+            processed_fields, bgc_variable_info = self._process_data(
+                processed_fields, type="bgc"
             )
 
         d_meta = get_variable_metadata()
@@ -106,7 +105,9 @@ class InitialConditions:
 
         ds = self._add_global_metadata(ds)
 
-        self._validate(ds)
+        if self.bgc_source is not None:
+            variable_info = {**variable_info, **bgc_variable_info}
+        self._validate(ds, variable_info)
 
         # substitute NaNs over land by a fill value to avoid blow-up of ROMS
         for var_name in ds.data_vars:
@@ -184,7 +185,12 @@ class InitialConditions:
                     {"time": processed_fields["temp"]["time"]}
                 )
 
-        return processed_fields
+        for var_name in processed_fields.keys():
+            processed_fields[var_name] = transpose_dimensions(
+                processed_fields[var_name]
+            )
+
+        return processed_fields, variable_info
 
     def _input_checks(self):
 
@@ -259,11 +265,21 @@ class InitialConditions:
         - `vector_pair`: For vector variables, this indicates the associated variable that forms the vector (e.g., 'u' and 'v').
         - `is_3d`: Indicates whether the variable is 3D (True for variables like 'temp' and 'salt') or 2D (False for 'zeta').
 
+        Parameters
+        ----------
+        data : object
+            The data object which contains variable names for the "bgc" type variables.
+
+        type : str, optional, default="physics"
+            The type of variable metadata to return. Can be one of:
+            - "physics": for physical variables such as temperature, salinity, and velocity components.
+            - "bgc": for biogeochemical variables (like ALK).
+
         Returns
         -------
         dict
             A dictionary where the keys are variable names and the values are dictionaries of metadata
-            about each variable, including 'location', 'is_vector', 'vector_pair', and 'is_3d'.
+            about each variable, including 'location', 'is_vector', 'vector_pair', 'is_3d', and 'validate'.
         """
         default_info = {
             "location": "rho",
@@ -272,7 +288,6 @@ class InitialConditions:
             "is_3d": True,
         }
 
-        # Define a dictionary for variable names and their associated information
         if type == "physics":
             variable_info = {
                 "zeta": {
@@ -280,38 +295,54 @@ class InitialConditions:
                     "is_vector": False,
                     "vector_pair": None,
                     "is_3d": False,
+                    "validate": True,
                 },
-                "temp": default_info,
-                "salt": default_info,
+                "temp": {**default_info, "validate": False},
+                "salt": {**default_info, "validate": False},
                 "u": {
                     "location": "u",
                     "is_vector": True,
                     "vector_pair": "v",
                     "is_3d": True,
+                    "validate": False,
                 },
                 "v": {
                     "location": "v",
                     "is_vector": True,
                     "vector_pair": "u",
                     "is_3d": True,
+                    "validate": False,
                 },
                 "ubar": {
                     "location": "u",
                     "is_vector": True,
                     "vector_pair": "vbar",
                     "is_3d": False,
+                    "validate": False,
                 },
                 "vbar": {
                     "location": "v",
                     "is_vector": True,
                     "vector_pair": "ubar",
                     "is_3d": False,
+                    "validate": False,
+                },
+                "w": {
+                    "location": "rho",
+                    "is_vector": False,
+                    "vector_pair": None,
+                    "is_3d": True,
+                    "validate": False,
                 },
             }
-        elif type == "bgc":
+
+        if type == "bgc":
             variable_info = {}
             for var_name in data.var_names.keys():
-                variable_info[var_name] = default_info
+                if var_name == "ALK":
+                    variable_info[var_name] = {**default_info, "validate": True}
+                else:
+                    variable_info[var_name] = {**default_info, "validate": False}
 
         return variable_info
 
@@ -373,7 +404,7 @@ class InitialConditions:
 
         return ds
 
-    def _validate(self, ds):
+    def _validate(self, ds, variable_info):
         """Validates the dataset by checking for NaN values in SSH at wet points, which
         would indicate missing raw data coverage over the target domain.
 
@@ -381,6 +412,8 @@ class InitialConditions:
         ----------
         ds : xarray.Dataset
             The dataset to validate.
+        variable_info : dict
+            A dictionary containing metadata about the variables, including whether to validate them.
 
         Raises
         ------
@@ -393,8 +426,11 @@ class InitialConditions:
         This check is only applied to the 2D variable SSH to improve performance.
         """
 
-        ds["zeta"].load()
-        nan_check(ds["zeta"].squeeze(), self.grid.ds.mask_rho)
+        for var_name in variable_info:
+            # Only validate variables based on "validate" flag if use_dask is False
+            if not self.use_dask or variable_info[var_name]["validate"]:
+                ds[var_name].load()
+                nan_check(ds[var_name].squeeze(), self.grid.ds.mask_rho)
 
     def _add_global_metadata(self, ds):
 
@@ -425,6 +461,7 @@ class InitialConditions:
         xi=None,
         depth_contours=False,
         layer_contours=False,
+        ax=None,
     ) -> None:
         """Plot the initial conditions field for a given eta-, xi-, or s_rho- slice.
 
@@ -492,6 +529,8 @@ class InitialConditions:
             be added to the plot. This is particularly useful in vertical sections to
             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 horizontal plots that display the eta- and xi-dimensions at the same time.
 
         Returns
         -------
@@ -518,7 +557,12 @@ class InitialConditions:
         ):
             raise ValueError("For 2D fields, specify either eta or xi, not both.")
 
-        self.ds[var_name].load()
+        if self.use_dask:
+            from dask.diagnostics import ProgressBar
+
+            with ProgressBar():
+                self.ds[var_name].load()
+
         field = self.ds[var_name].squeeze()
 
         if all(dim in field.dims for dim in ["eta_rho", "xi_rho"]):
@@ -634,13 +678,17 @@ class InitialConditions:
 
             if len(field.dims) == 2:
                 _section_plot(
-                    field, interface_depth=interface_depth, title=title, kwargs=kwargs
+                    field,
+                    interface_depth=interface_depth,
+                    title=title,
+                    kwargs=kwargs,
+                    ax=ax,
                 )
             else:
                 if "s_rho" in field.dims:
-                    _profile_plot(field, title=title)
+                    _profile_plot(field, title=title, ax=ax)
                 else:
-                    _line_plot(field, title=title)
+                    _line_plot(field, title=title, ax=ax)
 
     def save(
         self, filepath: Union[str, Path], np_eta: int = None, np_xi: int = None
@@ -681,7 +729,13 @@ class InitialConditions:
         if filepath.suffix == ".nc":
             filepath = filepath.with_suffix("")
 
-        dataset_list = [self.ds.load()]
+        if self.use_dask:
+            from dask.diagnostics import ProgressBar
+
+            with ProgressBar():
+                self.ds.load()
+
+        dataset_list = [self.ds]
         output_filenames = [str(filepath)]
 
         saved_filenames = save_datasets(
@@ -719,15 +773,18 @@ class InitialConditions:
 
         initial_conditions_data = {
             "InitialConditions": {
-                "source": self.source,
                 "ini_time": self.ini_time.isoformat(),
-                "model_reference_date": self.model_reference_date.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,
@@ -737,7 +794,7 @@ class InitialConditions:
             # Write header
             file.write(header)
             # Write YAML data
-            yaml.dump(yaml_data, file, default_flow_style=False)
+            yaml.dump(yaml_data, file, default_flow_style=False, sort_keys=False)
 
     @classmethod
     def from_yaml(

roms_tools/setup/plot.py

@@ -12,6 +12,32 @@ def _plot(
     title="",
     kwargs={},
 ):
+    """Plots a grid or field on a map with optional depth contours.
+
+    This function plots a map using Cartopy projections. It supports plotting a grid, a field, and adding depth contours if desired.
+    The projection can be customized, and the grid can be adjusted for domains straddling the 180° meridian.
+
+    Parameters
+    ----------
+    grid_ds : xarray.Dataset
+        The grid dataset containing coordinates (`lon_rho`, `lat_rho`).
+    field : xarray.DataArray, optional
+        The field to plot. If None, only the grid is plotted.
+    depth_contours : bool, optional
+        If True, adds depth contours to the plot.
+    straddle : bool, optional
+        If True, adjusts longitude values to straddle across the 180° meridian.
+    c : str, optional
+        Color for the boundary plot (default is 'red').
+    title : str, optional
+        Title of the plot.
+    kwargs : dict, optional
+        Additional keyword arguments to pass to `pcolormesh` (e.g., colormap or color limits).
+
+    Notes
+    -----
+    The function raises a `NotImplementedError` if the domain contains the North or South Pole.
+    """
 
     if field is None:
         lon_deg = grid_ds["lon_rho"]
@@ -84,12 +110,11 @@ def _plot(
         cs = ax.contour(lon_deg, lat_deg, field.layer_depth, transform=proj, colors="k")
         ax.clabel(cs, inline=True, fontsize=10)
 
-    return fig
 
+def _section_plot(field, interface_depth=None, title="", kwargs={}, ax=None):
 
-def _section_plot(field, interface_depth=None, title="", kwargs={}):
-
-    fig, ax = plt.subplots(1, 1, figsize=(9, 5))
+    if ax is None:
+        fig, ax = plt.subplots(1, 1, figsize=(9, 5))
 
     dims_to_check = ["eta_rho", "eta_u", "eta_v", "xi_rho", "xi_u", "xi_v"]
     try:
@@ -132,7 +157,23 @@ def _section_plot(field, interface_depth=None, title="", kwargs={}):
     ax.set_title(title)
 
 
-def _profile_plot(field, title=""):
+def _profile_plot(field, title="", ax=None):
+    """Plots a profile of the given field against depth.
+
+    Parameters
+    ----------
+    field : xarray.DataArray
+        Data to plot.
+    title : str, optional
+        Title of the plot.
+    ax : matplotlib.axes.Axes, optional
+        Axes to plot on. If None, a new figure is created.
+
+    Raises
+    ------
+    ValueError
+        If no expected depth coordinate is found in the field.
+    """
 
     depths_to_check = [
         "layer_depth_rho",
@@ -153,16 +194,33 @@ def _profile_plot(field, title=""):
             "None of the expected coordinates (layer_depth_rho, layer_depth_u, layer_depth_v, interface_depth_rho, interface_depth_u, interface_depth_v) found in field.coords"
         )
 
-    fig, ax = plt.subplots(1, 1, figsize=(4, 7))
+    if ax is None:
+        fig, ax = plt.subplots(1, 1, figsize=(4, 7))
     kwargs = {"y": depth_label, "yincrease": False}
     field.plot(**kwargs)
     ax.set_title(title)
     ax.grid()
 
 
-def _line_plot(field, title=""):
-
-    fig, ax = plt.subplots(1, 1, figsize=(7, 4))
-    field.plot()
+def _line_plot(field, title="", ax=None):
+    """Plots a line graph of the given field.
+
+    Parameters
+    ----------
+    field : xarray.DataArray
+        Data to plot.
+    title : str, optional
+        Title of the plot.
+    ax : matplotlib.axes.Axes, optional
+        Axes to plot on. If None, a new figure is created.
+
+    Returns
+    -------
+    None
+        Modifies the plot in-place.
+    """
+    if ax is None:
+        fig, ax = plt.subplots(1, 1, figsize=(7, 4))
+    field.plot(ax=ax)
     ax.set_title(title)
     ax.grid()

roms_tools/setup/surface_forcing.py

@@ -133,14 +133,7 @@ class SurfaceForcing:
 
         ds = self._write_into_dataset(processed_fields, data, d_meta)
 
-        if self.use_coarse_grid:
-            mask = self.grid.ds["mask_coarse"].rename(
-                {"eta_coarse": "eta_rho", "xi_coarse": "xi_rho"}
-            )
-        else:
-            mask = self.grid.ds["mask_rho"]
-
-        self._validate(ds, mask)
+        self._validate(ds, target_coords["mask"], variable_info)
 
         # substitute NaNs over land by a fill value to avoid blow-up of ROMS
         for var_name in ds.data_vars:
@@ -233,28 +226,34 @@ class SurfaceForcing:
         # Define a dictionary for variable names and their associated information
         if self.type == "physics":
             variable_info = {
-                "swrad": default_info,
-                "lwrad": default_info,
-                "Tair": default_info,
-                "qair": default_info,
-                "rain": default_info,
+                "swrad": {**default_info, "validate": True},
+                "lwrad": {**default_info, "validate": False},
+                "Tair": {**default_info, "validate": False},
+                "qair": {**default_info, "validate": True},
+                "rain": {**default_info, "validate": False},
                 "uwnd": {
                     "location": "u",
                     "is_vector": True,
                     "vector_pair": "vwnd",
                     "is_3d": False,
+                    "validate": True,
                 },
                 "vwnd": {
                     "location": "v",
                     "is_vector": True,
                     "vector_pair": "uwnd",
                     "is_3d": False,
+                    "validate": True,
                 },
             }
         elif self.type == "bgc":
             variable_info = {}
-            for var in data.var_names.keys():
-                variable_info[var] = default_info
+            for var_name in data.var_names.keys():
+                variable_info[var_name] = default_info
+                if var_name == "pco2_air":
+                    variable_info[var_name] = {**default_info, "validate": True}
+                else:
+                    variable_info[var_name] = {**default_info, "validate": False}
 
         return variable_info
 
@@ -300,9 +299,6 @@ class SurfaceForcing:
             ds[var_name].attrs["long_name"] = d_meta[var_name]["long_name"]
             ds[var_name].attrs["units"] = d_meta[var_name]["units"]
 
-        if self.use_coarse_grid:
-            ds = ds.rename({"eta_coarse": "eta_rho", "xi_coarse": "xi_rho"})
-
         ds = self._add_global_metadata(ds)
 
         # Convert the time coordinate to the format expected by ROMS
@@ -369,7 +365,7 @@ class SurfaceForcing:
 
         return ds
 
-    def _validate(self, ds, mask):
+    def _validate(self, ds, mask, variable_info):
         """Validates the dataset by checking for NaN values at wet points, which would
         indicate missing raw data coverage over the target domain.
 
@@ -379,6 +375,10 @@ class SurfaceForcing:
             The dataset to validate.
         mask : xarray.DataArray
             Land mask (1=ocean, 0=land) to determine wet points in the domain.
+        variable_info : dict
+            A dictionary containing metadata about each variable (e.g., location,
+            whether it's a 3D variable, etc.). Used to retrieve information for
+            validating each variable.
 
         Raises
         ------
@@ -392,7 +392,9 @@ class SurfaceForcing:
         """
 
         for var_name in ds.data_vars:
-            nan_check(ds[var_name].isel(time=0), mask)
+            # Only validate variables based on "validate" flag if use_dask is False
+            if not self.use_dask or variable_info[var_name]["validate"]:
+                nan_check(ds[var_name].isel(time=0), mask)
 
     def _add_global_metadata(self, ds=None):
 
@@ -462,15 +464,16 @@ class SurfaceForcing:
         if var_name not in self.ds:
             raise ValueError(f"Variable '{var_name}' is not found in dataset.")
 
-        field = self.ds[var_name].isel(time=time).load()
+        field = self.ds[var_name].isel(time=time)
+        if self.use_dask:
+            from dask.diagnostics import ProgressBar
+
+            with ProgressBar():
+                field = field.load()
+
         title = field.long_name
 
-        # assign lat / lon
-        if self.use_coarse_grid:
-            field = field.rename({"eta_rho": "eta_coarse", "xi_rho": "xi_coarse"})
-            field = field.where(self.grid.ds.mask_coarse)
-        else:
-            field = field.where(self.grid.ds.mask_rho)
+        field = field.where(self.target_coords["mask"])
 
         field = field.assign_coords(
             {"lon": self.target_coords["lon"], "lat": self.target_coords["lat"]}
@@ -502,24 +505,26 @@ class SurfaceForcing:
         )
 
     def save(
-        self, filepath: Union[str, Path], np_eta: int = None, np_xi: int = None
+        self,
+        filepath: Union[str, Path],
+        np_eta: int = None,
+        np_xi: int = None,
+        group: bool = False,
     ) -> None:
-        """Save the surface forcing fields to netCDF4 files.
+        """Save the surface forcing fields to one or more netCDF4 files.
 
-        This method saves the dataset by grouping it into subsets based on the data frequency. The subsets are then written
-        to one or more netCDF4 files. The filenames of the output files reflect the temporal coverage of the data.
+        This method saves the dataset either as a single file or as multiple files depending on the partitioning and grouping options.
+        The dataset can be saved in two modes:
 
-        There are two modes of saving the dataset:
+        1. **Single File Mode (default)**:
+            - If both `np_eta` and `np_xi` are `None`, the entire dataset is saved as a single netCDF4 file.
+            - The file is named based on the `filepath`, with `.nc` automatically appended.
 
-          1. **Single File Mode (default)**:
+        2. **Partitioned Mode**:
+            - If either `np_eta` or `np_xi` is specified, the dataset is partitioned into spatial tiles along the `eta` and `xi` axes.
+            - Each tile is saved as a separate netCDF4 file, and filenames are modified with an index (e.g., `"filepath_YYYYMM.0.nc"`, `"filepath_YYYYMM.1.nc"`).
 
-            If both `np_eta` and `np_xi` are `None`, the entire dataset, divided by temporal subsets, is saved as a single netCDF4 file
-            with the base filename specified by `filepath.nc`.
-
-          2. **Partitioned Mode**:
-
-            - If either `np_eta` or `np_xi` is specified, the dataset is divided into spatial tiles along the eta-axis and xi-axis.
-            - Each spatial tile is saved as a separate netCDF4 file.
+        Additionally, if `group` is set to `True`, the dataset is first grouped into temporal subsets, resulting in multiple grouped files before partitioning and saving.
 
         Parameters
         ----------
@@ -531,6 +536,8 @@ class SurfaceForcing:
             The number of partitions along the `eta` direction. If `None`, no spatial partitioning is performed.
         np_xi : int, optional
             The number of partitions along the `xi` direction. If `None`, no spatial partitioning is performed.
+        group: bool, optional
+            If `True`, groups the dataset into multiple files based on temporal data frequency. Defaults to `False`.
 
         Returns
         -------
@@ -545,7 +552,18 @@ class SurfaceForcing:
         if filepath.suffix == ".nc":
             filepath = filepath.with_suffix("")
 
-        dataset_list, output_filenames = group_dataset(self.ds.load(), str(filepath))
+        if self.use_dask:
+            from dask.diagnostics import ProgressBar
+
+            with ProgressBar():
+                self.ds.load()
+
+        if group:
+            dataset_list, output_filenames = group_dataset(self.ds, str(filepath))
+        else:
+            dataset_list = [self.ds]
+            output_filenames = [str(filepath)]
+
         saved_filenames = save_datasets(
             dataset_list, output_filenames, np_eta=np_eta, np_xi=np_xi
         )
@@ -603,7 +621,7 @@ class SurfaceForcing:
             # Write header
             file.write(header)
             # Write YAML data
-            yaml.dump(yaml_data, file, default_flow_style=False)
+            yaml.dump(yaml_data, file, default_flow_style=False, sort_keys=False)
 
     @classmethod
     def from_yaml(