revise show

Author: foldfelis

src/fourier.jl

@@ -28,15 +28,15 @@ end
 
 ```jldoctest
 julia> SpectralConv(2=>5, (16, ))
-SpectralConv(2 => 5, (16,), σ=identity)
+SpectralConv(2 => 5, (16,), σ=identity, permuted=false)
 
 julia> using Flux
 
 julia> SpectralConv(2=>5, (16, ), relu)
-SpectralConv(2 => 5, (16,), σ=relu)
+SpectralConv(2 => 5, (16,), σ=relu, permuted=false)
 
 julia> SpectralConv(2=>5, (16, ), relu, permuted=true)
-SpectralConv(2 => 5, (16,), σ=relu)
+SpectralConv(2 => 5, (16,), σ=relu, permuted=true)
 ```
 """
 function SpectralConv(
@@ -60,8 +60,10 @@ Flux.@functor SpectralConv
 
 Base.ndims(::SpectralConv{P,N}) where {P,N} = N
 
-function Base.show(io::IO, l::SpectralConv)
-    print(io, "SpectralConv($(l.in_channel) => $(l.out_channel), $(l.modes), σ=$(string(l.σ)))")
+permuted(::SpectralConv{P}) where {P} = P
+
+function Base.show(io::IO, l::SpectralConv{P}) where {P}
+    print(io, "SpectralConv($(l.in_channel) => $(l.out_channel), $(l.modes), σ=$(string(l.σ)), permuted=$P)")
 end
 
 function spectral_conv(m::SpectralConv, 𝐱::AbstractArray)
@@ -114,36 +116,15 @@ end
 
 ```jldoctest
 julia> FourierOperator(2=>5, (16, ))
-Chain(
-  Parallel(
-    +,
-    Dense(2, 5),                        # 15 parameters
-    SpectralConv(2 => 5, (16,), σ=identity),  # 160 parameters
-  ),
-  NeuralOperators.var"#activation_func#14"{typeof(identity)}(identity),
-)                   # Total: 3 arrays, 175 parameters, 1.668 KiB.
+FourierOperator(2 => 5, (16,), σ=identity, permuted=false)
 
 julia> using Flux
 
 julia> FourierOperator(2=>5, (16, ), relu)
-Chain(
-  Parallel(
-    +,
-    Dense(2, 5),                        # 15 parameters
-    SpectralConv(2 => 5, (16,), σ=identity),  # 160 parameters
-  ),
-  NeuralOperators.var"#activation_func#14"{typeof(relu)}(NNlib.relu),
-)                   # Total: 3 arrays, 175 parameters, 1.668 KiB.
+FourierOperator(2 => 5, (16,), σ=relu, permuted=false)
 
 julia> FourierOperator(2=>5, (16, ), relu, permuted=true)
-Chain(
-  Parallel(
-    +,
-    Conv((1,), 2 => 5),                 # 15 parameters
-    SpectralConvPerm(2 => 5, (16,), σ=identity),  # 160 parameters
-  ),
-  NeuralOperators.var"#activation_func#14"{typeof(relu)}(NNlib.relu),
-)                   # Total: 3 arrays, 175 parameters, 1.871 KiB.
+FourierOperator(2 => 5, (16,), σ=relu, permuted=true)
 ```
 """
 function FourierOperator(
@@ -160,6 +141,10 @@ end
 
 Flux.@functor FourierOperator
 
+function Base.show(io::IO, l::FourierOperator)
+    print(io, "FourierOperator($(l.conv.in_channel) => $(l.conv.out_channel), $(l.conv.modes), σ=$(string(l.σ)), permuted=$(permuted(l.conv)))")
+end
+
 function (m::FourierOperator)(𝐱)
     return m.σ.(m.linear(𝐱) + m.conv(𝐱))
 end

test/fourier.jl

@@ -7,7 +7,7 @@
         SpectralConv(ch, modes)
     )
     @test ndims(SpectralConv(ch, modes)) == 1
-    @test repr(SpectralConv(ch, modes)) == "SpectralConv(64 => 128, (16,), σ=identity)"
+    @test repr(SpectralConv(ch, modes)) == "SpectralConv(64 => 128, (16,), σ=identity, permuted=false)"
 
     𝐱, _ = get_burgers_data(n=5)
     @test size(m(𝐱)) == (128, 1024, 5)
@@ -26,6 +26,7 @@ end
         SpectralConv(ch, modes, permuted=true)
     )
     @test ndims(SpectralConv(ch, modes, permuted=true)) == 1
+    @test repr(SpectralConv(ch, modes, permuted=true)) == "SpectralConv(64 => 128, (16,), σ=identity, permuted=true)"
 
     𝐱, _ = get_burgers_data(n=5)
     𝐱 = permutedims(𝐱, (2, 1, 3))
@@ -44,6 +45,7 @@ end
         Dense(2, 64),
         FourierOperator(ch, modes)
     )
+    @test repr(FourierOperator(ch, modes)) == "FourierOperator(64 => 128, (16,), σ=identity, permuted=false)"
 
     𝐱, _ = get_burgers_data(n=5)
     @test size(m(𝐱)) == (128, 1024, 5)
@@ -61,6 +63,7 @@ end
         Conv((1, ), 2=>64),
         FourierOperator(ch, modes, permuted=true)
     )
+    @test repr(FourierOperator(ch, modes, permuted=true)) == "FourierOperator(64 => 128, (16,), σ=identity, permuted=true)"
 
     𝐱, _ = get_burgers_data(n=5)
     𝐱 = permutedims(𝐱, (2, 1, 3))