Use pauli_rep to compute qml.matrix() whenever possible

[Qottmann]

Use pauli_rep whenever possible to compute the matrix of a Sum operator more efficiently.

Builds on / is blocked by #5301

[sc-58850]

[trbromley]

Do we have any data on how much faster this is? Maybe we can use @obliviateandsurrender's benchmark.

[github-actions]

Hello. You may have forgotten to update the changelog!
Please edit doc/releases/changelog-dev.md with:

• A one-to-two sentence description of the change. You may include a small working example for new features.
• A link back to this PR.
• Your name (or GitHub username) in the contributors section.

[Qottmann]

Some numbers, the speedup is actually not really significant 🤔

op = qml.sum(*(0.5 * (X(i) @ X(i+1)) for i in range(10)))
op += qml.sum(*(0.7 * (Z(i) @ Y(i+1)) for i in range(10)))
op += qml.sum(*(1.5 * (Z(i) @ Y(i+3)) for i in range(10)))

%timeit a = op.matrix() # this branch
2.41 s ± 31.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit a = op.matrix() # master
2.84 s ± 168 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

[Qottmann]

qml.matrix(Hamiltonian) is fast because it is using the sparse_matrix method. I just pushed an upgrade that also enables using that for Sum and Prod instances with some considerable speed - up :)

op = qml.sum(*(0.5 * (X(i) @ X(i+1)) for i in range(10)))
op += qml.sum(*(0.7 * (Z(i) @ Y(i+1)) for i in range(10)))
op += qml.sum(*(1.5 * (Z(i) @ Y(i+3)) for i in range(10)))

%timeit a = qml.matrix(op) # this branch
109 ms ± 8.92 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit a = qml.matrix(op) # master branch
3.65 s ± 176 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

import pyscf
import numpy as np
h3_molecule = [["H", (0, 0, 0)], ["H", (0, 1.2, 1.2)], ["H", (0, 0, 2.4)]]
molecule, basis, symm, charge, spin, tol = h3_molecule, "631g", None, 0, 1, 1e-6
mol = pyscf.gto.M(atom=molecule, basis=basis, symmetry=symm, charge=charge, spin=spin)
myhf = pyscf.scf.ROHF(mol).run(verbose=0)
core_constant = np.array([mol.energy_nuc()])
# molecular integrals in AO basis
one_ao = mol.intor_symmetric("int1e_kin") + mol.intor_symmetric("int1e_nuc")
two_ao = mol.intor("int2e_sph")
# rotate to MO basis
one_mo = np.einsum("pi,pq,qj->ij", myhf.mo_coeff, one_ao, myhf.mo_coeff)
two_mo = pyscf.ao2mo.incore.full(two_ao, myhf.mo_coeff)
# physicist ordering convention
two_mo = np.swapaxes(two_mo, 1, 3)
h_ferm = qml.qchem.fermionic_observable(core_constant, one_mo, two_mo)

with qml.operation.disable_new_opmath_cm():
    H = qml.qchem.qubit_observable(h_ferm)
    print(len(H.terms()[1]), type(H))
    %timeit qml.matrix(H)

with qml.operation.enable_new_opmath_cm():
    H = qml.qchem.qubit_observable(h_ferm)
    print(len(H.terms()[1]), type(H))
    %timeit qml.matrix(H)

923 <class 'pennylane.ops.qubit.hamiltonian.Hamiltonian'>
1.15 s ± 17.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
923 <class 'pennylane.ops.op_math.sum.Sum'>
139 ms ± 914 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

[albi3ro]

qml.matrix(Hamiltonian) is fast because it is using the sparse_matrix method. I just pushed an upgrade that also enables using that for Sum and Prod instances with some considerable speed - up :)

The issue is that Operator.sparse_matrix is not differentiable, but Operator.matrix() can be, so I would say switching to use the sparse matrix instead of the dense matrix to be a feature regression. We used Hamiltonian.sparse_matrix because we never got around to adding Hamiltonian.matrix, so we had no alternative.

I do think we could leverage PauliWord._to_sparse_matrix to improve the performance of calculating a dense matrix, but we need to make sure that dense matrix continues to be fully differentaible.

[Qottmann]

I am using the pauli_rep.to_matrix() function now but getting some unexpected failures in other parts of the codebase (e.g. diagonalize_qwc_pauli_words). Will have to dig in later

Edit: The failures seem to come from #5301 This PR is currently blocked by this

[trbromley]

There's currently a CI check failing, but bringing in @maliasadi @vincentmr here as this is a performance issue related to op arithmetic. Thanks!

[vincentmr]

There's currently a CI check failing, but bringing in @maliasadi @vincentmr here as this is a performance issue related to op arithmetic. Thanks!

Maybe we should bring in @mudit2812 too, since he's the instigator of #5301 .

[albi3ro]

Would it be possible to update PauliSentence.matrix() to take advantage of the sparse speed up of the non-differentiable PauliWord.matrix(), but maintain the differentiability of the dense matrix?

I'd like to keep the dense matrix differentiable if at all possible.

[mudit2812]

Maybe we should bring in @mudit2812 too, since he's the instigator of #5301 .

I fixed all failing CI in my PR. I will update branch and let's see if CI passes. Once CI passes, I will run lightning CI against this branch as well to validate that tests pass. I'll update this PR accordingly, but if lightning tests break because of this PR, we should have a PR on the lightning side ready to be merged right after this.

[trbromley]

Also keeping @mlxd in the loop 👍

[Qottmann]

In the current implementation that uses pauli_rep.to_mat(format="csr"), indeed coefficients parametrizing an operator, for which we compute the matrix via qml.matrix(), is not differentiable right now (you can still differentiate "through it" but not coefficient inside of it).

I.e.

def f(x):
    mat = qml.matrix(x * X(0))
    return jnp.sum(mat)

x = jnp.array(0.5, dtype=complex)
f(x), jax.grad(f, holomorphic=True)(x)

(Array(1.+0.j, dtype=complex128), Array(0.+0.j, dtype=complex128))

This is why the pulse tests fail (the "true" results are computed with explicit matrices evaluated via qml.matrix(), using H(params, t).matrix() would solve that)

@albi3ro I didnt fully grasp your suggestion, could you elaborate which specific steps you would take?

[mlxd]

Is CSR what we want to use for this?

[Qottmann]

For context, this came up in switching to new opmath because some generators and qchem operators where now slower in computing qml.matrix(), as they were no longer ops.Hamiltonian instances and therefore didnt hit the op.sparse_matrix(..).toarray() method below anymore. The idea was to just allow still doing that via the pauli rep in this PR; but as @albi3ro pointed out this breaks differentiability. Afaik christina has been cooking up an alternate solution 👩‍🍳

[Qottmann]

Closing in favor of #5578 that ensures differentiability