Add concurrency support for effective_hamiltonian in labs

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

@@ -277,3 +277,39 @@ 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.zeros(shape=(3, 3)), 1: np.zeros(shape=(3, 3)), 2: np.zeros(shape=(3, 3))}
+
+    n_frags = len(fragments)
+
+    first_order = ProductFormula(list(range(n_frags)), coeffs=[delta / r] * n_frags) ** r
+    ham = effective_hamiltonian(
+        first_order, fragments, order=2, num_workers=num_workers, backend=backend
+    )
+
+    assert isinstance(ham, np.ndarray)

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,13 +91,33 @@ 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()
+            ],
+        )
+
+    return sum(partial_sum)
 
-    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(
@@ -100,7 +133,6 @@ def _insert_fragments(
     )
 
 
-# pylint: disable=too-many-arguments, too-many-positional-arguments
 def perturbation_error(
     product_formula: ProductFormula,
     fragments: dict[Hashable, Fragment],