Refactor Snapshot class to avoid the set_geometry step