Atomizer/tools/adaptive-isogrid/src/nx/solve_and_extract.py

"""
NXOpen script — remesh full plate, solve, and export results.json.

Outputs:
{
  "nodes_xy": [[x, y], ...],
  "stress_values": [...],
  "disp_values": [...],
  "strain_values": [...],
  "mass": 0.0
}
"""

from __future__ import annotations

import argparse
import csv
import json
from pathlib import Path
from typing import Any, Dict, List, Sequence, Tuple


def remesh_full_plate(sim_part: Any) -> None:
    """
    Trigger mesh regeneration for all FEM collectors associated with the plate.
    """
    # This hook should iterate mesh managers/collectors in your SIM/FEM template.
    # API often used: FEModel.UpdateFemodel(), mesh manager GenerateMesh/Update.
    fe_model = sim_part.FindObject("FEModel")
    fe_model.UpdateFemodel()


def solve_active_solution(sim_part: Any, solution_name: str | None = None) -> Any:
    """
    Solve the requested solution (or first available).
    """
    import NXOpen  # type: ignore
    import NXOpen.CAE  # type: ignore

    simulation = sim_part.FindObject("Simulation")

    target_solution = None
    if solution_name:
        try:
            target_solution = simulation.FindObject(f"Solution[{solution_name}]")
        except Exception:
            target_solution = None

    if target_solution is None:
        target_solution = simulation.FindObject("Solution[Solution 1]")

    solve_mgr = NXOpen.CAE.SimSolveManager.GetSimSolveManager(NXOpen.Session.GetSession())
    solve_mgr.SubmitSolves([target_solution])
    return target_solution


def _parse_csv_results(path: Path) -> Tuple[List[List[float]], List[float]]:
    """
    Parse generic CSV with columns: x,y,value (header-insensitive).
    """
    coords: List[List[float]] = []
    values: List[float] = []
    with path.open(newline="") as f:
        reader = csv.DictReader(f)
        for row in reader:
            keys = {k.lower(): k for k in row.keys()}
            x = float(row[keys.get("x", keys.get("x_mm", "x"))])
            y = float(row[keys.get("y", keys.get("y_mm", "y"))])
            v = float(row[keys.get("value", keys.get("von_mises", "value"))])
            coords.append([x, y])
            values.append(v)
    return coords, values


def extract_results_nxopen(sim_part: Any, sandbox_ids: List[str]) -> Dict[str, Any]:
    """
    Preferred extractor: NXOpen post-processing API.

    NOTE: Result object names/components vary by template. Keep this function as the
    primary integration point for project-specific API wiring.
    """
    raise NotImplementedError(
        "Wire NXOpen post-processing calls here (nodal stress/displacement, elemental strain, mass)."
    )


def extract_results_csv_fallback(work_dir: Path) -> Dict[str, Any]:
    """
    Fallback extractor: parse Simcenter-exported CSV files in work_dir.

    Expected files:
      - nodal_stress.csv
      - nodal_disp.csv
      - elemental_strain.csv
      - mass.json (optional: {"mass": ...})
    """
    stress_coords, stress_vals = _parse_csv_results(work_dir / "nodal_stress.csv")
    disp_coords, disp_vals = _parse_csv_results(work_dir / "nodal_disp.csv")
    _, strain_vals = _parse_csv_results(work_dir / "elemental_strain.csv")

    # Use stress nodal coordinates as canonical nodes_xy
    nodes_xy = stress_coords if stress_coords else disp_coords

    mass = 0.0
    mass_file = work_dir / "mass.json"
    if mass_file.exists():
        mass = float(json.loads(mass_file.read_text()).get("mass", 0.0))

    return {
        "nodes_xy": nodes_xy,
        "stress_values": stress_vals,
        "disp_values": disp_vals,
        "strain_values": strain_vals,
        "mass": mass,
    }


def run_in_nx(
    work_dir: Path,
    result_path: Path,
    sandbox_ids: List[str],
    solution_name: str | None = None,
    use_csv_fallback: bool = False,
) -> Dict[str, Any]:
    import NXOpen  # type: ignore

    session = NXOpen.Session.GetSession()
    sim_part = session.Parts.BaseWork
    if sim_part is None:
        raise RuntimeError("No active NX SIM/FEM work part.")

    remesh_full_plate(sim_part)
    solve_active_solution(sim_part, solution_name=solution_name)

    if use_csv_fallback:
        results = extract_results_csv_fallback(work_dir)
    else:
        results = extract_results_nxopen(sim_part, sandbox_ids)

    result_path.parent.mkdir(parents=True, exist_ok=True)
    result_path.write_text(json.dumps(results, indent=2))
    print(f"[solve_and_extract] wrote {result_path}")
    return results


def main(argv: Sequence[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Remesh, solve, and export results.json")
    parser.add_argument("--work-dir", default=".", help="Working directory for CSV fallback artifacts")
    parser.add_argument("--results", default="results.json", help="Output JSON path")
    parser.add_argument("--sandbox-id", action="append", default=[], help="Sandbox id filter (repeatable)")
    parser.add_argument("--solution", default=None, help="NX solution name (defaults to Solution 1)")
    parser.add_argument("--csv-fallback", action="store_true", help="Parse CSV files instead of NXOpen post API")
    args = parser.parse_args(argv)

    run_in_nx(
        work_dir=Path(args.work_dir),
        result_path=Path(args.results),
        sandbox_ids=args.sandbox_id,
        solution_name=args.solution,
        use_csv_fallback=args.csv_fallback,
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())