Repository: pyjanitor-devs/pyjanitor

Issue:
explode_levels
# Brief Description

I would like to propose an `explode_levels` function, as an inverse to `collapse_levels` - essentially converting a single column index to a MultiIndex. 

This offers a chaining alternative to `df.columns = df.columns.str.split(names_sep, expand=True)`

# Example API

```python
df = {'max_speed_mean': [267.3333333333333, 50.5],
 'max_speed_median': [389.0, 50.5]}
df = pd.DataFrame(df)
df
   max_speed_mean  max_speed_median
0      267.333333             389.0
1       50.500000              50.5

        max       
        speed       
         mean median
0  267.333333  389.0
1   50.500000   50.5

```



Discussion:

Resolving patch:
diff --git a/janitor/functions/__init__.py b/janitor/functions/__init__.py
--- a/janitor/functions/__init__.py
+++ b/janitor/functions/__init__.py
@@ -41,6 +41,7 @@
 from .encode_categorical import encode_categorical
 from .expand_column import expand_column
 from .expand_grid import expand_grid
+from .explode_index import explode_index
 from .factorize_columns import factorize_columns
 from .fill import fill_direction, fill_empty
 from .filter import filter_column_isin, filter_date, filter_on, filter_string
@@ -108,6 +109,7 @@
     "encode_categorical",
     "expand_column",
     "expand_grid",
+    "explode_index",
     "factorize_columns",
     "fill_direction",
     "fill_empty",
diff --git a/janitor/functions/collapse_levels.py b/janitor/functions/collapse_levels.py
--- a/janitor/functions/collapse_levels.py
+++ b/janitor/functions/collapse_levels.py
@@ -39,7 +39,7 @@ class  max_speed    type
         2    bird         24  parrot
         3  mammal         80    Lion
         4  mammal         21  Monkey
-        >>> grouped_df = df.groupby("class").agg(["mean", "median"])
+        >>> grouped_df = df.groupby("class")[['max_speed']].agg(["mean", "median"])
         >>> grouped_df  # doctest: +NORMALIZE_WHITESPACE
                  max_speed
                       mean median
diff --git a/janitor/functions/explode_index.py b/janitor/functions/explode_index.py
new file mode 100644
--- /dev/null
+++ b/janitor/functions/explode_index.py
@@ -0,0 +1,112 @@
+"""Implementation of the `explode_index` function."""
+
+from __future__ import annotations
+
+import re
+from typing import Union
+
+import pandas as pd
+import pandas_flavor as pf
+
+from janitor.utils import check
+
+
+@pf.register_dataframe_method
+def explode_index(
+    df: pd.DataFrame,
+    names_sep: Union[str, None] = None,
+    names_pattern: Union[str, None] = None,
+    axis: str = "columns",
+    level_names: list = None,
+) -> pd.DataFrame:
+    """Explode a single index DataFrame into a MultiIndex DataFrame.
+
+    This method does not mutate the original DataFrame.
+
+    Examples:
+        >>> import pandas as pd
+        >>> import janitor
+        >>> df = pd.DataFrame(
+        ...          {'max_speed_mean': [267.3333333333333, 50.5],
+        ...           'max_speed_median': [389.0, 50.5]})
+        >>> df
+           max_speed_mean  max_speed_median
+        0      267.333333             389.0
+        1       50.500000              50.5
+        >>> df.explode_index(names_sep='_',axis='columns')  # doctest: +NORMALIZE_WHITESPACE
+                  max
+                speed
+                 mean median
+        0  267.333333  389.0
+        1   50.500000   50.5
+        >>> df.explode_index(names_pattern=r"(.+speed)_(.+)",axis='columns') # doctest: +NORMALIZE_WHITESPACE
+            max_speed
+                 mean median
+        0  267.333333  389.0
+        1   50.500000   50.5
+        >>> df.explode_index(
+        ...     names_pattern=r"(?P<measurement>.+speed)_(?P<aggregation>.+)",
+        ...     axis='columns'
+        ... ) # doctest: +NORMALIZE_WHITESPACE
+        measurement   max_speed
+        aggregation        mean median
+        0            267.333333  389.0
+        1             50.500000   50.5
+        >>> df.explode_index(
+        ...     names_sep='_',
+        ...     axis='columns',
+        ...     level_names = ['min or max', 'measurement','aggregation']
+        ... ) # doctest: +NORMALIZE_WHITESPACE
+        min or max          max
+        measurement       speed
+        aggregation        mean median
+        0            267.333333  389.0
+        1             50.500000   50.5
+
+    Args:
+        df: A pandas DataFrame.
+        names_sep: string or compiled regex used to split the column/index into levels.
+        names_pattern: regex to extract new levels from the column/index.
+        axis: 'index/columns'. Determines which axis to explode.
+        level_names: names of the levels in the MultiIndex.
+
+    Returns:
+        A pandas DataFrame with a MultiIndex.
+    """  # noqa: E501
+    check("axis", axis, [str])
+    if axis not in {"index", "columns"}:
+        raise ValueError("axis should be either index or columns.")
+    if (names_sep is None) and (names_pattern is None):
+        raise ValueError(
+            "Provide argument for either names_sep or names_pattern."
+        )
+    if (names_sep is not None) and (names_pattern is not None):
+        raise ValueError(
+            "Provide argument for either names_sep or names_pattern, not both."
+        )
+    if names_sep is not None:
+        check("names_sep", names_sep, [str])
+    if names_pattern is not None:
+        check("names_pattern", names_pattern, [str])
+    if level_names is not None:
+        check("level_names", level_names, [list])
+
+    new_index = getattr(df, axis)
+    if isinstance(new_index, pd.MultiIndex):
+        return df
+    # avoid a copy - Index is immutable; a slice is safe to use.
+    df = df[:]
+    if names_sep:
+        new_index = new_index.str.split(names_sep, expand=True)
+    else:
+        named_groups = re.compile(names_pattern).groupindex
+        if named_groups and not level_names:
+            level_names = list(named_groups)
+        new_index = new_index.str.extract(names_pattern)
+        new_index = [arr.array for _, arr in new_index.items()]
+        new_index = pd.MultiIndex.from_arrays(new_index)
+    if level_names:
+        new_index.names = level_names
+
+    setattr(df, axis, new_index)
+    return df


Test patch:
diff --git a/tests/functions/test_explode_index.py b/tests/functions/test_explode_index.py
new file mode 100644
--- /dev/null
+++ b/tests/functions/test_explode_index.py
@@ -0,0 +1,104 @@
+import pandas as pd
+import pytest
+from pandas.testing import assert_frame_equal
+
+
+@pytest.fixture
+def df_checks():
+    """fixture dataframe"""
+    return pd.DataFrame(
+        {
+            "fam_id": [1, 1, 1, 2, 2, 2, 3, 3, 3],
+        }
+    )
+
+
+@pytest.fixture
+def df_multi():
+    """MultiIndex dataframe fixture."""
+    return pd.DataFrame(
+        {
+            ("name", "a"): {0: "Wilbur", 1: "Petunia", 2: "Gregory"},
+            ("names", "aa"): {0: 67, 1: 80, 2: 64},
+            ("more_names", "aaa"): {0: 56, 1: 90, 2: 50},
+        }
+    )
+
+
+def test_type_axis(df_checks):
+    """Raise TypeError if wrong type is provided for axis."""
+    with pytest.raises(TypeError, match="axis should be one of.+"):
+        df_checks.explode_index(axis=1)
+
+
+def test_axis_values(df_checks):
+    """Raise ValueError if wrong value is provided for axis."""
+    msg = "axis should be either index or columns."
+    with pytest.raises(ValueError, match=msg):
+        df_checks.explode_index(axis="INDEX")
+
+
+def test_names_sep_pattern(df_checks):
+    """Raise ValueError if both names_sep and names_pattern is provided."""
+    msg = "Provide argument for either names_sep or names_pattern, not both."
+    with pytest.raises(ValueError, match=msg):
+        df_checks.explode_index(
+            axis="columns", names_sep="_", names_pattern=r"(.+)_(.+)"
+        )
+
+
+def test_names_sep_pattern_both_none(df_checks):
+    """Raise ValueError if neither names_sep nor names_pattern is provided."""
+    msg = "Provide argument for either names_sep or names_pattern."
+    with pytest.raises(ValueError, match=msg):
+        df_checks.explode_index(
+            axis="columns", names_sep=None, names_pattern=None
+        )
+
+
+def test_names_sep_typeerror(df_checks):
+    """Raise TypeError if names_sep is a wrong type."""
+    with pytest.raises(TypeError, match="names_sep should be one of.+"):
+        df_checks.explode_index(axis="columns", names_sep=1)
+
+
+def test_names_pattern_typeerror(df_checks):
+    """Raise TypeError if names_pattern is a wrong type."""
+    with pytest.raises(TypeError, match="names_pattern should be one of.+"):
+        df_checks.explode_index(names_pattern=1)
+
+
+def test_level_names_typeerror(df_checks):
+    """Raise TypeError if level_names is a wrong type."""
+    with pytest.raises(TypeError, match="level_names should be one of.+"):
+        df_checks.explode_index(names_sep="_", level_names="new_level")
+
+
+def test_multiindex(df_multi):
+    """Test output if df.columns is a multiindex"""
+    actual = df_multi.explode_index(names_sep="_")
+    assert_frame_equal(df_multi, actual)
+
+
+def test_names_sep(df_checks):
+    """test output if names_sep"""
+    actual = df_checks.explode_index(names_sep="_", level_names=["a", "b"])
+    expected = pd.DataFrame(
+        {
+            ("fam", "id"): [1, 1, 1, 2, 2, 2, 3, 3, 3],
+        }
+    )
+    expected.columns.names = ["a", "b"]
+    assert_frame_equal(actual, expected)
+
+
+def test_names_pattern(df_checks):
+    """test output if names_pattern"""
+    actual = df_checks.explode_index(names_pattern=r"(?P<a>.+)_(?P<b>.+)")
+    expected = pd.DataFrame(
+        {
+            ("fam", "id"): [1, 1, 1, 2, 2, 2, 3, 3, 3],
+        }
+    )
+    expected.columns.names = ["a", "b"]
+    assert_frame_equal(actual, expected)