How to solve a quadratic matrix equation in Python?

Question

I am trying to solve a quadratic equation in Python that is a matrix equation. I am looking for a matrix Ax that is 2x2 such that it will satisfy

M_1 Ax^2 - M_2 A_x - M_3 = 0

where M_1, M_2 and M_3 are known 2x2 numpy arrays of constants and Ax^2 is the matrix multiplication of Ax by itself.

My code so far:

import numpy as np
from scipy.optimize import fsolve

# Calibration
gamma = 0.36
beta = 0.98
b_bar = 0.1
psi = 0.001
g = 0.66
rho_x = 0.11
rho_z = 0.95
c_star = 1 + b_bar * np.exp(-g) * (1 - beta * np.exp(-g * (gamma - 1)))

# Matrices
M1 = np.array([[gamma, 0], [0,0]])
M2 = np.array([[gamma, psi * b_bar/(beta * np.exp(-gamma * g))],
              [c_star, b_bar * (beta * np.exp(-gamma * g) - psi * b_bar)]])
M3 = np.array([[0,0],
              [0, - b_bar * np.exp(-g)]])


fsolve(lambda Ax:  M1 @ Ax @ Ax - M2 @ Ax - M3, np.eye(2))

I get the following error:

ValueError: matmul: Input operand 1 has a mismatch in its core dimension 0, with gufunc signature (n?,k),(k,m?)->(n?,m?) (size 4 is different from 2)

Evaluating the function out of fsolve seems to work fine. I don't know what is going on. What am I missing? What is the best way to solve this system in Python? Thanks!

question from:https://stackoverflow.com/questions/65909004/how-to-solve-a-quadratic-matrix-equation-in-python

Answer 1

"You need to pass in 1D vectors and sufficient data points to solve." – Mad Physicist

This would require some higher form of magic to do in a lambda expression. It is much more straightforward to implement such a wrapper in an explicit function definition

def fun(x):
    Ax = x.reshape([2,2]) 
    R = M1 @ Ax @ Ax - M2 @ Ax - M3
    return R.flatten() 

fsolve(fun, np.eye(2).flatten())