0

I've got a simple nested list comprehension but can't wrap my head around how I'm supposed to write that into clean numpy code. (Which I have to, for speed improvements)

Here is my list comprehension:

import numpy as np

all_as = np.arange(-1, 1+0.01, 0.01)
all_cs = np.arange(-1, 1+0.01, 0.01)

out = [(a,c) for c in all_cs for a in all_as]

out = np.array(out) # obviously not very efficient...

I'm sure there is a way to not first create a list, but I don't see how.

I know there've been a lot similar questions, but none of them really helped me.

NewNewton
  • 1,015
  • 1
  • 10
  • 22
  • I don't know if it changes anything in term of performance but you can delete one list declaration and do `out = [(a,c) for c in all_as for a in all_as]` which seems equal to the previous expression. – Clément Dec 12 '19 at 09:19
  • Or even [this one](https://stackoverflow.com/questions/11144513/cartesian-product-of-x-and-y-array-points-into-single-array-of-2d-points?noredirect=1&lq=1) that has a pretty good canonical answer (not the checked answer, the second one), and an even faster one by @PaulPanzer below that. – Daniel F Dec 12 '19 at 09:54

2 Answers2

1

You should use product from itertools.

from itertools import product
out = list(product(all_as, all_cs))

Example

>>> data1 = np.arange(0,3)
>>> data2 = np.arange(4,7)

>>> data1
array([0, 1, 2])
>>> data2
array([4, 5, 6])

>>> list(product(data2,data1))
[(4, 0), (4, 1), (4, 2), (5, 0), (5, 1), (5, 2), (6, 0), (6, 1), (6, 2)]

>>> list(product(data1,data2))
[(0, 4), (0, 5), (0, 6), (1, 4), (1, 5), (1, 6), (2, 4), (2, 5), (2, 6)]
abhilb
  • 5,639
  • 2
  • 20
  • 26
1

You can try the following Scikit-Learn utility function.

from sklearn.utils.extmath import cartesian

cartesian([all_as, all_cs])

which returns

array([[-1.  , -1.  ],
       [-1.  , -0.99],
       [-1.  , -0.98],
       ...,
       [ 1.  ,  0.98],
       [ 1.  ,  0.99],
       [ 1.  ,  1.  ]])
Josmoor98
  • 1,721
  • 10
  • 27