AmberTorchPB: A

Unified Framework for Poisson–Boltzmann-Based Reaction Field Energy Calculation via Tensor Computation

Authors: Yongxian Wu, Qiankang Wang, Robin Jiang, Ray Luo

Abstract

Electrostatic interactions are pivotal to understanding biomolecular structure and function, with the Poisson–Boltzmann (PB) equation serving as a cornerstone for modeling these phenomena in ionic solutions; however, the application of PB solvers to large-scale macromolecular assemblies is currently impeded by significant computational bottlenecks and a fragmented software ecosystem rooted in legacy architectures, which collectively struggle to exploit the capabilities of modern high-performance computing (HPC). While traditional methods grapple with scalability and hardware adaptation, tensor abstraction utilized in contemporary deep learning has emerged as a transformative paradigm for efficiently managing hardware heterogeneity, memory optimization, and mixed-precision arithmetic. Capitalizing on this advancement, we introduce AmberTorchPB, a unified, extensible, and accelerator-aware framework built upon LibTorch designed to modernize biomolecular electrostatics. By abstracting low-level data management, AmberTorchPB enables a single algorithmic implementation to seamlessly support diverse sparse matrix layouts, numerical precisions, and computing devices. We demonstrate the framework’s versatility by implementing and benchmarking a suite of iterative solvers, thereby providing a robust C++ backend that facilitates rapid prototyping, rigorous benchmarking, and the deployment of high-fidelity electrostatic simulations on heterogeneous architectures.

Introduction

Electrostatic interactions are fundamental to understanding the structure, dynamics, and function of molecules across a wide range of scientific disciplines, including but not limited to computational biochemistry and biophysics. These interactions are critical not only for the stability and formation of biomolecular structures but also for regulating essential biological processes such as protein–ligand binding, , enzyme catalysis, −

and DNA transcription. −

A key tool for modeling these interactions is the Poisson–Boltzmann (PB) equation, which describes the electrostatic potential in and around charged biomolecules immersed in an ionic solution. , The accurate and efficient solution of the PB equation is vital for elucidating the underlying electrostatic mechanisms that govern biomolecular behavior. For instance, in drug discovery, precise electrostatic calculations enable the reliable prediction of binding affinity and specificity, which in turn impacts the efficacy and safety of therapeutic candidates. ,

However, solving the PB equation for large and complex biomolecular systems remains computationally demanding, primarily due to the intricate geometries involved and the necessity of accurately modeling the mobile ionic environment. To address this, numerous numerical techniques leveraging finite difference −

and finite element methods , have been proposed to solve PB equations with high fidelity. −

In addition, dedicated treatments of charge singularities, including removal and regularization strategies for smooth solute–solvent interfaces, have further improved PB robustness and accuracy. , Among these approaches, the finite difference method is frequently employed due to its relative simplicity and has been integrated into widely used solvers such as AMBER PBSA, ,,−

Delphi, , APBS, , MIBPB, , and CHARMM PBEQ. Within these computational frameworks, various iterative solvers have been adopted for biomolecular electrostatics, including successive over-relaxation (SOR), the generalized minimal residual method (GMRES), and the biconjugate gradient stabilized method (BiCGSTAB), alongside standard conjugate gradient techniques. Despite substantial progress in enhancing numerical stability and accuracy across a range of PB solvers, − ,−

the inherent computational complexity of the PB equation remains a significant bottleneck. This challenge is particularly acute when scaling to expansive biomolecular systems comprising tens of thousands to millions of atoms. As the system size increases, the requirements for fine spatial resolution and the rigorous treatment of complex dielectric interfaces impose prohibitive memory and runtime demands, rendering high-fidelity electrostatic analysis computationally expensive for large-scale macromolecular assemblies.

Compounding this computational burden is the fragmentation of the current software ecosystem. Historically, PB solvers have been developed independently over decades, resulting in disparate implementations often rooted in legacy codebases. This siloed development creates significant barriers to maintaining a generalizable framework capable of adapting to modern high-performance computing (HPC) paradigms. Consequently, adapting these distinct, legacy architectures to leverage advanced computational techniques, such as massive multithreading and hardware acceleration via CUDA, proves to be an arduous engineering challenge, limiting the community’s ability to fully exploit modern hardware resources.

Furthermore, the absence of a unified framework precludes rigorous benchmarking under the contemporary computational conditions. Without a standardized platform, it is difficult to systematically evaluate how different solvers behave when subjected to modern hardware constraints, particularly regarding their scalability on heterogeneous accelerators and multicore architectures. This lack of comparative insight hinders the identification of optimal solving strategies for specific hardware configurations.

Recent advances in deep learning, which inherently rely on intensive floating-point operations, offer a promising paradigm for improving computational efficiency in biomolecular systems. , To better leverage the growing availability of heterogeneous high-performance hardware, such as graphics processing units (GPUs) and tensor processing units (TPUs), the tensor abstraction has been proposed to manage device placement, flexible floating-point operators, mixed-precision arithmetic, and batch-wise operations. , This abstraction layer enables researchers to focus on algorithm development by abstracting away the complexities of data management and low-level operations, thereby facilitating the deployment of computation-intensive and high-throughput models, such as large language models.

Motivated by these limitations regarding computational scalability, software fragmentation, and the lack of modern benchmarking infrastructure, we introduce AmberTorchPB, a framework built upon the tensor abstraction provided by LibTorch, a robust tensor computation library. This unified, extensible, and accelerator-aware framework is designed to overcome the rigidity of existing approaches, providing a modern testbed for high-performance biomolecular electrostatics. In this study, we deliberately focus on finite difference discretizations of the linear PB equation, whereas finite element formulations, nonlinear PB models, and ion-size-modified PB variants are beyond the current scope and are reserved for future extensions within the same tensor-based framework.

Our method introduces several key properties and innovations embodied in AmberTorchPB. First, by leveraging the tensor abstraction in LibTorch, we propose a unified PB solver framework that supports multiple sparse matrix layouts, numerical precision types, and hardware acceleration strategies, all while requiring only a single implementation of each algorithm. In addition to explicit sparse CSR/COO tensors, AmberTorchPB incorporates a matrix-free stencil layout in which the linear system is evaluated directly through tensor slicing operations, thereby avoiding explicit materialization of the full matrix parameter and improving both memory efficiency and computational performance on CPU and GPU architectures. Second, AmberTorchPB is designed for extensibility; to demonstrate this flexibility, we have implemented a diverse set of solvers within the framework, including CG, BiCG, SOR, GMRES, and red-black successive over-relaxation (RB-SOR), each utilizing the same unified infrastructure. Third, the framework provides a consistent benchmarking environment for systematically examining solver behavior under varying conditions such as different numerical precisions and accelerator configurations. Finally, we provide a general C++ backend library that defines a unified interface for PB solvers, enabling external software to adopt AmberTorchPB as a computational backend. By offering this unified abstraction layer, our framework allows researchers to rapidly prototype new solvers and deploy novel algorithms without the burden of adapting implementations across heterogeneous accelerator platforms.

Method

In this section, we first provide a concise review of the essential background on PB free energy calculations, followed by a detailed description of the proposed AmberTorchPB method.

PB Electrostatic Solvation Free Energy

To characterize electrostatic potentials in biomolecular systems, the PB model is widely employed. In this formulation, the electrostatic potential is defined over two the interior solute domain, Ωint, and the exterior solvent domain, Ωext, where dissolved ions are modeled by using a Boltzmann distribution. Within the solute region, fixed atomic charges {q ~ i ~ | i = N ~ a ~} are located at their corresponding atomic coordinates {r ~ i ~ | i = N ~ a ~}, where N ~ a ~ denotes the total number of atoms. These two domains are separated by a dielectric interface Γ, typically represented by the solvent-excluded or molecular surface. An illustration of the PB model is provided in Figure .

Under this setup, the electrostatic potential ϕ­(r) satisfies the nonlinear PB equation 1∇·(ϵ(r)∇ϕ(r))=−ρ(r)−λ(r)∑iniqiexp[−qiϕ(r)kT]where ϵ­(r) is the spatially varying dielectric constant, ρ­(r) is the fixed charge density, and the second term on the right-hand side represents the mobile ion charge density in the solvent. Here, λ­(r) denotes the Stern-layer masking function, n ~ i ~ and q ~ i ~ are the number density and charge of ion species i, respectively, k is the Boltzmann constant, and T is the temperature. The dielectric constant is defined as2ϵ(r)={ϵ1,r∈Ωint,ϵ2,r∈Ωext.

When the electrostatic potential is small relative to the thermal energy (i.e., |q ~ i ~ϕ­(r)/kT| ≪ 1), the nonlinear PB equation can be approximated by its linearized form ,,

3−∇·(ϵ(r)∇ϕ(r))+κ2(r)ϕ(r)=ρ(r)where κ2(r)=λ(r)∑iniqi2kT . The PB electrostatic solvation free energy, ΔG, is subsequently evaluated as4ΔG=12∑i=1Naqiϕrxn(ri)where ϕrxn(r ~ i ~) denotes the reaction field potential, which can be obtained using a singularity-free reformulation of eq . ,

Finite Difference Method for Solving the PB

Equation

Finite difference methods (FDMs) are widely employed to discretize and solve the PB equation, as presented in eq , over a structured grid. −

This computational domain is characterized as a three-dimensional Cartesian mesh, where each point is defined by the indices (i, j, and k) corresponding to the spatial coordinates (x ~ i , y ~ j , and z ~ k ), utilizing a uniform grid spacing of h = Δx = Δy = Δz. Based on this structured grid, the discretized equation is51h2[ϵi+1/2ϕi+1+ϵi−1/2ϕi−1+ϵj+1/2ϕj+1+ϵj−1/2ϕj−1+ϵk+1/2ϕk+1+ϵk−1/2ϕk−1]−[Σϵh2+κijk2]ϕijk=ρijkwhere Σϵ=ϵi+1/2+ϵi−1/2+ϵj+1/2+ϵj−1/2+ϵk+1/2+ϵk−1/2 is the sum of the 6 surrounding half-step permittivities, ϵ i+1/2 represents the permittivity at the midpoint between nodes i and i + 1. By flattening the three-dimensional grid into a one-dimensional vector representation, this formulation yields a large-scale linear system6Ax=bwhere N g = N ~ x ~ × N ~ y ~ × N ~ z ~ denotes the total number of points within the structured grid. The system matrix A∈RNg×Ng is a square matrix whose entries are determined by the parameters ϵ, h, and κ, while the vector b encodes the source charge density ρ ijk ~ at each corresponding grid point.

The primary computational challenges arise from the high dimensionality of the linear system, where N a typically exceeds 10^6^, and its extreme sparsity, as the nonzero entries of A are confined to the diagonal and the six adjacent stencil points. To address these complexities, numerous iterative solvers have been employed in biomolecular electrostatics, including successive over-relaxation, the generalized minimal residual method, and the biconjugate gradient stabilized method, alongside standard conjugate gradient techniques. Typically, these iterative methods terminate upon reaching a maximum iteration count or satisfying a preset tolerance ϵ, defined as7|Ax−b||b|<ϵwhere |b| represents the l1-norm of the vector b. Collectively, these approaches facilitate efficient computation of the electrostatic potential ϕ­(r).

Despite substantial progress in solving PB equations, − ,−

the computational cost of the PB equation remains a major bottleneck, particularly for systems exceeding millions of atoms where memory and runtime demands scale prohibitively. This challenge is exacerbated by fragmentation of the current software ecosystem. Because existing solvers were developed independently over decades using disparate legacy implementations, it remains difficult to integrate them into a general framework capable of leveraging modern HPC techniques, such as massive multithreading and CUDA acceleration. Furthermore, the absence of a unified platform prevents rigorous benchmarking of these algorithms on heterogeneous hardware, obscuring their performance characteristics under modern computational constraints. To address these limitations in scalability, software adaptability, and benchmarking, we introduce AmberTorchPB, a unified and accelerator-aware solver framework.

AmberTorchPB: A Tensor-Based Framework for

PB Reaction Field Energy Calculations

Motivated by these limitations concerning computational scalability, software fragmentation, and the scarcity of modern benchmarking infrastructure, we introduce AmberTorchPB. This framework is founded upon the tensor abstraction provided by LibTorch, a robust library for tensor computation. An architectural overview of AmberTorchPB is depicted in Figure .

Framework Overview

The tensor is an abstraction object widely adopted within deep learning communities, initially proposed by Lush and Torch, further developed by Theano, and reaching maturity in TensorFlow and PyTorch. The fundamental motivation underlying this abstraction is to enable researchers to represent multidimensional arrays efficiently. Crucially, this representation encapsulates metainformation regarding the array, including memory allocation across heterogeneous devices and flexible data precision settings.

Beyond data representation, tensor abstraction incorporates distinct interfaces designed to facilitate efficient tensor operations. A primary capability is automatic operator dispatching, which enables the execution of a unified code across diverse platforms by adaptation to the underlying device environment. This feature significantly streamlines tensor computation within HPC infrastructures; developers are required only to implement the platform-specific tensor operator, after which the dispatching mechanism autonomously resolves the appropriate device for execution based on environmental constraints and user directives. Furthermore, the framework supports tensor conversion, allowing researchers to instantiate tensors with varying numerical precisions and memory layouts through a consolidated abstraction. This flexibility enables users to tailor data representation to the specific layout and precision requirements of their workflows.

In AmberTorchPB, we propose leveraging this tensor abstraction to develop a unified PB solver framework that supports multiple sparse matrix layouts, numerical precision types, and hardware acceleration strategies, all while requiring only a single implementation of each algorithm. This approach effectively circumvents the aforementioned limitations regarding computational scalability, software fragmentation, and the lack of a modern benchmarking infrastructure for PB equation solvers, thereby providing a modern testbed for high-performance biomolecular electrostatics.

AmberTorchPB Configurations

The architecture of AmberTorchPB features a backend C++ library constructed upon the tensor abstraction provided by LibTorch. This design exposes a unified interface that enables a caller function, such as the PB solving routine in AMBER PBSA, to inject custom runtime configurations. These parameters facilitate precise control over the solver selection, hardware device targeting, floating-point format selection, and sparse matrix layout. An illustrative YAML configuration for invoking AmberTorchPB solvers is presented in Listing 1. The abbreviated “···” in Listing 1 indicates additional hyperparameters, which are provided in full in Listing S1 in Section S5.

Solver Type

Within AmberTorchPB, we leveraged the tensor object to implement five standard sparse linear system solvers for the PB equation, including the conjugate gradient (CG) method, the biconjugate gradient stabilized (BiCG) method, the generalized minimal residual (GMRES) method, successive over-relaxation (SOR), and red-black successive over-relaxation (RB-SOR). All implemented solvers consistently converge to the specified tolerance, yielding accurate PB free energy values. A detailed discussion regarding the implementation strategies and efficient utilization of tensor computation is provided in Section .

Device Option

A primary advantage of tensor abstraction is algorithmic portability, which allows a single implementation to be deployed across heterogeneous hardware. Consequently, AmberTorchPB supports solver execution on both CPUs and GPUs. To ensure compatibility, we incorporated an automated detection mechanism that verifies whether the runtime environment meets the necessary prerequisites for GPU acceleration. For CPU-based execution, the system defaults to utilizing LibTorch’s native multithreading capabilities to expedite the PB equation solving process.

Precision Type

AmberTorchPB offers flexible options for users to define the numerical precision of their PB linear systems. This selection dictates the floating-point format used to express A, x, and b within the system. While float32 and float64 serve as standard representations for CPU and GPU computations, with float64 being essential for applications requiring high numerical fidelity, we have also extended support to float16 and bfloat16 formats. Although these lower-precision formats are increasingly prevalent in large-scale neural network training, they prove to be equally advantageous for solving the PB equation in large-scale systems. Specifically, they theoretically halve memory consumption and enhance throughput on hardware accelerators, such as GPUs.

Layout Type

To address the representation of large-scale sparse linear systems within matrix A, we implemented two widely adopted sparse tensor storage compressed sparse row (CSR) and coordinate format (COO). To further improve the scalability, we additionally implement a matrix-free stencil layout, in which matrix-vector products are evaluated directly through tensor slicing operations without explicitly storing A. The matrix-free layout is exposed through the same general matrix-vector product interface as the sparse layouts, enabling all implemented solvers to operate with either explicit sparse tensors or matrix-free operators without modification of solver-level formulations.

Preconditioner Type

Preconditioning is essential for PB linear systems because the discretized operators are large, sparse, and often ill-conditioned, and Krylov iterations may stagnate or converge slowly without appropriate spectral conditioning. To improve both convergence robustness and runtime efficiency, AmberTorchPB implements three widely adopted block Jacobi, incomplete Cholesky (ICC), and the algebraic multigrid (AMG). Block Jacobi is selected for its high parallel efficiency and low synchronization overhead on modern multicore and accelerator hardware, ICC is included because sparse approximate Cholesky factors are effective for symmetric positive-definite systems at a modest setup cost, and AMG is provided for its strong scalability and near-mesh-independent convergence behavior on elliptic operators. Benefiting from the tensor abstraction in AmberTorchPB, these preconditioners are exposed through a unified runtime interface across layouts and devices, allowing users to select performant and broadly validated preconditioning strategies without solver-level code changes.

Solver Implementations

To facilitate flexible runtime management, we implemented a base class, TorchPBSolver, which encapsulates all of the shared routines for the derived classes specific to each solver. These routines include initializing linear system parameters within the desired sparse layout, assigning the computational device, retrieving the current solution, and defining virtual functions to be overridden by derived classes. An illustration of this architectural design is presented in Figure .

Regarding the CG, BiCG, and GMRES algorithms, the primary tensor operation required is sparse matrix-vector multiplication (SpMV), denoted as Ax ^ t ^. Here, A represents a large sparse matrix, and x ^ t ^ denotes the dense vector at the t-th iteration of the solver. The tensor library within LibTorch provides robust support for this operator across various devices and floating-point formats. Beyond explicit sparse SpMV, AmberTorchPB provides a matrix-free operator that evaluates Ax ^ t ^ directly through tensor slicing operations, thereby eliminating the explicit materialization of A. The sparse matrix and matrix-free backends are unified under a general matrix-vector product interface in TorchPBSolver, allowing all derived solvers to share a common solver implementation while selecting layouts at runtime.

SpMV-Based SOR

Implementation of SOR presents unique challenges, as its general formulation necessitates a forward-substitution process that is not natively supported by standard tensor operators. Specifically, we decompose the matrix A into three components8A=D+L+Uwhere D represents the diagonal entries of A, while L and U denote the matrices comprising the strictly lower triangular and upper triangular entries of A, respectively. This factorization facilitates the reformulation of the original linear equation as follows9(D+ωL)x=ωb−[ωU+(ω−1)D]xwhere ω > 1 serves as the relaxation factor. The iterative update for SOR is defined as10x(t+1)=(D+ωL)−1{ωb−[ωU+(ω−1)D]x(t)}Explicitly calculating the inverse of the matrix (D + ωL) is computationally intractable for large, sparse matrices. Furthermore, while forward substitution offers a sequential alternative to direct inversion, it inherently presents significant challenges for efficient parallelization. Consequently, we reformulate the update equation to avoid direct inversion11(D+ωL)x(t+1)=ωb−[ωU+(ω−1)D]x(t)To obtain the new solution x ^ t+1^, we leverage the property that D + ωL is a well-defined lower triangular matrix. Solving such triangular linear systems is conducive to parallel execution across different platforms. Regarding the right-hand side, both ωb and the term ωU + (ω – 1) D can be precomputed and cached to eliminate redundant calculations. Furthermore, the term [ωU + (ω – 1)D] x ^(t)^ can be computed using the standard SpMV operator. This methodology allows for the implementation of SOR utilizing exclusively pure tensor operators.

Red-Black SOR

Although the SpMV-based SOR formulation avoids explicit matrix inversion, its dependence on triangular solutions still incurs substantial computational complexity on both CPUs and GPUs because of strong sequential data dependencies, which leads to poor efficiency for large molecular systems and their associated linear systems. To address this limitation, we further implement a RB-SOR solver based on the classical red-black ordering strategy. Specifically, we partition grid points into two complementary color sets and update them in alternating half-sweeps such that all points within each color set can be updated concurrently when conditioned on the opposite set. Benefiting from the tensor abstraction in AmberTorchPB, these red-black updates are expressed through tensor masking and indexing operators, enabling highly parallel execution on both multicore CPUs and GPUs in LibTorch without explicit for-loop traversal. This design preserves the relaxation behavior of SOR while providing an improved parallel performance in our unified framework.

Results and Discussion

Experimental Setup

Data Set Curation

To evaluate the predictive performance of AmberTorchPB for PB free energies, we curated biomolecules drawn from the AMBER PBSA benchmark suite, which encompasses a broad spectrum of biomolecular systems, including both proteins and nucleic acids. The PBSA data set contains 570 proteins and 353 nucleic acid structures, selected to capture diverse functional properties and, in the case of nucleic acids, greater structural rigidity compared with proteins. The data set spans a substantial range of system sizes, from small nucleic acid structures of approximately 250 atoms to large proteins containing up to 8,254 atoms, thereby providing a rigorous test of AmberTorchPB’s scalability and robustness across structurally heterogeneous biomolecular environments.

For the electrostatic PB energy calculations in AMBER PBSA, the PB energy of each molecule was obtained by using the numerical solver (i.e., CG) implemented in the AMBER PBSA program, with the atomic cavity radii taken from the topology files and a solvent probe radius of 1.4 Å. The convergence tolerance was set to 10^–4^ by default, with tighter thresholds of 10^–5^ and 10^–6^ applied when required. Four uniform grid spacings, 0.35, 0.55, 0.75, and 0.95 Å, were used for all molecules, and all other simulation settings were matched to those of the AMBER 23 PBSA module. , All simulations were performed using float32 precision.

To benchmark AmberTorchPB against AMBER PBSA, we implemented an adapter interface that integrates AmberTorchPB directly into the AMBER PBSA program. Through this interface, AmberTorchPB operates as a plug-in, providing numerical solutions to the PB linear system in place of the original AMBER PBSA solver. All experiments were conducted by using the same simulation parameters as those employed in AMBER PBSA to ensure a consistent and fair comparison. In addition to this baseline evaluation, we further examined several distinctive capabilities of AmberTorchPB, including support for multiple solvers, configurable precision options, and cross-platform execution. We observed that both CSR and COO layouts demonstrate accurate convergence. For the sake of brevity, we adopt the CSR layout for the experiments and present the comparative results between CSR and COO in Section S3.

Evaluation Metrics

To quantify the accuracy of the solver implementation in AmberTorchPB, we employed two evaluation the coefficient of determination (R ^2^) and the mean absolute percentage error (MAPE). Specifically, R ^2^ measures the proportion of variance in the ground-truth PB electrostatic solvation free energy ΔG reported by AMBER PBSA that is explained by the values ΔĜ calculated by AmberTorchPB12R2=1−∑i=1n(ΔGi−ΔĜi)2∑i=1n(ΔGi−ΔG̅)2where n is the number of molecules evaluated and ΔG̅ is the mean of the ground-truth values.

To compare the result variations of the same molecule system under different parameters, we also leveraged the MAPE to describe the level of variations. In detail, the MAPE measures the relative difference of the values ΔĜ calculated by AmberTorchPB against the ground-truth PB electrostatic solvation free energy ΔG reported by AMBER PBSA13MAPE=100·1n∑i=1n|ΔĜi−ΔGiΔGi|

Accuracy Analysis of AmberTorchPB in Solving

PB Free Energy

In this experiment, we evaluated the performance of AmberTorchPB in computing PB free energies across two molecule splits, nucleic acids and proteins. Our assessment focused on three key (1) whether AmberTorchPB, built upon a tensor computation abstraction, can achieve accurate energy predictions comparable to the well-studied numerical solver implemented in AMBER PBSA; (2) how its accuracy is maintained across different computational devices, including both CPUs and GPUs; and (3) whether AmberTorchPB reproduces the exact performance trends observed in AMBER PBSA under variations of key parameters, such as grid spacing and convergence tolerance. To quantify prediction accuracy, we employed standard regression metrics, including the R ^2^ and MAPE. Together, these metrics evaluate both the absolute and the relative accuracy of the predicted PB free energies, providing a comprehensive characterization of AmberTorchPB’s predictive performance.

AmberTorchPB Accurately Solves PB Free Energies

across Diverse Biomolecular Types

We first conducted experiments to evaluate the performance of AmberTorchPB in solving PB free energy across two molecule splits, namely, nucleic acids and proteins. Specifically, we compared the energies produced by AmberTorchPB with those obtained from the well-studied numerical PB solvers (i.e., CG and BiCG) implemented in AMBER PBSA. For both software frameworks, we adopted identical solving parameters; atomic cavity radii were taken from the topology files, the solvent probe radius was set to 1.4 Å, the convergence tolerance was fixed at 10^–4^, and the grid spacing was set to 0.55 Å. The results are summarized in Figure . As shown, AmberTorchPB exhibits a strong agreement with the numerical PB solvers in calculating PB free energies. In particular, the regression analysis yields an R ^2^ value of 1.00, indicating excellent predictive accuracy. This conclusion is further corroborated by the similarity in marginal distributions between the calculated energies and the ground-truth values from AMBER PBSA across both biomolecular categories.

AmberTorchPB Achieves High Correspondence

on PB Free Energies across Devices

One advantage of AmberTorchPB is that the solver is implemented once but can be executed seamlessly across different hardware platforms, including CPUs and GPUs. Having already established that AmberTorchPB produces accurate PB free energies compared with AMBER PBSA, we next evaluated whether AmberTorchPB maintains consistent accuracy when running on GPUs. To enable GPU execution, the only required modification is to set the device_opt to 1 in the configuration file, upon which AmberTorchPB automatically detects GPU availability and offloads computation accordingly. We tested GPU execution of AmberTorchPB using both the CG and BiCG solvers on nucleic acid and protein splits with a fixed convergence tolerance of 10^–4^ and a grid spacing of 0.55 Å. The results are presented in Figure . As shown, AmberTorchPB running on GPUs exhibits strong agreement with the PB free energies obtained on CPUs. In particular, the regression analysis achieves an R ^2^ value of 1.00, indicating an excellent predictive correspondence. This conclusion is further supported by the highly similar marginal distributions of energies computed across devices.

AmberTorchPB Closely Tracks AMBER PBSA under

Different Solving Conditions

Having examined the accuracy of AmberTorchPB across different molecule splits and hardware platforms, we next assessed its behavior under varying numerical solving conditions. In this experiment, we evaluate whether AmberTorchPB can reliably capture the PB free energy trends produced by AMBER PBSA when the grid spacing and convergence tolerance are systematically varied. Using both CG and BiCG, we conducted two sets of (1) varying the grid spacing across 0.35, 0.55, 0.75, and 0.95 Å and (2) tightening the convergence tolerance from 10^–4^ to 10^–5^ and 10^–6^. Four representative molecules of different sizes were selected from the two splits, including 1P4Z and 351D from the nucleic acid split and 1b16_A and 1ds1_A from the protein split.

For the grid spacing experiments, the results for CG are shown in Figure , where Figure a reports the PB free energies produced by AMBER PBSA and Figure b presents the corresponding trends from AmberTorchPB under identical conditions. The two curves exhibit nearly identical behaviors across all grid spacings, indicating that AmberTorchPB faithfully tracks the numerical trend of AMBER PBSA and functions as a robust plug-in replacement. The analogous BiCG results are listed in Figure S1.

We further analyzed the MAPE of AmberTorchPB when operating under different convergence tolerances, as summarized in Figure . The MAPE was computed by using the PB free energy obtained by AMBER PBSA at a tolerance of 10^–6^ as the reference. The results show that AmberTorchPB remains highly stable across tolerance settings, with relative differences consistently below 0.01% compared with the 10^–6^ baseline. This confirms that AmberTorchPB maintains robust and accurate PB free energy predictions even when the convergence tolerance is varied.

Extensibility of AmberTorchPB in Implementing

Different Solvers

In addition to the conjugate gradient-based methods (i.e., CG and BiCG) traditionally used for solving PB equations, other iterative solvers such as GMRES and SOR offer broader applicability to matrices with more general structural properties. Unlike CG and BiCG, which rely primarily on SpMV operations, GMRES and SOR require additional computational primitives, including vector rotations and triangular system solvers. To evaluate the extensibility of AmberTorchPB, we extended its implementation to support both GMRES and SOR, and we assessed their convergence behavior in comparison with those of CG and BiCG.

AmberTorchPB is Accurately Extendable to

GMRES and SOR

We first conducted experiments to assess the accuracy of the GMRES and SOR solvers implemented in AmberTorchPB for computing PB free energies across two molecular splits, namely, nucleic acids and proteins. Consistent with the accuracy analysis presented in the previous section, we compared the energies produced by AmberTorchPB’s GMRES and SOR implementations with those obtained from the well-established numerical PB solver in AMBER PBSA. For both software frameworks, we employed identical solving parameters; atomic cavity radii were taken from the topology files, the solvent probe radius was set to 1.4 Å, the convergence tolerance was fixed at 10^–4^, and the grid spacing was set to 0.55 Å. The results are summarized in Figure . As shown in the figure, GMRES and SOR are less stable than CG, exhibiting occasional convergence failures, consistent with observations reported in the existing literature. Nevertheless, for the molecules on which these solvers do converge, AmberTorchPB demonstrates strong agreement with the numerical PB solver in computing the PB free energies. In particular, the regression analysis yields an R ^2^ value of 1.00, indicating excellent predictive accuracy. This conclusion is further supported by the close similarity in the marginal distributions between the predicted energies and the ground-truth values from AMBER PBSA across both biomolecular categories.

Convergence Analysis of Different Solvers

Given the accurate implementation of GMRES and SOR, one novelty of AmberTorchPB is that it provides a modern and unified framework for analyzing the convergence behavior of the four commonly used solvers for PB equations. In particular, AmberTorchPB enables a systematic examination of solver properties while leveraging the latest hardware accelerations within a single, coherent software environment.

We first analyze the number of iterations each solver requires to reach the desired tolerance. In this experiment, the convergence tolerance was fixed at 10^–4^, and the grid spacing was set to 0.55 Å. All four solvers in AmberTorchPB were used to solve PB equations across the two molecule splits, namely, nucleic acids and proteins. We then examined the proportional histogram of iteration counts, as summarized in Figure . As shown in the figure, GMRES requires the fewest iterations to converge on both splits, typically below 1 × 10^2^ iterations, in contrast to approximately 1 × 10^3^ iterations for CG and BiCG. SOR also converges with fewer iterations than CG and BiCG. It is important to note, however, that fewer iterations do not necessarily imply a shorter runtime because each solver incurs different computational costs per iteration on different devices. We also observe that BiCG exhibits a much broader distribution of iteration counts due to its oscillatory and occasionally spiking convergence trajectory. Representative convergence paths for the four molecules are shown in Figure .

Convergence Relationship of Different Solvers

We then analyzed the relationship between the required number of iterations and the atom count of each biomolecular system, where the atom count serves as a proxy for the molecular size. The results are shown in Figure . We observe a strong linear trend for all solvers, with GMRES exhibiting the smallest slope among the four. This indicates that GMRES requires fewer iterations as the system size increases, with SOR showing slightly higher iteration counts but still maintaining favorable scaling. In contrast, CG and BiCG display a more pronounced scaling behavior, where larger systems generally require substantially more iterations for convergence. These observations demonstrate that the total runtime (i.e., the product of iteration count and per-iteration cost) reflects a balance between the molecular size and solver complexity; when the atom count becomes sufficiently large, the increased system size can amortize the computational cost of each iteration.

AmberTorchPB is Efficiently Extendable to

Preconditioners across Solvers

We extended AmberTorchPB with a unified preconditioner interface and evaluated three representative strategies: block Jacobi, incomplete Cholesky (ICC), and algebraic multigrid (AMG), which are classical choices for sparse elliptic systems. −

We then benchmarked preconditioned CG convergence on four representative molecules under the same numerical setting used in prior convergence analyses (a grid spacing of 0.55 Å and a tolerance of 1 × 10^–4^).

As shown in Figure , all three preconditioners substantially accelerate convergence relative to the nonpreconditioned baseline. AMG consistently reaches the target residual in the fewest iterations across all four systems, while block Jacobi and ICC also provide clear and stable iteration reductions. Although the relative ranking between block Jacobi and ICC can vary slightly with molecular system characteristics, both methods markedly improve the convergence behavior and mitigate the long-tail iteration growth observed without preconditioning.

These results demonstrate that AmberTorchPB is not only extensible at the solver level but also efficiently extensible to modern preconditioning components through the same tensor-based abstraction. Importantly, we observe the same qualitative preconditioner pattern on other Krylov solvers implemented in AmberTorchPB, including BiCG and GMRES; AMG provides the strongest acceleration, and block Jacobi/ICC consistently improves convergence over the nonpreconditioned setting. Collectively, this supports the practical applicability of AmberTorchPB for large-scale PB systems and highlights its flexibility for incorporating advanced linear algebra features without solver-specific code duplication.

AmberTorchPB is Efficiently Extendable to

the Matrix-Free Layout across Solvers

We further extended AmberTorchPB with a matrix-free stencil layout that avoids materializing A. Leveraging the tensor abstraction in AmberTorchPB, matrix vector products are computed on the fly through parallel tensor slicing and element-wise operations rather than explicit sparse matrix traversal loops. We evaluated this layout under the same numerical setting used in our previous solver experiments (a grid spacing of 0.55 Å and a tolerance of 10^–4^) and compared it against the CSR sparse matrix layout across CG, BiCG, GMRES, and SOR. All matrix-free layout experiments in this subsection were conducted on the representative molecule 1P4Z.

As shown in Figure , the convergence trajectories from the matrix-free layout closely overlap those from the sparse matrix layout for all four solvers. The solver-dependent convergence pattern is preserved; GMRES reaches the tolerance in the fewest iterations, SOR converges rapidly after an initial phase, and both CG and BiCG maintain the same residual decay behavior as their sparse matrix counterparts. These observations indicate that replacing explicit sparse storage with matrix-free evaluation does not compromise the convergence quality or numerical stability.

We next examined memory consumption using the CG solver across four molecules on the GPU, with atom numbers ranging from approximately 800 to 5900, at a grid spacing of 0.35 Å and a convergence tolerance of 10^–4^. As shown in Figure , matrix-free execution consistently uses less GPU memory than CSR across all evaluated molecules, and the reduction becomes more pronounced as the atom number increases. For the largest tested system, the matrix-free layout lowers the peak memory by more than 500 MiB and by over 30% relative to sparse matrix storage, demonstrating clear scalability benefits for large PB systems.

Overall, these results show that AmberTorchPB is efficiently extendable not only to new iterative solvers and preconditioners but also to modern matrix-free operator layouts through the same tensor-based programming interface. This design enables rapid integration of contemporary numerical features while preserving the solver correctness and improving practical resource efficiency.

AmberTorchPB Enables PB Calculations across

Data Precisions

AmberTorchPB implements an abstraction layer that supports solving PB equations under different data precisions to satisfy varying computational requirements. By default, AmberTorchPB operates in the float32 floating-point format, as it is the most common choice for PB equation solvers and is well supported across both CPUs and GPUs. For higher precision using float64, we observe a convergence behavior consistent with that of float32, as shown in Figure S3. In addition, we examine float16 and bfloat16, which have become increasingly popular in deep learning and machine learning applications due to their substantially reduced memory footprint. This reduction is particularly beneficial for PB calculations, where large biomolecular systems and extremely sparse linear systems demand memory-efficient computation.

We conducted experiments using AmberTorchPB with float16 and bfloat16 on GPUs across all four solvers. However, we found that low-precision execution consistently introduces numerical instability when solving PB equations. This instability arises because the values in the PB linear system are often extremely small and close to zero, making them highly susceptible to underflow in float16 and bfloat16. These issues are further amplified in more complex solvers, such as BiCG, GMRES, and SOR, all of which fail to converge under low-precision formats.

In contrast, we observed that CG can successfully converge using bfloat16, as summarized in Figure . In this experiment, the convergence tolerance was fixed at 10^–4^ and the grid spacing was set to 0.55 Å. Although a small offset is present in the converged PB free energies produced by AmberTorchPB, the overall regression trends remain strong, yielding an R ^2^ value of 1.00 on both molecule splits. We also observed that float16 causes the CG to fail to converge, likely because bfloat16 provides a more fine-grained representation of values near zero.

Representative convergence trajectories for four molecules across different floating-point formats using CG are shown in Figure . As illustrated, float16 initially reduces the relative residual but later diverges due to numerical instability, whereas bfloat16 successfully converges, albeit with a slightly higher iteration count.

Efficiency Analysis of AmberTorchPB

AmberTorchPB enables researchers to implement a solver once and execute it seamlessly across multiple hardware platforms including CPUs and GPUs. This capability provides a substantial benefit to the research community as the conventional approach often requires maintaining separate implementations for each hardware platform. Such duplication not only increases development and maintenance costs but also can introduce subtle discrepancies between implementations, creating a major barrier for software reliability and long-term sustainability. By leveraging the operator-dispatching framework built on LibTorch, AmberTorchPB allows a single solver implementation to execute efficiently using multithreading and various hardware accelerators. In this section, we present experiments evaluating the performance of AmberTorchPB on both CPUs and GPUs and compare its results against those of AMBER PBSA, which maintains both a CPU implementation and a CUDA-specific GPU implementation. We then examine the efficiency of AmberTorchPB’s five solvers as well as the impact of using different floating-point formats.

AmberTorchPB Shows Improved Efficiency Compared

to AMBER PBSA

We first examined the efficiency of AmberTorchPB by comparing its elapsed time against the CG implementations on CPUs and GPUs in AMBER PBSA. All CPU-based timing experiments were performed on a multisocket machine equipped with Intel­(R) Xeon­(R) Gold 6230 processors running at 2.10 GHz, consisting of four sockets with 20 physical cores each. GPU-based experiments were conducted on a system equipped with an NVIDIA GeForce RTX 4090 GPU with 24 GB of memory. The results are shown in Figure . From the figure, we observe that AmberTorchPB achieves more than a 2-fold speedup over AMBER PBSA on CPUs, primarily due to its native multithreading support. Moreover, the matrix-free variant further reduces the CPU elapsed time substantially relative to the sparse layout version, yielding the best overall CPU efficiency among all compared configurations. By contrast, although AMBER PBSA can partially leverage array operations in Fortran, its dominant spMV routine remains serialized, limiting the overall performance. We also observe a near 2-fold speedup compared to AMBER PBSA on GPUs. On the GPU platform, matrix-free execution remains competitive with the sparse layout implementation while preserving a strong acceleration over AMBER PBSA. It is worth noting that the CUDA version of AMBER PBSA relies on manually optimized CUDA kernels from a previous work, whereas AmberTorchPB employs the same unified implementation without GPU-specific optimization. These results demonstrate that AmberTorchPB not only provides flexibility and accuracy for solving PB equations but also achieves a favorable balance between portability and performance, maintaining a competitive speed relative to platform-specific, manually optimized implementations.

We further evaluated the scalability of the native multithreading capability in AmberTorchPB across varying thread counts (1, 2, 4, 6, 8, 16, 32, and 64). The elapsed time under 64 threads served as the reference baseline for each molecule; the relative speedup was subsequently calculated by dividing the reference time by the elapsed time at each thread count. These results are presented in Figure . As shown in Figure a, we examined the speedup across the two data set splits. The results indicate that the speedup increases consistently in both splits as the thread count rises from 1 to 8, validating the effective multithreading scalability of AmberTorchPB. However, when the thread count exceeds 16, the performance in the nucleic acid split peaks and subsequently declines. This trend likely reflects a trade-off between parallel computation gains and the overhead of data synchronization, particularly for molecules of a relatively small size. In contrast, for larger molecules within the protein split, employing 64 threads generally continues to yield computational acceleration. To further investigate the influence of the molecular size on multithreading efficiency, we categorized the molecules from both splits into three Small (<500 atoms), Medium (500–5000 atoms), and Large (>5000 atoms). The performance of these groups is detailed in Figure b. For the Small group, our analysis suggests that 8 threads provide an optimal balance between computational efficiency and the multithreading overhead. Conversely, for both the Medium and Large groups, utilizing a higher thread count proves to be more efficient, with the speedup being most pronounced in the Large group.

Efficiency Analysis of AmberTorchPB on Different

Solvers

We further evaluated the performance of the five solvers implemented in AmberTorchPB across different hardware platforms. The results are summarized in Figure . From the figure, CG, BiCG, and RB-SOR exhibit the shortest elapsed times on both CPUs and GPUs. In particular, compared with classical SOR, RB-SOR consistently achieves a several-fold lower elapsed time while preserving the same relaxation-based solver family, indicating that red-black ordering substantially improves parallel execution efficiency and mitigates the serial dependency bottleneck of SOR. This advantage is especially evident on GPUs, where RB-SOR remains close to CG and BiCG in runtime, whereas SOR is still noticeably slower. Although GMRES and SOR require substantially fewer iterations to reach convergence, their per-iteration computational cost remains higher than those of CG, BiCG, and RB-SOR, resulting in a longer overall runtime. When comparing the same solver across hardware platforms, GPU execution yields roughly a 10-fold speedup relative to CPU execution. This acceleration further highlights the flexibility of AmberTorchPB, which enables a single solver implementation to run efficiently across diverse hardware accelerators without additional platform-specific optimization.

Efficiency Analysis of AmberTorchPB on Different

Data Precisions

Finally, we evaluated the efficiency of AmberTorchPB across multiple floating-point formats, as shown in Figure . From the figure, we observe no substantial difference in the elapsed time on either CPUs or GPUs when running the solver with different numerical precisions. On GPUs, although the peak FLOPs of float64 arithmetic are significantly lower than those of float32 (e.g., on an NVIDIA GeForce RTX 4090, the float64 throughput is far below that of float32), the resulting wall-clock time remains similar, with float64 being only slightly slower. This behavior can be attributed to the extreme sparsity of the PB linear system, where the communication overhead dominates the runtime rather than the floating-point throughput. A similar performance pattern is also observed for bfloat16.

We further analyzed peak GPU memory consumption under differing data precisions using the CG solver as a representative example. Four molecules varying in size from approximately 1000 to 6000 atoms were selected for evaluation. The results, as summarized in Figure , indicate that the float64 format yields the highest memory cost among the three precisions tested. Furthermore, we observed that bfloat16 achieves significant GPU memory savings compared to float32, exceeding 30% as atom counts increase. It is noteworthy that bfloat16 does not realize the theoretical 50% memory reduction relative to float32; this is attributable to the design of AmberTorchPB, which retains the solution at the current step in float32 to maintain compatibility with the caller memory layout. Consequently, in summary, utilizing bfloat16 substantially mitigates memory requirements relative to those of float32, especially for larger molecular systems.

Conclusion

In this work, we introduced AmberTorchPB, a unified and extensible framework leveraging tensor abstraction to modernize the solution of the PB equation for biomolecular electrostatics. By building upon the LibTorch ecosystem, we addressed the long-standing challenges of software fragmentation and legacy code maintenance that have historically impeded the adoption of modern high-performance computing in this domain. Our comprehensive empirical evaluation demonstrates that AmberTorchPB not only replicates the high-fidelity accuracy of the AMBER PBSA solver across diverse nucleic acid and protein systems but also does so with superior computational efficiency. Specifically, our framework achieves more than a 2-fold speedup on both CPUs and GPUs compared to AMBER PBSA, validating the effectiveness of native multithreading and operator-level hardware acceleration. Furthermore, AmberTorchPB exhibits remarkable versatility through its “implement once, deploy everywhere” architecture. We successfully integrated a diverse suite of iterative solvers, including CG, BiCG, GMRES, SOR, and RB-SOR, and systematically benchmarked their convergence behaviors and scaling properties within a single computational environment. The framework also enabled a rigorous exploration of mixed-precision arithmetic; while reduced precision formats like float16 generally introduced instability for complex solvers, the successful convergence of CG under bfloat16 highlights promising avenues for future memory-efficient optimizations. At the same time, the current implementation is intentionally scoped to finite difference linear PB computations and does not yet include finite element discretizations, nonlinear PB solvers, ion-size-modified PB models, or specialized interface and singularity treatments. Nevertheless, we have already demonstrated the flexibility and extensibility of AmberTorchPB by integrating multiple iterative solvers, a matrix-free stencil formulation, mixed-precision execution, and preconditioning options within a unified tensor-centric infrastructure. These results establish a practical foundation for incorporating modern nonlinear PB algorithms and additional numerical advances in future work. Ultimately, AmberTorchPB serves as a robust, accelerator-aware testbed that decouples algorithmic innovation from hardware complexity. By providing a standardized interface for rapid prototyping and rigorous benchmarking, AmberTorchPB paves the way for the next generation of scalable, high-performance tools in computational biochemistry and biophysics.

Supplementary Material