MTT69/TurbulentChannel

1#!/usr/bin/env python32"""3Paper-ready validation figures: clean + noisy data on the same axes.4 5Open symbols = Case A (ideal conditions)6Filled symbols = Case B (degraded, SNR ~8)7DNS reference = solid black line with 95% CI band8 9Produces:10  1. mean_velocity_comparison.png    — U+ vs y+11  2. stresses_comparison.png         — 1x3 subplots (uu+, vv+, -uv+)12  3. combined_stresses_comparison.png — all stresses on one axis13                                        (+ ww+ when stereo data is present)14 15Run indices are auto-detected: statistics are computed only for the final16PIV pass, so each mean_stats.mat holds exactly one populated run. The17ensemble uses its final pass. Coordinates are used as calibrated — no18y-offset fudges.19"""20 21import numpy as np22import scipy.io as sio23from scipy.interpolate import interp1d24import matplotlib.pyplot as plt25import matplotlib as mpl26from pathlib import Path27 28# ── Publication font setup: match LaTeX body text ─────────────────────────────29mpl.rcParams.update({30    'font.family': 'serif',31    'font.serif': ['CMU Serif', 'Computer Modern Roman', 'DejaVu Serif'],32    'mathtext.fontset': 'cm',33    'axes.unicode_minus': False,34    'text.usetex': False,35    'axes.labelsize': 11,36    'axes.titlesize': 11,37    'legend.fontsize': 9,38    'xtick.labelsize': 10,39    'ytick.labelsize': 10,40    'lines.linewidth': 1.5,41    'figure.dpi': 600,42    'savefig.dpi': 600,43    'savefig.bbox': 'tight',44    'savefig.pad_inches': 0.05,45})46 47# ── Okabe-Ito colourblind-safe palette ────────────────────────────────────────48COLORS = {49    'Instantaneous': '#0072B2',50    'Ensemble':      '#D55E00',51    'Stereo':        '#009E73',52}53MARKERS = {54    'Instantaneous': 'o',55    'Ensemble':      's',56    'Stereo':        '^',57}58DNS_COLOR = 'k'59 60 61# =============================================================================62# Data loading (reuse from benchmark_comparison)63# =============================================================================64 65def _load_gt(gt_dir):66    from benchmark_comparison import load_wall_units, load_ground_truth67    gt_dir = Path(gt_dir)68    for name in ('wall_units.mat', 'diagnostics.mat', 'direct_stats.mat'):69        p = gt_dir / name70        if p.exists():71            wu = load_wall_units(p)72            break73    for name in ('profiles.mat', 'ensemble_statistics_full.mat', 'direct_stats.mat'):74        p = gt_dir / name75        if p.exists():76            gt = load_ground_truth(p, wall_units_path=gt_dir / 'direct_stats.mat')77            break78    gt_plus = {79        'y_plus': gt['y_plus'], 'U_plus': gt['U_plus'],80        'uu_plus': gt['uu_plus'], 'vv_plus': gt['vv_plus'], 'uv_plus': gt['uv_plus'],81    }82    for key in ('ww_plus',83                'uu_plus_ci_lo', 'uu_plus_ci_hi', 'vv_plus_ci_lo', 'vv_plus_ci_hi',84                'ww_plus_ci_lo', 'ww_plus_ci_hi',85                'uv_plus_ci_lo', 'uv_plus_ci_hi', 'U_plus_ci_lo', 'U_plus_ci_hi'):86        if key in gt:87            gt_plus[key] = gt[key]88    return gt_plus, wu89 90 91def _populated_run_idx(mat_path, result_key='piv_result'):92    """Return the index of the single populated run in a stats .mat file.93 94    Statistics are computed only for the final PIV pass, so earlier run slots95    are zero-size placeholders. Fails loudly if zero or >1 runs hold data.96    """97    d = sio.loadmat(str(mat_path), squeeze_me=True, struct_as_record=False)98    runs = np.atleast_1d(d[result_key])99    populated = [i for i, r in enumerate(runs) if np.atleast_2d(r.ux).size > 1]100    if len(populated) != 1:101        raise ValueError(102            f'{mat_path}: expected exactly one populated run, '103            f'found {populated or "none"} of {len(runs)} slots')104    return populated[0]105 106 107def _load_inst(stats_path, wu):108    from benchmark_comparison import (109        load_piv_statistics, compute_piv_profiles, convert_to_wall_units)110    run_idx = _populated_run_idx(stats_path)111    piv = load_piv_statistics(Path(stats_path), run_idx=run_idx)112    prof = compute_piv_profiles(piv, x_exclude_vectors=4)113    return convert_to_wall_units(prof, wu)114 115 116def _load_ens(ens_path, coords_path, wu, run_idx=-1):117    from benchmark_comparison import (118        load_ensemble_statistics, compute_piv_profiles, convert_to_wall_units)119    ens = sio.loadmat(str(ens_path), squeeze_me=True, struct_as_record=False)120    passes = np.atleast_1d(ens['ensemble_result'])121    run_idx = range(len(passes))[run_idx]  # normalise negative index122    ws = np.atleast_1d(passes[run_idx].window_size)123    print(f'  Ensemble pass {run_idx}: window {ws[0]}x{ws[1]} px')124    piv = load_ensemble_statistics(Path(ens_path), Path(coords_path), run_idx=run_idx)125    prof = compute_piv_profiles(piv, x_exclude_vectors=4)126    return convert_to_wall_units(prof, wu)127 128 129def _load_stereo(stats_path, wu, trim_top=10):130    from benchmark_comparison import convert_to_wall_units131    run_idx = _populated_run_idx(stats_path)132    stats = sio.loadmat(str(stats_path), squeeze_me=True, struct_as_record=False)133    piv_s = stats['piv_result'][run_idx]134    coords_s = stats['coordinates'][run_idx]135    x, y = coords_s.x, coords_s.y136    valid_cols = np.any(~np.isnan(y), axis=0)137    col_indices = np.where(valid_cols)[0]138    mid_col = col_indices[len(col_indices) // 2]139    y_unique = y[:, mid_col]140    valid_rows = ~np.isnan(y_unique)141    y_unique = y_unique[valid_rows]142    if trim_top > 0:143        if y_unique[0] > y_unique[-1]:144            y_unique = y_unique[trim_top:]145            vi = np.where(valid_rows)[0][trim_top:]146        else:147            y_unique = y_unique[:-trim_top]148            vi = np.where(valid_rows)[0][:-trim_top]149        tm = np.zeros(valid_rows.shape, dtype=bool)150        tm[vi] = True151        valid_rows = tm152    xs = col_indices[0] + 4153    xe = col_indices[-1] - 3154    x_mask = np.zeros(x.shape[1], dtype=bool)155    x_mask[xs:xe] = True156    prof = {157        'y_mm': y_unique,158        'U': np.nanmean(piv_s.ux[valid_rows][:, x_mask] * 1000, axis=1),159        'V': np.nanmean(piv_s.uy[valid_rows][:, x_mask] * 1000, axis=1),160        'uu': np.nanmean(piv_s.uu[valid_rows][:, x_mask] * 1e6, axis=1),161        'vv': np.nanmean(piv_s.vv[valid_rows][:, x_mask] * 1e6, axis=1),162        'ww': np.nanmean(piv_s.ww[valid_rows][:, x_mask] * 1e6, axis=1),163        'uv': np.nanmean(piv_s.uv[valid_rows][:, x_mask] * 1e6, axis=1),164    }165    return convert_to_wall_units(prof, wu)166 167 168def _trim(plus, n=1):169    """Remove n near-wall points."""170    if n <= 0:171        return plus172    yp = plus['y_plus']173    if yp[0] > yp[-1]:174        sl = slice(None, -n)175    else:176        sl = slice(n, None)177    return {k: (v[sl] if isinstance(v, np.ndarray) and len(v) > n else v)178            for k, v in plus.items()}179 180 181# =============================================================================182# Plotting helpers183# =============================================================================184 185def _ci_band(ax, yp, lo, hi, sign=1):186    if sign == -1:187        lo, hi = hi, lo188    ax.fill_between(yp, sign * lo, sign * hi,189                    color=DNS_COLOR, alpha=0.10, linewidth=0)190 191 192def _plot_method(ax, yp, vals, method, filled=True, label=None, ms=3.5, alpha=0.7):193    """Plot a single method series — filled or open markers."""194    color = COLORS[method]195    marker = MARKERS[method]196    if filled:197        ax.plot(yp, vals, color=color, marker=marker, markersize=ms,198                alpha=alpha, linestyle='none', label=label, zorder=5)199    else:200        ax.plot(yp, vals, marker=marker, markersize=ms, alpha=alpha,201                linestyle='none', label=label, zorder=4,202                markerfacecolor='none', markeredgecolor=color, markeredgewidth=0.8)203 204 205# =============================================================================206# Figure 1: Mean velocity207# =============================================================================208 209def plot_velocity(gt_plus, clean, noisy, wu, output_dir):210    Re_tau = wu['Re_tau']211    fig, ax = plt.subplots(figsize=(7, 5))212 213    # DNS + CI214    if 'U_plus_ci_lo' in gt_plus:215        _ci_band(ax, gt_plus['y_plus'], gt_plus['U_plus_ci_lo'], gt_plus['U_plus_ci_hi'])216    ax.semilogx(gt_plus['y_plus'], gt_plus['U_plus'], color=DNS_COLOR,217                linewidth=2, label='DNS', zorder=10)218 219    # Clean (open)220    for method, plus in clean.items():221        _plot_method(ax, plus['y_plus'], plus['U_plus'], method, filled=False,222                     label=f'{method} — Case A')223 224    # Noisy (filled)225    for method, plus in noisy.items():226        _plot_method(ax, plus['y_plus'], plus['U_plus'], method, filled=True,227                     label=f'{method} — Case B')228 229    ax.set_xlabel(r'$y^+$')230    ax.set_ylabel(r'$U^+$')231    ax.set_xlim(1, Re_tau)232    ax.set_ylim(0, 25)233    ax.grid(True, alpha=0.25, linewidth=0.5)234    ax.legend(loc='lower right', framealpha=0.9)235 236    fig.tight_layout()237    out = Path(output_dir)238    out.mkdir(parents=True, exist_ok=True)239    fig.savefig(out / 'mean_velocity_comparison.png')240    fig.savefig(out / 'mean_velocity_comparison.pdf')241    plt.close(fig)242    print(f'  Saved: {out / "mean_velocity_comparison.png"}')243 244 245# =============================================================================246# Figure 2: Stresses — 1x3 subplots247# =============================================================================248 249def plot_stresses_subplots(gt_plus, clean, noisy, wu, output_dir):250    Re_tau = wu['Re_tau']251    has_ci = 'uu_plus_ci_lo' in gt_plus252 253    panels = [254        ('uu_plus', r"$\overline{u'u'}^+$", 1),255        ('vv_plus', r"$\overline{v'v'}^+$", 1),256        ('uv_plus', r"$-\overline{u'v'}^+$", -1),257    ]258 259    fig, axes = plt.subplots(1, 3, figsize=(7, 2.8))260 261    for ax, (var, ylabel, sign) in zip(axes, panels):262        # CI band263        ci_lo_key, ci_hi_key = f'{var}_ci_lo', f'{var}_ci_hi'264        if has_ci and ci_lo_key in gt_plus:265            _ci_band(ax, gt_plus['y_plus'], gt_plus[ci_lo_key], gt_plus[ci_hi_key], sign=sign)266 267        # DNS268        ax.plot(gt_plus['y_plus'], sign * gt_plus[var], color=DNS_COLOR,269                linewidth=1.8, label='DNS', zorder=10)270 271        # Clean (open)272        for method, plus in clean.items():273            _plot_method(ax, plus['y_plus'], sign * plus[var], method,274                         filled=False, label=f'{method} — A', ms=2.5, alpha=0.65)275 276        # Noisy (filled)277        for method, plus in noisy.items():278            _plot_method(ax, plus['y_plus'], sign * plus[var], method,279                         filled=True, label=f'{method} — B', ms=2.5, alpha=0.65)280 281        ax.set_xlabel(r'$y^+$')282        ax.set_ylabel(ylabel)283        ax.set_xscale('log')284        ax.set_xlim(1, Re_tau)285        ax.grid(True, alpha=0.25, linewidth=0.5)286 287    # Shared legend288    handles, labels = axes[0].get_legend_handles_labels()289    fig.legend(handles, labels, loc='upper center', ncol=4,290               bbox_to_anchor=(0.5, 1.05), framealpha=0.9, fontsize=8)291 292    fig.tight_layout()293    fig.subplots_adjust(top=0.82)294    out = Path(output_dir)295    out.mkdir(parents=True, exist_ok=True)296    fig.savefig(out / 'stresses_comparison.png')297    fig.savefig(out / 'stresses_comparison.pdf')298    plt.close(fig)299    print(f'  Saved: {out / "stresses_comparison.png"}')300 301 302# =============================================================================303# Figure 3: Combined stresses — single axis304# =============================================================================305 306def plot_combined_stresses(gt_plus, clean, noisy, wu, output_dir):307    Re_tau = wu['Re_tau']308    has_ci = 'uu_plus_ci_lo' in gt_plus309 310    component_styles = {311        'uu_plus': {'ls': '-',  'tex': r"$\overline{u'u'}^+$",  'sign': 1},312        'vv_plus': {'ls': '--', 'tex': r"$\overline{v'v'}^+$",  'sign': 1},313        'ww_plus': {'ls': '-.', 'tex': r"$\overline{w'w'}^+$",  'sign': 1},314        'uv_plus': {'ls': ':',  'tex': r"$-\overline{u'v'}^+$", 'sign': -1},315    }316    # Only draw components at least one method provides (ww is stereo-only)317    available = {var for plus in list(clean.values()) + list(noisy.values())318                 for var in plus}319    component_styles = {var: csty for var, csty in component_styles.items()320                        if var in available and var in gt_plus}321 322    fig, ax = plt.subplots(figsize=(7, 5))323 324    # DNS reference lines + CI bands325    for var, csty in component_styles.items():326        sign = csty['sign']327        ci_lo, ci_hi = f'{var}_ci_lo', f'{var}_ci_hi'328        if has_ci and ci_lo in gt_plus:329            _ci_band(ax, gt_plus['y_plus'], gt_plus[ci_lo], gt_plus[ci_hi], sign=sign)330        ax.plot(gt_plus['y_plus'], sign * gt_plus[var],331                color=DNS_COLOR, linewidth=1.8, linestyle=csty['ls'], zorder=10)332 333    # Clean (open) — all components334    for method, plus in clean.items():335        for var, csty in component_styles.items():336            if var not in plus:337                continue338            _plot_method(ax, plus['y_plus'], csty['sign'] * plus[var], method,339                         filled=False, ms=2.5, alpha=0.55)340 341    # Noisy (filled) — all components342    for method, plus in noisy.items():343        for var, csty in component_styles.items():344            if var not in plus:345                continue346            _plot_method(ax, plus['y_plus'], csty['sign'] * plus[var], method,347                         filled=True, ms=2.5, alpha=0.55)348 349    # ── Two-part legend ──────────────────────────────────────────────────350    # Part 1: method + condition351    method_handles = [352        plt.Line2D([], [], color=DNS_COLOR, linewidth=1.8, linestyle='-', label='DNS')353    ]354    for method in list(clean.keys()) + [m for m in noisy if m not in clean]:355        c = COLORS[method]356        m = MARKERS[method]357        # Open (Case A)358        method_handles.append(359            plt.Line2D([], [], color=c, marker=m, markersize=5, linestyle='none',360                       markerfacecolor='none', markeredgecolor=c, markeredgewidth=0.8,361                       label=f'{method} — Case A'))362        # Filled (Case B)363        method_handles.append(364            plt.Line2D([], [], color=c, marker=m, markersize=5, linestyle='none',365                       label=f'{method} — Case B'))366 367    # Part 2: component line styles368    comp_handles = []369    for var, csty in component_styles.items():370        comp_handles.append(371            plt.Line2D([], [], color='gray', linewidth=1.5,372                       linestyle=csty['ls'], label=csty['tex']))373 374    leg1 = ax.legend(handles=method_handles, loc='upper right',375                     framealpha=0.9, title='Method')376    ax.add_artist(leg1)377    ax.legend(handles=comp_handles, loc='upper left',378              framealpha=0.9, title='Component')379 380    ax.set_xlabel(r'$y^+$')381    ax.set_ylabel(r'Stress$^+$')382    ax.set_xscale('log')383    ax.set_xlim(1, Re_tau)384    ax.grid(True, alpha=0.25, linewidth=0.5)385 386    fig.tight_layout()387    out = Path(output_dir)388    out.mkdir(parents=True, exist_ok=True)389    fig.savefig(out / 'combined_stresses_comparison.png')390    fig.savefig(out / 'combined_stresses_comparison.pdf')391    plt.close(fig)392    print(f'  Saved: {out / "combined_stresses_comparison.png"}')393 394 395# =============================================================================396# Main397# =============================================================================398 399if __name__ == '__main__':400    import argparse401    parser = argparse.ArgumentParser(402        description='Generate combined clean+noisy paper figures. '403                    'All path arguments accept either the Case A (clean) or Case B (noisy) '404                    'variant; if only noisy paths are supplied, only Case B is plotted.')405    parser.add_argument('--output-dir', '-o', type=str, required=True,406                        help='Output directory for generated figures')407 408    # Ground truth (required — at least one of the two must be provided)409    parser.add_argument('--gt-clean-dir', type=str, default=None,410                        help='Directory containing direct_stats.mat for Case A (clean / 85k particles)')411    parser.add_argument('--gt-noisy-dir', type=str, default=None,412                        help='Directory containing direct_stats.mat for Case B (noisy / 22k particles)')413 414    # Case A (clean) PIV results415    parser.add_argument('--inst-clean-stats', type=str, default=None,416                        help='Path to instantaneous mean_stats.mat for Case A')417    parser.add_argument('--ens-clean-dir', type=str, default=None,418                        help='Directory containing ensemble_result.mat + coordinates.mat for Case A')419    parser.add_argument('--stereo-clean-stats', type=str, default=None,420                        help='Path to stereo mean_stats.mat for Case A')421 422    # Case B (noisy) PIV results423    parser.add_argument('--inst-noisy-stats', type=str, default=None,424                        help='Path to instantaneous mean_stats.mat for Case B')425    parser.add_argument('--ens-noisy-dir', type=str, default=None,426                        help='Directory containing ensemble_result.mat + coordinates.mat for Case B')427    parser.add_argument('--stereo-noisy-stats', type=str, default=None,428                        help='Path to stereo mean_stats.mat for Case B')429 430    parser.add_argument('--ens-pass', type=int, default=-1,431                        help='Ensemble pass index to plot (default -1 = final pass)')432 433    args = parser.parse_args()434    output_dir = Path(args.output_dir)435 436    if not args.gt_clean_dir and not args.gt_noisy_dir:437        parser.error('At least one of --gt-clean-dir / --gt-noisy-dir must be provided')438 439    # Ground truth: prefer clean (85k particles, tighter CI) for reference axes440    gt_dir_primary = args.gt_clean_dir or args.gt_noisy_dir441    gt_plus, wu = _load_gt(Path(gt_dir_primary))442    print(f"DNS: Re_tau={wu['Re_tau']:.0f}")443 444    # ── Case A (clean) ───────────────────────────────────────────────────445    clean = {}446    if args.inst_clean_stats or args.ens_clean_dir or args.stereo_clean_stats:447        print("\nLoading Case A (ideal)...")448        if args.inst_clean_stats:449            inst_clean = _trim(_load_inst(Path(args.inst_clean_stats), wu=wu))450            print(f"  Instantaneous: y+={inst_clean['y_plus'].min():.1f}-{inst_clean['y_plus'].max():.1f}")451            clean['Instantaneous'] = inst_clean452        if args.ens_clean_dir:453            ens_dir = Path(args.ens_clean_dir)454            ens_clean = _load_ens(455                ens_dir / 'ensemble_result.mat',456                ens_dir / 'coordinates.mat', wu=wu, run_idx=args.ens_pass)457            print(f"  Ensemble:      y+={ens_clean['y_plus'].min():.1f}-{ens_clean['y_plus'].max():.1f}")458            clean['Ensemble'] = ens_clean459        if args.stereo_clean_stats:460            stereo_clean = _trim(_load_stereo(Path(args.stereo_clean_stats), wu=wu))461            print(f"  Stereo:        y+={stereo_clean['y_plus'].min():.1f}-{stereo_clean['y_plus'].max():.1f}")462            clean['Stereo'] = stereo_clean463 464    # ── Case B (noisy) ───────────────────────────────────────────────────465    noisy = {}466    if args.inst_noisy_stats or args.ens_noisy_dir or args.stereo_noisy_stats:467        print("\nLoading Case B (degraded, SNR ~8)...")468        wu_n = wu  # fallback to clean wu469        if args.gt_noisy_dir:470            _, wu_n = _load_gt(Path(args.gt_noisy_dir))471        if args.inst_noisy_stats:472            inst_noisy = _trim(_load_inst(Path(args.inst_noisy_stats), wu=wu_n))473            print(f"  Instantaneous: y+={inst_noisy['y_plus'].min():.1f}-{inst_noisy['y_plus'].max():.1f}")474            noisy['Instantaneous'] = inst_noisy475        if args.ens_noisy_dir:476            ens_dir_n = Path(args.ens_noisy_dir)477            ens_noisy = _load_ens(478                ens_dir_n / 'ensemble_result.mat',479                ens_dir_n / 'coordinates.mat', wu=wu_n, run_idx=args.ens_pass)480            print(f"  Ensemble:      y+={ens_noisy['y_plus'].min():.1f}-{ens_noisy['y_plus'].max():.1f}")481            noisy['Ensemble'] = ens_noisy482        if args.stereo_noisy_stats:483            stereo_noisy = _trim(_load_stereo(Path(args.stereo_noisy_stats), wu=wu_n))484            print(f"  Stereo:        y+={stereo_noisy['y_plus'].min():.1f}-{stereo_noisy['y_plus'].max():.1f}")485            noisy['Stereo'] = stereo_noisy486 487    if not clean and not noisy:488        parser.error('At least one PIV result path must be provided (--*-clean-* or --*-noisy-*)')489 490    # ── Generate figures ─────────────────────────────────────────────────491    print("\nGenerating figures...")492    plot_velocity(gt_plus, clean, noisy, wu, output_dir)493    plot_stresses_subplots(gt_plus, clean, noisy, wu, output_dir)494    plot_combined_stresses(gt_plus, clean, noisy, wu, output_dir)495    print("Done.")496