Maped torch implementation

[henrygbell]

What problem does this PR address?

The existing MAPED class (#169) runs on NumPy/CPU, which makes the alignment + merge pipeline slow on the large 4D-STEM tilt/precession stacks it's meant for. This PR adds MAPEDTorch, a PyTorch port of the full MAPED workflow that runs on GPU (or CPU) and operates directly on torch.Tensor 4D-STEM volumes.

[Describe the problem this PR solves. Focus on inputs/outputs.
Attach screenshots, plots, and before/after comparisons.]

[Show the function signatures, class interfaces, and how a scientist
would use this in a notebook. The API is what a scientist actually
types. Get this right first.]

This PR implements the following methods:

1. preprocess()
   • Computes bright field and mean diffraction patterns for all datasets
2. diffraction_origin()
   • Computes the origin of the mean diffraction patterns
3. diffraction_align()
   • Uses cross-correlation based registration to align the mean diffraction patterns and saves the shifts
4. real_space_align()
   • Uses cross-correlation based registration to align the bright field images and saves the shifts
5. merge_datasets()
   • Apply the shifts computed in the last step and merge the datasets. Returns a dataset_4DSTEM object with the merged data

API — how a scientist uses it in a notebook:

from quantem.diffraction import MAPEDTorch

# datasets: list of torch.Tensor, each shape (R, C, H, W), all same dtype + device
maped = MAPEDTorch.from_datasets(datasets)

maped.preprocess(plot_summary=True)                      # -> mean BF
maped.diffraction_origin(sigma=None, plot_origins=True)  # find DP origins
maped.dscan_align(iterations=5)                          # optional descan correction

maped.diffraction_align(num_iter=30, edge_blend=2.0, upsample_factor=100)
maped.real_space_align(num_iter=30, edge_blend=5, hanning_filter=True,
                       pad_val="median", shift_method="bilinear")

merged_ds = maped.merge_datasets(shift_method="bilinear", batch_size=4)  # -> Dataset4dstem

What should the reviewer(s) do?

Attached is a validation notebook which executes the whole pipeline against ground truth and is the easiest way to review correctness: maped_torch_validation.ipynb

Validation notebook workflow:

1. It builds N synthetic datasets with known real-space and diffraction-space shifts applied per dataset.
2. It runs the pipeline and checks that MAPEDTorch recovers those shifts. The residual center(applied + recovered) should be ≈ 0 for every dataset; the notebook asserts a max residual < 0.2 px for both diffraction_shifts and real_space_shifts.
3. It merges the aligned stack and confirms (via shift-invariant max NCC) that the merged result recovers the underlying structure.

Current results of this PR's MAPED code

# Simulated dataset parameters
N = 5               # number of datasets
NOISE = 0.05        # fraction of peak amplitude added as gaussian noise 
shift_tol = 0.2      # tolerance in pixels for shift residual
ncc_tol = 0.95      # tolerance for merged dataset to ground truth ncc

# Validation results
Diffraction alignment : PASS (residual 0.0715 px)
Real-space alignment  : PASS (residual 0.1012 px)
Merge structure recov.: PASS (NCC 0.9989)

ALL CHECKS PASSED 

Please use this notebook for future updates to the MAPED pipeline to assess correctness.

Please look at:

• The validation notebook above
• GPU device handling

If possible:

• Test on your own data

[bobleesj]

@henrygbell beautiful github workflow. clean simple, thanks for the tests. I will review later this week!

[gvarnavi]

Quick comment on the class naming: Presumably the ambition is to make most (all) quantem functionality torch-first, so I would avoid class names like MAPEDTorch. If you want to keep the existing numpy class functionality, then I would rename that one something like MAPEDNumpy instead.

[bobleesj]

@henrygbell left comments for strengthening API doc to guide humans and some minor formatting

[bobleesj]

we need a good documentation on which parameters (and magnitude) to use. are there any chances where these parameters matter? these human lessons should be documented in general.

p.s. the goal is to not using/writing any parameters when possible - all automated if possible. but that's not really possible in practice.

[bobleesj]

generally for all other function as well - scientists wll run the pubic API in jupyter notebook, they find their results are bad, then they will revert to this API doc,

we want this API doc to provide clues/lessons for them. - smae goes for LLM/AI.

[bobleesj]

as @gvarnavi mentioned - MAPED, don't think we have to call it Maped, etc. MAPED

[bobleesj]

try not to have any extra empty lice spaces - we want code to read like paragraphs, each paragraph, each thesis/reason to exist.

[bobleesj]

numRow, row_coords, etc. keep naming consistent throughout

[bobleesj]

yup, as mentioned, this is a good code block where no extra line spaces, but only add one when there is a reason to add.

[henrygbell]

@bobleesj Thanks for your comments. They have been addressed in the recent commits, I also updated the validation notebook for the MAPEDTorch -> MAPED change:

maped_torch_validation.ipynb