MPI Collective communication along axes with uneven data distribution per rank

Question

I am attempting to implement a method in MPI for a well established particle simulation program that involves image processing. The program runs a loop for millions of iterations that performs a calculation on the particles in 3D space and then projects them onto a 2D image with a Gaussian distribution before calculations are performed on the image.

As an example with a 3x3x3 set of ranks, reductions are performed along the third axis until the "top" face of 3x3 cells has all of the image information for their respective "pencils". This local information for the pencil is sufficient to do the calculations on the image.

In the simplest case of an even data distribution where each rank has the same size of cell and therefore section of image, I can think of a few ways to do this. It seems like creating a Cartesian topology (and then using MPI_Cart_sub) or hand constructed derived communicators for each of the pencils to be reduced across would work and be relatively efficient.

The issue comes into play here in the case where each cell is non uniform. In particular, the cells are divided into a kd-tree in three dimensions to improve load balancing for cases where the distribution of particles is extremely non-homogeneous.

Graphical example

The question then becomes how to do this same kind of operation where there is a reduction of dimensionality. This above method falls apart because cells are no longer arranged into evenly sized pencils. I've extensively looked through the MPI documentation and tried to find similar problems that may give me some insight into the most efficient way to do this. Storing a full copy of the image on each rank and then all reducing over the entire process space is too expensive with the number of nodes that are being used. I suspect there is some combination of derived communicators and MPI types that could solve this problem in a relatively elegant way, but I haven't been able to think of anything. I would greatly appreciate any help in thinking through this and developing a good method.