Using numpy to build an array of all combinations of two arrays
Cartesian product function by Stackoverflow user, pv.
See Answer here.
Created
June 8, 2014 19:10
-
-
Save hernamesbarbara/68d073f551565de02ac5 to your computer and use it in GitHub Desktop.
A numpy implementation ~5× faster than using itertools.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import numpy as np | |
| def cartesian(arrays, out=None): | |
| """ | |
| Generate a cartesian product of input arrays. | |
| Parameters | |
| ---------- | |
| arrays : list of array-like | |
| 1-D arrays to form the cartesian product of. | |
| out : ndarray | |
| Array to place the cartesian product in. | |
| Returns | |
| ------- | |
| out : ndarray | |
| 2-D array of shape (M, len(arrays)) containing cartesian products | |
| formed of input arrays. | |
| Examples | |
| -------- | |
| >>> cartesian(([1, 2, 3], [4, 5], [6, 7])) | |
| array([[1, 4, 6], | |
| [1, 4, 7], | |
| [1, 5, 6], | |
| [1, 5, 7], | |
| [2, 4, 6], | |
| [2, 4, 7], | |
| [2, 5, 6], | |
| [2, 5, 7], | |
| [3, 4, 6], | |
| [3, 4, 7], | |
| [3, 5, 6], | |
| [3, 5, 7]]) | |
| """ | |
| arrays = [np.asarray(x) for x in arrays] | |
| dtype = arrays[0].dtype | |
| n = np.prod([x.size for x in arrays]) | |
| if out is None: | |
| out = np.zeros([n, len(arrays)], dtype=dtype) | |
| m = n / arrays[0].size | |
| out[:,0] = np.repeat(arrays[0], m) | |
| if arrays[1:]: | |
| cartesian(arrays[1:], out=out[0:m,1:]) | |
| for j in xrange(1, arrays[0].size): | |
| out[j*m:(j+1)*m,1:] = out[0:m,1:] | |
| return out |
Author
👍 that's awesome.
A loop-based version that can be jitted with numba
@njit(cache=True) def cartesian_jit(arrays): """ Generate a cartesian product of input arrays. Parameters ---------- arrays : list or tuple of arrays 1-D arrays to form the cartesian product of. Returns ------- out : ndarray 2-D array of shape (M, len(arrays)) containing cartesian products formed of input arrays. Examples -------- >>> cartesian(([1, 2, 3], [4, 5], [6, 7])) array([[1, 4, 6], [1, 4, 7], [1, 5, 6], [1, 5, 7], [2, 4, 6], [2, 4, 7], [2, 5, 6], [2, 5, 7], [3, 4, 6], [3, 4, 7], [3, 5, 6], [3, 5, 7]]) """ n = 1 for x in arrays: n *= x.size out = np.zeros((n, len(arrays))) for i in range(len(arrays)): m = int(n / arrays[i].size) out[:n, i] = np.repeat(arrays[i], m) n //= arrays[i].size n = arrays[-1].size for k in range(len(arrays)-2, -1, -1): n *= arrays[k].size m = int(n / arrays[k].size) for j in range(1, arrays[k].size): out[j*m:(j+1)*m,k+1:] = out[0:m,k+1:] return out
It doesn't work,
python arrays don't have .size attribute.
@aero-108 Yes, numpy arrays do have a .size attribute...
a = np.array([1,2,3])
a.size
Out[2]: 3
A loop-based version that can be jitted with numba
@njit(cache=True) def cartesian_jit(arrays): """ Generate a cartesian product of input arrays. Parameters ---------- arrays : list or tuple of arrays 1-D arrays to form the cartesian product of. Returns ------- out : ndarray 2-D array of shape (M, len(arrays)) containing cartesian products formed of input arrays. Examples -------- >>> cartesian(([1, 2, 3], [4, 5], [6, 7])) array([[1, 4, 6], [1, 4, 7], [1, 5, 6], [1, 5, 7], [2, 4, 6], [2, 4, 7], [2, 5, 6], [2, 5, 7], [3, 4, 6], [3, 4, 7], [3, 5, 6], [3, 5, 7]]) """ n = 1 for x in arrays: n *= x.size out = np.zeros((n, len(arrays))) for i in range(len(arrays)): m = int(n / arrays[i].size) out[:n, i] = np.repeat(arrays[i], m) n //= arrays[i].size n = arrays[-1].size for k in range(len(arrays)-2, -1, -1): n *= arrays[k].size m = int(n / arrays[k].size) for j in range(1, arrays[k].size): out[j*m:(j+1)*m,k+1:] = out[0:m,k+1:] return out
Is it possible to make a generator function using Numba?
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
A loop-based version that can be jitted with numba