dataeval 0.61.0__py3-none-any.whl → 0.63.0__py3-none-any.whl

dataeval/_internal/metrics/divergence.py

@@ -7,46 +7,23 @@ from typing import Any, Callable, Dict, Literal
 
 import numpy as np
 
+from dataeval._internal.functional.divergence import divergence_fnn, divergence_mst
+from dataeval._internal.interop import ArrayLike, to_numpy
 from dataeval._internal.metrics.base import EvaluateMixin, MethodsMixin
 
-from .utils import compute_neighbors, minimum_spanning_tree
-
-
-def _mst(data: np.ndarray, labels: np.ndarray) -> int:
-    mst = minimum_spanning_tree(data).toarray()
-    edgelist = np.transpose(np.nonzero(mst))
-    errors = np.sum(labels[edgelist[:, 0]] != labels[edgelist[:, 1]])
-    return errors
-
-
-def _fnn(data: np.ndarray, labels: np.ndarray) -> int:
-    nn_indices = compute_neighbors(data, data)
-    errors = np.sum(np.abs(labels[nn_indices] - labels))
-    return errors
-
-
 _METHODS = Literal["MST", "FNN"]
 _FUNCTION = Callable[[np.ndarray, np.ndarray], int]
 
 
 class Divergence(EvaluateMixin, MethodsMixin[_METHODS, _FUNCTION]):
     """
-    Calculates the estimated divergence between two datasets
+    Calculates the estimated HP divergence between two datasets
 
     Parameters
     ----------
-    data_a : np.ndarray
-        Array of images or image embeddings to compare
-    data_b : np.ndarray
-        Array of images or image embeddings to compare
     method : Literal["MST, "FNN"], default "MST"
         Method used to estimate dataset divergence
 
-    See Also
-    --------
-        For more information about this divergence, its formal definition,
-        and its associated estimators see https://arxiv.org/abs/1412.6534.
-
     Warning
     -------
         MST is very slow in this implementation, this is unlike matlab where
@@ -55,38 +32,69 @@ class Divergence(EvaluateMixin, MethodsMixin[_METHODS, _FUNCTION]):
         Source of slowdown:
         conversion to and from CSR format adds ~10% of the time diff between
         1nn and scipy mst function the remaining 90%
+
+    References
+    ----------
+    For more information about this divergence, its formal definition,
+    and its associated estimators see https://arxiv.org/abs/1412.6534.
+
+    Examples
+    --------
+    Initialize the Divergence class:
+
+    >>> divert = Divergence()
+
+    Specify the method:
+
+    >>> divert = Divergence(method="FNN")
     """
 
-    def __init__(
-        self,
-        data_a: np.ndarray,
-        data_b: np.ndarray,
-        method: _METHODS = "MST",
-    ) -> None:
-        self.data_a = data_a
-        self.data_b = data_b
+    def __init__(self, method: _METHODS = "MST") -> None:
         self._set_method(method)
 
     @classmethod
     def _methods(cls) -> Dict[str, _FUNCTION]:
-        return {"FNN": _fnn, "MST": _mst}
+        return {"FNN": divergence_fnn, "MST": divergence_mst}
 
-    def evaluate(self) -> Dict[str, Any]:
+    def evaluate(self, data_a: ArrayLike, data_b: ArrayLike) -> Dict[str, Any]:
         """
         Calculates the divergence and any errors between the datasets
 
+        Parameters
+        ----------
+        data_a : ArrayLike, shape - (N, P)
+            A dataset in an ArrayLike format to compare.
+            Function expects the data to have 2 dimensions, N number of observations in a P-dimesionial space.
+        data_b : ArrayLike, shape - (N, P)
+            A dataset in an ArrayLike format to compare.
+            Function expects the data to have 2 dimensions, N number of observations in a P-dimesionial space.
+
         Returns
         -------
         Dict[str, Any]
-            dp : float
+            divergence : float
                 divergence value between 0.0 and 1.0
-            errors : int
-                the number of differing edges
+            error : int
+                the number of differing edges between the datasets
+
+        Notes
+        -----
+        The divergence value indicates how similar the 2 datasets are
+        with 0 indicating approximately identical data distributions.
+
+        Examples
+        --------
+        Evaluate the datasets:
+
+        >>> divert.evaluate(datasetA, datasetB)
+        {'divergence': 0.28, 'error': 36.0}
         """
-        N = self.data_a.shape[0]
-        M = self.data_b.shape[0]
+        a = to_numpy(data_a)
+        b = to_numpy(data_b)
+        N = a.shape[0]
+        M = b.shape[0]
 
-        stacked_data = np.vstack((self.data_a, self.data_b))
+        stacked_data = np.vstack((a, b))
         labels = np.vstack([np.zeros([N, 1]), np.ones([M, 1])])
 
         errors = self._method(stacked_data, labels)