diff --git a/essos/augmented_lagrangian.py b/essos/augmented_lagrangian.py index f2393e90..27cac2a3 100644 --- a/essos/augmented_lagrangian.py +++ b/essos/augmented_lagrangian.py @@ -758,7 +758,35 @@ class ALM(NamedTuple): -#Inner solve with optimistix L-BFGS; a (lower, upper) box is enforced by a tanh reparametrization +def _bound_reparametrization(params, bounds): + """Finite intervals use tanh; one-sided bounds reflect at their finite endpoint.""" + lower, upper = (jnp.asarray(bound, dtype=params.dtype) for bound in bounds) + has_lower, has_upper = jnp.isfinite(lower), jnp.isfinite(upper) + both = has_lower & has_upper + # Keep inactive branches finite too: their derivatives are traced by JAX. + lo, hi = jnp.where(has_lower, lower, 0.0), jnp.where(has_upper, upper, 0.0) + fixed = both & (lo == hi) + mid = (lo + hi) / 2 + half = jnp.where(both & ~fixed, (hi - lo) / 2, 1.0) + margin = jnp.maximum(1e-12, jnp.finfo(params.dtype).eps) + box_initial = jnp.arctanh(jnp.clip((params - mid) / half, -1 + margin, 1 - margin)) + distance = jnp.where(has_lower, params - lo, hi - params) + initial = jnp.where(both, box_initial, + jnp.where(has_lower | has_upper, distance, params)) + initial = jnp.where(fixed, 0.0, initial) + + def to_params(z): + # Reflection includes the endpoint exactly and lets a boundary start + # move inward, unlike an exponential/softplus inverse at zero distance. + one_sided = jnp.where(has_lower, lo + jnp.abs(z), hi - jnp.abs(z)) + result = jnp.where(both, mid + half * jnp.tanh(z), + jnp.where(has_lower | has_upper, one_sided, z)) + return jnp.where(fixed, lo, result) + + return to_params, initial + + +#Inner solve with optimistix L-BFGS; bounds are enforced by reparametrization def ALM_model_lbfgs(constraints: BaseConstraint,#List of constraints loss= lambda x: 0., #function which represents the loss (Callable, default 0.) model_lagrangian='Standard', @@ -812,8 +840,7 @@ def update_fn(params, lag_state,grad,info,beta=beta,mu_max=mu_max,alpha=alpha,ga if bounds is None: to_x, z0 = (lambda z: z), main_params else: - mid, half = (bounds[1] + bounds[0]) / 2, (bounds[1] - bounds[0]) / 2 - to_x, z0 = (lambda z: mid + half * jnp.tanh(z)), jnp.arctanh(jnp.clip((main_params - mid) / half, -1 + 1e-12, 1 - 1e-12)) + to_x, z0 = _bound_reparametrization(main_params, bounds) solver = optimistix.LBFGS(rtol=omega_min, atol=omega_min, history_length=history_size) sol = optimistix.minimise(lambda z, lp: lagrangian(to_x(z), lp, **kargs), solver, z0, lagrange_params, has_aux=True, max_steps=inner_maxiter, throw=False) diff --git a/tests/test_augmented_lagrangian.py b/tests/test_augmented_lagrangian.py index d03382f6..04b896dd 100644 --- a/tests/test_augmented_lagrangian.py +++ b/tests/test_augmented_lagrangian.py @@ -225,3 +225,33 @@ def test_ALM_lbfgs_respects_bounds(self): if __name__ == "__main__": pytest.main([__file__]) + + +@pytest.mark.parametrize('bounds,start,target', [ + ((0., jnp.inf), [1.], [3.]), + ((0., jnp.inf), [0.], [3.]), + ((-jnp.inf, 0.), [-1.], [-3.]), + ((-jnp.inf, 0.), [0.], [-3.]), + ((-jnp.inf, jnp.inf), [0.], [150.]), + ((2., 2.), [2.], [2.]), + ((jnp.array([0., -jnp.inf, -jnp.inf, 2., -1.]), + jnp.array([jnp.inf, 0., jnp.inf, 2., 1.])), + [1., -1., 0., 2., 0.], [3., -3., 7., 2., .5]), +]) +def test_alm_lbfgs_supports_infinite_and_fixed_bounds(bounds, start, target): + # These bounds were forwarded to JAXopt by the former public entry point. + # Check actual ALM updates, including mixed bound types in one parameter array. + from essos.augmented_lagrangian import ALM_model_jaxopt_lbfgsb + target = jnp.asarray(target) + constraint = eq(lambda x: x - target) + x = jnp.asarray(start) + params = x, constraint.init(x) + solver = ALM_model_jaxopt_lbfgsb(constraint, bounds=bounds) + state, gradient, info = solver.init(params) + for _ in range(10): + params, state, gradient, info = solver.update(params, state, gradient, info) + assert bool(jnp.all(jnp.isfinite(params[0]))) + assert bool(jnp.all(jnp.isfinite(gradient[0]))) + assert bool(jnp.all(params[0] >= bounds[0])) + assert bool(jnp.all(params[0] <= bounds[1])) + assert bool(jnp.allclose(params[0], target, rtol=1e-3, atol=1e-3))