Repository: pyscf/pyscf — pull request #2354 (python)
Title: Fix the transformation matrix of Mole.decontracted_basis

fix #2348

Resolved issues:
title: `mole.to_uncontracted_cartesian_basis` behavior change
body: The behavior of `mole.to_uncontracted_cartesian_basis` has chaged since v2.1. The coefficients returned no longer have the size of the number of shells of the original Mole object. This breaks some code for some basis sets, e.g., this [line]( https://github.com/pyscf/pyscf/blob/ca8c7c1680defdfee2380eda3af3a28d9fb375cb/pyscf/pbc/dft/multigrid/multigrid.py#L903).
I think we should modify the docstring and maybe also add a warning.

### Minimal API referenced by tests (typed signatures)
_Added only when tests call brand‑new symbols introduced by the fix; types reflect annotations present in the repo or `Unknown` if not annotated._

- Function in `pyscf/gto/mole.py`: `def decontract_basis(mol, atoms, to_cart, aggregate) -> Union[list, tuple]: ...`
  **Inputs:** mol: gto.Mole, atoms: Union[list, tuple], to_cart: bool, aggregate: bool
  **Output:** Union[list, tuple]

```python
def decontract_basis(mol, atoms, to_cart, aggregate) -> Union[list, tuple]: ...
```

Hints:

Fix patch:
diff --git a/pyscf/gto/mole.py b/pyscf/gto/mole.py
index 0f2ec506f9..24319af8cc 100644
--- a/pyscf/gto/mole.py
+++ b/pyscf/gto/mole.py
@@ -572,7 +572,7 @@ def to_uncontracted_cartesian_basis(mol):
     '''
     return decontract_basis(mol, to_cart=True)
 
-def decontract_basis(mol, atoms=None, to_cart=False):
+def decontract_basis(mol, atoms=None, to_cart=False, aggregate=False):
     '''Decontract the basis of a Mole or a Cell.  Returns a Mole (Cell) object
     with the uncontracted basis environment and a list of coefficients that
     transform the uncontracted basis to the original basis. Each element in
@@ -584,6 +584,9 @@ def decontract_basis(mol, atoms=None, to_cart=False):
             are decontracted
         to_cart: bool
             Decontract basis and transfer to Cartesian basis
+        aggregate: bool
+            Whether to aggregate the transformation coefficients into a giant
+            transformation matrix
 
     Examples:
 
@@ -602,17 +605,27 @@ def decontract_basis(mol, atoms=None, to_cart=False):
     bas_exps = mol.bas_exps()
     def _to_full_contraction(mol, bas_idx):
         es = numpy.hstack([bas_exps[i] for i in bas_idx])
-        cs = scipy.linalg.block_diag(*[mol._libcint_ctr_coeff(ib) for ib in bas_idx])
-
-        es, e_idx, rev_idx = numpy.unique(es.round(9), True, True)
-        cs_new = numpy.zeros((es.size, cs.shape[1]))
-        for i, j in enumerate(rev_idx):
-            cs_new[j] += cs[i]
-        return es[::-1], cs_new[::-1]
+        _, e_idx, rev_idx = numpy.unique(es.round(9), True, True)
+        if aggregate:
+            cs = scipy.linalg.block_diag(
+                *[mol._libcint_ctr_coeff(i) for i in bas_idx])
+            if len(es) != len(e_idx):
+                cs_new = numpy.zeros((e_idx.size, cs.shape[1]))
+                for i, j in enumerate(rev_idx):
+                    cs_new[j] += cs[i]
+                es = es[e_idx][::-1]
+                cs = cs_new[::-1]
+            yield es, cs
+        else:
+            if len(es) != len(e_idx):
+                raise RuntimeError('Duplicated pGTOs across shells')
+            for i in bas_idx:
+                yield bas_exps[i], mol._libcint_ctr_coeff(i)
 
     _bas = []
-    _env = mol._env.copy()
+    env = [mol._env.copy()]
     contr_coeff = []
+    pexp = env[0].size
 
     lmax = mol._bas[:,ANG_OF].max()
     if mol.cart:
@@ -648,7 +661,7 @@ def _to_full_contraction(mol, bas_idx):
                 continue
 
         lmax = mol._bas[ib0:ib1,ANG_OF].max()
-        pexp = mol._bas[ib0,PTR_EXP]
+
         for l in range(lmax+1):
             bas_idx = ib0 + numpy.where(mol._bas[ib0:ib1,ANG_OF] == l)[0]
             if len(bas_idx) == 0:
@@ -656,30 +669,27 @@ def _to_full_contraction(mol, bas_idx):
             if bas_idx[0] + len(bas_idx) != bas_idx[-1] + 1:
                 raise NotImplementedError('Discontinuous bases of same angular momentum')
 
-            mol_exps, b_coeff = _to_full_contraction(mol, bas_idx)
-            nprim, nc = b_coeff.shape
-            bs = numpy.zeros((nprim, BAS_SLOTS), dtype=numpy.int32)
-            bs[:,ATOM_OF] = ia
-            bs[:,ANG_OF ] = l
-            bs[:,NCTR_OF] = bs[:,NPRIM_OF] = 1
-            norm = gto_norm(l, mol_exps)
-            if atoms is None:
+            for mol_exps, b_coeff in _to_full_contraction(mol, bas_idx):
+                nprim, nc = b_coeff.shape
+                bs = numpy.zeros((nprim, BAS_SLOTS), dtype=numpy.int32)
+                bs[:,ATOM_OF] = ia
+                bs[:,ANG_OF ] = l
+                bs[:,NCTR_OF] = bs[:,NPRIM_OF] = 1
                 bs[:,PTR_EXP] = pexp + numpy.arange(nprim)
                 bs[:,PTR_COEFF] = pexp + numpy.arange(nprim, nprim*2)
-                _env[pexp:pexp+nprim] = mol_exps
-                _env[pexp+nprim:pexp+nprim*2] = norm
+                norm = gto_norm(l, mol_exps)
+                env.append(mol_exps)
+                env.append(norm)
                 pexp += nprim * 2
-            else:
-                bs[:,PTR_EXP] = _env.size + numpy.arange(nprim)
-                bs[:,PTR_COEFF] = _env.size + numpy.arange(nprim, nprim*2)
-                _env = np.hstack([_env, mol_exps, norm])
-            _bas.append(bs)
+                _bas.append(bs)
 
-            c = numpy.einsum('pi,p,xm->pxim', b_coeff, 1./norm, c2s[l])
-            contr_coeff.append(c.reshape(nprim * c2s[l].shape[0], -1))
+                c = numpy.einsum('pi,p,xm->pxim', b_coeff, 1./norm, c2s[l])
+                contr_coeff.append(c.reshape(nprim * c2s[l].shape[0], -1))
 
     pmol._bas = numpy.asarray(numpy.vstack(_bas), dtype=numpy.int32)
-    pmol._env = _env
+    pmol._env = numpy.hstack(env)
+    if aggregate:
+        contr_coeff = scipy.linalg.block_diag(*contr_coeff)
     return pmol, contr_coeff
 
 def format_ecp(ecp_tab):
diff --git a/pyscf/pbc/df/ft_ao.py b/pyscf/pbc/df/ft_ao.py
index 4215251fdd..b85524abb7 100644
--- a/pyscf/pbc/df/ft_ao.py
+++ b/pyscf/pbc/df/ft_ao.py
@@ -526,8 +526,9 @@ def assign(type_code):
         assign(SMOOTH_BASIS)
         return ao_type
 
-    def decontract_basis(self, to_cart=True):
-        pcell, ctr_coeff = self.ref_cell.decontract_basis(to_cart=to_cart)
+    def decontract_basis(self, to_cart=True, aggregate=False):
+        pcell, ctr_coeff = self.ref_cell.decontract_basis(to_cart=to_cart,
+                                                          aggregate=aggregate)
         pcell = pcell.view(self.__class__)
         pcell.ref_cell = None
 
diff --git a/pyscf/pbc/dft/multigrid/multigrid.py b/pyscf/pbc/dft/multigrid/multigrid.py
index 23f2bbb83f..2bcb533302 100644
--- a/pyscf/pbc/dft/multigrid/multigrid.py
+++ b/pyscf/pbc/dft/multigrid/multigrid.py
@@ -601,16 +601,16 @@ def _eval_rhoG(mydf, dm_kpts, hermi=1, kpts=numpy.zeros((1,3)), deriv=0,
             dms_ht = numpy.asarray(dms[:,:,idx_h[:,None],idx_t], order='C')
             dms_lh = numpy.asarray(dms[:,:,idx_l[:,None],idx_h], order='C')
 
-            t_cell = h_cell + grids_sparse.cell
+            l_cell = grids_sparse.cell
+            h_pcell, h_coeff = h_cell.decontract_basis(to_cart=True, aggregate=True)
+            l_pcell, l_coeff = l_cell.decontract_basis(to_cart=True, aggregate=True)
+            t_cell = h_pcell + l_pcell
+            t_coeff = scipy.linalg.block_diag(h_coeff, l_coeff)
+
             nshells_h = _pgto_shells(h_cell)
             nshells_t = _pgto_shells(t_cell)
-            t_cell, t_coeff = t_cell.to_uncontracted_cartesian_basis()
 
             if deriv == 0:
-                h_coeff = scipy.linalg.block_diag(*t_coeff[:h_cell.nbas])
-                l_coeff = scipy.linalg.block_diag(*t_coeff[h_cell.nbas:])
-                t_coeff = scipy.linalg.block_diag(*t_coeff)
-
                 if hermi == 1:
                     naol, naoh = dms_lh.shape[2:]
                     dms_ht[:,:,:,naoh:] += dms_lh.transpose(0,1,3,2)
@@ -635,10 +635,6 @@ def _eval_rhoG(mydf, dm_kpts, hermi=1, kpts=numpy.zeros((1,3)), deriv=0,
                                          'LDA', kpts, grids_dense, ignore_imag, log)
 
             elif deriv == 1:
-                h_coeff = scipy.linalg.block_diag(*t_coeff[:h_cell.nbas])
-                l_coeff = scipy.linalg.block_diag(*t_coeff[h_cell.nbas:])
-                t_coeff = scipy.linalg.block_diag(*t_coeff)
-
                 pgto_dms = lib.einsum('nkij,pi,qj->nkpq', dms_ht, h_coeff, t_coeff)
                 shls_slice = (0, nshells_h, 0, nshells_t)
                 #:rho = eval_rho(t_cell, pgto_dms, shls_slice, 0, 'GGA', kpts,
@@ -895,13 +891,13 @@ def _get_j_pass2(mydf, vG, hermi=1, kpts=numpy.zeros((1,3)), verbose=None):
 
             h_cell = grids_dense.cell
             l_cell = grids_sparse.cell
-            t_cell = h_cell + l_cell
-            t_cell, coeff = t_cell.to_uncontracted_cartesian_basis()
+            h_pcell, h_coeff = h_cell.decontract_basis(to_cart=True, aggregate=True)
+            l_pcell, l_coeff = l_cell.decontract_basis(to_cart=True, aggregate=True)
+            t_cell = h_pcell + l_pcell
+            t_coeff = scipy.linalg.block_diag(h_coeff, l_coeff)
+
             nshells_h = _pgto_shells(h_cell)
             nshells_t = _pgto_shells(t_cell)
-
-            h_coeff = scipy.linalg.block_diag(*coeff[:h_cell.nbas])
-            t_coeff = scipy.linalg.block_diag(*coeff)
             shls_slice = (0, nshells_h, 0, nshells_t)
             vp = eval_mat(t_cell, vR, shls_slice, 1, 0, 'LDA', kpts)
             # Imaginary part may contribute
@@ -919,7 +915,6 @@ def _get_j_pass2(mydf, vG, hermi=1, kpts=numpy.zeros((1,3)), verbose=None):
                 vj_kpts[:,:,idx_l[:,None],idx_h] += \
                         vp[:,:,:,naoh:].transpose(0,1,3,2).conj()
             else:
-                l_coeff = scipy.linalg.block_diag(*coeff[h_cell.nbas:])
                 shls_slice = (nshells_h, nshells_t, 0, nshells_h)
                 vp = eval_mat(t_cell, vR, shls_slice, 1, 0, 'LDA', kpts)
                 # Imaginary part may contribute
@@ -996,15 +991,13 @@ def _get_gga_pass2(mydf, vG, hermi=1, kpts=numpy.zeros((1,3)), verbose=None):
 
             h_cell = grids_dense.cell
             l_cell = grids_sparse.cell
-            t_cell = h_cell + l_cell
-            t_cell, coeff = t_cell.to_uncontracted_cartesian_basis()
+            h_pcell, h_coeff = h_cell.decontract_basis(to_cart=True, aggregate=True)
+            l_pcell, l_coeff = l_cell.decontract_basis(to_cart=True, aggregate=True)
+            t_cell = h_pcell + l_pcell
+            t_coeff = scipy.linalg.block_diag(h_coeff, l_coeff)
+
             nshells_h = _pgto_shells(h_cell)
             nshells_t = _pgto_shells(t_cell)
-
-            h_coeff = scipy.linalg.block_diag(*coeff[:h_cell.nbas])
-            l_coeff = scipy.linalg.block_diag(*coeff[h_cell.nbas:])
-            t_coeff = scipy.linalg.block_diag(*coeff)
-
             shls_slice = (0, nshells_h, 0, nshells_t)
             vpR = eval_mat(t_cell, vR, shls_slice, 1, 0, 'GGA', kpts)
             vp = vpR = lib.einsum('nkpq,pi,qj->nkij', vpR, h_coeff, t_coeff)
diff --git a/pyscf/x2c/x2c.py b/pyscf/x2c/x2c.py
index 291cd8e40e..4bd83298d3 100644
--- a/pyscf/x2c/x2c.py
+++ b/pyscf/x2c/x2c.py
@@ -810,8 +810,7 @@ def to_gpu(self):
 
 def _uncontract_mol(mol, xuncontract=None, exp_drop=0.2):
     '''mol._basis + uncontracted steep functions'''
-    pmol, contr_coeff = mol.decontract_basis(atoms=xuncontract)
-    contr_coeff = scipy.linalg.block_diag(*contr_coeff)
+    pmol, contr_coeff = mol.decontract_basis(atoms=xuncontract, aggregate=True)
     return pmol, contr_coeff
 
 


Test patch:
diff --git a/pyscf/gto/test/test_mole.py b/pyscf/gto/test/test_mole.py
index 1a91682db0..b0ffeb993c 100644
--- a/pyscf/gto/test/test_mole.py
+++ b/pyscf/gto/test/test_mole.py
@@ -1037,14 +1037,27 @@ def test_decontract_basis(self):
         mol._bas = mol._bas[:5]
         pmol, c = mol.decontract_basis()
         self.assertEqual(pmol.nbas, 14)
+        self.assertEqual(len(c), 5)
 
         mol = gto.M(atom='He',
                     basis=('ccpvdz', [[0, [5, 1]], [1, [3, 1]]]))
         pmol, contr_coeff = mol.decontract_basis()
+        self.assertEqual(len(contr_coeff), 5)
         contr_coeff = scipy.linalg.block_diag(*contr_coeff)
         s = contr_coeff.T.dot(pmol.intor('int1e_ovlp')).dot(contr_coeff)
         self.assertAlmostEqual(abs(s - mol.intor('int1e_ovlp')).max(), 0, 12)
 
+        mol = gto.M(atom='H 0 0 0; F 0 0 1', basis=[[0, (2, .5), (1, .5)],
+                                                    [0, (2, .1), (1, .9)],
+                                                    [0, (4., 1)]])
+        with self.assertRaises(RuntimeError):
+            mol.decontract_basis(aggregate=False)
+        pmol, c = mol.decontract_basis(aggregate=True)
+        self.assertEqual(pmol.nbas, 6)
+        self.assertEqual(c.shape, (6, 6))
+        s = c.T.dot(pmol.intor('int1e_ovlp')).dot(c)
+        self.assertAlmostEqual(abs(s - mol.intor('int1e_ovlp')).max(), 0, 12)
+
     def test_ao_rotation_matrix(self):
         mol = gto.M(atom='O 0 0 0.2; H1 0 -.8 -.5; H2 0 .8 -.5', basis='ccpvdz')
         numpy.random.seed(1)