Numpy broadcasting to the 4th dimension: … vs. : v

Question

In a montecarlo simulation I have the following 7 pokercards for 2 players and 3 different montecarlo runs.

self.cards:

array([[[  6.,  12.],
    [  1.,   6.],
    [  3.,   3.],
    [  8.,   8.],
    [  1.,   1.],
    [  4.,   4.],
    [  2.,   2.]],

   [[  6.,   7.],
    [  1.,   1.],
    [  3.,   3.],
    [  2.,   2.],
    [ 12.,  12.],
    [  5.,   5.],
    [ 10.,  10.]],

   [[  6.,   3.],
    [  1.,  11.],
    [  2.,   2.],
    [  6.,   6.],
    [ 12.,  12.],
    [  6.,   6.],
    [  7.,   7.]]])

The corresponding suits are:

self.suits

array([[[ 2.,  1.],
    [ 1.,  2.],
    [ 2.,  2.],
    [ 2.,  2.],
    [ 1.,  1.],
    [ 2.,  2.],
    [ 2.,  2.]],

   [[ 2.,  0.],
    [ 1.,  3.],
    [ 2.,  2.],
    [ 0.,  0.],
    [ 1.,  1.],
    [ 1.,  1.],
    [ 1.,  1.]],

   [[ 2.,  2.],
    [ 1.,  0.],
    [ 3.,  3.],
    [ 2.,  2.],
    [ 1.,  1.],
    [ 1.,  1.],
    [ 1.,  1.]]])

Now I would like to 'merge' the arrays in a way that the cards array is expanded to the 4th dimension having a size of 4: 0 containing all suits==1, 1 all suits==2, 2 all suits==3 and 3 all suits ==4

I can easily create 4 different arrays:

club_cards=(self.suits == 1) * self.cards 
diamond_cards=(self.suits == 2) * self.cards 
heart_cards=(self.suits == 3) * self.cards 
spade_cards=(self.suits == 4) * self.cards

and then stack them together:

stacked_array=np.stack((club_cards,diamond_cards, heart_cards, spade_cards),axis=0)

The result as expected has a shape of (4, 3, 8, 2)

array([[[[  1.,  12.],
         [  1.,   1.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [-11.,   0.]],

        [[ 12.,  12.],
         [ 10.,  10.],
         [  5.,   5.],
         [  1.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.]],

        [[ 12.,  12.],
         [  7.,   7.],
         [  6.,   6.],
         [  1.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.]]],


       [[[  8.,   8.],
         [  6.,   6.],
         [  4.,   4.],
         [  3.,   3.],
         [  2.,   2.],
         [  0.,   0.],
         [  0.,   0.],
         [ -4.,  -4.]],

        [[  6.,   3.],
         [  3.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [ -6.,  -9.]],

        [[  6.,   6.],
         [  6.,   3.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [ -6.,  -6.]]],


       [[[  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [-12., -12.]],

        [[  0.,   1.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [-12., -11.]],

        [[  2.,   2.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [-10., -10.]]],


       [[[  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [-12., -12.]],

        [[  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [-12., -12.]],

        [[  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [  0.,   0.],
         [-12., -12.]]]])

While this might make sense in the above case, it is not always possible, especially if there are more than 4 cases that need to be stacked together, which brings me to my question:

How can I do this with broadcasting? Below my specific questions:

I have tried a few things.

1. Let's focus on the first step to get the booleans in doing suits==np.arange(4) (the second step is just a multiplication with the cards which will need to be broadcast in the same way as the suits). My understanding is that we want to add a dimension for the suits array, so shouldn't we signal this with the 3 dot notation: self.suits[...,:,:,:]==np.arange(4)? Instead the following seems to almost work: self.suits[:,:,:,None]==np.arange(4) (except that it adds the dimension at the wrong place). The following doesn't work either: self.suits[None,:,:,:]==np.arange(4). How can I extend the array in the first dimension so the results are the same as in the above stack?

2. In what circumstances do I need the ... and when the None? I would expect to use the ... as this would signal that this dimension needs to be expanded as necessary (in this case to a size of 4)? Why does this seem to be incorrect and a None is used instead?

Answer 1

You are stacking the indivdual card results along axis=0. So, when porting to a broadcasting based solution, we can create a range array of those scalars 1, 2, 3, 4 in a 4D array with all axes being singleton dimensions (dims with length = 1) except the first one. There could be different ways to create such a 4D array. One way would be : np.arange(1,5)[:,None,None,None], where we create a 1D array with np.arange and simply add three singleton dims as the last three ones with np.newaxis/None.

We perform equality comparison with this 4D array against b, which would allow internally broadcasting of b elements along the last three dims. Then, we multiply it with a as also done in the original code and get the desired output.

Thus, the implementation would be -

out = a*(b == np.arange(1,5)[:,None,None,None]) 

When/how to use ...(ellipsis) :

We use ...(ellipsis), when trying to add new axes into a multi-dimensional array and we don't want to specify colons per dim. Thus, to make a a 4D array with the last dim being a singleton, we would do : a[:,:,:,None]. Too much of typing! So, we use ... there to help us out : a[...,None]. Please note that this ... notation is used irrespective of the number of dimensions. So, if a were a 5D array and to add a new axis into it as the last one, we would do a[:,:,:,:,:,None] or simply with ellipsis : a[...,None]. Neat huh!