Upgrade qml.lie_closure to handle dense matrices

doc/code/qml_pauli.rst

@@ -167,3 +167,15 @@ See our `introduction to Dynamical Lie Algebras for quantum practitioners <https
     ~lie_closure
     ~structure_constants
     ~center
+
+Utility functions
+^^^^^^^^^^^^^^^^^
+
+Utility functions for handling Pauli operators in the context of DLA functionality.
+
+.. currentmodule:: pennylane.pauli
+
+.. autosummary::
+    :toctree: api
+
+    ~trace_inner_product

doc/releases/changelog-dev.md

@@ -4,6 +4,10 @@
 
 <h3>New features since last release</h3>
 
+* ``qml.lie_closure`` now accepts and outputs dense inputs using the ``dense`` keyword.
+  Also added ``qml.pauli.trace_inner_product`` that can handle batches of dense matrices.
+  [(#6811)](https://github.com/PennyLaneAI/pennylane/pull/6811)
+
 <h3>Improvements 🛠</h3>
 
 * `QNode` objects now have an `update` method that allows for re-configuring settings like `diff_method`, `mcm_method`, and more. This allows for easier on-the-fly adjustments to workflows. Any arguments not specified will retain their original value.
@@ -44,6 +48,11 @@
 * An informative error is raised when a `QNode` with `diff_method=None` is differentiated.
   [(#6770)](https://github.com/PennyLaneAI/pennylane/pull/6770)
 
+<h3>Labs: a place for unified and rapid prototyping of research software 🧪</h3>
+
+* ``pennylane.labs.dla.lie_closure_dense`` is removed and integrated into ``qml.lie_closure`` using the new ``dense`` keyword.
+  [(#6811)](https://github.com/PennyLaneAI/pennylane/pull/6811)
+
 <h3>Breaking changes 💔</h3>
 
 * `MultiControlledX` no longer accepts strings as control values.
@@ -104,5 +113,6 @@ Yushao Chen,
 Diksha Dhawan,
 Pietropaolo Frisoni,
 Marcus Gisslén,
+Korbinian Kottmann,
 Christina Lee,
 Andrija Paurevic

pennylane/labs/dla/__init__.py

@@ -20,7 +20,6 @@
 .. autosummary::
     :toctree: api
 
-    ~lie_closure_dense
     ~structure_constants_dense
     ~cartan_decomp
     ~recursive_cartan_decomp
@@ -38,7 +37,6 @@
 
     ~adjvec_to_op
     ~op_to_adjvec
-    ~trace_inner_product
     ~orthonormalize
     ~pauli_coefficients
     ~batched_pauli_decompose
@@ -82,7 +80,6 @@
 
 """
 
-from .lie_closure_dense import lie_closure_dense
 from .structure_constants_dense import structure_constants_dense
 from .cartan import (
     cartan_decomp,
@@ -97,7 +94,6 @@
     check_orthonormal,
     pauli_coefficients,
     batched_pauli_decompose,
-    trace_inner_product,
     op_to_adjvec,
     orthonormalize,
 )

pennylane/labs/dla/cartan_subalgebra.py

@@ -151,13 +151,13 @@ def cartan_subalgebra(
         We start by computing these ingredients using :func:`~cartan_decomp` and :func:`~structure_constants`.
         As an example, we take the Lie algebra of the Heisenberg model with generators :math:`\{X_i X_{i+1}, Y_i Y_{i+1}, Z_i Z_{i+1}\}`.
 
-        >>> from pennylane.labs.dla import lie_closure_dense, cartan_decomp
+        >>> from pennylane.labs.dla import cartan_decomp
         >>> from pennylane import X, Y, Z
         >>> n = 3
         >>> gens = [X(i) @ X(i+1) for i in range(n-1)]
         >>> gens += [Y(i) @ Y(i+1) for i in range(n-1)]
         >>> gens += [Z(i) @ Z(i+1) for i in range(n-1)]
-        >>> g = lie_closure_dense(gens)
+        >>> g = qml.lie_closure(gens, dense=True)
 
         Taking the Heisenberg Lie algebra, we can perform the Cartan decomposition. We take the :func:`~even_odd_involution` as a valid Cartan involution.
         The resulting vertical and horizontal subspaces :math:`\mathfrak{k}` and :math:`\mathfrak{m}` need to fulfill the commutation relations

pennylane/labs/dla/dense_util.py

@@ -23,7 +23,7 @@
 import pennylane as qml
 from pennylane.operation import Operator
 from pennylane.ops.qubit.matrix_ops import _walsh_hadamard_transform
-from pennylane.pauli import PauliSentence, PauliVSpace, PauliWord
+from pennylane.pauli import PauliSentence, PauliVSpace, PauliWord, trace_inner_product
 from pennylane.typing import TensorLike
 
 
@@ -513,9 +513,9 @@ def check_orthonormal(g: Iterable[Union[PauliSentence, Operator]], inner_product
 
     **Example**
 
-    >>> from pennylane.labs.dla import orthonormalize, check_orthonormal, trace_inner_product
+    >>> from pennylane.labs.dla import orthonormalize, check_orthonormal
     >>> ops = [qml.X(0), qml.X(0) + qml.Y(0), qml.Y(0) + qml.Z(0)]
-    >>> check_orthonormal(ops, trace_inner_product)
+    >>> check_orthonormal(ops, qml.pauli.trace_inner_product)
     False
     >>> ops_orth = orthonormalize(ops)
     >>> check_orthonormal(ops_orth, trace_inner_product)
@@ -530,62 +530,6 @@ def check_orthonormal(g: Iterable[Union[PauliSentence, Operator]], inner_product
     return True
 
 
-def trace_inner_product(
-    A: Union[PauliSentence, Operator, np.ndarray], B: Union[PauliSentence, Operator, np.ndarray]
-):
-    r"""Implementation of the trace inner product :math:`\langle A, B \rangle = \text{tr}\left(A B\right)/\text{dim}(A)` between two Hermitian operators :math:`A` and :math:`B`.
-
-    If the inputs are ``np.ndarray``, leading broadcasting axes are supported for either or both
-    inputs.
-
-    Args:
-        A (Union[PauliSentence, Operator, np.ndarray]): First operator
-        B (Union[PauliSentence, Operator, np.ndarray]): Second operator
-
-    Returns:
-        Union[float, np.ndarray]: Result is either a single float or a batch of floats.
-
-    **Example**
-
-    >>> from pennylane.labs.dla import trace_inner_product
-    >>> trace_inner_product(qml.X(0) + qml.Y(0), qml.Y(0) + qml.Z(0))
-    1.0
-
-    If both operators are dense arrays, a leading batch dimension is broadcasted.
-
-    >>> batch = 10
-    >>> ops1 = np.random.rand(batch, 16, 16)
-    >>> op2 = np.random.rand(16, 16)
-    >>> trace_inner_product(ops1, op2).shape
-    (10,)
-    >>> trace_inner_product(op2, ops1).shape
-    (10,)
-
-    We can also have both arguments broadcasted.
-
-    >>> trace_inner_product(ops1, ops1).shape
-    (10, 10)
-
-    """
-    if getattr(A, "pauli_rep", None) is not None and getattr(B, "pauli_rep", None) is not None:
-        return (A.pauli_rep @ B.pauli_rep).trace()
-
-    if all(isinstance(op, np.ndarray) for op in A) and all(isinstance(op, np.ndarray) for op in B):
-        A = np.array(A)
-        B = np.array(B)
-
-    if not isinstance(A, type(B)):
-        raise TypeError("Both input operators need to be of the same type")
-
-    if isinstance(A, np.ndarray):
-        assert A.shape[-2:] == B.shape[-2:]
-        # The axes of the first input are switched, compared to tr[A@B], because we need to
-        # transpose A.
-        return np.tensordot(A, B, axes=[[-1, -2], [-2, -1]]) / A.shape[-1]
-
-    raise NotImplementedError
-
-
 def change_basis_ad_rep(adj: np.ndarray, basis_change: np.ndarray):
     r"""Apply a ``basis_change`` between bases of operators to the adjoint representation ``adj``.
 

pennylane/labs/dla/structure_constants_dense.py

@@ -24,15 +24,15 @@ def structure_constants_dense(g: TensorLike, is_orthonormal: bool = True) -> Ten
     Compute the structure constants that make up the adjoint representation of a Lie algebra.
 
     This function computes the structure constants of a Lie algebra provided by their dense matrix representation,
-    obtained from, e.g., :func:`~lie_closure_dense`.
+    obtained from, e.g., :func:`~lie_closure`.
     This is sometimes more efficient than using the sparse Pauli representations of :class:`~PauliWord` and
     :class:`~PauliSentence` that are employed in :func:`~structure_constants`, e.g., when there are few generators
     that are sums of many Paulis.
 
     .. seealso:: For details on the mathematical definitions, see :func:`~structure_constants` and the section "Lie algebra basics" in our `g-sim demo <https://pennylane.ai/qml/demos/tutorial_liesim/#lie-algebra-basics>`__.
 
     Args:
-        g (np.array): The (dynamical) Lie algebra provided as dense matrices, as generated from :func:`~lie_closure_dense`.
+        g (np.array): The (dynamical) Lie algebra provided as dense matrices, as generated from :func:`~lie_closure`.
             ``g`` should have shape ``(d, 2**n, 2**n)`` where ``d`` is the dimension of the algebra and ``n`` is the number of qubits. Each matrix ``g[i]`` should be Hermitian.
         is_orthonormal (bool): Whether or not the matrices in ``g`` are orthonormal with respect to the Hilbert-Schmidt inner product on
             (skew-)Hermitian matrices. If the inputs are orthonormal, it is recommended to set ``is_orthonormal`` to ``True`` to reduce
@@ -44,12 +44,11 @@ def structure_constants_dense(g: TensorLike, is_orthonormal: bool = True) -> Ten
 
     **Example**
 
-    Let us generate the DLA of the transverse field Ising model using :func:`~lie_closure_dense`.
+    Let us generate the DLA of the transverse field Ising model using :func:`~lie_closure`.
 
-    >>> from pennylane.labs.dla import lie_closure_dense
     >>> n = 4
     >>> gens = [qml.X(i) @ qml.X(i+1) + qml.Y(i) @ qml.Y(i+1) + qml.Z(i) @ qml.Z(i+1) for i in range(n-1)]
-    >>> g = lie_closure_dense(gens)
+    >>> g = qml.lie_closure(gens, dense=True)
     >>> g.shape
     (12, 16, 16)
 

pennylane/labs/tests/dla/test_cartan.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for pennylane/dla/lie_closure_dense.py functionality"""
+"""Tests for pennylane/dla/cartan.py functionality"""
 # pylint: disable=no-self-use,too-few-public-methods,missing-class-docstring
 import pytest