CODEWITHSUNDEEP
pythonmatplotlib
mauve
mauve

Reputation: 2763

Add colorbar to scatter plot or change the plot type

I am plotting some data that includes spatial (x, y) components as well as a z component, which is the value of the measurement at that point in space. I was looking at the gallery, and I'm just not getting it. I think that what I want is a pcolormesh, but I don't understand what I need to put in for arguments. I finally had success getting a scatter plot to do basically what I want, but it's less pretty than I want. If I could figure out a way to make the points in the scatter plot bigger, I would be a lot happier with my plot. Furthermore, I am stuck on trying to add a legend - I only need the colorbar portion, since the end user doesn't really care about the X and Y dimensions. Looking at the colorbar example, it seems that I need to add an axis, but I don't understand how I'm telling it that the axis I need is the Z axis. 

x_vals = list(first_array[data_loc_dictionary['x_coord_index']][:])
y_vals = list(first_array[data_loc_dictionary['y_coord_index']][:])
y_vals = [-i for i in y_vals]
z_vals = list(first_array[data_loc_dictionary['value_index']][:])

plt.scatter(x_vals, y_vals, s = len(x_vals)^2, c = z_vals, cmap = 'rainbow')
plt.show()

Here is an example of what I am trying to duplicate: enter image description here And here is what the code above produces: enter image description here

Good catch with the ^2 - enter image description here

Upvotes: 0

Views: 819

Answers (2)

mwaskom
mwaskom

Reputation: 49032

In case you want to know how to replicate your initial example image with pcolormesh, I would do:

import numpy as np
import matplotlib.pyplot as plt

f, ax = plt.subplots(figsize=(6, 5))
grid = np.arange(-5, 6)
x, y = np.meshgrid(grid, grid)

z = np.random.randn(len(x), len(y)) 
mask = (np.abs(x) + np.abs(y)) > 4
z = np.ma.masked_array(z, mask)

mesh = ax.pcolormesh(x - .5, y - .5, z, cmap="coolwarm", vmin=-3, vmax=3)
plt.colorbar(mesh)

To produce:

enter image description here

Upvotes: 1

wflynny
wflynny

Reputation: 18541

What about this basic example:

# generate random data
In [63]: x = np.random.rand(20)
In [64]: y = np.random.rand(20)
In [65]: z = np.random.rand(20)

# plot it with square markers: marker='s'
In [66]: plt.scatter(x, y, s=len(x)**2, c=z, cmap='rainbow', marker='s')
Out[66]: <matplotlib.collections.PathCollection at 0x39e6c90>

# colorbar
In [67]: c = plt.colorbar(orientation='horizontal')
In [68]: c.set_label('This is a colorbar')
In [69]: plt.show()

enter image description here

The Size of the points is given by

s : scalar or array_like, shape (n, ), optional, default: 20

    size in points^2.

I see no reason why s=len(x)**2 is a good choice by default. I would play around with it according to your preference.

Upvotes: 2

Related Questions