reducing test times for DGM

Author: ayushinav

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
@@ -84,8 +87,8 @@ 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);
+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;
@@ -97,7 +100,9 @@ callback = function (p, l)
     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.01); callback = callback, maxiters = 500)
 phi = discretization.phi;
 
 u_predict= [first(phi([t, x], res.minimizer)) for t in ts, x in xs]

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]