Skip to content

Instantly share code, notes, and snippets.

@trcook
Last active April 18, 2020 11:03
Show Gist options
  • Save trcook/a60ef43d44dc7ef197601b9ff72dd9d8 to your computer and use it in GitHub Desktop.
Save trcook/a60ef43d44dc7ef197601b9ff72dd9d8 to your computer and use it in GitHub Desktop.
get expand.grid -like functionality in python (3) -- expanded from
# to get expand-grid like functionality in python 3
# from http://stackoverflow.com/questions/12864445/numpy-meshgrid-points
import numpy as np
def expandgrid(*arrs):
arrs = tuple(reversed(arrs))
lens =[i for i in map(len, arrs)]
dim = len(arrs)
sz = 1
for s in lens:
sz *= s
ans = []
for i, arr in enumerate(arrs):
slc = [1]*dim
slc[i] = lens[i]
arr2 = np.asarray(arr).reshape(slc)
for j, sz in enumerate(lens):
if j != i:
arr2 = arr2.repeat(sz, axis=j)
ans.append(arr2)
out=np.transpose(np.vstack(map(np.ravel, tuple(ans))))
return out
# alternatively:
import itertools as it
[i for i in it.product(["a","b","c"],[1,2,3])]
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment