add plot function to the original version