Add polar transform and pair distribution function analysis

[ehrhardtkm]

What does this PR do?

Implementation of polar coordinate transforms and computing pair distribution functions (PDFs) for 4DSTEM diffraction data. The PDF relies on a polar transformed data structure, which has been implemented in Polar4dstem. The code will be updated in the future with other useful tools, like beam stop masking and elliptical corrections.

The workflow takes a 4DSTEM dataset through the following steps:

1. Origin finding — Automatically locate the center of each diffraction pattern by minimizing angular intensity variation (auto_origin_id)
2. Polar transform — Convert diffraction patterns from Cartesian to polar coordinates using torch grid_sample for GPU acceleration (dataset4dstem_polar_transform)
3. PDF calculation — Compute the reduced PDF G(r) from the azimuthally averaged intensity, including background fitting, structure factor extraction, and a windowed sine transform (calculate_Gr)
4. Density estimation (optional) — Iteratively estimate atomic number density and dampen unphysical low-r oscillations using the Yoshimoto-Omote method (estimate_density, calculate_gr)

New files

• polar4dstem.py — Polar4dstem dataset class for polar-coordinate 4D-STEM data, plus the origin finding and polar transform functions
• polar.py — Contains PairDistributionFunction class. Handles the full PDF pipeline: radial averaging, background fitting, structure factor calculation, windowed sine transform to G(r), iterative density estimation, and conversion to g(r)
• tests/diffraction/test_polar.py — Tests covering construction, transforms, background fitting, PDF calculation, density estimation, and end-to-end workflows

Modified files

• dataset4dstem.py — Added polar_transform() convenience method on Dataset4dstem
• __init__.py files — Exposed new classes and modules

A companion tutorial notebook (PDFanalysis_simulatedTa.ipynb) demonstrating an example workflow on simulated amorphous Ta data is available.

---

What should the reviewer(s) do?

Provide feedback on the API design and implementation. Key areas for review:

• Polar4dstem class design and its relationship to Dataset4dstem
• PDF pipeline correctness (background subtraction, structure factor normalization, Fourier transform)
• Ensure the example notebook runs correctly
• Note any design patterns that were not followed or potential bugs

Checklist items:

• This PR introduces a public-facing change (e.g., API, CLI input/output).
  • 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).

---

Reviewer checklist

• Tests pass and are appropriate for the changes
• Documentation and examples are sufficient and clear
• Imported or adapted code has a compatible license
• The implementation matches the stated intent of the contribution

[bobleesj]

I'm happy to review this PR since this would naturally extend along Show4D widget that you could extend and I was a PDF group for x-ray.

If possible, if you could copy/paste API design here as a comment and attach screenshots (plots, before and after) - these can be helpful as well. API design, as a public facing component, is becoming more important since those are what humans mostly these days.

The attached notebook is always good for the reviewers, but PRs contain global URLs so it's a great place to market/demo to the whole group and the world.

[ehrhardtkm]

Step 1: Find diffraction centers (method on Dataset4dstem)

# Coarse-to-fine origin finding by minimizing angular std
origins = ds.auto_origin_id(
    *, ellipse_params=None,             # (a, b, theta_deg) distortion correction
    num_annular_bins=180,
    radial_min=0.0,                     # pixels
    radial_max=None,                    # pixels (defaults to max radius)
    radial_step=1.0,                    # pixels
    two_fold_rotation_symmetry=False,
    device="cpu",
)
# -> NDArray (scan_y, scan_x, 2) of (row, col) origins in pixels

Step 2: Polar transform (method on Dataset4dstem)

# Cartesian -> polar via torch grid_sample
polar = ds.polar_transform(
    origin_array=None,                  # None | (2,) | (scan_y, scan_x, 2)
    ellipse_params=None,                # (a, b, theta_deg) distortion correction
    num_annular_bins=180,
    radial_min=0.0,                     # pixels
    radial_max=None,                    # pixels
    radial_step=1.0,                    # pixels
    two_fold_rotation_symmetry=False,
    name=None, signal_units=None,
    scan_pos=None,                      # if set, returns a torch.Tensor from a single scan position
    device="cpu",
)
# -> Polar4dstem (or torch.Tensor if scan_pos given)

Step 3: PDF construction

from quantem.diffraction import PairDistributionFunction

pdf = PairDistributionFunction.from_data(
    data,                               # NDArray (2D/4D) | Dataset2d | Dataset3d | Dataset4dstem | Polar4dstem
    *, find_origin=True,
    origin_row=None, origin_col=None,
    ellipse_params=None,                # (a, b, theta_deg) distortion correction
    num_annular_bins=180,
    radial_min=0.0, radial_max=None, radial_step=1.0,
    two_fold_rotation_symmetry=False,
    device="cpu",
)

Step 4: PDF pipeline

# Azimuthal average -> I(k)
pdf.calculate_radial_mean(mask_realspace=None, returnval=False)

# (optional, called automatically by calculate_Gr if needed)
pdf.fit_bg(Ik, kmin, kmax)  # -> (bg, f)

# F(k) + windowed sine transform -> G(r)
# calls calculate_radial_mean + fit_bg internally if not already run
pdf.calculate_Gr(
    k_min_fit=None, k_max_fit=None,       # bg fitting range (1/A)
    k_min_window=None, k_max_window=None,  # Lorch window range; defaults to fit range (1/A)
    k_lowpass=None, k_highpass=None,        # (1/A)
    r_min=0.0, r_max=20.0, r_step=0.02,    # (A)
    mask_realspace=None,
    damp_origin_oscillations=False,         # runs estimate_density if True
    density=None, r_cut=0.8,
    returnval=False,
)  # -> [r, G(r)] if returnval

# (optional) Yoshimoto & Omote iterative density estimation
pdf.estimate_density(density=None, r_cut=0.8, max_iter=40, tol_percent=1e-4)
# -> (rho0, Fk_damped, G_corrected)

# G(r) -> g(r)
pdf.calculate_gr(density=None, r_cut=0.8, set_pdf_positive=False, returnval=False)
# -> [r, g(r)] if returnval

Plotting

pdf.plot_pdf_results(which=("reduced_pdf",), qmin=None, qmax=None,
                     rmin=None, rmax=None, figsize=(6,4), returnfigs=False)
# which: "radial_mean" | "background_fits" | "reduced_sf" |
#        "reduced_pdf" | "pdf" | "oscillation_damping"

# Individual plots (all accept qmin/qmax or rmin/rmax, figsize, returnfig)
pdf.plot_radial_mean()
pdf.plot_background_fits()
pdf.plot_reduced_sf()
pdf.plot_reduced_pdf()
pdf.plot_pdf()
pdf.plot_oscillation_damping()

---

Typical usage

import quantem as em
from quantem.diffraction import PairDistributionFunction

# Load 4DSTEM data
ds = em.core.io.read_4dstem("path/to/data.h5", file_type="arina")

# Separate steps:
origins = ds.auto_origin_id()
polar = ds.polar_transform(origin_array=origins)
pdf = PairDistributionFunction.from_data(polar)

# Convenient one step alternative (auto origin + transform done internally):
# pdf = PairDistributionFunction.from_data(ds)

# PDF analysis
r, Gr = pdf.calculate_Gr(k_min_fit=0.02, k_max_fit=2.0, returnval=True)
r, gr = pdf.calculate_gr(returnval=True)
pdf.plot_pdf_results(which=["reduced_pdf", "pdf"])

Updated notebook is available: PDFanalysis_simulatedTa.ipynb

[ehrhardtkm]

The auto origin finding can be evaluated by plotting the polar transform:


When calculating the G(r), the user can either let the I(k) and background fitting run on it’s own behind the scenes by directly calling rdf.calculate_Gr() or compute these steps separately. Either way, they will produce the same result, given the same parameters.

Both cells above yield the same output, which can be visualized with the plotting functions.


The small oscillations near 0 of the G(r) can be damped by setting damp_origin_oscillation=True. The changes made by the oscillation damping process can be visualized with the ”oscillation_damping” plotting option.