As you are using dolfinx, consider for instance Speeding up time-dependent diffusion equation - #4 by dokken