With the introduction of IrregSpacedGrids in #211 and ShockGrids in #206 we started using jnp.searchsorted to calculate coordinates on the grids. Using a binary search introduces a step into the coordinate calculation, that has a time complexity of O(log(n)) and could therefore negatively impact the runtime for larger grids. In a first small experiment with grids up to 200 points in the Mahler-Yum Model, the runtime stayed the same but this could be different for larger grids or different models that use more interpolation.
As using jnp.searchsorted can reduce the complexity of the code, it would good to perform a bigger set of benchmarks with IrregSpacedGrids and also have a look at the JAX computation graph to get a better understanding what optimizations are done on the coordinate calculation and how important it actually is for the overall runtime.
With the introduction of
IrregSpacedGridsin #211 andShockGridsin #206 we started usingjnp.searchsortedto calculate coordinates on the grids. Using a binary search introduces a step into the coordinate calculation, that has a time complexity ofO(log(n))and could therefore negatively impact the runtime for larger grids. In a first small experiment with grids up to 200 points in the Mahler-Yum Model, the runtime stayed the same but this could be different for larger grids or different models that use more interpolation.As using
jnp.searchsortedcan reduce the complexity of the code, it would good to perform a bigger set of benchmarks withIrregSpacedGridsand also have a look at the JAX computation graph to get a better understanding what optimizations are done on the coordinate calculation and how important it actually is for the overall runtime.