dataeval 0.72.1__py3-none-any.whl → 0.72.2__py3-none-any.whl

dataeval/{_internal/metrics → metrics/bias}/balance.py

@@ -1,36 +1,99 @@
 from __future__ import annotations
 
+__all__ = ["BalanceOutput", "balance"]
+
 import warnings
 from dataclasses import dataclass
-from typing import Mapping
+from typing import Any, Mapping
 
 import numpy as np
 from numpy.typing import ArrayLike, NDArray
 from sklearn.feature_selection import mutual_info_classif, mutual_info_regression
 
-from dataeval._internal.metrics.utils import entropy, preprocess_metadata
-from dataeval._internal.output import OutputMetadata, set_metadata
+from dataeval.metrics.bias.metadata import entropy, heatmap, preprocess_metadata
+from dataeval.output import OutputMetadata, set_metadata
 
 
 @dataclass(frozen=True)
 class BalanceOutput(OutputMetadata):
     """
-    Output class for :func:`balance` :term:`Bias` metric
+    Output class for :func:`balance` bias metric
 
     Attributes
     ----------
     balance : NDArray[np.float64]
-        Estimate of :term:`mutual information<Mutual Information (MI)>` between metadata factors and class label
+        Estimate of mutual information between metadata factors and class label
     factors : NDArray[np.float64]
         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[np.int64]
+        Class labels for each value in the dataset
+    metadata_names: list[str]
+        Names of each metadata factor
     """
 
     balance: NDArray[np.float64]
     factors: NDArray[np.float64]
     classwise: NDArray[np.float64]
 
+    class_list: NDArray[np.int64]
+    metadata_names: list[str]
+
+    def plot(
+        self,
+        row_labels: NDArray[Any] | None = None,
+        col_labels: NDArray[Any] | None = None,
+        plot_classwise: bool = False,
+    ) -> None:
+        """
+        Plot a heatmap of balance information
+
+        Parameters
+        ----------
+        row_labels: NDArray | None, default None
+            Array containing the labels for rows in the histogram
+        col_labels: NDArray | None, default None
+            Array containing the labels for columns in the histogram
+        plot_classwise: bool, default False
+            Whether to plot per-class balance instead of global balance
+
+        """
+        if plot_classwise:
+            if row_labels is None:
+                row_labels = np.unique(self.class_list)
+            if col_labels is None:
+                col_labels = np.concatenate((["class"], self.metadata_names))
+
+            heatmap(
+                self.classwise,
+                row_labels,
+                col_labels,
+                xlabel="Factors",
+                ylabel="Class",
+                cbarlabel="Normalized Mutual Information",
+            )
+        else:
+            data = np.concatenate([self.balance[np.newaxis, 1:], self.factors], axis=0)
+            # Create a mask for the upper triangle of the symmetrical array, ignoring the diagonal
+            mask = np.triu(data + 1, k=0) < 1
+            # 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))
+
+            if row_labels is None:
+                row_labels = heat_labels[:-1]
+            if col_labels is None:
+                col_labels = heat_labels[1:]
+
+            heatmap(
+                heat_data,
+                row_labels,
+                col_labels,
+                cbarlabel="Normalized Mutual Information",
+            )
+
 
 def validate_num_neighbors(num_neighbors: int) -> int:
     if not isinstance(num_neighbors, (int, float)):
@@ -55,7 +118,7 @@ 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:
     """
-    :term:`Mutual information (MI)` between factors (class label, metadata, label/image properties)
+    Mutual information (MI) between factors (class label, metadata, label/image properties)
 
     Parameters
     ----------
@@ -70,7 +133,7 @@ def balance(class_labels: ArrayLike, metadata: Mapping[str, ArrayLike], num_neig
     Returns
     -------
     BalanceOutput
-        (num_factors+1) x (num_factors+1) estimate of :term:`mutual information<Mutual Information (MI)>`
+        (num_factors+1) x (num_factors+1) estimate of mutual information
         between num_factors metadata factors and class label. Symmetry is enforced.
 
     Note
@@ -83,7 +146,7 @@ def balance(class_labels: ArrayLike, metadata: Mapping[str, ArrayLike], num_neig
 
     Example
     -------
-    Return :term:`balance<Balance>` (:term:`mutual information<Mutual Information (MI)>`) of factors with class_labels
+    Return balance (mutual information) of factors with class_labels
 
     >>> bal = balance(class_labels, metadata)
     >>> bal.balance
@@ -114,6 +177,9 @@ def balance(class_labels: ArrayLike, metadata: Mapping[str, ArrayLike], num_neig
     mi = np.empty((num_factors, num_factors))
     mi[:] = np.nan
 
+    class_idx = names.index("class_label")
+    class_lbl = np.array(data[:, class_idx], dtype=int)
+
     for idx in range(num_factors):
         tgt = data[:, idx].astype(int)
 
@@ -174,4 +240,4 @@ def balance(class_labels: ArrayLike, metadata: Mapping[str, ArrayLike], num_neig
     norm_factor = 0.5 * np.add.outer(ent_tgt_bin, ent_all) + 1e-6
     classwise = classwise_mi / norm_factor
 
-    return BalanceOutput(balance, factors, classwise)
+    return BalanceOutput(balance, factors, classwise, class_lbl, list(metadata.keys()))

dataeval/{_internal/metrics → metrics/bias}/coverage.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+__all__ = ["CoverageOutput", "coverage"]
+
 import math
 from dataclasses import dataclass
 from typing import Literal
@@ -8,9 +10,9 @@ import numpy as np
 from numpy.typing import ArrayLike, NDArray
 from scipy.spatial.distance import pdist, squareform
 
-from dataeval._internal.interop import to_numpy
-from dataeval._internal.metrics.utils import flatten
-from dataeval._internal.output import OutputMetadata, set_metadata
+from dataeval.interop import to_numpy
+from dataeval.output import OutputMetadata, set_metadata
+from dataeval.utils.shared import flatten
 
 
 @dataclass(frozen=True)
@@ -33,7 +35,7 @@ class CoverageOutput(OutputMetadata):
     critical_value: float
 
 
-@set_metadata("dataeval.metrics")
+@set_metadata()
 def coverage(
     embeddings: ArrayLike,
     radius_type: Literal["adaptive", "naive"] = "adaptive",

dataeval/{_internal/metrics → metrics/bias}/diversity.py

@@ -1,13 +1,16 @@
 from __future__ import annotations
 
+__all__ = ["DiversityOutput", "diversity"]
+
 from dataclasses import dataclass
-from typing import Literal, Mapping
+from typing import Any, Literal, Mapping
 
 import numpy as np
 from numpy.typing import ArrayLike, NDArray
 
-from dataeval._internal.metrics.utils import entropy, get_counts, get_method, get_num_bins, preprocess_metadata
-from dataeval._internal.output import OutputMetadata, set_metadata
+from dataeval.metrics.bias.metadata import entropy, get_counts, get_num_bins, heatmap, preprocess_metadata
+from dataeval.output import OutputMetadata, set_metadata
+from dataeval.utils.shared import get_method
 
 
 @dataclass(frozen=True)
@@ -21,18 +24,52 @@ class DiversityOutput(OutputMetadata):
         :term:`Diversity` index for classes and factors
     classwise : NDArray[np.float64]
         Classwise diversity index [n_class x n_factor]
+    class_list: NDArray[np.int64]
+        Class labels for each value in the dataset
+    metadata_names: list[str]
+        Names of each metadata factor
     """
 
     diversity_index: NDArray[np.float64]
     classwise: NDArray[np.float64]
 
+    class_list: NDArray[np.int64]
+    metadata_names: list[str]
+
+    method: Literal["shannon", "simpson"]
+
+    def plot(self, row_labels: NDArray[Any] | None = None, col_labels: NDArray[Any] | None = None) -> None:
+        """
+        Plot a heatmap of diversity information
+
+        Parameters
+        ----------
+        row_labels: NDArray | None, default None
+            Array containing the labels for rows in the histogram
+        col_labels: NDArray | None, default None
+            Array containing the labels for columns in the histogram
+        """
+        if row_labels is None:
+            row_labels = np.unique(self.class_list)
+        if col_labels is None:
+            col_labels = np.array(self.metadata_names)
+
+        heatmap(
+            self.classwise,
+            row_labels,
+            col_labels,
+            xlabel="Factors",
+            ylabel="Class",
+            cbarlabel=f"Normalized {self.method.title()} Index",
+        )
+
 
 def diversity_shannon(
-    data: NDArray,
+    data: NDArray[Any],
     names: list[str],
     is_categorical: list[bool],
     subset_mask: NDArray[np.bool_] | None = None,
-) -> NDArray:
+) -> NDArray[np.float64]:
     """
     Compute :term:`diversity<Diversity>` for discrete/categorical variables and, through standard
     histogram binning, for continuous variables.
@@ -79,11 +116,11 @@ def diversity_shannon(
 
 
 def diversity_simpson(
-    data: NDArray,
+    data: NDArray[Any],
     names: list[str],
     is_categorical: list[bool],
     subset_mask: NDArray[np.bool_] | None = None,
-) -> NDArray:
+) -> NDArray[np.float64]:
     """
     Compute :term:`diversity<Diversity>` for discrete/categorical variables and, through standard
     histogram binning, for continuous variables.
@@ -139,10 +176,7 @@ def diversity_simpson(
     return ev_index
 
 
-DIVERSITY_FN_MAP = {"simpson": diversity_simpson, "shannon": diversity_shannon}
-
-
-@set_metadata("dataeval.metrics")
+@set_metadata()
 def diversity(
     class_labels: ArrayLike, metadata: Mapping[str, ArrayLike], method: Literal["shannon", "simpson"] = "simpson"
 ) -> DiversityOutput:
@@ -202,12 +236,12 @@ def diversity(
     --------
     numpy.histogram
     """
-    diversity_fn = get_method(DIVERSITY_FN_MAP, method)
+    diversity_fn = get_method({"simpson": diversity_simpson, "shannon": diversity_shannon}, method)
     data, names, is_categorical = preprocess_metadata(class_labels, metadata)
     diversity_index = diversity_fn(data, names, is_categorical, None).astype(np.float64)
 
     class_idx = names.index("class_label")
-    class_lbl = data[:, class_idx]
+    class_lbl = np.array(data[:, class_idx], dtype=int)
 
     u_classes = np.unique(class_lbl)
     num_factors = len(names)
@@ -218,4 +252,4 @@ def diversity(
         diversity[idx, :] = diversity_fn(data, names, is_categorical, subset_mask)
     div_no_class = np.concatenate((diversity[:, :class_idx], diversity[:, (class_idx + 1) :]), axis=1)
 
-    return DiversityOutput(diversity_index, div_no_class)
+    return DiversityOutput(diversity_index, div_no_class, class_lbl, list(metadata.keys()), method)

dataeval/metrics/bias/metadata.py

@@ -0,0 +1,275 @@
+from __future__ import annotations
+
+__all__ = []
+
+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
+
+
+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]]:
+    # convert class_labels and dict of lists to matrix of metadata values
+    preprocessed_metadata = {"class_label": np.asarray(class_labels, dtype=int)}
+
+    # 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 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
+
+
+def heatmap(
+    data: NDArray[Any],
+    row_labels: NDArray[Any],
+    col_labels: NDArray[Any],
+    xlabel: str = "",
+    ylabel: str = "",
+    cbarlabel: str = "",
+) -> None:
+    """
+    Plots a formatted heatmap
+
+    Parameters
+    ----------
+    data: NDArray
+        Array containing numerical values for factors to plot
+    row_labels: NDArray
+        Array containing the labels for rows in the histogram
+    col_labels: NDArray
+        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()
+    plt.show()
+
+
+# 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", "")

dataeval/{_internal/metrics → metrics/bias}/parity.py

@@ -1,15 +1,17 @@
 from __future__ import annotations
 
+__all__ = ["ParityOutput", "parity", "label_parity"]
+
 import warnings
 from dataclasses import dataclass
-from typing import Generic, Mapping, TypeVar
+from typing import Any, Generic, Mapping, TypeVar
 
 import numpy as np
 from numpy.typing import ArrayLike, NDArray
 from scipy.stats import chi2_contingency, chisquare
 
-from dataeval._internal.interop import to_numpy
-from dataeval._internal.output import OutputMetadata, set_metadata
+from dataeval.interop import to_numpy
+from dataeval.output import OutputMetadata, set_metadata
 
 TData = TypeVar("TData", np.float64, NDArray[np.float64])
 
@@ -31,7 +33,7 @@ class ParityOutput(Generic[TData], OutputMetadata):
     p_value: TData
 
 
-def digitize_factor_bins(continuous_values: NDArray, bins: int, factor_name: str) -> NDArray:
+def digitize_factor_bins(continuous_values: NDArray[Any], bins: int, factor_name: str) -> NDArray[np.intp]:
     """
     Digitizes a list of values into a given number of bins.
 
@@ -64,8 +66,8 @@ def digitize_factor_bins(continuous_values: NDArray, bins: int, factor_name: str
 
 
 def format_discretize_factors(
-    data_factors: Mapping[str, NDArray], continuous_factor_bincounts: Mapping[str, int]
-) -> dict[str, NDArray]:
+    data_factors: Mapping[str, NDArray[Any]], continuous_factor_bincounts: Mapping[str, int]
+) -> dict[str, NDArray[Any]]:
     """
     Sets up the internal list of metadata factors.
 
@@ -115,7 +117,7 @@ def format_discretize_factors(
     return metadata_factors
 
 
-def normalize_expected_dist(expected_dist: NDArray, observed_dist: NDArray) -> NDArray:
+def normalize_expected_dist(expected_dist: NDArray[Any], observed_dist: NDArray[Any]) -> NDArray[Any]:
     """
     Normalize the expected label distribution to match the total number of labels in the observed distribution.
 
@@ -162,7 +164,7 @@ def normalize_expected_dist(expected_dist: NDArray, observed_dist: NDArray) -> N
     return expected_dist
 
 
-def validate_dist(label_dist: NDArray, label_name: str):
+def validate_dist(label_dist: NDArray[Any], label_name: str) -> None:
     """
     Verifies that the given label distribution has labels and checks if
     any labels have frequencies less than 5.
@@ -191,7 +193,7 @@ def validate_dist(label_dist: NDArray, label_name: str):
         )
 
 
-@set_metadata("dataeval.metrics")
+@set_metadata()
 def label_parity(
     expected_labels: ArrayLike,
     observed_labels: ArrayLike,
@@ -279,7 +281,7 @@ def label_parity(
     return ParityOutput(cs, p)
 
 
-@set_metadata("dataeval.metrics")
+@set_metadata()
 def parity(
     class_labels: ArrayLike,
     data_factors: Mapping[str, ArrayLike],

dataeval/metrics/estimators/__init__.py

@@ -2,8 +2,8 @@
 Estimators calculate performance bounds and the statistical distance between datasets.
 """
 
-from dataeval._internal.metrics.ber import BEROutput, ber
-from dataeval._internal.metrics.divergence import DivergenceOutput, divergence
-from dataeval._internal.metrics.uap import UAPOutput, uap
+from dataeval.metrics.estimators.ber import BEROutput, ber
+from dataeval.metrics.estimators.divergence import DivergenceOutput, divergence
+from dataeval.metrics.estimators.uap import UAPOutput, uap
 
 __all__ = ["ber", "divergence", "uap", "BEROutput", "DivergenceOutput", "UAPOutput"]