Pytorch coupling

.github/workflows/build.yml

@@ -48,13 +48,14 @@ jobs:
       - name: Build Firedrake
         run: |
           cd ..
-          ./firedrake/scripts/firedrake-install $COMPLEX --venv-name build --tinyasm --disable-ssh --minimal-petsc --slepc --documentation-dependencies --install thetis --install gusto --install icepack --install irksome --install femlium --no-package-manager || (cat firedrake-install.log && /bin/false)
+          ./firedrake/scripts/firedrake-install $COMPLEX --venv-name build --tinyasm --disable-ssh --minimal-petsc --slepc --documentation-dependencies --install thetis --install gusto --install icepack --install irksome --install femlium --no-package-manager --package-branch pyadjoint adjoint-1-forms || (cat firedrake-install.log && /bin/false)
       - name: Install test dependencies
         run: |
           . ../build/bin/activate
           python $(which firedrake-clean)
           python -m pip install pytest-cov pytest-timeout pytest-xdist pytest-timeout
           python -m pip list
+          python -m pip install torch --extra-index-url https://download.pytorch.org/whl/cpu
       - name: Test Firedrake
         run: |
           . ../build/bin/activate

firedrake/__init__.py

@@ -103,6 +103,7 @@
 from firedrake.ensemble import *
 from firedrake.randomfunctiongen import *
 from firedrake.progress_bar import ProgressBar  # noqa: F401
+from firedrake.pytorch_coupling import *
 
 from firedrake.logging import *
 # Set default log level

firedrake/pytorch_coupling/__init__.py

@@ -0,0 +1,2 @@
+from firedrake.pytorch_coupling.backends import get_backend                    # noqa: F401
+from firedrake.pytorch_coupling.pytorch_custom_operator import torch_operator  # noqa: F401

firedrake/pytorch_coupling/backends.py

@@ -0,0 +1,120 @@
+from firedrake.function import Function
+from firedrake.vector import Vector
+from firedrake.constant import Constant
+
+import firedrake.utils as utils
+
+
+class AbstractMLBackend(object):
+
+    def backend(self):
+        raise NotImplementedError
+
+    def to_ml_backend(self, x):
+        r"""Convert from Firedrake to ML backend.
+
+           :arg x: Firedrake object
+        """
+        raise NotImplementedError
+
+    def from_ml_backend(self, x, V):
+        r"""Convert from ML backend to Firedrake.
+
+           :arg x: ML backend object
+        """
+        raise NotImplementedError
+
+    def get_function_space(self, x):
+        """Get function space out of x"""
+        if isinstance(x, Function):
+            return x.function_space()
+        elif isinstance(x, Vector):
+            return self.get_function_space(x.function)
+        elif isinstance(x, float):
+            return None
+        else:
+            raise ValueError("Cannot infer the function space of %s" % x)
+
+
+class PytorchBackend(AbstractMLBackend):
+
+    def __init__(self):
+        try:
+            import torch
+            self._backend = torch
+        except ImportError:
+            self._backend = None
+
+    @property
+    def backend(self):
+        if self:
+            return self._backend
+        else:
+            raise ImportError("Error when trying to import PyTorch.")
+
+    def __bool__(self):
+        return self._backend is not None
+
+    @utils.cached_property
+    def custom_operator(self):
+        from firedrake.pytorch_coupling.pytorch_custom_operator import FiredrakeTorchOperator
+        return FiredrakeTorchOperator().apply
+
+    def to_ml_backend(self, x, gather=False, batched=True, **kwargs):
+        r"""Convert a Firedrake object `x` into a PyTorch tensor.
+
+            :arg x: Firedrake object (Function, Vector, Constant)
+            :kwarg gather: if True, gather data from all processes
+            :kwarg batched: if True, add a batch dimension to the tensor
+            :kwarg kwargs: additional arguments to be passed to torch.Tensor constructor
+                - device: device on which the tensor is allocated (default: "cpu")
+                - dtype: the desired data type of returned tensor (default: type of x.dat.data)
+                - requires_grad: if the tensor should be annotated (default: False)
+        """
+        if isinstance(x, (Function, Vector)):
+            if gather:
+                # Gather data from all processes
+                x_P = self.backend.tensor(x.vector().gather(), **kwargs)
+            else:
+                # Use local data
+                x_P = self.backend.tensor(x.vector().get_local(), **kwargs)
+            if batched:
+                # Default behaviour: add batch dimension after converting to PyTorch
+                return x_P[None, :]
+            return x_P
+        elif isinstance(x, Constant):
+            return self.backend.tensor(x.values(), **kwargs)
+        elif isinstance(x, (float, int)):
+            if isinstance(x, float):
+                # Set double-precision
+                kwargs['dtype'] = self.backend.double
+            return self.backend.tensor(x, **kwargs)
+        else:
+            raise ValueError("Cannot convert %s to a torch tensor" % str(type(x)))
+
+    def from_ml_backend(self, x, V=None):
+        r"""Convert a PyTorch tensor `x` into a Firedrake object.
+
+            :arg x: PyTorch tensor (torch.Tensor)
+            :kwarg V: function space of the corresponding Function or None when `x` is to be mapped to a Constant
+        """
+        if x.device.type != "cpu":
+            raise NotImplementedError("Firedrake does not support GPU tensors")
+
+        if V is None:
+            val = x.detach().numpy()
+            if val.shape == (1,):
+                val = val[0]
+            return Constant(val)
+        else:
+            x = x.detach().numpy()
+            x_F = Function(V, dtype=x.dtype)
+            x_F.vector().set_local(x)
+            return x_F
+
+
+def get_backend(backend_name='pytorch'):
+    if backend_name == 'pytorch':
+        return PytorchBackend()
+    else:
+        raise NotImplementedError("The backend: %s is not supported." % backend_name)

firedrake/pytorch_coupling/pytorch_custom_operator.py

@@ -0,0 +1,86 @@
+import collections
+from functools import partial
+
+from firedrake.pytorch_coupling import get_backend
+from firedrake.function import Function
+
+from pyadjoint.reduced_functional import ReducedFunctional
+
+
+backend = get_backend("pytorch")
+
+if backend:
+    # PyTorch is installed
+    BackendFunction = backend.backend.autograd.Function
+else:
+    class BackendFunction(object):
+        """Dummy class that exceptions on instantiation."""
+        def __init__(self):
+            raise ImportError("PyTorch is not installed and is required to use the FiredrakeTorchOperator.")
+
+
+class FiredrakeTorchOperator(BackendFunction):
+    """
+    PyTorch custom operator representing a set of Firedrake operations expressed as a ReducedFunctional F.
+    `FiredrakeTorchOperator` is a wrapper around `torch.autograd.Function` that executes forward and backward
+    passes by directly calling the reduced functional F.
+
+    Inputs:
+        metadata: dictionary used to stash Firedrake objects.
+        *ω: PyTorch tensors representing the inputs to the Firedrake operator F
+
+    Outputs:
+        y: PyTorch tensor representing the output of the Firedrake operator F
+    """
+
+    def __init__(self):
+        super(FiredrakeTorchOperator, self).__init__()
+
+    # This method is wrapped by something cancelling annotation (probably 'with torch.no_grad()')
+    @staticmethod
+    def forward(ctx, metadata, *ω):
+        """Forward pass of the PyTorch custom operator."""
+        F = metadata['F']
+        V = metadata['V_controls']
+        # Convert PyTorch input (i.e. controls) to Firedrake
+        ω_F = [backend.from_ml_backend(ωi, Vi) for ωi, Vi in zip(ω, V)]
+        # Forward operator: delegated to pyadjoint.ReducedFunctional which recomputes the blocks on the tape
+        y_F = F(ω_F)
+        # Stash metadata to the PyTorch context
+        ctx.metadata.update(metadata)
+        # Convert Firedrake output to PyTorch
+        y = backend.to_ml_backend(y_F)
+        return y.detach()
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        """Backward pass of the PyTorch custom operator."""
+        F = ctx.metadata['F']
+        V = ctx.metadata['V_output']
+        # Convert PyTorch gradient to Firedrake
+        adj_input = backend.from_ml_backend(grad_output, V)
+        if isinstance(adj_input, Function):
+            adj_input = adj_input.vector()
+
+        # Compute adjoint model of `F`: delegated to pyadjoint.ReducedFunctional
+        Δω = F.derivative(adj_input=adj_input)
+
+        # Tuplify adjoint output
+        Δω = (Δω,) if not isinstance(Δω, collections.abc.Sequence) else Δω
+
+        # None is for metadata arg in `forward`
+        return None, *[backend.to_ml_backend(Δωi) for Δωi in Δω]
+
+
+def torch_operator(F):
+    """Operator that converts a pyadjoint.ReducedFunctional into a firedrake.FiredrakeTorchOperator
+       whose inputs and outputs are PyTorch tensors.
+    """
+    if not isinstance(F, ReducedFunctional):
+        raise ValueError("F must be a ReducedFunctional")
+
+    V_output = backend.get_function_space(F.functional)
+    V_controls = [c.control.function_space() for c in F.controls]
+    metadata = {'F': F, 'V_controls': V_controls, 'V_output': V_output}
+    φ = partial(backend.custom_operator, metadata)
+    return φ

tests/conftest.py

@@ -44,6 +44,9 @@ def pytest_configure(config):
     config.addinivalue_line(
         "markers",
         "skipcomplexnoslate: mark as skipped in complex mode due to lack of Slate")
+    config.addinivalue_line(
+        "markers",
+        "skiptorch: mark as skipped if PyTorch is not installed")
 
 
 @pytest.fixture(autouse=True)
@@ -79,6 +82,10 @@ def pytest_runtest_call(item):
 
 def pytest_collection_modifyitems(session, config, items):
     from firedrake.utils import SLATE_SUPPORTS_COMPLEX
+    from firedrake.pytorch_coupling import get_backend
+
+    backend = get_backend("pytorch")
+
     for item in items:
         if complex_mode:
             if item.get_closest_marker("skipcomplex") is not None:
@@ -89,6 +96,10 @@ def pytest_collection_modifyitems(session, config, items):
             if item.get_closest_marker("skipreal") is not None:
                 item.add_marker(pytest.mark.skip(reason="Test makes no sense unless in complex mode"))
 
+        if not backend:
+            if item.get_closest_marker("skiptorch") is not None:
+                item.add_marker(pytest.mark.skip(reason="Test makes no sense if PyTorch is not installed"))
+
 
 @pytest.fixture(scope="module", autouse=True)
 def check_empty_tape(request):