Repository: wsp-sag/cheval

Issue:
LinkedDataFrame linkage indexes do not update on groupby
### Overview
GGHM demand model code makes extensive use of `LinkedDataFrame.groupby`, which worked under old versions of pandas, but are broken in pandas >=0.25.

### Reproduction: 
Install cheval with pandas >=0.25, and run the following
``` python
df1_dict = {
    "df2_id": [1,     2,   3,   1,   2,   3],
    "cats":   ["M", "F", "M", "F", "M", "F"],
}

df2_dict = {
    "col1":   ["a", "b", "c"],
}

df1 = LinkedDataFrame(pd.DataFrame(df1_dict))
df2 = LinkedDataFrame(pd.DataFrame(df2_dict))

df1.link_to(df2, "df2", on_self="df2_id")

for cat, group_df in df1.groupby("cats"):
    group_df.df2.col1 # ValueError: Length of values (6) does not match length of index (3)
```

This is happening because the private method interception to update the linkage indexes (`_take` and `_take_with_is_copy`) is no longer getting called in newer versions of pandas.

### Recommended fix: 
Find a less hack-ey method to update indexes. Either: 
- Move to having all updates happen in a stable public endpoint (likely `__finalize__`)
- Have all linkages check indexes when they are accessed and update them on-the-fly if necessary (would have to be careful, as a naïve approach may subtly break code if the indexes are the same length)

If performance poses a concern, we could do both and use a stable endpoind, but do lazy update of indexes.


#### Note: 
In pandas 1.3, there is a separate bug which causes LinkedDataFrames to not be properly initialized during groupbys: see [pandas issue #45363](https://github.com/pandas-dev/pandas/pull/45363). This would result in an infinite recursion on `LinkedDataFrame.__getattr__` when doing basically anything on the grouped `LinkedDataFrame` objects, and so the above code would issue a `RecursionError`.

This issue is basically unfixable on our end; it was fixed in pandas 1.4, but we may want cheval to issue a warning if it detects this incompatible version of pandas.

Discussion:

Patch:
diff --git a/cheval/api.py b/cheval/api.py
index 3e02d40..aa93f81 100644
--- a/cheval/api.py
+++ b/cheval/api.py
@@ -1,18 +1,16 @@
+from typing import Dict, Union, Set, TYPE_CHECKING, List, Generator, Hashable, Optional
 import abc
 from collections import deque
-from typing import (TYPE_CHECKING, Dict, Generator, Hashable, List, Optional,
-                    Set, Union)
 
-import attr
-import numpy as np
 import pandas as pd
+import numpy as np
+import attr
 
-from .exceptions import ModelNotReadyError
-from .ldf import LinkedDataFrame
 from .parsing.expr_items import ChainTuple
 from .parsing.expressions import Expression
+from .ldf import LinkedDataFrame
+from .exceptions import ModelNotReadyError
 from .utils import convert_series, to_numpy
-
 if TYPE_CHECKING:
     from .model import ChoiceModel
 
diff --git a/cheval/core.py b/cheval/core.py
index fd151d6..cb34dca 100644
--- a/cheval/core.py
+++ b/cheval/core.py
@@ -1,15 +1,9 @@
 """Collection of pure-Numba functions for sampling & probability computation"""
 from typing import Tuple
-
 import numpy as np
-from numba import boolean as nbool
-from numba import float32 as nfloat
-from numba import float64 as ndouble
-from numba import int64 as nlong
-from numba import njit
-from numba import optional as maybe
-from numba import prange, void
 from numpy import ndarray
+from numba import (prange, njit, float64 as ndouble, int64 as nlong, void, float32 as nfloat, boolean as nbool,
+                   optional as maybe)
 
 try:
     from numba.core.types import Tuple as NTuple
@@ -131,9 +125,9 @@ def simple_probabilities(weights: ndarray) -> ndarray:
 
 
 @njit([
-    NTuple((ndouble[:], ndouble))(ndouble[:], nbool), NTuple((ndouble[:], ndouble))(nfloat[:], nbool)
+    NTuple((ndouble[:], ndouble))(ndouble[:]), NTuple((ndouble[:], ndouble))(nfloat[:])
 ], nogil=True)
-def multinomial_probabilities(utilities: ndarray, check_infeasible=True) -> Tuple[ndarray, float]:
+def multinomial_probabilities(utilities: ndarray) -> Tuple[ndarray, float]:
     """Computes probabilities given a multinomial logit model formulation."""
     n_cols = len(utilities)
     p = np.zeros(n_cols, dtype=np.float64)  # Return value
@@ -145,23 +139,20 @@ def multinomial_probabilities(utilities: ndarray, check_infeasible=True) -> Tupl
         p[i] = expu
 
     if ls <= 0:
-        if check_infeasible:
-            raise UtilityBoundsError("MNL utilities all exceeded minimum value (Logsum == 0.0)")
-        for i in range(n_cols):
-            p[i] = 0  # ensure that `p` will return all zeros
-    else:
-        for i in range(n_cols):
-            p[i] = p[i] / ls
+        raise UtilityBoundsError("MNL utilities all exceeded minimum value (Logsum == 0.0)")
+
+    for i in range(n_cols):
+        p[i] = p[i] / ls
 
     return p, ls
 
 
 @njit([
-    NTuple((ndouble[:], ndouble))(ndouble[:], nlong[:], nlong[:], ndouble[:], nbool[:], nbool, nbool),
-    NTuple((ndouble[:], ndouble))(nfloat[:], nlong[:], nlong[:], ndouble[:], nbool[:], nbool, nbool)
+    NTuple((ndouble[:], ndouble))(ndouble[:], nlong[:], nlong[:], ndouble[:], nbool[:], nbool),
+    NTuple((ndouble[:], ndouble))(nfloat[:], nlong[:], nlong[:], ndouble[:], nbool[:], nbool)
 ], nogil=True)
-def nested_probabilities(utilities: ndarray, hierarchy, levels, logsum_scales, bottom_flags, scale_utilities=True,
-                         check_infeasible=True) -> Tuple[ndarray, float]:
+def nested_probabilities(utilities: ndarray, hierarchy, levels, logsum_scales, bottom_flags, scale_utilities=True
+                         ) -> Tuple[ndarray, float]:
     """Probability evaluation of a nested logit model, without needing a tree structure or any recursion."""
 
     n_cells = len(utilities)
@@ -203,8 +194,7 @@ def nested_probabilities(utilities: ndarray, hierarchy, levels, logsum_scales, b
         current_level -= 1
 
     if top_logsum <= 0:
-        if check_infeasible:
-            raise UtilityBoundsError("Nested logit top-level logsum is 0.0")
+        raise UtilityBoundsError("Nested logit top-level logsum is 0.0")
 
     # Step 2: Use logsums to compute conditional probabilities
     for index, parent in enumerate(hierarchy):
@@ -256,7 +246,7 @@ def simple_multisample(weights: ndarray, n: int, seed: int, out: ndarray = None)
 ], nogil=True)
 def multinomial_sample(utilities: ndarray, r: float) -> Tuple[int, float]:
     """Samples once from an array of multinomial logit utilities, from an existing random draw"""
-    p_array, ls = multinomial_probabilities(utilities, check_infeasible=True)  # Check infeasible req'd (for sampling)
+    p_array, ls = multinomial_probabilities(utilities)
     return sample_once(p_array, r), ls
 
 
@@ -266,7 +256,7 @@ def multinomial_sample(utilities: ndarray, r: float) -> Tuple[int, float]:
 ], nogil=True)
 def multinomial_multisample(utilities: ndarray, n: int, seed: int, out: ndarray = None) -> Tuple[np.ndarray, float]:
     """Samples multiple times from an array of multinomial logit utilities, based on a random seed. Thread-safe."""
-    p_array, ls = multinomial_probabilities(utilities, check_infeasible=True)  # Check infeasible req'd (for sampling)
+    p_array, ls = multinomial_probabilities(utilities)
     return sample_multi(p_array, n, seed, out), ls
 
 
@@ -277,10 +267,8 @@ def multinomial_multisample(utilities: ndarray, n: int, seed: int, out: ndarray
 def nested_sample(utilities: ndarray, r: float, parents, levels, ls_scales, bottom_flags, scale_utilities=True
                   ) -> Tuple[int, float]:
     """Samples once from an array of nested logit utilities, from an existing random draw"""
-    p_array, ls = nested_probabilities(
-        utilities, parents, levels, ls_scales, bottom_flags, scale_utilities=scale_utilities,
-        check_infeasible=True  # Check infeasible req'd (for sampling)
-    )
+    p_array, ls = nested_probabilities(utilities, parents, levels, ls_scales, bottom_flags,
+                                       scale_utilities=scale_utilities)
     return sample_once(p_array, r), ls
 
 
@@ -291,9 +279,8 @@ def nested_sample(utilities: ndarray, r: float, parents, levels, ls_scales, bott
 def nested_multisample(utilities: ndarray, parents, levels, ls_scales, bottom_flags, n: int, seed: int,
                        out: ndarray = None, scale_utilities=True) -> Tuple[ndarray, float]:
     """Samples multiple times from an array of nested logit utilities, based on a random seed. Thread-safe."""
-    p_array, ls = nested_probabilities(
-        utilities, parents, levels, ls_scales, bottom_flags, scale_utilities=scale_utilities, check_infeasible=True
-    )  # Check infeasible req'd (for sampling)
+    p_array, ls = nested_probabilities(utilities, parents, levels, ls_scales, bottom_flags,
+                                       scale_utilities=scale_utilities)
     return sample_multi(p_array, n, seed, out), ls
 
 
@@ -351,10 +338,10 @@ def worker_multinomial_sample(utilities: ndarray, n: int, seed: int) -> Tuple[nd
 
 
 @njit([
-    NTuple((ndouble[:, :], ndouble[:]))(ndouble[:, :], nbool),
-    NTuple((ndouble[:, :], ndouble[:]))(nfloat[:, :], nbool)
+    NTuple((ndouble[:, :], ndouble[:]))(ndouble[:, :]),
+    NTuple((ndouble[:, :], ndouble[:]))(nfloat[:, :])
 ], parallel=True, nogil=True)
-def worker_multinomial_probabilities(utilities: ndarray, check_infeasible=True) -> Tuple[ndarray, ndarray]:
+def worker_multinomial_probabilities(utilities: ndarray) -> Tuple[ndarray, ndarray]:
     """Runs multinomial_probabilities in parallel"""
     n_rows, n_cols = utilities.shape
     result = np.zeros((n_rows, n_cols), dtype=np.float64)
@@ -362,7 +349,7 @@ def worker_multinomial_probabilities(utilities: ndarray, check_infeasible=True)
 
     for i in prange(n_rows):
         utility_row = utilities[i, :]
-        p_array, ls = multinomial_probabilities(utility_row, check_infeasible)
+        p_array, ls = multinomial_probabilities(utility_row)
         result[i, :] = p_array
         ls_array[i] = ls
 
@@ -402,11 +389,11 @@ def worker_nested_sample(utilities: ndarray, parents, levels, ls_scales, bottom_
 
 
 @njit([
-    NTuple((ndouble[:, :], ndouble[:]))(ndouble[:, :], nlong[:], nlong[:], ndouble[:], nbool[:], nbool, nbool),
-    NTuple((ndouble[:, :], ndouble[:]))(nfloat[:, :], nlong[:], nlong[:], ndouble[:], nbool[:], nbool, nbool)
+    NTuple((ndouble[:, :], ndouble[:]))(ndouble[:, :], nlong[:], nlong[:], ndouble[:], nbool[:], nbool),
+    NTuple((ndouble[:, :], ndouble[:]))(nfloat[:, :], nlong[:], nlong[:], ndouble[:], nbool[:], nbool)
 ], parallel=True, nogil=True)
-def worker_nested_probabilities(utilities: ndarray, parents, levels, ls_scales, bottom_flags, scale_utilities=True,
-                                check_infeasible=True) -> Tuple[ndarray, ndarray]:
+def worker_nested_probabilities(utilities: ndarray, parents, levels, ls_scales, bottom_flags, scale_utilities=True
+                                ) -> Tuple[ndarray, ndarray]:
     """Runs nested_probabilities in parallel"""
     n_rows, n_cols = utilities.shape
     result = np.zeros((n_rows, n_cols), dtype=np.float64)
@@ -415,7 +402,7 @@ def worker_nested_probabilities(utilities: ndarray, parents, levels, ls_scales,
     for i in prange(n_rows):
         utility_row = utilities[i, :]
         p_array, ls = nested_probabilities(utility_row, parents, levels, ls_scales, bottom_flags,
-                                           scale_utilities=scale_utilities, check_infeasible=check_infeasible)
+                                           scale_utilities=scale_utilities)
         result[i, :] = p_array
         ls_array[i] = ls
 
diff --git a/cheval/ldf/api.py b/cheval/ldf/api.py
index 099cccc..c824cd7 100644
--- a/cheval/ldf/api.py
+++ b/cheval/ldf/api.py
@@ -613,14 +613,34 @@ class LinkedDataFrame(DataFrame):
         if not isinstance(other, LinkedDataFrame):
             return self
 
-        # Copy the link metadata
-        for link_name, link in other._links.items():
-            self._links[link_name] = link.copy_meta()
-
-        self._identified_links = other._identified_links.copy()
+        # Copy the link metadata from the parent dataframe
+        # Note that as of pandas 1.5, the names of the methods is not considered stable
+        # and may change in future versions.
+        if method == "groupby":
+            self._copy_links(other, reindex=True)
+        else:
+            for link_name, link in other._links.items():
+                self._links[link_name] = link.copy_meta()
+            self._identified_links = other._identified_links.copy()
 
         return self
 
+    def _copy_links(self, parent: 'LinkedDataFrame', reindex: bool = False):
+        """Copies linkage information from a parent dataframe
+
+        If `reindex` is True, computes an indexer based on the parent DataFrame's and
+        this DataFrame's indexes and passes it to `_LinkMeta.copy` to modify the linkage.
+        """
+        if reindex:
+            indexer = parent.index.get_indexer(self.index)
+        else:
+            indexer = None
+        link_names = set()
+        for link_name, link in parent._links.items():
+            self._links[link_name] = link.copy(indexer)
+            link_names.add(link_name)
+        self._identified_links = link_names
+
     def __subset_links(self, target: 'LinkedDataFrame', indexer: np.ndarray):
         for link_name, link_entry in self._links.items():
             target._links[link_name] = link_entry.copy(indexer)
@@ -768,72 +788,44 @@ class LinkedDataFrame(DataFrame):
             if refresh or entry.flat_indexer is None:
                 entry.precompute()
 
-    def pivot_table(self, values=None, index=None, columns=None, aggfunc='mean', fill_value=None, margins=False,
-                    dropna=True, margins_name='All'):
-        temp_columns = []
-        try:
-            new_index, temp_flags = self._make_temp_col(index)
-            if new_index[0] is None:
-                new_index = None
-            else:
-                for col_name, is_temp in zip(new_index, temp_flags):
-                    if is_temp:
-                        temp_columns.append(col_name)
-                if len(new_index) == 1:
-                    new_index = new_index[0]
-
-            new_columns, temp_flags = self._make_temp_col(columns)
-            if new_columns[0] is None:
-                new_columns = None
-            else:
-                for col_name, is_temp in zip(new_columns, temp_flags):
-                    if is_temp:
-                        temp_columns.append(col_name)
-                if len(new_columns) == 1:
-                    new_columns = new_columns[0]
-
-            new_values, temp_flags = self._make_temp_col(values)
-            if new_values[0] is None:
-                new_values = None
-            else:
-                for col_name, is_temp in zip(new_values, temp_flags):
-                    if is_temp:
-                        temp_columns.append(col_name)
-                if len(new_values) == 1:
-                    new_values = new_values[0]
-
-            # TODO: determine why margins=True causes things to crash...
-
-            return super().pivot_table(index=new_index, columns=new_columns, values=new_values, aggfunc=aggfunc,
-                                       fill_value=fill_value, margins=margins, dropna=dropna,
-                                       margins_name=margins_name)
-        finally:
-            for c in temp_columns:
-                del self[c]
-
-    def _make_temp_col(self, item: Union[str, List[str]]) -> Tuple[List[Optional[str]], List[bool]]:
-        if item is None:
-            return [None], [False]
-
-        item_is_single = isinstance(item, str)
-        items = [item] if item_is_single else item
-
-        new_columns, flags = [], []
-        for sub_item in items:
-            if sub_item in self:
-                new_columns.append(sub_item)
-                flags.append(False)
+    def pivot_table(self, values=None, index=None, columns=None, **kwargs):
+        # Construct a new DataFrame from any of the columns or lookups specified
+        # in values, index, or columns, and then pivot that
+        new_columns = {}
+        value_cols = self._resolve_columns(values)
+        new_columns.update(value_cols)
+        index_cols = self._resolve_columns(index)
+        new_columns.update(index_cols)
+        column_cols = self._resolve_columns(columns)
+        new_columns.update(column_cols)
+        pivot_df = pd.DataFrame(new_columns, index=self.index)
+
+        return pivot_df.pivot_table(index=index, columns=columns, values=values, **kwargs)
+
+    def _resolve_columns(self, lookup_items) -> Dict[str, pd.Series]:
+        """Resolves potential column names or lookups to a dict of Series
+
+        Helper for pivot_table to pull any relevant columns or linkage lookups
+        into a clean DataFrame.
+        Because pivot_table can take a wide variety of inputs as index, columns, and values,
+        this attempts to be very flexible with whether and how it looks up provided items.
+
+        Any items which get resolved as columns or linkage lookups are returned as a
+        Series in a key: value pair in the returned dict.
+        """
+        new_columns = {}
+        if isinstance(lookup_items, str):
+            if lookup_items in self.columns:
+                new_columns[lookup_items] = self[lookup_items]
             else:
-                new_column_values = self.evaluate(sub_item)  # Evaluate the column to find links
-                counter = 0
-                new_column_name = f"temp_{counter}"
-                while new_column_name in self:
-                    counter += 1
-                    new_column_name = f"temp_{counter}"
-                self[new_column_name] = new_column_values
-                new_columns.append(new_column_name)
-                flags.append(True)
-        return new_columns, flags
+                new_columns[lookup_items] = self.evaluate(lookup_items)
+        elif (lookup_items is not None) and not isinstance(lookup_items, (pd.Series, np.ndarray)):
+            for item in lookup_items:
+                if lookup_items in self.columns:
+                    new_columns[item] = self[item]
+                else:
+                    new_columns[item] = self.evaluate(item)
+        return new_columns
 
     def get_link_target(self, name) -> Union[DataFrame, 'LinkedDataFrame']:
         """Gets the referenced target ("other") of an outbound link. This is useful when the original reference to the
diff --git a/cheval/model.py b/cheval/model.py
index 8430806..36ee5c4 100644
--- a/cheval/model.py
+++ b/cheval/model.py
@@ -1,24 +1,21 @@
+from typing import Iterable, Dict, Union, Tuple, Set, Hashable, Optional, Any
 from itertools import chain as iter_chain
-from logging import Logger
 from multiprocessing import cpu_count
-from typing import Any, Dict, Hashable, Iterable, Optional, Set, Tuple, Union
+from logging import Logger
 
-import numba as nb
-import numexpr as ne
-import numpy as np
+from pandas import Series, DataFrame, Index, MultiIndex
 import pandas as pd
 from numpy import ndarray
-from pandas import DataFrame, Index, MultiIndex, Series
-
-from .api import (AbstractSymbol, ChoiceNode, ExpressionGroup,
-                  ExpressionSubGroup, MatrixSymbol, NumberSymbol, TableSymbol,
-                  VectorSymbol)
-from .core import (UtilityBoundsError, fast_indexed_add,
-                   worker_multinomial_probabilities, worker_multinomial_sample,
-                   worker_nested_probabilities, worker_nested_sample)
+import numpy as np
+import numexpr as ne
+import numba as nb
+
+from .api import (AbstractSymbol, ExpressionGroup, ChoiceNode, NumberSymbol, VectorSymbol, TableSymbol, MatrixSymbol,
+                  ExpressionSubGroup)
 from .exceptions import ModelNotReadyError
-from .parsing.constants import (NAN_STR, NEG_INF_STR, NEG_INF_VAL, OUT_STR,
-                                RESERVED_WORDS)
+from .core import (worker_nested_probabilities, worker_nested_sample, worker_multinomial_probabilities,
+                   worker_multinomial_sample, fast_indexed_add, UtilityBoundsError)
+from .parsing.constants import NAN_STR, NEG_INF_STR, NEG_INF_VAL, OUT_STR, RESERVED_WORDS
 
 
 class ChoiceModel(object):
@@ -539,24 +536,20 @@ class ChoiceModel(object):
         return retval
 
     def run_stochastic(self, n_threads: int = 1, clear_scope: bool = True, logger: Logger = None,
-                       group: Hashable = None, scale_utilities: bool = True,
-                       check_infeasible: bool = True) -> Tuple[DataFrame, Series]:
+                       group: Hashable = None, scale_utilities: bool = True) -> Tuple[DataFrame, Series]:
         """For each record, compute the probability distribution of the logit model. A DataFrame will be returned whose
         columns match the sorted list of node names (alternatives) in the model. Probabilities over all alternatives for
         each record will sum to 1.0.
 
         Args:
-            n_threads (int): The number of threads to be used in the computation. Must be >= 1.
-            clear_scope (bool): If True data stored in the scope for utility computation will be released, freeing up
-                memory. Turning this off is of limited use.
-            logger (Logger): Optional Logger instance which reports expressions being evaluated
-            group (Hashable): Evaluate only the specified utility group. Raises KeyError if the group name is not
-                defined.
-            scale_utilities (bool): For a nested model, if True then lower-level utilities will be divided by the logsum
-                scale of the parent nest. If False, no scaling is performed. This is entirely dependant on the reported
-                form of estimated model parameters.
-            check_infeasible (bool): Flag to check for infeasible solutions (i.e., probabilities are 0 for all choices).
-                Set to ``False`` to disable the check and allow for cases where no solutions exist.
+            n_threads: The number of threads to be used in the computation. Must be >= 1.
+            clear_scope: If True data stored in the scope for utility computation will be released, freeing up memory.
+                Turning this off is of limited use.
+            logger: Optional Logger instance which reports expressions being evaluated
+            group: Evaluate only the specified utility group. Raises KeyError if the group name is not defined.
+            scale_utilities: For a nested model, if True then lower-level utilities will be divided by the logsum scale
+                of the parent nest. If False, no scaling is performed. This is entirely dependant on the reported form
+                of estimated model parameters.
 
         Returns:
             Tuple[DataFrame, Series]: The first item returned is always the results of the model evaluation,
@@ -580,11 +573,10 @@ class ChoiceModel(object):
         if nested:
             hierarchy, levels, logsum_scales, bottom_flags = self._flatten()
             raw_result, logsum = worker_nested_probabilities(utility_table, hierarchy, levels, logsum_scales,
-                                                             bottom_flags, scale_utilities=scale_utilities,
-                                                             check_infeasible=check_infeasible)
+                                                             bottom_flags, scale_utilities=scale_utilities)
             result_frame = self._build_nested_stochastic_frame(raw_result)
         else:
-            raw_result, logsum = worker_multinomial_probabilities(utility_table, check_infeasible=check_infeasible)
+            raw_result, logsum = worker_multinomial_probabilities(utility_table)
             result_frame = DataFrame(raw_result, index=self.decision_units, columns=self.choices)
         logsum = Series(logsum, index=self.decision_units)
 
diff --git a/cheval/parsing/__init__.py b/cheval/parsing/__init__.py
index e69de29..119fabb 100644
--- a/cheval/parsing/__init__.py
+++ b/cheval/parsing/__init__.py
@@ -0,0 +1,1 @@
+from .ast_transformer import NAN_STR
diff --git a/cheval/parsing/ast_transformer.py b/cheval/parsing/ast_transformer.py
index e4eaab4..65d3da5 100644
--- a/cheval/parsing/ast_transformer.py
+++ b/cheval/parsing/ast_transformer.py
@@ -1,15 +1,16 @@
 """Module for AST Transformer subclass which specially-parses utility expressions"""
+from typing import Dict, Tuple, Set, Union
+
 import ast
 from collections import deque
-from typing import Dict, Set, Tuple, Union
 
 import astor
 import numpy as np
 from numexpr import expressions as nee
 
-from .constants import NAN_STR
 from .exceptions import UnsupportedSyntaxError
 from .expr_items import ChainedSymbol, EvaluationMode
+from .constants import *
 
 # Only nodes used in expressions are included, due to the limited parsing
 _UNSUPPORTED_NODES: Tuple[type] = (ast.Load, ast.Store, ast.Del, ast.IfExp, ast.Subscript, ast.ListComp, ast.DictComp,
diff --git a/cheval/parsing/expr_items.py b/cheval/parsing/expr_items.py
index f94b627..113da44 100644
--- a/cheval/parsing/expr_items.py
+++ b/cheval/parsing/expr_items.py
@@ -1,7 +1,7 @@
 """Collection of common classes to store more complex substitutions (e.g. LinkedDataFrame lookup)"""
 from collections import deque, namedtuple
-from enum import Enum
 from typing import Dict, Tuple
+from enum import Enum
 
 ChainTuple = namedtuple("ChainTuple", ['chain', 'func', 'args', 'withfunc'])
 
diff --git a/cheval/parsing/expressions.py b/cheval/parsing/expressions.py
index 1dc6468..48d5f84 100644
--- a/cheval/parsing/expressions.py
+++ b/cheval/parsing/expressions.py
@@ -1,14 +1,14 @@
+from typing import Dict, Set, Tuple, Optional
 import ast
-from typing import Dict, Optional, Set, Tuple
 
-import astor
 import attr
-import numpy as np
 import pandas as pd
+import numpy as np
+import astor
 
+from .expr_items import ChainedSymbol
 from .ast_transformer import ExpressionParser
 from .exceptions import UnsupportedSyntaxError
-from .expr_items import ChainedSymbol
 
 
 def _split_filter(e: str) -> Tuple[str, Optional[str]]:
diff --git a/cheval/utils.py b/cheval/utils.py
index 19d8ea0..443736e 100644
--- a/cheval/utils.py
+++ b/cheval/utils.py
@@ -1,8 +1,7 @@
-from typing import Union
-
-import numpy as np
 import pandas as pd
 from pandas.api.types import is_integer_dtype
+import numpy as np
+from typing import Union
 
 _USE_TO_NUMPY = hasattr(pd.Series, 'to_numpy')
 


Test patch:
diff --git a/testing/test_ldf.py b/testing/test_ldf.py
index f4adf35..128dac7 100644
--- a/testing/test_ldf.py
+++ b/testing/test_ldf.py
@@ -105,6 +105,87 @@ def test_pivot_table():
         'apartment': {'Ford': 1, 'Honda': 0, 'Toyota': 0}, 'house': {'Ford': 1, 'Honda': 2, 'Toyota': 1}
     })
     expected_result.index.name = 'manufacturer'
-    expected_result.columns.name = 'temp_0'
+    expected_result.columns.name = 'household.dwelling_type'
 
     assert_frame_equal(test_result, expected_result)
+
+
+def test_pivot_table_multiindex():
+    """Test that pivot_table can correctly handle tables with a MultiIndex
+
+    Creates LinkedDataFrames with MultiIndexes to ensure pivot_table is not getting
+    tripped up trying to use the new pivoted index to reindex linkages which use
+    the original table indexes.
+    Generally, pivoting a LinkedDataFrame should produce a table which does not
+    contain any linkages.
+
+    This catches a specific error which arose in the GGHM demand model Python 3 conversion.
+    """
+    # We want the indexes for each table to have different dimensions,
+    # so that indexing logic errors are made explicit as raised exceptions
+    # We're using the following MultiIndexes:
+    # For the vehicles table: (hhid, veh_id, dummy)
+    # For the households table: (hhid, dummy, dummy, dummy)
+    # The output pivot table should have an index of (manufacturer)
+    # A (manufacturer, household.dwelling_type) index may be created by pandas
+    # as an internal implementation detail
+    n_rows_veh = len(vehicles_data["household_id"])
+    veh_index = pd.MultiIndex.from_arrays([vehicles_data["household_id"], vehicles_data["vehicle_id"],
+                                           [0]*n_rows_veh
+                                           ])
+    vehicles = LinkedDataFrame(vehicles_data, index=veh_index)
+
+    n_rows_hh = len(households_data["household_id"])
+    hh_index = pd.MultiIndex.from_arrays([households_data["household_id"],
+                                          [0]*n_rows_hh, [0]*n_rows_hh, [0]*n_rows_hh
+                                          ])
+    households = LinkedDataFrame(households_data, index=hh_index)
+
+    vehicles.link_to(households, 'household', on=['household_id'])
+    households.link_to(vehicles, 'vehicles', on=['household_id'])
+
+    test_result = vehicles.pivot_table(values='vehicle_id', index='manufacturer', columns='household.dwelling_type',
+                                       aggfunc="sum", fill_value=0.0)
+    test_result.columns = test_result.columns.astype(str)
+
+    expected_result = pd.DataFrame({
+        'apartment': {'Ford': 0, 'Honda': 0, 'Toyota': 0}, 'house': {'Ford': 1, 'Honda': 0, 'Toyota': 0}
+    })
+    expected_result.index.name = 'manufacturer'
+    expected_result.columns.name = 'household.dwelling_type'
+
+    assert_frame_equal(test_result, expected_result)
+
+
+def test_groupby():
+    """Verifies basic groupby functionality for LinkedDataFrames
+
+    Most groupby functionality should be handled by pandas, but the internal
+    implementation may differ version-to-version, which may impact the ability of
+    LinkedDataFrame objects to correctly update the linkages.
+
+    This tests that groupby can be run on LinkedDataFrames, and that the linked dataframe
+    is accessible from each group and that it contains the correct data.
+    """
+
+    df1_dict = {
+        "df2_id": [  0,   1,   2,   0,   1,   2,   0],
+        "cats":   ["F", "M", "F", "M", "F", "M", "F"],
+    }
+
+    df2_dict = {
+        "col1":   ["a", "b", "c"],
+    }
+
+    df1 = LinkedDataFrame(pd.DataFrame(df1_dict))
+    df2 = LinkedDataFrame(pd.DataFrame(df2_dict))
+
+    df1.link_to(df2, "df2", on_self="df2_id")
+
+    groups = df1.groupby("cats")
+
+    for cat, group_df in groups:
+        if cat == "F":
+            assert list(sorted(group_df.df2.col1)) == ["a", "a", "b", "c"]
+        else:
+            assert list(sorted(group_df.df2.col1)) == ["a", "b", "c"]