google-deepmind / torax

TORAX: Tokamak transport simulation in JAX
https://torax.readthedocs.io
Other
345 stars 31 forks source link

Modify `InterpolatedVarTimeRho` to be jittable. #341

Closed copybara-service[bot] closed 1 month ago

copybara-service[bot] commented 1 month ago

Modify InterpolatedVarTimeRho to be jittable.

Performance test on iterhybrid_rampup Before: Compile time 12.58s Simulated 80.67 seconds of physics in 9.25s of wall clock time Sim and params build time: 3.04s

After: Compile time: 12.93s Simulated 80.67 seconds of physics in 9.25s of wall clock time. Sim and params build time: 5.63s