support n_images <-> n_matrices for slice extraction

[McHaillet]

I think we are now only supporting extraction of multiple slices from a single image/volume. For tilt axis refinement it would be handy to supply a batch of n images together with a batch of n matrices, so that each image gets a slice extracted based on its own matrix.

[alisterburt]

yeah, the current API is

• volume: (d, h, w)
• rotation matrices: (..., 3, 3)
• result: (..., h, w)

We already mentioned wanting to add support for extracting the same slices from multiple volumes (c.f. teamtomo/torch-image-interpolation#31) this would add a channel dim...

• volume: (c, d, h, w)
• rotation_matrices: (..., 3, 3)
• result: (..., c, h, w)

If we add support for batches too then we end up with

• volume: (b, c, d, h, w)
• rotation_matrices: (..., b, 3, 3)
• result: (..., b, c, h, w)

This does create a little ambiguity in one reasonable case, multiple slices from multiple volumes:

• volume: (10, 1, d, h, w)
• rotation_matrices: (10, 3, 3)
• result: (10, 1, h, w)

People might naively do

• volumes: (10, d, h, w) (batch of volumes in channel dim)
• rotation_matrices: (10, 3, 3)
• result: (10, 10, h, w)

What do you think about this? The whole API feels clean enough to me but it has this footgun...

[McHaillet]

hmm I get what you mean

People might naively do

* volumes: `(10, d, h, w)` (batch of volumes in channel dim)

* rotation_matrices: `(10, 3, 3)`

* result: `(10, 10, h, w)`

This does feel confusing. Although that also comes from the fact that I am still not fully used to working with channel dimensions. For people used to pytorch that might be more natural.

[McHaillet]

In that case it might be easiest to put statements like this at the start of the functions:

def ..._2d(image):
  if image.ndim != 4:
    raise RuntimeError("Expected image input to have 4 dimensions.. bla bla")

def ..._3d(volume):
  if volume.ndim != 5:
    raise RuntimeError("Expected image input to have 5 dimensions.. bla bla")

Then you know you need to be explicit about it.

But not sure if this is desirable? Makes calling a bit less clean for simple cases. But the added functionality is nice

[McHaillet]

From private discussion:

Alister Burt: I’ll have to let it roll around my head a little - something that comes up when I think about it is that the API in torch-fourier-shift allows you to easily do the same op over a batch which is effectively the same as multichannel support for a single image…
That API does feel simpler… so the question is really: do we need batch support AND channel support, is just batch support enough

Fourier slice with batching only would be

volume: (…, d, h, w)
rotation_matrices: (…, 3, 3)
output: (…, h, w)

If you supply a (d, h, w) volume and multiple matrices you get multiple slices from the same volume, if you supply multiple volumes and one matrix you would get the same slice from multiple volumes
No ambiguity…

edit:
Multiple volumes + multiple matrices would apply each matrix to each volume. (so batch dimension would need to be the same)

[McHaillet]

I realized that the following would be handy to have (and is already support by torch-fourier-shift):

tilts = torch.rand((4, 35, 400, 400))
shifts = torch.rand((35, 2))
aligned = fourier_shift_image_2d(tilts, shifts)  # already supported
# aligned = (4, 35, 400, 400)  

rotations = torch.rand((35, 3, 3))
reconstructions = backproject_2d_to_3d(aligned, rotations)
# expected = (4, 400, 400, 400)

This is good for simultaneously reconstructing multiple sub tilt series with the same alignment parameters.

[alisterburt]

ooh, need to think this through the API for reconstruction... this might be an opportunity to add support for shifts to the API here too

current reconstruction API:

• images (..., h, w)
• rotations (..., 3, 3)
• volume (d, h, w)

maybe batched reconstruction API:

• images (a, b, h, w)
• rotations (b, 3, 3)
• volume (a, d, h, w)

where a and b can both be arbitrary ... shapes?

[alisterburt]

I'm not 100% on that API actually, it doesn't match broadcasting semantics which would imply that rotations get promoted to (a, b, 3, 3) then you still get a (d h w) output

[alisterburt]

what you're arguing for is essentially channel axis support:

• images (..., c, h, w)
• rotations (..., 3, 3)
• volume (c, d, h, w)

Supporting this introduces the same ambiguity/confusion for people who don't know about channels, they use (b, h, w) images with (b, 3, 3) rotations, expect it to work and it either has to

• explicitly fail because shapes aren't strictly compatible (... doesn't match in each input)
• follow broadcasting rules and end up with a (b, d, h, w) reconstruction, each from the same image with b different rotations...

neither feels great to me

[alisterburt]

It really isn't obvious to me how to handle this: an explicit batch dimension also feels weird, it has the same problem as the channel dim

• images (..., b, h, w)
• rotations (..., b, 3, 3)
• volume (b, d, h, w)

[alisterburt]

I'm leaning towards a separate backproject_2d_to_3d_batched function which does

• images (..., b, h, w)
• rotations (..., b, 3, 3)
• volume (b, d, h, w)

It seems that the only way to keep the simple case simple and still provide access to the batched functionality is to have this in a separate function, what do you think @McHaillet ?

[McHaillet]

I absolutely get what you mean with the ambiguity, and we should avoid that. The functionality of batching is just really good to have in here!

I'm leaning towards a separate backproject_2d_to_3d_batched

That makes so much sense, I think we were over complicating this a lot by trying to make a single function that adapts to everything. Let's go for this!

[McHaillet]

Reading back, my current implementation still does not support this specific case

• images (..., b, h, w)

• rotations (..., b, 3, 3)

• volume (b, d, h, w)

Rather only this:

• images (b, ..., h, w)

• rotations (..., 3, 3)

• volume (b, d, h, w)

I.e. the same matrices are applied to each example in the batch.

I think I will leave it for now though, because this use case of separate operations on each volume in the batch feels fairly niche. So far I only see it for the tilt axis angle refinement , but that can easily be mitigated with a for-loop.

The batching as I wrote it now, on the other hand, is immediately handy for all the alignment optimization stuff.

[alisterburt]

Any chance we could switch to

• images (..., c, h, w)
• rotations (..., 3, 3)
• volume (c, d, h, w)

And call it multichannel rather than batch? Makes the semantics easier to reason about

[McHaillet]

My PR #23 made the update to channels!