search

patrick-kidger / optimistix

Nonlinear optimisation (root-finding, least squares, ...) in JAX+Equinox. https://docs.kidger.site/optimistix/
Apache License 2.0
265 stars 12 forks source link

Will constrained optimization be supported? #45

Open zhouxian opened 4 months ago

zhouxian commented 4 months ago

Hi, thank you for the amazing library!

I was wondering if minimizing with user-specified bounds, or algorithms like projected gradient descent are supported? If not, what would be the best practice you suggest if we are trying to solve things like argmin(norm(Ax+b)) s.t. x >=0?

Thanks in advance!

patrick-kidger commented 4 months ago

So right now we do support hypercube constraints in our root-finding methods (e.g. Newton or Bisection). And moreover it is sort-of-possible to cast a minimisation problem as a root-finding problem, by seeking a root of the gradient of the function. (You have to be careful as this may also converge to any other external point.) So I think the particular example you've given could be tackled today, although I'm not sure how efficient it would be as compared to a dedicated solver.

Other than that we don't currently support any kind of constraints in our solves. We may do at some point, but as always developer time is the limiting factor!

mjo22 commented 4 months ago

Hello! Was also curious on this front. In particular, I am wondering if there are plans to support box constraints in BFGS, or if this can already be done in the current framework. jaxopt seems to have an implementation of this, but I’m not sure how well maintained it is.

patrick-kidger commented 4 months ago

For something simple like box constraints, then it should be fairly easy to add that in. We'd be happy to take a PR adding that. (Note that box constraints are already implemented for the root-finders, so you could follow a similar pattern to there.)

f0uriest commented 4 months ago

If box constraints are supported it shouldn't be too much work to add an augmented Lagrangian type method for general constraints. Could maybe be like a "meta solver" where you can wrap another solver to manage the unconstrained sub problems? I can hopefully take a look at it once I finish up a few other projects.