Repository: scikit-hep/cabinetry

Issue:
Uncertainties per sample in yield tables
The current yield table implementation (after #158) does not include a calculation of yield uncertainties per sample. These numbers can be more difficult to interpret due to correlations between samples, but would nevertheless be useful to provide. Additional calculations are required to obtain these uncertainties, so they may be implemented as optional.

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/scikit-hep/cabinetry/pull/356?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scikit-hep) Report
> Merging [#356](https://codecov.io/gh/scikit-hep/cabinetry/pull/356?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scikit-hep) (c6b651f) into [master](https://codecov.io/gh/scikit-hep/cabinetry/commit/aee8689cbdb56a0cf147bd38f4feb6f70a2345f5?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scikit-hep) (aee8689) will **decrease** coverage by `0.09%`.
> The diff coverage is `89.47%`.

```diff
@@             Coverage Diff             @@
##            master     #356      +/-   ##
===========================================
- Coverage   100.00%   99.90%   -0.10%     
===========================================
  Files           23       23              
  Lines         1999     2003       +4     
  Branches       326      326              
===========================================
+ Hits          1999     2001       +2     
- Misses           0        2       +2     
```


| [Impacted Files](https://codecov.io/gh/scikit-hep/cabinetry/pull/356?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scikit-hep) | Coverage Δ | |
|---|---|---|
| [src/cabinetry/visualize/\_\_init\_\_.py](https://codecov.io/gh/scikit-hep/cabinetry/pull/356/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scikit-hep#diff-c3JjL2NhYmluZXRyeS92aXN1YWxpemUvX19pbml0X18ucHk=) | `100.00% <ø> (ø)` | |
| [src/cabinetry/model\_utils.py](https://codecov.io/gh/scikit-hep/cabinetry/pull/356/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scikit-hep#diff-c3JjL2NhYmluZXRyeS9tb2RlbF91dGlscy5weQ==) | `98.85% <88.88%> (-1.15%)` | :arrow_down: |
| [src/cabinetry/tabulate.py](https://codecov.io/gh/scikit-hep/cabinetry/pull/356/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scikit-hep#diff-c3JjL2NhYmluZXRyeS90YWJ1bGF0ZS5weQ==) | `100.00% <100.00%> (ø)` | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scikit-hep). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scikit-hep)


Patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index a9ad8f1..9f6fb8c 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -39,7 +39,7 @@ repos:
         # configuration duplicated in setup.cfg
         args: ["--match=(?!setup|example).*\\.py'", "--match-dir='^(?!(tests|utils|docs)).*'", "--convention=google"]
 -   repo: https://github.com/codespell-project/codespell
-    rev: v2.1.0
+    rev: v2.2.1
     hooks:
     -   id: codespell
         args: ["-L", "hist", "src", "tests", "utils", "docs/*rst"]
diff --git a/pyproject.toml b/pyproject.toml
index 16cf156..fede21c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -37,6 +37,7 @@ python_version = "3.7"
 module = [
     "uproot",
     "awkward",
+    "awkward._v2",
     "pyhf",
     "matplotlib.*",
     "iminuit",
diff --git a/setup.cfg b/setup.cfg
index a663ffe..491a9dd 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -25,7 +25,7 @@ python_requires = >=3.7
 install_requires =
     pyhf[minuit]~=0.7.0rc1  # model.config.suggested_fixed API change
     boost_histogram>=1.0.0  # subclassing with family, 1.02 for stdev scaling fix (currently not needed)
-    awkward>=1.0  # new API
+    awkward>=1.8  # _v2 API in submodule
     tabulate>=0.8.1  # multiline text
     matplotlib
     typing_extensions;python_version<"3.8"  # for Literal type
diff --git a/src/cabinetry/model_utils.py b/src/cabinetry/model_utils.py
index 8d013d3..95e0df1 100644
--- a/src/cabinetry/model_utils.py
+++ b/src/cabinetry/model_utils.py
@@ -5,7 +5,12 @@ import json
 import logging
 from typing import Any, DefaultDict, Dict, List, NamedTuple, Optional, Tuple, Union
 
-import awkward as ak
+try:
+    # use awkward v2
+    import awkward._v2 as ak
+except ModuleNotFoundError:  # pragma: no cover
+    # fallback if the _v2 submodule disappears after the full v2 release
+    import awkward as ak  # pragma: no cover
 import numpy as np
 import pyhf
 
@@ -16,7 +21,7 @@ log = logging.getLogger(__name__)
 
 
 # cache holding results from yield uncertainty calculations
-_YIELD_STDEV_CACHE: Dict[Any, Tuple[List[List[float]], List[float]]] = {}
+_YIELD_STDEV_CACHE: Dict[Any, Tuple[List[List[List[float]]], List[List[float]]]] = {}
 
 
 class ModelPrediction(NamedTuple):
@@ -26,17 +31,18 @@ class ModelPrediction(NamedTuple):
         model (pyhf.pdf.Model): model to which prediction corresponds to
         model_yields (List[List[List[float]]]): yields per sample, channel and bin,
             indices: channel, sample, bin
-        total_stdev_model_bins (List[List[float]]): total yield uncertainty per channel
-            and per bin, indices: channel, bin
-        total_stdev_model_channels (List[float]): total yield uncertainty per channel,
-            index: channel
+        total_stdev_model_bins (List[List[List[float]]]): total yield uncertainty per
+            channel, sample and bin, indices: channel, sample, bin (last sample: sum
+            over samples)
+        total_stdev_model_channels (List[List[float]]): total yield uncertainty per
+            channel and sample, indices: channel, sample (last sample: sum over samples)
         label (str): label for the prediction, e.g. "pre-fit" or "post-fit"
     """
 
     model: pyhf.pdf.Model
     model_yields: List[List[List[float]]]
-    total_stdev_model_bins: List[List[float]]
-    total_stdev_model_channels: List[float]
+    total_stdev_model_bins: List[List[List[float]]]
+    total_stdev_model_channels: List[List[float]]
     label: str
 
 
@@ -193,14 +199,15 @@ def yield_stdev(
     parameters: np.ndarray,
     uncertainty: np.ndarray,
     corr_mat: np.ndarray,
-) -> Tuple[List[List[float]], List[float]]:
-    """Calculates symmetrized yield standard deviation of a model, per bin and channel.
+) -> Tuple[List[List[List[float]]], List[List[float]]]:
+    """Calculates symmetrized model yield standard deviation per channel / sample / bin.
 
-    Returns both the uncertainties per bin (in a list of channels), and the uncertainty
-    of the total yield per channel (again, for a list of channels). To calculate the
-    uncertainties for the total yield, the function internally treats the sum of yields
-    per channel like another channel with one bin. The results of this function are
-    cached to speed up subsequent calls with the same arguments.
+    Returns both the uncertainties per bin (in a list of channels and samples), and the
+    uncertainty of the total yield per channel (again, for a list of channels and
+    samples). To calculate the uncertainties for the total yield per channel, the
+    function internally treats the sum of yields per channel like another channel with
+    one bin. Similarly, the sum over samples is treated as another sample. The results
+    of this function are cached to speed up subsequent calls with the same arguments.
 
     Args:
         model (pyhf.pdf.Model): the model for which to calculate the standard deviations
@@ -210,9 +217,13 @@ def yield_stdev(
         corr_mat (np.ndarray): correlation matrix
 
     Returns:
-        Tuple[List[List[float]], List[float]]:
-            - list of channels, each channel is a list of standard deviations per bin
-            - list of standard deviations per channel
+        Tuple[List[List[List[float]]], List[List[float]]]:
+            - list of channels, each channel is a list of samples, and each sample a
+              list of standard deviations per bin (the last sample corresponds to a sum
+              over all samples)
+            - list of standard deviations per channel, each channel is a list containing
+              the standard deviations per sample (the last sample corresponds to a sum
+              over all samples)
     """
     # check whether results are already stored in cache
     cached_results = _YIELD_STDEV_CACHE.get(
@@ -230,10 +241,12 @@ def yield_stdev(
 
     # the lists up_variations and down_variations will contain the model distributions
     # with all parameters varied individually within uncertainties
-    # indices: variation, channel, bin
+    # indices: variation, channel, sample, bin
     # following the channels contained in the model, there are additional entries with
     # yields summed per channel (internally treated like additional channels) to get the
     # per-channel uncertainties
+    # in the same way, the total model prediction is following the list of individual
+    # samples (and then treated like an additional sample)
     up_variations = []
     down_variations = []
 
@@ -246,28 +259,42 @@ def yield_stdev(
         down_pars = parameters.copy().astype(float)
         down_pars[i_par] -= uncertainty[i_par]
 
-        # total model distribution with this parameter varied up
+        # model distribution per sample with this parameter varied up
         up_comb = pyhf.tensorlib.to_numpy(
-            model.expected_data(up_pars, include_auxdata=False)
+            model.main_model.expected_data(up_pars, return_by_sample=True)
         )
-        # turn into list of channels
+        # attach another entry with the total model prediction (sum over all samples)
+        # indices: sample, bin
+        up_comb = np.vstack((up_comb, np.sum(up_comb, axis=0)))
+        # turn into list of channels (keep all samples, select correct bins per channel)
+        # indices: channel, sample, bin
         up_yields = [
-            up_comb[model.config.channel_slices[ch]] for ch in model.config.channels
+            up_comb[:, model.config.channel_slices[ch]] for ch in model.config.channels
         ]
         # append list of yields summed per channel
-        up_yields += [np.asarray([sum(chan_yields)]) for chan_yields in up_yields]
+        up_yields += [
+            np.asarray(np.sum(chan_yields, axis=-1, keepdims=True))
+            for chan_yields in up_yields
+        ]
+        # indices: variation, channel, sample, bin
         up_variations.append(up_yields)
 
-        # total model distribution with this parameter varied down
+        # model distribution per sample with this parameter varied down
         down_comb = pyhf.tensorlib.to_numpy(
-            model.expected_data(down_pars, include_auxdata=False)
+            model.main_model.expected_data(down_pars, return_by_sample=True)
         )
+        # add total prediction (sum over samples)
+        down_comb = np.vstack((down_comb, np.sum(down_comb, axis=0)))
         # turn into list of channels
         down_yields = [
-            down_comb[model.config.channel_slices[ch]] for ch in model.config.channels
+            down_comb[:, model.config.channel_slices[ch]]
+            for ch in model.config.channels
         ]
         # append list of yields summed per channel
-        down_yields += [np.asarray([sum(chan_yields)]) for chan_yields in down_yields]
+        down_yields += [
+            np.asarray(np.sum(chan_yields, axis=-1, keepdims=True))
+            for chan_yields in down_yields
+        ]
         down_variations.append(down_yields)
 
     # convert to awkward arrays for further processing
@@ -275,13 +302,14 @@ def yield_stdev(
     down_variations_ak = ak.from_iter(down_variations)
 
     # calculate symmetric uncertainties for all components
+    # indices: variation, channel (last entries sums), sample (last entry sum), bin
     sym_uncs = (up_variations_ak - down_variations_ak) / 2
 
-    # calculate total variance, indexed by channel and bin (per-channel numbers act like
-    # additional channels with one bin each)
+    # calculate total variance, indexed by channel, sample, bin (per-channel numbers act
+    # like additional channels with one bin each)
     if np.count_nonzero(corr_mat - np.diagflat(np.ones_like(parameters))) == 0:
         # no off-diagonal contributions from correlation matrix (e.g. pre-fit)
-        total_variance = np.sum(np.power(sym_uncs, 2), axis=0)
+        total_variance = ak.sum(np.power(sym_uncs, 2), axis=0)
     else:
         # full calculation including off-diagonal contributions
         # with v as vector of variations (each element contains yields under variation)
@@ -289,31 +317,33 @@ def yield_stdev(
         # variance = sum_i sum_j v[i] * M[i, j] * v[j]
         # where the product between elements of v again is elementwise (multiplying bin
         # yields), and the final variance shape is the same as element of v (yield
-        # uncertainties per bin and per channel)
+        # uncertainties per bin, sample, channel)
 
         # possible optimizations that could be considered here:
         #   - skipping staterror-staterror terms for per-bin calculation (orthogonal)
         #   - taking advantage of correlation matrix symmetry
         #   - (optional) skipping combinations with correlations below threshold
 
-        # calculate M[i, j] * v[j] first, indices: pars (i), pars (j), channel, bin
-        m_times_v = corr_mat[..., np.newaxis, np.newaxis] * sym_uncs[np.newaxis, ...]
-        # now multiply by v[i] as well, indices: pars(i), pars(j), channel, bin
-        v_times_m_times_v = sym_uncs[:, np.newaxis, ...] * m_times_v
-        # finally perform sums over i and j, remaining indices: channel, bin
-        # ak.flatten due to https://github.com/scikit-hep/awkward-1.0/issues/1283
-        assert v_times_m_times_v.ndim == 4, (
-            "unexpected dimensionality of array, please report if you run into this: "
-            "https://github.com/scikit-hep/cabinetry/issues/326"
+        # calculate M[i, j] * v[j] first
+        # indices: pars (i), pars (j), channel, sample, bin
+        m_times_v = (
+            corr_mat[..., np.newaxis, np.newaxis, np.newaxis]
+            * sym_uncs[np.newaxis, ...]
         )
-        total_variance = np.sum(ak.flatten(v_times_m_times_v, axis=1), axis=0)
+        # now multiply by v[i] as well, indices: pars(i), pars(j), channel, sample, bin
+        v_times_m_times_v = sym_uncs[:, np.newaxis, ...] * m_times_v
+        # finally perform sums over i and j, remaining indices: channel, sample, bin
+        total_variance = ak.sum(ak.sum(v_times_m_times_v, axis=1), axis=0)
 
     # convert to standard deviations per bin and per channel
     n_channels = len(model.config.channels)
+    # indices: (channel, sample, bin)
     total_stdev_per_bin = np.sqrt(total_variance[:n_channels])
-    total_stdev_per_channel = ak.flatten(np.sqrt(total_variance[n_channels:]))
-    log.debug(f"total stdev is {total_stdev_per_bin}")
-    log.debug(f"total stdev per channel is {total_stdev_per_channel}")
+    # indices: (channel, sample)
+    total_stdev_per_channel = ak.sum(np.sqrt(total_variance[n_channels:]), axis=-1)
+    # log total stdev per bin / channel (-1 index for sample sum)
+    log.debug(f"total stdev is {total_stdev_per_bin[:, -1, :]}")
+    log.debug(f"total stdev per channel is {total_stdev_per_channel[:, -1]}")
 
     # convert to lists
     total_stdev_per_bin = ak.to_list(total_stdev_per_bin)
diff --git a/src/cabinetry/tabulate.py b/src/cabinetry/tabulate.py
index c8f8e48..1ec0ce9 100644
--- a/src/cabinetry/tabulate.py
+++ b/src/cabinetry/tabulate.py
@@ -87,7 +87,7 @@ def _save_tables(
 def _yields_per_bin(
     model: pyhf.pdf.Model,
     model_yields: List[List[List[float]]],
-    total_stdev_model: List[List[float]],
+    total_stdev_model: List[List[List[float]]],
     data: List[List[float]],
     channels: List[str],
     label: str,
@@ -97,8 +97,8 @@ def _yields_per_bin(
     Args:
         model (pyhf.pdf.Model): the model which the table corresponds to
         model_yields (List[List[List[float]]]): yields per channel, sample, and bin
-        total_stdev_model (List[List[float]]): total model standard deviation per
-            channel and bin
+        total_stdev_model (List[List[List[float]]]): total model standard deviation per
+            channel, sample and bin
         data (List[List[float]]): data yield per channel and bin
         channels (List[str]): names of channels to use
         label (str): label for model prediction to include in log
@@ -121,7 +121,8 @@ def _yields_per_bin(
                     {
                         _header_name(
                             channel_name, i_bin
-                        ): f"{model_yields[i_chan][i_sam][i_bin]:.2f}"
+                        ): f"{model_yields[i_chan][i_sam][i_bin]:.2f} "
+                        f"\u00B1 {total_stdev_model[i_chan][i_sam][i_bin]:.2f}"
                     }
                 )
         table.append(sample_dict)
@@ -140,7 +141,7 @@ def _yields_per_bin(
                 {
                     header_name: (
                         f"{total_model[i_bin]:.2f} "
-                        f"\u00B1 {total_stdev_model[i_chan][i_bin]:.2f}"
+                        f"\u00B1 {total_stdev_model[i_chan][-1][i_bin]:.2f}"
                     )
                 }
             )
@@ -157,7 +158,7 @@ def _yields_per_bin(
 def _yields_per_channel(
     model: pyhf.pdf.Model,
     model_yields: List[List[float]],
-    total_stdev_model: List[float],
+    total_stdev_model: List[List[float]],
     data: List[float],
     channels: List[str],
     label: str,
@@ -167,7 +168,8 @@ def _yields_per_channel(
     Args:
         model (pyhf.pdf.Model): the model which the table corresponds to
         model_yields (List[List[float]]): yields per channel and sample
-        total_stdev_model (List[float]): total model standard deviation per channel
+        total_stdev_model (List[List[float]]): total model standard deviation per
+            channel and sample
         data (List[float]): data yield per channel
         channels (List[str]): names of channels to use
         label (str): label for model prediction to include in log
@@ -184,7 +186,12 @@ def _yields_per_channel(
     for i_sam, sample_name in enumerate(model.config.samples):
         sample_dict = {"sample": sample_name}  # one dict per sample
         for i_chan, channel_name in zip(channel_indices, channels):
-            sample_dict.update({channel_name: f"{model_yields[i_chan][i_sam]:.2f}"})
+            sample_dict.update(
+                {
+                    channel_name: f"{model_yields[i_chan][i_sam]:.2f} "
+                    f"\u00B1 {total_stdev_model[i_chan][i_sam]:.2f}"
+                }
+            )
         table.append(sample_dict)
 
     # dicts for total model prediction and data
@@ -193,7 +200,10 @@ def _yields_per_channel(
     for i_chan, channel_name in zip(channel_indices, channels):
         total_model = np.sum(model_yields[i_chan], axis=0)  # sum over samples
         total_dict.update(
-            {channel_name: f"{total_model:.2f} \u00B1 {total_stdev_model[i_chan]:.2f}"}
+            {
+                channel_name: f"{total_model:.2f} "
+                f"\u00B1 {total_stdev_model[i_chan][-1]:.2f}"
+            }
         )
         data_dict.update({channel_name: f"{data[i_chan]:.2f}"})
     table += [total_dict, data_dict]
diff --git a/src/cabinetry/visualize/__init__.py b/src/cabinetry/visualize/__init__.py
index aed3b43..034acd5 100644
--- a/src/cabinetry/visualize/__init__.py
+++ b/src/cabinetry/visualize/__init__.py
@@ -235,7 +235,7 @@ def data_mc(
 
         fig = plot_model.data_mc(
             histogram_dict_list,
-            np.asarray(model_prediction.total_stdev_model_bins[i_chan]),
+            np.asarray(model_prediction.total_stdev_model_bins[i_chan][-1]),
             bin_edges,
             figure_path=figure_path,
             log_scale=log_scale,


Test patch:
diff --git a/tests/test_integration.py b/tests/test_integration.py
index 7e869d2..376aebb 100644
--- a/tests/test_integration.py
+++ b/tests/test_integration.py
@@ -162,18 +162,36 @@ def test_integration(tmp_path, ntuple_creator, caplog):
     model_prefit = cabinetry.model_utils.prediction(model)
     assert np.allclose(
         model_prefit.total_stdev_model_bins,
-        [[69.040789, 58.343328, 38.219599, 45.296964]],
+        [
+            [
+                [69.040789, 58.329118, 37.973787, 45.137157],  # background
+                [0.0, 0.100772, 1.487773, 1.620867],  # signal
+                [69.040789, 58.343328, 38.219599, 45.296964],  # sum over samples
+            ]
+        ],
+    )
+    assert np.allclose(
+        model_prefit.total_stdev_model_channels, [[136.368732, 2.851565, 136.791978]]
     )
-    assert np.allclose(model_prefit.total_stdev_model_channels, [136.791978])
     _ = cabinetry.visualize.data_mc(model_prefit, data, close_figure=True)
 
     model_postfit = cabinetry.model_utils.prediction(model, fit_results=fit_results)
     assert np.allclose(
         model_postfit.total_stdev_model_bins,
-        [[11.898333, 7.283185, 7.414715, 7.687922]],
+        [
+            [
+                [11.898551, 7.513216, 21.002006, 24.284847],  # background
+                [0.0, 1.467646, 22.137293, 22.269200],  # signal
+                [11.898551, 7.283171, 7.414715, 7.687966],  # sum over samples
+            ]
+        ],
         rtol=1e-4,
     )
-    assert np.allclose(model_postfit.total_stdev_model_channels, [20.439750], atol=5e-4)
+    assert np.allclose(
+        model_postfit.total_stdev_model_channels,
+        [[41.043814, 45.814417, 20.439575]],
+        atol=2e-3,
+    )
     _ = cabinetry.visualize.data_mc(model_postfit, data, close_figure=True)
 
     # nuisance parameter ranking
diff --git a/tests/test_model_utils.py b/tests/test_model_utils.py
index 9175429..a5a27a3 100644
--- a/tests/test_model_utils.py
+++ b/tests/test_model_utils.py
@@ -13,8 +13,8 @@ from cabinetry.fit.results_containers import FitResults
 def test_ModelPrediction(example_spec):
     model = pyhf.Workspace(example_spec).model()
     model_yields = [[[10.0]]]
-    total_stdev_model_bins = [[2.0]]
-    total_stdev_model_channels = [2.0]
+    total_stdev_model_bins = [[[2.0], [2.0]]]
+    total_stdev_model_channels = [[2.0, 2.0]]
     label = "abc"
     model_prediction = model_utils.ModelPrediction(
         model, model_yields, total_stdev_model_bins, total_stdev_model_channels, label
@@ -158,8 +158,8 @@ def test_yield_stdev(example_spec, example_spec_multibin):
     total_stdev_bin, total_stdev_chan = model_utils.yield_stdev(
         model, parameters, uncertainty, corr_mat
     )
-    assert np.allclose(total_stdev_bin, [[8.03150606]])
-    assert np.allclose(total_stdev_chan, [8.03150606])
+    assert np.allclose(total_stdev_bin, [[[8.03150606], [8.03150606]]])
+    assert np.allclose(total_stdev_chan, [[8.03150606, 8.03150606]])
 
     # pre-fit
     parameters = np.asarray([1.0, 1.0])
@@ -168,8 +168,8 @@ def test_yield_stdev(example_spec, example_spec_multibin):
     total_stdev_bin, total_stdev_chan = model_utils.yield_stdev(
         model, parameters, uncertainty, diag_corr_mat
     )
-    assert np.allclose(total_stdev_bin, [[2.56754823]])  # the staterror
-    assert np.allclose(total_stdev_chan, [2.56754823])
+    assert np.allclose(total_stdev_bin, [[[2.56754823], [2.56754823]]])  # the staterror
+    assert np.allclose(total_stdev_chan, [[2.56754823, 2.56754823]])
 
     # multiple channels, bins, staterrors
     model = pyhf.Workspace(example_spec_multibin).model()
@@ -186,9 +186,12 @@ def test_yield_stdev(example_spec, example_spec_multibin):
     total_stdev_bin, total_stdev_chan = model_utils.yield_stdev(
         model, parameters, uncertainty, corr_mat
     )
-    expected_stdev_bin = [[8.056054, 1.670629], [2.775377]]
-    expected_stdev_chan = [9.596340, 2.775377]
-    for i_reg in range(2):
+    expected_stdev_bin = [
+        [[8.056054, 1.670629], [8.056054, 1.670629]],
+        [[2.775377], [2.775377]],
+    ]
+    expected_stdev_chan = [[9.596340, 9.596340], [2.775377, 2.775377]]
+    for i_reg in range(2):  # two channels
         assert np.allclose(total_stdev_bin[i_reg], expected_stdev_bin[i_reg])
         assert np.allclose(total_stdev_chan[i_reg], expected_stdev_chan[i_reg])
 
@@ -214,9 +217,12 @@ def test_yield_stdev(example_spec, example_spec_multibin):
 @mock.patch(
     "cabinetry.model_utils.yield_stdev",
     side_effect=[
-        ([[5.0, 2.0], [1.0]], [5.38516481, 1.0]),
-        ([[0.3]], [0.3]),
-        ([[0.3]], [0.3]),
+        (
+            [[[5.0, 2.0], [5.0, 2.0]], [[1.0], [1.0]]],
+            [[5.38516481, 5.38516481], [1.0, 1.0]],
+        ),
+        ([[[0.3], [0.3]]], [[0.3, 0.3]]),
+        ([[[0.3], [0.3]]], [[0.3, 0.3]]),
     ],
 )
 @mock.patch(
@@ -242,8 +248,13 @@ def test_prediction(
     assert model_pred.model == model
     # yields from pyhf expected_data call, per-bin / per-channel uncertainty from mock
     assert model_pred.model_yields == [[[25.0, 5.0]], [[8.0]]]
-    assert model_pred.total_stdev_model_bins == [[5.0, 2.0], [1.0]]
-    assert np.allclose(model_pred.total_stdev_model_channels, [5.38516481, 1.0])
+    assert model_pred.total_stdev_model_bins == [
+        [[5.0, 2.0], [5.0, 2.0]],
+        [[1.0], [1.0]],
+    ]
+    assert np.allclose(
+        model_pred.total_stdev_model_channels, [[5.38516481, 5.38516481], [1.0, 1.0]]
+    )
     assert model_pred.label == "pre-fit"
 
     # Asimov parameter calculation and pre-fit uncertainties
@@ -272,8 +283,8 @@ def test_prediction(
     model_pred = model_utils.prediction(model, fit_results=fit_results)
     assert model_pred.model == model
     assert np.allclose(model_pred.model_yields, [[[57.54980000]]])  # new par value
-    assert model_pred.total_stdev_model_bins == [[0.3]]  # from mock
-    assert model_pred.total_stdev_model_channels == [0.3]  # from mock
+    assert model_pred.total_stdev_model_bins == [[[0.3], [0.3]]]  # from mock
+    assert model_pred.total_stdev_model_channels == [[0.3, 0.3]]  # from mock
     assert model_pred.label == "post-fit"
     assert "parameter names in fit results and model do not match" not in [
         rec.message for rec in caplog.records
diff --git a/tests/test_tabulate.py b/tests/test_tabulate.py
index 3a7f5f1..5c7e82a 100644
--- a/tests/test_tabulate.py
+++ b/tests/test_tabulate.py
@@ -22,7 +22,7 @@ def test__yields_per_bin(example_spec_multibin, example_spec_with_background, ca
     # multiple channels
     model = pyhf.Workspace(example_spec_multibin).model()
     yields = [[[25.0, 5.0]], [[8.0]]]
-    total_stdev = [[5.0, 2.0], [1.0]]
+    total_stdev = [[[5.0, 2.0], [5.0, 2.0]], [[1.0], [1.0]]]
     data = [[35, 8], [10]]
     channels = model.config.channels
     label = "pre-fit"
@@ -34,9 +34,9 @@ def test__yields_per_bin(example_spec_multibin, example_spec_with_background, ca
     assert yield_table == [
         {
             "sample": "Signal",
-            "region_1\nbin 1": "25.00",
-            "region_1\nbin 2": "5.00",
-            "region_2\nbin 1": "8.00",
+            "region_1\nbin 1": "25.00 \u00B1 5.00",
+            "region_1\nbin 2": "5.00 \u00B1 2.00",
+            "region_2\nbin 1": "8.00 \u00B1 1.00",
         },
         {
             "sample": "total",
@@ -57,7 +57,7 @@ def test__yields_per_bin(example_spec_multibin, example_spec_with_background, ca
     # multiple samples
     model = pyhf.Workspace(example_spec_with_background).model()
     yields = [[[150.0], [50.0]]]
-    total_stdev = [[8.60]]
+    total_stdev = [[[6.45], [2.15], [8.60]]]
     data = [[160]]
     channels = model.config.channels
 
@@ -65,8 +65,8 @@ def test__yields_per_bin(example_spec_multibin, example_spec_with_background, ca
         model, yields, total_stdev, data, channels, label
     )
     assert yield_table == [
-        {"sample": "Background", "Signal Region\nbin 1": "150.00"},
-        {"sample": "Signal", "Signal Region\nbin 1": "50.00"},
+        {"sample": "Background", "Signal Region\nbin 1": "150.00 \u00B1 6.45"},
+        {"sample": "Signal", "Signal Region\nbin 1": "50.00 \u00B1 2.15"},
         {"sample": "total", "Signal Region\nbin 1": "200.00 \u00B1 8.60"},
         {"sample": "data", "Signal Region\nbin 1": "160.00"},
     ]
@@ -77,18 +77,18 @@ def test__save_tables(tmp_path):
         "yields_per_bin": [
             {
                 "sample": "Background",
-                "Signal_region\nbin 1": "111.00",
-                "Signal_region\nbin 2": "116.00",
+                "Signal_region\nbin 1": "111.00 \u00B1 9.50",
+                "Signal_region\nbin 2": "116.00 \u00B1 10.00",
             },
             {
                 "sample": "Signal",
-                "Signal_region\nbin 1": "1.00",
-                "Signal_region\nbin 2": "5.00",
+                "Signal_region\nbin 1": "1.00 \u00B1 0.50",
+                "Signal_region\nbin 2": "5.00 \u00B1 1.00",
             },
             {
                 "sample": "total",
-                "Signal_region\nbin 1": "112.00 ± 10.00",
-                "Signal_region\nbin 2": "121.00 ± 11.00",
+                "Signal_region\nbin 1": "112.00 \u00B1 10.00",
+                "Signal_region\nbin 2": "121.00 \u00B1 11.00",
             },
             {
                 "sample": "data",
@@ -97,9 +97,9 @@ def test__save_tables(tmp_path):
             },
         ],
         "yields_per_channel": [
-            {"sample": "Background", "Signal_region": "227.00"},
-            {"sample": "Signal", "Signal_region": "6.00"},
-            {"sample": "total", "Signal_region": "233.00 ± 17.00"},
+            {"sample": "Background", "Signal_region": "227.00 \u00B1 16.00"},
+            {"sample": "Signal", "Signal_region": "6.00 \u00B1 1.00"},
+            {"sample": "total", "Signal_region": "233.00 \u00B1 17.00"},
             {"sample": "data", "Signal_region": "235.00"},
         ],
     }
@@ -113,17 +113,17 @@ def test__save_tables(tmp_path):
         "sample      Signal_region    Signal_region\n"
         "            bin 1            bin 2\n"
         "----------  ---------------  ---------------\n"
-        "Background  111.00           116.00\n"
-        "Signal      1.00             5.00\n"
-        "total       112.00 ± 10.00   121.00 ± 11.00\n"
+        "Background  111.00 \u00B1 9.50    116.00 \u00B1 10.00\n"
+        "Signal      1.00 \u00B1 0.50      5.00 \u00B1 1.00\n"
+        "total       112.00 \u00B1 10.00   121.00 \u00B1 11.00\n"
         "data        112.00           123.00\n"
     )
     assert fname_channel.read_text() == (
         "sample      Signal_region\n"
         "----------  ---------------\n"
-        "Background  227.00\n"
-        "Signal      6.00\n"
-        "total       233.00 ± 17.00\n"
+        "Background  227.00 \u00B1 16.00\n"
+        "Signal      6.00 \u00B1 1.00\n"
+        "total       233.00 \u00B1 17.00\n"
         "data        235.00\n"
     )
 
@@ -147,7 +147,7 @@ def test__yields_per_channel(
     # multiple channels
     model = pyhf.Workspace(example_spec_multibin).model()
     yields = [[30], [8.0]]
-    total_stdev = [5.39, 1.0]
+    total_stdev = [[5.39, 5.39], [1.0, 1.0]]
     data = [43, 10]
     channels = model.config.channels
     label = "pre-fit"
@@ -157,7 +157,11 @@ def test__yields_per_channel(
         model, yields, total_stdev, data, channels, label
     )
     assert yield_table == [
-        {"sample": "Signal", "region_1": "30.00", "region_2": "8.00"},
+        {
+            "sample": "Signal",
+            "region_1": "30.00 \u00B1 5.39",
+            "region_2": "8.00 \u00B1 1.00",
+        },
         {
             "sample": "total",
             "region_1": "30.00 \u00B1 5.39",
@@ -173,7 +177,7 @@ def test__yields_per_channel(
     # multiple samples
     model = pyhf.Workspace(example_spec_with_background).model()
     yields = [[150.0, 50]]
-    total_stdev = [8.60]
+    total_stdev = [[6.45, 2.15, 8.60]]
     data = [160]
     channels = model.config.channels
 
@@ -181,8 +185,8 @@ def test__yields_per_channel(
         model, yields, total_stdev, data, channels, label
     )
     assert yield_table == [
-        {"sample": "Background", "Signal Region": "150.00"},
-        {"sample": "Signal", "Signal Region": "50.00"},
+        {"sample": "Background", "Signal Region": "150.00 \u00B1 6.45"},
+        {"sample": "Signal", "Signal Region": "50.00 \u00B1 2.15"},
         {"sample": "total", "Signal Region": "200.00 \u00B1 8.60"},
         {"sample": "data", "Signal Region": "160.00"},
     ]
@@ -205,14 +209,16 @@ def test_yields(
     # to be able to check that they are correctly propagated to the output
     caplog.set_level(logging.DEBUG)
     model = pyhf.Workspace(example_spec).model()
-    model_pred = model_utils.ModelPrediction(model, [[[10.0]]], [[0.3]], [0.3], "pred")
+    model_pred = model_utils.ModelPrediction(
+        model, [[[10.0]]], [[[0.3], [0.3]]], [[0.3, 0.3]], "pred"
+    )
     data = [12.0, 1.0]  # with auxdata to strip via mock
 
     table_dict = tabulate.yields(model_pred, data)
     assert mock_data.call_args_list == [((model, data), {})]
     assert mock_filter.call_args_list == [((model, None), {})]
     assert mock_bin.call_args_list == [
-        ((model, [[[10.0]]], [[0.3]], [[12.0]], ["SR"], "pred"), {})
+        ((model, [[[10.0]]], [[[0.3], [0.3]]], [[12.0]], ["SR"], "pred"), {})
     ]
     assert mock_channel.call_count == 0
     assert mock_save.call_count == 1
@@ -239,7 +245,7 @@ def test_yields(
     )
     assert mock_bin.call_count == 1  # one call from before
     assert mock_channel.call_args_list == [
-        ((model, [[10.0]], [0.3], [12.0], ["SR"], "pred"), {})
+        ((model, [[10.0]], [[0.3, 0.3]], [12.0], ["SR"], "pred"), {})
     ]
     assert table_dict == {"yields_per_channel": [{"yields_per_channel": None}]}
     assert mock_save.call_count == 1  # one call from before, no new call
diff --git a/tests/visualize/test_visualize.py b/tests/visualize/test_visualize.py
index e98ed68..b1bdc92 100644
--- a/tests/visualize/test_visualize.py
+++ b/tests/visualize/test_visualize.py
@@ -149,7 +149,7 @@ def test_data_mc(mock_data, mock_filter, mock_dict, mock_bins, mock_draw, exampl
     figure_folder = "tmp"
     model = pyhf.Workspace(example_spec).model()
     model_pred = model_utils.ModelPrediction(
-        model, [[[10.0]]], [[0.3]], [0.3], "pre-fit"
+        model, [[[10.0]]], [[[0.3], [0.3]]], [[0.3, 0.3]], "pre-fit"
     )
     data = [12.0, 1.0]
 
@@ -198,7 +198,7 @@ def test_data_mc(mock_data, mock_filter, mock_dict, mock_bins, mock_draw, exampl
     # do not save figure, histogram input mode: no binning or variable specified (via
     # side effect)
     model_pred = model_utils.ModelPrediction(
-        model, [[[11.0]]], [[0.2]], [0.2], "post-fit"
+        model, [[[11.0]]], [[[0.2], [0.2]]], [[0.2, 0.2]], "post-fit"
     )
     _ = visualize.data_mc(
         model_pred,