Update binomial_weights function for tomo fitters (#1075)

<!--
⚠️ If you do not respect this template, your pull request will be
closed.
⚠️ Your pull request title should be short detailed and understandable
for all.
⚠️ Also, please add it in the CHANGELOG file under Unreleased section.
⚠️ If your pull request fixes an open issue, please link to the issue.

✅ I have added the tests to cover my changes.
✅ I have updated the documentation accordingly.
✅ I have read the CONTRIBUTING document.
-->

### Summary

Updates `binomial_weights` function to use the dirichlet distribution
for computing weights, and to work with the conditional tomography
experiment changes to allow weights being an array of the same shape as
outcome data.

### Details and comments

qiskit_experiments/library/tomography/fitters/cvxpy_lstsq.py

@@ -365,6 +365,7 @@ def cvxpy_gaussian_lstsq(
     trace_preserving: Union[None, bool, str] = "auto",
     partial_trace: Optional[np.ndarray] = None,
     outcome_prior: Union[np.ndarray, int] = 0.5,
+    max_weight: float = 1e10,
     **kwargs,
 ) -> Dict:
     r"""Constrained Gaussian linear least-squares tomography fitter.
@@ -429,6 +430,8 @@ def cvxpy_gaussian_lstsq(
                        trace to POVM matrices.
         outcome_prior: The Baysian prior :math:`\alpha` to use computing Gaussian
             weights. See additional information.
+        max_weight: Set the maximum value allowed for weights vector computed from
+            tomography data variance.
         kwargs: kwargs for cvxpy solver.
 
     Raises:
@@ -440,14 +443,13 @@ def cvxpy_gaussian_lstsq(
     """
     t_start = time.time()
 
-    _, variance = lstsq_utils.dirichlet_mean_and_var(
+    weights = lstsq_utils.binomial_weights(
         outcome_data,
         shot_data=shot_data,
         outcome_prior=outcome_prior,
+        max_weight=max_weight,
     )
 
-    weights = 1.0 / np.sqrt(variance)
-
     fits, metadata = cvxpy_linear_lstsq(
         outcome_data,
         shot_data,

qiskit_experiments/library/tomography/fitters/lininv.py

@@ -38,6 +38,7 @@ def linear_inversion(
     preparation_qubits: Optional[Tuple[int, ...]] = None,
     conditional_measurement_indices: Optional[np.ndarray] = None,
     conditional_preparation_indices: Optional[np.ndarray] = None,
+    atol: float = 1e-8,
 ) -> Tuple[np.ndarray, Dict]:
     r"""Linear inversion tomography fitter.
 
@@ -106,6 +107,7 @@ def linear_inversion(
         conditional_preparation_indices: Optional, conditional preparation data
             indices. If set this will return a list of fitted states conditioned
             on a fixed basis preparation of these qubits.
+        atol: truncate any probabilities below this value to zero.
 
     Raises:
         AnalysisError: If the fitted vector is not a square matrix
@@ -242,10 +244,10 @@ def linear_inversion(
 
             # Get probabilities and optional measurement basis component
             for outcome, freq in enumerate(outcomes):
-                if freq == 0:
+                prob = freq / shots
+                if np.isclose(prob, 0, atol=atol):
                     # Skip component with zero probability
                     continue
-                prob = freq / shots
 
                 # Get component on non-conditional bits
                 outcome_meas = f_meas_outcome(outcome)

qiskit_experiments/library/tomography/fitters/lstsq_utils.py

@@ -16,7 +16,7 @@
 from typing import Optional, Tuple, Callable, Sequence, Union
 import functools
 import numpy as np
-from qiskit.utils import deprecate_function
+from qiskit.utils import deprecate_arguments
 from qiskit_experiments.exceptions import AnalysisError
 from qiskit_experiments.library.tomography.basis import (
     MeasurementBasis,
@@ -182,6 +182,55 @@ def lstsq_data(
     return basis_mat, probs, prob_weights
 
 
+@deprecate_arguments({"beta": "outcome_prior"})
+def binomial_weights(
+    outcome_data: np.ndarray,
+    shot_data: Optional[Union[np.ndarray, int]] = None,
+    outcome_prior: Union[np.ndarray, int] = 0.5,
+    max_weight: float = 1e10,
+) -> np.ndarray:
+    r"""Compute weights from tomography data variance.
+
+    This is computed via a Bayesian update of a Dirichlet distribution
+    with observed outcome data frequences :math:`f_i(s)`, and Dirichlet
+    prior :math:`\alpha_i(s)` for tomography basis index `i` and
+    measurement outcome `s`.
+
+    The mean posterior probabilities are computed as
+
+    .. math:
+        p_i(s) &= \frac{f_i(s) + \alpha_i(s)}{\bar{\alpha}_i + N_i} \\
+        Var[p_i(s)] &= \frac{p_i(s)(1-p_i(s))}{\bar{\alpha}_i + N_i + 1}
+
+    where :math:`N_i = \sum_s f_i(s)` is the total number of shots, and
+    :math:`\bar{\alpha}_i = \sum_s \alpha_i(s)` is the norm of the prior
+    vector for basis index `i`.
+
+    Args:
+        outcome_data: measurement outcome frequency data.
+        shot_data: Optional, basis measurement total shot data. If not
+            provided this will be inferred from the sum of outcome data
+            for each basis index.
+        outcome_prior: measurement outcome Dirichlet distribution prior.
+        max_weight: Set the maximum value allowed for weights vector computed from
+            tomography data variance.
+
+    Returns:
+        The array of weights computed from the tomography data.
+    """
+    # Compute variance
+    _, variance = dirichlet_mean_and_var(
+        outcome_data,
+        shot_data=shot_data,
+        outcome_prior=outcome_prior,
+    )
+    # Use max weights to determin a min variance value and clip variance
+    min_variance = 1 / (max_weight**2)
+    variance = np.clip(variance, min_variance, None)
+    weights = 1.0 / np.sqrt(variance)
+    return weights
+
+
 def dirichlet_mean_and_var(
     outcome_data: np.ndarray,
     shot_data: Optional[Union[np.ndarray, int]] = None,
@@ -237,55 +286,6 @@ def dirichlet_mean_and_var(
     return mean_probs, variance
 
 
-# pylint: disable = bad-docstring-quotes
-@deprecate_function(
-    "The binomial_weights function is deprecated and will "
-    "be removed in the 0.6 release. Use the `dirichlet_mean_and_var` "
-    "function instead."
-)
-def binomial_weights(
-    outcome_data: np.ndarray,
-    shot_data: np.ndarray,
-    beta: float = 0,
-) -> np.ndarray:
-    r"""Compute weights vector from the binomial distribution.
-    The returned weights are given by :math:`w_i = 1 / \sigma_i` where
-    the standard deviation :math:`\sigma_i` is estimated as
-    :math:`\sigma_i = \sqrt{p_i(1-p_i) / n_i}`. To avoid dividing
-    by zero the probabilities are hedged using the *add-beta* rule
-    .. math:
-        p_i = \frac{f_i + \beta}{n_i + K \beta}
-    where :math:`f_i` is the observed frequency, :math:`n_i` is the
-    number of shots, and :math:`K` is the number of possible measurement
-    outcomes.
-    Args:
-        outcome_data: measurement outcome frequency data.
-        shot_data: basis measurement total shot data.
-        beta: Hedging parameter for converting frequencies to
-              probabilities. If 0 hedging is disabled.
-    Returns:
-        The weight vector.
-    """
-    size = outcome_data.size
-    num_data, num_outcomes = outcome_data.shape
-
-    # Compute hedged probabilities where the "add-beta" rule ensures
-    # there are no zero or 1 values so we don't have any zero variance
-    probs = np.zeros(size, dtype=float)
-    prob_shots = np.zeros(size, dtype=int)
-    idx = 0
-    for i in range(num_data):
-        shots = shot_data[i]
-        denom = shots + num_outcomes * beta
-        freqs = outcome_data[i]
-        for outcome in range(num_outcomes):
-            probs[idx] = (freqs[outcome] + beta) / denom
-            prob_shots[idx] = shots
-            idx += 1
-    variance = probs * (1 - probs)
-    return np.sqrt(prob_shots / variance)
-
-
 @functools.lru_cache(None)
 def _partial_outcome_function(indices: Tuple[int]) -> Callable:
     """Return function for computing partial outcome of specified indices"""

qiskit_experiments/library/tomography/fitters/scipy_lstsq.py

@@ -190,6 +190,7 @@ def scipy_gaussian_lstsq(
     measurement_qubits: Optional[Tuple[int, ...]] = None,
     preparation_qubits: Optional[Tuple[int, ...]] = None,
     outcome_prior: Union[np.ndarray, int] = 0.5,
+    max_weight: float = 1e10,
     **kwargs,
 ) -> Dict:
     r"""Gaussian linear least-squares tomography fitter.
@@ -239,6 +240,8 @@ def scipy_gaussian_lstsq(
             If None they are assumed to be ``[0, ..., N-1]`` for N preparated qubits.
         outcome_prior: The Baysian prior :math:`\alpha` to use computing Gaussian
             weights. See additional information.
+        max_weight: Set the maximum value allowed for weights vector computed from
+            tomography data variance.
         kwargs: additional kwargs for :func:`scipy.linalg.lstsq`.
 
     Raises:
@@ -248,12 +251,14 @@ def scipy_gaussian_lstsq(
         The fitted matrix rho that maximizes the least-squares likelihood function.
     """
     t_start = time.time()
-    _, variance = lstsq_utils.dirichlet_mean_and_var(
+
+    weights = lstsq_utils.binomial_weights(
         outcome_data,
         shot_data=shot_data,
         outcome_prior=outcome_prior,
+        max_weight=max_weight,
     )
-    weights = 1.0 / np.sqrt(variance)
+
     fits, metadata = scipy_linear_lstsq(
         outcome_data,
         shot_data,