Possible error in Fourier layer

[PaulaGarciaMolina]

Hi! I have been working with your code to try to understand FNO and I found a possible issue in the Fourier layer.

def FourierLayer(X, stage):
    
    # Name definition
    FFT_name = 'fft-layer' + str(stage)
    RFFT_name = 'ifft-layer' + str(stage)
    
    
    Res_X = X
    
    
    X = rfft2d(X, name = FFT_name)
    
    
    out_ft = tf.zeros((1, X.shape[1], X.shape[2], X.shape[3]),dtype=tf.dtypes.complex64)

    
    out_ft_first = tf.Variable(np.array(out_ft.shape))
    out_ft_first = tf.math.multiply(X[:, :, :4, :4], tf.cast(Res_X[:, :, :4, :4], tf.complex64))

    out_ft_end  = tf.Variable(np.array(out_ft.shape))
    out_ft_end = tf.math.multiply(X[:, :, -4:, :4], tf.cast(Res_X[:, :, -4:, :4], tf.complex64))


    o1 = tf.concat([out_ft_first,out_ft[:,:,:4,4:]],axis=3)
    o2 = tf.concat([out_ft[:,:,-4:,4:],out_ft_end],axis=3)
    
    out_ft = tf.concat([o1,o2],axis=2)

    
    X = irfft2d(X,name = RFFT_name)
    
    
    X = Add()([X, Res_X])
    
    X = Activation('relu')(X)

    
    return X

As you can see, the frequencies are cut and this leads to out_ft. Then, this should be the input argument of the irfft2, however, X (which is the Fourier transform of the original data) is given. This way you aren't reducing the parameters in momentum space, it is like performing the identity. I believe this is a typo but I may have misunderstood something.

Thank you!

Best,
Paula