from petsc4py import PETSc
import numpy as np
from scipy.sparse.linalg import gmres, lgmres

data = np.load("data.npz")
M_np = data['M']
rhs_np = data['rhs']

M = PETSc.Mat().createDense(M_np.shape, array=M_np)
rhs = PETSc.Vec().createWithArray(rhs_np)
x = M.getVecLeft()

def monitor(ksp, its, rnorm):
    print(its, rnorm)

ksp = PETSc.KSP().create()
ksp.setType("lgmres")
ksp.setTolerances(1e-10, 1e-50)
ksp.setMonitor(monitor)
ksp.setOperators(M, None)
pc = ksp.getPC()
pc.setType("none")

ksp.solve(rhs, x)
ksp.view()
print("res:", np.linalg.norm(M_np.dot(x.array) - rhs_np))