roms-tools 0.0.6__py3-none-any.whl → 0.1.0__py3-none-any.whl

roms_tools/setup/tides.py

@@ -0,0 +1,676 @@
+from datetime import datetime
+import xarray as xr
+import numpy as np
+from dataclasses import dataclass, field
+from roms_tools.setup.grid import Grid
+from roms_tools.setup.plot import _plot
+import os
+import hashlib
+
+
+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
+
+
+def egbert_correction(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.
+
+    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
+    )  # 2N2
+
+    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.043 + 0.414 * cosn  # Mf
+    pf[9] = 1.0 - 0.130 * cosn  # Mm
+    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="nc")
+
+    putmp = (
+        np.arctan(
+            (-0.03731 * sinn + 0.00052 * sin2n)
+            / (1.0 - 0.03731 * cosn + 0.00052 * cos2n)
+        )
+        * deg
+    )  # 2N2
+
+    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] = -23.7 * sinn + 2.7 * sin2n - 0.4 * sin3n  # Mf
+    pu[9] = 0.0  # Mm
+    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="nc")
+    # 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="nc",
+    )
+
+    return pf, pu, aa
+
+
+@dataclass(frozen=True, kw_only=True)
+class TPXO:
+    """
+    Represents TPXO tidal atlas.
+
+    Parameters
+    ----------
+    filename : str
+        The path to the TPXO dataset.
+
+    Attributes
+    ----------
+    ds : xr.Dataset
+        The xarray Dataset containing TPXO tidal model data.
+
+    Notes
+    -----
+    This class provides a convenient interface to work with TPXO tidal atlas.
+
+    Examples
+    --------
+    >>> tpxo = TPXO()
+    >>> tpxo.load_data("tpxo_data.nc")
+    >>> print(tpxo.ds)
+    <xarray.Dataset>
+    Dimensions:  ...
+    """
+
+    filename: str
+    ds: xr.Dataset = field(init=False, repr=False)
+
+    def __post_init__(self):
+        ds = self.load_data(self.filename)
+        # Lon_r is constant along ny, i.e., is only a function of nx
+        ds["nx"] = ds["lon_r"].isel(ny=0)
+        # Lat_r is constant along nx, i.e., is only a function of ny
+        ds["ny"] = ds["lat_r"].isel(nx=0)
+
+        object.__setattr__(self, "ds", ds)
+
+    def get_corrected_tides(self, model_reference_date, allan_factor):
+        # Get equilibrium tides
+        tpc = self.compute_equilibrium_tide(self.ds["lon_r"], self.ds["lat_r"])
+        # Correct for SAL
+        tsc = allan_factor * (self.ds["sal_Re"] + 1j * self.ds["sal_Im"])
+        tpc = tpc - tsc
+
+        # Elevations and transports
+        thc = self.ds["h_Re"] + 1j * self.ds["h_Im"]
+        tuc = self.ds["u_Re"] + 1j * self.ds["u_Im"]
+        tvc = self.ds["v_Re"] + 1j * self.ds["v_Im"]
+
+        # Apply correction for phases and amplitudes
+        pf, pu, aa = egbert_correction(model_reference_date)
+        tpxo_reference_date = datetime(1992, 1, 1)
+        dt = (model_reference_date - tpxo_reference_date).days * 3600 * 24
+
+        thc = pf * thc * np.exp(1j * (self.ds["omega"] * dt + pu + aa))
+        tuc = pf * tuc * np.exp(1j * (self.ds["omega"] * dt + pu + aa))
+        tvc = pf * tvc * np.exp(1j * (self.ds["omega"] * dt + pu + aa))
+        tpc = pf * tpc * np.exp(1j * (self.ds["omega"] * dt + pu + aa))
+
+        tides = {"ssh": thc, "u": tuc, "v": tvc, "pot": tpc, "omega": self.ds["omega"]}
+
+        return tides
+
+    @staticmethod
+    def load_data(filename):
+        """
+        Load tidal forcing data from the specified file.
+
+        Parameters
+        ----------
+        filename : str
+            The path to the tidal dataset file.
+
+        Returns
+        -------
+        ds : xr.Dataset
+            The loaded xarray Dataset containing the tidal forcing data.
+
+        Raises
+        ------
+        FileNotFoundError
+            If the specified file does not exist.
+        ValueError
+            If the file checksum does not match the expected value.
+
+        Notes
+        -----
+        This method performs basic file existence and checksum checks to ensure the integrity of the loaded dataset.
+
+        """
+        # Check if the file exists
+        if not os.path.exists(filename):
+            raise FileNotFoundError(f"File '{filename}' not found.")
+
+        # Calculate the checksum of the file
+        expected_checksum = (
+            "306956d8769737ba39040118d8d08f467187fe453e02a5651305621d095bce6e"
+        )
+        with open(filename, "rb") as file:
+            actual_checksum = hashlib.sha256(file.read()).hexdigest()
+
+        # Compare the checksums
+        if actual_checksum != expected_checksum:
+            raise ValueError(
+                "Checksum mismatch. The file may be corrupted or tampered with."
+            )
+
+        # Load the dataset
+        ds = xr.open_dataset(filename)
+
+        return ds
+
+    @staticmethod
+    def compute_equilibrium_tide(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 computes the equilibrium tide complex amplitude for given longitudes
+        and latitudes. It considers 15 tidal constituents and their corresponding
+        amplitudes and elasticity factors. The types of tides are classified as follows:
+            - 2: semidiurnal
+            - 1: diurnal
+            - 0: long-term
+
+        """
+
+        # Amplitudes and elasticity factors for 15 tidal constituents
+        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.042041,  # Mf
+                    0.022191,  # Mm
+                    0.0,  # M4
+                    0.0,  # Mn4
+                    0.0,  # Ms4
+                    0.006141,  # 2n2
+                    0.000764,  # S1
+                ]
+            ),
+            dims="nc",
+        )
+        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,  # Mf
+                    0.693,  # Mm
+                    0.693,  # M4
+                    0.693,  # Mn4
+                    0.693,  # Ms4
+                    0.693,  # 2n2
+                    0.693,  # S1
+                ]
+            ),
+            dims="nc",
+        )
+
+        # types: 2 = semidiurnal, 1 = diurnal, 0 = long-term
+        ityp = xr.DataArray(
+            data=np.array([2, 2, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0, 0, 2, 1]), dims="nc"
+        )
+
+        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)
+
+        return tpc
+
+    @staticmethod
+    def concatenate_across_dateline(field):
+        """
+        Concatenate a field across the dateline for TPXO atlas.
+
+        Parameters
+        ----------
+        field : xr.DataArray
+            The field to be concatenated across the dateline.
+
+        Returns
+        -------
+        field_concatenated : xr.DataArray
+            The field concatenated across the dateline.
+
+        Notes
+        -----
+        The TPXO atlas has a minimum longitude of 0.167 and a maximum longitude of 360.0.
+        This method concatenates the field along the dateline on the lower end, considering
+        the discontinuity in longitudes.
+
+        """
+        lon = field["nx"]
+        lon_minus360 = lon - 360
+        lon_concatenated = xr.concat([lon_minus360, lon], dim="nx")
+        field_concatenated = xr.concat([field, field], dim="nx")
+        field_concatenated["nx"] = lon_concatenated
+
+        return field_concatenated
+
+
+@dataclass(frozen=True, kw_only=True)
+class TidalForcing:
+    """
+    Represents tidal forcing data used in ocean modeling.
+
+    Parameters
+    ----------
+    grid : Grid
+        The grid object representing the ROMS grid associated with the tidal forcing data.
+    filename: str
+        The path to the native tidal dataset.
+    nc : int, optional
+        Number of constituents to consider. Maximum number is 14. Default is 10.
+    model_reference_date : datetime, optional
+        The reference date for the ROMS simulation. Default is datetime(2000, 1, 1).
+    source : str, optional
+        The source of the tidal data. Default is "tpxo".
+    allan_factor : float, optional
+        The Allan factor used in tidal model computation. Default is 2.0.
+
+    Attributes
+    ----------
+    ds : xr.Dataset
+        The xarray Dataset containing the tidal forcing data.
+
+    Notes
+    -----
+    This class represents tidal forcing data used in ocean modeling. It provides
+    functionality to load and process tidal data for use in numerical simulations.
+    The tidal forcing data is loaded from a TPXO dataset and processed to generate
+    tidal elevation, tidal potential, and tidal velocity fields.
+
+    Examples
+    --------
+    >>> grid = Grid(...)
+    >>> tidal_forcing = TidalForcing(grid)
+    >>> print(tidal_forcing.ds)
+    """
+
+    grid: Grid
+    filename: str
+    nc: int = 10
+    model_reference_date: datetime = datetime(2000, 1, 1)
+    source: str = "tpxo"
+    allan_factor: float = 2.0
+    ds: xr.Dataset = field(init=False, repr=False)
+
+    def __post_init__(self):
+        if self.source == "tpxo":
+            tpxo = TPXO(filename=self.filename)
+
+            tides = tpxo.get_corrected_tides(
+                self.model_reference_date, self.allan_factor
+            )
+
+            # rename dimension and select desired number of constituents
+            for k in tides.keys():
+                tides[k] = tides[k].rename({"nc": "ntides"})
+                tides[k] = tides[k].isel(ntides=slice(None, self.nc))
+
+            # make sure interpolation works across dateline
+            for key in ["ssh", "pot", "u", "v"]:
+                tides[key] = tpxo.concatenate_across_dateline(tides[key])
+
+            # interpolate onto desired grid
+            ssh_tide = (
+                tides["ssh"]
+                .interp(nx=self.grid.ds.lon_rho, ny=self.grid.ds.lat_rho)
+                .drop_vars(["nx", "ny"])
+            )
+            pot_tide = (
+                tides["pot"]
+                .interp(nx=self.grid.ds.lon_rho, ny=self.grid.ds.lat_rho)
+                .drop_vars(["nx", "ny"])
+            )
+            u = (
+                tides["u"]
+                .interp(nx=self.grid.ds.lon_rho, ny=self.grid.ds.lat_rho)
+                .drop_vars(["nx", "ny"])
+            )
+            v = (
+                tides["v"]
+                .interp(nx=self.grid.ds.lon_rho, ny=self.grid.ds.lat_rho)
+                .drop_vars(["nx", "ny"])
+            )
+
+        # Rotate to grid orientation
+        u_tide = u * np.cos(self.grid.ds.angle) + v * np.sin(self.grid.ds.angle)
+        v_tide = v * np.cos(self.grid.ds.angle) - u * np.sin(self.grid.ds.angle)
+
+        # Convert to barotropic velocity
+        u_tide = u_tide / self.grid.ds.h
+        v_tide = v_tide / self.grid.ds.h
+
+        # Interpolate from rho- to velocity points
+        u_tide = (
+            (u_tide + u_tide.shift(xi_rho=1))
+            .isel(xi_rho=slice(1, None))
+            .drop_vars(["lat_rho", "lon_rho"])
+        )
+        u_tide = u_tide.swap_dims({"xi_rho": "xi_u"})
+        v_tide = (
+            (v_tide + v_tide.shift(eta_rho=1))
+            .isel(eta_rho=slice(1, None))
+            .drop_vars(["lat_rho", "lon_rho"])
+        )
+        v_tide = v_tide.swap_dims({"eta_rho": "eta_v"})
+
+        # save in new dataset
+        ds = xr.Dataset()
+
+        ds["omega"] = tides["omega"]
+
+        ds["ssh_Re"] = ssh_tide.real
+        ds["ssh_Im"] = ssh_tide.imag
+        ds["ssh_Re"].attrs["long_name"] = "Tidal elevation, real part"
+        ds["ssh_Im"].attrs["long_name"] = "Tidal elevation, complex part"
+        ds["ssh_Re"].attrs["units"] = "m"
+        ds["ssh_Im"].attrs["units"] = "m"
+
+        ds["pot_Re"] = pot_tide.real
+        ds["pot_Im"] = pot_tide.imag
+        ds["pot_Re"].attrs["long_name"] = "Tidal potential, real part"
+        ds["pot_Im"].attrs["long_name"] = "Tidal potential, complex part"
+        ds["pot_Re"].attrs["units"] = "m"
+        ds["pot_Im"].attrs["units"] = "m"
+
+        ds["u_Re"] = u_tide.real
+        ds["u_Im"] = u_tide.imag
+        ds["u_Re"].attrs["long_name"] = "Tidal velocity in x-direction, real part"
+        ds["u_Im"].attrs["long_name"] = "Tidal velocity in x-direction, complex part"
+        ds["u_Re"].attrs["units"] = "m/s"
+        ds["u_Im"].attrs["units"] = "m/s"
+
+        ds["v_Re"] = v_tide.real
+        ds["v_Im"] = v_tide.imag
+        ds["v_Re"].attrs["long_name"] = "Tidal velocity in y-direction, real part"
+        ds["v_Im"].attrs["long_name"] = "Tidal velocity in y-direction, complex part"
+        ds["v_Re"].attrs["units"] = "m/s"
+        ds["v_Im"].attrs["units"] = "m/s"
+
+        ds.attrs["source"] = self.source
+        ds.attrs["model_reference_date"] = self.model_reference_date
+        ds.attrs["allan_factor"] = self.allan_factor
+
+        object.__setattr__(self, "ds", ds)
+
+    def plot(self, var, nc=0) -> None:
+        """
+        Plot the specified tidal forcing variable for a given tidal constituent.
+
+        Parameters
+        ----------
+        var : str
+            The tidal forcing variable to plot. Options include:
+            - "ssh_Re": Real part of tidal elevation.
+            - "ssh_Im": Imaginary part of tidal elevation.
+            - "pot_Re": Real part of tidal potential.
+            - "pot_Im": Imaginary part of tidal potential.
+            - "u_Re": Real part of tidal velocity in the x-direction.
+            - "u_Im": Imaginary part of tidal velocity in the x-direction.
+            - "v_Re": Real part of tidal velocity in the y-direction.
+            - "v_Im": Imaginary part of tidal velocity in the y-direction.
+        nc : int, optional
+            The index of the tidal constituent to plot. Default is 0, which corresponds
+            to the first constituent.
+
+        Returns
+        -------
+        None
+            This method does not return any value. It generates and displays a plot.
+
+        Raises
+        ------
+        ValueError
+            If the specified field is not one of the valid options.
+
+
+        Examples
+        --------
+        >>> tidal_forcing = TidalForcing(grid)
+        >>> tidal_forcing.plot("ssh_Re", nc=0)
+        """
+
+        vmax = max(
+            self.ds[var].isel(ntides=nc).max(), -self.ds[var].isel(ntides=nc).min()
+        )
+        kwargs = {"cmap": "RdBu_r", "vmax": vmax, "vmin": -vmax}
+
+        _plot(
+            self.ds,
+            field=self.ds[var].isel(ntides=nc),
+            straddle=self.grid.straddle,
+            c="g",
+            kwargs=kwargs,
+        )
+
+    def save(self, filepath: str) -> None:
+        """
+        Save the tidal forcing information to a netCDF4 file.
+
+        Parameters
+        ----------
+        filepath
+        """
+        self.ds.to_netcdf(filepath)