From 50ef06bb0e6b28bdc605c25459ff645561741a2f Mon Sep 17 00:00:00 2001 From: Tom Bendall Date: Mon, 18 Nov 2024 08:11:04 +0000 Subject: [PATCH] add and implement augmentation with vorticity transport --- gusto/core/domain.py | 9 +- gusto/spatial_methods/__init__.py | 1 + gusto/spatial_methods/augmentation.py | 159 ++++++++++++++++++ .../explicit_runge_kutta.py | 35 ++-- .../implicit_runge_kutta.py | 8 +- .../time_discretisation.py | 109 ++++++++++-- 6 files changed, 285 insertions(+), 36 deletions(-) create mode 100644 gusto/spatial_methods/augmentation.py diff --git a/gusto/core/domain.py b/gusto/core/domain.py index 0cf23fb7f..6d873d962 100644 --- a/gusto/core/domain.py +++ b/gusto/core/domain.py @@ -6,9 +6,10 @@ from gusto.core.coordinates import Coordinates from gusto.core.function_spaces import Spaces, check_degree_args -from firedrake import (Constant, SpatialCoordinate, sqrt, CellNormal, cross, - inner, grad, VectorFunctionSpace, Function, FunctionSpace, - perp) +from firedrake import ( + Constant, SpatialCoordinate, sqrt, CellNormal, cross, inner, grad, + VectorFunctionSpace, Function, FunctionSpace, perp, curl +) import numpy as np @@ -113,12 +114,14 @@ def __init__(self, mesh, dt, family, degree=None, V = VectorFunctionSpace(mesh, "DG", sphere_degree) self.outward_normals = Function(V).interpolate(CellNormal(mesh)) self.perp = lambda u: cross(self.outward_normals, u) + self.divperp = lambda u: inner(self.outward_normals, curl(u)) else: kvec = [0.0]*dim kvec[dim-1] = 1.0 self.k = Constant(kvec) if dim == 2: self.perp = perp + self.divperp = lambda u: -u[0].dx(1) + u[1].dx(0) # -------------------------------------------------------------------- # # Construct information relating to height/radius diff --git a/gusto/spatial_methods/__init__.py b/gusto/spatial_methods/__init__.py index baa1e0788..349094cdf 100644 --- a/gusto/spatial_methods/__init__.py +++ b/gusto/spatial_methods/__init__.py @@ -2,3 +2,4 @@ from gusto.spatial_methods.diffusion_methods import * # noqa from gusto.spatial_methods.transport_methods import * # noqa from gusto.spatial_methods.limiters import * # noqa +from gusto.spatial_methods.augmentation import * # noqa \ No newline at end of file diff --git a/gusto/spatial_methods/augmentation.py b/gusto/spatial_methods/augmentation.py new file mode 100644 index 000000000..d78b3ce07 --- /dev/null +++ b/gusto/spatial_methods/augmentation.py @@ -0,0 +1,159 @@ +from abc import ABCMeta +from firedrake import ( + MixedFunctionSpace, Function, TestFunctions, split, inner, dx, grad, + LinearVariationalProblem, LinearVariationalSolver, lhs, rhs, dot, + ds_b, ds_v, ds_t, ds, FacetNormal, TestFunction, TrialFunction, + transpose, nabla_grad, outer, dS, dS_h, dS_v, sign, jump, div, + Constant, sqrt, cross, curl +) +from firedrake.fml import subject +from gusto import ( + time_derivative, transport, transporting_velocity, TransportEquationType, + logger +) + + +class Augmentation(object, metaclass=ABCMeta): + """ + Augments an equation with another equation to be solved simultaneously. + """ + + +class VorticityTransport(Augmentation): + """ + Solves the transport of a velocity field, simultaneously with the vorticity. + """ + + ### An argument to time discretisation or spatial method?? + # TODO: this all needs to be generalised + + def __init__(self, domain, V_vel, V_vort, transpose_commutator=True, + supg=False, min_dx=None): + + self.fs = MixedFunctionSpace((V_vel, V_vort)) + self.X = Function(self.fs) + self.tests = TestFunctions(self.fs) + + u = Function(V_vel) + F, Z = split(self.X) + test_F, test_Z = self.tests + + if hasattr(domain.mesh, "_base_mesh"): + self.ds = ds_b + ds_t + ds_v + self.dS = dS_v + dS_h + else: + self.ds = ds + self.dS = dS + + n = FacetNormal(domain.mesh) + sign_u = 0.5*(sign(dot(u, n)) + 1) + upw = lambda f: (sign_u('+')*f('+') + sign_u('-')*f('-')) + + if domain.mesh.topological_dimension() == 2: + mix_test = test_F - domain.perp(grad(test_Z)) + F_cross_u = Z*domain.perp(u) + elif domain.mesh.topological_dimension == 3: + mix_test = test_F - curl(test_Z) + F_cross_u = cross(Z, u) + + time_deriv_form = inner(F, test_F)*dx + inner(Z, test_Z)*dx + + # Standard vector invariant transport form ----------------------------- + transport_form = ( + # vorticity term + inner(mix_test, F_cross_u)*dx + + inner(n, test_Z*Z*u)*self.ds + # 0.5*grad(v . F) + - 0.5 * div(mix_test) * inner(u, F)*dx + + 0.5 * inner(mix_test, n) * inner(u, F)*self.ds + ) + + # Communtator of tranpose gradient terms ------------------------------- + # This is needed for general vector transport + if transpose_commutator: + u_dot_nabla_F = dot(u, transpose(nabla_grad(F))) + transport_form += ( + - inner(n, test_Z*domain.perp(u_dot_nabla_F))*self.ds + # + 0.5*grad(F).v + - 0.5 * dot(F, div(outer(u, mix_test)))*dx + + 0.5 * inner(mix_test('+'), n('+'))*dot(jump(u), upw(F))*self.dS + # - 0.5*grad(v).F + + 0.5 * dot(u, div(outer(F, mix_test)))*dx + - 0.5 * inner(mix_test('+'), n('+'))*dot(jump(F), upw(u))*self.dS + ) + + # SUPG terms ----------------------------------------------------------- + # Add the vorticity residual to the transported vorticity, + # which damps enstrophy + if supg: + if min_dx is not None: + lamda = Constant(0.5) + #TODO: decide on expression here + # tau = 0.5 / (lamda/domain.dt + sqrt(dot(u, u))/Constant(min_dx)) + tau = 0.5*domain.dt*(1.0 + sqrt(dot(u, u))*domain.dt/Constant(min_dx)) + else: + tau = 0.5*domain.dt + + dxqp = dx(degree=3) + + if domain.mesh.topological_dimension() == 2: + time_deriv_form -= inner(mix_test, tau*Z*domain.perp(u)/domain.dt)*dxqp + transport_form -= inner( + mix_test, tau*domain.perp(u)*domain.divperp(Z*domain.perp(u)) + )*dxqp + if transpose_commutator: + transport_form -= inner( + mix_test, + tau*domain.perp(u)*domain.divperp(u_dot_nabla_F) + )*dxqp + elif domain.mesh.topological_dimension() == 3: + time_deriv_form -= inner(mix_test, tau*cross(Z, u)/domain.dt)*dxqp + transport_form -= inner( + mix_test, tau*cross(curl(Z*u), u) + )*dxqp + if transpose_commutator: + transport_form -= inner( + mix_test, + tau*cross(curl(u_dot_nabla_F), u) + )*dxqp + + residual = ( + time_derivative(time_deriv_form) + + transport( + transport_form, TransportEquationType.vector_invariant + ) + ) + residual = transporting_velocity(residual, u) + + self.residual = subject(residual, self.X) + + self.x_in = Function(self.fs) + self.Z_in = Function(V_vort) + self.x_out = Function(self.fs) + + vort_test = TestFunction(V_vort) + vort_trial = TrialFunction(V_vort) + + F_in, _ = split(self.x_in) + + eqn = ( + inner(vort_trial, vort_test)*dx + + inner(domain.perp(grad(vort_test)), F_in)*dx + + vort_test*inner(n, domain.perp(F_in))*self.ds + ) + problem = LinearVariationalProblem( + lhs(eqn), rhs(eqn), self.Z_in, constant_jacobian=True + ) + self.solver = LinearVariationalSolver(problem) + + def pre_apply(self, x_in): + self.x_in.subfunctions[0].assign(x_in) + + def post_apply(self, x_out): + x_out.assign(self.x_out.subfunctions[0]) + + def update(self, x_in_mixed): + self.x_in.subfunctions[0].assign(x_in_mixed.subfunctions[0]) + logger.info('Vorticity solve') + self.solver.solve() + self.x_in.subfunctions[1].assign(self.Z_in) diff --git a/gusto/time_discretisation/explicit_runge_kutta.py b/gusto/time_discretisation/explicit_runge_kutta.py index dabdf93ec..ee712027c 100644 --- a/gusto/time_discretisation/explicit_runge_kutta.py +++ b/gusto/time_discretisation/explicit_runge_kutta.py @@ -89,7 +89,8 @@ class ExplicitRungeKutta(ExplicitTimeDiscretisation): def __init__(self, domain, butcher_matrix, field_name=None, fixed_subcycles=None, subcycle_by_courant=None, rk_formulation=RungeKuttaFormulation.increment, - solver_parameters=None, limiter=None, options=None): + solver_parameters=None, limiter=None, options=None, + augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -123,7 +124,8 @@ def __init__(self, domain, butcher_matrix, field_name=None, fixed_subcycles=fixed_subcycles, subcycle_by_courant=subcycle_by_courant, solver_parameters=solver_parameters, - limiter=limiter, options=options) + limiter=limiter, options=options, + augmentation=augmentation) self.butcher_matrix = butcher_matrix self.nbutcher = int(np.shape(self.butcher_matrix)[0]) self.rk_formulation = rk_formulation @@ -210,7 +212,7 @@ def lhs(self): if self.rk_formulation == RungeKuttaFormulation.increment: l = self.residual.label_map( lambda t: t.has_label(time_derivative), - map_if_true=replace_subject(self.x_out, self.idx), + map_if_true=replace_subject(self.x_out, old_idx=self.idx, new_idx=self.new_idx), map_if_false=drop) return l.form @@ -220,7 +222,7 @@ def lhs(self): for stage in range(self.nStages): l = self.residual.label_map( lambda t: t.has_label(time_derivative), - map_if_true=replace_subject(self.field_i[stage+1], self.idx), + map_if_true=replace_subject(self.field_i[stage+1], old_idx=self.idx, new_idx=self.new_idx), map_if_false=drop) lhs_list.append(l) @@ -229,7 +231,7 @@ def lhs(self): if self.rk_formulation == RungeKuttaFormulation.linear: l = self.residual.label_map( lambda t: t.has_label(time_derivative), - map_if_true=replace_subject(self.x1, self.idx), + map_if_true=replace_subject(self.x1, old_idx=self.idx, new_idx=self.new_idx), map_if_false=drop) return l.form @@ -246,7 +248,7 @@ def rhs(self): if self.rk_formulation == RungeKuttaFormulation.increment: r = self.residual.label_map( all_terms, - map_if_true=replace_subject(self.x1, old_idx=self.idx)) + map_if_true=replace_subject(self.x1, old_idx=self.idx, new_idx=self.new_idx)) r = r.label_map( lambda t: t.has_label(time_derivative), @@ -273,7 +275,7 @@ def rhs(self): for stage in range(self.nStages): r = self.residual.label_map( all_terms, - map_if_true=replace_subject(self.field_i[0], old_idx=self.idx)) + map_if_true=replace_subject(self.field_i[0], old_idx=self.idx, new_idx=self.new_idx)) r = r.label_map( lambda t: t.has_label(time_derivative), @@ -284,7 +286,7 @@ def rhs(self): r_i = self.residual.label_map( lambda t: t.has_label(time_derivative), map_if_true=drop, - map_if_false=replace_subject(self.field_i[i], old_idx=self.idx) + map_if_false=replace_subject(self.field_i[i], old_idx=self.idx, new_idx=self.new_idx) ) r -= self.butcher_matrix[stage, i]*self.dt*r_i @@ -297,8 +299,8 @@ def rhs(self): r = self.residual.label_map( lambda t: t.has_label(time_derivative), - map_if_true=replace_subject(self.x0, old_idx=self.idx), - map_if_false=replace_subject(self.field_rhs, old_idx=self.idx) + map_if_true=replace_subject(self.x0, old_idx=self.idx, new_idx=self.new_idx), + map_if_false=replace_subject(self.field_rhs, old_idx=self.idx, new_idx=self.new_idx) ) r = r.label_map( lambda t: t.has_label(time_derivative), @@ -325,8 +327,8 @@ def rhs(self): ) r_all_but_last = r_all_but_last.label_map( lambda t: t.has_label(time_derivative), - map_if_true=replace_subject(self.x0, old_idx=self.idx), - map_if_false=replace_subject(self.field_rhs, old_idx=self.idx) + map_if_true=replace_subject(self.x0, old_idx=self.idx, new_idx=self.new_idx), + map_if_false=replace_subject(self.field_rhs, old_idx=self.idx, new_idx=self.new_idx) ) r_all_but_last = r_all_but_last.label_map( lambda t: t.has_label(time_derivative), @@ -468,6 +470,9 @@ def apply_cycle(self, x_out, x_in): x_out (:class:`Function`): the output field to be computed. """ + if self.augmentation is not None: + self.augmentation.update(x_in) + # TODO: is this limiter application necessary? if self.limiter is not None: self.limiter.apply(x_in) @@ -546,7 +551,8 @@ def __init__( self, domain, field_name=None, fixed_subcycles=None, subcycle_by_courant=None, rk_formulation=RungeKuttaFormulation.increment, - solver_parameters=None, limiter=None, options=None + solver_parameters=None, limiter=None, options=None, + augmentation=None ): """ Args: @@ -586,7 +592,8 @@ def __init__( subcycle_by_courant=subcycle_by_courant, rk_formulation=rk_formulation, solver_parameters=solver_parameters, - limiter=limiter, options=options) + limiter=limiter, options=options, + augmentation=augmentation) class RK4(ExplicitRungeKutta): diff --git a/gusto/time_discretisation/implicit_runge_kutta.py b/gusto/time_discretisation/implicit_runge_kutta.py index 8eb1c75ab..b220ef0be 100644 --- a/gusto/time_discretisation/implicit_runge_kutta.py +++ b/gusto/time_discretisation/implicit_runge_kutta.py @@ -56,7 +56,7 @@ class ImplicitRungeKutta(TimeDiscretisation): # --------------------------------------------------------------------------- def __init__(self, domain, butcher_matrix, field_name=None, - solver_parameters=None, options=None,): + solver_parameters=None, options=None, augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -75,7 +75,7 @@ def __init__(self, domain, butcher_matrix, field_name=None, """ super().__init__(domain, field_name=field_name, solver_parameters=solver_parameters, - options=options) + options=options, augmentation=augmentation) self.butcher_matrix = butcher_matrix self.nStages = int(np.shape(self.butcher_matrix)[1]) @@ -165,7 +165,7 @@ class ImplicitMidpoint(ImplicitRungeKutta): y^(n+1) = y^n + dt*k0 \n """ def __init__(self, domain, field_name=None, solver_parameters=None, - options=None): + options=None, augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -182,7 +182,7 @@ def __init__(self, domain, field_name=None, solver_parameters=None, butcher_matrix = np.array([[0.5], [1.]]) super().__init__(domain, butcher_matrix, field_name, solver_parameters=solver_parameters, - options=options) + options=options, augmentation=augmentation) class QinZhang(ImplicitRungeKutta): diff --git a/gusto/time_discretisation/time_discretisation.py b/gusto/time_discretisation/time_discretisation.py index df108a615..35bfe7bfd 100644 --- a/gusto/time_discretisation/time_discretisation.py +++ b/gusto/time_discretisation/time_discretisation.py @@ -30,7 +30,17 @@ def wrapper_apply(original_apply): """Decorator to add steps for using a wrapper around the apply method.""" def get_apply(self, x_out, x_in): - if self.wrapper is not None: + if self.augmentation is not None: + + def new_apply(self, x_out, x_in): + + self.augmentation.pre_apply(x_in) + original_apply(self, self.augmentation.x_out, self.augmentation.x_in) + self.augmentation.post_apply(x_out) + + return new_apply(self, x_out, x_in) + + elif self.wrapper is not None: def new_apply(self, x_out, x_in): @@ -51,7 +61,7 @@ class TimeDiscretisation(object, metaclass=ABCMeta): """Base class for time discretisation schemes.""" def __init__(self, domain, field_name=None, solver_parameters=None, - limiter=None, options=None): + limiter=None, options=None, augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -78,6 +88,7 @@ def __init__(self, domain, field_name=None, solver_parameters=None, self.options = options self.limiter = limiter self.courant_max = None + self.augmentation = augmentation if options is not None: self.wrapper_name = options.name @@ -146,6 +157,14 @@ def setup(self, equation, apply_bcs=True, *active_labels): self.fs = equation.function_space self.idx = None + if self.augmentation is not None: + self.fs = self.augmentation.fs + self.residual = self.augmentation.residual + self.idx = None + self.new_idx = None + else: + self.new_idx = None + bcs = equation.bcs[self.field_name] if len(active_labels) > 0: @@ -266,7 +285,7 @@ def lhs(self): """Set up the discretisation's left hand side (the time derivative).""" l = self.residual.label_map( lambda t: t.has_label(time_derivative), - map_if_true=replace_subject(self.x_out, old_idx=self.idx), + map_if_true=replace_subject(self.x_out, old_idx=self.idx, new_idx=self.new_idx), map_if_false=drop) return l.form @@ -276,7 +295,7 @@ def rhs(self): """Set up the time discretisation's right hand side.""" r = self.residual.label_map( all_terms, - map_if_true=replace_subject(self.x1, old_idx=self.idx)) + map_if_true=replace_subject(self.x1, old_idx=self.idx, new_idx=self.new_idx)) r = r.label_map( lambda t: t.has_label(time_derivative), @@ -316,7 +335,7 @@ class ExplicitTimeDiscretisation(TimeDiscretisation): def __init__(self, domain, field_name=None, fixed_subcycles=None, subcycle_by_courant=None, solver_parameters=None, limiter=None, - options=None): + options=None, augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -343,7 +362,8 @@ def __init__(self, domain, field_name=None, fixed_subcycles=None, """ super().__init__(domain, field_name, solver_parameters=solver_parameters, - limiter=limiter, options=options) + limiter=limiter, options=options, + augmentation=augmentation) if fixed_subcycles is not None and subcycle_by_courant is not None: raise ValueError('Cannot specify both subcycle and subcycle_by ' @@ -438,11 +458,12 @@ def apply(self, x_out, x_in): """ # If doing adaptive subcycles, update dt and ncycles here if self.subcycle_by_courant is not None: - self.ncycles = math.ceil(float(self.courant_max)/self.subcycle_by_courant) + self.ncycles = min(12, math.ceil(float(self.courant_max)/self.subcycle_by_courant)) self.dt.assign(self.original_dt/self.ncycles) self.x0.assign(x_in) for i in range(self.ncycles): + logger.info(f'Subcycle {i}') self.subcycle_idx = i self.apply_cycle(self.x1, self.x0) self.x0.assign(self.x1) @@ -539,7 +560,9 @@ class ThetaMethod(TimeDiscretisation): """ def __init__(self, domain, theta, field_name=None, - solver_parameters=None, options=None): + solver_parameters=None, options=None, + fixed_subcycles=None, subcycle_by_courant=None, + augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -569,12 +592,43 @@ def __init__(self, domain, theta, field_name=None, 'pc_type': 'bjacobi', 'sub_pc_type': 'ilu'} + if fixed_subcycles is not None and subcycle_by_courant is not None: + raise ValueError('Cannot specify both subcycle and subcycle_by ' + + 'arguments to a time discretisation') + self.fixed_subcycles = fixed_subcycles + self.subcycle_by_courant = subcycle_by_courant + super().__init__(domain, field_name, solver_parameters=solver_parameters, - options=options) + options=options, + augmentation=augmentation) self.theta = theta + def setup(self, equation, apply_bcs=True, *active_labels): + """ + Set up the time discretisation based on the equation. + + Args: + equation (:class:`PrognosticEquation`): the model's equation. + apply_bcs (bool, optional): whether boundary conditions are to be + applied. Defaults to True. + *active_labels (:class:`Label`): labels indicating which terms of + the equation to include. + """ + super().setup(equation, apply_bcs, *active_labels) + + # if user has specified a number of fixed subcycles, then save this + # and rescale dt accordingly; else perform just one cycle using dt + if self.fixed_subcycles is not None: + self.dt.assign(self.dt/self.fixed_subcycles) + self.ncycles = self.fixed_subcycles + else: + self.dt = self.dt + self.ncycles = 1 + self.x0 = Function(self.fs) + self.x1 = Function(self.fs) + @cached_property def lhs(self): """Set up the discretisation's left hand side (the time derivative).""" @@ -597,21 +651,44 @@ def rhs(self): return r.form - @wrapper_apply - def apply(self, x_out, x_in): + def apply_cycle(self, x_out, x_in): """ - Apply the time discretisation to advance one whole time step. + Apply the time discretisation through a single sub-step. Args: x_out (:class:`Function`): the output field to be computed. x_in (:class:`Function`): the input field. """ + if self.augmentation is not None: + self.augmentation.update(x_in) self.x1.assign(x_in) # Set initial solver guess self.x_out.assign(x_in) self.solver.solve() x_out.assign(self.x_out) + @wrapper_apply + def apply(self, x_out, x_in): + """ + Apply the time discretisation to advance one whole time step. + + Args: + x_out (:class:`Function`): the output field to be computed. + x_in (:class:`Function`): the input field. + """ + # If doing adaptive subcycles, update dt and ncycles here + if self.subcycle_by_courant is not None: + self.ncycles = min(12, math.ceil(float(self.courant_max)/self.subcycle_by_courant)) + self.dt.assign(self.original_dt/self.ncycles) + + self.x0.assign(x_in) + for i in range(self.ncycles): + logger.info(f'Subcycle {i}') + self.subcycle_idx = i + self.apply_cycle(self.x1, self.x0) + self.x0.assign(self.x1) + x_out.assign(self.x1) + class TrapeziumRule(ThetaMethod): """ @@ -624,7 +701,8 @@ class TrapeziumRule(ThetaMethod): """ def __init__(self, domain, field_name=None, solver_parameters=None, - options=None): + options=None, fixed_subcycles=None, subcycle_by_courant=None, + augmentation=None): """ Args: domain (:class:`Domain`): the model's domain object, containing the @@ -640,8 +718,9 @@ def __init__(self, domain, field_name=None, solver_parameters=None, """ super().__init__(domain, 0.5, field_name, solver_parameters=solver_parameters, - options=options) - + options=options, fixed_subcycles=fixed_subcycles, + subcycle_by_courant=subcycle_by_courant, + augmentation=augmentation) class TR_BDF2(TimeDiscretisation): """