Also improve efficiency from cubic to quadratic by avoiding taking the trace of a mat-mat multiplication, and rather just summing the formula for the diagonal entries. Include a unit test to avoid regressions.