roms-tools 1.1.0__py3-none-any.whl → 1.3.0__py3-none-any.whl

roms_tools/setup/initial_conditions.py

@@ -3,16 +3,20 @@ import numpy as np
 import yaml
 import importlib.metadata
 from dataclasses import dataclass, field, asdict
-from typing import Optional, Dict, Union
+from typing import Dict, Union, List, Optional
 from roms_tools.setup.grid import Grid
 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,
+    save_datasets,
 )
 from roms_tools.setup.mixins import ROMSToolsMixins
 from roms_tools.setup.plot import _plot, _section_plot, _profile_plot, _line_plot
 import matplotlib.pyplot as plt
+from pathlib import Path
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -26,18 +30,22 @@ class InitialConditions(ROMSToolsMixins):
         Object representing the grid information used for the model.
     ini_time : datetime
         The date and time at which the initial conditions are set.
-    physics_source : Dict[str, Union[str, None]]
+    source : Dict[str, Union[str, Path, List[Union[str, Path]]], bool]
         Dictionary specifying the source of the physical initial condition data:
         - "name" (str): Name of the data source (e.g., "GLORYS").
-        - "path" (str): Path to the physical data file. Can contain wildcards.
+        - "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 physical data is climatology data. Defaults to False.
-    bgc_source : Optional[Dict[str, Union[str, None]]]
+    bgc_source : Optional[Dict[str, Union[str, Path, List[Union[str, Path]]], bool]]
         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.
+        - "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 BGC data is climatology data. Defaults to True.
     model_reference_date : datetime, optional
         The reference date for the model. Defaults to 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
     ----------
@@ -49,7 +57,7 @@ class InitialConditions(ROMSToolsMixins):
     >>> initial_conditions = InitialConditions(
     ...     grid=grid,
     ...     ini_time=datetime(2022, 1, 1),
-    ...     physics_source={"name": "GLORYS", "path": "physics_data.nc"},
+    ...     source={"name": "GLORYS", "path": "physics_data.nc"},
     ...     bgc_source={
     ...         "name": "CESM_REGRIDDED",
     ...         "path": "bgc_data.nc",
@@ -60,9 +68,10 @@ class InitialConditions(ROMSToolsMixins):
 
     grid: Grid
     ini_time: datetime
-    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]]]]
+    bgc_source: Optional[Dict[str, Union[str, Path, List[Union[str, Path]]]]] = None
     model_reference_date: datetime = datetime(2000, 1, 1)
+    use_dask: bool = False
 
     ds: xr.Dataset = field(init=False, repr=False)
 
@@ -82,7 +91,7 @@ class InitialConditions(ROMSToolsMixins):
         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)
+        data_vars = super().process_velocities(data_vars, angle, "u", "v")
 
         if self.bgc_source is not None:
             bgc_data = self._get_bgc_data()
@@ -98,10 +107,7 @@ class InitialConditions(ROMSToolsMixins):
             bgc_data_vars = super().regrid_data(bgc_data, vars_2d, vars_3d, lon, lat)
 
             # Ensure time coordinate matches if climatology is applied in one case but not the other
-            if (
-                not self.physics_source["climatology"]
-                and self.bgc_source["climatology"]
-            ):
+            if not self.source["climatology"] and self.bgc_source["climatology"]:
                 for var in bgc_data_vars.keys():
                     bgc_data_vars[var] = bgc_data_vars[var].assign_coords(
                         {"time": data_vars["temp"]["time"]}
@@ -110,29 +116,34 @@ class InitialConditions(ROMSToolsMixins):
             # Combine data variables from physical and biogeochemical sources
             data_vars.update(bgc_data_vars)
 
-        d_meta = super().get_variable_metadata()
+        d_meta = get_variable_metadata()
         ds = self._write_into_dataset(data_vars, d_meta)
 
         ds = self._add_global_metadata(ds)
 
         ds["zeta"].load()
+        # NaN values at wet points indicate that the raw data did not cover the domain, and the following will raise a ValueError
         nan_check(ds["zeta"].squeeze(), self.grid.ds.mask_rho)
 
+        # 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):
 
-        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
+        if "name" not in self.source.keys():
+            raise ValueError("`source` must include a 'name'.")
+        if "path" not in self.source.keys():
+            raise ValueError("`source` must include a 'path'.")
+        # set self.source["climatology"] to False if not provided
         object.__setattr__(
             self,
-            "physics_source",
+            "source",
             {
-                **self.physics_source,
-                "climatology": self.physics_source.get("climatology", False),
+                **self.source,
+                "climatology": self.source.get("climatology", False),
             },
         )
         if self.bgc_source is not None:
@@ -156,34 +167,34 @@ class InitialConditions(ROMSToolsMixins):
 
     def _get_data(self):
 
-        if self.physics_source["name"] == "GLORYS":
+        if self.source["name"] == "GLORYS":
             data = GLORYSDataset(
-                filename=self.physics_source["path"],
+                filename=self.source["path"],
                 start_time=self.ini_time,
-                climatology=self.physics_source["climatology"],
+                climatology=self.source["climatology"],
+                use_dask=self.use_dask,
             )
         else:
-            raise ValueError(
-                'Only "GLORYS" is a valid option for physics_source["name"].'
-            )
+            raise ValueError('Only "GLORYS" is a valid option for source["name"].')
         return data
 
     def _get_bgc_data(self):
 
         if self.bgc_source["name"] == "CESM_REGRIDDED":
 
-            bgc_data = CESMBGCDataset(
+            data = CESMBGCDataset(
                 filename=self.bgc_source["path"],
                 start_time=self.ini_time,
                 climatology=self.bgc_source["climatology"],
+                use_dask=self.use_dask,
             )
-            bgc_data.post_process()
+            data.post_process()
         else:
             raise ValueError(
                 'Only "CESM_REGRIDDED" is a valid option for bgc_source["name"].'
             )
 
-        return bgc_data
+        return data
 
     def _write_into_dataset(self, data_vars, d_meta):
 
@@ -202,19 +213,43 @@ class InitialConditions(ROMSToolsMixins):
         ds["w"].attrs["long_name"] = d_meta["w"]["long_name"]
         ds["w"].attrs["units"] = d_meta["w"]["units"]
 
+        variables_to_drop = [
+            "s_rho",
+            "lat_rho",
+            "lon_rho",
+            "layer_depth_rho",
+            "interface_depth_rho",
+            "lat_u",
+            "lon_u",
+            "lat_v",
+            "lon_v",
+        ]
+        existing_vars = [var for var in variables_to_drop if var in ds]
+        ds = ds.drop_vars(existing_vars)
+
+        ds["Cs_r"] = self.grid.ds["Cs_r"]
+        ds["Cs_w"] = self.grid.ds["Cs_w"]
+
+        # Preserve absolute time coordinate for readability
+        ds = ds.assign_coords({"abs_time": ds["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)
+        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)}"
+        ds["ocean_time"].attrs["units"] = "seconds"
+        ds = ds.swap_dims({"time": "ocean_time"})
+        ds = ds.drop_vars("time")
+
         return ds
 
     def _add_global_metadata(self, ds):
 
-        ds = ds.assign_coords(
-            {
-                "layer_depth_u": self.grid.ds["layer_depth_u"],
-                "layer_depth_v": self.grid.ds["layer_depth_v"],
-                "interface_depth_u": self.grid.ds["interface_depth_u"],
-                "interface_depth_v": self.grid.ds["interface_depth_v"],
-            }
-        )
-
         ds.attrs["title"] = "ROMS initial conditions file created by ROMS-Tools"
         # Include the version of roms-tools
         try:
@@ -224,28 +259,13 @@ class InitialConditions(ROMSToolsMixins):
         ds.attrs["roms_tools_version"] = roms_tools_version
         ds.attrs["ini_time"] = str(self.ini_time)
         ds.attrs["model_reference_date"] = str(self.model_reference_date)
-        ds.attrs["physical_source"] = self.physics_source["name"]
+        ds.attrs["source"] = self.source["name"]
         if self.bgc_source is not None:
             ds.attrs["bgc_source"] = self.bgc_source["name"]
 
-        # 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)
-        ocean_time = (ds["time"] - model_reference_date).astype("float64") * 1e-9
-        ds = ds.assign_coords(ocean_time=("time", np.float32(ocean_time)))
-        ds["ocean_time"].attrs[
-            "long_name"
-        ] = f"seconds since {np.datetime_as_string(model_reference_date, unit='s')}"
-        ds["ocean_time"].attrs["units"] = "seconds"
-
         ds.attrs["theta_s"] = self.grid.ds.attrs["theta_s"]
         ds.attrs["theta_b"] = self.grid.ds.attrs["theta_b"]
         ds.attrs["hc"] = self.grid.ds.attrs["hc"]
-        ds["sc_r"] = self.grid.ds["sc_r"]
-        ds["Cs_r"] = self.grid.ds["Cs_r"]
-
-        ds = ds.drop_vars(["s_rho"])
 
         return ds
 
@@ -306,13 +326,23 @@ class InitialConditions(ROMSToolsMixins):
             - "diazP": Diazotroph Phosphorus (mmol/m³).
             - "diazFe": Diazotroph Iron (mmol/m³).
         s : int, optional
-            The index of the vertical layer to plot. Default is None.
+            The index of the vertical layer (`s_rho`) to plot. If not specified, the plot
+            will represent a horizontal slice (eta- or xi- plane). Default is None.
         eta : int, optional
-            The eta-index to plot. Default is None.
+            The eta-index to plot. Used for vertical sections or horizontal slices.
+            Default is None.
         xi : int, optional
-            The xi-index to plot. Default is None.
+            The xi-index to plot. Used for vertical sections or horizontal slices.
+            Default is None.
         depth_contours : bool, optional
-            Whether to include depth contours in the plot. Default is False.
+            If True, depth contours will be overlaid on the plot, showing lines of constant
+            depth. This is typically used for plots that show a single vertical layer.
+            Default is False.
+        layer_contours : bool, optional
+            If True, contour lines representing the boundaries between vertical layers will
+            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.
 
         Returns
         -------
@@ -325,6 +355,7 @@ class InitialConditions(ROMSToolsMixins):
             If the specified `varname` is not one of the valid options.
             If the field specified by `varname` is 3D and none of `s`, `eta`, or `xi` are specified.
             If the field specified by `varname` is 2D and both `eta` and `xi` are specified.
+
         """
 
         if len(self.ds[varname].squeeze().dims) == 3 and not any(
@@ -343,17 +374,38 @@ class InitialConditions(ROMSToolsMixins):
         field = self.ds[varname].squeeze()
 
         if all(dim in field.dims for dim in ["eta_rho", "xi_rho"]):
-            interface_depth = self.ds.interface_depth_rho
+            interface_depth = self.grid.ds.interface_depth_rho
+            layer_depth = self.grid.ds.layer_depth_rho
+            field = field.where(self.grid.ds.mask_rho)
+            field = field.assign_coords(
+                {"lon": self.grid.ds.lon_rho, "lat": self.grid.ds.lat_rho}
+            )
+
         elif all(dim in field.dims for dim in ["eta_rho", "xi_u"]):
-            interface_depth = self.ds.interface_depth_u
+            interface_depth = self.grid.ds.interface_depth_u
+            layer_depth = self.grid.ds.layer_depth_u
+            field = field.where(self.grid.ds.mask_u)
+            field = field.assign_coords(
+                {"lon": self.grid.ds.lon_u, "lat": self.grid.ds.lat_u}
+            )
+
         elif all(dim in field.dims for dim in ["eta_v", "xi_rho"]):
-            interface_depth = self.ds.interface_depth_v
+            interface_depth = self.grid.ds.interface_depth_v
+            layer_depth = self.grid.ds.layer_depth_v
+            field = field.where(self.grid.ds.mask_v)
+            field = field.assign_coords(
+                {"lon": self.grid.ds.lon_v, "lat": self.grid.ds.lat_v}
+            )
+        else:
+            ValueError("provided field does not have two horizontal dimension")
 
         # slice the field as desired
         title = field.long_name
         if s is not None:
             title = title + f", s_rho = {field.s_rho[s].item()}"
             field = field.isel(s_rho=s)
+            layer_depth = layer_depth.isel(s_rho=s)
+            field = field.assign_coords({"layer_depth": layer_depth})
         else:
             depth_contours = False
 
@@ -361,10 +413,14 @@ class InitialConditions(ROMSToolsMixins):
             if "eta_rho" in field.dims:
                 title = title + f", eta_rho = {field.eta_rho[eta].item()}"
                 field = field.isel(eta_rho=eta)
+                layer_depth = layer_depth.isel(eta_rho=eta)
+                field = field.assign_coords({"layer_depth": layer_depth})
                 interface_depth = interface_depth.isel(eta_rho=eta)
             elif "eta_v" in field.dims:
                 title = title + f", eta_v = {field.eta_v[eta].item()}"
                 field = field.isel(eta_v=eta)
+                layer_depth = layer_depth.isel(eta_v=eta)
+                field = field.assign_coords({"layer_depth": layer_depth})
                 interface_depth = interface_depth.isel(eta_v=eta)
             else:
                 raise ValueError(
@@ -374,10 +430,14 @@ class InitialConditions(ROMSToolsMixins):
             if "xi_rho" in field.dims:
                 title = title + f", xi_rho = {field.xi_rho[xi].item()}"
                 field = field.isel(xi_rho=xi)
+                layer_depth = layer_depth.isel(xi_rho=xi)
+                field = field.assign_coords({"layer_depth": layer_depth})
                 interface_depth = interface_depth.isel(xi_rho=xi)
             elif "xi_u" in field.dims:
                 title = title + f", xi_u = {field.xi_u[xi].item()}"
                 field = field.isel(xi_u=xi)
+                layer_depth = layer_depth.isel(xi_u=xi)
+                field = field.assign_coords({"layer_depth": layer_depth})
                 interface_depth = interface_depth.isel(xi_u=xi)
             else:
                 raise ValueError(
@@ -430,25 +490,63 @@ class InitialConditions(ROMSToolsMixins):
                 else:
                     _line_plot(field, title=title)
 
-    def save(self, filepath: str) -> None:
+    def save(
+        self, filepath: Union[str, Path], np_eta: int = None, np_xi: int = None
+    ) -> None:
         """
         Save the initial conditions information to a netCDF4 file.
 
+        This method supports saving the dataset in two modes:
+
+        1. **Single File Mode (default)**:
+           - If both `np_eta` and `np_xi` are `None`, the entire dataset is saved as a single file at the specified `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.
+           - The files are saved as `filepath.0.nc`, `filepath.1.nc`, ..., where the numbering corresponds to the partition index.
+
         Parameters
         ----------
-        filepath
+        filepath : Union[str, Path]
+            The base path or filename where the dataset should be saved.
+        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
+        -------
+        List[Path]
+            A list of Path objects for the filenames that were saved.
         """
-        self.ds.to_netcdf(filepath)
 
-    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 = [self.ds.load()]
+        output_filenames = [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
@@ -467,7 +565,7 @@ class InitialConditions(ROMSToolsMixins):
 
         initial_conditions_data = {
             "InitialConditions": {
-                "physics_source": self.physics_source,
+                "source": self.source,
                 "ini_time": self.ini_time.isoformat(),
                 "model_reference_date": self.model_reference_date.isoformat(),
             }
@@ -481,29 +579,34 @@ class InitialConditions(ROMSToolsMixins):
             **initial_conditions_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) -> "InitialConditions":
+    def from_yaml(
+        cls, filepath: Union[str, Path], use_dask: bool = False
+    ) -> "InitialConditions":
         """
         Create an instance of the InitialConditions 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
         -------
         InitialConditions
             An instance of the InitialConditions 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
@@ -533,7 +636,4 @@ class InitialConditions(ROMSToolsMixins):
         grid = Grid.from_yaml(filepath)
 
         # Create and return an instance of InitialConditions
-        return cls(
-            grid=grid,
-            **initial_conditions_data,
-        )
+        return cls(grid=grid, **initial_conditions_data, use_dask=use_dask)