Esempio n. 1
0
    def advection_term(self, q):

        if self.state.mesh.topological_dimension() == 3:
            # <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)

            L = (inner(q, curl(cross(self.ubar, self.test))) * dx -
                 inner(both(self.Upwind * q),
                       both(cross(self.n, cross(self.ubar, self.test)))) *
                 self.dS)

        else:

            if self.ibp == "once":
                L = (-inner(
                    self.gradperp(inner(self.test, self.perp(self.ubar))), q) *
                     dx - inner(
                         jump(inner(self.test, self.perp(self.ubar)), self.n),
                         self.perp_u_upwind(q)) * self.dS)
            else:
                L = (
                    (-inner(self.test,
                            div(self.perp(q)) * self.perp(self.ubar))) * dx -
                    inner(jump(inner(self.test, self.perp(self.ubar)), self.n),
                          self.perp_u_upwind(q)) * self.dS + jump(
                              inner(self.test, self.perp(self.ubar)) *
                              self.perp(q), self.n) * self.dS)

        L -= 0.5 * div(self.test) * inner(q, self.ubar) * dx

        return L
Esempio n. 2
0
    def __init__(self, state, V, *, ibp="once", solver_params=None):
        super().__init__(state=state,
                         V=V,
                         ibp=ibp,
                         solver_params=solver_params)

        self.Upwind = 0.5 * (sign(dot(self.ubar, self.n)) + 1)

        if self.state.mesh.topological_dimension() == 2:
            self.perp = state.perp
            if state.on_sphere:
                outward_normals = CellNormal(state.mesh)
                self.perp_u_upwind = lambda q: self.Upwind('+') * cross(
                    outward_normals('+'), q('+')) + self.Upwind('-') * cross(
                        outward_normals('-'), q('-'))
            else:
                self.perp_u_upwind = lambda q: self.Upwind('+') * state.perp(
                    q('+')) + self.Upwind('-') * state.perp(q('-'))
            self.gradperp = lambda u: state.perp(grad(u))
        elif self.state.mesh.topological_dimension() == 3:
            if self.ibp == "twice":
                raise NotImplementedError(
                    "ibp=twice is not implemented for 3d problems")
        else:
            raise RuntimeError("topological mesh dimension must be 2 or 3")
Esempio n. 3
0
def vector_invariant_form(state, test, q, ibp=IntegrateByParts.ONCE):

    Vu = state.spaces("HDiv")
    dS_ = (dS_v + dS_h) if Vu.extruded else dS
    ubar = Function(Vu)
    n = FacetNormal(state.mesh)
    Upwind = 0.5 * (sign(dot(ubar, n)) + 1)

    if state.mesh.topological_dimension() == 3:

        if ibp != IntegrateByParts.ONCE:
            raise NotImplementedError

        # <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)

        L = (inner(q, curl(cross(ubar, test))) * dx -
             inner(both(Upwind * q), both(cross(n, cross(ubar, test)))) * dS_)

    else:

        perp = state.perp
        if state.on_sphere:
            outward_normals = CellNormal(state.mesh)
            perp_u_upwind = lambda q: Upwind('+') * cross(
                outward_normals('+'), q('+')) + Upwind('-') * cross(
                    outward_normals('-'), q('-'))
        else:
            perp_u_upwind = lambda q: Upwind('+') * perp(q('+')) + Upwind(
                '-') * perp(q('-'))

        if ibp == IntegrateByParts.ONCE:
            L = (-inner(perp(grad(inner(test, perp(ubar)))), q) * dx -
                 inner(jump(inner(test, perp(ubar)), n), perp_u_upwind(q)) *
                 dS_)
        else:
            L = ((-inner(test,
                         div(perp(q)) * perp(ubar))) * dx -
                 inner(jump(inner(test, perp(ubar)), n), perp_u_upwind(q)) *
                 dS_ + jump(inner(test, perp(ubar)) * perp(q), n) * dS_)

    L -= 0.5 * div(test) * inner(q, ubar) * dx

    form = transporting_velocity(L, ubar)

    return transport(form, TransportEquationType.vector_invariant)
Esempio n. 4
0
    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)
Esempio n. 5
0
    def advection_term(self, q):

        n = FacetNormal(self.state.mesh)
        Upwind = 0.5 * (sign(dot(self.ubar, n)) + 1)

        if self.state.mesh.topological_dimension() == 3:
            # <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)

            L = (inner(q, curl(cross(self.ubar, self.test))) * dx -
                 inner(both(Upwind * q),
                       both(cross(n, cross(self.ubar, self.test)))) * self.dS)

        else:

            perp = self.state.perp
            if self.state.on_sphere:
                outward_normals = CellNormal(self.state.mesh)
                perp_u_upwind = lambda q: Upwind('+') * cross(
                    outward_normals('+'), q('+')) + Upwind('-') * cross(
                        outward_normals('-'), q('-'))
            else:
                perp_u_upwind = lambda q: Upwind('+') * perp(q('+')) + Upwind(
                    '-') * perp(q('-'))

            if self.ibp == IntegrateByParts.ONCE:
                L = (-inner(perp(grad(inner(self.test, perp(self.ubar)))), q) *
                     dx - inner(jump(inner(self.test, perp(self.ubar)), n),
                                perp_u_upwind(q)) * self.dS)
            else:
                L = ((-inner(self.test,
                             div(perp(q)) * perp(self.ubar))) * dx -
                     inner(jump(inner(self.test, perp(self.ubar)), n),
                           perp_u_upwind(q)) * self.dS +
                     jump(inner(self.test, perp(self.ubar)) * perp(q), n) *
                     self.dS)

        L -= 0.5 * div(self.test) * inner(q, self.ubar) * dx

        return L
Esempio n. 6
0
    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)
Esempio n. 7
0
    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)
Esempio n. 8
0
    def setup(self, state, vorticity_type=None):
        """Solver for vorticity.

        :arg state: The state containing model.
        :arg vorticity_type: must be "relative", "absolute" or "potential"
        """
        if not self._initialised:
            vorticity_types = ["relative", "absolute", "potential"]
            if vorticity_type not in vorticity_types:
                raise ValueError("vorticity type must be one of %s, not %s" % (vorticity_types, vorticity_type))
            try:
                space = state.spaces("CG")
            except AttributeError:
                dgspace = state.spaces("DG")
                cg_degree = dgspace.ufl_element().degree() + 2
                space = FunctionSpace(state.mesh, "CG", cg_degree)
            super().setup(state, space=space)
            u = state.fields("u")
            gamma = TestFunction(space)
            q = TrialFunction(space)

            if vorticity_type == "potential":
                D = state.fields("D")
                a = q*gamma*D*dx
            else:
                a = q*gamma*dx

            if state.on_sphere:
                cell_normals = CellNormal(state.mesh)
                gradperp = lambda psi: cross(cell_normals, grad(psi))
                L = (- inner(gradperp(gamma), u))*dx
            else:
                raise NotImplementedError("The vorticity diagnostics have only been implemented for 2D spherical geometries.")

            if vorticity_type != "relative":
                f = state.fields("coriolis")
                L += gamma*f*dx

            problem = LinearVariationalProblem(a, L, self.field)
            self.solver = LinearVariationalSolver(problem, solver_parameters={"ksp_type": "cg"})
Esempio n. 9
0
 def coriolis_term(self):
     u0 = split(self.x0)[0]
     return -inner(self.test, cross(2 * self.state.Omega, u0)) * dx
Esempio n. 10
0
    def __init__(self,
                 mesh,
                 dt,
                 output=None,
                 parameters=None,
                 diagnostics=None,
                 diagnostic_fields=None):

        if output is None:
            raise RuntimeError(
                "You must provide a directory name for dumping results")
        else:
            self.output = output
        self.parameters = parameters

        if diagnostics is not None:
            self.diagnostics = diagnostics
        else:
            self.diagnostics = Diagnostics()
        if diagnostic_fields is not None:
            self.diagnostic_fields = diagnostic_fields
        else:
            self.diagnostic_fields = []

        # The mesh
        self.mesh = mesh

        self.spaces = SpaceCreator(mesh)

        if self.output.dumplist is None:

            self.output.dumplist = []

        self.fields = StateFields(*self.output.dumplist)

        self.dumpdir = None
        self.dumpfile = None
        self.to_pickup = None

        # figure out if we're on a sphere
        try:
            self.on_sphere = (mesh._base_mesh.geometric_dimension() == 3
                              and mesh._base_mesh.topological_dimension() == 2)
        except AttributeError:
            self.on_sphere = (mesh.geometric_dimension() == 3
                              and mesh.topological_dimension() == 2)

        #  build the vertical normal and define perp for 2d geometries
        dim = mesh.topological_dimension()
        if self.on_sphere:
            x = SpatialCoordinate(mesh)
            R = sqrt(inner(x, x))
            self.k = interpolate(x / R, mesh.coordinates.function_space())
            if dim == 2:
                outward_normals = CellNormal(mesh)
                self.perp = lambda u: cross(outward_normals, u)
        else:
            kvec = [0.0] * dim
            kvec[dim - 1] = 1.0
            self.k = Constant(kvec)
            if dim == 2:
                self.perp = lambda u: as_vector([-u[1], u[0]])

        # setup logger
        logger.setLevel(output.log_level)
        set_log_handler(mesh.comm)
        if parameters is not None:
            logger.info("Physical parameters that take non-default values:")
            logger.info(", ".join("%s: %s" % (k, float(v))
                                  for (k, v) in vars(parameters).items()))

        #  Constant to hold current time
        self.t = Constant(0.0)
        if type(dt) is Constant:
            self.dt = dt
        elif type(dt) in (float, int):
            self.dt = Constant(dt)
        else:
            raise TypeError(
                f'dt must be a Constant, float or int, not {type(dt)}')
Esempio n. 11
0
    Dn.interpolate(Dexpr)
b.interpolate(bexpr)

if COMM_WORLD.Get_rank() == 0:
    print("Finished setting up functions at", ctime())

# Build perp and upwind perp (including 2D version for test purposes)
n = FacetNormal(mesh)
s = lambda u: 0.5 * (sign(dot(u, n)) + 1)
uw = lambda u, v: (s(u)('+') * v('+') + s(u)('-') * v('-'))

if mesh.geometric_dimension() == 2:
    perp = lambda u: as_vector([-u[1], u[0]])
    p_uw = lambda u, v: perp(uw(u, v))
else:
    perp = lambda u: cross(CellNormal(mesh), u)
    out_n = CellNormal(mesh)
    p_uw = lambda u, v: (s(u)('+') * cross(out_n('+'), v('+')) + s(u)
                         ('-') * cross(out_n('-'), v('-')))

# Initial solve for vorticity for output
eta = TestFunction(W2)
q_ = TrialFunction(W2)
q_eqn = eta * q_ * Dn * dx + inner(perp(grad(eta)), un) * dx - eta * f * dx
q_p = LinearVariationalProblem(lhs(q_eqn), rhs(q_eqn), qn)
q_solver = LinearVariationalSolver(q_p,
                                   solver_parameters={
                                       "ksp_type": "preonly",
                                       "pc_type": "lu"
                                   })
q_solver.solve()
Esempio n. 12
0
def perp(u):
    return fd.cross(outward_normals, u)
Esempio n. 13
0
    def __init__(self,
                 mesh,
                 vertical_degree=None,
                 horizontal_degree=1,
                 family="RT",
                 Coriolis=None,
                 sponge_function=None,
                 hydrostatic=None,
                 timestepping=None,
                 output=None,
                 parameters=None,
                 diagnostics=None,
                 fieldlist=None,
                 diagnostic_fields=None,
                 u_bc_ids=None):

        self.family = family
        self.vertical_degree = vertical_degree
        self.horizontal_degree = horizontal_degree
        self.Omega = Coriolis
        self.mu = sponge_function
        self.hydrostatic = hydrostatic
        self.timestepping = timestepping
        if output is None:
            raise RuntimeError(
                "You must provide a directory name for dumping results")
        else:
            self.output = output
        self.parameters = parameters
        if fieldlist is None:
            raise RuntimeError(
                "You must provide a fieldlist containing the names of the prognostic fields"
            )
        else:
            self.fieldlist = fieldlist
        if diagnostics is not None:
            self.diagnostics = diagnostics
        else:
            self.diagnostics = Diagnostics(*fieldlist)
        if diagnostic_fields is not None:
            self.diagnostic_fields = diagnostic_fields
        else:
            self.diagnostic_fields = []
        if u_bc_ids is not None:
            self.u_bc_ids = u_bc_ids
        else:
            self.u_bc_ids = []

        # The mesh
        self.mesh = mesh

        # Build the spaces
        self._build_spaces(mesh, vertical_degree, horizontal_degree, family)

        # Allocate state
        self._allocate_state()
        if self.output.dumplist is None:
            self.output.dumplist = fieldlist
        self.fields = FieldCreator(fieldlist, self.xn, self.output.dumplist)

        # set up bcs
        V = self.fields('u').function_space()
        self.bcs = []
        if V.extruded:
            self.bcs.append(DirichletBC(V, 0.0, "bottom"))
            self.bcs.append(DirichletBC(V, 0.0, "top"))
        for id in self.u_bc_ids:
            self.bcs.append(DirichletBC(V, 0.0, id))

        self.dumpfile = None

        # figure out if we're on a sphere
        try:
            self.on_sphere = (mesh._base_mesh.geometric_dimension() == 3
                              and mesh._base_mesh.topological_dimension() == 2)
        except AttributeError:
            self.on_sphere = (mesh.geometric_dimension() == 3
                              and mesh.topological_dimension() == 2)

        #  build the vertical normal and define perp for 2d geometries
        dim = mesh.topological_dimension()
        if self.on_sphere:
            x = SpatialCoordinate(mesh)
            R = sqrt(inner(x, x))
            self.k = interpolate(x / R, mesh.coordinates.function_space())
            if dim == 2:
                outward_normals = CellNormal(mesh)
                self.perp = lambda u: cross(outward_normals, u)
        else:
            kvec = [0.0] * dim
            kvec[dim - 1] = 1.0
            self.k = Constant(kvec)
            if dim == 2:
                self.perp = lambda u: as_vector([-u[1], u[0]])

        # project test function for hydrostatic case
        if self.hydrostatic:
            self.h_project = lambda u: u - self.k * inner(u, self.k)
        else:
            self.h_project = lambda u: u

        #  Constant to hold current time
        self.t = Constant(0.0)

        # setup logger
        logger.setLevel(output.log_level)
        set_log_handler(mesh.comm)
        logger.info("Timestepping parameters that take non-default values:")
        logger.info(", ".join("%s: %s" % item
                              for item in vars(timestepping).items()))
        if parameters is not None:
            logger.info("Physical parameters that take non-default values:")
            logger.info(", ".join("%s: %s" % item
                                  for item in vars(parameters).items()))
Esempio n. 14
0
    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')
Esempio n. 15
0
    def __init__(
            self,
            state,
            family,
            degree,
            Omega=None,
            terms_to_linearise={
                'u': [time_derivative],
                'p': [time_derivative],
                'b': [time_derivative, transport]
            },
            u_transport_option="vector_invariant_form",
            no_normal_flow_bc_ids=None,
            active_tracers=None):

        field_names = ['u', 'p', 'b']

        if active_tracers is not None:
            raise NotImplementedError(
                'Tracers not implemented for Boussinesq equations')

        if active_tracers is None:
            active_tracers = []

        super().__init__(field_names,
                         state,
                         family,
                         degree,
                         terms_to_linearise=terms_to_linearise,
                         no_normal_flow_bc_ids=no_normal_flow_bc_ids,
                         active_tracers=active_tracers)

        w, phi, gamma = self.tests[0:3]
        u, p, b = split(self.X)
        bbar = state.fields("bbar", space=state.spaces("theta"), dump=False)
        bbar = state.fields("pbar", space=state.spaces("DG"), dump=False)

        # -------------------------------------------------------------------- #
        # Time Derivative Terms
        # -------------------------------------------------------------------- #
        mass_form = self.generate_mass_terms()

        # -------------------------------------------------------------------- #
        # Transport Terms
        # -------------------------------------------------------------------- #
        # Velocity transport term -- depends on formulation
        if u_transport_option == "vector_invariant_form":
            u_adv = prognostic(vector_invariant_form(state, w, u), "u")
        elif u_transport_option == "vector_advection_form":
            u_adv = prognostic(advection_form(state, w, u), "u")
        elif u_transport_option == "vector_manifold_advection_form":
            u_adv = prognostic(vector_manifold_advection_form(state, w, u),
                               "u")
        elif u_transport_option == "circulation_form":
            ke_form = kinetic_energy_form(state, w, u)
            ke_form = transport.remove(ke_form)
            ke_form = ke_form.label_map(
                lambda t: t.has_label(transporting_velocity), lambda t: Term(
                    ufl.replace(t.form, {t.get(transporting_velocity): u}), t.
                    labels))
            ke_form = transporting_velocity.remove(ke_form)
            u_adv = advection_equation_circulation_form(state, w, u) + ke_form
        else:
            raise ValueError("Invalid u_transport_option: %s" %
                             u_transport_option)

        # Buoyancy transport
        b_adv = prognostic(advection_form(state, gamma, b), "b")
        linear_b_adv = linear_advection_form(state, gamma, bbar).label_map(
            lambda t: t.has_label(transporting_velocity), lambda t: Term(
                ufl.replace(t.form, {
                    t.get(transporting_velocity): self.trials[0]
                }), t.labels))
        b_adv = linearisation(b_adv, linear_b_adv)

        adv_form = subject(u_adv + b_adv, self.X)

        # Add transport of tracers
        if len(active_tracers) > 0:
            adv_form += self.generate_tracer_transport_terms(
                state, active_tracers)

        # -------------------------------------------------------------------- #
        # Pressure Gradient Term
        # -------------------------------------------------------------------- #
        pressure_gradient_form = subject(prognostic(-div(w) * p * dx, "u"),
                                         self.X)

        # -------------------------------------------------------------------- #
        # Gravitational Term
        # -------------------------------------------------------------------- #
        gravity_form = subject(prognostic(-b * inner(w, state.k) * dx, "u"),
                               self.X)

        # -------------------------------------------------------------------- #
        # Divergence Term
        # -------------------------------------------------------------------- #
        # This enforces that div(u) = 0
        # The p features here so that the div(u) evaluated in the "forcing" step
        # replaces the whole pressure field, rather than merely providing an
        # increment to it.
        divergence_form = name(
            subject(prognostic(phi * (p - div(u)) * dx, "p"), self.X),
            "incompressibility")

        residual = (mass_form + adv_form + divergence_form +
                    pressure_gradient_form + gravity_form)

        # -------------------------------------------------------------------- #
        # Extra Terms (Coriolis)
        # -------------------------------------------------------------------- #
        if Omega is not None:
            residual += subject(
                prognostic(inner(w, cross(2 * Omega, u)) * dx, "u"), self.X)

        # -------------------------------------------------------------------- #
        # Linearise equations
        # -------------------------------------------------------------------- #
        # TODO: add linearisation states for variables
        # Add linearisations to equations
        self.residual = self.generate_linear_terms(residual,
                                                   self.terms_to_linearise)
Esempio n. 16
0
    def __init__(
            self,
            state,
            family,
            degree,
            Omega=None,
            sponge=None,
            extra_terms=None,
            terms_to_linearise={
                'u': [time_derivative],
                'rho': [time_derivative, transport],
                'theta': [time_derivative, transport]
            },
            u_transport_option="vector_invariant_form",
            diffusion_options=None,
            no_normal_flow_bc_ids=None,
            active_tracers=None):

        field_names = ['u', 'rho', 'theta']

        if active_tracers is None:
            active_tracers = []

        super().__init__(field_names,
                         state,
                         family,
                         degree,
                         terms_to_linearise=terms_to_linearise,
                         no_normal_flow_bc_ids=no_normal_flow_bc_ids,
                         active_tracers=active_tracers)

        g = state.parameters.g
        cp = state.parameters.cp

        w, phi, gamma = self.tests[0:3]
        u, rho, theta = split(self.X)[0:3]
        rhobar = state.fields("rhobar", space=state.spaces("DG"), dump=False)
        thetabar = state.fields("thetabar",
                                space=state.spaces("theta"),
                                dump=False)
        zero_expr = Constant(0.0) * theta
        exner = exner_pressure(state.parameters, rho, theta)
        n = FacetNormal(state.mesh)

        # -------------------------------------------------------------------- #
        # Time Derivative Terms
        # -------------------------------------------------------------------- #
        mass_form = self.generate_mass_terms()

        # -------------------------------------------------------------------- #
        # Transport Terms
        # -------------------------------------------------------------------- #
        # Velocity transport term -- depends on formulation
        if u_transport_option == "vector_invariant_form":
            u_adv = prognostic(vector_invariant_form(state, w, u), "u")
        elif u_transport_option == "vector_advection_form":
            u_adv = prognostic(advection_form(state, w, u), "u")
        elif u_transport_option == "vector_manifold_advection_form":
            u_adv = prognostic(vector_manifold_advection_form(state, w, u),
                               "u")
        elif u_transport_option == "circulation_form":
            ke_form = kinetic_energy_form(state, w, u)
            ke_form = transport.remove(ke_form)
            ke_form = ke_form.label_map(
                lambda t: t.has_label(transporting_velocity), lambda t: Term(
                    ufl.replace(t.form, {t.get(transporting_velocity): u}), t.
                    labels))
            ke_form = transporting_velocity.remove(ke_form)
            u_adv = advection_equation_circulation_form(state, w, u) + ke_form
        else:
            raise ValueError("Invalid u_transport_option: %s" %
                             u_transport_option)

        # Density transport (conservative form)
        rho_adv = prognostic(continuity_form(state, phi, rho), "rho")
        # Transport term needs special linearisation
        if transport in terms_to_linearise['rho']:
            linear_rho_adv = linear_continuity_form(
                state, phi, rhobar).label_map(
                    lambda t: t.has_label(transporting_velocity),
                    lambda t: Term(
                        ufl.replace(t.form, {
                            t.get(transporting_velocity):
                            self.trials[0]
                        }), t.labels))
            rho_adv = linearisation(rho_adv, linear_rho_adv)

        # Potential temperature transport (advective form)
        theta_adv = prognostic(advection_form(state, gamma, theta), "theta")
        # Transport term needs special linearisation
        if transport in terms_to_linearise['theta']:
            linear_theta_adv = linear_advection_form(
                state, gamma, thetabar).label_map(
                    lambda t: t.has_label(transporting_velocity),
                    lambda t: Term(
                        ufl.replace(t.form, {
                            t.get(transporting_velocity):
                            self.trials[0]
                        }), t.labels))
            theta_adv = linearisation(theta_adv, linear_theta_adv)

        adv_form = subject(u_adv + rho_adv + theta_adv, self.X)

        # Add transport of tracers
        if len(active_tracers) > 0:
            adv_form += self.generate_tracer_transport_terms(
                state, active_tracers)

        # -------------------------------------------------------------------- #
        # Pressure Gradient Term
        # -------------------------------------------------------------------- #
        # First get total mass of tracers
        tracer_mr_total = zero_expr
        for tracer in active_tracers:
            if tracer.variable_type == TracerVariableType.mixing_ratio:
                idx = self.field_names.index(tracer.name)
                tracer_mr_total += split(self.X)[idx]
            else:
                raise NotImplementedError(
                    'Only mixing ratio tracers are implemented')
        theta_v = theta / (Constant(1.0) + tracer_mr_total)

        pressure_gradient_form = name(
            subject(
                prognostic(
                    cp * (-div(theta_v * w) * exner * dx +
                          jump(theta_v * w, n) * avg(exner) * dS_v), "u"),
                self.X), "pressure_gradient")

        # -------------------------------------------------------------------- #
        # Gravitational Term
        # -------------------------------------------------------------------- #
        gravity_form = subject(
            prognostic(Term(g * inner(state.k, w) * dx), "u"), self.X)

        residual = (mass_form + adv_form + pressure_gradient_form +
                    gravity_form)

        # -------------------------------------------------------------------- #
        # Moist Thermodynamic Divergence Term
        # -------------------------------------------------------------------- #
        if len(active_tracers) > 0:
            cv = state.parameters.cv
            c_vv = state.parameters.c_vv
            c_pv = state.parameters.c_pv
            c_pl = state.parameters.c_pl
            R_d = state.parameters.R_d
            R_v = state.parameters.R_v

            # Get gas and liquid moisture mixing ratios
            mr_l = zero_expr
            mr_v = zero_expr

            for tracer in active_tracers:
                if tracer.is_moisture:
                    if tracer.variable_type == TracerVariableType.mixing_ratio:
                        idx = self.field_names.index(tracer.name)
                        if tracer.phase == Phases.gas:
                            mr_v += split(self.X)[idx]
                        elif tracer.phase == Phases.liquid:
                            mr_l += split(self.X)[idx]
                    else:
                        raise NotImplementedError(
                            'Only mixing ratio tracers are implemented')

            c_vml = cv + mr_v * c_vv + mr_l * c_pl
            c_pml = cp + mr_v * c_pv + mr_l * c_pl
            R_m = R_d + mr_v * R_v

            residual += subject(
                prognostic(
                    gamma * theta * div(u) * (R_m / c_vml - (R_d * c_pml) /
                                              (cp * c_vml)) * dx, "theta"),
                self.X)

        # -------------------------------------------------------------------- #
        # Extra Terms (Coriolis, Sponge, Diffusion and others)
        # -------------------------------------------------------------------- #
        if Omega is not None:
            residual += subject(
                prognostic(inner(w, cross(2 * Omega, u)) * dx, "u"), self.X)

        if sponge is not None:
            W_DG = FunctionSpace(state.mesh, "DG", 2)
            x = SpatialCoordinate(state.mesh)
            z = x[len(x) - 1]
            H = sponge.H
            zc = sponge.z_level
            assert float(zc) < float(
                H
            ), "you have set the sponge level above the height of your domain"
            mubar = sponge.mubar
            muexpr = conditional(
                z <= zc, 0.0,
                mubar * sin((pi / 2.) * (z - zc) / (H - zc))**2)
            self.mu = Function(W_DG).interpolate(muexpr)

            residual += name(
                subject(
                    prognostic(
                        self.mu * inner(w, state.k) * inner(u, state.k) * dx,
                        "u"), self.X), "sponge")

        if diffusion_options is not None:
            for field, diffusion in diffusion_options:
                idx = self.field_names.index(field)
                test = self.tests[idx]
                fn = split(self.X)[idx]
                residual += subject(
                    prognostic(
                        interior_penalty_diffusion_form(
                            state, test, fn, diffusion), field), self.X)

        if extra_terms is not None:
            for field, term in extra_terms:
                idx = self.field_names.index(field)
                test = self.tests[idx]
                residual += subject(prognostic(inner(test, term) * dx, field),
                                    self.X)

        # -------------------------------------------------------------------- #
        # Linearise equations
        # -------------------------------------------------------------------- #
        # TODO: add linearisation states for variables
        # Add linearisations to equations
        self.residual = self.generate_linear_terms(residual,
                                                   self.terms_to_linearise)