diff --git a/examples/diffusion_fourier/exp_bilinear_scale/bilin.jl b/examples/diffusion_fourier/exp_bilinear_scale/bilin.jl index 6362cd2f..80d93b6d 100644 --- a/examples/diffusion_fourier/exp_bilinear_scale/bilin.jl +++ b/examples/diffusion_fourier/exp_bilinear_scale/bilin.jl @@ -39,15 +39,15 @@ include("../datagen.jl") N = 128 # problem size K1 = 32 # ν-samples K2 = 32 # f-samples -E = 200 # epochs +E = 100 # epochs rng = Random.default_rng() Random.seed!(rng, 199) # datagen _V, _data, _, _ = datagen1D(rng, N, K1, K2) # train -# V_, data_, _, _ = datagen1D(rng, N, K1, K2) # train -V_, data_, _, _ = datagen1D(rng, N, K1, K2; mode = :test) # test +V_, data_, _, _ = datagen1D(rng, N, K1, K2) # train +# V_, data_, _, _ = datagen1D(rng, N, K1, K2; mode = :test) # test __data = combine_data1D(_data) data__ = combine_data1D(data_) @@ -100,7 +100,7 @@ w2 = 16 # width linear wo = 8 # width project m = (32,) # modes -split = SplitRows(1:2, 3) +split = SplitRows(1:2, 3) nonlin = Chain(PermutedBatchNorm(c1, 3), Dense(c1, w1, tanh), OpKernel(w1, w1, m, tanh)) linear = Dense(c2, w2, use_bias = false) bilin = OpConvBilinear(w1, w2, o, m) @@ -109,13 +109,15 @@ bilin = OpConvBilinear(w1, w2, o, m) NN = linear_nonlinear(split, nonlin, linear, bilin) opt = Optimisers.Adam() -batchsize = size(__data[1])[end] +batchsize = 256 # size(__data[1])[end] # 1024 learning_rates = (1f-3,) -nepochs = E .* (1.00,) .|> Int +nepochs = E .* (1.00,) .|> Int +# learning_rates = (1f-3, 5f-4, 2.5f-4, 1.25f-4,) +# nepochs = E .* (0.25, 0.25, 0.25, 0.25,) .|> Int dir = joinpath(@__DIR__, "exp_FNO_linear_nonlinear") -device = Lux.gpu +device = Lux.cpu -model, _ = train_model(rng, NN, __data, data__, _V, opt; +model, ST = train_model(rng, NN, __data, data__, _V, opt; batchsize, learning_rates, nepochs, dir, cbstep = 1, device) end diff --git a/examples/diffusion_fourier/exp_bilinear_scale/exp_FNO_linear_nonlinear/plt_r2_test.png b/examples/diffusion_fourier/exp_bilinear_scale/exp_FNO_linear_nonlinear/plt_r2_test.png index cb64ca43..0f2cd0e1 100644 Binary files a/examples/diffusion_fourier/exp_bilinear_scale/exp_FNO_linear_nonlinear/plt_r2_test.png and b/examples/diffusion_fourier/exp_bilinear_scale/exp_FNO_linear_nonlinear/plt_r2_test.png differ diff --git a/examples/diffusion_fourier/exp_bilinear_scale/exp_FNO_linear_nonlinear/plt_r2_train.png b/examples/diffusion_fourier/exp_bilinear_scale/exp_FNO_linear_nonlinear/plt_r2_train.png index e9fe0f23..70d3b74a 100644 Binary files a/examples/diffusion_fourier/exp_bilinear_scale/exp_FNO_linear_nonlinear/plt_r2_train.png and b/examples/diffusion_fourier/exp_bilinear_scale/exp_FNO_linear_nonlinear/plt_r2_train.png differ diff --git a/examples/diffusion_fourier/exp_bilinear_scale/exp_FNO_linear_nonlinear/plt_training.png b/examples/diffusion_fourier/exp_bilinear_scale/exp_FNO_linear_nonlinear/plt_training.png index f0cdd9eb..8a0413a6 100644 Binary files a/examples/diffusion_fourier/exp_bilinear_scale/exp_FNO_linear_nonlinear/plt_training.png and b/examples/diffusion_fourier/exp_bilinear_scale/exp_FNO_linear_nonlinear/plt_training.png differ diff --git a/examples/diffusion_fourier/exp_bilinear_scale/exp_FNO_linear_nonlinear/plt_traj_test.png b/examples/diffusion_fourier/exp_bilinear_scale/exp_FNO_linear_nonlinear/plt_traj_test.png index 09856463..08fb4c1e 100644 Binary files a/examples/diffusion_fourier/exp_bilinear_scale/exp_FNO_linear_nonlinear/plt_traj_test.png and b/examples/diffusion_fourier/exp_bilinear_scale/exp_FNO_linear_nonlinear/plt_traj_test.png differ diff --git a/examples/diffusion_fourier/exp_bilinear_scale/exp_FNO_linear_nonlinear/plt_traj_train.png b/examples/diffusion_fourier/exp_bilinear_scale/exp_FNO_linear_nonlinear/plt_traj_train.png index cda1bcd6..c426c9ab 100644 Binary files a/examples/diffusion_fourier/exp_bilinear_scale/exp_FNO_linear_nonlinear/plt_traj_train.png and b/examples/diffusion_fourier/exp_bilinear_scale/exp_FNO_linear_nonlinear/plt_traj_train.png differ diff --git a/examples/diffusion_fourier/exp_bilinear_scale/exp_FNO_linear_nonlinear/statistics.txt b/examples/diffusion_fourier/exp_bilinear_scale/exp_FNO_linear_nonlinear/statistics.txt index 0859f38c..9444980f 100644 --- a/examples/diffusion_fourier/exp_bilinear_scale/exp_FNO_linear_nonlinear/statistics.txt +++ b/examples/diffusion_fourier/exp_bilinear_scale/exp_FNO_linear_nonlinear/statistics.txt @@ -1,17 +1,19 @@ - TRAIN LOSS: 0.0259869 TEST LOSS: 281.77728 + TRAIN LOSS: 0.02262737 TEST LOSS: 0.03570642 #======================# TRAIN STATS -R² score: 0.9733658 -mean SQR error: 0.02598686 -mean ABS error: 0.12003216 -max ABS error: 1.0625379 +R² score: 0.9730648 +MSE (mean SQR error): 0.02510134 +RMSE (root mean SQR error): 0.15843402 +MAE (mean ABS error): 0.11758655 +maxAE (max ABS error) 1.0394883 #======================# #======================# TEST STATS -R² score: 0.9745712 -mean SQR error: 281.7768 -mean ABS error: 10.313357 -max ABS error: 168.86465 +R² score: 0.97130257 +MSE (mean SQR error): 0.03251882 +RMSE (root mean SQR error): 0.18032974 +MAE (mean ABS error): 0.13228929 +maxAE (max ABS error) 1.1780059 #======================# diff --git a/examples/pdebench/burgers1d.jl b/examples/pdebench/burgers1d.jl index 078bdeda..73de7ec0 100644 --- a/examples/pdebench/burgers1d.jl +++ b/examples/pdebench/burgers1d.jl @@ -45,7 +45,7 @@ K_ = 64 # get data dir = @__DIR__ -filename = joinpath(dir, "1D_Burgers_Sols_Nu0.01.hdf5") +filename = joinpath(dir, "1D_Burgers_Sols_Nu0.001.hdf5") include(joinpath(dir, "pdebench.jl")) _data, data_ = burgers1D(filename, _K, K_, rng) @@ -83,7 +83,7 @@ nepochs = E .* (0.25, 0.25, 0.25, 0.25) .|> Int # learning_rates = (1f-3, 5f-4, 2.5f-4, 1.25f-4) # nepochs = E .* (0.25, 0.25, 0.25, 0.25) .|> Int -dir = joinpath(@__DIR__, "model_burgers1D") +dir = joinpath(@__DIR__, "dump") device = Lux.cpu model, ST = train_model(rng, NN, _data, data_, V, opt; diff --git a/examples/pdebench/model_burgers1D/plt_r2_test.png b/examples/pdebench/model_burgers1D/plt_r2_test.png deleted file mode 100644 index 7b6d1b46..00000000 Binary files a/examples/pdebench/model_burgers1D/plt_r2_test.png and /dev/null differ diff --git a/examples/pdebench/model_burgers1D/plt_r2_train.png b/examples/pdebench/model_burgers1D/plt_r2_train.png deleted file mode 100644 index 42607206..00000000 Binary files a/examples/pdebench/model_burgers1D/plt_r2_train.png and /dev/null differ diff --git a/examples/pdebench/model_burgers1D/plt_training.png b/examples/pdebench/model_burgers1D/plt_training.png deleted file mode 100644 index 848454ae..00000000 Binary files a/examples/pdebench/model_burgers1D/plt_training.png and /dev/null differ diff --git a/examples/pdebench/model_burgers1D/plt_traj_test.png b/examples/pdebench/model_burgers1D/plt_traj_test.png deleted file mode 100644 index 4aebc7b8..00000000 Binary files a/examples/pdebench/model_burgers1D/plt_traj_test.png and /dev/null differ diff --git a/examples/pdebench/model_burgers1D/plt_traj_train.png b/examples/pdebench/model_burgers1D/plt_traj_train.png deleted file mode 100644 index 9b816bd3..00000000 Binary files a/examples/pdebench/model_burgers1D/plt_traj_train.png and /dev/null differ diff --git a/examples/pdebench/model_burgers1D/statistics.txt b/examples/pdebench/model_burgers1D/statistics.txt deleted file mode 100644 index 08ca6ee1..00000000 --- a/examples/pdebench/model_burgers1D/statistics.txt +++ /dev/null @@ -1,19 +0,0 @@ - TRAIN LOSS: 0.00030949 TEST LOSS: 0.00574172 -#======================# -TRAIN STATS -R² score: 0.9982961 -MSE (mean SQR error): 0.00057452 -RMSE (root mean SQR error): 0.02396922 -MAE (mean ABS error): 0.01535009 -maxAE (max ABS error) 0.30036545 - -#======================# -#======================# -TEST STATS -R² score: 0.9795042 -MSE (mean SQR error): 0.00701486 -RMSE (root mean SQR error): 0.08375475 -MAE (mean ABS error): 0.05496668 -maxAE (max ABS error) 0.48964748 - -#======================# diff --git a/src/operator.jl b/src/operator.jl index db159260..3fa2371c 100644 --- a/src/operator.jl +++ b/src/operator.jl @@ -22,9 +22,12 @@ function OpKernel(ch_in::Int, ch_out::Int, modes::NTuple{D, Int}, conv = OpConv(ch_in, ch_out, modes; transform, init) loc = Dense(ch_in, ch_out; init_weight = init, use_bias) + # Parallel(+, loc, conv), # x -> lox(x) + conv(x) + Chain( - Lux.Parallel(+, loc, conv), - Lux.WrappedFunction(activation), + BranchLayer(loc, conv), # x -> (loc(x), conv(x)) + WrappedFunction(sum), # (x1, x2) -> x1 + x2 + WrappedFunction(activation), # x -> act(x) ) end @@ -37,12 +40,15 @@ function OpKernelBilinear(ch_in1::Int, ch_in2::Int, ch_out::Int, activation = fastify(activation) + null = NoOpLayer() conv = OpConvBilinear(ch_in1, ch_in2, ch_out, modes; transform, init) loc = Bilinear((ch_in1, ch_in2) => ch_out; init_weight = init, use_bias = false) + # Parallel(+, null, null), Chain( - Lux.Parallel(.+, loc, conv), - Lux.WrappedFunction(activation), + BranchLayer(loc, conv), # x -> (loc(x), conv(x)) + WrappedFunction(sum), # (x1, x2) -> x1 + x2 + WrappedFunction(activation), # x -> act(x) ) end @@ -66,9 +72,9 @@ x2 → linear → y2 ↗ function linear_nonlinear(split, nonlin, linear, bilinear, project = NoOpLayer()) Chain( - split, - Parallel(nothing, nonlin, linear), - bilinear, + split, # x -> (x1, x2) + Parallel(nothing, nonlin, linear), # (x1, x2) -> (f(x1), g(x2)) + bilinear, # (f(x1), g(x2)) -> y project, ) end