From f236e7c6ff63070274a0d338739a317187f04eea Mon Sep 17 00:00:00 2001
From: Vladimir <vldud.off@gmail.com>
Date: Tue, 7 Nov 2023 11:32:51 +0300
Subject: [PATCH 1/2] Add the example

---
 ...isson_aligned_multioutput_pideeponet_2d.py | 69 +++++++++++++++++++
 1 file changed, 69 insertions(+)
 create mode 100644 examples/operator/poisson_aligned_multioutput_pideeponet_2d.py

diff --git a/examples/operator/poisson_aligned_multioutput_pideeponet_2d.py b/examples/operator/poisson_aligned_multioutput_pideeponet_2d.py
new file mode 100644
index 000000000..34d2c7742
--- /dev/null
+++ b/examples/operator/poisson_aligned_multioutput_pideeponet_2d.py
@@ -0,0 +1,69 @@
+"""
+Poisson-like 2D problem
+Supported backend: tensorflow.compat.v1, tensorflow
+"""
+import os
+
+os.environ["DDEBACKEND"] = "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",
+    num_outputs=2,
+)
+
+# Train model
+model = dde.Model(data, net)
+model.compile("adam", lr=0.001)
+model.train(iterations=5000)

From 66001594fa6eb7b08df2151db3c039e334822756 Mon Sep 17 00:00:00 2001
From: Vladimir <vldud.off@gmail.com>
Date: Sun, 24 Dec 2023 20:09:53 +0300
Subject: [PATCH 2/2] Use split_branch strategy

---
 ...oisson_aligned_multioutput_pideeponet_2d.py | 18 +++++++-----------
 1 file changed, 7 insertions(+), 11 deletions(-)

diff --git a/examples/operator/poisson_aligned_multioutput_pideeponet_2d.py b/examples/operator/poisson_aligned_multioutput_pideeponet_2d.py
index 34d2c7742..241b67c92 100644
--- a/examples/operator/poisson_aligned_multioutput_pideeponet_2d.py
+++ b/examples/operator/poisson_aligned_multioutput_pideeponet_2d.py
@@ -2,22 +2,15 @@
 Poisson-like 2D problem
 Supported backend: tensorflow.compat.v1, tensorflow
 """
-import os
-
-os.environ["DDEBACKEND"] = "tensorflow.compat.v1"
-
 import numpy as np
-import deepxde as dde
-from deepxde.backend import tf
-import matplotlib.pyplot as plt
 
+import deepxde as dde
 
 
 # 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]
+    # A = y[:, 0:1] and 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]
@@ -54,13 +47,16 @@ def equation(x, y, f):
     num_function=10,
 )
 
-# Define DeepONet with two outputs
+# Define DeepONet with two outputs.
+# Use `split_branch` strategy. The output size of the trunk net is equal
+# to the output size of the branch net divided by the number of outputs.
 net = dde.nn.DeepONetCartesianProd(
     [evaluation_points.shape[0], 100, 100],
-    [geomtime.dim, 100, 100],
+    [geomtime.dim, 100, 50],
     activation="tanh",
     kernel_initializer="Glorot normal",
     num_outputs=2,
+    multi_output_strategy="split_branch",
 )
 
 # Train model