Coverage for src/lsqfitgp/_signature.py: 87%

1# lsqfitgp/_signature.py 

2# 

3# Copyright (c) 2023, 2024, 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"""Define Signature.""" 

21 

22import inspect 3 ctx1cab

23 

24import jax 3 ctx1cab

25from numpy.lib import _function_base_impl as function_base 3 ctx1cab

26 

27 

28class Signature: 3 ctx1cab

29 """Class to parse a numpy gufunc signature.""" 

30 

31 def __init__(self, signature): 3 ctx1cab

32 self.signature = signature 3 ctx1cab

33 self.incores, self.outcores = function_base._parse_gufunc_signature(signature) 3 ctx1cab

34 

35 @classmethod 3 ctx1cab

36 def from_tuples(cls, incores, outcores): 3 ctx1cab

37 self = cls.__new__(cls) 3 ctx1cab

38 tuplestr = lambda t: '(' + ','.join(map(str, t)) + ')' 3 ctx1cab

39 self.signature = ( 3 ctx1cab

40 ','.join(map(tuplestr, incores)) + '->' + ','.join(map(tuplestr, outcores)) 

41 ) 

42 self.incores = incores 3 ctx1cab

43 self.outcores = outcores 3 ctx1cab

44 return self 3 ctx1cab

45 

46 def __repr__(self): 3 ctx1cab

47 return self.signature 

48 

49 def check_nargs(self, func): 3 ctx1cab

50 """Check that the function has the correct number of arguments.""" 

51 sig = inspect.signature(func) 3 ctx1cab

52 if any( 3 ctx1cab

53 p.kind == inspect.Parameter.VAR_POSITIONAL for p in sig.parameters.values() 

54 ): 

55 return 2 ctx1ab

56 if len(sig.parameters) != len(self.incores): 56 ↛ 57line 56 didn't jump to line 57 because the condition on line 56 was never true3 ctx1cab

57 msg = ( 

58 f'function {func} has {len(sig.parameters)} ' 

59 f'arguments, but signature {self.signature} ' 

60 f'requires {len(self.incores)}' 

61 ) 

62 raise ValueError(msg) 

63 

64 @property 3 ctx1cab

65 def nin(self): 3 ctx1cab

66 return len(self.incores) 3 ctx1cab

67 

68 @property 3 ctx1cab

69 def nout(self): 3 ctx1cab

70 return len(self.outcores) 

71 

72 def eval(self, *args): 3 ctx1cab

73 """ 

74 Evaluate the signature with the given arguments. 

75 

76 Parameters 

77 ---------- 

78 *args : sequence of numpy.ndarray or None 

79 A missing argument can be replaced with None, provided the other 

80 arguments are sufficient to infer all dimension sizes. 

81 

82 Returns 

83 ------- 

84 sig : EvaluatedSignature 

85 An object with attributes `broadcast_shape`, `sizes`, 

86 `core_out_shapes`, `out_shapes`, `core_in_shapes`, `in_shapes`. 

87 """ 

88 return self.EvaluatedSignature(self, *args) 3 ctx1cab

89 

90 class EvaluatedSignature: 3 ctx1cab

91 def __init__(self, sig, *args): 3 ctx1cab

92 

93 assert len(args) == len(sig.incores) 3 ctx1cab

94 

95 known_args = [] 3 ctx1cab

96 known_cores = [] 3 ctx1cab

97 missing_cores = [] 3 ctx1cab

98 for arg, core in zip(args, sig.incores, strict=True): 3 ctx1cab

99 if arg is None: 3 ctx1cab

100 missing_cores.append(core) 3 ctx1cab

101 else: 

102 known_args.append(jax.ShapeDtypeStruct(arg.shape, 'd')) 3 ctx1cab

103 known_cores.append(core) 3 ctx1cab

104 

105 self.broadcast_shape, self.sizes = function_base._parse_input_dimensions( 3 ctx1cab

106 known_args, known_cores 

107 ) 

108 

109 missing_indices = set(sum(missing_cores, ())) 3 ctx1cab

110 missing_indices.difference_update(self.sizes) 3 ctx1cab

111 if missing_indices: 111 ↛ 112line 111 didn't jump to line 112 because the condition on line 111 was never true3 ctx1cab

112 msg = f'cannot infer sizes of dimesions {missing_indices} from signature {sig.signature}' 

113 raise ValueError(msg) 

114 

115 self.core_out_shapes, self.out_shapes = self._compute_shapes(sig.outcores) 3 ctx1cab

116 self.core_in_shapes, self.in_shapes = self._compute_shapes(sig.incores) 3 ctx1cab

117 

118 def _compute_shapes(self, cores): 3 ctx1cab

119 coreshapes = [] 3 ctx1cab

120 shapes = [] 3 ctx1cab

121 for core in cores: 3 ctx1cab

122 core = tuple(self.sizes[i] for i in core) 3 ctx1cab

123 coreshapes.append(core) 3 ctx1cab

124 shapes.append(self.broadcast_shape + core) 3 ctx1cab

125 return tuple(coreshapes), tuple(shapes) 3 ctx1cab

126 

127 def _repr(self, shapes): 3 ctx1cab

128 return ','.join(map(str, shapes)).replace(' ', '') 

129 

130 def __repr__(self): 3 ctx1cab

131 return self._repr(self.in_shapes) + '->' + self._repr(self.out_shapes) 

132 

133 

134# I use numpy's internals to parse the signature, but these do not correspond to 

135# the description in 

136# https://numpy.org/doc/stable/reference/c-api/generalized-ufuncs.html. In 

137# particular: 

138# - integers are parsed like identifiers 

139# - ? is not accepted 

140# Looking at the issues, it seems a long standing issue amongst many with 

141# vectorize and is not going to be solved. 

142# See https://github.com/HypothesisWorks/hypothesis/blob/4e675dee1a4cba9d6902290bbc5719fd72072ec7/hypothesis-python/src/hypothesis/extra/_array_helpers.py#L289 

143# for a more complete implementation