Differentiability wrt. to mesh vertex positions

[andrinr]

I'm currently trying to figure out whether phiflow supports differentiability wrt to mesh vertex positions. My initial assessment of the way the mesh is setup and initial failed experiments make me believe this is not supported.

[holl-]

When using the Mesh class, the vertex positions are variable by default. So a function like loss_function(mesh: Mesh) can be differentiated w.r.t. the vertex positions. When you pass it as part of a Field, make sure to set variable_attrs=('values', 'geometry'), else the vertex positions are taken as constant.

[dionhaefner]

I'm trying to adapt the cylinder FVM example to include a simple derivative wrt. mesh vertices, but running into several issues.

• Import Error: Cannot import name 'getitem_dataclass' #184 triggers for both install via pip and from master. Going with development branches for now.
• Running the example with JAX instead of torch (from phi.flow import *) diverges immediately with this exception:

File ~/.virtualenvs/science/lib/python3.12/site-packages/phiml/math/_optimize.py:199, in SolveInfo.convergence_check(self, only_warn)
    197             warnings.warn(self.msg, ConvergenceWarning)
    198         else:
--> 199             raise Diverged(self)
    200 if not self.converged.trajectory[-1].all:
    201     if NotConverged not in self.solve.suppress:

Diverged: Solve diverged within -1 iterations using scipy.sparse.linalg.spsolve.

Any suggestions on how to fix it, or another example that we could use for prototyping that uses an unstructured mesh?

[holl-]

Hey @dionhaefner Could you update to the develop branch?

pip uninstall phiflow
pip install git+https://github.com/tum-pbs/PhiFlow@develop
pip uninstall phiml
pip install git+https://github.com/tum-pbs/PhiML@develop

That should hopefully fix the errors.

[dionhaefner]

It fixes the errors at import but not the runtime errors in the cylinder FVM example.

• With phi.flow I get the divergence posted above.
• With phi.torch.flow I get a kernel crash.
• With phi.jax.flow I get an infinite loop (which I guess is also because of a solver divergence, just an unhandled one).

(We'd preferably use JAX so a workaround for that would be most appreciated.)

[holl-]

To verify that the gradient w.r.t. vertex positions works in principle, you can start with this:

mesh = geom.load_gmsh('cylinder.msh', ('y-', 'x+', 'y+', 'x-', 'cyl+', 'cyl-'))

def loss_function(mesh: Mesh):
    return math.l2_loss(mesh.center)

d_mesh: Mesh = gradient(loss_function, 'mesh', False)(mesh)
show(PointCloud(mesh.vertices, -.05 * d_mesh.vertices.center))

Yeah, the divergence error looks like a numerics problem.
To get it running now, you can employ the following hack

v, p = fluid.make_incompressible(v, (), Solve('scipy-direct', suppress=[Diverged, NotConverged]))

Not sure why the solve doesn't converge, it should be easily solveable with a ZERO_GRADIENT boundary condition...