Source code for nannyml.plots.components.figure

#  Author:   Niels Nuyttens  <niels@nannyml.com>
#
#  License: Apache Software License 2.0
from typing import Any, Dict, List, Optional, Union

import matplotlib.colors
import numpy as np
import pandas as pd
import plotly.graph_objs as go

from nannyml.exceptions import InvalidArgumentsException
from nannyml.plots.colors import Colors
from nannyml.plots.components.hover import Hover
from nannyml.plots.components.step_plot import alert as step_plot_alert
from nannyml.plots.components.step_plot import metric as step_plot_metric
from nannyml.plots.util import add_artificial_endpoint, check_and_convert, is_time_based_x_axis


[docs]class Figure(go.Figure): """Extending the Plotly Figure class functionality.""" SUPPORTED_METRIC_STYLES = ['step'] def __init__( self, title: str, x_axis_title: str, y_axis_title: str, y_axis_limit: Optional[List] = None, metric_style: str = 'step', subplot_args: Optional[Dict[str, Any]] = None, **kwargs, ): """Creates a new Figure.""" layout = go.Layout( title=title, xaxis=dict( title=x_axis_title, linecolor=Colors.INDIGO_PERSIAN, showgrid=False, mirror=True, zeroline=False ), yaxis=dict( title=y_axis_title, linecolor=Colors.INDIGO_PERSIAN, showgrid=False, range=y_axis_limit, mirror=True, zeroline=False, ), paper_bgcolor='rgba(255,255,255,1)', plot_bgcolor='rgba(255,255,255,1)', hoverlabel=dict(bgcolor="white", font_size=14), **kwargs, ) super().__init__(layout=layout) if subplot_args is not None: self.set_subplots(**subplot_args) if metric_style not in self.SUPPORTED_METRIC_STYLES: raise InvalidArgumentsException( f"metric style '{metric_style}' is not supported. " "Please provide one of the following values: " f"{str.join(',', self.SUPPORTED_METRIC_STYLES)}" ) self._metric_style: str = metric_style
[docs] def add_metric( self, data: Union[np.ndarray, pd.Series], name: str, color: str, indices: Optional[Union[np.ndarray, pd.Series]] = None, start_dates: Optional[Union[np.ndarray, pd.Series]] = None, end_dates: Optional[Union[np.ndarray, pd.Series]] = None, hover: Optional[Hover] = None, subplot_args: Optional[Dict[str, Any]] = None, **kwargs, ): if self._metric_style not in self.SUPPORTED_METRIC_STYLES: raise InvalidArgumentsException( f"metric style '{self._metric_style}' is not supported. " "Please provide one of the following values: " f"{str.join(',', self.SUPPORTED_METRIC_STYLES)}" ) if self._metric_style == 'step': step_plot_metric( figure=self, data=data, chunk_indices=indices, chunk_start_dates=start_dates, chunk_end_dates=end_dates, name=name, color=color, hover=hover, subplot_args=subplot_args, **kwargs, )
[docs] def add_period_separator(self, x, color=Colors.GRAY_DARK, subplot_args: Optional[Dict[str, Any]] = None, **kwargs): if subplot_args is not None: kwargs.update(dict(col=subplot_args.get('col'), row=subplot_args.get('row'))) self.add_vline(x=x, line=dict(color=color, width=1), layer='below', **kwargs)
[docs] def add_threshold( self, data, name, indices=None, start_dates=None, end_dates=None, color=Colors.RED_IMPERIAL, with_additional_endpoint: bool = False, subplot_args: Optional[Dict[str, Any]] = None, **kwargs, ): data, start_dates, end_dates, indices = check_and_convert(data, start_dates, end_dates, indices) x = start_dates if is_time_based_x_axis(start_dates, end_dates) else indices if with_additional_endpoint: x, data = add_artificial_endpoint(indices, start_dates, end_dates, data) show_in_legend = True if subplot_args is None: subplot_args = {} else: if 'showlegend' in kwargs: is_first_subplot = subplot_args.get('row', 0) == 1 and subplot_args.get('col', 0) == 1 show_in_legend = kwargs['showlegend'] and is_first_subplot kwargs['showlegend'] = show_in_legend self.add_trace( go.Scatter( name=name, mode='lines', x=x, y=data, line=dict(color=color, width=2, dash='dash'), hoverinfo='skip', **kwargs, ), **subplot_args, )
[docs] def add_confidence_band( self, upper_confidence_boundaries, lower_confidence_boundaries, name, indices=None, start_dates=None, end_dates=None, color=Colors.RED_IMPERIAL, with_additional_endpoint: bool = False, subplot_args: Optional[Dict[str, Any]] = None, **kwargs, ): data, start_dates, end_dates, indices = check_and_convert( [upper_confidence_boundaries, lower_confidence_boundaries], start_dates, end_dates, indices ) x = start_dates if is_time_based_x_axis(start_dates, end_dates) else indices if with_additional_endpoint: x, data = add_artificial_endpoint(indices, start_dates, end_dates, data) show_in_legend = kwargs.get('showlegend', True) is_single_plot = subplot_args and subplot_args.get('row') is None and subplot_args.get('col') is None if subplot_args is None or is_single_plot: subplot_args = {} del kwargs['showlegend'] self.add_trace( go.Scatter( name=name, mode='lines', x=x, y=data[0], line=dict(shape='hv', color='rgba(0,0,0,0)'), hoverinfo='skip', showlegend=False, connectgaps=True, **kwargs, ), **subplot_args, ) self.add_trace( go.Scatter( name=name, mode='lines', x=x, y=data[1], line=dict(shape='hv', color='rgba(0,0,0,0)'), fill='tonexty', fillcolor='rgba{}'.format(matplotlib.colors.to_rgba(matplotlib.colors.to_rgb(color), alpha=0.2)), hoverinfo='skip', showlegend=show_in_legend, connectgaps=True, **kwargs, ), **subplot_args, )
[docs] def add_alert( self, data: Union[np.ndarray, pd.Series], name: str, color: str = Colors.RED_IMPERIAL, indices: Optional[Union[np.ndarray, pd.Series]] = None, start_dates: Optional[Union[np.ndarray, pd.Series]] = None, end_dates: Optional[Union[np.ndarray, pd.Series]] = None, subplot_args: Optional[Dict[str, Any]] = None, **kwargs, ): if self._metric_style not in self.SUPPORTED_METRIC_STYLES: raise InvalidArgumentsException( f"metric style '{self._metric_style}' is not supported. " "Please provide one of the following values: " f"{str.join(',', self.SUPPORTED_METRIC_STYLES)}" ) if self._metric_style == 'step': step_plot_alert( figure=self, data=data, chunk_indices=indices, chunk_start_dates=start_dates, chunk_end_dates=end_dates, name=name, color=color, subplot_args=subplot_args, **kwargs, )