Add concurrency support for effective_hamiltonian in labs (#8081)

**Context:**
Implements concurrency in `effective_hamiltonian` to support running
parallel jobs.

**Description of the Change:**

**Benefits:**

**Possible Drawbacks:**

**Related GitHub Issues:**

Author: soranjh

doc/releases/changelog-dev.md

@@ -515,9 +515,12 @@
   [(#8156)](https://github.com/PennyLaneAI/pennylane/pull/8156)
 <h3>Labs: a place for unified and rapid prototyping of research software 🧪</h3>
 
-  * Fixed a queueing issue in `ResourceOperator` tests.
+* Added concurrency support for `effective_hamiltonian` in labs.
+  [(#8081)](https://github.com/PennyLaneAI/pennylane/pull/8081)
+
+* Fixed a queueing issue in `ResourceOperator` tests.
   [(#8204)](https://github.com/PennyLaneAI/pennylane/pull/8204)
-  
+
 * The module `qml.labs.zxopt` has been removed as its functionalities are now available in the
   submodule :mod:`~.transforms.zx`. The same functions are available, but their signature
   may have changed.

pennylane/labs/tests/trotter_error/product_formulas/test_bch.py

@@ -277,3 +277,45 @@ def test_against_matrix_log(fragments, product_formula):
     log_error = log - (1j * t * ham) / t**order
 
     assert np.allclose(np.linalg.norm(bch_error - log_error), 0)
+
+
+params = [
+    ("serial", 1),
+    ("mp_pool", 1),
+    ("mp_pool", 2),
+    ("cf_procpool", 1),
+    ("cf_procpool", 2),
+    ("cf_threadpool", 1),
+    ("cf_threadpool", 2),
+    ("mpi4py_pool", 1),
+    ("mpi4py_pool", 2),
+    ("mpi4py_comm", 1),
+    ("mpi4py_comm", 2),
+]
+
+
+@pytest.mark.parametrize("backend, num_workers", params)
+def test_effective_hamiltonian_backend(backend, num_workers, mpi4py_support):
+    """Test that effective_hamiltonian function runs without errors for different backends."""
+
+    if backend in {"mpi4py_pool", "mpi4py_comm"} and not mpi4py_support:
+        pytest.skip(f"Skipping test: '{backend}' requires mpi4py, which is not installed.")
+
+    r, delta = 1, 0.5
+
+    fragments = {0: np.random.random(size=(3, 3)), 1: np.random.random(size=(3, 3))}
+    n_frags = len(fragments)
+
+    first_order = ProductFormula(list(range(n_frags)), coeffs=[delta / r] * n_frags) ** r
+    actual = effective_hamiltonian(
+        first_order, fragments, order=2, num_workers=num_workers, backend=backend
+    )
+
+    ham = sum(fragments.values(), np.zeros(shape=(3, 3)))
+    expected = np.zeros(shape=(3, 3), dtype=np.complex128)
+    for j in range(n_frags - 1):
+        for k in range(j + 1, n_frags):
+            expected += (1 / 2) * nested_commutator([fragments[j], fragments[k]])
+    expected *= 1j / r * delta
+
+    assert np.allclose(1j * delta * (expected + ham), actual)

pennylane/labs/tests/trotter_error/product_formulas/test_error.py

@@ -32,8 +32,6 @@
 def test_perturbation_error(backend, parallel_mode, mpi4py_support):
     """Test that perturbation error function runs without errors for different backends."""
 
-    print(f"{backend}, {mpi4py_support}")
-
     if backend in {"mpi4py_pool", "mpi4py_comm"} and not mpi4py_support:
         pytest.skip(f"Skipping test: '{backend}' requires mpi4py, which is not installed.")
 

pennylane/labs/trotter_error/product_formulas/error.py

@@ -23,6 +23,8 @@
 from pennylane.labs.trotter_error.product_formulas.bch import bch_expansion
 from pennylane.labs.trotter_error.product_formulas.product_formula import ProductFormula
 
+# pylint: disable=too-many-arguments, too-many-positional-arguments
+
 
 class _AdditiveIdentity:
     """Only used to initialize accumulators for summing Fragments"""
@@ -39,14 +41,25 @@ def effective_hamiltonian(
     fragments: dict[Hashable, Fragment],
     order: int,
     timestep: float = 1.0,
+    num_workers: int = 1,
+    backend: str = "serial",
 ):
-    r"""Compute the effective Hamiltonian :math:`\hat{H}_{eff} = \hat{H} + \hat{\epsilon}` that corresponds to a given product formula.
+    r"""Compute the effective Hamiltonian :math:`\hat{H}_{eff} = \hat{H} + \hat{\epsilon}` that
+    corresponds to a given product formula.
 
     Args:
-        product_formula (ProductFormula): A product formula used to approximate the time-evolution operator for a Hamiltonian.
-        fragments (Dict[Hashable, :class:`~.pennylane.labs.trotter_error.Fragment`): The fragments that sum to the Hamiltonian. The keys in the dictionary must match the labels used to build the :class:`~.pennylane.labs.trotter_error.ProductFormula` object.
+        product_formula (ProductFormula): A product formula used to approximate the time-evolution
+            operator for a Hamiltonian.
+        fragments (Dict[Hashable, :class:`~.pennylane.labs.trotter_error.Fragment`): The fragments
+            that sum to the Hamiltonian. The keys in the dictionary must match the labels used to
+            build the :class:`~.pennylane.labs.trotter_error.ProductFormula` object.
         order (int): The order of the approximatation.
         timestep (float): The timestep for simulation.
+        num_workers (int): the number of concurrent units used for the computation. Default value is
+            set to 1.
+        backend (string): the executor backend from the list of supported backends.
+            Available options : "mp_pool", "cf_procpool", "cf_threadpool", "serial", "mpi4py_pool",
+            "mpi4py_comm". Default value is set to "serial".
 
     **Example**
 
@@ -78,15 +91,38 @@ def effective_hamiltonian(
         raise ValueError("Fragments do not match product formula")
 
     bch = bch_expansion(product_formula(1j * timestep), order)
-    eff = _AdditiveIdentity()
 
-    for ith_order in bch:
-        for commutator, coeff in ith_order.items():
-            eff += coeff * nested_commutator(_insert_fragments(commutator, fragments))
+    executor = concurrency.backends.get_executor(backend)
+    with executor(max_workers=num_workers) as ex:
+        partial_sum = ex.starmap(
+            _eval_commutator,
+            (
+                (commutator, coeff, fragments)
+                for ith_order in bch
+                for commutator, coeff in ith_order.items()
+            ),
+        )
 
+    eff = _AdditiveIdentity()
+    for term in partial_sum:
+        eff += term
     return eff
 
 
+def _eval_commutator(commutator, coeff, fragments):
+    r"""Computes a commutator after replacing symbols in the commutator with concrete fragments.
+
+    Args:
+        commutator: commutator to be evaluated
+        coeff (complex): coefficient associated with the commutator
+        fragments (dict): dictionary representing a sequence of fragments
+
+    Returns:
+        ndarray: the evaluated form of the commutator
+    """
+    return coeff * nested_commutator(_insert_fragments(commutator, fragments))
+
+
 def _insert_fragments(
     commutator: tuple[Hashable], fragments: dict[Hashable, Fragment]
 ) -> tuple[Fragment]:
@@ -100,7 +136,6 @@ def _insert_fragments(
     )
 
 
-# pylint: disable=too-many-arguments, too-many-positional-arguments
 def perturbation_error(
     product_formula: ProductFormula,
     fragments: dict[Hashable, Fragment],