Coverage for /usr/local/lib/python3.10/dist-packages/Adifpy/differentiate/evaluator.py: 98%

51 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2022-12-07 00:47 -0500

1"""Automatic Differentiation object""" 

2 

3from typing import Callable 

4 

5import numpy as np 

6 

7from Adifpy.differentiate.forward_mode import forward_mode 

8from Adifpy.differentiate.reverse_mode import reverse_mode 

9from Adifpy.differentiate.helpers import isscalar_or_array, isscalar 

10 

11 

12class Evaluator: 

13 """AD evaluation object 

14 

15 >>> my_evaluator = Evaluator(lambda x: x * x) 

16 >>> my_evaluator.eval(1) 

17 (1, 2) 

18 >>> my_evaluator.eval(3) 

19 (9, 6) 

20 """ 

21 

22 def __init__(self, fn: Callable): 

23 self.fn = fn 

24 

25 self.input_dim = None 

26 self.output_dim = None 

27 

28 def eval(self, pt, **kwargs): 

29 """Perform AD on this Evaluator's function, at this point 

30 

31 Args: 

32 pt (float | iterable): the point or vector at which to evaluate the function 

33 seed_vector (iterable, optional): the seed vector, if the function has vector input 

34 force_mode (str, optional): either 'forward' or 'reverse' for forcing AD mode 

35 

36 Returns: 

37 If the function's output space is R, a tuple of the value and directional derivative. 

38 Otherwise, a tuple of two lists: the values and directional derivatives for each component 

39 """ 

40 if isinstance(pt, list): 

41 pt = np.array(pt) 

42 

43 # Ensure the evaluation point is valid 

44 if not isscalar_or_array(pt): 

45 raise TypeError(f'Evaluation point must be a scalar or NumPy array, not {type(pt)}.') 

46 

47 shape = np.shape(pt) 

48 pt_shape = 1 if shape == () else shape[0] 

49 

50 if self.input_dim is None: 

51 self.input_dim = pt_shape 

52 elif self.input_dim != pt_shape: 

53 print(f'WARNING: Expected point in R{self.input_dim}, but got point in R{pt_shape}') 

54 

55 # Ensure that a seed vector is provided for vector functions 

56 if self.input_dim > 1 and 'seed_vector' not in kwargs: 

57 raise AttributeError('For vector functions, `seed_vector` argument is required') 

58 elif 'seed_vector' not in kwargs: 

59 # Set the default seed vector for functions that map from R 

60 kwargs['seed_vector'] = 1 

61 elif self.input_dim > 1: 

62 # Ensure the seed vector is valid (and throw an error if it is not) 

63 kwargs['seed_vector'] = np.array(kwargs['seed_vector']) 

64 else: 

65 raise TypeError('seed_vector argument should not be provided for functions from R -> R') 

66 

67 # Ensure the seed vector is of the expected dimensionality 

68 shape = np.shape(kwargs['seed_vector']) 

69 seed_dim = 1 if shape == () else shape[0] 

70 assert seed_dim == self.input_dim, \ 

71 f'Evaluation point has {self.input_dim} dimensions, but seed vector has {seed_dim} dimensions' 

72 

73 # Set the output dimension (and ensure the function is valid) 

74 try: 

75 fn_output = self.fn(pt) 

76 

77 if not isscalar_or_array(fn_output): 

78 print(type(fn_output), fn_output) 

79 raise TypeError('Output must not be None') 

80 

81 self.output_dim = 1 if isscalar(fn_output) else len(fn_output) 

82 except Exception as error: 

83 raise RuntimeError('Evaluator function failed') from error 

84 

85 # Decide which AD mode to use, either depending on forced user input or optimized for performance 

86 differentiator = forward_mode 

87 if 'force_mode' in kwargs: 

88 match kwargs['force_mode']: 

89 case 'forward': 

90 differentiator = forward_mode 

91 case 'reverse': 

92 if self.input_dim != 1: 

93 raise NotImplementedError('Reverse mode only supports functions from R -> R') 

94 differentiator = reverse_mode 

95 case _: 

96 raise ValueError('`force_mode` argument must be either `forward` or `reverse`') 

97 

98 return differentiator(func=self.fn, pt=pt, seed_vector=kwargs['seed_vector'])