def setup(self, state): if not self._initialised: space = state.fields("theta").function_space() broken_space = FunctionSpace(state.mesh, BrokenElement(space.ufl_element())) boundary_method = Boundary_Method.physics if (state.vertical_degree == 0 and state.horizontal_degree == 0) else None super().setup(state, space=space) # now let's attach all of our fields self.u = state.fields("u") self.rho = state.fields("rho") self.theta = state.fields("theta") self.rho_averaged = Function(space) self.recoverer = Recoverer(self.rho, self.rho_averaged, VDG=broken_space, boundary_method=boundary_method) try: self.r_v = state.fields("water_v") except NotImplementedError: self.r_v = Constant(0.0) try: self.r_c = state.fields("water_c") except NotImplementedError: self.r_c = Constant(0.0) try: self.rain = state.fields("rain") except NotImplementedError: self.rain = Constant(0.0) # now let's store the most common expressions self.pi = thermodynamics.pi(state.parameters, self.rho_averaged, self.theta) self.T = thermodynamics.T(state.parameters, self.theta, self.pi, r_v=self.r_v) self.p = thermodynamics.p(state.parameters, self.pi) self.r_l = self.r_c + self.rain self.r_t = self.r_v + self.r_c + self.rain
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)
class ThermodynamicDiagnostic(DiagnosticField): name = "thermodynamic_diagnostic" def setup(self, state): if not self._initialised: space = state.fields("theta").function_space() broken_space = FunctionSpace(state.mesh, BrokenElement(space.ufl_element())) h_deg = space.ufl_element().degree()[0] v_deg = space.ufl_element().degree()[1] - 1 boundary_method = Boundary_Method.physics if ( v_deg == 0 and h_deg == 0) else None super().setup(state, space=space) # now let's attach all of our fields self.u = state.fields("u") self.rho = state.fields("rho") self.theta = state.fields("theta") self.rho_averaged = Function(space) self.recoverer = Recoverer(self.rho, self.rho_averaged, VDG=broken_space, boundary_method=boundary_method) try: self.r_v = state.fields("vapour_mixing_ratio") except NotImplementedError: self.r_v = Constant(0.0) try: self.r_c = state.fields("cloud_liquid_mixing_ratio") except NotImplementedError: self.r_c = Constant(0.0) try: self.rain = state.fields("rain_mixing_ratio") except NotImplementedError: self.rain = Constant(0.0) # now let's store the most common expressions self.exner = thermodynamics.exner_pressure(state.parameters, self.rho_averaged, self.theta) self.T = thermodynamics.T(state.parameters, self.theta, self.exner, r_v=self.r_v) self.p = thermodynamics.p(state.parameters, self.exner) self.r_l = self.r_c + self.rain self.r_t = self.r_v + self.r_c + self.rain def compute(self, state): self.recoverer.project()
def unsaturated_hydrostatic_balance(state, theta_d, H, exner0=None, top=False, exner_boundary=Constant(1.0), max_outer_solve_count=40, max_inner_solve_count=20): """ Given vertical profiles for dry potential temperature and relative humidity compute hydrostatically balanced virtual potential temperature, dry density and water vapour profiles. The general strategy is to split up the solving into two steps: 1) finding rho to balance the theta profile 2) finding theta_v and r_v to get back theta_d and H We iteratively solve these steps until we (hopefully) converge to a solution. :arg state: The :class:`State` object. :arg theta_d: The initial dry potential temperature profile. :arg H: The relative humidity profile. :arg exner0: Optional function to put exner pressure into. :arg top: If True, set a boundary condition at the top, otherwise it will be at the bottom. :arg exner_boundary: The value of exner on the specified boundary. :arg max_outer_solve_count: Max number of iterations for outer loop of balance solver. :arg max_inner_solve_count: Max number of iterations for inner loop of balanace solver. """ theta0 = state.fields('theta') rho0 = state.fields('rho') mr_v0 = state.fields('vapour_mixing_ratio') # Calculate hydrostatic exner pressure Vt = theta0.function_space() Vr = rho0.function_space() R_d = state.parameters.R_d R_v = state.parameters.R_v epsilon = R_d / R_v VDG = state.spaces("DG") if any(deg > 2 for deg in VDG.ufl_element().degree()): logger.warning( "default quadrature degree most likely not sufficient for this degree element" ) # apply first guesses theta0.assign(theta_d * 1.01) mr_v0.assign(0.01) v_deg = Vr.ufl_element().degree()[1] if v_deg == 0: method = Boundary_Method.physics else: method = None rho_h = Function(Vr) rho_averaged = Function(Vt) Vt_broken = FunctionSpace(state.mesh, BrokenElement(Vt.ufl_element())) rho_recoverer = Recoverer(rho0, rho_averaged, VDG=Vt_broken, boundary_method=method) w_h = Function(Vt) delta = 1.0 # make expressions for determining mr_v0 exner = thermodynamics.exner_pressure(state.parameters, rho_averaged, theta0) p = thermodynamics.p(state.parameters, exner) T = thermodynamics.T(state.parameters, theta0, exner, mr_v0) r_v_expr = thermodynamics.r_v(state.parameters, H, T, p) # make expressions to evaluate residual exner_ev = thermodynamics.exner_pressure(state.parameters, rho_averaged, theta0) p_ev = thermodynamics.p(state.parameters, exner_ev) T_ev = thermodynamics.T(state.parameters, theta0, exner_ev, mr_v0) RH_ev = thermodynamics.RH(state.parameters, mr_v0, T_ev, p_ev) RH = Function(Vt) for i in range(max_outer_solve_count): # solve for rho with theta_vd and w_v guesses compressible_hydrostatic_balance(state, theta0, rho_h, top=top, exner_boundary=exner_boundary, mr_t=mr_v0, solve_for_rho=True) # damp solution rho0.assign(rho0 * (1 - delta) + delta * rho_h) # calculate averaged rho rho_recoverer.project() RH.assign(RH_ev) if errornorm(RH, H) < 1e-10: break # now solve for r_v for j in range(max_inner_solve_count): w_h.interpolate(r_v_expr) mr_v0.assign(mr_v0 * (1 - delta) + delta * w_h) # compute theta_vd theta0.assign(theta_d * (1 + mr_v0 / epsilon)) # test quality of solution by re-evaluating expression RH.assign(RH_ev) if errornorm(RH, H) < 1e-10: break if i == max_outer_solve_count: raise RuntimeError( 'Hydrostatic balance solve has not converged within %i' % i, 'iterations') if exner0 is not None: exner = thermodynamics.exner_pressure(state.parameters, rho0, theta0) exner0.interpolate(exner) # do one extra solve for rho compressible_hydrostatic_balance(state, theta0, rho0, top=top, exner_boundary=exner_boundary, mr_t=mr_v0, solve_for_rho=True)
def saturated_hydrostatic_balance(state, theta_e, mr_t, exner0=None, top=False, exner_boundary=Constant(1.0), max_outer_solve_count=40, max_theta_solve_count=5, max_inner_solve_count=3): """ Given a wet equivalent potential temperature, theta_e, and the total moisture content, mr_t, compute a hydrostatically balance virtual potential temperature, dry density and water vapour profile. The general strategy is to split up the solving into two steps: 1) finding rho to balance the theta profile 2) finding theta_v and r_v to get back theta_e and saturation We iteratively solve these steps until we (hopefully) converge to a solution. :arg state: The :class:`State` object. :arg theta_e: The initial wet equivalent potential temperature profile. :arg mr_t: The total water pseudo-mixing ratio profile. :arg exner0: Optional function to put exner pressure into. :arg top: If True, set a boundary condition at the top, otherwise it will be at the bottom. :arg exner_boundary: The value of exner on the specified boundary. :arg max_outer_solve_count: Max number of outer iterations for balance solver. :arg max_theta_solve_count: Max number of iterations for theta solver (middle part of solve). :arg max_inner_solve_count: Max number of iterations on the inner most loop for the water vapour solver. """ theta0 = state.fields('theta') rho0 = state.fields('rho') mr_v0 = state.fields('vapour_mixing_ratio') # Calculate hydrostatic exner pressure Vt = theta0.function_space() Vr = rho0.function_space() VDG = state.spaces("DG") if any(deg > 2 for deg in VDG.ufl_element().degree()): logger.warning( "default quadrature degree most likely not sufficient for this degree element" ) theta0.interpolate(theta_e) mr_v0.interpolate(mr_t) v_deg = Vr.ufl_element().degree()[1] if v_deg == 0: boundary_method = Boundary_Method.physics else: boundary_method = None rho_h = Function(Vr) Vt_broken = FunctionSpace(state.mesh, BrokenElement(Vt.ufl_element())) rho_averaged = Function(Vt) rho_recoverer = Recoverer(rho0, rho_averaged, VDG=Vt_broken, boundary_method=boundary_method) w_h = Function(Vt) theta_h = Function(Vt) theta_e_test = Function(Vt) delta = 0.8 # expressions for finding theta0 and mr_v0 from theta_e and mr_t exner = thermodynamics.exner_pressure(state.parameters, rho_averaged, theta0) p = thermodynamics.p(state.parameters, exner) T = thermodynamics.T(state.parameters, theta0, exner, mr_v0) r_v_expr = thermodynamics.r_sat(state.parameters, T, p) theta_e_expr = thermodynamics.theta_e(state.parameters, T, p, mr_v0, mr_t) for i in range(max_outer_solve_count): # solve for rho with theta_vd and w_v guesses compressible_hydrostatic_balance(state, theta0, rho_h, top=top, exner_boundary=exner_boundary, mr_t=mr_t, solve_for_rho=True) # damp solution rho0.assign(rho0 * (1 - delta) + delta * rho_h) theta_e_test.assign(theta_e_expr) if errornorm(theta_e_test, theta_e) < 1e-8: break # calculate averaged rho rho_recoverer.project() # now solve for r_v for j in range(max_theta_solve_count): theta_h.interpolate(theta_e / theta_e_expr * theta0) theta0.assign(theta0 * (1 - delta) + delta * theta_h) # break when close enough if errornorm(theta_e_test, theta_e) < 1e-6: break for k in range(max_inner_solve_count): w_h.interpolate(r_v_expr) mr_v0.assign(mr_v0 * (1 - delta) + delta * w_h) # break when close enough theta_e_test.assign(theta_e_expr) if errornorm(theta_e_test, theta_e) < 1e-6: break if i == max_outer_solve_count: raise RuntimeError( 'Hydrostatic balance solve has not converged within %i' % i, 'iterations') if exner0 is not None: exner = thermodynamics.exner(state.parameters, rho0, theta0) exner0.interpolate(exner) # do one extra solve for rho compressible_hydrostatic_balance(state, theta0, rho0, top=top, exner_boundary=exner_boundary, mr_t=mr_t, solve_for_rho=True)
def __init__(self, state, iterations=1): super().__init__(state) self.iterations = iterations # obtain our fields self.theta = state.fields('theta') self.water_v = state.fields('water_v') self.water_c = state.fields('water_c') rho = state.fields('rho') try: rain = state.fields('rain') water_l = self.water_c + rain except NotImplementedError: water_l = self.water_c # declare function space Vt = self.theta.function_space() # make rho variables # we recover rho into theta space if state.vertical_degree == 0 and state.horizontal_degree == 0: boundary_method = Boundary_Method.physics else: boundary_method = None Vt_broken = FunctionSpace(state.mesh, BrokenElement(Vt.ufl_element())) rho_averaged = Function(Vt) self.rho_recoverer = Recoverer(rho, rho_averaged, VDG=Vt_broken, boundary_method=boundary_method) # define some parameters as attributes dt = state.timestepping.dt R_d = state.parameters.R_d cp = state.parameters.cp cv = state.parameters.cv c_pv = state.parameters.c_pv c_pl = state.parameters.c_pl c_vv = state.parameters.c_vv R_v = state.parameters.R_v # make useful fields Pi = thermodynamics.pi(state.parameters, rho_averaged, self.theta) T = thermodynamics.T(state.parameters, self.theta, Pi, r_v=self.water_v) p = thermodynamics.p(state.parameters, Pi) L_v = thermodynamics.Lv(state.parameters, T) R_m = R_d + R_v * self.water_v c_pml = cp + c_pv * self.water_v + c_pl * water_l c_vml = cv + c_vv * self.water_v + c_pl * water_l # use Teten's formula to calculate w_sat w_sat = thermodynamics.r_sat(state.parameters, T, p) # make appropriate condensation rate dot_r_cond = ((self.water_v - w_sat) / (dt * (1.0 + ((L_v**2.0 * w_sat) / (cp * R_v * T**2.0))))) # make cond_rate function, that needs to be the same for all updates in one time step cond_rate = Function(Vt) # adjust cond rate so negative concentrations don't occur self.lim_cond_rate = Interpolator( conditional(dot_r_cond < 0, max_value(dot_r_cond, -self.water_c / dt), min_value(dot_r_cond, self.water_v / dt)), cond_rate) # tell the prognostic fields what to update to self.water_v_new = Interpolator(self.water_v - dt * cond_rate, Vt) self.water_c_new = Interpolator(self.water_c + dt * cond_rate, Vt) self.theta_new = Interpolator( self.theta * (1.0 + dt * cond_rate * (cv * L_v / (c_vml * cp * T) - R_v * cv * c_pml / (R_m * cp * c_vml))), Vt)
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
def __init__(self, state): super().__init__(state) # obtain our fields self.theta = state.fields('theta') self.water_v = state.fields('water_v') self.rain = state.fields('rain') rho = state.fields('rho') try: water_c = state.fields('water_c') water_l = self.rain + water_c except NotImplementedError: water_l = self.rain # declare function space Vt = self.theta.function_space() # make rho variables # we recover rho into theta space if state.vertical_degree == 0 and state.horizontal_degree == 0: boundary_method = Boundary_Method.physics else: boundary_method = None Vt_broken = FunctionSpace(state.mesh, BrokenElement(Vt.ufl_element())) rho_averaged = Function(Vt) self.rho_recoverer = Recoverer(rho, rho_averaged, VDG=Vt_broken, boundary_method=boundary_method) # define some parameters as attributes dt = state.timestepping.dt R_d = state.parameters.R_d cp = state.parameters.cp cv = state.parameters.cv c_pv = state.parameters.c_pv c_pl = state.parameters.c_pl c_vv = state.parameters.c_vv R_v = state.parameters.R_v # make useful fields Pi = thermodynamics.pi(state.parameters, rho_averaged, self.theta) T = thermodynamics.T(state.parameters, self.theta, Pi, r_v=self.water_v) p = thermodynamics.p(state.parameters, Pi) L_v = thermodynamics.Lv(state.parameters, T) R_m = R_d + R_v * self.water_v c_pml = cp + c_pv * self.water_v + c_pl * water_l c_vml = cv + c_vv * self.water_v + c_pl * water_l # use Teten's formula to calculate w_sat w_sat = thermodynamics.r_sat(state.parameters, T, p) # expression for ventilation factor a = Constant(1.6) b = Constant(124.9) c = Constant(0.2046) C = a + b * (rho_averaged * self.rain)**c # make appropriate condensation rate f = Constant(5.4e5) g = Constant(2.55e6) h = Constant(0.525) dot_r_evap = (((1 - self.water_v / w_sat) * C * (rho_averaged * self.rain)**h) / (rho_averaged * (f + g / (p * w_sat)))) # make evap_rate function, needs to be the same for all updates in one time step evap_rate = Function(Vt) # adjust evap rate so negative rain doesn't occur self.lim_evap_rate = Interpolator( conditional( dot_r_evap < 0, 0.0, conditional(self.rain < 0.0, 0.0, min_value(dot_r_evap, self.rain / dt))), evap_rate) # tell the prognostic fields what to update to self.water_v_new = Interpolator(self.water_v + dt * evap_rate, Vt) self.rain_new = Interpolator(self.rain - dt * evap_rate, Vt) self.theta_new = Interpolator( self.theta * (1.0 - dt * evap_rate * (cv * L_v / (c_vml * cp * T) - R_v * cv * c_pml / (R_m * cp * c_vml))), Vt)
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)
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
class Condensation(Physics): """ The process of condensation of water vapour into liquid water and evaporation of liquid water into water vapour, with the associated latent heat changes. The parametrization follows that used in Bryan and Fritsch (2002). :arg state: :class:`.State.` object. :arg iterations: number of iterations to do of condensation scheme per time step. """ def __init__(self, state, iterations=1): super().__init__(state) self.iterations = iterations # obtain our fields self.theta = state.fields('theta') self.water_v = state.fields('vapour_mixing_ratio') self.water_c = state.fields('cloud_liquid_mixing_ratio') rho = state.fields('rho') try: # TODO: use the phase flag for the tracers here rain = state.fields('rain_mixing_ratio') water_l = self.water_c + rain except NotImplementedError: water_l = self.water_c # declare function space Vt = self.theta.function_space() # make rho variables # we recover rho into theta space h_deg = rho.function_space().ufl_element().degree()[0] v_deg = rho.function_space().ufl_element().degree()[1] if v_deg == 0 and h_deg == 0: boundary_method = Boundary_Method.physics else: boundary_method = None Vt_broken = FunctionSpace(state.mesh, BrokenElement(Vt.ufl_element())) rho_averaged = Function(Vt) self.rho_recoverer = Recoverer(rho, rho_averaged, VDG=Vt_broken, boundary_method=boundary_method) # define some parameters as attributes dt = state.dt R_d = state.parameters.R_d cp = state.parameters.cp cv = state.parameters.cv c_pv = state.parameters.c_pv c_pl = state.parameters.c_pl c_vv = state.parameters.c_vv R_v = state.parameters.R_v # make useful fields exner = thermodynamics.exner_pressure(state.parameters, rho_averaged, self.theta) T = thermodynamics.T(state.parameters, self.theta, exner, r_v=self.water_v) p = thermodynamics.p(state.parameters, exner) L_v = thermodynamics.Lv(state.parameters, T) R_m = R_d + R_v * self.water_v c_pml = cp + c_pv * self.water_v + c_pl * water_l c_vml = cv + c_vv * self.water_v + c_pl * water_l # use Teten's formula to calculate w_sat w_sat = thermodynamics.r_sat(state.parameters, T, p) # make appropriate condensation rate dot_r_cond = ((self.water_v - w_sat) / (dt * (1.0 + ((L_v ** 2.0 * w_sat) / (cp * R_v * T ** 2.0))))) # make cond_rate function, that needs to be the same for all updates in one time step cond_rate = Function(Vt) # adjust cond rate so negative concentrations don't occur self.lim_cond_rate = Interpolator(conditional(dot_r_cond < 0, max_value(dot_r_cond, - self.water_c / dt), min_value(dot_r_cond, self.water_v / dt)), cond_rate) # tell the prognostic fields what to update to self.water_v_new = Interpolator(self.water_v - dt * cond_rate, Vt) self.water_c_new = Interpolator(self.water_c + dt * cond_rate, Vt) self.theta_new = Interpolator(self.theta * (1.0 + dt * cond_rate * (cv * L_v / (c_vml * cp * T) - R_v * cv * c_pml / (R_m * cp * c_vml))), Vt) def apply(self): self.rho_recoverer.project() for i in range(self.iterations): self.lim_cond_rate.interpolate() self.theta.assign(self.theta_new.interpolate()) self.water_v.assign(self.water_v_new.interpolate()) self.water_c.assign(self.water_c_new.interpolate())
class Evaporation(Physics): """ The process of evaporation of rain into water vapour with the associated latent heat change. This parametrization comes from Klemp and Wilhelmson (1978). :arg state: :class:`.State.` object. """ def __init__(self, state): super().__init__(state) # obtain our fields self.theta = state.fields('theta') self.water_v = state.fields('vapour_mixing_ratio') self.rain = state.fields('rain_mixing_ratio') rho = state.fields('rho') try: water_c = state.fields('cloud_liquid_mixing_ratio') water_l = self.rain + water_c except NotImplementedError: water_l = self.rain # declare function space Vt = self.theta.function_space() # make rho variables # we recover rho into theta space h_deg = rho.function_space().ufl_element().degree()[0] v_deg = rho.function_space().ufl_element().degree()[1] if v_deg == 0 and h_deg == 0: boundary_method = Boundary_Method.physics else: boundary_method = None Vt_broken = FunctionSpace(state.mesh, BrokenElement(Vt.ufl_element())) rho_averaged = Function(Vt) self.rho_recoverer = Recoverer(rho, rho_averaged, VDG=Vt_broken, boundary_method=boundary_method) # define some parameters as attributes dt = state.dt R_d = state.parameters.R_d cp = state.parameters.cp cv = state.parameters.cv c_pv = state.parameters.c_pv c_pl = state.parameters.c_pl c_vv = state.parameters.c_vv R_v = state.parameters.R_v # make useful fields exner = thermodynamics.exner_pressure(state.parameters, rho_averaged, self.theta) T = thermodynamics.T(state.parameters, self.theta, exner, r_v=self.water_v) p = thermodynamics.p(state.parameters, exner) L_v = thermodynamics.Lv(state.parameters, T) R_m = R_d + R_v * self.water_v c_pml = cp + c_pv * self.water_v + c_pl * water_l c_vml = cv + c_vv * self.water_v + c_pl * water_l # use Teten's formula to calculate w_sat w_sat = thermodynamics.r_sat(state.parameters, T, p) # expression for ventilation factor a = Constant(1.6) b = Constant(124.9) c = Constant(0.2046) C = a + b * (rho_averaged * self.rain) ** c # make appropriate condensation rate f = Constant(5.4e5) g = Constant(2.55e6) h = Constant(0.525) dot_r_evap = (((1 - self.water_v / w_sat) * C * (rho_averaged * self.rain) ** h) / (rho_averaged * (f + g / (p * w_sat)))) # make evap_rate function, needs to be the same for all updates in one time step evap_rate = Function(Vt) # adjust evap rate so negative rain doesn't occur self.lim_evap_rate = Interpolator(conditional(dot_r_evap < 0, 0.0, conditional(self.rain < 0.0, 0.0, min_value(dot_r_evap, self.rain / dt))), evap_rate) # tell the prognostic fields what to update to self.water_v_new = Interpolator(self.water_v + dt * evap_rate, Vt) self.rain_new = Interpolator(self.rain - dt * evap_rate, Vt) self.theta_new = Interpolator(self.theta * (1.0 - dt * evap_rate * (cv * L_v / (c_vml * cp * T) - R_v * cv * c_pml / (R_m * cp * c_vml))), Vt) def apply(self): self.rho_recoverer.project() self.lim_evap_rate.interpolate() self.theta.assign(self.theta_new.interpolate()) self.water_v.assign(self.water_v_new.interpolate()) self.rain.assign(self.rain_new.interpolate())