Refactor pso
Add test coverage to pso