add matrix formulation

UB4/UB4.py

@@ -110,6 +110,29 @@ def IE_mat(y0, t, dt, A):
     
     return y
 
+def EE_mat(y0, t, dt, A):
+    assert A.shape[0] == A.shape[1]
+    no_of_steps = int(t // dt)
+    y = np.zeros((no_of_steps, *y0.shape))
+    y[0] = y0
+
+    for i in range(1, no_of_steps):
+        y[i] = np.dot(np.identity(A.shape[0]) + dt * A, y[i-1])
+    
+    return y
+
+def IM_mat(y0, t, dt, A):
+    assert A.shape[0] == A.shape[1]
+    no_of_steps = int(t // dt)
+    y = np.zeros((no_of_steps, *y0.shape))
+    y[0] = y0
+
+    for i in range(1, no_of_steps):
+        propagation_mat = np.dot(np.linalg.inv(np.identity(A.shape[0]) - dt/2 * A), np.identity(A.shape[0]) + dt/2 * A)
+        y[i] = np.dot(propagation_mat, y[i-1])
+    
+    return y
+
 # Ex2: dy/dt = lambda * y, y(0) = 1, lambda = -1
 y0 = np.array(1)
 lamda = -1
@@ -130,10 +153,9 @@ for i in range(len(dt)):
     # plt.plot(implicit_euler(f, y0, t[0], dt[i]), '--', label = "implicit euler (IE)")
     # plt.plot(explicit_euler(f, y0, t[0], dt[i]), 'x', label = "explicit euler (EE)")
     # plt.plot(implicit_midpoint(f, y0, t[0], dt[i]), '>', label = "implicit midpoint (IM)")
-    # plt.plot(IE(y0, t[1], dt[i]), '--', label = "implicit euler (IE)")
-    # plt.plot(EE(y0, t[1], dt[i]), 'x', label = "explicit euler (EE)")
-    # plt.plot(IM(y0, t[1], dt[i]), '>', label = "implicit midpoint (IM)")
-    plt.plot(IE_mat(y0, t[0], dt[i], np.array(lamda).reshape(1,1)), '--', label = "implicit euler (IE)")
+    plt.plot(IE_mat(y0, t[0], dt[i], A = np.array(lamda).reshape(1,1)), '--', label = "implicit euler (IE)")
+    plt.plot(EE_mat(y0, t[0], dt[i], A = np.array(lamda).reshape(1,1)), 'x', label = "explicit euler (EE)")
+    plt.plot(IM_mat(y0, t[0], dt[i], A = np.array(lamda).reshape(1,1)), '>', label = "implicit midpoint (IM)")
     time = np.arange(int(t[0] // dt[i]))
     plt.plot(exact_solution(time), '3', label="exact solution (ES)")
     plt.legend()