Numerical calculation of differential equations with TensorFlow 2.0

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:

Recommended Posts

Numerical calculation of differential equations with TensorFlow 2.0
Numerical analysis of ordinary differential equations with Scipy's odeint and ode
Numerical calculation with Python
TensorFlow Tutorial-Partial Differential Equations (Translation)
Numerical calculation of partial differential equations with singularity (for example, asymptotic behavior analysis of Hardy-Hénon type heat equation)
[Numerical calculation method, python] Solving ordinary differential equations by Eular method
Error-free calculation with big.Float of golang
[Science / technical calculation by Python] Numerical solution of first-order ordinary differential equations, initial value problem, numerical calculation
[Scientific / technical calculation by Python] Numerical solution of second-order ordinary differential equations, initial value problem, numerical calculation
(First post) A story about numerical calculation of influenza and new pneumonia coronavirus with Tensorflow
Find the numerical solution of the second-order ordinary differential equation with scipy
Solve the initial value problem of ordinary differential equations with JModelica
1. Statistics learned with Python 1-3. Calculation of various statistics (statistics)
Real-time calculation of mean values with coroutines
I replaced the numerical calculation of Python with Rust and compared the speed
Shuffle hundreds of thousands of images evenly with tensorflow.
Solving ordinary differential equations with Python ~ Universal gravitation
[Scientific / technical calculation by Python] Solving the boundary value problem of ordinary differential equations in matrix format, numerical calculation
[Scientific / technical calculation by Python] Solving second-order ordinary differential equations by Numerov method, numerical calculation
[Scientific / technical calculation by Python] Numerical calculation to find the value of derivative (differential)
1. Statistics learned with Python 1-2. Calculation of various statistics (Numpy)
Numpy leave? !! Partial differential of matrix with Sympy
Zundokokiyoshi with TensorFlow
Breakout with Tensorflow
Sequential calculation of mean value with online algorithm
Calculation of mutual information (continuous value) with numpy
[Scientific / technical calculation by Python] Numerical solution of one-dimensional and two-dimensional wave equations by FTCS method (explicit method), hyperbolic partial differential equations
Numerical calculation of compressible fluid by finite volume method
Performs high-speed calculation of only specific descriptors with mordred
Maximum likelihood estimation of mean and variance with TensorFlow
Solve simultaneous ordinary differential equations with Python and SymPy.
Start numerical calculation in Python (with Homebrew and pip)
Build a numerical calculation environment with pyenv and miniconda3
Numerical summary of data
Reading data with TensorFlow
Kyotei forecast with TensorFlow
Try regression with TensorFlow
I tried to find the average of the sequence with TensorFlow
List of array (ndarray) operations of Python's numerical calculation library "Numpy"
Numerical approximation method when the calculation of the derivative is troublesome
Numerical calculation with Python
Start numerical calculation in Python (with Homebrew and pip)
Build a numerical calculation environment with pyenv and miniconda3
I replaced the numerical calculation of Python with Rust and compared the speed
Perform DFT calculation with ASE and GPAW
Numerical calculation of differential equations with TensorFlow 2.0