diff --git a/src/disk_checker/plotting/plots.py b/src/disk_checker/plotting/plots.py new file mode 100644 index 0000000..270eefc --- /dev/null +++ b/src/disk_checker/plotting/plots.py @@ -0,0 +1,162 @@ +"""Plotting module for disk latency visualization.""" + +import matplotlib +matplotlib.use("Agg") # Non-interactive backend for server environments +import matplotlib.pyplot as plt +from matplotlib.ticker import FuncFormatter +import numpy as np +from pathlib import Path +from typing import List, Optional + +from disk_checker.core.converter import to_unit, UNIT_LABELS +from disk_checker.core.helpers import human_count + + +def draw_histogram( + ax, + df, + metrics: List[str], + unit: str = "us", + log_y: bool = False, + log_x: bool = False, + normalize: bool = False, +) -> None: + """ + Draw a histogram showing the distribution of latencies for each metric. + + Uses a step plot (where="post") to show the discrete nature of the data. + If normalize=True, Y-axis shows probability density instead of counts. + + Args: + ax: Matplotlib axis object + df: Long-format DataFrame from load_latency_csv + metrics: List of metric names to plot + unit: X-axis unit (ns, us, ms, s) + log_y: Use logarithmic Y scale + log_x: Use logarithmic X scale + normalize: Normalize to probability density + """ + for metric in metrics: + g = df[df["metric"] == metric].sort_values("latency_ns") + if g.empty: + continue + + x = to_unit(g["latency_ns"].values, unit) + y = g["count"].values.astype(float) + + if normalize: + s = y.sum() + if s > 0: + y = y / s + + ax.step(x, y, where="post", lw=1.4, label=metric) + ax.fill_between(x, y, step="post", alpha=0.25) + + ax.set_xlabel(f"Latency ({UNIT_LABELS[unit]})") + ax.set_ylabel("Probability" if normalize else "Number of I/Os") + ax.set_title("Completion latency distribution") + ax.grid(True, which="both", ls=":", alpha=0.5) + + if not normalize: + ax.yaxis.set_major_formatter(FuncFormatter(human_count)) + if log_y: + ax.set_yscale("log") + if log_x: + ax.set_xscale("log") + if len(metrics) > 1: + ax.legend(title="metric", fontsize=9) + + +def draw_cdf(ax, df, metrics: List[str], unit: str = "us", log_x: bool = False) -> None: + """ + Draw a Cumulative Distribution Function (CDF). + + Shows the percentage of I/Os completed within each latency threshold. + Each metric gets its own line; Y-axis ranges from 0% to 100%. + + Args: + ax: Matplotlib axis object + df: Long-format DataFrame + metrics: List of metric names + unit: X-axis unit + log_x: Use logarithmic X scale + """ + for metric in metrics: + g = df[df["metric"] == metric].sort_values("latency_ns") + if g.empty: + continue + + x = to_unit(g["latency_ns"].values, unit) + counts = g["count"].values.astype(float) + total = counts.sum() + if total <= 0: + continue + + cdf = np.cumsum(counts) / total * 100.0 + ax.plot(x, cdf, lw=1.6, label=f"{metric} (n={int(total):,})") + + ax.set_xlabel(f"Latency ({UNIT_LABELS[unit]})") + ax.set_ylabel("Cumulative % of I/Os") + ax.set_title("Cumulative latency distribution") + ax.set_ylim(0, 100) + ax.grid(True, which="both", ls=":", alpha=0.5) + ax.yaxis.set_major_formatter(FuncFormatter(lambda v, _: f"{v:.0f}%")) + + if log_x: + ax.set_xscale("log") + ax.legend(fontsize=9) + + +def generate_plot( + df, + metrics: List[str], + output_path: Path, + unit: str = "us", + log_y: bool = False, + log_x: bool = False, + normalize: bool = False, + show_cdf: bool = False, + cdf_only: bool = False, + title: str = "", + dpi: int = 150, +) -> Path: + """ + Generate a complete plot with histogram and/or CDF. + + Creates the appropriate subplot layout based on the options provided. + + Args: + df: Long-format DataFrame + metrics: Metrics to plot + output_path: Where to save the plot + unit: X-axis unit + log_y: Logarithmic Y scale + log_x: Logarithmic X scale + normalize: Normalize histogram + show_cdf: Show CDF below histogram + cdf_only: Show only CDF + title: Plot title + dpi: Output resolution + + Returns: + Path to saved plot + """ + if cdf_only: + fig, ax = plt.subplots(figsize=(11, 6)) + draw_cdf(ax, df, metrics, unit, log_x) + elif show_cdf: + fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(11, 10)) + draw_histogram( + ax1, df, metrics, unit, log_y, log_x, normalize + ) + draw_cdf(ax2, df, metrics, unit, log_x) + else: + fig, ax = plt.subplots(figsize=(11, 6)) + draw_histogram(ax, df, metrics, unit, log_y, log_x, normalize) + + if title: + fig.suptitle(title, fontsize=13, y=0.995) + fig.tight_layout() + fig.savefig(output_path, dpi=dpi, bbox_inches="tight") + plt.close(fig) + return output_path \ No newline at end of file diff --git a/tests/test_plots.py b/tests/test_plots.py new file mode 100644 index 0000000..3aaa6d7 --- /dev/null +++ b/tests/test_plots.py @@ -0,0 +1,134 @@ +"""Tests for the plotting module.""" + +import pytest +from pathlib import Path +from unittest.mock import MagicMock, patch +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +from disk_checker.plotting.plots import draw_histogram, draw_cdf, generate_plot +import pandas as pd + + +class TestDrawHistogram: + """Tests for draw_histogram function.""" + + @pytest.fixture + def sample_df(self): + """Create a sample DataFrame for testing.""" + return pd.DataFrame({ + "metric": ["read_clat", "read_clat", "write_clat"], + "latency_ns": [1000, 2000, 1500], + "count": [1, 2, 1], + }) + + def test_draws_histogram(self, sample_df, tmp_path): + """Test that histogram is drawn without errors.""" + fig, ax = plt.subplots() + draw_histogram(ax, sample_df, ["read_clat"]) + assert ax.get_title() == "Completion latency distribution" + plt.close(fig) + + def test_multiple_metrics(self, sample_df, tmp_path): + """Test drawing multiple metrics.""" + fig, ax = plt.subplots() + draw_histogram(ax, sample_df, ["read_clat", "write_clat"]) + assert len(ax.collections) > 0 + plt.close(fig) + + def test_log_scale(self, sample_df, tmp_path): + """Test logarithmic scale.""" + fig, ax = plt.subplots() + draw_histogram(ax, sample_df, ["read_clat"], log_x=True, log_y=True) + assert ax.get_xscale() == "log" + assert ax.get_yscale() == "log" + plt.close(fig) + + +class TestDrawCdf: + """Tests for draw_cdf function.""" + + @pytest.fixture + def sample_df(self): + """Create a sample DataFrame for testing.""" + return pd.DataFrame({ + "metric": ["read_clat", "read_clat"], + "latency_ns": [1000, 2000], + "count": [1, 3], + }) + + def test_draws_cdf(self, sample_df, tmp_path): + """Test that CDF is drawn without errors.""" + fig, ax = plt.subplots() + draw_cdf(ax, sample_df, ["read_clat"]) + assert ax.get_title() == "Cumulative latency distribution" + assert ax.get_ylim() == (0, 100) + plt.close(fig) + + def test_log_scale(self, sample_df, tmp_path): + """Test logarithmic X scale.""" + fig, ax = plt.subplots() + draw_cdf(ax, sample_df, ["read_clat"], log_x=True) + assert ax.get_xscale() == "log" + plt.close(fig) + + +class TestGeneratePlot: + """Tests for generate_plot function.""" + + @pytest.fixture + def sample_df(self): + """Create a sample DataFrame for testing.""" + return pd.DataFrame({ + "metric": ["read_clat", "read_clat"], + "latency_ns": [1000, 2000], + "count": [1, 3], + }) + + def test_generates_histogram(self, sample_df, tmp_path): + """Test generating histogram plot.""" + output = tmp_path / "test.png" + result = generate_plot( + df=sample_df, + metrics=["read_clat"], + output_path=output, + ) + assert result == output + assert output.exists() + + def test_generates_cdf_only(self, sample_df, tmp_path): + """Test generating CDF-only plot.""" + output = tmp_path / "test_cdf.png" + result = generate_plot( + df=sample_df, + metrics=["read_clat"], + output_path=output, + cdf_only=True, + ) + assert result == output + assert output.exists() + + def test_generates_histogram_with_cdf(self, sample_df, tmp_path): + """Test generating histogram with CDF below.""" + output = tmp_path / "test_both.png" + result = generate_plot( + df=sample_df, + metrics=["read_clat"], + output_path=output, + show_cdf=True, + ) + assert result == output + assert output.exists() + + def test_custom_dpi(self, sample_df, tmp_path): + """Test custom DPI setting.""" + output = tmp_path / "test_dpi.png" + result = generate_plot( + df=sample_df, + metrics=["read_clat"], + output_path=output, + dpi=300, + ) + assert result == output + assert output.exists() \ No newline at end of file