diff --git a/examples/compressible/dcmip_3_1_meanflow_quads.py b/examples/compressible/dcmip_3_1_meanflow_quads.py index a0f93c33b..40cbd20fd 100644 --- a/examples/compressible/dcmip_3_1_meanflow_quads.py +++ b/examples/compressible/dcmip_3_1_meanflow_quads.py @@ -11,6 +11,9 @@ from firedrake import exp, acos, cos, sin, pi, sqrt, asin, atan_2 import sys +# ---------------------------------------------------------------------------- # +# Test case parameters +# ---------------------------------------------------------------------------- # nlayers = 10 # Number of vertical layers refinements = 3 # Number of horiz. refinements @@ -24,7 +27,6 @@ tmax = 3600.0 dumpfreq = int(tmax / (4*dt)) - parameters = CompressibleParameters() a_ref = 6.37122e6 # Radius of the Earth (m) X = 125.0 # Reduced-size Earth reduction factor @@ -43,18 +45,24 @@ 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. +z_top = 1.0e4 # Height position of the model top + +# ---------------------------------------------------------------------------- # +# Set up model objects +# ---------------------------------------------------------------------------- # +# Domain # 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") +domain = Domain(mesh, dt, "RTCF", 1) x = SpatialCoordinate(mesh) + # Create polar coordinates: # Since we use a CG1 field, this is constant on layers W_Q1 = FunctionSpace(mesh, "CG", 1) @@ -64,29 +72,41 @@ lat = Function(W_Q1).interpolate(lat_expr) lon = Function(W_Q1).interpolate(atan_2(x[1], x[0])) -dirname = 'dcmip_3_1_meanflow' +# Equation +eqns = CompressibleEulerEquations(domain, parameters) +# I/O +dirname = 'dcmip_3_1_meanflow' output = OutputParameters(dirname=dirname, dumpfreq=dumpfreq, - perturbation_fields=['theta', 'rho'], log_level='INFO') +diagnostic_fields = [Perturbation('theta'), Perturbation('rho')] +io = IO(domain, output, diagnostic_fields=diagnostic_fields) -state = State(mesh, - dt=dt, - output=output, - parameters=parameters) +# Transport schemes +transported_fields = [ImplicitMidpoint(domain, "u"), + SSPRK3(domain, "rho", subcycles=2), + SSPRK3(domain, "theta", options=SUPGOptions(), subcycles=2)] -eqns = CompressibleEulerEquations(state, "RTCF", 1) +# Linear solver +linear_solver = CompressibleSolver(eqns) +# Time stepper +stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + linear_solver=linear_solver) + +# ---------------------------------------------------------------------------- # # Initial conditions -u0 = state.fields.u -theta0 = state.fields.theta -rho0 = state.fields.rho +# ---------------------------------------------------------------------------- # + +u0 = stepper.fields('u') +theta0 = stepper.fields('theta') +rho0 = stepper.fields('rho') # spaces -Vu = state.spaces("HDiv") -Vt = state.spaces("theta") -Vr = state.spaces("DG") +Vu = domain.spaces("HDiv") +Vt = domain.spaces("theta") +Vr = domain.spaces("DG") # Initial conditions with u0 uexpr = as_vector([-u_max*x[1]/a, u_max*x[0]/a, 0.0]) @@ -122,7 +142,7 @@ theta0.interpolate(theta_b) # Compute the balanced density -compressible_hydrostatic_balance(state, +compressible_hydrostatic_balance(eqns, theta_b, rho_b, top=False, @@ -131,20 +151,11 @@ theta0 += theta_b rho0.assign(rho_b) -state.initialise([('u', u0), ('rho', rho0), ('theta', theta0)]) -state.set_reference_profiles([('rho', rho_b), ('theta', theta_b)]) +stepper.set_reference_profiles([('rho', rho_b), ('theta', theta_b)]) -# Set up transport schemes -transported_fields = [ImplicitMidpoint(state, "u"), - SSPRK3(state, "rho", subcycles=2), - SSPRK3(state, "theta", options=SUPGOptions(), subcycles=2)] - -# Set up linear solver -linear_solver = CompressibleSolver(state, eqns) - -# Build time stepper -stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields, - linear_solver=linear_solver) +# ---------------------------------------------------------------------------- # +# Run +# ---------------------------------------------------------------------------- # # Run! stepper.run(t=0, tmax=tmax) diff --git a/examples/compressible/dry_bryan_fritsch.py b/examples/compressible/dry_bryan_fritsch.py index f7775bb21..842064165 100644 --- a/examples/compressible/dry_bryan_fritsch.py +++ b/examples/compressible/dry_bryan_fritsch.py @@ -13,7 +13,14 @@ FunctionSpace, VectorFunctionSpace) import sys +# ---------------------------------------------------------------------------- # +# Test case parameters +# ---------------------------------------------------------------------------- # + dt = 1.0 +L = 10000. +H = 10000. + if '--running-tests' in sys.argv: deltax = 1000. tmax = 5. @@ -24,43 +31,70 @@ degree = 0 dirname = 'dry_bryan_fritsch' -# make mesh -L = 10000. -H = 10000. +# ---------------------------------------------------------------------------- # +# Set up model objects +# ---------------------------------------------------------------------------- # + +# Domain nlayers = int(H/deltax) ncolumns = int(L/deltax) m = IntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) +domain = Domain(mesh, dt, "CG", degree) +# Equation +params = CompressibleParameters() +u_transport_option = "vector_advection_form" +eqns = CompressibleEulerEquations(domain, params, + u_transport_option=u_transport_option, + no_normal_flow_bc_ids=[1, 2]) +# I/O output = OutputParameters(dirname=dirname, dumpfreq=int(tmax / (5*dt)), dumplist=['u'], - perturbation_fields=['theta'], log_level='INFO') +diagnostic_fields = [Perturbation('theta')] +io = IO(domain, output, diagnostic_fields=diagnostic_fields) -params = CompressibleParameters() +# Transport schemes -- set up options for using recovery wrapper +VDG1 = domain.spaces("DG1_equispaced") +VCG1 = FunctionSpace(mesh, "CG", 1) +Vu_DG1 = VectorFunctionSpace(mesh, VDG1.ufl_element()) +Vu_CG1 = VectorFunctionSpace(mesh, "CG", 1) -state = State(mesh, - dt=dt, - output=output, - parameters=params) +u_opts = RecoveryOptions(embedding_space=Vu_DG1, + recovered_space=Vu_CG1, + boundary_method=BoundaryMethod.taylor) +rho_opts = RecoveryOptions(embedding_space=VDG1, + recovered_space=VCG1, + boundary_method=BoundaryMethod.taylor) +theta_opts = RecoveryOptions(embedding_space=VDG1, + recovered_space=VCG1) -u_transport_option = "vector_advection_form" +transported_fields = [SSPRK3(domain, "rho", options=rho_opts), + SSPRK3(domain, "theta", options=theta_opts), + SSPRK3(domain, "u", options=u_opts)] -eqns = CompressibleEulerEquations(state, "CG", degree, - u_transport_option=u_transport_option, - no_normal_flow_bc_ids=[1, 2]) +# Linear solver +linear_solver = CompressibleSolver(eqns) +# Time stepper +stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + linear_solver=linear_solver) + +# ---------------------------------------------------------------------------- # # Initial conditions -u0 = state.fields("u") -rho0 = state.fields("rho") -theta0 = state.fields("theta") +# ---------------------------------------------------------------------------- # + +u0 = stepper.fields("u") +rho0 = stepper.fields("rho") +theta0 = stepper.fields("theta") # spaces -Vu = state.spaces("HDiv") -Vt = state.spaces("theta") -Vr = state.spaces("DG") +Vu = domain.spaces("HDiv") +Vt = domain.spaces("theta") +Vr = domain.spaces("DG") x, z = SpatialCoordinate(mesh) # Define constant theta_e and water_t @@ -68,7 +102,7 @@ theta_b = Function(Vt).interpolate(Constant(Tsurf)) # Calculate hydrostatic fields -compressible_hydrostatic_balance(state, theta_b, rho0, solve_for_rho=True) +compressible_hydrostatic_balance(eqns, theta_b, rho0, solve_for_rho=True) # make mean fields rho_b = Function(Vr).assign(rho0) @@ -95,33 +129,11 @@ rho_solver = LinearVariationalSolver(rho_problem) rho_solver.solve() -state.set_reference_profiles([('rho', rho_b), - ('theta', theta_b)]) - -# Set up transport schemes -VDG1 = state.spaces("DG1_equispaced") -VCG1 = FunctionSpace(mesh, "CG", 1) -Vu_DG1 = VectorFunctionSpace(mesh, VDG1.ufl_element()) -Vu_CG1 = VectorFunctionSpace(mesh, "CG", 1) - -u_opts = RecoveryOptions(embedding_space=Vu_DG1, - recovered_space=Vu_CG1, - boundary_method=BoundaryMethod.taylor) -rho_opts = RecoveryOptions(embedding_space=VDG1, - recovered_space=VCG1, - boundary_method=BoundaryMethod.taylor) -theta_opts = RecoveryOptions(embedding_space=VDG1, - recovered_space=VCG1) - -transported_fields = [SSPRK3(state, "rho", options=rho_opts), - SSPRK3(state, "theta", options=theta_opts), - SSPRK3(state, "u", options=u_opts)] +stepper.set_reference_profiles([('rho', rho_b), + ('theta', theta_b)]) -# Set up linear solver -linear_solver = CompressibleSolver(state, eqns) - -# build time stepper -stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields, - linear_solver=linear_solver) +# ---------------------------------------------------------------------------- # +# Run +# ---------------------------------------------------------------------------- # stepper.run(t=0, tmax=tmax) diff --git a/examples/compressible/moist_bryan_fritsch.py b/examples/compressible/moist_bryan_fritsch.py index dd5977846..7e3e65493 100644 --- a/examples/compressible/moist_bryan_fritsch.py +++ b/examples/compressible/moist_bryan_fritsch.py @@ -15,7 +15,14 @@ LinearVariationalProblem, LinearVariationalSolver) import sys +# ---------------------------------------------------------------------------- # +# Test case parameters +# ---------------------------------------------------------------------------- # + dt = 1.0 +L = 10000. +H = 10000. + if '--running-tests' in sys.argv: deltax = 1000. tmax = 5. @@ -23,47 +30,65 @@ deltax = 200 tmax = 1000. -L = 10000. -H = 10000. +# ---------------------------------------------------------------------------- # +# Set up model objects +# ---------------------------------------------------------------------------- # + +# Domain nlayers = int(H/deltax) ncolumns = int(L/deltax) m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) degree = 1 +domain = Domain(mesh, dt, 'CG', degree) -dirname = 'moist_bryan_fritsch' +# Equation +params = CompressibleParameters() +tracers = [WaterVapour(), CloudWater()] +eqns = CompressibleEulerEquations(domain, params, active_tracers=tracers) +# I/O +dirname = 'moist_bryan_fritsch' output = OutputParameters(dirname=dirname, dumpfreq=int(tmax / (5*dt)), dumplist=['u'], - perturbation_fields=[], log_level='INFO') +diagnostic_fields = [Theta_e(eqns)] +io = IO(domain, output, diagnostic_fields=diagnostic_fields) -params = CompressibleParameters() -diagnostic_fields = [Theta_e()] -tracers = [WaterVapour(), CloudWater()] +# Transport schemes +transported_fields = [SSPRK3(domain, "rho"), + SSPRK3(domain, "theta", options=EmbeddedDGOptions()), + SSPRK3(domain, "water_vapour", options=EmbeddedDGOptions()), + SSPRK3(domain, "cloud_water", options=EmbeddedDGOptions()), + ImplicitMidpoint(domain, "u")] + +# Linear solver +linear_solver = CompressibleSolver(eqns) -state = State(mesh, - dt=dt, - output=output, - parameters=params, - diagnostic_fields=diagnostic_fields) +# Physics schemes (condensation/evaporation) +physics_schemes = [(SaturationAdjustment(eqns), ForwardEuler(domain))] -eqns = CompressibleEulerEquations(state, "CG", degree, active_tracers=tracers) +# Time stepper +stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + linear_solver=linear_solver, + physics_schemes=physics_schemes) +# ---------------------------------------------------------------------------- # # Initial conditions -u0 = state.fields("u") -rho0 = state.fields("rho") -theta0 = state.fields("theta") -water_v0 = state.fields("water_vapour") -water_c0 = state.fields("cloud_water") -moisture = ["water_vapour", "cloud_water"] +# ---------------------------------------------------------------------------- # + +u0 = stepper.fields("u") +rho0 = stepper.fields("rho") +theta0 = stepper.fields("theta") +water_v0 = stepper.fields("water_vapour") +water_c0 = stepper.fields("cloud_water") # spaces -Vu = state.spaces("HDiv") -Vt = state.spaces("theta") -Vr = state.spaces("DG") +Vu = domain.spaces("HDiv") +Vt = domain.spaces("theta") +Vr = domain.spaces("DG") x, z = SpatialCoordinate(mesh) quadrature_degree = (4, 4) dxp = dx(degree=(quadrature_degree)) @@ -75,15 +100,15 @@ water_t = Function(Vt).assign(total_water) # Calculate hydrostatic fields -saturated_hydrostatic_balance(state, theta_e, water_t) +saturated_hydrostatic_balance(eqns, stepper.fields, theta_e, water_t) # make mean fields theta_b = Function(Vt).assign(theta0) rho_b = Function(Vr).assign(rho0) water_vb = Function(Vt).assign(water_v0) water_cb = Function(Vt).assign(water_t - water_vb) -exner_b = thermodynamics.exner_pressure(state.parameters, rho_b, theta_b) -Tb = thermodynamics.T(state.parameters, theta_b, exner_b, r_v=water_vb) +exner_b = thermodynamics.exner_pressure(eqns.parameters, rho_b, theta_b) +Tb = thermodynamics.T(eqns.parameters, theta_b, exner_b, r_v=water_vb) # define perturbation xc = L / 2 @@ -115,10 +140,10 @@ rho_recoverer = Recoverer(rho0, rho_averaged) rho_recoverer.project() -exner = thermodynamics.exner_pressure(state.parameters, rho_averaged, theta0) -p = thermodynamics.p(state.parameters, exner) -T = thermodynamics.T(state.parameters, theta0, exner, r_v=w_v) -w_sat = thermodynamics.r_sat(state.parameters, T, p) +exner = thermodynamics.exner_pressure(eqns.parameters, rho_averaged, theta0) +p = thermodynamics.p(eqns.parameters, exner) +T = thermodynamics.T(eqns.parameters, theta0, exner, r_v=w_v) +w_sat = thermodynamics.r_sat(eqns.parameters, T, p) w_functional = (phi * w_v * dxp - phi * w_sat * dxp) w_problem = NonlinearVariationalProblem(w_functional, w_v) @@ -128,29 +153,13 @@ water_v0.assign(w_v) water_c0.assign(water_t - water_v0) -state.set_reference_profiles([('rho', rho_b), - ('theta', theta_b), - ('water_vapour', water_vb)]) - -rho_opts = None -theta_opts = EmbeddedDGOptions() -u_transport = ImplicitMidpoint(state, "u") +stepper.set_reference_profiles([('rho', rho_b), + ('theta', theta_b), + ('water_vapour', water_vb), + ('cloud_water', water_cb)]) -transported_fields = [SSPRK3(state, "rho", options=rho_opts), - SSPRK3(state, "theta", options=theta_opts), - SSPRK3(state, "water_vapour", options=theta_opts), - SSPRK3(state, "cloud_water", options=theta_opts), - u_transport] - -# Set up linear solver -linear_solver = CompressibleSolver(state, eqns, moisture=moisture) - -# define condensation -physics_schemes = [(SaturationAdjustment(eqns, params), ForwardEuler(state))] - -# build time stepper -stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields, - linear_solver=linear_solver, - physics_schemes=physics_schemes) +# ---------------------------------------------------------------------------- # +# Run +# ---------------------------------------------------------------------------- # stepper.run(t=0, tmax=tmax) diff --git a/examples/compressible/mountain_hydrostatic.py b/examples/compressible/mountain_hydrostatic.py index 944e101b7..d97744bd3 100644 --- a/examples/compressible/mountain_hydrostatic.py +++ b/examples/compressible/mountain_hydrostatic.py @@ -9,7 +9,13 @@ exp, pi, cos, Function, conditional, Mesh, op2, sqrt) import sys +# ---------------------------------------------------------------------------- # +# Test case parameters +# ---------------------------------------------------------------------------- # + dt = 5.0 +L = 240000. # Domain length +H = 50000. # Height position of the model top if '--running-tests' in sys.argv: tmax = dt @@ -21,13 +27,15 @@ dumpfreq = int(tmax / (5*dt)) +# ---------------------------------------------------------------------------- # +# Set up model objects +# ---------------------------------------------------------------------------- # + +# Domain nlayers = res*20 # horizontal layers columns = res*12 # number of columns -L = 240000. m = PeriodicIntervalMesh(columns, L) -# build volume mesh -H = 50000. # Height position of the model top ext_mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) Vc = VectorFunctionSpace(ext_mesh, "DG", 2) coord = SpatialCoordinate(ext_mesh) @@ -38,42 +46,74 @@ hm = 1. zs = hm*a**2/((x-xc)**2 + a**2) -dirname = 'hydrostatic_mountain' zh = 5000. xexpr = as_vector([x, conditional(z < zh, z + cos(0.5*pi*z/zh)**6*zs, z)]) new_coords = Function(Vc).interpolate(xexpr) mesh = Mesh(new_coords) +domain = Domain(mesh, dt, "CG", 1) + +# Equation +parameters = CompressibleParameters(g=9.80665, cp=1004.) +sponge = SpongeLayerParameters(H=H, z_level=H-20000, mubar=0.3/dt) +eqns = CompressibleEulerEquations(domain, parameters, sponge=sponge) +# I/O +dirname = 'hydrostatic_mountain' output = OutputParameters(dirname=dirname, dumpfreq=dumpfreq, dumplist=['u'], - perturbation_fields=['theta', 'rho'], log_level='INFO') +diagnostic_fields = [CourantNumber(), VelocityZ(), HydrostaticImbalance(eqns), + Perturbation('theta'), Perturbation('rho')] +io = IO(domain, output, diagnostic_fields=diagnostic_fields) -parameters = CompressibleParameters(g=9.80665, cp=1004.) -diagnostic_fields = [CourantNumber(), VelocityZ(), HydrostaticImbalance()] +# Transport schemes +theta_opts = SUPGOptions() +transported_fields = [ImplicitMidpoint(domain, "u"), + SSPRK3(domain, "rho"), + SSPRK3(domain, "theta", options=theta_opts)] -state = State(mesh, - dt=dt, - output=output, - parameters=parameters, - diagnostic_fields=diagnostic_fields) +# Linear solver +params = {'mat_type': 'matfree', + 'ksp_type': 'preonly', + 'pc_type': 'python', + 'pc_python_type': 'firedrake.SCPC', + # Velocity mass operator is singular in the hydrostatic case. + # So for reconstruction, we eliminate rho into u + 'pc_sc_eliminate_fields': '1, 0', + 'condensed_field': {'ksp_type': 'fgmres', + 'ksp_rtol': 1.0e-8, + 'ksp_atol': 1.0e-8, + 'ksp_max_it': 100, + 'pc_type': 'gamg', + 'pc_gamg_sym_graph': True, + 'mg_levels': {'ksp_type': 'gmres', + 'ksp_max_it': 5, + 'pc_type': 'bjacobi', + 'sub_pc_type': 'ilu'}}} -# sponge function -sponge = SpongeLayerParameters(H=H, z_level=H-20000, mubar=0.3/dt) +alpha = 0.51 # off-centering parameter +linear_solver = CompressibleSolver(eqns, alpha, solver_parameters=params, + overwrite_solver_parameters=True) -eqns = CompressibleEulerEquations(state, "CG", 1, sponge=sponge) +# Time stepper +stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + linear_solver=linear_solver, + alpha=alpha) +# ---------------------------------------------------------------------------- # # Initial conditions -u0 = state.fields("u") -rho0 = state.fields("rho") -theta0 = state.fields("theta") +# ---------------------------------------------------------------------------- # + +u0 = stepper.fields("u") +rho0 = stepper.fields("rho") +theta0 = stepper.fields("theta") # spaces -Vu = state.spaces("HDiv") -Vt = state.spaces("theta") -Vr = state.spaces("DG") +Vu = domain.spaces("HDiv") +Vt = domain.spaces("theta") +Vr = domain.spaces("DG") # Thermodynamic constants required for setting initial conditions # and reference profiles @@ -107,7 +147,7 @@ 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'}} -compressible_hydrostatic_balance(state, theta_b, rho_b, exner, +compressible_hydrostatic_balance(eqns, theta_b, rho_b, exner, top=True, exner_boundary=0.5, params=exner_params) @@ -123,13 +163,13 @@ def minimum(f): p0 = minimum(exner) -compressible_hydrostatic_balance(state, theta_b, rho_b, exner, +compressible_hydrostatic_balance(eqns, theta_b, rho_b, exner, top=True, params=exner_params) p1 = minimum(exner) alpha = 2.*(p1-p0) beta = p1-alpha exner_top = (1.-beta)/alpha -compressible_hydrostatic_balance(state, theta_b, rho_b, exner, +compressible_hydrostatic_balance(eqns, theta_b, rho_b, exner, top=True, exner_boundary=exner_top, solve_for_rho=True, params=exner_params) @@ -138,41 +178,11 @@ def minimum(f): u0.project(as_vector([20.0, 0.0])) remove_initial_w(u0) -state.set_reference_profiles([('rho', rho_b), - ('theta', theta_b)]) - -# Set up transport schemes -theta_opts = SUPGOptions() -transported_fields = [ImplicitMidpoint(state, "u"), - SSPRK3(state, "rho"), - SSPRK3(state, "theta", options=theta_opts)] - -# Set up linear solver -params = {'mat_type': 'matfree', - 'ksp_type': 'preonly', - 'pc_type': 'python', - 'pc_python_type': 'firedrake.SCPC', - # Velocity mass operator is singular in the hydrostatic case. - # So for reconstruction, we eliminate rho into u - 'pc_sc_eliminate_fields': '1, 0', - 'condensed_field': {'ksp_type': 'fgmres', - 'ksp_rtol': 1.0e-8, - 'ksp_atol': 1.0e-8, - 'ksp_max_it': 100, - 'pc_type': 'gamg', - 'pc_gamg_sym_graph': True, - 'mg_levels': {'ksp_type': 'gmres', - 'ksp_max_it': 5, - 'pc_type': 'bjacobi', - 'sub_pc_type': 'ilu'}}} - -alpha = 0.51 # off-centering parameter -linear_solver = CompressibleSolver(state, eqns, alpha, solver_parameters=params, - overwrite_solver_parameters=True) +stepper.set_reference_profiles([('rho', rho_b), + ('theta', theta_b)]) -# build time stepper -stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields, - linear_solver=linear_solver, - alpha=alpha) +# ---------------------------------------------------------------------------- # +# Run +# ---------------------------------------------------------------------------- # stepper.run(t=0, tmax=tmax) diff --git a/examples/compressible/mountain_nonhydrostatic.py b/examples/compressible/mountain_nonhydrostatic.py index 56509715b..b8e6656be 100644 --- a/examples/compressible/mountain_nonhydrostatic.py +++ b/examples/compressible/mountain_nonhydrostatic.py @@ -9,7 +9,14 @@ exp, pi, cos, Function, conditional, Mesh, op2) import sys +# ---------------------------------------------------------------------------- # +# Test case parameters +# ---------------------------------------------------------------------------- # + dt = 5.0 +L = 144000. # Domain length +H = 35000. # Height position of the model top + if '--running-tests' in sys.argv: tmax = dt dumpfreq = 1 @@ -21,11 +28,12 @@ nlayers = 70 # horizontal layers columns = 180 # number of columns -L = 144000. -m = PeriodicIntervalMesh(columns, L) +# ---------------------------------------------------------------------------- # +# Set up model objects +# ---------------------------------------------------------------------------- # -# build volume mesh -H = 35000. # Height position of the model top +# Domain +m = PeriodicIntervalMesh(columns, L) ext_mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) Vc = VectorFunctionSpace(ext_mesh, "DG", 2) coord = SpatialCoordinate(ext_mesh) @@ -36,42 +44,52 @@ hm = 1. zs = hm*a**2/((x-xc)**2 + a**2) -dirname = 'nonhydrostatic_mountain' zh = 5000. xexpr = as_vector([x, conditional(z < zh, z + cos(0.5*pi*z/zh)**6*zs, z)]) new_coords = Function(Vc).interpolate(xexpr) mesh = Mesh(new_coords) +domain = Domain(mesh, dt, "CG", 1) +# Equation +parameters = CompressibleParameters(g=9.80665, cp=1004.) +sponge = SpongeLayerParameters(H=H, z_level=H-10000, mubar=0.15/dt) +eqns = CompressibleEulerEquations(domain, parameters, sponge=sponge) + +# I/O +dirname = 'nonhydrostatic_mountain' output = OutputParameters(dirname=dirname, dumpfreq=dumpfreq, dumplist=['u'], - perturbation_fields=['theta', 'rho'], log_level='INFO') +diagnostic_fields = [CourantNumber(), VelocityZ(), Perturbation('theta'), Perturbation('rho')] +io = IO(domain, output, diagnostic_fields=diagnostic_fields) -parameters = CompressibleParameters(g=9.80665, cp=1004.) -diagnostic_fields = [CourantNumber(), VelocityZ()] - -state = State(mesh, - dt=dt, - output=output, - parameters=parameters, - diagnostic_fields=diagnostic_fields) +# Transport schemes +theta_opts = SUPGOptions() +transported_fields = [ImplicitMidpoint(domain, "u"), + SSPRK3(domain, "rho"), + SSPRK3(domain, "theta", options=theta_opts)] -# sponge function -sponge = SpongeLayerParameters(H=H, z_level=H-10000, mubar=0.15/dt) +# Linear solver +linear_solver = CompressibleSolver(eqns) -eqns = CompressibleEulerEquations(state, "CG", 1, sponge=sponge) +# Time stepper +stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + linear_solver=linear_solver) +# ---------------------------------------------------------------------------- # # Initial conditions -u0 = state.fields("u") -rho0 = state.fields("rho") -theta0 = state.fields("theta") +# ---------------------------------------------------------------------------- # + +u0 = stepper.fields("u") +rho0 = stepper.fields("rho") +theta0 = stepper.fields("theta") # spaces -Vu = state.spaces("HDiv") -Vt = state.spaces("theta") -Vr = state.spaces("DG") +Vu = domain.spaces("HDiv") +Vt = domain.spaces("theta") +Vr = domain.spaces("DG") # Thermodynamic constants required for setting initial conditions # and reference profiles @@ -103,7 +121,7 @@ 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'}} -compressible_hydrostatic_balance(state, theta_b, rho_b, exner, +compressible_hydrostatic_balance(eqns, theta_b, rho_b, exner, top=True, exner_boundary=0.5, params=exner_params) @@ -119,13 +137,13 @@ def minimum(f): p0 = minimum(exner) -compressible_hydrostatic_balance(state, theta_b, rho_b, exner, +compressible_hydrostatic_balance(eqns, theta_b, rho_b, exner, top=True, params=exner_params) p1 = minimum(exner) alpha = 2.*(p1-p0) beta = p1-alpha exner_top = (1.-beta)/alpha -compressible_hydrostatic_balance(state, theta_b, rho_b, exner, +compressible_hydrostatic_balance(eqns, theta_b, rho_b, exner, top=True, exner_boundary=exner_top, solve_for_rho=True, params=exner_params) @@ -134,20 +152,11 @@ def minimum(f): u0.project(as_vector([10.0, 0.0])) remove_initial_w(u0) -state.set_reference_profiles([('rho', rho_b), - ('theta', theta_b)]) +stepper.set_reference_profiles([('rho', rho_b), + ('theta', theta_b)]) -# Set up transport schemes -theta_opts = SUPGOptions() -transported_fields = [ImplicitMidpoint(state, "u"), - SSPRK3(state, "rho"), - SSPRK3(state, "theta", options=theta_opts)] - -# Set up linear solver -linear_solver = CompressibleSolver(state, eqns) - -# build time stepper -stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields, - linear_solver=linear_solver) +# ---------------------------------------------------------------------------- # +# Run +# ---------------------------------------------------------------------------- # stepper.run(t=0, tmax=tmax) diff --git a/examples/compressible/robert_bubble.py b/examples/compressible/robert_bubble.py index 1c4cf397a..52c305a65 100644 --- a/examples/compressible/robert_bubble.py +++ b/examples/compressible/robert_bubble.py @@ -9,6 +9,10 @@ Constant, pi, cos, Function, sqrt, conditional) import sys +# ---------------------------------------------------------------------------- # +# Test case parameters +# ---------------------------------------------------------------------------- # + dt = 1. L = 1000. H = 1000. @@ -24,37 +28,54 @@ nlayers = int(H/10.) ncolumns = int(L/10.) +# ---------------------------------------------------------------------------- # +# Set up model objects +# ---------------------------------------------------------------------------- # + +# Domain m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) +domain = Domain(mesh, dt, "CG", 1) -dirname = 'robert_bubble' +# Equation +parameters = CompressibleParameters() +eqns = CompressibleEulerEquations(domain, parameters) +# I/O +dirname = 'robert_bubble' output = OutputParameters(dirname=dirname, dumpfreq=dumpfreq, dumplist=['u'], - perturbation_fields=['theta', 'rho'], log_level='INFO') +diagnostic_fields = [CourantNumber(), Perturbation('theta'), Perturbation('rho')] +io = IO(domain, output, diagnostic_fields=diagnostic_fields) -parameters = CompressibleParameters() -diagnostic_fields = [CourantNumber()] +# Transport schemes +theta_opts = EmbeddedDGOptions() +transported_fields = [] +transported_fields.append(ImplicitMidpoint(domain, "u")) +transported_fields.append(SSPRK3(domain, "rho")) +transported_fields.append(SSPRK3(domain, "theta", options=theta_opts)) -state = State(mesh, - dt=dt, - output=output, - parameters=parameters, - diagnostic_fields=diagnostic_fields) +# Linear solver +linear_solver = CompressibleSolver(eqns) -eqns = CompressibleEulerEquations(state, "CG", 1) +# Time stepper +stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + linear_solver=linear_solver) +# ---------------------------------------------------------------------------- # # Initial conditions -u0 = state.fields("u") -rho0 = state.fields("rho") -theta0 = state.fields("theta") +# ---------------------------------------------------------------------------- # + +u0 = stepper.fields("u") +rho0 = stepper.fields("rho") +theta0 = stepper.fields("theta") # spaces -Vu = state.spaces("HDiv") -Vt = state.spaces("theta") -Vr = state.spaces("DG") +Vu = domain.spaces("HDiv") +Vt = domain.spaces("theta") +Vr = domain.spaces("DG") # Isentropic background state Tsurf = Constant(300.) @@ -63,7 +84,7 @@ rho_b = Function(Vr) # Calculate hydrostatic exner -compressible_hydrostatic_balance(state, theta_b, rho_b, solve_for_rho=True) +compressible_hydrostatic_balance(eqns, theta_b, rho_b, solve_for_rho=True) x = SpatialCoordinate(mesh) xc = 500. @@ -75,21 +96,11 @@ theta0.interpolate(theta_b + theta_pert) rho0.interpolate(rho_b) -state.set_reference_profiles([('rho', rho_b), - ('theta', theta_b)]) +stepper.set_reference_profiles([('rho', rho_b), + ('theta', theta_b)]) -# Set up transport schemes -theta_opts = EmbeddedDGOptions() -transported_fields = [] -transported_fields.append(ImplicitMidpoint(state, "u")) -transported_fields.append(SSPRK3(state, "rho")) -transported_fields.append(SSPRK3(state, "theta", options=theta_opts)) - -# Set up linear solver -linear_solver = CompressibleSolver(state, eqns) - -# build time stepper -stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields, - linear_solver=linear_solver) +# ---------------------------------------------------------------------------- # +# Run +# ---------------------------------------------------------------------------- # stepper.run(t=0, tmax=tmax) diff --git a/examples/compressible/skamarock_klemp_hydrostatic.py b/examples/compressible/skamarock_klemp_hydrostatic.py index 9f3ba1770..7e2e12255 100644 --- a/examples/compressible/skamarock_klemp_hydrostatic.py +++ b/examples/compressible/skamarock_klemp_hydrostatic.py @@ -10,6 +10,10 @@ ExtrudedMesh, exp, sin, Function, pi) import sys +# ---------------------------------------------------------------------------- # +# Test case parameters +# ---------------------------------------------------------------------------- # + dt = 25. if '--running-tests' in sys.argv: nlayers = 5 # horizontal layers @@ -22,45 +26,59 @@ tmax = 60000.0 dumpfreq = int(tmax / (2*dt)) +L = 6.0e6 # Length of domain +H = 1.0e4 # Height position of the model top -L = 6.0e6 -m = PeriodicRectangleMesh(columns, 1, L, 1.e4, quadrilateral=True) +# ---------------------------------------------------------------------------- # +# Set up model objects +# ---------------------------------------------------------------------------- # -# build volume mesh -H = 1.0e4 # Height position of the model top +# Domain -- 3D volume mesh +m = PeriodicRectangleMesh(columns, 1, L, 1.e4, quadrilateral=True) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) +domain = Domain(mesh, dt, "RTCF", 1) -dirname = 'skamarock_klemp_hydrostatic' +# Equation +parameters = CompressibleParameters() +Omega = as_vector((0., 0., 0.5e-4)) +balanced_pg = as_vector((0., -1.0e-4*20, 0.)) +eqns = CompressibleEulerEquations(domain, parameters, Omega=Omega, + extra_terms=[("u", balanced_pg)]) +# I/O +dirname = 'skamarock_klemp_hydrostatic' output = OutputParameters(dirname=dirname, dumpfreq=dumpfreq, dumplist=['u'], - perturbation_fields=['theta', 'rho'], log_level='INFO') +diagnostic_fields = [CourantNumber(), Perturbation('theta'), Perturbation('rho')] +io = IO(domain, output, diagnostic_fields=diagnostic_fields) -parameters = CompressibleParameters() -diagnostic_fields = [CourantNumber()] +# Transport schemes +transported_fields = [] +transported_fields.append(ImplicitMidpoint(domain, "u")) +transported_fields.append(SSPRK3(domain, "rho")) +transported_fields.append(SSPRK3(domain, "theta", options=SUPGOptions())) -state = State(mesh, - dt=dt, - output=output, - parameters=parameters, - diagnostic_fields=diagnostic_fields) +# Linear solver +linear_solver = CompressibleSolver(eqns) -Omega = as_vector((0., 0., 0.5e-4)) -balanced_pg = as_vector((0., -1.0e-4*20, 0.)) -eqns = CompressibleEulerEquations(state, "RTCF", 1, Omega=Omega, - extra_terms=[("u", balanced_pg)]) +# Time stepper +stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + linear_solver=linear_solver) +# ---------------------------------------------------------------------------- # # Initial conditions -u0 = state.fields("u") -rho0 = state.fields("rho") -theta0 = state.fields("theta") +# ---------------------------------------------------------------------------- # + +u0 = stepper.fields("u") +rho0 = stepper.fields("rho") +theta0 = stepper.fields("theta") # spaces -Vu = state.spaces("HDiv") -Vt = state.spaces("theta") -Vr = state.spaces("DG") +Vu = domain.spaces("HDiv") +Vt = domain.spaces("theta") +Vr = domain.spaces("DG") # Thermodynamic constants required for setting initial conditions # and reference profiles @@ -85,26 +103,17 @@ theta_pert = deltaTheta*sin(pi*z/H)/(1 + (x - L/2)**2/a**2) theta0.interpolate(theta_b + theta_pert) -compressible_hydrostatic_balance(state, theta_b, rho_b, +compressible_hydrostatic_balance(eqns, theta_b, rho_b, solve_for_rho=True) rho0.assign(rho_b) u0.project(as_vector([20.0, 0.0, 0.0])) -state.set_reference_profiles([('rho', rho_b), - ('theta', theta_b)]) - -# Set up transport schemes -transported_fields = [] -transported_fields.append(ImplicitMidpoint(state, "u")) -transported_fields.append(SSPRK3(state, "rho")) -transported_fields.append(SSPRK3(state, "theta", options=SUPGOptions())) +stepper.set_reference_profiles([('rho', rho_b), + ('theta', theta_b)]) -# Set up linear solver -linear_solver = CompressibleSolver(state, eqns) - -# build time stepper -stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields, - linear_solver=linear_solver) +# ---------------------------------------------------------------------------- # +# Run +# ---------------------------------------------------------------------------- # stepper.run(t=0, tmax=tmax) diff --git a/examples/compressible/skamarock_klemp_nonlinear.py b/examples/compressible/skamarock_klemp_nonlinear.py index 984fd4e06..15d947a68 100644 --- a/examples/compressible/skamarock_klemp_nonlinear.py +++ b/examples/compressible/skamarock_klemp_nonlinear.py @@ -13,7 +13,14 @@ import numpy as np import sys +# ---------------------------------------------------------------------------- # +# Test case parameters +# ---------------------------------------------------------------------------- # + dt = 6. +L = 3.0e5 # Domain length +H = 1.0e4 # Height position of the model top + if '--running-tests' in sys.argv: nlayers = 5 columns = 30 @@ -25,53 +32,61 @@ tmax = 3600. dumpfreq = int(tmax / (2*dt)) +# ---------------------------------------------------------------------------- # +# Set up model objects +# ---------------------------------------------------------------------------- # -L = 3.0e5 +# Domain -- 3D volume mesh 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) +domain = Domain(mesh, dt, "CG", 1) +# Equation +Tsurf = 300. +parameters = CompressibleParameters() +eqns = CompressibleEulerEquations(domain, parameters) + +# I/O points_x = np.linspace(0., L, 100) points_z = [H/2.] points = np.array([p for p in itertools.product(points_x, points_z)]) - dirname = 'skamarock_klemp_nonlinear' - output = OutputParameters(dirname=dirname, dumpfreq=dumpfreq, pddumpfreq=dumpfreq, dumplist=['u'], - perturbation_fields=['theta', 'rho'], point_data=[('theta_perturbation', points)], log_level='INFO') +diagnostic_fields = [CourantNumber(), Gradient("u"), Perturbation('theta'), + Gradient("theta_perturbation"), Perturbation('rho'), + RichardsonNumber("theta", parameters.g/Tsurf), Gradient("theta")] +io = IO(domain, output, diagnostic_fields=diagnostic_fields) -parameters = CompressibleParameters() -g = parameters.g -Tsurf = 300. - -diagnostic_fields = [CourantNumber(), Gradient("u"), - Gradient("theta_perturbation"), - RichardsonNumber("theta", g/Tsurf), Gradient("theta")] +# Transport schemes +theta_opts = SUPGOptions() +transported_fields = [ImplicitMidpoint(domain, "u"), + SSPRK3(domain, "rho"), + SSPRK3(domain, "theta", options=theta_opts)] -state = State(mesh, - dt=dt, - output=output, - parameters=parameters, - diagnostic_fields=diagnostic_fields) +# Linear solver +linear_solver = CompressibleSolver(eqns) -eqns = CompressibleEulerEquations(state, "CG", 1) +# Time stepper +stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + linear_solver=linear_solver) +# ---------------------------------------------------------------------------- # # Initial conditions -u0 = state.fields("u") -rho0 = state.fields("rho") -theta0 = state.fields("theta") +# ---------------------------------------------------------------------------- # + +u0 = stepper.fields("u") +rho0 = stepper.fields("rho") +theta0 = stepper.fields("theta") # spaces -Vu = state.spaces("HDiv") -Vt = state.spaces("theta") -Vr = state.spaces("DG") +Vu = domain.spaces("HDiv") +Vt = domain.spaces("theta") +Vr = domain.spaces("DG") # Thermodynamic constants required for setting initial conditions # and reference profiles @@ -81,14 +96,13 @@ x, z = SpatialCoordinate(mesh) # N^2 = (g/theta)dtheta/dz => dtheta/dz = theta N^2g => theta=theta_0exp(N^2gz) -Tsurf = 300. thetab = Tsurf*exp(N**2*z/g) theta_b = Function(Vt).interpolate(thetab) rho_b = Function(Vr) # Calculate hydrostatic exner -compressible_hydrostatic_balance(state, theta_b, rho_b) +compressible_hydrostatic_balance(eqns, theta_b, rho_b) a = 5.0e3 deltaTheta = 1.0e-2 @@ -97,20 +111,11 @@ rho0.assign(rho_b) u0.project(as_vector([20.0, 0.0])) -state.set_reference_profiles([('rho', rho_b), - ('theta', theta_b)]) - -# Set up transport schemes -theta_opts = SUPGOptions() -transported_fields = [ImplicitMidpoint(state, "u"), - SSPRK3(state, "rho"), - SSPRK3(state, "theta", options=theta_opts)] +stepper.set_reference_profiles([('rho', rho_b), + ('theta', theta_b)]) -# Set up linear solver -linear_solver = CompressibleSolver(state, eqns) - -# build time stepper -stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields, - linear_solver=linear_solver) +# ---------------------------------------------------------------------------- # +# Run +# ---------------------------------------------------------------------------- # stepper.run(t=0, tmax=tmax) diff --git a/examples/compressible/straka_bubble.py b/examples/compressible/straka_bubble.py index 14e7cbda4..972ab491e 100644 --- a/examples/compressible/straka_bubble.py +++ b/examples/compressible/straka_bubble.py @@ -10,6 +10,10 @@ conditional) import sys +# ---------------------------------------------------------------------------- # +# Test case parameters +# ---------------------------------------------------------------------------- # + if '--running-tests' in sys.argv: res_dt = {800.: 4.} tmax = 4. @@ -27,45 +31,65 @@ for delta, dt in res_dt.items(): - dirname = "straka_dx%s_dt%s" % (delta, dt) + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain nlayers = int(H/delta) # horizontal layers columns = int(L/delta) # number of columns - m = PeriodicIntervalMesh(columns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) + domain = Domain(mesh, dt, "CG", 1) + # Equation + parameters = CompressibleParameters() + diffusion_options = [ + ("u", DiffusionParameters(kappa=75., mu=10./delta)), + ("theta", DiffusionParameters(kappa=75., mu=10./delta))] + eqns = CompressibleEulerEquations(domain, parameters, + diffusion_options=diffusion_options) + + # I/O + dirname = "straka_dx%s_dt%s" % (delta, dt) dumpfreq = int(tmax / (ndumps*dt)) output = OutputParameters(dirname=dirname, dumpfreq=dumpfreq, dumplist=['u'], - perturbation_fields=['theta', 'rho'], log_level='INFO') + diagnostic_fields = [CourantNumber(), Perturbation('theta'), Perturbation('rho')] + io = IO(domain, output, diagnostic_fields=diagnostic_fields) - parameters = CompressibleParameters() - diagnostic_fields = [CourantNumber()] + # Transport schemes + theta_opts = SUPGOptions() + transported_fields = [ImplicitMidpoint(domain, "u"), + SSPRK3(domain, "rho"), + SSPRK3(domain, "theta", options=theta_opts)] - state = State(mesh, - dt=dt, - output=output, - parameters=parameters, - diagnostic_fields=diagnostic_fields) + # Linear solver + linear_solver = CompressibleSolver(eqns) - diffusion_options = [ - ("u", DiffusionParameters(kappa=75., mu=10./delta)), - ("theta", DiffusionParameters(kappa=75., mu=10./delta))] + # Diffusion schemes + diffusion_schemes = [BackwardEuler(domain, "u"), + BackwardEuler(domain, "theta")] - eqns = CompressibleEulerEquations(state, "CG", 1, - diffusion_options=diffusion_options) + # Time stepper + stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + linear_solver=linear_solver, + diffusion_schemes=diffusion_schemes) + # ------------------------------------------------------------------------ # # Initial conditions - u0 = state.fields("u") - rho0 = state.fields("rho") - theta0 = state.fields("theta") + # ------------------------------------------------------------------------ # + + u0 = stepper.fields("u") + rho0 = stepper.fields("rho") + theta0 = stepper.fields("theta") # spaces - Vu = state.spaces("HDiv") - Vt = state.spaces("theta") - Vr = state.spaces("DG") + Vu = domain.spaces("HDiv") + Vt = domain.spaces("theta") + Vr = domain.spaces("DG") # Isentropic background state Tsurf = Constant(300.) @@ -75,7 +99,7 @@ exner = Function(Vr) # Calculate hydrostatic exner - compressible_hydrostatic_balance(state, theta_b, rho_b, exner0=exner, + compressible_hydrostatic_balance(eqns, theta_b, rho_b, exner0=exner, solve_for_rho=True) x = SpatialCoordinate(mesh) @@ -89,24 +113,11 @@ theta0.interpolate(theta_b + T_pert*exner) rho0.assign(rho_b) - state.set_reference_profiles([('rho', rho_b), - ('theta', theta_b)]) + stepper.set_reference_profiles([('rho', rho_b), + ('theta', theta_b)]) - # Set up transport schemes - theta_opts = SUPGOptions() - transported_fields = [ImplicitMidpoint(state, "u"), - SSPRK3(state, "rho"), - SSPRK3(state, "theta", options=theta_opts)] - - # Set up linear solver - linear_solver = CompressibleSolver(state, eqns) - - diffusion_schemes = [BackwardEuler(state, "u"), - BackwardEuler(state, "theta")] - - # build time stepper - stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields, - linear_solver=linear_solver, - diffusion_schemes=diffusion_schemes) + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # stepper.run(t=0, tmax=tmax) diff --git a/examples/compressible/unsaturated_bubble.py b/examples/compressible/unsaturated_bubble.py index 3de8dddec..b07eaa34e 100644 --- a/examples/compressible/unsaturated_bubble.py +++ b/examples/compressible/unsaturated_bubble.py @@ -14,6 +14,10 @@ from firedrake.slope_limiter.vertex_based_limiter import VertexBasedLimiter import sys +# ---------------------------------------------------------------------------- # +# Test case parameters +# ---------------------------------------------------------------------------- # + dt = 1.0 if '--running-tests' in sys.argv: deltax = 240. @@ -28,63 +32,101 @@ h = 2400. nlayers = int(h/deltax) ncolumns = int(L/deltax) +degree = 0 +# ---------------------------------------------------------------------------- # +# Set up model objects +# ---------------------------------------------------------------------------- # + +# Domain m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=h/nlayers) -degree = 0 +domain = Domain(mesh, dt, "CG", degree) +# Equation +params = CompressibleParameters() +tracers = [WaterVapour(), CloudWater(), Rain()] +eqns = CompressibleEulerEquations(domain, params, + active_tracers=tracers) + +# I/O dirname = 'unsaturated_bubble' output = OutputParameters(dirname=dirname, dumpfreq=tdump, - perturbation_fields=['theta', 'water_vapour', 'rho'], log_level='INFO') -params = CompressibleParameters() -diagnostic_fields = [RelativeHumidity()] -tracers = [WaterVapour(), CloudWater(), Rain()] +diagnostic_fields = [RelativeHumidity(eqns), Perturbation('theta'), + Perturbation('water_vapour'), Perturbation('rho')] +io = IO(domain, output, diagnostic_fields=diagnostic_fields) -state = State(mesh, - dt=dt, - output=output, - parameters=params, - diagnostic_fields=diagnostic_fields) +# Transport schemes -- specify options for using recovery wrapper +VDG1 = domain.spaces("DG1_equispaced") +VCG1 = FunctionSpace(mesh, "CG", 1) +Vu_DG1 = VectorFunctionSpace(mesh, VDG1.ufl_element()) +Vu_CG1 = VectorFunctionSpace(mesh, "CG", 1) -eqns = CompressibleEulerEquations(state, "CG", degree, - active_tracers=tracers) +u_opts = RecoveryOptions(embedding_space=Vu_DG1, + recovered_space=Vu_CG1, + boundary_method=BoundaryMethod.taylor) +rho_opts = RecoveryOptions(embedding_space=VDG1, + recovered_space=VCG1, + boundary_method=BoundaryMethod.taylor) +theta_opts = RecoveryOptions(embedding_space=VDG1, recovered_space=VCG1) +limiter = VertexBasedLimiter(VDG1) +transported_fields = [SSPRK3(domain, "u", options=u_opts), + SSPRK3(domain, "rho", options=rho_opts), + SSPRK3(domain, "theta", options=theta_opts), + SSPRK3(domain, "water_vapour", options=theta_opts, limiter=limiter), + SSPRK3(domain, "cloud_water", options=theta_opts, limiter=limiter), + SSPRK3(domain, "rain", options=theta_opts, limiter=limiter)] + +# Linear solver +linear_solver = CompressibleSolver(eqns) + +# Physics schemes +# NB: to use wrapper options with Fallout, need to pass field name to time discretisation +physics_schemes = [(Fallout(eqns, 'rain', domain), SSPRK3(domain, field_name='rain', options=theta_opts, limiter=limiter)), + (Coalescence(eqns), ForwardEuler(domain)), + (EvaporationOfRain(eqns), ForwardEuler(domain)), + (SaturationAdjustment(eqns), ForwardEuler(domain))] + +# Time stepper +stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + linear_solver=linear_solver, + physics_schemes=physics_schemes) + +# ---------------------------------------------------------------------------- # # Initial conditions -u0 = state.fields("u") -rho0 = state.fields("rho") -theta0 = state.fields("theta") -water_v0 = state.fields("water_vapour") -water_c0 = state.fields("cloud_water") -rain0 = state.fields("rain", theta0.function_space()) -moisture = ["water_vapour", "cloud_water", "rain"] +# ---------------------------------------------------------------------------- # + +u0 = stepper.fields("u") +rho0 = stepper.fields("rho") +theta0 = stepper.fields("theta") +water_v0 = stepper.fields("water_vapour") +water_c0 = stepper.fields("cloud_water") +rain0 = stepper.fields("rain") # spaces -Vu = state.spaces("HDiv") -Vt = state.spaces("theta") -Vr = state.spaces("DG") +Vu = domain.spaces("HDiv") +Vt = domain.spaces("theta") +Vr = domain.spaces("DG") x, z = SpatialCoordinate(mesh) quadrature_degree = (4, 4) dxp = dx(degree=(quadrature_degree)) -VDG1 = state.spaces("DG1_equispaced") -VCG1 = FunctionSpace(mesh, "CG", 1) -Vu_DG1 = VectorFunctionSpace(mesh, VDG1.ufl_element()) -Vu_CG1 = VectorFunctionSpace(mesh, "CG", 1) physics_boundary_method = BoundaryMethod.extruded # Define constant theta_e and water_t Tsurf = 283.0 psurf = 85000. -exner_surf = (psurf / state.parameters.p_0) ** state.parameters.kappa +exner_surf = (psurf / eqns.parameters.p_0) ** eqns.parameters.kappa humidity = 0.2 S = 1.3e-5 -theta_surf = thermodynamics.theta(state.parameters, Tsurf, psurf) +theta_surf = thermodynamics.theta(eqns.parameters, Tsurf, psurf) theta_d = Function(Vt).interpolate(theta_surf * exp(S*z)) H = Function(Vt).assign(humidity) # Calculate hydrostatic fields -unsaturated_hydrostatic_balance(state, theta_d, H, +unsaturated_hydrostatic_balance(eqns, stepper.fields, theta_d, H, exner_boundary=Constant(exner_surf)) # make mean fields @@ -116,21 +158,21 @@ w_h = Function(Vt) delta = 1.0 -R_d = state.parameters.R_d -R_v = state.parameters.R_v +R_d = eqns.parameters.R_d +R_v = eqns.parameters.R_v epsilon = R_d / R_v # make expressions for determining water_v0 -exner = thermodynamics.exner_pressure(state.parameters, rho_averaged, theta0) -p = thermodynamics.p(state.parameters, exner) -T = thermodynamics.T(state.parameters, theta0, exner, water_v0) -r_v_expr = thermodynamics.r_v(state.parameters, H, T, p) +exner = thermodynamics.exner_pressure(eqns.parameters, rho_averaged, theta0) +p = thermodynamics.p(eqns.parameters, exner) +T = thermodynamics.T(eqns.parameters, theta0, exner, water_v0) +r_v_expr = thermodynamics.r_v(eqns.parameters, H, T, p) # make expressions to evaluate residual -exner_ev = thermodynamics.exner_pressure(state.parameters, rho_averaged, theta0) -p_ev = thermodynamics.p(state.parameters, exner_ev) -T_ev = thermodynamics.T(state.parameters, theta0, exner_ev, water_v0) -RH_ev = thermodynamics.RH(state.parameters, water_v0, T_ev, p_ev) +exner_ev = thermodynamics.exner_pressure(eqns.parameters, rho_averaged, theta0) +p_ev = thermodynamics.p(eqns.parameters, exner_ev) +T_ev = thermodynamics.T(eqns.parameters, theta0, exner_ev, water_v0) +RH_ev = thermodynamics.RH(eqns.parameters, water_v0, T_ev, p_ev) RH = Function(Vt) # set-up rho problem to keep exner constant @@ -175,40 +217,11 @@ raise RuntimeError('Hydrostatic balance solve has not converged within %i' % i, 'iterations') # initialise fields -state.set_reference_profiles([('rho', rho_b), - ('theta', theta_b), - ('water_vapour', water_vb)]) - -# Set up transport schemes -u_opts = RecoveryOptions(embedding_space=Vu_DG1, - recovered_space=Vu_CG1, - boundary_method=BoundaryMethod.taylor) -rho_opts = RecoveryOptions(embedding_space=VDG1, - recovered_space=VCG1, - boundary_method=BoundaryMethod.taylor) -theta_opts = RecoveryOptions(embedding_space=VDG1, recovered_space=VCG1) -limiter = VertexBasedLimiter(VDG1) - -transported_fields = [SSPRK3(state, "u", options=u_opts), - SSPRK3(state, "rho", options=rho_opts), - SSPRK3(state, "theta", options=theta_opts), - SSPRK3(state, "water_vapour", options=theta_opts, limiter=limiter), - SSPRK3(state, "cloud_water", options=theta_opts, limiter=limiter), - SSPRK3(state, "rain", options=theta_opts, limiter=limiter)] - -# Set up linear solver -linear_solver = CompressibleSolver(state, eqns, moisture=moisture) - -# define physics schemes -# NB: to use wrapper options with Fallout, need to pass field name to time discretisation -physics_schemes = [(Fallout(eqns, 'rain', state), SSPRK3(state, field_name='rain', options=theta_opts, limiter=limiter)), - (Coalescence(eqns), ForwardEuler(state)), - (EvaporationOfRain(eqns, params), ForwardEuler(state)), - (SaturationAdjustment(eqns, params), ForwardEuler(state))] - -# build time stepper -stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields, - linear_solver=linear_solver, - physics_schemes=physics_schemes) +stepper.set_reference_profiles([('rho', rho_b), + ('theta', theta_b), + ('water_vapour', water_vb)]) +# ---------------------------------------------------------------------------- # +# Run +# ---------------------------------------------------------------------------- # stepper.run(t=0, tmax=tmax) diff --git a/examples/incompressible/skamarock_klemp_incompressible.py b/examples/incompressible/skamarock_klemp_incompressible.py index 8b05c1ded..e2296967e 100644 --- a/examples/incompressible/skamarock_klemp_incompressible.py +++ b/examples/incompressible/skamarock_klemp_incompressible.py @@ -10,7 +10,14 @@ sin, SpatialCoordinate, Function, pi) import sys +# ---------------------------------------------------------------------------- # +# Test case parameters +# ---------------------------------------------------------------------------- # + dt = 6. +L = 3.0e5 # Domain length +H = 1.0e4 # Height position of the model top + if '--running-tests' in sys.argv: tmax = dt dumpfreq = 1 @@ -23,38 +30,47 @@ columns = 300 # number of columns nlayers = 10 # horizontal layers -# set up mesh -L = 3.0e5 +# ---------------------------------------------------------------------------- # +# Set up model objects +# ---------------------------------------------------------------------------- # + +# Domain m = PeriodicIntervalMesh(columns, L) -H = 1.0e4 # Height position of the model top mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) +domain = Domain(mesh, dt, 'CG', 1) -# output parameters +# Equation +parameters = CompressibleParameters() +eqns = IncompressibleBoussinesqEquations(domain, parameters) + +# I/O output = OutputParameters(dirname='skamarock_klemp_incompressible', dumpfreq=dumpfreq, dumplist=['u'], - perturbation_fields=['b'], log_level='INFO') - -# physical parameters -parameters = CompressibleParameters() - # list of diagnostic fields, each defined in a class in diagnostics.py -diagnostic_fields = [CourantNumber(), Divergence()] +diagnostic_fields = [CourantNumber(), Divergence(), Perturbation('b')] +io = IO(domain, output, diagnostic_fields=diagnostic_fields) -# setup state -state = State(mesh, - dt=dt, - output=output, - parameters=parameters, - diagnostic_fields=diagnostic_fields) +# Transport schemes +b_opts = SUPGOptions() +transported_fields = [ImplicitMidpoint(domain, "u"), + SSPRK3(domain, "b", options=b_opts)] -eqns = IncompressibleBoussinesqEquations(state, "CG", 1) +# Linear solver +linear_solver = IncompressibleSolver(eqns) +# Time stepper +stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + linear_solver=linear_solver) + +# ---------------------------------------------------------------------------- # # Initial conditions -u0 = state.fields("u") -b0 = state.fields("b") -p0 = state.fields("p") +# ---------------------------------------------------------------------------- # + +u0 = stepper.fields("u") +b0 = stepper.fields("b") +p0 = stepper.fields("p") # spaces Vb = b0.function_space() @@ -75,25 +91,17 @@ # interpolate the expression to the function b0.interpolate(b_b + b_pert) -incompressible_hydrostatic_balance(state, b_b, p0) +incompressible_hydrostatic_balance(eqns, b_b, p0) uinit = (as_vector([20.0, 0.0])) u0.project(uinit) # set the background buoyancy -state.set_reference_profiles([('b', b_b)]) +stepper.set_reference_profiles([('b', b_b)]) -# Set up transport schemes -b_opts = SUPGOptions() -transported_fields = [ImplicitMidpoint(state, "u"), - SSPRK3(state, "b", options=b_opts)] - -# Set up linear solver for the timestepping scheme -linear_solver = IncompressibleSolver(state, eqns) - -# build time stepper -stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields, - linear_solver=linear_solver) +# ---------------------------------------------------------------------------- # +# Run +# ---------------------------------------------------------------------------- # # Run! stepper.run(t=0, tmax=tmax) diff --git a/examples/shallow_water/linear_williamson_2.py b/examples/shallow_water/linear_williamson_2.py index 88c4ba768..75930ac65 100644 --- a/examples/shallow_water/linear_williamson_2.py +++ b/examples/shallow_water/linear_williamson_2.py @@ -9,6 +9,10 @@ from firedrake import IcosahedralSphereMesh, SpatialCoordinate, as_vector, pi import sys +# ---------------------------------------------------------------------------- # +# Test case parameters +# ---------------------------------------------------------------------------- # + dt = 3600. day = 24.*60.*60. if '--running-tests' in sys.argv: @@ -23,32 +27,43 @@ R = 6371220. H = 2000. +# ---------------------------------------------------------------------------- # +# Set up model objects +# ---------------------------------------------------------------------------- # + +# Domain mesh = IcosahedralSphereMesh(radius=R, refinement_level=refinements, degree=3) x = SpatialCoordinate(mesh) mesh.init_cell_orientations(x) +domain = Domain(mesh, dt, 'BDM', 1) + +# Equation +parameters = ShallowWaterParameters(H=H) +Omega = parameters.Omega +x = SpatialCoordinate(mesh) +fexpr = 2*Omega*x[2]/R +eqns = LinearShallowWaterEquations(domain, parameters, fexpr=fexpr) +# I/O output = OutputParameters(dirname='linear_williamson_2', dumpfreq=dumpfreq, - steady_state_error_fields=['u', 'D'], log_level='INFO') -parameters = ShallowWaterParameters(H=H) +diagnostic_fields = [SteadyStateError('u'), SteadyStateError('D')] +io = IO(domain, output, diagnostic_fields=diagnostic_fields) -state = State(mesh, - dt=dt, - output=output, - parameters=parameters) +# Transport schemes +transport_schemes = [ForwardEuler(domain, "D")] -# Coriolis expression -Omega = parameters.Omega -x = SpatialCoordinate(mesh) -fexpr = 2*Omega*x[2]/R -eqns = LinearShallowWaterEquations(state, "BDM", 1, fexpr=fexpr) +# Time stepper +stepper = SemiImplicitQuasiNewton(eqns, io, transport_schemes) + +# ---------------------------------------------------------------------------- # +# Initial conditions +# ---------------------------------------------------------------------------- # -# interpolate initial conditions -# Initial/current conditions -u0 = state.fields("u") -D0 = state.fields("D") +u0 = stepper.fields("u") +D0 = stepper.fields("D") u_max = 2*pi*R/(12*day) # Maximum amplitude of the zonal wind (m/s) uexpr = as_vector([-u_max*x[1]/R, u_max*x[0]/R, 0.0]) g = parameters.g @@ -56,9 +71,11 @@ u0.project(uexpr) D0.interpolate(Dexpr) -transport_schemes = [ForwardEuler(state, "D")] +Dbar = Function(D0.function_space()).assign(H) +stepper.set_reference_profiles([('D', Dbar)]) -# build time stepper -stepper = SemiImplicitQuasiNewton(eqns, state, transport_schemes) +# ---------------------------------------------------------------------------- # +# Run +# ---------------------------------------------------------------------------- # stepper.run(t=0, tmax=tmax) diff --git a/examples/shallow_water/williamson_2.py b/examples/shallow_water/williamson_2.py index 9ff00b3c4..1053798d5 100644 --- a/examples/shallow_water/williamson_2.py +++ b/examples/shallow_water/williamson_2.py @@ -10,6 +10,10 @@ from firedrake import IcosahedralSphereMesh, SpatialCoordinate, as_vector, pi import sys +# ---------------------------------------------------------------------------- # +# Test case parameters +# ---------------------------------------------------------------------------- # + day = 24.*60.*60. if '--running-tests' in sys.argv: ref_dt = {3: 3000.} @@ -30,38 +34,51 @@ for ref_level, dt in ref_dt.items(): - dirname = "williamson_2_ref%s_dt%s" % (ref_level, dt) + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain mesh = IcosahedralSphereMesh(radius=R, refinement_level=ref_level, degree=1) x = SpatialCoordinate(mesh) global_normal = x mesh.init_cell_orientations(x) + domain = Domain(mesh, dt, 'BDM', 1) + + # Equation + Omega = parameters.Omega + fexpr = 2*Omega*x[2]/R + eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr) + # I/O + dirname = "williamson_2_ref%s_dt%s" % (ref_level, dt) dumpfreq = int(tmax / (ndumps*dt)) output = OutputParameters(dirname=dirname, dumpfreq=dumpfreq, dumplist_latlon=['D', 'D_error'], - steady_state_error_fields=['D', 'u'], log_level='INFO') diagnostic_fields = [RelativeVorticity(), PotentialVorticity(), ShallowWaterKineticEnergy(), - ShallowWaterPotentialEnergy(), - ShallowWaterPotentialEnstrophy()] + ShallowWaterPotentialEnergy(parameters), + ShallowWaterPotentialEnstrophy(), + SteadyStateError('u'), SteadyStateError('D')] + io = IO(domain, output, diagnostic_fields=diagnostic_fields) - state = State(mesh, - dt=dt, - output=output, - parameters=parameters, - diagnostic_fields=diagnostic_fields) + # Transport schemes + transported_fields = [ImplicitMidpoint(domain, "u"), + SSPRK3(domain, "D", subcycles=2)] - Omega = parameters.Omega - fexpr = 2*Omega*x[2]/R - eqns = ShallowWaterEquations(state, "BDM", 1, fexpr=fexpr) + # Time stepper + stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # - # interpolate initial conditions - u0 = state.fields("u") - D0 = state.fields("D") + u0 = stepper.fields("u") + D0 = stepper.fields("D") x = SpatialCoordinate(mesh) u_max = 2*pi*R/(12*day) # Maximum amplitude of the zonal wind (m/s) uexpr = as_vector([-u_max*x[1]/R, u_max*x[0]/R, 0.0]) @@ -71,10 +88,11 @@ u0.project(uexpr) D0.interpolate(Dexpr) - transported_fields = [ImplicitMidpoint(state, "u"), - SSPRK3(state, "D", subcycles=2)] + Dbar = Function(D0.function_space()).assign(H) + stepper.set_reference_profiles([('D', Dbar)]) - # build time stepper - stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields) + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # stepper.run(t=0, tmax=tmax) diff --git a/examples/shallow_water/williamson_5.py b/examples/shallow_water/williamson_5.py index 8bfc1d144..db586d6db 100644 --- a/examples/shallow_water/williamson_5.py +++ b/examples/shallow_water/williamson_5.py @@ -10,6 +10,10 @@ as_vector, pi, sqrt, Min) import sys +# ---------------------------------------------------------------------------- # +# Test case parameters +# ---------------------------------------------------------------------------- # + day = 24.*60.*60. if '--running-tests' in sys.argv: ref_dt = {3: 3000.} @@ -30,26 +34,18 @@ for ref_level, dt in ref_dt.items(): - dirname = "williamson_5_ref%s_dt%s" % (ref_level, dt) + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain mesh = IcosahedralSphereMesh(radius=R, refinement_level=ref_level, degree=1) x = SpatialCoordinate(mesh) mesh.init_cell_orientations(x) + domain = Domain(mesh, dt, 'BDM', 1) - dumpfreq = int(tmax / (ndumps*dt)) - output = OutputParameters(dirname=dirname, - dumplist_latlon=['D'], - dumpfreq=dumpfreq, - log_level='INFO') - - diagnostic_fields = [Sum('D', 'topography')] - - state = State(mesh, - dt=dt, - output=output, - parameters=parameters, - diagnostic_fields=diagnostic_fields) - + # Equation Omega = parameters.Omega fexpr = 2*Omega*x[2]/R theta, lamda = latlon_coords(mesh) @@ -62,11 +58,31 @@ rsq = Min(R0sq, lsq+thsq) r = sqrt(rsq) bexpr = 2000 * (1 - r/R0) - eqns = ShallowWaterEquations(state, "BDM", 1, fexpr=fexpr, bexpr=bexpr) + eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, bexpr=bexpr) + + # I/O + dirname = "williamson_5_ref%s_dt%s" % (ref_level, dt) + dumpfreq = int(tmax / (ndumps*dt)) + output = OutputParameters(dirname=dirname, + dumplist_latlon=['D'], + dumpfreq=dumpfreq, + log_level='INFO') + diagnostic_fields = [Sum('D', 'topography')] + io = IO(domain, output, diagnostic_fields=diagnostic_fields) + + # Transport schemes + transported_fields = [ImplicitMidpoint(domain, "u"), + SSPRK3(domain, "D")] + + # Time stepper + stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # - # interpolate initial conditions - u0 = state.fields('u') - D0 = state.fields('D') + u0 = stepper.fields('u') + D0 = stepper.fields('D') u_max = 20. # Maximum amplitude of the zonal wind (m/s) uexpr = as_vector([-u_max*x[1]/R, u_max*x[0]/R, 0.0]) g = parameters.g @@ -76,10 +92,11 @@ u0.project(uexpr) D0.interpolate(Dexpr) - transported_fields = [ImplicitMidpoint(state, "u"), - SSPRK3(state, "D")] + Dbar = Function(D0.function_space()).assign(H) + stepper.set_reference_profiles([('D', Dbar)]) - # build time stepper - stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields) + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # stepper.run(t=0, tmax=tmax) diff --git a/gusto/__init__.py b/gusto/__init__.py index 2b1d6c40c..dabfcdf2a 100644 --- a/gusto/__init__.py +++ b/gusto/__init__.py @@ -1,19 +1,19 @@ from gusto.active_tracers import * # noqa from gusto.configuration import * # noqa +from gusto.domain import * # noqa from gusto.diagnostics import * # noqa from gusto.diffusion import * # noqa from gusto.equations import * # noqa -from gusto.fields import * # noqa from gusto.fml import * # noqa from gusto.forcing import * # noqa from gusto.initialisation_tools import * # noqa +from gusto.io import * # noqa from gusto.labels import * # noqa from gusto.limiters import * # noqa from gusto.linear_solvers import * # noqa from gusto.physics import * # noqa from gusto.preconditioners import * # noqa from gusto.recovery import * # noqa -from gusto.state import * # noqa from gusto.time_discretisation import * # noqa from gusto.timeloop import * # noqa from gusto.transport_forms import * # noqa diff --git a/gusto/configuration.py b/gusto/configuration.py index 6b97638ac..dc6efa5e7 100644 --- a/gusto/configuration.py +++ b/gusto/configuration.py @@ -116,10 +116,6 @@ class OutputParameters(Configuration): #: TODO: Should the output fields be interpolated or projected to #: a linear space? Default is interpolation. project_fields = False - #: List of fields to dump error fields for steady state simulation - steady_state_error_fields = [] - #: List of fields for computing perturbations from the initial state - perturbation_fields = [] #: List of ordered pairs (name, points) where name is the field # name and points is the points at which to dump them point_data = [] diff --git a/gusto/diagnostics.py b/gusto/diagnostics.py index e6a8168b3..34ac90fd4 100644 --- a/gusto/diagnostics.py +++ b/gusto/diagnostics.py @@ -1,14 +1,18 @@ """Common diagnostic fields.""" from firedrake import op2, assemble, dot, dx, FunctionSpace, Function, sqrt, \ - TestFunction, TrialFunction, Constant, grad, inner, \ + TestFunction, TrialFunction, Constant, grad, inner, curl, \ LinearVariationalProblem, LinearVariationalSolver, FacetNormal, \ - ds, ds_b, ds_v, ds_t, dS_v, div, avg, jump, DirichletBC, \ - TensorFunctionSpace, SpatialCoordinate, VectorFunctionSpace, as_vector + ds_b, ds_v, ds_t, dS_v, div, avg, jump, \ + TensorFunctionSpace, SpatialCoordinate, as_vector, \ + Projector, Interpolator +from firedrake.assign import Assigner from abc import ABCMeta, abstractmethod, abstractproperty -from gusto import thermodynamics +import gusto.thermodynamics as tde from gusto.recovery import Recoverer, BoundaryMethod +from gusto.equations import CompressibleEulerEquations +from gusto.active_tracers import TracerVariableType, Phases import numpy as np __all__ = ["Diagnostics", "CourantNumber", "VelocityX", "VelocityZ", "VelocityY", "Gradient", @@ -130,301 +134,330 @@ def total(f): class DiagnosticField(object, metaclass=ABCMeta): """Base object to represent diagnostic fields for outputting.""" - def __init__(self, required_fields=()): + def __init__(self, space=None, method='interpolate', required_fields=()): """ Args: + space (:class:`FunctionSpace`, optional): the function space to + evaluate the diagnostic field in. Defaults to None, in which + case a default space will be chosen for this diagnostic. + method (str, optional): a string specifying the method of evaluation + for this diagnostic. Valid options are 'interpolate', 'project', + 'assign' and 'solve'. Defaults to 'interpolate'. required_fields (tuple, optional): tuple of names of the fields that are required for the computation of this diagnostic field. Defaults to (). """ + assert method in ['interpolate', 'project', 'solve', 'assign'], \ + f'Invalid evaluation method {self.method} for diagnostic {self.name}' + self._initialised = False self.required_fields = required_fields + self.space = space + self.method = method + self.expr = None + + # Property to allow graceful failures if solve method not valid + if not hasattr(self, "solve_implemented"): + self.solve_implemented = False + + if method == 'solve' and not self.solve_implemented: + raise NotImplementedError(f'Solve method has not been implemented for diagnostic {self.name}') @abstractproperty def name(self): """The name of this diagnostic field""" pass - def setup(self, state, space=None): + @abstractmethod + def setup(self, domain, state_fields, space=None): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. space (:class:`FunctionSpace`, optional): the function space for the diagnostic field to be computed in. Defaults to None, in which case the space will be DG0. """ + if not self._initialised: - if space is None: - space = state.spaces("DG0", "DG", 0) - self.field = state.fields(self.name, space, pickup=False) - self._initialised = True + if self.space is None: + if space is None: + space = domain.spaces("DG0", "DG", 0) + self.space = space + else: + space = self.space - @abstractmethod - def compute(self, state): - """ - Compute the diagnostic field from the current state. + # Add space to domain + assert space.name is not None, \ + f'Diagnostics {self.name} is using a function space which does not have a name' + domain.spaces(space.name, V=space) - Args: - state (:class:`State`): the model's state. - """ - pass + self.field = state_fields(self.name, space=space, dump=True, pickup=False) - def __call__(self, state): - """ - Compute the diagnostic field from the current state. + if self.method != 'solve': + assert self.expr is not None, \ + f"The expression for diagnostic {self.name} has not been specified" - Args: - state (:class:`State`): the model's state. - """ - return self.compute(state) + # Solve method must be declared in diagnostic's own setup routine + if self.method == 'interpolate': + self.evaluator = Interpolator(self.expr, self.field) + elif self.method == 'project': + self.evaluator = Projector(self.expr, self.field) + elif self.method == 'assign': + self.evaluator = Assigner(self.field, self.expr) + + self._initialised = True + + def compute(self): + """Compute the diagnostic field from the current state.""" + + if self.method == 'interpolate': + self.evaluator.interpolate() + elif self.method == 'assign': + self.evaluator.assign() + elif self.method == 'project': + self.evaluator.project() + elif self.method == 'solve': + self.evaluator.solve() + + def __call__(self): + """Return the diagnostic field computed from the current state.""" + self.compute() + return self.field class CourantNumber(DiagnosticField): """Dimensionless Courant number diagnostic field.""" name = "CourantNumber" - def setup(self, state): + def setup(self, domain, state_fields): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - if not self._initialised: - super(CourantNumber, self).setup(state) - # set up area computation - V = state.spaces("DG0") - test = TestFunction(V) - self.area = Function(V) - assemble(test*dx, tensor=self.area) - def compute(self, state): - """ - Compute and return the diagnostic field from the current state. + # set up area computation + V = domain.spaces("DG0", "DG", 0) + test = TestFunction(V) + self.area = Function(V) + assemble(test*dx, tensor=self.area) + u = state_fields("u") - Args: - state (:class:`State`): the model's state. + self.expr = sqrt(dot(u, u))/sqrt(self.area)*domain.dt - Returns: - :class:`Function`: the diagnostic field. - """ - u = state.fields("u") - dt = Constant(state.dt) - return self.field.project(sqrt(dot(u, u))/sqrt(self.area)*dt) + super(CourantNumber, self).setup(domain, state_fields) +# TODO: unify all component diagnostics class VelocityX(DiagnosticField): """The geocentric Cartesian X component of the velocity field.""" name = "VelocityX" - def setup(self, state): + def setup(self, domain, state_fields): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - if not self._initialised: - space = state.spaces("CG1", "CG", 1) - super(VelocityX, self).setup(state, space=space) - - def compute(self, state): - """ - Compute and return the diagnostic field from the current state. - - Args: - state (:class:`State`): the model's state. - - Returns: - :class:`Function`: the diagnostic field. - """ - u = state.fields("u") - uh = u[0] - return self.field.interpolate(uh) + u = state_fields("u") + self.expr = u[0] + super(VelocityX, self).setup(domain, state_fields) class VelocityZ(DiagnosticField): """The geocentric Cartesian Z component of the velocity field.""" name = "VelocityZ" - def setup(self, state): + def setup(self, domain, state_fields): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. - """ - if not self._initialised: - space = state.spaces("CG1", "CG", 1) - super(VelocityZ, self).setup(state, space=space) - - def compute(self, state): + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - Compute and return the diagnostic field from the current state. - - Args: - state (:class:`State`): the model's state. - - Returns: - :class:`Function`: the diagnostic field. - """ - u = state.fields("u") - w = u[u.geometric_dimension() - 1] - return self.field.interpolate(w) + u = state_fields("u") + self.expr = u[u.geometric_dimension() - 1] + super(VelocityZ, self).setup(domain, state_fields) class VelocityY(DiagnosticField): """The geocentric Cartesian Y component of the velocity field.""" name = "VelocityY" - def setup(self, state): + def setup(self, domain, state_fields): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - if not self._initialised: - space = state.spaces("CG1", "CG", 1) - super(VelocityY, self).setup(state, space=space) - - def compute(self, state): - """ - Compute and return the diagnostic field from the current state. - - Args: - state (:class:`State`): the model's state. - - Returns: - :class:`Function`: the diagnostic field. - """ - u = state.fields("u") - v = u[1] - return self.field.interpolate(v) + u = state_fields("u") + self.expr = u[1] + super(VelocityY, self).setup(domain, state_fields) class Gradient(DiagnosticField): """Diagnostic for computing the gradient of fields.""" - def __init__(self, name): + def __init__(self, name, space=None, method='solve'): """ Args: name (str): name of the field to compute the gradient of. + space (:class:`FunctionSpace`, optional): the function space to + evaluate the diagnostic field in. Defaults to None, in which + case a default space will be chosen for this diagnostic. + method (str, optional): a string specifying the method of evaluation + for this diagnostic. Valid options are 'interpolate', 'project', + 'assign' and 'solve'. Defaults to 'solve'. """ - super().__init__() self.fname = name + self.solve_implemented = True + super().__init__(space=space, method=method, required_fields=(name,)) @property def name(self): """Gives the name of this diagnostic field.""" return self.fname+"_gradient" - def setup(self, state): + def setup(self, domain, state_fields): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. - """ - if not self._initialised: - mesh_dim = state.mesh.geometric_dimension() - try: - field_dim = state.fields(self.fname).ufl_shape[0] - except IndexError: - field_dim = 1 - shape = (mesh_dim, ) * field_dim - space = TensorFunctionSpace(state.mesh, "CG", 1, shape=shape) - super().setup(state, space=space) - - f = state.fields(self.fname) - test = TestFunction(space) - trial = TrialFunction(space) - n = FacetNormal(state.mesh) - a = inner(test, trial)*dx - L = -inner(div(test), f)*dx - if space.extruded: - L += dot(dot(test, n), f)*(ds_t + ds_b) - prob = LinearVariationalProblem(a, L, self.field) - self.solver = LinearVariationalSolver(prob) - - def compute(self, state): + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - Compute and return the diagnostic field from the current state. + f = state_fields(self.fname) - Args: - state (:class:`State`): the model's state. + mesh_dim = domain.mesh.geometric_dimension() + try: + field_dim = state_fields(self.fname).ufl_shape[0] + except IndexError: + field_dim = 1 + shape = (mesh_dim, ) * field_dim + space = TensorFunctionSpace(domain.mesh, "CG", 1, shape=shape, name=f'Tensor{field_dim}_CG1') - Returns: - :class:`Function`: the diagnostic field. - """ - self.solver.solve() - return self.field + if self.method != 'solve': + self.expr = grad(f) + + super().setup(domain, state_fields, space=space) + + # Set up problem now that self.field has been set up + if self.method == 'solve': + test = TestFunction(space) + trial = TrialFunction(space) + n = FacetNormal(domain.mesh) + a = inner(test, trial)*dx + L = -inner(div(test), f)*dx + if space.extruded: + L += dot(dot(test, n), f)*(ds_t + ds_b) + prob = LinearVariationalProblem(a, L, self.field) + self.evaluator = LinearVariationalSolver(prob) class Divergence(DiagnosticField): """Diagnostic for computing the divergence of vector-valued fields.""" - name = "Divergence" - - def setup(self, state): + def __init__(self, name='u', space=None, method='interpolate'): """ - Sets up the :class:`Function` for the diagnostic field. - Args: - state (:class:`State`): the model's state. + name (str, optional): name of the field to compute the gradient of. + Defaults to 'u', in which case this takes the divergence of the + wind field. + space (:class:`FunctionSpace`, optional): the function space to + evaluate the diagnostic field in. Defaults to None, in which + case the default space is the domain's DG space. + method (str, optional): a string specifying the method of evaluation + for this diagnostic. Valid options are 'interpolate', 'project', + 'assign' and 'solve'. Defaults to 'interpolate'. """ - if not self._initialised: - space = state.spaces("DG1", "DG", 1) - super(Divergence, self).setup(state, space=space) + self.fname = name + super().__init__(space=space, method=method, required_fields=(self.fname,)) + + @property + def name(self): + """Gives the name of this diagnostic field.""" + return self.fname+"_divergence" - def compute(self, state): + def setup(self, domain, state_fields): """ - Compute and return the diagnostic field from the current state. + Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. - - Returns: - :class:`Function`: the diagnostic field. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - u = state.fields("u") - return self.field.interpolate(div(u)) + f = state_fields(self.fname) + self.expr = div(f) + space = domain.spaces("DG") + super(Divergence, self).setup(domain, state_fields, space=space) class SphericalComponent(DiagnosticField): """Base diagnostic for computing spherical-polar components of fields.""" - def __init__(self, name): + def __init__(self, name, space=None, method='interpolate'): """ Args: name (str): name of the field to compute the component of. + space (:class:`FunctionSpace`, optional): the function space to + evaluate the diagnostic field in. Defaults to None, in which + case the default space is the domain's DG space. + method (str, optional): a string specifying the method of evaluation + for this diagnostic. Valid options are 'interpolate', 'project', + 'assign' and 'solve'. Defaults to 'interpolate'. """ - super().__init__() self.fname = name + super().__init__(space=space, method=method, required_fields=(name,)) - def setup(self, state): + # TODO: these routines must be moved to somewhere more available generally + # (e.g. initialisation tools?) + def _spherical_polar_unit_vectors(self, domain): """ - Sets up the :class:`Function` for the diagnostic field. + Generate ufl expressions for the spherical polar unit vectors. Args: - state (:class:`State`): the model's state. + domain (:class:`Domain`): the model's domain, containing its mesh. + + Returns: + tuple of (:class:`ufl.Expr`): the zonal, meridional and radial unit + vectors. """ - if not self._initialised: - # check geometric dimension is 3D - if state.mesh.geometric_dimension() != 3: - raise ValueError('Spherical components only work when the geometric dimension is 3!') - space = FunctionSpace(state.mesh, "CG", 1) - super().setup(state, space=space) - - V = VectorFunctionSpace(state.mesh, "CG", 1) - self.x, self.y, self.z = SpatialCoordinate(state.mesh) - self.x_hat = Function(V).interpolate(Constant(as_vector([1.0, 0.0, 0.0]))) - self.y_hat = Function(V).interpolate(Constant(as_vector([0.0, 1.0, 0.0]))) - self.z_hat = Function(V).interpolate(Constant(as_vector([0.0, 0.0, 1.0]))) - self.R = sqrt(self.x**2 + self.y**2) # distance from z axis - self.r = sqrt(self.x**2 + self.y**2 + self.z**2) # distance from origin - self.f = state.fields(self.fname) - if np.prod(self.f.ufl_shape) != 3: + x, y, z = SpatialCoordinate(domain.mesh) + x_hat = Constant(as_vector([1.0, 0.0, 0.0])) + y_hat = Constant(as_vector([0.0, 1.0, 0.0])) + z_hat = Constant(as_vector([0.0, 0.0, 1.0])) + R = sqrt(x**2 + y**2) # distance from z axis + r = sqrt(x**2 + y**2 + z**2) # distance from origin + + lambda_hat = (x * y_hat - y * x_hat) / R + phi_hat = (-x*z/R * x_hat - y*z/R * y_hat + R * z_hat) / r + r_hat = (x * x_hat + y * y_hat + z * z_hat) / r + + return lambda_hat, phi_hat, r_hat + + def _check_args(self, domain, field): + """ + Checks the validity of the domain and field for taking the spherical + component diagnostic. + + Args: + domain (:class:`Domain`): the model's domain object. + field (:class:`Function`): the field to take the component of. + """ + + # check geometric dimension is 3D + if domain.mesh.geometric_dimension() != 3: + raise ValueError('Spherical components only work when the geometric dimension is 3!') + + if np.prod(field.ufl_shape) != 3: raise ValueError('Components can only be found of a vector function space in 3D.') @@ -435,20 +468,19 @@ def name(self): """Gives the name of this diagnostic field.""" return self.fname+"_meridional" - def compute(self, state): + def setup(self, domain, state_fields): """ - Compute and return the diagnostic field from the current state. + Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. - - Returns: - :class:`Function`: the diagnostic field. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - lambda_hat = (-self.x * self.z * self.x_hat / self.R - - self.y * self.z * self.y_hat / self.R - + self.R * self.z_hat) / self.r - return self.field.project(dot(self.f, lambda_hat)) + f = state_fields(self.fname) + self._check_args(domain, f) + _, phi_hat, _ = self._spherical_polar_unit_vectors(domain) + self.expr = dot(f, phi_hat) + super().setup(domain, state_fields) class ZonalComponent(SphericalComponent): @@ -458,18 +490,19 @@ def name(self): """Gives the name of this diagnostic field.""" return self.fname+"_zonal" - def compute(self, state): + def setup(self, domain, state_fields): """ - Compute and return the diagnostic field from the current state. + Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. - - Returns: - :class:`Function`: the diagnostic field. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - phi_hat = (self.x * self.y_hat - self.y * self.x_hat) / self.R - return self.field.project(dot(self.f, phi_hat)) + f = state_fields(self.fname) + self._check_args(domain, f) + lambda_hat, _, _ = self._spherical_polar_unit_vectors(domain) + self.expr = dot(f, lambda_hat) + super().setup(domain, state_fields) class RadialComponent(SphericalComponent): @@ -479,67 +512,65 @@ def name(self): """Gives the name of this diagnostic field.""" return self.fname+"_radial" - def compute(self, state): + def setup(self, domain, state_fields): """ - Compute and return the diagnostic field from the current state. + Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. - - Returns: - :class:`Function`: the diagnostic field. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - r_hat = (self.x * self.x_hat + self.y * self.y_hat + self.z * self.z_hat) / self.r - return self.field.project(dot(self.f, r_hat)) + f = state_fields(self.fname) + self._check_args(domain, f) + _, _, r_hat = self._spherical_polar_unit_vectors(domain) + self.expr = dot(f, r_hat) + super().setup(domain, state_fields) class RichardsonNumber(DiagnosticField): """Dimensionless Richardson number diagnostic field.""" name = "RichardsonNumber" - def __init__(self, density_field, factor=1.): + def __init__(self, density_field, factor=1., space=None, method='interpolate'): u""" Args: - density_field (:class:`Function`): the density field. + density_field (str): the name of the density field. factor (float, optional): a factor to multiply the Brunt-Väisälä frequency by. Defaults to 1. - """ - super().__init__(required_fields=(density_field, "u_gradient")) + space (:class:`FunctionSpace`, optional): the function space to + evaluate the diagnostic field in. Defaults to None, in which + case a default space will be chosen for this diagnostic. + method (str, optional): a string specifying the method of evaluation + for this diagnostic. Valid options are 'interpolate', 'project', + 'assign' and 'solve'. Defaults to 'interpolate'. + """ + super().__init__(space=space, method=method, required_fields=(density_field, "u_gradient")) self.density_field = density_field self.factor = Constant(factor) - def setup(self, state): + def setup(self, domain, state_fields): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ rho_grad = self.density_field+"_gradient" - super().setup(state) - self.grad_density = state.fields(rho_grad) - self.gradu = state.fields("u_gradient") - - def compute(self, state): - """ - Compute and return the diagnostic field from the current state. + grad_density = state_fields(rho_grad) + gradu = state_fields("u_gradient") - Args: - state (:class:`State`): the model's state. - - Returns: - :class:`Function`: the diagnostic field. - """ denom = 0. - z_dim = state.mesh.geometric_dimension() - 1 - u_dim = state.fields("u").ufl_shape[0] + z_dim = domain.mesh.geometric_dimension() - 1 + u_dim = state_fields("u").ufl_shape[0] for i in range(u_dim-1): - denom += self.gradu[i, z_dim]**2 - Nsq = self.factor*self.grad_density[z_dim] - self.field.interpolate(Nsq/denom) - return self.field + denom += gradu[i, z_dim]**2 + Nsq = self.factor*grad_density[z_dim] + self.expr = Nsq/denom + super().setup(domain, state_fields) +# TODO: unify all energy diagnostics -- should be based on equation class Energy(DiagnosticField): """Base diagnostic field for computing energy density fields.""" def kinetic(self, u, factor=None): @@ -565,70 +596,123 @@ class KineticEnergy(Energy): """Diagnostic kinetic energy density.""" name = "KineticEnergy" - def compute(self, state): + def __init__(self, space=None, method='interpolate'): """ - Compute and return the diagnostic field from the current state. - Args: - state (:class:`State`): the model's state. + space (:class:`FunctionSpace`, optional): the function space to + evaluate the diagnostic field in. Defaults to None, in which + case a default space will be chosen for this diagnostic. + method (str, optional): a string specifying the method of evaluation + for this diagnostic. Valid options are 'interpolate', 'project', + 'assign' and 'solve'. Defaults to 'interpolate'. + """ + super().__init__(space=space, method=method, required_fields=("u")) - Returns: - :class:`Function`: the diagnostic field. + def setup(self, domain, state_fields): """ - u = state.fields("u") - energy = self.kinetic(u) - return self.field.interpolate(energy) + Sets up the :class:`Function` for the diagnostic field. + + Args: + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. + """ + u = state_fields("u") + self.expr = self.kinetic(u) + super().setup(domain, state_fields) class ShallowWaterKineticEnergy(Energy): """Diagnostic shallow-water kinetic energy density.""" name = "ShallowWaterKineticEnergy" - def compute(self, state): + def __init__(self, space=None, method='interpolate'): """ - Compute and return the diagnostic field from the current state. - Args: - state (:class:`State`): the model's state. + space (:class:`FunctionSpace`, optional): the function space to + evaluate the diagnostic field in. Defaults to None, in which + case a default space will be chosen for this diagnostic. + method (str, optional): a string specifying the method of evaluation + for this diagnostic. Valid options are 'interpolate', 'project', + 'assign' and 'solve'. Defaults to 'interpolate'. + """ + super().__init__(space=space, method=method, required_fields=("D", "u")) - Returns: - :class:`Function`: the diagnostic field. + def setup(self, domain, state_fields): """ - u = state.fields("u") - D = state.fields("D") - energy = self.kinetic(u, D) - return self.field.interpolate(energy) + Sets up the :class:`Function` for the diagnostic field. + + Args: + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. + """ + u = state_fields("u") + D = state_fields("D") + self.expr = self.kinetic(u, D) + super().setup(domain, state_fields) class ShallowWaterPotentialEnergy(Energy): """Diagnostic shallow-water potential energy density.""" name = "ShallowWaterPotentialEnergy" - def compute(self, state): + def __init__(self, parameters, space=None, method='interpolate'): """ - Compute and return the diagnostic field from the current state. - Args: - state (:class:`State`): the model's state. + parameters (:class:`ShallowWaterParameters`): the configuration + object containing the physical parameters for this equation. + space (:class:`FunctionSpace`, optional): the function space to + evaluate the diagnostic field in. Defaults to None, in which + case a default space will be chosen for this diagnostic. + method (str, optional): a string specifying the method of evaluation + for this diagnostic. Valid options are 'interpolate', 'project', + 'assign' and 'solve'. Defaults to 'interpolate'. + """ + self.parameters = parameters + super().__init__(space=space, method=method, required_fields=("D")) - Returns: - :class:`Function`: the diagnostic field. + def setup(self, domain, state_fields): """ - g = state.parameters.g - D = state.fields("D") - energy = 0.5*g*D**2 - return self.field.interpolate(energy) + Sets up the :class:`Function` for the diagnostic field. + + Args: + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. + """ + g = self.parameters.g + D = state_fields("D") + self.expr = 0.5*g*D**2 + super().setup(domain, state_fields) class ShallowWaterPotentialEnstrophy(DiagnosticField): """Diagnostic (dry) compressible kinetic energy density.""" - def __init__(self, base_field_name="PotentialVorticity"): + def __init__(self, base_field_name="PotentialVorticity", space=None, + method='interpolate'): """ Args: base_field_name (str, optional): the base potential vorticity field to compute the enstrophy from. Defaults to "PotentialVorticity". - """ - super().__init__() + space (:class:`FunctionSpace`, optional): the function space to + evaluate the diagnostic field in. Defaults to None, in which + case a default space will be chosen for this diagnostic. + method (str, optional): a string specifying the method of evaluation + for this diagnostic. Valid options are 'interpolate', 'project', + 'assign' and 'solve'. Defaults to 'interpolate'. + """ + base_enstrophy_names = ["PotentialVorticity", "RelativeVorticity", "AbsoluteVorticity"] + if base_field_name not in base_enstrophy_names: + raise ValueError( + f"Don't know how to compute enstrophy with base_field_name={base_field_name};" + + f"base_field_name should be one of {base_enstrophy_names}") + # Work out required fields + if base_field_name in ["PotentialVorticity", "AbsoluteVorticity"]: + required_fields = (base_field_name, "D") + elif base_field_name == "RelativeVorticity": + required_fields = (base_field_name, "D", "coriolis") + else: + raise NotImplementedError(f'Enstrophy with vorticity {base_field_name} not implemented') + + super().__init__(space=space, method=method, required_fields=required_fields) self.base_field_name = base_field_name @property @@ -637,211 +721,210 @@ def name(self): base_name = "SWPotentialEnstrophy" return "_from_".join((base_name, self.base_field_name)) - def compute(self, state): + def setup(self, domain, state_fields): """ - Compute and return the diagnostic field from the current state. + Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. - - Returns: - :class:`Function`: the diagnostic field. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ if self.base_field_name == "PotentialVorticity": - pv = state.fields("PotentialVorticity") - D = state.fields("D") - enstrophy = 0.5*pv**2*D + pv = state_fields("PotentialVorticity") + D = state_fields("D") + self.expr = 0.5*pv**2*D elif self.base_field_name == "RelativeVorticity": - zeta = state.fields("RelativeVorticity") - D = state.fields("D") - f = state.fields("coriolis") - enstrophy = 0.5*(zeta + f)**2/D + zeta = state_fields("RelativeVorticity") + D = state_fields("D") + f = state_fields("coriolis") + self.expr = 0.5*(zeta + f)**2/D elif self.base_field_name == "AbsoluteVorticity": - zeta_abs = state.fields("AbsoluteVorticity") - D = state.fields("D") - enstrophy = 0.5*(zeta_abs)**2/D + zeta_abs = state_fields("AbsoluteVorticity") + D = state_fields("D") + self.expr = 0.5*(zeta_abs)**2/D else: - raise ValueError("Don't know how to compute enstrophy with base_field_name=%s; base_field_name should be %s %s or %s." % (self.base_field_name, "RelativeVorticity", "AbsoluteVorticity", "PotentialVorticity")) - return self.field.interpolate(enstrophy) + raise NotImplementedError(f'Enstrophy with {self.base_field_name} not implemented') + super().setup(domain, state_fields) class CompressibleKineticEnergy(Energy): """Diagnostic (dry) compressible kinetic energy density.""" name = "CompressibleKineticEnergy" - def compute(self, state): + def __init__(self, space=None, method='interpolate'): + """ + Args: + space (:class:`FunctionSpace`, optional): the function space to + evaluate the diagnostic field in. Defaults to None, in which + case a default space will be chosen for this diagnostic. + method (str, optional): a string specifying the method of evaluation + for this diagnostic. Valid options are 'interpolate', 'project', + 'assign' and 'solve'. Defaults to 'interpolate'. + """ + super().__init__(space=space, method=method, required_fields=("rho", "u")) + + def compute(self, eqn): """ Compute and return the diagnostic field from the current state. Args: - state (:class:`State`): the model's state. + eqn (:class:`PrognosticEquation`): the model's equation. Returns: :class:`Function`: the diagnostic field. """ - u = state.fields("u") - rho = state.fields("rho") + u = eqn.fields("u") + rho = eqn.fields("rho") energy = self.kinetic(u, rho) return self.field.interpolate(energy) class Exner(DiagnosticField): """The diagnostic Exner pressure field.""" - def __init__(self, reference=False): + def __init__(self, parameters, reference=False, space=None, method='interpolate'): """ Args: + parameters (:class:`CompressibleParameters`): the configuration + object containing the physical parameters for this equation. reference (bool, optional): whether to compute the reference Exner pressure field or not. Defaults to False. - """ - super(Exner, self).__init__() + space (:class:`FunctionSpace`, optional): the function space to + evaluate the diagnostic field in. Defaults to None, in which + case a default space will be chosen for this diagnostic. + method (str, optional): a string specifying the method of evaluation + for this diagnostic. Valid options are 'interpolate', 'project', + 'assign' and 'solve'. Defaults to 'interpolate'. + """ + self.parameters = parameters self.reference = reference if reference: - self.rho_name = "rhobar" - self.theta_name = "thetabar" + self.rho_name = "rho_bar" + self.theta_name = "theta_bar" else: self.rho_name = "rho" self.theta_name = "theta" + super().__init__(space=space, method=method, required_fields=(self.rho_name, self.theta_name)) @property def name(self): """Gives the name of this diagnostic field.""" if self.reference: - return "Exnerbar" + return "Exner_bar" else: return "Exner" - def setup(self, state): + def setup(self, domain, state_fields): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - if not self._initialised: - space = state.spaces("CG1", "CG", 1) - super(Exner, self).setup(state, space=space) - - def compute(self, state): - """ - Compute and return the diagnostic field from the current state. - - Args: - state (:class:`State`): the model's state. - - Returns: - :class:`Function`: the diagnostic field. - """ - rho = state.fields(self.rho_name) - theta = state.fields(self.theta_name) - exner = thermodynamics.exner_pressure(state.parameters, rho, theta) - return self.field.interpolate(exner) + rho = state_fields(self.rho_name) + theta = state_fields(self.theta_name) + self.expr = tde.exner_pressure(self.parameters, rho, theta) + super().setup(domain, state_fields) class Sum(DiagnosticField): """Base diagnostic for computing the sum of two fields.""" - def __init__(self, field1, field2): + def __init__(self, field_name1, field_name2): """ Args: - field1 (:class:`Function`): one field to be added. - field2 (:class:`Function`): the other field to be added. + field_name1 (str): the name of one field to be added. + field_name2 (str): the name of the other field to be added. """ - super().__init__(required_fields=(field1, field2)) - self.field1 = field1 - self.field2 = field2 + super().__init__(method='assign', required_fields=(field_name1, field_name2)) + self.field_name1 = field_name1 + self.field_name2 = field_name2 @property def name(self): """Gives the name of this diagnostic field.""" - return self.field1+"_plus_"+self.field2 + return self.field_name1+"_plus_"+self.field_name2 - def setup(self, state): + def setup(self, domain, state_fields): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - if not self._initialised: - space = state.fields(self.field1).function_space() - super(Sum, self).setup(state, space=space) - - def compute(self, state): - """ - Compute and return the diagnostic field from the current state. - - Args: - state (:class:`State`): the model's state. - - Returns: - :class:`Function`: the diagnostic field. - """ - field1 = state.fields(self.field1) - field2 = state.fields(self.field2) - return self.field.assign(field1 + field2) + field1 = state_fields(self.field_name1) + field2 = state_fields(self.field_name2) + space = field1.function_space() + self.expr = field1 + field2 + super(Sum, self).setup(domain, state_fields, space=space) class Difference(DiagnosticField): """Base diagnostic for calculating the difference between two fields.""" - def __init__(self, field1, field2): + def __init__(self, field_name1, field_name2): """ Args: - field1 (:class:`Function`): the field to be subtracted from. - field2 (:class:`Function`): the field to be subtracted. + field_name1 (str): the name of the field to be subtracted from. + field_name2 (str): the name of the field to be subtracted. """ - super().__init__(required_fields=(field1, field2)) - self.field1 = field1 - self.field2 = field2 + super().__init__(method='assign', required_fields=(field_name1, field_name2)) + self.field_name1 = field_name1 + self.field_name2 = field_name2 @property def name(self): """Gives the name of this diagnostic field.""" - return self.field1+"_minus_"+self.field2 + return self.field_name1+"_minus_"+self.field_name2 - def setup(self, state): + def setup(self, domain, state_fields): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. - """ - if not self._initialised: - space = state.fields(self.field1).function_space() - super(Difference, self).setup(state, space=space) - - def compute(self, state): + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - Compute and return the diagnostic field from the current state. - Args: - state (:class:`State`): the model's state. - - Returns: - :class:`Function`: the diagnostic field. - """ - field1 = state.fields(self.field1) - field2 = state.fields(self.field2) - return self.field.assign(field1 - field2) + field1 = state_fields(self.field_name1) + field2 = state_fields(self.field_name2) + self.expr = field1 - field2 + space = field1.function_space() + super(Difference, self).setup(domain, state_fields, space=space) class SteadyStateError(Difference): """Base diagnostic for computing the steady-state error in a field.""" - def __init__(self, state, name): + def __init__(self, name): """ Args: - state (:class:`State`): the model's state. - name (str): name of the field to take the perturbation of. + name (str): name of the field to take the steady-state error of. """ - DiagnosticField.__init__(self) - self.field1 = name - self.field2 = name+'_init' - field1 = state.fields(name) - field2 = state.fields(self.field2, field1.function_space()) + self.field_name1 = name + self.field_name2 = name+'_init' + DiagnosticField.__init__(self, method='assign', required_fields=(name)) + + def setup(self, domain, state_fields): + """ + Sets up the :class:`Function` for the diagnostic field. + + Args: + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. + """ + # Create and store initial field + field1 = state_fields(self.field_name1) + field2 = state_fields(self.field_name2, space=field1.function_space(), dump=False) + + # TODO: when checkpointing, the initial field should either be picked up + # or computed again (picking up can be easily specified if we change the line above) field2.assign(field1) + super(SteadyStateError, self).setup(domain, state_fields) + @property def name(self): """Gives the name of this diagnostic field.""" - return self.field1+"_error" + return self.field_name1+"_error" class Perturbation(Difference): @@ -851,465 +934,549 @@ def __init__(self, name): Args: name (str): name of the field to take the perturbation of. """ - self.field1 = name - self.field2 = name+'bar' - DiagnosticField.__init__(self, required_fields=(self.field1, self.field2)) + field_name1 = name + field_name2 = name+'_bar' + Difference.__init__(self, field_name1, field_name2) @property def name(self): """Gives the name of this diagnostic field.""" - return self.field1+"_perturbation" + return self.field_name1+"_perturbation" +# TODO: unify thermodynamic diagnostics class ThermodynamicDiagnostic(DiagnosticField): """Base thermodynamic diagnostic field, computing many common fields.""" - def setup(self, state): + def __init__(self, equations, space=None, method='interpolate'): + """ + Args: + equations (:class:`PrognosticEquationSet`): the equation set being + solved by the model. + space (:class:`FunctionSpace`, optional): the function space to + evaluate the diagnostic field in. Defaults to None, in which + case a default space will be chosen for this diagnostic. + method (str, optional): a string specifying the method of evaluation + for this diagnostic. Valid options are 'interpolate', 'project', + 'assign' and 'solve'. Defaults to 'interpolate'. + """ + self.equations = equations + self.parameters = equations.parameters + # Work out required fields + if isinstance(equations, CompressibleEulerEquations): + required_fields = ['rho', 'theta'] + if equations.active_tracers is not None: + for active_tracer in equations.active_tracers: + if active_tracer.chemical == 'H2O': + required_fields.append(active_tracer.name) + else: + raise NotImplementedError(f'Thermodynamic diagnostics not implemented for {type(equations)}') + super().__init__(space=space, method=method, required_fields=tuple(required_fields)) + + def _setup_thermodynamics(self, domain, state_fields): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. - """ - if not self._initialised: - space = state.fields("theta").function_space() - h_deg = space.ufl_element().degree()[0] - v_deg = space.ufl_element().degree()[1]-1 - boundary_method = BoundaryMethod.extruded if (v_deg == 0 and h_deg == 0) else None - super().setup(state, space=space) - - # now let's attach all of our fields - self.u = state.fields("u") - self.rho = state.fields("rho") - self.theta = state.fields("theta") - self.rho_averaged = Function(space) - self.recoverer = Recoverer(self.rho, self.rho_averaged, boundary_method=boundary_method) - try: - self.r_v = state.fields("water_vapour") - except NotImplementedError: - self.r_v = Constant(0.0) - try: - self.r_c = state.fields("cloud_water") - except NotImplementedError: - self.r_c = Constant(0.0) - try: - self.rain = state.fields("rain") - except NotImplementedError: - self.rain = Constant(0.0) - - # now let's store the most common expressions - self.exner = thermodynamics.exner_pressure(state.parameters, self.rho_averaged, self.theta) - self.T = thermodynamics.T(state.parameters, self.theta, self.exner, r_v=self.r_v) - self.p = thermodynamics.p(state.parameters, self.exner) - self.r_l = self.r_c + self.rain - self.r_t = self.r_v + self.r_c + self.rain - - def compute(self, state): - """ - Compute thermodynamic auxiliary fields commonly used by diagnostics. - - Args: - state (:class:`State`): the model's state. - """ + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. + """ + + self.Vtheta = domain.spaces('theta') + h_deg = self.Vtheta.ufl_element().degree()[0] + v_deg = self.Vtheta.ufl_element().degree()[1]-1 + boundary_method = BoundaryMethod.extruded if (v_deg == 0 and h_deg == 0) else None + + # Extract all fields + self.rho = state_fields("rho") + self.theta = state_fields("theta") + # Rho must be averaged to Vtheta + self.rho_averaged = Function(self.Vtheta) + self.recoverer = Recoverer(self.rho, self.rho_averaged, boundary_method=boundary_method) + + zero_expr = Constant(0.0)*self.theta + self.r_v = zero_expr # Water vapour + self.r_l = zero_expr # Liquid water + self.r_t = zero_expr # All water mixing ratios + for active_tracer in self.equations.active_tracers: + if active_tracer.chemical == "H2O": + if active_tracer.variable_type != TracerVariableType.mixing_ratio: + raise NotImplementedError('Only mixing ratio tracers are implemented') + if active_tracer.phase == Phases.gas: + self.r_v += state_fields(active_tracer.name) + elif active_tracer.phase == Phases.liquid: + self.r_l += state_fields(active_tracer.name) + self.r_t += state_fields(active_tracer.name) + + # Store the most common expressions + self.exner = tde.exner_pressure(self.parameters, self.rho_averaged, self.theta) + self.T = tde.T(self.parameters, self.theta, self.exner, r_v=self.r_v) + self.p = tde.p(self.parameters, self.exner) + + def compute(self): + """Compute the thermodynamic diagnostic.""" self.recoverer.project() + super().compute() class Theta_e(ThermodynamicDiagnostic): """The moist equivalent potential temperature diagnostic field.""" name = "Theta_e" - def compute(self, state): + def setup(self, domain, state_fields): """ - Compute and return the diagnostic field from the current state. + Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. - - Returns: - :class:`Function`: the diagnostic field. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - super().compute(state) - - return self.field.interpolate(thermodynamics.theta_e(state.parameters, self.T, self.p, self.r_v, self.r_t)) + self._setup_thermodynamics(domain, state_fields) + self.expr = tde.theta_e(self.parameters, self.T, self.p, self.r_v, self.r_t) + super().setup(domain, state_fields, space=self.Vtheta) class InternalEnergy(ThermodynamicDiagnostic): """The moist compressible internal energy density.""" name = "InternalEnergy" - def compute(self, state): + def setup(self, domain, state_fields): """ - Compute and return the diagnostic field from the current state. + Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. - - Returns: - :class:`Function`: the diagnostic field. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - super().compute(state) - - return self.field.interpolate(thermodynamics.internal_energy(state.parameters, self.rho_averaged, self.T, r_v=self.r_v, r_l=self.r_l)) + self._setup_thermodynamics(domain, state_fields) + self.expr = tde.internal_energy(self.parameters, self.rho_averaged, self.T, r_v=self.r_v, r_l=self.r_l) + super().setup(domain, state_fields, space=self.Vtheta) class PotentialEnergy(ThermodynamicDiagnostic): """The moist compressible potential energy density.""" name = "PotentialEnergy" - def setup(self, state): + def setup(self, domain, state_fields): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. - """ - super().setup(state) - self.x = SpatialCoordinate(state.mesh) - - def compute(self, state): + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - Compute and return the diagnostic field from the current state. - - Args: - state (:class:`State`): the model's state. - - Returns: - :class:`Function`: the diagnostic field. - """ - super().compute(state) - - return self.field.interpolate(self.rho_averaged * (1 + self.r_t) * state.parameters.g * dot(self.x, state.k)) + x = SpatialCoordinate(domain.mesh) + self.expr = self.rho_averaged * (1 + self.r_t) * self.parameters.g * dot(x, domain.k) + super().setup(domain, state_fields, space=domain.spaces("DG")) +# TODO: this needs consolidating with energy diagnostics class ThermodynamicKineticEnergy(ThermodynamicDiagnostic): """The moist compressible kinetic energy density.""" name = "ThermodynamicKineticEnergy" - def compute(self, state): + def __init__(self, equations, space=None, method='interpolate'): + """ + Args: + equations (:class:`PrognosticEquationSet`): the equation set being + solved by the model. + space (:class:`FunctionSpace`, optional): the function space to + evaluate the diagnostic field in. Defaults to None, in which + case a default space will be chosen for this diagnostic. + method (str, optional): a string specifying the method of evaluation + for this diagnostic. Valid options are 'interpolate', 'project', + 'assign' and 'solve'. Defaults to 'interpolate'. + """ + self.equations = equations + self.parameters = equations.parameters + # Work out required fields + if isinstance(equations, CompressibleEulerEquations): + required_fields = ['rho', 'u'] + if equations.active_tracers is not None: + for active_tracer in equations.active_tracers: + if active_tracer.chemical == 'H2O': + required_fields.append(active_tracer.name) + else: + raise NotImplementedError(f'Thermodynamic K.E. not implemented for {type(equations)}') + super().__init__(space=space, method=method, required_fields=tuple(required_fields)) + + def setup(self, domain, state_fields): """ - Compute and return the diagnostic field from the current state. + Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. - - Returns: - :class:`Function`: the diagnostic field. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - super().compute(state) - - return self.field.project(0.5 * self.rho_averaged * (1 + self.r_t) * dot(self.u, self.u)) + u = state_fields('u') + self.expr = 0.5 * self.rho_averaged * (1 + self.r_t) * dot(u, u) + super().setup(domain, state_fields, space=domain.spaces("DG")) class Dewpoint(ThermodynamicDiagnostic): """The dewpoint temperature diagnostic field.""" name = "Dewpoint" - def compute(self, state): + def setup(self, domain, state_fields): """ - Compute and return the diagnostic field from the current state. + Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. - - Returns: - :class:`Function`: the diagnostic field. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - super().compute(state) - - return self.field.interpolate(thermodynamics.T_dew(state.parameters, self.p, self.r_v)) + self._setup_thermodynamics(domain, state_fields) + self.expr = tde.T_dew(self.parameters, self.p, self.r_v) + super().setup(domain, state_fields, space=self.Vtheta) class Temperature(ThermodynamicDiagnostic): """The absolute temperature diagnostic field.""" name = "Temperature" - def compute(self, state): + def setup(self, domain, state_fields): """ - Compute and return the diagnostic field from the current state. + Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. - - Returns: - :class:`Function`: the diagnostic field. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - super().compute(state) - - return self.field.assign(self.T) + self._setup_thermodynamics(domain, state_fields) + self.expr = self.T + super().setup(domain, state_fields, space=self.Vtheta) class Theta_d(ThermodynamicDiagnostic): """The dry potential temperature diagnostic field.""" name = "Theta_d" - def compute(self, state): + def setup(self, domain, state_fields): """ - Compute and return the diagnostic field from the current state. + Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. - - Returns: - :class:`Function`: the diagnostic field. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - super().compute(state) - - return self.field.interpolate(self.theta / (1 + self.r_v * state.parameters.R_v / state.parameters.R_d)) + self._setup_thermodynamics(domain, state_fields) + self.expr = self.theta / (1 + self.r_v * self.parameters.R_v / self.parameters.R_d) + super().setup(domain, state_fields, space=self.Vtheta) class RelativeHumidity(ThermodynamicDiagnostic): """The relative humidity diagnostic field.""" name = "RelativeHumidity" - def compute(self, state): + def setup(self, domain, state_fields): """ - Compute and return the diagnostic field from the current state. + Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. - - Returns: - :class:`Function`: the diagnostic field. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - super().compute(state) - - return self.field.interpolate(thermodynamics.RH(state.parameters, self.r_v, self.T, self.p)) + self._setup_thermodynamics(domain, state_fields) + self.expr = tde.RH(self.parameters, self.r_v, self.T, self.p) + super().setup(domain, state_fields, space=self.Vtheta) class Pressure(ThermodynamicDiagnostic): """The pressure field computed in the 'theta' space.""" name = "Pressure_Vt" - def compute(self, state): + def setup(self, domain, state_fields): """ - Compute and return the diagnostic field from the current state. + Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. - - Returns: - :class:`Function`: the diagnostic field. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - super().compute(state) - - return self.field.assign(self.p) + self._setup_thermodynamics(domain, state_fields) + self.expr = self.p + super().setup(domain, state_fields, space=self.Vtheta) class Exner_Vt(ThermodynamicDiagnostic): """The Exner pressure field computed in the 'theta' space.""" name = "Exner_Vt" - def compute(self, state): + def setup(self, domain, state_fields): """ - Compute and return the diagnostic field from the current state. + Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. - - Returns: - :class:`Function`: the diagnostic field. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - super().compute(state) - - return self.field.assign(self.exner) + self._setup_thermodynamics(domain, state_fields) + self.expr = self.exner + super().setup(domain, state_fields, space=self.Vtheta) +# TODO: this doesn't contain the effects of moisture +# TODO: this has not been implemented for other equation sets class HydrostaticImbalance(DiagnosticField): """Hydrostatic imbalance diagnostic field.""" name = "HydrostaticImbalance" - def setup(self, state): + def __init__(self, equations, space=None, method='interpolate'): + """ + Args: + equations (:class:`PrognosticEquationSet`): the equation set being + solved by the model. + space (:class:`FunctionSpace`, optional): the function space to + evaluate the diagnostic field in. Defaults to None, in which + case a default space will be chosen for this diagnostic. + method (str, optional): a string specifying the method of evaluation + for this diagnostic. Valid options are 'interpolate', 'project', + 'assign' and 'solve'. Defaults to 'interpolate'. + """ + # Work out required fields + if isinstance(equations, CompressibleEulerEquations): + required_fields = ['rho', 'theta', 'rho_bar', 'theta_bar'] + if equations.active_tracers is not None: + for active_tracer in equations.active_tracers: + if active_tracer.chemical == 'H2O': + required_fields.append(active_tracer.name) + self.equations = equations + self.parameters = equations.parameters + else: + raise NotImplementedError(f'Hydrostatic Imbalance not implemented for {type(equations)}') + super().__init__(space=space, method=method, required_fields=required_fields) + + def setup(self, domain, state_fields): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. - """ - if not self._initialised: - Vu = state.spaces("HDiv") - space = FunctionSpace(state.mesh, Vu.ufl_element()._elements[-1]) - super().setup(state, space=space) - rho = state.fields("rho") - rhobar = state.fields("rhobar") - theta = state.fields("theta") - thetabar = state.fields("thetabar") - exner = thermodynamics.exner_pressure(state.parameters, rho, theta) - exnerbar = thermodynamics.exner_pressure(state.parameters, rhobar, thetabar) - - cp = Constant(state.parameters.cp) - n = FacetNormal(state.mesh) - - F = TrialFunction(space) - w = TestFunction(space) - a = inner(w, F)*dx - L = (- cp*div((theta-thetabar)*w)*exnerbar*dx - + cp*jump((theta-thetabar)*w, n)*avg(exnerbar)*dS_v - - cp*div(thetabar*w)*(exner-exnerbar)*dx - + cp*jump(thetabar*w, n)*avg(exner-exnerbar)*dS_v) - - bcs = [DirichletBC(space, 0.0, "bottom"), - DirichletBC(space, 0.0, "top")] - - imbalanceproblem = LinearVariationalProblem(a, L, self.field, bcs=bcs) - self.imbalance_solver = LinearVariationalSolver(imbalanceproblem) - - def compute(self, state): + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - Compute and return the diagnostic field from the current state. + Vu = domain.spaces("HDiv") + rho = state_fields("rho") + rhobar = state_fields("rho_bar") + theta = state_fields("theta") + thetabar = state_fields("theta_bar") + exner = tde.exner_pressure(self.parameters, rho, theta) + exnerbar = tde.exner_pressure(self.parameters, rhobar, thetabar) - Args: - state (:class:`State`): the model's state. + cp = Constant(self.parameters.cp) + n = FacetNormal(domain.mesh) - Returns: - :class:`Function`: the diagnostic field. + # TODO: not sure about this expression! + # Gravity does not appear, and why are there reference profiles? + F = TrialFunction(Vu) + w = TestFunction(Vu) + imbalance = Function(Vu) + a = inner(w, F)*dx + L = (- cp*div((theta-thetabar)*w)*exnerbar*dx + + cp*jump((theta-thetabar)*w, n)*avg(exnerbar)*dS_v + - cp*div(thetabar*w)*(exner-exnerbar)*dx + + cp*jump(thetabar*w, n)*avg(exner-exnerbar)*dS_v) + + bcs = self.equations.bcs['u'] + + imbalanceproblem = LinearVariationalProblem(a, L, imbalance, bcs=bcs) + self.imbalance_solver = LinearVariationalSolver(imbalanceproblem) + self.expr = dot(imbalance, domain.k) + super().setup(domain, state_fields) + + def compute(self): + """Compute and return the diagnostic field from the current state. """ self.imbalance_solver.solve() - return self.field[1] + super().compute() class Precipitation(DiagnosticField): """The total precipitation falling through the domain's bottom surface.""" name = "Precipitation" - def setup(self, state): - """ - Sets up the :class:`Function` for the diagnostic field. - - Args: - state (:class:`State`): the model's state. - """ - if not self._initialised: - space = state.spaces("DG0", "DG", 0) - super().setup(state, space=space) - - rain = state.fields('rain') - rho = state.fields('rho') - v = state.fields('rainfall_velocity') - self.phi = TestFunction(space) - flux = TrialFunction(space) - n = FacetNormal(state.mesh) - un = 0.5 * (dot(v, n) + abs(dot(v, n))) - self.flux = Function(space) - - a = self.phi * flux * dx - L = self.phi * rain * un * rho - if space.extruded: - L = L * (ds_b + ds_t + ds_v) - else: - L = L * ds - - # setup solver - problem = LinearVariationalProblem(a, L, self.flux) - self.solver = LinearVariationalSolver(problem) + def __init__(self): + self.solve_implemented = True + required_fields = ('rain', 'rainfall_velocity', 'rho') + super().__init__(method='solve', required_fields=required_fields) - def compute(self, state): + def setup(self, domain, state_fields): """ - Compute and return the diagnostic field from the current state. + Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. - - Returns: - :class:`Function`: the diagnostic field. - """ + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. + """ + space = domain.spaces("DG0", "DG", 0) + assert space.extruded, 'Cannot compute precipitation on a non-extruded mesh' + rain = state_fields('rain') + rho = state_fields('rho') + v = state_fields('rainfall_velocity') + # Set up problem + self.phi = TestFunction(space) + flux = TrialFunction(space) + n = FacetNormal(domain.mesh) + un = 0.5 * (dot(v, n) + abs(dot(v, n))) + self.flux = Function(space) + + a = self.phi * flux * dx + L = self.phi * rain * un * rho * (ds_b + ds_t + ds_v) + + # setup solver + problem = LinearVariationalProblem(a, L, self.flux) + self.solver = LinearVariationalSolver(problem) + self.space = space + self.field = state_fields(self.name, space=space, dump=True, pickup=False) + # TODO: might we want to pick up this field? Otherwise initialise to zero + self.field.assign(0.0) + + def compute(self): + """Compute the diagnostic field from the current state.""" self.solver.solve() self.field.assign(self.field + assemble(self.flux * self.phi * dx)) - return self.field class Vorticity(DiagnosticField): """Base diagnostic field class for shallow-water vorticity variables.""" - def setup(self, state, vorticity_type=None): + def setup(self, domain, state_fields, vorticity_type=None): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. vorticity_type (str, optional): denotes which type of vorticity to be computed ('relative', 'absolute' or 'potential'). Defaults to None. """ - if not self._initialised: - vorticity_types = ["relative", "absolute", "potential"] - if vorticity_type not in vorticity_types: - raise ValueError("vorticity type must be one of %s, not %s" % (vorticity_types, vorticity_type)) - try: - space = state.spaces("CG") - except ValueError: - dgspace = state.spaces("DG") - cg_degree = dgspace.ufl_element().degree() + 2 - space = FunctionSpace(state.mesh, "CG", cg_degree) - super().setup(state, space=space) - u = state.fields("u") + + vorticity_types = ["relative", "absolute", "potential"] + if vorticity_type not in vorticity_types: + raise ValueError(f"vorticity type must be one of {vorticity_types}, not {vorticity_type}") + try: + space = domain.spaces("CG") + except ValueError: + dgspace = domain.spaces("DG") + # TODO: should this be degree + 1? + cg_degree = dgspace.ufl_element().degree() + 2 + space = FunctionSpace(domain.mesh, "CG", cg_degree, name=f"CG{cg_degree}") + + u = state_fields("u") + if vorticity_type in ["absolute", "potential"]: + f = state_fields("coriolis") + if vorticity_type == "potential": + D = state_fields("D") + + if self.method != 'solve': + if vorticity_type == "potential": + self.expr = curl(u + f) / D + elif vorticity_type == "absolute": + self.expr = curl(u + f) + elif vorticity_type == "relative": + self.expr = curl(u) + + super().setup(domain, state_fields, space=space) + + # Set up problem now that self.field has been set up + if self.method == 'solve': gamma = TestFunction(space) q = TrialFunction(space) if vorticity_type == "potential": - D = state.fields("D") a = q*gamma*D*dx else: a = q*gamma*dx - L = (- inner(state.perp(grad(gamma)), u))*dx + L = (- inner(domain.perp(grad(gamma)), u))*dx if vorticity_type != "relative": - f = state.fields("coriolis") + f = state_fields("coriolis") L += gamma*f*dx problem = LinearVariationalProblem(a, L, self.field) - self.solver = LinearVariationalSolver(problem, solver_parameters={"ksp_type": "cg"}) - - def compute(self, state): - """ - Compute and return the diagnostic field from the current state. - - Args: - state (:class:`State`): the model's state. - - Returns: - :class:`Function`: the diagnostic field. - """ - self.solver.solve() - return self.field + self.evaluator = LinearVariationalSolver(problem, solver_parameters={"ksp_type": "cg"}) class PotentialVorticity(Vorticity): u"""Diagnostic field for shallow-water potential vorticity, q=(∇×(u+f))/D""" name = "PotentialVorticity" - def setup(self, state): + def __init__(self, space=None, method='solve'): + """ + Args: + space (:class:`FunctionSpace`, optional): the function space to + evaluate the diagnostic field in. Defaults to None, in which + case a default space will be chosen for this diagnostic. + method (str, optional): a string specifying the method of evaluation + for this diagnostic. Valid options are 'interpolate', 'project', + 'assign' and 'solve'. Defaults to 'solve'. + """ + self.solve_implemented = True + super().__init__(space=space, method=method, + required_fields=('u', 'D', 'coriolis')) + + def setup(self, domain, state_fields): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - super().setup(state, vorticity_type="potential") + super().setup(domain, state_fields, vorticity_type="potential") class AbsoluteVorticity(Vorticity): u"""Diagnostic field for absolute vorticity, ζ=∇×(u+f)""" name = "AbsoluteVorticity" - def setup(self, state): + def __init__(self, space=None, method='solve'): + """ + Args: + space (:class:`FunctionSpace`, optional): the function space to + evaluate the diagnostic field in. Defaults to None, in which + case a default space will be chosen for this diagnostic. + method (str, optional): a string specifying the method of evaluation + for this diagnostic. Valid options are 'interpolate', 'project', + 'assign' and 'solve'. Defaults to 'solve'. + """ + self.solve_implemented = True + super().__init__(space=space, method=method, required_fields=('u', 'coriolis')) + + def setup(self, domain, state_fields): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - super().setup(state, vorticity_type="absolute") + super().setup(domain, state_fields, vorticity_type="absolute") class RelativeVorticity(Vorticity): u"""Diagnostic field for relative vorticity, ζ=∇×u""" name = "RelativeVorticity" - def setup(self, state): + def __init__(self, space=None, method='solve'): + """ + Args: + space (:class:`FunctionSpace`, optional): the function space to + evaluate the diagnostic field in. Defaults to None, in which + case a default space will be chosen for this diagnostic. + method (str, optional): a string specifying the method of evaluation + for this diagnostic. Valid options are 'interpolate', 'project', + 'assign' and 'solve'. Defaults to 'solve'. + """ + self.solve_implemented = True + super().__init__(space=space, method=method, required_fields=('u',)) + + def setup(self, domain, state_fields): """ Sets up the :class:`Function` for the diagnostic field. Args: - state (:class:`State`): the model's state. + domain (:class:`Domain`): the model's domain object. + state_fields (:class:`StateFields`): the model's field container. """ - super().setup(state, vorticity_type="relative") + super().setup(domain, state_fields, vorticity_type="relative") diff --git a/gusto/diffusion.py b/gusto/diffusion.py index 621ada6bd..175038e85 100644 --- a/gusto/diffusion.py +++ b/gusto/diffusion.py @@ -8,7 +8,7 @@ __all__ = ["interior_penalty_diffusion_form"] -def interior_penalty_diffusion_form(state, test, q, parameters): +def interior_penalty_diffusion_form(domain, test, q, parameters): u""" Form for the interior penalty discretisation of a diffusion term, ∇.(κ∇q) @@ -16,7 +16,8 @@ def interior_penalty_diffusion_form(state, test, q, parameters): weight function. Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. test (:class:`TestFunction`): the equation's test function. q (:class:`Function`): the variable being diffused. parameters (:class:`DiffusionParameters`): object containing metadata @@ -26,11 +27,11 @@ def interior_penalty_diffusion_form(state, test, q, parameters): :class:`ufl.Form`: the diffusion form. """ - dS_ = (dS_v + dS_h) if state.mesh.extruded else dS + dS_ = (dS_v + dS_h) if domain.mesh.extruded else dS kappa = parameters.kappa mu = parameters.mu - n = FacetNormal(state.mesh) + n = FacetNormal(domain.mesh) form = inner(grad(test), grad(q)*kappa)*dx diff --git a/gusto/domain.py b/gusto/domain.py new file mode 100644 index 000000000..301b4c81e --- /dev/null +++ b/gusto/domain.py @@ -0,0 +1,87 @@ +""" +The Domain object that is provided in this module contains the model's mesh and +the set of compatible function spaces defined upon it. It also contains the +model's time interval. +""" + +from gusto.function_spaces import Spaces, check_degree_args +from firedrake import (Constant, SpatialCoordinate, sqrt, CellNormal, cross, + as_vector, inner, interpolate) + + +class Domain(object): + """ + The Domain holds the model's mesh and its compatible function spaces. + + The compatible function spaces are given by the de Rham complex, and are + specified here through the family of the HDiv velocity space and the degree + of the DG space. + + For extruded meshes, it is possible to seperately specify the horizontal and + vertical degrees of the elements. Alternatively, if these degrees should be + the same then this can be specified through the "degree" argument. + """ + def __init__(self, mesh, dt, family, degree=None, + horizontal_degree=None, vertical_degree=None): + """ + Args: + mesh (:class:`Mesh`): the model's mesh. + dt (:class:`Constant`): the time taken to perform a single model + step. If a float or int is passed, it will be cast to a + :class:`Constant`. + family (str): the finite element space family used for the velocity + field. This determines the other finite element spaces used via + the de Rham complex. + degree (int, optional): the element degree used for the DG space + Defaults to None, in which case the horizontal degree must be provided. + horizontal_degree (int, optional): the element degree used for the + horizontal part of the DG space. Defaults to None. + vertical_degree (int, optional): the element degree used for the + vertical part of the DG space. Defaults to None. + + Raises: + ValueError: if incompatible degrees are specified (e.g. specifying + both "degree" and "horizontal_degree"). + """ + + if type(dt) is Constant: + self.dt = dt + elif type(dt) in (float, int): + self.dt = Constant(dt) + else: + raise TypeError(f'dt must be a Constant, float or int, not {type(dt)}') + + check_degree_args('Domain', mesh, degree, horizontal_degree, vertical_degree) + + # Get degrees + self.horizontal_degree = degree if horizontal_degree is None else horizontal_degree + self.vertical_degree = degree if vertical_degree is None else vertical_degree + + self.mesh = mesh + self.family = family + self.spaces = Spaces(mesh) + # Build and store compatible spaces + self.compatible_spaces = [space for space in self.spaces.build_compatible_spaces(self.family, self.horizontal_degree, self.vertical_degree)] + + # Figure out if we're on a sphere + # TODO: could we run on other domains that could confuse this? + if hasattr(mesh, "_base_mesh") and hasattr(mesh._base_mesh, 'geometric_dimension'): + self.on_sphere = (mesh._base_mesh.geometric_dimension() == 3 and mesh._base_mesh.topological_dimension() == 2) + else: + self.on_sphere = (mesh.geometric_dimension() == 3 and mesh.topological_dimension() == 2) + + # build the vertical normal and define perp for 2d geometries + dim = mesh.topological_dimension() + if self.on_sphere: + x = SpatialCoordinate(mesh) + R = sqrt(inner(x, x)) + self.k = interpolate(x/R, mesh.coordinates.function_space()) + if dim == 2: + outward_normals = CellNormal(mesh) + self.perp = lambda u: cross(outward_normals, u) + else: + kvec = [0.0]*dim + kvec[dim-1] = 1.0 + self.k = Constant(kvec) + if dim == 2: + self.perp = lambda u: as_vector([-u[1], u[0]]) diff --git a/gusto/equations.py b/gusto/equations.py index 6597be2fc..4c80ad6d7 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -6,6 +6,7 @@ TrialFunction, FacetNormal, jump, avg, dS_v, DirichletBC, conditional, SpatialCoordinate, split, Constant, action) +from gusto.fields import PrescribedFields from gusto.fml.form_manipulation_labelling import Term, all_terms, keep, drop, Label from gusto.labels import (subject, time_derivative, transport, prognostic, transporting_velocity, replace_subject, linearisation, @@ -28,30 +29,27 @@ class PrognosticEquation(object, metaclass=ABCMeta): """Base class for prognostic equations.""" - def __init__(self, state, function_space, field_name): + def __init__(self, domain, function_space, field_name): """ Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. function_space (:class:`FunctionSpace`): the function space that the equation's prognostic is defined on. field_name (str): name of the prognostic field. """ - self.state = state + self.domain = domain self.function_space = function_space + self.X = Function(function_space) self.field_name = field_name self.bcs = {} + self.prescribed_fields = PrescribedFields() if len(function_space) > 1: assert hasattr(self, "field_names") - state.fields(field_name, function_space, - subfield_names=self.field_names, pickup=True) for fname in self.field_names: - state.diagnostics.register(fname) self.bcs[fname] = [] - else: - state.fields(field_name, function_space) - state.diagnostics.register(field_name) self.bcs[field_name] = [] @@ -72,98 +70,85 @@ def label_terms(self, term_filter, label): class AdvectionEquation(PrognosticEquation): u"""Discretises the advection equation, ∂q/∂t + (u.∇)q = 0""" - def __init__(self, state, function_space, field_name, - ufamily=None, udegree=None, Vu=None, **kwargs): + def __init__(self, domain, function_space, field_name, Vu=None, **kwargs): """ Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. function_space (:class:`FunctionSpace`): the function space that the equation's prognostic is defined on. field_name (str): name of the prognostic field. - ufamily (str, optional): the family of the function space to use - for the velocity field. Only used if `Vu` is not provided. - Defaults to None. - udegree (int, optional): the degree of the function space to use for - the velocity field. Only used if `Vu` is not provided. Defaults - to None. Vu (:class:`FunctionSpace`, optional): the function space for the - velocity field. If this is Defaults to None. + velocity field. If this is not specified, uses the HDiv spaces + set up by the domain. Defaults to None. **kwargs: any keyword arguments to be passed to the advection form. """ - super().__init__(state, function_space, field_name) + super().__init__(domain, function_space, field_name) + + if Vu is not None: + V = domain.spaces("HDiv", V=Vu, overwrite_space=True) + else: + V = domain.spaces("HDiv") + self.prescribed_fields("u", V) - if not hasattr(state.fields, "u"): - if Vu is not None: - V = state.spaces("HDiv", V=Vu) - else: - assert ufamily is not None, "Specify the family for u" - assert udegree is not None, "Specify the degree of the u space" - V = state.spaces("HDiv", ufamily, udegree) - state.fields("u", V) test = TestFunction(function_space) q = Function(function_space) mass_form = time_derivative(inner(q, test)*dx) self.residual = subject( - mass_form + advection_form(state, test, q, **kwargs), q + mass_form + advection_form(domain, test, q, **kwargs), q ) class ContinuityEquation(PrognosticEquation): u"""Discretises the continuity equation, ∂q/∂t + ∇(u*q) = 0""" - def __init__(self, state, function_space, field_name, - ufamily=None, udegree=None, Vu=None, **kwargs): + def __init__(self, domain, function_space, field_name, Vu=None, **kwargs): """ Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. function_space (:class:`FunctionSpace`): the function space that the equation's prognostic is defined on. field_name (str): name of the prognostic field. - ufamily (str, optional): the family of the function space to use - for the velocity field. Only used if `Vu` is not provided. - Defaults to None. - udegree (int, optional): the degree of the function space to use for - the velocity field. Only used if `Vu` is not provided. Defaults - to None. Vu (:class:`FunctionSpace`, optional): the function space for the - velocity field. If this is Defaults to None. + velocity field. If this is not specified, uses the HDiv spaces + set up by the domain. Defaults to None. **kwargs: any keyword arguments to be passed to the advection form. """ - super().__init__(state, function_space, field_name) + super().__init__(domain, function_space, field_name) + + if Vu is not None: + V = domain.spaces("HDiv", V=Vu, overwrite_space=True) + else: + V = domain.spaces("HDiv") + self.prescribed_fields("u", V) - if not hasattr(state.fields, "u"): - if Vu is not None: - V = state.spaces("HDiv", V=Vu) - else: - assert ufamily is not None, "Specify the family for u" - assert udegree is not None, "Specify the degree of the u space" - V = state.spaces("HDiv", ufamily, udegree) - state.fields("u", V) test = TestFunction(function_space) q = Function(function_space) mass_form = time_derivative(inner(q, test)*dx) self.residual = subject( - mass_form + continuity_form(state, test, q, **kwargs), q + mass_form + continuity_form(domain, test, q, **kwargs), q ) class DiffusionEquation(PrognosticEquation): u"""Discretises the diffusion equation, ∂q/∂t = ∇.(κ∇q)""" - def __init__(self, state, function_space, field_name, + def __init__(self, domain, function_space, field_name, diffusion_parameters): """ Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. function_space (:class:`FunctionSpace`): the function space that the equation's prognostic is defined on. field_name (str): name of the prognostic field. diffusion_parameters (:class:`DiffusionParameters`): parameters describing the diffusion to be applied. """ - super().__init__(state, function_space, field_name) + super().__init__(domain, function_space, field_name) test = TestFunction(function_space) q = Function(function_space) @@ -172,28 +157,22 @@ def __init__(self, state, function_space, field_name, self.residual = subject( mass_form + interior_penalty_diffusion_form( - state, test, q, diffusion_parameters), q + domain, test, q, diffusion_parameters), q ) class AdvectionDiffusionEquation(PrognosticEquation): u"""The advection-diffusion equation, ∂q/∂t + (u.∇)q = ∇.(κ∇q)""" - def __init__(self, state, function_space, field_name, - ufamily=None, udegree=None, Vu=None, diffusion_parameters=None, - **kwargs): + def __init__(self, domain, function_space, field_name, Vu=None, + diffusion_parameters=None, **kwargs): """ Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. function_space (:class:`FunctionSpace`): the function space that the equation's prognostic is defined on. field_name (str): name of the prognostic field. - ufamily (str, optional): the family of the function space to use - for the velocity field. Only used if `Vu` is not provided. - Defaults to None. - udegree (int, optional): the degree of the function space to use for - the velocity field. Only used if `Vu` is not provided. Defaults - to None. Vu (:class:`FunctionSpace`, optional): the function space for the velocity field. If this is Defaults to None. diffusion_parameters (:class:`DiffusionParameters`, optional): @@ -201,25 +180,23 @@ def __init__(self, state, function_space, field_name, **kwargs: any keyword arguments to be passed to the advection form. """ - super().__init__(state, function_space, field_name) + super().__init__(domain, function_space, field_name) + + if Vu is not None: + V = domain.spaces("HDiv", V=Vu, overwrite_space=True) + else: + V = domain.spaces("HDiv") + self.prescribed_fields("u", V) - if not hasattr(state.fields, "u"): - if Vu is not None: - V = state.spaces("HDiv", V=Vu) - else: - assert ufamily is not None, "Specify the family for u" - assert udegree is not None, "Specify the degree of the u space" - V = state.spaces("HDiv", ufamily, udegree) - state.fields("u", V) test = TestFunction(function_space) q = Function(function_space) mass_form = time_derivative(inner(q, test)*dx) self.residual = subject( mass_form - + advection_form(state, test, q, **kwargs) + + advection_form(domain, test, q, **kwargs) + interior_penalty_diffusion_form( - state, test, q, diffusion_parameters), q + domain, test, q, diffusion_parameters), q ) @@ -232,18 +209,14 @@ class PrognosticEquationSet(PrognosticEquation, metaclass=ABCMeta): contains common routines for these equation sets. """ - def __init__(self, field_names, state, family, degree, - linearisation_map=None, no_normal_flow_bc_ids=None, - active_tracers=None): + def __init__(self, field_names, domain, linearisation_map=None, + no_normal_flow_bc_ids=None, active_tracers=None): """ Args: field_names (list): a list of strings for names of the prognostic variables for the equation set. - state (:class:`State`): the model's state object. - family (str): the finite element space family used for the velocity - field. This determines the other finite element spaces used via - the de Rham complex. - degree (int): the element degree used for the velocity space. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. linearisation_map (func, optional): a function specifying which terms in the equation set to linearise. Defaults to None. no_normal_flow_bc_ids (list, optional): a list of IDs of domain @@ -259,33 +232,32 @@ def __init__(self, field_names, state, family, degree, self.linearisation_map = lambda t: False if linearisation_map is None else linearisation_map(t) # Build finite element spaces - self.spaces = [space for space in self._build_spaces(state, family, degree)] + # TODO: this implies order of spaces matches order of variables + # we should not assume this and should instead specify which variable + # is in which space + self.spaces = [space for space in domain.compatible_spaces] # Add active tracers to the list of prognostics if active_tracers is None: active_tracers = [] - self.add_tracers_to_prognostics(state, active_tracers) + self.add_tracers_to_prognostics(domain, active_tracers) # Make the full mixed function space W = MixedFunctionSpace(self.spaces) # Can now call the underlying PrognosticEquation full_field_name = "_".join(self.field_names) - super().__init__(state, W, full_field_name) + super().__init__(domain, W, full_field_name) # Set up test functions, trials and prognostics self.tests = TestFunctions(W) self.trials = TrialFunction(W) - self.X = Function(W) self.X_ref = Function(W) # Set up no-normal-flow boundary conditions if no_normal_flow_bc_ids is None: no_normal_flow_bc_ids = [] - self.set_no_normal_flow_bcs(state, no_normal_flow_bc_ids) - - def _build_spaces(self, state, family, degree): - return state.spaces.build_compatible_spaces(family, degree) + self.set_no_normal_flow_bcs(domain, no_normal_flow_bc_ids) # ======================================================================== # # Set up time derivative / mass terms @@ -384,7 +356,7 @@ def linearise_equation_set(self): # Boundary Condition Routines # ======================================================================== # - def set_no_normal_flow_bcs(self, state, no_normal_flow_bc_ids): + def set_no_normal_flow_bcs(self, domain, no_normal_flow_bc_ids): """ Sets up the boundary conditions for no-normal flow at domain boundaries. @@ -393,7 +365,8 @@ def set_no_normal_flow_bcs(self, state, no_normal_flow_bc_ids): a velocity variable named 'u' to apply the boundary conditions to. Args: - state (:class:`State`): the model's state. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. no_normal_flow_bc_ids (list): A list of IDs of the domain boundaries at which no normal flow will be enforced. @@ -407,7 +380,7 @@ def set_no_normal_flow_bcs(self, state, no_normal_flow_bc_ids): 'No-normal-flow boundary conditions can only be applied ' + 'when there is a variable called "u" and none was found') - Vu = state.spaces("HDiv") + Vu = domain.spaces("HDiv") if Vu.extruded: self.bcs['u'].append(DirichletBC(Vu, 0.0, "bottom")) self.bcs['u'].append(DirichletBC(Vu, 0.0, "top")) @@ -425,12 +398,13 @@ def set_no_normal_flow_bcs(self, state, no_normal_flow_bc_ids): # Active Tracer Routines # ======================================================================== # - def add_tracers_to_prognostics(self, state, active_tracers): + def add_tracers_to_prognostics(self, domain, active_tracers): """ Augments the equation set with specified active tracer variables. Args: - state (:class:`State`): the model's state. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. active_tracers (list): A list of :class:`ActiveTracer` objects that encode the metadata for the active tracers. @@ -446,16 +420,17 @@ def add_tracers_to_prognostics(self, state, active_tracers): self.field_names.append(tracer.name) else: raise ValueError(f'There is already a field named {tracer.name}') - self.spaces.append(state.spaces(tracer.space)) + self.spaces.append(domain.spaces(tracer.space)) else: raise TypeError(f'Tracers must be ActiveTracer objects, not {type(tracer)}') - def generate_tracer_transport_terms(self, state, active_tracers): + def generate_tracer_transport_terms(self, domain, active_tracers): """ Adds the transport forms for the active tracers to the equation set. Args: - state (:class:`State`): the model's state. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. active_tracers (list): A list of :class:`ActiveTracer` objects that encode the metadata for the active tracers. @@ -478,9 +453,9 @@ def generate_tracer_transport_terms(self, state, active_tracers): tracer_prog = split(self.X)[idx] tracer_test = self.tests[idx] if tracer.transport_eqn == TransportEquationType.advective: - tracer_adv = prognostic(advection_form(state, tracer_test, tracer_prog), tracer.name) + tracer_adv = prognostic(advection_form(domain, tracer_test, tracer_prog), tracer.name) elif tracer.transport_eqn == TransportEquationType.conservative: - tracer_adv = prognostic(continuity_form(state, tracer_test, tracer_prog), tracer.name) + tracer_adv = prognostic(continuity_form(domain, tracer_test, tracer_prog), tracer.name) else: raise ValueError(f'Transport eqn {tracer.transport_eqn} not recognised') @@ -495,38 +470,53 @@ def generate_tracer_transport_terms(self, state, active_tracers): class ForcedAdvectionEquation(PrognosticEquationSet): - - def __init__(self, state, function_space, field_name, - ufamily=None, udegree=None, Vu=None, active_tracers=None, - **kwargs): + u""" + Discretises the advection equation with a source/sink term, + ∂q/∂t + (u.∇)q = F, + which can also be augmented with active tracers. + """ + def __init__(self, domain, function_space, field_name, Vu=None, + active_tracers=None, **kwargs): + """ + Args: + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + function_space (:class:`FunctionSpace`): the function space that the + equation's prognostic is defined on. + field_name (str): name of the prognostic field. + Vu (:class:`FunctionSpace`, optional): the function space for the + velocity field. If this is not specified, uses the HDiv spaces + set up by the domain. Defaults to None. + active_tracers (list, optional): a list of `ActiveTracer` objects + that encode the metadata for any active tracers to be included + in the equations. Defaults to None. + **kwargs: any keyword arguments to be passed to the advection form. + """ self.field_names = [field_name] self.active_tracers = active_tracers self.terms_to_linearise = {} # Build finite element spaces - self.spaces = [state.spaces("tracer", V=function_space)] + self.spaces = [domain.spaces("tracer", V=function_space)] # Add active tracers to the list of prognostics if active_tracers is None: active_tracers = [] - self.add_tracers_to_prognostics(state, active_tracers) + self.add_tracers_to_prognostics(domain, active_tracers) # Make the full mixed function space W = MixedFunctionSpace(self.spaces) # Can now call the underlying PrognosticEquation full_field_name = "_".join(self.field_names) - PrognosticEquation.__init__(self, state, W, full_field_name) + PrognosticEquation.__init__(self, domain, W, full_field_name) - if not hasattr(state.fields, "u"): - if Vu is not None: - V = state.spaces("HDiv", V=Vu) - else: - assert ufamily is not None, "Specify the family for u" - assert udegree is not None, "Specify the degree of the u space" - V = state.spaces("HDiv", ufamily, udegree) - state.fields("u", V) + if Vu is not None: + V = domain.spaces("HDiv", V=Vu, overwrite_space=True) + else: + V = domain.spaces("HDiv") + self.prescribed_fields("u", V) self.tests = TestFunctions(W) self.X = Function(W) @@ -534,7 +524,7 @@ def __init__(self, state, function_space, field_name, mass_form = self.generate_mass_terms() self.residual = subject( - mass_form + advection_form(state, self.tests[0], split(self.X)[0], **kwargs), self.X + mass_form + advection_form(domain, self.tests[0], split(self.X)[0], **kwargs), self.X ) # ============================================================================ # @@ -551,17 +541,16 @@ class ShallowWaterEquations(PrognosticEquationSet): for Coriolis parameter 'f' and bottom surface 'b'. """ - def __init__(self, state, family, degree, fexpr=None, bexpr=None, + def __init__(self, domain, parameters, fexpr=None, bexpr=None, linearisation_map='default', u_transport_option='vector_invariant_form', no_normal_flow_bc_ids=None, active_tracers=None): """ Args: - state (:class:`State`): the model's state object. - family (str): the finite element space family used for the velocity - field. This determines the other finite element spaces used via - the de Rham complex. - degree (int): the element degree used for the velocity space. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + parameters (:class:`Configuration`, optional): an object containing + the model's physical parameters. fexpr (:class:`ufl.Expr`, optional): an expression for the Coroilis parameter. Defaults to None. bexpr (:class:`ufl.Expr`, optional): an expression for the bottom @@ -594,18 +583,19 @@ def __init__(self, state, family, degree, fexpr=None, bexpr=None, if linearisation_map == 'default': # Default linearisation is time derivatives, pressure gradient and - # transport term from depth equation + # transport term from depth equation. Don't include active tracers linearisation_map = lambda t: \ - (any(t.has_label(time_derivative, pressure_gradient)) - or (t.get(prognostic) == "D" and t.has_label(transport))) - - super().__init__(field_names, state, family, degree, + t.get(prognostic) in ["u", "D"] \ + and (any(t.has_label(time_derivative, pressure_gradient)) + or (t.get(prognostic) == "D" and t.has_label(transport))) + super().__init__(field_names, domain, linearisation_map=linearisation_map, no_normal_flow_bc_ids=no_normal_flow_bc_ids, active_tracers=active_tracers) - g = state.parameters.g - H = state.parameters.H + self.parameters = parameters + g = parameters.g + H = parameters.H w, phi = self.tests[0:2] u, D = split(self.X)[0:2] @@ -621,28 +611,28 @@ def __init__(self, state, family, degree, fexpr=None, bexpr=None, # -------------------------------------------------------------------- # # Velocity transport term -- depends on formulation if u_transport_option == "vector_invariant_form": - u_adv = prognostic(vector_invariant_form(state, w, u), "u") + u_adv = prognostic(vector_invariant_form(domain, w, u), "u") elif u_transport_option == "vector_advection_form": - u_adv = prognostic(advection_form(state, w, u), "u") + u_adv = prognostic(advection_form(domain, w, u), "u") elif u_transport_option == "vector_manifold_advection_form": - u_adv = prognostic(vector_manifold_advection_form(state, w, u), "u") + u_adv = prognostic(vector_manifold_advection_form(domain, w, u), "u") elif u_transport_option == "circulation_form": - ke_form = prognostic(kinetic_energy_form(state, w, u), "u") + ke_form = prognostic(kinetic_energy_form(domain, w, u), "u") ke_form = transport.remove(ke_form) ke_form = ke_form.label_map( lambda t: t.has_label(transporting_velocity), lambda t: Term(ufl.replace( t.form, {t.get(transporting_velocity): u}), t.labels)) ke_form = transporting_velocity.remove(ke_form) - u_adv = prognostic(advection_equation_circulation_form(state, w, u), "u") + ke_form + u_adv = prognostic(advection_equation_circulation_form(domain, w, u), "u") + ke_form else: raise ValueError("Invalid u_transport_option: %s" % u_transport_option) # Depth transport term - D_adv = prognostic(continuity_form(state, phi, D), "D") + D_adv = prognostic(continuity_form(domain, phi, D), "D") # Transport term needs special linearisation if self.linearisation_map(D_adv.terms[0]): - linear_D_adv = linear_continuity_form(state, phi, H).label_map( + linear_D_adv = linear_continuity_form(domain, phi, H).label_map( lambda t: t.has_label(transporting_velocity), lambda t: Term(ufl.replace( t.form, {t.get(transporting_velocity): u_trial}), t.labels)) @@ -653,7 +643,7 @@ def __init__(self, state, family, degree, fexpr=None, bexpr=None, # Add transport of tracers if len(active_tracers) > 0: - adv_form += self.generate_tracer_transport_terms(state, active_tracers) + adv_form += self.generate_tracer_transport_terms(domain, active_tracers) # -------------------------------------------------------------------- # # Pressure Gradient Term @@ -666,37 +656,33 @@ def __init__(self, state, family, degree, fexpr=None, bexpr=None, # -------------------------------------------------------------------- # # Extra Terms (Coriolis and Topography) # -------------------------------------------------------------------- # + # TODO: Is there a better way to store the Coriolis / topography fields? + # The current approach is that these are prescribed fields, stored in + # the equation, and initialised when the equation is + if fexpr is not None: - V = FunctionSpace(state.mesh, "CG", 1) - f = state.fields("coriolis", space=V) - f.interpolate(fexpr) + V = FunctionSpace(domain.mesh, "CG", 1) + f = self.prescribed_fields("coriolis", V).interpolate(fexpr) coriolis_form = perp( coriolis( subject(prognostic(f*inner(u, w)*dx, "u"), self.X) - ), state.perp) + ), domain.perp) # Add linearisation linear_coriolis = perp( coriolis( subject(prognostic(f*inner(u_trial, w)*dx, "u"), self.X) - ), state.perp) + ), domain.perp) coriolis_form = linearisation(coriolis_form, linear_coriolis) residual += coriolis_form if bexpr is not None: - b = state.fields("topography", state.spaces("DG")) - b.interpolate(bexpr) + b = self.prescribed_fields("topography", domain.spaces("DG")).interpolate(bexpr) topography_form = subject(prognostic(-g*div(w)*b*dx, "u"), self.X) residual += topography_form # -------------------------------------------------------------------- # # Linearise equations # -------------------------------------------------------------------- # - u_ref, D_ref = self.X_ref.split()[0:2] - # Linearise about D = H - # TODO: add interface to update linearisation state - D_ref.assign(Constant(H)) - u_ref.assign(Constant(0.0)) - # Add linearisations to equations self.residual = self.generate_linear_terms(residual, self.linearisation_map) @@ -713,17 +699,16 @@ class LinearShallowWaterEquations(ShallowWaterEquations): which is then linearised. """ - def __init__(self, state, family, degree, fexpr=None, bexpr=None, + def __init__(self, domain, parameters, fexpr=None, bexpr=None, linearisation_map='default', u_transport_option="vector_invariant_form", no_normal_flow_bc_ids=None, active_tracers=None): """ Args: - state (:class:`State`): the model's state object. - family (str): the finite element space family used for the velocity - field. This determines the other finite element spaces used via - the de Rham complex. - degree (int): the element degree used for the velocity space. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + parameters (:class:`Configuration`, optional): an object containing + the model's physical parameters. fexpr (:class:`ufl.Expr`, optional): an expression for the Coroilis parameter. Defaults to None. bexpr (:class:`ufl.Expr`, optional): an expression for the bottom @@ -753,7 +738,7 @@ def __init__(self, state, family, degree, fexpr=None, bexpr=None, (any(t.has_label(time_derivative, pressure_gradient, coriolis)) or (t.get(prognostic) == "D" and t.has_label(transport))) - super().__init__(state, family, degree, fexpr=fexpr, bexpr=bexpr, + super().__init__(domain, parameters, fexpr=fexpr, bexpr=bexpr, linearisation_map=linearisation_map, u_transport_option=u_transport_option, no_normal_flow_bc_ids=no_normal_flow_bc_ids, @@ -765,7 +750,7 @@ def __init__(self, state, family, degree, fexpr=None, bexpr=None, # D transport term is a special case -- add facet term _, D = split(self.X) _, phi = self.tests - D_adv = prognostic(linear_continuity_form(state, phi, D, facet_term=True), "D") + D_adv = prognostic(linear_continuity_form(domain, phi, D, facet_term=True), "D") self.residual = self.residual.label_map( lambda t: t.has_label(transport) and t.get(prognostic) == "D", map_if_true=lambda t: Term(D_adv.form, t.labels) @@ -785,7 +770,7 @@ class CompressibleEulerEquations(PrognosticEquationSet): pressure. """ - def __init__(self, state, family, degree, Omega=None, sponge=None, + def __init__(self, domain, parameters, Omega=None, sponge=None, extra_terms=None, linearisation_map='default', u_transport_option="vector_invariant_form", diffusion_options=None, @@ -793,11 +778,10 @@ def __init__(self, state, family, degree, Omega=None, sponge=None, active_tracers=None): """ Args: - state (:class:`State`): the model's state object. - family (str): the finite element space family used for the velocity - field. This determines the other finite element spaces used via - the de Rham complex. - degree (int): the element degree used for the velocity space. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + parameters (:class:`Configuration`, optional): an object containing + the model's physical parameters. Omega (:class:`ufl.Expr`, optional): an expression for the planet's rotation vector. Defaults to None. sponge (:class:`ufl.Expr`, optional): an expression for a sponge @@ -835,26 +819,27 @@ def __init__(self, state, family, degree, Omega=None, sponge=None, if linearisation_map == 'default': # Default linearisation is time derivatives and scalar transport terms + # Don't include active tracers linearisation_map = lambda t: \ - (t.has_label(time_derivative) - or (t.get(prognostic) != "u" and t.has_label(transport))) - - super().__init__(field_names, state, family, degree, + t.get(prognostic) in ['u', 'rho', 'theta'] \ + and (t.has_label(time_derivative) + or (t.get(prognostic) != 'u' and t.has_label(transport))) + super().__init__(field_names, domain, linearisation_map=linearisation_map, no_normal_flow_bc_ids=no_normal_flow_bc_ids, active_tracers=active_tracers) - g = state.parameters.g - cp = state.parameters.cp + self.parameters = parameters + g = parameters.g + cp = parameters.cp w, phi, gamma = self.tests[0:3] u, rho, theta = split(self.X)[0:3] u_trial = split(self.trials)[0] - rhobar = state.fields("rhobar", space=state.spaces("DG"), dump=False) - thetabar = state.fields("thetabar", space=state.spaces("theta"), dump=False) + _, rho_bar, theta_bar = split(self.X_ref)[0:3] zero_expr = Constant(0.0)*theta - exner = exner_pressure(state.parameters, rho, theta) - n = FacetNormal(state.mesh) + exner = exner_pressure(parameters, rho, theta) + n = FacetNormal(domain.mesh) # -------------------------------------------------------------------- # # Time Derivative Terms @@ -866,38 +851,38 @@ def __init__(self, state, family, degree, Omega=None, sponge=None, # -------------------------------------------------------------------- # # Velocity transport term -- depends on formulation if u_transport_option == "vector_invariant_form": - u_adv = prognostic(vector_invariant_form(state, w, u), "u") + u_adv = prognostic(vector_invariant_form(domain, w, u), "u") elif u_transport_option == "vector_advection_form": - u_adv = prognostic(advection_form(state, w, u), "u") + u_adv = prognostic(advection_form(domain, w, u), "u") elif u_transport_option == "vector_manifold_advection_form": - u_adv = prognostic(vector_manifold_advection_form(state, w, u), "u") + u_adv = prognostic(vector_manifold_advection_form(domain, w, u), "u") elif u_transport_option == "circulation_form": - ke_form = prognostic(kinetic_energy_form(state, w, u), "u") + ke_form = prognostic(kinetic_energy_form(domain, w, u), "u") ke_form = transport.remove(ke_form) ke_form = ke_form.label_map( lambda t: t.has_label(transporting_velocity), lambda t: Term(ufl.replace( t.form, {t.get(transporting_velocity): u}), t.labels)) ke_form = transporting_velocity.remove(ke_form) - u_adv = prognostic(advection_equation_circulation_form(state, w, u), "u") + ke_form + u_adv = prognostic(advection_equation_circulation_form(domain, w, u), "u") + ke_form else: raise ValueError("Invalid u_transport_option: %s" % u_transport_option) # Density transport (conservative form) - rho_adv = prognostic(continuity_form(state, phi, rho), "rho") + rho_adv = prognostic(continuity_form(domain, phi, rho), "rho") # Transport term needs special linearisation if self.linearisation_map(rho_adv.terms[0]): - linear_rho_adv = linear_continuity_form(state, phi, rhobar).label_map( + linear_rho_adv = linear_continuity_form(domain, phi, rho_bar).label_map( lambda t: t.has_label(transporting_velocity), lambda t: Term(ufl.replace( t.form, {t.get(transporting_velocity): u_trial}), t.labels)) rho_adv = linearisation(rho_adv, linear_rho_adv) # Potential temperature transport (advective form) - theta_adv = prognostic(advection_form(state, gamma, theta), "theta") + theta_adv = prognostic(advection_form(domain, gamma, theta), "theta") # Transport term needs special linearisation if self.linearisation_map(theta_adv.terms[0]): - linear_theta_adv = linear_advection_form(state, gamma, thetabar).label_map( + linear_theta_adv = linear_advection_form(domain, gamma, theta_bar).label_map( lambda t: t.has_label(transporting_velocity), lambda t: Term(ufl.replace( t.form, {t.get(transporting_velocity): u_trial}), t.labels)) @@ -907,7 +892,7 @@ def __init__(self, state, family, degree, Omega=None, sponge=None, # Add transport of tracers if len(active_tracers) > 0: - adv_form += self.generate_tracer_transport_terms(state, active_tracers) + adv_form += self.generate_tracer_transport_terms(domain, active_tracers) # -------------------------------------------------------------------- # # Pressure Gradient Term @@ -929,7 +914,7 @@ def __init__(self, state, family, degree, Omega=None, sponge=None, # -------------------------------------------------------------------- # # Gravitational Term # -------------------------------------------------------------------- # - gravity_form = subject(prognostic(Term(g*inner(state.k, w)*dx), "u"), self.X) + gravity_form = subject(prognostic(Term(g*inner(domain.k, w)*dx), "u"), self.X) residual = (mass_form + adv_form + pressure_gradient_form + gravity_form) @@ -937,12 +922,12 @@ def __init__(self, state, family, degree, Omega=None, sponge=None, # Moist Thermodynamic Divergence Term # -------------------------------------------------------------------- # if len(active_tracers) > 0: - cv = state.parameters.cv - c_vv = state.parameters.c_vv - c_pv = state.parameters.c_pv - c_pl = state.parameters.c_pl - R_d = state.parameters.R_d - R_v = state.parameters.R_v + cv = parameters.cv + c_vv = parameters.c_vv + c_pv = parameters.c_pv + c_pl = parameters.c_pl + R_d = parameters.R_d + R_v = parameters.R_v # Get gas and liquid moisture mixing ratios mr_l = zero_expr @@ -976,8 +961,8 @@ def __init__(self, state, family, degree, Omega=None, sponge=None, inner(w, cross(2*Omega, u))*dx, "u"), self.X) if sponge is not None: - W_DG = FunctionSpace(state.mesh, "DG", 2) - x = SpatialCoordinate(state.mesh) + W_DG = FunctionSpace(domain.mesh, "DG", 2) + x = SpatialCoordinate(domain.mesh) z = x[len(x)-1] H = sponge.H zc = sponge.z_level @@ -986,10 +971,10 @@ def __init__(self, state, family, degree, Omega=None, sponge=None, muexpr = conditional(z <= zc, 0.0, mubar*sin((pi/2.)*(z-zc)/(H-zc))**2) - self.mu = Function(W_DG).interpolate(muexpr) + self.mu = self.prescribed_fields("sponge", W_DG).interpolate(muexpr) residual += name(subject(prognostic( - self.mu*inner(w, state.k)*inner(u, state.k)*dx, "u"), self.X), "sponge") + self.mu*inner(w, domain.k)*inner(u, domain.k)*dx, "u"), self.X), "sponge") if diffusion_options is not None: for field, diffusion in diffusion_options: @@ -998,7 +983,7 @@ def __init__(self, state, family, degree, Omega=None, sponge=None, fn = split(self.X)[idx] residual += subject( prognostic(interior_penalty_diffusion_form( - state, test, fn, diffusion), field), self.X) + domain, test, fn, diffusion), field), self.X) if extra_terms is not None: for field, term in extra_terms: @@ -1010,7 +995,6 @@ def __init__(self, state, family, degree, Omega=None, sponge=None, # -------------------------------------------------------------------- # # Linearise equations # -------------------------------------------------------------------- # - # TODO: add linearisation states for variables # Add linearisations to equations self.residual = self.generate_linear_terms(residual, self.linearisation_map) @@ -1033,7 +1017,7 @@ class HydrostaticCompressibleEulerEquations(CompressibleEulerEquations): equations. """ - def __init__(self, state, family, degree, Omega=None, sponge=None, + def __init__(self, domain, parameters, Omega=None, sponge=None, extra_terms=None, linearisation_map='default', u_transport_option="vector_invariant_form", diffusion_options=None, @@ -1041,11 +1025,10 @@ def __init__(self, state, family, degree, Omega=None, sponge=None, active_tracers=None): """ Args: - state (:class:`State`): the model's state object. - family (str): the finite element space family used for the velocity - field. This determines the other finite element spaces used via - the de Rham complex. - degree (int): the element degree used for the velocity space. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + parameters (:class:`Configuration`, optional): an object containing + the model's physical parameters. Omega (:class:`ufl.Expr`, optional): an expression for the planet's rotation vector. Defaults to None. sponge (:class:`ufl.Expr`, optional): an expression for a sponge @@ -1076,7 +1059,7 @@ def __init__(self, state, family, degree, Omega=None, sponge=None, NotImplementedError: only mixing ratio tracers are implemented. """ - super().__init__(state, family, degree, Omega=Omega, sponge=sponge, + super().__init__(domain, parameters, Omega=Omega, sponge=sponge, extra_terms=extra_terms, linearisation_map=linearisation_map, u_transport_option=u_transport_option, @@ -1089,7 +1072,7 @@ def __init__(self, state, family, degree, Omega=None, sponge=None, map_if_true=lambda t: hydrostatic(t, self.hydrostatic_projection(t)) ) - k = self.state.k + k = self.domain.k u = split(self.X)[0] self.residual += name( subject( @@ -1116,8 +1099,8 @@ def hydrostatic_projection(self, t): """ # TODO: make this more general, i.e. should work on the sphere - assert not self.state.on_sphere, "the hydrostatic projection is not yet implemented for spherical geometry" - k = Constant((*self.state.k, 0, 0)) + assert not self.domain.on_sphere, "the hydrostatic projection is not yet implemented for spherical geometry" + k = Constant((*self.domain.k, 0, 0)) X = t.get(subject) new_subj = X - k * inner(X, k) @@ -1125,7 +1108,6 @@ def hydrostatic_projection(self, t): class IncompressibleBoussinesqEquations(PrognosticEquationSet): - # TODO: check that these are correct """ Class for the incompressible Boussinesq equations, which evolve the velocity 'u', the pressure 'p' and the buoyancy 'b'. @@ -1138,18 +1120,17 @@ class IncompressibleBoussinesqEquations(PrognosticEquationSet): where k is the vertical unit vector and, Ω is the planet's rotation vector. """ - def __init__(self, state, family, degree, Omega=None, + def __init__(self, domain, parameters, Omega=None, linearisation_map='default', u_transport_option="vector_invariant_form", no_normal_flow_bc_ids=None, active_tracers=None): """ Args: - state (:class:`State`): the model's state object. - family (str): the finite element space family used for the velocity - field. This determines the other finite element spaces used via - the de Rham complex. - degree (int): the element degree used for the velocity space. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. + parameters (:class:`Configuration`, optional): an object containing + the model's physical parameters. Omega (:class:`ufl.Expr`, optional): an expression for the planet's rotation vector. Defaults to None. linearisation_map (func, optional): a function specifying which @@ -1183,20 +1164,23 @@ def __init__(self, state, family, degree, Omega=None, if linearisation_map == 'default': # Default linearisation is time derivatives and scalar transport terms + # Don't include active tracers linearisation_map = lambda t: \ - (t.has_label(time_derivative) - or (t.get(prognostic) not in ["u", "p"] and t.has_label(transport))) + t.get(prognostic) in ['u', 'p', 'b'] \ + and (t.has_label(time_derivative) + or (t.get(prognostic) not in ['u', 'p'] and t.has_label(transport))) - super().__init__(field_names, state, family, degree, + super().__init__(field_names, domain, linearisation_map=linearisation_map, no_normal_flow_bc_ids=no_normal_flow_bc_ids, active_tracers=active_tracers) + self.parameters = parameters + w, phi, gamma = self.tests[0:3] u, p, b = split(self.X) u_trial = split(self.trials)[0] - bbar = state.fields("bbar", space=state.spaces("theta"), dump=False) - bbar = state.fields("pbar", space=state.spaces("DG"), dump=False) + b_bar = split(self.X_ref)[2] # -------------------------------------------------------------------- # # Time Derivative Terms @@ -1208,27 +1192,27 @@ def __init__(self, state, family, degree, Omega=None, # -------------------------------------------------------------------- # # Velocity transport term -- depends on formulation if u_transport_option == "vector_invariant_form": - u_adv = prognostic(vector_invariant_form(state, w, u), "u") + u_adv = prognostic(vector_invariant_form(domain, w, u), "u") elif u_transport_option == "vector_advection_form": - u_adv = prognostic(advection_form(state, w, u), "u") + u_adv = prognostic(advection_form(domain, w, u), "u") elif u_transport_option == "vector_manifold_advection_form": - u_adv = prognostic(vector_manifold_advection_form(state, w, u), "u") + u_adv = prognostic(vector_manifold_advection_form(domain, w, u), "u") elif u_transport_option == "circulation_form": - ke_form = prognostic(kinetic_energy_form(state, w, u), "u") + ke_form = prognostic(kinetic_energy_form(domain, w, u), "u") ke_form = transport.remove(ke_form) ke_form = ke_form.label_map( lambda t: t.has_label(transporting_velocity), lambda t: Term(ufl.replace( t.form, {t.get(transporting_velocity): u}), t.labels)) ke_form = transporting_velocity.remove(ke_form) - u_adv = prognostic(advection_equation_circulation_form(state, w, u), "u") + ke_form + u_adv = prognostic(advection_equation_circulation_form(domain, w, u), "u") + ke_form else: raise ValueError("Invalid u_transport_option: %s" % u_transport_option) # Buoyancy transport - b_adv = prognostic(advection_form(state, gamma, b), "b") + b_adv = prognostic(advection_form(domain, gamma, b), "b") if self.linearisation_map(b_adv.terms[0]): - linear_b_adv = linear_advection_form(state, gamma, bbar).label_map( + linear_b_adv = linear_advection_form(domain, gamma, b_bar).label_map( lambda t: t.has_label(transporting_velocity), lambda t: Term(ufl.replace( t.form, {t.get(transporting_velocity): u_trial}), t.labels)) @@ -1238,7 +1222,7 @@ def __init__(self, state, family, degree, Omega=None, # Add transport of tracers if len(active_tracers) > 0: - adv_form += self.generate_tracer_transport_terms(state, active_tracers) + adv_form += self.generate_tracer_transport_terms(domain, active_tracers) # -------------------------------------------------------------------- # # Pressure Gradient Term @@ -1248,7 +1232,7 @@ def __init__(self, state, family, degree, Omega=None, # -------------------------------------------------------------------- # # Gravitational Term # -------------------------------------------------------------------- # - gravity_form = subject(prognostic(-b*inner(w, state.k)*dx, "u"), self.X) + gravity_form = subject(prognostic(-b*inner(w, domain.k)*dx, "u"), self.X) # -------------------------------------------------------------------- # # Divergence Term @@ -1274,6 +1258,5 @@ def __init__(self, state, family, degree, Omega=None, # -------------------------------------------------------------------- # # Linearise equations # -------------------------------------------------------------------- # - # TODO: add linearisation states for variables # Add linearisations to equations self.residual = self.generate_linear_terms(residual, self.linearisation_map) diff --git a/gusto/fields.py b/gusto/fields.py index ad231f39d..f3e674f73 100644 --- a/gusto/fields.py +++ b/gusto/fields.py @@ -1,6 +1,6 @@ -from firedrake import Function +from firedrake import Function, MixedElement, functionspaceimpl -__all__ = ["TimeLevelFields", "StateFields"] +__all__ = ["PrescribedFields", "TimeLevelFields", "StateFields"] class Fields(object): @@ -17,8 +17,7 @@ def __init__(self, equation): def add_field(self, name, space, subfield_names=None): """ - Adds a new field to the :class:`FieldCreator`. - + Adds a new field to the :class:`Fields` object. Args: name (str): the name of the prognostic variable. space (:class:`FunctionSpace`): the space to create the field in. @@ -42,15 +41,46 @@ def add_field(self, name, space, subfield_names=None): def __call__(self, name): """ - Returns a specified field from the :class:`FieldCreator`. + Returns a specified field from the :class:`Fields` object. + Args: + name (str): the name of the field. + Returns: + :class:`Function`: the desired field. + """ + return getattr(self, name) + + def __iter__(self): + """Returns an iterable of the contained fields.""" + return iter(self.fields) + + +class PrescribedFields(Fields): + """Object to hold and create a specified set of prescribed fields.""" + def __init__(self): + self.fields = [] + + def __call__(self, name, space=None): + """ + Returns a specified field from the :class:`PrescribedFields`. If a named + field does not yet exist in the :class:`PrescribedFields` object, then + the space argument must be specified so that it can be created and added + to the object. Args: name (str): the name of the field. + space (:class:`FunctionSpace`, optional): the function space to + create the field in. Defaults to None. Returns: :class:`Function`: the desired field. """ - return getattr(self, name) + if hasattr(self, name): + # Field already exists in object, so return it + return getattr(self, name) + else: + # Create field + self.add_field(name, space) + return getattr(self, name) def __iter__(self): """Returns an iterable of the contained fields.""" @@ -58,20 +88,50 @@ def __iter__(self): class StateFields(Fields): - """Creates the prognostic fields for the :class:`State` object.""" + """ + Container for all of the model's fields. + + The `StateFields` are a container for all the fields to be used by a time + stepper. In the case of the prognostic fields, these are pointers to the + time steppers' fields at the (n+1) time level. Prescribed fields are + pointers to the respective equation sets, while diagnostic fields are + created here. + """ - def __init__(self, *fields_to_dump): + def __init__(self, prognostic_fields, prescribed_fields, *fields_to_dump): """ Args: + prognostic_fields (:class:`Fields`): the (n+1) time level fields. + prescribed_fields (iter): an iterable of (name, function_space) + tuples, that are used to create the prescribed fields. *fields_to_dump (str): the names of fields to be dumped. """ self.fields = [] - self.output_specified = len(fields_to_dump) > 0 + output_specified = len(fields_to_dump) > 0 self.to_dump = set((fields_to_dump)) self.to_pickup = set(()) + self._field_types = [] + self._field_names = [] + + # Add pointers to prognostic fields + for field in prognostic_fields.fields: + # Don't add the mixed field + if type(field.ufl_element()) is not MixedElement: + # If fields_to_dump not specified, dump by default + to_dump = field.name() in fields_to_dump or not output_specified + self.__call__(field.name(), field=field, dump=to_dump, + pickup=True, field_type="prognostic") + else: + self.__call__(field.name(), field=field, dump=False, + pickup=False, field_type="prognostic") - def __call__(self, name, space=None, subfield_names=None, dump=True, - pickup=False): + for field in prescribed_fields.fields: + to_dump = field.name() in fields_to_dump + self.__call__(field.name(), field=field, dump=to_dump, + pickup=True, field_type="prescribed") + + def __call__(self, name, field=None, space=None, dump=True, pickup=False, + field_type=None): """ Returns a field from or adds a field to the :class:`StateFields`. @@ -79,13 +139,16 @@ def __call__(self, name, space=None, subfield_names=None, dump=True, the optional arguments must be specified so that it can be created and added to the :class:`StateFields`. + If "field" is specified, then the pointer to the field is added to the + :class:`StateFields` object. If "space" is specified, then the field + itself is created. + Args: name (str): name of the field to be returned/added. + field (:class:`Function`, optional): an existing field to be added + to the :class:`StateFields` object. space (:class:`FunctionSpace`, optional): the function space to create the field in. Defaults to None. - subfield_names (list, optional): a list of names of the constituent - prognostic variables to be created, if the provided space is - actually a :class:`MixedFunctionSpace`. Defaults to None. dump (bool, optional): whether the created field should be outputted. Defaults to True. pickup (bool, optional): whether the created field should be picked @@ -94,22 +157,68 @@ def __call__(self, name, space=None, subfield_names=None, dump=True, Returns: :class:`Function`: the specified field. """ - try: + if hasattr(self, name): + # Field already exists in object, so return it return getattr(self, name) - except AttributeError: - self.add_field(name, space, subfield_names) + else: + # Field does not yet exist in StateFields + if field is None and space is None: + raise ValueError(f'Field {name} does not exist in StateFields. ' + + 'Either field or space argument must be ' + + 'specified to add this field to StateFields') + elif field is not None and space is not None: + raise ValueError('Cannot specify both field and space to StateFields') + + if field is not None: + # Field pointer, so just add existing field to StateFields + assert isinstance(field, Function), \ + f'field argument for creating field {name} must be a Function, not {type(field)}' + setattr(self, name, field) + self.fields.append(field) + else: + # Create field + assert isinstance(space, functionspaceimpl.WithGeometry), \ + f'space argument for creating field {name} must be FunctionSpace, not {type(space)}' + self.add_field(name, space) + if dump: - if subfield_names is not None: - self.to_dump.update(subfield_names) - else: - self.to_dump.add(name) + self.to_dump.add(name) if pickup: - if subfield_names is not None: - self.to_pickup.update(subfield_names) + self.to_pickup.add(name) + + # Work out field type + if field_type is None: + # Prognostics can only be specified through __init__ + if pickup: + field_type = "prescribed" + elif dump: + field_type = "diagnostic" else: - self.to_pickup.add(name) + field_type = "derived" + else: + permitted_types = ["prognostic", "prescribed", "diagnostic", "derived"] + assert field_type in permitted_types, \ + f'field_type {field_type} not in permitted types {permitted_types}' + self._field_types.append(field_type) + self._field_names.append(name) + return getattr(self, name) + def field_type(self, field_name): + """ + Returns the type (e.g. prognostic/diagnostic) of a field held in the + :class:`StateFields`. + + Args: + field_name (str): name of the field to return the type of. + + Returns: + str: a string describing the type (e.g. prognostic) of the field. + """ + assert hasattr(self, field_name), f'StateFields has no field {field_name}' + idx = self._field_names.index(field_name) + return self._field_types[idx] + class TimeLevelFields(object): """Creates the fields required in the :class:`Timestepper` object.""" @@ -149,14 +258,15 @@ def add_fields(self, equation, levels=None): except AttributeError: setattr(self, level, Fields(equation)) - def initialise(self, state): + def initialise(self, state_fields): """ - Initialises the time fields from those currently in state + Initialises the time fields from those currently in the equation. - Args: state (:class:`State`): the model state object + Args: + state_fields (:class:`StateFields`): the model's field container. """ for field in self.n: - field.assign(state.fields(field.name())) + field.assign(state_fields(field.name())) self.np1(field.name()).assign(field) def update(self): diff --git a/gusto/forcing.py b/gusto/forcing.py index ee9cd8378..f09a77efc 100644 --- a/gusto/forcing.py +++ b/gusto/forcing.py @@ -33,7 +33,7 @@ def __init__(self, equation, alpha): self.field_name = equation.field_name implicit_terms = ["incompressibility", "sponge"] - dt = equation.state.dt + dt = equation.domain.dt W = equation.function_space self.x0 = Function(W) diff --git a/gusto/function_spaces.py b/gusto/function_spaces.py new file mode 100644 index 000000000..0e3a82e70 --- /dev/null +++ b/gusto/function_spaces.py @@ -0,0 +1,321 @@ +""" +This module contains routines to generate the compatible function spaces to be +used by the model. +""" + +from firedrake import (HDiv, FunctionSpace, FiniteElement, TensorProductElement, + interval) + +# TODO: there is danger here for confusion about degree, particularly for the CG +# spaces -- does a "CG" space with degree = 1 mean the "CG" space in the de Rham +# complex of degree 1 ("CG3"), or "CG1"? +# TODO: would it be better to separate creation of specific named spaces from +# the creation of the de Rham complex spaces? +# TODO: how do we create HCurl spaces if we want them? + + +class Spaces(object): + """Object to create and hold the model's finite element spaces.""" + def __init__(self, mesh): + """ + Args: + mesh (:class:`Mesh`): the model's mesh. + """ + self.mesh = mesh + self.extruded_mesh = hasattr(mesh, "_base_mesh") + self._initialised_base_spaces = False + + def __call__(self, name, family=None, degree=None, + horizontal_degree=None, vertical_degree=None, + V=None, overwrite_space=False): + """ + Returns a space, and also creates it if it is not created yet. + + If a space needs creating, it may be that more arguments (such as the + family and degree) need to be provided. Alternatively a space can be + passed in to be stored in the space container. + + For extruded meshes, it is possible to seperately specify the horizontal + and vertical degrees of the elements. Alternatively, if these degrees + should be the same then this can be specified through the "degree" + argument. + + Args: + name (str): the name of the space. + family (str, optional): name of the finite element family to be + created. Defaults to None. + degree (int, optional): the element degree used for the space. + Defaults to None, in which case the horizontal degree must be + provided. + horizontal_degree (int, optional): the horizontal degree of the + finite element space to be created. Defaults to None. + vertical_degree (int, optional): the vertical degree of the + finite element space to be created. Defaults to None. + V (:class:`FunctionSpace`, optional): an existing space, to be + stored in the creator object. If this is provided, it will be + added to the creator and no other action will be taken. This + space will be returned. Defaults to None. + overwrite_space (bool, optional): Logical to allow space existing in + container to be overwritten by an incoming space. Defaults to + False. + + Returns: + :class:`FunctionSpace`: the desired function space. + """ + + if hasattr(self, name) and (V is None or not overwrite_space): + # We have requested a space that should already have been created + if V is not None: + assert getattr(self, name) == V, \ + f'There is a conflict between the space {name} already ' + \ + 'existing in the space container, and the space being passed to it' + return getattr(self, name) + + else: + # Space does not exist in creator + if V is not None: + # The space itself has been provided (to add it to the creator) + value = V + + elif name == "DG1_equispaced": + # Special case as no degree arguments need providing + value = self.build_dg_space(1, 1, variant='equispaced', name='DG1_equispaced') + + else: + check_degree_args('Spaces', self.mesh, degree, horizontal_degree, vertical_degree) + + # Convert to horizontal and vertical degrees + horizontal_degree = degree if horizontal_degree is None else horizontal_degree + vertical_degree = degree if vertical_degree is None else vertical_degree + + # Loop through name and family combinations + if name == "HDiv" and family in ["BDM", "RT", "CG", "RTCF"]: + value = self.build_hdiv_space(family, horizontal_degree, vertical_degree) + elif name == "theta": + value = self.build_theta_space(horizontal_degree, vertical_degree) + elif family == "DG": + value = self.build_dg_space(horizontal_degree, vertical_degree, name=name) + elif family == "CG": + value = self.build_cg_space(horizontal_degree, vertical_degree, name=name) + else: + raise ValueError(f'There is no space corresponding to {name}') + setattr(self, name, value) + return value + + def build_compatible_spaces(self, family, horizontal_degree, + vertical_degree=None): + """ + Builds the sequence of compatible finite element spaces for the mesh. + + If the mesh is not extruded, this builds and returns the spaces: + (HDiv, DG). + If the mesh is extruded, this builds and returns the following spaces: + (HDiv, DG, theta). + The 'theta' space corresponds to the vertical component of the velocity. + + Args: + family (str): the family of the horizontal part of the HDiv space. + horizontal_degree (int): the polynomial degree of the horizontal + part of the DG space. + vertical_degree (int, optional): the polynomial degree of the + vertical part of the DG space. Defaults to None. Must be + specified if the mesh is extruded. + + Returns: + tuple: the created compatible :class:`FunctionSpace` objects. + """ + if self.extruded_mesh and not self._initialised_base_spaces: + self.build_base_spaces(family, horizontal_degree, vertical_degree) + Vu = self.build_hdiv_space(family, horizontal_degree, vertical_degree) + setattr(self, "HDiv", Vu) + Vdg = self.build_dg_space(horizontal_degree, vertical_degree, name='DG') + setattr(self, "DG", Vdg) + Vth = self.build_theta_space(horizontal_degree, vertical_degree) + setattr(self, "theta", Vth) + return Vu, Vdg, Vth + else: + Vu = self.build_hdiv_space(family, horizontal_degree+1) + setattr(self, "HDiv", Vu) + Vdg = self.build_dg_space(horizontal_degree, vertical_degree, name='DG') + setattr(self, "DG", Vdg) + return Vu, Vdg + + def build_base_spaces(self, family, horizontal_degree, vertical_degree): + """ + Builds the :class:`FiniteElement` objects for the base mesh. + + Args: + family (str): the family of the horizontal part of the HDiv space. + horizontal_degree (int): the polynomial degree of the horizontal + part of the DG space. + vertical_degree (int): the polynomial degree of the vertical part of + the DG space. + """ + cell = self.mesh._base_mesh.ufl_cell().cellname() + + # horizontal base spaces + self.S1 = FiniteElement(family, cell, horizontal_degree+1) + self.S2 = FiniteElement("DG", cell, horizontal_degree) + + # vertical base spaces + self.T0 = FiniteElement("CG", interval, vertical_degree+1) + self.T1 = FiniteElement("DG", interval, vertical_degree) + + self._initialised_base_spaces = True + + def build_hdiv_space(self, family, horizontal_degree, vertical_degree=None): + """ + Builds and returns the HDiv :class:`FunctionSpace`. + + Args: + family (str): the family of the horizontal part of the HDiv space. + horizontal_degree (int): the polynomial degree of the horizontal + part of the DG space from the de Rham complex. + vertical_degree (int, optional): the polynomial degree of the + vertical part of the the DG space from the de Rham complex. + Defaults to None. Must be specified if the mesh is extruded. + + Returns: + :class:`FunctionSpace`: the HDiv space. + """ + if self.extruded_mesh: + if not self._initialised_base_spaces: + if vertical_degree is None: + raise ValueError('vertical_degree must be specified to create HDiv space on an extruded mesh') + self.build_base_spaces(family, horizontal_degree, vertical_degree) + Vh_elt = HDiv(TensorProductElement(self.S1, self.T1)) + Vt_elt = TensorProductElement(self.S2, self.T0) + Vv_elt = HDiv(Vt_elt) + V_elt = Vh_elt + Vv_elt + else: + cell = self.mesh.ufl_cell().cellname() + V_elt = FiniteElement(family, cell, horizontal_degree) + return FunctionSpace(self.mesh, V_elt, name='HDiv') + + def build_dg_space(self, horizontal_degree, vertical_degree=None, variant=None, name='DG'): + """ + Builds and returns the DG :class:`FunctionSpace`. + + Args: + horizontal_degree (int): the polynomial degree of the horizontal + part of the DG space. + vertical_degree (int, optional): the polynomial degree of the + vertical part of the DG space. Defaults to None. Must be + specified if the mesh is extruded. + variant (str, optional): the variant of the underlying + :class:`FiniteElement` to use. Defaults to None, which will call + the default variant. + name (str, optional): name to assign to the function space. Default + is "DG". + + Returns: + :class:`FunctionSpace`: the DG space. + """ + assert not hasattr(self, name), f'There already exists a function space with name {name}' + + if self.extruded_mesh: + if vertical_degree is None: + raise ValueError('vertical_degree must be specified to create DG space on an extruded mesh') + if not self._initialised_base_spaces or self.T1.degree() != vertical_degree or self.T1.variant() != variant: + cell = self.mesh._base_mesh.ufl_cell().cellname() + S2 = FiniteElement("DG", cell, horizontal_degree, variant=variant) + T1 = FiniteElement("DG", interval, vertical_degree, variant=variant) + else: + S2 = self.S2 + T1 = self.T1 + V_elt = TensorProductElement(S2, T1) + else: + cell = self.mesh.ufl_cell().cellname() + V_elt = FiniteElement("DG", cell, horizontal_degree, variant=variant) + + return FunctionSpace(self.mesh, V_elt, name=name) + + def build_theta_space(self, horizontal_degree, vertical_degree): + """ + Builds and returns the 'theta' space. + + This corresponds to the non-Piola mapped space of the vertical component + of the velocity. The space will be discontinuous in the horizontal but + continuous in the vertical. + + Args: + horizontal_degree (int): the polynomial degree of the horizontal + part of the DG space from the de Rham complex. + vertical_degree (int): the polynomial degree of the vertical part of + the DG space from the de Rham complex. + + Raises: + AssertionError: the mesh is not extruded. + + Returns: + :class:`FunctionSpace`: the 'theta' space. + """ + assert self.extruded_mesh, 'Cannot create theta space if mesh is not extruded' + if not self._initialised_base_spaces: + cell = self.mesh._base_mesh.ufl_cell().cellname() + self.S2 = FiniteElement("DG", cell, horizontal_degree) + self.T0 = FiniteElement("CG", interval, vertical_degree+1) + V_elt = TensorProductElement(self.S2, self.T0) + return FunctionSpace(self.mesh, V_elt, name='theta') + + def build_cg_space(self, horizontal_degree, vertical_degree=None, name='CG'): + """ + Builds the continuous scalar space at the top of the de Rham complex. + + Args: + horizontal_degree (int): the polynomial degree of the horizontal + part of the CG space. + vertical_degree (int, optional): the polynomial degree of the + vertical part of the the CG space. Defaults to None. Must be + specified if the mesh is extruded. + name (str, optional): name to assign to the function space. Default + is "CG". + + Returns: + :class:`FunctionSpace`: the continuous space. + """ + assert not hasattr(self, name), f'There already exists a function space with name {name}' + + if self.extruded_mesh: + if vertical_degree is None: + raise ValueError('vertical_degree must be specified to create CG space on an extruded mesh') + cell = self.mesh._base_mesh.ufl_cell().cellname() + CG_hori = FiniteElement("CG", cell, horizontal_degree) + CG_vert = FiniteElement("CG", interval, vertical_degree) + V_elt = TensorProductElement(CG_hori, CG_vert) + else: + cell = self.mesh.ufl_cell().cellname() + V_elt = FiniteElement("CG", cell, horizontal_degree) + + return FunctionSpace(self.mesh, V_elt, name=name) + + +def check_degree_args(name, mesh, degree, horizontal_degree, vertical_degree): + """ + Check the degree arguments passed to either the :class:`Domain` or the + :class:`Spaces` object. This will raise errors if the arguments are not + appropriate. + + Args: + name (str): name of object to print out. + mesh (:class:`Mesh`): the model's mesh. + degree (int): the element degree. + horizontal_degree (int): the element degree used for the horizontal part + of a space. + vertical_degree (int): the element degree used for the vertical part + of a space. + """ + + extruded_mesh = hasattr(mesh, "_base_mesh") + + # Checks on degree arguments + if degree is None and horizontal_degree is None: + raise ValueError(f'Either "degree" or "horizontal_degree" must be passed to {name}') + if extruded_mesh and degree is None and vertical_degree is None: + raise ValueError(f'For extruded meshes, either "degree" or "vertical_degree" must be passed to {name}') + if degree is not None and horizontal_degree is not None: + raise ValueError(f'Cannot pass both "degree" and "horizontal_degree" to {name}') + if extruded_mesh and degree is not None and vertical_degree is not None: + raise ValueError(f'Cannot pass both "degree" and "vertical_degree" to {name}') + if not extruded_mesh and vertical_degree is not None: + raise ValueError(f'Cannot pass "vertical_degree" to {name} if mesh is not extruded') diff --git a/gusto/initialisation_tools.py b/gusto/initialisation_tools.py index 1c8aa9b32..5807686a1 100644 --- a/gusto/initialisation_tools.py +++ b/gusto/initialisation_tools.py @@ -60,7 +60,7 @@ def sphere_to_cartesian(mesh, u_zonal, u_merid): return as_vector((cartesian_u_expr, cartesian_v_expr, cartesian_w_expr)) -def incompressible_hydrostatic_balance(state, b0, p0, top=False, params=None): +def incompressible_hydrostatic_balance(equation, b0, p0, top=False, params=None): """ Gives a pressure field in hydrostatic-balance for the Incompressible eqns. @@ -70,7 +70,7 @@ def incompressible_hydrostatic_balance(state, b0, p0, top=False, params=None): with zero flow enforced at one of the boundaries. Args: - state (:class:`State`): the model's state. + equation (:class:`PrognosticEquation`): the model's equation object. b0 (:class:`ufl.Expr`): the input buoyancy field. p0 (:class:`Function`): the pressure to be returned. top (bool, optional): whether the no-flow boundary condition is enforced @@ -81,8 +81,9 @@ def incompressible_hydrostatic_balance(state, b0, p0, top=False, params=None): """ # get F - Vu = state.spaces("HDiv") - Vv = FunctionSpace(state.mesh, Vu.ufl_element()._elements[-1]) + domain = equation.domain + Vu = domain.spaces("HDiv") + Vv = FunctionSpace(equation.domain.mesh, Vu.ufl_element()._elements[-1]) v = TrialFunction(Vv) w = TestFunction(Vv) @@ -94,13 +95,13 @@ def incompressible_hydrostatic_balance(state, b0, p0, top=False, params=None): bcs = [DirichletBC(Vv, 0.0, bstring)] a = inner(w, v)*dx - L = inner(state.k, w)*b0*dx + L = inner(equation.domain.k, w)*b0*dx F = Function(Vv) solve(a == L, F, bcs=bcs) # define mixed function space - VDG = state.spaces("DG") + VDG = domain.spaces("DG") WV = (Vv)*(VDG) # get pprime @@ -136,7 +137,7 @@ def incompressible_hydrostatic_balance(state, b0, p0, top=False, params=None): p0.project(pprime) -def compressible_hydrostatic_balance(state, theta0, rho0, exner0=None, +def compressible_hydrostatic_balance(equation, theta0, rho0, exner0=None, top=False, exner_boundary=Constant(1.0), mr_t=None, solve_for_rho=False, @@ -151,7 +152,7 @@ def compressible_hydrostatic_balance(state, theta0, rho0, exner0=None, procedure for solving the resulting discrete systems. Args: - state (:class:`State`): the model's state. + equation (:class:`PrognosticEquation`): the model's equation object. theta0 (:class:`ufl.Expr`): the input (dry) potential temperature field. rho0 (:class:`Function`): the hydrostatically-balanced density to be found. @@ -175,16 +176,18 @@ def compressible_hydrostatic_balance(state, theta0, rho0, exner0=None, """ # Calculate hydrostatic Pi - VDG = state.spaces("DG") - Vu = state.spaces("HDiv") - Vv = FunctionSpace(state.mesh, Vu.ufl_element()._elements[-1]) + domain = equation.domain + parameters = equation.parameters + VDG = domain.spaces("DG") + Vu = domain.spaces("HDiv") + Vv = FunctionSpace(equation.domain.mesh, Vu.ufl_element()._elements[-1]) W = MixedFunctionSpace((Vv, VDG)) v, exner = TrialFunctions(W) dv, dexner = TestFunctions(W) - n = FacetNormal(state.mesh) + n = FacetNormal(equation.domain.mesh) - cp = state.parameters.cp + cp = parameters.cp # add effect of density of water upon theta theta = theta0 @@ -207,9 +210,9 @@ def compressible_hydrostatic_balance(state, theta0, rho0, exner0=None, arhs = -cp*inner(dv, n)*theta*exner_boundary*bmeasure # Possibly make g vary with spatial coordinates? - g = state.parameters.g + g = parameters.g - arhs -= g*inner(dv, state.k)*dx + arhs -= g*inner(dv, equation.domain.k)*dx bcs = [DirichletBC(W.sub(0), zero(), bstring)] @@ -239,16 +242,16 @@ def compressible_hydrostatic_balance(state, theta0, rho0, exner0=None, if solve_for_rho: w1 = Function(W) v, rho = w1.split() - rho.interpolate(thermodynamics.rho(state.parameters, theta0, exner)) + rho.interpolate(thermodynamics.rho(parameters, theta0, exner)) v, rho = split(w1) dv, dexner = TestFunctions(W) - exner = thermodynamics.exner_pressure(state.parameters, rho, theta0) + exner = thermodynamics.exner_pressure(parameters, rho, theta0) F = ( (cp*inner(v, dv) - cp*div(dv*theta)*exner)*dx + dexner*div(theta0*v)*dx + cp*inner(dv, n)*theta*exner_boundary*bmeasure ) - F += g*inner(dv, state.k)*dx + F += g*inner(dv, equation.domain.k)*dx rhoproblem = NonlinearVariationalProblem(F, w1, bcs=bcs) rhosolver = NonlinearVariationalSolver(rhoproblem, solver_parameters=params, options_prefix="rhosolver") @@ -256,7 +259,7 @@ def compressible_hydrostatic_balance(state, theta0, rho0, exner0=None, v, rho_ = w1.split() rho0.assign(rho_) else: - rho0.interpolate(thermodynamics.rho(state.parameters, theta0, exner)) + rho0.interpolate(thermodynamics.rho(parameters, theta0, exner)) def remove_initial_w(u): @@ -276,8 +279,9 @@ def remove_initial_w(u): u.assign(uin) -def saturated_hydrostatic_balance(state, theta_e, mr_t, exner0=None, - top=False, exner_boundary=Constant(1.0), +def saturated_hydrostatic_balance(equation, state_fields, theta_e, mr_t, + exner0=None, top=False, + exner_boundary=Constant(1.0), max_outer_solve_count=40, max_theta_solve_count=5, max_inner_solve_count=3): @@ -296,8 +300,8 @@ def saturated_hydrostatic_balance(state, theta_e, mr_t, exner0=None, converge to a solution. Args: - state (:class:`State`): the model's state object, through which the - prognostic variables are accessed. + equation (:class:`PrognosticEquation`): the model's equation object. + state_fields (:class:`StateFields`): the model's field container. theta_e (:class:`ufl.Expr`): expression for the desired wet equivalent potential temperature field. mr_t (:class:`ufl.Expr`): expression for the total moisture content. @@ -324,15 +328,17 @@ def saturated_hydrostatic_balance(state, theta_e, mr_t, exner0=None, number of iterations. """ - theta0 = state.fields('theta') - rho0 = state.fields('rho') - mr_v0 = state.fields('water_vapour') + theta0 = state_fields('theta') + rho0 = state_fields('rho') + mr_v0 = state_fields('water_vapour') # Calculate hydrostatic exner pressure + domain = equation.domain + parameters = equation.parameters Vt = theta0.function_space() Vr = rho0.function_space() - VDG = state.spaces("DG") + VDG = domain.spaces("DG") if any(deg > 2 for deg in VDG.ufl_element().degree()): logger.warning("default quadrature degree most likely not sufficient for this degree element") @@ -353,15 +359,15 @@ def saturated_hydrostatic_balance(state, theta_e, mr_t, exner0=None, delta = 0.8 # expressions for finding theta0 and mr_v0 from theta_e and mr_t - exner = thermodynamics.exner_pressure(state.parameters, rho_averaged, theta0) - p = thermodynamics.p(state.parameters, exner) - T = thermodynamics.T(state.parameters, theta0, exner, mr_v0) - r_v_expr = thermodynamics.r_sat(state.parameters, T, p) - theta_e_expr = thermodynamics.theta_e(state.parameters, T, p, mr_v0, mr_t) + exner = thermodynamics.exner_pressure(parameters, rho_averaged, theta0) + p = thermodynamics.p(parameters, exner) + T = thermodynamics.T(parameters, theta0, exner, mr_v0) + r_v_expr = thermodynamics.r_sat(parameters, T, p) + theta_e_expr = thermodynamics.theta_e(parameters, T, p, mr_v0, mr_t) for i in range(max_outer_solve_count): # solve for rho with theta_vd and w_v guesses - compressible_hydrostatic_balance(state, theta0, rho_h, top=top, + compressible_hydrostatic_balance(equation, theta0, rho_h, top=top, exner_boundary=exner_boundary, mr_t=mr_t, solve_for_rho=True) @@ -396,17 +402,18 @@ def saturated_hydrostatic_balance(state, theta_e, mr_t, exner0=None, raise RuntimeError('Hydrostatic balance solve has not converged within %i' % i, 'iterations') if exner0 is not None: - exner = thermodynamics.exner(state.parameters, rho0, theta0) + exner = thermodynamics.exner(parameters, rho0, theta0) exner0.interpolate(exner) # do one extra solve for rho - compressible_hydrostatic_balance(state, theta0, rho0, top=top, + compressible_hydrostatic_balance(equation, theta0, rho0, top=top, exner_boundary=exner_boundary, mr_t=mr_t, solve_for_rho=True) -def unsaturated_hydrostatic_balance(state, theta_d, H, exner0=None, - top=False, exner_boundary=Constant(1.0), +def unsaturated_hydrostatic_balance(equation, state_fields, theta_d, H, + exner0=None, top=False, + exner_boundary=Constant(1.0), max_outer_solve_count=40, max_inner_solve_count=20): """ @@ -423,8 +430,8 @@ def unsaturated_hydrostatic_balance(state, theta_d, H, exner0=None, These steps are iterated until we (hopefully) converge to a solution. Args: - state (:class:`State`): the model's state object, through which the - prognostic variables are accessed. + equation (:class:`PrognosticEquation`): the model's equation object. + state_fields (:class:`StateFields`): the model's field container. theta_d (:class:`ufl.Expr`): the specified dry potential temperature field. H (:class:`ufl.Expr`): the specified relative humidity field. @@ -449,18 +456,20 @@ def unsaturated_hydrostatic_balance(state, theta_d, H, exner0=None, number of iterations. """ - theta0 = state.fields('theta') - rho0 = state.fields('rho') - mr_v0 = state.fields('water_vapour') + theta0 = state_fields('theta') + rho0 = state_fields('rho') + mr_v0 = state_fields('water_vapour') # Calculate hydrostatic exner pressure + domain = equation.domain + parameters = equation.parameters Vt = theta0.function_space() Vr = rho0.function_space() - R_d = state.parameters.R_d - R_v = state.parameters.R_v + R_d = parameters.R_d + R_v = parameters.R_v epsilon = R_d / R_v - VDG = state.spaces("DG") + VDG = domain.spaces("DG") if any(deg > 2 for deg in VDG.ufl_element().degree()): logger.warning("default quadrature degree most likely not sufficient for this degree element") @@ -480,21 +489,21 @@ def unsaturated_hydrostatic_balance(state, theta_d, H, exner0=None, delta = 1.0 # make expressions for determining mr_v0 - exner = thermodynamics.exner_pressure(state.parameters, rho_averaged, theta0) - p = thermodynamics.p(state.parameters, exner) - T = thermodynamics.T(state.parameters, theta0, exner, mr_v0) - r_v_expr = thermodynamics.r_v(state.parameters, H, T, p) + exner = thermodynamics.exner_pressure(parameters, rho_averaged, theta0) + p = thermodynamics.p(parameters, exner) + T = thermodynamics.T(parameters, theta0, exner, mr_v0) + r_v_expr = thermodynamics.r_v(parameters, H, T, p) # make expressions to evaluate residual - exner_ev = thermodynamics.exner_pressure(state.parameters, rho_averaged, theta0) - p_ev = thermodynamics.p(state.parameters, exner_ev) - T_ev = thermodynamics.T(state.parameters, theta0, exner_ev, mr_v0) - RH_ev = thermodynamics.RH(state.parameters, mr_v0, T_ev, p_ev) + exner_ev = thermodynamics.exner_pressure(parameters, rho_averaged, theta0) + p_ev = thermodynamics.p(parameters, exner_ev) + T_ev = thermodynamics.T(parameters, theta0, exner_ev, mr_v0) + RH_ev = thermodynamics.RH(parameters, mr_v0, T_ev, p_ev) RH = Function(Vt) for i in range(max_outer_solve_count): # solve for rho with theta_vd and w_v guesses - compressible_hydrostatic_balance(state, theta0, rho_h, top=top, + compressible_hydrostatic_balance(equation, theta0, rho_h, top=top, exner_boundary=exner_boundary, mr_t=mr_v0, solve_for_rho=True) @@ -525,10 +534,10 @@ def unsaturated_hydrostatic_balance(state, theta_d, H, exner0=None, raise RuntimeError('Hydrostatic balance solve has not converged within %i' % i, 'iterations') if exner0 is not None: - exner = thermodynamics.exner_pressure(state.parameters, rho0, theta0) + exner = thermodynamics.exner_pressure(parameters, rho0, theta0) exner0.interpolate(exner) # do one extra solve for rho - compressible_hydrostatic_balance(state, theta0, rho0, top=top, + compressible_hydrostatic_balance(equation, theta0, rho0, top=top, exner_boundary=exner_boundary, mr_t=mr_v0, solve_for_rho=True) diff --git a/gusto/state.py b/gusto/io.py similarity index 58% rename from gusto/state.py rename to gusto/io.py index b4c4ed2cb..974c16db5 100644 --- a/gusto/state.py +++ b/gusto/io.py @@ -1,234 +1,27 @@ -""" -Provides the model's state object, which controls IO and other core functions. - -The model's :class:`State` object is defined in this module. It controls various -input/output (IO) aspects, as well as setting up the compatible finite element -spaces and holding the mesh. In some ways it acts as a bucket, holding core -parts of the model. -""" +"""Provides the model's IO, which controls input, output and diagnostics.""" + from os import path, makedirs import itertools from netCDF4 import Dataset import sys import time -from gusto.diagnostics import Diagnostics, Perturbation, SteadyStateError -from firedrake import (FiniteElement, TensorProductElement, HDiv, - FunctionSpace, VectorFunctionSpace, - interval, Function, Mesh, functionspaceimpl, - File, SpatialCoordinate, sqrt, Constant, inner, - op2, DumbCheckpoint, FILE_CREATE, FILE_READ, interpolate, - CellNormal, cross, as_vector) +from gusto.diagnostics import Diagnostics +from firedrake import (FiniteElement, TensorProductElement, VectorFunctionSpace, + interval, Function, Mesh, functionspaceimpl, File, + op2, DumbCheckpoint, FILE_CREATE, FILE_READ) import numpy as np from gusto.configuration import logger, set_log_handler -from gusto.fields import StateFields - -__all__ = ["State"] - - -class SpaceCreator(object): - """Object to create and hold the model's finite element spaces.""" - def __init__(self, mesh): - """ - Args: - mesh (:class:`Mesh`): the model's mesh. - """ - self.mesh = mesh - self.extruded_mesh = hasattr(mesh, "_base_mesh") - self._initialised_base_spaces = False - - def __call__(self, name, family=None, degree=None, V=None): - """ - Returns a space, and also creates it if it is not created yet. - - If a space needs creating, it may be that more arguments (such as the - family and degree) need to be provided. Alternatively a space can be - passed in to be stored in the creator. - - Args: - name (str): the name of the space. - family (str, optional): name of the finite element family to be - created. Defaults to None. - degree (int, optional): the degree of the finite element space to be - created. Defaults to None. - V (:class:`FunctionSpace`, optional): an existing space, to be - stored in the creator object. If this is provided, it will be - added to the creator and no other action will be taken. This - space will be returned. Defaults to None. - - Returns: - :class:`FunctionSpace`: the desired function space. - """ - - try: - return getattr(self, name) - except AttributeError: - if V is not None: - value = V - elif name == "HDiv" and family in ["BDM", "RT", "CG", "RTCF"]: - value = self.build_hdiv_space(family, degree) - elif name == "theta": - value = self.build_theta_space(degree) - elif name == "DG1_equispaced": - value = self.build_dg_space(1, variant='equispaced') - elif family == "DG": - value = self.build_dg_space(degree) - elif family == "CG": - value = self.build_cg_space(degree) - else: - raise ValueError(f'State has no space corresponding to {name}') - setattr(self, name, value) - return value - - def build_compatible_spaces(self, family, degree): - """ - Builds the sequence of compatible finite element spaces for the mesh. - - If the mesh is not extruded, this builds and returns the spaces: - (HDiv, DG). - If the mesh is extruded, this builds and returns the following spaces: - (HDiv, DG, theta). - The 'theta' space corresponds to the vertical component of the velocity. - - Args: - family (str): the family of the horizontal part of the HDiv space. - degree (int): the polynomial degree of the DG space. - - Returns: - tuple: the created compatible :class:`FunctionSpace` objects. - """ - if self.extruded_mesh and not self._initialised_base_spaces: - self.build_base_spaces(family, degree) - Vu = self.build_hdiv_space(family, degree) - setattr(self, "HDiv", Vu) - Vdg = self.build_dg_space(degree) - setattr(self, "DG", Vdg) - Vth = self.build_theta_space(degree) - setattr(self, "theta", Vth) - return Vu, Vdg, Vth - else: - Vu = self.build_hdiv_space(family, degree) - setattr(self, "HDiv", Vu) - Vdg = self.build_dg_space(degree) - setattr(self, "DG", Vdg) - return Vu, Vdg - - def build_base_spaces(self, family, degree): - """ - Builds the :class:`FiniteElement` objects for the base mesh. - - Args: - family (str): the family of the horizontal part of the HDiv space. - degree (int): the polynomial degree of the DG space. - """ - cell = self.mesh._base_mesh.ufl_cell().cellname() - - # horizontal base spaces - self.S1 = FiniteElement(family, cell, degree+1) - self.S2 = FiniteElement("DG", cell, degree) - - # vertical base spaces - self.T0 = FiniteElement("CG", interval, degree+1) - self.T1 = FiniteElement("DG", interval, degree) - - self._initialised_base_spaces = True - - def build_hdiv_space(self, family, degree): - """ - Builds and returns the HDiv :class:`FunctionSpace`. - - Args: - family (str): the family of the horizontal part of the HDiv space. - degree (int): the polynomial degree of the space. - - Returns: - :class:`FunctionSpace`: the HDiv space. - """ - if self.extruded_mesh: - if not self._initialised_base_spaces: - self.build_base_spaces(family, degree) - Vh_elt = HDiv(TensorProductElement(self.S1, self.T1)) - Vt_elt = TensorProductElement(self.S2, self.T0) - Vv_elt = HDiv(Vt_elt) - V_elt = Vh_elt + Vv_elt - else: - cell = self.mesh.ufl_cell().cellname() - V_elt = FiniteElement(family, cell, degree+1) - return FunctionSpace(self.mesh, V_elt, name='HDiv') - - def build_dg_space(self, degree, variant=None): - """ - Builds and returns the DG :class:`FunctionSpace`. - - Args: - degree (int): the polynomial degree of the space. - variant (str): the variant of the underlying :class:`FiniteElement` - to use. Defaults to None, which will call the default variant. - - Returns: - :class:`FunctionSpace`: the DG space. - """ - if self.extruded_mesh: - if not self._initialised_base_spaces or self.T1.degree() != degree or self.T1.variant() != variant: - cell = self.mesh._base_mesh.ufl_cell().cellname() - S2 = FiniteElement("DG", cell, degree, variant=variant) - T1 = FiniteElement("DG", interval, degree, variant=variant) - else: - S2 = self.S2 - T1 = self.T1 - V_elt = TensorProductElement(S2, T1) - else: - cell = self.mesh.ufl_cell().cellname() - V_elt = FiniteElement("DG", cell, degree, variant=variant) - name = f'DG{degree}_equispaced' if variant == 'equispaced' else f'DG{degree}' - return FunctionSpace(self.mesh, V_elt, name=name) - - def build_theta_space(self, degree): - """ - Builds and returns the 'theta' space. - - This corresponds to the non-Piola mapped space of the vertical component - of the velocity. The space will be discontinuous in the horizontal but - continuous in the vertical. - - Args: - degree (int): degree of the corresponding density space. - - Raises: - AssertionError: the mesh is not extruded. - - Returns: - :class:`FunctionSpace`: the 'theta' space. - """ - assert self.extruded_mesh - if not self._initialised_base_spaces: - cell = self.mesh._base_mesh.ufl_cell().cellname() - self.S2 = FiniteElement("DG", cell, degree) - self.T0 = FiniteElement("CG", interval, degree+1) - V_elt = TensorProductElement(self.S2, self.T0) - return FunctionSpace(self.mesh, V_elt, name='Vtheta') - - def build_cg_space(self, degree): - """ - Builds the continuous scalar space at the top of the de Rham complex. - - Args: - degree (int): degree of the continuous space. - Returns: - :class:`FunctionSpace`: the continuous space. - """ - return FunctionSpace(self.mesh, "CG", degree, name=f'CG{degree}') +__all__ = ["IO"] class PointDataOutput(object): """Object for outputting field point data.""" - def __init__(self, filename, ndt, field_points, description, + def __init__(self, filename, field_points, description, field_creator, comm, tolerance=None, create=True): """ Args: filename (str): name of file to output to. - ndt (int): number of time points to output at. TODO: remove as this - is unused. field_points (list): some iterable of pairs, matching fields with arrays of evaluation points: (field_name, evaluation_points). description (str): a description of the simulation to be included in @@ -337,17 +130,17 @@ def __init__(self, filename, diagnostics, description, comm, create=True): for diagnostic in diagnostics.available_diagnostics: group.createVariable(diagnostic, np.float64, ("time", )) - def dump(self, state, t): + def dump(self, state_fields, t): """ Output the global diagnostics. - state (:class:`State`): the model's state object. + state_fields (:class:`StateFields`): the model's field container. t (float): simulation time at which the output occurs. """ diagnostics = [] for fname in self.diagnostics.fields: - field = state.fields(fname) + field = state_fields(fname) for dname in self.diagnostics.available_diagnostics: diagnostic = getattr(self.diagnostics, dname) diagnostics.append((fname, dname, diagnostic(field))) @@ -362,24 +155,16 @@ def dump(self, state, t): var[idx:idx + 1] = value -class State(object): - """Keeps the model's mesh and variables, and controls its IO.""" +class IO(object): + """Controls the model's input, output and diagnostics.""" - def __init__(self, mesh, dt, - output=None, - parameters=None, - diagnostics=None, - diagnostic_fields=None): + def __init__(self, domain, output, diagnostics=None, diagnostic_fields=None): """ Args: - mesh (:class:`Mesh`): the model's mesh. - dt (:class:`Constant`): the time taken to perform a single model - step. If a float or int is passed, it will be cast to a - :class:`Constant`. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. output (:class:`OutputParameters`, optional): holds and describes the options for outputting. Defaults to None. - parameters (:class:`Configuration`, optional): an object containing - the model's physical parameters. Defaults to None. diagnostics (:class:`Diagnostics`, optional): object holding and controlling the model's diagnostics. Defaults to None. diagnostic_fields (list, optional): an iterable of `DiagnosticField` @@ -389,13 +174,9 @@ def __init__(self, mesh, dt, RuntimeError: if no output is provided. TypeError: if `dt` cannot be cast to a :class:`Constant`. """ - - if output is None: - # TODO: maybe this shouldn't be an optional argument then? - raise RuntimeError("You must provide a directory name for dumping results") - else: - self.output = output - self.parameters = parameters + self.domain = domain + self.mesh = domain.mesh + self.output = output if diagnostics is not None: self.diagnostics = diagnostics @@ -406,78 +187,55 @@ def __init__(self, mesh, dt, else: self.diagnostic_fields = [] - # The mesh - self.mesh = mesh - - self.spaces = SpaceCreator(mesh) - if self.output.dumplist is None: - self.output.dumplist = [] - self.fields = StateFields(*self.output.dumplist) - self.dumpdir = None self.dumpfile = None self.to_pickup = None - # figure out if we're on a sphere - try: - self.on_sphere = (mesh._base_mesh.geometric_dimension() == 3 and mesh._base_mesh.topological_dimension() == 2) - except AttributeError: - self.on_sphere = (mesh.geometric_dimension() == 3 and mesh.topological_dimension() == 2) - - # build the vertical normal and define perp for 2d geometries - dim = mesh.topological_dimension() - if self.on_sphere: - x = SpatialCoordinate(mesh) - R = sqrt(inner(x, x)) - self.k = interpolate(x/R, mesh.coordinates.function_space()) - if dim == 2: - outward_normals = CellNormal(mesh) - self.perp = lambda u: cross(outward_normals, u) - else: - kvec = [0.0]*dim - kvec[dim-1] = 1.0 - self.k = Constant(kvec) - if dim == 2: - self.perp = lambda u: as_vector([-u[1], u[0]]) - # setup logger logger.setLevel(output.log_level) - set_log_handler(mesh.comm) - if parameters is not None: + set_log_handler(self.mesh.comm) + + def log_parameters(self, equation): + """ + Logs an equation's physical parameters that take non-default values. + + Args: + equation (:class:`PrognosticEquation`): the model's equation which + contains any physical parameters used in the model run. + """ + if hasattr(equation, 'parameters') and equation.parameters is not None: logger.info("Physical parameters that take non-default values:") - logger.info(", ".join("%s: %s" % (k, float(v)) for (k, v) in vars(parameters).items())) - - # Constant to hold current time - self.t = Constant(0.0) - if type(dt) is Constant: - self.dt = dt - elif type(dt) in (float, int): - self.dt = Constant(dt) - else: - raise TypeError(f'dt must be a Constant, float or int, not {type(dt)}') + logger.info(", ".join("%s: %s" % (k, float(v)) for (k, v) in vars(equation.parameters).items())) - def setup_diagnostics(self): - """Concatenates the various types of diagnostic field.""" - for name in self.output.perturbation_fields: - f = Perturbation(name) - self.diagnostic_fields.append(f) + def setup_diagnostics(self, state_fields): + """ + Prepares the I/O for computing the model's global diagnostics and + diagnostic fields. - for name in self.output.steady_state_error_fields: - f = SteadyStateError(self, name) - self.diagnostic_fields.append(f) + Args: + state_fields (:class:`StateFields`): the model's field container. + """ - fields = set([f.name() for f in self.fields]) - field_deps = [(d, sorted(set(d.required_fields).difference(fields),)) for d in self.diagnostic_fields] + diagnostic_names = [diagnostic.name for diagnostic in self.diagnostic_fields] + non_diagnostics = [fname for fname in state_fields._field_names if state_fields.field_type(fname) != "diagnostic" or fname not in diagnostic_names] + # Filter out non-diagnostic fields + field_deps = [(d, sorted(set(d.required_fields).difference(non_diagnostics),)) for d in self.diagnostic_fields] schedule = topo_sort(field_deps) self.diagnostic_fields = schedule for diagnostic in self.diagnostic_fields: - diagnostic.setup(self) + diagnostic.setup(self.domain, state_fields) self.diagnostics.register(diagnostic.name) - def setup_dump(self, t, tmax, pickup=False): + # Register fields for global diagnostics + # TODO: it should be possible to specify which global diagnostics are used + for fname in state_fields._field_names: + if fname in state_fields.to_dump: + self.diagnostics.register(fname) + + def setup_dump(self, state_fields, t, tmax, pickup=False): """ Sets up a series of things used for outputting. @@ -487,6 +245,7 @@ def setup_dump(self, t, tmax, pickup=False): checkpointing file. Args: + state_fields (:class:`StateFields`): the model's field container. t (float): the current model time. tmax (float): the end time of the model's simulation. pickup (bool, optional): whether to pick up the model's initial @@ -520,7 +279,7 @@ def setup_dump(self, t, tmax, pickup=False): comm=self.mesh.comm) # make list of fields to dump - self.to_dump = [f for f in self.fields if f.name() in self.fields.to_dump] + self.to_dump = [f for f in state_fields.fields if f.name() in state_fields.to_dump] # make dump counter self.dumpcount = itertools.count() @@ -537,7 +296,7 @@ def setup_dump(self, t, tmax, pickup=False): # make functions on latlon mesh, as specified by dumplist_latlon self.to_dump_latlon = [] for name in self.output.dumplist_latlon: - f = self.fields(name) + f = state_fields(name) field = Function( functionspaceimpl.WithGeometry.create( f.function_space(), mesh_ll), @@ -557,11 +316,10 @@ def setup_dump(self, t, tmax, pickup=False): if len(self.output.point_data) > 0: # set up point data output pointdata_filename = self.dumpdir+"/point_data.nc" - ndt = int(tmax/float(self.dt)) - self.pointdata_output = PointDataOutput(pointdata_filename, ndt, + self.pointdata_output = PointDataOutput(pointdata_filename, self.output.point_data, self.output.dirname, - self.fields, + state_fields, self.mesh.comm, self.output.tolerance, create=not pickup) @@ -582,21 +340,21 @@ def setup_dump(self, t, tmax, pickup=False): mode=FILE_CREATE) # make list of fields to pickup (this doesn't include # diagnostic fields) - self.to_pickup = [f for f in self.fields if f.name() in self.fields.to_pickup] + self.to_pickup = [state_fields(f) for f in state_fields.to_pickup] # if we want to checkpoint then make a checkpoint counter if self.output.checkpoint: self.chkptcount = itertools.count() # dump initial fields - self.dump(t) + self.dump(state_fields, t) - def pickup_from_checkpoint(self): + def pickup_from_checkpoint(self, state_fields): """Picks up the model's variables from a checkpoint file.""" # TODO: this duplicates some code from setup_dump. Can this be avoided? # It is because we don't know if we are picking up or setting dump first if self.to_pickup is None: - self.to_pickup = [f for f in self.fields if f.name() in self.fields.to_pickup] + self.to_pickup = [state_fields(f) for f in state_fields.to_pickup] # Set dumpdir if has not been done already if self.dumpdir is None: self.dumpdir = path.join("results", self.output.dirname) @@ -619,7 +377,7 @@ def pickup_from_checkpoint(self): return t - def dump(self, t): + def dump(self, state_fields, t): """ Dumps all of the required model output. @@ -628,6 +386,7 @@ def dump(self, t): a checkpoint file if specified. Args: + state_fields (:class:`StateFields`): the model's field container. t (float): the simulation's current time. """ output = self.output @@ -635,15 +394,15 @@ def dump(self, t): # Diagnostics: # Compute diagnostic fields for field in self.diagnostic_fields: - field(self) + field.compute() if output.dump_diagnostics: # Output diagnostic data - self.diagnostic_output.dump(self, t) + self.diagnostic_output.dump(state_fields, t) if len(output.point_data) > 0 and (next(self.pddumpcount) % output.pddumpfreq) == 0: # Output pointwise data - self.pointdata_output.dump(self.fields, t) + self.pointdata_output.dump(state_fields, t) # Dump all the fields to the checkpointing file (backup version) if output.checkpoint and (next(self.chkptcount) % output.chkptfreq) == 0: @@ -659,51 +418,13 @@ def dump(self, t): if len(output.dumplist_latlon) > 0: self.dumpfile_ll.write(*self.to_dump_latlon) - def initialise(self, initial_conditions): - """ - Initialise the state's prognostic variables. - - Args: - initial_conditions (list): an iterable of pairs: (field_name, expr), - where 'field_name' is the string giving the name of the - prognostic field and expr is the :class:`ufl.Expr` whose value - is used to set the initial field. - """ - for name, ic in initial_conditions: - f_init = getattr(self.fields, name) - f_init.assign(ic) - f_init.rename(name) - - def set_reference_profiles(self, reference_profiles): - """ - Initialise the state's reference profiles. - - reference_profiles (list): an iterable of pairs: (field_name, expr), - where 'field_name' is the string giving the name of the - reference profile field expr is the :class:`ufl.Expr` whose - value is used to set the reference field. - """ - for name, profile in reference_profiles: - if name+'bar' in self.fields: - # For reference profiles already added to state, allow - # interpolation from expressions - ref = self.fields(name+'bar') - elif isinstance(profile, Function): - # Need to add reference profile to state so profile must be - # a Function - ref = self.fields(name+'bar', space=profile.function_space(), dump=False) - else: - raise ValueError(f'When initialising reference profile {name}' - + ' the passed profile must be a Function') - ref.interpolate(profile) - def get_latlon_mesh(mesh): """ Construct a planar latitude-longitude mesh from a spherical mesh. Args: - mesh (:class:`State`): the mesh on which the simulation is performed. + mesh (:class:`Mesh`): the mesh on which the simulation is performed. """ coords_orig = mesh.coordinates coords_fs = coords_orig.function_space() @@ -779,6 +500,7 @@ def topo_sort(field_deps): name2field = dict((f.name, f) for f, _ in field_deps) # map node: (input_deps, output_deps) graph = dict((f.name, (list(deps), [])) for f, deps in field_deps) + roots = [] for f, input_deps in field_deps: if len(input_deps) == 0: diff --git a/gusto/linear_solvers.py b/gusto/linear_solvers.py index 3ddd5b13d..547fe99d7 100644 --- a/gusto/linear_solvers.py +++ b/gusto/linear_solvers.py @@ -14,6 +14,7 @@ from firedrake.petsc import flatten_parameters from pyop2.profiling import timed_function, timed_region +from gusto.active_tracers import TracerVariableType from gusto.configuration import logger, DEBUG from gusto.labels import linearisation, time_derivative, hydrostatic from gusto import thermodynamics @@ -28,11 +29,10 @@ class TimesteppingSolver(object, metaclass=ABCMeta): """Base class for timestepping linear solvers for Gusto.""" - def __init__(self, state, equations, alpha=0.5, solver_parameters=None, + def __init__(self, equations, alpha=0.5, solver_parameters=None, overwrite_solver_parameters=False): """ Args: - state (:class:`State`): the model's state object. equations (:class:`PrognosticEquation`): the model's equation. alpha (float, optional): the semi-implicit off-centring factor. Defaults to 0.5. A value of 1 is fully-implicit. @@ -44,8 +44,8 @@ def __init__(self, state, equations, alpha=0.5, solver_parameters=None, update the default parameters with the `solver_parameters` passed in. Defaults to False. """ - self.state = state self.equations = equations + self.dt = equations.domain.dt self.alpha = alpha if solver_parameters is not None: @@ -122,12 +122,11 @@ class CompressibleSolver(TimesteppingSolver): 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'}}} - def __init__(self, state, equations, alpha=0.5, + def __init__(self, equations, alpha=0.5, quadrature_degree=None, solver_parameters=None, - overwrite_solver_parameters=False, moisture=None): + overwrite_solver_parameters=False): """ Args: - state (:class:`State`): the model's state object. equations (:class:`PrognosticEquation`): the model's equation. alpha (float, optional): the semi-implicit off-centring factor. Defaults to 0.5. A value of 1 is fully-implicit. @@ -141,15 +140,13 @@ def __init__(self, state, equations, alpha=0.5, `solver_parameters` that have been passed in. If False then update the default parameters with the `solver_parameters` passed in. Defaults to False. - moisture (list, optional): list of names of moisture fields. - Defaults to None. """ - self.moisture = moisture + self.equations = equations if quadrature_degree is not None: self.quadrature_degree = quadrature_degree else: - dgspace = state.spaces("DG") + dgspace = equations.domain.spaces("DG") if any(deg > 2 for deg in dgspace.ufl_element().degree()): logger.warning("default quadrature degree most likely not sufficient for this degree element") self.quadrature_degree = (5, 5) @@ -164,20 +161,20 @@ def __init__(self, state, equations, alpha=0.5, # Turn monitor on for the trace system self.solver_parameters["condensed_field"]["ksp_monitor_true_residual"] = None - super().__init__(state, equations, alpha, solver_parameters, + super().__init__(equations, alpha, solver_parameters, overwrite_solver_parameters) @timed_function("Gusto:SolverSetup") def _setup_solver(self): - state = self.state - dt = state.dt + equations = self.equations + dt = self.dt beta_ = dt*self.alpha - cp = state.parameters.cp - Vu = state.spaces("HDiv") - Vu_broken = FunctionSpace(state.mesh, BrokenElement(Vu.ufl_element())) - Vtheta = state.spaces("theta") - Vrho = state.spaces("DG") + cp = equations.parameters.cp + Vu = equations.domain.spaces("HDiv") + Vu_broken = FunctionSpace(equations.domain.mesh, BrokenElement(Vu.ufl_element())) + Vtheta = equations.domain.spaces("theta") + Vrho = equations.domain.spaces("DG") # Store time-stepping coefficients as UFL Constants beta = Constant(beta_) @@ -185,7 +182,7 @@ def _setup_solver(self): h_deg = Vrho.ufl_element().degree()[0] v_deg = Vrho.ufl_element().degree()[1] - Vtrace = FunctionSpace(state.mesh, "HDiv Trace", degree=(h_deg, v_deg)) + Vtrace = FunctionSpace(equations.domain.mesh, "HDiv Trace", degree=(h_deg, v_deg)) # Split up the rhs vector (symbolically) self.xrhs = Function(self.equations.function_space) @@ -196,17 +193,16 @@ def _setup_solver(self): w, phi, dl = TestFunctions(M) u, rho, l0 = TrialFunctions(M) - n = FacetNormal(state.mesh) + n = FacetNormal(equations.domain.mesh) # Get background fields - thetabar = state.fields("thetabar") - rhobar = state.fields("rhobar") - exnerbar = thermodynamics.exner_pressure(state.parameters, rhobar, thetabar) - exnerbar_rho = thermodynamics.dexner_drho(state.parameters, rhobar, thetabar) - exnerbar_theta = thermodynamics.dexner_dtheta(state.parameters, rhobar, thetabar) + _, rhobar, thetabar = split(equations.X_ref)[0:3] + exnerbar = thermodynamics.exner_pressure(equations.parameters, rhobar, thetabar) + exnerbar_rho = thermodynamics.dexner_drho(equations.parameters, rhobar, thetabar) + exnerbar_theta = thermodynamics.dexner_dtheta(equations.parameters, rhobar, thetabar) # Analytical (approximate) elimination of theta - k = state.k # Upward pointing unit vector + k = equations.domain.k # Upward pointing unit vector theta = -dot(k, u)*dot(k, grad(thetabar))*beta + theta_in # Only include theta' (rather than exner') in the vertical @@ -231,13 +227,21 @@ def V(u): ds_tbp = (ds_t(degree=(self.quadrature_degree)) + ds_b(degree=(self.quadrature_degree))) - # Add effect of density of water upon theta - if self.moisture is not None: - water_t = Function(Vtheta).assign(0.0) - for water in self.moisture: - water_t += self.state.fields(water) - theta_w = theta / (1 + water_t) - thetabar_w = thetabar / (1 + water_t) + # Add effect of density of water upon theta, using moisture reference profiles + # TODO: Explore if this is the right thing to do for the linear problem + if equations.active_tracers is not None: + mr_t = Constant(0.0)*thetabar + for tracer in equations.active_tracers: + if tracer.chemical == 'H2O': + if tracer.variable_type == TracerVariableType.mixing_ratio: + idx = equations.field_names.index(tracer.name) + mr_bar = split(equations.X_ref)[idx] + mr_t += mr_bar + else: + raise NotImplementedError('Only mixing ratio tracers are implemented') + + theta_w = theta / (1 + mr_t) + thetabar_w = thetabar / (1 + mr_t) else: theta_w = theta thetabar_w = thetabar @@ -260,18 +264,12 @@ def L_tr(f): rho_avg_prb = LinearVariationalProblem(a_tr, L_tr(rhobar), rhobar_avg) exner_avg_prb = LinearVariationalProblem(a_tr, L_tr(exnerbar), exnerbar_avg) - rho_avg_solver = LinearVariationalSolver(rho_avg_prb, - solver_parameters=cg_ilu_parameters, - options_prefix='rhobar_avg_solver') - exner_avg_solver = LinearVariationalSolver(exner_avg_prb, - solver_parameters=cg_ilu_parameters, - options_prefix='exnerbar_avg_solver') - - with timed_region("Gusto:HybridProjectRhobar"): - rho_avg_solver.solve() - - with timed_region("Gusto:HybridProjectExnerbar"): - exner_avg_solver.solve() + self.rho_avg_solver = LinearVariationalSolver(rho_avg_prb, + solver_parameters=cg_ilu_parameters, + options_prefix='rhobar_avg_solver') + self.exner_avg_solver = LinearVariationalSolver(exner_avg_prb, + solver_parameters=cg_ilu_parameters, + options_prefix='exnerbar_avg_solver') # "broken" u, rho, and trace system # NOTE: no ds_v integrals since equations are defined on @@ -306,6 +304,7 @@ def L_tr(f): + dl*dot(u, n)*(ds_tbp + ds_vp) ) + # TODO: can we get this term using FML? # contribution of the sponge term if hasattr(self.equations, "mu"): eqn += dt*self.equations.mu*inner(w, k)*inner(u, k)*dx @@ -364,6 +363,13 @@ def solve(self, xrhs, dy): """ self.xrhs.assign(xrhs) + # TODO: can we avoid computing these each time the solver is called? + with timed_region("Gusto:HybridProjectRhobar"): + self.rho_avg_solver.solve() + + with timed_region("Gusto:HybridProjectExnerbar"): + self.exner_avg_solver.solve() + # Solve the hybridized system self.hybridized_solver.solve() @@ -405,13 +411,6 @@ class IncompressibleSolver(TimesteppingSolver): (1) Analytically eliminate b (introduces error near topography) (2) Solve resulting system for (u,p) using a hybrid-mixed method (3) Reconstruct b - - :arg state: a :class:`.State` object containing everything else. - :arg solver_parameters: (optional) Solver parameters. - :arg overwrite_solver_parameters: boolean, if True use only the - solver_parameters that have been passed in, if False then - update the default solver parameters with the solver_parameters - passed in. """ solver_parameters = { @@ -430,12 +429,12 @@ class IncompressibleSolver(TimesteppingSolver): @timed_function("Gusto:SolverSetup") def _setup_solver(self): - state = self.state # just cutting down line length a bit - dt = state.dt + equation = self.equations # just cutting down line length a bit + dt = self.dt beta_ = dt*self.alpha - Vu = state.spaces("HDiv") - Vb = state.spaces("theta") - Vp = state.spaces("DG") + Vu = equation.domain.spaces("HDiv") + Vb = equation.domain.spaces("theta") + Vp = equation.domain.spaces("DG") # Store time-stepping coefficients as UFL Constants beta = Constant(beta_) @@ -450,10 +449,10 @@ def _setup_solver(self): u, p = TrialFunctions(M) # Get background fields - bbar = state.fields("bbar") + bbar = split(equation.X_ref)[2] # Analytical (approximate) elimination of theta - k = state.k # Upward pointing unit vector + k = equation.domain.k # Upward pointing unit vector b = -dot(k, u)*dot(k, grad(bbar))*beta + b_in # vertical projection @@ -570,7 +569,7 @@ def __init__(self, equation, alpha): lambda t: Term(t.get(linearisation).form, t.labels), drop) - dt = equation.state.dt + dt = equation.domain.dt W = equation.function_space beta = dt*alpha diff --git a/gusto/physics.py b/gusto/physics.py index 7685e4dc7..11b221a05 100644 --- a/gusto/physics.py +++ b/gusto/physics.py @@ -29,8 +29,20 @@ class Physics(object, metaclass=ABCMeta): """Base class for the parametrisation of physical processes for Gusto.""" - def __init__(self): - pass + def __init__(self, equation, parameters=None): + """ + Args: + equation (:class:`PrognosticEquationSet`): the model's equation. + parameters (:class:`Configuration`, optional): parameters containing + the values of gas constants. Defaults to None, in which case the + parameters are obtained from the equation. + """ + + self.equation = equation + if parameters is None and hasattr(equation, 'parameters'): + self.parameters = equation.parameters + else: + self.parameters = parameters @abstractmethod def evaluate(self): @@ -56,25 +68,28 @@ class SaturationAdjustment(Physics): A filter is applied to prevent generation of negative mixing ratios. """ - def __init__(self, equation, parameters, vapour_name='water_vapour', - cloud_name='cloud_water', latent_heat=True): + def __init__(self, equation, vapour_name='water_vapour', + cloud_name='cloud_water', latent_heat=True, parameters=None): """ Args: equation (:class:`PrognosticEquationSet`): the model's equation. - parameters (:class:`Configuration`): an object containing the - model's physical parameters. vapour_name (str, optional): name of the water vapour variable. Defaults to 'water_vapour'. cloud_name (str, optional): name of the cloud water variable. Defaults to 'cloud_water'. latent_heat (bool, optional): whether to have latent heat exchange feeding back from the phase change. Defaults to True. + parameters (:class:`Configuration`, optional): parameters containing + the values of gas constants. Defaults to None, in which case the + parameters are obtained from the equation. Raises: NotImplementedError: currently this is only implemented for the CompressibleEulerEquations. """ + super().__init__(equation, parameters=parameters) + # TODO: make a check on the variable type of the active tracers # if not a mixing ratio, we need to convert to mixing ratios # this will be easier if we change equations to have dictionary of @@ -85,8 +100,8 @@ def __init__(self, equation, parameters, vapour_name='water_vapour', assert cloud_name in equation.field_names, f"Field {cloud_name} does not exist in the equation set" # Make prognostic for physics scheme + parameters = self.parameters self.X = Function(equation.X.function_space()) - self.equation = equation self.latent_heat = latent_heat # Vapour and cloud variables are needed for every form of this scheme @@ -109,9 +124,7 @@ def __init__(self, equation, parameters, vapour_name='water_vapour', V_idxs.append(theta_idx) # need to evaluate rho at theta-points, and do this via recovery - # TODO: make this bit of code neater if possible using domain object - v_deg = V.ufl_element().degree()[1] - boundary_method = BoundaryMethod.extruded if v_deg == 1 else None + boundary_method = BoundaryMethod.extruded if equation.domain.vertical_degree == 0 else None rho_averaged = Function(V) self.rho_recoverer = Recoverer(rho, rho_averaged, boundary_method=boundary_method) @@ -242,13 +255,14 @@ class Fallout(Physics): for Cartesian geometry. """ - def __init__(self, equation, rain_name, state, moments=AdvectedMoments.M3): + def __init__(self, equation, rain_name, domain, moments=AdvectedMoments.M3): """ Args: equation (:class:`PrognosticEquationSet`): the model's equation. rain_name (str, optional): name of the rain variable. Defaults to 'rain'. - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. moments (int, optional): an :class:`AdvectedMoments` enumerator, representing the number of moments of the size distribution of raindrops to be transported. Defaults to `AdvectedMoments.M3`. @@ -265,14 +279,15 @@ def __init__(self, equation, rain_name, state, moments=AdvectedMoments.M3): rain = self.X.split()[rain_idx] test = equation.tests[rain_idx] - Vu = state.spaces("HDiv") - v = state.fields('rainfall_velocity', Vu) + Vu = domain.spaces("HDiv") + # TODO: there must be a better way than forcing this into the equation + v = equation.prescribed_fields(name='rainfall_velocity', space=Vu) # -------------------------------------------------------------------- # # Create physics term -- which is actually a transport term # -------------------------------------------------------------------- # - adv_term = advection_form(state, test, rain, outflow=True) + adv_term = advection_form(domain, test, rain, outflow=True) # Add rainfall velocity by replacing transport_velocity in term adv_term = adv_term.label_map(identity, map_if_true=lambda t: Term( @@ -289,7 +304,7 @@ def __init__(self, equation, rain_name, state, moments=AdvectedMoments.M3): if moments == AdvectedMoments.M0: # all rain falls at terminal velocity terminal_velocity = Constant(5) # in m/s - v.project(-terminal_velocity*state.k) + v.project(-terminal_velocity*domain.k) elif moments == AdvectedMoments.M3: # this advects the third moment M3 of the raindrop # distribution, which corresponds to the mean mass @@ -323,7 +338,7 @@ def __init__(self, equation, rain_name, state, moments=AdvectedMoments.M3): raise NotImplementedError('Currently we only have implementations for zero and one moment schemes for rainfall. Valid options are AdvectedMoments.M0 and AdvectedMoments.M3') if moments != AdvectedMoments.M0: - self.determine_v = Projector(-v_expression*state.k, v) + self.determine_v = Projector(-v_expression*domain.k, v) def evaluate(self, x_in, dt): """ @@ -450,24 +465,28 @@ class EvaporationOfRain(Physics): is the virtual dry potential temperature. """ - def __init__(self, equation, parameters, rain_name='rain', - vapour_name='water_vapour', latent_heat=True): + def __init__(self, equation, rain_name='rain', vapour_name='water_vapour', + latent_heat=True, parameters=None): """ Args: equation (:class:`PrognosticEquationSet`): the model's equation. - parameters (:class:`Configuration`): an object containing the - model's physical parameters. cloud_name (str, optional): name of the rain variable. Defaults to 'rain'. vapour_name (str, optional): name of the water vapour variable. Defaults to 'water_vapour'. latent_heat (bool, optional): whether to have latent heat exchange feeding back from the phase change. Defaults to True. + parameters (:class:`Configuration`, optional): parameters containing + the values of gas constants. Defaults to None, in which case the + parameters are obtained from the equation. Raises: NotImplementedError: currently this is only implemented for the CompressibleEulerEquations. """ + + super().__init__(equation, parameters=parameters) + # TODO: make a check on the variable type of the active tracers # if not a mixing ratio, we need to convert to mixing ratios # this will be easier if we change equations to have dictionary of @@ -479,7 +498,7 @@ def __init__(self, equation, parameters, rain_name='rain', # Make prognostic for physics scheme self.X = Function(equation.X.function_space()) - self.equation = equation + parameters = self.parameters self.latent_heat = latent_heat # Vapour and cloud variables are needed for every form of this scheme @@ -502,9 +521,7 @@ def __init__(self, equation, parameters, rain_name='rain', V_idxs.append(theta_idx) # need to evaluate rho at theta-points, and do this via recovery - # TODO: make this bit of code neater if possible using domain object - v_deg = V.ufl_element().degree()[1] - boundary_method = BoundaryMethod.extruded if v_deg == 1 else None + boundary_method = BoundaryMethod.extruded if equation.domain.vertical_degree == 1 else None rho_averaged = Function(V) self.rho_recoverer = Recoverer(rho, rho_averaged, boundary_method=boundary_method) @@ -604,7 +621,7 @@ def evaluate(self, x_in, dt): interpolator.interpolate() -class InstantRain(object): +class InstantRain(Physics): """ The process of converting vapour above the saturation curve to rain. @@ -616,8 +633,8 @@ class InstantRain(object): """ def __init__(self, equation, saturation_curve, vapour_name="water_vapour", - rain_name=None, parameters=None, convective_feedback=False, - set_tau_to_dt=False): + rain_name=None, convective_feedback=False, set_tau_to_dt=False, + parameters=None): """ Args: equation (:class: 'PrognosticEquationSet'): the model's equation. @@ -627,17 +644,20 @@ def __init__(self, equation, saturation_curve, vapour_name="water_vapour", Defaults to "water_vapour". rain_name (str, optional): name of the rain variable. Defaults to None. - parameters (:class: 'Configuration', optional): an object - containing the model's physical parameters. Defaults to None - but required if convective_feedback is True. convective_feedback (bool, optional): True if the conversion of vapour affects the height equation. Defaults to False. set_tau_to_dt (bool, optional): True if the timescale for the conversion is equal to the timestep and False if not. If False then the user must provide a timescale, tau, that gets passed to the parameters list. + parameters (:class:`Configuration`, optional): parameters containing + the values of gas constants. Defaults to None, in which case the + parameters are obtained from the equation. """ + super().__init__(equation, parameters=None) + + parameters = self.parameters self.convective_feedback = convective_feedback self.set_tau_to_dt = set_tau_to_dt @@ -650,7 +670,7 @@ def __init__(self, equation, saturation_curve, vapour_name="water_vapour", if self.convective_feedback: assert "D" in equation.field_names, "Depth field must exist for convective feedback" - assert parameters is not None, "You must provide parameters for convective feedback" + assert parameters.gamma is not None, "If convective feedback is used, gamma parameter must be specified" # obtain function space and functions; vapour needed for all cases W = equation.function_space diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 54a85742c..e5406df66 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -73,11 +73,12 @@ def new_apply(self, x_out, x_in): class TimeDiscretisation(object, metaclass=ABCMeta): """Base class for time discretisation schemes.""" - def __init__(self, state, field_name=None, solver_parameters=None, + def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None): """ Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. field_name (str, optional): name of the field to be evolved. Defaults to None. solver_parameters (dict, optional): dictionary of parameters to @@ -89,10 +90,11 @@ def __init__(self, state, field_name=None, solver_parameters=None, to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. """ - self.state = state + self.domain = domain self.field_name = field_name + self.equation = None - self.dt = self.state.dt + self.dt = domain.dt self.limiter = limiter @@ -125,11 +127,12 @@ def setup(self, equation, uadv=None, apply_bcs=True, *active_labels): *active_labels (:class:`Label`): labels indicating which terms of the equation to include. """ + self.equation = equation self.residual = equation.residual - if self.field_name is not None: + if self.field_name is not None and hasattr(equation, "field_names"): self.idx = equation.field_names.index(self.field_name) - self.fs = self.state.fields(self.field_name).function_space() + self.fs = equation.spaces[self.idx] self.residual = self.residual.label_map( lambda t: t.get(prognostic) == self.field_name, lambda t: Term( @@ -172,7 +175,7 @@ def setup(self, equation, uadv=None, apply_bcs=True, *active_labels): # construct the embedding space if not specified if options.embedding_space is None: V_elt = BrokenElement(self.fs.ufl_element()) - self.fs = FunctionSpace(self.state.mesh, V_elt) + self.fs = FunctionSpace(self.domain.mesh, V_elt) else: self.fs = options.embedding_space self.xdg_in = Function(self.fs) @@ -180,7 +183,7 @@ def setup(self, equation, uadv=None, apply_bcs=True, *active_labels): if self.idx is None: self.x_projected = Function(equation.function_space) else: - self.x_projected = Function(self.state.fields(self.field_name).function_space()) + self.x_projected = Function(equation.spaces[self.idx]) new_test = TestFunction(self.fs) parameters = {'ksp_type': 'cg', 'pc_type': 'bjacobi', @@ -204,7 +207,7 @@ def setup(self, equation, uadv=None, apply_bcs=True, *active_labels): if self.discretisation_option == "supg": # construct tau, if it is not specified - dim = self.state.mesh.topological_dimension() + dim = self.domain.mesh.topological_dimension() if options.tau is not None: # if tau is provided, check that is has the right size tau = options.tau @@ -253,7 +256,10 @@ def setup(self, equation, uadv=None, apply_bcs=True, *active_labels): if self.discretisation_option == "recovered": # set up the necessary functions - self.x_in = Function(self.state.fields(self.field_name).function_space()) + if self.idx is not None: + self.x_in = Function(equation.spaces[self.idx]) + else: + self.x_in = Function(equation.function_space) # Operator to recover to higher discontinuous space self.x_recoverer = ReversibleRecoverer(self.x_in, self.xdg_in, options) @@ -371,14 +377,17 @@ def replace_transport_term(self): # Do the options specify a different ibp to the old transport term? if old_transport_term.labels['ibp'] != self.options.ibp: # Set up a new transport term - field = self.state.fields(self.field_name) + if self.idx is not None: + field = self.equation.X.split()[self.idx] + else: + field = self.equation.X test = TestFunction(self.fs) # Set up new transport term (depending on the type of transport equation) if old_transport_term.labels['transport'] == TransportEquationType.advective: - new_transport_term = advection_form(self.state, test, field, ibp=self.options.ibp) + new_transport_term = advection_form(self.domain, test, field, ibp=self.options.ibp) elif old_transport_term.labels['transport'] == TransportEquationType.conservative: - new_transport_term = continuity_form(self.state, test, field, ibp=self.options.ibp) + new_transport_term = continuity_form(self.domain, test, field, ibp=self.options.ibp) else: raise NotImplementedError(f'Replacement of transport term not implemented yet for {old_transport_term.labels["transport"]}') @@ -434,11 +443,12 @@ def apply(self, x_out, x_in): class ExplicitTimeDiscretisation(TimeDiscretisation): """Base class for explicit time discretisations.""" - def __init__(self, state, field_name=None, subcycles=None, + def __init__(self, domain, field_name=None, subcycles=None, solver_parameters=None, limiter=None, options=None): """ Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. field_name (str, optional): name of the field to be evolved. Defaults to None. subcycles (int, optional): the number of sub-steps to perform. @@ -452,7 +462,7 @@ def __init__(self, state, field_name=None, subcycles=None, to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. """ - super().__init__(state, field_name, + super().__init__(domain, field_name, solver_parameters=solver_parameters, limiter=limiter, options=options) @@ -774,11 +784,12 @@ class BackwardEuler(TimeDiscretisation): The backward Euler method for operator F is the most simple implicit scheme: y^(n+1) = y^n + dt*F[y^(n+1)]. """ - def __init__(self, state, field_name=None, solver_parameters=None, + def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None): """ Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. field_name (str, optional): name of the field to be evolved. Defaults to None. subcycles (int, optional): the number of sub-steps to perform. @@ -797,7 +808,7 @@ def __init__(self, state, field_name=None, solver_parameters=None, """ if isinstance(options, (EmbeddedDGOptions, RecoveryOptions)): raise NotImplementedError("Only SUPG advection options have been implemented for this time discretisation") - super().__init__(state=state, field_name=field_name, + super().__init__(domain=domain, field_name=field_name, solver_parameters=solver_parameters, limiter=limiter, options=options) @@ -844,11 +855,12 @@ class ThetaMethod(TimeDiscretisation): for off-centring parameter theta. """ - def __init__(self, state, field_name=None, theta=None, + def __init__(self, domain, field_name=None, theta=None, solver_parameters=None, options=None): """ Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. field_name (str, optional): name of the field to be evolved. Defaults to None. theta (float, optional): the off-centring parameter. theta = 1 @@ -876,7 +888,7 @@ def __init__(self, state, field_name=None, theta=None, 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'} - super().__init__(state, field_name, + super().__init__(domain, field_name, solver_parameters=solver_parameters, options=options) @@ -926,11 +938,12 @@ class ImplicitMidpoint(ThetaMethod): It is equivalent to the "theta" method with theta = 1/2. """ - def __init__(self, state, field_name=None, solver_parameters=None, + def __init__(self, domain, field_name=None, solver_parameters=None, options=None): """ Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. field_name (str, optional): name of the field to be evolved. Defaults to None. solver_parameters (dict, optional): dictionary of parameters to @@ -940,7 +953,7 @@ def __init__(self, state, field_name=None, solver_parameters=None, to control the "wrapper" methods, such as Embedded DG or a recovery method. Defaults to None. """ - super().__init__(state, field_name, theta=0.5, + super().__init__(domain, field_name, theta=0.5, solver_parameters=solver_parameters, options=options) @@ -948,11 +961,12 @@ def __init__(self, state, field_name=None, solver_parameters=None, class MultilevelTimeDiscretisation(TimeDiscretisation): """Base class for multi-level timesteppers""" - def __init__(self, state, field_name=None, solver_parameters=None, + def __init__(self, domain, field_name=None, solver_parameters=None, limiter=None, options=None): """ Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. field_name (str, optional): name of the field to be evolved. Defaults to None. solver_parameters (dict, optional): dictionary of parameters to @@ -966,7 +980,7 @@ def __init__(self, state, field_name=None, solver_parameters=None, """ if isinstance(options, (EmbeddedDGOptions, RecoveryOptions)): raise NotImplementedError("Only SUPG advection options have been implemented for this time discretisation") - super().__init__(state=state, field_name=field_name, + super().__init__(domain=domain, field_name=field_name, solver_parameters=solver_parameters, limiter=limiter, options=options) self.initial_timesteps = 0 diff --git a/gusto/timeloop.py b/gusto/timeloop.py index 988677f70..6de21b00b 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -1,15 +1,16 @@ """Classes for controlling the timestepping loop.""" from abc import ABCMeta, abstractmethod, abstractproperty -from firedrake import Function, Projector +from firedrake import Function, Projector, Constant from pyop2.profiling import timed_stage from gusto.configuration import logger +from gusto.equations import PrognosticEquationSet from gusto.forcing import Forcing from gusto.fml.form_manipulation_labelling import drop from gusto.labels import (transport, diffusion, time_derivative, linearisation, prognostic, physics) from gusto.linear_solvers import LinearTimesteppingSolver -from gusto.fields import TimeLevelFields +from gusto.fields import TimeLevelFields, StateFields from gusto.time_discretisation import ExplicitTimeDiscretisation __all__ = ["Timestepper", "SplitPhysicsTimestepper", "SemiImplicitQuasiNewton", @@ -19,19 +20,24 @@ class BaseTimestepper(object, metaclass=ABCMeta): """Base class for timesteppers.""" - def __init__(self, equation, state): + def __init__(self, equation, io): """ Args: equation (:class:`PrognosticEquation`): the prognostic equation. - state (:class:`State`): the model's state object + io (:class:`IO`): the model's object for controlling input/output. """ self.equation = equation - self.state = state + self.io = io + self.dt = self.equation.domain.dt + self.t = Constant(0.0) + self.reference_profiles_initialised = False self.setup_fields() self.setup_scheme() + self.io.log_parameters(equation) + @abstractproperty def transporting_velocity(self): return NotImplementedError @@ -61,39 +67,65 @@ def run(self, t, tmax, pickup=False): pickup: (bool): specify whether to pickup from a previous run """ - state = self.state - if pickup: - t = state.pickup_from_checkpoint() + t = self.io.pickup_from_checkpoint(self.fields) - state.setup_diagnostics() + self.io.setup_diagnostics(self.fields) with timed_stage("Dump output"): - state.setup_dump(t, tmax, pickup) - - state.t.assign(t) + self.io.setup_dump(self.fields, t, tmax, pickup) - self.x.initialise(state) + self.t.assign(t) - while float(state.t) < tmax - 0.5*float(state.dt): - logger.info(f'at start of timestep, t={float(state.t)}, dt={float(state.dt)}') + while float(self.t) < tmax - 0.5*float(self.dt): + logger.info(f'at start of timestep, t={float(self.t)}, dt={float(self.dt)}') self.x.update() self.timestep() - for field in self.x.np1: - state.fields(field.name()).assign(field) - - state.t.assign(state.t + state.dt) + self.t.assign(self.t + self.dt) with timed_stage("Dump output"): - state.dump(float(state.t)) + self.io.dump(self.fields, float(self.t)) - if state.output.checkpoint: - state.chkpt.close() + if self.io.output.checkpoint: + self.io.chkpt.close() - logger.info(f'TIMELOOP complete. t={float(state.t)}, tmax={tmax}') + logger.info(f'TIMELOOP complete. t={float(self.t)}, tmax={tmax}') + + def set_reference_profiles(self, reference_profiles): + """ + Initialise the model's reference profiles. + + reference_profiles (list): an iterable of pairs: (field_name, expr), + where 'field_name' is the string giving the name of the reference + profile field expr is the :class:`ufl.Expr` whose value is used to + set the reference field. + """ + for field_name, profile in reference_profiles: + if field_name+'_bar' in self.fields: + # For reference profiles already added to state, allow + # interpolation from expressions + ref = self.fields(field_name+'_bar') + elif isinstance(profile, Function): + # Need to add reference profile to state so profile must be + # a Function + ref = self.fields(field_name+'_bar', space=profile.function_space(), dump=False) + else: + raise ValueError(f'When initialising reference profile {field_name}' + + ' the passed profile must be a Function') + ref.interpolate(profile) + + # Assign profile to X_ref belonging to equation + if isinstance(self.equation, PrognosticEquationSet): + assert field_name in self.equation.field_names, \ + f'Cannot set reference profile as field {field_name} not found' + idx = self.equation.field_names.index(field_name) + X_ref = self.equation.X_ref.split()[idx] + X_ref.assign(ref) + + self.reference_profiles_initialised = True class Timestepper(BaseTimestepper): @@ -101,16 +133,16 @@ class Timestepper(BaseTimestepper): Implements a timeloop by applying a scheme to a prognostic equation. """ - def __init__(self, equation, scheme, state): + def __init__(self, equation, scheme, io): """ Args: equation (:class:`PrognosticEquation`): the prognostic equation scheme (:class:`TimeDiscretisation`): the scheme to use to timestep the prognostic equation - state (:class:`State`): the model's state object + io (:class:`IO`): the model's object for controlling input/output. """ self.scheme = scheme - super().__init__(equation=equation, state=state) + super().__init__(equation=equation, io=io) @property def transporting_velocity(self): @@ -118,6 +150,8 @@ def transporting_velocity(self): def setup_fields(self): self.x = TimeLevelFields(self.equation, self.scheme.nlevels) + self.fields = StateFields(self.x.np1, self.equation.prescribed_fields, + *self.io.output.dumplist) def setup_scheme(self): self.scheme.setup(self.equation, self.transporting_velocity) @@ -140,13 +174,13 @@ class SplitPhysicsTimestepper(Timestepper): scheme to be applied to the physics terms than the prognostic equation. """ - def __init__(self, equation, scheme, state, physics_schemes=None): + def __init__(self, equation, scheme, io, physics_schemes=None): """ Args: equation (:class:`PrognosticEquation`): the prognostic equation scheme (:class:`TimeDiscretisation`): the scheme to use to timestep the prognostic equation - state (:class:`State`): the model's state object + io (:class:`IO`): the model's object for controlling input/output. physics_schemes: (list, optional): a list of :class:`Physics` and :class:`TimeDiscretisation` options describing physical parametrisations and timestepping schemes to use for each. @@ -156,7 +190,9 @@ def __init__(self, equation, scheme, state, physics_schemes=None): self.equation = equation self.scheme = scheme - self.state = state + self.io = io + self.dt = self.equation.domain.dt + self.t = Constant(0.0) self.setup_fields() self.setup_scheme() @@ -178,6 +214,8 @@ def transporting_velocity(self): def setup_fields(self): self.x = TimeLevelFields(self.equation, self.scheme.nlevels) + self.fields = StateFields(self.x.np1, self.equation.prescribed_fields, + *self.io.output.dumplist) def setup_scheme(self): from gusto.labels import coriolis, pressure_gradient @@ -199,14 +237,14 @@ def timestep(self): class SemiImplicitQuasiNewton(BaseTimestepper): """ Implements a semi-implicit quasi-Newton discretisation, - with Strang splitting and auxilliary semi-Lagrangian transport. + with Strang splitting and auxiliary semi-Lagrangian transport. The timestep consists of an outer loop applying the transport and an inner loop to perform the quasi-Newton interations for the fast-wave terms. """ - def __init__(self, equation_set, state, transport_schemes, + def __init__(self, equation_set, io, transport_schemes, auxiliary_equations_and_schemes=None, linear_solver=None, diffusion_schemes=None, @@ -216,7 +254,7 @@ def __init__(self, equation_set, state, transport_schemes, Args: equation_set (:class:`PrognosticEquationSet`): the prognostic equation set to be solved - state (:class:`State`) the model's state object + io (:class:`IO`): the model's object for controlling input/output. transport_schemes: iterable of ``(field_name, scheme)`` pairs indicating the name of the field (str) to transport, and the :class:`TimeDiscretisation` to use @@ -265,17 +303,21 @@ def __init__(self, equation_set, state, transport_schemes, assert scheme.field_name in equation_set.field_names self.diffusion_schemes.append((scheme.field_name, scheme)) - super().__init__(equation_set, state) - if auxiliary_equations_and_schemes is not None: for eqn, scheme in auxiliary_equations_and_schemes: - self.x.add_fields(eqn) - scheme.setup(eqn, self.transporting_velocity) + assert not hasattr(eqn, "field_names"), 'Cannot use auxiliary schemes with multiple fields' self.auxiliary_schemes = [ (eqn.field_name, scheme) for eqn, scheme in auxiliary_equations_and_schemes] else: + auxiliary_equations_and_schemes = [] self.auxiliary_schemes = [] + self.auxiliary_equations_and_schemes = auxiliary_equations_and_schemes + + super().__init__(equation_set, io) + + for aux_eqn, aux_scheme in self.auxiliary_equations_and_schemes: + aux_scheme.setup(aux_eqn, self.transporting_velocity) self.tracers_to_copy = [] for name in equation_set.field_names: @@ -319,6 +361,13 @@ def setup_fields(self): """Sets up time levels n, star, p and np1""" self.x = TimeLevelFields(self.equation, 1) self.x.add_fields(self.equation, levels=("star", "p")) + for aux_eqn, _ in self.auxiliary_equations_and_schemes: + self.x.add_fields(aux_eqn) + # Prescribed fields for auxiliary eqns should come from prognostics of + # other equations, so only the prescribed fields of the main equation + # need passing to StateFields + self.fields = StateFields(self.x.np1, self.equation.prescribed_fields, + *self.io.output.dumplist) def setup_scheme(self): """Sets up transport, diffusion and physics schemes""" @@ -379,7 +428,7 @@ def timestep(self): xrhs -= xnp1(self.field_name) with timed_stage("Implicit solve"): - self.linear_solver.solve(xrhs, dy) # solves linear system and places result in state.dy + self.linear_solver.solve(xrhs, dy) # solves linear system and places result in dy xnp1X = xnp1(self.field_name) xnp1X += dy @@ -401,19 +450,34 @@ def timestep(self): for _, scheme in self.physics_schemes: scheme.apply(xnp1(scheme.field_name), xnp1(scheme.field_name)) + def run(self, t, tmax, pickup=False): + """ + Runs the model for the specified time, from t to tmax. + + Args: + t (float): the start time of the run + tmax (float): the end time of the run + pickup: (bool): specify whether to pickup from a previous run + """ + + assert self.reference_profiles_initialised, \ + 'Reference profiles for must be initialised to use Semi-Implicit Timestepper' + + super().run(t, tmax, pickup=pickup) + class PrescribedTransport(Timestepper): """ Implements a timeloop with a prescibed transporting velocity """ - def __init__(self, equation, scheme, state, physics_schemes=None, + def __init__(self, equation, scheme, io, physics_schemes=None, prescribed_transporting_velocity=None): """ Args: equation (:class:`PrognosticEquation`): the prognostic equation scheme (:class:`TimeDiscretisation`): the scheme to use to timestep the prognostic equation - state (:class:`State`): the model's state object + io (:class:`IO`): the model's object for controlling input/output. physics_schemes: (list, optional): a list of :class:`Physics` and :class:`TimeDiscretisation` options describing physical parametrisations and timestepping schemes to use for each. @@ -427,7 +491,7 @@ def __init__(self, equation, scheme, state, physics_schemes=None, updated. Defaults to None. """ - super().__init__(equation, scheme, state) + super().__init__(equation, scheme, io) if physics_schemes is not None: self.physics_schemes = physics_schemes @@ -442,17 +506,19 @@ def __init__(self, equation, scheme, state, physics_schemes=None, if prescribed_transporting_velocity is not None: self.velocity_projection = Projector( - prescribed_transporting_velocity(self.state.t), - self.state.fields('u')) + prescribed_transporting_velocity(self.t), + self.fields('u')) else: self.velocity_projection = None @property def transporting_velocity(self): - return self.state.fields('u') + return self.fields('u') def setup_fields(self): self.x = TimeLevelFields(self.equation, self.scheme.nlevels) + self.fields = StateFields(self.x.np1, self.equation.prescribed_fields, + *self.io.output.dumplist) def setup_scheme(self): self.scheme.setup(self.equation, self.transporting_velocity) diff --git a/gusto/transport_forms.py b/gusto/transport_forms.py index 3ea6a468b..ae9d14d7b 100644 --- a/gusto/transport_forms.py +++ b/gusto/transport_forms.py @@ -14,12 +14,13 @@ "advection_equation_circulation_form", "linear_continuity_form"] -def linear_advection_form(state, test, qbar): +def linear_advection_form(domain, test, qbar): """ The form corresponding to the linearised advective transport operator. Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. test (:class:`TestFunction`): the test function. qbar (:class:`ufl.Expr`): the variable to be transported. @@ -27,22 +28,23 @@ def linear_advection_form(state, test, qbar): :class:`LabelledForm`: a labelled transport form. """ - ubar = Function(state.spaces("HDiv")) + ubar = Function(domain.spaces("HDiv")) # TODO: why is there a k here? - L = test*dot(ubar, state.k)*dot(state.k, grad(qbar))*dx + L = test*dot(ubar, domain.k)*dot(domain.k, grad(qbar))*dx form = transporting_velocity(L, ubar) return transport(form, TransportEquationType.advective) -def linear_continuity_form(state, test, qbar, facet_term=False): +def linear_continuity_form(domain, test, qbar, facet_term=False): """ The form corresponding to the linearised continuity transport operator. Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. test (:class:`TestFunction`): the test function. qbar (:class:`ufl.Expr`): the variable to be transported. facet_term (bool, optional): whether to include interior facet terms. @@ -52,14 +54,14 @@ def linear_continuity_form(state, test, qbar, facet_term=False): :class:`LabelledForm`: a labelled transport form. """ - Vu = state.spaces("HDiv") + Vu = domain.spaces("HDiv") ubar = Function(Vu) L = qbar*test*div(ubar)*dx if facet_term: - n = FacetNormal(state.mesh) - Vu = state.spaces("HDiv") + n = FacetNormal(domain.mesh) + Vu = domain.spaces("HDiv") dS_ = (dS_v + dS_h) if Vu.extruded else dS L += jump(ubar*test, n)*avg(qbar)*dS_ @@ -68,7 +70,7 @@ def linear_continuity_form(state, test, qbar, facet_term=False): return transport(form, TransportEquationType.conservative) -def advection_form(state, test, q, ibp=IntegrateByParts.ONCE, outflow=False): +def advection_form(domain, test, q, ibp=IntegrateByParts.ONCE, outflow=False): u""" The form corresponding to the advective transport operator. @@ -77,7 +79,8 @@ def advection_form(state, test, q, ibp=IntegrateByParts.ONCE, outflow=False): form is integrated by parts. Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. test (:class:`TestFunction`): the test function. q (:class:`ufl.Expr`): the variable to be transported. ibp (:class:`IntegrateByParts`, optional): an enumerator representing @@ -96,7 +99,7 @@ def advection_form(state, test, q, ibp=IntegrateByParts.ONCE, outflow=False): if outflow and ibp == IntegrateByParts.NEVER: raise ValueError("outflow is True and ibp is None are incompatible options") - Vu = state.spaces("HDiv") + Vu = domain.spaces("HDiv") dS_ = (dS_v + dS_h) if Vu.extruded else dS ubar = Function(Vu) @@ -106,7 +109,7 @@ def advection_form(state, test, q, ibp=IntegrateByParts.ONCE, outflow=False): L = inner(outer(test, ubar), grad(q))*dx if ibp != IntegrateByParts.NEVER: - n = FacetNormal(state.mesh) + n = FacetNormal(domain.mesh) un = 0.5*(dot(ubar, n) + abs(dot(ubar, n))) L += dot(jump(test), (un('+')*q('+') - un('-')*q('-')))*dS_ @@ -116,7 +119,7 @@ def advection_form(state, test, q, ibp=IntegrateByParts.ONCE, outflow=False): + inner(test('-'), dot(ubar('-'), n('-'))*q('-')))*dS_ if outflow: - n = FacetNormal(state.mesh) + n = FacetNormal(domain.mesh) un = 0.5*(dot(ubar, n) + abs(dot(ubar, n))) L += test*un*q*(ds_v + ds_t + ds_b) @@ -125,7 +128,7 @@ def advection_form(state, test, q, ibp=IntegrateByParts.ONCE, outflow=False): return ibp_label(transport(form, TransportEquationType.advective), ibp) -def continuity_form(state, test, q, ibp=IntegrateByParts.ONCE, outflow=False): +def continuity_form(domain, test, q, ibp=IntegrateByParts.ONCE, outflow=False): u""" The form corresponding to the continuity transport operator. @@ -134,7 +137,8 @@ def continuity_form(state, test, q, ibp=IntegrateByParts.ONCE, outflow=False): form is integrated by parts. Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. test (:class:`TestFunction`): the test function. q (:class:`ufl.Expr`): the variable to be transported. ibp (:class:`IntegrateByParts`, optional): an enumerator representing @@ -153,7 +157,7 @@ def continuity_form(state, test, q, ibp=IntegrateByParts.ONCE, outflow=False): if outflow and ibp == IntegrateByParts.NEVER: raise ValueError("outflow is True and ibp is None are incompatible options") - Vu = state.spaces("HDiv") + Vu = domain.spaces("HDiv") dS_ = (dS_v + dS_h) if Vu.extruded else dS ubar = Function(Vu) @@ -163,7 +167,7 @@ def continuity_form(state, test, q, ibp=IntegrateByParts.ONCE, outflow=False): L = inner(test, div(outer(q, ubar)))*dx if ibp != IntegrateByParts.NEVER: - n = FacetNormal(state.mesh) + n = FacetNormal(domain.mesh) un = 0.5*(dot(ubar, n) + abs(dot(ubar, n))) L += dot(jump(test), (un('+')*q('+') - un('-')*q('-')))*dS_ @@ -173,7 +177,7 @@ def continuity_form(state, test, q, ibp=IntegrateByParts.ONCE, outflow=False): + inner(test('-'), dot(ubar('-'), n('-'))*q('-')))*dS_ if outflow: - n = FacetNormal(state.mesh) + n = FacetNormal(domain.mesh) un = 0.5*(dot(ubar, n) + abs(dot(ubar, n))) L += test*un*q*(ds_v + ds_t + ds_b) @@ -182,7 +186,7 @@ def continuity_form(state, test, q, ibp=IntegrateByParts.ONCE, outflow=False): return ibp_label(transport(form, TransportEquationType.conservative), ibp) -def vector_manifold_advection_form(state, test, q, ibp=IntegrateByParts.ONCE, outflow=False): +def vector_manifold_advection_form(domain, test, q, ibp=IntegrateByParts.ONCE, outflow=False): """ Form for advective transport operator including vector manifold correction. @@ -192,7 +196,8 @@ def vector_manifold_advection_form(state, test, q, ibp=IntegrateByParts.ONCE, ou is based on that of Bernard, Remacle et al (2009). Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. test (:class:`TestFunction`): the test function. q (:class:`ufl.Expr`): the variable to be transported. ibp (:class:`IntegrateByParts`, optional): an enumerator representing @@ -205,13 +210,13 @@ def vector_manifold_advection_form(state, test, q, ibp=IntegrateByParts.ONCE, ou class:`LabelledForm`: a labelled transport form. """ - L = advection_form(state, test, q, ibp, outflow) + L = advection_form(domain, test, q, ibp, outflow) # TODO: there should maybe be a restriction on IBP here - Vu = state.spaces("HDiv") + Vu = domain.spaces("HDiv") dS_ = (dS_v + dS_h) if Vu.extruded else dS ubar = Function(Vu) - n = FacetNormal(state.mesh) + n = FacetNormal(domain.mesh) un = 0.5*(dot(ubar, n) + abs(dot(ubar, n))) L += un('+')*inner(test('-'), n('+')+n('-'))*inner(q('+'), n('+'))*dS_ L += un('-')*inner(test('+'), n('+')+n('-'))*inner(q('-'), n('-'))*dS_ @@ -219,7 +224,7 @@ def vector_manifold_advection_form(state, test, q, ibp=IntegrateByParts.ONCE, ou return L -def vector_manifold_continuity_form(state, test, q, ibp=IntegrateByParts.ONCE, outflow=False): +def vector_manifold_continuity_form(domain, test, q, ibp=IntegrateByParts.ONCE, outflow=False): """ Form for continuity transport operator including vector manifold correction. @@ -229,7 +234,8 @@ def vector_manifold_continuity_form(state, test, q, ibp=IntegrateByParts.ONCE, o is based on that of Bernard, Remacle et al (2009). Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. test (:class:`TestFunction`): the test function. q (:class:`ufl.Expr`): the variable to be transported. ibp (:class:`IntegrateByParts`, optional): an enumerator representing @@ -242,12 +248,12 @@ def vector_manifold_continuity_form(state, test, q, ibp=IntegrateByParts.ONCE, o class:`LabelledForm`: a labelled transport form. """ - L = continuity_form(state, test, q, ibp, outflow) + L = continuity_form(domain, test, q, ibp, outflow) - Vu = state.spaces("HDiv") + Vu = domain.spaces("HDiv") dS_ = (dS_v + dS_h) if Vu.extruded else dS ubar = Function(Vu) - n = FacetNormal(state.mesh) + n = FacetNormal(domain.mesh) un = 0.5*(dot(ubar, n) + abs(dot(ubar, n))) L += un('+')*inner(test('-'), n('+')+n('-'))*inner(q('+'), n('+'))*dS_ L += un('-')*inner(test('+'), n('+')+n('-'))*inner(q('-'), n('-'))*dS_ @@ -257,7 +263,7 @@ def vector_manifold_continuity_form(state, test, q, ibp=IntegrateByParts.ONCE, o return transport(form) -def vector_invariant_form(state, test, q, ibp=IntegrateByParts.ONCE): +def vector_invariant_form(domain, test, q, ibp=IntegrateByParts.ONCE): u""" The form corresponding to the vector invariant transport operator. @@ -273,7 +279,8 @@ def vector_invariant_form(state, test, q, ibp=IntegrateByParts.ONCE): when integrating by parts. Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. test (:class:`TestFunction`): the test function. q (:class:`ufl.Expr`): the variable to be transported. ibp (:class:`IntegrateByParts`, optional): an enumerator representing @@ -287,13 +294,13 @@ def vector_invariant_form(state, test, q, ibp=IntegrateByParts.ONCE): class:`LabelledForm`: a labelled transport form. """ - Vu = state.spaces("HDiv") + Vu = domain.spaces("HDiv") dS_ = (dS_v + dS_h) if Vu.extruded else dS ubar = Function(Vu) - n = FacetNormal(state.mesh) + n = FacetNormal(domain.mesh) Upwind = 0.5*(sign(dot(ubar, n))+1) - if state.mesh.topological_dimension() == 3: + if domain.mesh.topological_dimension() == 3: if ibp != IntegrateByParts.ONCE: raise NotImplementedError @@ -313,9 +320,9 @@ def vector_invariant_form(state, test, q, ibp=IntegrateByParts.ONCE): else: - perp = state.perp - if state.on_sphere: - outward_normals = CellNormal(state.mesh) + perp = domain.perp + if domain.on_sphere: + outward_normals = CellNormal(domain.mesh) perp_u_upwind = lambda q: Upwind('+')*cross(outward_normals('+'), q('+')) + Upwind('-')*cross(outward_normals('-'), q('-')) else: perp_u_upwind = lambda q: Upwind('+')*perp(q('+')) + Upwind('-')*perp(q('-')) @@ -342,7 +349,7 @@ def vector_invariant_form(state, test, q, ibp=IntegrateByParts.ONCE): return transport(form, TransportEquationType.vector_invariant) -def kinetic_energy_form(state, test, q): +def kinetic_energy_form(domain, test, q): u""" The form corresponding to the kinetic energy term. @@ -351,7 +358,8 @@ def kinetic_energy_form(state, test, q): (1/2)∇(u.q). Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. test (:class:`TestFunction`): the test function. q (:class:`ufl.Expr`): the variable to be transported. @@ -359,7 +367,7 @@ def kinetic_energy_form(state, test, q): class:`LabelledForm`: a labelled transport form. """ - ubar = Function(state.spaces("HDiv")) + ubar = Function(domain.spaces("HDiv")) L = 0.5*div(test)*inner(q, ubar)*dx form = transporting_velocity(L, ubar) @@ -367,7 +375,7 @@ def kinetic_energy_form(state, test, q): return transport(form, TransportEquationType.vector_invariant) -def advection_equation_circulation_form(state, test, q, +def advection_equation_circulation_form(domain, test, q, ibp=IntegrateByParts.ONCE): u""" The circulation term in the transport of a vector-valued field. @@ -384,7 +392,8 @@ def advection_equation_circulation_form(state, test, q, term. An an upwind discretisation is used when integrating by parts. Args: - state (:class:`State`): the model's state object. + domain (:class:`Domain`): the model's domain object, containing the + mesh and the compatible function spaces. test (:class:`TestFunction`): the test function. q (:class:`ufl.Expr`): the variable to be transported. ibp (:class:`IntegrateByParts`, optional): an enumerator representing @@ -399,8 +408,8 @@ def advection_equation_circulation_form(state, test, q, """ form = ( - vector_invariant_form(state, test, q, ibp=ibp) - - kinetic_energy_form(state, test, q) + vector_invariant_form(domain, test, q, ibp=ibp) + - kinetic_energy_form(domain, test, q) ) return form diff --git a/integration-tests/balance/test_compressible_balance.py b/integration-tests/balance/test_compressible_balance.py index bfd7829e4..f8efda306 100644 --- a/integration-tests/balance/test_compressible_balance.py +++ b/integration-tests/balance/test_compressible_balance.py @@ -12,54 +12,60 @@ def setup_balance(dirname): - # set up grid and time stepping parameters + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Parameters dt = 1. tmax = 5. deltax = 400 L = 2000. H = 10000. - nlayers = int(H/deltax) ncolumns = int(L/deltax) + # Domain m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) + domain = Domain(mesh, dt, "CG", 1) - output = OutputParameters(dirname=dirname+'/dry_balance', dumpfreq=10, dumplist=['u']) + # Equation parameters = CompressibleParameters() + eqns = CompressibleEulerEquations(domain, parameters) - state = State(mesh, - dt=dt, - output=output, - parameters=parameters) + # I/O + output = OutputParameters(dirname=dirname+'/dry_balance', dumpfreq=10, dumplist=['u']) + io = IO(domain, output) - eqns = CompressibleEulerEquations(state, "CG", 1) + # Set up transport schemes + transported_fields = [ImplicitMidpoint(domain, "u"), + SSPRK3(domain, "rho"), + SSPRK3(domain, "theta", options=EmbeddedDGOptions())] + + # Set up linear solver + linear_solver = CompressibleSolver(eqns) + + # build time stepper + stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, + linear_solver=linear_solver) + # ------------------------------------------------------------------------ # # Initial conditions - rho0 = state.fields("rho") - theta0 = state.fields("theta") + # ------------------------------------------------------------------------ # + + rho0 = stepper.fields("rho") + theta0 = stepper.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())] + compressible_hydrostatic_balance(eqns, theta0, rho0, solve_for_rho=True) - # Set up linear solver - linear_solver = CompressibleSolver(state, eqns) - - # build time stepper - stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields, - linear_solver=linear_solver) + stepper.set_reference_profiles([('rho', rho0), + ('theta', theta0)]) return stepper, tmax diff --git a/integration-tests/balance/test_saturated_balance.py b/integration-tests/balance/test_saturated_balance.py index e5c044c35..41dce0f18 100644 --- a/integration-tests/balance/test_saturated_balance.py +++ b/integration-tests/balance/test_saturated_balance.py @@ -14,68 +14,43 @@ def setup_saturated(dirname, recovered): - # set up grid and time stepping parameters + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Parameters dt = 1. tmax = 3. deltax = 400. L = 2000. H = 10000. - nlayers = int(H/deltax) ncolumns = int(L/deltax) + degree = 0 if recovered else 1 + # Domain m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) + domain = Domain(mesh, dt, "CG", degree) - # 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) - + # Equation tracers = [WaterVapour(), CloudWater()] - if recovered: u_transport_option = "vector_advection_form" else: u_transport_option = "vector_invariant_form" + parameters = CompressibleParameters() eqns = CompressibleEulerEquations( - state, "CG", degree, u_transport_option=u_transport_option, active_tracers=tracers) - - # Initial conditions - rho0 = state.fields("rho") - theta0 = state.fields("theta") - water_v0 = state.fields("water_vapour") - water_c0 = state.fields("cloud_water") - moisture = ['water_vapour', 'cloud_water'] - - # spaces - Vt = theta0.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) + domain, parameters, u_transport_option=u_transport_option, active_tracers=tracers) - state.set_reference_profiles([('rho', rho0), - ('theta', theta0)]) + # I/O + output = OutputParameters(dirname=dirname+'/saturated_balance', dumpfreq=1, dumplist=['u']) + diagnostic_fields = [Theta_e(eqns)] + io = IO(domain, output, diagnostic_fields=diagnostic_fields) # Set up transport schemes if recovered: - VDG1 = state.spaces("DG1_equispaced") + VDG1 = domain.spaces("DG1_equispaced") VCG1 = FunctionSpace(mesh, "CG", 1) Vu_DG1 = VectorFunctionSpace(mesh, VDG1.ufl_element()) Vu_CG1 = VectorFunctionSpace(mesh, "CG", 1) @@ -99,26 +74,51 @@ def setup_saturated(dirname, recovered): wv_opts = EmbeddedDGOptions() wc_opts = EmbeddedDGOptions() - transported_fields = [SSPRK3(state, 'rho', options=rho_opts), - SSPRK3(state, 'theta', options=theta_opts), - SSPRK3(state, 'water_vapour', options=wv_opts), - SSPRK3(state, 'cloud_water', options=wc_opts)] + transported_fields = [SSPRK3(domain, 'rho', options=rho_opts), + SSPRK3(domain, 'theta', options=theta_opts), + SSPRK3(domain, 'water_vapour', options=wv_opts), + SSPRK3(domain, 'cloud_water', options=wc_opts)] if recovered: - transported_fields.append(SSPRK3(state, 'u', options=u_opts)) + transported_fields.append(SSPRK3(domain, 'u', options=u_opts)) else: - transported_fields.append(ImplicitMidpoint(state, 'u')) + transported_fields.append(ImplicitMidpoint(domain, 'u')) - linear_solver = CompressibleSolver(state, eqns, moisture=moisture) + # Linear solver + linear_solver = CompressibleSolver(eqns) - # add physics - physics_schemes = [(SaturationAdjustment(eqns, parameters), ForwardEuler(state))] + # Physics schemes + physics_schemes = [(SaturationAdjustment(eqns), ForwardEuler(domain))] - # build time stepper - stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields, + # Time stepper + stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, linear_solver=linear_solver, physics_schemes=physics_schemes) + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + rho0 = stepper.fields("rho") + theta0 = stepper.fields("theta") + water_v0 = stepper.fields("water_vapour") + water_c0 = stepper.fields("cloud_water") + + # spaces + Vt = theta0.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(eqns, stepper.fields, theta_e, water_t) + water_c0.assign(water_t - water_v0) + + stepper.set_reference_profiles([('rho', rho0), ('theta', theta0)]) + return stepper, tmax diff --git a/integration-tests/balance/test_unsaturated_balance.py b/integration-tests/balance/test_unsaturated_balance.py index d6d2dd5c5..e2377ceed 100644 --- a/integration-tests/balance/test_unsaturated_balance.py +++ b/integration-tests/balance/test_unsaturated_balance.py @@ -14,7 +14,11 @@ def setup_unsaturated(dirname, recovered): - # set up grid and time stepping parameters + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Parameters dt = 1. tmax = 3. deltax = 400 @@ -24,53 +28,32 @@ def setup_unsaturated(dirname, recovered): nlayers = int(H/deltax) ncolumns = int(L/deltax) - m = PeriodicIntervalMesh(ncolumns, L) - mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) - degree = 0 if recovered else 1 - output = OutputParameters(dirname=dirname+'/unsaturated_balance', dumpfreq=1) - parameters = CompressibleParameters() - diagnostic_fields = [Theta_d(), RelativeHumidity()] - - state = State(mesh, - dt=dt, - output=output, - parameters=parameters, - diagnostic_fields=diagnostic_fields) + # Domain + m = PeriodicIntervalMesh(ncolumns, L) + mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) + domain = Domain(mesh, dt, "CG", degree) + # Equation tracers = [WaterVapour(), CloudWater()] if recovered: u_transport_option = "vector_advection_form" else: u_transport_option = "vector_invariant_form" + parameters = CompressibleParameters() eqns = CompressibleEulerEquations( - state, "CG", degree, u_transport_option=u_transport_option, active_tracers=tracers) + domain, parameters, u_transport_option=u_transport_option, active_tracers=tracers) - # Initial conditions - rho0 = state.fields("rho") - theta0 = state.fields("theta") - moisture = ['water_vapour', 'cloud_water'] - - # spaces - Vt = theta0.function_space() - - # Isentropic background state - Tsurf = Constant(300.) - humidity = Constant(0.5) - theta_d = Function(Vt).interpolate(Tsurf) - RH = Function(Vt).interpolate(humidity) - - # Calculate hydrostatic exner - unsaturated_hydrostatic_balance(state, theta_d, RH) - - state.set_reference_profiles([('rho', rho0), - ('theta', theta0)]) + # I/O + output = OutputParameters(dirname=dirname+'/unsaturated_balance', dumpfreq=1) + diagnostic_fields = [Theta_d(eqns), RelativeHumidity(eqns)] + io = IO(domain, output, diagnostic_fields=diagnostic_fields) - # Set up transport schemes + # Transport schemes if recovered: - VDG1 = state.spaces("DG1_equispaced") + VDG1 = domain.spaces("DG1_equispaced") VCG1 = FunctionSpace(mesh, "CG", 1) Vu_DG1 = VectorFunctionSpace(mesh, VDG1.ufl_element()) Vu_CG1 = VectorFunctionSpace(mesh, "CG", 1) @@ -87,25 +70,47 @@ def setup_unsaturated(dirname, recovered): rho_opts = None theta_opts = EmbeddedDGOptions() - transported_fields = [SSPRK3(state, "rho", options=rho_opts), - SSPRK3(state, "theta", options=theta_opts), - SSPRK3(state, "water_vapour", options=theta_opts), - SSPRK3(state, "cloud_water", options=theta_opts)] + transported_fields = [SSPRK3(domain, "rho", options=rho_opts), + SSPRK3(domain, "theta", options=theta_opts), + SSPRK3(domain, "water_vapour", options=theta_opts), + SSPRK3(domain, "cloud_water", options=theta_opts)] if recovered: - transported_fields.append(SSPRK3(state, "u", options=u_opts)) + transported_fields.append(SSPRK3(domain, "u", options=u_opts)) else: - transported_fields.append(ImplicitMidpoint(state, "u")) + transported_fields.append(ImplicitMidpoint(domain, "u")) - linear_solver = CompressibleSolver(state, eqns, moisture=moisture) + # Linear solver + linear_solver = CompressibleSolver(eqns) # Set up physics - physics_schemes = [(SaturationAdjustment(eqns, parameters), ForwardEuler(state))] + physics_schemes = [(SaturationAdjustment(eqns), ForwardEuler(domain))] - # build time stepper - stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields, + # Time stepper + stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, linear_solver=linear_solver, physics_schemes=physics_schemes) + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + rho0 = stepper.fields("rho") + theta0 = stepper.fields("theta") + + # spaces + Vt = theta0.function_space() + + # Isentropic background state + Tsurf = Constant(300.) + humidity = Constant(0.5) + theta_d = Function(Vt).interpolate(Tsurf) + RH = Function(Vt).interpolate(humidity) + + # Calculate hydrostatic exner + unsaturated_hydrostatic_balance(eqns, stepper.fields, theta_d, RH) + + stepper.set_reference_profiles([('rho', rho0), ('theta', theta0)]) + return stepper, tmax diff --git a/integration-tests/conftest.py b/integration-tests/conftest.py index 883c10909..92ae1deca 100644 --- a/integration-tests/conftest.py +++ b/integration-tests/conftest.py @@ -9,7 +9,7 @@ from collections import namedtuple import pytest -opts = ('state', 'tmax', 'f_init', 'f_end', 'family', 'degree', +opts = ('domain', 'tmax', 'io', 'f_init', 'f_end', 'degree', 'uexpr', 'umax', 'radius', 'tol') TracerSetup = namedtuple('TracerSetup', opts) TracerSetup.__new__.__defaults__ = (None,)*len(opts) @@ -29,7 +29,8 @@ def tracer_sphere(tmpdir, degree): dt = pi/3. * 0.02 output = OutputParameters(dirname=str(tmpdir), dumpfreq=15) - state = State(mesh, dt=dt, output=output) + domain = Domain(mesh, dt, family="BDM", degree=degree) + io = IO(domain, output) umax = 1.0 uexpr = as_vector([- umax * x[1] / radius, umax * x[0] / radius, 0.0]) @@ -40,7 +41,7 @@ def tracer_sphere(tmpdir, degree): tol = 0.05 - return TracerSetup(state, tmax, f_init, f_end, "BDM", degree, + return TracerSetup(domain, tmax, io, f_init, f_end, degree, uexpr, umax, radius, tol) @@ -56,7 +57,8 @@ def tracer_slice(tmpdir, degree): dt = 0.01 tmax = 0.75 output = OutputParameters(dirname=str(tmpdir), dumpfreq=25) - state = State(mesh, dt=dt, output=output) + domain = Domain(mesh, dt, family="CG", degree=degree) + io = IO(domain, output) uexpr = as_vector([2.0, 0.0]) @@ -73,7 +75,7 @@ def tracer_slice(tmpdir, degree): tol = 0.12 - return TracerSetup(state, tmax, f_init, f_end, "CG", degree, uexpr, tol=tol) + return TracerSetup(domain, tmax, io, f_init, f_end, degree, uexpr, tol=tol) def tracer_blob_slice(tmpdir, degree): @@ -83,13 +85,14 @@ def tracer_blob_slice(tmpdir, degree): mesh = ExtrudedMesh(m, layers=10, layer_height=1.) output = OutputParameters(dirname=str(tmpdir), dumpfreq=25) - state = State(mesh, dt=dt, output=output) + domain = Domain(mesh, dt, family="CG", degree=degree) + io = IO(domain, 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) + return TracerSetup(domain, tmax, io, f_init, degree=degree) @pytest.fixture() diff --git a/integration-tests/data/dry_compressible_chkpt.h5 b/integration-tests/data/dry_compressible_chkpt.h5 index a3fbe43e1..d03b5ce5d 100644 Binary files a/integration-tests/data/dry_compressible_chkpt.h5 and b/integration-tests/data/dry_compressible_chkpt.h5 differ diff --git a/integration-tests/data/incompressible_chkpt.h5 b/integration-tests/data/incompressible_chkpt.h5 index fa9425bac..3b3be3e7f 100644 Binary files a/integration-tests/data/incompressible_chkpt.h5 and b/integration-tests/data/incompressible_chkpt.h5 differ diff --git a/integration-tests/data/moist_compressible_chkpt.h5 b/integration-tests/data/moist_compressible_chkpt.h5 index 22ed1b5bb..0f62ecd6b 100644 Binary files a/integration-tests/data/moist_compressible_chkpt.h5 and b/integration-tests/data/moist_compressible_chkpt.h5 differ diff --git a/integration-tests/diffusion/test_diffusion.py b/integration-tests/diffusion/test_diffusion.py index 5be082e25..9cc43b38f 100644 --- a/integration-tests/diffusion/test_diffusion.py +++ b/integration-tests/diffusion/test_diffusion.py @@ -8,18 +8,16 @@ import pytest -def run(equation, diffusion_scheme, state, tmax): - - timestepper = Timestepper(equation, diffusion_scheme, state) +def run(timestepper, tmax): timestepper.run(0., tmax) - return timestepper.state.fields("f") + return timestepper.fields("f") @pytest.mark.parametrize("DG", [True, False]) def test_scalar_diffusion(tmpdir, DG, tracer_setup): setup = tracer_setup(tmpdir, geometry="slice", blob=True) - state = setup.state + domain = setup.domain f_init = setup.f_init tmax = setup.tmax tol = 5.e-2 @@ -28,20 +26,20 @@ def test_scalar_diffusion(tmpdir, DG, tracer_setup): f_end_expr = (1/(1+4*tmax))*f_init**(1/(1+4*tmax)) if DG: - V = state.spaces("DG", "DG", 1) + V = domain.spaces("DG", "DG", degree=1) else: - V = state.spaces("theta", degree=1) + V = domain.spaces("theta", degree=1) mu = 5. diffusion_params = DiffusionParameters(kappa=kappa, mu=mu) - eqn = DiffusionEquation(state, V, "f", - diffusion_parameters=diffusion_params) - - diffusion_scheme = BackwardEuler(state) + eqn = DiffusionEquation(domain, V, "f", diffusion_parameters=diffusion_params) + diffusion_scheme = BackwardEuler(domain) + timestepper = Timestepper(eqn, diffusion_scheme, setup.io) - state.fields("f").interpolate(f_init) - f_end = run(eqn, diffusion_scheme, state, tmax) + # Initial conditions + timestepper.fields("f").interpolate(f_init) + f_end = run(timestepper, tmax) assert errornorm(f_end_expr, f_end) < tol @@ -49,7 +47,7 @@ def test_scalar_diffusion(tmpdir, DG, tracer_setup): def test_vector_diffusion(tmpdir, DG, tracer_setup): setup = tracer_setup(tmpdir, geometry="slice", blob=True) - state = setup.state + domain = setup.domain f_init = setup.f_init tmax = setup.tmax tol = 3.e-2 @@ -59,24 +57,24 @@ def test_vector_diffusion(tmpdir, DG, tracer_setup): kappa = Constant([[kappa, 0.], [0., kappa]]) if DG: - V = VectorFunctionSpace(state.mesh, "DG", 1) + V = VectorFunctionSpace(domain.mesh, "DG", 1) else: - V = state.spaces("HDiv", "CG", 1) + V = domain.spaces("HDiv", "CG", 1) f_init = as_vector([f_init, 0.]) f_end_expr = as_vector([f_end_expr, 0.]) mu = 5. diffusion_params = DiffusionParameters(kappa=kappa, mu=mu) - eqn = DiffusionEquation(state, V, "f", - diffusion_parameters=diffusion_params) + eqn = DiffusionEquation(domain, V, "f", diffusion_parameters=diffusion_params) + diffusion_scheme = BackwardEuler(domain) + timestepper = Timestepper(eqn, diffusion_scheme, setup.io) + # Initial conditions if DG: - state.fields("f").interpolate(f_init) + timestepper.fields("f").interpolate(f_init) else: - state.fields("f").project(f_init) - - diffusion_scheme = BackwardEuler(state) + timestepper.fields("f").project(f_init) - f_end = run(eqn, diffusion_scheme, state, tmax) + f_end = run(timestepper, tmax) assert errornorm(f_end_expr, f_end) < tol diff --git a/integration-tests/equations/test_advection_diffusion.py b/integration-tests/equations/test_advection_diffusion.py index ad4201ca2..20e7076b5 100644 --- a/integration-tests/equations/test_advection_diffusion.py +++ b/integration-tests/equations/test_advection_diffusion.py @@ -10,22 +10,36 @@ def run_advection_diffusion(tmpdir): - # Mesh, state and equation - L = 10 - mesh = PeriodicIntervalMesh(20, L) + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain dt = 0.02 tmax = 1.0 + L = 10 + mesh = PeriodicIntervalMesh(20, L) + domain = Domain(mesh, dt, "CG", 1) + # Equation diffusion_params = DiffusionParameters(kappa=0.75, mu=5) - output = OutputParameters(dirname=str(tmpdir), dumpfreq=25) - state = State(mesh, dt=dt, output=output) - V = state.spaces("DG", "DG", 1) + V = domain.spaces("DG") Vu = VectorFunctionSpace(mesh, "CG", 1) - equation = AdvectionDiffusionEquation(state, V, "f", Vu=Vu, + equation = AdvectionDiffusionEquation(domain, V, "f", Vu=Vu, diffusion_parameters=diffusion_params) + # I/O + output = OutputParameters(dirname=str(tmpdir), dumpfreq=25) + io = IO(domain, output) + + # Time stepper + stepper = PrescribedTransport(equation, SSPRK3(domain), io) + + # ------------------------------------------------------------------------ # # Initial conditions + # ------------------------------------------------------------------------ # + x = SpatialCoordinate(mesh) xc_init = 0.25*L xc_end = 0.75*L @@ -45,15 +59,18 @@ def run_advection_diffusion(tmpdir): f_init_expr = f_init*exp(-(x_init / f_width_init)**2) f_end_expr = f_end*exp(-(x_end / f_width_end)**2) - state.fields('f').interpolate(f_init_expr) - state.fields('u').interpolate(as_vector([Constant(umax)])) - f_end = state.fields('f_end', V).interpolate(f_end_expr) + stepper.fields('f').interpolate(f_init_expr) + stepper.fields('u').interpolate(as_vector([Constant(umax)])) + f_end = stepper.fields('f_end', space=V) + f_end.interpolate(f_end_expr) - # Time stepper - timestepper = PrescribedTransport(equation, SSPRK3(state), state) - timestepper.run(0, tmax=tmax) + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # + + stepper.run(0, tmax=tmax) - error = norm(state.fields('f') - f_end) / norm(f_end) + error = norm(stepper.fields('f') - f_end) / norm(f_end) return error diff --git a/integration-tests/equations/test_dry_compressible.py b/integration-tests/equations/test_dry_compressible.py index 272e7d796..edb06e6f9 100644 --- a/integration-tests/equations/test_dry_compressible.py +++ b/integration-tests/equations/test_dry_compressible.py @@ -7,11 +7,16 @@ from gusto import * from gusto import thermodynamics as tde from firedrake import (SpatialCoordinate, PeriodicIntervalMesh, exp, - sqrt, ExtrudedMesh, norm) + sqrt, ExtrudedMesh, norm, as_vector) def run_dry_compressible(tmpdir): + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain dt = 6.0 tmax = 2*dt nlayers = 10 # horizontal layers @@ -20,23 +25,39 @@ def run_dry_compressible(tmpdir): Lz = 1000.0 m = PeriodicIntervalMesh(ncols, Lx) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=Lz/nlayers) + domain = Domain(mesh, dt, "CG", 1) + + # Equation + parameters = CompressibleParameters() + eqn = CompressibleEulerEquations(domain, parameters) + # I/O output = OutputParameters(dirname=tmpdir+"/dry_compressible", dumpfreq=2, chkptfreq=2) - parameters = CompressibleParameters() - R_d = parameters.R_d - g = parameters.g + io = IO(domain, output) + + # Transport schemes + transported_fields = [ImplicitMidpoint(domain, "u"), + SSPRK3(domain, "rho"), + SSPRK3(domain, "theta")] - state = State(mesh, - dt=dt, - output=output, - parameters=parameters) + # Linear solver + linear_solver = CompressibleSolver(eqn) - eqn = CompressibleEulerEquations(state, "CG", 1) + # Time stepper + stepper = SemiImplicitQuasiNewton(eqn, io, transported_fields, + linear_solver=linear_solver) + # ------------------------------------------------------------------------ # # Initial conditions - rho0 = state.fields("rho") - theta0 = state.fields("theta") + # ------------------------------------------------------------------------ # + + R_d = parameters.R_d + g = parameters.g + + rho0 = stepper.fields("rho") + theta0 = stepper.fields("theta") + u0 = stepper.fields("u") # Approximate hydrostatic balance x, z = SpatialCoordinate(mesh) @@ -46,50 +67,44 @@ def run_dry_compressible(tmpdir): theta0.interpolate(tde.theta(parameters, T, p)) rho0.interpolate(p / (R_d * T)) - state.set_reference_profiles([('rho', rho0), - ('theta', theta0)]) + # Add horizontal translation to ensure some transport happens + u0.project(as_vector([0.5, 0.0])) + + stepper.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 = SemiImplicitQuasiNewton(eqn, state, transported_fields, - linear_solver=linear_solver) - + # ------------------------------------------------------------------------ # # Run + # ------------------------------------------------------------------------ # + stepper.run(t=0, tmax=tmax) - # State for checking checkpoints + # IO for checking checkpoints checkpoint_name = 'dry_compressible_chkpt' new_path = join(abspath(dirname(__file__)), '..', f'data/{checkpoint_name}') + check_eqn = CompressibleEulerEquations(domain, parameters) 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 = SemiImplicitQuasiNewton(check_eqn, check_state, []) + check_io = IO(domain, check_output) + check_stepper = SemiImplicitQuasiNewton(check_eqn, check_io, []) + check_stepper.set_reference_profiles([]) check_stepper.run(t=0, tmax=0, pickup=True) - return state, check_state + return stepper, check_stepper def test_dry_compressible(tmpdir): dirname = str(tmpdir) - state, check_state = run_dry_compressible(dirname) + stepper, check_stepper = run_dry_compressible(dirname) for variable in ['u', 'rho', 'theta']: - new_variable = state.fields(variable) - check_variable = check_state.fields(variable) + new_variable = stepper.fields(variable) + check_variable = check_stepper.fields(variable) error = norm(new_variable - check_variable) / norm(check_variable) # Slack values chosen to be robust to different platforms diff --git a/integration-tests/equations/test_forced_advection.py b/integration-tests/equations/test_forced_advection.py index 9bada6329..e9e4574b9 100644 --- a/integration-tests/equations/test_forced_advection.py +++ b/integration-tests/equations/test_forced_advection.py @@ -8,14 +8,18 @@ """ from gusto import * -from firedrake import (PeriodicIntervalMesh, SpatialCoordinate, FunctionSpace, - VectorFunctionSpace, conditional, acos, cos, pi, +from firedrake import (PeriodicIntervalMesh, SpatialCoordinate, + VectorFunctionSpace, conditional, acos, cos, pi, sin, as_vector, errornorm) def run_forced_advection(tmpdir): - # mesh, state and equation + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain Lx = 100 delta_x = 2.0 nx = int(Lx/delta_x) @@ -23,20 +27,12 @@ def run_forced_advection(tmpdir): x = SpatialCoordinate(mesh)[0] dt = 0.2 - output = OutputParameters(dirname=str(tmpdir), dumpfreq=1) - diagnostic_fields = [CourantNumber()] + domain = Domain(mesh, dt, "CG", 1) - state = State(mesh, - dt=dt, - output=output, - parameters=None, - diagnostics=None, - diagnostic_fields=diagnostic_fields) - - VD = FunctionSpace(mesh, "DG", 1) + VD = domain.spaces("DG") Vu = VectorFunctionSpace(mesh, "CG", 1) - # set up parameters and initial conditions + # Equation u_max = 1 C0 = 0.6 K0 = 0.3 @@ -49,22 +45,34 @@ def run_forced_advection(tmpdir): msat = Function(VD) msat.interpolate(msat_expr) - # initial moisture profile - mexpr = C0 + K0*cos((2*pi*x)/Lx) - rain = Rain(space='tracer', transport_eqn=TransportEquationType.no_transport) - meqn = ForcedAdvectionEquation(state, VD, field_name="water_vapour", Vu=Vu, + meqn = ForcedAdvectionEquation(domain, VD, field_name="water_vapour", Vu=Vu, active_tracers=[rain]) physics_schemes = [(InstantRain(meqn, msat, rain_name="rain", - set_tau_to_dt=True), ForwardEuler(state))] + set_tau_to_dt=True, parameters=None), ForwardEuler(domain))] - state.fields("u").project(as_vector([u_max])) - qv = state.fields("water_vapour") - qv.project(mexpr) + # I/O + output = OutputParameters(dirname=str(tmpdir), dumpfreq=1) + diagnostic_fields = [CourantNumber()] + io = IO(domain, output, diagnostic_fields=diagnostic_fields) + + # Time Stepper + stepper = PrescribedTransport(meqn, RK4(domain), io, + physics_schemes=physics_schemes) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + # initial moisture profile + mexpr = C0 + K0*cos((2*pi*x)/Lx) + + stepper.fields("u").project(as_vector([u_max])) + stepper.fields("water_vapour").project(mexpr) # exact rainfall profile (analytically) - r_exact = state.fields("r_exact", VD) + r_exact = stepper.fields("r_exact", space=VD) lim1 = Lx/(2*pi) * acos((C0 + K0 - Csat)/Ksat) lim2 = Lx/2 coord = (Ksat*cos(2*pi*x/Lx) + Csat - C0)/K0 @@ -72,13 +80,13 @@ def run_forced_advection(tmpdir): r_expr = conditional(x < lim2, conditional(x > lim1, exact_expr, 0), 0) r_exact.interpolate(r_expr) - # build time stepper - stepper = PrescribedTransport(meqn, RK4(state), state, - physics_schemes=physics_schemes) + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # stepper.run(0, tmax=tmax) - error = errornorm(r_exact, state.fields("rain")) + error = errornorm(r_exact, stepper.fields("rain")) return error diff --git a/integration-tests/equations/test_incompressible.py b/integration-tests/equations/test_incompressible.py index 0ae835a72..9bdf9f358 100644 --- a/integration-tests/equations/test_incompressible.py +++ b/integration-tests/equations/test_incompressible.py @@ -6,11 +6,16 @@ from os.path import join, abspath, dirname from gusto import * from firedrake import (SpatialCoordinate, PeriodicIntervalMesh, exp, - sqrt, ExtrudedMesh, Function, norm) + sqrt, ExtrudedMesh, Function, norm, as_vector) def run_incompressible(tmpdir): + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain dt = 6.0 tmax = 2*dt nlayers = 10 # horizontal layers @@ -19,21 +24,39 @@ def run_incompressible(tmpdir): Lz = 1000.0 m = PeriodicIntervalMesh(ncols, Lx) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=Lz/nlayers) + domain = Domain(mesh, dt, "CG", 1) + + # Equation + parameters = CompressibleParameters() + eqn = IncompressibleBoussinesqEquations(domain, parameters) + # I/O output = OutputParameters(dirname=tmpdir+"/incompressible", dumpfreq=2, chkptfreq=2) - parameters = CompressibleParameters() + io = IO(domain, output) - state = State(mesh, - dt=dt, - output=output, - parameters=parameters) + # Transport Schemes + b_opts = SUPGOptions() + transported_fields = [ImplicitMidpoint(domain, "u"), + SSPRK3(domain, "b", options=b_opts)] + + # Linear solver + linear_solver = IncompressibleSolver(eqn) - eqns = IncompressibleBoussinesqEquations(state, "CG", 1) + # Time stepper + stepper = SemiImplicitQuasiNewton(eqn, io, transported_fields, + linear_solver=linear_solver) + # ------------------------------------------------------------------------ # # Initial conditions - p0 = state.fields("p") - b0 = state.fields("b") + # ------------------------------------------------------------------------ # + + p0 = stepper.fields("p") + b0 = stepper.fields("b") + u0 = stepper.fields("u") + + # Add horizontal translation to ensure some transport happens + u0.project(as_vector([0.5, 0.0])) # z.grad(bref) = N**2 x, z = SpatialCoordinate(mesh) @@ -41,51 +64,42 @@ def run_incompressible(tmpdir): bref = z*(N**2) b_b = Function(b0.function_space()).interpolate(bref) - incompressible_hydrostatic_balance(state, b_b, p0) - state.initialise([('p', p0), - ('b', b_b)]) + incompressible_hydrostatic_balance(eqn, b_b, p0) + stepper.set_reference_profiles([('p', p0), ('b', b_b)]) # Add perturbation r = sqrt((x-Lx/2)**2 + (z-Lz/2)**2) b_pert = 0.1*exp(-(r/(Lx/5)**2)) b0.interpolate(b_b + b_pert) - # Set up transport schemes - b_opts = SUPGOptions() - transported_fields = [ImplicitMidpoint(state, "u"), - SSPRK3(state, "b", options=b_opts)] - - # Set up linear solver for the timestepping scheme - linear_solver = IncompressibleSolver(state, eqns) - - # build time stepper - stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields, - linear_solver=linear_solver) - + # ------------------------------------------------------------------------ # # Run + # ------------------------------------------------------------------------ # + stepper.run(t=0, tmax=tmax) # State for checking checkpoints checkpoint_name = 'incompressible_chkpt' new_path = join(abspath(dirname(__file__)), '..', f'data/{checkpoint_name}') + check_eqn = IncompressibleBoussinesqEquations(domain, parameters) check_output = OutputParameters(dirname=tmpdir+"/incompressible", checkpoint_pickup_filename=new_path) - check_state = State(mesh, dt=dt, output=check_output) - check_eqn = IncompressibleBoussinesqEquations(check_state, "CG", 1) - check_stepper = SemiImplicitQuasiNewton(check_eqn, check_state, []) + check_io = IO(domain, check_output) + check_stepper = SemiImplicitQuasiNewton(check_eqn, check_io, []) + check_stepper.set_reference_profiles([]) check_stepper.run(t=0, tmax=0, pickup=True) - return state, check_state + return stepper, check_stepper def test_incompressible(tmpdir): dirname = str(tmpdir) - state, check_state = run_incompressible(dirname) + stepper, check_stepper = run_incompressible(dirname) for variable in ['u', 'b', 'p']: - new_variable = state.fields(variable) - check_variable = check_state.fields(variable) + new_variable = stepper.fields(variable) + check_variable = check_stepper.fields(variable) error = norm(new_variable - check_variable) / norm(check_variable) # Slack values chosen to be robust to different platforms diff --git a/integration-tests/equations/test_moist_compressible.py b/integration-tests/equations/test_moist_compressible.py index e7c636bc3..91d55d1d8 100644 --- a/integration-tests/equations/test_moist_compressible.py +++ b/integration-tests/equations/test_moist_compressible.py @@ -7,11 +7,16 @@ from gusto import * import gusto.thermodynamics as tde from firedrake import (SpatialCoordinate, PeriodicIntervalMesh, exp, - sqrt, ExtrudedMesh, norm) + sqrt, ExtrudedMesh, norm, as_vector) def run_moist_compressible(tmpdir): + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain dt = 6.0 tmax = 2*dt nlayers = 10 # horizontal layers @@ -20,27 +25,45 @@ def run_moist_compressible(tmpdir): Lz = 1000.0 m = PeriodicIntervalMesh(ncols, Lx) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=Lz/nlayers) + domain = Domain(mesh, dt, "CG", 1) + + # Equation + parameters = CompressibleParameters() + tracers = [WaterVapour(name='vapour_mixing_ratio'), CloudWater(name='cloud_liquid_mixing_ratio')] + eqn = CompressibleEulerEquations(domain, parameters, active_tracers=tracers) + # I/O output = OutputParameters(dirname=tmpdir+"/moist_compressible", dumpfreq=2, chkptfreq=2) - parameters = CompressibleParameters() - R_d = parameters.R_d - R_v = parameters.R_v - g = parameters.g + io = IO(domain, output) - tracers = [WaterVapour(name='vapour_mixing_ratio'), CloudWater(name='cloud_liquid_mixing_ratio')] + # Transport schemes + transported_fields = [ImplicitMidpoint(domain, "u"), + SSPRK3(domain, "rho"), + SSPRK3(domain, "theta")] - state = State(mesh, - dt=dt, - output=output, - parameters=parameters) + # Linear solver + linear_solver = CompressibleSolver(eqn) - eqn = CompressibleEulerEquations(state, "CG", 1, active_tracers=tracers) + # Time stepper + stepper = SemiImplicitQuasiNewton(eqn, io, transported_fields, + linear_solver=linear_solver) + # ------------------------------------------------------------------------ # # Initial conditions - rho0 = state.fields("rho") - theta0 = state.fields("theta") - m_v0 = state.fields("vapour_mixing_ratio") + # ------------------------------------------------------------------------ # + + R_d = parameters.R_d + R_v = parameters.R_v + g = parameters.g + + rho0 = stepper.fields("rho") + theta0 = stepper.fields("theta") + m_v0 = stepper.fields("vapour_mixing_ratio") + u0 = stepper.fields("u") + + # Add horizontal translation to ensure some transport happens + u0.project(as_vector([0.5, 0.0])) # Approximate hydrostatic balance x, z = SpatialCoordinate(mesh) @@ -52,50 +75,42 @@ def run_moist_compressible(tmpdir): theta0.interpolate(tde.theta(parameters, T_vd, p)) rho0.interpolate(p / (R_d * T)) - state.set_reference_profiles([('rho', rho0), - ('theta', theta0)]) + stepper.set_reference_profiles([('rho', rho0), ('theta', theta0), + ('vapour_mixing_ratio', m_v0)]) # 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, moisture=['vapour_mixing_ratio']) - - # build time stepper - stepper = SemiImplicitQuasiNewton(eqn, state, transported_fields, - linear_solver=linear_solver) - + # ------------------------------------------------------------------------ # # Run + # ------------------------------------------------------------------------ # + stepper.run(t=0, tmax=tmax) # State for checking checkpoints checkpoint_name = 'moist_compressible_chkpt' new_path = join(abspath(dirname(__file__)), '..', f'data/{checkpoint_name}') + check_eqn = CompressibleEulerEquations(domain, parameters, active_tracers=tracers) check_output = OutputParameters(dirname=tmpdir+"/moist_compressible", checkpoint_pickup_filename=new_path) - check_state = State(mesh, dt=dt, output=check_output, parameters=parameters) - check_eqn = CompressibleEulerEquations(check_state, "CG", 1, active_tracers=tracers) - check_stepper = SemiImplicitQuasiNewton(check_eqn, check_state, []) + check_io = IO(domain, output=check_output) + check_stepper = SemiImplicitQuasiNewton(check_eqn, check_io, []) + check_stepper.set_reference_profiles([]) check_stepper.run(t=0, tmax=0, pickup=True) - return state, check_state + return stepper, check_stepper def test_moist_compressible(tmpdir): dirname = str(tmpdir) - state, check_state = run_moist_compressible(dirname) + stepper, check_stepper = run_moist_compressible(dirname) for variable in ['u', 'rho', 'theta', 'vapour_mixing_ratio']: - new_variable = state.fields(variable) - check_variable = check_state.fields(variable) + new_variable = stepper.fields(variable) + check_variable = check_stepper.fields(variable) error = norm(new_variable - check_variable) / norm(check_variable) # Slack values chosen to be robust to different platforms diff --git a/integration-tests/equations/test_sw_linear_triangle.py b/integration-tests/equations/test_sw_linear_triangle.py index 57ffe22be..908e182b7 100644 --- a/integration-tests/equations/test_sw_linear_triangle.py +++ b/integration-tests/equations/test_sw_linear_triangle.py @@ -11,6 +11,13 @@ def setup_sw(dirname): + + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain + dt = 3600. refinements = 3 # number of horizontal cells = 20*(4^refinements) R = 6371220. @@ -22,25 +29,31 @@ def setup_sw(dirname): x = SpatialCoordinate(mesh) mesh.init_cell_orientations(x) - dt = 3600. - output = OutputParameters(dirname=dirname+"/sw_linear_w2", steady_state_error_fields=['u', 'D'], dumpfreq=12) - parameters = ShallowWaterParameters(H=H) - - state = State(mesh, - dt=dt, - output=output, - parameters=parameters) + domain = Domain(mesh, dt, "BDM", degree=1) - # Coriolis + # Equation + parameters = ShallowWaterParameters(H=H) Omega = parameters.Omega fexpr = 2*Omega*x[2]/R + eqns = LinearShallowWaterEquations(domain, parameters, fexpr=fexpr) + + # I/O + diagnostic_fields = [SteadyStateError('u'), SteadyStateError('D')] + output = OutputParameters(dirname=dirname+"/sw_linear_w2", dumpfreq=12) + io = IO(domain, output, diagnostic_fields=diagnostic_fields) + + # Transport schemes + transport_schemes = [ForwardEuler(domain, "D")] - eqns = LinearShallowWaterEquations(state, "BDM", 1, fexpr=fexpr) + # Time stepper + stepper = SemiImplicitQuasiNewton(eqns, io, transport_schemes) - # interpolate initial conditions - # Initial/current conditions - u0 = state.fields("u") - D0 = state.fields("D") + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + u0 = stepper.fields("u") + D0 = stepper.fields("D") u_max = 2*pi*R/(12*day) # Maximum amplitude of the zonal wind (m/s) uexpr = as_vector([-u_max*x[1]/R, u_max*x[0]/R, 0.0]) g = parameters.g @@ -48,10 +61,8 @@ def setup_sw(dirname): u0.project(uexpr) D0.interpolate(Dexpr) - transport_schemes = [ForwardEuler(state, "D")] - - # build time stepper - stepper = SemiImplicitQuasiNewton(eqns, state, transport_schemes) + Dbar = Function(D0.function_space()).assign(H) + stepper.set_reference_profiles([('D', Dbar)]) return stepper, 2*day diff --git a/integration-tests/equations/test_sw_triangle.py b/integration-tests/equations/test_sw_triangle.py index 86a20256e..d2ce2303a 100644 --- a/integration-tests/equations/test_sw_triangle.py +++ b/integration-tests/equations/test_sw_triangle.py @@ -22,15 +22,27 @@ def setup_sw(dirname, dt, u_transport_option): + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # Domain refinements = 3 # number of horizontal cells = 20*(4^refinements) mesh = IcosahedralSphereMesh(radius=R, refinement_level=refinements) + domain = Domain(mesh, dt, family="BDM", degree=1) x = SpatialCoordinate(mesh) mesh.init_cell_orientations(x) - output = OutputParameters(dirname=dirname+"/sw", dumplist_latlon=['D', 'D_error'], steady_state_error_fields=['D', 'u']) + # Equation parameters = ShallowWaterParameters(H=H) + Omega = parameters.Omega + fexpr = 2*Omega*x[2]/R + eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, + u_transport_option=u_transport_option) + + # I/O diagnostic_fields = [RelativeVorticity(), AbsoluteVorticity(), PotentialVorticity(), ShallowWaterPotentialEnstrophy('RelativeVorticity'), @@ -48,42 +60,44 @@ def setup_sw(dirname, dt, u_transport_option): 'SWPotentialEnstrophy_from_AbsoluteVorticity'), MeridionalComponent('u'), ZonalComponent('u'), - RadialComponent('u')] + RadialComponent('u'), + SteadyStateError('D'), + SteadyStateError('u')] + output = OutputParameters(dirname=dirname+"/sw", dumplist_latlon=['D', 'D_error']) + io = IO(domain, output, diagnostic_fields=diagnostic_fields) - state = State(mesh, - dt=dt, - output=output, - parameters=parameters, - diagnostic_fields=diagnostic_fields) + return domain, eqns, io - Omega = parameters.Omega - fexpr = 2*Omega*x[2]/R - eqns = ShallowWaterEquations(state, family="BDM", degree=1, - fexpr=fexpr, - u_transport_option=u_transport_option) + +def set_up_initial_conditions(domain, equation, stepper): + + x = SpatialCoordinate(domain.mesh) # interpolate initial conditions - u0 = state.fields("u") - D0 = state.fields("D") + u0 = stepper.fields("u") + D0 = stepper.fields("D") uexpr = as_vector([-u_max*x[1]/R, u_max*x[0]/R, 0.0]) - g = parameters.g + g = equation.parameters.g + Omega = equation.parameters.Omega Dexpr = H - ((R * Omega * u_max + u_max*u_max/2.0)*(x[2]*x[2]/(R*R)))/g u0.project(uexpr) D0.interpolate(Dexpr) - vspace = FunctionSpace(state.mesh, "CG", 3) + Dbar = Function(D0.function_space()).assign(H) + stepper.set_reference_profiles([('D', Dbar)]) + + vspace = FunctionSpace(domain.mesh, "CG", 3) vexpr = (2*u_max/R)*x[2]/R - f = state.fields("coriolis") - vrel_analytical = state.fields("AnalyticalRelativeVorticity", vspace) + + f = stepper.fields("coriolis") + vrel_analytical = stepper.fields("AnalyticalRelativeVorticity", space=vspace) vrel_analytical.interpolate(vexpr) - vabs_analytical = state.fields("AnalyticalAbsoluteVorticity", vspace) + vabs_analytical = stepper.fields("AnalyticalAbsoluteVorticity", space=vspace) vabs_analytical.interpolate(vexpr + f) - pv_analytical = state.fields("AnalyticalPotentialVorticity", vspace) + pv_analytical = stepper.fields("AnalyticalPotentialVorticity", space=vspace) pv_analytical.interpolate((vexpr+f)/D0) - return state, eqns - def check_results(dirname): filename = path.join(dirname, "sw/diagnostics.nc") @@ -139,12 +153,20 @@ def test_sw_setup(tmpdir, u_transport_option): dirname = str(tmpdir) dt = 1500 - state, eqns = setup_sw(dirname, dt, u_transport_option) + domain, eqns, io = setup_sw(dirname, dt, u_transport_option) + # Transport schemes transported_fields = [] - transported_fields.append((ImplicitMidpoint(state, "u"))) - transported_fields.append((SSPRK3(state, "D"))) - stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields) + transported_fields.append((ImplicitMidpoint(domain, "u"))) + transported_fields.append((SSPRK3(domain, "D"))) + + # Time stepper + stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields) + + # Initial conditions + set_up_initial_conditions(domain, eqns, stepper) + + # Run stepper.run(t=0, tmax=0.25*day) check_results(dirname) @@ -157,9 +179,14 @@ def test_sw_ssprk3(tmpdir, u_transport_option): dirname = str(tmpdir) dt = 100 - state, eqns = setup_sw(dirname, dt, u_transport_option) + domain, eqns, io = setup_sw(dirname, dt, u_transport_option) + + stepper = Timestepper(eqns, SSPRK3(domain), io) + + # Initial conditions + set_up_initial_conditions(domain, eqns, stepper) - stepper = Timestepper(eqns, SSPRK3(state), state) + # Run stepper.run(t=0, tmax=0.01*day) check_results(dirname) diff --git a/integration-tests/model/test_checkpointing.py b/integration-tests/model/test_checkpointing.py index e4fc4ddb7..e478aeb4e 100644 --- a/integration-tests/model/test_checkpointing.py +++ b/integration-tests/model/test_checkpointing.py @@ -20,43 +20,42 @@ def setup_checkpointing(dirname): # build volume mesh H = 1.0e4 # Height position of the model top mesh = ExtrudedMesh(m, layers=nlayers, layer_height=H/nlayers) + domain = Domain(mesh, dt, "CG", 1) - output = OutputParameters(dirname=dirname, dumpfreq=1, - chkptfreq=2, log_level='INFO') parameters = CompressibleParameters() + eqns = CompressibleEulerEquations(domain, parameters) - state = State(mesh, - dt=dt, - output=output, - parameters=parameters) - - eqns = CompressibleEulerEquations(state, "CG", 1) + output = OutputParameters(dirname=dirname, dumpfreq=1, + chkptfreq=2, log_level='INFO') + io = IO(domain, output) # Set up transport schemes transported_fields = [] - transported_fields.append(SSPRK3(state, "u")) - transported_fields.append(SSPRK3(state, "rho")) - transported_fields.append(SSPRK3(state, "theta")) + transported_fields.append(SSPRK3(domain, "u")) + transported_fields.append(SSPRK3(domain, "rho")) + transported_fields.append(SSPRK3(domain, "theta")) # Set up linear solver - linear_solver = CompressibleSolver(state, eqns) + linear_solver = CompressibleSolver(eqns) # build time stepper - stepper = SemiImplicitQuasiNewton(eqns, state, transported_fields, + stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields, linear_solver=linear_solver) - return state, stepper, dt + initialise_fields(eqns, stepper) + + return stepper, dt -def initialise_fields(state): +def initialise_fields(eqns, stepper): L = 1.e5 H = 1.0e4 # Height position of the model top # Initial conditions - u0 = state.fields("u") - rho0 = state.fields("rho") - theta0 = state.fields("theta") + u0 = stepper.fields("u") + rho0 = stepper.fields("rho") + theta0 = stepper.fields("theta") # spaces Vt = theta0.function_space() @@ -64,11 +63,11 @@ def initialise_fields(state): # Thermodynamic constants required for setting initial conditions # and reference profiles - g = state.parameters.g - N = state.parameters.N + g = eqns.parameters.g + N = eqns.parameters.N # N^2 = (g/theta)dtheta/dz => dtheta/dz = theta N^2g => theta=theta_0exp(N^2gz) - x, z = SpatialCoordinate(state.mesh) + x, z = SpatialCoordinate(eqns.domain.mesh) Tsurf = 300. thetab = Tsurf*exp(N**2*z/g) @@ -76,7 +75,7 @@ def initialise_fields(state): rho_b = Function(Vr) # Calculate hydrostatic exner - compressible_hydrostatic_balance(state, theta_b, rho_b) + compressible_hydrostatic_balance(eqns, theta_b, rho_b) a = 5.0e3 deltaTheta = 1.0e-2 @@ -85,34 +84,32 @@ def initialise_fields(state): rho0.assign(rho_b) u0.project(as_vector([20.0, 0.0])) - state.set_reference_profiles([('rho', rho_b), ('theta', theta_b)]) + stepper.set_reference_profiles([('rho', rho_b), ('theta', theta_b)]) def test_checkpointing(tmpdir): dirname_1 = str(tmpdir)+'/checkpointing_1' dirname_2 = str(tmpdir)+'/checkpointing_2' - state_1, stepper_1, dt = setup_checkpointing(dirname_1) - state_2, stepper_2, dt = setup_checkpointing(dirname_2) + stepper_1, dt = setup_checkpointing(dirname_1) + stepper_2, dt = setup_checkpointing(dirname_2) # ------------------------------------------------------------------------ # # Run for 4 time steps and store values # ------------------------------------------------------------------------ # - initialise_fields(state_1) stepper_1.run(t=0.0, tmax=4*dt) # ------------------------------------------------------------------------ # # Start again, run for 2 time steps, checkpoint and then run for 2 more # ------------------------------------------------------------------------ # - initialise_fields(state_2) stepper_2.run(t=0.0, tmax=2*dt) # Wipe fields, then pickup - state_2.fields('u').project(as_vector([-10.0, 0.0])) - state_2.fields('rho').interpolate(Constant(0.0)) - state_2.fields('theta').interpolate(Constant(0.0)) + stepper_2.fields('u').project(as_vector([-10.0, 0.0])) + stepper_2.fields('rho').interpolate(Constant(0.0)) + stepper_2.fields('theta').interpolate(Constant(0.0)) stepper_2.run(t=2*dt, tmax=4*dt, pickup=True) @@ -124,14 +121,14 @@ def test_checkpointing(tmpdir): # This is the best way to compare fields from different meshes for field_name in ['u', 'rho', 'theta']: with DumbCheckpoint(dirname_1+'/chkpt', mode=FILE_READ) as chkpt: - field_1 = Function(state_1.fields(field_name).function_space(), + field_1 = Function(stepper_1.fields(field_name).function_space(), name=field_name) chkpt.load(field_1) # These are preserved in the comments for when we can use CheckpointFile # mesh = chkpt.load_mesh(name='firedrake_default_extruded') # field_1 = chkpt.load_function(mesh, name=field_name) with DumbCheckpoint(dirname_2+'/chkpt', mode=FILE_READ) as chkpt: - field_2 = Function(state_1.fields(field_name).function_space(), + field_2 = Function(stepper_1.fields(field_name).function_space(), name=field_name) chkpt.load(field_2) # These are preserved in the comments for when we can use CheckpointFile diff --git a/integration-tests/model/test_passive_tracer.py b/integration-tests/model/test_passive_tracer.py index c5ba70a56..63f1f5d5f 100644 --- a/integration-tests/model/test_passive_tracer.py +++ b/integration-tests/model/test_passive_tracer.py @@ -15,13 +15,16 @@ def run_tracer(setup): + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + # Get initial conditions from shared config - state = setup.state - mesh = state.mesh - dt = state.dt - output = state.output + domain = setup.domain + mesh = domain.mesh + io = setup.io - x = SpatialCoordinate(state.mesh) + x = SpatialCoordinate(mesh) H = 0.1 parameters = ShallowWaterParameters(H=H) Omega = parameters.Omega @@ -30,41 +33,44 @@ def run_tracer(setup): R = setup.radius fexpr = 2*Omega*x[2]/R - # Need to create a new state containing parameters - state = State(mesh, dt=dt, output=output, parameters=parameters) - # Equations - eqns = LinearShallowWaterEquations(state, setup.family, - setup.degree, fexpr=fexpr) - tracer_eqn = AdvectionEquation(state, state.spaces("DG"), "tracer") + eqns = LinearShallowWaterEquations(domain, parameters, fexpr=fexpr) + tracer_eqn = AdvectionEquation(domain, domain.spaces("DG"), "tracer") + + # set up transport schemes + transport_schemes = [ForwardEuler(domain, "D")] + + # Set up tracer transport + tracer_transport = [(tracer_eqn, SSPRK3(domain))] + + # build time stepper + stepper = SemiImplicitQuasiNewton( + eqns, io, transport_schemes, + auxiliary_equations_and_schemes=tracer_transport) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # # Specify initial prognostic fields - u0 = state.fields("u") - D0 = state.fields("D") - tracer0 = state.fields("tracer", D0.function_space()) + u0 = stepper.fields("u") + D0 = stepper.fields("D") + tracer0 = stepper.fields("tracer") tracer_end = Function(D0.function_space()) # Expressions for initial fields corresponding to Williamson 2 test case Dexpr = H - ((R * Omega * umax)*(x[2]*x[2]/(R*R))) / g u0.project(setup.uexpr) D0.interpolate(Dexpr) + Dbar = Function(D0.function_space()).assign(H) tracer0.interpolate(setup.f_init) tracer_end.interpolate(setup.f_end) - # set up transport schemes - transport_schemes = [ForwardEuler(state, "D")] - - # Set up tracer transport - tracer_transport = [(tracer_eqn, SSPRK3(state))] - - # build time stepper - stepper = SemiImplicitQuasiNewton( - eqns, state, transport_schemes, - auxiliary_equations_and_schemes=tracer_transport) + stepper.set_reference_profiles([('D', Dbar)]) stepper.run(t=0, tmax=setup.tmax) - error = norm(state.fields("tracer") - tracer_end) / norm(tracer_end) + error = norm(stepper.fields("tracer") - tracer_end) / norm(tracer_end) return error diff --git a/integration-tests/model/test_prescribed_transport.py b/integration-tests/model/test_prescribed_transport.py index a881c2f18..758a6349d 100644 --- a/integration-tests/model/test_prescribed_transport.py +++ b/integration-tests/model/test_prescribed_transport.py @@ -8,38 +8,38 @@ from firedrake import sin, cos, norm, pi, as_vector -def run(eqn, transport_scheme, state, tmax, f_end, prescribed_u): - timestepper = PrescribedTransport(eqn, transport_scheme, state, - prescribed_transporting_velocity=prescribed_u) +def run(timestepper, tmax, f_end): timestepper.run(0, tmax) - return norm(state.fields("f") - f_end) / norm(f_end) + return norm(timestepper.fields("f") - f_end) / norm(f_end) def test_prescribed_transport_setup(tmpdir, tracer_setup): - # Make mesh and state using routine from conftest + # Make domain using routine from conftest geometry = "slice" setup = tracer_setup(tmpdir, geometry, degree=1) - state = setup.state - _, z = SpatialCoordinate(state.mesh) + domain = setup.domain + _, z = SpatialCoordinate(domain.mesh) - V = state.spaces("DG", "DG", 1) + V = domain.spaces("DG") # Make equation - eqn = AdvectionEquation(state, V, "f", - ufamily=setup.family, udegree=1) + eqn = AdvectionEquation(domain, V, "f") # Initialise fields def u_evaluation(t): return as_vector([2.0*cos(2*pi*t/setup.tmax), sin(2*pi*t/setup.tmax)*sin(pi*z)]) - state.fields("f").interpolate(setup.f_init) - state.fields("u").project(u_evaluation(Constant(0.0))) + transport_scheme = SSPRK3(domain) - transport_scheme = SSPRK3(state) + timestepper = PrescribedTransport(eqn, transport_scheme, setup.io, + prescribed_transporting_velocity=u_evaluation) + + # Initial conditions + timestepper.fields("f").interpolate(setup.f_init) + timestepper.fields("u").project(u_evaluation(Constant(0.0))) # Run and check error - error = run(eqn, transport_scheme, state, setup.tmax, - setup.f_init, u_evaluation) + error = run(timestepper, setup.tmax, setup.f_init) assert error < setup.tol, \ 'The transport error is greater than the permitted tolerance' diff --git a/integration-tests/model/test_time_discretisation.py b/integration-tests/model/test_time_discretisation.py index c6172432e..0fe7c7f12 100644 --- a/integration-tests/model/test_time_discretisation.py +++ b/integration-tests/model/test_time_discretisation.py @@ -3,10 +3,9 @@ import pytest -def run(eqn, transport_scheme, state, tmax, f_end): - timestepper = PrescribedTransport(eqn, transport_scheme, state) +def run(timestepper, tmax, f_end): timestepper.run(0, tmax) - return norm(state.fields("f") - f_end) / norm(f_end) + return norm(timestepper.fields("f") - f_end) / norm(f_end) @pytest.mark.parametrize("scheme", ["ssprk", "implicit_midpoint", @@ -14,23 +13,26 @@ def run(eqn, transport_scheme, state, tmax, f_end): def test_time_discretisation(tmpdir, scheme, tracer_setup): geometry = "sphere" setup = tracer_setup(tmpdir, geometry) - state = setup.state - V = state.spaces("DG", "DG", 1) + domain = setup.domain + V = domain.spaces("DG") - eqn = AdvectionEquation(state, V, "f", ufamily=setup.family, - udegree=setup.degree) - - state.fields("f").interpolate(setup.f_init) - state.fields("u").project(setup.uexpr) + eqn = AdvectionEquation(domain, V, "f") if scheme == "ssprk": - transport_scheme = SSPRK3(state) + transport_scheme = SSPRK3(domain) elif scheme == "implicit_midpoint": - transport_scheme = ImplicitMidpoint(state) + transport_scheme = ImplicitMidpoint(domain) elif scheme == "RK4": - transport_scheme = RK4(state) + transport_scheme = RK4(domain) elif scheme == "Heun": - transport_scheme = Heun(state) + transport_scheme = Heun(domain) elif scheme == "BDF2": - transport_scheme = BDF2(state) - assert run(eqn, transport_scheme, state, setup.tmax, setup.f_end) < setup.tol + transport_scheme = BDF2(domain) + + timestepper = PrescribedTransport(eqn, transport_scheme, setup.io) + + # Initial conditions + timestepper.fields("f").interpolate(setup.f_init) + timestepper.fields("u").project(setup.uexpr) + + assert run(timestepper, setup.tmax, setup.f_end) < setup.tol diff --git a/integration-tests/physics/test_condensation.py b/integration-tests/physics/test_condensation.py index 9fdff1bb2..0b92d6f11 100644 --- a/integration-tests/physics/test_condensation.py +++ b/integration-tests/physics/test_condensation.py @@ -15,42 +15,59 @@ def run_cond_evap(dirname, process): + + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + dt = 2.0 + # declare grid shape, with length L and height H L = 1000. H = 1000. nlayers = int(H / 10.) ncolumns = int(L / 10.) - # make mesh + # make mesh and domain m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers)) + domain = Domain(mesh, dt, "CG", 1) + x, z = SpatialCoordinate(mesh) - dt = 2.0 + # spaces + + # Set up equation + tracers = [WaterVapour(), CloudWater()] + parameters = CompressibleParameters() + eqn = CompressibleEulerEquations(domain, parameters, active_tracers=tracers) + + # I/O output = OutputParameters(dirname=dirname+"/cond_evap", dumpfreq=1, dumplist=['u']) - parameters = CompressibleParameters() + io = IO(domain, output, diagnostic_fields=[Sum('water_vapour', 'cloud_water')]) - state = State(mesh, - dt=dt, - output=output, - parameters=parameters, - diagnostic_fields=[Sum('water_vapour', 'cloud_water')]) + # Physics scheme + physics_schemes = [(SaturationAdjustment(eqn, parameters=parameters), ForwardEuler(domain))] - # spaces - Vt = state.spaces("theta", degree=1) - Vr = state.spaces("DG", "DG", degree=1) + # Time stepper + scheme = ForwardEuler(domain) + stepper = SplitPhysicsTimestepper(eqn, scheme, io, + physics_schemes=physics_schemes) - # Set up equation -- use compressible to set up these spaces - tracers = [WaterVapour(), CloudWater()] - eqn = CompressibleEulerEquations(state, "CG", 1, active_tracers=tracers) + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + Vt = domain.spaces("theta", degree=1) + Vr = domain.spaces("DG", "DG", degree=1) # Declare prognostic fields - rho0 = state.fields("rho") - theta0 = state.fields("theta") - water_v0 = state.fields("water_vapour", Vt) - water_c0 = state.fields("cloud_water", Vt) + rho0 = stepper.fields("rho") + theta0 = stepper.fields("theta") + water_v0 = stepper.fields("water_vapour") + water_c0 = stepper.fields("cloud_water") # Set a background state with constant pressure and temperature pressure = Function(Vr).interpolate(Constant(100000.)) @@ -90,29 +107,26 @@ def run_cond_evap(dirname, process): eqn.residual = eqn.residual.label_map(lambda t: t.has_label(time_derivative), map_if_true=identity, map_if_false=drop) - physics_schemes = [(SaturationAdjustment(eqn, parameters), ForwardEuler(state))] - - # build time stepper - scheme = ForwardEuler(state) - stepper = SplitPhysicsTimestepper(eqn, scheme, state, - physics_schemes=physics_schemes) + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # stepper.run(t=0, tmax=dt) - return state, mv_true, mc_true, theta_d_true, mc_init + return eqn, stepper, mv_true, mc_true, theta_d_true, mc_init @pytest.mark.parametrize("process", ["evaporation", "condensation"]) def test_cond_evap(tmpdir, process): dirname = str(tmpdir) - state, mv_true, mc_true, theta_d_true, mc_init = run_cond_evap(dirname, process) + eqn, stepper, mv_true, mc_true, theta_d_true, mc_init = run_cond_evap(dirname, process) - water_v = state.fields('water_vapour') - water_c = state.fields('cloud_water') - theta_vd = state.fields('theta') + water_v = stepper.fields('water_vapour') + water_c = stepper.fields('cloud_water') + theta_vd = stepper.fields('theta') theta_d = Function(theta_vd.function_space()) - theta_d.interpolate(theta_vd/(1 + water_v * state.parameters.R_v / state.parameters.R_d)) + theta_d.interpolate(theta_vd/(1 + water_v * eqn.parameters.R_v / eqn.parameters.R_d)) # Check that water vapour is approximately equal to saturation amount assert norm(water_v - mv_true) / norm(mv_true) < 0.01, \ diff --git a/integration-tests/physics/test_instant_rain.py b/integration-tests/physics/test_instant_rain.py index 9d2cb0cf2..bfb65cd70 100644 --- a/integration-tests/physics/test_instant_rain.py +++ b/integration-tests/physics/test_instant_rain.py @@ -14,42 +14,54 @@ def run_instant_rain(dirname): - # set up mesh + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + + # set up mesh and domain L = 10 nx = 10 mesh = PeriodicSquareMesh(nx, nx, L) + dt = 0.1 + domain = Domain(mesh, dt, "BDM", 1) x, y = SpatialCoordinate(mesh) # parameters H = 30 g = 10 fexpr = Constant(0) - dt = 0.1 - output = OutputParameters(dirname=dirname+"/instant_rain", - dumpfreq=1, - dumplist=['vapour', "rain"]) + # Equation + vapour = WaterVapour(name="water_vapour", space='DG') + rain = Rain(name="rain", space="DG", + transport_eqn=TransportEquationType.no_transport) parameters = ShallowWaterParameters(H=H, g=g) + eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, + active_tracers=[vapour, rain]) + # I/O + output = OutputParameters(dirname=dirname+"/instant_rain", + dumpfreq=1, + dumplist=['vapour', "rain"]) diagnostic_fields = [CourantNumber()] + io = IO(domain, output, diagnostic_fields=diagnostic_fields) - state = State(mesh, - dt=dt, - output=output, - diagnostic_fields=diagnostic_fields, - parameters=parameters) - - vapour = WaterVapour(name="water_vapour", space='DG') - rain = Rain(name="rain", space="DG", - transport_eqn=TransportEquationType.no_transport) + # Physics schemes + # define saturation function + saturation = Constant(0.5) + physics_schemes = [(InstantRain(eqns, saturation, rain_name="rain", + set_tau_to_dt=True), ForwardEuler(domain))] - VD = FunctionSpace(mesh, "DG", 1) + # Time stepper + stepper = PrescribedTransport(eqns, RK4(domain), io, + physics_schemes=physics_schemes) - eqns = ShallowWaterEquations(state, "BDM", 1, fexpr=fexpr, - active_tracers=[vapour, rain]) + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # - vapour0 = state.fields("water_vapour") + vapour0 = stepper.fields("water_vapour") # set up vapour xc = L/2 @@ -60,31 +72,27 @@ def run_instant_rain(dirname): vapour0.interpolate(vapour_expr) + VD = FunctionSpace(mesh, "DG", 1) initial_vapour = Function(VD).interpolate(vapour_expr) - # define saturation function - saturation = Constant(0.5) - # define expected solutions; vapour should be equal to saturation and rain # should be (initial vapour - saturation) vapour_true = Function(VD).interpolate(saturation) rain_true = Function(VD).interpolate(vapour0 - saturation) - physics_schemes = [(InstantRain(eqns, saturation, rain_name="rain", - set_tau_to_dt=True), ForwardEuler(state))] - - stepper = PrescribedTransport(eqns, RK4(state), state, - physics_schemes=physics_schemes) + # ------------------------------------------------------------------------ # + # Run + # ------------------------------------------------------------------------ # stepper.run(t=0, tmax=5*dt) - return state, saturation, initial_vapour, vapour_true, rain_true + return stepper, saturation, initial_vapour, vapour_true, rain_true def test_instant_rain_setup(tmpdir): dirname = str(tmpdir) - state, saturation, initial_vapour, vapour_true, rain_true = run_instant_rain(dirname) - v = state.fields("water_vapour") - r = state.fields("rain") + stepper, saturation, initial_vapour, vapour_true, rain_true = run_instant_rain(dirname) + v = stepper.fields("water_vapour") + r = stepper.fields("rain") # check that the maximum of the vapour field is equal to the saturation assert v.dat.data.max() - saturation.values() < 0.001, "The maximum of the final vapour field should be equal to saturation" diff --git a/integration-tests/physics/test_precipitation.py b/integration-tests/physics/test_precipitation.py index d4d9ca679..aded26c70 100644 --- a/integration-tests/physics/test_precipitation.py +++ b/integration-tests/physics/test_precipitation.py @@ -13,38 +13,47 @@ def setup_fallout(dirname): + # ------------------------------------------------------------------------ # + # Set up model objects + # ------------------------------------------------------------------------ # + # declare grid shape, with length L and height H + dt = 0.1 L = 10. H = 10. nlayers = 10 ncolumns = 10 - # make mesh + # Domain m = PeriodicIntervalMesh(ncolumns, L) mesh = ExtrudedMesh(m, layers=nlayers, layer_height=(H / nlayers)) + domain = Domain(mesh, dt, "CG", 1) x = SpatialCoordinate(mesh) - dt = 0.1 - output = OutputParameters(dirname=dirname+"/fallout", - dumpfreq=10, - dumplist=['rain']) - parameters = CompressibleParameters() + # Equation + Vrho = domain.spaces("DG1_equispaced") + active_tracers = [Rain(space='DG1_equispaced')] + eqn = ForcedAdvectionEquation(domain, Vrho, "rho", active_tracers=active_tracers) + + # I/O + output = OutputParameters(dirname=dirname+"/fallout", dumpfreq=10, dumplist=['rain']) diagnostic_fields = [Precipitation()] - state = State(mesh, - dt=dt, - output=output, - parameters=parameters, - diagnostic_fields=diagnostic_fields) + io = IO(domain, output, diagnostic_fields=diagnostic_fields) - Vrho = state.spaces("DG1_equispaced") - active_tracers = [Rain(space='DG1_equispaced')] - eqn = ForcedAdvectionEquation(state, Vrho, "rho", ufamily="CG", udegree=1, - active_tracers=active_tracers) - scheme = ForwardEuler(state) - state.fields("rho").assign(1.) + # Physics schemes + physics_schemes = [(Fallout(eqn, 'rain', domain), SSPRK3(domain, 'rain'))] - physics_schemes = [(Fallout(eqn, 'rain', state), SSPRK3(state, 'rain'))] - rain0 = state.fields("rain") + # build time stepper + scheme = ForwardEuler(domain) + stepper = PrescribedTransport(eqn, scheme, io, + physics_schemes=physics_schemes) + + # ------------------------------------------------------------------------ # + # Initial conditions + # ------------------------------------------------------------------------ # + + stepper.fields("rho").assign(1.) + rain0 = stepper.fields("rain") # set up rain xc = L / 2 @@ -55,10 +64,6 @@ def setup_fallout(dirname): rain0.interpolate(rain_expr) - # build time stepper - stepper = PrescribedTransport(eqn, scheme, state, - physics_schemes=physics_schemes) - return stepper, 10.0 diff --git a/integration-tests/transport/test_dg_transport.py b/integration-tests/transport/test_dg_transport.py index 46e48372e..9b4abe75a 100644 --- a/integration-tests/transport/test_dg_transport.py +++ b/integration-tests/transport/test_dg_transport.py @@ -8,29 +8,32 @@ import pytest -def run(eqn, transport_scheme, state, tmax, f_end): - timestepper = PrescribedTransport(eqn, transport_scheme, state) +def run(timestepper, tmax, f_end): timestepper.run(0, tmax) - return norm(state.fields("f") - f_end) / norm(f_end) + return norm(timestepper.fields("f") - f_end) / norm(f_end) @pytest.mark.parametrize("geometry", ["slice", "sphere"]) @pytest.mark.parametrize("equation_form", ["advective", "continuity"]) def test_dg_transport_scalar(tmpdir, geometry, equation_form, tracer_setup): setup = tracer_setup(tmpdir, geometry) - state = setup.state - V = state.spaces("DG", "DG", 1) + domain = setup.domain + V = domain.spaces("DG") + if equation_form == "advective": - eqn = AdvectionEquation(state, V, "f", ufamily=setup.family, - udegree=setup.degree) + eqn = AdvectionEquation(domain, V, "f") else: - eqn = ContinuityEquation(state, V, "f", ufamily=setup.family, - udegree=setup.degree) - state.fields("f").interpolate(setup.f_init) - state.fields("u").project(setup.uexpr) + eqn = ContinuityEquation(domain, V, "f") + + transport_scheme = SSPRK3(domain) + + timestepper = PrescribedTransport(eqn, transport_scheme, setup.io) + + # Initial conditions + timestepper.fields("f").interpolate(setup.f_init) + timestepper.fields("u").project(setup.uexpr) - transport_scheme = SSPRK3(state) - error = run(eqn, transport_scheme, state, setup.tmax, setup.f_end) + error = run(timestepper, setup.tmax, setup.f_end) assert error < setup.tol, \ 'The transport error is greater than the permitted tolerance' @@ -39,20 +42,23 @@ def test_dg_transport_scalar(tmpdir, geometry, equation_form, tracer_setup): @pytest.mark.parametrize("equation_form", ["advective", "continuity"]) def test_dg_transport_vector(tmpdir, geometry, equation_form, tracer_setup): setup = tracer_setup(tmpdir, geometry) - state = setup.state - gdim = state.mesh.geometric_dimension() + domain = setup.domain + gdim = domain.mesh.geometric_dimension() f_init = as_vector([setup.f_init]*gdim) - V = VectorFunctionSpace(state.mesh, "DG", 1) + V = VectorFunctionSpace(domain.mesh, "DG", 1) if equation_form == "advective": - eqn = AdvectionEquation(state, V, "f", ufamily=setup.family, - udegree=setup.degree) + eqn = AdvectionEquation(domain, V, "f") else: - eqn = ContinuityEquation(state, V, "f", ufamily=setup.family, - udegree=setup.degree) - state.fields("f").interpolate(f_init) - state.fields("u").project(setup.uexpr) - transport_schemes = SSPRK3(state) + eqn = ContinuityEquation(domain, V, "f") + + transport_scheme = SSPRK3(domain) + + timestepper = PrescribedTransport(eqn, transport_scheme, setup.io) + + # Initial conditions + timestepper.fields("f").interpolate(f_init) + timestepper.fields("u").project(setup.uexpr) f_end = as_vector([setup.f_end]*gdim) - error = run(eqn, transport_schemes, state, setup.tmax, f_end) + error = run(timestepper, setup.tmax, f_end) assert error < setup.tol, \ 'The transport error is greater than the permitted tolerance' diff --git a/integration-tests/transport/test_embedded_dg_advection.py b/integration-tests/transport/test_embedded_dg_advection.py index 1df300601..d760a82f6 100644 --- a/integration-tests/transport/test_embedded_dg_advection.py +++ b/integration-tests/transport/test_embedded_dg_advection.py @@ -8,10 +8,9 @@ import pytest -def run(eqn, transport_schemes, state, tmax, f_end): - timestepper = PrescribedTransport(eqn, transport_schemes, state) +def run(timestepper, tmax, f_end): timestepper.run(0, tmax) - return norm(state.fields("f") - f_end) / norm(f_end) + return norm(timestepper.fields("f") - f_end) / norm(f_end) @pytest.mark.parametrize("ibp", [IntegrateByParts.ONCE, IntegrateByParts.TWICE]) @@ -20,25 +19,26 @@ def run(eqn, transport_schemes, state, tmax, f_end): def test_embedded_dg_advection_scalar(tmpdir, ibp, equation_form, space, tracer_setup): setup = tracer_setup(tmpdir, geometry="slice") - state = setup.state - V = state.spaces("theta", degree=1) + domain = setup.domain + V = domain.spaces("theta") if space == "broken": opts = EmbeddedDGOptions() elif space == "dg": - opts = EmbeddedDGOptions(embedding_space=state.spaces("DG1", "DG", 1)) + opts = EmbeddedDGOptions(embedding_space=domain.spaces("DG")) if equation_form == "advective": - eqn = AdvectionEquation(state, V, "f", ufamily=setup.family, - udegree=setup.degree, ibp=ibp) + eqn = AdvectionEquation(domain, V, "f", ibp=ibp) else: - eqn = ContinuityEquation(state, V, "f", ufamily=setup.family, - udegree=setup.degree, ibp=ibp) - state.fields("f").interpolate(setup.f_init) - state.fields("u").project(setup.uexpr) + eqn = ContinuityEquation(domain, V, "f", ibp=ibp) - transport_schemes = SSPRK3(state, options=opts) + transport_schemes = SSPRK3(domain, options=opts) + timestepper = PrescribedTransport(eqn, transport_schemes, setup.io) - error = run(eqn, transport_schemes, state, setup.tmax, setup.f_end) + # Initial conditions + timestepper.fields("f").interpolate(setup.f_init) + timestepper.fields("u").project(setup.uexpr) + + error = run(timestepper, setup.tmax, setup.f_end) assert error < setup.tol, \ 'The transport error is greater than the permitted tolerance' diff --git a/integration-tests/transport/test_limiters.py b/integration-tests/transport/test_limiters.py index 5672c8f1c..e98a0842e 100644 --- a/integration-tests/transport/test_limiters.py +++ b/integration-tests/transport/test_limiters.py @@ -1,5 +1,5 @@ """ -This tests three limiter options for different transport schemes. +This tests limiter options for different transport schemes. A sharp bubble of warm air is generated in a vertical slice and then transported by a prescribed transport scheme. If the limiter is working, the transport should have produced no new maxima or minima. @@ -26,46 +26,76 @@ def setup_limiters(dirname, space): rotations = 0.25 # ------------------------------------------------------------------------ # - # Mesh and spaces + # Build model objects # ------------------------------------------------------------------------ # + # Domain m = PeriodicIntervalMesh(20, Ld) mesh = ExtrudedMesh(m, layers=20, layer_height=(Ld/20)) - output = OutputParameters(dirname=dirname+'/limiters', - dumpfreq=1, dumplist=['u', 'tracer', 'true_tracer']) - parameters = CompressibleParameters() + degree = 0 if space in ['DG0', 'Vtheta_degree_0'] else 1 - state = State(mesh, - dt=dt, - output=output, - parameters=parameters) + domain = Domain(mesh, dt, family="CG", degree=degree) if space == 'DG0': - V = state.spaces('DG', 'DG', 0) + V = domain.spaces('DG') VCG1 = FunctionSpace(mesh, 'CG', 1) - VDG1 = state.spaces('DG1_equispaced') + VDG1 = domain.spaces('DG1_equispaced') + elif space == 'DG1': + V = domain.spaces('DG') elif space == 'DG1_equispaced': - V = state.spaces('DG1_equispaced') + V = domain.spaces('DG1_equispaced') elif space == 'Vtheta_degree_0': - V = state.spaces('theta', degree=0) + V = domain.spaces('theta') VCG1 = FunctionSpace(mesh, 'CG', 1) - VDG1 = state.spaces('DG1_equispaced') + VDG1 = domain.spaces('DG1_equispaced') elif space == 'Vtheta_degree_1': - V = state.spaces('theta', degree=1) + V = domain.spaces('theta') else: raise NotImplementedError - Vpsi = FunctionSpace(mesh, 'CG', 2) + Vpsi = domain.spaces('CG', 'CG', degree+1) + + # Equation + eqn = AdvectionEquation(domain, V, 'tracer') + + # I/O + output = OutputParameters(dirname=dirname+'/limiters', + dumpfreq=1, dumplist=['u', 'tracer', 'true_tracer']) + io = IO(domain, output) + + # ------------------------------------------------------------------------ # + # Set up transport scheme + # ------------------------------------------------------------------------ # + + if space in ['DG0', 'Vtheta_degree_0']: + opts = RecoveryOptions(embedding_space=VDG1, + recovered_space=VCG1, + project_low_method='recover', + boundary_method=BoundaryMethod.taylor) + transport_schemes = SSPRK3(domain, options=opts, + limiter=VertexBasedLimiter(VDG1)) + + elif space == 'DG1': + transport_schemes = SSPRK3(domain, limiter=DG1Limiter(V)) + + elif space == 'DG1_equispaced': + transport_schemes = SSPRK3(domain, limiter=VertexBasedLimiter(V)) + + elif space == 'Vtheta_degree_1': + opts = EmbeddedDGOptions() + transport_schemes = SSPRK3(domain, options=opts, limiter=ThetaLimiter(V)) + else: + raise NotImplementedError - # set up the equation - eqn = AdvectionEquation(state, V, 'tracer', ufamily='CG', udegree=1) + # Build time stepper + stepper = PrescribedTransport(eqn, transport_schemes, io) # ------------------------------------------------------------------------ # # Initial condition # ------------------------------------------------------------------------ # - tracer0 = state.fields('tracer', V) - true_field = state.fields('true_tracer', V) + tracer0 = stepper.fields('tracer', V) + true_field = stepper.fields('true_tracer', space=V) x, z = SpatialCoordinate(mesh) @@ -125,7 +155,7 @@ def setup_limiters(dirname, space): # ------------------------------------------------------------------------ # psi = Function(Vpsi) - u = state.fields('u') + u = stepper.fields('u') # set up solid body rotation for transport # we do this slightly complicated stream function to make the velocity 0 at edges @@ -149,42 +179,18 @@ def setup_limiters(dirname, space): gradperp = lambda v: as_vector([-v.dx(1), v.dx(0)]) u.project(gradperp(psi)) - # ------------------------------------------------------------------------ # - # Set up transport scheme - # ------------------------------------------------------------------------ # - - if space in ['DG0', 'Vtheta_degree_0']: - opts = RecoveryOptions(embedding_space=VDG1, - recovered_space=VCG1, - project_low_method='recover', - boundary_method=BoundaryMethod.taylor) - transport_schemes = SSPRK3(state, options=opts, - limiter=VertexBasedLimiter(VDG1)) - - elif space == 'DG1_equispaced': - transport_schemes = SSPRK3(state, limiter=VertexBasedLimiter(V)) - - elif space == 'Vtheta_degree_1': - opts = EmbeddedDGOptions() - transport_schemes = SSPRK3(state, options=opts, limiter=ThetaLimiter(V)) - else: - raise NotImplementedError - - # build time stepper - stepper = PrescribedTransport(eqn, transport_schemes, state) - - return stepper, tmax, state, true_field + return stepper, tmax, true_field @pytest.mark.parametrize('space', ['Vtheta_degree_0', 'Vtheta_degree_1', - 'DG0', 'DG1_equispaced']) + 'DG0', 'DG1', 'DG1_equispaced']) def test_limiters(tmpdir, space): # Setup and run dirname = str(tmpdir) - stepper, tmax, state, true_field = setup_limiters(dirname, space) + stepper, tmax, true_field = setup_limiters(dirname, space) stepper.run(t=0, tmax=tmax) - final_field = state.fields('tracer') + final_field = stepper.fields('tracer') # Check tracer is roughly in the correct place assert norm(true_field - final_field) / norm(true_field) < 0.05, \ diff --git a/integration-tests/transport/test_recovered_transport.py b/integration-tests/transport/test_recovered_transport.py index 35b3881b2..87160db95 100644 --- a/integration-tests/transport/test_recovered_transport.py +++ b/integration-tests/transport/test_recovered_transport.py @@ -9,41 +9,41 @@ import pytest -def run(eqn, transport_scheme, state, tmax, f_end): - timestepper = PrescribedTransport(eqn, transport_scheme, state) +def run(timestepper, tmax, f_end): timestepper.run(0, tmax) - return norm(state.fields("f") - f_end) / norm(f_end) + return norm(timestepper.fields("f") - f_end) / norm(f_end) @pytest.mark.parametrize("geometry", ["slice", "sphere"]) def test_recovered_space_setup(tmpdir, geometry, tracer_setup): - # Make mesh and state using routine from conftest + # Make domain using routine from conftest setup = tracer_setup(tmpdir, geometry, degree=0) - state = setup.state - mesh = state.mesh + domain = setup.domain + mesh = domain.mesh # Spaces for recovery VDG0 = FunctionSpace(mesh, "DG", 0) - VDG1 = state.spaces("DG1_equispaced") + VDG1 = domain.spaces("DG1_equispaced") VCG1 = FunctionSpace(mesh, "CG", 1) # Make equation - eqn = ContinuityEquation(state, VDG0, "f", - ufamily=setup.family, udegree=1) - - # Initialise fields - state.fields("f").interpolate(setup.f_init) - state.fields("u").project(setup.uexpr) + eqn = ContinuityEquation(domain, VDG0, "f") # Declare transport scheme recovery_opts = RecoveryOptions(embedding_space=VDG1, recovered_space=VCG1, boundary_method=BoundaryMethod.taylor) - transport_scheme = SSPRK3(state, options=recovery_opts) + transport_scheme = SSPRK3(domain, options=recovery_opts) + + timestepper = PrescribedTransport(eqn, transport_scheme, setup.io) + + # Initialise fields + timestepper.fields("f").interpolate(setup.f_init) + timestepper.fields("u").project(setup.uexpr) # Run and check error - error = run(eqn, transport_scheme, state, setup.tmax, setup.f_end) + error = run(timestepper, setup.tmax, setup.f_end) assert error < setup.tol, \ 'The transport error is greater than the permitted tolerance' diff --git a/integration-tests/transport/test_subcycling.py b/integration-tests/transport/test_subcycling.py index 7a019c586..09915ba18 100644 --- a/integration-tests/transport/test_subcycling.py +++ b/integration-tests/transport/test_subcycling.py @@ -8,28 +8,30 @@ import pytest -def run(eqn, transport_scheme, state, tmax, f_end): - timestepper = PrescribedTransport(eqn, transport_scheme, state) +def run(timestepper, tmax, f_end): timestepper.run(0, tmax) - return norm(state.fields("f") - f_end) / norm(f_end) + return norm(timestepper.fields("f") - f_end) / norm(f_end) @pytest.mark.parametrize("equation_form", ["advective", "continuity"]) def test_subcyling(tmpdir, equation_form, tracer_setup): geometry = "slice" setup = tracer_setup(tmpdir, geometry) - state = setup.state - V = state.spaces("DG", "DG", 1) + domain = setup.domain + V = domain.spaces("DG") if equation_form == "advective": - eqn = AdvectionEquation(state, V, "f", ufamily=setup.family, - udegree=setup.degree) + eqn = AdvectionEquation(domain, V, "f") else: - eqn = ContinuityEquation(state, V, "f", ufamily=setup.family, - udegree=setup.degree) - state.fields("f").interpolate(setup.f_init) - state.fields("u").project(setup.uexpr) + eqn = ContinuityEquation(domain, V, "f") - transport_scheme = SSPRK3(state, subcycles=2) - error = run(eqn, transport_scheme, state, setup.tmax, setup.f_end) + transport_scheme = SSPRK3(domain, subcycles=2) + + timestepper = PrescribedTransport(eqn, transport_scheme, setup.io) + + # Initial conditions + timestepper.fields("f").interpolate(setup.f_init) + timestepper.fields("u").project(setup.uexpr) + + error = run(timestepper, setup.tmax, setup.f_end) assert error < setup.tol, \ 'The transport error is greater than the permitted tolerance' diff --git a/integration-tests/transport/test_supg_transport.py b/integration-tests/transport/test_supg_transport.py index dab44a551..9475234f0 100644 --- a/integration-tests/transport/test_supg_transport.py +++ b/integration-tests/transport/test_supg_transport.py @@ -3,15 +3,14 @@ to the correct position. """ -from firedrake import norm, VectorFunctionSpace, as_vector +from firedrake import norm, FunctionSpace, VectorFunctionSpace, as_vector from gusto import * import pytest -def run(eqn, transport_scheme, state, tmax, f_end): - timestepper = PrescribedTransport(eqn, transport_scheme, state) +def run(timestepper, tmax, f_end): timestepper.run(0, tmax) - return norm(state.fields("f") - f_end) / norm(f_end) + return norm(timestepper.fields("f") - f_end) / norm(f_end) @pytest.mark.parametrize("equation_form", ["advective", "continuity"]) @@ -21,31 +20,34 @@ def test_supg_transport_scalar(tmpdir, equation_form, scheme, space, tracer_setup): setup = tracer_setup(tmpdir, geometry="slice") - state = setup.state + domain = setup.domain if space == "CG": - V = state.spaces("CG1", "CG", 1) + V = FunctionSpace(domain.mesh, "CG", 1) ibp = IntegrateByParts.NEVER else: - V = state.spaces("theta", degree=1) + V = domain.spaces("theta") ibp = IntegrateByParts.TWICE opts = SUPGOptions(ibp=ibp) if equation_form == "advective": - eqn = AdvectionEquation(state, V, "f", ufamily=setup.family, - udegree=setup.degree) + eqn = AdvectionEquation(domain, V, "f") else: - eqn = ContinuityEquation(state, V, "f", ufamily=setup.family, - udegree=setup.degree) - state.fields("f").interpolate(setup.f_init) - state.fields("u").project(setup.uexpr) + eqn = ContinuityEquation(domain, V, "f") + if scheme == "ssprk": - transport_scheme = SSPRK3(state, options=opts) + transport_scheme = SSPRK3(domain, options=opts) elif scheme == "implicit_midpoint": - transport_scheme = ImplicitMidpoint(state, options=opts) + transport_scheme = ImplicitMidpoint(domain, options=opts) + + timestepper = PrescribedTransport(eqn, transport_scheme, setup.io) + + # Initial conditions + timestepper.fields("f").interpolate(setup.f_init) + timestepper.fields("u").project(setup.uexpr) - error = run(eqn, transport_scheme, state, setup.tmax, setup.f_end) + error = run(timestepper, setup.tmax, setup.f_end) assert error < setup.tol, \ 'The transport error is greater than the permitted tolerance' @@ -57,37 +59,40 @@ def test_supg_transport_vector(tmpdir, equation_form, scheme, space, tracer_setup): setup = tracer_setup(tmpdir, geometry="slice") - state = setup.state + domain = setup.domain - gdim = state.mesh.geometric_dimension() + gdim = domain.mesh.geometric_dimension() f_init = as_vector([setup.f_init]*gdim) if space == "CG": - V = VectorFunctionSpace(state.mesh, "CG", 1) + V = VectorFunctionSpace(domain.mesh, "CG", 1) ibp = IntegrateByParts.NEVER else: - V = state.spaces("HDiv", setup.family, setup.degree) + V = domain.spaces("HDiv") ibp = IntegrateByParts.TWICE opts = SUPGOptions(ibp=ibp) if equation_form == "advective": - eqn = AdvectionEquation(state, V, "f", ufamily=setup.family, - udegree=setup.degree) + eqn = AdvectionEquation(domain, V, "f") else: - eqn = ContinuityEquation(state, V, "f", ufamily=setup.family, - udegree=setup.degree) - f = state.fields("f") + eqn = ContinuityEquation(domain, V, "f") + + if scheme == "ssprk": + transport_scheme = SSPRK3(domain, options=opts) + elif scheme == "implicit_midpoint": + transport_scheme = ImplicitMidpoint(domain, options=opts) + + timestepper = PrescribedTransport(eqn, transport_scheme, setup.io) + + # Initial conditions + f = timestepper.fields("f") if space == "CG": f.interpolate(f_init) else: f.project(f_init) - state.fields("u").project(setup.uexpr) - if scheme == "ssprk": - transport_scheme = SSPRK3(state, options=opts) - elif scheme == "implicit_midpoint": - transport_scheme = ImplicitMidpoint(state, options=opts) + timestepper.fields("u").project(setup.uexpr) f_end = as_vector([setup.f_end]*gdim) - error = run(eqn, transport_scheme, state, setup.tmax, f_end) + error = run(timestepper, setup.tmax, f_end) assert error < setup.tol, \ 'The transport error is greater than the permitted tolerance' diff --git a/integration-tests/transport/test_vector_recovered_space.py b/integration-tests/transport/test_vector_recovered_space.py index 3fd91e5dd..4b4c15221 100644 --- a/integration-tests/transport/test_vector_recovered_space.py +++ b/integration-tests/transport/test_vector_recovered_space.py @@ -9,26 +9,25 @@ import pytest -def run(eqn, transport_scheme, state, tmax, f_end): - timestepper = PrescribedTransport(eqn, transport_scheme, state) +def run(timestepper, tmax, f_end): timestepper.run(0, tmax) - return norm(state.fields("f") - f_end) / norm(f_end) + return norm(timestepper.fields("f") - f_end) / norm(f_end) @pytest.mark.parametrize("geometry", ["slice"]) def test_vector_recovered_space_setup(tmpdir, geometry, tracer_setup): - # Make mesh and state using routine from conftest + # Make domain using routine from conftest setup = tracer_setup(tmpdir, geometry, degree=0) - state = setup.state - mesh = state.mesh - gdim = state.mesh.geometric_dimension() + domain = setup.domain + mesh = domain.mesh + gdim = mesh.geometric_dimension() # Spaces for recovery - Vu = state.spaces("HDiv", family=setup.family, degree=setup.degree) + Vu = domain.spaces("HDiv") if geometry == "slice": - VDG1 = state.spaces("DG1_equispaced") + VDG1 = domain.spaces("DG1_equispaced") Vec_DG1 = VectorFunctionSpace(mesh, VDG1.ufl_element(), name='Vec_DG1') Vec_CG1 = VectorFunctionSpace(mesh, "CG", 1, name='Vec_CG1') @@ -40,19 +39,18 @@ def test_vector_recovered_space_setup(tmpdir, geometry, tracer_setup): f'Recovered spaces for geometry {geometry} have not been implemented') # Make equation - eqn = AdvectionEquation(state, Vu, "f", - ufamily=setup.family, udegree=1) + eqn = AdvectionEquation(domain, Vu, "f") + transport_scheme = SSPRK3(domain, options=rec_opts) + timestepper = PrescribedTransport(eqn, transport_scheme, setup.io) # Initialise fields f_init = as_vector([setup.f_init]*gdim) - state.fields("f").project(f_init) - state.fields("u").project(setup.uexpr) - - transport_scheme = SSPRK3(state, options=rec_opts) + timestepper.fields("f").project(f_init) + timestepper.fields("u").project(setup.uexpr) f_end = as_vector([setup.f_end]*gdim) # Run and check error - error = run(eqn, transport_scheme, state, setup.tmax, f_end) + error = run(timestepper, setup.tmax, f_end) assert error < setup.tol, \ 'The transport error is greater than the permitted tolerance'