update run scripts

examples/diffusion_fourier/exp_2d/exp_FNO_linear_nonlinear/statistics.txt

@@ -1,17 +1,17 @@
-	 TRAIN LOSS: 0.0106195	 TEST LOSS: 0.0122039
+	 TRAIN LOSS: 0.01120188	 TEST LOSS: 0.01220754
 #======================#
 TRAIN STATS
-R² score:       0.83026416
-mean SQR error: 0.01107069
-mean ABS error: 0.08505448
-max  ABS error: 0.46613726
+R² score:       0.82948816
+mean SQR error: 0.01120188
+mean ABS error: 0.085977
+max  ABS error: 0.41472676
 
 #======================#
 #======================#
 TEST  STATS
-R² score:       0.84946002
-mean SQR error: 0.01146579
-mean ABS error: 0.08442349
-max  ABS error: 0.54226281
+R² score:       0.8304977
+mean SQR error: 0.01220754
+mean ABS error: 0.08744246
+max  ABS error: 0.48108292
 
 #======================#

examples/diffusion_fourier/exp_2d/exp_FNO_nonlin/statistics.txt

@@ -1,17 +1,17 @@
-	 TRAIN LOSS: 0.00027291	 TEST LOSS: 0.02129042
+	 TRAIN LOSS: 0.00580268	 TEST LOSS: 0.01472381
 #======================#
 TRAIN STATS
-R² score:       0.99641058
-mean SQR error: 0.00026939
-mean ABS error: 0.01296109
-max  ABS error: 0.14936594
+R² score:       0.9087739
+mean SQR error: 0.00580268
+mean ABS error: 0.06006256
+max  ABS error: 0.4061329
 
 #======================#
 #======================#
 TEST  STATS
-R² score:       0.66283454
-mean SQR error: 0.02190537
-mean ABS error: 0.11673716
-max  ABS error: 0.80765662
+R² score:       0.7960545
+mean SQR error: 0.01472381
+mean ABS error: 0.09616495
+max  ABS error: 0.58538777
 
 #======================#

examples/diffusion_fourier/exp_2d/model.jl

@@ -47,33 +47,35 @@ include("../datagen.jl")
 BSON.@load joinpath(@__DIR__, "..", "data2D_N$(N).bson") _data data_
 
 rng = Random.default_rng()
-Random.seed!(rng, 345)
+Random.seed!(rng, 123)
 
-__data = combine_data2D(_data)
-data__ = combine_data2D(data_)
+__data = combine_data2D(_data, 128)
+data__ = combine_data2D(data_, 128)
 
 ###
 # FNO model
 ###
-if true
+if false
 
-w = 12        # width
-m = (12, 12,) # modes
+w = 16        # width
+m = (16, 16,) # modes
 c = size(__data[1], 1) # in  channels
 o = size(__data[2], 1) # out channels
 
 NN = Lux.Chain(
-    Dense(c , w, tanh),
+    PermutedBatchNorm(c, 4),
+    Dense(c, w, tanh),
     OpKernel(w, w, m, tanh),
     OpKernel(w, w, m, tanh),
     OpKernel(w, w, m, tanh),
-    Dense(w , o)
+    OpKernel(w, w, m, tanh),
+    Dense(w, o),
 )
 
 opt = Optimisers.Adam()
 batchsize = 128
 learning_rates = (1f-2, 1f-3,)
-nepochs  = E .* (0.10, 0.90,) .|> Int
+nepochs  = E .* (0.20, 0.80,) .|> Int
 dir = joinpath(@__DIR__, "exp_FNO_nonlin")
 device = Lux.gpu
 
@@ -86,18 +88,18 @@ end
 # Bilinear (linear / nonlin) model
 ###
 
-if false
+if true
 
 # fixed params
 c1 = 3     # in  channel nonlin
 c2 = 1     # in  channel linear
-o  = size(__data[2]   , 1) # out channel
+o  = size(__data[2], 1) # out channel
 
 # hyper params
-w1 = 8       # width nonlin
-w2 = 8       # width linear
-wo = 8        # width project
-m = (12, 12,) # modes
+w1 = 32       # width nonlin
+w2 = 32       # width linear
+wo = 1        # width project
+m = (16, 16,) # modes
 
 split = SplitRows(1:3, 4)
 nonlin = Chain(
@@ -107,26 +109,26 @@ nonlin = Chain(
     )
 linear = Chain(
         Dense(c2, w2, use_bias = false),
-        OpConv(w2, w2, m)
     )
 bilin  = OpConvBilinear(w1, w2, wo, m)
 # bilin  = OpKernelBilinear(w1, w2, o, m) # errors
 
 project = Dense(wo, o, use_bias = false)
+project = NoOpLayer()
 
 NN = linear_nonlinear(split, nonlin, linear, bilin, project)
 
 opt = Optimisers.Adam()
 batchsize = 128
-learning_rates = (1f-3,)
+learning_rates = (1f-5,)
 nepochs  = E .* (1.00,) .|> Int
 dir = joinpath(@__DIR__, "exp_FNO_linear_nonlinear")
 device = Lux.gpu
 
-FNO_bl, _ = train_model(rng, NN, __data, data__, V, opt;
+FNO_bl, ST = train_model(rng, NN, __data, data__, V, opt;
         batchsize, learning_rates, nepochs, dir, device)
 
-end
+plt = plot_training(ST...)
 
-nothing
+end
 #