How to compute gradient of function that mixes loops

[cdelv]

Hi,

I'm currently triying to compute gradient of a potential energy definition for a mesh. The idea is to use this gradient to calculate forces in the vertices. The problem I have is that the potential energy is given by:

U = ev*np.power(mesh.volume/V0 - 1.0, 2) + ea*np.sum(np.power(mesh.area_faces/A0 - 1.0, 2))

Therefore, the gradient requires adding a value to the result of the sumation. However, the volume also depends on the vertices as it's calculated as a triple scalar product of the vertex's positions per face. So my code looks something like

vol = 0
U = 0
for i in faces:
    # compute area
    # compute partial volume
    vol += partial_volume
    U += area_potential_energy

U += volume potential energy.

However, this violates the 3rd rule for automatic gradients in Taichi. I have tried doing lots of things to separate the gradient calculations but I'm unable to express it only with loops.

I thought of computing only the gradient of the volume and then multiplying by the corresponding factor and then add it to the other gradient but this also does not work as vertices.grad gets overriten by the U computation. A solution would be to copy the result first and then compute the other one, but this is very inefficient.

I appreciate the help.

PSD: Is it possible to call a gradient inside a taichi function?

Best regards.

[hughperkins]

Hi, I'm just answering as a learning exercise. I literally wrote my first taichi program ~4 hours ago. So apologies if this answer is both 100% naive and 100% wrong.

I started by AI'ing a program that attempts to reproduce your problem, and ended up with:

import taichi as ti
import numpy as np

ti.init(arch=ti.gpu)

V0 = 1.0
A0 = 1.0
ev = 1.0
ea = 1.0

n_vertices = 4
n_faces    = 4
vertices = ti.Vector.field(3, float, n_vertices, needs_grad=True)
faces    = ti.Vector.field(3, int,   n_faces)
forces   = ti.Vector.field(3, float, n_vertices)
U        = ti.field(float, shape=(), needs_grad=True)

vertices.from_numpy(np.array([
    [ 1,  1,  1],
    [-1, -1,  1],
    [-1,  1, -1],
    [ 1, -1, -1]], dtype=np.float32))
faces.from_numpy(np.array([
    [0, 1, 2],
    [0, 1, 3],
    [0, 2, 3],
    [1, 2, 3]], dtype=np.int32))

@ti.kernel
def compute_energy():
    vol = 0.0
    U[None] = 0.0
    for fi in range(n_faces):
        idxs = faces[fi]
        p0 = vertices[idxs[0]]
        p1 = vertices[idxs[1]]
        p2 = vertices[idxs[2]]

        area = 0.5 * (p1 - p0).cross(p2 - p0).norm()
        U[None] += ea * (area / A0 - 1.0) ** 2

        vol += p0.dot((p1 - p0).cross(p2 - p0)) / 6.0

    U[None] += ev * (vol / V0 - 1.0) ** 2


with ti.ad.Tape(loss=U):
    compute_energy()
for i in range(n_vertices):
    forces[i] = -vertices.grad[i]

print("Forces:", forces.to_numpy())

This gives the error:

[E 04/25/25 22:39:08.162 348700] [reverse_segments.cpp:reverse_segments@74] Invalid program input for autodiff: Mixed usage of for-loops and statements without looping.
Please split them into two kernels and check the documentation for more details:
https://docs.taichi-lang.org/docs/differentiable_programming

This sounds like the problem you are describing?

So, then I attempted - likely very naively - to split into multiple kernels, as follows:

import taichi as ti
import numpy as np

ti.init(arch=ti.gpu)

V0 = 1.0
A0 = 1.0
ev = 1.0
ea = 1.0

n_vertices = 4
n_faces    = 4
vertices = ti.Vector.field(3, float, n_vertices, needs_grad=True)
faces    = ti.Vector.field(3, int,   n_faces)
forces   = ti.Vector.field(3, float, n_vertices)
U        = ti.field(float, shape=(), needs_grad=True)
vol = ti.field(float, shape=(), needs_grad=True)

vertices.from_numpy(np.array([
    [ 1,  1,  1],
    [-1, -1,  1],
    [-1,  1, -1],
    [ 1, -1, -1]], dtype=np.float32))
faces.from_numpy(np.array([
    [0, 1, 2],
    [0, 1, 3],
    [0, 2, 3],
    [1, 2, 3]], dtype=np.int32))

@ti.kernel
def initialize():
    vol[None] = 0.0
    U[None] = 0.0

@ti.kernel
def compute_energy_1():
    for fi in range(n_faces):
        idxs = faces[fi]
        p0 = vertices[idxs[0]]
        p1 = vertices[idxs[1]]
        p2 = vertices[idxs[2]]

        area = 0.5 * (p1 - p0).cross(p2 - p0).norm()
        U[None] += ea * (area / A0 - 1.0) ** 2

        vol[None] += p0.dot((p1 - p0).cross(p2 - p0)) / 6.0

@ti.kernel
def compute_energy_2():
    U[None] += ev * (vol[None] / V0 - 1.0) ** 2

with ti.ad.Tape(loss=U):
    initialize()
    compute_energy_1()
    compute_energy_2()
for i in range(n_vertices):
    forces[i] = -vertices.grad[i]

print("Forces:", forces.to_numpy())

This gives the following output:

(taichi) ~/git/taichi-play $ python 8634_sol1.py
[Taichi] version 1.7.3, llvm 15.0.7, commit 5ec301be, osx, python 3.10.16
[Taichi] Starting on arch=metal
Forces: [[-11.381197 -12.71453  -11.381197]
 [ 11.381197  10.047864 -11.381197]
 [ 11.381197 -12.714531  11.381197]
 [-11.381197  10.047864  11.381197]]

Thoughts? To what extent does this address the issue you raise?