Improve expand_matrix method

[AlbertMitjans]

Change dense matrix expansion method, which improves performance. We use the same method that is used for sparse matrices. Consequently, we have merged the two methods (they had a lot of duplicated lines of code).

The current method does the following:

• Compute subset of wire_order that contains all labels in wires + expand mat if this subset is larger than the matrix wires.
• Compute the needed permutations to change the wire order of mat to match subset_wire_order.
• We reshape the matrix to a "binary" shape to easily permute wires, we use qml.math.transpose to permute the wires and then we reshape back.
• We expand the matrix even further if needed.

[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.

[AlbertMitjans]

• Before:

• After:
  

[codecov]

Codecov Report

Merging #3064 (934d40b) into master (e79f264) will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##           master    #3064   +/-   ##
=======================================
  Coverage   99.68%   99.68%           
=======================================
  Files         273      273           
  Lines       23671    23683   +12     
=======================================
+ Hits        23597    23609   +12     
  Misses         74       74           

Impacted Files	Coverage Δ
pennylane/math/matrix_manipulation.py	100.00% <100.00%> (ø)

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

[AlbertMitjans]

Here is the script I use for benchmarks:

import functools
import random
from time import time

import matplotlib.pyplot as plt
import numpy as np

import pennylane as qml

random.seed(123)
NUM_WIRES = 20
NUM_OPERANDS = 20


def build_new_sum_op(num: int):
    summands = [qml.s_prod(0.5, qml.PauliX(i)) for i in num]
    return qml.op_sum(*summands)


def build_old_sum_op(num: int):
    coeffs = [0.5 for _ in num]
    ops = [qml.PauliX(i) for i in num]
    return qml.Hamiltonian(coeffs=coeffs, observables=ops)


def build_new_prod_op(num: int):
    factors = [qml.s_prod(0.5, qml.PauliX(i)) for i in num]
    return qml.prod(*factors)


def build_old_prod_op(num: int):
    tensor = functools.reduce(lambda x, y: x @ y, [qml.PauliX(i) for i in num])
    return qml.Hamiltonian(coeffs=[0.5], observables=[tensor])


x = np.arange(2, NUM_OPERANDS)
prod_old_times = []
prod_new_times = []
sum_old_times = []
sum_new_times = []

for n in x:
    print(n)
    random_numbers = [random.randint(0, NUM_WIRES) for _ in range(n)]

    new_prod_op = build_new_prod_op(random_numbers)
    old_prod_op = build_old_prod_op(random_numbers)
    new_sum_op = build_new_sum_op(random_numbers)
    old_sum_op = build_old_sum_op(random_numbers)

    start = time()
    new_prod_op.matrix()
    stop = time()
    prod_new_times.append(stop - start)

    start = time()
    qml.matrix(old_prod_op)
    stop = time()
    prod_old_times.append(stop - start)

    start = time()
    new_sum_op.matrix()
    stop = time()
    sum_new_times.append(stop - start)

    start = time()
    qml.matrix(old_sum_op)
    stop = time()
    sum_old_times.append(stop - start)


fig, (ax1, ax2) = plt.subplots(1, 2)
fig.suptitle(f"Matrix (num_wires = {NUM_WIRES})")
ax1.set_title("Prod(PauliX(0), PauliX(1), ...) (random)")
ax1.set_ylabel("new_time / old_time")
ax1.set_xlabel("# factors")
ax1.bar(x, np.array(prod_new_times) / np.array(prod_old_times), width=0.8)
ax1.set_xticks(x[::2])
ax2.set_title("Sum(PauliX(0), PauliX(1), ...) (random)")
ax2.set_ylabel("new_time / old_time")
ax2.set_xlabel("# summands")
ax2.bar(x, np.array(sum_new_times) / np.array(sum_old_times), width=0.8)
ax2.set_xticks(x[::2])
plt.tight_layout()
plt.show()

[Jaybsoni]

Looks good! Just some clarifying questions for me to better understand the implementation. Also Missing some tests! Other then that, looks good 👍🏼

[albi3ro]

The provided benchmarking script seems to still have a bit of random noise and variation, even though a seed is set. But very impressive improvement for sums with large numbers of wires. A bit cleaner and easier to follow now too 👍

[AmintorDusko]

Hi @AlbertMitjans, Nice job here!
I'm happy that this branch outperformed the master in all cases considered in the benchmark.
I'm attaching here a comparison.

[AmintorDusko]

One more thing. Could you please check here?
I think that this sentence should be The method works with scipy sparse matrices, for example:.