google-deepmind / torax

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

Precompute rho interpolation for interpolated parameters and use `InterpolatedVarSingleAxis` logic to do interpolation. #340

Closed copybara-service[bot] closed 1 month ago

copybara-service[bot] commented 1 month ago

Precompute rho interpolation for interpolated parameters and use InterpolatedVarSingleAxis logic to do interpolation.

Now that all the profile conditions that are InterpolatedVarTimeRho are only being computed once it's worth precomputing what we can in the constructor.

Performance test on iterhybrid_rampup Before: Sim and params build time: 2.88s Simulated 80.00s of physics in 7.40s of wall clock time. compile time: 12.73s

After: Sim and params build time: 3.47s Simulated 80.00s of physics in 6.77s of wall clock time. Compile time 13.23s