pyextal.blochNumpy¶
NumPy Bloch wave backend.
Builds the Bloch-wave structure matrix for every tilt on the CPU and solves the eigenproblems with numpy.linalg.eig on a BLAS-pinned thread pool. Parsing, structure factors and beam selection come from the Fortran extension (pyextal.blochwave), so tilts, beam ordering and ndiag match the Fortran backend exactly.
No matrix padding: each tilt is solved with its exact (ndiag × ndiag) matrix.
No MPI or GPU dependency: runs on any machine with NumPy.
Scales over CPU cores via the ncores parameter.
- Intensity formula (standard Bloch wave):
psi[g, t] = sum_j C[ib[g], j] * (C^{-1} @ phi_0)[j] * exp(-i*pi*gamma[j]*t / bigkn) I[g, t] = |psi[g, t]|^2
Public API (drop-in replacements)¶
bloch_numpy(fname, …) → SimParams (replaces bloch()) bloch_run_numpy(param, …) → SimParams (replaces bloch_run()) tilt_run_numpy(param, tilts, …) → SimParams (replaces tilt_run()) larbed_numpy(param, thickness) → np.ndarray (replaces LARBED_tilt())
Functions¶
|
Return the number of physical CPU cores, falling back to logical count. |
|
Eigendecompose a single (ndiag, ndiag) complex matrix with NumPy. |
Pin BLAS/LAPACK to one thread for the duration of the context. |
|
|
list(map(fn, items)) on a thread pool with BLAS pinned to one thread. |
|
Copy all simulation parameters from Fortran modules into a plain dict. |
|
Generate the k-space tilt grid. |
|
Build the Bloch-wave structure matrix A for every tilt (CPU). |
|
Eigendecompose each matrix in A_list using a thread pool. |
|
Assemble Data objects and attach them to a SimParams instance. |
|
NumPy backend: generate tilt grid and run eigensolve. |
|
Full Bloch wave simulation using the NumPy backend. |
|
NumPy eigensolve for a pre-specified set of tilts. |
|
Compute LARBED intensities from the eigensolve results in param.tilts. |
Module Contents¶
- pyextal.blochNumpy._physical_core_count() int¶
Return the number of physical CPU cores, falling back to logical count.
- pyextal.blochNumpy._eig_worker(A: numpy.ndarray) Tuple[numpy.ndarray, numpy.ndarray]¶
Eigendecompose a single (ndiag, ndiag) complex matrix with NumPy.
- pyextal.blochNumpy._single_threaded_blas()¶
Pin BLAS/LAPACK to one thread for the duration of the context.
With the BLAS library’s default threading, every LAPACK/matmul call inside the worker pool spawns its own team of BLAS threads; the resulting oversubscription erases the pool speedup entirely (measured with OpenBLAS: 8-thread pool ≈ serial, 12-thread pool slower than serial). One BLAS thread per worker restores near-linear scaling.
threadpoolctl (a required dependency) handles OpenBLAS, MKL and BLIS; the mkl-service fallback only matters if threadpoolctl is missing.
- pyextal.blochNumpy._thread_map(fn, items, ncores: int = 0) list¶
list(map(fn, items)) on a thread pool with BLAS pinned to one thread.
NumPy releases the GIL inside LAPACK, matmul and large ufunc loops, so plain threads give real parallelism without pickling. ncores: 0 → physical core count; 1 → serial in the calling thread.
- pyextal.blochNumpy._extract_params() dict¶
Copy all simulation parameters from Fortran modules into a plain dict.
Must be called after bloch_parse() has populated the Fortran modules. Diagonal elements of ughr/ughi are zeroed here; they are set per-tilt during A-matrix construction.
- pyextal.blochNumpy._generate_tilt_grid(param: pyextal.callBloch.SimParams, pixsiz: float, xaxis: numpy.ndarray, subaper: int, subnx: numpy.ndarray, subny: numpy.ndarray) Tuple[numpy.ndarray, numpy.ndarray]¶
Generate the k-space tilt grid.
- Returns:
net ((ntilt, 2) integer pixel offsets (ix, iy))
tilts ((ntilt, 3) float32 tilt vectors)
- pyextal.blochNumpy._build_all_A_matrices(tilts: numpy.ndarray, p: dict, alpha: float, ncores: int = 0) dict¶
Build the Bloch-wave structure matrix A for every tilt (CPU).
- Implements the exact algorithm from parabloch.f90 including:
excitation errors (2KS_g)
gn normalization factors
Bethe perturbation correction
beam selection via Fortran beamcheck (guarantees identical ig ordering)
Two-pass design for performance¶
- Pass 1 (batched): sg/gn/bigkn vectorized over all tilts; beam selection via
the reentrant Fortran beamcheck_batch, chunked over a thread pool (the routine releases the GIL — no OpenMP/MPI runtime involved). Falls back to the serial per-tilt beamcheck for older compiled extensions.
- Pass 2 (batched): group tilts by ndiag and split into chunks; within a
chunk each tilt’s Ugh sub-matrices are gathered with batched fancy indexing and the Bethe matmuls run as one (n_chunk, ndiag, npert)@(n_chunk, npert, ndiag) call. Chunks run on a BLAS-pinned thread pool. All assembly is float32/complex64 — same precision as the Fortran backend, and the eigensolver input dtype.
- returns:
A_list : list of (ndiag, ndiag) complex64 arrays ndiag_list : list of int ig_list : list of (nbeams,) int32 arrays (0-based) ib_list : list of (nout,) int32 arrays (position of output beams in ig) bigkn_list : list of float gn_list : list of (ndiag,) float32 arrays (for eigenvector post-scaling)
- rtype:
dict with keys
- pyextal.blochNumpy._numpy_eigensolve(A_list: List[numpy.ndarray], ncores: int = 0) Tuple[List[numpy.ndarray], List[numpy.ndarray]]¶
Eigendecompose each matrix in A_list using a thread pool.
Matrices are NOT padded — each tilt uses its exact (ndiag_i × ndiag_i) matrix.
np.linalg.eig releases the GIL during LAPACK, so threads give true parallelism without the pickle/IPC overhead of multiprocessing.Pool. BLAS-internal threading is pinned to 1 inside the pool (see _single_threaded_blas) — one LAPACK thread per worker.
- Parameters:
A_list (list of (ndiag_i, ndiag_i) complex64 arrays, one per tilt)
ncores (number of threads (0 → auto: physical core count; 1 → serial))
- Returns:
eigenvalues_list (list of (ndiag_i,) complex arrays)
eigenvectors_list (list of (ndiag_i, ndiag_i) complex arrays)
- pyextal.blochNumpy._assemble_simparams(param: pyextal.callBloch.SimParams, net: numpy.ndarray, tilts: numpy.ndarray, eigenvalues_list, eigenvectors_list, precomputed: dict) pyextal.callBloch.SimParams¶
Assemble Data objects and attach them to a SimParams instance.
Applies gn post-scaling to eigenvectors (matching parabloch.f90 line vector(i,j)=vector(i,j)*gn(k)) and recomputes cinv from the scaled C. Everything stays complex64 — the same precision as the Fortran backend and the eigensolver input.
- Parameters:
param (empty SimParams (ntilt=0, tilts=[]) from bloch_parse)
net ((ntilt, 2) integer pixel grid)
tilts ((ntilt, 3) float tilt vectors)
eigenvalues_list (per-tilt (ndiag_i,) complex eigenvalues)
eigenvectors_list (per-tilt (ndiag_i, ndiag_i) complex eigenvectors)
precomputed (dict from _build_all_A_matrices)
- pyextal.blochNumpy.bloch_run_numpy(param: pyextal.callBloch.SimParams, subaper: int = 0, subnx: numpy.ndarray | None = None, subny: numpy.ndarray | None = None, pixsiz: float = 0.0, xaxis: numpy.ndarray | None = None, ncores: int = 0) pyextal.callBloch.SimParams¶
NumPy backend: generate tilt grid and run eigensolve.
Drop-in replacement for callBloch.bloch_run(backend=’numpy’). Requires that bloch_parse() has been called beforehand.
- Parameters:
param (SimParams from bloch_parse())
subaper (aperture type (0=square, 1=circular, 2=custom-r, 3=rect))
subnx (aperture x bounds)
subny (aperture y bounds)
pixsiz (pixel size; 0 → auto from param)
xaxis ((3,) horizontal reciprocal direction; None → param.gg)
ncores (number of parallel worker processes for eigensolve)
- Return type:
param with .ntilt and .tilts populated
- pyextal.blochNumpy.bloch_numpy(fname: str, t_x: float = 0.0, t_y: float = 0.0, reflections=None, values=None, HKL: bool = False, subaper: int = 0, subnx: numpy.ndarray | None = None, subny: numpy.ndarray | None = None, pixsiz: float = 0.0, xaxis: numpy.ndarray | None = None, ncores: int = 0) pyextal.callBloch.SimParams¶
Full Bloch wave simulation using the NumPy backend.
Drop-in replacement for callBloch.bloch(backend=’numpy’).
- pyextal.blochNumpy.tilt_run_numpy(param: pyextal.callBloch.SimParams, tilts: numpy.ndarray, indices: numpy.ndarray | None = None, ncores: int = 0) pyextal.callBloch.SimParams¶
NumPy eigensolve for a pre-specified set of tilts.
Drop-in replacement for callBloch.tilt_run(backend=’numpy’). Reads the current Ugh matrix from the Fortran module so it reflects any structure-factor updates applied before this call.
- Parameters:
param (SimParams initialised by bloch_parse())
tilts ((ntilt, 3) float32 tilt vectors)
indices ((nsel,) int array — beam row indices to put first in eigenvector,) – matching Fortran tilt_run ordering for LARBED_tilt consistency.
ncores (number of parallel worker processes)
- Return type:
param with .ntilt set and .tilts populated
- pyextal.blochNumpy.larbed_numpy(param: pyextal.callBloch.SimParams, thickness, nout: int | None = None) numpy.ndarray¶
Compute LARBED intensities from the eigensolve results in param.tilts.
Pure-Python replacement for callBloch.LARBED_tilt() (backend=’numpy’). Reads param.tilts (populated by tilt_run_numpy / bloch_run_numpy).
- Parameters:
param (SimParams with .tilts populated by tilt_run_numpy / bloch_run_numpy)
thickness (scalar or (ntilt,) array — sample thickness in Å)
nout (number of output beams to return (first nout rows of eigenvector).) – None → all param.nout beams. Mirrors the nout arg of LARBED_tilt.
- Returns:
intensities
- Return type: