Numpy matrix power/exponent with modulo?

In order to prevent overflow, you can use the fact that you get the same result if you first take the modulo of each of your input numbers; in fact:

(M**k) mod p = ([M mod p]**k) mod p,

for a matrix M. This comes from the following two fundamental identities, which are valid for integers x and y (and a positive power p):

(x+y) mod p = ([x mod p]+[y mod p]) mod p  # All additions can be done on numbers *modulo p*
(x*y) mod p = ([x mod p]*[y mod p]) mod p  # All multiplications can be done on numbers *modulo p*

The same identities hold for matrices as well, since matrix addition and multiplication can be expressed through scalar addition and multiplication. With this, you only exponentiate small numbers (n mod p is generally much smaller than n) and are much less likely to get overflows. In NumPy, you would therefore simply do

((arr % p)**k) % p

in order to get (arr**k) mod p.

If this is still not enough (i.e., if there is a risk that [n mod p]**k causes overflow despite n mod p being small), you can break up the exponentiation into multiple exponentiations. The fundamental identities above yield

(n**[a+b]) mod p = ([{n mod p}**a mod p] * [{n mod p}**b mod p]) mod p

and

(n**[a*b]) mod p = ([n mod p]**a mod p)**b mod p.

Thus, you can break up power k as a+b+… or a*b*… or any combination thereof. The identities above allow you to perform only exponentiations of small numbers by small numbers, which greatly lowers the risk of integer overflows.

Using the implementation from Numpy:

https://github.com/numpy/numpy/blob/master/numpy/matrixlib/defmatrix.py#L98

I adapted it by adding a modulo term. HOWEVER, there is a bug, in that if an overflow occurs, no OverflowError or any other sort of exception is raised. From that point on, the solution will be wrong. There is a bug report here.

Here is the code. Use with care:

from numpy.core.numeric import concatenate, isscalar, binary_repr, identity, asanyarray, dot
from numpy.core.numerictypes import issubdtype    
def matrix_power(M, n, mod_val):
    # Implementation shadows numpy's matrix_power, but with modulo included
    M = asanyarray(M)
    if len(M.shape) != 2 or M.shape[0] != M.shape[1]:
        raise ValueError("input  must be a square array")
    if not issubdtype(type(n), int):
        raise TypeError("exponent must be an integer")

    from numpy.linalg import inv

    if n==0:
        M = M.copy()
        M[:] = identity(M.shape[0])
        return M
    elif n<0:
        M = inv(M)
        n *= -1

    result = M % mod_val
    if n <= 3:
        for _ in range(n-1):
            result = dot(result, M) % mod_val
        return result

    # binary decompositon to reduce the number of matrix
    # multiplications for n > 3
    beta = binary_repr(n)
    Z, q, t = M, 0, len(beta)
    while beta[t-q-1] == '0':
        Z = dot(Z, Z) % mod_val
        q += 1
    result = Z
    for k in range(q+1, t):
        Z = dot(Z, Z) % mod_val
        if beta[t-k-1] == '1':
            result = dot(result, Z) % mod_val
    return result % mod_val