Modify code to add the p.A and A^2/2 terms (when using a non-uniform vector potential) to the kinetic energy output.