Small additions to code