Repository: bambinos/formulae

Issue:
Add `as_dataframe` method for group level design matrix
Simulated data:
```python
groups = 3
samples_per_group = 10
N = groups * samples_per_group
x = np.tile(np.linspace(0, 10, samples_per_group), groups)
group = np.repeat(np.arange(groups), samples_per_group)
y = 2 + -0.5 * group + rng.normal(0, 0.1, N)
df = pd.DataFrame({"x": x, "group": group, "y": y})
```

If we model this with `y ~ 1 + (1 | C(group))` and were interested in the design matrices, then at the moment we _can_ do this:

```python
dm = design_matrices("y ~ 1 + (1 | C(group))", df)
dm.common.as_dataframe()
```

But we cannot call the `as_dataframe` method on the group level design matrix. Ie. we can only call

```python
np.array(dm.group)
```
which result in this
![Screenshot 2023-06-18 at 09 59 31](https://github.com/bambinos/formulae/assets/6765047/cb5863d0-99e4-4f0a-b742-58f37c4b34b1)

but it would be nice to be able to call

```python
dm.group.as_dataframe()
```

Not knowing how the internals work, I believe it would just be a matter of building a dataframe with the correct column labels, which should be `1|C(group)[0]`, `1|C(group)[1]`, and `1|C(group)[2]`.

This would be particularly useful for teaching and learning purposes.

Discussion:
@tomicapretto: Interesting idea indeed!@lucianopaz: @tomicapretto, I just ran into this issue and would like to try to add the functionality. I've written this. Could you tell me if it's OK?

```python
import formulae
import pandas as pd


def get_group_dataframe(group: formulae.matrices.GroupEffectsMatrix) -> pd.DataFrame:
    assert isinstance(group, formulae.matrices.GroupEffectsMatrix)
    columns = []
    for term in group.terms.values():
        columns.extend(term.labels)
    return pd.DataFrame(group.design_matrix, columns=columns)
```@tomicapretto: @lucianopaz that is correct. Now that you did the work, I'm going to add it as a method :) 

Solution patch:
diff --git a/formulae/matrices.py b/formulae/matrices.py
index 68fde88..fddd777 100644
--- a/formulae/matrices.py
+++ b/formulae/matrices.py
@@ -428,6 +428,13 @@ class GroupEffectsMatrix:
         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.
 


Test patch:
diff --git a/tests/test_design_matrices.py b/tests/test_design_matrices.py
index b1178ae..f040559 100644
--- a/tests/test_design_matrices.py
+++ b/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 = (