CODEWITHSUNDEEP
pythonmatplotlibscatter-plot
Nils
Nils

Reputation: 930

Connecting dots in a 2D scatterplot with a color as a third dimension

Let's say I have the following dataset:

x = np.arange(150000,550000,100000)
y = np.random.rand(7*4)
z = [0.6,0.6,0.6,0.6,0.7,0.7,0.7,0.7,0.8,0.8,0.8,0.8,0.9,0.9,0.9,0.9,1.0,1.0,1.0,1.0,1.1,1.1,1.1,1.1,1.2,1.2,1.2,1.2]

x_ = np.hstack([x,x,x,x,x,x,x])

and I am doing a scatter plot:

plt.figure()
plt.scatter(x_,y,c=z)
plt.colorbar()
plt.set_cmap('jet')
plt.xlim(100000,500000)
plt.show()

However, I would like to connect the dots of the same color. I tried just using plt.plot with the same variables, however it connects all the dots instead of just yellow dots with yellow dots.

Any help is appreciated.

EDIT:

The z-axis is discrete and I know the values beforehand. Also I know I could just call the plt.plot() method multiple times to draw the lines but I am hoping that there is maybe another way to do it.

Upvotes: 1

Views: 1375

Answers (2)

ImportanceOfBeingErnest
ImportanceOfBeingErnest

Reputation: 339250

Much easier, using a LineCollection:

import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
import numpy as np; np.random.seed(42)

x = np.arange(150000,550000,100000)
y = np.random.rand(7*4)
z = [0.6,0.6,0.6,0.6,0.7,0.7,0.7,0.7,0.8,0.8,0.8,0.8,0.9,0.9,0.9,0.9,1.0,1.0,1.0,1.0,1.1,1.1,1.1,1.1,1.2,1.2,1.2,1.2]

x_ = np.tile(x, 7)
segs = np.stack((x_, y), axis=1).reshape(7, 4, 2)


plt.figure()
sc = plt.scatter(x_,y,c=z, cmap="plasma")
plt.colorbar()

lc = LineCollection(segs, cmap="plasma", array=np.unique(z))
plt.gca().add_collection(lc)

plt.show()

enter image description here

Upvotes: 1

JohanC
JohanC

Reputation: 80329

If I understand correctly, you have a list of x-values, and for each x some y's are associated, with each x,y having a particular z-value. The z-values belong to a limited set (or can be rounded to make sure there only is a limited set).

So, I created copies of the x, y and z, and sorted them simultaneously via z. Then, looping through the z-array, collect the x,y's, and every time z changes, all lines belonging to that color can be drawn. In order not to need a special step to draw the last set of lines, I append a sentinel.

import numpy as np
import matplotlib.pyplot as plt

x = np.arange(150000,550000,100000)
y = np.random.rand(7*4)
z = [0.6,0.6,0.6,0.6,0.7,0.7,0.7,0.7,0.8,0.8,0.8,0.8,0.9,0.9,0.9,0.9,1.0,1.0,1.0,1.0,1.1,1.1,1.1,1.1,1.2,1.2,1.2,1.2]
z_min = min(z)
z_max = max(z)

x_ = np.hstack([x,x,x,x,x,x,x])

zs = [round(c, 1) for c in z]  # make sure almost equal z-values are exactly equal
zs, xs, ys = zip( *sorted( zip(z, x_, y) ) )  # sort the x,y via z
zs += (1000,) # add a sentinel at the end that can be used to stop the line drawing
xs += (None, )
ys += (None, )

plt.set_cmap('plasma')
cmap = plt.get_cmap()  # get the color map of the current plot call with `plt.get_cmap('jet')`
norm = plt.Normalize(z_min, z_max) # needed to map the z-values between 0 and 1

plt.scatter(x_, y, c=z, zorder=10)  # z-order: plot the scatter dots on top of the lines
prev_x, prev_y, prev_z = None, None, None
x1s, y1s, x2s, y2s = [], [], [], []
for x0, y0, z0 in zip(xs, ys, zs):
    if z0 == prev_z:
        x1s.append(prev_x)
        y1s.append(prev_y)
        x2s.append(x0)
        y2s.append(y0)
    elif prev_z is not None:  # the z changed, draw the lines belonging to the previous z
        print(x1s, y1s, x2s, y2s)
        plt.plot(x1s, y1s, x2s, y2s, color=cmap(norm(prev_z)))
        x1s, y1s, x2s, y2s = [], [], [], []
    prev_x, prev_y, prev_z = x0, y0, z0

plt.colorbar()
plt.show()

Is this what you meant? the resulting plot

Upvotes: 1

Related Questions