roms-tools 1.6.1__py3-none-any.whl → 1.7.0__py3-none-any.whl

roms_tools/setup/surface_forcing.py

@@ -1,8 +1,6 @@
 import xarray as xr
-import pandas as pd
-import yaml
 import importlib.metadata
-from dataclasses import dataclass, field, asdict
+from dataclasses import dataclass, field
 from roms_tools.setup.grid import Grid
 from datetime import datetime
 import numpy as np
@@ -22,6 +20,9 @@ from roms_tools.setup.utils import (
     save_datasets,
     get_target_coords,
     rotate_velocities,
+    convert_to_roms_time,
+    _to_yaml,
+    _from_yaml,
 )
 from roms_tools.setup.plot import _plot
 import matplotlib.pyplot as plt
@@ -133,14 +134,7 @@ class SurfaceForcing:
 
         ds = self._write_into_dataset(processed_fields, data, d_meta)
 
-        if self.use_coarse_grid:
-            mask = self.grid.ds["mask_coarse"].rename(
-                {"eta_coarse": "eta_rho", "xi_coarse": "xi_rho"}
-            )
-        else:
-            mask = self.grid.ds["mask_rho"]
-
-        self._validate(ds, mask)
+        self._validate(ds, target_coords["mask"], variable_info)
 
         # substitute NaNs over land by a fill value to avoid blow-up of ROMS
         for var_name in ds.data_vars:
@@ -233,28 +227,34 @@ class SurfaceForcing:
         # Define a dictionary for variable names and their associated information
         if self.type == "physics":
             variable_info = {
-                "swrad": default_info,
-                "lwrad": default_info,
-                "Tair": default_info,
-                "qair": default_info,
-                "rain": default_info,
+                "swrad": {**default_info, "validate": True},
+                "lwrad": {**default_info, "validate": False},
+                "Tair": {**default_info, "validate": False},
+                "qair": {**default_info, "validate": True},
+                "rain": {**default_info, "validate": False},
                 "uwnd": {
                     "location": "u",
                     "is_vector": True,
                     "vector_pair": "vwnd",
                     "is_3d": False,
+                    "validate": True,
                 },
                 "vwnd": {
                     "location": "v",
                     "is_vector": True,
                     "vector_pair": "uwnd",
                     "is_3d": False,
+                    "validate": True,
                 },
             }
         elif self.type == "bgc":
             variable_info = {}
-            for var in data.var_names.keys():
-                variable_info[var] = default_info
+            for var_name in data.var_names.keys():
+                variable_info[var_name] = default_info
+                if var_name == "pco2_air":
+                    variable_info[var_name] = {**default_info, "validate": True}
+                else:
+                    variable_info[var_name] = {**default_info, "validate": False}
 
         return variable_info
 
@@ -300,44 +300,12 @@ class SurfaceForcing:
             ds[var_name].attrs["long_name"] = d_meta[var_name]["long_name"]
             ds[var_name].attrs["units"] = d_meta[var_name]["units"]
 
-        if self.use_coarse_grid:
-            ds = ds.rename({"eta_coarse": "eta_rho", "xi_coarse": "xi_rho"})
-
         ds = self._add_global_metadata(ds)
 
         # Convert the time coordinate to the format expected by ROMS
-        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
-            sfc_time = xr.DataArray(
-                (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"]})
-
-            sfc_time = (
-                (ds["time"] - np.datetime64(self.model_reference_date)).astype(
-                    "float64"
-                )
-                / 3600
-                / 24
-                * 1e-9
-            )
+        ds, sfc_time = convert_to_roms_time(
+            ds, self.model_reference_date, data.climatology
+        )
 
         if self.type == "physics":
             time_coords = ["time"]
@@ -351,14 +319,6 @@ class SurfaceForcing:
             ]
         for time_coord in time_coords:
             ds = ds.assign_coords({time_coord: sfc_time})
-            ds[time_coord].attrs[
-                "long_name"
-            ] = f"days since {str(self.model_reference_date)}"
-            ds[time_coord].encoding["units"] = "days"
-            ds[time_coord].attrs["units"] = "days"
-            if data.climatology:
-                ds[time_coord].attrs["cycle_length"] = 365.25
-        ds.encoding["unlimited_dims"] = "time"
 
         if self.type == "bgc":
             ds = ds.drop_vars(["time"])
@@ -369,7 +329,7 @@ class SurfaceForcing:
 
         return ds
 
-    def _validate(self, ds, mask):
+    def _validate(self, ds, mask, variable_info):
         """Validates the dataset by checking for NaN values at wet points, which would
         indicate missing raw data coverage over the target domain.
 
@@ -379,6 +339,10 @@ class SurfaceForcing:
             The dataset to validate.
         mask : xarray.DataArray
             Land mask (1=ocean, 0=land) to determine wet points in the domain.
+        variable_info : dict
+            A dictionary containing metadata about each variable (e.g., location,
+            whether it's a 3D variable, etc.). Used to retrieve information for
+            validating each variable.
 
         Raises
         ------
@@ -392,7 +356,9 @@ class SurfaceForcing:
         """
 
         for var_name in ds.data_vars:
-            nan_check(ds[var_name].isel(time=0), mask)
+            # Only validate variables based on "validate" flag if use_dask is False
+            if not self.use_dask or variable_info[var_name]["validate"]:
+                nan_check(ds[var_name].isel(time=0), mask)
 
     def _add_global_metadata(self, ds=None):
 
@@ -471,12 +437,7 @@ class SurfaceForcing:
 
         title = field.long_name
 
-        # assign lat / lon
-        if self.use_coarse_grid:
-            field = field.rename({"eta_rho": "eta_coarse", "xi_rho": "xi_coarse"})
-            field = field.where(self.grid.ds.mask_coarse)
-        else:
-            field = field.where(self.grid.ds.mask_rho)
+        field = field.where(self.target_coords["mask"])
 
         field = field.assign_coords(
             {"lon": self.target_coords["lon"], "lat": self.target_coords["lat"]}
@@ -582,49 +543,8 @@ class SurfaceForcing:
         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
-        grid_data.pop("straddle", None)
 
-        # Include the version of roms-tools
-        try:
-            roms_tools_version = importlib.metadata.version("roms-tools")
-        except importlib.metadata.PackageNotFoundError:
-            roms_tools_version = "unknown"
-
-        # Create header
-        header = f"---\nroms_tools_version: {roms_tools_version}\n---\n"
-
-        # Create YAML data for Grid and optional attributes
-        grid_yaml_data = {"Grid": grid_data}
-
-        # Combine all sections
-        surface_forcing_data = {
-            "SurfaceForcing": {
-                "start_time": self.start_time.isoformat(),
-                "end_time": self.end_time.isoformat(),
-                "source": self.source,
-                "type": self.type,
-                "correct_radiation": self.correct_radiation,
-                "use_coarse_grid": self.use_coarse_grid,
-                "model_reference_date": self.model_reference_date.isoformat(),
-            }
-        }
-
-        # Merge YAML data while excluding empty sections
-        yaml_data = {
-            **grid_yaml_data,
-            **surface_forcing_data,
-        }
-
-        with filepath.open("w") as file:
-            # Write header
-            file.write(header)
-            # Write YAML data
-            yaml.dump(yaml_data, file, default_flow_style=False, sort_keys=False)
+        _to_yaml(self, filepath)
 
     @classmethod
     def from_yaml(
@@ -645,33 +565,8 @@ class SurfaceForcing:
             An instance of the SurfaceForcing class.
         """
         filepath = Path(filepath)
-        # Read the entire file content
-        with filepath.open("r") as file:
-            file_content = file.read()
-
-        # Split the content into YAML documents
-        documents = list(yaml.safe_load_all(file_content))
-
-        surface_forcing_data = None
-
-        # Process the YAML documents
-        for doc in documents:
-            if doc is None:
-                continue
-            if "SurfaceForcing" in doc:
-                surface_forcing_data = doc["SurfaceForcing"]
-
-        if surface_forcing_data is None:
-            raise ValueError("No SurfaceForcing configuration found in the YAML file.")
-
-        # Convert from string to datetime
-        for date_string in ["model_reference_date", "start_time", "end_time"]:
-            surface_forcing_data[date_string] = datetime.fromisoformat(
-                surface_forcing_data[date_string]
-            )
 
-        # Create Grid instance from the YAML file
         grid = Grid.from_yaml(filepath)
+        params = _from_yaml(cls, filepath)
 
-        # Create and return an instance of SurfaceForcing
-        return cls(grid=grid, **surface_forcing_data, use_dask=use_dask)
+        return cls(grid=grid, **params, use_dask=use_dask)

roms_tools/setup/tides.py

@@ -1,11 +1,9 @@
 from datetime import datetime
 import xarray as xr
 import numpy as np
-import yaml
-import importlib.metadata
 from typing import Dict, Union, List
-
-from dataclasses import dataclass, field, asdict
+import importlib.metadata
+from dataclasses import dataclass, field
 from roms_tools.setup.grid import Grid
 from roms_tools.setup.plot import _plot
 from roms_tools.setup.datasets import TPXODataset
@@ -19,6 +17,8 @@ from roms_tools.setup.utils import (
     get_target_coords,
     rotate_velocities,
     get_vector_pairs,
+    _to_yaml,
+    _from_yaml,
 )
 from roms_tools.setup.regrid import LateralRegrid
 import matplotlib.pyplot as plt
@@ -129,7 +129,7 @@ class TidalForcing:
 
         ds = self._add_global_metadata(ds)
 
-        self._validate(ds)
+        self._validate(ds, variable_info)
 
         # substitute NaNs over land by a fill value to avoid blow-up of ROMS
         for var_name in ds.data_vars:
@@ -176,33 +176,37 @@ class TidalForcing:
 
         # Define a dictionary for variable names and their associated information
         variable_info = {
-            "ssh_Re": default_info,
-            "ssh_Im": default_info,
-            "pot_Re": default_info,
-            "pot_Im": default_info,
+            "ssh_Re": {**default_info, "validate": True},
+            "ssh_Im": {**default_info, "validate": False},
+            "pot_Re": {**default_info, "validate": False},
+            "pot_Im": {**default_info, "validate": False},
             "u_Re": {
                 "location": "u",
                 "is_vector": True,
                 "vector_pair": "v_Re",
                 "is_3d": False,
+                "validate": True,
             },
             "v_Re": {
                 "location": "v",
                 "is_vector": True,
                 "vector_pair": "u_Re",
                 "is_3d": False,
+                "validate": True,
             },
             "u_Im": {
                 "location": "u",
                 "is_vector": True,
                 "vector_pair": "v_Im",
                 "is_3d": False,
+                "validate": False,
             },
             "v_Im": {
                 "location": "v",
                 "is_vector": True,
                 "vector_pair": "u_Im",
                 "is_3d": False,
+                "validate": False,
             },
         }
 
@@ -239,14 +243,17 @@ class TidalForcing:
 
         return ds
 
-    def _validate(self, ds):
-        """Validates the dataset by checking for NaN values at wet points, which would
-        indicate missing raw data coverage over the target domain.
+    def _validate(self, ds, variable_info):
+        """Validates the dataset by checking for NaN values at wet points for specified
+        variables, which would indicate missing raw data coverage over the target
+        domain.
 
         Parameters
         ----------
         ds : xarray.Dataset
             The dataset to validate, containing tidal variables and a mask for wet points.
+        variable_info : dict
+            A dictionary containing metadata about the variables, including whether to validate them.
 
         Raises
         ------
@@ -260,7 +267,9 @@ class TidalForcing:
         The method utilizes `self.grid.ds.mask_rho` to determine the wet points in the domain.
         """
         for var_name in ds.data_vars:
-            nan_check(ds[var_name].isel(ntides=0), self.grid.ds.mask_rho)
+            # only validate variables based on "validate" flag if use_dask is false
+            if not self.use_dask or variable_info[var_name]["validate"]:
+                nan_check(ds[var_name].isel(ntides=0), self.grid.ds.mask_rho)
 
     def plot(self, var_name, ntides=0) -> None:
         """Plot the specified tidal forcing variable for a given tidal constituent.
@@ -409,42 +418,8 @@ class TidalForcing:
         filepath : Union[str, Path]
             The path to the YAML file where the parameters will be saved.
         """
-        filepath = Path(filepath)
-
-        grid_data = asdict(self.grid)
-        grid_data.pop("ds", None)  # Exclude non-serializable fields
-        grid_data.pop("straddle", None)
-
-        # Include the version of roms-tools
-        try:
-            roms_tools_version = importlib.metadata.version("roms-tools")
-        except importlib.metadata.PackageNotFoundError:
-            roms_tools_version = "unknown"
-
-        # Create header
-        header = f"---\nroms_tools_version: {roms_tools_version}\n---\n"
 
-        # Extract grid data
-        grid_yaml_data = {"Grid": grid_data}
-
-        # Extract tidal forcing data
-        tidal_forcing_data = {
-            "TidalForcing": {
-                "source": self.source,
-                "ntides": self.ntides,
-                "allan_factor": self.allan_factor,
-                "model_reference_date": self.model_reference_date.isoformat(),
-            }
-        }
-
-        # Combine both sections
-        yaml_data = {**grid_yaml_data, **tidal_forcing_data}
-
-        with filepath.open("w") as file:
-            # Write header
-            file.write(header)
-            # Write YAML data
-            yaml.dump(yaml_data, file, default_flow_style=False, sort_keys=False)
+        _to_yaml(self, filepath)
 
     @classmethod
     def from_yaml(
@@ -465,36 +440,9 @@ class TidalForcing:
             An instance of the TidalForcing class.
         """
         filepath = Path(filepath)
-        # Read the entire file content
-        with filepath.open("r") as file:
-            file_content = file.read()
-
-        # Split the content into YAML documents
-        documents = list(yaml.safe_load_all(file_content))
-
-        tidal_forcing_data = None
-
-        # Process the YAML documents
-        for doc in documents:
-            if doc is None:
-                continue
-            if "TidalForcing" in doc:
-                tidal_forcing_data = doc["TidalForcing"]
-                break
-
-        if tidal_forcing_data is None:
-            raise ValueError("No TidalForcing configuration found in the YAML file.")
-
-        # Convert the model_reference_date from string to datetime
-        tidal_forcing_params = tidal_forcing_data
-        tidal_forcing_params["model_reference_date"] = datetime.fromisoformat(
-            tidal_forcing_params["model_reference_date"]
-        )
 
-        # Create Grid instance from the YAML file
         grid = Grid.from_yaml(filepath)
-
-        # Create and return an instance of TidalForcing
+        tidal_forcing_params = _from_yaml(cls, filepath)
         return cls(grid=grid, **tidal_forcing_params, use_dask=use_dask)
 
     def _correct_tides(self, data):