Skip to content

Torch-native PDF pipeline + origin finder improvements - #248

Merged
bobleesj merged 15 commits into
electronmicroscopy:devfrom
ehrhardtkm:pdf-torch
Oct 16, 2026
Merged

bobleesj merged 15 commits into
electronmicroscopy:devfrom
ehrhardtkm:pdf-torch

Conversation

@ehrhardtkm

Copy link
Copy Markdown
Contributor

What problem this PR addresses

Follow-on to #177 (the polar transform + PDF pipeline that merged into nanobeam). This PR brings that work onto dev and lands the changes I've made since.

Conceptually unchanged — same Polar4dstem datastructure, same PairDistributionFunction shape, same overall workflow (origin → polar → PDF). The public API should be the same, except for a few minor changes.

Changes since #177:

  • Fully torch-native, dual-backend pipeline. The whole pipeline accepts numpy or torch input on the existing Dataset, device is inferred from the input, and there are no GPU↔CPU round-trips inside the pipeline. numpy vs tensor input give identical results.

  • Background fitting. Rewritten with an analytic Jacobian that is noticeably faster than the previous implementation. Added a batched fit_bg_batched path for full-dataset use, which the downstream PDF widget (PR3) needs to recompute backgrounds interactively. Also includes some small fixes I picked up while running this on real data.

  • Origin finding. The existing finder's hardcoded margin requirement caused it to error out on smaller detectors, so I made it the margin dynamic and added a 2D-quadratic subpixel refinement (renamed to find_origin_angular_grid). I also added a secondary option, find_origin_angular_descent, a margin-free, COM-anchored descent finder. It's slower but I've found it slightly more reliable on noisier data. Both are dispatched from Polar4dstem.from_data(origin_method="grid"|"descent").

  • Tests in tests/diffraction/test_polar.py are parametrized over both origin finders and add coverage for the new code paths.

Out of scope (follow-on PRs to shortly follow):

  • Elliptical-distortion fit + reciprocal-pixel calibration → PR2, stacks on this branch.
  • ShowPDF widget vendor → PR3.

What should the reviewer(s) do

Merge after review. Reviewing commit-by-commit is likely easier than the unified diff — each commit is a single conceptual change.

  • This PR introduces a public-facing change (e.g., figures, CLI input/output, API).
    • For functional and algorithmic changes, tests are written or updated.
    • Documentation (e.g., tutorials, examples, README) has been updated.
    • A tracking issue or plan to update documentation exists.
  • This PR affects internal functionality only (no user-facing change).

@bobleesj

Copy link
Copy Markdown
Collaborator

@ehrhardtkm before I review, it seems like being merged to upstream/dev instead of upstream/nanobeam. Has upstream/nanobeam containing pdf code merged to upstream/dev?

@bobleesj

bobleesj commented Jun 12, 2026 •

Copy link
Copy Markdown
Collaborator

Another question I had was from the PR comment was - what's approx speed/improvement given various sizes? Mostly likely it will be either 512x512x192x192 or 256x256x192x192 to be practical, just ballpark numbers.

Ref:

@ehrhardtkm

Copy link
Copy Markdown
Contributor Author

@ehrhardtkm before I review, it seems like being merged to upstream/dev instead of upstream/nanobeam. Has upstream/nanobeam containing pdf code merged to upstream/dev?

@bobleesj upstream/nanobeam hasn't been merged into upstream/dev and no PDF code exists on dev yet. The issue with merging nanobeam into dev is that nanobeam has diverged significantly from dev (50 commits ahead and over 400 commits behind` and including other new modules (MAPED, strain mapping) that aren't all fully torch-native.

This PR carries just the PDF code from PR #177 plus improvements/fixes since then, which are listed in this PR's description. If nanobeam ever gets merged into dev, we should take dev's version of the PDF files.

@bobleesj

Copy link
Copy Markdown
Collaborator

@ehrhardtkm ah-ha divergence! Often inevitable esp for scientific code - maybe before the hackathon, if possible, we shall have nanobeam merged to dev. What do you think? MAPED, strain, PDF, these are super important for the upcoming.

@ehrhardtkm

Copy link
Copy Markdown
Contributor Author

@bobleesj I don’t think nanobeam can go into dev as a plain merge. Apart from the PDF code it carries 5k lines (MAPED, strain mapping, a file reader, imaging_utils changes). Also, from what I see, its imaging_utils fixes the same sub-pixel cross-correlation bug as dev’s #281, plus open PRs #284 and #97 are also on newer version of the strain code. It would be cleaner for at least the strain and PDF code to go into dev in these separate merges, I think.

@bobleesj
bobleesj self-requested a review October 15, 2026 07:51
@bobleesj

Copy link
Copy Markdown
Collaborator

@bobleesj I don’t think nanobeam can go into dev as a plain merge.

happy to merge this PR to dev and we can incrementally port things over to dev - it doesn't have to all come from nanobeam into dev. Coding is landscape is changing.

@arthurmccray could you please check if I have permission for merging?

Screenshot 2026-10-15 at 12 51 42 AM

@arthurmccray

Copy link
Copy Markdown
Collaborator

@bobleesj I just made you a maintainer, so you should have merge abilities now. Lemme know if it doesn't work for some reason

@bobleesj

Copy link
Copy Markdown
Collaborator

@arthurmccray thank you.

@ehrhardtkm ah, we merge conflicts!

@ehrhardtkm

Copy link
Copy Markdown
Contributor Author

Adding some more updates since the initial PR description, for posterity:

Low-r correction now takes the density as an input. Below the nearest-neighbour distance there are no atom pairs, so G(r) should follow the straight line −4πρ₀r there and errors in S(k) show up as spurious oscillations instead. damp_origin_oscillations corrects S(k) until G(r) follows that line on [0, r_cut]. ρ₀ is now supplied by the user, and calculate_gr (g(r) = 1 + G/(4πρ₀r)) needs it too.

PairDistributionFunction caches the radial mean I(k) and the background fit. Calling calculate_Gr now it refits the background every time, reuses I(k) only for the same real-space mask, and clears an old g(r) to avoid using stale results.

Because all scan positions are resampled onto one polar grid, the k range ends where the pattern closest to the detector edge runs out. One wrongly found center could therefore shorten it for the whole dataset. It now uses the 0.1th percentile of the positions instead of the worst one, and warns if positions lose data.

Elliptical distortion fit (new). fit_elliptical_distortion fits an amorphous ring with an ellipse; passing the fitted ellipse_params to from_data makes the polar transform sample along that ellipse. this follows py4dstem's algorithm.

Also, some timing benchamrking that @bobleesj asked for, for the full PDF pipeline (origin finding → polar transform → PDF) Before is PR#177:

Size GPU (L40S): before / after CPU (32 threads): before / after
256×256×192×192 6.8 s / 11.8 s 137 s / 85 s
512×512×192×192 91 s / 47 s 560 s / 315 s

Will work on the merge conflicts shortly

@ehrhardtkm

Copy link
Copy Markdown
Contributor Author

@bobleesj Let me know if you have any other comments or concerns!

@bobleesj bobleesj left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Testing review

@bobleesj

Copy link
Copy Markdown
Collaborator

@arthurmccray

Screenshot 2026-10-15 at 9 33 22 PM

Could you please check my status? Seems like my prev review rights were revoked.

@arthurmccray

Copy link
Copy Markdown
Collaborator

can you not do a review and approve it? I made you a maintainer so you should be able to?
image

@bobleesj bobleesj left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

testing approve

@bobleesj

Copy link
Copy Markdown
Collaborator

can you not do a review and approve it? I made you a maintainer so you should be able to?

oh thanks. Sorry Arthur - didn't know we had to submit review for PR approval. All good.

@bobleesj bobleesj left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Merging since it has no side effects on other sets of algos. Thank you!

@bobleesj
bobleesj merged commit 6d9474c into electronmicroscopy:dev Oct 16, 2026
4 checks passed
bobleesj pushed a commit to electronmicroscopy/quantem.widget that referenced this pull request Oct 19, 2026
* feat: add ShowPDF for pair distribution function analysis

Add ShowPDF, an interactive front end to quantem's PairDistributionFunction
(electronmicroscopy/quantem#248) for ePDF analysis of 4D-STEM scans.

- Select scan positions with a painted mask, a square probe, a line band, or a
  line scan with one curve per position, and see I(k) with its background,
  F(k), G(r) or g(r) of the selection.
- The settings are passed to calculate_Gr (fit range, k window, filters, r grid,
  low-r correction). The number density is an input, never estimated.
- Line widths are odd, so a horizontal or vertical band of width w covers
  exactly w rows or columns. A line scan has a position at every scan pixel
  along the line's longer axis.
- Standalone HTML export follows the html-export protocol: "full" keeps the
  k-space and correction controls live; "uint8" also embeds every position's
  radial profile, so the selection stays live too.
- Saved notebooks follow the save-state contract: by default they keep a PNG
  preview instead of the scan image, mask and line-scan stack; save_state=True
  keeps everything. Exports always embed them.
- The Python tests skip until quantem provides the #248 API; the tests on
  exported HTML run in the browser tier (QUANTEM_WIDGET_BROWSER=1).
- docs/tutorials/showpdf.ipynb and datasets.showpdf_ta are included but not in
  the docs sidebar yet: the Ta scan waits for a quantem-data pull request.

* chore: pin the ShowPDF tutorial data to its quantem-data commit

The simulated amorphous Ta scan is now on bobleesj/quantem-data
(4dstem/ta_amorphous_sim, added by its pull request 5 in commit fe907499).
showpdf_ta and showpdf_ta_ground_truth read that commit by default, and the
docs and changelog no longer say the data is unpublished.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants