Setonix singularity runner

runner.sh

#!/bin/bash -l
 
#SBATCH --account=pawsey0407
#SBATCH --job-name=uw2131
#SBATCH --ntasks=2
#SBATCH --ntasks-per-node=24
#SBATCH --time=00:20:00
#SBATCH --export=none
 
# load Singularity
module load singularity/3.8.6
 
# Define the container to use
export pawseyRepository=/scratch/pawsey0407/software/singularity/
export containerImage=$pawseyRepository/underworld2_2.13.1b.sif
 
# model name
#export model="05_Rayleigh_Taylor.py"
export model="Test.py"

# as per
# https://support.pawsey.org.au/documentation/pages/viewpage.action?pageId=116131367#UsewithSingularity-RunningPythonandR
# we unset all the host python-related ENV vars
unset $( env | grep ^PYTHON | cut -d = -f 1 | xargs )

# execute
srun --export=all -u -n $SLURM_NTASKS singularity exec -B ${PWD}:/work $containerImage bash -c "cd /work/; python3 $model"

Test.py

#!/usr/bin/env python
# coding: utf-8

# In[3]:


import underworld as uw
from underworld import function as fn
import underworld.visualisation as vis
uw.utils.matplotlib_inline()

import matplotlib.pyplot as pyplot
pyplot.ion()  # this is need so that we don't hang on show() for pure python runs
import numpy as np
import math


res = 16
mesh = uw.mesh.FeMesh_Cartesian( elementType = "Q1/dQ0", 
                                 elementRes  = (res, res), 
                                 minCoord    = (0., 0.), 
                                 maxCoord    = (1., 1.),
                                 periodic    = [True, False]  ) 

velocityField    = mesh.add_variable(         nodeDofCount=mesh.dim )
pressureField    = mesh.subMesh.add_variable( nodeDofCount=1 )

velocityField.data[:] = [0.,0.]
pressureField.data[:] = 0.

iWalls = mesh.specialSets["MinI_VertexSet"] + mesh.specialSets["MaxI_VertexSet"]
jWalls = mesh.specialSets["MinJ_VertexSet"] + mesh.specialSets["MaxJ_VertexSet"]

shearVelocity = 0.05

for index in mesh.specialSets["MinJ_VertexSet"]:
    velocityField.data[index] = [0., 0.]
for index in mesh.specialSets["MaxJ_VertexSet"]:
    velocityField.data[index] = [shearVelocity, 0.]
    
periodicBC = uw.conditions.DirichletCondition( variable        = velocityField, 
                                               indexSetsPerDof = ( jWalls, jWalls) ) 

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

materialIndex  = swarm.add_variable( dataType="int",    count=1 )

materialViscous        = 0
materialViscoelastic   = 1
materialViscoelastic2  = 2

xCoordFn = fn.input()[0]
conditions = [ ( xCoordFn > 0.6 , materialViscoelastic2),
               ( xCoordFn > 0.4 , materialViscoelastic ),
               ( xCoordFn < 0.4 , materialViscoelastic2)]

materialIndex.data[:]  = fn.branching.conditional( conditions ).evaluate(swarm)

# initialise swarm variables for viscoelastic rheology & analysis
previousStress         = swarm.add_variable( dataType="double", count=3 )
previousStress.data[:] = [0., 0., 0.]

# add another variable to track a single particle
markerVariable = swarm.add_variable( dataType="int", count=1)
# go ahead and mark a single particle on proc 0
markerVariable.data[:] = 0
if uw.mpi.rank==0:
    markerVariable.data[0] = 1
# also let's wrap in min_max so we can pull it out later.
# we pass in the previousStress as the aux function which 
# will be evaluated at the max value of the markerVariable
# (ie the marked particle).
minmax_marker = fn.view.min_max(markerVariable,fn_auxiliary=previousStress)

maxT    = 1.0    # max time for shearing velocity BC
eta     = 1.0e2  # viscosity
mu      = 1.0e2  # elastic modulus

alpha   = eta / mu                         # viscoelastic relaxation time
dt_e    = alpha / 10.                      # elastic time step
eta_eff = ( eta * dt_e ) / (alpha + dt_e)  # effective viscosity

nsteps       = int(10/dt_e*3.)+1     # number of steps to reach t = 10
velBCstep    = int(maxT / (dt_e/3.)) # timestep of maxT


# define viscosity
mappingDictViscosity = { materialViscous       : eta, 
                         materialViscoelastic  : eta_eff, 
                         materialViscoelastic2 : eta_eff }
viscosityMapFn       = fn.branching.map( fn_key=materialIndex, mapping=mappingDictViscosity )

# define strain rate tensor
strainRate = fn.tensor.symmetric( velocityField.fn_gradient )
strainRate_2ndInvariant = fn.tensor.second_invariant(strainRate)

# define stress tensor.
# tauHistoryFn will be passed into Stokes for the ve force term
viscousStressFn        = 2. * viscosityMapFn * strainRate 
tauHistoryFn           = eta_eff / ( mu * dt_e ) * previousStress
viscoelasticeStressFn  = viscousStressFn + tauHistoryFn

mappingDictStress = { materialViscous       : viscousStressFn, 
                      materialViscoelastic  : viscoelasticeStressFn, 
                      materialViscoelastic2 : viscoelasticeStressFn }
stressMapFn = fn.branching.map( fn_key=materialIndex, mapping=mappingDictStress )

# buoyancy force term
buoyancyFn   = ( 0.0, -1.0 )


stokes = uw.systems.Stokes( velocityField = velocityField, 
                               pressureField = pressureField,
                               voronoi_swarm = swarm, 
                               conditions    = [periodicBC,],
                               fn_viscosity  = viscosityMapFn, 
                               fn_bodyforce  = buoyancyFn,
                               fn_stresshistory = tauHistoryFn)

solver = uw.systems.Solver( stokes )

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


# define an update function
def update():
    # Retrieve the maximum possible timestep for the advection system.
    dt = advector.get_max_dt()
    if dt > ( dt_e / 3. ):
        dt = dt_e / 3.  

    # Advect using this timestep size.
    advector.integrate(dt)
    
    # smoothed stress history for use in (t + 1) timestep   
    phi = dt / dt_e;
    stressMapFn_data = stressMapFn.evaluate(swarm)
    previousStress.data[:] = ( phi*stressMapFn_data[:] + ( 1.-phi )*previousStress.data[:] )
    
    return time+dt, step+1

# Stepping. Initialise time and timestep.
time = 0.
step = 0

tTracer           = np.zeros(nsteps)
previousStress_xy = np.zeros(nsteps)

import os

outputPath = os.path.join(os.path.abspath("."),"test/")
#outputPath = 'output/'

if uw.mpi.rank==0:
    if not os.path.exists(outputPath):
        os.makedirs(outputPath)
        

while step < nsteps :
    # solve stokes problem
    solver.solve()
    
    # output for analysis               
    tTracer[step] = time

    # keep record of the marked particle's shear stress.
    minmax_marker.reset()         # reset minimax so that we will get a new aux_fn evaluation
    minmax_marker.evaluate(swarm) # evaluated across all particles. minmax will be record
    previousStress_xy[step] = minmax_marker.max_global_auxiliary()[:,2]   # extract evaluated value
    # We are finished with current timestep, update.
    time, step = update()
    
    swarm.save(outputPath+"swarm-"+str(step).zfill(4)+".h5")
    mesh.save(outputPath+"mesh-"+str(step).zfill(4)+".h5")
    
    # change BC if time > 1.0, then watch stress decay
    if step >= velBCstep:
        for index in mesh.specialSets["MaxJ_VertexSet"]:
            velocityField.data[index] = [0.0, 0.]