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

roms_tools/setup/grid.py

@@ -137,7 +137,6 @@ class Grid:
         # Check if the Greenwich meridian goes through the domain.
         self._straddle()
 
-        ds = _add_lat_lon_at_velocity_points(self.ds, self.straddle)
         object.__setattr__(self, "ds", ds)
 
         # Update the grid by adding grid variables that are coarsened versions of the original grid variables
@@ -805,7 +804,6 @@ def _make_grid_ds(
 
     # rotate coordinate system
     rotated_lon_lat_vars = _rotate(*initial_lon_lat_vars, rot)
-    lon, *_ = rotated_lon_lat_vars
 
     # translate coordinate system
     translated_lon_lat_vars = _translate(*rotated_lon_lat_vars, center_lat, center_lon)
@@ -817,7 +815,28 @@ def _make_grid_ds(
     # compute angle of local grid positive x-axis relative to east
     ang = _compute_angle(lon, lonu, latu, lonq)
 
-    ds = _create_grid_ds(lon, lat, pm, pn, ang, rot, center_lon, center_lat)
+    # make sure lons are in [0, 360] range
+    lon[lon < 0] = lon[lon < 0] + 2 * np.pi
+    lonu[lonu < 0] = lonu[lonu < 0] + 2 * np.pi
+    lonv[lonv < 0] = lonv[lonv < 0] + 2 * np.pi
+    lonq[lonq < 0] = lonq[lonq < 0] + 2 * np.pi
+
+    ds = _create_grid_ds(
+        lon,
+        lat,
+        lonu,
+        latu,
+        lonv,
+        latv,
+        lonq,
+        latq,
+        pm,
+        pn,
+        ang,
+        rot,
+        center_lon,
+        center_lat,
+    )
 
     ds = _add_global_metadata(ds, size_x, size_y, center_lon, center_lat, rot)
 
@@ -1096,15 +1115,18 @@ def _compute_angle(lon, lonu, latu, lonq):
     ang[:, 0] = ang[:, 1]
     ang[:, -1] = ang[:, -2]
 
-    lon[lon < 0] = lon[lon < 0] + 2 * np.pi
-    lonq[lonq < 0] = lonq[lonq < 0] + 2 * np.pi
-
     return ang
 
 
 def _create_grid_ds(
     lon,
     lat,
+    lonu,
+    latu,
+    lonv,
+    latv,
+    lonq,
+    latq,
     pm,
     pn,
     angle,
@@ -1124,8 +1146,49 @@ def _create_grid_ds(
         dims=["eta_rho", "xi_rho"],
         attrs={"long_name": "latitude of rho-points", "units": "degrees North"},
     )
+    lon_u = xr.Variable(
+        data=lonu * 180 / np.pi,
+        dims=["eta_rho", "xi_u"],
+        attrs={"long_name": "longitude of u-points", "units": "degrees East"},
+    )
+    lat_u = xr.Variable(
+        data=latu * 180 / np.pi,
+        dims=["eta_rho", "xi_u"],
+        attrs={"long_name": "latitude of u-points", "units": "degrees North"},
+    )
+    lon_v = xr.Variable(
+        data=lonv * 180 / np.pi,
+        dims=["eta_v", "xi_rho"],
+        attrs={"long_name": "longitude of v-points", "units": "degrees East"},
+    )
+    lat_v = xr.Variable(
+        data=latv * 180 / np.pi,
+        dims=["eta_v", "xi_rho"],
+        attrs={"long_name": "latitude of v-points", "units": "degrees North"},
+    )
+    lon_q = xr.Variable(
+        data=lonq * 180 / np.pi,
+        dims=["eta_psi", "xi_psi"],
+        attrs={"long_name": "longitude of psi-points", "units": "degrees East"},
+    )
+    lat_q = xr.Variable(
+        data=latq * 180 / np.pi,
+        dims=["eta_psi", "xi_psi"],
+        attrs={"long_name": "latitude of psi-points", "units": "degrees North"},
+    )
 
-    ds = ds.assign_coords({"lat_rho": lat_rho, "lon_rho": lon_rho})
+    ds = ds.assign_coords(
+        {
+            "lat_rho": lat_rho,
+            "lon_rho": lon_rho,
+            "lat_u": lat_u,
+            "lon_u": lon_u,
+            "lat_v": lat_v,
+            "lon_v": lon_v,
+            "lat_psi": lat_q,
+            "lon_psi": lon_q,
+        }
+    )
 
     ds["angle"] = xr.Variable(
         data=angle,
@@ -1242,7 +1305,28 @@ def _f2c_xdir(f):
 
 
 def _add_lat_lon_at_velocity_points(ds, straddle):
-
+    """
+    Adds latitude and longitude coordinates at velocity points (u and v points) to the dataset.
+    This function computes approximate latitude and longitude values at u and v velocity points
+    based on the rho points (cell centers). If the grid straddles the Greenwich meridian, it adjusts
+    the longitudes to avoid jumps from 360 to 0 degrees. The computed coordinates are added to the
+    dataset as new variables with appropriate metadata.
+    Parameters
+    ----------
+    ds : xarray.Dataset
+        The input dataset containing rho point coordinates ("lat_rho", "lon_rho").
+    straddle : bool
+        Indicates whether the grid straddles the Greenwich meridian. If True, longitudes are adjusted
+        to avoid discontinuities.
+    Returns
+    -------
+    ds : xarray.Dataset
+        The dataset with added coordinates for u and v points ("lat_u", "lon_u", "lat_v", "lon_v").
+    Notes
+    -----
+    This function only computes approximate latitude and longitude values. It should only be used if
+    more accurate values are not available from grid generation.
+    """
     if straddle:
         # avoid jump from 360 to 0 in interpolation
         lon_rho = xr.where(ds["lon_rho"] > 180, ds["lon_rho"] - 360, ds["lon_rho"])

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/topography.py

@@ -61,6 +61,10 @@ def _add_topography_and_mask(
 
     # fill enclosed basins with land
     mask = _fill_enclosed_basins(mask.values)
+
+    # adjust mask boundaries by copying values from adjacent cells
+    mask = _handle_boundaries(mask)
+
     ds["mask_rho"] = xr.DataArray(mask.astype(np.int32), dims=("eta_rho", "xi_rho"))
     ds["mask_rho"].attrs = {
         "long_name": "Mask at rho-points",
@@ -233,10 +237,7 @@ def _smooth_topography_locally(h, hmin=5, rmax=0.2):
         )
 
         # No gradient at the domain boundaries
-        h_log[0, :] = h_log[1, :]
-        h_log[-1, :] = h_log[-2, :]
-        h_log[:, 0] = h_log[:, 1]
-        h_log[:, -1] = h_log[:, -2]
+        h_log = _handle_boundaries(h_log)
 
         # Update h
         h = hmin * np.exp(h_log)
@@ -252,6 +253,31 @@ def _smooth_topography_locally(h, hmin=5, rmax=0.2):
     return h
 
 
+def _handle_boundaries(field):
+    """
+    Adjust the boundaries of a 2D field by copying values from adjacent cells.
+
+    Parameters
+    ----------
+    field : numpy.ndarray or xarray.DataArray
+        A 2D array representing a field (e.g., topography or mask) whose boundary values
+        need to be adjusted.
+
+    Returns
+    -------
+    field : numpy.ndarray or xarray.DataArray
+        The input field with adjusted boundary values.
+
+    """
+
+    field[0, :] = field[1, :]
+    field[-1, :] = field[-2, :]
+    field[:, 0] = field[:, 1]
+    field[:, -1] = field[:, -2]
+
+    return field
+
+
 def _compute_rfactor(h):
     """
     Computes slope parameter (or r-factor) r = |Delta h| / 2h in both horizontal grid directions.