Development of stratified stochastic tensor contraction method for applications in electronic structure theory

Calculation of high-rank tensor contractions plays a central role in computational physics, quantum chemistry, and computer science. The ability to perform a tensor contraction within a given computational budget poses a challenge to feasibility and restricts the types of systems that can be investigated computationally. We present the development and implementation of the Stratified Stochastic Tensor Contraction (SSTC) method as an alternative to sequential evaluation of tensor contraction. The SSTC approach partitions the tensor indices into stratified segments and employs composite index mappings to guide efficient sampling. By leveraging intrinsic structure within tensors, SSTC achieves rapid convergence to exact results with controllable sampling errors. We demonstrate the method on representative cases, two-point Coulomb integrals, four-point kernels, and correlation energy calculations using 2nd order Møller–Plesset perturbation theory, all of which are relevant to explicitly correlated electronic structure theory. In all examples, the SSTC method provides an accurate estimate of the tensor contraction with systematically reducible errors. We present the theory of exact tensor contractions as a foundation for the stochastic approach, followed by the mathematical derivation of the SSTC framework. Key components include the construction of a composite index to enable loop compression, and the application of functional minimization techniques for variance reduction. A discussion comparing sequential and stochastic contraction strategies, analysis of the distribution of sampling points through cumulative distribution functions and correlations between sampling assignments and kernel values are presented. A description of the SSTC method’s impact on computational clock times and dependence on segment number is included. Additional considerations include sampling with vs without replacement and the computational advantages of index mapping. These results highlight the effectiveness of the SSTC method as a versatile and scalable alternative to the conventional sequential tensor contraction approach, with broad applicability to quantum chemistry, many-body physics, stochastic quantum mechanics, and tensor-based machine learning.