Matrix exponentials

import numpy as np

The function below calculates the matrix exponential $\exp(A)=\sum_{k=0}^{\infty}A^k/k!$.

def matrix_exp(A):
    """Compute the matrix exponential of A using the Taylor series method.
    
    Parameters:
        A : a square matrix, of shape n x n

    Returns:
        The matrix exponential of A, i.e. the sum of A^k / k! for k = 0, 1, 2, ...
    
    """
    n = A.shape[0] # The number of rows in A.  
    # We do not check directly that A is square, but the next steps will generate
    # errors if it is not.

    # In the loop below, T will be A^k / k!.  Here k is implicitly zero,
    # so T is the identity matrix of the same size as A.
    T = np.eye(n) 
    # In the loop below, S will be the sum of A^j / j! for j = 0, 1, 2, ..., k.
    # Here k is implicitly zero, so S is the identity matrix of the same size as A.
    S = np.eye(n)
    # We will stop the loop when the largest entry in T is less than epsilon.
    epsilon = 1e-10
    # The main stopping criterion is that the largest entry in T should be less 
    # than epsilon, but we also stop if we have computed 100 terms in the series.
    for k in range(1, 100):
        T = T @ A / k # Update the value of T to A^k / k!
        S = S + T # Update the value of S by adding one more term
        if np.max(np.abs(T)) < epsilon:
            break
    return S

Here is an example where $\exp(A)$ has a simple pattern.

n = 5
A = np.zeros((n, n))
for i in range(n-1):
    A[i, i+1] = 1
B = matrix_exp(A)
print("A")
print(A)
print("")
print("exp(A)")
print(B)

Here is a random example.

A = np.random.random(16).reshape((4,4))
B = matrix_exp(A)
print("A")
print(A)
print("")
print("exp(A)")
print(B)

It should work out that $\exp(A)\exp(-A)=I$. We can use this to check the correctness of our function: we put $E=\exp(A)\exp(-A)-I$ (which should be zero) and then use linalg.norm() to measure the size of $E$.

n = 4
A = np.random.random(n*n).reshape((n,n))
print("A")
print(A)
print("")
E = matrix_exp(A) @ matrix_exp(-A) - np.eye(n)
print("Size of error: " + str(np.linalg.norm(E)))