diff --git a/examples/shallow_water/mm_ot_sw_galewsky_jet.py b/examples/shallow_water/mm_ot_sw_galewsky_jet.py new file mode 100644 index 000000000..33db53f0e --- /dev/null +++ b/examples/shallow_water/mm_ot_sw_galewsky_jet.py @@ -0,0 +1,223 @@ +from gusto import * +from firedrake import IcosahedralSphereMesh, \ + Constant, ge, le, exp, cos, \ + conditional, interpolate, SpatialCoordinate, VectorFunctionSpace, \ + Function, assemble, dx, pi, CellNormal + +import numpy as np +import sys + +day = 24.*60.*60. +dt = 240. + +if '--running-tests' in sys.argv: + ref_level = 3 + dt = 480. + tmax = 1440. + cubed_sphere = False +else: + # setup resolution and timestepping parameters + cubed_sphere = True + if cubed_sphere: + ncells = 24 + dt = 200. + else: + ref_level = 4 + dt = 240. + tmax = 6*day + +# setup shallow water parameters +R = 6371220. +H = 10000. +parameters = ShallowWaterParameters(H=H) + +perturb = True +if perturb: + dirname = "mm_ot_sw_galewsky_jet_perturbed_dt%s" % dt +else: + dirname = "mm_ot_sw_galewsky_jet_unperturbed" + +# Domain +if cubed_sphere: + dirname += "_cs%s" % ncells + mesh = GeneralCubedSphereMesh(radius=R, + num_cells_per_edge_of_panel=ncells, + degree=2) + family = "RTCF" +else: + dirname += "_is%s" % ref_level + mesh = IcosahedralSphereMesh(radius=R, + refinement_level=ref_level, degree=2) + family = "BDM" +domain = Domain(mesh, dt, family, 1, move_mesh=True) + +# Equation +Omega = parameters.Omega +x = SpatialCoordinate(domain.mesh) +fexpr = 2*Omega*x[2]/R +eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, + u_transport_option="vector_invariant_form") + +# I/O +output = OutputParameters(dirname=dirname, + dumplist_latlon=['D', 'PotentialVorticity', + 'RelativeVorticity']) + +pv = PotentialVorticity() +diagnostic_fields = [pv, RelativeVorticity()] +io = IO(domain, output, diagnostic_fields=diagnostic_fields) + +# Transport schemes +transported_fields = [TrapeziumRule(domain, "u"), + SSPRK3(domain, "D")] + +transport_methods = [DGUpwind(eqns, "u"), DGUpwind(eqns, "D")] + + +# Mesh movement +def update_pv(): + pv() + + +def reinterpolate_coriolis(): + domain.k = interpolate(x/R, domain.mesh.coordinates.function_space()) + domain.outward_normals.interpolate(CellNormal(domain.mesh)) + eqns.prescribed_fields("coriolis").interpolate(fexpr) + + +monitor = MonitorFunction("PotentialVorticity", adapt_to="gradient") +mesh_generator = OptimalTransportMeshGenerator(domain.mesh, + monitor, + pre_meshgen_callback=update_pv, + post_meshgen_callback=reinterpolate_coriolis) + +# Time stepper +stepper = MeshMovement(eqns, io, transported_fields, + spatial_methods=transport_methods, + mesh_generator=mesh_generator) + +# initial conditions +u0 = stepper.fields("u") +D0 = stepper.fields("D") + +# get lat lon coordinates +theta, lamda = latlon_coords(mesh) + +# expressions for meridional and zonal velocity +u_max = 80.0 +theta0 = pi/7. +theta1 = pi/2. - theta0 +en = np.exp(-4./((theta1-theta0)**2)) +u_zonal_expr = (u_max/en)*exp(1/((theta - theta0)*(theta - theta1))) +u_zonal = conditional(ge(theta, theta0), conditional(le(theta, theta1), u_zonal_expr, 0.), 0.) +u_merid = 0.0 + +# get cartesian components of velocity +uexpr = sphere_to_cartesian(mesh, u_zonal, u_merid) +u0.project(uexpr, form_compiler_parameters={'quadrature_degree': 12}) + +Rc = Constant(R) +g = Constant(parameters.g) + + +def D_integrand(th): + # Initial D field is calculated by integrating D_integrand w.r.t. theta + # Assumes the input is between theta0 and theta1. + # Note that this function operates on vectorized input. + from scipy import exp, sin, tan + f = 2.0*parameters.Omega*sin(th) + u_zon = (80.0/en)*exp(1.0/((th - theta0)*(th - theta1))) + return u_zon*(f + tan(th)*u_zon/R) + + +def Dval(X): + # Function to return value of D at X + from scipy import integrate + + # Preallocate output array + val = np.zeros(len(X)) + + angles = np.zeros(len(X)) + + # Minimize work by only calculating integrals for points with + # theta between theta_0 and theta_1. + # For theta <= theta_0, the integral is 0 + # For theta >= theta_1, the integral is constant. + + # Precalculate this constant: + poledepth, _ = integrate.fixed_quad(D_integrand, theta0, theta1, n=64) + poledepth *= -R/parameters.g + + angles[:] = np.arcsin(X[:, 2]/R) + + for ii in range(len(X)): + if angles[ii] <= theta0: + val[ii] = 0.0 + elif angles[ii] >= theta1: + val[ii] = poledepth + else: + # Fixed quadrature with 64 points gives absolute errors below 1e-13 + # for a quantity of order 1e-3. + v, _ = integrate.fixed_quad(D_integrand, theta0, angles[ii], n=64) + val[ii] = -(R/parameters.g)*v + + return val + + +# Get coordinates to pass to Dval function +W = VectorFunctionSpace(mesh, D0.ufl_element()) +X = interpolate(mesh.coordinates, W) +D0.dat.data[:] = Dval(X.dat.data_ro) + +# Adjust mean value of initial D +C = Function(D0.function_space()).assign(Constant(1.0)) +area = assemble(C*dx) +Dmean = assemble(D0*dx)/area +D0 -= Dmean +D0 += Constant(parameters.H) + +# optional perturbation +if perturb: + alpha = Constant(1/3.) + beta = Constant(1/15.) + Dhat = Constant(120.) + theta2 = Constant(pi/4.) + g = Constant(parameters.g) + D_pert = Function(D0.function_space()).interpolate(Dhat*cos(theta)*exp(-(lamda/alpha)**2)*exp(-((theta2 - theta)/beta)**2)) + D0 += D_pert + + +def initialise_fn(): + u0 = stepper.fields("u") + D0 = stepper.fields("D") + + u0.project(uexpr, form_compiler_parameters={'quadrature_degree': 12}) + + X = interpolate(domain.mesh.coordinates, W) + D0.dat.data[:] = Dval(X.dat.data_ro) + area = assemble(C*dx) + Dmean = assemble(D0*dx)/area + D0 -= Dmean + D0 += parameters.H + if perturb: + theta, lamda = latlon_coords(domain.mesh) + D_pert.interpolate(Dhat*cos(theta)*exp(-(lamda/alpha)**2)*exp(-((theta2 - theta)/beta)**2)) + D0 += D_pert + domain.k = interpolate(x/R, domain.mesh.coordinates.function_space()) + domain.outward_normals.interpolate(CellNormal(domain.mesh)) + eqns.prescribed_fields("coriolis").interpolate(fexpr) + pv() + + +pv.setup(domain, stepper.fields) +mesh_generator.get_first_mesh(initialise_fn) + +domain.k = interpolate(x/R, domain.mesh.coordinates.function_space()) +domain.outward_normals.interpolate(CellNormal(domain.mesh)) +eqns.prescribed_fields("coriolis").interpolate(fexpr) +pv() + +Dbar = Function(D0.function_space()).assign(H) +stepper.set_reference_profiles([('D', Dbar)]) + +stepper.run(t=0, tmax=tmax) diff --git a/gusto/__init__.py b/gusto/__init__.py index 772e94656..cf9547c24 100644 --- a/gusto/__init__.py +++ b/gusto/__init__.py @@ -32,6 +32,7 @@ def perp(self, o, a): from gusto.limiters import * # noqa from gusto.linear_solvers import * # noqa from gusto.meshes import * # noqa +from gusto.moving_mesh import * # noqa from gusto.numerical_integrator import * # noqa from gusto.physics import * # noqa from gusto.preconditioners import * # noqa diff --git a/gusto/diagnostics.py b/gusto/diagnostics.py index b8e1212af..a3d1176ca 100644 --- a/gusto/diagnostics.py +++ b/gusto/diagnostics.py @@ -1452,7 +1452,7 @@ def setup(self, domain, state_fields, vorticity_type=None): f = state_fields("coriolis") L += gamma*f*dx - problem = LinearVariationalProblem(a, L, self.field) + problem = LinearVariationalProblem(a, L, self.field, constant_jacobian=False) self.evaluator = LinearVariationalSolver(problem, solver_parameters={"ksp_type": "cg"}) diff --git a/gusto/domain.py b/gusto/domain.py index 0528381ad..bbf126356 100644 --- a/gusto/domain.py +++ b/gusto/domain.py @@ -25,7 +25,8 @@ class Domain(object): the same then this can be specified through the "degree" argument. """ def __init__(self, mesh, dt, family, degree=None, - horizontal_degree=None, vertical_degree=None): + horizontal_degree=None, vertical_degree=None, + move_mesh=False): """ Args: mesh (:class:`Mesh`): the model's mesh. @@ -59,6 +60,9 @@ def __init__(self, mesh, dt, family, degree=None, else: raise TypeError(f'dt must be a Constant, float or int, not {type(dt)}') + # store whether we are moving the mesh + self.move_mesh = move_mesh + # -------------------------------------------------------------------- # # Build compatible function spaces # -------------------------------------------------------------------- # diff --git a/gusto/equations.py b/gusto/equations.py index dd4a9ceb4..5a10b51db 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -10,7 +10,8 @@ from gusto.fml import (Term, all_terms, keep, drop, Label, subject, name, replace_subject, replace_trial_function) from gusto.labels import (time_derivative, transport, prognostic, hydrostatic, - linearisation, pressure_gradient, coriolis) + linearisation, pressure_gradient, coriolis, + transporting_velocity) from gusto.thermodynamics import exner_pressure from gusto.common_forms import (advection_form, continuity_form, vector_invariant_form, kinetic_energy_form, @@ -236,6 +237,7 @@ def __init__(self, field_names, domain, space_names, # Make the full mixed function space W = MixedFunctionSpace(self.spaces) + self.W = W # Can now call the underlying PrognosticEquation full_field_name = "_".join(self.field_names) @@ -668,14 +670,29 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, # -------------------------------------------------------------------- # # Transport Terms # -------------------------------------------------------------------- # + + # Mesh movement requires the circulation form, and an + # additional modification + if domain.move_mesh: + assert u_transport_option == "vector_invariant_form" + # Velocity transport term -- depends on formulation if u_transport_option == "vector_invariant_form": u_adv = prognostic(vector_invariant_form(domain, w, u, u), 'u') elif u_transport_option == "vector_advection_form": u_adv = prognostic(advection_form(w, u, u), 'u') + if domain.move_mesh: + ke_form = prognostic(kinetic_energy_form(w, u, u), "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 -= ke_form elif u_transport_option == "circulation_form": - ke_form = prognostic(kinetic_energy_form(w, u, u), 'u') - u_adv = prognostic(advection_equation_circulation_form(domain, w, u, u), 'u') + ke_form + ke_form = prognostic(kinetic_energy_form(w, u, u), "u") + u_adv = prognostic(advection_equation_circulation_form(domain, w, u, u), "u") + ke_form else: raise ValueError("Invalid u_transport_option: %s" % u_transport_option) diff --git a/gusto/forcing.py b/gusto/forcing.py index 662463515..cfeb28947 100644 --- a/gusto/forcing.py +++ b/gusto/forcing.py @@ -90,11 +90,13 @@ def __init__(self, equation, alpha): # now we can set up the explicit and implicit problems explicit_forcing_problem = LinearVariationalProblem( - a.form, L_explicit.form, self.xF, bcs=bcs + a.form, L_explicit.form, self.xF, bcs=bcs, + constant_jacobian=False ) implicit_forcing_problem = LinearVariationalProblem( - a.form, L_implicit.form, self.xF, bcs=bcs + a.form, L_implicit.form, self.xF, bcs=bcs, + constant_jacobian=False ) self.solvers = {} diff --git a/gusto/linear_solvers.py b/gusto/linear_solvers.py index 68ddf1477..5bdc6f3d2 100644 --- a/gusto/linear_solvers.py +++ b/gusto/linear_solvers.py @@ -599,7 +599,8 @@ def __init__(self, equation, alpha): bcs = [DirichletBC(W.sub(0), bc.function_arg, bc.sub_domain) for bc in equation.bcs['u']] problem = LinearVariationalProblem(aeqn.form, action(Leqn.form, self.xrhs), - self.dy, bcs=bcs) + self.dy, bcs=bcs, + constant_jacobian=False) self.solver = LinearVariationalSolver(problem, solver_parameters=self.solver_parameters, diff --git a/gusto/moving_mesh/__init__.py b/gusto/moving_mesh/__init__.py new file mode 100644 index 000000000..85449fc4a --- /dev/null +++ b/gusto/moving_mesh/__init__.py @@ -0,0 +1,3 @@ +from gusto.moving_mesh.monge_ampere import * # noqa +from gusto.moving_mesh.monitor import * # noqa +from gusto.moving_mesh.utility_functions import * # noqa diff --git a/gusto/moving_mesh/monge_ampere.py b/gusto/moving_mesh/monge_ampere.py new file mode 100644 index 000000000..1b05385f2 --- /dev/null +++ b/gusto/moving_mesh/monge_ampere.py @@ -0,0 +1,276 @@ +from firedrake import * +from gusto.logging import logger +import numpy as np + + +class OptimalTransportMeshGenerator(object): + + def __init__(self, mesh_in, monitor, initial_tol=1.e-4, tol=1.e-2, pre_meshgen_callback=None, post_meshgen_callback=None): + + self.mesh_in = mesh_in + self.monitor = monitor + self.initial_tol = initial_tol + self.tol = tol + self.pre_meshgen_fn = pre_meshgen_callback + self.post_meshgen_fn = post_meshgen_callback + + def setup(self): + mesh_in = self.mesh_in + cellname = mesh_in.ufl_cell().cellname() + quads = (cellname == "quadrilateral") + + if hasattr(mesh_in, '_radius'): + self.meshtype = "sphere" + self.R = mesh_in._radius + self.Rc = Constant(self.R) + dxdeg = dx(degree=8) # quadrature degree for nasty terms + else: + self.meshtype = "plane" + dxdeg = dx + + # Set up internal 'computational' mesh for own calculations + # This will be a unit sphere; make sure to scale appropriately + # when bringing coords in and out. + new_coords = Function(VectorFunctionSpace(mesh_in, "Q" if quads else "P", 2)) + new_coords.interpolate(SpatialCoordinate(mesh_in)) + new_coords.dat.data[:] /= np.linalg.norm(new_coords.dat.data, axis=1).reshape(-1, 1) + self.mesh = Mesh(new_coords) + + # Set up copy of passed-in mesh coordinate field for returning to + # external functions + self.output_coords = Function(mesh_in.coordinates) + self.own_output_coords = Function(VectorFunctionSpace(self.mesh, "Q" if quads else "P", mesh_in.coordinates.ufl_element().degree())) + + # get mesh area + self.total_area = assemble(Constant(1.0)*dx(self.mesh)) + + # set up functions + P2 = FunctionSpace(self.mesh, "Q" if quads else "P", 2) + TensorP2 = TensorFunctionSpace(self.mesh, "Q" if quads else "P", 2) + MixedSpace = P2*TensorP2 + P1 = FunctionSpace(self.mesh, "Q" if quads else "P", 1) # for representing m + W_cts = VectorFunctionSpace(self.mesh, "Q" if quads else "P", 1 if self.meshtype == "plane" else 2) # guaranteed-continuous version of coordinate field + + self.phisigma = Function(MixedSpace) + self.phi, self.sigma = split(self.phisigma) + self.phisigma_temp = Function(MixedSpace) + self.phi_temp, self.sigma_temp = split(self.phisigma_temp) + + # mesh coordinates + self.x = Function(self.mesh.coordinates) # for 'physical' coords + self.xi = Function(self.mesh.coordinates) # for 'computational' coords + + # monitor function mesh coords + self.x_old = Function(self.monitor.mesh.coordinates) + self.x_new = Function(self.monitor.mesh.coordinates) + + self.m = Function(P1) + self.theta = Constant(0.0) + + # define mesh equations + v, tau = TestFunctions(MixedSpace) + + if self.meshtype == "plane": + I = Identity(2) + F_mesh = inner(self.sigma, tau)*dx + dot(div(tau), grad(self.phi))*dx - (self.m*det(I + self.sigma) - self.theta)*v*dx + self.thetaform = self.m*det(I + self.sigma_temp)*dx + + elif self.meshtype == "sphere": + modgphi = sqrt(dot(grad(self.phi), grad(self.phi)) + 1e-12) + expxi = self.xi*cos(modgphi) + grad(self.phi)*sin(modgphi)/modgphi + projxi = Identity(3) - outer(self.xi, self.xi) + + modgphi_temp = sqrt(dot(grad(self.phi_temp), grad(self.phi_temp)) + 1e-12) + expxi_temp = self.xi*cos(modgphi_temp) + grad(self.phi_temp)*sin(modgphi_temp)/modgphi_temp + + F_mesh = inner(self.sigma, tau)*dxdeg + dot(div(tau), expxi)*dxdeg - (self.m*det(outer(expxi, self.xi) + dot(self.sigma, projxi)) - self.theta)*v*dxdeg + self.thetaform = self.m*det(outer(expxi_temp, self.xi) + dot(self.sigma_temp, projxi))*dxdeg + + # Define a solver for obtaining grad(phi) by L^2 projection + u_cts = TrialFunction(W_cts) + v_cts = TestFunction(W_cts) + + self.gradphi_cts = Function(W_cts) + + a_cts = dot(v_cts, u_cts)*dx + L_gradphi = dot(v_cts, grad(self.phi_temp))*dx + + probgradphi = LinearVariationalProblem(a_cts, L_gradphi, self.gradphi_cts) + self.solvgradphi = LinearVariationalSolver(probgradphi, solver_parameters={'ksp_type': 'cg'}) + + if self.meshtype == "plane": + self.gradphi_dg = Function(mesh.coordinates).assign(0) + elif self.meshtype == "sphere": + self.gradphi_cts2 = Function(W_cts) # extra, as gradphi_cts not necessarily tangential + + # set up initial sigma (needed on sphere) + sigma_ = TrialFunction(TensorP2) + tau_ = TestFunction(TensorP2) + sigma_ini = Function(TensorP2) + + asigmainit = inner(sigma_, tau_)*dxdeg + if self.meshtype == "plane": + Lsigmainit = -dot(div(tau_), grad(self.phi))*dx + else: + Lsigmainit = -dot(div(tau_), expxi)*dxdeg + + solve(asigmainit == Lsigmainit, sigma_ini, solver_parameters={'ksp_type': 'cg'}) + + self.phisigma.sub(1).assign(sigma_ini) + + # solver options for mesh generation + phi__, sigma__ = TrialFunctions(MixedSpace) + v__, tau__ = TestFunctions(MixedSpace) + + # Custom preconditioning matrix + Jp = inner(sigma__, tau__)*dx + phi__*v__*dx + dot(grad(phi__), grad(v__))*dx + + self.mesh_prob = NonlinearVariationalProblem(F_mesh, self.phisigma, Jp=Jp) + V1_nullspace = VectorSpaceBasis(constant=True) + self.nullspace = MixedVectorSpaceBasis(MixedSpace, [V1_nullspace, MixedSpace.sub(1)]) + + self.params = {"ksp_type": "gmres", + "pc_type": "fieldsplit", + "pc_fieldsplit_type": "multiplicative", + "pc_fieldsplit_off_diag_use_amat": True, + "fieldsplit_0_pc_type": "gamg", + "fieldsplit_0_ksp_type": "preonly", + # "fieldsplit_0_mg_levels_ksp_type": "chebyshev", + # "fieldsplit_0_mg_levels_ksp_chebyshev_estimate_eigenvalues": True, + # "fieldsplit_0_mg_levels_ksp_chebyshev_estimate_eigenvalues_random": True, + "fieldsplit_0_mg_levels_ksp_max_it": 5, + "fieldsplit_0_mg_levels_pc_type": "bjacobi", + "fieldsplit_0_mg_levels_sub_ksp_type": "preonly", + "fieldsplit_0_mg_levels_sub_pc_type": "ilu", + "fieldsplit_1_ksp_type": "preonly", + "fieldsplit_1_pc_type": "bjacobi", + "fieldsplit_1_sub_ksp_type": "preonly", + "fieldsplit_1_sub_pc_type": "ilu", + "ksp_max_it": 100, + "snes_max_it": 50, + "ksp_gmres_restart": 100, + "snes_rtol": self.initial_tol, + # "snes_atol": 1e-12, + "snes_linesearch_type": "l2", + "snes_linesearch_max_it": 5, + "snes_linesearch_maxstep": 1.05, + "snes_linesearch_damping": 0.8, + # "ksp_monitor": None, + "snes_monitor": None, + # "snes_linesearch_monitor": None, + "snes_lag_preconditioner": -2} + + self.mesh_solv = NonlinearVariationalSolver( + self.mesh_prob, + nullspace=self.nullspace, + transpose_nullspace=self.nullspace, + pre_jacobian_callback=self.update_mxtheta, + pre_function_callback=self.update_mxtheta, + solver_parameters=self.params) + + self.mesh_solv.snes.setMonitor(self.fakemonitor) + + def update_mxtheta(self, cursol): + with self.phisigma_temp.dat.vec as v: + cursol.copy(v) + + # Obtain continous version of grad phi. + self.mesh.coordinates.assign(self.xi) + self.solvgradphi.solve() + + # "Fix grad(phi) on sphere" + if self.meshtype == "sphere": + # Ensures that grad(phi).x = 0, assuming |x| = 1 + # v_new = v - (v.x)x + self.gradphi_cts2.interpolate(self.gradphi_cts - dot(self.gradphi_cts, self.mesh.coordinates)*self.mesh.coordinates) + + # Generate coordinates + if self.meshtype == "plane": + self.gradphi_dg.interpolate(self.gradphi_cts) + self.x.assign(self.xi + self.gradphi_dg) # x = xi + grad(phi) + else: + # Generate new coordinate field using exponential map + # x = cos(|v|)*x + sin(|v|)*(v/|v|) + gradphinorm = sqrt(dot(self.gradphi_cts2, self.gradphi_cts2)) + 1e-12 + self.x.interpolate(cos(gradphinorm)*self.xi + sin(gradphinorm)*(self.gradphi_cts2/gradphinorm)) + + if self.initial_mesh: + # self.mesh.coordinates.assign(self.x) + self.own_output_coords.interpolate(Constant(self.R)*self.x) + self.output_coords.dat.data[:] = self.own_output_coords.dat.data_ro[:] + self.mesh_in.coordinates.assign(self.output_coords) + self.initialise_fn() + self.monitor.mesh.coordinates.dat.data[:] = self.own_output_coords.dat.data_ro[:] + self.monitor.update_monitor() + self.m.dat.data[:] = self.monitor.m.dat.data_ro[:] + else: + self.own_output_coords.interpolate(Constant(self.R)*self.x) + self.x_new.dat.data[:] = self.own_output_coords.dat.data_ro[:] + self.monitor.get_monitor_on_new_mesh(self.monitor.m, self.x_old, self.x_new) + self.m.dat.data[:] = self.monitor.m.dat.data_ro[:] + + self.mesh.coordinates.assign(self.xi) + theta_new = assemble(self.thetaform)/self.total_area + self.theta.assign(theta_new) + + def fakemonitor(self, snes, it, rnorm): + cursol = snes.getSolution() + self.update_mxtheta(cursol) # updates m, x, and theta + + def get_first_mesh(self, initialise_fn): + + """ + This function is used to generate a mesh adapted to the initial state. + :arg initialise_fn: a user-specified Python function that sets the + initial condition + """ + self.initial_mesh = True + self.initialise_fn = initialise_fn + self.mesh_solv.solve() + + # remake mesh solver with new tolerance + self.params["snes_rtol"] = self.tol + # self.params["snes_linesearch_type"] = "bt" + self.params["snes_max_it"] = 15 + + self.mesh_solv = NonlinearVariationalSolver(self.mesh_prob, + nullspace=self.nullspace, + transpose_nullspace=self.nullspace, + pre_jacobian_callback=self.update_mxtheta, + pre_function_callback=self.update_mxtheta, + solver_parameters=self.params) + + self.mesh_solv.snes.setMonitor(self.fakemonitor) + self.initial_mesh = False + + def get_new_mesh(self): + # Back up the current mesh + self.x_old.dat.data[:] = self.mesh_in.coordinates.dat.data_ro[:] + + # Make monitor function + # TODO: should I just pass in the 'coords to use' to update_monitor? + self.monitor.mesh.coordinates.dat.data[:] = self.mesh_in.coordinates.dat.data_ro[:] + self.monitor.update_monitor() + + # Back this up + self.monitor.m_old.assign(self.monitor.m) + + # Generate new mesh, coords put in self.x + try: + self.mesh_solv.solve() + except ConvergenceError: + logger.warning( + 'mesh solver did not converge - oh well, continuing anyway!') + + # Move data from internal mesh to output mesh. + self.own_output_coords.interpolate(Constant(self.R)*self.x) + self.output_coords.dat.data[:] = self.own_output_coords.dat.data_ro[:] + return self.output_coords + + def pre_meshgen_callback(self): + if self.pre_meshgen_fn: + self.pre_meshgen_fn() + + def post_meshgen_callback(self): + if self.post_meshgen_fn: + self.post_meshgen_fn() diff --git a/gusto/moving_mesh/monitor.py b/gusto/moving_mesh/monitor.py new file mode 100644 index 000000000..856c0a4f9 --- /dev/null +++ b/gusto/moving_mesh/monitor.py @@ -0,0 +1,153 @@ +from firedrake import * +import numpy as np +from mpi4py import MPI + +__all__ = ["MonitorFunction"] + + +class MonitorFunction(object): + + def __init__(self, f_name, adapt_to="function", avg_weight=0.5, + max_min_cap=4.0): + + assert adapt_to in ("function", "gradient", "hessian") + assert 0.0 <= avg_weight < 1.0 + assert max_min_cap >= 1.0 or max_min_cap == 0.0 + + self.f_name = f_name + self.adapt_to = adapt_to + self.avg_weight = avg_weight + self.max_min_cap = max_min_cap + + def setup(self, state_fields): + f = state_fields(self.f_name) + cellname = f.ufl_element().cell().cellname() + quads = (cellname == "quadrilateral") + + # Set up internal mesh for monitor function calculations + new_coords = Function(f.function_space().mesh().coordinates) + self.mesh = Mesh(new_coords) + self.f = Function(FunctionSpace(self.mesh, f.ufl_element())) # make "own copy" of f on internal mesh + self.user_f = f + + # Set up function spaces + P1 = FunctionSpace(self.mesh, "Q" if quads else "P", 1) # for representing m + DP1 = FunctionSpace(self.mesh, "DQ" if quads else "DP", 1) # for advection + VectorP1 = VectorFunctionSpace(self.mesh, "Q" if quads else "P", 1) + self.limiter = VertexBasedLimiter(DP1) + + self.gradq = Function(VectorP1) + if self.adapt_to == "hessian": + TensorP1 = TensorFunctionSpace(self.mesh, "Q" if quads else "P", 1) + self.hessq = Function(TensorP1) + + # get mesh area + self.total_area = assemble(Constant(1.0)*dx(self.mesh)) + + self.m = Function(P1) + self.m_prereg = Function(P1) + self.m_old = Function(P1) + self.m_dg = Function(DP1) + self.dm = Function(DP1) + self.m_int_form = self.m_prereg*dx + + # define monitor function in terms of q + if False: # for plane + v_ones = as_vector(np.ones(2)) + else: + v_ones = as_vector(np.ones(3)) + # Obtain weak gradient + # u_vp1 = TrialFunction(VectorP1) + v_vp1 = TestFunction(VectorP1) + self.a_vp1_lumped = dot(v_vp1, v_ones)*dx + self.L_vp1 = -div(v_vp1)*self.f*dx + + if self.adapt_to == "hessian": + if False: # for plane + t_ones = as_vector(np.ones((2, 2))) + else: + t_ones = as_vector(np.ones((3, 3))) + # Obtain approximation to hessian + # u_tp1 = TrialFunction(TensorP1) + v_tp1 = TestFunction(TensorP1) + self.a_tp1_lumped = inner(v_tp1, t_ones)*dx + self.L_tp1 = inner(v_tp1, grad(self.gradq))*dx + + # Define forms for lumped project of monitor function into P1 + v_p1 = TestFunction(P1) + self.a_p1_lumped = v_p1*dx + if True: # not uniform: + if self.adapt_to == "gradient": + self.L_monitor = v_p1*sqrt(dot(self.gradq, self.gradq))*dx + elif self.adapt_to == "hessian": + self.L_monitor = v_p1*sqrt(inner(self.hessq, self.hessq))*dx + else: + self.L_monitor = v_p1*dx + + u_dg = TrialFunction(DP1) + v_dg = TestFunction(DP1) + n = FacetNormal(self.mesh) + self.mesh_adv_vel = Function(self.mesh.coordinates) + vn = 0.5*(dot(self.mesh_adv_vel, n) + abs(dot(self.mesh_adv_vel, n))) + a_dg = v_dg*u_dg*dx + L_madv = self.m_dg*div(v_dg*self.mesh_adv_vel)*dx - jump(v_dg)*jump(vn*self.m_dg)*dS + + prob_madv = LinearVariationalProblem(a_dg, L_madv, self.dm, constant_jacobian=False) + self.solv_madv = LinearVariationalSolver(prob_madv, + solver_parameters={"ksp_type": "preonly", + "pc_type": "bjacobi", + "sub_ksp_type": "preonly", + "sub_pc_type": "ilu"}) + + def update_monitor(self): + + self.f.dat.data[:] = self.user_f.dat.data[:] + + # TODO don't need to do this if adapting to function + assemble(self.L_vp1, tensor=self.gradq) + self.gradq.dat /= assemble(self.a_vp1_lumped).dat # obtain weak gradient + if self.adapt_to == "hessian": + assemble(self.L_tp1, tensor=self.hessq) + self.hessq.dat /= assemble(self.a_tp1_lumped).dat + + # obtain m by lumped projection + self.m_prereg.interpolate(assemble(self.L_monitor)/assemble(self.a_p1_lumped)) + + # calculate average of m + m_int = assemble(self.m_int_form) + m_avg = m_int/self.total_area + + # cap max-to-avg ratio + self.m_prereg.dat.data[:] = np.fmin(self.m_prereg.dat.data[:], (self.max_min_cap - 1.0)*m_avg) + + # use Beckett+Mackenzie regularization + self.m.assign(Constant(self.avg_weight)*self.m_prereg + Constant(1.0 - self.avg_weight)*Constant(m_avg)) + + assert (self.m.dat.data >= 0.0).all() + + # make m O(1) + m_min = self.m.comm.allreduce(self.m.dat.data_ro.min(), op=MPI.MIN) + self.m.dat.data[:] /= m_min + + # mmax_pre = max(self.m.dat.data)/min(self.m.dat.data) + + def get_monitor_on_new_mesh(self, m, x_old, x_new): + # We have the function m on old mesh: m_old. We want to represent + # this on the trial mesh. Do this by using the same values + # (equivalent to advection by +v), then do an advection step of -v. + + self.mesh.coordinates.assign(x_new) + self.limiter.centroid_solver = self.limiter._construct_centroid_solver() + self.mesh_adv_vel.assign(x_old - x_new) + + # Make discontinuous m + self.m_dg.interpolate(self.m_old) + + # Advect this by -v + for ii in range(10): + self.solv_madv.solve() + self.m_dg.assign(self.m_dg + Constant(1.0/10.)*self.dm) + self.limiter.apply(self.m_dg) + + project(self.m_dg, self.m) # project discontinuous m back into CG + # mmax_post = max(self.m.dat.data)/min(self.m.dat.data) diff --git a/gusto/moving_mesh/utility_functions.py b/gusto/moving_mesh/utility_functions.py new file mode 100644 index 000000000..39f327f91 --- /dev/null +++ b/gusto/moving_mesh/utility_functions.py @@ -0,0 +1,21 @@ +import numpy as np + + +__all__ = ["spherical_logarithm"] + + +def spherical_logarithm(X0, X1, v, R): + """ + Find vector function v such that X1 = exp(v)X0 on + a sphere of radius R, centre the origin. + """ + + v.assign(X1 - X0) + + # v <- v - X0(v.X0/R^2); make v orthogonal to X0 + v.dat.data[:] -= X0.dat.data_ro*np.einsum('ij,ij->i', v.dat.data_ro, X0.dat.data_ro).reshape(-1, 1)/R**2 + + # v <- theta*R*v-hat, where theta is the angle between X0 and X1 + # fmin(theta, 1.0) is used to avoid silly floating point errors + # fmax(|v|, 1e-16*R) is used to avoid division by zero + v.dat.data[:] = np.arccos(np.fmin(np.einsum('ij,ij->i', X0.dat.data_ro, X1.dat.data_ro)/R**2, 1.0)).reshape(-1, 1)*R*v.dat.data_ro[:]/np.fmax(np.linalg.norm(v.dat.data_ro, axis=1), R*1e-16).reshape(-1, 1) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 98e8d53f5..b01209bd3 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -1365,7 +1365,6 @@ def apply(self, x_out, *x_in): """ if self.initial_timesteps < self.nlevels-1: self.initial_timesteps += 1 - print(self.initial_timesteps) solver = self.solver0 else: solver = self.solver diff --git a/gusto/timeloop.py b/gusto/timeloop.py index 3fba00429..d8d303833 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -1,24 +1,27 @@ """Classes for controlling the timestepping loop.""" from abc import ABCMeta, abstractmethod, abstractproperty -from firedrake import Function, Projector, Constant, split +from firedrake import Function, Projector, Constant, TrialFunction, \ + TestFunction, assemble, inner, dx, split +from firedrake.petsc import PETSc from pyop2.profiling import timed_stage -from gusto.equations import PrognosticEquationSet -from gusto.fml import drop, Label, Term +from gusto.equations import PrognosticEquationSet, ShallowWaterEquations from gusto.fields import TimeLevelFields, StateFields from gusto.forcing import Forcing +from gusto.fml import drop, Label, Term from gusto.labels import ( transport, diffusion, time_derivative, linearisation, prognostic, physics, transporting_velocity ) from gusto.linear_solvers import LinearTimesteppingSolver from gusto.logging import logger +from gusto.moving_mesh.utility_functions import spherical_logarithm from gusto.time_discretisation import ExplicitTimeDiscretisation from gusto.transport_methods import TransportMethod import ufl __all__ = ["Timestepper", "SplitPhysicsTimestepper", "SemiImplicitQuasiNewton", - "PrescribedTransport"] + "PrescribedTransport", "MeshMovement"] class BaseTimestepper(object, metaclass=ABCMeta): @@ -668,3 +671,209 @@ def timestep(self): self.velocity_projection.project() super().timestep() + + +class MeshMovement(SemiImplicitQuasiNewton): + + def __init__(self, equation_set, io, transport_schemes, + spatial_methods, + auxiliary_equations_and_schemes=None, + linear_solver=None, + diffusion_schemes=None, + physics_schemes=None, + mesh_generator=None, + **kwargs): + + Vu = equation_set.domain.spaces("HDiv") + self.uadv = Function(Vu) + + assert isinstance(equation_set, ShallowWaterEquations), \ + 'Mesh movement only works with the shallow water equation set' + + super().__init__(equation_set=equation_set, + io=io, + transport_schemes=transport_schemes, + spatial_methods=spatial_methods, + auxiliary_equations_and_schemes=auxiliary_equations_and_schemes, + linear_solver=linear_solver, + diffusion_schemes=diffusion_schemes, + physics_schemes=physics_schemes, + **kwargs) + + self.mesh_generator = mesh_generator + mesh = self.equation.domain.mesh + self.mesh = mesh + self.X0 = Function(mesh.coordinates) + self.X1 = Function(mesh.coordinates) + from firedrake import File, FunctionSpace + Vf = FunctionSpace(mesh, "DG", 1) + f = Function(Vf) + coords_out = File("out.pvd") + mesh.coordinates.assign(self.X0) + coords_out.write(f) + mesh.coordinates.assign(self.X1) + coords_out.write(f) + + self.on_sphere = self.equation.domain.on_sphere + + self.v = Function(mesh.coordinates.function_space()) + self.v_V1 = Function(Vu) + self.v1 = Function(mesh.coordinates.function_space()) + self.v1_V1 = Function(Vu) + + # Set up diagnostics - also called in the run method, is it ok + # to do this twice (I suspect not) + self.io.setup_diagnostics(self.fields) + + mesh_generator.monitor.setup(self.fields) + mesh_generator.setup() + + self.tests = {} + for name in self.equation.field_names: + self.tests[name] = TestFunction(self.x.n(name).function_space()) + self.trials = {} + for name in self.equation.field_names: + self.trials[name] = TrialFunction(self.x.n(name).function_space()) + self.ksp = {} + for name in self.equation.field_names: + self.ksp[name] = PETSc.KSP().create() + self.Lvec = {} + + def setup_fields(self): + super().setup_fields() + self.x.add_fields(self.equation, levels=("mid",)) + + @property + def transporting_velocity(self): + return self.uadv + + def timestep(self): + """Defines the timestep""" + xn = self.x.n + xnp1 = self.x.np1 + xstar = self.x.star + xmid = self.x.mid + xp = self.x.p + xrhs = self.xrhs + dy = self.dy + un = xn("u") + unp1 = xnp1("u") + X0 = self.X0 + X1 = self.X1 + + # both X0 and X1 hold the mesh coordinates at the start of the timestep + X1.assign(self.mesh.coordinates) + X0.assign(X1) + + # this recomputes the fields that the monitor function depends on + self.mesh_generator.pre_meshgen_callback() + + # solve for new mesh and store this in X1 + with timed_stage("Mesh generation"): + X1.assign(self.mesh_generator.get_new_mesh()) + + # still on the old mesh, compute explicit forcing from xn and + # store in xstar + with timed_stage("Apply forcing terms"): + self.forcing.apply(xn, xn, xstar(self.field_name), "explicit") + + # set xp to xstar + xp(self.field_name).assign(xstar(self.field_name)) + + # Compute v (mesh velocity w.r.t. initial mesh) and + # v1 (mesh velocity w.r.t. final mesh) + # TODO: use Projectors below! + if self.on_sphere: + spherical_logarithm(X0, X1, self.v, self.mesh._radius) + self.v /= self.dt + spherical_logarithm(X1, X0, self.v1, self.mesh._radius) + self.v1 /= -self.dt + + self.mesh.coordinates.assign(X0) + self.v_V1.project(self.v) + + self.mesh.coordinates.assign(X1) + self.v1_V1.project(self.v1) + + else: + self.mesh.coordinates.assign(X0) + self.v_V1.project((X1 - X0)/self.dt) + + self.mesh.coordinates.assign(X1) + self.v1_V1.project(-(X0 - X1)/self.dt) + + with timed_stage("Transport"): + for name, scheme in self.active_transport: + # transport field from xstar to xmid on old mesh + self.mesh.coordinates.assign(X0) + self.uadv.assign(0.5*(un - self.v_V1)) + scheme.apply(xmid(name), xstar(name)) + + # assemble the rhs of the projection operator on the + # old mesh + # TODO this should only be for fields that are in the HDiv + # space or are transported using the conservation form + rhs = inner(xmid(name), self.tests[name])*dx + with assemble(rhs).dat.vec as v: + self.Lvec[name] = v + + # put mesh_new into mesh + self.mesh.coordinates.assign(X1) + + # this should reinterpolate any necessary fields + # (e.g. Coriolis and cell normals) + self.mesh_generator.post_meshgen_callback() + + # now create the matrix for the projection and solve + # TODO this should only be for fields that are in the HDiv + # space or are transported using the conservation form + for name, scheme in self.active_transport: + lhs = inner(self.trials[name], self.tests[name])*dx + amat = assemble(lhs) + a = amat.M.handle + self.ksp[name].setOperators(a) + self.ksp[name].setFromOptions() + + with xmid(name).dat.vec as x_: + self.ksp[name].solve(self.Lvec[name], x_) + + for k in range(self.maxk): + + for name, scheme in self.active_transport: + # transport field from xmid to xp on new mesh + self.uadv.assign(0.5*(unp1 - self.v1_V1)) + scheme.apply(xp(name), xmid(name)) + + # xrhs is the residual which goes in the linear solve + xrhs.assign(0.) + + for i in range(self.maxi): + + with timed_stage("Apply forcing terms"): + self.forcing.apply(xp, xnp1, xrhs, "implicit") + + xrhs -= xnp1(self.field_name) + + with timed_stage("Implicit solve"): + # solves linear system and places result in dy + self.linear_solver.solve(xrhs, dy) + + xnp1X = xnp1(self.field_name) + xnp1X += dy + + # Update xnp1 values for active tracers not included in the linear solve + self.copy_active_tracers(xp, xnp1) + + self._apply_bcs() + + for name, scheme in self.auxiliary_schemes: + # transports a field from xn and puts result in xnp1 + scheme.apply(xnp1(name), xn(name)) + + with timed_stage("Diffusion"): + for name, scheme in self.diffusion_schemes: + scheme.apply(xnp1(name), xnp1(name)) + + with timed_stage("Physics"): + for _, scheme in self.physics_schemes: + scheme.apply(xnp1(scheme.field_name), xnp1(scheme.field_name)) diff --git a/unit-tests/fml_tests/test_replace_perp.py b/unit-tests/fml_tests/test_replace_perp.py index f095df545..14c0e5a6d 100644 --- a/unit-tests/fml_tests/test_replace_perp.py +++ b/unit-tests/fml_tests/test_replace_perp.py @@ -1,5 +1,5 @@ # The perp routine should come from UFL when it is fully implemented there -from gusto import perp +from gusto.perp import perp from gusto.fml import subject, replace_subject, all_terms from firedrake import (UnitSquareMesh, FunctionSpace, MixedFunctionSpace, TestFunctions, Function, split, inner, dx, errornorm,