roms-tools 1.2.0__py3-none-any.whl → 1.4.0__py3-none-any.whl

ci/environment.yml

@@ -3,20 +3,11 @@ channels:
   - conda-forge
   - nodefaults
 dependencies:
-  - python>=3.9
-  - numpy
-  - scipy
-  - matplotlib
-  - cartopy
-  - xarray
-  - xarray-datatree
-  - netcdf4
-  - pooch
-  - packaging
-  - numba
-  - gcm_filters
+  - python>=3.10
   # testing
+  - zarr
   - pytest
+  - pytest-xdist
   - flake8
   - black
   - pre-commit

roms_tools/_version.py

@@ -1,2 +1,2 @@
 # Do not change! Do not track in version control!
-__version__ = "1.2.0"
+__version__ = "1.4.0"

roms_tools/setup/boundary_forcing.py

@@ -2,9 +2,8 @@ import xarray as xr
 import numpy as np
 import pandas as pd
 import yaml
-from datatree import DataTree
 import importlib.metadata
-from typing import Dict, Union, Optional
+from typing import Dict, Union, List
 from dataclasses import dataclass, field, asdict
 from roms_tools.setup.grid import Grid
 from roms_tools.setup.mixins import ROMSToolsMixins
@@ -12,17 +11,21 @@ from datetime import datetime
 from roms_tools.setup.datasets import GLORYSDataset, CESMBGCDataset
 from roms_tools.setup.utils import (
     nan_check,
+    substitute_nans_by_fillvalue,
+    get_variable_metadata,
+    get_boundary_info,
+    group_dataset,
+    save_datasets,
 )
 from roms_tools.setup.plot import _section_plot, _line_plot
-import calendar
-import dask
 import matplotlib.pyplot as plt
+from pathlib import Path
 
 
 @dataclass(frozen=True, kw_only=True)
 class BoundaryForcing(ROMSToolsMixins):
     """
-    Represents boundary forcing for ROMS.
+    Represents boundary forcing input data for ROMS.
 
     Parameters
     ----------
@@ -34,23 +37,21 @@ class BoundaryForcing(ROMSToolsMixins):
         End time of the desired boundary forcing data.
     boundaries : Dict[str, bool], optional
         Dictionary specifying which boundaries are forced (south, east, north, west). Default is all True.
-    physics_source : Dict[str, Union[str, None]]
-        Dictionary specifying the source of the physical boundary forcing data:
+    source : Dict[str, Union[str, Path, List[Union[str, Path]]], bool]
+        Dictionary specifying the source of the boundary forcing data:
         - "name" (str): Name of the data source (e.g., "GLORYS").
-        - "path" (str): Path to the physical data file. Can contain wildcards.
-        - "climatology" (bool): Indicates if the physical data is climatology data. Defaults to False.
-    bgc_source : Optional[Dict[str, Union[str, None]]]
-        Dictionary specifying the source of the biogeochemical (BGC) initial condition data:
-        - "name" (str): Name of the BGC data source (e.g., "CESM_REGRIDDED").
-        - "path" (str): Path to the BGC data file. Can contain wildcards.
-        - "climatology" (bool): Indicates if the BGC data is climatology data. Defaults to True.
+        - "path" (Union[str, Path, List[Union[str, Path]]]): The path to the raw data file(s). Can be a single string (with or without wildcards),
+          a single Path object, or a list of strings or Path objects containing multiple files.
+        - "climatology" (bool): Indicates if the data is climatology data. Defaults to False.
     model_reference_date : datetime, optional
         Reference date for the model. Default is January 1, 2000.
+    use_dask: bool, optional
+        Indicates whether to use dask for processing. If True, data is processed with dask; if False, data is processed eagerly. Defaults to False.
 
     Attributes
     ----------
     ds : xr.Dataset
-        Xarray Dataset containing the atmospheric forcing data.
+        Xarray Dataset containing the boundary forcing data.
 
     Examples
     --------
@@ -59,12 +60,8 @@ class BoundaryForcing(ROMSToolsMixins):
     ...     boundaries={"south": True, "east": True, "north": False, "west": True},
     ...     start_time=datetime(2022, 1, 1),
     ...     end_time=datetime(2022, 1, 2),
-    ...     physics_source={"name": "GLORYS", "path": "physics_data.nc"},
-    ...     bgc_source={
-    ...         "name": "CESM_REGRIDDED",
-    ...         "path": "bgc_data.nc",
-    ...         "climatology": True,
-    ...     },
+    ...     source={"name": "GLORYS", "path": "glorys_data.nc"},
+    ...     type="physics",
     ... )
     """
 
@@ -79,9 +76,10 @@ class BoundaryForcing(ROMSToolsMixins):
             "west": True,
         }
     )
-    physics_source: Dict[str, Union[str, None]]
-    bgc_source: Optional[Dict[str, Union[str, None]]] = None
+    source: Dict[str, Union[str, Path, List[Union[str, Path]]]]
+    type: str = "physics"
     model_reference_date: datetime = datetime(2000, 1, 1)
+    use_dask: bool = False
 
     ds: xr.Dataset = field(init=False, repr=False)
 
@@ -98,107 +96,89 @@ class BoundaryForcing(ROMSToolsMixins):
             straddle=straddle,
         )
 
-        vars_2d = ["zeta"]
-        vars_3d = ["temp", "salt", "u", "v"]
-        data_vars = super().regrid_data(data, vars_2d, vars_3d, lon, lat)
-        data_vars = super().process_velocities(data_vars, angle, "u", "v")
-        object.__setattr__(data, "data_vars", data_vars)
+        if self.type == "physics":
+            vars_2d = ["zeta"]
+            vars_3d = ["temp", "salt", "u", "v"]
+        elif self.type == "bgc":
+            vars_2d = []
+            vars_3d = data.var_names.keys()
 
-        if self.bgc_source is not None:
-            bgc_data = self._get_bgc_data()
-            bgc_data.choose_subdomain(
-                latitude_range=[lat.min().values, lat.max().values],
-                longitude_range=[lon.min().values, lon.max().values],
-                margin=2,
-                straddle=straddle,
-            )
+        data_vars = super().regrid_data(data, vars_2d, vars_3d, lon, lat)
 
-            vars_2d = []
-            vars_3d = bgc_data.var_names.keys()
-            data_vars = super().regrid_data(bgc_data, vars_2d, vars_3d, lon, lat)
-            object.__setattr__(bgc_data, "data_vars", data_vars)
-        else:
-            bgc_data = None
+        if self.type == "physics":
+            data_vars = super().process_velocities(data_vars, angle, "u", "v")
+        object.__setattr__(data, "data_vars", data_vars)
 
-        d_meta = super().get_variable_metadata()
-        bdry_coords = super().get_boundary_info()
+        d_meta = get_variable_metadata()
+        bdry_coords = get_boundary_info()
 
-        ds = self._write_into_datatree(data, bgc_data, d_meta, bdry_coords)
+        ds = self._write_into_dataset(data, d_meta, bdry_coords)
 
+        # NaN values at wet points indicate that the raw data did not cover the domain, and the following will raise a ValueError
+        # this check works only for 2D fields because for 3D I extrapolate to bottom which eliminates NaNs
         for direction in ["south", "east", "north", "west"]:
             if self.boundaries[direction]:
-                nan_check(
-                    ds["physics"][f"zeta_{direction}"].isel(bry_time=0),
-                    self.grid.ds.mask_rho.isel(**bdry_coords["rho"][direction]),
-                )
+                if type == "physics":
+                    nan_check(
+                        ds[f"zeta_{direction}"].isel(bry_time=0),
+                        self.grid.ds.mask_rho.isel(**bdry_coords["rho"][direction]),
+                    )
+                elif type == "bgc":
+                    nan_check(
+                        ds[f"ALK_{direction}"].isel(bry_time=0, s_rho=-1),
+                        self.grid.ds.mask_rho.isel(**bdry_coords["rho"][direction]),
+                    )
+
+        # substitute NaNs over land by a fill value to avoid blow-up of ROMS
+        for var in ds.data_vars:
+            ds[var] = substitute_nans_by_fillvalue(ds[var])
 
         object.__setattr__(self, "ds", ds)
 
     def _input_checks(self):
+        # Validate the 'type' parameter
+        if self.type not in ["physics", "bgc"]:
+            raise ValueError("`type` must be either 'physics' or 'bgc'.")
+
+        # Ensure 'source' dictionary contains required keys
+        if "name" not in self.source:
+            raise ValueError("`source` must include a 'name'.")
+        if "path" not in self.source:
+            raise ValueError("`source` must include a 'path'.")
 
-        if "name" not in self.physics_source.keys():
-            raise ValueError("`physics_source` must include a 'name'.")
-        if "path" not in self.physics_source.keys():
-            raise ValueError("`physics_source` must include a 'path'.")
-        # set self.physics_source["climatology"] to False if not provided
+        # Set 'climatology' to False if not provided in 'source'
         object.__setattr__(
             self,
-            "physics_source",
-            {
-                **self.physics_source,
-                "climatology": self.physics_source.get("climatology", False),
-            },
+            "source",
+            {**self.source, "climatology": self.source.get("climatology", False)},
         )
-        if self.bgc_source is not None:
-            if "name" not in self.bgc_source.keys():
-                raise ValueError(
-                    "`bgc_source` must include a 'name' if it is provided."
-                )
-            if "path" not in self.bgc_source.keys():
-                raise ValueError(
-                    "`bgc_source` must include a 'path' if it is provided."
-                )
-            # set self.bgc_source["climatology"] to True if not provided
-            object.__setattr__(
-                self,
-                "bgc_source",
-                {
-                    **self.bgc_source,
-                    "climatology": self.bgc_source.get("climatology", True),
-                },
-            )
 
     def _get_data(self):
 
-        if self.physics_source["name"] == "GLORYS":
-            data = GLORYSDataset(
-                filename=self.physics_source["path"],
-                start_time=self.start_time,
-                end_time=self.end_time,
-                climatology=self.physics_source["climatology"],
-            )
-        else:
-            raise ValueError(
-                'Only "GLORYS" is a valid option for physics_source["name"].'
-            )
-
-        return data
+        data_dict = {
+            "filename": self.source["path"],
+            "start_time": self.start_time,
+            "end_time": self.end_time,
+            "climatology": self.source["climatology"],
+            "use_dask": self.use_dask,
+        }
 
-    def _get_bgc_data(self):
+        if self.type == "physics":
+            if self.source["name"] == "GLORYS":
+                data = GLORYSDataset(**data_dict)
+            else:
+                raise ValueError(
+                    'Only "GLORYS" is a valid option for source["name"] when type is "physics".'
+                )
 
-        if self.bgc_source["name"] == "CESM_REGRIDDED":
+        elif self.type == "bgc":
+            if self.source["name"] == "CESM_REGRIDDED":
 
-            data = CESMBGCDataset(
-                filename=self.bgc_source["path"],
-                start_time=self.start_time,
-                end_time=self.end_time,
-                climatology=self.bgc_source["climatology"],
-            )
-            data.post_process()
-        else:
-            raise ValueError(
-                'Only "CESM_REGRIDDED" is a valid option for bgc_source["name"].'
-            )
+                data = CESMBGCDataset(**data_dict)
+            else:
+                raise ValueError(
+                    'Only "CESM_REGRIDDED" is a valid option for source["name"] when type is "bgc".'
+                )
 
         return data
 
@@ -253,60 +233,58 @@ class BoundaryForcing(ROMSToolsMixins):
         existing_vars = [var for var in variables_to_drop if var in ds]
         ds = ds.drop_vars(existing_vars)
 
-        # Preserve absolute time coordinate for readability
-        ds = ds.assign_coords({"abs_time": ds["time"]})
+        ds = self._add_global_metadata(ds)
 
         # Convert the time coordinate to the format expected by ROMS
         if data.climatology:
+            ds.attrs["climatology"] = str(True)
+            # Preserve absolute time coordinate for readability
+            ds = ds.assign_coords(
+                {"abs_time": np.datetime64(self.model_reference_date) + ds["time"]}
+            )
             # Convert to pandas TimedeltaIndex
             timedelta_index = pd.to_timedelta(ds["time"].values)
+
             # Determine the start of the year for the base_datetime
             start_of_year = datetime(self.model_reference_date.year, 1, 1)
+
             # Calculate the offset from midnight of the new year
             offset = self.model_reference_date - start_of_year
+
+            # Convert the timedelta to nanoseconds first, then to days
             bry_time = xr.DataArray(
-                timedelta_index - offset,
+                (timedelta_index - offset).view("int64") / 3600 / 24 * 1e-9,
                 dims="time",
             )
+
         else:
+            # Preserve absolute time coordinate for readability
+            ds = ds.assign_coords({"abs_time": ds["time"]})
             # TODO: Check if we need to convert from 12:00:00 to 00:00:00 as in matlab scripts
-            bry_time = ds["time"] - np.datetime64(self.model_reference_date)
+            bry_time = (
+                (ds["time"] - np.datetime64(self.model_reference_date)).astype(
+                    "float64"
+                )
+                / 3600
+                / 24
+                * 1e-9
+            )
 
         ds = ds.assign_coords({"bry_time": bry_time})
         ds["bry_time"].attrs[
             "long_name"
-        ] = f"nanoseconds since {np.datetime_as_string(np.datetime64(self.model_reference_date), unit='ns')}"
-        ds["bry_time"].encoding["units"] = "nanoseconds"
+        ] = f"days since {str(self.model_reference_date)}"
+        ds["bry_time"].encoding["units"] = "days"
+        ds["bry_time"].attrs["units"] = "days"
         ds = ds.swap_dims({"time": "bry_time"})
         ds = ds.drop_vars("time")
+        ds.encoding["unlimited_dims"] = "bry_time"
 
         if data.climatology:
             ds["bry_time"].attrs["cycle_length"] = 365.25
 
         return ds
 
-    def _write_into_datatree(self, data, bgc_data, d_meta, bdry_coords):
-
-        ds = self._add_global_metadata()
-        ds["sc_r"] = self.grid.ds["sc_r"]
-        ds["Cs_r"] = self.grid.ds["Cs_r"]
-
-        ds = DataTree(name="root", data=ds)
-
-        ds_physics = self._write_into_dataset(data, d_meta, bdry_coords)
-        ds_physics = self._add_global_metadata(ds_physics)
-        ds_physics.attrs["physics_source"] = self.physics_source["name"]
-
-        ds_physics = DataTree(name="physics", parent=ds, data=ds_physics)
-
-        if bgc_data:
-            ds_bgc = self._write_into_dataset(bgc_data, d_meta, bdry_coords)
-            ds_bgc = self._add_global_metadata(ds_bgc)
-            ds_bgc.attrs["bgc_source"] = self.bgc_source["name"]
-            ds_bgc = DataTree(name="bgc", parent=ds, data=ds_bgc)
-
-        return ds
-
     def _get_coordinates(self, direction, point):
         """
         Retrieve layer and interface depth coordinates for a specified grid boundary.
@@ -332,7 +310,7 @@ class BoundaryForcing(ROMSToolsMixins):
             boundary and point type.
         """
 
-        bdry_coords = super().get_boundary_info()
+        bdry_coords = get_boundary_info()
 
         layer_depth = self.grid.ds[f"layer_depth_{point}"].isel(
             **bdry_coords[point][direction]
@@ -356,6 +334,7 @@ class BoundaryForcing(ROMSToolsMixins):
         ds.attrs["roms_tools_version"] = roms_tools_version
         ds.attrs["start_time"] = str(self.start_time)
         ds.attrs["end_time"] = str(self.end_time)
+        ds.attrs["source"] = self.source["name"]
         ds.attrs["model_reference_date"] = str(self.model_reference_date)
 
         ds.attrs["theta_s"] = self.grid.ds.attrs["theta_s"]
@@ -376,7 +355,7 @@ class BoundaryForcing(ROMSToolsMixins):
         Parameters
         ----------
         varname : str
-            The name of the initial conditions field to plot. Options include:
+            The name of the boundary forcing field to plot. Options include:
             - "temp_{direction}": Potential temperature, where {direction} can be one of ["south", "east", "north", "west"].
             - "salt_{direction}": Salinity, where {direction} can be one of ["south", "east", "north", "west"].
             - "zeta_{direction}": Sea surface height, where {direction} can be one of ["south", "east", "north", "west"].
@@ -434,17 +413,10 @@ class BoundaryForcing(ROMSToolsMixins):
             If the specified varname is not one of the valid options.
         """
 
-        if varname in self.ds["physics"]:
-            ds = self.ds["physics"]
-        else:
-            if "bgc" in self.ds and varname in self.ds["bgc"]:
-                ds = self.ds["bgc"]
-            else:
-                raise ValueError(
-                    f"Variable '{varname}' is not found in 'physics' or 'bgc' datasets."
-                )
+        if varname not in self.ds:
+            raise ValueError(f"Variable '{varname}' is not found in dataset.")
 
-        field = ds[varname].isel(bry_time=time).load()
+        field = self.ds[varname].isel(bry_time=time).load()
         title = field.long_name
 
         if "s_rho" in field.dims:
@@ -493,92 +465,67 @@ class BoundaryForcing(ROMSToolsMixins):
         else:
             _line_plot(field, title=title)
 
-    def save(self, filepath: str, time_chunk_size: int = 1) -> None:
+    def save(
+        self, filepath: Union[str, Path], np_eta: int = None, np_xi: int = None
+    ) -> None:
         """
-        Save the interpolated boundary forcing fields to netCDF4 files.
+        Save the boundary forcing fields to 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.
+
+        There are two modes of saving the dataset:
 
-        This method groups the dataset by year and month, chunks the data by the specified
-        time chunk size, and saves each chunked subset to a separate netCDF4 file named
-        according to the year, month, and day range if not a complete month of data is included.
+        1. **Single File Mode (default)**:
+           - 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.
 
         Parameters
         ----------
-        filepath : str
-            The base path and filename for the output files. The files will be named with
-            the format "filepath.YYYYMM.nc" if a full month of data is included, or
-            "filepath.YYYYMMDD-DD.nc" otherwise.
-        time_chunk_size : int, optional
-            Number of time slices to include in each chunk along the time dimension. Default is 1,
-            meaning each chunk contains one time slice.
+        filepath : Union[str, Path]
+            The base path and filename for the output files. The format of the filenames depends on whether partitioning is used
+            and the temporal range of the data. For partitioned datasets, files will be named with an additional index, e.g.,
+            `"filepath_YYYYMM.0.nc"`, `"filepath_YYYYMM.1.nc"`, etc.
+        np_eta : int, optional
+            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.
 
         Returns
         -------
-        None
+        List[Path]
+            A list of Path objects for the filenames that were saved.
         """
-        datasets = []
-        filenames = []
-        writes = []
-
-        for node in ["physics", "bgc"]:
-            if node in self.ds:
-                ds = self.ds[node].to_dataset()
-                # copy vertical coordinate variables from parent to children because I believe this is info that ROMS needs
-                for var in self.ds.data_vars:
-                    ds[var] = self.ds[var]
-                if hasattr(ds["bry_time"], "cycle_length"):
-                    filename = f"{filepath}_{node}_clim.nc"
-                    filenames.append(filename)
-                    datasets.append(ds)
-                else:
-                    # Group dataset by year
-                    gb = ds.groupby("abs_time.year")
-
-                    for year, group_ds in gb:
-                        # Further group each yearly group by month
-                        sub_gb = group_ds.groupby("abs_time.month")
-
-                        for month, ds in sub_gb:
-                            # Chunk the dataset by the specified time chunk size
-                            ds = ds.chunk({"bry_time": time_chunk_size})
-                            datasets.append(ds)
-
-                            # Determine the number of days in the month
-                            num_days_in_month = calendar.monthrange(year, month)[1]
-                            first_day = ds.abs_time.dt.day.values[0]
-                            last_day = ds.abs_time.dt.day.values[-1]
-
-                            # Create filename based on whether the dataset contains a full month
-                            if first_day == 1 and last_day == num_days_in_month:
-                                # Full month format: "filepath_physics_YYYYMM.nc"
-                                year_month_str = f"{year}{month:02}"
-                                filename = f"{filepath}_{node}_{year_month_str}.nc"
-                            else:
-                                # Partial month format: "filepath_physics_YYYYMMDD-DD.nc"
-                                year_month_day_str = (
-                                    f"{year}{month:02}{first_day:02}-{last_day:02}"
-                                )
-                                filename = f"{filepath}_{node}_{year_month_day_str}.nc"
-                            filenames.append(filename)
-
-        print("Saving the following files:")
-        for ds, filename in zip(datasets, filenames):
-            print(filename)
-            # Prepare the dataset for writing to a netCDF file without immediately computing
-            write = ds.to_netcdf(filename, compute=False)
-            writes.append(write)
-
-        # Perform the actual write operations in parallel
-        dask.compute(*writes)
-
-    def to_yaml(self, filepath: str) -> None:
+
+        # Ensure filepath is a Path object
+        filepath = Path(filepath)
+
+        # Remove ".nc" suffix if present
+        if filepath.suffix == ".nc":
+            filepath = filepath.with_suffix("")
+
+        dataset_list, output_filenames = group_dataset(self.ds.load(), str(filepath))
+        saved_filenames = save_datasets(
+            dataset_list, output_filenames, np_eta=np_eta, np_xi=np_xi
+        )
+
+        return saved_filenames
+
+    def to_yaml(self, filepath: Union[str, Path]) -> None:
         """
         Export the parameters of the class to a YAML file, including the version of roms-tools.
 
         Parameters
         ----------
-        filepath : str
+        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
@@ -600,8 +547,8 @@ class BoundaryForcing(ROMSToolsMixins):
                 "start_time": self.start_time.isoformat(),
                 "end_time": self.end_time.isoformat(),
                 "boundaries": self.boundaries,
-                "physics_source": self.physics_source,
-                "bgc_source": self.bgc_source,
+                "source": self.source,
+                "type": self.type,
                 "model_reference_date": self.model_reference_date.isoformat(),
             }
         }
@@ -611,29 +558,34 @@ class BoundaryForcing(ROMSToolsMixins):
             **boundary_forcing_data,
         }
 
-        with open(filepath, "w") as file:
+        with filepath.open("w") as file:
             # Write header
             file.write(header)
             # Write YAML data
             yaml.dump(yaml_data, file, default_flow_style=False)
 
     @classmethod
-    def from_yaml(cls, filepath: str) -> "BoundaryForcing":
+    def from_yaml(
+        cls, filepath: Union[str, Path], use_dask: bool = False
+    ) -> "BoundaryForcing":
         """
         Create an instance of the BoundaryForcing class from a YAML file.
 
         Parameters
         ----------
-        filepath : str
+        filepath : Union[str, Path]
             The path to the YAML file from which the parameters will be read.
+        use_dask: bool, optional
+            Indicates whether to use dask for processing. If True, data is processed with dask; if False, data is processed eagerly. Defaults to False.
 
         Returns
         -------
         BoundaryForcing
             An instance of the BoundaryForcing class.
         """
+        filepath = Path(filepath)
         # Read the entire file content
-        with open(filepath, "r") as file:
+        with filepath.open("r") as file:
             file_content = file.read()
 
         # Split the content into YAML documents
@@ -661,7 +613,4 @@ class BoundaryForcing(ROMSToolsMixins):
         grid = Grid.from_yaml(filepath)
 
         # Create and return an instance of InitialConditions
-        return cls(
-            grid=grid,
-            **boundary_forcing_data,
-        )
+        return cls(grid=grid, **boundary_forcing_data, use_dask=use_dask)