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)
.
Computational Science Stack Exchange is a question and answer site for scientists using computers to solve scientific problems. It only takes a minute to sign up.
Sign up to join this communityHave 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)
.
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]])
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)
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