Add example/test of manual reordering

qstack/orcaio.py

@@ -175,6 +175,38 @@ def read_basis(MAX_PRIMITIVES=37):
         return coefficients_ab, energies_ab, occupations_ab, ls
 
 
+def _get_indices(mol, ls_from_orca):
+    """ Get coefficient needed to reorder the AO read from Orca.
+
+    Args:
+        mol (pyscf Mole): pyscf Mole object.
+        ls_from_orca : dict of {int : [int]} with a list of basis functions
+                       angular momenta for those atoms (not elements!)
+                       whose basis functions are *not* sorted wrt to angular momenta.
+                       The lists represent the Orca order.
+
+    Returns:
+        numpy int 1darray of (nao,) containing the indices to be used as
+                c_reordered = c_orca[indices]
+    """
+    if ls_from_orca is None:
+        return None
+    ao_limits = mol.offset_ao_by_atom()[:,2:]
+    indices_full = np.arange(mol.nao)
+    for iat, ls in ls_from_orca.items():
+        indices = []
+        i = 0
+        for il, l in enumerate(ls):
+            indices.append((l, il, i + np.arange(2*l+1)))
+            i += 2*l+1
+        indices = sorted(indices, key=lambda x: (x[0], x[1]))
+        indices = np.array([j for i in indices for j in i[2]])
+        atom_slice = np.s_[ao_limits[iat][0]:ao_limits[iat][1]]
+        indices_full[atom_slice] = indices[:] + ao_limits[iat][0]
+    assert np.all(sorted(indices_full)==np.arange(mol.nao))
+    return indices_full
+
+
 def reorder_coeff_inplace(c_full, mol, reorder_dest='pyscf', ls_from_orca=None):
     """ Reorder coefficient read from ORCA .gbw
 
@@ -192,22 +224,10 @@ def reorder_coeff_inplace(c_full, mol, reorder_dest='pyscf', ls_from_orca=None):
     def _reorder_coeff(c):
         # In ORCA, at least def2-SVP and def2-TZVP for 3d metals
         # are not sorted wrt angular momenta
-        if ls_from_orca is not None:
-            indices_full = np.arange(mol.nao)
-            for iat, ls in ls_from_orca.items():
-                indices = []
-                i = 0
-                for il, l in enumerate(ls):
-                    indices.append((l, il, i + np.arange(2*l+1)))
-                    i += 2*l+1
-                indices = sorted(indices, key=lambda x: (x[0], x[1]))
-                indices = np.array([j for i in indices for j in i[2]])
-                ao_limits = mol.offset_ao_by_atom()[iat][2:]
-                atom_slice = np.s_[ao_limits[0]:ao_limits[1]]
-                indices_full[atom_slice] = indices[:] + ao_limits[0]
-            for i in range(len(c)):
-                c[:,i] = c[indices_full,i]
+        idx = _get_indices(mol, ls_from_orca)
         for i in range(len(c)):
+            if idx is not None:
+                c[:,i] = c[idx,i]
             c[:,i] = reorder_ao(mol, c[:,i], src='orca', dest=reorder_dest)
     for i in range(c_full.shape[0]):
         _reorder_coeff(c_full[i])

tests/test_orca.py

@@ -2,8 +2,10 @@
 
 import os
 import numpy as np
-import qstack
 from pyscf.data import elements
+import qstack
+from qstack.tools import reorder_ao
+
 
 def _dipole_moment(mol, dm):
     charges = mol.atom_charges()
@@ -64,9 +66,8 @@ def compare_MO(c0, c1):
 
 def test_orca_gbw_reader_def2tzvp():
     path = os.path.dirname(os.path.realpath(__file__))
-    mol = qstack.compound.xyz_to_mol(path+'/data/orca/CEHZOF/CEHZOF.xyz', 'def2tzvp', charge=1, spin=1, ecp='def2tzvp')
+    mol = qstack.compound.xyz_to_mol(path+'/data/orca/CEHZOF/CEHZOF.xyz', 'def2tzvp', ecp='def2tzvp')
     c504, e504, occ504 = qstack.orcaio.read_gbw(mol, path+'/data/orca/CEHZOF/CEHZOF_1_SPE.gbw')
-
     dm = np.zeros_like(c504)
     for i, (ci, occi) in enumerate(zip(c504, occ504)):
         dm[i,:,:] = (ci[:,occi>0] * occi[occi>0]) @ ci[:,occi>0].T
@@ -85,8 +86,26 @@ def test_orca_input_reader():
     assert np.allclose(mol.atom_coords(), mol0.atom_coords())
 
 
+def test_orca_density_reader_def2tzvp():
+    path = os.path.dirname(os.path.realpath(__file__))
+    mol = qstack.compound.xyz_to_mol(path+'/data/orca/CEHZOF/CEHZOF.xyz', 'def2tzvp', ecp='def2tzvp')
+    c, e, occ = qstack.orcaio.read_gbw(mol, path+'/data/orca/CEHZOF/CEHZOF_1_SPE.gbw')
+    c = c.squeeze()
+    occ = occ.squeeze()
+    dm0 = (c[:,occ>0] * occ[occ>0]) @ c[:,occ>0].T
+
+    dm = qstack.orcaio.read_density(mol, 'CEHZOF_1_SPE', directory=path+'/data/orca/CEHZOF', version=504, reorder_dest=None)
+    Co_idx = [mol.atom_symbol(i) for i in range(mol.natm)].index('Co')
+    ls_from_orca = {Co_idx : [0]*6 + [1]*3 + [2]*3 + [1, 2, 3]}
+    idx = qstack.orcaio._get_indices(mol, ls_from_orca)
+    dm = dm[np.ix_(idx,idx)]
+    dm = reorder_ao(mol, dm, src='orca', dest='pyscf')
+    assert np.linalg.norm(dm-dm0) < 1e-14
+
+
 if __name__ == '__main__':
     test_orca_input_reader()
     test_orca_density_reader()
     test_orca_gbw_reader()
     test_orca_gbw_reader_def2tzvp()
+    test_orca_density_reader_def2tzvp()