cleaned up code, made plotting easier