Numpy meshgrid 保持底层数组

Question

I want to make a meshgrid, but keep the bottom layer an array.

When I do this:

One = np.array([["1A1","1A2"],["1B1","1B2"]])
Two = np.array([["2A1","2A2"],["2B1","2B2"]])
np.array(np.meshgrid(One, Two))

I get:

array([[['1A1', '1A2', '1B1', '1B2'],
        ['1A1', '1A2', '1B1', '1B2'],
        ['1A1', '1A2', '1B1', '1B2'],
        ['1A1', '1A2', '1B1', '1B2']],

       [['2A1', '2A1', '2A1', '2A1'],
        ['2A2', '2A2', '2A2', '2A2'],
        ['2B1', '2B1', '2B1', '2B1'],
        ['2B2', '2B2', '2B2', '2B2']]], dtype='<U3')

But I don't want to get all the Strings separately. I want the basic structure of ["1A1","1A2"] not to be split up.

Basically, what I am looking for is something, that results in this:

array([[[['1A1', '1A2'],['1B1', '1B2']],
        [['1A1', '1A2'],['1B1', '1B2']]],

       [[['2A1', '2A2'],['2A1', '2A2']],
        [['2B1', '2B2'],['2B1', '2B2']]]], dtype='<U3')

Is there a function or other way to achieve that?

EDIT: The strings are just there to make it easier to understand. I will use this with ints if there are any other methods.

Answer 1

Another approach is to make object dtype arrays with the units you want to keep together as elements:

In [82]: One_ = np.empty(2,object)                                                                   
In [84]: One_[:] = One.tolist()                                                                      
In [85]: One_                                                                                        
Out[85]: array([list(['1A1', '1A2']), list(['1B1', '1B2'])], dtype=object)

In [86]: Two_ = np.empty(2,object)                                                                   
In [87]: Two_[:] = Two.tolist()                                                                      

Now meshgrid can combine them, just as it would with numeric or string dtypes:

In [88]: np.meshgrid(One_, Two_)                                                                     
Out[88]: 
[array([[list(['1A1', '1A2']), list(['1B1', '1B2'])],
        [list(['1A1', '1A2']), list(['1B1', '1B2'])]], dtype=object),
 array([[list(['2A1', '2A2']), list(['2A1', '2A2'])],
        [list(['2B1', '2B2']), list(['2B1', '2B2'])]], dtype=object)]

and as one array:

In [89]: np.stack(_)                                                                                 
Out[89]: 
array([[[list(['1A1', '1A2']), list(['1B1', '1B2'])],
        [list(['1A1', '1A2']), list(['1B1', '1B2'])]],

       [[list(['2A1', '2A2']), list(['2A1', '2A2'])],
        [list(['2B1', '2B2']), list(['2B1', '2B2'])]]], dtype=object)

And conversion back to string array:

In [90]: np.array(_.tolist())                                                                        
Out[90]: 
array([[[['1A1', '1A2'],
         ['1B1', '1B2']],

        [['1A1', '1A2'],
         ['1B1', '1B2']]],


       [[['2A1', '2A2'],
         ['2A1', '2A2']],

        [['2B1', '2B2'],
         ['2B1', '2B2']]]], dtype='<U3')

Or use a structured array view to group the two strings:

In [92]: One.view([('f0','U3',(2,))])                                                                
Out[92]: 
array([[(['1A1', '1A2'],)],
       [(['1B1', '1B2'],)]], dtype=[('f0', '<U3', (2,))])
In [93]: Two.view([('f0','U3',(2,))])                                                                
Out[93]: 
array([[(['2A1', '2A2'],)],
       [(['2B1', '2B2'],)]], dtype=[('f0', '<U3', (2,))])
In [94]: np.meshgrid(_92, _93)                                                                       
Out[94]: 
[array([[(['1A1', '1A2'],), (['1B1', '1B2'],)],
        [(['1A1', '1A2'],), (['1B1', '1B2'],)]],
       dtype=[('f0', '<U3', (2,))]),
 array([[(['2A1', '2A2'],), (['2A1', '2A2'],)],
        [(['2B1', '2B2'],), (['2B1', '2B2'],)]],
       dtype=[('f0', '<U3', (2,))])]
In [95]: np.stack(_)                                                                                 
Out[95]: 
array([[[(['1A1', '1A2'],), (['1B1', '1B2'],)],
        [(['1A1', '1A2'],), (['1B1', '1B2'],)]],

       [[(['2A1', '2A2'],), (['2A1', '2A2'],)],
        [(['2B1', '2B2'],), (['2B1', '2B2'],)]]],
      dtype=[('f0', '<U3', (2,))])
In [96]: import numpy.lib.recfunctions as rf                                                         
In [97]: rf.structured_to_unstructured(_95)                                                          
Out[97]: 
array([[[['1A1', '1A2'],
         ['1B1', '1B2']],

        [['1A1', '1A2'],
         ['1B1', '1B2']]],


       [[['2A1', '2A2'],
         ['2A1', '2A2']],

        [['2B1', '2B2'],
         ['2B1', '2B2']]]], dtype='<U3')

Answer 2

Here is my attempt. I split into four of the "basic structures" that you wanted, applied meshgrid separately to pairs, and then put them all into one array:

import numpy as np

one = np.array(["1A1","1A2"])
two = np.array(["1B1","1B2"])
three = np.array(["2A1","2A2"])
four = np.array(["2B1","2B2"])

y = np.array(np.meshgrid(one,two))
y_2 = np.array(np.meshgrid(three,four))
y_3 = np.array([y,y_2])

print(y_3, type(y_3))

Here is the output. I believe it is an array and the inner elements should be dtype '<U3'.

Answer 3

Your layout is a bit unusual, but this repeated stack seems to do the job:

In [79]: np.stack((np.stack((One,One)), np.stack((Two,Two))))                                        
Out[79]: 
array([[[['1A1', '1A2'],
         ['1B1', '1B2']],

        [['1A1', '1A2'],
         ['1B1', '1B2']]],


       [[['2A1', '2A2'],
         ['2B1', '2B2']],

        [['2A1', '2A2'],
         ['2B1', '2B2']]]], dtype='<U3')

Answer 4

The two comments(@hpaulj, @Richard K Yu) that were made both had a small mistake, but fortunately they made a different mistake, so I could merge their ideas to create a solution. My solution is:

np.array([np.stack((One,One)),np.array(np.meshgrid(*Two))])

However, if someone knows a more elegant solution, I would be very happy, because it is kind of clumpy at the moment.