From 726516fc7b31089a0d1e6ff997b05673d9a43638 Mon Sep 17 00:00:00 2001 From: "Garth N. Wells" Date: Tue, 19 Nov 2024 21:13:11 +0000 Subject: [PATCH 1/6] Update Basix for mixed elements --- python/basix/ufl.py | 27 +++++++++++++++++++++++++-- 1 file changed, 25 insertions(+), 2 deletions(-) diff --git a/python/basix/ufl.py b/python/basix/ufl.py index 4b611b6a5..56f159b24 100644 --- a/python/basix/ufl.py +++ b/python/basix/ufl.py @@ -5,12 +5,14 @@ # SPDX-License-Identifier: MIT """Functions to directly wrap Basix elements in UFL.""" +import collections import hashlib as _hashlib import itertools as _itertools import typing as _typing from abc import abstractmethod as _abstractmethod from abc import abstractproperty as _abstractproperty from warnings import warn as _warn +import functools import numpy as np import numpy.typing as _npt @@ -1983,6 +1985,7 @@ def _compute_signature(element: _basix.finite_element.FiniteElement) -> str: return signature +@functools.singledisptach def element( family: _typing.Union[_basix.ElementFamily, str], cell: _typing.Union[_basix.CellType, str], @@ -2066,6 +2069,22 @@ def element( return blocked_element(ufl_e, shape=shape, symmetry=symmetry) +@element.register(collections.Sequence) +def _(elements: list[_ElementBase]) -> _ElementBase: + """Create a UFL compatible mixed element from a list of elements. + + Args: + elements: List of elements. + + Returns: + A mixed finite element. + """ + if len(elements) > 1: + return _MixedElement(elements) + else: + return elements + + def enriched_element( elements: list[_ElementBase], map_type: _typing.Optional[_basix.MapType] = None, @@ -2229,12 +2248,16 @@ def mixed_element(elements: list[_ElementBase]) -> _ElementBase: """Create a UFL compatible mixed element from a list of elements. Args: - elements: The list of elements + elements: List of elements. Returns: A mixed finite element. """ - return _MixedElement(elements) + return element(elements) + # if len(elements) > 1: + # return _MixedElement(elements) + # else: + # return elements def quadrature_element( From ad38d3a7ed2116d76549ad7a7b7b108149befef1 Mon Sep 17 00:00:00 2001 From: "Garth N. Wells" Date: Tue, 19 Nov 2024 21:21:00 +0000 Subject: [PATCH 2/6] Typing update --- python/basix/ufl.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/basix/ufl.py b/python/basix/ufl.py index 56f159b24..a29ef8c94 100644 --- a/python/basix/ufl.py +++ b/python/basix/ufl.py @@ -2070,7 +2070,7 @@ def element( @element.register(collections.Sequence) -def _(elements: list[_ElementBase]) -> _ElementBase: +def _(elements: _typing.Union[_ElementBase, list[_ElementBase]]) -> _ElementBase: """Create a UFL compatible mixed element from a list of elements. Args: From 679bce765bace9cd95c1792e52c86d6df198f1c5 Mon Sep 17 00:00:00 2001 From: "Garth N. Wells" Date: Tue, 19 Nov 2024 21:22:52 +0000 Subject: [PATCH 3/6] Fix typo --- python/basix/ufl.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/basix/ufl.py b/python/basix/ufl.py index a29ef8c94..f31557487 100644 --- a/python/basix/ufl.py +++ b/python/basix/ufl.py @@ -1985,7 +1985,7 @@ def _compute_signature(element: _basix.finite_element.FiniteElement) -> str: return signature -@functools.singledisptach +@functools.singledispatch def element( family: _typing.Union[_basix.ElementFamily, str], cell: _typing.Union[_basix.CellType, str], From a2f6235e9fe85cbc01336b5a0d75638aad8219e6 Mon Sep 17 00:00:00 2001 From: "Garth N. Wells" Date: Tue, 19 Nov 2024 21:29:53 +0000 Subject: [PATCH 4/6] Remove single dispatch --- python/basix/ufl.py | 28 ++++------------------------ 1 file changed, 4 insertions(+), 24 deletions(-) diff --git a/python/basix/ufl.py b/python/basix/ufl.py index f31557487..152c03825 100644 --- a/python/basix/ufl.py +++ b/python/basix/ufl.py @@ -5,14 +5,12 @@ # SPDX-License-Identifier: MIT """Functions to directly wrap Basix elements in UFL.""" -import collections import hashlib as _hashlib import itertools as _itertools import typing as _typing from abc import abstractmethod as _abstractmethod from abc import abstractproperty as _abstractproperty from warnings import warn as _warn -import functools import numpy as np import numpy.typing as _npt @@ -1985,7 +1983,6 @@ def _compute_signature(element: _basix.finite_element.FiniteElement) -> str: return signature -@functools.singledispatch def element( family: _typing.Union[_basix.ElementFamily, str], cell: _typing.Union[_basix.CellType, str], @@ -2069,22 +2066,6 @@ def element( return blocked_element(ufl_e, shape=shape, symmetry=symmetry) -@element.register(collections.Sequence) -def _(elements: _typing.Union[_ElementBase, list[_ElementBase]]) -> _ElementBase: - """Create a UFL compatible mixed element from a list of elements. - - Args: - elements: List of elements. - - Returns: - A mixed finite element. - """ - if len(elements) > 1: - return _MixedElement(elements) - else: - return elements - - def enriched_element( elements: list[_ElementBase], map_type: _typing.Optional[_basix.MapType] = None, @@ -2253,11 +2234,10 @@ def mixed_element(elements: list[_ElementBase]) -> _ElementBase: Returns: A mixed finite element. """ - return element(elements) - # if len(elements) > 1: - # return _MixedElement(elements) - # else: - # return elements + if len(elements) > 1: + return _MixedElement(elements) + else: + return elements def quadrature_element( From fbb37f883a6d995ed5e3f8d1d3012fc573dd59ed Mon Sep 17 00:00:00 2001 From: "Garth N. Wells" Date: Tue, 19 Nov 2024 21:39:52 +0000 Subject: [PATCH 5/6] Logic fix --- python/basix/ufl.py | 2 +- test/test_ufl_wrapper.py | 12 ++++++++++++ 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/python/basix/ufl.py b/python/basix/ufl.py index 152c03825..eb5f80287 100644 --- a/python/basix/ufl.py +++ b/python/basix/ufl.py @@ -2237,7 +2237,7 @@ def mixed_element(elements: list[_ElementBase]) -> _ElementBase: if len(elements) > 1: return _MixedElement(elements) else: - return elements + return elements[0] def quadrature_element( diff --git a/test/test_ufl_wrapper.py b/test/test_ufl_wrapper.py index fe533adce..7252b3cbc 100644 --- a/test/test_ufl_wrapper.py +++ b/test/test_ufl_wrapper.py @@ -98,11 +98,23 @@ def test_enriched_element(elements): "e,space0,space1", [ (basix.ufl.element("Lagrange", basix.CellType.triangle, 2), "H1", basix.SobolevSpace.H1), + ( + basix.ufl.mixed_element([basix.ufl.element("Lagrange", basix.CellType.triangle, 2)]), + "H1", + basix.SobolevSpace.H1, + ), ( basix.ufl.element("Discontinuous Lagrange", basix.CellType.triangle, 0), "L2", basix.SobolevSpace.L2, ), + ( + basix.ufl.mixed_elementelement( + [basix.ufl.element("Discontinuous Lagrange", basix.CellType.triangle, 0)] + ), + "L2", + basix.SobolevSpace.L2, + ), ( basix.ufl.mixed_element( [ From 4e5c21874532cbeff6d04853d594317984d5c2ff Mon Sep 17 00:00:00 2001 From: "Garth N. Wells" Date: Tue, 19 Nov 2024 21:40:43 +0000 Subject: [PATCH 6/6] Test fix --- test/test_ufl_wrapper.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/test_ufl_wrapper.py b/test/test_ufl_wrapper.py index 7252b3cbc..a812c1d30 100644 --- a/test/test_ufl_wrapper.py +++ b/test/test_ufl_wrapper.py @@ -109,7 +109,7 @@ def test_enriched_element(elements): basix.SobolevSpace.L2, ), ( - basix.ufl.mixed_elementelement( + basix.ufl.mixed_element( [basix.ufl.element("Discontinuous Lagrange", basix.CellType.triangle, 0)] ), "L2",