import sympy as sp
def main() -> None:
rho = sp.symbols("rho", integer=True, positive=True)
t_vesin = 27 * rho - 1
t_ours = (rho - 1) / 2 + 13 * rho
ratio = sp.simplify(t_vesin / t_ours)
difference = sp.factor(t_vesin - 2 * t_ours)
mul, add = sp.symbols("mul add", integer=True, positive=True)
vesin_shift = 9 * mul + 6 * add
vesin_delta = 3 * add + 3 * add vesin_dot = 3 * mul + 2 * add
flops_vesin = sp.simplify(vesin_shift + vesin_delta + vesin_dot)
flops_ours = 3 * add + 3 * mul + 2 * add
flop_ratio = sp.simplify(
sp.simplify(flops_vesin.subs({mul: 1, add: 1}))
* t_vesin
/ (sp.simplify(flops_ours.subs({mul: 1, add: 1})) * t_ours)
)
print("per atom, rho atoms in the home cell")
print(f" vesin tests = {t_vesin}")
print(f" ours tests = {t_ours}")
print(f" vesin / ours = {ratio}")
print(f" vesin - 2*ours = {difference}")
print(f" vesin flops / candidate = {flops_vesin}")
print(f" ours flops / candidate = {flops_ours}")
print(f" flop ratio (vesin / ours) = {flop_ratio}")
n_bin = sp.Integer(4)
atoms = n_bin**3 * 64
rho_c = sp.Integer(64)
hits = 729088 tests_v = int(t_vesin.subs(rho, rho_c) * atoms)
tests_u = int(t_ours.subs(rho, rho_c) * atoms)
fv = int(flops_vesin.subs({mul: 1, add: 1}))
fo = int(flops_ours.subs({mul: 1, add: 1}))
print()
print(f"18 Å cube, cutoff 4 Å, {atoms} atoms, {rho_c} per cell")
print(f" vesin distance tests = {tests_v}")
print(f" ours distance tests = {tests_u}")
print(f" vesin flops = {tests_v * fv}")
print(f" ours flops = {tests_u * fo}")
bytes_vesin = hits * (8 * 2 + 4 * 3)
bytes_pair = hits * 40
bytes_scratch = (hits // 2) * (4 + 8 + 8)
print(f" full-list rows = {hits}")
print(f" vesin shift bytes = {bytes_vesin}")
print(f" pair bytes = {bytes_pair}")
print(f" scratch bytes = {bytes_scratch}")
print(
" extra traffic is the scratch and the 40-byte row; "
"the distance loop is the smaller one"
)
if __name__ == "__main__":
main()