From 946a7a53202dc2b4ebf35f0b09d3f9e063321b8f Mon Sep 17 00:00:00 2001 From: Patrick Avery Date: Tue, 28 Jul 2026 12:21:07 -0500 Subject: [PATCH] Rewrite ApplySymmetryToTensor3 with NumPy/BLAS (11x speedup) Profiling the free-energy Hessian of a 128-atom PdH supercell in python-sscha showed 82% of the run inside this method: the Fortran symph routines apply the symmetries element by element in serial. The three steps are redone with numpy array operations: the permutation symmetry as an average over the 6 index orderings, the translation average computed once per group of equivalent atoms, and each symmetry at gamma applied as one matrix product over the cartesian components. On the PdH case: 181 s -> 16 s, matching the old result to 1e-16. Also checked on two lower-symmetry cells (120 and 216 atoms). Peak memory grows because of numpy temporaries: on the 216-atom cell about 11 GB against 6.5 GB, with the time going 211 s -> 34 s. The matrix products use numpy's threaded BLAS but do not depend on it: on the 120-atom cell the Fortran path takes 20.4 s and the numpy one 6.5 s on a single thread (3.8 s with threads on a 24-core machine). On shared nodes without core pinning, cap the threads with OPENBLAS_NUM_THREADS. The Fortran routines (permute_v3, trans_v3, sym_v3) no longer have callers here and can be removed in a follow-up if preferred. Signed-off-by: Patrick Avery --- cellconstructor/symmetries.py | 60 ++++++++++++++++++++++++++++++----- 1 file changed, 52 insertions(+), 8 deletions(-) diff --git a/cellconstructor/symmetries.py b/cellconstructor/symmetries.py index 8c9be14a..65a250f0 100644 --- a/cellconstructor/symmetries.py +++ b/cellconstructor/symmetries.py @@ -420,20 +420,64 @@ def ApplySymmetryToTensor3(self, v3, initialize_symmetries = True): the parser to copy the matrix when doing the symmetrization in Fortran. initialize_symmetries : bool If True the symmetries will be initialized using spglib. Otherwise - the already present symmetries will be use. Use it False at your own risk! + the already present symmetries will be use. Use it False at your own risk! (It can crash with seg fault if symmetries are not properly initialized) """ if initialize_symmetries: self.SetupFromSPGLIB() - # Apply the permutation symmetry - symph.permute_v3(v3) - - # Apply the translational symmetries - symph.trans_v3(v3, self.QE_translations_irt) + nat = self.QE_nat - # Apply all the symmetries at gamma - symph.sym_v3(v3, self.QE_at, self.QE_s, self.QE_irt, self.QE_nsymq) + # Apply the permutation symmetry (average over the 6 orderings of + # the three indexes, same as permute_v3) + b = (v3 + v3.transpose(0, 2, 1) + v3.transpose(1, 0, 2) + + v3.transpose(1, 2, 0) + v3.transpose(2, 0, 1) + + v3.transpose(2, 1, 0)) / 6.0 + + # Reshape so the three atom indexes come first and the 27 cartesian + # components of each atom triple sit together in memory + b = b.reshape(nat, 3, nat, 3, nat, 3).transpose(0, 2, 4, 1, 3, 5) + b = np.ascontiguousarray(b).reshape(nat, nat, nat, 27) + + # Apply the translational symmetries (same as trans_v3). + # After the average over the translations, atoms that are equivalent + # by translation carry the same values (with the other two indexes + # relabeled), so the average is only computed for the first atom of + # each equivalent group and then copied to the others. + perms = np.asarray(self.QE_translations_irt, dtype=int).T - 1 # (nr, nat) + nr = perms.shape[0] + first_equiv = perms.min(axis=0) + tr_to_first = perms.argmin(axis=0) + firsts = np.unique(first_equiv) + acc = np.zeros((len(firsts),) + b.shape[1:], dtype=b.dtype) + for p in perms: + acc += b[np.ix_(p[firsts], p, p)] + acc /= nr + for i in range(nat): + p = perms[tr_to_first[i]] + b[i] = acc[np.searchsorted(firsts, first_equiv[i])][np.ix_(p, p)] + + # Apply all the symmetries at gamma (same operations sym_v3 applied). + # For each symmetry: relabel the atoms with irt, then rotate the 27 + # cartesian components of every atom triple in a single matrix + # product. R is the cartesian rotation matrix of the symmetry, + # built from QE_s and the cell vectors, and np.kron(np.kron(R, R), R) + # rotates all three indexes at once. + at = np.asarray(self.QE_at, dtype=np.double) + bg = np.linalg.inv(at.T) + s_all = np.asarray(self.QE_s) + irt = np.asarray(self.QE_irt, dtype=int) + out = np.zeros_like(b) + for isym in range(self.QE_nsymq): + R = bg @ s_all[:, :, isym] @ at.T + R3 = np.kron(np.kron(R, R), R) + p = irt[isym, :nat] - 1 + out += (b[np.ix_(p, p, p)].reshape(-1, 27) @ R3.T).reshape(b.shape) + out /= self.QE_nsymq + + # Back to the (3*nat, 3*nat, 3*nat) layout, overwriting v3 in place + out = out.reshape(nat, nat, nat, 3, 3, 3).transpose(0, 3, 1, 4, 2, 5) + v3[...] = out.reshape(3 * nat, 3 * nat, 3 * nat) def ApplySymmetryToEffCharge(self, eff_charges): """