Add values from a dictionary to subarrays using numpy

Question

I'm working with numpy arrays and dictionaries, the keys of the dictionary are coordinates in the numpy array and the values of the dictionary are a list of values I need to add in those coordinates, I also have a 3D list of coordinates that I use for reference, I was able to do it, but I'm creating unnecessary copies of some things to do it. I believe there is an easy way, but I really don't know how to do it, this is my code:

import numpy as np

arr = np.array([[[  0., 448.,  94., 111., 118.],
                 [  0.,   0.,   0.,   0.,   0.],
                 [  0.,   6.,   0.,   6.,   9.],
                 [  0.,  99.,   4.,   0.,   0.],
                 [  0.,  31.,   9.,   0.,   0.]],

                [[  0., 496.,  99.,  41.,  20.],
                 [  0.,   0.,   0.,   0.,   0.],
                 [  0.,  41.,   0.,   1.,   6.],
                 [  0.,  34.,   2.,   0.,   0.],
                 [  0.,  91.,   4.,   0.,   0.]],

                [[  0., 411.,  53.,  75.,  32.],
                 [  0.,   0.,   0.,   0.,   0.],
                 [  0.,  45.,   0.,   3.,   0.],
                 [  0.,  10.,   3.,   0.,   7.],
                 [  0.,  38.,   0.,   9.,   0.]],

                [[  0., 433.,  67.,  57.,  23.],
                 [  0.,   0.,   0.,   0.,   0.],
                 [  0.,  56.,   0.,   4.,   0.],
                 [  0.,   7.,   5.,   0.,   6.],
                 [  0., 101.,   0.,   6.,   0.]]])

#The first list  in reference are the coordinates for the subarray [:,2:,2:] of the first two arrays in arr
#The second list in reference are the coordinates for the subarray [:,2:,2:] of the second two arrays in arr
reference = [[[2, 3], [2, 4], [3, 2], [4, 2]], [[2, 3], [3, 2], [3, 4], [4, 3]]]

#Dictionary whose keys matches the coordinates in the reference list
mydict = {(2, 3): [5, 1], (2, 4): [14, 16], (3, 2): [19, 1], (3, 4): [14, 30], (4, 2): [16, 9], (4, 3): [6, 2]}

#I extract the values of the dict if the key matches the reference and created a 3D list with the values
listvalues = [[mydict.get(tuple(v), v) for v in row] for row in reference] 

#Output
listvalues = [[[5, 1], [14, 16], [19, 1], [16, 9]], [[5, 1], [19, 1], [14, 30], [6, 2]]]

#Then I create a numpy array with my aux list and transpose.
newvalues = np.array(listvalues).transpose(0, 2, 1)

newvalues = [[[ 5, 14, 19, 16],
              [ 1, 16,  1,  9]],

             [[ 5, 19, 14,  6],
              [ 1,  1, 30,  2]]]

What I need is to get a copy of arr (arr shape is (4, 5, 5) and then the copy of arr which I will call newarr will have a shape of (8, 5, 5)) then I need to use the array [5 14 19 16] in newvalues to add the numbers in the corresponding coordinates in the first two arrays of newarr and then the values [5 19 14 6] in the next two arrays in newarr, then (here the copy starts) add the values of [ 1 16 1 9] in the next two arrays of newarr and finally add the values of [ 1 1 30 2] in the final two arrays. Here is the rest of the code.

newarr = np.tile(arr, (2, 1, 1)) #Here I repeat my original array

price = np.reshape(newvalues, (4, 4), order='F') #Here I reshape my 3D array of values to 2D and the order change

final = np.repeat(price, 2, axis =0) #And here I repeat the price so newarr and price have the same dimension in axis = 0

#And finally since they have the dimension in axis = 0 I add the values in the subarray.

index = newarr[:, 2:, 2:] #This is the slice of the subarray
index[index.astype('bool')] = index[index.astype('bool')] + np.array(final).ravel() #And this add values to the right places.

print(newarr)

Output

newarr=[[[  0., 448.,  94., 111., 118.],
         [  0.,   0.,   0.,   0.,   0.],
         [  0.,   6.,   0.,  11.,  23.],
         [  0.,  99.,  23.,   0.,   0.],
         [  0.,  31.,  25.,   0.,   0.]],
                                          #In these two add the values of [5 14 19 16]
        [[  0., 496.,  99.,  41.,  20.],
         [  0.,   0.,   0.,   0.,   0.],
         [  0.,  41.,   0.,   6.,  20.],
         [  0.,  34.,  21.,   0.,   0.],
         [  0.,  91.,  20.,   0.,   0.]],

        [[  0., 411.,  53.,  75.,  32.],
         [  0.,   0.,   0.,   0.,   0.],
         [  0.,  45.,   0.,   8.,   0.],
         [  0.,  10.,  22.,   0.,  21.],
         [  0.,  38.,   0.,  15.,   0.]],
                                           #In these two add the values of [5 19 14 6]
        [[  0., 433.,  67.,  57.,  23.],
         [  0.,   0.,   0.,   0.,   0.],
         [  0.,  56.,   0.,   9.,   0.],
         [  0.,   7.,  24.,   0.,  20.],
         [  0., 101.,   0.,  12.,   0.]],
                                            #<-Here starts the copy of my original array
        [[  0., 448.,  94., 111., 118.],
         [  0.,   0.,   0.,   0.,   0.],
         [  0.,   6.,   0.,   7.,  25.],
         [  0.,  99.,   5.,   0.,   0.],
         [  0.,  31.,  18.,   0.,   0.]],
                                           #In these two add the values of [ 1 16 1  9]
        [[  0., 496.,  99.,  41.,  20.],
         [  0.,   0.,   0.,   0.,   0.],
         [  0.,  41.,   0.,   2.,  22.],
         [  0.,  34.,   3.,   0.,   0.],
         [  0.,  91.,  13.,   0.,   0.]],

        [[  0., 411.,  53.,  75.,  32.],
         [  0.,   0.,   0.,   0.,   0.],
         [  0.,  45.,   0.,   4.,   0.],
         [  0.,  10.,   4.,   0.,  37.],
         [  0.,  38.,   0.,  11.,   0.]],
                                           #And finally in these two add the values of [ 1 1 30 2]
        [[  0., 433.,  67.,  57.,  23.],
         [  0.,   0.,   0.,   0.,   0.],
         [  0.,  56.,   0.,   5.,   0.],
         [  0.,   7.,   6.,   0.,  36.],
         [  0., 101.,   0.,   8.,   0.]],

I mean it does what I need, but like I said, I think there are some unnecessary copies that I don't need, and it's ugly code, I believe there should be an easy way, exploiting the possibilities of the dictionary and the numpy array, but I just can't see it. Any help will be appreciated, this is just an example to see what's going on, but the arr can have more arrays and the list values of the the dictionary can be bigger.

Answer 1

This seems to so what you want, but is specific to the example and code you gave.

There are two pairs of subarrays in arr and two different sets of indices and data to add to the subarrays. So there are four combinations. These get figured out by the values of i, j, and 'k'.

Because the data to be added is sparse, I'm going to use scipy.sparse.coo_matrix() to build arrays from reference and mydict.

The line data = ... converts the information in mydict and reference into a list of three tuples. data[0] is the values to be added, data[1] are the row coordinates, and data[2] are the col coordinates.

m = coo_matrix(...) builds the sparse matrix and converts it to an numpy.array.

x = arr[2*j:2*j+2] + m uses the numpy array broadcasting rules to add m to the subarrays of the arr slice. So x is a pair of subarrays with the values added to the selected coordinates.

All of the x arrays are gathered in a list newarr, and are vertically stacked at the end.

import numpy as np
from scipy.sparse import coo_matrix

newarr = []

for k in range(4):
    i,j = divmod(k,2)
    
    data = [*zip(*((mydict[tuple(coord)][i], *coord) for coord in reference[j]))]

    m = coo_matrix((data[0],(data[1], data[2]))).toarray()
    
    x = arr[2*j:2*j+2] + m

    newarr.append(x)
    
newarr = np.vstack(newarr)