Add features to PGM solver classes

[bwohlberg]

Add features to PGM solver classes that were recently added to ADMM solver classes:

• Timing of solver iterations
• Iteration counter that maintains state between solve calls
• Optional callback function argument of solve method

Also cleanup up flake8 warnings, primarily related to over-length lines.

[lukepfister]

I'd be wary of adding extra state to the optimizers. Jax doesn't play nice with side effects.

In fact, we discussed eventually making the internals to PGM and ADMM entirely stateless at some point-- you can see bits of that with the way that the step size is always passed as an argument to the step function, even though it is an attribute of the PGM instance.

[bwohlberg]

I'd be wary of adding extra state to the optimizers. Jax doesn't play nice with side effects.

In fact, we discussed eventually making the internals to PGM and ADMM entirely stateless at some point-- you can see bits of that with the way that the step size is always passed as an argument to the step function, even though it is an attribute of the PGM instance.

Noted, but the same mechanism is already in ADMM. I'll go ahead and merge this soon, but let's discuss when possible, and make an issue for appropriate changes to both ADMM and PGM.