Neural Operator fix and addition

* Building FNO for 1D/2D/3D data
* Fixing bug in trunk/branch net in DeepONEt
* Fixing type check bug in spectral conv
* Adding tests for FNO
* Fixing bug in Fourier Layer (conv1d/2d/3d)
This commit is contained in:
Dario Coscia
2023-09-09 22:09:34 +02:00
committed by Nicola Demo
parent 83ecdb0eab
commit 603f56d264
6 changed files with 315 additions and 33 deletions

View File

@@ -108,8 +108,10 @@ class DeepONet(torch.nn.Module):
check_consistency(trunk_net, torch.nn.Module)
# check trunk branch nets consistency
trunk_out_dim = trunk_net.layers[-1].out_features
branch_out_dim = branch_net.layers[-1].out_features
input_trunk = torch.rand(10, len(input_indeces_trunk_net))
input_branch = torch.rand(10, len(input_indeces_branch_net))
trunk_out_dim = trunk_net(input_trunk).shape[-1]
branch_out_dim = branch_net(input_branch).shape[-1]
if trunk_out_dim != branch_out_dim:
raise ValueError('Branch and trunk networks have not the same '
'output dimension.')