Tidied up some of the code in the model class