cg.py

import numpy as np

def conjugate_gradient(A, b, iters, eps):
  "Approximate A^{-1} b for positive definite A using conjugate gradient."
  v = np.zeros_like(b)           # initial approximate solution
  r = b                          # initial residual vector
  rho = rho_prev = np.dot(r, r)  # initial residual norm squared
  tol = eps * np.sqrt(rho)       # residual norm stop criterion
  for k in range(iters):
    if np.sqrt(rho) < tol: break
    w = A(b)
    alpha = rho / np.dot(b, w)
    v = v + alpha * b
    r = r - alpha * w
    rho, rho_prev = np.dot(r, r), rho
    b = r + (rho / rho_prev) * b
  return v