diff --git a/.github/workflows/pyop2.yml b/.github/workflows/pyop2.yml index 36065accf1..9047362d69 100644 --- a/.github/workflows/pyop2.yml +++ b/.github/workflows/pyop2.yml @@ -1,14 +1,10 @@ -name: PyOP2 +name: Test PyOP2 and TSFC -# Trigger the workflow on push or pull request, -# but only for the master branch on: push: branches: - master pull_request: - branches: - - master jobs: test: @@ -88,14 +84,14 @@ jobs: make make install - - name: Checkout PyOP2 + - name: Checkout Firedrake uses: actions/checkout@v4 with: - path: PyOP2 + path: firedrake - name: Install PyOP2 dependencies shell: bash - working-directory: PyOP2 + working-directory: firedrake run: | source ../venv/bin/activate python -m pip install -U pip @@ -103,14 +99,22 @@ jobs: - name: Install PyOP2 shell: bash - working-directory: PyOP2 + working-directory: firedrake run: | source ../venv/bin/activate python -m pip install -v ".[test]" - - name: Run tests + - name: Run TSFC tests + shell: bash + working-directory: firedrake + run: | + source ../venv/bin/activate + pytest --tb=native --timeout=480 --timeout-method=thread -o faulthandler_timeout=540 -v tests/tsfc + timeout-minutes: 10 + + - name: Run PyOP2 tests shell: bash - working-directory: PyOP2 + working-directory: firedrake run: | source ../venv/bin/activate # Running parallel test cases separately works around a bug in pytest-mpi diff --git a/Makefile b/Makefile index 033504d3ef..69c1a6a0a1 100644 --- a/Makefile +++ b/Makefile @@ -24,6 +24,8 @@ lint: @python -m flake8 $(FLAKE8_FORMAT) pyop2 @echo " Linting PyOP2 scripts" @python -m flake8 $(FLAKE8_FORMAT) pyop2/scripts --filename=* + @echo " Linting TSFC" + @python -m flake8 $(FLAKE8_FORMAT) tsfc actionlint: @echo " Pull latest actionlint image" diff --git a/firedrake/scripts/firedrake-zenodo b/firedrake/scripts/firedrake-zenodo index 6194d9df1d..d609a81326 100755 --- a/firedrake/scripts/firedrake-zenodo +++ b/firedrake/scripts/firedrake-zenodo @@ -16,35 +16,35 @@ import datetime # Change this to https://sandbox.zenodo.org/api for testing ZENODO_URL = "https://zenodo.org/api" -# TODO: Remove "petsc4py" once all users have switched to "petsc/src/binding/petsc4py". -# And the same for slepc4py. descriptions = OrderedDict([ ("firedrake", "an automated finite element system"), - ("tsfc", "The Two Stage Form Compiler"), ("ufl", "The Unified Form Language"), - ("FInAT", "a smarter library of finite elements"), ("fiat", "The Finite Element Automated Tabulator"), ("petsc", "Portable, Extensible Toolkit for Scientific Computation"), - ("petsc4py", "The Python interface to PETSc"), ("loopy", "Transformation-Based Generation of High-Performance CPU/GPU Code"), ("slepc", "Scalable Library for Eigenvalue Problem Computations"), - ("slepc4py", "The Python interface to SLEPc")]) - -projects = dict( - [("firedrake", "firedrakeproject"), - ("tsfc", "firedrakeproject"), - ("ufl", "firedrakeproject"), - ("FInAT", "FInAT"), - ("fiat", "firedrakeproject"), - ("petsc", "firedrakeproject"), - ("petsc4py", "firedrakeproject"), - ("loopy", "firedrakeproject"), - ("slepc", "firedrakeproject"), - ("slepc4py", "firedrakeproject")]) + # removed components, left so old Firedrake versions are archivable + ("PyOP2", "Framework for performance-portable parallel computations on unstructured meshes"), + ("tsfc", "The Two Stage Form Compiler"), + ("FInAT", "a smarter library of finite elements"), +]) + +projects = dict([ + ("firedrake", "firedrakeproject"), + ("ufl", "firedrakeproject"), + ("fiat", "firedrakeproject"), + ("petsc", "firedrakeproject"), + ("loopy", "firedrakeproject"), + ("slepc", "firedrakeproject"), + # removed components, left so old Firedrake versions are archivable + ("PyOP2", "OP2"), + ("tsfc", "firedrakeproject"), + ("FInAT", "FInAT"), +]) components = list(descriptions.keys()) -optional_components = ("slepc", "slepc4py", "petsc4py") +optional_components = ("slepc", "PyOP2", "tsfc", "FInAT") parser = ArgumentParser(description="""Create Zenodo DOIs for specific versions of Firedrake components. diff --git a/pyproject.toml b/pyproject.toml index 69af99a483..240772b7c9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,8 +31,6 @@ dependencies = [ "sympy", "fenics-ufl @ git+https://github.com/firedrakeproject/ufl.git", "fenics-fiat @ git+https://github.com/firedrakeproject/fiat.git", - "finat @ git+https://github.com/FInAT/FInAT.git", - "tsfc @ git+https://github.com/firedrakeproject/tsfc.git", "pyadjoint-ad @ git+https://github.com/dolfin-adjoint/pyadjoint.git", "loopy @ git+https://github.com/firedrakeproject/loopy.git@main", ] diff --git a/requirements-git.txt b/requirements-git.txt index c037220ed1..93e6df0f47 100644 --- a/requirements-git.txt +++ b/requirements-git.txt @@ -1,7 +1,5 @@ git+https://github.com/firedrakeproject/ufl.git#egg=fenics-ufl git+https://github.com/firedrakeproject/fiat.git#egg=fenics-fiat -git+https://github.com/FInAT/FInAT.git#egg=finat -git+https://github.com/firedrakeproject/tsfc.git#egg=tsfc git+https://github.com/dolfin-adjoint/pyadjoint.git#egg=pyadjoint-ad git+https://github.com/firedrakeproject/loopy.git@main#egg=loopy git+https://github.com/firedrakeproject/pytest-mpi.git@main#egg=pytest-mpi diff --git a/scripts/firedrake-install b/scripts/firedrake-install index a83b04a451..1637971d28 100755 --- a/scripts/firedrake-install +++ b/scripts/firedrake-install @@ -1941,6 +1941,19 @@ else: if args.rebuild: pip_uninstall("petsc4py") + # Handle archived repositories + archived_repos = [("PyOP2", "firedrake", "PyOP2"), + ("tsfc", "firedrake+fiat", "tsfc"), + ("FInAT", "fiat", "FInAT")] + for src_repo, dest_repo, pip_pkg_name in archived_repos: + archived_path = os.path.join(firedrake_env, "src", src_repo) + if os.path.exists(archived_path): + log.warning("%s has been moved into %s, renaming to %s_old and pip uninstalling.\n" + % (src_repo, dest_repo, src_repo)) + pip_uninstall(pip_pkg_name) + new_path = os.path.join(firedrake_env, "src", "%s_old" % src_repo) + os.rename(archived_path, new_path) + # If there is a petsc package to remove, then we're an old installation which will need to rebuild PETSc. update_petsc = pip_uninstall("petsc") or update_petsc diff --git a/tests/tsfc/test_codegen.py b/tests/tsfc/test_codegen.py new file mode 100644 index 0000000000..8d0bc79655 --- /dev/null +++ b/tests/tsfc/test_codegen.py @@ -0,0 +1,30 @@ +import pytest + +from gem import impero_utils +from gem.gem import Index, Indexed, IndexSum, Product, Variable + + +def test_loop_fusion(): + i = Index() + j = Index() + Ri = Indexed(Variable('R', (6,)), (i,)) + + def make_expression(i, j): + A = Variable('A', (6,)) + s = IndexSum(Indexed(A, (j,)), (j,)) + return Product(Indexed(A, (i,)), s) + + e1 = make_expression(i, j) + e2 = make_expression(i, i) + + def gencode(expr): + impero_c = impero_utils.compile_gem([(Ri, expr)], (i, j)) + return impero_c.tree + + assert len(gencode(e1).children) == len(gencode(e2).children) + + +if __name__ == "__main__": + import os + import sys + pytest.main(args=[os.path.abspath(__file__)] + sys.argv[1:]) diff --git a/tests/tsfc/test_coffee_optimise.py b/tests/tsfc/test_coffee_optimise.py new file mode 100644 index 0000000000..065bb22ee4 --- /dev/null +++ b/tests/tsfc/test_coffee_optimise.py @@ -0,0 +1,78 @@ +import pytest + +from gem.gem import Index, Indexed, Product, Variable, Division, Literal, Sum +from gem.optimise import replace_division +from tsfc.coffee_mode import optimise_expressions + + +def test_replace_div(): + i = Index() + A = Variable('A', ()) + B = Variable('B', (6,)) + Bi = Indexed(B, (i,)) + d = Division(Bi, A) + result, = replace_division([d]) + expected = Product(Bi, Division(Literal(1.0), A)) + + assert result == expected + + +def test_loop_optimise(): + I = 20 + J = K = 10 + i = Index('i', I) + j = Index('j', J) + k = Index('k', K) + + A1 = Variable('a1', (I,)) + A2 = Variable('a2', (I,)) + A3 = Variable('a3', (I,)) + A1i = Indexed(A1, (i,)) + A2i = Indexed(A2, (i,)) + A3i = Indexed(A3, (i,)) + + B = Variable('b', (J,)) + C = Variable('c', (J,)) + Bj = Indexed(B, (j,)) + Cj = Indexed(C, (j,)) + + E = Variable('e', (K,)) + F = Variable('f', (K,)) + G = Variable('g', (K,)) + Ek = Indexed(E, (k,)) + Fk = Indexed(F, (k,)) + Gk = Indexed(G, (k,)) + + Z = Variable('z', ()) + + # Bj*Ek + Bj*Fk => (Ek + Fk)*Bj + expr = Sum(Product(Bj, Ek), Product(Bj, Fk)) + result, = optimise_expressions([expr], (j, k)) + expected = Product(Sum(Ek, Fk), Bj) + assert result == expected + + # Bj*Ek + Bj*Fk + Bj*Gk + Cj*Ek + Cj*Fk => + # (Ek + Fk + Gk)*Bj + (Ek+Fk)*Cj + expr = Sum(Sum(Sum(Sum(Product(Bj, Ek), Product(Bj, Fk)), Product(Bj, Gk)), + Product(Cj, Ek)), Product(Cj, Fk)) + result, = optimise_expressions([expr], (j, k)) + expected = Sum(Product(Sum(Sum(Ek, Fk), Gk), Bj), Product(Sum(Ek, Fk), Cj)) + assert result == expected + + # Z*A1i*Bj*Ek + Z*A2i*Bj*Ek + A3i*Bj*Ek + Z*A1i*Bj*Fk => + # Bj*(Ek*(Z*A1i + Z*A2i) + A3i) + Z*A1i*Fk) + + expr = Sum(Sum(Sum(Product(Z, Product(A1i, Product(Bj, Ek))), + Product(Z, Product(A2i, Product(Bj, Ek)))), + Product(A3i, Product(Bj, Ek))), + Product(Z, Product(A1i, Product(Bj, Fk)))) + result, = optimise_expressions([expr], (j, k)) + expected = Product(Sum(Product(Ek, Sum(Sum(Product(Z, A1i), Product(Z, A2i)), A3i)), + Product(Fk, Product(Z, A1i))), Bj) + assert result == expected + + +if __name__ == "__main__": + import os + import sys + pytest.main(args=[os.path.abspath(__file__)] + sys.argv[1:]) diff --git a/tests/tsfc/test_create_fiat_element.py b/tests/tsfc/test_create_fiat_element.py new file mode 100644 index 0000000000..f8a7d6efc4 --- /dev/null +++ b/tests/tsfc/test_create_fiat_element.py @@ -0,0 +1,150 @@ +import pytest + +import FIAT +from FIAT.discontinuous_lagrange import DiscontinuousLagrange as FIAT_DiscontinuousLagrange + +import ufl +import finat.ufl +from tsfc.finatinterface import create_element as _create_element + + +supported_elements = { + # These all map directly to FIAT elements + "Brezzi-Douglas-Marini": FIAT.BrezziDouglasMarini, + "Brezzi-Douglas-Fortin-Marini": FIAT.BrezziDouglasFortinMarini, + "Lagrange": FIAT.Lagrange, + "Nedelec 1st kind H(curl)": FIAT.Nedelec, + "Nedelec 2nd kind H(curl)": FIAT.NedelecSecondKind, + "Raviart-Thomas": FIAT.RaviartThomas, + "Regge": FIAT.Regge, +} +"""A :class:`.dict` mapping UFL element family names to their +FIAT-equivalent constructors.""" + + +def create_element(ufl_element): + """Create a FIAT element given a UFL element.""" + finat_element = _create_element(ufl_element) + return finat_element.fiat_equivalent + + +@pytest.fixture(params=["BDM", + "BDFM", + "Lagrange", + "N1curl", + "N2curl", + "RT", + "Regge"]) +def triangle_names(request): + return request.param + + +@pytest.fixture +def ufl_element(triangle_names): + return finat.ufl.FiniteElement(triangle_names, ufl.triangle, 2) + + +def test_triangle_basic(ufl_element): + element = create_element(ufl_element) + assert isinstance(element, supported_elements[ufl_element.family()]) + + +@pytest.fixture(params=["CG", "DG", "DG L2"], scope="module") +def tensor_name(request): + return request.param + + +@pytest.fixture(params=[ufl.interval, ufl.triangle, + ufl.quadrilateral], + ids=lambda x: x.cellname(), + scope="module") +def ufl_A(request, tensor_name): + return finat.ufl.FiniteElement(tensor_name, request.param, 1) + + +@pytest.fixture +def ufl_B(tensor_name): + return finat.ufl.FiniteElement(tensor_name, ufl.interval, 1) + + +def test_tensor_prod_simple(ufl_A, ufl_B): + tensor_ufl = finat.ufl.TensorProductElement(ufl_A, ufl_B) + + tensor = create_element(tensor_ufl) + A = create_element(ufl_A) + B = create_element(ufl_B) + + assert isinstance(tensor, FIAT.TensorProductElement) + + assert tensor.A is A + assert tensor.B is B + + +@pytest.mark.parametrize(('family', 'expected_cls'), + [('P', FIAT.GaussLobattoLegendre), + ('DP', FIAT.GaussLegendre), + ('DP L2', FIAT.GaussLegendre)]) +def test_interval_variant_default(family, expected_cls): + ufl_element = finat.ufl.FiniteElement(family, ufl.interval, 3) + assert isinstance(create_element(ufl_element), expected_cls) + + +@pytest.mark.parametrize(('family', 'variant', 'expected_cls'), + [('P', 'equispaced', FIAT.Lagrange), + ('P', 'spectral', FIAT.GaussLobattoLegendre), + ('DP', 'equispaced', FIAT_DiscontinuousLagrange), + ('DP', 'spectral', FIAT.GaussLegendre), + ('DP L2', 'equispaced', FIAT_DiscontinuousLagrange), + ('DP L2', 'spectral', FIAT.GaussLegendre)]) +def test_interval_variant(family, variant, expected_cls): + ufl_element = finat.ufl.FiniteElement(family, ufl.interval, 3, variant=variant) + assert isinstance(create_element(ufl_element), expected_cls) + + +def test_triangle_variant_spectral(): + ufl_element = finat.ufl.FiniteElement('DP', ufl.triangle, 2, variant='spectral') + create_element(ufl_element) + + +def test_triangle_variant_spectral_l2(): + ufl_element = finat.ufl.FiniteElement('DP L2', ufl.triangle, 2, variant='spectral') + create_element(ufl_element) + + +def test_quadrilateral_variant_spectral_q(): + element = create_element(finat.ufl.FiniteElement('Q', ufl.quadrilateral, 3, variant='spectral')) + assert isinstance(element.element.A, FIAT.GaussLobattoLegendre) + assert isinstance(element.element.B, FIAT.GaussLobattoLegendre) + + +def test_quadrilateral_variant_spectral_dq(): + element = create_element(finat.ufl.FiniteElement('DQ', ufl.quadrilateral, 1, variant='spectral')) + assert isinstance(element.element.A, FIAT.GaussLegendre) + assert isinstance(element.element.B, FIAT.GaussLegendre) + + +def test_quadrilateral_variant_spectral_dq_l2(): + element = create_element(finat.ufl.FiniteElement('DQ L2', ufl.quadrilateral, 1, variant='spectral')) + assert isinstance(element.element.A, FIAT.GaussLegendre) + assert isinstance(element.element.B, FIAT.GaussLegendre) + + +def test_quadrilateral_variant_spectral_rtcf(): + element = create_element(finat.ufl.FiniteElement('RTCF', ufl.quadrilateral, 2, variant='spectral')) + assert isinstance(element.element._elements[0].A, FIAT.GaussLobattoLegendre) + assert isinstance(element.element._elements[0].B, FIAT.GaussLegendre) + assert isinstance(element.element._elements[1].A, FIAT.GaussLegendre) + assert isinstance(element.element._elements[1].B, FIAT.GaussLobattoLegendre) + + +def test_cache_hit(ufl_element): + A = create_element(ufl_element) + B = create_element(ufl_element) + + assert A is B + + +if __name__ == "__main__": + import os + import sys + pytest.main(args=[os.path.abspath(__file__)] + sys.argv[1:]) diff --git a/tests/tsfc/test_create_finat_element.py b/tests/tsfc/test_create_finat_element.py new file mode 100644 index 0000000000..c0f34c292c --- /dev/null +++ b/tests/tsfc/test_create_finat_element.py @@ -0,0 +1,138 @@ +import pytest + +import ufl +import finat.ufl +import finat +from tsfc.finatinterface import create_element, supported_elements + + +@pytest.fixture(params=["BDM", + "BDFM", + "Lagrange", + "N1curl", + "N2curl", + "RT", + "Regge"]) +def triangle_names(request): + return request.param + + +@pytest.fixture +def ufl_element(triangle_names): + return finat.ufl.FiniteElement(triangle_names, ufl.triangle, 2) + + +def test_triangle_basic(ufl_element): + element = create_element(ufl_element) + assert isinstance(element, supported_elements[ufl_element.family()]) + + +@pytest.fixture +def ufl_vector_element(triangle_names): + return finat.ufl.VectorElement(triangle_names, ufl.triangle, 2) + + +def test_triangle_vector(ufl_element, ufl_vector_element): + scalar = create_element(ufl_element) + vector = create_element(ufl_vector_element) + + assert isinstance(vector, finat.TensorFiniteElement) + assert scalar == vector.base_element + + +@pytest.fixture(params=["CG", "DG", "DG L2"]) +def tensor_name(request): + return request.param + + +@pytest.fixture(params=[ufl.interval, ufl.triangle, + ufl.quadrilateral], + ids=lambda x: x.cellname()) +def ufl_A(request, tensor_name): + return finat.ufl.FiniteElement(tensor_name, request.param, 1) + + +@pytest.fixture +def ufl_B(tensor_name): + return finat.ufl.FiniteElement(tensor_name, ufl.interval, 1) + + +def test_tensor_prod_simple(ufl_A, ufl_B): + tensor_ufl = finat.ufl.TensorProductElement(ufl_A, ufl_B) + + tensor = create_element(tensor_ufl) + A = create_element(ufl_A) + B = create_element(ufl_B) + + assert isinstance(tensor, finat.TensorProductElement) + + assert tensor.factors == (A, B) + + +@pytest.mark.parametrize(('family', 'expected_cls'), + [('P', finat.GaussLobattoLegendre), + ('DP', finat.GaussLegendre), + ('DP L2', finat.GaussLegendre)]) +def test_interval_variant_default(family, expected_cls): + ufl_element = finat.ufl.FiniteElement(family, ufl.interval, 3) + assert isinstance(create_element(ufl_element), expected_cls) + + +@pytest.mark.parametrize(('family', 'variant', 'expected_cls'), + [('P', 'equispaced', finat.Lagrange), + ('P', 'spectral', finat.GaussLobattoLegendre), + ('DP', 'equispaced', finat.DiscontinuousLagrange), + ('DP', 'spectral', finat.GaussLegendre), + ('DP L2', 'equispaced', finat.DiscontinuousLagrange), + ('DP L2', 'spectral', finat.GaussLegendre)]) +def test_interval_variant(family, variant, expected_cls): + ufl_element = finat.ufl.FiniteElement(family, ufl.interval, 3, variant=variant) + assert isinstance(create_element(ufl_element), expected_cls) + + +def test_triangle_variant_spectral(): + ufl_element = finat.ufl.FiniteElement('DP', ufl.triangle, 2, variant='spectral') + create_element(ufl_element) + + +def test_triangle_variant_spectral_l2(): + ufl_element = finat.ufl.FiniteElement('DP L2', ufl.triangle, 2, variant='spectral') + create_element(ufl_element) + + +def test_quadrilateral_variant_spectral_q(): + element = create_element(finat.ufl.FiniteElement('Q', ufl.quadrilateral, 3, variant='spectral')) + assert isinstance(element.product.factors[0], finat.GaussLobattoLegendre) + assert isinstance(element.product.factors[1], finat.GaussLobattoLegendre) + + +def test_quadrilateral_variant_spectral_dq(): + element = create_element(finat.ufl.FiniteElement('DQ', ufl.quadrilateral, 1, variant='spectral')) + assert isinstance(element.product.factors[0], finat.GaussLegendre) + assert isinstance(element.product.factors[1], finat.GaussLegendre) + + +def test_quadrilateral_variant_spectral_dq_l2(): + element = create_element(finat.ufl.FiniteElement('DQ L2', ufl.quadrilateral, 1, variant='spectral')) + assert isinstance(element.product.factors[0], finat.GaussLegendre) + assert isinstance(element.product.factors[1], finat.GaussLegendre) + + +def test_cache_hit(ufl_element): + A = create_element(ufl_element) + B = create_element(ufl_element) + + assert A is B + + +def test_cache_hit_vector(ufl_vector_element): + A = create_element(ufl_vector_element) + B = create_element(ufl_vector_element) + + assert A is B + + +if __name__ == "__main__": + import os + import sys + pytest.main(args=[os.path.abspath(__file__)] + sys.argv[1:]) diff --git a/tests/tsfc/test_delta_elimination.py b/tests/tsfc/test_delta_elimination.py new file mode 100644 index 0000000000..3ac3f8dc28 --- /dev/null +++ b/tests/tsfc/test_delta_elimination.py @@ -0,0 +1,26 @@ +import pytest + +from gem.gem import Delta, Identity, Index, Indexed, one +from gem.optimise import delta_elimination, remove_componenttensors + + +def test_delta_elimination(): + i = Index() + j = Index() + k = Index() + I = Identity(3) + + sum_indices = (i, j) + factors = [Delta(i, j), Delta(i, k), Indexed(I, (j, k))] + + sum_indices, factors = delta_elimination(sum_indices, factors) + factors = remove_componenttensors(factors) + + assert sum_indices == [] + assert factors == [one, one, Indexed(I, (k, k))] + + +if __name__ == "__main__": + import os + import sys + pytest.main(args=[os.path.abspath(__file__)] + sys.argv[1:]) diff --git a/tests/tsfc/test_dual_evaluation.py b/tests/tsfc/test_dual_evaluation.py new file mode 100644 index 0000000000..b4f6e9770a --- /dev/null +++ b/tests/tsfc/test_dual_evaluation.py @@ -0,0 +1,62 @@ +import pytest +import ufl +import finat.ufl +from tsfc.finatinterface import create_element +from tsfc import compile_expression_dual_evaluation + + +def test_ufl_only_simple(): + mesh = ufl.Mesh(finat.ufl.VectorElement("P", ufl.triangle, 1)) + V = ufl.FunctionSpace(mesh, finat.ufl.FiniteElement("P", ufl.triangle, 2)) + v = ufl.Coefficient(V) + expr = ufl.inner(v, v) + W = V + to_element = create_element(W.ufl_element()) + kernel = compile_expression_dual_evaluation(expr, to_element, W.ufl_element()) + assert kernel.needs_external_coords is False + + +def test_ufl_only_spatialcoordinate(): + mesh = ufl.Mesh(finat.ufl.VectorElement("P", ufl.triangle, 1)) + V = ufl.FunctionSpace(mesh, finat.ufl.FiniteElement("P", ufl.triangle, 2)) + x, y = ufl.SpatialCoordinate(mesh) + expr = x*y - y**2 + x + W = V + to_element = create_element(W.ufl_element()) + kernel = compile_expression_dual_evaluation(expr, to_element, W.ufl_element()) + assert kernel.needs_external_coords is True + + +def test_ufl_only_from_contravariant_piola(): + mesh = ufl.Mesh(finat.ufl.VectorElement("P", ufl.triangle, 1)) + V = ufl.FunctionSpace(mesh, finat.ufl.FiniteElement("RT", ufl.triangle, 1)) + v = ufl.Coefficient(V) + expr = ufl.inner(v, v) + W = ufl.FunctionSpace(mesh, finat.ufl.FiniteElement("P", ufl.triangle, 2)) + to_element = create_element(W.ufl_element()) + kernel = compile_expression_dual_evaluation(expr, to_element, W.ufl_element()) + assert kernel.needs_external_coords is True + + +def test_ufl_only_to_contravariant_piola(): + mesh = ufl.Mesh(finat.ufl.VectorElement("P", ufl.triangle, 1)) + V = ufl.FunctionSpace(mesh, finat.ufl.FiniteElement("P", ufl.triangle, 2)) + v = ufl.Coefficient(V) + expr = ufl.as_vector([v, v]) + W = ufl.FunctionSpace(mesh, finat.ufl.FiniteElement("RT", ufl.triangle, 1)) + to_element = create_element(W.ufl_element()) + kernel = compile_expression_dual_evaluation(expr, to_element, W.ufl_element()) + assert kernel.needs_external_coords is True + + +def test_ufl_only_shape_mismatch(): + mesh = ufl.Mesh(finat.ufl.VectorElement("P", ufl.triangle, 1)) + V = ufl.FunctionSpace(mesh, finat.ufl.FiniteElement("RT", ufl.triangle, 1)) + v = ufl.Coefficient(V) + expr = ufl.inner(v, v) + assert expr.ufl_shape == () + W = V + to_element = create_element(W.ufl_element()) + assert to_element.value_shape == (2,) + with pytest.raises(ValueError): + compile_expression_dual_evaluation(expr, to_element, W.ufl_element()) diff --git a/tests/tsfc/test_estimated_degree.py b/tests/tsfc/test_estimated_degree.py new file mode 100644 index 0000000000..e4842f2ee6 --- /dev/null +++ b/tests/tsfc/test_estimated_degree.py @@ -0,0 +1,35 @@ +import logging + +import pytest + +import ufl +import finat.ufl +from tsfc import compile_form +from tsfc.logging import logger + + +class MockHandler(logging.Handler): + def emit(self, record): + raise RuntimeError() + + +def test_estimated_degree(): + cell = ufl.tetrahedron + mesh = ufl.Mesh(finat.ufl.VectorElement('P', cell, 1)) + V = ufl.FunctionSpace(mesh, finat.ufl.FiniteElement('P', cell, 1)) + f = ufl.Coefficient(V) + u = ufl.TrialFunction(V) + v = ufl.TestFunction(V) + a = u * v * ufl.tanh(ufl.sqrt(ufl.sinh(f) / ufl.sin(f**f))) * ufl.dx + + handler = MockHandler() + logger.addHandler(handler) + with pytest.raises(RuntimeError): + compile_form(a) + logger.removeHandler(handler) + + +if __name__ == "__main__": + import os + import sys + pytest.main(args=[os.path.abspath(__file__)] + sys.argv[1:]) diff --git a/tests/tsfc/test_firedrake_972.py b/tests/tsfc/test_firedrake_972.py new file mode 100644 index 0000000000..e308c7986d --- /dev/null +++ b/tests/tsfc/test_firedrake_972.py @@ -0,0 +1,44 @@ +import numpy +import pytest + +from ufl import (Mesh, FunctionSpace, + Coefficient, TestFunction, interval, indices, dx) +from finat.ufl import VectorElement, TensorElement +from ufl.classes import IndexSum, Product, MultiIndex + +from tsfc import compile_form + + +def count_flops(n): + mesh = Mesh(VectorElement('CG', interval, 1)) + tfs = FunctionSpace(mesh, TensorElement('DG', interval, 1, shape=(n, n))) + vfs = FunctionSpace(mesh, VectorElement('DG', interval, 1, dim=n)) + + ensemble_f = Coefficient(vfs) + ensemble2_f = Coefficient(vfs) + phi = TestFunction(tfs) + + i, j = indices(2) + nc = 42 # magic number + L = ((IndexSum(IndexSum(Product(nc * phi[i, j], Product(ensemble_f[i], ensemble_f[i])), + MultiIndex((i,))), MultiIndex((j,))) * dx) + + (IndexSum(IndexSum(Product(nc * phi[i, j], Product(ensemble2_f[j], ensemble2_f[j])), + MultiIndex((i,))), MultiIndex((j,))) * dx) + - (IndexSum(IndexSum(2 * nc * Product(phi[i, j], Product(ensemble_f[i], ensemble2_f[j])), + MultiIndex((i,))), MultiIndex((j,))) * dx)) + + kernel, = compile_form(L, parameters=dict(mode='spectral')) + return kernel.flop_count + + +def test_convergence(): + ns = [10, 20, 40, 80, 100] + flops = [count_flops(n) for n in ns] + rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(ns)) + assert (rates < 2).all() # only quadratic operation count, not more + + +if __name__ == "__main__": + import os + import sys + pytest.main(args=[os.path.abspath(__file__)] + sys.argv[1:]) diff --git a/tests/tsfc/test_flexibly_indexed.py b/tests/tsfc/test_flexibly_indexed.py new file mode 100644 index 0000000000..551e4b042b --- /dev/null +++ b/tests/tsfc/test_flexibly_indexed.py @@ -0,0 +1,103 @@ +from itertools import product + +import numpy +import pytest + +import gem +import tsfc +import tsfc.loopy + + +parameters = tsfc.loopy.LoopyContext() +parameters.names = {} + + +def convert(expression, multiindex): + assert not expression.free_indices + element = gem.Indexed(expression, multiindex) + element, = gem.optimise.remove_componenttensors((element,)) + subscript = tsfc.loopy.expression(element, parameters) + # Convert a pymbolic subscript expression to a rank tuple. For example + # the subscript: + # + # Subscript(Variable('A'), (Sum((3,)), Sum((5,)))) + # + # will yield a rank of (3, 5). + return sum((idx.children for idx in subscript.index), start=()) + + +@pytest.fixture(scope='module') +def vector(): + return gem.Variable('u', (12,)) + + +@pytest.fixture(scope='module') +def matrix(): + return gem.Variable('A', (10, 12)) + + +def test_reshape(vector): + expression = gem.reshape(vector, (3, 4)) + assert expression.shape == (3, 4) + + actual = [convert(expression, multiindex) + for multiindex in numpy.ndindex(expression.shape)] + + assert [(i,) for i in range(12)] == actual + + +def test_view(matrix): + expression = gem.view(matrix, slice(3, 8), slice(5, 12)) + assert expression.shape == (5, 7) + + actual = [convert(expression, multiindex) + for multiindex in numpy.ndindex(expression.shape)] + + assert list(product(range(3, 8), range(5, 12))) == actual + + +def test_view_view(matrix): + expression = gem.view(gem.view(matrix, slice(3, 8), slice(5, 12)), + slice(4), slice(3, 6)) + assert expression.shape == (4, 3) + + actual = [convert(expression, multiindex) + for multiindex in numpy.ndindex(expression.shape)] + + assert list(product(range(3, 7), range(8, 11))) == actual + + +def test_view_reshape(vector): + expression = gem.view(gem.reshape(vector, (3, 4)), slice(2), slice(1, 3)) + assert expression.shape == (2, 2) + + actual = [convert(expression, multiindex) + for multiindex in numpy.ndindex(expression.shape)] + + assert [(1,), (2,), (5,), (6,)] == actual + + +def test_reshape_shape(vector): + expression = gem.reshape(gem.view(vector, slice(5, 11)), (3, 2)) + assert expression.shape == (3, 2) + + actual = [convert(expression, multiindex) + for multiindex in numpy.ndindex(expression.shape)] + + assert [(i,) for i in range(5, 11)] == actual + + +def test_reshape_reshape(vector): + expression = gem.reshape(gem.reshape(vector, (4, 3)), (2, 2), (3,)) + assert expression.shape == (2, 2, 3) + + actual = [convert(expression, multiindex) + for multiindex in numpy.ndindex(expression.shape)] + + assert [(i,) for i in range(12)] == actual + + +if __name__ == "__main__": + import os + import sys + pytest.main(args=[os.path.abspath(__file__)] + sys.argv[1:]) diff --git a/tests/tsfc/test_flop_count.py b/tests/tsfc/test_flop_count.py new file mode 100644 index 0000000000..1925fbaf28 --- /dev/null +++ b/tests/tsfc/test_flop_count.py @@ -0,0 +1,64 @@ +import pytest +import gem.gem as gem +from gem.flop_count import count_flops +from gem.impero_utils import preprocess_gem +from gem.impero_utils import compile_gem + + +def test_count_flops(expression): + expr, expected = expression + flops = count_flops(expr) + assert flops == expected + + +@pytest.fixture(params=("expr1", "expr2", "expr3", "expr4")) +def expression(request): + if request.param == "expr1": + expr = gem.Sum(gem.Product(gem.Variable("a", ()), gem.Literal(2)), + gem.Division(gem.Literal(3), gem.Variable("b", ()))) + C = gem.Variable("C", (1,)) + i, = gem.indices(1) + Ci = C[i] + expr, = preprocess_gem([expr]) + assignments = [(Ci, expr)] + expr = compile_gem(assignments, (i,)) + # C += a*2 + 3/b + expected = 1 + 3 + elif request.param == "expr2": + expr = gem.Comparison(">=", gem.MaxValue(gem.Literal(1), gem.Literal(2)), + gem.MinValue(gem.Literal(3), gem.Literal(1))) + C = gem.Variable("C", (1,)) + i, = gem.indices(1) + Ci = C[i] + expr, = preprocess_gem([expr]) + assignments = [(Ci, expr)] + expr = compile_gem(assignments, (i,)) + # C += max(1, 2) >= min(3, 1) + expected = 1 + 3 + elif request.param == "expr3": + expr = gem.Solve(gem.Identity(3), gem.Inverse(gem.Identity(3))) + C = gem.Variable("C", (3, 3)) + i, j = gem.indices(2) + Cij = C[i, j] + expr, = preprocess_gem([expr[i, j]]) + assignments = [(Cij, expr)] + expr = compile_gem(assignments, (i, j)) + # C += solve(Id(3x3), Id(3x3)^{-1}) + expected = 9 + 18 + 54 + 54 + elif request.param == "expr4": + A = gem.Variable("A", (10, 15)) + B = gem.Variable("B", (8, 10)) + i, j, k = gem.indices(3) + Aij = A[i, j] + Bki = B[k, i] + Cjk = gem.IndexSum(Aij * Bki, (i,)) + expr = Cjk + expr, = preprocess_gem([expr]) + assignments = [(gem.Variable("C", (15, 8))[j, k], expr)] + expr = compile_gem(assignments, (j, k)) + # Cjk += \sum_i Aij * Bki + expected = 2 * 10 * 8 * 15 + + else: + raise ValueError("Unexpected expression") + return expr, expected diff --git a/tests/tsfc/test_gem_failure.py b/tests/tsfc/test_gem_failure.py new file mode 100644 index 0000000000..0a79a39ab8 --- /dev/null +++ b/tests/tsfc/test_gem_failure.py @@ -0,0 +1,45 @@ +from ufl import (triangle, tetrahedron, FunctionSpace, Mesh, + TrialFunction, TestFunction, inner, grad, dx, dS) +from finat.ufl import FiniteElement, VectorElement +from tsfc import compile_form +from FIAT.hdiv_trace import TraceError +import pytest + + +@pytest.mark.parametrize('cell', [triangle, tetrahedron]) +@pytest.mark.parametrize('degree', range(3)) +def test_cell_error(cell, degree): + """Test that tabulating the trace element deliberatly on the + cell triggers `gem.Failure` to raise the TraceError exception. + """ + trace_element = FiniteElement("HDiv Trace", cell, degree) + domain = Mesh(VectorElement("Lagrange", cell, 1)) + space = FunctionSpace(domain, trace_element) + lambdar = TrialFunction(space) + gammar = TestFunction(space) + + with pytest.raises(TraceError): + compile_form(lambdar * gammar * dx) + + +@pytest.mark.parametrize('cell', [triangle, tetrahedron]) +@pytest.mark.parametrize('degree', range(3)) +def test_gradient_error(cell, degree): + """Test that tabulating gradient evaluations of the trace + element triggers `gem.Failure` to raise the TraceError + exception. + """ + trace_element = FiniteElement("HDiv Trace", cell, degree) + domain = Mesh(VectorElement("Lagrange", cell, 1)) + space = FunctionSpace(domain, trace_element) + lambdar = TrialFunction(space) + gammar = TestFunction(space) + + with pytest.raises(TraceError): + compile_form(inner(grad(lambdar('+')), grad(gammar('+'))) * dS) + + +if __name__ == "__main__": + import os + import sys + pytest.main(args=[os.path.abspath(__file__)] + sys.argv[1:]) diff --git a/tests/tsfc/test_geometry.py b/tests/tsfc/test_geometry.py new file mode 100644 index 0000000000..458f1b5657 --- /dev/null +++ b/tests/tsfc/test_geometry.py @@ -0,0 +1,76 @@ +import pytest +import numpy as np + +from FIAT.reference_element import UFCInterval, UFCTriangle, UFCTetrahedron +from FIAT.reference_element import UFCQuadrilateral, TensorProductCell + +from tsfc.fem import make_cell_facet_jacobian + +interval = UFCInterval() +triangle = UFCTriangle() +quadrilateral = UFCQuadrilateral() +tetrahedron = UFCTetrahedron() +interval_x_interval = TensorProductCell(interval, interval) +triangle_x_interval = TensorProductCell(triangle, interval) +quadrilateral_x_interval = TensorProductCell(quadrilateral, interval) + + +@pytest.mark.parametrize(('cell', 'cell_facet_jacobian'), + [(interval, [[], + []]), + (triangle, [[-1, 1], + [0, 1], + [1, 0]]), + (quadrilateral, [[0, 1], + [0, 1], + [1, 0], + [1, 0]]), + (tetrahedron, [[-1, -1, 1, 0, 0, 1], + [0, 0, 1, 0, 0, 1], + [1, 0, 0, 0, 0, 1], + [1, 0, 0, 1, 0, 0]])]) +def test_cell_facet_jacobian(cell, cell_facet_jacobian): + facet_dim = cell.get_spatial_dimension() - 1 + for facet_number in range(len(cell.get_topology()[facet_dim])): + actual = make_cell_facet_jacobian(cell, facet_dim, facet_number) + expected = np.reshape(cell_facet_jacobian[facet_number], actual.shape) + assert np.allclose(expected, actual) + + +@pytest.mark.parametrize(('cell', 'cell_facet_jacobian'), + [(interval_x_interval, [1, 0]), + (triangle_x_interval, [1, 0, 0, 1, 0, 0]), + (quadrilateral_x_interval, [[1, 0, 0, 1, 0, 0]])]) +def test_cell_facet_jacobian_horiz(cell, cell_facet_jacobian): + dim = cell.get_spatial_dimension() + + actual = make_cell_facet_jacobian(cell, (dim - 1, 0), 0) # bottom facet + assert np.allclose(np.reshape(cell_facet_jacobian, actual.shape), actual) + + actual = make_cell_facet_jacobian(cell, (dim - 1, 0), 1) # top facet + assert np.allclose(np.reshape(cell_facet_jacobian, actual.shape), actual) + + +@pytest.mark.parametrize(('cell', 'cell_facet_jacobian'), + [(interval_x_interval, [[0, 1], + [0, 1]]), + (triangle_x_interval, [[-1, 0, 1, 0, 0, 1], + [0, 0, 1, 0, 0, 1], + [1, 0, 0, 0, 0, 1]]), + (quadrilateral_x_interval, [[0, 0, 1, 0, 0, 1], + [0, 0, 1, 0, 0, 1], + [1, 0, 0, 0, 0, 1], + [1, 0, 0, 0, 0, 1]])]) +def test_cell_facet_jacobian_vert(cell, cell_facet_jacobian): + dim = cell.get_spatial_dimension() + vert_dim = (dim - 2, 1) + for facet_number in range(len(cell.get_topology()[vert_dim])): + actual = make_cell_facet_jacobian(cell, vert_dim, facet_number) + expected = np.reshape(cell_facet_jacobian[facet_number], actual.shape) + assert np.allclose(expected, actual) + + +if __name__ == "__main__": + import os + import sys + pytest.main(args=[os.path.abspath(__file__)] + sys.argv[1:]) diff --git a/tests/tsfc/test_idempotency.py b/tests/tsfc/test_idempotency.py new file mode 100644 index 0000000000..cf554f4641 --- /dev/null +++ b/tests/tsfc/test_idempotency.py @@ -0,0 +1,72 @@ +import ufl +import finat.ufl +from tsfc import compile_form +import loopy +import pytest + + +@pytest.fixture(params=[ufl.interval, + ufl.triangle, + ufl.quadrilateral, + ufl.tetrahedron], + ids=lambda x: x.cellname()) +def cell(request): + return request.param + + +@pytest.fixture(params=[1, 2], + ids=lambda x: "P%d-coords" % x) +def coord_degree(request): + return request.param + + +@pytest.fixture +def mesh(cell, coord_degree): + c = finat.ufl.VectorElement("CG", cell, coord_degree) + return ufl.Mesh(c) + + +@pytest.fixture(params=[finat.ufl.FiniteElement, + finat.ufl.VectorElement, + finat.ufl.TensorElement], + ids=["FE", "VE", "TE"]) +def V(request, mesh): + return ufl.FunctionSpace(mesh, request.param("CG", mesh.ufl_cell(), 2)) + + +@pytest.fixture(params=["cell", "ext_facet", "int_facet"]) +def itype(request): + return request.param + + +@pytest.fixture(params=["functional", "1-form", "2-form"]) +def form(V, itype, request): + if request.param == "functional": + u = ufl.Coefficient(V) + v = ufl.Coefficient(V) + elif request.param == "1-form": + u = ufl.Coefficient(V) + v = ufl.TestFunction(V) + elif request.param == "2-form": + u = ufl.TrialFunction(V) + v = ufl.TestFunction(V) + + if itype == "cell": + return ufl.inner(u, v)*ufl.dx + elif itype == "ext_facet": + return ufl.inner(u, v)*ufl.ds + elif itype == "int_facet": + return ufl.inner(u('+'), v('-'))*ufl.dS + + +def test_idempotency(form): + k1 = compile_form(form)[0] + k2 = compile_form(form)[0] + + assert loopy.generate_code_v2(k1.ast).device_code() == loopy.generate_code_v2(k2.ast).device_code() + + +if __name__ == "__main__": + import os + import sys + pytest.main(args=[os.path.abspath(__file__)] + sys.argv[1:]) diff --git a/tests/tsfc/test_impero_loopy_flop_counts.py b/tests/tsfc/test_impero_loopy_flop_counts.py new file mode 100644 index 0000000000..76bdad96e3 --- /dev/null +++ b/tests/tsfc/test_impero_loopy_flop_counts.py @@ -0,0 +1,66 @@ +""" +Tests impero flop counts against loopy. +""" +import pytest +import numpy +import loopy +from tsfc import compile_form +from ufl import (FunctionSpace, Mesh, TestFunction, + TrialFunction, dx, grad, inner, + interval, triangle, quadrilateral, + TensorProductCell) +from finat.ufl import FiniteElement, VectorElement +from tsfc.parameters import target + + +def count_loopy_flops(kernel): + name = kernel.name + program = kernel.ast + program = program.with_kernel( + program[name].copy( + target=target, + silenced_warnings=["insn_count_subgroups_upper_bound", + "get_x_map_guessing_subgroup_size"]) + ) + op_map = loopy.get_op_map(program + .with_entrypoints(kernel.name), + subgroup_size=1) + return op_map.filter_by(name=['add', 'sub', 'mul', 'div', + 'func:abs'], + dtype=[float]).eval_and_sum({}) + + +@pytest.fixture(params=[interval, triangle, quadrilateral, + TensorProductCell(triangle, interval)], + ids=lambda cell: cell.cellname()) +def cell(request): + return request.param + + +@pytest.fixture(params=[{"mode": "vanilla"}, + {"mode": "spectral"}], + ids=["vanilla", "spectral"]) +def parameters(request): + return request.param + + +def test_flop_count(cell, parameters): + mesh = Mesh(VectorElement("P", cell, 1)) + loopy_flops = [] + new_flops = [] + for k in range(1, 5): + V = FunctionSpace(mesh, FiniteElement("P", cell, k)) + u = TrialFunction(V) + v = TestFunction(V) + a = inner(u, v)*dx + inner(grad(u), grad(v))*dx + kernel, = compile_form(a, prefix="form", + parameters=parameters) + # Record new flops here, and compare asymptotics and + # approximate order of magnitude. + new_flops.append(kernel.flop_count) + loopy_flops.append(count_loopy_flops(kernel)) + + new_flops = numpy.asarray(new_flops) + loopy_flops = numpy.asarray(loopy_flops) + + assert all(new_flops == loopy_flops) diff --git a/tests/tsfc/test_interpolation_factorisation.py b/tests/tsfc/test_interpolation_factorisation.py new file mode 100644 index 0000000000..b3d4e3288b --- /dev/null +++ b/tests/tsfc/test_interpolation_factorisation.py @@ -0,0 +1,64 @@ +from functools import partial +import numpy +import pytest + +from ufl import (Mesh, FunctionSpace, Coefficient, + interval, quadrilateral, hexahedron) +from finat.ufl import FiniteElement, VectorElement, TensorElement + +from tsfc import compile_expression_dual_evaluation +from tsfc.finatinterface import create_element + + +@pytest.fixture(params=[interval, quadrilateral, hexahedron], + ids=lambda x: x.cellname()) +def mesh(request): + return Mesh(VectorElement("P", request.param, 1)) + + +@pytest.fixture(params=[FiniteElement, VectorElement, TensorElement], + ids=lambda x: x.__name__) +def element(request, mesh): + if mesh.ufl_cell() == interval: + family = "DP" + else: + family = "DQ" + return partial(request.param, family, mesh.ufl_cell()) + + +def flop_count(mesh, source, target): + Vtarget = FunctionSpace(mesh, target) + Vsource = FunctionSpace(mesh, source) + to_element = create_element(Vtarget.ufl_element()) + expr = Coefficient(Vsource) + kernel = compile_expression_dual_evaluation(expr, to_element, Vtarget.ufl_element()) + return kernel.flop_count + + +def test_sum_factorisation(mesh, element): + # Interpolation between sum factorisable elements should cost + # O(p^{d+1}) + degrees = numpy.asarray([2**n - 1 for n in range(2, 9)]) + flops = [] + for lo, hi in zip(degrees - 1, degrees): + flops.append(flop_count(mesh, element(int(lo)), element(int(hi)))) + flops = numpy.asarray(flops) + rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees)) + assert (rates < (mesh.topological_dimension()+1)).all() + + +def test_sum_factorisation_scalar_tensor(mesh, element): + # Interpolation into tensor elements should cost value_shape + # more than the equivalent scalar element. + degree = 2**7 - 1 + source = element(degree - 1) + target = element(degree) + tensor_flops = flop_count(mesh, source, target) + expect = FunctionSpace(mesh, target).value_size + if isinstance(target, FiniteElement): + scalar_flops = tensor_flops + else: + target = target.sub_elements[0] + source = source.sub_elements[0] + scalar_flops = flop_count(mesh, source, target) + assert numpy.allclose(tensor_flops / scalar_flops, expect, rtol=1e-2) diff --git a/tests/tsfc/test_pickle_gem.py b/tests/tsfc/test_pickle_gem.py new file mode 100644 index 0000000000..beb101f912 --- /dev/null +++ b/tests/tsfc/test_pickle_gem.py @@ -0,0 +1,31 @@ +import pickle +import gem +import numpy +import pytest + + +@pytest.mark.parametrize('protocol', range(3)) +def test_pickle_gem(protocol): + f = gem.VariableIndex(gem.Indexed(gem.Variable('facet', (2,), dtype=gem.uint_type), (1,))) + q = gem.Index() + r = gem.Index() + _1 = gem.Indexed(gem.Literal(numpy.random.rand(3, 6, 8)), (f, q, r)) + _2 = gem.Indexed(gem.view(gem.Variable('w', (None, None)), slice(8), slice(1)), (r, 0)) + expr = gem.ComponentTensor(gem.IndexSum(gem.Product(_1, _2), (r,)), (q,)) + + unpickled = pickle.loads(pickle.dumps(expr, protocol)) + assert repr(expr) == repr(unpickled) + + +@pytest.mark.parametrize('protocol', range(3)) +def test_listtensor(protocol): + expr = gem.ListTensor([gem.Variable('x', ()), gem.Zero()]) + + unpickled = pickle.loads(pickle.dumps(expr, protocol)) + assert expr == unpickled + + +if __name__ == "__main__": + import os + import sys + pytest.main(args=[os.path.abspath(__file__)] + sys.argv[1:]) diff --git a/tests/tsfc/test_refactorise.py b/tests/tsfc/test_refactorise.py new file mode 100644 index 0000000000..8c8251dfa9 --- /dev/null +++ b/tests/tsfc/test_refactorise.py @@ -0,0 +1,66 @@ +from functools import partial + +import pytest + +import gem +from gem.node import traversal +from gem.refactorise import ATOMIC, COMPOUND, OTHER, Monomial, collect_monomials + + +def test_refactorise(): + f = gem.Variable('f', (3,)) + u = gem.Variable('u', (3,)) + v = gem.Variable('v', ()) + + i = gem.Index() + f_i = gem.Indexed(f, (i,)) + u_i = gem.Indexed(u, (i,)) + + def classify(atomics_set, expression): + if expression in atomics_set: + return ATOMIC + + for node in traversal([expression]): + if node in atomics_set: + return COMPOUND + + return OTHER + classifier = partial(classify, {u_i, v}) + + # \sum_i 5*(2*u_i + -1*v)*(u_i + v*f) + expr = gem.IndexSum( + gem.Product( + gem.Literal(5), + gem.Product( + gem.Sum(gem.Product(gem.Literal(2), u_i), + gem.Product(gem.Literal(-1), v)), + gem.Sum(u_i, gem.Product(v, f_i)) + ) + ), + (i,) + ) + + expected = [ + Monomial((i,), + (u_i, u_i), + gem.Literal(10)), + Monomial((i,), + (u_i, v), + gem.Product(gem.Literal(5), + gem.Sum(gem.Product(f_i, gem.Literal(2)), + gem.Literal(-1)))), + Monomial((), + (v, v), + gem.Product(gem.Literal(5), + gem.IndexSum(gem.Product(f_i, gem.Literal(-1)), + (i,)))), + ] + + actual, = collect_monomials([expr], classifier) + assert expected == list(actual) + + +if __name__ == "__main__": + import os + import sys + pytest.main(args=[os.path.abspath(__file__)] + sys.argv[1:]) diff --git a/tests/tsfc/test_simplification.py b/tests/tsfc/test_simplification.py new file mode 100644 index 0000000000..e5fe3d66e5 --- /dev/null +++ b/tests/tsfc/test_simplification.py @@ -0,0 +1,31 @@ +import pytest + +from gem.gem import Variable, Zero, Conditional, \ + LogicalAnd, Index, Indexed, Product + + +def test_conditional_simplification(): + a = Variable("A", ()) + b = Variable("B", ()) + + expr = Conditional(LogicalAnd(b, a), a, a) + + assert expr == a + + +def test_conditional_zero_folding(): + b = Variable("B", ()) + a = Variable("A", (3, )) + i = Index() + expr = Conditional(LogicalAnd(b, b), + Product(Indexed(a, (i, )), + Zero()), + Zero()) + + assert expr == Zero() + + +if __name__ == "__main__": + import os + import sys + pytest.main(args=[os.path.abspath(__file__)] + sys.argv[1:]) diff --git a/tests/tsfc/test_sum_factorisation.py b/tests/tsfc/test_sum_factorisation.py new file mode 100644 index 0000000000..3e785c5b26 --- /dev/null +++ b/tests/tsfc/test_sum_factorisation.py @@ -0,0 +1,174 @@ +import numpy +import pytest + +from ufl import (Mesh, FunctionSpace, TestFunction, TrialFunction, + TensorProductCell, dx, action, interval, triangle, + quadrilateral, curl, dot, div, grad) +from finat.ufl import (FiniteElement, VectorElement, EnrichedElement, + TensorProductElement, HCurlElement, HDivElement) + +from tsfc import compile_form + + +def helmholtz(cell, degree): + m = Mesh(VectorElement('CG', cell, 1)) + V = FunctionSpace(m, FiniteElement('CG', cell, degree)) + u = TrialFunction(V) + v = TestFunction(V) + return (u*v + dot(grad(u), grad(v)))*dx + + +def split_mixed_poisson(cell, degree): + m = Mesh(VectorElement('CG', cell, 1)) + if cell.cellname() in ['interval * interval', 'quadrilateral']: + hdiv_element = FiniteElement('RTCF', cell, degree) + elif cell.cellname() == 'triangle * interval': + U0 = FiniteElement('RT', triangle, degree) + U1 = FiniteElement('DG', triangle, degree - 1) + V0 = FiniteElement('CG', interval, degree) + V1 = FiniteElement('DG', interval, degree - 1) + Wa = HDivElement(TensorProductElement(U0, V1)) + Wb = HDivElement(TensorProductElement(U1, V0)) + hdiv_element = EnrichedElement(Wa, Wb) + elif cell.cellname() == 'quadrilateral * interval': + hdiv_element = FiniteElement('NCF', cell, degree) + RT = FunctionSpace(m, hdiv_element) + DG = FunctionSpace(m, FiniteElement('DQ', cell, degree - 1)) + sigma = TrialFunction(RT) + u = TrialFunction(DG) + tau = TestFunction(RT) + v = TestFunction(DG) + return [dot(sigma, tau) * dx, div(tau) * u * dx, div(sigma) * v * dx] + + +def split_vector_laplace(cell, degree): + m = Mesh(VectorElement('CG', cell, 1)) + if cell.cellname() in ['interval * interval', 'quadrilateral']: + hcurl_element = FiniteElement('RTCE', cell, degree) + elif cell.cellname() == 'triangle * interval': + U0 = FiniteElement('RT', triangle, degree) + U1 = FiniteElement('CG', triangle, degree) + V0 = FiniteElement('CG', interval, degree) + V1 = FiniteElement('DG', interval, degree - 1) + Wa = HCurlElement(TensorProductElement(U0, V0)) + Wb = HCurlElement(TensorProductElement(U1, V1)) + hcurl_element = EnrichedElement(Wa, Wb) + elif cell.cellname() == 'quadrilateral * interval': + hcurl_element = FiniteElement('NCE', cell, degree) + RT = FunctionSpace(m, hcurl_element) + CG = FunctionSpace(m, FiniteElement('Q', cell, degree)) + sigma = TrialFunction(CG) + u = TrialFunction(RT) + tau = TestFunction(CG) + v = TestFunction(RT) + return [dot(u, grad(tau))*dx, dot(grad(sigma), v)*dx, dot(curl(u), curl(v))*dx] + + +def count_flops(form): + kernel, = compile_form(form, parameters=dict(mode='spectral')) + flops = kernel.flop_count + return flops + + +@pytest.mark.parametrize(('cell', 'order'), + [(quadrilateral, 5), + (TensorProductCell(interval, interval), 5), + (TensorProductCell(triangle, interval), 7), + (TensorProductCell(quadrilateral, interval), 7)]) +def test_lhs(cell, order): + degrees = list(range(3, 8)) + if cell == TensorProductCell(triangle, interval): + degrees = list(range(3, 6)) + flops = [count_flops(helmholtz(cell, degree)) + for degree in degrees] + rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees)) + assert (rates < order).all() + + +@pytest.mark.parametrize(('cell', 'order'), + [(quadrilateral, 3), + (TensorProductCell(interval, interval), 3), + (TensorProductCell(triangle, interval), 5), + (TensorProductCell(quadrilateral, interval), 4)]) +def test_rhs(cell, order): + degrees = list(range(3, 8)) + if cell == TensorProductCell(triangle, interval): + degrees = list(range(3, 6)) + flops = [count_flops(action(helmholtz(cell, degree))) + for degree in degrees] + rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees)) + assert (rates < order).all() + + +@pytest.mark.parametrize(('cell', 'order'), + [(quadrilateral, 5), + (TensorProductCell(interval, interval), 5), + (TensorProductCell(triangle, interval), 7), + (TensorProductCell(quadrilateral, interval), 7) + ]) +def test_mixed_poisson(cell, order): + degrees = numpy.arange(3, 8) + if cell == TensorProductCell(triangle, interval): + degrees = numpy.arange(3, 6) + flops = [[count_flops(form) + for form in split_mixed_poisson(cell, int(degree))] + for degree in degrees] + rates = numpy.diff(numpy.log(flops).T) / numpy.diff(numpy.log(degrees)) + assert (rates < order).all() + + +@pytest.mark.parametrize(('cell', 'order'), + [(quadrilateral, 3), + (TensorProductCell(interval, interval), 3), + (TensorProductCell(triangle, interval), 5), + (TensorProductCell(quadrilateral, interval), 4) + ]) +def test_mixed_poisson_action(cell, order): + degrees = numpy.arange(3, 8) + if cell == TensorProductCell(triangle, interval): + degrees = numpy.arange(3, 6) + flops = [[count_flops(action(form)) + for form in split_mixed_poisson(cell, int(degree))] + for degree in degrees] + rates = numpy.diff(numpy.log(flops).T) / numpy.diff(numpy.log(degrees)) + assert (rates < order).all() + + +@pytest.mark.parametrize(('cell', 'order'), + [(quadrilateral, 5), + (TensorProductCell(interval, interval), 5), + (TensorProductCell(triangle, interval), 7), + (TensorProductCell(quadrilateral, interval), 7) + ]) +def test_vector_laplace(cell, order): + degrees = numpy.arange(3, 8) + if cell == TensorProductCell(triangle, interval): + degrees = numpy.arange(3, 6) + flops = [[count_flops(form) + for form in split_vector_laplace(cell, int(degree))] + for degree in degrees] + rates = numpy.diff(numpy.log(flops).T) / numpy.diff(numpy.log(degrees)) + assert (rates < order).all() + + +@pytest.mark.parametrize(('cell', 'order'), + [(quadrilateral, 3), + (TensorProductCell(interval, interval), 3), + (TensorProductCell(triangle, interval), 5), + (TensorProductCell(quadrilateral, interval), 4) + ]) +def test_vector_laplace_action(cell, order): + degrees = numpy.arange(3, 8) + if cell == TensorProductCell(triangle, interval): + degrees = numpy.arange(3, 6) + flops = [[count_flops(action(form)) + for form in split_vector_laplace(cell, int(degree))] + for degree in degrees] + rates = numpy.diff(numpy.log(flops).T) / numpy.diff(numpy.log(degrees)) + assert (rates < order).all() + + +if __name__ == "__main__": + import os + import sys + pytest.main(args=[os.path.abspath(__file__)] + sys.argv[1:]) diff --git a/tests/tsfc/test_syntax_sugar.py b/tests/tsfc/test_syntax_sugar.py new file mode 100644 index 0000000000..56bbc4f4ac --- /dev/null +++ b/tests/tsfc/test_syntax_sugar.py @@ -0,0 +1,41 @@ +import pytest +import gem + + +def test_expressions(): + x = gem.Variable("x", (3, 4)) + y = gem.Variable("y", (4, )) + i, j = gem.indices(2) + + xij = x[i, j] + yj = y[j] + + assert xij == gem.Indexed(x, (i, j)) + assert yj == gem.Indexed(y, (j, )) + + assert xij + yj == gem.Sum(xij, yj) + assert xij * yj == gem.Product(xij, yj) + assert xij - yj == gem.Sum(xij, gem.Product(gem.Literal(-1), yj)) + assert xij / yj == gem.Division(xij, yj) + + assert xij + 1 == gem.Sum(xij, gem.Literal(1)) + assert 1 + xij == gem.Sum(gem.Literal(1), xij) + + assert (xij + y).shape == (4, ) + + assert (x @ y).shape == (3, ) + + assert x.T.shape == (4, 3) + + with pytest.raises(ValueError): + xij.T @ y + + with pytest.raises(ValueError): + xij + "foo" + + +def test_as_gem(): + with pytest.raises(ValueError): + gem.as_gem([1, 2]) + + assert gem.as_gem(1) == gem.Literal(1) diff --git a/tests/tsfc/test_tensor.py b/tests/tsfc/test_tensor.py new file mode 100644 index 0000000000..9d09a467fb --- /dev/null +++ b/tests/tsfc/test_tensor.py @@ -0,0 +1,116 @@ +import numpy +import pytest + +from ufl import (Mesh, FunctionSpace, + Coefficient, TestFunction, TrialFunction, dx, div, + inner, interval, triangle, tetrahedron, dot, grad) +from finat.ufl import FiniteElement, VectorElement + +from tsfc import compile_form + + +def mass(cell, degree): + m = Mesh(VectorElement('CG', cell, 1)) + V = FunctionSpace(m, FiniteElement('CG', cell, degree)) + u = TrialFunction(V) + v = TestFunction(V) + return u*v*dx + + +def poisson(cell, degree): + m = Mesh(VectorElement('CG', cell, 1)) + V = FunctionSpace(m, FiniteElement('CG', cell, degree)) + u = TrialFunction(V) + v = TestFunction(V) + return dot(grad(u), grad(v))*dx + + +def helmholtz(cell, degree): + m = Mesh(VectorElement('CG', cell, 1)) + V = FunctionSpace(m, FiniteElement('CG', cell, degree)) + u = TrialFunction(V) + v = TestFunction(V) + return (u*v + dot(grad(u), grad(v)))*dx + + +def elasticity(cell, degree): + m = Mesh(VectorElement('CG', cell, 1)) + V = FunctionSpace(m, VectorElement('CG', cell, degree)) + u = TrialFunction(V) + v = TestFunction(V) + + def eps(u): + return 0.5*(grad(u) + grad(u).T) + return inner(eps(u), eps(v))*dx + + +def count_flops(form): + kernel, = compile_form(form, parameters=dict(mode='tensor')) + return kernel.flop_count + + +@pytest.mark.parametrize('form', [mass, poisson, helmholtz, elasticity]) +@pytest.mark.parametrize(('cell', 'order'), + [(interval, 2), + (triangle, 4), + (tetrahedron, 6)]) +def test_bilinear(form, cell, order): + degrees = numpy.arange(1, 9 - 2 * cell.topological_dimension()) + flops = [count_flops(form(cell, int(degree))) + for degree in degrees] + rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees + 1)) + assert (rates < order).all() + + +@pytest.mark.parametrize(('cell', 'order'), + [(interval, 1), + (triangle, 2), + (tetrahedron, 3)]) +def test_linear(cell, order): + def form(cell, degree): + m = Mesh(VectorElement('CG', cell, 1)) + V = FunctionSpace(m, FiniteElement('CG', cell, degree)) + v = TestFunction(V) + return v*dx + + degrees = numpy.arange(2, 9 - 1.5 * cell.topological_dimension()) + flops = [count_flops(form(cell, int(degree))) + for degree in degrees] + rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees + 1)) + assert (rates < order).all() + + +@pytest.mark.parametrize(('cell', 'order'), + [(interval, 1), + (triangle, 2), + (tetrahedron, 3)]) +def test_functional(cell, order): + def form(cell, degree): + m = Mesh(VectorElement('CG', cell, 1)) + V = FunctionSpace(m, VectorElement('CG', cell, degree)) + f = Coefficient(V) + return div(f)*dx + + dim = cell.topological_dimension() + degrees = numpy.arange(2, 8 - dim) + (3 - dim) + flops = [count_flops(form(cell, int(degree))) + for degree in degrees] + rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees + 1)) + assert (rates < order).all() + + +def test_mini(): + m = Mesh(VectorElement('CG', triangle, 1)) + P1 = FiniteElement('Lagrange', triangle, 1) + B = FiniteElement("Bubble", triangle, 3) + V = FunctionSpace(m, VectorElement(P1 + B)) + u = TrialFunction(V) + v = TestFunction(V) + a = inner(grad(u), grad(v))*dx + count_flops(a) + + +if __name__ == "__main__": + import os + import sys + pytest.main(args=[os.path.abspath(__file__)] + sys.argv[1:]) diff --git a/tests/tsfc/test_tsfc_182.py b/tests/tsfc/test_tsfc_182.py new file mode 100644 index 0000000000..556a6bafb0 --- /dev/null +++ b/tests/tsfc/test_tsfc_182.py @@ -0,0 +1,35 @@ +import pytest + +from ufl import Coefficient, TestFunction, dx, inner, tetrahedron, Mesh, FunctionSpace +from finat.ufl import FiniteElement, MixedElement, VectorElement + +from tsfc import compile_form + + +@pytest.mark.parametrize('mode', ['vanilla', 'coffee', 'spectral']) +def test_delta_elimination(mode): + # Code sample courtesy of Marco Morandini: + # https://github.com/firedrakeproject/tsfc/issues/182 + scheme = "default" + degree = 3 + + element_lambda = FiniteElement("Quadrature", tetrahedron, degree, + quad_scheme=scheme) + element_eps_p = VectorElement("Quadrature", tetrahedron, degree, + dim=6, quad_scheme=scheme) + + element_chi_lambda = MixedElement(element_eps_p, element_lambda) + domain = Mesh(VectorElement("Lagrange", tetrahedron, 1)) + space = FunctionSpace(domain, element_chi_lambda) + + chi_lambda = Coefficient(space) + delta_chi_lambda = TestFunction(space) + + L = inner(delta_chi_lambda, chi_lambda) * dx(degree=degree, scheme=scheme) + kernel, = compile_form(L, parameters={'mode': mode}) + + +if __name__ == "__main__": + import os + import sys + pytest.main(args=[os.path.abspath(__file__)] + sys.argv[1:]) diff --git a/tests/tsfc/test_tsfc_204.py b/tests/tsfc/test_tsfc_204.py new file mode 100644 index 0000000000..89f1481590 --- /dev/null +++ b/tests/tsfc/test_tsfc_204.py @@ -0,0 +1,34 @@ +from tsfc import compile_form +from ufl import (Coefficient, FacetNormal, + FunctionSpace, Mesh, as_matrix, + dot, dS, ds, dx, facet, grad, inner, outer, split, triangle) +from finat.ufl import BrokenElement, FiniteElement, MixedElement, VectorElement + + +def test_physically_mapped_facet(): + mesh = Mesh(VectorElement("P", triangle, 1)) + + # set up variational problem + U = FiniteElement("Morley", mesh.ufl_cell(), 2) + V = FiniteElement("P", mesh.ufl_cell(), 1) + R = FiniteElement("P", mesh.ufl_cell(), 1) + Vv = VectorElement(BrokenElement(V)) + Qhat = VectorElement(BrokenElement(V[facet]), dim=2) + Vhat = VectorElement(V[facet], dim=2) + Z = FunctionSpace(mesh, MixedElement(U, Vv, Qhat, Vhat, R)) + + z = Coefficient(Z) + u, d, qhat, dhat, lam = split(z) + + s = FacetNormal(mesh) + trans = as_matrix([[1, 0], [0, 1]]) + mat = trans*grad(grad(u))*trans + outer(d, d) * u + J = (u**2*dx + + u**3*dx + + u**4*dx + + inner(mat, mat)*dx + + inner(grad(d), grad(d))*dx + + dot(s, d)**2*ds) + L_match = inner(qhat, dhat - d) + L = J + inner(lam, inner(d, d)-1)*dx + (L_match('+') + L_match('-'))*dS + L_match*ds + compile_form(L) diff --git a/tests/tsfc/test_tsfc_274.py b/tests/tsfc/test_tsfc_274.py new file mode 100644 index 0000000000..453d8746e8 --- /dev/null +++ b/tests/tsfc/test_tsfc_274.py @@ -0,0 +1,41 @@ +import gem +import numpy +from finat.point_set import PointSet +from gem.interpreter import evaluate +from tsfc.finatinterface import create_element +from ufl import quadrilateral +from finat.ufl import FiniteElement, RestrictedElement + + +def test_issue_274(): + # See https://github.com/firedrakeproject/tsfc/issues/274 + ufl_element = RestrictedElement( + FiniteElement("Q", quadrilateral, 2), restriction_domain="facet" + ) + ps = PointSet([[0.5]]) + finat_element = create_element(ufl_element) + evaluations = [] + for eid in range(4): + (val,) = finat_element.basis_evaluation(0, ps, (1, eid)).values() + evaluations.append(val) + + i = gem.Index() + j = gem.Index() + (expr,) = evaluate( + [ + gem.ComponentTensor( + gem.Indexed(gem.select_expression(evaluations, i), (j,)), + (*ps.indices, i, j), + ) + ] + ) + + (expect,) = evaluate( + [ + gem.ComponentTensor( + gem.Indexed(gem.ListTensor(evaluations), (i, j)), (*ps.indices, i, j) + ) + ] + ) + + assert numpy.allclose(expr.arr, expect.arr) diff --git a/tests/tsfc/test_underintegration.py b/tests/tsfc/test_underintegration.py new file mode 100644 index 0000000000..24221e05a7 --- /dev/null +++ b/tests/tsfc/test_underintegration.py @@ -0,0 +1,106 @@ +from functools import reduce + +import numpy +import pytest + +from ufl import (Mesh, FunctionSpace, TestFunction, TrialFunction, TensorProductCell, dx, + action, interval, quadrilateral, dot, grad) +from finat.ufl import FiniteElement, VectorElement + +from FIAT import ufc_cell +from FIAT.quadrature import GaussLobattoLegendreQuadratureLineRule, GaussLegendreQuadratureLineRule + +from finat.point_set import GaussLobattoLegendrePointSet, GaussLegendrePointSet +from finat.quadrature import QuadratureRule, TensorProductQuadratureRule + +from tsfc import compile_form + + +def gll_quadrature_rule(cell, elem_deg): + fiat_cell = ufc_cell("interval") + fiat_rule = GaussLobattoLegendreQuadratureLineRule(fiat_cell, elem_deg + 1) + line_rules = [QuadratureRule(GaussLobattoLegendrePointSet(fiat_rule.get_points()), + fiat_rule.get_weights()) + for _ in range(cell.topological_dimension())] + finat_rule = reduce(lambda a, b: TensorProductQuadratureRule([a, b]), line_rules) + return finat_rule + + +def gl_quadrature_rule(cell, elem_deg): + fiat_cell = ufc_cell("interval") + fiat_rule = GaussLegendreQuadratureLineRule(fiat_cell, elem_deg + 1) + line_rules = [QuadratureRule(GaussLegendrePointSet(fiat_rule.get_points()), + fiat_rule.get_weights()) + for _ in range(cell.topological_dimension())] + finat_rule = reduce(lambda a, b: TensorProductQuadratureRule([a, b]), line_rules) + return finat_rule + + +def mass_cg(cell, degree): + m = Mesh(VectorElement('Q', cell, 1)) + V = FunctionSpace(m, FiniteElement('Q', cell, degree, variant='spectral')) + u = TrialFunction(V) + v = TestFunction(V) + return u*v*dx(scheme=gll_quadrature_rule(cell, degree)) + + +def mass_dg(cell, degree): + m = Mesh(VectorElement('Q', cell, 1)) + V = FunctionSpace(m, FiniteElement('DQ', cell, degree, variant='spectral')) + u = TrialFunction(V) + v = TestFunction(V) + return u*v*dx(scheme=gl_quadrature_rule(cell, degree)) + + +def laplace(cell, degree): + m = Mesh(VectorElement('Q', cell, 1)) + V = FunctionSpace(m, FiniteElement('Q', cell, degree, variant='spectral')) + u = TrialFunction(V) + v = TestFunction(V) + return dot(grad(u), grad(v))*dx(scheme=gll_quadrature_rule(cell, degree)) + + +def count_flops(form): + kernel, = compile_form(form, parameters=dict(mode='spectral')) + return kernel.flop_count + + +@pytest.mark.parametrize('form', [mass_cg, mass_dg]) +@pytest.mark.parametrize(('cell', 'order'), + [(quadrilateral, 2), + (TensorProductCell(interval, interval), 2), + (TensorProductCell(quadrilateral, interval), 3)]) +def test_mass(form, cell, order): + degrees = numpy.arange(4, 10) + flops = [count_flops(form(cell, int(degree))) for degree in degrees] + rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees + 1)) + assert (rates < order).all() + + +@pytest.mark.parametrize('form', [mass_cg, mass_dg]) +@pytest.mark.parametrize(('cell', 'order'), + [(quadrilateral, 2), + (TensorProductCell(interval, interval), 2), + (TensorProductCell(quadrilateral, interval), 3)]) +def test_mass_action(form, cell, order): + degrees = numpy.arange(4, 10) + flops = [count_flops(action(form(cell, int(degree)))) for degree in degrees] + rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees + 1)) + assert (rates < order).all() + + +@pytest.mark.parametrize(('cell', 'order'), + [(quadrilateral, 4), + (TensorProductCell(interval, interval), 4), + (TensorProductCell(quadrilateral, interval), 5)]) +def test_laplace(cell, order): + degrees = numpy.arange(4, 10) + flops = [count_flops(laplace(cell, int(degree))) for degree in degrees] + rates = numpy.diff(numpy.log(flops)) / numpy.diff(numpy.log(degrees + 1)) + assert (rates < order).all() + + +if __name__ == "__main__": + import os + import sys + pytest.main(args=[os.path.abspath(__file__)] + sys.argv[1:]) diff --git a/tsfc/__init__.py b/tsfc/__init__.py new file mode 100644 index 0000000000..f9075c71a6 --- /dev/null +++ b/tsfc/__init__.py @@ -0,0 +1,22 @@ +from tsfc.driver import compile_form, compile_expression_dual_evaluation # noqa: F401 +from tsfc.parameters import default_parameters # noqa: F401 + +try: + from firedrake_citations import Citations + Citations().add("Kirby2006", """ +@Article{Kirby2006, + author = {Kirby, Robert C. and Logg, Anders}, + title = {A Compiler for Variational Forms}, + journal = {ACM Trans. Math. Softw.}, + year = 2006, + volume = 32, + number = 3, + pages = {417--444}, + month = sep, + numpages = 28, + doi = {10.1145/1163641.1163644}, + acmid = 1163644, +}""") + del Citations +except ImportError: + pass diff --git a/tsfc/coffee_mode.py b/tsfc/coffee_mode.py new file mode 100644 index 0000000000..632b915b41 --- /dev/null +++ b/tsfc/coffee_mode.py @@ -0,0 +1,81 @@ +from functools import partial, reduce + +from gem.node import traversal, Memoizer +from gem.gem import Failure, Sum, index_sum +from gem.optimise import replace_division, unroll_indexsum +from gem.refactorise import collect_monomials +from gem.unconcatenate import unconcatenate +from gem.coffee import optimise_monomial_sum +from gem.utils import groupby + +import tsfc.spectral as spectral + + +def Integrals(expressions, quadrature_multiindex, argument_multiindices, parameters): + """Constructs an integral representation for each GEM integrand + expression. + + :arg expressions: integrand multiplied with quadrature weight; + multi-root GEM expression DAG + :arg quadrature_multiindex: quadrature multiindex (tuple) + :arg argument_multiindices: tuple of argument multiindices, + one multiindex for each argument + :arg parameters: parameters dictionary + + :returns: list of integral representations + """ + # Unroll + max_extent = parameters["unroll_indexsum"] + if max_extent: + def predicate(index): + return index.extent <= max_extent + expressions = unroll_indexsum(expressions, predicate=predicate) + # Integral representation: just a GEM expression + return replace_division([index_sum(e, quadrature_multiindex) for e in expressions]) + + +def flatten(var_reps, index_cache): + """Flatten mode-specific intermediate representation to a series of + assignments. + + :arg var_reps: series of (return variable, [integral representation]) pairs + :arg index_cache: cache :py:class:`dict` for :py:func:`unconcatenate` + + :returns: series of (return variable, GEM expression root) pairs + """ + assignments = unconcatenate([(variable, reduce(Sum, reps)) + for variable, reps in var_reps], + cache=index_cache) + + def group_key(assignment): + variable, expression = assignment + return variable.free_indices + + for argument_indices, assignment_group in groupby(assignments, group_key): + variables, expressions = zip(*assignment_group) + expressions = optimise_expressions(expressions, argument_indices) + for var, expr in zip(variables, expressions): + yield (var, expr) + + +finalise_options = dict(remove_componenttensors=False) + + +def optimise_expressions(expressions, argument_indices): + """Perform loop optimisations on GEM DAGs + + :arg expressions: list of GEM DAGs + :arg argument_indices: tuple of argument indices + + :returns: list of optimised GEM DAGs + """ + # Skip optimisation for if Failure node is present + for n in traversal(expressions): + if isinstance(n, Failure): + return expressions + + # Apply argument factorisation unconditionally + classifier = partial(spectral.classify, set(argument_indices), + delta_inside=Memoizer(spectral._delta_inside)) + monomial_sums = collect_monomials(expressions, classifier) + return [optimise_monomial_sum(ms, argument_indices) for ms in monomial_sums] diff --git a/tsfc/driver.py b/tsfc/driver.py new file mode 100644 index 0000000000..6e3c3baaf3 --- /dev/null +++ b/tsfc/driver.py @@ -0,0 +1,331 @@ +import collections +import time +import sys +from itertools import chain +from finat.physically_mapped import DirectlyDefinedElement, PhysicallyMappedElement + +import ufl +from ufl.algorithms import extract_arguments, extract_coefficients +from ufl.algorithms.analysis import has_type +from ufl.classes import Form, GeometricQuantity +from ufl.domain import extract_unique_domain + +import gem +import gem.impero_utils as impero_utils + +import finat + +from tsfc import fem, ufl_utils +from tsfc.logging import logger +from tsfc.parameters import default_parameters, is_complex +from tsfc.ufl_utils import apply_mapping, extract_firedrake_constants +import tsfc.kernel_interface.firedrake_loopy as firedrake_interface_loopy + +# To handle big forms. The various transformations might need a deeper stack +sys.setrecursionlimit(3000) + + +TSFCIntegralDataInfo = collections.namedtuple("TSFCIntegralDataInfo", + ["domain", "integral_type", "subdomain_id", "domain_number", + "arguments", + "coefficients", "coefficient_numbers"]) +TSFCIntegralDataInfo.__doc__ = """ + Minimal set of objects for kernel builders. + + domain - The mesh. + integral_type - The type of integral. + subdomain_id - What is the subdomain id for this kernel. + domain_number - Which domain number in the original form + does this kernel correspond to (can be used to index into + original_form.ufl_domains() to get the correct domain). + coefficients - A list of coefficients. + coefficient_numbers - A list of which coefficients from the + form the kernel needs. + + This is a minimal set of objects that kernel builders need to + construct a kernel from :attr:`integrals` of :class:`~ufl.IntegralData`. + """ + + +def compile_form(form, prefix="form", parameters=None, interface=None, diagonal=False, log=False): + """Compiles a UFL form into a set of assembly kernels. + + :arg form: UFL form + :arg prefix: kernel name will start with this string + :arg parameters: parameters object + :arg diagonal: Are we building a kernel for the diagonal of a rank-2 element tensor? + :arg log: bool if the Kernel should be profiled with Log events + :returns: list of kernels + """ + cpu_time = time.time() + + assert isinstance(form, Form) + + GREEN = "\033[1;37;32m%s\033[0m" + + # Determine whether in complex mode: + complex_mode = parameters and is_complex(parameters.get("scalar_type")) + fd = ufl_utils.compute_form_data(form, complex_mode=complex_mode) + logger.info(GREEN % "compute_form_data finished in %g seconds.", time.time() - cpu_time) + + kernels = [] + for integral_data in fd.integral_data: + start = time.time() + kernel = compile_integral(integral_data, fd, prefix, parameters, interface=interface, diagonal=diagonal, log=log) + if kernel is not None: + kernels.append(kernel) + logger.info(GREEN % "compile_integral finished in %g seconds.", time.time() - start) + + logger.info(GREEN % "TSFC finished in %g seconds.", time.time() - cpu_time) + return kernels + + +def compile_integral(integral_data, form_data, prefix, parameters, interface, *, diagonal=False, log=False): + """Compiles a UFL integral into an assembly kernel. + + :arg integral_data: UFL integral data + :arg form_data: UFL form data + :arg prefix: kernel name will start with this string + :arg parameters: parameters object + :arg interface: backend module for the kernel interface + :arg diagonal: Are we building a kernel for the diagonal of a rank-2 element tensor? + :arg log: bool if the Kernel should be profiled with Log events + :returns: a kernel constructed by the kernel interface + """ + parameters = preprocess_parameters(parameters) + if interface is None: + interface = firedrake_interface_loopy.KernelBuilder + scalar_type = parameters["scalar_type"] + integral_type = integral_data.integral_type + if integral_type.startswith("interior_facet") and diagonal: + raise NotImplementedError("Sorry, we can't assemble the diagonal of a form for interior facet integrals") + mesh = integral_data.domain + arguments = form_data.preprocessed_form.arguments() + kernel_name = f"{prefix}_{integral_type}_integral" + # Dict mapping domains to index in original_form.ufl_domains() + domain_numbering = form_data.original_form.domain_numbering() + domain_number = domain_numbering[integral_data.domain] + coefficients = [form_data.function_replace_map[c] for c in integral_data.integral_coefficients] + # This is which coefficient in the original form the + # current coefficient is. + # Consider f*v*dx + g*v*ds, the full form contains two + # coefficients, but each integral only requires one. + coefficient_numbers = tuple(form_data.original_coefficient_positions[i] + for i, (_, enabled) in enumerate(zip(form_data.reduced_coefficients, integral_data.enabled_coefficients)) + if enabled) + integral_data_info = TSFCIntegralDataInfo(domain=integral_data.domain, + integral_type=integral_data.integral_type, + subdomain_id=integral_data.subdomain_id, + domain_number=domain_number, + arguments=arguments, + coefficients=coefficients, + coefficient_numbers=coefficient_numbers) + builder = interface(integral_data_info, + scalar_type, + diagonal=diagonal) + builder.set_coordinates(mesh) + builder.set_cell_sizes(mesh) + builder.set_coefficients(integral_data, form_data) + # TODO: We do not want pass constants to kernels that do not need them + # so we should attach the constants to integral data instead + builder.set_constants(form_data.constants) + ctx = builder.create_context() + for integral in integral_data.integrals: + params = parameters.copy() + params.update(integral.metadata()) # integral metadata overrides + integrand = ufl.replace(integral.integrand(), form_data.function_replace_map) + integrand_exprs = builder.compile_integrand(integrand, params, ctx) + integral_exprs = builder.construct_integrals(integrand_exprs, params) + builder.stash_integrals(integral_exprs, params, ctx) + return builder.construct_kernel(kernel_name, ctx, log) + + +def preprocess_parameters(parameters): + if parameters is None: + parameters = default_parameters() + else: + _ = default_parameters() + _.update(parameters) + parameters = _ + # Remove these here, they're handled later on. + if parameters.get("quadrature_degree") in ["auto", "default", None, -1, "-1"]: + del parameters["quadrature_degree"] + if parameters.get("quadrature_rule") in ["auto", "default", None]: + del parameters["quadrature_rule"] + return parameters + + +def compile_expression_dual_evaluation(expression, to_element, ufl_element, *, + domain=None, interface=None, + parameters=None, log=False): + """Compile a UFL expression to be evaluated against a compile-time known reference element's dual basis. + + Useful for interpolating UFL expressions into e.g. N1curl spaces. + + :arg expression: UFL expression + :arg to_element: A FInAT element for the target space + :arg ufl_element: The UFL element of the target space. + :arg domain: optional UFL domain the expression is defined on (required when expression contains no domain). + :arg interface: backend module for the kernel interface + :arg parameters: parameters object + :arg log: bool if the Kernel should be profiled with Log events + :returns: Loopy-based ExpressionKernel object. + """ + if parameters is None: + parameters = default_parameters() + else: + _ = default_parameters() + _.update(parameters) + parameters = _ + + # Determine whether in complex mode + complex_mode = is_complex(parameters["scalar_type"]) + + if isinstance(to_element, (PhysicallyMappedElement, DirectlyDefinedElement)): + raise NotImplementedError("Don't know how to interpolate onto zany spaces, sorry") + + orig_expression = expression + + # Map into reference space + expression = apply_mapping(expression, ufl_element, domain) + + # Apply UFL preprocessing + expression = ufl_utils.preprocess_expression(expression, + complex_mode=complex_mode) + + # Initialise kernel builder + if interface is None: + # Delayed import, loopy is a runtime dependency + from tsfc.kernel_interface.firedrake_loopy import ExpressionKernelBuilder as interface + + builder = interface(parameters["scalar_type"]) + arguments = extract_arguments(expression) + argument_multiindices = tuple(builder.create_element(arg.ufl_element()).get_indices() + for arg in arguments) + + # Replace coordinates (if any) unless otherwise specified by kwarg + if domain is None: + domain = extract_unique_domain(expression) + assert domain is not None + + # Collect required coefficients and determine numbering + coefficients = extract_coefficients(expression) + orig_coefficients = extract_coefficients(orig_expression) + coefficient_numbers = tuple(orig_coefficients.index(c) for c in coefficients) + builder.set_coefficient_numbers(coefficient_numbers) + + needs_external_coords = False + if has_type(expression, GeometricQuantity) or any(fem.needs_coordinate_mapping(c.ufl_element()) for c in coefficients): + # Create a fake coordinate coefficient for a domain. + coords_coefficient = ufl.Coefficient(ufl.FunctionSpace(domain, domain.ufl_coordinate_element())) + builder.domain_coordinate[domain] = coords_coefficient + builder.set_cell_sizes(domain) + coefficients = [coords_coefficient] + coefficients + needs_external_coords = True + builder.set_coefficients(coefficients) + + constants = extract_firedrake_constants(expression) + builder.set_constants(constants) + + # Split mixed coefficients + expression = ufl_utils.split_coefficients(expression, builder.coefficient_split) + + # Set up kernel config for translation of UFL expression to gem + kernel_cfg = dict(interface=builder, + ufl_cell=domain.ufl_cell(), + # FIXME: change if we ever implement + # interpolation on facets. + integral_type="cell", + argument_multiindices=argument_multiindices, + index_cache={}, + scalar_type=parameters["scalar_type"]) + + # Allow interpolation onto QuadratureElements to refer to the quadrature + # rule they represent + if isinstance(to_element, finat.QuadratureElement): + kernel_cfg["quadrature_rule"] = to_element._rule + + # Create callable for translation of UFL expression to gem + fn = DualEvaluationCallable(expression, kernel_cfg) + + # Get the gem expression for dual evaluation and corresponding basis + # indices needed for compilation of the expression + evaluation, basis_indices = to_element.dual_evaluation(fn) + + # Build kernel body + return_indices = basis_indices + tuple(chain(*argument_multiindices)) + return_shape = tuple(i.extent for i in return_indices) + return_var = gem.Variable('A', return_shape) + return_expr = gem.Indexed(return_var, return_indices) + + # TODO: one should apply some GEM optimisations as in assembly, + # but we don't for now. + evaluation, = impero_utils.preprocess_gem([evaluation]) + impero_c = impero_utils.compile_gem([(return_expr, evaluation)], return_indices) + index_names = dict((idx, "p%d" % i) for (i, idx) in enumerate(basis_indices)) + # Handle kernel interface requirements + builder.register_requirements([evaluation]) + builder.set_output(return_var) + # Build kernel tuple + return builder.construct_kernel(impero_c, index_names, needs_external_coords, log=log) + + +class DualEvaluationCallable(object): + """ + Callable representing a function to dual evaluate. + + When called, this takes in a + :class:`finat.point_set.AbstractPointSet` and returns a GEM + expression for evaluation of the function at those points. + + :param expression: UFL expression for the function to dual evaluate. + :param kernel_cfg: A kernel configuration for creation of a + :class:`GemPointContext` or a :class:`PointSetContext` + + Not intended for use outside of + :func:`compile_expression_dual_evaluation`. + """ + def __init__(self, expression, kernel_cfg): + self.expression = expression + self.kernel_cfg = kernel_cfg + + def __call__(self, ps): + """The function to dual evaluate. + + :param ps: The :class:`finat.point_set.AbstractPointSet` for + evaluating at + :returns: a gem expression representing the evaluation of the + input UFL expression at the given point set ``ps``. + For point set points with some shape ``(*value_shape)`` + (i.e. ``()`` for scalar points ``(x)`` for vector points + ``(x, y)`` for tensor points etc) then the gem expression + has shape ``(*value_shape)`` and free indices corresponding + to the input :class:`finat.point_set.AbstractPointSet`'s + free indices alongside any input UFL expression free + indices. + """ + + if not isinstance(ps, finat.point_set.AbstractPointSet): + raise ValueError("Callable argument not a point set!") + + # Avoid modifying saved kernel config + kernel_cfg = self.kernel_cfg.copy() + + if isinstance(ps, finat.point_set.UnknownPointSet): + # Run time known points + kernel_cfg.update(point_indices=ps.indices, point_expr=ps.expression) + # GemPointContext's aren't allowed to have quadrature rules + kernel_cfg.pop("quadrature_rule", None) + translation_context = fem.GemPointContext(**kernel_cfg) + else: + # Compile time known points + kernel_cfg.update(point_set=ps) + translation_context = fem.PointSetContext(**kernel_cfg) + + gem_expr, = fem.compile_ufl(self.expression, translation_context, point_sum=False) + # In some cases ps.indices may be dropped from expr, but nothing + # new should now appear + argument_multiindices = kernel_cfg["argument_multiindices"] + assert set(gem_expr.free_indices) <= set(chain(ps.indices, *argument_multiindices)) + + return gem_expr diff --git a/tsfc/fem.py b/tsfc/fem.py new file mode 100644 index 0000000000..abc8bc7cb8 --- /dev/null +++ b/tsfc/fem.py @@ -0,0 +1,777 @@ +"""Functions to translate UFL finite element objects and reference +geometric quantities into GEM expressions.""" + +import collections +import itertools +from functools import singledispatch + +import gem +import numpy +import ufl +from FIAT.orientation_utils import Orientation as FIATOrientation +from FIAT.reference_element import UFCHexahedron, UFCSimplex, make_affine_mapping +from FIAT.reference_element import TensorProductCell +from finat.physically_mapped import (NeedsCoordinateMappingElement, + PhysicalGeometry) +from finat.point_set import PointSet, PointSingleton +from finat.quadrature import make_quadrature +from gem.node import traversal +from gem.optimise import constant_fold_zero, ffc_rounding +from gem.unconcatenate import unconcatenate +from gem.utils import cached_property +from ufl.classes import (Argument, CellCoordinate, CellEdgeVectors, + CellFacetJacobian, CellOrientation, CellOrigin, + CellVertices, CellVolume, Coefficient, FacetArea, + FacetCoordinate, GeometricQuantity, Jacobian, + JacobianDeterminant, NegativeRestricted, + PositiveRestricted, QuadratureWeight, + ReferenceCellEdgeVectors, ReferenceCellVolume, + ReferenceFacetVolume, ReferenceNormal, + SpatialCoordinate) +from ufl.corealg.map_dag import map_expr_dag, map_expr_dags +from ufl.corealg.multifunction import MultiFunction +from ufl.domain import extract_unique_domain + +from tsfc import ufl2gem +from tsfc.finatinterface import as_fiat_cell, create_element +from tsfc.kernel_interface import ProxyKernelInterface +from tsfc.modified_terminals import (analyse_modified_terminal, + construct_modified_terminal) +from tsfc.parameters import is_complex +from tsfc.ufl_utils import (ModifiedTerminalMixin, PickRestriction, + TSFCConstantMixin, entity_avg, one_times, + preprocess_expression, simplify_abs) + + +class ContextBase(ProxyKernelInterface): + """Common UFL -> GEM translation context.""" + + keywords = ('ufl_cell', + 'fiat_cell', + 'integral_type', + 'integration_dim', + 'entity_ids', + 'argument_multiindices', + 'facetarea', + 'index_cache', + 'scalar_type') + + def __init__(self, interface, **kwargs): + ProxyKernelInterface.__init__(self, interface) + + invalid_keywords = set(kwargs.keys()) - set(self.keywords) + if invalid_keywords: + raise ValueError("unexpected keyword argument '{0}'".format(invalid_keywords.pop())) + self.__dict__.update(kwargs) + + @cached_property + def fiat_cell(self): + return as_fiat_cell(self.ufl_cell) + + @cached_property + def integration_dim(self): + return self.fiat_cell.get_dimension() + + entity_ids = [0] + + @cached_property + def epsilon(self): + return numpy.finfo(self.scalar_type).resolution + + @cached_property + def complex_mode(self): + return is_complex(self.scalar_type) + + def entity_selector(self, callback, restriction): + """Selects code for the correct entity at run-time. Callback + generates code for a specified entity. + + This function passes ``callback`` the entity number. + + :arg callback: A function to be called with an entity number + that generates code for that entity. + :arg restriction: Restriction of the modified terminal, used + for entity selection. + """ + if len(self.entity_ids) == 1: + return callback(self.entity_ids[0]) + else: + f = self.entity_number(restriction) + return gem.select_expression(list(map(callback, self.entity_ids)), f) + + argument_multiindices = () + + @cached_property + def index_cache(self): + return {} + + @cached_property + def translator(self): + # NOTE: reference cycle! + return Translator(self) + + @cached_property + def use_canonical_quadrature_point_ordering(self): + return isinstance(self.fiat_cell, UFCHexahedron) and self.integral_type in ['exterior_facet', 'interior_facet'] + + +class CoordinateMapping(PhysicalGeometry): + """Callback class that provides physical geometry to FInAT elements. + + Required for elements whose basis transformation requires physical + geometry such as Argyris and Hermite. + + :arg mt: The modified terminal whose element will be tabulated. + :arg interface: The interface carrying information (generally a + :class:`PointSetContext`). + """ + def __init__(self, mt, interface): + super().__init__() + self.mt = mt + self.interface = interface + + def preprocess(self, expr, context): + """Preprocess a UFL expression for translation. + + :arg expr: A UFL expression + :arg context: The translation context. + :returns: A new UFL expression + """ + ifacet = self.interface.integral_type.startswith("interior_facet") + return preprocess_expression(expr, complex_mode=context.complex_mode, + do_apply_restrictions=ifacet) + + @property + def config(self): + config = {name: getattr(self.interface, name) + for name in ["ufl_cell", "index_cache", "scalar_type"]} + config["interface"] = self.interface + return config + + def cell_size(self): + return self.interface.cell_size(self.mt.restriction) + + def jacobian_at(self, point): + ps = PointSingleton(point) + expr = Jacobian(extract_unique_domain(self.mt.terminal)) + assert ps.expression.shape == (extract_unique_domain(expr).topological_dimension(), ) + if self.mt.restriction == '+': + expr = PositiveRestricted(expr) + elif self.mt.restriction == '-': + expr = NegativeRestricted(expr) + config = {"point_set": PointSingleton(point)} + config.update(self.config) + context = PointSetContext(**config) + expr = self.preprocess(expr, context) + return map_expr_dag(context.translator, expr) + + def detJ_at(self, point): + expr = JacobianDeterminant(extract_unique_domain(self.mt.terminal)) + if self.mt.restriction == '+': + expr = PositiveRestricted(expr) + elif self.mt.restriction == '-': + expr = NegativeRestricted(expr) + config = {"point_set": PointSingleton(point)} + config.update(self.config) + context = PointSetContext(**config) + expr = self.preprocess(expr, context) + return map_expr_dag(context.translator, expr) + + def reference_normals(self): + cell = self.interface.fiat_cell + sd = cell.get_spatial_dimension() + num_faces = len(cell.get_topology()[sd-1]) + + return gem.Literal(numpy.asarray([cell.compute_normal(i) for i in range(num_faces)])) + + def reference_edge_tangents(self): + cell = self.interface.fiat_cell + num_edges = len(cell.get_topology()[1]) + return gem.Literal(numpy.asarray([cell.compute_edge_tangent(i) for i in range(num_edges)])) + + def physical_tangents(self): + cell = self.interface.fiat_cell + sd = cell.get_spatial_dimension() + num_edges = len(cell.get_topology()[1]) + els = self.physical_edge_lengths() + rts = gem.ListTensor([cell.compute_tangents(1, i)[0] / els[i] for i in range(num_edges)]) + jac = self.jacobian_at(cell.make_points(sd, 0, sd+1)[0]) + + return rts @ jac.T + + def physical_normals(self): + cell = self.interface.fiat_cell + if not (isinstance(cell, UFCSimplex) and cell.get_dimension() == 2): + raise NotImplementedError("Can't do physical normals on that cell yet") + + num_edges = len(cell.get_topology()[1]) + pts = self.physical_tangents() + return gem.ListTensor([[pts[i, 1], -1*pts[i, 0]] for i in range(num_edges)]) + + def physical_edge_lengths(self): + expr = ufl.classes.CellEdgeVectors(extract_unique_domain(self.mt.terminal)) + if self.mt.restriction == '+': + expr = PositiveRestricted(expr) + elif self.mt.restriction == '-': + expr = NegativeRestricted(expr) + + cell = self.interface.fiat_cell + sd = cell.get_spatial_dimension() + num_edges = len(cell.get_topology()[1]) + expr = ufl.as_vector([ufl.sqrt(ufl.dot(expr[i, :], expr[i, :])) for i in range(num_edges)]) + config = {"point_set": PointSingleton(cell.make_points(sd, 0, sd+1)[0])} + config.update(self.config) + context = PointSetContext(**config) + expr = self.preprocess(expr, context) + return map_expr_dag(context.translator, expr) + + def physical_points(self, point_set, entity=None): + """Converts point_set from reference to physical space""" + expr = SpatialCoordinate(extract_unique_domain(self.mt.terminal)) + point_shape, = point_set.expression.shape + if entity is not None: + e, _ = entity + assert point_shape == e + else: + assert point_shape == extract_unique_domain(expr).topological_dimension() + if self.mt.restriction == '+': + expr = PositiveRestricted(expr) + elif self.mt.restriction == '-': + expr = NegativeRestricted(expr) + config = {"point_set": point_set} + config.update(self.config) + if entity is not None: + config.update({name: getattr(self.interface, name) + for name in ["integration_dim", "entity_ids"]}) + context = PointSetContext(**config) + expr = self.preprocess(expr, context) + mapped = map_expr_dag(context.translator, expr) + indices = tuple(gem.Index() for _ in mapped.shape) + return gem.ComponentTensor(gem.Indexed(mapped, indices), point_set.indices + indices) + + def physical_vertices(self): + vs = PointSet(self.interface.fiat_cell.vertices) + return self.physical_points(vs) + + +def needs_coordinate_mapping(element): + """Does this UFL element require a CoordinateMapping for translation?""" + if element.family() == 'Real': + return False + else: + return isinstance(create_element(element), NeedsCoordinateMappingElement) + + +class PointSetContext(ContextBase): + """Context for compile-time known evaluation points.""" + + keywords = ContextBase.keywords + ( + 'quadrature_degree', + 'quadrature_rule', + 'point_set', + 'weight_expr', + ) + + @cached_property + def integration_cell(self): + return self.fiat_cell.construct_subelement(self.integration_dim) + + @cached_property + def quadrature_rule(self): + return make_quadrature(self.integration_cell, self.quadrature_degree) + + @cached_property + def point_set(self): + return self.quadrature_rule.point_set + + @cached_property + def point_indices(self): + return self.point_set.indices + + @cached_property + def point_expr(self): + return self.point_set.expression + + @cached_property + def weight_expr(self): + return self.quadrature_rule.weight_expression + + def basis_evaluation(self, finat_element, mt, entity_id): + return finat_element.basis_evaluation(mt.local_derivatives, + self.point_set, + (self.integration_dim, entity_id), + coordinate_mapping=CoordinateMapping(mt, self)) + + +class GemPointContext(ContextBase): + """Context for evaluation at arbitrary reference points.""" + + keywords = ContextBase.keywords + ( + 'point_indices', + 'point_expr', + 'weight_expr', + ) + + def basis_evaluation(self, finat_element, mt, entity_id): + return finat_element.point_evaluation(mt.local_derivatives, + self.point_expr, + (self.integration_dim, entity_id)) + + +class Translator(MultiFunction, ModifiedTerminalMixin, ufl2gem.Mixin): + """Multifunction for translating UFL -> GEM. Incorporates ufl2gem.Mixin, and + dispatches on terminal type when reaching modified terminals.""" + + def __init__(self, context): + # MultiFunction.__init__ does not call further __init__ + # methods, but ufl2gem.Mixin must be initialised. + # (ModifiedTerminalMixin requires no initialisation.) + MultiFunction.__init__(self) + ufl2gem.Mixin.__init__(self) + + # Need context during translation! + self.context = context + + # We just use the provided quadrature rule to + # perform the integration. + # Can't put these in the ufl2gem mixin, since they (unlike + # everything else) want access to the translation context. + def cell_avg(self, o): + if self.context.integral_type != "cell": + # Need to create a cell-based quadrature rule and + # translate the expression using that (c.f. CellVolume + # below). + raise NotImplementedError("CellAvg on non-cell integrals not yet implemented") + integrand, = o.ufl_operands + domain = extract_unique_domain(o) + measure = ufl.Measure(self.context.integral_type, domain=domain) + integrand, degree, argument_multiindices = entity_avg(integrand / CellVolume(domain), measure, self.context.argument_multiindices) + + config = {name: getattr(self.context, name) + for name in ["ufl_cell", "index_cache", "scalar_type"]} + config.update(quadrature_degree=degree, interface=self.context, + argument_multiindices=argument_multiindices) + expr, = compile_ufl(integrand, PointSetContext(**config), point_sum=True) + return expr + + def facet_avg(self, o): + if self.context.integral_type == "cell": + raise ValueError("Can't take FacetAvg in cell integral") + integrand, = o.ufl_operands + domain = extract_unique_domain(o) + measure = ufl.Measure(self.context.integral_type, domain=domain) + integrand, degree, argument_multiindices = entity_avg(integrand / FacetArea(domain), measure, self.context.argument_multiindices) + + config = {name: getattr(self.context, name) + for name in ["ufl_cell", "index_cache", "scalar_type", + "integration_dim", "entity_ids", + "integral_type"]} + config.update(quadrature_degree=degree, interface=self.context, + argument_multiindices=argument_multiindices) + expr, = compile_ufl(integrand, PointSetContext(**config), point_sum=True) + return expr + + def modified_terminal(self, o): + """Overrides the modified terminal handler from + :class:`ModifiedTerminalMixin`.""" + mt = analyse_modified_terminal(o) + return translate(mt.terminal, mt, self.context) + + +@singledispatch +def translate(terminal, mt, ctx): + """Translates modified terminals into GEM. + + :arg terminal: terminal, for dispatching + :arg mt: analysed modified terminal + :arg ctx: translator context + :returns: GEM translation of the modified terminal + """ + raise AssertionError("Cannot handle terminal type: %s" % type(terminal)) + + +@translate.register(QuadratureWeight) +def translate_quadratureweight(terminal, mt, ctx): + return ctx.weight_expr + + +@translate.register(GeometricQuantity) +def translate_geometricquantity(terminal, mt, ctx): + raise NotImplementedError("Cannot handle geometric quantity type: %s" % type(terminal)) + + +@translate.register(CellOrientation) +def translate_cell_orientation(terminal, mt, ctx): + return ctx.cell_orientation(mt.restriction) + + +@translate.register(ReferenceCellVolume) +def translate_reference_cell_volume(terminal, mt, ctx): + return gem.Literal(ctx.fiat_cell.volume()) + + +@translate.register(ReferenceFacetVolume) +def translate_reference_facet_volume(terminal, mt, ctx): + assert ctx.integral_type != "cell" + # Sum of quadrature weights is entity volume + return gem.optimise.aggressive_unroll(gem.index_sum(ctx.weight_expr, + ctx.point_indices)) + + +@translate.register(CellFacetJacobian) +def translate_cell_facet_jacobian(terminal, mt, ctx): + cell = ctx.fiat_cell + facet_dim = ctx.integration_dim + assert facet_dim != cell.get_dimension() + + def callback(entity_id): + return gem.Literal(make_cell_facet_jacobian(cell, facet_dim, entity_id)) + return ctx.entity_selector(callback, mt.restriction) + + +def make_cell_facet_jacobian(cell, facet_dim, facet_i): + facet_cell = cell.construct_subelement(facet_dim) + xs = facet_cell.get_vertices() + ys = cell.get_vertices_of_subcomplex(cell.get_topology()[facet_dim][facet_i]) + + # Use first 'dim' points to make an affine mapping + dim = cell.get_spatial_dimension() + A, b = make_affine_mapping(xs[:dim], ys[:dim]) + + for x, y in zip(xs[dim:], ys[dim:]): + # The rest of the points are checked to make sure the + # mapping really *is* affine. + assert numpy.allclose(y, A.dot(x) + b) + + return A + + +@translate.register(ReferenceNormal) +def translate_reference_normal(terminal, mt, ctx): + def callback(facet_i): + n = ctx.fiat_cell.compute_reference_normal(ctx.integration_dim, facet_i) + return gem.Literal(n) + return ctx.entity_selector(callback, mt.restriction) + + +@translate.register(ReferenceCellEdgeVectors) +def translate_reference_cell_edge_vectors(terminal, mt, ctx): + from FIAT.reference_element import \ + TensorProductCell as fiat_TensorProductCell + fiat_cell = ctx.fiat_cell + if isinstance(fiat_cell, fiat_TensorProductCell): + raise NotImplementedError("ReferenceCellEdgeVectors not implemented on TensorProductElements yet") + + nedges = len(fiat_cell.get_topology()[1]) + vecs = numpy.vstack(map(fiat_cell.compute_edge_tangent, range(nedges))) + assert vecs.shape == terminal.ufl_shape + return gem.Literal(vecs) + + +@translate.register(CellCoordinate) +def translate_cell_coordinate(terminal, mt, ctx): + if ctx.integration_dim == ctx.fiat_cell.get_dimension(): + return ctx.point_expr + + # This destroys the structure of the quadrature points, but since + # this code path is only used to implement CellCoordinate in facet + # integrals, hopefully it does not matter much. + ps = ctx.point_set + point_shape = tuple(index.extent for index in ps.indices) + + def callback(entity_id): + t = ctx.fiat_cell.get_entity_transform(ctx.integration_dim, entity_id) + data = numpy.asarray(list(map(t, ps.points))) + return gem.Literal(data.reshape(point_shape + data.shape[1:])) + + return gem.partial_indexed(ctx.entity_selector(callback, mt.restriction), + ps.indices) + + +@translate.register(FacetCoordinate) +def translate_facet_coordinate(terminal, mt, ctx): + assert ctx.integration_dim != ctx.fiat_cell.get_dimension() + return ctx.point_expr + + +@translate.register(SpatialCoordinate) +def translate_spatialcoordinate(terminal, mt, ctx): + # Replace terminal with a Coefficient + terminal = ctx.coordinate(extract_unique_domain(terminal)) + # Get back to reference space + terminal = preprocess_expression(terminal, complex_mode=ctx.complex_mode) + # Rebuild modified terminal + expr = construct_modified_terminal(mt, terminal) + # Translate replaced UFL snippet + return ctx.translator(expr) + + +class CellVolumeKernelInterface(ProxyKernelInterface): + # Since CellVolume is evaluated as a cell integral, we must ensure + # that the right restriction is applied when it is used in an + # interior facet integral. This proxy diverts coefficient + # translation to use a specified restriction. + + def __init__(self, wrapee, restriction): + ProxyKernelInterface.__init__(self, wrapee) + self.restriction = restriction + + def coefficient(self, ufl_coefficient, r): + assert r is None + return self._wrapee.coefficient(ufl_coefficient, self.restriction) + + +@translate.register(CellVolume) +def translate_cellvolume(terminal, mt, ctx): + integrand, degree = one_times(ufl.dx(domain=extract_unique_domain(terminal))) + interface = CellVolumeKernelInterface(ctx, mt.restriction) + + config = {name: getattr(ctx, name) + for name in ["ufl_cell", "index_cache", "scalar_type"]} + config.update(interface=interface, quadrature_degree=degree) + expr, = compile_ufl(integrand, PointSetContext(**config), point_sum=True) + return expr + + +@translate.register(FacetArea) +def translate_facetarea(terminal, mt, ctx): + assert ctx.integral_type != 'cell' + domain = extract_unique_domain(terminal) + integrand, degree = one_times(ufl.Measure(ctx.integral_type, domain=domain)) + + config = {name: getattr(ctx, name) + for name in ["ufl_cell", "integration_dim", "scalar_type", + "entity_ids", "index_cache"]} + config.update(interface=ctx, quadrature_degree=degree) + expr, = compile_ufl(integrand, PointSetContext(**config), point_sum=True) + return expr + + +@translate.register(CellOrigin) +def translate_cellorigin(terminal, mt, ctx): + domain = extract_unique_domain(terminal) + coords = SpatialCoordinate(domain) + expression = construct_modified_terminal(mt, coords) + point_set = PointSingleton((0.0,) * domain.topological_dimension()) + + config = {name: getattr(ctx, name) + for name in ["ufl_cell", "index_cache", "scalar_type"]} + config.update(interface=ctx, point_set=point_set) + context = PointSetContext(**config) + return context.translator(expression) + + +@translate.register(CellVertices) +def translate_cell_vertices(terminal, mt, ctx): + coords = SpatialCoordinate(extract_unique_domain(terminal)) + ufl_expr = construct_modified_terminal(mt, coords) + ps = PointSet(numpy.array(ctx.fiat_cell.get_vertices())) + + config = {name: getattr(ctx, name) + for name in ["ufl_cell", "index_cache", "scalar_type"]} + config.update(interface=ctx, point_set=ps) + context = PointSetContext(**config) + expr = context.translator(ufl_expr) + + # Wrap up point (vertex) index + c = gem.Index() + return gem.ComponentTensor(gem.Indexed(expr, (c,)), ps.indices + (c,)) + + +@translate.register(CellEdgeVectors) +def translate_cell_edge_vectors(terminal, mt, ctx): + # WARNING: Assumes straight edges! + coords = CellVertices(extract_unique_domain(terminal)) + ufl_expr = construct_modified_terminal(mt, coords) + cell_vertices = ctx.translator(ufl_expr) + + e = gem.Index() + c = gem.Index() + expr = gem.ListTensor([ + gem.Sum(gem.Indexed(cell_vertices, (u, c)), + gem.Product(gem.Literal(-1), + gem.Indexed(cell_vertices, (v, c)))) + for _, (u, v) in sorted(ctx.fiat_cell.get_topology()[1].items()) + ]) + return gem.ComponentTensor(gem.Indexed(expr, (e,)), (e, c)) + + +def fiat_to_ufl(fiat_dict, order): + # All derivative multiindices must be of the same dimension. + dimension, = set(len(alpha) for alpha in fiat_dict.keys()) + + # All derivative tables must have the same shape. + shape, = set(table.shape for table in fiat_dict.values()) + sigma = tuple(gem.Index(extent=extent) for extent in shape) + + # Convert from FIAT to UFL format + eye = numpy.eye(dimension, dtype=int) + tensor = numpy.empty((dimension,) * order, dtype=object) + for multiindex in numpy.ndindex(tensor.shape): + alpha = tuple(eye[multiindex, :].sum(axis=0)) + tensor[multiindex] = gem.Indexed(fiat_dict[alpha], sigma) + delta = tuple(gem.Index(extent=dimension) for _ in range(order)) + if order > 0: + tensor = gem.Indexed(gem.ListTensor(tensor), delta) + else: + tensor = tensor[()] + return gem.ComponentTensor(tensor, sigma + delta) + + +@translate.register(Argument) +def translate_argument(terminal, mt, ctx): + argument_multiindex = ctx.argument_multiindices[terminal.number()] + sigma = tuple(gem.Index(extent=d) for d in mt.expr.ufl_shape) + element = ctx.create_element(terminal.ufl_element(), restriction=mt.restriction) + + def callback(entity_id): + finat_dict = ctx.basis_evaluation(element, mt, entity_id) + # Filter out irrelevant derivatives + filtered_dict = {alpha: table + for alpha, table in finat_dict.items() + if sum(alpha) == mt.local_derivatives} + + # Change from FIAT to UFL arrangement + square = fiat_to_ufl(filtered_dict, mt.local_derivatives) + + # A numerical hack that FFC used to apply on FIAT tables still + # lives on after ditching FFC and switching to FInAT. + return ffc_rounding(square, ctx.epsilon) + table = ctx.entity_selector(callback, mt.restriction) + if ctx.use_canonical_quadrature_point_ordering: + quad_multiindex = ctx.quadrature_rule.point_set.indices + quad_multiindex_permuted = _make_quad_multiindex_permuted(mt, ctx) + mapper = gem.node.MemoizerArg(gem.optimise.filtered_replace_indices) + table = mapper(table, tuple(zip(quad_multiindex, quad_multiindex_permuted))) + return gem.ComponentTensor(gem.Indexed(table, argument_multiindex + sigma), sigma) + + +@translate.register(TSFCConstantMixin) +def translate_constant_value(terminal, mt, ctx): + return ctx.constant(terminal) + + +@translate.register(Coefficient) +def translate_coefficient(terminal, mt, ctx): + vec = ctx.coefficient(terminal, mt.restriction) + + if terminal.ufl_element().family() == 'Real': + assert mt.local_derivatives == 0 + return vec + + element = ctx.create_element(terminal.ufl_element(), restriction=mt.restriction) + + # Collect FInAT tabulation for all entities + per_derivative = collections.defaultdict(list) + for entity_id in ctx.entity_ids: + finat_dict = ctx.basis_evaluation(element, mt, entity_id) + for alpha, table in finat_dict.items(): + # Filter out irrelevant derivatives + if sum(alpha) == mt.local_derivatives: + # A numerical hack that FFC used to apply on FIAT + # tables still lives on after ditching FFC and + # switching to FInAT. + table = ffc_rounding(table, ctx.epsilon) + per_derivative[alpha].append(table) + + # Merge entity tabulations for each derivative + if len(ctx.entity_ids) == 1: + def take_singleton(xs): + x, = xs # asserts singleton + return x + per_derivative = {alpha: take_singleton(tables) + for alpha, tables in per_derivative.items()} + else: + f = ctx.entity_number(mt.restriction) + per_derivative = {alpha: gem.select_expression(tables, f) + for alpha, tables in per_derivative.items()} + + # Coefficient evaluation + ctx.index_cache.setdefault(terminal.ufl_element(), element.get_indices()) + beta = ctx.index_cache[terminal.ufl_element()] + zeta = element.get_value_indices() + vec_beta, = gem.optimise.remove_componenttensors([gem.Indexed(vec, beta)]) + value_dict = {} + for alpha, table in per_derivative.items(): + table_qi = gem.Indexed(table, beta + zeta) + summands = [] + for var, expr in unconcatenate([(vec_beta, table_qi)], ctx.index_cache): + indices = tuple(i for i in var.index_ordering() if i not in ctx.unsummed_coefficient_indices) + value = gem.IndexSum(gem.Product(expr, var), indices) + summands.append(gem.optimise.contraction(value)) + optimised_value = gem.optimise.make_sum(summands) + value_dict[alpha] = gem.ComponentTensor(optimised_value, zeta) + + # Change from FIAT to UFL arrangement + result = fiat_to_ufl(value_dict, mt.local_derivatives) + assert result.shape == mt.expr.ufl_shape + assert set(result.free_indices) - ctx.unsummed_coefficient_indices <= set(ctx.point_indices) + + # Detect Jacobian of affine cells + if not result.free_indices and all(numpy.count_nonzero(node.array) <= 2 + for node in traversal((result,)) + if isinstance(node, gem.Literal)): + result = gem.optimise.aggressive_unroll(result) + + if ctx.use_canonical_quadrature_point_ordering: + quad_multiindex = ctx.quadrature_rule.point_set.indices + quad_multiindex_permuted = _make_quad_multiindex_permuted(mt, ctx) + mapper = gem.node.MemoizerArg(gem.optimise.filtered_replace_indices) + result = mapper(result, tuple(zip(quad_multiindex, quad_multiindex_permuted))) + return result + + +def _make_quad_multiindex_permuted(mt, ctx): + quad_rule = ctx.quadrature_rule + # Note that each quad index here represents quad points on a physical + # cell axis, but the table is indexed by indices representing the points + # on each reference cell axis, so we need to apply permutation based on the orientation. + cell = quad_rule.ref_el + quad_multiindex = quad_rule.point_set.indices + if isinstance(cell, TensorProductCell): + for comp in set(cell.cells): + extents = set(q.extent for c, q in zip(cell.cells, quad_multiindex) if c == comp) + if len(extents) != 1: + raise ValueError("Must have the same number of quadrature points in each symmetric axis") + quad_multiindex_permuted = [] + o = ctx.entity_orientation(mt.restriction) + if not isinstance(o, FIATOrientation): + raise ValueError(f"Expecting an instance of FIATOrientation : got {o}") + eo = cell.extract_extrinsic_orientation(o) + eo_perm_map = gem.Literal(quad_rule.extrinsic_orientation_permutation_map, dtype=gem.uint_type) + for ref_axis in range(len(quad_multiindex)): + io = cell.extract_intrinsic_orientation(o, ref_axis) + io_perm_map = gem.Literal(quad_rule.intrinsic_orientation_permutation_map_tuple[ref_axis], dtype=gem.uint_type) + # Effectively swap axes if needed. + ref_index = tuple((phys_index, gem.Indexed(eo_perm_map, (eo, ref_axis, phys_axis))) for phys_axis, phys_index in enumerate(quad_multiindex)) + quad_index_permuted = gem.VariableIndex(gem.FlexiblyIndexed(io_perm_map, ((0, ((io, 1), )), (0, ref_index)))) + quad_multiindex_permuted.append(quad_index_permuted) + return tuple(quad_multiindex_permuted) + + +def compile_ufl(expression, context, interior_facet=False, point_sum=False): + """Translate a UFL expression to GEM. + + :arg expression: The UFL expression to compile. + :arg context: translation context - either a :class:`GemPointContext` + or :class:`PointSetContext` + :arg interior_facet: If ``true``, treat expression as an interior + facet integral (default ``False``) + :arg point_sum: If ``true``, return a `gem.IndexSum` of the final + gem expression along the ``context.point_indices`` (if present). + """ + + # Abs-simplification + expression = simplify_abs(expression, context.complex_mode) + if interior_facet: + expressions = [] + for rs in itertools.product(("+", "-"), repeat=len(context.argument_multiindices)): + expressions.append(map_expr_dag(PickRestriction(*rs), expression)) + else: + expressions = [expression] + + # Translate UFL to GEM, lowering finite element specific nodes + result = map_expr_dags(context.translator, expressions) + if point_sum: + result = [gem.index_sum(expr, context.point_indices) for expr in result] + return constant_fold_zero(result) diff --git a/tsfc/finatinterface.py b/tsfc/finatinterface.py new file mode 100644 index 0000000000..b7e3d0ad72 --- /dev/null +++ b/tsfc/finatinterface.py @@ -0,0 +1,365 @@ +# This file was modified from FFC +# (http://bitbucket.org/fenics-project/ffc), copyright notice +# reproduced below. +# +# Copyright (C) 2009-2013 Kristian B. Oelgaard and Anders Logg +# +# This file is part of FFC. +# +# FFC is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# FFC is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Lesser General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with FFC. If not, see . + +import weakref +from functools import singledispatch + +import FIAT +import finat +import finat.ufl +import ufl + +__all__ = ("as_fiat_cell", "create_base_element", + "create_element", "supported_elements") + + +supported_elements = { + # These all map directly to FInAT elements + "Bernstein": finat.Bernstein, + "Bernardi-Raugel": finat.BernardiRaugel, + "Bernardi-Raugel Bubble": finat.BernardiRaugelBubble, + "Brezzi-Douglas-Marini": finat.BrezziDouglasMarini, + "Brezzi-Douglas-Fortin-Marini": finat.BrezziDouglasFortinMarini, + "Bubble": finat.Bubble, + "FacetBubble": finat.FacetBubble, + "Crouzeix-Raviart": finat.CrouzeixRaviart, + "Discontinuous Lagrange": finat.DiscontinuousLagrange, + "Discontinuous Raviart-Thomas": lambda c, d: finat.DiscontinuousElement(finat.RaviartThomas(c, d)), + "Discontinuous Taylor": finat.DiscontinuousTaylor, + "Gauss-Legendre": finat.GaussLegendre, + "Gauss-Lobatto-Legendre": finat.GaussLobattoLegendre, + "HDiv Trace": finat.HDivTrace, + "Hellan-Herrmann-Johnson": finat.HellanHerrmannJohnson, + "Johnson-Mercier": finat.JohnsonMercier, + "Nonconforming Arnold-Winther": finat.ArnoldWintherNC, + "Conforming Arnold-Winther": finat.ArnoldWinther, + "Hu-Zhang": finat.HuZhang, + "Hermite": finat.Hermite, + "Kong-Mulder-Veldhuizen": finat.KongMulderVeldhuizen, + "Argyris": finat.Argyris, + "Hsieh-Clough-Tocher": finat.HsiehCloughTocher, + "QuadraticPowellSabin6": finat.QuadraticPowellSabin6, + "QuadraticPowellSabin12": finat.QuadraticPowellSabin12, + "Reduced-Hsieh-Clough-Tocher": finat.ReducedHsiehCloughTocher, + "Mardal-Tai-Winther": finat.MardalTaiWinther, + "Alfeld-Sorokina": finat.AlfeldSorokina, + "Arnold-Qin": finat.ArnoldQin, + "Reduced-Arnold-Qin": finat.ReducedArnoldQin, + "Christiansen-Hu": finat.ChristiansenHu, + "Guzman-Neilan 1st kind H1": finat.GuzmanNeilanFirstKindH1, + "Guzman-Neilan 2nd kind H1": finat.GuzmanNeilanSecondKindH1, + "Guzman-Neilan Bubble": finat.GuzmanNeilanBubble, + "Guzman-Neilan H1(div)": finat.GuzmanNeilanH1div, + "Morley": finat.Morley, + "Bell": finat.Bell, + "Lagrange": finat.Lagrange, + "Nedelec 1st kind H(curl)": finat.Nedelec, + "Nedelec 2nd kind H(curl)": finat.NedelecSecondKind, + "Raviart-Thomas": finat.RaviartThomas, + "Regge": finat.Regge, + "Gopalakrishnan-Lederer-Schoberl 1st kind": finat.GopalakrishnanLedererSchoberlFirstKind, + "Gopalakrishnan-Lederer-Schoberl 2nd kind": finat.GopalakrishnanLedererSchoberlSecondKind, + "BDMCE": finat.BrezziDouglasMariniCubeEdge, + "BDMCF": finat.BrezziDouglasMariniCubeFace, + # These require special treatment below + "DQ": None, + "Q": None, + "RTCE": None, + "RTCF": None, + "NCE": None, + "NCF": None, + "Real": finat.Real, + "DPC": finat.DPC, + "S": finat.Serendipity, + "SminusF": finat.TrimmedSerendipityFace, + "SminusDiv": finat.TrimmedSerendipityDiv, + "SminusE": finat.TrimmedSerendipityEdge, + "SminusCurl": finat.TrimmedSerendipityCurl, + "DPC L2": finat.DPC, + "Discontinuous Lagrange L2": finat.DiscontinuousLagrange, + "Gauss-Legendre L2": finat.GaussLegendre, + "DQ L2": None, + "Direct Serendipity": finat.DirectSerendipity, +} +"""A :class:`.dict` mapping UFL element family names to their +FInAT-equivalent constructors. If the value is ``None``, the UFL +element is supported, but must be handled specially because it doesn't +have a direct FInAT equivalent.""" + + +def as_fiat_cell(cell): + """Convert a ufl cell to a FIAT cell. + + :arg cell: the :class:`ufl.Cell` to convert.""" + if not isinstance(cell, ufl.AbstractCell): + raise ValueError("Expecting a UFL Cell") + return FIAT.ufc_cell(cell) + + +@singledispatch +def convert(element, **kwargs): + """Handler for converting UFL elements to FInAT elements. + + :arg element: The UFL element to convert. + + Do not use this function directly, instead call + :func:`create_element`.""" + if element.family() in supported_elements: + raise ValueError("Element %s supported, but no handler provided" % element) + raise ValueError("Unsupported element type %s" % type(element)) + + +cg_interval_variants = { + "fdm": finat.FDMLagrange, + "fdm_ipdg": finat.FDMLagrange, + "fdm_quadrature": finat.FDMQuadrature, + "fdm_broken": finat.FDMBrokenH1, + "fdm_hermite": finat.FDMHermite, +} + + +dg_interval_variants = { + "fdm": finat.FDMDiscontinuousLagrange, + "fdm_quadrature": finat.FDMDiscontinuousLagrange, + "fdm_ipdg": lambda *args: finat.DiscontinuousElement(finat.FDMLagrange(*args)), + "fdm_broken": finat.FDMBrokenL2, +} + + +# Base finite elements first +@convert.register(finat.ufl.FiniteElement) +def convert_finiteelement(element, **kwargs): + cell = as_fiat_cell(element.cell) + if element.family() == "Quadrature": + degree = element.degree() + scheme = element.quadrature_scheme() + if degree is None or scheme is None: + raise ValueError("Quadrature scheme and degree must be specified!") + + return finat.make_quadrature_element(cell, degree, scheme), set() + lmbda = supported_elements[element.family()] + if element.family() == "Real" and element.cell.cellname() in {"quadrilateral", "hexahedron"}: + lmbda = None + element = finat.ufl.FiniteElement("DQ", element.cell, 0) + if lmbda is None: + if element.cell.cellname() == "quadrilateral": + # Handle quadrilateral short names like RTCF and RTCE. + element = element.reconstruct(cell=quadrilateral_tpc) + elif element.cell.cellname() == "hexahedron": + # Handle hexahedron short names like NCF and NCE. + element = element.reconstruct(cell=hexahedron_tpc) + else: + raise ValueError("%s is supported, but handled incorrectly" % + element.family()) + finat_elem, deps = _create_element(element, **kwargs) + return finat.FlattenedDimensions(finat_elem), deps + + finat_kwargs = {} + kind = element.variant() + if kind is None: + kind = 'spectral' # default variant + + if element.family() == "Lagrange": + if kind == 'spectral': + lmbda = finat.GaussLobattoLegendre + elif element.cell.cellname() == "interval" and kind in cg_interval_variants: + lmbda = cg_interval_variants[kind] + elif any(map(kind.startswith, ['integral', 'demkowicz', 'fdm'])): + lmbda = finat.IntegratedLegendre + finat_kwargs["variant"] = kind + elif kind in ['mgd', 'feec', 'qb', 'mse']: + degree = element.degree() + shift_axes = kwargs["shift_axes"] + restriction = kwargs["restriction"] + deps = {"shift_axes", "restriction"} + return finat.RuntimeTabulated(cell, degree, variant=kind, shift_axes=shift_axes, restriction=restriction), deps + else: + # Let FIAT handle the general case + lmbda = finat.Lagrange + finat_kwargs["variant"] = kind + + elif element.family() in ["Discontinuous Lagrange", "Discontinuous Lagrange L2"]: + if kind == 'spectral': + lmbda = finat.GaussLegendre + elif element.cell.cellname() == "interval" and kind in dg_interval_variants: + lmbda = dg_interval_variants[kind] + elif any(map(kind.startswith, ['integral', 'demkowicz', 'fdm'])): + lmbda = finat.Legendre + finat_kwargs["variant"] = kind + elif kind in ['mgd', 'feec', 'qb', 'mse']: + degree = element.degree() + shift_axes = kwargs["shift_axes"] + restriction = kwargs["restriction"] + deps = {"shift_axes", "restriction"} + return finat.RuntimeTabulated(cell, degree, variant=kind, shift_axes=shift_axes, restriction=restriction, continuous=False), deps + else: + # Let FIAT handle the general case + lmbda = finat.DiscontinuousLagrange + finat_kwargs["variant"] = kind + + elif element.variant() is not None: + finat_kwargs["variant"] = element.variant() + + return lmbda(cell, element.degree(), **finat_kwargs), set() + + +# Element modifiers and compound element types +@convert.register(finat.ufl.BrokenElement) +def convert_brokenelement(element, **kwargs): + finat_elem, deps = _create_element(element._element, **kwargs) + return finat.DiscontinuousElement(finat_elem), deps + + +@convert.register(finat.ufl.EnrichedElement) +def convert_enrichedelement(element, **kwargs): + elements, deps = zip(*[_create_element(elem, **kwargs) + for elem in element._elements]) + return finat.EnrichedElement(elements), set.union(*deps) + + +@convert.register(finat.ufl.NodalEnrichedElement) +def convert_nodalenrichedelement(element, **kwargs): + elements, deps = zip(*[_create_element(elem, **kwargs) + for elem in element._elements]) + return finat.NodalEnrichedElement(elements), set.union(*deps) + + +@convert.register(finat.ufl.MixedElement) +def convert_mixedelement(element, **kwargs): + elements, deps = zip(*[_create_element(elem, **kwargs) + for elem in element.sub_elements]) + return finat.MixedElement(elements), set.union(*deps) + + +@convert.register(finat.ufl.VectorElement) +@convert.register(finat.ufl.TensorElement) +def convert_tensorelement(element, **kwargs): + inner_elem, deps = _create_element(element.sub_elements[0], **kwargs) + shape = element.reference_value_shape + shape = shape[:len(shape) - len(inner_elem.value_shape)] + shape_innermost = kwargs["shape_innermost"] + return (finat.TensorFiniteElement(inner_elem, shape, not shape_innermost), + deps | {"shape_innermost"}) + + +@convert.register(finat.ufl.TensorProductElement) +def convert_tensorproductelement(element, **kwargs): + cell = element.cell + if type(cell) is not ufl.TensorProductCell: + raise ValueError("TensorProductElement not on TensorProductCell?") + shift_axes = kwargs["shift_axes"] + dim_offset = 0 + elements = [] + deps = set() + for elem in element.sub_elements: + kwargs["shift_axes"] = shift_axes + dim_offset + dim_offset += elem.cell.topological_dimension() + finat_elem, ds = _create_element(elem, **kwargs) + elements.append(finat_elem) + deps.update(ds) + return finat.TensorProductElement(elements), deps + + +@convert.register(finat.ufl.HDivElement) +def convert_hdivelement(element, **kwargs): + finat_elem, deps = _create_element(element._element, **kwargs) + return finat.HDivElement(finat_elem), deps + + +@convert.register(finat.ufl.HCurlElement) +def convert_hcurlelement(element, **kwargs): + finat_elem, deps = _create_element(element._element, **kwargs) + return finat.HCurlElement(finat_elem), deps + + +@convert.register(finat.ufl.WithMapping) +def convert_withmapping(element, **kwargs): + return _create_element(element.wrapee, **kwargs) + + +@convert.register(finat.ufl.RestrictedElement) +def convert_restrictedelement(element, **kwargs): + finat_elem, deps = _create_element(element._element, **kwargs) + return finat.RestrictedElement(finat_elem, element.restriction_domain()), deps + + +hexahedron_tpc = ufl.TensorProductCell(ufl.interval, ufl.interval, ufl.interval) +quadrilateral_tpc = ufl.TensorProductCell(ufl.interval, ufl.interval) +_cache = weakref.WeakKeyDictionary() + + +def create_element(ufl_element, shape_innermost=True, shift_axes=0, restriction=None): + """Create a FInAT element (suitable for tabulating with) given a UFL element. + + :arg ufl_element: The UFL element to create a FInAT element from. + :arg shape_innermost: Vector/tensor indices come after basis function indices + :arg restriction: cell restriction in interior facet integrals + (only for runtime tabulated elements) + """ + finat_element, deps = _create_element(ufl_element, + shape_innermost=shape_innermost, + shift_axes=shift_axes, + restriction=restriction) + return finat_element + + +def _create_element(ufl_element, **kwargs): + """A caching wrapper around :py:func:`convert`. + + Takes a UFL element and an unspecified set of parameter options, + and returns the converted element with the set of keyword names + that were relevant for conversion. + """ + # Look up conversion in cache + try: + cache = _cache[ufl_element] + except KeyError: + _cache[ufl_element] = {} + cache = _cache[ufl_element] + + for key, finat_element in cache.items(): + # Cache hit if all relevant parameter values match. + if all(kwargs[param] == value for param, value in key): + return finat_element, set(param for param, value in key) + + # Convert if cache miss + if ufl_element.cell is None: + raise ValueError("Don't know how to build element when cell is not given") + + finat_element, deps = convert(ufl_element, **kwargs) + + # Store conversion in cache + key = frozenset((param, kwargs[param]) for param in deps) + cache[key] = finat_element + + # Forward result + return finat_element, deps + + +def create_base_element(ufl_element, **kwargs): + """Create a "scalar" base FInAT element given a UFL element. + Takes a UFL element and an unspecified set of parameter options, + and returns the converted element. + """ + finat_element = create_element(ufl_element, **kwargs) + if isinstance(finat_element, finat.TensorFiniteElement): + finat_element = finat_element.base_element + return finat_element diff --git a/tsfc/kernel_args.py b/tsfc/kernel_args.py new file mode 100644 index 0000000000..a397f0f937 --- /dev/null +++ b/tsfc/kernel_args.py @@ -0,0 +1,62 @@ +import abc + + +class KernelArg(abc.ABC): + """Abstract base class wrapping a loopy argument. + + Defining this type system allows Firedrake (or other libraries) to + prepare arguments for the kernel without needing to worry about their + ordering. Instead this can be offloaded to tools such as + :func:`functools.singledispatch`. + """ + + def __init__(self, arg): + self.loopy_arg = arg + + @property + def dtype(self): + return self.loopy_arg.dtype + + +class OutputKernelArg(KernelArg): + ... + + +class CoordinatesKernelArg(KernelArg): + ... + + +class CoefficientKernelArg(KernelArg): + ... + + +class ConstantKernelArg(KernelArg): + ... + + +class CellOrientationsKernelArg(KernelArg): + ... + + +class CellSizesKernelArg(KernelArg): + ... + + +class TabulationKernelArg(KernelArg): + ... + + +class ExteriorFacetKernelArg(KernelArg): + ... + + +class InteriorFacetKernelArg(KernelArg): + ... + + +class ExteriorFacetOrientationKernelArg(KernelArg): + ... + + +class InteriorFacetOrientationKernelArg(KernelArg): + ... diff --git a/tsfc/kernel_interface/__init__.py b/tsfc/kernel_interface/__init__.py new file mode 100644 index 0000000000..5114263848 --- /dev/null +++ b/tsfc/kernel_interface/__init__.py @@ -0,0 +1,51 @@ +from abc import ABCMeta, abstractmethod, abstractproperty + +from gem.utils import make_proxy_class + + +class KernelInterface(metaclass=ABCMeta): + """Abstract interface for accessing the GEM expressions corresponding + to kernel arguments.""" + + @abstractmethod + def coordinate(self, ufl_domain): + """A function that maps :class:`ufl.Domain`s to coordinate + :class:`ufl.Coefficient`s.""" + + @abstractmethod + def coefficient(self, ufl_coefficient, restriction): + """A function that maps :class:`ufl.Coefficient`s to GEM + expressions.""" + + @abstractmethod + def constant(self, const): + """Return the GEM expression corresponding to the constant.""" + + @abstractmethod + def cell_orientation(self, restriction): + """Cell orientation as a GEM expression.""" + + @abstractmethod + def cell_size(self, restriction): + """Mesh cell size as a GEM expression. Shape (nvertex, ) in FIAT vertex ordering.""" + + @abstractmethod + def entity_number(self, restriction): + """Facet or vertex number as a GEM index.""" + + @abstractmethod + def entity_orientation(self, restriction): + """Entity orientation as a GEM index.""" + + @abstractmethod + def create_element(self, element, **kwargs): + """Create a FInAT element (suitable for tabulating with) given + a UFL element.""" + + @abstractproperty + def unsummed_coefficient_indices(self): + """A set of indices that coefficient evaluation should not sum over. + Used for macro-cell integration.""" + + +ProxyKernelInterface = make_proxy_class('ProxyKernelInterface', KernelInterface) diff --git a/tsfc/kernel_interface/common.py b/tsfc/kernel_interface/common.py new file mode 100644 index 0000000000..df7e879f09 --- /dev/null +++ b/tsfc/kernel_interface/common.py @@ -0,0 +1,545 @@ +import collections +import operator +import string +from functools import reduce +from itertools import chain, product + +import gem +import gem.impero_utils as impero_utils +import numpy +from FIAT.reference_element import TensorProductCell +from finat.cell_tools import max_complex +from finat.quadrature import AbstractQuadratureRule, make_quadrature +from gem.node import traversal +from gem.optimise import constant_fold_zero +from gem.optimise import remove_componenttensors as prune +from gem.utils import cached_property +from numpy import asarray +from tsfc import fem, ufl_utils +from tsfc.finatinterface import as_fiat_cell, create_element +from tsfc.kernel_interface import KernelInterface +from tsfc.logging import logger +from ufl.utils.sequences import max_degree + + +class KernelBuilderBase(KernelInterface): + """Helper class for building local assembly kernels.""" + + def __init__(self, scalar_type, interior_facet=False): + """Initialise a kernel builder. + + :arg interior_facet: kernel accesses two cells + """ + assert isinstance(interior_facet, bool) + self.scalar_type = scalar_type + self.interior_facet = interior_facet + + self.prepare = [] + self.finalise = [] + + # Coordinates + self.domain_coordinate = {} + + # Coefficients + self.coefficient_map = collections.OrderedDict() + + # Constants + self.constant_map = collections.OrderedDict() + + @cached_property + def unsummed_coefficient_indices(self): + return frozenset() + + def coordinate(self, domain): + return self.domain_coordinate[domain] + + def coefficient(self, ufl_coefficient, restriction): + """A function that maps :class:`ufl.Coefficient`s to GEM + expressions.""" + kernel_arg = self.coefficient_map[ufl_coefficient] + if ufl_coefficient.ufl_element().family() == 'Real': + return kernel_arg + elif not self.interior_facet: + return kernel_arg + else: + return kernel_arg[{'+': 0, '-': 1}[restriction]] + + def constant(self, const): + return self.constant_map[const] + + def cell_orientation(self, restriction): + """Cell orientation as a GEM expression.""" + f = {None: 0, '+': 0, '-': 1}[restriction] + # Assume self._cell_orientations tuple is set up at this point. + co_int = self._cell_orientations[f] + return gem.Conditional(gem.Comparison("==", co_int, gem.Literal(1)), + gem.Literal(-1), + gem.Conditional(gem.Comparison("==", co_int, gem.Zero()), + gem.Literal(1), + gem.Literal(numpy.nan))) + + def cell_size(self, restriction): + if not hasattr(self, "_cell_sizes"): + raise RuntimeError("Haven't called set_cell_sizes") + if self.interior_facet: + return self._cell_sizes[{'+': 0, '-': 1}[restriction]] + else: + return self._cell_sizes + + def entity_number(self, restriction): + """Facet or vertex number as a GEM index.""" + # Assume self._entity_number dict is set up at this point. + return self._entity_number[restriction] + + def entity_orientation(self, restriction): + """Facet orientation as a GEM index.""" + # Assume self._entity_orientation dict is set up at this point. + return self._entity_orientation[restriction] + + def apply_glue(self, prepare=None, finalise=None): + """Append glue code for operations that are not handled in the + GEM abstraction. + + Current uses: mixed interior facet mess + + :arg prepare: code snippets to be prepended to the kernel + :arg finalise: code snippets to be appended to the kernel + """ + if prepare is not None: + self.prepare.extend(prepare) + if finalise is not None: + self.finalise.extend(finalise) + + def register_requirements(self, ir): + """Inspect what is referenced by the IR that needs to be + provided by the kernel interface. + + :arg ir: multi-root GEM expression DAG + """ + # Nothing is required by default + pass + + +class KernelBuilderMixin(object): + """Mixin for KernelBuilder classes.""" + + def compile_integrand(self, integrand, params, ctx): + """Compile UFL integrand. + + :arg integrand: UFL integrand. + :arg params: a dict containing "quadrature_rule". + :arg ctx: context created with :meth:`create_context` method. + + See :meth:`create_context` for typical calling sequence. + """ + # Split Coefficients + if self.coefficient_split: + integrand = ufl_utils.split_coefficients(integrand, self.coefficient_split) + # Compile: ufl -> gem + info = self.integral_data_info + functions = list(info.arguments) + [self.coordinate(info.domain)] + list(info.coefficients) + set_quad_rule(params, info.domain.ufl_cell(), info.integral_type, functions) + quad_rule = params["quadrature_rule"] + config = self.fem_config() + config['argument_multiindices'] = self.argument_multiindices + config['quadrature_rule'] = quad_rule + config['index_cache'] = ctx['index_cache'] + expressions = fem.compile_ufl(integrand, + fem.PointSetContext(**config), + interior_facet=self.interior_facet) + ctx['quadrature_indices'].extend(quad_rule.point_set.indices) + return expressions + + def construct_integrals(self, integrand_expressions, params): + """Construct integrals from integrand expressions. + + :arg integrand_expressions: gem expressions for integrands. + :arg params: a dict containing "mode" and "quadrature_rule". + + integrand_expressions must be indexed with :attr:`argument_multiindices`; + these gem expressions are obtained by calling :meth:`compile_integrand` + method or by modifying the gem expressions returned by + :meth:`compile_integrand`. + + See :meth:`create_context` for typical calling sequence. + """ + mode = pick_mode(params["mode"]) + return mode.Integrals(integrand_expressions, + params["quadrature_rule"].point_set.indices, + self.argument_multiindices, + params) + + def stash_integrals(self, reps, params, ctx): + """Stash integral representations in ctx. + + :arg reps: integral representations. + :arg params: a dict containing "mode". + :arg ctx: context in which reps are stored. + + See :meth:`create_context` for typical calling sequence. + """ + mode = pick_mode(params["mode"]) + mode_irs = ctx['mode_irs'] + mode_irs.setdefault(mode, collections.OrderedDict()) + for var, rep in zip(self.return_variables, reps): + mode_irs[mode].setdefault(var, []).append(rep) + + def compile_gem(self, ctx): + """Compile gem representation of integrals to impero_c. + + :arg ctx: the context containing the gem representation of integrals. + :returns: a tuple of impero_c, oriented, needs_cell_sizes, tabulations + required to finally construct a kernel in :meth:`construct_kernel`. + + See :meth:`create_context` for typical calling sequence. + """ + # Finalise mode representations into a set of assignments + mode_irs = ctx['mode_irs'] + + assignments = [] + for mode, var_reps in mode_irs.items(): + assignments.extend(mode.flatten(var_reps.items(), ctx['index_cache'])) + + if assignments: + return_variables, expressions = zip(*assignments) + else: + return_variables = [] + expressions = [] + expressions = constant_fold_zero(expressions) + + # Need optimised roots + options = dict(reduce(operator.and_, + [mode.finalise_options.items() + for mode in mode_irs.keys()])) + expressions = impero_utils.preprocess_gem(expressions, **options) + + # Let the kernel interface inspect the optimised IR to register + # what kind of external data is required (e.g., cell orientations, + # cell sizes, etc.). + oriented, needs_cell_sizes, tabulations = self.register_requirements(expressions) + + # Extract Variables that are actually used + active_variables = gem.extract_type(expressions, gem.Variable) + # Construct ImperoC + assignments = list(zip(return_variables, expressions)) + index_ordering = get_index_ordering(ctx['quadrature_indices'], return_variables) + try: + impero_c = impero_utils.compile_gem(assignments, index_ordering, remove_zeros=True) + except impero_utils.NoopError: + impero_c = None + return impero_c, oriented, needs_cell_sizes, tabulations, active_variables + + def fem_config(self): + """Return a dictionary used with fem.compile_ufl. + + One needs to update this dictionary with "argument_multiindices", + "quadrature_rule", and "index_cache" before using this with + fem.compile_ufl. + """ + info = self.integral_data_info + integral_type = info.integral_type + cell = info.domain.ufl_cell() + fiat_cell = as_fiat_cell(cell) + integration_dim, entity_ids = lower_integral_type(fiat_cell, integral_type) + return dict(interface=self, + ufl_cell=cell, + integral_type=integral_type, + integration_dim=integration_dim, + entity_ids=entity_ids, + scalar_type=self.fem_scalar_type) + + def create_context(self): + """Create builder context. + + *index_cache* + + Map from UFL FiniteElement objects to multiindices. + This is so we reuse Index instances when evaluating the same + coefficient multiple times with the same table. + + We also use the same dict for the unconcatenate index cache, + which maps index objects to tuples of multiindices. These two + caches shall never conflict as their keys have different types + (UFL finite elements vs. GEM index objects). + + *quadrature_indices* + + List of quadrature indices used. + + *mode_irs* + + Dict for mode representations. + + For each set of integrals to make a kernel for (i,e., + `integral_data.integrals`), one must first create a ctx object by + calling :meth:`create_context` method. + This ctx object collects objects associated with the integrals that + are eventually used to construct the kernel. + The following is a typical calling sequence: + + .. code-block:: python3 + + builder = KernelBuilder(...) + params = {"mode": "spectral"} + ctx = builder.create_context() + for integral in integral_data.integrals: + integrand = integral.integrand() + integrand_exprs = builder.compile_integrand(integrand, params, ctx) + integral_exprs = builder.construct_integrals(integrand_exprs, params) + builder.stash_integrals(integral_exprs, params, ctx) + kernel = builder.construct_kernel(kernel_name, ctx) + + """ + return {'index_cache': {}, + 'quadrature_indices': [], + 'mode_irs': collections.OrderedDict()} + + +def set_quad_rule(params, cell, integral_type, functions): + # Check if the integral has a quad degree attached, otherwise use + # the estimated polynomial degree attached by compute_form_data + try: + quadrature_degree = params["quadrature_degree"] + except KeyError: + quadrature_degree = params["estimated_polynomial_degree"] + function_degrees = [f.ufl_function_space().ufl_element().degree() + for f in functions] + if all((asarray(quadrature_degree) > 10 * asarray(degree)).all() + for degree in function_degrees): + logger.warning("Estimated quadrature degree %s more " + "than tenfold greater than any " + "argument/coefficient degree (max %s)", + quadrature_degree, max_degree(function_degrees)) + quad_rule = params.get("quadrature_rule", "default") + if isinstance(quad_rule, str): + scheme = quad_rule + fiat_cell = as_fiat_cell(cell) + finat_elements = set(create_element(f.ufl_element()) for f in functions + if f.ufl_element().family() != "Real") + fiat_cells = [fiat_cell] + [finat_el.complex for finat_el in finat_elements] + fiat_cell = max_complex(fiat_cells) + + integration_dim, _ = lower_integral_type(fiat_cell, integral_type) + integration_cell = fiat_cell.construct_subcomplex(integration_dim) + quad_rule = make_quadrature(integration_cell, quadrature_degree, scheme=scheme) + params["quadrature_rule"] = quad_rule + + if not isinstance(quad_rule, AbstractQuadratureRule): + raise ValueError("Expected to find a QuadratureRule object, not a %s" % + type(quad_rule)) + + +def get_index_ordering(quadrature_indices, return_variables): + split_argument_indices = tuple(chain(*[var.index_ordering() + for var in return_variables])) + return tuple(quadrature_indices) + split_argument_indices + + +def get_index_names(quadrature_indices, argument_multiindices, index_cache): + index_names = [] + + def name_index(index, name): + index_names.append((index, name)) + if index in index_cache: + for multiindex, suffix in zip(index_cache[index], + string.ascii_lowercase): + name_multiindex(multiindex, name + suffix) + + def name_multiindex(multiindex, name): + if len(multiindex) == 1: + name_index(multiindex[0], name) + else: + for i, index in enumerate(multiindex): + name_index(index, name + str(i)) + + name_multiindex(quadrature_indices, 'ip') + for multiindex, name in zip(argument_multiindices, ['j', 'k']): + name_multiindex(multiindex, name) + return index_names + + +def lower_integral_type(fiat_cell, integral_type): + """Lower integral type into the dimension of the integration + subentity and a list of entity numbers for that dimension. + + :arg fiat_cell: FIAT reference cell + :arg integral_type: integral type (string) + """ + vert_facet_types = ['exterior_facet_vert', 'interior_facet_vert'] + horiz_facet_types = ['exterior_facet_bottom', 'exterior_facet_top', 'interior_facet_horiz'] + + dim = fiat_cell.get_dimension() + if integral_type == 'cell': + integration_dim = dim + elif integral_type in ['exterior_facet', 'interior_facet']: + if isinstance(fiat_cell, TensorProductCell): + raise ValueError("{} integral cannot be used with a TensorProductCell; need to distinguish between vertical and horizontal contributions.".format(integral_type)) + integration_dim = dim - 1 + elif integral_type == 'vertex': + integration_dim = 0 + elif integral_type in vert_facet_types + horiz_facet_types: + # Extrusion case + if not isinstance(fiat_cell, TensorProductCell): + raise ValueError("{} integral requires a TensorProductCell.".format(integral_type)) + basedim, extrdim = dim + assert extrdim == 1 + + if integral_type in vert_facet_types: + integration_dim = (basedim - 1, 1) + elif integral_type in horiz_facet_types: + integration_dim = (basedim, 0) + else: + raise NotImplementedError("integral type %s not supported" % integral_type) + + if integral_type == 'exterior_facet_bottom': + entity_ids = [0] + elif integral_type == 'exterior_facet_top': + entity_ids = [1] + else: + entity_ids = list(range(len(fiat_cell.get_topology()[integration_dim]))) + + return integration_dim, entity_ids + + +def pick_mode(mode): + "Return one of the specialized optimisation modules from a mode string." + try: + from firedrake_citations import Citations + cites = {"vanilla": ("Homolya2017", ), + "coffee": ("Luporini2016", "Homolya2017", ), + "spectral": ("Luporini2016", "Homolya2017", "Homolya2017a"), + "tensor": ("Kirby2006", "Homolya2017", )} + for c in cites[mode]: + Citations().register(c) + except ImportError: + pass + if mode == "vanilla": + import tsfc.vanilla as m + elif mode == "coffee": + import tsfc.coffee_mode as m + elif mode == "spectral": + import tsfc.spectral as m + elif mode == "tensor": + import tsfc.tensor as m + else: + raise ValueError("Unknown mode: {}".format(mode)) + return m + + +def check_requirements(ir): + """Look for cell orientations, cell sizes, and collect tabulations + in one pass.""" + cell_orientations = False + cell_sizes = False + rt_tabs = {} + for node in traversal(ir): + if isinstance(node, gem.Variable): + if node.name == "cell_orientations": + cell_orientations = True + elif node.name == "cell_sizes": + cell_sizes = True + elif node.name.startswith("rt_"): + rt_tabs[node.name] = node.shape + return cell_orientations, cell_sizes, tuple(sorted(rt_tabs.items())) + + +def prepare_constant(constant, number): + """Bridges the kernel interface and the GEM abstraction for + Constants. + + :arg constant: Firedrake Constant + :arg number: Value to uniquely identify the constant + :returns: (funarg, expression) + expression - GEM expression referring to the Constant value(s) + """ + value_size = numpy.prod(constant.ufl_shape, dtype=int) + return gem.reshape(gem.Variable(f"c_{number}", (value_size,)), + constant.ufl_shape) + + +def prepare_coefficient(coefficient, name, interior_facet=False): + """Bridges the kernel interface and the GEM abstraction for + Coefficients. + + :arg coefficient: UFL Coefficient + :arg name: unique name to refer to the Coefficient in the kernel + :arg interior_facet: interior facet integral? + :returns: (funarg, expression) + expression - GEM expression referring to the Coefficient + values + """ + assert isinstance(interior_facet, bool) + + if coefficient.ufl_element().family() == 'Real': + # Constant + value_size = coefficient.ufl_function_space().value_size + expression = gem.reshape(gem.Variable(name, (value_size,)), + coefficient.ufl_shape) + return expression + + finat_element = create_element(coefficient.ufl_element()) + shape = finat_element.index_shape + size = numpy.prod(shape, dtype=int) + + if not interior_facet: + expression = gem.reshape(gem.Variable(name, (size,)), shape) + else: + varexp = gem.Variable(name, (2 * size,)) + plus = gem.view(varexp, slice(size)) + minus = gem.view(varexp, slice(size, 2 * size)) + expression = (gem.reshape(plus, shape), gem.reshape(minus, shape)) + return expression + + +def prepare_arguments(arguments, multiindices, interior_facet=False, diagonal=False): + """Bridges the kernel interface and the GEM abstraction for + Arguments. Vector Arguments are rearranged here for interior + facet integrals. + + :arg arguments: UFL Arguments + :arg multiindices: Argument multiindices + :arg interior_facet: interior facet integral? + :arg diagonal: Are we assembling the diagonal of a rank-2 element tensor? + :returns: (funarg, expression) + expressions - GEM expressions referring to the argument + tensor + """ + assert isinstance(interior_facet, bool) + + if len(arguments) == 0: + # No arguments + expression = gem.Indexed(gem.Variable("A", (1,)), (0,)) + return (expression, ) + + elements = tuple(create_element(arg.ufl_element()) for arg in arguments) + shapes = tuple(element.index_shape for element in elements) + + if diagonal: + if len(arguments) != 2: + raise ValueError("Diagonal only for 2-forms") + try: + element, = set(elements) + except ValueError: + raise ValueError("Diagonal only for diagonal blocks (test and trial spaces the same)") + + elements = (element, ) + shapes = tuple(element.index_shape for element in elements) + multiindices = multiindices[:1] + + def expression(restricted): + return gem.Indexed(gem.reshape(restricted, *shapes), + tuple(chain(*multiindices))) + + u_shape = numpy.array([numpy.prod(shape, dtype=int) for shape in shapes]) + if interior_facet: + c_shape = tuple(2 * u_shape) + slicez = [[slice(r * s, (r + 1) * s) + for r, s in zip(restrictions, u_shape)] + for restrictions in product((0, 1), repeat=len(arguments))] + else: + c_shape = tuple(u_shape) + slicez = [[slice(s) for s in u_shape]] + + varexp = gem.Variable("A", c_shape) + expressions = [expression(gem.view(varexp, *slices)) for slices in slicez] + return tuple(prune(expressions)) diff --git a/tsfc/kernel_interface/firedrake_loopy.py b/tsfc/kernel_interface/firedrake_loopy.py new file mode 100644 index 0000000000..0969854cd3 --- /dev/null +++ b/tsfc/kernel_interface/firedrake_loopy.py @@ -0,0 +1,455 @@ +import numpy +from collections import namedtuple, OrderedDict +from functools import partial + +from ufl import Coefficient, FunctionSpace +from ufl.domain import extract_unique_domain +from finat.ufl import MixedElement as ufl_MixedElement, FiniteElement + +import gem +from gem.flop_count import count_flops + +import loopy as lp + +from tsfc import kernel_args, fem +from tsfc.finatinterface import create_element +from tsfc.kernel_interface.common import KernelBuilderBase as _KernelBuilderBase, KernelBuilderMixin, get_index_names, check_requirements, prepare_coefficient, prepare_arguments, prepare_constant +from tsfc.loopy import generate as generate_loopy + + +# Expression kernel description type +ExpressionKernel = namedtuple('ExpressionKernel', ['ast', 'oriented', 'needs_cell_sizes', + 'coefficient_numbers', + 'needs_external_coords', + 'tabulations', 'name', 'arguments', + 'flop_count', 'event']) + + +def make_builder(*args, **kwargs): + return partial(KernelBuilder, *args, **kwargs) + + +class Kernel: + __slots__ = ("ast", "arguments", "integral_type", "oriented", "subdomain_id", + "domain_number", "needs_cell_sizes", "tabulations", + "coefficient_numbers", "name", "flop_count", "event", + "__weakref__") + """A compiled Kernel object. + + :kwarg ast: The loopy kernel object. + :kwarg integral_type: The type of integral. + :kwarg oriented: Does the kernel require cell_orientations. + :kwarg subdomain_id: What is the subdomain id for this kernel. + :kwarg domain_number: Which domain number in the original form + does this kernel correspond to (can be used to index into + original_form.ufl_domains() to get the correct domain). + :kwarg coefficient_numbers: A list of which coefficients from the + form the kernel needs. + :kwarg tabulations: The runtime tabulations this kernel requires + :kwarg needs_cell_sizes: Does the kernel require cell sizes. + :kwarg name: The name of this kernel. + :kwarg flop_count: Estimated total flops for this kernel. + :kwarg event: name for logging event + """ + def __init__(self, ast=None, arguments=None, integral_type=None, oriented=False, + subdomain_id=None, domain_number=None, + coefficient_numbers=(), + needs_cell_sizes=False, + tabulations=None, + flop_count=0, + name=None, + event=None): + # Defaults + self.ast = ast + self.arguments = arguments + self.integral_type = integral_type + self.oriented = oriented + self.domain_number = domain_number + self.subdomain_id = subdomain_id + self.coefficient_numbers = coefficient_numbers + self.needs_cell_sizes = needs_cell_sizes + self.tabulations = tabulations + self.flop_count = flop_count + self.name = name + self.event = event + + +class KernelBuilderBase(_KernelBuilderBase): + + def __init__(self, scalar_type, interior_facet=False): + """Initialise a kernel builder. + + :arg interior_facet: kernel accesses two cells + """ + super().__init__(scalar_type=scalar_type, interior_facet=interior_facet) + + # Cell orientation + if self.interior_facet: + cell_orientations = gem.Variable("cell_orientations", (2,), dtype=gem.uint_type) + self._cell_orientations = (gem.Indexed(cell_orientations, (0,)), + gem.Indexed(cell_orientations, (1,))) + else: + cell_orientations = gem.Variable("cell_orientations", (1,), dtype=gem.uint_type) + self._cell_orientations = (gem.Indexed(cell_orientations, (0,)),) + + def _coefficient(self, coefficient, name): + """Prepare a coefficient. Adds glue code for the coefficient + and adds the coefficient to the coefficient map. + + :arg coefficient: :class:`ufl.Coefficient` + :arg name: coefficient name + :returns: GEM expression representing the coefficient + """ + expr = prepare_coefficient(coefficient, name, interior_facet=self.interior_facet) + self.coefficient_map[coefficient] = expr + return expr + + def set_cell_sizes(self, domain): + """Setup a fake coefficient for "cell sizes". + + :arg domain: The domain of the integral. + + This is required for scaling of derivative basis functions on + physically mapped elements (Argyris, Bell, etc...). We need a + measure of the mesh size around each vertex (hence this lives + in P1). + + Should the domain have topological dimension 0 this does + nothing. + """ + if domain.ufl_cell().topological_dimension() > 0: + # Can't create P1 since only P0 is a valid finite element if + # topological_dimension is 0 and the concept of "cell size" + # is not useful for a vertex. + f = Coefficient(FunctionSpace(domain, FiniteElement("P", domain.ufl_cell(), 1))) + expr = prepare_coefficient(f, "cell_sizes", interior_facet=self.interior_facet) + self._cell_sizes = expr + + def create_element(self, element, **kwargs): + """Create a FInAT element (suitable for tabulating with) given + a UFL element.""" + return create_element(element, **kwargs) + + def generate_arg_from_variable(self, var, dtype=None): + """Generate kernel arg from a :class:`gem.Variable`. + + :arg var: a :class:`gem.Variable` + :arg dtype: dtype of the kernel arg + :returns: kernel arg + """ + return lp.GlobalArg(var.name, dtype=dtype or self.scalar_type, shape=var.shape) + + def generate_arg_from_expression(self, expr, dtype=None): + """Generate kernel arg from gem expression(s). + + :arg expr: gem expression(s) representing a coefficient or the output tensor + :arg dtype: dtype of the kernel arg + :returns: kernel arg + """ + var, = gem.extract_type(expr if isinstance(expr, tuple) else (expr, ), gem.Variable) + return self.generate_arg_from_variable(var, dtype=dtype or self.scalar_type) + + +class ExpressionKernelBuilder(KernelBuilderBase): + """Builds expression kernels for UFL interpolation in Firedrake.""" + + def __init__(self, scalar_type): + super(ExpressionKernelBuilder, self).__init__(scalar_type=scalar_type) + self.oriented = False + self.cell_sizes = False + + def set_coefficients(self, coefficients): + """Prepare the coefficients of the expression. + + :arg coefficients: UFL coefficients from Firedrake + """ + self.coefficient_split = {} + + for i, coefficient in enumerate(coefficients): + if type(coefficient.ufl_element()) == ufl_MixedElement: + subcoeffs = coefficient.subfunctions # Firedrake-specific + self.coefficient_split[coefficient] = subcoeffs + for j, subcoeff in enumerate(subcoeffs): + self._coefficient(subcoeff, f"w_{i}_{j}") + else: + self._coefficient(coefficient, f"w_{i}") + + def set_constants(self, constants): + for i, const in enumerate(constants): + gemexpr = prepare_constant(const, i) + self.constant_map[const] = gemexpr + + def set_coefficient_numbers(self, coefficient_numbers): + """Store the coefficient indices of the original form. + + :arg coefficient_numbers: Iterable of indices describing which coefficients + from the input expression need to be passed in to the kernel. + """ + self.coefficient_numbers = coefficient_numbers + + def register_requirements(self, ir): + """Inspect what is referenced by the IR that needs to be + provided by the kernel interface.""" + self.oriented, self.cell_sizes, self.tabulations = check_requirements(ir) + + def set_output(self, o): + """Produce the kernel return argument""" + loopy_arg = lp.GlobalArg(o.name, dtype=self.scalar_type, shape=o.shape) + self.output_arg = kernel_args.OutputKernelArg(loopy_arg) + + def construct_kernel(self, impero_c, index_names, needs_external_coords, log=False): + """Constructs an :class:`ExpressionKernel`. + + :arg impero_c: gem.ImperoC object that represents the kernel + :arg index_names: pre-assigned index names + :arg needs_external_coords: If ``True``, the first argument to + the kernel is an externally provided coordinate field. + :arg log: bool if the Kernel should be profiled with Log events + + :returns: :class:`ExpressionKernel` object + """ + args = [self.output_arg] + if self.oriented: + funarg = self.generate_arg_from_expression(self._cell_orientations, dtype=numpy.int32) + args.append(kernel_args.CellOrientationsKernelArg(funarg)) + if self.cell_sizes: + funarg = self.generate_arg_from_expression(self._cell_sizes) + args.append(kernel_args.CellSizesKernelArg(funarg)) + for _, expr in self.coefficient_map.items(): + # coefficient_map is OrderedDict. + funarg = self.generate_arg_from_expression(expr) + args.append(kernel_args.CoefficientKernelArg(funarg)) + + # now constants + for gemexpr in self.constant_map.values(): + funarg = self.generate_arg_from_expression(gemexpr) + args.append(kernel_args.ConstantKernelArg(funarg)) + + for name_, shape in self.tabulations: + tab_loopy_arg = lp.GlobalArg(name_, dtype=self.scalar_type, shape=shape) + args.append(kernel_args.TabulationKernelArg(tab_loopy_arg)) + + loopy_args = [arg.loopy_arg for arg in args] + + name = "expression_kernel" + loopy_kernel, event = generate_loopy(impero_c, loopy_args, self.scalar_type, + name, index_names, log=log) + return ExpressionKernel(loopy_kernel, self.oriented, self.cell_sizes, + self.coefficient_numbers, needs_external_coords, + self.tabulations, name, args, count_flops(impero_c), event) + + +class KernelBuilder(KernelBuilderBase, KernelBuilderMixin): + """Helper class for building a :class:`Kernel` object.""" + + def __init__(self, integral_data_info, scalar_type, + dont_split=(), diagonal=False): + """Initialise a kernel builder.""" + integral_type = integral_data_info.integral_type + super(KernelBuilder, self).__init__(scalar_type, integral_type.startswith("interior_facet")) + self.fem_scalar_type = scalar_type + + self.diagonal = diagonal + self.local_tensor = None + self.coefficient_split = {} + self.coefficient_number_index_map = OrderedDict() + self.dont_split = frozenset(dont_split) + + # Facet number + if integral_type in ['exterior_facet', 'exterior_facet_vert']: + facet = gem.Variable('facet', (1,), dtype=gem.uint_type) + self._entity_number = {None: gem.VariableIndex(gem.Indexed(facet, (0,)))} + facet_orientation = gem.Variable('facet_orientation', (1,), dtype=gem.uint_type) + self._entity_orientation = {None: gem.OrientationVariableIndex(gem.Indexed(facet_orientation, (0,)))} + elif integral_type in ['interior_facet', 'interior_facet_vert']: + facet = gem.Variable('facet', (2,), dtype=gem.uint_type) + self._entity_number = { + '+': gem.VariableIndex(gem.Indexed(facet, (0,))), + '-': gem.VariableIndex(gem.Indexed(facet, (1,))) + } + facet_orientation = gem.Variable('facet_orientation', (2,), dtype=gem.uint_type) + self._entity_orientation = { + '+': gem.OrientationVariableIndex(gem.Indexed(facet_orientation, (0,))), + '-': gem.OrientationVariableIndex(gem.Indexed(facet_orientation, (1,))) + } + elif integral_type == 'interior_facet_horiz': + self._entity_number = {'+': 1, '-': 0} + facet_orientation = gem.Variable('facet_orientation', (1,), dtype=gem.uint_type) # base mesh entity orientation + self._entity_orientation = { + '+': gem.OrientationVariableIndex(gem.Indexed(facet_orientation, (0,))), + '-': gem.OrientationVariableIndex(gem.Indexed(facet_orientation, (0,))) + } + + self.set_arguments(integral_data_info.arguments) + self.integral_data_info = integral_data_info + + def set_arguments(self, arguments): + """Process arguments. + + :arg arguments: :class:`ufl.Argument`s + :returns: GEM expression representing the return variable + """ + argument_multiindices = tuple(create_element(arg.ufl_element()).get_indices() + for arg in arguments) + if self.diagonal: + # Error checking occurs in the builder constructor. + # Diagonal assembly is obtained by using the test indices for + # the trial space as well. + a, _ = argument_multiindices + argument_multiindices = (a, a) + return_variables = prepare_arguments(arguments, + argument_multiindices, + interior_facet=self.interior_facet, + diagonal=self.diagonal) + self.return_variables = return_variables + self.argument_multiindices = argument_multiindices + + def set_coordinates(self, domain): + """Prepare the coordinate field. + + :arg domain: :class:`ufl.Domain` + """ + # Create a fake coordinate coefficient for a domain. + f = Coefficient(FunctionSpace(domain, domain.ufl_coordinate_element())) + self.domain_coordinate[domain] = f + self._coefficient(f, "coords") + + def set_coefficients(self, integral_data, form_data): + """Prepare the coefficients of the form. + + :arg integral_data: UFL integral data + :arg form_data: UFL form data + """ + # enabled_coefficients is a boolean array that indicates which + # of reduced_coefficients the integral requires. + n, k = 0, 0 + for i in range(len(integral_data.enabled_coefficients)): + if integral_data.enabled_coefficients[i]: + original = form_data.reduced_coefficients[i] + coefficient = form_data.function_replace_map[original] + if type(coefficient.ufl_element()) == ufl_MixedElement: + if original in self.dont_split: + self.coefficient_split[coefficient] = [coefficient] + self._coefficient(coefficient, f"w_{k}") + self.coefficient_number_index_map[coefficient] = (n, 0) + k += 1 + else: + self.coefficient_split[coefficient] = [] + for j, element in enumerate(coefficient.ufl_element().sub_elements): + c = Coefficient(FunctionSpace(extract_unique_domain(coefficient), element)) + self.coefficient_split[coefficient].append(c) + self._coefficient(c, f"w_{k}") + self.coefficient_number_index_map[c] = (n, j) + k += 1 + else: + self._coefficient(coefficient, f"w_{k}") + self.coefficient_number_index_map[coefficient] = (n, 0) + k += 1 + n += 1 + + def set_constants(self, constants): + for i, const in enumerate(constants): + gemexpr = prepare_constant(const, i) + self.constant_map[const] = gemexpr + + def register_requirements(self, ir): + """Inspect what is referenced by the IR that needs to be + provided by the kernel interface.""" + return check_requirements(ir) + + def construct_kernel(self, name, ctx, log=False): + """Construct a fully built :class:`Kernel`. + + This function contains the logic for building the argument + list for assembly kernels. + + :arg name: kernel name + :arg ctx: kernel builder context to get impero_c from + :arg log: bool if the Kernel should be profiled with Log events + :returns: :class:`Kernel` object + """ + impero_c, oriented, needs_cell_sizes, tabulations, active_variables = self.compile_gem(ctx) + if impero_c is None: + return self.construct_empty_kernel(name) + info = self.integral_data_info + # In the following funargs are only generated + # for gem expressions that are actually used; + # see `generate_arg_from_expression()` method. + # Specifically, funargs are not generated for + # unused components of mixed coefficients. + # Problem solving environment, such as Firedrake, + # will know which components have been included + # in the list of kernel arguments by investigating + # `Kernel.coefficient_numbers`. + # Add return arg + funarg = self.generate_arg_from_expression(self.return_variables) + args = [kernel_args.OutputKernelArg(funarg)] + # Add coordinates arg + coord = self.domain_coordinate[info.domain] + expr = self.coefficient_map[coord] + funarg = self.generate_arg_from_expression(expr) + args.append(kernel_args.CoordinatesKernelArg(funarg)) + if oriented: + funarg = self.generate_arg_from_expression(self._cell_orientations, dtype=numpy.int32) + args.append(kernel_args.CellOrientationsKernelArg(funarg)) + if needs_cell_sizes: + funarg = self.generate_arg_from_expression(self._cell_sizes) + args.append(kernel_args.CellSizesKernelArg(funarg)) + coefficient_indices = OrderedDict() + for coeff, (number, index) in self.coefficient_number_index_map.items(): + a = coefficient_indices.setdefault(number, []) + expr = self.coefficient_map[coeff] + var, = gem.extract_type(expr if isinstance(expr, tuple) else (expr, ), gem.Variable) + if var in active_variables: + funarg = self.generate_arg_from_expression(expr) + args.append(kernel_args.CoefficientKernelArg(funarg)) + a.append(index) + + # now constants + for gemexpr in self.constant_map.values(): + funarg = self.generate_arg_from_expression(gemexpr) + args.append(kernel_args.ConstantKernelArg(funarg)) + + coefficient_indices = tuple(tuple(v) for v in coefficient_indices.values()) + assert len(coefficient_indices) == len(info.coefficient_numbers) + if info.integral_type in ["exterior_facet", "exterior_facet_vert"]: + ext_loopy_arg = lp.GlobalArg("facet", numpy.uint32, shape=(1,)) + args.append(kernel_args.ExteriorFacetKernelArg(ext_loopy_arg)) + elif info.integral_type in ["interior_facet", "interior_facet_vert"]: + int_loopy_arg = lp.GlobalArg("facet", numpy.uint32, shape=(2,)) + args.append(kernel_args.InteriorFacetKernelArg(int_loopy_arg)) + # Will generalise this in the submesh PR. + if fem.PointSetContext(**self.fem_config()).use_canonical_quadrature_point_ordering: + if info.integral_type == "exterior_facet": + ext_ornt_loopy_arg = lp.GlobalArg("facet_orientation", gem.uint_type, shape=(1,)) + args.append(kernel_args.ExteriorFacetOrientationKernelArg(ext_ornt_loopy_arg)) + elif info.integral_type == "interior_facet": + int_ornt_loopy_arg = lp.GlobalArg("facet_orientation", gem.uint_type, shape=(2,)) + args.append(kernel_args.InteriorFacetOrientationKernelArg(int_ornt_loopy_arg)) + for name_, shape in tabulations: + tab_loopy_arg = lp.GlobalArg(name_, dtype=self.scalar_type, shape=shape) + args.append(kernel_args.TabulationKernelArg(tab_loopy_arg)) + index_names = get_index_names(ctx['quadrature_indices'], self.argument_multiindices, ctx['index_cache']) + ast, event_name = generate_loopy(impero_c, [arg.loopy_arg for arg in args], + self.scalar_type, name, index_names, log=log) + flop_count = count_flops(impero_c) # Estimated total flops for this kernel. + return Kernel(ast=ast, + arguments=tuple(args), + integral_type=info.integral_type, + subdomain_id=info.subdomain_id, + domain_number=info.domain_number, + coefficient_numbers=tuple(zip(info.coefficient_numbers, coefficient_indices)), + oriented=oriented, + needs_cell_sizes=needs_cell_sizes, + tabulations=tabulations, + flop_count=flop_count, + name=name, + event=event_name) + + def construct_empty_kernel(self, name): + """Return None, since Firedrake needs no empty kernels. + + :arg name: function name + :returns: None + """ + return None diff --git a/tsfc/logging.py b/tsfc/logging.py new file mode 100644 index 0000000000..6bced525f4 --- /dev/null +++ b/tsfc/logging.py @@ -0,0 +1,6 @@ +"""Logging for TSFC.""" + +import logging + +logger = logging.getLogger('tsfc') +logger.addHandler(logging.StreamHandler()) diff --git a/tsfc/loopy.py b/tsfc/loopy.py new file mode 100644 index 0000000000..6826f0b672 --- /dev/null +++ b/tsfc/loopy.py @@ -0,0 +1,578 @@ +"""Generate loopy kernel from ImperoC tuple data. + +This is the final stage of code generation in TSFC.""" + +from numbers import Integral +import numpy +from functools import singledispatch +from collections import defaultdict, OrderedDict + +from gem import gem, impero as imp +from gem.node import Memoizer + +import islpy as isl +import loopy as lp + +import pymbolic.primitives as p +from loopy.symbolic import SubArrayRef + +from pytools import UniqueNameGenerator + +from tsfc.parameters import is_complex + +from contextlib import contextmanager +from tsfc.parameters import target + + +def profile_insns(kernel_name, instructions, log=False): + if log: + event_name = "Log_Event_" + kernel_name + event_id_var_name = "ID_" + event_name + # Logging registration + # The events are registered in PyOP2 and the event id is passed onto the dll + preamble = "PetscLogEvent "+event_id_var_name+" = -1;" + # Profiling + prepend = [lp.CInstruction("", "PetscLogEventBegin("+event_id_var_name+",0,0,0,0);")] + append = [lp.CInstruction("", "PetscLogEventEnd("+event_id_var_name+",0,0,0,0);")] + instructions = prepend + instructions + append + return instructions, event_name, [(str(2**31-1)+"_"+kernel_name, preamble)] + else: + return instructions, None, None + + +@singledispatch +def _assign_dtype(expression, self): + return set.union(*map(self, expression.children)) + + +@_assign_dtype.register(gem.Terminal) +def _assign_dtype_terminal(expression, self): + return {expression.dtype or self.scalar_type} + + +@_assign_dtype.register(gem.Variable) +def _assign_dtype_variable(expression, self): + return {expression.dtype or self.scalar_type} + + +@_assign_dtype.register(gem.Zero) +@_assign_dtype.register(gem.Identity) +@_assign_dtype.register(gem.Delta) +def _assign_dtype_real(expression, self): + return {expression.dtype or self.real_type} + + +@_assign_dtype.register(gem.Literal) +def _assign_dtype_identity(expression, self): + return {expression.array.dtype} + + +@_assign_dtype.register(gem.Power) +def _assign_dtype_power(expression, self): + # Conservative + return {expression.dtype or self.scalar_type} + + +@_assign_dtype.register(gem.MathFunction) +def _assign_dtype_mathfunction(expression, self): + if expression.name in {"abs", "real", "imag"}: + return {expression.dtype or self.real_type} + elif expression.name == "sqrt": + return {expression.dtype or self.scalar_type} + else: + return set.union(*map(self, expression.children)) + + +@_assign_dtype.register(gem.MinValue) +@_assign_dtype.register(gem.MaxValue) +def _assign_dtype_minmax(expression, self): + # UFL did correctness checking + return {expression.dtype or self.real_type} + + +@_assign_dtype.register(gem.Conditional) +def _assign_dtype_conditional(expression, self): + return set.union(*map(self, expression.children[1:])) + + +@_assign_dtype.register(gem.Comparison) +@_assign_dtype.register(gem.LogicalNot) +@_assign_dtype.register(gem.LogicalAnd) +@_assign_dtype.register(gem.LogicalOr) +def _assign_dtype_logical(expression, self): + return {expression.dtype or numpy.int8} + + +def assign_dtypes(expressions, scalar_type): + """Assign numpy data types to expressions. + + Used for declaring temporaries when converting from Impero to lower level code. + + :arg expressions: List of GEM expressions. + :arg scalar_type: Default scalar type. + + :returns: list of tuples (expression, dtype).""" + mapper = Memoizer(_assign_dtype) + mapper.scalar_type = scalar_type + mapper.real_type = numpy.finfo(scalar_type).dtype + return [(e, numpy.result_type(*mapper(e))) for e in expressions] + + +class LoopyContext(object): + def __init__(self, target=None): + self.indices = {} # indices for declarations and referencing values, from ImperoC + self.active_indices = {} # gem index -> pymbolic variable + self.index_extent = OrderedDict() # pymbolic variable for indices -> extent + self.gem_to_pymbolic = {} # gem node -> pymbolic variable + self.name_gen = UniqueNameGenerator() + self.target = target + self.loop_priorities = set() # used to avoid disadvantageous loop interchanges + + def fetch_multiindex(self, multiindex): + indices = [] + for index in multiindex: + if isinstance(index, gem.Index): + indices.append(self.active_indices[index]) + elif isinstance(index, gem.VariableIndex): + indices.append(expression(index.expression, self)) + else: + assert isinstance(index, int) + indices.append(index) + return tuple(indices) + + # Generate index from gem multiindex + def gem_to_pym_multiindex(self, multiindex): + indices = [] + for index in multiindex: + assert index.extent + if not index.name: + name = self.name_gen(self.index_names[index]) + else: + name = index.name + self.index_extent[name] = index.extent + indices.append(p.Variable(name)) + return tuple(indices) + + # Generate index from shape + def pymbolic_multiindex(self, shape): + indices = [] + for extent in shape: + name = self.name_gen(self.index_names[extent]) + self.index_extent[name] = extent + indices.append(p.Variable(name)) + return tuple(indices) + + # Generate pym variable from gem + def pymbolic_variable_and_destruct(self, node): + pym = self.pymbolic_variable(node) + if isinstance(pym, p.Subscript): + return pym.aggregate, pym.index_tuple + else: + return pym, () + + # Generate pym variable or subscript + def pymbolic_variable(self, node): + pym = self._gem_to_pym_var(node) + if node in self.indices: + indices = self.fetch_multiindex(self.indices[node]) + if indices: + return p.Subscript(pym, indices) + return pym + + def _gem_to_pym_var(self, node): + try: + pym = self.gem_to_pymbolic[node] + except KeyError: + name = self.name_gen(node.name) + pym = p.Variable(name) + self.gem_to_pymbolic[node] = pym + return pym + + def active_inames(self): + # Return all active indices + return frozenset([i.name for i in self.active_indices.values()]) + + def save_loop_ordering(self): + """Save the active loops to prevent loop reordering.""" + priority = tuple(map(str, self.active_indices.values())) + if len(priority) > 1: + self.loop_priorities.add(priority) + + +@contextmanager +def active_indices(mapping, ctx): + """Push active indices onto context. + :arg mapping: dict mapping gem indices to pymbolic index expressions + :arg ctx: code generation context. + :returns: new code generation context.""" + ctx.active_indices.update(mapping) + ctx.save_loop_ordering() + yield ctx + for key in mapping: + ctx.active_indices.pop(key) + + +def generate(impero_c, args, scalar_type, kernel_name="loopy_kernel", index_names=[], + return_increments=True, log=False): + """Generates loopy code. + + :arg impero_c: ImperoC tuple with Impero AST and other data + :arg args: list of loopy.GlobalArgs + :arg scalar_type: type of scalars as C typename string + :arg kernel_name: function name of the kernel + :arg index_names: pre-assigned index names + :arg return_increments: Does codegen for Return nodes increment the lvalue, or assign? + :arg log: bool if the Kernel should be profiled with Log events + :returns: loopy kernel + """ + ctx = LoopyContext(target=target) + ctx.indices = impero_c.indices + ctx.index_names = defaultdict(lambda: "i", index_names) + ctx.epsilon = numpy.finfo(scalar_type).resolution + ctx.scalar_type = scalar_type + ctx.return_increments = return_increments + + # Create arguments + data = list(args) + for i, (temp, dtype) in enumerate(assign_dtypes(impero_c.temporaries, scalar_type)): + name = "t%d" % i + if isinstance(temp, gem.Constant): + data.append(lp.TemporaryVariable(name, shape=temp.shape, dtype=dtype, initializer=temp.array, address_space=lp.AddressSpace.LOCAL, read_only=True)) + else: + shape = tuple([i.extent for i in ctx.indices[temp]]) + temp.shape + data.append(lp.TemporaryVariable(name, shape=shape, dtype=dtype, initializer=None, address_space=lp.AddressSpace.LOCAL, read_only=False)) + ctx.gem_to_pymbolic[temp] = p.Variable(name) + + # Create instructions + instructions = statement(impero_c.tree, ctx) + + # add a no-op touching all kernel arguments to make sure they + # are not silently dropped + noop = lp.CInstruction( + (), "", read_variables=frozenset({a.name for a in args}), + within_inames=frozenset(), within_inames_is_final=True) + instructions.append(noop) + + # Profile the instructions + instructions, event_name, preamble = profile_insns(kernel_name, instructions, log) + + # Create domains + domains = create_domains(ctx.index_extent.items()) + + # Create loopy kernel + knl = lp.make_kernel( + domains, + instructions, + data, + name=kernel_name, + target=target, + seq_dependencies=True, + silenced_warnings=["summing_if_branches_ops"], + lang_version=(2018, 2), + preambles=preamble, + loop_priority=frozenset(ctx.loop_priorities), + ) + + return knl, event_name + + +def create_domains(indices): + """ Create ISL domains from indices + + :arg indices: iterable of (index_name, extent) pairs + :returns: A list of ISL sets representing the iteration domain of the indices.""" + + domains = [] + for idx, extent in indices: + inames = isl.make_zero_and_vars([idx]) + domains.append(((inames[0].le_set(inames[idx])) & (inames[idx].lt_set(inames[0] + extent)))) + + if not domains: + domains = [isl.BasicSet("[] -> {[]}")] + return domains + + +@singledispatch +def statement(tree, ctx): + """Translates an Impero (sub)tree into a loopy instructions corresponding + to a C statement. + + :arg tree: Impero (sub)tree + :arg ctx: miscellaneous code generation data + :returns: list of loopy instructions + """ + raise AssertionError("cannot generate loopy from %s" % type(tree)) + + +@statement.register(imp.Block) +def statement_block(tree, ctx): + from itertools import chain + return list(chain(*(statement(child, ctx) for child in tree.children))) + + +@statement.register(imp.For) +def statement_for(tree, ctx): + extent = tree.index.extent + assert extent + idx = ctx.name_gen(ctx.index_names[tree.index]) + ctx.index_extent[idx] = extent + with active_indices({tree.index: p.Variable(idx)}, ctx) as ctx_active: + return statement(tree.children[0], ctx_active) + + +@statement.register(imp.Initialise) +def statement_initialise(leaf, ctx): + return [lp.Assignment(expression(leaf.indexsum, ctx), 0.0, within_inames=ctx.active_inames())] + + +@statement.register(imp.Accumulate) +def statement_accumulate(leaf, ctx): + lhs = expression(leaf.indexsum, ctx) + rhs = lhs + expression(leaf.indexsum.children[0], ctx) + return [lp.Assignment(lhs, rhs, within_inames=ctx.active_inames())] + + +@statement.register(imp.Return) +def statement_return(leaf, ctx): + lhs = expression(leaf.variable, ctx) + rhs = expression(leaf.expression, ctx) + if ctx.return_increments: + rhs = lhs + rhs + return [lp.Assignment(lhs, rhs, within_inames=ctx.active_inames())] + + +@statement.register(imp.ReturnAccumulate) +def statement_returnaccumulate(leaf, ctx): + lhs = expression(leaf.variable, ctx) + rhs = lhs + expression(leaf.indexsum.children[0], ctx) + return [lp.Assignment(lhs, rhs, within_inames=ctx.active_inames())] + + +@statement.register(imp.Evaluate) +def statement_evaluate(leaf, ctx): + expr = leaf.expression + if isinstance(expr, gem.ListTensor): + ops = [] + var, index = ctx.pymbolic_variable_and_destruct(expr) + for multiindex, value in numpy.ndenumerate(expr.array): + ops.append(lp.Assignment(p.Subscript(var, index + multiindex), expression(value, ctx), within_inames=ctx.active_inames())) + return ops + elif isinstance(expr, gem.Constant): + return [] + elif isinstance(expr, gem.ComponentTensor): + idx = ctx.gem_to_pym_multiindex(expr.multiindex) + var, sub_idx = ctx.pymbolic_variable_and_destruct(expr) + lhs = p.Subscript(var, idx + sub_idx) + with active_indices(dict(zip(expr.multiindex, idx)), ctx) as ctx_active: + return [lp.Assignment(lhs, expression(expr.children[0], ctx_active), within_inames=ctx_active.active_inames())] + elif isinstance(expr, gem.Inverse): + idx = ctx.pymbolic_multiindex(expr.shape) + var = ctx.pymbolic_variable(expr) + lhs = (SubArrayRef(idx, p.Subscript(var, idx)),) + + idx_reads = ctx.pymbolic_multiindex(expr.children[0].shape) + var_reads = ctx.pymbolic_variable(expr.children[0]) + reads = (SubArrayRef(idx_reads, p.Subscript(var_reads, idx_reads)),) + rhs = p.Call(p.Variable("inverse"), reads) + + return [lp.CallInstruction(lhs, rhs, within_inames=ctx.active_inames())] + elif isinstance(expr, gem.Solve): + idx = ctx.pymbolic_multiindex(expr.shape) + var = ctx.pymbolic_variable(expr) + lhs = (SubArrayRef(idx, p.Subscript(var, idx)),) + + reads = [] + for child in expr.children: + idx_reads = ctx.pymbolic_multiindex(child.shape) + var_reads = ctx.pymbolic_variable(child) + reads.append(SubArrayRef(idx_reads, p.Subscript(var_reads, idx_reads))) + rhs = p.Call(p.Variable("solve"), tuple(reads)) + + return [lp.CallInstruction(lhs, rhs, within_inames=ctx.active_inames())] + else: + return [lp.Assignment(ctx.pymbolic_variable(expr), expression(expr, ctx, top=True), within_inames=ctx.active_inames())] + + +def expression(expr, ctx, top=False): + """Translates GEM expression into a pymbolic expression + + :arg expr: GEM expression + :arg ctx: miscellaneous code generation data + :arg top: do not generate temporary reference for the root node + :returns: pymbolic expression + """ + if not top and expr in ctx.gem_to_pymbolic: + return ctx.pymbolic_variable(expr) + else: + return _expression(expr, ctx) + + +@singledispatch +def _expression(expr, ctx): + raise AssertionError("cannot generate expression from %s" % type(expr)) + + +@_expression.register(gem.Failure) +def _expression_failure(expr, ctx): + raise expr.exception + + +@_expression.register(gem.Product) +def _expression_product(expr, ctx): + return p.Product(tuple(expression(c, ctx) for c in expr.children)) + + +@_expression.register(gem.Sum) +def _expression_sum(expr, ctx): + return p.Sum(tuple(expression(c, ctx) for c in expr.children)) + + +@_expression.register(gem.Division) +def _expression_division(expr, ctx): + return p.Quotient(*(expression(c, ctx) for c in expr.children)) + + +@_expression.register(gem.FloorDiv) +def _expression_floordiv(expr, ctx): + return p.FloorDiv(*(expression(c, ctx) for c in expr.children)) + + +@_expression.register(gem.Remainder) +def _expression_remainder(expr, ctx): + return p.Remainder(*(expression(c, ctx) for c in expr.children)) + + +@_expression.register(gem.Power) +def _expression_power(expr, ctx): + return p.Variable("pow")(*(expression(c, ctx) for c in expr.children)) + + +@_expression.register(gem.MathFunction) +def _expression_mathfunction(expr, ctx): + if expr.name.startswith('cyl_bessel_'): + # Bessel functions + if is_complex(ctx.scalar_type): + raise NotImplementedError("Bessel functions for complex numbers: " + "missing implementation") + nu, arg = expr.children + nu_ = expression(nu, ctx) + arg_ = expression(arg, ctx) + if isinstance(ctx.target, lp.target.c.CWithGNULibcTarget): + # Generate right functions calls to gnulibc bessel functions + # cyl_bessel_{jy} -> bessel_{jy} + name = expr.name[4:] + return p.Variable(f"{name}n")(int(nu_), arg_) + else: + # Modified Bessel functions (C++ only) + # These mappings work for FEniCS only, and fail with Firedrake + # since no Boost available. + # Is this actually still supported/has ever been used by anyone? + if expr.name in {'cyl_bessel_i', 'cyl_bessel_k'}: + name = 'boost::math::' + expr.name + return p.Variable(name)(nu_, arg_) + else: + # cyl_bessel_{jy} -> {jy} + name = expr.name[-1:] + if nu == gem.Zero(): + return p.Variable(f"{name}0")(arg_) + elif nu == gem.one: + return p.Variable(f"{name}1")(arg_) + else: + return p.Variable(f"{name}n")(nu_, arg_) + else: + if expr.name == "ln": + name = "log" + else: + name = expr.name + # Not all mathfunctions apply to complex numbers, but this + # will be picked up in loopy. This way we allow erf(real(...)) + # in complex mode (say). + return p.Variable(name)(*(expression(c, ctx) for c in expr.children)) + + +@_expression.register(gem.MinValue) +def _expression_minvalue(expr, ctx): + return p.Variable("min")(*[expression(c, ctx) for c in expr.children]) + + +@_expression.register(gem.MaxValue) +def _expression_maxvalue(expr, ctx): + return p.Variable("max")(*[expression(c, ctx) for c in expr.children]) + + +@_expression.register(gem.Comparison) +def _expression_comparison(expr, ctx): + left, right = [expression(c, ctx) for c in expr.children] + return p.Comparison(left, expr.operator, right) + + +@_expression.register(gem.LogicalNot) +def _expression_logicalnot(expr, ctx): + child, = expr.children + return p.LogicalNot(expression(child, ctx)) + + +@_expression.register(gem.LogicalAnd) +def _expression_logicaland(expr, ctx): + return p.LogicalAnd(tuple([expression(c, ctx) for c in expr.children])) + + +@_expression.register(gem.LogicalOr) +def _expression_logicalor(expr, ctx): + return p.LogicalOr(tuple([expression(c, ctx) for c in expr.children])) + + +@_expression.register(gem.Conditional) +def _expression_conditional(expr, ctx): + return p.If(*[expression(c, ctx) for c in expr.children]) + + +@_expression.register(gem.Constant) +def _expression_scalar(expr, ctx): + assert not expr.shape + v = expr.value + if numpy.isnan(v): + return p.Variable("NAN") + r = numpy.round(v, 1) + if r and numpy.abs(v - r) < ctx.epsilon: + return r + return v + + +@_expression.register(gem.Variable) +def _expression_variable(expr, ctx): + return ctx.pymbolic_variable(expr) + + +@_expression.register(gem.Indexed) +def _expression_indexed(expr, ctx): + rank = ctx.fetch_multiindex(expr.multiindex) + var = expression(expr.children[0], ctx) + if isinstance(var, p.Subscript): + rank = var.index + rank + var = var.aggregate + return p.Subscript(var, rank) + + +@_expression.register(gem.FlexiblyIndexed) +def _expression_flexiblyindexed(expr, ctx): + var = expression(expr.children[0], ctx) + + rank = [] + for off, idxs in expr.dim2idxs: + rank_ = [expression(off, ctx)] + for index, stride in idxs: + if isinstance(index, gem.Index): + rank_.append(p.Product((ctx.active_indices[index], expression(stride, ctx)))) + elif isinstance(index, gem.VariableIndex): + rank_.append(p.Product((expression(index.expression, ctx), expression(stride, ctx)))) + else: + raise ValueError(f"Expecting Index or VariableIndex, not {type(index)}") + rank.append(p.Sum(tuple(rank_))) + + return p.Subscript(var, tuple(rank)) + + +@_expression.register(Integral) +def _expression_numbers_integral(expr, ctx): + return expr diff --git a/tsfc/modified_terminals.py b/tsfc/modified_terminals.py new file mode 100644 index 0000000000..8c5162bf97 --- /dev/null +++ b/tsfc/modified_terminals.py @@ -0,0 +1,179 @@ +# -*- coding: utf-8 -*- +# Copyright (C) 2011-2015 Martin Sandve Alnæs +# +# This file is part of UFLACS. +# +# UFLACS is free software: you can redistribute it and/or modify +# it under the terms of the GNU Lesser General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# UFLACS is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Lesser General Public License for more details. +# +# You should have received a copy of the GNU Lesser General Public License +# along with UFLACS. If not, see . +# +# Modified by Miklós Homolya, 2016. + +"""Definitions of 'modified terminals', a core concept in uflacs.""" + +from ufl.classes import (ReferenceValue, ReferenceGrad, + NegativeRestricted, PositiveRestricted, + Restricted, ConstantValue, + Jacobian, SpatialCoordinate, Zero) +from ufl.checks import is_cellwise_constant +from ufl.domain import extract_unique_domain + + +class ModifiedTerminal(object): + + """A modified terminal expression is an object of a Terminal subtype, wrapped in terminal modifier types. + + The variables of this class are: + + expr - The original UFL expression + + terminal - the underlying Terminal object + local_derivatives - tuple of ints, each meaning derivative in that local direction + reference_value - bool, whether this is represented in reference frame + restriction - None, '+' or '-' + """ + + def __init__(self, expr, terminal, local_derivatives, restriction, reference_value): + # The original expression + self.expr = expr + + # The underlying terminal expression + self.terminal = terminal + + # Components + self.reference_value = reference_value + self.restriction = restriction + + # Derivatives + self.local_derivatives = local_derivatives + + def as_tuple(self): + t = self.terminal + rv = self.reference_value + ld = self.local_derivatives + r = self.restriction + return (t, rv, ld, r) + + def __hash__(self): + return hash(self.as_tuple()) + + def __eq__(self, other): + return isinstance(other, ModifiedTerminal) and self.as_tuple() == other.as_tuple() + + def __lt__(self, other): + return self.as_tuple() < other.as_tuple() + + def __str__(self): + s = [] + s += ["terminal: {0}".format(self.terminal)] + s += ["local_derivatives: {0}".format(self.local_derivatives)] + s += ["restriction: {0}".format(self.restriction)] + return '\n'.join(s) + + +def is_modified_terminal(v): + "Check if v is a terminal or a terminal wrapped in terminal modifier types." + while not v._ufl_is_terminal_: + if v._ufl_is_terminal_modifier_: + v = v.ufl_operands[0] + else: + return False + return True + + +def strip_modified_terminal(v): + "Extract core Terminal from a modified terminal or return None." + while not v._ufl_is_terminal_: + if v._ufl_is_terminal_modifier_: + v = v.ufl_operands[0] + else: + return None + return v + + +def analyse_modified_terminal(expr): + """Analyse a so-called 'modified terminal' expression and return its properties in more compact form. + + A modified terminal expression is an object of a Terminal subtype, wrapped in terminal modifier types. + + The wrapper types can include 0-* Grad or ReferenceGrad objects, + and 0-1 ReferenceValue, 0-1 Restricted, and 0-1 Indexed objects. + """ + # Data to determine + local_derivatives = 0 + reference_value = None + restriction = None + + # Start with expr and strip away layers of modifiers + t = expr + while not t._ufl_is_terminal_: + if isinstance(t, ReferenceValue): + assert reference_value is None, "Got twice pulled back terminal!" + reference_value = True + t, = t.ufl_operands + + elif isinstance(t, ReferenceGrad): + local_derivatives += 1 + t, = t.ufl_operands + + elif isinstance(t, Restricted): + assert restriction is None, "Got twice restricted terminal!" + restriction = t._side + t, = t.ufl_operands + + elif t._ufl_terminal_modifiers_: + raise ValueError("Missing handler for terminal modifier type %s, object is %s." % (type(t), repr(t))) + + else: + raise ValueError("Unexpected type %s object %s." % (type(t), repr(t))) + + # Make reference_value true or false + if reference_value is None: + reference_value = False + + # Consistency check + if isinstance(t, (SpatialCoordinate, Jacobian)): + pass + else: + if local_derivatives and not reference_value: + raise ValueError("Local derivatives of non-local value?") + + return ModifiedTerminal(expr, t, local_derivatives, restriction, reference_value) + + +def construct_modified_terminal(mt, terminal): + """Construct a modified terminal given terminal modifiers from an + analysed modified terminal and a terminal.""" + expr = terminal + + if mt.reference_value: + expr = ReferenceValue(expr) + + dim = extract_unique_domain(expr).topological_dimension() + for n in range(mt.local_derivatives): + # Return zero if expression is trivially constant. This has to + # happen here because ReferenceGrad has no access to the + # topological dimension of a literal zero. + if is_cellwise_constant(expr): + expr = Zero(expr.ufl_shape + (dim,), expr.ufl_free_indices, + expr.ufl_index_dimensions) + else: + expr = ReferenceGrad(expr) + + # No need to apply restrictions to ConstantValue terminals + if not isinstance(expr, ConstantValue): + if mt.restriction == '+': + expr = PositiveRestricted(expr) + elif mt.restriction == '-': + expr = NegativeRestricted(expr) + + return expr diff --git a/tsfc/parameters.py b/tsfc/parameters.py new file mode 100644 index 0000000000..1277713ad5 --- /dev/null +++ b/tsfc/parameters.py @@ -0,0 +1,36 @@ +import numpy +from loopy.target.c import CWithGNULibcTarget + + +PARAMETERS = { + "quadrature_rule": "auto", + "quadrature_degree": "auto", + + # Default mode + "mode": "spectral", + + # Maximum extent to unroll index sums. Default is 3, so that loops + # over geometric dimensions are unrolled; this improves assembly + # performance. Can be disabled by setting it to None, False or 0; + # that makes compilation time much shorter. + "unroll_indexsum": 3, + + # Scalar type numpy dtype + "scalar_type": numpy.dtype(numpy.float64), + + # So that tests pass (needs to match scalar_type) + "scalar_type_c": "double", +} + + +target = CWithGNULibcTarget() + + +def default_parameters(): + return PARAMETERS.copy() + + +def is_complex(scalar_type): + """Decides complex mode based on scalar type.""" + return scalar_type and (isinstance(scalar_type, numpy.dtype) and scalar_type.kind == 'c') \ + or (isinstance(scalar_type, str) and "complex" in scalar_type) diff --git a/tsfc/spectral.py b/tsfc/spectral.py new file mode 100644 index 0000000000..8ee838481d --- /dev/null +++ b/tsfc/spectral.py @@ -0,0 +1,193 @@ +from collections import OrderedDict, defaultdict, namedtuple +from functools import partial, reduce +from itertools import chain, zip_longest + +from gem.gem import Delta, Indexed, Sum, index_sum, one +from gem.node import Memoizer +from gem.optimise import delta_elimination as _delta_elimination +from gem.optimise import remove_componenttensors, replace_division, unroll_indexsum +from gem.refactorise import ATOMIC, COMPOUND, OTHER, MonomialSum, collect_monomials +from gem.unconcatenate import unconcatenate +from gem.coffee import optimise_monomial_sum +from gem.utils import groupby + + +Integral = namedtuple('Integral', ['expression', + 'quadrature_multiindex', + 'argument_indices']) + + +def Integrals(expressions, quadrature_multiindex, argument_multiindices, parameters): + """Constructs an integral representation for each GEM integrand + expression. + + :arg expressions: integrand multiplied with quadrature weight; + multi-root GEM expression DAG + :arg quadrature_multiindex: quadrature multiindex (tuple) + :arg argument_multiindices: tuple of argument multiindices, + one multiindex for each argument + :arg parameters: parameters dictionary + + :returns: list of integral representations + """ + # Rewrite: a / b => a * (1 / b) + expressions = replace_division(expressions) + + # Unroll + max_extent = parameters["unroll_indexsum"] + if max_extent: + def predicate(index): + return index.extent <= max_extent + expressions = unroll_indexsum(expressions, predicate=predicate) + + expressions = [index_sum(e, quadrature_multiindex) for e in expressions] + argument_indices = tuple(chain(*argument_multiindices)) + return [Integral(e, quadrature_multiindex, argument_indices) for e in expressions] + + +def _delta_inside(node, self): + """Does node contain a Delta?""" + return any(isinstance(child, Delta) or self(child) + for child in node.children) + + +def flatten(var_reps, index_cache): + quadrature_indices = OrderedDict() + + pairs = [] # assignment pairs + for variable, reps in var_reps: + # Extract argument indices + argument_indices, = set(r.argument_indices for r in reps) + assert set(variable.free_indices) == set(argument_indices) + + # Extract and verify expressions + expressions = [r.expression for r in reps] + assert all(set(e.free_indices) <= set(argument_indices) + for e in expressions) + + # Save assignment pair + pairs.append((variable, reduce(Sum, expressions))) + + # Collect quadrature_indices + for r in reps: + quadrature_indices.update(zip_longest(r.quadrature_multiindex, ())) + + # Split Concatenate nodes + pairs = unconcatenate(pairs, cache=index_cache) + + def group_key(pair): + variable, expression = pair + return frozenset(variable.free_indices) + + delta_inside = Memoizer(_delta_inside) + # Variable ordering after delta cancellation + narrow_variables = OrderedDict() + # Assignments are variable -> MonomialSum map + delta_simplified = defaultdict(MonomialSum) + # Group assignment pairs by argument indices + for free_indices, pair_group in groupby(pairs, group_key): + variables, expressions = zip(*pair_group) + # Argument factorise expressions + classifier = partial(classify, set(free_indices), delta_inside=delta_inside) + monomial_sums = collect_monomials(expressions, classifier) + # For each monomial, apply delta cancellation and insert + # result into delta_simplified. + for variable, monomial_sum in zip(variables, monomial_sums): + for monomial in monomial_sum: + var, s, a, r = delta_elimination(variable, *monomial) + narrow_variables.setdefault(var) + delta_simplified[var].add(s, a, r) + + # Final factorisation + for variable in narrow_variables: + monomial_sum = delta_simplified[variable] + # Collect sum indices applicable to the current MonomialSum + sum_indices = set().union(*[m.sum_indices for m in monomial_sum]) + # Put them in a deterministic order + sum_indices = [i for i in quadrature_indices if i in sum_indices] + # Sort for increasing index extent, this obtains the good + # factorisation for triangle x interval cells. Python sort is + # stable, so in the common case when index extents are equal, + # the previous deterministic ordering applies which is good + # for getting smaller temporaries. + sum_indices = sorted(sum_indices, key=lambda index: index.extent) + # Apply sum factorisation combined with COFFEE technology + expression = sum_factorise(variable, sum_indices, monomial_sum) + yield (variable, expression) + + +finalise_options = dict(replace_delta=False) + + +def classify(argument_indices, expression, delta_inside): + """Classifier for argument factorisation""" + n = len(argument_indices.intersection(expression.free_indices)) + if n == 0: + return OTHER + elif n == 1: + if isinstance(expression, (Delta, Indexed)) and not delta_inside(expression): + return ATOMIC + else: + return COMPOUND + else: + return COMPOUND + + +def delta_elimination(variable, sum_indices, args, rest): + """IndexSum-Delta cancellation for monomials.""" + factors = list(args) + [variable, rest] # construct factors + + def prune(factors): + # Skip last factor (``rest``, see above) which can be + # arbitrarily complicated, so its pruning may be expensive, + # and its early pruning brings no advantages. + result = remove_componenttensors(factors[:-1]) + result.append(factors[-1]) + return result + + # Cancel sum indices + sum_indices, factors = _delta_elimination(sum_indices, factors) + factors = prune(factors) + + # Cancel variable indices + var_indices, factors = _delta_elimination(variable.free_indices, factors) + factors = prune(factors) + + # Destructure factors after cancellation + rest = factors.pop() + variable = factors.pop() + args = [f for f in factors if f != one] + + assert set(var_indices) == set(variable.free_indices) + return variable, sum_indices, args, rest + + +def sum_factorise(variable, tail_ordering, monomial_sum): + if tail_ordering: + key_ordering = OrderedDict() + sub_monosums = defaultdict(MonomialSum) + for sum_indices, atomics, rest in monomial_sum: + # Pull out those sum indices that are not contained in the + # tail ordering, together with those atomics which do not + # share free indices with the tail ordering. + # + # Based on this, split the monomial sum, then recursively + # optimise each sub monomial sum with the first tail index + # removed. + tail_indices = tuple(i for i in sum_indices if i in tail_ordering) + tail_atomics = tuple(a for a in atomics + if set(tail_indices) & set(a.free_indices)) + head_indices = tuple(i for i in sum_indices if i not in tail_ordering) + head_atomics = tuple(a for a in atomics if a not in tail_atomics) + key = (head_indices, head_atomics) + key_ordering.setdefault(key) + sub_monosums[key].add(tail_indices, tail_atomics, rest) + sub_monosums = [(k, sub_monosums[k]) for k in key_ordering] + + monomial_sum = MonomialSum() + for (sum_indices, atomics), monosum in sub_monosums: + new_rest = sum_factorise(variable, tail_ordering[1:], monosum) + monomial_sum.add(sum_indices, atomics, new_rest) + + # Use COFFEE algorithm to optimise the monomial sum + return optimise_monomial_sum(monomial_sum, variable.index_ordering()) diff --git a/tsfc/tensor.py b/tsfc/tensor.py new file mode 100644 index 0000000000..8902b470a2 --- /dev/null +++ b/tsfc/tensor.py @@ -0,0 +1,93 @@ +from collections import defaultdict +from functools import partial, reduce +from itertools import count + +import numpy + +import gem +from gem.optimise import remove_componenttensors, unroll_indexsum +from gem.refactorise import ATOMIC, COMPOUND, OTHER, collect_monomials +from gem.unconcatenate import flatten as concatenate + + +def einsum(factors, sum_indices): + """Evaluates a tensor product at compile time. + + :arg factors: iterable of indexed GEM literals + :arg sum_indices: indices to sum over + :returns: a single indexed GEM literal + """ + # Maps the problem onto numpy.einsum + index2letter = defaultdict(partial(lambda c: chr(ord('i') + next(c)), count())) + operands = [] + subscript_parts = [] + for factor in factors: + literal, = factor.children + selectors = [] + letters = [] + for index in factor.multiindex: + if isinstance(index, int): + selectors.append(index) + else: + selectors.append(slice(None)) + letters.append(index2letter[index]) + operands.append(literal.array.__getitem__(tuple(selectors))) + subscript_parts.append(''.join(letters)) + + result_pairs = sorted((letter, index) + for index, letter in index2letter.items() + if index not in sum_indices) + + subscripts = ','.join(subscript_parts) + '->' + ''.join(l for l, i in result_pairs) + tensor = numpy.einsum(subscripts, *operands) + return gem.Indexed(gem.Literal(tensor), tuple(i for l, i in result_pairs)) + + +def Integrals(expressions, quadrature_multiindex, argument_multiindices, parameters): + # Concatenate + expressions = concatenate(expressions) + + # Unroll + max_extent = parameters["unroll_indexsum"] + if max_extent: + def predicate(index): + return index.extent <= max_extent + expressions = unroll_indexsum(expressions, predicate=predicate) + + # Refactorise + def classify(quadrature_indices, expression): + if not quadrature_indices.intersection(expression.free_indices): + return OTHER + elif isinstance(expression, gem.Indexed) and isinstance(expression.children[0], gem.Literal): + return ATOMIC + else: + return COMPOUND + classifier = partial(classify, set(quadrature_multiindex)) + + result = [] + for expr, monomial_sum in zip(expressions, collect_monomials(expressions, classifier)): + # Select quadrature indices that are present + quadrature_indices = set(index for index in quadrature_multiindex + if index in expr.free_indices) + + products = [] + for sum_indices, factors, rest in monomial_sum: + # Collapse quadrature literals for each monomial + if factors or quadrature_indices: + replacement = einsum(remove_componenttensors(factors), quadrature_indices) + else: + replacement = gem.Literal(1) + # Rebuild expression + products.append(gem.IndexSum(gem.Product(replacement, rest), sum_indices)) + result.append(reduce(gem.Sum, products, gem.Zero())) + return result + + +def flatten(var_reps, index_cache): + for variable, reps in var_reps: + expressions = reps + for expression in expressions: + yield (variable, expression) + + +finalise_options = {} diff --git a/tsfc/ufl2gem.py b/tsfc/ufl2gem.py new file mode 100644 index 0000000000..e283fb1414 --- /dev/null +++ b/tsfc/ufl2gem.py @@ -0,0 +1,160 @@ +"""Translation of UFL tensor-algebra into GEM tensor-algebra.""" + +import collections +import ufl + +from gem import (Literal, Zero, Identity, Sum, Product, Division, + Power, MathFunction, MinValue, MaxValue, Comparison, + LogicalNot, LogicalAnd, LogicalOr, Conditional, + Index, Indexed, ComponentTensor, IndexSum, + ListTensor) + + +class Mixin(object): + """A mixin to be used with a UFL MultiFunction to translate UFL + algebra into GEM tensor-algebra. This node types translate pretty + straightforwardly to GEM. Other node types are not handled in + this mixin.""" + + def __init__(self): + self.index_map = collections.defaultdict(Index) + """A map for translating UFL free indices into GEM free + indices.""" + + def scalar_value(self, o): + return Literal(o.value()) + + def identity(self, o): + return Identity(o._dim) + + def zero(self, o): + return Zero(o.ufl_shape) + + def sum(self, o, *ops): + if o.ufl_shape: + indices = tuple(Index() for i in range(len(o.ufl_shape))) + return ComponentTensor(Sum(*[Indexed(op, indices) for op in ops]), indices) + else: + return Sum(*ops) + + def product(self, o, *ops): + assert o.ufl_shape == () + return Product(*ops) + + def division(self, o, numerator, denominator): + return Division(numerator, denominator) + + def real(self, o, expr): + if o.ufl_shape: + indices = tuple(Index() for i in range(len(o.ufl_shape))) + return ComponentTensor(MathFunction('real', Indexed(expr, indices)), indices) + else: + return MathFunction('real', expr) + + def imag(self, o, expr): + if o.ufl_shape: + indices = tuple(Index() for i in range(len(o.ufl_shape))) + return ComponentTensor(MathFunction('imag', Indexed(expr, indices)), indices) + else: + return MathFunction('imag', expr) + + def conj(self, o, expr): + if o.ufl_shape: + indices = tuple(Index() for i in range(len(o.ufl_shape))) + return ComponentTensor(MathFunction('conj', Indexed(expr, indices)), indices) + else: + return MathFunction('conj', expr) + + def abs(self, o, expr): + if o.ufl_shape: + indices = tuple(Index() for i in range(len(o.ufl_shape))) + return ComponentTensor(MathFunction('abs', Indexed(expr, indices)), indices) + else: + return MathFunction('abs', expr) + + def power(self, o, base, exponent): + return Power(base, exponent) + + def math_function(self, o, expr): + return MathFunction(o._name, expr) + + def atan2(self, o, y, x): + return MathFunction("atan2", y, x) + + def bessel_i(self, o, nu, arg): + return MathFunction(o._name, nu, arg) + + def bessel_j(self, o, nu, arg): + return MathFunction(o._name, nu, arg) + + def bessel_k(self, o, nu, arg): + return MathFunction(o._name, nu, arg) + + def bessel_y(self, o, nu, arg): + return MathFunction(o._name, nu, arg) + + def min_value(self, o, *ops): + return MinValue(*ops) + + def max_value(self, o, *ops): + return MaxValue(*ops) + + def binary_condition(self, o, left, right): + return Comparison(o._name, left, right) + + def not_condition(self, o, expr): + return LogicalNot(expr) + + def and_condition(self, o, *ops): + return LogicalAnd(*ops) + + def or_condition(self, o, *ops): + return LogicalOr(*ops) + + def conditional(self, o, condition, then, else_): + assert condition.shape == () + if o.ufl_shape: + indices = tuple(Index() for i in range(len(o.ufl_shape))) + return ComponentTensor(Conditional(condition, Indexed(then, indices), + Indexed(else_, indices)), + indices) + else: + return Conditional(condition, then, else_) + + def multi_index(self, o): + indices = [] + for i in o: + if isinstance(i, ufl.classes.FixedIndex): + indices.append(int(i)) + elif isinstance(i, ufl.classes.Index): + indices.append(self.index_map[i.count()]) + return tuple(indices) + + def indexed(self, o, aggregate, index): + return Indexed(aggregate, index) + + def list_tensor(self, o, *ops): + return ListTensor(ops) + + def component_tensor(self, o, expression, index): + return ComponentTensor(expression, index) + + def index_sum(self, o, summand, indices): + # ufl.IndexSum technically has a MultiIndex, but it must have + # exactly one index in it. + index, = indices + + if o.ufl_shape: + indices = tuple(Index() for i in range(len(o.ufl_shape))) + return ComponentTensor(IndexSum(Indexed(summand, indices), (index,)), indices) + else: + return IndexSum(summand, (index,)) + + def variable(self, o, expression, label): + # Only used by UFL AD, at this point, the bare expression is + # what we want. + return expression + + def label(self, o): + # Only used by UFL AD, don't need it at this point. + pass diff --git a/tsfc/ufl_utils.py b/tsfc/ufl_utils.py new file mode 100644 index 0000000000..c789459cdf --- /dev/null +++ b/tsfc/ufl_utils.py @@ -0,0 +1,497 @@ +"""Utilities for preprocessing UFL objects.""" + +from functools import singledispatch + +import numpy + +import ufl +from ufl import as_tensor, indices, replace +from ufl.algorithms import compute_form_data as ufl_compute_form_data +from ufl.algorithms import estimate_total_polynomial_degree +from ufl.algorithms.analysis import extract_arguments, extract_type +from ufl.algorithms.apply_function_pullbacks import apply_function_pullbacks +from ufl.algorithms.apply_algebra_lowering import apply_algebra_lowering +from ufl.algorithms.apply_derivatives import apply_derivatives +from ufl.algorithms.apply_geometry_lowering import apply_geometry_lowering +from ufl.algorithms.apply_restrictions import apply_restrictions +from ufl.algorithms.comparison_checker import do_comparison_check +from ufl.algorithms.remove_complex_nodes import remove_complex_nodes +from ufl.corealg.map_dag import map_expr_dag +from ufl.corealg.multifunction import MultiFunction +from ufl.geometry import QuadratureWeight +from ufl.geometry import Jacobian, JacobianDeterminant, JacobianInverse +from ufl.classes import (Abs, Argument, CellOrientation, Coefficient, + ComponentTensor, Expr, FloatValue, Division, + MultiIndex, Product, + ScalarValue, Sqrt, Zero, CellVolume, FacetArea) +from ufl.domain import extract_unique_domain +from finat.ufl import MixedElement +from ufl.utils.sorting import sorted_by_count + +from gem.node import MemoizerArg + +from tsfc.modified_terminals import (is_modified_terminal, + analyse_modified_terminal, + construct_modified_terminal) + + +preserve_geometry_types = (CellVolume, FacetArea) + + +def compute_form_data(form, + do_apply_function_pullbacks=True, + do_apply_integral_scaling=True, + do_apply_geometry_lowering=True, + preserve_geometry_types=preserve_geometry_types, + do_apply_restrictions=True, + do_estimate_degrees=True, + complex_mode=False): + """Preprocess UFL form in a format suitable for TSFC. Return + form data. + + This is merely a wrapper to UFL compute_form_data with default + kwargs overriden in the way TSFC needs it and is provided for + other form compilers based on TSFC. + """ + fd = ufl_compute_form_data( + form, + do_apply_function_pullbacks=do_apply_function_pullbacks, + do_apply_integral_scaling=do_apply_integral_scaling, + do_apply_geometry_lowering=do_apply_geometry_lowering, + preserve_geometry_types=preserve_geometry_types, + do_apply_restrictions=do_apply_restrictions, + do_estimate_degrees=do_estimate_degrees, + complex_mode=complex_mode + ) + constants = extract_firedrake_constants(form) + fd.constants = constants + return fd + + +def extract_firedrake_constants(a): + """Build a sorted list of all constants in a""" + return sorted_by_count(extract_type(a, TSFCConstantMixin)) + + +def one_times(measure): + # Workaround for UFL issue #80: + # https://bitbucket.org/fenics-project/ufl/issues/80 + form = 1 * measure + fd = compute_form_data(form, do_estimate_degrees=False) + itg_data, = fd.integral_data + integral, = itg_data.integrals + integrand = integral.integrand() + + # UFL considers QuadratureWeight a geometric quantity, and the + # general handler for geometric quantities estimates the degree of + # the coordinate element. This would unnecessarily increase the + # estimated degree, so we drop QuadratureWeight instead. + expression = replace(integrand, {QuadratureWeight(itg_data.domain): 1}) + + # Now estimate degree for the preprocessed form + degree = estimate_total_polynomial_degree(expression) + + return integrand, degree + + +def entity_avg(integrand, measure, argument_multiindices): + arguments = extract_arguments(integrand) + if len(arguments) == 1: + a, = arguments + integrand = ufl.replace(integrand, {a: ufl.Argument(a.function_space(), + number=0, + part=a.part())}) + argument_multiindices = (argument_multiindices[a.number()], ) + + degree = estimate_total_polynomial_degree(integrand) + form = integrand * measure + fd = compute_form_data(form, do_estimate_degrees=False, + do_apply_function_pullbacks=False) + itg_data, = fd.integral_data + integral, = itg_data.integrals + integrand = integral.integrand() + return integrand, degree, argument_multiindices + + +def preprocess_expression(expression, complex_mode=False, + do_apply_restrictions=False): + """Imitates the compute_form_data processing pipeline. + + :arg complex_mode: Are we in complex UFL mode? + :arg do_apply_restrictions: Propogate restrictions to terminals? + + Useful, for example, to preprocess non-scalar expressions, which + are not and cannot be forms. + """ + if complex_mode: + expression = do_comparison_check(expression) + else: + expression = remove_complex_nodes(expression) + expression = apply_algebra_lowering(expression) + expression = apply_derivatives(expression) + expression = apply_function_pullbacks(expression) + expression = apply_geometry_lowering(expression, preserve_geometry_types) + expression = apply_derivatives(expression) + expression = apply_geometry_lowering(expression, preserve_geometry_types) + expression = apply_derivatives(expression) + if not complex_mode: + expression = remove_complex_nodes(expression) + if do_apply_restrictions: + expression = apply_restrictions(expression) + return expression + + +class ModifiedTerminalMixin(object): + """Mixin to use with MultiFunctions that operate on modified + terminals.""" + + def unexpected(self, o): + assert False, "Not expected %r at this stage." % o + + # global derivates should have been pulled back + grad = unexpected + div = unexpected + curl = unexpected + + # div and curl should have been algebraically lowered + reference_div = unexpected + reference_curl = unexpected + + def _modified_terminal(self, o): + assert is_modified_terminal(o) + return self.modified_terminal(o) + + # Unlike UFL, we do not regard Indexed as a terminal modifier. + # indexed = _modified_terminal + + positive_restricted = _modified_terminal + negative_restricted = _modified_terminal + + reference_grad = _modified_terminal + reference_value = _modified_terminal + + terminal = _modified_terminal + + +class CoefficientSplitter(MultiFunction, ModifiedTerminalMixin): + def __init__(self, split): + MultiFunction.__init__(self) + self._split = split + + expr = MultiFunction.reuse_if_untouched + + def modified_terminal(self, o): + mt = analyse_modified_terminal(o) + terminal = mt.terminal + + if not isinstance(terminal, Coefficient): + # Only split coefficients + return o + + if type(terminal.ufl_element()) != MixedElement: + # Only split mixed coefficients + return o + + # Reference value expected + assert mt.reference_value + + # Derivative indices + beta = indices(mt.local_derivatives) + + components = [] + for subcoeff in self._split[terminal]: + # Apply terminal modifiers onto the subcoefficient + component = construct_modified_terminal(mt, subcoeff) + # Collect components of the subcoefficient + for alpha in numpy.ndindex(subcoeff.ufl_element().reference_value_shape): + # New modified terminal: component[alpha + beta] + components.append(component[alpha + beta]) + # Repack derivative indices to shape + c, = indices(1) + return ComponentTensor(as_tensor(components)[c], MultiIndex((c,) + beta)) + + +def split_coefficients(expression, split): + """Split mixed coefficients, so mixed elements need not be + implemented. + + :arg split: A :py:class:`dict` mapping each mixed coefficient to a + sequence of subcoefficients. If None, calling this + function is a no-op. + """ + if split is None: + return expression + + splitter = CoefficientSplitter(split) + return map_expr_dag(splitter, expression) + + +class PickRestriction(MultiFunction, ModifiedTerminalMixin): + """Pick out parts of an expression with specified restrictions on + the arguments. + + :arg test: The restriction on the test function. + :arg trial: The restriction on the trial function. + + Returns those parts of the expression that have the requested + restrictions, or else :class:`ufl.classes.Zero` if no such part + exists. + """ + def __init__(self, test=None, trial=None): + self.restrictions = {0: test, 1: trial} + MultiFunction.__init__(self) + + expr = MultiFunction.reuse_if_untouched + + def multi_index(self, o): + return o + + def modified_terminal(self, o): + mt = analyse_modified_terminal(o) + t = mt.terminal + r = mt.restriction + if isinstance(t, Argument) and r != self.restrictions[t.number()]: + return Zero(o.ufl_shape, o.ufl_free_indices, o.ufl_index_dimensions) + else: + return o + + +def ufl_reuse_if_untouched(o, *ops): + """Reuse object if operands are the same objects.""" + if all(a is b for a, b in zip(o.ufl_operands, ops)): + return o + else: + return o._ufl_expr_reconstruct_(*ops) + + +@singledispatch +def _simplify_abs(o, self, in_abs): + """Single-dispatch function to simplify absolute values. + + :arg o: UFL node + :arg self: Callback handler for recursion + :arg in_abs: Is ``o`` inside an absolute value? + + When ``in_abs`` we must return a non-negative value, potentially + by wrapping the returned node with ``Abs``. + """ + raise AssertionError("UFL node expected, not %s" % type(o)) + + +@_simplify_abs.register(Expr) +def _simplify_abs_expr(o, self, in_abs): + # General case, only wrap the outer expression (if necessary) + operands = [self(op, False) for op in o.ufl_operands] + result = ufl_reuse_if_untouched(o, *operands) + if in_abs: + result = Abs(result) + return result + + +@_simplify_abs.register(Sqrt) +def _simplify_abs_sqrt(o, self, in_abs): + result = ufl_reuse_if_untouched(o, self(o.ufl_operands[0], False)) + if self.complex_mode and in_abs: + return Abs(result) + else: + return result + + +@_simplify_abs.register(ScalarValue) +def _simplify_abs_(o, self, in_abs): + if not in_abs: + return o + # Inline abs(constant) + return ufl.as_ufl(abs(o._value)) + + +@_simplify_abs.register(CellOrientation) +def _simplify_abs_cellorientation(o, self, in_abs): + if not in_abs: + return o + # Cell orientation is +-1 + return FloatValue(1) + + +@_simplify_abs.register(Division) +@_simplify_abs.register(Product) +def _simplify_abs_product(o, self, in_abs): + if not in_abs: + # Just reconstruct + ops = [self(op, False) for op in o.ufl_operands] + return ufl_reuse_if_untouched(o, *ops) + + # Visit children, distributing Abs + ops = [self(op, True) for op in o.ufl_operands] + + # Strip Abs off again (we will put it outside now) + stripped = False + strip_ops = [] + for op in ops: + if isinstance(op, Abs): + stripped = True + strip_ops.append(op.ufl_operands[0]) + else: + strip_ops.append(op) + + # Rebuild, and wrap with Abs if necessary + result = ufl_reuse_if_untouched(o, *strip_ops) + if stripped: + result = Abs(result) + return result + + +@_simplify_abs.register(Abs) +def _simplify_abs_abs(o, self, in_abs): + return self(o.ufl_operands[0], True) + + +def simplify_abs(expression, complex_mode): + """Simplify absolute values in a UFL expression. Its primary + purpose is to "neutralise" CellOrientation nodes that are + surrounded by absolute values and thus not at all necessary.""" + mapper = MemoizerArg(_simplify_abs) + mapper.complex_mode = complex_mode + return mapper(expression, False) + + +def apply_mapping(expression, element, domain): + """Apply the inverse of the pullback for element to an expression. + + :arg expression: An expression in physical space + :arg element: The element we're going to interpolate into, whose + value_shape must match the shape of the expression, and will + advertise the pullback to apply. + :arg domain: Optional domain to provide in case expression does + not contain a domain (used for constructing geometric quantities). + :returns: A new UFL expression with shape element.reference_value_shape + :raises NotImplementedError: If we don't know how to apply the + inverse of the pullback. + :raises ValueError: If we get shape mismatches. + + The following is borrowed from the UFC documentation: + + Let g be a field defined on a physical domain T with physical + coordinates x. Let T_0 be a reference domain with coordinates + X. Assume that F: T_0 -> T such that + + x = F(X) + + Let J be the Jacobian of F, i.e J = dx/dX and let K denote the + inverse of the Jacobian K = J^{-1}. Then we (currently) have the + following four types of mappings: + + 'identity' mapping for g: + + G(X) = g(x) + + For vector fields g: + + 'contravariant piola' mapping for g: + + G(X) = det(J) K g(x) i.e G_i(X) = det(J) K_ij g_j(x) + + 'covariant piola' mapping for g: + + G(X) = J^T g(x) i.e G_i(X) = J^T_ij g(x) = J_ji g_j(x) + + 'double covariant piola' mapping for g: + + G(X) = J^T g(x) J i.e. G_il(X) = J_ji g_jk(x) J_kl + + 'double contravariant piola' mapping for g: + + G(X) = det(J)^2 K g(x) K^T i.e. G_il(X)=(detJ)^2 K_ij g_jk K_lk + + 'covariant contravariant piola' mapping for g: + + G(X) = det(J) J^T g(x) K^T i.e. G_il(X) = det(J) J_ji g_jk(x) K_lk + + If 'contravariant piola' or 'covariant piola' (or their double + variants) are applied to a matrix-valued function, the appropriate + mappings are applied row-by-row. + """ + mesh = extract_unique_domain(expression) + if mesh is None: + mesh = domain + if domain is not None and mesh != domain: + raise NotImplementedError("Multiple domains not supported") + pvs = element.pullback.physical_value_shape(element, mesh) + if expression.ufl_shape != pvs: + raise ValueError(f"Mismatching shapes, got {expression.ufl_shape}, expected {pvs}") + mapping = element.mapping().lower() + if mapping == "identity": + rexpression = expression + elif mapping == "covariant piola": + J = Jacobian(mesh) + *k, i, j = indices(len(expression.ufl_shape) + 1) + kj = (*k, j) + rexpression = as_tensor(J[j, i] * expression[kj], (*k, i)) + elif mapping == "l2 piola": + detJ = JacobianDeterminant(mesh) + rexpression = expression * detJ + elif mapping == "contravariant piola": + K = JacobianInverse(mesh) + detJ = JacobianDeterminant(mesh) + *k, i, j = indices(len(expression.ufl_shape) + 1) + kj = (*k, j) + rexpression = as_tensor(detJ * K[i, j] * expression[kj], (*k, i)) + elif mapping == "double covariant piola": + J = Jacobian(mesh) + *k, i, j, m, n = indices(len(expression.ufl_shape) + 2) + kmn = (*k, m, n) + rexpression = as_tensor(J[m, i] * expression[kmn] * J[n, j], (*k, i, j)) + elif mapping == "double contravariant piola": + K = JacobianInverse(mesh) + detJ = JacobianDeterminant(mesh) + *k, i, j, m, n = indices(len(expression.ufl_shape) + 2) + kmn = (*k, m, n) + rexpression = as_tensor(detJ**2 * K[i, m] * expression[kmn] * K[j, n], (*k, i, j)) + elif mapping == "covariant contravariant piola": + J = Jacobian(mesh) + K = JacobianInverse(mesh) + detJ = JacobianDeterminant(mesh) + *k, i, j, m, n = indices(len(expression.ufl_shape) + 2) + kmn = (*k, m, n) + rexpression = as_tensor(detJ * J[m, i] * expression[kmn] * K[j, n], (*k, i, j)) + elif mapping == "symmetries": + # This tells us how to get from the pieces of the reference + # space expression to the physical space one. + # We're going to apply the inverse of the physical to + # reference space mapping. + fcm = element.flattened_sub_element_mapping() + sub_elem = element.sub_elements[0] + shape = expression.ufl_shape + flat = ufl.as_vector([expression[i] for i in numpy.ndindex(shape)]) + vs = sub_elem.pullback.physical_value_shape(sub_elem, mesh) + rvs = sub_elem.reference_value_shape + seen = set() + rpieces = [] + gm = int(numpy.prod(vs, dtype=int)) + for gi, ri in enumerate(fcm): + # For each unique piece in reference space + if ri in seen: + continue + seen.add(ri) + # Get the physical space piece + piece = [flat[gm*gi + j] for j in range(gm)] + piece = as_tensor(numpy.asarray(piece).reshape(vs)) + # get into reference space + piece = apply_mapping(piece, sub_elem, mesh) + assert piece.ufl_shape == rvs + # Concatenate with the other pieces + rpieces.extend([piece[idx] for idx in numpy.ndindex(rvs)]) + # And reshape + rexpression = as_tensor(numpy.asarray(rpieces).reshape(element.reference_value_shape)) + else: + raise NotImplementedError(f"Don't know how to handle mapping type {mapping} for expression of rank {ufl.FunctionSpace(mesh, element).value_shape}") + if rexpression.ufl_shape != element.reference_value_shape: + raise ValueError(f"Mismatching reference shapes, got {rexpression.ufl_shape} expected {element.reference_value_shape}") + return rexpression + + +class TSFCConstantMixin: + """ Mixin class to identify Constants """ + + def __init__(self): + pass diff --git a/tsfc/vanilla.py b/tsfc/vanilla.py new file mode 100644 index 0000000000..ecf24cd5e4 --- /dev/null +++ b/tsfc/vanilla.py @@ -0,0 +1,47 @@ +from functools import reduce + +from gem import index_sum, Sum +from gem.optimise import unroll_indexsum +from gem.unconcatenate import unconcatenate + + +def Integrals(expressions, quadrature_multiindex, argument_multiindices, parameters): + """Constructs an integral representation for each GEM integrand + expression. + + :arg expressions: integrand multiplied with quadrature weight; + multi-root GEM expression DAG + :arg quadrature_multiindex: quadrature multiindex (tuple) + :arg argument_multiindices: tuple of argument multiindices, + one multiindex for each argument + :arg parameters: parameters dictionary + + :returns: list of integral representations + """ + # Unroll + max_extent = parameters["unroll_indexsum"] + if max_extent: + def predicate(index): + return index.extent <= max_extent + expressions = unroll_indexsum(expressions, predicate=predicate) + # Integral representation: just a GEM expression + return [index_sum(e, quadrature_multiindex) for e in expressions] + + +def flatten(var_reps, index_cache): + """Flatten mode-specific intermediate representation to a series of + assignments. + + :arg var_reps: series of (return variable, [integral representation]) pairs + :arg index_cache: cache :py:class:`dict` for :py:func:`unconcatenate` + + :returns: series of (return variable, GEM expression root) pairs + """ + return unconcatenate([(variable, reduce(Sum, reps)) + for variable, reps in var_reps], + cache=index_cache) + + +finalise_options = dict(remove_componenttensors=False) +"""To avoid duplicate work, these options that are safe to pass to +:py:func:`gem.impero_utils.preprocess_gem`."""