heat map using matplotlib

Question

I have a dataset generated in this way:

 aa = linspace(A - 5, A + 5, n_points)
 bb = linspace(B - 1.5, B + 1.5, n_points)
 z = []
 for a in aa:
     for b in bb:
         z.append(cost([a, b]))

I would like and head map where z define the color in the point (a,b) . I need this to analyze local minimum.

I am using matplotlib but I do not know exactly how to proceed.

Answer 1

Typically you'd use imshow or pcolormesh for this.

For example:

import numpy as np
import matplotlib.pyplot as plt

n_points = 10
aa = np.linspace(-5, 5, n_points)
bb = np.linspace(-1.5, 1.5, n_points)

def cost(a, b):
    return a + b

z = []
for a in aa:
    for b in bb:
        z.append(cost(a, b))

z = np.reshape(z, [len(aa), len(bb)])

fig, ax = plt.subplots()
im = ax.pcolormesh(aa, bb, z)
fig.colorbar(im)

ax.axis('tight')
plt.show()

However, it would be better to write your example code as:

import numpy as np
import matplotlib.pyplot as plt

n_points = 10
a = np.linspace(-5, 5, n_points)
b = np.linspace(-1.5, 1.5, n_points)
a, b = np.meshgrid(b, a)

z = a + b # Vectorize your cost function

fig, ax = plt.subplots()
im = ax.pcolormesh(a, b, z)
fig.colorbar(im)

ax.axis('tight')
plt.show()

Or, even more compactly:

import numpy as np
import matplotlib.pyplot as plt

npoints = 10
b, a = np.mgrid[-5:5:npoints*1j, -1.5:1.5:npoints*1j]

z = a + b

fig, ax = plt.subplots()
im = ax.pcolormesh(a, b, z)
fig.colorbar(im)

ax.axis('tight')
plt.show()

Answer 2

I just did something similar, and I used a Scatter plot.

plt.scatter(x_vals, y_vals, s = 100,  c = z_vals, cmap = 'rainbow')
c = plt.colorbar()