看到fem99.py这个代码里面的自动求导部分,有点困惑,其中的中间变量如D_i, phi_i等,它们不是taichi的域,也不会被追踪导数,那么链式法则在这里是如何生效的呢?taichi的编译器会自动替换掉中间变量嘛?
pos = ti.Vector.field(2, float, NV, needs_grad=True)
vel = ti.Vector.field(2, float, NV)
f2v = ti.Vector.field(3, int, NF) # ids of three vertices of each face
B = ti.Matrix.field(2, 2, float, NF)
F = ti.Matrix.field(2, 2, float, NF, needs_grad=True)
V = ti.field(float, NF)
phi = ti.field(float, NF) # potential energy of each face (Neo-Hookean)
U = ti.field(float, (), needs_grad=True) # total potential energy
@ti.kernel
def update_U():
for i in range(NF):
ia, ib, ic = f2v[i]
a, b, c = pos[ia], pos[ib], pos[ic]
V[i] = abs((a - c).cross(b - c))
D_i = ti.Matrix.cols([a - c, b - c])
F[i] = D_i @ B[i]
for i in range(NF):
F_i = F[i]
log_J_i = ti.log(F_i.determinant())
phi_i = mu / 2 * ((F_i.transpose() @ F_i).trace() - 2)
phi_i -= mu * log_J_i
phi_i += lam / 2 * log_J_i**2
phi[i] = phi_i
U[None] += V[i] * phi_i