from typing import Any, Callable, Dict, List, Optional, Tuple, Union
import matplotlib as mpl
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
import iccas as ic
from iccas.charts.common import DEFAULT_CMAP, legend
from iccas.i18n import translated
from iccas.i18n.lib import NullTranslation, Translation
[docs]def add_labels_to_bars(
rects: List[plt.Rectangle],
fmt: str = '{:.1f}',
ax: Optional[plt.Axes] = None,
y_offset: int = 7, # in points
**kwargs,
) -> Tuple[List[plt.Annotation], Callable[[], None]]:
"""
Attach a text label above each bar in ``rects``, displaying its height.
Returns a tuple with the list of labels and a function to update the labels
(useful for animations)
"""
ax = ax or plt.gca()
labels = []
for rect in rects:
height = rect.get_height()
label = ax.annotate(
fmt.format(height),
xy=(rect.get_x() + rect.get_width() / 2, height),
ha='center', va='bottom',
xytext=(0, y_offset), textcoords="offset points",
**kwargs,
)
labels.append(label)
def update_labels():
for label, rect, in zip(labels, rects):
height = rect.get_height()
label.set_text(fmt.format(height))
label.xy = label.xy[0], height
label.xyann = (0, y_offset)
return labels, update_labels
@translated()
class AgeDistributionBarChart:
def __init__(
self,
counts: pd.DataFrame,
variable: str = "cases",
*,
strings: Translation = NullTranslation(),
normalize: bool = True,
age_group_size: int = 10,
window: int = 14,
population_distribution: pd.Series = None,
ax: plt.Axes = None,
resample_kwargs: Dict[str, Any] = {},
):
"""
Shows the age distribution of cases/deaths at a given date.
Usable either to draw a static chart or to generate an animation.
Args:
counts:
variable:
normalize:
age_group_size:
window:
population_distribution:
ax:
lang:
"""
self.ax = ax = ax or plt.gca()
s = strings
data = counts[variable].drop(columns="unknown")
data = ic.running_average(data, window=window, **resample_kwargs)
data = ic.aggregate_age_groups(data, cuts=age_group_size)
if normalize:
data = data.divide(data.sum(axis=1), axis=0)
ax.yaxis.set_major_formatter(
mpl.ticker.PercentFormatter(xmax=1.0, decimals=0)
)
self.data = data
if normalize and population_distribution is not None:
population = ic.aggregate_age_groups(
population_distribution, cuts=age_group_size
)
sns.barplot(
ax=ax,
label=s["istat_population_data_label"],
x=population.index,
y=population,
facecolor='#a1c9f4',
hatch="/",
edgecolor="white",
)
age_groups = data.columns
label = s[f"{variable}_label"]
self.bars = ax.bar(
x=age_groups,
height=[0] * len(age_groups),
label=label,
facecolor='#4878d0',
alpha=0.7,
)
self.labels, self.update_labels = add_labels_to_bars(
self.bars,
fmt='{:.0%}' if normalize else '{:n}',
fontsize=10,
ax=ax,
color='white',
bbox=dict(
boxstyle=mpl.patches.BoxStyle("Round", pad=0.3),
color='black', alpha=.35,
),
)
ymax = 1.05 * data.max().max()
ax.set_ylim(0, ymax)
ax.set_xticklabels(age_groups)
ax.set_xlabel(s["age"])
ax.set_ylabel("")
ax.grid(False, which='both', axis='x')
if normalize:
title = s.get(f"running_{variable}_age_distribution", count=window)
else:
title = s.get(f"running_average_{variable}_title", count=window)
ax.set_title(title, fontsize=14)
if normalize and population_distribution is not None:
ax.yaxis.set_major_formatter(
mpl.ticker.PercentFormatter(xmax=1.0, decimals=0)
)
ax.legend()
self.date_text = ax.text(
0.5,
0.93,
"",
ha="center",
va="center",
transform=ax.transAxes,
color=(1, 1, 1, 0.9),
bbox=dict(
boxstyle=mpl.patches.BoxStyle("Round", pad=0.4),
color=(0, 0, 0, 0.30),
),
fontsize=14,
fontweight="semibold",
)
self.artists = [
self.date_text,
*self.bars,
*self.labels
]
def show(self, date: pd.Timestamp):
d = self.data.loc[date]
if isinstance(d, pd.DataFrame):
if len(d) > 1:
raise ValueError(f'"{date}" matches more than one date')
d = d.stack("age_group")
for bar, h in zip(self.bars, d):
bar.set_height(h)
formatted_date = date.strftime("%d %B %Y")
self.date_text.set_text(formatted_date)
self.update_labels()
return self.artists
def animation(self, **kwargs) -> mpl.animation.FuncAnimation:
kwargs = {"blit": True, "repeat": False, "interval": 80, **kwargs}
return mpl.animation.FuncAnimation(
self.ax.figure, self.show, frames=self.data.index, **kwargs
)
[docs]def age_dist_bar_chart(
counts: pd.DataFrame,
date: pd.Timestamp,
variable: str = "cases",
*,
normalize: bool = True,
age_group_size: int = 10,
window: int = 14,
population_distribution: pd.Series = None,
ax: plt.Axes = None,
lang: Optional[str] = None,
):
if ax is None:
fig, ax = plt.subplots(figsize=(10, 7))
AgeDistributionBarChart(
ax=ax,
counts=counts,
variable=variable,
normalize=normalize,
age_group_size=age_group_size,
window=window,
population_distribution=population_distribution,
lang=lang, # type: ignore
).show(date)
[docs]@translated()
def average_by_period_bar_chart(
counts: pd.DataFrame,
variable: str,
*,
strings: Translation,
freq: Union[str, int] = 7,
normalize: bool = False,
age_group_size: int = 20,
stacked: bool = True,
ylim: float = None,
ax: Optional[plt.Axes] = None,
figsize: Tuple[float, float] = (12, 7),
) -> plt.Axes:
if ax is None:
_, ax = plt.subplots(figsize=figsize)
# Data preparation
d = counts[variable].drop(columns='unknown')
d = ic.aggregate_age_groups(d, age_group_size)
d = ic.count_by_period(d, freq=freq)
if normalize:
d = d.div(d.sum(axis=1), axis=0)
# Plot
ax = d.plot.bar(stacked=stacked, cmap=DEFAULT_CMAP, ax=ax, width=.85)
legend(ax=ax, title=strings['age'])
# Axes setup
if freq == 'M':
xlabel = strings['month']
date_fmt = '%B'
xtick_rotation = 0
else:
xlabel = strings['period_end_date']
date_fmt = '%d %b'
xtick_rotation = 45
ax.set_xticklabels(d.index.strftime(date_fmt), rotation=xtick_rotation)
ax.set_xlabel(xlabel)
if normalize:
if stacked:
ax.set_ylim([0, 1.0])
ax.yaxis.set_major_formatter(mpl.ticker.PercentFormatter(xmax=1.0, decimals=0))
else:
ax.yaxis.set_major_formatter(mpl.ticker.StrMethodFormatter('{x:n}'))
if ylim:
ax.set_ylim(0, ylim * ax.get_ylim()[1])
# Title
base_title = strings[
f'{variable}_age_distribution' if normalize else f'average_daily_{variable}'
]
if freq == 'M':
period_label = strings['month_by_month']
elif freq in {'W', 'W-SUN'}:
period_label = strings['week_by_week']
elif isinstance(freq, int):
period_label = strings.get('by_periods_of_n_days', count=freq)
title = f"{base_title} {period_label}"
ax.set_title(title)
return ax
if __name__ == '__main__':
plt.style.use('seaborn')
ic.set_locale('it')
df = ic.fix_monotonicity(ic.get())
anim = AgeDistributionBarChart(
df, window=1,
population_distribution=ic.get_population_by_age_group().percentage
).animation()
plt.show()