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/data/triple.py

@@ -29,6 +29,11 @@ def __init__(self, X_train, y_train, X_test, y_test):
         self.train_sampler = BatchSampler(len(self.train_y), shuffle=True)
 
     def losses(self, targets, outputs, loss_fn, inputs, model, aux=None):
+        if isinstance(loss_fn, list):
+            losses = []
+            for fn in loss_fn:
+                losses.append(fn(targets, outputs))
+            return losses
         return loss_fn(targets, outputs)
 
     def train_next_batch(self, batch_size=None):
@@ -74,14 +79,21 @@ def __init__(self, X_train, y_train, X_test, y_test):
         self.trunk_sampler = BatchSampler(len(X_train[1]), shuffle=True)
 
     def losses(self, targets, outputs, loss_fn, inputs, model, aux=None):
+        if isinstance(loss_fn, list):
+            losses = []
+            for fn in loss_fn:
+                losses.append(fn(targets, outputs))
+            return losses
         return loss_fn(targets, outputs)
 
     def train_next_batch(self, batch_size=None):
         if batch_size is None:
             return self.train_x, self.train_y
         if not isinstance(batch_size, (tuple, list)):
             indices = self.branch_sampler.get_next(batch_size)
-            return (self.train_x[0][indices], self.train_x[1]), self.train_y[indices]
+            return (self.train_x[0][indices], self.train_x[1]), self.train_y[
+                indices
+            ]
         indices_branch = self.branch_sampler.get_next(batch_size[0])
         indices_trunk = self.trunk_sampler.get_next(batch_size[1])
         return (

deepxde/nn/tensorflow/__init__.py

@@ -1,6 +1,13 @@
 """Package for tensorflow NN modules."""
 
-__all__ = ["DeepONet", "DeepONetCartesianProd", "FNN", "NN", "PFNN", "PODDeepONet"]
+__all__ = [
+    "DeepONet",
+    "DeepONetCartesianProd",
+    "FNN",
+    "NN",
+    "PFNN",
+    "PODDeepONet",
+]
 
 from .deeponet import DeepONet, DeepONetCartesianProd, PODDeepONet
 from .fnn import FNN, PFNN