Ising-ish model with central gravitational force

central.py

import taichi as ti
import taichi.math as tm
import time
import random

ti.init(arch=ti.gpu)

w = 1440//2
h = 1440//2
# pixels = ti.field(dtype=float, shape=(h, w))
pixels = color_field = ti.Vector.field(3, dtype=ti.f32, shape=(w, h))
state = ti.field(dtype=ti.i32, shape=(w, h))
newstate = ti.field(dtype=ti.i32, shape=(w, h))

m = ti.field(dtype=ti.f32, shape=(4, 4))
density = ti.field(dtype=ti.f32, shape=(4,))
colors = ti.Vector.field(3, dtype=ti.f32, shape=(4,))
probs = ti.field(dtype=ti.f32, shape=(4,))

@ti.func
def gravity(i: ti.i32, j: ti.i32):
  dx = i - 0.5 * w
  dy = j - 0.5 * h
  n = ti.sqrt(dx * dx + dy * dy)
  return ti.Vector([dx / n, dy / n])

@ti.kernel
def init_m():
    m[0, 0] = -1
    m[0, 1] = 1
    m[0, 2] = 1
    m[0, 3] = 0

    m[1, 1] = -1
    m[1, 2] = -1
    m[1, 3] = -1

    m[2, 2] = 0.5
    m[2, 3] = -1

    m[3, 3] = 0

    m[1, 0] = m[0, 1]
    m[2, 0] = m[0, 2]
    m[2, 1] = m[1, 2]
    m[3, 0] = m[0, 3]
    m[3, 1] = m[1, 3]
    m[3, 2] = m[2, 3]

    density[0] = 1.5
    density[1] = -1.5
    density[2] = 3.0
    density[3] = 0.0

    colors[0] = ti.Vector([0, 1.0, 0.3])
    colors[1] = ti.Vector([0., 0., 0])
    colors[2] = ti.Vector([0.9, 0.1, 0])
    colors[3] = ti.Vector([1, 1, 1])

    probs[0] = 0.5
    probs[1] = 0.2
    probs[2] = 0.3
    probs[3] = 0.0

@ti.kernel
def init():
    for i, j in state:
        r = ti.random(dtype=ti.f32)
        c = 0
        for k in range(4):
            r -= probs[k]
            if r <= 0 or k == 3:
              c = k
              break
        state[i, j] = c

@ti.kernel
def sim_v(beta:ti.f32, k:ti.int32, l:ti.int32):
           for i0 in range(0, w // 4):
               for j0 in range(0, h // 4):
                 if ti.random(ti.i32) % 2 == 0:
                    i = 4 * i0 + k
                    j = 4 * j0 + l

                    energy = 0.

                    if i < w and j + 1 < h:
                        # e1:    e2:
                        #
                        #  d      d
                        # bvf    buf
                        # aue    ave
                        #  c      c

                        u = state[i, j]
                        v = state[i, j + 1]

                        if i - 1 >= 0:
                            a = state[i - 1, j]
                            energy += m[u, a] - m[v, a]
                        if j - 1 >= 0:
                            c = state[i, j - 1]
                            energy += m[u, c] - m[v, c]
                        if i + 1 < w:
                            e = state[i + 1, j]
                            energy += m[u, e] - m[v, e]
                        if i - 1 >= 0 and j + 1 < h:
                            b = state[i - 1, j + 1]
                            energy += m[v, b] - m[u, b]
                        if j + 2 < h:
                            d = state[i, j + 2]
                            energy += m[v, d] - m[u, d]
                        if i + 1 < w and j + 1 < h:
                            f = state[i + 1, j + 1]
                            energy += m[v, f] - m[u, f]

                        g = gravity(i, j)
                        energy += (density[u] - density[v]) * g[1];

                        if ti.log(ti.random(dtype=ti.f32)) < (-beta * energy):
                            state[i, j] = v
                            state[i, j + 1] = u

@ti.kernel
def sim_h(beta:ti.f32, k:ti.int32, l:ti.int32):
           for i0 in range(0, w // 4):
               for j0 in range(0, h // 4):
                 if ti.random(ti.i32) % 2 == 0:
                    i = 4 * i0 + k
                    j = 4 * j0 + l

                    energy = 0.

                    if i + 1 < w and j < h:
                        # e1:    e2:
                        #
                        #  ab    ab
                        # cuvd  cvud
                        #  ef    ef

                        u = state[i, j]
                        v = state[i + 1, j]

                        if j - 1 >= 0:
                            a = state[i, j - 1]
                            energy += m[u, a] - m[v, a]
                        if i - 1 >= 0:
                            c = state[i - 1, j]
                            energy += m[u, c] - m[v, c]
                        if j + 1 < h:
                            e = state[i, j + 1]
                            energy += m[u, e] - m[v, e]
                        if j - 1 >= 0 and i + 1 < w:
                            b = state[i + 1, j - 1]
                            energy += m[v, b] - m[u, b]
                        if i + 2 < w:
                            d = state[i + 2, j]
                            energy += m[v, d] - m[u, d]
                        if j + 1 < h and i + 1 < w:
                            f = state[i + 1, j + 1]
                            energy += m[v, f] - m[u, f]

                        g = gravity(i, j)
                        energy += (density[u] - density[v]) * g[0];

                        if ti.log(ti.random(dtype=ti.f32)) < (-beta * energy):
                            state[i, j] = v
                            state[i + 1, j] = u

@ti.kernel
def transfer_state_to_pixels():
    for i, j in pixels:
        pixels[i, j] = colors[state[i, j]]

window = ti.ui.Window("Stoch", fps_limit=200, res=(1440,1440), vsync=True)
canvas = window.get_canvas()

init_m()
init()
i = 0
beta = 1.25
while window.running:
    if window.is_pressed(ti.ui.LMB):
      x, y = window.get_cursor_pos()
      if y > 0 and y < 1:
        beta = 1. / (4. * y)
        print(beta)

    for subframe in range(1):

      for k in range(4):
          for l in range(4):
              sim_v(beta, k, l)

      for k in range(4):
          for l in range(4):
              sim_h(beta, k, l)

    transfer_state_to_pixels()
    canvas.set_image(pixels)
    window.show()