Find the diagonal of a linear operator in scipy

Question

I have a function, A(x), that I have turned into a linear operator using scipy.sparse.linalg.LinearOperator:

A_operator = scipy.sparse.linalg.LinearOperator((ny, nx), matvec=A).

I have an Ax=b problem where I am using the conjugate gradient method (scipy.sparse.linalg.cg) to solve, however it is diverging and so I think my function, A, may be ill-conditioned.

I need to pre-condition by taking the inverse of the diagonal of A.

How do you easily find the diagonal of a linear operator or function? I can only find functions that calculate the diagonal of a matrix, and my problem is too large to explicitly formulate A as a matrix.

Answer 1

I calculated this by hand to find the diagonal of the function if it were a matrix.