Creating this issue as reference point for working on this issue - [ ] Test in larger dataset - [ ] Explore whether test coverage can be improved - [ ] Add documentation