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

dataeval/detectors/ood/base_torch.py

@@ -0,0 +1,109 @@
+"""
+Source code derived from Alibi-Detect 0.11.4
+https://github.com/SeldonIO/alibi-detect/tree/v0.11.4
+
+Original code Copyright (c) 2023 Seldon Technologies Ltd
+Licensed under Apache Software License (Apache 2.0)
+"""
+
+from __future__ import annotations
+
+from typing import Callable, cast
+
+import torch
+from numpy.typing import ArrayLike
+
+from dataeval.detectors.drift.torch import get_device
+from dataeval.detectors.ood.base import OODBaseMixin, OODFitMixin, OODGMMMixin
+from dataeval.interop import to_numpy
+from dataeval.utils.torch.gmm import gmm_params
+from dataeval.utils.torch.trainer import trainer
+
+
+class OODBase(OODBaseMixin[torch.nn.Module], OODFitMixin[Callable[..., torch.nn.Module], torch.optim.Optimizer]):
+    def __init__(self, model: torch.nn.Module, device: str | torch.device | None = None) -> None:
+        self.device: torch.device = get_device(device)
+        super().__init__(model)
+
+    def fit(
+        self,
+        x_ref: ArrayLike,
+        threshold_perc: float,
+        loss_fn: Callable[..., torch.nn.Module] | None,
+        optimizer: torch.optim.Optimizer | None,
+        epochs: int,
+        batch_size: int,
+        verbose: bool,
+    ) -> None:
+        """
+        Train the model and infer the threshold value.
+
+        Parameters
+        ----------
+        x_ref : ArrayLike
+            Training data.
+        threshold_perc : float, default 100.0
+            Percentage of reference data that is normal.
+        loss_fn : Callable | None, default None
+            Loss function used for training.
+        optimizer : Optimizer, default keras.optimizers.Adam
+            Optimizer used for training.
+        epochs : int, default 20
+            Number of training epochs.
+        batch_size : int, default 64
+            Batch size used for training.
+        verbose : bool, default True
+            Whether to print training progress.
+        """
+
+        # Train the model
+        trainer(
+            model=self.model,
+            x_train=to_numpy(x_ref),
+            y_train=None,
+            loss_fn=loss_fn,
+            optimizer=optimizer,
+            preprocess_fn=None,
+            epochs=epochs,
+            batch_size=batch_size,
+            device=self.device,
+            verbose=verbose,
+        )
+
+        # Infer the threshold values
+        self._ref_score = self.score(x_ref, batch_size)
+        self._threshold_perc = threshold_perc
+
+
+class OODBaseGMM(OODBase, OODGMMMixin[torch.Tensor]):
+    def fit(
+        self,
+        x_ref: ArrayLike,
+        threshold_perc: float,
+        loss_fn: Callable[..., torch.nn.Module] | None,
+        optimizer: torch.optim.Optimizer | None,
+        epochs: int,
+        batch_size: int,
+        verbose: bool,
+    ) -> None:
+        # Train the model
+        trainer(
+            model=self.model,
+            x_train=to_numpy(x_ref),
+            y_train=None,
+            loss_fn=loss_fn,
+            optimizer=optimizer,
+            preprocess_fn=None,
+            epochs=epochs,
+            batch_size=batch_size,
+            device=self.device,
+            verbose=verbose,
+        )
+
+        # Calculate the GMM parameters
+        _, z, gamma = cast(tuple[torch.Tensor, torch.Tensor, torch.Tensor], self.model(x_ref))
+        self._gmm_params = gmm_params(z, gamma)
+
+        # Infer the threshold values
+        self._ref_score = self.score(x_ref, batch_size)
+        self._threshold_perc = threshold_perc

dataeval/detectors/ood/llr.py

@@ -16,7 +16,7 @@ from typing import TYPE_CHECKING, Callable
 import numpy as np
 from numpy.typing import ArrayLike, NDArray
 
-from dataeval.detectors.ood.base import OODBase, OODScoreOutput
+from dataeval.detectors.ood.base import OODBaseMixin, OODScoreOutput
 from dataeval.interop import to_numpy
 from dataeval.utils.lazy import lazyload
 from dataeval.utils.tensorflow._internal.trainer import trainer
@@ -96,7 +96,7 @@ def _mutate_categorical(
     return tf.cast(X, tf.float32)  # type: ignore
 
 
-class OOD_LLR(OODBase):
+class OOD_LLR(OODBaseMixin[tf_models.PixelCNN]):
     """
     Likelihood Ratios based outlier detector.
 

dataeval/detectors/ood/metadata_ks_compare.py

@@ -2,17 +2,45 @@ from __future__ import annotations
 
 import numbers
 import warnings
-from typing import Any, Mapping
+from dataclasses import dataclass
+from typing import Any, Mapping, NamedTuple
 
 import numpy as np
 from numpy.typing import NDArray
 from scipy.stats import iqr, ks_2samp
 from scipy.stats import wasserstein_distance as emd
 
+from dataeval.output import OutputMetadata, set_metadata
 
+
+class MetadataKSResult(NamedTuple):
+    statistic: float
+    statistic_location: float
+    shift_magnitude: float
+    pvalue: float
+
+
+@dataclass(frozen=True)
+class KSOutput(OutputMetadata):
+    """
+    Output class for results of ks_2samp featurewise comparisons of new metadata to reference metadata.
+
+    Attributes
+    ----------
+    mdc : dict[str, dict[str, float]]
+        dict keyed by metadata feature names. Each value contains four floats, which are the KS statistic itself, its
+        location within the range of the reference metadata, the shift of new metadata relative to reference, the
+        p-value from the KS two-sample test.
+
+    """
+
+    mdc: dict[str, MetadataKSResult]
+
+
+@set_metadata()
 def meta_distribution_compare(
     md0: Mapping[str, list[Any] | NDArray[Any]], md1: Mapping[str, list[Any] | NDArray[Any]]
-) -> dict[str, dict[str, float]]:
+) -> KSOutput:
     """Measures the featurewise distance between two metadata distributions, and computes a p-value to evaluate its
         significance.
 
@@ -43,27 +71,29 @@ def meta_distribution_compare(
     >>> import numpy
     >>> md0 = {"time": [1.2, 3.4, 5.6], "altitude": [235, 6789, 101112]}
     >>> md1 = {"time": [7.8, 9.10, 11.12], "altitude": [532, 9876, 211101]}
-    >>> md_out = meta_distribution_compare(md0, md1)
+    >>> md_out = meta_distribution_compare(md0, md1).mdc
     >>> for k, v in md_out.items():
     >>>     print(k)
     >>>     for kv in v:
     >>>         print("\t", f"{kv}: {v[kv]:.3f}")
     time
-             statistic_location: 0.444
-             shift_magnitude: 2.700
-             pvalue: 0.000
+            statistic: 1.000
+            statistic_location: 0.444
+            shift_magnitude: 2.700
+            pvalue: 0.000
     altitude
-             statistic_location: 0.478
-             shift_magnitude: 0.749
-             pvalue: 0.944
+            statistic: 0.333
+            statistic_location: 0.478
+            shift_magnitude: 0.749
+            pvalue: 0.944
     """
 
     if (metadata_keys := md0.keys()) != md1.keys():
         raise ValueError(f"Both sets of metadata keys must be identical: {list(md0)}, {list(md1)}")
 
-    mdc_dict = {}  # output dict
+    mdc = {}  # output dict
     for k in metadata_keys:
-        mdc_dict.update({k: {}})
+        mdc.update({k: {}})
 
         x0, x1 = list(md0[k]), list(md1[k])
 
@@ -81,7 +111,9 @@ def meta_distribution_compare(
 
         xmin, xmax = min(allx), max(allx)
         if xmin == xmax:  # only one value in this feature, so fill in the obvious results for feature k
-            mdc_dict[k].update({"statistic_location": 0.0, "shift_magnitude": 0.0, "pvalue": 1.0})
+            mdc[k] = MetadataKSResult(
+                **{"statistic": 0.0, "statistic_location": 0.0, "shift_magnitude": 0.0, "pvalue": 1.0}
+            )
             continue
 
         ks_result = ks_2samp(x0, x1, method="asymp")
@@ -94,6 +126,13 @@ def meta_distribution_compare(
 
         drift = emd(x0, x1) / dX
 
-        mdc_dict[k].update({"statistic_location": loc, "shift_magnitude": drift, "pvalue": ks_result.pvalue})  #  pyright: ignore
+        mdc[k] = MetadataKSResult(
+            **{
+                "statistic": ks_result.statistic,  #  pyright: ignore
+                "statistic_location": loc,
+                "shift_magnitude": drift,
+                "pvalue": ks_result.pvalue,  #  pyright: ignore
+            }
+        )
 
-    return mdc_dict
+    return KSOutput(mdc)

dataeval/detectors/ood/vae.py

@@ -15,7 +15,8 @@ from typing import TYPE_CHECKING, Callable
 import numpy as np
 from numpy.typing import ArrayLike
 
-from dataeval.detectors.ood.base import OODBase, OODScoreOutput
+from dataeval.detectors.ood.base import OODScoreOutput
+from dataeval.detectors.ood.base_tf import OODBase
 from dataeval.interop import to_numpy
 from dataeval.utils.lazy import lazyload
 from dataeval.utils.tensorflow._internal.loss import Elbo
@@ -67,7 +68,7 @@ class OOD_VAE(OODBase):
         self,
         x_ref: ArrayLike,
         threshold_perc: float = 100.0,
-        loss_fn: Callable[..., tf.Tensor] = Elbo(0.05),
+        loss_fn: Callable[..., tf.Tensor] | None = Elbo(0.05),
         optimizer: keras.optimizers.Optimizer | None = None,
         epochs: int = 20,
         batch_size: int = 64,

dataeval/detectors/ood/vaegmm.py

@@ -15,7 +15,8 @@ from typing import TYPE_CHECKING, Callable
 import numpy as np
 from numpy.typing import ArrayLike
 
-from dataeval.detectors.ood.base import OODGMMBase, OODScoreOutput
+from dataeval.detectors.ood.base import OODScoreOutput
+from dataeval.detectors.ood.base_tf import OODBaseGMM
 from dataeval.interop import to_numpy
 from dataeval.utils.lazy import lazyload
 from dataeval.utils.tensorflow._internal.gmm import gmm_energy
@@ -33,7 +34,7 @@ else:
     tf_models = lazyload("dataeval.utils.tensorflow._internal.models")
 
 
-class OOD_VAEGMM(OODGMMBase):
+class OOD_VAEGMM(OODBaseGMM):
     """
     VAE with Gaussian Mixture Model based outlier detector.
 
@@ -53,7 +54,7 @@ class OOD_VAEGMM(OODGMMBase):
         self,
         x_ref: ArrayLike,
         threshold_perc: float = 100.0,
-        loss_fn: Callable[..., tf.Tensor] = LossGMM(elbo=Elbo(0.05)),
+        loss_fn: Callable[..., tf.Tensor] | None = LossGMM(elbo=Elbo(0.05)),
         optimizer: keras.optimizers.Optimizer | None = None,
         epochs: int = 20,
         batch_size: int = 64,
@@ -69,7 +70,7 @@ class OOD_VAEGMM(OODGMMBase):
         _, z, _ = predict_batch(X_samples, self.model, batch_size=batch_size)
 
         # compute average energy for samples
-        energy, _ = gmm_energy(z, self.gmm_params, return_mean=False)
+        energy, _ = gmm_energy(z, self._gmm_params, return_mean=False)
         energy_samples = energy.numpy().reshape((-1, self.samples))  # type: ignore
         iscore = np.mean(energy_samples, axis=-1)
         return OODScoreOutput(iscore)

dataeval/metrics/bias/__init__.py

@@ -6,6 +6,7 @@ representation which may impact model performance.
 from dataeval.metrics.bias.balance import BalanceOutput, balance
 from dataeval.metrics.bias.coverage import CoverageOutput, coverage
 from dataeval.metrics.bias.diversity import DiversityOutput, diversity
+from dataeval.metrics.bias.metadata_preprocessing import MetadataOutput, metadata_preprocessing
 from dataeval.metrics.bias.parity import ParityOutput, label_parity, parity
 
 __all__ = [
@@ -14,8 +15,10 @@ __all__ = [
     "diversity",
     "label_parity",
     "parity",
+    "metadata_preprocessing",
     "BalanceOutput",
     "CoverageOutput",
     "DiversityOutput",
     "ParityOutput",
+    "MetadataOutput",
 ]

dataeval/metrics/bias/balance.py

@@ -5,13 +5,15 @@ __all__ = ["BalanceOutput", "balance"]
 import contextlib
 import warnings
 from dataclasses import dataclass
-from typing import Any, Mapping
+from typing import Any
 
 import numpy as np
-from numpy.typing import ArrayLike, NDArray
+import scipy as sp
+from numpy.typing import NDArray
 from sklearn.feature_selection import mutual_info_classif, mutual_info_regression
 
-from dataeval.metrics.bias.metadata import entropy, heatmap, preprocess_metadata
+from dataeval.metrics.bias.metadata_preprocessing import MetadataOutput
+from dataeval.metrics.bias.metadata_utils import get_counts, heatmap
 from dataeval.output import OutputMetadata, set_metadata
 
 with contextlib.suppress(ImportError):
@@ -31,17 +33,17 @@ class BalanceOutput(OutputMetadata):
         Estimate of inter/intra-factor mutual information
     classwise : NDArray[np.float64]
         Estimate of mutual information between metadata factors and individual class labels
-    class_list: NDArray
-        Array of the class labels present in the dataset
-    metadata_names: list[str]
+    factor_names : list[str]
         Names of each metadata factor
+    class_list : NDArray
+        Array of the class labels present in the dataset
     """
 
     balance: NDArray[np.float64]
     factors: NDArray[np.float64]
     classwise: NDArray[np.float64]
+    factor_names: list[str]
     class_list: NDArray[Any]
-    metadata_names: list[str]
 
     def plot(
         self,
@@ -54,9 +56,9 @@ class BalanceOutput(OutputMetadata):
 
         Parameters
         ----------
-        row_labels : ArrayLike | None, default None
+        row_labels : ArrayLike or None, default None
             List/Array containing the labels for rows in the histogram
-        col_labels : ArrayLike | None, default None
+        col_labels : ArrayLike or None, default None
             List/Array containing the labels for columns in the histogram
         plot_classwise : bool, default False
             Whether to plot per-class balance instead of global balance
@@ -65,7 +67,7 @@ class BalanceOutput(OutputMetadata):
             if row_labels is None:
                 row_labels = self.class_list
             if col_labels is None:
-                col_labels = np.concatenate((["class"], self.metadata_names))
+                col_labels = self.factor_names
 
             fig = heatmap(
                 self.classwise,
@@ -83,7 +85,7 @@ class BalanceOutput(OutputMetadata):
             # Finalize the data for the plot, last row is last factor x last factor so it gets dropped
             heat_data = np.where(mask, np.nan, data)[:-1]
             # Creating label array for heat map axes
-            heat_labels = np.concatenate((["class"], self.metadata_names))
+            heat_labels = self.factor_names
 
             if row_labels is None:
                 row_labels = heat_labels[:-1]
@@ -95,7 +97,7 @@ class BalanceOutput(OutputMetadata):
         return fig
 
 
-def validate_num_neighbors(num_neighbors: int) -> int:
+def _validate_num_neighbors(num_neighbors: int) -> int:
     if not isinstance(num_neighbors, (int, float)):
         raise TypeError(
             f"Variable {num_neighbors} is not real-valued numeric type."
@@ -116,19 +118,17 @@ def validate_num_neighbors(num_neighbors: int) -> int:
 
 
 @set_metadata("dataeval.metrics")
-def balance(class_labels: ArrayLike, metadata: Mapping[str, ArrayLike], num_neighbors: int = 5) -> BalanceOutput:
+def balance(
+    metadata: MetadataOutput,
+    num_neighbors: int = 5,
+) -> BalanceOutput:
     """
     Mutual information (MI) between factors (class label, metadata, label/image properties)
 
     Parameters
     ----------
-    class_labels: ArrayLike
-        List of class labels for each image
-    metadata: Mapping[str, ArrayLike]
-        Dict of lists of metadata factors for each image
-    num_neighbors: int, default 5
-        Number of nearest neighbors to use for computing MI between discrete
-        and continuous variables.
+    metadata : MetadataOutput
+        Output after running `metadata_preprocessing`
 
     Returns
     -------
@@ -140,30 +140,34 @@ def balance(class_labels: ArrayLike, metadata: Mapping[str, ArrayLike], num_neig
     ----
     We use `mutual_info_classif` from sklearn since class label is categorical.
     `mutual_info_classif` outputs are consistent up to O(1e-4) and depend on a random
-    seed. MI is computed differently for categorical and continuous variables, and
-    we attempt to infer whether a variable is categorical by the fraction of unique
-    values in the dataset.
+    seed. MI is computed differently for categorical and continuous variables.
 
     Example
     -------
     Return balance (mutual information) of factors with class_labels
 
-    >>> bal = balance(class_labels, metadata)
+    >>> bal = balance(metadata)
     >>> bal.balance
-    array([0.99999822, 0.13363788, 0.04505382, 0.02994455])
+    array([0.9999982 , 0.2494567 , 0.02994455, 0.13363788, 0.        ,
+           0.        ])
 
     Return intra/interfactor balance (mutual information)
 
     >>> bal.factors
-    array([[0.99999843, 0.04133555, 0.09725766],
-           [0.04133555, 0.08433558, 0.1301489 ],
-           [0.09725766, 0.1301489 , 0.99999856]])
+    array([[0.99999935, 0.31360499, 0.26925848, 0.85201924, 0.36653548],
+           [0.31360499, 0.99999856, 0.09725766, 0.15836905, 1.98031993],
+           [0.26925848, 0.09725766, 0.99999846, 0.03713108, 0.01544656],
+           [0.85201924, 0.15836905, 0.03713108, 0.47450653, 0.25509664],
+           [0.36653548, 1.98031993, 0.01544656, 0.25509664, 1.06260686]])
 
     Return classwise balance (mutual information) of factors with individual class_labels
 
     >>> bal.classwise
-    array([[0.99999822, 0.13363788, 0.        , 0.        ],
-           [0.99999822, 0.13363788, 0.        , 0.        ]])
+    array([[0.9999982 , 0.2494567 , 0.02994455, 0.13363788, 0.        ,
+            0.        ],
+           [0.9999982 , 0.2494567 , 0.02994455, 0.13363788, 0.        ,
+            0.        ]])
+
 
     See Also
     --------
@@ -171,69 +175,78 @@ def balance(class_labels: ArrayLike, metadata: Mapping[str, ArrayLike], num_neig
     sklearn.feature_selection.mutual_info_regression
     sklearn.metrics.mutual_info_score
     """
-    num_neighbors = validate_num_neighbors(num_neighbors)
-    data, names, is_categorical, unique_labels = preprocess_metadata(class_labels, metadata)
-    num_factors = len(names)
-    mi = np.empty((num_factors, num_factors))
-    mi[:] = np.nan
+    num_neighbors = _validate_num_neighbors(num_neighbors)
+
+    num_factors = metadata.total_num_factors
+    is_discrete = [True] * (len(metadata.discrete_factor_names) + 1) + [False] * len(metadata.continuous_factor_names)
+    mi = np.full((num_factors, num_factors), np.nan, dtype=np.float32)
+    data = np.hstack((metadata.class_labels[:, np.newaxis], metadata.discrete_data))
+    discretized_data = data
+    if metadata.continuous_data is not None:
+        data = np.hstack((data, metadata.continuous_data))
+        discrete_idx = [metadata.discrete_factor_names.index(name) for name in metadata.continuous_factor_names]
+        discretized_data = np.hstack((discretized_data, metadata.discrete_data[:, discrete_idx]))
 
     for idx in range(num_factors):
-        tgt = data[:, idx].astype(int)
-
-        if is_categorical[idx]:
-            mi[idx, :] = mutual_info_classif(
+        if idx >= len(metadata.discrete_factor_names) + 1:
+            mi[idx, :] = mutual_info_regression(
                 data,
-                tgt,
-                discrete_features=is_categorical,  # type: ignore
+                data[:, idx],
+                discrete_features=is_discrete,  # type: ignore
                 n_neighbors=num_neighbors,
                 random_state=0,
             )
         else:
-            mi[idx, :] = mutual_info_regression(
+            mi[idx, :] = mutual_info_classif(
                 data,
-                tgt,
-                discrete_features=is_categorical,  # type: ignore
+                data[:, idx],
+                discrete_features=is_discrete,  # type: ignore
                 n_neighbors=num_neighbors,
                 random_state=0,
             )
 
-    ent_all = entropy(data, names, is_categorical, normalized=False)
-    norm_factor = 0.5 * np.add.outer(ent_all, ent_all) + 1e-6
+    # Normalization via entropy
+    bin_cnts = get_counts(discretized_data)
+    ent_factor = sp.stats.entropy(bin_cnts, axis=0)
+    norm_factor = 0.5 * np.add.outer(ent_factor, ent_factor) + 1e-6
+
     # in principle MI should be symmetric, but it is not in practice.
     nmi = 0.5 * (mi + mi.T) / norm_factor
     balance = nmi[0]
     factors = nmi[1:, 1:]
 
-    # unique class labels
-    class_idx = names.index("class_label")
-    u_cls = np.unique(data[:, class_idx])
-    num_classes = len(u_cls)
-
     # assume class is a factor
-    classwise_mi = np.empty((num_classes, num_factors))
-    classwise_mi[:] = np.nan
-
-    # categorical variables, excluding class label
-    cat_mask = np.concatenate((is_categorical[:class_idx], is_categorical[(class_idx + 1) :]), axis=0).astype(int)
+    num_classes = metadata.class_names.size
+    classwise_mi = np.full((num_classes, num_factors), np.nan, dtype=np.float32)
 
-    tgt_bin = np.stack([data[:, class_idx] == cls for cls in u_cls]).T.astype(int)
-    ent_tgt_bin = entropy(
-        tgt_bin, names=[str(idx) for idx in range(num_classes)], is_categorical=[True for idx in range(num_classes)]
-    )
+    # classwise targets
+    classes = np.unique(metadata.class_labels)
+    tgt_bin = data[:, 0][:, None] == classes
 
     # classification MI for discrete/categorical features
     for idx in range(num_classes):
-        # tgt = class_data == cls
         # units: nat
         classwise_mi[idx, :] = mutual_info_classif(
             data,
             tgt_bin[:, idx],
-            discrete_features=cat_mask,  # type: ignore
+            discrete_features=is_discrete,  # type: ignore
             n_neighbors=num_neighbors,
             random_state=0,
         )
 
-    norm_factor = 0.5 * np.add.outer(ent_tgt_bin, ent_all) + 1e-6
+    # Classwise normalization via entropy
+    classwise_bin_cnts = get_counts(tgt_bin)
+    ent_tgt_bin = sp.stats.entropy(classwise_bin_cnts, axis=0)
+    norm_factor = 0.5 * np.add.outer(ent_tgt_bin, ent_factor) + 1e-6
     classwise = classwise_mi / norm_factor
 
-    return BalanceOutput(balance, factors, classwise, unique_labels, list(metadata.keys()))
+    # Grabbing factor names for plotting function
+    factor_names = ["class"]
+    for name in metadata.discrete_factor_names:
+        if name in metadata.continuous_factor_names:
+            name = name + "-discrete"
+        factor_names.append(name)
+    for name in metadata.continuous_factor_names:
+        factor_names.append(name + "-continuous")
+
+    return BalanceOutput(balance, factors, classwise, factor_names, metadata.class_names)

dataeval/metrics/bias/coverage.py

@@ -5,14 +5,14 @@ __all__ = ["CoverageOutput", "coverage"]
 import contextlib
 import math
 from dataclasses import dataclass
-from typing import Any, Literal
+from typing import Literal
 
 import numpy as np
 from numpy.typing import ArrayLike, NDArray
 from scipy.spatial.distance import pdist, squareform
 
 from dataeval.interop import to_numpy
-from dataeval.metrics.bias.metadata import coverage_plot
+from dataeval.metrics.bias.metadata_utils import coverage_plot
 from dataeval.output import OutputMetadata, set_metadata
 from dataeval.utils.shared import flatten
 
@@ -27,9 +27,9 @@ class CoverageOutput(OutputMetadata):
 
     Attributes
     ----------
-    indices : NDArray
+    indices : NDArray[np.intp]
         Array of uncovered indices
-    radii : NDArray
+    radii : NDArray[np.float64]
         Array of critical value radii
     critical_value : float
         Radius for :term:`coverage<Coverage>`
@@ -39,11 +39,7 @@ class CoverageOutput(OutputMetadata):
     radii: NDArray[np.float64]
     critical_value: float
 
-    def plot(
-        self,
-        images: NDArray[Any],
-        top_k: int = 6,
-    ) -> Figure:
+    def plot(self, images: ArrayLike, top_k: int = 6) -> Figure:
         """
         Plot the top k images together for visualization
 
@@ -53,6 +49,10 @@ class CoverageOutput(OutputMetadata):
             Original images (not embeddings) in (N, C, H, W) or (N, H, W) format
         top_k : int, default 6
             Number of images to plot (plotting assumes groups of 3)
+
+        Returns
+        -------
+        matplotlib.figure.Figure
         """
         # Determine which images to plot
         highest_uncovered_indices = self.indices[:top_k]
@@ -82,12 +82,12 @@ def coverage(
     embeddings : ArrayLike, shape - (N, P)
         A dataset in an ArrayLike format.
         Function expects the data to have 2 dimensions, N number of observations in a P-dimesionial space.
-    radius_type : Literal["adaptive", "naive"], default "adaptive"
+    radius_type : {"adaptive", "naive"}, default "adaptive"
         The function used to determine radius.
-    k: int, default 20
+    k : int, default 20
         Number of observations required in order to be covered.
         [1] suggests that a minimum of 20-50 samples is necessary.
-    percent: float, default 0.01
+    percent : float, default 0.01
         Percent of observations to be considered uncovered. Only applies to adaptive radius.
 
     Returns