Atomizer/optimization_engine/op2_extractor.py

"""
Robust OP2 Extraction - Handles pyNastran FATAL flag issues gracefully.

This module provides a more robust OP2 extraction that:
1. Catches pyNastran FATAL flag exceptions
2. Checks if eigenvalues were actually extracted despite the flag
3. Falls back to F06 extraction if OP2 fails
4. Logs detailed failure information

Usage:
    from optimization_engine.op2_extractor import robust_extract_first_frequency

    frequency = robust_extract_first_frequency(
        op2_file=Path("results.op2"),
        mode_number=1,
        f06_file=Path("results.f06"),  # Optional fallback
        verbose=True
    )
"""

from pathlib import Path
from typing import Optional, Tuple
import numpy as np


def robust_extract_first_frequency(
    op2_file: Path,
    mode_number: int = 1,
    f06_file: Optional[Path] = None,
    verbose: bool = False
) -> float:
    """
    Robustly extract natural frequency from OP2 file, handling pyNastran issues.

    This function attempts multiple strategies:
    1. Standard pyNastran OP2 reading
    2. Force reading with debug=False to ignore FATAL flags
    3. Partial OP2 reading (extract eigenvalues even if FATAL flag exists)
    4. Fallback to F06 file parsing (if provided)

    Args:
        op2_file: Path to OP2 output file
        mode_number: Mode number to extract (1-based index)
        f06_file: Optional F06 file for fallback extraction
        verbose: Print detailed extraction information

    Returns:
        Natural frequency in Hz

    Raises:
        ValueError: If frequency cannot be extracted by any method
    """
    from pyNastran.op2.op2 import OP2

    if not op2_file.exists():
        raise FileNotFoundError(f"OP2 file not found: {op2_file}")

    # Strategy 1: Try standard OP2 reading
    try:
        if verbose:
            print(f"[OP2 EXTRACT] Attempting standard read: {op2_file.name}")

        model = OP2()
        model.read_op2(str(op2_file))

        if hasattr(model, 'eigenvalues') and len(model.eigenvalues) > 0:
            frequency = _extract_frequency_from_model(model, mode_number)
            if verbose:
                print(f"[OP2 EXTRACT] ✓ Success (standard read): {frequency:.6f} Hz")
            return frequency
        else:
            raise ValueError("No eigenvalues found in OP2 file")

    except Exception as e:
        if verbose:
            print(f"[OP2 EXTRACT] ✗ Standard read failed: {str(e)[:100]}")

        # Check if this is a FATAL flag issue
        is_fatal_flag = 'FATAL' in str(e) and 'op2_reader' in str(e.__class__.__module__)

        if is_fatal_flag:
            # Strategy 2: Try reading with more lenient settings
            if verbose:
                print(f"[OP2 EXTRACT] Detected pyNastran FATAL flag issue")
                print(f"[OP2 EXTRACT] Attempting partial extraction...")

            try:
                model = OP2()
                # Try to read with debug=False and skip_undefined_matrices=True
                model.read_op2(
                    str(op2_file),
                    debug=False,
                    skip_undefined_matrices=True
                )

                # Check if eigenvalues were extracted despite FATAL
                if hasattr(model, 'eigenvalues') and len(model.eigenvalues) > 0:
                    frequency = _extract_frequency_from_model(model, mode_number)
                    if verbose:
                        print(f"[OP2 EXTRACT] ✓ Success (lenient mode): {frequency:.6f} Hz")
                        print(f"[OP2 EXTRACT] Note: pyNastran reported FATAL but data is valid!")
                    return frequency

            except Exception as e2:
                if verbose:
                    print(f"[OP2 EXTRACT] ✗ Lenient read also failed: {str(e2)[:100]}")

        # Strategy 3: Fallback to F06 parsing
        if f06_file and f06_file.exists():
            if verbose:
                print(f"[OP2 EXTRACT] Falling back to F06 extraction: {f06_file.name}")

            try:
                frequency = extract_frequency_from_f06(f06_file, mode_number, verbose=verbose)
                if verbose:
                    print(f"[OP2 EXTRACT] ✓ Success (F06 fallback): {frequency:.6f} Hz")
                return frequency

            except Exception as e3:
                if verbose:
                    print(f"[OP2 EXTRACT] ✗ F06 extraction failed: {str(e3)}")

        # All strategies failed
        raise ValueError(
            f"Could not extract frequency from OP2 file: {op2_file.name}. "
            f"Original error: {str(e)}"
        )


def _extract_frequency_from_model(model, mode_number: int) -> float:
    """Extract frequency from loaded OP2 model."""
    if not hasattr(model, 'eigenvalues') or len(model.eigenvalues) == 0:
        raise ValueError("No eigenvalues found in model")

    # Get first subcase
    subcase = list(model.eigenvalues.keys())[0]
    eig_obj = model.eigenvalues[subcase]

    # Check if mode exists
    if mode_number > len(eig_obj.eigenvalues):
        raise ValueError(
            f"Mode {mode_number} not found. "
            f"Only {len(eig_obj.eigenvalues)} modes available"
        )

    # Extract eigenvalue and convert to frequency
    eigenvalue = eig_obj.eigenvalues[mode_number - 1]
    angular_freq = np.sqrt(abs(eigenvalue))  # Use abs to handle numerical precision issues
    frequency_hz = angular_freq / (2 * np.pi)

    return float(frequency_hz)


def extract_frequency_from_f06(
    f06_file: Path,
    mode_number: int = 1,
    verbose: bool = False
) -> float:
    """
    Extract natural frequency from F06 text file (fallback method).

    Parses the F06 file to find eigenvalue results table and extracts frequency.

    Args:
        f06_file: Path to F06 output file
        mode_number: Mode number to extract (1-based index)
        verbose: Print extraction details

    Returns:
        Natural frequency in Hz

    Raises:
        ValueError: If frequency cannot be found in F06
    """
    if not f06_file.exists():
        raise FileNotFoundError(f"F06 file not found: {f06_file}")

    with open(f06_file, 'r', encoding='latin-1', errors='ignore') as f:
        content = f.read()

    # Look for eigenvalue table
    # Nastran F06 format has eigenvalue results like:
    #                                             R E A L   E I G E N V A L U E S
    #    MODE    EXTRACTION      EIGENVALUE            RADIANS             CYCLES            GENERALIZED         GENERALIZED
    #     NO.       ORDER                                                                       MASS              STIFFNESS
    #         1         1       -6.602743E+04        2.569656E+02        4.089338E+01        1.000000E+00        6.602743E+04

    lines = content.split('\n')

    # Find eigenvalue table
    eigenvalue_section_start = None
    for i, line in enumerate(lines):
        if 'R E A L   E I G E N V A L U E S' in line:
            eigenvalue_section_start = i
            break

    if eigenvalue_section_start is None:
        raise ValueError("Eigenvalue table not found in F06 file")

    # Parse eigenvalue table (starts a few lines after header)
    for i in range(eigenvalue_section_start + 3, min(eigenvalue_section_start + 100, len(lines))):
        line = lines[i].strip()

        if not line or line.startswith('1'):  # Page break
            continue

        # Parse line with mode data
        parts = line.split()
        if len(parts) >= 5:
            try:
                mode_num = int(parts[0])
                if mode_num == mode_number:
                    # Frequency is in column 5 (CYCLES)
                    frequency = float(parts[4])
                    if verbose:
                        print(f"[F06 EXTRACT] Found mode {mode_num}: {frequency:.6f} Hz")
                    return frequency
            except (ValueError, IndexError):
                continue

    raise ValueError(f"Mode {mode_number} not found in F06 eigenvalue table")


def validate_op2_file(op2_file: Path, f06_file: Optional[Path] = None) -> Tuple[bool, str]:
    """
    Validate if an OP2 file contains usable eigenvalue data.

    Args:
        op2_file: Path to OP2 file
        f06_file: Optional F06 file for cross-reference

    Returns:
        (is_valid, message): Tuple of validation status and explanation
    """
    if not op2_file.exists():
        return False, f"OP2 file does not exist: {op2_file}"

    if op2_file.stat().st_size == 0:
        return False, "OP2 file is empty"

    # Try to extract first frequency
    try:
        frequency = robust_extract_first_frequency(
            op2_file,
            mode_number=1,
            f06_file=f06_file,
            verbose=False
        )
        return True, f"Valid OP2 file (first frequency: {frequency:.6f} Hz)"

    except Exception as e:
        return False, f"Cannot extract data from OP2: {str(e)}"


# Convenience function (same signature as old function for backward compatibility)
def extract_first_frequency(op2_file: Path, mode_number: int = 1) -> float:
    """
    Extract first natural frequency (backward compatible with old function).

    This is the simple version - just use robust_extract_first_frequency directly
    for more control.

    Args:
        op2_file: Path to OP2 file
        mode_number: Mode number (1-based)

    Returns:
        Frequency in Hz
    """
    # Try to find F06 file in same directory
    f06_file = op2_file.with_suffix('.f06')

    return robust_extract_first_frequency(
        op2_file,
        mode_number=mode_number,
        f06_file=f06_file if f06_file.exists() else None,
        verbose=False
    )