fix: rename to centered_truncation

Author: avik-pal

docs/src/tutorials/burgers_deeponet.md

@@ -131,9 +131,10 @@ draw(
         :data_sequence => L"u(x)";
         color=:label => "",
         layout=:sequence => nonnumeric,
+        linestyle=:label => "",
     ) *
-    visual(Lines),
-    scales(; Color=(; palette=:tab10));
+    visual(Lines; linewidth=4),
+    scales(; Color=(; palette=:tab10), LineStyle = (; palette = [:solid, :dash]));
     figure=(;
         size=(1024, 1024),
         title="Using DeepONet to solve the Burgers equation",

docs/src/tutorials/burgers_fno.md

@@ -46,10 +46,7 @@ x_data = reshape(T.(collect(read(file, "a")[1:N, 1:Δsamples:end])), N, :)
 y_data = reshape(T.(collect(read(file, "u")[1:N, 1:Δsamples:end])), N, :)
 close(file)
 
-x_data = hcat(
-    repeat(reshape(collect(T, range(0, 1; length=grid_size)), :, 1, 1), 1, 1, N),
-    reshape(permutedims(x_data, (2, 1)), grid_size, 1, N)
-);
+x_data = reshape(permutedims(x_data, (2, 1)), grid_size, 1, N);
 y_data = reshape(permutedims(y_data, (2, 1)), grid_size, 1, N);
 ```
 
@@ -62,9 +59,7 @@ const cdev = cpu_device()
 const xdev = reactant_device(; force=true)
 
 fno = FourierNeuralOperator(
-    gelu;
-    chs = (2, 32, 32, 32, 1),
-    modes = (16,)
+    (16,), 2, 1, 32; activation=gelu, stabilizer=tanh, centered_truncation=true
 )
 ps, st = Lux.setup(Random.default_rng(), fno) |> xdev;
 ```
@@ -74,13 +69,16 @@ ps, st = Lux.setup(Random.default_rng(), fno) |> xdev;
 ```@example burgers_fno
 dataloader = DataLoader((x_data, y_data); batchsize=128, shuffle=true) |> xdev;
 
-function train_model!(model, ps, st, dataloader; epochs=5000)
+function train_model!(model, ps, st, dataloader; epochs=1000)
     train_state = Training.TrainState(model, ps, st, Adam(0.0001f0))
 
-    for epoch in 1:epochs, data in dataloader
-        (_, loss, _, train_state) = Training.single_train_step!(
-            AutoEnzyme(), MAELoss(), data, train_state; return_gradients=Val(false)
-        )
+    for epoch in 1:epochs
+        loss = -Inf
+        for data in dataloader
+            (_, loss, _, train_state) = Training.single_train_step!(
+                AutoEnzyme(), MAELoss(), data, train_state; return_gradients=Val(false)
+            )
+        end
 
         if epoch % 100 == 1 || epoch == epochs
             @printf("Epoch %d: loss = %.6e\n", epoch, loss)
@@ -90,7 +88,7 @@ function train_model!(model, ps, st, dataloader; epochs=5000)
     return train_state.parameters, train_state.states
 end
 
-(ps_trained, st_trained) = train_model!(fno, ps, st, dataloader)
+ps_trained, st_trained = train_model!(fno, ps, st, dataloader)
 nothing #hide
 ```
 
@@ -104,7 +102,7 @@ AoG.set_aog_theme!()
 x_data_dev = x_data |> xdev;
 y_data_dev = y_data |> xdev;
 
-grid = x_data[:, 1, :]
+grid = range(0, 1; length=grid_size)
 pred = first(
     Reactant.with_config(;
         convolution_precision=PrecisionConfig.HIGH,
@@ -116,7 +114,7 @@ pred = first(
 
 data_sequence, sequence, repeated_grid, label = Float32[], Int[], Float32[], String[]
 for i in 1:16
-    append!(repeated_grid, vcat(grid[:, i], grid[:, i]))
+    append!(repeated_grid, repeat(grid, 2))
     append!(sequence, repeat([i], grid_size * 2))
     append!(label, repeat(["Ground Truth"], grid_size))
     append!(label, repeat(["Predictions"], grid_size))
@@ -135,7 +133,7 @@ draw(
         linestyle=:label => "",
     ) *
     visual(Lines; linewidth=4),
-    scales(; Color=(; palette=:tab10), LineStyle = (; palette = [:solid, :dash, :dot]));
+    scales(; Color=(; palette=:tab10), LineStyle = (; palette = [:solid, :dash]));
     figure=(;
         size=(1024, 1024),
         title="Using FNO to solve the Burgers equation",

src/layers.jl

@@ -89,9 +89,11 @@ julia> SpectralConv(2 => 5, (16,));
 ```
 """
 function SpectralConv(
-    ch::Pair{<:Integer,<:Integer}, modes::Dims; shift::Bool=false, kwargs...
+    ch::Pair{<:Integer,<:Integer}, modes::Dims; centered_truncation::Bool=false, kwargs...
 )
-    return OperatorConv(ch, modes, FourierTransform{ComplexF32}(modes, shift); kwargs...)
+    return OperatorConv(
+        ch, modes, FourierTransform{ComplexF32}(modes, centered_truncation); kwargs...
+    )
 end
 
 """
@@ -129,8 +131,8 @@ function OperatorKernel(
     stabilizer=identity,
     complex_data::Bool=false,
     fno_skip::Symbol=:linear,
-    channel_mlp_skip::Symbol=:linear,
-    use_channel_mlp::Bool=true,
+    channel_mlp_skip::Symbol=:soft_gating,
+    use_channel_mlp::Bool=false,
     channel_mlp_expansion::Real=0.5,
     kwargs...,
 ) where {N}
@@ -205,10 +207,14 @@ julia> SpectralKernel(2 => 5, (16,));
 ```
 """
 function SpectralKernel(
-    ch::Pair{<:Integer,<:Integer}, modes::Dims, act=identity; shift::Bool=false, kwargs...
+    ch::Pair{<:Integer,<:Integer},
+    modes::Dims,
+    act=identity;
+    centered_truncation::Bool=false,
+    kwargs...,
 )
     return OperatorKernel(
-        ch, modes, FourierTransform{ComplexF32}(modes, shift), act; kwargs...
+        ch, modes, FourierTransform{ComplexF32}(modes, centered_truncation), act; kwargs...
     )
 end
 

src/models/fno.jl

@@ -81,7 +81,7 @@ function FourierNeuralOperator(
     fno_skip::Symbol=:linear,
     complex_data::Bool=false,
     stabilizer=tanh,
-    shift::Bool=false,
+    centered_truncation::Bool=false,
 ) where {N}
     lifting_channels = hidden_channels * lifting_channel_ratio
     projection_channels = out_channels * projection_channel_ratio
@@ -114,7 +114,7 @@ function FourierNeuralOperator(
                 modes,
                 activation;
                 stabilizer,
-                shift,
+                centered_truncation,
                 use_channel_mlp,
                 channel_mlp_expansion,
                 channel_mlp_skip,

src/transform.jl

@@ -19,30 +19,30 @@ function truncate_modes end
 function inverse end
 
 """
-    FourierTransform{T}(modes, shift::Bool=false)
+    FourierTransform{T}(modes, centered_truncation::Bool=false)
 
 A concrete implementation of `AbstractTransform` for Fourier transforms.
 """
 struct FourierTransform{T,M} <: AbstractTransform{T}
     modes::M
-    shift::Bool
+    centered_truncation::Bool
 end
 
-function FourierTransform{T}(modes::Dims, shift::Bool=false) where {T}
-    return FourierTransform{T,typeof(modes)}(modes, shift)
+function FourierTransform{T}(modes::Dims, centered_truncation::Bool=false) where {T}
+    return FourierTransform{T,typeof(modes)}(modes, centered_truncation)
 end
 
 function Base.show(io::IO, ft::FourierTransform)
     print(io, "FourierTransform{", eltype(ft), "}(")
-    print(io, ft.modes, ", shift=", ft.shift, ")")
+    print(io, ft.modes, ", centered_truncation=", ft.centered_truncation, ")")
     return nothing
 end
 
 Base.ndims(T::FourierTransform) = length(T.modes)
 
 function transform(ft::FourierTransform, x::AbstractArray)
     res = Lux.Utils.eltype(x) <: Complex ? fft(x, 1:ndims(ft)) : rfft(x, 1:ndims(ft))
-    if ft.shift && ndims(ft) > 1
+    if ft.centered_truncation && ndims(ft) > 1
         res = fftshift(res, 1:ndims(ft))
     end
     return res
@@ -57,7 +57,7 @@ truncate_modes(ft::FourierTransform, x_fft::AbstractArray) = low_pass(ft, x_fft)
 function inverse(
     ft::FourierTransform, x_fft::AbstractArray{T,N}, x::AbstractArray{T2,N}
 ) where {T,T2,N}
-    if ft.shift && ndims(ft) > 1
+    if ft.centered_truncation && ndims(ft) > 1
         x_fft = fftshift(x_fft, 1:ndims(ft))
     end
 

test/fno_tests.jl

@@ -7,26 +7,27 @@
             chs=(2, 64, 64, 64, 64, 64, 128, 1),
             x_size=(1024, 2, 5),
             y_size=(1024, 1, 5),
-            shift=false,
+            centered_truncation=false,
         ),
         (
             modes=(16, 16),
             chs=(2, 64, 64, 64, 64, 64, 128, 4),
             x_size=(32, 32, 2, 5),
             y_size=(32, 32, 4, 5),
-            shift=false,
+            centered_truncation=false,
         ),
         (
             modes=(16, 16),
             chs=(2, 64, 64, 64, 64, 64, 128, 4),
             x_size=(32, 32, 2, 5),
             y_size=(32, 32, 4, 5),
-            shift=true,
+            centered_truncation=true,
         ),
     ]
 
-    @testset "$(length(setup.modes))D | shift=$(setup.shift)" for setup in setups
-        fno = FourierNeuralOperator(; setup.chs, setup.modes, setup.shift)
+    @testset "$(length(setup.modes))D | centered_truncation=$(setup.centered_truncation)" for setup in
+                                                                                              setups
+        fno = FourierNeuralOperator(; setup.chs, setup.modes, setup.centered_truncation)
         display(fno)
         ps, st = Lux.setup(rng, fno)
 

test/layers_tests.jl

@@ -3,22 +3,33 @@
 
     opconv = [SpectralConv, SpectralKernel]
     setups = [
-        (; m=(16,), x_size=(1024, 2, 5), y_size=(1024, 16, 5), shift=false),
-        (; m=(10, 10), x_size=(22, 22, 1, 5), y_size=(22, 22, 16, 5), shift=false),
-        (; m=(10, 10), x_size=(22, 22, 1, 5), y_size=(22, 22, 16, 5), shift=true),
+        (; m=(16,), x_size=(1024, 2, 5), y_size=(1024, 16, 5), centered_truncation=false),
+        (;
+            m=(10, 10),
+            x_size=(22, 22, 1, 5),
+            y_size=(22, 22, 16, 5),
+            centered_truncation=false,
+        ),
+        (;
+            m=(10, 10),
+            x_size=(22, 22, 1, 5),
+            y_size=(22, 22, 16, 5),
+            centered_truncation=true,
+        ),
     ]
 
     rdev = reactant_device()
 
-    @testset "$(op) $(length(setup.m))D | shift=$(setup.shift)" for op in opconv,
+    @testset "$(op) $(length(setup.m))D | centered_truncation=$(setup.centered_truncation)" for op in
+                                                                                                opconv,
         setup in setups
 
         in_chs = setup.x_size[end - 1]
         out_chs = setup.y_size[end - 1]
         ch = 4 => out_chs
 
         l1 = Conv(ntuple(_ -> 1, length(setup.m)), in_chs => first(ch))
-        m = Chain(l1, op(ch, setup.m; setup.shift))
+        m = Chain(l1, op(ch, setup.m; setup.centered_truncation))
         display(m)
         ps, st = Lux.setup(rng, m)