{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "#w=np.array([-.1,-.1,-.1,-.1,-.1, 0])\n",
    "#w=np.array([-1,-1,-1,-1,-1, 0])\n",
    "#c=np.array([  1,  1,  1,.01,100, 1e10])*1e-6\n",
    "e1e2t=214298.32535430687\n",
    "e3t=np.array([25.70331479, 26.28684983, 26.59728865, 26.75965336, \n",
    "              26.84381704, 26.88724213, 26.90959407, 26.92108493,\n",
    "              26.9269885,  26.93002054, 26.93157752, 26.93237697,])/20\n",
    "c=np.array([0.31565186, 0.30992895, 0.3301549,  0.423929,\n",
    "            0.54217166, 0.01, 0.01, 0.01,\n",
    "            0.01 , 0.0099, 100,   0.])\n",
    "w=-1*np.ones(np.shape(c))*2.8e-4\n",
    "w[-1]=0\n",
    "tmask=np.ones(np.shape(c))\n",
    "tmask[-1]=0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 0.00000000e+00, -8.83825208e-05, -8.67801060e-05, -9.24433720e-05,\n",
       "       -1.18700120e-04, -1.51808065e-04, -2.80000000e-06, -2.80000000e-06,\n",
       "       -2.80000000e-06, -2.80000000e-06, -2.77200000e-06,  0.00000000e+00])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#1. bottom value zero\n",
    "zwz=np.zeros(np.shape(c))\n",
    "#2. upstream advection\n",
    "zwz[1:]=w[1:]*c[:-1]\n",
    "zwz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "dt=40*2\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([3.15651834e-01, 3.09928950e-01, 3.30154898e-01, 4.23928993e-01,\n",
       "       5.42171651e-01, 1.00000414e-02, 1.00000000e-02, 1.00000000e-02,\n",
       "       1.00000000e-02, 9.90000001e-03, 1.00000000e+02, 0.00000000e+00])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ztra=np.zeros(np.shape(c))\n",
    "ztra[:-1]=-1*(zwz[:-1]-zwz[1:])/e1e2t\n",
    "delc_upwind=ztra.copy() # ** added to tra at this step\n",
    "zwi=(e3t*c+dt*ztra)/e3t*tmask # set fse3t's to 1\n",
    "zwi # upwind increases low point from .01 to .0298 in one time step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 3. antidiffusive flux\n",
    "zwz_sav=zwz.copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "jnzts=20 # same as SSC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "cb=c.copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "ca=np.zeros(np.shape(c))\n",
    "zwzts=np.zeros(np.shape(c))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\n"
     ]
    }
   ],
   "source": [
    "for jl in range(0,jnzts):\n",
    "    if jl==0:\n",
    "        jtaken=np.mod(jnzts+1,2)\n",
    "        print(jtaken)\n",
    "        zts=dt/jnzts # z_rzts=1/jnzts\n",
    "        cn=c.copy() # this makes first step euler\n",
    "    elif jl==1:\n",
    "        zts=2*dt/jnzts\n",
    "    zwz[1:]=.5*w[1:]*(cn[1:]+cn[:-1])\n",
    "    if jtaken==0:\n",
    "        zwzts[1:]=zwzts[1:]+zwz[1:]*zts\n",
    "    jtaken=np.mod(jtaken+1,2) # switch on/off\n",
    "    \n",
    "    zbtr=1/(e1e2t*e3t)\n",
    "    ztra[1:-1]=-1*zbtr[1:-1]*(zwz[1:-1]-zwz[2:])\n",
    "    ca[1:-1]=cb[1:-1]+zts*ztra[1:-1] # step forward\n",
    "    # swap for next step:\n",
    "    cb=cn.copy()\n",
    "    cn=ca.copy()\n",
    "    ca=np.zeros(np.shape(c))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.00000000e+00, 3.09928937e-01, 3.30154896e-01, 4.23928992e-01,\n",
       "       5.42171676e-01, 1.00000207e-02, 1.00000000e-02, 1.00000000e-02,\n",
       "       1.00000000e-02, 9.89611895e-03, 1.00000004e+02, 0.00000000e+00])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 0.00000000e+00, -3.47120417e-03, -7.16893903e-03, -8.44573961e-03,\n",
       "       -1.08203274e-02, -6.18432280e-03, -2.24000116e-04, -2.24000000e-04,\n",
       "       -2.24000000e-04, -2.22858266e-04, -1.12011088e+00,  0.00000000e+00])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "zwzts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# anti-diffusive vertical flux using average flux from sub-timestepping\n",
    "zwz=zwzts/dt-zwz_sav"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 0.00000000e+00,  4.49924686e-05, -2.83163186e-06, -1.31283731e-05,\n",
       "       -1.65539729e-05,  7.45040298e-05, -1.44821411e-12, -1.86347248e-20,\n",
       "       -2.75210745e-16,  1.42716732e-08, -1.39986140e-02,  0.00000000e+00])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "zwz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "zbig=1e40\n",
    "zrtrn=1e-15"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 3.15651860e-01,  3.09928950e-01,  3.30154900e-01,  4.23929000e-01,\n",
       "        5.42171660e-01,  1.00000414e-02,  1.00000000e-02,  1.00000000e-02,\n",
       "        1.00000000e-02,  9.90000001e-03,  1.00000000e+02, -1.00000000e+40])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 4. nonosc(pbef=ptb,zwx,zwy,pcc=zwz,paft=zwi,p2dt)\n",
    "# zwi is after value based on upwind\n",
    "pbef=c.copy()\n",
    "paft=zwi.copy()\n",
    "# search local extrema\n",
    "zbup=np.maximum(pbef*tmask-zbig*(1-tmask),paft*tmask-zbig*(1-tmask))\n",
    "zbup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([3.15651834e-01, 3.09928950e-01, 3.30154898e-01, 4.23928993e-01,\n",
       "       5.42171651e-01, 1.00000000e-02, 1.00000000e-02, 1.00000000e-02,\n",
       "       1.00000000e-02, 9.90000000e-03, 1.00000000e+02, 1.00000000e+40])"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "zbdo=np.minimum(pbef*tmask+zbig*(1-tmask),paft*tmask+zbig*(1-tmask))\n",
    "zbdo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([3.15651860e-01, 3.30154900e-01, 4.23929000e-01, 5.42171660e-01,\n",
       "       5.42171660e-01, 5.42171660e-01, 1.00000414e-02, 1.00000000e-02,\n",
       "       1.00000000e-02, 1.00000000e+02, 1.00000000e+02, 0.00000000e+00])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "zup=np.zeros(np.shape(c))\n",
    "zdo=zup.copy()\n",
    "zup[0]=np.maximum(zbup[0],zbup[1])\n",
    "zup[1:-1]=np.maximum(np.maximum(zbup[:-2],zbup[1:-1]),zbup[2:])\n",
    "#zup[-1]=np.maximum(zbup[-2],zbup[-1]) this is never calculated\n",
    "zup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.30992895, 0.30992895, 0.30992895, 0.3301549 , 0.01      ,\n",
       "       0.01      , 0.01      , 0.01      , 0.0099    , 0.0099    ,\n",
       "       0.0099    , 0.        ])"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "zdo[0]=np.minimum(zbdo[0],zbdo[1])\n",
    "zdo[1:-1]=np.minimum(np.minimum(zbdo[:-2],zbdo[1:-1]),zbdo[2:])\n",
    "#zdo[-1]=np.minimum(zbdo[-2],zbdo[-1])\n",
    "zdo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "zeros=np.zeros(np.shape(c))\n",
    "pcc=zwz.copy()\n",
    "zpos=np.zeros(np.shape(c))\n",
    "zneg=zpos.copy()\n",
    "zpos[:-1]=np.maximum(zeros[1:],pcc[1:])-np.minimum(zeros[:-1],pcc[:-1])\n",
    "zneg[:-1]=np.maximum(zeros[:-1],pcc[:-1])-np.minimum(zeros[1:],pcc[1:])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([4.49924686e-05, 0.00000000e+00, 2.83163186e-06, 1.31283731e-05,\n",
       "       9.10580028e-05, 0.00000000e+00, 1.44821411e-12, 1.86347248e-20,\n",
       "       1.42716734e-08, 0.00000000e+00, 1.39986140e-02, 0.00000000e+00])"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "zpos"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.00000000e+00, 4.78241005e-05, 1.31283731e-05, 1.65539729e-05,\n",
       "       0.00000000e+00, 7.45040313e-05, 1.86347248e-20, 2.75210745e-16,\n",
       "       0.00000000e+00, 1.39986283e-02, 0.00000000e+00, 0.00000000e+00])"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "zneg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "#! up & down beta terms\n",
    "zbt = e1e2t*e3t / dt\n",
    "zbetup = ( zup - paft) / ( zpos + zrtrn ) * zbt\n",
    "zbetdo = ( paft - zdo ) / ( zneg + zrtrn ) * zbt\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1.96438478e+00, 7.12108645e+16, 1.17972908e+08, 3.22806733e+07,\n",
       "       3.63591818e-01, 1.91644677e+18, 1.02905386e+08, 0.00000000e+00,\n",
       "       0.00000000e+00, 3.60655436e+20, 0.00000000e+00, 0.00000000e+00])"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "zbetup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1.97016635e+16, 3.35064296e-02, 5.48825328e+06, 2.03029860e+07,\n",
       "       1.91335167e+18, 2.00000003e+00, 0.00000000e+00, 0.00000000e+00,\n",
       "       3.60650534e+14, 2.00019628e-06, 3.60676287e+20, 0.00000000e+00])"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "zbetdo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "#! monotonic flux in the k direction, i.e. pcc\n",
    "ones=np.ones(np.shape(c))\n",
    "za=zeros.copy()\n",
    "zb=zeros.copy()\n",
    "zc=zeros.copy()\n",
    "za[:-1] = np.minimum(np.minimum( ones[1:], zbetdo[1:]), zbetup[:-1] )\n",
    "zb[:-1] = np.minimum(np.minimum( ones[1:], zbetup[1:]), zbetdo[:-1] )\n",
    "zc[:-1]=np.where(pcc[1:]<0,0,1)\n",
    "#zc[:-1]=.5+np.abs(pcc[1:])/pcc[1:]*.5\n",
    "#zc =       ( 0.5  + SIGN( 0.5 , pcc(ji,jj,jk+1) ) )\n",
    "# SIGN(A,B) returns value of A with sign of B\n",
    "pcc[1:]=pcc[1:]*(zc[:-1]*za[:-1]+(1-zc[:-1])*zb[:-1])\n",
    "#pcc(ji,jj,jk+1) = pcc(ji,jj,jk+1) * ( zc * za + ( 1._wp - zc) * zb )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 0.00000000e+00,  1.50753698e-06, -9.48778735e-08, -1.31283731e-05,\n",
       "       -6.01888912e-06,  2.70890557e-05, -1.44821411e-12, -0.00000000e+00,\n",
       "       -0.00000000e+00,  0.00000000e+00, -0.00000000e+00,  0.00000000e+00])"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pcc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 0.00000000e+00,  4.49924686e-05, -2.83163186e-06, -1.31283731e-05,\n",
       "       -1.65539729e-05,  7.45040298e-05, -1.44821411e-12, -1.86347248e-20,\n",
       "       -2.75210745e-16,  1.42716732e-08, -1.39986140e-02,  0.00000000e+00])"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# before nonosc, zwz contains difference between centered substepped leapfrog and upwind:\n",
    "zwz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "# after nonosc, zwz (pcc) is adjusted so that only one element changes from upwind "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 5. final trend with corrected fluxes\n",
    "ztra=0*ztra\n",
    "ztra2=ztra.copy()\n",
    "ztra[:-1]=-1*zbtr[:-1]*(pcc[:-1]-pcc[1:])\n",
    "ztra2[:-1]=-1*zbtr[:-1]*(zwz[:-1]-zwz[1:])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 5.47381429e-12, -5.68915331e-12, -4.57335326e-11,  2.47952685e-11,\n",
       "        1.15106320e-10, -9.40283574e-11,  5.02269547e-18, -0.00000000e+00,\n",
       "        0.00000000e+00, -0.00000000e+00,  0.00000000e+00,  0.00000000e+00])"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# total change is ztra+delc, ztra is after nonosc adjustment\n",
    "ztra"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 1.63366087e-10, -1.69792884e-10, -3.61304733e-11, -1.19472339e-11,\n",
       "        3.16581162e-10, -2.58609654e-10,  5.02269541e-18, -9.54013762e-22,\n",
       "        4.94650364e-14, -4.85132104e-08,  4.85103562e-08,  0.00000000e+00])"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ztra2 # this is before nonosc version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([-4.12427492e-10,  7.47749567e-12, -2.64270194e-11, -1.22524280e-10,\n",
       "       -1.54494650e-10,  6.95330048e-10, -0.00000000e+00, -0.00000000e+00,\n",
       "       -0.00000000e+00,  1.30658977e-13,  1.29352387e-11,  0.00000000e+00])"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "delc_upwind"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([-4.06953677e-10,  1.78834236e-12, -7.21605520e-11, -9.77290114e-11,\n",
       "       -3.93883298e-11,  6.01301690e-10,  5.02269547e-18, -0.00000000e+00,\n",
       "        0.00000000e+00,  1.30658977e-13,  1.29352387e-11,  0.00000000e+00])"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "delc_upwind+ztra"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([-2.49061405e-10, -1.62315388e-10, -6.25574927e-11, -1.34471514e-10,\n",
       "        1.62086512e-10,  4.36720394e-10,  5.02269541e-18, -9.54013762e-22,\n",
       "        4.94650364e-14, -4.85130797e-08,  4.85232915e-08,  0.00000000e+00])"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "delc_upwind+ztra2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([3.1565186e-01, 3.0992895e-01, 3.3015490e-01, 4.2392900e-01,\n",
       "       5.4217166e-01, 1.0000000e-02, 1.0000000e-02, 1.0000000e-02,\n",
       "       1.0000000e-02, 9.9000000e-03, 1.0000000e+02, 0.0000000e+00])"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "# nonosc and diatoms???? maybe code shouldn't have to be monotonic"
   ]
  }
 ],
 "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}