I’m attempting to convert a grouped Pandas DataFrame to a 3D numpy array, where every DataFrame cell is mapped to a 1D numpy array. I’m expecting a 3D array of shape (4, 2, 3), but instead I get a shape (4, 2) 2D array of shape (3,) 1D arrays.

Simplified example:

example_df = pd.DataFrame(
    [["A", True, True], 
     ["A", True, False], 
     ["B", False, True], 
     ["B", False, False]],
    columns=["grp", "col1", "col2"],
)


def mapper(val):
    return np.ones(3) if val is True else np.zeros(3)


result = (
    example_df.groupby(["grp"])
    .apply(lambda df: df.map(mapper), include_groups=False)
    .to_numpy()
)

print(type(result))
print(type(result[0]))
print(type(result[0][0]))
print(type(result[0][0][0]))
print(result.shape)
print(result[0, 0].shape)

Output:

<class 'numpy.ndarray'>
<class 'numpy.ndarray'>
<class 'numpy.ndarray'>
<class 'numpy.float64'>
(4, 2)
(3,)