def generate_tracer_transport_terms(self, state, active_tracers): """ Adds the transport forms for the active tracers to the equation. :arg state: The `State` object. :arg active_tracers: A list of `ActiveTracer` objects that encode the metadata for the active tracers. """ # By default return None if no tracers are to be transported adv_form = None no_tracer_transported = True for i, tracer in enumerate(active_tracers): if tracer.transport_flag: idx = self.field_names.index(tracer.name) tracer_prog = split(self.X)[idx] tracer_test = self.tests[idx] if tracer.transport_eqn == TransportEquationType.advective: tracer_adv = prognostic( advection_form(state, tracer_test, tracer_prog), tracer.name) elif tracer.transport_eqn == TransportEquationType.conservative: tracer_adv = prognostic( continuity_form(state, tracer_test, tracer_prog), tracer.name) else: raise ValueError( f'Transport eqn {tracer.transport_eqn} not recognised') if no_tracer_transported: # We arrive here for the first tracer to be transported adv_form = subject(tracer_adv, self.X) no_tracer_transported = False else: adv_form += subject(tracer_adv, self.X) return adv_form
def generate_mass_terms(self): """ Builds the weak time derivative terms ("mass terms") for all the prognostic variables of the equation set. Returns a :class:`LabelledForm` containing all of these terms. """ for i, (test, field_name) in enumerate(zip(self.tests, self.field_names)): prog = split(self.X)[i] mass = subject(prognostic(inner(prog, test) * dx, field_name), self.X) if i == 0: mass_form = time_derivative(mass) else: mass_form += time_derivative(mass) return mass_form
def __init__( self, state, family, degree, Omega=None, sponge=None, extra_terms=None, terms_to_linearise={ 'u': [time_derivative], 'rho': [time_derivative, transport], 'theta': [time_derivative, transport] }, u_transport_option="vector_invariant_form", diffusion_options=None, no_normal_flow_bc_ids=None, active_tracers=None): super().__init__(state, family, degree, Omega=Omega, sponge=sponge, extra_terms=extra_terms, terms_to_linearise=terms_to_linearise, u_transport_option=u_transport_option, diffusion_options=diffusion_options, no_normal_flow_bc_ids=no_normal_flow_bc_ids, active_tracers=active_tracers) self.residual = self.residual.label_map( lambda t: t.has_label(time_derivative), map_if_true=lambda t: hydrostatic(t, self.hydrostatic_projection(t) )) k = self.state.k u = split(self.X)[0] self.residual += name( subject( prognostic(-inner(k, self.tests[0]) * inner(k, u) * dx, "u"), self.X), "hydrostatic_form")
def __init__( self, state, family, degree, fexpr=None, bexpr=None, terms_to_linearise={ 'D': [time_derivative, transport], 'u': [time_derivative, pressure_gradient, coriolis] }, u_transport_option="vector_invariant_form", no_normal_flow_bc_ids=None, active_tracers=None): super().__init__(state, family, degree, fexpr=fexpr, bexpr=bexpr, terms_to_linearise=terms_to_linearise, u_transport_option=u_transport_option, no_normal_flow_bc_ids=no_normal_flow_bc_ids, active_tracers=active_tracers) # Use the underlying routine to do a first linearisation of the equations self.linearise_equation_set() # D transport term is a special case -- add facet term _, D = split(self.X) _, phi = self.tests D_adv = prognostic( linear_continuity_form(state, phi, D, facet_term=True), "D") self.residual = self.residual.label_map( lambda t: t.has_label(transport) and t.get(prognostic) == "D", map_if_true=lambda t: Term(D_adv.form, t.labels))
def __init__( self, state, family, degree, Omega=None, terms_to_linearise={ 'u': [time_derivative], 'p': [time_derivative], 'b': [time_derivative, transport] }, u_transport_option="vector_invariant_form", no_normal_flow_bc_ids=None, active_tracers=None): field_names = ['u', 'p', 'b'] if active_tracers is not None: raise NotImplementedError( 'Tracers not implemented for Boussinesq equations') if active_tracers is None: active_tracers = [] super().__init__(field_names, state, family, degree, terms_to_linearise=terms_to_linearise, no_normal_flow_bc_ids=no_normal_flow_bc_ids, active_tracers=active_tracers) w, phi, gamma = self.tests[0:3] u, p, b = split(self.X) bbar = state.fields("bbar", space=state.spaces("theta"), dump=False) bbar = state.fields("pbar", space=state.spaces("DG"), dump=False) # -------------------------------------------------------------------- # # Time Derivative Terms # -------------------------------------------------------------------- # mass_form = self.generate_mass_terms() # -------------------------------------------------------------------- # # Transport Terms # -------------------------------------------------------------------- # # Velocity transport term -- depends on formulation if u_transport_option == "vector_invariant_form": u_adv = prognostic(vector_invariant_form(state, w, u), "u") elif u_transport_option == "vector_advection_form": u_adv = prognostic(advection_form(state, w, u), "u") elif u_transport_option == "vector_manifold_advection_form": u_adv = prognostic(vector_manifold_advection_form(state, w, u), "u") elif u_transport_option == "circulation_form": ke_form = kinetic_energy_form(state, w, u) ke_form = transport.remove(ke_form) ke_form = ke_form.label_map( lambda t: t.has_label(transporting_velocity), lambda t: Term( ufl.replace(t.form, {t.get(transporting_velocity): u}), t. labels)) ke_form = transporting_velocity.remove(ke_form) u_adv = advection_equation_circulation_form(state, w, u) + ke_form else: raise ValueError("Invalid u_transport_option: %s" % u_transport_option) # Buoyancy transport b_adv = prognostic(advection_form(state, gamma, b), "b") linear_b_adv = linear_advection_form(state, gamma, bbar).label_map( lambda t: t.has_label(transporting_velocity), lambda t: Term( ufl.replace(t.form, { t.get(transporting_velocity): self.trials[0] }), t.labels)) b_adv = linearisation(b_adv, linear_b_adv) adv_form = subject(u_adv + b_adv, self.X) # Add transport of tracers if len(active_tracers) > 0: adv_form += self.generate_tracer_transport_terms( state, active_tracers) # -------------------------------------------------------------------- # # Pressure Gradient Term # -------------------------------------------------------------------- # pressure_gradient_form = subject(prognostic(-div(w) * p * dx, "u"), self.X) # -------------------------------------------------------------------- # # Gravitational Term # -------------------------------------------------------------------- # gravity_form = subject(prognostic(-b * inner(w, state.k) * dx, "u"), self.X) # -------------------------------------------------------------------- # # Divergence Term # -------------------------------------------------------------------- # # This enforces that div(u) = 0 # The p features here so that the div(u) evaluated in the "forcing" step # replaces the whole pressure field, rather than merely providing an # increment to it. divergence_form = name( subject(prognostic(phi * (p - div(u)) * dx, "p"), self.X), "incompressibility") residual = (mass_form + adv_form + divergence_form + pressure_gradient_form + gravity_form) # -------------------------------------------------------------------- # # Extra Terms (Coriolis) # -------------------------------------------------------------------- # if Omega is not None: residual += subject( prognostic(inner(w, cross(2 * Omega, u)) * dx, "u"), self.X) # -------------------------------------------------------------------- # # Linearise equations # -------------------------------------------------------------------- # # TODO: add linearisation states for variables # Add linearisations to equations self.residual = self.generate_linear_terms(residual, self.terms_to_linearise)
def __init__( self, state, family, degree, Omega=None, sponge=None, extra_terms=None, terms_to_linearise={ 'u': [time_derivative], 'rho': [time_derivative, transport], 'theta': [time_derivative, transport] }, u_transport_option="vector_invariant_form", diffusion_options=None, no_normal_flow_bc_ids=None, active_tracers=None): field_names = ['u', 'rho', 'theta'] if active_tracers is None: active_tracers = [] super().__init__(field_names, state, family, degree, terms_to_linearise=terms_to_linearise, no_normal_flow_bc_ids=no_normal_flow_bc_ids, active_tracers=active_tracers) g = state.parameters.g cp = state.parameters.cp w, phi, gamma = self.tests[0:3] u, rho, theta = split(self.X)[0:3] rhobar = state.fields("rhobar", space=state.spaces("DG"), dump=False) thetabar = state.fields("thetabar", space=state.spaces("theta"), dump=False) zero_expr = Constant(0.0) * theta exner = exner_pressure(state.parameters, rho, theta) n = FacetNormal(state.mesh) # -------------------------------------------------------------------- # # Time Derivative Terms # -------------------------------------------------------------------- # mass_form = self.generate_mass_terms() # -------------------------------------------------------------------- # # Transport Terms # -------------------------------------------------------------------- # # Velocity transport term -- depends on formulation if u_transport_option == "vector_invariant_form": u_adv = prognostic(vector_invariant_form(state, w, u), "u") elif u_transport_option == "vector_advection_form": u_adv = prognostic(advection_form(state, w, u), "u") elif u_transport_option == "vector_manifold_advection_form": u_adv = prognostic(vector_manifold_advection_form(state, w, u), "u") elif u_transport_option == "circulation_form": ke_form = kinetic_energy_form(state, w, u) ke_form = transport.remove(ke_form) ke_form = ke_form.label_map( lambda t: t.has_label(transporting_velocity), lambda t: Term( ufl.replace(t.form, {t.get(transporting_velocity): u}), t. labels)) ke_form = transporting_velocity.remove(ke_form) u_adv = advection_equation_circulation_form(state, w, u) + ke_form else: raise ValueError("Invalid u_transport_option: %s" % u_transport_option) # Density transport (conservative form) rho_adv = prognostic(continuity_form(state, phi, rho), "rho") # Transport term needs special linearisation if transport in terms_to_linearise['rho']: linear_rho_adv = linear_continuity_form( state, phi, rhobar).label_map( lambda t: t.has_label(transporting_velocity), lambda t: Term( ufl.replace(t.form, { t.get(transporting_velocity): self.trials[0] }), t.labels)) rho_adv = linearisation(rho_adv, linear_rho_adv) # Potential temperature transport (advective form) theta_adv = prognostic(advection_form(state, gamma, theta), "theta") # Transport term needs special linearisation if transport in terms_to_linearise['theta']: linear_theta_adv = linear_advection_form( state, gamma, thetabar).label_map( lambda t: t.has_label(transporting_velocity), lambda t: Term( ufl.replace(t.form, { t.get(transporting_velocity): self.trials[0] }), t.labels)) theta_adv = linearisation(theta_adv, linear_theta_adv) adv_form = subject(u_adv + rho_adv + theta_adv, self.X) # Add transport of tracers if len(active_tracers) > 0: adv_form += self.generate_tracer_transport_terms( state, active_tracers) # -------------------------------------------------------------------- # # Pressure Gradient Term # -------------------------------------------------------------------- # # First get total mass of tracers tracer_mr_total = zero_expr for tracer in active_tracers: if tracer.variable_type == TracerVariableType.mixing_ratio: idx = self.field_names.index(tracer.name) tracer_mr_total += split(self.X)[idx] else: raise NotImplementedError( 'Only mixing ratio tracers are implemented') theta_v = theta / (Constant(1.0) + tracer_mr_total) pressure_gradient_form = name( subject( prognostic( cp * (-div(theta_v * w) * exner * dx + jump(theta_v * w, n) * avg(exner) * dS_v), "u"), self.X), "pressure_gradient") # -------------------------------------------------------------------- # # Gravitational Term # -------------------------------------------------------------------- # gravity_form = subject( prognostic(Term(g * inner(state.k, w) * dx), "u"), self.X) residual = (mass_form + adv_form + pressure_gradient_form + gravity_form) # -------------------------------------------------------------------- # # Moist Thermodynamic Divergence Term # -------------------------------------------------------------------- # if len(active_tracers) > 0: cv = state.parameters.cv c_vv = state.parameters.c_vv c_pv = state.parameters.c_pv c_pl = state.parameters.c_pl R_d = state.parameters.R_d R_v = state.parameters.R_v # Get gas and liquid moisture mixing ratios mr_l = zero_expr mr_v = zero_expr for tracer in active_tracers: if tracer.is_moisture: if tracer.variable_type == TracerVariableType.mixing_ratio: idx = self.field_names.index(tracer.name) if tracer.phase == Phases.gas: mr_v += split(self.X)[idx] elif tracer.phase == Phases.liquid: mr_l += split(self.X)[idx] else: raise NotImplementedError( 'Only mixing ratio tracers are implemented') c_vml = cv + mr_v * c_vv + mr_l * c_pl c_pml = cp + mr_v * c_pv + mr_l * c_pl R_m = R_d + mr_v * R_v residual += subject( prognostic( gamma * theta * div(u) * (R_m / c_vml - (R_d * c_pml) / (cp * c_vml)) * dx, "theta"), self.X) # -------------------------------------------------------------------- # # Extra Terms (Coriolis, Sponge, Diffusion and others) # -------------------------------------------------------------------- # if Omega is not None: residual += subject( prognostic(inner(w, cross(2 * Omega, u)) * dx, "u"), self.X) if sponge is not None: W_DG = FunctionSpace(state.mesh, "DG", 2) x = SpatialCoordinate(state.mesh) z = x[len(x) - 1] H = sponge.H zc = sponge.z_level assert float(zc) < float( H ), "you have set the sponge level above the height of your domain" mubar = sponge.mubar muexpr = conditional( z <= zc, 0.0, mubar * sin((pi / 2.) * (z - zc) / (H - zc))**2) self.mu = Function(W_DG).interpolate(muexpr) residual += name( subject( prognostic( self.mu * inner(w, state.k) * inner(u, state.k) * dx, "u"), self.X), "sponge") if diffusion_options is not None: for field, diffusion in diffusion_options: idx = self.field_names.index(field) test = self.tests[idx] fn = split(self.X)[idx] residual += subject( prognostic( interior_penalty_diffusion_form( state, test, fn, diffusion), field), self.X) if extra_terms is not None: for field, term in extra_terms: idx = self.field_names.index(field) test = self.tests[idx] residual += subject(prognostic(inner(test, term) * dx, field), self.X) # -------------------------------------------------------------------- # # Linearise equations # -------------------------------------------------------------------- # # TODO: add linearisation states for variables # Add linearisations to equations self.residual = self.generate_linear_terms(residual, self.terms_to_linearise)
def __init__( self, state, family, degree, fexpr=None, bexpr=None, terms_to_linearise={ 'D': [time_derivative, transport], 'u': [time_derivative, pressure_gradient] }, u_transport_option="vector_invariant_form", no_normal_flow_bc_ids=None, active_tracers=None): field_names = ["u", "D"] if active_tracers is not None: raise NotImplementedError( 'Tracers not implemented for shallow water equations') if active_tracers is None: active_tracers = [] super().__init__(field_names, state, family, degree, terms_to_linearise=terms_to_linearise, no_normal_flow_bc_ids=no_normal_flow_bc_ids, active_tracers=active_tracers) g = state.parameters.g H = state.parameters.H w, phi = self.tests u, D = split(self.X) # -------------------------------------------------------------------- # # Time Derivative Terms # -------------------------------------------------------------------- # mass_form = self.generate_mass_terms() # -------------------------------------------------------------------- # # Transport Terms # -------------------------------------------------------------------- # # Velocity transport term -- depends on formulation if u_transport_option == "vector_invariant_form": u_adv = prognostic(vector_invariant_form(state, w, u), "u") elif u_transport_option == "vector_advection_form": u_adv = prognostic(advection_form(state, w, u), "u") elif u_transport_option == "vector_manifold_advection_form": u_adv = prognostic(vector_manifold_advection_form(state, w, u), "u") elif u_transport_option == "circulation_form": ke_form = kinetic_energy_form(state, w, u) ke_form = transport.remove(ke_form) ke_form = ke_form.label_map( lambda t: t.has_label(transporting_velocity), lambda t: Term( ufl.replace(t.form, {t.get(transporting_velocity): u}), t. labels)) ke_form = transporting_velocity.remove(ke_form) u_adv = advection_equation_circulation_form(state, w, u) + ke_form else: raise ValueError("Invalid u_transport_option: %s" % u_transport_option) # Depth transport term D_adv = prognostic(continuity_form(state, phi, D), "D") # Transport term needs special linearisation if transport in terms_to_linearise['D']: linear_D_adv = linear_continuity_form(state, phi, H).label_map( lambda t: t.has_label(transporting_velocity), lambda t: Term( ufl.replace(t.form, { t.get(transporting_velocity): self.trials[0] }), t.labels)) # Add linearisation to D_adv D_adv = linearisation(D_adv, linear_D_adv) adv_form = subject(u_adv + D_adv, self.X) # Add transport of tracers if len(active_tracers) > 0: adv_form += self.generate_tracer_transport_terms( state, active_tracers) # -------------------------------------------------------------------- # # Pressure Gradient Term # -------------------------------------------------------------------- # pressure_gradient_form = pressure_gradient( subject(prognostic(-g * div(w) * D * dx, "u"), self.X)) residual = (mass_form + adv_form + pressure_gradient_form) # -------------------------------------------------------------------- # # Extra Terms (Coriolis and Topography) # -------------------------------------------------------------------- # if fexpr is not None: V = FunctionSpace(state.mesh, "CG", 1) f = state.fields("coriolis", space=V) f.interpolate(fexpr) coriolis_form = coriolis( subject(prognostic(f * inner(state.perp(u), w) * dx, "u"), self.X)) residual += coriolis_form if bexpr is not None: b = state.fields("topography", state.spaces("DG")) b.interpolate(bexpr) topography_form = subject(prognostic(-g * div(w) * b * dx, "u"), self.X) residual += topography_form # -------------------------------------------------------------------- # # Linearise equations # -------------------------------------------------------------------- # u, D = self.X.split() # Linearise about D = H # TODO: add linearisation state for u D.assign(Constant(H)) u, D = split(self.X) # Add linearisations to equations self.residual = self.generate_linear_terms(residual, self.terms_to_linearise)