dataeval 0.76.1__py3-none-any.whl → 0.81.0__py3-none-any.whl

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

@@ -5,16 +5,17 @@ __all__ = []
 import contextlib
 import warnings
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, Literal, overload
 
 import numpy as np
 import scipy as sp
 from numpy.typing import NDArray
 from sklearn.feature_selection import mutual_info_classif, mutual_info_regression
 
-from dataeval.output import Output, set_metadata
-from dataeval.utils.metadata import Metadata, get_counts
-from dataeval.utils.plot import heatmap
+from dataeval._output import Output, set_metadata
+from dataeval.utils._bin import get_counts
+from dataeval.utils._plot import heatmap
+from dataeval.utils.data import Metadata
 
 with contextlib.suppress(ImportError):
     from matplotlib.figure import Figure
@@ -23,8 +24,8 @@ with contextlib.suppress(ImportError):
 @dataclass(frozen=True)
 class BalanceOutput(Output):
     """
-    Output class for :func:`balance` :term:`bias<Bias>` metric.
-    
+    Output class for :func:`.balance` :term:`bias<Bias>` metric.
+
     Attributes
     ----------
     balance : NDArray[np.float64]
@@ -35,21 +36,62 @@ class BalanceOutput(Output):
         Estimate of mutual information between metadata factors and individual class labels
     factor_names : list[str]
         Names of each metadata factor
-    class_list : NDArray
-        Array of the class labels present in the dataset
+    class_names : list[str]
+        List 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]
+    class_names: list[str]
+
+    @overload
+    def _by_factor_type(
+        self,
+        attr: Literal["factor_names"],
+        factor_type: Literal["discrete", "continuous", "both"],
+    ) -> list[str]: ...
+
+    @overload
+    def _by_factor_type(
+        self,
+        attr: Literal["balance", "factors", "classwise"],
+        factor_type: Literal["discrete", "continuous", "both"],
+    ) -> NDArray[np.float64]: ...
+
+    def _by_factor_type(
+        self,
+        attr: Literal["balance", "factors", "classwise", "factor_names"],
+        factor_type: Literal["discrete", "continuous", "both"],
+    ) -> NDArray[np.float64] | list[str]:
+        # if not filtering by factor_type then just return the requested attribute without mask
+        if factor_type == "both":
+            return getattr(self, attr)
+
+        # create the mask for the selected factor_type
+        mask_lambda = (
+            (lambda x: "-continuous" not in x) if factor_type == "discrete" else (lambda x: "-discrete" not in x)
+        )
+
+        # return the masked attribute
+        if attr == "factor_names":
+            return [x.replace(f"-{factor_type}", "") for x in self.factor_names if mask_lambda(x)]
+        else:
+            factor_type_mask = [mask_lambda(x) for x in self.factor_names]
+            if attr == "factors":
+                return self.factors[factor_type_mask[1:]][:, factor_type_mask[1:]]
+            elif attr == "balance":
+                return self.balance[factor_type_mask]
+            elif attr == "classwise":
+                return self.classwise[:, factor_type_mask]
 
     def plot(
         self,
         row_labels: list[Any] | NDArray[Any] | None = None,
         col_labels: list[Any] | NDArray[Any] | None = None,
         plot_classwise: bool = False,
+        factor_type: Literal["discrete", "continuous", "both"] = "discrete",
     ) -> Figure:
         """
         Plot a heatmap of balance information
@@ -62,15 +104,17 @@ class BalanceOutput(Output):
             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
+        factor_type : "discrete", "continuous", or "both", default "discrete"
+            Whether to plot discretized values, continuous values, or to include both
         """
         if plot_classwise:
             if row_labels is None:
-                row_labels = self.class_list
+                row_labels = self.class_names
             if col_labels is None:
-                col_labels = self.factor_names
+                col_labels = self._by_factor_type("factor_names", factor_type)
 
             fig = heatmap(
-                self.classwise,
+                self._by_factor_type("classwise", factor_type),
                 row_labels,
                 col_labels,
                 xlabel="Factors",
@@ -79,13 +123,19 @@ class BalanceOutput(Output):
             )
         else:
             # Combine balance and factors results
-            data = np.concatenate([self.balance[np.newaxis, 1:], self.factors], axis=0)
+            data = np.concatenate(
+                [
+                    self._by_factor_type("balance", factor_type)[np.newaxis, 1:],
+                    self._by_factor_type("factors", factor_type),
+                ],
+                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 = self.factor_names
+            heat_labels = self._by_factor_type("factor_names", factor_type)
 
             if row_labels is None:
                 row_labels = heat_labels[:-1]
@@ -128,7 +178,7 @@ def balance(
     Parameters
     ----------
     metadata : Metadata
-        Preprocessed metadata from :func:`dataeval.utils.metadata.preprocess`
+        Preprocessed metadata
     num_neighbors : int, default 5
         Number of points to consider as neighbors
 
@@ -184,7 +234,7 @@ def balance(
     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:
+    if len(metadata.continuous_data):
         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]))
@@ -218,7 +268,7 @@ def balance(
     factors = nmi[1:, 1:]
 
     # assume class is a factor
-    num_classes = metadata.class_names.size
+    num_classes = len(metadata.class_names)
     classwise_mi = np.full((num_classes, num_factors), np.nan, dtype=np.float32)
 
     # classwise targets

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

@@ -8,12 +8,12 @@ from dataclasses import dataclass
 from typing import Any, Literal
 
 import numpy as np
-from numpy.typing import ArrayLike, NDArray
+from numpy.typing import NDArray
 from scipy.spatial.distance import pdist, squareform
 
-from dataeval.interop import to_numpy
-from dataeval.output import Output, set_metadata
-from dataeval.utils.shared import flatten
+from dataeval._output import Output, set_metadata
+from dataeval.typing import ArrayLike
+from dataeval.utils._array import ensure_embeddings, flatten, to_numpy
 
 with contextlib.suppress(ImportError):
     from matplotlib.figure import Figure
@@ -71,21 +71,21 @@ def _plot(images: NDArray[Any], num_images: int) -> Figure:
 @dataclass(frozen=True)
 class CoverageOutput(Output):
     """
-    Output class for :func:`coverage` :term:`bias<Bias>` metric.
+    Output class for :func:`.coverage` :term:`bias<Bias>` metric.
 
     Attributes
     ----------
-    indices : NDArray[np.intp]
+    uncovered_indices : NDArray[np.intp]
         Array of uncovered indices
-    radii : NDArray[np.float64]
+    critical_value_radii : NDArray[np.float64]
         Array of critical value radii
-    critical_value : float
+    coverage_radius : float
         Radius for :term:`coverage<Coverage>`
     """
 
-    indices: NDArray[np.intp]
-    radii: NDArray[np.float64]
-    critical_value: float
+    uncovered_indices: NDArray[np.intp]
+    critical_value_radii: NDArray[np.float64]
+    coverage_radius: float
 
     def plot(self, images: ArrayLike, top_k: int = 6) -> Figure:
         """
@@ -102,8 +102,9 @@ class CoverageOutput(Output):
         -------
         matplotlib.figure.Figure
         """
+
         # Determine which images to plot
-        highest_uncovered_indices = self.indices[:top_k]
+        highest_uncovered_indices = self.uncovered_indices[:top_k]
 
         # Grab the images
         images = to_numpy(images)
@@ -119,7 +120,7 @@ class CoverageOutput(Output):
 def coverage(
     embeddings: ArrayLike,
     radius_type: Literal["adaptive", "naive"] = "adaptive",
-    k: int = 20,
+    num_observations: int = 20,
     percent: float = 0.01,
 ) -> CoverageOutput:
     """
@@ -128,11 +129,11 @@ def coverage(
     Parameters
     ----------
     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.
+        Dataset embeddings as unit interval [0, 1].
+        Function expects the data to have 2 dimensions, N number of observations in a P-dimensional space.
     radius_type : {"adaptive", "naive"}, default "adaptive"
         The function used to determine radius.
-    k : int, default 20
+    num_observations : 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
@@ -146,7 +147,9 @@ def coverage(
     Raises
     ------
     ValueError
-        If length of :term:`embeddings<Embeddings>` is less than or equal to k
+        If embeddings are not unit interval [0-1]
+    ValueError
+        If length of :term:`embeddings<Embeddings>` is less than or equal to num_observations
     ValueError
         If radius_type is unknown
 
@@ -157,10 +160,10 @@ def coverage(
     Example
     -------
     >>> results = coverage(embeddings)
-    >>> results.indices
+    >>> results.uncovered_indices
     array([447, 412,   8,  32,  63])
-    >>> results.critical_value
-    0.8459038956941765
+    >>> results.coverage_radius
+    0.17592147193757596
 
     Reference
     ---------
@@ -169,26 +172,29 @@ def coverage(
     [1] Seymour Sudman. 1976. Applied sampling. Academic Press New York (1976).
     """
 
-    # Calculate distance matrix, look at the (k+1)th farthest neighbor for each image.
-    embeddings = to_numpy(embeddings)
-    n = len(embeddings)
-    if n <= k:
+    # Calculate distance matrix, look at the (num_observations + 1)th farthest neighbor for each image.
+    embeddings = ensure_embeddings(embeddings, dtype=np.float64, unit_interval=True)
+    len_embeddings = len(embeddings)
+    if len_embeddings <= num_observations:
         raise ValueError(
-            f"Number of observations n={n} is less than or equal to the specified number of neighbors k={k}."
+            f"Length of embeddings ({len_embeddings}) is less than or equal to the specified number of \
+                observations ({num_observations})."
         )
-    mat = squareform(pdist(flatten(embeddings))).astype(np.float64)
-    sorted_dists = np.sort(mat, axis=1)
-    crit = sorted_dists[:, k + 1]
+    embeddings_matrix = squareform(pdist(flatten(embeddings))).astype(np.float64)
+    sorted_dists = np.sort(embeddings_matrix, axis=1)
+    critical_value_radii = sorted_dists[:, num_observations + 1]
 
     d = embeddings.shape[1]
     if radius_type == "naive":
-        rho = (1 / math.sqrt(math.pi)) * ((2 * k * math.gamma(d / 2 + 1)) / (n)) ** (1 / d)
-        pvals = np.where(crit > rho)[0]
+        coverage_radius = (1 / math.sqrt(math.pi)) * (
+            (2 * num_observations * math.gamma(d / 2 + 1)) / (len_embeddings)
+        ) ** (1 / d)
+        uncovered_indices = np.where(critical_value_radii > coverage_radius)[0]
     elif radius_type == "adaptive":
-        # Use data adaptive cutoff as rho
-        selection = int(max(n * percent, 1))
-        pvals = np.argsort(crit)[::-1][:selection]
-        rho = float(np.mean(np.sort(crit)[::-1][selection - 1 : selection + 1]))
+        # Use data adaptive cutoff as coverage_radius
+        selection = int(max(len_embeddings * percent, 1))
+        uncovered_indices = np.argsort(critical_value_radii)[::-1][:selection]
+        coverage_radius = float(np.mean(np.sort(critical_value_radii)[::-1][selection - 1 : selection + 1]))
     else:
         raise ValueError(f"{radius_type} is an invalid radius type. Expected 'adaptive' or 'naive'")
-    return CoverageOutput(pvals, crit, rho)
+    return CoverageOutput(uncovered_indices, critical_value_radii, coverage_radius)

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

@@ -8,12 +8,14 @@ from typing import Any, Literal
 
 import numpy as np
 import scipy as sp
-from numpy.typing import ArrayLike, NDArray
+from numpy.typing import NDArray
 
-from dataeval.output import Output, set_metadata
-from dataeval.utils.metadata import Metadata, get_counts
-from dataeval.utils.plot import heatmap
-from dataeval.utils.shared import get_method
+from dataeval._output import Output, set_metadata
+from dataeval.typing import ArrayLike
+from dataeval.utils._bin import get_counts
+from dataeval.utils._method import get_method
+from dataeval.utils._plot import heatmap
+from dataeval.utils.data import Metadata
 
 with contextlib.suppress(ImportError):
     from matplotlib.figure import Figure
@@ -37,7 +39,7 @@ def _plot(labels: NDArray[Any], bar_heights: NDArray[Any]) -> Figure:
     """
     import matplotlib.pyplot as plt
 
-    fig, ax = plt.subplots(figsize=(10, 10))
+    fig, ax = plt.subplots(figsize=(8, 8))
 
     ax.bar(labels, bar_heights)
     ax.set_xlabel("Factors")
@@ -51,7 +53,7 @@ def _plot(labels: NDArray[Any], bar_heights: NDArray[Any]) -> Figure:
 @dataclass(frozen=True)
 class DiversityOutput(Output):
     """
-    Output class for :func:`diversity` :term:`bias<Bias>` metric.
+    Output class for :func:`.diversity` :term:`bias<Bias>` metric.
 
     Attributes
     ----------
@@ -61,14 +63,14 @@ class DiversityOutput(Output):
         Classwise diversity index [n_class x n_factor]
     factor_names : list[str]
         Names of each metadata factor
-    class_list : NDArray[Any]
+    class_names : list[str]
         Class labels for each value in the dataset
     """
 
     diversity_index: NDArray[np.double]
     classwise: NDArray[np.double]
     factor_names: list[str]
-    class_list: NDArray[Any]
+    class_names: list[str]
 
     def plot(
         self,
@@ -90,7 +92,7 @@ class DiversityOutput(Output):
         """
         if plot_classwise:
             if row_labels is None:
-                row_labels = self.class_list
+                row_labels = self.class_names
             if col_labels is None:
                 col_labels = self.factor_names
 
@@ -191,6 +193,9 @@ def diversity_simpson(
     return ev_index
 
 
+_DIVERSITY_FN_MAP = {"simpson": diversity_simpson, "shannon": diversity_shannon}
+
+
 @set_metadata
 def diversity(
     metadata: Metadata,
@@ -210,7 +215,7 @@ def diversity(
     Parameters
     ----------
     metadata : Metadata
-        Preprocessed metadata from :func:`dataeval.utils.metadata.preprocess`
+        Preprocessed metadata
     method : "simpson" or "shannon", default "simpson"
         The methodology used for defining diversity
 
@@ -251,7 +256,7 @@ def diversity(
     --------
     scipy.stats.entropy
     """
-    diversity_fn = get_method({"simpson": diversity_simpson, "shannon": diversity_shannon}, method)
+    diversity_fn = get_method(_DIVERSITY_FN_MAP, method)
     discretized_data = np.hstack((metadata.class_labels[:, np.newaxis], metadata.discrete_data))
     cnts = get_counts(discretized_data)
     num_bins = np.bincount(np.nonzero(cnts)[1])

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

@@ -2,40 +2,86 @@ from __future__ import annotations
 
 __all__ = []
 
+import contextlib
 import warnings
 from dataclasses import dataclass
 from typing import Any, Generic, TypeVar
 
 import numpy as np
-from numpy.typing import ArrayLike, NDArray
+from numpy.typing import NDArray
 from scipy.stats import chisquare
 from scipy.stats.contingency import chi2_contingency, crosstab
 
-from dataeval.interop import as_numpy, to_numpy
-from dataeval.output import Output, set_metadata
-from dataeval.utils.metadata import Metadata
+from dataeval._output import Output, set_metadata
+from dataeval.typing import ArrayLike
+from dataeval.utils._array import as_numpy
+from dataeval.utils.data import Metadata
+
+with contextlib.suppress(ImportError):
+    import pandas as pd
 
 TData = TypeVar("TData", np.float64, NDArray[np.float64])
 
 
 @dataclass(frozen=True)
-class ParityOutput(Generic[TData], Output):
+class BaseParityOutput(Generic[TData], Output):
+    score: TData
+    p_value: TData
+
+    def to_dataframe(self) -> pd.DataFrame:
+        """
+        Exports the parity output results to a pandas DataFrame.
+
+        Returns
+        -------
+        pd.DataFrame
+        """
+        import pandas as pd
+
+        return pd.DataFrame(
+            index=self.factor_names,  # type: ignore - list[str] is documented as acceptable index type
+            data={
+                "score": self.score.round(2),
+                "p-value": self.p_value.round(2),
+            },
+        )
+
+
+@dataclass(frozen=True)
+class LabelParityOutput(BaseParityOutput[np.float64]):
+    """
+    Output class for :func:`.label_parity` :term:`bias<Bias>` metrics.
+
+    Attributes
+    ----------
+    score : np.float64
+        chi-squared score(s) of the test
+    p_value : np.float64
+        p-value(s) of the test
+    """
+
+
+@dataclass(frozen=True)
+class ParityOutput(BaseParityOutput[NDArray[np.float64]]):
     """
-    Output class for :func:`parity` and :func:`label_parity` :term:`bias<Bias>` metrics.
+    Output class for :func:`.parity` :term:`bias<Bias>` metrics.
 
     Attributes
     ----------
-    score : np.float64 | NDArray[np.float64]
+    score : NDArray[np.float64]
         chi-squared score(s) of the test
-    p_value : np.float64 | NDArray[np.float64]
+    p_value : NDArray[np.float64]
         p-value(s) of the test
-    metadata_names : list[str] | None
+    factor_names : list[str]
         Names of each metadata factor
+    insufficient_data: dict
+        Dictionary of metadata factors with less than 5 class occurrences per value
     """
 
-    score: TData
-    p_value: TData
-    metadata_names: list[str] | None
+    # score: NDArray[np.float64]
+    # p_value: NDArray[np.float64]
+    factor_names: list[str]
+    insufficient_data: dict[str, dict[int, dict[str, int]]]
 
 
 def normalize_expected_dist(expected_dist: NDArray[Any], observed_dist: NDArray[Any]) -> NDArray[Any]:
@@ -109,7 +155,7 @@ def validate_dist(label_dist: NDArray[Any], label_name: str) -> None:
         raise ValueError(f"No labels found in the {label_name} dataset")
     if np.any(label_dist < 5):
         warnings.warn(
-            f"Labels {np.where(label_dist<5)[0]} in {label_name}"
+            f"Labels {np.where(label_dist < 5)[0]} in {label_name}"
             " dataset have frequencies less than 5. This may lead"
             " to invalid chi-squared evaluation.",
             UserWarning,
@@ -121,7 +167,7 @@ def label_parity(
     expected_labels: ArrayLike,
     observed_labels: ArrayLike,
     num_classes: int | None = None,
-) -> ParityOutput[np.float64]:
+) -> LabelParityOutput:
     """
     Calculate the chi-square statistic to assess the :term:`parity<Parity>` \
     between expected and observed label distributions.
@@ -142,7 +188,7 @@ def label_parity(
 
     Returns
     -------
-    ParityOutput[np.float64]
+    LabelParityOutput
         chi-squared score and :term`P-Value` of the test
 
     Raises
@@ -171,7 +217,7 @@ def label_parity(
     >>> expected_labels = rng.choice([0, 1, 2, 3, 4], (100))
     >>> observed_labels = rng.choice([2, 3, 0, 4, 1], (100))
     >>> label_parity(expected_labels, observed_labels)
-    ParityOutput(score=14.007374204742625, p_value=0.0072715574616218, metadata_names=None)
+    LabelParityOutput(score=14.007374204742625, p_value=0.0072715574616218)
     """
 
     # Calculate
@@ -179,8 +225,8 @@ def label_parity(
         num_classes = 0
 
     # Calculate the class frequencies associated with the datasets
-    observed_dist = np.bincount(to_numpy(observed_labels), minlength=num_classes)
-    expected_dist = np.bincount(to_numpy(expected_labels), minlength=num_classes)
+    observed_dist = np.bincount(as_numpy(observed_labels), minlength=num_classes)
+    expected_dist = np.bincount(as_numpy(expected_labels), minlength=num_classes)
 
     # Validate
     validate_dist(observed_dist, "observed")
@@ -202,11 +248,11 @@ def label_parity(
         )
 
     cs, p = chisquare(f_obs=observed_dist, f_exp=expected_dist)
-    return ParityOutput(cs, p, None)
+    return LabelParityOutput(cs, p)
 
 
 @set_metadata
-def parity(metadata: Metadata) -> ParityOutput[NDArray[np.float64]]:
+def parity(metadata: Metadata) -> ParityOutput:
     """
     Calculate chi-square statistics to assess the linear relationship \
     between multiple factors and class labels.
@@ -218,7 +264,7 @@ def parity(metadata: Metadata) -> ParityOutput[NDArray[np.float64]]:
     Parameters
     ----------
     metadata : Metadata
-        Preprocessed metadata from :func:`dataeval.utils.metadata.preprocess`
+        Preprocessed metadata
 
     Returns
     -------
@@ -250,22 +296,21 @@ def parity(metadata: Metadata) -> ParityOutput[NDArray[np.float64]]:
     --------
     Randomly creating some "continuous" and categorical variables using ``np.random.default_rng``
 
-    >>> from dataeval.utils.metadata import preprocess
-    >>> rng = np.random.default_rng(175)
-    >>> labels = rng.choice([0, 1, 2], (100))
-    >>> metadata_dict = {
-    ...         "age": list(rng.choice([25, 30, 35, 45], (100))),
-    ...         "income": list(rng.choice([50000, 65000, 80000], (100))),
-    ...         "gender": list(rng.choice(["M", "F"], (100))),
-    ... }
-    >>> continuous_factor_bincounts = {"age": 4, "income": 3}
-    >>> metadata = preprocess(metadata_dict, labels, continuous_factor_bincounts)
+    >>> metadata = generate_random_metadata(
+    ...     labels=["doctor", "artist", "teacher"],
+    ...     factors={
+    ...         "age": [25, 30, 35, 45],
+    ...         "income": [50000, 65000, 80000],
+    ...         "gender": ["M", "F"]},
+    ...     length=100,
+    ...     random_seed=175)
+    >>> metadata.continuous_factor_bins = {"age": 4, "income": 3}
     >>> parity(metadata)
-    ParityOutput(score=array([7.35731943, 5.46711299, 0.51506212]), p_value=array([0.28906231, 0.24263543, 0.77295762]), metadata_names=['age', 'income', 'gender'])
+    ParityOutput(score=array([7.35731943, 5.46711299, 0.51506212]), p_value=array([0.28906231, 0.24263543, 0.77295762]), factor_names=['age', 'income', 'gender'], insufficient_data={'age': {3: {'artist': 4}, 4: {'artist': 4, 'teacher': 3}}, 'income': {1: {'artist': 3}}})
     """  # noqa: E501
     chi_scores = np.zeros(metadata.discrete_data.shape[1])
     p_values = np.zeros_like(chi_scores)
-    not_enough_data = {}
+    insufficient_data = {}
     for i, col_data in enumerate(metadata.discrete_data.T):
         # Builds a contingency matrix where entry at index (r,c) represents
         # the frequency of current_factor_name achieving value unique_factor_values[r]
@@ -279,14 +324,14 @@ def parity(metadata: Metadata) -> ParityOutput[NDArray[np.float64]]:
         current_factor_name = metadata.discrete_factor_names[i]
         for int_factor, int_class in zip(counts[0], counts[1]):
             if contingency_matrix[int_factor, int_class] > 0:
-                factor_category = unique_factor_values[int_factor]
-                if current_factor_name not in not_enough_data:
-                    not_enough_data[current_factor_name] = {}
-                if factor_category not in not_enough_data[current_factor_name]:
-                    not_enough_data[current_factor_name][factor_category] = []
-                not_enough_data[current_factor_name][factor_category].append(
-                    (metadata.class_names[int_class], int(contingency_matrix[int_factor, int_class]))
-                )
+                factor_category = unique_factor_values[int_factor].item()
+                if current_factor_name not in insufficient_data:
+                    insufficient_data[current_factor_name] = {}
+                if factor_category not in insufficient_data[current_factor_name]:
+                    insufficient_data[current_factor_name][factor_category] = {}
+                class_name = metadata.class_names[int_class]
+                class_count = contingency_matrix[int_factor, int_class].item()
+                insufficient_data[current_factor_name][factor_category][class_name] = class_count
 
         # This deletes rows containing only zeros,
         # because scipy.stats.chi2_contingency fails when there are rows containing only zeros.
@@ -299,24 +344,7 @@ def parity(metadata: Metadata) -> ParityOutput[NDArray[np.float64]]:
         chi_scores[i] = chi2
         p_values[i] = p
 
-    if not_enough_data:
-        factor_msg = []
-        for factor, fact_dict in not_enough_data.items():
-            stacked_msg = []
-            for key, value in fact_dict.items():
-                msg = []
-                for item in value:
-                    msg.append(f"label {item[0]}: {item[1]} occurrences")
-                flat_msg = "\n\t\t".join(msg)
-                stacked_msg.append(f"value {key} - {flat_msg}\n\t")
-            factor_msg.append(factor + " - " + "".join(stacked_msg))
-
-        message = "\n".join(factor_msg)
-
-        warnings.warn(
-            f"The following factors did not meet the recommended 5 occurrences for each value-label combination. \n\
-            Recommend rerunning parity after adjusting the following factor-value-label combinations: \n{message}",
-            UserWarning,
-        )
+    if insufficient_data:
+        warnings.warn("Some factors did not meet the recommended 5 occurrences for each value-label combination.")
 
-    return ParityOutput(chi_scores, p_values, metadata.discrete_factor_names)
+    return ParityOutput(chi_scores, p_values, metadata.discrete_factor_names, insufficient_data)