diff --git a/linearmodels/iv/data.py b/linearmodels/iv/data.py index 14c1ad2ed5..86d0fb4bd8 100644 --- a/linearmodels/iv/data.py +++ b/linearmodels/iv/data.py @@ -38,6 +38,11 @@ def convert_columns(s: pd.Series, drop_first: bool) -> linearmodels.typing.AnyPa out = pd.get_dummies(s, drop_first=drop_first) # TODO: Remove once pandas typing fixed out.columns = pd.Index([str(s.name) + "." + str(c) for c in out]) + if s.isna().any(): + used = out.ne(0).any(axis=0) + if used.any(): + out = out.loc[:, used] + out = out.astype(float).where(s.notna(), axis=0) return out return s diff --git a/linearmodels/panel/data.py b/linearmodels/panel/data.py index 456f43ed01..d7587fb45a 100644 --- a/linearmodels/panel/data.py +++ b/linearmodels/panel/data.py @@ -164,6 +164,11 @@ def convert_columns(s: Series, drop_first: bool) -> linearmodels.typing.AnyPanda out = get_dummies(s, drop_first=drop_first) # TODO: Remove once pandas typing fixed out.columns = Index([str(s.name) + "." + str(c) for c in out]) + if s.isna().any(): + used = out.ne(0).any(axis=0) + if used.any(): + out = out.loc[:, used] + out = out.astype(float).where(s.notna(), axis=0) return out return s diff --git a/linearmodels/tests/iv/test_data.py b/linearmodels/tests/iv/test_data.py index 73da5beec1..6b0c671fd8 100644 --- a/linearmodels/tests/iv/test_data.py +++ b/linearmodels/tests/iv/test_data.py @@ -250,3 +250,26 @@ def test_duplicate_column_names(): x = pd.DataFrame(np.ones((3, 2)), columns=["x", "x"]) with pytest.raises(ValueError, match=r"DataFrame contains duplicate column names"): IVData(x) + + +@pytest.mark.parametrize("drop_first", [True, False]) +def test_categorical_missing_values(drop_first: bool) -> None: + values = pd.Series(pd.Categorical(["a", None, "b", "a"]), name="group") + data = IVData(values, drop_first=drop_first) + assert_equal(data.isnull.to_numpy(), [False, True, False, False]) + assert np.isnan(data.ndarray[1]).all() + assert_equal(data.pandas["group.b"].iloc[[0, 2, 3]].to_numpy(), [0, 1, 0]) + data.drop(data.isnull) + assert data.rows == [0, 2, 3] + assert not data.isnull.any() + + +@pytest.mark.parametrize("drop_first", [True, False]) +def test_categorical_missing_values_with_unused_category(drop_first: bool) -> None: + values = pd.Series( + pd.Categorical(["a", None, "b", "a"], categories=["a", "b", "unused"]), + name="group", + ) + data = IVData(values, drop_first=drop_first) + assert "group.unused" not in data.cols + assert_equal(data.isnull.to_numpy(), [False, True, False, False]) diff --git a/linearmodels/tests/panel/test_data.py b/linearmodels/tests/panel/test_data.py index 1e20a4f9d3..7fb3430ee9 100644 --- a/linearmodels/tests/panel/test_data.py +++ b/linearmodels/tests/panel/test_data.py @@ -958,3 +958,35 @@ def test_fake_panel_properties(mi_df): assert_index_equal(panel.minor_axis, mi_df.index.levels[0]) df = panel.to_frame() assert_frame_equal(df, mi_df) + + +@pytest.mark.parametrize("drop_first", [True, False]) +def test_categorical_missing_values(drop_first): + index = MultiIndex.from_product([["a", "b"], [0, 1]]) + values = DataFrame( + {"group": Categorical(["first", None, "second", "first"])}, index=index + ) + data = PanelData(values, drop_first=drop_first) + assert_equal(data.isnull.to_numpy(), [False, True, False, False]) + assert np.isnan(data.values2d[1]).all() + assert_equal(data.dataframe["group.second"].iloc[[0, 2, 3]].to_numpy(), [0, 1, 0]) + data.drop(data.isnull) + assert data.dataframe.index.tolist() == [("a", 0), ("b", 0), ("b", 1)] + assert not data.isnull.any() + + +@pytest.mark.parametrize("drop_first", [True, False]) +def test_categorical_missing_values_with_unused_category(drop_first): + index = MultiIndex.from_product([["a", "b"], [0, 1]]) + values = DataFrame( + { + "group": Categorical( + ["first", None, "second", "first"], + categories=["first", "second", "unused"], + ) + }, + index=index, + ) + data = PanelData(values, drop_first=drop_first) + assert "group.unused" not in data.dataframe + assert_equal(data.isnull.to_numpy(), [False, True, False, False])