dataeval 0.86.0__py3-none-any.whl → 0.86.2__py3-none-any.whl

dataeval/outputs/_bias.py

@@ -4,7 +4,7 @@ __all__ = []
 
 import contextlib
 from dataclasses import asdict, dataclass
-from typing import Any, Literal, TypeVar, overload
+from typing import Any, TypeVar
 
 import numpy as np
 import pandas as pd
@@ -128,33 +128,30 @@ class CoverageOutput(Output):
 
         import matplotlib.pyplot as plt
 
+        images = Images(images) if isinstance(images, Dataset) else images
+        if np.max(self.uncovered_indices) > len(images):
+            raise ValueError(
+                f"Uncovered indices {self.uncovered_indices} specify images "
+                f"unavailable in the provided number of images {len(images)}."
+            )
+
         # Determine which images to plot
         selected_indices = self.uncovered_indices[:top_k]
 
-        images = Images(images) if isinstance(images, Dataset) else images
-
         # Plot the images
         num_images = min(top_k, len(selected_indices))
 
         rows = int(np.ceil(num_images / 3))
-        fig, axs = plt.subplots(rows, 3, figsize=(9, 3 * rows))
-
-        if rows == 1:
-            for j in range(3):
-                if j >= len(selected_indices):
-                    continue
-                image = channels_first_to_last(as_numpy(images[selected_indices[j]]))
-                axs[j].imshow(image)
-                axs[j].axis("off")
-        else:
-            for i in range(rows):
-                for j in range(3):
-                    i_j = i * 3 + j
-                    if i_j >= len(selected_indices):
-                        continue
-                    image = channels_first_to_last(as_numpy(images[selected_indices[i_j]]))
-                    axs[i, j].imshow(image)
-                    axs[i, j].axis("off")
+        cols = min(3, num_images)
+        fig, axs = plt.subplots(rows, cols, figsize=(3 * cols, 3 * rows))
+
+        for image, ax in zip(images[:num_images], axs.flat):
+            image = channels_first_to_last(as_numpy(image))
+            ax.imshow(image)
+            ax.axis("off")
+
+        for ax in axs.flat[num_images:]:
+            ax.axis("off")
 
         fig.tight_layout()
         return fig
@@ -202,52 +199,11 @@ class BalanceOutput(Output):
     factor_names: list[str]
     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 = np.asarray([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.
@@ -260,8 +216,6 @@ 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
 
         Returns
         -------
@@ -275,10 +229,10 @@ class BalanceOutput(Output):
             if row_labels is None:
                 row_labels = self.class_names
             if col_labels is None:
-                col_labels = self._by_factor_type("factor_names", factor_type)
+                col_labels = self.factor_names
 
             fig = heatmap(
-                self._by_factor_type("classwise", factor_type),
+                self.classwise,
                 row_labels,
                 col_labels,
                 xlabel="Factors",
@@ -289,8 +243,8 @@ class BalanceOutput(Output):
             # Combine balance and factors results
             data = np.concatenate(
                 [
-                    self._by_factor_type("balance", factor_type)[np.newaxis, 1:],
-                    self._by_factor_type("factors", factor_type),
+                    self.balance[np.newaxis, 1:],
+                    self.factors,
                 ],
                 axis=0,
             )
@@ -299,7 +253,7 @@ class BalanceOutput(Output):
             # 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._by_factor_type("factor_names", factor_type)
+            heat_labels = self.factor_names
 
             if row_labels is None:
                 row_labels = heat_labels[:-1]
@@ -379,7 +333,7 @@ class DiversityOutput(Output):
             import matplotlib.pyplot as plt
 
             fig, ax = plt.subplots(figsize=(8, 8))
-            heat_labels = np.concatenate((["class"], self.factor_names))
+            heat_labels = ["class_labels"] + self.factor_names
             ax.bar(heat_labels, self.diversity_index)
             ax.set_xlabel("Factors")
             plt.setp(ax.get_xticklabels(), rotation=45, ha="right", rotation_mode="anchor")

dataeval/outputs/_drift.py

@@ -103,19 +103,13 @@ class DriftMVDCOutput(PerMetricResult):
         metric = Metric(display_name="Domain Classifier", column_name="domain_classifier_auroc")
         super().__init__(results_data, [metric])
 
-    def plot(self, showme: bool = True) -> Figure:
+    def plot(self) -> Figure:
         """
         Render the roc_auc metric over the train/test data in relation to the threshold.
 
-        Parameters
-        ----------
-        showme : bool, default True
-            Option to display the figure.
-
         Returns
         -------
         matplotlib.figure.Figure
-
         """
         import matplotlib.pyplot as plt
 
@@ -146,6 +140,4 @@ class DriftMVDCOutput(PerMetricResult):
             ax.set_ylabel("ROC AUC", fontsize=7)
             ax.set_xlabel("Chunk Index", fontsize=7)
             ax.set_ylim((0.0, 1.1))
-            if showme:
-                plt.show()
         return fig

dataeval/outputs/_linters.py

@@ -43,10 +43,12 @@ class DuplicatesOutput(Output, Generic[TIndexCollection]):
     near: list[TIndexCollection]
 
 
-def _reorganize_by_class_and_metric(result: IndexIssueMap, lstats: LabelStatsOutput):
+def _reorganize_by_class_and_metric(
+    result: IndexIssueMap, lstats: LabelStatsOutput
+) -> tuple[dict[str, list[int]], dict[str, dict[str, int]]]:
     """Flip result from grouping by image to grouping by class and metric"""
-    metrics = {}
-    class_wise = {label: {} for label in lstats.class_names}
+    metrics: dict[str, list[int]] = {}
+    class_wise: dict[str, dict[str, int]] = {label: {} for label in lstats.class_names}
 
     # Group metrics and calculate class-wise counts
     for img, group in result.items():
@@ -59,7 +61,7 @@ def _reorganize_by_class_and_metric(result: IndexIssueMap, lstats: LabelStatsOut
     return metrics, class_wise
 
 
-def _create_table(metrics, class_wise):
+def _create_table(metrics: dict[str, list[int]], class_wise: dict[str, dict[str, int]]) -> list[str]:
     """Create table for displaying the results"""
     max_class_length = max(len(str(label)) for label in class_wise) + 2
     max_total = max(len(metrics[group]) for group in metrics) + 2
@@ -69,7 +71,7 @@ def _create_table(metrics, class_wise):
         + [f"{group:^{max(5, len(str(group))) + 2}}" for group in sorted(metrics.keys())]
         + [f"{'Total':<{max_total}}"]
     )
-    table_rows = []
+    table_rows: list[str] = []
 
     for class_cat, results in class_wise.items():
         table_value = [f"{class_cat:>{max_class_length}}"]
@@ -81,15 +83,14 @@ def _create_table(metrics, class_wise):
         table_value.append(f"{total:^{max_total}}")
         table_rows.append(" | ".join(table_value))
 
-    table = [table_header] + table_rows
-    return table
+    return [table_header] + table_rows
 
 
-def _create_pandas_dataframe(class_wise):
+def _create_pandas_dataframe(class_wise: dict[str, dict[str, int]]) -> list[dict[str, str | int]]:
     """Create data for pandas dataframe"""
     data = []
     for label, metrics_dict in class_wise.items():
-        row = {"Class": label}
+        row: dict[str, str | int] = {"Class": label}
         total = sum(metrics_dict.values())
         row.update(metrics_dict)  # Add metric counts
         row["Total"] = total
@@ -118,8 +119,7 @@ class OutliersOutput(Output, Generic[TIndexIssueMap]):
     def __len__(self) -> int:
         if isinstance(self.issues, dict):
             return len(self.issues)
-        else:
-            return sum(len(d) for d in self.issues)
+        return sum(len(d) for d in self.issues)
 
     def to_table(self, labelstats: LabelStatsOutput) -> str:
         """

dataeval/outputs/_stats.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 __all__ = []
 
 from dataclasses import dataclass
-from typing import Any, Iterable, NamedTuple, Optional, Union
+from typing import TYPE_CHECKING, Any, Iterable, NamedTuple, Optional, Sequence, Union
 
 import numpy as np
 import pandas as pd
@@ -13,10 +13,16 @@ from typing_extensions import TypeAlias
 from dataeval.outputs._base import Output
 from dataeval.utils._plot import channel_histogram_plot, histogram_plot
 
+if TYPE_CHECKING:
+    from matplotlib.figure import Figure
+
 OptionalRange: TypeAlias = Optional[Union[int, Iterable[int]]]
 
 SOURCE_INDEX = "source_index"
-BOX_COUNT = "box_count"
+OBJECT_COUNT = "object_count"
+IMAGE_COUNT = "image_count"
+
+BASE_ATTRS = (SOURCE_INDEX, OBJECT_COUNT, IMAGE_COUNT)
 
 
 class SourceIndex(NamedTuple):
@@ -51,17 +57,24 @@ class BaseStatsOutput(Output):
     ----------
     source_index : List[SourceIndex]
         Mapping from statistic to source image, box and channel index
-    box_count : NDArray[np.uint16]
+    object_count : NDArray[np.uint16]
+        The number of detected objects in each image
     """
 
     source_index: list[SourceIndex]
-    box_count: NDArray[np.uint16]
+    object_count: NDArray[np.uint16]
+    image_count: int
 
     def __post_init__(self) -> None:
-        length = len(self.source_index)
-        bad = {k: len(v) for k, v in self.data().items() if k not in [SOURCE_INDEX, BOX_COUNT] and len(v) != length}
-        if bad:
-            raise ValueError(f"All values must have the same length as source_index. Bad values: {str(bad)}.")
+        si_length = len(self.source_index)
+        mismatch = {k: len(v) for k, v in self.data().items() if k not in BASE_ATTRS and len(v) != si_length}
+        if mismatch:
+            raise ValueError(f"All values must have the same length as source_index. Bad values: {str(mismatch)}.")
+        oc_length = len(self.object_count)
+        if oc_length != self.image_count:
+            raise ValueError(
+                f"Total object counts per image does not match image count. {oc_length} != {self.image_count}."
+            )
 
     def get_channel_mask(
         self,
@@ -123,21 +136,64 @@ class BaseStatsOutput(Output):
 
         return max_channels, ch_mask
 
-    def factors(self) -> dict[str, NDArray[Any]]:
+    def factors(
+        self,
+        filter: str | Sequence[str] | None = None,  # noqa: A002
+        exclude_constant: bool = False,
+    ) -> dict[str, NDArray[Any]]:
+        """
+        Returns all 1-dimensional data as a dictionary of numpy arrays.
+
+        Parameters
+        ----------
+        filter : str, Sequence[str] or None, default None:
+            If provided, only returns keys that match the filter.
+        exclude_constant : bool, default False
+            If True, exclude arrays that contain only a single unique value.
+
+        Returns
+        -------
+        dict[str, NDArray[Any]]
+        """
+        filter_ = [filter] if isinstance(filter, str) else filter
         return {
             k: v
             for k, v in self.data().items()
-            if k not in (SOURCE_INDEX, BOX_COUNT) and isinstance(v, np.ndarray) and v[v != 0].size > 0 and v.ndim == 1
+            if k not in BASE_ATTRS
+            and (filter_ is None or k in filter_)
+            and isinstance(v, np.ndarray)
+            and v.ndim == 1
+            and (not exclude_constant or len(np.unique(v)) > 1)
         }
 
     def plot(
         self, log: bool, channel_limit: int | None = None, channel_index: int | Iterable[int] | None = None
-    ) -> None:
+    ) -> Figure:
+        """
+        Plots the statistics as a set of histograms.
+
+        Parameters
+        ----------
+        log : bool
+            If True, plots the histograms on a logarithmic scale.
+        channel_limit : int or None
+            The maximum number of channels to plot. If None, all channels are plotted.
+        channel_index : int, Iterable[int] or None
+            The index or indices of the channels to plot. If None, all channels are plotted.
+
+        Returns
+        -------
+        matplotlib.Figure
+        """
+        from matplotlib.figure import Figure
+
         max_channels, ch_mask = self._get_channels(channel_limit, channel_index)
+        factors = self.factors(exclude_constant=True)
+        if not factors:
+            return Figure()
         if max_channels == 1:
-            histogram_plot(self.factors(), log)
-        else:
-            channel_histogram_plot(self.factors(), log, max_channels, ch_mask)
+            return histogram_plot(factors, log)
+        return channel_histogram_plot(factors, log, max_channels, ch_mask)
 
 
 @dataclass(frozen=True)
@@ -147,9 +203,9 @@ class DimensionStatsOutput(BaseStatsOutput):
 
     Attributes
     ----------
-    left : NDArray[np.int32]
+    offset_x : NDArray[np.int32]
         Offsets from the left edge of images in pixels
-    top : NDArray[np.int32]
+    offset_y : NDArray[np.int32]
         Offsets from the top edge of images in pixels
     width : NDArray[np.uint32]
         Width of the images in pixels
@@ -160,25 +216,28 @@ class DimensionStatsOutput(BaseStatsOutput):
     size : NDArray[np.uint32]
         Size of the images in pixels
     aspect_ratio : NDArray[np.float16]
-        :term:`ASspect Ratio<Aspect Ratio>` of the images (width/height)
+        :term:`Aspect Ratio<Aspect Ratio>` of the images (width/height)
     depth : NDArray[np.uint8]
         Color depth of the images in bits
-    center : NDArray[np.uint16]
+    center : NDArray[np.uint32]
         Offset from center in [x,y] coordinates of the images in pixels
-    distance : NDArray[np.float16]
+    distance_center : NDArray[np.float32]
         Distance in pixels from center
+    distance_edge : NDArray[np.uint32]
+        Distance in pixels from nearest edge
     """
 
-    left: NDArray[np.int32]
-    top: NDArray[np.int32]
+    offset_x: NDArray[np.int32]
+    offset_y: NDArray[np.int32]
     width: NDArray[np.uint32]
     height: NDArray[np.uint32]
     channels: NDArray[np.uint8]
     size: NDArray[np.uint32]
     aspect_ratio: NDArray[np.float16]
     depth: NDArray[np.uint8]
-    center: NDArray[np.int16]
-    distance: NDArray[np.float16]
+    center: NDArray[np.int32]
+    distance_center: NDArray[np.float32]
+    distance_edge: NDArray[np.uint32]
 
 
 @dataclass(frozen=True)

dataeval/outputs/_workflows.py

@@ -154,10 +154,10 @@ def calc_params(p_i: NDArray[Any], n_i: NDArray[Any], niter: int) -> NDArray[Any
         Array of parameters to recreate line of best fit
     """
 
-    def is_valid(f_new, x_new, f_old, x_old):
+    def is_valid(f_new, x_new, f_old, x_old) -> bool:  # noqa: ANN001
         return f_new != np.nan
 
-    def f(x):
+    def f(x) -> float:  # noqa: ANN001
         try:
             return np.sum(np.square(p_i - f_out(n_i, x)))
         except RuntimeWarning:

dataeval/utils/_array.py

@@ -23,7 +23,7 @@ T = TypeVar("T", ArrayLike, np.ndarray, torch.Tensor)
 _np_dtype = TypeVar("_np_dtype", bound=np.generic)
 
 
-def _try_import(module_name) -> ModuleType | None:
+def _try_import(module_name: str) -> ModuleType | None:
     if module_name in _MODULE_CACHE:
         return _MODULE_CACHE[module_name]
 
@@ -148,8 +148,7 @@ def ensure_embeddings(
 
     if dtype is None:
         return embeddings
-    else:
-        return arr
+    return arr
 
 
 @overload
@@ -174,10 +173,9 @@ def flatten(array: ArrayLike) -> NDArray[Any] | torch.Tensor:
     if isinstance(array, np.ndarray):
         nparr = as_numpy(array)
         return nparr.reshape((nparr.shape[0], -1))
-    elif isinstance(array, torch.Tensor):
+    if isinstance(array, torch.Tensor):
         return torch.flatten(array, start_dim=1)
-    else:
-        raise TypeError(f"Unsupported array type {type(array)}.")
+    raise TypeError(f"Unsupported array type {type(array)}.")
 
 
 _TArray = TypeVar("_TArray", bound=Array)
@@ -199,7 +197,6 @@ def channels_first_to_last(array: _TArray) -> _TArray:
     """
     if isinstance(array, np.ndarray):
         return np.transpose(array, (1, 2, 0))
-    elif isinstance(array, torch.Tensor):
+    if isinstance(array, torch.Tensor):
         return torch.permute(array, (1, 2, 0))
-    else:
-        raise TypeError(f"Unsupported array type {type(array)}.")
+    raise TypeError(f"Unsupported array type {type(array)}.")

dataeval/utils/_bin.py

@@ -195,5 +195,4 @@ def bin_by_clusters(data: NDArray[np.number[Any]]) -> NDArray[np.float64]:
     if extend_bins:
         bin_edges = np.concatenate([bin_edges, extend_bins])
 
-    bin_edges = np.sort(bin_edges)
-    return bin_edges
+    return np.sort(bin_edges)

dataeval/utils/_clusterer.py

@@ -4,6 +4,7 @@ __all__ = []
 
 import warnings
 from dataclasses import dataclass
+from typing import Any
 
 import numba
 import numpy as np
@@ -30,7 +31,9 @@ from dataeval.utils._fast_mst import calculate_neighbor_distances, minimum_spann
 
 
 @numba.njit(parallel=True, locals={"i": numba.types.int32})
-def compare_links_to_cluster_std(mst, clusters):
+def compare_links_to_cluster_std(
+    mst: NDArray[np.float32], clusters: NDArray[np.intp]
+) -> tuple[NDArray[np.int32], NDArray[np.int32]]:
     cluster_ids = np.unique(clusters)
     cluster_grouping = np.full(mst.shape[0], -1, dtype=np.int16)
 
@@ -79,7 +82,7 @@ def cluster(data: ArrayLike) -> ClusterData:
     cluster_selection_epsilon = 0.0
     # cluster_selection_method = "eom"
 
-    x = flatten(to_numpy(data))
+    x: NDArray[Any] = flatten(to_numpy(data))
     samples, features = x.shape  # Due to flatten(), we know shape has a length of 2
     if samples < 2:
         raise ValueError(f"Data should have at least 2 samples; got {samples}")
@@ -125,9 +128,9 @@ def cluster(data: ArrayLike) -> ClusterData:
     return ClusterData(clusters, mst, linkage_tree, condensed_tree, membership_strengths, kneighbors, kdistances)
 
 
-def sorted_union_find(index_groups):
+def sorted_union_find(index_groups: NDArray[np.int32]) -> list[list[np.int32]]:
     """Merges and sorts groups of indices that share any common index"""
-    groups = [[np.int32(x) for x in range(0)] for y in range(0)]
+    groups: list[list[np.int32]] = [[np.int32(x) for x in range(0)] for y in range(0)]
     uniques, inverse = np.unique(index_groups, return_inverse=True)
     inverse = inverse.flatten()
     disjoint_set = ds_rank_create(uniques.size)

dataeval/utils/_fast_mst.py

@@ -6,9 +6,11 @@
 __all__ = []
 
 import warnings
+from typing import Any
 
 import numba
 import numpy as np
+from numpy.typing import NDArray
 from sklearn.neighbors import NearestNeighbors
 
 with warnings.catch_warnings():
@@ -17,24 +19,26 @@ with warnings.catch_warnings():
 
 
 @numba.njit()
-def _ds_union_by_rank(disjoint_set, point, nbr):
+def _ds_union_by_rank(disjoint_set: tuple[NDArray[np.int32], NDArray[np.int32]], point: int, nbr: int) -> int:
     y = ds_find(disjoint_set, point)
     x = ds_find(disjoint_set, nbr)
 
     if x == y:
         return 0
 
-    if disjoint_set.rank[x] < disjoint_set.rank[y]:
+    if disjoint_set[1][x] < disjoint_set[1][y]:
         x, y = y, x
 
-    disjoint_set.parent[y] = x
-    if disjoint_set.rank[x] == disjoint_set.rank[y]:
-        disjoint_set.rank[x] += 1
+    disjoint_set[0][y] = x
+    if disjoint_set[1][x] == disjoint_set[1][y]:
+        disjoint_set[1][x] += 1
     return 1
 
 
 @numba.njit(locals={"i": numba.types.uint32, "nbr": numba.types.uint32, "dist": numba.types.float32})
-def _init_tree(n_neighbors, n_distance):
+def _init_tree(
+    n_neighbors: NDArray[np.intp], n_distance: NDArray[np.float32]
+) -> tuple[NDArray[np.float32], int, tuple[NDArray[np.int32], NDArray[np.int32]], NDArray[np.uint32]]:
     # Initial graph to hold tree connections
     tree = np.zeros((n_neighbors.size - 1, 3), dtype=np.float32)
     disjoint_set = ds_rank_create(n_neighbors.size)
@@ -56,7 +60,13 @@ def _init_tree(n_neighbors, n_distance):
 
 
 @numba.njit(locals={"i": numba.types.uint32, "nbr": numba.types.uint32})
-def _update_tree_by_distance(tree, int_tree, disjoint_set, n_neighbors, n_distance):
+def _update_tree_by_distance(
+    tree: NDArray[np.float32],
+    int_tree: int,
+    disjoint_set: tuple[NDArray[np.int32], NDArray[np.int32]],
+    n_neighbors: NDArray[np.uint32],
+    n_distance: NDArray[np.float32],
+) -> tuple[NDArray[np.float32], int, tuple[NDArray[np.int32], NDArray[np.int32]], NDArray[np.uint32]]:
     cluster_points = np.empty(n_neighbors.size, dtype=np.uint32)
     sort_dist = np.argsort(n_distance)
     dist_sorted = n_distance[sort_dist]
@@ -80,9 +90,9 @@ def _update_tree_by_distance(tree, int_tree, disjoint_set, n_neighbors, n_distan
 
 
 @numba.njit(locals={"i": numba.types.uint32})
-def _cluster_edges(tracker, last_idx, cluster_distances):
+def _cluster_edges(tracker: NDArray[Any], last_idx: int, cluster_distances: NDArray[Any]) -> list[NDArray[np.intp]]:
     cluster_ids = np.unique(tracker)
-    edge_points = []
+    edge_points: list[NDArray[np.intp]] = []
     for idx in range(cluster_ids.size):
         cluster_points = np.nonzero(tracker == cluster_ids[idx])[0]
         cluster_size = cluster_points.size
@@ -102,14 +112,16 @@ def _cluster_edges(tracker, last_idx, cluster_distances):
     return edge_points
 
 
-def _compute_nn(dataA, dataB, k):
+def _compute_nn(dataA: NDArray[Any], dataB: NDArray[Any], k: int) -> tuple[NDArray[np.int32], NDArray[np.float32]]:
     distances, neighbors = NearestNeighbors(n_neighbors=k + 1, algorithm="brute").fit(dataA).kneighbors(dataB)
     neighbors = np.array(neighbors[:, 1 : k + 1], dtype=np.int32)
     distances = np.array(distances[:, 1 : k + 1], dtype=np.float32)
     return neighbors, distances
 
 
-def _calculate_cluster_neighbors(data, groups, point_array):
+def _calculate_cluster_neighbors(
+    data: NDArray[Any], groups: list[NDArray[np.intp]], point_array: NDArray[Any]
+) -> tuple[NDArray[np.uint32], NDArray[np.float32]]:
     """Rerun nearest neighbor based on clusters"""
     cluster_neighbors = np.zeros(point_array.size, dtype=np.uint32)
     cluster_nbr_distances = np.full(point_array.size, np.inf, dtype=np.float32)
@@ -126,7 +138,9 @@ def _calculate_cluster_neighbors(data, groups, point_array):
     return cluster_neighbors, cluster_nbr_distances
 
 
-def minimum_spanning_tree(data, neighbors, distances):
+def minimum_spanning_tree(
+    data: NDArray[Any], neighbors: NDArray[np.int32], distances: NDArray[np.float32]
+) -> NDArray[np.float32]:
     # Transpose arrays to get number of samples along a row
     k_neighbors = neighbors.T.astype(np.uint32).copy()
     k_distances = distances.T.astype(np.float32).copy()
@@ -168,7 +182,7 @@ def minimum_spanning_tree(data, neighbors, distances):
     return tree
 
 
-def calculate_neighbor_distances(data: np.ndarray, k: int = 10):
+def calculate_neighbor_distances(data: np.ndarray, k: int = 10) -> tuple[NDArray[np.int32], NDArray[np.float32]]:
     # Have the potential to add in other distance calculations - supported calculations:
     # https://github.com/lmcinnes/pynndescent/blob/master/pynndescent/pynndescent_.py#L524
     try: