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.