Split State into Domain and IO

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)

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,51 +31,78 @@
 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
 Tsurf = 300.0
 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)