diff --git a/.github/workflows/run_ci.sh b/.github/workflows/run_ci.sh index 5531b7103..dc426c98f 100755 --- a/.github/workflows/run_ci.sh +++ b/.github/workflows/run_ci.sh @@ -9,6 +9,7 @@ sudo apt-get -qq install \ curl python -m pip install --upgrade pip +pip install "scipy<1.16" pip install pytest pip install . diff --git a/pyscf/tools/trexio.py b/pyscf/tools/trexio.py index 33cfc9794..b91c72e64 100644 --- a/pyscf/tools/trexio.py +++ b/pyscf/tools/trexio.py @@ -21,16 +21,21 @@ from pyscf import gto from pyscf import scf from pyscf import pbc +from pyscf import mcscf from pyscf import fci import trexio -def to_trexio(obj, filename, backend='h5'): +def to_trexio(obj, filename, backend='h5', ci_threshold=None, chunk_size=None): with trexio.File(filename, 'u', back_end=_mode(backend)) as tf: if isinstance(obj, gto.Mole): _mol_to_trexio(obj, tf) elif isinstance(obj, scf.hf.SCF): _scf_to_trexio(obj, tf) + elif isinstance(obj, mcscf.casci.CASCI) or isinstance(obj, mcscf.CASSCF): + ci_threshold = ci_threshold if ci_threshold is not None else 0. + chunk_size = chunk_size if chunk_size is not None else 100000 + _mcscf_to_trexio(obj, tf, ci_threshold=ci_threshold, chunk_size=chunk_size) else: raise NotImplementedError(f'Conversion function for {obj.__class__}') @@ -248,8 +253,39 @@ def _scf_to_trexio(mf, trexio_file): def _cc_to_trexio(cc_obj, trexio_file): raise NotImplementedError -def _mcscf_to_trexio(cas_obj, trexio_file): - raise NotImplementedError +def _mcscf_to_trexio(cas_obj, trexio_file, ci_threshold=0., chunk_size=100000): + mol = cas_obj.mol + _mol_to_trexio(mol, trexio_file) + mo_energy_cas = cas_obj.mo_energy + mo_cas = cas_obj.mo_coeff + num_mo = mo_energy_cas.size + spin_cas = np.zeros(mo_energy_cas.size, dtype=int) + mo_type_cas = 'CAS' + trexio.write_mo_type(trexio_file, mo_type_cas) + idx = _order_ao_index(mol) + trexio.write_mo_num(trexio_file, num_mo) + trexio.write_mo_coefficient(trexio_file, mo_cas[idx].T.ravel()) + trexio.write_mo_energy(trexio_file, mo_energy_cas) + trexio.write_mo_spin(trexio_file, spin_cas) + + ncore = cas_obj.ncore + ncas = cas_obj.ncas + mo_classes = np.array(["Virtual"] * num_mo, dtype=str) # Initialize all MOs as Virtual + mo_classes[:ncore] = "Core" + mo_classes[ncore:ncore + ncas] = "Active" + trexio.write_mo_class(trexio_file, list(mo_classes)) + + occupation = np.zeros(num_mo) + occupation[:ncore] = 2.0 + rdm1 = cas_obj.fcisolver.make_rdm1(cas_obj.ci, ncas, cas_obj.nelecas) + natural_occ = np.linalg.eigh(rdm1)[0] + occupation[ncore:ncore + ncas] = natural_occ[::-1] + occupation[ncore + ncas:] = 0.0 + trexio.write_mo_occupation(trexio_file, occupation) + + total_elec_cas = sum(cas_obj.nelecas) + + det_to_trexio(cas_obj, ncas, total_elec_cas, trexio_file, ci_threshold, chunk_size) def mol_from_trexio(filename): mol = gto.Mole() @@ -329,19 +365,25 @@ def write_eri(eri, filename, backend='h5'): idx[:,:,2:] = idx_pair[None,:,:] idx = idx[np.tril_indices(npair)] + # Physicist notation + idx=idx.reshape((num_integrals,4)) + for i in range(num_integrals): + idx[i,1],idx[i,2]=idx[i,2],idx[i,1] + + idx=idx.flatten() + with trexio.File(filename, 'w', back_end=_mode(backend)) as tf: trexio.write_mo_2e_int_eri(tf, 0, num_integrals, idx, eri.ravel()) def read_eri(filename): - '''Read ERIs in AO basis, 8-fold symmetry is assumed''' with trexio.File(filename, 'r', back_end=trexio.TREXIO_AUTO) as tf: nmo = trexio.read_mo_num(tf) nao_pair = nmo * (nmo+1) // 2 eri_size = nao_pair * (nao_pair+1) // 2 idx, data, n_read, eof_flag = trexio.read_mo_2e_int_eri(tf, 0, eri_size) eri = np.zeros(eri_size) - x = idx[:,0]*(idx[:,0]+1)//2 + idx[:,1] - y = idx[:,2]*(idx[:,2]+1)//2 + idx[:,3] + x = idx[:,0]*(idx[:,0]+1)//2 + idx[:,2] + y = idx[:,1]*(idx[:,1]+1)//2 + idx[:,3] eri[x*(x+1)//2+y] = data return eri @@ -410,17 +452,18 @@ def get_occsa_and_occsb(mcscf, norb, nelec, ci_threshold=0.): return occsa_sorted, occsb_sorted, ci_values_sorted, num_determinants -def det_to_trexio(mcscf, norb, nelec, filename, backend='h5', ci_threshold=0., chunk_size=100000): +def det_to_trexio(mcscf, norb, nelec, trexio_file, ci_threshold=0., chunk_size=100000): from trexio_tools.group_tools import determinant as trexio_det - mo_num = norb - int64_num = int((mo_num - 1) / 64) + 1 + ncore = mcscf.ncore + int64_num = trexio.get_int64_num(trexio_file) + occsa, occsb, ci_values, num_determinants = get_occsa_and_occsb(mcscf, norb, nelec, ci_threshold) det_list = [] for a, b, coeff in zip(occsa, occsb, ci_values): - occsa_upshifted = [orb + 1 for orb in a] - occsb_upshifted = [orb + 1 for orb in b] + occsa_upshifted = [orb for orb in range(ncore)] + [orb+ncore for orb in a] + occsb_upshifted = [orb for orb in range(ncore)] + [orb+ncore for orb in b] det_tmp = [] det_tmp += trexio_det.to_determinant_list(occsa_upshifted, int64_num) det_tmp += trexio_det.to_determinant_list(occsb_upshifted, int64_num) @@ -431,24 +474,19 @@ def det_to_trexio(mcscf, norb, nelec, filename, backend='h5', ci_threshold=0., c else: n_chunks = 1 - with trexio.File(filename, 'u', back_end=_mode(backend)) as tf: - if trexio.has_determinant(tf): - trexio.delete_determinant(tf) - trexio.write_mo_num(tf, mo_num) - trexio.write_electron_up_num(tf, len(a)) - trexio.write_electron_dn_num(tf, len(b)) - trexio.write_electron_num(tf, len(a) + len(b)) + if trexio.has_determinant(trexio_file): + trexio.delete_determinant(trexio_file) - offset_file = 0 - for i in range(n_chunks): - start = i * chunk_size - end = min((i + 1) * chunk_size, num_determinants) - current_chunk_size = end - start - - if current_chunk_size > 0: - trexio.write_determinant_list(tf, offset_file, current_chunk_size, det_list[start:end]) - trexio.write_determinant_coefficient(tf, offset_file, current_chunk_size, ci_values[start:end]) - offset_file += current_chunk_size + offset_file = 0 + for i in range(n_chunks): + start = i * chunk_size + end = min((i + 1) * chunk_size, num_determinants) + current_chunk_size = end - start + + if current_chunk_size > 0: + trexio.write_determinant_list(trexio_file, offset_file, current_chunk_size, det_list[start:end]) + trexio.write_determinant_coefficient(trexio_file, offset_file, current_chunk_size, ci_values[start:end]) + offset_file += current_chunk_size def read_det_trexio(filename): with trexio.File(filename, 'r', back_end=trexio.TREXIO_AUTO) as tf: