[labs] [WIP] CartanSubAlgebra

pennylane/labs/dla/__init__.py

@@ -38,6 +38,8 @@
     ~pauli_coefficients
     ~check_cartan_decomp
     ~check_commutation
+    ~check_all_commuting
+    ~project
     ~apply_basis_change
 
 Involutions
@@ -71,11 +73,26 @@
     concurrence_involution,
     recursive_cartan_decomposition,
 )
+
 from .dense_util import (
     pauli_decompose,
     pauli_coefficients,
     check_cartan_decomp,
     check_commutation,
+    check_all_commuting,
     apply_basis_change,
 )
-from .involutions import khaneja_glaser_involution, AI, AII, AIII, BDI, CI, CII, DIII, ClassB
+
+from .involutions import (
+    khaneja_glaser_involution,
+    AI,
+    AII,
+    AIII,
+    BDI,
+    CI,
+    CII,
+    DIII,
+    ClassB,
+)
+
+from .cartan_subalgebra import cartan_subalgebra, adjvec_to_op, op_to_adjvec

pennylane/labs/dla/cartan_subalgebra.py

@@ -0,0 +1,291 @@
+# Copyright 2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+"""Functionality to compute the Cartan subalgebra"""
+# pylint: disable=too-many-arguments
+import numpy as np
+from scipy.linalg import null_space
+
+import pennylane as qml
+from pennylane.operation import Operator
+from pennylane.pauli import PauliSentence, PauliVSpace
+from pennylane.pauli.dla.center import _intersect_bases
+from pennylane.typing import TensorLike
+
+
+def _gram_schmidt(X):
+    """Orthogonalize basis of column vectors in X"""
+    _, R = np.linalg.qr(X.T, mode="complete")
+    return R.T
+
+
+def _is_independent(v, A, tol=1e-14):
+    v /= np.linalg.norm(v)
+    v = v - A @ qml.math.linalg.solve(qml.math.conj(A.T) @ A, A.conj().T) @ v
+    return np.linalg.norm(v) > tol
+
+
+def _orthogonal_complement_basis(h, m, tol):
+    """find mtilde = m - h"""
+    # Step 1: Find the span of h
+    h = np.array(h)
+    m = np.array(m)
+
+    # Compute the orthonormal basis of h using QR decomposition
+    Q, _ = np.linalg.qr(h.T)
+
+    # Step 2: Project each vector in m onto the orthogonal complement of span(h)
+    projections = m - np.dot(np.dot(m, Q), Q.T)
+
+    # Step 3: Find a basis for the non-zero projections
+    # We'll use SVD to find the basis
+    U, S, _ = np.linalg.svd(projections.T)
+
+    # Choose columns of U corresponding to non-zero singular values
+    rank = np.sum(S > tol)
+    basis = U[:, :rank]
+
+    return basis.T  # Transpose to get row vectors
+
+
+def cartan_subalgebra(g, k, m, ad, which=0, tol=1e-14, verbose=0):
+    r"""
+    Compute Cartan subalgebra of g given m, the odd parity subspace. I.e. the maximal Abelian subalgebra in m.
+
+    .. seealso:: :func:`~cartan_decomposition`, :func:`~structure_constants`
+
+    Args:
+        g (List[Union[PauliSentence, np.ndarray]]): Lie algebra :math:`\mathfrak{g}`, which is assumed to be ordered as :math:`\mathfrak{g} = \mathfrak{k} \oplus \mathfrak{m}`
+        k (List[Union[PauliSentence, np.ndarray]]): Vertical space :math:`\mathfrak{m}` from Cartan decomposition :math:`\mathfrak{g} = \mathfrak{k} \oplus \mathfrak{m}`
+        m (List[Union[PauliSentence, np.ndarray]]): Horizontal space :math:`\mathfrak{m}` from Cartan decomposition :math:`\mathfrak{g} = \mathfrak{k} \oplus \mathfrak{m}`
+        ad (Array): The math:`|\mathfrak{g}| \times |\mathfrak{g}| \times |\mathfrak{g}|` dimensional adjoint representation of :math:`\mathfrak{g}`
+
+    Returns:
+        np_newg, np_k, np_mtilde, np_h, new_adj, basis_change
+        np.ndarray: Adjoint vectors for new ordered :math:`\mathfrak{g}` of dimension :math:`|\mathfrak{g}| \times |\mathfrak{g}|`, ordered as :math:`\mathfrak{g} = \mathfrak{k} \oplus \tilde{\mathfrak{m}} \oplus \mathfrak{h}`
+        np.ndarray: Adjoint vectors for :math:`\mathfrak{k}` of dimension :math:`|\mathfrak{k}| \times |\mathfrak{g}|`.
+        np.ndarray: Adjoint vectors for :math:`\tilde{\mathfrak{m}}` of dimension :math:`|\tilde{\mathfrak{m}}| \times |\mathfrak{g}|`.
+        np.ndarray: Adjoint vectors for :math:`\mathfrak{h}` of dimension :math:`|\mathfrak{h}| \times |\mathfrak{g}|`.
+        np.ndarray: The new :math:`|\mathfrak{g}| \times |\mathfrak{g}| \times |\mathfrak{g}|` dimensional adjoint representation of :math:`\mathfrak{g}_\text{new}`, according to the ordering of new vectors.
+
+    **Example**
+
+    A quick example computing a Cartan subalgebra of :math:`\mathfrak{su}(4)` using the Cartan involution :func:`~even_odd_involution`.
+
+    >>> import pennylane as qml
+    >>> from pennylane.labs.dla import cartan_decomposition, cartan_subalgebra, even_odd_involution
+    >>> g = list(qml.pauli.pauli_group(2)) # u(4)
+    >>> g = g[1:]       # remove identity -> su(4)
+    >>> k, m = cartan_decomposition(g, even_odd_involution)
+    >>> g = k + m # re-order g to separate k and m
+    >>> adj = qml.structure_constants(g)
+    >>> np_newg, np_k, np_mtilde, np_h, new_adj = cartan_subalgebra(g, k, m, adj)
+
+    All results are what we call adjoint vectors of dimension :math:`|\mathfrak{g}|`, where each component corresponds to an entry in (the ordered) ``g``.
+    The adjoint vectors for the Cartan subalgebra are in ``np_h``.
+    We can reconstruct an operator by computing :math:`\hat{O}_v = \sum_i v_i g_i` for an adjoint vector :math:`v` and :math:`g_i \in \mathfrak{g}`.
+
+    >>> v = np_h[0]
+    >>> op = sum(v_i * g_i for v_i, g_i in zip(v, g))
+    >>> op = qml.simplify(op)
+    >>> op
+    Z(0) @ Z(1)
+
+    For convenience, we provide a helper function :func:`~adjvec_to_op` for the collections of adjoint vectors in the returns.
+
+    >>> from pennylane.labs.dla import adjvec_to_op
+    >>> h = adjvec_to_op(np_h, g)
+    >>> h
+    [Z(0) @ Z(1), (-1+0j) * (Y(0) @ X(1)), (-1+0j) * (X(0) @ Y(1))]
+
+    .. details::
+        :title: Usage Details
+
+        Let us walk through an example of computing the Cartan subalgebra. The basis for computing
+        the Cartan subalgebra is having the Lie algebra :math:`\mathfrak{g}`, a Cartan decomposition
+        :math:`\mathfrak{g} = \mathfrak{k} \oplus \mathfrak{m}` and its adjoint representation.
+
+        We start by computing these ingredients using :func:`~cartan_decomposition` and :func:`~structure_constants`.
+        As an example, we take the Lie algebra of the Heisenberg model with generators :math:`\{X_i X_j, Y_i Y_j, Z_i, Z_j\}`.
+
+        >>> from pennylane.labs.dla import lie_closure_dense, cartan_decomposition
+        >>> 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)
+
+        Taking the Heisenberg Lie algebra, can perform the Cartan decomposition. We take the :math:`~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
+        :math:`[\mathfrak{k}, \mathfrak{k}] \subeq \mathfrak{k}`, :math:`[\mathfrak{k}, \mathfrak{m}] \subeq \mathfrak{m}` and :math:`[\mathfrak{m}, \mathfrak{m}] \subeq \mathfrak{k}`,
+        which we can check using the helper function :math:`~check_cartan_decomp`.
+
+        >>> from pennylane.labs.dla import even_odd_involution, check_cartan_decomp
+        >>> k, m = cartan_decomposition(g, even_odd_involution)
+        >>> assert check_cartan_decomp(k, m) # check commutation relations to be valid Cartan decomposition
+
+        Our life is easier when we use a canonical ordering of the operators. This is why we re-define ``g`` with the new ordering in terms of operators in :math:`\mathfrak{k}` first, and then
+        all remaining operators from :math:`\mathfrak{m}`.
+
+        >>> g = np.vstack([k, m]) # re-order g to separate k and m operators
+        >>> adj = structure_constants_dense(g) # compute adjoint representation of g
+
+        Finally, we can compute a Cartan subalgebra :math:`\mathfrak{h}`, a maximal Abelian subalgebra of :math:`\mathfrak{m}`.
+
+        >>> np_newg, np_k, np_mtilde, np_h, new_adj, basis_change = cartan_subalgebra(g, k, m, adj, which=3)
+
+        The final results are all returned as what we call adjoint representation vectors.
+        These are dense vector representations of dimension :math:`|\mathfrak{g}|`, in which each entry corresponds to the respective operator in :math:`\mathfrak{g}`.
+        Given an adjoint representation vector :math:`v`, we can reconstruct the respective operator by simply computing :math:`\hat{O}_v = \sum_i v_i g_i`, where
+        :math:`g_i \in \mathfrak{g}` (hence the need for a canonical ordering).
+
+        We provide a convenience function :func:`~adjvec_to_op` that works with both ``g`` represented as dense matrices or PL operators.
+        Because we used dense matrices in this example, we transform the operators back to PennyLane operators using :func:`~pauli_decompose`.
+
+        >>> from pennylane.labs.dla import adjvec_to_op
+        >>> h = adjvec_to_op(np_h, g)
+        >>> h_op = [qml.pauli_decompose(op).pauli_rep for op in h]
+        >>> h_op
+        [0.35355339059327384 * Y(1) @ Y(2),
+         0.3535533905932738 * Y(0) @ Y(2),
+         -0.3535533905932738 * Y(0) @ Y(1)]
+
+        In that case we chose a Cartan subalgebra from which we can readily see that it is commuting, but we also provide a small helper function to check that.
+
+        >>> from pennylane.labs.dla import check_all_commuting
+        >>> assert check_all_commuting(h_op)
+
+        Last but not least, the adjoint representation ``new_adj`` is updated to represent the new basis and its ordering of ``g``.
+        In particular, we can compute commutators between
+    """
+
+    np_m = op_to_adjvec(m, g)
+    np_h = op_to_adjvec([m[which]], g)
+
+    iteration = 1
+    while True:
+        if verbose:
+            print(f"iteration: {iteration}")
+        kernel_intersection = np_m.T
+        for h_i in np_h:
+
+            # obtain adjoint rep of candidate h_i
+            adjoint_of_h_i = np.einsum("gab,a->gb", ad, h_i)
+            # compute kernel of adjoint
+            new_kernel = null_space(adjoint_of_h_i, rcond=tol)
+
+            # intersect kernel to stay in m
+            kernel_intersection = _intersect_bases(kernel_intersection, new_kernel, rcond=tol)
+
+        if kernel_intersection.shape[1] == len(np_h):
+            # No new vector was added from all the kernels
+            break
+
+        kernel_intersection = _gram_schmidt(kernel_intersection)  # orthogonalize
+        for vec in kernel_intersection.T:
+            if _is_independent(vec, np.array(np_h).T, tol):
+                np_h = np.vstack([np_h, vec])
+                break
+
+        iteration += 1
+
+    np_h = _gram_schmidt(np_h.T).T  # orthogonalize Abelian subalgebra
+    np_k = op_to_adjvec(k, g)  # adjoint vectors of k space for re-ordering
+
+    np_mtilde = _orthogonal_complement_basis(np_h, np_m, tol=tol)  # the "rest" of m without h
+    np_newg = np.vstack([np_k, np_mtilde, np_h])
+
+    # Instead of recomputing the adjoint representation, take the basis transformation oldg -> newg and transform the adjoint representation accordingly
+    # implementation to be tested:
+    basis_change = np.linalg.lstsq(np.vstack([np_k, np_m]).T, np_newg.T, rcond=None)[
+        0
+    ]  # basis change old g @ X = new g => adj_new = contact(adj, X, X, X)
+    new_adj = np.einsum("kmn,ki,mj,nl->ijl", ad, basis_change, basis_change, basis_change)
+
+    return np_newg, np_k, np_mtilde, np_h, new_adj
+
+
+def adjvec_to_op(adj_vecs, basis):
+    """Transform adjoint vectors representing operators back into a basis of choice
+
+    Args:
+        adj_vecs (np.ndarray): collection of vectors with shape ``(batch, dim_basis)``
+        basis (List[Union[PauliSentence, Operator, np.ndarray]]): collection of basis operators
+
+    """
+    res = []
+    # currently agnostic to dense or op input, but could be vectorized in the dense case (TODO?)
+    if all(isinstance(op, PauliSentence) for op in basis):
+        for vec in adj_vecs:
+            op_j = sum(c * op for c, op in zip(vec, basis))
+            op_j.simplify()
+            res.append(op_j)
+        return res
+
+    if all(isinstance(op, Operator) for op in basis):
+        for vec in adj_vecs:
+            op_j = sum(c * op for c, op in zip(vec, basis))
+            op_j = qml.simplify(op_j)
+            res.append(op_j)
+        return res
+
+    if all(isinstance(op, np.ndarray) for op in basis):
+        res = np.einsum("id,dmn->imn", adj_vecs, basis)  # i, dimg | dimg m n
+        return res
+
+    return NotImplementedError
+
+
+def op_to_adjvec(ops, basis):
+    """Project a batch of ops onto a given basis"""
+    if isinstance(basis, PauliVSpace):
+        basis = basis.basis
+
+    if all(isinstance(op, Operator) for op in ops) and all(
+        isinstance(op, Operator) for op in basis
+    ):
+        ops = [op.pauli_rep for op in ops]
+        basis = [op.pauli_rep for op in basis]
+
+    # PauliSentence branch
+    if all(isinstance(op, PauliSentence) for op in ops) and all(
+        isinstance(op, PauliSentence) for op in basis
+    ):
+        res = []
+        for op in ops:
+            rep = np.zeros((len(basis),), dtype=complex)
+            for i, basis_i in enumerate(basis):
+                # v = ∑ (v · e_j / ||e_j||^2) * e_j
+                value = (basis_i @ op).trace()
+                value = value / (basis_i @ basis_i).trace()
+
+                rep[i] = value
+
+            res.append(rep)
+
+        return np.array(res)
+
+    # dense branch
+    if all(isinstance(op, TensorLike) for op in ops) and all(
+        isinstance(op, TensorLike) for op in basis
+    ):
+        basis = np.array(basis)
+        # if len(ops.shape) == 2:
+        #     ops = np.array([ops])
+        ops = np.array(ops)
+
+        res = np.einsum("bij,cji->bc", ops, basis)
+
+        return res
+
+    return NotImplemented

pennylane/labs/dla/dense_util.py

@@ -12,12 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Utility tools for dense Lie algebra representations"""
-from itertools import product
+from itertools import combinations, product
 from typing import List
 
 import numpy as np
 
 import pennylane as qml
+from pennylane.operation import Operator
 from pennylane.ops.qubit.matrix_ops import _walsh_hadamard_transform
 from pennylane.pauli import PauliSentence, PauliWord
 
@@ -218,6 +219,34 @@ def check_commutation(ops1, ops2, vspace):
     return all(assert_vals)
 
 
+def check_all_commuting(ops: List[PauliSentence]):
+    """Helper function to check if all operators in a set of operators commute"""
+    res = []
+    if all(isinstance(op, PauliSentence) for op in ops):
+        for oi, oj in combinations(ops, 2):
+            com = oj.commutator(oi)
+            com.simplify()
+            res.append(len(com) == 0)
+
+        return all(res)
+
+    if all(isinstance(op, Operator) for op in ops):
+        for oi, oj in combinations(ops, 2):
+            com = qml.simplify(qml.commutator(oj, oi))
+            res.append(qml.equal(com, 0 * qml.Identity()))
+
+        return all(res)
+
+    if all(isinstance(op, PauliSentence) for op in ops):
+        for oi, oj in combinations(ops, 2):
+            com = oj @ oi - oi @ oj
+            res.append(np.allclose(com, np.zeros_like(com)))
+
+        return all(res)
+
+    return NotImplemented
+
+
 def check_cartan_decomp(k: List[PauliSentence], m: List[PauliSentence], verbose=True):
     """Helper function to check the validity of a Cartan decomposition by checking its commutation relations"""
     if any(isinstance(op, np.ndarray) for op in k):
@@ -240,7 +269,8 @@ def check_cartan_decomp(k: List[PauliSentence], m: List[PauliSentence], verbose=
 
 
 def apply_basis_change(change_op, targets):
-    if single_target := (np.ndim(targets) == 2):
+    """Helper function for recursive Cartan decompositions"""
+    if single_target := np.ndim(targets) == 2:
         targets = [targets]
     if isinstance(targets, list):
         targets = np.array(targets)