语法糖
别名
有时为命名繁琐的全局变量和函数创建别名可以提高可读性。 在 Taichi 中,这可以通过将 ti.static()
分配给 kernel 和函数本地变量实现,这会迫使 Taichi 使用标准的 python 指针分配。
例如,观察下面这个简单的内核:
@ti.kernel
def my_kernel():
for i, j in field_a:
field_b[i, j] = some_function(field_a[i, j])
fields 和函数可以通过 ti.static
重命名新的别名:
@ti.kernel
def my_kernel():
a, b, fun = ti.static(field_a, field_b, some_function)
for i, j in a:
b[i, j] = fun(a[i, j])
还可以为类成员和方法创建别名,这有助于避免出现面向对象编程代码塞满 self
的情况。
例如,考虑使用类内核来计算某个 field 的二维拉普拉斯算子:
@ti.kernel
def compute_laplacian(self):
for i, j in a:
self.b[i, j] = (self.a[i + 1, j] - 2.0*self.a[i, j] + self.a[i-1, j])/(self.dx**2) \
+ (self.a[i, j + 1] - 2.0*self.a[i, j] + self.a[i, j-1])/(self.dy**2)
使用 ti.static()
,这可以简化为:
@ti.kernel
def compute_laplacian(self):
a, b, dx, dy = ti.static(self.a, self.b, self.dx, self.dy)
for i, j in a:
b[i, j] = (a[i+1, j] - 2.0*a[i, j] + a[i-1, j])/(dx**2) \
+ (a[i, j+1] - 2.0*a[i, j] + a[i, j-1])/(dy**2)