feat(plotting): extract plotting module with histogram and CDF

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

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()