MAINT: Clean-up code to accomodate a single x axis

Fixes #4