some improvements to code and text
added figures