import numpy as np
from ..util.log import Handle
from ..util.plot.axes import init_axes
from ..util.plot.style import linekwargs, scatterkwargs
logger = Handle(__name__)
[docs]def stem(x, y, ax=None, orientation="horizontal", **kwargs):
"""
Create a stem (or 'lollipop') plot, with optional orientation.
Parameters
-----------
x, y : :class:`numpy.ndarray`
1D arrays for independent and dependent axes.
ax : :class:`matplotlib.axes.Axes`, :code:`None`
The subplot to draw on.
orientation : :class:`str`
Orientation of the plot (horizontal or vertical).
Returns
-------
:class:`matplotlib.axes.Axes`
Axes on which the stem diagram is plotted.
"""
ax = init_axes(ax=ax, **kwargs)
orientation = orientation.lower()
xs, ys = np.array([x, x]), np.array([np.zeros_like(y), y])
positivey = (y > 0 | ~np.isfinite(y)).all() | np.allclose(y, 0)
kwargs["color"] = kwargs.get("color") or "0.5"
if "h" in orientation:
ax.plot(xs, ys, **linekwargs(kwargs))
ax.scatter(x, y, **scatterkwargs(kwargs))
if positivey:
ax.set_ylim(0, ax.get_ylim()[1])
else:
ax.plot(ys, xs, **linekwargs(kwargs))
ax.scatter(y, x, **scatterkwargs(kwargs))
if positivey:
ax.set_xlim(0, ax.get_xlim()[1])
return ax