warp VS prysm

About to embark on some physics simulation experiments and am hoping to get some input on available options for making use of my GPU through Python: Currently reading the docs for NVIDIA Warp, and CUDA python but would appreciate any other pointers on available packages or red flags on packages that are more hassle than they are worth to learn.

I would recommend using NanoBind, the follow up of PyBind11 by the same author (Wensel Jakob), and move as much performance critical code to C or C++. https://github.com/wjakob/nanobind

If you really care about performance called from Python, consider something like NVIDIA Warp (Preview). Warp jits and runs your code on CUDA or CPU. Although Warp targets physics simulation, geometry processing, and procedural animation, it can be used for other tasks as well. https://github.com/NVIDIA/warp

Jax is another option, by Google, jitting and vectorizing code for TPU, GPU or CPU. https://github.com/google/jax

My current concept is to just combine zernike polynomials with a random factor and calculate the PSF from that, which can be somewhat easily be done with the prysm library. These PSFs can then be convolved with circular and gaussian kernels for modelling additional defocus and accounting for other stuff like the AA filter. Then I'd add chromatic aberration by offseting/scaling the PSFs for each channel. Some generated kernels already look pretty good when comparing them to stars in astrophotography images, but others not so much.

https://github.com/brandondube/prysm (caveat emptor: mine)

You can trace about 1 billion raysurfaces per second in pure python with CuPy, or a few million raysurfaces per second on CPU.

New paper from /u/BDube_Lensman using prysm to model NASA's Roman Coronagraph

This free book is what this free code is based on

Prysm Originally for diffraction type optics but seems to able to handle...everything? Performance as a priamary concern, GPU acceleration, proven JPL heritage :) Raytracing is however still experimental and without docs, generally whilst the library looks excellent if you're an optics person already I think I lack a bit of the base fundamental knowledge to really use it powerfully from just the API reference. I can see BDube has some raytracing example code in some of the issues I could probably adapt and muddle my way through at least. No guis is mildly annoying for a noob like myself, but I can work my way around matplotlib-ing just fine instead i'm sure.

You may want to steal my shim set since it lets you hot swap Numpy<-->cupy at runtime

I do computational diffraction with large manycore servers and GPUs at a FFRDC. The difference between MKL and not MKL is the difference between hitting enter and getting a result in an hour or two vs tomorrow.