Handle infinite and fixed bounds in ALM inner solves

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)

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))