Add as_dataframe() method to GroupEffectsMatrix (#115)

* Update python-version in test.yml

* Add as_dataframe() method to GroupEffectsMatrix

.pylintrc

@@ -61,9 +61,7 @@ disable=missing-docstring,
         too-many-locals,
         too-many-branches,
         too-many-statements,
-        no-self-use,
         too-few-public-methods,
-        bad-continuation,
         invalid-name
 
 
@@ -132,13 +130,6 @@ max-line-length=100
 # Maximum number of lines in a module
 max-module-lines=1000
 
-# List of optional constructs for which whitespace checking is disabled. `dict-
-# separator` is used to allow tabulation in dicts, etc.: {1  : 1,\n222: 2}.
-# `trailing-comma` allows a space between comma and closing bracket: (a, ).
-# `empty-line` allows space-only lines.
-no-space-check=trailing-comma,
-               dict-separator
-
 # Allow the body of a class to be on the same line as the declaration if body
 # contains single statement.
 single-line-class-stmt=no
@@ -500,4 +491,4 @@ min-public-methods=2
 
 # Exceptions that will emit a warning when being caught. Defaults to
 # "Exception"
-overgeneral-exceptions=Exception
+overgeneral-exceptions=builtins.Exception

formulae/matrices.py

@@ -428,6 +428,13 @@ def evaluate_new_data(self, data):
         new_instance.evaluated = True
         return new_instance
 
+    def as_dataframe(self):
+        """Returns `self.design_matrix` as a pandas.DataFrame."""
+        columns = []
+        for term in self.terms.values():
+            columns.extend(term.labels)
+        return pd.DataFrame(self.design_matrix, columns=columns)
+
     def __getitem__(self, term):
         """Get the sub-matrix that corresponds to a given term.
 

tests/test_design_matrices.py

@@ -1052,6 +1052,22 @@ def test_group_specific_as_array(data):
     assert np.asarray(group).shape == (20, 4)
 
 
+def test_group_specific_as_data_frame(data):
+    _, _, group = design_matrices("y ~ 1 + (x1|g) + (x1:x2|g) + (h|g)", data)
+    group_specific_dataframe = group.as_dataframe()
+    assert group_specific_dataframe.shape == (20, 8)
+    assert group_specific_dataframe.columns.tolist() == [
+        "1|g[A]",
+        "1|g[B]",
+        "x1|g[A]",
+        "x1|g[B]",
+        "x1:x2|g[A]",
+        "x1:x2|g[B]",
+        "h[B]|g[A]",
+        "h[B]|g[B]",
+    ]
+
+
 def test_group_specific_repr_and_str(data):
     _, _, group = design_matrices("y ~ 1 + (x1|g) + (h|g)", data)
     text = (