Automatic differentiation to calculate Jacobians

[mikeingold]

The codebase currently uses an inelegant but reasonable finite difference approximation to calculate the Jacobian of a geometry's parametric function.

https://github.com/mikeingold/MeshIntegrals.jl/blob/66ed919150b3774bf09dabab60476869c5514493/src/utils.jl#L5-L15

At time of writing, there didn't appear to be an automatic differentiation package that could compute the same thing out-of-the-box. These parametric functions return a Point rather than a simple number. I believe this could be addressed by defining chain rules in Meshes.jl, but I'm open to other solutions.

[mikeingold]

Meshes.jl now uses Zygote.jl as an AD-of-choice. Ideally, all Geometrys in Meshes.jl would have diff rules established so that this Jacobian function would flow from that.

[mikeingold]

Merged with #37

[mikeingold]

Since differential was generalized in #101, I'm going to un-merge this Issue from #37.

[JoshuaLampert]

In #115 we agreed that it would be best to implement the AD support over in Meshes.jl (we could also think about transferring the issue?).
I haven't tested anything AD related in the context of Meshes.jl and MeshIntegrals.jl (I only have a bit of experience using Zygote.jl, ForwardDiff.jl and Enzyme.jl in other contexts), but I see that the fact that the parametric functions of the geometries return Points complicates the issue and probably necessitates implementing custom AD rules via ChainRules.jl.

I’m also not sure if writing these directly for parametric functions is the best path, or if we can write them for lower-level functions and have that support roll up automatically.

I would opt for minimizing the number of rules to implement, i.e. not needing to implement a rule for each geometry, but rather write something more low level. This seems to be easier to maintain and extend.
These are just my initial thoughts without doing any experimenting myself. My familiarity with ChainRules.jl is pretty low. So I'm not sure if this approach is feasible.
Getting some advice from @kylebeggs would be awesome.

[mikeingold]

... but I see that the fact that the parametric functions of the geometries return Points complicates the issue and probably necessitates implementing custom AD rules via ChainRules.jl.

Yes, I'm definitely not an AD expert but this is one of the reasons that my intuition tells me we'll need custom AD rules. Parametric functions map from floating point tuples/vectors to Meshes.Point, but their derivatives should in my opinion map from floating point tuples/vectors to Meshes.Vec.

Again, though, I'm not sure if there's a simpler way to define rules like the derivative of Point(...) is Vec(...), etc and get those to propagate automatically. This would be especially beneficial in that it wouldn't require AD rule updates any time the corresponding functions are updated.

[kylebeggs]

@JoshuaLampert is spot on. I highly encourage rules to be implemented at the lowest level possible. Usually, once you have those, everything else will "just work" and then the only reason to write more rules is for efficiency (i.e. you know some computations in the graph for autodiff can be skipped or something via knowledge of the math behind the operation). I think for Meshes.jl, the hardest part will really be defining what it means for some of these types/methods to be differentiated, as @mikeingold alluded to.

Although the rules should go in Meshes.jl, we can use MeshIntegrals.jl as a catalyst for figuring out what rules even need to be defined to begin with. We can create a matrix of to-dos like there was in the early stages of the package in the README.

[kylebeggs]

I suppose that would be just any types <:Geometry

[mikeingold]

We can create a matrix of to-dos like there was in the early stages of the package in the README.

For reference, the Support Matrix you’re probably referring to is now hosted on our Documenter site.