Using the Autodif module

The autodif module is designed to make gradient computation trivial for any mathematical function of an arbitrary number of parameters.

To use the autodif module from python, we have to import it.

In [1]:
import godot.core.autodif as ad

import numpy as np

Elementary Example

We construct 2 autodiff scalars:

In [2]:
p = ad.Scalar(1.0, "p")
q = ad.Scalar(1.0, "q")

Define a function

In [3]:
def f( p, q ):
    return ad.sin(p) / (p + q)

Now evaluate the function, to get the value and the gradients:

In [4]:
r = f(p, q)
print(r)
         Value |             p |             q
  4.207355e-01 |  5.978341e-02 | -2.103677e-01

We get the value and the gradients as numpy arrays as follows:

In [5]:
r.value()
Out[5]:
0.42073549240394825
In [6]:
r.at(["p", "q"])
Out[6]:
array([ 0.05978341, -0.21036775])

To test, we can numerically compute the derivative with respect to p

In [7]:
df_dp_ad = r.at(["p"])[0]

def num_diff(h):
    r1 = f(p + h , q).value()
    r2 = f(p - h, q ).value()
    return (r1 - r2 ) / (2 * h)

def compare_result(h):
    df_dp_num = num_diff(h)
    print(f"Epsilon = {h:6.3g}    Numerical = {df_dp_num:10.6f}    Autodiff = {df_dp_ad:10.6f}    Error = {100 - df_dp_ad / df_dp_num*100:8.2f} %")
    

compare_result(1.0)
compare_result(0.1)
compare_result(0.001)
compare_result(1e-6)
compare_result(1e-12)
compare_result(1e-15)
Epsilon =      1    Numerical =   0.151550    Autodiff =   0.059783    Error =    60.55 %
Epsilon =    0.1    Numerical =   0.060536    Autodiff =   0.059783    Error =     1.24 %
Epsilon =  0.001    Numerical =   0.059783    Autodiff =   0.059783    Error =     0.00 %
Epsilon =  1e-06    Numerical =   0.059783    Autodiff =   0.059783    Error =    -0.00 %
Epsilon =  1e-12    Numerical =   0.059758    Autodiff =   0.059783    Error =    -0.04 %
Epsilon =  1e-15    Numerical =   0.111022    Autodiff =   0.059783    Error =    46.15 %

Numerical Integration

We can also use autodiff in more complex situations. As an example, we define a damped oscillator ODE

$$\ddot{x} + b \dot{x} + k * x = 0$$

We write this as a first order vector ODE:

$$\dot{x} = v$$$$\dot{v} = - b * v - k * x$$
In [8]:
k = 1.0
b = 0.1
def f(t, x):
    v = x[1]
    a =  -k * x[0] - b * x[1]
    return ad.concatenate(v, a)

We write a Runge Kutte 4th order stepper, and integrate the equation above:

In [9]:
# Runge-Kutte Order 4 stepper
def rk4(t ,x, f, h):
    k1 = f(t, x)
    k2 = f(t + h/2, x + h * k1 / 2)
    k3 = f(t + h/2, x + h * k2 / 2)
    k4 = f(t + h, x + h * k3)
    return x + h * (k1 + k2 + k3 + k4) / 6

# propagate the oscillator state
def propagate(x0, f, t0, t1, h):
    t = np.arange(t0, t1, h)
    x = [x0]
    for i in t[1:]:
        x.append(rk4(i, x[-1], f, h))
    return t, x

h = 0.1
x0 = ad.Vector([0 , 1], "x0")
t, x = propagate(x0, f, 0, 100, h)

Now we plot the results

In [10]:
import matplotlib.pyplot as plt
plt.rcParams['figure.figsize'] = [12, 6]

def plot(times,state):
    pos = [s[0] for s in state]
    vel = [s[1] for s in state]
    fig, (ax1, ax2) = plt.subplots(1,2)
    ax1.grid()
    ax1.plot(times,[p.value() for p in pos],label='pos')
    ax1.plot(times,[p.value() for p in vel],label='vel')
    ax1.set_xlabel('Time')
    ax1.axhline(0.1, ls='--', lw='1', color='r')
    ax1.axvline(20, ls='--', lw='1', color='r')
    ax1.legend()

    ax2.grid()
    ax2.plot(
        [p.value() for p in pos],
        [p.value() for p in vel]
    )
    ax2.set_xlabel('Pos')
    ax2.set_ylabel('Vel')
    ax2.set_aspect('equal', adjustable='box')
plot(t,x)

But we also have the state transition matrix computed:

In [11]:
x_final = x[-1]
print(x_final)
         Value |         x0[0]         x0[1]
 -4.961353e-03 | -1.103039e-02 -4.961353e-03
 -1.053425e-02 |  4.961353e-03 -1.053425e-02

Now, I redefine b to be an autodiff variable and recompute:

In [12]:
b = ad.Scalar(0.1, "b")
t, x = propagate(x0, f, 0, 100, h)
x_final = x[-1]
print(x_final)
         Value |         x0[0]         x0[1] |             b
 -4.961353e-03 | -1.103039e-02 -4.961353e-03 |  1.944537e-01
 -1.053425e-02 |  4.961353e-03 -1.053425e-02 |  3.368484e-01

In [13]:
# print the value
value = x_final.value()
print( f"The value is: \n{value}\n" )

# print the state transition matrix
phi = x_final.at( 'x0' )
print( f"The STM is \n{phi}\n" )

# print the state transition matrix
psi = x_final.at( 'b' )
print( f"The STM with respect to parameters is \n{psi}\n" )

# print the full state transition matrix
phi2 = x_final.at( ['x0', 'b'] )
print( f"The full STM is \n{phi2}\n" )
The value is: 
[-0.00496135 -0.01053425]

The STM is 
[[-0.01103039 -0.00496135]
 [ 0.00496135 -0.01053425]]

The STM with respect to parameters is 
[[0.19445371]
 [0.33684845]]

The full STM is 
[[-0.01103039 -0.00496135  0.19445371]
 [ 0.00496135 -0.01053425  0.33684845]]

We can use this to compute an impulse at some time to control the state at a future time, for example.

We can perform optimisations, estimations, etc.

As a simple example, we control the damping parameter 'b' to fix the amplitude at time t = 20.

In [14]:
for i in range(5):

    # get the index of t = 20
    index = int(20/h)
    x[index]

    # fix the target amplitude value
    x_target = 0.1

    # extract the gradient of the amplitude at the index with respect to b.
    dx_db = x[index][0].at('b')

    # perform the linear update to b
    db = ( x_target - x[index].value() ) / dx_db
    b+=db[0]

    # re-propagate and plot
    t,x = propagate(x0,f,0,100,h)

    # check the value of x at the index
    print(f"The amplitude is {x[index].value()[0]} at t = 20")

plot(t,x)
The amplitude is 0.14279024048657457 at t = 20
The amplitude is 0.1053282787170966 at t = 20
The amplitude is 0.10011824624432575 at t = 20
The amplitude is 0.10000006193695035 at t = 20
The amplitude is 0.10000000000001473 at t = 20
In [ ]: