Merge pull request #128 from lcmd-epfl/bugfix

- Fix existing and potential errors when Mole.atom_charges() are used to determine atom number
- Fix xyz_to_mol(ignore=True) bug
- Fix number of electrons correction
- Refactor: remove singleatom_basis_enumerator()
- Remove legacy code from sym.py and refactor

Author: briling

qstack/compound.py

@@ -7,7 +7,7 @@
 from pyscf import gto, data
 from qstack import constants
 from qstack.reorder import get_mrange
-from qstack.mathutils.array import stack_padding
+from qstack.mathutils.array import stack_padding, loadtxtvar
 from qstack.mathutils.rotation_matrix import rotate_euler
 from qstack.tools import Cursor
 
@@ -144,8 +144,9 @@ def xyz_to_mol(inp, basis="def2-svp", charge=None, spin=None, ignore=False, unit
     if ignore:
         if charge not in (0, None) or spin not in (0, None):
             warnings.warn("Spin and charge values are overwritten", RuntimeWarning, stacklevel=2)
+        atoms = [int(q) if q.isdigit() else data.elements.ELEMENTS_PROTON[q] for q in loadtxtvar(molxyz, dtype='str', usecols=0)]
         mol.spin = 0
-        mol.charge = - sum(mol.atom_charges())%2
+        mol.charge = -(sum(atoms)%2)
     else:
         if charge is not None:
             mol.charge = charge
@@ -162,8 +163,7 @@ def xyz_to_mol(inp, basis="def2-svp", charge=None, spin=None, ignore=False, unit
             mol.spin = 0
 
     mol.build()
-    species_charges = [data.elements.charge(z) for z in mol.elements]
-    if mol.basis == 'minao' and ecp is None and (np.array(species_charges) > 36).any():
+    if mol.basis == 'minao' and ecp is None and (numbers(mol) > 36).any():
         msg = f"{mol.basis} basis set requires the use of effective core potentials for atoms with Z>36"
         raise RuntimeError(msg)
     return mol
@@ -303,64 +303,31 @@ def make_atom(q, basis, ecp=None):
     return mol
 
 
-def singleatom_basis_enumerator(basis):
-    """Enumerate the different tensors of atomic orbitals within a 1-atom basis set.
-
-    Each tensor is a 2l+1-sized group of orbitals that share a radial function and l value.
-
-    Args:
-        basis (list): Basis set definition in pyscf format.
-
-    Returns:
-        tuple: A tuple containing:
-        - l_per_bas (list): Angular momentum quantum number l for each basis shell.
-        - n_per_bas (list): Radial function counter n (starting at 0) for each basis shell.
-        - ao_starts (list): Starting index in AO array for each basis shell.
-    """
-    ao_starts = []
-    l_per_bas = []
-    n_per_bas = []
-    cursor = Cursor(action='ranger')
-    cursor_per_l = []
-    for bas in basis:
-        # shape of `bas`, l, then another optional constant, then lists [exp, coeff, coeff, coeff]
-        # that make a matrix between the number of functions (number of coeff per list)
-        # and the number of primitive gaussians (one per list)
-        l = bas[0]
-        while len(cursor_per_l) <= l:
-            cursor_per_l.append(Cursor(action='ranger'))
-        n_count = len(bas[-1])-1
-        l_per_bas += [l] * n_count
-        n_per_bas.extend(cursor_per_l[l].add(n_count))
-        msize = 2*l+1
-        ao_starts.extend(cursor.add(msize*n_count)[::msize])
-    return l_per_bas, n_per_bas, ao_starts
-
-
 def basis_flatten(mol, return_both=True, return_shells=False):
     """Flatten a basis set definition for AOs.
 
     Args:
         mol (pyscf.gto.Mole): pyscf Mole object.
         return_both (bool): Whether to return both AO info and primitive Gaussian info. Defaults to True.
-        return_shells (bool): Whether to return angular momenta per shell. Defaults to False.
+        return_shells (bool): Whether to return angular momenta and starting indices per shell
+            (2l+1-sized group of orbitals that share a radial function and l value). Defaults to False.
 
     Returns:
         - numpy.ndarray: 3×mol.nao int array where each column corresponds to an AO and rows are:
-            - 0: atom index
-            - 1: angular momentum quantum number l
-            - 2: magnetic quantum number m
+            - 0: atom index,
+            - 1: angular momentum quantum number l,
+            - 2: magnetic quantum number m.
         If return_both is True, also returns:
         - numpy.ndarray: 2×mol.nao×max_n float array where index (i,j,k) means:
-            - i: 0 for exponent, 1 for contraction coefficient of a primitive Gaussian
-            - j: AO index
-            - k: radial function index (padded with zeros if necessary)
-        If return_shell is True, also returns:
-        - numpy.ndarray: angular momentum quantum number for each shell
-
+            - i: 0 for exponent, 1 for contraction coefficient of a primitive Gaussian,
+            - j: AO index,
+            - k: radial function index (padded with zeros if necessary).
+        If return_shells is True, also returns:
+        - numpy.ndarray: starting AO indices for each shell.
     """
     x = []
-    L = []
+    ao_starts = []
+    cursor = Cursor(action='ranger')
     y = np.zeros((3, mol.nao), dtype=int)
     i = Cursor(action='slicer')
     a = mol.bas_exps()
@@ -376,11 +343,26 @@ def basis_flatten(mol, return_both=True, return_shells=False):
                     x.extend([ac]*msize)
             y[:,i(msize*n)] = np.vstack((np.array([[iat, l]]*msize*n).T, [*get_mrange(l)]*n))
             if return_shells:
-                L.extend([l]*n)
+                ao_starts.extend(cursor.add(msize*n)[::msize])
 
     ret = [y]
     if return_both:
         ret.append(stack_padding(x).transpose((1,0,2)))
     if return_shells:
-        ret.append(np.array(L))
+        ret.append(np.array(ao_starts))
     return ret[0] if len(ret)==1 else ret
+
+
+def numbers(mol):
+    """Get atom numbers of a molecule.
+
+    Use this function to get atomic NUMBERS to index elements.
+    Use `mol.atom_charges()` to get CHARGES (it returns effective charges when ECP are used).
+
+    Args:
+        mol (pyscf.gto.Mole): pyscf Mole object.
+
+    Returns:
+        numpy.ndarray: Array of atomic numbers.
+    """
+    return np.array([data.elements.charge(q) for q in mol.elements])

qstack/equio.py

@@ -4,11 +4,10 @@
 from functools import reduce
 from types import SimpleNamespace
 import numpy as np
-from pyscf import data
 import metatensor
 from qstack.tools import Cursor
 from qstack.reorder import get_mrange, pyscf2gpr_l1_order
-from qstack.compound import singleatom_basis_enumerator
+from qstack.compound import basis_flatten, numbers
 
 
 vector_label_names = SimpleNamespace(
@@ -35,9 +34,14 @@ def _get_llist(mol):
         mol (pyscf.gto.Mole): pyscf Mole object.
 
     Returns:
-        dict: Dictionary with atom numbers as keys and List of angular momentum quantum numbers for each basis function as values.
+        dict: Dictionary with atom numbers as keys and List of angular momentum quantum numbers for each shell as values.
     """
-    return {int(q): singleatom_basis_enumerator(mol._basis[data.elements.ELEMENTS[q]])[0] for q in np.unique(mol.atom_charges())}
+    ao, ao_start = basis_flatten(mol, return_both=False, return_shells=True)
+    iat_shells, l_shells, _ = ao[:,ao_start]
+    llist = {}
+    for q, iat in zip(*np.unique(numbers(mol), return_index=True), strict=True):
+        llist[q] = l_shells[np.where(iat_shells==iat)].tolist()
+    return llist
 
 
 def _get_tsize(tensor):
@@ -78,7 +82,7 @@ def vector_to_tensormap(mol, c):
     Returns:
         metatensor.TensorMap: Tensor map representation of the vector.
     """
-    atom_charges = mol.atom_charges()
+    atom_charges = numbers(mol)
 
     tm_label_vals = []
     block_prop_label_vals = {}
@@ -161,7 +165,7 @@ def tensormap_to_vector(mol, tensor):
         raise RuntimeError(f'Tensor size mismatch ({nao} instead of {mol.nao})')
 
     c = np.zeros(mol.nao)
-    atom_charges = mol.atom_charges()
+    atom_charges = numbers(mol)
     llists = _get_llist(mol)
     i = 0
     for iat, q in enumerate(atom_charges):
@@ -191,7 +195,7 @@ def matrix_to_tensormap(mol, dm):
     Returns:
         metatensor.TensorMap: Tensor map representation of the matrix.
     """
-    atom_charges = mol.atom_charges()
+    atom_charges = numbers(mol)
     elements, counts = np.unique(atom_charges, return_counts=True)
     counts = dict(zip(elements, counts, strict=True))
     element_indices = {q: np.where(atom_charges==q)[0] for q in elements}
@@ -313,7 +317,7 @@ def tensormap_to_matrix(mol, tensor):
         raise RuntimeError(f'Tensor size mismatch ({nao2} instead of {mol.nao*mol.nao})')
 
     dm = np.zeros((mol.nao, mol.nao))
-    atom_charges = mol.atom_charges()
+    atom_charges = numbers(mol)
     llists = _get_llist(mol)
     i1 = 0
     for iat1, q1 in enumerate(atom_charges):

qstack/fields/decomposition.py

@@ -203,7 +203,7 @@ def correct_N(mol, c0, N=None, mode='Lagrange', metric='u'):
         numpy ndarray: Corrected decomposition coefficients (1D array).
     """
     mode = mode.lower()
-    q = moments.r2_c(mol, None, moments=[0])
+    q = moments.r2_c(mol, None, moments=[0])[0]
     N0 = c0 @ q
 
     if N is None:

qstack/mathutils/array.py

@@ -1,5 +1,6 @@
 """Array manipulation utility functions."""
 
+from io import StringIO
 import numpy as np
 from qstack.tools import slice_generator
 
@@ -94,3 +95,19 @@ def vstack_padding(xs):
         slices = (s0, *(np.s_[0:s] for s in x.shape[1:]))
         X[slices] = x
     return X
+
+
+def loadtxtvar(var, *kargs, **kwargs):
+    """Load a numpy array from a string variable.
+
+    Args:
+        var (str): String variable containing the data.
+        *kargs: Additional positional arguments for numpy.loadtxt.
+        **kwargs: Additional keyword arguments for numpy.loadtxt.
+
+    Returns:
+        numpy.ndarray: Loaded array.
+    """
+    with StringIO(var) as f:
+        x = np.loadtxt(f, *kargs, **kwargs)
+    return x

qstack/reorder.py

@@ -98,9 +98,9 @@ def reorder_ao(mol, vector, src='pyscf', dest='gpr'):
         NotImplementedError: If the specified convention is not implemented.
         ValueError: If vector dimension is not 1 or 2.
     """
-    def get_idx(L, m, convention):
+    def get_idx(l_shells, m, convention):
         convention = convention.lower()
-        l_slices = slice_generator(L, inc=lambda l: 2*l+1)
+        l_slices = slice_generator(l_shells, inc=lambda l: 2*l+1)
         if convention == 'gpr':
             return np.arange(len(m)), np.ones_like(m)
         elif convention == 'pyscf':
@@ -113,9 +113,11 @@ def get_idx(L, m, convention):
 
     from .compound import basis_flatten
 
-    (_, _, m), L = basis_flatten(mol, return_both=False, return_shells=True)
-    idx_src, sign_src  = get_idx(L, m, src)
-    idx_dest, sign_dest = get_idx(L, m, dest)
+    (_, l, m), shell_start = basis_flatten(mol, return_both=False, return_shells=True)
+    l_shells = l[shell_start]
+
+    idx_src, sign_src  = get_idx(l_shells, m, src)
+    idx_dest, sign_dest = get_idx(l_shells, m, dest)
 
     if vector.ndim == 2:
         sign_src  = np.einsum('i,j->ij', sign_src, sign_src)