I have been working on a small Python package to solve a class of PDEs using scipy.integrate.solve_ivp. As discretizations are made finer, runtime becomes a bottleneck—especially when I need to solve the PDE for a large number of different initial conditions. I would like to make use of GPU acceleration to speed things up, and I am planning on how to best integrate GPU-based computations into my current implementation. Is it feasible to keep using solve_ivp but make the ODE function transfer data to the GPU, perform a forward step with something like CuPy, and then transfer back to CPU? Or would the transfer overhead be too large for practical benefit? I also considered batched solve_ivp by concatenating many ICs into a single array, then reshaping in the ODE function, performing a forward step for all solutions in parallel, and then reshaping back. Would this work or would it clash with adaptive time-stepping methods? More generally, if I need to solve the PDE for many ICs, what would be the best (or easiest) way to make use of GPU acceleration (preferably using Python)? Any recommendations or best practices would be greatly appreciated. Continue reading...
