plotting with meshgrid and imshow

Question

Imshow and meshgrid are not working the way I thought. I have some function defined for a given (x,y) point in 2D that returns a scalar f(x,y). I want to visualize the function f using imshow.

x = np.linspace(0,4)
y = np.linspace(0,1)

X,Y = np.meshgrid(x,y)
Z = np.zeros((50,50))

for i in range(50):
   for j in range(50):
       Z[i,j] = f(X[i,j],Y[i,j])

fig = plt.figure()
plt.imshow(Z,extent=[0,4,1,0])
plt.show()

This works as expected except in the extent I think it should be [0,4,0,1]... Am I defining the Z[i,j] to each (x,y) pair incorrectly? An explanation for how this works would be great! Thanks!

Answer 1

As far as I am aware, the imshow is normally used to display an image. The extent is then used to define how large it should be, say you might want to give an image as the background of the plot.

Instead I think you will find it more intuitive to use pcolor, a demo can be found here. It works much the same as imshow so you can just supply Z. However, you can also give it the X and Y arrays. This way you can really check if your supplying the values correctly:

import matplotlib.pyplot as plt
import numpy as np

x = np.linspace(0,4)
y = np.linspace(0,1)

def f(x, y):
    return y * np.sin(x) 

X, Y = np.meshgrid(x,y)
Z = np.zeros((50,50))

for i in range(50):
   for j in range(50):
       Z[i,j] = f(X[i,j],Y[i,j])

plt.pcolor(X, Y, Z)
plt.show()

I have added a function to show it works. Note that if your function is able to handle numpy arrays you can replace the initialisation of Z and the nested for loops with

X, Y = np.meshgrid(x,y)
Z = f(X, Y)

This is cleaner and will be faster to compute.