Как добавить цветную строку с дополнительной информацией
Я собираюсь построить график «Путешествие подвески x Расстояние поезда». У меня есть набор данных с этой информацией, а также с информацией о том, движется ли поезд по прямому пути или по кривой. Я хочу имитировать следующее изображение, но мне нужно добавить некоторую дополнительную информацию (например, расположение моста и туннеля). Проблема в том, что то, что я пробовал, занимает почти четыре минуты.
Сюжет, который хочу скопировать
# Plot of the line I want to add
def transition_line(xmin, xmax):
for i in range((df_irv['Distance'] - xmin).abs().idxmin(), (df_irv['Distance'] - xmax).abs().idxmin()):
plt.plot([df_irv['Distance'][i], df_irv['Distance'][i+1]], [max(df_irv['SuspTravel']) + 0.5]*2, color='red' if df_irv['Element'][i] == 'CURVA' else 'blue', linewidth=10, alpha=0.5)
# Function to plot the data with adjustable x-axis limits
def plot_graph(xmin, xmax, sensors):
plt.figure(figsize=(10, 5))
plt.plot(df_irv['Distance'], df_irv[sensors], label='Suspension Sensor')
plt.xlim(xmin, xmax)
plt.xlabel('Distance (Km)')
plt.ylabel('Suspension Sensor')
plt.title('Suspension Sensor vs Distance')
plt.legend()
plt.grid(True)
transition_line(xmin, xmax)
plt.show()
# Create sliders for x-axis limits
xmin_slider = IntSlider(value=0, min=0, max=df_irv['Distance'].max(), step=1, description='X min')
xmax_slider = IntSlider(value=20, min=0, max=df_irv['Distance'].max(), step=1, description='X max')
# Interactive plot
interact(plot_graph, xmin=xmin_slider, xmax=xmax_slider, sensors = ['SuspTravel', 'Roll', 'Bounce'])
Изображение создано моей попыткой
Я попробовал, но это все равно занимает слишком много времени. И если я использую минимальное разрешение для расстояния, анализ будет бесполезен.
если он слишком медленный, возможно, вместо Python придется использовать другой язык.
Ответы 1
Вызов plt.plot() много раз в цикле может быть медленным. (Кроме того, можно избежать вычисления одного и того же max(df_irv['SuspTravel']) для каждого шага цикла, вычислив его один раз перед началом цикла.)
Чтобы ускорить рисование коротких линий, можно использовать аналогичный подход, как в коде разноцветных линий из учебника matplotlib. Вместо цикла массивы numpy работают намного быстрее (массивы реализованы в оптимизированном коде C).
Вот как может выглядеть код:
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
import pandas as pd
import numpy as np
# create some dummy test data
df_irv = pd.DataFrame({'Distance': np.random.randint(10, 100, 1000).cumsum(),
'SuspTravel': np.random.randn(1000).cumsum() * 100,
'Element': np.random.choice(['CURVA', 'RECTA'], 1000, p=[.1, .9])})
fig, ax = plt.subplots()
# add plot
ax.plot('Distance', 'SuspTravel', data=df_irv)
# add "transition line"
xmin = df_irv['Distance'].min()
xmax = df_irv['Distance'].max()
id_xmin = (df_irv['Distance'] - xmin).abs().idxmin()
id_xmax = (df_irv['Distance'] - xmax).abs().idxmin()
xvals = df_irv['Distance'][id_xmin:id_xmax + 1]
yvals = np.full(id_xmax - id_xmin, max(df_irv['SuspTravel']) + 0.5)
colors = df_irv['Element'][id_xmin:id_xmax].map({'CURVA': 'red', 'RECTA': 'blue'})
segments = np.c_[xvals[:-1], yvals, xvals[1:], yvals].reshape(-1, 2, 2)
lines = LineCollection(segments, colors=colors)
lines.set_linewidth(10)
line = ax.add_collection(lines)
plt.show()
Спасибо, попробую и вернусь с отзывом.
возможно, уменьшите данные - не отображайте все, кроме каждого второго или каждого десятого элемента.