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

roms_tools/setup/grid.py

@@ -7,12 +7,14 @@ import matplotlib.pyplot as plt
 import yaml
 import importlib.metadata
 
+from typing import Union
 from roms_tools.setup.topography import _add_topography_and_mask, _add_velocity_masks
 from roms_tools.setup.plot import _plot, _section_plot, _profile_plot, _line_plot
 from roms_tools.setup.utils import interpolate_from_rho_to_u, interpolate_from_rho_to_v
 from roms_tools.setup.vertical_coordinate import sigma_stretch, compute_depth
-
+from roms_tools.setup.utils import extract_single_value, save_datasets
 import warnings
+from pathlib import Path
 
 RADIUS_OF_EARTH = 6371315.0  # in m
 
@@ -236,10 +238,17 @@ class Grid:
                 object.__setattr__(self, "ds", ds)
             else:
                 self.ds[coarse_var] = coarse_field
+
         self.ds["mask_coarse"] = xr.where(self.ds["mask_coarse"] > 0.5, 1, 0).astype(
             np.int32
         )
 
+        for fine_var, coarse_var in d.items():
+            self.ds[coarse_var].attrs[
+                "long_name"
+            ] = f"{self.ds[fine_var].attrs['long_name']} on coarsened grid"
+            self.ds[coarse_var].attrs["units"] = self.ds[fine_var].attrs["units"]
+
     def update_vertical_coordinate(self, N, theta_s, theta_b, hc) -> None:
         """
         Create vertical coordinate variables for the ROMS grid.
@@ -275,7 +284,7 @@ class Grid:
             "interface_depth_rho",
             "interface_depth_u",
             "interface_depth_v",
-            "sc_r",
+            "Cs_w",
             "Cs_r",
         ]
 
@@ -290,14 +299,14 @@ class Grid:
         cs_w, sigma_w = sigma_stretch(theta_s, theta_b, N, "w")
         zw = compute_depth(h * 0, h, hc, cs_w, sigma_w)
 
-        ds["sc_r"] = sigma_r.astype(np.float32)
-        ds["sc_r"].attrs["long_name"] = "S-coordinate at rho-points"
-        ds["sc_r"].attrs["units"] = "nondimensional"
-
         ds["Cs_r"] = cs_r.astype(np.float32)
         ds["Cs_r"].attrs["long_name"] = "S-coordinate stretching curves at rho-points"
         ds["Cs_r"].attrs["units"] = "nondimensional"
 
+        ds["Cs_w"] = cs_w.astype(np.float32)
+        ds["Cs_w"].attrs["long_name"] = "S-coordinate stretching curves at w-points"
+        ds["Cs_w"].attrs["units"] = "nondimensional"
+
         ds.attrs["theta_s"] = np.float32(theta_s)
         ds.attrs["theta_b"] = np.float32(theta_b)
         ds.attrs["hc"] = np.float32(hc)
@@ -515,24 +524,60 @@ class Grid:
                 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 grid 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)
+
+        # 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
 
     @classmethod
-    def from_file(cls, filepath: str) -> "Grid":
+    def from_file(cls, filepath: Union[str, Path]) -> "Grid":
         """
         Create a Grid instance from an existing file.
 
         Parameters
         ----------
-        filepath : str
+        filepath : Union[str, Path]
             Path to the file containing the grid information.
 
         Returns
@@ -571,8 +616,14 @@ class Grid:
         ):
             grid._coarsen()
 
+        # Move variables to coordinates if necessary
+        for var in ["lat_rho", "lon_rho", "lat_coarse", "lon_coarse"]:
+            if var not in ds.coords:
+                ds = grid.ds.set_coords(var)
+                object.__setattr__(grid, "ds", ds)
+
         # Update vertical coordinate if necessary
-        if not all(var in grid.ds for var in ["sc_r", "Cs_r"]):
+        if not all(var in grid.ds for var in ["Cs_r", "Cs_w"]):
             N = 100
             theta_s = 5.0
             theta_b = 2.0
@@ -581,13 +632,45 @@ class Grid:
             grid.update_vertical_coordinate(
                 N=N, theta_s=theta_s, theta_b=theta_b, hc=hc
             )
+        else:
+            object.__setattr__(grid, "theta_s", ds.attrs["theta_s"].item())
+            object.__setattr__(grid, "theta_b", ds.attrs["theta_b"].item())
+            object.__setattr__(grid, "hc", ds.attrs["hc"].item())
+            object.__setattr__(grid, "N", len(ds.s_rho))
 
         # Manually set the remaining attributes by extracting parameters from dataset
         object.__setattr__(grid, "nx", ds.sizes["xi_rho"] - 2)
         object.__setattr__(grid, "ny", ds.sizes["eta_rho"] - 2)
-        object.__setattr__(grid, "center_lon", ds.attrs["center_lon"])
-        object.__setattr__(grid, "center_lat", ds.attrs["center_lat"])
-        object.__setattr__(grid, "rot", ds.attrs["rot"])
+        if "center_lon" in ds.attrs:
+            center_lon = ds.attrs["center_lon"]
+        elif "tra_lon" in ds:
+            center_lon = extract_single_value(ds["tra_lon"])
+        else:
+            raise ValueError(
+                "Missing grid information: 'center_lon' attribute or 'tra_lon' variable "
+                "must be present in the dataset."
+            )
+        object.__setattr__(grid, "center_lon", center_lon)
+        if "center_lat" in ds.attrs:
+            center_lat = ds.attrs["center_lat"]
+        elif "tra_lat" in ds:
+            center_lat = extract_single_value(ds["tra_lat"])
+        else:
+            raise ValueError(
+                "Missing grid information: 'center_lat' attribute or 'tra_lat' variable "
+                "must be present in the dataset."
+            )
+        object.__setattr__(grid, "center_lat", center_lat)
+        if "rot" in ds.attrs:
+            rot = ds.attrs["rot"]
+        elif "rotate" in ds:
+            rot = extract_single_value(ds["rotate"])
+        else:
+            raise ValueError(
+                "Missing grid information: 'rot' attribute or 'rotate' variable "
+                "must be present in the dataset."
+            )
+        object.__setattr__(grid, "rot", rot)
 
         for attr in [
             "size_x",
@@ -596,19 +679,25 @@ class Grid:
             "hmin",
         ]:
             if attr in ds.attrs:
-                object.__setattr__(grid, attr, ds.attrs[attr])
+                a = ds.attrs[attr]
+            else:
+                a = None
+            object.__setattr__(grid, attr, a)
 
         return grid
 
-    def to_yaml(self, filepath: str) -> None:
+    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)
+
         data = asdict(self)
         data.pop("ds", None)
         data.pop("straddle", None)
@@ -625,20 +714,20 @@ class Grid:
         # Use the class name as the top-level key
         yaml_data = {self.__class__.__name__: 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) -> "Grid":
+    def from_yaml(cls, filepath: Union[str, Path]) -> "Grid":
         """
         Create an instance of the 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.
 
         Returns
@@ -646,8 +735,10 @@ class Grid:
         Grid
             An instance of the Grid 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

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,23 @@ 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]]
+        If no exact match is found, the closest time entry to `ini_time` within the time range [ini_time, ini_time + 24 hours] is selected.
+    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.
-        - "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 BGC data is climatology data. Defaults to False.
     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,20 +58,21 @@ 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",
-    ...         "climatology": True,
+    ...         "climatology": False,
     ...     },
     ... )
     """
 
     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)
 
@@ -97,42 +107,43 @@ class InitialConditions(ROMSToolsMixins):
             vars_3d = bgc_data.var_names.keys()
             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"]
-            ):
-                for var in bgc_data_vars.keys():
-                    bgc_data_vars[var] = bgc_data_vars[var].assign_coords(
-                        {"time": data_vars["temp"]["time"]}
-                    )
+            # Ensure time coordinate matches that of physical variables
+            for var in bgc_data_vars.keys():
+                bgc_data_vars[var] = bgc_data_vars[var].assign_coords(
+                    {"time": data_vars["temp"]["time"]}
+                )
 
             # 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:
@@ -144,28 +155,27 @@ class InitialConditions(ROMSToolsMixins):
                 raise ValueError(
                     "`bgc_source` must include a 'path' if it is provided."
                 )
-            # set self.bgc_source["climatology"] to True if not provided
+            # set self.bgc_source["climatology"] to False if not provided
             object.__setattr__(
                 self,
                 "bgc_source",
                 {
                     **self.bgc_source,
-                    "climatology": self.bgc_source.get("climatology", True),
+                    "climatology": self.bgc_source.get("climatology", False),
                 },
             )
 
     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):
@@ -176,8 +186,8 @@ class InitialConditions(ROMSToolsMixins):
                 filename=self.bgc_source["path"],
                 start_time=self.ini_time,
                 climatology=self.bgc_source["climatology"],
+                use_dask=self.use_dask,
             )
-            data.post_process()
         else:
             raise ValueError(
                 'Only "CESM_REGRIDDED" is a valid option for bgc_source["name"].'
@@ -216,8 +226,8 @@ class InitialConditions(ROMSToolsMixins):
         existing_vars = [var for var in variables_to_drop if var in ds]
         ds = ds.drop_vars(existing_vars)
 
-        ds["sc_r"] = self.grid.ds["sc_r"]
         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"]})
@@ -227,10 +237,10 @@ class InitialConditions(ROMSToolsMixins):
 
         # 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 = ds.assign_coords(ocean_time=("time", ocean_time.data.astype("float64")))
         ds["ocean_time"].attrs[
             "long_name"
-        ] = f"seconds since {np.datetime_as_string(model_reference_date, unit='s')}"
+        ] = 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")
@@ -248,7 +258,7 @@ 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"]
 
@@ -479,25 +489,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
@@ -516,7 +564,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(),
             }
@@ -530,29 +578,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
@@ -582,7 +635,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)