Get `ProximalProjection` to work with augmented lagrangian

[f0uriest]

Needed when we want to exactly enforce equilibrium constraint which may be useful for stability optimization

[dpanici]

What needs to be done for this? something like allowing ProximalProjection to accept nonlinear constraints other than the eq ones, or just wrapping the nonlinear constraints that are not eq constraints in some projection wrapper as well?

[dpanici]

@f0uriest bump on my question?

[f0uriest]

So basically right now we have the ProximalProjection act like a regular objective, with compute_scaled_error, jac_scaled_error etc. We would want to have it also act like a constraint, the difficulty being that both objective and constraint have the same method names. We probably don't want a second projection wrapper since that would be more expensive (re-solving twice) and may lead to the two equilibria going out of sync or something.

Here's a rough sketch of what I was thinking

class ProximalProjection():

    def __init__(self, eq, objective, eq_constraint, nonlinear_constraints):
        ...

    def build():
        self.wrapped_objective = namedtuple("wrapped_objective", ["compute_scaled_error", "jac_scaled_error"])(self.compute_scaled_error, self.jac_scaled_error)
        self.wrapped_constraint = namedtuple("wrapped_constraint", ["compute_scaled_error", "jac_scaled_error"])(self.con_compute_scaled_error, self.con_jac_scaled_error)

    def compute_scaled_error(self, x):
        x = self.update_eq(x)
        return self.objective.compute_scaled_error(x)

    def con_compute_scaled_error(self, x):
        x = self.update_eq(x)
        return self.nonlinear_constraints.compute_scaled_error(x)

The basic idea is we add the stuff for the constraints, with different names, then use named tuples as a sort of local namespace for the objective and constraints. then when we actually call the low level optimizer eg lsq-auglag we pass objective=prox.wrapped_objective, constraint=prox.wrapped_constraint instead of the prox itself.