Merge branch 'master' into pbrubeck/fix/base-form-tensor

Author: pbrubeck

docs/source/images/sin_mesh.png

@@ -83,6 +83,30 @@ Or simply:
 
    new_mesh = Mesh(Function(f))
 
+Immersing a mesh in higher dimensional space
+--------------------------------------------
+
+A useful special case of creating a mesh on modified coordinates is to immerse
+a lower dimensional mesh in a higher dimension, for example to create a mesh of
+a two-dimensional manifold immersed in 3D.
+
+This is accomplished by setting the value dimension of the new
+:py:func:`~.VectorFunctionSpace` to that of the space in which it should be
+immersed. For example, a mesh of square bent into a sine wave using
+linear (flat) elements can be created with:
+
+.. literalinclude:: ../../tests/firedrake/regression/test_mesh_generation.py
+   :language: python3
+   :dedent:
+   :start-after: start_immerse
+   :end-before: end_immerse
+
+
+.. figure:: images/sin_mesh.png
+  :align: center
+
+  A sine-wave shaped triangle mesh immersed in three-dimensional space.
+
 
 Replacing the mesh geometry of an existing function
 ---------------------------------------------------

docs/source/mesh-coordinates.rst

@@ -47,21 +47,22 @@ def __init__(self, *operands, function_space, derivatives=None, argument_slots=(
         Parameters
         ----------
         *operands : ufl.core.expr.Expr or ufl.form.BaseForm
-                    Operands of the external operator.
+            Operands of the external operator.
         function_space : firedrake.functionspaceimpl.WithGeometryBase
-                         The function space the external operator is mapping to.
+            The function space the external operator is mapping to.
         derivatives : tuple
-                      Tuple specifiying the derivative multiindex.
+            Tuple specifiying the derivative multiindex.
         *argument_slots : ufl.coefficient.BaseCoefficient or ufl.argument.BaseArgument
-                          Tuple containing the arguments of the linear form associated with the external operator,
-                          i.e. the arguments with respect to which the external operator is linear. Those arguments
-                          can be ufl.Argument objects, as a result of differentiation, or ufl.Coefficient objects,
-                          as a result of taking the action on a given function.
+            Tuple containing the arguments of the linear form associated with the external operator,
+            i.e. the arguments with respect to which the external operator is linear. Those arguments can
+            be ``ufl.argument.BaseArgument`` objects, as a result of differentiation, or both
+            ``ufl.coefficient.BaseCoefficient`` and ``ufl.argument.BaseArgument`` object, as a result
+            of taking the action on a given function.
         operator_data : dict
-                        Dictionary containing the data of the external operator, i.e. the external data
-                        specific to the external operator subclass considered. This dictionary will be passed on
-                        over the UFL symbolic reconstructions making the operator data accessible to the external operators
-                        arising from symbolic operations on the original operator, such as the Jacobian of the external operator.
+            Dictionary containing the data of the external operator, i.e. the external data
+            specific to the external operator subclass considered. This dictionary will be passed on
+            over the UFL symbolic reconstructions making the operator data accessible to the external operators
+            arising from symbolic operations on the original operator, such as the Jacobian of the external operator.
         """
         from firedrake_citations import Citations
         Citations().register("Bouziani2021")

firedrake/external_operators/abstract_external_operators.py

@@ -15,22 +15,25 @@ def __init__(self, *operands, function_space, derivatives=None, argument_slots=(
         Parameters
         ----------
         *operands : ufl.core.expr.Expr or ufl.form.BaseForm
-                    Operands of the ML operator.
+            Operands of the ML operator.
         function_space : firedrake.functionspaceimpl.WithGeometryBase
-                         The function space the ML operator is mapping to.
+            The function space the ML operator is mapping to.
         derivatives : tuple
-                      Tuple specifiying the derivative multiindex.
+            Tuple specifiying the derivative multiindex.
         *argument_slots : ufl.coefficient.BaseCoefficient or ufl.argument.BaseArgument
-                          Tuple containing the arguments of the linear form associated with the ML operator,
-                          i.e. the arguments with respect to which the ML operator is linear. Those arguments
-                          can be ufl.Argument objects, as a result of differentiation, or ufl.Coefficient objects,
-                          as a result of taking the action on a given function.
+            Tuple containing the arguments of the linear form associated with the ML operator,
+            i.e. the arguments with respect to which the ML operator is linear. Those arguments can
+            be ``ufl.argument.BaseArgument`` objects, as a result of differentiation, or both
+            ``ufl.coefficient.BaseCoefficient`` and ``ufl.argument.BaseArgument`` object, as a result
+            of taking the action on a given function. If argument slots are not provided, then they will
+            be generated in the :class:`.AbstractExternalOperator` constructor.
         operator_data : dict
-                        Dictionary to stash external data specific to the ML operator. This dictionary must
-                        at least contain the following:
-                        (i) 'model': The machine learning model implemented in the ML framework considered.
-                        (ii) 'inputs_format': The format of the inputs to the ML model: `0` for models acting globally on the inputs, `1` when acting locally/pointwise on the inputs.
-                        Other strategies can also be considered by subclassing the :class:`.MLOperator` class.
+            Dictionary to stash external data specific to the ML operator. This dictionary must
+            at least contain the following:
+            (i) 'model': The machine learning model implemented in the ML framework considered.
+            (ii) 'inputs_format': The format of the inputs to the ML model: ``0`` for models acting globally
+            on the inputs, ``1`` when acting locally/pointwise on the inputs.
+            Other strategies can also be considered by subclassing the :class:`.MLOperator` class.
         """
         AbstractExternalOperator.__init__(self, *operands, function_space=function_space, derivatives=derivatives,
                                           argument_slots=argument_slots, operator_data=operator_data)

firedrake/external_operators/ml_operator.py

@@ -92,7 +92,7 @@ def bwd(self, _, grad_output: "jax.Array") -> "jax.Array":
             adj_input = float(adj_input)
 
         # Compute adjoint model of `F`: delegated to pyadjoint.ReducedFunctional
-        adj_output = self.F.derivative(adj_input=adj_input)
+        adj_output = self.F.derivative(adj_input=adj_input, options={'riesz_representation': None})
 
         # Tuplify adjoint output
         adj_output = (adj_output,) if not isinstance(adj_output, collections.abc.Sequence) else adj_output

firedrake/ml/jax/fem_operator.py

@@ -39,36 +39,40 @@ def __init__(
             *operands: Union[ufl.core.expr.Expr, ufl.form.BaseForm],
             function_space: WithGeometryBase,
             derivatives: Optional[tuple] = None,
-            argument_slots: Optional[tuple[Union[ufl.coefficient.BaseCoefficient, ufl.argument.BaseArgument]]],
+            argument_slots: tuple[Union[ufl.coefficient.BaseCoefficient, ufl.argument.BaseArgument]] = (),
             operator_data: Optional[dict] = {}
     ):
-        """External operator class representing machine learning models implemented in JAX.
+        """
+        External operator class representing machine learning models implemented in JAX.
 
         The :class:`.JaxOperator` allows users to embed machine learning models implemented in JAX
-        into PDE systems implemented in Firedrake. The actual evaluation of the :class:`.JaxOperator` is
-        delegated to the specified JAX model. Similarly, differentiation through the :class:`.JaxOperator`
-        class is achieved using JAX differentiation on the JAX model associated with the :class:`.JaxOperator` object.
+        into PDE systems implemented in Firedrake. The actual evaluation of the :class:`.JaxOperator`
+        is delegated to the specified JAX model. Similarly, differentiation through the
+        :class:`.JaxOperator` is achieved using JAX differentiation on the associated JAX model.
 
         Parameters
         ----------
         *operands
-                    Operands of the :class:`.JaxOperator`.
+            Operands of the :class:`.JaxOperator`.
         function_space
-                         The function space the ML operator is mapping to.
+            The function space the ML operator is mapping to.
         derivatives
-                      Tuple specifiying the derivative multiindex.
-        *argument_slots
-                          Tuple containing the arguments of the linear form associated with the ML operator,
-                          i.e. the arguments with respect to which the ML operator is linear. Those arguments
-                          can be ufl.Argument objects, as a result of differentiation, or ufl.Coefficient objects,
-                          as a result of taking the action on a given function.
+            Tuple specifying the derivative multi-index.
+        argument_slots
+            Tuple containing the arguments of the linear form associated with the ML operator,
+            i.e., the arguments with respect to which the ML operator is linear. These arguments
+            can be ``ufl.argument.BaseArgument`` objects, as a result of differentiation,
+            or both ``ufl.coefficient.BaseCoefficient`` and ``ufl.argument.BaseArgument`` objects,
+            as a result of taking the action on a given function.
         operator_data
-                        Dictionary to stash external data specific to the ML operator. This dictionary must
-                        at least contain the following:
-                        (i) 'model': The machine learning model implemented in JaX
-                        (ii) 'inputs_format': The format of the inputs to the ML model: `0` for models acting globally on the inputs, `1` when acting locally/pointwise on the inputs.
-                        Other strategies can also be considered by subclassing the :class:`.JaxOperator` class.
+            Dictionary to stash external data specific to the ML operator. This dictionary must
+            contain the following:
+            (i) ``'model'`` : The machine learning model implemented in JaX.
+            (ii) ``'model'`` : The format of the inputs to the ML model: ``0`` for models acting
+            globally on the inputs. ``1`` for models acting locally/pointwise on the inputs.
+            Other strategies can also be considered by subclassing the :class:`.JaxOperator` class.
         """
+
         MLOperator.__init__(self, *operands, function_space=function_space, derivatives=derivatives,
                             argument_slots=argument_slots, operator_data=operator_data)
 
@@ -90,8 +94,7 @@ def _pre_forward_callback(self, *operands: Union[Function, Cofunction], unsqueez
 
     def _post_forward_callback(self, y_P: "jax.Array") -> Union[Function, Cofunction]:
         """Callback function to convert the JAX output of the ML model to a Firedrake function."""
-        space = self.ufl_function_space()
-        return from_jax(y_P, space)
+        return from_jax(y_P, self.ufl_function_space())
 
     # -- JAX routines for computing AD-based quantities -- #
 

firedrake/ml/jax/ml_operator.py

@@ -83,7 +83,7 @@ def backward(ctx, grad_output):
             adj_input = float(adj_input)
 
         # Compute adjoint model of `F`: delegated to pyadjoint.ReducedFunctional
-        adj_output = F.derivative(adj_input=adj_input, options={"riesz_representation": "l2"})
+        adj_output = F.derivative(adj_input=adj_input, options={"riesz_representation": None})
 
         # Tuplify adjoint output
         adj_output = (adj_output,) if not isinstance(adj_output, collections.abc.Sequence) else adj_output

firedrake/ml/pytorch/fem_operator.py

@@ -40,22 +40,24 @@ class is achieved via the `torch.autograd` module, which provides automatic diff
         Parameters
         ----------
         *operands : ufl.core.expr.Expr or ufl.form.BaseForm
-                    Operands of the :class:`.PytorchOperator`.
+            Operands of the :class:`.PytorchOperator`.
         function_space : firedrake.functionspaceimpl.WithGeometryBase
-                         The function space the ML operator is mapping to.
+            The function space the ML operator is mapping to.
         derivatives : tuple
-                      Tuple specifiying the derivative multiindex.
+            Tuple specifiying the derivative multiindex.
         *argument_slots : ufl.coefficient.BaseCoefficient or ufl.argument.BaseArgument
-                          Tuple containing the arguments of the linear form associated with the ML operator,
-                          i.e. the arguments with respect to which the ML operator is linear. Those arguments
-                          can be ufl.Argument objects, as a result of differentiation, or ufl.Coefficient objects,
-                          as a result of taking the action on a given function.
+            Tuple containing the arguments of the linear form associated with the ML operator, i.e. the
+            arguments with respect to which the ML operator is linear. Those arguments can be
+            ``ufl.argument.BaseArgument`` objects, as a result of differentiation, or both
+            ``ufl.coefficient.BaseCoefficient`` and ``ufl.argument.BaseArgument`` object, as a result
+            of taking the action on a given function.
         operator_data : dict
-                        Dictionary to stash external data specific to the ML operator. This dictionary must
-                        at least contain the following:
-                        (i) 'model': The machine learning model implemented in PyTorch.
-                        (ii) 'inputs_format': The format of the inputs to the ML model: `0` for models acting globally on the inputs, `1` when acting locally/pointwise on the inputs.
-                        Other strategies can also be considered by subclassing the :class:`.PytorchOperator` class.
+            Dictionary to stash external data specific to the ML operator. This dictionary must
+            at least contain the following:
+            (i) ``'model'``: The machine learning model implemented in PyTorch.
+            (ii) ``'inputs_format'``: The format of the inputs to the ML model: ``0`` for models acting globally
+            on the inputs, ``1`` when acting locally/pointwise on the inputs.
+            Other strategies can also be considered by subclassing the :class:`.PytorchOperator` class.
         """
         MLOperator.__init__(self, *operands, function_space=function_space, derivatives=derivatives,
                             argument_slots=argument_slots, operator_data=operator_data)
@@ -98,8 +100,7 @@ def _pre_forward_callback(self, *operands, unsqueeze=False):
 
     def _post_forward_callback(self, y_P):
         """Callback function to convert the PyTorch output of the ML model to a Firedrake function."""
-        space = self.ufl_function_space()
-        return from_torch(y_P, space)
+        return from_torch(y_P, self.ufl_function_space())
 
     # -- PyTorch routines for computing AD based quantities via `torch.autograd.functional` -- #
 

firedrake/ml/pytorch/ml_operator.py

@@ -155,7 +155,7 @@ def sniff_compiler(exe, comm=mpi.COMM_WORLD):
         # Find the name of the compiler family
         if output.startswith("gcc") or output.startswith("g++"):
             name = "GNU"
-        elif output.startswith("clang"):
+        elif output.startswith("clang") or output.startswith("Homebrew clang"):
             name = "clang"
         elif output.startswith("Apple LLVM") or output.startswith("Apple clang"):
             name = "clang"

pyop2/compilation.py

@@ -1303,7 +1303,7 @@ def build_and_install_jax():
     """Install JAX for a CPU or CUDA backend."""
     log.info("Installing JAX (backend: %s)" % args.jax)
     version_name = "jax" if args.jax == "cpu" else "jax[cuda12]"
-    run_pip_install([version_name])
+    run_pip_install([version_name] + ["jaxlib"] + ["ml_dtypes"] + ["opt_einsum"])
 
 
 def build_and_install_slepc():