Coverage for src/lsqfitgp/_linalg/_stdcplx.py: 44%

1# lsqfitgp/_linalg/_stdcplx.py 

2# 

3# Copyright (c) 2022, 2023, 2026, Giacomo Petrillo 

4# 

5# This file is part of lsqfitgp. 

6# 

7# lsqfitgp is free software: you can redistribute it and/or modify 

8# it under the terms of the GNU General Public License as published by 

9# the Free Software Foundation, either version 3 of the License, or 

10# (at your option) any later version. 

11# 

12# lsqfitgp is distributed in the hope that it will be useful, 

13# but WITHOUT ANY WARRANTY; without even the implied warranty of 

14# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

15# GNU General Public License for more details. 

16# 

17# You should have received a copy of the GNU General Public License 

18# along with lsqfitgp. If not, see <http://www.gnu.org/licenses/>. 

19 

20"""Module to estimate the time taken by standard linear algebra operations.""" 

21 

22import inspect 3 ctx1abc

23import timeit 3 ctx1abc

24 

25from jax import numpy as jnp 3 ctx1abc

26from jax import random 3 ctx1abc

27from jax.scipy import linalg as jlinalg 3 ctx1abc

28 

29 

30def benchmark(func, *args, **kwargs): 3 ctx1abc

31 timer = timeit.Timer('func(*args, **kwargs)', globals=locals()) 

32 n, _ = timer.autorange() 

33 times = timer.repeat(5, n) 

34 return min(times) / n 

35 

36 

37# the lambdas fix the number of arguments, counted with `inspect.signature` 

38ops = { 3 ctx1abc

39 'chol': [ 

40 lambda x: jnp.linalg.cholesky(x), # noqa: PLW0108, the operation 

41 lambda s: s[0] ** 3, # complexity in terms of arguments' shapes 

42 ], 

43 'eigh': [lambda x: jnp.linalg.eigh(x), lambda s: s[0] ** 3], # noqa: PLW0108 

44 'qr-red': [ 

45 lambda x: jnp.linalg.qr(x, mode='reduced'), 

46 lambda s: min(s) ** 2 * max(s), 

47 ], 

48 'qr-full': [lambda x: jnp.linalg.qr(x, mode='full'), lambda s: max(s) ** 3], 

49 'svd-red': [ 

50 lambda x: jnp.linalg.svd(x, full_matrices=False), 

51 lambda s: min(s) ** 2 * max(s), 

52 ], 

53 'svd-full': [ 

54 lambda x: jnp.linalg.svd(x, full_matrices=True), 

55 lambda s: max(s) ** 3, 

56 ], 

57 'solve_triangular': [ 

58 lambda x, y: jlinalg.solve_triangular(x, y), # noqa: PLW0108 

59 lambda s, t: s[0] ** 2 * t[1], 

60 ], 

61 'matmul': [lambda x, y: jnp.matmul(x, y), lambda s, t: s[0] * s[1] * t[1]], # noqa: PLW0108 

62} 

63 

64 

65def gen_ops_factors(n): # pragma: no cover 3 ctx1abc

66 key = random.key(202208101236) 

67 factors = {} 

68 for op, (job, est) in ops.items(): 

69 print(f'{op}({n})... ', end='', flush=True) # noqa: T201 

70 nparams = len(inspect.signature(job).parameters) 

71 key, subkey = random.split(key) 

72 m = random.normal(subkey, (nparams, n, n), jnp.float32) 

73 args = m @ jnp.swapaxes(m, -2, -1) 

74 time = benchmark(job, *args) 

75 print(f'{time:.2g} s') # noqa: T201 

76 factors[op] = time / est(*(a.shape for a in args)) 

77 return factors 

78 

79 

80ops_factors = { 3 ctx1abc

81 'chol': 6.03470915928483e-12, 

82 'eigh': 1.824986875290051e-10, 

83 'qr-red': 1.1241237493231893e-10, 

84 'qr-full': 1.2058762495871633e-10, 

85 'svd-red': 4.468000000342727e-10, 

86 'svd-full': 4.2561762500554324e-10, 

87 'solve_triangular': 4.1634716559201486e-12, 

88 'matmul': 5.6301691802218555e-12, 

89} # = gen_ops_factors(1000) 

90 

91ops_consts = { 3 ctx1abc

92 'chol': 1.810961455339566e-06, 

93 'eigh': 2.390482500195503e-06, 

94 'qr-red': 2.6676162518560884e-06, 

95 'qr-full': 2.6932845800183714e-06, 

96 'svd-red': 3.7152979196980598e-06, 

97 'svd-full': 3.663789590355009e-06, 

98 'solve_triangular': 2.170706249307841e-06, 

99 'matmul': 1.718031665077433e-06, 

100} # = gen_ops_factors(1) 

101 

102 

103def predtime(op, shapes, types): 3 ctx1abc

104 """ 

105 Estimate the time taken by a linear algebra operation. 

106 

107 Parameters 

108 ---------- 

109 op : str 

110 The identifier of the operation, see `listops`. 

111 shapes : sequence of tuples of integers 

112 The shapes of the arguments. 

113 types : sequence of numpy data types 

114 The types of the arguments. They are promoted according to JAX rules. 

115 

116 Returns 

117 ------- 

118 time : float 

119 The estimated time. Not accurate. The unit of measure is seconds on a 

120 particular laptop cpu used to calibrate the estimate with 1000x1000 

121 matrices. 

122 """ 

123 _, est = ops[op] 

124 factor = ops_factors[op] 

125 const = ops_consts[op] 

126 dt = jnp.sin(jnp.empty(0, jnp.result_type(*types))).dtype 

127 if dt == jnp.float64: 

128 factor *= 2 

129 return const + factor * est(*shapes) 

130 

131 

132def listops(): 3 ctx1abc

133 """ 

134 List available linear algebra operations. 

135 

136 Returns 

137 ------- 

138 ops : dict 

139 A dictionary operation identifier -> number of arguments. 

140 """ 

141 return {op: len(inspect.signature(job).parameters) for op, (job, _) in ops.items()}