refactor: dataset accepts torch (cpu/gpu) array-like for streamlined 4DSTEM algo & workflows

[bobleesj]

Problem

The main goal of the Dataset class is to containerize and retrieve the underlying data, including metadata.

Currently it is forced to numpy. This deviates from the main workflow for 4D-STEM reconstructions, which primarily rely
on torch or cupy - both for using quantem directly and for algo dev in forks.

Ref:

if not isinstance(arr, np.ndarray):
    raise TypeError(                                                                                                     
        "Dataset requires a NumPy array (CuPy is not supported on this branch)."
    )                                                                                                                    

Two cases where I don't get to use Dataset4dstem:

Case i) When a 4D-STEM dataset is loaded via the CUDA path (#135),
wrapping it in Dataset4dstem requires converting from VRAM to CPU RAM. With a duck-typed dataset, the data can stay in
VRAM and the widget (Show4DSTEM) reads it directly, eliminating the RAM↔VRAM round trip.

Case ii) Live SSB. Detector frames stream into VRAM and are processed on the GPU - they never touch CPU RAM.

Proposed solution

Remove isinstance(arr, np.ndarray) and replace it with a minimal duck-type check.

Accept any object exposing shape, ndim, dtype, and the standard reduction methods (.sum, .mean, .max, .min).

[bobleesj]

@gvarnavi tagging - I believe you've mentioned ducktyping back then so formalizing so that I get to use the data structure architecture provided in quantem.

[arthurmccray]

Thanks Bob! With regards to this and #225 (@bobleesj @cedriclim1), I think it's time to reconsider building Dataset around torch tensors rather than numpy arrays. It seems like there are significant performance benefits to loading directly from disk -> VRAM, and should always be fine to have torch tensors on CPU for folks who don't have a GPU. I am against further fragmenting our dataclasses, e.g. I don't like the idea of a dataset4dstem_torch or similar, so this seems like the only solution.

Let's plan to talk about this next week. We can then write up a little design doc/plan to share on the Slack, and at next dev meeting make a formal decision and figure out who precisely is going to be doing what.

Until then, I have one clarifying question with regards to speedups and loading:

• it seems to be much faster to read from disk directly into VRAM rather than disk -> RAM, correct? If you're going to be doing stuff in RAM anyways, is there any benefit of doing disk -> VRAM -> RAM? Or would that be slower than direct disk -> RAM?
• I guess also, are there any other situations that would currently require lots of transfers between VRAM and RAM that having our Dataset be torch-native wouldn't solve?

And my rough thoughts on what it would look like to have our Dataset be built around torch:

• dset.array would be a torch.Tensor (I don't think we would rename .array to .tensor but we could... it would certainly be more explicit)
  • We would have a .numpy() method (mirroring torch syntax) that returns a plain numpy array of the data
  • the @array.setter would do the conversion to torch when setting, so we could do something like dset.array = some_numpy_function(dset.numpy()) for cases that still require numpy functions in user code.
• The relevant metadata currently stored as numpy arrays (such as .sampling) would now be tensors as well
• All internal tensors would be on a new .device attribute
• we should consider inheriting from nn.module, and using buffers for .sampling and others.
  • It would also be fine to just keep everything together by writing our own .to method, and I'm not sure which is the correct answer. I am currently leaning towards not inheriting from nn.module as it can be trickier to integrate with Autoserialize, but by using torch tensors we might have burned that bridge already.
• All the various methods (e.g. .bin()) will have to be double-checked and modified
• Testing for this PR will be a pain because it touches literally everything, but my hope would be that it actually doesn't break many modules, cuz our current workflow (of moving things to torch where needed) won't actually be affected.

Also tagging other core dev members so that they're up-to-date and can weigh in if interested at this point @cophus @gvarnavi @smribet

[bobleesj]

Thanks @arthurmccray for the great insights. My responses to the points mentioned:

read from disk directly into VRAM rather than disk -> RAM, correct?

RAM is slightly used to stage the compressed 4DSTEM data. The workflow is disk (1-2 GB) -> host RAM (1-2 GB) -> VRAM (19.3 GB, cupy)

is there any benefit of doing disk -> VRAM -> RAM? or Would that be slower than direct disk -> RAM?

VRAM->RAM can be useful for memory management: Given a single GPU (RTX6000), use VRAM to load/bin dozens of 4DSTEM datasets from disk to VRAM under a min, then stage them to RAM so that VRAM memory is saved.

are there any other situations that would currently require lots of transfers between VRAM and RAM that having our Dataset be torch-native wouldn't solve?

I haven't encountered a situation yet. Show4DSTEM` does a lot of (torch) VRAM->RAM transfer (extract one DP at a time) and it remains real-time.

Testing for this PR will be a pain because it touches literally everything,

Regarding this, I'd like this experience to be less painful for everyone.

How about rolling out incremental PRs where initial PRs focus on the two constraints?

1. don't break existing notebooks and workflows
2. don't change much but do allow dataset4dstem to hold torch so that MAPED/PDF/widgets can benefit immediately and potentially use them on the microscope.

[arthurmccray]

How about rolling out incremental PRs where initial PRs focus on the two constraints?

1. don't break existing notebooks and workflows

2. don't change much but do allow `dataset4dstem` to hold torch so that MAPED/PDF/widgets can benefit immediately and potentially use them on the microscope.

That seems like a good idea, we could add the device and to methods and type everything for NDarray | torch.Tensor. The main issue would be the bin and similar methods handling either, though I guess they could just be written for torch and, if given a numpy array, just convert to and back at the end.

Part of me thinks that, if everyone is on board, we could get away with going fully numpy -> torch all in one PR, but it would definitely require everyone's participation.

[cedriclim1]

Thanks for the comments @arthurmccray, and also just weighing in too:

dset.array would be a torch.Tensor (I don't think we would rename .array to .tensor but we could... it would certainly be more explicit)

I think the end product should always be a numpy array, so I think .tensor makes a lot of sense here IMO. I think for general users, they still will probably prefer having the numpy operations and have more familiarity with this. Maybe what we can have is a def from_torch/def from_tensor class method in Dataset4dstem? This way .array will still always return a numpy array, we will just have to do some checking if .tensor is not empty.

I think this way preserves and minimally impacts all the features that have been built so far i.e., Ptycho, MAPED, and other 4DSTEM related projects. This also allows whatever Bob wants to do in terms of I/O speed, and others can start implementing this in their own workflows if I/O is a problem. Eventually, I hope everyone just starts using tensor's off the bat and we can revisit making everything in pure torch. The only hiccup I guess is that everything has to be called to .tensor, but I like being more explicit!

We would have a .numpy() method (mirroring torch syntax) that returns a plain numpy array of the data
the @array.setter would do the conversion to torch when setting, so we could do something like dset.array = some_numpy_function(dset.numpy()) for cases that still require numpy functions in user code.

I think for this specific example, I'm liking the .tensor property more.

The relevant metadata currently stored as numpy arrays (such as .sampling) would now be tensors as well

Why does the metadata have to be stored as a tensor? I think for metadata related stuff that isn't required for torch can still remain as numpy or generic Python data types.

All internal tensors would be on a new .device attribute

Having the class method def from_torch would also help with this I think.

we should consider inheriting from nn.module, and using buffers for .sampling and others.

I think we shouldn't do this and not overcomplicate the data structure. All downstream stuff that require ML dataset models should just inherit nn.Module there.

Let me know if this sounds crazy.

[cedriclim1]

I guess also additionally, preserving .array as a numpy array will keep all the functionality intact from the base class dataset.

[arthurmccray]

Well another option (maybe what you're alluding to) is to maintain a .array property that always creates and returns a numpy version of the .tensor attribute. An issue with this is that it isn't a view into the underlying dataset (could cause confusion), but it would be an easy way of maintaining backwards compatibility while also enabling everything to be torch-native under the hood.

Why does the metadata have to be stored as a tensor? I think for metadata related stuff that isn't required for torch can still remain as numpy or generic Python data types.

I was thinking that metadata should be the same type as the underlying data because of how it's used, e.g. subtracting off .origin from coordinates (which will also be tensors) or multiplying .sampling by a coordinate grid to get it in calibrated units. If everything we're doing is in torch by default, it doesn't make a ton of sense to have to do torch.from_numpy(self.sampling, device=self.device) all over the place.

In any case, I don't think anything is set in stone yet, and we should definitely brainstorm this out on a whiteboard next week :)

[cedriclim1]

Re: silently failing on .array, we can just use .flags.writeable to make sure inplace operations won't silently fail.

@property
def array(self):
    arr = self.tensor.detach().cpu().numpy()
    arr.flags.writeable = False
    return arr

[arthurmccray]

Plan of PRs after discussing with @bobleesj and @cedriclim1

1. non-breaking just adding basic torch support
   1. in base Dataset dset.tensor # is None if not set
      1. Idea is that Dataset will have .tensor and .array, but only one of them will be set and the other will be None.
      2. Subclasses (e.g. Dataset3d) for which we aren't implementing torch stuff yet, they will only have their .from_array classmethod that sets .array but leaves .tensor empty, so it shouldn't break anything.
   2. in Dataset4dstem: dset.from_tensor classmethod
   3. base Dataset: 1. dset.array being a property/return thing above
      2. dset.device (just return self.tensor.device)
      3. dset.to (just moving dset.tensor)
      4. dset.numpy() method (will be required in the future, make it now so people get used to it)
   4. make sure Autoserialize still works...
2. Filling in functionality of methods for torch without breaking
   1. bin() crop() etc. making a torch-verion, in parallel with the numpy version (non breaking still)
   2. Move fourier_sampling() somewhere else (who put it in Dataset?)
   3. add deprecation warnings that things will become torch native, basically only on .array switch to .numpy() for getting a read-only numpy version of the array, or .tensor for the view into the dataset.
3. Making everything fully torch
   1. only thing breaking is in place operations on .array,
   2. add the arr.flags_writable=False
   3. removing numpy methods for bin and related
   4. add a from_data classmethod that accepts tensors or numpy array (and converts to a tensor)
   5. move metadata items like sampling to torch.Tensor if not already
   6. again make sure Autoserialize still works 4. make sure everything above works for a bit....
4. remove .array and require .numpy()

Plan! @bobleesj will do step 1 and we'll go on from there!

[bobleesj]

@arthurmccray Thank you Arthur for the clear plan. Very excited to get this rolling based on the blueprint and see where things end up overall.