diff --git a/gusto/diagnostics/diagnostics.py b/gusto/diagnostics/diagnostics.py index 90192f022..00272edac 100644 --- a/gusto/diagnostics/diagnostics.py +++ b/gusto/diagnostics/diagnostics.py @@ -1,14 +1,15 @@ """Common diagnostic fields.""" -from firedrake import (assemble, dot, dx, Function, sqrt, TestFunction, +from firedrake import (dot, dx, Function, sqrt, TestFunction, TrialFunction, Constant, grad, inner, FacetNormal, LinearVariationalProblem, LinearVariationalSolver, ds_b, ds_v, ds_t, dS_h, dS_v, ds, dS, div, avg, pi, TensorFunctionSpace, SpatialCoordinate, as_vector, - Projector, Interpolator, FunctionSpace, FiniteElement, + Projector, assemble, FunctionSpace, FiniteElement, TensorProductElement) from firedrake.assign import Assigner +from firedrake.__future__ import interpolate from ufl.domain import extract_unique_domain from abc import ABCMeta, abstractmethod, abstractproperty @@ -193,7 +194,7 @@ def setup(self, domain, state_fields, space=None): # Solve method must be declared in diagnostic's own setup routine if self.method == 'interpolate': - self.evaluator = Interpolator(self.expr, self.field) + self.evaluator = interpolate(self.expr, self.space) elif self.method == 'project': self.evaluator = Projector(self.expr, self.field) elif self.method == 'assign': @@ -207,7 +208,7 @@ def compute(self): logger.debug(f'Computing diagnostic {self.name} with {self.method} method') if self.method == 'interpolate': - self.evaluator.interpolate() + self.field.assign(assemble(self.evaluator)) elif self.method == 'assign': self.evaluator.assign() elif self.method == 'project': @@ -294,7 +295,7 @@ def setup(self, domain, state_fields, space=None): # Solve method must be declared in diagnostic's own setup routine if self.method == 'interpolate': - self.evaluator = Interpolator(self.expr, self.field) + self.evaluator = interpolate(self.expr, self.space) elif self.method == 'project': self.evaluator = Projector(self.expr, self.field) elif self.method == 'assign': diff --git a/gusto/diagnostics/shallow_water_diagnostics.py b/gusto/diagnostics/shallow_water_diagnostics.py index eb2dc4a32..024e58ecf 100644 --- a/gusto/diagnostics/shallow_water_diagnostics.py +++ b/gusto/diagnostics/shallow_water_diagnostics.py @@ -2,9 +2,10 @@ from firedrake import ( - dx, TestFunction, TrialFunction, grad, inner, curl, Function, Interpolator, + dx, TestFunction, TrialFunction, grad, inner, curl, Function, assemble, LinearVariationalProblem, LinearVariationalSolver, conditional ) +from firedrake.__future__ import interpolate from gusto.diagnostics.diagnostics import DiagnosticField, Energy __all__ = ["ShallowWaterKineticEnergy", "ShallowWaterPotentialEnergy", @@ -321,14 +322,14 @@ def setup(self, domain, state_fields): qsat_expr = self.equation.compute_saturation(state_fields.X( self.equation.field_name)) - self.qsat_interpolator = Interpolator(qsat_expr, self.qsat_func) + self.qsat_interpolate = interpolate(qsat_expr, space) self.expr = conditional(q_t < self.qsat_func, q_t, self.qsat_func) super().setup(domain, state_fields, space=space) def compute(self): """Performs the computation of the diagnostic field.""" - self.qsat_interpolator.interpolate() + self.qsat_func.assign(assemble(self.qsat_interpolate)) super().compute() @@ -371,7 +372,7 @@ def setup(self, domain, state_fields): qsat_expr = self.equation.compute_saturation(state_fields.X( self.equation.field_name)) - self.qsat_interpolator = Interpolator(qsat_expr, self.qsat_func) + self.qsat_interpolate = interpolate(qsat_expr, space) vapour = conditional(q_t < self.qsat_func, q_t, self.qsat_func) self.expr = q_t - vapour @@ -379,5 +380,5 @@ def setup(self, domain, state_fields): def compute(self): """Performs the computation of the diagnostic field.""" - self.qsat_interpolator.interpolate() + self.qsat_func.assign(assemble(self.qsat_interpolate)) super().compute() diff --git a/gusto/physics/microphysics.py b/gusto/physics/microphysics.py index e8313ab11..a5cbe3da5 100644 --- a/gusto/physics/microphysics.py +++ b/gusto/physics/microphysics.py @@ -4,9 +4,10 @@ """ from firedrake import ( - Interpolator, conditional, Function, dx, min_value, max_value, Constant, pi, - Projector + conditional, Function, dx, min_value, max_value, Constant, pi, + Projector, assemble ) +from firedrake.__future__ import interpolate from firedrake.fml import identity, Term, subject from gusto.equations import Phases, TracerVariableType from gusto.recovery import Recoverer, BoundaryMethod @@ -170,8 +171,7 @@ def __init__(self, equation, vapour_name='water_vapour', # Add terms to equations and make interpolators # -------------------------------------------------------------------- # self.source = [Function(V) for factor in factors] - self.source_interpolators = [Interpolator(sat_adj_expr*factor, source) - for factor, source in zip(factors, self.source)] + self.source_interpolate = [interpolate(sat_adj_expr*factor, V) for factor in factors] tests = [equation.tests[idx] for idx in V_idxs] @@ -195,8 +195,8 @@ def evaluate(self, x_in, dt): if isinstance(self.equation, CompressibleEulerEquations): self.rho_recoverer.project() # Evaluate the source - for interpolator in self.source_interpolators: - interpolator.interpolate() + for interpolator, src in zip(self.source_interpolate, self.source): + src.assign(assemble(interpolator)) class AdvectedMoments(Enum): @@ -440,7 +440,7 @@ def __init__(self, equation, cloud_name='cloud_water', rain_name='rain', min_value(accu_rate, self.cloud_water / self.dt), min_value(accr_rate + accu_rate, self.cloud_water / self.dt)))) - self.source_interpolator = Interpolator(rain_expr, self.source) + self.source_interpolate = interpolate(rain_expr, Vt) # Add term to equation's residual test_cl = equation.tests[self.cloud_idx] @@ -464,7 +464,7 @@ def evaluate(self, x_in, dt): self.rain.assign(x_in.subfunctions[self.rain_idx]) self.cloud_water.assign(x_in.subfunctions[self.cloud_idx]) # Evaluate the source - self.source.assign(self.source_interpolator.interpolate()) + self.source.assign(assemble(self.source_interpolate)) class EvaporationOfRain(PhysicsParametrisation): @@ -609,8 +609,7 @@ def __init__(self, equation, rain_name='rain', vapour_name='water_vapour', # Add terms to equations and make interpolators # -------------------------------------------------------------------- # self.source = [Function(V) for factor in factors] - self.source_interpolators = [Interpolator(evap_rate*factor, source) - for factor, source in zip(factors, self.source)] + self.source_interpolate = [interpolate(evap_rate*factor, V) for factor in factors] tests = [equation.tests[idx] for idx in V_idxs] @@ -634,5 +633,5 @@ def evaluate(self, x_in, dt): if isinstance(self.equation, CompressibleEulerEquations): self.rho_recoverer.project() # Evaluate the source - for interpolator in self.source_interpolators: - interpolator.interpolate() + for interpolator, src in zip(self.source_interpolate, self.source): + src.assign(assemble(interpolator)) diff --git a/gusto/physics/physics_parametrisation.py b/gusto/physics/physics_parametrisation.py index 2d43dae48..c304c23f6 100644 --- a/gusto/physics/physics_parametrisation.py +++ b/gusto/physics/physics_parametrisation.py @@ -8,7 +8,8 @@ """ from abc import ABCMeta, abstractmethod -from firedrake import Interpolator, Function, dx, Projector +from firedrake import Function, dx, Projector, assemble +from firedrake.__future__ import interpolate from firedrake.fml import subject from gusto.core.labels import PhysicsLabel from gusto.core.logging import logger @@ -117,14 +118,14 @@ def __init__(self, equation, variable_name, rate_expression, # Handle method of evaluating source/sink if self.method == 'interpolate': - self.source_interpolator = Interpolator(expression, V) + self.source_interpolate = interpolate(expression, V) else: self.source_projector = Projector(expression, V) # If not time-varying, evaluate for the first time here if not self.time_varying: if self.method == 'interpolate': - self.source.assign(self.source_interpolator.interpolate()) + self.source.assign(assemble(self.source_interpolate)) else: self.source.assign(self.source_projector.project()) @@ -140,7 +141,7 @@ def evaluate(self, x_in, dt): if self.time_varying: logger.info(f'Evaluating physics parametrisation {self.label.label}') if self.method == 'interpolate': - self.source.assign(self.source_interpolator.interpolate()) + self.source.assign(assemble(self.source_interpolate)) else: self.source.assign(self.source_projector.project()) else: diff --git a/gusto/physics/shallow_water_microphysics.py b/gusto/physics/shallow_water_microphysics.py index bd3f8a02e..7f1b91227 100644 --- a/gusto/physics/shallow_water_microphysics.py +++ b/gusto/physics/shallow_water_microphysics.py @@ -3,8 +3,9 @@ """ from firedrake import ( - Interpolator, conditional, Function, dx, min_value, max_value, Constant + conditional, Function, dx, min_value, max_value, Constant, assemble ) +from firedrake.__future__ import interpolate from firedrake.fml import subject from gusto.core.logging import logger from gusto.physics.physics_parametrisation import PhysicsParametrisation @@ -134,7 +135,7 @@ def __init__(self, equation, saturation_curve, self.evaluate) # interpolator does the conversion of vapour to rain - self.source_interpolator = Interpolator(conditional( + self.source_interpolate = interpolate(conditional( self.water_v > self.saturation_curve, (1/self.tau)*gamma_r*(self.water_v - self.saturation_curve), 0), Vv) @@ -159,7 +160,7 @@ def evaluate(self, x_in, dt): if self.set_tau_to_dt: self.tau.assign(dt) self.water_v.assign(x_in.subfunctions[self.Vv_idx]) - self.source.assign(self.source_interpolator.interpolate()) + self.source.assign(assemble(self.source_interpolate)) class SWSaturationAdjustment(PhysicsParametrisation): @@ -321,8 +322,8 @@ def __init__(self, equation, saturation_curve, # Add terms to equations and make interpolators # sources have the same order as V_idxs and factors self.source = [Function(Vc) for factor in factors] - self.source_interpolators = [Interpolator(sat_adj_expr*factor, source) - for factor, source in zip(factors, self.source)] + self.source_interpolate = [interpolate(sat_adj_expr*factor, Vc) + for factor in factors] # test functions have the same order as factors and sources (vapour, # cloud, depth, buoyancy) so that the correct test function multiplies @@ -359,5 +360,5 @@ def evaluate(self, x_in, dt): self.cloud.assign(x_in.subfunctions[self.Vc_idx]) if self.time_varying_gamma_v: self.gamma_v.interpolate(self.gamma_v_computation(x_in)) - for interpolator in self.source_interpolators: - interpolator.interpolate() + for interpolator, src in zip(self.source_interpolate, self.source): + src.assign(assemble(interpolator)) diff --git a/gusto/recovery/recovery.py b/gusto/recovery/recovery.py index b88c6ad6a..9d186dd48 100644 --- a/gusto/recovery/recovery.py +++ b/gusto/recovery/recovery.py @@ -13,7 +13,8 @@ Function, FunctionSpace, Interpolator, Projector, SpatialCoordinate, TensorProductElement, VectorFunctionSpace, as_vector, function, interval, - VectorElement) + VectorElement, assemble) +from firedrake.__future__ import interpolate from gusto.recovery import Averager from .recovery_kernels import (BoundaryRecoveryExtruded, BoundaryRecoveryHCurl, BoundaryGaussianElimination) @@ -144,14 +145,14 @@ def __init__(self, x_inout, method=BoundaryMethod.extruded, eff_coords=None): V_broken = FunctionSpace(mesh, BrokenElement(V_inout.ufl_element())) self.x_DG1_wrong = Function(V_broken) self.x_DG1_correct = Function(V_broken) - self.interpolator = Interpolator(self.x_inout, self.x_DG1_wrong) + self.interpolate = interpolate(self.x_inout, V_broken) self.averager = Averager(self.x_DG1_correct, self.x_inout) self.kernel = BoundaryGaussianElimination(V_broken) def apply(self): """Applies the boundary recovery process.""" if self.method == BoundaryMethod.taylor: - self.interpolator.interpolate() + self.x_DG1_wrong.assign(assemble(self.interpolate)) self.kernel.apply(self.x_DG1_wrong, self.x_DG1_correct, self.act_coords, self.eff_coords, self.num_ext) self.averager.project() @@ -275,7 +276,7 @@ def __init__(self, x_in, x_out, method='interpolate', boundary_method=None): self.boundary_recoverers.append(BoundaryRecoverer(x_out_scalars[i], method=BoundaryMethod.taylor, eff_coords=eff_coords[i])) - self.interpolate_to_vector = Interpolator(as_vector(x_out_scalars), self.x_out) + self.interpolate_to_vector = interpolate(as_vector(x_out_scalars), V_out) def project(self): """Perform the whole recovery step.""" @@ -294,7 +295,7 @@ def project(self): # Correct at boundaries boundary_recoverer.apply() # Combine the components to obtain the vector field - self.interpolate_to_vector.interpolate() + self.x_out.assign(assemble(self.interpolate_to_vector)) else: # Extrapolate at boundaries self.boundary_recoverer.apply() diff --git a/gusto/recovery/reversible_recovery.py b/gusto/recovery/reversible_recovery.py index fc8fbd332..d87449817 100644 --- a/gusto/recovery/reversible_recovery.py +++ b/gusto/recovery/reversible_recovery.py @@ -4,7 +4,8 @@ """ from gusto.core.conservative_projection import ConservativeProjector -from firedrake import (Projector, Function, Interpolator) +from firedrake import (Projector, Function, assemble) +from firedrake.__future__ import interpolate from .recovery import Recoverer __all__ = ["ReversibleRecoverer", "ConservativeRecoverer"] @@ -52,7 +53,7 @@ def __init__(self, source_field, target_field, reconstruct_opts): elif self.opts.project_high_method == 'project': self.projector_high = Projector(self.q_recovered, self.q_rec_high) elif self.opts.project_high_method == 'interpolate': - self.projector_high = Interpolator(self.q_recovered, self.q_rec_high) + self.projector_high = interpolate(self.q_recovered, target_field.function_space()) self.interp_high = True else: raise ValueError(f'Method {self.opts.project_high_method} ' @@ -68,7 +69,7 @@ def __init__(self, source_field, target_field, reconstruct_opts): elif self.opts.project_low_method == 'project': self.projector_low = Projector(self.q_rec_high, self.q_corr_low) elif self.opts.project_low_method == 'interpolate': - self.projector_low = Interpolator(self.q_rec_high, self.q_corr_low) + self.projector_low = interpolate(self.q_rec_high, source_field.function_space()) self.interp_low = True else: raise ValueError(f'Method {self.opts.project_low_method} ' @@ -84,17 +85,17 @@ def __init__(self, source_field, target_field, reconstruct_opts): elif self.opts.injection_method == 'project': self.injector = Projector(self.q_corr_low, self.q_corr_high) elif self.opts.injection_method == 'interpolate': - self.injector = Interpolator(self.q_corr_low, self.q_corr_high) + self.injector = interpolate(self.q_corr_low, target_field.function_space()) self.interp_inj = True else: raise ValueError(f'Method {self.opts.injection_method} for injection not valid') def project(self): self.recoverer.project() - self.projector_high.interpolate() if self.interp_high else self.projector_high.project() - self.projector_low.interpolate() if self.interp_low else self.projector_low.project() + self.q_rec_high.assign(assemble(self.projector_high)) if self.interp_high else self.projector_high.project() + self.q_corr_low.assign(assemble(self.projector_low)) if self.interp_low else self.projector_low.project() self.q_corr_low.assign(self.q_low - self.q_corr_low) - self.injector.interpolate() if self.interp_inj else self.injector.project() + self.q_corr_high.assign(assemble(self.injector)) if self.interp_inj else self.injector.project() self.q_high.assign(self.q_corr_high + self.q_rec_high) diff --git a/gusto/solvers/linear_solvers.py b/gusto/solvers/linear_solvers.py index 1f44cc950..ca684c6af 100644 --- a/gusto/solvers/linear_solvers.py +++ b/gusto/solvers/linear_solvers.py @@ -11,10 +11,11 @@ rhs, FacetNormal, div, dx, jump, avg, dS, dS_v, dS_h, ds_v, ds_t, ds_b, ds_tb, inner, action, dot, grad, Function, VectorSpaceBasis, cross, BrokenElement, FunctionSpace, MixedFunctionSpace, DirichletBC, as_vector, - Interpolator, conditional + assemble, conditional ) from firedrake.fml import Term, drop from firedrake.petsc import flatten_parameters +from firedrake.__future__ import interpolate from pyop2.profiling import timed_function, timed_region from gusto.equations.active_tracers import TracerVariableType @@ -660,11 +661,11 @@ def _setup_solver(self): qtbar = split(equation.X_ref)[3] # set up interpolators that use the X_ref values for D and b_e - self.q_sat_expr_interpolator = Interpolator( - equation.compute_saturation(equation.X_ref), self.q_sat_func) - self.q_v_interpolator = Interpolator( + self.q_sat_expr_interpolate = interpolate( + equation.compute_saturation(equation.X_ref), VD) + self.q_v_interpolate = interpolate( conditional(qtbar < self.q_sat_func, qtbar, self.q_sat_func), - self.qvbar) + VD) # bbar was be_bar and here we correct to become bbar bbar += equation.parameters.beta2 * self.qvbar @@ -729,8 +730,8 @@ def trace_nullsp(T): @timed_function("Gusto:UpdateReferenceProfiles") def update_reference_profiles(self): if self.equations.equivalent_buoyancy: - self.q_sat_expr_interpolator.interpolate() - self.q_v_interpolator.interpolate() + self.q_sat_func.assign(assemble(self.q_sat_expr_interpolate)) + self.qvbar.assign(assemble(self.q_v_interpolate)) @timed_function("Gusto:LinearSolve") def solve(self, xrhs, dy): diff --git a/gusto/time_discretisation/wrappers.py b/gusto/time_discretisation/wrappers.py index f0aae843e..18381bc1f 100644 --- a/gusto/time_discretisation/wrappers.py +++ b/gusto/time_discretisation/wrappers.py @@ -6,9 +6,10 @@ from abc import ABCMeta, abstractmethod from firedrake import ( - FunctionSpace, Function, BrokenElement, Projector, Interpolator, - Constant, as_ufl, dot, grad, TestFunction, MixedFunctionSpace + FunctionSpace, Function, BrokenElement, Projector, Constant, as_ufl, dot, grad, + TestFunction, MixedFunctionSpace, assemble ) +from firedrake.__future__ import interpolate from firedrake.fml import Term from gusto.core.configuration import EmbeddedDGOptions, RecoveryOptions, SUPGOptions from gusto.recovery import Recoverer, ReversibleRecoverer, ConservativeRecoverer @@ -264,7 +265,7 @@ def setup(self, original_space, post_apply_bcs): # Operators for projecting back self.interp_back = (self.options.project_low_method == 'interpolate') if self.options.project_low_method == 'interpolate': - self.x_out_projector = Interpolator(self.x_out, self.x_projected) + self.x_out_projector = interpolate(self.x_out, self.original_space) elif self.options.project_low_method == 'project': self.x_out_projector = Projector(self.x_out, self.x_projected, bcs=post_apply_bcs) @@ -302,7 +303,7 @@ def post_apply(self, x_out): """ if self.interp_back: - self.x_out_projector.interpolate() + self.x_projected.assign(assemble(self.x_out_projector)) else: self.x_out_projector.project() x_out.assign(self.x_projected) diff --git a/gusto/timestepping/semi_implicit_quasi_newton.py b/gusto/timestepping/semi_implicit_quasi_newton.py index c5f8a7ece..19142dbaa 100644 --- a/gusto/timestepping/semi_implicit_quasi_newton.py +++ b/gusto/timestepping/semi_implicit_quasi_newton.py @@ -4,10 +4,11 @@ """ from firedrake import ( - Function, Constant, TrialFunctions, DirichletBC, div, Interpolator, + Function, Constant, TrialFunctions, DirichletBC, div, assemble, LinearVariationalProblem, LinearVariationalSolver ) from firedrake.fml import drop, replace_subject +from firedrake.__future__ import interpolate from pyop2.profiling import timed_stage from gusto.core import TimeLevelFields, StateFields from gusto.core.labels import (transport, diffusion, time_derivative, @@ -234,9 +235,8 @@ def __init__(self, equation_set, io, transport_schemes, spatial_methods, V_DG = equation_set.domain.spaces('DG') self.predictor_field_in = Function(V_DG) div_factor = Constant(1.0) - (Constant(1.0) - self.alpha)*self.dt*div(self.x.n('u')) - self.predictor_interpolator = Interpolator( - self.x.star(predictor)*div_factor, self.predictor_field_in - ) + self.predictor_interpolate = interpolate( + self.x.star(predictor)*div_factor, V_DG) def _apply_bcs(self): """ @@ -334,7 +334,7 @@ def transport_fields(self, outer, xstar, xp): # Pre-multiply this variable by (1 - dt*beta*div(u)) V = xstar(name).function_space() field_out = Function(V) - self.predictor_interpolator.interpolate() + self.predictor_field_in.assign(assemble(self.predictor_interpolate)) scheme.apply(field_out, self.predictor_field_in) # xp is xstar plus the increment from the transported predictor