qadence 1.1.1__py3-none-any.whl → 1.2.1__py3-none-any.whl

qadence/engines/jax/__init__.py

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+from jax import config
+
+from .differentiable_backend import DifferentiableBackend
+from .differentiable_expectation import DifferentiableExpectation
+
+config.update("jax_enable_x64", True)

qadence/engines/jax/differentiable_backend.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+from qadence.backend import Backend, ConvertedCircuit, ConvertedObservable
+from qadence.engines.differentiable_backend import (
+    DifferentiableBackend as DifferentiableBackendInterface,
+)
+from qadence.engines.jax.differentiable_expectation import DifferentiableExpectation
+from qadence.measurements import Measurements
+from qadence.mitigations import Mitigations
+from qadence.noise import Noise
+from qadence.types import ArrayLike, DiffMode, Endianness, Engine, ParamDictType
+
+
+class DifferentiableBackend(DifferentiableBackendInterface):
+    """A class which wraps a QuantumBackend with the automatic differentation engine JAX.
+
+    Arguments:
+        backend: An instance of the QuantumBackend type perform execution.
+        diff_mode: A differentiable mode supported by the differentiation engine.
+        **psr_args: Arguments that will be passed on to `DifferentiableExpectation`.
+    """
+
+    def __init__(
+        self,
+        backend: Backend,
+        diff_mode: DiffMode = DiffMode.AD,
+        **psr_args: int | float | None,
+    ) -> None:
+        super().__init__(backend=backend, engine=Engine.JAX, diff_mode=diff_mode)
+        self.psr_args = psr_args
+
+    def expectation(
+        self,
+        circuit: ConvertedCircuit,
+        observable: list[ConvertedObservable] | ConvertedObservable,
+        param_values: ParamDictType = {},
+        state: ArrayLike | None = None,
+        measurement: Measurements | None = None,
+        noise: Noise | None = None,
+        mitigation: Mitigations | None = None,
+        endianness: Endianness = Endianness.BIG,
+    ) -> ArrayLike:
+        """Compute the expectation value of the `circuit` with the given `observable`.
+
+        Arguments:
+            circuit: A converted circuit as returned by `backend.circuit`.
+            observable: A converted observable as returned by `backend.observable`.
+            param_values: _**Already embedded**_ parameters of the circuit. See
+                [`embedding`][qadence.blocks.embedding.embedding] for more info.
+            state: Initial state.
+            measurement: Optional measurement protocol. If None, use
+                exact expectation value with a statevector simulator.
+            noise: A noise model to use.
+            mitigation: The error mitigation to use.
+            endianness: Endianness of the resulting bit strings.
+        """
+        observable = observable if isinstance(observable, list) else [observable]
+
+        if self.diff_mode == DiffMode.AD:
+            expectation = self.backend.expectation(circuit, observable, param_values, state)
+        else:
+            expectation = DifferentiableExpectation(
+                backend=self.backend,
+                circuit=circuit,
+                observable=observable,
+                param_values=param_values,
+                state=state,
+                measurement=measurement,
+                noise=noise,
+                mitigation=mitigation,
+                endianness=endianness,
+            ).psr()
+        return expectation

qadence/engines/jax/differentiable_expectation.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Tuple
+
+import jax.numpy as jnp
+from jax import Array, custom_vjp, vmap
+
+from qadence.backend import Backend as QuantumBackend
+from qadence.backend import ConvertedCircuit, ConvertedObservable
+from qadence.backends.jax_utils import (
+    tensor_to_jnp,
+)
+from qadence.blocks.utils import uuid_to_eigen
+from qadence.measurements import Measurements
+from qadence.mitigations import Mitigations
+from qadence.noise import Noise
+from qadence.types import Endianness, Engine, ParamDictType
+
+
+def compute_single_gap(eigen_vals: Array, default_val: float = 2.0) -> Array:
+    eigen_vals = eigen_vals.reshape(1, 2)
+    gaps = jnp.abs(jnp.tril(eigen_vals.T - eigen_vals))
+    return jnp.unique(jnp.where(gaps > 0.0, gaps, default_val), size=1)
+
+
+@dataclass
+class DifferentiableExpectation:
+    """A handler for differentiating expectation estimation using various engines."""
+
+    backend: QuantumBackend
+    circuit: ConvertedCircuit
+    observable: list[ConvertedObservable] | ConvertedObservable
+    param_values: ParamDictType
+    state: Array | None = None
+    measurement: Measurements | None = None
+    noise: Noise | None = None
+    mitigation: Mitigations | None = None
+    endianness: Endianness = Endianness.BIG
+    engine: Engine = Engine.JAX
+
+    def psr(self) -> Any:
+        n_obs = len(self.observable)
+
+        def expectation_fn(state: Array, values: ParamDictType, psr_params: ParamDictType) -> Array:
+            return self.backend.expectation(
+                circuit=self.circuit, observable=self.observable, param_values=values, state=state
+            )
+
+        @custom_vjp
+        def expectation(state: Array, values: ParamDictType, psr_params: ParamDictType) -> Array:
+            return expectation_fn(state, values, psr_params)
+
+        uuid_to_eigs = {
+            k: tensor_to_jnp(v) for k, v in uuid_to_eigen(self.circuit.abstract.block).items()
+        }
+        self.psr_params = {
+            k: self.param_values[k] for k in uuid_to_eigs.keys()
+        }  # Subset of params on which to perform PSR.
+
+        def expectation_fwd(state: Array, values: ParamDictType, psr_params: ParamDictType) -> Any:
+            return expectation_fn(state, values, psr_params), (
+                state,
+                values,
+                psr_params,
+            )
+
+        def expectation_bwd(res: Tuple[Array, ParamDictType, ParamDictType], tangent: Array) -> Any:
+            state, values, psr_params = res
+            grads = {}
+            # FIXME Hardcoding the single spectral_gap to 2.
+            spectral_gap = 2.0
+            shift = jnp.pi / 2
+
+            def shift_circ(param_name: str, values: dict) -> Array:
+                shifted_values = values.copy()
+                shiftvals = jnp.array(
+                    [shifted_values[param_name] + shift, shifted_values[param_name] - shift]
+                )
+
+                def _expectation(val: Array) -> Array:
+                    shifted_values[param_name] = val
+                    return expectation(state, shifted_values, psr_params)
+
+                return vmap(_expectation, in_axes=(0,))(shiftvals)
+
+            for param_name, _ in psr_params.items():
+                f_plus, f_min = shift_circ(param_name, values)
+                grad = spectral_gap * (f_plus - f_min) / (4.0 * jnp.sin(spectral_gap * shift / 2.0))
+                grads[param_name] = jnp.sum(tangent * grad, axis=1) if n_obs > 1 else tangent * grad
+            return None, None, grads
+
+        expectation.defvjp(expectation_fwd, expectation_bwd)
+        return expectation(self.state, self.param_values, self.psr_params)

qadence/engines/torch/__init__.py

@@ -0,0 +1,4 @@
+from __future__ import annotations
+
+from .differentiable_backend import DifferentiableBackend
+from .differentiable_expectation import DifferentiableExpectation

qadence/engines/torch/differentiable_backend.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+from functools import partial
+
+from qadence.backend import Backend as QuantumBackend
+from qadence.backend import ConvertedCircuit, ConvertedObservable
+from qadence.engines.differentiable_backend import (
+    DifferentiableBackend as DifferentiableBackendInterface,
+)
+from qadence.engines.torch.differentiable_expectation import DifferentiableExpectation
+from qadence.extensions import get_gpsr_fns
+from qadence.measurements import Measurements
+from qadence.mitigations import Mitigations
+from qadence.noise import Noise
+from qadence.types import ArrayLike, DiffMode, Endianness, Engine, ParamDictType
+
+
+class DifferentiableBackend(DifferentiableBackendInterface):
+    """A class which wraps a QuantumBackend with the automatic differentation engine TORCH.
+
+    Arguments:
+        backend: An instance of the QuantumBackend type perform execution.
+        diff_mode: A differentiable mode supported by the differentiation engine.
+        **psr_args: Arguments that will be passed on to `DifferentiableExpectation`.
+    """
+
+    def __init__(
+        self,
+        backend: QuantumBackend,
+        diff_mode: DiffMode = DiffMode.AD,
+        **psr_args: int | float | None,
+    ) -> None:
+        super().__init__(backend=backend, engine=Engine.TORCH, diff_mode=diff_mode)
+        self.psr_args = psr_args
+
+    def expectation(
+        self,
+        circuit: ConvertedCircuit,
+        observable: list[ConvertedObservable] | ConvertedObservable,
+        param_values: ParamDictType = {},
+        state: ArrayLike | None = None,
+        measurement: Measurements | None = None,
+        noise: Noise | None = None,
+        mitigation: Mitigations | None = None,
+        endianness: Endianness = Endianness.BIG,
+    ) -> ArrayLike:
+        """Compute the expectation value of the `circuit` with the given `observable`.
+
+        Arguments:
+            circuit: A converted circuit as returned by `backend.circuit`.
+            observable: A converted observable as returned by `backend.observable`.
+            param_values: _**Already embedded**_ parameters of the circuit. See
+                [`embedding`][qadence.blocks.embedding.embedding] for more info.
+            state: Initial state.
+            measurement: Optional measurement protocol. If None, use
+                exact expectation value with a statevector simulator.
+            noise: A noise model to use.
+            mitigation: The error mitigation to use.
+            endianness: Endianness of the resulting bit strings.
+        """
+        observable = observable if isinstance(observable, list) else [observable]
+        differentiable_expectation = DifferentiableExpectation(
+            backend=self.backend,
+            circuit=circuit,
+            observable=observable,
+            param_values=param_values,
+            state=state,
+            measurement=measurement,
+            noise=noise,
+            mitigation=mitigation,
+            endianness=endianness,
+        )
+
+        if self.diff_mode == DiffMode.AD:
+            expectation = differentiable_expectation.ad
+        elif self.diff_mode == DiffMode.ADJOINT:
+            expectation = differentiable_expectation.adjoint
+        else:
+            try:
+                fns = get_gpsr_fns()
+                psr_fn = fns[self.diff_mode]
+            except KeyError:
+                raise ValueError(f"{self.diff_mode} differentiation mode is not supported")
+            expectation = partial(differentiable_expectation.psr, psr_fn=psr_fn, **self.psr_args)
+        return expectation()

qadence/extensions.py

@@ -2,24 +2,34 @@ from __future__ import annotations
 
 import importlib
 from string import Template
-from typing import TypeVar
 
 from qadence.backend import Backend
-from qadence.blocks import (
-    AbstractBlock,
-)
+from qadence.blocks.abstract import TAbstractBlock
 from qadence.logger import get_logger
-from qadence.types import BackendName, DiffMode
-
-TAbstractBlock = TypeVar("TAbstractBlock", bound=AbstractBlock)
+from qadence.types import BackendName, DiffMode, Engine
 
 backends_namespace = Template("qadence.backends.$name")
 
 logger = get_logger(__name__)
 
 
+def _available_engines() -> dict:
+    """Returns a dictionary of currently installed, native qadence engines."""
+    res = {}
+    for engine in Engine.list():
+        module_path = f"qadence.engines.{engine}.differentiable_backend"
+        try:
+            module = importlib.import_module(module_path)
+            DifferentiableBackendCls = getattr(module, "DifferentiableBackend")
+            res[engine] = DifferentiableBackendCls
+        except (ImportError, ModuleNotFoundError):
+            pass
+    logger.info(f"Found engines: {res.keys()}")
+    return res
+
+
 def _available_backends() -> dict:
-    """Fallback function for native Qadence available backends if extensions is not present."""
+    """Returns a dictionary of currently installed, native qadence backends."""
     res = {}
     for backend in BackendName.list():
         module_path = f"qadence.backends.{backend}.backend"
@@ -29,11 +39,12 @@ def _available_backends() -> dict:
             res[backend] = BackendCls
         except (ImportError, ModuleNotFoundError):
             pass
+    logger.info(f"Found backends: {res.keys()}")
     return res
 
 
 def _supported_gates(name: BackendName | str) -> list[TAbstractBlock]:
-    """Fallback function for native Qadence backend supported gates if extensions is not present."""
+    """Returns a list of supported gates for the queried backend 'name'."""
     from qadence import operations
 
     name = str(BackendName(name).name.lower())
@@ -107,6 +118,7 @@ try:
     set_backend_config = getattr(module, "set_backend_config")
 except ModuleNotFoundError:
     available_backends = _available_backends
+    available_engines = _available_engines
     supported_gates = _supported_gates
     get_gpsr_fns = _gpsr_fns
     set_backend_config = _set_backend_config

qadence/finitediff.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from typing import Callable
+
+import torch
+from torch import Tensor
+
+
+def finitediff(
+    f: Callable,
+    x: Tensor,
+    derivative_indices: tuple[int, ...],
+    eps: float = None,
+) -> Tensor:
+    """
+    Arguments:
+
+        f: Function to differentiate
+        x: Input of shape `(batch_size, input_size)`
+        derivative_indices: which *input* to differentiate (i.e. which variable x[:,i])
+        eps: finite difference spacing (uses `torch.finfo(x.dtype).eps ** (1 / (2 + order))` as a
+            default)
+    """
+
+    if eps is None:
+        order = len(derivative_indices)
+        eps = torch.finfo(x.dtype).eps ** (1 / (2 + order))
+
+    # compute derivative direction vector(s)
+    eps = torch.as_tensor(eps, dtype=x.dtype)
+    _eps = 1 / eps  # type: ignore[operator]
+    ev = torch.zeros_like(x)
+    i = derivative_indices[0]
+    ev[:, i] += eps
+
+    # recursive finite differencing for higher order than 3 / mixed derivatives
+    if len(derivative_indices) > 3 or len(set(derivative_indices)) > 1:
+        di = derivative_indices[1:]
+        return (finitediff(f, x + ev, di) - finitediff(f, x - ev, di)) * _eps / 2
+    elif len(derivative_indices) == 3:
+        return (f(x + 2 * ev) - 2 * f(x + ev) + 2 * f(x - ev) - f(x - 2 * ev)) * _eps**3 / 2
+    elif len(derivative_indices) == 2:
+        return (f(x + ev) + f(x - ev) - 2 * f(x)) * _eps**2
+    elif len(derivative_indices) == 1:
+        return (f(x + ev) - f(x - ev)) * _eps / 2
+    else:
+        raise ValueError("If you see this error there is a bug in the `finitediff` function.")

qadence/mitigations/readout.py

@@ -6,17 +6,50 @@ from functools import reduce
 import numpy as np
 import numpy.typing as npt
 import torch
+from numpy.linalg import inv, matrix_rank, pinv
 from scipy.linalg import norm
 from scipy.optimize import LinearConstraint, minimize
 
 from qadence.mitigations.protocols import Mitigations
 from qadence.noise.protocols import Noise
+from qadence.types import ReadOutOptimization
 
 
 def corrected_probas(p_corr: npt.NDArray, T: npt.NDArray, p_raw: npt.NDArray) -> np.double:
     return norm(T @ p_corr.T - p_raw.T, ord=2) ** 2
 
 
+def mle_solve(p_raw: npt.NDArray) -> npt.NDArray:
+    """
+    Compute the MLE probability vector.
+
+    Algorithmic details can be found in https://arxiv.org/pdf/1106.5458.pdf Page(3).
+    """
+    # Sort p_raw by values while keeping track of indices.
+    index_sort = p_raw.argsort()
+    p_sort = p_raw[index_sort]
+    neg_sum = 0
+    breakpoint = len(p_sort) - 1
+
+    for i in range(len(p_sort)):
+        ## if neg_sum cannot be distributed among other probabilities, continue to accumulate
+        if p_sort[i] + neg_sum / (len(p_sort) - i) < 0:
+            neg_sum += p_sort[i]
+            p_sort[i] = 0
+        # set breakpoint to current index
+        else:
+            breakpoint = i
+            break
+    ## number of entries to which i can distribute(includes breakpoint)
+    size = len(p_sort) - breakpoint
+    p_sort[breakpoint:] += neg_sum / size
+
+    re_index_sort = index_sort.argsort()
+    p_corr = p_sort[re_index_sort]
+
+    return p_corr
+
+
 def renormalize_counts(corrected_counts: npt.NDArray, n_shots: int) -> npt.NDArray:
     """Renormalize counts rounding discrepancies."""
     total_counts = sum(corrected_counts)
@@ -25,51 +58,85 @@ def renormalize_counts(corrected_counts: npt.NDArray, n_shots: int) -> npt.NDArr
         corrected_counts -= counts_diff
         corrected_counts = np.where(corrected_counts < 0, 0, corrected_counts)
         sum_corrected_counts = sum(corrected_counts)
-        if sum_corrected_counts < n_shots:
-            renormalization_factor = max(sum_corrected_counts, n_shots) / min(
-                sum_corrected_counts, n_shots
-            )
-        else:
-            renormalization_factor = min(sum_corrected_counts, n_shots) / max(
-                sum_corrected_counts, n_shots
-            )
+
+        renormalization_factor = n_shots / sum_corrected_counts
         corrected_counts = np.rint(corrected_counts * renormalization_factor).astype(int)
     return corrected_counts
 
 
+def matrix_inv(K: npt.NDArray) -> npt.NDArray:
+    return inv(K) if matrix_rank(K) == K.shape[0] else pinv(K)
+
+
 def mitigation_minimization(
-    noise: Noise, mitigation: Mitigations, samples: list[Counter]
+    noise: Noise,
+    mitigation: Mitigations,
+    samples: list[Counter],
 ) -> list[Counter]:
     """Minimize a correction matrix subjected to stochasticity constraints.
 
     See Equation (5) in https://arxiv.org/pdf/2001.09980.pdf.
+    See Page(3) in https://arxiv.org/pdf/1106.5458.pdf for MLE implementation
+
+    Args:
+        noise: Specifies confusion matrix and default error probability
+        mitigation: Selects additional mitigation options based on noise choice.
+        For readout we have the following mitigation options for optimization
+        1. constrained 2. mle. Default : mle
+        samples: List of samples to be mitigated
+
+    Returns:
+        Mitigated counts computed by the algorithm
     """
     noise_matrices = noise.options.get("noise_matrix", noise.options["confusion_matrices"])
+    optimization_type = mitigation.options.get("optimization_type", ReadOutOptimization.MLE)
     n_qubits = len(list(samples[0].keys())[0])
     n_shots = sum(samples[0].values())
-    # Build the whole T matrix.
-    T_matrix = reduce(torch.kron, noise_matrices).detach().numpy()
     corrected_counters: list[Counter] = []
+
+    if optimization_type == ReadOutOptimization.CONSTRAINED:
+        # Build the whole T matrix.
+        T_matrix = reduce(torch.kron, noise_matrices).detach().numpy()
+
+    if optimization_type == ReadOutOptimization.MLE:
+        # Check if matrix is singular and use appropriate inverse.
+        noise_matrices_inv = list(map(matrix_inv, noise_matrices.numpy()))
+        T_inv = reduce(np.kron, noise_matrices_inv)
+
     for sample in samples:
         bitstring_length = 2**n_qubits
         # List of bitstrings in lexicographical order.
         ordered_bitstrings = [f"{i:0{n_qubits}b}" for i in range(bitstring_length)]
         # Array of raw probabilites.
         p_raw = np.array([sample[bs] for bs in ordered_bitstrings]) / n_shots
-        # Initial random guess in [0,1].
-        p_corr0 = np.random.rand(bitstring_length)
-        # Stochasticity constraints.
-        normality_constraint = LinearConstraint(
-            np.ones(bitstring_length).astype(int), lb=1.0, ub=1.0
-        )
-        positivity_constraint = LinearConstraint(
-            np.eye(bitstring_length).astype(int), lb=0.0, ub=1.0
-        )
-        constraints = [normality_constraint, positivity_constraint]
-        # Minimize the corrected probabilities.
-        res = minimize(corrected_probas, p_corr0, args=(T_matrix, p_raw), constraints=constraints)
-        # breakpoint()
-        corrected_counts = np.rint(res.x * n_shots).astype(int)
+
+        if optimization_type == ReadOutOptimization.CONSTRAINED:
+            # Initial random guess in [0,1].
+            p_corr0 = np.random.rand(bitstring_length)
+            # Stochasticity constraints.
+            normality_constraint = LinearConstraint(
+                np.ones(bitstring_length).astype(int), lb=1.0, ub=1.0
+            )
+            positivity_constraint = LinearConstraint(
+                np.eye(bitstring_length).astype(int), lb=0.0, ub=1.0
+            )
+            constraints = [normality_constraint, positivity_constraint]
+            # Minimize the corrected probabilities.
+            res = minimize(
+                corrected_probas, p_corr0, args=(T_matrix, p_raw), constraints=constraints
+            )
+            p_corr = res.x
+
+        elif optimization_type == ReadOutOptimization.MLE:
+            # Compute corrected inverse using matrix inversion and run MLE.
+            p_corr = mle_solve(T_inv @ p_raw)
+        else:
+            raise NotImplementedError(
+                f"Requested method {optimization_type} does not match supported protocols."
+            )
+
+        corrected_counts = np.rint(p_corr * n_shots).astype(int)
+
         # Renormalize if total counts differs from n_shots.
         corrected_counts = renormalize_counts(corrected_counts=corrected_counts, n_shots=n_shots)
         # At this point, the count should be off by at most 2, added or substracted to/from the

qadence/ml_tools/models.py

@@ -178,9 +178,16 @@ class TransformedModule(torch.nn.Module):
         if isinstance(self.model, (QuantumModel, QNN)):
             if not isinstance(x, dict):
                 x = self._format_to_dict(x)
-            return {
-                key: self._input_scaling * (val + self._input_shifting) for key, val in x.items()
-            }
+            if self.in_features == 1:
+                return {
+                    key: self._input_scaling * (val + self._input_shifting)
+                    for key, val in x.items()
+                }
+            else:
+                return {
+                    key: self._input_scaling[idx] * (val + self._input_shifting[idx])
+                    for idx, (key, val) in enumerate(x.items())
+                }
 
         else:
             assert isinstance(self.model, torch.nn.Module) and isinstance(x, Tensor)