Create Image using Matplotlib imshow meshgrid and custom colors

Question

I am trying to create an image where the x axis is the width, and y axis is the height of the image. And where each point can be given a color based on a RBG mapping. From looking at imshow() from Matplotlib I guess I need to create a meshgrid on the form (NxMx3) where 3 is a tuple or something similar with the rbg colors.

But so far I have not managed to understand how to do that. Lets say I have this example:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap

x_min = 1 
x_max = 5
y_min = 1 
y_max = 5
Nx = 5 #number of steps for x axis
Ny = 5 #number of steps for y axis

x = np.linspace(x_min, x_max, Nx)
y = np.linspace(y_min, y_max, Ny)

#Can then create a meshgrid using this to get the x and y axis system
xx, yy = np.meshgrid(x, y)

#imagine I have some funcion that does someting based on the x and y values
def somefunc(x_value, y_value):
    #do something and return rbg based on that
    return x_value + y_value

res = somefunc(xx, yy)

cmap = LinearSegmentedColormap.from_list('mycmap', ['white', 'blue', 'black'])
plt.figure(dpi=100)
plt.imshow(res, cmap=cmap, interpolation='bilinear')

plt.show()

And this creates a plot, but what would I have to do if my goal was to give spesific rbg values based on x and y values inside somefunc and make the resulting numpy array into a N x M x 3 array

I tried to make the somefunc function return a tuple of rbg values to use (r, b g) but that does not seem to work

Answer 1

It will of course completely depend on what you want to do with the values you supply to the function. So let's assume you just want to put the x values as the red channel and the y values as the blue channel, this could look like

def somefunc(x_value, y_value):
    return np.dstack((x_value/5., np.zeros_like(x_value), y_value/5.))

Complete example:

import numpy as np
import matplotlib.pyplot as plt

x_min = 1 
x_max = 5
y_min = 1 
y_max = 5
Nx = 5 #number of steps for x axis
Ny = 5 #number of steps for y axis

x = np.linspace(x_min, x_max, Nx)
y = np.linspace(y_min, y_max, Ny)

#Can then create a meshgrid using this to get the x and y axis system
xx, yy = np.meshgrid(x, y)

#imagine I have some funcion that does someting based on the x and y values
def somefunc(x_value, y_value):
    return np.dstack((x_value/5., np.zeros_like(x_value), y_value/5.))


res = somefunc(xx, yy)

plt.figure(dpi=100)
plt.imshow(res)

plt.show()

If you already have a (more complicated) function that returns an RGB tuple you may loop over the grid to fill an empty array with the values of the function.

#If you already have some function that returns an RGB tuple
def somefunc(x_value, y_value):
    if x_value > 2 and y_value < 3:
        return np.array(((y_value+1)/4., (y_value+2)/5., 0.43))
    elif x_value <=2:
        return np.array((y_value/5., (x_value+3)/5., 0.0))
    else:
        return np.array((x_value/5., (y_value+5)/10., 0.89))
# you may loop over the grid to fill a new array with those values
res = np.zeros((xx.shape[0],xx.shape[1],3))
for i in range(xx.shape[0]):
    for j in range(xx.shape[1]):
        res[i,j,:] = somefunc(xx[i,j],yy[i,j])



plt.figure(dpi=100)
plt.imshow(res)