Refactor visualization code into Python viz package