From 23cb507ff24d572ad213f2bfe9b152c101e9f294 Mon Sep 17 00:00:00 2001 From: JingYu Ning Date: Sun, 15 Aug 2021 00:38:33 +0800 Subject: [PATCH 01/14] extend SpectralConv to n-dim --- src/fourier.jl | 44 ++++++++++++++++++++++++-------------------- src/model.jl | 2 +- test/fourier.jl | 8 ++++---- 3 files changed, 29 insertions(+), 25 deletions(-) diff --git a/src/fourier.jl b/src/fourier.jl index fddda1c2..90e40b1f 100644 --- a/src/fourier.jl +++ b/src/fourier.jl @@ -1,53 +1,57 @@ export - SpectralConv1d, + SpectralConv, FourierOperator -struct SpectralConv1d{T, S} +struct SpectralConv{T, S, N} weight::T in_channel::S out_channel::S - modes::S + modes::NTuple{N, S} + ndim::S Οƒ end c_glorot_uniform(dims...) = Flux.glorot_uniform(dims...) + Flux.glorot_uniform(dims...)*im -function SpectralConv1d( - ch::Pair{<:Integer, <:Integer}, - modes::Integer, +function SpectralConv( + ch::Pair{S, S}, + modes::NTuple{N, S}, Οƒ=identity; init=c_glorot_uniform, T::DataType=ComplexF32 -) +) where {S<:Integer, N} in_chs, out_chs = ch scale = one(T) / (in_chs * out_chs) - weights = scale * init(out_chs, in_chs, modes) + weights = scale * init(out_chs, in_chs, modes...) - return SpectralConv1d(weights, in_chs, out_chs, modes, Οƒ) + return SpectralConv(weights, in_chs, out_chs, modes, N, Οƒ) end -Flux.@functor SpectralConv1d +Flux.@functor SpectralConv -spectral_conv(𝐱₁, 𝐱₂) = @tullio 𝐲[m, o, b] := 𝐱₁[m, i, b] * 𝐱₂[o, i, m] +spectral_conv(𝐱₁, 𝐱₂) = @tullio 𝐲[m, o, b] := 𝐱₁[m, i, b] * 𝐱₂[o, i, m] # TODO: extend `m` to n-dim -function (m::SpectralConv1d)(𝐱::AbstractArray) - 𝐱ᡀ = permutedims(Zygote.hook(real, 𝐱), (2, 1, 3)) # [x, in_chs, batch] <- [in_chs, x, batch] - 𝐱_fft = fft(𝐱ᡀ, 1) # [x, in_chs, batch] +function (m::SpectralConv)(𝐱::AbstractArray) + 𝐱ᡀ = permutedims(Zygote.hook(real, 𝐱), (m.ndim+1, 1:m.ndim..., m.ndim+2)) # [x, in_chs, batch] <- [in_chs, x, batch] + 𝐱_fft = fft(𝐱ᡀ, 1:m.ndim) # [x, in_chs, batch] # [modes, out_chs, batch] <- [modes, in_chs, batch] * [out_chs, in_chs, modes] - 𝐱_weighted = spectral_conv(view(𝐱_fft, 1:m.modes, :, :), m.weight) + ranges = [1:dim_modes for dim_modes in m.modes] + 𝐱_weighted = spectral_conv(view(𝐱_fft, ranges..., :, :), m.weight) + # [x, out_chs, batch] <- [modes, out_chs, batch] - 𝐱_padded = cat(𝐱_weighted, zeros(ComplexF32, size(𝐱_fft, 1)-m.modes, Base.tail(size(𝐱_weighted))...), dims=1) + pad = zeros(ComplexF32, (collect(size(𝐱_fft)[1:m.ndim])-collect(m.modes))..., size(𝐱_weighted)[end-1:end]...) + 𝐱_padded = cat(𝐱_weighted, pad, dims=1:m.ndim) - 𝐱_out = ifft(𝐱_padded, 1) # [x, out_chs, batch] - 𝐱_outα΅€ = permutedims(real(𝐱_out), (2, 1, 3)) # [out_chs, x, batch] <- [x, out_chs, batch] + 𝐱_out = ifft(𝐱_padded, 1:m.ndim) # [x, out_chs, batch] + 𝐱_outα΅€ = permutedims(real(𝐱_out), (2:m.ndim+1..., 1, m.ndim+2)) # [out_chs, x, batch] <- [x, out_chs, batch] return m.Οƒ.(𝐱_outα΅€) end -function FourierOperator(ch::Pair{<:Integer, <:Integer}, modes::Integer, Οƒ=identity) +function FourierOperator(ch::Pair{S, S}, modes::NTuple{N, S}, Οƒ=identity) where {S<:Integer, N} return Chain( - Parallel(+, Dense(ch.first, ch.second), SpectralConv1d(ch, modes)), + Parallel(+, Dense(ch.first, ch.second), SpectralConv(ch, modes)), x -> Οƒ.(x) ) end diff --git a/src/model.jl b/src/model.jl index f68f7b4a..db4b9f9c 100644 --- a/src/model.jl +++ b/src/model.jl @@ -2,7 +2,7 @@ export FourierNeuralOperator function FourierNeuralOperator() - modes = 16 + modes = (16, ) ch = 64 => 64 Οƒ = relu diff --git a/test/fourier.jl b/test/fourier.jl index 00fec3df..c2296011 100644 --- a/test/fourier.jl +++ b/test/fourier.jl @@ -1,10 +1,10 @@ -@testset "SpectralConv1d" begin - modes = 16 +@testset "SpectralConv" begin + modes = (16, ) ch = 64 => 64 m = Chain( Dense(2, 64), - SpectralConv1d(ch, modes) + SpectralConv(ch, modes) ) 𝐱, _ = get_burgers_data(n=1000) @@ -17,7 +17,7 @@ end @testset "FourierOperator" begin - modes = 16 + modes = (16, ) ch = 64 => 64 m = Chain( From 7b24c31f1bd037a3c6b07d16b867567a4610d4fc Mon Sep 17 00:00:00 2001 From: JingYu Ning Date: Sun, 15 Aug 2021 12:08:11 +0800 Subject: [PATCH 02/14] refactor --- src/fourier.jl | 16 +++++++++------- test/fourier.jl | 1 + 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/src/fourier.jl b/src/fourier.jl index 90e40b1f..eb9a2b31 100644 --- a/src/fourier.jl +++ b/src/fourier.jl @@ -2,7 +2,7 @@ export SpectralConv, FourierOperator -struct SpectralConv{T, S, N} +struct SpectralConv{N, T, S} weight::T in_channel::S out_channel::S @@ -29,22 +29,24 @@ end Flux.@functor SpectralConv +Base.ndims(::SpectralConv{N}) where {N} = N + spectral_conv(𝐱₁, 𝐱₂) = @tullio 𝐲[m, o, b] := 𝐱₁[m, i, b] * 𝐱₂[o, i, m] # TODO: extend `m` to n-dim function (m::SpectralConv)(𝐱::AbstractArray) - 𝐱ᡀ = permutedims(Zygote.hook(real, 𝐱), (m.ndim+1, 1:m.ndim..., m.ndim+2)) # [x, in_chs, batch] <- [in_chs, x, batch] - 𝐱_fft = fft(𝐱ᡀ, 1:m.ndim) # [x, in_chs, batch] + 𝐱ᡀ = permutedims(Zygote.hook(real, 𝐱), (ndims(m)+1, 1:ndims(m)..., ndims(m)+2)) # [x, in_chs, batch] <- [in_chs, x, batch] + 𝐱_fft = fft(𝐱ᡀ, 1:ndims(m)) # [x, in_chs, batch] # [modes, out_chs, batch] <- [modes, in_chs, batch] * [out_chs, in_chs, modes] ranges = [1:dim_modes for dim_modes in m.modes] 𝐱_weighted = spectral_conv(view(𝐱_fft, ranges..., :, :), m.weight) # [x, out_chs, batch] <- [modes, out_chs, batch] - pad = zeros(ComplexF32, (collect(size(𝐱_fft)[1:m.ndim])-collect(m.modes))..., size(𝐱_weighted)[end-1:end]...) - 𝐱_padded = cat(𝐱_weighted, pad, dims=1:m.ndim) + pad = zeros(ComplexF32, (collect(size(𝐱_fft)[1:ndims(m)])-collect(m.modes))..., size(𝐱_weighted)[end-1:end]...) + 𝐱_padded = cat(𝐱_weighted, pad, dims=1:ndims(m)) - 𝐱_out = ifft(𝐱_padded, 1:m.ndim) # [x, out_chs, batch] - 𝐱_outα΅€ = permutedims(real(𝐱_out), (2:m.ndim+1..., 1, m.ndim+2)) # [out_chs, x, batch] <- [x, out_chs, batch] + 𝐱_out = ifft(𝐱_padded, 1:ndims(m)) # [x, out_chs, batch] + 𝐱_outα΅€ = permutedims(real(𝐱_out), (2:ndims(m)+1..., 1, ndims(m)+2)) # [out_chs, x, batch] <- [x, out_chs, batch] return m.Οƒ.(𝐱_outα΅€) end diff --git a/test/fourier.jl b/test/fourier.jl index c2296011..5d31ac56 100644 --- a/test/fourier.jl +++ b/test/fourier.jl @@ -6,6 +6,7 @@ Dense(2, 64), SpectralConv(ch, modes) ) + @test ndims(SpectralConv(ch, modes)) == 1 𝐱, _ = get_burgers_data(n=1000) @test size(m(𝐱)) == (64, 1024, 1000) From 35493440152883909b6d32c1264382ef6e4dcb1d Mon Sep 17 00:00:00 2001 From: JingYu Ning Date: Sun, 15 Aug 2021 12:56:48 +0800 Subject: [PATCH 03/14] skip NavierStokes --- src/data.jl | 32 ++++++++++++++++++++++++++++---- test/data.jl | 4 ++++ 2 files changed, 32 insertions(+), 4 deletions(-) diff --git a/src/data.jl b/src/data.jl index 5780c463..32eb28c5 100644 --- a/src/data.jl +++ b/src/data.jl @@ -1,9 +1,10 @@ export - get_burgers_data + get_burgers_data, + get_navier_stokes -function register_datasets() +function register_burgers() register(DataDep( - "BurgersR10", + "Burgers", """ Burgers' equation dataset from [fourier_neural_operator](https://github.com/zongyi-li/fourier_neural_operator) @@ -15,8 +16,27 @@ function register_datasets() )) end +# function register_navier_stokes() +# register(DataDep( +# "NavierStokes", +# """ +# Navier–Stokes equations dataset from +# [fourier_neural_operator](https://github.com/zongyi-li/fourier_neural_operator) +# """, +# "https://drive.google.com/file/d/1r3idxpsHa21ijhlu3QQ1hVuXcqnBTO7d/view?usp=sharing", +# "1a3b2893489dd1493923362bc74cd571e0b4f6ee290985eda060f4140df602d0", +# fetch_method=gdownload, +# post_fetch_method=unpack +# )) +# end + +function register_datasets() + register_burgers() + # register_navier_stokes() +end + function get_burgers_data(; n=2048, Ξ”samples=2^3, grid_size=div(2^13, Ξ”samples), T=Float32) - file = matopen(joinpath(datadep"BurgersR10", "burgers_data_R10.mat")) + file = matopen(joinpath(datadep"Burgers", "burgers_data_R10.mat")) x_data = T.(collect(read(file, "a")[1:n, 1:Ξ”samples:end]')) y_data = T.(collect(read(file, "u")[1:n, 1:Ξ”samples:end]')) close(file) @@ -27,3 +47,7 @@ function get_burgers_data(; n=2048, Ξ”samples=2^3, grid_size=div(2^13, Ξ”samples return x_loc_data, y_data end + +# function get_navier_stokes() +# file = matopen(joinpath(datadep"NavierStokes", "ns_V1e-3_N5000_T50.mat")) +# end diff --git a/test/data.jl b/test/data.jl index 41476f23..ba9033ae 100644 --- a/test/data.jl +++ b/test/data.jl @@ -4,3 +4,7 @@ @test size(xs) == (2, 1024, 1000) @test size(ys) == (1024, 1000) end + +@testset "get navier stokes data" begin + # get_navier_stokes() +end From c2a9d22ac6b49dd58ab52cafa0f9a70151ce5c73 Mon Sep 17 00:00:00 2001 From: JingYu Ning Date: Sun, 15 Aug 2021 16:59:06 +0800 Subject: [PATCH 04/14] darcy flow data --- Project.toml | 2 ++ src/NeuralOperators.jl | 1 + src/data.jl | 70 +++++++++++++++++++++++++++++++----------- test/data.jl | 18 +++++++++-- 4 files changed, 71 insertions(+), 20 deletions(-) diff --git a/Project.toml b/Project.toml index a200eee4..e29e3baa 100644 --- a/Project.toml +++ b/Project.toml @@ -12,6 +12,7 @@ Fetch = "bb354801-46f6-40b6-9c3d-d42d7a74c775" Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" MAT = "23992714-dd62-5051-b70f-ba57cb901cac" +StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" @@ -23,6 +24,7 @@ FFTW = "1.4" Flux = "0.12" KernelAbstractions = "0.7" MAT = "0.10" +StatsBase = "0.33" Tullio = "0.3" Zygote = "0.6" julia = "1.6" diff --git a/src/NeuralOperators.jl b/src/NeuralOperators.jl index bffee18f..7f2ab44f 100644 --- a/src/NeuralOperators.jl +++ b/src/NeuralOperators.jl @@ -2,6 +2,7 @@ module NeuralOperators using DataDeps using Fetch using MAT + using StatsBase using Flux using FFTW diff --git a/src/data.jl b/src/data.jl index 32eb28c5..070442ed 100644 --- a/src/data.jl +++ b/src/data.jl @@ -1,6 +1,28 @@ export + UnitGaussianNormalizer, + Encode, + Decode, get_burgers_data, - get_navier_stokes + get_darcy_flow_data + +struct UnitGaussianNormalizer{T} + mean::Array{T} + std::Array{T} + Ο΅::T +end + +function UnitGaussianNormalizer(𝐱; Ο΅=1f-5) + dims = 1:length(size(𝐱))-1 + + return UnitGaussianNormalizer(mean(𝐱, dims=dims), StatsBase.std(𝐱, dims=dims), Ο΅) +end + +struct Encode end +struct Decode end + +(n::UnitGaussianNormalizer)(𝐱, ::Type{Encode}) = @. (𝐱-n.mean) / (n.std+n.Ο΅) +(n::UnitGaussianNormalizer)(𝐱, ::Type{Decode}) = @. 𝐱 * (n.std+n.Ο΅) + n.mean + function register_burgers() register(DataDep( @@ -16,23 +38,23 @@ function register_burgers() )) end -# function register_navier_stokes() -# register(DataDep( -# "NavierStokes", -# """ -# Navier–Stokes equations dataset from -# [fourier_neural_operator](https://github.com/zongyi-li/fourier_neural_operator) -# """, -# "https://drive.google.com/file/d/1r3idxpsHa21ijhlu3QQ1hVuXcqnBTO7d/view?usp=sharing", -# "1a3b2893489dd1493923362bc74cd571e0b4f6ee290985eda060f4140df602d0", -# fetch_method=gdownload, -# post_fetch_method=unpack -# )) -# end +function register_darcy_flow() + register(DataDep( + "DarcyFlow", + """ + Darcy flow dataset from + [fourier_neural_operator](https://github.com/zongyi-li/fourier_neural_operator) + """, + "https://drive.google.com/file/d/1Z1uxG9R8AdAGJprG5STcphysjm56_0Jf/view?usp=sharing", + "802825de9da7398407296c99ca9ceb2371c752f6a3bdd1801172e02ce19edda4", + fetch_method=gdownload, + post_fetch_method=unpack + )) +end function register_datasets() register_burgers() - # register_navier_stokes() + register_darcy_flow() end function get_burgers_data(; n=2048, Ξ”samples=2^3, grid_size=div(2^13, Ξ”samples), T=Float32) @@ -48,6 +70,18 @@ function get_burgers_data(; n=2048, Ξ”samples=2^3, grid_size=div(2^13, Ξ”samples return x_loc_data, y_data end -# function get_navier_stokes() -# file = matopen(joinpath(datadep"NavierStokes", "ns_V1e-3_N5000_T50.mat")) -# end +function get_darcy_flow_data(; n=1024, Ξ”samples=5, T=Float32, test_data=true) + # size(training_data) == size(testing_data) == (1024, 421, 421) + file = test_data ? "piececonst_r421_N1024_smooth2.mat" : "piececonst_r421_N1024_smooth1.mat" + file = matopen(joinpath(datadep"DarcyFlow", file)) + x_data = T.(permutedims(read(file, "coeff")[1:n, 1:Ξ”samples:end, 1:Ξ”samples:end], (3, 2, 1))) + y_data = T.(permutedims(read(file, "sol")[1:n, 1:Ξ”samples:end, 1:Ξ”samples:end], (3, 2, 1))) + + x_dims = insert!([size(x_data)...], 3, 1) + y_dims = insert!([size(y_data)...], 3, 1) + x_data, y_data = reshape(x_data, x_dims...), reshape(y_data, y_dims...) + + x_normalizer, y_normalizer = UnitGaussianNormalizer(x_data), UnitGaussianNormalizer(y_data) + + return x_normalizer(x_data, Encode), y_normalizer(y_data, Encode), x_normalizer, y_normalizer +end diff --git a/test/data.jl b/test/data.jl index ba9033ae..3113e940 100644 --- a/test/data.jl +++ b/test/data.jl @@ -5,6 +5,20 @@ @test size(ys) == (1024, 1000) end -@testset "get navier stokes data" begin - # get_navier_stokes() +@testset "unit gaussian normalizer" begin + dims = (3, 3, 5, 6) + 𝐱 = rand(Float32, dims) + + n = UnitGaussianNormalizer(𝐱) + + @test size(n.mean) == size(n.std) + @test size(n(𝐱, Encode)) == dims + @test size(n(n(𝐱, Encode), Decode)) == dims +end + +@testset "get darcy flow data" begin + xs, ys, x_normalizer, y_normalizer = get_darcy_flow_data() + + @test size(xs) == (85, 85, 1, 1024) + @test size(ys) == (85, 85, 1, 1024) end From 4a341b2e01522b2c04c0dbf721221bec5f4701c5 Mon Sep 17 00:00:00 2001 From: JingYu Ning Date: Sun, 15 Aug 2021 17:34:41 +0800 Subject: [PATCH 05/14] deal with n-dim m --- src/fourier.jl | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/src/fourier.jl b/src/fourier.jl index eb9a2b31..6f4fa14c 100644 --- a/src/fourier.jl +++ b/src/fourier.jl @@ -22,7 +22,7 @@ function SpectralConv( ) where {S<:Integer, N} in_chs, out_chs = ch scale = one(T) / (in_chs * out_chs) - weights = scale * init(out_chs, in_chs, modes...) + weights = scale * init(out_chs, in_chs, prod(modes)) return SpectralConv(weights, in_chs, out_chs, modes, N, Οƒ) end @@ -31,7 +31,7 @@ Flux.@functor SpectralConv Base.ndims(::SpectralConv{N}) where {N} = N -spectral_conv(𝐱₁, 𝐱₂) = @tullio 𝐲[m, o, b] := 𝐱₁[m, i, b] * 𝐱₂[o, i, m] # TODO: extend `m` to n-dim +spectral_conv(𝐱₁, 𝐱₂) = @tullio 𝐲[m, o, b] := 𝐱₁[m, i, b] * 𝐱₂[o, i, m] function (m::SpectralConv)(𝐱::AbstractArray) 𝐱ᡀ = permutedims(Zygote.hook(real, 𝐱), (ndims(m)+1, 1:ndims(m)..., ndims(m)+2)) # [x, in_chs, batch] <- [in_chs, x, batch] @@ -39,11 +39,13 @@ function (m::SpectralConv)(𝐱::AbstractArray) # [modes, out_chs, batch] <- [modes, in_chs, batch] * [out_chs, in_chs, modes] ranges = [1:dim_modes for dim_modes in m.modes] - 𝐱_weighted = spectral_conv(view(𝐱_fft, ranges..., :, :), m.weight) + 𝐱_flattened = reshape(view(𝐱_fft, ranges..., :, :), prod(m.modes), size(𝐱_fft)[end-1:end]...) + 𝐱_weighted = spectral_conv(𝐱_flattened, m.weight) + 𝐱_shaped = reshape(𝐱_weighted, m.modes..., size(𝐱_weighted)[end-1:end]...) # [x, out_chs, batch] <- [modes, out_chs, batch] - pad = zeros(ComplexF32, (collect(size(𝐱_fft)[1:ndims(m)])-collect(m.modes))..., size(𝐱_weighted)[end-1:end]...) - 𝐱_padded = cat(𝐱_weighted, pad, dims=1:ndims(m)) + pad = zeros(ComplexF32, (collect(size(𝐱_fft)[1:ndims(m)])-collect(m.modes))..., size(𝐱_shaped)[end-1:end]...) + 𝐱_padded = cat(𝐱_shaped, pad, dims=1:ndims(m)) 𝐱_out = ifft(𝐱_padded, 1:ndims(m)) # [x, out_chs, batch] 𝐱_outα΅€ = permutedims(real(𝐱_out), (2:ndims(m)+1..., 1, ndims(m)+2)) # [out_chs, x, batch] <- [x, out_chs, batch] From 8e6e4a537fc3d0e6e53cb5b5c46cd4b7eb2dfc56 Mon Sep 17 00:00:00 2001 From: JingYu Ning Date: Sun, 15 Aug 2021 17:47:45 +0800 Subject: [PATCH 06/14] fix wrong dim --- src/data.jl | 4 ++-- test/data.jl | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/data.jl b/src/data.jl index 070442ed..d1a1fec3 100644 --- a/src/data.jl +++ b/src/data.jl @@ -77,8 +77,8 @@ function get_darcy_flow_data(; n=1024, Ξ”samples=5, T=Float32, test_data=true) x_data = T.(permutedims(read(file, "coeff")[1:n, 1:Ξ”samples:end, 1:Ξ”samples:end], (3, 2, 1))) y_data = T.(permutedims(read(file, "sol")[1:n, 1:Ξ”samples:end, 1:Ξ”samples:end], (3, 2, 1))) - x_dims = insert!([size(x_data)...], 3, 1) - y_dims = insert!([size(y_data)...], 3, 1) + x_dims = pushfirst!([size(x_data)...], 1) + y_dims = pushfirst!([size(y_data)...], 1) x_data, y_data = reshape(x_data, x_dims...), reshape(y_data, y_dims...) x_normalizer, y_normalizer = UnitGaussianNormalizer(x_data), UnitGaussianNormalizer(y_data) diff --git a/test/data.jl b/test/data.jl index 3113e940..4b4e6660 100644 --- a/test/data.jl +++ b/test/data.jl @@ -19,6 +19,6 @@ end @testset "get darcy flow data" begin xs, ys, x_normalizer, y_normalizer = get_darcy_flow_data() - @test size(xs) == (85, 85, 1, 1024) - @test size(ys) == (85, 85, 1, 1024) + @test size(xs) == (1, 85, 85, 1024) + @test size(ys) == (1, 85, 85, 1024) end From dc018ec0f8aac9994270279f626f9d9b61fe4d80 Mon Sep 17 00:00:00 2001 From: JingYu Ning Date: Sun, 15 Aug 2021 19:56:48 +0800 Subject: [PATCH 07/14] fix wrong dim on permutation and shrink test --- src/fourier.jl | 4 ++-- test/fourier.jl | 52 ++++++++++++++++++++++++++++++++++++++++--------- 2 files changed, 45 insertions(+), 11 deletions(-) diff --git a/src/fourier.jl b/src/fourier.jl index 6f4fa14c..5ac7d264 100644 --- a/src/fourier.jl +++ b/src/fourier.jl @@ -34,7 +34,7 @@ Base.ndims(::SpectralConv{N}) where {N} = N spectral_conv(𝐱₁, 𝐱₂) = @tullio 𝐲[m, o, b] := 𝐱₁[m, i, b] * 𝐱₂[o, i, m] function (m::SpectralConv)(𝐱::AbstractArray) - 𝐱ᡀ = permutedims(Zygote.hook(real, 𝐱), (ndims(m)+1, 1:ndims(m)..., ndims(m)+2)) # [x, in_chs, batch] <- [in_chs, x, batch] + 𝐱ᡀ = permutedims(Zygote.hook(real, 𝐱), (2:ndims(m)+1..., 1, ndims(m)+2)) # [x, in_chs, batch] <- [in_chs, x, batch] 𝐱_fft = fft(𝐱ᡀ, 1:ndims(m)) # [x, in_chs, batch] # [modes, out_chs, batch] <- [modes, in_chs, batch] * [out_chs, in_chs, modes] @@ -48,7 +48,7 @@ function (m::SpectralConv)(𝐱::AbstractArray) 𝐱_padded = cat(𝐱_shaped, pad, dims=1:ndims(m)) 𝐱_out = ifft(𝐱_padded, 1:ndims(m)) # [x, out_chs, batch] - 𝐱_outα΅€ = permutedims(real(𝐱_out), (2:ndims(m)+1..., 1, ndims(m)+2)) # [out_chs, x, batch] <- [x, out_chs, batch] + 𝐱_outα΅€ = permutedims(real(𝐱_out), (ndims(m)+1, 1:ndims(m)..., ndims(m)+2)) # [out_chs, x, batch] <- [x, out_chs, batch] return m.Οƒ.(𝐱_outα΅€) end diff --git a/test/fourier.jl b/test/fourier.jl index 5d31ac56..0898d755 100644 --- a/test/fourier.jl +++ b/test/fourier.jl @@ -1,4 +1,4 @@ -@testset "SpectralConv" begin +@testset "SpectralConv1d" begin modes = (16, ) ch = 64 => 64 @@ -8,16 +8,15 @@ ) @test ndims(SpectralConv(ch, modes)) == 1 - 𝐱, _ = get_burgers_data(n=1000) - @test size(m(𝐱)) == (64, 1024, 1000) + 𝐱, _ = get_burgers_data(n=5) + @test size(m(𝐱)) == (64, 1024, 5) - T = Float32 loss(x, y) = Flux.mse(m(x), y) - data = [(T.(𝐱[:, :, 1:5]), rand(T, 64, 1024, 5))] + data = [(𝐱, rand(Float32, 64, 1024, 5))] Flux.train!(loss, params(m), data, Flux.ADAM()) end -@testset "FourierOperator" begin +@testset "FourierOperator1d" begin modes = (16, ) ch = 64 => 64 @@ -26,10 +25,45 @@ end FourierOperator(ch, modes) ) - 𝐱, _ = get_burgers_data(n=1000) - @test size(m(𝐱)) == (64, 1024, 1000) + 𝐱, _ = get_burgers_data(n=5) + @test size(m(𝐱)) == (64, 1024, 5) loss(x, y) = Flux.mse(m(x), y) - data = [(Float32.(𝐱[:, :, 1:5]), rand(Float32, 64, 1024, 5))] + data = [(𝐱, rand(Float32, 64, 1024, 5))] + Flux.train!(loss, params(m), data, Flux.ADAM()) +end + +@testset "SpectralConv2d" begin + modes = (16, 16) + ch = 64 => 64 + + m = Chain( + Dense(1, 64), + SpectralConv(ch, modes) + ) + @test ndims(SpectralConv(ch, modes)) == 2 + + 𝐱, _, _, _ = get_darcy_flow_data(n=5, Ξ”samples=20) + @test size(m(𝐱)) == (64, 22, 22, 5) + + loss(x, y) = Flux.mse(m(x), y) + data = [(𝐱, rand(Float32, 64, 22, 22, 5))] + Flux.train!(loss, params(m), data, Flux.ADAM()) +end + +@testset "FourierOperator2d" begin + modes = (16, 16) + ch = 64 => 64 + + m = Chain( + Dense(1, 64), + FourierOperator(ch, modes) + ) + + 𝐱, _, _, _ = get_darcy_flow_data(n=5, Ξ”samples=20) + @test size(m(𝐱)) == (64, 22, 22, 5) + + loss(x, y) = Flux.mse(m(x), y) + data = [(𝐱, rand(Float32, 64, 22, 22, 5))] Flux.train!(loss, params(m), data, Flux.ADAM()) end From 58cf9947aa5e247cec8937a5ceb3086e4b634f0d Mon Sep 17 00:00:00 2001 From: JingYu Ning Date: Sun, 15 Aug 2021 20:18:07 +0800 Subject: [PATCH 08/14] revise api --- src/data.jl | 14 ++++++-------- test/data.jl | 6 +++--- 2 files changed, 9 insertions(+), 11 deletions(-) diff --git a/src/data.jl b/src/data.jl index d1a1fec3..f276519c 100644 --- a/src/data.jl +++ b/src/data.jl @@ -1,7 +1,7 @@ export UnitGaussianNormalizer, - Encode, - Decode, + encode, + decode, get_burgers_data, get_darcy_flow_data @@ -17,11 +17,8 @@ function UnitGaussianNormalizer(𝐱; Ο΅=1f-5) return UnitGaussianNormalizer(mean(𝐱, dims=dims), StatsBase.std(𝐱, dims=dims), Ο΅) end -struct Encode end -struct Decode end - -(n::UnitGaussianNormalizer)(𝐱, ::Type{Encode}) = @. (𝐱-n.mean) / (n.std+n.Ο΅) -(n::UnitGaussianNormalizer)(𝐱, ::Type{Decode}) = @. 𝐱 * (n.std+n.Ο΅) + n.mean +encode(n::UnitGaussianNormalizer, 𝐱::AbstractArray) = @. (𝐱-n.mean) / (n.std+n.Ο΅) +decode(n::UnitGaussianNormalizer, 𝐱::AbstractArray) = @. 𝐱 * (n.std+n.Ο΅) + n.mean function register_burgers() @@ -76,6 +73,7 @@ function get_darcy_flow_data(; n=1024, Ξ”samples=5, T=Float32, test_data=true) file = matopen(joinpath(datadep"DarcyFlow", file)) x_data = T.(permutedims(read(file, "coeff")[1:n, 1:Ξ”samples:end, 1:Ξ”samples:end], (3, 2, 1))) y_data = T.(permutedims(read(file, "sol")[1:n, 1:Ξ”samples:end, 1:Ξ”samples:end], (3, 2, 1))) + close(file) x_dims = pushfirst!([size(x_data)...], 1) y_dims = pushfirst!([size(y_data)...], 1) @@ -83,5 +81,5 @@ function get_darcy_flow_data(; n=1024, Ξ”samples=5, T=Float32, test_data=true) x_normalizer, y_normalizer = UnitGaussianNormalizer(x_data), UnitGaussianNormalizer(y_data) - return x_normalizer(x_data, Encode), y_normalizer(y_data, Encode), x_normalizer, y_normalizer + return encode(x_normalizer, x_data), encode(y_normalizer, y_data), x_normalizer, y_normalizer end diff --git a/test/data.jl b/test/data.jl index 4b4e6660..f20cae63 100644 --- a/test/data.jl +++ b/test/data.jl @@ -12,12 +12,12 @@ end n = UnitGaussianNormalizer(𝐱) @test size(n.mean) == size(n.std) - @test size(n(𝐱, Encode)) == dims - @test size(n(n(𝐱, Encode), Decode)) == dims + @test size(encode(n, 𝐱)) == dims + @test size(decode(n, encode(n, 𝐱))) == dims end @testset "get darcy flow data" begin - xs, ys, x_normalizer, y_normalizer = get_darcy_flow_data() + xs, ys, _, _ = get_darcy_flow_data() @test size(xs) == (1, 85, 85, 1024) @test size(ys) == (1, 85, 85, 1024) From c419e7254d62a79635c6aec5e461c50917a7c801 Mon Sep 17 00:00:00 2001 From: JingYu Ning Date: Sun, 15 Aug 2021 22:16:07 +0800 Subject: [PATCH 09/14] Apply suggestions from code review Co-authored-by: Peter --- src/fourier.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/fourier.jl b/src/fourier.jl index 5ac7d264..4ce6ceb0 100644 --- a/src/fourier.jl +++ b/src/fourier.jl @@ -34,17 +34,17 @@ Base.ndims(::SpectralConv{N}) where {N} = N spectral_conv(𝐱₁, 𝐱₂) = @tullio 𝐲[m, o, b] := 𝐱₁[m, i, b] * 𝐱₂[o, i, m] function (m::SpectralConv)(𝐱::AbstractArray) - 𝐱ᡀ = permutedims(Zygote.hook(real, 𝐱), (2:ndims(m)+1..., 1, ndims(m)+2)) # [x, in_chs, batch] <- [in_chs, x, batch] + 𝐱ᡀ = permutedims(Zygote.hook(real, 𝐱), (ntuple(i->i+1, ndims(m))..., 1, ndims(m)+2)) # [x, in_chs, batch] <- [in_chs, x, batch] 𝐱_fft = fft(𝐱ᡀ, 1:ndims(m)) # [x, in_chs, batch] # [modes, out_chs, batch] <- [modes, in_chs, batch] * [out_chs, in_chs, modes] ranges = [1:dim_modes for dim_modes in m.modes] 𝐱_flattened = reshape(view(𝐱_fft, ranges..., :, :), prod(m.modes), size(𝐱_fft)[end-1:end]...) 𝐱_weighted = spectral_conv(𝐱_flattened, m.weight) - 𝐱_shaped = reshape(𝐱_weighted, m.modes..., size(𝐱_weighted)[end-1:end]...) + 𝐱_shaped = reshape(𝐱_weighted, m.modes..., size(𝐱_weighted, xlen-1), size(𝐱_weighted, xlen)) # [x, out_chs, batch] <- [modes, out_chs, batch] - pad = zeros(ComplexF32, (collect(size(𝐱_fft)[1:ndims(m)])-collect(m.modes))..., size(𝐱_shaped)[end-1:end]...) + pad = zeros(ComplexF32, ntuple(i->size(𝐱_fft, i) - m.modes[i] , ndims(m))..., size(𝐱_shaped, xlen-1), size(𝐱_shaped, xlen)) 𝐱_padded = cat(𝐱_shaped, pad, dims=1:ndims(m)) 𝐱_out = ifft(𝐱_padded, 1:ndims(m)) # [x, out_chs, batch] From 483dfcb258c6ca35ebf595b87bec5d906bfd9e9a Mon Sep 17 00:00:00 2001 From: JingYu Ning Date: Sun, 15 Aug 2021 22:17:13 +0800 Subject: [PATCH 10/14] Apply suggestions from code review --- docs/Manifest.toml | 24 ++++++++++++------------ src/data.jl | 2 +- src/fourier.jl | 4 ++-- 3 files changed, 15 insertions(+), 15 deletions(-) diff --git a/docs/Manifest.toml b/docs/Manifest.toml index 8d1917da..8c333a77 100644 --- a/docs/Manifest.toml +++ b/docs/Manifest.toml @@ -73,10 +73,10 @@ uuid = "fa961155-64e5-5f13-b03f-caf6b980ea82" version = "0.4.1" [[CUDA]] -deps = ["AbstractFFTs", "Adapt", "BFloat16s", "CEnum", "CompilerSupportLibraries_jll", "DataStructures", "ExprTools", "GPUArrays", "GPUCompiler", "LLVM", "LazyArtifacts", "Libdl", "LinearAlgebra", "Logging", "Printf", "Random", "Random123", "RandomNumbers", "Reexport", "Requires", "SparseArrays", "SpecialFunctions", "TimerOutputs"] -git-tree-sha1 = "889889f1c13467406a126cd2789b4844487ddfc1" +deps = ["AbstractFFTs", "Adapt", "BFloat16s", "CEnum", "CompilerSupportLibraries_jll", "ExprTools", "GPUArrays", "GPUCompiler", "LLVM", "LazyArtifacts", "Libdl", "LinearAlgebra", "Logging", "Printf", "Random", "Random123", "RandomNumbers", "Reexport", "Requires", "SparseArrays", "SpecialFunctions", "TimerOutputs"] +git-tree-sha1 = "9303b20dfa74e4bcb4da425d351d551fbb5850be" uuid = "052768ef-5323-5732-b1bb-66c8b64840ba" -version = "3.3.5" +version = "3.4.0" [[CUDAKernels]] deps = ["Adapt", "CUDA", "Cassette", "KernelAbstractions", "SpecialFunctions", "StaticArrays"] @@ -91,9 +91,9 @@ version = "0.3.7" [[ChainRules]] deps = ["ChainRulesCore", "Compat", "LinearAlgebra", "Random", "Statistics"] -git-tree-sha1 = "11567f2471013449c2fcf119f674c681484a130e" +git-tree-sha1 = "6615deb51db68c3fa0cc7b34e5399c15f63fa97e" uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2" -version = "1.5.1" +version = "1.7.0" [[ChainRulesCore]] deps = ["Compat", "LinearAlgebra", "SparseArrays"] @@ -168,9 +168,9 @@ version = "1.0.3" [[DiffRules]] deps = ["NaNMath", "Random", "SpecialFunctions"] -git-tree-sha1 = "85d2d9e2524da988bffaf2a381864e20d2dae08d" +git-tree-sha1 = "3ed8fa7178a10d1cd0f1ca524f249ba6937490c0" uuid = "b552c78f-8df3-52c6-915a-8e097449b14b" -version = "1.2.1" +version = "1.3.0" [[Distributed]] deps = ["Random", "Serialization", "Sockets"] @@ -247,15 +247,15 @@ version = "0.2.3" [[GPUArrays]] deps = ["Adapt", "LinearAlgebra", "Printf", "Random", "Serialization", "Statistics"] -git-tree-sha1 = "8034b1a19f7a19743c53cda450fcc65d1b8f7ab5" +git-tree-sha1 = "8fac1cf7d6ce0f2249c7acaf25d22e1e85c4a07f" uuid = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7" -version = "8.0.1" +version = "8.0.2" [[GPUCompiler]] deps = ["ExprTools", "InteractiveUtils", "LLVM", "Libdl", "Logging", "TimerOutputs", "UUIDs"] -git-tree-sha1 = "f26f15d9c353f7091065390ea826df9e03917e58" +git-tree-sha1 = "4ed2616d5e656c8716736b64da86755467f26cf5" uuid = "61eb1bfa-7361-4325-ad38-22787b887f55" -version = "0.12.8" +version = "0.12.9" [[HDF5]] deps = ["Blosc", "Compat", "HDF5_jll", "Libdl", "Mmap", "Random", "Requires"] @@ -468,7 +468,7 @@ version = "0.3.5" uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908" [[NeuralOperators]] -deps = ["CUDA", "CUDAKernels", "DataDeps", "FFTW", "Fetch", "Flux", "KernelAbstractions", "MAT", "Tullio", "Zygote"] +deps = ["CUDA", "CUDAKernels", "DataDeps", "FFTW", "Fetch", "Flux", "KernelAbstractions", "MAT", "StatsBase", "Tullio", "Zygote"] path = ".." uuid = "ea5c82af-86e5-48da-8ee1-382d6ad7af4b" version = "0.1.0" diff --git a/src/data.jl b/src/data.jl index f276519c..3dc5e84f 100644 --- a/src/data.jl +++ b/src/data.jl @@ -12,7 +12,7 @@ struct UnitGaussianNormalizer{T} end function UnitGaussianNormalizer(𝐱; Ο΅=1f-5) - dims = 1:length(size(𝐱))-1 + dims = 1:ndims(𝐱)-1 return UnitGaussianNormalizer(mean(𝐱, dims=dims), StatsBase.std(𝐱, dims=dims), Ο΅) end diff --git a/src/fourier.jl b/src/fourier.jl index 4ce6ceb0..5f3356b7 100644 --- a/src/fourier.jl +++ b/src/fourier.jl @@ -38,8 +38,8 @@ function (m::SpectralConv)(𝐱::AbstractArray) 𝐱_fft = fft(𝐱ᡀ, 1:ndims(m)) # [x, in_chs, batch] # [modes, out_chs, batch] <- [modes, in_chs, batch] * [out_chs, in_chs, modes] - ranges = [1:dim_modes for dim_modes in m.modes] - 𝐱_flattened = reshape(view(𝐱_fft, ranges..., :, :), prod(m.modes), size(𝐱_fft)[end-1:end]...) + xlen = ndims(𝐱_fft) + 𝐱_flattened = reshape(view(𝐱_fft, map(d->1:d, m.modes)..., :, :), :, size(𝐱_fft, xlen-1), size(𝐱_fft, xlen)) 𝐱_weighted = spectral_conv(𝐱_flattened, m.weight) 𝐱_shaped = reshape(𝐱_weighted, m.modes..., size(𝐱_weighted, xlen-1), size(𝐱_weighted, xlen)) From 7225fe58812f52f6f9c47c82ff69da2281533740 Mon Sep 17 00:00:00 2001 From: JingYu Ning Date: Sun, 15 Aug 2021 22:58:22 +0800 Subject: [PATCH 11/14] fix dims for shaped array --- src/fourier.jl | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/src/fourier.jl b/src/fourier.jl index 5f3356b7..f1b2ae79 100644 --- a/src/fourier.jl +++ b/src/fourier.jl @@ -31,20 +31,21 @@ Flux.@functor SpectralConv Base.ndims(::SpectralConv{N}) where {N} = N +# [prod(m.modes), out_chs, batch] <- [prod(m.modes), in_chs, batch] * [out_chs, in_chs, prod(m.modes)] spectral_conv(𝐱₁, 𝐱₂) = @tullio 𝐲[m, o, b] := 𝐱₁[m, i, b] * 𝐱₂[o, i, m] function (m::SpectralConv)(𝐱::AbstractArray) + n_dims = ndims(𝐱) + 𝐱ᡀ = permutedims(Zygote.hook(real, 𝐱), (ntuple(i->i+1, ndims(m))..., 1, ndims(m)+2)) # [x, in_chs, batch] <- [in_chs, x, batch] 𝐱_fft = fft(𝐱ᡀ, 1:ndims(m)) # [x, in_chs, batch] - # [modes, out_chs, batch] <- [modes, in_chs, batch] * [out_chs, in_chs, modes] - xlen = ndims(𝐱_fft) - 𝐱_flattened = reshape(view(𝐱_fft, map(d->1:d, m.modes)..., :, :), :, size(𝐱_fft, xlen-1), size(𝐱_fft, xlen)) - 𝐱_weighted = spectral_conv(𝐱_flattened, m.weight) - 𝐱_shaped = reshape(𝐱_weighted, m.modes..., size(𝐱_weighted, xlen-1), size(𝐱_weighted, xlen)) + 𝐱_flattened = reshape(view(𝐱_fft, map(d->1:d, m.modes)..., :, :), :, size(𝐱_fft, n_dims-1), size(𝐱_fft, n_dims)) + 𝐱_weighted = spectral_conv(𝐱_flattened, m.weight) # [prod(m.modes), out_chs, batch], only 3-dims + 𝐱_shaped = reshape(𝐱_weighted, m.modes..., size(𝐱_weighted, 2), size(𝐱_weighted, 3)) # [x, out_chs, batch] <- [modes, out_chs, batch] - pad = zeros(ComplexF32, ntuple(i->size(𝐱_fft, i) - m.modes[i] , ndims(m))..., size(𝐱_shaped, xlen-1), size(𝐱_shaped, xlen)) + pad = zeros(ComplexF32, ntuple(i->size(𝐱_fft, i)-m.modes[i], ndims(m))..., size(𝐱_shaped, n_dims-1), size(𝐱_shaped, n_dims)) 𝐱_padded = cat(𝐱_shaped, pad, dims=1:ndims(m)) 𝐱_out = ifft(𝐱_padded, 1:ndims(m)) # [x, out_chs, batch] From 99d775d1460d30a90dccc5bf66d3272a057ce83d Mon Sep 17 00:00:00 2001 From: JingYu Ning Date: Sun, 15 Aug 2021 23:01:14 +0800 Subject: [PATCH 12/14] update deps for doc --- docs/Manifest.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/Manifest.toml b/docs/Manifest.toml index 8c333a77..569ac435 100644 --- a/docs/Manifest.toml +++ b/docs/Manifest.toml @@ -211,9 +211,9 @@ version = "3.3.9+8" [[Fetch]] deps = ["Base64", "HTTP", "JSON3", "Random", "StructTypes"] -git-tree-sha1 = "805a7f0edd71138f053b572613e918ef147625f0" +git-tree-sha1 = "84ba4219db49572bc3020589e77db293707aad51" uuid = "bb354801-46f6-40b6-9c3d-d42d7a74c775" -version = "0.1.0" +version = "0.1.1" [[FillArrays]] deps = ["LinearAlgebra", "Random", "SparseArrays", "Statistics"] From 3e06c3a0b33a32341c96efe6a0edcb1794d33771 Mon Sep 17 00:00:00 2001 From: JingYu Ning Date: Sun, 15 Aug 2021 23:46:30 +0800 Subject: [PATCH 13/14] update compat --- Project.toml | 1 + src/data.jl | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index e29e3baa..20c3fa4e 100644 --- a/Project.toml +++ b/Project.toml @@ -21,6 +21,7 @@ CUDA = "3.3" CUDAKernels = "0.3" DataDeps = "0.7" FFTW = "1.4" +Fetch = "0.1" Flux = "0.12" KernelAbstractions = "0.7" MAT = "0.10" diff --git a/src/data.jl b/src/data.jl index 3dc5e84f..08bcf568 100644 --- a/src/data.jl +++ b/src/data.jl @@ -67,7 +67,7 @@ function get_burgers_data(; n=2048, Ξ”samples=2^3, grid_size=div(2^13, Ξ”samples return x_loc_data, y_data end -function get_darcy_flow_data(; n=1024, Ξ”samples=5, T=Float32, test_data=true) +function get_darcy_flow_data(; n=1024, Ξ”samples=5, T=Float32, test_data=false) # size(training_data) == size(testing_data) == (1024, 421, 421) file = test_data ? "piececonst_r421_N1024_smooth2.mat" : "piececonst_r421_N1024_smooth1.mat" file = matopen(joinpath(datadep"DarcyFlow", file)) From 01167c51b5a8e7afcbefc86b91b22574ff405a90 Mon Sep 17 00:00:00 2001 From: JingYu Ning Date: Mon, 16 Aug 2021 00:30:48 +0800 Subject: [PATCH 14/14] diversion --- src/data.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/data.jl b/src/data.jl index 08bcf568..376baa1c 100644 --- a/src/data.jl +++ b/src/data.jl @@ -28,7 +28,7 @@ function register_burgers() Burgers' equation dataset from [fourier_neural_operator](https://github.com/zongyi-li/fourier_neural_operator) """, - "https://drive.google.com/file/d/16a8od4vidbiNR3WtaBPCSZ0T3moxjhYe/view?usp=sharing", + "https://drive.google.com/file/d/17MYsKzxUQVaLMWodzPbffR8hhDHoadPp/view?usp=sharing", "9cbbe5070556c777b1ba3bacd49da5c36ea8ed138ba51b6ee76a24b971066ecd", fetch_method=gdownload, post_fetch_method=unpack @@ -42,7 +42,7 @@ function register_darcy_flow() Darcy flow dataset from [fourier_neural_operator](https://github.com/zongyi-li/fourier_neural_operator) """, - "https://drive.google.com/file/d/1Z1uxG9R8AdAGJprG5STcphysjm56_0Jf/view?usp=sharing", + "https://drive.google.com/file/d/1zzVMuGhOG70EnR5L24LWqmX9-Wh_H5Wu/view?usp=sharing", "802825de9da7398407296c99ca9ceb2371c752f6a3bdd1801172e02ce19edda4", fetch_method=gdownload, post_fetch_method=unpack