[pastewka]

Step 3 of the plan: the optimization loop stays on the GPU. Stacked on #39, which it needs for the regularization's device path. Requires NuMPI 0.19.0 (muSpectre/NuMPI#106), so it can only be merged after that release.

What changes

• optimize_trust_region passes the design to tr_newton_bounded as an array of the fields' module (CuPy on a GPU). With NuMPI #106 the trust region iterates on that device, so the design, gradient and Hessian products never leave it. Only scalars reach the host (dot products and maxima, plus the objective). The result is still returned on the host. The callback receives device arrays, which the docstring now says.
• StressTargetProblem computes in the fields' module throughout and returns its gradient and Hessian product as the same kind of array it was given: device arrays for the trust region, host arrays (one copy each way) for L-BFGS and the tests. The cached design, ensure_state's comparison and the tolerance controller's gradient norm move to the device.
• NodalElementMap takes device input, and scatter/scatter_mean(..., on_device=True) return the gradient on the device.
• simulate.py's callback sums the iterate with rho.sum(), which works for host and device arrays.

What this removes

Per objective evaluation, two full-field copies (design in, gradient out). Per Hessian product, two more (direction in, product out). The trust region's own vector arithmetic (the optimization's "(other)", about 22 s of 842 in the last 4-rank run) also moves from host NumPy to the GPU. The gain is modest, but nothing that grows with the grid stays on the host.

Tests

• New GPU tests in test/test_gpu.py:

  • device input gives device output, matching host input;
  • the trust region iterates on device arrays (the callback checks this), and its objective history and result follow the CPU run over 5 iterations.

  They're skipped here (no CuPy) and need a run on MI300A.

• The CPU suite passes (249) against NuMPI #106.

🤖 Generated with Claude Code