Standard Map

标准映射 (Standard Map) 是研究哈密顿系统里的混沌过程常用工具，其背景可以参考 Standard Map。具体的迭代方程为：

$$ I_{n+1} \leftarrow I_n + K \sin(\theta_n), $$$$ \theta_{n+1} \leftarrow \theta_n + I_{n+1}.$$

这里 $I_{n+1}$ 和 $\theta_{n+1}$ 都会对 $2\pi$ 取模。

标准版程序

import numpy as np
from numpy import pi
import matplotlib.pyplot as plt
%matplotlib inline

def smap(In, thetan, K, depth):
    Iv = np.empty(depth)
    thetav = np.empty(depth)

    for i in range(depth):
        Iv[i] = In
        thetav[i] = thetan
        
        Inp1 = In + K*np.sin(thetan)
        Inp1 = Inp1 % (2*np.pi)
        
        thetanp1 = thetan + Inp1

        In = Inp1
        thetan = thetanp1 % (2*np.pi)
        
    return Iv, thetav

def smap_set(I0, theta0, width, grid, K, depth):
    I0v = np.linspace(I0-width/2.0, I0+width/2.0, grid)
    theta0v = np.linspace(theta0-width/2.0, theta0+width/2.0, grid)        

    I = np.empty((grid,grid,depth))
    theta = np.empty(I.shape)

    for i in range(grid):
        for j in range(grid):
            I00 = I0v[i]
            theta00 = theta0v[j]

            I1,theta1 = smap(I00, theta00, K, depth)
            
            I[i,j] = I1
            theta[i,j] = theta1

    return I, theta

if __name__ == '__main__':
    
    I0 = pi
    theta0 = pi
    width = 4*pi
    grid = 6
    K = 1.2
    depth = 1000

    I, theta = smap_set(I0, theta0, width, grid, K, depth)

    plt.close()

    fig,ax=plt.subplots(figsize=(8,6))

    for i in range(grid):
        for j in range(grid):
            ax.scatter(theta[i,j], I[i,j], marker='.', s=1)

    ax.set_xlim(0,2*np.pi)
    ax.set_ylim(0,2*np.pi)
    ax.set_xlabel(r'$\theta$')
    ax.set_ylabel(r'$I$')
            
    plt.show()

numba 加速版程序

import numpy as np
from numpy import pi
import matplotlib.pyplot as plt
import numba as nb
%matplotlib inline

@nb.njit
def smap_nb(In, thetan, K, depth):
    Iv = np.empty(depth)
    thetav = np.empty(depth)

    for i in range(depth):
        Iv[i] = In
        thetav[i] = thetan
        
        Inp1 = In + K*np.sin(thetan)
        Inp1 = Inp1 % (2*np.pi)
        
        thetanp1 = thetan + Inp1

        In = Inp1
        thetan = thetanp1 % (2*np.pi)
        
    return Iv, thetav

@nb.njit
def smap_set_nb(I0, theta0, width, grid, K, depth):
    I0v = np.linspace(I0-width/2.0, I0+width/2.0, grid)
    theta0v = np.linspace(theta0-width/2.0, theta0+width/2.0, grid)        

    I = np.empty((grid,grid,depth))
    theta = np.empty(I.shape)

    for i in range(grid):
        for j in range(grid):
            I00 = I0v[i]
            theta00 = theta0v[j]

            I1,theta1 = smap_nb(I00, theta00, K, depth)
            
            I[i,j] = I1
            theta[i,j] = theta1

    return I, theta

if __name__ == '__main__':
    
    I0 = pi
    theta0 = pi
    width = 4*pi
    grid = 6
    K = 1.2
    depth = 1000

    I, theta = smap_set_nb(I0, theta0, width, grid, K, depth)

    plt.close()

    fig,ax=plt.subplots(figsize=(8,6))

    for i in range(grid):
        for j in range(grid):
            ax.scatter(theta[i,j], I[i,j], marker='.', s=1)

    ax.set_xlim(0,2*np.pi)
    ax.set_ylim(0,2*np.pi)
    ax.set_xlabel(r'$\theta$')
    ax.set_ylabel(r'$I$')
            
    plt.show()

速度比较

%timeit smap_set(I0, theta0, width, grid, K, depth)

91.7 ms ± 1.11 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit smap_set_nb(I0, theta0, width, grid, K, depth)

1.34 ms ± 1.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

在这个例子里，numba 实现了约 90 倍的提速。