BUG: Preserve missing categorical observations in dummy conversion

linearmodels/iv/data.py

@@ -36,6 +36,8 @@ def convert_columns(s: pd.Series, drop_first: bool) -> linearmodels.typing.AnyPa
     """
     if isinstance(s.dtype, pd.CategoricalDtype):
         out = pd.get_dummies(s, drop_first=drop_first)
+        if s.hasnans:
+            out = out.astype(float).where(s.notna(), axis=0)
         # TODO: Remove once pandas typing fixed
         out.columns = pd.Index([str(s.name) + "." + str(c) for c in out])
         return out

linearmodels/panel/data.py

@@ -162,6 +162,8 @@ def convert_columns(s: Series, drop_first: bool) -> linearmodels.typing.AnyPanda
 
     if isinstance(s.dtype, pd.CategoricalDtype):
         out = get_dummies(s, drop_first=drop_first)
+        if s.hasnans:
+            out = out.astype(float).where(s.notna(), axis=0)
         # TODO: Remove once pandas typing fixed
         out.columns = Index([str(s.name) + "." + str(c) for c in out])
         return out

linearmodels/tests/iv/test_data.py

@@ -250,3 +250,15 @@ 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()

linearmodels/tests/panel/test_data.py

@@ -958,3 +958,18 @@ 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()