Merge pull request #84 from Shotgunosine/fixef_labels

fix fixed effects indices

Author: ejolly

pymer4/models/Lmer.py

@@ -748,13 +748,22 @@ def fit(
         # Cluster (e.g subject) level coefficients
         rstring = """
             function(model){
+            getIndex <- function(df){
+                    orignames <- names(df)
+                    df <- transform(df, index=row.names(df))
+                    names(df) <- append(orignames, c("index"))
+                    df
+                    }
             out <- coef(model)
+            out <- lapply(out, getIndex)
             out
             }
         """
         fixef_func = robjects.r(rstring)
         fixefs = fixef_func(self.model_obj)
-        fixefs = [pd.DataFrame(f) for f in fixefs]
+        fixefs = [
+            pd.DataFrame(e, index=e.index).drop(columns=["index"]) for e in fixefs
+        ]
         if len(fixefs) > 1:
             if self.coefs is not None:
                 f_corrected_order = []

pymer4/tests/test_models.py

@@ -105,6 +105,8 @@ def test_gaussian_lmm():
     assert np.allclose(model.coefs["Estimate"], estimates, atol=0.001)
 
     assert isinstance(model.fixef, list)
+    assert (model.fixef[0].index.astype(int) == df.Group.unique()).all()
+    assert (model.fixef[1].index.astype(float) == df.IV3.unique()).all()
     assert model.fixef[0].shape == (47, 3)
     assert model.fixef[1].shape == (3, 3)
 
@@ -133,6 +135,24 @@ def test_gaussian_lmm():
     # Smoketest for old_optimizer
     model.fit(summarize=False, old_optimizer=True)
 
+    # test fixef code for 1 fixed effect
+    model = Lmer("DV ~ IV3 + IV2 + (IV2|Group)", data=df)
+    model.fit(summarize=False, control=opt_opts)
+
+    assert (model.fixef.index.astype(int) == df.Group.unique()).all()
+    assert model.fixef.shape == (47, 3)
+    assert np.allclose(model.coefs.loc[:, "Estimate"], model.fixef.mean(), atol=0.01)
+
+    # test fixef code for 0 fixed effects
+    model = Lmer("DV ~ (IV2|Group) + (1|IV3)", data=df)
+    model.fit(summarize=False, control=opt_opts)
+
+    assert isinstance(model.fixef, list)
+    assert (model.fixef[0].index.astype(int) == df.Group.unique()).all()
+    assert (model.fixef[1].index.astype(float) == df.IV3.unique()).all()
+    assert model.fixef[0].shape == (47, 2)
+    assert model.fixef[1].shape == (3, 2)
+
 
 def test_contrasts():
     df = sns.load_dataset("gammas").rename(columns={"BOLD signal": "bold"})