Task: Fix dimension bug in spinor X2C code (pyscf/pyscf#2611)

org: pyscf
repo: pyscf
pr_number: 2611
issue_title: Fix dimension bug in spinor X2C code
issue_body: Fix #2605 
resolved_issues: number: 2605
title: spin-orbit contracted basis set
body: In PySCF's basis set internal format, the kappa value can be used to treat spin-orbit contracted basis set.

> kappa can have value 
>  (corresponding to spinors with 
> ), 
>  (corresponding to spinors with 
> ) or 0. When kappa is 0, both types of spinors are assumed in the basis. 

```
[[angular, kappa, [[exp, c_1, c_2, ..],
                   [exp, c_1, c_2, ..],
                   ... ]],
 [angular, kappa, [[exp, c_1, c_2, ..],
                   [exp, c_1, c_2, ..]
                   ... ]]]
```
There is a possible bug in treating contracted basis set with contraction coefficients for different kappa. The following example with a fake basis set of this kind.
```
from pyscf import gto
from pyscf import scf

mol = gto.M(atom='''
O 0  0  0
H 0  1  1
H 0 -1 -1''',
            basis={
                'O':
                [[
                    0, [61420.0, 9e-05, -2e-05], [9199.0, 0.000698, -0.000159],
                    [2091.0, 0.003664, -0.000829],
                    [590.9, 0.015218, -0.003508], [192.3, 0.052423, -0.012156],
                    [69.32, 0.145921, -0.036261], [26.97, 0.305258, -0.082992],
                    [11.1, 0.398508, -0.15209], [4.682, 0.21698, -0.115331]
                ], [0, [0.5547, 1.0]],
                 [1, -2, [63.42, 0.006044], [14.66, 0.041799], [4.459, 0.1611]],
                 [1, 1, [34.46, 0.015928], [7.749, 0.09974], [2.28, 0.31049]],
                 [1, 1, [1.531, 1.0]], [1, 0, [0.5302, 1.0]]],
                'H':
                '631g'
            })
print(mol._basis)
print(mol.nao_nr())
print(mol.nao_2c())
mf = scf.X2C(mol)
mf.kernel()
```
The calculation won't run through.

Reference solution:
fix_patch: diff --git a/examples/x2c/01-x2c1e.py b/examples/x2c/01-x2c1e.py
new file mode 100644
index 0000000000..a2f9061c27
--- /dev/null
+++ b/examples/x2c/01-x2c1e.py
@@ -0,0 +1,40 @@
+#!/usr/bin/env python
+#
+# Author: Qiming Sun <osirpt.sun@gmail.com>
+#
+
+'''
+Constructing the X matrix using one-electron approximation. This is the default
+configuration of the X2C module.
+'''
+
+import pyscf
+from pyscf.x2c import UHF as X2C_UHF
+from pyscf.x2c import dft as x2c_dft
+
+mol = pyscf.M(
+    verbose = 0,
+    atom = [["O" , (0. , 0.     , 0.)],
+            [1   , (0. , -0.757 , 0.587)],
+            [1   , (0. , 0.757  , 0.587)] ],
+    basis = 'ccpvdz-dk',
+)
+
+method = mol.RHF()
+enr = method.kernel()
+print('E(NR) = %.12g' % enr)
+
+method = X2C_UHF(mol)
+ex2c = method.kernel()
+print('E(X2C1E-UHF) = %.12g' % ex2c)
+
+# Customizing basis sets for X matrix
+method.with_x2c.basis = {'O': 'unc-ccpvqz', 'H':'unc-ccpvdz'}
+print('E(X2C1E-UHF) = %.12g' % method.kernel())
+
+# Constructing X matrix using single center approximation
+method.with_x2c.approx = 'atom1e'
+print('E(X2C1E-UHF) = %.12g' % method.kernel())
+
+method = x2c_dft.UKS(mol).run()
+print('E(X2C1E-UKS) = %.12g' % method.e_tot)
diff --git a/pyscf/hessian/uks.py b/pyscf/hessian/uks.py
index 74bdb85bac..dfb4592edc 100644
--- a/pyscf/hessian/uks.py
+++ b/pyscf/hessian/uks.py
@@ -41,7 +41,7 @@ def partial_hess_elec(hessobj, mo_energy=None, mo_coeff=None, mo_occ=None,
     mf = hessobj.base
     ni = mf._numint
     if mf.do_nlc():
-        raise NotImplementedError('RKS Hessian for NLC functional')
+        raise NotImplementedError('UKS Hessian for NLC functional')
 
     if mo_energy is None: mo_energy = mf.mo_energy
     if mo_occ is None:    mo_occ = mf.mo_occ
diff --git a/pyscf/x2c/x2c.py b/pyscf/x2c/x2c.py
index d1b0e34ed6..1754f3cd7c 100644
--- a/pyscf/x2c/x2c.py
+++ b/pyscf/x2c/x2c.py
@@ -61,7 +61,10 @@ def get_xmol(self, mol=None):
             xmol.build(False, False, basis=self.basis)
             return xmol, None
         elif self.xuncontract:
-            xmol, contr_coeff = _uncontract_mol(mol, self.xuncontract)
+            if all(mol._bas[:,mole.KAPPA_OF] == 0):
+                xmol, contr_coeff = _uncontract_mol(mol, self.xuncontract)
+            else:
+                xmol, contr_coeff = _decontract_spinor(mol, self.xuncontract)
             return xmol, contr_coeff
         else:
             return mol, None
@@ -113,9 +116,13 @@ def get_hcore(self, mol=None):
             h1 = reduce(numpy.dot, (c.T.conj(), h1, c))
         elif self.xuncontract:
             np, nc = contr_coeff_nr.shape
-            contr_coeff = numpy.zeros((np*2,nc*2))
-            contr_coeff[0::2,0::2] = contr_coeff_nr
-            contr_coeff[1::2,1::2] = contr_coeff_nr
+            if np == h1.shape[0]:
+                # spinor adapted contraction matrix generated by _decontract_spinor
+                contr_coeff = contr_coeff_nr
+            else:
+                contr_coeff = numpy.zeros((np*2,nc*2))
+                contr_coeff[0::2,0::2] = contr_coeff_nr
+                contr_coeff[1::2,1::2] = contr_coeff_nr
             h1 = reduce(numpy.dot, (contr_coeff.T.conj(), h1, contr_coeff))
         return h1
 
@@ -986,30 +993,89 @@ def _sigma_dot(mat):
     nao = mat.shape[-1] * 2
     return lib.einsum('sxy,spq->xpyq', quaternion, mat).reshape(nao, nao)
 
-
-if __name__ == '__main__':
-    from pyscf.x2c import dft
-    mol = mole.Mole()
-    mol.build(
-        verbose = 0,
-        atom = [["O" , (0. , 0.     , 0.)],
-                [1   , (0. , -0.757 , 0.587)],
-                [1   , (0. , 0.757  , 0.587)] ],
-        basis = 'ccpvdz-dk',
-    )
-
-    method = hf.RHF(mol)
-    enr = method.kernel()
-    print('E(NR) = %.12g' % enr)
-
-    method = UHF(mol)
-    ex2c = method.kernel()
-    print('E(X2C1E) = %.12g' % ex2c)
-    method.with_x2c.basis = {'O': 'unc-ccpvqz', 'H':'unc-ccpvdz'}
-    print('E(X2C1E) = %.12g' % method.kernel())
-    method.with_x2c.approx = 'atom1e'
-    print('E(X2C1E) = %.12g' % method.kernel())
-
-    method = dft.UKS(mol)
-    ex2c = method.kernel()
-    print('E(X2C1E) = %.12g' % ex2c)
+def _decontract_spinor(mol, atoms=None):
+    pmol = mol.copy(deep=False)
+
+    # Some input basis may be segmented basis from a general contracted set.
+    # This may lead to duplicated pGTOs. First contract all basis to remove
+    # duplicated primitive functions.
+    bas_exps = mol.bas_exps()
+    _bas = []
+    env = [mol._env.copy()]
+    contr_coeff = []
+    pexp = env[0].size
+    ao_loc = mol.ao_loc_2c()
+
+    aoslices = mol.aoslice_by_atom()
+    for ia, (ib0, ib1) in enumerate(aoslices[:,:2]):
+        if ib0 == ib1: # No basis on atom ia
+            continue
+
+        if atoms is not None:
+            if isinstance(atoms, str):
+                to_apply = ((atoms == mol.atom_pure_symbol(ia)) or
+                            (atoms == mol.atom_symbol(ia)))
+            elif isinstance(atoms, (tuple, list)):
+                to_apply = ((mol.atom_pure_symbol(ia) in atoms) or
+                            (mol.atom_symbol(ia) in atoms) or
+                            (ia in atoms))
+            else:
+                to_apply = True
+            if not to_apply:
+                _bas.append(mol._bas[ib0:ib1])
+                contr_coeff.append(numpy.eye(ao_loc[ib1] - ao_loc[ib0]))
+                continue
+
+        lmax = mol._bas[ib0:ib1,mole.ANG_OF].max()
+        for l in range(lmax+1):
+            bas_idx = ib0 + numpy.where(mol._bas[ib0:ib1,mole.ANG_OF] == l)[0]
+            if len(bas_idx) == 0:
+                continue
+            if bas_idx[0] + len(bas_idx) != bas_idx[-1] + 1:
+                raise NotImplementedError('Discontinuous bases of same angular momentum')
+
+            es = numpy.hstack([bas_exps[i] for i in bas_idx])
+            _, e_idx, rev_idx = numpy.unique(es.round(9), True, True)
+            cs = scipy.linalg.block_diag(
+                *[mol._libcint_ctr_coeff(i) for i in bas_idx])
+            nprim = len(e_idx)
+            if len(es) != nprim:
+                cs_new = numpy.zeros((nprim, cs.shape[1]))
+                for i, j in enumerate(rev_idx):
+                    cs_new[j] += cs[i]
+                es = es[e_idx][::-1]
+                cs = cs_new[::-1]
+
+            bs = numpy.zeros((nprim, mole.BAS_SLOTS), dtype=numpy.int32)
+            bs[:,mole.ATOM_OF] = ia
+            bs[:,mole.ANG_OF ] = l
+            bs[:,mole.NCTR_OF] = bs[:,mole.NPRIM_OF] = 1
+            bs[:,mole.PTR_EXP] = pexp + numpy.arange(nprim)
+            bs[:,mole.PTR_COEFF] = pexp + numpy.arange(nprim, nprim*2)
+            norm = mole.gto_norm(l, es)
+            env.append(es)
+            env.append(norm)
+            pexp += nprim * 2
+            _bas.append(bs)
+
+            degen = 4 * l + 2
+            c = numpy.einsum('pi,p,xm->pxim', cs, 1./norm, numpy.eye(degen))
+            c = c.reshape(nprim*degen, -1)
+            full_idx = numpy.arange(c.shape[1]).reshape(-1, degen)
+            spinor_idx = []
+            p0 = p1 = 0
+            for ib in bas_idx:
+                p0, p1 = p1, p1 + mol.bas_nctr(ib)
+                kappa = mol.bas_kappa(ib)
+                if kappa == 0:
+                    spinor_idx.append(full_idx[p0:p1].ravel())
+                elif kappa > 0:
+                    spinor_idx.append(full_idx[p0:p1,:l*2].ravel())
+                else:
+                    spinor_idx.append(full_idx[p0:p1,l*2:].ravel())
+            contr_coeff.append(c[:,numpy.hstack(spinor_idx)])
+
+    pmol._bas = numpy.asarray(numpy.vstack(_bas), dtype=numpy.int32)
+    pmol._env = numpy.hstack(env)
+    contr_coeff = scipy.linalg.block_diag(*contr_coeff)
+    return pmol, contr_coeff