Add a function that returns a diagonal matrix and use it in a few places. Currently only a matrix whose diagonal elements are all the same is implemented.