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

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

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

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

同樣地，在計(jì)算機(jī)模擬下，也有不少場景需要在空間中隨機(jī)、均勻地生成點(diǎn)。

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

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

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

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

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

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

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

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

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

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

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

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

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

但這種算法的效率很低。

因?yàn)殡S著得到的點(diǎn)的個(gè)數(shù)增加，沖突的概率越來越大，獲得新的點(diǎn)所需的時(shí)間也越來越長，每次比較當(dāng)前點(diǎn)和所有已有點(diǎn)之間的距離也會(huì)降低效率。

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

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

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

這個(gè)動(dòng)畫是使用Taichi生成的，運(yùn)行在單個(gè)CPU線程上，除去編譯的時(shí)間計(jì)算，耗時(shí)僅在0.7s多一點(diǎn)，而同樣的代碼翻譯成NumPy要耗時(shí)70s左右。[2]

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

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

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

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

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

這個(gè)采樣過程很難并行化，因?yàn)楫?dāng)一個(gè)線程“偷偷”加入一個(gè)新的點(diǎn)的時(shí)候，會(huì)改變其它所有線程對距離的判斷。所以Taichi僅使用單個(gè)CPU線程來執(zhí)行這個(gè)算法：

ti.init(arch=ti.cpu)

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

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

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

安裝和導(dǎo)入Taichi

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

pip install taichi==1.0.3

然后，在代碼開頭寫上：

import taichi as ti
import taichi.math as tm

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

準(zhǔn)備工作

在泊松采樣算法中，采樣點(diǎn)之間的距離有一個(gè)下界r。

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

于是任何小方格中至多包含一個(gè)點(diǎn)。如下圖所示：

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

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

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

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

于是我們的初始設(shè)置可以這樣寫：

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

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

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

grid.fill(-1)

如何生成新的點(diǎn)

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

這里需要注意的是：

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

檢查一個(gè)點(diǎn)是否和已有點(diǎn)沖突的邏輯我們單獨(dú)寫成一個(gè)函數(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是需要檢查點(diǎn)的坐標(biāo)，index=(x, y)是p所在的方格的下標(biāo)。

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

完成了準(zhǔn)備工作，我們可以開始正式的循環(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ū)域的中心，即坐標(biāo)為(0.5,0.5)的點(diǎn)選擇為初始點(diǎn)，讓它作為“種子”將隨機(jī)點(diǎn)逐漸擴(kuò)散到整個(gè)區(qū)域。

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

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

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

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

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

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

展示動(dòng)畫效果

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

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()

這里我們控制動(dòng)畫的速度為每生成300個(gè)點(diǎn)就繪制一幀。

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

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

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

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

相關(guān)代碼可以在文末的原文鏈接中找到。

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

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

項(xiàng)目地址：
https://github.com/taichi-dev/poisson-sampling-homework

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

— 完 —

量子位 QbitAI · 頭條號簽約

關(guān)注我們，第一時(shí)間獲知前沿科技動(dòng)態(tài)