Tensorflow 1.x backend: multiple outputs extension of DeepONet

deepxde/data/pde_operator.py

@@ -72,6 +72,8 @@ def __init__(
 
     def losses(self, targets, outputs, loss_fn, inputs, model, aux=None):
         f = []
+        if hasattr(model.net, 'prepare_multiple_outputs'):
+            outputs = model.net.prepare_multiple_outputs(outputs)
         if self.pde.pde is not None:
             f = self.pde.pde(inputs[1], outputs, model.net.auxiliary_vars)
             if not isinstance(f, (list, tuple)):
@@ -237,7 +239,8 @@ def _losses(self, outputs, loss_fn, inputs, model, num_func):
         losses = []
         for i in range(num_func):
             out = outputs[i][:, None]
-
+            if hasattr(model.net, 'prepare_multiple_outputs'):
+                out = model.net.prepare_multiple_outputs(out)
             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

@@ -42,6 +42,7 @@ def __init__(
         stacked=False,
         trainable_branch=True,
         trainable_trunk=True,
+        output_count=1
     ):
         super().__init__()
         if isinstance(trainable_trunk, (list, tuple)):
@@ -50,6 +51,7 @@ def __init__(
 
         self.layer_size_func = layer_sizes_branch
         self.layer_size_loc = layer_sizes_trunk
+        self.output_count = output_count
         if isinstance(activation, dict):
             self.activation_branch = activations.get(activation["branch"])
             self.activation_trunk = activations.get(activation["trunk"])
@@ -94,13 +96,33 @@ def _feed_dict_inputs(self, inputs):
             inputs = [np.tile(self._X_func_default, (n, 1)), inputs]
         return dict(zip([self.X_func, self.X_loc], inputs))
 
+    def prepare_multiple_outputs(self, outputs):
+        if self.output_count > 1:
+            return tf.reshape(outputs, (tf.shape(outputs)[0], tf.shape(outputs)[2]))
+        return outputs
+
     @timing
     def build(self):
         print("Building DeepONet...")
         self.X_func = tf.placeholder(config.real(tf), [None, self.layer_size_func[0]])
         self.X_loc = tf.placeholder(config.real(tf), [None, self.layer_size_loc[0]])
         self._inputs = [self.X_func, self.X_loc]
 
+        if self.output_count == 1:
+            self.y = self._build_vanilla_deeponet()
+        else:
+            ys = []
+            for _ in range(self.output_count):
+                ys.append(self._build_vanilla_deeponet())
+            self.y = tf.stack(ys, axis=2)
+
+        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.output_count])
+        self.built = True
+
+    def _build_vanilla_deeponet(self):
         # Branch net to encode the input function
         y_func = self.X_func
         if callable(self.layer_size_func[1]):
@@ -162,18 +184,17 @@ def build(self):
             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)
+        y = tf.einsum(
+            "bi,bi->b",
+            y_func,
+            y_loc
+        )
+        y = tf.expand_dims(y, axis=1)
         # Add bias
         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)
-
-        self.target = tf.placeholder(config.real(tf), [None, 1])
-        self.built = True
+            y += b
+        return y
 
     def _dense(
         self,
@@ -195,7 +216,7 @@ def _dense(
         )
 
     def _stacked_dense(
-        self, inputs, units, stack_size, activation=None, use_bias=True, trainable=True
+            self, inputs, units, stack_size, activation=None, use_bias=True, trainable=True
     ):
         """Stacked densely-connected NN layer.
 
@@ -268,10 +289,12 @@ def __init__(
         activation,
         kernel_initializer,
         regularization=None,
+        output_count=1
     ):
         super().__init__()
         self.layer_size_func = layer_size_branch
         self.layer_size_loc = layer_size_trunk
+        self.output_count = output_count
         if isinstance(activation, dict):
             self.activation_branch = activations.get(activation["branch"])
             self.activation_trunk = activations.get(activation["trunk"])
@@ -294,14 +317,33 @@ def outputs(self):
     def targets(self):
         return self.target
 
+    def prepare_multiple_outputs(self, outputs):
+        if self.output_count > 1:
+            return tf.reshape(outputs, (tf.shape(outputs)[0], tf.shape(outputs)[2]))
+        return outputs
+
     @timing
     def build(self):
         print("Building DeepONetCartesianProd...")
         self.X_func = tf.placeholder(config.real(tf), [None, self.layer_size_func[0]])
         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
+        if self.output_count == 1:
+            self.y = self._build_vanilla_deeponet()
+        else:
+            ys = []
+            for _ in range(0, self.output_count):
+                ys.append(self._build_vanilla_deeponet())
+            self.y = tf.stack(ys, axis=2)
+
+        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_vanilla_deeponet(self):
         y_func = self.X_func
         if callable(self.layer_size_func[1]):
             # User-defined network
@@ -336,18 +378,16 @@ def build(self):
                 kernel_regularizer=self.regularizer,
             )
 
-        # 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)
+        y = tf.einsum(
+            "bi,ni->bn",
+            y_func,
+            y_loc
+        )
         # Add 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)
-
-        self.target = tf.placeholder(config.real(tf), [None, None])
-        self.built = True
+        y += b
+        return y

examples/operator/poisson_aligned_2d.py

@@ -0,0 +1,72 @@
+"""
+Poisson-like 2D problem
+Backend supported: tensorflow.compat.v1
+"""
+import numpy as np
+import deepxde as dde
+from deepxde.backend import tf
+import matplotlib.pyplot as plt
+
+
+# Two target variables: A and B
+# Equations: dA_xx = f, dB_tt = f
+def equation(x, y, f):
+    A = y[:, 0:1]
+    B = y[:, 1:2]
+    dA_xx = dde.grad.hessian(y, x, component=0, i=0, j=0)
+    dB_tt = dde.grad.hessian(y, x, component=1, i=1, j=1)
+    return [
+        dA_xx - f,
+        dB_tt - f
+    ]
+
+
+# Define space/time geometry
+geomtime = dde.geometry.GeometryXTime(
+    dde.geometry.Interval(0, 1),
+    dde.geometry.TimeDomain(0, 1)
+)
+
+# Boundary conditions for A and B
+A_bc = dde.icbc.DirichletBC(
+    geomtime,
+    lambda _: 0,
+    lambda _, on_boundary: on_boundary and np.isclose(_[0], 0),
+    component=0
+)
+B_bc = dde.icbc.DirichletBC(
+    geomtime,
+    lambda _: 0,
+    lambda _, on_boundary: on_boundary and np.isclose(_[0], 1),
+    component=1
+)
+
+space = dde.data.GRF2D()
+evaluation_points = geomtime.uniform_points(10)
+
+data = dde.data.PDEOperatorCartesianProd(
+    dde.data.TimePDE(
+        geomtime,
+        equation,
+        [A_bc, B_bc],
+        num_domain=1000,
+        num_boundary=10
+    ),
+    space,
+    evaluation_points,
+    num_function=10
+)
+
+# Define DeepONet with two outputs
+net = dde.nn.DeepONetCartesianProd(
+    [evaluation_points.shape[0], 100, 100],
+    [geomtime.dim, 100, 100],
+    activation="tanh",
+    kernel_initializer="Glorot normal",
+    output_count=2
+)
+
+# Train model
+model = dde.Model(data, net)
+model.compile("adam", lr=0.001)
+model.train(iterations=5000)