timeplot.py

"""時系列マルチプロットモジュール。

pandas.DataFrame を入力として、matplotlib(pyplot) または plotly で
複数系列の重ね書き・複数段(サブプロット)の時系列グラフを作成する。

README.md で示された想定シグネチャ:

    def timeplot(
        df: pd.Dataflame,
        xplot,     # dfのx軸を示すカラム名
        yplot={...},
        puttyppe,  # pyplotかplotlyを選ぶ
    ):

これに沿って、yplot は「1段にまとめて重ね書きしたい列名のリストのリスト」
(フラットなリストなら単段) として受け取る。
"""

from __future__ import annotations

from typing import List, Literal, Optional, Sequence, Tuple, Union

import pandas as pd

PlotType = Literal["pyplot", "plotly"]
YPlotSpec = Union[Sequence[str], Sequence[Sequence[str]]]


def _normalize_yplot(yplot: YPlotSpec) -> List[List[str]]:
    """yplot 指定を「段(row)ごとの列名リスト」の形に正規化する。

    - ["a", "b"]            -> [["a", "b"]]   (1段に a, b を重ね書き)
    - [["a", "b"], ["c"]]   -> そのまま        (2段: 1段目 a,b 重ね書き / 2段目 c)
    """
    yplot = list(yplot)
    if len(yplot) == 0:
        raise ValueError("yplot must not be empty")
    if isinstance(yplot[0], str):
        return [list(yplot)]
    return [list(group) for group in yplot]


def timeplot(
    df: pd.DataFrame,
    xplot: str,
    yplot: YPlotSpec,
    plot_type: PlotType = "pyplot",
    title: Optional[str] = None,
    ylabel: Union[str, Sequence[Optional[str]], None] = None,
    figsize: Tuple[float, float] = (10.0, 3.0),
):
    """複数段の時系列マルチプロットを作成する。

    Parameters
    ----------
    df:
        プロット対象のデータフレーム。
    xplot:
        x軸に使うカラム名 (時刻カラムなど)。
    yplot:
        y軸カラムの指定。
        - フラットなリスト ``["col_a", "col_b"]`` を渡すと1段に重ね書き。
        - 入れ子リスト ``[["col_a", "col_b"], ["col_c"]]`` を渡すと
          段ごとに分けてプロットする (各段は x 軸を共有)。
    plot_type:
        ``"pyplot"`` (matplotlib) か ``"plotly"`` を選択。
    title:
        グラフ全体のタイトル。
    ylabel:
        各段の y 軸ラベル。文字列1つなら全段共通、リストなら段ごとに指定。
    figsize:
        matplotlib 使用時の1段あたりの (幅, 高さ)。

    Returns
    -------
    matplotlib.figure.Figure | plotly.graph_objects.Figure
        ``plot_type`` に応じた Figure オブジェクト。呼び出し側で
        ``fig.savefig(...)`` / ``fig.show()`` / ``fig.write_html(...)`` などを行う。
    """
    groups = _normalize_yplot(yplot)

    missing = [c for group in groups for c in group if c not in df.columns]
    if xplot not in df.columns:
        missing.append(xplot)
    if missing:
        raise KeyError(f"df に存在しないカラムが指定されました: {missing}")

    if ylabel is None or isinstance(ylabel, str):
        ylabels: List[Optional[str]] = [ylabel] * len(groups)
    else:
        ylabels = list(ylabel)
        if len(ylabels) != len(groups):
            raise ValueError("ylabel のリスト長は yplot の段数と一致させてください")

    if plot_type == "pyplot":
        return _timeplot_pyplot(df, xplot, groups, title, ylabels, figsize)
    if plot_type == "plotly":
        return _timeplot_plotly(df, xplot, groups, title, ylabels)
    raise ValueError(f"未対応の plot_type です: {plot_type!r} ('pyplot' か 'plotly' を指定)")


_JP_FONT_CANDIDATES = (
    "Yu Gothic",
    "Meiryo",
    "MS Gothic",
    "Noto Sans JP",
    "Noto Sans CJK JP",
    "Hiragino Sans",
)


def _set_jp_font_if_available() -> None:
    """日本語ラベルの文字化け(tofu)を避けるため、利用可能な日本語フォントを設定する。

    見つからない環境では何もしない (デフォルトフォントのまま)。
    """
    import matplotlib.font_manager as fm
    import matplotlib.pyplot as plt

    available = {f.name for f in fm.fontManager.ttflist}
    for name in _JP_FONT_CANDIDATES:
        if name in available:
            plt.rcParams["font.family"] = name
            plt.rcParams["axes.unicode_minus"] = False
            return


def _timeplot_pyplot(df, xplot, groups, title, ylabels, figsize):
    import matplotlib.pyplot as plt

    _set_jp_font_if_available()

    n = len(groups)
    fig, axes = plt.subplots(
        n, 1, sharex=True, figsize=(figsize[0], figsize[1] * n), squeeze=False
    )
    axes = axes[:, 0]
    for ax, group, ylabel in zip(axes, groups, ylabels):
        for col in group:
            ax.plot(df[xplot], df[col], label=col)
        ax.legend(loc="upper right")
        if ylabel:
            ax.set_ylabel(ylabel)
        ax.grid(True, alpha=0.3)
    axes[-1].set_xlabel(xplot)
    if title:
        fig.suptitle(title)
    fig.tight_layout()
    return fig


def _timeplot_plotly(df, xplot, groups, title, ylabels):
    import plotly.graph_objects as go
    from plotly.subplots import make_subplots

    n = len(groups)
    fig = make_subplots(rows=n, cols=1, shared_xaxes=True, vertical_spacing=0.08)
    for row, (group, ylabel) in enumerate(zip(groups, ylabels), start=1):
        for col in group:
            fig.add_trace(
                go.Scatter(x=df[xplot], y=df[col], mode="lines", name=col),
                row=row,
                col=1,
            )
        if ylabel:
            fig.update_yaxes(title_text=ylabel, row=row, col=1)
    fig.update_xaxes(title_text=xplot, row=n, col=1)
    fig.update_layout(title=title, height=300 * n)
    return fig


def main() -> None:
    """モジュール単体実行時のサンプル (簡易テストを兼ねる)。"""
    import numpy as np

    n = 200
    t = pd.date_range("2026-01-01", periods=n, freq="1min")
    phase = np.linspace(0, 4 * np.pi, n)
    df = pd.DataFrame(
        {
            "time": t,
            "angle_deg": 180 * np.sin(phase),
            "angular_velocity_dps": np.gradient(180 * np.sin(phase)),
            "temperature": 25 + np.random.default_rng(0).normal(0, 0.5, n),
        }
    )

    fig1 = timeplot(
        df,
        xplot="time",
        yplot=[["angle_deg"], ["angular_velocity_dps", "temperature"]],
        plot_type="pyplot",
        title="angle & velocity (pyplot)",
        ylabel=["deg", "dps / degC"],
    )
    out1 = "timeplot_sample.png"
    fig1.savefig(out1)
    print(f"saved: {out1}")

    fig2 = timeplot(
        df,
        xplot="time",
        yplot=["angle_deg", "temperature"],
        plot_type="plotly",
        title="angle & temperature (plotly)",
    )
    out2 = "timeplot_sample.html"
    fig2.write_html(out2)
    print(f"saved: {out2}")


if __name__ == "__main__":
    main()