Repository: alteryx/woodwork (python)

Issue:
Documentation - Update Understanding Types and Tags Guide
Update the Understanding Types and Tags guide to use new accessor approach: https://woodwork.alteryx.com/en/stable/guides/understanding_types_and_tags.html

Pull request:
Update types and tags guide to use accessor
- Update types and tags guide to use accessor
- Closes #605

Patch:
diff --git a/docs/source/guides/understanding_types_and_tags.ipynb b/docs/source/guides/understanding_types_and_tags.ipynb
index 52bf3a1..5f2c9a6 100644
--- a/docs/source/guides/understanding_types_and_tags.ipynb
+++ b/docs/source/guides/understanding_types_and_tags.ipynb
@@ -20,7 +20,7 @@
    "source": [
     "## Definitions of Types and Tags\n",
     "\n",
-    "Woodwork has been designed to allow users to easily specify additional information about data contained in a DataTable while providing interfaces to access the underlying data based on that additional information. Because a single DataTable might store various types of data like numbers, text, or dates in different columns, the additional information is defined on a per-column basis.\n",
+    "Woodwork has been designed to allow users to easily specify additional typing information for a DataFrame while providing the ability to interface with the data based on the typing information. Because a single DataFrame might store various types of data like numbers, text, or dates in different columns, the additional information is defined on a per-column basis.\n",
     "\n",
     "There are 3 main ways that Woodwork stores additional information about user data:\n",
     "\n",
@@ -36,7 +36,7 @@
     "### Physical Types\n",
     "Physical types define how the data is stored on disk or in memory. You might also see the physical type for a column referred to as the column’s `dtype`.\n",
     "\n",
-    "For example, typical Pandas dtypes often used include `object`, `int64`, `float64` and `datetime64[ns]`, though there are many more. In Woodwork, there are 7 different physical types that are used, each corresponding to a Pandas dtype. When a DataTable is created, the dtype of the underlying data is converted to one of these values, if it isn't already one of these types:\n",
+    "For example, typical Pandas dtypes often used include `object`, `int64`, `float64` and `datetime64[ns]`, though there are many more. In Woodwork, there are 7 different physical types that are used, each corresponding to a Pandas dtype. When Woodwork is initialized on a DataFrame, the dtype of the underlying data is converted to one of these values, if it isn't already one of these types:\n",
     "\n",
     "* `boolean`\n",
     "* `category`\n",
@@ -46,7 +46,9 @@
     "* `string`\n",
     "* `timedelta64[ns]`\n",
     "\n",
-    "The physical type conversion is done based on the `LogicalType` that has been specified or inferred for a given column."
+    "The physical type conversion is done based on the `LogicalType` that has been specified or inferred for a given column.\n",
+    "\n",
+    "When using Woodwork with a Koalas DataFrame, the physical types used may be different than those listed above. For more information, refer to the guide [Using Woodwork with Dask and Koalas DataFrames](https://woodwork.alteryx.com/en/stable/guides/using_woodwork_with_dask_and_koalas.html#Notes-on-Koalas-Dtype-Conversions)."
    ]
   },
   {
@@ -76,7 +78,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "In the table, notice that each logical type has a specific `pandas_dtype` value associated with it. Any time a logical type is set for a column, the physical type of the underlying data is converted to the type shown in the `pandas_dtype` column. There is only one physical type associated with each logical type."
+    "In the table, notice that each logical type has a specific `physical_type` value associated with it. Any time a logical type is set for a column, the physical type of the underlying data is converted to the type shown in the `physical_type` column. There is only one physical type associated with each logical type."
    ]
   },
   {
@@ -129,11 +131,11 @@
    "source": [
     "## Working with Logical Types\n",
     "\n",
-    "When creating a DataTable, users have the option to specify the logical types for all, some, or none of the columns in the underlying dataframe. If logical types are defined for all of the columns, these logical types are used directly, provided the data is compatible with the specified logical type. You can't, for example, use a logical type of `Integer` on a column that contains text values that can't be converted to integers.\n",
+    "When initializing Woodwork, users have the option to specify the logical types for all, some, or none of the columns in the underlying DataFrame. If logical types are defined for all of the columns, these logical types are used directly, provided the data is compatible with the specified logical type. You can't, for example, use a logical type of `Integer` on a column that contains text values that can't be converted to integers.\n",
     "\n",
-    "If users don't supply any logical type information during creation of the DataTable, Woodwork infers the logical types based on the physical type of the column and the information contained in the columns. If the user passes information for some of the columns, the logical types are inferred for any columns not specified.\n",
+    "If users don't supply any logical type information during initialization, Woodwork infers the logical types based on the physical type of the column and the information contained in the columns. If the user passes information for some of the columns, the logical types are inferred for any columns not specified.\n",
     "\n",
-    "These scenarios are illustrated in this section. To start, create a simple dataframe to use for this example."
+    "These scenarios are illustrated in this section. To start, create a simple DataFrame to use for this example."
    ]
   },
   {
@@ -143,21 +145,21 @@
    "outputs": [],
    "source": [
     "import pandas as pd\n",
-    "from woodwork import DataTable\n",
+    "import woodwork as ww\n",
     "\n",
-    "data = pd.DataFrame({\n",
+    "df = pd.DataFrame({\n",
     "    'integers': [-2, 30, 20],\n",
     "    'bools': [True, False, True],\n",
     "    'names': [\"Jane Doe\", \"Bill Smith\", \"John Hancock\"]\n",
     "})\n",
-    "data"
+    "df"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Now that you've created the data to use for the example, create a DataTable object, assigning logical type values to each of the columns. Then view the types stored for each column by using the `DataTable.types` property."
+    "Importing Woodwork creates a special namespace on the DataFrame, called `ww`, that can be used to initialize and modify Woodwork information for a DataFrame.  Now that you've created the data to use for the example, you can initialize Woodwork on this DataFrame, assigning logical type values to each of the columns. Then view the types stored for each column by using the `DataFrame.ww.types` property."
    ]
   },
   {
@@ -172,8 +174,8 @@
     "    'names': 'FullName'\n",
     "}\n",
     "\n",
-    "dt = DataTable(data, logical_types=logical_types)\n",
-    "dt"
+    "df.ww.init(logical_types=logical_types)\n",
+    "df.ww.types"
    ]
   },
   {
@@ -190,10 +192,10 @@
    "outputs": [],
    "source": [
     "logical_types = {\n",
-    "    'names': 'FullName',\n",
+    "    'names': 'FullName'\n",
     "}\n",
-    "dt = DataTable(data, logical_types=logical_types)\n",
-    "dt"
+    "df.ww.init(logical_types=logical_types)\n",
+    "df.ww"
    ]
   },
   {
@@ -211,8 +213,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "dt = DataTable(data)\n",
-    "dt"
+    "df.ww.init()\n",
+    "df.ww"
    ]
   },
   {
@@ -230,8 +232,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "dt = dt.set_types(logical_types={'names': 'FullName'})\n",
-    "dt"
+    "df.ww.set_types(logical_types={'names': 'FullName'})\n",
+    "df.ww"
    ]
   },
   {
@@ -242,7 +244,7 @@
     "\n",
     "When setting the LogicalType for a column, the type can be specified by passing a string representing the camel-case name of the LogicalType class as you have done in previous examples. Alternatively, you can pass the class directly instead of a string or the snake-case name of the string. All of these would be valid values to use for setting the FullName Logical type: `FullName`, `\"FullName\"` or `\"full_name\"`. \n",
     "\n",
-    "Note—in order to use the class name, firs you have to import the class."
+    "Note—in order to use the class name, first you have to import the class."
    ]
   },
   {
@@ -254,22 +256,12 @@
     "Woodwork provides several methods for working with semantic types. You can add and remove specific tags, or you can reset the tags to their default values. In this section, you learn how to use those methods."
    ]
   },
-  {
-   "cell_type": "raw",
-   "metadata": {
-    "raw_mimetype": "text/restructuredtext"
-   },
-   "source": [
-    ".. note::\n",
-    "    When calling any of the methods that modify the logical type or semantic tags for a column, a new DataTable object is returned and the original DataTable remains unchanged. If you need to use the modified table, assign the returned DataTable to a variable that allows access to it."
-   ]
-  },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "### Standard Tags\n",
-    "As mentioned above,  Woodwork applies default semantic tags to columns by default, based on the logical type that was specified or inferred. If this behavior is undesirable, it can be controlled by setting the parameter `use_standard_tags` to `False` when creating the DataTable."
+    "As mentioned above,  Woodwork applies default semantic tags to columns by default, based on the logical type that was specified or inferred. If this behavior is undesirable, it can be controlled by setting the parameter `use_standard_tags` to `False` when initializing Woodwork."
    ]
   },
   {
@@ -278,17 +270,17 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "dt = DataTable(data, use_standard_tags=False)\n",
-    "dt"
+    "df.ww.init(use_standard_tags=False)\n",
+    "df.ww"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "As can be seen in the table above, when creating a DataTable with the `use_standard_tags` set to `False`, all semantic tags are empty. The only exception to this is if the index or time index column were set. We discuss that in more detail later on.\n",
+    "As can be seen in the output above, when initializing Woodwork with `use_standard_tags` set to `False`, all semantic tags are empty. The only exception to this is if the index or time index column were set. We discuss that in more detail later on.\n",
     "\n",
-    "Create a new table with the standard tags, and specify some additional user-defined semantic tags during creation."
+    "Create a new Woodwork DataFrame with the standard tags, and specify some additional user-defined semantic tags during creation."
    ]
   },
   {
@@ -301,17 +293,17 @@
     "    'bools': 'user_status',\n",
     "    'names': 'legal_name'\n",
     "}\n",
-    "dt = DataTable(data, semantic_tags=semantic_tags)\n",
-    "dt"
+    "df.ww.init(semantic_tags=semantic_tags)\n",
+    "df.ww"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Woodwork has applied the tags we specified along with any standard tags to the columns in our DataTable.\n",
+    "Woodwork has applied the tags we specified along with any standard tags to the columns in our DataFrame.\n",
     "\n",
-    "After creating the table, you have changed your mind and decided you don't like the tag of `user_status` that you applied to the `bools` column. Now you want to remove it. You can do that with the `remove_semantic_tags` method."
+    "After initializing Woodwork, you have changed your mind and decided you don't like the tag of `user_status` that you applied to the `bools` column. Now you want to remove it. You can do that with the `remove_semantic_tags` method."
    ]
   },
   {
@@ -320,8 +312,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "dt = dt.remove_semantic_tags({'bools':'user_status'})\n",
-    "dt"
+    "df.ww.remove_semantic_tags({'bools':'user_status'})\n",
+    "df.ww"
    ]
   },
   {
@@ -339,8 +331,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "dt = dt.add_semantic_tags({'bools':['tag1', 'tag2']})\n",
-    "dt"
+    "df.ww.add_semantic_tags({'bools':['tag1', 'tag2']})\n",
+    "df.ww"
    ]
   },
   {
@@ -349,8 +341,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "dt = dt.remove_semantic_tags({'bools':['tag1', 'tag2']})\n",
-    "dt"
+    "df.ww.remove_semantic_tags({'bools':['tag1', 'tag2']})\n",
+    "df.ww"
    ]
   },
   {
@@ -366,15 +358,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "dt = dt.reset_semantic_tags()\n",
-    "dt"
+    "df.ww.reset_semantic_tags()\n",
+    "df.ww"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "In this case, since you created our DataTable with the default behavior of using standard tags, calling `reset_semantic_tags` resulted in all of our semantic tags being reset to the standard tags for each column."
+    "In this case, since you initialized Woodwork with the default behavior of using standard tags, calling `reset_semantic_tags` resulted in all of our semantic tags being reset to the standard tags for each column."
    ]
   },
   {
@@ -383,13 +375,13 @@
    "source": [
     "### Index and Time Index Tags\n",
     "\n",
-    "When creating a DataTable, you have the option to specify which column represents the index and which column represents the time index. If these columns are specified, semantic tags of `index` and `time_index` are applied to the specified columns. Behind the scenes, Woodwork is performing additional validation checks on the columns to make sure they are appropriate. For example, index columns must be unique, and time index columns must contain datetime values or values that can be converted to datetimes.\n",
+    "When initializing Woodwork, you have the option to specify which column represents the index and which column represents the time index. If these columns are specified, semantic tags of `index` and `time_index` are applied to the specified columns. Behind the scenes, Woodwork is performing additional validation checks on the columns to make sure they are appropriate. For example, index columns must be unique, and time index columns must contain datetime values or numeric values.\n",
     "\n",
-    "Because of the need for these validation checks, you can't set the `index` or `time_index` tags directly on a column. In order to designate a column as the index, the `set_index` method should be used. Similarly, in order to set the time index column, the `set_time_index` method should be used. Optionally, these can be specified when initially creating the DataTable.\n",
+    "Because of the need for these validation checks, you can't set the `index` or `time_index` tags directly on a column. In order to designate a column as the index, the `set_index` method should be used. Similarly, in order to set the time index column, the `set_time_index` method should be used. Optionally, these can be specified when initializing Woodwork by using the `index` or `time_index` parameters.\n",
     "\n",
-    "#### Setting or changing the index\n",
+    "#### Setting the index\n",
     "\n",
-    "Create a new sample data set that contains columns that can be used as index and time index columns by using the `index` or `time_index` parameters."
+    "Create a new sample DataFrame that contains columns that can be used as index and time index columns and initialize Woodwork."
    ]
   },
   {
@@ -398,22 +390,22 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "data = pd.DataFrame({\n",
+    "df = pd.DataFrame({\n",
     "    'index': [0, 1, 2],\n",
     "    'id': [1, 2, 3],\n",
     "    'times': pd.to_datetime(['2020-09-01', '2020-09-02', '2020-09-03']),\n",
     "    'numbers': [10, 20, 30]\n",
     "})\n",
     "\n",
-    "dt = DataTable(data)\n",
-    "dt"
+    "df.ww.init()\n",
+    "df.ww"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Without specifying an index or time index column during creation of the DataTable, Woodwork has inferred that the `index` and `id` columns are integers and the numeric semantic tag has been applied. You can now set the index column with the `set_index` method."
+    "Without specifying an index or time index column during initialization, Woodwork has inferred that the `index` and `id` columns are integers and the numeric semantic tag has been applied. You can now set the index column with the `set_index` method."
    ]
   },
   {
@@ -422,15 +414,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "dt = dt.set_index('index')\n",
-    "dt"
+    "df.ww.set_index('index')\n",
+    "df.ww"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Inspecting the types now reveals that the `index` semantic tag has been added to the `index` column, and the `numeric` standard tag has been removed. You can also check that the index has been set correctly by checking the value of the `DataTable.index` attribute."
+    "Inspecting the types now reveals that the `index` semantic tag has been added to the `index` column, and the `numeric` standard tag has been removed. You can also check that the index has been set correctly by checking the value of the `DataFrame.ww.index` attribute."
    ]
   },
   {
@@ -439,7 +431,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "dt.index"
+    "df.ww.index"
    ]
   },
   {
@@ -455,19 +447,19 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "dt = dt.set_index('id')\n",
-    "dt"
+    "df.ww.set_index('id')\n",
+    "df.ww"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "The `index` tag has been removed from the `index` column and added to the `id` column. The `numeric` standard tag that was originally present on the `index` column has not been added back. If this tag is desired, the user must currently add it back using the `add_semantic_tags` method.\n",
+    "The `index` tag has been removed from the `index` column and added to the `id` column. The `numeric` standard tag that was originally present on the `index` column has been added back.\n",
     "\n",
-    "#### Setting or changing the time index\n",
+    "#### Setting the time index\n",
     "\n",
-    "Setting the time index works similarly to setting the index. You can now set the time index for the DataTable with the `set_time_index` method."
+    "Setting the time index works similarly to setting the index. You can now set the time index with the `set_time_index` method."
    ]
   },
   {
@@ -476,15 +468,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "dt = dt.set_time_index('times')\n",
-    "dt"
+    "df.ww.set_time_index('times')\n",
+    "df.ww"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "After calling `set_time_index`, the `time_index` semantic tag has been added to the `times` column of the DataTable."
+    "After calling `set_time_index`, the `time_index` semantic tag has been added to the semantic tags for `times` column."
    ]
   }
  ],
@@ -505,7 +497,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.7.4"
+   "version": "3.8.5"
   }
  },
  "nbformat": 4,
diff --git a/docs/source/release_notes.rst b/docs/source/release_notes.rst
index 4fdb7b2..061781d 100644
--- a/docs/source/release_notes.rst
+++ b/docs/source/release_notes.rst
@@ -54,8 +54,10 @@ Release Notes
         * Clean up dtype usage across Woodwork (:pr:`682`)
         * Improve error when calling accessor properties or methods before init (:pr:`683`)
         * Remove dtype from Schema dictionary (:pr:`685`)
+        * Add ``include_index`` param and allow unique columns in Accessor mutual information (:pr:`699`)
     * Documentation Changes
         * Update README.md and Get Started guide to use accessor (:pr:`655`)
+        * Update Understanding Types and Tags guide to use accessor (:pr:`657`)
         * Update docstrings and API Reference page (:pr:`660`)
     * Testing Changes
         * Update branch reference in tests to run on main (:pr:`641`)
diff --git a/woodwork/statistics_utils.py b/woodwork/statistics_utils.py
index 81a12e0..38b7e76 100644
--- a/woodwork/statistics_utils.py
+++ b/woodwork/statistics_utils.py
@@ -164,7 +164,7 @@ def _make_categorical_for_mutual_info(schema, data, num_bins):
     return data
 
 
-def _get_mutual_information_dict(dataframe, num_bins=10, nrows=None):
+def _get_mutual_information_dict(dataframe, num_bins=10, nrows=None, include_index=False):
     """Calculates mutual information between all pairs of columns in the DataFrame that
     support mutual information. Logical Types that support mutual information are as
     follows:  Boolean, Categorical, CountryCode, Datetime, Double, Integer, Ordinal,
@@ -178,6 +178,10 @@ def _get_mutual_information_dict(dataframe, num_bins=10, nrows=None):
         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`,
@@ -186,8 +190,10 @@ def _get_mutual_information_dict(dataframe, num_bins=10, nrows=None):
         (perfect dependency).
         """
     valid_types = get_valid_mi_types()
-    valid_columns = [col_name for col_name, col in dataframe.ww.columns.items() if (
-        col_name != dataframe.ww.index and _get_ltype_class(col['logical_type']) in valid_types)]
+    valid_columns = [col_name for col_name, col in dataframe.ww.columns.items() if _get_ltype_class(col['logical_type']) in valid_types]
+
+    if not include_index and dataframe.ww.index is not None:
+        valid_columns.remove(dataframe.ww.index)
 
     data = dataframe[valid_columns]
     if dd and isinstance(data, dd.DataFrame):
@@ -203,10 +209,6 @@ def _get_mutual_information_dict(dataframe, num_bins=10, nrows=None):
     not_null_cols = data.columns[data.notnull().any()]
     if set(not_null_cols) != set(valid_columns):
         data = data.loc[:, not_null_cols]
-    # remove columns that are unique
-    not_unique_cols = [col for col in data.columns if not data[col].is_unique]
-    if set(not_unique_cols) != set(valid_columns):
-        data = data.loc[:, not_unique_cols]
 
     data = _replace_nans_for_mutual_info(dataframe.ww.schema, data)
     data = _make_categorical_for_mutual_info(dataframe.ww.schema, data, num_bins)
diff --git a/woodwork/table_accessor.py b/woodwork/table_accessor.py
index 43dce58..a213f96 100644
--- a/woodwork/table_accessor.py
+++ b/woodwork/table_accessor.py
@@ -667,7 +667,7 @@ class WoodworkTableAccessor:
         new_df.ww.init(schema=new_schema)
         return new_df
 
-    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 DataFrame that
         support mutual information. Logical Types that support mutual information are as
@@ -678,8 +678,12 @@ class WoodworkTableAccessor:
             num_bins (int): Determines number of bins to use for converting
                 numeric features into categorical.
             nrows (int): The number of rows to sample for when determining mutual info.
-            If specified, samples the desired number of rows from the data.
+                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`,
@@ -689,9 +693,9 @@ class WoodworkTableAccessor:
         """
         if self._schema is None:
             _raise_init_error()
-        return _get_mutual_information_dict(self._dataframe, num_bins=num_bins, nrows=nrows)
+        return _get_mutual_information_dict(self._dataframe, num_bins=num_bins, nrows=nrows, include_index=include_index)
 
-    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 DataFrame that
         support mutual information. Logical Types that support mutual information are as
         follows:  Boolean, Categorical, CountryCode, Datetime, Double, Integer, Ordinal,
@@ -703,6 +707,10 @@ class WoodworkTableAccessor:
             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`,
@@ -710,7 +718,7 @@ class WoodworkTableAccessor:
             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 describe_dict(self, include=None):


Test patch:
diff --git a/woodwork/tests/accessor/test_statistics.py b/woodwork/tests/accessor/test_statistics.py
index c6d2132..b35938e 100644
--- a/woodwork/tests/accessor/test_statistics.py
+++ b/woodwork/tests/accessor/test_statistics.py
@@ -127,7 +127,8 @@ def test_mutual_info(df_mi):
     pd.testing.assert_frame_equal(mi, mi_many_rows)
 
     mi = df_mi.ww.mutual_information(nrows=1)
-    assert mi.shape[0] == 0
+    assert mi.shape[0] == 10
+    assert (mi['mutual_info'] == 1.0).all()
 
     mi = df_mi.ww.mutual_information(num_bins=2)
     assert mi.shape[0] == 10
@@ -141,11 +142,14 @@ def test_mutual_info(df_mi):
     pd.testing.assert_frame_equal(to_pandas(df_mi), to_pandas(original_df))
 
 
-def test_mutual_info_does_not_include_index(sample_df):
+def test_mutual_info_on_index(sample_df):
     sample_df.ww.init(index='id')
     mi = sample_df.ww.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 = sample_df.ww.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):
@@ -177,8 +181,8 @@ def test_mutual_info_unique_cols(df_mi_unique):
     mi = df_mi_unique.ww.mutual_information()
 
     cols_used = set(np.unique(mi[['column_1', 'column_2']].values))
-    assert 'unique' not in cols_used
-    assert 'unique_with_one_nan' not in cols_used
+    assert 'unique' in cols_used
+    assert 'unique_with_one_nan' in cols_used
     assert 'unique_with_nans' in cols_used
     assert 'ints' in cols_used