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

roms_tools/setup/initial_conditions.py

@@ -30,6 +30,7 @@ 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.
+        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").
@@ -41,7 +42,7 @@ class InitialConditions(ROMSToolsMixins):
         - "name" (str): Name of the BGC data source (e.g., "CESM_REGRIDDED").
         - "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.
+        - "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
@@ -61,7 +62,7 @@ class InitialConditions(ROMSToolsMixins):
     ...     bgc_source={
     ...         "name": "CESM_REGRIDDED",
     ...         "path": "bgc_data.nc",
-    ...         "climatology": True,
+    ...         "climatology": False,
     ...     },
     ... )
     """
@@ -106,12 +107,11 @@ 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.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)
@@ -155,13 +155,13 @@ 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),
                 },
             )
 
@@ -188,7 +188,6 @@ class InitialConditions(ROMSToolsMixins):
                 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"].'

roms_tools/setup/mixins.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 from roms_tools.setup.grid import Grid
-from roms_tools.setup.fill import fill_and_interpolate
+from roms_tools.setup.fill import LateralFill
 from roms_tools.setup.utils import (
     extrapolate_deepest_to_bottom,
     interpolate_from_rho_to_u,
@@ -106,18 +106,25 @@ class ROMSToolsMixins:
         # interpolate onto desired grid
         data_vars = {}
 
+        # Set up solver that does lateral fill
+        lateral_fill = LateralFill(
+            data.ds["mask"],
+            [data.dim_names["latitude"], data.dim_names["longitude"]],
+        )
+
         # 2d interpolation
-        fill_dims = [data.dim_names["latitude"], data.dim_names["longitude"]]
         coords = {data.dim_names["latitude"]: lat, data.dim_names["longitude"]: lon}
         for var in vars_2d:
-            mask = xr.where(data.ds[data.var_names[var]].isel(time=0).isnull(), 0, 1)
-
-            data_vars[var] = fill_and_interpolate(
-                data.ds[data.var_names[var]].astype(np.float64),
-                mask,
-                fill_dims=fill_dims,
-                coords=coords,
-                method="linear",
+            # Propagate ocean values into land via lateral fill
+            data.ds[data.var_names[var]] = lateral_fill.apply(
+                data.ds[data.var_names[var]].astype(np.float64)
+            )
+
+            # Regrid
+            data_vars[var] = (
+                data.ds[data.var_names[var]]
+                .interp(coords, method="linear")
+                .drop_vars(list(coords.keys()))
             )
 
         if vars_3d:
@@ -127,25 +134,26 @@ class ROMSToolsMixins:
                 data.dim_names["latitude"]: lat,
                 data.dim_names["longitude"]: lon,
             }
-        # extrapolate deepest value all the way to bottom ("flooding")
+        # extrapolate deepest value all the way to bottom
         for var in vars_3d:
             data.ds[data.var_names[var]] = extrapolate_deepest_to_bottom(
                 data.ds[data.var_names[var]], data.dim_names["depth"]
             )
-            mask = xr.where(data.ds[data.var_names[var]].isel(time=0).isnull(), 0, 1)
+            # Propagate ocean values into land via lateral fill
+            data.ds[data.var_names[var]] = lateral_fill.apply(
+                data.ds[data.var_names[var]].astype(np.float64)
+            )
 
-            # setting fillvalue_interp to None means that we allow extrapolation in the
+            # Regrid
+            # setting the fill value to None means that we allow extrapolation in the
             # interpolation step to avoid NaNs at the surface if the lowest depth in original
             # data is greater than zero
-
-            data_vars[var] = fill_and_interpolate(
-                data.ds[data.var_names[var]].astype(np.float64),
-                mask,
-                fill_dims=fill_dims,
-                coords=coords,
-                method="linear",
-                fillvalue_interp=None,
+            data_vars[var] = (
+                data.ds[data.var_names[var]]
+                .interp(coords, method="linear", kwargs={"fill_value": None})
+                .drop_vars(list(coords.keys()))
             )
+
             if data.dim_names["time"] != "time":
                 data_vars[var] = data_vars[var].rename({data.dim_names["time"]: "time"})
 

roms_tools/setup/surface_forcing.py

@@ -127,13 +127,12 @@ class SurfaceForcing(ROMSToolsMixins):
                 correction_data.choose_subdomain(coords_correction, straddle=straddle)
                 # apply mask from ERA5 data
                 if "mask" in data.var_names.keys():
-                    mask = xr.where(
-                        data.ds[data.var_names["mask"]].isel(time=0).isnull(), 0, 1
-                    )
+                    mask = data.ds["mask"]
                     for var in correction_data.ds.data_vars:
                         correction_data.ds[var] = xr.where(
                             mask == 1, correction_data.ds[var], np.nan
                         )
+                    correction_data.ds["mask"] = mask
                 vars_2d = ["swr_corr"]
                 vars_3d = []
                 # spatial interpolation
@@ -203,7 +202,6 @@ class SurfaceForcing(ROMSToolsMixins):
         if self.type == "physics":
             if self.source["name"] == "ERA5":
                 data = ERA5Dataset(**data_dict)
-                data.post_process()
             else:
                 raise ValueError(
                     'Only "ERA5" is a valid option for source["name"] when type is "physics".'
@@ -213,7 +211,6 @@ class SurfaceForcing(ROMSToolsMixins):
             if self.source["name"] == "CESM_REGRIDDED":
 
                 data = CESMBGCSurfaceForcingDataset(**data_dict)
-                data.post_process()
             else:
                 raise ValueError(
                     'Only "CESM_REGRIDDED" is a valid option for source["name"] when type is "bgc".'
@@ -301,6 +298,7 @@ class SurfaceForcing(ROMSToolsMixins):
             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"])

roms_tools/setup/tides.py

@@ -8,7 +8,6 @@ from typing import Dict, Union, List
 from dataclasses import dataclass, field, asdict
 from roms_tools.setup.grid import Grid
 from roms_tools.setup.plot import _plot
-from roms_tools.setup.fill import fill_and_interpolate
 from roms_tools.setup.datasets import TPXODataset
 from roms_tools.setup.utils import (
     nan_check,
@@ -82,17 +81,12 @@ class TidalForcing(ROMSToolsMixins):
             straddle=straddle,
         )
 
-        tides = self._get_corrected_tides(data)
-
         # select desired number of constituents
-        for k in tides.keys():
-            tides[k] = tides[k].isel(ntides=slice(None, self.ntides))
+        object.__setattr__(data, "ds", data.ds.isel(ntides=slice(None, self.ntides)))
 
-        # interpolate onto desired grid
-        coords = {"latitude": lat, "longitude": lon}
-        mask = xr.where(data.ds.depth > 0, 1, 0)
+        self._correct_tides(data)
 
-        varnames = [
+        vars_2d = [
             "ssh_Re",
             "ssh_Im",
             "pot_Re",
@@ -102,16 +96,9 @@ class TidalForcing(ROMSToolsMixins):
             "v_Re",
             "v_Im",
         ]
-        data_vars = {}
-
-        for var in varnames:
-            data_vars[var] = fill_and_interpolate(
-                tides[var],
-                mask,
-                list(coords.keys()),
-                coords,
-                method="linear",
-            )
+        vars_3d = []
+
+        data_vars = super().regrid_data(data, vars_2d, vars_3d, lon, lat)
 
         data_vars = super().process_velocities(
             data_vars, angle, "u_Re", "v_Re", interpolate=False
@@ -132,7 +119,7 @@ class TidalForcing(ROMSToolsMixins):
 
         d_meta = get_variable_metadata()
         ds = self._write_into_dataset(data_vars, d_meta)
-        ds["omega"] = tides["omega"]
+        ds["omega"] = data.ds["omega"]
 
         ds = self._add_global_metadata(ds)
 
@@ -412,13 +399,28 @@ class TidalForcing(ROMSToolsMixins):
         # Create and return an instance of TidalForcing
         return cls(grid=grid, **tidal_forcing_params, use_dask=use_dask)
 
-    def _get_corrected_tides(self, data):
+    def _correct_tides(self, data):
+        """
+        Apply tidal corrections to the dataset.
+        This method corrects the dataset for equilibrium tides, self-attraction and loading (SAL) effects, and
+        adjusts phases and amplitudes of tidal elevations and transports using Egbert's correction.
+        Parameters
+        ----------
+        data : Dataset
+            The dataset containing tidal data, including variables for sea surface height (ssh), zonal and meridional
+            currents (u, v), and self-attraction and loading corrections (sal).
+        Returns
+        -------
+        None
+            The dataset is modified in-place with corrected real and imaginary components for ssh, u, v, and the
+            potential field ('pot_Re', 'pot_Im').
+        """
 
         # Get equilibrium tides
         tpc = compute_equilibrium_tide(
             data.ds[data.dim_names["longitude"]], data.ds[data.dim_names["latitude"]]
         )
-        tpc = tpc.isel(**{data.dim_names["ntides"]: data.ds[data.dim_names["ntides"]]})
+        tpc = tpc.isel(ntides=data.ds["ntides"])
         # Correct for SAL
         tsc = self.allan_factor * (
             data.ds[data.var_names["sal_Re"]] + 1j * data.ds[data.var_names["sal_Im"]]
@@ -432,9 +434,9 @@ class TidalForcing(ROMSToolsMixins):
 
         # Apply correction for phases and amplitudes
         pf, pu, aa = egbert_correction(self.model_reference_date)
-        pf = pf.isel(**{data.dim_names["ntides"]: data.ds[data.dim_names["ntides"]]})
-        pu = pu.isel(**{data.dim_names["ntides"]: data.ds[data.dim_names["ntides"]]})
-        aa = aa.isel(**{data.dim_names["ntides"]: data.ds[data.dim_names["ntides"]]})
+        pf = pf.isel(ntides=data.ds["ntides"])
+        pu = pu.isel(ntides=data.ds["ntides"])
+        aa = aa.isel(ntides=data.ds["ntides"])
 
         dt = (self.model_reference_date - data.reference_date).days * 3600 * 24
 
@@ -443,22 +445,18 @@ class TidalForcing(ROMSToolsMixins):
         tvc = pf * tvc * np.exp(1j * (data.ds["omega"] * dt + pu + aa))
         tpc = pf * tpc * np.exp(1j * (data.ds["omega"] * dt + pu + aa))
 
-        tides = {
-            "ssh_Re": thc.real,
-            "ssh_Im": thc.imag,
-            "u_Re": tuc.real,
-            "u_Im": tuc.imag,
-            "v_Re": tvc.real,
-            "v_Im": tvc.imag,
-            "pot_Re": tpc.real,
-            "pot_Im": tpc.imag,
-            "omega": data.ds["omega"],
-        }
-
-        for k in tides.keys():
-            tides[k] = tides[k].rename({data.dim_names["ntides"]: "ntides"})
+        data.ds[data.var_names["ssh_Re"]] = thc.real
+        data.ds[data.var_names["ssh_Im"]] = thc.imag
+        data.ds[data.var_names["u_Re"]] = tuc.real
+        data.ds[data.var_names["u_Im"]] = tuc.imag
+        data.ds[data.var_names["v_Re"]] = tvc.real
+        data.ds[data.var_names["v_Im"]] = tvc.imag
+        data.ds["pot_Re"] = tpc.real
+        data.ds["pot_Im"] = tpc.imag
 
-        return tides
+        # Update var_names dictionary
+        var_names = {**data.var_names, "pot_Re": "pot_Re", "pot_Im": "pot_Im"}
+        object.__setattr__(data, "var_names", var_names)
 
 
 def modified_julian_days(year, month, day, hour=0):
@@ -611,7 +609,7 @@ def egbert_correction(date):
     pf[12] = pftmp**2  # Ms4
     pf[13] = pftmp  # 2n2
     pf[14] = 1.0  # S1
-    pf = xr.DataArray(pf, dims="nc")
+    pf = xr.DataArray(pf, dims="ntides")
 
     putmp = (
         np.arctan(
@@ -647,7 +645,7 @@ def egbert_correction(date):
     pu[12] = putmp  # Ms4
     pu[13] = putmp  # 2n2
     pu[14] = 0.0  # S1
-    pu = xr.DataArray(pu, dims="nc")
+    pu = xr.DataArray(pu, dims="ntides")
     # convert from degrees to radians
     pu = pu * rad
 
@@ -671,7 +669,7 @@ def egbert_correction(date):
                 0.0,  # S1
             ]
         ),
-        dims="nc",
+        dims="ntides",
     )
 
     return pf, pu, aa
@@ -725,7 +723,7 @@ def compute_equilibrium_tide(lon, lat):
                 0.000764,  # S1
             ]
         ),
-        dims="nc",
+        dims="ntides",
     )
     B = xr.DataArray(
         data=np.array(
@@ -747,12 +745,12 @@ def compute_equilibrium_tide(lon, lat):
                 0.693,  # S1
             ]
         ),
-        dims="nc",
+        dims="ntides",
     )
 
     # types: 2 = semidiurnal, 1 = diurnal, 0 = long-term
     ityp = xr.DataArray(
-        data=np.array([2, 2, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0, 0, 2, 1]), dims="nc"
+        data=np.array([2, 2, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0, 0, 2, 1]), dims="ntides"
     )
 
     d2r = np.pi / 180

roms_tools/setup/utils.py

@@ -165,27 +165,31 @@ def interpolate_from_rho_to_v(field, method="additive"):
 
 def extrapolate_deepest_to_bottom(field: xr.DataArray, dim: str) -> xr.DataArray:
     """
-    Extrapolate the deepest non-NaN values to the bottom along a specified dimension.
+    Extrapolates the deepest non-NaN values to the bottom along the specified dimension using forward fill.
+
+    This function assumes that the specified dimension is ordered from top to bottom (e.g., a vertical dimension like 'depth').
+    It fills `NaN` values below the deepest valid (non-NaN) entry along the given dimension by carrying forward the last valid value.
 
     Parameters
     ----------
     field : xr.DataArray
-        The input data array containing NaN values that need to be filled. This array
-        should have at least one dimension named by `dim`.
+        The input `xarray.DataArray` containing potential `NaN` values to be filled.
+        This array must have at least one dimension corresponding to `dim`, typically
+        a vertical axis such as 'depth' or 'height'.
     dim : str
-        The name of the dimension along which to perform the interpolation and extrapolation.
-        Typically, this would be a vertical dimension such as 'depth' or 's_rho'.
+        The name of the dimension along which to perform the forward fill operation.
+        The function assumes that this dimension is ordered from top to bottom, with
+        larger index values representing deeper or lower levels.
 
     Returns
     -------
-    field_interpolated : xr.DataArray
-        A new data array with NaN values along the specified dimension filled by nearest
-        neighbor interpolation and extrapolation to the bottom. The original data array is not modified.
+    xr.DataArray
+        A new `xarray.DataArray` with the `NaN` values along the specified dimension
+        filled by forward filling the deepest valid values down to the bottom.
+        The original input data remains unmodified.
 
     """
-    field_interpolated = field.interpolate_na(
-        dim=dim, method="nearest", fill_value="extrapolate"
-    )
+    field_interpolated = field.ffill(dim=dim)
 
     return field_interpolated
 
@@ -300,19 +304,19 @@ def interpolate_from_climatology(
         raise TypeError("Input 'field' must be an xarray.DataArray or xarray.Dataset.")
 
 
-def is_cftime_datetime(data_array: xr.DataArray) -> bool:
+def get_time_type(data_array: xr.DataArray) -> str:
     """
-    Checks if the xarray DataArray contains cftime datetime objects.
+    Determines the type of time values in the xarray DataArray.
 
     Parameters
     ----------
     data_array : xr.DataArray
-        The xarray DataArray to be checked for cftime datetime objects.
+        The xarray DataArray to be checked for time data types.
 
     Returns
     -------
-    bool
-        True if the DataArray contains cftime datetime objects, False otherwise.
+    str
+        A string indicating the type of the time data: 'cftime', 'datetime', or 'int'.
 
     Raises
     ------
@@ -324,16 +328,26 @@ def is_cftime_datetime(data_array: xr.DataArray) -> bool:
         cftime.DatetimeNoLeap,
         cftime.DatetimeJulian,
         cftime.DatetimeGregorian,
+        cftime.Datetime360Day,
+        cftime.DatetimeProlepticGregorian,
     )
 
-    # Check if any of the coordinate values are of cftime type
+    # Check if any of the coordinate values are of cftime, datetime, or integer type
     if isinstance(data_array.values, (np.ndarray, list)):
-        # Check the dtype of the array; numpy datetime64 indicates it's not cftime
+        # Check if the data type is numpy datetime64, indicating standard datetime objects
         if data_array.values.dtype == "datetime64[ns]":
-            return False
+            return "datetime"
 
-        # Check if any of the values in the array are instances of cftime types
-        return any(isinstance(value, cftime_types) for value in data_array.values)
+        # Check if any values in the array are instances of cftime types
+        if any(isinstance(value, cftime_types) for value in data_array.values):
+            return "cftime"
+
+        # Check if all values are of integer type (e.g., for indices or time steps)
+        if np.issubdtype(data_array.values.dtype, np.integer):
+            return "int"
+
+        # If none of the above conditions are met, raise a ValueError
+        raise ValueError("Unsupported data type for time values in input dataset.")
 
     # Handle unexpected types
     raise TypeError("DataArray values must be of type numpy.ndarray or list.")
@@ -599,8 +613,12 @@ def group_dataset(ds, filepath):
         if len(ds["abs_time"]) > 2:
             # Determine the frequency of the data
             abs_time_freq = pd.infer_freq(ds["abs_time"].to_index())
-
-            if abs_time_freq in ["D", "H", "T", "S"]:  # Daily or higher frequency
+            if abs_time_freq.lower() in [
+                "d",
+                "h",
+                "t",
+                "s",
+            ]:  # Daily or higher frequency
                 dataset_list, output_filenames = group_by_month(ds, filepath)
             else:
                 dataset_list, output_filenames = group_by_year(ds, filepath)

roms_tools/tests/test_setup/test_boundary_forcing.py

@@ -199,7 +199,7 @@ def test_roundtrip_yaml(bdry_forcing_fixture, request, tmp_path, use_dask):
 
 def test_files_have_same_hash(boundary_forcing, tmp_path, use_dask):
 
-    yaml_filepath = tmp_path / "test_yaml"
+    yaml_filepath = tmp_path / "test_yaml_.yaml"
     filepath1 = tmp_path / "test1.nc"
     filepath2 = tmp_path / "test2.nc"
 

roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/.zattrs

@@ -3,7 +3,7 @@
     "end_time": "2021-06-30 00:00:00",
     "hc": 250.0,
     "model_reference_date": "2000-01-01 00:00:00",
-    "roms_tools_version": "0.1.dev98",
+    "roms_tools_version": "0.1.dev152+dirty",
     "source": "CESM_REGRIDDED",
     "start_time": "2021-06-29 00:00:00",
     "theta_b": 2.0,

roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/.zmetadata

@@ -5,7 +5,7 @@
             "end_time": "2021-06-30 00:00:00",
             "hc": 250.0,
             "model_reference_date": "2000-01-01 00:00:00",
-            "roms_tools_version": "0.1.dev98",
+            "roms_tools_version": "0.1.dev152+dirty",
             "source": "CESM_REGRIDDED",
             "start_time": "2021-06-29 00:00:00",
             "theta_b": 2.0,