Tensorflow 1.x backend: multiple outputs extension of DeepONet

deepxde/data/pde_operator.py

@@ -236,8 +236,7 @@ def _losses(self, outputs, loss_fn, inputs, model, num_func):
 
         losses = []
         for i in range(num_func):
-            out = outputs[i][:, None]
-
+            out = outputs[i] if model.net.num_outputs > 1 else outputs[i][:, None]
             f = []
             if self.pde.pde is not None:
                 f = self.pde.pde(inputs[1], out, model.net.auxiliary_vars[i][:, None])

deepxde/nn/tensorflow_compat_v1/deeponet.py

@@ -7,6 +7,131 @@
 from ... import config
 from ...backend import tf
 from ...utils import timing
+from abc import ABC, abstractmethod
+
+
+class DeepONetStrategy(ABC):
+    """DeepONet building strategy.
+
+    See the section 3.1.6. in
+    L. Lu, X. Meng, S. Cai, Z. Mao, S. Goswami, Z. Zhang, & G. Karniadakis.
+    A comprehensive and fair comparison of two neural operators
+    (with practical extensions) based on FAIR data.
+    Computer Methods in Applied Mechanics and Engineering, 393, 114778, 2022.
+    """
+
+    def __init__(self, net):
+        self.net = net
+
+    def _build_branch_and_trunk(self):
+        # Branch net to encode the input function
+        branch = self.net.build_branch_net()
+        # Trunk net to encode the domain of the output function
+        trunk = self.net.build_trunk_net()
+        return branch, trunk
+
+    @abstractmethod
+    def build(self):
+        pass
+
+
+class VanillaStrategy(DeepONetStrategy):
+    def build(self):
+        branch, trunk = self._build_branch_and_trunk()
+        if branch.shape[-1] != trunk.shape[-1]:
+            raise AssertionError(
+                "Output sizes of branch net and trunk net do not match."
+            )
+        y = self.net.merge(branch, trunk)
+        return y
+
+
+class IndependentStrategy(DeepONetStrategy):
+    """Directly use n independent DeepONets,
+    and each DeepONet outputs only one function.
+    """
+
+    def build(self):
+        vanilla_strategy = VanillaStrategy(self.net)
+        ys = []
+        for _ in range(self.net.num_outputs):
+            ys.append(vanilla_strategy.build())
+        return self.net.concatenate_outputs(ys)
+
+
+class SplitBothStrategy(DeepONetStrategy):
+    """Split the outputs of both the branch net and the trunk net into n groups,
+    and then the kth group outputs the kth solution.
+
+    For example, if n = 2 and both the branch and trunk nets have 100 output neurons,
+    then the dot product between the first 50 neurons of
+    the branch and trunk nets generates the first function,
+    and the remaining 50 neurons generate the second function.
+    """
+
+    def build(self):
+        branch, trunk = self._build_branch_and_trunk()
+        if branch.shape[-1] != trunk.shape[-1]:
+            raise AssertionError(
+                "Output sizes of branch net and trunk net do not match."
+            )
+        if branch.shape[-1] % self.net.num_outputs != 0:
+            raise AssertionError(
+                f"Output size of the branch net is not evenly divisible by {self.net.num_outputs}."
+            )
+        branch_groups = tf.split(
+            branch, num_or_size_splits=self.net.num_outputs, axis=1
+        )
+        trunk_groups = tf.split(trunk, num_or_size_splits=self.net.num_outputs, axis=1)
+        ys = []
+        for i in range(self.net.num_outputs):
+            y = self.net.merge(branch_groups[i], trunk_groups[i])
+            ys.append(y)
+        return self.net.concatenate_outputs(ys)
+
+
+class SplitBranchStrategy(DeepONetStrategy):
+    """Split the branch net and share the trunk net."""
+
+    def build(self):
+        branch, trunk = self._build_branch_and_trunk()
+        if branch.shape[-1] % self.net.num_outputs != 0:
+            raise AssertionError(
+                f"Output size of the branch net is not evenly divisible by {self.net.num_outputs}."
+            )
+        if branch.shape[-1] / self.net.num_outputs != trunk.shape[-1]:
+            raise AssertionError(
+                f"Output size of the trunk net does not equal to {branch.shape[-1] // self.net.num_outputs}."
+            )
+        branch_groups = tf.split(
+            branch, num_or_size_splits=self.net.num_outputs, axis=1
+        )
+        ys = []
+        for i in range(self.net.num_outputs):
+            y = self.net.merge(branch_groups[i], trunk)
+            ys.append(y)
+        return self.net.concatenate_outputs(ys)
+
+
+class SplitTrunkStrategy(DeepONetStrategy):
+    """Split the trunk net and share the branch net."""
+
+    def build(self):
+        branch, trunk = self._build_branch_and_trunk()
+        if trunk.shape[-1] % self.net.num_outputs != 0:
+            raise AssertionError(
+                f"Output size of the trunk net is not evenly divisible by {self.net.num_outputs}."
+            )
+        if trunk.shape[-1] / self.net.num_outputs != branch.shape[-1]:
+            raise AssertionError(
+                f"Output size of the branch net does not equal to {trunk.shape[-1] // self.net.num_outputs}."
+            )
+        trunk_groups = tf.split(trunk, num_or_size_splits=self.net.num_outputs, axis=1)
+        ys = []
+        for i in range(self.net.num_outputs):
+            y = self.net.merge(branch, trunk_groups[i])
+            ys.append(y)
+        return self.net.concatenate_outputs(ys)
 
 
 class DeepONet(NN):
@@ -29,6 +154,8 @@ class DeepONet(NN):
             `activation["branch"]`.
         trainable_branch: Boolean.
         trainable_trunk: Boolean or a list of booleans.
+        num_outputs (integer): number of outputs.
+        strategy (str): "vanilla", "independent", "split_both", "split_branch" or "split_trunk".
     """
 
     def __init__(
@@ -42,6 +169,8 @@ def __init__(
         stacked=False,
         trainable_branch=True,
         trainable_trunk=True,
+        num_outputs=1,
+        strategy="independent",
     ):
         super().__init__()
         if isinstance(trainable_trunk, (list, tuple)):
@@ -69,6 +198,22 @@ def __init__(
         self._inputs = None
         self._X_func_default = None
 
+        self.num_outputs = num_outputs
+        if self.num_outputs == 1:
+            if strategy != "vanilla":
+                strategy = "vanilla"
+                print('Strategy is forcibly changed to "vanilla".')
+        elif strategy == "vanilla":
+            strategy = "independent"
+            print('Strategy is forcibly changed to "independent".')
+        self.strategy = {
+            "independent": IndependentStrategy,
+            "split_both": SplitBothStrategy,
+            "split_branch": SplitBranchStrategy,
+            "split_trunk": SplitTrunkStrategy,
+            "vanilla": VanillaStrategy,
+        }.get(strategy, IndependentStrategy)(self)
+
     @property
     def inputs(self):
         return self._inputs
@@ -101,7 +246,14 @@ def build(self):
         self.X_loc = tf.placeholder(config.real(tf), [None, self.layer_size_loc[0]])
         self._inputs = [self.X_func, self.X_loc]
 
-        # Branch net to encode the input function
+        self.y = self.strategy.build()
+        if self._output_transform is not None:
+            self.y = self._output_transform(self._inputs, self.y)
+
+        self.target = tf.placeholder(config.real(tf), [None, self.num_outputs])
+        self.built = True
+
+    def build_branch_net(self):
         y_func = self.X_func
         if callable(self.layer_size_func[1]):
             # User-defined network
@@ -141,8 +293,9 @@ def build(self):
                 regularizer=self.regularizer,
                 trainable=self.trainable_branch,
             )
+        return y_func
 
-        # Trunk net to encode the domain of the output function
+    def build_trunk_net(self):
         y_loc = self.X_loc
         if self._input_transform is not None:
             y_loc = self._input_transform(y_loc)
@@ -156,24 +309,20 @@ def build(self):
                 if isinstance(self.trainable_trunk, (list, tuple))
                 else self.trainable_trunk,
             )
+        return y_loc
 
+    def merge(self, branch, trunk):
         # Dot product
-        if y_func.shape[-1] != y_loc.shape[-1]:
-            raise AssertionError(
-                "Output sizes of branch net and trunk net do not match."
-            )
-        self.y = tf.einsum("bi,bi->b", y_func, y_loc)
-        self.y = tf.expand_dims(self.y, axis=1)
-        # Add bias
+        y = tf.einsum("bi,bi->b", branch, trunk)
+        y = tf.expand_dims(y, axis=1)
         if self.use_bias:
             b = tf.Variable(tf.zeros(1, dtype=config.real(tf)))
-            self.y += b
-
-        if self._output_transform is not None:
-            self.y = self._output_transform(self._inputs, self.y)
+            y += b
+        return y
 
-        self.target = tf.placeholder(config.real(tf), [None, 1])
-        self.built = True
+    @staticmethod
+    def concatenate_outputs(ys):
+        return tf.concat(ys, axis=1)
 
     def _dense(
         self,
@@ -259,6 +408,8 @@ class DeepONetCartesianProd(NN):
             both trunk and branch nets. If `activation` is a ``dict``, then the trunk
             net uses the activation `activation["trunk"]`, and the branch net uses
             `activation["branch"]`.
+        num_outputs (integer): number of outputs.
+        strategy (str): "vanilla", "independent", "split_both", "split_branch" or "split_trunk".
     """
 
     def __init__(
@@ -268,6 +419,8 @@ def __init__(
         activation,
         kernel_initializer,
         regularization=None,
+        num_outputs=1,
+        strategy="independent",
     ):
         super().__init__()
         self.layer_size_func = layer_size_branch
@@ -279,9 +432,24 @@ def __init__(
             self.activation_branch = self.activation_trunk = activations.get(activation)
         self.kernel_initializer = initializers.get(kernel_initializer)
         self.regularizer = regularizers.get(regularization)
-
         self._inputs = None
 
+        self.num_outputs = num_outputs
+        if self.num_outputs == 1:
+            if strategy != "vanilla":
+                strategy = "vanilla"
+                print('Strategy is forcibly changed to "vanilla".')
+        elif strategy == "vanilla":
+            strategy = "independent"
+            print('Strategy is forcibly changed to "independent".')
+        self.strategy = {
+            "independent": IndependentStrategy,
+            "split_both": SplitBothStrategy,
+            "split_branch": SplitBranchStrategy,
+            "split_trunk": SplitTrunkStrategy,
+            "vanilla": VanillaStrategy,
+        }.get(strategy, IndependentStrategy)(self)
+
     @property
     def inputs(self):
         return self._inputs
@@ -301,7 +469,14 @@ def build(self):
         self.X_loc = tf.placeholder(config.real(tf), [None, self.layer_size_loc[0]])
         self._inputs = [self.X_func, self.X_loc]
 
-        # Branch net to encode the input function
+        self.y = self.strategy.build()
+        if self._output_transform is not None:
+            self.y = self._output_transform(self._inputs, self.y)
+
+        self.target = tf.placeholder(config.real(tf), [None, None])
+        self.built = True
+
+    def build_branch_net(self):
         y_func = self.X_func
         if callable(self.layer_size_func[1]):
             # User-defined network
@@ -322,7 +497,9 @@ def build(self):
                 kernel_initializer=self.kernel_initializer,
                 kernel_regularizer=self.regularizer,
             )
+        return y_func
 
+    def build_trunk_net(self):
         # Trunk net to encode the domain of the output function
         y_loc = self.X_loc
         if self._input_transform is not None:
@@ -335,19 +512,15 @@ def build(self):
                 kernel_initializer=self.kernel_initializer,
                 kernel_regularizer=self.regularizer,
             )
+        return y_loc
 
-        # Dot product
-        if y_func.shape[-1] != y_loc.shape[-1]:
-            raise AssertionError(
-                "Output sizes of branch net and trunk net do not match."
-            )
-        self.y = tf.einsum("bi,ni->bn", y_func, y_loc)
+    def merge(self, branch, trunk):
+        y = tf.einsum("bi,ni->bn", branch, trunk)
         # Add bias
         b = tf.Variable(tf.zeros(1, dtype=config.real(tf)))
-        self.y += b
+        y += b
+        return y
 
-        if self._output_transform is not None:
-            self.y = self._output_transform(self._inputs, self.y)
-
-        self.target = tf.placeholder(config.real(tf), [None, None])
-        self.built = True
+    @staticmethod
+    def concatenate_outputs(ys):
+        return tf.stack(ys, axis=2)