How to use the underworld.swarm function in underworld

try:
            _fn_diffusivity = uw.function.Function.convert(fn_diffusivity)
        except Exception as e:
            raise uw._prepend_message_to_exception(e, "Exception encountered. Note that provided 'fn_diffusivity' must be of or convertible to 'Function' class.\nEncountered exception message:\n")

        if not fn_bodyforce:
            if pressureField.mesh.dim == 2:
                fn_bodyforce = (0.,0.)
            else:
                fn_bodyforce = (0.,0.,0.)
        try:
            _fn_bodyforce = uw.function.Function.convert(fn_bodyforce)
        except Exception as e:
            raise uw._prepend_message_to_exception(e, "Exception encountered. Note that provided 'fn_bodyforce' must be of or convertible to 'Function' class.\nEncountered exception message:\n")

        if voronoi_swarm and not isinstance(voronoi_swarm, uw.swarm.Swarm):
            raise TypeError( "Provided 'swarm' must be of 'Swarm' class." )
        self._swarm = voronoi_swarm
        if len(numpy.unique(voronoi_swarm.owningCell.data[:])) != len(voronoi_swarm.owningCell.data[:]):
            import warnings
            warnings.warn("It is not advised to fill any cell with more than one particle, as the Q1 shape function cannot capture material interfaces. Use at your own risk.")

        if velocityField and not isinstance( velocityField, uw.mesh.MeshVariable):
            raise TypeError( "Provided 'velocityField' must be of 'MeshVariable' class." )
        self._velocityField = velocityField

        if swarmVarVelocity and not isinstance(swarmVarVelocity, uw.swarm.SwarmVariable):
            raise TypeError("Provided 'swarmVarVelocity' must be of 'SwarmVariable' class.")
        self._swarmVarVelocity = swarmVarVelocity

        if voronoi_swarm and pressureField.mesh.elementType=='Q2':
            import warnings

mesh = self.phiField.mesh

        if self._mswarm == None:
            mswarm = uw.swarm.Swarm(mesh, particleEscape=True)
            mswarm_map = mswarm.add_variable(dataType="int", count=1)
            mswarm_home_pts = mswarm.add_variable(dataType="double", count=mesh.dim)
            mswarm_phiStar = mswarm.add_variable(dataType="float", count=1)

            local_nId = -1 * np.ones(mesh.nodesGlobal, dtype=np.int)
            for i, gId in enumerate(mesh.data_nodegId):
                local_nId[gId] = i

            # print("{} - building mswarm".format(uw.mpi.rank), flush=True )

            layout = uw.swarm.layouts.PerCellRandomLayout(mswarm, particlesPerCell=mesh.data_elementNodes[0].shape[0])
            mswarm.populate_using_layout(layout)

            # element_centroids = mesh.data[local_nId[mesh.data_elementNodes]].mean(axis=1)
            # element_centroids2 = element_centroids.reshape(tuple((*element_centroids.shape, 1)))
            # element_coords = mesh.data[local_nId[mesh.data_elementNodes]].transpose(0,2,1)
            # swarm_coords   = (element_coords - element_centroids2) * ratio + element_centroids2
            # swarm_coords2  = swarm_coords.transpose(0,2,1).reshape(-1, mesh.dim)

            # This is not optimised for the element loop
            # But there eliminates the initial search issues
            # associated with adding points to an empty swarm.

            ## print("{} - adding {} particles".format(uw.mpi.rank, swarm_coords2.shape[0]), flush=True )

            with mswarm.deform_swarm(update_owners=True):
                for el in range(0, mesh.elementsLocal):

_fn_bodyforce = uw.function.Function.convert(fn_bodyforce)
        except Exception as e:
            raise uw._prepend_message_to_exception(e, "Exception encountered. Note that provided 'fn_bodyforce' must be of or convertible to 'Function' class.\nEncountered exception message:\n")

        if voronoi_swarm and not isinstance(voronoi_swarm, uw.swarm.Swarm):
            raise TypeError( "Provided 'swarm' must be of 'Swarm' class." )
        self._swarm = voronoi_swarm
        if len(numpy.unique(voronoi_swarm.owningCell.data[:])) != len(voronoi_swarm.owningCell.data[:]):
            import warnings
            warnings.warn("It is not advised to fill any cell with more than one particle, as the Q1 shape function cannot capture material interfaces. Use at your own risk.")

        if velocityField and not isinstance( velocityField, uw.mesh.MeshVariable):
            raise TypeError( "Provided 'velocityField' must be of 'MeshVariable' class." )
        self._velocityField = velocityField

        if swarmVarVelocity and not isinstance(swarmVarVelocity, uw.swarm.SwarmVariable):
            raise TypeError("Provided 'swarmVarVelocity' must be of 'SwarmVariable' class.")
        self._swarmVarVelocity = swarmVarVelocity

        if voronoi_swarm and pressureField.mesh.elementType=='Q2':
            import warnings
            warnings.warn("Voronoi integration may yield unsatisfactory results for Q2 element types. Q2 element types can also result in spurious velocity calculations.")

        if not isinstance( _removeBCs, bool):
            raise TypeError( "Provided '_removeBCs' must be of type bool." )
        self._removeBCs = _removeBCs

        # error check dcs 'dirichlet conditions' and ncs 'neumann cond.' FeMesh_IndexSets
        nbc  = None
        mesh = pressureField.mesh

        if not isinstance(conditions,(list,tuple)):

indexSetsPerDof = (iWalls, base) )

for index in mesh.specialSets["MinI_VertexSet"]:
    velocityField.data[index] = [minXv, 0.]
for index in mesh.specialSets["MaxI_VertexSet"]:
    velocityField.data[index] = [maxXv, 0.]
    


# ### Setup the material swarm
# 
# This is used for tracking deformation and history dependence of the rheology

# In[7]:

swarm         = uw.swarm.Swarm( mesh=mesh )
swarmLayout   = uw.swarm.layouts.GlobalSpaceFillerLayout( swarm=swarm, particlesPerCell=20 )
swarm.populate_using_layout( layout=swarmLayout )


# ### Create a particle advection system
# 
# Note that we need to set up one advector systems for each particle swarm (our global swarm and a separate one if we add passive tracers).

# In[8]:

advector        = uw.systems.SwarmAdvector( swarm=swarm,        velocityField=velocityField, order=2 )


# ### Add swarm variables
# 
# We are using a single material with a single rheology. We need to track the plastic strain in order to have some manner of strain-related softening (e.g. of the cohesion or the friction coefficient). For visualisation of swarm data we need an actual swarm variable and not just the computation.

# ### Setup the material swarm and passive tracers
# 
# The material swarm is used for tracking deformation and history dependence of the rheology
# 
# Passive swarms can track all sorts of things but lack all the machinery for integration and re-population

# In[5]:

swarm  = uw.swarm.Swarm( mesh=mesh, particleEscape=True )
swarmLayout = uw.swarm.layouts.GlobalSpaceFillerLayout( swarm=swarm, particlesPerCell=50 )
swarm.populate_using_layout( layout=swarmLayout )

# For population control, make this an integration swarm
popSwarm = uw.swarm.VoronoiIntegrationSwarm(swarm, particlesPerCell=50)

surfaceSwarm = uw.swarm.Swarm( mesh=mesh, particleEscape=True )
deformationSwarm = uw.swarm.Swarm ( mesh=mesh, particleEscape=True )


# ### Create a particle advection system
# 
# Note that we need to set up one advector systems for each particle swarm (our global swarm and a separate one if we add passive tracers).

# In[6]:

advector        = uw.systems.SwarmAdvector( swarm=swarm,            velocityField=velocityField, order=2 )
advector2       = uw.systems.SwarmAdvector( swarm=surfaceSwarm,     velocityField=velocityField, order=2 )
advector3       = uw.systems.SwarmAdvector( swarm=deformationSwarm, velocityField=velocityField, order=2 )

# populated.

        self.empty = False

        # Should do some checking first

        self.mesh = mesh
        self.velocity = velocityField
        self.thickness = fthickness
        self.ID = fID
        self.insidePt = insidePt
        self.director = None

        # Set up the swarm and variables on all procs

        self.swarm = uw.swarm.Swarm( mesh=self.mesh, particleEscape=True )
        self.director = self.swarm.add_variable( dataType="double", count=2)
        self._swarm_advector = uw.systems.SwarmAdvector( swarm=self.swarm,
                                                         velocityField=self.velocity, order=2 )

        self.swarm.add_particles_with_coordinates(np.stack((pointsX, pointsY)).T)
        self.director.data[...] = 0.0

        self._update_kdtree()
        self._update_surface_normals()

        return

# create force vectors
        self._fvector = sle.AssembledVector(temperatureField, tEqNums )

        # and matrices
        self._kmatrix = sle.AssembledMatrix( self._solutionVector, self._solutionVector, rhs=self._fvector )

        # we will use voronoi if that has been requested by the user, else use
        # gauss integration.
        if self._swarm:
            intswarm = self._swarm._voronoi_swarm
            # need to ensure voronoi is populated now, as assembly terms will call
            # initial test functions which may require a valid voronoi swarm
            self._swarm._voronoi_swarm.repopulate()
        else:
            intswarm = uw.swarm.GaussIntegrationSwarm(mesh)


        self._kMatTerm = sle.MatrixAssemblyTerm_NA_i__NB_i__Fn(  integrationSwarm=intswarm,
                                                                 assembledObject=self._kmatrix,
                                                                 fn=_fn_diffusivity)
        self._forceVecTerm = sle.VectorAssemblyTerm_NA__Fn(   integrationSwarm=intswarm,
                                                              assembledObject=self._fvector,
                                                              fn=fn_heating)
        # search for neumann conditions
        for cond in conditions:
            if isinstance( cond, uw.conditions.NeumannCondition ):
                #NOTE many NeumannConditions can be used but the _sufaceFluxTerm only records the last

                ### -VE flux because of the FEM discretisation method of the initial equation
                negativeCond = uw.conditions.NeumannCondition( fn_flux=cond.fn_flux,
                                                               variable=cond.variable,

# Create a particle swarm
# ------
# 
# 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.

# **Setup a swarm**
# 
# To set up a swarm of particles the following steps are needed:
# 1. Initialise and name a swarm, here called ``swarm``.
# 2. Define data variable (``materialIndex``) to store an index that will state what material a given particle is.
# 3. Populate the swarm over the whole domain using the layout command, here this is used to allocate 20 particles in each element.

# In[6]:

# Create the swarm. 
swarm = uw.swarm.Swarm( mesh=mesh )

# Add a data variable which will store an index to determine material.
materialIndex = swarm.add_variable( dataType="int", count=1 )

# Create a layout object that will populate the swarm across the whole domain.
swarmLayout = uw.swarm.layouts.GlobalSpaceFillerLayout( swarm=swarm, particlesPerCell=20 )

# Go ahead and populate the swarm.
swarm.populate_using_layout( layout=swarmLayout )


# **Define a shape**
# 
# Define a python function that mathematically describes a shape, in this case a circle offset to the centre of the sinker. Note that this returns an underworld function, which can be used in the branching condition function below.

# In[7]:

# Create utilities
        self.MaterialIndexFieldFloat = self.swarm.add_variable(dataType="double",  count=1 )
        self.DensityVar = uw.mesh.MeshVariable(self.mesh, nodeDofCount=1)
        self.MaterialVar = uw.mesh.MeshVariable(self.mesh, nodeDofCount=1)
        self.projectorDensity = uw.utils.MeshVariable_Projection(self.DensityVar,
                                                                 self.densityFn,
                                                                 type=0 )
        self.projectorMaterial = uw.utils.MeshVariable_Projection(self.MaterialVar,
                                         self.MaterialIndexFieldFloat, type=0 )

        if not self.mesh._cself.isRegular:
            raise TypeError("You are using an irregular mesh: \
                             isostasy module only works with regular meshes")

        if self.surface is not None and not isinstance(self.surface, uw.swarm._swarm.Swarm):
            raise TypeError("'surface' must be a tuple'")
       
        self.initialized = True

def __init__(self, meshVariable=None, fn=None, voronoi_swarm=None, type=0, **kwargs):

        if not meshVariable:
            raise ValueError("You must specify a mesh variable via the 'meshVariable' parameter.")
        if not isinstance( meshVariable, uw.mesh.MeshVariable):
            raise TypeError( "Provided 'meshVariable' must be of 'MeshVariable' class." )
        self._meshVariable = meshVariable

        if not fn:
            raise ValueError("You must specify a function via the 'fn' parameter.")
        try:
            _fn = uw.function.Function.convert(fn)
        except Exception as e:
            raise uw._prepend_message_to_exception(e, "Exception encountered. Note that provided 'fn' must be of or convertible to 'Function' class.\nEncountered exception message:\n")

        if voronoi_swarm and not isinstance(voronoi_swarm, uw.swarm.Swarm):
            raise TypeError( "Provided 'swarm' must be of 'Swarm' class." )
        self._swarm = voronoi_swarm
        if voronoi_swarm and meshVariable.mesh.elementType=='Q2':
            import warnings
            warnings.warn("Voronoi integration may yield unsatisfactory results for Q2 mesh.")
            


        if not type in [0,1]:
            raise ValueError( "Provided 'type' must take a value of 0 (weighted average) or 1 (weighted residual)." )
        self.type = type

        eqNum = sle.EqNumber(meshVariable)
        # create force vectors
        self._fvector = sle.AssembledVector(meshVariable, eqNum)