diff --git a/pyscf/tools/test/test_trexio.py b/pyscf/tools/test/test_trexio.py index fe0cdce..020498f 100644 --- a/pyscf/tools/test/test_trexio.py +++ b/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() \ No newline at end of file + test_mf_ccecp_ccpvqz() diff --git a/pyscf/tools/trexio.py b/pyscf/tools/trexio.py index cce6927..2139909 100644 --- a/pyscf/tools/trexio.py +++ b/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