Coverage for src/lsqfitgp/_Kernel/_util.py: 100%

1# lsqfitgp/_Kernel/_util.py 

2# 

3# Copyright (c) 2020, 2022, 2023, 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 numbers 6 ctx1efabcd

21import operator 6 ctx1efabcd

22 

23import numpy 6 ctx1efabcd

24import jax 6 ctx1efabcd

25from jax import numpy as jnp 6 ctx1efabcd

26from jax import tree_util 6 ctx1efabcd

27 

28from .. import _array 6 ctx1efabcd

29 

30def is_numerical_scalar(x): 6 ctx1efabcd

31 return ( 6 ctx1efabcd

32 isinstance(x, numbers.Number) or 

33 (isinstance(x, (numpy.ndarray, jnp.ndarray)) and x.ndim == 0) 

34 ) 

35 # do not use jnp.isscalar because it returns False for strongly 

36 # typed 0-dim arrays; do not use jnp.ndim(•) == 0 because it accepts 

37 # non-numerical types 

38 

39def is_nonnegative_integer_scalar(x): 6 ctx1efabcd

40 if isinstance(x, numbers.Integral) and x >= 0: 4 ctx1abcd

41 # python scalars and numpy scalars 

42 return True 4 ctx1abcd

43 if isinstance(x, numpy.ndarray) and x.ndim == 0 and numpy.issubdtype(x.dtype, numpy.integer) and x.item() >= 0: 4 ctx1abcd

44 # 0-dim numpy arrays 

45 return True 4 ctx1abcd

46 if isinstance(x, jnp.ndarray) and x.ndim == 0 and jnp.issubdtype(x.dtype, jnp.integer): 4 ctx1abcd

47 try: 4 ctx1abcd

48 # concrete jax arrays 

49 return x.item() >= 0 4 ctx1abcd

50 except jax.errors.ConcretizationTypeError: 4 ctx1abcd

51 # jax tracers 

52 return jnp.issubdtype(x.dtype, jnp.unsignedinteger) 4 ctx1abcd

53 return False 4 ctx1abcd

54 

55def is_scalar_cond_trueontracer(x, cond): 6 ctx1efabcd

56 if isinstance(x, numbers.Number) and cond(x): 6 ctx1efabcd

57 # python scalars and numpy scalars 

58 return True 6 ctx1efabcd

59 if isinstance(x, numpy.ndarray) and x.ndim == 0 and numpy.issubdtype(x.dtype, numpy.number) and cond(x.item()): 6 ctx1efabcd

60 # 0-dim numpy arrays 

61 return True 4 ctx1abcd

62 if isinstance(x, jnp.ndarray) and x.ndim == 0 and jnp.issubdtype(x.dtype, jnp.number): 6 ctx1efabcd

63 try: 6 ctx1efabcd

64 # concrete jax arrays 

65 return cond(x.item()) 6 ctx1efabcd

66 except jax.errors.ConcretizationTypeError: 6 ctx1efabcd

67 # jax tracers 

68 return True 6 ctx1efabcd

69 return False 4 ctx1abcd

70 

71def is_nonnegative_scalar_trueontracer(x): 6 ctx1efabcd

72 return is_scalar_cond_trueontracer(x, lambda x: x >= 0) 6 ctx1efabcd

73 

74# TODO reimplement with tree_reduce, closuring ndim to recognize shaped fields 

75def _reduce_recurse_dtype(fun, args, reductor, npreductor, jnpreductor, **kw): 6 ctx1efabcd

76 x = args[0] 6 ctx1efabcd

77 if x.dtype.names is None: 6 ctx1efabcd

78 return fun(*args, **kw) 6 ctx1efabcd

79 else: 

80 acc = None 6 ctx1efabcd

81 for name in x.dtype.names: 6 ctx1efabcd

82 recargs = tuple(arg[name] for arg in args) 6 ctx1efabcd

83 result = _reduce_recurse_dtype(fun, recargs, reductor, npreductor, jnpreductor, **kw) 6 ctx1efabcd

84 

85 dtype = x.dtype[name] 6 ctx1efabcd

86 if dtype.ndim: 6 ctx1efabcd

87 axis = tuple(range(-dtype.ndim, 0)) 6 ctx1efabcd

88 red = jnpreductor if isinstance(result, jnp.ndarray) else npreductor 6 ctx1efabcd

89 result = red(result, axis=axis) 6 ctx1efabcd

90 

91 if acc is None: 6 ctx1efabcd

92 acc = result 6 ctx1efabcd

93 else: 

94 acc = reductor(acc, result) 5 ctx1eabcd

95 

96 assert acc.shape == _array.broadcast(*args).shape 6 ctx1efabcd

97 return acc 6 ctx1efabcd

98 

99def sum_recurse_dtype(fun, *args, **kw): 6 ctx1efabcd

100 return _reduce_recurse_dtype(fun, args, operator.add, numpy.sum, jnp.sum, **kw) 6 ctx1efabcd

101 

102def prod_recurse_dtype(fun, *args, **kw): 6 ctx1efabcd

103 return _reduce_recurse_dtype(fun, args, operator.mul, numpy.prod, jnp.prod, **kw) 6 ctx1efabcd

104 

105def ufunc_recurse_dtype(ufunc, x, *args): 6 ctx1efabcd

106 """ apply an ufunc to all the leaf fields """ 

107 

108 allargs = (x, *args) 6 ctx1efabcd

109 expected_shape = jnp.broadcast_shapes(*(x.shape for x in allargs)) 6 ctx1efabcd

110 

111 if x.dtype.names is None: 6 ctx1efabcd

112 out = ufunc(*allargs) 6 ctx1efabcd

113 else: 

114 args = map(_array.StructuredArray, allargs) 6 ctx1efabcd

115 out = tree_util.tree_map(ufunc, *args) 6 ctx1efabcd

116 

117 assert out.shape == expected_shape 6 ctx1efabcd

118 return out 6 ctx1efabcd