Atomizer/optimization_engine/extractors/extract_principal_stress.py

"""
Extract Principal Stresses from Structural Analysis
====================================================

Extracts principal stresses (sigma1, sigma2, sigma3) from OP2 files.
Useful for failure criteria beyond von Mises (e.g., Tresca, max principal).

Pattern: solid_stress
Element Types: CTETRA, CHEXA, CQUAD4, CTRIA3
Result Type: principal stress
API: model.op2_results.stress.ctetra_stress[subcase]

Phase 2 Task 2.3 - NX Open Automation Roadmap
"""

from pathlib import Path
from typing import Dict, Any, Optional, List, Literal
import numpy as np
from pyNastran.op2.op2 import OP2


# Column indices for principal stresses by element type
# Based on pyNastran stress data layout
STRESS_COLUMNS = {
    # Solid elements (CTETRA, CHEXA, CPENTA): 10 columns
    # [o1, o2, o3, ovm, omax_shear, ...] format varies
    'ctetra': {'o1': 0, 'o2': 1, 'o3': 2, 'von_mises': 9, 'max_shear': 8},
    'chexa': {'o1': 0, 'o2': 1, 'o3': 2, 'von_mises': 9, 'max_shear': 8},
    'cpenta': {'o1': 0, 'o2': 1, 'o3': 2, 'von_mises': 9, 'max_shear': 8},
    # Shell elements (CQUAD4, CTRIA3): 8 columns per surface
    # [fiber_dist, oxx, oyy, txy, angle, o1, o2, von_mises]
    'cquad4': {'o1': 5, 'o2': 6, 'von_mises': 7},
    'ctria3': {'o1': 5, 'o2': 6, 'von_mises': 7},
}


def extract_principal_stress(
    op2_file: Path,
    subcase: int = 1,
    element_type: str = 'ctetra',
    principal: Literal['max', 'mid', 'min', 'all'] = 'max'
) -> Dict[str, Any]:
    """
    Extract principal stresses from solid or shell elements.

    Principal stresses are the eigenvalues of the stress tensor,
    ordered as: sigma1 >= sigma2 >= sigma3 (o1 >= o2 >= o3)

    Args:
        op2_file: Path to .op2 file
        subcase: Subcase ID (default 1)
        element_type: Element type ('ctetra', 'chexa', 'cquad4', 'ctria3')
        principal: Which principal stress to return:
            - 'max': Maximum principal (sigma1, tension positive)
            - 'mid': Middle principal (sigma2)
            - 'min': Minimum principal (sigma3, compression negative)
            - 'all': Return all three principals

    Returns:
        dict: {
            'max_principal': Maximum principal stress value,
            'min_principal': Minimum principal stress value,
            'mid_principal': Middle principal stress (if applicable),
            'max_element': Element ID with maximum principal,
            'min_element': Element ID with minimum principal,
            'von_mises_max': Max von Mises for comparison,
            'element_type': Element type used,
            'subcase': Subcase ID,
            'units': 'MPa (model units)'
        }

    Example:
        >>> result = extract_principal_stress('model.op2', element_type='ctetra')
        >>> print(f"Max tension: {result['max_principal']:.2f} MPa")
        >>> print(f"Max compression: {result['min_principal']:.2f} MPa")
    """
    model = OP2()
    model.read_op2(str(op2_file))

    # Map element type to stress attribute
    stress_attr_map = {
        'ctetra': 'ctetra_stress',
        'chexa': 'chexa_stress',
        'cpenta': 'cpenta_stress',
        'cquad4': 'cquad4_stress',
        'ctria3': 'ctria3_stress'
    }

    element_type_lower = element_type.lower()
    stress_attr = stress_attr_map.get(element_type_lower)
    if not stress_attr:
        raise ValueError(f"Unknown element type: {element_type}. "
                        f"Supported: {list(stress_attr_map.keys())}")

    # Access stress through op2_results container
    stress_dict = None
    if hasattr(model, 'op2_results') and hasattr(model.op2_results, 'stress'):
        stress_container = model.op2_results.stress
        if hasattr(stress_container, stress_attr):
            stress_dict = getattr(stress_container, stress_attr)

    if stress_dict is None or not stress_dict:
        available = [a for a in dir(model) if 'stress' in a.lower()]
        raise ValueError(f"No {element_type} stress results in OP2. "
                        f"Available: {available}")

    # Get subcase
    available_subcases = list(stress_dict.keys())
    if subcase in available_subcases:
        actual_subcase = subcase
    else:
        actual_subcase = available_subcases[0]

    stress = stress_dict[actual_subcase]

    # Get column indices for this element type
    cols = STRESS_COLUMNS.get(element_type_lower)
    if cols is None:
        # Fallback: assume standard layout
        cols = {'o1': 0, 'o2': 1, 'o3': 2 if element_type_lower in ['ctetra', 'chexa', 'cpenta'] else None}

    itime = 0  # First time step (static analysis)

    # Extract principal stresses
    o1 = stress.data[itime, :, cols['o1']]  # Max principal
    o2 = stress.data[itime, :, cols['o2']]  # Mid principal

    # Solid elements have 3 principals, shells have 2
    if 'o3' in cols and cols['o3'] is not None:
        o3 = stress.data[itime, :, cols['o3']]  # Min principal
    else:
        o3 = None

    # Get element IDs
    element_ids = np.array([eid for (eid, node) in stress.element_node])

    # Find extremes
    max_o1_idx = np.argmax(o1)
    min_o1_idx = np.argmin(o1)

    result = {
        'max_principal': float(np.max(o1)),
        'max_principal_element': int(element_ids[max_o1_idx]),
        'min_principal_o1': float(np.min(o1)),
        'min_principal_o1_element': int(element_ids[min_o1_idx]),
        'mean_principal': float(np.mean(o1)),
        'element_type': element_type,
        'subcase': actual_subcase,
        'units': 'MPa (model units)',
        'num_elements': len(element_ids),
    }

    # Add mid principal stats
    result['max_mid_principal'] = float(np.max(o2))
    result['min_mid_principal'] = float(np.min(o2))

    # Add minimum principal (sigma3) for solid elements
    if o3 is not None:
        min_o3_idx = np.argmin(o3)  # Most negative = max compression
        max_o3_idx = np.argmax(o3)
        result['max_min_principal'] = float(np.max(o3))
        result['min_min_principal'] = float(np.min(o3))  # Most compressive
        result['min_principal_element'] = int(element_ids[min_o3_idx])
        result['has_three_principals'] = True
    else:
        result['has_three_principals'] = False

    # Add von Mises for comparison
    if 'von_mises' in cols:
        vm = stress.data[itime, :, cols['von_mises']]
        result['von_mises_max'] = float(np.max(vm))
        result['von_mises_mean'] = float(np.mean(vm))

    return result


def extract_max_principal_stress(
    op2_file: Path,
    subcase: int = 1,
    element_type: str = 'ctetra'
) -> float:
    """
    Convenience function to extract maximum principal stress.

    Args:
        op2_file: Path to .op2 file
        subcase: Subcase ID
        element_type: Element type

    Returns:
        Maximum principal stress value (tension positive)
    """
    result = extract_principal_stress(op2_file, subcase, element_type)
    return result['max_principal']


def extract_min_principal_stress(
    op2_file: Path,
    subcase: int = 1,
    element_type: str = 'ctetra'
) -> float:
    """
    Convenience function to extract minimum principal stress.

    For solid elements, returns sigma3 (most compressive).
    For shell elements, returns sigma2.

    Args:
        op2_file: Path to .op2 file
        subcase: Subcase ID
        element_type: Element type

    Returns:
        Minimum principal stress value (compression negative)
    """
    result = extract_principal_stress(op2_file, subcase, element_type)
    if result['has_three_principals']:
        return result['min_min_principal']
    else:
        return result['min_mid_principal']


if __name__ == '__main__':
    import sys
    if len(sys.argv) > 1:
        op2_file = Path(sys.argv[1])
        element_type = sys.argv[2] if len(sys.argv) > 2 else 'ctetra'

        result = extract_principal_stress(op2_file, element_type=element_type)

        print(f"\nPrincipal Stress Results ({element_type}):")
        print(f"  Max Principal (σ1): {result['max_principal']:.2f} MPa")
        print(f"    Element: {result['max_principal_element']}")
        if result.get('has_three_principals'):
            print(f"  Min Principal (σ3): {result['min_min_principal']:.2f} MPa")
            print(f"    Element: {result['min_principal_element']}")
        if 'von_mises_max' in result:
            print(f"  Von Mises Max: {result['von_mises_max']:.2f} MPa")
        print(f"  Elements analyzed: {result['num_elements']}")
    else:
        print(f"Usage: python {sys.argv[0]} <op2_file> [element_type]")