BUG: Preserve missing categorical observations in dummy conversion

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
 

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
 

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])

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])