def les_smagorinsky_eddy_viscosity():
   from firedrake_fluids.les import LES
   
   errors = []
   for n in [2, 4, 8, 16, 32]:
      mesh = UnitSquareMesh(n, n)
      smagorinsky_coefficient = 2.0
      filter_width = CellVolume(mesh)**(1.0/2.0) # Square root of element area
      
      fs_exact = FunctionSpace(mesh, "CG", 3)
      fs = FunctionSpace(mesh, "CG", 1)
      vfs = VectorFunctionSpace(mesh, "CG", 1)
      
      u = Function(vfs).interpolate(Expression(('sin(x[0])', 'cos(x[0])')))
      density = Function(fs).interpolate(Expression("1.0"))
      
      exact_solution = project(Expression('pow(%f, 2) * sqrt(2.0*cos(x[0])*cos(x[0]) + 0.5*sin(x[0])*sin(x[0]) + 0.5*sin(x[0])*sin(x[0]))' % smagorinsky_coefficient), fs_exact)

      les = LES(mesh, fs, u, density, (smagorinsky_coefficient/filter_width))
      # Since the eddy viscosity depends on the filter_width, we need to provide smagorinsky_coefficient/filter_width here 
      # for the convergence test because we want to compare against the same exact solution throughout.
      les.solve()
      
      eddy_viscosity = les.eddy_viscosity
      print eddy_viscosity.vector().array()
      errors.append(sqrt(assemble(dot(eddy_viscosity - exact_solution, eddy_viscosity - exact_solution) * dx)))

   return errors
Exemple #2
0
def les_smagorinsky_eddy_viscosity():
    from firedrake_fluids.les import LES

    errors = []
    for n in [2, 4, 8, 16, 32]:
        mesh = UnitSquareMesh(n, n)
        smagorinsky_coefficient = 2.0
        filter_width = CellVolume(mesh)**(1.0 / 2.0
                                          )  # Square root of element area

        fs_exact = FunctionSpace(mesh, "CG", 3)
        fs = FunctionSpace(mesh, "CG", 1)
        vfs = VectorFunctionSpace(mesh, "CG", 1)

        u = Function(vfs).interpolate(Expression(('sin(x[0])', 'cos(x[0])')))
        density = Function(fs).interpolate(Expression("1.0"))

        exact_solution = project(
            Expression(
                'pow(%f, 2) * sqrt(2.0*cos(x[0])*cos(x[0]) + 0.5*sin(x[0])*sin(x[0]) + 0.5*sin(x[0])*sin(x[0]))'
                % smagorinsky_coefficient), fs_exact)

        les = LES(mesh, fs, u, density,
                  (smagorinsky_coefficient / filter_width))
        # Since the eddy viscosity depends on the filter_width, we need to provide smagorinsky_coefficient/filter_width here
        # for the convergence test because we want to compare against the same exact solution throughout.
        les.solve()

        eddy_viscosity = les.eddy_viscosity
        print eddy_viscosity.vector().array()
        errors.append(
            sqrt(
                assemble(
                    dot(eddy_viscosity - exact_solution,
                        eddy_viscosity - exact_solution) * dx)))

    return errors
   def run(self):
      """ Perform the simulation! """

      # Time-stepping parameters and constants
      LOG.info("Setting up a few constants...")
      T = self.options["T"]
      t = self.options["t"]
      theta = self.options["theta"]
      dt = self.options["dt"]
      dimension = self.options["dimension"]
      g_magnitude = self.options["g_magnitude"]

      # Get the function spaces
      U = self.function_spaces["VelocityFunctionSpace"]
      H = self.function_spaces["FreeSurfaceFunctionSpace"]
      
      # Is the Velocity field represented by a discontinous function space?
      dg = (U.ufl_element().family() == "Discontinuous Lagrange")
      
      # Weight u and h by theta to obtain the theta time-stepping scheme.
      assert(theta >= 0.0 and theta <= 1.0)
      LOG.info("Time-stepping scheme using theta = %g" % (theta))
      u_mid = (1.0 - theta) * self.u0 + theta * self.u
      h_mid = (1.0 - theta) * self.h0 + theta * self.h
         
      # The total height of the free surface.
      self.h_total = self.h_mean + self.h0

      # Non-linear approximation to the velocity
      u_nl = Function(U).assign(self.u0)
      
      # Second-order Adams-Bashforth velocity
      u_bash = (3.0/2.0)*self.u0 - (1.0/2.0)*self.u00
      
      # Simple P1 function space, to be used in the stabilisation routines (if applicable).
      P1 = FunctionSpace(self.mesh, "CG", 1)
      cellsize = CellSize(self.mesh)

      ###########################################################
      ################# Tentative velocity step #################
      ###########################################################
      
      # The collection of all the individual terms in their weak form.
      LOG.info("Constructing form...")
      F = 0

      # Mass term
      if(self.options["have_momentum_mass"]):
         LOG.debug("Momentum equation: Adding mass term...")
         M_momentum = (1.0/dt)*(inner(self.w, self.u) - inner(self.w, self.u0))*dx
         F += M_momentum
         
      # Append any Expression objects for weak BCs here.
      weak_bc_expressions = []
      
      # Advection term
      if(self.options["have_momentum_advection"]):
         LOG.debug("Momentum equation: Adding advection term...")

         if(self.options["integrate_advection_term_by_parts"]):
            outflow = (dot(self.u0, self.n) + abs(dot(self.u0, self.n)))/2.0

            A_momentum = -inner(dot(u_nl, grad(self.w)), u_bash)*dx - inner(dot(u_bash, grad(u_nl)), self.w)*dx
            A_momentum += inner(self.w, outflow*u_mid)*ds
            if(dg):
               # Only add interior facet integrals if we are dealing with a discontinous Galerkin discretisation.
               A_momentum += dot(outflow('+')*u_mid('+') - outflow('-')*u_mid('-'), jump(self.w))*dS

         else:
            A_momentum = inner(dot(grad(self.u), u_nl), self.w)*dx
         F += A_momentum

      # Viscous stress term. Note that the viscosity is kinematic (not dynamic).
      if(self.options["have_momentum_stress"]):
         LOG.debug("Momentum equation: Adding stress term...")
         
         viscosity = Function(H)
         
         # Background viscosity
         background_viscosity = Function(H).interpolate(Expression(libspud.get_option("/system/equations/momentum_equation/stress_term/scalar_field::Viscosity/value/constant")))
         viscosity.assign(background_viscosity)

         # Eddy viscosity
         if(self.options["have_turbulence_parameterisation"]):
            LOG.debug("Momentum equation: Adding turbulence parameterisation...")
            base_option_path = "/system/equations/momentum_equation/turbulence_parameterisation"
            # Large eddy simulation (LES)
            if(libspud.have_option(base_option_path + "/les")):
               les = LES(self.mesh, H)
               density = Constant(1.0) # We divide through by density in the momentum equation, so just set this to 1.0 for now.
               smagorinsky_coefficient = Constant(libspud.get_option(base_option_path + "/les/smagorinsky/smagorinsky_coefficient"))
               
               eddy_viscosity = Function(H)
               eddy_viscosity_lhs, eddy_viscosity_rhs = les.eddy_viscosity(u_mid, density, smagorinsky_coefficient)
               eddy_viscosity_problem = LinearVariationalProblem(eddy_viscosity_lhs, eddy_viscosity_rhs, eddy_viscosity, bcs=[])
               eddy_viscosity_solver = LinearVariationalSolver(eddy_viscosity_problem)
               
            # Add on eddy viscosity
            viscosity += eddy_viscosity

         # Stress tensor: tau = grad(u) + transpose(grad(u)) - (2/3)*div(u)
         if(not dg):
            # Perform a double dot product of the stress tensor and grad(w).
            K_momentum = -viscosity*inner(grad(self.u) + grad(self.u).T, grad(self.w))*dx
            #K_momentum += viscosity*(2.0/3.0)*inner(div(self.u)*Identity(dimension), grad(self.w))*dx
         else:
            # Interior penalty method
            cellsize = Constant(0.2) # In general, we should use CellSize(self.mesh) instead.
            alpha = 1/cellsize # Penalty parameter.
            
            K_momentum = -viscosity('+')*inner(grad(u_mid), grad(self.w))*dx
            for dim in range(self.options["dimension"]):
               K_momentum += -viscosity('+')*(alpha('+')/cellsize('+'))*dot(jump(self.w[dim], self.n), jump(u_mid[dim], self.n))*dS
               K_momentum += viscosity('+')*dot(avg(grad(self.w[dim])), jump(u_mid[dim], self.n))*dS + viscosity('+')*dot(jump(self.w[dim], self.n), avg(grad(u_mid[dim])))*dS

         F -= K_momentum # Negative sign here because we are bringing the stress term over from the RHS.

      # The gradient of the height of the free surface, h
      LOG.debug("Momentum equation: Adding gradient term...")
      C_momentum = -g_magnitude*inner(self.w, grad(self.h0))*dx
      F -= C_momentum

      # Quadratic drag term in the momentum equation
      if(self.options["have_drag"]):
         LOG.debug("Momentum equation: Adding drag term...")
         
         base_option_path = "/system/equations/momentum_equation/drag_term"
         
         # Get the bottom drag/friction coefficient.
         LOG.debug("Momentum equation: Adding bottom drag contribution...")
         bottom_drag = ExpressionFromOptions(path=base_option_path+"/scalar_field::BottomDragCoefficient/value", t=t).get_expression()
         bottom_drag = Function(H).interpolate(bottom_drag)
         
         # Add on the turbine drag, if provided.
         self.array = None
         
         # Magnitude of the velocity field
         magnitude = sqrt(dot(self.u0, self.u0))
         
         # Form the drag term
         array = sw.get_turbine_array()
         if(array):
            LOG.debug("Momentum equation: Adding turbine drag contribution...")
            drag_coefficient = bottom_drag + array.turbine_drag()
         else:
            drag_coefficient = bottom_drag
         D_momentum = -inner(self.w, (drag_coefficient*magnitude/self.h_total)*self.u)*dx
         F -= D_momentum

      # Add in any source terms
      if(self.options["have_momentum_source"]):
         LOG.debug("Momentum equation: Adding source term...")
         momentum_source_expression = ExpressionFromOptions(path = "/system/equations/momentum_equation/source_term/vector_field::Source/value", t=t).get_expression()
         momentum_source_function = Function(U)
         F -= inner(self.w, momentum_source_function.interpolate(momentum_source_expression))*dx
         
      # Add in any SU stabilisation
      if(self.options["have_su_stabilisation"]):
         LOG.debug("Momentum equation: Adding streamline-upwind stabilisation term...")
         stabilisation = Stabilisation(self.mesh, P1, cellsize)
         
         magnitude = magnitude_vector(self.u0, P1)

         # Bound the values for the magnitude below by 1.0e-9 for numerical stability reasons.
         u_nodes = magnitude.vector()
         near_zero = numpy.array([1.0e-9 for i in range(len(u_nodes))])
         u_nodes.set_local(numpy.maximum(u_nodes.array(), near_zero))

         diffusivity = ExpressionFromOptions(path = "/system/equations/momentum_equation/stress_term/scalar_field::Viscosity/value", t=self.options["t"]).get_expression()
         diffusivity = Function(H).interpolate(diffusivity) # Background viscosity
         grid_pe = grid_peclet_number(diffusivity, magnitude, P1, cellsize)
   
         # Bound the values for grid_pe below by 1.0e-9 for numerical stability reasons. 
         grid_pe_nodes = grid_pe.vector()
         values = numpy.array([1.0e-9 for i in range(len(grid_pe_nodes))])
         grid_pe_nodes.set_local(numpy.maximum(grid_pe_nodes.array(), values))

         F += stabilisation.streamline_upwind(self.w, self.u0, magnitude, grid_pe)

      LOG.info("Form construction complete.")

      ##########################################################
      ################ Pressure correction step ################
      ##########################################################   
      u_tent = Function(U)
      u_tent_nl = theta*u_tent + (1.0-theta)*self.u0
      F_h_corr = inner(self.v, (self.h - self.h0))*dx \
                 + g_magnitude*(dt**2)*(theta**2)*self.h_total*inner(grad(self.v), grad(self.h - self.h0))*dx \
                 + dt*self.v*div(self.h_total*u_tent_nl)*dx
                 
      ##########################################################
      ################ Velocity correction step ################
      ##########################################################   
      h1 = Function(H)
      u1 = Function(U)
      F_u_corr = (1.0/dt)*inner(self.w, self.u - u_tent)*dx + g_magnitude*theta*inner(self.w, grad(h1 - self.h0))*dx
                 
      LOG.info("Applying strong Dirichlet boundary conditions...")
      # Get all the Dirichlet boundary conditions for the Velocity field
      bcs_u = []; bcs_u2 = []
      bcs_h = []
      bc_expressions = []
      for i in range(0, libspud.option_count("/system/core_fields/vector_field::Velocity/boundary_condition")):
         if(libspud.have_option("/system/core_fields/vector_field::Velocity/boundary_condition[%d]/type::dirichlet" % i) and
            not libspud.have_option("/system/core_fields/vector_field::Velocity/boundary_condition[%d]/type::dirichlet/apply_weakly" % i)):
            expr = ExpressionFromOptions(path = ("/system/core_fields/vector_field::Velocity/boundary_condition[%d]/type::dirichlet" % i), t=t).get_expression()
            # Surface IDs on the domain boundary
            surface_ids = libspud.get_option("/system/core_fields/vector_field::Velocity/boundary_condition[%d]/surface_ids" % i)
            method = ("geometric" if dg else "topological")
            bc = DirichletBC(U, expr, surface_ids, method=method)
            bcs_u.append(bc)
            bc_expressions.append(expr)
            LOG.debug("Applying Velocity BC #%d strongly to surface IDs: %s" % (i, surface_ids))

      for i in range(0, libspud.option_count("/system/core_fields/vector_field::Velocity/boundary_condition")):
         if(libspud.have_option("/system/core_fields/vector_field::Velocity/boundary_condition[%d]/type::dirichlet" % i) and
            not libspud.have_option("/system/core_fields/vector_field::Velocity/boundary_condition[%d]/type::dirichlet/apply_weakly" % i)):
            expr = ExpressionFromOptions(path = ("/system/core_fields/vector_field::Velocity/boundary_condition[%d]/type::dirichlet" % i), t=t).get_expression()
            # Surface IDs on the domain boundary
            surface_ids = libspud.get_option("/system/core_fields/vector_field::Velocity/boundary_condition[%d]/surface_ids" % i)
            method = ("geometric" if dg else "topological")
            bc = DirichletBC(U, expr, surface_ids, method=method)
            bcs_u2.append(bc)
            bc_expressions.append(expr)
            LOG.debug("Applying Velocity BC #%d strongly to surface IDs: %s" % (i, surface_ids))


      for i in range(0, libspud.option_count("/system/core_fields/scalar_field::FreeSurfacePerturbation/boundary_condition/type::dirichlet")):
         if(libspud.have_option("/system/core_fields/scalar_field::FreeSurfacePerturbation/boundary_condition[%d]/type::dirichlet" % i) and
            not(libspud.have_option("/system/core_fields/scalar_field::FreeSurfacePerturbation/boundary_condition[%d]/type::dirichlet/apply_weakly" % i))):
            expr = ExpressionFromOptions(path = ("/system/core_fields/scalar_field::FreeSurfacePerturbation/boundary_condition[%d]/type::dirichlet" % i), t=t).get_expression()
            # Surface IDs on the domain boundary
            surface_ids = libspud.get_option("/system/core_fields/scalar_field::FreeSurfacePerturbation/boundary_condition[%d]/surface_ids" % i)
            method = ("geometric" if dg else "topological")
            bc = DirichletBC(H, expr, surface_ids, method=method)
            bcs_h.append(bc)
            bc_expressions.append(expr)
            LOG.debug("Applying FreeSurfacePerturbation BC #%d strongly to surface IDs: %s" % (i, surface_ids))
            
      # Prepare solver_parameters dictionary
      LOG.debug("Defining solver_parameters dictionary...")
      solver_parameters = {'ksp_monitor': True, 'ksp_view': False, 'pc_view': False, 'snes_type': 'ksponly', 'ksp_max_it':10000} # NOTE: use 'snes_type': 'newtonls' for production runs.
      
      # KSP (iterative solver) options
      solver_parameters["ksp_type"] = libspud.get_option("/system/solver/iterative_method/name")
      solver_parameters["ksp_rtol"] = libspud.get_option("/system/solver/relative_error")
      
      solver_parameters['ksp_converged_reason'] = True
      solver_parameters['ksp_monitor_true_residual'] = True
      
      # Preconditioner options
      solver_parameters["pc_type"] = libspud.get_option("/system/solver/preconditioner/name")
      # Fieldsplit sub-options
      if(solver_parameters["pc_type"] == "fieldsplit"):
         LOG.debug("Setting up the fieldsplit preconditioner...")
         solver_parameters["pc_fieldsplit_type"] = libspud.get_option("/system/solver/preconditioner::fieldsplit/type/name")
         if(solver_parameters["pc_fieldsplit_type"] == "schur"):
            solver_parameters["pc_fieldsplit_schur_fact_type"] = libspud.get_option("/system/solver/preconditioner::fieldsplit/type::schur/fact_type/name")
         solver_parameters["fieldsplit_0_ksp_type"] = libspud.get_option("/system/solver/preconditioner::fieldsplit/block_0_ksp_type/iterative_method/name")
         solver_parameters["fieldsplit_1_ksp_type"] = libspud.get_option("/system/solver/preconditioner::fieldsplit/block_1_ksp_type/iterative_method/name")

         if(libspud.get_option("/system/solver/preconditioner::fieldsplit/block_0_pc_type/preconditioner/name") != "ilu"):
            solver_parameters["fieldsplit_0_pc_type"] = libspud.get_option("/system/solver/preconditioner::fieldsplit/block_0_pc_type/preconditioner/name")
            solver_parameters["fieldsplit_1_pc_type"] = libspud.get_option("/system/solver/preconditioner::fieldsplit/block_1_pc_type/preconditioner/name")
            
         # Enable inner iteration monitors.
         solver_parameters["fieldsplit_0_ksp_monitor"] = True
         solver_parameters["fieldsplit_1_ksp_monitor"] = True
         solver_parameters["fieldsplit_0_pc_factor_shift_type"] = 'INBLOCKS'
         solver_parameters["fieldsplit_1_pc_factor_shift_type"] = 'INBLOCKS'
         
      # Construct the solver objects
      problem_tent = LinearVariationalProblem(lhs(F), rhs(F), u_tent, bcs=bcs_u)
      solver_tent = LinearVariationalSolver(problem_tent, solver_parameters={'ksp_monitor': False, 
                                                                              'ksp_view': False, 
                                                                              'pc_view': False, 
                                                                              'pc_type': 'sor',
                                                                              'ksp_type': 'gmres',
                                                                              'ksp_rtol': 1.0e-7})

      problem_h_corr = LinearVariationalProblem(lhs(F_h_corr), rhs(F_h_corr), h1, bcs=bcs_h)
      solver_h_corr = LinearVariationalSolver(problem_h_corr, solver_parameters={'ksp_monitor': False, 
                                                                                 'ksp_view': False, 
                                                                                 'pc_view': False, 
                                                                                 'pc_type': 'sor',
                                                                                 'ksp_type': 'gmres',
                                                                                 'ksp_rtol': 1.0e-7})
      
     
      problem_u_corr = LinearVariationalProblem(lhs(F_u_corr), rhs(F_u_corr), u1, bcs=bcs_u2)
      solver_u_corr = LinearVariationalSolver(problem_u_corr, solver_parameters={'ksp_monitor': False, 
                                                                                 'ksp_view': False, 
                                                                                 'pc_view': False, 
                                                                                 'pc_type': 'sor',
                                                                                 'ksp_type': 'gmres',
                                                                                 'ksp_rtol': 1.0e-7})
      
      
      t += dt
      iterations_since_dump = 1
      iterations_since_checkpoint = 1
      
      # PETSc solver run-times
      from petsc4py import PETSc
      main_solver_stage = PETSc.Log.Stage('Main block-coupled system solve')

      total_solver_time = 0.0
      # The time-stepping loop
      LOG.info("Entering the time-stepping loop...")
      EPSILON = 1.0e-14
      
      while t <= T + EPSILON: # A small value EPSILON is added here in case of round-off error.
         LOG.info("t = %g" % t)

         while True:
            ## Update any time-dependent Functions and Expressions.
            
            # Re-compute the velocity magnitude and grid Peclet number fields.
            if(self.options["have_su_stabilisation"]):
               magnitude.assign(magnitude_vector(self.u0, P1))

               # Bound the values for the magnitude below by 1.0e-9 for numerical stability reasons.
               u_nodes = magnitude.vector()
               near_zero = numpy.array([1.0e-9 for i in range(len(u_nodes))])
               u_nodes.set_local(numpy.maximum(u_nodes.array(), near_zero))
               
               grid_pe.assign(grid_peclet_number(diffusivity, magnitude, P1, cellsize))
         
               # Bound the values for grid_pe below by 1.0e-9 for numerical stability reasons. 
               grid_pe_nodes = grid_pe.vector()
               values = numpy.array([1.0e-9 for i in range(len(grid_pe_nodes))])
               grid_pe_nodes.set_local(numpy.maximum(grid_pe_nodes.array(), values))

            if(self.options["have_turbulence_parameterisation"]):
               eddy_viscosity_solver.solve()
               viscosity.assign(background_viscosity + eddy_viscosity)

            # Time-dependent source terms
            if(self.options["have_momentum_source"]):
               momentum_source_expression.t = t
               momentum_source_function.interpolate(momentum_source_expression)
            if(self.options["have_continuity_source"]):
               continuity_source_expression.t = t
               continuity_source_function.interpolate(continuity_source_expression)
               
            # Update any time-varying DirichletBC objects.
            for expr in bc_expressions:
               expr.t = t
            for expr in weak_bc_expressions:
               expr.t = t
                      
            # Solve the system of equations!
            start_solver_time = mpi4py.MPI.Wtime()
            main_solver_stage.push()
            LOG.debug("Solving the system of equations...")
            
            solver_tent.solve()
            solver_h_corr.solve()
            solver_u_corr.solve()
            
            main_solver_stage.pop()
            end_solver_time = mpi4py.MPI.Wtime()
            total_solver_time += (end_solver_time - start_solver_time)

            # Move to next time step  
            if(steady_state(u1, u_nl, 1e-7)):
               break
            u_nl.assign(u1)

         self.u00.assign(self.u0)
         self.u0.assign(u1)
         self.h0.assign(h1)
         t += dt
         iterations_since_dump += 1
         iterations_since_checkpoint += 1
         LOG.debug("Moving to next time level...")

         # Write the solution to file.
         if((self.options["dump_period"] is not None) and (dt*iterations_since_dump >= self.options["dump_period"])):
            LOG.debug("Writing data to file...")
            self.output_functions["Velocity"].assign(u1)
            self.output_files["Velocity"] << self.output_functions["Velocity"]
            self.output_functions["FreeSurfacePerturbation"].assign(h1)
            self.output_files["FreeSurfacePerturbation"] << self.output_functions["FreeSurfacePerturbation"]
            iterations_since_dump = 0 # Reset the counter.

         # Print out the total power generated by turbines.
         if(self.options["have_drag"] and self.array is not None):
            LOG.info("Power = %.2f" % self.array.power(u1, density=1000))
            
         # Checkpointing
         if((self.options["checkpoint_period"] is not None) and (dt*iterations_since_checkpoint >= self.options["checkpoint_period"])):
            LOG.debug("Writing checkpoint data to file...")
            self.solution.dat.save("checkpoint")
            iterations_since_checkpoint = 0 # Reset the counter.
            
         # Check whether a steady-state has been reached.
         if(steady_state(u1, self.u0, self.options["steady_state_tolerance"]) and steady_state(h1, self.h0, self.options["steady_state_tolerance"])):
            LOG.info("Steady-state attained. Exiting the time-stepping loop...")
            break

         self.compute_diagnostics()            


      LOG.info("Out of the time-stepping loop.")

      LOG.debug("Total solver time: %.2f" % (total_solver_time))

      return u1, h1
   def run(self, array=None, annotate=False, checkpoint=None):
      """ Perform the simulation! """
      
      # The solution field defined on the mixed function space
      solution = Function(self.W, name="Solution", annotate=annotate)
      # The solution from the previous time-step. At t=0, this holds the initial conditions.
      solution_old = Function(self.W, name="SolutionOld", annotate=annotate)
      
      # Assign the initial condition
      initial_condition = self.get_initial_condition(checkpoint=checkpoint)
      solution_old.assign(initial_condition, annotate=annotate)

      # Get the test functions
      test_functions = TestFunctions(self.W)
      w = test_functions[0]; v = test_functions[1]
      LOG.info("Test functions created.")
      
      # These are like the TrialFunctions, but are just regular Functions here because we want to solve a non-linear problem
      # 'u' and 'h' are the velocity and free surface perturbation, respectively.
      functions = split(solution)
      u = functions[0]; h = functions[1]
      LOG.info("Trial functions created.")

      functions_old = split(solution_old)
      u_old = functions_old[0]; h_old = functions_old[1]

      # Write initial conditions to file
      LOG.info("Writing initial conditions to file...")
      self.output_functions["Velocity"].assign(solution_old.split()[0], annotate=False)
      self.output_files["Velocity"] << self.output_functions["Velocity"]
      self.output_functions["FreeSurfacePerturbation"].assign(solution_old.split()[1], annotate=False)
      self.output_files["FreeSurfacePerturbation"] << self.output_functions["FreeSurfacePerturbation"]
      
      # Construct the collection of all the individual terms in their weak form.
      LOG.info("Constructing form...")
      F = 0
      
      theta = self.options["theta"]
      dt = self.options["dt"]
      dimension = self.options["dimension"]
      g_magnitude = self.options["g_magnitude"]
      
      # Is the Velocity field represented by a discontinous function space?
      dg = (self.W.sub(0).ufl_element().family() == "Discontinuous Lagrange")

      # Mean free surface height
      h_mean = Function(self.W.sub(1), name="FreeSurfaceMean", annotate=False)
      h_mean.interpolate(ExpressionFromOptions(path = "/system/core_fields/scalar_field::FreeSurfaceMean/value").get_expression())

      # Weight u and h by theta to obtain the theta time-stepping scheme.
      assert(theta >= 0.0 and theta <= 1.0)
      LOG.info("Time-stepping scheme using theta = %g" % (theta))
      u_mid = (1.0 - theta) * u_old + theta * u
      h_mid = (1.0 - theta) * h_old + theta * h
         
      # The total height of the free surface.
      H = h_mean + h
      
      # Simple P1 function space, to be used in the stabilisation routines (if applicable).
      P1 = FunctionSpace(self.mesh, "CG", 1)
      cellsize = CellSize(self.mesh)

      # Normal vector to each element facet
      n = FacetNormal(self.mesh)
      
      # Mass term
      if(self.options["have_momentum_mass"]):
         LOG.debug("Momentum equation: Adding mass term...")
         M_momentum = (1.0/dt)*(inner(w, u) - inner(w, u_old))*dx
         F += M_momentum
      
      # Advection term
      if(self.options["have_momentum_advection"]):
         LOG.debug("Momentum equation: Adding advection term...")

         if(self.options["integrate_advection_term_by_parts"]):
            outflow = (dot(u_mid, n) + abs(dot(u_mid, n)))/2.0

            A_momentum = -inner(dot(u_mid, grad(w)), u_mid)*dx - inner(dot(u_mid, grad(u_mid)), w)*dx
            A_momentum += inner(w, outflow*u_mid)*ds
            if(dg):
               # Only add interior facet integrals if we are dealing with a discontinous Galerkin discretisation.
               A_momentum += dot(outflow('+')*u_mid('+') - outflow('-')*u_mid('-'), jump(w))*dS

         else:
            A_momentum = inner(dot(grad(u_mid), u_mid), w)*dx
         F += A_momentum
         
      # Viscous stress term. Note that the viscosity is kinematic (not dynamic).
      if(self.options["have_momentum_stress"]):
         LOG.debug("Momentum equation: Adding stress term...")
         
         viscosity = Function(self.W.sub(1))
         
         # Background viscosity
         background_viscosity = Function(self.W.sub(1)).interpolate(Expression(libspud.get_option("/system/equations/momentum_equation/stress_term/scalar_field::Viscosity/value/constant")))
         viscosity.assign(background_viscosity)

         # Eddy viscosity
         if(self.options["have_turbulence_parameterisation"]):
            LOG.debug("Momentum equation: Adding turbulence parameterisation...")
            base_option_path = "/system/equations/momentum_equation/turbulence_parameterisation"
            # Large eddy simulation (LES)
            if(libspud.have_option(base_option_path + "/les")):
               
               density = Constant(1.0) # We divide through by density in the momentum equation, so just set this to 1.0 for now.
               smagorinsky_coefficient = Constant(libspud.get_option(base_option_path + "/les/smagorinsky/smagorinsky_coefficient"))
               les = LES(self.mesh, self.W.sub(1), u_mid, density, smagorinsky_coefficient)

               # Add on eddy viscosity
               viscosity += les.eddy_viscosity

         # Stress tensor: tau = grad(u) + transpose(grad(u)) - (2/3)*div(u)
         if(not dg):
            # Perform a double dot product of the stress tensor and grad(w).
            K_momentum = -viscosity*inner(grad(u_mid) + grad(u_mid).T, grad(w))*dx
            K_momentum += viscosity*(2.0/3.0)*inner(div(u_mid)*Identity(dimension), grad(w))*dx
         else:
            # Interior penalty method
            cellsize = Constant(0.2) # In general, we should use CellSize(self.mesh) instead.
            alpha = 1/cellsize # Penalty parameter.
            
            K_momentum = -viscosity('+')*inner(grad(u_mid), grad(w))*dx
            for dim in range(self.options["dimension"]):
               K_momentum += -viscosity('+')*(alpha('+')/cellsize('+'))*dot(jump(w[dim], n), jump(u_mid[dim], n))*dS
               K_momentum += viscosity('+')*dot(avg(grad(w[dim])), jump(u_mid[dim], n))*dS + viscosity('+')*dot(jump(w[dim], n), avg(grad(u_mid[dim])))*dS

         F -= K_momentum # Negative sign here because we are bringing the stress term over from the RHS.

      # The gradient of the height of the free surface, h
      LOG.debug("Momentum equation: Adding gradient term...")
      C_momentum = -g_magnitude*inner(w, grad(h_mid))*dx
      F -= C_momentum
      
      # Quadratic drag term in the momentum equation
      if(self.options["have_drag"]):
         LOG.debug("Momentum equation: Adding drag term...")
         
         base_option_path = "/system/equations/momentum_equation/drag_term"
         
         # Get the bottom drag/friction coefficient.
         LOG.debug("Momentum equation: Adding bottom drag contribution...")
         bottom_drag = ExpressionFromOptions(path=base_option_path+"/scalar_field::BottomDragCoefficient/value", t=0).get_expression()
         bottom_drag = Function(self.W.sub(1)).interpolate(bottom_drag)
         
         # Magnitude of the velocity field
         magnitude = sqrt(dot(u_old, u_old))
         
         # Form the drag term
         if(array):
            LOG.debug("Momentum equation: Adding turbine drag contribution...")
            drag_coefficient = bottom_drag + array.turbine_drag()
         else:
            drag_coefficient = bottom_drag
         D_momentum = -inner(w, (drag_coefficient*magnitude/H)*u_mid)*dx
         F -= D_momentum
      
      # The mass term in the shallow water continuity equation 
      # (i.e. an advection equation for the free surface height, h)
      if(self.options["have_continuity_mass"]):
         LOG.debug("Continuity equation: Adding mass term...")
         M_continuity = (1.0/dt)*(inner(v, h) - inner(v, h_old))*dx
         F += M_continuity

      # Append any Expression objects for weak BCs here.
      weak_bc_expressions = []
      
      # Divergence term in the shallow water continuity equation
      LOG.debug("Continuity equation: Adding divergence term...")
      if(self.options["integrate_continuity_equation_by_parts"]):
         LOG.debug("The divergence term is being integrated by parts.")
         Ct_continuity = - H*inner(u_mid, grad(v))*dx
         if(dg):
            Ct_continuity += inner(jump(v, n), avg(H*u_mid))*dS
                           
         # Add in the surface integrals, but check to see if any boundary conditions need to be applied weakly here.
         boundary_markers = self.mesh.exterior_facets.unique_markers
         for marker in boundary_markers:
            marker = int(marker) # ds() will not accept markers of type 'numpy.int32', so convert it to type 'int' here.
            
            bc_type = None
            for i in range(0, libspud.option_count("/system/core_fields/vector_field::Velocity/boundary_condition")):
               bc_path = "/system/core_fields/vector_field::Velocity/boundary_condition[%d]" % i
               if(not (marker in libspud.get_option(bc_path + "/surface_ids"))):
                  # This BC is not associated with this marker, so skip it.
                  continue
                  
               # Determine the BC type.
               if(libspud.have_option(bc_path + "/type::no_normal_flow")):
                  bc_type = "no_normal_flow"
               elif(libspud.have_option(bc_path + "/type::dirichlet")):
                  if(libspud.have_option(bc_path + "/type::dirichlet/apply_weakly")):
                     bc_type = "weak_dirichlet"
                  else:
                     bc_type = "dirichlet"
               elif(libspud.have_option(bc_path + "/type::flather")):
                  bc_type = "flather"
                  
               # Apply the boundary condition...
               try:
                  LOG.debug("Applying Velocity BC of type '%s' to surface ID %d..." % (bc_type, marker))
                  if(bc_type == "flather"):
                     # The known exterior value for the Velocity.
                     u_ext = ExpressionFromOptions(path = (bc_path + "/type::flather/exterior_velocity"), t=0).get_expression()
                     Ct_continuity += H*inner(Function(self.W.sub(0)).interpolate(u_ext), n)*v*ds(int(marker))
                     
                     # The known exterior value for the FreeSurfacePerturbation.
                     h_ext = ExpressionFromOptions(path = (bc_path + "/type::flather/exterior_free_surface_perturbation"), t=0).get_expression()
                     Ct_continuity += H*sqrt(g_magnitude/H)*(h_mid - Function(self.W.sub(1)).interpolate(h_ext))*v*ds(int(marker))
                     
                     weak_bc_expressions.append(u_ext)
                     weak_bc_expressions.append(h_ext)
                     
                  elif(bc_type == "weak_dirichlet"):
                     u_bdy = ExpressionFromOptions(path = (bc_path + "/type::dirichlet"), t=0).get_expression()
                     Ct_continuity += H*(dot(Function(self.W.sub(0)).interpolate(u_bdy), n))*v*ds(int(marker))
                     
                     weak_bc_expressions.append(u_bdy)
                     
                  elif(bc_type == "dirichlet"):
                     # Add in the surface integral as it is here. The BC will be applied strongly later using a DirichletBC object.
                     Ct_continuity += H * inner(u_mid, n) * v * ds(int(marker))
                  elif(bc_type == "no_normal_flow"):
                     # Do nothing here since dot(u, n) is zero.
                     continue
                  else:
                     raise ValueError("Unknown boundary condition type!")
                     
               except ValueError as e:
                  LOG.exception(e)
                  sys.exit()
                  
            # If no boundary condition has been applied, include the surface integral as it is.
            if(bc_type is None):
               Ct_continuity += H * inner(u_mid, n) * v * ds(int(marker))

      else:
         Ct_continuity = inner(v, div(H*u_mid))*dx
      F += Ct_continuity

      # Add in any source terms
      if(self.options["have_momentum_source"]):
         LOG.debug("Momentum equation: Adding source term...")
         momentum_source_expression = ExpressionFromOptions(path = "/system/equations/momentum_equation/source_term/vector_field::Source/value", t=0).get_expression()
         momentum_source_function = Function(self.W.sub(0), annotate=False)
         F -= inner(w, momentum_source_function.interpolate(momentum_source_expression))*dx

      if(self.options["have_continuity_source"]):
         LOG.debug("Continuity equation: Adding source term...")
         continuity_source_expression = ExpressionFromOptions(path = "/system/equations/continuity_equation/source_term/scalar_field::Source/value", t=0).get_expression()
         continuity_source_function = Function(self.W.sub(1), annotate=False)
         F -= inner(v, continuity_source_function.interpolate(continuity_source_expression))*dx
         
      # Add in any SU stabilisation
      if(self.options["have_su_stabilisation"]):
         LOG.debug("Momentum equation: Adding streamline-upwind stabilisation term...")
         stabilisation = Stabilisation(self.mesh, P1, cellsize)
         
         magnitude = magnitude_vector(solution_old.split()[0], P1)

         # Bound the values for the magnitude below by 1.0e-9 for numerical stability reasons.
         u_nodes = magnitude.vector()
         near_zero = numpy.array([1.0e-9 for i in range(len(u_nodes))])
         u_nodes.set_local(numpy.maximum(u_nodes.array(), near_zero))

         diffusivity = ExpressionFromOptions(path = "/system/equations/momentum_equation/stress_term/scalar_field::Viscosity/value", t=self.options["t"]).get_expression()
         diffusivity = Function(self.W.sub(1)).interpolate(diffusivity) # Background viscosity
         grid_pe = grid_peclet_number(diffusivity, magnitude, P1, cellsize)
   
         # Bound the values for grid_pe below by 1.0e-9 for numerical stability reasons. 
         grid_pe_nodes = grid_pe.vector()
         values = numpy.array([1.0e-9 for i in range(len(grid_pe_nodes))])
         grid_pe_nodes.set_local(numpy.maximum(grid_pe_nodes.array(), values))

         F += stabilisation.streamline_upwind(w, u, magnitude, grid_pe)

      LOG.info("Form construction complete.")
      
      bcs, bc_expressions = self.get_dirichlet_boundary_conditions()
      
      # Prepare solver_parameters dictionary
      solver_parameters = self.get_solver_parameters()
      
      # Construct the solver objects
      problem = NonlinearVariationalProblem(F, solution, bcs=bcs)
      solver = NonlinearVariationalSolver(problem, solver_parameters=solver_parameters)
      LOG.debug("Variational problem solver created.")

      # PETSc solver run-times
      from petsc4py import PETSc
      main_solver_stage = PETSc.Log.Stage('Main block-coupled system solve')

      total_solver_time = 0.0

      # Time-stepping parameters and constants
      T = self.options["T"]
      t = self.options["t"]

      t += dt
      iterations_since_dump = 1
      iterations_since_checkpoint = 1
       
      # The time-stepping loop
      LOG.info("Entering the time-stepping loop...")
      #if annotate: adj_start_timestep(time=t)
      
      EPSILON = 1.0e-14
      while t <= T + EPSILON: # A small value EPSILON is added here in case of round-off error.
         LOG.info("t = %g" % t)

         ## Update any time-dependent Functions and Expressions.
         
         # Re-compute the velocity magnitude and grid Peclet number fields.
         if(self.options["have_su_stabilisation"]):
            magnitude.assign(magnitude_vector(solution_old.split()[0], P1))

            # Bound the values for the magnitude below by 1.0e-9 for numerical stability reasons.
            u_nodes = magnitude.vector()
            near_zero = numpy.array([1.0e-9 for i in range(len(u_nodes))])
            u_nodes.set_local(numpy.maximum(u_nodes.array(), near_zero))
            
            grid_pe.assign(grid_peclet_number(diffusivity, magnitude, P1, cellsize))
      
            # Bound the values for grid_pe below by 1.0e-9 for numerical stability reasons. 
            grid_pe_nodes = grid_pe.vector()
            values = numpy.array([1.0e-9 for i in range(len(grid_pe_nodes))])
            grid_pe_nodes.set_local(numpy.maximum(grid_pe_nodes.array(), values))

         if(self.options["have_turbulence_parameterisation"]):
            les.solve()

         # Time-dependent source terms
         if(self.options["have_momentum_source"]):
            momentum_source_expression.t = t
            momentum_source_function.interpolate(momentum_source_expression)
         if(self.options["have_continuity_source"]):
            continuity_source_expression.t = t
            continuity_source_function.interpolate(continuity_source_expression)

         # Update any time-varying DirichletBC objects.
         for expr in bc_expressions:
            expr.t = t
         for expr in weak_bc_expressions:
            expr.t = t
         
         # Solve the system of equations!
         start_solver_time = mpi4py.MPI.Wtime()
         main_solver_stage.push()
         
         LOG.debug("Solving the system of equations...")
         solver.solve(annotate=annotate)
         main_solver_stage.pop()
         end_solver_time = mpi4py.MPI.Wtime()
         total_solver_time += (end_solver_time - start_solver_time)     

         # Write the solution to file.
         if((self.options["dump_period"] is not None) and (dt*iterations_since_dump >= self.options["dump_period"])):
            LOG.debug("Writing data to file...")
            self.output_functions["Velocity"].assign(solution.split()[0], annotate=False)
            self.output_files["Velocity"] << self.output_functions["Velocity"]
            self.output_functions["FreeSurfacePerturbation"].assign(solution.split()[1], annotate=False)
            self.output_files["FreeSurfacePerturbation"] << self.output_functions["FreeSurfacePerturbation"]
            iterations_since_dump = 0 # Reset the counter.

         # Print out the total power generated by turbines.
         if(array):
            LOG.info("Power = %.2f" % array.power(u, density=1000))
            
         # Checkpointing
         if((self.options["checkpoint_period"] is not None) and (dt*iterations_since_checkpoint >= self.options["checkpoint_period"])):
            LOG.debug("Writing checkpoint data to file...")
            solution.dat.save("checkpoint")
            iterations_since_checkpoint = 0 # Reset the counter.
            
         # Check whether a steady-state has been reached.
         if(steady_state(solution.split()[0], solution_old.split()[0], self.options["steady_state_tolerance"]) and steady_state(solution.split()[1], solution_old.split()[1], self.options["steady_state_tolerance"])):
            LOG.info("Steady-state attained. Exiting the time-stepping loop...")
            break

         self.compute_diagnostics()

         # Move to next time step
         solution_old.assign(solution, annotate=annotate)  
         #array.turbine_drag.assign(project(array.turbine_drag, array.turbine_drag.function_space(), annotate=annotate), annotate=annotate)
         t += dt
         #if annotate: adj_inc_timestep(time=t, finished=t>T)
         
         iterations_since_dump += 1
         iterations_since_checkpoint += 1
         LOG.debug("Moving to next time level...")
      
      LOG.info("Out of the time-stepping loop.")

      LOG.debug("Total solver time: %.2f" % (total_solver_time))

      return solution
    def run(self):
        """ Perform the simulation! """

        # Time-stepping parameters and constants
        LOG.info("Setting up a few constants...")
        T = self.options["T"]
        t = self.options["t"]
        theta = self.options["theta"]
        dt = self.options["dt"]
        dimension = self.options["dimension"]
        g_magnitude = self.options["g_magnitude"]

        # Get the function spaces
        U = self.function_spaces["VelocityFunctionSpace"]
        H = self.function_spaces["FreeSurfaceFunctionSpace"]

        # Is the Velocity field represented by a discontinous function space?
        dg = (U.ufl_element().family() == "Discontinuous Lagrange")

        # Weight u and h by theta to obtain the theta time-stepping scheme.
        assert (theta >= 0.0 and theta <= 1.0)
        LOG.info("Time-stepping scheme using theta = %g" % (theta))
        u_mid = (1.0 - theta) * self.u0 + theta * self.u
        h_mid = (1.0 - theta) * self.h0 + theta * self.h

        # The total height of the free surface.
        self.h_total = self.h_mean + self.h0

        # Non-linear approximation to the velocity
        u_nl = Function(U).assign(self.u0)

        # Second-order Adams-Bashforth velocity
        u_bash = (3.0 / 2.0) * self.u0 - (1.0 / 2.0) * self.u00

        # Simple P1 function space, to be used in the stabilisation routines (if applicable).
        P1 = FunctionSpace(self.mesh, "CG", 1)
        cellsize = CellSize(self.mesh)

        ###########################################################
        ################# Tentative velocity step #################
        ###########################################################

        # The collection of all the individual terms in their weak form.
        LOG.info("Constructing form...")
        F = 0

        # Mass term
        if (self.options["have_momentum_mass"]):
            LOG.debug("Momentum equation: Adding mass term...")
            M_momentum = (1.0 / dt) * (inner(self.w, self.u) -
                                       inner(self.w, self.u0)) * dx
            F += M_momentum

        # Append any Expression objects for weak BCs here.
        weak_bc_expressions = []

        # Advection term
        if (self.options["have_momentum_advection"]):
            LOG.debug("Momentum equation: Adding advection term...")

            if (self.options["integrate_advection_term_by_parts"]):
                outflow = (dot(self.u0, self.n) +
                           abs(dot(self.u0, self.n))) / 2.0

                A_momentum = -inner(dot(u_nl, grad(self.w)),
                                    u_bash) * dx - inner(
                                        dot(u_bash, grad(u_nl)), self.w) * dx
                A_momentum += inner(self.w, outflow * u_mid) * ds
                if (dg):
                    # Only add interior facet integrals if we are dealing with a discontinous Galerkin discretisation.
                    A_momentum += dot(
                        outflow('+') * u_mid('+') - outflow('-') * u_mid('-'),
                        jump(self.w)) * dS

            else:
                A_momentum = inner(dot(grad(self.u), u_nl), self.w) * dx
            F += A_momentum

        # Viscous stress term. Note that the viscosity is kinematic (not dynamic).
        if (self.options["have_momentum_stress"]):
            LOG.debug("Momentum equation: Adding stress term...")

            viscosity = Function(H)

            # Background viscosity
            background_viscosity = Function(H).interpolate(
                Expression(
                    libspud.get_option(
                        "/system/equations/momentum_equation/stress_term/scalar_field::Viscosity/value/constant"
                    )))
            viscosity.assign(background_viscosity)

            # Eddy viscosity
            if (self.options["have_turbulence_parameterisation"]):
                LOG.debug(
                    "Momentum equation: Adding turbulence parameterisation...")
                base_option_path = "/system/equations/momentum_equation/turbulence_parameterisation"
                # Large eddy simulation (LES)
                if (libspud.have_option(base_option_path + "/les")):
                    les = LES(self.mesh, H)
                    density = Constant(
                        1.0
                    )  # We divide through by density in the momentum equation, so just set this to 1.0 for now.
                    smagorinsky_coefficient = Constant(
                        libspud.get_option(
                            base_option_path +
                            "/les/smagorinsky/smagorinsky_coefficient"))

                    eddy_viscosity = Function(H)
                    eddy_viscosity_lhs, eddy_viscosity_rhs = les.eddy_viscosity(
                        u_mid, density, smagorinsky_coefficient)
                    eddy_viscosity_problem = LinearVariationalProblem(
                        eddy_viscosity_lhs,
                        eddy_viscosity_rhs,
                        eddy_viscosity,
                        bcs=[])
                    eddy_viscosity_solver = LinearVariationalSolver(
                        eddy_viscosity_problem)

                # Add on eddy viscosity
                viscosity += eddy_viscosity

            # Stress tensor: tau = grad(u) + transpose(grad(u)) - (2/3)*div(u)
            if (not dg):
                # Perform a double dot product of the stress tensor and grad(w).
                K_momentum = -viscosity * inner(
                    grad(self.u) + grad(self.u).T, grad(self.w)) * dx
                #K_momentum += viscosity*(2.0/3.0)*inner(div(self.u)*Identity(dimension), grad(self.w))*dx
            else:
                # Interior penalty method
                cellsize = Constant(
                    0.2
                )  # In general, we should use CellSize(self.mesh) instead.
                alpha = 1 / cellsize  # Penalty parameter.

                K_momentum = -viscosity('+') * inner(grad(u_mid), grad(
                    self.w)) * dx
                for dim in range(self.options["dimension"]):
                    K_momentum += -viscosity('+') * (
                        alpha('+') / cellsize('+')) * dot(
                            jump(self.w[dim], self.n), jump(
                                u_mid[dim], self.n)) * dS
                    K_momentum += viscosity('+') * dot(
                        avg(grad(self.w[dim])), jump(u_mid[dim], self.n)
                    ) * dS + viscosity('+') * dot(jump(self.w[dim], self.n),
                                                  avg(grad(u_mid[dim]))) * dS

            F -= K_momentum  # Negative sign here because we are bringing the stress term over from the RHS.

        # The gradient of the height of the free surface, h
        LOG.debug("Momentum equation: Adding gradient term...")
        C_momentum = -g_magnitude * inner(self.w, grad(self.h0)) * dx
        F -= C_momentum

        # Quadratic drag term in the momentum equation
        if (self.options["have_drag"]):
            LOG.debug("Momentum equation: Adding drag term...")

            base_option_path = "/system/equations/momentum_equation/drag_term"

            # Get the bottom drag/friction coefficient.
            LOG.debug("Momentum equation: Adding bottom drag contribution...")
            bottom_drag = ExpressionFromOptions(
                path=base_option_path +
                "/scalar_field::BottomDragCoefficient/value",
                t=t).get_expression()
            bottom_drag = Function(H).interpolate(bottom_drag)

            # Add on the turbine drag, if provided.
            self.array = None

            # Magnitude of the velocity field
            magnitude = sqrt(dot(self.u0, self.u0))

            # Form the drag term
            array = sw.get_turbine_array()
            if (array):
                LOG.debug(
                    "Momentum equation: Adding turbine drag contribution...")
                drag_coefficient = bottom_drag + array.turbine_drag()
            else:
                drag_coefficient = bottom_drag
            D_momentum = -inner(
                self.w,
                (drag_coefficient * magnitude / self.h_total) * self.u) * dx
            F -= D_momentum

        # Add in any source terms
        if (self.options["have_momentum_source"]):
            LOG.debug("Momentum equation: Adding source term...")
            momentum_source_expression = ExpressionFromOptions(
                path=
                "/system/equations/momentum_equation/source_term/vector_field::Source/value",
                t=t).get_expression()
            momentum_source_function = Function(U)
            F -= inner(
                self.w,
                momentum_source_function.interpolate(
                    momentum_source_expression)) * dx

        # Add in any SU stabilisation
        if (self.options["have_su_stabilisation"]):
            LOG.debug(
                "Momentum equation: Adding streamline-upwind stabilisation term..."
            )
            stabilisation = Stabilisation(self.mesh, P1, cellsize)

            magnitude = magnitude_vector(self.u0, P1)

            # Bound the values for the magnitude below by 1.0e-9 for numerical stability reasons.
            u_nodes = magnitude.vector()
            near_zero = numpy.array([1.0e-9 for i in range(len(u_nodes))])
            u_nodes.set_local(numpy.maximum(u_nodes.array(), near_zero))

            diffusivity = ExpressionFromOptions(
                path=
                "/system/equations/momentum_equation/stress_term/scalar_field::Viscosity/value",
                t=self.options["t"]).get_expression()
            diffusivity = Function(H).interpolate(
                diffusivity)  # Background viscosity
            grid_pe = grid_peclet_number(diffusivity, magnitude, P1, cellsize)

            # Bound the values for grid_pe below by 1.0e-9 for numerical stability reasons.
            grid_pe_nodes = grid_pe.vector()
            values = numpy.array([1.0e-9 for i in range(len(grid_pe_nodes))])
            grid_pe_nodes.set_local(
                numpy.maximum(grid_pe_nodes.array(), values))

            F += stabilisation.streamline_upwind(self.w, self.u0, magnitude,
                                                 grid_pe)

        LOG.info("Form construction complete.")

        ##########################################################
        ################ Pressure correction step ################
        ##########################################################
        u_tent = Function(U)
        u_tent_nl = theta * u_tent + (1.0 - theta) * self.u0
        F_h_corr = inner(self.v, (self.h - self.h0))*dx \
                   + g_magnitude*(dt**2)*(theta**2)*self.h_total*inner(grad(self.v), grad(self.h - self.h0))*dx \
                   + dt*self.v*div(self.h_total*u_tent_nl)*dx

        ##########################################################
        ################ Velocity correction step ################
        ##########################################################
        h1 = Function(H)
        u1 = Function(U)
        F_u_corr = (1.0 / dt) * inner(
            self.w, self.u - u_tent) * dx + g_magnitude * theta * inner(
                self.w, grad(h1 - self.h0)) * dx

        LOG.info("Applying strong Dirichlet boundary conditions...")
        # Get all the Dirichlet boundary conditions for the Velocity field
        bcs_u = []
        bcs_u2 = []
        bcs_h = []
        bc_expressions = []
        for i in range(
                0,
                libspud.option_count(
                    "/system/core_fields/vector_field::Velocity/boundary_condition"
                )):
            if (libspud.have_option(
                    "/system/core_fields/vector_field::Velocity/boundary_condition[%d]/type::dirichlet"
                    % i
            ) and not libspud.have_option(
                    "/system/core_fields/vector_field::Velocity/boundary_condition[%d]/type::dirichlet/apply_weakly"
                    % i)):
                expr = ExpressionFromOptions(path=(
                    "/system/core_fields/vector_field::Velocity/boundary_condition[%d]/type::dirichlet"
                    % i),
                                             t=t).get_expression()
                # Surface IDs on the domain boundary
                surface_ids = libspud.get_option(
                    "/system/core_fields/vector_field::Velocity/boundary_condition[%d]/surface_ids"
                    % i)
                method = ("geometric" if dg else "topological")
                bc = DirichletBC(U, expr, surface_ids, method=method)
                bcs_u.append(bc)
                bc_expressions.append(expr)
                LOG.debug(
                    "Applying Velocity BC #%d strongly to surface IDs: %s" %
                    (i, surface_ids))

        for i in range(
                0,
                libspud.option_count(
                    "/system/core_fields/vector_field::Velocity/boundary_condition"
                )):
            if (libspud.have_option(
                    "/system/core_fields/vector_field::Velocity/boundary_condition[%d]/type::dirichlet"
                    % i
            ) and not libspud.have_option(
                    "/system/core_fields/vector_field::Velocity/boundary_condition[%d]/type::dirichlet/apply_weakly"
                    % i)):
                expr = ExpressionFromOptions(path=(
                    "/system/core_fields/vector_field::Velocity/boundary_condition[%d]/type::dirichlet"
                    % i),
                                             t=t).get_expression()
                # Surface IDs on the domain boundary
                surface_ids = libspud.get_option(
                    "/system/core_fields/vector_field::Velocity/boundary_condition[%d]/surface_ids"
                    % i)
                method = ("geometric" if dg else "topological")
                bc = DirichletBC(U, expr, surface_ids, method=method)
                bcs_u2.append(bc)
                bc_expressions.append(expr)
                LOG.debug(
                    "Applying Velocity BC #%d strongly to surface IDs: %s" %
                    (i, surface_ids))

        for i in range(
                0,
                libspud.option_count(
                    "/system/core_fields/scalar_field::FreeSurfacePerturbation/boundary_condition/type::dirichlet"
                )):
            if (libspud.have_option(
                    "/system/core_fields/scalar_field::FreeSurfacePerturbation/boundary_condition[%d]/type::dirichlet"
                    % i
            ) and not (libspud.have_option(
                    "/system/core_fields/scalar_field::FreeSurfacePerturbation/boundary_condition[%d]/type::dirichlet/apply_weakly"
                    % i))):
                expr = ExpressionFromOptions(path=(
                    "/system/core_fields/scalar_field::FreeSurfacePerturbation/boundary_condition[%d]/type::dirichlet"
                    % i),
                                             t=t).get_expression()
                # Surface IDs on the domain boundary
                surface_ids = libspud.get_option(
                    "/system/core_fields/scalar_field::FreeSurfacePerturbation/boundary_condition[%d]/surface_ids"
                    % i)
                method = ("geometric" if dg else "topological")
                bc = DirichletBC(H, expr, surface_ids, method=method)
                bcs_h.append(bc)
                bc_expressions.append(expr)
                LOG.debug(
                    "Applying FreeSurfacePerturbation BC #%d strongly to surface IDs: %s"
                    % (i, surface_ids))

        # Prepare solver_parameters dictionary
        LOG.debug("Defining solver_parameters dictionary...")
        solver_parameters = {
            'ksp_monitor': True,
            'ksp_view': False,
            'pc_view': False,
            'snes_type': 'ksponly',
            'ksp_max_it': 10000
        }  # NOTE: use 'snes_type': 'newtonls' for production runs.

        # KSP (iterative solver) options
        solver_parameters["ksp_type"] = libspud.get_option(
            "/system/solver/iterative_method/name")
        solver_parameters["ksp_rtol"] = libspud.get_option(
            "/system/solver/relative_error")

        solver_parameters['ksp_converged_reason'] = True
        solver_parameters['ksp_monitor_true_residual'] = True

        # Preconditioner options
        solver_parameters["pc_type"] = libspud.get_option(
            "/system/solver/preconditioner/name")
        # Fieldsplit sub-options
        if (solver_parameters["pc_type"] == "fieldsplit"):
            LOG.debug("Setting up the fieldsplit preconditioner...")
            solver_parameters["pc_fieldsplit_type"] = libspud.get_option(
                "/system/solver/preconditioner::fieldsplit/type/name")
            if (solver_parameters["pc_fieldsplit_type"] == "schur"):
                solver_parameters[
                    "pc_fieldsplit_schur_fact_type"] = libspud.get_option(
                        "/system/solver/preconditioner::fieldsplit/type::schur/fact_type/name"
                    )
            solver_parameters["fieldsplit_0_ksp_type"] = libspud.get_option(
                "/system/solver/preconditioner::fieldsplit/block_0_ksp_type/iterative_method/name"
            )
            solver_parameters["fieldsplit_1_ksp_type"] = libspud.get_option(
                "/system/solver/preconditioner::fieldsplit/block_1_ksp_type/iterative_method/name"
            )

            if (libspud.get_option(
                    "/system/solver/preconditioner::fieldsplit/block_0_pc_type/preconditioner/name"
            ) != "ilu"):
                solver_parameters["fieldsplit_0_pc_type"] = libspud.get_option(
                    "/system/solver/preconditioner::fieldsplit/block_0_pc_type/preconditioner/name"
                )
                solver_parameters["fieldsplit_1_pc_type"] = libspud.get_option(
                    "/system/solver/preconditioner::fieldsplit/block_1_pc_type/preconditioner/name"
                )

            # Enable inner iteration monitors.
            solver_parameters["fieldsplit_0_ksp_monitor"] = True
            solver_parameters["fieldsplit_1_ksp_monitor"] = True
            solver_parameters["fieldsplit_0_pc_factor_shift_type"] = 'INBLOCKS'
            solver_parameters["fieldsplit_1_pc_factor_shift_type"] = 'INBLOCKS'

        # Construct the solver objects
        problem_tent = LinearVariationalProblem(lhs(F),
                                                rhs(F),
                                                u_tent,
                                                bcs=bcs_u)
        solver_tent = LinearVariationalSolver(problem_tent,
                                              solver_parameters={
                                                  'ksp_monitor': False,
                                                  'ksp_view': False,
                                                  'pc_view': False,
                                                  'pc_type': 'sor',
                                                  'ksp_type': 'gmres',
                                                  'ksp_rtol': 1.0e-7
                                              })

        problem_h_corr = LinearVariationalProblem(lhs(F_h_corr),
                                                  rhs(F_h_corr),
                                                  h1,
                                                  bcs=bcs_h)
        solver_h_corr = LinearVariationalSolver(problem_h_corr,
                                                solver_parameters={
                                                    'ksp_monitor': False,
                                                    'ksp_view': False,
                                                    'pc_view': False,
                                                    'pc_type': 'sor',
                                                    'ksp_type': 'gmres',
                                                    'ksp_rtol': 1.0e-7
                                                })

        problem_u_corr = LinearVariationalProblem(lhs(F_u_corr),
                                                  rhs(F_u_corr),
                                                  u1,
                                                  bcs=bcs_u2)
        solver_u_corr = LinearVariationalSolver(problem_u_corr,
                                                solver_parameters={
                                                    'ksp_monitor': False,
                                                    'ksp_view': False,
                                                    'pc_view': False,
                                                    'pc_type': 'sor',
                                                    'ksp_type': 'gmres',
                                                    'ksp_rtol': 1.0e-7
                                                })

        t += dt
        iterations_since_dump = 1
        iterations_since_checkpoint = 1

        # PETSc solver run-times
        from petsc4py import PETSc
        main_solver_stage = PETSc.Log.Stage('Main block-coupled system solve')

        total_solver_time = 0.0
        # The time-stepping loop
        LOG.info("Entering the time-stepping loop...")
        EPSILON = 1.0e-14

        while t <= T + EPSILON:  # A small value EPSILON is added here in case of round-off error.
            LOG.info("t = %g" % t)

            while True:
                ## Update any time-dependent Functions and Expressions.

                # Re-compute the velocity magnitude and grid Peclet number fields.
                if (self.options["have_su_stabilisation"]):
                    magnitude.assign(magnitude_vector(self.u0, P1))

                    # Bound the values for the magnitude below by 1.0e-9 for numerical stability reasons.
                    u_nodes = magnitude.vector()
                    near_zero = numpy.array(
                        [1.0e-9 for i in range(len(u_nodes))])
                    u_nodes.set_local(numpy.maximum(u_nodes.array(),
                                                    near_zero))

                    grid_pe.assign(
                        grid_peclet_number(diffusivity, magnitude, P1,
                                           cellsize))

                    # Bound the values for grid_pe below by 1.0e-9 for numerical stability reasons.
                    grid_pe_nodes = grid_pe.vector()
                    values = numpy.array(
                        [1.0e-9 for i in range(len(grid_pe_nodes))])
                    grid_pe_nodes.set_local(
                        numpy.maximum(grid_pe_nodes.array(), values))

                if (self.options["have_turbulence_parameterisation"]):
                    eddy_viscosity_solver.solve()
                    viscosity.assign(background_viscosity + eddy_viscosity)

                # Time-dependent source terms
                if (self.options["have_momentum_source"]):
                    momentum_source_expression.t = t
                    momentum_source_function.interpolate(
                        momentum_source_expression)
                if (self.options["have_continuity_source"]):
                    continuity_source_expression.t = t
                    continuity_source_function.interpolate(
                        continuity_source_expression)

                # Update any time-varying DirichletBC objects.
                for expr in bc_expressions:
                    expr.t = t
                for expr in weak_bc_expressions:
                    expr.t = t

                # Solve the system of equations!
                start_solver_time = mpi4py.MPI.Wtime()
                main_solver_stage.push()
                LOG.debug("Solving the system of equations...")

                solver_tent.solve()
                solver_h_corr.solve()
                solver_u_corr.solve()

                main_solver_stage.pop()
                end_solver_time = mpi4py.MPI.Wtime()
                total_solver_time += (end_solver_time - start_solver_time)

                # Move to next time step
                if (steady_state(u1, u_nl, 1e-7)):
                    break
                u_nl.assign(u1)

            self.u00.assign(self.u0)
            self.u0.assign(u1)
            self.h0.assign(h1)
            t += dt
            iterations_since_dump += 1
            iterations_since_checkpoint += 1
            LOG.debug("Moving to next time level...")

            # Write the solution to file.
            if ((self.options["dump_period"] is not None) and
                (dt * iterations_since_dump >= self.options["dump_period"])):
                LOG.debug("Writing data to file...")
                self.output_functions["Velocity"].assign(u1)
                self.output_files["Velocity"] << self.output_functions[
                    "Velocity"]
                self.output_functions["FreeSurfacePerturbation"].assign(h1)
                self.output_files[
                    "FreeSurfacePerturbation"] << self.output_functions[
                        "FreeSurfacePerturbation"]
                iterations_since_dump = 0  # Reset the counter.

            # Print out the total power generated by turbines.
            if (self.options["have_drag"] and self.array is not None):
                LOG.info("Power = %.2f" % self.array.power(u1, density=1000))

            # Checkpointing
            if ((self.options["checkpoint_period"] is not None)
                    and (dt * iterations_since_checkpoint >=
                         self.options["checkpoint_period"])):
                LOG.debug("Writing checkpoint data to file...")
                self.solution.dat.save("checkpoint")
                iterations_since_checkpoint = 0  # Reset the counter.

            # Check whether a steady-state has been reached.
            if (steady_state(u1, self.u0,
                             self.options["steady_state_tolerance"])
                    and steady_state(h1, self.h0,
                                     self.options["steady_state_tolerance"])):
                LOG.info(
                    "Steady-state attained. Exiting the time-stepping loop...")
                break

            self.compute_diagnostics()

        LOG.info("Out of the time-stepping loop.")

        LOG.debug("Total solver time: %.2f" % (total_solver_time))

        return u1, h1