fix setting of index

Author: MarcoGorelli

pandas/core/groupby/generic.py

@@ -956,7 +956,8 @@ def aggregate(self, func=None, *args, **kwargs):
                     result = self._aggregate_frame(func)
                 else:
                     result.columns = Index(
-                        result.columns.levels[0], name=self._selected_obj.columns.name
+                        [i[:-1] if len(i) > 2 else i[0] for i in result.columns],
+                        name=self._selected_obj.columns.name,
                     )
 
         if not self.as_index:

pandas/tests/groupby/aggregate/test_aggregate.py

@@ -691,6 +691,17 @@ def test_agg_relabel_multiindex_duplicates():
     tm.assert_frame_equal(result, expected)
 
 
+def test_multiindex_custom_func():
+    # GH 31777
+    df = pd.DataFrame(
+        np.random.rand(10, 4), columns=pd.MultiIndex.from_product([[1, 2], [3, 4]])
+    )
+    grp = df.groupby(np.r_[np.ones(5), np.zeros(5)])
+    result = grp.agg(lambda s: s.mean())
+    expected = grp.agg("mean")
+    tm.assert_frame_equal(result, expected)
+
+
 def myfunc(s):
     return np.percentile(s, q=0.90)