JuliaSpace / SatelliteToolbox.jl

A toolbox for satellite analysis written in julia language.
MIT License
248 stars 33 forks source link

Compatibility of transformations code with Automatic Differentiation tools (Zygote) #89

Open srikumarks opened 1 year ago

srikumarks commented 1 year ago

Zygote is the most used Automatic Differentiation tool and it doesn't support static arrays due to their use for mutations [^1]. DCM and Quaternion structures used in frame conversion code, which I'm trying to use with Zygote, use SMatrix and SVector for efficiency reasons as stated in the code comments. However that makes them incompatible with Zygote.

Having the transformation code be AD compatible is perhaps useful (broadly true for most libraries) and so posting this issue in that spirit. One way to address this in the library is to perhaps also support a third rotation representation that uses ordinary vectors/matrices (could be DCM based representation, but just uses a vector of 9 numbers or an ordinary Matrix instead of an SMatrix{3,3}).

This impacts the angle_to_rot functions defined in the ReferenceFrameRotations package - ReferenceFrameRotations/uTarN/src/conversions/angle_to_rot.jl. It looks at the moment that having angle_to_rot support the ordinary matrix as a return type should suffice to get all the r_eci_to_eci and r_eci_to_ecef and such functions to also support them when overloaded with the first argument being the type Matrix. Will also have to check whether the "multiply with vector to transform it" steps will also work properly.

(I might implement this, but posting an issue in advance for any comments or tips that could help with it ... with the best case being "it already exists"!)

[^1]: Apparently. I'm not conversant with all the details of why.

ronisbr commented 1 year ago

Hi @srikumarks !

Thanks for the information! I will check what can we do. The incompatibility with Zygote is caused because DCM and Quaternion are immutable? What is the source of the problem?

srikumarks commented 1 year ago

The incompatibility with Zygote is caused because DCM and Quaternion are immutable? What is the source of the problem?

Not because they're immutable (being immutable is great), but because they use static arrays - SMatrix and SVector as the underlying representation. These don't get adjoints in Zygote and I haven't had much luck rolling one on my own too. Nothing I try in terms of a custom adjoint seems to work. I referred to this thread - https://github.com/FluxML/Zygote.jl/issues/570 .