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

dataeval/_internal/detectors/linter.py

@@ -1,6 +1,7 @@
-from typing import Literal, Optional, Sequence, Union
+from typing import Iterable, Literal, Optional, Sequence, Union
 
 import numpy as np
+from numpy.typing import ArrayLike
 
 from dataeval._internal.flags import ImageProperty, ImageVisuals, LinterFlags
 from dataeval._internal.metrics.stats import ImageStats
@@ -30,25 +31,79 @@ def _get_outlier_mask(
 
 
 class Linter:
-    """
-    Calculates statistical outliers of a dataset using various statistical
-    tests applied to each image
+    r"""
+    Calculates statistical outliers of a dataset using various statistical tests applied to each image
+
+    Parameters
+    ----------
+    flags : [ImageProperty | ImageStatistics | ImageVisuals], default None
+        Metric(s) to calculate for each image - calculates all metrics if None
+    outlier_method : ["modzscore" | "zscore" | "iqr"], optional - default "modzscore"
+        Statistical method used to identify outliers
+    outlier_threshold : float, optional - default None
+        Threshold value for the given ``outlier_method``, above which data is considered an outlier.
+        Uses method specific default if `None`
+
+    Attributes
+    ----------
+    stats : ImageStats
+        Class to hold the value of each metric for each image
+
+    See Also
+    --------
+    Duplicates
+
+    Notes
+    ------
+    There are 3 different statistical methods:
+
+    - zscore
+    - modzscore
+    - iqr
+
+    | The z score method is based on the difference between the data point and the mean of the data.
+        The default threshold value for `zscore` is 3.
+    | Z score = :math:`|x_i - \mu| / \sigma`
+
+    | The modified z score method is based on the difference between the data point and the median of the data.
+        The default threshold value for `modzscore` is 3.5.
+    | Modified z score = :math:`0.6745 * |x_i - x̃| / MAD`, where :math:`MAD` is the median absolute deviation
+
+    | The interquartile range method is based on the difference between the data point and
+        the difference between the 75th and 25th qartile. The default threshold value for `iqr` is 1.5.
+    | Interquartile range = :math:`threshold * (Q_3 - Q_1)`
+
+    Examples
+    --------
+    Initialize the Linter class:
+
+    >>> lint = Linter()
+
+    Specifying specific metrics to analyze:
+
+    >>> lint = Linter(flags=[ImageProperty.SIZE, ImageVisuals.ALL])
+
+    Specifying an outlier method:
+
+    >>> lint = Linter(outlier_method="iqr")
+
+    Specifying an outlier method and threshold:
+
+    >>> lint = Linter(outlier_method="zscore", outlier_threshold=2.5)
     """
 
     def __init__(
         self,
-        images: np.ndarray,
         flags: Optional[Union[LinterFlags, Sequence[LinterFlags]]] = None,
+        outlier_method: Literal["zscore", "modzscore", "iqr"] = "modzscore",
+        outlier_threshold: Optional[float] = None,
     ):
         flags = flags if flags is not None else (ImageProperty.ALL, ImageVisuals.ALL)
         self.stats = ImageStats(flags)
-        self.images = images
+        self.outlier_method: Literal["zscore", "modzscore", "iqr"] = outlier_method
+        self.outlier_threshold = outlier_threshold
 
-    def _get_outliers(
-        self,
-        outlier_method: Literal["zscore", "modzscore", "iqr"] = "modzscore",
-        outlier_threshold: Optional[float] = None,
-    ) -> dict:
+    def _get_outliers(self) -> dict:
         flagged_images = {}
 
         for stat, values in self.results.items():
@@ -56,23 +111,37 @@ class Linter:
                 continue
 
             if values.ndim == 1 and np.std(values) != 0:
-                mask = _get_outlier_mask(values, outlier_method, outlier_threshold)
+                mask = _get_outlier_mask(values, self.outlier_method, self.outlier_threshold)
                 indices = np.flatnonzero(mask)
                 for i, value in zip(indices, values[mask]):
                     flagged_images.setdefault(i, {}).update({stat: np.round(value, 2)})
 
         return dict(sorted(flagged_images.items()))
 
-    def evaluate(self) -> dict:
+    def evaluate(self, images: Iterable[ArrayLike]) -> dict:
         """
-        Returns indices of outliers with and the issues identified for each
+        Returns indices of outliers with the issues identified for each
+
+        Parameters
+        ----------
+        images : Iterable[ArrayLike], shape - (N, C, H, W)
+            A dataset in an ArrayLike format.
+            Function expects the data to have 3 dimensions, CxHxW.
 
         Returns
         -------
         Dict[int, Dict[str, float]]
-            Dictionary containing the indices of outliers and a dictionary issues and calculated values
+            Dictionary containing the indices of outliers and a dictionary showing
+            the issues and calculated values for the given index.
+
+        Example
+        -------
+        Evaluate the dataset:
+
+        >>> lint.evaluate(images)
+        {18: {'brightness': 0.78}, 25: {'brightness': 0.98}}
         """
         self.stats.reset()
-        self.stats.update(self.images)
+        self.stats.update(images)
         self.results = self.stats.compute()
         return self._get_outliers()

dataeval/_internal/detectors/ood/ae.py

@@ -10,8 +10,10 @@ from typing import Callable
 
 import keras
 import numpy as np
+from numpy.typing import ArrayLike
 
 from dataeval._internal.detectors.ood.base import OODBase, OODScore
+from dataeval._internal.interop import to_numpy
 from dataeval._internal.models.tensorflow.autoencoder import AE
 from dataeval._internal.models.tensorflow.utils import predict_batch
 
@@ -30,7 +32,7 @@ class OOD_AE(OODBase):
 
     def fit(
         self,
-        x_ref: np.ndarray,
+        x_ref: ArrayLike,
         threshold_perc: float = 100.0,
         loss_fn: Callable = keras.losses.MeanSquaredError(),
         optimizer: keras.optimizers.Optimizer = keras.optimizers.Adam,
@@ -43,7 +45,7 @@ class OOD_AE(OODBase):
 
         Parameters
         ----------
-        x_ref : np.ndarray
+        x_ref : ArrayLike
             Training batch.
         threshold_perc : float, default 100.0
             Percentage of reference data that is normal.
@@ -58,10 +60,10 @@ class OOD_AE(OODBase):
         verbose : bool, default True
             Whether to print training progress.
         """
-        super().fit(x_ref, threshold_perc, loss_fn, optimizer, epochs, batch_size, verbose)
+        super().fit(to_numpy(x_ref), threshold_perc, loss_fn, optimizer, epochs, batch_size, verbose)
 
-    def score(self, X: np.ndarray, batch_size: int = int(1e10)) -> OODScore:
-        self._validate(X)
+    def score(self, X: ArrayLike, batch_size: int = int(1e10)) -> OODScore:
+        self._validate(X := to_numpy(X))
 
         # reconstruct instances
         X_recon = predict_batch(X, self.model, batch_size=batch_size)

dataeval/_internal/detectors/ood/aegmm.py

@@ -9,9 +9,10 @@ Licensed under Apache Software License (Apache 2.0)
 from typing import Callable
 
 import keras
-import numpy as np
+from numpy.typing import ArrayLike
 
 from dataeval._internal.detectors.ood.base import OODGMMBase, OODScore
+from dataeval._internal.interop import to_numpy
 from dataeval._internal.models.tensorflow.autoencoder import AEGMM
 from dataeval._internal.models.tensorflow.gmm import gmm_energy
 from dataeval._internal.models.tensorflow.losses import LossGMM
@@ -32,7 +33,7 @@ class OOD_AEGMM(OODGMMBase):
 
     def fit(
         self,
-        x_ref: np.ndarray,
+        x_ref: ArrayLike,
         threshold_perc: float = 100.0,
         loss_fn: Callable = LossGMM(),
         optimizer: keras.optimizers.Optimizer = keras.optimizers.Adam,
@@ -45,7 +46,7 @@ class OOD_AEGMM(OODGMMBase):
 
         Parameters
         ----------
-        x_ref : np.ndarray
+        x_ref : ArrayLike
             Training batch.
         threshold_perc : float, default 100.0
             Percentage of reference data that is normal.
@@ -62,8 +63,8 @@ class OOD_AEGMM(OODGMMBase):
         """
         super().fit(x_ref, threshold_perc, loss_fn, optimizer, epochs, batch_size, verbose)
 
-    def score(self, X: np.ndarray, batch_size: int = int(1e10)) -> OODScore:
-        self._validate(X)
+    def score(self, X: ArrayLike, batch_size: int = int(1e10)) -> OODScore:
+        self._validate(X := to_numpy(X))
         _, z, _ = predict_batch(X, self.model, batch_size=batch_size)
         energy, _ = gmm_energy(z, self.gmm_params, return_mean=False)
         return OODScore(energy.numpy())  # type: ignore

dataeval/_internal/detectors/ood/base.py

@@ -12,7 +12,9 @@ from typing import Callable, Dict, List, Literal, NamedTuple, Optional, Tuple, c
 import keras
 import numpy as np
 import tensorflow as tf
+from numpy.typing import ArrayLike
 
+from dataeval._internal.interop import to_numpy
 from dataeval._internal.models.tensorflow.gmm import GaussianMixtureModelParams, gmm_params
 from dataeval._internal.models.tensorflow.trainer import trainer
 
@@ -66,13 +68,13 @@ class OODBase(ABC):
         self._validate(X)
 
     @abstractmethod
-    def score(self, X: np.ndarray, batch_size: int = int(1e10)) -> OODScore:
+    def score(self, X: ArrayLike, batch_size: int = int(1e10)) -> OODScore:
         """
         Compute instance and (optionally) feature level outlier scores.
 
         Parameters
         ----------
-        X : np.ndarray
+        X : ArrayLike
             Batch of instances.
         batch_size : int, default int(1e10)
             Batch size used when making predictions with the autoencoder.
@@ -87,7 +89,7 @@ class OODBase(ABC):
 
     def fit(
         self,
-        x_ref: np.ndarray,
+        x_ref: ArrayLike,
         threshold_perc: float,
         loss_fn: Callable,
         optimizer: keras.optimizers.Optimizer,
@@ -100,7 +102,7 @@ class OODBase(ABC):
 
         Parameters
         ----------
-        x_ref: : np.ndarray
+        x_ref: : ArrayLike
             Training batch.
         threshold_perc : float
             Percentage of reference data that is normal.
@@ -119,7 +121,7 @@ class OODBase(ABC):
         trainer(
             model=self.model,
             loss_fn=loss_fn,
-            x_train=x_ref,
+            x_train=to_numpy(x_ref),
             optimizer=optimizer,
             epochs=epochs,
             batch_size=batch_size,
@@ -132,7 +134,7 @@ class OODBase(ABC):
 
     def predict(
         self,
-        X: np.ndarray,
+        X: ArrayLike,
         batch_size: int = int(1e10),
         ood_type: Literal["feature", "instance"] = "instance",
     ) -> Dict[str, np.ndarray]:
@@ -141,18 +143,18 @@ class OODBase(ABC):
 
         Parameters
         ----------
-        X
+        X : ArrayLike
             Batch of instances.
-        ood_type
-            Predict out-of-distribution at the 'feature' or 'instance' level.
-        batch_size
+        batch_size : int, default int(1e10)
             Batch size used when making predictions with the autoencoder.
+        ood_type : Literal["feature", "instance"], default "instance"
+            Predict out-of-distribution at the 'feature' or 'instance' level.
 
         Returns
         -------
         Dictionary containing the outlier predictions and both feature and instance level outlier scores.
         """
-        self._validate_state(X)
+        self._validate_state(X := to_numpy(X))
         # compute outlier scores
         score = self.score(X, batch_size=batch_size)
         ood_pred = (score.get(ood_type) > self._threshold_score(ood_type)).astype(int)
@@ -171,7 +173,7 @@ class OODGMMBase(OODBase):
 
     def fit(
         self,
-        x_ref: np.ndarray,
+        x_ref: ArrayLike,
         threshold_perc: float,
         loss_fn: Callable[[tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor], tf.Tensor],
         optimizer: keras.optimizers.Optimizer,
@@ -183,7 +185,7 @@ class OODGMMBase(OODBase):
         trainer(
             model=self.model,
             loss_fn=loss_fn,
-            x_train=x_ref,
+            x_train=to_numpy(x_ref),
             optimizer=optimizer,
             epochs=epochs,
             batch_size=batch_size,

dataeval/_internal/detectors/ood/llr.py

@@ -14,8 +14,10 @@ import numpy as np
 import tensorflow as tf
 from keras.layers import Input
 from keras.models import Model
+from numpy.typing import ArrayLike
 
 from dataeval._internal.detectors.ood.base import OODBase, OODScore
+from dataeval._internal.interop import to_numpy
 from dataeval._internal.models.tensorflow.pixelcnn import PixelCNN
 from dataeval._internal.models.tensorflow.trainer import trainer
 from dataeval._internal.models.tensorflow.utils import predict_batch
@@ -125,7 +127,7 @@ class OOD_LLR(OODBase):
 
     def fit(
         self,
-        x_ref: np.ndarray,
+        x_ref: ArrayLike,
         threshold_perc: float = 100.0,
         loss_fn: Optional[Callable] = None,
         optimizer: keras.optimizers.Optimizer = keras.optimizers.Adam,
@@ -141,7 +143,7 @@ class OOD_LLR(OODBase):
 
         Parameters
         ----------
-        x_ref : np.ndarray
+        x_ref : ArrayLike
             Training batch.
         threshold_perc : float, default 100.0
             Percentage of reference data that is normal.
@@ -163,6 +165,7 @@ class OOD_LLR(OODBase):
         mutate_batch_size: int, default int(1e10)
             Batch size used to generate the mutations for the background dataset.
         """
+        x_ref = to_numpy(x_ref)
         input_shape = x_ref.shape[1:]
         optimizer = optimizer() if isinstance(optimizer, type) else optimizer
         # Separate into two separate optimizers, one for semantic model and one for background model
@@ -178,7 +181,7 @@ class OOD_LLR(OODBase):
 
         # create background data
         mutate_fn = partial(mutate_fn, **mutate_fn_kwargs)
-        X_back = predict_batch(x_ref, mutate_fn, batch_size=mutate_batch_size, dtype=x_ref.dtype)
+        X_back = predict_batch(x_ref, mutate_fn, batch_size=mutate_batch_size, dtype=x_ref.dtype)  # type: ignore
 
         # prepare sequential data
         if self.sequential and not self.has_log_prob:
@@ -275,10 +278,10 @@ class OOD_LLR(OODBase):
 
     def score(
         self,
-        X: np.ndarray,
+        X: ArrayLike,
         batch_size: int = int(1e10),
     ) -> OODScore:
-        self._validate(X)
+        self._validate(X := to_numpy(X))
         fscore = -self._llr(X, True, batch_size=batch_size)
         iscore = -self._llr(X, False, batch_size=batch_size)
         return OODScore(iscore, fscore)

dataeval/_internal/detectors/ood/vae.py

@@ -10,8 +10,10 @@ from typing import Callable
 
 import keras
 import numpy as np
+from numpy.typing import ArrayLike
 
 from dataeval._internal.detectors.ood.base import OODBase, OODScore
+from dataeval._internal.interop import to_numpy
 from dataeval._internal.models.tensorflow.autoencoder import VAE
 from dataeval._internal.models.tensorflow.losses import Elbo
 from dataeval._internal.models.tensorflow.utils import predict_batch
@@ -34,7 +36,7 @@ class OOD_VAE(OODBase):
 
     def fit(
         self,
-        x_ref: np.ndarray,
+        x_ref: ArrayLike,
         threshold_perc: float = 100.0,
         loss_fn: Callable = Elbo(0.05),
         optimizer: keras.optimizers.Optimizer = keras.optimizers.Adam,
@@ -47,7 +49,7 @@ class OOD_VAE(OODBase):
 
         Parameters
         ----------
-        x_ref : np.ndarray
+        x_ref : ArrayLike
             Training batch.
         threshold_perc : float, default 100.0
             Percentage of reference data that is normal.
@@ -64,8 +66,8 @@ class OOD_VAE(OODBase):
         """
         super().fit(x_ref, threshold_perc, loss_fn, optimizer, epochs, batch_size, verbose)
 
-    def score(self, X: np.ndarray, batch_size: int = int(1e10)) -> OODScore:
-        self._validate(X)
+    def score(self, X: ArrayLike, batch_size: int = int(1e10)) -> OODScore:
+        self._validate(X := to_numpy(X))
 
         # sample reconstructed instances
         X_samples = np.repeat(X, self.samples, axis=0)

dataeval/_internal/detectors/ood/vaegmm.py

@@ -10,8 +10,10 @@ from typing import Callable
 
 import keras
 import numpy as np
+from numpy.typing import ArrayLike
 
 from dataeval._internal.detectors.ood.base import OODGMMBase, OODScore
+from dataeval._internal.interop import to_numpy
 from dataeval._internal.models.tensorflow.autoencoder import VAEGMM
 from dataeval._internal.models.tensorflow.gmm import gmm_energy
 from dataeval._internal.models.tensorflow.losses import Elbo, LossGMM
@@ -35,7 +37,7 @@ class OOD_VAEGMM(OODGMMBase):
 
     def fit(
         self,
-        x_ref: np.ndarray,
+        x_ref: ArrayLike,
         threshold_perc: float = 100.0,
         loss_fn: Callable = LossGMM(elbo=Elbo(0.05)),
         optimizer: keras.optimizers.Optimizer = keras.optimizers.Adam,
@@ -48,7 +50,7 @@ class OOD_VAEGMM(OODGMMBase):
 
         Parameters
         ----------
-        X : np.ndarray
+        X : ArrayLike
             Training batch.
         threshold_perc : float, default 100.0
             Percentage of reference data that is normal.
@@ -65,8 +67,8 @@ class OOD_VAEGMM(OODGMMBase):
         """
         super().fit(x_ref, threshold_perc, loss_fn, optimizer, epochs, batch_size, verbose)
 
-    def score(self, X: np.ndarray, batch_size: int = int(1e10)) -> OODScore:
-        self._validate(X)
+    def score(self, X: ArrayLike, batch_size: int = int(1e10)) -> OODScore:
+        self._validate(X := to_numpy(X))
 
         # draw samples from latent space
         X_samples = np.repeat(X, self.samples, axis=0)

dataeval/_internal/interop.py

@@ -0,0 +1,43 @@
+from importlib import import_module
+from typing import Iterable, Optional
+
+import numpy as np
+from numpy.typing import ArrayLike
+
+module_cache = {}
+
+
+def try_import(module_name):
+    if module_name in module_cache:
+        return module_cache[module_name]
+
+    try:
+        module = import_module(module_name)
+    except ImportError:  # pragma: no cover - covered by test_mindeps.py
+        module = None
+
+    module_cache[module_name] = module
+    return module
+
+
+def to_numpy(array: Optional[ArrayLike]) -> np.ndarray:
+    if array is None:
+        return np.ndarray([])
+
+    if isinstance(array, np.ndarray):
+        return array
+
+    tf = try_import("tensorflow")
+    if tf and tf.is_tensor(array):
+        return array.numpy()  # type: ignore
+
+    torch = try_import("torch")
+    if torch and isinstance(array, torch.Tensor):
+        return array.detach().cpu().numpy()  # type: ignore
+
+    return np.asarray(array)
+
+
+def to_numpy_iter(iterable: Iterable[ArrayLike]):
+    for array in iterable:
+        yield to_numpy(array)

dataeval/_internal/metrics/balance.py

@@ -0,0 +1,180 @@
+import warnings
+from typing import Dict, List, NamedTuple, Sequence
+
+import numpy as np
+from numpy.typing import NDArray
+from sklearn.feature_selection import mutual_info_classif, mutual_info_regression
+
+from dataeval._internal.metrics.utils import entropy, preprocess_metadata
+
+
+class BalanceOutput(NamedTuple):
+    """
+    Attributes
+    ----------
+    mutual_information : NDArray[np.float64]
+        Estimate of mutual information between metadata factors and class label
+    """
+
+    mutual_information: NDArray[np.float64]
+
+
+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."
+            "num_neighbors should be an int, greater than 0 and less than"
+            "the number of samples in the dataset"
+        )
+    if num_neighbors < 1:
+        raise ValueError(
+            f"Invalid value for {num_neighbors}."
+            "Choose a value greater than 0 and less than number of samples"
+            "in the dataset."
+        )
+    if isinstance(num_neighbors, float):
+        num_neighbors = int(num_neighbors)
+        warnings.warn(f"Variable {num_neighbors} is currently type float and will be truncated to type int.")
+
+    return num_neighbors
+
+
+def balance(class_labels: Sequence[int], metadata: List[Dict], num_neighbors: int = 5) -> BalanceOutput:
+    """
+    Mutual information (MI) between factors (class label, metadata, label/image properties)
+
+    Parameters
+    ----------
+    class_labels: Sequence[int]
+        List of class labels for each image
+    metadata: List[Dict]
+        List 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.
+
+    Returns
+    -------
+    BalanceOutput
+        (num_factors+1) x (num_factors+1) estimate of mutual information
+        between num_factors metadata factors and class label. Symmetry is enforced.
+
+    Notes
+    -----
+    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.
+
+    See Also
+    --------
+    sklearn.feature_selection.mutual_info_classif
+    sklearn.feature_selection.mutual_info_regression
+    sklearn.metrics.mutual_info_score
+    """
+    num_neighbors = validate_num_neighbors(num_neighbors)
+    data, names, is_categorical = preprocess_metadata(class_labels, metadata)
+    num_factors = len(names)
+    mi = np.empty((num_factors, num_factors))
+    mi[:] = np.nan
+
+    for idx in range(num_factors):
+        tgt = data[:, idx]
+
+        if is_categorical[idx]:
+            # categorical target
+            mi[idx, :] = mutual_info_classif(
+                data,
+                tgt,
+                discrete_features=is_categorical,  # type: ignore
+                n_neighbors=num_neighbors,
+            )
+        else:
+            # continuous variables
+            mi[idx, :] = mutual_info_regression(
+                data,
+                tgt,
+                discrete_features=is_categorical,  # type: ignore
+                n_neighbors=num_neighbors,
+            )
+
+    ent_all = entropy(data, names, is_categorical, normalized=False)
+    norm_factor = 0.5 * np.add.outer(ent_all, ent_all) + 1e-6
+    # in principle MI should be symmetric, but it is not in practice.
+    nmi = 0.5 * (mi + mi.T) / norm_factor
+
+    return BalanceOutput(nmi)
+
+
+def balance_classwise(class_labels: Sequence[int], metadata: List[Dict], num_neighbors: int = 5) -> BalanceOutput:
+    """
+    Compute mutual information (analogous to correlation) between metadata factors
+    (class label, metadata, label/image properties) with individual class labels.
+
+    Parameters
+    ----------
+    class_labels: Sequence[int]
+        List of class labels for each image
+    metadata: List[Dict]
+        List 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.
+
+    Notes
+    -----
+    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, so we
+    have to specify with is_categorical.
+
+    Returns
+    -------
+    BalanceOutput
+        (num_classes x num_factors) estimate of mutual information between
+        num_factors metadata factors and individual class labels.
+
+    See Also
+    --------
+    sklearn.feature_selection.mutual_info_classif
+    sklearn.feature_selection.mutual_info_regression
+    sklearn.metrics.mutual_info_score
+    compute_mutual_information
+    """
+    num_neighbors = validate_num_neighbors(num_neighbors)
+    data, names, is_categorical = preprocess_metadata(class_labels, metadata)
+    num_factors = len(names)
+    # unique class labels
+    class_idx = names.index("class_label")
+    class_data = data[:, class_idx]
+    u_cls = np.unique(class_data)
+    num_classes = len(u_cls)
+
+    data_no_class = np.concatenate((data[:, :class_idx], data[:, (class_idx + 1) :]), axis=1)
+
+    # assume class is a factor
+    mi = np.empty((num_classes, num_factors - 1))
+    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)
+
+    # classification MI for discrete/categorical features
+    for idx, cls in enumerate(u_cls):
+        tgt = class_data == cls
+        # units: nat
+        mi[idx, :] = mutual_info_classif(
+            data_no_class,
+            tgt,
+            discrete_features=cat_mask,  # type: ignore
+            n_neighbors=num_neighbors,
+        )
+
+    # let this recompute for all features including class label
+    ent_all = entropy(data, names, is_categorical)
+    ent_tgt = ent_all[class_idx]
+    ent_all = np.concatenate((ent_all[:class_idx], ent_all[(class_idx + 1) :]), axis=0)
+    norm_factor = 0.5 * np.add.outer(ent_tgt, ent_all) + 1e-6
+    nmi = mi / norm_factor
+    return BalanceOutput(nmi)