How to use the underworld.mesh.MeshVariable function in underworld

# In[23]:

# build a non-partitioned mesh with same box size
mesh0 = uw.mesh.FeMesh_Cartesian( elementType = ("Q1/dQ0"), 
                                  elementRes  = (res, res), 
                                  minCoord    = (0., 0.), 
                                  maxCoord    = (boxLength, boxHeight),
                                  partitioned = False ) 
# load previous mesh coordinate data onto new non-partitioned mesh
mesh0.load(outputPath+'mesh.h5')

# load T, P and V data onto the new mesh
# note that pressure is always on the submesh
temperatureField0  = uw.mesh.MeshVariable( mesh=mesh0, nodeDofCount=1 ) 
pressureField0     = uw.mesh.MeshVariable( mesh=mesh0.subMesh, nodeDofCount=1 ) 
velocityField0     = uw.mesh.MeshVariable( mesh=mesh0, nodeDofCount=dim ) 
temperatureField0.load(outputPath+"tempfield.h5")
pressureField0.load(outputPath+"presfield.h5")
velocityField0.load(outputPath+"velfield.h5")


# **Temperature gradient**
# 
# The final benchmarks in the Blankenbach paper involve the temperature gradient in the vertical direction ($\frac{\partial T}{\partial z}$). This is easy to find using the underworld functions, as shown below.
# 

# In[24]:

if(uw.rank()==0):
    tempgradField = temperatureField0.fn_gradient
    vertTGradField = - boxHeight * tempgradField[1] / tempMax # scaled for direct benchmarking below

# old = self.data
            (x,y) = (self.data[:,1], self.data[:,2] )

            # Move nodes away from the axis
            r = np.hypot(x,y)
            x[ r == 0.0 ] += 0.000001 * self._cyl_size[2]

            theta = np.arctan2(y,x)
            scale  = np.maximum(np.cos(theta%(np.pi*0.5)), np.sin(theta%(np.pi*0.5)))

            self.data[:,0] = self.data[:,0] * (top - base) + base
            self.data[:,1] *= scale
            self.data[:,2] *= scale

        # add a boundary MeshVariable - 1 if nodes is on the boundary(ie 'AllWalls_VertexSet'), 0 if node is internal
        self.bndMeshVariable = uw.mesh.MeshVariable(self, 1)
        self.bndMeshVariable.data[:] = 0.
        self.bndMeshVariable.data[self.specialSets["AllWalls_VertexSet"].data] = 1.0


        self._boundaryNodeFn = fn.branching.conditional(
                                        [  ( self.bndMeshVariable != 0.0 , 1. ),
                                        (                            True, 0. )   ] )

        """
        Rotation documentation.
        We will create 3 basis vectors that will rotate the (x,y,z) problem to be a
        (r, theta, z) natural coordinate system for this mesh

        The rotations are performed on the local element level using the existing machinery
        provided by UW2. As such only elements on the domain boundary need rotation, all internal
        elements can continue with the (x,y,z) representation.

# 
# Read (``savedRes`` $\times$ ``savedRes``) resolution data for $P$, $v$ and $T$ fields as well as existing summary statistics data. These are converted into lists so that the main time loop below will append with new values.
# 

# In[ ]:

if(LoadFromFile == True):
# set up mesh for savedRes*savedRes data file
    meshSaved = uw.mesh.FeMesh_Cartesian( elementType = ("Q1/dQ0"), 
                                       elementRes  = (savedRes, savedRes), 
                                       minCoord    = (0., 0.), 
                                       maxCoord    = (1., 1.) )
    temperatureFieldSaved    = uw.mesh.MeshVariable( mesh=meshSaved,         nodeDofCount=1 )
    temperatureDotFieldSaved = uw.mesh.MeshVariable( mesh=meshSaved,         nodeDofCount=1 )
    pressureFieldSaved       = uw.mesh.MeshVariable( mesh=meshSaved.subMesh, nodeDofCount=1 )
    velocityFieldSaved       = uw.mesh.MeshVariable( mesh=meshSaved,         nodeDofCount=2 )

    temperatureFieldSaved.load(    inputPath+'temperatureField.h5' )
    temperatureDotFieldSaved.load( inputPath+'temperatureDotField.h5' )
    pressureFieldSaved.load(       inputPath+'pressureField.h5')
    velocityFieldSaved.load(       inputPath+'velocityField.h5')
    
    if(res==savedRes): # no remeshing needed, copy data directly
        temperatureField.data[:]    = temperatureFieldSaved.data[:]
        pressureField.data[:]       = pressureFieldSaved.data[:]
        velocityField.data[:]       = velocityFieldSaved.data[:]
        temperatureDotField.data[:] = temperatureDotFieldSaved.data[:]
    else: # remeshing needed
        temperatureField.data[:]    = temperatureFieldSaved.evaluate( mesh )
        pressureField.data[:]       = pressureFieldSaved.evaluate( mesh.subMesh )
        velocityField.data[:]       = velocityFieldSaved.evaluate( mesh )
        temperatureDotField.data[:] = temperatureDotFieldSaved.evaluate( mesh )

minCoord = tuple([nd(val) for val in self.minCoord])
        maxCoord = tuple([nd(val) for val in self.maxCoord])

        self.mesh = uw.mesh.FeMesh_Cartesian(elementType=self.elementType,
                                             elementRes=self.elementRes,
                                             minCoord=minCoord,
                                             maxCoord=maxCoord,
                                             periodic=self.periodic)
        
        # Add common fields
        self.temperature = None
        self.pressureField = uw.mesh.MeshVariable(mesh=self.mesh.subMesh, nodeDofCount=1)
        self.velocityField = uw.mesh.MeshVariable(mesh=self.mesh, nodeDofCount=self.mesh.dim)
        self.tractionField = uw.mesh.MeshVariable(mesh=self.mesh, nodeDofCount=self.mesh.dim)
        self.strainRateField = uw.mesh.MeshVariable(mesh=self.mesh, nodeDofCount=1)
        self.pressureField.data[...] = 0.
        self.velocityField.data[...] = 0.
        self.tractionField.data[...] = 0.
       
        # symmetric component of the gradient of the flow velocityField.
        self.strainRate_default = strainRate_default
        self.solutionExist = False
        self.strainRate = fn.tensor.symmetric(self.velocityField.fn_gradient)
        self.strainRate_2ndInvariant = fn.tensor.second_invariant(self.strainRate)
 
        # Create the material swarm
        self.swarm = uw.swarm.Swarm(mesh=self.mesh, particleEscape=True)
        self.swarmLayout = swarmLayout
        self.population_control = populationControl

        self.materials = []

super(AnnulusConvection, self).__init__(**kwargs)

        # Because we must build the mesh first, before the fields, this wrapper is difficult to seperate into
        # a straight themo-mechanical wrapper for any mesh discretisation

        # create the FEM mesh
        self.mesh = annulus = uw.mesh._FeMesh_Annulus(elementRes=elRes,
                                      radialLengths=radialLengths, angularExtent=(0.,360.),
                                      periodic = [False, True])

        # create the fields
        self.fields = dict()
        
        # stokes fields
        vField = self.fields['velocity']    = uw.mesh.MeshVariable(annulus, nodeDofCount=2)
        self.fields['pressure']             = uw.mesh.MeshVariable(annulus.subMesh, nodeDofCount=1)
        # heat equation fields
        tField = self.fields['temperature'] = uw.mesh.MeshVariable(annulus, nodeDofCount=1)
        self.fields['tDot']                 = uw.mesh.MeshVariable(annulus, nodeDofCount=1)

        # fields to checkpoint by default
        # self.checkpoint_fields = dict(self.fields.items()) # possibility to chkp all fields
        self.checkpoint_fields = dict( temperature = tField,
                                       tDot        = self.fields['tDot'],
                                       velocity    = vField             )

        # create the nodes Sets
        self.meshSets=dict()
        outer = self.meshSets['outerNodes'] = annulus.specialSets["MaxI_VertexSet"]
        inner = self.meshSets['innerNodes'] = annulus.specialSets["MinI_VertexSet"]
        self.meshSets['boundaryNodes'] = self.meshSets['outerNodes']+self.meshSets['innerNodes']

# Create mesh and variables
# ------

# In[5]:

mesh = uw.mesh.FeMesh_Cartesian( elementType = ("Q1/dQ0"), 
                                 elementRes  = (res, res), 
                                 minCoord    = (0., 0.), 
                                 maxCoord    = (boxLength, boxHeight))

velocityField       = uw.mesh.MeshVariable( mesh=mesh,         nodeDofCount=dim )
pressureField       = uw.mesh.MeshVariable( mesh=mesh.subMesh, nodeDofCount=1 )
temperatureField    = uw.mesh.MeshVariable( mesh=mesh,         nodeDofCount=1 )
temperatureDotField = uw.mesh.MeshVariable( mesh=mesh,         nodeDofCount=1 )

# initialise velocity, pressure and temperatureDot field
velocityField.data[:]       = [0.,0.]
pressureField.data[:]       = 0.
temperatureField.data[:]    = 0.
temperatureDotField.data[:] = 0.


# Set up material parameters and functions
# -----
# 
# Set values and functions for viscosity, density and buoyancy force.

# In[6]:

# Set a constant viscosity.

mesh = uw.mesh.FeMesh_Cartesian( elementType = ("Q1/dQ0"), 
                                 elementRes  = (int(boxLength*res), res), 
                                 minCoord    = (0., 0.), 
                                 maxCoord    = (boxLength, boxHeight))

Tmesh = mesh


# Create mesh variables.  Note the pressure field uses the sub-mesh, the temperature uses the same mesh at this stage. The `temperatureDotField` variable is a work array for the advection-diffusion solver.


velocityField       = uw.mesh.MeshVariable( mesh=mesh,         nodeDofCount=dim )
pressureField       = uw.mesh.MeshVariable( mesh=mesh.subMesh, nodeDofCount=1 )

temperatureField    = uw.mesh.MeshVariable( mesh=Tmesh,         nodeDofCount=1 )
temperatureDotField = uw.mesh.MeshVariable( mesh=Tmesh,         nodeDofCount=1 )

# Initialise values
velocityField.data[:]       = [0.,0.]
pressureField.data[:]       = 0.
temperatureDotField.data[:] = 0.


# Set up material parameters and functions
# -----
# 
# Set functions for viscosity, density and buoyancy force. These functions and variables only need to be defined at the beginning of the simulation, not each timestep.


# Set viscosity to be a constant.
viscosity = 1.

def __init__(self, variable, indexSetsPerDof):
        if not isinstance( variable, uw.mesh.MeshVariable ):
            raise TypeError("Provided variable must be of class 'MeshVariable'.")
        self._variable = variable

        if isinstance( indexSetsPerDof, uw.container.IndexSet ):
            indexSets = ( indexSetsPerDof, )
        elif isinstance( indexSetsPerDof, (list,tuple)):
            indexSets = indexSetsPerDof
        else:
            raise TypeError("You must provide the required 'indexSetsPerDof' item\n"+
                             "as a list or tuple of 'IndexSet' items.")
        for guy in indexSets:
            if not isinstance( guy, (uw.container.IndexSet,type(None)) ):
                raise TypeError("Provided list must only contain objects of 'NoneType' or type 'IndexSet'.")
        self._indexSets = indexSets

        if variable.nodeDofCount != len(self._indexSets):

# Add Common Swarm Variables
        self.materialField = self.swarm.add_variable(dataType="int", count=1)
        self.plasticStrain = self.swarm.add_variable(dataType="double", count=1)
        self.viscosityField = self.swarm.add_variable(dataType="double", count=1)
        self.densityField = self.swarm.add_variable(dataType="double", count=1)
        self.plasticStrain.data[...] = 0.0
        self.viscosityField.data[...] = 0.
        self.densityField.data[...] = 0.
 

        # Create a bunch of tools to project swarmVariable onto the mesh
        self._projMaterialField  = uw.mesh.MeshVariable( mesh=self.mesh,nodeDofCount=1)
        self._materialFieldProjector = uw.utils.MeshVariable_Projection(self._projMaterialField, self.materialField, type=0)

        self._projViscosityField = uw.mesh.MeshVariable(mesh=self.mesh, nodeDofCount=1)
        self._viscosityFieldProjector = uw.utils.MeshVariable_Projection(self._projViscosityField, self.viscosityField, type=0)

        self._projPlasticStrain = uw.mesh.MeshVariable(mesh=self.mesh, nodeDofCount=1)
        self._plasticStrainProjector = uw.utils.MeshVariable_Projection(self._projPlasticStrain,
                                                                     self.plasticStrain, type=0)
        
        self._projStrainRate = uw.mesh.MeshVariable(mesh=self.mesh, nodeDofCount=1)
        self._strainRateProjector = uw.utils.MeshVariable_Projection(self._projStrainRate,
                                                     self.strainRate_2ndInvariant, type=0)
        
        self._projDensityField = uw.mesh.MeshVariable(mesh=self.mesh, nodeDofCount=1)
        self._densityFieldProjector = uw.utils.MeshVariable_Projection(self._projDensityField, self.densityField, type=0)