Reshaping num_rx_ant and num_tx_ant dimensions of channel tensor into geometric components for single and multi panel arrays

[naperman]

I'm doing some PUSCH uplink processing and I'm trying to reshape the output of the UMi channel model(channel coefficients 'a') as follows:

For the channels coefficients tensor a:
[batch size, num_rx, num_rx_ant, num_tx, num_tx_ant, num_paths, num_time_steps]

I want to unravel the num_rx_ant dimension and reshape [..., num_rx_ant, ...] to

[..., num_rows, num_cols, num_polarizations, ...] for single-panel

and

[..., num_rows_per_panel, num_cols_per_panel, num_rx, num_rows, num_cols, num_polarizations, ...] for multi-panel

I'm using einops to perform the reshaping, so a readable syntax like the following would work fine:

a_reshaped = rearrange(a, '... (num_rows num_cols num_polarizations) ... -> ... num_rows num_cols num_polarizations ...',
num_rows=4, num_cols=8, num_polarizations=2)

I just want to make sure if i'm getting the order of unravelled dimensions correct for the single and multi-panel cases. The order doesn't have to be the specific ones i showed in the above examples, but just something that accurately reflects the ordering that sionna uses, which can be split into the above factors.

Looking at the documentations(figures reproduced below), the ordering of the polarization dimension is not clear in single panel arrays.

For multi-panel arrays, it looks more complicated because the ordering of elements cuts across panels and not done panel-by-panel.