GPU acceleration

[seb-0-0]

gpu acceleration

Implement functionalities to Speed up the calculation

Task1 : matrix diagonalization

Accelerate the matrix diagonlization and have a single test file that :

• runs successfully with the numpy library
• runs successfully with the cupy library
• runs for a case for which the gpu demonstrates acceleration

[biochem-fan]

I emailed this in January but to share it more widely:

If we want to calculate only at a single thickness, matrix diagonalization is less efficient and not necessary. We can directly calculate matrix exponentials.

Ultramicroscopy "Stacked-Bloch-wave electron diffraction simulations using GPU acceleration"
https://doi.org/10.1016/j.ultramic.2014.03.003

This calculates the matrix exponential directly, without diagonalization. The paper claims this approach is faster and more suitable for GPU. This is very interesting. Although the paper does not mention derivative calculation, it should also be possible.

The matrix exponential is available in SciPy as scipy.linalg.expm and its derivative as scipy.linalg.expm_frechet. PyTorch and JAX also have theirs with GPU support.

[ronandrevon]

That sounds super interesting thank you for sharing this. We ll have a look at implementing it.

[ronandrevon]

I have tested matrix diagonalization on a remote gpu and it appears that I was starting to get acceleration for 2000x2000 matrix diagonalization. So to get the blochwave code to run larger matrices like that you have 2 options :

• increase the excitation error threshold Smax so more beams in the vicinity of the ewald sphere are included in the matrix
• increase the index Nmax so larger beams may also be included in the eigen matrix.

misc

You can make regular commits but ideally the new commits should leave the code in a working state.