Simplify code and fix computation of projected gradient