太極圖形60行代碼實現(xiàn)經(jīng)典論文，0.7秒搞定泊松盤采樣，比Numpy實現(xiàn)快100倍

蕭簫 2022-06-15 21:23:05 來源：量子位

編輯整理自 太極圖形
量子位 | 公眾號 QbitAI

由隨機均勻的點組成的圖案，在動植物身上已經(jīng)很常見了。

像楊梅、草莓、荔枝、紅毛丹這樣的水果，表面都有顆粒或者毛發(fā)狀的結(jié)構(gòu)，它們隨機、均勻地散布在水果表面：

類似的圖案在動物身上也有，比如大家都愛涮的毛肚：

同樣地，在計算機模擬下，也有不少場景需要在空間中隨機、均勻地生成點。

像生成動物毛發(fā)時的毛孔位置、多人對戰(zhàn)游戲中的玩家出生位置、生成森林時的樹木位置等等。

這些場景的共同特點是，都需要讓任何兩點之間的距離大于等于一個下界（這個下界是預設的，改變它就可以控制生成點之間的間隔）。

但如果直接使用完全隨機生成的點，大概率會獲得一個很不均勻的分布結(jié)果，有些地方“扎堆”、有些地方稀疏：

如果用這樣的點來模擬毛發(fā)等位置生成，效果就很差。

所以，需要在生成點的過程中加入一個距離判斷，來剔除那些不合要求的點。

此前，用NumPy生成這樣一個效果，往往需要70s左右，非常不劃算。

現(xiàn)在，太極圖形基于Taichi實現(xiàn)了一個超快算法，同樣的效果運行在單個CPU線程上，只需要0.7s就能生成這樣的圖案，快了100倍左右。

一起來看看他們是怎么做的。

采用Bridson算法實現(xiàn)

此前，有一種常見算法是dart throwing?（像一個人蒙上眼睛胡亂扔飛鏢的樣子）：

每次在區(qū)域內(nèi)隨機選擇一個點，并檢查該點與所有已經(jīng)得到的點之間是否存在“沖突”。

若該點與某個已得到的點的最小距離小于指定的下界，就拋棄這個點，否則這就是一個合格的點，把它加入已有點的集合。

重復這個操作直到獲得了足夠多的點，或者連續(xù)失敗了N次為止（N是某個設定的正整數(shù)）。

但這種算法的效率很低。

因為隨著得到的點的個數(shù)增加，沖突的概率越來越大，獲得新的點所需的時間也越來越長，每次比較當前點和所有已有點之間的距離也會降低效率。

相比之下，Bridson算法則要更加高效。

這個算法的原理來自于Robert Bridson發(fā)表于2007年的論文”Fast Poisson Disk Sampling in Arbitrary Dimensions”[1]（論文非常短，只有一頁A4紙），如果再去掉標題、引言的話，真正的算法內(nèi)容只有一小段話。

開頭這個動圖，演示了Bridson圓盤采樣算法在一個400×400的網(wǎng)格區(qū)域上的運行效果，算法嘗試獲得100K個均勻散布的點，實際生成了大約53.7K個：

這個動畫是使用Taichi生成的，運行在單個CPU線程上，除去編譯的時間計算，耗時僅在0.7s多一點，而同樣的代碼翻譯成NumPy要耗時70s左右。[2]

從上面的動畫效果可見，Bridson算法很像包菜的生長過程：我們從一個種子點開始，一層一層地向外添加新的點。

每一次我們添加的新的點，都位于最外層的點的周圍，并且盡可能地包住最外層。

為了避免每次都檢查和所有已有點之間的距離，Taichi采用了所謂網(wǎng)格化的技巧：

將整個空間劃分為網(wǎng)格，對一個需要檢查的點，只要找到它所在的網(wǎng)格，然后檢查它和臨近網(wǎng)格中的點之間的最小距離即可。

只要這個距離大于指定的下界，更遠處的點就不必再檢查了。這個技巧在圖形學和物理仿真中是非常常用的。

這個采樣過程很難并行化，因為當一個線程“偷偷”加入一個新的點的時候，會改變其它所有線程對距離的判斷。所以Taichi僅使用單個CPU線程來執(zhí)行這個算法：

ti.init(arch=ti.cpu)

上面的代碼中通過指定arch=ti.cpu來讓程序運行在CPU上。

你可能會想，既然是單線程+CPU，那為什么不直接寫純Python呢？別著急，我們的計算部分會放在ti.kernel函數(shù)中，這種函數(shù)并不運行在Python虛擬機中，而是會被Taichi編譯執(zhí)行，所以會比純Python的實現(xiàn)快很多倍！

在我們介紹Bridson算法的具體實現(xiàn)之前，你不妨猜猜這個Taichi程序包含多少行代碼？

安裝和導入Taichi

首先推薦大家使用最新的Taichi發(fā)布版本，這樣可以使用更豐富的功能，在不同平臺上的支持也更穩(wěn)定。截止本文寫作時最新版本是1.0.3:

pip install taichi==1.0.3

然后，在代碼開頭寫上：

import taichi as ti
import taichi.math as tm

這樣會導入Taichi以及Taichi的math模塊。math模塊除了包含常用的數(shù)學函數(shù)之外，還提供了非常方便的向量運算。

準備工作

在泊松采樣算法中，采樣點之間的距離有一個下界r。

我們假設整個區(qū)域是由N×N個同樣大小的方格組成的網(wǎng)格區(qū)域，使得每個小方格的對角線長度正好是r，即網(wǎng)格的邊長是r/√2。

于是任何小方格中至多包含一個點。如下圖所示：

這就是我們前面提到的網(wǎng)格化方法，即對于任何一個點p，設它所在的方格是D，則任何與p的距離小于等于r的點必然位于以D中心的、由5×5個方格組成的正方形區(qū)域中。

在檢查距離時，我們只要針對這個子區(qū)域進行計算即可。

我們用一個一維數(shù)組samples和一個N×N的二維數(shù)組grid來記錄已經(jīng)得到的采樣點：

1. samples保存當前所有已經(jīng)采樣點的坐標，它的每個元素是一個二維坐標（x,y）。
2. grid[i, j]是一個整數(shù)，它存儲的是第（i, j）個方格中采樣點在數(shù)組samples中的下標。grid[i, j] = -1表示這個方格中沒有采樣點。

于是我們的初始設置可以這樣寫：

grid_n = 400
res = (grid_n, grid_n)
dx = 1 / res[0]
inv_dx = res[0]
radius = dx * ti.sqrt(2)
desired_samples = 100000
grid = ti.field(dtype=int, shape=res)
samples = ti.Vector.field(2, float, shape=desired_samples)

這里網(wǎng)格大小設置為400×400，它占據(jù)的平面區(qū)域是[0,1]×[0,1]，所以網(wǎng)格的步長是dx = 1/400。采樣的最小間隔是每個小方格對角線的長度，即radius = sqrt(2)*dx。

我們把采樣點的目標個數(shù)設置為desired_examples=100000，這是一個目測值，因為400×400的網(wǎng)格包含160000個小方格，考慮到每個方格中至多只有一個點，我們能得到的滿足距離約束的點的最大數(shù)目肯定少于160000。

初始時網(wǎng)格中沒有任何點，所以需要將grid中的值都置為-1:

grid.fill(-1)

如何生成新的點

在加入新的隨機點時，我們總是從已有點的附近隨機選擇一個位置，然后比較它和已知點是否滿足最小距離約束，是的話就將其加入已有點，否則就將其拋棄然后重新選擇點。

這里需要注意的是：

1. 當一個已有點附近已經(jīng)被填滿時，我們后面再加入新的點時就不必考慮它的附近了，可以用一個下標head來記錄這一點。我們約定samples數(shù)組中下標< head的點附近都已經(jīng)被填滿，從而不必再考慮它們，只考慮下標>= head的點即可。初始時head = 0。
2. samples是一個長度為100K的數(shù)組，這不代表我們真的能取到這么多點，但具體個數(shù)是多少無法事先確定，所以我們還需要用一個下標tail來記錄目前已經(jīng)獲得的點的個數(shù)。初始時tail = 1，因為我們將選擇區(qū)域中心作為第一個點。當然這個初始點的位置可以是任意的。
3. 正如前面提到的，當我們檢查一個點p是否與已有點滿足最小距離約束時，沒有必要遍歷檢查所有的點。只要檢查以p所在方格為中心，由5×5個方格組成的正方形區(qū)域即可。

檢查一個點是否和已有點沖突的邏輯我們單獨寫成一個函數(shù)：

@ti.func
def check_collision(p, index):
    x, y = index
    collision = False
    for i in range(max(0, x - 2), min(grid_n, x + 3)):
        for j in range(max(0, y - 2), min(grid_n, y + 3)):
            if grid[i, j] != -1:
                q = samples[grid[i, j]]
                if (q - p).norm() < radius - 1e-6:
                    collision = True
    return collision

其中p是需要檢查點的坐標，index=(x, y)是p所在的方格的下標。

我們遍歷所有滿足x-2 <= i <= x+2和y-2 <= j <= y+2的下標(i, j)，檢查方格(i, j)中是否已經(jīng)有點，即 grid[i, j]是否等于-1。有的話它與p的距離是否小于radius，然后返回對應的判斷。

完成了準備工作，我們可以開始正式的循環(huán)了：

@ti.kernel
def poisson_disk_sample(desired_samples: int) -> int:
    samples[0] = tm.vec2(0.5)
    grid[int(grid_n * 0.5), int(grid_n * 0.5)] = 0
    head, tail = 0, 1
    while head < tail and head < desired_samples:
        source_x = samples[head]
        head += 1

        for _ in range(100):
            theta = ti.random() * 2 * tm.pi
            offset = tm.vec2(tm.cos(theta), tm.sin(theta)) * (1 + ti.random()) * radius
            new_x = source_x + offset
            new_index = int(new_x * inv_dx)

            if 0 <= new_x[0] < 1 and 0 <= new_x[1] < 1:
                collision = check_collision(new_x, new_index)
                if not collision and tail < desired_samples:
                    samples[tail] = new_x
                    grid[new_index] = tail
                    tail += 1
    return tail

首先我們把區(qū)域的中心，即坐標為(0.5,0.5)的點選擇為初始點，讓它作為“種子”將隨機點逐漸擴散到整個區(qū)域。

接下來的while循環(huán)是算法的主體，這個循環(huán)是串行執(zhí)行的，只占用一個線程。

我們每次找到第一個需要考慮的點samples[head]，然后在以它為中心，半徑為[radius, 2*raidus]的圓環(huán)中隨機選擇100個點，逐個檢查這100個點是否不超出[0,1]×[0,1]的區(qū)域范圍，以及是否和已有點沖突。

如果都滿足的話它就是一個合格的點，我們將它的坐標和方格下標更新到samples和grid中，并將已有點的個數(shù)tail增加1。在這100個點都檢查完后，可能有多個點會被加入已有點的集合。

注意在半徑為[radius, 2*raidus]的圓環(huán)中采樣可以讓我們得到的點在滿足最小距離約束的同時距離已有點也不會太遠。

當這100個點都檢查完畢后，我們可以認為samples[head]這個點的周圍已經(jīng)沒有空白區(qū)域可以放置新的點，所以將head增加1，并重新檢查下一個samples[head] 附近的區(qū)域。

當所有的點周圍的空間都已經(jīng)被填滿，即head = tail時，或者我們已經(jīng)獲得了desired_samples個點，即tail = desired_samples時循環(huán)結(jié)束。這時samples中下標在0~tail-1范圍內(nèi)的點就是全部的已有點。

展示動畫效果

我們可以只用幾行代碼，就把整個采樣的過程用動畫的形式顯示出來：

num_samples = poisson_disk_sample(desired_samples)
gui = ti.GUI("Poisson Disk Sampling", res=800, background_color=0xFFFFFF)
count = 0
speed = 300
while gui.running:
    gui.circles(samples.to_numpy()[:min(count * speed, num_samples)],
                color=0x000000,
                radius=1.5)
    count += 1
    gui.show()

這里我們控制動畫的速度為每生成300個點就繪制一幀。

至此我們已經(jīng)介紹完了程序的所有要點，把各部分組合起來：

import taichi as ti
import taichi.math as tm
ti.init(arch=ti.cpu)

grid_n = 400
res = (grid_n, grid_n)
dx = 1 / res[0]
inv_dx = res[0]
radius = dx * ti.sqrt(2)
desired_samples = 100000
grid = ti.field(dtype=int, shape=res)
samples = ti.Vector.field(2, float, shape=desired_samples)

grid.fill(-1)

@ti.func
def check_collision(p, index):
    x, y = index
    collision = False
    for i in range(max(0, x - 2), min(grid_n, x + 3)):
        for j in range(max(0, y - 2), min(grid_n, y + 3)):
            if grid[i, j] != -1:
                q = samples[grid[i, j]]
                if (q - p).norm() < radius - 1e-6:
                    collision = True
    return collision

@ti.kernel
def poisson_disk_sample(desired_samples: int) -> int:
    samples[0] = tm.vec2(0.5)
    grid[int(grid_n * 0.5), int(grid_n * 0.5)] = 0
    head, tail = 0, 1
    while head < tail and head < desired_samples:
        source_x = samples[head]
        head += 1

        for _ in range(100):
            theta = ti.random() * 2 * tm.pi
            offset = tm.vec2(tm.cos(theta), tm.sin(theta)) * (1 + ti.random()) * radius
            new_x = source_x + offset
            new_index = int(new_x * inv_dx)

            if 0 <= new_x[0] < 1 and 0 <= new_x[1] < 1:
                collision = check_collision(new_x, new_index)
                if not collision and tail < desired_samples:
                    samples[tail] = new_x
                    grid[new_index] = tail
                    tail += 1
    return tail

num_samples = poisson_disk_sample(desired_samples)
gui = ti.GUI("Poisson Disk Sampling", res=800, background_color=0xFFFFFF)
count = 0
speed = 300
while gui.running:
    gui.circles(samples.to_numpy()[:min(count * speed, num_samples)],
                color=0x000000,
                radius=1.5)
    count += 1
    gui.show()

代碼總行數(shù)：60。

One More Thing

具體來說，這篇代碼實現(xiàn)了兩個操作：

1. 在60行代碼中實現(xiàn)了一個完整的泊松采樣動畫。
2. 在一個400×400的網(wǎng)格中采集了53k個點，但耗時不到1秒。

相關代碼可以在文末的原文鏈接中找到。

嚴格來說，本文實現(xiàn)的算法和Bridson論文里描述的算法有一點點不一樣的地方（更簡單一些），但是效果卻差不多。

你能看出是哪里不一樣嗎？（TIP：可以關注一下原論文Step 2中“active list”的處理方式）

項目地址：
https://github.com/taichi-dev/poisson-sampling-homework

參考資料：
[1]Robert Bridson的原論文見Fast Poisson Disk Sampling in Arbitrary Dimensions.
[2]Poisson采樣用Taichi, Numpy, Numba實現(xiàn)的benchmark比較見GitHub