Autodif workshop demo¶

This demo is developed for the GODOT workshop and demonstrates how to carry out simple computations with GODOT automatic differentiation module.

In [ ]:

import numpy as np

We can write simple functions using numpy

In [ ]:

def f(x, y):
    return np.sin(x) / (x + y)

a = 1.0
b = 2.0
c = f(a, b)

print(c)

Elementary Example¶

We can create two scalar ad variables in godot:

In [ ]:

from godot.core import autodif as ad

a = ad.Scalar(1.0, "a")
b = ad.Scalar(2.0, "b")

print(a)

We can do operations between these

In [ ]:

a + b

In [ ]:

a * b

And call mathematical functions on them

In [ ]:

ad.sin(a)

In [ ]:

ad.atan2(a, b)

We can redefine our function with ad support.

In [ ]:

# def f(x, y):
#     return np.sin(x) / (x + y)

def g(x, y):
    return ad.sin(x) / (x + y)

c = g(a, b)

print(c)

You can also call the new function with normal floating point.

In [ ]:

g(1.0, 2.0)

We can inspect the value and gradient object separately as well as specific partials

In [ ]:

print(c.value())

print(c.gradient())

print(c.gradient()["a"])

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

In [ ]:

def num_diff(x, y, h):
    return (f(x + h, y) - f(x - h, y)) / (2 * h)

eps = [10.0**(-n) for n in np.arange(0.5, 15, 0.1)]
res = [abs(num_diff(a.value(), b.value(), h) - c.gradient()["a"][0]) for h in eps]

import matplotlib.pyplot as plt
plt.loglog(eps, res)
plt.xlabel("Numerical difference step")
plt.ylabel("Relative error")
plt.grid()
plt.show()

We can also do the same with vectors:

In [ ]:

v = ad.Vector([1, 2, 3], "v")
w = ad.Vector([-2, 1, -1], "w")

print(v)

We can do operations like dot product

In [ ]:

ad.dot(v, w)

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} + kx = 0$$

We write this as a first order vector ODE:

$$\dot{x} = v$$$$\dot{v} = - bv - kx$$

In [ ]:

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 [ ]:

# 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

# Fixed step size propagation of 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 [ ]:

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 [ ]:

x_final = x[-1]
print(x_final)

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

In [ ]:

k = ad.Scalar(1.0, "k")
b = ad.Scalar(0.1, "b")
t, x = propagate(x0, f, 0, 100, h)
x_final = x[-1]
print(x_final)

In [ ]:

# 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(["k", "b"])
print( f"The STM with respect to parameters is \n{psi}\n" )

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

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 [ ]:

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)