Search code examples
pythonpandasdataframenumpy3d

Creating 3D python data from index sums of 2D data


I have a symmetrical dataframe with roughly 400 rows & columns. I would like to turn it into a 3-dimensional object (cube) by taking the sums of each pair between the three X,Y,Z indices.

For example, say the coordinates on the cube were A-B-C, I would take the sum of the following cells in the 2D dataframe/array: A-B,B-C,A-C. Remember again since it is symmetrical, whether the first or second index is the row or column does not matter here.

Here's a 3x3 snapshot of the data:

        #   A   B   C
df_in = [[  1,  5,  2], # A
         [  5,  1,  3], # B
         [  2,  3,  1]] # C

And then these would be the 3 planes that make up the outputted cube:

          #   A   B   C
outputA = [[  3, 11,  5], # A
           [ 11, 11, 10], # B
           [  5, 10,  5]] # C
          #   A   B   C
outputB = [[ 11, 11, 10], # A
           [ 11,  3,  7], # B
           [ 10,  7,  7]] # C
          #   A   B   C
outputC = [[  5, 10,  5], # A
           [ 10,  7,  7], # B
           [  5,  7,  3]] # C

So finally the cube would be the combination of these three arrays stacked along the Z-axis. Given the size, I'm evidently attempting to avoid loops but I'm coming up blank on thinking of efficient, or even just not-super-slow ways to complete this.


Solution

  • Found a fairly quick one-line solution:

    cube_data = df_in.values[:, :, None] + df_in.values[:, None, :] + df_in.values[None, :, :]