Example_Reynolds_Strouhal_CylinderInFlow.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import underworld as uw\n",
    "from underworld import function as fn\n",
    "from underworld import visualisation as glucifer\n",
    "import numpy as np\n",
    "import math\n",
    "import warnings\n",
    "import os,csv\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "from underworld.scaling import units as u\n",
    "from underworld.scaling import non_dimensionalise as nd\n",
    "\n",
    "outputPath = os.path.join(os.path.abspath(\".\"),\"Reynolds_Tube_FixBall2/\")\n",
    "if uw.mpi.rank == 0:\n",
    "    if not os.path.exists(outputPath):\n",
    "        os.makedirs(outputPath)\n",
    "uw.mpi.barrier()\n",
    "\n",
    "scaling_coefficients = uw.scaling.get_coefficients()\n",
    "\n",
    "tempMin = 273.*u.degK \n",
    "tempMax = (1400.+ 273.)*u.degK\n",
    "bodyforce = 3300 * u.kilogram / u.metre**3 * 9.81 * u.meter / u.second**2\n",
    "velocity = 1.0 *u.meter/u.second\n",
    "\n",
    "KL = 0.1*u.meter\n",
    "Kt = KL/velocity\n",
    "KT = tempMax \n",
    "KM = bodyforce * KL**2 * Kt**2\n",
    "K  = 1.*u.mole\n",
    "\n",
    "scaling_coefficients[\"[length]\"] = KL\n",
    "scaling_coefficients[\"[time]\"] = Kt\n",
    "scaling_coefficients[\"[mass]\"]= KM\n",
    "scaling_coefficients[\"[temperature]\"] = KT\n",
    "scaling_coefficients[\"[substance]\"] = K\n",
    "\n",
    "\n",
    "gravity = 0*nd(9.81 * u.meter / u.second**2)\n",
    "R = nd(8.3144621 * u.joule / u.mole / u.degK)\n",
    "\n",
    "time_factor=nd(1*u.second)\n",
    "vis_factor = nd(1*u.pascal*u.second)\n",
    "vel_factor = nd(1*u.meter/u.second) \n",
    "strainRate_factor = nd(1/u.second)\n",
    "length_factor = nd(1*u.centimeter)\n",
    "stress_factor = nd(1*u.pascal)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Setup parameters\n",
    "-----\n",
    "\n",
    "Set simulation parameters for the test and position of the cylinder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the resolution.\n",
    "resx = 64\n",
    "resy = 128\n",
    "# Set size and position of dense sphere.\n",
    "sphereRadius = nd(2.*1.27/2*u.centimeter)\n",
    "sphereCentre = (0.,nd(80.*u.centimeter))\n",
    "minX = nd(-12*u.centimeter)\n",
    "maxX = nd(12*u.centimeter)\n",
    "minY = 0.\n",
    "maxY = nd(120.*u.centimeter)\n",
    "\n",
    "Vy  =  -nd(1.0*u.meter/u.second)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "vis = nd(0.031*u.pascal*u.second)\n",
    "rho = nd(876.*u.kilogram/u.meter**3)\n",
    "D = nd(1.27*u.centimeter)\n",
    "v = nd(2*u.meter/u.second)\n",
    "Re = rho*v*D/vis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check the desinged Reynolds nunber"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "717.7548387096776\n"
     ]
    }
   ],
   "source": [
    "print(Re)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create mesh and finite element variables\n",
    "------"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\tGlobal element size: 64x128\n",
      "\tLocal offset of rank 0: 0x0\n",
      "\tLocal range of rank 0: 64x128\n"
     ]
    }
   ],
   "source": [
    "mesh = uw.mesh.FeMesh_Cartesian( elementType = (\"Q1/dQ0\"), \n",
    "                                 elementRes  = (resx, resy), \n",
    "                                 minCoord    = (minX, minY), \n",
    "                                 maxCoord    = (maxX, maxY),\n",
    "                                 periodic    =  [False, False] )\n",
    "\n",
    "velocityField    = mesh.add_variable(         nodeDofCount=2 )\n",
    "pressureField    = mesh.subMesh.add_variable( nodeDofCount=1 )\n",
    "stressField      = mesh.add_variable(         nodeDofCount=2 )\n",
    "step = 0 \n",
    "if uw.mpi.rank == 0:\n",
    "\n",
    "    dicMesh = { 'elements' : mesh.elementRes, \n",
    "                'minCoord' : mesh.minCoord,\n",
    "                'maxCoord' : mesh.maxCoord}\n",
    "\n",
    "    fo = open(outputPath+\"dicMesh\"+str(step).zfill(4),'w')\n",
    "    fo.write(str(dicMesh))\n",
    "    fo.close()  \n",
    "\n",
    "uw.mpi.barrier()\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set initial conditions and boundary conditions\n",
    "----------\n",
    "\n",
    "**Initial and boundary conditions**\n",
    "\n",
    "Initialise the velocity and pressure fields to zero."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "velocityField.data[:] = [0.,0.]\n",
    "pressureField.data[:] = 0."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Conditions on the boundaries**\n",
    "\n",
    "Construct sets for the both horizontal and vertical walls to define conditons for underworld solvers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "iWalls = mesh.specialSets[\"MinI_VertexSet\"] + mesh.specialSets[\"MaxI_VertexSet\"]\n",
    "jWalls = mesh.specialSets[\"MinJ_VertexSet\"] + mesh.specialSets[\"MaxJ_VertexSet\"]\n",
    "top = mesh.specialSets[\"MaxJ_VertexSet\"]\n",
    "base = mesh.specialSets[\"MinJ_VertexSet\"]\n",
    "\n",
    "Ball = mesh.specialSets['Empty']\n",
    "#make the ball/cylinder fixed\n",
    "for index, coord in enumerate(mesh.data):\n",
    "    x = coord[0]\n",
    "    y = coord[1]\n",
    "    sphere =  (x - sphereCentre[0])**2+(y - sphereCentre[1])**2\n",
    "\n",
    "    if sphere < sphereRadius**2:\n",
    "        Ball+=index\n",
    "        \n",
    "for index in mesh.specialSets[\"MaxJ_VertexSet\"]:\n",
    "    velocityField.data[index] = Vy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Boundary condition for the example of a falling cylinder "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# freeslipBC = uw.conditions.DirichletCondition( variable      = velocityField, \n",
    "#                                                indexSetsPerDof = (iWalls,iWalls+jWalls) )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Boundary condition for the example of  a fixed cylinder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "freeslipBC = uw.conditions.DirichletCondition( variable      = velocityField, \n",
    "                                               indexSetsPerDof = (iWalls,Ball+jWalls) )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a particle swarm\n",
    "------\n",
    "\n",
    "Swarms refer to (large) groups of particles which can advect with the fluid flow. These can be used to determine 'materials' as they can carry local information such as the fluid density and viscosity."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Setup a swarm**\n",
    "\n",
    "To set up a swarm of particles the following steps are needed:\n",
    "1. Initialise and name a swarm, here called ``swarm``.\n",
    "2. Define data variable (``materialIndex``) to store an index that will state what material a given particle is.\n",
    "3. Populate the swarm over the whole domain using the layout command, here this is used to allocate 20 particles in each element."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the swarm and an advector associated with it\n",
    "swarm = uw.swarm.Swarm( mesh=mesh,particleEscape=True )\n",
    "advector = uw.systems.SwarmAdvector( swarm=swarm, velocityField=velocityField, order=2 )\n",
    "\n",
    "# Add a data variable which will store an index to determine material.\n",
    "materialIndex = swarm.add_variable( dataType=\"int\", count=1 )\n",
    "previousVm =  swarm.add_variable( dataType=\"double\", count=2 )\n",
    "velSwarm =  swarm.add_variable( dataType=\"double\", count=2 )\n",
    "densitySwarm =  swarm.add_variable( dataType=\"double\", count=1 )\n",
    "\n",
    "# Create a layout object that will populate the swarm across the whole domain.\n",
    "swarmLayout = uw.swarm.layouts.PerCellGaussLayout( swarm=swarm, gaussPointCount=5 )\n",
    "pop_control = uw.swarm.PopulationControl(swarm,aggressive=True,particlesPerCell=20)\n",
    "\n",
    "# Go ahead and populate the swarm.\n",
    "swarm.populate_using_layout( layout=swarmLayout )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Define a shape**\n",
    " 1. Define an underworld `function` that descripes the geometry of a shape (a sphere), called `fn_sphere`.\n",
    " 2. Set up a `fn.branching.conditional` to run `fn_sphere` and return a given index - either `materialLightIndex` or `materialHeavyIndex`.\n",
    " 3. Execute the above underworld functions on the swarm we created and save the result on the `materialIndex` swarm variable.\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a function for a sphere. returns `True` if query is inside sphere, `False` otherwise.\n",
    "coord = fn.input() - sphereCentre\n",
    "fn_sphere = fn.math.dot( coord, coord ) < sphereRadius**2\n",
    "\n",
    "# define some names for our index \n",
    "materialLightIndex = 0\n",
    "materialHeavyIndex = 1\n",
    "\n",
    "# set up the condition for being in a sphere. If not in sphere then will return light index.\n",
    "conditions = [ ( fn_sphere , materialHeavyIndex), \n",
    "               ( True      , materialLightIndex) ]\n",
    "\n",
    "# Execute the branching conditional function, evaluating at the location of each particle in the swarm.\n",
    "# The results are copied into the materialIndex swarm variable.\n",
    "materialIndex.data[:] = fn.branching.conditional( conditions ).evaluate(swarm)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Branching conditional function**\n",
    "\n",
    "For more information on the `fn.branching.conditional` see the Functions user guide [here](../user_guide/05_Functions.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Define minimum y coordinate function**\n",
    "\n",
    " 1. Define a new swarm called `tracerSwarm`, with one particle at the base of the sphere. (This swarm behaves as a passive tracer swarm).\n",
    " 2. Define a function that finds the minimum y coordinate value of the `tracerSwarm` in a parallel safe way."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# build a tracer swarm with one particle\n",
    "tracerSwarm = uw.swarm.Swarm(mesh,particleEscape=True)\n",
    "advector_tracer = uw.systems.SwarmAdvector( swarm=tracerSwarm, velocityField=velocityField, order=2 )\n",
    "\n",
    "# build a numpy array with one particle, specifying it's exact location\n",
    "x_pos = sphereCentre[0]\n",
    "y_pos = sphereCentre[1]#-sphereRadius\n",
    "coord_array = np.array(object=(x_pos,y_pos),ndmin=2)\n",
    "tracerSwarm.add_particles_with_coordinates(coord_array)\n",
    "\n",
    "# define a y coordinate `min_max` function\n",
    "fn_ycoord = fn.view.min_max( fn.coord()[1] )\n",
    "\n",
    "def GetSwarmYMin():\n",
    "    fn_ycoord.reset()\n",
    "    fn_ycoord.evaluate(tracerSwarm),velocityField.evaluate(tracerSwarm.particleCoordinates.data)\n",
    "    return fn_ycoord.min_global(),velocityField[1].evaluate(tracerSwarm.particleCoordinates.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 0.,  8.]])"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tracerSwarm.data[:]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Design band swarm to visualize flow patterns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "bandSwarm = uw.swarm.Swarm(mesh,particleEscape=True)\n",
    "advector_band = uw.systems.SwarmAdvector( swarm=bandSwarm, velocityField=velocityField, order=2 )\n",
    "material_band = bandSwarm.add_variable( dataType=\"int\", count=1 )\n",
    "#swarmLayout_band = uw.swarm.layouts.PerCellGaussLayout( swarm=bandSwarm, gaussPointCount=1 )\n",
    "swarmLayout_band = uw.swarm.layouts.PerCellSpaceFillerLayout( swarm=bandSwarm, particlesPerCell=8 )\n",
    "pop_control_band = uw.swarm.PopulationControl(bandSwarm,aggressive=True,particlesPerCell=8)\n",
    "\n",
    "bandSwarm.populate_using_layout( layout=swarmLayout_band )\n",
    "\n",
    "y  = fn.input()[1]\n",
    "x  = fn.input()[0]\n",
    "nband = 19\n",
    "bW = (maxX-minX)/nband/2\n",
    "x_c = np.linspace(minX,maxX,nband)\n",
    "# conditions_band = [ (fn_sphere, 3),\n",
    "#                     ( (y > sphereCentre[1]+sphereRadius)&(fn.math.abs(x-x_c[1])<bW) , 1),\n",
    "#                     ( (y > sphereCentre[1]+sphereRadius)&(fn.math.abs(x-x_c[3])<bW) , 2),\n",
    "#                     ( (y > sphereCentre[1]+sphereRadius)&(fn.math.abs(x-x_c[5])<bW) , 3),\n",
    "#                     ( (y > sphereCentre[1]+sphereRadius)&(fn.math.abs(x-x_c[7])<bW) , 1),\n",
    "#                     ( (y > sphereCentre[1]+sphereRadius)&(fn.math.abs(x-x_c[9])<bW) , 2),\n",
    "#                     ( (y > sphereCentre[1]+sphereRadius)&(fn.math.abs(x-x_c[10])<bW) ,3),\n",
    "#                     ( (y > sphereCentre[1]+sphereRadius)&(fn.math.abs(x-x_c[13])<bW) ,1),\n",
    "#                     ( (y > sphereCentre[1]+sphereRadius)&(fn.math.abs(x-x_c[15])<bW) ,2),\n",
    "#                     ( (y > sphereCentre[1]+sphereRadius)&(fn.math.abs(x-x_c[17])<bW) ,3),\n",
    "#                      (True,      0) ]\n",
    "\n",
    "conditions_band = [ (fn_sphere, 4),\n",
    "                    ((fn.math.abs(x-x_c[1])<bW) , 1),\n",
    "                    ((fn.math.abs(x-x_c[3])<bW) , 2),\n",
    "                    ((fn.math.abs(x-x_c[5])<bW) , 3),\n",
    "                    ((fn.math.abs(x-x_c[7])<bW) , 1),\n",
    "                    ((fn.math.abs(x-x_c[9])<bW) , 2),\n",
    "                    ((fn.math.abs(x-x_c[11])<bW) , 3),\n",
    "                    ((fn.math.abs(x-x_c[13])<bW) , 1),\n",
    "                    ((fn.math.abs(x-x_c[15])<bW) , 2),\n",
    "                    ((fn.math.abs(x-x_c[17])<bW) , 3),\n",
    "                     (True,      0) ]\n",
    "\n",
    "material_band.data[:] =  fn.branching.conditional( conditions_band ).evaluate(bandSwarm)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "** Test minimum y coordinate function**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Minimum y value for sinker = 8.000e+00,velocity = 8.000e+00)\n"
     ]
    }
   ],
   "source": [
    "outPutTrack = GetSwarmYMin()\n",
    "print('Minimum y value for sinker = {0:.3e},velocity = {0:.3e})'.format(outPutTrack[0],outPutTrack[1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Plot the particles by material**\n",
    "\n",
    "Plot the initial positions of all swarm particles coloured by their material indices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'material' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[33], line 2\u001b[0m\n\u001b[1;32m      1\u001b[0m fig1 \u001b[38;5;241m=\u001b[39m glucifer\u001b[38;5;241m.\u001b[39mFigure( figsize\u001b[38;5;241m=\u001b[39m(\u001b[38;5;241m600\u001b[39m,\u001b[38;5;241m1400\u001b[39m) )\n\u001b[0;32m----> 2\u001b[0m fig1\u001b[38;5;241m.\u001b[39mPoints(swarm, \u001b[43mmaterial\u001b[49m, colourBar\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m, pointSize\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1.3\u001b[39m)\n\u001b[1;32m      3\u001b[0m fig1\u001b[38;5;241m.\u001b[39mVectorArrows(mesh, velocityField,scale\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1e3\u001b[39m)\n\u001b[1;32m      4\u001b[0m fig1\u001b[38;5;241m.\u001b[39mshow()\n",
      "\u001b[0;31mNameError\u001b[0m: name 'material' is not defined"
     ]
    }
   ],
   "source": [
    "fig1 = glucifer.Figure( figsize=(600,1400) )\n",
    "fig1.Points(swarm, material, colourBar=True, pointSize=1.3)\n",
    "fig1.VectorArrows(mesh, velocityField,scale=1e3)\n",
    "fig1.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set up material parameters and functions\n",
    "-----\n",
    "\n",
    "Here the functions for density and viscosity are set using the ``map`` function. This function evaluates a key function (here the material index), and the result (i.e. the key) is used to determine which function to evaluate to obtain the actual result (such as the particle density). \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set constants for the viscosity and density of the sinker.\n",
    "viscSphere = nd(1e6*u.pascal*u.second)\n",
    "densitySphere = nd(7780. * u.kilogram / u.metre**3 )\n",
    "\n",
    "# Here we set a viscosity value of '1.' for both materials  nd(0.31e-1*u.pascal*u.second)\n",
    "mappingDictViscosity = { materialLightIndex:nd(0.31e-1*u.pascal*u.second), materialHeavyIndex:viscSphere }\n",
    "# Create the viscosity map function.\n",
    "viscosityMapFn = fn.branching.map( fn_key=materialIndex, mapping=mappingDictViscosity )\n",
    "# Here we set a density of '0.' for the lightMaterial, and '1.' for the heavymaterial.\n",
    "mappingDictDensity = { materialLightIndex:nd(876. * u.kilogram / u.metre**3 ), materialHeavyIndex:densitySphere }\n",
    "# Create the density map function.\n",
    "densityFn = fn.branching.map( fn_key=materialIndex, mapping=mappingDictDensity )\n",
    "\n",
    "# And the final buoyancy force function.\n",
    "z_hat = ( 0.0, -0.0 )\n",
    "buoyancyFn = densityFn * z_hat * gravity\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "previousVm.data[:] = 0.\n",
    "dt_e = fn.misc.constant(nd(1e-1*u.second))\n",
    "\n",
    "LHS_fn = densityFn/dt_e \n",
    "RHS_fn = densityFn*previousVm/dt_e"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "stokes0 = uw.systems.Stokes(    velocityField = velocityField, \n",
    "                               pressureField = pressureField,\n",
    "                               #voronoi_swarm = swarm, \n",
    "                               conditions    = [freeslipBC,],\n",
    "                               fn_viscosity  = viscosityMapFn, \n",
    "                               #fn_bodyforce  = buoyancyFn)\n",
    "                               #fn_stresshistory = inertiaFn,\n",
    "                               fn_bodyforce  = buoyancyFn+RHS_fn)\n",
    "                               #fn_one_on_lambda  = oneonlambdaFn)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "massMatrixTerm = uw.systems.sle.MatrixAssemblyTerm_NA__NB__Fn(\n",
    "                        assembledObject  = stokes0._kmatrix,\n",
    "                        integrationSwarm = stokes0._constitMatTerm._integrationSwarm,\n",
    "                        fn   = LHS_fn,\n",
    "                        mesh = mesh)\n",
    "\n",
    "solver0 = uw.systems.Solver( stokes0 )\n",
    "solver0.set_inner_method('mumps') \n",
    "solver0.set_penalty(1e3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "top = mesh.specialSets[\"MaxJ_VertexSet\"]\n",
    "surfaceArea = uw.utils.Integral(fn=1.0,mesh=mesh, integrationType='surface', surfaceIndexSet=top)\n",
    "surfacePressureIntegral = uw.utils.Integral(fn=pressureField, mesh=mesh, integrationType='surface', surfaceIndexSet=top)\n",
    "\n",
    "# a callback function to calibrate the pressure - will pass to solver later\n",
    "def pressure_calibrate():\n",
    "    (area,) = surfaceArea.evaluate()\n",
    "    (p0,) = surfacePressureIntegral.evaluate()\n",
    "    offset = p0/area\n",
    "    #print(\"Zeroing pressure using mean upper surface pressure {}\".format( offset ))\n",
    "    pressureField.data[:] -= offset\n",
    "    #velSwarm.data[:] = velocityField.evaluate(swarm)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Analysis tools\n",
    "-----\n",
    "\n",
    "**RMS velocity**\n",
    "\n",
    "Set up integrals used to calculate the RMS velocity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "vdotv = fn.math.dot( velocityField, velocityField )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Main simulation loop\n",
    "-----\n",
    "\n",
    "The main time stepping loop begins here. Before this the time and timestep are initialised to zero and the output statistics arrays are set up. Also the frequency of outputting basic statistics to the screen is set in steps_output.\n",
    "\n",
    "Note that there are two ``advector.integrate`` steps, one for each swarm, that need to be done each time step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define an update function\n",
    "def update():\n",
    "    # Retrieve the maximum possible timestep for the advection system.\n",
    "    \n",
    "    dt1 = advector.get_max_dt()\n",
    "    dt2 = dt_e.value\n",
    "    dt = min(dt1,dt2)\n",
    "    \n",
    "    dt_e.value = dt\n",
    "    print (\"dt_adv=\",dt1/nd(1.*u.second),\"dt_e=\",dt_e.value/nd(1.*u.second),\"dt=\",dt/nd(1.*u.second))        \n",
    "\n",
    "    previousVm.data[:] =  velocityField.evaluate(swarm) #( phi*previousVm_data[:] + ( 1.-phi )*previousVm.data[:] )\n",
    "    \n",
    "    \n",
    "    #densitySwarm = densityFn.evaluate(swarm)\n",
    "    \n",
    "    # Advect using this timestep size.    \n",
    "    advector.integrate(dt)\n",
    "    advector_tracer.integrate(dt)\n",
    "    advector_band.integrate(dt)\n",
    "    \n",
    "    pop_control.repopulate()\n",
    "    pop_control_band.repopulate()\n",
    "    \n",
    "    return time+dt, step+1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Linear solver (OO4DX61N__system-execute) \n",
      "\n",
      "BSSCR -- Block Stokes Schur Compliment Reduction Solver \n",
      "\n",
      "\n",
      "-----  K2_GMG  ------\n",
      "\n",
      "AUGMENTED LAGRANGIAN K2 METHOD - Penalty = 1000.000000\n",
      "\n",
      "\n",
      "\t* K+p*K2 in time: 0.011886 seconds\n",
      "\n",
      "  Setting schur_pc to \"gkgdiag\" \n",
      "\n",
      "\n",
      "SCR Solver Summary:\n",
      "\n",
      "  RHS V Solve:            = 0.02989 secs / 1 its\n",
      "  Pressure Solve:         = 206.2 secs / 10000 its\n",
      "  Final V Solve:          = 0.02037 secs / 1 its\n",
      "\n",
      "  Total BSSCR Linear solve time: 210.372666 seconds\n",
      "\n",
      "Linear solver (OO4DX61N__system-execute), solution time 2.103912e+02 (secs)\n",
      "step =      0; time = 0.000e+00; v_rms = 1.196e+05; height = 8.00000e+01\n",
      "Vy-max= 628270.848263\n",
      "dt_adv= 2.9837138618742044e-09 dt_e= 2.9837138618742044e-09 dt= 2.9837138618742044e-09\n",
      "Time Integration\n",
      "\t2nd order:                 EX9JD1SA__integrand -    0.1462 [min] /    0.1462 [max] (secs)\n",
      "Time Integration - 0.146288 [min] / 0.146288 [max] (secs)\n",
      "Time Integration\n",
      "\t2nd order:                 90CDJ7S1__integrand -    0.0000 [min] /    0.0000 [max] (secs)\n",
      "Time Integration - 5.88894e-05 [min] / 5.88894e-05 [max] (secs)\n",
      "Time Integration\n",
      "\t2nd order:                 KE2416J4__integrand -    0.0521 [min] /    0.0521 [max] (secs)\n",
      "Time Integration - 0.0521569 [min] / 0.0521569 [max] (secs)\n",
      "In func WeightsCalculator_CalculateAll(): for swarm \"EILB2G4H__swarm\"\n",
      "\tdone 33% (2731 cells)...\n",
      "\tdone 67% (5462 cells)...\n",
      "\tdone 100% (8192 cells)...\n",
      "WeightsCalculator_CalculateAll(): finished update of weights for swarm \"EILB2G4H__swarm\"\n",
      "In func WeightsCalculator_CalculateAll(): for swarm \"AOFBISSH__swarm\"\n",
      "\tdone 33% (2731 cells)...\n",
      "\tdone 67% (5462 cells)...\n",
      "\tdone 100% (8192 cells)...\n",
      "WeightsCalculator_CalculateAll(): finished update of weights for swarm \"AOFBISSH__swarm\"\n"
     ]
    }
   ],
   "source": [
    "# Stepping. Initialise time and timestep.\n",
    "time = 0.\n",
    "step = 0\n",
    "nsteps = 1\n",
    "\n",
    "tSinker = np.zeros(nsteps)\n",
    "ySinker = np.zeros([nsteps,3])\n",
    "steps_output = 20\n",
    "\n",
    "if step == 0:\n",
    "    title = ['step','time','Y','Vy']\n",
    "    with open(outputPath+'Sample.csv', 'w') as f:\n",
    "        csv_write = csv.writer(f)\n",
    "        csv_write.writerow(title)\n",
    "        \n",
    "fn_Vy = fn.view.min_max(velocityField[1])  \n",
    "\n",
    "# Perform 10 steps\n",
    "while step<nsteps:\n",
    "\n",
    "    \n",
    "    solver0.solve()\n",
    "    if step % steps_output == 0 or step == nsteps-1:\n",
    "        mesh.save(outputPath+\"mesh\"+str(step).zfill(4)) \n",
    "        swarm.save(outputPath+\"swarm\"+str(step).zfill(4))\n",
    "        materialIndex.save(outputPath+\"materialVariable\"+str(step).zfill(4))\n",
    "        velocityField.save(outputPath+\"velocityField\"+str(step).zfill(4))\n",
    "        bandSwarm.save(outputPath+\"bandSwarm\"+str(step).zfill(4))\n",
    "        material_band.save(outputPath+\"materialBand\"+str(step).zfill(4))\n",
    "        \n",
    "    outPutTrack = GetSwarmYMin()\n",
    "    # Calculate the RMS velocity\n",
    "    vrms = math.sqrt( mesh.integrate(vdotv)[0] / mesh.integrate(1.)[0] )\n",
    "    if uw.mpi.rank==0:\n",
    "        ySinker[step,0] = outPutTrack[0]\n",
    "        ySinker[step,1] = outPutTrack[1]\n",
    "        tSinker[step] = time\n",
    "        print('step = {0:6d}; time = {1:.3e}; v_rms = {2:.3e}; height = {3:.5e}'\n",
    "              .format(step,time/time_factor,vrms/vel_factor,outPutTrack[0]/length_factor))\n",
    "        print  (\"Vy-max=\",np.max(-velocityField.data[:,1]/nd(1.*u.meter/u.second)))\n",
    "        fn_ycoord.evaluate(tracerSwarm)\n",
    "    fn_Vy.evaluate(tracerSwarm)\n",
    "    velocityField.evaluate(tracerSwarm.particleCoordinates.data)\n",
    "    \n",
    "    ySinker[step,0] = fn_ycoord.min_global()\n",
    "    ySinker[step,1] = fn_Vy.min_global()\n",
    "    ySinker[step,2] = time\n",
    "    tSinker[step] = time       \n",
    "    \n",
    "    if uw.mpi.rank == 0:\n",
    "  \n",
    "        dicMesh = { 'elements' : mesh.elementRes, \n",
    "                    'minCoord' : mesh.minCoord,\n",
    "                    'maxCoord' : mesh.maxCoord}\n",
    "\n",
    "        fo = open(outputPath+\"dicMesh\"+str(step).zfill(4),'w')\n",
    "        fo.write(str(dicMesh))\n",
    "        fo.close()  \n",
    "        \n",
    "    uw.mpi.barrier()\n",
    "    \n",
    "    if uw.mpi.rank== 0:\n",
    "        SP_output = [step,time, ySinker[step,0],ySinker[step,1]]\n",
    "        with open(outputPath+'Sample.csv', 'a') as f:\n",
    "            csv_write = csv.writer(f)  \n",
    "            csv_write.writerow(SP_output)\n",
    "    uw.mpi.barrier()        \n",
    "    \n",
    "    # update\n",
    "    time, step = update()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "underworld",
   "language": "python",
   "name": "underworld"
  },
  "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
}