Updated UML to make better use of model classes