Atomizer/optimization_engine/solve_simulation.py

"""
NX Journal Script to Solve Simulation in Batch Mode

This script handles BOTH single-part simulations AND multi-part assembly FEMs.

=============================================================================
MULTI-PART ASSEMBLY FEM WORKFLOW (for .afm-based simulations)
=============================================================================

Based on recorded NX journal from interactive session (Nov 28, 2025).

The correct workflow for assembly FEM updates:

1. LOAD PARTS
   - Open ASSY_M1.prt and M1_Blank_fem1_i.prt to have geometry loaded
   - Find and switch to M1_Blank part for expression editing

2. UPDATE EXPRESSIONS
   - Switch to modeling application
   - Edit expressions with units
   - Call MakeUpToDate() on modified expressions
   - Call DoUpdate() to rebuild geometry

3. SWITCH TO SIM AND UPDATE FEM COMPONENTS
   - Open the .sim file
   - Navigate component hierarchy via RootComponent.FindObject()
   - For each component FEM:
     - SetWorkComponent() to make it the work part
     - FindObject("FEModel").UpdateFemodel()

4. MERGE DUPLICATE NODES (critical for assembly FEM!)
   - Switch to assembly FEM component
   - CreateDuplicateNodesCheckBuilder()
   - Set MergeOccurrenceNodes = True
   - IdentifyDuplicateNodes() then MergeDuplicateNodes()

5. RESOLVE LABEL CONFLICTS
   - CreateAssemblyLabelManagerBuilder()
   - SetFEModelOccOffsets() for each occurrence
   - Commit()

6. SOLVE
   - SetWorkComponent(Null) to return to sim level
   - SolveChainOfSolutions()

=============================================================================
"""

import sys
import os
import NXOpen
import NXOpen.Assemblies
import NXOpen.CAE


def main(args):
    """
    Main entry point for NX journal.

    Args:
        args: Command line arguments
              args[0]: .sim file path
              args[1]: solution_name (optional, or "None" for default)
              args[2+]: expression updates as "name=value" pairs
    """
    if len(args) < 1:
        print("ERROR: No .sim file path provided")
        print("Usage: run_journal.exe solve_simulation.py <sim_file_path> [solution_name] [expr1=val1] ...")
        return False

    sim_file_path = args[0]
    solution_name = args[1] if len(args) > 1 and args[1] != 'None' else None

    # Parse expression updates
    expression_updates = {}
    for arg in args[2:]:
        if '=' in arg:
            name, value = arg.split('=', 1)
            expression_updates[name] = float(value)

    # Get working directory
    working_dir = os.path.dirname(os.path.abspath(sim_file_path))
    sim_filename = os.path.basename(sim_file_path)

    print(f"[JOURNAL] " + "="*60)
    print(f"[JOURNAL] NX SIMULATION SOLVER (Assembly FEM Workflow)")
    print(f"[JOURNAL] " + "="*60)
    print(f"[JOURNAL] Simulation: {sim_filename}")
    print(f"[JOURNAL] Working directory: {working_dir}")
    print(f"[JOURNAL] Solution: {solution_name or 'Solution 1'}")
    print(f"[JOURNAL] Expression updates: {len(expression_updates)}")
    for name, value in expression_updates.items():
        print(f"[JOURNAL]   {name} = {value}")

    try:
        theSession = NXOpen.Session.GetSession()

        # Set load options
        theSession.Parts.LoadOptions.LoadLatest = False
        theSession.Parts.LoadOptions.ComponentLoadMethod = NXOpen.LoadOptions.LoadMethod.FromDirectory
        theSession.Parts.LoadOptions.SetSearchDirectories([working_dir], [True])
        theSession.Parts.LoadOptions.ComponentsToLoad = NXOpen.LoadOptions.LoadComponents.All
        theSession.Parts.LoadOptions.PartLoadOption = NXOpen.LoadOptions.LoadOption.FullyLoad
        theSession.Parts.LoadOptions.SetInterpartData(True, NXOpen.LoadOptions.Parent.All)
        theSession.Parts.LoadOptions.AbortOnFailure = False

        # Close any open parts
        try:
            theSession.Parts.CloseAll([NXOpen.BasePart.CloseWholeTree])
        except:
            pass

        # Check for assembly FEM files
        afm_files = [f for f in os.listdir(working_dir) if f.endswith('.afm')]
        is_assembly = len(afm_files) > 0

        if is_assembly and expression_updates:
            print(f"[JOURNAL] ")
            print(f"[JOURNAL] DETECTED: Multi-part Assembly FEM")
            print(f"[JOURNAL] Using ASSEMBLY FEM WORKFLOW")
            print(f"[JOURNAL] ")
            return solve_assembly_fem_workflow(
                theSession, sim_file_path, solution_name, expression_updates, working_dir
            )
        else:
            print(f"[JOURNAL] ")
            print(f"[JOURNAL] Using SIMPLE WORKFLOW (no expression updates or single-part)")
            print(f"[JOURNAL] ")
            return solve_simple_workflow(
                theSession, sim_file_path, solution_name, expression_updates, working_dir
            )

    except Exception as e:
        print(f"[JOURNAL] FATAL ERROR: {e}")
        import traceback
        traceback.print_exc()
        return False


def solve_assembly_fem_workflow(theSession, sim_file_path, solution_name, expression_updates, working_dir):
    """
    Full assembly FEM workflow based on recorded NX journal.

    This is the correct workflow for multi-part assembly FEMs.
    """
    sim_filename = os.path.basename(sim_file_path)

    # ==========================================================================
    # STEP 1: LOAD REQUIRED PARTS
    # ==========================================================================
    print(f"[JOURNAL] STEP 1: Loading required parts...")

    # Load ASSY_M1.prt (to have the geometry assembly available)
    assy_prt_path = os.path.join(working_dir, "ASSY_M1.prt")
    if os.path.exists(assy_prt_path):
        print(f"[JOURNAL]   Loading ASSY_M1.prt...")
        markId1 = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Visible, "Load Part")
        part1, partLoadStatus1 = theSession.Parts.Open(assy_prt_path)
        partLoadStatus1.Dispose()
    else:
        print(f"[JOURNAL]   WARNING: ASSY_M1.prt not found, continuing anyway...")

    # Load M1_Blank_fem1_i.prt (idealized geometry)
    idealized_prt_path = os.path.join(working_dir, "M1_Blank_fem1_i.prt")
    if os.path.exists(idealized_prt_path):
        print(f"[JOURNAL]   Loading M1_Blank_fem1_i.prt...")
        markId2 = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Visible, "Load Part")
        part2, partLoadStatus2 = theSession.Parts.Open(idealized_prt_path)
        partLoadStatus2.Dispose()

    # ==========================================================================
    # STEP 2: UPDATE EXPRESSIONS IN M1_BLANK
    # ==========================================================================
    print(f"[JOURNAL] STEP 2: Updating expressions in M1_Blank...")

    # Find and switch to M1_Blank part
    try:
        part3 = theSession.Parts.FindObject("M1_Blank")
        markId3 = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Visible, "Change Displayed Part")
        status1, partLoadStatus3 = theSession.Parts.SetActiveDisplay(
            part3,
            NXOpen.DisplayPartOption.AllowAdditional,
            NXOpen.PartDisplayPartWorkPartOption.UseLast
        )
        partLoadStatus3.Dispose()

        # Switch to modeling application for expression editing
        theSession.ApplicationSwitchImmediate("UG_APP_MODELING")

        workPart = theSession.Parts.Work

        # Create undo mark for expressions
        markId4 = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Visible, "Start")
        theSession.SetUndoMarkName(markId4, "Expressions Dialog")

        # Write expressions to a temp file and import (more reliable than editing one by one)
        exp_file_path = os.path.join(working_dir, "_temp_expressions.exp")
        with open(exp_file_path, 'w') as f:
            for expr_name, expr_value in expression_updates.items():
                # Determine unit
                if 'angle' in expr_name.lower() or 'vertical' in expr_name.lower():
                    unit_str = "Degrees"
                else:
                    unit_str = "MilliMeter"
                f.write(f"[{unit_str}]{expr_name}={expr_value}\n")
                print(f"[JOURNAL]     {expr_name} = {expr_value} ({unit_str})")

        print(f"[JOURNAL]   Importing expressions from file...")
        markId_import = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Visible, "Import Expressions")

        try:
            expModified, errorMessages = workPart.Expressions.ImportFromFile(
                exp_file_path,
                NXOpen.ExpressionCollection.ImportMode.Replace
            )
            print(f"[JOURNAL]   Expressions imported: {expModified} modified")
            if errorMessages:
                print(f"[JOURNAL]   Import errors: {errorMessages}")

            # Update geometry after import
            print(f"[JOURNAL]   Rebuilding geometry...")
            markId_update = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Invisible, "NX update")
            nErrs = theSession.UpdateManager.DoUpdate(markId_update)
            theSession.DeleteUndoMark(markId_update, "NX update")
            print(f"[JOURNAL]   Geometry rebuilt ({nErrs} errors)")

            updated_expressions = list(expression_updates.keys())

        except Exception as e:
            print(f"[JOURNAL]   ERROR importing expressions: {e}")
            updated_expressions = []

        # Clean up temp file
        try:
            os.remove(exp_file_path)
        except:
            pass

        theSession.SetUndoMarkName(markId4, "Expressions")

    except Exception as e:
        print(f"[JOURNAL]   ERROR updating expressions: {e}")

    # ==========================================================================
    # STEP 3: OPEN SIM AND UPDATE COMPONENT FEMs
    # ==========================================================================
    print(f"[JOURNAL] STEP 3: Opening sim and updating component FEMs...")

    # Try to find the sim part first (like the recorded journal does)
    # This ensures we're working with the same loaded sim part context
    sim_part_name = os.path.splitext(sim_filename)[0]  # e.g., "ASSY_M1_assyfem1_sim1"
    print(f"[JOURNAL]   Looking for sim part: {sim_part_name}")

    markId_sim = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Visible, "Change Displayed Part")

    try:
        # First try to find it among loaded parts (like recorded journal)
        simPart1 = theSession.Parts.FindObject(sim_part_name)
        status_sim, partLoadStatus = theSession.Parts.SetActiveDisplay(
            simPart1,
            NXOpen.DisplayPartOption.AllowAdditional,
            NXOpen.PartDisplayPartWorkPartOption.UseLast
        )
        partLoadStatus.Dispose()
        print(f"[JOURNAL]   Found and activated existing sim part")
    except:
        # Fallback: Open fresh if not found
        print(f"[JOURNAL]   Sim part not found, opening fresh: {sim_filename}")
        basePart, partLoadStatus = theSession.Parts.OpenActiveDisplay(
            sim_file_path,
            NXOpen.DisplayPartOption.AllowAdditional
        )
        partLoadStatus.Dispose()

    workSimPart = theSession.Parts.BaseWork
    displaySimPart = theSession.Parts.BaseDisplay
    theSession.ApplicationSwitchImmediate("UG_APP_SFEM")
    theSession.Post.UpdateUserGroupsFromSimPart(workSimPart)

    # Navigate component hierarchy
    try:
        rootComponent = workSimPart.ComponentAssembly.RootComponent
        component1 = rootComponent.FindObject("COMPONENT ASSY_M1_assyfem1 1")

        # Update M1_Blank_fem1
        print(f"[JOURNAL]   Updating M1_Blank_fem1...")
        try:
            component2 = component1.FindObject("COMPONENT M1_Blank_fem1 1")
            markId_fem1 = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Visible, "Make Work Part")
            partLoadStatus5 = theSession.Parts.SetWorkComponent(
                component2,
                NXOpen.PartCollection.RefsetOption.Entire,
                NXOpen.PartCollection.WorkComponentOption.Visible
            )
            workFemPart = theSession.Parts.BaseWork
            partLoadStatus5.Dispose()

            markId_update1 = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Visible, "Update FE Model")
            fEModel1 = workFemPart.FindObject("FEModel")
            fEModel1.UpdateFemodel()
            print(f"[JOURNAL]     M1_Blank_fem1 updated")
        except Exception as e:
            print(f"[JOURNAL]     WARNING: M1_Blank_fem1: {e}")

        # Update M1_Vertical_Support_Skeleton_fem1
        print(f"[JOURNAL]   Updating M1_Vertical_Support_Skeleton_fem1...")
        try:
            component3 = component1.FindObject("COMPONENT M1_Vertical_Support_Skeleton_fem1 3")
            markId_fem2 = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Visible, "Make Work Part")
            partLoadStatus6 = theSession.Parts.SetWorkComponent(
                component3,
                NXOpen.PartCollection.RefsetOption.Entire,
                NXOpen.PartCollection.WorkComponentOption.Visible
            )
            workFemPart = theSession.Parts.BaseWork
            partLoadStatus6.Dispose()

            markId_update2 = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Visible, "Update FE Model")
            fEModel2 = workFemPart.FindObject("FEModel")
            fEModel2.UpdateFemodel()
            print(f"[JOURNAL]     M1_Vertical_Support_Skeleton_fem1 updated")
        except Exception as e:
            print(f"[JOURNAL]     WARNING: M1_Vertical_Support_Skeleton_fem1: {e}")

    except Exception as e:
        print(f"[JOURNAL]   ERROR navigating component hierarchy: {e}")

    # ==========================================================================
    # STEP 4: MERGE DUPLICATE NODES
    # ==========================================================================
    print(f"[JOURNAL] STEP 4: Merging duplicate nodes...")

    try:
        # Switch to assembly FEM
        partLoadStatus8 = theSession.Parts.SetWorkComponent(
            component1,
            NXOpen.PartCollection.RefsetOption.Entire,
            NXOpen.PartCollection.WorkComponentOption.Visible
        )
        workAssyFemPart = theSession.Parts.BaseWork
        displaySimPart = theSession.Parts.BaseDisplay
        partLoadStatus8.Dispose()
        print(f"[JOURNAL]   Switched to assembly FEM: {workAssyFemPart.Name}")

        markId_merge = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Visible, "Start")

        caePart1 = workAssyFemPart
        duplicateNodesCheckBuilder1 = caePart1.ModelCheckMgr.CreateDuplicateNodesCheckBuilder()

        # Set tolerance
        unit_tol = duplicateNodesCheckBuilder1.Tolerance.Units
        duplicateNodesCheckBuilder1.Tolerance.Units = unit_tol
        duplicateNodesCheckBuilder1.Tolerance.SetFormula("0.01")
        print(f"[JOURNAL]   Tolerance: 0.01 mm")

        # Enable occurrence node merge - CRITICAL for assembly FEM
        duplicateNodesCheckBuilder1.MergeOccurrenceNodes = True
        print(f"[JOURNAL]   MergeOccurrenceNodes: True")

        theSession.SetUndoMarkName(markId_merge, "Duplicate Nodes Dialog")

        # Configure display settings
        displaysettings1 = NXOpen.CAE.ModelCheck.DuplicateNodesCheckBuilder.DisplaySettings()
        displaysettings1.ShowDuplicateNodes = True
        displaysettings1.ShowMergedNodeLabels = False
        displaysettings1.ShowRetainedNodeLabels = False
        displaysettings1.KeepNodesColor = displaySimPart.Colors.Find("Blue")
        displaysettings1.MergeNodesColor = displaySimPart.Colors.Find("Yellow")
        displaysettings1.UnableToMergeNodesColor = displaySimPart.Colors.Find("Red")
        duplicateNodesCheckBuilder1.DisplaySettingsData = displaysettings1

        # Check scope
        duplicateNodesCheckBuilder1.CheckScopeOption = NXOpen.CAE.ModelCheck.CheckScope.Displayed
        print(f"[JOURNAL]   CheckScope: Displayed")

        # Identify duplicates
        print(f"[JOURNAL]   Identifying duplicate nodes...")
        numDuplicates = duplicateNodesCheckBuilder1.IdentifyDuplicateNodes()
        print(f"[JOURNAL]   Found {numDuplicates} duplicate node sets")

        # Merge duplicates
        if numDuplicates > 0:
            print(f"[JOURNAL]   Merging duplicate nodes...")
            numMerged = duplicateNodesCheckBuilder1.MergeDuplicateNodes()
            print(f"[JOURNAL]   Merged {numMerged} duplicate node sets")
        else:
            print(f"[JOURNAL]   WARNING: No duplicate nodes found to merge!")
            print(f"[JOURNAL]   This may indicate mesh update didn't work properly")

        theSession.SetUndoMarkName(markId_merge, "Duplicate Nodes")
        duplicateNodesCheckBuilder1.Destroy()
        theSession.DeleteUndoMark(markId_merge, None)

    except Exception as e:
        print(f"[JOURNAL]   WARNING: Node merge: {e}")
        import traceback
        traceback.print_exc()

    # ==========================================================================
    # STEP 5: RESOLVE LABEL CONFLICTS
    # ==========================================================================
    print(f"[JOURNAL] STEP 5: Resolving label conflicts...")

    try:
        markId_labels = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Visible, "Start")

        assyFemPart1 = workAssyFemPart
        assemblyLabelManagerBuilder1 = assyFemPart1.CreateAssemblyLabelManagerBuilder()

        theSession.SetUndoMarkName(markId_labels, "Assembly Label Manager Dialog")

        markId_labels2 = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Invisible, "Assembly Label Manager")

        # Set offsets for each FE model occurrence
        # These offsets ensure unique node/element labels across components
        entitytypes = [
            NXOpen.CAE.AssemblyLabelManagerBuilder.EntityType.Node,
            NXOpen.CAE.AssemblyLabelManagerBuilder.EntityType.Element,
            NXOpen.CAE.AssemblyLabelManagerBuilder.EntityType.Csys,
            NXOpen.CAE.AssemblyLabelManagerBuilder.EntityType.Physical,
            NXOpen.CAE.AssemblyLabelManagerBuilder.EntityType.Group,
            NXOpen.CAE.AssemblyLabelManagerBuilder.EntityType.Ply,
            NXOpen.CAE.AssemblyLabelManagerBuilder.EntityType.Ssmo,
        ]

        # Apply offsets to each occurrence (values from recorded journal)
        occurrence_offsets = [
            ("FEModelOccurrence[3]", 2),
            ("FEModelOccurrence[4]", 74),
            ("FEModelOccurrence[5]", 146),
            ("FEModelOccurrence[7]", 218),
        ]

        for occ_name, offset_val in occurrence_offsets:
            try:
                fEModelOcc = workAssyFemPart.FindObject(occ_name)
                offsets = [offset_val] * 7
                assemblyLabelManagerBuilder1.SetFEModelOccOffsets(fEModelOcc, entitytypes, offsets)
            except:
                pass  # Some occurrences may not exist

        nXObject1 = assemblyLabelManagerBuilder1.Commit()

        theSession.DeleteUndoMark(markId_labels2, None)
        theSession.SetUndoMarkName(markId_labels, "Assembly Label Manager")
        assemblyLabelManagerBuilder1.Destroy()

        print(f"[JOURNAL]   Label conflicts resolved")

    except Exception as e:
        print(f"[JOURNAL]   WARNING: Label management: {e}")

    # ==========================================================================
    # STEP 6: SOLVE
    # ==========================================================================
    print(f"[JOURNAL] STEP 6: Solving simulation...")

    try:
        # Return to sim level by setting null component
        partLoadStatus9 = theSession.Parts.SetWorkComponent(
            NXOpen.Assemblies.Component.Null,
            NXOpen.PartCollection.RefsetOption.Entire,
            NXOpen.PartCollection.WorkComponentOption.Visible
        )
        workSimPart = theSession.Parts.BaseWork
        partLoadStatus9.Dispose()

        # Set up solve
        markId_solve = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Visible, "Start")
        theSession.SetUndoMarkName(markId_solve, "Solve Dialog")

        markId_solve2 = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Invisible, "Solve")

        theCAESimSolveManager = NXOpen.CAE.SimSolveManager.GetSimSolveManager(theSession)

        simSimulation1 = workSimPart.FindObject("Simulation")
        sol_name = solution_name if solution_name else "Solution 1"
        simSolution1 = simSimulation1.FindObject(f"Solution[{sol_name}]")

        psolutions1 = [simSolution1]

        print(f"[JOURNAL]   Solving: {sol_name} (Foreground mode)")
        numsolved, numfailed, numskipped = theCAESimSolveManager.SolveChainOfSolutions(
            psolutions1,
            NXOpen.CAE.SimSolution.SolveOption.Solve,
            NXOpen.CAE.SimSolution.SetupCheckOption.CompleteCheckAndOutputErrors,
            NXOpen.CAE.SimSolution.SolveMode.Foreground  # Use Foreground to ensure OP2 is complete
        )

        theSession.DeleteUndoMark(markId_solve2, None)
        theSession.SetUndoMarkName(markId_solve, "Solve")

        print(f"[JOURNAL] Solve completed: {numsolved} solved, {numfailed} failed, {numskipped} skipped")

        return numfailed == 0

    except Exception as e:
        print(f"[JOURNAL] ERROR solving: {e}")
        import traceback
        traceback.print_exc()
        return False


def solve_simple_workflow(theSession, sim_file_path, solution_name, expression_updates, working_dir):
    """
    Simple workflow for single-part simulations or when no expression updates needed.
    """
    print(f"[JOURNAL] Opening simulation: {sim_file_path}")

    # Open the .sim file
    basePart1, partLoadStatus1 = theSession.Parts.OpenActiveDisplay(
        sim_file_path,
        NXOpen.DisplayPartOption.AllowAdditional
    )
    partLoadStatus1.Dispose()

    workSimPart = theSession.Parts.BaseWork
    theSession.ApplicationSwitchImmediate("UG_APP_SFEM")
    theSession.Post.UpdateUserGroupsFromSimPart(workSimPart)

    # Set up solve
    markId_solve = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Visible, "Start")
    theSession.SetUndoMarkName(markId_solve, "Solve Dialog")

    markId_solve2 = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Invisible, "Solve")

    theCAESimSolveManager = NXOpen.CAE.SimSolveManager.GetSimSolveManager(theSession)

    simSimulation1 = workSimPart.FindObject("Simulation")
    sol_name = solution_name if solution_name else "Solution 1"
    simSolution1 = simSimulation1.FindObject(f"Solution[{sol_name}]")

    psolutions1 = [simSolution1]

    print(f"[JOURNAL]   Solving: {sol_name}")
    numsolved, numfailed, numskipped = theCAESimSolveManager.SolveChainOfSolutions(
        psolutions1,
        NXOpen.CAE.SimSolution.SolveOption.Solve,
        NXOpen.CAE.SimSolution.SetupCheckOption.CompleteCheckAndOutputErrors,
        NXOpen.CAE.SimSolution.SolveMode.Background
    )

    theSession.DeleteUndoMark(markId_solve2, None)
    theSession.SetUndoMarkName(markId_solve, "Solve")

    print(f"[JOURNAL] Solve completed: {numsolved} solved, {numfailed} failed, {numskipped} skipped")

    # Save
    try:
        partSaveStatus = workSimPart.Save(
            NXOpen.BasePart.SaveComponents.TrueValue,
            NXOpen.BasePart.CloseAfterSave.FalseValue
        )
        partSaveStatus.Dispose()
        print(f"[JOURNAL] Saved!")
    except:
        pass

    return numfailed == 0


if __name__ == '__main__':
    success = main(sys.argv[1:])
    sys.exit(0 if success else 1)