Minor improvements to proof of concept