def test_physics_recovery_kernels(boundary): m = IntervalMesh(3, 3) mesh = ExtrudedMesh(m, layers=3, layer_height=1.0) cell = m.ufl_cell().cellname() hori_elt = FiniteElement("DG", cell, 0) vert_elt = FiniteElement("CG", interval, 1) theta_elt = TensorProductElement(hori_elt, vert_elt) Vt = FunctionSpace(mesh, theta_elt) Vt_brok = FunctionSpace(mesh, BrokenElement(theta_elt)) initial_field = Function(Vt) true_field = Function(Vt_brok) new_field = Function(Vt_brok) initial_field, true_field = setup_values(boundary, initial_field, true_field) kernel = kernels.PhysicsRecoveryTop( ) if boundary == "top" else kernels.PhysicsRecoveryBottom() kernel.apply(new_field, initial_field) tolerance = 1e-12 index = 11 if boundary == "top" else 6 assert abs(true_field.dat.data[index] - new_field.dat.data[index]) < tolerance, \ "Value at %s from physics recovery is not correct" % boundary
def mesh(geometry, element): Lx = 100. Ly = 100. H = 100. deltax = Lx / 5. deltay = Ly / 5. deltaz = H / 5. ncolumnsy = int(Ly/deltay) ncolumnsx = int(Lx/deltax) nlayers = int(H/deltaz) quadrilateral = True if element == "quadrilateral" else False if geometry == "periodic-in-both": m = PeriodicRectangleMesh(ncolumnsx, ncolumnsy, Lx, Ly, direction='both', quadrilateral=quadrilateral) elif geometry == "periodic-in-x": m = PeriodicRectangleMesh(ncolumnsx, ncolumnsy, Lx, Ly, direction='x', quadrilateral=quadrilateral) elif geometry == "periodic-in-y": m = PeriodicRectangleMesh(ncolumnsx, ncolumnsy, Lx, Ly, direction='y', quadrilateral=quadrilateral) elif geometry == "non-periodic": m = RectangleMesh(ncolumnsx, ncolumnsy, Lx, Ly, quadrilateral=quadrilateral) extruded_mesh = ExtrudedMesh(m, layers=nlayers, layer_height=deltaz) return extruded_mesh
def run(problem, tensor, size_factor, degree): formname = problem.__name__ cellname = 'cube' if tensor else 'simplex' PETSc.Sys.Print("%s: %s, degree=%d" % (formname, cellname, degree)) num_cells = COMM_WORLD.size * max(1, 1e7 * size_factor / (degree + 1)**{'spectral': 4, 'coffee': 6}[args.mode]) h = int(floor(cbrt(num_cells / COMM_WORLD.size))) w = int(floor(sqrt(num_cells / h))) d = int(round(num_cells / (w * h))) num_cells = w * d * h if tensor: mesh = ExtrudedMesh(UnitSquareMesh(w, d, quadrilateral=True), h) else: mesh = UnitCubeMesh(w, d, h) comm = mesh.comm J = problem(mesh, int(degree)) # Warmup and allocate A = assemble(J, mat_type="matfree", form_compiler_parameters={'mode': args.mode}) A.force_evaluation() Ap = A.petscmat x, y = Ap.createVecs() assert x.size == y.size num_dofs = x.size Ap.mult(x, y) stage = PETSc.Log.Stage("%s(%d) %s" % (formname, degree, cellname)) with stage: assemble(J, mat_type="matfree", form_compiler_parameters={'mode': args.mode}, tensor=A) A.force_evaluation() Ap = A.petscmat for _ in range(num_matvecs): Ap.mult(x, y) parloop = parloop_event.getPerfInfo() matmult = matmult_event.getPerfInfo() assert parloop["count"] == 2 * num_matvecs assert matmult["count"] == num_matvecs parloop_time = comm.allreduce(parloop["time"], op=MPI.SUM) / (comm.size * num_matvecs) matmult_time = comm.allreduce(matmult["time"], op=MPI.SUM) / (comm.size * num_matvecs) matfree_overhead = (1 - parloop_time / matmult_time) PETSc.Sys.Print("Matrix-free action overhead: %.1f%%" % (matfree_overhead * 100,)) if COMM_WORLD.rank == 0: header = not os.path.exists(filepath) data = {"num_cells": num_cells, "num_dofs": num_dofs, "num_procs": comm.size, "tsfc_mode": args.mode, "problem": formname, "cell_type": cellname, "degree": degree, "matmult_time": matmult_time, "parloop_time": parloop_time} df = pandas.DataFrame(data, index=[0]) df.to_csv(filepath, index=False, mode='a', header=header)
def setup_balance(dirname): # set up grid and time stepping parameters dt = 1. tmax = 5. deltax = 400 L = 2000. H = 10000. nlayers = int(H / deltax) ncolumns = int(L / deltax) m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers) output = OutputParameters(dirname=dirname + '/dry_balance', dumpfreq=10, dumplist=['u']) parameters = CompressibleParameters() state = State(mesh, dt=dt, output=output, parameters=parameters) eqns = CompressibleEulerEquations(state, "CG", 1) # Initial conditions rho0 = state.fields("rho") theta0 = state.fields("theta") # Isentropic background state Tsurf = Constant(300.) theta0.interpolate(Tsurf) # Calculate hydrostatic exner compressible_hydrostatic_balance(state, theta0, rho0, solve_for_rho=True) state.set_reference_profiles([('rho', rho0), ('theta', theta0)]) # Set up transport schemes transported_fields = [ ImplicitMidpoint(state, "u"), SSPRK3(state, "rho"), SSPRK3(state, "theta", options=EmbeddedDGOptions()) ] # Set up linear solver linear_solver = CompressibleSolver(state, eqns) # build time stepper stepper = CrankNicolson(state, eqns, transported_fields, linear_solver=linear_solver) return stepper, tmax
def setup_IPdiffusion(dirname, vector, DG): dt = 0.01 L = 10. m = PeriodicIntervalMesh(50, L) mesh = ExtrudedMesh(m, layers=50, layer_height=0.2) fieldlist = ['u', 'D'] timestepping = TimesteppingParameters(dt=dt) parameters = CompressibleParameters() output = OutputParameters(dirname=dirname) state = State(mesh, vertical_degree=1, horizontal_degree=1, family="CG", timestepping=timestepping, parameters=parameters, output=output, fieldlist=fieldlist) x = SpatialCoordinate(mesh) if vector: kappa = Constant([[0.05, 0.], [0., 0.05]]) if DG: Space = VectorFunctionSpace(mesh, "DG", 1) else: Space = state.spaces("HDiv") fexpr = as_vector([exp(-(L / 2. - x[0])**2 - (L / 2. - x[1])**2), 0.]) else: kappa = 0.05 if DG: Space = state.spaces("DG") else: Space = state.spaces("HDiv_v") fexpr = exp(-(L / 2. - x[0])**2 - (L / 2. - x[1])**2) f = state.fields("f", Space) try: f.interpolate(fexpr) except NotImplementedError: # finat elements raise NotImplementedError if they don't # advertise a dual for interpolation. f.project(fexpr) mu = 5. f_diffusion = InteriorPenalty(state, f.function_space(), kappa=kappa, mu=mu) diffused_fields = [("f", f_diffusion)] stepper = AdvectionDiffusion(state, diffused_fields=diffused_fields) return stepper
def tracer_blob_slice(tmpdir, degree): dt = 0.01 L = 10. m = PeriodicIntervalMesh(10, L) mesh = ExtrudedMesh(m, layers=10, layer_height=1.) output = OutputParameters(dirname=str(tmpdir), dumpfreq=25) state = State(mesh, dt=dt, output=output) tmax = 1. x = SpatialCoordinate(mesh) f_init = exp(-((x[0] - 0.5 * L)**2 + (x[1] - 0.5 * L)**2)) return TracerSetup(state, tmax, f_init, family="CG", degree=degree)
def setup_gw(dirname): nlayers = 10 # horizontal layers columns = 30 # number of columns L = 1.e5 m = PeriodicRectangleMesh(columns, 1, L, 1.e4, quadrilateral=True) dt = 6.0 # build volume mesh H = 1.0e4 # Height position of the model top mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers) fieldlist = ['u', 'p', 'b'] timestepping = TimesteppingParameters(dt=dt) output = OutputParameters(dirname=dirname + "/gw_incompressible", dumplist=['u'], dumpfreq=5) parameters = CompressibleParameters() state = State(mesh, vertical_degree=1, horizontal_degree=1, family="RTCF", timestepping=timestepping, output=output, parameters=parameters, fieldlist=fieldlist) # Initial conditions u0 = state.fields("u") p0 = state.fields("p") b0 = state.fields("b") # z.grad(bref) = N**2 x, y, z = SpatialCoordinate(mesh) N = parameters.N bref = z * (N**2) b_b = Function(b0.function_space()).interpolate(bref) b0.interpolate(b_b) incompressible_hydrostatic_balance(state, b0, p0) state.initialise([('u', u0), ('p', p0), ('b', b0)]) # Set up forcing forcing = IncompressibleForcing(state) return state, forcing
def setup_fallout(dirname): # declare grid shape, with length L and height H L = 10. H = 10. nlayers = 10 ncolumns = 10 # make mesh m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers)) x = SpatialCoordinate(mesh) dt = 0.1 output = OutputParameters(dirname=dirname + "/fallout", dumpfreq=10, dumplist=['rain_mixing_ratio']) parameters = CompressibleParameters() diagnostic_fields = [Precipitation()] state = State(mesh, dt=dt, output=output, parameters=parameters, diagnostic_fields=diagnostic_fields) Vrho = state.spaces("DG1_equispaced") problem = [(AdvectionEquation(state, Vrho, "rho", ufamily="CG", udegree=1), ForwardEuler(state))] state.fields("rho").assign(1.) physics_list = [Fallout(state)] rain0 = state.fields("rain_mixing_ratio") # set up rain xc = L / 2 zc = H / 2 rc = H / 4 r = sqrt((x[0] - xc)**2 + (x[1] - zc)**2) rain_expr = conditional(r > rc, 0., 1e-3 * (cos(pi * r / (rc * 2)))**2) rain0.interpolate(rain_expr) # build time stepper stepper = PrescribedTransport(state, problem, physics_list=physics_list) return stepper, 10.0
def mesh(geometry): L = 100. H = 100. deltax = L / 5. deltaz = H / 5. nlayers = int(H / deltaz) ncolumns = int(L / deltax) if geometry == "periodic": m = PeriodicIntervalMesh(ncolumns, L) elif geometry == "non-periodic": m = IntervalMesh(ncolumns, L) extruded_mesh = ExtrudedMesh(m, layers=nlayers, layer_height=deltaz) return extruded_mesh
def state(tmpdir, geometry): """ returns an instance of the State class, having set up either spherical geometry or 2D vertical slice geometry """ output = OutputParameters(dirname=str(tmpdir), dumplist=["f"], dumpfreq=15) if geometry == "sphere": mesh = IcosahedralSphereMesh(radius=1, refinement_level=3, degree=1) x = SpatialCoordinate(mesh) mesh.init_cell_orientations(x) family = "BDM" vertical_degree = None fieldlist = ["u", "D"] dt = pi / 3. * 0.01 uexpr = as_vector([-x[1], x[0], 0.0]) if geometry == "slice": m = PeriodicIntervalMesh(15, 1.) mesh = ExtrudedMesh(m, layers=15, layer_height=1. / 15.) family = "CG" vertical_degree = 1 fieldlist = ["u", "rho", "theta"] dt = 0.01 x = SpatialCoordinate(mesh) uexpr = as_vector([1.0, 0.0]) timestepping = TimesteppingParameters(dt=dt) state = State(mesh, vertical_degree=vertical_degree, horizontal_degree=1, family=family, timestepping=timestepping, output=output, fieldlist=fieldlist) u0 = state.fields("u") u0.project(uexpr) return state
def setup_checkpointing(dirname): nlayers = 5 # horizontal layers columns = 15 # number of columns L = 3.e5 m = PeriodicIntervalMesh(columns, L) dt = 2.0 # build volume mesh H = 1.0e4 # Height position of the model top mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers) output = OutputParameters(dirname=dirname, dumpfreq=1, chkptfreq=2, log_level='INFO') parameters = CompressibleParameters() state = State(mesh, dt=dt, output=output, parameters=parameters) eqns = CompressibleEulerEquations(state, "CG", 1) # Set up transport schemes transported_fields = [] transported_fields.append(SSPRK3(state, "u")) transported_fields.append(SSPRK3(state, "rho")) transported_fields.append(SSPRK3(state, "theta")) # Set up linear solver linear_solver = CompressibleSolver(state, eqns) # build time stepper stepper = CrankNicolson(state, eqns, transported_fields, linear_solver=linear_solver) return state, stepper, dt
def tracer_slice(tmpdir, degree): n = 30 if degree == 0 else 15 m = PeriodicIntervalMesh(n, 1.) mesh = ExtrudedMesh(m, layers=n, layer_height=1. / n) # Parameters chosen so that dt != 1 and u != 1 # Gaussian is translated by 1.5 times width of domain to demonstrate # that transport is working correctly dt = 0.01 tmax = 0.75 output = OutputParameters(dirname=str(tmpdir), dumpfreq=25) state = State(mesh, dt=dt, output=output) uexpr = as_vector([2.0, 0.0]) x = SpatialCoordinate(mesh) width = 1. / 10. f0 = 0.5 fmax = 2.0 xc_init = 0.25 xc_end = 0.75 r_init = sqrt((x[0] - xc_init)**2 + (x[1] - 0.5)**2) r_end = sqrt((x[0] - xc_end)**2 + (x[1] - 0.5)**2) f_init = f0 + (fmax - f0) * exp(-(r_init / width)**2) f_end = f0 + (fmax - f0) * exp(-(r_end / width)**2) tol = 0.12 return TracerSetup(state, tmax, f_init, f_end, "CG", degree, uexpr, tol=tol)
def setup_fallout(dirname): # declare grid shape, with length L and height H L = 10. H = 10. nlayers = 10 ncolumns = 10 # make mesh m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers)) x = SpatialCoordinate(mesh) fieldlist = ['u', 'rho', 'theta', 'rain'] timestepping = TimesteppingParameters(dt=0.1, maxk=4, maxi=1) output = OutputParameters(dirname=dirname + "/fallout", dumpfreq=10, dumplist=['rain']) parameters = CompressibleParameters() diagnostic_fields = [Precipitation()] state = State(mesh, vertical_degree=1, horizontal_degree=1, family="CG", timestepping=timestepping, output=output, parameters=parameters, fieldlist=fieldlist, diagnostic_fields=diagnostic_fields) # declare initial fields u0 = state.fields("u") rho0 = state.fields("rho") theta0 = state.fields("theta") # spaces Vt = theta0.function_space() # declare tracer field and a background field rain0 = state.fields("rain", Vt) # set up rain xc = L / 2 zc = H / 2 rc = H / 4 r = sqrt((x[0] - xc)**2 + (x[1] - zc)**2) rain_expr = conditional(r > rc, 0., 1e-3 * (cos(pi * r / (rc * 2)))**2) rain0.interpolate(rain_expr) rho0.assign(1.0) state.initialise([('u', u0), ('rho', rho0), ('rain', rain0)]) # build advection dictionary advected_fields = [] advected_fields.append(("u", NoAdvection(state, u0, None))) advected_fields.append(("rho", NoAdvection(state, rho0, None))) advected_fields.append(("rain", NoAdvection(state, rain0, None))) physics_list = [Fallout(state)] # build time stepper stepper = AdvectionDiffusion(state, advected_fields, physics_list=physics_list) return stepper, 10.0
def test_limit_midpoints(profile): # ------------------------------------------------------------------------ # # Set up meshes and spaces # ------------------------------------------------------------------------ # m = IntervalMesh(3, 3) mesh = ExtrudedMesh(m, layers=1, layer_height=3.0) cell = m.ufl_cell().cellname() DG_hori_elt = FiniteElement("DG", cell, 1, variant='equispaced') DG_vert_elt = FiniteElement("DG", interval, 1, variant='equispaced') Vt_vert_elt = FiniteElement("CG", interval, 2) DG_elt = TensorProductElement(DG_hori_elt, DG_vert_elt) theta_elt = TensorProductElement(DG_hori_elt, Vt_vert_elt) Vt_brok = FunctionSpace(mesh, BrokenElement(theta_elt)) DG1 = FunctionSpace(mesh, DG_elt) new_field = Function(Vt_brok) init_field = Function(Vt_brok) DG1_field = Function(DG1) # ------------------------------------------------------------------------ # # Initial conditions # ------------------------------------------------------------------------ # _, z = SpatialCoordinate(mesh) if profile == 'undershoot': # A quadratic whose midpoint is lower than the top and bottom values init_expr = (80. / 9.) * z**2 - 20. * z + 300. elif profile == 'overshoot': # A quadratic whose midpoint is higher than the top and bottom values init_expr = (-80. / 9.) * z**2 + (100. / 3) * z + 300. elif profile == 'linear': # Linear profile which must be unchanged init_expr = (20. / 3.) * z + 300. else: raise NotImplementedError # Linear DG field has the same values at top and bottom as quadratic DG_expr = (20. / 3.) * z + 300. init_field.interpolate(init_expr) DG1_field.interpolate(DG_expr) # ------------------------------------------------------------------------ # # Apply kernel # ------------------------------------------------------------------------ # kernel = kernels.LimitMidpoints(Vt_brok) kernel.apply(new_field, DG1_field, init_field) # ------------------------------------------------------------------------ # # Check values # ------------------------------------------------------------------------ # tol = 1e-12 assert np.max(new_field.dat.data) <= np.max(init_field.dat.data) + tol, \ 'LimitMidpoints kernel is giving an overshoot' assert np.min(new_field.dat.data) >= np.min(init_field.dat.data) - tol, \ 'LimitMidpoints kernel is giving an undershoot' if profile == 'linear': assert np.allclose(init_field.dat.data, new_field.dat.data), \ 'For a profile with no maxima or minima, the LimitMidpoints ' + \ 'kernel should leave the field unchanged'
def setup_3d_recovery(dirname): L = 100. H = 10. W = 1. deltax = L / 5. deltay = W / 5. deltaz = H / 5. nlayers = int(H / deltaz) ncolumnsx = int(L / deltax) ncolumnsy = int(W / deltay) m = RectangleMesh(ncolumnsx, ncolumnsy, L, W, quadrilateral=True) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers) x, y, z = SpatialCoordinate(mesh) # horizontal base spaces cell = mesh._base_mesh.ufl_cell().cellname() u_hori = FiniteElement("RTCF", cell, 1) w_hori = FiniteElement("DG", cell, 0) # vertical base spaces u_vert = FiniteElement("DG", interval, 0) w_vert = FiniteElement("CG", interval, 1) # build elements u_element = HDiv(TensorProductElement(u_hori, u_vert)) w_element = HDiv(TensorProductElement(w_hori, w_vert)) theta_element = TensorProductElement(w_hori, w_vert) v_element = u_element + w_element # spaces VDG0 = FunctionSpace(mesh, "DG", 0) VCG1 = FunctionSpace(mesh, "CG", 1) VDG1 = FunctionSpace(mesh, "DG", 1) Vt = FunctionSpace(mesh, theta_element) Vt_brok = FunctionSpace(mesh, BrokenElement(theta_element)) Vu = FunctionSpace(mesh, v_element) VuCG1 = VectorFunctionSpace(mesh, "CG", 1) VuDG1 = VectorFunctionSpace(mesh, "DG", 1) # set up initial conditions np.random.seed(0) expr = np.random.randn( ) + np.random.randn() * x + np.random.randn() * y + np.random.randn( ) * z + np.random.randn() * x * y + np.random.randn( ) * x * z + np.random.randn() * y * z + np.random.randn() * x * y * z # our actual theta and rho and v rho_CG1_true = Function(VCG1).interpolate(expr) theta_CG1_true = Function(VCG1).interpolate(expr) v_CG1_true = Function(VuCG1).interpolate(as_vector([expr, expr, expr])) rho_Vt_true = Function(Vt).interpolate(expr) # make the initial fields by projecting expressions into the lowest order spaces rho_DG0 = Function(VDG0).interpolate(expr) rho_CG1 = Function(VCG1) theta_Vt = Function(Vt).interpolate(expr) theta_CG1 = Function(VCG1) v_Vu = Function(Vu).project(as_vector([expr, expr, expr])) v_CG1 = Function(VuCG1) rho_Vt = Function(Vt) # make the recoverers and do the recovery rho_recoverer = Recoverer(rho_DG0, rho_CG1, VDG=VDG1, boundary_method=Boundary_Method.dynamics) theta_recoverer = Recoverer(theta_Vt, theta_CG1, VDG=VDG1, boundary_method=Boundary_Method.dynamics) v_recoverer = Recoverer(v_Vu, v_CG1, VDG=VuDG1, boundary_method=Boundary_Method.dynamics) rho_Vt_recoverer = Recoverer(rho_DG0, rho_Vt, VDG=Vt_brok, boundary_method=Boundary_Method.physics) rho_recoverer.project() theta_recoverer.project() v_recoverer.project() rho_Vt_recoverer.project() rho_diff = errornorm(rho_CG1, rho_CG1_true) / norm(rho_CG1_true) theta_diff = errornorm(theta_CG1, theta_CG1_true) / norm(theta_CG1_true) v_diff = errornorm(v_CG1, v_CG1_true) / norm(v_CG1_true) rho_Vt_diff = errornorm(rho_Vt, rho_Vt_true) / norm(rho_Vt_true) return (rho_diff, theta_diff, v_diff, rho_Vt_diff)
def setup_sk(dirname): nlayers = 10 # horizontal layers columns = 30 # number of columns L = 1.e5 m = PeriodicIntervalMesh(columns, L) dt = 6.0 # build volume mesh H = 1.0e4 # Height position of the model top mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers) # Set up points for output at the centre of the domain, edges and corners. # The point at x=L*(13.0/30) is in the halo region for a two-way MPI decomposition points_x = [0.0, L * (13.0 / 30), L / 2.0, L] points_z = [0.0, H / 2.0, H] points = np.array([p for p in itertools.product(points_x, points_z)]) fieldlist = ['u', 'rho', 'theta'] timestepping = TimesteppingParameters(dt=dt) output = OutputParameters(dirname=dirname + "/sk_nonlinear", dumplist=['u'], dumpfreq=5, log_level=INFO, point_data=[('rho', points), ('u', points)]) parameters = CompressibleParameters() diagnostic_fields = [CourantNumber()] state = State(mesh, vertical_degree=1, horizontal_degree=1, family="CG", timestepping=timestepping, output=output, parameters=parameters, fieldlist=fieldlist, diagnostic_fields=diagnostic_fields) # Initial conditions u0 = state.fields("u") rho0 = state.fields("rho") theta0 = state.fields("theta") # spaces Vu = u0.function_space() Vt = theta0.function_space() Vr = rho0.function_space() # Thermodynamic constants required for setting initial conditions # and reference profiles g = parameters.g N = parameters.N # N^2 = (g/theta)dtheta/dz => dtheta/dz = theta N^2g => theta=theta_0exp(N^2gz) x, z = SpatialCoordinate(mesh) Tsurf = 300. thetab = Tsurf * exp(N**2 * z / g) theta_b = Function(Vt).interpolate(thetab) rho_b = Function(Vr) # Calculate hydrostatic Pi compressible_hydrostatic_balance(state, theta_b, rho_b) a = 5.0e3 deltaTheta = 1.0e-2 theta_pert = deltaTheta * sin(np.pi * z / H) / (1 + (x - L / 2)**2 / a**2) theta0.interpolate(theta_b + theta_pert) rho0.assign(rho_b) u0.project(as_vector([20.0, 0.0])) state.initialise([('u', u0), ('rho', rho0), ('theta', theta0)]) state.set_reference_profiles([('rho', rho_b), ('theta', theta_b)]) # Set up advection schemes ueqn = EulerPoincare(state, Vu) rhoeqn = AdvectionEquation(state, Vr, equation_form="continuity") thetaeqn = SUPGAdvection(state, Vt) advected_fields = [] advected_fields.append(("u", ThetaMethod(state, u0, ueqn))) advected_fields.append(("rho", SSPRK3(state, rho0, rhoeqn))) advected_fields.append(("theta", SSPRK3(state, theta0, thetaeqn))) # Set up linear solver linear_solver = CompressibleSolver(state) # Set up forcing compressible_forcing = CompressibleForcing(state) # build time stepper stepper = CrankNicolson(state, advected_fields, linear_solver, compressible_forcing) return stepper, 2 * dt
def run_profliler(hybridization, model_degree, model_family, mesh_degree, cfl, refinements, layers, debug, rtol, flexsolver=True, stronger_smoother=False, suppress_data_output=False): nlayers = layers # Number of vertical layers refinements = refinements # Number of horiz. cells = 20*(4^refinements) hybrid = bool(hybridization) # Set up problem parameters parameters = CompressibleParameters() a_ref = 6.37122e6 # Radius of the Earth (m) X = 1.0 # Reduced-size Earth reduction factor a = a_ref / X # Scaled radius of planet (m) g = parameters.g # Acceleration due to gravity (m/s^2) N = parameters.N # Brunt-Vaisala frequency (1/s) p_0 = parameters.p_0 # Reference pressure (Pa, not hPa) c_p = parameters.cp # SHC of dry air at const. pressure (J/kg/K) R_d = parameters.R_d # Gas constant for dry air (J/kg/K) kappa = parameters.kappa # R_d/c_p T_eq = 300.0 # Isothermal atmospheric temperature (K) p_eq = 1000.0 * 100.0 # Reference surface pressure at the equator u_0 = 20.0 # Maximum amplitude of the zonal wind (m/s) d = 5000.0 # Width parameter for Theta' lamda_c = 2.0 * np.pi / 3.0 # Longitudinal centerpoint of Theta' phi_c = 0.0 # Lat. centerpoint of Theta' (equator) deltaTheta = 1.0 # Maximum amplitude of Theta' (K) L_z = 20000.0 # Vert. wave length of the Theta' perturb. gamma = (1 - kappa) / kappa cs = sqrt(c_p * T_eq / gamma) # Speed of sound in an air parcel if model_family == "RTCF": # Cubed-sphere mesh m = CubedSphereMesh(radius=a, refinement_level=refinements, degree=mesh_degree) elif model_family == "RT" or model_family == "BDFM": m = IcosahedralSphereMesh(radius=a, refinement_level=refinements, degree=mesh_degree) else: raise ValueError("Unknown family: %s" % model_family) cell_vs = interpolate(CellVolume(m), FunctionSpace(m, "DG", 0)) a_max = fmax(cell_vs) dx_max = sqrt(a_max) u_max = u_0 PETSc.Sys.Print("\nDetermining Dt from specified horizontal CFL: %s" % cfl) dt = int(cfl * (dx_max / cs)) # Height position of the model top (m) z_top = 1.0e4 deltaz = z_top / nlayers vertical_cfl = dt * (cs / deltaz) PETSc.Sys.Print(""" Problem parameters:\n Profiling linear solver for the compressible Euler equations.\n Speed of sound in compressible atmosphere: %s,\n Hybridized compressible solver: %s,\n Model degree: %s,\n Model discretization: %s,\n Mesh degree: %s,\n Horizontal refinements: %s,\n Vertical layers: %s,\n Dx (max, m): %s,\n Dz (m): %s,\n Dt (s): %s,\n horizontal CFL: %s,\n vertical CFL: %s. """ % (cs, hybrid, model_degree, model_family, mesh_degree, refinements, nlayers, dx_max, deltaz, dt, cfl, vertical_cfl)) # Build volume mesh mesh = ExtrudedMesh(m, layers=nlayers, layer_height=deltaz, extrusion_type="radial") x = SpatialCoordinate(mesh) # Create polar coordinates: # Since we use a CG1 field, this is constant on layers W_Q1 = FunctionSpace(mesh, "CG", 1) z_expr = sqrt(x[0] * x[0] + x[1] * x[1] + x[2] * x[2]) - a z = Function(W_Q1).interpolate(z_expr) lat_expr = asin(x[2] / sqrt(x[0] * x[0] + x[1] * x[1] + x[2] * x[2])) lat = Function(W_Q1).interpolate(lat_expr) lon = Function(W_Q1).interpolate(atan_2(x[1], x[0])) fieldlist = ['u', 'rho', 'theta'] timestepping = TimesteppingParameters(dt=dt, maxk=1, maxi=1) dirname = 'meanflow_ref' if hybrid: dirname += '_hybridization' # No output output = OutputParameters(dumpfreq=3600, dirname=dirname, perturbation_fields=['theta', 'rho'], dump_vtus=False, dump_diagnostics=False, checkpoint=False, log_level='INFO') diagnostics = Diagnostics(*fieldlist) state = State(mesh, vertical_degree=model_degree, horizontal_degree=model_degree, family=model_family, timestepping=timestepping, output=output, parameters=parameters, diagnostics=diagnostics, fieldlist=fieldlist) # Initial conditions u0 = state.fields.u theta0 = state.fields.theta rho0 = state.fields.rho # spaces Vu = u0.function_space() Vt = theta0.function_space() Vr = rho0.function_space() x = SpatialCoordinate(mesh) # Random velocity field CG2 = VectorFunctionSpace(mesh, "CG", 2) urand = Function(CG2) urand.dat.data[:] += np.random.randn(*urand.dat.data.shape) u0.project(urand) # Surface temperature G = g**2 / (N**2 * c_p) Ts_expr = G + (T_eq - G) * exp(-(u_max * N**2 / (4 * g * g)) * u_max * (cos(2.0 * lat) - 1.0)) Ts = Function(W_Q1).interpolate(Ts_expr) # Surface pressure ps_expr = p_eq * exp((u_max / (4.0 * G * R_d)) * u_max * (cos(2.0 * lat) - 1.0)) * (Ts / T_eq)**(1.0 / kappa) ps = Function(W_Q1).interpolate(ps_expr) # Background pressure p_expr = ps * (1 + G / Ts * (exp(-N**2 * z / g) - 1))**(1.0 / kappa) p = Function(W_Q1).interpolate(p_expr) # Background temperature Tb_expr = G * (1 - exp(N**2 * z / g)) + Ts * exp(N**2 * z / g) Tb = Function(W_Q1).interpolate(Tb_expr) # Background potential temperature thetab_expr = Tb * (p_0 / p)**kappa thetab = Function(W_Q1).interpolate(thetab_expr) theta_b = Function(theta0.function_space()).interpolate(thetab) rho_b = Function(rho0.function_space()) sin_tmp = sin(lat) * sin(phi_c) cos_tmp = cos(lat) * cos(phi_c) r = a * acos(sin_tmp + cos_tmp * cos(lon - lamda_c)) s = (d**2) / (d**2 + r**2) theta_pert = deltaTheta * s * sin(2 * np.pi * z / L_z) theta0.interpolate(theta_pert) # Compute the balanced density PETSc.Sys.Print("Computing balanced density field...\n") # Use vert. hybridization preconditioner for initialization pi_params = { 'ksp_type': 'preonly', 'pc_type': 'python', 'mat_type': 'matfree', 'pc_python_type': 'gusto.VerticalHybridizationPC', 'vert_hybridization': { 'ksp_type': 'gmres', 'pc_type': 'gamg', 'pc_gamg_sym_graph': True, 'ksp_rtol': 1e-12, 'ksp_atol': 1e-12, 'mg_levels': { 'ksp_type': 'richardson', 'ksp_max_it': 3, 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu' } } } if debug: pi_params['vert_hybridization']['ksp_monitor_true_residual'] = None compressible_hydrostatic_balance(state, theta_b, rho_b, top=False, pi_boundary=(p / p_0)**kappa, solve_for_rho=False, params=pi_params) # Random potential temperature perturbation theta0.assign(0.0) theta0.dat.data[:] += np.random.randn(len(theta0.dat.data)) # Random density field rho0.assign(0.0) rho0.dat.data[:] += np.random.randn(len(rho0.dat.data)) state.initialise([('u', u0), ('rho', rho0), ('theta', theta0)]) state.set_reference_profiles([('rho', rho_b), ('theta', theta_b)]) # Set up advection schemes ueqn = EulerPoincare(state, Vu) rhoeqn = AdvectionEquation(state, Vr, equation_form="continuity") thetaeqn = SUPGAdvection(state, Vt, equation_form="advective") advected_fields = [] advected_fields.append(("u", ThetaMethod(state, u0, ueqn))) advected_fields.append(("rho", SSPRK3(state, rho0, rhoeqn, subcycles=2))) advected_fields.append( ("theta", SSPRK3(state, theta0, thetaeqn, subcycles=2))) # Set up linear solver if hybrid: outer_solver_type = "Hybrid_SCPC" PETSc.Sys.Print(""" Setting up hybridized solver on the traces.""") if flexsolver: inner_solver_type = "fgmres_gamg_gmres_smoother" inner_parameters = { 'ksp_type': 'fgmres', 'ksp_rtol': rtol, 'ksp_max_it': 500, 'ksp_gmres_restart': 30, 'pc_type': 'gamg', 'pc_gamg_sym_graph': None, 'mg_levels': { 'ksp_type': 'gmres', 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu', 'ksp_max_it': 3 } } else: inner_solver_type = "fgmres_ml_richardson" inner_parameters = { 'ksp_type': 'fgmres', 'ksp_rtol': rtol, 'ksp_max_it': 500, 'ksp_gmres_restart': 30, 'pc_type': 'ml', 'pc_mg_cycles': 1, 'pc_ml_maxNlevels': 25, 'mg_levels': { 'ksp_type': 'richardson', 'ksp_richardson_scale': 0.8, 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu', 'ksp_max_it': 3 } } if stronger_smoother: inner_parameters['mg_levels']['ksp_max_it'] = 5 inner_solver_type += "_stronger" if debug: PETSc.Sys.Print("""Debugging on.""") inner_parameters['ksp_monitor_true_residual'] = None PETSc.Sys.Print("Inner solver: %s" % inner_solver_type) # Use Firedrake static condensation interface solver_parameters = { 'mat_type': 'matfree', 'pmat_type': 'matfree', 'ksp_type': 'preonly', 'pc_type': 'python', 'pc_python_type': 'firedrake.SCPC', 'pc_sc_eliminate_fields': '0, 1', 'condensed_field': inner_parameters } linear_solver = HybridizedCompressibleSolver( state, solver_parameters=solver_parameters, overwrite_solver_parameters=True) else: outer_solver_type = "gmres_SchurPC" PETSc.Sys.Print(""" Setting up GCR fieldsplit solver with Schur complement PC.""") solver_parameters = { 'pc_type': 'fieldsplit', 'pc_fieldsplit_type': 'schur', 'ksp_type': 'fgmres', 'ksp_max_it': 100, 'ksp_rtol': rtol, 'pc_fieldsplit_schur_fact_type': 'FULL', 'pc_fieldsplit_schur_precondition': 'selfp', 'fieldsplit_0': { 'ksp_type': 'preonly', 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu' }, 'fieldsplit_1': { 'ksp_type': 'preonly', 'ksp_max_it': 30, 'ksp_monitor_true_residual': None, 'pc_type': 'hypre', 'pc_hypre_type': 'boomeramg', 'pc_hypre_boomeramg_max_iter': 1, 'pc_hypre_boomeramg_agg_nl': 0, 'pc_hypre_boomeramg_coarsen_type': 'Falgout', 'pc_hypre_boomeramg_smooth_type': 'Euclid', 'pc_hypre_boomeramg_eu_bj': 1, 'pc_hypre_boomeramg_interptype': 'classical', 'pc_hypre_boomeramg_P_max': 0, 'pc_hypre_boomeramg_agg_nl': 0, 'pc_hypre_boomeramg_strong_threshold': 0.25, 'pc_hypre_boomeramg_max_levels': 25, 'pc_hypre_boomeramg_no_CF': False } } inner_solver_type = "hypre" if debug: solver_parameters['ksp_monitor_true_residual'] = None linear_solver = CompressibleSolver(state, solver_parameters=solver_parameters, overwrite_solver_parameters=True) # Set up forcing compressible_forcing = CompressibleForcing(state) param_info = ParameterInfo(dt=dt, deltax=dx_max, deltaz=deltaz, horizontal_courant=cfl, vertical_courant=vertical_cfl, family=model_family, model_degree=model_degree, mesh_degree=mesh_degree, solver_type=outer_solver_type, inner_solver_type=inner_solver_type) # Build profiler profiler = Profiler(parameterinfo=param_info, state=state, advected_fields=advected_fields, linear_solver=linear_solver, forcing=compressible_forcing, suppress_data_output=suppress_data_output) PETSc.Sys.Print("Starting profiler...\n") profiler.run(t=0, tmax=dt)
def setup_condens(dirname): # declare grid shape, with length L and height H L = 1000. H = 1000. nlayers = int(H / 100.) ncolumns = int(L / 100.) # make mesh m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers)) x = SpatialCoordinate(mesh) fieldlist = ['u', 'rho', 'theta'] timestepping = TimesteppingParameters(dt=1.0, maxk=4, maxi=1) output = OutputParameters(dirname=dirname + "/condens", dumpfreq=1, dumplist=['u'], perturbation_fields=['theta', 'rho']) parameters = CompressibleParameters() state = State(mesh, vertical_degree=1, horizontal_degree=1, family="CG", timestepping=timestepping, output=output, parameters=parameters, fieldlist=fieldlist, diagnostic_fields=[Sum('water_v', 'water_c')]) # declare initial fields u0 = state.fields("u") rho0 = state.fields("rho") theta0 = state.fields("theta") # spaces Vpsi = FunctionSpace(mesh, "CG", 2) Vt = theta0.function_space() Vr = rho0.function_space() # make a gradperp gradperp = lambda u: as_vector([-u.dx(1), u.dx(0)]) # declare tracer field and a background field water_v0 = state.fields("water_v", Vt) water_c0 = state.fields("water_c", Vt) # Isentropic background state Tsurf = Constant(300.) theta_b = Function(Vt).interpolate(Tsurf) rho_b = Function(Vr) # Calculate initial rho compressible_hydrostatic_balance(state, theta_b, rho_b, solve_for_rho=True) # set up water_v xc = 500. zc = 350. rc = 250. r = sqrt((x[0] - xc)**2 + (x[1] - zc)**2) w_expr = conditional(r > rc, 0., 0.25 * (1. + cos((pi / rc) * r))) # set up velocity field u_max = 10.0 psi_expr = ((-u_max * L / pi) * sin(2 * pi * x[0] / L) * sin(pi * x[1] / L)) psi0 = Function(Vpsi).interpolate(psi_expr) u0.project(gradperp(psi0)) theta0.interpolate(theta_b) rho0.interpolate(rho_b) water_v0.interpolate(w_expr) state.initialise([('u', u0), ('rho', rho0), ('theta', theta0), ('water_v', water_v0), ('water_c', water_c0)]) state.set_reference_profiles([('rho', rho_b), ('theta', theta_b)]) # set up advection schemes rhoeqn = AdvectionEquation(state, Vr, equation_form="continuity") thetaeqn = SUPGAdvection(state, Vt, supg_params={"dg_direction": "horizontal"}, equation_form="advective") # build advection dictionary advected_fields = [] advected_fields.append(("u", NoAdvection(state, u0, None))) advected_fields.append(("rho", SSPRK3(state, rho0, rhoeqn))) advected_fields.append(("theta", SSPRK3(state, theta0, thetaeqn))) advected_fields.append(("water_v", SSPRK3(state, water_v0, thetaeqn))) advected_fields.append(("water_c", SSPRK3(state, water_c0, thetaeqn))) physics_list = [Condensation(state)] # build time stepper stepper = AdvectionDiffusion(state, advected_fields, physics_list=physics_list) return stepper, 5.0
def run_dry_compressible(tmpdir): dt = 6.0 tmax = 2 * dt nlayers = 10 # horizontal layers ncols = 10 # number of columns Lx = 1000.0 Lz = 1000.0 m = PeriodicIntervalMesh(ncols, Lx) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=Lz / nlayers) output = OutputParameters(dirname=tmpdir + "/dry_compressible", dumpfreq=2, chkptfreq=2) parameters = CompressibleParameters() R_d = parameters.R_d g = parameters.g state = State(mesh, dt=dt, output=output, parameters=parameters) eqn = CompressibleEulerEquations(state, "CG", 1) # Initial conditions rho0 = state.fields("rho") theta0 = state.fields("theta") # Approximate hydrostatic balance x, z = SpatialCoordinate(mesh) T = Constant(300.0) zH = R_d * T / g p = Constant(100000.0) * exp(-z / zH) theta0.interpolate(tde.theta(parameters, T, p)) rho0.interpolate(p / (R_d * T)) state.set_reference_profiles([('rho', rho0), ('theta', theta0)]) # Add perturbation r = sqrt((x - Lx / 2)**2 + (z - Lz / 2)**2) theta_pert = 1.0 * exp(-(r / (Lx / 5))**2) theta0.interpolate(theta0 + theta_pert) # Set up transport schemes transported_fields = [ ImplicitMidpoint(state, "u"), SSPRK3(state, "rho"), SSPRK3(state, "theta") ] # Set up linear solver for the timestepping scheme linear_solver = CompressibleSolver(state, eqn) # build time stepper stepper = CrankNicolson(state, eqn, transported_fields, linear_solver=linear_solver) # Run stepper.run(t=0, tmax=tmax) # State for checking checkpoints checkpoint_name = 'dry_compressible_chkpt' new_path = join(abspath(dirname(__file__)), '..', f'data/{checkpoint_name}') check_output = OutputParameters(dirname=tmpdir + "/dry_compressible", checkpoint_pickup_filename=new_path) check_state = State(mesh, dt=dt, output=check_output, parameters=parameters) check_eqn = CompressibleEulerEquations(check_state, "CG", 1) check_stepper = CrankNicolson(check_state, check_eqn, []) check_stepper.run(t=0, tmax=0, pickup=True) return state, check_state
def setup_saturated(dirname): # set up grid and time stepping parameters dt = 1. tmax = 3. deltax = 400. L = 2000. H = 10000. nlayers = int(H / deltax) ncolumns = int(L / deltax) m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers) # option to easily change between recovered and not if necessary # default should be to use lowest order set of spaces recovered = True degree = 0 if recovered else 1 fieldlist = ['u', 'rho', 'theta'] timestepping = TimesteppingParameters(dt=dt, maxk=4, maxi=1) output = OutputParameters(dirname=dirname + '/saturated_balance', dumpfreq=1, dumplist=['u'], perturbation_fields=['water_v']) parameters = CompressibleParameters() diagnostics = Diagnostics(*fieldlist) diagnostic_fields = [Theta_e()] state = State(mesh, vertical_degree=degree, horizontal_degree=degree, family="CG", timestepping=timestepping, output=output, parameters=parameters, diagnostics=diagnostics, fieldlist=fieldlist, diagnostic_fields=diagnostic_fields) # Initial conditions u0 = state.fields("u") rho0 = state.fields("rho") theta0 = state.fields("theta") water_v0 = state.fields("water_v", theta0.function_space()) water_c0 = state.fields("water_c", theta0.function_space()) moisture = ['water_v', 'water_c'] # spaces Vu = u0.function_space() Vt = theta0.function_space() Vr = rho0.function_space() # Isentropic background state Tsurf = Constant(300.) total_water = Constant(0.02) theta_e = Function(Vt).interpolate(Tsurf) water_t = Function(Vt).interpolate(total_water) # Calculate hydrostatic Pi saturated_hydrostatic_balance(state, theta_e, water_t) water_c0.assign(water_t - water_v0) state.initialise([('u', u0), ('rho', rho0), ('theta', theta0), ('water_v', water_v0), ('water_c', water_c0)]) state.set_reference_profiles([('rho', rho0), ('theta', theta0), ('water_v', water_v0)]) # Set up advection schemes if recovered: VDG1 = FunctionSpace(mesh, "DG", 1) VCG1 = FunctionSpace(mesh, "CG", 1) Vt_brok = FunctionSpace(mesh, BrokenElement(Vt.ufl_element())) Vu_DG1 = VectorFunctionSpace(mesh, "DG", 1) Vu_CG1 = VectorFunctionSpace(mesh, "CG", 1) u_opts = RecoveredOptions(embedding_space=Vu_DG1, recovered_space=Vu_CG1, broken_space=Vu, boundary_method=Boundary_Method.dynamics) rho_opts = RecoveredOptions(embedding_space=VDG1, recovered_space=VCG1, broken_space=Vr, boundary_method=Boundary_Method.dynamics) theta_opts = RecoveredOptions(embedding_space=VDG1, recovered_space=VCG1, broken_space=Vt_brok) ueqn = EmbeddedDGAdvection(state, Vu, equation_form="advective", options=u_opts) rhoeqn = EmbeddedDGAdvection(state, Vr, equation_form="continuity", options=rho_opts) thetaeqn = EmbeddedDGAdvection(state, Vt, equation_form="advective", options=theta_opts) else: ueqn = EulerPoincare(state, Vu) rhoeqn = AdvectionEquation(state, Vr, equation_form="continuity") thetaeqn = EmbeddedDGAdvection(state, Vt, equation_form="advective", options=EmbeddedDGOptions()) advected_fields = [('rho', SSPRK3(state, rho0, rhoeqn)), ('theta', SSPRK3(state, theta0, thetaeqn)), ('water_v', SSPRK3(state, water_v0, thetaeqn)), ('water_c', SSPRK3(state, water_c0, thetaeqn))] if recovered: advected_fields.append(('u', SSPRK3(state, u0, ueqn))) else: advected_fields.append(('u', ThetaMethod(state, u0, ueqn))) linear_solver = CompressibleSolver(state, moisture=moisture) # Set up forcing if recovered: compressible_forcing = CompressibleForcing(state, moisture=moisture, euler_poincare=False) else: compressible_forcing = CompressibleForcing(state, moisture=moisture) # add physics physics_list = [Condensation(state)] # build time stepper stepper = CrankNicolson(state, advected_fields, linear_solver, compressible_forcing, physics_list=physics_list) return stepper, tmax
T_eq = 300.0 # Isothermal atmospheric temperature (K) p_eq = 1000.0 * 100.0 # Reference surface pressure at the equator u_max = 20.0 # Maximum amplitude of the zonal wind (m/s) d = 5000.0 # Width parameter for Theta' lamda_c = 2.0 * pi / 3.0 # Longitudinal centerpoint of Theta' phi_c = 0.0 # Latitudinal centerpoint of Theta' (equator) deltaTheta = 1.0 # Maximum amplitude of Theta' (K) L_z = 20000.0 # Vertical wave length of the Theta' perturb. # Cubed-sphere horizontal mesh m = CubedSphereMesh(radius=a, refinement_level=refinements, degree=2) # Build volume mesh z_top = 1.0e4 # Height position of the model top mesh = ExtrudedMesh(m, layers=nlayers, layer_height=z_top / nlayers, extrusion_type="radial") x = SpatialCoordinate(mesh) # Create polar coordinates: # Since we use a CG1 field, this is constant on layers W_Q1 = FunctionSpace(mesh, "CG", 1) z_expr = sqrt(x[0] * x[0] + x[1] * x[1] + x[2] * x[2]) - a z = Function(W_Q1).interpolate(z_expr) lat_expr = asin(x[2] / sqrt(x[0] * x[0] + x[1] * x[1] + x[2] * x[2])) lat = Function(W_Q1).interpolate(lat_expr) lon = Function(W_Q1).interpolate(atan_2(x[1], x[0])) fieldlist = ['u', 'rho', 'theta'] timestepping = TimesteppingParameters(dt=dt, maxk=4, maxi=1) dirname = 'meanflow_ref_hybridization'
form, parameters=parameters)[0] indices = IndexDict( {idx: sympy.symbols(name) for idx, name in index_names}) expr = expression(impero_kernel.tree, impero_kernel.temporaries, indices, top=True) p1 = sympy.symbols("p") + 1 '''Currently assume p+1 quad points in each direction.''' return expr.subs([(i, p1) for i in indices.values()]).expand() m = ExtrudedMesh(UnitSquareMesh(2, 2, quadrilateral=True), 2) mass = form.mass(m, 6) poisson = form.poisson(m, 6) hyperelasticity = form.hyperelasticity(m, 6) curl_curl = form.curl_curl(m, 6) parameters = firedrake.parameters['form_compiler'].copy() parameters['return_impero'] = True parameters['mode'] = 'spectral' for mode, action in (("assembly", False), ("action", True)): print(mode) print(" mass: ", complexity(mass, parameters, action)) print(" laplacian: ", complexity(poisson, parameters, action)) print(" hyperelasticity:", complexity(hyperelasticity, parameters,
def setup_tracer(dirname, hybridization): # declare grid shape, with length L and height H L = 1000. H = 1000. nlayers = int(H / 100.) ncolumns = int(L / 100.) # make mesh m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers)) fieldlist = ['u', 'rho', 'theta'] timestepping = TimesteppingParameters(dt=10.0, maxk=4, maxi=1) output = OutputParameters(dirname=dirname+"/tracer", dumpfreq=1, dumplist=['u'], perturbation_fields=['theta', 'rho']) parameters = CompressibleParameters() state = State(mesh, vertical_degree=1, horizontal_degree=1, family="CG", timestepping=timestepping, output=output, parameters=parameters, fieldlist=fieldlist, diagnostic_fields=[Difference('theta', 'tracer')]) # declare initial fields u0 = state.fields("u") rho0 = state.fields("rho") theta0 = state.fields("theta") # spaces Vu = u0.function_space() Vt = theta0.function_space() Vr = rho0.function_space() # declare tracer field and a background field tracer0 = state.fields("tracer", Vt) # Isentropic background state Tsurf = Constant(300.) theta_b = Function(Vt).interpolate(Tsurf) rho_b = Function(Vr) # Calculate initial rho compressible_hydrostatic_balance(state, theta_b, rho_b, solve_for_rho=True) # set up perturbation to theta xc = 500. zc = 350. rc = 250. x = SpatialCoordinate(mesh) r = sqrt((x[0]-xc)**2 + (x[1]-zc)**2) theta_pert = conditional(r > rc, 0., 0.25*(1. + cos((pi/rc)*r))) theta0.interpolate(theta_b + theta_pert) rho0.interpolate(rho_b) tracer0.interpolate(theta0) state.initialise([('u', u0), ('rho', rho0), ('theta', theta0), ('tracer', tracer0)]) state.set_reference_profiles([('rho', rho_b), ('theta', theta_b)]) # set up advection schemes ueqn = EulerPoincare(state, Vu) rhoeqn = AdvectionEquation(state, Vr, equation_form="continuity") thetaeqn = SUPGAdvection(state, Vt, supg_params={"dg_direction": "horizontal"}, equation_form="advective") # build advection dictionary advected_fields = [] advected_fields.append(("u", ThetaMethod(state, u0, ueqn))) advected_fields.append(("rho", SSPRK3(state, rho0, rhoeqn))) advected_fields.append(("theta", SSPRK3(state, theta0, thetaeqn))) advected_fields.append(("tracer", SSPRK3(state, tracer0, thetaeqn))) # Set up linear solver if hybridization: linear_solver = HybridizedCompressibleSolver(state) else: linear_solver = CompressibleSolver(state) compressible_forcing = CompressibleForcing(state) # build time stepper stepper = CrankNicolson(state, advected_fields, linear_solver, compressible_forcing) return stepper, 100.0
def make_mesh(nx, ny): return ExtrudedMesh(UnitIntervalMesh(nx), ny, layer_height=1.0 / ny)
def setup_saturated(dirname, recovered): # set up grid and time stepping parameters dt = 1. tmax = 3. deltax = 400. L = 2000. H = 10000. nlayers = int(H / deltax) ncolumns = int(L / deltax) m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers) # option to easily change between recovered and not if necessary # default should be to use lowest order set of spaces degree = 0 if recovered else 1 output = OutputParameters(dirname=dirname + '/saturated_balance', dumpfreq=1, dumplist=['u']) parameters = CompressibleParameters() diagnostic_fields = [Theta_e()] state = State(mesh, dt=dt, output=output, parameters=parameters, diagnostic_fields=diagnostic_fields) tracers = [WaterVapour(), CloudWater()] if recovered: u_transport_option = "vector_advection_form" else: u_transport_option = "vector_invariant_form" eqns = CompressibleEulerEquations(state, "CG", degree, u_transport_option=u_transport_option, active_tracers=tracers) # Initial conditions u0 = state.fields("u") rho0 = state.fields("rho") theta0 = state.fields("theta") water_v0 = state.fields("vapour_mixing_ratio") water_c0 = state.fields("cloud_liquid_mixing_ratio") moisture = ['vapour_mixing_ratio', 'cloud_liquid_mixing_ratio'] # spaces Vu = u0.function_space() Vt = theta0.function_space() Vr = rho0.function_space() # Isentropic background state Tsurf = Constant(300.) total_water = Constant(0.02) theta_e = Function(Vt).interpolate(Tsurf) water_t = Function(Vt).interpolate(total_water) # Calculate hydrostatic exner saturated_hydrostatic_balance(state, theta_e, water_t) water_c0.assign(water_t - water_v0) state.set_reference_profiles([('rho', rho0), ('theta', theta0)]) # Set up transport schemes if recovered: VDG1 = state.spaces("DG1_equispaced") VCG1 = FunctionSpace(mesh, "CG", 1) Vt_brok = FunctionSpace(mesh, BrokenElement(Vt.ufl_element())) Vu_DG1 = VectorFunctionSpace(mesh, VDG1.ufl_element()) Vu_CG1 = VectorFunctionSpace(mesh, "CG", 1) u_opts = RecoveredOptions(embedding_space=Vu_DG1, recovered_space=Vu_CG1, broken_space=Vu, boundary_method=Boundary_Method.dynamics) rho_opts = RecoveredOptions(embedding_space=VDG1, recovered_space=VCG1, broken_space=Vr, boundary_method=Boundary_Method.dynamics) theta_opts = RecoveredOptions(embedding_space=VDG1, recovered_space=VCG1, broken_space=Vt_brok) wv_opts = RecoveredOptions(embedding_space=VDG1, recovered_space=VCG1, broken_space=Vt_brok) wc_opts = RecoveredOptions(embedding_space=VDG1, recovered_space=VCG1, broken_space=Vt_brok) else: rho_opts = None theta_opts = EmbeddedDGOptions() wv_opts = EmbeddedDGOptions() wc_opts = EmbeddedDGOptions() transported_fields = [ SSPRK3(state, 'rho', options=rho_opts), SSPRK3(state, 'theta', options=theta_opts), SSPRK3(state, 'vapour_mixing_ratio', options=wv_opts), SSPRK3(state, 'cloud_liquid_mixing_ratio', options=wc_opts) ] if recovered: transported_fields.append(SSPRK3(state, 'u', options=u_opts)) else: transported_fields.append(ImplicitMidpoint(state, 'u')) linear_solver = CompressibleSolver(state, eqns, moisture=moisture) # add physics physics_list = [Condensation(state)] # build time stepper stepper = CrankNicolson(state, eqns, transported_fields, linear_solver=linear_solver, physics_list=physics_list) return stepper, tmax
def make_mesh(nx, ny, nz, quadrilateral=False): return ExtrudedMesh(UnitSquareMesh(nx, ny, quadrilateral=quadrilateral), nz, layer_height=1.0 / nz)
def setup_balance(dirname): # set up grid and time stepping parameters dt = 1. tmax = 5. deltax = 400 L = 2000. H = 10000. nlayers = int(H / deltax) ncolumns = int(L / deltax) m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers) fieldlist = ['u', 'rho', 'theta'] timestepping = TimesteppingParameters(dt=dt, maxk=4, maxi=1) output = OutputParameters(dirname=dirname + '/dry_balance', dumpfreq=10, dumplist=['u']) parameters = CompressibleParameters() diagnostics = Diagnostics(*fieldlist) diagnostic_fields = [] state = State(mesh, vertical_degree=1, horizontal_degree=1, family="CG", timestepping=timestepping, output=output, parameters=parameters, diagnostics=diagnostics, fieldlist=fieldlist, diagnostic_fields=diagnostic_fields) # Initial conditions u0 = state.fields("u") rho0 = state.fields("rho") theta0 = state.fields("theta") # spaces Vu = u0.function_space() Vt = theta0.function_space() Vr = rho0.function_space() # Isentropic background state Tsurf = Constant(300.) theta0.interpolate(Tsurf) # Calculate hydrostatic Pi compressible_hydrostatic_balance(state, theta0, rho0, solve_for_rho=True) state.initialise([('u', u0), ('rho', rho0), ('theta', theta0)]) state.set_reference_profiles([('rho', rho0), ('theta', theta0)]) # Set up advection schemes ueqn = VectorInvariant(state, Vu) rhoeqn = AdvectionEquation(state, Vr, equation_form="continuity") thetaeqn = EmbeddedDGAdvection(state, Vt, equation_form="advective", options=EmbeddedDGOptions()) advected_fields = [("u", ThetaMethod(state, u0, ueqn)), ("rho", SSPRK3(state, rho0, rhoeqn)), ("theta", SSPRK3(state, theta0, thetaeqn))] # Set up linear solver linear_solver = CompressibleSolver(state) # Set up forcing compressible_forcing = CompressibleForcing(state) # build time stepper stepper = CrankNicolson(state, advected_fields, linear_solver, compressible_forcing) return stepper, tmax
def setup_recovered_space(dirname): # declare grid shape, with length L and height H L = 400. H = 400. nlayers = int(H / 20.) ncolumns = int(L / 20.) # make mesh m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers)) x, z = SpatialCoordinate(mesh) fieldlist = ['u', 'rho'] timestepping = TimesteppingParameters(dt=1.0, maxk=4, maxi=1) output = OutputParameters(dirname=dirname + "/recovered_space_test", dumpfreq=5, dumplist=['u']) parameters = CompressibleParameters() state = State(mesh, vertical_degree=1, horizontal_degree=1, family="CG", timestepping=timestepping, output=output, parameters=parameters, fieldlist=fieldlist) # declare initial fields u0 = state.fields("u") # spaces Vpsi = FunctionSpace(mesh, "CG", 2) VDG0 = FunctionSpace(mesh, "DG", 0) VDG1 = FunctionSpace(mesh, "DG", 1) VCG1 = FunctionSpace(mesh, "CG", 1) # set up tracer field tracer0 = state.fields("tracer", VDG0) # make a gradperp gradperp = lambda u: as_vector([-u.dx(1), u.dx(0)]) # set up bubble xc = 200. zc = 200. rc = 100. tracer0.interpolate( conditional( sqrt((x - xc)**2.0) < rc, conditional( sqrt((z - zc)**2.0) < rc, Constant(0.2), Constant(0.0)), Constant(0.0))) # set up velocity field u_max = Constant(10.0) psi_expr = -u_max * z psi0 = Function(Vpsi).interpolate(psi_expr) u0.project(gradperp(psi0)) state.initialise([('u', u0), ('tracer', tracer0)]) # set up advection schemes recovered_opts = RecoveredOptions(embedding_space=VDG1, recovered_space=VCG1, broken_space=VDG0, boundary_method=Boundary_Method.dynamics) tracereqn = EmbeddedDGAdvection(state, VDG0, equation_form="continuity", options=recovered_opts) # build advection dictionary advected_fields = [] advected_fields.append(('tracer', SSPRK3(state, tracer0, tracereqn))) # build time stepper stepper = AdvectionDiffusion(state, advected_fields) return stepper, 100.0
else: tmax = 3600. if '--hybridization' in sys.argv: hybridization = True else: hybridization = False nlayers = 50 # horizontal layers columns = 50 # number of columns L = 3.0e5 m = PeriodicIntervalMesh(columns, L) # build volume mesh H = 1.0e4 # Height position of the model top mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H / nlayers) fieldlist = ['u', 'rho', 'theta'] timestepping = TimesteppingParameters(dt=dt) dirname = 'sk_linear' if hybridization: dirname += '_hybridization' output = OutputParameters(dirname=dirname, dumplist=['u'], perturbation_fields=['theta', 'rho']) parameters = CompressibleParameters() state = State(mesh,
def setup_hori_limiters(dirname): # declare grid shape L = 400. H = L ncolumns = int(L / 10.) nlayers = ncolumns # make mesh m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers)) x, z = SpatialCoordinate(mesh) fieldlist = ['u'] timestepping = TimesteppingParameters(dt=1.0, maxk=4, maxi=1) output = OutputParameters(dirname=dirname + "/limiting_hori", dumpfreq=5, dumplist=['u'], perturbation_fields=['theta0', 'theta1']) parameters = CompressibleParameters() state = State(mesh, vertical_degree=1, horizontal_degree=1, family="CG", timestepping=timestepping, output=output, parameters=parameters, fieldlist=fieldlist) # make elements # v is continuous in vertical, h is horizontal cell = mesh._base_mesh.ufl_cell().cellname() DG0_element = FiniteElement("DG", cell, 0) CG1_element = FiniteElement("CG", cell, 1) DG1_element = FiniteElement("DG", cell, 1) CG2_element = FiniteElement("CG", cell, 2) V0_element = TensorProductElement(DG0_element, CG1_element) V1_element = TensorProductElement(DG1_element, CG2_element) # spaces Vpsi = FunctionSpace(mesh, "CG", 2) VDG1 = FunctionSpace(mesh, "DG", 1) VCG1 = FunctionSpace(mesh, "CG", 1) V0 = FunctionSpace(mesh, V0_element) V1 = FunctionSpace(mesh, V1_element) V0_brok = FunctionSpace(mesh, BrokenElement(V0.ufl_element())) V0_spaces = (VDG1, VCG1, V0_brok) # declare initial fields u0 = state.fields("u") theta0 = state.fields("theta0", V0) theta1 = state.fields("theta1", V1) # make a gradperp gradperp = lambda u: as_vector([-u.dx(1), u.dx(0)]) # Isentropic background state Tsurf = 300. thetab = Constant(Tsurf) theta_b1 = Function(V1).interpolate(thetab) theta_b0 = Function(V0).interpolate(thetab) # set up bubble xc = 200. zc = 200. rc = 100. theta_expr = conditional( sqrt((x - xc)**2.0) < rc, conditional(sqrt((z - zc)**2.0) < rc, Constant(2.0), Constant(0.0)), Constant(0.0)) theta_pert1 = Function(V1).interpolate(theta_expr) theta_pert0 = Function(V0).interpolate(theta_expr) # set up velocity field u_max = Constant(10.0) psi_expr = -u_max * z psi0 = Function(Vpsi).interpolate(psi_expr) u0.project(gradperp(psi0)) theta0.interpolate(theta_b0 + theta_pert0) theta1.interpolate(theta_b1 + theta_pert1) state.initialise([('u', u0), ('theta1', theta1), ('theta0', theta0)]) state.set_reference_profiles([('theta1', theta_b1), ('theta0', theta_b0)]) # set up advection schemes thetaeqn1 = EmbeddedDGAdvection(state, V1, equation_form="advective") thetaeqn0 = EmbeddedDGAdvection(state, V0, equation_form="advective", recovered_spaces=V0_spaces) # build advection dictionary advected_fields = [] advected_fields.append(('u', NoAdvection(state, u0, None))) advected_fields.append(('theta1', SSPRK3(state, theta1, thetaeqn1, limiter=ThetaLimiter(thetaeqn1)))) advected_fields.append(('theta0', SSPRK3(state, theta0, thetaeqn0, limiter=VertexBasedLimiter(VDG1)))) # build time stepper stepper = AdvectionDiffusion(state, advected_fields) return stepper, 40.0