Merge pull request #28 from CMBSciPol/docs

Docs

Author: ASKabalan

docs/minimization.md

@@ -95,6 +95,50 @@ This prevents numerical under/overflow across the extreme dynamic range between
 
 The step size α is capped at the distance to the nearest bound (α_max), then a line search finds the optimal α in [0, α_max]. If a parameter hits a bound, it becomes an active constraint.
 
+## Termination Control (ADABK solvers)
+
+Four parameters control when ADABK solvers decide to stop. Pass them via the `options` dict of `minimize()`:
+
+```python
+final_params, state = minimize(
+    fn=my_loss_fn,
+    init_params=params,
+    solver_name="ADABK0",
+    lower_bound=lower,
+    upper_bound=upper,
+    options={
+        "cooldown": 50,              # default: 20
+        "min_steps": 200,            # default: 10
+        "verbose_print": True,       # default: False
+        "max_linesearch_steps": 100, # default: 50
+    },
+)
+```
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `cooldown` | `20` | Steps to suppress termination after a constraint is released. Prevents premature convergence caused by a transient function spike when a bound constraint opens. |
+| `min_steps` | `10` | Minimum iterations before termination is ever considered. Useful when the initial gradient is near zero but the landscape is not yet explored. |
+| `verbose_print` | `False` | Print per-step diagnostics: current `f`, `f_diff`, `best_f`, cooldown status, and termination decision. Uses `jax.debug.print` so it is JIT-compatible. |
+| `max_linesearch_steps` | `50` | Maximum bounded line-search steps per iteration. |
+
+### How termination is decided
+
+Termination requires **all** of the following to hold simultaneously:
+
+1. `f_diff = |f_current − f_prev| < atol + rtol × max(1, |best_f|)` — spike-immune f-change check
+2. Cauchy-convergence in y-space (base Optimistix check)
+3. Step count ≥ `min_steps`
+4. Not inside the cooldown window after the last constraint release
+
+### CLI equivalents
+
+When using `kmeans-model` or `ptep-model`, the same parameters are available as flags:
+
+```bash
+kmeans-model ... --cooldown 50 --min-steps 200 --verbose
+```
+
 ## Conditioning
 
 Conditioning (preconditioning) transforms the optimization problem to improve convergence. It applies two transformations before optimization:

src/furax_cs/optim/minimize.py

@@ -220,8 +220,20 @@ def minimize(
         Box constraints.
     precondition : bool
         Whether to apply parameter transformation and output scaling.
-    solver_options : dict, optional
-        Additional arguments passed to the solver factory (get_solver).
+    options : dict, optional
+        Extra arguments passed to the solver factory (get_solver).
+        For active-set solvers (``ADABK{N}`` family) the recognised keys are:
+
+        * ``cooldown`` (int, default 20) — steps to suppress termination
+          after a constraint release.
+        * ``min_steps`` (int, default 10) — minimum iterations before
+          termination is considered.
+        * ``verbose_print`` (bool, default False) — print per-step debug
+          info via ``jax.debug.print`` (JIT-compatible).
+        * ``max_linesearch_steps`` (int, default 50) — maximum line-search
+          steps per iteration (active-set and ``optax_lbfgs`` solvers).
+        * ``linesearch`` (str) — linesearch variant for ``optax_lbfgs``
+          (``"zoom"`` or ``"backtracking"``).
     **fn_kwargs
         Additional arguments passed to fn.
 

src/furax_cs/optim/solvers.py

@@ -371,6 +371,18 @@ def get_solver(
         Lower bounds for box projection (optax solvers only).
     upper : PyTree, optional
         Upper bounds for box projection (optax solvers only).
+    verbose_print : bool
+        If True, print per-step termination diagnostics for active-set
+        solvers via ``jax.debug.print`` (JIT-compatible).
+    min_steps : int
+        Minimum iterations before termination is considered
+        (active-set solvers only).
+    cooldown : int
+        Steps to suppress termination after a constraint release
+        (active-set solvers only).
+    max_linesearch_steps : int
+        Maximum line-search steps per iteration (active-set and
+        ``optax_lbfgs`` solvers).
 
     Returns
     -------