code cleanup
fix a bug (closes #7)