Coverage for /usr/lib/python3/dist-packages/sympy/core/multidimensional.py: 23%

1""" 

2Provides functionality for multidimensional usage of scalar-functions. 

3 

4Read the vectorize docstring for more details. 

5""" 

6 

7from functools import wraps 

8 

9 

10def apply_on_element(f, args, kwargs, n): 

11 """ 

12 Returns a structure with the same dimension as the specified argument, 

13 where each basic element is replaced by the function f applied on it. All 

14 other arguments stay the same. 

15 """ 

16 # Get the specified argument. 

17 if isinstance(n, int): 

18 structure = args[n] 

19 is_arg = True 

20 elif isinstance(n, str): 

21 structure = kwargs[n] 

22 is_arg = False 

23 

24 # Define reduced function that is only dependent on the specified argument. 

25 def f_reduced(x): 

26 if hasattr(x, "__iter__"): 

27 return list(map(f_reduced, x)) 

28 else: 

29 if is_arg: 

30 args[n] = x 

31 else: 

32 kwargs[n] = x 

33 return f(*args, **kwargs) 

34 

35 # f_reduced will call itself recursively so that in the end f is applied to 

36 # all basic elements. 

37 return list(map(f_reduced, structure)) 

38 

39 

40def iter_copy(structure): 

41 """ 

42 Returns a copy of an iterable object (also copying all embedded iterables). 

43 """ 

44 return [iter_copy(i) if hasattr(i, "__iter__") else i for i in structure] 

45 

46 

47def structure_copy(structure): 

48 """ 

49 Returns a copy of the given structure (numpy-array, list, iterable, ..). 

50 """ 

51 if hasattr(structure, "copy"): 

52 return structure.copy() 

53 return iter_copy(structure) 

54 

55 

56class vectorize: 

57 """ 

58 Generalizes a function taking scalars to accept multidimensional arguments. 

59 

60 Examples 

61 ======== 

62 

63 >>> from sympy import vectorize, diff, sin, symbols, Function 

64 >>> x, y, z = symbols('x y z') 

65 >>> f, g, h = list(map(Function, 'fgh')) 

66 

67 >>> @vectorize(0) 

68 ... def vsin(x): 

69 ... return sin(x) 

70 

71 >>> vsin([1, x, y]) 

72 [sin(1), sin(x), sin(y)] 

73 

74 >>> @vectorize(0, 1) 

75 ... def vdiff(f, y): 

76 ... return diff(f, y) 

77 

78 >>> vdiff([f(x, y, z), g(x, y, z), h(x, y, z)], [x, y, z]) 

79 [[Derivative(f(x, y, z), x), Derivative(f(x, y, z), y), Derivative(f(x, y, z), z)], [Derivative(g(x, y, z), x), Derivative(g(x, y, z), y), Derivative(g(x, y, z), z)], [Derivative(h(x, y, z), x), Derivative(h(x, y, z), y), Derivative(h(x, y, z), z)]] 

80 """ 

81 def __init__(self, *mdargs): 

82 """ 

83 The given numbers and strings characterize the arguments that will be 

84 treated as data structures, where the decorated function will be applied 

85 to every single element. 

86 If no argument is given, everything is treated multidimensional. 

87 """ 

88 for a in mdargs: 

89 if not isinstance(a, (int, str)): 

90 raise TypeError("a is of invalid type") 

91 self.mdargs = mdargs 

92 

93 def __call__(self, f): 

94 """ 

95 Returns a wrapper for the one-dimensional function that can handle 

96 multidimensional arguments. 

97 """ 

98 @wraps(f) 

99 def wrapper(*args, **kwargs): 

100 # Get arguments that should be treated multidimensional 

101 if self.mdargs: 

102 mdargs = self.mdargs 

103 else: 

104 mdargs = range(len(args)) + kwargs.keys() 

105 

106 arglength = len(args) 

107 

108 for n in mdargs: 

109 if isinstance(n, int): 

110 if n >= arglength: 

111 continue 

112 entry = args[n] 

113 is_arg = True 

114 elif isinstance(n, str): 

115 try: 

116 entry = kwargs[n] 

117 except KeyError: 

118 continue 

119 is_arg = False 

120 if hasattr(entry, "__iter__"): 

121 # Create now a copy of the given array and manipulate then 

122 # the entries directly. 

123 if is_arg: 

124 args = list(args) 

125 args[n] = structure_copy(entry) 

126 else: 

127 kwargs[n] = structure_copy(entry) 

128 result = apply_on_element(wrapper, args, kwargs, n) 

129 return result 

130 return f(*args, **kwargs) 

131 return wrapper