Skip to content

Commit

Permalink
Fixes required due to UFL changes
Browse files Browse the repository at this point in the history 
*** Makes thwaites depend on g-adopt ***
Any bit of code that required fixing but is duplicated in g-adopt is now
imported from g-adopt, specifically: extend_function_to_3d,
is_continuous, normal_is_continuous. VertexBasedP1DGLimiter is
subclassed to add special treatment of squeezed triangles.

In test_*.py files please do not add any code that's run when imported
as a module, as pytest will automatically import all of these to
discover what files there are, and thus that code will already be run
before pytest itself runs any tests. If you want to add code when the
test file is run as a script, add a `if __name__ == "__main__"` guard.

One of the UFL changes seems to have broken RTCE on tensor product
meshes (firedrakeproject/firedrake#3243). Have
implemented a fix in a fork of FInAT
(https://github.com/thwaitesproject/FInAT) that should now be pulled in
the CI.
  • Loading branch information
@stephankramer
stephankramer committed Nov 20, 2023
1 parent f98392f commit c420600
Show file tree
Hide file tree
Showing 5 changed files with 11 additions and 288 deletions.
3 changes: 2 additions & 1 deletion requirements.txt
View file
Original file line number Diff line number Diff line change
@@ -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
2 changes: 0 additions & 2 deletions tests/adjoint/test_adjoint.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 2 additions & 2 deletions thwaites/interpolate.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
217 changes: 6 additions & 211 deletions thwaites/limiter.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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)
Expand All @@ -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) {
Expand Down Expand Up @@ -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)
73 changes: 1 addition & 72 deletions thwaites/utility.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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))
Expand Down Expand Up @@ -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.
Expand Down

0 comments on commit c420600

Please sign in to comment.