{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "King / Blankenbach Benchmark Case 1\n",
    "======\n",
    "\n",
    "Isoviscous thermal convection using EBA formulation.\n",
    "----\n",
    "\n",
    "Two-dimensional, incompressible, bottom heated, steady isoviscous thermal convection in a 1 x 1 box, see case 1 of King *et al.* 2009 / Blankenbach *et al.* 1989 for details.\n",
    "\n",
    "\n",
    "**This example introduces:**\n",
    "1. Extended Boussinesq Approximation, EBA, formulation for Stokes Flow.\n",
    "\n",
    "**Keywords:** Stokes system, EBA, advective diffusive systems, analysis tools\n",
    "\n",
    "**References**\n",
    "\n",
    "Scott D. King, Changyeol Lee, Peter E. Van Keken, Wei Leng, Shijie Zhong, Eh Tan, Nicola Tosi, Masanori C. Kameyama, A community benchmark for 2-D Cartesian compressible convection in the Earth's mantle, Geophysical Journal International, Volume 180, Issue 1, January 2010, Pages 73–87, https://doi.org/10.1111/j.1365-246X.2009.04413.x\n",
    "\n",
    "B. Blankenbach, F. Busse, U. Christensen, L. Cserepes, D. Gunkel, U. Hansen, H. Harder, G. Jarvis, M. Koch, G. Marquart, D. Moore, P. Olson, H. Schmeling and T. Schnaubelt. A benchmark comparison for mantle convection codes. Geophysical Journal International, 98, 1, 23–38, 1989\n",
    "http://onlinelibrary.wiley.com/doi/10.1111/j.1365-246X.1989.tb05511.x/abstract\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import underworld as uw\n",
    "from underworld import function as fn\n",
    "import underworld.visualisation as vis\n",
    "import math\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from underworld.scaling import units as u"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "10000.000000000002 dimensionless\n"
     ]
    }
   ],
   "source": [
    "# The physical S.I. units from the Blankenbach paper\n",
    "# Sanity check for the Rayleigh number.\n",
    "# In this implementation the equations are non-dimensionalised with Ra\n",
    "\n",
    "# Ra = a*g*dT*h**3 / (eta0*dif)\n",
    "h  = 1e6 * u.m\n",
    "dT = 1e3 * u.degK\n",
    "a  = 2.5e-5 * u.degK**-1\n",
    "g  = 10  * u.m * u.s**-2\n",
    "diff = 1e-6 * u.m**2 * u.s**-1\n",
    "eta = 1e23 * u.kg * u.s**-1 * u.m**-1\n",
    "rho = 4000 * u.kg * u.m**-3 # reference density, only for units\n",
    "\n",
    "Ra = (a*g*dT*h**3)/(eta/rho*diff)\n",
    "print(Ra.to_compact())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Setup parameters\n",
    "-----"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "boxHeight = 1.0\n",
    "boxLength = 1.0\n",
    "# Set grid resolution.\n",
    "res = 64\n",
    "# Set max & min temperautres\n",
    "tempMin = 0.0\n",
    "tempMax = 1.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Choose which Rayleigh number, see case 1 of Blankenbach *et al.* 1989 for details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "Di = 0.5\n",
    "Ra = 1.e4\n",
    "eta0 = 1.e23"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set input and output file directory "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "outputPath = 'EBA/'\n",
    "# Make output directory if necessary.\n",
    "if uw.mpi.rank==0:\n",
    "    import os\n",
    "    if not os.path.exists(outputPath):\n",
    "        os.makedirs(outputPath)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create mesh and variables\n",
    "------"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "mesh = uw.mesh.FeMesh_Cartesian( elementType = (\"Q1/dQ0\"), \n",
    "                                 elementRes  = (res, res), \n",
    "                                 minCoord    = (0., 0.), \n",
    "                                 maxCoord    = (boxLength, boxHeight))\n",
    "\n",
    "velocityField       = mesh.add_variable(         nodeDofCount=2 )\n",
    "pressureField       = mesh.subMesh.add_variable( nodeDofCount=1 )\n",
    "temperatureField    = mesh.add_variable(         nodeDofCount=1 )\n",
    "temperatureDotField = mesh.add_variable(         nodeDofCount=1 )\n",
    "\n",
    "# initialise velocity, pressure and temperatureDot field\n",
    "velocityField.data[:]       = [0.,0.]\n",
    "pressureField.data[:]       = 0.\n",
    "temperatureField.data[:]    = 0.\n",
    "temperatureDotField.data[:] = 0."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set up material parameters and functions\n",
    "-----\n",
    "\n",
    "Set values and functions for viscosity, density and buoyancy force."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set a constant viscosity.\n",
    "viscosity = 1.\n",
    "\n",
    "# Create our density function.\n",
    "densityFn = Ra * temperatureField\n",
    "\n",
    "# Define our vertical unit vector using a python tuple (this will be automatically converted to a function).\n",
    "z_hat = ( 0.0, 1.0 )\n",
    "\n",
    "# A buoyancy function.\n",
    "buoyancyFn = densityFn * z_hat"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set initial temperature field\n",
    "-----\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Use a sinusodial perturbation**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "temperatureField.data[:] = 0.\n",
    "pertStrength = 0.1\n",
    "deltaTemp = tempMax - tempMin\n",
    "for index, coord in enumerate(mesh.data):\n",
    "    pertCoeff = math.cos( math.pi * coord[0]/boxLength ) * math.sin( math.pi * coord[1]/boxLength )\n",
    "    temperatureField.data[index] = tempMin + deltaTemp*(boxHeight - coord[1]) + pertStrength * pertCoeff\n",
    "    temperatureField.data[index] = max(tempMin, min(tempMax, temperatureField.data[index]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Show initial temperature field**\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# fig = vis.Figure()\n",
    "# fig.append( vis.objects.Surface(mesh, temperatureField) )\n",
    "# fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create boundary conditions\n",
    "----------\n",
    "\n",
    "Set temperature boundary conditions on the bottom ( ``MinJ`` ) and top ( ``MaxJ`` )."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "for index in mesh.specialSets[\"MinJ_VertexSet\"]:\n",
    "    temperatureField.data[index] = tempMax\n",
    "for index in mesh.specialSets[\"MaxJ_VertexSet\"]:\n",
    "    temperatureField.data[index] = tempMin"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Construct sets for the both horizontal and vertical walls. Combine the sets of vertices to make the ``I`` (left and right side walls) and ``J`` (top and bottom walls) sets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "iWalls = mesh.specialSets[\"MinI_VertexSet\"] + mesh.specialSets[\"MaxI_VertexSet\"]\n",
    "jWalls = mesh.specialSets[\"MinJ_VertexSet\"] + mesh.specialSets[\"MaxJ_VertexSet\"]\n",
    "\n",
    "freeslipBC = uw.conditions.DirichletCondition( variable        = velocityField, \n",
    "                                               indexSetsPerDof = (iWalls, jWalls) )\n",
    "tempBC     = uw.conditions.DirichletCondition( variable        = temperatureField, \n",
    "                                               indexSetsPerDof = (jWalls,) )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "System setup\n",
    "-----\n",
    "\n",
    "**Setup a Stokes system**\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "stokes = uw.systems.Stokes(    velocityField = velocityField, \n",
    "                               pressureField = pressureField,\n",
    "                               conditions    = [freeslipBC,],\n",
    "                               fn_viscosity  = viscosity, \n",
    "                               fn_bodyforce  = buoyancyFn )\n",
    "# get the default stokes equation solver\n",
    "solver = uw.systems.Solver( stokes )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# a function for the 2nd invariant strain rate tensor\n",
    "fn_sr2Inv = fn.tensor.second_invariant(fn.tensor.symmetric( velocityField.fn_gradient ))\n",
    "\n",
    "# a function for viscous dissipation, i.e.\n",
    "# the contraction of dev. stress tensor with strain rate tensor.\n",
    "vd =  2 * viscosity * 2 * fn_sr2Inv**2\n",
    "\n",
    "# function for adiabatic heating\n",
    "adiabatic_heating = Di * velocityField[1]*(temperatureField)\n",
    "\n",
    "# combine viscous dissipation and adiabatic heating\n",
    "# terms to the energy equation, via the argument 'fn_source'\n",
    "fn_source = Di/Ra * vd - adiabatic_heating\n",
    "\n",
    "### As discussed by King et al. (JI09) the volume integral of the viscous dissipation and \n",
    "### the adiabatic heating should balance.\n",
    "\n",
    "int_vd = uw.utils.Integral([Di/Ra*vd,adiabatic_heating], mesh)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Create an advection diffusion system**\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "advDiff = uw.systems.AdvectionDiffusion( phiField       = temperatureField, \n",
    "                                         phiDotField    = temperatureDotField, \n",
    "                                         velocityField  = velocityField, \n",
    "                                         fn_diffusivity = 1.0,\n",
    "                                         fn_sourceTerm  = fn_source,\n",
    "                                         conditions     = [tempBC,] )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Analysis tools\n",
    "-----"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "**Nusselt number**\n",
    "\n",
    "The Nusselt number is the ratio between convective and conductive heat transfer\n",
    "\n",
    "\\\\[\n",
    "Nu = -h \\frac{ \\int_0^l \\partial_z T (x, z=h) dx}{ \\int_0^l T (x, z=0) dx}\n",
    "\\\\]\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "nuTop    = uw.utils.Integral( fn=temperatureField.fn_gradient[1], \n",
    "                              mesh=mesh, integrationType='Surface', \n",
    "                              surfaceIndexSet=mesh.specialSets[\"MaxJ_VertexSet\"])\n",
    "\n",
    "nuBottom = uw.utils.Integral( fn=temperatureField,               \n",
    "                              mesh=mesh, integrationType='Surface', \n",
    "                              surfaceIndexSet=mesh.specialSets[\"MinJ_VertexSet\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Nusselt number = 1.000000\n"
     ]
    }
   ],
   "source": [
    "nu = - nuTop.evaluate()[0]/nuBottom.evaluate()[0]\n",
    "if uw.mpi.rank == 0 : print('Nusselt number = {0:.6f}'.format(nu))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**RMS velocity**\n",
    "\n",
    "The root mean squared velocity is defined by intergrating over the entire simulation domain via\n",
    "\n",
    "\\\\[\n",
    "\\begin{aligned}\n",
    "v_{rms}  =  \\sqrt{ \\frac{ \\int_V (\\mathbf{v}.\\mathbf{v}) dV } {\\int_V dV} }\n",
    "\\end{aligned}\n",
    "\\\\]\n",
    "\n",
    "where $V$ denotes the volume of the box."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initial vrms = 0.000\n"
     ]
    }
   ],
   "source": [
    "vrms = stokes.velocity_rms()\n",
    "if uw.mpi.rank == 0 : print('Initial vrms = {0:.3f}'.format(vrms))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Main simulation loop\n",
    "-----"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "#initialise time, step, output arrays\n",
    "time = 0.\n",
    "step = 0\n",
    "timeVal = []\n",
    "vrmsVal = []\n",
    "step_end = 30\n",
    "\n",
    "# output frequency\n",
    "step_output = max(1,min(100, step_end/10)) # reasonable automatic choice\n",
    "epsilon     = 1.e-8\n",
    "\n",
    "velplotmax = 0.0\n",
    "nuLast = -1.0\n",
    "rerr   = 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define an update function\n",
    "def update():\n",
    "    # Determining the maximum timestep for advancing the a-d system.\n",
    "    dt = advDiff.get_max_dt()\n",
    "    # Advect using this timestep size. \n",
    "    advDiff.integrate(dt)\n",
    "    return time+dt, step+1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "v_old = velocityField.copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "steps =      0; time = 0.000e+00; v_rms = 17.899; Nu = 1.000; Rel change = 2.000e+00 vChange = 1.000e+00\n",
      "steps =     10; time = 1.221e-03; v_rms = 21.672; Nu = 1.059; Rel change = 6.079e-03 vChange = 2.010e-02\n",
      "steps =     20; time = 2.441e-03; v_rms = 26.347; Nu = 1.137; Rel change = 8.114e-03 vChange = 1.939e-02\n",
      "steps =     30; time = 3.662e-03; v_rms = 31.759; Nu = 1.252; Rel change = 1.083e-02 vChange = 1.838e-02\n",
      "velocity relative tolerance is: 0.018\n"
     ]
    }
   ],
   "source": [
    "tol = 1e-8\n",
    "# Perform steps.\n",
    "while step<=step_end and rerr > tol:\n",
    "    \n",
    "    # copy to previous\n",
    "    v_old.data[:] = velocityField.data[:]\n",
    "    \n",
    "    # Solving the Stokes system.\n",
    "    solver.solve()\n",
    "     \n",
    "    aerr = uw.utils._nps_2norm(v_old.data-velocityField.data)\n",
    "    magV = uw.utils._nps_2norm(v_old.data)\n",
    "    rerr = ( aerr/magV if magV>1e-8 else 1)       # calculate relative variation\n",
    "\n",
    "    # Calculate & store the RMS velocity and Nusselt number.\n",
    "    vrms = stokes.velocity_rms()\n",
    "    nu   = - nuTop.evaluate()[0]/nuBottom.evaluate()[0]\n",
    "    vrmsVal.append(vrms)\n",
    "    timeVal.append(time)\n",
    "    velplotmax = max(vrms, velplotmax)\n",
    "\n",
    "    # print output statistics \n",
    "    if step%(step_end/step_output) == 0:\n",
    "\n",
    "#         mH = mesh.save(outputPath+\"mesh-{}.h5\".format(step))\n",
    "#         tH = temperatureField.save(outputPath+\"t-{}.h5\".format(step), mH)\n",
    "#         vH = velocityField.save(outputPath+\"v-{}.h5\".format(step), mH)\n",
    "#         velocityField.xdmf(outputPath+\"v-{}.xdmf\".format(step), vH, \"velocity\", mH, \"mesh\")\n",
    "#         temperatureField.xdmf(outputPath+\"t-{}.xdmf\".format(step), tH, \"temperature\", mH, \"mesh\" )\n",
    "\n",
    "        if(uw.mpi.rank==0):\n",
    "            print('steps = {0:6d}; time = {1:.3e}; v_rms = {2:.3f}; Nu = {3:.3f}; Rel change = {4:.3e} vChange = {5:.3e}'\n",
    "              .format(step, time, vrms, nu, abs((nu - nuLast)/nu), rerr))\n",
    "            \n",
    "#     # Check loop break conditions.\n",
    "#     if(abs((nu - nuLast)/nu) < epsilon):\n",
    "#         if(uw.mpi.rank==0):\n",
    "#             print('steps = {0:6d}; time = {1:.3e}; v_rms = {2:.3f}; Nu = {3:.3f}; Rel change = {4:.3e}'\n",
    "#               .format(step, time, vrms, nu, abs((nu - nuLast)/nu)))\n",
    "#         break\n",
    "    nuLast = nu\n",
    "    \n",
    "    # update\n",
    "    time, step = update()\n",
    "    \n",
    "print(\"velocity relative tolerance is: {:.3f}\".format(rerr))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Post analysis\n",
    "-----\n",
    "\n",
    "**Benchmark values**\n",
    "\n",
    "We can check the volume integral of viscous dissipation and adibatic heating are equal\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "vd, ad = int_vd.evaluate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "# error if >2% difference in vd and ad\n",
    "if not np.isclose(vd,ad, rtol=2e-2):\n",
    "    if uw.mpi.rank == 0: print('vd = {0:.3e}, ad = {1:.3e}'.format(vd,ad))\n",
    "    raise RuntimeError(\"The volume integral of viscous dissipation and adiabatic heating should be approximately equal\")"
   ]
  }
 ],
 "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}