From fd35bf90d197e254616b0c7f8bccdd81b6266bf3 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Mon, 10 Jul 2023 18:35:44 +0100 Subject: [PATCH] serial vertical slice script --- .../serial/vertical_slice/slice_mountain.py | 265 ++++++++++++++++++ 1 file changed, 265 insertions(+) create mode 100644 examples/serial/vertical_slice/slice_mountain.py diff --git a/examples/serial/vertical_slice/slice_mountain.py b/examples/serial/vertical_slice/slice_mountain.py new file mode 100644 index 00000000..310ab9cf --- /dev/null +++ b/examples/serial/vertical_slice/slice_mountain.py @@ -0,0 +1,265 @@ +import firedrake as fd +from math import pi +from utils.diagnostics import convective_cfl_calculator +from utils.serial import SerialMiniApp +from utils.vertical_slice import hydrostatic_rho, \ + get_form_mass, get_form_function, maximum +from firedrake.petsc import PETSc + +PETSc.Sys.Print("Setting up problem") + +comm = fd.COMM_WORLD + +# set up the mesh + +nt = 5 +dt = 2.5 + +nlayers = 140 # horizontal layers +base_columns = 360 # number of columns +L = 144e3 +H = 35e3 # Height position of the model top + +distribution_parameters = { + "partition": True, + "overlap_type": (fd.DistributedMeshOverlapType.VERTEX, 2) +} + +# surface mesh of ground +base_mesh = fd.PeriodicIntervalMesh(base_columns, L, + distribution_parameters=distribution_parameters, + comm=comm) + +# volume mesh of the slice +mesh = fd.ExtrudedMesh(base_mesh, layers=nlayers, layer_height=H/nlayers) +n = fd.FacetNormal(mesh) + +g = fd.Constant(9.810616) +N = fd.Constant(0.01) # Brunt-Vaisala frequency (1/s) +cp = fd.Constant(1004.5) # SHC of dry air at const. pressure (J/kg/K) +R_d = fd.Constant(287.) # Gas constant for dry air (J/kg/K) +kappa = fd.Constant(2.0/7.0) # R_d/c_p +p_0 = fd.Constant(1000.0*100.0) # reference pressure (Pa, not hPa) +cv = fd.Constant(717.) # SHC of dry air at const. volume (J/kg/K) +T_0 = fd.Constant(273.15) # ref. temperature + +dT = fd.Constant(dt) + +# making a mountain out of a molehill +a = 10000. +xc = L/2. +x, z = fd.SpatialCoordinate(mesh) +hm = 1. +zs = hm*a**2/((x-xc)**2 + a**2) + +smooth_z = True +name = "mountain_nh" +if smooth_z: + name += '_smootherz' + zh = 5000. + xexpr = fd.as_vector([x, fd.conditional(z < zh, z + fd.cos(0.5*pi*z/zh)**6*zs, z)]) +else: + xexpr = fd.as_vector([x, z + ((H-z)/H)*zs]) +mesh.coordinates.interpolate(xexpr) + +horizontal_degree = 1 +vertical_degree = 1 + +S1 = fd.FiniteElement("CG", fd.interval, horizontal_degree+1) +S2 = fd.FiniteElement("DG", fd.interval, horizontal_degree) + +# vertical base spaces +T0 = fd.FiniteElement("CG", fd.interval, vertical_degree+1) +T1 = fd.FiniteElement("DG", fd.interval, vertical_degree) + +# build spaces V2, V3, Vt +V2h_elt = fd.HDiv(fd.TensorProductElement(S1, T1)) +V2t_elt = fd.TensorProductElement(S2, T0) +V3_elt = fd.TensorProductElement(S2, T1) +V2v_elt = fd.HDiv(V2t_elt) +V2_elt = V2h_elt + V2v_elt + +V1 = fd.FunctionSpace(mesh, V2_elt, name="Velocity") +V2 = fd.FunctionSpace(mesh, V3_elt, name="Pressure") +Vt = fd.FunctionSpace(mesh, V2t_elt, name="Temperature") +Vv = fd.FunctionSpace(mesh, V2v_elt, name="Vv") + +W = V1 * V2 * Vt # velocity, density, temperature + +PETSc.Sys.Print(f"DoFs: {W.dim()}") +PETSc.Sys.Print(f"DoFs/core: {W.dim()/comm.size}") + +Un = fd.Function(W) + +x, z = fd.SpatialCoordinate(mesh) + +# N^2 = (g/theta)dtheta/dz => dtheta/dz = theta N^2g => theta=theta_0exp(N^2gz) +Tsurf = fd.Constant(300.) +thetab = Tsurf*fd.exp(N**2*z/g) + +cp = fd.Constant(1004.5) # SHC of dry air at const. pressure (J/kg/K) +Up = fd.as_vector([fd.Constant(0.0), fd.Constant(1.0)]) # up direction + +un = Un.subfunctions[0] +rhon = Un.subfunctions[1] +thetan = Un.subfunctions[2] +un.project(fd.as_vector([10.0, 0.0])) +thetan.interpolate(thetab) +theta_back = fd.Function(Vt).assign(thetan) +rhon.assign(1.0e-5) + +PETSc.Sys.Print("Calculating hydrostatic state") + +Pi = fd.Function(V2) + +hydrostatic_rho(Vv, V2, mesh, thetan, rhon, pi_boundary=fd.Constant(0.02), + cp=cp, R_d=R_d, p_0=p_0, kappa=kappa, g=g, Up=Up, + top=True, Pi=Pi) +p0 = maximum(Pi) + +hydrostatic_rho(Vv, V2, mesh, thetan, rhon, pi_boundary=fd.Constant(0.05), + cp=cp, R_d=R_d, p_0=p_0, kappa=kappa, g=g, Up=Up, + top=True, Pi=Pi) +p1 = maximum(Pi) +alpha = 2.*(p1-p0) +beta = p1-alpha +pi_top = (1.-beta)/alpha + +hydrostatic_rho(Vv, V2, mesh, thetan, rhon, pi_boundary=fd.Constant(pi_top), + cp=cp, R_d=R_d, p_0=p_0, kappa=kappa, g=g, Up=Up, + top=True) + +rho_back = fd.Function(V2).assign(rhon) + +zc = H-10000. +mubar = 0.15/dt +mu_top = fd.conditional(z <= zc, 0.0, mubar*fd.sin((pi/2.)*(z-zc)/(H-zc))**2) +mu = fd.Function(V2).interpolate(mu_top/dT) + +form_function = get_form_function(n, Up, c_pen=2.0**(-7./2), + cp=cp, g=g, R_d=R_d, + p_0=p_0, kappa=kappa, mu=mu) + +form_mass = get_form_mass() + +zv = fd.as_vector([fd.Constant(0.), fd.Constant(0.)]) +bcs = [fd.DirichletBC(W.sub(0), zv, "bottom"), + fd.DirichletBC(W.sub(0), zv, "top")] + +for bc in bcs: + bc.apply(Un) + +# Parameters for the newton iterations +solver_parameters = { + "snes": { + "monitor": None, + "converged_reason": None, + "stol": 1e-12, + "atol": 1e-6, + "rtol": 1e-6, + "ksp_ew": None, + "ksp_ew_version": 1, + "ksp_ew_threshold": 1e-5, + "ksp_ew_rtol0": 1e-3, + }, + "ksp_type": "gmres", + "ksp": { + "monitor": None, + "converged_reason": None, + "atol": 1e-7, + }, + "pc_type": "python", + "pc_python_type": "firedrake.AssembledPC", + "assembled": { + "pc_type": "python", + "pc_python_type": "firedrake.ASMVankaPC", + "pc_vanka": { + "construct_dim": 0, + "sub_sub_pc_type": "lu", + "sub_sub_pc_factor_mat_solver_type": 'mumps', + }, + }, +} + +theta = 0.5 + +miniapp = SerialMiniApp(dt=dt, theta=theta, w_initial=Un, + form_mass=form_mass, + form_function=form_function, + solver_parameters=solver_parameters, + bcs=bcs) + +PETSc.Sys.Print("Solving problem") + +uout = fd.Function(V1, name='velocity') +thetaout = fd.Function(Vt, name='temperature') +rhoout = fd.Function(V2, name='density') + +ofile = fd.File('output/slice_mountain.pvd', + comm=comm) + + +def assign_out_functions(): + uout.assign(miniapp.w0.subfunctions[0]) + rhoout.assign(miniapp.w0.subfunctions[1]) + thetaout.assign(miniapp.w0.subfunctions[2]) + + rhoout.assign(rhoout - rho_back) + thetaout.assign(thetaout - theta_back) + + +def write_to_file(time): + ofile.write(uout, rhoout, thetaout, t=time) + + +assign_out_functions() +write_to_file(time=0) + +PETSc.Sys.Print('### === --- Timestepping loop --- === ###') +linear_its = 0 +nonlinear_its = 0 + +cfl_calc = convective_cfl_calculator(mesh) +cfl_series = [] + + +def max_cfl(u, dt): + with cfl_calc(u, dt).dat.vec_ro as v: + return v.max()[1] + + +def preproc(app, it, time): + PETSc.Sys.Print('') + PETSc.Sys.Print(f'### === --- Calculating time-step {it} --- === ###') + PETSc.Sys.Print('') + + +def postproc(app, it, time): + global linear_its + global nonlinear_its + + linear_its += miniapp.nlsolver.snes.getLinearSolveIterations() + nonlinear_its += miniapp.nlsolver.snes.getIterationNumber() + + assign_out_functions() + write_to_file(time=time) + PETSc.Sys.Print('') + + cfl = max_cfl(uout, dt) + cfl_series.append(cfl) + PETSc.Sys.Print(f'Time = {time}') + PETSc.Sys.Print(f'Maximum CFL = {cfl}') + + +# solve for each window +miniapp.solve(nt=nt, + preproc=preproc, + postproc=postproc) + +PETSc.Sys.Print('') +PETSc.Sys.Print('### === --- Iteration counts --- === ###') +PETSc.Sys.Print('') + +PETSc.Sys.Print(f'linear iterations: {linear_its} | iterations per timestep: {linear_its/nt}') +PETSc.Sys.Print(f'nonlinear iterations: {nonlinear_its} | iterations per timestep: {nonlinear_its/nt}') +PETSc.Sys.Print('')