Wrong calculation result in 2-depth for-loop

[QiruiFU]

Describe the bug
I'm simulating the temperature changing in a 2D scene. For every direction, I use laplacian += new_lap to compute the laplacian of temperature filed. For four points who are symmetrical to each other, this value is expected to be same. But in fact they are not.

To Reproduce

# sample code here
import sys
import numpy as np

import taichi as ti
import taichi.math as tm

from matplotlib import cm

# ti.init(arch=ti.cpu, cpu_max_num_threads = 1, debug = True)
ti.init(arch=ti.cuda)

benchmark = 2

@ti.data_oriented
class lbm_solver:
    def __init__(
        self,
        name # name of the flow case
        ):
        self.name = name
        self.nx = 301
        self.ny = 301 
       
        self.old_temperature = ti.field(float, shape=(self.nx, self.ny))
        self.new_temperature = ti.field(float, shape=(self.nx, self.ny))

        self.w = ti.field(float, shape=9)
        self.e = ti.Vector.field(2, int, shape=9)

        self.image = ti.Vector.field(3, dtype=ti.f32, shape=(3 * self.nx, self.ny))  # RGB image


    @ti.kernel
    def init(self):
        self.e[0] = ti.Vector([0, 0])
        self.e[1] = ti.Vector([1, 0])
        self.e[2] = ti.Vector([0, 1])
        self.e[3] = ti.Vector([-1, 0])
        self.e[4] = ti.Vector([0, -1])
        self.e[5] = ti.Vector([1, 1])
        self.e[6] = ti.Vector([-1, 1])
        self.e[7] = ti.Vector([-1, -1])
        self.e[8] = ti.Vector([1, -1])

        self.w[0] = 4.0 / 9.0
        self.w[1] = 1.0 / 9.0
        self.w[2] = 1.0 / 9.0
        self.w[3] = 1.0 / 9.0
        self.w[4] = 1.0 / 9.0
        self.w[5] = 1.0 / 36.0
        self.w[6] = 1.0 / 36.0
        self.w[7] = 1.0 / 36.0
        self.w[8] = 1.0 / 36.0

        for i, j in self.old_temperature:
            self.old_temperature[i, j] = self.new_temperature[i, j] = 373.15



    @ti.func
    def GetTempPos(self, x, y):
        res = 0.0
        inf_positive = 873.15
        if y <= 0 or y >= self.ny - 1 or x <= 0 or x >= self.nx - 1:
            res = inf_positive
        else:
            res = self.old_temperature[x, y]
        
        return res
            
    
    @ti.func
    def GetTemp(self, x, y, k):
        res = 0.0
        nx, ny = (x + self.e[k][0]) % self.nx, (y + self.e[k][1]) % self.ny
        # if self.is_INB(x, y, k):
        #     bx, by = (x - self.e[k][0]) % self.nx, (y - self.e[k][1]) % self.ny
        #     T_i = 373.15
        #     T2 = self.GetTempPos(bx, by)

        #     u = (0.5 - self.phi[x, y]) / (self.phi[nx, ny] - self.phi[x, y])
        #     res = (2.0 * T_i + (u - 1.0) * T2) / (1 + u)
        # else:
        #     res = self.GetTempPos(nx, ny)
        res = self.GetTempPos(nx, ny)
        
        return res

    
    @ti.kernel
    def UpdateTemp(self):
        for i, j in ti.ndrange(self.nx, self.ny):
            X = 0.05

            laplacian = 0.0
            Tgrad = ti.Vector([0.0, 0.0])

            inv = ti.Vector([0, 3, 4, 1, 2, 7, 8, 5, 6])

            for k in range(1, 9):
                new_grad = 1.5 * self.w[k] * (self.GetTemp(i, j, k) - self.GetTemp(i, j, inv[k])) * self.e[k]
                new_lap = 3.0 * self.w[k] * (self.GetTemp(i, j, k) + self.GetTemp(i, j, inv[k]) - 2.0 * self.GetTempPos(i, j))
                Tgrad += new_grad
                laplacian += new_lap

                if (i == 0 and j == 8) or \
                    (i == 300 and j == 8) or \
                    (i == 0 and j == 292) or \
                    (i == 300 and j == 292):

                    print("break", i, j, k, new_lap, laplacian, new_grad, Tgrad)

            self.new_temperature[i, j] = self.old_temperature[i, j] + X * laplacian
            
        for i, j in ti.ndrange(self.nx, self.ny):
            self.old_temperature[i, j] = self.new_temperature[i, j]



    @ti.kernel
    def UpdateImage(self):
        for i, j in self.old_temperature:
            temp_color = (self.old_temperature[i, j] - 300) / 500.0
            self.image[i + 2 * self.nx, j] = ti.Vector([temp_color, 0.0, 0.0])

    
    @ti.kernel
    def CheckSym(self):
        check_temp = True

        for i0, j0 in self.old_temperature:
            i1, j1 = 300 - i0, 300 - j0

            if self.old_temperature[i0, j0]!=self.old_temperature[i0, j1] or \
                self.old_temperature[i0, j0]!=self.old_temperature[i1, j1] or \
                self.old_temperature[i0, j0]!=self.old_temperature[i1, j0]:
                check_temp = False
                print("diff temp, ", i0, j0)

        print("check sym ", check_temp)
    

    def solve(self):
        gui = ti.GUI(self.name, (3 * self.nx, self.ny))
        self.init()
        frame = 0

        while not gui.get_event(ti.GUI.ESCAPE, ti.GUI.EXIT) and frame < 100:
            sys.stdout.flush()
            frame += 1
            for substep in range(1):
                self.UpdateTemp()
                self.CheckSym()

            
            self.UpdateImage()
            gui.set_image(self.image)
            gui.show()
        


if __name__ == '__main__':
    lbm = lbm_solver(
        name = "LBM"
    )
    lbm.solve()

Log/Screenshots

[Taichi] version 1.7.0, llvm 15.0.4, commit 2fd24490, linux, python 3.10.12
[Taichi] Starting on arch=cuda
...

if you print all outputs into a file, from line 444 to line 463, you will see (the empty lines are added by myself) :

Columns directed by blue arrow: left column is new_lap and right column is laplacian.

Additional comments

[Taichi] version 1.7.0, llvm 15.0.4, commit 2fd24490, linux, python 3.10.12

*******************************************
**      Taichi Programming Language      **
*******************************************

Docs:   https://docs.taichi-lang.org/
GitHub: https://github.com/taichi-dev/taichi/
Forum:  https://forum.taichi.graphics/

Taichi system diagnose:

python: 3.10.12 (main, Feb  4 2025, 14:57:36) [GCC 11.4.0]
system: linux
executable: /usr/bin/python3
platform: Linux-6.8.0-52-generic-x86_64-with-glibc2.35
architecture: 64bit ELF
uname: uname_result(system='Linux', node='xxxx', release='6.8.0-52-generic', version='#53~22.04.1-Ubuntu SMP PREEMPT_DYNAMIC Wed Jan 15 19:18:46 UTC 2', machine='x86_64')
locale: en_US.UTF-8
PATH: /home/xxx/.cargo/bin:/usr/local/cuda-12.3/bin:/home/xxx/.local/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/games:/usr/local/games:/snap/bin:/snap/bin
PYTHONPATH: ['/home/xxx/.local/bin', '/usr/lib/python310.zip', '/usr/lib/python3.10', '/usr/lib/python3.10/lib-dynload', '/home/xxx/.local/lib/python3.10/site-packages', '/usr/local/lib/python3.10/dist-packages', '/usr/lib/python3/dist-packages']

No LSB modules are available.
Distributor ID:	Ubuntu
Description:	Ubuntu 22.04.5 LTS
Release:	22.04
Codename:	jammy



import: <module 'taichi' from '/home/xxx/.local/lib/python3.10/site-packages/taichi/__init__.py'>

cpu: True
metal: False
opengl: True
cuda: True
vulkan: True

`glewinfo` not available: [Errno 2] No such file or directory: 'glewinfo'

Thu Feb 27 22:02:59 2025       
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 550.120                Driver Version: 550.120        CUDA Version: 12.4     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA GeForce RTX 4060 ...    Off |   00000000:01:00.0 Off |                  N/A |
| N/A   37C    P8              1W /   80W |     109MiB /   8188MiB |      0%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+
                                                                                         
+-----------------------------------------------------------------------------------------+
| Processes:                                                                              |
|  GPU   GI   CI        PID   Type   Process name                              GPU Memory |
|        ID   ID                                                               Usage      |
|=========================================================================================|
|    0   N/A  N/A      2243      G   /usr/lib/xorg/Xorg                              4MiB |
|    0   N/A  N/A      2974    C+G   ...libexec/gnome-remote-desktop-daemon         94MiB |
+-----------------------------------------------------------------------------------------+

[Taichi] version 1.7.0, llvm 15.0.4, commit 2fd24490, linux, python 3.10.12

[Taichi] version 1.7.0, llvm 15.0.4, commit 2fd24490, linux, python 3.10.12
[Taichi] Starting on arch=x64

[Taichi] version 1.7.0, llvm 15.0.4, commit 2fd24490, linux, python 3.10.12
[Taichi] Starting on arch=opengl

[Taichi] version 1.7.0, llvm 15.0.4, commit 2fd24490, linux, python 3.10.12
[Taichi] Starting on arch=cuda

[Taichi] version 1.7.0, llvm 15.0.4, commit 2fd24490, linux, python 3.10.12

*******************************************
**      Taichi Programming Language      **
*******************************************

Docs:   https://docs.taichi-lang.org/
GitHub: https://github.com/taichi-dev/taichi/
Forum:  https://forum.taichi.graphics/

                                TAICHI EXAMPLES                                 
 ────────────────────────────────────────────────────────────────────────────── 
  0: ad_gravity             25:                       50: patterns              
                            karman_vortex_street                                
  1: circle_packing_image   26: keyboard              51: pbf2d                 
  2: comet                  27: laplace               52: physarum              
  3: cornell_box            28: laplace_equation      53:                       
                                                      poisson_disk_sampling     
  4: diff_sph               29: mandelbrot_zoom       54: print_offset          
  5:                        30: marching_squares      55: rasterizer            
  differential_evolution                                                        
  6: euler                  31: mass_spring_3d_ggui   56: regression            
  7: eulerfluid2d           32: mass_spring_game      57: sdf_renderer          
  8: explicit_activation    33:                       58: simple_derivative     
                            mass_spring_game_ggui                               
  9: export_mesh            34: mciso_advanced        59: simple_texture        
  10: export_ply            35: mgpcg                 60: simple_uv             
  11: export_videos         36: mgpcg_advanced        61: snow_phaseField       
  12: fem128                37: minimal               62: stable_fluid          
  13: fem128_ggui           38: minimization          63: stable_fluid_ggui     
  14: fem99                 39: mpm128                64: stable_fluid_graph    
  15: fractal               40: mpm128_ggui           65: taichi_bitmasked      
  16: fractal3d_ggui        41: mpm3d                 66: taichi_dynamic        
  17: fullscreen            42: mpm3d_ggui            67: taichi_logo           
  18: game_of_life          43: mpm88                 68: taichi_ngp            
  19: gui_image_io          44: mpm88_graph           69: taichi_sparse         
  20: gui_widgets           45: mpm99                 70: texture_graph         
  21: implicit_fem          46:                       71: tutorial              
                            mpm_lagrangian_forces                               
  22:                       47: nbody                 72:                       
  implicit_mass_spring                                two_stream_instability    
  23:                       48: odop_solar            73: vortex_rings          
  initial_value_problem                                                         
  24: jacobian              49: oit_renderer          74: waterwave             
 ────────────────────────────────────────────────────────────────────────────── 
42
Running example minimal ...
[Taichi] Starting on arch=x64
42.0
>>> Running time: 0.32s

Consider attaching this log when maintainers ask about system information.
>>> Running time: 6.38s