4

I have two object arrays not necessarily of the same length: 

import numpy as np

class Obj_A:
    def __init__(self,n):
        self.type = 'a'+str(n)
    def __eq__(self,other):
        return self.type==other.type

class Obj_B:
    def __init__(self,n):
        self.type = 'b'+str(n)
    def __eq__(self,other):
        return self.type==other.type

a = np.array([Obj_A(n) for n in range(2)])
b = np.array([Obj_B(n) for n in range(3)])

I would like to generate the matrix 

mat = np.array([[[a[0],b[0]],[a[0],b[1]],[a[0],b[2]]],
                [[a[1],b[0]],[a[1],b[1]],[a[1],b[2]]]])

this matrix has shape (len(a),len(b),2). Its elements are 

mat[i,j] = [a[i],b[j]]

A solution is

mat = np.empty((len(a),len(b),2),dtype='object')
for i,aa in enumerate(a):
    for j,bb in enumerate(b): 
        mat[i,j] = np.array([aa,bb],dtype='object')

but this is too expensive for my problem, which has O(len(a)) = O(len(b)) = 1e5.

I suspect there is a clean numpy solution involving np.repeat, np.tile and np.transpose, similar to the accepted answer here, but the output in this case does not simply reshape to the desired result.

kevinkayaks
  • 2,636
  • 1
  • 14
  • 30

1 Answers1

4

I would suggest using np.meshgrid(), which takes two input arrays and repeats both along different axes so that looking at corresponding positions of the outputs gets you all possible combinations. For example:

>>> x, y = np.meshgrid([1, 2, 3], [4, 5])
>>> x
array([[1, 2, 3],
       [1, 2, 3]])
>>> y
array([[4, 4, 4],
       [5, 5, 5]])

In your case, you can put the two arrays together and transpose them into the proper configuration. Based on some experimentation I think this should work for you:

>>> np.transpose(np.meshgrid(a, b), (2, 1, 0))
David Z
  • 128,184
  • 27
  • 255
  • 279
  • Thanks David. This runs fast and will scale to the problem sizes I'm interested in. I don't fully understand the (2,1,0). Is there any easy way to see what's happening there? – kevinkayaks Aug 18 '18 at 04:31
  • 1
    That just tells `np.transpose()` how to permute the dimensions of the array. I'd suggest looking at [the documentation](https://docs.scipy.org/doc/numpy/reference/generated/numpy.transpose.html) and playing around with a few examples, e.g. make a 3D array `arr` and try `arr.transpose(0, 2, 1)`, `arr.transpose(2, 0, 1)`, and so on to see how it works. – David Z Aug 18 '18 at 04:33
  • 1
    For a 3d array, (2,1,0) is the default transpose order, so `np.transpose(np.meshgrid(a,b))` should be enough. Another way to express it is `np.array(np.meshgrid(a,b)).T` – hpaulj Aug 18 '18 at 06:07