这个神器，让你的代码运行快上100倍！

文 | 闲欢

来源：Python 技术「ID: pythonall」

Python 已经得到了全球程序员的喜爱，连续多期稳坐编程语言排行榜第一把交椅。但是还是遭到一些人的诟病，原因之一就是认为它运行缓慢。

于是，大家都在想尽各种办法来提高 Python 代码的运行速度，大多数体现在写代码的习惯优化以及代码优化上。但是平时写代码注意太多这些方面可能会有糟糕的体验，甚至会不利于我们的工作效率。

要是有一款能够自动优化我们代码的神器该有多好啊！

今天就给大家带来这样的一款神器——taichi！

taichi 是何方神圣

Taichi 起步于 MIT 的计算机科学与人工智能实验室（CSAIL），设计初衷是便利计算机图形学研究人员的日常工作，帮助他们快速实现适用于 GPU 的视觉计算和物理模拟算法。

说人话就是 Taichi 是一个基于 Python 的领域特定语言，专为高性能能并行计算设计。

本来是服务于学术界的一款 DSL ，但是我们也可以拿来用在我们这些凡夫俗子的代码中（虽然有点大材小用）！

安装

Taichi 是一个 PyPI 包，所以使用 pip 命令即可安装：

pip install taichi

注意 taichi 安装的先决条件是：

• Python: 3.7/3.8/3.9/3.10 (64-bit)

• OS: Windows, OS X, and Linux (64-bit)

在使用命令安装的时候，如果遇到错误，可以使用管理员模式命令行进行安装。

一个小栗子

我们先来用一个小栗子，感受一下它的鬼斧神工！

import time
 
def is_prime(n):
    result = True
    for k in range(2, int(n**0.5) + 1):
        if n % k == 0:
            result = False
            break
    return result

def count_primes(n: int) -> int:
    count = 0
    for k in range(2, n):
        if is_prime(k):
            count += 1
    
    return count

t0 = time.time()
print(count_primes(100000))
t1 = time.time()

print(t1-t0)

这个是我们以前经常用来做例子的统计质数个数。求100000以内速度比较快，但是到了1000000，运行时间就明显慢了下来，竟然需要3.38秒。

Python 的大型 for 循环或嵌套 for 循环总是导致运行时性能不佳。

我们只需导入 Taichi 或切换到 Taichi 的 GPU 后端，就能看到整体性能的大幅提升：

import time
import taichi as ti

ti.init()

@ti.func
def is_prime(n):
    result = True
    for k in range(2, int(n**0.5) + 1):
        if n % k == 0:
            result = False
            break
    return result

@ti.kernel
def count_primes(n: int) -> int:
    count = 0
    for k in range(2, n):
        if is_prime(k):
            count += 1
    
    return count

t0 = time.time()
print(count_primes(1000000))
t1 = time.time()

print(t1-t0)

在这里，我们只需要引入 taichi 库，然后加两个注解，速度直接飙到了0.1秒，速度提升了30多倍。如果我们把数字再扩大，速度提升会更明显！

没有使用之前，统计10000000以内质数使用 90 秒，使用之后，并且更改为 GPU 运行，使用 0.1秒。

我们还可以将 Taichi 的后端从 CPU 更改为 GPU 运行：

ti.init(arch=ti.gpu)

用 Taichi 进行物理模拟

上面的动图很好地模拟了一块布料落到一个球体上。动图中的布料建模使用了弹簧质点系统，其中包含 10,000 多个质点和大约 100,000个弹簧。模拟如此大规模的物理系统并实时渲染绝不是一项容易的任务。

Taichi 让物理模拟程序变得更易读和直观，同时仍然达到与 C++ 或 CUDA 相当的性能。只需拥有基本 Python 编程技能，就可以使用 Taichi 用更少的代码编写高性能并行程序，从而关注较高层次的算法本身，把诸如性能优化的任务交由 Taichi 处理。

我们直接上源代码：

import taichi as ti
ti.init(arch=ti.vulkan)  # Alternatively, ti.init(arch=ti.cpu)

n = 128
quad_size = 1.0 / n
dt = 4e-2 / n
substeps = int(1 / 60 // dt)

gravity = ti.Vector([0, -9.8, 0])
spring_Y = 3e4
dashpot_damping = 1e4
drag_damping = 1

ball_radius = 0.3
ball_center = ti.Vector.field(3, dtype=float, shape=(1, ))
ball_center[0] = [0, 0, 0]

x = ti.Vector.field(3, dtype=float, shape=(n, n))
v = ti.Vector.field(3, dtype=float, shape=(n, n))

num_triangles = (n - 1) * (n - 1) * 2
indices = ti.field(int, shape=num_triangles * 3)
vertices = ti.Vector.field(3, dtype=float, shape=n * n)
colors = ti.Vector.field(3, dtype=float, shape=n * n)

bending_springs = False

@ti.kernel
def initialize_mass_points():
    random_offset = ti.Vector([ti.random() - 0.5, ti.random() - 0.5]) * 0.1

    for i, j in x:
        x[i, j] = [
            i * quad_size - 0.5 + random_offset[0], 0.6,
            j * quad_size - 0.5 + random_offset[1]
        ]
        v[i, j] = [0, 0, 0]


@ti.kernel
def initialize_mesh_indices():
    for i, j in ti.ndrange(n - 1, n - 1):
        quad_id = (i * (n - 1)) + j
        # 1st triangle of the square
        indices[quad_id * 6 + 0] = i * n + j
        indices[quad_id * 6 + 1] = (i + 1) * n + j
        indices[quad_id * 6 + 2] = i * n + (j + 1)
        # 2nd triangle of the square
        indices[quad_id * 6 + 3] = (i + 1) * n + j + 1
        indices[quad_id * 6 + 4] = i * n + (j + 1)
        indices[quad_id * 6 + 5] = (i + 1) * n + j

    for i, j in ti.ndrange(n, n):
        if (i // 4 + j // 4) % 2 == 0:
            colors[i * n + j] = (0.22, 0.72, 0.52)
        else:
            colors[i * n + j] = (1, 0.334, 0.52)

initialize_mesh_indices()

spring_offsets = []
if bending_springs:
    for i in range(-1, 2):
        for j in range(-1, 2):
            if (i, j) != (0, 0):
                spring_offsets.append(ti.Vector([i, j]))

else:
    for i in range(-2, 3):
        for j in range(-2, 3):
            if (i, j) != (0, 0) and abs(i) + abs(j) <= 2:
                spring_offsets.append(ti.Vector([i, j]))

@ti.kernel
def substep():
    for i in ti.grouped(x):
        v[i] += gravity * dt

    for i in ti.grouped(x):
        force = ti.Vector([0.0, 0.0, 0.0])
        for spring_offset in ti.static(spring_offsets):
            j = i + spring_offset
            if 0 <= j[0] < n and 0 <= j[1] < n:
                x_ij = x[i] - x[j]
                v_ij = v[i] - v[j]
                d = x_ij.normalized()
                current_dist = x_ij.norm()
                original_dist = quad_size * float(i - j).norm()
                # Spring force
                force += -spring_Y * d * (current_dist / original_dist - 1)
                # Dashpot damping
                force += -v_ij.dot(d) * d * dashpot_damping * quad_size

        v[i] += force * dt

    for i in ti.grouped(x):
        v[i] *= ti.exp(-drag_damping * dt)
        offset_to_center = x[i] - ball_center[0]
        if offset_to_center.norm() <= ball_radius:
            # Velocity projection
            normal = offset_to_center.normalized()
            v[i] -= min(v[i].dot(normal), 0) * normal
        x[i] += dt * v[i]

@ti.kernel
def update_vertices():
    for i, j in ti.ndrange(n, n):
        vertices[i * n + j] = x[i, j]

window = ti.ui.Window("Taichi Cloth Simulation on GGUI", (1024, 1024),
                      vsync=True)
canvas = window.get_canvas()
canvas.set_background_color((1, 1, 1))
scene = ti.ui.Scene()
camera = ti.ui.make_camera()

current_t = 0.0
initialize_mass_points()

while window.running:
    if current_t > 1.5:
        # Reset
        initialize_mass_points()
        current_t = 0

    for i in range(substeps):
        substep()
        current_t += dt
    update_vertices()

    camera.position(0.0, 0.0, 3)
    camera.lookat(0.0, 0.0, 0)
    scene.set_camera(camera)

    scene.point_light(pos=(0, 1, 2), color=(1, 1, 1))
    scene.ambient_light((0.5, 0.5, 0.5))
    scene.mesh(vertices,
               indices=indices,
               per_vertex_color=colors,
               two_sided=True)

    # Draw a smaller ball to avoid visual penetration
    scene.particles(ball_center, radius=ball_radius * 0.95, color=(0.5, 0.42, 0.8))
    canvas.scene(scene)
    window.show()

感兴趣的可以具体看看代码的实现过程，如果不加 taichi 库，这段代码运行起来会有点吃力，但是上了 taichi 之后，运行效果是如此丝滑！

总结

大家看这个 taichi 是不是会联想到“太极”二字？

没错，这个库是中国人发明的，它就是毕业于清华大学，后来去麻省理工学院进修的胡渊鸣！

他是太极图形创始人，北京太琦图形科技有限公司联合创始人、CEO。总之，他是个天才，让国人也在世界编程舞台上绽放异彩！

PS：Python技术交流群（技术交流、摸鱼、白嫖课程为主）又不定时开放了，感兴趣的朋友，可以在下方公号内回复：666，即可进入，一起 100 天计划！

老规矩，酱友们还记得么，右下角的 “在看” 点一下，如果感觉文章内容不错的话，记得分享朋友圈让更多的人知道！

【代码获取方式】

识别文末二维码，回复：闲欢