Finish changing nbody code for spartan

Add `set_diagonal` function.