roms-tools 1.2.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)
 
@@ -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,16 +167,15 @@ 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):
@@ -176,6 +186,7 @@ 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:
@@ -216,8 +227,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 +238,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 +259,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 +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
@@ -516,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(),
             }
@@ -530,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
@@ -582,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)

roms_tools/setup/mixins.py

@@ -207,13 +207,6 @@ class ROMSToolsMixins:
             data_vars[vname] = v_rot
 
         if "s_rho" in data_vars[uname].dims and "s_rho" in data_vars[vname].dims:
-            # 3D masks for ROMS domain
-            umask = self.grid.ds.mask_u.expand_dims({"s_rho": data_vars[uname].s_rho})
-            vmask = self.grid.ds.mask_v.expand_dims({"s_rho": data_vars[vname].s_rho})
-
-            data_vars[uname] = data_vars[uname] * umask
-            data_vars[vname] = data_vars[vname] * vmask
-
             # Compute barotropic velocity
             dz = -self.grid.ds["interface_depth_rho"].diff(dim="s_w")
             dz = dz.rename({"s_w": "s_rho"})
@@ -228,169 +221,3 @@ class ROMSToolsMixins:
             ).transpose("time", "eta_v", "xi_rho")
 
         return data_vars
-
-    def get_variable_metadata(self):
-        """
-        Retrieves metadata for commonly used variables in the dataset.
-
-        This method returns a dictionary containing the metadata for various variables, including long names
-        and units for each variable.
-
-        Returns
-        -------
-        dict of str: dict
-            Dictionary where keys are variable names and values are dictionaries with "long_name" and "units" keys.
-
-        """
-
-        d = {
-            "ssh_Re": {"long_name": "Tidal elevation, real part", "units": "m"},
-            "ssh_Im": {"long_name": "Tidal elevation, complex part", "units": "m"},
-            "pot_Re": {"long_name": "Tidal potential, real part", "units": "m"},
-            "pot_Im": {"long_name": "Tidal potential, complex part", "units": "m"},
-            "u_Re": {
-                "long_name": "Tidal velocity in x-direction, real part",
-                "units": "m/s",
-            },
-            "u_Im": {
-                "long_name": "Tidal velocity in x-direction, complex part",
-                "units": "m/s",
-            },
-            "v_Re": {
-                "long_name": "Tidal velocity in y-direction, real part",
-                "units": "m/s",
-            },
-            "v_Im": {
-                "long_name": "Tidal velocity in y-direction, complex part",
-                "units": "m/s",
-            },
-            "uwnd": {"long_name": "10 meter wind in x-direction", "units": "m/s"},
-            "vwnd": {"long_name": "10 meter wind in y-direction", "units": "m/s"},
-            "swrad": {
-                "long_name": "downward short-wave (solar) radiation",
-                "units": "W/m^2",
-            },
-            "lwrad": {
-                "long_name": "downward long-wave (thermal) radiation",
-                "units": "W/m^2",
-            },
-            "Tair": {"long_name": "air temperature at 2m", "units": "degrees Celsius"},
-            "qair": {"long_name": "absolute humidity at 2m", "units": "kg/kg"},
-            "rain": {"long_name": "total precipitation", "units": "cm/day"},
-            "temp": {"long_name": "potential temperature", "units": "degrees Celsius"},
-            "salt": {"long_name": "salinity", "units": "PSU"},
-            "zeta": {"long_name": "sea surface height", "units": "m"},
-            "u": {"long_name": "u-flux component", "units": "m/s"},
-            "v": {"long_name": "v-flux component", "units": "m/s"},
-            "w": {"long_name": "w-flux component", "units": "m/s"},
-            "ubar": {
-                "long_name": "vertically integrated u-flux component",
-                "units": "m/s",
-            },
-            "vbar": {
-                "long_name": "vertically integrated v-flux component",
-                "units": "m/s",
-            },
-            "PO4": {"long_name": "dissolved inorganic phosphate", "units": "mmol/m^3"},
-            "NO3": {"long_name": "dissolved inorganic nitrate", "units": "mmol/m^3"},
-            "SiO3": {"long_name": "dissolved inorganic silicate", "units": "mmol/m^3"},
-            "NH4": {"long_name": "dissolved ammonia", "units": "mmol/m^3"},
-            "Fe": {"long_name": "dissolved inorganic iron", "units": "mmol/m^3"},
-            "Lig": {"long_name": "iron binding ligand", "units": "mmol/m^3"},
-            "O2": {"long_name": "dissolved oxygen", "units": "mmol/m^3"},
-            "DIC": {"long_name": "dissolved inorganic carbon", "units": "mmol/m^3"},
-            "DIC_ALT_CO2": {
-                "long_name": "dissolved inorganic carbon, alternative CO2",
-                "units": "mmol/m^3",
-            },
-            "ALK": {"long_name": "alkalinity", "units": "meq/m^3"},
-            "ALK_ALT_CO2": {
-                "long_name": "alkalinity, alternative CO2",
-                "units": "meq/m^3",
-            },
-            "DOC": {"long_name": "dissolved organic carbon", "units": "mmol/m^3"},
-            "DON": {"long_name": "dissolved organic nitrogen", "units": "mmol/m^3"},
-            "DOP": {"long_name": "dissolved organic phosphorus", "units": "mmol/m^3"},
-            "DOCr": {
-                "long_name": "refractory dissolved organic carbon",
-                "units": "mmol/m^3",
-            },
-            "DONr": {
-                "long_name": "refractory dissolved organic nitrogen",
-                "units": "mmol/m^3",
-            },
-            "DOPr": {
-                "long_name": "refractory dissolved organic phosphorus",
-                "units": "mmol/m^3",
-            },
-            "zooC": {"long_name": "zooplankton carbon", "units": "mmol/m^3"},
-            "spChl": {
-                "long_name": "small phytoplankton chlorophyll",
-                "units": "mg/m^3",
-            },
-            "spC": {"long_name": "small phytoplankton carbon", "units": "mmol/m^3"},
-            "spP": {
-                "long_name": "small phytoplankton phosphorous",
-                "units": "mmol/m^3",
-            },
-            "spFe": {"long_name": "small phytoplankton iron", "units": "mmol/m^3"},
-            "spCaCO3": {"long_name": "small phytoplankton CaCO3", "units": "mmol/m^3"},
-            "diatChl": {"long_name": "diatom chloropyll", "units": "mg/m^3"},
-            "diatC": {"long_name": "diatom carbon", "units": "mmol/m^3"},
-            "diatP": {"long_name": "diatom phosphorus", "units": "mmol/m^3"},
-            "diatFe": {"long_name": "diatom iron", "units": "mmol/m^3"},
-            "diatSi": {"long_name": "diatom silicate", "units": "mmol/m^3"},
-            "diazChl": {"long_name": "diazotroph chloropyll", "units": "mg/m^3"},
-            "diazC": {"long_name": "diazotroph carbon", "units": "mmol/m^3"},
-            "diazP": {"long_name": "diazotroph phosphorus", "units": "mmol/m^3"},
-            "diazFe": {"long_name": "diazotroph iron", "units": "mmol/m^3"},
-            "pco2_air": {"long_name": "atmospheric pCO2", "units": "ppmv"},
-            "pco2_air_alt": {
-                "long_name": "atmospheric pCO2, alternative CO2",
-                "units": "ppmv",
-            },
-            "iron": {"long_name": "iron decomposition", "units": "nmol/cm^2/s"},
-            "dust": {"long_name": "dust decomposition", "units": "kg/m^2/s"},
-            "nox": {"long_name": "NOx decomposition", "units": "kg/m^2/s"},
-            "nhy": {"long_name": "NHy decomposition", "units": "kg/m^2/s"},
-        }
-        return d
-
-    def get_boundary_info(self):
-
-        """
-        This method provides information about the boundary points for the rho, u, and v
-        variables on the grid, specifying the indices for the south, east, north, and west
-        boundaries.
-
-        Returns
-        -------
-        dict
-            A dictionary where keys are variable types ("rho", "u", "v"), and values
-            are nested dictionaries mapping directions ("south", "east", "north", "west")
-            to the corresponding boundary coordinates.
-        """
-
-        # Boundary coordinates
-        bdry_coords = {
-            "rho": {
-                "south": {"eta_rho": 0},
-                "east": {"xi_rho": -1},
-                "north": {"eta_rho": -1},
-                "west": {"xi_rho": 0},
-            },
-            "u": {
-                "south": {"eta_rho": 0},
-                "east": {"xi_u": -1},
-                "north": {"eta_rho": -1},
-                "west": {"xi_u": 0},
-            },
-            "v": {
-                "south": {"eta_v": 0},
-                "east": {"xi_rho": -1},
-                "north": {"eta_v": -1},
-                "west": {"xi_rho": 0},
-            },
-        }
-
-        return bdry_coords

roms_tools/setup/plot.py

@@ -71,7 +71,14 @@ def _plot(
 
     if field is not None:
         p = ax.pcolormesh(lon_deg, lat_deg, field, transform=proj, **kwargs)
-        plt.colorbar(p, label=f"{field.long_name} [{field.units}]")
+        if hasattr(field, "long_name"):
+            label = f"{field.long_name} [{field.units}]"
+        elif hasattr(field, "Long_name"):
+            # this is the case for matlab generated grids
+            label = f"{field.Long_name} [{field.units}]"
+        else:
+            label = ""
+        plt.colorbar(p, label=label)
 
     if depth_contours:
         cs = ax.contour(lon_deg, lat_deg, field.layer_depth, transform=proj, colors="k")