☯

https://xpil.eu/gkt 📋

Yin oraz Yang to koncept dualizmu świata pochodzący z chińskiej filozofii. Jest to również znak rozpoznawczy T'ai Chi, a na polski tłumaczy się go bodajże jako "schemat najwyższej ostateczności".

Ale ja dziś nie o tym. Zamiast gadać o taoizmie i starodawnych Chinach, chciałbym się dziś skupić na nowoodkrytej przeze mnie niedawno bibliotece Pythona o wdzięcznej i dźwięcznej nazwie jak w tytule, czyli taichi.

Cóż to za stwór?

A no jest to całkiem pokaźny kawałek kodu umożliwiający pisanie aplikacji na platformy przetwarzania równoległego (procesory wielowątkowe Intela i AMD, CUDA, Vulkan, OpenGL, Metal i co tam, panie, jeszcze) za pomocą składni Python. Biorąc pod uwagę nieszczęsne, od lat denerwujące programistów blokowanie wielowątkowości w Pythonie, taichi daje całkiem solidnego kopa. W szczególności pozwala na pisanie procedur w Pythonie, które będą wykonywane bezpośrednio na chipie karty graficznej; jeżeli procedura taka zawiera pętlę for, będzie ona automatycznie "zrównoleglona". Sama biblioteka jest napisana w C, ale dzięki dekoratorom my jako jej użytkownicy wcale nie musimy o tym wiedzieć.

Ponieważ taichi zainteresowałem się względnie niedawno i wiem o tej bibliotece naprawdę niewiele, pokażę dziś na skromnym (mniej niż 100 linii kodu) przykładzie jak wykorzystać taichi do napisania prostego symulatora planetarnego. A więc: centralna gwiazda (lub gwiazdy), do tego turlające się dookoła planety o rozmaitych masach. Dla uproszczenia - wszystko w dwóch wymiarach.

Uwaga: kod jest raczej prymitywny i nie nadaje się do "porządnych" naukowych symulacji zachowania ciał niebieskich w przestrzeni. W szczególności nie uwzględnia zderzeń (jeżeli dwa ciała niebieskie znajdą się w tym samym miejscu, "przenikną" przez siebie). Nie uwzględnia też ewentualnych nieregularności w rozkładach mas wewnątrz planet / gwiazd (każde ciało niebieskie traktowane jest jako matematyczny punkt o zadanej masie) ani efektów relatywistycznych. Aha, ponieważ nie umiem jeszcze sterować zoomem ani perspektywą w GUI taichi, a planety prędzej czy później zaczynają uciekać poza krawędzie ekranu, dodałem dodatkowo gumowe brzegi, więc jeżeli cokolwiek próbuje uciec za ekran, odbije się od krawędzi 🙂 Ponadto zamiast użyć wbudowanych w taichi gradientów, liczę te gradienty na piechotę, bo póki co nijak nie potrafię załapać zero-wymiarowych obiektów wymaganych przy liczeniu gradientów metodami natywnymi.

Sam kod wygląda jak widać poniżej. Proszę się nie śmiać, to moje pierwsze wypociny z taichi i pewnie dałoby się to napisać pierdylion razy lepiej:

import taichi as ti

ti.init()

ile_planet = 5
kwant_czasu = 3e-5
rozdzielczość_x = 1200
rozdzielczość_y = 800
# planety: [x, y]
planety = ti.Vector.field(2, dtype=ti.f32, shape=ile_planet)
# prędkości: [vx, vy]
prędkości = ti.Vector.field(2, dtype=ti.f32, shape=ile_planet)
# masy: [m]
masy = ti.field(ti.f32, shape=ile_planet)
# gradient siły [fx, fy]
wektory_sił = ti.Vector.field(2, dtype=ti.f32, shape=ile_planet)

@ti.kernel
def policz_wektory_sił():
    wektory_sił.fill([0, 0])
    for p1, p2 in ti.ndrange(ile_planet, ile_planet):
        if(p1 != p2):
            odległość = (planety[p1]-planety[p2]).norm(1e-3)
            zmiana = masy[p1]*masy[p2]*(planety[p2]-planety[p1])/(odległość**3)
            wektory_sił[p1] += zmiana


@ti.kernel
def przesuń():
    for i in planety:
        prędkości[i] += kwant_czasu / masy[i] * wektory_sił[i]

        planety[i] += kwant_czasu*prędkości[i]

        # bounce, bounce
        if(planety[i][0] < 0 or planety[i][0] > 1):
            prędkości[i][0] *= -1
        if(planety[i][1] < 0 or planety[i][1] > 1):
            prędkości[i][1] *= -1


def krok():
    policz_wektory_sił()
    przesuń()


@ti.kernel
def init():
    for p in planety:
        planety[p] = [ti.random()/3+0.35, ti.random()/3+0.35]
        prędkości[p] = [105-ti.random()*20, 10-ti.random()*20]
        masy[p] = ti.random()*50
    planety[0] = [0.4, 0.4]
    prędkości[0] = [0, 0]
    masy[0] = 5000
    # planety[1] = [0.6, 0.4]
    # prędkości[1] = [0, 0]
    # masy[1] = 5000
    # planety[2] = [0.4, 0.6]
    # prędkości[2] = [0, 0]
    # masy[2] = 5000


def main():
    init()
    gui = ti.GUI('Planety', res=(rozdzielczość_x, rozdzielczość_y))
    while gui.running:
        krok()
        pn = planety.to_numpy()
        mn = masy.to_numpy()
        for i, p in enumerate(pn):
            rad = max(1, min(16, int(mn[i]/5)))
            gui.circle(p, radius=rad)  # , color=kolor)
        gui.show()


if __name__ == "__main__":
    main()

Efekt działania jest taki, że na ekranie pokazuje się kilka/-naście/-dziesiąt/-set(?) kółek o rozmaitych wielkościach (=masach), które następnie zaczynają się poruszać pod wpływem grawitacji.

Na moim starutkiej, szrotowatej karcie grafiki popróbowałem różnych liczb planet (linia #5 w powyższym kodzie) i do mniej więcej 80 planet miałem 60fps, a potem zaczyna zwalniać (aż do marnych 3fps dla 2000 planet), przy czym zmniejszanie okienka pozwala oczywiście na poprawę płynności przy utracie rozdzielczości. Coś za coś.

Omówmy sobie teraz poszczególne kawałki kodu:

import taichi as ti

ti.init()

Na początek importujemy bibliotekę i inicjalizujemy ją.

ile_planet = 5
kwant_czasu = 3e-5
rozdzielczość_x = 1200
rozdzielczość_y = 800
# planety: [x, y]
planety = ti.Vector.field(2, dtype=ti.f32, shape=ile_planet)
# prędkości: [vx, vy]
prędkości = ti.Vector.field(2, dtype=ti.f32, shape=ile_planet)
# masy: [m]
masy = ti.field(ti.f32, shape=ile_planet)
# gradient siły [fx, fy]
wektory_sił = ti.Vector.field(2, dtype=ti.f32, shape=ile_planet)

W liniach 5-16 inicjalizujemy kilka zmiennych globalnych, przy czym o ile linie 5-8 są nudne, o tyle dalej robi się interesująco: w linii 10 mamy zmienną planety typu ti.Vector.field - jest to ni mniej ni więcej jak tablica punktów (x, y) reprezentujących poszczególne planety, poindeksowana liczbami naturalnymi od zera. Jest to wewnętrzny typ biblioteki taichi, który jednak daje się łatwo przekształcić na "zwykłą" Pythonową listę jak też listę numpy. W linii 12 mamy prędkości poszczególnych planet (również wektory (x,y) gdzie x to składowa pozioma a y - pionowa prędkości), potem w linii 14 masy typu ti.field czyli lista liczb (również indeksowana od zera) wreszcie w linii 16 lista wektorów sił oddziałujących na każdą planetę (też dwie składowe x i y), którą będziemy w każdym kwancie czasu przeliczać od nowa, a która umożliwi ustalenie w którą stronę dana planeta jest w danym momencie ciągnięta i jak mocno.

Zauważamy mimochodem, że wymiarowość (czyli liczba składowych każdego elementu) typu ti.Vector.field podajemy jako pierwszy parametr konstruktora (tutaj: 2), natomiast parametry dtype oraz shape definiują odpowiednio typ pojedynczej składowej (tutaj: f32 czyli 32-bitowa liczba zmiennoprzecinkowa) oraz długość - rozmiar tablicy (czyli liczba par).

@ti.kernel
def policz_wektory_sił():
    wektory_sił.fill([0, 0])
    for p1, p2 in ti.ndrange(ile_planet, ile_planet):
        if(p1 != p2):
            odległość = (planety[p1]-planety[p2]).norm(1e-3)
            zmiana = masy[p1]*masy[p2]*(planety[p2]-planety[p1])/(odległość**2)
            wektory_sił[p1] += zmiana

W liniach 18-25 definiujemy sobie funkcję, która przeliczy nam dla każdej planety jaka siła na nią działa (suma sił przyciągania wszystkich pozostałych planet / gwiazd) oraz w którą stronę. Zaczynamy od użycia dekoratora @ti.kernel, który mówi Pythonowi, żeby teraz przekazać "władzę" bibliotece taichi - kod tej funkcji / procedury zostanie wewnętrznie przełożony na "język" karty graficznej (np. CUDA) a także, o ile funkcja posiada pętlę for, odpowiednio "zrównoleglony". To również narzuca nam pewne ograniczenia - wewnątrz takiej funkcji powinniśmy używać wyłącznie odwołań do elementów zadeklarowanych wcześniej z przedrostkiem ti., przy czym nie jest to stuprocentowa prawda (da się używać zmiennych lokalnych pythonowych) tylko takie ogólne zalecenie. Jeżeli użyjemy tu czegoś "niedozwolonego", zostaniemy o tym powiadomieni już na etapie kompilacji.

W linii 20 zerujemy wektory sił dla wszystkich planet (metoda fill wypełnia listę zadaną wartością, tutaj [0,0]), następnie uruchamiamy pętlę for, która leci po wszystkich planetach i to aż dwukrotnie. Metoda ndrange zwraca n-wymiarową kostkę liczb, tutaj n=2 czyli kwadrat o boku długości liczby planet. W zasadzie można by wykonać to w postaci dwóch pętli, jednej zewnętrznej, drugiej wewnętrznej, ale wtedy biblioteka taichi "zrównolegliłaby" tylko tę zewnętrzną.

Prościej mówiąc, powyższa pętla przelatuje przez wszystkie kombinacje dwuelementowe planet, przy czym p1 i p2 to są naturalne indeksy.

I teraz tak: jeżeli indeksy te są różne (czyli mamy dwie różne planety), liczymy odległość między nimi - wyrażenie w nawiasie (planety[p1]-planety[p2]) daje nam wektor, którego długość liczymy wbudowaną funkcją norm, która dodatkowo pozwala na wymuszenie minimalnej dozwolonej długości (będziemy przez nią później dzielić, głupio by było, gdyby wyszło zero). Tutaj - 1e-3 czyli jedna tysięczna.

W linii #24 wyliczamy sobie siłę działającą między planetami p1 i p2, ze wzoru newtonowskiego:

$$F=\frac{G m_1 m_2}{r^2}$$

, przy czym G dla uproszczenia pomijamy, \(m_1\) oraz \(m_2\) to masy obydwu planet, a r - odległość między nimi środkami ich mas. Tym samym zmienna zmiana jest wektorem [dfx, dfy] określającym składowe poziomą i pionową oddziaływania między tymi dwiema konkretnymi planetami.

W kolejnej linii dodajemy ten wektor do wszystkich już policzonych sił dla tej planety - w efekcie po wykonaniu pętli zmienna wektory_sił zawiera wszystkie aktualne wypadkowe wektory sił oddziałujących na każdą planetę / gwiazdę.

@ti.kernel
def przesuń():
    for i in planety:
        prędkości[i] += kwant_czasu / masy[i] * wektory_sił[i]

        planety[i] += kwant_czasu*prędkości[i]
        # bounce, bounce
        if(planety[i][0] < 0 or planety[i][0] > 1):
            prędkości[i][0] *= -1
        if(planety[i][1] < 0 or planety[i][1] > 1):
            prędkości[i][1] *= -1

Procedura przesuń robi to, co napisano na opakowaniu. Najpierw (linie 30-31) dla każdej planety zmienia jej prędkość (która jest wektorem [vx, vy]) proporcjonalnie do długości kwantu czasu oraz wektora siły dla tej planety, a odwrotnie proporcjonalnie do jej masy. Dzięki temu ostatniemu planety większe (tudzież gwiazdy) poruszają się "niemrawo" podczas gdy mniejsze kamyki wariują jak ratlerek na dźwięk dzwonka. A potem (linie 32-33) modyfikujemy lokalizację każdej planety zgodnie z jej wektorem prędkości. Linie 35-38 to prymitywna "proteza" czyli gumowy płotek sprawiający, że ciała niebieskie nie uciekają za ekran tylko odbijają się od jego krawędzi.

I to w zasadzie wszystko jeśli chodzi o samo gęste, dalej już tylko budujemy z powyższych klocków. Czyli najpierw:

def krok():
    policz_wektory_sił()
    przesuń()

procedura krok najpierw przelicza wektory sił, a potem przesuwa planety.

@ti.kernel
def init():
    for p in planety:
        planety[p] = [ti.random()/3+0.35, ti.random()/3+0.35]
        prędkości[p] = [105-ti.random()*20, 10-ti.random()*20]
        masy[p] = ti.random()*50
    planety[0] = [0.4, 0.4]
    prędkości[0] = [0, 0]
    masy[0] = 5000
    # planety[1] = [0.6, 0.4]
    # prędkości[1] = [0, 0]
    # masy[1] = 5000
    # planety[2] = [0.4, 0.6]
    # prędkości[2] = [0, 0]
    # masy[2] = 5000

init ustala warunki początkowe, czyli najpierw losuje położenia, prędkości i masy wszystkich planet (linie 49-52), a potem ustawia planetę numer 0 w miarę na środku, stopuje ją i nadaje jej wielką masę, czyli promuje ją do roli gwiazdy. Zakomentowane linie 56-61 to na wypadek gdybyśmy chcieli więcej gwiazd. Aha, pole widzenia rozciąga się od (0,0) w dolnym lewym rogu do (1,1) w górnym prawym. Pewnie da się to jakoś przeskalować, ale nie chciało mi się już szukać jak konkretnie.

def main():
    init()
    gui = ti.GUI('Planety', res=(rozdzielczość_x, rozdzielczość_y))
    while gui.running:
        krok()
        pn = planety.to_numpy()
        mn = masy.to_numpy()
        for i, p in enumerate(pn):
            rad = max(1, min(16, int(mn[i]/5)))
            gui.circle(p, radius=rad)  # , color=kolor)
        gui.show()
if __name__ == "__main__":
    main()

Końcówka, czyli linie 64-76 to procedura główna - najpierw wołamy init, potem inicjalizacja GUI (okienko graficzne), wreszcie w nieskończonej pętli (a tak naprawdę dopóki ktoś nie zamknie okienka otwartego w poprzedniej linii) wykonujemy krok, a potem rysujemy planety za pomocą metody circles, przy czym wielkości kółek ustawiamy pi x drzwi proporcjonalnie do mas planet, ale z ograniczeniem z obu stron żeby nie wyszły za małe albo za wielkie.

Efekt jest może nie jakoś kosmicznie (nomen-omen) powalający, ale jednak całkiem sympatyczny. Niestety zabrakło mi już rozpędu na wyeksportowanie animacji do jakiegoś pliku video, więc każdy będzie sobie musiał sam programik uruchomić i obejrzeć.

Dokumentacja taichi dostępna jest tutaj: https://docs.taichi-lang.org/docs.

Polecam pobawić się, zwłaszcza jeżeli interesujesz się przetwarzaniem równoległym.

https://xpil.eu/gkt 📋