The code for sde simulation seems to very slow. I should do a profiling and maybe reimplement this in jax.