pyextal.blochNumpy ================== .. py:module:: pyextal.blochNumpy .. autoapi-nested-parse:: 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 --------- .. autoapisummary:: pyextal.blochNumpy._physical_core_count pyextal.blochNumpy._eig_worker pyextal.blochNumpy._single_threaded_blas pyextal.blochNumpy._thread_map pyextal.blochNumpy._extract_params pyextal.blochNumpy._generate_tilt_grid pyextal.blochNumpy._build_all_A_matrices pyextal.blochNumpy._numpy_eigensolve pyextal.blochNumpy._assemble_simparams pyextal.blochNumpy.bloch_run_numpy pyextal.blochNumpy.bloch_numpy pyextal.blochNumpy.tilt_run_numpy pyextal.blochNumpy.larbed_numpy Module Contents --------------- .. py:function:: _physical_core_count() -> int Return the number of physical CPU cores, falling back to logical count. .. py:function:: _eig_worker(A: numpy.ndarray) -> Tuple[numpy.ndarray, numpy.ndarray] Eigendecompose a single (ndiag, ndiag) complex matrix with NumPy. .. py:function:: _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. .. py:function:: _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. .. py:function:: _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. .. py:function:: _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*) .. py:function:: _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 .. py:function:: _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. :param A_list: :type A_list: list of (ndiag_i, ndiag_i) complex64 arrays, one per tilt :param ncores: :type 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*) .. py:function:: _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. :param param: :type param: empty SimParams (ntilt=0, tilts=[]) from bloch_parse :param net: :type net: (ntilt, 2) integer pixel grid :param tilts: :type tilts: (ntilt, 3) float tilt vectors :param eigenvalues_list: :type eigenvalues_list: per-tilt (ndiag_i,) complex eigenvalues :param eigenvectors_list: :type eigenvectors_list: per-tilt (ndiag_i, ndiag_i) complex eigenvectors :param precomputed: :type precomputed: dict from _build_all_A_matrices .. py:function:: bloch_run_numpy(param: pyextal.callBloch.SimParams, subaper: int = 0, subnx: Optional[numpy.ndarray] = None, subny: Optional[numpy.ndarray] = None, pixsiz: float = 0.0, xaxis: Optional[numpy.ndarray] = 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. :param param: :type param: SimParams from bloch_parse() :param subaper: :type subaper: aperture type (0=square, 1=circular, 2=custom-r, 3=rect) :param subnx: :type subnx: aperture x bounds :param subny: :type subny: aperture y bounds :param pixsiz: :type pixsiz: pixel size; 0 → auto from param :param xaxis: :type xaxis: (3,) horizontal reciprocal direction; None → param.gg :param ncores: :type ncores: number of parallel worker processes for eigensolve :rtype: param with .ntilt and .tilts populated .. py:function:: 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: Optional[numpy.ndarray] = None, subny: Optional[numpy.ndarray] = None, pixsiz: float = 0.0, xaxis: Optional[numpy.ndarray] = None, ncores: int = 0) -> pyextal.callBloch.SimParams Full Bloch wave simulation using the NumPy backend. Drop-in replacement for callBloch.bloch(backend='numpy'). .. py:function:: tilt_run_numpy(param: pyextal.callBloch.SimParams, tilts: numpy.ndarray, indices: Optional[numpy.ndarray] = 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. :param param: :type param: SimParams initialised by bloch_parse() :param tilts: :type tilts: (ntilt, 3) float32 tilt vectors :param indices: matching Fortran tilt_run ordering for LARBED_tilt consistency. :type indices: (nsel,) int array — beam row indices to put first in eigenvector, :param ncores: :type ncores: number of parallel worker processes :rtype: param with .ntilt set and .tilts populated .. py:function:: larbed_numpy(param: pyextal.callBloch.SimParams, thickness, nout: Optional[int] = 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). :param param: :type param: SimParams with .tilts populated by tilt_run_numpy / bloch_run_numpy :param thickness: :type thickness: scalar or (ntilt,) array — sample thickness in Å :param nout: None → all param.nout beams. Mirrors the nout arg of LARBED_tilt. :type nout: number of output beams to return (first nout rows of eigenvector). :returns: **intensities** :rtype: (nout, ntilt) float64 — same shape as LARBED_tilt output