qadence 1.11.1__py3-none-any.whl → 1.11.3__py3-none-any.whl

qadence/backend.py

@@ -25,7 +25,14 @@ from qadence.measurements import Measurements
 from qadence.mitigations import Mitigations
 from qadence.noise import NoiseHandler
 from qadence.parameters import stringify
-from qadence.types import ArrayLike, BackendName, DiffMode, Endianness, Engine, ParamDictType
+from qadence.types import (
+    ArrayLike,
+    BackendName,
+    DiffMode,
+    Endianness,
+    Engine,
+    ParamDictType,
+)
 
 logger = getLogger(__name__)
 
@@ -54,11 +61,18 @@ class BackendConfiguration:
         conf_msg = ""
         for _field in fields(self):
             if not _field.name.startswith("_"):
-                conf_msg += (
-                    f"Name: {_field.name} - Type: {_field.type} - Default value: {_field.default}\n"
-                )
+                conf_msg += f"Name: {_field.name} - Type: {_field.type} - Current value: {getattr(self, _field.name)} - Default value: {_field.default}\n"
         return conf_msg
 
+    def change_config(self, new_config: dict) -> None:
+        """Change configuration with the input."""
+
+        for key, value in new_config.items():
+            if hasattr(self, key):
+                setattr(self, key, value)
+            else:
+                raise ValueError(f"Warning: '{key}' is not a valid configuration attribute.")
+
     @classmethod
     def from_dict(cls, values: dict) -> BackendConfiguration:
         field_names = {field.name for field in fields(cls)}
@@ -208,7 +222,7 @@ class Backend(ABC):
         if observable is not None:
             observable = observable if isinstance(observable, list) else [observable]
             conv_obs = []
-            obs_embedding_fn_list = []
+            obs_embedding_fns = []
 
             for obs in observable:
                 obs = check_observable(obs)
@@ -217,13 +231,18 @@ class Backend(ABC):
                     c_obs.abstract, self.config._use_gate_params, self.engine
                 )
                 params.update(obs_params)
-                obs_embedding_fn_list.append(obs_embedding_fn)
+                obs_embedding_fns.append(obs_embedding_fn)
                 conv_obs.append(c_obs)
 
             def embedding_fn_dict(a: dict, b: dict) -> dict:
-                embedding_dict = circ_embedding_fn(a, b)
-                for o in obs_embedding_fn_list:
-                    embedding_dict.update(o(a, b))
+                if "circuit" in b or "observables" in b:
+                    embedding_dict = {"circuit": circ_embedding_fn(a, b), "observables": dict()}
+                    for obs_embedding_fn in obs_embedding_fns:
+                        embedding_dict["observables"].update(obs_embedding_fn(a, b))
+                else:
+                    embedding_dict = circ_embedding_fn(a, b)
+                    for obs_embedding_fn in obs_embedding_fns:
+                        embedding_dict.update(obs_embedding_fn(a, b))
                 return embedding_dict
 
             return Converted(conv_circ, conv_obs, embedding_fn_dict, params)
@@ -309,7 +328,11 @@ class Backend(ABC):
         raise NotImplementedError
 
     @abstractmethod
-    def assign_parameters(self, circuit: ConvertedCircuit, param_values: dict[str, Tensor]) -> Any:
+    def assign_parameters(
+        self,
+        circuit: ConvertedCircuit,
+        param_values: dict[str, Tensor] | dict[str, dict[str, Tensor]],
+    ) -> Any:
         raise NotImplementedError
 
     @staticmethod

qadence/backends/agpsr_utils.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+from functools import lru_cache
+from typing import Callable
+
+import numpy as np
+import torch
+from scipy.optimize import minimize
+from torch import Tensor
+
+
+def variance(shifts: Tensor, spectral_gaps: Tensor) -> Tensor:
+    """Calculate the exact variance of deirivative estimation using aGPSR.
+
+    Args:
+        shifts (Tensor): shifts to apply for each spectral gap
+        spectral_gaps (Tensor): tensor containing spectral gap values
+
+    Returns:
+        Tensor: variance tensor
+    """
+
+    # calculate inverse of M (see: https://arxiv.org/pdf/2108.01218.pdf on p. 4 for definitions)
+    M = 4 * torch.sin(torch.outer(torch.as_tensor(shifts), spectral_gaps) / 2)
+    try:
+        # calculate the variance of derivative estimation by solving a linear equation system
+        a = torch.linalg.solve(M, spectral_gaps.reshape(-1, 1))
+        var = 2 * torch.matmul(a.T, a)
+    except RuntimeError:  # matrix M is singulkar
+        # fallback method of variance calculation using inverse matrix
+        M_inv = torch.linalg.pinv(M)
+        a = torch.matmul(spectral_gaps.reshape(1, -1), M_inv)
+        var = 2 * torch.matmul(a, a.T)
+
+    return var
+
+
+@lru_cache
+def calculate_optimal_shifts(
+    n_eqs: int,
+    spectral_gaps: Tensor,
+    lb: float,
+    ub: float,
+) -> Tensor:
+    """Calculates optimal shift values for GPSR algorithm.
+
+    Args:
+        n_eqs (int): number of equations in linear equation system for derivative estimation
+        spectral_gaps (Tensor): tensor containing spectral gap values
+        lb (float): lower bound of optimal shift value search interval
+        ub (float): upper bound of optimal shift value search interval
+
+    Returns:
+        Tensor: optimal shift values
+    """
+    if not (lb and ub):
+        raise ValueError("Both lower and upper bounds of optimization interval must be given.")
+
+    constraints = []
+
+    # specify solution bound constraints
+    for i in range(n_eqs):
+
+        def fn_lb(x, i=i):  # type: ignore [no-untyped-def]
+            return x[i] - lb
+
+        def fn_ub(x, i=i):  # type: ignore [no-untyped-def]
+            return ub - x[i]
+
+        constraints.append({"type": "ineq", "fun": fn_lb})
+        constraints.append({"type": "ineq", "fun": fn_ub})
+
+    # specify constraints for solutions to be unique
+    for i in range(n_eqs - 1, 0, -1):
+        for j in range(i):
+
+            def fn(x, i=i, j=j):  # type: ignore [no-untyped-def]
+                return np.abs(x[i] - x[j]) - 0.02
+
+            constraints.append({"type": "ineq", "fun": fn})
+
+    init_guess = torch.linspace(lb, ub, n_eqs)
+
+    def minimize_variance(
+        var_fn: Callable[[Tensor, Tensor], Tensor],
+    ) -> Tensor:
+        res = minimize(
+            fun=var_fn,
+            x0=init_guess,
+            args=(spectral_gaps,),
+            method="COBYLA",
+            constraints=constraints,
+        )
+        return torch.as_tensor(res.x)
+
+    return minimize_variance(variance)

qadence/backends/api.py

@@ -50,7 +50,14 @@ def backend_factory(
         # Wrap the quantum Backend in a DifferentiableBackend if a diff_mode is passed.
         if diff_mode is not None:
             diff_backend_cls = import_engine(backend_inst.engine)
-            backend_inst = diff_backend_cls(backend=backend_inst, diff_mode=DiffMode(diff_mode))  # type: ignore[operator]
+            psr_args = {
+                "n_eqs": configuration.n_eqs,  # type: ignore [attr-defined]
+                "shift_prefac": configuration.shift_prefac,  # type: ignore [attr-defined]
+                "gap_step": configuration.gap_step,  # type: ignore [attr-defined]
+                "lb": configuration.lb,  # type: ignore [attr-defined]
+                "ub": configuration.ub,  # type: ignore [attr-defined]
+            }
+            backend_inst = diff_backend_cls(backend=backend_inst, diff_mode=DiffMode(diff_mode), **psr_args)  # type: ignore[operator]
         return backend_inst
     except (BackendNotFoundError, EngineNotFoundError, ConfigNotFoundError) as e:
         logger.error(e.msg)

qadence/backends/horqrux/backend.py

@@ -7,7 +7,7 @@ from typing import Any
 
 import jax
 import jax.numpy as jnp
-from horqrux.utils import zero_state
+from horqrux.utils.operator_utils import zero_state
 from jax.typing import ArrayLike
 
 from qadence.backend import Backend as BackendInterface
@@ -27,7 +27,8 @@ from qadence.types import BackendName, Endianness, Engine, ParamDictType
 from qadence.utils import int_to_basis
 
 from .config import Configuration, default_passes
-from .convert_ops import HorqruxCircuit, convert_block, convert_observable
+from .convert_ops import convert_block, convert_observable
+from horqrux.circuit import QuantumCircuit as HorqruxCircuit
 
 logger = getLogger(__name__)
 
@@ -58,7 +59,7 @@ class Backend(BackendInterface):
             circuit = transpile(*passes)(circuit)
         ops = convert_block(circuit.block, n_qubits=circuit.n_qubits, config=self.config)
         return ConvertedCircuit(
-            native=HorqruxCircuit(ops), abstract=circuit, original=original_circ
+            native=HorqruxCircuit(circuit.n_qubits, ops), abstract=circuit, original=original_circ
         )
 
     def observable(self, observable: AbstractBlock, n_qubits: int) -> ConvertedObservable:
@@ -97,7 +98,7 @@ class Backend(BackendInterface):
             state = zero_state(n_qubits)
         else:
             state = horqify(state) if horqify_state else state
-        state = circuit.native.forward(state, param_values)
+        state = circuit.native(state, param_values)
         if endianness != self.native_endianness:
             state = jnp.reshape(state, (1, 2**n_qubits))  # batch_size is always 1
             ls = list(range(2**n_qubits))
@@ -133,19 +134,32 @@ class Backend(BackendInterface):
         Returns:
             A jax.Array of shape (batch_size, n_observables)
         """
-        observable = observable if isinstance(observable, list) else [observable]
-        batch_size = max([arr.size for arr in param_values.values()])
+        observables = observable if isinstance(observable, list) else [observable]
+        if "observables" in param_values or "circuit" in param_values:
+            raise NotImplementedError("The Horqrux backend does not support separated parameters.")
+        else:
+            merged_params = param_values
+            batch_size = max([arr.size for arr in param_values.values()])  # type: ignore[union-attr]
         n_obs = len(observable)
 
         def _expectation(params: ParamDictType) -> ArrayLike:
+            param_circuits = params["circuit"] if "circuit" in params else params
+            param_observables = params["observables"] if "observables" in params else params
             out_state = self.run(
-                circuit, params, state, endianness, horqify_state=True, unhorqify_state=False
+                circuit,
+                param_circuits,
+                state,
+                endianness,
+                horqify_state=True,
+                unhorqify_state=False,
+            )
+            return jnp.array(
+                [observable.native(out_state, param_observables) for observable in observables]
             )
-            return jnp.array([o.native.forward(out_state, params) for o in observable])
 
         if batch_size > 1:  # We vmap for batch_size > 1
-            expvals = jax.vmap(_expectation, in_axes=({k: 0 for k in param_values.keys()},))(
-                uniform_batchsize(param_values)
+            expvals = jax.vmap(_expectation, in_axes=({k: 0 for k in merged_params.keys()},))(
+                uniform_batchsize(merged_params)
             )
         else:
             expvals = _expectation(param_values)

qadence/backends/horqrux/config.py

@@ -32,4 +32,20 @@ def default_passes(config: Configuration) -> list[Callable]:
 
 @dataclass
 class Configuration(BackendConfiguration):
-    pass
+    n_eqs: int | None = None
+    """Number of equations to use in aGPSR calculations."""
+
+    shift_prefac: float = 0.5
+    """Prefactor governing the magnitude of parameter shift values.
+
+    Select smaller value if spectral gaps are large.
+    """
+
+    gap_step: float = 1.0
+    """Step between generated pseudo-gaps when using aGPSR algorithm."""
+
+    lb: float | None = None
+    """Lower bound of optimal shift value search interval."""
+
+    ub: float | None = None
+    """Upper bound of optimal shift value search interval."""

qadence/backends/horqrux/convert_ops.py

@@ -8,11 +8,15 @@ from typing import Any, Callable, Dict
 
 import jax.numpy as jnp
 from horqrux.analog import _HamiltonianEvolution as NativeHorqHEvo
-from horqrux.apply import apply_gate
-from horqrux.parametric import RX, RY, RZ
-from horqrux.primitive import NOT, SWAP, H, I, X, Y, Z
-from horqrux.primitive import Primitive as Gate
-from horqrux.utils import ControlQubits, TargetQubits, inner
+from horqrux.apply import apply_gates
+from horqrux.primitives.parametric import RX, RY, RZ
+from horqrux.primitives.primitive import NOT, SWAP, H, I, X, Y, Z
+from horqrux.primitives.primitive import Primitive as Gate
+from horqrux.composite.sequence import OpSequence as HorqruxSequence
+from horqrux.composite.compose import Scale as HorqScaleGate
+from horqrux.composite.compose import Add as HorqAddGate
+from horqrux.composite.compose import Observable as HorqruxObservable
+from horqrux.utils.operator_utils import ControlQubits, TargetQubits
 from jax import Array
 from jax.scipy.linalg import expm
 from jax.tree_util import register_pytree_node_class
@@ -36,6 +40,7 @@ from qadence.operations import (
     MCRY,
     MCRZ,
     MCZ,
+    CZ,
 )
 from qadence.operations import SWAP as QDSWAP
 from qadence.types import OpName, ParamDictType
@@ -53,6 +58,7 @@ ops_map: Dict[str, Callable] = {
     OpName.CRX: RX,
     OpName.CRY: RY,
     OpName.CRZ: RZ,
+    OpName.CZ: Z,
     OpName.CNOT: NOT,
     OpName.I: I,
     OpName.SWAP: SWAP,
@@ -61,37 +67,6 @@ ops_map: Dict[str, Callable] = {
 supported_gates = list(set(list(ops_map.keys())))
 
 
-@register_pytree_node_class
-@dataclass
-class HorqruxCircuit:
-    operators: list[Gate] = field(default_factory=list)
-
-    def tree_flatten(self) -> tuple[tuple[list[Any]], tuple[()]]:
-        children = (self.operators,)
-        aux_data = ()
-        return (children, aux_data)
-
-    @classmethod
-    def tree_unflatten(cls, aux_data: Any, children: Any) -> Any:
-        return cls(*children, *aux_data)
-
-    def _forward(self, state: Array, values: ParamDictType) -> Array:
-        return reduce(lambda state, gate: gate.forward(state, values), self.operators, state)
-
-    def forward(self, state: Array, values: ParamDictType) -> Array:
-        return self._forward(state, values)
-
-
-@register_pytree_node_class
-@dataclass
-class HorqruxObservable(HorqruxCircuit):
-    def __init__(self, operators: list[Gate]):
-        super().__init__(operators=operators)
-
-    def forward(self, state: Array, values: ParamDictType) -> Array:
-        return jnp.real(inner(state, self._forward(state, values)))
-
-
 def convert_observable(
     block: AbstractBlock, n_qubits: int, config: Configuration
 ) -> HorqruxObservable:
@@ -109,7 +84,7 @@ def convert_block(
     ops = []
     if isinstance(block, CompositeBlock):
         ops = list(flatten(*(convert_block(b, n_qubits, config) for b in block.blocks)))
-        ops = [HorqAddGate(ops)] if isinstance(block, AddBlock) else [HorqruxCircuit(ops)]
+        ops = [HorqAddGate(ops)] if isinstance(block, AddBlock) else [HorqruxSequence(ops)]
     elif isinstance(block, ScaleBlock):
         op = convert_block(block.block, n_qubits, config=config)[0]
         param_name = config.get_param_name(block)[0]
@@ -118,7 +93,7 @@ def convert_block(
         native_op_fn = ops_map[block.name]
         target, control = (
             (block.qubit_support[1], block.qubit_support[0])
-            if isinstance(block, (CNOT, CRX, CRY, CRZ, QDSWAP))
+            if isinstance(block, (CZ, CNOT, CRX, CRY, CRZ, QDSWAP))
             else (block.qubit_support[0], (None,))
         )
         native_gate: Gate
@@ -133,7 +108,7 @@ def convert_block(
                 native_gate = native_op_fn(block.qubit_support[::-1])
             else:
                 native_gate = native_op_fn(target=target, control=control)
-        ops = [HorqOperation(native_gate)]
+        ops = [native_gate]
 
     elif isinstance(block, (MCRX, MCRY, MCRZ, MCZ)):
         block_name = block.name[2:] if block.name.startswith("M") else block.name
@@ -146,7 +121,7 @@ def convert_block(
             native_gate = native_op_fn(param=param, target=target, control=control)
         else:
             native_gate = native_op_fn(target, control)
-        ops = [HorqOperation(native_gate)]
+        ops = [native_gate]
     elif isinstance(block, TimeEvolutionBlock):
         ops = [HorqHamiltonianEvolution(block, config)]
     else:
@@ -155,58 +130,6 @@ def convert_block(
     return ops
 
 
-@register_pytree_node_class
-class HorqAddGate(HorqruxCircuit):
-    def __init__(self, operations: list[Gate]):
-        self.operators = operations
-        self.name = "Add"
-
-    def forward(self, state: Array, values: ParamDictType = {}) -> Array:
-        return reduce(add, (gate.forward(state, values) for gate in self.operators))
-
-    def __repr__(self) -> str:
-        return self.name + f"({self.operators})"
-
-
-@register_pytree_node_class
-@dataclass
-class HorqOperation:
-    def __init__(self, native_gate: Gate):
-        self.native_gate = native_gate
-
-    def forward(self, state: Array, values: ParamDictType) -> Array:
-        return apply_gate(state, self.native_gate, values)
-
-    def tree_flatten(self) -> tuple[tuple[Gate], tuple[()]]:
-        children = (self.native_gate,)
-        aux_data = ()
-        return (children, aux_data)
-
-    @classmethod
-    def tree_unflatten(cls, aux_data: Any, children: Any) -> Any:
-        return cls(*children, *aux_data)
-
-
-@register_pytree_node_class
-@dataclass
-class HorqScaleGate:
-    def __init__(self, gate: HorqOperation, parameter_name: str):
-        self.gate = gate
-        self.parameter: str = parameter_name
-
-    def forward(self, state: Array, values: ParamDictType) -> Array:
-        return jnp.array(values[self.parameter]) * self.gate.forward(state, values)
-
-    def tree_flatten(self) -> tuple[tuple[HorqOperation], tuple[str]]:
-        children = (self.gate,)
-        aux_data = (self.parameter,)
-        return (children, aux_data)
-
-    @classmethod
-    def tree_unflatten(cls, aux_data: Any, children: Any) -> Any:
-        return cls(*children, *aux_data)
-
-
 @register_pytree_node_class
 @dataclass
 class HorqHamiltonianEvolution(NativeHorqHEvo):
@@ -216,7 +139,7 @@ class HorqHamiltonianEvolution(NativeHorqHEvo):
         config: Configuration,
     ):
         super().__init__("I", block.qubit_support, (None,))
-        self.qubit_support = block.qubit_support
+        self._qubit_support = block.qubit_support
         self.param_names = config.get_param_name(block)
         self.block = block
         self.hmat: Array
@@ -224,7 +147,7 @@ class HorqHamiltonianEvolution(NativeHorqHEvo):
         if isinstance(block.generator, AbstractBlock) and not block.generator.is_parametric:
             hmat = block_to_jax(
                 block.generator,
-                qubit_support=self.qubit_support,
+                qubit_support=self._qubit_support,
                 use_full_support=False,
             )
             self.hmat = hmat
@@ -236,7 +159,7 @@ class HorqHamiltonianEvolution(NativeHorqHEvo):
                 hmat = block_to_jax(
                     block.generator,  # type: ignore[arg-type]
                     values=values,
-                    qubit_support=self.qubit_support,
+                    qubit_support=self._qubit_support,
                     use_full_support=False,
                 )
                 return hmat
@@ -249,12 +172,12 @@ class HorqHamiltonianEvolution(NativeHorqHEvo):
         """The evolved operator given current parameter values for generator and time evolution."""
         return expm(self._hamiltonian(self, values) * (-1j * self._time_evolution(values)))
 
-    def forward(
+    def __call__(
         self,
         state: Array,
         values: dict[str, Array],
     ) -> Array:
-        return apply_gate(state, self, values)
+        return apply_gates(state, self, values)
 
     def tree_flatten(self) -> tuple[tuple[NativeHorqHEvo], tuple]:
         children = (self,)

qadence/backends/jax_utils.py

@@ -19,6 +19,7 @@ from qadence.blocks import (
 )
 from qadence.blocks.block_to_tensor import _gate_parameters
 from qadence.types import Endianness, ParamDictType
+from qadence.utils import merge_separate_params
 
 
 def jarr_to_tensor(arr: Array, dtype: Any = cdouble) -> Tensor:
@@ -52,9 +53,11 @@ def horqify(state: Array) -> Array:
 
 
 def uniform_batchsize(param_values: ParamDictType) -> ParamDictType:
-    max_batch_size = max(p.size for p in param_values.values())
+    if "observables" in param_values or "circuit" in param_values:
+        param_values = merge_separate_params(param_values)
+    max_batch_size = max(p.size for p in param_values.values())  # type: ignore[union-attr]
     batched_values = {
-        k: (v if v.size == max_batch_size else v.repeat(max_batch_size))
+        k: (v if v.size == max_batch_size else v.repeat(max_batch_size))  # type: ignore[union-attr]
         for k, v in param_values.items()
     }
     return batched_values

qadence/backends/{gpsr.py → parameter_shift_rules.py}

@@ -8,32 +8,45 @@ from torch import Tensor
 
 from qadence.types import PI
 from qadence.utils import _round_complex
+from qadence.backends.agpsr_utils import calculate_optimal_shifts
 
 
-def general_psr(spectrum: Tensor, n_eqs: int | None = None, shift_prefac: float = 0.5) -> Callable:
+def general_psr(
+    spectrum: Tensor,
+    n_eqs: int | None = None,
+    shift_prefac: float | None = 0.5,
+    gap_step: float = 1.0,
+    lb: float | None = None,
+    ub: float | None = None,
+) -> Callable:
     """Define whether single_gap_psr or multi_gap_psr is used.
 
     Args:
         spectrum (Tensor): Spectrum of the operation we apply PSR onto.
-        n_eqs (int | None, optional): Number of equations. Defaults to None.
-            If provided, we keep the n_eqs higher spectral gaps.
-        shift_prefac (float, optional): Shift prefactor. Defaults to 0.5.
+        n_eqs (int | None): Number of equations. Defaults to None.
+            If provided, aGPSR algorithm is effectively used.
+        shift_prefac (float | None): prefactor governing the magnitude of parameter shift values -
+            select smaller value if spectral gaps are large. Defaults to 0.5.
+        gap_step (float): Step between generated pseudo-gaps when using aGPSR algorithm. Defaults to 1.0.
+        lb (float | None): Lower bound of optimal shift value search interval. Defaults to None.
+        ub (float | None): Upper bound of optimal shift value search interval. Defaults to None.
 
     Returns:
-        Callable: single_gap_psr or multi_gap_psr function for
-            concerned operation.
+        Callable: single_gap_psr or multi_gap_psr function for concerned operation.
     """
+
     diffs = _round_complex(spectrum - spectrum.reshape(-1, 1))
-    sorted_unique_spectral_gaps = torch.unique(torch.abs(torch.tril(diffs)))
+    orig_unique_spectral_gaps = torch.unique(torch.abs(torch.tril(diffs)))
 
     # We have to filter out zeros
-    sorted_unique_spectral_gaps = sorted_unique_spectral_gaps[sorted_unique_spectral_gaps > 0]
-    n_eqs = (
-        len(sorted_unique_spectral_gaps)
-        if n_eqs is None
-        else min(n_eqs, len(sorted_unique_spectral_gaps))
-    )
-    sorted_unique_spectral_gaps = torch.tensor(list(sorted_unique_spectral_gaps)[:n_eqs])
+    orig_unique_spectral_gaps = orig_unique_spectral_gaps[orig_unique_spectral_gaps > 0]
+
+    if n_eqs is None:  # GPSR case
+        n_eqs = len(orig_unique_spectral_gaps)
+        sorted_unique_spectral_gaps = orig_unique_spectral_gaps
+    else:  # aGPSR case
+        sorted_unique_spectral_gaps = torch.arange(0, n_eqs) * gap_step
+        sorted_unique_spectral_gaps[0] = 0.001
 
     if n_eqs == 1:
         return partial(
@@ -46,6 +59,8 @@ def general_psr(spectrum: Tensor, n_eqs: int | None = None, shift_prefac: float
             multi_gap_psr,
             spectral_gaps=sorted_unique_spectral_gaps,
             shift_prefac=shift_prefac,
+            lb=lb,
+            ub=ub,
         )
 
 
@@ -60,8 +75,7 @@ def single_gap_psr(
 
     Args:
         expectation_fn (Callable[[dict[str, Tensor]], Tensor]): backend-dependent function
-        to calculate expectation value
-
+            to calculate expectation value
         param_dict (dict[str, Tensor]): dict storing parameters of parameterized blocks
         param_name (str): name of parameter with respect to that differentiation is performed
 
@@ -93,19 +107,22 @@ def multi_gap_psr(
     param_dict: dict[str, Tensor],
     param_name: str,
     spectral_gaps: Tensor,
-    shift_prefac: float = 0.5,
+    shift_prefac: float | None = 0.5,
+    lb: float | None = None,
+    ub: float | None = None,
 ) -> Tensor:
     """Implements multi-gap multi-qubit GPSR rule.
 
     Args:
         expectation_fn (Callable[[dict[str, Tensor]], Tensor]): backend-dependent function
-        to calculate expectation value
-
+            to calculate expectation value
         param_dict (dict[str, Tensor]): dict storing parameters values of parameterized blocks
         param_name (str): name of parameter with respect to that differentiation is performed
-        spectral_gaps (Tensor): tensor containing spectral gap values
+            spectral_gaps (Tensor): tensor containing spectral gap values
         shift_prefac (float): prefactor governing the magnitude of parameter shift values -
-        select smaller value if spectral gaps are large
+            select smaller value if spectral gaps are large
+        lb (float): lower bound of optimal shift value search interval
+        ub (float): upper bound of optimal shift value search interval
 
     Returns:
         Tensor: tensor containing derivative values
@@ -113,10 +130,15 @@ def multi_gap_psr(
     n_eqs = len(spectral_gaps)
     batch_size = max(t.size(0) for t in param_dict.values())
 
-    # get shift values
-    shifts = shift_prefac * torch.linspace(
-        PI / 2 - PI / 4, PI / 2 + PI / 5, n_eqs
-    )  # breaking the symmetry of sampling range around PI/2
+    # get shift values - values minimize the variance of expectation
+    if shift_prefac is not None:
+        # Set shift values manually by breaking the symmetry of sampling range
+        # around PI/2 to reduce the possibility that M is singular
+        shifts = shift_prefac * torch.linspace(PI / 2 - PI / 4, PI / 2 + PI / 5, n_eqs)
+    else:
+        # calculate optimal shift values
+        shifts = calculate_optimal_shifts(n_eqs, spectral_gaps, lb, ub)
+
     device = torch.device("cpu")
     try:
         device = [v.device for v in param_dict.values()][0]
@@ -127,7 +149,7 @@ def multi_gap_psr(
     # calculate F vector and M matrix
     # (see: https://arxiv.org/pdf/2108.01218.pdf on p. 4 for definitions)
     F = []
-    M = torch.empty((n_eqs, n_eqs)).to(device=device)
+    M = 4 * torch.sin(torch.outer(shifts, spectral_gaps) / 2).to(device=device)
     n_obs = 1
     for i in range(n_eqs):
         # + shift
@@ -142,10 +164,6 @@ def multi_gap_psr(
 
         F.append((f_plus - f_minus))
 
-        # calculate M matrix
-        for j in range(n_eqs):
-            M[i, j] = 4 * torch.sin(shifts[i] * spectral_gaps[j] / 2)
-
     # get number of observables from expectation value tensor
     if f_plus.numel() > 1:
         batch_size = F[0].shape[0]