google / google/jaxopt

Native optimization methods failing on sharp function which scipy wrapper succeeds on

Open
#430 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
1.1k
Forks
76
Avg merge
2d 21h
Merged PRs (30d)
1

Description

Hello!

I'm using Jax/Jaxopt for a different problem then is likely intended: optics simulations based on [transfer matrix methods](https://en.wikipedia.org/wiki/Transfer-matrix_method_(optics)). Effectively, each element of a system is modeled as a matrix, and simulating a full system results in the cascaded multiplication of several matrices, making jax and jaxopt very intriguing libraries to help speed this up.

The problem I'm working with boils down to minimizing a complicated function of one parameter that essentially just results in a Lorentzian resonance occurring at a parameter value of ~10E-6 with a width of ~1E-9, so for my MVE i'm just using a lorentzian function.

Now the problem. Ideally I want to minimize using native jaxopt so that I can easily `vmap` and `grad` the result, however all the native routines seem to completely fail at optimizing this function. Meanwhile, the scipy wrapper succeeds with ease. Here's some example code for minimizing a narrow lorentzian dip with a few methods given the same starting position and tolerance.

```python
from jax import jit,grad,vmap
import jax.numpy as jnp
import jaxopt as jopt
from jaxopt.projection import projection_box
import matplotlib.pyplot as plt
@jit
def lorentz(x,amp,mean,gamma,offset):
delta = x-mean
return amp * gamma**2/(gamma**2+4*delta**2) + offset

@jit
def lor(x):
return lorentz(x,-0.8,10.0E-6,10E-10,1.0)

x_init = 10.1E-6
x_min = 9.5E-6
x_max = 10.5E-6
tol = 1E-8

xs = jnp.linspace(x_min,x_max,10000)
plt.plot(xs,lor(xs))
plt.scatter(x_init,lor(x_init),label='initial')

# Projected Gradient
solver = jopt.ProjectedGradient(lor,projection_box,tol=tol)
params = solver.run(x_init,
hyperparams_proj=(x_min,x_max))
plt.scatter(params.params,lor(params.params),
label='projected',zorder=10)

# BFGS
# Stepsize avoids pinging off to large values
solver = jopt.BFGS(lor,tol=tol,stepsize=1E-9)
params = solver.run(x_init)
plt.scatter(params.params,lor(params.params),label='bfgs',zorder=10)

# Scipy Bounded LBFGSB
solver = jopt.ScipyBoundedMinimize(fun=lor,
tol=tol,
method="l-bfgs-b")
params = solver.run(x_init,bounds=(x_min,x_max))
plt.scatter(params.params,lor(params.params),label='scipy - lbfgsb')
plt.legend()

# Scipy BFGS
solver = jopt.ScipyMinimize(fun=lor,
tol=tol,
method="bfgs")
params = solver.run(x_init)
plt.scatter(params.params,
lor(params.params),
label='scipy - bfgs',
marker='x')
plt.legend()

# Scipy Newton
solver = jopt.ScipyMinimize(fun=lor,
tol=tol,
method="newton-cg")
params = solver.run(x_init)
plt.scatter(params.params,
lor(params.params),
label='scipy - newton',
marker='+')
plt.legend()
```
And the resulting plot
![scrot](https://github.com/google/jaxopt/assets/14968322/74321e23-6316-4ebc-932e-740389f0a31b)

Regardless of method scipy nails the dip, while the other methods blow right past it.

To note, I'm using projected gradient as ideally I want to be able to include bounds as the real problem becomes periodic in the parameter. However, the gradient should always be well behaved enough to point the solver to the nearest dip, so AFAIK I shouldn't need to use bounds, hence why bfgs/lbfgs is viable.

This might just be a question of needing to fine tune some parameters, but I haven't been able to get anything to work better than these examples demonstrate.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.