for循环的语法对ndarray不起作用吗?

我想用for i, j in xxx的语法遍历一个ndarray,但是显示错误。
我的代码(Julia分形示例改版,换用了GGUI和ndarray):

import taichi as ti
import taichi.math as tm
import taichi.types as tt
import pygame

ti.init(arch=ti.cpu)

n = 320
pixels = tt.ndarray(dtype=float, ndim=2)

@ti.func
def complex_sqr(z):  # complex square of a 2D vector
    return tm.vec2(z[0] * z[0] - z[1] * z[1], 2 * z[0] * z[1])

@ti.kernel
def paint(t: float, pixels: ti.template()): # type: ignore
    for i, j in pixels:  # Parallelized over all pixels
        c = tm.vec2(-0.8, tm.cos(t) * 0.2)
        z = tm.vec2(i / n - 1, j / n - 0.5) * 2
        iterations = 0
        while z.norm() < 20 and iterations < 50:
            z = complex_sqr(z) + c
            iterations += 1
        pixels[i, j] = 1 - iterations * 0.02

gui = ti.ui.Window("Julia Set", res=(int(n * 2), n)) # type: ignore
canvas = gui.get_canvas()

i = 1
clock = pygame.time.Clock()
while gui.running:
    paint(i * 0.03, pixels)
    canvas.set_image(pixels)
    gui.show()
    i += 1
    clock.tick(60)

报错:

Traceback (most recent call last):
  File "d:\vscode_projects\taichi_studies\taichi_demo.py", line 32, in <module>
    paint(i * 0.03, pixels)
  File "D:\Program Files\Python39\lib\site-packages\taichi\lang\kernel_impl.py", line 976, in wrapped
    raise type(e)('\n' + str(e)) from None
taichi.lang.exception.TaichiCompilationError:
File "d:\vscode_projects\taichi_studies\taichi_demo.py", line 17, in paint:
    for i, j in pixels:  # Parallelized over all pixels
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
Traceback (most recent call last):
  File "D:\Program Files\Python39\lib\site-packages\taichi\lang\ast\ast_transformer_utils.py", line 25, in __call__ 
    return method(ctx, node)
  File "D:\Program Files\Python39\lib\site-packages\taichi\lang\ast\ast_transformer.py", line 1347, in build_For    
    return ASTTransformer.build_struct_for(ctx,
  File "D:\Program Files\Python39\lib\site-packages\taichi\lang\ast\ast_transformer.py", line 1237, in build_struct_for
    impl.begin_frontend_struct_for(ctx.ast_builder, expr_group,
  File "D:\Program Files\Python39\lib\site-packages\taichi\lang\impl.py", line 88, in begin_frontend_struct_for     
    raise TypeError(
TypeError: Cannot loop over the object <class 'taichi.types.ndarray_type.NdarrayType'> in Taichi scope. Only Taichi fields (via template) or dense arrays (via types.ndarray) are supported.

因为你的 ndarray 是一个 type,并不是一个实际的 array。何况你的 pixels 里面也没有 ndarray 的大小信息?

下面这段代码是可以的:

import taichi as ti
import taichi.math as tm
import taichi.types as tt

ti.init(arch=ti.cpu)

n = 320
pixels = ti.ndarray(dtype=float, shape=(n * 2, n))

@ti.func
def complex_sqr(z):  # complex square of a 2D vector
    return tm.vec2(z[0] * z[0] - z[1] * z[1], 2 * z[0] * z[1])

@ti.kernel
def paint(t: float, pixels: ti.types.ndarray(field_dim=2)): # type: ignore
    for i, j in pixels:  # Parallelized over all pixels
        c = tm.vec2(-0.8, tm.cos(t) * 0.2)
        z = tm.vec2(i / n - 1, j / n - 0.5) * 2
        iterations = 0
        while z.norm() < 20 and iterations < 50:
            z = complex_sqr(z) + c
            iterations += 1
        pixels[i, j] = 1 - iterations * 0.02

gui = ti.GUI("Julia Set", res=(n * 2, n))

for i in range(1000000):
    paint(i * 0.03, pixels)
    gui.set_image(pixels.to_numpy())
    gui.show()

就是初始化的时候是ti.ndarray而不是ti.types.ndarray,同时指定shape能将各个元素初始化为0。
传入kernel的时候要指定类型为ti.types.ndarray()