diff --git a/statista/time_series.py b/statista/time_series.py index 9f32c34..c6594e7 100644 --- a/statista/time_series.py +++ b/statista/time_series.py @@ -8,7 +8,7 @@ BOX_MEAN_PROP = dict(marker="x", markeredgecolor="w", markerfacecolor="firebrick") -VIOLIN_PROP = dict(face="#27408B", edge="#DC143C", alpha=0.7) +VIOLIN_PROP = dict(face="#27408B", edge="black", alpha=0.7) class TimeSeries(DataFrame): @@ -588,3 +588,95 @@ def raincloud( # Display the plot plt.show() return fig, ax + + def histogram(self, bins=10, **kwargs) -> Tuple[Figure, Axes]: + """ + Plots a histogram of the time series data. + + Parameters + ---------- + bins : int, optional, default is 10. + Number of histogram bins. + **kwargs: dict, optional + fig: matplotlib.figure.Figure, optional + Existing figure to plot on. If None, a new figure is created. + ax: matplotlib.axes.Axes, optional + Existing axes to plot on. If None, a new axes is created. + grid: bool, optional + Whether to show grid lines. Default is True. + color: str, optional, default is None. + Colors to use for the plot elements. + title: str, optional + Title of the plot. Default is 'Box Plot'. + xlabel: str, optional + Label for the x-axis. Default is 'Index'. + ylabel: str, optional + Label for the y-axis. Default is 'Value'. + title_fontsize: int, optional + Font size of the title. + label_fontsize: int, optional + Font size of the title and labels. + tick_fontsize: int, optional + Font size of the tick labels. + legend: str, optional + Legend to display in the plot. + legend_fontsize: int, optional + Font size of the legend. + + Returns + ------- + fig : matplotlib.figure.Figure + The figure object containing the plot. + ax : matplotlib.axes.Axes + The axes object containing the plot. + + Examples + -------- + >>> ts = TimeSeries(np.random.randn(100)) + >>> fig, ax = ts.histogram() + """ + # plt.style.use('ggplot') + + fig, ax = self._get_ax_fig(fig=kwargs.get("fig"), ax=kwargs.get("ax")) + color = kwargs.get("color") if "color" in kwargs else VIOLIN_PROP + ax.hist( + self.values, + bins=bins, + color=color.get("face"), + edgecolor=color.get("edge"), + alpha=color.get("alpha"), + ) + + # Set title and labels with larger font sizes + ax.set_title( + kwargs.get("title", "Histogram"), + fontsize=kwargs.get("title_fontsize", 18), + fontweight="bold", + ) + ax.set_xlabel( + kwargs.get("xlabel", "X-axis Label"), + fontsize=kwargs.get("label_fontsize", 14), + ) + ax.set_ylabel( + kwargs.get("ylabel", "Y-axis Label"), + fontsize=kwargs.get("label_fontsize", 14), + ) + + ax.grid(kwargs.get("grid"), axis="both", linestyle="-.", linewidth=0.3) + + # Customize ticks and their labels + ax.tick_params( + axis="both", which="major", labelsize=kwargs.get("tick_fontsize", 12) + ) + + # Add a legend if needed + if "legend" in kwargs: + ax.legend( + [kwargs.get("legend")], fontsize=kwargs.get("legend_fontsize", 12) + ) + + # Adjust layout for better spacing + plt.tight_layout() + + plt.show() + return fig, ax diff --git a/tests/test_time_series.py b/tests/test_time_series.py index 86cefc9..3056e9e 100644 --- a/tests/test_time_series.py +++ b/tests/test_time_series.py @@ -111,3 +111,70 @@ def test_raincloud(self, ts: TimeSeries, request): assert ax2 is ax, "If ax is provided, plot_box should use it." if ts.shape[1] > 1: assert len(ax.get_xticklabels()) == 3 + + +class TestHistogram: + + def test_default(self): + # Test with default parameters + ts = TimeSeries(np.array([1, 2, 2, 3, 3, 3, 4, 4, 4, 4])) + fig, ax = ts.histogram() + assert ax.get_title() == "Histogram" + assert ax.get_xlabel() == "X-axis Label" + assert ax.get_ylabel() == "Y-axis Label" + plt.close() + + def test_custom_labels(self): + # Test with custom title and labels + ts = TimeSeries(np.array([1, 2, 3, 4, 5])) + fig, ax = ts.histogram( + title="Custom Title", xlabel="Custom X", ylabel="Custom Y" + ) + + assert ax.get_title() == "Custom Title" + assert ax.get_xlabel() == "Custom X" + assert ax.get_ylabel() == "Custom Y" + plt.close() + + def test_custom_colors(self): + # Test with custom colors + ts = TimeSeries(np.array([1, 2, 3, 4, 5])) + fig, ax = ts.histogram(color=dict(face="green", edge="red", alpha=0.5)) + + patches = ax.patches + assert patches[0].get_facecolor() == ( + 0.0, + 0.5019607843137255, + 0.0, + 0.5, + ) # RGBA for green with alpha 0.5 + assert patches[0].get_edgecolor() == (1.0, 0.0, 0.0, 0.5) # RGBA for red + plt.close() + + def test_legend(self): + # Test with a legend + ts = TimeSeries(np.array([1, 2, 3, 4, 5])) + fig, ax = ts.histogram(legend="Sample Legend") + + legend = ax.get_legend() + assert legend is not None + assert legend.get_texts()[0].get_text() == "Sample Legend" + plt.close() + + def test_bins(self): + # Test with different number of bins + ts = TimeSeries(np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])) + fig, ax = ts.histogram(bins=5) + + # Number of bars should match the number of bins + assert len(ax.patches) == 5 + plt.close() + + def test_grid_and_ticks(self): + # Test grid and tick customization + ts = TimeSeries(np.array([1, 2, 3, 4, 5])) + fig, ax = ts.histogram(tick_fontsize=16) + + for tick in ax.get_xticklabels() + ax.get_yticklabels(): + assert tick.get_fontsize() == 16 # Check the fontsize of the ticks + plt.close()