ast: Improve MatrixMap simplification

Allow a full form MatrixMap to be simplified to a diagonal form
MatrixMap if all the off-diagonal values are zero.