Cleaned up code.

Closes #2.