breaking OSMesa on raijin

gistfile1.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## A simple analytical solution for slab detachment - model 2\n",
    "Stefan M. Schmalholz, Earth and Planetary Science Letters 304 (2011) 45–54"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# maybe this helps the mpi setup\n",
    "import mpi4py.rc\n",
    "mpi4py.rc.threaded = False\n",
    "mpi4py.rc.thread_level = \"single\"\n",
    "\n",
    "# import all python modules\n",
    "from mpi4py import MPI\n",
    "import os\n",
    "import numpy as np\n",
    "\n",
    "import underworld as uw\n",
    "import glucifer\n",
    "\n",
    "import underworld.function as fn\n",
    "from underworld.scaling import units as u \n",
    "from underworld.scaling import non_dimensionalise as nd\n",
    "\n",
    "comm = MPI.COMM_WORLD\n",
    "RANK = uw.mpi.rank\n",
    "\n",
    "# set output path\n",
    "outputPath='model_b'\n",
    "outputPath = os.path.abspath(outputPath)+'/'\n",
    "# build output path on root proc\n",
    "if RANK == 0:\n",
    "    if not os.path.exists(outputPath):\n",
    "        print(\"Making outputPath, \", outputPath)\n",
    "        os.makedirs(outputPath)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def checkpoint( mesh, fieldDict, swarm, swarmDict, outputDir='./',\n",
    "                    meshName='mesh', swarmName='swarm', time=0.,\n",
    "                    enable_xdmf=True, with_deform_mesh=False):\n",
    "\n",
    "    if 'counter' not in checkpoint.__dict__:\n",
    "        checkpoint.counter = 0\n",
    "        \n",
    "    prefix = outputDir\n",
    "    # Check the prefix is valid\n",
    "    if not os.path.exists(prefix) and uw.mpi.rank==0:\n",
    "        raise RuntimeError(\"Error checkpointing, can't find {}\".format(prefix))\n",
    "    if not isinstance(time, (float,int)):\n",
    "        raise TypeError(\"'index' is not a number\")\n",
    "    ii = str(checkpoint.counter).zfill(5)\n",
    "\n",
    "    # only is there's a mesh then check for fields to save\n",
    "    if mesh is not None:\n",
    "\n",
    "        # Error check the mesh and fields\n",
    "        if not isinstance(mesh, uw.mesh.FeMesh):\n",
    "            raise TypeError(\"'mesh' is not of type uw.mesh.FeMesh\")\n",
    "        if not isinstance(fieldDict, dict):\n",
    "            raise TypeError(\"'fieldDict' is not of type dict\")\n",
    "        for key, value in fieldDict.items():\n",
    "            if not isinstance( value, uw.mesh.MeshVariable ):\n",
    "                raise TypeError(\"'fieldDict' must contain uw.mesh.MeshVariable elements\")\n",
    "\n",
    "        # Save the mesh each call of checkpoint() if with_deform_mesh is enabled\n",
    "        if with_deform_mesh is True:\n",
    "            mh = mesh.save(prefix+meshName+\".h5\")\n",
    "        else:\n",
    "            # see if we have already saved the mesh. It only needs to be saved once\n",
    "            if 'mH' not in checkpoint.__dict__:\n",
    "                checkpoint.mH = mesh.save(prefix+meshName+\".h5\")\n",
    "            mh = checkpoint.mH\n",
    "\n",
    "        # save xdmf files\n",
    "        for key,value in fieldDict.items():\n",
    "            filename = prefix+key+'-'+ii\n",
    "            handle   = value.save(filename+'.h5', mh)\n",
    "            if enable_xdmf: value.xdmf(filename, handle, key, mh, meshName, time)\n",
    "\n",
    "    # is there a swarm\n",
    "    if swarm is not None:\n",
    "\n",
    "        # Error check the swarms\n",
    "        if not isinstance(swarm, uw.swarm.Swarm):\n",
    "            raise TypeError(\"'swarm' is not of type uw.swarm.Swarm\")\n",
    "        if not isinstance(swarmDict, dict):\n",
    "            raise TypeError(\"'swarmDict' is not of type dict\")\n",
    "        for key, value in swarmDict.items():\n",
    "            if not isinstance( value, uw.swarm.SwarmVariable ):\n",
    "                raise TypeError(\"'fieldDict' must contain uw.swarm.SwarmVariable elements\")\n",
    "\n",
    "        # save xdmf files\n",
    "        sH = swarm.save(prefix+swarmName+\"-\"+ii+\".h5\")\n",
    "        for key,value in swarmDict.items():\n",
    "            filename = prefix+key+'-'+ii\n",
    "            handle   = value.save(filename+'.h5')\n",
    "            if enable_xdmf: value.xdmf(filename, handle, key, sH, swarmName, time)\n",
    "    \n",
    "    checkpoint.counter += 1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# desired scaling\n",
    "K_length    = 1 * u.kilometer\n",
    "K_viscosity = 1e18 * u.pascal * u.second\n",
    "K_density   = 1e3 * u.kilogram / (u.meter)**3\n",
    "\n",
    "# resultant scaling\n",
    "K_mass        = K_density * K_length**3\n",
    "K_seconds     = K_mass / ( K_length * K_viscosity )\n",
    "K_substance   = 1. * u.mole\n",
    "K_temperature = 1. * u.kelvin\n",
    "\n",
    "# define uw scaling coefficients\n",
    "scaling_coeff = uw.scaling.get_coefficients()\n",
    "scaling_coeff[\"[length]\"]      = K_length.to_base_units()\n",
    "scaling_coeff[\"[temperature]\"] = K_temperature.to_base_units()\n",
    "scaling_coeff[\"[time]\"]        = K_seconds.to_base_units()\n",
    "scaling_coeff[\"[mass]\"]        = K_mass.to_base_units()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## complete set of model parameters\n",
    "\n",
    "# general params\n",
    "max_time      = nd( 22.56 * u.megayear)\n",
    "min_viscosity = nd( 1e18 * u.Pa * u.sec)\n",
    "max_viscosity = nd( 1e25 * u.Pa * u.sec)\n",
    "gravity       = nd( 9.81 * u.meter / u.second**2)\n",
    "\n",
    "# geometry params\n",
    "Lx          = nd(1000 * u.kilometer)\n",
    "Ly          = nd( 660 * u.kilometer)\n",
    "air_depth   = nd(  40 * u.kilometer)\n",
    "plate_depth = nd(  80 * u.kilometer)\n",
    "slab_length = nd( 250 * u.kilometer)\n",
    "slab_width  = nd(  80 * u.kilometer)\n",
    "\n",
    "# testing for absolute density values\n",
    "# mantle_density    = nd(3150 * u.kilogram / u.meter**3)\n",
    "# slab_density      = nd(3300 * u.kilogram / u.meter**3)\n",
    "\n",
    "# mantle params: density & linear viscosity\n",
    "mantle_density     = nd(0 * u.kilogram / u.meter**3)\n",
    "mantle_viscosity   = nd(1e21 * u.Pa * u.sec)\n",
    "\n",
    "# slab params: density & nonlinear viscosity\n",
    "slab_density       = nd(150 * u.kilogram / u.meter**3)\n",
    "slab_exponent      = 4.\n",
    "mu                 = 4.75e11\n",
    "_mu                = nd(mu * u.Pa*u.sec**(1/slab_exponent))\n",
    "\n",
    "DEFAULT_STRAINRATE = nd(1e-16 / u.seconds)\n",
    "\n",
    "# numerical parameters\n",
    "resFactor = 2\n",
    "resX      = resFactor * 100\n",
    "resY      = resFactor * 66\n",
    "elType    = \"Q1/dq0\"\n",
    "\n",
    "# solve setup\n",
    "STOKES_NL_TOL    = 1e-2\n",
    "STOKES_OUTER_TOL = 1e-5\n",
    "STOKES_INNER_TOL = 1e-6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# FEM mesh and variable setup\n",
    "elRes    = (resX, resY)\n",
    "minCoord = [0.,0.]\n",
    "maxCoord = [Lx, Ly]\n",
    "mesh     = uw.mesh.FeMesh_Cartesian(maxCoord=maxCoord, \n",
    "                                    minCoord=minCoord,\n",
    "                                    elementRes=elRes, \n",
    "                                    elementType=elType)\n",
    "vField   = mesh.add_variable(2)\n",
    "pField   = mesh.subMesh.add_variable(1)\n",
    "etaField = mesh.add_variable(1)\n",
    "\n",
    "# derivative functions of FE variables\n",
    "vdotv     = fn.math.dot(vField, vField)\n",
    "srFn      = fn.tensor.symmetric(vField.fn_gradient)\n",
    "sr_2InvFn = fn.tensor.second_invariant(srFn) # 2nd invariant of strain rate \n",
    "\n",
    "# lagrangian swarm setup - stores material and viscosity\n",
    "swarm        = uw.swarm.Swarm(mesh)\n",
    "advector     = uw.systems.SwarmAdvector(vField, swarm, order=2)\n",
    "materialVar  = swarm.add_variable(count=1, dataType='int')\n",
    "viscosityVar = swarm.add_variable(count=1, dataType='double')\n",
    "\n",
    "# populate the swarm\n",
    "layout = uw.swarm.layouts.PerCellSpaceFillerLayout(swarm, particlesPerCell=20)\n",
    "swarm.populate_using_layout(layout)\n",
    "\n",
    "# enable population control\n",
    "population_control = uw.swarm.PopulationControl(swarm, \n",
    "                                                aggressive=False,\n",
    "                                                splitThreshold=0.15, \n",
    "                                                maxDeletions=2,\n",
    "                                                maxSplits=10,\n",
    "                                                particlesPerCell=20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# No-slip side walls, free-slip top and bottom\n",
    "bottomWall = mesh.specialSets[\"Bottom_VertexSet\"]\n",
    "topWall    = mesh.specialSets[\"Top_VertexSet\"]\n",
    "leftWall   = mesh.specialSets[\"Left_VertexSet\"]\n",
    "rightWall  = mesh.specialSets[\"Right_VertexSet\"]\n",
    "iWalls     = leftWall + rightWall\n",
    "allWalls   = iWalls + topWall + bottomWall\n",
    "\n",
    "vField.data[...] = 0.0\n",
    "\n",
    "# dirichlet conditions for model, no slip sides walls, free slip top and bottom\n",
    "stokes_dBC = uw.conditions.DirichletCondition( variable        = vField, \n",
    "                                               indexSetsPerDof = (iWalls, allWalls) )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## build geometries\n",
    "\n",
    "# plate geometry\n",
    "p_h   = mesh.maxCoord[1] - plate_depth\n",
    "array = np.array([(mesh.minCoord[0], p_h), (mesh.maxCoord[0], p_h),\n",
    "                  (mesh.maxCoord[0], mesh.maxCoord[1]), (mesh.minCoord[0], mesh.maxCoord[1])])\n",
    "\n",
    "slab2_polygon = uw.function.shape.Polygon(vertices=array)\n",
    "\n",
    "# slab geometry\n",
    "slab_x0 = mesh.maxCoord[0]/2 - slab_width/2\n",
    "slab_x1 = slab_x0 + slab_width\n",
    "s_h     = p_h - slab_length\n",
    "array   = np.array([(slab_x0, s_h), (slab_x1, s_h),\n",
    "                    (slab_x1, Ly ), (slab_x0, Ly )])\n",
    "slab_polygon = uw.function.shape.Polygon(vertices=array)\n",
    "\n",
    "# assign geometries to swarm via polygon class, this is a proxy for \n",
    "# the 2 materials - mantle material,slab material.\n",
    "# As such we can add (or) the polygon (boolean) evaluation.\n",
    "# to the materialVar swarm variable, i.e. (0,1) = (mantle material,slab material)\n",
    "materialVar.data[:] = slab2_polygon.evaluate(swarm) + slab_polygon.evaluate(swarm)\n",
    "\n",
    "# define for fn mappings later on\n",
    "(mantleIndex,slabIndex) = (0,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## create analysis swarms\n",
    "\n",
    "# Ds - slab tip depth single tracking particle\n",
    "ts_ds      = uw.swarm.Swarm(mesh)\n",
    "ts_ds_ad   = uw.systems.SwarmAdvector(vField, ts_ds, order=2)\n",
    "depthVar   = ts_ds.add_variable(count=1, dataType='double')\n",
    "\n",
    "p_coords      = np.ndarray(shape=(1, 2))\n",
    "p_coords[:,0] = Lx/2.\n",
    "p_coords[:,1] = slab_length\n",
    "\n",
    "ts_ds.add_particles_with_coordinates(p_coords)\n",
    "y_coord          = fn.coord()[1]\n",
    "fn_ds            = fn.misc.constant(1.0) - y_coord/Ly\n",
    "depthVar.data[:] = fn_ds.evaluate(ts_ds)\n",
    "\n",
    "# width tracers - 100 points positioned down the side edges of the slab. 50 points per side.\n",
    "ts_ws      = uw.swarm.Swarm(mesh)\n",
    "ts_ws_ad   = uw.systems.SwarmAdvector(vField, ts_ws, order=2)\n",
    "dummVar    = ts_ws.add_variable(count=1, dataType='int') # dummy var required to save\n",
    "\n",
    "p_coords           = np.ndarray(shape=(100, 2))\n",
    "p_coords[0:50,0]   = slab_x0\n",
    "p_coords[0:50,1]   = np.linspace(s_h, p_h, num=50)\n",
    "p_coords[50:100,0] = slab_x1\n",
    "p_coords[50:100,1] = np.linspace(s_h, p_h, num=50)\n",
    "\n",
    "ignore = ts_ws.add_particles_with_coordinates(p_coords)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# setup velocity vis\n",
    "figVel = glucifer.Figure(title=\"Velocity arrows and 2nd inv strain rate\")\n",
    "figVel.append(glucifer.objects.VectorArrows(mesh, vField))\n",
    "figVel.append(glucifer.objects.Surface(mesh, sr_2InvFn, logScale=True))\n",
    "# figVel.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# assign Fns for phyical properties\n",
    "\n",
    "# Density\n",
    "densityFn     = uw.function.branching.map( fn_key=materialVar, \n",
    "                                           mapping={ slabIndex: slab_density, mantleIndex: mantle_density } )\n",
    "# Force\n",
    "forceFn       = gravity*densityFn * (0.0,-1.)\n",
    "\n",
    "# Viscosity\n",
    "dummy_sr_2Inv = fn.misc.constant(DEFAULT_STRAINRATE)\n",
    "\n",
    "# strain rate and power-law viscosity\n",
    "sr      = fn.misc.max(sr_2InvFn, dummy_sr_2Inv)\n",
    "eta_raw = _mu*fn.math.pow(sr, (1./slab_exponent - 1.))\n",
    "\n",
    "# viscous limiting\n",
    "maxBound       = fn.misc.min(eta_raw,  max_viscosity)\n",
    "finalViscosity = fn.misc.max(maxBound, min_viscosity)\n",
    "\n",
    "viscosityFn = uw.function.branching.map( fn_key=materialVar, \n",
    "                                         mapping={ slabIndex: finalViscosity, mantleIndex: mantle_viscosity } )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# setup solver\n",
    "stokes = uw.systems.Stokes( velocityField = vField, \n",
    "                            pressureField = pField,\n",
    "                            conditions    = [stokes_dBC],\n",
    "                            fn_viscosity  = viscosityFn, \n",
    "                            fn_bodyforce  = forceFn )\n",
    "solver = uw.systems.Solver( stokes )\n",
    "\n",
    "solver.set_outer_rtol(STOKES_OUTER_TOL)\n",
    "# solver.set_inner_method(\"mumps\") # direct solve for speed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "figEta = glucifer.Figure(title=\"Viscosity [Pa*s]\")\n",
    "factor = 1/nd(1*u.Pa*u.s)\n",
    "figEta.append(glucifer.objects.Points(swarm, factor*viscosityFn, logScale=True, fn_size=2.0))\n",
    "\n",
    "figEps = glucifer.Figure(title=\"Strain Rate 2nd inv. [1/s]\")\n",
    "factor = 1/nd(1*u.s**-1)\n",
    "figEps.append(glucifer.objects.Points(swarm, factor*sr_2InvFn, logScale=True, fn_size=2.0))\n",
    "\n",
    "figRho = glucifer.Figure(title=\"Density [kg/m**3]\")\n",
    "factor = 1/nd(1*u.kg/u.m**3)\n",
    "figRho.append(glucifer.objects.Points(swarm, factor*densityFn))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "figRho.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Get solution\n",
    "# solver.solve(nonLinearIterate=True,nonLinearTolerance=1e-2)\n",
    "# projection = uw.utils.MeshVariable_Projection(etaField, stokes.fn_viscosity)\n",
    "# projection.solve()\n",
    "\n",
    "# fieldDict = {'velocity':vField, 'viscosityField':etaField}\n",
    "# checkpoint( mesh, fieldDict, None, None, 0, outputDir=outputPath) # for initial analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# redefine\n",
    "swarmDict = {'viscosity':viscosityVar}\n",
    "fieldDict = {'velocity':vField}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "# simulation parameters\n",
    "steps         = 0\n",
    "finalTimestep = 5\n",
    "curTime       = 0\n",
    "cp_i          = 0\n",
    "chkpt_every   = 1 # in numeric steps\n",
    "\n",
    "# Wall time timers\n",
    "timer_0, timer_1 = 0,0\n",
    "\n",
    "if RANK == 0:\n",
    "    outfile = open(outputPath+'buildMount.txt', 'w+')\n",
    "    string = \"steps, timestep, time, vrms, slab_tip_depth, Wtime\"\n",
    "    print(string)\n",
    "    outfile.write( string+\"\\n\")\n",
    "\n",
    "# initialise loop\n",
    "dt = -1\n",
    "\n",
    "timer_0 = MPI.Wtime()\n",
    "while steps<finalTimestep:\n",
    "# while curTime < max_time:\n",
    "    \n",
    "    depthVar.data[:] = fn_ds.evaluate(ts_ds)\n",
    "    \n",
    "    # Get solution\n",
    "    solver.solve(nonLinearIterate=True,nonLinearTolerance=STOKES_NL_TOL)\n",
    "    \n",
    "    viscosityVar.data[:] = stokes.fn_viscosity.evaluate(swarm)\n",
    "    \n",
    "    if steps % chkpt_every == 0:\n",
    "        checkpoint( mesh, fieldDict, swarm, swarmDict, \n",
    "                    outputDir=outputPath, time=curTime)\n",
    "        cp_i += 1 # due to paraview 5.4 oddity\n",
    "    ts_ws.save(outputPath+'/width_swarm.'+str(steps).zfill(4)+'.h5')\n",
    "    \n",
    "    # calculate metrics\n",
    "    v2int = mesh.integrate(fn=vdotv)[0]\n",
    "    vol   = mesh.integrate(fn=1.0)[0]\n",
    "    \n",
    "    s_d = 0\n",
    "    if ts_ds.particleLocalCount>0: \n",
    "        s_d = depthVar.data[0][0]\n",
    "    s_d = comm.allreduce(s_d)\n",
    "    \n",
    "    # get time step first time around\n",
    "    if dt < 0:\n",
    "        dt = advector.get_max_dt()\n",
    "    \n",
    "    curTime += dt\n",
    "    timer_1 = MPI.Wtime()\n",
    "    string = \"{}, {:.5e}, {:.5e}, {:.5e}, {:.5e}, {:.5e}\".format(steps,\n",
    "                                 uw.scaling.dimensionalise(dt, u.year),\n",
    "                                 uw.scaling.dimensionalise(curTime, u.year),\n",
    "                                 uw.scaling.dimensionalise(np.sqrt(v2int/vol), u.cm/u.year), \n",
    "                                 s_d,\n",
    "                                 (timer_1-timer_0) )\n",
    "    timer_0 = timer_1\n",
    "    if RANK == 0:\n",
    "        print(string)\n",
    "        outfile.write(string+\"\\n\")\n",
    "        \n",
    "    figEta.save(outputPath+\"eta-\"+str(steps)+\".png\")\n",
    "    figEps.save(outputPath+\"eps-\"+str(steps)+\".png\")\n",
    "    \n",
    "    # Advect particles   \n",
    "    advector.integrate(dt)\n",
    "    ts_ws_ad.integrate(dt)\n",
    "    ts_ds_ad.integrate(dt)\n",
    "    # population control\n",
    "    population_control.repopulate()  \n",
    "    \n",
    "    steps += 1\n",
    "\n",
    "if RANK==0:\n",
    "    outfile.close()"
   ]
  }
 ],
 "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.5.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}