Repository: wsp-sag/cheval

Issue:
Agg on bool columns broken
Issue came up in GGHM demand model Python 3 update. `LinkedDataFrame.evaluate` fails on expressions of the form `link.sum(attribute == 3)`, which should be a valid expression to obtain the number of linked rows for which `attribute` satisfies some boolean expression.

The error happens when the link manager tries to apply a NaN missing values mask to a column with a bool datatype in preparation for the aggregation. Bool does not support NaN, so even if there are no missing values, pandas throws an error. There is logic in the LDF code to apply a compatible mask for integer columns, but bool columns are not included in this.

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 c824cd7..0bba2b4 100644
--- a/cheval/ldf/api.py
+++ b/cheval/ldf/api.py
@@ -550,7 +550,7 @@ class LinkedDataFrame(DataFrame):
 
                 # A fill value of NaN is only disallowed for integer types
                 fill_value = np.nan
-                if series_type == PandasDtype.INT_NAME:
+                if series_type in {PandasDtype.INT_NAME, PandasDtype.UINT_NAME, PandasDtype.BOOL_NAME}:
                     fill_value = int_fill
                 elif series_type == PandasDtype.TIME_NAME:
                     raise NotImplementedError("Haven't found a way to instantiate NaT filler")
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 128dac7..d278834 100644
--- a/testing/test_ldf.py
+++ b/testing/test_ldf.py
@@ -189,3 +189,25 @@ def test_groupby():
             assert list(sorted(group_df.df2.col1)) == ["a", "a", "b", "c"]
         else:
             assert list(sorted(group_df.df2.col1)) == ["a", "b", "c"]
+
+
+def test_bool_columns():
+    """Test that bool columns get correctly handled in aggregations on linkages
+
+    Errors with bool columns originally came up when attempting to aggregate on
+    a bool expression. It is hoped that this will catch similar errors with bool
+    columns.
+    """
+
+    vehicles = LinkedDataFrame(vehicles_data)
+    households = LinkedDataFrame(households_data)
+
+    vehicles.link_to(households, 'household', on='household_id')
+    households.link_to(vehicles, 'vehicles', on='household_id')
+
+    expression = "vehicles.sum(manufacturer=='Honda')"
+    test_result = households.evaluate(expression)
+
+    expected_result = pd.Series([1, 0, 0, 1], index=households.index)
+
+    assert_series_equal(test_result, expected_result)