MAPED data merging with Torch by henrygbell · Pull Request #204 · electronmicroscopy/quantem

henrygbell · 2026-04-02T22:47:01Z

What does this PR do?

This PR is the torchified version of PR-169. For my data it speeds the merge_datasets method up by ~10x or more and adds batching.

Example notebook

maped_test_nb_torch.ipynb

Data for that notebook

Download here

For reviewers:

Please test the code on your own datasets! Check the code for accuracy. Let me know if you have any suggestions

TODO next:

Can we optimize this torch version?
dscan alignment

…eeds to be overloaded is set_optimizer for PPLR cases

…to do the matching in set_optimizer instead of parsing in optimizer_params maybe?

… to check: Look at object_models.py and see how the optimizer matching should be handled. It seems like set_optimizers doesn't really do what it's supposed to do.

… probably have to do TV loss computation within the model?

…es. Also overloaded reconnecting optimizers

…well. Only things to ask Corneel about is multiscale res since this adds a significant amount of compute. Should I be doing variable num_samples_per_ray?

…t DDP, clean-up KPlanes, fix up object_models.py since it's insanely cluttered now

bobleesj · 2026-04-23T00:23:13Z

Happy to review and also address the remaining bottleneck. @henrygbell

bobleesj · 2026-04-23T02:45:36Z

@gvarnavi seems like CI tests aren't running on branch PRs, could we also add automated workflows? I want to ensure that this torch refactoring doesn't break the existing maped pytest.

bobleesj

@henrygbell I spent 15-20 mins reviewing. Before I run with real data, I noticed that new utility torch functions were added. It is not required to write unit tests for each (if you want, better if you compare against scipy/numpy) but at least we need to know that we maintain parity after refactoring. Please see the drift PR below where torching was done but no API/accuracy was changed

Ref: https://github.com/electronmicroscopy/quantem/pull/206/changes#r3040892988.

After your feedback, I will test with the data we collected at snsf.

bobleesj · 2026-04-23T02:43:56Z

-    pad_val: str | float = 0.0,
-    shift_method: str = "bilinear",
-):
+class MAPEDTorch(AutoSerialize):


general comment - we need an end-to-end test either with synthetic data or data that used before refactoring w/ torch to ensure nothing broke in this PR. Then I would be more comfortable reviewing this.

second comment is that did the API or is it just torching without any user behavior modified? Please indicate that so that we don't need to worry about backward compatability in your PR comment

bobleesj · 2026-04-23T02:47:25Z

+
+        Parameters
+        ----------
+        edge_blend


pls follow numpy docstring - parameter type missing.

important for doc api doc rendering automated

bobleesj · 2026-04-23T02:49:48Z

+
+        Stores
+        ------
+        self.scales : torch.tensor


yeah here i see paramters

bobleesj · 2026-04-23T02:51:18Z

+    - compute mean BF and mean DP summaries,
+    - choose/find diffraction origins,
+    - align diffraction space and real space,
+    - merge datasets into a single composite Dataset4dstem.


bobleesj · 2026-04-23T02:52:07Z

-            stack[ind] = np.fft.ifft2(F * ramp).real
+        Parameters
+        ----------
+        origins


numpy doc, same comment

bobleesj · 2026-04-23T02:52:14Z

-                mode="constant",
-                cval=0.0,
-                prefilter=False,
+        Stores


stores or returns?

bobleesj · 2026-04-23T02:55:48Z

+        return w
+    if alpha >= 1:
+        return torch.hann_window(N, device=device, dtype=dtype)
+


just code comment - no extra whitespace needed, fewer lines better in general so that we don't have to scroll too much. I

…ositionModel ABC, make sure to have a property for which kind of tensor decomposition method is being used. SO3Params are moved to a different file, thinking of making a kplanes_utils.py. Starting reorganization of object_models.py to have ObjectINR and ObjectTensorDecomp

…p-level Tomography

… of parameters now that helps with type-setting. The main reason for having model_base.py as is it is right now is if we ever wanted to go do TensoRF or something just to validate

…ry parsing

…numpy version.

This reverts commit 140c563.

…nobeam

…on to stop forced switch to CPU memory

bobleesj · 2026-06-03T01:13:28Z

@henrygbell seems like we are having some PR synching issues here as well.

A few things to check

upstream/namoeam is getting old, it is diverging from upstream/dev. Why don't you just make MAPED code merged to upstream/dev after making upstream/nanobeam merged to upstream/dev`? Please chat with @cophus

This `upstream/nanobeam' is getting a bit too old where problems are starting to occur

#235 (comment)

henrygbell and others added 19 commits April 2, 2026 15:06

Converted MAPED code to torch, added batching

ff2dcfc

Cleaned up comments, getting ready for PR

991c9cb

Removed some imports that we don't need

788cf55

Fixed hann_filter line in real_space_align

a551193

Added descan alignment using cross correlation

df581dc

Added k-planes model

2148679

Added PPLR stuff

c299ca8

object_models optimization setting is working well. Only thing that n…

9ac27a9

…eeds to be overloaded is set_optimizer for PPLR cases

Optimizing, set_optimizer is just default to Adam now, probably need …

5f40d5a

…to do the matching in set_optimizer instead of parsing in optimizer_params maybe?

KPlanes Tilted claude implementation, need to talk to Corneel. Things…

7f51dc5

… to check: Look at object_models.py and see how the optimizer matching should be handled. It seems like set_optimizers doesn't really do what it's supposed to do.

Added TV loss for PPLR models, I don't like this solution though will…

af140be

… probably have to do TV loss computation within the model?

Merge branch 'electronmicroscopy:dev' into optmixin_pplr

04bd1dc

object_models.py now has tv_loss for both KPlanes and INR architectur…

73c0d30

…es. Also overloaded reconnecting optimizers

KPlanes with R9+SVD parameterization, everything seems to be working …

35157fb

…well. Only things to ask Corneel about is multiscale res since this adds a significant amount of compute. Should I be doing variable num_samples_per_ray?

Everything seems to be working; only things to do is to take a look a…

1c89ad5

…t DDP, clean-up KPlanes, fix up object_models.py since it's insanely cluttered now

New TV loss function

4215d6e

TV volume -- needs significant refactoring everywhere

aebf801

DDP Fixes for PPLR stuff

2579847

Changes

a59baac

bobleesj reviewed Apr 23, 2026

View reviewed changes

cedriclim1 added 7 commits April 23, 2026 11:31

Removed some the _unwrap dependencies. Added ObjectTensorDecomp on to…

e33fec1

…p-level Tomography

Revamped model_base.py to cover type-hinting stuff. KPlanes has a set…

15bdcd4

… of parameters now that helps with type-setting. The main reason for having model_base.py as is it is right now is if we ever wanted to go do TensoRF or something just to validate

Final changes prior to draft PR

0c71a5d

Fixed typo

99fdd88

Small change to set_optimizer in OptimizerMixin to allow for dictiona…

4a44826

…ry parsing

Pretraining warning on ObjectTensorDecomp

ede9d56

bobleesj mentioned this pull request Apr 24, 2026

Enable pytest CI workflows in <upstream/branch> to ensure API and algo remain solid #215

Closed

cophus and others added 29 commits June 2, 2026 10:57

adding docstrings

649e308

Converted MAPED code to torch, added batching

9c01697

Cleaned up comments, getting ready for PR

2d1d0e1

Removed some imports that we don't need

9992c45

Fixed hann_filter line in real_space_align

1f640ae

Added descan alignment using cross correlation

494ebc7

Changed docstrings to all be in numpy format

9a0254e

Changed the plotting of MAPED.real_space_align to be consistent with …

b185ba2

…numpy version.

initial construction of polar4dstem class

95be310

Changing sampling direction for polar4dstem

b16746f

initial commit for RDF class

3bd18c9

initial pdf code + f parameters

698c852

origin finding and pdf updates

12d80a2

pixel calibration bug fix

b0e4af5

make torch native and general cleanup

6af80ce

additional cleanup

f156f06

added comments for increased clarity

142f0d6

minor fixes

6f744b1

origin finding speedup

d0b6680

make calculate_Gr use precomputed bg

db08dcf

refactor and row/col convention cleanup

df48ac8

docstring/API cleanup and origing finding speedup

47694f2

remove duplicate gaussian code and restore tomo utils

47ff589

Autocorrelation dscan alignment added

7f8843d

Add fast vectorized autocorrelation dscan alignment.

b3d76a8

MAPED direct fitting descan correction implemented.

140c563

Revert "MAPED direct fitting descan correction implemented."

05f6b01

This reverts commit 140c563.

Merge remote-tracking branch 'henrygbell/nanobeam' into henrygbell-na…

82f17f7

…nobeam

Updates the output merged dataset to use the from_tensor initializati…

ff8ea70

…on to stop forced switch to CPU memory

Conversation

henrygbell commented Apr 2, 2026

What does this PR do?

Example notebook

Data for that notebook

For reviewers:

Uh oh!

bobleesj commented Apr 23, 2026

Uh oh!

bobleesj commented Apr 23, 2026

Uh oh!

bobleesj left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bobleesj commented Jun 3, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

9 participants