Multi-dimensional data for operator learning?

[cayobro]

I have a very simple questions, but I have been struggling to find the answer for days!
In contrast to all the examples, my input and output data has more than one dimension. How can I handle this with the deepXDE framework?

branch input: n_s X n_t X 3
trunk input: n_t X 1
output: n_s X n_t X 12

I create my data via data.TripleCartesianProd with X=(branch_input, trunk_input), y=output.
My model is nn.DeepONetCartesianProd.

I always get dimension mismatches, no matter what I do. What is the dimension convention for this sort of multi-dimensional data?

Help would be greatly appreciated.
Carina

Example:

import deepxde as dde
from sklearn.model_selection import train_test_split
import numpy as np

dim_x = 3
dim_y = 12
n_s = 1000 # samples
n_t = 100 # time points
time_vec = np.linspace(0, 1, n_t).astype(np.float32)
input = np.random.uniform(size=(n_s, n_t, dim_x)).astype(np.float32)
output = np.random.uniform(size=(n_s, n_t, dim_y)).astype(np.float32)

print(f"Input shape: {input.shape}")
print(f"Output shape: {output.shape}")

X_train, X_test, y_train, y_test = train_test_split(input, output, test_size=0.2, random_state=42)

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

net = dde.nn.DeepONetCartesianProd(
    [dim_x, 128, 128],
    [1, 128, 128],
    "tanh",
    "Glorot uniform"
)

# 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)

[vl-dud]

Could you show the whole code?

[cayobro]

Sure, here a minimal example. This is with the raw data dimensions, no flattening or anything.

import deepxde as dde
from sklearn.model_selection import train_test_split
import numpy as np

dim_x = 3
dim_y = 12
n_s = 1000 # samples
n_t = 100 # time points
time_vec = np.linspace(0, 1, n_t).astype(np.float32)
input = np.random.uniform(size=(n_s, n_t, dim_x)).astype(np.float32)
output = np.random.uniform(size=(n_s, n_t, dim_y)).astype(np.float32)

print(f"Input shape: {input.shape}")
print(f"Output shape: {output.shape}")

X_train, X_test, y_train, y_test = train_test_split(input, output, test_size=0.2, random_state=42)

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

net = dde.nn.DeepONetCartesianProd(
    [dim_x, 128, 128],
    [1, 128, 128],
    "tanh",
    "Glorot uniform"
)

# 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)

[vl-dud]

import deepxde as dde
import numpy as np

dim_1 = 100
dim_2 = 1
# samples
n_s_1 = 1000
n_s_2 = 150
input_1 = np.random.uniform(size=(n_s_1, dim_1)).astype(np.float32)
input_2 = np.random.uniform(size=(n_s_2, dim_2)).astype(np.float32)
output = np.random.uniform(size=(n_s_1, n_s_2)).astype(np.float32)

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

net = dde.nn.DeepONetCartesianProd(
    [dim_1, 128, 128],
    [dim_2, 128, 128],
    "tanh",
    "Glorot uniform"
)

# 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)

[cayobro]

Your code results in

Input 1 shape: (1000, 100)
Input 2 shape: (150, 1)
Output shape: (1000, 150)

I don't think see how this relates to my problem; can you explain it?
My input represents n_s trajectories with n_t time points each, and the dimension of the function is dim_x = 3.
My output represents the resulting n_s trajectories with n_t time points each, and the dimension of this output function is dim_y = 12 (for example).

[vl-dud]

I don't think see how this relates to my problem; can you explain it?

Yes, it wasn't related to your problem, I just showed the correct size of the inputs and outputs.

and the dimension of the function is dim_x = 3.

I'm not sure if this is possible in the current implementation.

and the dimension of this output function is dim_y = 12

Use num_outputs argument of DeepONetCartesianProd.