gistfile1.py

from firedrake import *
from firedrake.petsc import PETSc
m = UnitSquareMesh(1, 1)

V = FunctionSpace(m, 'CG', 1)

u = TrialFunction(V)
v = TestFunction(V)
A = assemble(u*v*dx)

mat = A.M.handle

diag = mat.getDiagonal()

# in place
diag.reciprocal()

new_mat = mat.duplicate(copy=True)

new_mat.diagonalScale(diag)

mmat = new_mat.matMult(mat)

# mmat := mat - mmat

mmat.axpy(-1, mat)

solver = LinearVariationalSolver(...)

class PC(object):
    def apply(self, pc, x, y):
        # x is input vec
        # y is output vec
        mmat.mult(x, y)

pc = PC()

solver.snes.ksp.pc = pc

solver.solve()