def _setup_solver(self): state = self.state # just cutting down line length a bit dt = state.timestepping.dt beta = dt*state.timestepping.alpha mu = state.mu # Split up the rhs vector (symbolically) u_in, p_in, b_in = split(state.xrhs) # Build the reduced function space for u,p M = MixedFunctionSpace((state.V[0], state.V[1])) w, phi = TestFunctions(M) u, p = TrialFunctions(M) # Get background fields bbar = state.bbar # Analytical (approximate) elimination of theta k = state.k # Upward pointing unit vector b = -dot(k,u)*dot(k,grad(bbar))*beta + b_in # vertical projection def V(u): return k*inner(u,k) eqn = ( inner(w, (u - u_in))*dx - beta*div(w)*p*dx - beta*inner(w,k)*b*dx + phi*div(u)*dx ) 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 p solver self.up = Function(M) # Boundary conditions (assumes extruded mesh) dim = M.sub(0).ufl_element().value_shape()[0] bc = ("0.0",)*dim bcs = [DirichletBC(M.sub(0), Expression(bc), "bottom"), DirichletBC(M.sub(0), Expression(bc), "top")] # preconditioner equation L = self.L Ap = ( inner(w,u) + L*L*div(w)*div(u) + phi*p/L/L )*dx # Solver for u, p up_problem = LinearVariationalProblem( aeqn, Leqn, self.up, bcs=bcs, aP=Ap) nullspace = MixedVectorSpaceBasis(M, [M.sub(0), VectorSpaceBasis(constant=True)]) self.up_solver = LinearVariationalSolver(up_problem, solver_parameters=self.params, nullspace=nullspace) # Reconstruction of b b = TrialFunction(state.V[2]) gamma = TestFunction(state.V[2]) u, p = self.up.split() self.b = Function(state.V[2]) b_eqn = gamma*(b - b_in + dot(k,u)*dot(k,grad(bbar))*beta)*dx b_problem = LinearVariationalProblem(lhs(b_eqn), rhs(b_eqn), self.b) self.b_solver = LinearVariationalSolver(b_problem)
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 # 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((state.V[0], state.V[1])) w, phi = TestFunctions(M) u, rho = TrialFunctions(M) n = FacetNormal(state.mesh) # Get background fields thetabar = state.thetabar rhobar = state.rhobar pibar = exner(thetabar, rhobar, state) pibar_rho = exner_rho(thetabar, rhobar, state) pibar_theta = exner_theta(thetabar, rhobar, state) # 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) eqn = ( inner(w, (u - u_in))*dx - beta*cp*div(theta*V(w))*pibar*dx # 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)*pi*dx + beta*cp*jump(thetabar*w,n)*avg(pi)*dS_v + (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) dim = M.sub(0).ufl_element().value_shape()[0] bc = ("0.0",)*dim bcs = [DirichletBC(M.sub(0), Expression(bc), "bottom"), DirichletBC(M.sub(0), Expression(bc), "top")] # Solver for u, rho urho_problem = LinearVariationalProblem( aeqn, Leqn, self.urho, bcs=bcs) self.urho_solver = LinearVariationalSolver(urho_problem, solver_parameters=self.params, options_prefix='ImplicitSolver') # Reconstruction of theta theta = TrialFunction(state.V[2]) gamma = TestFunction(state.V[2]) u, rho = self.urho.split() self.theta = Function(state.V[2]) 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), 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. """ # Calculate hydrostatic Pi W = MixedFunctionSpace((state.Vv,state.V[1])) v, pi = TrialFunctions(W) dv, dpi = TestFunctions(W) n = FacetNormal(state.mesh) cp = state.parameters.cp alhs = ( (cp*inner(v,dv) - cp*div(dv*theta0)*pi)*dx + dpi*div(theta0*v)*dx ) if top: bmeasure = ds_t bstring = "bottom" else: bmeasure = ds_b bstring = "top" arhs = -cp*inner(dv,n)*theta0*pi_boundary*bmeasure if state.parameters.geopotential: Phi = state.Phi arhs += div(dv)*Phi*dx - inner(dv,n)*Phi*bmeasure else: g = state.parameters.g arhs -= g*inner(dv,state.k)*dx if(state.mesh.geometric_dimension() == 2): bcs = [DirichletBC(W.sub(0), Expression(("0.", "0.")), bstring)] elif(state.mesh.geometric_dimension() == 3): bcs = [DirichletBC(W.sub(0), Expression(("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_estimate_eigenvalues': True, 'fieldsplit_1_mg_levels_ksp_chebyshev_estimate_eigenvalues_random': 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) kappa = state.parameters.kappa R_d = state.parameters.R_d p_0 = state.parameters.p_0 if solve_for_rho: w1 = Function(W) v, rho = w1.split() rho.interpolate(p_0*(Pi**((1-kappa)/kappa))/R_d/theta0) v, rho = split(w1) dv, dpi = TestFunctions(W) pi = ((R_d/p_0)*rho*theta0)**(kappa/(1.-kappa)) F = ( (cp*inner(v,dv) - cp*div(dv*theta0)*pi)*dx + dpi*div(theta0*v)*dx + cp*inner(dv,n)*theta0*pi_boundary*bmeasure ) if state.parameters.geopotential: F += - div(dv)*Phi*dx + inner(dv,n)*Phi*bmeasure else: 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(p_0*(Pi**((1-kappa)/kappa))/R_d/theta0)
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 _setup_solver(self): state = self.state # just cutting down line length a bit dt = state.dt beta_ = dt * self.alpha Vu = state.spaces("HDiv") Vb = state.spaces("theta") Vp = state.spaces("DG") # Store time-stepping coefficients as UFL Constants beta = Constant(beta_) # Split up the rhs vector (symbolically) self.xrhs = Function(self.equations.function_space) u_in, p_in, b_in = split(self.xrhs) # Build the reduced function space for u,p M = MixedFunctionSpace((Vu, Vp)) w, phi = TestFunctions(M) u, p = TrialFunctions(M) # Get background fields bbar = state.fields("bbar") # Analytical (approximate) elimination of theta k = state.k # Upward pointing unit vector b = -dot(k, u) * dot(k, grad(bbar)) * beta + b_in # vertical projection def V(u): return k * inner(u, k) eqn = (inner(w, (u - u_in)) * dx - beta * div(w) * p * dx - beta * inner(w, k) * b * dx + phi * div(u) * dx) if hasattr(self.equations, "mu"): eqn += dt * self.equations.mu * inner(w, k) * inner(u, k) * dx aeqn = lhs(eqn) Leqn = rhs(eqn) # Place to put result of u p solver self.up = Function(M) # Boundary conditions (assumes extruded mesh) # BCs are declared for the plain velocity space. As we need them in # a mixed problem, we replicate the BCs but for subspace of M bcs = [ DirichletBC(M.sub(0), bc.function_arg, bc.sub_domain) for bc in self.equations.bcs['u'] ] # Solver for u, p up_problem = LinearVariationalProblem(aeqn, Leqn, self.up, bcs=bcs) # Provide callback for the nullspace of the trace system def trace_nullsp(T): return VectorSpaceBasis(constant=True) appctx = {"trace_nullspace": trace_nullsp} self.up_solver = LinearVariationalSolver( up_problem, solver_parameters=self.solver_parameters, appctx=appctx) # Reconstruction of b b = TrialFunction(Vb) gamma = TestFunction(Vb) u, p = self.up.split() self.b = Function(Vb) b_eqn = gamma * (b - b_in + dot(k, u) * dot(k, grad(bbar)) * beta) * dx b_problem = LinearVariationalProblem(lhs(b_eqn), rhs(b_eqn), self.b) self.b_solver = LinearVariationalSolver(b_problem)
def compressible_hydrostatic_balance(state, theta0, rho0, exner0=None, top=False, exner_boundary=Constant(1.0), mr_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 exner_boundary: a field or expression to use as boundary data for exner on the top or bottom as specified. :arg mr_t: the initial total water mixing ratio field. """ # Calculate hydrostatic Pi VDG = state.spaces("DG") Vu = state.spaces("HDiv") Vv = FunctionSpace(state.mesh, Vu.ufl_element()._elements[-1]) W = MixedFunctionSpace((Vv, VDG)) v, exner = TrialFunctions(W) dv, dexner = TestFunctions(W) n = FacetNormal(state.mesh) cp = state.parameters.cp # add effect of density of water upon theta theta = theta0 if mr_t is not None: theta = theta0 / (1 + mr_t) alhs = ((cp * inner(v, dv) - cp * div(dv * theta) * exner) * dx + dexner * div(theta * v) * dx) if top: bmeasure = ds_t bstring = "bottom" else: bmeasure = ds_b bstring = "top" arhs = -cp * inner(dv, n) * theta * exner_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), zero(), bstring)] w = Function(W) exner_problem = 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' } } exner_solver = LinearVariationalSolver(exner_problem, solver_parameters=params, options_prefix="exner_solver") exner_solver.solve() v, exner = w.split() if exner0 is not None: exner0.assign(exner) if solve_for_rho: w1 = Function(W) v, rho = w1.split() rho.interpolate(thermodynamics.rho(state.parameters, theta0, exner)) v, rho = split(w1) dv, dexner = TestFunctions(W) exner = thermodynamics.exner_pressure(state.parameters, rho, theta0) F = ((cp * inner(v, dv) - cp * div(dv * theta) * exner) * dx + dexner * div(theta0 * v) * dx + cp * inner(dv, n) * theta * exner_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, exner))
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))
def moist_hydrostatic_balance(state, theta_e, water_t, pi_boundary=Constant(1.0)): """ 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. :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 pi_boundary: the value of pi on the lower boundary of the domain. """ 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() Vv = state.fields('u').function_space() n = FacetNormal(state.mesh) g = state.parameters.g cp = state.parameters.cp R_d = state.parameters.R_d p_0 = state.parameters.p_0 VDG = state.spaces("DG") if any(deg > 2 for deg in VDG.ufl_element().degree()): state.logger.warning("default quadrature degree most likely not sufficient for this degree element") quadrature_degree = (5, 5) params = {'ksp_type': 'preonly', 'ksp_monitor_true_residual': True, 'ksp_converged_reason': True, 'snes_converged_reason': True, 'ksp_max_it': 100, 'mat_type': 'aij', 'pc_type': 'lu', 'pc_factor_mat_solver_type': 'mumps'} theta0.interpolate(theta_e) water_v0.interpolate(water_t) Pi = Function(Vr) epsilon = 0.9 # relaxation constant # set up mixed space Z = MixedFunctionSpace((Vt, Vt)) z = Function(Z) gamma, phi = TestFunctions(Z) theta_v, w_v = z.split() # give first guesses for trial functions theta_v.assign(theta0) w_v.assign(water_v0) theta_v, w_v = split(z) # define variables T = thermodynamics.T(state.parameters, theta_v, Pi, r_v=w_v) p = thermodynamics.p(state.parameters, Pi) w_sat = thermodynamics.r_sat(state.parameters, T, p) dxp = dx(degree=(quadrature_degree)) # set up weak form of theta_e and w_sat equations F = (-gamma * theta_e * dxp + gamma * thermodynamics.theta_e(state.parameters, T, p, w_v, water_t) * dxp - phi * w_v * dxp + phi * w_sat * dxp) problem = NonlinearVariationalProblem(F, z) solver = NonlinearVariationalSolver(problem, solver_parameters=params) theta_v, w_v = z.split() Pi_h = Function(Vr).interpolate((p / p_0) ** (R_d / cp)) # solve for Pi with theta_v and w_v constant # then solve for theta_v and w_v with Pi constant for i in range(5): compressible_hydrostatic_balance(state, theta0, rho0, pi0=Pi_h, water_t=water_t) Pi.assign(Pi * (1 - epsilon) + epsilon * Pi_h) solver.solve() theta0.assign(theta0 * (1 - epsilon) + epsilon * theta_v) water_v0.assign(water_v0 * (1 - epsilon) + epsilon * w_v) # now begin on Newton solver, setup up new mixed space Z = MixedFunctionSpace((Vt, Vt, Vr, Vv)) z = Function(Z) gamma, phi, psi, w = TestFunctions(Z) theta_v, w_v, pi, v = z.split() # use previous values as first guesses for newton solver theta_v.assign(theta0) w_v.assign(water_v0) pi.assign(Pi) theta_v, w_v, pi, v = split(z) # define variables T = thermodynamics.T(state.parameters, theta_v, pi, r_v=w_v) p = thermodynamics.p(state.parameters, pi) w_sat = thermodynamics.r_sat(state.parameters, T, p) F = (-gamma * theta_e * dxp + gamma * thermodynamics.theta_e(state.parameters, T, p, w_v, water_t) * dxp - phi * w_v * dxp + phi * w_sat * dxp + cp * inner(v, w) * dxp - cp * div(w * theta_v / (1.0 + water_t)) * pi * dxp + psi * div(theta_v * v / (1.0 + water_t)) * dxp + cp * inner(w, n) * pi_boundary * theta_v / (1.0 + water_t) * ds_b + g * inner(w, state.k) * dxp) bcs = [DirichletBC(Z.sub(3), 0.0, "top")] problem = NonlinearVariationalProblem(F, z, bcs=bcs) solver = NonlinearVariationalSolver(problem, solver_parameters=params) solver.solve() theta_v, w_v, pi, v = z.split() # assign final values theta0.assign(theta_v) water_v0.assign(w_v) # find rho compressible_hydrostatic_balance(state, theta0, rho0, 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 mu = state.mu Vu = state.spaces("HDiv") Vb = state.spaces("HDiv_v") Vp = state.spaces("DG") # Split up the rhs vector (symbolically) u_in, p_in, b_in = split(state.xrhs) # Build the reduced function space for u,p M = MixedFunctionSpace((Vu, Vp)) w, phi = TestFunctions(M) u, p = TrialFunctions(M) # Get background fields bbar = state.fields("bbar") # Analytical (approximate) elimination of theta k = state.k # Upward pointing unit vector b = -dot(k, u)*dot(k, grad(bbar))*beta + b_in # vertical projection def V(u): return k*inner(u, k) eqn = ( inner(w, (u - u_in))*dx - beta*div(w)*p*dx - beta*inner(w, k)*b*dx + phi*div(u)*dx ) 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 p solver self.up = 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, p up_problem = LinearVariationalProblem(aeqn, Leqn, self.up, bcs=bcs) # Provide callback for the nullspace of the trace system def trace_nullsp(T): return VectorSpaceBasis(constant=True) appctx = {"trace_nullspace": trace_nullsp} self.up_solver = LinearVariationalSolver(up_problem, solver_parameters=self.solver_parameters, appctx=appctx) # Reconstruction of b b = TrialFunction(Vb) gamma = TestFunction(Vb) u, p = self.up.split() self.b = Function(Vb) b_eqn = gamma*(b - b_in + dot(k, u)*dot(k, grad(bbar))*beta)*dx b_problem = LinearVariationalProblem(lhs(b_eqn), rhs(b_eqn), self.b) self.b_solver = LinearVariationalSolver(b_problem)