{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Bickley jet\n",
    "\n",
    "The Bickley jet flow is a kinematic idealized model of a meandering zonal jet\n",
    "flanked above and below by counterrotating vortices. It was introduced by\n",
    "[Rypina et al.](https://doi.org/10.1175/JAS4036.1); cf. also [del‐Castillo‐Negrete and Morrison](https://doi.org/10.1063/1.858639).\n",
    "\n",
    "The Bickley jet is described by a time-dependent velocity field arising from a\n",
    "stream-function. The corresponding velocity field can be conveniently defined\n",
    "using the `@velo_from_stream` macro from [`StreamMacros.jl`](https://github.com/CoherentStructures/StreamMacros.jl):"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "using Distributed\n",
    "nprocs() == 1 && addprocs()\n",
    "\n",
    "@everywhere using CoherentStructures, StreamMacros\n",
    "const bickley = @velo_from_stream psi begin\n",
    "    psi  = psi₀ + psi₁\n",
    "    psi₀ = - U₀ * L₀ * tanh(y / L₀)\n",
    "    psi₁ =   U₀ * L₀ * sech(y / L₀)^2 * re_sum_term\n",
    "\n",
    "    re_sum_term = Σ₁ + Σ₂ + Σ₃\n",
    "\n",
    "    Σ₁ = ε₁ * cos(k₁*(x - c₁*t))\n",
    "    Σ₂ = ε₂ * cos(k₂*(x - c₂*t))\n",
    "    Σ₃ = ε₃ * cos(k₃*(x - c₃*t))\n",
    "\n",
    "    k₁ = 2/r₀    ; k₂ = 4/r₀   ; k₃ = 6/r₀\n",
    "    ε₁ = 0.0075  ; ε₂ = 0.15   ; ε₃ = 0.3\n",
    "    c₂ = 0.205U₀ ; c₃ = 0.461U₀; c₁ = c₃ + (√5-1)*(c₂-c₃)\n",
    "    U₀ = 62.66e-6; L₀ = 1770e-3; r₀ = 6371e-3\n",
    "end"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Now, `bickley` is a callable function with the standard `OrdinaryDiffEq`\n",
    "signature `(u, p, t)` with state `u`, (unused) parameter `p` and time `t`.\n",
    "\n",
    "## Geodesic vortices\n",
    "\n",
    "Here we briefly demonstrate how to find material barriers to diffusive transport;\n",
    "see Geodesic elliptic material vortices for references and details."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "@everywhere using OrdinaryDiffEq, Tensors\n",
    "q = 81\n",
    "const tspan = range(0., stop=3456000., length=q)\n",
    "ny = 61\n",
    "nx = (22ny) ÷ 6\n",
    "xmin, xmax, ymin, ymax = 0.0 - 2.0, 6.371π + 2.0, -3.0, 3.0\n",
    "xspan = range(xmin, stop=xmax, length=nx)\n",
    "yspan = range(ymin, stop=ymax, length=ny)\n",
    "P = tuple.(xspan, yspan')\n",
    "const δ = 1.e-6\n",
    "const D = SymmetricTensor{2,2}((2., 0., 1/2))\n",
    "mCG_tensor = u -> av_weighted_CG_tensor(bickley, u, tspan, δ; D=(_ -> D), tolerance=1e-6, solver=Tsit5())\n",
    "\n",
    "C̅ = pmap(mCG_tensor, P; batch_size=ceil(Int, length(P)/nprocs()^2))\n",
    "p = LCSParameters(2.0)\n",
    "vortices, singularities = ellipticLCS(C̅, xspan, yspan, p)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "The result is visualized as follows:"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "using Plots\n",
    "trace = tensor_invariants(C̅)[5]\n",
    "fig = plot_vortices(vortices, singularities, (xmin, ymin), (xmax, ymax);\n",
    "    bg=trace, xspan=xspan, yspan=yspan, title=\"DBS field and transport barriers\", showlabel=true)\n",
    "Plots.plot(fig)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "## FEM-based Methods"
   ],
   "metadata": {}
  },
  {
   "cell_type": "markdown",
   "source": [
    "Assume we have setup the `bickley` function using the `@velo_from_stream` macro\n",
    "as described above. We are working on a periodic domain in one direction:"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "using Distances\n",
    "LL = (0.0, -3.0); UR = (6.371π, 3.0)\n",
    "ctx, _ = regularP2TriangularGrid((50, 15), LL, UR, quadrature_order=2)\n",
    "predicate = (p1, p2) -> peuclidean(p1, p2, [6.371π, Inf]) < 1e-10\n",
    "bdata = BoundaryData(ctx, predicate, []);"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "Using a FEM-based method to compute coherent structures:"
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "cgfun = x -> mean_diff_tensor(bickley, x, range(0.0, stop=40*3600*24, length=81), 1.e-8; tolerance=1e-5)\n",
    "\n",
    "K = assembleStiffnessMatrix(ctx, cgfun, bdata=bdata)\n",
    "M = assembleMassMatrix(ctx, bdata=bdata)\n",
    "λ, v = CoherentStructures.get_smallest_eigenpairs(K, M, 10)\n",
    "\n",
    "import Plots\n",
    "fig_spectrum = plot_real_spectrum(λ)\n",
    "\n",
    "Plots.plot(fig_spectrum)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "K-means clustering yields the coherent vortices."
   ],
   "metadata": {}
  },
  {
   "outputs": [],
   "cell_type": "code",
   "source": [
    "using Clustering\n",
    "ctx2, _ = regularTriangularGrid((200, 60), LL, UR)\n",
    "v_upsampled = sample_to(v, ctx, ctx2, bdata=bdata)\n",
    "\n",
    "function iterated_kmeans(numiterations, args...)\n",
    "    best = kmeans(args...)\n",
    "    for i in 1:(numiterations - 1)\n",
    "        cur = kmeans(args...)\n",
    "        if cur.totalcost < best.totalcost\n",
    "            best = cur\n",
    "        end\n",
    "    end\n",
    "    return best\n",
    "end\n",
    "\n",
    "n_partition = 8\n",
    "res = iterated_kmeans(20, permutedims(v_upsampled[:,2:n_partition]), n_partition)\n",
    "u = kmeansresult2LCS(res)\n",
    "u_combined = sum([u[:,i] * i for i in 1:n_partition])\n",
    "fig = plot_u(ctx2, u_combined, 400, 400;\n",
    "    color=:rainbow, colorbar=:none, title=\"$n_partition-partition of Bickley jet\")\n",
    "\n",
    "Plots.plot(fig)"
   ],
   "metadata": {},
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "source": [
    "---\n",
    "\n",
    "*This notebook was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).*"
   ],
   "metadata": {}
  }
 ],
 "nbformat_minor": 3,
 "metadata": {
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.8.5"
  },
  "kernelspec": {
   "name": "julia-1.8",
   "display_name": "Julia 1.8.5",
   "language": "julia"
  }
 },
 "nbformat": 4
}