Merge pull request #826 from ayushinav/dgm

reducing test times for DGM

Author: ChrisRackauckas

docs/src/tutorials/dgm.md

@@ -32,7 +32,10 @@ $$
 $$
 
 defined over
-$$ t \in [0, 1], x \in [-1, 1] $$
+
+$$ 
+t \in [0, 1], x \in [-1, 1] 
+$$
 
 with boundary conditions
 ```math
@@ -44,7 +47,7 @@ u(t, 1) & = 0
 ```
 
 ### Copy- Pasteable code
-```julia
+```@example dgm
 using NeuralPDE
 using ModelingToolkit, Optimization, OptimizationOptimisers
 import Lux: tanh, identity
@@ -55,15 +58,15 @@ using MethodOfLines, OrdinaryDiffEq
 @parameters x t
 @variables u(..)
 
-Dt= Differential(t)
-Dx= Differential(x)
-Dxx= Dx^2
-α = 0.05;
+Dt = Differential(t)
+Dx = Differential(x)
+Dxx = Dx^2
+α = 0.05
 # Burger's equation
-eq= Dt(u(t,x)) + u(t,x) * Dx(u(t,x)) - α * Dxx(u(t,x)) ~ 0 
+eq = Dt(u(t,x)) + u(t,x) * Dx(u(t,x)) - α * Dxx(u(t,x)) ~ 0 
 
 # boundary conditions
-bcs= [
+bcs = [
     u(0.0, x) ~ - sin(π*x),
     u(t, -1.0) ~ 0.0,
     u(t, 1.0) ~ 0.0
@@ -72,7 +75,7 @@ bcs= [
 domains = [t ∈ Interval(0.0, 1.0), x ∈ Interval(-1.0, 1.0)]
 
 # MethodOfLines, for FD solution
-dx= 0.01
+dx = 0.01
 order = 2
 discretization = MOLFiniteDifference([x => dx], t, saveat = 0.01)
 @named pde_system = PDESystem(eq, bcs, domains, [t, x], [u(t,x)])
@@ -84,25 +87,27 @@ xs = sol[x]
 u_MOL = sol[u(t,x)]
 
 # NeuralPDE, using Deep Galerkin Method
-strategy = QuasiRandomTraining(4_000, minibatch= 500);
-discretization= DeepGalerkin(2, 1, 50, 5, tanh, tanh, identity, strategy);
-@named pde_system = PDESystem(eq, bcs, domains, [t, x], [u(t,x)]);
-prob = discretize(pde_system, discretization);
-global iter = 0;
+strategy = QuasiRandomTraining(256, minibatch= 32)
+discretization = DeepGalerkin(2, 1, 50, 5, tanh, tanh, identity, strategy)
+@named pde_system = PDESystem(eq, bcs, domains, [t, x], [u(t,x)])
+prob = discretize(pde_system, discretization)
+global iter = 0
 callback = function (p, l)
-    global iter += 1;
+    global iter += 1
     if iter%20 == 0
         println("$iter => $l")
     end
     return false
 end
 
-res = Optimization.solve(prob, Adam(0.01); callback = callback, maxiters = 300);
-phi = discretization.phi;
+res = Optimization.solve(prob, Adam(0.1); callback = callback, maxiters = 100)
+prob = remake(prob, u0 = res.u)
+res = Optimization.solve(prob, Adam(0.01); callback = callback, maxiters = 500)
+phi = discretization.phi
 
 u_predict= [first(phi([t, x], res.minimizer)) for t in ts, x in xs]
 
-diff_u = abs.(u_predict .- u_MOL);
+diff_u = abs.(u_predict .- u_MOL)
 
 using Plots
 p1 = plot(tgrid, xgrid, u_MOL', linetype = :contourf, title = "FD");

test/dgm_test.jl

@@ -19,8 +19,8 @@ import Lux: tanh, identity
     # Space and time domains
     domains = [x ∈ Interval(0.0, 1.0), y ∈ Interval(0.0, 1.0)]
 
-    strategy = QuasiRandomTraining(4_000, minibatch= 500);
-    discretization= DeepGalerkin(2, 1, 30, 3, tanh, tanh, identity, strategy);
+    strategy = QuasiRandomTraining(256, minibatch= 32);
+    discretization= DeepGalerkin(2, 1, 20, 3, tanh, tanh, identity, strategy);
 
     @named pde_system = PDESystem(eq, bcs, domains, [x, y], [u(x, y)])
     prob = discretize(pde_system, discretization)
@@ -71,8 +71,8 @@ end
 
     domains = [t ∈ Interval(0.0, T), x ∈ Interval(0.0, S * S_multiplier)]
 
-    strategy = QuasiRandomTraining(4_000, minibatch= 500);
-    discretization= DeepGalerkin(2, 1, 30, 3, tanh, tanh, identity, strategy);
+    strategy = QuasiRandomTraining(128, minibatch= 32);
+    discretization= DeepGalerkin(2, 1, 40, 3, tanh, tanh, identity, strategy);
 
     @named pde_system = PDESystem(eq, bcs, domains, [t, x], [g(t,x)])
     prob = discretize(pde_system, discretization)
@@ -86,9 +86,9 @@ end
         return false
     end
 
-    res = Optimization.solve(prob, Adam(0.01); callback = callback, maxiters = 300)
+    res = Optimization.solve(prob, Adam(0.1); callback = callback, maxiters = 100)
     prob = remake(prob, u0 = res.u)
-    res = Optimization.solve(prob, Adam(0.001); callback = callback, maxiters = 300)
+    res = Optimization.solve(prob, Adam(0.01); callback = callback, maxiters = 500)
     phi = discretization.phi
 
     function analytical_soln(t, x, K, σ, T)
@@ -138,7 +138,7 @@ end
     u_MOL = sol[u(t,x)]
 
     # NeuralPDE
-    strategy = QuasiRandomTraining(4_000, minibatch= 500);
+    strategy = QuasiRandomTraining(256, minibatch= 32);
     discretization= DeepGalerkin(2, 1, 50, 5, tanh, tanh, identity, strategy);
     @named pde_system = PDESystem(eq, bcs, domains, [t, x], [u(t,x)]);
     prob = discretize(pde_system, discretization);
@@ -151,7 +151,9 @@ end
         return false
     end
 
-    res = Optimization.solve(prob, Adam(0.01); callback = callback, maxiters = 300);
+    res = Optimization.solve(prob, Adam(0.01); callback = callback, maxiters = 200);
+    prob = remake(prob, u0 = res.u);
+    res = Optimization.solve(prob, Adam(0.001); callback = callback, maxiters = 100);
     phi = discretization.phi;
 
     u_predict= [first(phi([t, x], res.u)) for t in ts, x in xs]