dataeval 0.73.0__py3-none-any.whl → 0.74.0__py3-none-any.whl

{dataeval-0.73.0.dist-info → dataeval-0.74.0.dist-info}/RECORD

@@ -1,36 +1,40 @@
-dataeval/__init__.py,sha256=cAgMAbawI3EC6HdLfV_g_mMpH5Y-zy-n2qzrRKBH_6s,641
-dataeval/detectors/__init__.py,sha256=xdp8LYOFjV5tVbAwu0Y03KU9EajHkSFy_M3raqbxpDc,383
-dataeval/detectors/drift/__init__.py,sha256=MRPWFOaoVoqAHW36nA5F3wk7QXJU4oecND2RbtgG9oY,757
-dataeval/detectors/drift/base.py,sha256=0S-0MFpIFaJ4_8IGreFKSmyna2L50FBn7DVaoNWmw8E,14509
+dataeval/__init__.py,sha256=bwKFegCsdGFydqDvza_wSvJgRGr-0pQ59UpcePQ1mNs,601
+dataeval/detectors/__init__.py,sha256=mwAyY54Hvp6N4D57cde3_besOinK8jVF43k0Mw4XZi8,363
+dataeval/detectors/drift/__init__.py,sha256=BSXm21y7cAawHep-ZldCJ5HOvzYjPzYGKGrmoEs3i0E,737
+dataeval/detectors/drift/base.py,sha256=xwI6C-PEH0ZjpSqP6No6WDZp42DnE16OHi_mXe2JSvI,14499
 dataeval/detectors/drift/cvm.py,sha256=kc59w2_wtxFGNnLcaJRvX5v_38gPXiebSGNiFVdunEQ,4142
 dataeval/detectors/drift/ks.py,sha256=gcpe1WIQeNeZdLYkdMZCFLXUp1bHMQUxwJE6-RLVOXs,4229
 dataeval/detectors/drift/mmd.py,sha256=TqGOnUNYKwpS0GQPV3dSl-_qRa0g2flmoQ-dxzW_JfY,7586
-dataeval/detectors/drift/torch.py,sha256=D46J72OPW8-PpP3w9ODMBfcDSdailIgVjgHVFpbYfws,11649
+dataeval/detectors/drift/torch.py,sha256=igEQ2DV9JmcpTdUKCOHBi5LxtoNeCAslJS2Ldulg1hw,7585
 dataeval/detectors/drift/uncertainty.py,sha256=Xz2yzJjtJfw1vLag234jwRvaa_HK36nMajGx8bQaNRs,5322
 dataeval/detectors/drift/updates.py,sha256=UJ0z5hlunRi7twnkLABfdJG3tT2EqX4y9IGx8_USYvo,1780
 dataeval/detectors/linters/__init__.py,sha256=BvpaB1RUpkEhhXk3Mqi5NYoOcJKZRFSBOJCmQOIfYRU,483
-dataeval/detectors/linters/clusterer.py,sha256=OtBE5rglAGdTTQRmKUHP6J-uWmnh2E3lZxeqJCnc87U,21014
+dataeval/detectors/linters/clusterer.py,sha256=sau5A9YcQ6VDjbZGOIaCaRHW_63opaA31pqHo5Rm-hQ,21018
 dataeval/detectors/linters/duplicates.py,sha256=tOD43rJkvheIA3mznbUqHhft2yD3xRZQdCt61daIca4,5665
 dataeval/detectors/linters/merged_stats.py,sha256=X-bDTwjyR8RuVmzxLaHZmQ5nI3oOWvsqVlitdSncapk,1355
 dataeval/detectors/linters/outliers.py,sha256=BUVvtbKHo04KnRmrgb84MBr0l1gtcY3-xNCHjetFrEQ,10117
-dataeval/detectors/ood/__init__.py,sha256=FVyVuaxVKAOgSTaaBf-j2OXXDarSBFcJ7CTlMV6w88s,661
-dataeval/detectors/ood/ae.py,sha256=XQ_rCsf0VWg_2YXt33XGe6ZgxEud1PfIl7TmBVP1GkM,2347
-dataeval/detectors/ood/aegmm.py,sha256=6UKv0uJYWAzu1F-cITFGly4w9y_t7wqg3OmVyCN365o,2041
-dataeval/detectors/ood/base.py,sha256=a_d52pJMWVmduSt8OvUWYwHE8mpCaI6pIAE4_ib_GOs,8841
-dataeval/detectors/ood/llr.py,sha256=TwUk1RsZhnM5tUssGVMBhWggCW2izs_Asy9QPHkTJaU,10615
-dataeval/detectors/ood/metadata_ks_compare.py,sha256=jH7uDwyyBIIcTrRhQEdnLAdrwf7LfNczKBw0CpJyF5c,4282
+dataeval/detectors/ood/__init__.py,sha256=XckkWVhYbbg9iWVsCPEQN-t7FFSt2a4jmCwAAempkM4,793
+dataeval/detectors/ood/ae.py,sha256=km7buF8LbMmwsyfu1xMOI5CJDnQX1x8_-c04zTGMXRI,2389
+dataeval/detectors/ood/ae_torch.py,sha256=pO9w5221bXR9lEBkE7oakXeE7PXUUR--xcTpmHvOCSk,2142
+dataeval/detectors/ood/aegmm.py,sha256=CI2HEkRMJSEFTVLZEhz4CStkaS7i66yTPtnbkbCqTes,2084
+dataeval/detectors/ood/base.py,sha256=u9S7z7zJ8wuPqrtn63ePdAa8DdI579EbCy8Tn0M3XI8,6983
+dataeval/detectors/ood/base_tf.py,sha256=ppj8rAjXjHEab2oGfQO2olXyN4aGZH8_QHIEghOoeFQ,3297
+dataeval/detectors/ood/base_torch.py,sha256=yFbSfQsBMwZeVf8mrixmkZYBGChhV5oAHtkgzWnMzsA,3405
+dataeval/detectors/ood/llr.py,sha256=IrOam-kqUU4bftolR3MvhcEq-NNj2euyI-lYvMuXYn8,10645
+dataeval/detectors/ood/metadata_ks_compare.py,sha256=Ka6MABdJH5ZlHF66mENpSOLCE8H9xdQ_wWNwMYVO_Q0,5352
 dataeval/detectors/ood/metadata_least_likely.py,sha256=nxMCXUOjOfWHDTGT2SLE7OYBCydRq8zHLd8t17k7hMM,5193
 dataeval/detectors/ood/metadata_ood_mi.py,sha256=KLay2BmgHrStBV92VpIs_B1yEfQKllsMTgzOQEng01I,4065
-dataeval/detectors/ood/vae.py,sha256=UKrQNFdHcnxAY0fAFbLrXasY8Z6qg138BXxqwc1hlts,3154
-dataeval/detectors/ood/vaegmm.py,sha256=_wwmT37URs0MyhbORk91XJExClv-4e15LH_Bj60Pw1w,2409
+dataeval/detectors/ood/vae.py,sha256=yjK4p-XYhnH3wWPiwAclb3eyZE0wpTazLLuKhzurcWY,3203
+dataeval/detectors/ood/vaegmm.py,sha256=FhPJBzs7wyEPQUUMxOMsdPpCdAZwN82vztjt05cSrds,2459
 dataeval/interop.py,sha256=TZCkZo844DvzHoxuRo-YsBhT6GvKmyQTHtUEQZPly1M,1728
 dataeval/metrics/__init__.py,sha256=fPBNLd-T6mCErZBBJrxWmXIL0jCk7fNUYIcNEBkMa80,238
-dataeval/metrics/bias/__init__.py,sha256=puf645-hAO5hFHNHlZ239TPopqWIoN-uLGXFB8-hA_o,599
-dataeval/metrics/bias/balance.py,sha256=Uz7RHf3UuiAxfYlZpKMg4jMzXwXcEfYj7BUnUjzgkw0,8579
-dataeval/metrics/bias/coverage.py,sha256=eB8PacN_uJ19pMd5SVI3N98NC2KJMgE3tgI-DJFNHYs,4497
-dataeval/metrics/bias/diversity.py,sha256=v9fiuySovMajW9Re0EH_FdbuJryAAdVKkvOuNngO5nc,9618
-dataeval/metrics/bias/metadata.py,sha256=OZB9BzPW6JMq2kTp_a9ucqRNcPpfqOexINax1jH5vVQ,11318
-dataeval/metrics/bias/parity.py,sha256=vfGnt_GoGMjMfWgY1FjqNV-gjqVq13tsTTmVkNtRfDM,17120
+dataeval/metrics/bias/__init__.py,sha256=dYiPHenS8J7pgRMMW2jNkTBmTbPoYTxT04fZu9PFats,747
+dataeval/metrics/bias/balance.py,sha256=BH644D_xN7rRUdJMNgVcGHWq3TTnehYjSBhSMhmAFyY,9154
+dataeval/metrics/bias/coverage.py,sha256=LBrNG6GIrvMJjZckr72heyCTMCke_p5BT8NJWi-noEY,4546
+dataeval/metrics/bias/diversity.py,sha256=__7I934sVoymXqgHoneXglJhIU5iHRIuklFwC2ks84w,7719
+dataeval/metrics/bias/metadata_preprocessing.py,sha256=DbtzsiHjkCxs411okb6s2B_H2TqfvwJ4xyt9m_OsqJo,12266
+dataeval/metrics/bias/metadata_utils.py,sha256=HmTjlRRTdM9566oKUDDdVMJ8luss4DYykFOiS2FQzhM,6558
+dataeval/metrics/bias/parity.py,sha256=lLa2zN0AK-zWzlXmvLCbMxTZFodAKLs8wSGl_YZdNFo,12765
 dataeval/metrics/estimators/__init__.py,sha256=O6ocxJq8XDkfJWwXeJnnnzbOyRnFPKF4kTIVTTZYOA8,380
 dataeval/metrics/estimators/ber.py,sha256=SVT-BIC_GLs0l2l2NhWu4OpRbgn96w-OwTSoPHTnQbE,5037
 dataeval/metrics/estimators/divergence.py,sha256=pImaa216-YYTgGWDCSTcpJrC-dfl7150yVrPfW_TyGc,4293
@@ -46,14 +50,15 @@ dataeval/metrics/stats/pixelstats.py,sha256=x90O10IqVjEORtYwueFLvJnVYTxhPBOOx5HM
 dataeval/metrics/stats/visualstats.py,sha256=y0xIvst7epcajk8vz2jngiAiz0T7DZC-M97Rs1-vV9I,4950
 dataeval/output.py,sha256=jWXXNxFNBEaY1rN7Z-6LZl6bQT-I7z_wqr91Rhrdt_0,3061
 dataeval/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-dataeval/utils/__init__.py,sha256=Qr-D0yHnDE8qit0-Wf6xmdMX9Wle2p_mXKgTueTy5GA,753
+dataeval/utils/__init__.py,sha256=FZLWDA7nMbHOcdg3701cVJpQmUp1Wxxk8h_qIrUQQjY,713
+dataeval/utils/gmm.py,sha256=YuLsJKsVWgH_wHr1u_hSRH5Yeexdj8exht8h99L7bLo,561
 dataeval/utils/image.py,sha256=KgC_1nW__nGN5q6bVZNvG4U_qIBdjcPATz9qe8f2XuA,1928
 dataeval/utils/lazy.py,sha256=M0iBHuJh4UPrSJPHZ0jhFwRSZhyjHJQx_KEf1OCkHD8,588
-dataeval/utils/metadata.py,sha256=A6VN7KbdiOA6rUQvUGKwDcvtOyjBer8bRW_wFxNhmW0,8556
-dataeval/utils/shared.py,sha256=BvEeYPMNQTmx4LSaImGeC0VkvcbEY3Byqtxa-jQ3xgc,3623
-dataeval/utils/split_dataset.py,sha256=IopyxwC3FaZwgVriW4OXze-mDMpOlvRr83OADA5Jydk,19454
+dataeval/utils/metadata.py,sha256=0A--iru0zEmi044mKz5P35q69KrI30yoiRSlvs7TSdQ,9418
+dataeval/utils/shared.py,sha256=xvF3VLfyheVwJtdtDrneOobkKf7t-JTmf_w91FWXmqo,3616
+dataeval/utils/split_dataset.py,sha256=Ot1ZJhbIhVfcShYXF9MkWXak5odBXyuBdRh-noXh-MI,19555
 dataeval/utils/tensorflow/__init__.py,sha256=l4OjIA75JJXeNWDCkST1xtDMVYsw97lZ-9JXFBlyuYg,539
-dataeval/utils/tensorflow/_internal/gmm.py,sha256=RIFx8asEpi2kMf8JVzq9M3aAvNe9fjpJPf3BzWE-aeE,3787
+dataeval/utils/tensorflow/_internal/gmm.py,sha256=XvjhWM3ppP-R9nCZGs80WphmQR3u7wb-VtoCQYeXZlQ,3404
 dataeval/utils/tensorflow/_internal/loss.py,sha256=TFhoNPgqeJtdpIHYobZPyzMpeWjzlFqzu5LCtthEUi4,4463
 dataeval/utils/tensorflow/_internal/models.py,sha256=TzQYRrFe5XomhnPw05v-HBODQdFIqWg21WH1xS0XBlg,59868
 dataeval/utils/tensorflow/_internal/trainer.py,sha256=uBFTnAy9o2T_FoT3RSX-AA7T-2FScyOdYEg9_7Dpd28,4314
@@ -61,13 +66,14 @@ dataeval/utils/tensorflow/_internal/utils.py,sha256=lr5hKkAPbjMCUNIzMUIqbEddwbWQ
 dataeval/utils/tensorflow/loss/__init__.py,sha256=Q-66vt91Oe1ByYfo28tW32zXDq2MqQ2gngWgmIVmof8,227
 dataeval/utils/torch/__init__.py,sha256=lpkqfgyARUxgrV94cZESQv8PIP2p-UnwItZ_wIr0XzQ,675
 dataeval/utils/torch/blocks.py,sha256=HVhBTMMD5NA4qheMUgyol1KWiKZDIuc8k5j4RcMKmhk,1466
-dataeval/utils/torch/datasets.py,sha256=9YV9-Uhq6NCMuu1hPhMnQXjmeI-Ld8ve1z_haxre88o,15023
-dataeval/utils/torch/models.py,sha256=0BsXmLK8W1OZ8nnEGb1f9LzIeCgtevQC37dvKS1v1vA,3236
-dataeval/utils/torch/trainer.py,sha256=EraOKiXxiMNiycStZNMR5yRz3ehgp87d9ewR9a9dV4w,5559
-dataeval/utils/torch/utils.py,sha256=FI4LJ6DvXFQJVff8fxSCP7LRkp8H9BIUgYX0kk7_Cuo,1537
-dataeval/workflows/__init__.py,sha256=x2JnOoKmLUCZOsB6RNPqMdVvxEb6Hpda5GPJnD_k0v0,310
+dataeval/utils/torch/datasets.py,sha256=10elNgLuH_FDX_CHE3y2Z215JN4-PQovQm5brcIJOeM,15021
+dataeval/utils/torch/gmm.py,sha256=VbLlUQohwToApT493_tjQBWy2UM5R-3ppS9Dp-eP7BA,3240
+dataeval/utils/torch/models.py,sha256=sdGeo7a8vshCTGA4lYyVxxb_aDWUlxdtIVxrddS-_ls,8542
+dataeval/utils/torch/trainer.py,sha256=8BEXr6xtk-CHJTcNxOBnWgkFWfJUAiBy28cEdBhLMRU,7883
+dataeval/utils/torch/utils.py,sha256=nWRcT6z6DbFVrL1RyxCOX3DPoCrv9G0B-VI_9LdGCQQ,5784
+dataeval/workflows/__init__.py,sha256=ef1MiVL5IuhlDXXbwsiAfafhnr7tD3TXF9GRusy9_O8,290
 dataeval/workflows/sufficiency.py,sha256=1jSYhH9i4oesmJYs5PZvWS1LGXf8ekOgNhpFtMPLPXk,18552
-dataeval-0.73.0.dist-info/LICENSE.txt,sha256=Kpzcfobf1HlqafF-EX6dQLw9TlJiaJzfgvLQFukyXYw,1060
-dataeval-0.73.0.dist-info/METADATA,sha256=YVw0z5C5BZs-9gCxCmbo4aNIN7Ph3rZsel7FofFrMKY,4714
-dataeval-0.73.0.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-dataeval-0.73.0.dist-info/RECORD,,
+dataeval-0.74.0.dist-info/LICENSE.txt,sha256=Kpzcfobf1HlqafF-EX6dQLw9TlJiaJzfgvLQFukyXYw,1060
+dataeval-0.74.0.dist-info/METADATA,sha256=OPnkHZTm8R1LHqLxcSnOHjqj5GuHmjUVI3dddTVsBAc,4680
+dataeval-0.74.0.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+dataeval-0.74.0.dist-info/RECORD,,

dataeval/metrics/bias/metadata.py

@@ -1,358 +0,0 @@
-from __future__ import annotations
-
-__all__ = []
-
-import contextlib
-from typing import Any, Mapping
-
-import numpy as np
-from numpy.typing import ArrayLike, NDArray
-from scipy.stats import entropy as sp_entropy
-
-from dataeval.interop import to_numpy
-
-with contextlib.suppress(ImportError):
-    from matplotlib.figure import Figure
-
-CLASS_LABEL = "class_label"
-
-
-def get_counts(
-    data: NDArray[np.int_], names: list[str], is_categorical: list[bool], subset_mask: NDArray[np.bool_] | None = None
-) -> tuple[dict[str, NDArray[np.int_]], dict[str, NDArray[np.int_]]]:
-    """
-    Initialize dictionary of histogram counts --- treat categorical values
-    as histogram bins.
-
-    Parameters
-    ----------
-    subset_mask: NDArray[np.bool_] | None
-        Boolean mask of samples to bin (e.g. when computing per class).  True -> include in histogram counts
-
-    Returns
-    -------
-    counts: Dict
-        histogram counts per metadata factor in `factors`.  Each
-        factor will have a different number of bins.  Counts get reused
-        across metrics, so hist_counts are cached but only if computed
-        globally, i.e. without masked samples.
-    """
-
-    hist_counts, hist_bins = {}, {}
-    # np.where needed to satisfy linter
-    mask = np.where(subset_mask if subset_mask is not None else np.ones(data.shape[0], dtype=bool))
-
-    for cdx, fn in enumerate(names):
-        # linter doesn't like double indexing
-        col_data = data[mask, cdx].squeeze()
-        if is_categorical[cdx]:
-            # if discrete, use unique values as bins
-            bins, cnts = np.unique(col_data, return_counts=True)
-        else:
-            bins = hist_bins.get(fn, "auto")
-            cnts, bins = np.histogram(col_data, bins=bins, density=True)
-
-        hist_counts[fn] = cnts
-        hist_bins[fn] = bins
-
-    return hist_counts, hist_bins
-
-
-def entropy(
-    data: NDArray[Any],
-    names: list[str],
-    is_categorical: list[bool],
-    normalized: bool = False,
-    subset_mask: NDArray[np.bool_] | None = None,
-) -> NDArray[np.float64]:
-    """
-    Meant for use with :term:`bias<Bias>` metrics, :term:`balance<Balance>`, :term:`diversity<Diversity>`,
-    ClasswiseBalance, and Classwise Diversity.
-
-    Compute entropy for discrete/categorical variables and for continuous variables through standard
-    histogram binning.
-
-    Parameters
-    ----------
-    normalized: bool
-        Flag that determines whether or not to normalize entropy by log(num_bins)
-    subset_mask: NDArray[np.bool_] | None
-        Boolean mask of samples to bin (e.g. when computing per class).  True -> include in histogram counts
-
-    Note
-    ----
-    For continuous variables, histogram bins are chosen automatically.  See
-    numpy.histogram for details.
-
-    Returns
-    -------
-    ent: NDArray[np.float64]
-        Entropy estimate per column of X
-
-    See Also
-    --------
-    numpy.histogram
-    scipy.stats.entropy
-    """
-
-    num_factors = len(names)
-    hist_counts, _ = get_counts(data, names, is_categorical, subset_mask)
-
-    ev_index = np.empty(num_factors)
-    for col, cnts in enumerate(hist_counts.values()):
-        # entropy in nats, normalizes counts
-        ev_index[col] = sp_entropy(cnts)
-        if normalized:
-            if len(cnts) == 1:
-                # log(0)
-                ev_index[col] = 0
-            else:
-                ev_index[col] /= np.log(len(cnts))
-    return ev_index
-
-
-def get_num_bins(
-    data: NDArray[Any], names: list[str], is_categorical: list[bool], subset_mask: NDArray[np.bool_] | None = None
-) -> NDArray[np.float64]:
-    """
-    Number of bins or unique values for each metadata factor, used to
-    normalize entropy/:term:`diversity<Diversity>`.
-
-    Parameters
-    ----------
-    subset_mask: NDArray[np.bool_] | None
-        Boolean mask of samples to bin (e.g. when computing per class).  True -> include in histogram counts
-
-    Returns
-    -------
-    NDArray[np.float64]
-    """
-    # likely cached
-    hist_counts, _ = get_counts(data, names, is_categorical, subset_mask)
-    num_bins = np.empty(len(hist_counts))
-    for idx, cnts in enumerate(hist_counts.values()):
-        num_bins[idx] = len(cnts)
-
-    return num_bins
-
-
-def infer_categorical(arr: NDArray[Any], threshold: float = 0.2) -> NDArray[Any]:
-    """
-    Compute fraction of feature values that are unique --- intended to be used
-    for inferring whether variables are categorical.
-    """
-    if arr.ndim == 1:
-        arr = np.expand_dims(arr, axis=1)
-    num_samples = arr.shape[0]
-    pct_unique = np.empty(arr.shape[1])
-    for col in range(arr.shape[1]):  # type: ignore
-        uvals = np.unique(arr[:, col], axis=0)
-        pct_unique[col] = len(uvals) / num_samples
-    return pct_unique < threshold
-
-
-def preprocess_metadata(
-    class_labels: ArrayLike, metadata: Mapping[str, ArrayLike], cat_thresh: float = 0.2
-) -> tuple[NDArray[Any], list[str], list[bool], NDArray[np.str_]]:
-    # if class_labels is not numeric
-    class_array = to_numpy(class_labels)
-    if not np.issubdtype(class_array.dtype, np.number):
-        unique_classes, numerical_labels = np.unique(class_array, return_inverse=True)
-    else:
-        numerical_labels = np.asarray(class_array, dtype=int)
-        unique_classes = np.unique(class_array)
-
-    # convert class_labels and dict of lists to matrix of metadata values
-    preprocessed_metadata = {CLASS_LABEL: numerical_labels}
-
-    # map columns of dict that are not numeric (e.g. string) to numeric values
-    # that mutual information and diversity functions can accommodate.  Each
-    # unique string receives a unique integer value.
-    for k, v in metadata.items():
-        if k == CLASS_LABEL:
-            k = "label_class"
-        # if not numeric
-        v = to_numpy(v)
-        if not np.issubdtype(v.dtype, np.number):
-            _, mapped_vals = np.unique(v, return_inverse=True)
-            preprocessed_metadata[k] = mapped_vals
-        else:
-            preprocessed_metadata[k] = v
-
-    data = np.stack(list(preprocessed_metadata.values()), axis=-1)
-    names = list(preprocessed_metadata.keys())
-    is_categorical = [infer_categorical(preprocessed_metadata[var], cat_thresh)[0] for var in names]
-
-    return data, names, is_categorical, unique_classes
-
-
-def heatmap(
-    data: NDArray[Any],
-    row_labels: list[str] | NDArray[Any],
-    col_labels: list[str] | NDArray[Any],
-    xlabel: str = "",
-    ylabel: str = "",
-    cbarlabel: str = "",
-) -> Figure:
-    """
-    Plots a formatted heatmap
-
-    Parameters
-    ----------
-    data : NDArray
-        Array containing numerical values for factors to plot
-    row_labels : ArrayLike
-        List/Array containing the labels for rows in the histogram
-    col_labels : ArrayLike
-        List/Array containing the labels for columns in the histogram
-    xlabel : str, default ""
-        X-axis label
-    ylabel : str, default ""
-        Y-axis label
-    cbarlabel : str, default ""
-        Label for the colorbar
-    """
-    import matplotlib
-    import matplotlib.pyplot as plt
-
-    fig, ax = plt.subplots(figsize=(10, 10))
-
-    # Plot the heatmap
-    im = ax.imshow(data, vmin=0, vmax=1.0)
-
-    # Create colorbar
-    cbar = fig.colorbar(im, shrink=0.5)
-    cbar.set_ticks([0.0, 0.25, 0.5, 0.75, 1.0])
-    cbar.set_ticklabels(["0.0", "0.25", "0.5", "0.75", "1.0"])
-    cbar.set_label(cbarlabel, loc="center")
-
-    # Show all ticks and label them with the respective list entries.
-    ax.set_xticks(np.arange(data.shape[1]), labels=col_labels)
-    ax.set_yticks(np.arange(data.shape[0]), labels=row_labels)
-
-    ax.tick_params(top=False, bottom=True, labeltop=False, labelbottom=True)
-    # Rotate the tick labels and set their alignment.
-    plt.setp(ax.get_xticklabels(), rotation=45, ha="right", rotation_mode="anchor")
-
-    # Turn spines off and create white grid.
-    ax.spines[:].set_visible(False)
-
-    ax.set_xticks(np.arange(data.shape[1] + 1) - 0.5, minor=True)
-    ax.set_yticks(np.arange(data.shape[0] + 1) - 0.5, minor=True)
-    ax.grid(which="minor", color="w", linestyle="-", linewidth=3)
-    ax.tick_params(which="minor", bottom=False, left=False)
-
-    if xlabel:
-        ax.set_xlabel(xlabel)
-    if ylabel:
-        ax.set_ylabel(ylabel)
-
-    valfmt = matplotlib.ticker.FuncFormatter(format_text)  # type: ignore
-
-    # Normalize the threshold to the images color range.
-    threshold = im.norm(1.0) / 2.0
-
-    # Set default alignment to center, but allow it to be
-    # overwritten by textkw.
-    kw = {"horizontalalignment": "center", "verticalalignment": "center"}
-
-    # Loop over the data and create a `Text` for each "pixel".
-    # Change the text's color depending on the data.
-    textcolors = ("white", "black")
-    texts = []
-    for i in range(data.shape[0]):
-        for j in range(data.shape[1]):
-            kw.update(color=textcolors[int(im.norm(data[i, j]) > threshold)])
-            text = im.axes.text(j, i, valfmt(data[i, j], None), **kw)  # type: ignore
-            texts.append(text)
-
-    fig.tight_layout()
-    return fig
-
-
-# Function to define how the text is displayed in the heatmap
-def format_text(*args: str) -> str:
-    """
-    Helper function to format text for heatmap()
-
-    Parameters
-    ----------
-    *args: Tuple (str, str)
-        Text to be formatted. Second element is ignored, but is a
-        mandatory pass-through argument as per matplotlib.ticket.FuncFormatter
-
-    Returns
-    -------
-    str
-        Formatted text
-    """
-    x = args[0]
-    return f"{x:.2f}".replace("0.00", "0").replace("0.", ".").replace("nan", "")
-
-
-def diversity_bar_plot(labels: NDArray[Any], bar_heights: NDArray[Any]) -> Figure:
-    """
-    Plots a formatted bar plot
-
-    Parameters
-    ----------
-    labels : NDArray
-        Array containing the labels for each bar
-    bar_heights : NDArray
-        Array containing the values for each bar
-    """
-    import matplotlib.pyplot as plt
-
-    fig, ax = plt.subplots(figsize=(10, 10))
-
-    ax.bar(labels, bar_heights)
-    ax.set_xlabel("Factors")
-
-    plt.setp(ax.get_xticklabels(), rotation=45, ha="right", rotation_mode="anchor")
-
-    fig.tight_layout()
-    return fig
-
-
-def coverage_plot(images: NDArray[Any], num_images: int) -> Figure:
-    """
-    Creates a single plot of all of the provided images
-
-    Parameters
-    ----------
-    images : NDArray
-        Array containing only the desired images to plot
-    """
-    import matplotlib.pyplot as plt
-
-    num_images = min(num_images, len(images))
-
-    if images.ndim == 4:
-        images = np.moveaxis(images, 1, -1)
-    elif images.ndim == 3:
-        images = np.repeat(images[:, :, :, np.newaxis], 3, axis=-1)
-    else:
-        raise ValueError(
-            f"Expected a (N,C,H,W) or a (N, H, W) set of images, but got a {images.ndim}-dimensional set of images."
-        )
-
-    rows = np.ceil(num_images / 3).astype(int)
-    fig, axs = plt.subplots(rows, 3, figsize=(9, 3 * rows))
-
-    if rows == 1:
-        for j in range(3):
-            if j >= len(images):
-                continue
-            axs[j].imshow(images[j])
-            axs[j].axis("off")
-    else:
-        for i in range(rows):
-            for j in range(3):
-                i_j = i * 3 + j
-                if i_j >= len(images):
-                    continue
-                axs[i, j].imshow(images[i_j])
-                axs[i, j].axis("off")
-
-    fig.tight_layout()
-    return fig