qstack.spahm.rho.dmb_rep_bond

Functions for SPAHM(b) computation.

qstack.spahm.rho.dmb_rep_bond.bonds_dict_init(qqs, M)

Initialize storage for bond representations.

Creates a dictionary with zero-initialized arrays for each bond type, with array sizes matching the corresponding metric matrix dimensions.

Parameters:

• qqs (list) – List of bond pair names (e.g., [‘CC’, ‘CH’, ‘OH’]).

• M (dict) – Dictionary mapping bond names to metric matrices.

Returns:

(mybonds, N) where:

• mybonds (dict): Bond name to zero-initialized numpy array mapping.

• N (int): Total number of basis functions across all bond types.

Return type:

tuple

qstack.spahm.rho.dmb_rep_bond.fit_dm(dm, mol, mybasis, ri0, ri1)

Fit density matrix using auxiliary basis functions at bond center.

Decomposes the bond density matrix into auxiliary basis coefficients centered at the bond midpoint, then splits coefficients by angular momentum.

Parameters:

• dm (numpy.ndarray) – Density matrix for the bond.

• mol (pyscf.gto.Mole) – Molecule object containing the bond.

• mybasis (dict) – Basis set dictionary for the bond type.

• ri0 (numpy.ndarray) – Coordinates of first atom in Å.

• ri1 (numpy.ndarray) – Coordinates of second atom in Å.

Returns:

Coefficients split by angular momentum quantum number [(l, coeff), …].

Return type:

list

qstack.spahm.rho.dmb_rep_bond.get_basis_info(qqs, mybasis, only_m0, printlevel)

Compute basis indices and metric matrices for bond pairs.

Parameters:

• qqs (list) – List of bond pair names (e.g., [‘CC’, ‘CH’, ‘OH’]).

• mybasis (dict) – Dictionary mapping bond names to basis set dictionaries.

• only_m0 (bool) – If True, use only m=0 angular momentum components.

• printlevel (int) – Verbosity level.

Returns:

(idx, M) where: - idx (dict): Pair indices for each bond type (list of [i, j] pairs) - M (dict): Metric matrices for each bond type (numpy 2D ndarray)

Return type:

tuple

qstack.spahm.rho.dmb_rep_bond.get_element_pairs(elements)

Generate all possible element pair combinations.

Creates complete list of bond types assuming all elements can bond with each other.

Parameters:

elements (list) – List of element symbols.

Returns:

(qqs, qqs4q) where: - qqs (list): Sorted list of unique bond pair names - qqs4q (dict): Maps each element to its list of possible bond partners

Return type:

tuple

qstack.spahm.rho.dmb_rep_bond.get_element_pairs_cutoff(elements, mols, cutoff, align=False)

Determine element pairs based on actual distances in molecules.

Identifies which element pairs actually form bonds within the distance cutoff by scanning molecular geometries.

Parameters:

• elements (list) – List of element symbols to consider.

• mols (list) – List of pyscf Mole objects.

• cutoff (float) – Maximum bond distance in Å.

• align (bool) – If True, includes all element pairs regardless of distance. Defaults to False.

Returns:

(qqs, qqs4q) where: - qqs (list): Sorted list of bond pair names found within cutoff - qqs4q (dict): Maps each element to its list of bond partners

Return type:

tuple

qstack.spahm.rho.dmb_rep_bond.make_bname(q0, q1)

Create canonical bond name from two element symbols.

Parameters:

• q0 (str) – First element symbol.

• q1 (str) – Second element symbol.

Returns:

Concatenated element symbols in alphabetical order (e.g., ‘CH’, ‘CC’, ‘NO’).

Return type:

str

qstack.spahm.rho.dmb_rep_bond.read_basis_wrapper(mols, bpath, only_m0, printlevel, cutoff=None, elements=None, pairfile=None, dump_and_exit=False, same_basis=False)

Read basis sets for all element pairs present in molecules.

Determines which element pairs exist (either all possible or within cutoff distance), loads corresponding basis sets, and computes metric matrices and indices. Can cache pair information to file for subsequent runs.

Parameters:

• mols (list) – List of pyscf Mole objects to analyze.

• bpath (str) – Directory path containing basis set files.

• only_m0 (bool) – If True, use only m=0 angular momentum components.

• printlevel (int) – Verbosity level for output (>1 for detailed printing).

• cutoff (float, optional) – Maximum bond distance in Å for pair detection. If None, considers all element combinations. Defaults to None.

• elements (list, optional) – List of element symbols to consider. If None, extracts all elements from molecules. Defaults to None.

• pairfile (str, optional) – Path to save/load element pair information. Defaults to None.

• dump_and_exit (bool) – If True, saves pair information and exits. Defaults to False.

• same_basis (bool) – If True, uses generic CC.bas for all pairs. Defaults to False.

Returns:

(elements, mybasis, qqs, qqs4q, idx, M) where:

• elements (list): Sorted list of element symbols.

• mybasis (dict): Bond name to basis set dictionary mapping.

• qqs (dict): Maps each element to list of all bond pair names.

• qqs4q (dict): Maps each element to its specific bond partners.

• idx (dict): Pair indices (list of [i, j] pairs) for each bond type.

• M (dict): Metric matrices (2D numpy ndarray) for each bond type.

Return type:

tuple

qstack.spahm.rho.dmb_rep_bond.read_basis_wrapper_pairs(mols, bondidx, bpath, only_m0, printlevel, same_basis=False)

Read basis sets and computes metric matrices for specified bond pairs.

Processes bond pairs from molecular structures and loads their corresponding basis sets from disk, then computes basis indices and metric matrices.

Parameters:

• mols (list) – List of pyscf Mole objects.

• bondidx (list) – List of bond index pairs [(i0, i1), …] for each molecule.

• bpath (str) – Directory path containing basis set files.

• only_m0 (bool) – If True, use only m=0 angular momentum components.

• printlevel (int) – Verbosity level for output (>1 for detailed printing).

• same_basis (bool) – If True, uses generic CC.bas for all pairs. Defaults to False.

Returns:

(mybasis, idx, M) where:

• mybasis (dict): Bond name to basis set dictionary mapping,

• idx (dict): Pair indices (list of [i, j] pairs) for each bond type,

• M (dict): Metric matrices (2D numpy ndarray) for each bond type.

Return type:

tuple

qstack.spahm.rho.dmb_rep_bond.read_df_basis(bnames, bpath, same_basis=False)

Load bond-optimized basis sets from .bas files.

Parameters:

• bnames (list) – List of bond pair names (e.g., [‘CC’, ‘CH’]).

• bpath (str) – Directory path containing .bas files.

• same_basis (bool) – If True, uses generic CC.bas for all pairs. Defaults to False.

Returns:

Dictionary mapping bond names to basis set dictionaries.

Return type:

dict

qstack.spahm.rho.dmb_rep_bond.repr_for_bond(i0, i1, L, mybasis, idx, q, r, cutoff)

Compute bond representation for a specific atom pair.

Extracts bond density, fits it with basis functions at the bond center, and symmetrizes the representation from both atom perspectives.

Parameters:

• i0 (int) – Index of first atom.

• i1 (int) – Index of second atom.

• L (Lowdin_split) – Löwdin-split density matrix object.

• mybasis (dict) – Bond basis sets keyed by bond names.

• idx (dict) – Pair indices for symmetrization.

• q (list) – Element symbols for all atoms.

• r (numpy ndarray) – Atomic coordinates in Å.

• cutoff (float) – Maximum bond distance.

Returns:

([v0, v1], bname) where: - v0: Representation from atom i0’s perspective. - v1: Representation from atom i1’s perspective. - bname: Bond name (e.g., ‘CH’). Returns (None, None) if distance exceeds cutoff.

Return type:

tuple

qstack.spahm.rho.dmb_rep_bond.repr_for_mol(mol, dm, qqs, M, mybasis, idx, maxlen, cutoff, only_z=None)

Compute SPAHM(b) representations for all atoms in a molecule.

Constructs bond-based atomic representations by summing contributions from all bonds of the same type within cutoff distance.

Parameters:

• mol (pyscf.gto.Mole) – Molecule object.

• dm (numpy.ndarray) – Molecular density matrix.

• qqs (dict) – Maps each element to list of bond pair names it can form.

• M (dict) – Metric matrices for each bond type.

• mybasis (dict) – Bond basis sets keyed by bond names.

• idx (dict) – Pair indices for symmetrization of each bond type.

• maxlen (int) – Maximum representation length for zero-padding.

• cutoff (float) – Maximum bond distance in Å to consider.

• only_z (list, optional) – If provided, compute representations only for atoms with these element symbols. Defaults to None (all atoms).

Returns:

Array of shape (n_atoms, maxlen) containing atom representations,

where each row is a zero-padded SPAHM(b) vector.

Return type:

numpy.ndarray

qstack.spahm.rho.dmb_rep_bond.vec_from_cs(z, cs, lmax, idx)

Rotate basis coefficients to bond axis and creates vectorized representation.

Applies Wigner D-matrix rotation to align coefficients with the bond vector (same as pretending the bond is along the z-axis), ensuring rotational invariance, then vectorizes using symmetry indices.

Parameters:

• z (numpy.ndarray) – Bond vector (displacement from one atom to another).

• cs (list) – Angular momentum decomposed coefficients [(l, coeff), …].

• lmax (int) – Maximum angular momentum quantum number.

• idx (dict) – Pair indices mapping for symmetrization.

Returns:

Rotationally invariant vectorized representation of the bond.

Return type:

numpy.ndarray