'i=1 is not valid'

[DKupetsky]

Hello, I'm trying to solve the 6 DOF harmonic oscillator

xdd + b*xd + cx - F(t) = 0

for x(t) (and F(t)) in 6 spatial dimensions, and b,c constants, using a deeponet mapping from F(t) -> x(t). My strategy is to write an ODE function with 6 ODEs:

'

M = np.random.random((6,6))
A = np.random.random((6,6))
B = np.random.random((6,6))
C = np.random.random((6,6))

dde.config.disable_xla_jit()

def matmul(i,v,Mat):
    x = 0
    for j in range(6):
        x += Mat[i][j]*v[i]
    return(x)

geom = dde.geometry.TimeDomain(0,1)
def pde(t,y,F):
    x0,x1,x2,x3,x4,x5 = y[:,0:1], y[:,1:2], y[:,2:3], y[:,3:4], y[:,4:5], y[:,5:6]
    F0,F1,F2,F3,F4,F5 = F[:,0:1], F[:,1:2], F[:,2:3], F[:,3:4], F[:,4:5], F[:,5:6]

    dx0dt = dde.grad.jacobian(y,t,i=0)
    ddx0dt = dde.grad.hessian(y,t,i=0)

    dx1dt = dde.grad.jacobian(y,t,i=1)
    ddx1dt = dde.grad.hessian(y,t,i=1)

    dx2dt = dde.grad.jacobian(y,t,i=2)
    ddx2dt = dde.grad.hessian(y,t,i=2)

    dx3dt = dde.grad.jacobian(y,t,i=3)
    ddx3dt = dde.grad.hessian(y,t,i=3)

    dx4dt = dde.grad.jacobian(y,t,i=4)
    ddx4dt = dde.grad.hessian(y,t,i=4)

    dx5dt = dde.grad.jacobian(y,t,i=5)
    ddx5dt = dde.grad.hessian(y,t,i=5)

    dif0 = matmul(0,ddx0dt,M+A) + matmul(0,dx0dt,B) + matmul(0,x0,C) - F0
    dif1 = matmul(1,ddx1dt,M+A) + matmul(1,dx1dt,B) + matmul(1,x1,C) - F1
    dif2 = matmul(2,ddx2dt,M+A) + matmul(2,dx2dt,B) + matmul(2,x2,C) - F2
    dif3 = matmul(3,ddx3dt,M+A) + matmul(3,dx3dt,B) + matmul(3,x3,C) - F3
    dif4 = matmul(4,ddx4dt,M+A) + matmul(4,dx4dt,B) + matmul(4,x4,C) - F4
    dif5 = matmul(5,ddx5dt,M+A) + matmul(5,dx5dt,B) + matmul(5,x5,C) - F5
    return torch.concat(dif0,dif1,dif2,dif3,dif4,dif5)

'

And then setting initial condition x_i(0) = 0, also sampling F from GRFs and setting up the deeponet:

`

ic = dde.icbc.IC(geom, lambda _: 0, lambda _, on_initial: on_initial)

pde = dde.data.PDE(geom, pde, ic, num_domain=20, num_boundary=6, num_test=40)

# Function space
func_space = dde.data.GRF(length_scale=0.2)

# DataZ
eval_pts = np.linspace(0, 1, num=36)[:, None]

data = dde.data.PDEOperatorCartesianProd(
    pde, func_space, eval_pts, 1000, function_variables=[0], num_test=100, batch_size=100
)

# Net
net = dde.nn.DeepONetCartesianProd(
    [36, 128, 128, 128],
    [1, 128, 128, 128],
    "tanh",
    "Glorot normal",
)



model = dde.Model(data, net)
model.compile("adam", lr=0.0005)
losshistory, train_state = model.train(epochs=40000)

`

However, depending on whether I run the first block of code I listed before the second (I'm using jupyter), I get the following error:

_

Compiling model...
'compile' took 0.000457 s

Warning: epochs is deprecated and will be removed in a future version. Use iterations instead.
Training model...

---

ValueError Traceback (most recent call last)
Cell In[130], line 27
25 model = dde.Model(data, net)
26 model.compile("adam", lr=0.0005)
---> 27 losshistory, train_state = model.train(epochs=40000)

File ~\anaconda3\Lib\site-packages\deepxde\utils\internal.py:22, in timing..wrapper(*args, **kwargs)
19 @wraps(f)
20 def wrapper(*args, **kwargs):
21 ts = timeit.default_timer()
---> 22 result = f(*args, **kwargs)
23 te = timeit.default_timer()
24 if config.rank == 0:

File ~\anaconda3\Lib\site-packages\deepxde\model.py:622, in Model.train(self, iterations, batch_size, display_every, disregard_previous_best, callbacks, model_restore_path, model_save_path, epochs)
620 self.train_state.set_data_train(*self.data.train_next_batch(self.batch_size))
621 self.train_state.set_data_test(*self.data.test())
--> 622 self._test()
623 self.callbacks.on_train_begin()
624 if optimizers.is_external_optimizer(self.opt_name):

File ~\anaconda3\Lib\site-packages\deepxde\model.py:811, in Model._test(self)
806 def _test(self):
807 # TODO Now only print the training loss in rank 0. The correct way is to print the average training loss of all ranks.
808 (
809 self.train_state.y_pred_train,
810 self.train_state.loss_train,
--> 811 ) = self._outputs_losses(
812 True,
813 self.train_state.X_train,
814 self.train_state.y_train,
815 self.train_state.train_aux_vars,
816 )
817 self.train_state.y_pred_test, self.train_state.loss_test = self._outputs_losses(
818 False,
819 self.train_state.X_test,
820 self.train_state.y_test,
821 self.train_state.test_aux_vars,
822 )
824 if isinstance(self.train_state.y_test, (list, tuple)):

File ~\anaconda3\Lib\site-packages\deepxde\model.py:532, in Model.outputs_losses(self, training, inputs, targets, auxiliary_vars)
530 elif backend_name == "pytorch":
531 self.net.requires_grad(requires_grad=False)
--> 532 outs = outputs_losses(inputs, targets, auxiliary_vars)
533 self.net.requires_grad_()
534 elif backend_name == "jax":
535 # TODO: auxiliary_vars

File ~\anaconda3\Lib\site-packages\deepxde\model.py:309, in Model._compile_pytorch..outputs_losses_train(inputs, targets, auxiliary_vars)
308 def outputs_losses_train(inputs, targets, auxiliary_vars):
--> 309 return outputs_losses(
310 True, inputs, targets, auxiliary_vars, self.data.losses_train
311 )

File ~\anaconda3\Lib\site-packages\deepxde\model.py:297, in Model.compile_pytorch..outputs_losses(training, inputs, targets, auxiliary_vars, losses_fn)
295 if targets is not None:
296 targets = torch.as_tensor(targets)
--> 297 losses = losses_fn(targets, outputs, loss_fn, inputs, self)
298 if not isinstance(losses, list):
299 losses = [losses]

File ~\anaconda3\Lib\site-packages\deepxde\data\pde_operator.py:270, in PDEOperatorCartesianProd.losses_train(self, targets, outputs, loss_fn, inputs, model, aux)
268 def losses_train(self, targets, outputs, loss_fn, inputs, model, aux=None):
269 num_func = self.num_func if self.batch_size is None else self.batch_size
--> 270 return self._losses(outputs, loss_fn, inputs, model, num_func)

File ~\anaconda3\Lib\site-packages\deepxde\data\pde_operator.py:243, in PDEOperatorCartesianProd._losses(self, outputs, loss_fn, inputs, model, num_func)
241 f = []
242 if self.pde.pde is not None:
--> 243 f = self.pde.pde(inputs[1], out, model.net.auxiliary_vars[i][:, None])
244 if not isinstance(f, (list, tuple)):
245 f = [f]

Cell In[129], line 22, in pde(t, y, F)
19 dx0dt = dde.grad.jacobian(y,t,i=0)
20 ddx0dt = dde.grad.hessian(y,t,i=0)
---> 22 dx1dt = dde.grad.jacobian(y,t,i=1)
23 ddx1dt = dde.grad.hessian(y,t,i=1)
25 dx2dt = dde.grad.jacobian(y,t,i=2)

File ~\anaconda3\Lib\site-packages\deepxde\gradients.py:181, in jacobian(ys, xs, i, j)
160 def jacobian(ys, xs, i=0, j=None):
161 """Compute Jacobian matrix J: J[i][j] = dy_i / dx_j, where i = 0, ..., dim_y - 1 and
162 j = 0, ..., dim_x - 1.
163
(...)
179 y_i, i.e., J[i].
180 """
--> 181 return jacobian._Jacobians(ys, xs, i=i, j=j)

File ~\anaconda3\Lib\site-packages\deepxde\gradients.py:153, in Jacobians.call(self, ys, xs, i, j)
151 if key not in self.Js:
152 self.Js[key] = Jacobian(ys, xs)
--> 153 return self.Js[key](i, j)

File ~\anaconda3\Lib\site-packages\deepxde\gradients.py:39, in Jacobian.call(self, i, j)
35 """Returns J[i][j]. If j is None, returns the gradient of y_i, i.e.,
36 J[i].
37 """
38 if not 0 <= i < self.dim_y:
---> 39 raise ValueError("i={} is not valid.".format(i))
40 if j is not None and not 0 <= j < self.dim_x:
41 raise ValueError("j={} is not valid.".format(j))

ValueError: i=1 is not valid.

If I run the first before the second, and if I run the second again after this error I get:

`Compiling model...
'compile' took 0.000302 s

Warning: epochs is deprecated and will be removed in a future version. Use iterations instead.
Training model...

---

TypeError Traceback (most recent call last)
Cell In[131], line 27
25 model = dde.Model(data, net)
26 model.compile("adam", lr=0.0005)
---> 27 losshistory, train_state = model.train(epochs=40000)

File ~\anaconda3\Lib\site-packages\deepxde\utils\internal.py:22, in timing..wrapper(*args, **kwargs)
19 @wraps(f)
20 def wrapper(*args, **kwargs):
21 ts = timeit.default_timer()
---> 22 result = f(*args, **kwargs)
23 te = timeit.default_timer()
24 if config.rank == 0:

File ~\anaconda3\Lib\site-packages\deepxde\model.py:622, in Model.train(self, iterations, batch_size, display_every, disregard_previous_best, callbacks, model_restore_path, model_save_path, epochs)
620 self.train_state.set_data_train(*self.data.train_next_batch(self.batch_size))
621 self.train_state.set_data_test(*self.data.test())
--> 622 self._test()
623 self.callbacks.on_train_begin()
624 if optimizers.is_external_optimizer(self.opt_name):

File ~\anaconda3\Lib\site-packages\deepxde\model.py:811, in Model._test(self)
806 def _test(self):
807 # TODO Now only print the training loss in rank 0. The correct way is to print the average training loss of all ranks.
808 (
809 self.train_state.y_pred_train,
810 self.train_state.loss_train,
--> 811 ) = self._outputs_losses(
812 True,
813 self.train_state.X_train,
814 self.train_state.y_train,
815 self.train_state.train_aux_vars,
816 )
817 self.train_state.y_pred_test, self.train_state.loss_test = self._outputs_losses(
818 False,
819 self.train_state.X_test,
820 self.train_state.y_test,
821 self.train_state.test_aux_vars,
822 )
824 if isinstance(self.train_state.y_test, (list, tuple)):

File ~\anaconda3\Lib\site-packages\deepxde\model.py:532, in Model.outputs_losses(self, training, inputs, targets, auxiliary_vars)
530 elif backend_name == "pytorch":
531 self.net.requires_grad(requires_grad=False)
--> 532 outs = outputs_losses(inputs, targets, auxiliary_vars)
533 self.net.requires_grad_()
534 elif backend_name == "jax":
535 # TODO: auxiliary_vars

File ~\anaconda3\Lib\site-packages\deepxde\model.py:309, in Model._compile_pytorch..outputs_losses_train(inputs, targets, auxiliary_vars)
308 def outputs_losses_train(inputs, targets, auxiliary_vars):
--> 309 return outputs_losses(
310 True, inputs, targets, auxiliary_vars, self.data.losses_train
311 )

File ~\anaconda3\Lib\site-packages\deepxde\model.py:297, in Model.compile_pytorch..outputs_losses(training, inputs, targets, auxiliary_vars, losses_fn)
295 if targets is not None:
296 targets = torch.as_tensor(targets)
--> 297 losses = losses_fn(targets, outputs, loss_fn, inputs, self)
298 if not isinstance(losses, list):
299 losses = [losses]

File ~\anaconda3\Lib\site-packages\deepxde\data\pde_operator.py:270, in PDEOperatorCartesianProd.losses_train(self, targets, outputs, loss_fn, inputs, model, aux)
268 def losses_train(self, targets, outputs, loss_fn, inputs, model, aux=None):
269 num_func = self.num_func if self.batch_size is None else self.batch_size
--> 270 return self._losses(outputs, loss_fn, inputs, model, num_func)

File ~\anaconda3\Lib\site-packages\deepxde\data\pde_operator.py:243, in PDEOperatorCartesianProd._losses(self, outputs, loss_fn, inputs, model, num_func)
241 f = []
242 if self.pde.pde is not None:
--> 243 f = self.pde.pde(inputs[1], out, model.net.auxiliary_vars[i][:, None])
244 if not isinstance(f, (list, tuple)):
245 f = [f]

TypeError: 'PDE' object is not callable`

Am I misunderstanding how to train the deeponet? I figured I didn't need any spatial geometry since I was able to get away without one in a 1D case, and all derivatives are with respect to time. Do I need 6 inputs in the trunk net? I found two people with a similar issue in this forum but neither of them was followed up on.

Thanks!