diff --git a/pipit/trace.py b/pipit/trace.py index cc867a40..54460c1b 100644 --- a/pipit/trace.py +++ b/pipit/trace.py @@ -952,6 +952,15 @@ def plot_message_histogram(self, bins=20, *args, **kwargs): # Return the Bokeh plot return plot_message_histogram(data, *args, **kwargs) + def plot_comm_over_time(self, output="size", message_type="send", *args, **kwargs): + from .vis import plot_comm_over_time + + # Generate the data + data = self.comm_over_time(output=output, message_type=message_type, *args, **kwargs) + + # Return the Bokeh plot + return plot_comm_over_time(data, message_type=message_type, output=output) + def plot_timeline(self, *args, **kwargs): from .vis import plot_timeline diff --git a/pipit/vis/__init__.py b/pipit/vis/__init__.py index 5119ca11..733cc0f0 100644 --- a/pipit/vis/__init__.py +++ b/pipit/vis/__init__.py @@ -1,2 +1,6 @@ -from .core import plot_comm_matrix, plot_message_histogram # noqa: F401 +from .core import ( + plot_comm_matrix, + plot_message_histogram, + plot_comm_over_time, +) # noqa: F401 from .timeline import plot_timeline # noqa: F401 diff --git a/pipit/vis/core.py b/pipit/vis/core.py index 7b413a12..e1bd511e 100644 --- a/pipit/vis/core.py +++ b/pipit/vis/core.py @@ -14,6 +14,8 @@ get_size_hover_formatter, get_size_tick_formatter, show, + get_time_tick_formatter, + get_time_hover_formatter, ) @@ -146,3 +148,52 @@ def plot_message_histogram( # Return plot return show(p, return_fig=return_fig) + + +def plot_comm_over_time(data, output, message_type, return_fig=False): + """Plots the trace's communication over time. + + Args: + data (hist, edges): Histogram and edges + output (str): Specifies whether the matrix contains "size" or "count" values. + message_type (str): Specifies whether the message is "send" or "receive". + return_fig (bool, optional): Specifies whether to return the Bokeh figure + object. Defaults to False, which displays the result and returns nothing. + + Returns: + Bokeh figure object if return_fig, None otherwise + """ + + hist, edges = data + is_size = output == "size" + + p = figure( + x_axis_label="Time", + y_axis_label="Total volume sent" if is_size else "Number of messages", + tools="hover,save", + ) + p.y_range.start = 0 + p.xaxis.formatter = get_time_tick_formatter() + p.yaxis.formatter = get_size_tick_formatter() + + p.quad(top=hist, bottom=0, left=edges[:-1], right=edges[1:], line_color="white") + + hover = p.select(HoverTool) + hover.tooltips = ( + [ + ("Bin", "@left{custom} - @right{custom}"), + ("Total volume sent:", "@top{custom}"), + ] + if is_size + else [ + ("Bin", "@left{custom} - @right{custom}"), + ("number of messages:", "@top"), + ] + ) + hover.formatters = { + "@left": get_time_hover_formatter(), + "@right": get_time_hover_formatter(), + "@top": get_size_hover_formatter(), + } + + return show(p, return_fig=return_fig)