Assign same value to different ndarray indices - Python

Question

The question may not be clear, so I you explain with the code:

@staticmethod
def generate_data_uncertainty_heat_map(data_property, data_uncertainty, file_path):
    plt.figure()
    uncertainty = numpy.zeros((data_property.rows, data_property.cols, 4), 'uint8')
    uncertainty[..., 0] = uncertainty[..., 1] = uncertainty[..., 2] = numpy.uint8(data_property * 255)
    uncertainty[..., 3] = numpy.uint8(data_uncertainty)
    fig = plt.imshow(uncertainty, extent=[0, data_property.cols, data_property.rows, 0])
    plt.colorbar(fig)
    plt.savefig(file_path + '.png')
    plt.close()

What this does is take two (n,n) ndarrays to form an RGBA image with matploblib. To do such thing, I use the data_property parameter as my RGB and the data_uncertainty as my opacity.

I basically would like to know if I can write

uncertainty[..., 0] = uncertainty[..., 1] = uncertainty[..., 2]

in another way, saying data for uncertainty[..., 0 or 1 or 2] the value should be numpy.uint8(data_property * 255).

Thank you in advance.

Answer 1

You can add a new axis on the right using data_property[..., np.newaxis]. Then the assignment can be done like this:

uncertainty[..., :3] = numpy.uint8(data_property * 255)[..., np.newaxis]

---

For example,

In [49]: x = np.arange(6).reshape(2,3)

In [50]: x
Out[50]: 
array([[0, 1, 2],
       [3, 4, 5]])

In [51]: y = np.zeros((2,3,4))

In [52]: y[...,:3] = x[...,np.newaxis]

In [53]: y
Out[53]: 
array([[[ 0.,  0.,  0.,  0.],
        [ 1.,  1.,  1.,  0.],
        [ 2.,  2.,  2.,  0.]],

       [[ 3.,  3.,  3.,  0.],
        [ 4.,  4.,  4.,  0.],
        [ 5.,  5.,  5.,  0.]]])

In [54]: np.allclose(y[...,0], x)
Out[54]: True

In [55]: np.allclose(y[...,1], x)
Out[55]: True

In [56]: np.allclose(y[...,2], x)
Out[56]: True

Answer 2

Yes you can, as long as data_property has the same shape as uncertainty[..., 0]