roms-tools 2.6.1__py3-none-any.whl → 2.7.0__py3-none-any.whl

roms_tools/setup/datasets.py

@@ -6,10 +6,12 @@ import numpy as np
 from typing import Dict, Optional, Union, List
 from pathlib import Path
 import logging
+from roms_tools.constants import R_EARTH
 from roms_tools.utils import _load_data
 from roms_tools.setup.utils import (
     assign_dates_to_climatology,
     interpolate_from_climatology,
+    interpolate_cyclic_time,
     get_time_type,
     convert_cftime_to_datetime,
     one_dim_fill,
@@ -19,6 +21,7 @@ from roms_tools.download import (
     download_correction_data,
     download_topo,
     download_river_data,
+    download_sal_data,
 )
 from roms_tools.setup.fill import LateralFill
 
@@ -44,8 +47,15 @@ class Dataset:
         Dictionary specifying the names of dimensions in the dataset.
     var_names: Dict[str, str]
         Dictionary of variable names that are required in the dataset.
+    opt_var_names: Dict[str, str], optional
+        Dictionary of variable names that are optional in the dataset.
+        Defaults to an empty dictionary.
     climatology : bool
         Indicates whether the dataset is climatological. Defaults to False.
+    needs_lateral_fill: bool, optional
+        Indicates whether land values require lateral filling. If `True`, ocean values will be extended into land areas
+        to replace NaNs or non-ocean values (such as atmospheric values in ERA5 data). If `False`, it is assumed that
+        land values are already correctly assigned, and lateral filling will be skipped. Defaults to `True`.
     use_dask: bool
         Indicates whether to use dask for chunking. If True, data is loaded with dask; if False, data is loaded eagerly. Defaults to False.
     apply_post_processing: bool
@@ -78,7 +88,9 @@ class Dataset:
         }
     )
     var_names: Dict[str, str]
+    opt_var_names: Optional[Dict[str, str]] = field(default_factory=dict)
     climatology: Optional[bool] = False
+    needs_lateral_fill: Optional[bool] = True
     use_dask: Optional[bool] = False
     apply_post_processing: Optional[bool] = True
 
@@ -109,6 +121,9 @@ class Dataset:
         ds = self.clean_up(ds)
         self.check_dataset(ds)
 
+        # Select relevant fields
+        ds = self.select_relevant_fields(ds)
+
         # Select relevant times
         if "time" in self.dim_names and self.start_time is not None:
             ds = self.add_time_info(ds)
@@ -117,9 +132,6 @@ class Dataset:
             if self.dim_names["time"] != "time":
                 ds = ds.rename({self.dim_names["time"]: "time"})
 
-        # Select relevant fields
-        ds = self.select_relevant_fields(ds)
-
         # Make sure that latitude is ascending
         ds = self.ensure_dimension_is_ascending(ds, dim="latitude")
         # Make sure there are no 360 degree jumps in longitude
@@ -158,7 +170,7 @@ class Dataset:
 
         return ds
 
-    def clean_up(self, ds: xr.Dataset) -> xr.Dataset:
+    def clean_up(self, ds: xr.Dataset, **kwargs) -> xr.Dataset:
         """Dummy method to be overridden by child classes to clean up the dataset.
 
         This method is intended as a placeholder and should be implemented in subclasses
@@ -172,9 +184,9 @@ class Dataset:
         Returns
         -------
         xr.Dataset
-            The xarray Dataset cleaned up (as implemented by child classes).
+            The cleaned-up xarray Dataset (as implemented by child classes).
         """
-        return ds
+        return ds  # Default behavior (no-op, subclasses should override)
 
     def check_dataset(self, ds: xr.Dataset) -> None:
         """Check if the dataset contains the specified variables and dimensions.
@@ -208,7 +220,11 @@ class Dataset:
         """
 
         for var in ds.data_vars:
-            if var not in self.var_names.values():
+            if (
+                var not in self.var_names.values()
+                and var not in self.opt_var_names.values()
+                and var != "mask"
+            ):
                 ds = ds.drop_vars(var)
 
         return ds
@@ -372,7 +388,6 @@ class Dataset:
             and longitude differences, in meters.
         """
 
-        r_earth = 6371315.0
         lat_dim = self.dim_names["latitude"]
         lon_dim = self.dim_names["longitude"]
 
@@ -385,11 +400,11 @@ class Dataset:
         lon_diff = np.abs(np.diff(longitudes)).min()  # Minimal longitude spacing
 
         # Latitude spacing is constant at all longitudes
-        min_lat_spacing = (2 * np.pi * r_earth * lat_diff) / 360
+        min_lat_spacing = (2 * np.pi * R_EARTH * lat_diff) / 360
 
         # Longitude spacing varies with latitude
         min_lon_spacing = (
-            2 * np.pi * r_earth * lon_diff * np.cos(np.radians(latitudes.min()))
+            2 * np.pi * R_EARTH * lon_diff * np.cos(np.radians(latitudes.min()))
         ) / 360
 
         # The minimal spacing is the smaller of the two
@@ -591,6 +606,7 @@ class Dataset:
         lon = subdomain[self.dim_names["longitude"]]
 
         if self.is_global:
+            concats = []
             # Concatenate only if necessary
             if lon_max + margin > lon.max():
                 # See if shifting by +360 degrees helps
@@ -600,10 +616,7 @@ class Dataset:
                     subdomain[self.dim_names["longitude"]] = lon + 360
                     lon = subdomain[self.dim_names["longitude"]]
                 else:
-                    subdomain = self.concatenate_longitudes(
-                        subdomain, end="upper", verbose=verbose
-                    )
-                    lon = subdomain[self.dim_names["longitude"]]
+                    concats.append("upper")
             if lon_min - margin < lon.min():
                 # See if shifting by -360 degrees helps
                 if (lon_min - margin > (lon - 360).min()) and (
@@ -612,10 +625,14 @@ class Dataset:
                     subdomain[self.dim_names["longitude"]] = lon - 360
                     lon = subdomain[self.dim_names["longitude"]]
                 else:
-                    subdomain = self.concatenate_longitudes(
-                        subdomain, end="lower", verbose=verbose
-                    )
-                    lon = subdomain[self.dim_names["longitude"]]
+                    concats.append("lower")
+
+            if concats:
+                end = "both" if len(concats) == 2 else concats[0]
+                subdomain = self.concatenate_longitudes(
+                    subdomain, end=end, verbose=False
+                )
+                lon = subdomain[self.dim_names["longitude"]]
 
         else:
             # Adjust longitude range if needed to match the expected range
@@ -639,6 +656,7 @@ class Dataset:
                         lon_min += 360
                         lon_max += 360
         # Select the subdomain in longitude direction
+
         subdomain = subdomain.sel(
             **{
                 self.dim_names["longitude"]: slice(lon_min - margin, lon_max + margin),
@@ -682,38 +700,54 @@ class Dataset:
 
         Notes
         -----
+        This method assumes that the variables in the dataset use a dimension
+        ordering where latitude comes before longitude, i.e., ('latitude', 'longitude').
+        Ensure that this convention is followed to avoid unexpected behavior.
+
         Looping over `self.ds.data_vars` instead of `self.var_names` ensures that each
         dataset variable is filled only once, even if multiple entries in `self.var_names`
         point to the same variable in the dataset.
         """
-        lateral_fill = LateralFill(
-            self.ds["mask"],
-            [self.dim_names["latitude"], self.dim_names["longitude"]],
-        )
 
-        separate_fill_for_velocities = False
-        if "mask_vel" in self.ds.data_vars:
-            lateral_fill_vel = LateralFill(
-                self.ds["mask_vel"],
+        if self.needs_lateral_fill:
+            logging.info(
+                "Applying 2D horizontal fill to the source data before regridding."
+            )
+
+            lateral_fill = LateralFill(
+                self.ds["mask"],
                 [self.dim_names["latitude"], self.dim_names["longitude"]],
             )
-            separate_fill_for_velocities = True
-
-        for var_name in self.ds.data_vars:
-            if var_name.startswith("mask"):
-                # Skip variables that are mask types
-                continue
-            elif (
-                separate_fill_for_velocities
-                and "u" in self.var_names
-                and "v" in self.var_names
-                and var_name in [self.var_names["u"], self.var_names["v"]]
-            ):
-                # Apply lateral fill with velocity mask for velocity variables if present
-                self.ds[var_name] = lateral_fill_vel.apply(self.ds[var_name])
-            else:
-                # Apply standard lateral fill for other variables
-                self.ds[var_name] = lateral_fill.apply(self.ds[var_name])
+
+            separate_fill_for_velocities = False
+            # TODO: Replace hardcoded mask detection with a dictionary-based mask selection.
+            # This dictionary could assign which mask to use for what fields.
+            if "mask_vel" in self.ds.data_vars:
+                lateral_fill_vel = LateralFill(
+                    self.ds["mask_vel"],
+                    [self.dim_names["latitude"], self.dim_names["longitude"]],
+                )
+                separate_fill_for_velocities = True
+
+            for var_name in self.ds.data_vars:
+                if var_name.startswith("mask"):
+                    # Skip variables that are mask types
+                    continue
+                elif (
+                    separate_fill_for_velocities
+                    and "u" in self.var_names
+                    and "v" in self.var_names
+                    and var_name in [self.var_names["u"], self.var_names["v"]]
+                ):
+                    # Apply lateral fill with velocity mask for velocity variables if present
+                    self.ds[var_name] = lateral_fill_vel.apply(self.ds[var_name])
+                else:
+                    # Apply standard lateral fill for other variables
+                    self.ds[var_name] = lateral_fill.apply(self.ds[var_name])
+        else:
+            logging.info(
+                "2D horizontal fill is skipped because source data already contains filled values."
+            )
 
     def extrapolate_deepest_to_bottom(self):
         """Extrapolate deepest non-NaN values to fill bottom NaNs along the depth
@@ -775,169 +809,226 @@ class TPXODataset(Dataset):
     ----------
     filename : str
         The path to the TPXO dataset file.
-    var_names : Dict[str, str], optional
-        Dictionary of variable names required in the dataset. Defaults to:
-        {
-            "h_Re": "h_Re",
-            "h_Im": "h_Im",
-            "sal_Re": "sal_Re",
-            "sal_Im": "sal_Im",
-            "u_Re": "u_Re",
-            "u_Im": "u_Im",
-            "v_Re": "v_Re",
-            "v_Im": "v_Im",
-            "depth": "depth"
-        }
+    grid_filename : str
+        The path to the TPXO grid file.
+    location : str
+        "h", "u", "v"
+    var_names : Dict[str, str]
+        Dictionary of variable names required in the dataset.
     dim_names : Dict[str, str], optional
         Dictionary specifying the names of dimensions in the dataset. Defaults to:
         {"longitude": "ny", "latitude": "nx", "ntides": "nc"}.
+    tolerate_coord_mismatch : bool, optional
+        If True, allows mismatched latitude/longitude coordinates between the TPXO dataset and the TPXO grid
+        by selecting the nearest grid-aligned coordinates. Default is False.
 
     Attributes
     ----------
     ds : xr.Dataset
         The xarray Dataset containing the TPXO tidal model data, loaded from the specified file.
-    reference_date : datetime
-        The reference date for the TPXO data. Default is datetime(1992, 1, 1).
     """
 
     filename: str
-    var_names: Dict[str, str] = field(
-        default_factory=lambda: {
-            "ssh_Re": "h_Re",
-            "ssh_Im": "h_Im",
-            "sal_Re": "sal_Re",
-            "sal_Im": "sal_Im",
-            "u_Re": "u_Re",
-            "u_Im": "u_Im",
-            "v_Re": "v_Re",
-            "v_Im": "v_Im",
-            "depth": "depth",
-        }
-    )
+    grid_filename: str
+    location: str
+    var_names: Dict[str, str]
     dim_names: Dict[str, str] = field(
-        default_factory=lambda: {"longitude": "ny", "latitude": "nx", "ntides": "nc"}
+        default_factory=lambda: {"longitude": "nx", "latitude": "ny", "ntides": "nc"}
     )
+    tolerate_coord_mismatch: bool = False
     ds: xr.Dataset = field(init=False, repr=False)
-    reference_date: datetime = datetime(1992, 1, 1)
 
     def clean_up(self, ds: xr.Dataset) -> xr.Dataset:
-        """Clean up and standardize the dimensions and coordinates of the dataset for
-        further processing.
-
-        This method performs the following operations:
-        - Assigns new coordinate variables for 'omega', 'longitude', and 'latitude' based on existing dataset variables.
-          - 'omega' is retained as it is.
-          - 'longitude' is derived from 'lon_r', assuming it is constant along the 'ny' dimension.
-          - 'latitude' is derived from 'lat_r', assuming it is constant along the 'nx' dimension.
-        - Renames the dimensions 'nx' and 'ny' to 'longitude' and 'latitude', respectively, for consistency.
-        - Renames the tidal dimension to 'ntides' for standardization.
-        - Updates the `dim_names` attribute of the object to reflect the new dimension names: 'longitude', 'latitude', and 'ntides'.
+        """Standardize the dataset's dimensions and coordinates for further processing.
+
+            This method performs the following operations:
+            - Assigns 'longitude' and 'latitude' coordinates from existing dataset variables.
+            - Adds a `mask` variable indicating valid data points based on grid conditions.
+            - Renames dimensions:
+              - 'nx' → 'longitude'
+              - 'ny' → 'latitude'
+              - Tidal dimension → 'ntides'
+            - Updates `dim_names` to reflect the new dimension names.
+            - If coordinates do not match and `self.tolerate_coord_mismatch=True`, the dataset is reindexed
+        to the grid's coordinates using nearest-neighbor selection. Otherwise, a ValueError is raised.
+
+            Parameters
+            ----------
+            ds : xr.Dataset
+                The dataset to be standardized.
+
+            Returns
+            -------
+            xr.Dataset
+                The cleaned dataset with updated coordinates and dimensions.
+
+            Raises
+            ------
+            ValueError
+                If longitude or latitude values do not match the grid.
+        """
 
-        Parameters
-        ----------
-        ds : xr.Dataset
-            The input dataset to be cleaned and standardized. It should contain the coordinates 'omega', 'lon_r', 'lat_r', and the tidal dimension.
+        ds_grid = _load_data(self.grid_filename, self.dim_names, self.use_dask)
+
+        # Define mask and coordinate names based on location
+        if self.location == "h":
+            mask_name = "mz"
+            lon_name = "lon_z"
+            lat_name = "lat_z"
+        elif self.location == "u":
+            mask_name = "mu"
+            lon_name = "lon_u"
+            lat_name = "lat_u"
+        elif self.location == "v":
+            mask_name = "mv"
+            lon_name = "lon_v"
+            lat_name = "lat_v"
+
+        # Assign and rename coordinates
+        ds_grid = ds_grid.assign_coords(
+            {
+                "nx": ds_grid[lon_name].isel(
+                    ny=0
+                ),  # lon is constant along ny, i.e., is only a function of nx
+                "ny": ds_grid[lat_name].isel(
+                    nx=0
+                ),  # lat is constant along nx, i.e., is only a function of ny
+            }
+        )
+        ds_grid = ds_grid.rename({"nx": "longitude", "ny": "latitude"})
 
-        Returns
-        -------
-        ds : xr.Dataset
-            A cleaned and standardized `xarray.Dataset` with updated coordinates and dimensions.
-        """
+        # Drop all dimensions except 'longitude' and 'latitude'
+        dims_to_keep = {"longitude", "latitude"}
+        dims_to_drop = [dim for dim in ds_grid.dims if dim not in dims_to_keep]
+        if dims_to_drop:
+            ds_grid = ds_grid.isel({dim: 0 for dim in dims_to_drop})
+
+        # Ensure correct dimension order
+        ds_grid = ds_grid.transpose("latitude", "longitude")
+
+        dims_to_keep = {"longitude", "latitude"}
+        dims_to_drop = set(ds_grid.dims) - dims_to_keep
+        ds_grid = (
+            ds_grid.isel({dim: 0 for dim in dims_to_drop}) if dims_to_drop else ds_grid
+        )
+        # Bring dimensions in correct order
+        ds_grid = ds_grid.transpose("latitude", "longitude")
+
+        ds = ds.rename({"con": "nc"})
         ds = ds.assign_coords(
             {
-                "omega": ds["omega"],
-                "nx": ds["lon_r"].isel(
+                "nc": ("nc", [c.strip() for c in ds["nc"].values]),  # Strip padding
+                "nx": ds[lon_name].isel(
                     ny=0
-                ),  # lon_r is constant along ny, i.e., is only a function of nx
-                "ny": ds["lat_r"].isel(
+                ),  # lon is constant along ny, i.e., is only a function of nx
+                "ny": ds[lat_name].isel(
                     nx=0
-                ),  # lat_r is constant along nx, i.e., is only a function of ny
+                ),  # lat is constant along nx, i.e., is only a function of ny
             }
         )
         ds = ds.rename(
             {"nx": "longitude", "ny": "latitude", self.dim_names["ntides"]: "ntides"}
         )
 
+        ds = ds.transpose("latitude", "longitude", "ntides")
+
         self.dim_names = {
             "latitude": "latitude",
             "longitude": "longitude",
             "ntides": "ntides",
         }
 
+        if self.tolerate_coord_mismatch:
+            # Reindex dataset to match grid lat/lon using nearest-neighbor method
+            ds = ds.sel(
+                {
+                    "latitude": ds_grid["latitude"],
+                    "longitude": ds_grid["longitude"],
+                },
+                method="nearest",
+            )
+
+        # Validate matching lat/lon shapes and values
+        for coord in [lon_name, lat_name]:
+            if ds[coord].shape != ds_grid[coord].shape:
+                raise ValueError(
+                    f"Mismatch in {coord} array sizes. Dataset: {ds[coord].shape}, Grid: {ds_grid[coord].shape}"
+                )
+            if not np.allclose(ds[coord].values, ds_grid[coord].values):
+                raise ValueError(
+                    f"{coord.capitalize()} values from dataset do not match grid. Dataset: {ds[coord].values}, Grid: {ds_grid[coord].values}"
+                )
+
+        # Add mask
+        mask = ds_grid[mask_name].isnull()
+        # Create a fresh xarray.DataArray for the mask using only the dimension names.
+        # This avoids issues that can arise from small coordinate mismatches when assigning directly.
+        ds["mask"] = xr.DataArray(
+            ds_grid[mask_name].isnull().values, dims=list(mask.dims)
+        )
+
         return ds
 
-    def check_number_constituents(self, ntides: int):
-        """Checks if the number of constituents in the dataset is at least `ntides`.
+    def select_constituents(self, ntides: int, omega: Dict[str, float]):
+        """Selects the first `ntides` tidal constituents based on the provided omega
+        values.
+
+        This method filters the dataset to retain only the tidal constituents that match
+        the first `ntides` from the provided `omega` dictionary. It ensures that the dataset
+        contains the expected constituents before proceeding with the selection. If the dataset
+        does not contain the required constituents, an error is raised.
 
         Parameters
         ----------
         ntides : int
-            The required number of tidal constituents.
+            The number of tidal constituents to retain from the omega values. The method selects
+            the first `ntides` constituents based on the provided omega dictionary and maps their
+            corresponding values to the dataset.
+
+        omega : dict
+            A dictionary where keys are tidal constituent names and values are their associated omega
+            values. The first `ntides` keys from this dictionary will be used to filter the tidal
+            constituents in the dataset.
 
         Raises
         ------
         ValueError
-            If the number of constituents in the dataset is less than `ntides`.
+            If the dataset does not contain all required tidal constituents from the first `ntides`
+            selected from the `omega` dictionary, a `ValueError` is raised, indicating the mismatch
+            between the expected and present constituents in the dataset.
         """
-        if len(self.ds[self.dim_names["ntides"]]) < ntides:
-            raise ValueError(
-                f"The dataset contains fewer than {ntides} tidal constituents."
-            )
 
-    def post_process(self):
-        """Apply a depth-based mask to the dataset, ensuring only positive depths are
-        retained.
+        # Expected constituents based on the first 'ntides' from the omega dictionary
+        expected_constituents = list(omega.keys())[:ntides]
 
-        This method checks if the 'depth' variable is present in the dataset. If found, a mask is created where
-        depths greater than 0 are considered valid (mask value of 1). This mask is applied to all data variables
-        in the dataset, replacing values at invalid depths (depth ≤ 0) with NaN. The mask itself is also stored
-        in the dataset under the variable 'mask'.
+        # Extract the current tidal constituents from the dataset
+        dataset_constituents = [
+            c.decode("utf-8").strip() for c in self.ds["ntides"].values
+        ]
 
-        Returns
-        -------
-        None
-            The dataset is modified in-place by applying the mask to each variable.
-        """
+        # Check if the dataset contains the expected constituents
+        if not all(c in dataset_constituents for c in expected_constituents):
+            raise ValueError(
+                f"The dataset contains tidal constituents {dataset_constituents} that do not match the first {ntides} required constituents "
+                f"from the TPXO dataset: {expected_constituents}. "
+                "Ensure the dataset includes the required constituents or reduce the 'ntides' parameter."
+            )
 
-        if "depth" in self.var_names.keys():
-            ds = self.ds
-            mask = xr.where(self.ds["depth"] > 0, 1, 0)
-            ds["mask"] = mask
-            ds = ds.drop_vars(["depth"])
+        # Select only the expected constituents from the dataset
+        filtered_constituents = [
+            c for c in dataset_constituents if c in expected_constituents
+        ]
 
-            self.ds = ds
+        # Encode the filtered constituents back to byte strings before selecting in xarray
+        filtered_constituents_bytes = [c.encode("utf-8") for c in filtered_constituents]
 
-        # Remove "depth" from var_names
-        updated_var_names = {**self.var_names}  # Create a copy of the dictionary
-        updated_var_names.pop("depth", None)  # Remove "depth" if it exists
-        self.var_names = updated_var_names
+        # Update the dataset with the filtered constituents
+        self.ds = self.ds.sel(ntides=filtered_constituents_bytes)
 
 
 @dataclass(kw_only=True)
 class GLORYSDataset(Dataset):
-    """Represents GLORYS data on original grid.
-
-    Parameters
-    ----------
-    filename : str
-        The path to the data files. Can contain wildcards.
-    start_time : Optional[datetime], optional
-        The start time for selecting relevant data. If not provided, the data is not filtered by start time.
-    end_time : Optional[datetime], optional
-        The end time for selecting relevant data. If not provided, only data at the start_time is selected if start_time is provided,
-        or no filtering is applied if start_time is not provided.
-    var_names: Dict[str, str], optional
-        Dictionary of variable names that are required in the dataset.
-    dim_names: Dict[str, str], optional
-        Dictionary specifying the names of dimensions in the dataset.
-    climatology : bool
-        Indicates whether the dataset is climatological. Defaults to False.
-
-    Attributes
-    ----------
-    ds : xr.Dataset
-        The xarray Dataset containing the GLORYS data on its original grid.
-    """
+    """Represents GLORYS data on original grid."""
 
     var_names: Dict[str, str] = field(
         default_factory=lambda: {
@@ -993,31 +1084,159 @@ class GLORYSDataset(Dataset):
 
 
 @dataclass(kw_only=True)
-class CESMDataset(Dataset):
-    """Represents CESM data on original grid.
+class UnifiedDataset(Dataset):
+    """Represents unified BGC data on original grid.
 
-    Parameters
-    ----------
-    filename : str
-        The path to the data files. Can contain wildcards.
-    start_time : Optional[datetime], optional
-        The start time for selecting relevant data. If not provided, the data is not filtered by start time.
-    end_time : Optional[datetime], optional
-        The end time for selecting relevant data. If not provided, only data at the start_time is selected if start_time is provided,
-        or no filtering is applied if start_time is not provided.
-    var_names: Dict[str, str], optional
-        Dictionary of variable names that are required in the dataset.
-    dim_names: Dict[str, str], optional
-        Dictionary specifying the names of dimensions in the dataset.
-    climatology : bool
-        Indicates whether the dataset is climatological. Defaults to False.
-
-    Attributes
-    ----------
-    ds : xr.Dataset
-        The xarray Dataset containing the CESM data on its original grid.
+    Notes
+    -----
+    Pierre has already addressed lateral filling during preprocessing,
+    and since the dataset does not contain a mask, the `needs_lateral_fill`
+    attribute is set to `False`.
     """
 
+    needs_lateral_fill: Optional[bool] = False
+
+    # overwrite clean_up method from parent class
+    def clean_up(self, ds: xr.Dataset) -> xr.Dataset:
+        """Ensure the dataset's time dimension is correctly defined and standardized.
+
+        This method verifies that the time dimension exists in the dataset and assigns it appropriately. If the "time" dimension is missing, the method attempts to assign an existing "time" or "month" dimension. If neither exists, it expands the dataset to include a "time" dimension with a size of one.
+
+        Returns
+        -------
+        ds : xr.Dataset
+            The xarray Dataset with the correct time dimension assigned or added.
+        """
+
+        ds = ds.rename({"lon": "longitude", "lat": "latitude", "dep": "depth"})
+        ds = ds.assign_coords(
+            {
+                "latitude": ds["latitude"],
+                "longitude": ds["longitude"],
+                "depth": ds["depth"],
+            }
+        )
+
+        self.dim_names = {
+            "latitude": "latitude",
+            "longitude": "longitude",
+            "depth": "depth",
+        }
+
+        # Handle time dimension
+        if "time" not in self.dim_names:
+            if "month" in ds.dims or "season" in ds.dims:
+                time_dim = "month" if "month" in ds.dims else "season"
+
+                if time_dim == "month":
+                    # Interpolate season to month so all variables have the same time dimension
+                    for var_name in list(self.var_names.values()) + list(
+                        self.opt_var_names.values()
+                    ):
+                        if var_name in ds.data_vars and "season" in ds[var_name].dims:
+                            ds[var_name] = interpolate_cyclic_time(
+                                ds[var_name],
+                                time_dim_name="season",
+                                day_of_year=ds["month"],
+                            )
+
+                # Rename dimension and convert from float64 days to timedelta
+                ds = ds.rename({time_dim: "time"})
+                self.dim_names["time"] = "time"
+
+                ds["time"] = xr.DataArray(
+                    (ds["time"].values * 86400 * 1e9).astype("timedelta64[ns]"),
+                    dims="time",
+                )
+
+            else:
+                # Handle case where all variables are time-invariant
+                ds = ds.expand_dims(time=1)
+                self.dim_names["time"] = "time"
+
+        return ds
+
+
+@dataclass(kw_only=True)
+class UnifiedBGCDataset(UnifiedDataset):
+    dim_names: Dict[str, str] = field(
+        default_factory=lambda: {
+            "longitude": "lon",
+            "latitude": "lat",
+            "depth": "dep",
+        }
+    )
+    var_names: Dict[str, str] = field(
+        default_factory=lambda: {
+            "PO4": "PO4",
+            "NO3": "NO3",
+            "SiO3": "SiO3",
+            "Fe": "Fe",
+            "O2": "O2",
+            "DIC": "DIC",
+            "ALK": "Alk",
+        }
+    )
+    opt_var_names: Dict[str, str] = field(
+        default_factory=lambda: {
+            "NH4": "NH4",
+            "Lig": "Lig",
+            "DIC_ALT_CO2": "DIC_ALT_CO2",
+            "ALK_ALT_CO2": "Alk_ALT_CO2",
+            "DOC": "DOC",
+            "DON": "DON",
+            "DOP": "DOP",
+            "DOPr": "DOPr",
+            "DONr": "DONr",
+            "DOCr": "DOCr",
+            "spChl": "spChl",
+            "spC": "spC",
+            "spP": "spP",
+            "spFe": "spFe",
+            "diatChl": "diatChl",
+            "diatC": "diatC",
+            "diatP": "diatP",
+            "diatFe": "diatFe",
+            "diatSi": "diatSi",
+            "diazChl": "diazChl",
+            "diazC": "diazC",
+            "diazP": "diazP",
+            "diazFe": "diazFe",
+            "spCaCO3": "spCaCO3",
+            "zooC": "zooC",
+            "CHL": "CHL",
+        }
+    )
+
+    climatology: Optional[bool] = True
+
+
+@dataclass(kw_only=True)
+class UnifiedBGCSurfaceDataset(UnifiedDataset):
+    dim_names: Dict[str, str] = field(
+        default_factory=lambda: {
+            "longitude": "lon",
+            "latitude": "lat",
+        }
+    )
+    var_names: Dict[str, str] = field(
+        default_factory=lambda: {"pco2_air": "pco2_air", "dust": "dust", "iron": "iron"}
+    )
+    opt_var_names: Dict[str, str] = field(
+        default_factory=lambda: {
+            "pco2_air_alt": "pco2_air_alt",
+            "nox": "nox",
+            "nhy": "nhy",
+        }
+    )
+
+    climatology: Optional[bool] = True
+
+
+@dataclass(kw_only=True)
+class CESMDataset(Dataset):
+    """Represents CESM data on original grid."""
+
     # overwrite clean_up method from parent class
     def clean_up(self, ds: xr.Dataset) -> xr.Dataset:
         """Ensure the dataset's time dimension is correctly defined and standardized.
@@ -1080,29 +1299,7 @@ class CESMDataset(Dataset):
 
 @dataclass(kw_only=True)
 class CESMBGCDataset(CESMDataset):
-    """Represents CESM BGC data on original grid.
-
-    Parameters
-    ----------
-    filename : str
-        The path to the data files. Can contain wildcards.
-    start_time : Optional[datetime], optional
-        The start time for selecting relevant data. If not provided, the data is not filtered by start time.
-    end_time : Optional[datetime], optional
-        The end time for selecting relevant data. If not provided, only data at the start_time is selected if start_time is provided,
-        or no filtering is applied if start_time is not provided.
-    var_names: Dict[str, str], optional
-        Dictionary of variable names that are required in the dataset.
-    dim_names: Dict[str, str], optional
-        Dictionary specifying the names of dimensions in the dataset.
-    climatology : bool
-        Indicates whether the dataset is climatological. Defaults to False.
-
-    Attributes
-    ----------
-    ds : xr.Dataset
-        The xarray Dataset containing the CESM data on its original grid.
-    """
+    """Represents CESM BGC data on original grid."""
 
     var_names: Dict[str, str] = field(
         default_factory=lambda: {
@@ -1200,29 +1397,7 @@ class CESMBGCDataset(CESMDataset):
 
 @dataclass(kw_only=True)
 class CESMBGCSurfaceForcingDataset(CESMDataset):
-    """Represents CESM BGC surface forcing data on original grid.
-
-    Parameters
-    ----------
-    filename : str
-        The path to the data files. Can contain wildcards.
-    start_time : Optional[datetime], optional
-        The start time for selecting relevant data. If not provided, the data is not filtered by start time.
-    end_time : Optional[datetime], optional
-        The end time for selecting relevant data. If not provided, only data at the start_time is selected if start_time is provided,
-        or no filtering is applied if start_time is not provided.
-    var_names: Dict[str, str], optional
-        Dictionary of variable names that are required in the dataset.
-    dim_names: Dict[str, str], optional
-        Dictionary specifying the names of dimensions in the dataset.
-    climatology : bool
-        Indicates whether the dataset is climatological. Defaults to False.
-
-    Attributes
-    ----------
-    ds : xr.Dataset
-        The xarray Dataset containing the CESM data on its original grid.
-    """
+    """Represents CESM BGC surface forcing data on original grid."""
 
     var_names: Dict[str, str] = field(
         default_factory=lambda: {
@@ -1269,29 +1444,7 @@ class CESMBGCSurfaceForcingDataset(CESMDataset):
 
 @dataclass(kw_only=True)
 class ERA5Dataset(Dataset):
-    """Represents ERA5 data on original grid.
-
-    Parameters
-    ----------
-    filename : str
-        The path to the data files. Can contain wildcards.
-    start_time : Optional[datetime], optional
-        The start time for selecting relevant data. If not provided, the data is not filtered by start time.
-    end_time : Optional[datetime], optional
-        The end time for selecting relevant data. If not provided, only data at the start_time is selected if start_time is provided,
-        or no filtering is applied if start_time is not provided.
-    var_names: Dict[str, str], optional
-        Dictionary of variable names that are required in the dataset.
-    dim_names: Dict[str, str], optional
-        Dictionary specifying the names of dimensions in the dataset.
-    climatology : bool
-        Indicates whether the dataset is climatological. Defaults to False.
-
-    Attributes
-    ----------
-    ds : xr.Dataset
-        The xarray Dataset containing the ERA5 data on its original grid.
-    """
+    """Represents ERA5 data on original grid."""
 
     var_names: Dict[str, str] = field(
         default_factory=lambda: {
@@ -1389,27 +1542,10 @@ class ERA5Dataset(Dataset):
 
 @dataclass(kw_only=True)
 class ERA5Correction(Dataset):
-    """Global dataset to correct ERA5 radiation. The dataset contains multiplicative
-    correction factors for the ERA5 shortwave radiation, obtained by comparing the
-    COREv2 climatology to the ERA5 climatology.
+    """Global dataset to correct ERA5 radiation.
 
-    Parameters
-    ----------
-    filename : str, optional
-        The path to the correction files. Defaults to download_correction_data('SSR_correction.nc').
-    var_names: Dict[str, str], optional
-        Dictionary of variable names that are required in the dataset.
-        Defaults to {"swr_corr": "ssr_corr"}.
-    dim_names: Dict[str, str], optional
-        Dictionary specifying the names of dimensions in the dataset.
-        Defaults to {"longitude": "longitude", "latitude": "latitude", "time": "time"}.
-    climatology : bool, optional
-        Indicates if the correction data is a climatology. Defaults to True.
-
-    Attributes
-    ----------
-    ds : xr.Dataset
-        The loaded xarray Dataset containing the correction data.
+    The dataset contains multiplicative correction factors for the ERA5 shortwave
+    radiation, obtained by comparing the COREv2 climatology to the ERA5 climatology.
     """
 
     filename: str = field(
@@ -1429,8 +1565,6 @@ class ERA5Correction(Dataset):
     )
     climatology: Optional[bool] = True
 
-    ds: xr.Dataset = field(init=False, repr=False)
-
     def __post_init__(self):
 
         if not self.climatology:
@@ -1494,27 +1628,7 @@ class ERA5Correction(Dataset):
 
 @dataclass(kw_only=True)
 class ETOPO5Dataset(Dataset):
-    """Represents topography data on the original grid from the ETOPO5 dataset.
-
-    Parameters
-    ----------
-    filename : str, optional
-        The path to the ETOPO5 dataset file. If not provided, the dataset will be downloaded
-        automatically via the `pooch` library.
-    var_names : Dict[str, str], optional
-        Dictionary of variable names required in the dataset. Defaults to:
-        {
-            "topo": "topo",
-        }
-    dim_names : Dict[str, str], optional
-        Dictionary specifying the names of dimensions in the dataset. Defaults to:
-        {"longitude": "lon", "latitude": "lat"}.
-
-    Attributes
-    ----------
-    ds : xr.Dataset
-        The xarray Dataset containing the ETOPO5 data, loaded from the specified file.
-    """
+    """Represents topography data on the original grid from the ETOPO5 dataset."""
 
     filename: str = field(default_factory=lambda: download_topo("etopo5.nc"))
     var_names: Dict[str, str] = field(
@@ -1525,7 +1639,6 @@ class ETOPO5Dataset(Dataset):
     dim_names: Dict[str, str] = field(
         default_factory=lambda: {"longitude": "lon", "latitude": "lat"}
     )
-    ds: xr.Dataset = field(init=False, repr=False)
 
     def clean_up(self, ds: xr.Dataset) -> xr.Dataset:
         """Assign lat and lon as coordinates.
@@ -1551,28 +1664,8 @@ class ETOPO5Dataset(Dataset):
 
 @dataclass(kw_only=True)
 class SRTM15Dataset(Dataset):
-    """Represents topography data on the original grid from the SRTM15 dataset.
+    """Represents topography data on the original grid from the SRTM15 dataset."""
 
-    Parameters
-    ----------
-    filename : str
-        The path to the SRTM15 dataset file.
-    var_names : Dict[str, str], optional
-        Dictionary of variable names required in the dataset. Defaults to:
-        {
-            "topo": "z",
-        }
-    dim_names : Dict[str, str], optional
-        Dictionary specifying the names of dimensions in the dataset. Defaults to:
-        {"longitude": "lon", "latitude": "lat"}.
-
-    Attributes
-    ----------
-    ds : xr.Dataset
-        The xarray Dataset containing the SRTM15 data, loaded from the specified file.
-    """
-
-    filename: str
     var_names: Dict[str, str] = field(
         default_factory=lambda: {
             "topo": "z",
@@ -1581,7 +1674,6 @@ class SRTM15Dataset(Dataset):
     dim_names: Dict[str, str] = field(
         default_factory=lambda: {"longitude": "lon", "latitude": "lat"}
     )
-    ds: xr.Dataset = field(init=False, repr=False)
 
 
 # river datasets
@@ -1962,37 +2054,11 @@ class RiverDataset:
 
 @dataclass(kw_only=True)
 class DaiRiverDataset(RiverDataset):
-    """Represents river data from the Dai river dataset.
-
-    Parameters
-    ----------
-    filename : Union[str, Path, List[Union[str, Path]]], optional
-        The path to the Dai River dataset file. If not provided, the dataset will be downloaded
-        automatically via the `pooch` library.
-    start_time : datetime
-        The start time for selecting relevant data.
-    end_time : datetime
-        The end time for selecting relevant data.
-    dim_names: Dict[str, str], optional
-        Dictionary specifying the names of dimensions in the dataset.
-    var_names: Dict[str, str], optional
-        Dictionary of variable names that are required in the dataset.
-    opt_var_names: Dict[str, str], optional
-        Dictionary of variable names that are optional in the dataset.
-    climatology : bool
-        Indicates whether the dataset is climatological. Defaults to False.
-
-    Attributes
-    ----------
-    ds : xr.Dataset
-        The xarray Dataset containing the forcing data on its original grid.
-    """
+    """Represents river data from the Dai river dataset."""
 
     filename: Union[str, Path, List[Union[str, Path]]] = field(
         default_factory=lambda: download_river_data("dai_trenberth_may2019.nc")
     )
-    start_time: datetime
-    end_time: datetime
     dim_names: Dict[str, str] = field(
         default_factory=lambda: {
             "station": "station",
@@ -2014,7 +2080,6 @@ class DaiRiverDataset(RiverDataset):
         }
     )
     climatology: Optional[bool] = False
-    ds: xr.Dataset = field(init=False, repr=False)
 
     def add_time_info(self, ds: xr.Dataset) -> xr.Dataset:
         """Adds time information to the dataset based on the climatology flag and
@@ -2053,6 +2118,497 @@ class DaiRiverDataset(RiverDataset):
         return ds
 
 
+@dataclass
+class TPXOManager:
+    """Manages multiple TPXODataset instances and selects and processes tidal
+    constituents from the TPXO dataset.
+
+    This class handles multiple tidal constituents following the TPXO9v2a standard.
+    The self-attraction and loading (SAL) correction data is sourced internally from TPXO9v2a.
+
+    Parameters
+    ----------
+    filenames : dict
+        Dictionary containing paths to TPXO dataset files. Expected keys:
+        - "h" : Path to the elevation file.
+        - "u" : Path to the u-velocity component file.
+        - "grid" : Path to the grid file.
+
+    ntides : int
+        Number of tidal constituents to select for processing.
+
+    reference_date : datetime, optional
+        Reference date for the TPXO data. Defaults to January 1, 1992.
+        Used as the baseline for tidal time series calculations.
+
+    allan_factor : float, optional
+        Factor used in tidal model computations. Defaults to 2.0.
+
+    use_dask : bool, optional
+        Whether to use Dask for chunking. If True, data is loaded lazily; if False, data is loaded eagerly. Defaults to False.
+
+    Notes
+    -----
+    In TPXO products newer then 9v2a, the order of tidal constituents may change beyond the 10th constituent.
+    Before selecting the first `ntides` constituents, newer products are reordered to match the TPXO9v2a standard.
+    However, this reordering has minimal impact, as constituents beyond the 10th have negligible amplitudes.
+    """
+
+    filenames: dict
+    ntides: int
+    reference_date: datetime = datetime(1992, 1, 1)
+    allan_factor: float = 2.0
+    use_dask: Optional[bool] = False
+
+    def __post_init__(self):
+
+        fname_sal = download_sal_data("sal_tpxo9.v2a.nc")
+
+        # Initialize the data_dict with TPXODataset instances
+        data_dict = {
+            "h": TPXODataset(
+                filename=self.filenames["h"],
+                grid_filename=self.filenames["grid"],
+                location="h",
+                var_names={"ssh_Re": "hRe", "ssh_Im": "hIm"},
+                use_dask=self.use_dask,
+            ),
+            "sal": TPXODataset(
+                filename=fname_sal,
+                grid_filename=self.filenames["grid"],
+                location="h",
+                var_names={"sal_Re": "hRe", "sal_Im": "hIm"},
+                use_dask=self.use_dask,
+                tolerate_coord_mismatch=True,  # Allow coordinate mismatch since SAL is from TPXO9v2a and may not align exactly with newer grids
+            ),
+            "u": TPXODataset(
+                filename=self.filenames["u"],
+                grid_filename=self.filenames["grid"],
+                location="u",
+                var_names={"u_Re": "URe", "u_Im": "UIm"},
+                use_dask=self.use_dask,
+            ),
+            "v": TPXODataset(
+                filename=self.filenames["u"],
+                grid_filename=self.filenames["grid"],
+                location="v",
+                var_names={"v_Re": "VRe", "v_Im": "VIm"},
+                use_dask=self.use_dask,
+            ),
+        }
+
+        omega = self.get_omega()
+
+        for data in data_dict.values():
+            data.select_constituents(self.ntides, omega)
+
+        data_dict["omega"] = xr.DataArray(
+            data=list(omega.values())[: self.ntides],
+            dims="ntides",
+            attrs={"long_name": "angular frequency", "units": "radians per second"},
+        )
+
+        object.__setattr__(self, "datasets", data_dict)
+
+    def get_omega(self):
+        """Retrieve angular frequencies (omega) for tidal constituents from the TPXO9.v2
+        atlas.
+
+        This method returns the angular frequencies (in radians per second) for 15 tidal constituents,
+        sourced from the TPXO tidal model and defined in the OTPSnc `constit.h` file, see https://www.tpxo.net/otps.
+        These values are essential for tidal modeling and analysis.
+
+        Returns
+        -------
+        dict
+            A dictionary where the keys are tidal constituent labels (str) and the values
+            are their respective angular frequencies (float, in radians per second).
+        """
+        omega = {
+            "m2": 1.405189e-04,  # Principal lunar semidiurnal
+            "s2": 1.454441e-04,  # Principal solar semidiurnal
+            "n2": 1.378797e-04,  # Larger lunar elliptic semidiurnal
+            "k2": 1.458423e-04,  # Lunisolar semidiurnal
+            "k1": 7.292117e-05,  # Lunar diurnal
+            "o1": 6.759774e-05,  # Lunar diurnal
+            "p1": 7.252295e-05,  # Solar diurnal
+            "q1": 6.495854e-05,  # Larger lunar elliptic diurnal
+            "mm": 0.026392e-04,  # Lunar monthly
+            "mf": 0.053234e-04,  # Lunar fortnightly
+            "m4": 2.810377e-04,  # Shallow water overtide of M2
+            "mn4": 2.783984e-04,  # Shallow water quarter diurnal
+            "ms4": 2.859630e-04,  # Shallow water quarter diurnal
+            "2n2": 1.352405e-04,  # Shallow water semidiurnal
+            "s1": 7.2722e-05,  # Solar diurnal
+        }
+        return omega
+
+    def compute_equilibrium_tide(self, lon, lat):
+        """Compute equilibrium tide for given longitudes and latitudes.
+
+        Parameters
+        ----------
+        lon : xr.DataArray
+            Longitudes in degrees.
+        lat : xr.DataArray
+            Latitudes in degrees.
+
+        Returns
+        -------
+        tpc : xr.DataArray
+            Equilibrium tide complex amplitude.
+
+        Notes
+        -----
+        This method calculates the equilibrium tide complex amplitude for specified
+        longitudes and latitudes, considering 15 tidal constituents and their corresponding
+        amplitudes and elasticity factors. The order of the tidal constituents corresponds
+        to the order in `self.get_omega()`, which must remain consistent for future use.
+
+        The tidal constituents are categorized as follows:
+        - **2**: Semidiurnal
+        - **1**: Diurnal
+        - **0**: Long-period
+
+        The amplitudes and elasticity factors are sourced from the `constit.h` file in the OTPSnc package.
+        """
+
+        # Amplitudes for 15 tidal constituents (from variable amp_d in constit.h of OTPSnc package)
+        A = xr.DataArray(
+            data=np.array(
+                [
+                    0.242334,  # M2
+                    0.112743,  # S2
+                    0.046397,  # N2
+                    0.030684,  # K2
+                    0.141565,  # K1
+                    0.100661,  # O1
+                    0.046848,  # P1
+                    0.019273,  # Q1
+                    0.022191,  # Mm
+                    0.042041,  # Mf
+                    0.0,  # M4
+                    0.0,  # Mn4
+                    0.0,  # Ms4
+                    0.006141,  # 2n2
+                    0.000764,  # S1
+                ]
+            ),
+            dims="ntides",
+        )
+
+        # Elasticity factors for 15 tidal constituents (from variable alpha_d in constit.h of OTPSnc package)
+        B = xr.DataArray(
+            data=np.array(
+                [
+                    0.693,  # M2
+                    0.693,  # S2
+                    0.693,  # N2
+                    0.693,  # K2
+                    0.736,  # K1
+                    0.695,  # O1
+                    0.706,  # P1
+                    0.695,  # Q1
+                    0.693,  # Mm
+                    0.693,  # Mf
+                    0.693,  # M4
+                    0.693,  # Mn4
+                    0.693,  # Ms4
+                    0.693,  # 2n2
+                    0.693,  # S1
+                ]
+            ),
+            dims="ntides",
+        )
+
+        # Tidal type (from variable ispec_d in constit.h of OTPSnc package)
+        # types: 2 = semidiurnal, 1 = diurnal, 0 = long-term
+        ityp = xr.DataArray(
+            data=np.array(
+                [
+                    2,  # M2
+                    2,  # S2
+                    2,  # N2
+                    2,  # K2
+                    1,  # K1
+                    1,  # O1
+                    1,  # P1
+                    1,  # Q1
+                    0,  # Mm
+                    0,  # Mf
+                    0,  # M4
+                    0,  # Mn4
+                    0,  # Ms4
+                    2,  # 2n2
+                    1,  # S1
+                ]
+            ),
+            dims="ntides",
+        )
+
+        d2r = np.pi / 180
+        coslat2 = np.cos(d2r * lat) ** 2
+        sin2lat = np.sin(2 * d2r * lat)
+
+        p_amp = (
+            xr.where(ityp == 2, 1, 0) * A * B * coslat2  # semidiurnal
+            + xr.where(ityp == 1, 1, 0) * A * B * sin2lat  # diurnal
+            + xr.where(ityp == 0, 1, 0) * A * B * (0.5 - 1.5 * coslat2)  # long-term
+        )
+        p_pha = (
+            xr.where(ityp == 2, 1, 0) * (-2 * lon * d2r)  # semidiurnal
+            + xr.where(ityp == 1, 1, 0) * (-lon * d2r)  # diurnal
+            + xr.where(ityp == 0, 1, 0) * xr.zeros_like(lon)  # long-term
+        )
+
+        tpc = p_amp * np.exp(-1j * p_pha)
+        tpc = tpc.isel(ntides=slice(None, self.ntides))
+
+        return tpc
+
+    def egbert_correction(self, date):
+        """Correct phases and amplitudes for real-time runs using parts of the post-
+        processing code from Egbert's & Erofeeva's (OSU) TPXO model.
+
+        Parameters
+        ----------
+        date : datetime.datetime
+            The date and time for which corrections are to be applied.
+
+        Returns
+        -------
+        pf : xr.DataArray
+            Amplitude scaling factor for each of the 15 tidal constituents.
+        pu : xr.DataArray
+            Phase correction [radians] for each of the 15 tidal constituents.
+        aa : xr.DataArray
+            Astronomical arguments [radians] associated with the corrections.
+
+        Notes
+        -----
+        The order of the tidal constituents corresponds
+        to the order in `self.get_omega()`, which must remain consistent for future use.
+
+        References
+        ----------
+        - Egbert, G.D., and S.Y. Erofeeva. "Efficient inverse modeling of barotropic ocean
+          tides." Journal of Atmospheric and Oceanic Technology 19, no. 2 (2002): 183-204.
+        """
+
+        year = date.year
+        month = date.month
+        day = date.day
+        hour = date.hour
+        minute = date.minute
+        second = date.second
+
+        rad = np.pi / 180.0
+        deg = 180.0 / np.pi
+        mjd = modified_julian_days(year, month, day)
+        tstart = mjd + hour / 24 + minute / (60 * 24) + second / (60 * 60 * 24)
+
+        # Determine nodal corrections pu & pf : these expressions are valid for period 1990-2010 (Cartwright 1990).
+        # Reset time origin for astronomical arguments to 4th of May 1860:
+        timetemp = tstart - 51544.4993
+
+        # mean longitude of lunar perigee
+        P = 83.3535 + 0.11140353 * timetemp
+        P = np.mod(P, 360.0)
+        if P < 0:
+            P = +360
+        P *= rad
+
+        # mean longitude of ascending lunar node
+        N = 125.0445 - 0.05295377 * timetemp
+        N = np.mod(N, 360.0)
+        if N < 0:
+            N = +360
+        N *= rad
+
+        sinn = np.sin(N)
+        cosn = np.cos(N)
+        sin2n = np.sin(2 * N)
+        cos2n = np.cos(2 * N)
+        sin3n = np.sin(3 * N)
+
+        pftmp = np.sqrt(
+            (1 - 0.03731 * cosn + 0.00052 * cos2n) ** 2
+            + (0.03731 * sinn - 0.00052 * sin2n) ** 2
+        )
+
+        pf = np.zeros(15)
+        pf[0] = pftmp  # M2
+        pf[1] = 1.0  # S2
+        pf[2] = pftmp  # N2
+        pf[3] = np.sqrt(
+            (1 + 0.2852 * cosn + 0.0324 * cos2n) ** 2
+            + (0.3108 * sinn + 0.0324 * sin2n) ** 2
+        )  # K2
+        pf[4] = np.sqrt(
+            (1 + 0.1158 * cosn - 0.0029 * cos2n) ** 2
+            + (0.1554 * sinn - 0.0029 * sin2n) ** 2
+        )  # K1
+        pf[5] = np.sqrt(
+            (1 + 0.189 * cosn - 0.0058 * cos2n) ** 2
+            + (0.189 * sinn - 0.0058 * sin2n) ** 2
+        )  # O1
+        pf[6] = 1.0  # P1
+        pf[7] = np.sqrt((1 + 0.188 * cosn) ** 2 + (0.188 * sinn) ** 2)  # Q1
+        pf[8] = 1.0 - 0.130 * cosn  # Mm
+        pf[9] = 1.043 + 0.414 * cosn  # Mf
+        pf[10] = pftmp**2  # M4
+        pf[11] = pftmp**2  # Mn4
+        pf[12] = pftmp**2  # Ms4
+        pf[13] = pftmp  # 2n2
+        pf[14] = 1.0  # S1
+        pf = xr.DataArray(pf, dims="ntides")
+
+        putmp = (
+            np.arctan(
+                (-0.03731 * sinn + 0.00052 * sin2n)
+                / (1.0 - 0.03731 * cosn + 0.00052 * cos2n)
+            )
+            * deg
+        )
+
+        pu = np.zeros(15)
+        pu[0] = putmp  # M2
+        pu[1] = 0.0  # S2
+        pu[2] = putmp  # N2
+        pu[3] = (
+            np.arctan(
+                -(0.3108 * sinn + 0.0324 * sin2n)
+                / (1.0 + 0.2852 * cosn + 0.0324 * cos2n)
+            )
+            * deg
+        )  # K2
+        pu[4] = (
+            np.arctan(
+                (-0.1554 * sinn + 0.0029 * sin2n)
+                / (1.0 + 0.1158 * cosn - 0.0029 * cos2n)
+            )
+            * deg
+        )  # K1
+        pu[5] = 10.8 * sinn - 1.3 * sin2n + 0.2 * sin3n  # O1
+        pu[6] = 0.0  # P1
+        pu[7] = np.arctan(0.189 * sinn / (1.0 + 0.189 * cosn)) * deg  # Q1
+        pu[8] = 0.0  # Mm
+        pu[9] = -23.7 * sinn + 2.7 * sin2n - 0.4 * sin3n  # Mf
+        pu[10] = putmp * 2.0  # M4
+        pu[11] = putmp * 2.0  # Mn4
+        pu[12] = putmp  # Ms4
+        pu[13] = putmp  # 2n2
+        pu[14] = 0.0  # S1
+        pu = xr.DataArray(pu, dims="ntides")
+        # convert from degrees to radians
+        pu = pu * rad
+
+        aa = xr.DataArray(
+            data=np.array(
+                [
+                    1.731557546,  # M2
+                    0.0,  # S2
+                    6.050721243,  # N2
+                    3.487600001,  # K2
+                    0.173003674,  # K1
+                    1.558553872,  # O1
+                    6.110181633,  # P1
+                    5.877717569,  # Q1
+                    1.964021610,  # Mm
+                    1.756042456,  # Mf
+                    3.463115091,  # M4
+                    1.499093481,  # Mn4
+                    1.731557546,  # Ms4
+                    4.086699633,  # 2n2
+                    0.0,  # S1
+                ]
+            ),
+            dims="ntides",
+        )
+        pf = pf.isel(ntides=slice(None, self.ntides))
+        pu = pu.isel(ntides=slice(None, self.ntides))
+        aa = aa.isel(ntides=slice(None, self.ntides))
+
+        return pf, pu, aa
+
+    def correct_tides(self, model_reference_date):
+        """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
+        ----------
+        model_reference_date : datetime
+            The reference date for the ROMS simulation.
+
+        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').
+        """
+
+        datasets = self.datasets
+        omega = self.datasets["omega"].isel(ntides=slice(None, self.ntides))
+
+        # Get equilibrium tides
+        lon = datasets["sal"].ds[datasets["sal"].dim_names["longitude"]]
+        lat = datasets["sal"].ds[datasets["sal"].dim_names["latitude"]]
+        tpc = self.compute_equilibrium_tide(lon, lat)
+
+        # Correct for SAL
+        tsc = self.allan_factor * (
+            datasets["sal"].ds[datasets["sal"].var_names["sal_Re"]]
+            + 1j * datasets["sal"].ds[datasets["sal"].var_names["sal_Im"]]
+        )
+        tpc = tpc - tsc
+
+        # Elevations and transports
+        thc = (
+            datasets["h"].ds[datasets["h"].var_names["ssh_Re"]]
+            + 1j * datasets["h"].ds[datasets["h"].var_names["ssh_Im"]]
+        )
+        tuc = (
+            datasets["u"].ds[datasets["u"].var_names["u_Re"]]
+            + 1j * datasets["u"].ds[datasets["u"].var_names["u_Im"]]
+        )
+        tvc = (
+            datasets["v"].ds[datasets["v"].var_names["v_Re"]]
+            + 1j * datasets["v"].ds[datasets["v"].var_names["v_Im"]]
+        )
+
+        # Apply correction for phases and amplitudes
+        pf, pu, aa = self.egbert_correction(model_reference_date)
+
+        dt = (model_reference_date - self.reference_date).days * 3600 * 24
+
+        thc = pf * thc * np.exp(1j * (omega * dt + pu + aa))
+        tuc = pf * tuc * np.exp(1j * (omega * dt + pu + aa))
+        tvc = pf * tvc * np.exp(1j * (omega * dt + pu + aa))
+        tpc = pf * tpc * np.exp(1j * (omega * dt + pu + aa))
+
+        datasets["h"].ds[datasets["h"].var_names["ssh_Re"]] = thc.real
+        datasets["h"].ds[datasets["h"].var_names["ssh_Im"]] = thc.imag
+        datasets["u"].ds[datasets["u"].var_names["u_Re"]] = tuc.real
+        datasets["u"].ds[datasets["u"].var_names["u_Im"]] = tuc.imag
+        datasets["v"].ds[datasets["v"].var_names["v_Re"]] = tvc.real
+        datasets["v"].ds[datasets["v"].var_names["v_Im"]] = tvc.imag
+        datasets["sal"].ds["pot_Re"] = tpc.real
+        datasets["sal"].ds["pot_Im"] = tpc.imag
+
+        object.__setattr__(self, "datasets", datasets)
+
+        # Update var_names dictionary
+        var_names = {
+            **datasets["sal"].var_names,
+            "pot_Re": "pot_Re",
+            "pot_Im": "pot_Im",
+        }
+        var_names.pop("sal_Re", None)  # Remove "sal_Re" if it exists
+        var_names.pop("sal_Im", None)  # Remove "sal_Im" if it exists
+        object.__setattr__(self.datasets["sal"], "var_names", var_names)
+
+
 # shared functions
 
 
@@ -2168,8 +2724,9 @@ def _select_relevant_times(
                     f"The dataset contains {len(ds[time_dim])} time steps, but the climatology flag is set to True, which requires exactly 12 time steps."
                 )
             if not end_time:
+                # Convert from timedelta64[ns] to fractional days
+                ds["time"] = ds["time"] / np.timedelta64(1, "D")
                 # Interpolate from climatology for initial conditions
-                ds["time"] = ds["time"].dt.days
                 ds = interpolate_from_climatology(ds, time_dim, start_time)
         else:
             time_type = get_time_type(ds[time_dim])
@@ -2254,3 +2811,61 @@ def decode_string(byte_array):
     )
 
     return decoded_string
+
+
+def modified_julian_days(year, month, day, hour=0):
+    """Calculate the Modified Julian Day (MJD) for a given date and time.
+
+    The Modified Julian Day (MJD) is a modified Julian day count starting from
+    November 17, 1858 AD. It is commonly used in astronomy and geodesy.
+
+    Parameters
+    ----------
+    year : int
+        The year.
+    month : int
+        The month (1-12).
+    day : int
+        The day of the month.
+    hour : float, optional
+        The hour of the day as a fractional number (0 to 23.999...). Default is 0.
+
+    Returns
+    -------
+    mjd : float
+        The Modified Julian Day (MJD) corresponding to the input date and time.
+
+    Notes
+    -----
+    The algorithm assumes that the input date (year, month, day) is within the
+    Gregorian calendar, i.e., after October 15, 1582. Negative MJD values are
+    allowed for dates before November 17, 1858.
+
+    References
+    ----------
+    - Wikipedia article on Julian Day: https://en.wikipedia.org/wiki/Julian_day
+    - Wikipedia article on Modified Julian Day: https://en.wikipedia.org/wiki/Modified_Julian_day
+
+    Examples
+    --------
+    >>> modified_julian_days(2024, 5, 20, 12)
+    58814.0
+    >>> modified_julian_days(1858, 11, 17)
+    0.0
+    >>> modified_julian_days(1582, 10, 4)
+    -141428.5
+    """
+
+    if month < 3:
+        year -= 1
+        month += 12
+
+    A = year // 100
+    B = A // 4
+    C = 2 - A + B
+    E = int(365.25 * (year + 4716))
+    F = int(30.6001 * (month + 1))
+    jd = C + day + hour / 24 + E + F - 1524.5
+    mjd = jd - 2400000.5
+
+    return mjd