Implementation of DeepONet architecture for a biological system that has multiple inputs and outputs

[RiazuddinNawaz]

Dear Community,
I am trying to implement the DeepONet architecture for a biological system that has multiple inputs and multiple outputs. The data is a Time Series data where I would like to send the input process parameters to branch network and time points (where the outputs are measured/evaluated) to trunk network. There are multiple inputs and multiple outputs where I want the network to map and predict these multiple outputs based on multiple inputs. I tried with DeepXDE library, but getting errors due to a mismatch in dimensions as there are multiple outputs. As I am completely new to this, could anybody suggest how to implement this architecture for this kind of system? Appreciate the help !! Thanks in advance. KR.

As a start, I was following this simple code (available as an example in the documentation) for implementation ("""Backend supported: tensorflow.compat.v1, tensorflow, pytorch, paddle"""
import deepxde as dde
import matplotlib.pyplot as plt
import numpy as np

Load dataset

d = np.load("antiderivative_aligned_train.npz", allow_pickle=True)
X_train = (d["X"][0].astype(np.float32), d["X"][1].astype(np.float32))
y_train = d["y"].astype(np.float32)
d = np.load("antiderivative_aligned_test.npz", allow_pickle=True)
X_test = (d["X"][0].astype(np.float32), d["X"][1].astype(np.float32))
y_test = d["y"].astype(np.float32)

data = dde.data.TripleCartesianProd(
X_train=X_train, y_train=y_train, X_test=X_test, y_test=y_test
)

Choose a network

m = 100
dim_x = 1
net = dde.nn.DeepONetCartesianProd(
[m, 40, 40],
[dim_x, 40, 40],
"relu",
"Glorot normal",
)

Define a Model

model = dde.Model(data, net)

Compile and Train

model.compile("adam", lr=0.001, metrics=["mean l2 relative error"])
losshistory, train_state = model.train(iterations=10000)

Plot the loss trajectory

dde.utils.plot_loss_history(losshistory)
plt.show())