cartesian products in numPy

Question

Have two arrays, in my case $X = \{1,2,\dots, n\}$ or X = np.arange(n).

How do I get $Y = X \times X = \{ [i,j]: 1 \leq i,j \leq n \}$ as a 2D array in numPy?

In numpy this would be a $n \times n \times 2$ array np.shape(Y)=(n,n,2).

Answer 1

You can use itertools.product:

>>> from numpy import *
>>> import itertools
>>> array([x for x in itertools.product(arange(2), arange(3))])
array([[0, 0],
       [0, 1],
       [0, 2],
       [1, 0],
       [1, 1],
       [1, 2]])

Answer 2

You want meshgrid(). Here is a quick example:

In [14]: a = arange(5)

In [18]: r = array(meshgrid(a,a))

In [20]: r.shape
Out[20]: (2, 5, 5)

In [21]: r
Out[21]: 
array([[[0, 1, 2, 3, 4],
        [0, 1, 2, 3, 4],
        [0, 1, 2, 3, 4],
        [0, 1, 2, 3, 4],
        [0, 1, 2, 3, 4]],

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

Since there is a lot of repeated values in there, you can call it with the sparse=True command to save mem and leverage some ndarray tricks.

If you actually want some simple function of (i,j) you can do the following trick which avoids a formation of an immidiate. I use it all the time but didn't see this in any of the other links:

In [3]: a = arange(5)

In [4]: a[:,newaxis] * a[newaxis,:]
Out[4]: 
array([[ 0,  0,  0,  0,  0],
       [ 0,  1,  2,  3,  4],
       [ 0,  2,  4,  6,  8],
       [ 0,  3,  6,  9, 12],
       [ 0,  4,  8, 12, 16]])

The example uses multiplication, but any number of binary functions can go in that place. This trick also works for higher dimensions thanks to the broadcasting rules:

In [6]: a = ones((2,5))

In [7]: a[:,newaxis] * a[newaxis,:]
Out[7]: 
array([[[ 1.,  1.,  1.,  1.,  1.],
        [ 1.,  1.,  1.,  1.,  1.]],

       [[ 1.,  1.,  1.,  1.,  1.],
        [ 1.,  1.,  1.,  1.,  1.]]])

In [8]: a.shape
Out[8]: (2, 5)

In [9]: b = a[:,newaxis] * a[newaxis,:]

In [10]: b.shape
Out[10]: (2, 2, 5)

Answer 3

A possibility would be broadcast your x as

xx = numpy.zeros(n) + x.reshape(n,1)

to obtain the 'x-grid'. Do the same with the 2nd component and then use numpy's dstack function to tile them in the 3rd dimension.

Answer 4

import numpy as np

def cartesian_cross_product(x,y):
    cross_product = np.transpose([np.tile(x, len(y)),np.repeat(y,len(x))])
    return cross_product