Clean up code

[ismael-mendoza]

• [ ] #71
• [ ] commented out code
• [ ] places where black formatting makes code particularly difficult to read (refactor into multiple lines) e.g. here
• [ ] anything else that might make code difficult to read or use?