postProcess/plotFootPrint.py

# Author: Vatsal Sanjay
# [email protected]
# CoMPhy Lab — Durham University
# Last updated: Nov 17, 2025

from __future__ import annotations

import argparse
import glob
import os
import re
import shutil
from dataclasses import dataclass
from typing import List, Optional, Tuple

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

# Configure matplotlib for publication-quality output with LaTeX fallback
matplotlib.rcParams['font.family'] = 'serif'
matplotlib.rcParams['font.serif'] = ['Computer Modern Roman']
if shutil.which('latex'):
    try:
        matplotlib.rcParams['text.usetex'] = True
        matplotlib.rcParams['text.latex.preamble'] = r'\usepackage{amsmath}'
    except Exception:
        matplotlib.rcParams['text.usetex'] = False
else:
    matplotlib.rcParams['text.usetex'] = False

# Font settings for 6x4 inch figure (smaller than default 12x9)
PLT_SETTINGS = {
    'LabelFont': 28,      # Axis labels
    'AxesFont': 20,       # Tick labels
    'LegendFont': 18,     # Legend entries
}


@dataclass(frozen=True)
class FootprintData:

    cutoff: float
    time: np.ndarray
    radius: np.ndarray

    def __repr__(self) -> str:
        return f"FootprintData(cutoff={self.cutoff:.4g}, n_points={len(self.time)})"


def parse_arguments() -> argparse.Namespace:

    parser = argparse.ArgumentParser(
        description="Create publication-quality footprint evolution plots."
    )
    parser.add_argument(
        "--resultsDir",
        type=str,
        default="simulationCases/1000/results",
        help="Path to the directory containing rFootvsTime_*.csv files.",
    )
    parser.add_argument(
        "--figsize",
        type=float,
        nargs=2,
        default=[6.0, 4.5],
        help="Figure size in inches (width height). Default: 6.0 4.5",
    )
    parser.add_argument(
        "--gui",
        action="store_true",
        help="Enable interactive GUI display after saving.",
    )
    return parser.parse_args()


def discover_csv_files(results_dir: str) -> List[str]:

    pattern = os.path.join(results_dir, "rFootvsTime_*.csv")
    files = glob.glob(pattern)
    if not files:
        raise FileNotFoundError(
            f"No rFootvsTime_*.csv files found in {results_dir}"
        )
    return sorted(files)


def extract_cutoff_from_filename(filename: str) -> Optional[float]:

    basename = os.path.basename(filename)
    match = re.search(r"rFootvsTime_([0-9.]+)\.csv", basename)
    if match:
        return float(match.group(1))
    return None


def load_footprint_data(filepath: str) -> FootprintData:

    cutoff = extract_cutoff_from_filename(filepath)
    if cutoff is None:
        raise ValueError(f"Could not extract cutoff from filename: {filepath}")

    try:
        df = pd.read_csv(filepath)
    except Exception as err:
        raise IOError(f"Failed to read {filepath}: {err}")

    if "time" not in df.columns or "rf" not in df.columns:
        raise ValueError(
            f"CSV file {filepath} must contain 'time' and 'rf' columns. "
            f"Found: {df.columns.tolist()}"
        )

    # Remove rows where rf is exactly zero (no footprint detected)
    df_filtered = df[df["rf"] > 0].copy()

    if len(df_filtered) == 0:
        print(f"Warning: No non-zero footprint data in {filepath}")
        return FootprintData(
            cutoff=cutoff,
            time=np.array([]),
            radius=np.array([])
        )

    return FootprintData(
        cutoff=cutoff,
        time=df_filtered["time"].values,
        radius=df_filtered["rf"].values,
    )


def setup_axes(ax: plt.Axes) -> None:

    # Thick axis spines
    for spine in ax.spines.values():
        spine.set_linewidth(2.5)

    # Major ticks
    ax.tick_params(
        axis='both',
        which='major',
        labelsize=PLT_SETTINGS['AxesFont'],
        width=2.5,
        length=8,
        direction='out',
        pad=6
    )

    # Minor ticks
    ax.tick_params(
        which='minor',
        width=1.5,
        length=4,
        direction='out'
    )

    ax.minorticks_on()


def create_footprint_plot(
    datasets: List[FootprintData],
    figsize: Tuple[float, float] = (6.0, 4.5),
) -> Tuple[plt.Figure, plt.Axes]:

    fig, ax = plt.subplots(figsize=figsize)
    fig.set_facecolor('white')

    # Use viridis colormap for different cutoffs
    n_datasets = len(datasets)
    colors = plt.cm.viridis(np.linspace(0.1, 0.9, n_datasets))

    # Plot each cutoff as a separate line
    for data, color in zip(datasets, colors):
        if len(data.time) == 0:
            continue

        # Format cutoff for legend (scientific notation)
        if data.cutoff < 0.01:
            label = rf"$x_{{\text{{cutoff}}}} = {data.cutoff:.1e}$ m"
        else:
            label = rf"$x_{{\text{{cutoff}}}} = {data.cutoff:.4f}$ m"

        ax.plot(
            data.time,
            data.radius,
            linewidth=2.5,
            color=color,
            label=label,
            marker='o',
            markersize=4,
            markevery=max(1, len(data.time) // 20),  # ~20 markers per line
            zorder=2
        )

    # Axis labels (R_0 = 1, tau_0 = 1)
    ax.set_xlabel(
        r'$t/\tau_0$',
        fontsize=PLT_SETTINGS['LabelFont'],
        labelpad=8
    )
    ax.set_ylabel(
        r'$r_{\text{foot}}/R_0$',
        fontsize=PLT_SETTINGS['LabelFont'],
        labelpad=8
    )

    # Legend
    ax.legend(
        loc='best',
        frameon=True,
        fontsize=PLT_SETTINGS['LegendFont'],
        framealpha=0.9,
        edgecolor='gray',
        fancybox=False,
    )

    # Light grid for readability
    ax.grid(True, which='major', alpha=0.3, linestyle='-', linewidth=0.5)
    ax.grid(True, which='minor', alpha=0.15, linestyle='-', linewidth=0.3)

    # Apply publication styling
    setup_axes(ax)

    plt.tight_layout()
    return fig, ax


def save_figure(
    fig: plt.Figure,
    output_path: str,
    dpi: int = 300,
) -> None:

    pdf_path = output_path if output_path.endswith('.pdf') else output_path + ".pdf"

    fig.savefig(pdf_path, bbox_inches='tight', dpi=dpi, format='pdf')
    print(f"Saved figure: {pdf_path}")


def main() -> None:
    args = parse_arguments()

    # Resolve and validate results directory
    results_dir = os.path.abspath(args.resultsDir)
    if not os.path.isdir(results_dir):
        raise FileNotFoundError(f"Results directory not found: {results_dir}")

    print(f"Searching for CSV files in: {results_dir}")

    # Discover and load all footprint data files
    csv_files = discover_csv_files(results_dir)
    print(f"Found {len(csv_files)} CSV files")

    datasets: List[FootprintData] = []
    for csv_file in csv_files:
        try:
            data = load_footprint_data(csv_file)
            datasets.append(data)
            print(f"Loaded: {data}")
        except Exception as err:
            print(f"Warning: Skipping {csv_file}: {err}")

    if not datasets:
        raise ValueError("No valid footprint data loaded")

    # Sort by cutoff value for consistent legend ordering
    datasets.sort(key=lambda d: d.cutoff)

    # Create plot
    print(f"\nCreating plot with {len(datasets)} cutoff values...")
    fig, ax = create_footprint_plot(datasets, figsize=tuple(args.figsize))

    # Set output path to footprint_evolution.pdf in results directory
    output_path = os.path.join(results_dir, "footprint_evolution.pdf")

    # Save figure
    save_figure(fig, output_path)

    # Display GUI only if --gui flag is passed
    if args.gui:
        print("\nDisplaying figure (close window to exit)...")
        plt.show()
    else:
        plt.close(fig)

    print("\nDone!")


if __name__ == "__main__":
    main()