class IncompressibleState(BaroclinicState): def __init__(self, mesh, vertical_degree=1, horizontal_degree=1, family="RT", z=None, k=None, Omega=None, mu=None, timestepping=None, output=None, parameters=None, diagnostics=None, fieldlist=None, diagnostic_fields=[], on_sphere=False): super(IncompressibleState, self).__init__(mesh=mesh, vertical_degree=vertical_degree, horizontal_degree=horizontal_degree, family=family, z=z, k=k, Omega=Omega, mu=mu, timestepping=timestepping, output=output, parameters=parameters, diagnostics=diagnostics, fieldlist=fieldlist, diagnostic_fields=diagnostic_fields) if parameters.geopotential: raise RuntimeError("geopotential formulation is not compatible with incompressible Boussinesq") def set_reference_profiles(self, b_ref): """ Initialise reference profiles :arg b_ref: :class:`.Function` object, reference bouyancy """ self.bbar = Function(self.V[2]) self.bbar.project(b_ref)
class CompressibleState(BaroclinicState): def set_reference_profiles(self, rho_ref, theta_ref): """ Initialise reference profiles :arg rho_ref: :class:`.Function` object, initial rho :arg theta_ref: :class:`.Function` object, initial theta """ self.rhobar = Function(self.V[1]) self.thetabar = Function(self.V[2]) self.rhobar.project(rho_ref) self.thetabar.project(theta_ref)
class Diagnostics(object): @staticmethod def min(f_in): fmin = op2.Global(1, [np.finfo(float).max], dtype=float) if len(f_in.ufl_shape) > 0: mesh = f_in.function_space().mesh() V = FunctionSpace(mesh, "DG", 1) f = Function(V).project(sqrt(inner(f_in, f_in))) else: f = f_in op2.par_loop( op2.Kernel( """void minify(double *a, double *b) { a[0] = a[0] > fabs(b[0]) ? fabs(b[0]) : a[0]; }""", "minify"), f.dof_dset.set, fmin(op2.MIN), f.dat(op2.READ)) return fmin.data[0] @staticmethod def max(f_in): fmax = op2.Global(1, [np.finfo(float).min], dtype=float) if len(f_in.ufl_shape) > 0: mesh = f_in.function_space().mesh() V = FunctionSpace(mesh, "DG", 1) f = Function(V).project(sqrt(inner(f_in, f_in))) else: f = f_in op2.par_loop( op2.Kernel( """void maxify(double *a, double *b) { a[0] = a[0] < fabs(b[0]) ? fabs(b[0]) : a[0]; }""", "maxify"), f.dof_dset.set, fmax(op2.MAX), f.dat(op2.READ)) return fmax.data[0] @staticmethod def l2(f): return sqrt(assemble(dot(f, f) * dx)) @staticmethod def energy(h, u, g): return assemble(0.5 * h * (inner(u, u) + g * h) * dx) def max_courant_number(self, u, dt): if not hasattr(self, "_area"): V = FunctionSpace(u.function_space().mesh(), "DG", 0) expr = TestFunction(V) * dx self._area = assemble(expr) self.Courant = Function(V) self.Courant.project(sqrt(inner(u, u)) / sqrt(self._area) * dt) return self.Courant.dat.data.max()
def _prepare_output(self, function, max_elem): from firedrake import FunctionSpace, VectorFunctionSpace, \ TensorFunctionSpace, Function from tsfc.finatinterface import create_element as create_finat_element name = function.name() # Need to project/interpolate? # If space is not the max element, we must do so. finat_elem = function.function_space().finat_element if finat_elem == create_finat_element(max_elem): return OFunction(array=get_array(function), name=name, function=function) # OK, let's go and do it. # Build appropriate space for output function. shape = function.ufl_shape if len(shape) == 0: V = FunctionSpace(function.ufl_domain(), max_elem) elif len(shape) == 1: if numpy.prod(shape) > 3: raise ValueError( "Can't write vectors with more than 3 components") V = VectorFunctionSpace(function.ufl_domain(), max_elem, dim=shape[0]) elif len(shape) == 2: if numpy.prod(shape) > 9: raise ValueError( "Can't write tensors with more than 9 components") V = TensorFunctionSpace(function.ufl_domain(), max_elem, shape=shape) else: raise ValueError("Unsupported shape %s" % (shape, )) output = Function(V) if self.project: output.project(function) else: output.interpolate(function) return OFunction(array=get_array(output), name=name, function=output)
class Advection(object, metaclass=ABCMeta): """ Base class for advection schemes. :arg state: :class:`.State` object. :arg field: field to be advected :arg equation: :class:`.Equation` object, specifying the equation that field satisfies :arg solver_parameters: solver_parameters :arg limiter: :class:`.Limiter` object. :arg options: :class:`.AdvectionOptions` object """ def __init__(self, state, field, equation=None, *, solver_parameters=None, limiter=None): if equation is not None: self.state = state self.field = field self.equation = equation # get ubar from the equation class self.ubar = self.equation.ubar self.dt = self.state.timestepping.dt # get default solver options if none passed in if solver_parameters is None: self.solver_parameters = equation.solver_parameters else: self.solver_parameters = solver_parameters if logger.isEnabledFor(DEBUG): self.solver_parameters["ksp_monitor_true_residual"] = True self.limiter = limiter if hasattr(equation, "options"): self.discretisation_option = equation.options.name self._setup(state, field, equation.options) else: self.discretisation_option = None self.fs = field.function_space() # setup required functions self.dq = Function(self.fs) self.q1 = Function(self.fs) def _setup(self, state, field, options): if options.name in ["embedded_dg", "recovered"]: self.fs = options.embedding_space self.xdg_in = Function(self.fs) self.xdg_out = Function(self.fs) self.x_projected = Function(field.function_space()) parameters = {'ksp_type': 'cg', 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'} self.Projector = Projector(self.xdg_out, self.x_projected, solver_parameters=parameters) if options.name == "recovered": # set up the necessary functions self.x_in = Function(field.function_space()) x_rec = Function(options.recovered_space) x_brok = Function(options.broken_space) # set up interpolators and projectors self.x_rec_projector = Recoverer(self.x_in, x_rec, VDG=self.fs, boundary_method=options.boundary_method) # recovered function self.x_brok_projector = Projector(x_rec, x_brok) # function projected back self.xdg_interpolator = Interpolator(self.x_in + x_rec - x_brok, self.xdg_in) if self.limiter is not None: self.x_brok_interpolator = Interpolator(self.xdg_out, x_brok) self.x_out_projector = Recoverer(x_brok, self.x_projected) def pre_apply(self, x_in, discretisation_option): """ Extra steps to advection if using an embedded method, which might be either the plain embedded method or the recovered space advection scheme. :arg x_in: the input set of prognostic fields. :arg discretisation option: string specifying which scheme to use. """ if discretisation_option == "embedded_dg": try: self.xdg_in.interpolate(x_in) except NotImplementedError: self.xdg_in.project(x_in) elif discretisation_option == "recovered": self.x_in.assign(x_in) self.x_rec_projector.project() self.x_brok_projector.project() self.xdg_interpolator.interpolate() def post_apply(self, x_out, discretisation_option): """ The projection steps, returning a field to its original space for an embedded DG advection scheme. For the case of the recovered scheme, there are two options dependent on whether the scheme is limited or not. :arg x_out: the outgoing field. :arg discretisation_option: string specifying which option to use. """ if discretisation_option == "embedded_dg": self.Projector.project() elif discretisation_option == "recovered": if self.limiter is not None: self.x_brok_interpolator.interpolate() self.x_out_projector.project() else: self.Projector.project() x_out.assign(self.x_projected) @abstractproperty def lhs(self): return self.equation.mass_term(self.equation.trial) @abstractproperty def rhs(self): return self.equation.mass_term(self.q1) - self.dt*self.equation.advection_term(self.q1) def update_ubar(self, xn, xnp1, alpha): un = xn.split()[0] unp1 = xnp1.split()[0] self.ubar.assign(un + alpha*(unp1-un)) @cached_property def solver(self): # setup solver using lhs and rhs defined in derived class problem = LinearVariationalProblem(self.lhs, self.rhs, self.dq) solver_name = self.field.name()+self.equation.__class__.__name__+self.__class__.__name__ return LinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name) @abstractmethod def apply(self, x_in, x_out): """ Function takes x as input, computes L(x) as defined by the equation, and returns x_out as output. :arg x: :class:`.Function` object, the input Function. :arg x_out: :class:`.Function` object, the output Function. """ pass
class TimeDiscretisation(object, metaclass=ABCMeta): """ Base class for time discretisation schemes. :arg state: :class:`.State` object. :arg field: field to be evolved :arg equation: :class:`.Equation` object, specifying the equation that field satisfies :arg solver_parameters: solver_parameters :arg limiter: :class:`.Limiter` object. :arg options: :class:`.DiscretisationOptions` object """ def __init__(self, state, field_name=None, solver_parameters=None, limiter=None, options=None): self.state = state self.field_name = field_name self.dt = self.state.dt self.limiter = limiter self.options = options if options is not None: self.discretisation_option = options.name else: self.discretisation_option = None # get default solver options if none passed in if solver_parameters is None: self.solver_parameters = {'ksp_type': 'cg', 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'} else: self.solver_parameters = solver_parameters if logger.isEnabledFor(DEBUG): self.solver_parameters["ksp_monitor_true_residual"] = None def setup(self, equation, uadv=None, apply_bcs=True, *active_labels): self.residual = equation.residual if self.field_name is not None: self.idx = equation.field_names.index(self.field_name) self.fs = self.state.fields(self.field_name).function_space() self.residual = self.residual.label_map( lambda t: t.get(prognostic) == self.field_name, lambda t: Term( split_form(t.form)[self.idx].form, t.labels), drop) bcs = equation.bcs[self.field_name] else: self.field_name = equation.field_name self.fs = equation.function_space self.idx = None if type(self.fs.ufl_element()) is MixedElement: bcs = [bc for _, bcs in equation.bcs.items() for bc in bcs] else: bcs = equation.bcs[self.field_name] if len(active_labels) > 0: self.residual = self.residual.label_map( lambda t: any(t.has_label(time_derivative, *active_labels)), map_if_false=drop) options = self.options # -------------------------------------------------------------------- # # Routines relating to transport # -------------------------------------------------------------------- # if hasattr(self.options, 'ibp'): self.replace_transport_term() self.replace_transporting_velocity(uadv) # -------------------------------------------------------------------- # # Wrappers for embedded / recovery methods # -------------------------------------------------------------------- # if self.discretisation_option in ["embedded_dg", "recovered"]: # construct the embedding space if not specified if options.embedding_space is None: V_elt = BrokenElement(self.fs.ufl_element()) self.fs = FunctionSpace(self.state.mesh, V_elt) else: self.fs = options.embedding_space self.xdg_in = Function(self.fs) self.xdg_out = Function(self.fs) if self.idx is None: self.x_projected = Function(equation.function_space) else: self.x_projected = Function(self.state.fields(self.field_name).function_space()) new_test = TestFunction(self.fs) parameters = {'ksp_type': 'cg', 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'} # -------------------------------------------------------------------- # # Make boundary conditions # -------------------------------------------------------------------- # if not apply_bcs: self.bcs = None elif self.discretisation_option in ["embedded_dg", "recovered"]: # Transfer boundary conditions onto test function space self.bcs = [DirichletBC(self.fs, bc.function_arg, bc.sub_domain) for bc in bcs] else: self.bcs = bcs # -------------------------------------------------------------------- # # Modify test function for SUPG methods # -------------------------------------------------------------------- # if self.discretisation_option == "supg": # construct tau, if it is not specified dim = self.state.mesh.topological_dimension() if options.tau is not None: # if tau is provided, check that is has the right size tau = options.tau assert as_ufl(tau).ufl_shape == (dim, dim), "Provided tau has incorrect shape!" else: # create tuple of default values of size dim default_vals = [options.default*self.dt]*dim # check for directions is which the space is discontinuous # so that we don't apply supg in that direction if is_cg(self.fs): vals = default_vals else: space = self.fs.ufl_element().sobolev_space() if space.name in ["HDiv", "DirectionalH"]: vals = [default_vals[i] if space[i].name == "H1" else 0. for i in range(dim)] else: raise ValueError("I don't know what to do with space %s" % space) tau = Constant(tuple([ tuple( [vals[j] if i == j else 0. for i, v in enumerate(vals)] ) for j in range(dim)]) ) self.solver_parameters = {'ksp_type': 'gmres', 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'} test = TestFunction(self.fs) new_test = test + dot(dot(uadv, tau), grad(test)) if self.discretisation_option is not None: # replace the original test function with one defined on # the embedding space, as this is the space where the # the problem will be solved self.residual = self.residual.label_map( all_terms, map_if_true=replace_test_function(new_test)) if self.discretisation_option == "embedded_dg": if self.limiter is None: self.x_out_projector = Projector(self.xdg_out, self.x_projected, solver_parameters=parameters) else: self.x_out_projector = Recoverer(self.xdg_out, self.x_projected) if self.discretisation_option == "recovered": # set up the necessary functions self.x_in = Function(self.state.fields(self.field_name).function_space()) x_rec = Function(options.recovered_space) x_brok = Function(options.broken_space) # set up interpolators and projectors self.x_rec_projector = Recoverer(self.x_in, x_rec, VDG=self.fs, boundary_method=options.boundary_method) # recovered function self.x_brok_projector = Projector(x_rec, x_brok) # function projected back self.xdg_interpolator = Interpolator(self.x_in + x_rec - x_brok, self.xdg_in) if self.limiter is not None: self.x_brok_interpolator = Interpolator(self.xdg_out, x_brok) self.x_out_projector = Recoverer(x_brok, self.x_projected) else: self.x_out_projector = Projector(self.xdg_out, self.x_projected) # setup required functions self.dq = Function(self.fs) self.q1 = Function(self.fs) def pre_apply(self, x_in, discretisation_option): """ Extra steps to discretisation if using an embedded method, which might be either the plain embedded method or the recovered space scheme. :arg x_in: the input set of prognostic fields. :arg discretisation option: string specifying which scheme to use. """ if discretisation_option == "embedded_dg": try: self.xdg_in.interpolate(x_in) except NotImplementedError: self.xdg_in.project(x_in) elif discretisation_option == "recovered": self.x_in.assign(x_in) self.x_rec_projector.project() self.x_brok_projector.project() self.xdg_interpolator.interpolate() def post_apply(self, x_out, discretisation_option): """ The projection steps, returning a field to its original space for an embedded DG scheme. For the case of the recovered scheme, there are two options dependent on whether the scheme is limited or not. :arg x_out: the outgoing field. :arg discretisation_option: string specifying which option to use. """ if discretisation_option == "recovered" and self.limiter is not None: self.x_brok_interpolator.interpolate() self.x_out_projector.project() x_out.assign(self.x_projected) @abstractproperty def lhs(self): l = self.residual.label_map( lambda t: t.has_label(time_derivative), map_if_true=replace_subject(self.dq, self.idx), map_if_false=drop) return l.form @abstractproperty def rhs(self): r = self.residual.label_map( all_terms, map_if_true=replace_subject(self.q1, self.idx)) r = r.label_map( lambda t: t.has_label(time_derivative), map_if_false=lambda t: -self.dt*t) return r.form def replace_transport_term(self): """ This routine allows the default transport term to be replaced with a different one, specified through the transport options. This is necessary because when the prognostic equations are declared, the whole transport """ # Extract transport term of equation old_transport_term_list = self.residual.label_map( lambda t: t.has_label(transport), map_if_false=drop) # If there are more transport terms, extract only the one for this variable if len(old_transport_term_list.terms) > 1: raise NotImplementedError('Cannot replace transport terms when there are more than one') # Then we should only have one transport term old_transport_term = old_transport_term_list.terms[0] # If the transport term has an ibp label, then it could be replaced if old_transport_term.has_label(ibp_label) and hasattr(self.options, 'ibp'): # Do the options specify a different ibp to the old transport term? if old_transport_term.labels['ibp'] != self.options.ibp: # Set up a new transport term field = self.state.fields(self.field_name) test = TestFunction(self.fs) # Set up new transport term (depending on the type of transport equation) if old_transport_term.labels['transport'] == TransportEquationType.advective: new_transport_term = advection_form(self.state, test, field, ibp=self.options.ibp) elif old_transport_term.labels['transport'] == TransportEquationType.conservative: new_transport_term = continuity_form(self.state, test, field, ibp=self.options.ibp) else: raise NotImplementedError(f'Replacement of transport term not implemented yet for {old_transport_term.labels["transport"]}') # Finally, drop the old transport term and add the new one self.residual = self.residual.label_map( lambda t: t.has_label(transport), map_if_true=drop) self.residual += subject(new_transport_term, field) def replace_transporting_velocity(self, uadv): # replace the transporting velocity in any terms that contain it if any([t.has_label(transporting_velocity) for t in self.residual]): assert uadv is not None if uadv == "prognostic": self.residual = self.residual.label_map( lambda t: t.has_label(transporting_velocity), map_if_true=lambda t: Term(ufl.replace( t.form, {t.get(transporting_velocity): split(t.get(subject))[0]}), t.labels) ) else: self.residual = self.residual.label_map( lambda t: t.has_label(transporting_velocity), map_if_true=lambda t: Term(ufl.replace( t.form, {t.get(transporting_velocity): uadv}), t.labels) ) self.residual = transporting_velocity.update_value(self.residual, uadv) @cached_property def solver(self): # setup solver using lhs and rhs defined in derived class problem = NonlinearVariationalProblem(self.lhs-self.rhs, self.dq, bcs=self.bcs) solver_name = self.field_name+self.__class__.__name__ return NonlinearVariationalSolver(problem, solver_parameters=self.solver_parameters, options_prefix=solver_name) @abstractmethod def apply(self, x_in, x_out): """ Function takes x as input, computes L(x) as defined by the equation, and returns x_out as output. :arg x: :class:`.Function` object, the input Function. :arg x_out: :class:`.Function` object, the output Function. """ pass