visualize.py

Go to the documentation of this file.

 
 
import sys
import os
import numpy as np
import matplotlib.pyplot as plt
from scipy.interpolate import griddata
 
 
STRESS_COMPONENTS = {
    'xx': (3, r'$\sigma_{xx}$'),
    'yy': (4, r'$\sigma_{yy}$'),
    'zz': (5, r'$\sigma_{zz}$'),
    'xy': (6, r'$\sigma_{xy}$'),
    'xz': (7, r'$\sigma_{xz}$'),
    'yz': (8, r'$\sigma_{yz}$'),
}
 
 
def read_file(filename):
    data = []
    with open(filename, 'r') as f:
        lines = f.readlines()
        for line in lines[2:]:  # Skip count and metadata header
            line = line.strip()
            if not line:
                continue
            try:
                parts = list(map(float, line.split()))
                if len(parts) >= 9:
                    data.append(parts[:9])
            except ValueError:
                continue
    return data
 
 
def plot_contour(data, output_fname, component_index, component_label):
    x = [row[0] for row in data]
    y = [row[1] for row in data]  # using x-y plane for plotting
    sig = [row[component_index] for row in data]
 
    xi = np.linspace(min(x), max(x), 300)
    yi = np.linspace(min(y), max(y), 300)
    xi, yi = np.meshgrid(xi, yi)
 
    sig_grid = griddata((x, y), sig, (xi, yi), method='linear')
 
    plt.figure(figsize=(6, 5))
    contour = plt.contourf(xi, yi, sig_grid, levels=100, cmap='RdBu_r')
    plt.colorbar(contour, label='Stress (GPa)')
    plt.xlabel('x')
    plt.ylabel('y')
    plt.title(component_label)
    plt.tight_layout()
    plt.savefig(output_fname)
    plt.close()
 
# Main script execution
if __name__ == "__main__":
    if len(sys.argv) != 3:
        print("Usage: python visualize.py <stress_file> <component: xx|yy|zz|xy|xz|yz>")
        sys.exit(1)
 
    filename = sys.argv[1]
    component_key = sys.argv[2].lower()
 
    if component_key not in STRESS_COMPONENTS:
        print("Component must be one of: xx, yy, zz, xy, xz, yz")
        sys.exit(1)
 
    component_index, component_label = STRESS_COMPONENTS[component_key]
 
    base, _ = os.path.splitext(filename)
    plt_fname = f"{base}_{component_key}.pdf"
 
    data = read_file(filename)
 
    # Convert stress to GPa
    for row in data:
        row[component_index] *= 160.0
 
    plot_contour(data, plt_fname, component_index, component_label)