"""Resumable dense triangle/FEMM waveforms and live, data-backed comparisons. Every plotted base sample is directly solved. Nested grids and a shifted grid test angular integration and periodic spectral linkage derivatives. No symmetry reconstruction, interpolated solver samples or historical scalar substitution. """ from __future__ import annotations import argparse import gzip import importlib.metadata import json import os import platform import shutil import sys import time import zipfile from contextlib import closing from dataclasses import asdict from pathlib import Path import numpy as np from pipeline.femm_utils.artifacts import atomic_json, file_digest from pipeline.femm_utils.repo_adapter import load_design, prepare from pipeline.simulation.studies import _canonical_sha256, load_reference REFERENCE = Path("FEMM/references/corrected-angle96-20261026/3115") DEFAULT_OUTPUT = Path("FEMM/output/triangle-dense-waveforms-20261028") def read(path): return json.loads(path.read_text()) def sampling_minimum(directory, plan): """Execution policy may reduce sample count, never convergence tolerances.""" policy_path = directory.parent / "sampling-policy.json" policy = read(policy_path) if policy_path.exists() else {} minimum = policy.get(directory.name, {}).get( "minimum_angles", plan["target_angles"] ) if minimum not in plan["angular_levels"]: raise ValueError("Sampling minimum must be a supported nested grid") request = directory / "refinement-request.json" return max(minimum, read(request)["minimum_angles"] if request.exists() else 0) def finite_json(value): if isinstance(value, dict): return {k: finite_json(v) for k, v in value.items()} if isinstance(value, list): return [finite_json(v) for v in value] return None if isinstance(value, float) and not np.isfinite(value) else value def same_solver_inputs(previous, current, solver): shared = ( "reference_sha256", "numerical_source_sha256", "n_theta", "boundary_subdivisions", ) specific = { "jax": ( "harmonics", "max_area_mm2", "controls", "angles_mechanical_rad", "currents_abc_a", "derivative_step_rad", ), "femm": ("femm_mesh_size_mm", "femm_precision"), } return all(previous[k] == current[k] for k in shared + specific[solver]) def snapshot(output): with zipfile.ZipFile(output / "source.zip", "w", zipfile.ZIP_DEFLATED) as archive: for directory in ( "src", "FEMM/pipeline", "config", "contracts", "material_datasheets", ): for path in sorted(Path(directory).rglob("*")): if path.is_file() and path.suffix in ( ".py", ".json", ".yaml", ".yml", ".js", ".html", ): archive.write(path, str(path)) for path in sorted(Path("FEMM").glob("*.py")): archive.write(path, str(path)) for path in ("main.py", "pyproject.toml", "uv.lock", "README.md"): archive.write(path, path) def initialize(output, args): import jax import jax.numpy as jnp from motor_sim_opt.environment.controller import six_step_phase_pattern from motor_sim_opt.jax_fea.context import FieldControls from motor_sim_opt.optimization.checkpoint import source_digest jax.config.update("jax_enable_x64", True) output.mkdir(parents=True, exist_ok=True) bundle = load_reference(REFERENCE) design = load_design(REFERENCE / "model.json", preliminary=True) poles = design.rotor.geometry.num_poles // 2 period = 2 * np.pi / poles count = args.max_angles if count not in (768, 1536): raise ValueError("max angles must be 768 or 1536") angles = (np.arange(count) + 0.5) * period / count patterns = np.asarray(six_step_phase_pattern(jnp.asarray(angles), poles)) config = { "format": "dense-triangle-femm-v1", "model": "3115", "max_angles": count, "target_angles": args.target_angles, "angular_levels": [96, 192, 384, 768] + ([1536] if count >= 1536 else []), "angles_mechanical_rad": angles.tolist(), "currents_abc_a": patterns.tolist(), "electrical_period_mechanical_rad": period, "pole_pairs": poles, "n_theta": args.nt, "harmonics": args.harmonics, "boundary_subdivisions": args.subdiv, "max_area_mm2": args.area, "femm_mesh_size_mm": args.femm_size, "femm_precision": 1e-10, "controls": asdict(FieldControls()), "reference_sha256": bundle["reference_sha256"], "numerical_source_sha256": source_digest(), "derivative_method": "periodic spectral derivative of directly solved open-circuit linkage, validated against independent central differences at h and h/2", "derivative_step_rad": period * 1e-4, "targets": { "kt_relative": 0.0025, "ke_relative": 0.0025, "torque_extrema_absolute_nm": 0.0001, "derivative_relative": 0.0025, }, "scope": "1 A six-step, one complete electrical period; many directly solved angles. No hardware/material calibration.", } identity = _canonical_sha256(config) if (output / "plan.json").exists(): if read(output / "plan.json")["identity"] != identity: raise ValueError( "configuration/source changed; use a fresh study directory" ) else: atomic_json(output / "plan.json", {"identity": identity, **config}) shutil.copytree(REFERENCE, output / "reference") snapshot(output) atomic_json( output / "environment.json", { "python": sys.version, "platform": platform.platform(), "packages": { p: importlib.metadata.version(p) for p in ( "jax", "jaxlib", "numpy", "scipy", "triangle", "plotly", "pydantic", ) }, }, ) return read(output / "plan.json") def task(plan, mode, index, *, offset=0.0, label=None): row = { "id": label or f"{mode}-{index:04d}", "mode": mode, "index": index, "angle_rad": plan["angles_mechanical_rad"][index] + offset, "phase_currents_a": plan["currents_abc_a"][index] if mode == "loaded" else [0.0, 0.0, 0.0], } row["request_sha256"] = _canonical_sha256(row) return row def grid_tasks(plan, count): stride = plan["max_angles"] // count return [ task(plan, mode, i) for i in range(0, plan["max_angles"], stride) for mode in ("zero", "loaded") ] def derivative_tasks(plan): tasks = [] for index in (0, plan["max_angles"] // 3, 2 * plan["max_angles"] // 3): for factor in (1.0, 0.5): for sign in (-1, 1): tasks.append( task( plan, "zero", index, offset=sign * factor * plan["derivative_step_rad"], label=f"fd-{index:04d}-{'h' if factor == 1 else 'half'}-{'minus' if sign < 0 else 'plus'}", ) ) return tasks def completed(directory, request, identity): path = directory / "tasks" / request["id"] / "result.json" if not path.exists(): return None row = read(path) if ( row.get("execution_sha256") != identity or row.get("request_sha256") != request["request_sha256"] ): raise ValueError(f"existing result identity mismatch: {path}") if row.get("status") != "accepted": return None if not all( k in row for k in ("torque_nm", "phase_linkage_wb_turn", "saved_artifacts") ): return None # FEMM helper publishes its intermediate readback before archiving completes. if not np.isfinite([row["torque_nm"], *row["phase_linkage_wb_turn"]]).all(): raise ValueError(f"nonfinite accepted evidence: {path}") for name, info in row.get("saved_artifacts", {}).items(): artifact = path.parent / name if not artifact.is_file() or artifact.stat().st_size != info["bytes"]: raise ValueError(f"missing or truncated artifact: {artifact}") return row def spectral_derivative(values, period): values = np.asarray(values) if values.ndim != 2 or not np.isfinite(values).all(): raise ValueError("finite (angle,phase) linkage array required") frequency = np.fft.fftfreq(len(values), d=period / len(values)) * 2 * np.pi return np.fft.ifft( 1j * frequency[:, None] * np.fft.fft(values, axis=0), axis=0 ).real def reduction(rows, plan): torque = np.asarray( [[r[mode]["torque_nm"] for r in rows] for mode in ("zero", "loaded")] ) linkage = np.asarray([r["zero"]["phase_linkage_wb_turn"] for r in rows]) patterns = np.asarray([r["loaded"]["phase_currents_a"] for r in rows]) phase_ke = spectral_derivative(linkage, plan["electrical_period_mechanical_rad"]) return { "kt": float(abs((torque[1] - torque[0]).mean())), "ke": float(abs(np.sum(phase_ke * patterns, axis=1).mean())), "phase_ke": phase_ke, "torque": torque, "linkage": linkage, "zero_extrema": [float(torque[0].min()), float(torque[0].max())], "loaded_extrema": [float(torque[1].min()), float(torque[1].max())], } def gather_grid(directory, plan, count): identity = read(directory / "execution.json")["identity"] pairs = [] for i in range(0, plan["max_angles"], plan["max_angles"] // count): pair = { m: completed(directory, task(plan, m, i), identity) for m in ("zero", "loaded") } if not all(pair.values()): return None pairs.append(pair) return pairs def convergence(directory, plan, count): current = gather_grid(directory, plan, count) coarse = gather_grid(directory, plan, count // 2) if current is None or coarse is None: return {"passed": False, "reason": "incomplete grid"} fine, old = reduction(current, plan), reduction(coarse, plan) shifted = reduction(current[1::2], plan) targets = plan["targets"] checks = { "kt_relative": abs(fine["kt"] - old["kt"]) / max(fine["kt"], 1e-12), "ke_relative": abs(fine["ke"] - old["ke"]) / max(fine["ke"], 1e-12), "shifted_kt_relative": abs(shifted["kt"] - old["kt"]) / max(fine["kt"], 1e-12), "shifted_ke_relative": abs(shifted["ke"] - old["ke"]) / max(fine["ke"], 1e-12), "zero_extrema_change_nm": float( np.max(abs(np.asarray(fine["zero_extrema"]) - old["zero_extrema"])) ), "loaded_extrema_change_nm": float( np.max(abs(np.asarray(fine["loaded_extrema"]) - old["loaded_extrema"])) ), "phase_ke_relative": float( np.linalg.norm(fine["phase_ke"][::2] - old["phase_ke"]) / max(np.linalg.norm(fine["phase_ke"][::2]), 1e-12) ), } passed = ( max(checks["kt_relative"], checks["shifted_kt_relative"]) <= targets["kt_relative"] and max( checks["ke_relative"], checks["shifted_ke_relative"], checks["phase_ke_relative"], ) <= targets["ke_relative"] and max(checks["zero_extrema_change_nm"], checks["loaded_extrema_change_nm"]) <= targets["torque_extrema_absolute_nm"] ) return { "passed": bool(passed), "coarse_angles": count // 2, "fine_angles": count, **checks, } def publish_status(directory, plan, phase, count=None): accepted = [ p for p in (directory / "tasks").glob("*/result.json") if read(p).get("status") == "accepted" ] payload = { "phase": phase, "accepted_tasks": len(accepted), "target_base_tasks": 2 * sampling_minimum(directory, plan), "current_grid": count, } if count: payload["convergence"] = convergence(directory, plan, count) atomic_json(directory / "status.json", payload) print(directory.name, json.dumps(payload), flush=True) def run_jax( output, plan, batch, probe_indices=None, probe_derivatives=False, request_driver=None, allow_retries=True, on_progress=None, ): import jax import jax.numpy as jnp from motor_sim_opt.core.materials import preliminary_library from motor_sim_opt.jax_fea import population as P from motor_sim_opt.jax_fea.context import FieldControls from motor_sim_opt.jax_fea.material_resolve import resolve_materials from motor_sim_opt.jax_fea.triangle_mesh import build_triangle_mesh jax.config.update("jax_platforms", "cpu") jax.config.update("jax_enable_x64", True) directory = output / "jax" directory.mkdir(exist_ok=True) execution = { "identity": _canonical_sha256( {"plan": plan["identity"], "solver": "jax-dense-v2-recorded-retries"} ), "batch": batch, "plan": plan["identity"], "retry_controls": [ {**plan["controls"], "nonlinear_iterations": 80, "maxiter": 40000}, { **plan["controls"], "nonlinear_iterations": 80, "maxiter": 40000, "nonlinear_relaxation": 0.5, }, ], } if not allow_retries: execution["retry_controls"] = [] execution["identity"] = _canonical_sha256( {"plan": plan["identity"], "solver": "jax-budget-test-no-retries"} ) if (directory / "execution.json").exists(): existing = read(directory / "execution.json") if existing != execution: previous = read(directory / "source-plan.json") if ( not same_solver_inputs(previous, plan, "jax") or existing["batch"] != batch or existing["retry_controls"] != execution["retry_controls"] ): raise ValueError("JAX physical inputs or execution controls changed") execution = existing execution["compatible_plan_identities"] = sorted( set( existing.get("compatible_plan_identities", [previous["identity"]]) + [plan["identity"]] ) ) atomic_json(directory / "execution.json", execution) if not (directory / "source-plan.json").exists(): atomic_json(directory / "source-plan.json", plan) design = load_design(output / "reference/model.json", preliminary=True) mesh = build_triangle_mesh( design, subdiv=plan["boundary_subdivisions"], n_theta=plan["n_theta"], max_area_mm2=plan["max_area_mm2"], ) inputs = P.build_population_inputs( [design], [mesh], [resolve_materials(design, preliminary_library())], subdiv=plan["boundary_subdivisions"], ) topology, motors = P.prepare_population_inputs( inputs, subdiv=plan["boundary_subdivisions"], n_harmonics=plan["harmonics"] ) motor = jax.tree.map(lambda x: x[0], motors) mesh_arrays = { "nodes_mm": np.asarray(mesh.nodes), "elements": mesh.quads, "region": np.asarray(mesh.region), "triangle_data": json.dumps(mesh.triangle_data), } if (directory / "mesh.npz").exists(): with np.load(directory / "mesh.npz", allow_pickle=False) as saved: if any( not np.array_equal(saved[key], value) for key, value in mesh_arrays.items() ): raise ValueError( "Triangulation changed on resume; old fields remain intact" ) else: np.savez_compressed(directory / "mesh.npz", **mesh_arrays) atomic_json( directory / "mesh-info.json", {"nodes": len(mesh.nodes), "triangles": len(mesh.quads)}, ) controls = FieldControls(**plan["controls"]) solve = jax.jit( jax.vmap(lambda a, c: P.sample_state(topology, motor, a, c, controls)) ) def make_retry(cfg): ctrl = FieldControls(**cfg) return jax.jit(lambda a, c: P.sample_state(topology, motor, a, c, ctrl)) retry_solvers = [make_retry(cfg) for cfg in execution["retry_controls"]] def execute(requests): pending = [ r for r in requests if completed(directory, r, execution["identity"]) is None ] for offset in range(0, len(pending), batch): chunk = pending[offset : offset + batch] padded = chunk + [chunk[-1]] * (batch - len(chunk)) start = time.monotonic() result = jax.device_get( solve( jnp.asarray([r["angle_rad"] for r in padded]), jnp.asarray([r["phase_currents_a"] for r in padded]), ) ) for lane, request in enumerate(chunk): flux, torque, solution = jax.tree.map( lambda x, lane=lane: x[lane], result ) target = directory / "tasks" / request["id"] target.mkdir(parents=True, exist_ok=True) attempts = [ { "controls": plan["controls"], "relative_residual": float(flux.relative_residual), "accepted": bool(flux.converged) and bool(torque.converged), } ] for retry, cfg in zip( retry_solvers, execution["retry_controls"], strict=True ): if bool(flux.converged) and bool(torque.converged): break if float( flux.relative_residual ) <= controls.nonlinear_tolerance and not bool( flux.material_characteristic_in_range ): break np.savez_compressed( target / f"attempt-{len(attempts)}-potentials.npz", xs=solution.xs, xr=solution.xr, as_ring=solution.as_ring, ar_ring=solution.ar_ring, ) flux, torque, solution = jax.device_get( retry( jnp.asarray(request["angle_rad"]), jnp.asarray(request["phase_currents_a"]), ) ) attempts.append( { "controls": cfg, "relative_residual": float(flux.relative_residual), "accepted": bool(flux.converged) and bool(torque.converged), } ) np.savez_compressed( target / "potentials.npz", xs=solution.xs, xr=solution.xr, as_ring=solution.as_ring, ar_ring=solution.ar_ring, ) ok = ( bool(flux.converged) and bool(torque.converged) and bool(flux.material_characteristic_in_range) ) row = { **request, "execution_sha256": execution["identity"], "status": "accepted" if ok else "rejected", "torque_nm": float(torque.value), "phase_linkage_wb_turn": np.asarray(flux.phase_linkage).tolist(), "relative_residual": float(flux.relative_residual), "material_in_range": bool(flux.material_characteristic_in_range), "nonlinear_relative_residual": float( flux.nonlinear_relative_residual ), "stator_iron_peak_t": float( np.max( np.where( np.asarray(motor.field.stator_iron) > 0, flux.stator_flux_density_t, 0, ) ) ), "rotor_iron_peak_t": float( np.max( np.where( np.asarray(motor.field.rotor_iron) > 0, flux.rotor_flux_density_t, 0, ) ) ), "batch_seconds_including_compilation": time.monotonic() - start, "attempts": attempts, "saved_artifacts": { "potentials.npz": { "bytes": (target / "potentials.npz").stat().st_size, "sha256": file_digest(target / "potentials.npz"), } }, } atomic_json(target / "result.json", finite_json(row)) print( "jax", request["id"], "accepted=", ok, "T=", row["torque_nm"], flush=True, ) if not ok: publish_status(directory, plan, "rejected state") raise ValueError( f"JAX field rejected; state retained: {request['id']}" ) publish_status(directory, plan, "solving") if on_progress is not None: on_progress() if request_driver is not None: request_driver(execute) return if probe_indices is not None: execute([task(plan, m, i) for i in probe_indices for m in ("zero", "loaded")]) if probe_derivatives: execute(derivative_tasks(plan)) publish_status(directory, plan, "probe complete") return minimum_angles = sampling_minimum(directory, plan) for count in plan["angular_levels"]: execute(grid_tasks(plan, count)) publish_status(directory, plan, "grid complete", count) if count >= minimum_angles and convergence(directory, plan, count)["passed"]: break execute(derivative_tasks(plan)) publish_status(directory, plan, "complete", count) def run_femm(output, plan, probe_indices=None, probe_derivatives=False): # Native children inherit this lower CPU priority; JAX remains at normal priority. os.nice(10) import jax from pipeline.femm_utils.backend import backend_fingerprint, open_native from pipeline.femm_utils.build_model import build from pipeline.femm_utils.solve import _solve_task jax.config.update("jax_platforms", "cpu") jax.config.update("jax_enable_x64", True) directory = (output / "femm").resolve() directory.mkdir(exist_ok=True) bundle = load_reference(output / "reference") if bundle["reference_sha256"] != plan["reference_sha256"]: raise ValueError("reference input identity changed") if not (directory / "case").exists(): prepare( load_design(output / "reference/model.json", preliminary=True), directory / "case", n_theta=plan["n_theta"], subdivisions=plan["boundary_subdivisions"], mesh_kind="mapped", ) if not (directory / "model").exists(): build( directory / "case", directory / "model", motion_band=True, mesh_size_mm=plan["femm_mesh_size_mm"], precision=plan["femm_precision"], simplify_collinear=True, ) from pipeline.femm_utils.artifacts import load_preparation spec, _ = load_preparation(directory / "case") if ( spec.design != bundle["model"]["design"] or spec.materials != bundle["model"]["resolved_materials"] or spec.n_theta != plan["n_theta"] or spec.radial_subdivisions != plan["boundary_subdivisions"] ): raise ValueError( "FEMM prepared model does not match the declared physical inputs" ) inputs = { "plan": plan["identity"], "backend": backend_fingerprint(), "base_fem_sha256": file_digest(directory / "model/base.fem"), "case_model_sha256": file_digest(directory / "case/model.json"), "compiler_sha256": file_digest(Path("FEMM/pipeline/femm_utils/build_model.py")), } execution = {"identity": _canonical_sha256(inputs), "inputs": inputs} if (directory / "execution.json").exists(): existing = read(directory / "execution.json") if existing != execution: if not same_solver_inputs( read(directory / "source-plan.json"), plan, "femm" ): raise ValueError("FEMM mesh or numerical source changed") physical = {k: v for k, v in inputs.items() if k != "plan"} previous = {k: v for k, v in existing["inputs"].items() if k != "plan"} if physical != previous: raise ValueError("FEMM model, backend or compiler changed") execution = existing execution["compatible_plan_identities"] = sorted( set( existing.get( "compatible_plan_identities", [existing["inputs"]["plan"]] ) + [plan["identity"]] ) ) atomic_json(directory / "execution.json", execution) if not (directory / "source-plan.json").exists(): atomic_json(directory / "source-plan.json", plan) with closing(open_native(timeout_seconds=1200)) as femm: femm.opendocument(str(directory / "model/base.fem")) femm.mi_createmesh() def execute(requests): for request in requests: if completed(directory, request, execution["identity"]) is not None: continue if shutil.disk_usage(output).free < 12 * 1024**3: raise OSError( "Fewer than 12 GiB free; checkpoint preserved instead of exhausting disk" ) start = time.monotonic() target = directory / "tasks" / request["id"] result = _solve_task(femm, request, target, execution["identity"]) ans = target / "base.ans" with ( ans.open("rb") as source, gzip.open( target / "base.ans.gz", "wb", compresslevel=1 ) as destination, ): shutil.copyfileobj(source, destination) compressed = target / "base.ans.gz" # Preserve a raw-answer hash and the compressed payload hash before deleting our own intermediate. row = { **result, **request, "torque_nm": result["outer_rotor_torque_nm"], "phase_linkage_wb_turn": result["phase_linkages_wb_turn"], "seconds": time.monotonic() - start, "saved_ans": "base.ans.gz", "saved_artifacts": { "base.ans.gz": { "bytes": compressed.stat().st_size, "sha256": file_digest(compressed), }, "base.fem": { "bytes": (target / "base.fem").stat().st_size, "sha256": file_digest(target / "base.fem"), }, }, } atomic_json(target / "result.json", row) ans.unlink() print( "femm", request["id"], "T=", row["torque_nm"], "s=", row["seconds"], flush=True, ) publish_status(directory, plan, "solving") if probe_indices is not None: execute( [task(plan, m, i) for i in probe_indices for m in ("zero", "loaded")] ) if probe_derivatives: execute(derivative_tasks(plan)) publish_status(directory, plan, "probe complete") return minimum_angles = sampling_minimum(directory, plan) for count in plan["angular_levels"]: execute(grid_tasks(plan, count)) publish_status(directory, plan, "grid complete", count) if ( count >= minimum_angles and convergence(directory, plan, count)["passed"] ): break execute(derivative_tasks(plan)) publish_status(directory, plan, "complete", count) def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT) parser.add_argument("--init", action="store_true") parser.add_argument("--jax", action="store_true") parser.add_argument("--femm", action="store_true") parser.add_argument("--report", action="store_true") parser.add_argument("--monitor", action="store_true") parser.add_argument("--nt", type=int, default=1344) parser.add_argument("--harmonics", type=int, default=630) parser.add_argument("--subdiv", type=int, default=64) parser.add_argument("--area", type=float, default=0.0125) parser.add_argument("--femm-size", type=float, default=0.0625) parser.add_argument("--max-angles", type=int, default=1536) parser.add_argument( "--target-angles", type=int, choices=(192, 384, 768, 1536), default=384 ) parser.add_argument("--batch", type=int, default=4) parser.add_argument("--probe-indices", type=int, nargs="+") parser.add_argument("--probe-derivatives", action="store_true") args = parser.parse_args() plan = ( initialize(args.output, args) if args.init else read(args.output / "plan.json") ) if args.jax: run_jax( args.output, plan, args.batch, args.probe_indices, args.probe_derivatives ) if args.femm: run_femm(args.output, plan, args.probe_indices, args.probe_derivatives) if args.report or args.monitor: from dense_waveform_report import write_report while True: write_report(args.output) if not args.monitor or all( (args.output / s / "status.json").exists() and read(args.output / s / "status.json")["phase"] == "complete" for s in ("jax", "femm") ): break time.sleep(30) if __name__ == "__main__": main()