Coverage for src/lsqfitgp/_gvarext/_ufunc.py: 94%

1# lsqfitgp/_gvarext/_ufunc.py 

2# 

3# Copyright (c) 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 

20import functools 3 ctx1bca

21import string 3 ctx1bca

22 

23import gvar 3 ctx1bca

24import jax 3 ctx1bca

25from jax import numpy as jnp 3 ctx1bca

26 

27from lsqfitgp import _signature 3 ctx1bca

28from lsqfitgp._gvarext._jacobian import from_jacobian, jacobian 3 ctx1bca

29 

30 

31def gvar_gufunc(func, *, signature=None): 3 ctx1bca

32 """ 

33 Wrap a jax-traceable generalized ufunc with one argument to support gvars. 

34 

35 Parameters 

36 ---------- 

37 func : callable 

38 A function from one array to one array. It must be a generalized ufunc, 

39 and differentiable one time with `jax`. 

40 signature : str, optional 

41 The signature of the generalized ufunc. If not specified, it is assumed 

42 to be scalar to scalar (normal ufunc). 

43 

44 Returns 

45 ------- 

46 decorated_func : callable 

47 A function that, in addition to numerical arrays, accepts gvars and 

48 returns gvars. 

49 

50 See Also 

51 -------- 

52 numpy.vectorize 

53 """ 

54 # parse signature 

55 if signature is None: 55 ↛ 56line 55 didn't jump to line 56 because the condition on line 55 was never true3 ctx1bca

56 signature = '()->()' 

57 sig = _signature.Signature(signature) 3 ctx1bca

58 (inp,) = sig.incores 3 ctx1bca

59 (out,) = sig.outcores 3 ctx1bca

60 jac_sig = _signature.Signature.from_tuples([inp], [out + inp]) 3 ctx1bca

61 

62 # make jacobian function 

63 deriv = jnp.vectorize(jax.jacfwd(func), signature=jac_sig.signature) 3 ctx1bca

64 

65 # get indices for summation 

66 ninp = len(inp) 3 ctx1bca

67 nout = len(out) 3 ctx1bca

68 head_indices = '...' 3 ctx1bca

69 out_indices = string.ascii_letters[:nout] 3 ctx1bca

70 in_indices = string.ascii_letters[nout : nout + ninp] 3 ctx1bca

71 gvar_indices = string.ascii_letters[nout + ninp] 3 ctx1bca

72 

73 # make summation formula 

74 jac_indices = head_indices + out_indices + in_indices 3 ctx1bca

75 in_jac_indices = head_indices + in_indices + gvar_indices 3 ctx1bca

76 out_indices = head_indices + out_indices + gvar_indices 3 ctx1bca

77 formula = f'{jac_indices},{in_jac_indices}->{out_indices}' 3 ctx1bca

78 

79 def gvar_function(x): 3 ctx1bca

80 

81 # unpack the gvars 

82 in_mean = gvar.mean(x) 3 ctx1bca

83 in_jac, indices = jacobian(x) 3 ctx1bca

84 

85 # apply function 

86 out_mean = func(in_mean) 3 ctx1bca

87 jac = deriv(in_mean) 3 ctx1bca

88 

89 # check shapes match 

90 head_ndim = jac.ndim - nout - ninp 3 ctx1bca

91 assert jac.shape[:head_ndim] == in_jac.shape[: in_jac.ndim - 1 - ninp] 3 ctx1bca

92 

93 # contract 

94 out_jac = jnp.einsum(formula, jac, in_jac) 3 ctx1bca

95 

96 # pack output 

97 return from_jacobian(out_mean, out_jac, indices) 3 ctx1bca

98 

99 @functools.wraps(func) 3 ctx1bca

100 def decorated_func(x): 3 ctx1bca

101 if isinstance(x, gvar.GVar): 3 ctx1bca

102 out = gvar_function(x) 3 ctx1bca

103 if not out.ndim: 103 ↛ 105line 103 didn't jump to line 105 because the condition on line 103 was always true3 ctx1bca

104 out = out.item() 3 ctx1bca

105 return out 3 ctx1bca

106 elif getattr(x, 'dtype', None) == object: 3 ctx1bca

107 return gvar_function(x) 1 ctx1a

108 else: 

109 return func(x) 3 ctx1bca

110 

111 return decorated_func 3 ctx1bca