Skip to content

Speed up the dynamical master-pattern simulation - #20

Merged
CSSFrancis merged 4 commits into
mainfrom
feat/master-speed
Oct 27, 2026
Merged

CSSFrancis merged 4 commits into
mainfrom
feat/master-speed

Conversation

@CSSFrancis

@CSSFrancis CSSFrancis commented Oct 27, 2026 •

Copy link
Copy Markdown
Member

Speeds up the dynamical master-pattern simulation (the bloch_prism backend and the symmetry fold), so masters at npx 300-512 are practical. Every change leaves the master the same: bit for bit, or to float32 rounding where a different but symmetry-equivalent direction is evaluated.

  • Depth integration without forming L_gh a depth at a time. The intensity is sum_n lambda_n psi_n^H S psi_n: the waves are stacked and multiplied by S once, instead of a batched outer product per depth. On a Cu master at 15 kV, complex64, npx 512: ~7 -> 2.7 s besides matrix_exp; the master the same to 1.3e-7.
  • The symmetry fold on the device. Each direction takes its image under the point group with the largest (z, y, x), rounded as _dedup_on_device rounds, instead of orix's in_fundamental_sector on the CPU: 0.5 s at npx 512 against 2.9 s. The same number of representatives as orix for Fm-3m, P6_3/mmc and Pnma, fewer for P-1; each is an image of its directions (test_every_representative_is_a_symmetry_image_and_none_is_spare).
  • Beam pairs gathered from small matrices. U(h_g - h_h), |U| and S(h_h - h_g) for every pair of beams are made once per intensities call (from the tables as they are then, so a u_lut swapped in for its gradient still is the one used) and gathered from, instead of building (directions, beams, beams, 3) hkl differences to look up in the 3D tables in every Bethe chunk and bucket. The same master bit for bit; the Bethe classification at npx 512 8.5 -> 5.3 s.

Not changed: matrix_exp stays. The Bloch matrices' |phase A|_1 is 14-18, so a Taylor series on the wave diverges at any useful length, and even at 6 terms it was no faster. Larger Bethe chunks were far slower (626 s against 7.9 s at 4096 a chunk).

The complex64 working precision (the backend's cdtype) is chosen by the caller, not changed here: Autopilot now passes BlochPrismBackend(cdtype=torch.complex64), whose master matches complex128's to a correlation of 1.00000000 (largest relative difference 5.6e-5) for 10.0 s against 15.5 at npx 150.

Wall time on a Titan X with the GPU otherwise idle, each a fresh process (imports, Monte Carlo and fold included), Cu at 15 kV with 9 energy bins:

npx main, complex128 main, complex64 this PR, complex64 this PR, complex128
150 14.9 s 9.7 s 8.6 s 13.4 s
300 46.2 s 16.6 s 12.3 s 35.5 s
512 122.5 s 36.2 s 23.1 s 92.9 s

The whole suite passes (396) apart from test_distribution_identity, which fails on a stale local install's metadata, unrelated.

… time

The prism backend accumulated L_gh = sum_n lambda_n conj(psi_g) psi_h as one
batched outer product per depth, then summed L_gh * S_gh. The intensity is
sum_n lambda_n psi_n^H S psi_n: the waves are stacked and multiplied by S once.
On a Cu master at 15 kV, complex64, npx 512: the depth integration besides
matrix_exp went from ~7 to 2.7 s, the master the same to 1.3e-7.
reduce_to_fundamental_sector folded with orix's in_fundamental_sector on the
CPU (2.9 s for an npx 512 grid, more on a first call). Each direction now takes
its image under the point group with the largest (z, y, x), rounded as the
de-duplication rounds, on the directions' device: 0.5 s at npx 512, the same
number of representatives as orix for Fm-3m, P6_3/mmc and Pnma and fewer for
P-1, each one an image of its directions.
…trices

The prism backend built (directions, beams, beams, 3) hkl differences in every
Bethe-classification chunk and every bucket, to look U and S up in the 3D
tables. Every pair of beams' U(h_g - h_h), |U| and S(h_h - h_g) is now a small
matrix made once per intensities call (from the tables as they are then, so a
u_lut swapped in for its gradient still is the one used) and gathered from. The
same master bit for bit; on a Cu master at npx 512 the classification went
from 8.5 to 5.3 s.
@CSSFrancis
CSSFrancis merged commit 874f6d2 into main Oct 27, 2026
4 checks passed
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.

1 participant