Support contracted Gaussians in trexio interface (fix #82)

pyscf/tools/test/test_trexio.py

@@ -18,6 +18,14 @@ def test_mol_ae_6_31g():
     s1 = mol.intor('int1e_ovlp')
     assert abs(ref - s1).max() < 1e-12
 
+def test_mol_general_contraction():
+    mol = pyscf.M(atom='C', basis='ano')
+    ref = mol.intor('int1e_ovlp')
+    trexio.to_trexio(mol, filename)
+    mol = trexio.mol_from_trexio(filename)
+    s1 = mol.intor('int1e_ovlp')
+    assert abs(ref - s1).max() < 1e-12
+
 def test_mol_ccecp_ccpvqz():
     filename = 'test_mol_ccecp_ccpvqz_sphe.h5'
     mol = pyscf.M(atom='H 0 0 0; F 0 0 1', basis='ccecp-cc-pVQZ', ecp='ccECP', cart=False)
@@ -62,7 +70,7 @@ def test_mf_ae_6_31g():
     filename = 'test_mf_ae_6_31g_sphe.h5'
     mol = pyscf.M(atom='H 0 0 0; F 0 0 1', basis='6-31g*', cart=False)
     mf = mol.RHF().run()
-    print(mf.mo_coeff)
+    #print(mf.mo_coeff)
     trexio.to_trexio(mf, filename)
     mf1 = trexio.scf_from_trexio(filename)
     assert abs(mf1.mo_coeff - mf.mo_coeff).max() < 1e-12
@@ -73,7 +81,7 @@ def test_mf_ae_6_31g():
     filename = 'test_mf_ae_6_31g_cart.h5'
     mol = pyscf.M(atom='H 0 0 0; F 0 0 1', basis='6-31g*', cart=True)
     mf = mol.RHF().run()
-    print(mf.mo_coeff)
+    #print(mf.mo_coeff)
     trexio.to_trexio(mf, filename)
     mf1 = trexio.scf_from_trexio(filename)
     assert abs(mf1.mo_coeff - mf.mo_coeff).max() < 1e-12
@@ -111,4 +119,4 @@ def test_mf_ccecp_ccpvqz():
     #test_mol_ccecp_ccpvqz()
     #test_mol_ae_ccpvdz()
     #test_mf_ae_6_31g()
-    test_mf_ccecp_ccpvqz()
+    test_mf_ccecp_ccpvqz()

pyscf/tools/trexio.py

@@ -14,6 +14,7 @@
 import re
 import math
 import numpy as np
+import scipy.linalg
 
 import pyscf
 from pyscf import lib
@@ -67,6 +68,8 @@ def _mol_to_trexio(mol, trexio_file):
 
     # 3.1 Basis set
     trexio.write_basis_type(trexio_file, 'Gaussian')
+    if any(mol._bas[:,gto.NCTR_OF] > 1):
+        mol = _to_segment_contraction(mol)
     trexio.write_basis_shell_num(trexio_file, mol.nbas)
     trexio.write_basis_prim_num(trexio_file, int(mol._bas[:,gto.NPRIM_OF].sum()))
     trexio.write_basis_nucleus_index(trexio_file, mol._bas[:,gto.ATOM_OF])
@@ -75,8 +78,6 @@ def _mol_to_trexio(mol, trexio_file):
     prim2sh = [[ib]*nprim for ib, nprim in enumerate(mol._bas[:,gto.NPRIM_OF])]
     trexio.write_basis_shell_index(trexio_file, np.hstack(prim2sh))
     trexio.write_basis_exponent(trexio_file, np.hstack(mol.bas_exps()))
-    if not all(mol._bas[:,gto.NCTR_OF] == 1):
-        raise NotImplementedError('The generalized contraction is not supported.')
     coef = [mol.bas_ctr_coeff(i).ravel() for i in range(mol.nbas)]
     trexio.write_basis_coefficient(trexio_file, np.hstack(coef))
     prim_norms = [gto.gto_norm(mol.bas_angular(i), mol.bas_exp(i)) for i in range(mol.nbas)]
@@ -438,3 +439,23 @@ def read_det_trexio(filename):
         det = trexio.read_determinant_list(tf, offset_file, num_det)
         return num_det, coeff, det
 
+def _to_segment_contraction(mol):
+    '''transform generally contracted basis to segment contracted basis
+    '''
+    _bas = []
+    for shell in mol._bas:
+        nctr = shell[gto.NCTR_OF]
+        if nctr == 1:
+            _bas.append(shell)
+            continue
+
+        nprim = shell[gto.NPRIM_OF]
+        pcoeff = shell[gto.PTR_COEFF]
+        bs = np.repeat(shell[np.newaxis], nctr, axis=0)
+        bs[:,gto.NCTR_OF] = 1
+        bs[:,gto.PTR_COEFF] = np.arange(pcoeff, pcoeff+nprim*nctr, nprim)
+        _bas.append(bs)
+
+    pmol = mol.copy()
+    pmol._bas = np.asarray(np.vstack(_bas), dtype=np.int32)
+    return pmol