Repository: PennyLaneAI/pennylane
Pull request #6846 (2025-01-16 19:48:21)

make sparse an optional feature
refactor sparse into optional feature

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/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.

Patch:
diff --git a/pennylane/labs/vibronic/hamiltonians/benchmark.py b/pennylane/labs/vibronic/hamiltonians/benchmark.py
index 94f350d17..fc669ea26 100644
--- a/pennylane/labs/vibronic/hamiltonians/benchmark.py
+++ b/pennylane/labs/vibronic/hamiltonians/benchmark.py
@@ -10,7 +10,7 @@ from pennylane.labs.vibronic import VibronicHamiltonian, coeffs
 N_STATES_1 = 8
 N_MODES_1 = 10
 
-vham1 = VibronicHamiltonian(N_STATES_1, N_MODES_1, *coeffs1())
+vham1 = VibronicHamiltonian(N_STATES_1, N_MODES_1, *coeffs1(), sparse=True)
 
 N_STATES_2 = 8
 N_MODES_2 = 21
diff --git a/pennylane/labs/vibronic/vibronic/vibronic_hamiltonian.py b/pennylane/labs/vibronic/vibronic/vibronic_hamiltonian.py
index dc10c386e..72e2ebf83 100644
--- a/pennylane/labs/vibronic/vibronic/vibronic_hamiltonian.py
+++ b/pennylane/labs/vibronic/vibronic/vibronic_hamiltonian.py
@@ -23,6 +23,7 @@ class VibronicHamiltonian:
         betas: np.ndarray,
         lambdas: np.ndarray,
         omegas: np.ndarray,
+        sparse: bool = False,
     ):
         if not is_pow_2(states) or states == 0:
             raise ValueError(f"States must be a positive power of 2, got {states} states.")
@@ -48,6 +49,7 @@ class VibronicHamiltonian:
         self.betas = betas
         self.lambdas = lambdas
         self.omegas = omegas
+        self.sparse = sparse
 
     def fragment(self, index: int) -> VibronicMatrix:
         """Return the fragment at the given index"""
@@ -79,11 +81,11 @@ class VibronicHamiltonian:
         term = VibronicTerm(("P", "P"), Node.tensor_node(np.diag(self.omegas) / 2))
         word = VibronicWord((term,))
         blocks = {(i, i): word for i in range(self.states)}
-        return VibronicMatrix(self.states, self.modes, blocks)
+        return VibronicMatrix(self.states, self.modes, blocks, sparse=self.sparse)
 
     def _fragment(self, i: int) -> VibronicMatrix:
         blocks = {(j, i ^ j): self.v_word(j, i ^ j) for j in range(self.states)}
-        return VibronicMatrix(self.states, self.modes, blocks)
+        return VibronicMatrix(self.states, self.modes, blocks, sparse=self.sparse)
 
     def block_operator(self) -> VibronicMatrix:
         """Return the block representation of the Hamiltonian"""
@@ -102,7 +104,7 @@ class VibronicHamiltonian:
         # pylint: disable=arguments-out-of-order
         """Compute the error matrix"""
         scalar = -(delta**2) / 24
-        epsilon = VibronicMatrix(self.states, self.modes)
+        epsilon = VibronicMatrix(self.states, self.modes, sparse=self.sparse)
 
         for i in range(self.states):
             for j in range(i + 1, self.states + 1):
@@ -133,7 +135,7 @@ class VibronicHamiltonian:
             term = VibronicTerm(("P", "P"), node)
             blocks[(j, m ^ j)] = VibronicWord((term,))
 
-        return VibronicMatrix(self.states, self.modes, blocks)
+        return VibronicMatrix(self.states, self.modes, blocks, sparse=self.sparse)
 
     def __add__(self, other: VibronicHamiltonian) -> VibronicHamiltonian:
         if not isinstance(other, VibronicHamiltonian):
diff --git a/pennylane/labs/vibronic/vibronic/vibronic_matrix.py b/pennylane/labs/vibronic/vibronic/vibronic_matrix.py
index c46b3b8e2..ca4235862 100644
--- a/pennylane/labs/vibronic/vibronic/vibronic_matrix.py
+++ b/pennylane/labs/vibronic/vibronic/vibronic_matrix.py
@@ -18,7 +18,11 @@ class VibronicMatrix:
     """The VibronicMatrix class"""
 
     def __init__(
-        self, states: int, modes: int, blocks: Dict[Tuple[int, int], VibronicWord] = None
+        self,
+        states: int,
+        modes: int,
+        blocks: Dict[Tuple[int, int], VibronicWord] = None,
+        sparse: bool = False,
     ) -> VibronicMatrix:
 
         if not is_pow_2(states) or states == 0:
@@ -30,6 +34,7 @@ class VibronicMatrix:
         self._blocks = blocks
         self.states = states
         self.modes = modes
+        self.sparse = sparse
 
     def block(self, row: int, col: int) -> VibronicWord:
         """Return the block indexed at (row, col)"""
@@ -106,27 +111,30 @@ class VibronicMatrix:
         if self.states != 1:
             raise RuntimeError("Base case called on VibronicMatrix with >1 state")
 
+        if self.sparse:
+            return self._norm_base_case_sparse(gridpoints)
+
         word = self.block(0, 0)
         norm = 0
 
         for term in word.terms:
             term_op_norm = math.prod(map(lambda op: op_norm(gridpoints) ** len(op), term.ops))
-            coeff_sum = sum(abs(term.coeffs.compute(index)) for index in term.coeffs.nonzero)
+            coeff_sum = sum(
+                abs(term.coeffs.compute(index))
+                for index in product(range(self.modes), repeat=len(term.ops))
+            )
             norm += coeff_sum * term_op_norm
 
         return norm
 
-    def _norm_base_case_bad(self, gridpoints: int) -> float:
-        if self.states != 1:
-            raise RuntimeError("Base case called on VibronicMatrix with >1 state")
-
+    def _norm_base_case_sparse(self, gridpoints: int) -> float:
         word = self.block(0, 0)
         norm = 0
 
         for term in word.terms:
             term_op_norm = math.prod(map(lambda op: op_norm(gridpoints) ** len(op), term.ops))
-            scalar = term.coeffs.compute_average() * (self.modes ** len(term.ops))
-            norm += scalar * term_op_norm
+            coeff_sum = sum(abs(term.coeffs.compute(index)) for index in term.coeffs.nonzero())
+            norm += coeff_sum * term_op_norm
 
         return norm
 
@@ -154,7 +162,9 @@ class VibronicMatrix:
         for key in r_keys.difference(l_keys):
             new_blocks[key] = other._blocks[key]
 
-        return VibronicMatrix(self.states, self.modes, new_blocks)
+        return VibronicMatrix(
+            self.states, self.modes, new_blocks, sparse=(self.sparse and other.sparse)
+        )
 
     def __sub__(self, other: VibronicMatrix) -> VibronicMatrix:
         if self.states != other.states:
@@ -181,14 +191,16 @@ class VibronicMatrix:
         for key in r_keys.difference(l_keys):
             new_blocks[key] = (-1) * other._blocks[key]
 
-        return VibronicMatrix(self.states, self.modes, new_blocks)
+        return VibronicMatrix(
+            self.states, self.modes, new_blocks, sparse=(self.sparse and other.sparse)
+        )
 
     def __mul__(self, scalar: float) -> VibronicMatrix:
         new_blocks = {}
         for key in self._blocks.keys():
             new_blocks[key] = scalar * self._blocks[key]
 
-        return VibronicMatrix(self.states, self.modes, new_blocks)
+        return VibronicMatrix(self.states, self.modes, new_blocks, sparse=self.sparse)
 
     __rmul__ = __mul__
 
@@ -209,7 +221,9 @@ class VibronicMatrix:
                 f"Cannot multiply VibronicMatrix on {self.modes} states with VibronicMatrix on {other.modes} states."
             )
 
-        product_matrix = VibronicMatrix(self.states, self.modes)
+        product_matrix = VibronicMatrix(
+            self.states, self.modes, sparse=(self.sparse and other.sparse)
+        )
 
         for i, j in product(range(self.states), repeat=2):
             block_products = [self.block(i, k) @ other.block(k, j) for k in range(self.states)]
@@ -234,10 +248,10 @@ class VibronicMatrix:
         # pylint: disable=chained-comparison
         half = self.states // 2
 
-        top_left = VibronicMatrix(half, self.modes, {})
-        top_right = VibronicMatrix(half, self.modes, {})
-        bottom_left = VibronicMatrix(half, self.modes, {})
-        bottom_right = VibronicMatrix(half, self.modes, {})
+        top_left = VibronicMatrix(half, self.modes, {}, sparse=self.sparse)
+        top_right = VibronicMatrix(half, self.modes, {}, sparse=self.sparse)
+        bottom_left = VibronicMatrix(half, self.modes, {}, sparse=self.sparse)
+        bottom_right = VibronicMatrix(half, self.modes, {}, sparse=self.sparse)
 
         for index, word in self._blocks.items():
             x, y = index
@@ -253,22 +267,6 @@ class VibronicMatrix:
 
         return top_left, top_right, bottom_left, bottom_right
 
-    def num_coeffs(self) -> int:
-        """Find number of coefficients"""
-
-        term_count = 0
-        zero_count = 0
-        coeffs = 0
-        for _, word in self._blocks.items():
-            print(word.is_zero, len(word.terms))
-            for term in word.terms:
-                coeffs += self.modes ** len(term.ops)
-                term_count += 1
-                if term.is_zero:
-                    zero_count += 1
-
-        print(f"{zero_count}/{term_count}", coeffs)
-
 
 def commutator(a: VibronicMatrix, b: VibronicMatrix):
     """Return the commutator [a, b] = ab - ba"""
diff --git a/pennylane/labs/vibronic/vibronic/vibronic_tree.py b/pennylane/labs/vibronic/vibronic/vibronic_tree.py
index 625bd1b8e..cdcd820bd 100644
--- a/pennylane/labs/vibronic/vibronic/vibronic_tree.py
+++ b/pennylane/labs/vibronic/vibronic/vibronic_tree.py
@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from enum import Enum
-from itertools import product, tee
+from itertools import product
 from typing import Iterator, Tuple
 
 from numpy import allclose, isclose, ndarray, zeros
@@ -28,7 +28,6 @@ class Node:  # pylint: disable=too-many-instance-attributes
     def __init__(  # pylint: disable=too-many-arguments
         self,
         node_type: NodeType,
-        nonzero: Iterator[Tuple[int]],
         l_child: Node = None,
         r_child: Node = None,
         l_shape: Tuple[int] = (),
@@ -48,7 +47,6 @@ class Node:  # pylint: disable=too-many-instance-attributes
         self.scalar = scalar
         self.value = value
         self.is_zero = is_zero
-        self.nonzero = nonzero
 
         if node_type == NodeType.SUM:
             self.shape = l_child.shape
@@ -71,15 +69,8 @@ class Node:  # pylint: disable=too-many-instance-attributes
                 f"Cannot add Node of shape {l_child.shape} with Node of shape {r_child.shape}."
             )
 
-        lnz1, lnz2 = tee(l_child.nonzero, 2)
-        rnz1, rnz2 = tee(r_child.nonzero, 2)
-
-        l_child.nonzero = lnz1
-        r_child.nonzero = rnz1
-
         return cls(
             node_type=NodeType.SUM,
-            nonzero=_uniq_chain(lnz2, rnz2),
             l_child=l_child,
             r_child=r_child,
             l_shape=l_child.shape,
@@ -95,15 +86,8 @@ class Node:  # pylint: disable=too-many-instance-attributes
     ) -> Node:
         """Construct a OUTER node"""
 
-        lnz1, lnz2 = tee(l_child.nonzero, 2)
-        rnz1, rnz2 = tee(r_child.nonzero, 2)
-
-        l_child.nonzero = lnz1
-        r_child.nonzero = rnz1
-
         return cls(
             node_type=NodeType.OUTER,
-            nonzero=_flatten_product(lnz2, rnz2),
             l_child=l_child,
             r_child=r_child,
             l_shape=l_child.shape,
@@ -142,14 +126,12 @@ class Node:  # pylint: disable=too-many-instance-attributes
         if len(tensor.shape):
             return cls(
                 node_type=NodeType.TENSOR,
-                nonzero=zip(*tensor.nonzero()),
                 tensor=tensor,
                 is_zero=allclose(tensor, zeros(tensor.shape)),
             )
 
         return cls(
             node_type=NodeType.FLOAT,
-            nonzero=((),) if tensor else iter(()),
             value=tensor,
             is_zero=(tensor == 0),
         )
@@ -158,12 +140,8 @@ class Node:  # pylint: disable=too-many-instance-attributes
     def scalar_node(cls, scalar: float, child: Node) -> Node:
         """Construct a SCALAR node"""
 
-        nz1, nz2 = tee(child.nonzero, 2)
-        child.nonzero = nz1
-
         return cls(
             node_type=NodeType.SCALAR,
-            nonzero=nz2 if scalar else iter(()),
             l_child=child,
             l_shape=child.shape,
             scalar=scalar,
@@ -298,6 +276,29 @@ class Node:  # pylint: disable=too-many-instance-attributes
 
         raise ValueError(f"Node was constructed with invalid NodeType {self.node_type}.")
 
+    def nonzero(self) -> Iterator[Tuple[int]]:
+        """Compute the nonzero indices"""
+
+        if self.node_type == NodeType.TENSOR:
+            return zip(*self.tensor.nonzero())
+
+        if self.node_type == NodeType.FLOAT:
+            return iter(((),)) if self.value else iter(())
+
+        if self.node_type == NodeType.SCALAR:
+            return self.l_child.nonzero() if self.scalar else iter(())
+
+        if self.node_type == NodeType.SUM:
+            return _uniq_chain(self.l_child.nonzero(), self.r_child.nonzero())
+
+        if self.node_type == NodeType.OUTER:
+            return _flatten_product(self.l_child.nonzero(), self.r_child.nonzero())
+
+        if self.node_type == NodeType.HADAMARD:
+            raise NotImplementedError
+
+        raise ValueError(f"Node was constructed with invalid NodeType {self.node_type}.")
+
     def _validate_index(self, index: Tuple[int]) -> bool:
         if len(index) != len(self.shape):
             return False


Test patch:
diff --git a/pennylane/labs/vibronic/tests/test_matrix.py b/pennylane/labs/vibronic/tests/test_matrix.py
index a9e435242..f71b1aae4 100644
--- a/pennylane/labs/vibronic/tests/test_matrix.py
+++ b/pennylane/labs/vibronic/tests/test_matrix.py
@@ -1,10 +1,14 @@
 """Test matrix operations"""
 
+from typing import Dict, Tuple
+
 import numpy as np
 import pytest
 
 from pennylane.labs.vibronic import Node, VibronicMatrix, VibronicTerm, VibronicWord, op_norm
 
+# pylint: disable=too-many-arguments
+
 vword0 = VibronicWord([VibronicTerm(tuple(), Node.tensor_node(np.array(0.5)))])
 blocks0 = {(0, 0): vword0}
 vmat0 = VibronicMatrix(1, 1, blocks0)
@@ -29,40 +33,78 @@ class TestMatrix:
     """Test properties of the VibronicMatrix class"""
 
     params = [
-        (vmat0, 2, 0.5),
-        (vmat0, 4, 0.5),
-        (vmat1, 2, op_norm(2)),
-        (vmat1, 4, op_norm(4)),
-        (vmat2, 2, 6 * op_norm(2) ** 2),
-        (vmat2, 4, 6 * op_norm(4) ** 2),
-        (vmat3, 2, np.sqrt(18 * op_norm(2) ** 3)),
-        (vmat3, 4, np.sqrt(18 * op_norm(4) ** 3)),
-        (vmat4, 2, 6 * op_norm(2) ** 2 + np.sqrt(18 * op_norm(2) ** 3)),
-        (vmat4, 4, 6 * op_norm(4) ** 2 + np.sqrt(18 * op_norm(4) ** 3)),
+        (blocks0, 2, 1, 1, False, 0.5),
+        (blocks0, 4, 1, 1, False, 0.5),
+        (blocks0, 2, 1, 1, True, 0.5),
+        (blocks0, 4, 1, 1, True, 0.5),
+        (blocks1, 2, 1, 1, False, op_norm(2)),
+        (blocks1, 4, 1, 1, False, op_norm(4)),
+        (blocks1, 2, 1, 1, True, op_norm(2)),
+        (blocks1, 4, 1, 1, True, op_norm(4)),
+        (blocks2, 2, 2, 2, False, 6 * op_norm(2) ** 2),
+        (blocks2, 4, 2, 2, False, 6 * op_norm(4) ** 2),
+        (blocks2, 2, 2, 2, True, 6 * op_norm(2) ** 2),
+        (blocks2, 4, 2, 2, True, 6 * op_norm(4) ** 2),
+        (blocks3, 2, 2, 2, False, np.sqrt(18 * op_norm(2) ** 3)),
+        (blocks3, 4, 2, 2, False, np.sqrt(18 * op_norm(4) ** 3)),
+        (blocks3, 2, 2, 2, True, np.sqrt(18 * op_norm(2) ** 3)),
+        (blocks3, 4, 2, 2, True, np.sqrt(18 * op_norm(4) ** 3)),
+        (blocks4, 2, 2, 2, False, 6 * op_norm(2) ** 2 + np.sqrt(18 * op_norm(2) ** 3)),
+        (blocks4, 4, 2, 2, False, 6 * op_norm(4) ** 2 + np.sqrt(18 * op_norm(4) ** 3)),
+        (blocks4, 2, 2, 2, True, 6 * op_norm(2) ** 2 + np.sqrt(18 * op_norm(2) ** 3)),
+        (blocks4, 4, 2, 2, True, 6 * op_norm(4) ** 2 + np.sqrt(18 * op_norm(4) ** 3)),
     ]
 
-    @pytest.mark.parametrize("vmatrix, gridpoints, expected", params)
-    def test_norm(self, vmatrix: VibronicMatrix, gridpoints: int, expected: float):
+    @pytest.mark.parametrize("blocks, gridpoints, states, modes, sparse, expected", params)
+    def test_norm(
+        self,
+        blocks: Dict[Tuple[int], VibronicWord],
+        gridpoints: int,
+        states: int,
+        modes: int,
+        sparse: bool,
+        expected: float,
+    ):
         """Test that the norm is correct"""
 
+        vmatrix = VibronicMatrix(states, modes, blocks, sparse=sparse)
+
         assert np.isclose(vmatrix.norm(gridpoints), expected)
 
     params = [
-        (vmat0, 2),
-        (vmat0, 4),
-        (vmat1, 2),
-        (vmat1, 4),
-        (vmat2, 2),
-        (vmat2, 4),
-        (vmat3, 2),
-        (vmat3, 4),
-        (vmat4, 2),
-        (vmat4, 4),
+        (blocks0, 2, 1, 1, False),
+        (blocks0, 4, 1, 1, False),
+        (blocks0, 2, 1, 1, True),
+        (blocks0, 4, 1, 1, True),
+        (blocks1, 2, 1, 1, False),
+        (blocks1, 4, 1, 1, False),
+        (blocks1, 2, 1, 1, True),
+        (blocks1, 4, 1, 1, True),
+        (blocks2, 2, 2, 2, False),
+        (blocks2, 4, 2, 2, False),
+        (blocks2, 2, 2, 2, True),
+        (blocks2, 4, 2, 2, True),
+        (blocks3, 2, 2, 2, False),
+        (blocks3, 4, 2, 2, False),
+        (blocks3, 2, 2, 2, True),
+        (blocks3, 4, 2, 2, True),
+        (blocks4, 2, 2, 2, False),
+        (blocks4, 4, 2, 2, False),
+        (blocks4, 2, 2, 2, True),
+        (blocks4, 4, 2, 2, True),
     ]
 
-    @pytest.mark.parametrize("vmatrix, gridpoints", params)
-    def test_norm_against_numpy(self, vmatrix: VibronicMatrix, gridpoints: int):
+    @pytest.mark.parametrize("blocks, gridpoints, states, modes, sparse", params)
+    def test_norm_against_numpy(
+        self,
+        blocks: Dict[Tuple[int], VibronicWord],
+        gridpoints: int,
+        states: int,
+        modes: int,
+        sparse: bool,
+    ):
         """Test that .norm is an upper bound on the true norm"""
+        vmatrix = VibronicMatrix(states, modes, blocks, sparse=sparse)
         upper_bound = vmatrix.norm(gridpoints)
         norm = np.max(np.linalg.eigvals(vmatrix.matrix(gridpoints)))
         assert norm <= upper_bound
diff --git a/pennylane/labs/vibronic/tests/test_norms.py b/pennylane/labs/vibronic/tests/test_norms.py
index 3bf6ecd99..c244e7bf1 100644
--- a/pennylane/labs/vibronic/tests/test_norms.py
+++ b/pennylane/labs/vibronic/tests/test_norms.py
@@ -9,7 +9,7 @@ from pennylane.labs.vibronic.hamiltonians.spin_vibronic_ham import get_coeffs as
 def test_against_10_mode():
     expected = 1.6621227513071748
 
-    vham = VibronicHamiltonian(8, 10, *ten_mode())
+    vham = VibronicHamiltonian(8, 10, *ten_mode(), sparse=True)
     ep = vham.epsilon(1)
     actual = ep.norm(4)