Fixing the demo to sample a single particle and improving plotting