cycling

cycling

Parameter cycling and sensitivity-based selection for Q2MM optimization.

Implements the upstream Q2MM grad-simp optimization loop, adapted for our scipy-based optimizer architecture. The key insight from Norrby & Liljefors (1998) and Quinn et al. (2022) is that the Nelder-Mead simplex converges well on ≤40 parameters but fails on larger sets; the upstream code therefore selected only the 2-4 most sensitive parameters for each simplex pass.

Sensitivity is measured by central differentiation of all parameters:

• d1 = (f(x+h) - f(x-h)) / 2 (1st derivative, unnormalised)
• d2 = f(x+h) + f(x-h) - 2·f(x) (2nd derivative, unnormalised)
• simp_var = d2 / d1² (upstream selection metric)

Low simp_var identifies parameters where the objective is steep but shallow-bottomed — gradient methods struggle here, but simplex can make progress. The upstream authors noted this criterion is imperfect (simplex.py:414: "Sorting based upon the 2nd derivative isn't such a good criterion"); we therefore also support |d1| (absolute first derivative) as an alternative.

References

• Norrby, P.-O.; Liljefors, T. J. Comput. Chem. 1998, 19, 1146-1166.
• Hansen, E.C. "Development and Applications of Q2MM" (PhD dissertation, University of Notre Dame, 2016).
• Quinn, T.R. et al. PLOS ONE 2022, 17, e0264960.
• Upstream code: github.com/nsf-c-cas/q2mm-2 (simplex.py, opt.py).

SensitivityResult dataclass

SensitivityResult(d1: ndarray, d2: ndarray, simp_var: ndarray, ranking: ndarray, metric: str, n_evals: int)

Result of parameter sensitivity analysis via central differentiation.

Attributes:

Name	Type	Description
d1	ndarray	First derivative (unnormalised) for each parameter.
d2	ndarray	Second derivative (unnormalised) for each parameter.
simp_var	ndarray	Upstream "simplex variable": d2 / d1**2 for each parameter.
ranking	ndarray	Parameter indices sorted by sensitivity (most sensitive first).
metric	str	Which metric was used for ranking.
n_evals	int	Number of objective function evaluations performed.

LoopResult dataclass

LoopResult(success: bool, initial_score: float, final_score: float, n_cycles: int, n_eval: int = 0, cycle_scores: list[float] = list(), selected_indices: list[list[int]] = list(), sensitivity_results: list[SensitivityResult] = list(), message: str = '')

Result of an :class:OptimizationLoop run.

Attributes:

Name	Type	Description
success	bool	True if converged before hitting max_cycles.
initial_score	float	Objective value before any optimisation.
final_score	float	Objective value after the last cycle.
n_cycles	int	Number of grad-simp cycles completed.
n_eval	int	Total objective function evaluations across all cycles (grad + sensitivity + simp).
cycle_scores	list[float]	Objective value at the end of each cycle.
selected_indices	list[list[int]]	Parameter indices selected for each simplex pass.
sensitivity_results	list[SensitivityResult]	Full sensitivity analysis for each cycle.
message	str	Human-readable summary.

improvement property

improvement: float

Fractional improvement: (initial - final) / initial.

Returns:

Name	Type	Description
float	float	Fractional improvement, or 0.0 if initial_score is zero.

summary

summary() -> str

Human-readable summary string.

Returns:

Name	Type	Description
str	str	Multi-line summary of the loop result.

Source code in q2mm/optimizers/cycling.py

def summary(self) -> str:
    """Human-readable summary string.

    Returns:
        str: Multi-line summary of the loop result.

    """
    lines = [
        f"OptimizationLoop: {'converged' if self.success else 'max cycles reached'}",
        f"  Cycles:      {self.n_cycles}",
        f"  Evaluations: {self.n_eval}",
        f"  Score:       {self.initial_score:.6f} → {self.final_score:.6f}",
        f"  Improvement: {self.improvement:.2%}",
    ]
    if self.message:
        lines.append(f"  Message:     {self.message}")
    return "\n".join(lines)

SubspaceObjective

SubspaceObjective(objective: ObjectiveFunction, active_indices: list[int] | ndarray, full_vector: ndarray)

Wraps an :class:ObjectiveFunction to optimise a parameter subset.

The wrapper accepts a sub-vector of length len(active_indices) and maps it into the full parameter vector before delegating to the underlying objective. This lets :func:scipy.optimize.minimize run Nelder-Mead (or any other method) on just the selected parameters while the rest stay fixed.

Parameters:

Name	Type	Description	Default
objective	ObjectiveFunction	The full objective function.	required
active_indices	list[int] | ndarray	Indices into the full parameter vector that are active.	required
full_vector	ndarray	The current full parameter vector (inactive params are taken from this snapshot).	required

Initialize the subspace objective wrapper.

Parameters:

Name	Type	Description	Default
objective	ObjectiveFunction	The full objective function.	required
active_indices	list[int] | ndarray	Indices into the full parameter vector that are active.	required
full_vector	ndarray	The current full parameter vector.	required

Raises:

Type	Description
ValueError	If active_indices is empty.

Source code in q2mm/optimizers/cycling.py

def __init__(
    self,
    objective: ObjectiveFunction,
    active_indices: list[int] | np.ndarray,
    full_vector: np.ndarray,
) -> None:
    """Initialize the subspace objective wrapper.

    Args:
        objective (ObjectiveFunction): The full objective function.
        active_indices (list[int] | np.ndarray): Indices into the
            full parameter vector that are active.
        full_vector (np.ndarray): The current full parameter vector.

    Raises:
        ValueError: If ``active_indices`` is empty.

    """
    self.objective = objective
    self.active_indices = np.asarray(active_indices, dtype=int)
    self._base_vector = full_vector.copy()
    if len(self.active_indices) == 0:
        raise ValueError("active_indices must not be empty")

build_full_vector

build_full_vector(sub_vector: ndarray) -> ndarray

Map a sub-vector back into the full parameter vector.

Parameters:

Name	Type	Description	Default
sub_vector	ndarray	Values for the active parameters.	required

Returns:

Type	Description
ndarray	np.ndarray: Full parameter vector with active slots replaced.

Source code in q2mm/optimizers/cycling.py

def build_full_vector(self, sub_vector: np.ndarray) -> np.ndarray:
    """Map a sub-vector back into the full parameter vector.

    Args:
        sub_vector (np.ndarray): Values for the active parameters.

    Returns:
        np.ndarray: Full parameter vector with active slots replaced.

    """
    full = self._base_vector.copy()
    full[self.active_indices] = sub_vector
    return full

__call__

__call__(sub_vector: ndarray) -> float

Evaluate the objective on sub_vector.

Parameters:

Name	Type	Description	Default
sub_vector	ndarray	Values for the active parameters.	required

Returns:

Name	Type	Description
float	float	Objective function score.

Source code in q2mm/optimizers/cycling.py

def __call__(self, sub_vector: np.ndarray) -> float:
    """Evaluate the objective on *sub_vector*.

    Args:
        sub_vector (np.ndarray): Values for the active parameters.

    Returns:
        float: Objective function score.

    """
    return float(self.objective(self.build_full_vector(sub_vector)))

residuals

residuals(sub_vector: ndarray) -> ndarray

Return the residual vector (for least_squares).

Parameters:

Name	Type	Description	Default
sub_vector	ndarray	Values for the active parameters.	required

Returns:

Type	Description
ndarray	np.ndarray: Weighted residual vector.

Source code in q2mm/optimizers/cycling.py

def residuals(self, sub_vector: np.ndarray) -> np.ndarray:
    """Return the residual vector (for ``least_squares``).

    Args:
        sub_vector (np.ndarray): Values for the active parameters.

    Returns:
        np.ndarray: Weighted residual vector.

    """
    return self.objective.residuals(self.build_full_vector(sub_vector))

get_initial_vector

get_initial_vector() -> ndarray

Return the sub-vector corresponding to current active parameters.

Returns:

Type	Description
ndarray	np.ndarray: Copy of active parameter values from the base vector.

Source code in q2mm/optimizers/cycling.py

def get_initial_vector(self) -> np.ndarray:
    """Return the sub-vector corresponding to current active parameters.

    Returns:
        np.ndarray: Copy of active parameter values from the base
            vector.

    """
    return self._base_vector[self.active_indices].copy()

get_bounds

get_bounds() -> list[tuple[float, float]]

Bounds for the active parameters only.

Returns:

Type	Description
list[tuple[float, float]]	list[tuple[float, float]]: Lower/upper bound pairs for each active parameter.

Source code in q2mm/optimizers/cycling.py

def get_bounds(self) -> list[tuple[float, float]]:
    """Bounds for the active parameters only.

    Returns:
        list[tuple[float, float]]: Lower/upper bound pairs for each
            active parameter.

    """
    all_bounds = self.objective.forcefield.get_bounds()
    return [all_bounds[i] for i in self.active_indices]

OptimizationLoop

OptimizationLoop(objective: ObjectiveFunction, *, max_params: int = 3, convergence: float = 0.01, max_cycles: int = 10, full_method: str = 'L-BFGS-B', simp_method: str = 'Nelder-Mead', full_maxiter: int = 200, simp_maxiter: int = 200, sensitivity_metric: Literal['simp_var', 'abs_d1'] = 'simp_var', eps: float = 0.001, full_jac: str | None = None, verbose: bool = True)

grad-simp cycling loop inspired by the upstream Q2MM workflow.

Each cycle

1. Full-space pass — run full_method (default L-BFGS-B) on all parameters.
2. Sensitivity analysis — central differentiation to rank params.
3. Subspace simplex — run simp_method (default Nelder-Mead) on only the top max_params most sensitive parameters.
4. Convergence check — stop if the fractional improvement in the objective falls below convergence.

Parameters:

Name	Type	Description	Default
objective	ObjectiveFunction	The objective function to minimise.	required
max_params	int	Number of parameters per simplex pass (upstream default: 3).	3
convergence	float	Stop when (score_before - score_after) / score_before < convergence.	0.01
max_cycles	int	Maximum number of grad-simp cycles.	10
full_method	str	Scipy method for the full-space pass.	'L-BFGS-B'
simp_method	str	Scipy method for the subspace pass.	'Nelder-Mead'
full_maxiter	int	Max iterations for the full-space pass.	200
simp_maxiter	int	Max iterations for the subspace pass.	200
sensitivity_metric	str	How to rank parameters for selection — "simp_var" or "abs_d1".	'simp_var'
eps	float	Finite-difference step size for the full-space optimizer.	0.001
verbose	bool	Whether to log progress.	True

References

Upstream loop.py:Loop.opt_loop() and simplex.py:Simplex.run().

Initialize the optimization loop.

Parameters:

Name	Type	Description	Default
objective	ObjectiveFunction	The objective function to minimise.	required
max_params	int	Number of parameters per simplex pass.	3
convergence	float	Fractional improvement threshold.	0.01
max_cycles	int	Maximum number of grad-simp cycles.	10
full_method	str	Scipy method for the full-space pass.	'L-BFGS-B'
simp_method	str	Scipy method for the subspace pass.	'Nelder-Mead'
full_maxiter	int	Max iterations for the full-space pass.	200
simp_maxiter	int	Max iterations for the subspace pass.	200
sensitivity_metric	str	Parameter ranking criterion.	'simp_var'
eps	float	Finite-difference step size.	0.001
full_jac	str | None	Jacobian strategy for the full-space pass. None uses scipy finite differences; 'analytical' uses :meth:ObjectiveFunction.gradient.	None
verbose	bool	Whether to log progress.	True

Source code in q2mm/optimizers/cycling.py

def __init__(
    self,
    objective: ObjectiveFunction,
    *,
    max_params: int = 3,
    convergence: float = 0.01,
    max_cycles: int = 10,
    full_method: str = "L-BFGS-B",
    simp_method: str = "Nelder-Mead",
    full_maxiter: int = 200,
    simp_maxiter: int = 200,
    sensitivity_metric: Literal["simp_var", "abs_d1"] = "simp_var",
    eps: float = 1e-3,
    full_jac: str | None = None,
    verbose: bool = True,
) -> None:
    """Initialize the optimization loop.

    Args:
        objective (ObjectiveFunction): The objective function to
            minimise.
        max_params (int): Number of parameters per simplex pass.
        convergence (float): Fractional improvement threshold.
        max_cycles (int): Maximum number of grad-simp cycles.
        full_method (str): Scipy method for the full-space pass.
        simp_method (str): Scipy method for the subspace pass.
        full_maxiter (int): Max iterations for the full-space pass.
        simp_maxiter (int): Max iterations for the subspace pass.
        sensitivity_metric (str): Parameter ranking criterion.
        eps (float): Finite-difference step size.
        full_jac (str | None): Jacobian strategy for the full-space
            pass.  ``None`` uses scipy finite differences;
            ``'analytical'`` uses :meth:`ObjectiveFunction.gradient`.
        verbose (bool): Whether to log progress.

    """
    self.objective = objective
    self.max_params = max_params
    self.convergence = convergence
    self.max_cycles = max_cycles
    self.full_method = full_method
    self.simp_method = simp_method
    self.full_maxiter = full_maxiter
    self.simp_maxiter = simp_maxiter
    self.sensitivity_metric = sensitivity_metric
    self.eps = eps
    self.full_jac = full_jac
    self.verbose = verbose

run

run() -> LoopResult

Execute the grad-simp cycling loop.

Returns:

Name	Type	Description
LoopResult	LoopResult	Contains convergence status, per-cycle scores, and selected parameter indices.

Source code in q2mm/optimizers/cycling.py

def run(self) -> LoopResult:
    """Execute the grad-simp cycling loop.

    Returns:
        LoopResult: Contains convergence status, per-cycle scores,
            and selected parameter indices.

    """
    from q2mm.optimizers.scipy_opt import ScipyOptimizer

    ff = self.objective.forcefield
    x0 = ff.get_param_vector().copy()

    # Enable penalty fallback for eigendecomposition failures during
    # optimisation.  This lets the optimizer retreat from pathological
    # parameter regions instead of crashing the entire run.
    prev_on_error = self.objective.on_error
    self.objective.on_error = "penalty"

    initial_score = float(self.objective(x0))

    n_eval_start = self.objective.n_eval  # track cumulative evals
    cycle_scores: list[float] = [initial_score]
    selected_indices: list[list[int]] = []
    sensitivity_results: list[SensitivityResult] = []
    converged = False

    # Detect optax: prefix for full-space optimizer
    use_optax = self.full_method.startswith("optax:")
    use_jaxopt = self.full_method.startswith("jaxopt:")
    use_basinhopping = self.full_method.startswith("basinhopping")
    use_multistart = self.full_method.startswith("multi:")

    try:
        if self.verbose:
            logger.info(
                "OptimizationLoop: initial score = %.6f, max_params = %d",
                initial_score,
                self.max_params,
            )

        for cycle in range(1, self.max_cycles + 1):
            score_before = cycle_scores[-1]

            # --- Step 1: Full-space optimisation ---
            if use_optax:
                from q2mm.optimizers.optax import OptaxOptimizer

                optax_spec = self.full_method.split(":", 1)[1]
                if "+" in optax_spec:
                    optax_name, schedule = optax_spec.split("+", 1)
                else:
                    optax_name = optax_spec
                    schedule = None
                full_opt = OptaxOptimizer(
                    optimizer=optax_name,
                    max_steps=self.full_maxiter,
                    schedule=schedule,
                    verbose=False,
                )
                full_result = full_opt.optimize(self.objective)
            elif use_jaxopt:
                from q2mm.optimizers.jaxopt_opt import JaxOptOptimizer

                jaxopt_method = self.full_method.split(":", 1)[1]
                full_opt = JaxOptOptimizer(
                    method=jaxopt_method,
                    maxiter=self.full_maxiter,
                    verbose=False,
                )
                full_result = full_opt.optimize(self.objective)
            elif use_basinhopping:
                from q2mm.optimizers.basinhopping import BasinHoppingOptimizer

                if ":" in self.full_method:
                    bh_spec = self.full_method.split(":", 1)[1].strip() or "L-BFGS-B"
                else:
                    bh_spec = "L-BFGS-B"
                full_opt = BasinHoppingOptimizer(
                    local_method=bh_spec,
                    local_maxiter=self.full_maxiter,
                    jac=self.full_jac,
                    verbose=False,
                )
                full_result = full_opt.optimize(self.objective)
            elif use_multistart:
                from q2mm.optimizers.multistart import MultiStartOptimizer

                ms_spec = self.full_method.split(":", 1)[1]
                inner_opt = ScipyOptimizer(
                    method=ms_spec,
                    maxiter=self.full_maxiter,
                    eps=self.eps,
                    jac=self.full_jac,
                    verbose=False,
                )
                full_opt = MultiStartOptimizer(
                    optimizer=inner_opt,
                    n_starts=5,
                    verbose=False,
                )
                full_result = full_opt.optimize(self.objective)
            else:
                full_opt = ScipyOptimizer(
                    method=self.full_method,
                    maxiter=self.full_maxiter,
                    eps=self.eps,
                    jac=self.full_jac,
                    verbose=False,
                )
                full_result = full_opt.optimize(self.objective)
            score_after_grad = full_result.final_score

            if self.verbose:
                logger.info(
                    "  Cycle %d GRAD (%s): %.6f → %.6f",
                    cycle,
                    self.full_method,
                    score_before,
                    score_after_grad,
                )

            # --- Step 2: Sensitivity analysis ---
            step_sizes = ff.get_step_sizes()
            bounds = ff.get_bounds()
            sens = compute_sensitivity(
                self.objective,
                step_sizes=step_sizes,
                metric=self.sensitivity_metric,
                bounds=bounds,
            )
            sensitivity_results.append(sens)

            # Select top max_params; ensure we have at least one active parameter
            n_active = min(self.max_params, ff.n_params)
            if n_active < 1:
                raise ValueError(
                    f"OptimizationLoop requires at least one active parameter, "
                    f"but max_params={self.max_params} and ff.n_params={ff.n_params}."
                )
            active = sens.ranking[:n_active].tolist()
            selected_indices.append(active)

            if self.verbose:
                labels = ff.get_param_type_labels()
                selected_labels = [f"{labels[i]}[{i}]" for i in active]
                logger.info(
                    "  Cycle %d sensitivity (%s): selected %s",
                    cycle,
                    self.sensitivity_metric,
                    ", ".join(selected_labels),
                )

            # --- Step 3: Subspace simplex ---
            current_full = ff.get_param_vector().copy()
            sub_obj = SubspaceObjective(self.objective, active, current_full)

            from scipy import optimize as sp_opt

            sub_x0 = sub_obj.get_initial_vector()
            sub_bounds = sub_obj.get_bounds()

            scipy_options: dict = {"maxiter": self.simp_maxiter}
            if self.simp_method == "Nelder-Mead":
                scipy_options["xatol"] = 1e-6
                scipy_options["fatol"] = 1e-8

            # Pass bounds when the method supports them
            bounded_methods = {"L-BFGS-B", "trust-constr", "SLSQP"}
            use_bounds = self.simp_method in bounded_methods

            scipy_result = sp_opt.minimize(
                sub_obj,
                sub_x0,
                method=self.simp_method,
                bounds=sub_bounds if use_bounds else None,
                options=scipy_options,
            )

            # Only accept the simplex result if it actually improved the score
            best_sub = scipy_result.x
            best_full = sub_obj.build_full_vector(best_sub)
            score_after_simp = float(self.objective(best_full))

            if score_after_simp < score_after_grad:
                ff.set_param_vector(best_full)
            else:
                # Simplex didn't improve — FF already has post-gradient
                # parameters (set by ScipyOptimizer.optimize), no restore
                # needed since objective evaluations are non-mutating.
                score_after_simp = score_after_grad
                if self.verbose:
                    logger.info(
                        "  Cycle %d SIMP: no improvement (%.6f ≥ %.6f), keeping GRAD result",
                        cycle,
                        score_after_simp,
                        score_after_grad,
                    )

            if self.verbose:
                logger.info(
                    "  Cycle %d SIMP (%s, %d params): %.6f → %.6f",
                    cycle,
                    self.simp_method,
                    n_active,
                    score_after_grad,
                    score_after_simp,
                )

            cycle_scores.append(score_after_simp)

            # --- Step 4: Convergence check ---
            if score_before > 0:
                change = (score_before - score_after_simp) / score_before
            else:
                change = 0.0

            if self.verbose:
                logger.info(
                    "  Cycle %d: %.2f%% improvement (threshold: %.2f%%)",
                    cycle,
                    change * 100,
                    self.convergence * 100,
                )

            if 0 <= change < self.convergence:
                converged = True
                if self.verbose:
                    logger.info("  Converged after %d cycles.", cycle)
                break

    finally:
        self.objective.on_error = prev_on_error

    final_score = cycle_scores[-1]
    n_cycles = len(cycle_scores) - 1  # exclude initial
    total_evals = self.objective.n_eval - n_eval_start

    return LoopResult(
        success=converged,
        initial_score=initial_score,
        final_score=final_score,
        n_cycles=n_cycles,
        n_eval=total_evals,
        cycle_scores=cycle_scores,
        selected_indices=selected_indices,
        sensitivity_results=sensitivity_results,
        message="converged" if converged else f"max cycles ({self.max_cycles}) reached",
    )

compute_sensitivity

compute_sensitivity(objective: ObjectiveFunction, step_sizes: ndarray | None = None, metric: Literal['simp_var', 'abs_d1'] = 'simp_var', bounds: list[tuple[float, float]] | None = None) -> SensitivityResult

Central differentiation to rank parameter sensitivity.

For each parameter i, evaluates f(x + h_i e_i) and f(x - h_i e_i) and computes:

• d1_i = (f_fwd - f_bwd) / 2
• d2_i = f_fwd + f_bwd - 2·f_0
• simp_var_i = d2_i / d1_i² (upstream criterion)

The ranking is by ascending simp_var (lowest = most suitable for simplex) or descending |d1| (largest gradient = most sensitive), depending on metric.

When bounds are provided, step sizes are shrunk so that x0 ± h_eff stays within [lower, upper] for each parameter. Parameters where no symmetric step is possible (i.e. x0 is at a bound) are skipped and assigned inf simp_var / zero d1.

When the engine supports :meth:~MMEngine.batched_energy and all references are energy-only, all required 2K + 1 evaluations (for the K parameters with nonzero step sizes, K ≤ N) are vectorised into a single call where possible (e.g. jax.vmap on GPU).

Cost: 2K + 1 objective evaluations (1 baseline + 2 per active parameter for central differentiation; in the worst case K = N, giving 2N + 1).

Parameters:

Name	Type	Description	Default
objective	ObjectiveFunction	Must already be evaluable (engine and molecules configured).	required
step_sizes	ndarray | None	Per-parameter step sizes. Defaults to :meth:ForceField.get_step_sizes if not provided.	None
metric	str	Ranking criterion — "simp_var" or "abs_d1".	'simp_var'
bounds	list[tuple[float, float]] | None	Per-parameter (lower, upper) bounds. When provided, step sizes are shrunk to the largest symmetric step that stays in bounds.	None

Returns:

Name	Type	Description
SensitivityResult	SensitivityResult	Derivatives, rankings, and evaluation count.

Raises:

Type	Description
ValueError	If step_sizes length does not match the parameter vector, or if metric is unknown.

Source code in q2mm/optimizers/cycling.py

def compute_sensitivity(
    objective: ObjectiveFunction,
    step_sizes: np.ndarray | None = None,
    metric: Literal["simp_var", "abs_d1"] = "simp_var",
    bounds: list[tuple[float, float]] | None = None,
) -> SensitivityResult:
    """Central differentiation to rank parameter sensitivity.

    For each parameter *i*, evaluates ``f(x + h_i e_i)`` and
    ``f(x - h_i e_i)`` and computes:

    - ``d1_i = (f_fwd - f_bwd) / 2``
    - ``d2_i = f_fwd + f_bwd - 2·f_0``
    - ``simp_var_i = d2_i / d1_i²``   (upstream criterion)

    The ranking is by *ascending* ``simp_var`` (lowest = most suitable
    for simplex) or *descending* ``|d1|`` (largest gradient = most
    sensitive), depending on *metric*.

    When *bounds* are provided, step sizes are shrunk so that
    ``x0 ± h_eff`` stays within ``[lower, upper]`` for each parameter.
    Parameters where no symmetric step is possible (i.e. ``x0`` is at a
    bound) are skipped and assigned ``inf`` simp_var / zero d1.

    When the engine supports :meth:`~MMEngine.batched_energy` and all
    references are energy-only, all required ``2K + 1`` evaluations
    (for the ``K`` parameters with nonzero step sizes, ``K ≤ N``) are
    vectorised into a single call where possible (e.g. ``jax.vmap`` on GPU).

    Cost: ``2K + 1`` objective evaluations (1 baseline + 2 per *active*
    parameter for central differentiation; in the worst case ``K = N``,
    giving ``2N + 1``).

    Args:
        objective (ObjectiveFunction): Must already be evaluable (engine
            and molecules configured).
        step_sizes (np.ndarray | None): Per-parameter step sizes.
            Defaults to :meth:`ForceField.get_step_sizes` if not
            provided.
        metric (str): Ranking criterion — ``"simp_var"`` or
            ``"abs_d1"``.
        bounds (list[tuple[float, float]] | None): Per-parameter
            ``(lower, upper)`` bounds.  When provided, step sizes are
            shrunk to the largest symmetric step that stays in bounds.

    Returns:
        SensitivityResult: Derivatives, rankings, and evaluation count.

    Raises:
        ValueError: If ``step_sizes`` length does not match the parameter
            vector, or if *metric* is unknown.

    """
    ff = objective.forcefield
    x0 = ff.get_param_vector().copy()
    n = len(x0)

    if step_sizes is None:
        step_sizes = ff.get_step_sizes()
    step_sizes = np.asarray(step_sizes, dtype=float)
    if len(step_sizes) != n:
        raise ValueError(f"step_sizes length {len(step_sizes)} != param vector length {n}")

    # Compute effective step sizes: shrink to stay within bounds
    eff_steps = step_sizes.copy()
    if bounds is not None:
        if len(bounds) != n:
            raise ValueError(f"bounds length {len(bounds)} != param vector length {n}")
        lower = np.array([b[0] for b in bounds])
        upper = np.array([b[1] for b in bounds])
        if np.any(lower > upper):
            raise ValueError("bounds must satisfy lower <= upper for every parameter")
        # Largest symmetric step that keeps x0 ± h within [lower, upper]
        max_up = upper - x0
        max_down = x0 - lower
        max_symmetric = np.minimum(max_up, max_down)
        # Shrink step to fit; zero out if no room
        eff_steps = np.where(
            max_symmetric > 0,
            np.minimum(np.abs(eff_steps), max_symmetric),
            0.0,
        )
        n_shrunk = int(np.sum((eff_steps < np.abs(step_sizes)) & (eff_steps > 0)))
        n_skipped = int(np.sum((step_sizes != 0) & (eff_steps == 0)))
        if n_shrunk > 0 or n_skipped > 0:
            logger.info(
                "compute_sensitivity: %d params shrunk to bounds, %d skipped (at bounds)",
                n_shrunk,
                n_skipped,
            )

    # Build the full (2K+1) parameter matrix: [x0, x0+h0, x0-h0, x0+h1, x0-h1, ...]
    param_rows = [x0]
    active_mask = eff_steps != 0
    for i in range(n):
        if not active_mask[i]:
            continue
        x_fwd = x0.copy()
        x_bwd = x0.copy()
        x_fwd[i] += eff_steps[i]
        x_bwd[i] -= eff_steps[i]
        param_rows.append(x_fwd)
        param_rows.append(x_bwd)

    param_matrix = np.array(param_rows)  # (2K+1, n) where K = active params
    n_evals = len(param_matrix)

    # Evaluate all parameter vectors — batched when possible
    all_scores = objective.batched_scores(param_matrix)

    # Unpack: first row is baseline, then pairs of (fwd, bwd)
    f0 = all_scores[0]
    d1 = np.zeros(n)
    d2 = np.zeros(n)
    pair_idx = 1
    for i in range(n):
        if not active_mask[i]:
            continue
        f_fwd = all_scores[pair_idx]
        f_bwd = all_scores[pair_idx + 1]
        pair_idx += 2
        d1[i] = (f_fwd - f_bwd) * 0.5
        d2[i] = f_fwd + f_bwd - 2.0 * f0

    # Compute simp_var, guarding against zero d1
    with np.errstate(divide="ignore", invalid="ignore"):
        simp_var = np.where(d1 != 0, d2 / (d1**2), np.inf)

    # Ranking
    if metric == "simp_var":
        ranking = np.argsort(simp_var)  # ascending: lowest = most suitable
    elif metric == "abs_d1":
        # Normalise by effective step size so ranking reflects true gradient
        # magnitude rather than being biased by per-type step size differences.
        normalised_d1 = np.where(eff_steps != 0, d1 / eff_steps, 0.0)
        ranking = np.argsort(-np.abs(normalised_d1))  # descending: largest = most sensitive
    else:
        raise ValueError(f"Unknown metric: {metric!r}")

    return SensitivityResult(
        d1=d1,
        d2=d2,
        simp_var=simp_var,
        ranking=ranking,
        metric=metric,
        n_evals=n_evals,
    )