Consider further vectorising XRD get_pattern method across hkl

I work on the materialscloud platform and am fairly interested in computing XRD patterns for a wide variety of materials (around 100k). The xrd get_pattern module does some np.array calculations however the hkl loop often by far dominates the total calculation time (see table). ``` hkl loop 0.017381 s build output 0.002513 s reci lattice search 0.001067 s lattice filtering 0.000125 s build site arrays 0.000123 s symmetry ref 0.000003 s ``` I think it should be possible to vectorise across the hkl planes and provide a significant speedup. My limited testing seems that following snippet should be 1:1 in function but significantly improve calculation time (3-4x speedup). This ofcourse comes at some cost of readability but worth discussion ``` hkl loop 0.000598 s build output 0.002347 s reci lattice search 0.001087 s lattice filtering 0.000153 s build site arrays 0.000107 s symmetry ref 0.000005 s peak grouping 0.000672 s # Peaks are grouped using a bin rather than in loop. ``` ``` py def get_pattern(self, structure: Structure, scaled=True, two_theta_range=(0, 90)): if self.symprec: finder = SpacegroupAnalyzer(structure, symprec=self.symprec) structure = finder.get_refined_structure() wavelength = self.wavelength latt = structure.lattice is_hex = latt.is_hexagonal() min_r, max_r = ( (0, 2 / wavelength) if two_theta_range is None else [2 * sin(radians(t / 2)) / wavelength for t in two_theta_range] ) recip_latt = latt.reciprocal_lattice_crystallographic recip_pts = recip_latt.get_points_in_sphere([[0, 0, 0]], [0, 0, 0], max_r) if min_r: recip_pts = [pt for pt in recip_pts if pt[1] >= min_r] # --- Build per-site arrays --- _zs, _coeffs, _fcoords, _occus, _dwfactors = [], [], [], [], [] for site in structure: for sp, occu in site.species.items(): _zs.append(sp.Z) try: c = ATOMIC_SCATTERING_PARAMS[sp.symbol] except KeyError: raise ValueError( f"Unable to calculate XRD pattern as there is no scattering " f"coefficients for {sp.symbol}." ) _coeffs.append(c) _dwfactors.append(self.debye_waller_factors.get(sp.symbol, 0)) _fcoords.append(site.frac_coords) _occus.append(occu) zs = np.array(_zs) # (N,) coeffs = np.array(_coeffs) # (N, 4, 2) fcoords = np.array(_fcoords) # (N, 3) occus = np.array(_occus) # (N,) dwfactors = np.array(_dwfactors) # (N,) # --- Unpack reciprocal points & filter g_hkl == 0 --- recip_pts_sorted = sorted(recip_pts, key=lambda i: (i[1], -i[0][0], -i[0][1], -i[0][2])) hkls_raw = np.array([pt[0] for pt in recip_pts_sorted]) # (M, 3) g_hkls = np.array([pt[1] for pt in recip_pts_sorted]) # (M,) nonzero = g_hkls != 0 hkls_raw = hkls_raw[nonzero] g_hkls = g_hkls[nonzero] hkls_int = np.round(hkls_raw).astype(int) # (M, 3) # --- Fully vectorized computation over all M hkl points --- # shapes: (M,) theta = np.arcsin(np.clip(wavelength * g_hkls / 2, -1, 1)) s2 = (g_hkls / 2) ** 2 # (M,) # Atomic scattering factors: (M, N) # fs[m, n] = zs[n] - 41.78214 * s2[m] * sum_k(coeffs[n,k,0] * exp(-coeffs[n,k,1]*s2[m])) # coeffs: (N, 4, 2) → broadcast s2: (M, 1, 1) s2_mnk = s2[:, None, None] # (M, 1, 1) gauss = np.sum( coeffs[None, :, :, 0] * np.exp(-coeffs[None, :, :, 1] * s2_mnk), axis=2, ) # (M, N) fs = zs[None, :] - 41.78214 * s2[:, None] * gauss # (M, N) # Debye-Waller per atom, per hkl: (M, N) dw = np.exp(-dwfactors[None, :] * s2[:, None]) # g·r for all hkl and all atoms: (M, N) g_dot_r = hkls_int.astype(float) @ fcoords.T # (M, N) # Structure factors: (M,) f_hkl = np.sum( fs * occus[None, :] * np.exp(2j * pi * g_dot_r) * dw, axis=1, ) i_hkl = (f_hkl * f_hkl.conjugate()).real # (M,) # Lorentz-polarization factor: (M,) cos2t = np.cos(2 * theta) sint = np.sin(theta) cost = np.cos(theta) lorentz = (1 + cos2t ** 2) / (sint ** 2 * cost) intensities = i_hkl * lorentz # (M,) two_thetas_arr = np.degrees(2 * theta) # (M,) # --- Merge peaks within TWO_THETA_TOL using rounding-based binning --- tol = AbstractDiffractionPatternCalculator.TWO_THETA_TOL bin_keys = np.round(two_thetas_arr / tol).astype(int) peaks: dict[int, list] = {} for m in range(len(g_hkls)): hkl = tuple(hkls_int[m]) if is_hex: hkl = (hkl[0], hkl[1], -hkl[0] - hkl[1], hkl[2]) key = bin_keys[m] d_hkl = 1.0 / g_hkls[m] if key in peaks: peaks[key][0] += float(intensities[m]) # np.float64 -> float peaks[key][1].append(hkl) else: peaks[key] = [intensities[m], [hkl], two_thetas_arr[m], d_hkl] # --- Build output --- max_intensity = max(v[0] for v in peaks.values()) tol_scaled = AbstractDiffractionPatternCalculator.SCALED_INTENSITY_TOL x, y, hkls_out, d_hkls_out = [], [], [], [] for key in sorted(peaks): v = peaks[key] fam = get_unique_families(v[1]) if v[0] / max_intensity * 100 > tol_scaled: x.append(float(v[2])) y.append(float(v[0])) hkls_out.append([{"hkl": hkl, "multiplicity": mult} for hkl, mult in fam.items()]) d_hkls_out.append(float(v[3])) xrd = DiffractionPattern(x, y, hkls_out, d_hkls_out) if scaled: xrd.normalize(mode="max", value=100) return xrd ```