def setup(self, state): if not self._initialised: space = state.spaces("Vv") super().setup(state, space=space) rho = state.fields("rho") rhobar = state.fields("rhobar") theta = state.fields("theta") thetabar = state.fields("thetabar") pi = thermodynamics.pi(state.parameters, rho, theta) pibar = thermodynamics.pi(state.parameters, rhobar, thetabar) cp = Constant(state.parameters.cp) n = FacetNormal(state.mesh) F = TrialFunction(space) w = TestFunction(space) a = inner(w, F)*dx L = (- cp*div((theta-thetabar)*w)*pibar*dx + cp*jump((theta-thetabar)*w, n)*avg(pibar)*dS_v - cp*div(thetabar*w)*(pi-pibar)*dx + cp*jump(thetabar*w, n)*avg(pi-pibar)*dS_v) bcs = [DirichletBC(space, 0.0, "bottom"), DirichletBC(space, 0.0, "top")] imbalanceproblem = LinearVariationalProblem(a, L, self.field, bcs=bcs) self.imbalance_solver = LinearVariationalSolver(imbalanceproblem)
def calculate_Pi0(state, theta0, rho0): # exner function Vr = rho0.function_space() Pi = Function(Vr).interpolate(thermodynamics.pi(state.parameters, rho0, theta0)) Pi0 = assemble(Pi*dx)/assemble(Constant(1)*dx(domain=state.mesh)) return Pi0
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 compressible_eady_initial_v(state, theta0, rho0, v): f = state.parameters.f cp = state.parameters.cp # exner function Vr = rho0.function_space() Pi = Function(Vr).interpolate(thermodynamics.pi(state.parameters, rho0, theta0)) # get Pi gradient Vu = state.spaces("HDiv") g = TrialFunction(Vu) wg = TestFunction(Vu) n = FacetNormal(state.mesh) a = inner(wg, g)*dx L = -div(wg)*Pi*dx + inner(wg, n)*Pi*ds_tb pgrad = Function(Vu) solve(a == L, pgrad) # get initial v m = TrialFunction(Vr) phi = TestFunction(Vr) a = phi*f*m*dx L = phi*cp*theta0*pgrad[0]*dx solve(a == L, v) return v
def compute(self, state): theta = state.fields('theta') rho = state.fields('rho') w_v = state.fields('water_v') w_c = state.fields('water_c') pi = thermodynamics.pi(state.parameters, rho, theta) T = thermodynamics.T(state.parameters, theta, pi, r_v=w_v) return self.field.interpolate(thermodynamics.internal_energy(state.parameters, rho, T, r_v=w_v, r_l=w_c))
def compute(self, state): theta = state.fields('theta') rho = state.fields('rho') w_v = state.fields('water_v') w_c = state.fields('water_c') w_t = w_c + w_v pi = thermodynamics.pi(state.parameters, rho, theta) p = thermodynamics.p(state.parameters, pi) T = thermodynamics.T(state.parameters, theta, pi, r_v=w_v) return self.field.interpolate(thermodynamics.theta_e(state.parameters, T, p, w_v, w_t))
def compute(self, state): x, y, z = SpatialCoordinate(state.mesh) g = state.parameters.g cp = state.parameters.cp cv = state.parameters.cv Pi0 = state.parameters.Pi0 rho = state.fields("rho") theta = state.fields("theta") Pi = thermodynamics.pi(state.parameters, rho, theta) potential = rho * (g * z + cv * Pi * theta - cp * Pi0 * theta) return self.field.interpolate(potential)
def forcing_term(self): # L = super(EadyForcing, self).forcing_term() L = Forcing.forcing_term(self) dthetady = self.state.parameters.dthetady Pi0 = self.state.parameters.Pi0 cp = self.state.parameters.cp _, rho0, theta0 = split(self.x0) Pi = thermodynamics.pi(self.state.parameters, rho0, theta0) Pi_0 = Constant(Pi0) L += self.scaling * cp * dthetady * (Pi - Pi_0) * inner( self.test, as_vector([0., 1., 0.])) * dx # Eady forcing return L
def pressure_gradient_term(self): u0, rho0, theta0 = split(self.x0) cp = self.state.parameters.cp n = FacetNormal(self.state.mesh) Vtheta = self.state.spaces("HDiv_v") # introduce new theta so it can be changed by moisture theta = theta0 # add effect of density of water upon theta if self.moisture is not None: water_t = Function(Vtheta).assign(0.0) for water in self.moisture: water_t += self.state.fields(water) theta = theta / (1 + water_t) pi = thermodynamics.pi(self.state.parameters, rho0, theta0) L = (+cp * div(theta * self.test) * pi * dx - cp * jump(self.test * theta, n) * avg(pi) * dS_v) return L
def compute(self, state): rho = state.fields(self.rho_name) theta = state.fields(self.theta_name) Pi = thermodynamics.pi(state.parameters, rho, theta) return self.field.interpolate(Pi)
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)
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_solver(self): import numpy as np state = self.state Dt = state.timestepping.dt beta_ = Dt*state.timestepping.alpha cp = state.parameters.cp mu = state.mu Vu = state.spaces("HDiv") Vu_broken = FunctionSpace(state.mesh, BrokenElement(Vu.ufl_element())) Vtheta = state.spaces("HDiv_v") Vrho = state.spaces("DG") # Store time-stepping coefficients as UFL Constants dt = Constant(Dt) beta = Constant(beta_) beta_cp = Constant(beta_ * cp) h_deg = state.horizontal_degree v_deg = state.vertical_degree Vtrace = FunctionSpace(state.mesh, "HDiv Trace", degree=(h_deg, v_deg)) # Split up the rhs vector (symbolically) u_in, rho_in, theta_in = split(state.xrhs) # Build the function space for "broken" u, rho, and pressure trace M = MixedFunctionSpace((Vu_broken, Vrho, Vtrace)) w, phi, dl = TestFunctions(M) u, rho, l0 = TrialFunctions(M) n = FacetNormal(state.mesh) # Get background fields thetabar = state.fields("thetabar") rhobar = state.fields("rhobar") pibar = thermodynamics.pi(state.parameters, rhobar, thetabar) pibar_rho = thermodynamics.pi_rho(state.parameters, rhobar, thetabar) pibar_theta = thermodynamics.pi_theta(state.parameters, rhobar, thetabar) # Analytical (approximate) elimination of theta k = state.k # Upward pointing unit vector theta = -dot(k, u)*dot(k, grad(thetabar))*beta + theta_in # Only include theta' (rather than pi') in the vertical # component of the gradient # The pi prime term (here, bars are for mean and no bars are # for linear perturbations) pi = pibar_theta*theta + pibar_rho*rho # Vertical projection def V(u): return k*inner(u, k) # Specify degree for some terms as estimated degree is too large dxp = dx(degree=(self.quadrature_degree)) dS_vp = dS_v(degree=(self.quadrature_degree)) dS_hp = dS_h(degree=(self.quadrature_degree)) ds_vp = ds_v(degree=(self.quadrature_degree)) ds_tbp = (ds_t(degree=(self.quadrature_degree)) + ds_b(degree=(self.quadrature_degree))) # Add effect of density of water upon theta if self.moisture is not None: water_t = Function(Vtheta).assign(0.0) for water in self.moisture: water_t += self.state.fields(water) theta_w = theta / (1 + water_t) thetabar_w = thetabar / (1 + water_t) else: theta_w = theta thetabar_w = thetabar _l0 = TrialFunction(Vtrace) _dl = TestFunction(Vtrace) a_tr = _dl('+')*_l0('+')*(dS_vp + dS_hp) + _dl*_l0*ds_vp + _dl*_l0*ds_tbp def L_tr(f): return _dl('+')*avg(f)*(dS_vp + dS_hp) + _dl*f*ds_vp + _dl*f*ds_tbp cg_ilu_parameters = {'ksp_type': 'cg', 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'} # Project field averages into functions on the trace space rhobar_avg = Function(Vtrace) pibar_avg = Function(Vtrace) rho_avg_prb = LinearVariationalProblem(a_tr, L_tr(rhobar), rhobar_avg) pi_avg_prb = LinearVariationalProblem(a_tr, L_tr(pibar), pibar_avg) rho_avg_solver = LinearVariationalSolver(rho_avg_prb, solver_parameters=cg_ilu_parameters, options_prefix='rhobar_avg_solver') pi_avg_solver = LinearVariationalSolver(pi_avg_prb, solver_parameters=cg_ilu_parameters, options_prefix='pibar_avg_solver') with timed_region("Gusto:HybridProjectRhobar"): rho_avg_solver.solve() with timed_region("Gusto:HybridProjectPibar"): pi_avg_solver.solve() # "broken" u, rho, and trace system # NOTE: no ds_v integrals since equations are defined on # a periodic (or sphere) base mesh. eqn = ( # momentum equation inner(w, (state.h_project(u) - u_in))*dx - beta_cp*div(theta_w*V(w))*pibar*dxp # following does nothing but is preserved in the comments # to remind us why (because V(w) is purely vertical). # + beta_cp*jump(theta_w*V(w), n=n)*pibar_avg('+')*dS_vp + beta_cp*jump(theta_w*V(w), n=n)*pibar_avg('+')*dS_hp + beta_cp*dot(theta_w*V(w), n)*pibar_avg*ds_tbp - beta_cp*div(thetabar_w*w)*pi*dxp # trace terms appearing after integrating momentum equation + beta_cp*jump(thetabar_w*w, n=n)*l0('+')*(dS_vp + dS_hp) + beta_cp*dot(thetabar_w*w, n)*l0*(ds_tbp + ds_vp) # mass continuity equation + (phi*(rho - rho_in) - beta*inner(grad(phi), u)*rhobar)*dx + beta*jump(phi*u, n=n)*rhobar_avg('+')*(dS_v + dS_h) # term added because u.n=0 is enforced weakly via the traces + beta*phi*dot(u, n)*rhobar_avg*(ds_tb + ds_v) # constraint equation to enforce continuity of the velocity # through the interior facets and weakly impose the no-slip # condition + dl('+')*jump(u, n=n)*(dS_vp + dS_hp) + dl*dot(u, n)*(ds_tbp + ds_vp) ) # contribution of the sponge term if mu is not None: eqn += dt*mu*inner(w, k)*inner(u, k)*dx aeqn = lhs(eqn) Leqn = rhs(eqn) # Function for the hybridized solutions self.urhol0 = Function(M) hybridized_prb = LinearVariationalProblem(aeqn, Leqn, self.urhol0) hybridized_solver = LinearVariationalSolver(hybridized_prb, solver_parameters=self.solver_parameters, options_prefix='ImplicitSolver') self.hybridized_solver = hybridized_solver # Project broken u into the HDiv space using facet averaging. # Weight function counting the dofs of the HDiv element: shapes = {"i": Vu.finat_element.space_dimension(), "j": np.prod(Vu.shape, dtype=int)} weight_kernel = """ for (int i=0; i<{i}; ++i) for (int j=0; j<{j}; ++j) w[i*{j} + j] += 1.0; """.format(**shapes) self._weight = Function(Vu) par_loop(weight_kernel, dx, {"w": (self._weight, INC)}) # Averaging kernel self._average_kernel = """ for (int i=0; i<{i}; ++i) for (int j=0; j<{j}; ++j) vec_out[i*{j} + j] += vec_in[i*{j} + j]/w[i*{j} + j]; """.format(**shapes) # HDiv-conforming velocity self.u_hdiv = Function(Vu) # Reconstruction of theta theta = TrialFunction(Vtheta) gamma = TestFunction(Vtheta) self.theta = Function(Vtheta) theta_eqn = gamma*(theta - theta_in + dot(k, self.u_hdiv)*dot(k, grad(thetabar))*beta)*dx theta_problem = LinearVariationalProblem(lhs(theta_eqn), rhs(theta_eqn), self.theta) self.theta_solver = LinearVariationalSolver(theta_problem, solver_parameters=cg_ilu_parameters, options_prefix='thetabacksubstitution') # Store boundary conditions for the div-conforming velocity to apply # post-solve self.bcs = self.state.bcs
# Define constant theta_e and water_t Tsurf = 320.0 total_water = 0.02 theta_e = Function(Vt).assign(Tsurf) water_t = Function(Vt).assign(total_water) # Calculate hydrostatic fields saturated_hydrostatic_balance(state, theta_e, water_t) # make mean fields theta_b = Function(Vt).assign(theta0) rho_b = Function(Vr).assign(rho0) water_vb = Function(Vt).assign(water_v0) water_cb = Function(Vt).assign(water_t - water_vb) pibar = thermodynamics.pi(state.parameters, rho_b, theta_b) Tb = thermodynamics.T(state.parameters, theta_b, pibar, r_v=water_vb) Ibar = state.fields("InternalEnergybar", Vt, dump=False) Ibar.interpolate( thermodynamics.internal_energy(state.parameters, rho_b, Tb, r_v=water_vb, r_l=water_cb)) # define perturbation xc = L / 2 zc = 2000. rc = 2000. Tdash = 2.0 r = sqrt((x - xc)**2 + (z - zc)**2)
def compressible_hydrostatic_balance(state, theta0, rho0, pi0=None, top=False, pi_boundary=Constant(1.0), water_t=None, solve_for_rho=False, params=None): """ Compute a hydrostatically balanced density given a potential temperature profile. By default, this uses a vertically-oriented hybridization procedure for solving the resulting discrete systems. :arg state: The :class:`State` object. :arg theta0: :class:`.Function`containing the potential temperature. :arg rho0: :class:`.Function` to write the initial density into. :arg top: If True, set a boundary condition at the top. Otherwise, set it at the bottom. :arg pi_boundary: a field or expression to use as boundary data for pi on the top or bottom as specified. :arg water_t: the initial total water mixing ratio field. """ # Calculate hydrostatic Pi VDG = state.spaces("DG") Vv = state.spaces("Vv") W = MixedFunctionSpace((Vv, VDG)) v, pi = TrialFunctions(W) dv, dpi = TestFunctions(W) n = FacetNormal(state.mesh) cp = state.parameters.cp # add effect of density of water upon theta theta = theta0 if water_t is not None: theta = theta0 / (1 + water_t) alhs = ( (cp*inner(v, dv) - cp*div(dv*theta)*pi)*dx + dpi*div(theta*v)*dx ) if top: bmeasure = ds_t bstring = "bottom" else: bmeasure = ds_b bstring = "top" arhs = -cp*inner(dv, n)*theta*pi_boundary*bmeasure # Possibly make g vary with spatial coordinates? g = state.parameters.g arhs -= g*inner(dv, state.k)*dx bcs = [DirichletBC(W.sub(0), Constant(0.0), bstring)] w = Function(W) PiProblem = LinearVariationalProblem(alhs, arhs, w, bcs=bcs) if params is None: params = {'ksp_type': 'preonly', 'pc_type': 'python', 'mat_type': 'matfree', 'pc_python_type': 'gusto.VerticalHybridizationPC', # Vertical trace system is only coupled vertically in columns # block ILU is a direct solver! 'vert_hybridization': {'ksp_type': 'preonly', 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'}} PiSolver = LinearVariationalSolver(PiProblem, solver_parameters=params, options_prefix="pisolver") PiSolver.solve() v, Pi = w.split() if pi0 is not None: pi0.assign(Pi) if solve_for_rho: w1 = Function(W) v, rho = w1.split() rho.interpolate(thermodynamics.rho(state.parameters, theta0, Pi)) v, rho = split(w1) dv, dpi = TestFunctions(W) pi = thermodynamics.pi(state.parameters, rho, theta0) F = ( (cp*inner(v, dv) - cp*div(dv*theta)*pi)*dx + dpi*div(theta0*v)*dx + cp*inner(dv, n)*theta*pi_boundary*bmeasure ) F += g*inner(dv, state.k)*dx rhoproblem = NonlinearVariationalProblem(F, w1, bcs=bcs) rhosolver = NonlinearVariationalSolver(rhoproblem, solver_parameters=params, options_prefix="rhosolver") rhosolver.solve() v, rho_ = w1.split() rho0.assign(rho_) else: rho0.interpolate(thermodynamics.rho(state.parameters, theta0, Pi))
def unsaturated_hydrostatic_balance(state, theta_d, H, pi0=None, top=False, pi_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 pi0: 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 pi_boundary: The value of pi 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') water_v0 = state.fields('water_v') # Calculate hydrostatic Pi 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) water_v0.assign(0.01) if state.vertical_degree == 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 water_v0 pie = thermodynamics.pi(state.parameters, rho_averaged, theta0) p = thermodynamics.p(state.parameters, pie) T = thermodynamics.T(state.parameters, theta0, pie, water_v0) r_v_expr = thermodynamics.r_v(state.parameters, H, T, p) # make expressions to evaluate residual pi_ev = thermodynamics.pi(state.parameters, rho_averaged, theta0) p_ev = thermodynamics.p(state.parameters, pi_ev) T_ev = thermodynamics.T(state.parameters, theta0, pi_ev, water_v0) RH_ev = thermodynamics.RH(state.parameters, water_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, pi_boundary=pi_boundary, water_t=water_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) water_v0.assign(water_v0 * (1 - delta) + delta * w_h) # compute theta_vd theta0.assign(theta_d * (1 + water_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 pi0 is not None: pie = thermodynamics.pi(state.parameters, rho0, theta0) pi0.interpolate(pie) # do one extra solve for rho compressible_hydrostatic_balance(state, theta0, rho0, top=top, pi_boundary=pi_boundary, water_t=water_v0, solve_for_rho=True)
def saturated_hydrostatic_balance(state, theta_e, water_t, pi0=None, top=False, pi_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, water_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 water_t: The total water pseudo-mixing ratio profile. :arg pi0: 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 pi_boundary: The value of pi 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') water_v0 = state.fields('water_v') # Calculate hydrostatic Pi 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) water_v0.interpolate(water_t) if state.vertical_degree == 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 water_v0 from theta_e and water_t pie = thermodynamics.pi(state.parameters, rho_averaged, theta0) p = thermodynamics.p(state.parameters, pie) T = thermodynamics.T(state.parameters, theta0, pie, water_v0) r_v_expr = thermodynamics.r_sat(state.parameters, T, p) theta_e_expr = thermodynamics.theta_e(state.parameters, T, p, water_v0, water_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, pi_boundary=pi_boundary, water_t=water_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) water_v0.assign(water_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 pi0 is not None: pie = thermodynamics.pi(state.parameters, rho0, theta0) pi0.interpolate(pie) # do one extra solve for rho compressible_hydrostatic_balance(state, theta0, rho0, top=top, pi_boundary=pi_boundary, water_t=water_t, solve_for_rho=True)
def _setup_solver(self): state = self.state # just cutting down line length a bit Dt = state.timestepping.dt beta_ = Dt*state.timestepping.alpha cp = state.parameters.cp mu = state.mu Vu = state.spaces("HDiv") Vtheta = state.spaces("HDiv_v") Vrho = state.spaces("DG") # Store time-stepping coefficients as UFL Constants dt = Constant(Dt) beta = Constant(beta_) beta_cp = Constant(beta_ * cp) # Split up the rhs vector (symbolically) u_in, rho_in, theta_in = split(state.xrhs) # Build the reduced function space for u,rho M = MixedFunctionSpace((Vu, Vrho)) w, phi = TestFunctions(M) u, rho = TrialFunctions(M) n = FacetNormal(state.mesh) # Get background fields thetabar = state.fields("thetabar") rhobar = state.fields("rhobar") pibar = thermodynamics.pi(state.parameters, rhobar, thetabar) pibar_rho = thermodynamics.pi_rho(state.parameters, rhobar, thetabar) pibar_theta = thermodynamics.pi_theta(state.parameters, rhobar, thetabar) # Analytical (approximate) elimination of theta k = state.k # Upward pointing unit vector theta = -dot(k, u)*dot(k, grad(thetabar))*beta + theta_in # Only include theta' (rather than pi') in the vertical # component of the gradient # the pi prime term (here, bars are for mean and no bars are # for linear perturbations) pi = pibar_theta*theta + pibar_rho*rho # vertical projection def V(u): return k*inner(u, k) # specify degree for some terms as estimated degree is too large dxp = dx(degree=(self.quadrature_degree)) dS_vp = dS_v(degree=(self.quadrature_degree)) # add effect of density of water upon theta if self.moisture is not None: water_t = Function(Vtheta).assign(0.0) for water in self.moisture: water_t += self.state.fields(water) theta_w = theta / (1 + water_t) thetabar_w = thetabar / (1 + water_t) else: theta_w = theta thetabar_w = thetabar eqn = ( inner(w, (state.h_project(u) - u_in))*dx - beta_cp*div(theta_w*V(w))*pibar*dxp # following does nothing but is preserved in the comments # to remind us why (because V(w) is purely vertical). # + beta_cp*jump(theta*V(w), n)*avg(pibar)*dS_v - beta_cp*div(thetabar_w*w)*pi*dxp + beta_cp*jump(thetabar_w*w, n)*avg(pi)*dS_vp + (phi*(rho - rho_in) - beta*inner(grad(phi), u)*rhobar)*dx + beta*jump(phi*u, n)*avg(rhobar)*(dS_v + dS_h) ) if mu is not None: eqn += dt*mu*inner(w, k)*inner(u, k)*dx aeqn = lhs(eqn) Leqn = rhs(eqn) # Place to put result of u rho solver self.urho = Function(M) # Boundary conditions (assumes extruded mesh) bcs = [DirichletBC(M.sub(0), 0.0, "bottom"), DirichletBC(M.sub(0), 0.0, "top")] # Solver for u, rho urho_problem = LinearVariationalProblem( aeqn, Leqn, self.urho, bcs=bcs) self.urho_solver = LinearVariationalSolver(urho_problem, solver_parameters=self.solver_parameters, options_prefix='ImplicitSolver') # Reconstruction of theta theta = TrialFunction(Vtheta) gamma = TestFunction(Vtheta) u, rho = self.urho.split() self.theta = Function(Vtheta) theta_eqn = gamma*(theta - theta_in + dot(k, u)*dot(k, grad(thetabar))*beta)*dx theta_problem = LinearVariationalProblem(lhs(theta_eqn), rhs(theta_eqn), self.theta) self.theta_solver = LinearVariationalSolver(theta_problem, options_prefix='thetabacksubstitution')
def compressible_hydrostatic_balance(state, theta0, rho0, pi0=None, top=False, pi_boundary=Constant(1.0), water_t=None, solve_for_rho=False, params=None): """ Compute a hydrostatically balanced density given a potential temperature profile. :arg state: The :class:`State` object. :arg theta0: :class:`.Function`containing the potential temperature. :arg rho0: :class:`.Function` to write the initial density into. :arg top: If True, set a boundary condition at the top. Otherwise, set it at the bottom. :arg pi_boundary: a field or expression to use as boundary data for pi on the top or bottom as specified. :arg water_t: the initial total water mixing ratio field. """ # Calculate hydrostatic Pi VDG = state.spaces("DG") Vv = state.spaces("Vv") W = MixedFunctionSpace((Vv, VDG)) v, pi = TrialFunctions(W) dv, dpi = TestFunctions(W) n = FacetNormal(state.mesh) cp = state.parameters.cp # add effect of density of water upon theta theta = theta0 if water_t is not None: theta = theta0 / (1 + water_t) alhs = ( (cp*inner(v, dv) - cp*div(dv*theta)*pi)*dx + dpi*div(theta*v)*dx ) if top: bmeasure = ds_t bstring = "bottom" else: bmeasure = ds_b bstring = "top" arhs = -cp*inner(dv, n)*theta*pi_boundary*bmeasure g = state.parameters.g arhs -= g*inner(dv, state.k)*dx bcs = [DirichletBC(W.sub(0), 0.0, bstring)] w = Function(W) PiProblem = LinearVariationalProblem(alhs, arhs, w, bcs=bcs) if params is None: params = {'pc_type': 'fieldsplit', 'pc_fieldsplit_type': 'schur', 'ksp_type': 'gmres', 'ksp_monitor_true_residual': True, 'ksp_max_it': 100, 'ksp_gmres_restart': 50, 'pc_fieldsplit_schur_fact_type': 'FULL', 'pc_fieldsplit_schur_precondition': 'selfp', 'fieldsplit_0_ksp_type': 'richardson', 'fieldsplit_0_ksp_max_it': 5, 'fieldsplit_0_pc_type': 'gamg', 'fieldsplit_1_pc_gamg_sym_graph': True, 'fieldsplit_1_mg_levels_ksp_type': 'chebyshev', 'fieldsplit_1_mg_levels_ksp_chebyshev_esteig': True, 'fieldsplit_1_mg_levels_ksp_max_it': 5, 'fieldsplit_1_mg_levels_pc_type': 'bjacobi', 'fieldsplit_1_mg_levels_sub_pc_type': 'ilu'} PiSolver = LinearVariationalSolver(PiProblem, solver_parameters=params) PiSolver.solve() v, Pi = w.split() if pi0 is not None: pi0.assign(Pi) if solve_for_rho: w1 = Function(W) v, rho = w1.split() rho.interpolate(thermodynamics.rho(state.parameters, theta0, Pi)) v, rho = split(w1) dv, dpi = TestFunctions(W) pi = thermodynamics.pi(state.parameters, rho, theta0) F = ( (cp*inner(v, dv) - cp*div(dv*theta)*pi)*dx + dpi*div(theta0*v)*dx + cp*inner(dv, n)*theta*pi_boundary*bmeasure ) F += g*inner(dv, state.k)*dx rhoproblem = NonlinearVariationalProblem(F, w1, bcs=bcs) rhosolver = NonlinearVariationalSolver(rhoproblem, solver_parameters=params) rhosolver.solve() v, rho_ = w1.split() rho0.assign(rho_) else: rho0.interpolate(thermodynamics.rho(state.parameters, theta0, Pi))
# set theta0 theta0.interpolate(theta_b + theta_pert) # calculate hydrostatic Pi rho_b = Function(Vr) compressible_hydrostatic_balance(state, theta_b, rho_b) compressible_hydrostatic_balance(state, theta0, rho0) # set Pi0 Pi0 = calculate_Pi0(state, theta0, rho0) state.parameters.Pi0 = Pi0 # set x component of velocity cp = state.parameters.cp dthetady = state.parameters.dthetady Pi = thermodynamics.pi(state.parameters, rho0, theta0) u = cp * dthetady / f * (Pi - Pi0) # set y component of velocity v = Function(Vr).assign(0.) compressible_eady_initial_v(state, theta0, rho0, v) # set initial u u_exp = as_vector([u, v, 0.]) u0.project(u_exp) # pass these initial conditions to the state.initialise method state.initialise([('u', u0), ('rho', rho0), ('theta', theta0)]) # set the background profiles state.set_reference_profiles([('rho', rho_b), ('theta', theta_b)])
def _setup_solver(self): from firedrake.assemble import create_assembly_callable import numpy as np state = self.state dt = state.timestepping.dt beta = dt*state.timestepping.alpha cp = state.parameters.cp mu = state.mu Vu = state.spaces("HDiv") Vu_broken = FunctionSpace(state.mesh, BrokenElement(Vu.ufl_element())) Vtheta = state.spaces("HDiv_v") Vrho = state.spaces("DG") h_deg = state.horizontal_degree v_deg = state.vertical_degree Vtrace = FunctionSpace(state.mesh, "HDiv Trace", degree=(h_deg, v_deg)) # Split up the rhs vector (symbolically) u_in, rho_in, theta_in = split(state.xrhs) # Build the function space for "broken" u and rho # and add the trace variable M = MixedFunctionSpace((Vu_broken, Vrho)) w, phi = TestFunctions(M) u, rho = TrialFunctions(M) l0 = TrialFunction(Vtrace) dl = TestFunction(Vtrace) n = FacetNormal(state.mesh) # Get background fields thetabar = state.fields("thetabar") rhobar = state.fields("rhobar") pibar = thermodynamics.pi(state.parameters, rhobar, thetabar) pibar_rho = thermodynamics.pi_rho(state.parameters, rhobar, thetabar) pibar_theta = thermodynamics.pi_theta(state.parameters, rhobar, thetabar) # Analytical (approximate) elimination of theta k = state.k # Upward pointing unit vector theta = -dot(k, u)*dot(k, grad(thetabar))*beta + theta_in # Only include theta' (rather than pi') in the vertical # component of the gradient # The pi prime term (here, bars are for mean and no bars are # for linear perturbations) pi = pibar_theta*theta + pibar_rho*rho # Vertical projection def V(u): return k*inner(u, k) # Specify degree for some terms as estimated degree is too large dxp = dx(degree=(self.quadrature_degree)) dS_vp = dS_v(degree=(self.quadrature_degree)) dS_hp = dS_h(degree=(self.quadrature_degree)) ds_vp = ds_v(degree=(self.quadrature_degree)) ds_tbp = ds_t(degree=(self.quadrature_degree)) + ds_b(degree=(self.quadrature_degree)) # Mass matrix for the trace space tM = assemble(dl('+')*l0('+')*(dS_v + dS_h) + dl*l0*ds_v + dl*l0*(ds_t + ds_b)) Lrhobar = Function(Vtrace) Lpibar = Function(Vtrace) rhopi_solver = LinearSolver(tM, solver_parameters={'ksp_type': 'cg', 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'}, options_prefix='rhobarpibar_solver') rhobar_avg = Function(Vtrace) pibar_avg = Function(Vtrace) def _traceRHS(f): return (dl('+')*avg(f)*(dS_v + dS_h) + dl*f*ds_v + dl*f*(ds_t + ds_b)) assemble(_traceRHS(rhobar), tensor=Lrhobar) assemble(_traceRHS(pibar), tensor=Lpibar) # Project averages of coefficients into the trace space with timed_region("Gusto:HybridProjectRhobar"): rhopi_solver.solve(rhobar_avg, Lrhobar) with timed_region("Gusto:HybridProjectPibar"): rhopi_solver.solve(pibar_avg, Lpibar) # Add effect of density of water upon theta if self.moisture is not None: water_t = Function(Vtheta).assign(0.0) for water in self.moisture: water_t += self.state.fields(water) theta_w = theta / (1 + water_t) thetabar_w = thetabar / (1 + water_t) else: theta_w = theta thetabar_w = thetabar # "broken" u and rho system Aeqn = (inner(w, (state.h_project(u) - u_in))*dx - beta*cp*div(theta_w*V(w))*pibar*dxp # following does nothing but is preserved in the comments # to remind us why (because V(w) is purely vertical). # + beta*cp*dot(theta_w*V(w), n)*pibar_avg('+')*dS_vp + beta*cp*dot(theta_w*V(w), n)*pibar_avg('+')*dS_hp + beta*cp*dot(theta_w*V(w), n)*pibar_avg*ds_tbp - beta*cp*div(thetabar_w*w)*pi*dxp + (phi*(rho - rho_in) - beta*inner(grad(phi), u)*rhobar)*dx + beta*dot(phi*u, n)*rhobar_avg('+')*(dS_v + dS_h)) if mu is not None: Aeqn += dt*mu*inner(w, k)*inner(u, k)*dx # Form the mixed operators using Slate # (A K)(X) = (X_r) # (K.T 0)(l) (0 ) # where X = ("broken" u, rho) A = Tensor(lhs(Aeqn)) X_r = Tensor(rhs(Aeqn)) # Off-diagonal block matrices containing the contributions # of the Lagrange multipliers (surface terms in the momentum equation) K = Tensor(beta*cp*dot(thetabar_w*w, n)*l0('+')*(dS_vp + dS_hp) + beta*cp*dot(thetabar_w*w, n)*l0*ds_vp + beta*cp*dot(thetabar_w*w, n)*l0*ds_tbp) # X = A.inv * (X_r - K * l), # 0 = K.T * X = -(K.T * A.inv * K) * l + K.T * A.inv * X_r, # so (K.T * A.inv * K) * l = K.T * A.inv * X_r # is the reduced equation for the Lagrange multipliers. # Right-hand side expression: (Forward substitution) Rexp = K.T * A.inv * X_r self.R = Function(Vtrace) # We need to rebuild R everytime data changes self._assemble_Rexp = create_assembly_callable(Rexp, tensor=self.R) # Schur complement operator: Smatexp = K.T * A.inv * K with timed_region("Gusto:HybridAssembleTraceOp"): S = assemble(Smatexp) S.force_evaluation() # Set up the Linear solver for the system of Lagrange multipliers self.lSolver = LinearSolver(S, solver_parameters=self.solver_parameters, options_prefix='lambda_solve') # Result function for the multiplier solution self.lambdar = Function(Vtrace) # Place to put result of u rho reconstruction self.urho = Function(M) # Reconstruction of broken u and rho u_, rho_ = self.urho.split() # Split operators for two-stage reconstruction _A = A.blocks _K = K.blocks _Xr = X_r.blocks A00 = _A[0, 0] A01 = _A[0, 1] A10 = _A[1, 0] A11 = _A[1, 1] K0 = _K[0, 0] Ru = _Xr[0] Rrho = _Xr[1] lambda_vec = AssembledVector(self.lambdar) # rho reconstruction Srho = A11 - A10 * A00.inv * A01 rho_expr = Srho.solve(Rrho - A10 * A00.inv * (Ru - K0 * lambda_vec), decomposition="PartialPivLU") self._assemble_rho = create_assembly_callable(rho_expr, tensor=rho_) # "broken" u reconstruction rho_vec = AssembledVector(rho_) u_expr = A00.solve(Ru - A01 * rho_vec - K0 * lambda_vec, decomposition="PartialPivLU") self._assemble_u = create_assembly_callable(u_expr, tensor=u_) # Project broken u into the HDiv space using facet averaging. # Weight function counting the dofs of the HDiv element: shapes = (Vu.finat_element.space_dimension(), np.prod(Vu.shape)) weight_kernel = """ for (int i=0; i<%d; ++i) { for (int j=0; j<%d; ++j) { w[i][j] += 1.0; }}""" % shapes self._weight = Function(Vu) par_loop(weight_kernel, dx, {"w": (self._weight, INC)}) # Averaging kernel self._average_kernel = """ for (int i=0; i<%d; ++i) { for (int j=0; j<%d; ++j) { vec_out[i][j] += vec_in[i][j]/w[i][j]; }}""" % shapes # HDiv-conforming velocity self.u_hdiv = Function(Vu) # Reconstruction of theta theta = TrialFunction(Vtheta) gamma = TestFunction(Vtheta) self.theta = Function(Vtheta) theta_eqn = gamma*(theta - theta_in + dot(k, self.u_hdiv)*dot(k, grad(thetabar))*beta)*dx theta_problem = LinearVariationalProblem(lhs(theta_eqn), rhs(theta_eqn), self.theta) self.theta_solver = LinearVariationalSolver(theta_problem, solver_parameters={'ksp_type': 'cg', 'pc_type': 'bjacobi', 'pc_sub_type': 'ilu'}, options_prefix='thetabacksubstitution') self.bcs = [DirichletBC(Vu, 0.0, "bottom"), DirichletBC(Vu, 0.0, "top")]