Improve chunking for calibration
The current batching strategy for calibration involves a flattening of all spatial dimensions. This makes it difficult to relate data during the calibration procedure to its spatial location. Furthermore, it necessitates reshaping and thus possibly incurs communication overhead. It would be good to improve the batching and chunking strategy to minimize communication and facilitate better parallelization.