Ising type model using Kawasaki type dynamics

defect.py

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

ti.init(arch=ti.gpu)

w = 256
h = 256
# 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=(h, w))
newstate = ti.field(dtype=ti.i32, shape=(h, w))

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.kernel
def init_m():
    m[0, 0] = -1
    m[0, 1] = 1
    m[0, 2] = 2
    m[0, 3] = 0

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

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

    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] = 0.0
    density[3] = 0.0

    colors[0] = ti.Vector([1, 0.25, 0])
    colors[1] = ti.Vector([0, 0.25, 1])
    colors[2] = ti.Vector([0, 0, 0])
    colors[3] = ti.Vector([1, 1, 1])

    probs[0] = 0.5
    probs[1] = 0.5
    probs[2] = 0.0
    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]

                        energy += density[u] - density[v];

                        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]

                        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=(4 * h, 4 * w), vsync=True)
canvas = window.get_canvas()

init_m()
init()
i = 0
beta = 1.0
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()