From 0789602c4f0402623e7bdb8489be7a9862574cfb Mon Sep 17 00:00:00 2001 From: jshipton Date: Thu, 16 Feb 2023 15:23:35 +0000 Subject: [PATCH 01/33] first attempt at mesh movement timeloop --- gusto/timeloop.py | 93 ++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 91 insertions(+), 2 deletions(-) diff --git a/gusto/timeloop.py b/gusto/timeloop.py index 020008d39..a38ee0fe7 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -1,7 +1,7 @@ """Classes for controlling the timestepping loop.""" from abc import ABCMeta, abstractmethod, abstractproperty -from firedrake import Function, Projector, Constant +from firedrake import Function, Projector, Constant, Mesh from pyop2.profiling import timed_stage from gusto.configuration import logger from gusto.equations import PrognosticEquationSet @@ -14,7 +14,7 @@ from gusto.time_discretisation import ExplicitTimeDiscretisation __all__ = ["Timestepper", "SplitPhysicsTimestepper", "SemiImplicitQuasiNewton", - "PrescribedTransport"] + "PrescribedTransport", "MeshMovement"] class BaseTimestepper(object, metaclass=ABCMeta): @@ -519,3 +519,92 @@ def timestep(self): with timed_stage("Physics"): for _, scheme in self.physics_schemes: scheme.apply(self.x.np1(scheme.field_name), self.x.np1(scheme.field_name)) + + +class MeshMovement(SemiImplicitQuasiNewton): + + def __init__(self, equation_set, io, transport_schemes, + auxiliary_equations_and_schemes=None, + linear_solver=None, + diffusion_schemes=None, + physics_schemes=None, + mesh_generator=None, + **kwargs): + + super().__init__(equation_set=equation_set, + io=io, + transport_schemes=transport_schemes, + 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.X0 = Function(mesh.coordinates) + self.X1 = Function(mesh.coordinates) + xold = Function(mesh.coordinates) + self.mesh_old = Mesh(xold) + + def timestep(self): + """Defines the timestep""" + xn = self.x.n + xnp1 = self.x.np1 + xstar = self.x.star + xp = self.x.p + xrhs = self.xrhs + dy = self.dy + + with timed_stage("Apply forcing terms"): + self.forcing.apply(xn, xn, xstar(self.field_name), "explicit") + + xp(self.field_name).assign(xstar(self.field_name)) + + for k in range(self.maxk): + + # This is used as the "old mesh" domain in projections + self.mesh_old.coordinates.dat.data[:] = self.X1.dat.data[:] + self.X0.assign(self.X1) + + self.mesh_generator.pre_meshgen_callback() + with timed_stage("Mesh generation"): + self.X1.assign(self.mesh_generator.get_new_mesh()) + + with timed_stage("Transport"): + for name, scheme in self.active_transport: + # transports a field from xstar and puts result in xp + scheme.apply(xp(name), xstar(name)) + + self.mesh_generator.post_meshgen_callback() + + xrhs.assign(0.) # xrhs is the residual which goes in the linear solve + + 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"): + self.linear_solver.solve(xrhs, dy) # solves linear system and places result in 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)) From a9f390a2f7692cf51abd3f1f9d603f6ccfd71405 Mon Sep 17 00:00:00 2001 From: jshipton Date: Thu, 16 Feb 2023 17:31:06 +0000 Subject: [PATCH 02/33] progress: moving mesh Galewsky script runs and produces an adapted initial mesh --- .../shallow_water/mm_ot_sw_galewsky_jet.py | 184 ++++++++++++++++++ gusto/__init__.py | 1 + gusto/moving_mesh/__init__.py | 2 + gusto/moving_mesh/monitor.py | 153 +++++++++++++++ gusto/timeloop.py | 5 +- 5 files changed, 343 insertions(+), 2 deletions(-) create mode 100644 examples/shallow_water/mm_ot_sw_galewsky_jet.py create mode 100644 gusto/moving_mesh/__init__.py create mode 100644 gusto/moving_mesh/monitor.py 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..d038aa71b --- /dev/null +++ b/examples/shallow_water/mm_ot_sw_galewsky_jet.py @@ -0,0 +1,184 @@ +from gusto import * +from firedrake import IcosahedralSphereMesh, Constant, ge, le, exp, cos, \ + conditional, interpolate, SpatialCoordinate, VectorFunctionSpace, \ + Function, assemble, dx, FunctionSpace,pi + +import numpy as np + +day = 24.*60.*60. +dt = 480. +tmax = 6*day + +# setup shallow water parameters +R = 6371220. +H = 10000. +parameters = ShallowWaterParameters(H=H) + +# Domain +mesh = IcosahedralSphereMesh(radius=R, + refinement_level=3, degree=1) +x = SpatialCoordinate(mesh) +mesh.init_cell_orientations(x) +domain = Domain(mesh, dt, 'BDM', 1) + +# Equation +Omega = parameters.Omega +fexpr = 2*Omega*x[2]/R +eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr) + +# I/O +perturb = True +if perturb: + dirname = "mm_ot_sw_galewsky_jet_perturbed" +else: + dirname = "mm_ot_sw_galewsky_jet_unperturbed" + +output = OutputParameters(dirname=dirname, + dumplist_latlon=['D', 'PotentialVorticity', + 'RelativeVorticity'], + log_level="INFO") +pv = PotentialVorticity() +diagnostic_fields = [pv, RelativeVorticity()] +io = IO(domain, output, diagnostic_fields=diagnostic_fields) + +# Transport schemes +transported_fields = [ImplicitMidpoint(domain, "u"), + SSPRK3(domain, "D")] + + +# Mesh movement + +def update_pv(): + pv() + +def reinterpolate_coriolis(): + eqns.prescribed_fields("coriolis").interpolate(fexpr) + +monitor = MonitorFunction("PotentialVorticity", adapt_to="gradient") +mesh_generator = OptimalTransportMeshGenerator(mesh, + monitor, + pre_meshgen_callback=update_pv, + post_meshgen_callback=reinterpolate_coriolis) + +# Time stepper +stepper = MeshMovement(eqns, io, transported_fields, + 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(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: + D_pert.interpolate(Dhat*cos(theta)*exp(-(lamda/alpha)**2)*exp(-((theta2 - theta)/beta)**2)) + D0 += D_pert + eqns.prescribed_fields("coriolis").interpolate(fexpr) + pv() + +pv.setup(domain, stepper.fields) +mesh_generator.get_first_mesh(initialise_fn) + +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 dabfcdf2a..28133a71a 100644 --- a/gusto/__init__.py +++ b/gusto/__init__.py @@ -11,6 +11,7 @@ from gusto.labels import * # noqa from gusto.limiters import * # noqa from gusto.linear_solvers import * # noqa +from gusto.moving_mesh import * # noqa from gusto.physics import * # noqa from gusto.preconditioners import * # noqa from gusto.recovery import * # noqa diff --git a/gusto/moving_mesh/__init__.py b/gusto/moving_mesh/__init__.py new file mode 100644 index 000000000..433e9c492 --- /dev/null +++ b/gusto/moving_mesh/__init__.py @@ -0,0 +1,2 @@ +from gusto.moving_mesh.monge_ampere import * # noqa +from gusto.moving_mesh.monitor import * # noqa diff --git a/gusto/moving_mesh/monitor.py b/gusto/moving_mesh/monitor.py new file mode 100644 index 000000000..1bc7064a8 --- /dev/null +++ b/gusto/moving_mesh/monitor.py @@ -0,0 +1,153 @@ +from firedrake import * +import math +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 + else: + 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/timeloop.py b/gusto/timeloop.py index a38ee0fe7..f16ec52d9 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -37,6 +37,7 @@ def __init__(self, equation, io): self.setup_scheme() self.io.log_parameters(equation) + self.io.setup_diagnostics(self.fields) @abstractproperty def transporting_velocity(self): @@ -70,8 +71,6 @@ def run(self, t, tmax, pickup=False): if pickup: t = self.io.pickup_from_checkpoint(self.fields) - self.io.setup_diagnostics(self.fields) - with timed_stage("Dump output"): self.io.setup_dump(self.fields, t, tmax, pickup) @@ -545,6 +544,8 @@ def __init__(self, equation_set, io, transport_schemes, self.X1 = Function(mesh.coordinates) xold = Function(mesh.coordinates) self.mesh_old = Mesh(xold) + mesh_generator.monitor.setup(self.fields) + mesh_generator.setup() def timestep(self): """Defines the timestep""" From f215579f50c45c20db0d5d1e2385523d1c5a3523 Mon Sep 17 00:00:00 2001 From: jshipton Date: Thu, 16 Feb 2023 22:29:48 +0000 Subject: [PATCH 03/33] Progress: runs but is blowing up. Much tidying needed --- gusto/equations.py | 1 + gusto/moving_mesh/__init__.py | 5 +- gusto/timeloop.py | 119 ++++++++++++++++++++++++++++++++-- 3 files changed, 116 insertions(+), 9 deletions(-) diff --git a/gusto/equations.py b/gusto/equations.py index aa66f10ac..b4581fdef 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -246,6 +246,7 @@ def __init__(self, field_names, domain, linearisation_map=None, # 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) diff --git a/gusto/moving_mesh/__init__.py b/gusto/moving_mesh/__init__.py index 433e9c492..85449fc4a 100644 --- a/gusto/moving_mesh/__init__.py +++ b/gusto/moving_mesh/__init__.py @@ -1,2 +1,3 @@ -from gusto.moving_mesh.monge_ampere import * # noqa -from gusto.moving_mesh.monitor import * # noqa +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/timeloop.py b/gusto/timeloop.py index f16ec52d9..c69b98aa0 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -1,7 +1,9 @@ """Classes for controlling the timestepping loop.""" from abc import ABCMeta, abstractmethod, abstractproperty -from firedrake import Function, Projector, Constant, Mesh +from firedrake import Function, Projector, Constant, Mesh, TrialFunction, \ + TestFunction, assemble, inner, dx +from firedrake.petsc import PETSc from pyop2.profiling import timed_stage from gusto.configuration import logger from gusto.equations import PrognosticEquationSet @@ -10,6 +12,7 @@ from gusto.labels import (transport, diffusion, time_derivative, linearisation, prognostic, physics) from gusto.linear_solvers import LinearTimesteppingSolver +from gusto.moving_mesh.utility_functions import spherical_logarithm from gusto.fields import TimeLevelFields, StateFields from gusto.time_discretisation import ExplicitTimeDiscretisation @@ -530,6 +533,9 @@ def __init__(self, equation_set, io, transport_schemes, mesh_generator=None, **kwargs): + Vu = equation_set.domain.spaces("HDiv") + self.uadv = Function(Vu) + super().__init__(equation_set=equation_set, io=io, transport_schemes=transport_schemes, @@ -538,23 +544,53 @@ def __init__(self, equation_set, io, transport_schemes, 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) + xold = Function(mesh.coordinates) self.mesh_old = Mesh(xold) + + 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) + mesh_generator.monitor.setup(self.fields) mesh_generator.setup() + 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 + um, Dm = Function(self.equation.W).split() + um_, Dm_ = Function(self.equation.W).split() + test_u = TestFunction(um.function_space()) + trial_u = TrialFunction(um.function_space()) + test_D = TestFunction(Dm.function_space()) + trial_D = TrialFunction(Dm.function_space()) with timed_stage("Apply forcing terms"): self.forcing.apply(xn, xn, xstar(self.field_name), "explicit") @@ -563,18 +599,87 @@ def timestep(self): for k in range(self.maxk): - # This is used as the "old mesh" domain in projections - self.mesh_old.coordinates.dat.data[:] = self.X1.dat.data[:] - self.X0.assign(self.X1) + X0.assign(X1) self.mesh_generator.pre_meshgen_callback() with timed_stage("Mesh generation"): - self.X1.assign(self.mesh_generator.get_new_mesh()) + X1.assign(self.mesh_generator.get_new_mesh()) + + # 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: - # transports a field from xstar and puts result in xp - scheme.apply(xp(name), xstar(name)) + # 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)) + print("before: ", name, xmid(name).dat.data.min(), xmid(name).dat.data.max()) + + if name == "u": + print("projecting u") + um_.assign(xmid("u")) + rhs = inner(um_, test_u)*dx + with assemble(rhs).dat.vec as v: + Lvec = v + + elif name == "D": + print("projecting D") + Dm_.assign(xmid("D")) + rhs = inner(Dm_, test_D)*dx + with assemble(rhs).dat.vec as v: + Lvec = v + + # put mesh_new into mesh + self.mesh.coordinates.assign(X1) + + if name == "u": + lhs = inner(trial_u, test_u)*dx + amat = assemble(lhs) + a = amat.M.handle + ksp = PETSc.KSP().create() + ksp.setOperators(a) + ksp.setFromOptions() + + with um.dat.vec as x_: + ksp.solve(Lvec, x_) + xmid(name).assign(um) + + elif name == "D": + lhs = inner(trial_D, test_D)*dx + amat = assemble(lhs) + a = amat.M.handle + ksp = PETSc.KSP().create() + ksp.setOperators(a) + ksp.setFromOptions() + + with Dm.dat.vec as x_: + ksp.solve(Lvec, x_) + xmid(name).assign(Dm) + print("after: ", name, xmid(name).dat.data.min(), xmid(name).dat.data.max()) + + # transport field from xmid to xp on new mesh + self.uadv.assign(0.5*(unp1 - self.v1_V1)) + scheme.apply(xp(name), xmid(name)) self.mesh_generator.post_meshgen_callback() From e314d47c3afb091092a21774c51d8392205a32b2 Mon Sep 17 00:00:00 2001 From: jshipton Date: Fri, 24 Feb 2023 12:06:57 +0000 Subject: [PATCH 04/33] oops, forgot these files --- gusto/moving_mesh/monge_ampere.py | 278 +++++++++++++++++++++++++ gusto/moving_mesh/utility_functions.py | 21 ++ 2 files changed, 299 insertions(+) create mode 100644 gusto/moving_mesh/monge_ampere.py create mode 100644 gusto/moving_mesh/utility_functions.py diff --git a/gusto/moving_mesh/monge_ampere.py b/gusto/moving_mesh/monge_ampere.py new file mode 100644 index 000000000..4c1c8b804 --- /dev/null +++ b/gusto/moving_mesh/monge_ampere.py @@ -0,0 +1,278 @@ +from firedrake import * +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.mesh_in.coordinates.assign(self.x) + self.initialise_fn() + #self.monitor.mesh.coordinates.dat.data[:] = self.mesh.coordinates.dat.data_ro[:] + 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.mesh.coordinates.assign(self.x) + 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: + print("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.mesh.coordinates.assign(self.x) + # self.output_coords.dat.data[:] = self.mesh.coordinates.dat.data_ro[:] + 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/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) From f8e282d0f5359967af8bbd34ce96f4d227c88bb1 Mon Sep 17 00:00:00 2001 From: jshipton Date: Thu, 9 Mar 2023 10:13:21 +0000 Subject: [PATCH 05/33] hack circulation form a bit for now --- gusto/equations.py | 2 +- gusto/transport_forms.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/gusto/equations.py b/gusto/equations.py index b4581fdef..37e0e2193 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -635,7 +635,7 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, lambda t: Term(ufl.replace( t.form, {t.get(transporting_velocity): u}), t.labels)) ke_form = transporting_velocity.remove(ke_form) - u_adv = prognostic(advection_equation_circulation_form(domain, w, u), "u") + ke_form + u_adv = prognostic(advection_equation_circulation_form(domain, w, u), "u") + 0.5*ke_form else: raise ValueError("Invalid u_transport_option: %s" % u_transport_option) diff --git a/gusto/transport_forms.py b/gusto/transport_forms.py index ae9d14d7b..f9ce9801b 100644 --- a/gusto/transport_forms.py +++ b/gusto/transport_forms.py @@ -368,7 +368,7 @@ def kinetic_energy_form(domain, test, q): """ ubar = Function(domain.spaces("HDiv")) - L = 0.5*div(test)*inner(q, ubar)*dx + L = div(test)*inner(q, ubar)*dx form = transporting_velocity(L, ubar) From 756be5eacdbac0b21df1165734b10b16b1ba1928 Mon Sep 17 00:00:00 2001 From: jshipton Date: Thu, 9 Mar 2023 10:53:43 +0000 Subject: [PATCH 06/33] oops, I forgot that the vector_invariant_form already included a grad(usq) term --- gusto/equations.py | 2 +- gusto/transport_forms.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/gusto/equations.py b/gusto/equations.py index 37e0e2193..b4581fdef 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -635,7 +635,7 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, lambda t: Term(ufl.replace( t.form, {t.get(transporting_velocity): u}), t.labels)) ke_form = transporting_velocity.remove(ke_form) - u_adv = prognostic(advection_equation_circulation_form(domain, w, u), "u") + 0.5*ke_form + u_adv = prognostic(advection_equation_circulation_form(domain, w, u), "u") + ke_form else: raise ValueError("Invalid u_transport_option: %s" % u_transport_option) diff --git a/gusto/transport_forms.py b/gusto/transport_forms.py index f9ce9801b..87af6f692 100644 --- a/gusto/transport_forms.py +++ b/gusto/transport_forms.py @@ -368,7 +368,7 @@ def kinetic_energy_form(domain, test, q): """ ubar = Function(domain.spaces("HDiv")) - L = div(test)*inner(q, ubar)*dx + L = -0.5*div(test)*inner(q, ubar)*dx form = transporting_velocity(L, ubar) @@ -409,7 +409,7 @@ def advection_equation_circulation_form(domain, test, q, form = ( vector_invariant_form(domain, test, q, ibp=ibp) - - kinetic_energy_form(domain, test, q) + + kinetic_energy_form(domain, test, q) ) return form From d90fb4de7503b121612b0faac484b3c0490c7caf Mon Sep 17 00:00:00 2001 From: jshipton Date: Sat, 18 Mar 2023 19:04:12 +0000 Subject: [PATCH 07/33] fix minus sign --- gusto/equations.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gusto/equations.py b/gusto/equations.py index b4581fdef..1c2d070d0 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -635,7 +635,7 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, lambda t: Term(ufl.replace( t.form, {t.get(transporting_velocity): u}), t.labels)) ke_form = transporting_velocity.remove(ke_form) - u_adv = prognostic(advection_equation_circulation_form(domain, w, u), "u") + ke_form + u_adv = prognostic(advection_equation_circulation_form(domain, w, u), "u") - ke_form else: raise ValueError("Invalid u_transport_option: %s" % u_transport_option) From 5f49d46dea7f320ec19f62086d300d09ba1416cc Mon Sep 17 00:00:00 2001 From: jshipton Date: Sat, 18 Mar 2023 19:18:27 +0000 Subject: [PATCH 08/33] a couple of hard coded constant_jacobian=False args --- gusto/diagnostics.py | 2 +- gusto/linear_solvers.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/gusto/diagnostics.py b/gusto/diagnostics.py index 6e93f1309..a9d3c2d7b 100644 --- a/gusto/diagnostics.py +++ b/gusto/diagnostics.py @@ -1393,7 +1393,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/linear_solvers.py b/gusto/linear_solvers.py index 547fe99d7..a998ea155 100644 --- a/gusto/linear_solvers.py +++ b/gusto/linear_solvers.py @@ -590,7 +590,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, From 905e4177994e0a4c4c5fe0f3cd1744bd31ccc8bf Mon Sep 17 00:00:00 2001 From: jshipton Date: Thu, 18 May 2023 17:24:46 +0100 Subject: [PATCH 09/33] these LVPs need constant_Jacobian=False --- gusto/forcing.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/gusto/forcing.py b/gusto/forcing.py index 0ad6fd995..10914bb4e 100644 --- a/gusto/forcing.py +++ b/gusto/forcing.py @@ -87,11 +87,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 ) solver_parameters = {} From 0d38a7907ffc3c033747db802367cd7f78b57fc6 Mon Sep 17 00:00:00 2001 From: jshipton Date: Thu, 18 May 2023 17:27:09 +0100 Subject: [PATCH 10/33] some fixes to setup script --- examples/shallow_water/mm_ot_sw_galewsky_jet.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/examples/shallow_water/mm_ot_sw_galewsky_jet.py b/examples/shallow_water/mm_ot_sw_galewsky_jet.py index d038aa71b..28ce4fbb2 100644 --- a/examples/shallow_water/mm_ot_sw_galewsky_jet.py +++ b/examples/shallow_water/mm_ot_sw_galewsky_jet.py @@ -1,7 +1,7 @@ from gusto import * from firedrake import IcosahedralSphereMesh, Constant, ge, le, exp, cos, \ conditional, interpolate, SpatialCoordinate, VectorFunctionSpace, \ - Function, assemble, dx, FunctionSpace,pi + Function, assemble, dx, FunctionSpace, pi, CellNormal import numpy as np @@ -16,15 +16,14 @@ # Domain mesh = IcosahedralSphereMesh(radius=R, - refinement_level=3, degree=1) + refinement_level=3, degree=2) x = SpatialCoordinate(mesh) -mesh.init_cell_orientations(x) domain = Domain(mesh, dt, 'BDM', 1) # Equation Omega = parameters.Omega fexpr = 2*Omega*x[2]/R -eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr) +eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, u_transport_option="circulation_form") # I/O perturb = True @@ -53,6 +52,7 @@ def update_pv(): def reinterpolate_coriolis(): eqns.prescribed_fields("coriolis").interpolate(fexpr) + domain.outward_normals.interpolate(CellNormal(domain.mesh)) monitor = MonitorFunction("PotentialVorticity", adapt_to="gradient") mesh_generator = OptimalTransportMeshGenerator(mesh, @@ -63,6 +63,7 @@ def reinterpolate_coriolis(): # Time stepper stepper = MeshMovement(eqns, io, transported_fields, mesh_generator=mesh_generator) +#stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields) # initial conditions u0 = stepper.fields("u") @@ -176,6 +177,7 @@ def initialise_fn(): mesh_generator.get_first_mesh(initialise_fn) eqns.prescribed_fields("coriolis").interpolate(fexpr) +domain.outward_normals.interpolate(CellNormal(domain.mesh)) pv() Dbar = Function(D0.function_space()).assign(H) From 54b1f05e538deec6e51fe28e3710cf46e336413e Mon Sep 17 00:00:00 2001 From: jshipton Date: Tue, 6 Jun 2023 15:42:20 +0100 Subject: [PATCH 11/33] IT WORKS! --- examples/shallow_water/mm_ot_sw_galewsky_jet.py | 16 +++++++++++----- gusto/time_discretisation.py | 2 +- gusto/timeloop.py | 14 +++++++++++--- 3 files changed, 23 insertions(+), 9 deletions(-) diff --git a/examples/shallow_water/mm_ot_sw_galewsky_jet.py b/examples/shallow_water/mm_ot_sw_galewsky_jet.py index 28ce4fbb2..e7c896ba3 100644 --- a/examples/shallow_water/mm_ot_sw_galewsky_jet.py +++ b/examples/shallow_water/mm_ot_sw_galewsky_jet.py @@ -17,11 +17,11 @@ # Domain mesh = IcosahedralSphereMesh(radius=R, refinement_level=3, degree=2) -x = SpatialCoordinate(mesh) domain = Domain(mesh, dt, 'BDM', 1) # Equation Omega = parameters.Omega +x = SpatialCoordinate(domain.mesh) fexpr = 2*Omega*x[2]/R eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, u_transport_option="circulation_form") @@ -51,11 +51,12 @@ def update_pv(): pv() def reinterpolate_coriolis(): - eqns.prescribed_fields("coriolis").interpolate(fexpr) + 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(mesh, +mesh_generator = OptimalTransportMeshGenerator(domain.mesh, monitor, pre_meshgen_callback=update_pv, post_meshgen_callback=reinterpolate_coriolis) @@ -161,23 +162,28 @@ def initialise_fn(): u0.project(uexpr, form_compiler_parameters={'quadrature_degree': 12}) - X = interpolate(mesh.coordinates, W) + 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() +# stepper.prescribed_uexpr = sphere_to_cartesian(mesh, u_zonal, u_merid) pv.setup(domain, stepper.fields) mesh_generator.get_first_mesh(initialise_fn) -eqns.prescribed_fields("coriolis").interpolate(fexpr) +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) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 1be341108..29e6b15dd 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -94,7 +94,7 @@ def __init__(self, domain, field_name=None, solver_parameters=None, self.field_name = field_name self.equation = None - self.dt = domain.dt + self.dt = 0.5*domain.dt self.limiter = limiter diff --git a/gusto/timeloop.py b/gusto/timeloop.py index ecdfe624a..afadd4f38 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -575,9 +575,14 @@ def __init__(self, equation_set, io, transport_schemes, self.mesh = mesh self.X0 = Function(mesh.coordinates) self.X1 = Function(mesh.coordinates) - - xold = Function(mesh.coordinates) - self.mesh_old = Mesh(xold) + 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 @@ -617,6 +622,8 @@ def timestep(self): test_D = TestFunction(Dm.function_space()) trial_D = TrialFunction(Dm.function_space()) + X1.assign(self.mesh.coordinates) + with timed_stage("Apply forcing terms"): self.forcing.apply(xn, xn, xstar(self.field_name), "explicit") @@ -629,6 +636,7 @@ def timestep(self): self.mesh_generator.pre_meshgen_callback() with timed_stage("Mesh generation"): X1.assign(self.mesh_generator.get_new_mesh()) + # X1.assign(X0) # Compute v (mesh velocity w.r.t. initial mesh) and # v1 (mesh velocity w.r.t. final mesh) From 34305f031143c12cb1a203200b1cefc13d55afd0 Mon Sep 17 00:00:00 2001 From: jshipton Date: Thu, 15 Jun 2023 10:36:24 +0100 Subject: [PATCH 12/33] remove debugging prints from timeloop --- gusto/timeloop.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/gusto/timeloop.py b/gusto/timeloop.py index afadd4f38..06ba9a74b 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -666,17 +666,14 @@ def timestep(self): self.mesh.coordinates.assign(X0) self.uadv.assign(0.5*(un - self.v_V1)) scheme.apply(xmid(name), xstar(name)) - print("before: ", name, xmid(name).dat.data.min(), xmid(name).dat.data.max()) if name == "u": - print("projecting u") um_.assign(xmid("u")) rhs = inner(um_, test_u)*dx with assemble(rhs).dat.vec as v: Lvec = v elif name == "D": - print("projecting D") Dm_.assign(xmid("D")) rhs = inner(Dm_, test_D)*dx with assemble(rhs).dat.vec as v: @@ -708,7 +705,6 @@ def timestep(self): with Dm.dat.vec as x_: ksp.solve(Lvec, x_) xmid(name).assign(Dm) - print("after: ", name, xmid(name).dat.data.min(), xmid(name).dat.data.max()) # transport field from xmid to xp on new mesh self.uadv.assign(0.5*(unp1 - self.v1_V1)) From 782ade1cf7f90ebc1b8417ae9a36a0542b77f03b Mon Sep 17 00:00:00 2001 From: jshipton Date: Thu, 15 Jun 2023 10:37:29 +0100 Subject: [PATCH 13/33] remove another debugging print --- gusto/time_discretisation.py | 1 - 1 file changed, 1 deletion(-) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index 29e6b15dd..ad1ea8fd6 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -1558,7 +1558,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 From 164c31aa5e4a6499203db910e711eb64247c7c2f Mon Sep 17 00:00:00 2001 From: jshipton Date: Thu, 15 Jun 2023 12:06:37 +0100 Subject: [PATCH 14/33] tidy code and fix memory issues --- gusto/timeloop.py | 62 ++++++++++++++++------------------------------- 1 file changed, 21 insertions(+), 41 deletions(-) diff --git a/gusto/timeloop.py b/gusto/timeloop.py index 06ba9a74b..2869d8ca5 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -594,6 +594,16 @@ def __init__(self, equation_set, io, transport_schemes, 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() + def setup_fields(self): super().setup_fields() self.x.add_fields(self.equation, levels=("mid",)) @@ -615,12 +625,6 @@ def timestep(self): unp1 = xnp1("u") X0 = self.X0 X1 = self.X1 - um, Dm = Function(self.equation.W).split() - um_, Dm_ = Function(self.equation.W).split() - test_u = TestFunction(um.function_space()) - trial_u = TrialFunction(um.function_space()) - test_D = TestFunction(Dm.function_space()) - trial_D = TrialFunction(Dm.function_space()) X1.assign(self.mesh.coordinates) @@ -636,7 +640,6 @@ def timestep(self): self.mesh_generator.pre_meshgen_callback() with timed_stage("Mesh generation"): X1.assign(self.mesh_generator.get_new_mesh()) - # X1.assign(X0) # Compute v (mesh velocity w.r.t. initial mesh) and # v1 (mesh velocity w.r.t. final mesh) @@ -667,44 +670,21 @@ def timestep(self): self.uadv.assign(0.5*(un - self.v_V1)) scheme.apply(xmid(name), xstar(name)) - if name == "u": - um_.assign(xmid("u")) - rhs = inner(um_, test_u)*dx - with assemble(rhs).dat.vec as v: - Lvec = v - - elif name == "D": - Dm_.assign(xmid("D")) - rhs = inner(Dm_, test_D)*dx - with assemble(rhs).dat.vec as v: - Lvec = v + rhs = inner(xmid(name), self.tests[name])*dx + with assemble(rhs).dat.vec as v: + Lvec = v # put mesh_new into mesh self.mesh.coordinates.assign(X1) - if name == "u": - lhs = inner(trial_u, test_u)*dx - amat = assemble(lhs) - a = amat.M.handle - ksp = PETSc.KSP().create() - ksp.setOperators(a) - ksp.setFromOptions() - - with um.dat.vec as x_: - ksp.solve(Lvec, x_) - xmid(name).assign(um) - - elif name == "D": - lhs = inner(trial_D, test_D)*dx - amat = assemble(lhs) - a = amat.M.handle - ksp = PETSc.KSP().create() - ksp.setOperators(a) - ksp.setFromOptions() - - with Dm.dat.vec as x_: - ksp.solve(Lvec, x_) - xmid(name).assign(Dm) + 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(Lvec, x_) # transport field from xmid to xp on new mesh self.uadv.assign(0.5*(unp1 - self.v1_V1)) From 04569b0e219beec9839df7e683490c4b30e9b746 Mon Sep 17 00:00:00 2001 From: jshipton Date: Mon, 19 Jun 2023 11:14:40 +0100 Subject: [PATCH 15/33] enable adapting to field --- gusto/moving_mesh/monitor.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/gusto/moving_mesh/monitor.py b/gusto/moving_mesh/monitor.py index 1bc7064a8..73f0d1fbe 100644 --- a/gusto/moving_mesh/monitor.py +++ b/gusto/moving_mesh/monitor.py @@ -79,10 +79,10 @@ def setup(self, state_fields): if True: # not uniform: if self.adapt_to == "gradient": self.L_monitor = v_p1*sqrt(dot(self.gradq, self.gradq))*dx - else: + elif self.adapt_to == "hessian": self.L_monitor = v_p1*sqrt(inner(self.hessq, self.hessq))*dx - else: - self.L_monitor = v_p1*dx + else: + self.L_monitor = v_p1*dx u_dg = TrialFunction(DP1) v_dg = TestFunction(DP1) From 839423580f7fa85bea7dabde6c0b036d8bcc95c0 Mon Sep 17 00:00:00 2001 From: jshipton Date: Mon, 19 Jun 2023 11:25:36 +0100 Subject: [PATCH 16/33] put some info in the directory name --- examples/shallow_water/mm_ot_sw_galewsky_jet.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/examples/shallow_water/mm_ot_sw_galewsky_jet.py b/examples/shallow_water/mm_ot_sw_galewsky_jet.py index e7c896ba3..837ee2f71 100644 --- a/examples/shallow_water/mm_ot_sw_galewsky_jet.py +++ b/examples/shallow_water/mm_ot_sw_galewsky_jet.py @@ -6,8 +6,9 @@ import numpy as np day = 24.*60.*60. -dt = 480. +dt = 240. tmax = 6*day +ref_level = 4 # setup shallow water parameters R = 6371220. @@ -16,7 +17,7 @@ # Domain mesh = IcosahedralSphereMesh(radius=R, - refinement_level=3, degree=2) + refinement_level=ref_level, degree=2) domain = Domain(mesh, dt, 'BDM', 1) # Equation @@ -28,7 +29,7 @@ # I/O perturb = True if perturb: - dirname = "mm_ot_sw_galewsky_jet_perturbed" + dirname = "mm_ot_sw_galewsky_jet_perturbed_ref%s_dt%s" % (ref_level, dt) else: dirname = "mm_ot_sw_galewsky_jet_unperturbed" From 4ce168324470ea4a32ca29b0ca6d536ebbc267be Mon Sep 17 00:00:00 2001 From: jshipton Date: Thu, 24 Aug 2023 14:16:35 +0100 Subject: [PATCH 17/33] remove factor of a half --- gusto/time_discretisation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gusto/time_discretisation.py b/gusto/time_discretisation.py index ad1ea8fd6..f91b5c751 100644 --- a/gusto/time_discretisation.py +++ b/gusto/time_discretisation.py @@ -94,7 +94,7 @@ def __init__(self, domain, field_name=None, solver_parameters=None, self.field_name = field_name self.equation = None - self.dt = 0.5*domain.dt + self.dt = domain.dt self.limiter = limiter From 71bcb2a00ef56e1cae4e0f7997627c3e438e75f9 Mon Sep 17 00:00:00 2001 From: jshipton Date: Thu, 24 Aug 2023 14:19:23 +0100 Subject: [PATCH 18/33] refactor mesh movement timeloop so that the new mesh is only computed once and the transport step on the old mesh is done outside of the iteration --- gusto/timeloop.py | 121 +++++++++++++++++++++++++++------------------- 1 file changed, 71 insertions(+), 50 deletions(-) diff --git a/gusto/timeloop.py b/gusto/timeloop.py index 2869d8ca5..c14435457 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -603,6 +603,8 @@ def __init__(self, equation_set, io, transport_schemes, self.ksp = {} for name in self.equation.field_names: self.ksp[name] = PETSc.KSP().create() + self.Lvec = {} + def setup_fields(self): super().setup_fields() @@ -626,73 +628,91 @@ def timestep(self): 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)) - for k in range(self.maxk): - - X0.assign(X1) + # 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_generator.pre_meshgen_callback() - with timed_stage("Mesh generation"): - X1.assign(self.mesh_generator.get_new_mesh()) + self.mesh.coordinates.assign(X0) + self.v_V1.project(self.v) - # 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(X1) + self.v1_V1.project(self.v1) - self.mesh.coordinates.assign(X0) - self.v_V1.project(self.v) + else: + self.mesh.coordinates.assign(X0) + self.v_V1.project((X1 - X0)/self.dt) - self.mesh.coordinates.assign(X1) - self.v1_V1.project(self.v1) + self.mesh.coordinates.assign(X1) + self.v1_V1.project(-(X0 - X1)/self.dt) - else: + 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.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)) - - rhs = inner(xmid(name), self.tests[name])*dx - with assemble(rhs).dat.vec as v: - Lvec = v - - # put mesh_new into mesh - self.mesh.coordinates.assign(X1) + 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() - 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() + # 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(Lvec, x_) + with xmid(name).dat.vec as x_: + self.ksp[name].solve(self.Lvec[name], x_) - # transport field from xmid to xp on new mesh - self.uadv.assign(0.5*(unp1 - self.v1_V1)) - scheme.apply(xp(name), xmid(name)) + for k in range(self.maxk): - self.mesh_generator.post_meshgen_callback() + 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.assign(0.) # xrhs is the residual which goes in the linear solve + # xrhs is the residual which goes in the linear solve + xrhs.assign(0.) for i in range(self.maxi): @@ -702,7 +722,8 @@ def timestep(self): xrhs -= xnp1(self.field_name) with timed_stage("Implicit solve"): - self.linear_solver.solve(xrhs, dy) # solves linear system and places result in dy + # solves linear system and places result in dy + self.linear_solver.solve(xrhs, dy) xnp1X = xnp1(self.field_name) xnp1X += dy From cd0232f73681b58674a807d63f61e0f97548323d Mon Sep 17 00:00:00 2001 From: jshipton Date: Thu, 24 Aug 2023 18:27:50 +0100 Subject: [PATCH 19/33] fix import error related to merge --- gusto/timeloop.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gusto/timeloop.py b/gusto/timeloop.py index 55cf7e77e..2c4c26aa7 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -8,7 +8,7 @@ from gusto.equations import PrognosticEquationSet from gusto.fields import TimeLevelFields, StateFields from gusto.forcing import Forcing -from gusto.fml.form_manipulation_labelling import drop, Label, Term +from gusto.fml import drop, Label, Term from gusto.labels import ( transport, diffusion, time_derivative, linearisation, prognostic, physics, transporting_velocity From af283f99fc7ea3345213794a67663b2f1c3de64d Mon Sep 17 00:00:00 2001 From: jshipton Date: Thu, 24 Aug 2023 19:23:32 +0100 Subject: [PATCH 20/33] fix some lint --- gusto/moving_mesh/monge_ampere.py | 14 +++++--------- gusto/moving_mesh/monitor.py | 14 +++++++------- 2 files changed, 12 insertions(+), 16 deletions(-) diff --git a/gusto/moving_mesh/monge_ampere.py b/gusto/moving_mesh/monge_ampere.py index 4c1c8b804..2fda50703 100644 --- a/gusto/moving_mesh/monge_ampere.py +++ b/gusto/moving_mesh/monge_ampere.py @@ -1,6 +1,7 @@ from firedrake import * 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): @@ -42,7 +43,7 @@ def setup(self): # get mesh area self.total_area = assemble(Constant(1.0)*dx(self.mesh)) - ### SET UP FUNCTIONS ### + # 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 @@ -65,7 +66,7 @@ def setup(self): self.m = Function(P1) self.theta = Constant(0.0) - ### DEFINE MESH EQUATIONS ### + # define mesh equations v, tau = TestFunctions(MixedSpace) if self.meshtype == "plane": @@ -101,7 +102,7 @@ def setup(self): elif self.meshtype == "sphere": self.gradphi_cts2 = Function(W_cts) # extra, as gradphi_cts not necessarily tangential - ### SET UP INITIAL SIGMA ### (needed on sphere) + # set up initial sigma (needed on sphere) sigma_ = TrialFunction(TensorP2) tau_ = TestFunction(TensorP2) sigma_ini = Function(TensorP2) @@ -116,7 +117,7 @@ def setup(self): self.phisigma.sub(1).assign(sigma_ini) - ### SOLVER OPTIONS FOR MESH GENERATION ### + # solver options for mesh generation phi__, sigma__ = TrialFunctions(MixedSpace) v__, tau__ = TestFunctions(MixedSpace) @@ -197,14 +198,11 @@ def update_mxtheta(self, cursol): 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.mesh_in.coordinates.assign(self.x) self.initialise_fn() - #self.monitor.mesh.coordinates.dat.data[:] = self.mesh.coordinates.dat.data_ro[:] 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.mesh.coordinates.assign(self.x) 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) @@ -264,8 +262,6 @@ def get_new_mesh(self): # Move data from internal mesh to output mesh. self.own_output_coords.interpolate(Constant(self.R)*self.x) - # self.mesh.coordinates.assign(self.x) - # self.output_coords.dat.data[:] = self.mesh.coordinates.dat.data_ro[:] self.output_coords.dat.data[:] = self.own_output_coords.dat.data_ro[:] return self.output_coords diff --git a/gusto/moving_mesh/monitor.py b/gusto/moving_mesh/monitor.py index 73f0d1fbe..856c0a4f9 100644 --- a/gusto/moving_mesh/monitor.py +++ b/gusto/moving_mesh/monitor.py @@ -1,10 +1,10 @@ from firedrake import * -import math 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, @@ -30,7 +30,7 @@ def setup(self, state_fields): 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 ### + # 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) @@ -43,7 +43,7 @@ def setup(self, state_fields): # 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) @@ -51,7 +51,7 @@ def setup(self, state_fields): self.dm = Function(DP1) self.m_int_form = self.m_prereg*dx - ### DEFINE MONITOR FUNCTION IN TERMS OF q ### + # define monitor function in terms of q if False: # for plane v_ones = as_vector(np.ones(2)) else: @@ -72,7 +72,7 @@ def setup(self, state_fields): 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 @@ -129,7 +129,7 @@ def update_monitor(self): 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) + # 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 @@ -150,4 +150,4 @@ def get_monitor_on_new_mesh(self, m, x_old, x_new): 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) + # mmax_post = max(self.m.dat.data)/min(self.m.dat.data) From f9ee380b115ad319537702f4426d0f28704d6570 Mon Sep 17 00:00:00 2001 From: jshipton Date: Thu, 24 Aug 2023 21:45:16 +0100 Subject: [PATCH 21/33] tiny lint fix --- gusto/timeloop.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/gusto/timeloop.py b/gusto/timeloop.py index 2c4c26aa7..aab6a26aa 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -1,7 +1,7 @@ """Classes for controlling the timestepping loop.""" from abc import ABCMeta, abstractmethod, abstractproperty -from firedrake import Function, Projector, Constant, Mesh, TrialFunction, \ +from firedrake import Function, Projector, Constant, TrialFunction, \ TestFunction, assemble, inner, dx, split from firedrake.petsc import PETSc from pyop2.profiling import timed_stage @@ -723,7 +723,6 @@ def __init__(self, equation_set, io, transport_schemes, self.ksp[name] = PETSc.KSP().create() self.Lvec = {} - def setup_fields(self): super().setup_fields() self.x.add_fields(self.equation, levels=("mid",)) From e2451b1c4a5cb74cfa57a5a3a16d4c470eea4940 Mon Sep 17 00:00:00 2001 From: jshipton Date: Fri, 25 Aug 2023 17:28:00 +0100 Subject: [PATCH 22/33] get perp from the right place - not sure how this works on main! --- unit-tests/fml_tests/test_replace_perp.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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, From 5f5b6a10627bc8bb187fe1fec1d959ea6f154c1a Mon Sep 17 00:00:00 2001 From: jshipton Date: Fri, 25 Aug 2023 17:53:47 +0100 Subject: [PATCH 23/33] removing hack to set up diagnostics earlier as this breaks tests --- gusto/timeloop.py | 1 - 1 file changed, 1 deletion(-) diff --git a/gusto/timeloop.py b/gusto/timeloop.py index aab6a26aa..2fa904860 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -44,7 +44,6 @@ def __init__(self, equation, io): self.setup_scheme() self.io.log_parameters(equation) - self.io.setup_diagnostics(self.fields) @abstractproperty def transporting_velocity(self): From 2c3652717d71f8901cac8e27f6e40197dbbdfa13 Mon Sep 17 00:00:00 2001 From: jshipton Date: Fri, 25 Aug 2023 18:30:57 +0100 Subject: [PATCH 24/33] I think this fixes the tests but might break mesh movement - check sign of ke term!! --- gusto/domain.py | 6 +++++- gusto/equations.py | 25 +++++++++++++++++-------- gusto/timeloop.py | 7 +++++++ 3 files changed, 29 insertions(+), 9 deletions(-) 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 e269972cc..4eb641679 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -669,20 +669,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 == "circulation_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') elif u_transport_option == "circulation_form": - ke_form = prognostic(kinetic_energy_form(domain, w, 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 = prognostic(advection_equation_circulation_form(domain, w, u), "u") - ke_form + ke_form = prognostic(kinetic_energy_form(w, u, u), "u") + if domain.move_mesh: + 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) + ke_form *= -1 + 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/timeloop.py b/gusto/timeloop.py index 2fa904860..c73a864cc 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -678,6 +678,9 @@ def __init__(self, equation_set, io, transport_schemes, 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, @@ -708,6 +711,10 @@ def __init__(self, equation_set, io, transport_schemes, 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() From 40e911609e2ea40894c3ddf33c98de85a1d659fc Mon Sep 17 00:00:00 2001 From: jshipton Date: Fri, 25 Aug 2023 20:06:56 +0100 Subject: [PATCH 25/33] more lint fixing --- gusto/equations.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/gusto/equations.py b/gusto/equations.py index 4eb641679..3448d829e 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, @@ -670,7 +671,6 @@ 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: From 389bde772474da6d5f1b25d6e3e83f7f4f9eecb6 Mon Sep 17 00:00:00 2001 From: jshipton Date: Fri, 25 Aug 2023 20:07:41 +0100 Subject: [PATCH 26/33] more fixes for merge with main --- examples/shallow_water/mm_ot_sw_galewsky_jet.py | 11 ++++++----- gusto/timeloop.py | 4 +++- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/examples/shallow_water/mm_ot_sw_galewsky_jet.py b/examples/shallow_water/mm_ot_sw_galewsky_jet.py index 837ee2f71..c799cfbb8 100644 --- a/examples/shallow_water/mm_ot_sw_galewsky_jet.py +++ b/examples/shallow_water/mm_ot_sw_galewsky_jet.py @@ -18,7 +18,7 @@ # Domain mesh = IcosahedralSphereMesh(radius=R, refinement_level=ref_level, degree=2) -domain = Domain(mesh, dt, 'BDM', 1) +domain = Domain(mesh, dt, 'BDM', 1, move_mesh=True) # Equation Omega = parameters.Omega @@ -35,19 +35,19 @@ output = OutputParameters(dirname=dirname, dumplist_latlon=['D', 'PotentialVorticity', - 'RelativeVorticity'], - log_level="INFO") + 'RelativeVorticity']) + pv = PotentialVorticity() diagnostic_fields = [pv, RelativeVorticity()] io = IO(domain, output, diagnostic_fields=diagnostic_fields) # Transport schemes -transported_fields = [ImplicitMidpoint(domain, "u"), +transported_fields = [TrapeziumRule(domain, "u"), SSPRK3(domain, "D")] +transport_methods = [DGUpwind(eqns, "u"), DGUpwind(eqns, "D")] # Mesh movement - def update_pv(): pv() @@ -64,6 +64,7 @@ def reinterpolate_coriolis(): # Time stepper stepper = MeshMovement(eqns, io, transported_fields, + spatial_methods=transport_methods, mesh_generator=mesh_generator) #stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields) diff --git a/gusto/timeloop.py b/gusto/timeloop.py index c73a864cc..e9280c5c2 100644 --- a/gusto/timeloop.py +++ b/gusto/timeloop.py @@ -5,7 +5,7 @@ TestFunction, assemble, inner, dx, split from firedrake.petsc import PETSc from pyop2.profiling import timed_stage -from gusto.equations import PrognosticEquationSet +from gusto.equations import PrognosticEquationSet, ShallowWaterEquations from gusto.fields import TimeLevelFields, StateFields from gusto.forcing import Forcing from gusto.fml import drop, Label, Term @@ -668,6 +668,7 @@ def timestep(self): class MeshMovement(SemiImplicitQuasiNewton): def __init__(self, equation_set, io, transport_schemes, + spatial_methods, auxiliary_equations_and_schemes=None, linear_solver=None, diffusion_schemes=None, @@ -684,6 +685,7 @@ def __init__(self, equation_set, io, transport_schemes, 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, From dc238c8a53a3d155f6ef3e8c8e7d3f84fdd0aaf2 Mon Sep 17 00:00:00 2001 From: jshipton Date: Fri, 25 Aug 2023 22:11:06 +0100 Subject: [PATCH 27/33] I think this works with the new definition of the transport forms --- examples/shallow_water/mm_ot_sw_galewsky_jet.py | 2 +- gusto/equations.py | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/examples/shallow_water/mm_ot_sw_galewsky_jet.py b/examples/shallow_water/mm_ot_sw_galewsky_jet.py index c799cfbb8..637eb2774 100644 --- a/examples/shallow_water/mm_ot_sw_galewsky_jet.py +++ b/examples/shallow_water/mm_ot_sw_galewsky_jet.py @@ -24,7 +24,7 @@ Omega = parameters.Omega x = SpatialCoordinate(domain.mesh) fexpr = 2*Omega*x[2]/R -eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, u_transport_option="circulation_form") +eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, u_transport_option="vector_invariant_form") # I/O perturb = True diff --git a/gusto/equations.py b/gusto/equations.py index 3448d829e..a046f54ff 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -674,15 +674,13 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, # Mesh movement requires the circulation form, and an # additional modification if domain.move_mesh: - assert u_transport_option == "circulation_form" + 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') - elif u_transport_option == "circulation_form": - ke_form = prognostic(kinetic_energy_form(w, u, u), "u") if domain.move_mesh: ke_form = transport.remove(ke_form) ke_form = ke_form.label_map( @@ -690,7 +688,9 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, lambda t: Term(ufl.replace( t.form, {t.get(transporting_velocity): u}), t.labels)) ke_form = transporting_velocity.remove(ke_form) - ke_form *= -1 + 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 else: raise ValueError("Invalid u_transport_option: %s" % u_transport_option) From 809ceedf2d4c5889da431931f94a30ad83708841 Mon Sep 17 00:00:00 2001 From: jshipton Date: Sat, 26 Aug 2023 15:05:52 +0100 Subject: [PATCH 28/33] make solver non-convergence warning go to logger --- gusto/moving_mesh/monge_ampere.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/gusto/moving_mesh/monge_ampere.py b/gusto/moving_mesh/monge_ampere.py index 2fda50703..1b05385f2 100644 --- a/gusto/moving_mesh/monge_ampere.py +++ b/gusto/moving_mesh/monge_ampere.py @@ -1,4 +1,5 @@ from firedrake import * +from gusto.logging import logger import numpy as np @@ -257,8 +258,9 @@ def get_new_mesh(self): # Generate new mesh, coords put in self.x try: self.mesh_solv.solve() - except: - print("mesh solver did not converge - oh well, continuing anyway!") + 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) From dda98e97ac00e7d9049276dcaeed12ad8accd2b0 Mon Sep 17 00:00:00 2001 From: jshipton Date: Sat, 26 Aug 2023 15:07:48 +0100 Subject: [PATCH 29/33] fix linting of moving mesh example file --- examples/shallow_water/mm_ot_sw_galewsky_jet.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/examples/shallow_water/mm_ot_sw_galewsky_jet.py b/examples/shallow_water/mm_ot_sw_galewsky_jet.py index 637eb2774..a0ff5a432 100644 --- a/examples/shallow_water/mm_ot_sw_galewsky_jet.py +++ b/examples/shallow_water/mm_ot_sw_galewsky_jet.py @@ -1,7 +1,7 @@ from gusto import * from firedrake import IcosahedralSphereMesh, Constant, ge, le, exp, cos, \ conditional, interpolate, SpatialCoordinate, VectorFunctionSpace, \ - Function, assemble, dx, FunctionSpace, pi, CellNormal + Function, assemble, dx, pi, CellNormal import numpy as np @@ -47,15 +47,18 @@ 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, @@ -66,7 +69,6 @@ def reinterpolate_coriolis(): stepper = MeshMovement(eqns, io, transported_fields, spatial_methods=transport_methods, mesh_generator=mesh_generator) -#stepper = SemiImplicitQuasiNewton(eqns, io, transported_fields) # initial conditions u0 = stepper.fields("u") @@ -158,6 +160,7 @@ def Dval(X): 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") @@ -179,7 +182,7 @@ def initialise_fn(): eqns.prescribed_fields("coriolis").interpolate(fexpr) pv() -# stepper.prescribed_uexpr = sphere_to_cartesian(mesh, u_zonal, u_merid) + pv.setup(domain, stepper.fields) mesh_generator.get_first_mesh(initialise_fn) From 2d9807d5b9f93a6c9780c7f51984504c733e7571 Mon Sep 17 00:00:00 2001 From: jshipton Date: Sat, 26 Aug 2023 15:08:14 +0100 Subject: [PATCH 30/33] not quite sure how this worked without the definition of the ke_form --- gusto/equations.py | 1 + 1 file changed, 1 insertion(+) diff --git a/gusto/equations.py b/gusto/equations.py index a046f54ff..5a10b51db 100644 --- a/gusto/equations.py +++ b/gusto/equations.py @@ -682,6 +682,7 @@ def __init__(self, domain, parameters, fexpr=None, bexpr=None, 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), From 17a7b5d2805298e84f5eaa4a79cb06ef8ed56e8f Mon Sep 17 00:00:00 2001 From: jshipton Date: Sat, 26 Aug 2023 22:39:56 +0100 Subject: [PATCH 31/33] don't kill the testing by adding a stupidly long example without the test-specific check at the top --- examples/shallow_water/mm_ot_sw_galewsky_jet.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/examples/shallow_water/mm_ot_sw_galewsky_jet.py b/examples/shallow_water/mm_ot_sw_galewsky_jet.py index a0ff5a432..a6705ff86 100644 --- a/examples/shallow_water/mm_ot_sw_galewsky_jet.py +++ b/examples/shallow_water/mm_ot_sw_galewsky_jet.py @@ -7,8 +7,16 @@ day = 24.*60.*60. dt = 240. -tmax = 6*day -ref_level = 4 + +if '--running-tests' in sys.argv: + ref_level = 3 + dt = 480. + tmax = 1440. +else: + # setup resolution and timestepping parameters for convergence test + ref_level = 4 + dt = 240. + tmax = 6*day # setup shallow water parameters R = 6371220. @@ -24,7 +32,8 @@ 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") +eqns = ShallowWaterEquations(domain, parameters, fexpr=fexpr, + u_transport_option="vector_invariant_form") # I/O perturb = True From 473d872eae86f0d6061a7a34bae4f3a94494a6c6 Mon Sep 17 00:00:00 2001 From: jshipton Date: Sun, 27 Aug 2023 14:44:05 +0100 Subject: [PATCH 32/33] import sys to fix running moving mesh example through pytest --- examples/shallow_water/mm_ot_sw_galewsky_jet.py | 1 + 1 file changed, 1 insertion(+) diff --git a/examples/shallow_water/mm_ot_sw_galewsky_jet.py b/examples/shallow_water/mm_ot_sw_galewsky_jet.py index a6705ff86..be54ed785 100644 --- a/examples/shallow_water/mm_ot_sw_galewsky_jet.py +++ b/examples/shallow_water/mm_ot_sw_galewsky_jet.py @@ -4,6 +4,7 @@ Function, assemble, dx, pi, CellNormal import numpy as np +import sys day = 24.*60.*60. dt = 240. From bd5eb9bcfd2be03abffef163b1c39308fa98ddd3 Mon Sep 17 00:00:00 2001 From: jshipton Date: Sun, 27 Aug 2023 15:56:13 +0100 Subject: [PATCH 33/33] set up cubed sphere galewsky with moving mesh --- .../shallow_water/mm_ot_sw_galewsky_jet.py | 42 +++++++++++++------ 1 file changed, 29 insertions(+), 13 deletions(-) diff --git a/examples/shallow_water/mm_ot_sw_galewsky_jet.py b/examples/shallow_water/mm_ot_sw_galewsky_jet.py index be54ed785..33db53f0e 100644 --- a/examples/shallow_water/mm_ot_sw_galewsky_jet.py +++ b/examples/shallow_water/mm_ot_sw_galewsky_jet.py @@ -1,5 +1,6 @@ from gusto import * -from firedrake import IcosahedralSphereMesh, Constant, ge, le, exp, cos, \ +from firedrake import IcosahedralSphereMesh, \ + Constant, ge, le, exp, cos, \ conditional, interpolate, SpatialCoordinate, VectorFunctionSpace, \ Function, assemble, dx, pi, CellNormal @@ -13,10 +14,16 @@ ref_level = 3 dt = 480. tmax = 1440. + cubed_sphere = False else: - # setup resolution and timestepping parameters for convergence test - ref_level = 4 - dt = 240. + # 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 @@ -24,10 +31,25 @@ 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 -mesh = IcosahedralSphereMesh(radius=R, - refinement_level=ref_level, degree=2) -domain = Domain(mesh, dt, 'BDM', 1, move_mesh=True) +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 @@ -37,12 +59,6 @@ u_transport_option="vector_invariant_form") # I/O -perturb = True -if perturb: - dirname = "mm_ot_sw_galewsky_jet_perturbed_ref%s_dt%s" % (ref_level, dt) -else: - dirname = "mm_ot_sw_galewsky_jet_unperturbed" - output = OutputParameters(dirname=dirname, dumplist_latlon=['D', 'PotentialVorticity', 'RelativeVorticity'])