class InteriorPenalty(Diffusion): """ Interior penalty diffusion method :arg state: :class:`.State` object. :arg V: Function space of diffused field :arg direction: list containing directions in which function space is discontinuous: 1 corresponds to vertical, 2 to horizontal. :arg params: dictionary containing the interior penalty parameters :mu and kappa where mu is the penalty weighting function, which is :recommended to be proportional to 1/dx """ def __init__(self, state, V, direction=[1,2], params=None): super(InteriorPenalty, self).__init__(state) dt = state.timestepping.dt kappa = params['kappa'] mu = params['mu'] gamma = TestFunction(V) phi = TrialFunction(V) self.phi1 = Function(V) n = FacetNormal(state.mesh) a = inner(gamma,phi)*dx + dt*inner(grad(gamma), grad(phi)*kappa)*dx def get_flux_form(dS, M): fluxes = (-inner(2*avg(outer(phi, n)), avg(grad(gamma)*M)) - inner(avg(grad(phi)*M), 2*avg(outer(gamma, n))) + mu*inner(2*avg(outer(phi, n)), 2*avg(outer(gamma, n)*kappa)))*dS return fluxes if 1 in direction: a += dt*get_flux_form(dS_v, kappa) if 2 in direction: a += dt*get_flux_form(dS_h, kappa) L = inner(gamma,phi)*dx problem = LinearVariationalProblem(a, action(L,self.phi1), self.phi1) self.solver = LinearVariationalSolver(problem) def apply(self, x_in, x_out): self.phi1.assign(x_in) self.solver.solve() x_out.assign(self.phi1)
class ShallowWaterForcing(Forcing): def __init__(self, state, linear=False): self.state = state g = state.parameters.g f = state.f Vu = state.V[0] W = state.W self.x0 = Function(W) # copy x to here u0, D0 = split(self.x0) n = FacetNormal(state.mesh) un = 0.5 * (dot(u0, n) + abs(dot(u0, n))) F = TrialFunction(Vu) w = TestFunction(Vu) self.uF = Function(Vu) outward_normals = CellNormal(state.mesh) perp = lambda u: cross(outward_normals, u) a = inner(w, F) * dx L = (-f * inner(w, perp(u0)) + g * div(w) * D0) * dx - g * inner( jump(w, n), un("+") * D0("+") - un("-") * D0("-") ) * dS if not linear: L -= 0.5 * div(w) * inner(u0, u0) * dx u_forcing_problem = LinearVariationalProblem(a, L, self.uF) self.u_forcing_solver = LinearVariationalSolver(u_forcing_problem) def apply(self, scaling, x_in, x_nl, x_out, **kwargs): self.x0.assign(x_nl) self.u_forcing_solver.solve() # places forcing in self.uF self.uF *= scaling uF, _ = x_out.split() x_out.assign(x_in) uF += self.uF
class CompressibleForcing(Forcing): """ Forcing class for compressible Euler equations. """ def __init__(self, state, linear=False): self.state = state self._build_forcing_solver(linear) def _build_forcing_solver(self, linear): """ Only put forcing terms into the u equation. """ state = self.state self.scaling = Constant(1.0) Vu = state.V[0] W = state.W self.x0 = Function(W) # copy x to here u0, rho0, theta0 = split(self.x0) F = TrialFunction(Vu) w = TestFunction(Vu) self.uF = Function(Vu) Omega = state.Omega cp = state.parameters.cp mu = state.mu n = FacetNormal(state.mesh) pi = exner(theta0, rho0, state) a = inner(w, F) * dx L = self.scaling * ( +cp * div(theta0 * w) * pi * dx # pressure gradient [volume] - cp * jump(w * theta0, n) * avg(pi) * dS_v # pressure gradient [surface] ) if state.parameters.geopotential: Phi = state.Phi L += self.scaling * div(w) * Phi * dx # gravity term else: g = state.parameters.g L -= self.scaling * g * inner(w, state.k) * dx # gravity term if not linear: L -= self.scaling * 0.5 * div(w) * inner(u0, u0) * dx if Omega is not None: L -= self.scaling * inner(w, cross(2 * Omega, u0)) * dx # Coriolis term if mu is not None: self.mu_scaling = Constant(1.0) L -= self.mu_scaling * mu * inner(w, state.k) * inner(u0, state.k) * dx bcs = [DirichletBC(Vu, 0.0, "bottom"), DirichletBC(Vu, 0.0, "top")] u_forcing_problem = LinearVariationalProblem(a, L, self.uF, bcs=bcs) self.u_forcing_solver = LinearVariationalSolver(u_forcing_problem) def apply(self, scaling, x_in, x_nl, x_out, **kwargs): self.x0.assign(x_nl) self.scaling.assign(scaling) if "mu_alpha" in kwargs and kwargs["mu_alpha"] is not None: self.mu_scaling.assign(kwargs["mu_alpha"]) self.u_forcing_solver.solve() # places forcing in self.uF u_out, _, _ = x_out.split() x_out.assign(x_in) u_out += self.uF
class SUPGAdvection(Advection): """ An SUPG advection scheme that can apply DG upwinding (in the direction specified by the direction arg) if the function space is only partially continuous. :arg state: :class:`.State` object. :arg V:class:`.FunctionSpace` object. The advected field function space. :arg direction: list containing the directions in which the function space is discontinuous. 1 corresponds to the vertical direction, 2 to the horizontal direction :arg supg_params: dictionary containing SUPG parameters tau for each direction. If not supplied tau is set to dt/sqrt(15.) """ def __init__(self, state, V, direction=[], supg_params=None): super(SUPGAdvection, self).__init__(state) dt = state.timestepping.dt params = supg_params.copy() if supg_params else {} params.setdefault('a0', dt/sqrt(15.)) params.setdefault('a1', dt/sqrt(15.)) gamma = TestFunction(V) theta = TrialFunction(V) self.theta0 = Function(V) # make SUPG test function taus = [params["a0"], params["a1"]] for i in direction: taus[i] = 0.0 tau = Constant(((taus[0], 0.), (0., taus[1]))) dgamma = dot(dot(self.ubar, tau), grad(gamma)) gammaSU = gamma + dgamma n = FacetNormal(state.mesh) un = 0.5*(dot(self.ubar, n) + abs(dot(self.ubar, n))) a_mass = gammaSU*theta*dx arhs = a_mass - dt*gammaSU*dot(self.ubar, grad(theta))*dx if 1 in direction: arhs -= ( dt*dot(jump(gammaSU), (un('+')*theta('+') - un('-')*theta('-')))*dS_v - dt*(gammaSU('+')*dot(self.ubar('+'), n('+'))*theta('+') + gammaSU('-')*dot(self.ubar('-'), n('-'))*theta('-'))*dS_v ) if 2 in direction: arhs -= ( dt*dot(jump(gammaSU), (un('+')*theta('+') - un('-')*theta('-')))*dS_h - dt*(gammaSU('+')*dot(self.ubar('+'), n('+'))*theta('+') + gammaSU('-')*dot(self.ubar('-'), n('-'))*theta('-'))*dS_h ) self.theta1 = Function(V) self.dtheta = Function(V) problem = LinearVariationalProblem(a_mass, action(arhs,self.theta1), self.dtheta) self.solver = LinearVariationalSolver(problem, options_prefix='SUPGAdvection') def apply(self, x_in, x_out): # SSPRK Stage 1 self.theta1.assign(x_in) self.solver.solve() self.theta1.assign(self.dtheta) # SSPRK Stage 2 self.solver.solve() self.theta1.assign(0.75*x_in + 0.25*self.dtheta) # SSPRK Stage 3 self.solver.solve() x_out.assign((1.0/3.0)*x_in + (2.0/3.0)*self.dtheta)
class IncompressibleForcing(Forcing): """ Forcing class for incompressible Euler Boussinesq equations. """ def __init__(self, state, linear=False): self.state = state self._build_forcing_solver(linear) def _build_forcing_solver(self, linear): """ Only put forcing terms into the u equation. """ state = self.state self.scaling = Constant(1.0) Vu = state.V[0] W = state.W self.x0 = Function(W) # copy x to here u0, p0, b0 = split(self.x0) F = TrialFunction(Vu) w = TestFunction(Vu) self.uF = Function(Vu) Omega = state.Omega mu = state.mu a = inner(w, F) * dx L = ( self.scaling * div(w) * p0 * dx # pressure gradient + self.scaling * b0 * inner(w, state.k) * dx # gravity term ) if not linear: L -= self.scaling * 0.5 * div(w) * inner(u0, u0) * dx if Omega is not None: L -= self.scaling * inner(w, cross(2 * Omega, u0)) * dx # Coriolis term if mu is not None: self.mu_scaling = Constant(1.0) L -= self.mu_scaling * mu * inner(w, state.k) * inner(u0, state.k) * dx bcs = [DirichletBC(Vu, 0.0, "bottom"), DirichletBC(Vu, 0.0, "top")] u_forcing_problem = LinearVariationalProblem(a, L, self.uF, bcs=bcs) self.u_forcing_solver = LinearVariationalSolver(u_forcing_problem) Vp = state.V[1] p = TrialFunction(Vp) q = TestFunction(Vp) self.divu = Function(Vp) a = p * q * dx L = q * div(u0) * dx divergence_problem = LinearVariationalProblem(a, L, self.divu) self.divergence_solver = LinearVariationalSolver(divergence_problem) def apply(self, scaling, x_in, x_nl, x_out, **kwargs): self.x0.assign(x_nl) self.scaling.assign(scaling) if "mu_alpha" in kwargs and kwargs["mu_alpha"] is not None: self.mu_scaling.assign(kwargs["mu_alpha"]) self.u_forcing_solver.solve() # places forcing in self.uF u_out, p_out, _ = x_out.split() x_out.assign(x_in) u_out += self.uF if "incompressible" in kwargs and kwargs["incompressible"]: self.divergence_solver.solve() p_out.assign(self.divu)
class DGAdvection(Advection): """ DG 3 step SSPRK advection scheme that can be applied to a scalar or vector field :arg state: :class:`.State` object. :arg V: function space of advected field - should be DG :arg continuity: optional boolean. If ``True``, the advection equation is of the form: :math: `D_t +\nabla \cdot(uD) = 0`. If ``False``, the advection equation is of the form: :math: `D_t + (u \cdot \nabla)D = 0`. """ def __init__(self, state, V, continuity=False): super(DGAdvection, self).__init__(state) element = V.fiat_element assert element.entity_dofs() == element.entity_closure_dofs(), "Provided space is not discontinuous" dt = state.timestepping.dt if V.extruded: surface_measure = (dS_h + dS_v) else: surface_measure = dS phi = TestFunction(V) D = TrialFunction(V) self.D1 = Function(V) self.dD = Function(V) n = FacetNormal(state.mesh) # ( dot(v, n) + |dot(v, n)| )/2.0 un = 0.5*(dot(self.ubar, n) + abs(dot(self.ubar, n))) a_mass = inner(phi,D)*dx if continuity: a_int = -inner(grad(phi), outer(D, self.ubar))*dx else: a_int = -inner(div(outer(phi,self.ubar)),D)*dx a_flux = (dot(jump(phi), un('+')*D('+') - un('-')*D('-')))*surface_measure arhs = a_mass - dt*(a_int + a_flux) DGproblem = LinearVariationalProblem(a_mass, action(arhs,self.D1), self.dD) self.DGsolver = LinearVariationalSolver(DGproblem, solver_parameters={ 'ksp_type':'preonly', 'pc_type':'bjacobi', 'sub_pc_type': 'ilu'}, options_prefix='DGAdvection') def apply(self, x_in, x_out): # SSPRK Stage 1 self.D1.assign(x_in) self.DGsolver.solve() self.D1.assign(self.dD) # SSPRK Stage 2 self.DGsolver.solve() self.D1.assign(0.75*x_in + 0.25*self.dD) # SSPRK Stage 3 self.DGsolver.solve() x_out.assign((1.0/3.0)*x_in + (2.0/3.0)*self.dD)
class EulerPoincareForm(Advection): def __init__(self, state, V): super(EulerPoincareForm, self).__init__(state) dt = state.timestepping.dt w = TestFunction(V) u = TrialFunction(V) self.u0 = Function(V) ustar = 0.5*(self.u0 + u) n = FacetNormal(state.mesh) Upwind = 0.5*(sign(dot(self.ubar, n))+1) if state.mesh.geometric_dimension() == 3: if V.extruded: surface_measure = (dS_h + dS_v) else: surface_measure = dS # <w,curl(u) cross ubar + grad( u.ubar)> # =<curl(u),ubar cross w> - <div(w), u.ubar> # =<u,curl(ubar cross w)> - # <<u_upwind, [[n cross(ubar cross w)cross]]>> both = lambda u: 2*avg(u) Eqn = ( inner(w, u-self.u0)*dx + dt*inner(ustar, curl(cross(self.ubar, w)))*dx - dt*inner(both(Upwind*ustar), both(cross(n, cross(self.ubar, w))))*surface_measure - dt*div(w)*inner(ustar, self.ubar)*dx ) # define surface measure and terms involving perp differently # for slice (i.e. if V.extruded is True) and shallow water # (V.extruded is False) else: if V.extruded: surface_measure = (dS_h + dS_v) perp = lambda u: as_vector([-u[1], u[0]]) perp_u_upwind = Upwind('+')*perp(ustar('+')) + Upwind('-')*perp(ustar('-')) else: surface_measure = dS outward_normals = CellNormal(state.mesh) perp = lambda u: cross(outward_normals, u) perp_u_upwind = Upwind('+')*cross(outward_normals('+'),ustar('+')) + Upwind('-')*cross(outward_normals('-'),ustar('-')) Eqn = ( (inner(w, u-self.u0) - dt*inner(w, div(perp(ustar))*perp(self.ubar)) - dt*div(w)*inner(ustar, self.ubar))*dx - dt*inner(jump(inner(w, perp(self.ubar)), n), perp_u_upwind)*surface_measure + dt*jump(inner(w, perp(self.ubar))*perp(ustar), n)*surface_measure ) a = lhs(Eqn) L = rhs(Eqn) self.u1 = Function(V) uproblem = LinearVariationalProblem(a, L, self.u1) self.usolver = LinearVariationalSolver(uproblem, options_prefix='EPAdvection') def apply(self, x_in, x_out): self.u0.assign(x_in) self.usolver.solve() x_out.assign(self.u1)
def run(steady=False): """ solve CdT/dt = S + div(k*grad(T)) => C*v*(dT/dt)/k*dx - S*v/k*dx + grad(v)*grad(T)*dx = v*dot(grad(T), n)*ds """ steps = 250 dt = 1e-10 timescale = (0, steps * dt) if steady: print('Running steady state.') else: print(f'Running with time step {dt:.2g}s on time interval: ' f'{timescale[0]:.2g}s - {timescale[1]:.2g}s') dt_invc = Constant(1 / dt) extent = [40e-6, 40e-6, 40e-6] mesh = BoxMesh(20, 20, 20, *extent) V = FunctionSpace(mesh, 'CG', 1) print(V.dim()) T = Function(V) # temperature at time i+1 (electron for now) T_ = Function(V) # temperature at time i v = TestFunction(V) # test function S = create_S(mesh, V, extent) C = create_heat_capacity(mesh, V, extent) k = create_conductivity(mesh, V, T) set_initial_value(mesh, T_, extent) # Mass matrix section M = C * T * dt_invc * v * dx M_ = C * T_ * dt_invc * v * dx # Stiffness matrix section A = k * dot(grad(T), grad(v)) * dx # function section f = S * v * dx # boundaries bcs, R, b = create_dirichlet_bounds(mesh, V, T, v, k, g=100, boundary=[1, 2, 3, 4, 5, 6]) # bcs += create_dirichlet_bounds(mesh, V, T, v, k, 500, [6])[0] # bcs, R, b = create_robin_bounds(mesh, T, v, k, 1e8/(100), 1e8) if steady: steps = 1 a = A + R L = f + b else: a = M + A + R L = M_ + f + b prob = NonlinearVariationalProblem(a - L, T, bcs=bcs) solver = NonlinearVariationalSolver(prob, solver_parameters=SOLVE_PARAMS) T.assign(T_) timestamp = datetime.now().strftime("%d-%b-%Y-%H-%M-%S") outfile = File(f'{timestamp}/first_output.pvd') outfile.write(T_, target_degree=1, target_continuity=H1) last_perc = 0 for i in range(steps): solver.solve() perc = int(100 * (i + 1) / steps) if perc > last_perc: print(f'{perc}%') last_perc = perc T_.assign(T) outfile.write(T_, target_degree=1, target_continuity=H1)
def _ad_copy(self): from firedrake import Function r = Function(self.function_space()) r.assign(self) return r
def _ad_add(self, other): from firedrake import Function r = Function(self.function_space()) Function.assign(r, self + other) return r
def _ad_mul(self, other): from firedrake import Function r = Function(self.function_space()) r.assign(self * other) return r
class Forcing(object, metaclass=ABCMeta): """ Base class for forcing terms for Gusto. :arg state: x :class:`.State` object. :arg euler_poincare: if True then the momentum equation is in Euler Poincare form and we need to add 0.5*grad(u^2) to the forcing term. If False then this term is not added. :arg linear: if True then we are solving a linear equation so nonlinear terms (namely the Euler Poincare term) should not be added. :arg extra_terms: extra terms to add to the u component of the forcing term - these will be multiplied by the appropriate test function. """ def __init__(self, state, euler_poincare=True, linear=False, extra_terms=None, moisture=None): self.state = state if linear: self.euler_poincare = False logger.warning( 'Setting euler_poincare to False because you have set linear=True' ) else: self.euler_poincare = euler_poincare # set up functions self.Vu = state.spaces("HDiv") # this is the function that the forcing term is applied to self.x0 = Function(state.W) self.test = TestFunction(self.Vu) self.trial = TrialFunction(self.Vu) # this is the function that contains the result of solving # <test, trial> = <test, F(x0)>, where F is the forcing term self.uF = Function(self.Vu) # find out which terms we need self.extruded = self.Vu.extruded self.coriolis = state.Omega is not None or hasattr( state.fields, "coriolis") self.sponge = state.mu is not None self.hydrostatic = state.hydrostatic self.topography = hasattr(state.fields, "topography") self.extra_terms = extra_terms self.moisture = moisture # some constants to use for scaling terms self.scaling = Constant(1.) self.impl = Constant(1.) self._build_forcing_solvers() def mass_term(self): return inner(self.test, self.trial) * dx def coriolis_term(self): u0 = split(self.x0)[0] return -inner(self.test, cross(2 * self.state.Omega, u0)) * dx def sponge_term(self): u0 = split(self.x0)[0] return self.state.mu * inner(self.test, self.state.k) * inner( u0, self.state.k) * dx def euler_poincare_term(self): u0 = split(self.x0)[0] return -0.5 * div(self.test) * inner(self.state.h_project(u0), u0) * dx def hydrostatic_term(self): u0 = split(self.x0)[0] return inner(u0, self.state.k) * inner(self.test, self.state.k) * dx @abstractmethod def pressure_gradient_term(self): pass def forcing_term(self): L = self.pressure_gradient_term() if self.extruded: L += self.gravity_term() if self.coriolis: L += self.coriolis_term() if self.euler_poincare: L += self.euler_poincare_term() if self.topography: L += self.topography_term() if self.extra_terms is not None: L += inner(self.test, self.extra_terms) * dx # scale L L = self.scaling * L # sponge term has a separate scaling factor as it is always implicit if self.sponge: L -= self.impl * self.state.timestepping.dt * self.sponge_term() # hydrostatic term has no scaling factor if self.hydrostatic: L += (2 * self.impl - 1) * self.hydrostatic_term() return L def _build_forcing_solvers(self): a = self.mass_term() L = self.forcing_term() if self.Vu.extruded: bcs = [ DirichletBC(self.Vu, 0.0, "bottom"), DirichletBC(self.Vu, 0.0, "top") ] else: bcs = None u_forcing_problem = LinearVariationalProblem(a, L, self.uF, bcs=bcs) solver_parameters = {} if self.state.output.log_level == DEBUG: solver_parameters["ksp_monitor_true_residual"] = True self.u_forcing_solver = LinearVariationalSolver( u_forcing_problem, solver_parameters=solver_parameters, options_prefix="UForcingSolver") def apply(self, scaling, x_in, x_nl, x_out, **kwargs): """ Function takes x as input, computes F(x_nl) and returns x_out = x + scale*F(x_nl) as output. :arg x_in: :class:`.Function` object :arg x_nl: :class:`.Function` object :arg x_out: :class:`.Function` object :arg implicit: forcing stage for sponge and hydrostatic terms, if present """ self.scaling.assign(scaling) self.x0.assign(x_nl) implicit = kwargs.get("implicit") if implicit is not None: self.impl.assign(int(implicit)) self.u_forcing_solver.solve() # places forcing in self.uF uF = x_out.split()[0] x_out.assign(x_in) uF += self.uF