fix #1776

let user pass an axis argument, rather than always picking z level, that allows for non cf use also pep8ed the file