Jake_functions.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.sparse import coo_matrix\n",
    "from scipy.sparse.linalg import eigs\n",
    "from scipy.integrate import cumtrapz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# solves G''(z) + (N^2(z) - omega^2)G(z)/c^2 = 0 \n",
    "#   subject to G'(0) = gG(0)/c^2 (free surface) & G(-D) = 0 (flat bottom)\n",
    "# G(z) is normalized so that the vertical integral of (G'(z))^2 is D\n",
    "# G' is dimensionless, G has dimensions of length\n",
    "\n",
    "# - N is buoyancy frequency [s^-1] (nX1 vector)\n",
    "# - depth [m] (maximum depth is considered the sea floor) (nX1 vector)\n",
    "# - omega is frequency [s^-1] (scalar)\n",
    "# - mmax is the highest baroclinic mode calculated\n",
    "# - m=0 is the barotropic mode\n",
    "# - 0 < m <= mmax are the baroclinic modes\n",
    "# - Modes are calculated by expressing in finite difference form 1) the\n",
    "#  governing equation for interior depths (rows 2 through n-1) and 2) the\n",
    "#  boundary conditions at the surface (1st row) and the bottome (last row).\n",
    "# - Solution is found by solving the eigenvalue system A*x = lambda*B*x\n",
    "def vertical_modes(N2_0, Depth, omega, mmax):\n",
    "    z = -1 * Depth\n",
    "\n",
    "    if np.size(np.shape(N2_0)) > 1:\n",
    "        N2 = np.nanmean(N2_0, axis=1)\n",
    "    else:\n",
    "        N2 = N2_0\n",
    "\n",
    "    n = np.size(z)\n",
    "    nm1 = n - 1\n",
    "    nm2 = n - 2\n",
    "    gravity = 9.82\n",
    "    # ----- vertical increments\n",
    "    dz = np.concatenate([[0], z[1:] - z[0:nm1]])  # depth increment [m]\n",
    "    dzm = np.concatenate([[0], 0.5 * (z[2:] - z[0:nm2]), [0]])  # depth increment between midpoints [m]\n",
    "    # ----- sparse matrices\n",
    "    # A = row pos, B = col pos, C = val  \n",
    "    A = np.concatenate([[0], [0], np.arange(1, nm1), np.arange(1, nm1), np.arange(1, nm1), [n - 1]])\n",
    "    B = np.concatenate([[0], [1], np.arange(1, nm1), np.arange(0, nm2), np.arange(2, n), [n - 1]])\n",
    "    C = np.concatenate(\n",
    "        [[-1 / dz[1]], [1 / dz[1]], (1 / dz[2:] + 1 / dz[1:nm1]) / dzm[1:nm1], -1 / (dz[1:nm1] * dzm[1:nm1]),\n",
    "         -1 / (dz[2:n] * dzm[1:nm1]), [-1]])\n",
    "        # [[-1 / dz[1]], [1 / dz[1]], (1 / dz[2:] + 1 / dz[1:nm1]) / dzm[1:nm1], -1 / (dz[1:nm1] * dzm[1:nm1]),\n",
    "        #  -1 / (dz[2:n] * dzm[1:nm1]), [-1]])\n",
    "    mat1 = coo_matrix((C, (A, B)), shape=(n, n))\n",
    "\n",
    "    D = np.concatenate([[0], np.arange(1, n)])\n",
    "    E = np.concatenate([[0], np.arange(1, n)])\n",
    "    F = np.concatenate([[gravity], N2[1:] - omega * omega])  # originially says N2[1:,10]\n",
    "    mat2 = coo_matrix((F, (D, E)), shape=(n, n))\n",
    "\n",
    "    # compute eigenvalues and vectors \n",
    "    vals, vecs = eigs(mat1, k=mmax + 1, M=mat2, sigma=0)\n",
    "    eigenvalue = np.real(vals)\n",
    "    wmodes = np.real(vecs)\n",
    "    s_ind = np.argsort(eigenvalue)\n",
    "    eigenvalue = eigenvalue[s_ind]\n",
    "    wmodes = wmodes[:, s_ind]\n",
    "    m = np.size(eigenvalue)\n",
    "    c = 1 / np.sqrt(eigenvalue)  # kelvin wave speed\n",
    "    # normalize mode (shapes)\n",
    "    Gz = np.zeros(np.shape(wmodes))\n",
    "    G = np.zeros(np.shape(wmodes))\n",
    "    for i in range(m):\n",
    "        dw_dz = np.nan * np.ones(np.shape(z))\n",
    "        dw_dz[0] = (wmodes[1, i] - wmodes[0, i]) / (z[1] - z[0])\n",
    "        dw_dz[-1] = (wmodes[-1, i] - wmodes[-2, i]) / (z[-1] - z[-2])\n",
    "        for j in range(1, len(z) - 1):\n",
    "            dw_dz[j] = (wmodes[j + 1, i] - wmodes[j - 1, i]) / (z[j + 1] - z[j - 1])\n",
    "        # dw_dz = np.gradient(wmodes[:, i], z)\n",
    "        norm_constant = np.sqrt(np.trapz((dw_dz * dw_dz), (-1 * z)) / (-1 * z[-1]))\n",
    "        # norm_constant = np.abs(np.trapz(dw_dz * dw_dz, z) / Depth.max())\n",
    "\n",
    "        if dw_dz[0] < 0:\n",
    "            norm_constant = -1 * norm_constant\n",
    "        Gz[:, i] = dw_dz / norm_constant\n",
    "\n",
    "        norm_constant_G = np.sqrt(np.trapz((wmodes[:, i] * wmodes[:, i]), (-1 * z)) / (-1 * z[-1]))\n",
    "        G[:, i] = wmodes[:, i] / norm_constant\n",
    "\n",
    "    epsilon = np.nan * np.zeros((5, 5, 5))  # barotropic and first 5 baroclinic\n",
    "    for i in range(0, 5):  # i modes\n",
    "        for j in range(0, 5):  # j modes\n",
    "            for m in range(0, 5):  # k modes\n",
    "                epsilon[i, j, m] = np.trapz((Gz[:, i] * Gz[:, j] * Gz[:, m]), -1.0*z) / (-1.0*z[-1])\n",
    "\n",
    "    return G, Gz, c, epsilon\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def vertical_modes_f(N2_0, depth, omega, mmax, bc_bot, ref_lat, slope):\n",
    "    z = -1 * depth\n",
    "    # ensure that N2 is 1-D array (take average if not)\n",
    "    if np.size(np.shape(N2_0)) > 1:\n",
    "        N2 = np.nanmean(N2_0, axis=1)\n",
    "    else:\n",
    "        N2 = N2_0  # buoyancy frequency squared\n",
    "\n",
    "    # make first element = 0 if not already so\n",
    "    if depth[0] != 0:\n",
    "        z = np.nan * np.zeros(len(depth))\n",
    "        z[0] = 0\n",
    "        z[1:] = -depth[0:-1]\n",
    "        N2 = np.concatenate([[0], N2[0:-1]])\n",
    "\n",
    "    n = np.size(z)\n",
    "    nm1 = n - 1\n",
    "    nm2 = n - 2\n",
    "    gravity = 9.82\n",
    "    # ---- fix N2 profiles with zeros\n",
    "    fixer = np.where(N2 == 0)[0]\n",
    "    fixer2 = fixer[(fixer > 10) & (fixer < len(depth) - 20)]\n",
    "    if len(fixer2) >= 1:\n",
    "        for i in range(len(fixer2)):\n",
    "            N2[fixer2[i]] = N2[fixer2[i] - 1] + (N2[fixer2[i] + 1] - N2[fixer2[i] - 1]) / 2\n",
    "    n2_first = np.where(N2 > 0)[0][0]\n",
    "    n2_last = np.where(N2 > 0)[0][-1]\n",
    "    N2[n2_last + 1:] = N2[n2_last]\n",
    "    N2[1:n2_first] = N2[n2_first]\n",
    "    # all of this removes N2 = 0 at the end of the array and leaves N2 = 0 only for the first element of the array\n",
    "\n",
    "    N2_inter = N2.copy()\n",
    "    N2_inter[0] = N2_inter[1]\n",
    "\n",
    "    # sloping bottom\n",
    "    om = 7.2921 * 10**(-5.)  # rotation rate of earth\n",
    "    f_ref = 2. * om * np.sin(np.deg2rad(ref_lat))\n",
    "    r_earth = 6371e3  # radius of earth\n",
    "    beta_ref_1 = (2.0 * om / r_earth) * np.cos(np.deg2rad(ref_lat))\n",
    "    alpha = slope\n",
    "    # k_wave_2 = (1 / 50000) ^ 2\n",
    "    Lacasce_bc = ((alpha * N2) / (beta_ref_1 * f_ref)) * 1.0\n",
    "\n",
    "    # 1st approximate (k(x)u'(x)) at halfway points\n",
    "    z12 = (z[1:] - z[0:-1]) / 2 + z[0:-1]\n",
    "    k = (f_ref**2) / N2_inter  # old 1 / N2_inter\n",
    "    k12 = np.flipud(np.interp(np.flipud(z12), np.flipud(z), np.flipud(k)))\n",
    "\n",
    "    # ----- vertical increments\n",
    "    dz = np.concatenate([[0], z[1:] - z[0:nm1]])  # depth increment [m]\n",
    "    dzm = np.concatenate([[0], 0.5 * (z[2:] - z[0:nm2]), [0]])  # depth increment between midpoints [m]\n",
    "    # ----- sparse matrices\n",
    "    # A = row pos, B = col pos, C = val\n",
    "    if bc_bot > 1:\n",
    "        A = np.concatenate([[0], [0], np.arange(1, nm1), np.arange(1, nm1), np.arange(1, nm1), [n - 1]])\n",
    "        B = np.concatenate([[0], [1], np.arange(1, nm1), np.arange(0, nm2), np.arange(2, n), [n - 1]])\n",
    "        # rough_bottom\n",
    "        C = np.concatenate(\n",
    "            [[-1 / dz[1]], [1 / dz[1]], (1 / dz[1:nm1] / dz[1:nm1]) * (k12[0:-1] + k12[1:]),\n",
    "             -1 * k12[0:-1] / (dz[1:nm1] * dz[1:nm1]), -1 * k12[1:] / (dz[2:n] * dz[1:nm1]), [-1]])\n",
    "    else:\n",
    "        A = np.concatenate([[0], [0], np.arange(1, nm1), np.arange(1, nm1), np.arange(1, nm1), [n - 1], [n - 1]])\n",
    "        B = np.concatenate([[0], [1], np.arange(1, nm1), np.arange(0, nm2), np.arange(2, n), [n - 2], [n - 1]])\n",
    "        # flat bottom\n",
    "        C = np.concatenate(\n",
    "            [[-1 / dz[1]], [1 / dz[1]], (1 / dz[1:nm1] / dz[1:nm1]) * (k12[0:-1] + k12[1:]),\n",
    "             -1 * k12[0:-1] / (dz[1:nm1] * dzm[1:nm1]), -1 * k12[1:] / (dz[2:n] * dzm[1:nm1]), [1 / dz[nm1]],\n",
    "             [-1 / dz[nm1]]])\n",
    "\n",
    "    mat1 = coo_matrix((C, (A, B)), shape=(n, n))\n",
    "\n",
    "    D = np.concatenate([[0], np.arange(1, nm1), [n - 1]])\n",
    "    E = np.concatenate([[0], np.arange(1, nm1), [n - 1]])\n",
    "    F = np.concatenate([[N2[0] / gravity], np.ones(len(np.arange(1, nm1))), [Lacasce_bc[-1]]])  # last term was 0\n",
    "    mat2 = coo_matrix((F, (D, E)), shape=(n, n))\n",
    "\n",
    "    # compute eigenvalues and vectors\n",
    "    vals, vecs = eigs(mat1, k=mmax + 1, M=mat2, sigma=0)\n",
    "    eigenvalue = np.real(vals)\n",
    "    wmodes = np.real(vecs)\n",
    "    s_ind = np.argsort(eigenvalue)\n",
    "    eigenvalue = eigenvalue[s_ind]\n",
    "    wmodes = wmodes[:, s_ind]\n",
    "    m = np.size(eigenvalue)\n",
    "    c = 1 / np.sqrt(eigenvalue)  # kelvin wave speed\n",
    "    # normalize mode (shapes)\n",
    "    F = np.zeros(np.shape(wmodes))\n",
    "    G = np.zeros(np.shape(wmodes))\n",
    "    for i in range(m):\n",
    "        norm_constant = np.sqrt(np.trapz(wmodes[:, i] * wmodes[:, i], z) / z[-1])\n",
    "        if wmodes[0, i] < 0:\n",
    "            norm_constant = -1 * norm_constant\n",
    "        F[:, i] = wmodes[:, i] / norm_constant\n",
    "        G[:, i] = cumtrapz(wmodes[:, i], x=z, initial=0) / norm_constant\n",
    "        # G[:, i] = (1 / N2_inter) * np.gradient(wmodes[:, i], z) / norm_constant\n",
    "        # G[:, i] = cumtrapz(wmodes[:, i], z, initial=0) / norm_constant\n",
    "\n",
    "    epsilon = np.nan * np.zeros((2, 3, 3))\n",
    "    for i in range(0, 2):  # i modes\n",
    "        for j in range(0, 3):  # j modes\n",
    "            for m in range(0, 3):  # k modes\n",
    "                epsilon[i, j, m] = np.trapz((F[:, i] * F[:, j] * F[:, m]), z) / (z[-1])\n",
    "\n",
    "    return G, F, c, norm_constant, epsilon\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def PE_Tide_GM(rho0, Depth, nmodes, N2, f_ref):\n",
    "    modenum = np.arange(0, nmodes)\n",
    "    Navg = np.trapz(np.nanmean(np.sqrt(N2), 1), Depth) / Depth[-1]  # mean N (where input is multiple profiles)\n",
    "    # exNavg = np.trapz(np.sqrt(N2), Depth) / Depth[-1]  # mean N (where input is a single profile)\n",
    "\n",
    "    TE_SD = (75 + 280 + 72) / (rho0 * Depth[-1])  # SD tidal energy [m^2/s^2] Hendry 1977\n",
    "    sigma_SD = 2.0 * np.pi / (12 * 3600)  # SD frequency [s^-1]\n",
    "    PE_SD = TE_SD * (sigma_SD ** 2 - f_ref ** 2) / (2.0 * sigma_SD ** 2)\n",
    "\n",
    "    # e-folding scale of N\n",
    "    def n_exp(p, n_e, z):\n",
    "        a = p[0]\n",
    "        b = p[1]\n",
    "        fsq = (n_e - a * np.exp((z / b))) ** 2\n",
    "        return fsq.sum()\n",
    "\n",
    "    from scipy.optimize import fmin\n",
    "    N2_f = np.nanmean(np.sqrt(N2), 1)\n",
    "    ins = np.transpose(np.concatenate([N2_f[:, None], -1.0 * Depth[:, None]], axis=1))\n",
    "    p = [0.001, 500.0]\n",
    "    coeffs = fmin(n_exp, p, args=(tuple(ins)), disp=0)\n",
    "\n",
    "    if np.abs(coeffs[1]) < 100:\n",
    "        bGM = 200.0  # GM e folding scale of N [m]\n",
    "    else:\n",
    "        bGM = np.abs(coeffs[1])  # GM e folding scale of N [m]\n",
    "    N0_GM = 0.0052  # GM N scale [s^-1];\n",
    "    jstar = 6.0  # GM vertical mode number scale\n",
    "    EGM = 0.000063  # GM energy level [no dimensions]\n",
    "    HHterm = 1.0 / (modenum[1:] * modenum[1:] + jstar * jstar)\n",
    "    HH = HHterm / np.sum(HHterm)\n",
    "\n",
    "    # compute energy over a range of frequencies and then integrate\n",
    "    omega = np.arange(np.round(f_ref, 5), np.round(Navg, 5), 0.00001)\n",
    "    if omega[0] < f_ref:\n",
    "        omega[0] = f_ref.copy() + 0.000001\n",
    "    BBterm = (2.0 / 3.14159) * (f_ref / omega) * (((omega ** 2) - (f_ref ** 2))**(-0.5))\n",
    "    BBint = np.trapz(BBterm, omega)\n",
    "    BBterm2 = (1.0 / BBint) * BBterm\n",
    "\n",
    "    EE = np.tile(BBterm2[:, None], (1, len(modenum) - 1)) * np.tile(HH[None, :], (len(omega), 1)) * EGM\n",
    "    omega_g = np.tile(omega[:, None], (1, len(modenum) - 1))\n",
    "    FPE = 0.5 * (Navg ** 2.0) * ((bGM ** 2.0) * N0_GM * (1.0 / Navg) * (omega_g ** 2 - f_ref ** 2) * (1 / (omega_g ** 2)) * EE)\n",
    "    FKE = 0.5 * bGM * bGM * N0_GM * Navg * ((omega_g ** 2) + (f_ref ** 2)) * (1.0 / (omega_g ** 2)) * EE\n",
    "\n",
    "    PE_GM_tot = bGM * bGM * N0_GM * Navg * EE # (bGM * bGM * N0_GM * Navg * HH * EGM/2  method charlie employed)\n",
    "\n",
    "    FPE_int = np.nan * np.ones(np.shape(FPE[0, :]))\n",
    "    FKE_int = np.nan * np.ones(np.shape(FPE[0, :]))\n",
    "    FPE_tot_int = np.nan * np.ones(np.shape(FPE[0, :]))\n",
    "    for i in range(len(modenum[1:])):\n",
    "        FPE_int[i] = np.trapz(FPE[:, i], omega)\n",
    "        FKE_int[i] = np.trapz(FKE[:, i], omega)\n",
    "        FPE_tot_int[i] = np.trapz(PE_GM_tot[:, i], omega)\n",
    "\n",
    "    return (PE_SD, FPE_tot_int, FPE_int, FKE_int, bGM)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def eta_fit(num_profs, grid, nmodes, n2, G, c, eta, eta_fit_dep_min, eta_fit_dep_max):\n",
    "    bvf = np.sqrt(n2)\n",
    "    Neta = np.nan * np.zeros(eta.shape)\n",
    "    NEta_m = np.nan * np.zeros(eta.shape)\n",
    "    Eta_m = np.nan * np.zeros(eta.shape)\n",
    "    AG = np.zeros([nmodes, num_profs])\n",
    "    PE_per_mass = np.nan * np.zeros([nmodes, num_profs])\n",
    "    for i in range(num_profs):\n",
    "        # fit to eta profiles\n",
    "        if len(np.shape(eta)) > 1:\n",
    "            this_eta = eta[:, i].copy()\n",
    "            eta_fs = eta[:, i].copy()  # ETA\n",
    "        else:\n",
    "            this_eta = eta.copy()\n",
    "            eta_fs = eta.copy()  # ETA\n",
    "        # obtain matrix of NEta\n",
    "        # old, Neta[:, i] = bvf * this_eta\n",
    "        # find indices within fitting range\n",
    "        iw = np.where((grid >= eta_fit_dep_min) & (grid <= eta_fit_dep_max))[0]\n",
    "        if iw.size > 1:\n",
    "            i_sh = np.where((grid < eta_fit_dep_min))\n",
    "            eta_fs[i_sh[0]] = grid[i_sh] * this_eta[iw[0]] / grid[iw[0]]\n",
    "\n",
    "            # i_dp = np.where((grid > eta_fit_dep_max))\n",
    "            # eta_fs[i_dp[0]] = (grid[i_dp] - grid[-1]) * this_eta[iw[0][-1]] / (grid[iw[0][-1]] - grid[-1])\n",
    "            # -- taper fit as z approaches -H\n",
    "            i_dp = np.where((grid > eta_fit_dep_max))[0]\n",
    "            lgs = grid[iw[-1]]\n",
    "            grid_ar = np.nan * np.ones(len(i_dp))\n",
    "            for oo in range(len(grid[i_dp])):\n",
    "                grid_ar[oo] = np.int(grid[i_dp[oo]])\n",
    "            eta_fs[i_dp] = (grid_ar - np.int(grid[-1])) * this_eta[iw[-1]] / (np.int(lgs) - grid[-1])\n",
    "\n",
    "            AG[1:, i] = np.squeeze(np.linalg.lstsq(G[:, 1:], eta_fs)[0])\n",
    "            # AG[1:, i] = np.squeeze(np.linalg.lstsq(G[:, 1:], np.transpose(np.atleast_2d(eta_fs)))[0])\n",
    "            Eta_m[:, i] = np.squeeze(np.matrix(G) * np.transpose(np.matrix(AG[:, i])))\n",
    "            NEta_m[:, i] = bvf * np.array(np.squeeze(np.matrix(G) * np.transpose(np.matrix(AG[:, i]))))\n",
    "            PE_per_mass[:, i] = 0.5 * AG[:, i] * AG[:, i] * c * c\n",
    "\n",
    "    return (AG, Eta_m, NEta_m, PE_per_mass)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "done\n"
     ]
    }
   ],
   "source": [
    "print('done')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}