Atomizer/optimization_engine/extractors/extract_strain_energy.py

"""
Extract Strain Energy from Structural Analysis
===============================================

Extracts element strain energy and strain energy density from OP2 files.
Strain energy is useful for topology optimization and structural efficiency metrics.

Pattern: strain_energy
Element Types: All structural elements
Result Type: strain energy
API: model.op2_results.strain_energy.*[subcase]

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

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


def extract_strain_energy(
    op2_file: Path,
    subcase: int = 1,
    element_type: Optional[str] = None,
    top_n: int = 10
) -> Dict[str, Any]:
    """
    Extract strain energy from structural elements.

    Strain energy (U) is a measure of the work done to deform the structure:
    U = 0.5 * integral(sigma * epsilon) dV

    High strain energy density indicates highly stressed regions.

    Args:
        op2_file: Path to .op2 file
        subcase: Subcase ID (default 1)
        element_type: Filter by element type (e.g., 'ctetra', 'chexa', 'cquad4')
                     If None, returns total from all elements
        top_n: Number of top elements to return by strain energy

    Returns:
        dict: {
            'total_strain_energy': Total strain energy (all elements),
            'mean_strain_energy': Mean strain energy per element,
            'max_strain_energy': Maximum element strain energy,
            'max_energy_element': Element ID with max strain energy,
            'top_elements': List of (element_id, energy) tuples,
            'energy_by_type': Dict of {element_type: total_energy},
            'num_elements': Total element count,
            'subcase': Subcase ID,
            'units': 'N-mm (model units)'
        }

    Example:
        >>> result = extract_strain_energy('model.op2')
        >>> print(f"Total strain energy: {result['total_strain_energy']:.2f} N-mm")
        >>> print(f"Highest energy element: {result['max_energy_element']}")
    """
    model = OP2()
    model.read_op2(str(op2_file))

    # Check for strain energy results
    # pyNastran stores strain energy in op2_results.strain_energy
    strain_energy_dict = None

    if hasattr(model, 'op2_results') and hasattr(model.op2_results, 'strain_energy'):
        se_container = model.op2_results.strain_energy

        # Strain energy is typically stored by element type
        # ctetra_strain_energy, chexa_strain_energy, etc.
        se_attrs = [a for a in dir(se_container) if 'strain_energy' in a.lower()]

        if not se_attrs:
            # Try direct access patterns
            if hasattr(se_container, 'strain_energy'):
                strain_energy_dict = se_container.strain_energy
    elif hasattr(model, 'strain_energy') and model.strain_energy:
        strain_energy_dict = model.strain_energy

    # Fallback: try legacy access pattern
    if strain_energy_dict is None and hasattr(model, 'ctetra_strain_energy'):
        strain_energy_dict = model.ctetra_strain_energy

    # Collect all strain energy data
    all_elements = []
    all_energies = []
    energy_by_type = {}

    # Search through all possible strain energy attributes
    se_attr_names = [
        'ctetra_strain_energy',
        'chexa_strain_energy',
        'cpenta_strain_energy',
        'cquad4_strain_energy',
        'ctria3_strain_energy',
        'cbar_strain_energy',
        'cbeam_strain_energy',
        'crod_strain_energy',
    ]

    found_any = False

    for attr_name in se_attr_names:
        se_dict = None

        # Try op2_results container first
        if hasattr(model, 'op2_results') and hasattr(model.op2_results, 'strain_energy'):
            se_container = model.op2_results.strain_energy
            if hasattr(se_container, attr_name):
                se_dict = getattr(se_container, attr_name)

        # Fallback to direct model attribute
        if se_dict is None and hasattr(model, attr_name):
            se_dict = getattr(model, attr_name)

        if se_dict is None or not se_dict:
            continue

        # Extract element type from attribute name
        etype = attr_name.replace('_strain_energy', '')

        # Filter by element type if specified
        if element_type is not None and etype.lower() != element_type.lower():
            continue

        # Get subcase
        available_subcases = list(se_dict.keys())
        if not available_subcases:
            continue

        if subcase in available_subcases:
            actual_subcase = subcase
        else:
            actual_subcase = available_subcases[0]

        se_result = se_dict[actual_subcase]
        found_any = True

        # Extract data
        # Strain energy typically stored as: data[itime, ielement, icolumn]
        # Column 0 is usually total strain energy
        itime = 0

        if hasattr(se_result, 'data'):
            energies = se_result.data[itime, :, 0]
            element_ids = se_result.element if hasattr(se_result, 'element') else np.arange(len(energies))

            all_elements.extend(element_ids.tolist())
            all_energies.extend(energies.tolist())

            energy_by_type[etype] = float(np.sum(energies))

    if not found_any or len(all_energies) == 0:
        # No strain energy found - this might not be requested in the analysis
        raise ValueError(
            "No strain energy results found in OP2 file. "
            "Ensure STRAIN=ALL or SEFINAL is in the Nastran case control."
        )

    # Convert to numpy for analysis
    all_energies = np.array(all_energies)
    all_elements = np.array(all_elements)

    # Statistics
    total_energy = float(np.sum(all_energies))
    mean_energy = float(np.mean(all_energies))
    max_idx = np.argmax(all_energies)
    max_energy = float(all_energies[max_idx])
    max_element = int(all_elements[max_idx])

    # Top N elements by strain energy
    top_indices = np.argsort(all_energies)[-top_n:][::-1]
    top_elements = [
        (int(all_elements[i]), float(all_energies[i]))
        for i in top_indices
    ]

    return {
        'total_strain_energy': total_energy,
        'mean_strain_energy': mean_energy,
        'max_strain_energy': max_energy,
        'max_energy_element': max_element,
        'min_strain_energy': float(np.min(all_energies)),
        'std_strain_energy': float(np.std(all_energies)),
        'top_elements': top_elements,
        'energy_by_type': energy_by_type,
        'num_elements': len(all_elements),
        'subcase': subcase,
        'units': 'N-mm (model units)',
    }


def extract_total_strain_energy(
    op2_file: Path,
    subcase: int = 1
) -> float:
    """
    Convenience function to extract total strain energy.

    Args:
        op2_file: Path to .op2 file
        subcase: Subcase ID

    Returns:
        Total strain energy (N-mm)
    """
    result = extract_strain_energy(op2_file, subcase)
    return result['total_strain_energy']


def extract_strain_energy_density(
    op2_file: Path,
    subcase: int = 1,
    element_type: str = 'ctetra'
) -> Dict[str, Any]:
    """
    Extract strain energy density (energy per volume).

    Strain energy density is useful for identifying critical regions
    and for material utilization optimization.

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

    Returns:
        dict: {
            'max_density': Maximum strain energy density,
            'mean_density': Mean strain energy density,
            'total_energy': Total strain energy,
            'units': 'N/mm^2 (MPa equivalent)'
        }

    Note:
        This requires element volume data which may not always be available.
        Falls back to energy-only metrics if volume is unavailable.
    """
    # For now, just return strain energy
    # Full implementation would require element volume from BDF or OP2
    result = extract_strain_energy(op2_file, subcase, element_type)

    # Without volume data, we can't compute true density
    # Return energy metrics with a note
    return {
        'max_strain_energy': result['max_strain_energy'],
        'mean_strain_energy': result['mean_strain_energy'],
        'total_strain_energy': result['total_strain_energy'],
        'max_element': result['max_energy_element'],
        'note': 'True density requires element volumes (not computed)',
        'units': 'N-mm (energy), density requires volume'
    }


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

        try:
            result = extract_strain_energy(op2_file)

            print(f"\nStrain Energy Results:")
            print(f"  Total: {result['total_strain_energy']:.4f} N-mm")
            print(f"  Mean:  {result['mean_strain_energy']:.4f} N-mm")
            print(f"  Max:   {result['max_strain_energy']:.4f} N-mm")
            print(f"    Element: {result['max_energy_element']}")
            print(f"\n  Energy by element type:")
            for etype, energy in result['energy_by_type'].items():
                print(f"    {etype}: {energy:.4f} N-mm")
            print(f"\n  Top 5 elements:")
            for eid, energy in result['top_elements'][:5]:
                print(f"    {eid}: {energy:.4f} N-mm")
            print(f"\n  Total elements: {result['num_elements']}")
        except ValueError as e:
            print(f"Error: {e}")
    else:
        print(f"Usage: python {sys.argv[0]} <op2_file>")