From 9d3f9f0d02d677660824b5e291e64a022387bc42 Mon Sep 17 00:00:00 2001
From: Sathvik Bhagavan <sathvik.bhagavan@juliahub.com>
Date: Thu, 14 Mar 2024 04:47:30 +0000
Subject: [PATCH] docs: fix dgm example

---
 docs/src/tutorials/dgm.md | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/docs/src/tutorials/dgm.md b/docs/src/tutorials/dgm.md
index 61d32cbe33..1917df8f43 100644
--- a/docs/src/tutorials/dgm.md
+++ b/docs/src/tutorials/dgm.md
@@ -7,6 +7,7 @@ Deep Galerkin Method is a meshless deep learning algorithm to solve high dimensi
 In the following example, we demonstrate computing the loss function using Quasi-Random Sampling, a sampling technique that uses quasi-Monte Carlo sampling to generate low discrepancy random sequences in high dimensional spaces.
 
 ### Algorithm
+
 The authors of DGM suggest a network composed of LSTM-type layers that works well for most of the parabolic and quasi-parabolic PDEs.
 
 ```math
@@ -47,13 +48,15 @@ u(t, 1) & = 0
 ```
 
 ### Copy- Pasteable code
+
 ```@example dgm
 using NeuralPDE
 using ModelingToolkit, Optimization, OptimizationOptimisers
-import Lux: tanh, identity
+using Lux: tanh, identity
 using Distributions
-import ModelingToolkit: Interval, infimum, supremum
+using ModelingToolkit: Interval, infimum, supremum
 using MethodOfLines, OrdinaryDiffEq
+using Plots
 
 @parameters x t
 @variables u(..)
@@ -108,8 +111,9 @@ 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)
+tgrid = 0.0:0.01:1.0
+xgrid = -1.0:0.01:1.0
 
-using Plots
 p1 = plot(tgrid, xgrid, u_MOL', linetype = :contourf, title = "FD");
 p2 = plot(tgrid, xgrid, u_predict', linetype = :contourf, title = "predict");
 p3 = plot(tgrid, xgrid, diff_u', linetype = :contourf, title = "error");