5DSTEM data structure API design discussion

[bobleesj]

Problem

Support 5D-STEM dataset (time-series, tilt, defocus, etc.)

Proposed solution

• Create Dataset5dstem inheriting Dataset5d, inheriting Dataset, inheriting AutoSerialize.
• Support slicing ex) dataset5d[0]
• Support tilt/time-series
• Support virtual images across the stack
• Support __str__ __repr__
• Supportbin, copy
• Support mean, max, min,
• ...More as needed - look forward to feedback from potential users including @TomaSusi (time-series) @smribet (MAPED), @ceobrero (times-series)

---

Ref:

I would be very interested in running SSB direct ptychography on a time series of 4D datasets!

Is there code to load such a dataset from a file? Nion Swift specifically, but getting a 5D numpy array is easy.

Originally posted by @TomaSusi in #119 (comment)

• https://onlinelibrary.wiley.com/doi/10.1111/jmi.13409
• https://github.com/TomaSusi/arina-ptycho

[bobleesj]

And of course tomography @cedriclim1

[bobleesj]

API design progress WIP:

# Construction
Dataset5dstem.from_file(file_path, file_type=None, **kwargs)
Dataset5dstem.from_array(array, name=None, origin=None, sampling=None, units=None,
                        signal_units="arb. units", stack_type="generic", stack_values=None)
Dataset5dstem.from_4dstem(datasets, stack_type="generic", stack_values=None, name=None)

# Properties
data.shape              # (frames, scan_row, scan_col, k_row, k_col)
data.array              # underlying numpy array
data.stack_type         # "time", "tilt", "energy", "dose", "focus", "generic"
data.stack_values       # explicit values for stack axis (or None)
data.virtual_images     # dict[str, Dataset3d] of cached virtual images
data.virtual_detectors  # dict[str, dict] of detector configs
data.origin             # origin for each dimension
data.sampling           # sampling for each dimension
data.units              # units for each dimension
len(data)               # number of frames

# Indexing - shape (4, 6, 7, 3, 5) example
data[2]              # -> Dataset4dstem (6, 7, 3, 5) single frame
data[1:3]            # -> Dataset5dstem (2, 6, 7, 3, 5) substack
data[:, 4, 1]        # -> Dataset3d (4, 3, 5) single scan position
data[:, 1:5, 2:6]    # -> Dataset5dstem (4, 4, 4, 3, 5) scan crop
data[..., 0:2, 1:4]  # -> Dataset5dstem (4, 6, 7, 2, 3) k-space crop

# Iteration
for frame in data:   # yields Dataset4dstem

# Stack reductions -> Dataset4dstem
data.stack_mean()
data.stack_sum()
data.stack_max()
data.stack_min()
data.stack_std()

# Virtual imaging -> Dataset3d (frames, scan_row, scan_col)
data.get_virtual_image(mask=None, mode=None, geometry=None, name="virtual_image", attach=True)
# mode="circle":  geometry=((cy, cx), radius) or (None, radius) for auto-fit
# mode="annular": geometry=((cy, cx), (r_inner, r_outer)) or (None, (r_inner, r_outer))
# Per-frame:      geometry=[((cy0, cx0), r0), ((cy1, cx1), r1), ...]

# Virtual image management
data.show_virtual_images(figsize=None, **kwargs)  # display stored virtual images
data.regenerate_virtual_images()                   # regenerate from stored detectors
data.update_virtual_detector(name, mask=None, mode=None, geometry=None)
data.clear_virtual_images()                        # clear images, keep detectors
data.clear_all_virtual_data()                      # clear images and detectors

# Visualization
data.show()  # raises NotImplementedError - use data[i].dp_mean.show() instead

# Binning
data.bin(2)                   # bins all dims including stack
data.bin(2, axes=(1,2,3,4))   # preserves stack

# Copy
data.copy()

# String representations
repr(data)  # Dataset5dstem(shape=..., dtype=..., stack_type='...')
str(data)   # multi-line formatted output

[bobleesj]

@TomaSusi Please feel free to suggest more if needed. Regarding SSB, you will have a CUDA version as demoed via slack, which will work with this dataset.

[gvarnavi]

Thanks for starting this discussion at @bobleesj !

Quick thoughts:

• Dataset already supports arbitrary indexing (though it creates a new Dataset by design)

import quantem as em
import numpy as np

shape = (6,5,4,3,2)
array = np.arange(np.prod(shape)).reshape(shape)
dataset = em.core.datastructures.Dataset.from_array(array)

dataset[0]
dataset[1:3]
dataset[:,4,1]
dataset[:,1:4,1:3]
dataset[...,0:2,1:3]

for frame in dataset:
    print(dataset)

• I wonder if there's an appropriate abstraction level here to ensure certain quantem functions / classes operate transparently on any Dataset-like object 🤔
• What's this SSB CUDA kernel? 👀 👀