BUCToolkit 1.0b2__tar.gz → 1.0.1__tar.gz

buctoolkit-1.0.1/BUCToolkit/Bases/BaseConstraints.py

@@ -0,0 +1,1207 @@
+#  Copyright (c) 2026.4.24, BUCToolkit.
+#  Authors: Pu Pengxin, Song Xin
+#  Version: 1.0.1
+#  File: BaseConstraints.py
+#  Environment: Python 3.12
+import inspect
+import math
+import os
+import warnings
+from typing import Iterable, Dict, Any, List, Literal, Optional, Callable, Sequence, Tuple  # noqa: F401
+
+import torch as th
+from torch import nn
+import numpy as np
+
+from BUCToolkit.utils._print_formatter import FLOAT_ARRAY_FORMAT, SCIENTIFIC_ARRAY_FORMAT
+from BUCToolkit.utils._Element_info import DTYPE
+from BUCToolkit.utils.grad_functions import bhvp
+from BUCToolkit.Bases.BaseMotion import BaseIO
+
+FLOAT_TYPE = os.environ.get('BT_FLOAT_TYPE', 'float32')
+FLOAT_TYPE = DTYPE.get(FLOAT_TYPE, th.float32)
+
+# Probe reference precision for the threshold estimate. Kept local to this
+# module: the map only decides which dtype is used to measure the round-off
+# of the user constraint function.
+_PROBE_DTYPE = {
+    th.float16: th.float32,
+    th.bfloat16: th.float32,
+    th.float32: th.float64,
+    th.float64: th.float64,
+}
+
+
+def _report_diagnostic(message: str, verbose: int, logger: Any) -> None:
+    """
+    Report a constraint diagnostic on the channel selected by verbosity.
+
+    Args:
+        message: Human-readable description of the diagnostic.
+        verbose: Diagnostic verbosity of the constraint component.
+        logger: Logger used when ``verbose > 0``.
+
+    Return: None.
+
+    Note:
+        At ``verbose == 0`` the message is emitted as a ``RuntimeWarning`` so
+        that it still reaches stderr after the standard logger has been
+        redirected to a file; at ``verbose > 0`` it goes to ``logger`` at
+        WARNING level. The default warning filter reports each call site once.
+    """
+    if verbose > 0:
+        logger.warning(message)
+    else:
+        warnings.warn(message, RuntimeWarning, stacklevel=2)
+
+
+class BaseConstr(BaseIO):
+    """
+    Base Constraints used for constraints MD or optimizations.
+    It can be combined into any normal gradient-based algorithms by applying `project1` to gradient or any other tangent space vectors,
+     and `project2` as a retraction mapping to coordinates or any other vectors on the manifold.
+
+     I strongly suggest composition instead of inheritance when BaseConstr is
+     embedded in another class: create the component as ``self._constr`` and
+     call its methods directly (for example ``self._constr._project2(...)``).
+     Do not reproduce a dynamic ``__getattr__`` proxy, which makes the host
+     class opaque to tooling and compilers.
+
+    Args:
+        constr_func: Callable, a tuple of Python functions as the constraint functions s_k(X) that map R^n -> R^k.
+            It takes one or more arguments, one of which must be a Tensor, and returns one Tensor with shape (k, ).
+            `None` for constant function that always return [0., ]. It should support auto-gradient ops. see example below.
+        constr_val: Callable[..., th.Tensor] | th.Tensor, the constraint value of `constr_func`,
+            i.e., constraints are `constr_func(X) = constr_val`.
+            A time-dependent callback may be `constr_val(t)` or
+            `constr_val(t, c0)`, where `t` is a scalar Tensor and `c0` is the
+            value of `constr_func(X_init)` computed at initialization.
+        constr_threshold: float | Literal['auto'], the threshold of constraint convergence (error of manifold violation).
+            ``'auto'`` estimates the achievable floor at initialization and
+            uses three times that value; it requires threshold detection.
+        require_fixman: bool, whether to calculate fixman
+        retraction_solver: Literal['newton', 'damped'], the position projection solver.
+            ``'newton'`` takes the default full Newton step; ``'damped'`` adds
+            a backtracking line search that is slower but can recover from an
+            ill-conditioned normal-space Jacobian. The solver is bound at
+            construction; reassign ``_retraction_map_solver`` to switch later.
+        device: str|torch.device, device that program rum on.
+        verbose: print level. 0 is silent, 1 permits scalar diagnostics only,
+            and 2 or greater also permits array diagnostics.
+            0 for silence, 1 for output Energy and Forces per step, 2 for output all structures.
+        _detect_dynamic_thres: Literal['None', 'Once'] | int, when to re-estimate
+            the achievable threshold floor. ``'Once'`` estimates it only during
+            initialization, ``'None'`` disables the estimate, and an ``int``
+            re-estimates it every that many projection calls. Each estimate
+            evaluates ``constr_func`` a few extra times for probing.
+
+    Examples for constrains:
+        def constr_func(X):
+            y = list()
+            # X: shape(N, D), Note the batch dimension would NOT be considered in constraints calculation of X.
+            # fix the distance between atoms (2, 4), (3, 7), (5, 8) into corresponding `constr_val[:3]`
+            y.append(th.linalg.norm(X[[2, 3, 5]] - X[[4, 7, 8]], dim=-1))
+
+            # fix the angle of atom7-atom5-atom8 and atom11-atom9-atom12 into corresponding `constr_val[3:6]`
+            x1 = X[[5, 9]]
+            x2 = X[[7, 11]]
+            x3 = X[[8, 12]]
+            y.append(
+                (
+                    th.sum((x2 - x1) * (x3 - x1))
+                ) / (th.linalg.norm(x2 - x1) * th.linalg.norm(x3 - x1))
+            )
+            z = th.cat(y)
+            return z
+
+    Methods:
+        run: run BatchMD.
+
+    """
+    def __init__(
+            self,
+            constr_func: Callable | None = None,
+            constr_val: Callable[..., th.Tensor|Tuple[th.Tensor]] | th.Tensor | None = None,
+            constr_threshold: float | Literal['auto'] = 'auto',
+            require_fixman: bool = False,
+            retraction_solver: Literal['newton', 'damped'] = 'newton',
+            device: str | th.device = 'cpu',
+            verbose: int = 0,
+            _detect_dynamic_thres: Literal['None', 'Once'] | int = 'Once',
+    ):
+        # var check
+        if retraction_solver not in ('newton', 'damped'):
+            raise ValueError(
+                f"`retraction_solver` must be 'newton' or 'damped', but got {retraction_solver!r}."
+            )
+        if _detect_dynamic_thres not in ('None', 'Once'):
+            _detect_dynamic_thres = int(_detect_dynamic_thres)
+            if _detect_dynamic_thres <= 0:
+                raise ValueError(f"`int(_detect_dynamic_thres)` must be > 0, but got {_detect_dynamic_thres}")
+        if constr_threshold != 'auto':
+            constr_threshold = float(constr_threshold)
+            if constr_threshold <= 0:
+                raise ValueError(f"`constr_threshold` must be > 0, but got {constr_threshold}")
+        elif _detect_dynamic_thres == 'None':
+            raise ValueError(
+                "`constr_threshold='auto'` requires threshold detection, but "
+                "`_detect_dynamic_thres='None'`."
+            )
+
+        if constr_func is None:
+            constr_func = lambda X: th.tensor(0., dtype=X.dtype, device=X.device)
+        self._lazy_calc_constr_val = False  # use the value of constr_func(X_init) to determine constr_val (and fixed).
+        self.device = th.device(device)
+
+        self.time_now = th.scalar_tensor(0., device=device, dtype=FLOAT_TYPE)
+        self.is_const_constr = False
+        self.constr_val_func_raw = None
+        self._constr_val_func = None
+        self._constr_val_uses_c0 = False
+        self._constr_val_c0 = None
+        self.time_step = 1.
+        self.require_fixman = require_fixman
+        self.retraction_solver = retraction_solver
+
+        if isinstance(constr_val, th.Tensor):
+            self.is_const_constr = True
+            self.constr_val_func_raw = None
+            self.constr_val_now = constr_val.to(self.device, dtype=FLOAT_TYPE)
+        elif callable(constr_val):
+            self.is_const_constr = False
+            self.constr_val_func_raw = constr_val
+            self.constr_val_now = None
+        elif constr_val is None:
+            self.is_const_constr = True
+            self.constr_val_func_raw = None
+            self._lazy_calc_constr_val = True
+        else:
+            raise TypeError(f'Expected `constr_val` is Callable or torch.Tensor, but got {type(constr_val)}')
+
+        if not isinstance(constr_func, Callable):
+            raise TypeError(f"`constr_func` must consist of callable, but got {type(constr_func)}")
+        self.constr_func = constr_func
+        self.Q = th.tensor([], device=self.device, dtype=FLOAT_TYPE)  # Q(n_batch, n_atoms * n_dim, n_constr)
+        self.jac = None # Jacobian Matrix
+        self.R = th.tensor([], device=self.device, dtype=FLOAT_TYPE)  # R(n_batch, n_constr, n_constr)
+        self.q = th.tensor([], device=self.device, dtype=FLOAT_TYPE)  # (n_batch, n_constr, 1), solution of `R^T q = d/dt constr_val(t)`
+        self.sqrtM = None  # M^1/2, (n_batch, n_atoms, n_dim)
+        self.negsqrtM = None  # M^-1/2
+        self.max_proj_iter = 10
+        self.retraction_stagnation_factor = 0.05
+        self.constr_thres = constr_threshold
+        self._is_auto_thres = (constr_threshold == 'auto')
+        self._detect_dynamic_thres = _detect_dynamic_thres
+        self._proj2_call_number: int = 0  # the count of `project2` call. used for auto-detection sheme.
+        self._thres_below = False  # flag indices whether current constr thres lower than the estimated infimum.
+        self.X_cache = None  # constr force, i.e., the mu of Lagrange multipler
+        self._compiled_jac = None  # pre-compiled Jacobian for replay
+
+        self._X_qr_cache = None  # check the consistency of X in project1, ensuring it is the new tangent space projection
+
+        super().__init__()
+        self.init_logger('Main.Constraints')
+        self.verbose = int(verbose)
+        # Bind the concrete retraction implementation once. The solver is
+        # selected at construction; reassign this attribute to switch later.
+        if retraction_solver == 'newton':
+            self._retraction_map_solver = self._retraction_newton
+        else:
+            self._retraction_map_solver = self._retraction_damp_newton
+
+    def initialize(
+            self,
+            func: Any | nn.Module,
+            X: th.Tensor,
+            Element_list: List[List[str]] | List[List[int]] | None,
+            masses: th.Tensor | None,
+            V_init: th.Tensor | None = None,
+            grad_func: Any | nn.Module = None,
+            func_args: Sequence = tuple(),
+            func_kwargs: Dict | None = None,
+            grad_func_args: Sequence = tuple(),
+            grad_func_kwargs: Dict | None = None,
+            is_grad_func_contain_y: bool = True,
+            require_grad: bool = False,
+            batch_indices: List[int] | Tuple[int, ...] | th.Tensor | np.ndarray | None = None,
+            fixed_atom_tensor: Optional[th.Tensor] = None,
+            is_fix_mass_center: bool = False,
+            compile_jacobian: bool = True,
+    ):
+        """Initialize masses, targets, and the tangent-space basis.
+
+        This method owns all state required by constraint projection. Proxy
+        users such as constrained MD and constrained FIRE should call it rather
+        than reproducing the setup or invoking ``_build_compiled_jacobian``
+        directly.
+
+        Args:
+            func: Parent algorithm objective. Accepted for compatibility with
+                the common motion-initialization hook; constraints do not use it.
+            X: Initial coordinates with shape
+                ``(n_batch, n_atom, n_dim)``.
+            Element_list: Atomic symbols or numbers supplied by the parent
+                algorithm. Constraint initialization does not use them.
+            masses: Per-coordinate masses broadcast to the shape of ``X``. If
+                ``None``, unit masses are used.
+            V_init: Initial velocities accepted by the common initialization
+                hook. Constraint initialization does not use them.
+            grad_func: Parent objective-gradient callable. Not used here.
+            func_args: Positional objective arguments. Not used here.
+            func_kwargs: Keyword objective arguments. Not used here.
+            grad_func_args: Positional gradient arguments. Not used here.
+            grad_func_kwargs: Keyword gradient arguments. Not used here.
+            is_grad_func_contain_y: Parent gradient-call convention. Not used
+                by constraint initialization.
+            require_grad: Parent objective autograd flag. Not used here.
+            batch_indices: Irregular-batch atom counts. Not used here because
+                ``constr_func`` is vectorized over the leading dimension of X.
+            fixed_atom_tensor: Parent selective-dynamics mask. Not used here.
+            is_fix_mass_center: Parent center-of-mass setting. Not used here.
+            compile_jacobian: Whether to compile the batched constraint
+                Jacobian. Compilation is enabled by default on both CPU and
+                CUDA. Unsupported constraint graphs automatically fall back to
+                the eager implementation.
+
+        Returns:
+            None. The method updates mass factors, the current constraint
+            target, Jacobian/QR state, compilation state, and projection cache.
+
+        Raises:
+            TypeError: If ``compile_jacobian`` is not ``bool``.
+                Also raised when a time-dependent ``constr_val`` callback has
+                no inspectable signature accepting ``(t)`` or ``(t, c0)``.
+            RuntimeError: If ``constr_val`` and ``constr_func(X)`` have
+                different shapes.
+            ValueError: If the batched constraint residual is not two-dimensional.
+
+        Note:
+            ``X``, ``masses``, and a tensor constraint target are converted to
+            the module-level ``FLOAT_TYPE`` (``BT_FLOAT_TYPE`` environment
+            variable, default float32) before use.
+        """
+        if not isinstance(compile_jacobian, bool):
+            raise TypeError(
+                f'`compile_jacobian` must be bool, but got {type(compile_jacobian)}.'
+            )
+        X = X.to(device=self.device, dtype=FLOAT_TYPE)
+        if masses is not None:
+            masses = masses.to(device=self.device, dtype=FLOAT_TYPE)
+        if masses is None: masses = th.ones_like(X)
+        self.sqrtM = th.sqrt(masses)  # M^1/2, (n_batch, n_atoms, n_dim)
+        self.negsqrtM = 1 / th.sqrt(masses)  # M^-1/2
+        _y_check = th.vmap(self.constr_func)(X)
+        if self._lazy_calc_constr_val:
+            self.constr_val_now = _y_check
+        elif not self.is_const_constr:
+            # Keep the initial target as a detached, immutable run reference;
+            # later time updates must not retain the coordinate autograd graph.
+            self._constr_val_c0 = _y_check.detach().clone()
+            try:
+                constr_val_signature = inspect.signature(self.constr_val_func_raw)
+            except (TypeError, ValueError) as error:
+                raise TypeError(
+                    '`constr_val` must have an inspectable signature accepting '
+                    '(t) or (t, c0).'
+                ) from error
+
+            try:
+                constr_val_signature.bind(self.time_now, self._constr_val_c0)
+            except TypeError:
+                try:
+                    constr_val_signature.bind(self.time_now)
+                except TypeError as unary_error:
+                    raise TypeError(
+                        '`constr_val` must accept either (t) or (t, c0).'
+                    ) from unary_error
+                self._constr_val_uses_c0 = False
+                self._constr_val_func = self.constr_val_func_raw
+            else:
+                self._constr_val_uses_c0 = True
+                self._constr_val_func = lambda t: self.constr_val_func_raw(t, self._constr_val_c0)
+            self._update_constr(self.time_now)
+        if self.verbose > 0:
+            self.logger.info(
+                f'Constraint values are now {np.array2string(self.constr_val_now.squeeze().numpy(force=True), **SCIENTIFIC_ARRAY_FORMAT)}'
+            )
+        # check constr_val shape
+        if _y_check.shape != self.constr_val_now.shape:
+            raise RuntimeError(
+                f'`constr_val` must have the same shape as what constr_func returned {self.constr_val_now.shape}, but got {_y_check.shape}.'
+            )
+        self.constr_val_now = self.constr_val_now.to(self.device, dtype=FLOAT_TYPE)
+        # A previous initialization may have compiled a different run shape or
+        # constraint callable.  Perform validation through the eager path, then
+        # build a fresh graph for the current inputs below.
+        self._compiled_jac = None
+        self._proj2_call_number = 0
+        self._thres_below = False
+        jac, y = self._jacobian(X)
+        if y.ndim != 2:
+            raise ValueError(f'`constr_func` must return a 2D tensor of shape (n_batch, n_constr), but got {y.shape}.')
+        if compile_jacobian:
+            self._build_compiled_jacobian(X)
+        self._do_qr(jac)
+        if self._detect_dynamic_thres != 'None':
+            # Numerical rank deficiency is a one-off initialization diagnostic;
+            # it is not repeated in the projection loop.
+            weighted_jac = th.flatten(jac * self.negsqrtM.unsqueeze(1), -2, -1)
+            spectrum = th.linalg.svdvals(weighted_jac)
+            eps = th.finfo(X.dtype).eps
+            sigma_min = float(spectrum.min())
+            sigma_max = float(spectrum.max())
+            if sigma_max > 0. and sigma_min <= sigma_max * eps ** 0.5:
+                _report_diagnostic(
+                    'The constraint Jacobian is numerically rank-deficient at '
+                    f'the working precision (sigma_min={sigma_min:.3e}, '
+                    f'sigma_max={sigma_max:.3e}, 1/sqrt(eps)={eps ** -0.5:.3e}). '
+                    'The projection may not converge. Please ensure that your constraints have a full-rank Jacobian '
+                    '(almost everywhere). If so, try to set `retraction_solver=\'damped\'`.',
+                    self.verbose,
+                    self.logger,
+                )
+        if self._detect_dynamic_thres == 'Once' or isinstance(self._detect_dynamic_thres, int):
+            self._apply_thres_estimate(self._detect_thres_inf(X, jac), X)
+        self.X_cache = th.empty_like(X)
+
+    def _constr_func_wrapped(self, X, constr_val_now):
+        """
+        Manage the constraint functions.
+        Converting constraint functions into s_k(X) = constr_func(X) - constr_val, thereby constraints are s_k(X) = 0.
+        Returns:
+
+        """
+        y = self.constr_func(X) - constr_val_now  # (n_batch, n_constr, )
+        y = th.atleast_1d(y)
+        return y, y  # repeat the output to use `has_aux` in th.func.jacrev to return the function values
+
+    def _constr_func_for_hvp(self, X):
+        """
+        Manage the constraint functions used for HVP.
+        Converting constraint functions into s_k(X) = constr_func(X) - constr_val, thereby constraints are s_k(X) = 0.
+        Returns:
+
+        """
+        y = self.constr_func(X) - self.constr_val_now  # (n_batch, n_constr, )
+        y = th.atleast_1d(y)
+        return y
+
+    def _constr_val_wrapped(self, t):
+        """
+        Manage the time-dependent constraint values.
+        Args:
+            t:
+
+        Returns:
+
+        """
+        if self.is_const_constr:
+            return None
+        else:
+            y = self._constr_val_func(t)
+            if isinstance(y, (Tuple, List)):
+                y = th.vstack(y).mT
+            else:
+                y = y.view(self._constr_val_c0.shape)
+
+            return y, y
+
+    def _update_constr(self, t: th.Tensor | float):
+        """Update a time-dependent constraint at an explicit evaluation time.
+
+        This method does not mutate ``self.time_now``. Integrators therefore
+        remain responsible for advancing the committed simulation clock and
+        may evaluate constraint stages such as ``time_now + time_step`` or a
+        half-step without changing the parent motion-loop state.
+
+        Args:
+            t: Scalar time at which ``constr_val_func_raw`` and its derivative
+                are evaluated. Python numeric values are converted to the same
+                device and dtype as ``self.time_now``.
+
+        Returns:
+            None. For a time-dependent constraint, ``self.constr_val_now`` and
+            ``self.d_constr`` are replaced by the value and time derivative at
+            ``t``. Constant constraints are unchanged.
+
+        Raises:
+            ValueError: If ``t`` is not scalar.
+        """
+        if not self.is_const_constr:
+            t = th.as_tensor(
+                t,
+                device=self.time_now.device,
+                dtype=self.time_now.dtype,
+            )
+            if t.numel() != 1:
+                raise ValueError(f'Constraint evaluation time must be scalar, but got shape {t.shape}.')
+            t = t.reshape(())
+            self.d_constr, self.constr_val_now = th.func.jacrev(
+                self._constr_val_wrapped,
+                has_aux=True,
+            )(t)  # d constr_val(t) / dt
+            # ensure the returns from udf function keeping the dtype and device.
+            self.d_constr = self.d_constr.to(self.device, dtype=FLOAT_TYPE)
+            self.constr_val_now = self.constr_val_now.to(self.device, dtype=FLOAT_TYPE)
+
+    def _build_compiled_jacobian(self, X_sample: th.Tensor):
+        """
+        Pre-compile the batched constraint Jacobian for a fixed run shape.
+
+        Each structure and its corresponding constraint target are mapped as a
+        pair.  The resulting function returns the same tensors as the eager
+        ``vmap(jacrev(...))`` path: a Jacobian with shape
+        ``(n_batch, n_constr, n_atom, n_dim)`` and residual values with shape
+        ``(n_batch, n_constr)``.  Molecular-dynamics runs keep these dimensions
+        fixed, which permits a static full-graph compilation and avoids graph
+        breaks inside every projection iteration.
+
+        Falls back to eager mode if compilation fails.
+
+        Args:
+            X_sample: Representative coordinates with shape
+                ``(n_batch, n_atom, n_dim)``.  Its shape, device, and dtype are
+                used to compile and warm up the Jacobian graph.
+
+        Returns:
+            None.  On success, ``self._compiled_jac`` stores the compiled
+            callable.  On failure it is reset to ``None`` so ``_jacobian`` uses
+            the eager implementation.
+        """
+        self._compiled_jac = None
+        try:
+            def _wrapped(X_single, constr_val):
+                y = self.constr_func(X_single) - constr_val
+                y = th.atleast_1d(y)
+                return y, y
+
+            jac_fn = th.func.vmap(
+                th.func.jacrev(_wrapped, argnums=0, has_aux=True),
+                # Pair X[i] with constr_val_now[i].  Broadcasting the complete
+                # target tensor here would introduce a second batch dimension.
+                in_dims=(0, 0)
+            )
+            compiled_jac = th.compile(jac_fn, fullgraph=True, dynamic=False)
+
+            # Warm up now so compilation latency and unsupported operations are
+            # handled before entering the MD loop.  Shape validation prevents a
+            # malformed compiled graph from silently replacing the eager path.
+            # Match the gradient mode used by _project2. Without this, a normal
+            # non-Fixman MD run warms the graph with gradients enabled and then
+            # recompiles it on first use inside the no_grad main loop.
+            with th.set_grad_enabled(self.require_fixman):
+                jac, y = compiled_jac(X_sample, self.constr_val_now)
+            expected_jac_shape = (
+                X_sample.shape[0], self.constr_val_now.shape[-1],
+                *X_sample.shape[1:]
+            )
+            expected_y_shape = self.constr_val_now.shape
+            if jac.shape != expected_jac_shape or y.shape != expected_y_shape:
+                raise RuntimeError(
+                    'Compiled constraint Jacobian returned invalid shapes: '
+                    f'expected {expected_jac_shape} and {expected_y_shape}, '
+                    f'but got {jac.shape} and {y.shape}.'
+                )
+            self._compiled_jac = compiled_jac
+
+            if self.verbose:
+                self.logger.info('Constraint Jacobian compiled successfully.')
+        except Exception as e:
+            if self.verbose:
+                self.logger.warning(f'Jacobian compilation failed ({type(e).__name__}): {e}. Falling back to eager.')
+            self._compiled_jac = None
+
+    def _jacobian(self, X: th.Tensor) -> Tuple[th.Tensor, th.Tensor]:
+        """
+        Compute the Jacobian for the current constraint target.
+
+        ``constr_val_now`` and ``d_constr`` must already correspond to the
+        current simulation time. Projection entry points refresh that state
+        before calling this method, which avoids differentiating a
+        time-dependent target in every nonlinear projection subiteration.
+
+        Args:
+            X: Coordinates with shape ``(n_batch, n_atom, n_dim)``.
+
+        Returns:
+            The Jacobian with shape
+            ``(n_batch, n_constr, n_atom, n_dim)`` and residual values with
+            shape ``(n_batch, n_constr)``.
+
+        """
+        if self._compiled_jac is not None:
+            jac, y = self._compiled_jac(X, self.constr_val_now)
+        else:
+            jac, y = th.vmap(th.func.jacrev(self._constr_func_wrapped, has_aux=True))(X, self.constr_val_now)
+        self.jac = jac
+        return jac, y
+
+    def _do_qr(self, jac: th.Tensor):
+        """
+        Do QR factorization and save/update self.Q.
+        Args:
+            jac: Jacobian matrix of s_k at X.
+
+        Returns: None
+
+        """
+        # weighted jac: J_w = J @ M^-1/2,
+        # (n_batch, n_constr, n_atoms, n_dim) * (n_batch, 1, n_atoms, n_dim) -flatten-> (n_batch, n_constr, n_atoms * n_dim)
+        _jac = th.flatten(jac * self.negsqrtM.unsqueeze(1), -2, -1)
+        _jacT = _jac.mT.contiguous()
+        self.Q, self.R = th.linalg.qr(_jacT)  # Q(n_batch, n_atoms * n_dim, n_constr) , R(n_batch, n_constr, n_constr)
+
+    def _detect_thres_inf(self, X: th.Tensor, jac: th.Tensor | None = None) -> float | None:
+        """
+        Estimate the smallest achievable constraint residual for this dtype.
+
+        The estimate combines the measured round-off of the user constraint
+        function with the coordinate round-off amplified by the normal-space
+        spectrum. The function is probed at a few deterministic ULP-scale
+        perturbations in the working dtype and, when possible, in the next
+        precision selected by the local ``_PROBE_DTYPE`` map.
+
+        Args:
+            X: Reference coordinates with shape ``(n_batch, n_atom, n_dim)``.
+            jac: Constraint Jacobian at ``X``. Recomputed when omitted.
+
+        Returns:
+            The estimated residual floor, or ``None`` when the constraint
+            function cannot be probed in the reference precision.
+        """
+        if jac is None:
+            jac, _ = self._jacobian(X)
+        work_dtype = X.dtype
+        reference_dtype = _PROBE_DTYPE.get(work_dtype, work_dtype)
+        with th.no_grad():
+            weighted_jac = th.flatten(jac * self.negsqrtM.unsqueeze(1), -2, -1)
+            sigma_max = float(th.linalg.svdvals(weighted_jac).max())
+            eps = th.finfo(work_dtype).eps
+            coordinate_scale = max(1., float(X.abs().max()))  # bound of input X, clamp to 1 to avoid too small values.
+            evaluation_floor = 0.
+            if reference_dtype != work_dtype:
+                try:
+                    rng = th.Generator(device=X.device).manual_seed(114514)  # the hard-code seed is only for reproduction
+                    noise = th.randn(X.shape, generator=rng, device=X.device, dtype=X.dtype)
+                    # 8 ULP of the largest coordinate; clamped to avoid underflow
+                    step = 8. * eps * X.abs().clamp_min(coordinate_scale)
+                    probes = (X, X + step * noise, X - step * noise)
+                    values_work = th.stack([th.vmap(self.constr_func)(p) for p in probes])
+                    values_reference = th.stack([
+                        th.vmap(self.constr_func)(p.to(reference_dtype)) for p in probes
+                    ])  # (3, batch, n_constr)
+
+                    if values_reference.dtype != reference_dtype:
+                        return None
+                    difference = (values_work.to(reference_dtype) - values_reference).abs()
+                    if not bool(th.isfinite(difference).all()):
+                        return None
+                    evaluation_floor = float(difference.max())
+                except Exception:
+                    return None
+            return float(
+                2. * evaluation_floor
+                + 2. * sigma_max * eps * coordinate_scale
+            )  # count both the inherent error from constr_function and the error from Jacobian
+
+    def _apply_thres_estimate(self, floor: float|None, X: th.Tensor) -> None:
+        """
+        Apply or warn about an estimated constraint-residual floor.
+
+        Args:
+            floor: Estimate from ``_detect_thres_inf``; ``None`` when it failed.
+            X: Coordinates used for the estimate.
+
+        Returns:
+            None. An ``'auto'`` threshold is replaced by three times the floor;
+            an explicit threshold is never changed and is only warned about
+            when it lies below half of the estimated floor.
+
+        Raises:
+            RuntimeError: If the threshold is ``'auto'`` and the estimate failed.
+        """
+        if floor is None:
+            if self._is_auto_thres:
+                raise RuntimeError(
+                    'Unable to estimate the constraint threshold floor for the '
+                    'given constraint function and dtype; set an explicit '
+                    '`constr_threshold` or make the constraint function '
+                    'evaluable in the probe precision.'
+                )
+            _report_diagnostic(
+                'Unable to estimate the constraint threshold floor for the '
+                'given constraint function; the user-supplied '
+                '`constr_threshold` is used unchanged.',
+                self.verbose,
+                self.logger,
+            )
+            return
+        if self._is_auto_thres:
+            self.constr_thres = max(3. * floor, th.finfo(X.dtype).eps)
+            return
+        is_below = float(self.constr_thres) < 0.5 * floor
+        if is_below and not self._thres_below:
+            _report_diagnostic(
+                f'`constr_threshold`={self.constr_thres} is below half of the '
+                f'estimated achievable floor ({floor:.3e} for dtype '
+                f'{X.dtype}); convergence may stagnate.',
+                self.verbose,
+                self.logger,
+            )
+        self._thres_below = is_below
+
+    def _project1(self, V:th.Tensor, X:th.Tensor | None = None, out: th.Tensor | None = None) -> th.Tensor:
+        """
+        Project X into the manifold defined by all constrains by QR factorization.
+        Usually used for update velocity V.
+        for Jacobian matrix J of constraints s_k at X, let J^T M^-1/2 = QR, thus s_k(X + v * dt) = s_k(X) + R^T Q^T M^1/2 v * dt + O(dt^2)
+        define the projector: P = M^-1/2 Q_2 Q_2^T M^1/2 = M^-1/2 (I - Q_1 Q_1^T) M^1/2, where Q = [Q_1, Q_2], Q_1 is from the financial QRF.
+
+        Args:
+            V: the input tensor of shape (n_batch, n_atom, n_dim). It might be coordinates, velocity, and higher deviations.
+            X: the point that calculate the projector. If None, stored Q R (maybe last step) will be used.
+            out: the output tensor. If None, return a new tensor.
+
+        Returns: th.Tensor, the projected X.
+
+        """
+        n_batch, n_atoms, n_dim = V.shape
+        sqrtM = self.sqrtM  # M^1/2, (n_batch, n_atoms, n_dim)
+        negsqrtM = self.negsqrtM  # M^-1/2
+        # Check consistency. For time-dependent constraints, callers must have
+        # explicitly selected the target stage through ``_update_constr(t)`` or
+        # a preceding ``_project2`` call.
+        if X is not None:
+            jac, y = self._jacobian(X)
+            self._do_qr(jac)
+        #self._X_qr_cache = X.clone()
+
+        # P = M^-1/2 @ (I - Q Q^T) @ M^1/2 = I - M^-1/2 @ Q @ Q^T @ M^1/2
+        Q = self.Q  # Q(n_batch, n_atoms * n_dim, n_constr)
+        Px = (sqrtM * V).reshape(n_batch, n_atoms * n_dim, 1).contiguous()  # (n_batch, n_atoms * n_dim, 1)
+        Px = (Q @ (Q.mT.contiguous() @ Px)).reshape(n_batch, n_atoms, n_dim)  # (n_batch, n_atoms, n_dim)
+        Px.mul_(negsqrtM)
+        Px = th.sub(V, Px, out=out)
+        # manifold veloc. correction
+        if not self.is_const_constr:  # time-dependent constraints
+            th.linalg.solve_triangular(self.R.mT.contiguous(), self.d_constr.unsqueeze(-1), upper=False, out=self.q)
+            Px.add_(self.negsqrtM * (self.Q @ self.q).reshape(n_batch, n_atoms, n_dim))
+
+        return Px
+
+    def _project_norm(self, V:th.Tensor, X:th.Tensor | None = None) -> th.Tensor:
+        """
+        Project V to normal space.
+        for Jacobian matrix J of constraints s_k at X, let J^T M^-1/2 = QR, thus s_k(X + v * dt) = s_k(X) + R^T Q^T M^1/2 v * dt + O(dt^2)
+        define the projector: P = M^-1/2 Q_2 Q_2^T M^1/2 = M^-1/2 (I - Q_1 Q_1^T) M^1/2, where Q = [Q_1, Q_2], Q_1 is from the financial QRF.
+        Args:
+            V: the input tensor of shape (n_batch, n_atom, n_dim). It might be coordinates, velocity, and higher deviations.
+
+        Returns: th.Tensor, the projected X.
+
+        """
+        n_batch, n_atoms, n_dim = V.shape
+        sqrtM = self.sqrtM  # M^1/2, (n_batch, n_atoms, n_dim)
+        negsqrtM = self.negsqrtM  # M^-1/2
+        # The normal-space basis depends on the coordinate Jacobian, not on the
+        # target value. Do not alter the caller-selected constraint time here.
+        if X is not None:
+            jac, y = self._jacobian(X)
+            self._do_qr(jac)
+
+        # P = M^-1/2 @ (I - Q Q^T) @ M^1/2 = I - M^-1/2 @ Q @ Q^T @ M^1/2
+        Q = self.Q  # Q(n_batch, n_atoms * n_dim, n_constr)
+        Px = (sqrtM * V).reshape(n_batch, n_atoms * n_dim, 1).contiguous()  # (n_batch, n_atoms * n_dim, 1)
+        Px = (Q @ (Q.mT.contiguous() @ Px)).reshape(n_batch, n_atoms, n_dim)  # (n_batch, n_atoms, n_dim)
+        Px.mul_(negsqrtM)
+
+        return Px
+
+    @staticmethod
+    def _retraction_newton(
+            X: th.Tensor,
+            X_work: th.Tensor,
+            sqrtM_Q: th.Tensor,
+            R: th.Tensor,
+            aux_multiplier: th.Tensor,
+            jacobian_func: Callable[[th.Tensor], Tuple[th.Tensor, th.Tensor]],
+            initial_constr_err: th.Tensor,
+            constr_threshold: float,
+            stagnation_threshold: float,
+            max_proj_iter: int,
+            time_step: float,
+            verbose: int,
+            logger: Any,
+    ) -> Tuple[th.Tensor, th.Tensor, th.Tensor]:
+        """Solve the retraction in the fixed initial normal space.
+
+        Args:
+            X: Coordinates before retraction. Updated in-place with the final
+                working coordinates when the solver exits.
+            X_work: Coordinate work buffer with the same shape as ``X``.
+            sqrtM_Q: Fixed mass-weighted normal basis.
+            R: Fixed upper-triangular QR factor for force recovery.
+            aux_multiplier: Accumulated multiplier updated in-place.
+            jacobian_func: Callable returning the current Jacobian and
+                residual.
+            initial_constr_err: Initial maximum constraint residual.
+            constr_threshold: Absolute constraint convergence threshold.
+            stagnation_threshold: Threshold for the existing early-stop
+                criterion.
+            max_proj_iter: Maximum number of solver iterations.
+            time_step: Integration time step used for force conversion.
+            verbose: Constraint diagnostic verbosity.
+            logger: Logger used for diagnostics.
+
+        Returns:
+            ``Fc``, the constraint force recovered once from the final
+            multiplier; ``dX``, the final coordinate displacement; and a
+            scalar Boolean tensor that is true only when the constraint
+            residual reached ``constr_threshold``.
+        """
+        n_batch, n_atoms, n_dim = X.shape
+        dX = th.zeros_like(X)
+        X_work.copy_(X)
+        constr_err_old = initial_constr_err
+
+        def _final_calc_Fc():
+            """
+            Re-solve the constraints force by triangular solve.
+            """
+            fc: th.Tensor = th.linalg.solve_triangular(
+                R,
+                aux_multiplier.unsqueeze(-1),
+                upper=True,
+            ).squeeze(-1)
+            fc *= 207.28617 / (time_step ** 2)
+            X.copy_(X_work)
+            return fc
+
+        for i in range(max_proj_iter):
+            jac, y = jacobian_func(X_work)
+            constr_err = th.max(th.abs(y))
+            if verbose > 1:
+                logger.info(f'{i: <3d} Constraint errors are now: {constr_err:.4e}')
+
+            is_converged = constr_err <= constr_threshold
+            is_fixed = th.abs(constr_err - constr_err_old) < stagnation_threshold
+            if is_converged or is_fixed:
+                if is_fixed and not is_converged:
+                    _report_diagnostic(
+                        'Constraint errors are not converged yet, but it has '
+                        'been fixed. This loop will be skipped.',
+                        verbose,
+                        logger,
+                    )
+                Fc = _final_calc_Fc()
+                return Fc, dX, is_converged
+            constr_err_old = constr_err
+
+            aux = th.bmm(jac.flatten(-2, -1), sqrtM_Q)
+            d_multiplier, _ = th.linalg.solve_ex(aux, -y)
+            aux_multiplier.add_(d_multiplier)
+            dX = th.einsum(
+                'bnc, bc -> bn',
+                sqrtM_Q, aux_multiplier
+            ).reshape(n_batch, n_atoms, n_dim)
+            th.add(X, dX, out=X_work)
+
+        _report_diagnostic(
+            "Projection of X to the manifold is not converged.", verbose, logger,
+        )
+        is_converged = th.zeros((), device=X.device, dtype=th.bool)
+        Fc = _final_calc_Fc()
+        return Fc, dX, is_converged
+
+    @staticmethod
+    def _retraction_damp_newton(
+            X: th.Tensor,
+            X_work: th.Tensor,
+            sqrtM_Q: th.Tensor,
+            R: th.Tensor,
+            aux_multiplier: th.Tensor,
+            jacobian_func: Callable[[th.Tensor], Tuple[th.Tensor, th.Tensor]],
+            initial_constr_err: th.Tensor,
+            constr_threshold: float,
+            stagnation_threshold: float,
+            max_proj_iter: int,
+            time_step: float,
+            verbose: int,
+            logger: Any,
+            max_backtrack: int = 8,
+            step_shrink: float = 0.5,
+    ) -> Tuple[th.Tensor, th.Tensor, th.Tensor]:
+        """Solve the retraction with a safeguarded Newton step.
+
+        The Newton direction is always formed in the fixed initial normal
+        space. A trial multiplier is accepted only when it lowers the maximum
+        absolute constraint residual; otherwise the same direction is reduced
+        geometrically. The accumulated multiplier is changed only after an
+        accepted trial, so the force and velocity correction retain their
+        existing physical meaning.
+
+        Args:
+            X: Coordinates before retraction. Updated in-place on exit.
+            X_work: Coordinate work buffer with the same shape as ``X``.
+            sqrtM_Q: Fixed mass-weighted normal basis.
+            R: Fixed upper-triangular QR factor for force recovery.
+            aux_multiplier: Accumulated multiplier updated in-place.
+            jacobian_func: Callable returning the current Jacobian and residual.
+            initial_constr_err: Initial maximum constraint residual.
+            constr_threshold: Absolute constraint convergence threshold.
+            stagnation_threshold: Threshold for the existing early-stop test.
+            max_proj_iter: Maximum number of solver iterations.
+            time_step: Integration time step used for force conversion.
+            verbose: Constraint diagnostic verbosity.
+            logger: Logger used for diagnostics.
+            max_backtrack: Maximum number of step reductions per iteration.
+            step_shrink: Multiplicative reduction applied to a rejected step.
+
+        Returns:
+            ``Fc``, ``dX``, and a scalar Boolean convergence status.
+        """
+        n_batch, n_atoms, n_dim = X.shape
+        dX = th.zeros_like(X)
+        X_work.copy_(X)
+        constr_err_old = initial_constr_err
+        is_converged = initial_constr_err <= constr_threshold
+
+        def _final_calc_Fc():
+            """
+            Re-solve the constraints force by triangular solve.
+            """
+            fc: th.Tensor = th.linalg.solve_triangular(
+                R,
+                aux_multiplier.unsqueeze(-1),
+                upper=True,
+            ).squeeze(-1)
+            fc *= 207.28617 / (time_step ** 2)
+            X.copy_(X_work)
+            return fc
+
+        for i in range(max_proj_iter):
+            jac, y = jacobian_func(X_work)
+            constr_err = th.max(th.abs(y))
+            if verbose > 1:
+                logger.info(f'{i: <3d} Constraint errors are now: {constr_err:.4e}')
+
+            aux = th.bmm(jac.flatten(-2, -1), sqrtM_Q)
+            aux_q, aux_r = th.linalg.qr(aux.mT)
+            aux_diag = th.abs(th.diagonal(aux_r, dim1=-2, dim2=-1))
+            aux_scale = th.max(aux_diag)
+            is_rank_deficient = aux_scale <= 0
+            if not is_rank_deficient:
+                is_rank_deficient = th.min(aux_diag) <= 1e-6 * aux_scale
+
+            is_converged = constr_err <= constr_threshold
+            is_fixed = th.abs(constr_err - constr_err_old) < stagnation_threshold
+            if is_fixed or is_converged:
+                if is_fixed and (not is_converged):
+                    if is_rank_deficient:
+                        _report_diagnostic(
+                            'Constraint errors are not converged yet, but the '
+                            'current normal-space Jacobian is degenerate. The '
+                            'remaining iterations will be skipped.',
+                            verbose,
+                            logger,
+                        )
+                    else:
+                        _report_diagnostic(
+                            'Constraint errors are not converged yet, but the '
+                            'residual improvement is below the stagnation '
+                            'threshold. The remaining iterations will be skipped.',
+                            verbose,
+                            logger,
+                        )
+                Fc = _final_calc_Fc()
+                return Fc, dX, is_converged
+            constr_err_old = constr_err
+
+            d_multiplier = th.linalg.solve_triangular(
+                aux_r.mT.contiguous(),
+                (-y).unsqueeze(-1),
+                upper=False,
+            )
+            d_multiplier = th.bmm(aux_q, d_multiplier).squeeze(-1)
+            accepted = False
+            step = 1.
+            candidate_multiplier = th.empty_like(aux_multiplier)
+            candidate_dX = th.empty_like(dX)
+            for _ in range(max_backtrack + 1):
+                candidate_multiplier.copy_(aux_multiplier).add_(d_multiplier, alpha=step)
+                candidate_dX = th.einsum(
+                    'bnc, bc -> bn',
+                    sqrtM_Q, candidate_multiplier
+                ).reshape(n_batch, n_atoms, n_dim)
+                th.add(X, candidate_dX, out=X_work)
+                _, y_trial = jacobian_func(X_work)
+                trial_err = th.max(th.abs(y_trial))
+                if trial_err < constr_err:
+                    accepted = True
+                    break
+                th.add(X, dX, out=X_work)
+                step *= step_shrink
+
+            if not accepted:
+                _report_diagnostic(
+                    'Damped Newton retraction step did not lower the constraint error.',
+                    verbose,
+                    logger,
+                )
+                Fc = _final_calc_Fc()
+                return Fc, dX, is_converged
+
+            aux_multiplier.copy_(candidate_multiplier)
+            dX.copy_(candidate_dX)
+
+        _report_diagnostic(
+            "Projection of X to the manifold is not converged.", verbose, logger,
+        )
+        Fc = _final_calc_Fc()
+        return Fc, dX, is_converged
+
+    def _project2(
+            self,
+            X: th.Tensor,
+            X_orig: th.Tensor | None = None,
+            V: th.Tensor | None = None,
+            constr_time: th.Tensor | float | None = None,
+    ) -> Tuple[th.Tensor, th.Tensor|None, th.Tensor|None]:
+        """
+        Continuously project the Jacobian of all constrains to the exact manifold.
+        ** Update X and V (if given) in-place.
+        If V is provided, it is updated in-place with the constraint force contribution
+        to velocity, using the converged multiplier mu to avoid subtraction cancellation.
+
+        Args:
+            X: (n_batch, n_atom, n_dim) coordinates to be projected. Modified in-place.
+            X_orig: reference coordinates for fixed Jacobian basis. Defaults to X.
+            V: (n_batch, n_atom, n_dim) velocities. If given, updated in-place by
+               dv = M^{-1/2} Q mu / time_step, where mu is the converged multiplier.
+            constr_time: Time at which the position constraint is enforced. If
+                ``None``, use ``self.time_now + self.time_step``, corresponding
+                to the end of the current outer MD step. Explicit values allow
+                integrators with other splitting schemes to select a half-step
+                or another stage without modifying the parent MD clock.
+
+        Returns:
+            Fc: constraint forces (eV/atom)
+            G: geometric Fixman force term, or ``None`` when
+                ``require_fixman=False``.
+            w: Fixman/Blue-Moon reweighting factor
+                ``1 / sqrt(det(J M^-1 J^T))`` for each batch sample, or
+                ``None`` when ``require_fixman=False``.
+        """
+        n_batch, n_atoms, n_dim = X.shape
+        self._proj2_call_number += 1
+        n_constr = self.R.shape[-1]
+        aux_multipler = th.zeros(n_batch, n_constr, device=self.device, dtype=FLOAT_TYPE)
+        if (X_orig is None) or (X_orig is X):
+            # A detached reference is required: the retraction rewrites ``X``
+            # in place at the end, which is illegal on a leaf that requires
+            # grad for the Fixman derivative.
+            X_orig = X.detach()
+        time_step = self.time_step
+
+        # The target value and its time derivative are constant throughout one
+        # nonlinear position projection. Refresh them once at the end-of-step
+        # stage by default, before the fixed-basis Jacobian and all subsequent
+        # Newton subiterations. The following velocity projection reuses this
+        # same staged state.
+        if constr_time is None:
+            constr_time = self.time_now + time_step
+        self._update_constr(constr_time)
+
+        with th.set_grad_enabled(self.require_fixman):
+            X_orig.requires_grad_(self.require_fixman)
+            jac, y = self._jacobian(X_orig)
+            check_freq = self._detect_dynamic_thres
+            if isinstance(check_freq, int):
+                if self._proj2_call_number % check_freq == 0:
+                    self._apply_thres_estimate(self._detect_thres_inf(X_orig, jac), X_orig)
+
+            constr_thres = self.constr_thres
+            self._do_qr(jac)
+            Q, R = self.Q, self.R
+            # Safeguard against degenerate constraint directions
+            # (e.g., erf soft-CN with vanishing gradients):
+            #   mask y for convergence — degenerate dirs ignored
+            #   clamp R diagonals — prevents solve_triangular blow-up
+            # Stop only when the residual improvement is small compared with
+            # the requested tolerance.  Five percent preserves
+            # the cheap stagnation escape without cutting off useful Newton
+            # progress too early.
+            degenerate_eps = self.retraction_stagnation_factor * constr_thres
+            #degenerate_mask = th.abs(R.diagonal(dim1=-2, dim2=-1)) < degenerate_eps  # (n_batch, n_constr)
+            #R.diagonal(dim1=-2, dim2=-1).clamp_min_(degenerate_eps)
+            sqrtM_Q = self.negsqrtM.reshape(n_batch, -1, 1) * self.Q  # (b, an, c)
+            if self.require_fixman:
+                # (n_batch, n_constr, n_constr) -> (n_batch, n_constr)
+                abs_diag_R = th.diagonal(R, dim1=-2, dim2=-1).abs()
+                log_sqrt_det_Z = abs_diag_R.log_().sum(dim=-1)  # (n_batch), "sqrt" is due to Z=R^T R, R is triangular
+                half_logdet_Z = log_sqrt_det_Z.sum() # scalar
+                # A linear constraint has a coordinate-independent Jacobian,
+                # hence a constant Fixman determinant and an exactly zero
+                # derivative. In that case ``half_logdet_Z`` contains no
+                # autograd graph and calling ``autograd.grad`` would raise even
+                # though the required result is well-defined. ``allow_unused``
+                # also covers CV implementations whose Jacobian graph is
+                # present but independent of some/all coordinate inputs.
+                if half_logdet_Z.requires_grad:
+                    half_dlnz_dx = th.autograd.grad(
+                        half_logdet_Z,
+                        X_orig,
+                        allow_unused=True,
+                    )[0]
+                    if half_dlnz_dx is None:
+                        half_dlnz_dx = th.zeros_like(X_orig)
+                else:
+                    half_dlnz_dx = th.zeros_like(X_orig)
+                half_dlnz_dx = half_dlnz_dx.reshape(n_batch, -1)
+                RG = th.einsum('bnc, bn -> bc', sqrtM_Q, half_dlnz_dx)
+                G = th.linalg.solve_triangular(R, RG.unsqueeze(-1), upper=True).squeeze(-1)
+                w = log_sqrt_det_Z.neg_().exp_()
+            else:
+                G = None
+                w = None
+        R.detach_()
+        sqrtM_Q.detach_()
+
+        Fc, dX, _is_converged = self._retraction_map_solver(
+            X=X,
+            X_work=self.X_cache,
+            sqrtM_Q=sqrtM_Q,
+            R=R,
+            aux_multiplier=aux_multipler,
+            jacobian_func=self._jacobian,
+            initial_constr_err=th.max(th.abs(y)),
+            constr_threshold=constr_thres,
+            stagnation_threshold=degenerate_eps,
+            max_proj_iter=self.max_proj_iter,
+            time_step=time_step,
+            verbose=self.verbose,
+            logger=self.logger,
+        )
+
+        if V is not None:
+            V.add_(dX.div_(time_step))
+
+        return Fc, G, w
+
+    def _jacobian_derivative(self, X, v: th.Tensor) -> th.Tensor:
+        """
+        Compute `dJ/dt = v^T H` by auto-differentiation.
+        Args:
+            X:
+            v:
+
+        Returns: v^T H with shape (n_batch, n_constr, n_atom, n_dim)
+
+        """
+        vH: th.Tensor = bhvp(self._constr_func_for_hvp, (X,), (v,), False)  # (n_batch, n_constr, n_atom, n_dim)
+        return vH.contiguous()
+
+    def christoffel_solver(
+            self,
+            X: th.Tensor,
+            V: th.Tensor,
+            F_tan: th.Tensor | None = None,
+            max_iter: int = 20,
+            tol: float = 1e-6,
+            omega: float = 0.7
+    ):
+        r"""
+        Implicit midpoint integration of the constrained geodesic equation:
+
+            dV/dt = -Gamma(V, V) + M^{-1} F_tan
+
+        where Gamma = M^{-1/2} Q R^{-T} (V^T H) V  is the Christoffel symbol
+        capturing the curvature of the constraint manifold, and F_tan is the
+        tangent-space projection of the physical forces.
+
+        Uses the Woodbury identity to reduce the Newton-step matrix inverse
+        from O(N_free^3) to O(k^3), where k << N_free is the number of constraints.
+
+        Args:
+            X:     current positions  (n_batch, n_atom, n_dim)
+            V:     current velocities (n_batch, n_atom, n_dim)
+            F_tan: tangent-projected forces, or None for pure-geodesic flow
+            max_iter: maximum Newton iterations
+            tol:      convergence tolerance on the residual norm
+            omega:    SOR relaxation factor (< 1 damps, =1 pure Newton)
+
+        Returns:
+            X_new, V_new  – updated in-place.
+        """
+        n_batch, n_atom, n_dim = X.shape
+        n_free = n_atom * n_dim
+        dt = self.time_step
+
+        # Precompute Woodbury blocks (independent of X_new, V_new)
+        B = self.negsqrtM.reshape(n_batch, -1, 1) * self.Q  # (B, n_free, k)
+        B_T = B.mT.contiguous()                               # (B, k, n_free)
+        eye_k = th.eye(self.R.shape[-1], device=X.device, dtype=X.dtype).expand(n_batch, -1, -1)
+
+        V_old = V.clone()
+        X_old = X.clone()
+        # Initial guess: explicit Euler half-step
+        X_new = X.add(V_old, alpha=dt)
+        V_new = V.clone()
+        if F_tan is not None:
+            # Include force: V_new = V_old + dt * acc, acc = M^{-1} F_tan (F_tan is already force)
+            acc = F_tan * self.negsqrtM * self.negsqrtM  # M^{-1} = negsqrtM^2
+            V_new.add_(acc, alpha=dt)
+
+        for it in range(max_iter):
+            X_mid = (X_new + X_old) * 0.5
+            V_mid = (V_new + V_old) * 0.5
+
+            # Compute Christoffel symbol: Gamma = B @ R^{-T} @ (J' V_mid)
+            dJ = self._jacobian_derivative(X_mid, V_mid)       # (B, k, N, D)
+            dJ_flat = dJ.flatten(-2, -1)                        # (B, k, n_free)
+            Jp_V = th.einsum('bkn, bn -> bk', dJ_flat, V_mid.flatten(-2, -1))  # J'·V_mid  (B, k)
+            _sol = th.linalg.solve_triangular(
+                self.R.mT.contiguous(), Jp_V.unsqueeze(-1), upper=False
+            )                                                    # R^{-T} J'·V  (B, k, 1)
+            Gamma = (B @ _sol).reshape(n_batch, n_atom, n_dim)  # Christoffel (B, N, D)
+
+            # Residuals for the implicit midpoint system
+            # R_X = X_new - X_old - dt * V_mid
+            R_X = X_new - X_old - dt * V_mid
+            # R_V = V_new - V_old - dt * (acc - Gamma)
+            if F_tan is not None:
+                R_V = V_new - V_old - dt * (acc - Gamma)
+            else:
+                R_V = V_new - V_old + dt * Gamma
+
+            # Convergence check
+            norm_RX = th.linalg.norm(R_X.reshape(n_batch, -1), dim=-1)
+            norm_RV = th.linalg.norm(R_V.reshape(n_batch, -1), dim=-1)
+            if th.all((norm_RX < tol) & (norm_RV < tol)):
+                V.copy_(V_new)
+                X.copy_(X_new)
+                return X_new, V_new
+
+            # --- Newton update with Woodbury identity ---
+            # dJ_flat^T = dJ.mT, shape (B, n_free, k)
+            dJ_T = dJ_flat.mT.contiguous()                     # (B, n_free, k)
+            dJ_T_B = th.bmm(dJ_T, B)                           # (B, k, k)
+            M_wood = eye_k + dt * dJ_T_B                       # I + dt * dJ^T B
+            rv_flat = R_V.reshape(n_batch, -1, 1)               # (B, n_free, 1)
+
+            dJ_T_rv = th.bmm(dJ_T, rv_flat)                    # dJ^T @ R_V  (B, k, 1)
+            sol_W = th.linalg.solve(M_wood, dJ_T_rv)           # (I + dt dJ^T B)^{-1} dJ^T R_V
+            P_RV = rv_flat - dt * th.bmm(B, sol_W)              # P @ R_V  (B, n_free, 1)
+
+            dV = P_RV.reshape(n_batch, n_atom, n_dim)
+            dX = (R_X + 0.5 * dt * dV)
+
+            # SOR relaxation
+            V_new = V_new - dV
+            V_new = omega * V_new + (1.0 - omega) * V
+            X_new = X_new - dX
+
+        # Did not converge – return best effort
+        V.copy_(V_new)
+        X.copy_(X_new)
+        return X_new, V_new