CODEWITHSUNDEEP
pythonmatplotlib
plaene
plaene

Reputation: 61

Add label to LineCollection plot in loop

I have a dataset from which I need to plot various, discontinuous, lines. Segments that belong to the same index within the dataset will have the same color. I have used the code provided in this topic as the base but due to the datastructure I have placed it in a loop:

import pylab as pl
from matplotlib import collections  as mc
color = ['r','g','b','c','y','k','m']
fig, ax = pl.subplots()
for i in indexset:
     lines=mc.LineCollection(data[i],colors=color[i])
     ax.addcollection(lines)

For each of the sets of lines, I would like to add a label/legend entry with the corresponding index i. Either a legend box with the index i next to the color of the line and/or the index as text within the graph. I tried adding a label=i argument within the loop, however then only the last index actually shows on the plot. How can I add a label for each index?

Upvotes: 1

Views: 2093

Answers (1)

JohanC
JohanC

Reputation: 80339

As your example doesn't contain test data, it is hard to guess the exact goal of the question. Here is an example that creates a legend of each set of lines. Note that data has 4 dimensions (7, 5, 2, 2):

  • the first, 7, for 7 sets of lines
  • the second, 5, telling each set has 5 lines (or curves)
  • the third, 2, telling that each line consists of 2 vertices
  • the fourth, 2, being one for the x and one for the y coordinate

The code uses matplotlib.pyplot, as the use of pylab is discouraged (already in 2012).

from matplotlib import collections  as mc
from matplotlib import pyplot  as plt
import numpy as np

colors = ['r', 'g', 'b', 'c', 'y', 'k', 'm']
data = np.random.randint(1, 10, size=(len(colors), 5, 2, 2))
labels = [f'lines {i} ({c})' for i, c in enumerate(colors, start=1)]
fig, ax = plt.subplots()
for d, col, label in zip(data, colors, labels):
    lines = mc.LineCollection(d, colors=col, label=label)
    ax.add_collection(lines)
ax.legend()
ax.autoscale()
plt.show()

example plot

Here is another example, with data created as:

data = []
for _ in colors:
    t = np.linspace(0, 2 * np.pi, np.random.randint(5, 12), endpoint=False)
    x0, y0 = np.random.uniform(1, 9, 2)
    data.append(np.vstack([x0 + np.cos(t), y0 + np.sin(t), x0 + .1 * np.cos(t), y0 + .1 * np.sin(t)]).T.reshape(-1, 2, 2))

star example

Upvotes: 3

Related Questions