Python: matplotlib: раскраска линий в случайном блуждании по размеру прыжка

Question

У меня есть код, который генерирует и строит случайную прогулку просто отлично. Тем не менее, я хочу раскрасить каждую линию в зависимости от того, насколько она велика. Это мой код:

import matplotlib.pyplot as plt
import numpy as np
import random

def randomwalk(N):
    x, y, bigness = np.zeros((N)), np.zeros((N)), np.zeros((N))
    for n in range(0,N):
        angle = random.random() * 2*np.pi
        jump = np.random.normal(0, 50)

        x[n] = x[n-1] + (np.cos(angle) * jump)
        y[n] = y[n-1] + (np.sin(angle) * jump)
        bigness[n] = abs(jump)      
    return x, y, bigness

x, y, bigness = randomwalk(100)

plt.plot(x, y)
plt.show()

Теперь, если я изменю предпоследнюю строку на

plt.scatter(x, y, c=bigness)

Затем я получаю кучу точек с требуемой окраской, но без линий, соединяющих их. И наоборот, функция "сюжет" не имеет возможности для отдельной раскраски.

Я хочу строки из функции plot, а раскраску из функции scatter. Как мне это сделать?

Answer 1

Вы можете вручную создать LineCollection и цвет каждого сегмента bigness:

import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
pairs = np.array(list(zip(x, y))).reshape(-1, 1, 2)
segments = np.concatenate((pairs[:-1], pairs[1:]), axis=1)
lc = mpl.collections.LineCollection(segments)
lc.set_array(bigness)
f, ax = plt.subplots()
ax.add_collection(lc)
ax.autoscale_view()
plt.show()

Чтобы указать цветовую карту и норму не по умолчанию, передайте cmap а также norm аргументы mpl.collections.LineCollection,

Answer 2

Вместо того, чтобы наносить на график всю линию, нарисуйте ее сегменты цветом, основанным на величине.

import matplotlib.pyplot as plt
import numpy as np
import random

def randomwalk(N):
    x, y, bigness = np.zeros((N)), np.zeros((N)), np.zeros((N))
    for n in range(0,N):
        angle = random.random() * 2*np.pi
        jump = np.random.normal(0, 50)

        x[n] = x[n-1] + (np.cos(angle) * jump)
        y[n] = y[n-1] + (np.sin(angle) * jump)
        bigness[n] = abs(jump)      
    return x, y, bigness

def random_color(bigness):
    return plt.cm.gist_ncar(bigness/100)

x, y, bigness = randomwalk(100)

xy = zip(x,y)
fig, ax = plt.subplots()
for start, stop, b in zip(xy[:-1], xy[1:], bigness):
    x, y = zip(start, stop)
    ax.plot(x, y, color=random_color(b))

plt.show()

Прекрасный грязный результат