Refactor unit defintions