RGI-Toolkit: Differentiable Restraints for Controllable Biomolecular Structure Prediction

Abstract

Diffusion-based models have substantially advanced biomolecular structure prediction. Within this framework, restraint-guided inference (RGI) corrects the denoiser output at each reverse-diffusion step using a differentiable restraint loss, allowing available experimental information, stereochemical requirements, or a specified target conformational state to be incorporated into predictions without retraining the model. Here, we present RGI-Toolkit, a general-purpose Python library that makes RGI available for a range of structure predictors. The toolkit consolidates restraint definition, atom selection, loss evaluation, and coordinate optimization into a shared, backend-agnostic engine, limiting predictor-specific code to a thin preprocessing adapter and a single hook in the denoising loop. Restraints on distances, angles, and dihedrals, on deviations from reference structures, and on user-defined reaction coordinates, together with ligand conformer restraints, are all specified in a common configuration format. The toolkit currently supports six predictors, including AlphaFold3. Across all six predictors, conformer restraints eliminated chirality and cis/trans errors in the predicted ligands and improved their local geometry, while geometric and custom restraints systematically shifted the sampled protein conformations toward the specified target states. RGI-Toolkit thus provides a unified framework for integrating experimental information and prior geometric knowledge into structure prediction and for generating biomolecular structures with specified conformational and stereochemical properties.