seal complex type between fft and ifft

Author: foldfelis

Project.toml

@@ -10,6 +10,7 @@ Fetch = "bb354801-46f6-40b6-9c3d-d42d7a74c775"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 MAT = "23992714-dd62-5051-b70f-ba57cb901cac"
 Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 julia = "1.6"

src/fourier.jl

@@ -1,14 +1,13 @@
 using Flux
 using FFTW
 using Tullio
+using Zygote
 
 export
     SpectralConv1d,
     FourierOperator,
     FNO
 
-c_glorot_uniform(dims...) = Flux.glorot_uniform(dims...) + Flux.glorot_uniform(dims...) * im
-
 struct SpectralConv1d{T, S}
     weight::T
     in_channel::S
@@ -17,8 +16,12 @@ struct SpectralConv1d{T, S}
     σ
 end
 
+function c_glorot_uniform(dims...)
+    return Flux.glorot_uniform(dims...) + Flux.glorot_uniform(dims...) * im
+end
+
 function SpectralConv1d(
-    ch::Pair{<:Integer,<:Integer},
+    ch::Pair{<:Integer, <:Integer},
     modes::Integer,
     σ=identity;
     init=c_glorot_uniform,
@@ -28,7 +31,10 @@ function SpectralConv1d(
     scale = one(T) / (in_chs * out_chs)
     weights = scale * init(out_chs, in_chs, modes)
 
-    return SpectralConv1d(weights, in_chs, out_chs, modes, σ)
+    return Chain(
+        x -> Zygote.hook(real, x),
+        SpectralConv1d(weights, in_chs, out_chs, modes, σ)
+    )
 end
 
 Flux.@functor SpectralConv1d
@@ -46,17 +52,17 @@ function (m::SpectralConv1d)(𝐱::AbstractArray)
 
     𝐱_out = ifft(𝐱_padded, 2)
 
-    return m.σ.(𝐱_out)
+    return m.σ.(real(𝐱_out))
 end
 
 function FourierOperator(
-    ch::Pair{<:Integer,<:Integer},
+    ch::Pair{<:Integer, <:Integer},
     modes::Integer,
     σ=identity
 )
     return Chain(
         Parallel(+,
-            Dense(ch.first, ch.second, init=c_glorot_uniform),
+            Dense(ch.first, ch.second),
             SpectralConv1d(ch, modes)
         ),
         x -> σ.(x)
@@ -66,16 +72,16 @@ end
 function FNO()
     modes = 16
     ch = 64 => 64
-    σ = x -> @. log(1 + exp(x))
+    σ = relu
 
     return Chain(
-        Dense(2, 64, init=c_glorot_uniform),
+        Dense(2, 64),
         FourierOperator(ch, modes, σ),
         FourierOperator(ch, modes, σ),
         FourierOperator(ch, modes, σ),
         FourierOperator(ch, modes),
-        Dense(64, 128, σ, init=c_glorot_uniform),
-        Dense(128, 1, init=c_glorot_uniform),
+        Dense(64, 128, σ),
+        Dense(128, 1),
         flatten
     )
 end

test/fourier.jl

@@ -5,15 +5,15 @@ using Flux
     ch = 64 => 64
 
     m = Chain(
-        Dense(2, 64, init=NeuralOperators.c_glorot_uniform),
+        Dense(2, 64),
         SpectralConv1d(ch, modes)
     )
 
     𝐱, _ = get_data()
     @test size(m(𝐱)) == (64, 1024, 1000)
 
     T = Float32
-    loss(x, y) = Flux.mse(real.(m(x)), y)
+    loss(x, y) = Flux.mse(m(x), y)
     data = [(T.(𝐱[:, :, 1:5]), rand(T, 64, 1024, 5))]
     Flux.train!(loss, params(m), data, Flux.ADAM())
 end
@@ -23,16 +23,15 @@ end
     ch = 64 => 64
 
     m = Chain(
-        Dense(2, 64, init=NeuralOperators.c_glorot_uniform),
+        Dense(2, 64),
         FourierOperator(ch, modes)
     )
 
     𝐱, _ = get_data()
     @test size(m(𝐱)) == (64, 1024, 1000)
 
-    T = Float32
-    loss(x, y) = Flux.mse(real.(m(x)), y)
-    data = [(T.(𝐱[:, :, 1:5]), rand(T, 64, 1024, 5))]
+    loss(x, y) = Flux.mse(m(x), y)
+    data = [(Float32.(𝐱[:, :, 1:5]), rand(Float32, 64, 1024, 5))]
     Flux.train!(loss, params(m), data, Flux.ADAM())
 end