Opnn supports user-defined branch net

Author: lululxvi

deepxde/maps/opnn.py

@@ -14,13 +14,17 @@
 
 
 class OpNN(Map):
-    """Operator neural networks.
+    """Deep operator network.
 
     Args:
+        layer_size_branch: A list of integers as the width of a fully connected network, or `(dim, f)` where `dim` is
+            the input dimension and `f` is a network function. The width of the last layer in the branch and trunk net
+            should be equal.
+        layer_size_trunk (list): A list of integers as the width of a fully connected network.
         activation: If `activation` is a ``string``, then the same activation is used in 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"]`.
-        trainable_branch (bool)
+        trainable_branch: Boolean.
         trainable_trunk: Boolean or a list of booleans.
     """
 
@@ -37,8 +41,6 @@ def __init__(
         trainable_trunk=True,
     ):
         super(OpNN, self).__init__()
-        if layer_size_branch[-1] != layer_size_trunk[-1]:
-            raise ValueError("Output sizes of branch net and trunk net do not match.")
         if isinstance(trainable_trunk, (list, tuple)):
             if len(trainable_trunk) != len(layer_size_trunk) - 1:
                 raise ValueError("trainable_trunk does not match layer_size_trunk.")
@@ -98,8 +100,11 @@ def build(self):
 
         # Branch net to encode the input function
         y_func = self.X_func
-        if self.stacked:
-            # Stacked
+        if callable(self.layer_size_func[1]):
+            # User-defined network
+            y_func = self.layer_size_func[1](y_func)
+        elif self.stacked:
+            # Stacked fully connected network
             stack_size = self.layer_size_func[-1]
             for i in range(1, len(self.layer_size_func) - 1):
                 y_func = self.stacked_dense(
@@ -117,7 +122,7 @@ def build(self):
                 trainable=self.trainable_branch,
             )
         else:
-            # Unstacked
+            # Unstacked fully connected network
             for i in range(1, len(self.layer_size_func) - 1):
                 y_func = self.dense(
                     y_func,
@@ -147,6 +152,10 @@ def build(self):
             )
 
         # Dot product
+        if y_func.get_shape().as_list()[-1] != y_loc.get_shape().as_list()[-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