Mapping z's in numpy array A = [[x0, y0, z0], [x1, y1, z1]] for 3rd column of array B = [[x1, y1, ?], [x0, y0, ?]] based off matching (x,y)?

Viewed 139

I have a numpy array T whose rows have the following column structure: [x, y, value], where x, y, value are integers. A sample T array would look like:

[[1, 0, 4],
 [0, 2, 3],
 [1, 2, 7]]

This data comes from a model where the third column specifies the value of a variable for the tuple (x, y). In the model, this tuple corresponds to a label for the value. For example, my label T_10 (subscript 10) has value 4, T_02 has value 3 and T_12 has value 7.

Now, I want to swap a pair of labels. For example, I want to replace all labels 2 with 1 (and vice versa), to get T_20, T_01 and T_21 for the previous examples respectively. So, this new data is

U = [[2, 0, ?],
     [0, 1, ?],
     [2, 1, ?]]

My issue is that I do not know how to make my new data look like this:

U = [[2, 0, -3]
     [0, 1, -4],
     [2, 1, -7]]

This new data should follow two rules:

First, it should correctly identify the row of T whose first and second columns (x, y) are the same as the new (x, y) in U. For each row of U, if the ordered pair (x, y) = (x, y) of T, then the appropriate '?' value in the third column of U should be the corresponding value in T.

Second: If, on the other hand, (x, y) of U = (y, x) of T, then it should be the negative of the corresponding value.

My attempt involved first extracting the columns of T, and then swapping the pair of labels using the following function:

def swap_indices(a, pair):
    for n, i in enumerate(a):
        if i == pair[0]: # check whether a0's element is = swap element 1
            a[n] = pair[1]
        elif i == pair[1]: 
            a[n] = pair[0]
    return a 

For example, I will swap label 0 with 1 and vice versa for column x and column y using:

pair = (0, 1)
a0 = swap_indices(T[:,0], pair) # column x  
a1 = swap_indices(T[:,1], pair) # column y 

Then I iterate over the number of rows of T; num_rows_of_T:

for k in range(num_rows_of_T):
    temp = np.where((T[k, 0] == a0[k]) & (T[k, 1] == a1[k]) | ((T[k, 0] == a1[k]) & (T[k, 1] == a0[k])))

Above, I am trying to get the index of the row where either (x, y) of U = (x, y) of T, or (x, y) of U = (y, x) of T. However this is where I get stuck. I don't think the above is correct. Also, this approach will not let me apply the second rule where I take the negative of the value of T if (x, y) = (y, x). I also tried using set() for starters (to get an unordered pair), but I cannot correctly find the corresponding value of T even then.

Basically, I want to find the values of T that match with the new labels in U. My data is nice, in that there may only exist one possible set of coordinates, and that there is always a bijective mapping between the (x,y) of T and U (given my two rules).

Any advice? Please help edit the question as necessary. It was very difficult for me to ask.

Here is a minimal working example:

import numpy as np

# swap index labels if match swap pair
def swap_indices(a, pair):
    for n, i in enumerate(a):
        if i == pair[0]: # check whether a0's element is = swap element 1
            a[n] = pair[1]
        elif i == pair[1]: 
            a[n] = pair[0]
    return a
        
def find_valid_swaps(The1 = np.array([1, 0, -1, 1, 0, 1]), headers = np.array(['10', '20', '21', '30', '31', '32'])):
 
    num_indices = len(The1)
    T = np.zeros((num_indices,3)); U = T;
    
    # match format given for T in question
    for i in range(num_indices):
        T[i,:] = [int(list(headers[i])[0]), int(list(headers[i])[1]), The1[i]]
    
    pair = (0, 1) # label pair to swap
    a0 = swap_indices(T[:, 0], pair) # column 0 of U
    a1 = swap_indices(T[:, 1], pair) # column 1 of U
    
    # try to extract correct 'value' from T based on new labels in U
    for k in range(num_indices):
        temp = np.where((T[k, 0] == a0[k]) & (T[k, 1] == a1[k]) | ((T[k, 0] == a1[k]) & (T[k, 1] == a0[k])))
        print("temp",temp[0][0])
        U[k, :] = [a0[k], a1[k], T[temp[0][0], 2]] # here, I would finally create the new U matrix, applying both rules

    print(U)

find_valid_swaps()

More involved example using @MadPhysicist 's answer:

# swap index labels if match swap pair
def swap_indices(a, pair):
    for n, i in enumerate(a):
        if i == pair[0]: # check whether a0's element is = swap element 1
            a[n] = pair[1]
        elif i == pair[1]: 
            a[n] = pair[0]
    return a
    
def key(arr, m):
    return arr[:, 0] * m + arr[:, 1]
    
def find_valid_swaps(Thetas1 = np.array([1, 1, 0, 0, -1, -1]), Thetas2 = np.array([1, 0, -1, 1, 0, 1]), num_bands = 4, headers = np.array(['10', '20', '21', '30', '31', '32'])):
    
    import itertools # for permutations: https://stackoverflow.com/questions/40092474/get-all-pairwise-combinations-from-a-list
    
    if (Thetas1==Thetas2).all():
        print("Warning: Input sets of indices are equal to each other. Will check other possible permutations regardless.")
    else: 
        print("Input sets of indices are unique. Will proceed checking other viable permutations.")

    num_indices = len(Thetas1)
    
    T = np.zeros((num_indices,3))
    U = np.zeros((num_indices,3))
    
    for i in range(num_indices):
        T[i,:] = [int(list(headers[i])[0]), int(list(headers[i])[1]), Thetas2[i]]
    
    print("input T")
    print(T)
    
    pair = (2,3)
    a0 = swap_indices(T[:,0], pair) # column 1  
    a1 = swap_indices(T[:,1], pair) # column 2 
    
    
    for k in range(num_indices):
        U[k, :] = [a0[k], a1[k], 0] 

    # below code due to @MadPhysicist from https://stackoverflow.com/questions/67223782/mapping-zs-in-numpy-array-a-x0-y0-z0-x1-y1-z1-for-3rd-column-of-ar/67235030?noredirect=1#67235030
    
    y_max = T[:, 1].max() + 1
    Tkey = key(T, y_max)
    s = np.argsort(Tkey)

    Ukey = key(U, y_max)
    i = np.searchsorted(Tkey, Ukey, sorter=s)
    i[i == len(i)] -= 1  # cleanup indices that won't match anyway
    mask = (Ukey == Tkey[s[i]])

    U2key = key(U[~mask, 1::-1], y_max)
    j = np.searchsorted(Tkey, U2key, sorter=s)
   
    U[mask, -1] = T[s[i[mask]], -1]
    U[~mask, -1] = -T[s[j], -1]
    
    print("reordered U")
    print(U)

The above gives output:

input T
[[ 1.,  0.,  1.]
 [ 2.,  0.,  0.]
 [ 2.,  1., -1.]
 [ 3.,  0.,  1.]
 [ 3.,  1.,  0.]
 [ 3.,  2.,  1.]]
reordered U
[[ 1.,  0.,  1.]
 [ 3.,  0.,  0.]
 [ 3.,  1., -1.]
 [ 2.,  0.,  1.]
 [ 2.,  1.,  0.]
 [ 2.,  3.,  1.]]
1 Answers

You can boil your algorithm down to three big steps:

  1. Sort Txy
  2. Binary search for Uxy in Txy
  3. Binary search for remaining Uyx in Txy

Combining the result is clearly pretty trivial. The whole operation should be quite doable in O(N log N) time because that's how long each step takes.

Since np.searchsorted is the prime candidate for steps 2 and 3, let's say that you can turn the first two columns into a unique key. For example, let's say that y <= y_max in all cases, and that y_max has a sane bound such that x * y_max + y <= 2**32-1 for all x. You can play with using np.int64 or using x_max instead of y_max at your leisure.

So now you do:

def key(arr, m):
    return arr[:, 0] * m + arr[:, 1]

y_max = T[:, :1].max(None) + 1
Tkey = key(T, y_max)
s = np.argsort(Tkey)

To find which elements of U match:

Ukey = key(U, y_max)
i = np.searchsorted(Tkey, Ukey, sorter=s)
i[i == len(i)] -= 1  # cleanup indices that won't match anyway
mask = (Ukey == Tkey[s[i]])

Now find the reverse index.

U2key = key(U[~mask, 1::-1], y_max)
j = np.searchsorted(Tkey, U2key, sorter=s)

Since the mapping is bijective, this step only searched for elements that are guaranteed to exist, and there's no need to verify the indices.

Now you can combine the indices. If U doesn't already have a third column, add one:

U = np.concatenate((U, np.empty_like(T[:, :1])), axis=1)

Using the indices we computed, extract the elements of Tsort that you want:

U[mask, -1] = T[s[i[mask]], -1]
U[~mask, -1] = -T[s[j], -1]

Now if you can't get a mapping like key working, things can be a bit more complicated. If nothing else works, first try

def key(arr):
    return arr[:, 0] + 1j * arr[:, 1]

The complex values will only be used as sort keys and nothing else. If that fails, you may have to define a structured datatype and view your array through that to get the search working. You can of course implement a hierarchical search, but I feel that that's out of scope here.


Here is a complete toy example based off your T, with a slightly modified U that shows both positive and negative numbers in the last column:

>>> T = np.array([[1, 0, 4],
                  [0, 2, 3],
                  [1, 2, 7]])
>>> U = np.array([[2, 1, 0],
                  [1, 0, 0],
                  [2, 0, 0]])
>>> def key(arr, m):
...     return arr[:, 0] * m + arr[:, 1]

>>> y_max = T[:, :1].max(None) + 1
>>> Tkey = key(T, y_max)
>>> s = np.argsort(Tkey)

>>> Ukey = key(U, y_max)
>>> i = np.searchsorted(Tkey, Ukey, sorter=s)
>>> i[i == len(i)] -= 1  # cleanup indices that won't match anyway
>>> mask = (Ukey == Tkey[s[i]])

>>> U2key = key(U[~mask, 1::-1], y_max)
>>> j = np.searchsorted(Tkey, U2key, sorter=s)

>>> U[mask, -1] = T[s[i[mask]], -1]
>>> U[~mask, -1] = -T[s[j], -1]
>>> print(U)
[[ 2  1 -7]
 [ 1  0  4]
 [ 2  0 -3]]
Related