matplotlib.lineCollection из pandas кадра данных. Низкая производительность текущего решения iterrows

Question

У меня есть большой фрейм данных, который содержит координаты со значением. Я хочу построить это в matplotlib другим цветом для каждого значения.

У меня есть рабочее решение, которое отображает это как lineCollection. Я использую итерроу, поскольку это легко для меня понять, но это очень медленно.

Я объединяюсь с другим df, который содержит цвет для каждого значения. Затем я через l oop и добавляю к текущему сегменту, если цвет такой же, как предыдущий. Если нет, я начинаю новый сегмент.

dff = df.merge(
        df_color,
        how='left',
        left_on='value',
        right_on='value'
    )

segments = []
segments_i = -1
colorlist = []
previous_color = None

for _, row in dff.iterrows():
        point = (row['x'], row['y'])
        color = row['color']

        if np.any(np.isnan(color)):
            color = (0,0,0,0)

        if color == previous_color:
            segments[segments_i].append(point)
            previous_color = color
        else:
            # add endpoint to current segment
            if segments_i > 0:
                segments[segments_i].append(point)

            # start new segment
            segments.append([point])
            colorlist.append(color)

            previous_color = color
            segments_i += 1

colorlist = np.asarray(colorlist)

lc = mc.LineCollection(segments, colors=colorlist)

ax.add_collection(lc)

Как я могу изменить это, чтобы повысить производительность?

Answer 1

Я изменил это:

for _, row in dff.iterrows():
    point = (row['x'], row['y'])
    color = row['color']

на это:

dff['point'] = list(zip(dff['x'], dff['y']))

for point, color in zip(dff['point'], dff['ctable']):
    ...

Это небольшое изменение заставило его работать в ~ 15 раз быстрее! На данный момент это достаточно хорошо для меня, но не стесняйтесь писать, если у вас есть лучшее решение.