Composition of affine with free-form deformations

[jacobhinkle]

Is your feature request related to a problem? Please describe.
Currently we can compose affine transforms using regular pytorch stuff, and we can compose displacement fields using lm.interp. It would be useful to compose combinations of those types though, in order to avoid an extra image interp by applying to each to an image, sequentially.

Describe the solution you'd like
I'd like to write differentiable functions for interpolation of "affine then displacement" and "displacement then affine". I have not really decided what I'd like these to be called. If at all possible I'd like the names to be short and clear, like interp and affine_interp.

Describe alternatives you've considered
A similar effect can be achieved currently, for example by lm.interp(lm.affine_interp(I, A, T), h). However, this performs two image interpolations, instead of adjusting the h-field by an affine transform. The other order would involve affine interpolation of a deformation field which is presumably smoother than an image, so this is numerically better in both cases.

Additional context
Note that this might be considered bloat. I haven't thought deeply about it but it seems to me that the kernels are pretty trivial and involve other functions we already have so it shouldn't increase compile time dramatically. In terms of a demonstration, I think it'd be nice to demo a deformation model that does image registration with affine and LDDMM jointly.

[jacobhinkle]

Note that one direction works just fine already; given a displacement field u, you can compose it with an affine transform via lm.affine_interp(u, A, T). So this issue really only applies in the setting where we want to let an affine transform act on an image, followed by a freeform deformation, in which case the transform looks something like the following

v(x) = A*(x + u(x)-o) + o + T - x

where o is the origin and u is a displacement field. This can be done directly in pytorch, perhaps a little inefficiently, using broadcasting to implement o and x, and coordinate-wise slices for the matrix multiply.

I'm willing to implement the pure pytorch version, but I need to think of a simple and clear name for that operation. So far we only have interp and affine_interp and those are pretty clear. One option is to just add this capability to affine_interp, where lm.affine_interp(u,A,T) gives one ordering and lm.affine_interp(A,T,u) interpolates in the other order. I'm a little hesitant to do that because it is a little too automatic, and I would imagine might lead to unexpected behavior. Also there is the issue that the displacement argument applies only for one ordering. The documentation would become quite muddled.

So I am leaning toward something like lm.precompose_affine for naming the function above.

[jacobhinkle]

Once this capability exists, I think it should be demonstrated in a jupyter notebook where we do affine atlas building, then use the affine transforms (without optimizing them) to precompose during LDDMM atlas building. That will do the atlas building with a loss that matches in the original data space. We could compare against interpolating the images with an affine transform then doing atlas building, as is currently done.