github.com/coyove/nj@v0.0.0-20221110084952-c7f8db1065c3/tests/bench/spec.py (about)

     1  from itertools import repeat
     2  from math import sqrt
     3  from multiprocessing import Pool
     4  from sys import argv
     5  import time
     6  
     7  
     8  def eval_A(i, j):
     9      ij = i + j
    10      return ij * (ij + 1) // 2 + i + 1
    11  
    12  
    13  def A_sum(u, i):
    14      return sum(u_j / eval_A(i, j) for j, u_j in enumerate(u))
    15  
    16  
    17  def At_sum(u, i):
    18      return sum(u_j / eval_A(j, i) for j, u_j in enumerate(u))
    19  
    20  
    21  def multiply_AtAv(u):
    22      r = range(len(u))
    23  
    24      tmp = pool.starmap(
    25          A_sum,
    26          zip(repeat(u), r)
    27      )
    28      return pool.starmap(
    29          At_sum,
    30          zip(repeat(tmp), r)
    31      )
    32  
    33  
    34  def main():
    35      n = int(argv[1])
    36      u = [1] * n
    37  
    38      for _ in range(10):
    39          v = multiply_AtAv(u)
    40          u = multiply_AtAv(v)
    41  
    42      vBv = vv = 0
    43  
    44      for ue, ve in zip(u, v):
    45          vBv += ue * ve
    46          vv  += ve * ve
    47  
    48      result = sqrt(vBv/vv)
    49      print("{0:.9f}".format(result))
    50  
    51  
    52  if __name__ == '__main__':
    53      start = time.time()
    54      with Pool(processes=1) as pool:
    55          main()
    56      print(time.time() - start)