Changes to set up

Changes to the code from class (up to 3.)