Python – Use numpy to generate pairs of matrices from two object vectors

I have two arrays of objects that are not necessarily 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 want to generate a 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]]]])

The shape of this matrix is (len(a), len(b),2). Its elements are

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

One 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, it 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 accepted answers here But the output in this case will not simply reshape the desired result.

I recommend using np.meshgrid(). , which takes two input arrays and repeats along different axes so that looking at the corresponding position of the output can get 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 two arrays together and transpose them to the correct configuration. Based on some experiments, I think this will work for you :

>>> np.transpose(np.meshgrid(a, b), (2, 1, 0))