Atomizer/optimization_engine/extractors/stiffness_calculator.py

"""
Generic Stiffness Calculator
============================

Reusable calculator for structural stiffness from FEA results.
Works with any structure: bracket, beam, plate, etc.

Stiffness (k) = Applied Force (F) / Displacement (δ)

Usage:
    calculator = StiffnessCalculator(
        field_file="export_field_dz.fld",
        op2_file="model.op2",
        force_component="fz",
        displacement_component="z"
    )
    results = calculator.calculate()
    stiffness = results['stiffness']
"""

from pathlib import Path
from typing import Dict, Any, Optional
from .field_data_extractor import FieldDataExtractor
from .op2_extractor import OP2Extractor


class StiffnessCalculator:
    """
    Generic stiffness calculator for structural analysis.

    Can be used for:
    - Bracket stiffness (force/deflection)
    - Beam bending stiffness
    - Plate flexural stiffness
    - Any structure where k = F/δ
    """

    def __init__(
        self,
        field_file: str,
        op2_file: str,
        force_component: str = "fz",  # fx, fy, fz, total
        displacement_component: str = "z",  # x, y, z
        displacement_aggregation: str = "max_abs",
        applied_force: Optional[float] = None,  # If known, skip OP2 extraction
    ):
        """
        Args:
            field_file: Path to displacement field file (.fld or .csv)
            op2_file: Path to OP2 file
            force_component: Force component to extract (fx, fy, fz, total)
            displacement_component: Displacement component (x, y, z)
            displacement_aggregation: How to aggregate displacement (max_abs, max, min)
            applied_force: Optional pre-known force value (skips OP2 extraction)
        """
        self.field_file = Path(field_file)
        self.op2_file = Path(op2_file)
        self.force_component = force_component
        self.displacement_component = displacement_component
        self.displacement_aggregation = displacement_aggregation
        self.applied_force = applied_force

    def calculate(self) -> Dict[str, Any]:
        """
        Calculate stiffness from FEA results.

        Returns:
            dict: {
                'stiffness': k value (N/mm or appropriate units),
                'displacement': max displacement (mm),
                'force': applied force (N),
                'compliance': 1/k (mm/N),
                'units': unit description
            }
        """
        # Extract displacement
        disp_extractor = FieldDataExtractor(
            field_file=str(self.field_file),
            result_column=f"{self.displacement_component}(mm)",
            aggregation=self.displacement_aggregation
        )
        disp_result = disp_extractor.extract()
        displacement = disp_result['value']  # mm

        if displacement == 0:
            raise ValueError(f"Zero displacement found. Check field file: {self.field_file}")

        # Extract force
        if self.applied_force is not None:
            force = self.applied_force
        else:
            op2_extractor = OP2Extractor(op2_file=str(self.op2_file))
            force_result = op2_extractor.extract_grid_point_forces(component=self.force_component)
            force = force_result['force']  # N (or mN, check units)

            # Note: Forces in OP2 might be in milli-Newtons (mN)
            # Check f06 header for units and convert if needed
            # For now, assuming Newtons

        if force == 0:
            raise ValueError(f"Zero force found. Check OP2 file: {self.op2_file}")

        # Calculate stiffness: k = F / δ
        stiffness = force / displacement  # N/mm

        # Calculate compliance (inverse of stiffness)
        compliance = displacement / force  # mm/N

        return {
            'stiffness': stiffness,
            'displacement': displacement,
            'force': force,
            'compliance': compliance,
            'units': {
                'stiffness': 'N/mm',
                'displacement': 'mm',
                'force': 'N (verify from f06)',
                'compliance': 'mm/N'
            },
            'field_file': str(self.field_file),
            'op2_file': str(self.op2_file),
            'displacement_stats': disp_result['stats'],
        }


def calculate_stiffness(
    field_file: str,
    op2_file: str,
    force_component: str = "fz",
    displacement_component: str = "z"
) -> float:
    """
    Convenience function to calculate stiffness.

    Args:
        field_file: Path to displacement field file
        op2_file: Path to OP2 file
        force_component: Force component (fx, fy, fz, total)
        displacement_component: Displacement component (x, y, z)

    Returns:
        Stiffness value (N/mm)
    """
    calculator = StiffnessCalculator(
        field_file=field_file,
        op2_file=op2_file,
        force_component=force_component,
        displacement_component=displacement_component
    )
    result = calculator.calculate()
    return result['stiffness']


if __name__ == "__main__":
    # Example usage
    import sys
    if len(sys.argv) > 2:
        field_file = sys.argv[1]
        op2_file = sys.argv[2]

        calculator = StiffnessCalculator(
            field_file=field_file,
            op2_file=op2_file,
            force_component="fz",
            displacement_component="z"
        )

        results = calculator.calculate()
        print(f"Stiffness: {results['stiffness']:.2f} N/mm")
        print(f"Displacement: {results['displacement']:.6f} mm")
        print(f"Force: {results['force']:.2f} N")
        print(f"Compliance: {results['compliance']:.6e} mm/N")