Repository: andersbogsnes/ml_tooling

Issue:
ENH: Plot for comparis of means of categorical targets
Being able to plot categorical targets in terms of the mean/median value of the features is a good way to see the correlation between the variables

**Describe the solution you'd like**
Add a `plot.target_feature_distribution`

**Additional context**
This would be a plot only relevant if y is categorical - do we need to "help" distinguish?


Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/andersbogsnes/ml_tooling/pull/636?src=pr&el=h1) Report
> Merging [#636](https://codecov.io/gh/andersbogsnes/ml_tooling/pull/636?src=pr&el=desc) into [develop](https://codecov.io/gh/andersbogsnes/ml_tooling/commit/bc0191ee781e539c9345835be6ffb880f8b1b818&el=desc) will **decrease** coverage by `1.51%`.
> The diff coverage is `0.00%`.

[![Impacted file tree graph](https://codecov.io/gh/andersbogsnes/ml_tooling/pull/636/graphs/tree.svg?width=650&height=150&src=pr&token=ucKtL3BlxO)](https://codecov.io/gh/andersbogsnes/ml_tooling/pull/636?src=pr&el=tree)

```diff
@@             Coverage Diff             @@
##           develop     #636      +/-   ##
===========================================
- Coverage   100.00%   98.48%   -1.52%     
===========================================
  Files           59       60       +1     
  Lines         1556     1580      +24     
  Branches       175      176       +1     
===========================================
  Hits          1556     1556              
- Misses           0       24      +24     
```


| [Impacted Files](https://codecov.io/gh/andersbogsnes/ml_tooling/pull/636?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [src/ml\_tooling/plots/mean\_categorical\_target.py](https://codecov.io/gh/andersbogsnes/ml_tooling/pull/636/diff?src=pr&el=tree#diff-c3JjL21sX3Rvb2xpbmcvcGxvdHMvbWVhbl9jYXRlZ29yaWNhbF90YXJnZXQucHk=) | `0.00% <0.00%> (ø)` | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/andersbogsnes/ml_tooling/pull/636?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/andersbogsnes/ml_tooling/pull/636?src=pr&el=footer). Last update [bc0191e...214d240](https://codecov.io/gh/andersbogsnes/ml_tooling/pull/636?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9954e31..8310c3d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -14,6 +14,7 @@
 - Fixed FuncTransform to have a default value
 - `load_estimator` now uses default storage if nothing is passed
 - `Model.bayessearc` is now `Model.bayesiansearch`
+- Added `target_feature_distribution`to `Dataset.plot`
 
 # v0.11.0
 - Added `load_demo_dataset` function
diff --git a/src/ml_tooling/plots/__init__.py b/src/ml_tooling/plots/__init__.py
index 835261f..7f4bd93 100644
--- a/src/ml_tooling/plots/__init__.py
+++ b/src/ml_tooling/plots/__init__.py
@@ -9,6 +9,7 @@ from .precision_recall_curve import plot_pr_curve
 from .learning_curve import plot_learning_curve
 from .validation_curve import plot_validation_curve
 from .missing_data import plot_missing_data
+from .target_feature_distribution import plot_target_feature_distribution
 
 __all__ = [
     "plot_confusion_matrix",
@@ -22,4 +23,5 @@ __all__ = [
     "plot_learning_curve",
     "plot_validation_curve",
     "plot_missing_data",
+    "plot_target_feature_distribution",
 ]
diff --git a/src/ml_tooling/plots/target_feature_distribution.py b/src/ml_tooling/plots/target_feature_distribution.py
new file mode 100644
index 0000000..28592ed
--- /dev/null
+++ b/src/ml_tooling/plots/target_feature_distribution.py
@@ -0,0 +1,63 @@
+from matplotlib import pyplot as plt
+from matplotlib.axes import Axes
+import numpy as np
+
+from ml_tooling.utils import DataType
+from ml_tooling.plots.utils import _plot_barh
+
+
+def plot_target_feature_distribution(
+    target: DataType,
+    feature: DataType,
+    title: str = "Target feature distribution",
+    method: str = "mean",
+    ax: plt.Axes = None,
+) -> Axes:
+    """
+    Creates a plot which compares the mean or median
+    of a binary target based on the given category features.
+    NaN values are ignored in the calculation.
+
+    Parameters
+    ----------
+    target: DataType
+        Target to aggregate per feature category
+    feature: DataType
+        Categorical feature to group by
+    title: str
+        Title of graph
+    method: str
+        Which method to compare with. One of 'median' or 'mean'.
+    ax: plt.Axes
+        Matplotlib axes to draw the graph on. Creates a new one by default
+    Returns
+    -------
+    plt.Axes
+
+    """
+    if ax is None:
+        fig, ax = plt.subplots()
+
+    agg_func_mapping = {"mean": np.nanmean, "median": np.nanmedian}
+
+    agg_func = agg_func_mapping[method]
+
+    feature_categories = np.sort(np.unique(feature))
+
+    data = np.asarray(
+        [agg_func(target[feature == category]) for category in feature_categories]
+    )
+
+    ax = _plot_barh(
+        values=data,
+        label_names=feature_categories,
+        add_label=True,
+        title=title,
+        x_label=f"{method} target",
+        y_label="Feature categories",
+        ax=ax,
+    )
+
+    ax.axvline(x=agg_func(target), linestyle="--")
+
+    return ax
diff --git a/src/ml_tooling/plots/viz/data_viz.py b/src/ml_tooling/plots/viz/data_viz.py
index 4ccac9c..3a7838f 100644
--- a/src/ml_tooling/plots/viz/data_viz.py
+++ b/src/ml_tooling/plots/viz/data_viz.py
@@ -4,8 +4,13 @@ import matplotlib.pyplot as plt
 from matplotlib.axes import Axes
 from sklearn.pipeline import Pipeline
 
+
+from ml_tooling.plots import (
+    plot_target_correlation,
+    plot_missing_data,
+    plot_target_feature_distribution,
+)
 from ml_tooling.config import MPL_STYLESHEET
-from ml_tooling.plots import plot_target_correlation, plot_missing_data
 
 
 class DataVisualize:
@@ -108,3 +113,39 @@ class DataVisualize:
 
         with plt.style.context(MPL_STYLESHEET):
             return plot_missing_data(df=x, ax=ax, top_n=top_n, bottom_n=bottom_n)
+
+    def target_feature_distribution(
+        self,
+        feature_name: str,
+        method: str = "mean",
+        ax: Optional[Axes] = None,
+        feature_pipeline: Optional[Pipeline] = None,
+    ) -> Axes:
+        """
+        Creates a plot which compares the mean or median
+        of a binary target based on the given category feature.
+
+        Parameters
+        ----------
+        feature_name: str
+            Which feature showcase in plot.
+        method: str
+            Which method to compare with. One of 'median' or 'mean'.
+        ax: plt.Axes
+            Matplotlib axes to draw the graph on. Creates a new one by default
+        feature_pipeline: Pipeline
+            A feature transformation pipeline to be applied before graphing the data
+        Returns
+        -------
+        plt.Axes
+        """
+
+        x = self.data.x[feature_name]
+
+        if feature_pipeline is not None:
+            x = feature_pipeline.fit_transform(x)
+
+        with plt.style.context(MPL_STYLESHEET):
+            return plot_target_feature_distribution(
+                feature=x, target=self.data.y, method=method, ax=ax,
+            )


Test patch:
diff --git a/tests/test_visualizations/test_target_feature_distribution.py b/tests/test_visualizations/test_target_feature_distribution.py
new file mode 100644
index 0000000..85f847d
--- /dev/null
+++ b/tests/test_visualizations/test_target_feature_distribution.py
@@ -0,0 +1,80 @@
+from unittest.mock import MagicMock
+import pytest
+
+from matplotlib import pyplot as plt
+import pandas as pd
+from sklearn.pipeline import Pipeline
+from sklearn.datasets import fetch_openml
+
+from ml_tooling.data import Dataset
+
+
+class TestTargetFeatureDistribution:
+    @pytest.fixture()
+    def australian_dataset(self):
+        australian_data = fetch_openml("Australian", as_frame=True)
+        australian_data_df = pd.DataFrame(
+            data=australian_data.data, columns=australian_data.feature_names
+        )
+        australian_data_df_target = australian_data.target.astype("int")
+
+        class AustralianData(Dataset):
+            def load_training_data(self):
+                return australian_data_df, australian_data_df_target
+
+            def load_prediction_data(self):
+                return australian_data_df
+
+        return AustralianData()
+
+    def test_target_correlation_uses_pipeline_when_passed(
+        self, train_iris_dataset: Dataset
+    ):
+        mock = MagicMock(spec=Pipeline)
+        mock.fit_transform.return_value = train_iris_dataset.x["petal length (cm)"]
+
+        train_iris_dataset.plot.target_feature_distribution(
+            feature_name="petal length (cm)", feature_pipeline=mock
+        )
+        mock.fit_transform.assert_called_once_with(
+            train_iris_dataset.x["petal length (cm)"]
+        )
+        plt.close()
+
+    def test_target_feature_distribution_works_as_expected(
+        self, australian_dataset: Dataset
+    ):
+
+        ax = australian_dataset.plot.target_feature_distribution(feature_name="A4")
+
+        assert [text.get_text() for text in ax.texts] == ["0.28", "0.50", "1.00"]
+
+        assert ax.title.get_text() == "Target feature distribution"
+        assert ax.get_xlabel() == "mean target"
+        assert ax.get_ylabel() == "Feature categories"
+        plt.close()
+
+    def test_target_feature_distribution_works_with_different_methods(
+        self, australian_dataset: Dataset
+    ):
+        ax = australian_dataset.plot.target_feature_distribution(
+            feature_name="A4", method="median"
+        )
+
+        assert [text.get_text() for text in ax.texts] == ["0.00", "0.00", "1.00"]
+
+        assert ax.title.get_text() == "Target feature distribution"
+        assert ax.get_xlabel() == "median target"
+        assert ax.get_ylabel() == "Feature categories"
+        plt.close()
+
+    def test_target_feature_distribution_plots_can_be_given_an_ax(
+        self, australian_dataset: Dataset
+    ):
+        fig, ax = plt.subplots()
+
+        test_ax = australian_dataset.plot.target_feature_distribution(
+            feature_name="A4", ax=ax
+        )
+        assert ax == test_ax
+        plt.close()