Atomizer/optimization_engine/solve_simulation.py

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

This script opens a .sim file, updates the FEM, and solves it through the NX API.
Usage: run_journal.exe solve_simulation.py <sim_file_path>

Based on recorded NX journal pattern for solving simulations.
"""

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


def main(args):
    """
    Open and solve a simulation file with updated expression values.

    Args:
        args: Command line arguments
              args[0]: .sim file path
              args[1]: tip_thickness value (optional)
              args[2]: support_angle value (optional)
    """
    if len(args) < 1:
        print("ERROR: No .sim file path provided")
        print("Usage: run_journal.exe solve_simulation.py <sim_file_path> [tip_thickness] [support_angle]")
        return False

    sim_file_path = args[0]

    # Parse expression values if provided
    tip_thickness = float(args[1]) if len(args) > 1 else None
    support_angle = float(args[2]) if len(args) > 2 else None

    print(f"[JOURNAL] Opening simulation: {sim_file_path}")
    if tip_thickness is not None:
        print(f"[JOURNAL] Will update tip_thickness = {tip_thickness}")
    if support_angle is not None:
        print(f"[JOURNAL] Will update support_angle = {support_angle}")

    try:
        theSession = NXOpen.Session.GetSession()

        # Close any currently open sim file to force reload from disk
        print("[JOURNAL] Checking for open parts...")
        try:
            current_work = theSession.Parts.BaseWork
            if current_work and hasattr(current_work, 'FullPath'):
                current_path = current_work.FullPath
                print(f"[JOURNAL] Closing currently open part: {current_path}")
                # Close without saving (we want to reload from disk)
                partCloseResponses1 = [NXOpen.BasePart.CloseWholeTree]
                theSession.Parts.CloseAll(partCloseResponses1)
                print("[JOURNAL] Parts closed")
        except Exception as e:
            print(f"[JOURNAL] No parts to close or error closing: {e}")

        # Open the .sim file (now will load fresh from disk with updated .prt files)
        print(f"[JOURNAL] Opening simulation fresh from disk...")
        basePart1, partLoadStatus1 = theSession.Parts.OpenActiveDisplay(
            sim_file_path,
            NXOpen.DisplayPartOption.AllowAdditional
        )

        workSimPart = theSession.Parts.BaseWork
        displaySimPart = theSession.Parts.BaseDisplay
        partLoadStatus1.Dispose()

        # Switch to simulation application
        theSession.ApplicationSwitchImmediate("UG_APP_SFEM")

        simPart1 = workSimPart
        theSession.Post.UpdateUserGroupsFromSimPart(simPart1)

        # STEP 1: Switch to Bracket.prt and update expressions, then update geometry
        print("[JOURNAL] STEP 1: Updating Bracket.prt geometry...")
        try:
            # Find the Bracket part
            bracketPart = theSession.Parts.FindObject("Bracket")
            if bracketPart:
                # Make Bracket the active display part
                status, partLoadStatus = theSession.Parts.SetActiveDisplay(
                    bracketPart,
                    NXOpen.DisplayPartOption.AllowAdditional,
                    NXOpen.PartDisplayPartWorkPartOption.UseLast
                )
                partLoadStatus.Dispose()

                workPart = theSession.Parts.Work

                # CRITICAL: Apply expression changes BEFORE updating geometry
                expressions_updated = []

                if tip_thickness is not None:
                    print(f"[JOURNAL]   Applying tip_thickness = {tip_thickness}")
                    expr_tip = workPart.Expressions.FindObject("tip_thickness")
                    if expr_tip:
                        unit_mm = workPart.UnitCollection.FindObject("MilliMeter")
                        workPart.Expressions.EditExpressionWithUnits(expr_tip, unit_mm, str(tip_thickness))
                        expressions_updated.append(expr_tip)
                        print(f"[JOURNAL]     tip_thickness updated")
                    else:
                        print(f"[JOURNAL]     WARNING: tip_thickness expression not found!")

                if support_angle is not None:
                    print(f"[JOURNAL]   Applying support_angle = {support_angle}")
                    expr_angle = workPart.Expressions.FindObject("support_angle")
                    if expr_angle:
                        unit_deg = workPart.UnitCollection.FindObject("Degrees")
                        workPart.Expressions.EditExpressionWithUnits(expr_angle, unit_deg, str(support_angle))
                        expressions_updated.append(expr_angle)
                        print(f"[JOURNAL]     support_angle updated")
                    else:
                        print(f"[JOURNAL]     WARNING: support_angle expression not found!")

                # Make expressions up to date
                if expressions_updated:
                    print(f"[JOURNAL]   Making {len(expressions_updated)} expression(s) up to date...")
                    for expr in expressions_updated:
                        markId_expr = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Invisible, "Make Up to Date")
                        objects1 = [expr]
                        theSession.UpdateManager.MakeUpToDate(objects1, markId_expr)
                        theSession.DeleteUndoMark(markId_expr, None)

                # CRITICAL: Update the geometry model - rebuilds features with new expressions
                print(f"[JOURNAL]   Rebuilding geometry with new expression values...")
                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]   Bracket geometry updated ({nErrs} errors)")
            else:
                print("[JOURNAL]   WARNING: Could not find Bracket part")
        except Exception as e:
            print(f"[JOURNAL]   ERROR updating Bracket.prt: {e}")
            import traceback
            traceback.print_exc()

        # STEP 2: Switch to Bracket_fem1 and update FE model
        print("[JOURNAL] STEP 2: Opening Bracket_fem1.fem...")
        try:
            # Find the FEM part
            femPart1 = theSession.Parts.FindObject("Bracket_fem1")
            if femPart1:
                # Make FEM the active display part
                status, partLoadStatus = theSession.Parts.SetActiveDisplay(
                    femPart1,
                    NXOpen.DisplayPartOption.AllowAdditional,
                    NXOpen.PartDisplayPartWorkPartOption.SameAsDisplay
                )
                partLoadStatus.Dispose()

                workFemPart = theSession.Parts.BaseWork

                # CRITICAL: Update FE Model - regenerates FEM with new geometry from Bracket.prt
                print("[JOURNAL]   Updating FE Model...")
                fEModel1 = workFemPart.FindObject("FEModel")
                if fEModel1:
                    fEModel1.UpdateFemodel()
                    print("[JOURNAL]   FE Model updated with new geometry!")
                else:
                    print("[JOURNAL]   WARNING: Could not find FEModel object")
            else:
                print("[JOURNAL]   WARNING: Could not find Bracket_fem1 part")
        except Exception as e:
            print(f"[JOURNAL]   ERROR updating FEM: {e}")
            import traceback
            traceback.print_exc()

        # STEP 3: Switch back to sim part
        print("[JOURNAL] STEP 3: Switching back to sim part...")
        try:
            status, partLoadStatus = theSession.Parts.SetActiveDisplay(
                simPart1,
                NXOpen.DisplayPartOption.AllowAdditional,
                NXOpen.PartDisplayPartWorkPartOption.UseLast
            )
            partLoadStatus.Dispose()
            workSimPart = theSession.Parts.BaseWork
            print("[JOURNAL]   Switched back to sim part")
        except Exception as e:
            print(f"[JOURNAL]   WARNING: Error switching to sim part: {e}")

        # Note: Old output files are deleted by nx_solver.py before calling this journal
        # This ensures NX performs a fresh solve

        # Solve the simulation
        print("[JOURNAL] Starting solve...")
        markId3 = theSession.SetUndoMark(NXOpen.Session.MarkVisibility.Visible, "Start")
        theSession.SetUndoMarkName(markId3, "Solve Dialog")

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

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

        # Get the first solution from the simulation
        simSimulation1 = workSimPart.FindObject("Simulation")
        simSolution1 = simSimulation1.FindObject("Solution[Solution 1]")

        psolutions1 = [simSolution1]

        # Solve in background mode
        numsolutionssolved1, numsolutionsfailed1, numsolutionsskipped1 = theCAESimSolveManager.SolveChainOfSolutions(
            psolutions1,
            NXOpen.CAE.SimSolution.SolveOption.Solve,
            NXOpen.CAE.SimSolution.SetupCheckOption.CompleteDeepCheckAndOutputErrors,
            NXOpen.CAE.SimSolution.SolveMode.Background
        )

        theSession.DeleteUndoMark(markId5, None)
        theSession.SetUndoMarkName(markId3, "Solve")

        print(f"[JOURNAL] Solve submitted!")
        print(f"[JOURNAL]   Solutions solved: {numsolutionssolved1}")
        print(f"[JOURNAL]   Solutions failed: {numsolutionsfailed1}")
        print(f"[JOURNAL]   Solutions skipped: {numsolutionsskipped1}")

        # NOTE: In Background mode, these values may not be accurate since the solve
        # runs asynchronously. The solve will continue after this journal finishes.
        # We rely on the Save operation and file existence checks to verify success.

        # Save the simulation to write all output files
        print("[JOURNAL] Saving simulation to ensure output files are written...")
        simPart2 = workSimPart
        partSaveStatus1 = simPart2.Save(
            NXOpen.BasePart.SaveComponents.TrueValue,
            NXOpen.BasePart.CloseAfterSave.FalseValue
        )
        partSaveStatus1.Dispose()
        print("[JOURNAL] Save complete!")

        return True

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


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