refactor

Author: truong-dang-gft

src/fdm_heat_eq_explicit_1D.py

@@ -1,6 +1,7 @@
 import math
 import numpy
-import matplotlib.pyplot as plt
+
+from src.utils import visu_1D
 
 """
 Original version was in Matlab. See [1] for further details.
@@ -9,23 +10,6 @@
 """
 
 
-def plot_solution(x, wn):
-    """
-
-    :param x:
-    :param wn:
-    :return:
-    """
-    xx = numpy.linspace(0, 1, 1000)
-    plt.plot(xx, numpy.exp(-time / 4 * math.pi ** 2) * numpy.sin(2 * math.pi * xx), label="exact solution")
-    plt.plot(x, wn, label="approximate solution")
-    plt.xlabel("x")
-    plt.ylabel("u")
-    plt.title("Heat eq")
-    plt.legend()
-    plt.show()
-
-
 if __name__ == "__main__":
     # FDM, forward difference, explicit scheme for the heat equation:
     # INPUT
@@ -44,5 +28,6 @@ def plot_solution(x, wn):
         wn[1:M + 1] = wn[1:M + 1] + mu * (wn[2:M + 2] - 2 * wn[1:M + 1] + wn[0:M])
         time += Dt
 
-    plot_solution(x, wn)
+    w = numpy.exp(-time / 4 * math.pi ** 2) * numpy.sin(2 * math.pi * x)
+    visu_1D.plot_solution(x, wn, w)
 

src/utils/__init__.py

@@ -0,0 +1,19 @@
+import matplotlib.pyplot as plt
+
+
+def plot_solution(grid, approximate_sol, exact_sol, title="Heat eq"):
+    """
+
+    :param grid:
+    :param approximate_sol:
+    :param exact_sol:
+    :param title:
+    :return:
+    """
+    plt.plot(grid, exact_sol, label="exact solution")
+    plt.plot(grid, approximate_sol, label="approximate solution")
+    plt.xlabel("x")
+    plt.ylabel("u")
+    plt.title(title)
+    plt.legend()
+    plt.show()