To speed up the computation of the reduced density matrix you might try Tullio, reshape your density matrix into a tensor and then sum over the indices you don’t need.
When you embed it back into the full Hilbert space do you mutate an existing matrix or do you create a new matrix?