Tosi et al., 2015 firedrake

tosi-case2.py

from firedrake import *
from firedrake.petsc import PETSc
from mpi4py import MPI
import math, numpy

#########################################################################################################
################################## Some important constants etc...: #####################################
#########################################################################################################

mesh   = UnitSquareMesh(40, 40) # This mesh is generated via a firedrake function
bottom_id, top_id = 3, 4 # When using firedrake generated mesh
left_id, right_id = 1, 2 # When using firedrake generated mesh

# Global Constants:
steady_state_tolerance = 1.0e-8
max_timesteps          = 1000
target_cfl_no          = 0.5
max_timestep           = 0.1

# Stokes related constants:
lambda_T     = math.log(10**5)
eta_star     = 10**-3
yield_stress = 1.
Ra           = Constant(100.0)
k            = Constant((0,1)) # Radial unit vector (in direction opposite to gravity):

# Temperature related constants:
delta_t = Constant(1e-3) # Initial time-step
kappa   = Constant(1.0)  # Thermal diffusivity

#### Print function to ensure log output is only written on processor zero (if running in parallel) ####
def log(*args):
    PETSc.Sys.Print(*args)

def log_params(f, str):
    if mesh.comm.rank == 0:
        f.write(str)
        f.flush()

# Stokes Equation Solver Parameters:
# direct solve:
# stokes_solver_parameters = {
#     'snes_type': 'ksponly',
#     'ksp_type': 'gmres',
#     'pc_type': 'lu',
#     'pc_factor_mat_solver_type': 'mumps',
#     'mat_type': 'aij'
# }
stokes_solver_parameters = {
    'snes_type': 'ksponly',
    'ksp_type': 'gmres',
    'pc_type': 'fieldsplit',
    'pc_fieldsplit_type': 'schur',
    'pc_fieldsplit_schur_fact_type': 'diag',
    'fieldsplit_0_ksp_type': 'preonly',
    'fieldsplit_0_pc_type': 'python',
    'fieldsplit_0_pc_python_type': 'firedrake.AssembledPC',
    'fieldsplit_0_assembled_pc_type': 'hypre',
    'fieldsplit_1_ksp_type': 'preonly',
    'fieldsplit_1_pc_type': 'python',
    'fieldsplit_1_pc_python_type': 'firedrake.MassInvPC',
    'fieldsplit_1_Mp_ksp_type': 'preonly',
    'fieldsplit_1_Mp_pc_type': 'lu',
    'fieldsplit_1_Mp_pc_factor_mat_solver_type': 'mumps',
    'mat_type': 'matfree',
    'ksp_monitor_true_residual': None,
}

# Temperature Equation Solver Parameters:
temperature_solver_parameters = {
    'snes_type': 'ksponly',
    'ksp_type': 'preonly',
    'pc_type': 'lu',
    'pc_factor_mat_solver_type': 'mumps',
    'mat_type': 'aij'
}

#########################################################################################################
################################## Geometry and Spatial Discretization: #################################
#########################################################################################################

# Set up function spaces - currently using the P2P1 element pair :
V    = VectorFunctionSpace(mesh, "CG", 2) # Velocity function space (vector)
W    = FunctionSpace(mesh, "CG", 1) # Pressure function space (scalar)
Q    = FunctionSpace(mesh, "CG", 1) # Temperature function space (scalar)
E    = FunctionSpace(mesh, "DG", 1) # Strain_rate second invariant function space

# Set up mixed function space and associated test functions:
Z    = MixedFunctionSpace([V, W])
N, M = TestFunctions(Z)
Y    = TestFunction(Q)
E_DG = TestFunction(E)

# Set up fields on these function spaces - split into each component so that they are easily accessible:
z    = Function(Z) # a field over the mixed function space Z.
u, p = split(z)
mu_field    = Function(W, name="Viscosity")
epsii_field = Function(E, name="Eps_II")

# Timestepping - CFL related stuff:
delta_x = sqrt(CellVolume(mesh))
ts_func = Function(Q) # Note that time stepping should be dictated by Temperature related mesh.

def compute_timestep(u):
    # A function to compute the timestep, based upon the CFL criterion
    ts_func.interpolate(delta_x / sqrt(dot(u,u)))
    ts_min = ts_func.dat.data.min()
    ts_min = mesh.comm.allreduce(ts_min, MPI.MIN)
    return min(ts_min*target_cfl_no,max_timestep)

#########################################################################################################
############################ T advection diffusion equation Prerequisites: ##############################
#########################################################################################################

# Set up temperature field and initialise based upon coordinates:
X       = SpatialCoordinate(mesh)
T_old   = Function(Q, name="OldTemperature")
T_old.interpolate((1 - X[1]) + 0.01 * cos(numpy.pi * X[0]) * sin(numpy.pi * X[1]))

T_new   = Function(Q, name="Temperature")
T_new.assign(T_old)

# Temporal discretisation - Using a Crank-Nicholson scheme where theta = 0.5:
theta   = 0.5
T_theta = theta * T_new + (1-theta) * T_old

#########################################################################################################
############################################ Setup Equations ############################################
#########################################################################################################

### Initially deal with Stokes equations ###

# Strain-rate:
epsilon    = 0.5 * (grad(u)+transpose(grad(u))) # Strain-rate
epsii_init = sqrt(inner(epsilon,epsilon)/2.) # 2nd invariant
epsii      = conditional(ge(epsii_init, 1e-40), epsii_init, 1e-40) # bounded
# Viscosity
#mu_lin     = exp(-lambda_T*T_new + (1 - X[1])) # with pressure-dependent part
mu_lin     = exp(-lambda_T * T_new)
mu_plast   = eta_star + (yield_stress / (numpy.sqrt(2) * epsii))
mu         = (2. / ((1./mu_lin) + (1./mu_plast)))
# Equation in weak (ufl) form - note that continuity equation is added here - need to better understand why:
tau        = 2. * mu * epsilon # Stress
F_stokes   = inner(grad(N), tau) * dx + dot(N, grad(p)) * dx - (dot(N,k)*Ra*T_theta) * dx
F_stokes  += dot(grad(M), u) * dx # Continuity equation

# Set up boundary conditions for Stokes: We first need to extract the velocity field from the mixed
# function space Z. This is done using Z.sub(0). We subsequently need to extract the x and y components
# of velocity. This is done using an additional .sub(0) and .sub(1), respectively. Note that the final arguments
# here are the physical boundary ID's.
bcvx = DirichletBC(Z.sub(0).sub(0), 0, (left_id,right_id))
bcvy = DirichletBC(Z.sub(0).sub(1), 0, (bottom_id,top_id))

### Next deal with Temperature advection-diffusion equation: ###
delta_x  = sqrt(CellVolume(mesh))
Y_SUPG   = Y + dot(u,grad(Y)) * (delta_x / (2*sqrt(dot(u,u))))
F_energy = Y_SUPG * ((T_new - T_old) / delta_t) * dx + Y_SUPG*dot(u,grad(T_theta)) * dx + dot(grad(Y),kappa*grad(T_theta)) * dx
bct_base = DirichletBC(Q, 1.0, bottom_id)
bct_top  = DirichletBC(Q, 0.0, top_id)

nullspace = MixedVectorSpaceBasis(Z, [Z.sub(0), VectorSpaceBasis(constant=True)])

# Write output files in VTK format:
u, p = z.split() # Do this first to extract individual velocity and pressure fields:
u.rename("Velocity")
p.rename("Pressure")
u_file     = File('velocity.pvd')
p_file     = File('pressure.pvd')
t_file     = File('temperature.pvd')
mu_file    = File('viscosity.pvd')
epsii_file = File('epsII.pvd')

period = 50

f = open('params.log', 'w')

# Now perform the time loop:
for timestep in range(0, max_timesteps):

    if(timestep != 0):
        delta_t.assign(compute_timestep(u)) # Compute adaptive time-step
    log("Timestep Number: ", timestep, " Timestep: ", float(delta_t))

    # Solve system - configured for solving non-linear systems, where everything is on the LHS (as above)
    # and the RHS == 0.
    solve(F_stokes==0, z, bcs=[bcvx,bcvy], solver_parameters=stokes_solver_parameters, nullspace=nullspace, appctx={'mu': mu})

    # Temperature system:
    solve(F_energy==0, T_new, bcs=[bct_base,bct_top], solver_parameters=temperature_solver_parameters)

    mu_field.interpolate(mu)
    epsii_field.interpolate(epsii)

    # Write output:
    # if period < 0 and timestep > 500:
    #     period = 10
    if period > 0 and (timestep % period == 0):
        u_file.write(u)
        p_file.write(p)
        t_file.write(T_new)
        mu_file.write(mu_field)
        epsii_file.write(epsii_field)

    # Compute diagnostics:
    domain_volume = assemble(1.*dx(domain=mesh))
    rms_velocity = numpy.sqrt(assemble(dot(u,u) * dx)) * numpy.sqrt(1./domain_volume)
    rms_velocity_surf = numpy.sqrt(assemble(u[0] ** 2 * ds(top_id)))
    max_velocity_surf = 0.

    n = FacetNormal(mesh)
    nusselt_number_base = assemble(dot(grad(T_new),n) * ds(bottom_id))
    nusselt_number_top  = assemble(dot(grad(T_new),n) * ds(top_id))
    energy_conservation_1 = abs(abs(nusselt_number_top) - abs(nusselt_number_base))

    average_temperature = assemble(T_new * dx) / domain_volume

    max_viscosity = mu_field.dat.data.max()
    min_viscosity = mu_field.dat.data.min()

#    rate_work_against_gravity = assemble(T_new * u[1] * dx) / domain_volume
#    rate_viscous_dissipation = assemble(inner(tau, epsii) * dx) / domain_volume
#    energy_conservation_2 = abs(rate_work_against_gravity - rate_viscous_dissipation / Ra.values()[0])

    log("RMS Velocity (domain, surf, surf_max): ", rms_velocity, rms_velocity_surf, max_velocity_surf)
    log("Nu (base, top, cons): ",nusselt_number_base,nusselt_number_top,energy_conservation_1)
    log("Average Temperature: ", average_temperature)
    log("mu (max, min)", max_viscosity, min_viscosity)
#    log("<W>, <phi>: ", rate_work_against_gravity, rate_viscous_dissipation)
#    log("<W>, <phi>, cons: ", rate_work_against_gravity, rate_viscous_dissipation, energy_conservation_2)

    # Calculate L2-norm of change in temperature:
    maxchange = math.sqrt(assemble((T_new - T_old)**2 * dx))

    # Leave if steady-state has been achieved:
    log("Maxchange:",maxchange," Targetting:",steady_state_tolerance)
    log()

    log_params(f, str(timestep))
    log_params(f, ' ' + str(maxchange))
    log_params(f, ' ' + str(rms_velocity))
    log_params(f, ' ' + str(rms_velocity_surf))
    log_params(f, ' ' + str(max_velocity_surf))
    log_params(f, ' ' + str(nusselt_number_base))
    log_params(f, ' ' + str(nusselt_number_top))
    log_params(f, ' ' + str(average_temperature))
    log_params(f, ' ' + str(max_viscosity))
    log_params(f, ' ' + str(min_viscosity))
#    log_params(f, ' ' + str(rate_work_against_gravity))
#    log_params(f, ' ' + str(rate_viscous_dissipation))
    log_params(f, '\n')
    if(maxchange < steady_state_tolerance):
        log("Steady-state achieved -- exiting time-step loop")
        break

    # Set T_old = T_new - assign the values of T_new to T_old
    T_old.assign(T_new)

f.close()