Further vectorized the diffusion calculations -- speeds things up.