generating matrix of pairs from two object vectors using numpy

Question

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.

Answer 1

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))