burg

examples/pdebench/advect1d.jl

@@ -60,7 +60,7 @@ m = (16,) # modes
 c = size(_data[1], 1) # in  channels
 o = size(_data[2], 1) # out channels
 
-NN = Lux.Chain(
+NN = Lux.Chain( # todo - double check parameter size
     Dense(c, w, tanh),
     OpKernel(w, w, m, tanh),
     OpKernel(w, w, m, tanh),

examples/pdebench/burgers_bilinear.jl

@@ -0,0 +1,109 @@
+#
+"""
+Learn solution to diffusion equation
+
+    -∇⋅ν∇u = f
+
+for constant ν₀, and variable f
+
+test bed for Fourier Neural Operator experiments where
+forcing is learned separately.
+"""
+
+using GeometryLearning
+
+# PDE stack
+using LinearAlgebra, FourierSpaces
+
+# ML stack
+using Lux, Random, Optimisers
+
+# vis/analysis, serialization
+using Plots, BSON
+
+# accelerator
+using CUDA, KernelAbstractions
+CUDA.allowscalar(false)
+import Lux: cpu, gpu, relu
+
+# misc
+using Tullio, Zygote
+
+using FFTW, LinearAlgebra
+BLAS.set_num_threads(4)
+FFTW.set_num_threads(8)
+
+rng = Random.default_rng()
+Random.seed!(rng, 983254)
+
+N = 1024
+E = 100
+
+# trajectories
+_K = 512
+K_ = 64
+
+# get data
+dir = @__DIR__
+filename = joinpath(dir, "1D_Burgers_Sols_Nu0.001.hdf5")
+include(joinpath(dir, "pdebench.jl"))
+_data, data_ = burgers1D(filename, _K, K_, rng)
+
+_data = (reshape(_data[1][2, :, :], 1, N, :), _data[2]) # omit x-coordinate
+data_ = (reshape(data_[1][2, :, :], 1, N, :), data_[2])
+
+V = FourierSpace(N)
+
+###
+# Bilin FNO model
+###
+
+#============================#
+w = 32    # width
+l = 8
+m = (128,) # modes
+c = size(_data[1], 1) # in  channels
+o = size(_data[2], 1) # out channels
+
+root = Chain(
+    Dense(c, w, relu),
+    OpKernel(w, w, m, relu),
+    OpKernel(w, w, m, relu),
+)
+branch = Chain(
+    OpKernel(w, w, m, relu), # use_bias = true
+    OpKernel(w, w, m, relu),
+)
+fuse = OpConvBilinear(w, w, l, m)
+
+project = Chain(
+    Dense(l, l, relu),
+    Dense(l, o),
+)
+
+NN = Chain(
+    root,
+    BranchLayer(deepcopy(branch), deepcopy(branch)),
+    fuse,
+    project,
+)
+
+#============================#
+
+opt = Optimisers.Adam()
+batchsize = 32
+
+learning_rates = (1f-2, 1f-3, 1f-4, 1f-5)
+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_nu0.001_bilinear")
+device = Lux.cpu
+
+model, ST = train_model(rng, NN, _data, data_, V, opt;
+    batchsize, learning_rates, nepochs, dir, device)
+
+plot_training(ST...)
+#

examples/pdebench/model_burgers1D_nu0.001/statistics.txt

@@ -0,0 +1,19 @@
+	 TRAIN LOSS: 0.00050868	 TEST LOSS: 0.00613378
+#======================#
+TRAIN STATS
+R² score:                   0.9975712
+MSE (mean SQR error):       0.00079163
+RMSE (root mean SQR error): 0.02813594
+MAE (mean ABS error):       0.01866633
+maxAE (max ABS error)       0.34589595
+
+#======================#
+#======================#
+TEST  STATS
+R² score:                   0.9755317
+MSE (mean SQR error):       0.00802541
+RMSE (root mean SQR error): 0.08958465
+MAE (mean ABS error):       0.06163863
+maxAE (max ABS error)       0.45869547
+
+#======================#

examples/pdebench/model_burgers1D_nu0.001_/statistics.txt

@@ -0,0 +1,19 @@
+	 TRAIN LOSS: 0.00495276	 TEST LOSS: 0.00816443
+#======================#
+TRAIN STATS
+R² score:                   0.9882567
+MSE (mean SQR error):       0.00397282
+RMSE (root mean SQR error): 0.06303029
+MAE (mean ABS error):       0.04488528
+maxAE (max ABS error)       0.37895635
+
+#======================#
+#======================#
+TEST  STATS
+R² score:                   0.9804525
+MSE (mean SQR error):       0.00676212
+RMSE (root mean SQR error): 0.0822321
+MAE (mean ABS error):       0.05994601
+maxAE (max ABS error)       0.4080286
+
+#======================#

examples/pdebench/model_burgers1D_nu0.001_bilinear/statistics.txt

@@ -0,0 +1,19 @@
+	 TRAIN LOSS: 0.00090023	 TEST LOSS: 0.00972281
+#======================#
+TRAIN STATS
+R² score:                   0.9961141
+MSE (mean SQR error):       0.00132547
+RMSE (root mean SQR error): 0.03640705
+MAE (mean ABS error):       0.02396612
+maxAE (max ABS error)       0.30083415
+
+#======================#
+#======================#
+TEST  STATS
+R² score:                   0.9676699
+MSE (mean SQR error):       0.01123419
+RMSE (root mean SQR error): 0.10599148
+MAE (mean ABS error):       0.07846318
+maxAE (max ABS error)       0.3812811
+
+#======================#

examples/pdebench/model_burgers1D_nu0.01/statistics.txt

@@ -0,0 +1,19 @@
+	 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
+
+#======================#