From acf130c0dc1361953c6bf82a4b35ab08691bf64f Mon Sep 17 00:00:00 2001
From: Purva Thakre <66048318+purva-thakre@users.noreply.github.com>
Date: Mon, 1 Jul 2024 12:50:38 -0500
Subject: [PATCH] Noise Scaling for LRE (#2347)

* add files associated with lre scaling

* change docstring for zne/layer_scaling

* num layers without measurements

* add scale factor vectors required for multivariate extrapolation

* all private functions required for lre

* lre noise scaling only for cirq circuits

* mypy

* change import cirq to from cirq import

* make num_chunks optional

* error message: degree and fold_multiplier

* init + apidoc

* docstring

* vincent's feedback: quick fixes

* add admonition

* decorator

* undo decorator

* cleanup

* alessandro's feedback: check for negative num_chunks, tests for chunking function, enumerate

* add raises error details to the docstrings

* alessandro's feedback

* alessandro's feedback - get rid of expected_chunks

* vincent's comments: cleanup docstrings
---
 docs/source/apidoc.md                         |   8 +
 docs/source/refs.bib                          |  10 +
 mitiq/lre/__init__.py                         |   8 +
 .../multivariate_scaling/layerwise_folding.py | 210 ++++++++++++
 mitiq/lre/tests/test_layerwise_folding.py     | 301 ++++++++++++++++++
 mitiq/zne/scaling/layer_scaling.py            |  17 +
 6 files changed, 554 insertions(+)
 create mode 100644 mitiq/lre/__init__.py
 create mode 100644 mitiq/lre/multivariate_scaling/layerwise_folding.py
 create mode 100644 mitiq/lre/tests/test_layerwise_folding.py

diff --git a/docs/source/apidoc.md b/docs/source/apidoc.md
index f0a572f2dc..be018e7c7a 100644
--- a/docs/source/apidoc.md
+++ b/docs/source/apidoc.md
@@ -87,6 +87,13 @@ See Ref. {cite}`Czarnik_2021_Quantum` for more details on these methods.
    :members:
 ```
 
+### Layerwise Richardson Extrapolation
+
+```{eval-rst}
+.. automodule:: mitiq.lre.multivariate_scaling.layerwise_folding
+   :members:
+```
+
 ### Pauli Twirling
 
 ```{eval-rst}
@@ -116,6 +123,7 @@ See Ref. {cite}`Czarnik_2021_Quantum` for more details on these methods.
    :members:
 ```
 
+
 #### Learning-based PEC
 
 ```{eval-rst}
diff --git a/docs/source/refs.bib b/docs/source/refs.bib
index 1280ac42f7..f7b6762353 100644
--- a/docs/source/refs.bib
+++ b/docs/source/refs.bib
@@ -731,6 +731,16 @@ @misc{Russo_2022_Testing
   year      = {2022},
   copyright = {arXiv.org perpetual, non-exclusive license},
 }
+
+@misc{Russo_2024_LRE,
+      title={Quantum error mitigation by layerwise Richardson extrapolation}, 
+      author={Vincent Russo and Andrea Mari},
+      year={2024},
+      eprint={2402.04000},
+      archivePrefix={arXiv},
+      primaryClass={quant-ph}
+}
+
 # Letter S
 
 @article{Sagastizabal_2019_PRA,
diff --git a/mitiq/lre/__init__.py b/mitiq/lre/__init__.py
new file mode 100644
index 0000000000..ba7acdf327
--- /dev/null
+++ b/mitiq/lre/__init__.py
@@ -0,0 +1,8 @@
+# Copyright (C) Unitary Fund
+#
+# This source code is licensed under the GPL license (v3) found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Methods for scaling noise in circuits by layers and using multivariate extrapolation."""
+
+from mitiq.lre.multivariate_scaling.layerwise_folding import multivariate_layer_scaling
\ No newline at end of file
diff --git a/mitiq/lre/multivariate_scaling/layerwise_folding.py b/mitiq/lre/multivariate_scaling/layerwise_folding.py
new file mode 100644
index 0000000000..9bd883fca2
--- /dev/null
+++ b/mitiq/lre/multivariate_scaling/layerwise_folding.py
@@ -0,0 +1,210 @@
+# Copyright (C) Unitary Fund
+#
+# This source code is licensed under the GPL license (v3) found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Functions for layerwise folding of input circuits to allow for multivariate
+extrapolation as defined in :cite:`Russo_2024_LRE`.
+"""
+
+import itertools
+from copy import deepcopy
+from typing import Any, Callable, List, Optional, Tuple
+
+import numpy as np
+from cirq import Circuit
+
+from mitiq import QPROGRAM
+from mitiq.utils import _append_measurements, _pop_measurements
+from mitiq.zne.scaling import fold_gates_at_random
+from mitiq.zne.scaling.folding import _check_foldable
+
+
+def _get_num_layers_without_measurements(input_circuit: Circuit) -> int:
+    """Checks if the circuit has non-terminal measurements and returns the
+    number of layers in the input circuit without the terminal measurements.
+
+        Args:
+            input_circuit: Circuit of interest.
+
+        Returns:
+            num_layers: the number of layers in the input circuit without the
+                terminal measurements.
+
+    """
+
+    _check_foldable(input_circuit)
+    circuit = deepcopy(input_circuit)
+    _pop_measurements(circuit)
+    return len(circuit)
+
+
+def _get_chunks(
+    input_circuit: Circuit, num_chunks: Optional[int] = None
+) -> List[Circuit]:
+    """Splits a circuit into approximately equal chunks.
+
+    Adapted from:
+    https://stackoverflow.com/questions/2130016/splitting-a-list-into-n-parts-of-approximately-equal-length
+
+        Args:
+            input_circuit: Circuit of interest.
+            num_chunks: Number of desired approximately equal chunks,
+                * when num_chunks == num_layers, the original circuit is
+                    returned.
+                * when num_chunks == 1, the entire circuit is chunked into 1
+                    layer.
+        Returns:
+            split_circuit: Circuit of interest split into approximately equal
+                chunks.
+
+        Raises:
+            ValueError:
+                When the number of chunks for the input circuit is larger than
+                    the number of layers in the input circuit.
+
+            ValueError:
+                When the number of chunks is less than 1.
+
+    """
+    num_layers = _get_num_layers_without_measurements(input_circuit)
+    if num_chunks is None:
+        num_chunks = num_layers
+
+    if num_chunks < 1:
+        raise ValueError(
+            "Number of chunks should be greater than or equal to 1."
+        )
+
+    if num_chunks > num_layers:
+        raise ValueError(
+            f"Number of chunks {num_chunks} cannot be greater than the number"
+            f" of layers {num_layers}."
+        )
+
+    k, m = divmod(num_layers, num_chunks)
+    return [
+        input_circuit[i * k + min(i, m) : (i + 1) * k + min(i + 1, m)]
+        for i in range(num_chunks)
+    ]
+
+
+def _get_scale_factor_vectors(
+    input_circuit: Circuit,
+    degree: int,
+    fold_multiplier: int,
+    num_chunks: Optional[int] = None,
+) -> List[Tuple[Any, ...]]:
+    """Returns the patterned scale factor vectors required for multivariate
+    extrapolation.
+
+        Args:
+            input_circuit: Circuit to be scaled.
+            degree: Degree of the multivariate polynomial.
+            fold_multiplier: Scaling gap required by unitary folding.
+            num_chunks: Number of desired approximately equal chunks.
+
+        Returns:
+            scale_factor_vectors: A vector of scale factors where each
+                component in the vector corresponds to the layer in the input
+                circuit.
+    """
+
+    circuit_chunks = _get_chunks(input_circuit, num_chunks)
+    num_layers = len(circuit_chunks)
+
+    # Find the exponents of all the monomial terms required for the folding
+    # pattern.
+    pattern_full = []
+    for i in range(degree + 1):
+        for j in itertools.combinations_with_replacement(range(num_layers), i):
+            pattern = np.zeros(num_layers, dtype=int)
+            # Get the monomial terms in graded lexicographic order.
+            for index in j:
+                pattern[index] += 1
+            # Use the fold multiplier on the folding pattern to determine which
+            # layers will be scaled.
+            pattern_full.append(tuple(fold_multiplier * pattern))
+
+    # Get the scale factor vectors.
+    # The layers are scaled as 2n+1 due to unitary folding.
+    return [
+        tuple(2 * num_folds + 1 for num_folds in pattern)
+        for pattern in pattern_full
+    ]
+
+
+def multivariate_layer_scaling(
+    input_circuit: Circuit,
+    degree: int,
+    fold_multiplier: int,
+    num_chunks: Optional[int] = None,
+    folding_method: Callable[
+        [QPROGRAM, float], QPROGRAM
+    ] = fold_gates_at_random,
+) -> List[Circuit]:
+    r"""
+    Defines the noise scaling function required for Layerwise Richardson
+    Extrapolation as defined in :cite:`Russo_2024_LRE`.
+
+    Note that this method only works for the multivariate extrapolation
+    methods. It does not allows a user to choose which layers in the input
+    circuit will be scaled.
+
+    .. seealso::
+
+        If you would prefer to choose the layers for unitary
+        folding, use :func:`mitiq.zne.scaling.layer_scaling.get_layer_folding`
+        instead.
+
+    Args:
+        input_circuit: Circuit to be scaled.
+        degree: Degree of the multivariate polynomial.
+        fold_multiplier: Scaling gap required by unitary folding.
+        num_chunks: Number of desired approximately equal chunks. When the
+            number of chunks is the same as the layers in the input circuit,
+            the input circuit is unchanged.
+        folding_method: Unitary folding method. Default is
+            :func:`fold_gates_at_random`.
+
+    Returns:
+        Multiple folded variations of the input circuit.
+
+    Raises:
+        ValueError:
+            When the degree for the multinomial is not greater than or
+                equal to 1; when the fold multiplier to scale the circuit is
+                greater than/equal to 1; when the number of chunks for a
+                large circuit is 0 or when the number of chunks in a circuit is
+                greater than the number of layers in the input circuit.
+
+    """
+    if degree < 1:
+        raise ValueError(
+            "Multinomial degree must be greater than or equal to 1."
+        )
+    if fold_multiplier < 1:
+        raise ValueError("Fold multiplier must be greater than or equal to 1.")
+    circuit_copy = deepcopy(input_circuit)
+    terminal_measurements = _pop_measurements(circuit_copy)
+
+    scaling_pattern = _get_scale_factor_vectors(
+        circuit_copy, degree, fold_multiplier, num_chunks
+    )
+
+    chunks = _get_chunks(circuit_copy, num_chunks)
+
+    multiple_folded_circuits = []
+    for scale_factor_vector in scaling_pattern:
+        folded_circuit = Circuit()
+        for chunk, scale_factor in zip(chunks, scale_factor_vector):
+            if scale_factor == 1:
+                folded_circuit += chunk
+            else:
+                chunks_circ = Circuit(chunk)
+                folded_chunk_circ = folding_method(chunks_circ, scale_factor)
+                folded_circuit += folded_chunk_circ
+        _append_measurements(folded_circuit, terminal_measurements)
+        multiple_folded_circuits.append(folded_circuit)
+
+    return multiple_folded_circuits
diff --git a/mitiq/lre/tests/test_layerwise_folding.py b/mitiq/lre/tests/test_layerwise_folding.py
new file mode 100644
index 0000000000..b1870450d8
--- /dev/null
+++ b/mitiq/lre/tests/test_layerwise_folding.py
@@ -0,0 +1,301 @@
+# Copyright (C) Unitary Fund
+#
+# This source code is licensed under the GPL license (v3) found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Unit tests for scaling noise by unitary folding of layers in the input
+circuit to allow for multivariate extrapolation."""
+
+from copy import deepcopy
+
+import pytest
+from cirq import Circuit, LineQubit, ops
+
+from mitiq.lre.multivariate_scaling.layerwise_folding import (
+    _get_chunks,
+    _get_num_layers_without_measurements,
+    _get_scale_factor_vectors,
+    multivariate_layer_scaling,
+)
+
+qreg1 = LineQubit.range(3)
+test_circuit1 = Circuit(
+    [ops.H.on_each(*qreg1)],
+    [ops.CNOT.on(qreg1[0], qreg1[1])],
+    [ops.X.on(qreg1[2])],
+    [ops.TOFFOLI.on(*qreg1)],
+)
+
+test_circuit1_with_measurements = deepcopy(test_circuit1)
+test_circuit1_with_measurements.append(ops.measure_each(*qreg1))
+
+
+def test_multivariate_layerwise_scaling():
+    """Checks if multiple scaled circuits are returned to fit the required
+    folding pattern for multivariate extrapolation."""
+    multiple_scaled_circuits = multivariate_layer_scaling(
+        test_circuit1, 2, 2, 3
+    )
+
+    assert len(multiple_scaled_circuits) == 10
+    folding_pattern = [
+        (1, 1, 1),
+        (5, 1, 1),
+        (1, 5, 1),
+        (1, 1, 5),
+        (9, 1, 1),
+        (5, 5, 1),
+        (5, 1, 5),
+        (1, 9, 1),
+        (1, 5, 5),
+        (1, 1, 9),
+    ]
+
+    for i, scale_factor_vector in enumerate(folding_pattern):
+        scale_layer1, scale_layer2, scale_layer3 = scale_factor_vector
+        expected_circuit = Circuit(
+            [ops.H.on_each(*qreg1)] * scale_layer1,
+            [ops.CNOT.on(qreg1[0], qreg1[1]), ops.X.on(qreg1[2])]
+            * scale_layer2,
+            [ops.TOFFOLI.on(*qreg1)] * scale_layer3,
+        )
+        assert expected_circuit == multiple_scaled_circuits[i]
+
+
+@pytest.mark.parametrize(
+    "test_input, expected",
+    [(test_circuit1, 3), (test_circuit1_with_measurements, 3)],
+)
+def test_get_num_layers(test_input, expected):
+    """Verifies function works as expected."""
+    calculated_num_layers = _get_num_layers_without_measurements(test_input)
+
+    assert calculated_num_layers == expected
+
+
+@pytest.mark.parametrize(
+    "test_input, test_chunks",
+    [
+        (test_circuit1 + test_circuit1 + test_circuit1, 3),
+        (test_circuit1 + test_circuit1 + test_circuit1, 1),
+        (test_circuit1 + test_circuit1 + test_circuit1, 5),
+    ],
+)
+def test_get_num_chunks(test_input, test_chunks):
+    """Verifies the chunking function works as expected."""
+    assert test_chunks == len(_get_chunks(test_input, test_chunks))
+
+
+def test_layers_with_chunking():
+    """Checks the order of moments in the input circuit is unchanged with
+    chunking."""
+
+    test_circuit = test_circuit1 + test_circuit1 + test_circuit1
+    calculated_circuit_chunks = _get_chunks(test_circuit, 4)
+    expected_chunks = [
+        test_circuit[0:3],
+        test_circuit[3:5],
+        test_circuit[5:7],
+        test_circuit[7:],
+    ]
+    assert calculated_circuit_chunks == expected_chunks
+
+
+@pytest.mark.parametrize(
+    "test_input, degree, test_fold_multiplier, expected_scale_factor_vectors",
+    [
+        (test_circuit1, 1, 1, [(1, 1, 1), (3, 1, 1), (1, 3, 1), (1, 1, 3)]),
+        (
+            test_circuit1,
+            2,
+            1,
+            [
+                (1, 1, 1),
+                (3, 1, 1),
+                (1, 3, 1),
+                (1, 1, 3),
+                (5, 1, 1),
+                (3, 3, 1),
+                (3, 1, 3),
+                (1, 5, 1),
+                (1, 3, 3),
+                (1, 1, 5),
+            ],
+        ),
+        (
+            test_circuit1,
+            2,
+            2,
+            [
+                (1, 1, 1),
+                (5, 1, 1),
+                (1, 5, 1),
+                (1, 1, 5),
+                (9, 1, 1),
+                (5, 5, 1),
+                (5, 1, 5),
+                (1, 9, 1),
+                (1, 5, 5),
+                (1, 1, 9),
+            ],
+        ),
+        (
+            test_circuit1,
+            2,
+            3,
+            [
+                (1, 1, 1),
+                (7, 1, 1),
+                (1, 7, 1),
+                (1, 1, 7),
+                (13, 1, 1),
+                (7, 7, 1),
+                (7, 1, 7),
+                (1, 13, 1),
+                (1, 7, 7),
+                (1, 1, 13),
+            ],
+        ),
+        (
+            test_circuit1_with_measurements,
+            1,
+            1,
+            [(1, 1, 1), (3, 1, 1), (1, 3, 1), (1, 1, 3)],
+        ),
+        (
+            test_circuit1_with_measurements,
+            2,
+            1,
+            [
+                (1, 1, 1),
+                (3, 1, 1),
+                (1, 3, 1),
+                (1, 1, 3),
+                (5, 1, 1),
+                (3, 3, 1),
+                (3, 1, 3),
+                (1, 5, 1),
+                (1, 3, 3),
+                (1, 1, 5),
+            ],
+        ),
+        (
+            test_circuit1_with_measurements,
+            2,
+            2,
+            [
+                (1, 1, 1),
+                (5, 1, 1),
+                (1, 5, 1),
+                (1, 1, 5),
+                (9, 1, 1),
+                (5, 5, 1),
+                (5, 1, 5),
+                (1, 9, 1),
+                (1, 5, 5),
+                (1, 1, 9),
+            ],
+        ),
+        (
+            test_circuit1_with_measurements,
+            2,
+            3,
+            [
+                (1, 1, 1),
+                (7, 1, 1),
+                (1, 7, 1),
+                (1, 1, 7),
+                (13, 1, 1),
+                (7, 7, 1),
+                (7, 1, 7),
+                (1, 13, 1),
+                (1, 7, 7),
+                (1, 1, 13),
+            ],
+        ),
+    ],
+)
+def test_get_scale_factor_vectors_no_chunking(
+    test_input, degree, test_fold_multiplier, expected_scale_factor_vectors
+):
+    """Verifies vectors of scale factors are calculated accurately."""
+    calculated_scale_factor_vectors = _get_scale_factor_vectors(
+        test_input, degree, test_fold_multiplier
+    )
+
+    assert calculated_scale_factor_vectors == expected_scale_factor_vectors
+
+
+@pytest.mark.parametrize(
+    "test_input, degree, test_fold_multiplier, test_chunks, expected_size",
+    [
+        (test_circuit1, 1, 1, 2, 3),
+        (test_circuit1, 2, 1, 3, 10),
+        (test_circuit1, 2, 3, 2, 6),
+    ],
+)
+def test_get_scale_factor_vectors_with_chunking(
+    test_input, degree, test_fold_multiplier, test_chunks, expected_size
+):
+    """Verifies vectors of scale factors are calculated accurately."""
+    calculated_scale_factor_vectors = _get_scale_factor_vectors(
+        test_input, degree, test_fold_multiplier, test_chunks
+    )
+
+    assert len(calculated_scale_factor_vectors) == expected_size
+
+
+@pytest.mark.parametrize(
+    "test_input, num_chunks, error_msg",
+    [
+        (
+            test_circuit1,
+            0,
+            "Number of chunks should be greater than or equal to 1.",
+        ),
+        (
+            test_circuit1,
+            5,
+            "Number of chunks 5 cannot be greater than the number of layers"
+            " 3.",
+        ),
+        (
+            test_circuit1,
+            -1,
+            "Number of chunks should be greater than or equal to 1.",
+        ),
+    ],
+)
+def test_invalid_num_chunks(test_input, num_chunks, error_msg):
+    """Ensures that the number of intended chunks in the input circuit raises
+    an error for an invalid value."""
+    with pytest.raises(ValueError, match=error_msg):
+        _get_scale_factor_vectors(test_input, 2, 2, num_chunks)
+
+
+@pytest.mark.parametrize(
+    "test_input, test_degree, test_fold_multiplier, error_msg",
+    [
+        (
+            test_circuit1,
+            0,
+            1,
+            "Multinomial degree must be greater than or equal to 1.",
+        ),
+        (
+            test_circuit1,
+            1,
+            0,
+            "Fold multiplier must be greater than or equal to 1.",
+        ),
+    ],
+)
+def test_invalid_degree_fold_multiplier(
+    test_input, test_degree, test_fold_multiplier, error_msg
+):
+    """Ensures that the args for the main noise scaling function raise
+    an error for an invalid value."""
+    with pytest.raises(ValueError, match=error_msg):
+        multivariate_layer_scaling(
+            test_input, test_degree, test_fold_multiplier
+        )
diff --git a/mitiq/zne/scaling/layer_scaling.py b/mitiq/zne/scaling/layer_scaling.py
index 24b47dd481..89f2c2c1eb 100644
--- a/mitiq/zne/scaling/layer_scaling.py
+++ b/mitiq/zne/scaling/layer_scaling.py
@@ -24,6 +24,23 @@ def layer_folding(
 ) -> cirq.Circuit:
     """Applies a variable amount of folding to select layers of a circuit.
 
+    Note that this method only works for the univariate extrapolation methods.
+    It allows a user to choose which layers in the input circuit will be
+    scaled.
+
+    .. seealso::
+
+        If you would prefer to
+        use a multivariate extrapolation method for unitary
+        folding, use
+        :func:`mitiq.lre.multivariate_scaling.layerwise_folding` instead.
+
+        The layerwise folding required for multivariate extrapolation is
+        different as the layers in the input circuit have to be scaled in
+        a specific pattern. The required specific pattern for multivariate
+        extrapolation does not allow a user to provide a choice of which
+        layers to fold.
+
     Args:
         circuit: The input circuit.
         layers_to_fold: A list with the index referring to the layer number,