Which script to use for 3D FNO and missing normalizer about neuraloperator HOT 8 CLOSED

The input has the shape (x, y, T_in); the output has the shape (x, y, T). In fourier_3d.py we treat T_in as the channel dimension and repeat the input in T's dimension, (x, y, T_in) -> (x, y, T, T_in) = (x, y, t, channel). Then we do FFT on x,y,t.

train_a = train_a.reshape(ntrain,S,S,1,T_in).repeat([1,1,1,T,1])
test_a = test_a.reshape(ntest,S,S,1,T_in).repeat([1,1,1,T,1])

We do FFT in T's dimension instead of T_in's dimension, because in many cases T_in=1 (we only have the initial condition), so it won't make sense to do FFT in T_in's dimension.

The prepared input now has the shape (x, y, T, T_in) = (x, y, t, channel), and it's constant in t-dimension, we add the grid to break the symmetry. Otherwise, the output will still be constant in t. On a higher level, it also helps to capture grid-related information (considering self.fc0).

Besides, we want to avoid using normalization. In many cases, it can be quite helpful, but if one wants to extrapolate to unseen time periods, normalization may not be applicable.

from neuraloperator.

Thanks for the kind words : )
In our setting, we use 3D FNO for the 2D Navier-Stokes problem and treated the time dimension as the third dimension to do convolution. In fourier_3d.py we input the first 10 time steps and output the next 40 time steps. In scripts/fourier_3d_time.py, we still do time-convolution, input the first 10 time steps, but output the next 10 time steps, and iterate four times to reach the 40 time steps.
Do you want to apply FNO on 3D Navier-Stokes problems?

from neuraloperator.

Hmm, probably self.fc0 along won't change it much. But after several layers of self.w0 ... self.w3 and F.relu, it may be something useful.

from neuraloperator.

Hi Nis, It was nice to talk to you this week. For your question

• Yes, we use fourier_3d.py to produce the results in Sec. 5.3, not scripts/fourier_3d_time.py.
• Yes, that's right.

from neuraloperator.

Thanks for the prompt reply. There seems to be a misconception on our side. We indeed try to apply the script to the 2D Navier-Stokes data and it took us some time to understand the data preparation for the 3D case. We ended with this snippet which should implement the same logic:

import torch

ntrain = 1
S = 3
T_in = 2
T = 4

train_a = torch.arange(ntrain * S * S * T_in)
train_a = train_a.reshape(ntrain, S, S, 1, T_in)
# T_in = 0    T_in = 1
# ┌──┬──┬──┐  ┌──┬──┬──┐
# | 0| 2| 4|  | 1| 3| 5|
# ├──┼──┼──┤  ├──┼──┼──┤
# | 6| 8|10|  | 7| 9|11|
# ├──┼──┼──┤  ├──┼──┼──┤
# |12|14|16|  |13|15|17|
# └──┴──┴──┘  └──┴──┴──┘
train_a = train_a.repeat([1, 1, 1, T, 1])

gridx = torch.linspace(0, 1, S, dtype=torch.float)
gridx = gridx.reshape(1, S, 1, 1, 1).repeat([1, 1, S, T, 1])
gridy = torch.linspace(0, 1, S, dtype=torch.float)
gridy = gridy.reshape(1, 1, S, 1, 1).repeat([1, S, 1, T, 1])
gridt = torch.linspace(0, 1, T + 1, dtype=torch.float)[1:]
gridt = gridt.reshape(1, 1, 1, T, 1).repeat([1, S, S, 1, 1])

train_a = torch.cat(
    (
        gridx.repeat([ntrain, 1, 1, 1, 1]),
        gridy.repeat([ntrain, 1, 1, 1, 1]),
        gridt.repeat([ntrain, 1, 1, 1, 1]),
        train_a,
    ),
    dim=-1,
)

As far as we understand the application of Fourier transformation, you apply a 3D Fourier transformation in x, y and t for each channel, i.e., 1D Fourier transformations for each y and t in x, then for each Fx and t in y and finally for each Fx and Fy in t (or in a different sequence which shouldn't effect the result).
Let's look at the transformation sequence in x:

print("\nx @ y=0, t=0")
print(train_a[0, :, 0, 0])

print("\nx @ y=0, t=1")
print(train_a[0, :, 0, 1])

print("\nx @ y=0, t=2")
print(train_a[0, :, 0, 2])

print("\nx @ y=1, t=0")
print(train_a[0, :, 1, 0])

print("\nx @ y=1, t=1")
print(train_a[0, :, 1, 1])

print("\nx @ y=1, t=2")
print(train_a[0, :, 1, 2])

>>> x @ y=0, t=0
>>> tensor([[ 0.0000,  0.0000,  0.2500,  0.0000,  1.0000],
>>>         [ 0.5000,  0.0000,  0.2500,  6.0000,  7.0000],
>>>         [ 1.0000,  0.0000,  0.2500, 12.0000, 13.0000]])
>>>
>>> x @ y=0, t=1
>>> tensor([[ 0.0000,  0.0000,  0.5000,  0.0000,  1.0000],
>>>         [ 0.5000,  0.0000,  0.5000,  6.0000,  7.0000],
>>>         [ 1.0000,  0.0000,  0.5000, 12.0000, 13.0000]])
>>>
>>> x @ y=0, t=2
>>> tensor([[ 0.0000,  0.0000,  0.7500,  0.0000,  1.0000],
>>>         [ 0.5000,  0.0000,  0.7500,  6.0000,  7.0000],
>>>         [ 1.0000,  0.0000,  0.7500, 12.0000, 13.0000]])
>>>
>>> x @ y=1, t=0
>>> tensor([[ 0.0000,  0.5000,  0.2500,  2.0000,  3.0000],
>>>         [ 0.5000,  0.5000,  0.2500,  8.0000,  9.0000],
>>>         [ 1.0000,  0.5000,  0.2500, 14.0000, 15.0000]])
>>>
>>> x @ y=1, t=1
>>> tensor([[ 0.0000,  0.5000,  0.5000,  2.0000,  3.0000],
>>>         [ 0.5000,  0.5000,  0.5000,  8.0000,  9.0000],
>>>         [ 1.0000,  0.5000,  0.5000, 14.0000, 15.0000]])
>>>
>>> x @ y=1, t=2
>>> tensor([[ 0.0000,  0.5000,  0.7500,  2.0000,  3.0000],
>>>         [ 0.5000,  0.5000,  0.7500,  8.0000,  9.0000],
>>>         [ 1.0000,  0.5000,  0.7500, 14.0000, 15.0000]])

In this example the first three values are augmented grid positions x, y, t and the last two are the values at T_in=0 and T_in=1. To be honest, we don't understand why a Fourier transformation of the grid, e.g.,

train_a[0, :, 0, 0, 0]
>>> tensor([0.0000, 0.5000, 1.0000])

could yield useful information. To be fair, the Fourier transformation is carried out after the fully connected layer self.fc0 but there is still a lot of redundant information?

Secondly, as far as we understand the code, the temporal Fourier transformation is not evaluated in T_in which we would naively expect but instead in copies of the same data, e.g.,

train_a[0, :, 0, 0, 3]
train_a[0, :, 0, 1, 3]
>>> tensor([ 0.,  6., 12.])
>>> tensor([ 0.,  6., 12.])

The example in scripts/fourier_3d_time looks more like as what we would expect: Here the Fourier transformation is evaluated in T_in where at each channel, i.e., for a fixed point in space and time, the value is a 4D vector: (x, y, t, v) where v is the vorticity. The missing piece here is the normalization though.

We might be utterly wrong here and hope that you can enlighten us:)

from neuraloperator.

Thanks for your help. Much appreciated!

I think we are almost on the same page. But, assuming that our snippet above is correct, the FFT in t for channels beyond the grid information, e.g, i=3, is based on a vector that is uniform before self.fc0:

# t @ x=0...2, y=0, channel=3
train_a[0, 0, 0, :, 3]
train_a[0, 1, 0, :, 3]
train_a[0, 2, 0, :, 3]
>>> tensor([0., 0., 0., 0.])
>>> tensor([6., 6., 6., 6.])
>>> tensor([12., 12., 12., 12.])

Do you see that self.fc0 significantly breaks this uniform distribution?

from neuraloperator.

Ok. I understand your point and your great results also speak for themselves.

Returning to my original questions this means:

• You use fourier_3d.py to produce the results in Sec. 5.3 and scripts/fourier_3d_time.py is not referred to in the paper (?)
• You use feature normalization in fourier_3d.py but not in scripts/fourier_3d_time.py which makes line 315 in the latter a copy & paste remnant (?)

Thanking you again:)

from neuraloperator.

The input has the shape (x, y, T_in); the output has the shape (x, y, T). In fourier_3d.py we treat T_in as the channel dimension and repeat the input in T's dimension, (x, y, T_in) -> (x, y, T, T_in) = (x, y, t, channel). Then we do FFT on x,y,t.

train_a = train_a.reshape(ntrain,S,S,1,T_in).repeat([1,1,1,T,1])
test_a = test_a.reshape(ntest,S,S,1,T_in).repeat([1,1,1,T,1])

We do FFT in T's dimension instead of T_in's dimension, because in many cases T_in=1 (we only have the initial condition), so it won't make sense to do FFT in T_in's dimension.

The prepared input now has the shape (x, y, T, T_in) = (x, y, t, channel), and it's constant in t-dimension, we add the grid to break the symmetry. Otherwise, the output will still be constant in t. On a higher level, it also helps to capture grid-related information (considering self.fc0).

Besides, we want to avoid using normalization. In many cases, it can be quite helpful, but if one wants to extrapolate to unseen time periods, normalization may not be applicable.

Can you explain why are you are doing this a bit more? What exactly is a channel?

from neuraloperator.