编程|胡渊鸣:import一个“太极”库 让Python代码提速100倍!( 二 )


文章图片

胡渊鸣电脑上的程序最快做到了0.9秒内完成 , 而换成用NumPy来实现 , 则需要476秒 , 差异达到了超500倍!
最后 , 我们再来一个不一样的例子 。
反应 - 扩散方程 , 效果惊人
自然界中 , 总有一些动物身上长着一些看起来无序但实则并非完全随机的花纹 。
编程|胡渊鸣:import一个“太极”库 让Python代码提速100倍!
文章图片

图灵机的发明者艾伦·图灵是第一个提出模型来描述这种现象的人 。
在该模型中 , 两种化学物质(U和V)来模拟图案的生成 。这两者之间的关系类似于猎物和捕食者 , 它们自行移动并有交互:
最初 , U和V随机分布在一个域上;
在每个时间步 , 它们逐渐扩散到邻近空间;
当U和V相遇时 , 一部分U被V吞噬 。因此 , V的浓度增加;
为了避免U被V根除 , 我们在每个时间步添加一定百分比 (f) 的U并删除一定百分比 (k) 的V 。
上面这个过程被概述为“反应-扩散方程”:
编程|胡渊鸣:import一个“太极”库 让Python代码提速100倍!
文章图片

其中有四个关键参数:Du(U的扩散速度) , Dv(V的扩散速度) , f(feed的缩写 , 控制U的加入)和k(kill的缩写 , 控制V的去除) 。
如果Taichi中实现这个方程 , 首先创建网格来表示域 , 用vec2表示每个网格中U, V的浓度值 。
拉普拉斯算子数值的计算需要访问相邻网格 。为了避免在同一循环中更新和读取数据 , 我们应该创建两个形状相同的网格W x H×2 。
每次从一个网格访问数据时 , 我们将更新的数据写入另一个网格 , 然后切换下一个网格 。那么数据结构设计就是这样:
编程|胡渊鸣:import一个“太极”库 让Python代码提速100倍!
文章图片

一开始 , 我们将U在网格中的浓度设置为 1 , 并将V放置在50个随机选择的位置:
编程|胡渊鸣:import一个“太极”库 让Python代码提速100倍!
文章图片

那么实际计算就可以用不到10行代码完成:
@ti.kerneldef compute(phase: int):    for i, j in ti.ndrange(W, H):        cen = uv[phase, i, j]        lapl = uv[phase, i + 1, j] + uv[phase, i, j + 1] + uv[phase, i - 1, j] + uv[phase, i, j - 1] - 4.0 * cen        du = Du * lapl[0] - cen[0] * cen[1] * cen[1] + feed * (1 - cen[0])        dv = Dv * lapl[1] + cen[0] * cen[1] * cen[1] - (feed + kill) * cen[1]        val = cen + 0.5 * tm.vec2(du, dv)        uv[1 - phase, i, j] = val在这里 , 我们使用整数相位(0或1)来控制我们从哪个网格读取数据 。
最后一步就是根据V的浓度对结果进行染色 , 就可以得到这样一个效果惊人的图案: