GENERAL EXAM: Fix plotting.

This is a hacky fix so that plotting works better for
my general exam. All of this code will be replaced soon
in a refactor.