add matrix formulation

UB4/UB4.py

@@ -98,6 +98,18 @@ def IM(y0, t, dt, lamda = -1):
     return y
 
 
+## Matrix formulation
+def IE_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.linalg.inv(np.identity(A.shape[0]) - dt * A), y[i-1])
+    
+    return y
+
 # Ex2: dy/dt = lambda * y, y(0) = 1, lambda = -1
 y0 = np.array(1)
 lamda = -1
@@ -118,10 +130,11 @@ 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)")
-    time = np.arange(int(t[1] // dt[i]))
+    # 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)")
+    time = np.arange(int(t[0] // dt[i]))
     plt.plot(exact_solution(time), '3', label="exact solution (ES)")
     plt.legend()
     plt.title(f't={t[0]}, dt={dt[i]}')