From 2327faf9c4c651ac2209724192287710860dc29d Mon Sep 17 00:00:00 2001 From: Shane Elipot Date: Thu, 26 Sep 2024 11:59:08 -0400 Subject: [PATCH] =?UTF-8?q?=E2=AD=90=20new=20`transfer.py`=20module=20(#38?= =?UTF-8?q?2)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add transfer function calculation and test cases * Ekman & Madsen cases for bld=inf * ruff format * complete code untested * Refactor code to improve performance and readability * Refactor transfer function calculations in transfer.py * ruff format * reformat * ruffing it up * Update Bessel functions in transfer.py and some tests * Add docstrings to besseltildes_noslip and besselktilde functions * formatting * more code, more tests * Add import for clouddrift.transfer module * Refactor wind_transfer_test_gradient method in transfer_tests.py * ruffed __init__.py * Refactor wind_transfer function and tests * rename test * added elipot method * sqrt -> numpy.lib.scimath.sqrt * Following @philippemiron's review and refactor some wind transfer test methods * reorder __all__ list * Refactor wind_transfer_free_slip function in transfer.py * new transfer functions (slab and kpp) * Rename slab_transfer to slab_wind_transfer * removed kpp, more test * ruffing it * Refactor plot_ragged function to improve clarity and documentation (#386) * 🔎 Development docs revision (#375) * Add docs focused on contributing and document several processes ----- Co-authored-by: Philippe Miron Co-authored-by: Shane Elipot * ⭐ Hurdat2 adapter (#378) * hurdat2 data adapter * use CF conventions * improve warning when it cannot be determined whether the remote file is new or not * increment revision --------- Co-authored-by: Shane Elipot * add condensed coverage to PR * Update default values for optional parameters in transfer.py * more tests * more tests * Refactor response_to_wind_stress function to include latitude and window parameters * upgrade to 3.10+ * remove top level transfer function for now --------- Co-authored-by: Kevin Santana Co-authored-by: Philippe Miron Co-authored-by: Philippe Miron --- .codecov.yml | 5 +- clouddrift/__init__.py | 2 + clouddrift/transfer.py | 875 +++++++++++++++++++++++++++++++++ tests/transfer_tests.py | 1008 +++++++++++++++++++++++++++++++++++++++ 4 files changed, 1889 insertions(+), 1 deletion(-) create mode 100644 clouddrift/transfer.py create mode 100644 tests/transfer_tests.py diff --git a/.codecov.yml b/.codecov.yml index 3fbef750..3ef58aa9 100644 --- a/.codecov.yml +++ b/.codecov.yml @@ -13,7 +13,10 @@ coverage: default: threshold: 10% -comment: off +comment: + require_changes: true + layout: "condensed_header, condensed_files, condensed_footer" + hide_project_coverage: TRUE # which folders/files to ignore ignore: diff --git a/clouddrift/__init__.py b/clouddrift/__init__.py index 0a873ec7..5d08e0a5 100644 --- a/clouddrift/__init__.py +++ b/clouddrift/__init__.py @@ -10,6 +10,7 @@ import clouddrift.ragged as ragged import clouddrift.signal as signal import clouddrift.sphere as sphere +import clouddrift.transfer as transfer import clouddrift.wavelet as wavelet from clouddrift.raggedarray import RaggedArray @@ -23,5 +24,6 @@ "ragged", "signal", "sphere", + "transfer", "wavelet", ] diff --git a/clouddrift/transfer.py b/clouddrift/transfer.py new file mode 100644 index 00000000..40c342cf --- /dev/null +++ b/clouddrift/transfer.py @@ -0,0 +1,875 @@ +""" +This module provides functions to calculate various transfer function from wind stress to oceanic +velocity. +""" + +import numpy as np +from numpy.lib.scimath import sqrt +from scipy.special import factorial, iv, kv # type: ignore + +from clouddrift.sphere import EARTH_DAY_SECONDS + + +def slab_wind_transfer( + omega: float | np.ndarray, + cor_freq: float, + friction: float, + bld: float, + density: float = 1025.0, +) -> np.ndarray: + """ + Compute the the transfer function in the case of a ocean slab layer. + + Parameters + ---------- + omega: array_like or float + Angular frequency of the wind stress forcing in radians per day. + cor_freq: float + Coriolis frequency, in radians per day. + friction: float + Friction coefficient, in s-1. + bld: float + Thickness of the slab layer, in meters. + density: float + Seawater density, in kg m-3. Default is 1025. + """ + # check that the boundary layer depth is positive + if bld < 0: + raise ValueError("Boundary layer depth bld must be positive.") + + # check that the density is positive + if density < 0: + raise ValueError("Density density must be positive.") + + # omega can be scalars or arrays, convert to arrays here + omega_ = np.atleast_1d(omega) + + # convert to radians per second + omega_ = omega_ / EARTH_DAY_SECONDS + cor_freq_ = cor_freq / EARTH_DAY_SECONDS + + G = 1 / (density * bld * (friction + 1j * (omega_ + cor_freq_))) + + return G + + +def wind_transfer( + omega: float | np.ndarray, + z: float | np.ndarray, + cor_freq: float, + delta: float, + mu: float, + bld: float, + boundary_condition: str = "no-slip", + method: str = "lilly", + density: float = 1025.0, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + Compute the transfer function from wind stress to oceanic velocity based on the physically-based + models of Elipot and Gille (2009) and Lilly and Elipot (2021). + + Parameters + ---------- + omega: array_like or float + Angular frequency of the wind stress forcing in radians per day. + z: array_like or float + Depth in the ocean, positive downwards, in meters. + cor_freq: float + Coriolis frequency, in radians per day. + delta: float + Ekman depth, in meters. + mu: float + Madsen depth, in meters. + bld: float + Boundary layer depth, in meters. + boundary_condition: str + Bottom boundary condition at the base of the ocean surface boundary layer. + Options are "no-slip" (Default) or "free-slip". + method: str + Method to compute the transfer function. Options are "lilly" (Default and preferred method) + or "elipot". + density: float + Seawater density, in kg m-3. Default is 1025. + + Returns + ------- + G: np.ndarray + The transfer function from wind stress to oceanic velocity in units of kg-1 m 2 s. + ddelta: np.ndarray + The gradient of transfer function from the rate of change of the Ekman depth. + dh: np.ndarray + The gradient of transfer function from to the rate of change of the boundary layer depth. + + Examples + -------- + To calculate the transfer function of Lilly and Elipot (2021) corresponding to the Ekman case + of infinite depth ocean surface boundary layer and constant vertical eddy viscosity, at the zero frequency: + + >>> omega = 0 + >>> z = np.linspace(0, 100, 100) + >>> cor_freq = 2 * np.pi * 1 + >>> K0 = 1e-4 + >>> delta = np.sqrt(2 * K0 / cor_freq / 86400) + >>> mu = 0 + >>> G = wind_transfer(omega, z, cor_freq, delta, mu, np.inf, "no-slip") + + Raises + ------ + ValueError + If the depth z is not positive. + + References + ---------- + [1] Elipot, S., and S. T. Gille (2009). Ekman layers in the Southern Ocean: spectral models and + observations, vertical viscosity and boundary layer depth Ocean Sci., 5, 115–139, 2009, doi:10.5194/os-5-115-2009. + + [2] Lilly, J. M. and S. Elipot (2021). A unifying perspective on transfer function solutions to + the unsteady Ekman problem. Fluids, 6 (2): 85, 1--36. doi:10.3390/fluids6020085. + + """ + # check that z is positive + if np.any(z < 0): + raise ValueError("Depth z must be positive.") + + # check that the boundary layer depth is positive + if bld < 0: + raise ValueError("Boundary layer depth bld must be positive.") + + # check that the Ekman depth is positive + if delta < 0: + raise ValueError("Ekman depth delta must be positive.") + + # check that the Madsen depth is positive + if mu < 0: + raise ValueError("Madsen depth mu must be positive.") + + # check that the density is positive + if density < 0: + raise ValueError("Density density must be positive.") + + # check that the boundary condition is valid + if boundary_condition not in ["no-slip", "free-slip"]: + raise ValueError("Boundary condition must be 'no-slip' or 'free-slip'.") + + # check that the method is valid + if method not in ["lilly", "elipot"]: + raise ValueError("Method must be 'lilly' or 'elipot'.") + + # z and omega can be scalars or arrays, convert to arrays here + z_ = np.atleast_1d(z) + omega_ = np.atleast_1d(omega) + + # convert to radians per second + omega_ = omega_ / EARTH_DAY_SECONDS + cor_freq_ = cor_freq / EARTH_DAY_SECONDS + + # Create the gridmesh of frequency and depth + [omega_grid, z_grid] = np.meshgrid(omega_, z_) + + if boundary_condition == "no-slip": + if method == "lilly": + G = _wind_transfer_no_slip( + omega_grid, + z_grid, + cor_freq_, + delta, + mu, + bld, + density, + ) + else: + G = _wind_transfer_elipot_no_slip( + omega_grid, + z_grid, + cor_freq_, + delta, + mu, + bld, + density, + ) + elif boundary_condition == "free-slip": + if method == "lilly": + G = _wind_transfer_free_slip( + omega_grid, + z_grid, + cor_freq_, + delta, + mu, + bld, + density, + ) + else: + G = _wind_transfer_elipot_free_slip( + omega_grid, + z_grid, + cor_freq_, + delta, + mu, + bld, + density, + ) + + # set G to nan where z > bld; may be mathematcially valid but not physically meaningful + G[z_grid > bld] = np.nan + + # analytical gradients of the transfer function for mu = 0 and free slip, lilly method + if boundary_condition == "no-slip" and method == "lilly" and mu == 0: + s = np.sign(cor_freq_) * np.sign(1 + omega_grid / cor_freq_) + Gamma = sqrt(2) * _rot(s * np.pi / 4) * sqrt(np.abs(1 + omega_grid / cor_freq_)) + ddelta1 = ( + (Gamma * (bld / delta) * np.tanh(Gamma * (bld / delta)) - 1) * G / delta + ) + + numer = np.exp(Gamma * (-z_grid / delta)) + np.exp( + -Gamma * (2 * bld - z_grid) / delta + ) + denom = 1 + np.exp(-Gamma * (2 * bld / delta)) + ddelta2 = ( + -2 + / (delta**2 * np.abs(cor_freq_) * density) + * (bld - z_grid) + / delta + * (numer / denom) + ) + + dG_ddelta = ddelta1 + ddelta2 + + dbld1 = -Gamma / delta * np.tanh(Gamma * (bld / delta)) * G + dbld2 = 2 / (delta**2 * np.abs(cor_freq_) * density) * (numer / denom) + dG_dbld = dbld1 + dbld2 + + else: + dG_ddelta = np.array([], dtype=float) # Empty array for consistency + dG_dbld = np.array([], dtype=float) + + return G, dG_ddelta, dG_dbld + + +def _wind_transfer_no_slip( + omega_grid: np.ndarray, + z_grid: np.ndarray, + coriolis_frequency: float, + delta: float, + mu: float, + bld: float, + density: float, +) -> np.ndarray: + """ + Transfer function from wind stress to oceanic velocity with no-slip boundary condition. + """ + + zo = np.divide(delta**2, mu) + s = np.sign(coriolis_frequency) * np.sign(1 + omega_grid / coriolis_frequency) + + xiz, xih, xi0 = _xis(s, zo, delta, z_grid, omega_grid, coriolis_frequency, bld) + + if bld == np.inf: + if mu == 0: + # Ekman solution + coeff = (sqrt(2) * _rot(-s * np.pi / 4)) / ( + delta * np.abs(coriolis_frequency) * density + ) + G = ( + coeff + * ( + np.exp( + -(1 + s * 1j) + * (z_grid / delta) + * sqrt(np.abs(1 + omega_grid / coriolis_frequency)) + ) + ) + / (sqrt(np.abs(1 + omega_grid / coriolis_frequency))) + ) + elif delta == 0: + # Madsen solution + coeff = 4 / (density * np.abs(coriolis_frequency) * mu) + G = coeff * kv( + 0, + 2 + * sqrt(2) + * _rot(s * np.pi / 4) + * sqrt((z_grid / mu) * np.abs(1 + omega_grid / coriolis_frequency)), + ) + else: + # mixed solution + k0z = kv(0, xiz) + k10 = kv(1, xi0) + coeff = (sqrt(2) * _rot(-s * np.pi / 4)) / ( + delta + * np.abs(coriolis_frequency) + * density + * sqrt(np.abs(1 + omega_grid / coriolis_frequency)) + ) + G = coeff * k0z / k10 + else: + if mu == 0: + # finite layer Ekman + coeff = (sqrt(2) * _rot(-s * np.pi / 4)) / ( + delta + * np.abs(coriolis_frequency) + * density + * sqrt(np.abs(1 + omega_grid / coriolis_frequency)) + ) + argh = ( + sqrt(2) + * _rot(s * np.pi / 4) + * (bld / delta) + * sqrt(np.abs(1 + omega_grid / coriolis_frequency)) + ) + argz = ( + sqrt(2) + * _rot(s * np.pi / 4) + * (z_grid / delta) + * sqrt(np.abs(1 + omega_grid / coriolis_frequency)) + ) + + numer = np.exp(-argz) - np.exp(argz) * np.exp(-2 * argh) + denom = 1 + np.exp(-2 * argh) + G = coeff * np.divide(numer, denom) + + # solution at inertial frequency + bool_idx = omega_grid == -coriolis_frequency + G[bool_idx] = 2 / ( + (density * np.abs(coriolis_frequency) * delta**2) + * (bld - z_grid[bool_idx]) + ) + + elif delta == 0: + # finite layer Madsen + coeff = 4 / (density * np.abs(coriolis_frequency) * mu) + argz = ( + 2 + * sqrt(2) + * _rot(s * np.pi / 4) + * sqrt((z_grid / mu) * np.abs(1 + omega_grid / coriolis_frequency)) + ) + argh = ( + 2 + * sqrt(2) + * _rot(s * np.pi / 4) + * sqrt((bld / mu) * np.abs(1 + omega_grid / coriolis_frequency)) + ) + k0z, i0z, k0h, i0h, _, _ = _bessels_noslip(argz, argh) + G = coeff * (k0z - i0z * k0h / i0h) + + # solution at inertial frequency + bool_idx = omega_grid == -coriolis_frequency + if isinstance(z_grid, np.ndarray): + G[bool_idx] = 0.5 * coeff * np.log(bld / z_grid[bool_idx]) + else: + G = 0.5 * coeff * np.log(bld / z_grid) + else: + # finite layer mixed + G = _wind_transfer_general_no_slip( + omega_grid, z_grid, coriolis_frequency, delta, mu, bld, density + ) + + return G + + +def _wind_transfer_general_no_slip( + omega: np.ndarray, + z: np.ndarray, + coriolis_frequency: float, + delta: float, + mu: float, + bld: float, + density: float, +) -> np.ndarray: + """ + Transfer function from wind stress to oceanic velocity with no-slip boundary condition, general form. + """ + zo = np.divide(delta**2, mu) + s = np.sign(coriolis_frequency) * np.sign(1 + omega / coriolis_frequency) + xiz, xih, xi0 = _xis(s, zo, delta, z, omega, coriolis_frequency, bld) + coeff = ( + sqrt(2) + * _rot(-s * np.pi / 4) + / ( + delta + * density + * np.abs(coriolis_frequency) + * sqrt(np.abs(1 + omega / coriolis_frequency)) + ) + ) + + k0z, i0z, k0h, i0h, k10, i10 = _bessels_noslip(xiz, xih, xi0=xi0) + numer = k0z / k10 - (k0h / k10) * (i0z / i0h) + denom = 1 + (k0h / k10) * (i10 / i0h) + G = coeff * np.divide(numer, denom) + + # large argument approximation + bool_idx = np.log10(np.abs(xiz)) > 2.9 + G[bool_idx] = _wind_transfer_general_no_slip_expansion( + omega[bool_idx], z[bool_idx], coriolis_frequency, delta, mu, bld, density + ) + + # inertial limit + G = _wind_transfer_inertiallimit( + G, omega, z, coriolis_frequency, delta, mu, bld, density + ) + + return G + + +def _wind_transfer_general_no_slip_expansion( + omega: np.ndarray, + z: np.ndarray, + coriolis_frequency: float, + delta: float, + mu: float, + bld: float, + density: float, +) -> np.ndarray: + """ + Compute the transfer function from wind stress to oceanic velocity with no-slip boundary, large argument approximation. + """ + zo = np.divide(delta**2, mu) + s = np.sign(coriolis_frequency) * np.sign(1 + omega / coriolis_frequency) + xiz, xih, xi0 = _xis(s, zo, delta, z, omega, coriolis_frequency, bld) + coeff = ( + sqrt(2) + * _rot(-s * np.pi / 4) + / ( + delta + * density + * np.abs(coriolis_frequency) + * sqrt(np.abs(1 + omega / coriolis_frequency)) + ) + ) + k0z, i0z, k0h, i0h, k10, i10 = _besseltildes_noslip(xiz, xih, xi0, 30) + + numer = np.exp(xi0 - xiz) * i0h * k0z - np.exp(xi0 + xiz - 2 * xih) * k0h * i0z + denom = i0h * k10 + np.exp(2 * xi0 - 2 * xih) * k0h * i10 + G = coeff * np.divide(numer, denom) + + bool_idx = omega == -coriolis_frequency + G[bool_idx] = ( + 4 + * zo + / ( + density + * np.abs(coriolis_frequency) + * delta**2 + * np.divide(sqrt(1 + bld / zo) - sqrt(1 + z / zo), (1 + z / zo) ** 0.25) + ) + ) + + return G + + +def _wind_transfer_inertiallimit( + G: np.ndarray, + omega: np.ndarray, + z: np.ndarray, + coriolis_freq: float, + delta: float, + mu: float, + bld: float, + density: float, +) -> np.ndarray: + """ + Transfer function from wind stress to oceanic velocity with no-slip boundary, inertial limit. + """ + zo = delta**2 / mu + bool_idx = omega == -coriolis_freq + + if np.any(bool_idx): + if not np.isinf(bld) and not np.isinf(zo): + G[bool_idx] = (2 / (density * np.abs(coriolis_freq) * mu)) * np.log( + (1 + bld / zo) / (1 + z[bool_idx] / zo) + ) + elif not np.isinf(bld) and np.isinf(zo): + G[bool_idx] = (2 / (density * np.abs(coriolis_freq) * delta**2)) * ( + bld - z[bool_idx] + ) + else: + G[bool_idx] = np.inf + + return G + + +def _wind_transfer_free_slip( + omega: np.ndarray, + z: np.ndarray, + coriolis_frequency: float, + delta: float, + mu: float, + bld: float, + density: float, +) -> np.ndarray: + """ + Transfer function from wind stress to oceanic velocity with free-slip boundary condition. + """ + + zo = np.divide(delta**2, mu) + + s = np.sign(coriolis_frequency) * np.sign(1 + omega / coriolis_frequency) + + xiz, xih, xi0 = _xis(s, zo, delta, z, omega, coriolis_frequency, bld) + + if delta != 0.0 and mu != 0.0: + coeff = ( + sqrt(2) + * _rot(-s * np.pi / 4) + / ( + delta + * density + * np.abs(coriolis_frequency) + * sqrt(np.abs(1 + omega / coriolis_frequency)) + ) + ) + k0z, i0z, k1h, i1h, k10, i10 = _bessels_freeslip(xiz, xih, xi0=xi0) + numer = i0z * k1h + i1h * k0z + denom = i1h * k10 - i10 * k1h + G = coeff * np.divide(numer, denom) + + elif mu == 0.0: + coeff = (sqrt(2) * _rot(-s * np.pi / 4)) / ( + delta + * np.abs(coriolis_frequency) + * density + * sqrt(np.abs(1 + omega / coriolis_frequency)) + ) + cosharg = ( + sqrt(2) + * _rot(s * np.pi / 4) + * (bld - z) + / delta + * sqrt(np.abs(1 + omega / coriolis_frequency)) + ) + sinharg = ( + sqrt(2) + * _rot(s * np.pi / 4) + * bld + / delta + * sqrt(np.abs(1 + omega / coriolis_frequency)) + ) + G = coeff * np.cosh(cosharg) / np.sinh(sinharg) + + elif delta == 0.0: + k0z, i0z, k1h, i1h, _, _ = _bessels_freeslip(xiz, xih) + + K1 = 0.5 * mu * np.abs(coriolis_frequency) + + coeff = 2 / (density * K1) + + G = coeff * (k0z + k1h * i0z / i1h) + + return G + + +def _wind_transfer_elipot_no_slip( + omega: np.ndarray, + z: np.ndarray, + coriolis_frequency: float, + delta: float, + mu: float, + bld: float, + density: float, +) -> np.ndarray: + """ + Transfer function from wind stress to oceanic velocity with no-slip boundary condition, Elipot method. + """ + + zo = np.divide(delta**2, mu) + K0 = 0.5 * delta**2 * np.abs(coriolis_frequency) + K1 = 0.5 * mu * np.abs(coriolis_frequency) + + delta1 = sqrt(2 * K0 / (omega + coriolis_frequency)) + delta2 = K1 / (omega + coriolis_frequency) + xiz = 2 * sqrt(1j * (zo + z) / delta2) + xi0 = 2 * sqrt(1j * zo / delta2) + xih = 2 * sqrt(1j * (zo + bld) / delta2) + + if bld == np.inf: + if K1 == 0: + # Ekman solution + coeff = 1 / (density * sqrt(1j * (omega + coriolis_frequency) * K0)) + G = coeff * np.exp(-z * (1 + 1j) / delta1) + elif K0 == 0: + # Madsen solution + coeff = 2 / (density * K1) + k0z = kv(0, 2 * sqrt(1j * z / delta2)) + G = coeff * k0z + else: + # Mixed solution + coeff = 1 / (density * sqrt(1j * (omega + coriolis_frequency) * K0)) + G = coeff * kv(0, xiz) / kv(1, xi0) + else: + if K1 == 0: + # Finite-layer Ekman + coeff = 1 / (density * sqrt(1j * (omega + coriolis_frequency) * K0)) + numer = np.sinh((1 + 1j) * (bld - z) / delta1) + denom = np.cosh((1 + 1j) * bld / delta1) + G = coeff * numer / denom + elif K0 == 0: + # Finite-layer Madsen + coeff = 2 / (density * K1) + k0z = kv(0, 2 * sqrt(1j * z / delta2)) + k0h = kv(0, 2 * sqrt(1j * bld / delta2)) + i0z = iv(0, 2 * sqrt(1j * z / delta2)) + i0h = iv(0, 2 * sqrt(1j * bld / delta2)) + G = coeff * (k0z - (k0h * i0z / i0h)) + else: + # Finite-layer mixed solution + coeff = 1 / (density * sqrt(1j * (omega + coriolis_frequency) * K0)) + k0z, i0z, k0h, i0h, k10, i10 = _bessels_noslip(xiz, xih, xi0) + numer = i0h * k0z - k0h * i0z + denom = i10 * k0h + k10 * i0h + G = coeff * numer / denom + + return G + + +def _wind_transfer_elipot_free_slip( + omega: np.ndarray, + z: np.ndarray, + coriolis_frequency: float, + delta: float, + mu: float, + bld: float, + density: float, +) -> np.ndarray: + """ + Transfer function from wind stress to oceanic velocity with free-slip boundary condition, Elipot method. + """ + + zo = np.divide(delta**2, mu) + K0 = 0.5 * delta**2 * np.abs(coriolis_frequency) + K1 = 0.5 * mu * np.abs(coriolis_frequency) + delta1 = sqrt(2 * K0 / (omega + coriolis_frequency)) + delta2 = K1 / (omega + coriolis_frequency) + xiz = 2 * sqrt(1j * (zo + z) / delta2) + xi0 = 2 * sqrt(1j * zo / delta2) + xih = 2 * sqrt(1j * (zo + bld) / delta2) + + if K1 == 0.0: + # Finite-layer Ekman + coeff = 1 / (density * sqrt(1j * (omega + coriolis_frequency) * K0)) + numer = np.cosh((1 + 1j) * (bld - z) / delta1) + denom = np.sinh((1 + 1j) * bld / delta1) + G = coeff * numer / denom + elif K0 == 0.0: + # Finite-layer Madsen + coeff = 2 / (density * K1) + k0z = kv(0, 2 * sqrt(1j * z / delta2)) + k0h = kv(0, 2 * sqrt(1j * bld / delta2)) + i0z = iv(0, 2 * sqrt(1j * z / delta2)) + i1h = iv(1, 2 * sqrt(1j * bld / delta2)) + G = coeff * (k0z - (k0h * i0z / i1h)) + else: + # Finite-layer mixed solution + coeff = 1 / (density * sqrt(1j * (omega + coriolis_frequency) * K0)) + k1h = kv(1, xih) + i0z = iv(0, xiz) + i1h = iv(1, xih) + k10 = kv(1, xi0) + i10 = iv(1, xi0) + k0z = kv(0, xiz) + numer = k1h * i0z + i1h * k0z + denom = i1h * k10 - i10 * k1h + G = coeff * numer / denom + + return G + + +def _bessels_freeslip( + xiz: float | np.ndarray, + xih: float | np.ndarray, + xi0: float | np.ndarray | None | None = None, +) -> tuple[ + float | np.ndarray, + float | np.ndarray, + float | np.ndarray, + float | np.ndarray, + float | np.ndarray, + float | np.ndarray, +]: + """ + Compute the Bessel functions for the free-slip boundary condition for the xsi(z), xsi(h), and xsi(0) functions. + """ + k0z = kv(0, xiz) + i0z = iv(0, xiz) + k1h = kv(1, xih) + i1h = iv(0, xih) + + if xi0 is not None: + k10 = kv(1, xi0) + i10 = iv(1, xi0) + return k0z, i0z, k1h, i1h, k10, i10 + else: + return k0z, i0z, k1h, i1h, np.nan, np.nan + + +def _bessels_noslip( + xiz: float | np.ndarray, + xih: float | np.ndarray, + xi0: float | np.ndarray | None = None, +) -> tuple[ + float | np.ndarray, + float | np.ndarray, + float | np.ndarray, + float | np.ndarray, + float | np.ndarray, + float | np.ndarray, +]: + """ + Compute the Bessel functions for the no-slip boundary condition for the xsi(z), xsi(h), and xsi(0) functions. + """ + k0z = kv(0, xiz) + i0z = iv(0, xiz) + k0h = kv(0, xih) + i0h = iv(0, xih) + + if xi0 is not None: + k10 = kv(1, xi0) + i10 = iv(1, xi0) + return k0z, i0z, k0h, i0h, k10, i10 + else: + return k0z, i0z, k0h, i0h, np.nan, np.nan + + +def _besseltildes_noslip( + xiz: float | np.ndarray, + xih: float | np.ndarray, + xi0: float | np.ndarray, + nterms: int = 30, +) -> tuple[ + float | np.ndarray, + float | np.ndarray, + float | np.ndarray, + float | np.ndarray, + float | np.ndarray, + float | np.ndarray, +]: + """ + Compute the n-term expansion about the large-argument exponential behavior of Bessel functions + for the xsi(z), xsi(h), and xsi(0) functions. + """ + k0z = kvtilde(0, xiz, nterms) + i0z = ivtilde(0, xiz, nterms) + k0h = kvtilde(0, xih, nterms) + i0h = ivtilde(0, xih, nterms) + k10 = kvtilde(1, xi0, nterms) + i10 = ivtilde(1, xi0, nterms) + + return k0z, i0z, k0h, i0h, k10, i10 + + +def kvtilde( + nu: int, + z: float | np.ndarray, + nterms: int = 30, +) -> float | np.ndarray: + """ + Compute the n-term expansion about the large-argument exponential behavior of the modified Bessel + function of the second kind of real order (kv). + """ + z = np.asarray( + z, dtype=np.complex128 + ) # Ensure z is an array for vectorized operations + sizez = z.shape + z = z.ravel() + + # Prepare zk for vectorized computation + zk = np.tile(z, (nterms, 1)).T + zk[:, 0] = 1 + zk = np.cumprod(zk, axis=1) + + k = np.arange(nterms) + ak = 4 * nu**2 - (2 * k - 1) ** 2 + ak[0] = 1 + ak = np.cumprod(ak) / (factorial(k) * (8**k)) + + # Handling potential non-finite values in ak + if not np.all(np.isfinite(ak)): + first_nonfinite = np.where(~np.isfinite(ak))[0][0] + ak = ak[:first_nonfinite] + zk = zk[:, :first_nonfinite] + + K = sqrt(np.pi / (2 * z)) * (np.dot(1.0 / zk, ak)) + K = K.reshape(sizez) + + return K + + +def ivtilde( + nu: int, + z: float | np.ndarray, + nterms: int = 30, +) -> float | np.ndarray: + """ + Compute the n-term expansion about the large-argument exponential behavior of the + Modified Bessel function of the first kind of real order (iv). + """ + z = np.asarray( + z, dtype=np.complex128 + ) # Ensure z is an array for vectorized operations + sizez = z.shape + z = z.ravel() + + # Prepare zk for vectorized computation with alternating signs for each term + zk = np.tile(-z, (nterms, 1)).T + zk[:, 0] = 1 + zk = np.cumprod(zk, axis=1) + + k = np.arange(nterms) + ak = 4 * nu**2 - (2 * k - 1) ** 2 + ak[0] = 1 + ak = np.cumprod(ak) / (factorial(k) * (8**k)) + + # Handling potential non-finite values in ak + if not np.all(np.isfinite(ak)): + first_nonfinite = np.where(~np.isfinite(ak))[0][0] + ak = ak[:first_nonfinite] + zk = zk[:, :first_nonfinite] + + I = 1 / sqrt(2 * z * np.pi) * (np.dot(1.0 / zk, ak)) + I = I.reshape(sizez) + + return I + + +def _xis( + s: float, + zo: float, + delta: float, + z: float | np.ndarray, + omega: float | np.ndarray, + coriolis_frequency: float, + bld: float, +) -> tuple[float | np.ndarray, float | np.ndarray, float | np.ndarray]: + """ + Compute the complex-valued "xsi" functions. + """ + xiz = ( + 2 + * sqrt(2) + * _rot(s * np.pi / 4) + * np.divide(zo, delta) + * sqrt((1 + np.divide(z, zo)) * np.abs(1 + omega / coriolis_frequency)) + ) + xih = ( + 2 + * sqrt(2) + * _rot(s * np.pi / 4) + * np.divide(zo, delta) + * sqrt((1 + np.divide(bld, zo)) * np.abs(1 + omega / coriolis_frequency)) + ) + xi0 = ( + 2 + * sqrt(2) + * _rot(s * np.pi / 4) + * np.divide(zo, delta) + * sqrt(np.abs(1 + omega / coriolis_frequency)) + ) + + return xiz, xih, xi0 + + +def _rot(x: float) -> complex: + """ + Compute the complex-valued rotation. + """ + return np.exp(1j * x) diff --git a/tests/transfer_tests.py b/tests/transfer_tests.py new file mode 100644 index 00000000..3cc16605 --- /dev/null +++ b/tests/transfer_tests.py @@ -0,0 +1,1008 @@ +import unittest + +import numpy as np +from numpy.lib.scimath import sqrt +from scipy.special import iv, kv # type: ignore + +from clouddrift.sphere import EARTH_DAY_SECONDS +from clouddrift.transfer import ( + _rot, + _xis, + ivtilde, + kvtilde, + wind_transfer, +) + +if __name__ == "__main__": + unittest.main() + + +class TransferFunctionTestShapes(unittest.TestCase): + def setUp(self): + self.omega = ( + 2 * np.pi * np.array([-2, -1.5, -0.5, 0, 0.5, 1, 1.5, 2]) + ) # Angular frequency in radians per day + self.z = np.array([0, 10, 20, 30]) # Depth in meters + self.cor_freq = 2 * np.pi * 1.5 # Coriolis frequency in radians per day + self.delta = 10 # Ekman depth in meters + self.mu = 0 # Madsen depth in meters + self.bld = 50 # Boundary layer depth in meters + self.density = 1025 # Seawater density in kg/m^3 + + def test_output_size_wind_transfer_no_slip_mu_is_0(self): + G, dG1, dG2 = wind_transfer( + self.omega, + self.z, + self.cor_freq, + self.delta, + self.mu, + self.bld, + boundary_condition="no-slip", + density=self.density, + ) + self.assertEqual(G.shape, (len(self.z), len(self.omega))) + self.assertEqual(dG1.shape, (len(self.z), len(self.omega))) + self.assertEqual(dG2.shape, (len(self.z), len(self.omega))) + + def test_output_size_wind_transfer_no_slip_mu_not_0(self): + G, dG1, dG2 = wind_transfer( + self.omega, + self.z, + self.cor_freq, + self.delta, + 5, + self.bld, + boundary_condition="no-slip", + density=self.density, + ) + self.assertEqual(G.shape, (len(self.z), len(self.omega))) + self.assertEqual(dG1.shape, (0,)) + self.assertEqual(dG2.shape, (0,)) + + G, dG1, dG2 = wind_transfer( + self.omega, + self.z, + self.cor_freq, + self.delta, + 5, + self.bld, + boundary_condition="no-slip", + density=self.density, + method="elipot", + ) + self.assertEqual(G.shape, (len(self.z), len(self.omega))) + self.assertEqual(dG1.shape, (0,)) + self.assertEqual(dG2.shape, (0,)) + + def test_output_size_wind_transfer_free_slip(self): + G, dG1, dG2 = wind_transfer( + self.omega, + self.z, + self.cor_freq, + self.delta, + 5, + self.bld, + boundary_condition="free-slip", + density=self.density, + ) + self.assertEqual(G.shape, (len(self.z), len(self.omega))) + self.assertEqual(dG1.shape, (0,)) + self.assertEqual(dG2.shape, (0,)) + + +class TransferFunctionTestInputs(unittest.TestCase): + def setUp(self): + self.omega = 2 * np.pi * np.array([-2, -1.5, -0.5, 0, 0.5, 1, 1.5, 2]) + self.z = np.array([0, 10, 20, 30]) + self.cor_freq = 2 * np.pi * 1.5 + self.delta = 10 + self.mu = 5 + self.bld = 50 + self.density = 1025 + + def test_wind_transfer_negative_z(self): + with self.assertRaises(ValueError): + wind_transfer( + self.omega, + -self.z, + self.cor_freq, + self.delta, + self.mu, + self.bld, + density=self.density, + ) + + def test_wind_transfer_negative_delta(self): + with self.assertRaises(ValueError): + wind_transfer( + self.omega, + self.z, + self.cor_freq, + -self.delta, + self.mu, + self.bld, + density=self.density, + ) + + def test_wind_transfer_negative_mu(self): + with self.assertRaises(ValueError): + wind_transfer( + self.omega, + self.z, + self.cor_freq, + self.delta, + -self.mu, + self.bld, + density=self.density, + ) + + def test_wind_transfer_negative_bld(self): + with self.assertRaises(ValueError): + wind_transfer( + self.omega, + self.z, + self.cor_freq, + self.delta, + self.mu, + -self.bld, + density=self.density, + ) + + def test_wind_transfer_negative_density(self): + with self.assertRaises(ValueError): + wind_transfer( + self.omega, + self.z, + self.cor_freq, + self.delta, + self.mu, + self.bld, + density=-self.density, + ) + + def test_wind_transfer_boundary_condition(self): + with self.assertRaises(ValueError): + wind_transfer( + self.omega, + self.z, + self.cor_freq, + self.delta, + self.mu, + self.bld, + boundary_condition="invalid", + density=self.density, + ) + + def test_wind_transfer_method(self): + with self.assertRaises(ValueError): + wind_transfer( + self.omega, + self.z, + self.cor_freq, + self.delta, + self.mu, + self.bld, + method="invalid", + density=self.density, + ) + + +class TransferFunctionSurfaceValues(unittest.TestCase): + def setUp(self): + self.omega = ( + 2 * np.pi * np.array([-2, -1.5, -0.5, 0, 0.5, 1, 1.5, 2]) + ) # Angular frequency in radians per day + self.z = np.array([0, 10, 20, 30]) # Depth in meters + self.cor_freq = 2 * np.pi * 1.5 # Coriolis frequency in radians per day + self.delta = 10 # Ekman depth in meters + self.mu = 0 # Madsen depth in meters + self.bld = 50 # Boundary layer depth in meters + self.density = 1025 # Seawater density in kg/m^3 + + def test_surface_angle_wind_transfer_no_slip_lilly_nh(self): + G, _, _ = wind_transfer( + self.omega, + self.z[0], + self.cor_freq, + self.delta, + self.mu, + self.bld, + boundary_condition="no-slip", + density=self.density, + ) + self.assertTrue( + np.allclose( + np.angle(G), + np.pi / 4 * np.array([1, 0, -1, -1, -1, -1, -1, -1]), + atol=1e-1, + ) + ) + + def test_surface_angle_wind_transfer_no_slip_elipot_nh(self): + G, _, _ = wind_transfer( + self.omega, + self.z[0], + self.cor_freq, + self.delta, + self.mu, + self.bld, + boundary_condition="no-slip", + density=self.density, + method="elipot", + ) + # except at inertial frequency which is not defined for elipot case + self.assertTrue( + np.allclose( + np.angle(np.delete(G, 1)), + np.pi / 4 * np.array([1, -1, -1, -1, -1, -1, -1]), + atol=1e-1, + ) + ) + + def test_surface_angle_wind_transfer_no_slip_lilly_sh(self): + G, _, _ = wind_transfer( + self.omega, + self.z[0], + -self.cor_freq, + self.delta, + self.mu, + self.bld, + boundary_condition="no-slip", + density=self.density, + ) + self.assertTrue( + np.allclose( + np.angle(G), + np.pi / 4 * np.array([1, 1, 1, 1, 1, 1, 0, -1]), + atol=1e-1, + ) + ) + + def test_surface_angle_wind_transfer_no_slip_elipot_sh(self): + G, _, _ = wind_transfer( + self.omega, + self.z[0], + -self.cor_freq, + self.delta, + self.mu, + self.bld, + boundary_condition="no-slip", + density=self.density, + method="elipot", + ) + # except at inertial frequency which is not defined for elipot case + self.assertTrue( + np.allclose( + np.angle(np.delete(G, 6)), + np.pi / 4 * np.array([1, 1, 1, 1, 1, 1, -1]), + atol=1e-1, + ) + ) + + +class TransferFunctionValues(unittest.TestCase): + # test for values reported in Lilly and Elipot 2021 + def setUp(self): + self.omega = np.arange(-10, 10 + 0.01, 0.01) * 2 * np.pi + self.z = np.arange(0, 50, 5) + self.cor_freq = 2 * np.pi * 1.5 + self.delta = 20 + self.z0 = 20 + self.mu = self.delta**2 / self.z0 + self.bld = 50 + self.density = 1025 + + def test_values_finite_bld(self): + idx = np.abs(self.omega + self.cor_freq).argmin() + G, _, _ = wind_transfer( + self.omega, + self.z, + self.cor_freq, + self.delta, + self.mu, + self.bld, + boundary_condition="no-slip", + density=self.density, + ) + + expected_values = np.array( + [ + 0.1252763 + 0 * 1j, + 0.10296194 + 0 * 1j, + 0.08472979 + 0 * 1j, + 0.06931472 + 0 * 1j, + 0.05596158 + 0 * 1j, + 0.04418328 + 0 * 1j, + 0.03364722 + 0 * 1j, + 0.02411621 + 0 * 1j, + 0.01541507 + 0 * 1j, + 0.0074108 + 0 * 1j, + ] + ) + Gp = G * np.abs(self.cor_freq / EARTH_DAY_SECONDS) * self.density + self.assertTrue(np.allclose(Gp[:, idx], expected_values, atol=1e-8)) + + def test_values_infinite_bld(self): + idx = np.abs(self.omega).argmin() + G, _, _ = wind_transfer( + self.omega, + self.z, + self.cor_freq, + self.delta, + self.mu, + np.inf, + boundary_condition="no-slip", + density=self.density, + ) + + expected_values = np.array( + [ + 0.04850551 - 0.0396564j, + 0.02829943 - 0.03744607j, + 0.01520417 - 0.03297259j, + 0.00668176 - 0.02797499j, + 0.00117548 - 0.02317678j, + -0.00230945 - 0.01886255j, + -0.00442761 - 0.01511781j, + -0.00561905 - 0.01193716j, + -0.0061842 - 0.00927553j, + -0.00633045 - 0.00707324j, + ] + ) + Gp = G * np.abs(self.cor_freq / EARTH_DAY_SECONDS) * self.density + self.assertTrue(np.allclose(Gp[:, idx], expected_values, atol=1e-8)) + + def test_values_infinite_bld_mu_is_zero(self): + idx = np.abs(self.omega).argmin() + G, _, _ = wind_transfer( + self.omega, + self.z, + self.cor_freq, + self.delta, + 0, + np.inf, + boundary_condition="no-slip", + density=self.density, + ) + + expected_values = np.array( + [ + 0.05 - 0.05j, + 0.02809557 - 0.04736341j, + 0.01207472 - 0.04115335j, + 0.0011821 - 0.03338043j, + -0.00553969 - 0.0254163j, + -0.00907736 - 0.0181115j, + -0.01033938 - 0.01191774j, + -0.01009828 - 0.00700083j, + -0.00896897 - 0.00333703j, + -0.00741087 - 0.00078996j, + ] + ) + Gp = G * np.abs(self.cor_freq / EARTH_DAY_SECONDS) * self.density + self.assertTrue(np.allclose(Gp[:, idx], expected_values, atol=1e-8)) + + def test_values_infinite_bld_delta_is_zero(self): + idx = np.abs(self.omega).argmin() + G, _, _ = wind_transfer( + self.omega, + self.z, + self.cor_freq, + 0, + self.mu, + np.inf, + boundary_condition="no-slip", + density=self.density, + ) + + expected_values = np.array( + [ + np.nan + np.nan * 1j, + 0.01603955 - 0.07145549j, + -0.0083329 - 0.04048001j, + -0.01373573 - 0.02367659j, + -0.01399476 - 0.01368128j, + -0.01260713 - 0.00748527j, + -0.01076378 - 0.0035753j, + -0.00891647 - 0.00110311j, + -0.00723577 + 0.00043968j, + -0.00577549 + 0.00137177j, + ] + ) + Gp = G * np.abs(self.cor_freq / EARTH_DAY_SECONDS) * self.density + self.assertTrue( + np.allclose(Gp[:, idx], expected_values, atol=1e-8, equal_nan=True) + ) + + def test_values_finite_bld_mu_is_zero(self): + idx = np.abs(self.omega).argmin() + G, _, _ = wind_transfer( + self.omega, + self.z, + self.cor_freq, + self.delta, + 0, + self.bld, + boundary_condition="no-slip", + density=self.density, + ) + + expected_values = np.array( + [ + 0.04916146 - 0.05044871j, + 0.02728561 - 0.04786423j, + 0.01135701 - 0.04180688j, + 0.00063926 - 0.03427553j, + -0.00579508 - 0.02661964j, + -0.00889455 - 0.01965344j, + -0.00952588 - 0.01377343j, + -0.00842319 - 0.00906617j, + -0.00617632 - 0.00539995j, + -0.00324643 - 0.00249871j, + ] + ) + Gp = G * np.abs(self.cor_freq / EARTH_DAY_SECONDS) * self.density + print(Gp[:, idx]) + self.assertTrue( + np.allclose(Gp[:, idx], expected_values, atol=1e-8, equal_nan=True) + ) + + def test_values_finite_bld_delta_is_zero(self): + idx = np.abs(self.omega).argmin() + G, _, _ = wind_transfer( + self.omega, + self.z, + self.cor_freq, + 0, + self.mu, + self.bld, + boundary_condition="no-slip", + density=self.density, + ) + + expected_values = np.array( + [ + np.nan + np.nan * 1j, + 0.0150223 - 0.07199065j, + -0.00915018 - 0.04153115j, + -0.01422336 - 0.02519934j, + -0.01402997 - 0.0156096j, + -0.0120768 - 0.00973228j, + -0.00956707 - 0.00603436j, + -0.00696735 - 0.00364922j, + -0.00446539 - 0.00205173j, + -0.00213448 - 0.0009083j, + ] + ) + Gp = G * np.abs(self.cor_freq / EARTH_DAY_SECONDS) * self.density + self.assertTrue( + np.allclose(Gp[:, idx], expected_values, atol=1e-8, equal_nan=True) + ) + + def test_values_free_slip(self): + idx = np.abs(self.omega).argmin() + G, _, _ = wind_transfer( + self.omega, + self.z, + self.cor_freq, + self.delta, + self.mu, + self.bld, + boundary_condition="free-slip", + density=self.density, + ) + + expected_values = np.array( + [ + 0.04653128 - 0.0365929j, + 0.02616164 - 0.03449013j, + 0.01266338 - 0.03030143j, + 0.00359447 - 0.02574129j, + -0.00253419 - 0.02152846j, + -0.00666497 - 0.01794904j, + -0.00940876 - 0.01509089j, + -0.01116722 - 0.01294936j, + -0.01220577 - 0.01147741j, + -0.01269919 - 0.01061011j, + ] + ) + Gp = G * np.abs(self.cor_freq / EARTH_DAY_SECONDS) * self.density + self.assertTrue( + np.allclose(Gp[:, idx], expected_values, atol=1e-8, equal_nan=True) + ) + + def test_values_free_slip_mu_is_zero(self): + idx = np.abs(self.omega).argmin() + G, _, _ = wind_transfer( + self.omega, + self.z, + self.cor_freq, + self.delta, + 0, + self.bld, + boundary_condition="free-slip", + density=self.density, + ) + + expected_values = np.array( + [ + 0.05083587 - 0.04953873j, + 0.02890206 - 0.0468502j, + 0.01278664 - 0.04048806j, + 0.00171537 - 0.03247495j, + -0.00529895 - 0.02420514j, + -0.00928086 - 0.01656615j, + -0.01118 - 0.01006565j, + -0.01180644 - 0.00494948j, + -0.01179884 - 0.00130261j, + -0.01161306 + 0.00087116j, + ] + ) + Gp = G * np.abs(self.cor_freq / EARTH_DAY_SECONDS) * self.density + self.assertTrue( + np.allclose(Gp[:, idx], expected_values, atol=1e-8, equal_nan=True) + ) + + # def test_values_free_slip_delta_is_zero(self): + # idx = np.abs(self.omega).argmin() + # G, _, _ = wind_transfer( + # self.omega, + # self.z, + # self.cor_freq, + # 0, + # self.mu, + # self.bld, + # boundary_condition="free-slip", + # density=self.density, + # ) + + # expected_values = np.array( + # [ + # np.nan + np.nan * 1j, + # 0.0150223 - 0.07199065j, + # -0.00915018 - 0.04153115j, + # -0.01422336 - 0.02519934j, + # -0.01402997 - 0.0156096j, + # -0.0120768 - 0.00973228j, + # -0.00956707 - 0.00603436j, + # -0.00696735 - 0.00364922j, + # -0.00446539 - 0.00205173j, + # -0.00213448 - 0.0009083j, + # ] + # ) + # Gp = G * np.abs(self.cor_freq / EARTH_DAY_SECONDS) * self.density + # self.assertTrue( + # np.allclose(Gp[:, idx], expected_values, atol=1e-8, equal_nan=True) + # ) + + +class TransferFunctionTestGradient(unittest.TestCase): + delta = 10 ** np.arange(-1, 0.05, 3) + bld = 10 ** np.arange(np.log10(15.15), 5, 0.05) + [delta_grid, bld_grid] = np.meshgrid(delta, bld) + + def test_gradient_ekman_case(self): + # Test the gradient of the transfer function, no-slip + omega = np.array([1e-4]) + z = 15 + cor_freq = 1e-4 + mu = 0 + delta_delta = 1e-6 + delta_bld = 1e-6 + # initialize transfer functions and gradients + wind_transfer_init = np.zeros((len(self.delta), len(self.bld)), dtype=complex) + wind_transfer_ddelta_plus = np.zeros( + (len(self.delta), len(self.bld)), dtype=complex + ) + wind_transfer_ddelta_minus = np.zeros( + (len(self.delta), len(self.bld)), dtype=complex + ) + wind_transfer_dbld_plus = np.zeros( + (len(self.delta), len(self.bld)), dtype=complex + ) + wind_transfer_dbld_minus = np.zeros( + (len(self.delta), len(self.bld)), dtype=complex + ) + dG_ddelta = np.zeros((len(self.delta), len(self.bld)), dtype=complex) + dG_dbld = np.zeros((len(self.delta), len(self.bld)), dtype=complex) + + for i in range(len(self.delta)): + for j in range(len(self.bld)): + wind_transfer_init[i, j], dG_ddelta[i, j], dG_dbld[i, j] = ( + wind_transfer( + omega=omega, + z=z, + cor_freq=cor_freq, + delta=self.delta[i], + mu=mu, + bld=self.bld[j], + ) + ) + wind_transfer_ddelta_plus[i, j], _, _ = wind_transfer( + omega=omega, + z=z, + cor_freq=cor_freq, + delta=self.delta[i] + delta_delta / 2, + mu=mu, + bld=self.bld[j], + ) + wind_transfer_ddelta_minus[i, j], _, _ = wind_transfer( + omega=omega, + z=z, + cor_freq=cor_freq, + delta=self.delta[i] - delta_delta / 2, + mu=mu, + bld=self.bld[j], + ) + wind_transfer_dbld_plus[i, j], _, _ = wind_transfer( + omega=omega, + z=z, + cor_freq=cor_freq, + delta=self.delta[i], + mu=mu, + bld=self.bld[j] + delta_bld / 2, + ) + wind_transfer_dbld_minus[i, j], _, _ = wind_transfer( + omega=omega, + z=z, + cor_freq=cor_freq, + delta=self.delta[i], + mu=mu, + bld=self.bld[j] - delta_bld / 2, + ) + + dG_ddelta_fd = ( + wind_transfer_ddelta_plus - wind_transfer_ddelta_minus + ) / delta_delta + dG_dbld_fd = (wind_transfer_dbld_plus - wind_transfer_dbld_minus) / delta_bld + + bool_indices = dG_ddelta_fd != 0 + # Calculate eps1 using numpy operations + eps1 = np.max( + ( + np.abs(dG_ddelta_fd[bool_indices] - dG_ddelta[bool_indices]) + / sqrt( + np.abs(dG_ddelta_fd[bool_indices]) ** 2 + + np.abs(dG_ddelta[bool_indices]) ** 2 + ) + ) + ) + + bool_indices = dG_dbld_fd != 0 + # Calculate eps2 using numpy operations + eps2 = np.max( + ( + np.abs(dG_dbld_fd[bool_indices] - dG_dbld[bool_indices]) + / sqrt( + np.abs(dG_dbld_fd[bool_indices]) ** 2 + + np.abs(dG_dbld[bool_indices]) ** 2 + ) + ) + ) + self.assertTrue( + np.log10(eps1) < -4 and np.log10(eps2) < -4, + "wind_transfer analytic and numerical gradients match for Ekman case", + ) + + +# class TransferFunctionTestMethods(unittest.TestCase): +# def setUp(self): +# self.z = 15.0 # np.arange(0.1, 101, 1) +# # need to also test bld = np.inf, K0 = 0 and K1 = 0 +# self.bld = [np.inf, 200.0] +# self.K0 = [0, 1/10, 1/10] +# self.K1 = [1,0,1] +# self.cor_freq = 2 * np.pi * 1.5 +# self.delta = sqrt(2 * np.array(self.K0) / self.cor_freq) +# self.mu = 2 * np.array(self.K1) / self.cor_freq +# self.omega = 2 * np.pi * np.array([0.5]) +# self.slipstr1 = "no-slip" +# self.slipstr2 = "free-slip" + +# def test_method_equivalence_no_slip(self): +# for s in [1, -1]: +# for i, (delta, mu) in enumerate(zip(self.delta, self.mu)): +# for j, bld in enumerate(self.bld): +# Ge, _, _ = wind_transfer( +# self.omega, +# self.z, +# self.cor_freq, +# delta, +# mu, +# bld, +# method="elipot", +# boundary_condition=self.slipstr1, +# ) +# Gl, _, _ = wind_transfer( +# self.omega, +# self.z, +# self.cor_freq, +# delta, +# mu, +# bld, +# method="lilly", +# boundary_condition=self.slipstr1, +# ) +# #bool_idx = Ge != np.nan and Gl != np.nan +# #print(Ge[bool_idx], Gl[bool_idx]) +# print(Ge, Gl) +# #bool1 = np.allclose(Ge[bool_idx], Gl[bool_idx], atol=1e-8) +# bool1 = np.allclose(Ge, Gl, atol=1e-8, equal_nan=True) +# self.assertTrue(bool1) + +# def test_method_equivalence_free_slip(self): +# for s in [1, -1]: +# for i, (delta, mu) in enumerate(zip(self.delta, self.mu)): +# for j, bld in enumerate(self.bld): +# Ge, _, _ = wind_transfer( +# self.omega, +# self.z, +# self.cor_freq, +# delta, +# mu, +# bld, +# method="elipot", +# boundary_condition=self.slipstr2, +# ) +# Gl, _, _ = wind_transfer( +# self.omega, +# self.z, +# self.cor_freq, +# delta, +# mu, +# bld, +# method="lilly", +# boundary_condition=self.slipstr2, +# ) +# #bool_idx = Ge != np.nan and Gl != np.nan +# #print(Ge, Gl) +# #bool1 = np.allclose(Ge[bool_idx], Gl[bool_idx], atol=1e-8) +# bool1 = np.allclose(Ge, Gl, atol=1e-8, equal_nan=True) +# self.assertTrue(bool1) + + +class TestKvTilde(unittest.TestCase): + def test_kvtilde(self): + atol = 1e-10 + for s in [1, -1]: + z = sqrt(s * 1j) * np.arange(15, 100.01, 0.01).reshape(-1, 1) + bk0 = np.zeros((len(z), 2), dtype=np.complex128) + bk = np.zeros_like(bk0) + + for i in range(2): + bk0[:, i] = (np.exp(z) * kv(i, z)).reshape(-1) + bk[:, i] = (kvtilde(i, z)).reshape(-1) + + if s == 1: + test_name = "kvTILDE for z with phase of pi/4" + else: + test_name = "kvTILDE for z with phase of -pi/4" + + with self.subTest(test_name=test_name): + self.assertTrue( + np.allclose(np.abs((bk0 - bk) / bk), 0, atol=atol), + msg=f"Failed: {test_name}", + ) + + for i in range(2): + bk[:, i] = kvtilde(i, z, 2).reshape(-1) + + if s == 1: + test_name = "2-term for z with phase of pi/4, order 0 and 1" + else: + test_name = "2-term for z with phase of -pi/4, order 0 and 1" + + with self.subTest(test_name=test_name): + self.assertTrue( + np.allclose(np.abs((bk0[:, :2] - bk) / bk), 0, atol=1e-3), + msg=f"Failed: {test_name}", + ) + + bk0 = sqrt(np.pi / (2 * z)) * (1 - 1 / (8 * z)) + bk1 = sqrt(np.pi / (2 * z)) * (1 + 3 / (8 * z)) + + if s == 1: + test_name = "2-term vs. analytic for z with phase of pi/4" + else: + test_name = "2-term vs. analytic for z with phase of -pi/4" + + with self.subTest(test_name=test_name): + self.assertTrue( + np.allclose( + np.abs((np.hstack([bk0, bk1]) - bk) / bk), 0, atol=1e-15 + ), + msg=f"Failed: {test_name}", + ) + + +class TestIvTilde(unittest.TestCase): + def test_ivtilde(self): + atol = 1e-10 + for s in [1, -1]: + z = sqrt(s * 1j) * np.arange(23.0, 100.0, 0.01).reshape(-1, 1) + bi0 = np.zeros((len(z), 2), dtype=np.complex128) + bi = np.zeros_like(bi0) + + for i in range(2): + bi0[:, i] = (np.exp(-z) * iv(i, z)).reshape(-1) + bi[:, i] = (ivtilde(i, z)).reshape(-1) + + if s == 1: + test_name = "ivTILDE for z with phase of pi/4" + else: + test_name = "ivTILDE for z with phase of -pi/4" + + with self.subTest(test_name=test_name): + self.assertTrue( + np.allclose(np.abs((bi0 - bi) / bi), 0, atol=atol), + msg=f"Failed: {test_name}", + ) + + for i in range(2): + bi[:, i] = ivtilde(i, z, 2).reshape(-1) + + if s == 1: + test_name = "2-term for z with phase of pi/4, order 0 and 1" + else: + test_name = "2-term for z with phase of -pi/4, order 0 and 1" + + with self.subTest(test_name=test_name): + self.assertTrue( + np.allclose(np.abs((bi0[:, :2] - bi) / bi), 0, atol=1e-3), + msg=f"Failed: {test_name}", + ) + + bi0 = sqrt(1 / (2 * np.pi * z)) * (1 + 1 / (8 * z)) + bi1 = sqrt(1 / (2 * np.pi * z)) * (1 - 3 / (8 * z)) + + if s == 1: + test_name = "2-term vs. analytic for z with phase of pi/4" + else: + test_name = "2-term vs. analytic for z with phase of -pi/4" + + with self.subTest(test_name=test_name): + self.assertTrue( + np.allclose( + np.abs((np.hstack([bi0, bi1]) - bi) / bi), 0, atol=1e-15 + ), + msg=f"Failed: {test_name}", + ) + + +class TestXis(unittest.TestCase): + def test_xis_case1(self): + s = 1.0 + delta = 30.0 + mu = 20.0 + zo = delta**2 / mu + z = 15.0 + omega = 5.0 + coriolis_frequency = 6.0 + bld = 100.0 + expected_xiz = ( + 2 + * sqrt(2) + * _rot(s * np.pi / 4) + * np.divide(zo, delta) + * sqrt((1 + np.divide(z, zo)) * np.abs(1 + omega / coriolis_frequency)) + ) + expected_xih = ( + 2 + * sqrt(2) + * _rot(s * np.pi / 4) + * np.divide(zo, delta) + * sqrt((1 + np.divide(bld, zo)) * np.abs(1 + omega / coriolis_frequency)) + ) + expected_xi0 = ( + 2 + * sqrt(2) + * _rot(s * np.pi / 4) + * np.divide(zo, delta) + * sqrt(np.abs(1 + omega / coriolis_frequency)) + ) + assert np.allclose( + _xis(s, zo, delta, z, omega, coriolis_frequency, bld), + (expected_xiz, expected_xih, expected_xi0), + ) + + def test_xis_case2(self): + s = -1.0 + delta = 30.0 + mu = 20.0 + zo = delta**2 / mu + z = np.array([1.0, 2.0, 3.0]) + omega = np.array([0.1, 0.2, 0.3]) + coriolis_frequency = 0.5 + bld = 7.0 + expected_xiz = ( + 2 + * sqrt(2) + * _rot(s * np.pi / 4) + * np.divide(zo, delta) + * sqrt((1 + np.divide(z, zo)) * np.abs(1 + omega / coriolis_frequency)) + ) + expected_xih = ( + 2 + * sqrt(2) + * _rot(s * np.pi / 4) + * np.divide(zo, delta) + * sqrt((1 + np.divide(bld, zo)) * np.abs(1 + omega / coriolis_frequency)) + ) + expected_xi0 = ( + 2 + * sqrt(2) + * _rot(s * np.pi / 4) + * np.divide(zo, delta) + * sqrt(np.abs(1 + omega / coriolis_frequency)) + ) + assert np.allclose( + _xis(s, zo, delta, z, omega, coriolis_frequency, bld), + (expected_xiz, expected_xih, expected_xi0), + ) + + def test_xis_case3(self): + s = 1 + delta = 1.0 + mu = 1.0 + zo = delta**2 / mu + z = 0.0 + omega = 0.0 + coriolis_frequency = 0.1 + bld = 10.0 + expected_xiz = ( + 2 + * sqrt(2) + * _rot(s * np.pi / 4) + * np.divide(zo, delta) + * sqrt((1 + np.divide(z, zo)) * np.abs(1 + omega / coriolis_frequency)) + ) + expected_xih = ( + 2 + * sqrt(2) + * _rot(s * np.pi / 4) + * np.divide(zo, delta) + * sqrt((1 + np.divide(bld, zo)) * np.abs(1 + omega / coriolis_frequency)) + ) + expected_xi0 = ( + 2 + * sqrt(2) + * _rot(s * np.pi / 4) + * np.divide(zo, delta) + * sqrt(np.abs(1 + omega / coriolis_frequency)) + ) + assert np.allclose( + _xis(s, zo, delta, z, omega, coriolis_frequency, bld), + (expected_xiz, expected_xih, expected_xi0), + ) + + +class TestRot(unittest.TestCase): + def test_rot_positive(self): + # Test with positive angle + angle = np.pi / 4 + expected_result = np.exp(1j * angle) + self.assertTrue(np.isclose(_rot(angle), expected_result)) + + def test_rot_negative(self): + # Test with negative angle + angle = -np.pi / 3 + expected_result = np.exp(1j * angle) + self.assertTrue(np.isclose(_rot(angle), expected_result)) + + def test_rot_zero(self): + # Test with zero angle + angle = 0 + expected_result = np.exp(1j * angle) + self.assertTrue(np.isclose(_rot(angle), expected_result)) + + def test_rot_random(self): + # Test with random angle + angle = np.random.uniform(-3 * np.pi, np.pi) + expected_result = np.exp(1j * angle) + self.assertTrue(np.isclose(_rot(angle), expected_result)) + + def test_rot_empty(self): + # Test with empty array + angles = np.array([]) + expected_results = np.array([]) + assert np.allclose(_rot(angles), expected_results)