"""Generate the initial collection and plot it."""

positions = np.array([0., 1., 2., 3., 5., 8., 13., 21.])

extra_positions = np.array([34., 55., 89.])

orientation = 'horizontal'

lineoffset = 1

linelength = .5

linewidth = 2

color = [1, 0, 0, 1]

linestyle = 'solid'

antialiased = True

coll = EventCollection(positions,

orientation=orientation,

lineoffset=lineoffset,

linelength=linelength,

linewidth=linewidth,

color=color,

linestyle=linestyle,

antialiased=antialiased

)

fig = plt.figure()

ax = fig.add_subplot(1, 1, 1)

ax.add_collection(coll)

ax.set_title('EventCollection: default')

props = {'positions': positions,

'extra_positions': extra_positions,

'orientation': orientation,

'lineoffset': lineoffset,

'linelength': linelength,

'linewidth': linewidth,

'color': color,

'linestyle': linestyle,

'antialiased': antialiased

}

ax.set_xlim(-1, 22)

ax.set_ylim(0, 2)

_, coll, props = generate_EventCollection_plot()

# check that the default segments have the correct coordinates

check_segments(coll,

props['positions'],

props['linelength'],

props['lineoffset'],

props['orientation'])

# check that the default positions match the input positions

np.testing.assert_array_equal(props['positions'], coll.get_positions())

# check that the default orientation matches the input orientation

assert props['orientation'] == coll.get_orientation()

# check that the default orientation matches the input orientation

assert coll.is_horizontal()

# check that the default linelength matches the input linelength

assert props['linelength'] == coll.get_linelength()

# check that the default lineoffset matches the input lineoffset

assert props['lineoffset'] == coll.get_lineoffset()

# check that the default linestyle matches the input linestyle

assert coll.get_linestyle() == [(0, None)]

# check that the default color matches the input color

for color in [coll.get_color(), *coll.get_colors()]:

splt, coll, props = generate_EventCollection_plot()

new_positions = np.hstack([props['positions'], props['extra_positions']])

coll.set_positions(new_positions)

np.testing.assert_array_equal(new_positions, coll.get_positions())

check_segments(coll, new_positions,

props['linelength'],

props['lineoffset'],

props['orientation'])

splt.set_title('EventCollection: set_positions')

splt, coll, props = generate_EventCollection_plot()

new_positions = np.hstack([props['positions'],

props['extra_positions'][0]])

coll.switch_orientation() # Test adding in the vertical orientation, too.

coll.add_positions(props['extra_positions'][0])

coll.switch_orientation()

np.testing.assert_array_equal(new_positions, coll.get_positions())

check_segments(coll,

new_positions,

props['linelength'],

props['lineoffset'],

props['orientation'])

splt.set_title('EventCollection: add_positions')

splt, coll, props = generate_EventCollection_plot()

new_positions = np.hstack([props['positions'],

props['extra_positions'][2]])

coll.append_positions(props['extra_positions'][2])

np.testing.assert_array_equal(new_positions, coll.get_positions())

check_segments(coll,

new_positions,

props['linelength'],

props['lineoffset'],

props['orientation'])

splt.set_title('EventCollection: append_positions')

splt, coll, props = generate_EventCollection_plot()

new_positions = np.hstack([props['positions'],

props['extra_positions'][1:]])

coll.extend_positions(props['extra_positions'][1:])

np.testing.assert_array_equal(new_positions, coll.get_positions())

check_segments(coll,

new_positions,

props['linelength'],

props['lineoffset'],

props['orientation'])

splt.set_title('EventCollection: extend_positions')

splt, coll, props = generate_EventCollection_plot()

new_orientation = 'vertical'

coll.switch_orientation()

assert new_orientation == coll.get_orientation()

assert not coll.is_horizontal()

new_positions = coll.get_positions()

check_segments(coll,

new_positions,

props['linelength'],

props['lineoffset'], new_orientation)

splt.set_title('EventCollection: switch_orientation')

splt.set_ylim(-1, 22)

"""

splt, coll, props = generate_EventCollection_plot()

coll.switch_orientation()

coll.switch_orientation()

new_positions = coll.get_positions()

assert props['orientation'] == coll.get_orientation()

assert coll.is_horizontal()

np.testing.assert_array_equal(props['positions'], new_positions)

check_segments(coll,

new_positions,

props['linelength'],

props['lineoffset'],

props['orientation'])

splt, coll, props = generate_EventCollection_plot()

new_orientation = 'vertical'

coll.set_orientation(new_orientation)

assert new_orientation == coll.get_orientation()

assert not coll.is_horizontal()

check_segments(coll,

props['positions'],

props['linelength'],

props['lineoffset'],

new_orientation)

splt.set_title('EventCollection: set_orientation')

splt.set_ylim(-1, 22)

splt, coll, props = generate_EventCollection_plot()

new_linelength = 15

coll.set_linelength(new_linelength)

assert new_linelength == coll.get_linelength()

check_segments(coll,

props['positions'],

new_linelength,

props['lineoffset'],

props['orientation'])

splt.set_title('EventCollection: set_linelength')

splt, coll, props = generate_EventCollection_plot()

new_lineoffset = -5.

coll.set_lineoffset(new_lineoffset)

assert new_lineoffset == coll.get_lineoffset()

check_segments(coll,

props['positions'],

props['linelength'],

new_lineoffset,

props['orientation'])

splt.set_title('EventCollection: set_lineoffset')

for prop, value, expected in [

('linestyle', 'dashed', [(0, (6.0, 6.0))]),

('linestyle', (0, (6., 6.)), [(0, (6.0, 6.0))]),

('linewidth', 5, 5),

]:

splt, coll, _ = generate_EventCollection_plot()

coll.set(**{prop: value})

assert plt.getp(coll, prop) == expected

splt, coll, _ = generate_EventCollection_plot()

new_color = np.array([0, 1, 1, 1])

coll.set_color(new_color)

for color in [coll.get_color(), *coll.get_colors()]:

np.testing.assert_array_equal(color, new_color)

"""

segments = coll.get_segments()

if (orientation.lower() == 'horizontal'

or orientation.lower() == 'none' or orientation is None):

# if horizontal, the position in is in the y-axis

pos1 = 1

pos2 = 0

elif orientation.lower() == 'vertical':

# if vertical, the position in is in the x-axis

pos1 = 0

pos2 = 1

else:

raise ValueError("orientation must be 'horizontal' or 'vertical'")

# test to make sure each segment is correct

for i, segment in enumerate(segments):

assert segment[0, pos1] == lineoffset + linelength / 2

assert segment[1, pos1] == lineoffset - linelength / 2

assert segment[0, pos2] == positions[i]

col = mcollections.PathCollection([])

col_data_lim = col.get_datalim(mtransforms.IdentityTransform())

assert_array_equal(col_data_lim.get_points(),

# Test if data limits are unchanged by adding an empty collection.

# GitHub issue #1490, pull #1497.

plt.figure()

ax = plt.axes()

coll = ax.scatter([0, 1], [0, 1])

ax.add_collection(coll)

bounds = ax.dataLim.bounds

coll = ax.scatter([], [])

ax = plt.axes()

x, y = np.arange(8), np.arange(10)

u = v = np.linspace(0, 10, 80).reshape(10, 8)

q = plt.quiver(x, y, u, v)

assert q.get_datalim(ax.transData).bounds == (0., 0., 7., 9.)

plt.figure()

ax = plt.axes()

x = np.linspace(-5, 10, 20)

y = np.linspace(-2, 4, 10)

y, x = np.meshgrid(y, x)

trans = mtransforms.Affine2D().translate(25, 32) + ax.transData

plt.quiver(x, y, np.sin(x), np.cos(y), transform=trans)

ax = plt.axes()

x = np.linspace(-5, 10, 20)

y = np.linspace(-2, 4, 10)

y, x = np.meshgrid(y, x)

trans = mtransforms.Affine2D().translate(25, 32) + ax.transData

plt.barbs(x, y, np.sin(x), np.cos(y), transform=trans)

# The calculated bounds are approximately the bounds of the original data,

# this is because the entire path is taken into account when updating the

# datalim.

assert_array_almost_equal(ax.dataLim.bounds, (20, 30, 15, 6),

# Test basic functionality

fig, ax = plt.subplots()

x = np.arange(4)

y = np.arange(3)

X, Y = np.meshgrid(x, y)

XY = np.vstack((X.ravel(), Y.ravel())).T

ww = X / x[-1]

hh = Y / y[-1]

aa = np.ones_like(ww) * 20 # first axis is 20 degrees CCW from x axis

ec = mcollections.EllipseCollection(ww, hh, aa,

units='x',

offsets=XY,

transOffset=ax.transData,

facecolors='none')

ax.add_collection(ec)

from mpl_toolkits.mplot3d import Axes3D

vertsQuad = [

[[0., 0.], [0., 1.], [1., 1.], [1., 0.]],

[[0., 1.], [2., 3.], [2., 2.], [1., 1.]],

[[2., 2.], [2., 3.], [4., 1.], [3., 1.]],

[[3., 0.], [3., 1.], [4., 1.], [4., 0.]]]

fig = plt.figure()

ax = Axes3D(fig)

colors = ['r', 'g', 'b', 'y', 'k']

zpos = list(range(5))

poly = mcollections.PolyCollection(

vertsQuad * len(zpos), linewidth=0.25)

poly.set_alpha(0.7)

# need to have a z-value for *each* polygon = element!

zs = []

cs = []

for z, c in zip(zpos, colors):

zs.extend([z] * len(vertsQuad))

cs.extend([c] * len(vertsQuad))

poly.set_color(cs)

ax.add_collection3d(poly, zs=zs, zdir='y')

# axis limit settings:

ax.set_xlim3d(0, 4)

ax.set_zlim3d(0, 3)

xx, yy = np.mgrid[:10, :10]

xy_points = np.transpose([xx.flatten(), yy.flatten()])

rotations = np.linspace(0, 2*np.pi, len(xy_points))

fig, ax = plt.subplots()

for xy, alpha in zip(xy_points, rotations):

col = mcollections.RegularPolyCollection(

4, sizes=(100,), rotation=alpha,

offsets=[xy], transOffset=ax.transData)

ax.add_collection(col, autolim=True)

# See issue #3860

class SquareCollection(mcollections.RegularPolyCollection):

def __init__(self, **kwargs):

super().__init__(4, rotation=np.pi/4., **kwargs)

def get_transform(self):

"""Return transform scaling circle areas to data space."""

ax = self.axes

pts2pixels = 72.0 / ax.figure.dpi

scale_x = pts2pixels * ax.bbox.width / ax.viewLim.width

scale_y = pts2pixels * ax.bbox.height / ax.viewLim.height

return mtransforms.Affine2D().scale(scale_x, scale_y)

fig, ax = plt.subplots()

xy = [(0, 0)]

# Unit square has a half-diagonal of `1/sqrt(2)`, so `pi * r**2` equals...

circle_areas = [np.pi / 2]

squares = SquareCollection(sizes=circle_areas, offsets=xy,

transOffset=ax.transData)

ax.add_collection(squares, autolim=True)

fig, ax = plt.subplots()

col = ax.scatter([0], [0], [1000], picker=True)

fig.savefig(io.BytesIO(), dpi=fig.dpi)

mouse_event = SimpleNamespace(x=325, y=240)

found, indices = col.contains(mouse_event)

assert found

fig, ax = plt.subplots()

widths, heights, angles = (10, 10), 10, 0

widths = 10, 10

coords = [(10, 10), (15, 15)]

e = mcollections.EllipseCollection(

widths, heights, angles,

units='xy',

offsets=coords,

transOffset=ax.transData)

ax.add_collection(e)

ax.set_xlim(0, 30)

# Should not fail break when faced with a

# non-zero indexed series

index = [11, 12, 13]

ec = fc = pd.Series(['red', 'blue', 'green'], index=index)

lw = pd.Series([1, 2, 3], index=index)

ls = pd.Series(['solid', 'dashed', 'dashdot'], index=index)

aa = pd.Series([True, False, True], index=index)

Collection(edgecolors=ec)

Collection(facecolors=fc)

Collection(linewidths=lw)

Collection(linestyles=ls)

col = mcollections.PathCollection([])

col.set_linestyles(['-', '-'])

col.set_linewidths([1, 2, 3])

assert col.get_linestyles() == [(0, None)] * 6

assert col.get_linewidths() == [1, 2, 3] * 2

col.set_linestyles(['-', '-', '-'])

assert col.get_linestyles() == [(0, None)] * 3

col = mcollections.PathCollection([], capstyle='round')

assert col.get_capstyle() == 'round'

col.set_capstyle('butt')

col = mcollections.PathCollection([], joinstyle='round')

assert col.get_joinstyle() == 'round'

col.set_joinstyle('miter')

fig = plt.figure()

ax = fig.add_subplot(1, 1, 1)

ax.set_xlim([-0.5, 1.5])

ax.set_ylim([-0.5, 2.5])

x = np.array([0.0, 1.0, 0.5])

ys = np.array([[0.0], [0.5], [1.0]]) + np.array([[0.0, 0.0, 1.0]])

segs = np.zeros((3, 3, 2))

segs[:, :, 0] = x

segs[:, :, 1] = ys

line_segments = LineCollection(segs, linewidth=[10, 15, 20])

line_segments.set_capstyle("round")

line_segments.set_joinstyle("miter")

ax.add_collection(line_segments)

fig, ax = plt.subplots()

sc = ax.scatter(range(5), range(5), c=range(5))

# this needs to be here to update internal state

fig.canvas.draw()

np.random.seed(19680801)

x, y = np.random.rand(2, 10)

y = np.random.rand(10)

c = np.random.randint(0, 5, size=10)

s = np.random.randint(10, 300, size=10)

fig, ax = plt.subplots()

sc = ax.scatter(x, y, c=c, s=s, cmap="jet", marker="o", linewidths=0)

h, l = sc.legend_elements(fmt="{x:g}")

assert len(h) == 5

assert_array_equal(np.array(l).astype(float), np.arange(5))

colors = np.array([line.get_color() for line in h])

colors2 = sc.cmap(np.arange(5)/4)

assert_array_equal(colors, colors2)

l1 = ax.legend(h, l, loc=1)

h2, lab2 = sc.legend_elements(num=9)

assert len(h2) == 9

l2 = ax.legend(h2, lab2, loc=2)

h, l = sc.legend_elements(prop="sizes", alpha=0.5, color="red")

alpha = np.array([line.get_alpha() for line in h])

assert_array_equal(alpha, 0.5)

color = np.array([line.get_markerfacecolor() for line in h])

assert_array_equal(color, "red")

l3 = ax.legend(h, l, loc=4)

h, l = sc.legend_elements(prop="sizes", num=4, fmt="{x:.2f}",

func=lambda x: 2*x)

actsizes = [line.get_markersize() for line in h]

labeledsizes = np.sqrt(np.array(l).astype(float)/2)

assert_array_almost_equal(actsizes, labeledsizes)

l4 = ax.legend(h, l, loc=3)

loc = mpl.ticker.MaxNLocator(nbins=9, min_n_ticks=9-1,

steps=[1, 2, 2.5, 3, 5, 6, 8, 10])

h5, lab5 = sc.legend_elements(num=loc)

assert len(h2) == len(h5)

levels = [-1, 0, 55.4, 260]

h6, lab6 = sc.legend_elements(num=levels, prop="sizes", fmt="{x:g}")

assert_array_equal(np.array(lab6).astype(float), levels[2:])

for l in [l1, l2, l3, l4]:

ax.add_artist(l)

# Check that EventCollection doesn't modify input in place

arr = np.array([3, 2, 1, 10])

coll = EventCollection(arr)

verts = np.arange(80, dtype=np.double).reshape(10, 4, 2)

col_arr = PolyCollection(verts)

col_list = PolyCollection(list(verts))

assert len(col_arr._paths) == len(col_list._paths)

for ap, lp in zip(col_arr._paths, col_list._paths):

assert np.array_equal(ap._vertices, lp._vertices)

a = [1, 2, 4]

height = .2

xy_pairs = np.column_stack([np.repeat(a, 2), np.tile([0, height], len(a))])

line_segs = xy_pairs.reshape([len(a), 2, 2])

f, ax = plt.subplots()

trans = mtransforms.blended_transform_factory(ax.transData, ax.transAxes)

ax.add_collection(LineCollection(line_segs, transform=trans))

ax.autoscale_view(scalex=True, scaley=False)

fig, ax = plt.subplots()

ax.scatter(0, 0)

np.testing.assert_allclose(ax.get_ylim(), [-0.06, 0.06])