def imshow(img: np.ndarray, axis: str = 'on', ax: plt.Axes = None, add_colorbar: bool = True,
**kwargs) -> (plt.Figure, plt.Axes):
"""Imshow wrapper to plot images in various image spaces. Constructs a figure object under the hood.
Arguments
---------
img: the input image
axis: whether to display the axis with ticks and everything by default
ax: if None, create new plt Axes, otherwise take this one for plotting
kwargs: those will be forwarded into the setup_plt_figure function
Returns
-------
fig, ax: a tuple of a matplotlib Figure and Axes object
"""
if ax is None:
fig, ax = setup_plt_figure(**kwargs)
else:
fig = ax.get_figure()
ax.set_title(kwargs.get('title', ''))
ax.axis(axis)
imshow_kwargs = {'cmap': kwargs['cmap'] if 'cmap' in kwargs else 'hot'}
for key in ['vmin', 'vmax']:
if key in kwargs:
imshow_kwargs[key] = kwargs.get(key)
im = ax.imshow(img, **imshow_kwargs)
if add_colorbar:
divider = make_axes_locatable(ax)
cax = divider.append_axes("right", size="2%", pad=0.05)
plt.colorbar(im, cax=cax)
font = {'family': 'normal',
'weight': 'bold',
'size': 12}
matplotlib.rc('font', **font)
return fig, ax
def __init__(self, ax: plt.Axes = None, **fig_kw):
if ax is None:
fig, ax = plt.subplots(nrows=1, ncols=1, **fig_kw)
self._fig = fig
self._ax = ax
else:
self._ax = ax
self._fig = ax.get_figure()
self._plot_data: pd.DataFrame = pd.DataFrame()
def _setup_plot(ax: plt.Axes = None, title: str = None) -> (plt.Figure, plt.Axes):
"""Initial setup of fig and ax to plot to."""
if not ax:
fig = plt.figure() # type: plt.Figure
ax = fig.add_subplot(1, 1, 1) # type: plt.Axes
if title:
ax.set_title(title)
return ax.get_figure(), ax
def _plotter(self,
df: DataFrame,
ax_case: plt.Axes,
ax_death: plt.Axes,
cases_label: str,
death_label: str,
case_color='b',
death_color='r') -> None:
x_dates = df['date'].values
y_cases = df['cases'].values
y_deaths = df['deaths'].values
ax_case.plot(x_dates, y_cases, c=case_color, label=cases_label)
ax_case.set_xlabel('Date')
ax_case.legend()
ax_case.xaxis_date()
ax_death.plot(x_dates, y_deaths, c=death_color, label=death_label)
ax_death.set_xlabel('Date')
ax_death.legend()
ax_death.xaxis_date()
ax_death.get_figure().autofmt_xdate()
def plot_matrix(ax: plt.Axes,
mtx,
labels,
title: str,
xlabel: str,
ylabel: str,
cmap='Pastel2'):
im = ax.imshow(mtx, interpolation='nearest', cmap=cmap)
# ax.figure.colorbar(im, ax=ax)
ax.set_xticks(np.arange(mtx.shape[1]))
ax.set_yticks(np.arange(mtx.shape[0]))
ax.set_xticklabels(labels)
ax.set_yticklabels(labels)
ax.set_title(title)
ax.set_xlabel(xlabel)
ax.set_ylabel(ylabel)
ax.set_autoscale_on(False)
# plt.setp(ax.get_xticklabels(), rotation=45, ha='right', rotation_mode='anchor')
for i in range(mtx.shape[0]):
for j in range(mtx.shape[1]):
ax.text(j, i, format(mtx[i, j], 'd'), ha='center', va='center')
ax.set_ylim(len(mtx) - .5, -.5)
ax.get_figure().tight_layout()
def __init__(
self,
ax: plt.Axes,
):
"""
:param ax: xxx
"""
canvas = ax.get_figure().canvas
canvas.mpl_connect("button_press_event", self._on_press)
canvas.mpl_connect("button_release_event", self._on_release)
canvas.mpl_connect("scroll_event", self._on_scroll_event)
self.ax = ax
self.canvas = canvas
self.last_point: Optional[Tuple[float, float]] = None
self.motion_handler: Optional[Callable] = None
def plot_trajectories_movie(trajectories: Sequence[np.array], dimension_description: DimensionDescription = "x,z",
delay: int = 100, n_previous: int = 10, ax: plt.Axes = None,
animation_kwargs: Dict = None, plot_kwargs: Dict = None,
update_fn: Callable[[int, List[plt.Artist]], List[plt.Artist]] = None)\
-> (Animation, plt.Figure, int):
"""
Creates a trajectory movie via a :class:`matplotlib.animation.Animation`. Plots each trajectory as a curve
of the last ``n_previous`` positions.
:param trajectories: A sequence of trajectories.
:param dimension_description: A dimension descriptionf of a pair or triple of dimensions.
:param delay: The delay between each animation frame.
:param n_previous: The number of previous positions to include in each curve. When equal to one, each curve
:param ax: (optional) a specific Axes instance to plot on. If not passed, a new Figure and Axes are created.
:param animation_kwargs: (optional) arguments to pass to :class:`matplotlib.animation.FuncAnimation`
:param plot_kwargs: (optional) arguments to pass to the initial plt.plot() calls that draw each line
:param update_fn: (optional) An update function to be called as the last step in each animation iteration.
Receives the current iteration number and the list of artists that the default update function returns.
Must return a list of artists, just like animation update functions.
:return: A 3-tuple of (created Animation instance, created Figure instance, number of animation frames).
"""
extractor, n = parse_multiple_dimensions_description(dimension_description, n=(2, 3))\
if type(dimension_description) is str else dimension_description
positions = [extractor(traj) for traj in trajectories]
maxlen = max(len(traj) for traj in trajectories)
if ax is None:
fig = plt.figure()
ax = fig.add_subplot(111) if n == 2 else fig.add_subplot(
111, projection='3d')
else:
fig = ax.get_figure()
empty_plot_arg = [[]] * n
lines = [
plt.plot(*empty_plot_arg, **(plot_kwargs or {}))[0]
for _ in range(len(trajectories))
]
labels = dimension_description.split(',')
for label_setter, label in zip(('set_xlabel', 'set_ylabel', 'set_zlabel'),
labels):
getattr(ax, label_setter)(label)
update_fn = update_fn or (
lambda _, artists: artists
) # if update_fn is not supplied: just return artists
def _init():
for lim_setter, idx in zip(('set_xlim', 'set_ylim', 'set_zlim'),
range(n)):
getattr(ax, lim_setter)(min(np.min(pos[idx]) for pos in positions),
max(np.max(pos[idx]) for pos in positions))
return lines
def _update(i):
for k, ln in enumerate(lines):
lower = max(0, i - n_previous - 1)
ln.set_data(positions[k][0][lower:i], positions[k][1][lower:i])
if n == 3:
ln.set_3d_properties(positions[k][2][lower:i])
return update_fn(i, lines)
animation_kwargs = animation_kwargs or {}
animation_kwargs = {
'frames': maxlen,
'interval': delay,
'repeat': True,
**animation_kwargs
}
ani = FuncAnimation(fig, _update, init_func=_init, **animation_kwargs)
return ani, fig, maxlen
def cases_and_deaths(
data: pd.DataFrame,
dates: bool = False,
ax: plt.Axes = None,
smooth: bool = True,
cases: str = "cases",
deaths: str = "deaths",
tight_layout=False,
**kwargs,
) -> plt.Axes:
"""
A simple chart showing observed new cases cases as vertical bars and
a smoothed out prediction of this curve.
Args:
data:
A dataframe with ["cases", "deaths"] columns.
dates:
If True, show dates instead of days in the x-axis.
ax:
An explicit matplotlib axes.
smooth:
If True, superimpose a plot of a smoothed-out version of the cases
curve.
cases:
deaths:
Name of the cases/deaths columns in the dataframe.
"""
if not dates:
data = data.reset_index(drop=True)
# Smoothed data
col_names = {cases: _("Cases"), deaths: _("Deaths")}
if smooth:
from pydemic import fitting as fit
smooth = pd.DataFrame(
{
_("{} (smooth)").format(col_names[cases]):
fit.smoothed_diff(data[cases]),
_("{} (smooth)").format(col_names[deaths]):
fit.smoothed_diff(data[deaths]),
},
index=data.index,
)
ax = smooth.plot(legend=False, lw=2, ax=ax)
# Prepare cases dataframe and plot it
kwargs.setdefault("alpha", 0.5)
new_cases = data.diff().fillna(0)
new_cases = new_cases.rename(col_names, axis=1)
if "ylim" not in kwargs:
deaths = new_cases.iloc[:, 1]
exp = np.log10(deaths[deaths > 0]).mean()
exp = min(10, int(exp / 2))
kwargs["ylim"] = (10**exp, None)
ax: plt.Axes = new_cases.plot.bar(width=1.0, ax=ax, **kwargs)
# Fix xticks
periods = 7 if dates else 10
xticks = ax.get_xticks()
labels = ax.get_xticklabels()
ax.set_xticks(xticks[::periods])
ax.set_xticklabels(labels[::periods])
ax.tick_params("x", rotation=0)
ax.set_ylim(1, None)
if tight_layout:
fig = ax.get_figure()
fig.tight_layout()
return ax
def animate_dataframes(frames: Iterable[pd.DataFrame], ax: plt.Axes,
lseries: Iterable[str], rseries: Iterable[str]=(),
xlim: Tuple[float]=None, ylim: Tuple[float]=None,
xlabel: str='', ylabel: str='', labels: Iterable[str]=None,
xscale: str='linear', yscale: str='linear', legend: bool=True,
fmt: Union[Iterable[Iterable[str]], Iterable, str]=(('-',), ('--',)),
anim_args: Dict[str, Any]={}) \
-> FuncAnimation:
"""
Make an animated line plot from a list of DataFrames.
Parameters
----------
`frames`:
A list of DataFrames to plot,
`ax`:
The axis on which to animate,
`lseries`:
List of column names to plot on the left,
`rseries`:
List of column names to plot on the right,
`xlim, ylim`:
Tuples of min/max axis values,
`xlabel, ylabel`:
String names of each axis. `xlabel` is a single string.
If `rseries` provided, `ylabel` must be a tuple of two labels,
`labels`:
Iterable of string labels for each frame in `frames`,
`xscale, yscale`:
Axis scales ('log', 'linear' etc.). If `rseries` provided, yscale
must be a tuple of two strings,
`legend`:
Whether to show legend on plot.
`fmt`:
A format string for lines. Either a single string for all lines, or
an iterable of strings applying to lseries and rseries in order, or an iterable
of iterables of strings - one iterable for each lseries and rseries.
`anim_args`:
A dictionary of arguments to pass to `matplotlib.animation.FuncAnimation`
class. See https://matplotlib.org/3.1.0/api/_as_gen/matplotlib.animation.FuncAnimation.html
for more details.
"""
# generate defaults
if ylim is None:
min1 = min(f[lseries].values.min() for f in frames)
max1 = max(f[lseries].values.max() for f in frames)
if rseries:
min2 = min(f[rseries].values.min() for f in frames)
max2 = max(f[rseries].values.max() for f in frames)
ylim = ((min1, max1), (min2, max2))
else:
ylim = ((min1, max1), )
elif isinstance(ylim, Iterable):
if len(ylim) == 2 and not any(
map(lambda x: isinstance(x, Iterable), ylim)):
ylim = (ylim, ylim)
if isinstance(yscale, str): yscale = (yscale, yscale)
if isinstance(ylabel, str): ylabel = (ylabel, ylabel)
if xlim is None:
xlim = (min(f.index.values.min() for f in frames),
max(f.index.values.max() for f in frames))
if rseries:
if isinstance(fmt, str):
raise TypeError('Provide format string for rseries.')
elif isinstance(fmt, Iterable) and all(
(map(lambda x: isinstance(x, str), fmt))):
fmt = ((*fmt[:len(lseries)]), (*fmt[len(lseries):]))
else:
if isinstance(fmt, str):
fmt = [fmt] * len(lseries)
elif isinstance(fmt, Iterable) and all(
map(lambda x: isinstance(x, Iterable), fmt)):
fmt = fmt[0] * len(lseries)
# set up figure and axes
fig = ax.get_figure()
ax1 = ax # left axis
ax1.set(ylabel=ylabel[0],
yscale=yscale[0],
ylim=ylim[0],
xlabel=xlabel,
xscale=xscale,
xlim=xlim)
fig.autofmt_xdate()
if rseries: # right axis
ax2 = ax1.twinx()
ax2.set(ylabel=ylabel[1], yscale=yscale[1], ylim=ylim[1])
lines1 = [ax1.plot([], [], f, label=s)[0] for s, f in \
zip_longest(lseries, fmt[0], fillvalue=fmt[0][-1])]
if rseries:
lines2 = [ax2.plot([], [], f, label=s)[0] for s, f in \
zip_longest(rseries, fmt[1], fillvalue=fmt[1][-1])]
else:
lines2 = []
if legend:
ax1.legend(lines1 + lines2, (*lseries, *rseries), loc='upper right')
# define animation start and frames
def init_func():
return (*lines1, *lines2)
def plot_func(i):
frame = frames[i]
if labels is not None:
ax1.set_title(labels[i])
for lines, series in zip((lines1, lines2), (lseries, rseries)):
for line, s in zip(lines, series):
line.set_data(frame.index, frame[s])
return (*lines1, *lines2)
anim = FuncAnimation(fig,
func=plot_func,
frames=len(frames),
init_func=init_func,
**anim_args)
return anim
