temp

Author: ksagiyam

ufl/coefficient.py

@@ -11,7 +11,9 @@
 # Modified by Cecile Daversin-Catty, 2018.
 # Modified by Ignacia Fierro-Piccardo 2023.
 
+import numpy as np
 from ufl.argument import Argument
+from ufl.core.expr import Expr
 from ufl.core.terminal import FormArgument
 from ufl.core.ufl_type import ufl_type
 from ufl.duals import is_dual, is_primal
@@ -20,6 +22,7 @@
 from ufl.split_functions import split
 from ufl.utils.counted import Counted
 
+
 # --- The Coefficient class represents a coefficient in a form ---
 
 
@@ -201,6 +204,74 @@ def __repr__(self):
         """Representation."""
         return self._repr
 
+    @Expr.traverse_dag
+    @Expr.traverse_dag_apply_coefficient_split_cache
+    def traverse_dag_apply_coefficient_split(
+        self,
+        coefficient_split,
+        reference_value=False,
+        reference_grad=0,
+        restricted=None,
+        cache=None,
+    ):
+        from ufl.classes import (
+            Terminal,
+            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:
+            return Terminal.traverse_dag_apply_coefficient_split(
+                self,
+                coefficient_split,
+                reference_value=reference_value,
+                reference_grad=reference_grad,
+                restricted=restricted,
+                cache=cache,
+            )
+        # 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

@@ -16,11 +16,15 @@
 # Modified by Anders Logg, 2008
 # Modified by Massimiliano Leoni, 2016
 
+import functools
 import warnings
 
 from ufl.core.ufl_type import UFLType, update_ufl_type_attributes
 
 
+PREFIX_TRAVERSE_DAG = "traverse_dag"
+
+
 class Expr(object, metaclass=UFLType):
     """Base class for all UFL expression types.
 
@@ -388,6 +392,84 @@ def __round__(self, n=None):
             val = NotImplemented
         return val
 
+    @staticmethod
+    def traverse_dag(f):
+        """Decorate DAG traversal methods."""
+        @functools.wraps(f)
+        def wrapper(self, *arg, **kwargs):
+            if not f.__name__.startswith(PREFIX_TRAVERSE_DAG):
+                raise RuntimeError(f"""
+                    DAG traversal method name must be prefixed by
+                    '{PREFIX_TRAVERSE_DAG}': got {f.__name__}
+                """)
+            if 'cache' not in kwargs:
+                raise RuntimeError("DAG traversal methods must have cache in kwargs")
+            if kwargs['cache'] is None:
+                # Make a new dict for caching if not passed by the caller.
+                kwargs['cache'] = {}
+            return f(self, *arg, **kwargs)
+        return wrapper
+
+    @traverse_dag
+    def traverse_dag_reuse_if_untouched(self, method_name: str, *args, **kwargs) -> "Expr":
+        """Reuse if untouched.
+
+        Args:
+            method_name: name of the method to be applied to children.
+
+        Returns:
+            Reconstructed Expr or `self` if children are untouched.
+
+        """
+        ops = [
+            getattr(op, method_name)(*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)
+
+    @staticmethod
+    def traverse_dag_apply_coefficient_split_cache(f):
+        """Use method specific key for caching."""
+        @functools.wraps(f)
+        def wrapper(
+            self,
+            coefficient_split,
+            reference_value=False,
+            reference_grad=0,
+            restricted=None,
+            cache=None,
+        ):
+            if cache is None:
+                raise RuntimeError(f"""
+                    Can not have cache=None.
+                    Must decorate {f} with ``Expr.traverse_dag``.
+                """)
+            key = (self, reference_value, reference_grad, restricted)
+            if key in cache:
+                return cache[key]
+            else:
+                result = f(
+                    self,
+                    coefficient_split,
+                    reference_value=reference_value,
+                    reference_grad=reference_grad,
+                    restricted=restricted,
+                    cache=cache,
+                )
+                cache[key] = result
+                return result
+        return wrapper
+
+    def __getattr__(self, name):
+        if name.startswith(PREFIX_TRAVERSE_DAG):
+            # Traverse DAG with reuse_if_untouched by default.
+            return functools.partial(self.traverse_dag_reuse_if_untouched, name)
+        else:
+            raise AttributeError(f"{type(self)} has no attribute {name}")
+
 
 # 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,49 @@ def __eq__(self, other):
         """Default comparison of terminals just compare repr strings."""
         return repr(self) == repr(other)
 
+    @Expr.traverse_dag
+    @Expr.traverse_dag_apply_coefficient_split_cache
+    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,42 @@ def __str__(self):
         """Format as a string."""
         return "reference_grad(%s)" % self.ufl_operands[0]
 
+    @Expr.traverse_dag
+    @Expr.traverse_dag_apply_coefficient_split_cache
+    def traverse_dag_apply_coefficient_split(
+        self,
+        coefficient_split,
+        reference_value=False,
+        reference_grad=0,
+        restricted=None,
+        cache=None,
+    ):
+        """Split mixed coefficients.
+
+        Args:
+            coefficient_split: `dict` that maps mixed coefficients to their components.
+            reference_value: If `ReferenceValue` has been applied.
+            reference_grad: Number of `ReferenceGrad`s that have been applied.
+            restricted: '+', '-', or None.
+            cache: `dict` for caching DAG nodes.
+
+        Returns:
+            This node wrapped with `ReferenceValue` (if ``reference_value``),
+            `ReferenceGrad` (``reference_grad`` times), and `Restricted` (if
+            ``restricted`` is '+' or '-'). The underlying terminal will be
+            decomposed into components according to ``coefficient_split``.
+
+        """
+        op, = self.ufl_operands
+        if not op._ufl_terminal_modifiers_:
+            raise ValueError(f"Must be a terminal modifier: {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

@@ -5,6 +5,7 @@
 #
 # SPDX-License-Identifier:    LGPL-3.0-or-later
 
+from ufl.core.expr import Expr
 from ufl.core.operator import Operator
 from ufl.core.terminal import FormArgument
 from ufl.core.ufl_type import ufl_type
@@ -34,3 +35,42 @@ def evaluate(self, x, mapping, component, index_values, derivatives=()):
     def __str__(self):
         """Format as a string."""
         return f"reference_value({self.ufl_operands[0]})"
+
+    @Expr.traverse_dag
+    @Expr.traverse_dag_apply_coefficient_split_cache
+    def traverse_dag_apply_coefficient_split(
+        self,
+        coefficient_split: dict,
+        reference_value: bool = False,
+        reference_grad: int = 0,
+        restricted: str | None = None,
+        cache=None,
+    ):
+        """Split mixed coefficients.
+
+        Args:
+            coefficient_split: `dict` that maps mixed coefficients to their components.
+            reference_value: If `ReferenceValue` has been applied.
+            reference_grad: Number of `ReferenceGrad`s that have been applied.
+            restricted: '+', '-', or None.
+            cache: `dict` for caching DAG nodes.
+
+        Returns:
+            This node wrapped with `ReferenceValue` (if ``reference_value``),
+            `ReferenceGrad` (``reference_grad`` times), and `Restricted` (if
+            ``restricted`` is '+' or '-'). The underlying terminal will be
+            decomposed into components according to ``coefficient_split``.
+
+        """
+        if reference_value:
+            raise RuntimeError(f"Can not apply ReferenceValue twice: {self}")
+        op, = self.ufl_operands
+        if not op._ufl_terminal_modifiers_:
+            raise ValueError(f"Must be a terminal modifier: {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

@@ -6,6 +6,7 @@
 # SPDX-License-Identifier:    LGPL-3.0-or-later
 
 from ufl.constantvalue import ConstantValue
+from ufl.core.expr import Expr
 from ufl.core.operator import Operator
 from ufl.core.ufl_type import ufl_type
 from ufl.precedence import parstr
@@ -48,6 +49,45 @@ def __str__(self):
         """Format as a string."""
         return f"{parstr(self.ufl_operands[0], self)}({self._side})"
 
+    @Expr.traverse_dag
+    @Expr.traverse_dag_apply_coefficient_split_cache
+    def traverse_dag_apply_coefficient_split(
+        self,
+        coefficient_split,
+        reference_value=False,
+        reference_grad=0,
+        restricted=None,
+        cache=None,
+    ):
+        """Split mixed coefficients.
+
+        Args:
+            coefficient_split: `dict` that maps mixed coefficients to their components.
+            reference_value: If `ReferenceValue` has been applied.
+            reference_grad: Number of `ReferenceGrad`s that have been applied.
+            restricted: '+', '-', or None.
+            cache: `dict` for caching DAG nodes.
+
+        Returns:
+            This node wrapped with `ReferenceValue` (if ``reference_value``),
+            `ReferenceGrad` (``reference_grad`` times), and `Restricted` (if
+            ``restricted`` is '+' or '-'). The underlying terminal will be
+            decomposed into components according to ``coefficient_split``.
+
+        """
+        if restricted is not None:
+            raise RuntimeError
+        op, = self.ufl_operands
+        if not op._ufl_terminal_modifiers_:
+            raise ValueError(f"Must be a terminal modifier: {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):