temp

ufl/coefficient.py

@@ -11,6 +11,7 @@
 # Modified by Cecile Daversin-Catty, 2018.
 # Modified by Ignacia Fierro-Piccardo 2023.
 
+import numpy as np
 from ufl.argument import Argument
 from ufl.core.terminal import FormArgument
 from ufl.core.ufl_type import ufl_type
@@ -20,6 +21,7 @@
 from ufl.split_functions import split
 from ufl.utils.counted import Counted
 
+
 # --- The Coefficient class represents a coefficient in a form ---
 
 
@@ -201,6 +203,82 @@ def __repr__(self):
         """Representation."""
         return self._repr
 
+    def traverse_dag_apply_coefficient_split(
+        self,
+        coefficient_split,
+        reference_value=False,
+        reference_grad=0,
+        restricted=None,
+        cache=None,
+    ):
+        from ufl.classes import (
+            ComponentTensor,
+            MultiIndex,
+            NegativeRestricted,
+            PositiveRestricted,
+            ReferenceGrad,
+            ReferenceValue,
+            Zero,
+        )
+        from ufl.core.multiindex import indices
+        from ufl.checks import is_cellwise_constant
+        from ufl.domain import extract_unique_domain
+        from ufl.tensors import as_tensor
+
+        if self not in coefficient_split:
+            c = self
+            if reference_value:
+                c = ReferenceValue(c)
+            for _ in range(reference_grad):
+                # 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(c):
+                    dim = extract_unique_domain(subcoeff).topological_dimension()
+                    c = Zero(c.ufl_shape + (dim,), c.ufl_free_indices, c.ufl_index_dimensions)
+                else:
+                    c = ReferenceGrad(c)
+            if restricted == "+":
+                c = PositiveRestricted(c)
+            elif restricted == "-":
+                c = NegativeRestricted(c)
+            elif restricted is not None:
+                raise RuntimeError(f"Got unknown restriction: {restricted}")
+            return c
+        # Reference value expected
+        if not reference_value:
+            raise RuntimeError(f"ReferenceValue expected: got {o}")
+        # Derivative indices
+        beta = indices(reference_grad)
+        components = []
+        for subcoeff in coefficient_split[self]:
+            c = subcoeff
+            # Apply terminal modifiers onto the subcoefficient
+            if reference_value:
+                c = ReferenceValue(c)
+            for _ in range(reference_grad):
+                # 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(c):
+                    dim = extract_unique_domain(subcoeff).topological_dimension()
+                    c = Zero(c.ufl_shape + (dim,), c.ufl_free_indices, c.ufl_index_dimensions)
+                else:
+                    c = ReferenceGrad(c)
+            if restricted == "+":
+                c = PositiveRestricted(c)
+            elif restricted == "-":
+                c = NegativeRestricted(c)
+            elif restricted is not None:
+                raise RuntimeError(f"Got unknown restriction: {restricted}")
+            # Collect components of the subcoefficient
+            for alpha in np.ndindex(subcoeff.ufl_element().reference_value_shape):
+                # New modified terminal: component[alpha + beta]
+                components.append(c[alpha + beta])
+        # Repack derivative indices to shape
+        i, = indices(1)
+        return ComponentTensor(as_tensor(components)[i], MultiIndex((i,) + beta))
+
 
 # --- Helper functions for subfunctions on mixed elements ---
 

ufl/core/expr.py

@@ -388,6 +388,16 @@ def __round__(self, n=None):
             val = NotImplemented
         return val
 
+    def traverse_dag_apply_coefficient_split(self, *args, **kwargs):
+        ops = [
+            op.traverse_dag_apply_coefficient_split(*args, **kwargs)
+            for op in self.ufl_operands
+        ]
+        if all(a is b for a, b in zip(self.ufl_operands, ops)):
+            return self
+        else:
+            return self._ufl_expr_reconstruct_(*ops)
+
 
 # Initializing traits here because Expr is not defined in the class
 # declaration

ufl/core/terminal.py

@@ -13,6 +13,7 @@
 
 import warnings
 
+import numpy as np
 from ufl.core.expr import Expr
 from ufl.core.ufl_type import ufl_type
 
@@ -92,6 +93,47 @@ def __eq__(self, other):
         """Default comparison of terminals just compare repr strings."""
         return repr(self) == repr(other)
 
+    def traverse_dag_apply_coefficient_split(
+        self,
+        coefficient_split,
+        reference_value=False,
+        reference_grad=0,
+        restricted=None,
+        cache=None,
+    ):
+        from ufl.classes import (
+            ComponentTensor,
+            MultiIndex,
+            NegativeRestricted,
+            PositiveRestricted,
+            ReferenceGrad,
+            ReferenceValue,
+            Zero,
+        )
+        from ufl.core.multiindex import indices
+        from ufl.checks import is_cellwise_constant
+        from ufl.domain import extract_unique_domain
+        from ufl.tensors import as_tensor
+
+        c = self
+        if reference_value:
+            c = ReferenceValue(c)
+        for _ in range(reference_grad):
+            # 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(c):
+                dim = extract_unique_domain(subcoeff).topological_dimension()
+                c = Zero(c.ufl_shape + (dim,), c.ufl_free_indices, c.ufl_index_dimensions)
+            else:
+                c = ReferenceGrad(c)
+        if restricted == "+":
+            c = PositiveRestricted(c)
+        elif restricted == "-":
+            c = NegativeRestricted(c)
+        elif restricted is not None:
+            raise RuntimeError(f"Got unknown restriction: {restricted}")
+        return c
 
 # --- Subgroups of terminals ---
 

ufl/differentiation.py

@@ -344,6 +344,24 @@ def __str__(self):
         """Format as a string."""
         return "reference_grad(%s)" % self.ufl_operands[0]
 
+    def traverse_dag_apply_coefficient_split(
+        self,
+        coefficient_split,
+        reference_value=False,
+        reference_grad=0,
+        restricted=None,
+        cache=None,
+    ):
+        op, = self.ufl_operands
+        if not op._ufl_terminal_modifiers_:
+            raise ValueError(f"Expecting a terminal modifier: got {op!r}.")
+        return op.traverse_dag_apply_coefficient_split(
+            coefficient_split,
+            reference_value=reference_value,
+            reference_grad=reference_grad + 1,
+            restricted=restricted,
+            cache=cache,
+        )
 
 @ufl_type(num_ops=1, inherit_indices_from_operand=0, is_terminal_modifier=True)
 class Div(CompoundDerivative):

ufl/referencevalue.py

@@ -34,3 +34,24 @@ def evaluate(self, x, mapping, component, index_values, derivatives=()):
     def __str__(self):
         """Format as a string."""
         return f"reference_value({self.ufl_operands[0]})"
+
+    def traverse_dag_apply_coefficient_split(
+        self,
+        coefficient_split,
+        reference_value=False,
+        reference_grad=0,
+        restricted=None,
+        cache=None,
+    ):
+        if reference_value:
+            raise RuntimeError
+        op, = self.ufl_operands
+        if not op._ufl_terminal_modifiers_:
+            raise ValueError(f"Expecting a terminal modifier: got {op!r}.")
+        return op.traverse_dag_apply_coefficient_split(
+            coefficient_split,
+            reference_value=True,
+            reference_grad=reference_grad,
+            restricted=restricted,
+            cache=cache,
+        )

ufl/restriction.py

@@ -48,6 +48,27 @@ def __str__(self):
         """Format as a string."""
         return f"{parstr(self.ufl_operands[0], self)}({self._side})"
 
+    def traverse_dag_apply_coefficient_split(
+        self,
+        coefficient_split,
+        reference_value=False,
+        reference_grad=0,
+        restricted=None,
+        cache=None,
+    ):
+        if restricted is not None:
+            raise RuntimeError
+        op, = self.ufl_operands
+        if not op._ufl_terminal_modifiers_:
+            raise ValueError(f"Expecting a terminal modifier: got {op!r}.")
+        return op.traverse_dag_apply_coefficient_split(
+            coefficient_split,
+            reference_value=reference_value,
+            reference_grad=reference_grad,
+            restricted=self._side,
+            cache=cache,
+        )
+
 
 @ufl_type(is_terminal_modifier=True)
 class PositiveRestricted(Restricted):