def heat(n, deg, time_stages, stage_type="deriv", splitting=IA): N = 2**n msh = UnitIntervalMesh(N) params = { "snes_type": "ksponly", "ksp_type": "preonly", "mat_type": "aij", "pc_type": "lu" } V = FunctionSpace(msh, "CG", deg) x, = SpatialCoordinate(msh) t = Constant(0.0) dt = Constant(2.0 / N) uexact = exp(-t) * sin(pi * x) rhs = expand_derivatives(diff(uexact, t)) - div(grad(uexact)) butcher_tableau = GaussLegendre(time_stages) u = project(uexact, V) v = TestFunction(V) F = (inner(Dt(u), v) * dx + inner(grad(u), grad(v)) * dx - inner(rhs, v) * dx) bc = DirichletBC(V, Constant(0), "on_boundary") stepper = TimeStepper(F, butcher_tableau, t, dt, u, bcs=bc, solver_parameters=params, stage_type=stage_type, splitting=splitting) while (float(t) < 1.0): if (float(t) + float(dt) > 1.0): dt.assign(1.0 - float(t)) stepper.advance() t.assign(float(t) + float(dt)) return errornorm(uexact, u) / norm(uexact)
def wave(n, deg, butcher_tableau, splitting=AI): N = 2**n msh = UnitIntervalMesh(N) params = { "snes_type": "ksponly", "ksp_type": "preonly", "mat_type": "aij", "pc_type": "lu" } V = FunctionSpace(msh, "CG", deg) W = FunctionSpace(msh, "DG", deg - 1) Z = V * W x, = SpatialCoordinate(msh) t = Constant(0.0) dt = Constant(2.0 / N) up = project(as_vector([0, sin(pi * x)]), Z) u, p = split(up) v, w = TestFunctions(Z) F = (inner(Dt(u), v) * dx + inner(u.dx(0), w) * dx + inner(Dt(p), w) * dx - inner(p, v.dx(0)) * dx) E = 0.5 * (inner(u, u) * dx + inner(p, p) * dx) stepper = TimeStepper(F, butcher_tableau, t, dt, up, solver_parameters=params, splitting=splitting) energies = [] while (float(t) < 1.0): if (float(t) + float(dt) > 1.0): dt.assign(1.0 - float(t)) stepper.advance() t.assign(float(t) + float(dt)) energies.append(assemble(E)) return np.array(energies)
class Forcing(object, metaclass=ABCMeta): """ Base class for forcing terms for Gusto. :arg state: x :class:`.State` object. :arg euler_poincare: if True then the momentum equation is in Euler Poincare form and we need to add 0.5*grad(u^2) to the forcing term. If False then this term is not added. :arg linear: if True then we are solving a linear equation so nonlinear terms (namely the Euler Poincare term) should not be added. :arg extra_terms: extra terms to add to the u component of the forcing term - these will be multiplied by the appropriate test function. """ def __init__(self, state, euler_poincare=True, linear=False, extra_terms=None, moisture=None): self.state = state if linear: self.euler_poincare = False logger.warning( 'Setting euler_poincare to False because you have set linear=True' ) else: self.euler_poincare = euler_poincare # set up functions self.Vu = state.spaces("HDiv") # this is the function that the forcing term is applied to self.x0 = Function(state.W) self.test = TestFunction(self.Vu) self.trial = TrialFunction(self.Vu) # this is the function that contains the result of solving # <test, trial> = <test, F(x0)>, where F is the forcing term self.uF = Function(self.Vu) # find out which terms we need self.extruded = self.Vu.extruded self.coriolis = state.Omega is not None or hasattr( state.fields, "coriolis") self.sponge = state.mu is not None self.hydrostatic = state.hydrostatic self.topography = hasattr(state.fields, "topography") self.extra_terms = extra_terms self.moisture = moisture # some constants to use for scaling terms self.scaling = Constant(1.) self.impl = Constant(1.) self._build_forcing_solvers() def mass_term(self): return inner(self.test, self.trial) * dx def coriolis_term(self): u0 = split(self.x0)[0] return -inner(self.test, cross(2 * self.state.Omega, u0)) * dx def sponge_term(self): u0 = split(self.x0)[0] return self.state.mu * inner(self.test, self.state.k) * inner( u0, self.state.k) * dx def euler_poincare_term(self): u0 = split(self.x0)[0] return -0.5 * div(self.test) * inner(self.state.h_project(u0), u0) * dx def hydrostatic_term(self): u0 = split(self.x0)[0] return inner(u0, self.state.k) * inner(self.test, self.state.k) * dx @abstractmethod def pressure_gradient_term(self): pass def forcing_term(self): L = self.pressure_gradient_term() if self.extruded: L += self.gravity_term() if self.coriolis: L += self.coriolis_term() if self.euler_poincare: L += self.euler_poincare_term() if self.topography: L += self.topography_term() if self.extra_terms is not None: L += inner(self.test, self.extra_terms) * dx # scale L L = self.scaling * L # sponge term has a separate scaling factor as it is always implicit if self.sponge: L -= self.impl * self.state.timestepping.dt * self.sponge_term() # hydrostatic term has no scaling factor if self.hydrostatic: L += (2 * self.impl - 1) * self.hydrostatic_term() return L def _build_forcing_solvers(self): a = self.mass_term() L = self.forcing_term() bcs = None if len(self.state.bcs) == 0 else self.state.bcs u_forcing_problem = LinearVariationalProblem(a, L, self.uF, bcs=bcs) solver_parameters = {} if logger.isEnabledFor(DEBUG): solver_parameters["ksp_monitor_true_residual"] = None self.u_forcing_solver = LinearVariationalSolver( u_forcing_problem, solver_parameters=solver_parameters, options_prefix="UForcingSolver") def apply(self, scaling, x_in, x_nl, x_out, **kwargs): """ Function takes x as input, computes F(x_nl) and returns x_out = x + scale*F(x_nl) as output. :arg x_in: :class:`.Function` object :arg x_nl: :class:`.Function` object :arg x_out: :class:`.Function` object :arg implicit: forcing stage for sponge and hydrostatic terms, if present """ self.scaling.assign(scaling) self.x0.assign(x_nl) implicit = kwargs.get("implicit") if implicit is not None: self.impl.assign(int(implicit)) self.u_forcing_solver.solve() # places forcing in self.uF uF = x_out.split()[0] x_out.assign(x_in) uF += self.uF
class Diffusion2D(Application): """ Application class containing the description of the diffusion problem. The spatial domain is a 10x10 square with periodic boundary conditions in each direction. The initial condition is a Gaussian in the centre of the domain. The spatial discretisation is P1 DG (piecewise linear discontinous elements) and uses an interior penalty method which penalises jumps at element interfaces. """ def __init__(self, mesh: object, kappa: float, comm_space: MPI.Comm, mu: float = 5., *args, **kwargs): """ Constructor :param mesh: spatial domain :param kappa: diffusion coefficient :param mu: penalty weighting function """ super(Diffusion2D, self).__init__(*args, **kwargs) # Spatial domain and function space self.mesh = mesh V = FunctionSpace(self.mesh, "DG", 1) self.function_space = V self.comm_space = comm_space # Placeholder for time step - will be updated in the update method self.dt = Constant(0.) # Things we need for the form gamma = TestFunction(V) phi = TrialFunction(V) self.f = Function(V) n = FacetNormal(mesh) # Set up the rhs and bilinear form of the equation a = (inner(gamma, phi) * dx + self.dt * (inner(grad(gamma), grad(phi) * kappa) * dx - inner(2 * avg(outer(phi, n)), avg(grad(gamma) * kappa)) * dS - inner(avg(grad(phi) * kappa), 2 * avg(outer(gamma, n))) * dS + mu * inner(2 * avg(outer(phi, n)), 2 * avg(outer(gamma, n) * kappa)) * dS)) rhs = inner(gamma, self.f) * dx # Function to hold the solution self.soln = Function(V) # Setup problem and solver prob = LinearVariationalProblem(a, rhs, self.soln) self.solver = NonlinearVariationalSolver(prob) # Set the data structure for any user-defined time point self.vector_template = VectorDiffusion2D(size=len(self.function_space), comm_space=self.comm_space) # Set initial condition: # Setting up a Gaussian blob in the centre of the domain. self.vector_t_start = VectorDiffusion2D(size=len(self.function_space), comm_space=self.comm_space) x = SpatialCoordinate(self.mesh) initial_tracer = exp(-((x[0] - 5)**2 + (x[1] - 5)**2)) tmp = Function(self.function_space) tmp.interpolate(initial_tracer) self.vector_t_start.set_values(np.copy(tmp.dat.data)) def step(self, u_start: VectorDiffusion2D, t_start: float, t_stop: float) -> VectorDiffusion2D: """ Time integration routine for 2D diffusion problem: Backward Euler :param u_start: approximate solution for the input time t_start :param t_start: time associated with the input approximate solution u_start :param t_stop: time to evolve the input approximate solution to :return: approximate solution at input time t_stop """ # Time-step size self.dt.assign(t_stop - t_start) # Get data from VectorDiffusion2D object u_start # and copy to Firedrake Function object tmp tmp = Function(self.function_space) for i in range(len(u_start.values)): tmp.dat.data[i] = u_start.values[i] self.f.assign(tmp) # Take Backward Euler step self.solver.solve() # Copy data from Firedrake Function object to VectorDiffusion2D object ret = VectorDiffusion2D(size=len(self.function_space), comm_space=self.comm_space) ret.set_values(np.copy(self.soln.dat.data)) return ret
class CompressibleForcing(Forcing): """ Forcing class for compressible Euler equations. """ def __init__(self, state, linear=False): self.state = state self._build_forcing_solver(linear) def _build_forcing_solver(self, linear): """ Only put forcing terms into the u equation. """ state = self.state self.scaling = Constant(1.0) Vu = state.V[0] W = state.W self.x0 = Function(W) # copy x to here u0, rho0, theta0 = split(self.x0) F = TrialFunction(Vu) w = TestFunction(Vu) self.uF = Function(Vu) Omega = state.Omega cp = state.parameters.cp mu = state.mu n = FacetNormal(state.mesh) pi = exner(theta0, rho0, state) a = inner(w, F) * dx L = self.scaling * ( +cp * div(theta0 * w) * pi * dx # pressure gradient [volume] - cp * jump(w * theta0, n) * avg(pi) * dS_v # pressure gradient [surface] ) if state.parameters.geopotential: Phi = state.Phi L += self.scaling * div(w) * Phi * dx # gravity term else: g = state.parameters.g L -= self.scaling * g * inner(w, state.k) * dx # gravity term if not linear: L -= self.scaling * 0.5 * div(w) * inner(u0, u0) * dx if Omega is not None: L -= self.scaling * inner(w, cross(2 * Omega, u0)) * dx # Coriolis term if mu is not None: self.mu_scaling = Constant(1.0) L -= self.mu_scaling * mu * inner(w, state.k) * inner(u0, state.k) * dx bcs = [DirichletBC(Vu, 0.0, "bottom"), DirichletBC(Vu, 0.0, "top")] u_forcing_problem = LinearVariationalProblem(a, L, self.uF, bcs=bcs) self.u_forcing_solver = LinearVariationalSolver(u_forcing_problem) def apply(self, scaling, x_in, x_nl, x_out, **kwargs): self.x0.assign(x_nl) self.scaling.assign(scaling) if "mu_alpha" in kwargs and kwargs["mu_alpha"] is not None: self.mu_scaling.assign(kwargs["mu_alpha"]) self.u_forcing_solver.solve() # places forcing in self.uF u_out, _, _ = x_out.split() x_out.assign(x_in) u_out += self.uF
class IncompressibleForcing(Forcing): """ Forcing class for incompressible Euler Boussinesq equations. """ def __init__(self, state, linear=False): self.state = state self._build_forcing_solver(linear) def _build_forcing_solver(self, linear): """ Only put forcing terms into the u equation. """ state = self.state self.scaling = Constant(1.0) Vu = state.V[0] W = state.W self.x0 = Function(W) # copy x to here u0, p0, b0 = split(self.x0) F = TrialFunction(Vu) w = TestFunction(Vu) self.uF = Function(Vu) Omega = state.Omega mu = state.mu a = inner(w, F) * dx L = ( self.scaling * div(w) * p0 * dx # pressure gradient + self.scaling * b0 * inner(w, state.k) * dx # gravity term ) if not linear: L -= self.scaling * 0.5 * div(w) * inner(u0, u0) * dx if Omega is not None: L -= self.scaling * inner(w, cross(2 * Omega, u0)) * dx # Coriolis term if mu is not None: self.mu_scaling = Constant(1.0) L -= self.mu_scaling * mu * inner(w, state.k) * inner(u0, state.k) * dx bcs = [DirichletBC(Vu, 0.0, "bottom"), DirichletBC(Vu, 0.0, "top")] u_forcing_problem = LinearVariationalProblem(a, L, self.uF, bcs=bcs) self.u_forcing_solver = LinearVariationalSolver(u_forcing_problem) Vp = state.V[1] p = TrialFunction(Vp) q = TestFunction(Vp) self.divu = Function(Vp) a = p * q * dx L = q * div(u0) * dx divergence_problem = LinearVariationalProblem(a, L, self.divu) self.divergence_solver = LinearVariationalSolver(divergence_problem) def apply(self, scaling, x_in, x_nl, x_out, **kwargs): self.x0.assign(x_nl) self.scaling.assign(scaling) if "mu_alpha" in kwargs and kwargs["mu_alpha"] is not None: self.mu_scaling.assign(kwargs["mu_alpha"]) self.u_forcing_solver.solve() # places forcing in self.uF u_out, p_out, _ = x_out.split() x_out.assign(x_in) u_out += self.uF if "incompressible" in kwargs and kwargs["incompressible"]: self.divergence_solver.solve() p_out.assign(self.divu)
def solver_CG(mesh, el, space, deg, T, dt=0.001, warm_up=False): """Solve the scalar wave equation on a unit square/cube using a CG FEM formulation with several different element types. Parameters ---------- mesh: Firedrake.mesh A utility mesh from the Firedrake package el: string The type of element either "tria" or "quad". `tria` in 3d implies tetrahedra and `quad` in 3d implies hexahedral elements. space: string The space of the FEM. Available options are: "CG": Continuous Galerkin Finite Elements, "KMV": Kong-Mulder-Veldhuzien higher-order mass lumped elements "S" (for Serendipity) (NB: quad/hexs only) "spectral": spectral elements using GLL quad points (NB: quads/hexs only) deg: int The spatial polynomial degree. T: float The simulation duration in simulation seconds. dt: float, optional Simulation timestep warm_up: boolean, optional Warm up symbolics by running one timestep and shutting down. Returns ------- u_n: Firedrake.Function The solution at time `T` """ sd = mesh.geometric_dimension() V = _build_space(mesh, el, space, deg) quad_rule1, quad_rule2 = _build_quad_rule(el, V, space) params = _select_params(space) # DEBUG # outfile = fd.File(os.getcwd() + "/results/simple_shots.pvd") # END DEBUG tot_dof = COMM_WORLD.allreduce(V.dof_dset.total_size, op=MPI.SUM) # if COMM_WORLD.rank == 0: # print("------------------------------------------") # print("The problem has " + str(tot_dof) + " degrees of freedom.") # print("------------------------------------------") nt = int(T / dt) # number of timesteps u = fd.TrialFunction(V) v = fd.TestFunction(V) u_np1 = fd.Function(V) # n+1 u_n = fd.Function(V) # n u_nm1 = fd.Function(V) # n-1 # constant speed c = Constant(1.5) m = ( (1.0 / (c * c)) * (u - 2.0 * u_n + u_nm1) / Constant(dt * dt) * v * dx(rule=quad_rule1) ) # mass-like matrix a = dot(grad(u_n), grad(v)) * dx(rule=quad_rule2) # stiffness matrix # injection of source into mesh ricker = Constant(0.0) source = Constant([0.5] * sd) coords = fd.SpatialCoordinate(mesh) F = m + a - delta_expr(source, *coords) * ricker * v * dx(rule=quad_rule2) a, r = fd.lhs(F), fd.rhs(F) A, R = fd.assemble(a), fd.assemble(r) solver = fd.LinearSolver(A, solver_parameters=params, options_prefix="") # timestepping loop results = [] t = 0.0 for step in range(nt): with PETSc.Log.Stage("{el}{deg}".format(el=el, deg=deg)): ricker.assign(RickerWavelet(t, freq=6)) R = fd.assemble(r, tensor=R) solver.solve(u_np1, R) snes = _get_time("SNESSolve") ksp = _get_time("KSPSolve") pcsetup = _get_time("PCSetUp") pcapply = _get_time("PCApply") jac = _get_time("SNESJacobianEval") residual = _get_time("SNESFunctionEval") sparsity = _get_time("CreateSparsity") results.append( [tot_dof, snes, ksp, pcsetup, pcapply, jac, residual, sparsity] ) if warm_up: # Warm up symbolics/disk cache solver.solve(u_np1, R) sys.exit("Warming up...") u_nm1.assign(u_n) u_n.assign(u_np1) t = step * float(dt) # if step % 10 == 0: # outfile.write(u_n) # print("Time is " + str(t), flush=True) results = np.asarray(results) if mesh.comm.rank == 0: with open( "data/scalar_wave.{el}.{deg}.{space}.csv".format( el=el, deg=deg, space=space ), "w", ) as f: np.savetxt( f, results, fmt=["%d"] + ["%e"] * 7, delimiter=",", header="tot_dof,SNESSolve,KSPSolve,PCSetUp,PCApply,SNESJacobianEval,SNESFunctionEval,CreateSparsity", comments="", ) return u_n
class DiagnosticEquations(object): """ An object setting up diagnostic equations and solvers for the stochastic Camassa-Holm equation. :arg diagnostic_variables: a DiagnosticVariables object. :arg prognostic_variables: a PrognosticVariables object. :arg outputting: an Outputting object. :arg simulation_parameters: a dictionary storing the simulation parameters. """ def __init__(self, diagnostic_variables, prognostic_variables, outputting, simulation_parameters): self.diagnostic_variables = diagnostic_variables self.prognostic_variables = prognostic_variables self.outputting = outputting self.simulation_parameters = simulation_parameters Dt = Constant(simulation_parameters['dt'][-1]) Ld = simulation_parameters['Ld'][-1] u = self.prognostic_variables.u Xi = self.prognostic_variables.dXi Vu = u.function_space() vector_u = True if Vu.ufl_element() == VectorElement else False ones = Function( VectorFunctionSpace(self.prognostic_variables.mesh, "CG", 1)).project(as_vector([Constant(1.0)])) self.to_update_constants = False self.interpolators = [] self.projectors = [] self.solvers = [] mesh = u.function_space().mesh() x, = SpatialCoordinate(mesh) alphasq = simulation_parameters['alphasq'][-1] periodic = simulation_parameters['periodic'][-1] # do peakon data checks here true_peakon_data = simulation_parameters['true_peakon_data'][-1] if true_peakon_data is not None: self.true_peakon_file = Dataset( 'results/' + true_peakon_data + '/data.nc', 'r') # check length of file is correct ndump = simulation_parameters['ndump'][-1] tmax = simulation_parameters['tmax'][-1] dt = simulation_parameters['dt'][-1] if len(self.true_peakon_file['time'][:]) != int(tmax / (ndump * dt)) + 1: raise ValueError( 'If reading in true peakon data, the dump frequency must be the same as that used for the true peakon data.' + ' Length of true peakon data as %i, but proposed length is %i' % (len(self.true_peakon_file['time'][:]), int(tmax / (ndump * dt)) + 1)) if self.true_peakon_file['p'][:].shape != (int(tmax / (ndump * dt)) + 1, ): raise ValueError( 'True peakon data shape %i must be the same shape as proposed data %i' % ((int(tmax / (ndump * dt)) + 1, ), self.true_peakon_file['p'][:].shape)) # do peakon data checks here true_mean_peakon_data = simulation_parameters['true_mean_peakon_data'][ -1] if true_mean_peakon_data is not None: self.true_mean_peakon_file = Dataset( 'results/' + true_mean_peakon_data + '/data.nc', 'r') # check length of file is correct ndump = simulation_parameters['ndump'][-1] tmax = simulation_parameters['tmax'][-1] dt = simulation_parameters['dt'][-1] if len(self.true_mean_peakon_file['time'][:]) != int(tmax / (ndump * dt)): raise ValueError( 'If reading in true peakon data, the dump frequency must be the same as that used for the true peakon data.' ) if self.true_mean_peakon_file['p'][:].shape != (int( tmax / (ndump * dt)), ): raise ValueError( 'True peakon data must have same shape as proposed data!') for key, value in self.diagnostic_variables.fields.items(): if key == 'uscalar': uscalar = self.diagnostic_variables.fields['uscalar'] u_interpolator = Interpolator(dot(ones, u), uscalar) self.interpolators.append(u_interpolator) elif key == 'Euscalar': Eu = self.prognostic_variables.Eu Euscalar = self.diagnostic_variables.fields['Euscalar'] Eu_interpolator = Interpolator(dot(ones, Eu), Euscalar) self.interpolators.append(Eu_interpolator) elif key == 'Xiscalar': Xi = self.prognostic_variables.dXi Xiscalar = self.diagnostic_variables.fields['Xiscalar'] Xi_interpolator = Interpolator(dot(ones, Xi), Xiscalar) self.interpolators.append(Xi_interpolator) elif key == 'du': if type(u.function_space().ufl_element()) == VectorElement: u_to_project = self.diagnostic_variables.fields['uscalar'] else: u_to_project = u du = self.diagnostic_variables.fields['du'] du_projector = Projector(u_to_project.dx(0), du) self.projectors.append(du_projector) elif key == 'jump_du': du = self.diagnostic_variables.fields['du'] jump_du = self.diagnostic_variables.fields['jump_du'] V = jump_du.function_space() jtrial = TrialFunction(V) psi = TestFunction(V) Lj = psi('+') * abs(jump(du)) * dS aj = psi('+') * jtrial('+') * dS jprob = LinearVariationalProblem(aj, Lj, jump_du) jsolver = LinearVariationalSolver(jprob) self.solvers.append(jsolver) elif key == 'du_smooth': du = self.diagnostic_variables.fields['du'] du_smooth = self.diagnostic_variables.fields['du_smooth'] projector = Projector(du, du_smooth) self.projectors.append(projector) elif key == 'u2_flux': gamma = simulation_parameters['gamma'][-1] u2_flux = self.diagnostic_variables.fields['u2_flux'] xis = self.prognostic_variables.pure_xi_list xis_x = [] xis_xxx = [] CG1 = FunctionSpace(mesh, "CG", 1) psi = TestFunction(CG1) for xi in xis: xis_x.append(Function(CG1).project(xi.dx(0))) for xi_x in xis_x: xi_xxx = Function(CG1) form = (psi * xi_xxx + psi.dx(0) * xi_x.dx(0)) * dx prob = NonlinearVariationalProblem(form, xi_xxx) solver = NonlinearVariationalSolver(prob) solver.solve() xis_xxx.append(xi_xxx) flux_expr = 0.0 * x for xi, xi_x, xi_xxx in zip(xis, xis_x, xis_xxx): flux_expr += (6 * u.dx(0) * xi + 12 * u * xi_x + gamma * xi_xxx) * (6 * u.dx(0) * xi + 24 * u * xi_x + gamma * xi_xxx) projector = Projector(flux_expr, u2_flux) self.projectors.append(projector) elif key == 'a': # find 6 * u_x * Xi + gamma * Xi_xxx mesh = u.function_space().mesh() gamma = simulation_parameters['gamma'][-1] a_flux = self.diagnostic_variables.fields['a'] xis = self.prognostic_variables.pure_xis xis_x = [] xis_xxx = [] CG1 = FunctionSpace(mesh, "CG", 1) psi = TestFunction(CG1) for xi in xis: xis_x.append(Function(CG1).project(xi.dx(0))) for xi_x in xis_x: xi_xxx = Function(CG1) form = (psi * xi_xxx + psi.dx(0) * xi_x.dx(0)) * dx prob = NonlinearVariationalProblem(form, xi_xxx) solver = NonlinearVariationalSolver(prob) solver.solve() xis_xxx.append(xi_xxx) x, = SpatialCoordinate(mesh) a_expr = 0.0 * x for xi, xi_x, xi_xxx in zip(xis, xis_x, xis_xxx): a_expr += 6 * u.dx(0) * xi + gamma * xi_xxx projector = Projector(a_expr, a_flux) self.projectors.append(projector) elif key == 'b': # find 12 * u * Xi_x mesh = u.function_space().mesh() gamma = simulation_parameters['gamma'][-1] b_flux = self.diagnostic_variables.fields['b'] xis = self.prognostic_variables.pure_xis x, = SpatialCoordinate(mesh) b_expr = 0.0 * x for xi, xi_x, xi_xxx in zip(xis, xis_x, xis_xxx): b_expr += 12 * u * xi.dx(0) projector = Projector(b_expr, b_flux) self.projectors.append(projector) elif key == 'kdv_1': # find the first part of the kdv form u0 = prognostic_variables.u0 uh = (u + u0) / 2 us = Dt * uh + sqrt(Dt) * Xi psi = TestFunction(Vu) du_1 = self.diagnostic_variables.fields['kdv_1'] eqn = psi * du_1 * dx - 6 * psi.dx(0) * uh * us * dx prob = NonlinearVariationalProblem(eqn, du_1) solver = NonlinearVariationalSolver(prob) self.solvers.append(solver) elif key == 'kdv_2': # find the second part of the kdv form u0 = prognostic_variables.u0 uh = (u + u0) / 2 us = Dt * uh + sqrt(Dt) * Xi psi = TestFunction(Vu) du_2 = self.diagnostic_variables.fields['kdv_2'] eqn = psi * du_2 * dx + 6 * psi * uh * us.dx(0) * dx prob = NonlinearVariationalProblem(eqn, du_2) solver = NonlinearVariationalSolver(prob) self.solvers.append(solver) elif key == 'kdv_3': # find the third part of the kdv form u0 = prognostic_variables.u0 uh = (u + u0) / 2 us = Dt * uh + sqrt(Dt) * Xi du_3 = self.diagnostic_variables.fields['kdv_3'] gamma = simulation_parameters['gamma'][-1] phi = TestFunction(Vu) F = Function(Vu) eqn = (phi * F * dx + phi.dx(0) * us.dx(0) * dx) prob = NonlinearVariationalProblem(eqn, F) solver = NonlinearVariationalSolver(prob) self.solvers.append(solver) self.projectors.append(Projector(-gamma * F.dx(0), du_3)) # nu = TestFunction(Vu) # back_eqn = nu * du_3 * dx - gamma * nu.dx(0) * F * dx # back_prob = NonlinearVariationalProblem(back_eqn, du_3) # back_solver = NonlinearVariationalSolver(back_prob) # self.solvers.append(solver) elif key == 'm': m = self.diagnostic_variables.fields['m'] phi = TestFunction(Vu) eqn = phi * m * dx - phi * u * dx - alphasq * phi.dx(0) * u.dx( 0) * dx prob = NonlinearVariationalProblem(eqn, m) solver = NonlinearVariationalSolver(prob) self.solvers.append(solver) elif key == 'u_xx': u_xx = self.diagnostic_variables.fields['u_xx'] phi = TestFunction(Vu) eqn = phi * u_xx * dx + phi.dx(0) * u_xx.dx(0) * dx prob = NonlinearVariationalProblem(eqn, u_xx) solver = NonlinearVariationalSolver(prob) self.solvers.append(solver) elif key == 'u_sde': self.to_update_constants = True self.Ld = Ld self.alphasq = alphasq self.p = Constant(1.0 * 0.5 * (1 + exp(-Ld / sqrt(alphasq))) / (1 - exp(-Ld / sqrt(alphasq)))) self.q = Constant(Ld / 2) u_sde = self.diagnostic_variables.fields['u_sde'] if periodic: expr = conditional( x < self.q - Ld / 2, self.p * ((exp(-(x - self.q + Ld) / sqrt(alphasq)) + exp(-Ld / sqrt(alphasq)) * exp( (x - self.q + Ld) / sqrt(alphasq))) / (1 - exp(-Ld / sqrt(alphasq)))), conditional( x < self.q + Ld / 2, self.p * ((exp(-sqrt((self.q - x)**2 / alphasq)) + exp(-Ld / sqrt(alphasq)) * exp(sqrt((self.q - x)**2 / alphasq))) / (1 - exp(-Ld / sqrt(alphasq)))), self.p * ((exp(-(self.q + Ld - x) / sqrt(alphasq)) + exp(-Ld / sqrt(alphasq) * exp( (self.q + Ld - x) / sqrt(alphasq)))) / (1 - exp(-Ld / sqrt(alphasq)))))) else: expr = conditional( x < self.q - Ld / 2, self.p * exp(-(x - self.q + Ld) / sqrt(alphasq)), conditional( x < self.q + Ld / 2, self.p * exp(-sqrt((self.q - x)**2 / alphasq)), self.p * exp(-(self.q + Ld - x) / sqrt(alphasq)))) self.interpolators.append(Interpolator(expr, u_sde)) elif key == 'u_sde_weak': u_sde = self.diagnostic_variables.fields['u_sde'] u_sde_weak = self.diagnostic_variables.fields['u_sde_weak'] psi = TestFunction(Vu) eqn = psi * u_sde_weak * dx - psi * (u - u_sde) * dx prob = NonlinearVariationalProblem(eqn, u_sde_weak) solver = NonlinearVariationalSolver(prob) self.solvers.append(solver) elif key == 'u_sde_mean': self.to_update_constants = True self.p = Constant(1.0) self.q = Constant(Ld / 2) if periodic: raise NotImplementedError( 'u_sde_mean not yet implemented for periodic peakon') u_sde = self.diagnostic_variables.fields['u_sde_mean'] expr = conditional( x < self.q - Ld / 2, self.p * exp(-(x - self.q + Ld) / sqrt(alphasq)), conditional( x < self.q + Ld / 2, self.p * exp(-sqrt((self.q - x)**2 / alphasq)), self.p * exp(-(self.q + Ld - x) / sqrt(alphasq)))) self.interpolators.append(Interpolator(expr, u_sde)) elif key == 'u_sde_weak_mean': u_sde = self.diagnostic_variables.fields['u_sde_mean'] u_sde_weak = self.diagnostic_variables.fields[ 'u_sde_weak_mean'] psi = TestFunction(Vu) eqn = psi * u_sde_weak * dx - psi * (u - u_sde) * dx prob = NonlinearVariationalProblem(eqn, u_sde_weak) solver = NonlinearVariationalSolver(prob) self.solvers.append(solver) elif key == 'pure_xi': pure_xi = 0.0 * x for xi in self.prognostic_variables.pure_xi_list: if vector_u: pure_xi += dot(ones, xi) else: pure_xi += xi Xiscalar = self.diagnostic_variables.fields['pure_xi'] Xi_interpolator = Interpolator(pure_xi, Xiscalar) self.interpolators.append(Xi_interpolator) elif key == 'pure_xi_x': pure_xi_x = 0.0 * x for xix in self.prognostic_variables.pure_xi_x_list: if vector_u: pure_xi_x += dot(ones, xix) else: pure_xi_x += xix Xiscalar = self.diagnostic_variables.fields['pure_xi_x'] Xi_interpolator = Interpolator(pure_xi_x, Xiscalar) self.interpolators.append(Xi_interpolator) elif key == 'pure_xi_xx': pure_xi_xx = 0.0 * x for xixx in self.prognostic_variables.pure_xi_xx_list: if vector_u: pure_xi_xx += dot(ones, xixx) else: pure_xi_xx += xixx Xiscalar = self.diagnostic_variables.fields['pure_xi_xx'] Xi_interpolator = Interpolator(pure_xi_xx, Xiscalar) self.interpolators.append(Xi_interpolator) elif key == 'pure_xi_xxx': pure_xi_xxx = 0.0 * x for xixxx in self.prognostic_variables.pure_xi_xxx_list: if vector_u: pure_xi_xxx += dot(ones, xixxx) else: pure_xi_xxx += xixxx Xiscalar = self.diagnostic_variables.fields['pure_xi_xxx'] Xi_interpolator = Interpolator(pure_xi_xxx, Xiscalar) self.interpolators.append(Xi_interpolator) elif key == 'pure_xi_xxxx': pure_xi_xxxx = 0.0 * x for xixxxx in self.prognostic_variables.pure_xi_xx_list: if vector_u: pure_xi_xxxx += dot(ones, xixxxx) else: pure_xi_xxxx += xixxxx Xiscalar = self.diagnostic_variables.fields['pure_xi_xxxx'] Xi_interpolator = Interpolator(pure_xi_xxxx, Xiscalar) self.interpolators.append(Xi_interpolator) else: raise NotImplementedError('Diagnostic %s not yet implemented' % key) def update_constants(self): """ Update p and q constants from true peakon data. """ self.p.assign(self.true_peakon_file['p'][self.outputting.t_idx] * 0.5 * (1 + exp(-self.Ld / sqrt(self.alphasq))) / (1 - exp(-self.Ld / sqrt(self.alphasq)))) self.q.assign(self.true_peakon_file['q'][self.outputting.t_idx]) def solve(self): """ Do interpolations, projections and solves. """ if self.to_update_constants: self.update_constants() for interpolator in self.interpolators: interpolator.interpolate() for projector in self.projectors: projector.project() for solver in self.solvers: solver.solve()
class TimeMixin: """ Class to extend Problems as a mixin. This mixin extends the problem to allow time stepping by adding the dT/dt term. To use, define a new class with this in the inheritance chain. i.e:: class NewProblem(TimeMixin, Problem): pass Attributes: T (firedrake.Function): The trial function for the problem. v (firedrake.Function): The test function for the problem. a (firedrake.Function): The section containing the combination of terms involving both T and v. T_ (firedrake.Function): A function used to hold the previous value for T. C (firedrake.Function): A function used to hold the heat capacity of the mesh. _delT (firedrake.Function): A function used as a placeholder for the time derivative. max_t (float): The time to iterate up to. dt (float): The size of each time step. _dt_invc (firedrake.Constant): Utility function to pass time step information to firedrake. steps (int): Number of iterations. steady_state (bool): Toggle whether solving steady state (dT/dt = 0) or time dependent problem. """ # pylint: disable=too-many-instance-attributes, no-member def __init__(self, *args, **kwargs): """ Initialiser for TimeMixin. Add M/dt to a and L. """ super().__init__(*args, **kwargs) self._add_function('T') self._add_function('a') self._add_function('T_') self._add_function('C') self._add_function('_delT') self.max_t = 1e-9 self.dt = 1e-10 self._dt_invc = Constant(0) self.steps = self.max_t / self.dt self.steady_state = True self.a += self._M() def set_method(self, method='BackwardEuler', **kwargs): """ Replace T with the correct substitution for the method, then replace the delT placeholder with the simple finite difference approximation: dT/dt ~ (T - T_)/delta_t Args: method (str, optional): The method to use. Defaults to 'BackwardEuler'. """ T = self.T iter_method = IterationMethod(self) substitution = iter_method.get_substitution(method, **kwargs) self._update_func('T', substitution) self.T = T delT = (self.T - self.T_) * self._dt_invc self._update_func('_delT', delT) def remove_timescale(self): """ Set 1/dt to 0 so that the M terms vanish. """ self.steady_state = True self._dt_invc.assign(0) def set_timescale(self, max_t=None, dt=None, steps=None): """ Set the time stepping variables (max_t, dt, and number of steps). This method is designed to be given 2 variables and calculate the third. If 3 variables are given it will check that they are consistent. Args: max_t (float, optional): The time to iterate until. Defaults to None. dt (float, optional): The increase in time for each iteration. Defaults to None. steps (int, optional): The number of steps to take. Defaults to None. Raises: ValueError: If less than 2 variables are provided. ValueError: If 3 values are provided which are inconsistent. RuntimeWarning: If steps is not an int. """ num_var_none = 0 if max_t is None: num_var_none += 1 if dt is None: num_var_none += 1 if steps is None: num_var_none += 1 if num_var_none > 1: raise ValueError('Must specify at least 2 of max_t, dt, and steps') if num_var_none == 0: calc_steps = int(max_t / dt) if calc_steps != max_t / dt: calc_steps += 1 if steps != calc_steps: raise ValueError('Conflicting arguments. Try specifying only 2' ' of max_t, dt, and steps.') if max_t is None: max_t = dt * steps if dt is None: dt = max_t / steps if steps is None: steps = max_t / dt if steps != int(steps): steps = int(steps + 1) LOGGER.warning("steps is not an integer, rounding up.") self.max_t = max_t self.dt = dt self._dt_invc.assign(1 / dt) self.steps = steps self.steady_state = False def _M(self): """ Create the mass matrix section. Returns: Function: The complete mass matrix section using delT. """ return self.C * self._delT * self.v * dx