Merge pull request lcmd-epfl#116 from lcmd-epfl/fix-spahm-a

Fix and refactor SPAHM(a)

Author: briling

qstack/basis_opt/opt.py

@@ -43,7 +43,6 @@ def energy(x):
             E += qbbt.energy_mol(newbasis, m)
         return E
 
-
     def gradient(x):
         """Compute total loss function (fitting error) and gradient for given exponents.
 
@@ -76,7 +75,6 @@ def gradient(x):
         dE_dx = dE_da * exponents
         return E, dE_dx
 
-
     def gradient_only(x):
         """Compute only the gradient of the loss function (wrapper for optimization algorithms).
 
@@ -88,7 +86,6 @@ def gradient_only(x):
         """
         return gradient(x)[1]
 
-
     def read_bases(basis_files):
         """Read basis set definitions from files or dicts.
 
@@ -117,7 +114,6 @@ def read_bases(basis_files):
                 basis.update(i)
         return basis
 
-
     def make_bf_start():
         """Create basis function index bounds for each element.
 
@@ -131,7 +127,6 @@ def make_bf_start():
             bf_bounds[q] = [start, start+nbf[i]]
         return bf_bounds
 
-
     def make_moldata(fname):
         """Create molecular data dictionary from file or dict.
 

qstack/compound.py

@@ -9,6 +9,7 @@
 from qstack.reorder import get_mrange
 from qstack.mathutils.array import stack_padding
 from qstack.mathutils.rotation_matrix import rotate_euler
+from qstack.tools import Cursor
 
 
 # detects a charge-spin line, containing only two ints (one positive or negative, the other positive and nonzero)
@@ -319,49 +320,49 @@ def singleatom_basis_enumerator(basis):
     ao_starts = []
     l_per_bas = []
     n_per_bas = []
-    cursor = 0
+    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(0)
-
+            cursor_per_l.append(Cursor(action='ranger'))
         n_count = len(bas[-1])-1
-        n_start = cursor_per_l[l]
-        cursor_per_l[l] += n_count
-
         l_per_bas += [l] * n_count
-        n_per_bas.extend(range(n_start, n_start+n_count))
+        n_per_bas.extend(cursor_per_l[l].add(n_count))
         msize = 2*l+1
-        ao_starts.extend(range(cursor, cursor+msize*n_count, msize))
-        cursor += msize*n_count
+        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):
+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.
 
     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)
+            - 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
+
     """
     x = []
+    L = []
     y = np.zeros((3, mol.nao), dtype=int)
-    i = 0
+    i = Cursor(action='slicer')
     a = mol.bas_exps()
     for iat in range(mol.natm):
         for bas_id in mol.atom_shell_ids(iat):
@@ -373,11 +374,13 @@ def basis_flatten(mol, return_both=True):
                 for c in cs.T:
                     ac = np.array([a[bas_id], c])
                     x.extend([ac]*msize)
-            y[:2,i:i+msize*n] = np.array([[iat, l]]*msize*n).T
-            y[2,i:i+msize*n] = [*get_mrange(l)]*n
-            i += msize*n
+            y[:,i(msize*n)] = np.vstack((np.array([[iat, l]]*msize*n).T, [*get_mrange(l)]*n))
+            if return_shells:
+                L.extend([l]*n)
+
+    ret = [y]
     if return_both:
-        x = stack_padding(x).transpose((1,0,2))
-        return y, x
-    else:
-        return y
+        ret.append(stack_padding(x).transpose((1,0,2)))
+    if return_shells:
+        ret.append(np.array(L))
+    return ret[0] if len(ret)==1 else ret

qstack/constants.py

@@ -3,10 +3,13 @@
 https://physics.nist.gov/cuu/Constants/
 https://physics.nist.gov/cuu/Constants/Table/allascii.txt
 """
+import math
+
+
 # Constants
 SPEED_LIGHT = 299792458.0
 PLANCK = 6.62607004e-34
-HBAR = PLANCK/(2*3.141592653589793)
+HBAR = PLANCK/(2*math.pi)
 FUND_CHARGE = 1.6021766208e-19
 MOL_NA = 6.022140857e23
 MASS_E = 9.10938356e-31
@@ -20,4 +23,4 @@
 BOHR2ANGS = 0.52917721092  # Angstroms
 HARTREE2J = HBAR**2/(MASS_E*(BOHR2ANGS*1e-10)**2)
 HARTREE2EV = 27.21138602
-AU2DEBYE = FUND_CHARGE * BOHR2ANGS*1e-10 / DEBYE # 2.541746
+AU2DEBYE = FUND_CHARGE * BOHR2ANGS*1e-10 / DEBYE  # 2.541746

qstack/equio.py

@@ -6,7 +6,8 @@
 import numpy as np
 from pyscf import data
 import metatensor
-from qstack.reorder import get_mrange
+from qstack.tools import Cursor
+from qstack.reorder import get_mrange, pyscf2gpr_l1_order
 from qstack.compound import singleatom_basis_enumerator
 
 
@@ -26,8 +27,6 @@
 
 _molid_name = 'mol_id'
 
-_pyscf2gpr_l1_order = [1,2,0]
-
 
 def _get_llist(mol):
     """Get list of angular momentum quantum numbers for basis functions of each element of a molecule.
@@ -50,7 +49,7 @@ def _get_tsize(tensor):
     Returns:
         int: Total size of the tensor (total number of elements).
     """
-    return sum([np.prod(tensor.block(key).values.shape) for key in tensor.keys])
+    return sum(np.prod(tensor.block(key).values.shape) for key in tensor.keys)
 
 
 def _labels_to_array(labels):
@@ -115,26 +114,24 @@ def vector_to_tensormap(mol, c):
     # Fill in the blocks
 
     iq = dict.fromkeys(llists.keys(), 0)
-    i = 0
+    i = Cursor(action='slicer')
     for q in atom_charges:
         if llists[q]==sorted(llists[q]):
             for l in set(llists[q]):
                 msize = 2*l+1
-                nsize = blocks[(l,q)].shape[-1]
-                cslice = c[i:i+nsize*msize].reshape(nsize,msize).T
+                nsize = blocks[l,q].shape[-1]
+                cslice = c[i(nsize*msize)].reshape(nsize,msize).T
                 if l==1:  # for l=1, the pyscf order is x,y,z (1,-1,0)
-                    cslice = cslice[_pyscf2gpr_l1_order]
-                blocks[(l,q)][iq[q],:,:] = cslice
-                i += msize*nsize
+                    cslice = cslice[pyscf2gpr_l1_order]
+                blocks[l,q][iq[q],:,:] = cslice
         else:
             il = dict.fromkeys(range(max(llists[q]) + 1), 0)
             for l in llists[q]:
                 msize = 2*l+1
-                cslice = c[i:i+msize]
+                cslice = c[i(msize)]
                 if l==1:  # for l=1, the pyscf order is x,y,z (1,-1,0)
-                    cslice = cslice[_pyscf2gpr_l1_order]
-                blocks[(l,q)][iq[q],:,il[l]] = cslice
-                i     += msize
+                    cslice = cslice[pyscf2gpr_l1_order]
+                blocks[l,q][iq[q],:,il[l]] = cslice
                 il[l] += 1
         iq[q] += 1
 
@@ -242,58 +239,54 @@ def matrix_to_tensormap(mol, dm):
 
     if all(llists[q]==sorted(llists[q]) for q in llists):
         iq1 = dict.fromkeys(elements, 0)
-        i1 = 0
+        i1 = Cursor(action='slicer')
         for iat1, q1 in enumerate(atom_charges):
             for l1 in set(llists[q1]):
                 msize1 = 2*l1+1
                 nsize1 = llists[q1].count(l1)
                 iq2 = dict.fromkeys(elements, 0)
-                i2 = 0
+                i1.add(nsize1*msize1)
+                i2 = Cursor(action='slicer')
                 for iat2, q2 in enumerate(atom_charges):
                     for l2 in set(llists[q2]):
                         msize2 = 2*l2+1
                         nsize2 = llists[q2].count(l2)
-                        dmslice = dm[i1:i1+nsize1*msize1,i2:i2+nsize2*msize2].reshape(nsize1,msize1,nsize2,msize2)
+                        dmslice = dm[i1(),i2(nsize2*msize2)].reshape(nsize1,msize1,nsize2,msize2)
                         dmslice = np.transpose(dmslice, axes=[1,3,0,2]).reshape(msize1,msize2,-1)
                         block = tensor_blocks[tm_label_vals.index((l1,l2,q1,q2))]
                         at_p = block.samples.position((iat1,iat2))
-                        blocks[(l1,l2,q1,q2)][at_p,:,:,:] = dmslice
-                        i2 += msize2*nsize2
+                        blocks[l1,l2,q1,q2][at_p,:,:,:] = dmslice
                     iq2[q2] += 1
-                i1 += msize1*nsize1
             iq1[q1] += 1
     else:
         iq1 = dict.fromkeys(elements, 0)
-        i1 = 0
+        i1 = Cursor(action='slicer')
         for iat1, q1 in enumerate(atom_charges):
             il1 = dict.fromkeys(range(max(llists[q1]) + 1), 0)
             for l1 in llists[q1]:
-                msize1 = 2*l1+1
+                i1.add(2*l1+1)
                 iq2 = dict.fromkeys(elements, 0)
-                i2 = 0
+                i2 = Cursor(action='slicer')
                 for iat2, q2 in enumerate(atom_charges):
                     il2 = dict.fromkeys(range(max(llists[q2]) + 1), 0)
                     for l2 in llists[q2]:
-                        msize2 = 2*l2+1
-                        dmslice = dm[i1:i1+msize1,i2:i2+msize2]
+                        dmslice = dm[i1(),i2(2*l2+1)]
                         block = tensor_blocks[tm_label_vals.index((l1, l2, q1, q2))]
                         at_p = block.samples.position((iat1, iat2))
                         n_p = block.properties.position((il1[l1], il2[l2]))
-                        blocks[(l1,l2,q1,q2)][at_p,:,:,n_p] = dmslice
-                        i2 += msize2
+                        blocks[l1,l2,q1,q2][at_p,:,:,n_p] = dmslice
                         il2[l2] += 1
                     iq2[q2] += 1
-                i1 += msize1
                 il1[l1] += 1
             iq1[q1] += 1
 
     # Fix the m order (for l=1, the pyscf order is x,y,z (1,-1,0))
     for key in blocks:
         l1,l2 = key[:2]
         if l1==1:
-            blocks[key] = np.ascontiguousarray(blocks[key][:,_pyscf2gpr_l1_order,:,:])
+            blocks[key] = np.ascontiguousarray(blocks[key][:,pyscf2gpr_l1_order,:,:])
         if l2==1:
-            blocks[key] = np.ascontiguousarray(blocks[key][:,:,_pyscf2gpr_l1_order,:])
+            blocks[key] = np.ascontiguousarray(blocks[key][:,:,pyscf2gpr_l1_order,:])
 
     # Build tensor map
     tensor_blocks = [metatensor.TensorBlock(values=blocks[key], samples=block_samp_labels[key], components=block_comp_labels[key], properties=block_prop_labels[key]) for key in tm_label_vals]
@@ -492,7 +485,7 @@ def split(tensor):
                 continue
             sampleidx = [t[0] for t in samples]
             samplelbl = [t[1] for t in samples]
-            #sampleidx = [block.samples.position(lbl) for lbl in samplelbl]
+            # sampleidx = [block.samples.position(lbl) for lbl in samplelbl]
 
             blocks[key] = block.values[sampleidx]
             block_samp_labels[key] = metatensor.Labels(tensor.sample_names[1:], np.array(samplelbl)[:,1:])

qstack/fields/dm.py

@@ -1,8 +1,9 @@
 """Density matrix manipulation and analysis functions."""
 
+import numpy as np
 from pyscf import dft
 from qstack import constants
-import numpy as np
+from qstack.tools import Cursor
 
 
 def get_converged_mf(mol, xc, dm0=None, verbose=False):
@@ -79,27 +80,26 @@ def sphericalize_density_matrix(mol, dm):
         A numpy ndarray with the sphericalized density matrix.
     """
     idx_by_l = [[] for i in range(constants.MAX_L)]
-    i0 = 0
+    i0 = Cursor(action='ranger')
     for ib in range(mol.nbas):
         l = mol.bas_angular(ib)
+        msize = 2*l+1
         nc = mol.bas_nctr(ib)
-        i1 = i0 + nc * (l*2+1)
-        idx_by_l[l].extend(range(i0, i1, l*2+1))
-        i0 = i1
+        idx_by_l[l].extend(i0(nc*msize)[::msize])
 
     spherical_dm = np.zeros_like(dm)
 
     for l in np.nonzero(idx_by_l)[0]:
+        msize = 2*l+1
         for idx in idx_by_l[l]:
             for jdx in idx_by_l[l]:
                 if l == 0:
                     spherical_dm[idx,jdx] = dm[idx,jdx]
                 else:
                     trace = 0
-                    for m in range(2*l+1):
+                    for m in range(msize):
                         trace += dm[idx+m,jdx+m]
-                    for m in range(2*l+1):
-                        spherical_dm[idx+m,jdx+m] = trace / (2*l+1)
+                    for m in range(msize):
+                        spherical_dm[idx+m,jdx+m] = trace / msize
 
     return spherical_dm
-

qstack/fields/dori.py

@@ -44,7 +44,7 @@ def eval_rho_dm(mol, ao, dm, deriv=2):
         if deriv==2:
             DM_dAO_dr_i = 2 * _dot_ao_dm(mol, dAO_dr[i], dm, None, None, None)
             for j in range(i, 3):
-                d2rho_dr2[i,j] = _contract_rho(dAO_dr[j], DM_dAO_dr_i) + 2.0*np.einsum('ip,ip->i', d2AO_dr2[triu_idx[(i,j)]], DM_AO)
+                d2rho_dr2[i,j] = _contract_rho(dAO_dr[j], DM_dAO_dr_i) + 2.0*np.einsum('ip,ip->i', d2AO_dr2[triu_idx[i,j]], DM_AO)
                 d2rho_dr2[j,i] = d2rho_dr2[i,j]
 
     if deriv==1:

qstack/fields/excited.py

@@ -120,7 +120,7 @@ def exciton_properties_dm(mol, hole, part):
     dist = np.linalg.norm(hole_r-part_r)
     hole_extent = np.sqrt(hole_r2-hole_r@hole_r)
     part_extent = np.sqrt(part_r2-part_r@part_r)
-    return(dist, hole_extent, part_extent)
+    return dist, hole_extent, part_extent
 
 
 def exciton_properties(mol, hole, part):

qstack/mathutils/array.py

@@ -1,6 +1,7 @@
 """Array manipulation utility functions."""
 
 import numpy as np
+from qstack.tools import slice_generator
 
 
 def scatter(values, indices):
@@ -89,9 +90,7 @@ def vstack_padding(xs):
     if len(np.unique(shapes_other_axes, axis=0))==1:
         return np.vstack(xs)
     X = np.zeros((shapes_common_axis.sum(), *shapes_other_axes.max(axis=0)))
-    idx = 0
-    for x in xs:
-        slices = (np.s_[idx:idx+x.shape[0]], *(np.s_[0:s] for s in x.shape[1:]))
+    for x, s0 in slice_generator(xs, inc=lambda x: x.shape[0]):
+        slices = (s0, *(np.s_[0:s] for s in x.shape[1:]))
         X[slices] = x
-        idx += x.shape[0]
     return X