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
« prev ^ index » next coverage.py v6.5.0, created at 2022-12-07 00:47 -0500
1"""Automatic Differentiation object"""
3from typing import Callable
5import numpy as np
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
12class Evaluator:
13 """AD evaluation object
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 """
22 def __init__(self, fn: Callable):
23 self.fn = fn
25 self.input_dim = None
26 self.output_dim = None
28 def eval(self, pt, **kwargs):
29 """Perform AD on this Evaluator's function, at this point
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
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)
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)}.')
47 shape = np.shape(pt)
48 pt_shape = 1 if shape == () else shape[0]
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}')
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')
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'
73 # Set the output dimension (and ensure the function is valid)
74 try:
75 fn_output = self.fn(pt)
77 if not isscalar_or_array(fn_output):
78 print(type(fn_output), fn_output)
79 raise TypeError('Output must not be None')
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
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`')
98 return differentiator(func=self.fn, pt=pt, seed_vector=kwargs['seed_vector'])