roms-tools 2.6.2__py3-none-any.whl → 3.0.0__py3-none-any.whl

roms_tools/setup/datasets.py

@@ -1,27 +1,31 @@
+import logging
 import time
-import xarray as xr
 from dataclasses import dataclass, field
 from datetime import datetime, timedelta
-import numpy as np
-from typing import Dict, Optional, Union, List
 from pathlib import Path
-import logging
-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,
-    gc_dist,
-)
+from typing import Dict, List, Optional, Union
+
+import numpy as np
+import xarray as xr
+
+from roms_tools.constants import R_EARTH
 from roms_tools.download import (
     download_correction_data,
-    download_topo,
     download_river_data,
+    download_sal_data,
+    download_topo,
 )
 from roms_tools.setup.fill import LateralFill
+from roms_tools.setup.utils import (
+    assign_dates_to_climatology,
+    convert_cftime_to_datetime,
+    gc_dist,
+    get_time_type,
+    interpolate_cyclic_time,
+    interpolate_from_climatology,
+    one_dim_fill,
+)
+from roms_tools.utils import _has_gcsfs, _load_data
 
 # lat-lon datasets
 
@@ -91,6 +95,7 @@ class Dataset:
     needs_lateral_fill: Optional[bool] = True
     use_dask: Optional[bool] = False
     apply_post_processing: Optional[bool] = True
+    read_zarr: Optional[bool] = False
 
     is_global: bool = field(init=False, repr=False)
     ds: xr.Dataset = field(init=False, repr=False)
@@ -164,11 +169,13 @@ class Dataset:
             If a list of files is provided but self.dim_names["time"] is not available or use_dask=False.
         """
 
-        ds = _load_data(self.filename, self.dim_names, self.use_dask)
+        ds = _load_data(
+            self.filename, self.dim_names, self.use_dask, read_zarr=self.read_zarr
+        )
 
         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
@@ -182,9 +189,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.
@@ -221,6 +228,7 @@ class Dataset:
             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)
 
@@ -385,7 +393,6 @@ class Dataset:
             and longitude differences, in meters.
         """
 
-        r_earth = 6371315.0
         lat_dim = self.dim_names["latitude"]
         lon_dim = self.dim_names["longitude"]
 
@@ -398,11 +405,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
@@ -698,6 +705,10 @@ 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.
@@ -801,119 +812,223 @@ class TPXODataset(Dataset):
 
     Parameters
     ----------
-    reference_date : datetime, optional
-        The reference date for the TPXO data. Default is datetime(1992, 1, 1).
+    filename : str
+        The path to the TPXO dataset file.
+    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.
     """
 
-    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",
-        }
-    )
+    filename: str
+    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"}
     )
-    reference_date: datetime = datetime(1992, 1, 1)
+    tolerate_coord_mismatch: bool = False
+    ds: xr.Dataset = field(init=False, repr=False)
 
     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)
@@ -1007,15 +1122,11 @@ class UnifiedDataset(Dataset):
             }
         )
 
-        object.__setattr__(
-            self,
-            "dim_names",
-            {
-                "latitude": "latitude",
-                "longitude": "longitude",
-                "depth": "depth",
-            },
-        )
+        self.dim_names = {
+            "latitude": "latitude",
+            "longitude": "longitude",
+            "depth": "depth",
+        }
 
         # Handle time dimension
         if "time" not in self.dim_names:
@@ -1434,6 +1545,36 @@ class ERA5Dataset(Dataset):
             self.var_names = var_names
 
 
+@dataclass(kw_only=True)
+class ERA5ARCODataset(ERA5Dataset):
+    # The GCP ARCO version of the dataset has different variable names
+
+    var_names: Dict[str, str] = field(
+        default_factory=lambda: {
+            "uwnd": "10m_u_component_of_wind",
+            "vwnd": "10m_v_component_of_wind",
+            "swrad": "surface_net_solar_radiation",
+            "lwrad": "surface_thermal_radiation_downwards",
+            "Tair": "2m_temperature",
+            "d2m": "2m_dewpoint_temperature",
+            "rain": "total_precipitation",
+            "mask": "sea_surface_temperature",
+        }
+    )
+
+    def __post_init__(self):
+        self.read_zarr = True
+        if not _has_gcsfs():
+            raise RuntimeError(
+                "To use cloud-based ERA5 data, GCSFS is required but not installed. Install it with:\n"
+                "  • `pip install roms-tools[stream]` or\n"
+                "  • `conda install gcsfs`\n"
+                "Alternatively, install `roms-tools` with conda to include all dependencies."
+            )
+
+        super().__post_init__()
+
+
 @dataclass(kw_only=True)
 class ERA5Correction(Dataset):
     """Global dataset to correct ERA5 radiation.
@@ -2012,6 +2153,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
 
 
@@ -2214,3 +2846,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