Repository: alteryx/woodwork (python)

Issue:
Add avoid_index parameter to mutual info to allow for more control
- We should add the following parameters to the mutual info function and associated defaults.  
    -  `avoid_index=False`
- This way the end-user can have more control of the logic for calculating MI  
- This is for the current DataTable implementation

Pull request:
Add include_index parameter to DataTable mutual info
- Allows users to include the index column in mutual info calculations, where the default is to not include the index column
- Closes #671 

Patch:
diff --git a/docs/source/release_notes.rst b/docs/source/release_notes.rst
index 62d15bf..8afe971 100644
--- a/docs/source/release_notes.rst
+++ b/docs/source/release_notes.rst
@@ -8,6 +8,7 @@ Release Notes
     * Changes
         * Restrict Koalas version to ``<1.7.0`` due to breaking changes (:pr:`674`)
         * Include unique columns in mutual information calculations (:pr:`687`) 
+        * Add parameter to include index column in mutual information calculations (:pr:`692`) 
     * Documentation Changes
         * Update to remove warning message from statistical insights guide (:pr:`690`)
     * Testing Changes
diff --git a/woodwork/datatable.py b/woodwork/datatable.py
index 989a60a..89dab4e 100644
--- a/woodwork/datatable.py
+++ b/woodwork/datatable.py
@@ -1000,7 +1000,7 @@ class DataTable(object):
             data[col_name] = new_col.cat.codes
         return data
 
-    def mutual_information_dict(self, num_bins=10, nrows=None):
+    def mutual_information_dict(self, num_bins=10, nrows=None, include_index=False):
         """
         Calculates mutual information between all pairs of columns in the DataTable that
         support mutual information. Logical Types that support mutual information are as
@@ -1013,7 +1013,10 @@ class DataTable(object):
             nrows (int): The number of rows to sample for when determining mutual info.
                 If specified, samples the desired number of rows from the data.
                 Defaults to using all rows.
-
+            include_index (bool): If True, the column specified as the index will be
+                included as long as its LogicalType is valid for mutual information calculations.
+                If False, the index column will not have mutual information calculated for it.
+                Defaults to False.
         Returns:
             list(dict): A list containing dictionaries that have keys `column_1`,
             `column_2`, and `mutual_info` that is sorted in decending order by mutual info.
@@ -1021,8 +1024,11 @@ class DataTable(object):
             (perfect dependency).
         """
         valid_types = get_valid_mi_types()
-        valid_columns = [col.name for col in self.columns.values() if (
-            col.name != self.index and _get_ltype_class(col.logical_type) in valid_types)]
+
+        valid_columns = [col.name for col in self.columns.values() if _get_ltype_class(col.logical_type) in valid_types]
+
+        if not include_index and self.index is not None:
+            valid_columns.remove(self.index)
 
         data = self._dataframe.loc[:, valid_columns]
         if dd and isinstance(data, dd.DataFrame):
@@ -1059,7 +1065,7 @@ class DataTable(object):
         mutual_info.sort(key=lambda mi: mi['mutual_info'], reverse=True)
         return mutual_info
 
-    def mutual_information(self, num_bins=10, nrows=None):
+    def mutual_information(self, num_bins=10, nrows=None, include_index=False):
         """
         Calculates mutual information between all pairs of columns in the DataTable that
         support mutual information. Logical Types that support mutual information are as
@@ -1072,6 +1078,10 @@ class DataTable(object):
             nrows (int): The number of rows to sample for when determining mutual info.
                 If specified, samples the desired number of rows from the data.
                 Defaults to using all rows.
+            include_index (bool): If True, the column specified as the index will be
+                included as long as its LogicalType is valid for mutual information calculations.
+                If False, the index column will not have mutual information calculated for it.
+                Defaults to False.
 
         Returns:
             pd.DataFrame: A Dataframe containing mutual information with columns `column_1`,
@@ -1079,7 +1089,7 @@ class DataTable(object):
             Mutual information values are between 0 (no mutual information) and 1
             (perfect dependency).
         """
-        mutual_info = self.mutual_information_dict(num_bins, nrows)
+        mutual_info = self.mutual_information_dict(num_bins, nrows, include_index)
         return pd.DataFrame(mutual_info)
 
     def to_dictionary(self):


Test patch:
diff --git a/woodwork/tests/datatable/test_datatable.py b/woodwork/tests/datatable/test_datatable.py
index e2d8765..d7be162 100644
--- a/woodwork/tests/datatable/test_datatable.py
+++ b/woodwork/tests/datatable/test_datatable.py
@@ -974,10 +974,14 @@ def test_datatable_mutual_information(df_mi):
     pd.testing.assert_frame_equal(to_pandas(dt.to_dataframe()), to_pandas(original_df))
 
 
-def test_mutual_info_does_not_include_index(sample_df):
+def test_mutual_info_on_index(sample_df):
+
     dt = DataTable(sample_df, index='id')
     mi = dt.mutual_information()
-    assert 'id' not in mi['column_1'].values
+    assert not ('id' in mi['column_1'].values or 'id' in mi['column_2'].values)
+
+    mi = dt.mutual_information(include_index=True)
+    assert 'id' in mi['column_1'].values or 'id' in mi['column_2'].values
 
 
 def test_mutual_info_returns_empty_df_properly(sample_df):