Note

Go to the end to download the full example code.

AutoMod: Automatic differentiation

Automatically generate sensitivities using automatic differentiation

Using automatic differentiation (jax or autograd) the derivatives are generated of a user-defined function. Sensitivities do not need to be implemented by hand, but are automatically generated using the pymoto.AutoMod module.

10 import numpy as np
11 import pymoto as pym
12
13
14 def my_new_function(x1, A, x2, val=1.3):
15         """ User-defined functionality which is to be differentiated.
16         This function calculates a response based on a multiple input values, here for example 3.
17         It is possible to use both vector and scalar inputs, as well as complex numbers.
18
19         Args:
20             x1: First vector
21             A: Matrix
22             x2: Second vector
23
24         Returns:
25             The results of the calculation
26         """
27         # E.g. test for incorrect data
28         if x1 is None or x2 is None:
29             raise RuntimeError("You forgot to set x1 and/or x2")
30
31         # Calculate response
32         v = x1 @ (A @ x2) + x1 * x2 + val
33
34         # Return the results
35         return v
36
37
38 if __name__ == "__main__":
39     print(__doc__)
40     print("_" * 80)
41
42     # Initialize signals we want to use in our program
43     x1 = pym.Signal("x1")
44     A = pym.Signal("A")
45     x2 = pym.Signal("x2")
46
47     which = 'vector'  # Choose 'scalar' or 'vector' to test the module
48     complex = True  # Choose True or False to test complex numbers
49     if which == 'scalar':
50         x1.state = 2.0
51         x2.state = 3.0
52         if complex:
53             x1.state = x1.state + 1.0*1j
54             x2.state = x2.state + 4.0*1j
55     elif which == 'vector':
56         x1.state = np.random.rand(4)
57         x2.state = np.random.rand(4)
58         A.state = np.random.rand(4, 4)
59         if complex:
60             x1.state = x1.state + 1j*np.random.rand(4)
61             x2.state = x2.state + 1j*np.random.rand(4)
62             A.state = A.state + 1j*np.random.rand(4, 4)
63
64     # The module can be instantiated using the constructor
65     y = pym.AutoMod(my_new_function)(x1, A, x2)
66     y.tag = "y"
67     print(f"The response is {y.tag} = {y.state}")
68
69     # Check the response values; they are the same as the original function
70     y_chk = my_new_function(x1.state, A.state, x2.state)
71     print(f"The expected response is = {y_chk}")
72     assert np.allclose(y.state, y_chk)
73
74     # Check the sensitivities; these are automatically calculated using the autodiff module
75     pym.finite_difference([x1, A, x2], y)

Gallery generated by Sphinx-Gallery