diff --git a/requirements.txt b/requirements.txt index cfc2509..4c89288 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,5 @@ pandas argparse rasterio -scipy == 1.9.1 +git+https://github.com/g-adopt/g-adopt@angus-g/ufl-update#egg=gadopt +git+https://github.com/thwaitesproject/FInAT.git@fix-hcurl-degree#egg=finat diff --git a/tests/adjoint/test_adjoint.py b/tests/adjoint/test_adjoint.py index 44e8c27..bdacdb0 100644 --- a/tests/adjoint/test_adjoint.py +++ b/tests/adjoint/test_adjoint.py @@ -766,8 +766,6 @@ def solve(self): assert np.allclose(tt, [2.0, 2.0, 2.0], rtol=5e-1) -test_2d_isomip_cavity_salfunctional(1) - def run_isomip(T, dump_flag=False, init_p_flag=True, mumps_pressure_projection=True, conditional_melt_flag=True, gammaT_control=False): # ISOMIP+ setup 2d slice with extuded mesh #Buoyancy driven overturning circulation diff --git a/thwaites/interpolate.py b/thwaites/interpolate.py index eb1b335..45b7162 100644 --- a/thwaites/interpolate.py +++ b/thwaites/interpolate.py @@ -83,8 +83,8 @@ def interpolate(f, Q, method='linear', y_transect=None): mesh = Q.mesh() element = Q.ufl_element() - if len(element.sub_elements()) > 0: - element = element.sub_elements()[0] + if len(element.sub_elements) > 0: + element = element.sub_elements[0] V = firedrake.VectorFunctionSpace(mesh, element) X = firedrake.interpolate(mesh.coordinates, V).dat.data_ro diff --git a/thwaites/limiter.py b/thwaites/limiter.py index 176a8cc..72f3ade 100644 --- a/thwaites/limiter.py +++ b/thwaites/limiter.py @@ -2,98 +2,11 @@ Slope limiters for discontinuous fields """ from __future__ import absolute_import -from firedrake import VertexBasedLimiter, FunctionSpace, TrialFunction, LinearSolver, TestFunction, dx, assemble +from firedrake import FunctionSpace, TestFunction, assemble from firedrake import Function, dS_v, ds_v, conditional, avg -from firedrake import TensorProductElement -import numpy as np -import ufl from pyop2.profiling import timed_region, timed_function, timed_stage # NOQA from pyop2 import op2 - - -def assert_function_space(fs, family, degree): - """ - Checks the family and degree of function space. - - Raises AssertionError if function space differs. - If the function space lies on an extruded mesh, checks both spaces of the - outer product. - - :arg fs: function space - :arg string family: name of element family - :arg int degree: polynomial degree of the function space - """ - fam_list = family - if not isinstance(family, list): - fam_list = [family] - ufl_elem = fs.ufl_element() - if isinstance(ufl_elem, ufl.VectorElement): - ufl_elem = ufl_elem.sub_elements()[0] - - if ufl_elem.family() == 'TensorProductElement': - # extruded mesh - A, B = ufl_elem.sub_elements() - assert A.family() in fam_list, \ - 'horizontal space must be one of {0:s}'.format(fam_list) - assert B.family() in fam_list, \ - 'vertical space must be {0:s}'.format(fam_list) - assert A.degree() == degree, \ - 'degree of horizontal space must be {0:d}'.format(degree) - assert B.degree() == degree, \ - 'degree of vertical space must be {0:d}'.format(degree) - else: - # assume 2D mesh - assert ufl_elem.family() in fam_list, \ - 'function space must be one of {0:s}'.format(fam_list) - assert ufl_elem.degree() == degree, \ - 'degree of function space must be {0:d}'.format(degree) - - -def get_extruded_base_element(ufl_element): - """ - Return UFL TensorProductElement of an extruded UFL element. - - In case of a non-extruded mesh, returns the element itself. - """ - if isinstance(ufl_element, ufl.HDivElement): - ufl_element = ufl_element._element - if isinstance(ufl_element, ufl.MixedElement): - ufl_element = ufl_element.sub_elements()[0] - if isinstance(ufl_element, ufl.VectorElement): - ufl_element = ufl_element.sub_elements()[0] # take the first component - if isinstance(ufl_element, ufl.EnrichedElement): - ufl_element = ufl_element._elements[0] - return ufl_element - - -def get_facet_mask(function_space, facet='bottom'): - """ - Returns the top/bottom nodes of extruded 3D elements. - - :arg function_space: Firedrake :class:`FunctionSpace` object - :kwarg str facet: 'top' or 'bottom' - - .. note:: - The definition of top/bottom depends on the direction of the extrusion. - Here we assume that the mesh has been extruded upwards (along positive - z axis). - """ - from tsfc.finatinterface import create_element as create_finat_element - - # get base element - elem = get_extruded_base_element(function_space.ufl_element()) - assert isinstance(elem, TensorProductElement), \ - f'function space must be defined on an extruded 3D mesh: {elem}' - # figure out number of nodes in sub elements - h_elt, v_elt = elem.sub_elements() - nb_nodes_h = create_finat_element(h_elt).space_dimension() - nb_nodes_v = create_finat_element(v_elt).space_dimension() - # compute top/bottom facet indices - # extruded dimension is the inner loop in index - # on interval elements, the end points are the first two dofs - offset = 0 if facet == 'bottom' else 1 - indices = np.arange(nb_nodes_h)*nb_nodes_v + offset - return indices +import gadopt class SqueezedDQ1Filter: @@ -125,7 +38,7 @@ def apply(self, u): u.assign(self.marker*self.u1 + (1-self.marker)*u) -class VertexBasedP1DGLimiter(VertexBasedLimiter): +class VertexBasedP1DGLimiter(gadopt.VertexBasedP1DGLimiter): """ Vertex based limiter for P1DG tracer fields, see Kuzmin (2010) @@ -143,126 +56,18 @@ def __init__(self, p1dg_space, squeezed_triangles=False): :arg squeezed_triangles: whether to deal with quads with squeezed vertical edges """ - assert_function_space(p1dg_space, ['Discontinuous Lagrange', 'DQ'], 1) - self.is_vector = p1dg_space.value_size > 1 - if self.is_vector: - p1dg_scalar_space = FunctionSpace(p1dg_space.mesh(), 'DG', 1) - super(VertexBasedP1DGLimiter, self).__init__(p1dg_scalar_space) - else: - super(VertexBasedP1DGLimiter, self).__init__(p1dg_space) - self.mesh = self.P0.mesh() - self.dim = self.mesh.geometric_dimension() - self.extruded = hasattr(self.mesh.ufl_cell(), 'sub_cells') - assert not self.extruded or len(p1dg_space.ufl_element().sub_elements()) > 0, \ - "Extruded mesh requires extruded function space" - assert not self.extruded or all(e.variant() == 'equispaced' for e in p1dg_space.ufl_element().sub_elements()), \ - "Extruded function space must be equivariant" - + super().__init__(p1dg_space) self.squeezed_triangles = squeezed_triangles if squeezed_triangles: self.squeezed_filter = SqueezedDQ1Filter(p1dg_space) - def _construct_centroid_solver(self): - """ - Constructs a linear problem for computing the centroids - - :return: LinearSolver instance - """ - u = TrialFunction(self.P0) - v = TestFunction(self.P0) - self.a_form = u * v * dx - a = assemble(self.a_form) - return LinearSolver(a, solver_parameters={'ksp_type': 'preonly', - 'pc_type': 'bjacobi', - 'sub_pc_type': 'ilu'}) - - def _update_centroids(self, field): - """ - Update centroid values - """ - b = assemble(TestFunction(self.P0) * field * dx) - self.centroid_solver.solve(self.centroids, b) - def compute_bounds(self, field): """ Re-compute min/max values of all neighbouring centroids :arg field: :class:`Function` to limit """ - # Call general-purpose bound computation. - super(VertexBasedP1DGLimiter, self).compute_bounds(field) - - # Add the average of lateral boundary facets to min/max fields - # NOTE this just computes the arithmetic mean of nodal values on the facet, - # which in general is not equivalent to the mean of the field over the bnd facet. - # This is OK for P1DG triangles, but not exact for the extruded case (quad facets) - from finat.finiteelementbase import entity_support_dofs - - if self.extruded: - entity_dim = (self.dim-2, 1) # get vertical facets - else: - entity_dim = self.dim-1 - boundary_dofs = entity_support_dofs(self.P1DG.finat_element, entity_dim) - local_facet_nodes = np.array([boundary_dofs[e] for e in sorted(boundary_dofs.keys())]) - n_bnd_nodes = local_facet_nodes.shape[1] - local_facet_idx = op2.Global(local_facet_nodes.shape, local_facet_nodes, dtype=np.int32, name='local_facet_idx') - code = """ - void my_kernel(double *qmax, double *qmin, double *field, unsigned int *facet, unsigned int *local_facet_idx) - { - double face_mean = 0.0; - for (int i = 0; i < %(nnodes)d; i++) { - unsigned int idx = local_facet_idx[facet[0]*%(nnodes)d + i]; - face_mean += field[idx]; - } - face_mean /= %(nnodes)d; - for (int i = 0; i < %(nnodes)d; i++) { - unsigned int idx = local_facet_idx[facet[0]*%(nnodes)d + i]; - qmax[idx] = fmax(qmax[idx], face_mean); - qmin[idx] = fmin(qmin[idx], face_mean); - } - }""" - bnd_kernel = op2.Kernel(code % {'nnodes': n_bnd_nodes}, 'my_kernel') - op2.par_loop(bnd_kernel, - self.P1DG.mesh().exterior_facets.set, - self.max_field.dat(op2.MAX, self.max_field.exterior_facet_node_map()), - self.min_field.dat(op2.MIN, self.min_field.exterior_facet_node_map()), - field.dat(op2.READ, field.exterior_facet_node_map()), - self.P1DG.mesh().exterior_facets.local_facet_dat(op2.READ), - local_facet_idx(op2.READ)) - if self.extruded: - # Add nodal values from surface/bottom boundaries - # NOTE calling firedrake par_loop with measure=ds_t raises an error - bottom_nodes = get_facet_mask(self.P1CG, 'bottom') - top_nodes = get_facet_mask(self.P1CG, 'top') - bottom_idx = op2.Global(len(bottom_nodes), bottom_nodes, dtype=np.int32, name='node_idx') - top_idx = op2.Global(len(top_nodes), top_nodes, dtype=np.int32, name='node_idx') - code = """ - void my_kernel(double *qmax, double *qmin, double *field, int *idx) { - double face_mean = 0; - for (int i=0; i<%(nnodes)d; i++) { - face_mean += field[idx[i]]; - } - face_mean /= %(nnodes)d; - for (int i=0; i<%(nnodes)d; i++) { - qmax[idx[i]] = fmax(qmax[idx[i]], face_mean); - qmin[idx[i]] = fmin(qmin[idx[i]], face_mean); - } - }""" - kernel = op2.Kernel(code % {'nnodes': len(bottom_nodes)}, 'my_kernel') - - op2.par_loop(kernel, self.mesh.cell_set, - self.max_field.dat(op2.MAX, self.max_field.function_space().cell_node_map()), - self.min_field.dat(op2.MIN, self.min_field.function_space().cell_node_map()), - field.dat(op2.READ, field.function_space().cell_node_map()), - bottom_idx(op2.READ), - iteration_region=op2.ON_BOTTOM) - - op2.par_loop(kernel, self.mesh.cell_set, - self.max_field.dat(op2.MAX, self.max_field.function_space().cell_node_map()), - self.min_field.dat(op2.MIN, self.min_field.function_space().cell_node_map()), - field.dat(op2.READ, field.function_space().cell_node_map()), - top_idx(op2.READ), - iteration_region=op2.ON_TOP) + super().compute_bounds(field) if self.squeezed_triangles: code = """ void my_kernel(double *qmax, double *qmin, double *marker) { @@ -309,14 +114,4 @@ def apply(self, field): with timed_stage('limiter'): if self.squeezed_triangles: self.squeezed_filter.apply(field) - - if self.is_vector: - tmp_func = self.P1DG.get_work_function() - fs = field.function_space() - for i in range(fs.value_size): - tmp_func.dat.data_with_halos[:] = field.dat.data_with_halos[:, i] - super(VertexBasedP1DGLimiter, self).apply(tmp_func) - field.dat.data_with_halos[:, i] = tmp_func.dat.data_with_halos[:] - self.P1DG.restore_work_function(tmp_func) - else: - super(VertexBasedP1DGLimiter, self).apply(field) + super().apply(field) diff --git a/thwaites/utility.py b/thwaites/utility.py index 70f8445..bf952aa 100644 --- a/thwaites/utility.py +++ b/thwaites/utility.py @@ -6,6 +6,7 @@ from firedrake import RW, READ, dx, par_loop, ExtrudedMesh import ufl import numpy as np +from gadopt.utility import extend_function_to_3d, is_continuous, normal_is_continuous # NOQA class CombinedSurfaceMeasure(ufl.Measure): @@ -32,55 +33,6 @@ def __rmul__(self, other): return other*self.ds_v + other*self.ds_t + other*self.ds_b -def _get_element(ufl_or_element): - if isinstance(ufl_or_element, ufl.FiniteElementBase): - return ufl_or_element - else: - return ufl_or_element.ufl_element() - - -def is_continuous(expr): - elem = _get_element(expr) - - family = elem.family() - if family == 'Lagrange': - return True - elif family == 'Discontinuous Lagrange' or family == 'DQ' or family == 'RTCE': - return False - elif isinstance(elem, ufl.HCurlElement) or isinstance(elem, ufl.HDivElement): - return False - elif family == 'TensorProductElement': - elem_h, elem_v = elem.sub_elements() - return is_continuous(elem_h) and is_continuous(elem_v) - elif family == 'EnrichedElement': - return all(is_continuous(e) for e in elem._elements) - else: - print(family) - raise NotImplementedError("Unknown finite element family") - - -def normal_is_continuous(expr): - elem = _get_element(expr) - - family = elem.family() - if family == 'Lagrange': - return True - elif family == 'Discontinuous Lagrange' or family == 'DQ' or family == 'RTCE': - return False - elif isinstance(elem, ufl.HCurlElement): - return False - elif isinstance(elem, ufl.HDivElement): - return True - elif family == 'TensorProductElement': - elem_h, elem_v = elem.sub_elements() - return normal_is_continuous(elem_h) and normal_is_continuous(elem_v) - elif family == 'EnrichedElement': - return all(normal_is_continuous(e) for e in elem._elements) - else: - print(family) - raise NotImplementedError("Unknown finite element family") - - def cell_size(mesh): if hasattr(mesh.ufl_cell(), 'sub_cells'): return sqrt(CellVolume(mesh)) @@ -168,29 +120,6 @@ def offset_backward_step_approx(x, k=1.0, x0=0.0): return 1.0 / (1.0 + exp(2.0*k*(x-x0))) -def extend_function_to_3d(func, mesh_extruded): - """ - Returns a 3D view of a 2D :class:`Function` on the extruded domain. - The 3D function resides in V x R function space, where V is the function - space of the source function. The 3D function shares the data of the 2D - function. - """ - fs = func.function_space() -# assert fs.mesh().geometric_dimension() == 2, 'Function must be in 2D space' - ufl_elem = fs.ufl_element() - family = ufl_elem.family() - degree = ufl_elem.degree() - name = func.name() - if isinstance(ufl_elem, ufl.VectorElement): - # vector function space - fs_extended = get_functionspace(mesh_extruded, family, degree, 'R', 0, dim=2, vector=True) - else: - fs_extended = get_functionspace(mesh_extruded, family, degree, 'R', 0) - func_extended = Function(fs_extended, name=name, val=func.dat._data) - func_extended.source = func - return func_extended - - class ExtrudedFunction(Function): """ A 2D :class:`Function` that provides a 3D view on the extruded domain.