Introduction

TensorFlow 2.0 is a very powerful framework that can be used for machine learning such as deep learning. TensorFlow treats numerical algorithms in the form of "calculation graphs" abstractly.

It's like an execution plan in SQL. This seems to be able to separate "definition of algorithm" and "execution of it". If used properly, it will be very useful not only for machine learning but also for general numerical calculations.

I tried it

Let's solve the Lotka-Volterra equation this time! An equation that models fluctuations in predator-prey and prey populations. It is also famous for reproducing the "periodic solution" that can be seen in nature.

import tensorflow as tf
import matplotlib.pyplot as plt
import timeit

a = 0.2
b = 0.04
c = 0.4
d = 0.05

d_esa = lambda esa, hoshokusha: a*esa - b*esa*hoshokusha
d_hoshokusha = lambda esa, hoshokusha: -c*hoshokusha + d*esa*hoshokusha

def runge_kutta(f, x, y, dt):
    k1 = dt*f(x,y)
    k2 = dt*f(x+0.5*k1, y+0.5*k1)
    k3 = dt*f(x+0.5*k2, y+0.5*k2)
    k4 = dt*f(x+k3, y+k3)
    return (k1+2.0*k2+2.0*k3+k4)/6.0

@tf.function
def lotka_volterra(t_max, t_num, esa_init, hoshokusha_init):
    esa_result = []
    hoshokusha_result = []
    t_result = []

    t = 0.0
    esa = esa_init
    hoshokusha = hoshokusha_init
    
    esa_result.append(esa)
    hoshokusha_result.append(hoshokusha)
    t_result.append(t)

    dt = t_max / t_num

    while t < t_max:
        t += dt
        esa += runge_kutta(d_esa, esa, hoshokusha, dt)
        hoshokusha += runge_kutta(d_hoshokusha, esa, hoshokusha, dt)

        esa_result.append(esa)
        hoshokusha_result.append(hoshokusha)
        t_result.append(t)
    
    return esa_result, hoshokusha_result, t_result

# warm up！！！！！！
esa_result, hoshokusha_result, t_result = lotka_volterra(100.0, 2000, 10, 1)
print(timeit.timeit(lambda: lotka_volterra(100.0, 2000, 10, 1), number=1))

plt.plot(t_result, esa_result)
plt.plot(t_result, hoshokusha_result)
plt.show()

By adding the tf.function decorator, the internal calculation will be calculated as a graph on TensorFlow. Starting with Tensorflow 2.0, it's also possible to "write with Python code" to some extent. I'm grateful. In this case, you can immediately execute it as "pure Python code" by taking tf.function.

In my personal environment, adding tf.function slowed down the execution time by about 4 times. (CPU) Of course, this time it's only overhead, so it's natural (´ ・ ω ・ `)

in conclusion

How was it?

It turns out that you don't need to use TensorFlow 2.0 at all for Lotka-Volterra equations!

TensorFlow may be useful for numerical calculations that are considerably heavier than this time, such as those that make heavy use of matrix (tensor) operations such as partial differential equations.

I tried to do it, but it didn't work because it was a secret with Santa! Have a nice year: dogeza:

[PYTHON] Numerical calculation of differential equations with TensorFlow 2.0

Introduction

I tried it

in conclusion