Cool, looks good! And makes sense to check if sqrt is enough since it seems to be MUCH faster (not surprisingly). How did you come up with the 1e-14 limit?

The 1e-14 limit is a bit arbitrary. I thought that sqrt should only be used if the matrix z2 is certainly diagonal, which is why I used a limit as small as possible. From numerical uncertainties, differences of few times 1e-16 are very common, so if I only consider factors of 10, the smallest reasonable limit might be 1e-15. But some numerical uncertainties might be a bit larger and reach 1e-15, which is why I chose 1e-14. I'm not sure if that's a good choice. I'm also not sure if one should discard the off-diagonal elements in z2 if sqrt is used (if the limit is as small as 1e-14 this shouldn't make a big difference).