Numerical error with JAX

[shawnwwimer]

I was using JAX for the backend and couldn't get some simulations to agree with an analytical form. I noticed that changing the backend to the CPU fixed this problem and found that JAX uses a low precision by default for some operations:

• Matrix multiplication precision API jax-ml/jax#2161
• Significant numerical differences of dot_general on TPU jax-ml/jax#16349

I don't have the time to test right now, but from that second issue: "Try setting jax.default_matmul_precision to float32". If anybody runs into a similar problem this may be the cause. If so, it may be good to note it in the readme.

[rafael-fuente]

There is a significant accuracy loss for diffraction calculations when using JAX float32 (which is the default when importing JAX)

Can you check if enabling JAX x64 solves your problem?

You need to add the following lines before importing diffractsim:

import jax
jax.config.update("jax_enable_x64", True)

[shawnwwimer]

Yes, that does fix it. The difference is definitely the precision. I brought it up here simply because I didn't see a related issue here and I was surprised to find out it was due to default JAX behavior.