plot_psi.py

#!/usr/bin/env python3
"""
plot_psi.py - Plot stream function contours from meshoutSQP.dat

Generates a PNG file showing PSI contour lines similar to OP5.png style:
- Blue streamlines in the flow region
- Red domain boundary
- Blue-filled stagnant region below the inner boundary

Usage:
    python plot_psi.py [meshfile] [output.png]

Default: meshoutSQP.dat -> solutionSQP.png
"""

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.tri as tri
import sys
import os

def read_tecplot_mesh(filename):
    """Read Tecplot FE mesh file and return points, triangles, and psi values."""
    with open(filename, 'r') as f:
        lines = f.readlines()

    # Parse header
    n_points = 0
    n_triangles = 0
    for line in lines:
        if 'ZONE' in line:
            parts = line.replace(',', ' ').split()
            for p in parts:
                if p.startswith('N='):
                    n_points = int(p[2:])
                elif p.startswith('E='):
                    n_triangles = int(p[2:])
            break

    if n_points == 0 or n_triangles == 0:
        raise ValueError(f"Could not parse Tecplot header from {filename}")

    # Find data start
    data_start = 0
    for i, line in enumerate(lines):
        if 'ZONE' in line:
            data_start = i + 1
            break

    # Read points (z, r, psi)
    z = np.zeros(n_points)
    r = np.zeros(n_points)
    psi = np.zeros(n_points)

    for i in range(n_points):
        parts = lines[data_start + i].split()
        z[i] = float(parts[0])
        r[i] = float(parts[1])
        psi[i] = float(parts[2])

    # Read triangles (1-based to 0-based)
    triangles = np.zeros((n_triangles, 3), dtype=int)
    for i in range(n_triangles):
        parts = lines[data_start + n_points + i].split()
        triangles[i] = [int(p) - 1 for p in parts[:3]]

    return z, r, psi, triangles

def read_gambit_boundary(filename):
    """Read inner boundary from gambit.jou file."""
    inner_z = []
    inner_r = []

    with open(filename, 'r') as f:
        for line in f:
            if 'vertex create coordinates' in line:
                parts = line.split()
                z_val = float(parts[3])
                r_val = float(parts[4])
                inner_z.append(z_val)
                inner_r.append(r_val)

    return np.array(inner_z), np.array(inner_r)

def plot_psi_contours(meshfile, outfile, n_contours=25):
    """Generate PSI contour plot similar to OP5.png."""

    print(f"Reading mesh from {meshfile}...")
    z, r, psi, triangles = read_tecplot_mesh(meshfile)
    print(f"  {len(z)} points, {len(triangles)} triangles")
    print(f"  z range: [{z.min():.3f}, {z.max():.3f}]")
    print(f"  r range: [{r.min():.3f}, {r.max():.3f}]")
    print(f"  psi range: [{psi.min():.6f}, {psi.max():.6f}]")

    # Try to read boundary from gambit.jou
    script_dir = os.path.dirname(os.path.abspath(meshfile))
    gambit_file = os.path.join(script_dir, "gambit.jou")

    # Create triangulation for contour plotting
    triang = tri.Triangulation(z, r, triangles)

    # Create figure
    fig, ax = plt.subplots(1, 1, figsize=(10, 6))

    # Domain geometry from code
    z_max = 4.0
    r_inlet_hub = 0.132
    r_inlet_shroud = 1.0
    r_outlet_hub = 0.55
    r_outlet_wall = 1.1
    z_cone_start = 3.06  # Where cone meets cylinder
    r_cone_top = 1.384

    # Read inner boundary from gambit.jou if available
    if os.path.exists(gambit_file):
        print(f"Reading boundary from {gambit_file}...")
        all_z, all_r = read_gambit_boundary(gambit_file)

        # Inner boundary: points where z increases monotonically from 0 to ~4
        # Detect end of inner boundary where z suddenly decreases (next section starts)
        inner_z = [all_z[0]]
        inner_r = [all_r[0]]
        for i in range(1, len(all_z)):
            if all_z[i] > inner_z[-1]:  # Still increasing
                inner_z.append(all_z[i])
                inner_r.append(all_r[i])
            else:
                break  # z decreased, inner boundary ended
        inner_z = np.array(inner_z)
        inner_r = np.array(inner_r)
        print(f"  Inner boundary: {len(inner_z)} points, z=[{inner_z[0]:.2f}, {inner_z[-1]:.2f}]")
    else:
        # Fallback: extract from mesh points near r minimum
        print("No gambit.jou found, extracting boundary from mesh...")
        # Get points that define the inner boundary (lowest r for each z)
        z_bins = np.linspace(z.min(), z.max(), 50)
        inner_z = []
        inner_r = []
        for i in range(len(z_bins) - 1):
            mask = (z >= z_bins[i]) & (z < z_bins[i+1])
            if np.any(mask):
                inner_z.append((z_bins[i] + z_bins[i+1]) / 2)
                inner_r.append(r[mask].min())
        inner_z = np.array(inner_z)
        inner_r = np.array(inner_r)

    # Plot PSI contours (streamlines) in blue - flow region only
    levels = np.linspace(psi.min() + 0.001, psi.max() - 0.001, n_contours)
    cs = ax.tricontour(triang, psi, levels=levels, colors='blue', linewidths=0.8, zorder=2)

    # Draw domain boundary in red - trace the perimeter properly
    lw = 2.0

    # Get actual boundary endpoints from inner boundary data
    z_inner_start, r_inner_start = inner_z[0], inner_r[0]  # Should be (0, 0.132)
    z_inner_end, r_inner_end = inner_z[-1], inner_r[-1]    # Should be (~3.9, ~0.52)

    # 1. Inner boundary (stagnant zone edge) - already have these points
    ax.plot(inner_z, inner_r, 'r-', linewidth=lw, zorder=3)

    # 2. Connect inner boundary end to outlet hub
    ax.plot([z_inner_end, z_max], [r_inner_end, r_outlet_hub], 'r-', linewidth=lw, zorder=3)

    # 3. Outlet: vertical line at z=z_max from r_outlet_hub to r_outlet_wall
    ax.plot([z_max, z_max], [r_outlet_hub, r_outlet_wall], 'r-', linewidth=lw, zorder=3)

    # 4. Upper boundary: cylinder portion (z_max to z_cone_start)
    ax.plot([z_max, z_cone_start], [r_outlet_wall, r_cone_top], 'r-', linewidth=lw, zorder=3)

    # 5. Upper boundary: cone portion (z_cone_start to inlet shroud at z=1)
    ax.plot([z_cone_start, 1.0], [r_cone_top, r_inlet_shroud], 'r-', linewidth=lw, zorder=3)

    # 6. Tilted inlet: from shroud (1, 1.0) down to inner boundary start
    ax.plot([1.0, z_inner_start], [r_inlet_shroud, r_inner_start], 'r-', linewidth=lw, zorder=3)

    # Labels and formatting
    ax.set_xlabel('dimensionless axial coordinate', fontsize=12)
    ax.set_ylabel('dimensionless radial coordinate', fontsize=12)
    ax.set_xlim(-0.1, z_max + 0.2)
    ax.set_ylim(0, r_cone_top + 0.15)
    ax.set_aspect('equal', adjustable='box')

    # Clean up axes
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)

    plt.tight_layout()
    plt.savefig(outfile, dpi=150, bbox_inches='tight', facecolor='white')
    print(f"Saved plot to {outfile}")
    plt.close()

    return True

def main():
    # Default files
    script_dir = os.path.dirname(os.path.abspath(__file__))
    meshfile = os.path.join(script_dir, "meshoutSQP.dat")
    outfile = os.path.join(script_dir, "solutionSQP.png")

    # Command line arguments
    if len(sys.argv) > 1:
        meshfile = sys.argv[1]
    if len(sys.argv) > 2:
        outfile = sys.argv[2]

    if not os.path.exists(meshfile):
        print(f"Error: Mesh file not found: {meshfile}")
        sys.exit(1)

    plot_psi_contours(meshfile, outfile)

if __name__ == "__main__":
    main()