Use ax in plotting calls. Inherit base docstring

Signed-off-by: Jonas Kalderstam <jonas@kalderstam.se>