Source code for tud_lbm.io.plotting.run_comparison

"""Cross-run comparison plots and batch CSV-export CLI entry point."""

from __future__ import annotations
import sys
from pathlib import Path
from typing import TYPE_CHECKING
from typing import TypedDict
from matplotlib.colors import TABLEAU_COLORS
from tud_lbm.io.plotting.simulation_csv import _CSV_FILENAME
from tud_lbm.io.plotting.simulation_csv import build_simulation_csv

if TYPE_CHECKING:
    import matplotlib.axes
    import pandas as pd
    from tud_lbm.config import SimulationConfig

_DIR_SPLIT_PARTS = 2
_LABEL_CA = r"$\mathrm{Ca}$"
_LABEL_IT_NORM = r"$\Delta\mathrm{t}/\mathrm{t}_{\mathrm{max}}$"
_LABEL_X_AVG_NORM = r"$X_{\mathrm{avg}}/R_0$"
_CONFIG_TOML = "config.toml"
_COMPARISON_DIR = "comparison_analysis"


class _PlotConfig(TypedDict, total=False):
    filename: str
    x: str
    x_pair: tuple[str, str]
    y: str
    y_pair: tuple[str, str]
    pair_labels: tuple[str, str]
    xlabel: str
    ylabel: str


_COMPARISON_PLOT_CONFIGS: list[_PlotConfig] = [
    {
        "filename": "01_Ca_vs_iteration.png",
        "x": "normalised_iteration",
        "y": "Ca",
        "xlabel": _LABEL_IT_NORM,
        "ylabel": _LABEL_CA,
    },
    {
        "filename": "02_Ca_vs_x_avg.png",
        "x": "avg_x_location_norm",
        "y": "Ca",
        "xlabel": _LABEL_X_AVG_NORM,
        "ylabel": _LABEL_CA,
    },
    {
        "filename": "03_Ca_contact_line_vs_iteration.png",
        "x": "normalised_iteration",
        "y_pair": ("Ca_cll_left", "Ca_cll_right"),
        "pair_labels": ("Trailing", "Leading"),
        "xlabel": _LABEL_IT_NORM,
        "ylabel": _LABEL_CA,
    },
    {
        "filename": "04_Ca_contact_line_vs_x_avg.png",
        "x": "avg_x_location_norm",
        "y_pair": ("Ca_cll_left", "Ca_cll_right"),
        "pair_labels": ("Trailing", "Leading"),
        "xlabel": _LABEL_X_AVG_NORM,
        "ylabel": _LABEL_CA,
    },
    {
        "filename": "05_Ca_cm_vs_iteration.png",
        "x": "normalised_iteration",
        "y": "Ca_cm",
        "xlabel": _LABEL_IT_NORM,
        "ylabel": _LABEL_CA,
    },
    {
        "filename": "06_Ca_cm_vs_x_avg.png",
        "x": "avg_x_location_norm",
        "y": "Ca_cm",
        "xlabel": _LABEL_X_AVG_NORM,
        "ylabel": _LABEL_CA,
    },
    {
        "filename": "07_contact_angles_vs_iteration.png",
        "x": "normalised_iteration",
        "y_pair": ("ca_left", "ca_right"),
        "pair_labels": ("Trailing", "Leading"),
        "xlabel": _LABEL_IT_NORM,
        "ylabel": "Contact angle (degrees)",
    },
    {
        "filename": "08_contact_angles_vs_x_avg.png",
        "x": "avg_x_location_norm",
        "y_pair": ("ca_left", "ca_right"),
        "pair_labels": ("Trailing", "Leading"),
        "xlabel": _LABEL_X_AVG_NORM,
        "ylabel": r"$\theta$ (degrees)",
    },
    {
        "filename": "09_contact_angle_vs_Ca_cll.png",
        "x_pair": ("Ca_cll_left", "Ca_cll_right"),
        "y_pair": ("ca_left", "ca_right"),
        "pair_labels": ("Trailing", "Leading"),
        "xlabel": r"$\mathrm{Ca}_{\mathrm{CL}}$ [−]",
        "ylabel": r"$\theta$ (degrees)",
    },
]

_COMPARISON_MARKERS = ["o", "s", "D", "^", "v", "<", ">", "p", "*", "h", "H", "+", "x"]


def _safe_load_config(toml_path: Path) -> SimulationConfig | None:
    from tud_lbm.io.readers import TomlAdapter

    try:
        return TomlAdapter().load(str(toml_path))
    except Exception as exc:  # noqa: BLE001
        print(f"  Skipped {toml_path.parent}: {exc}")
        return None


def _load_comparison_entries(parent_dir: Path) -> list[dict]:
    """Collect all processed run CSVs under *parent_dir*, sorted by advancing CA."""
    try:
        import pandas as pd
    except ImportError:
        return []

    entries = []
    for csv_path in sorted(parent_dir.rglob(_CSV_FILENAME)):
        run_dir = csv_path.parent
        if _COMPARISON_DIR in run_dir.parts:
            continue
        toml_path = run_dir / _CONFIG_TOML
        if not toml_path.exists():
            continue
        config = _safe_load_config(toml_path)
        if config is None:
            continue

        adv_ca = config.wetting_config.get("advancing_ca", 0) if config.wetting_config else 0
        label = config.simulation_name or _clean_dir_label(run_dir.name)
        entries.append(
            {
                "label": label,
                "sort_key": adv_ca,
                "data": pd.read_csv(csv_path),
            }
        )

    return sorted(entries, key=lambda e: e["sort_key"])


def _clean_dir_label(dir_name: str) -> str:
    parts = dir_name.split("_", 1)
    name = parts[1] if len(parts) == _DIR_SPLIT_PARTS and parts[0].replace("-", "").isdigit() else dir_name
    return name.replace("_", " ").capitalize()


def _plot_comparison_entry(
    ax: matplotlib.axes.Axes,
    df: pd.DataFrame,
    pc: _PlotConfig,
    x_col: str,
    color: str,
    marker: str,
    label: str,
) -> None:
    """Plot one entry from _COMPARISON_PLOT_CONFIGS for a single run."""
    if "y_pair" in pc:
        y1, y2 = pc["y_pair"]
        l1, l2 = pc["pair_labels"]
        if "x_pair" in pc:
            x1_col, x2_col = pc["x_pair"]
            if x1_col in df.columns and y1 in df.columns:
                ax.scatter(df[x1_col], df[y1], marker=marker, s=15, color=color, alpha=1.0, label=f"{label} ({l1})")
            if x2_col in df.columns and y2 in df.columns:
                ax.scatter(df[x2_col], df[y2], marker=marker, s=15, color=color, alpha=0.5, label=f"{label} ({l2})")
        else:
            if y1 in df.columns:
                ax.scatter(df[x_col], df[y1], marker=marker, s=15, color=color, alpha=1.0, label=f"{label} ({l1})")
            if y2 in df.columns:
                ax.scatter(df[x_col], df[y2], marker=marker, s=15, color=color, alpha=0.5, label=f"{label} ({l2})")
    else:
        y_col = pc["y"]
        if y_col in df.columns:
            ax.scatter(df[x_col], df[y_col], marker=marker, s=15, color=color, label=label)


[docs] def compare_runs(parent_dir: str | Path) -> None: """Generate comparison scatter plots across all processed runs. Reads ``simulation_data.csv`` files produced by :func:`build_simulation_csv` from every run directory under *parent_dir* and writes plots to ``<parent_dir>/comparison_analysis/``. Args: parent_dir: Parent directory that contains the individual run directories. """ import matplotlib.pyplot as plt from tud_lbm.io.plotting.figure_config import DEFAULT_STYLE parent_dir = Path(parent_dir) entries = _load_comparison_entries(parent_dir) if not entries: print("No processed simulation data found for comparison.") return out_dir = parent_dir / _COMPARISON_DIR out_dir.mkdir(parents=True, exist_ok=True) colors = list(TABLEAU_COLORS.values()) for pc in _COMPARISON_PLOT_CONFIGS: fig, ax = plt.subplots(figsize=DEFAULT_STYLE.comparison_figsize) for i, entry in enumerate(entries): color = str(colors[i % len(colors)]) marker = _COMPARISON_MARKERS[i % len(_COMPARISON_MARKERS)] x_col = pc["x_pair"][0] if "x_pair" in pc else pc["x"] if x_col not in entry["data"].columns: continue _plot_comparison_entry(ax, entry["data"], pc, x_col, color, marker, entry["label"]) ax.set_xlabel(pc["xlabel"], fontsize=DEFAULT_STYLE.comparison_axis_label_fontsize) ax.set_ylabel(pc["ylabel"], fontsize=DEFAULT_STYLE.comparison_axis_label_fontsize) ax.tick_params(axis="both", labelsize=DEFAULT_STYLE.comparison_tick_label_fontsize) handles, labels = ax.get_legend_handles_labels() if handles: ax.legend(handles, labels, fontsize=DEFAULT_STYLE.comparison_legend_fontsize, loc="best") ax.grid(False) plt.tight_layout() out_path = out_dir / pc["filename"] fig.savefig(out_path, dpi=DEFAULT_STYLE.dpi, bbox_inches="tight") plt.close(fig) print(f"Saved {out_path}")
[docs] def process_parent_dir( parent_dir: str | Path, fields: list[str] | None = None, ) -> tuple[int, int]: """Discover runs, export per-run CSVs, and generate comparison plots. Discovers runs by searching for ``config.toml`` files recursively. Skips directories whose path contains ``"init"`` or ``"comparison_analysis"``. Writes ``simulation_data.csv`` per run, then generates 8 comparison plots when at least one run produced CSV data. Args: parent_dir: Absolute or relative path to the parent directory. fields: Comparison operator names to run as per-run analysis plots. When ``None``, no per-run analysis plots are generated beyond the CSV. Returns: A tuple ``(n_runs_found, n_runs_with_csv)``. """ from tud_lbm.io.plotting.figure_builder import FigureBuilder parent = Path(parent_dir) skip_dirs = {"init", _COMPARISON_DIR} run_dirs: list[Path] = [] seen: set[Path] = set() for toml in sorted(parent.rglob(_CONFIG_TOML)): rd = toml.parent if rd in seen or any(s in str(rd).lower() for s in skip_dirs): continue seen.add(rd) run_dirs.append(rd) if not run_dirs: return 0, 0 print(f"Found {len(run_dirs)} simulation(s) to process.") n_ok = 0 for rd in run_dirs: config = _safe_load_config(rd / _CONFIG_TOML) if config is None: continue if build_simulation_csv(rd, config) is not None: n_ok += 1 if fields: FigureBuilder(config=config, run_dir=rd, fields=fields).build_analysis() if n_ok > 0: print("\nGenerating comparison plots...") compare_runs(parent) print(f"Done. Comparison plots in {parent / _COMPARISON_DIR}") return len(run_dirs), n_ok
[docs] def main(parent_dir: str | None = None) -> None: """Script entry point for batch CSV export and run comparison plots.""" if parent_dir is None: parent_dir = input("Enter parent directory (absolute path): ").strip() parent = Path(parent_dir) if not parent.is_dir(): print(f"Directory not found: {parent}") sys.exit(1) n_runs, n_ok = process_parent_dir(parent_dir) if n_runs == 0: print("No simulation run directories found.") sys.exit(1) if n_ok == 0: print("No runs produced CSV data. Check sim_type and data files.") sys.exit(1)
if __name__ == "__main__": main(sys.argv[1] if len(sys.argv) > 1 else None)