QuantumPackage/bin/qp_convert_h5_to_ezfio

#!/usr/bin/env python3
"""
convert hdf5 output (e.g. from PySCF) to ezfio

Usage:
    qp_convert_h5_to_ezfio [--noqmc] [--rmg] [-o EZFIO_DIR] FILE

Options:
    -o --output=EZFIO_DIR    Produced directory
                             by default is FILE.ezfio
    --noqmc                  don't include basis, cell, etc. for QMCPACK
    --rmg                      h5 contains cholesky decomposition informatin, these h5 result from RMG and the pyscf AFQMC converter of QMCPACK.

"""
from ezfio import ezfio
import h5py
import sys
import numpy as np
import os
from docopt import docopt
import gzip
#fname = sys.argv[1]
#qph5name = sys.argv[2]

def get_full_path(file_path):
    file_path = os.path.expanduser(file_path)
    file_path = os.path.expandvars(file_path)
#    file_path = os.path.abspath(file_path)
    return file_path


def flatten(l):
    res = []
    for i in l:
        if hasattr(i, "__iter__") and not isinstance(i, str):
            res.extend(flatten(i))
        else:
            res.append(i)
    return res

def save_array_do(ezfioname,subdir,data,chunksize=16384):
    dims = list(reversed(data.shape))
    rank = len(dims)
    flatdata = data.reshape(-1)
    dim_max = 1
    for i in dims:
        dim_max *= i
    with gzip.open(os.path.join(ezfioname,subdir)+'.gz','wb') as f:
        f.write(f'{rank:3d}\n'.encode())
        for d in dims:
            f.write(f'{d:20d} '.encode())
        f.write("\n".encode())
        fmtstring = chunksize*'{:24.15E}\n'
        for i in range(dim_max//chunksize):
            #f.write((chunksize*'{:24.15E}\n').format(*flatdata[i*chunksize:(i+1)*chunksize]).encode())
            #f.write(fmtstring.format(*flatdata[i*chunksize:(i+1)*chunksize]).encode())
            f.write((''.join("%24.15E\n" % xi for xi in flatdata[i*chunksize:(i+1)*chunksize])).encode())
            print(f'{100.*i/(dim_max//chunksize):7.3f}% complete')
        rem = dim_max%chunksize
        if rem:
            f.write((rem*'{:24.15E}\n').format(*flatdata[-rem:]).encode())
    return


def convert_mol(filename,qph5path):
    ezfio.set_file(filename)
    ezfio.set_nuclei_is_complex(False)

    with h5py.File(qph5path,'r') as qph5:
        nucl_num = qph5['nuclei'].attrs['nucl_num']
        ao_num = qph5['ao_basis'].attrs['ao_num']
        mo_num = qph5['mo_basis'].attrs['mo_num']
        elec_alpha_num = qph5['electrons'].attrs['elec_alpha_num']
        elec_beta_num = qph5['electrons'].attrs['elec_beta_num']

    ezfio.set_nuclei_nucl_num(nucl_num)

    ezfio.set_ao_basis_ao_num(ao_num)
    ezfio.set_mo_basis_mo_num(mo_num)
    ezfio.electrons_elec_alpha_num = elec_alpha_num
    ezfio.electrons_elec_beta_num  = elec_beta_num



    ##ao_num = mo_num
    ##Important !
    #import math
    #nelec_per_kpt = num_elec // n_kpts
    #nelec_alpha_per_kpt = int(math.ceil(nelec_per_kpt / 2.))
    #nelec_beta_per_kpt = int(math.floor(nelec_per_kpt / 2.))
    #
    #ezfio.electrons_elec_alpha_num = int(nelec_alpha_per_kpt * n_kpts)
    #ezfio.electrons_elec_beta_num = int(nelec_beta_per_kpt * n_kpts)

    #ezfio.electrons_elec_alpha_num = int(math.ceil(num_elec / 2.))
    #ezfio.electrons_elec_beta_num = int(math.floor(num_elec / 2.))

    #ezfio.set_utils_num_kpts(n_kpts)
    #ezfio.set_integrals_bielec_df_num(n_aux)

    #(old)Important
    #ezfio.set_nuclei_nucl_num(nucl_num)
    #ezfio.set_nuclei_nucl_charge([0.]*nucl_num)
    #ezfio.set_nuclei_nucl_coord( [ [0.], [0.], [0.] ]*nucl_num )
    #ezfio.set_nuclei_nucl_label( ['He'] * nucl_num )


    with h5py.File(qph5path,'r') as qph5:
        nucl_charge=qph5['nuclei/nucl_charge'][()].tolist()
        nucl_coord=qph5['nuclei/nucl_coord'][()].T.tolist()
        nucl_label=qph5['nuclei/nucl_label'][()].tolist()
        nuclear_repulsion = qph5['nuclei'].attrs['nuclear_repulsion']

    ezfio.set_nuclei_nucl_charge(nucl_charge)
    ezfio.set_nuclei_nucl_coord(nucl_coord)
    if isinstance(nucl_label[0],bytes):
        nucl_label = list(map(lambda x:x.decode(),nucl_label))
    ezfio.set_nuclei_nucl_label(nucl_label)

    ezfio.set_nuclei_io_nuclear_repulsion('Read')
    ezfio.set_nuclei_nuclear_repulsion(nuclear_repulsion)


    ##########################################
    #                                        #
    #                Basis                   #
    #                                        #
    ##########################################

    with h5py.File(qph5path,'r') as qph5:
        do_pseudo = qph5['pseudo'].attrs['do_pseudo']
        ezfio.set_pseudo_do_pseudo(do_pseudo)
        if (do_pseudo):
            ezfio.set_pseudo_pseudo_lmax(qph5['pseudo'].attrs['pseudo_lmax'])
            ezfio.set_pseudo_pseudo_klocmax(qph5['pseudo'].attrs['pseudo_klocmax'])
            ezfio.set_pseudo_pseudo_kmax(qph5['pseudo'].attrs['pseudo_kmax'])
            ezfio.set_pseudo_nucl_charge_remove(qph5['pseudo/nucl_charge_remove'][()].tolist())
            ezfio.set_pseudo_pseudo_n_k(qph5['pseudo/pseudo_n_k'][()].tolist())
            ezfio.set_pseudo_pseudo_n_kl(qph5['pseudo/pseudo_n_kl'][()].tolist())
            ezfio.set_pseudo_pseudo_v_k(qph5['pseudo/pseudo_v_k'][()].tolist())
            ezfio.set_pseudo_pseudo_v_kl(qph5['pseudo/pseudo_v_kl'][()].tolist())
            ezfio.set_pseudo_pseudo_dz_k(qph5['pseudo/pseudo_dz_k'][()].tolist())
            ezfio.set_pseudo_pseudo_dz_kl(qph5['pseudo/pseudo_dz_kl'][()].tolist())

    ##########################################
    #                                        #
    #                Basis                   #
    #                                        #
    ##########################################

    with h5py.File(qph5path,'r') as qph5:
        #coeftmp = qph5['ao_basis/ao_coef'][()]
        #expotmp = qph5['ao_basis/ao_expo'][()]
        ezfio.set_ao_basis_ao_basis(qph5['ao_basis'].attrs['ao_basis'])
        ezfio.set_ao_basis_ao_nucl(qph5['ao_basis/ao_nucl'][()].tolist())
        ezfio.set_ao_basis_ao_prim_num(qph5['ao_basis/ao_prim_num'][()].tolist())
        ezfio.set_ao_basis_ao_power(qph5['ao_basis/ao_power'][()].tolist())
        ezfio.set_ao_basis_ao_coef(qph5['ao_basis/ao_coef'][()].tolist())
        ezfio.set_ao_basis_ao_expo(qph5['ao_basis/ao_expo'][()].tolist())


    ##########################################
    #                                        #
    #               MO Coef                  #
    #                                        #
    ##########################################


    with h5py.File(qph5path,'r') as qph5:
        mo_coef = qph5['mo_basis/mo_coef'][()].tolist()
    ezfio.set_mo_basis_mo_coef(mo_coef)
    #maybe fix qp so we don't need this?
    #ezfio.set_mo_basis_mo_coef([[i for i in range(mo_num)] * ao_num])

    return

def convert_kpts_cd(filename,qph5path,qmcpack=True):
    import json

    ezfio.set_file(filename)
    ezfio.set_nuclei_is_complex(True)
    # Dummy atom since AFQMC h5 has no atom information
    nucl_num = 1
    ezfio.set_nuclei_nucl_num(nucl_num)
    ezfio.set_nuclei_nucl_charge([0.]*nucl_num)
    ezfio.set_nuclei_nucl_coord( [ [0.], [0.], [0.] ]*nucl_num )
    ezfio.set_nuclei_nucl_label( ['He'] * nucl_num )
    with h5py.File(qph5path,'r') as qph5:
        kpt_num = qph5['Hamiltonian']['KPoints'].shape[0]
        elec_alpha_num =  qph5['Hamiltonian']['dims'][4]
        elec_beta_num =  qph5['Hamiltonian']['dims'][5]
        orb_num =  qph5['Hamiltonian']['dims'][3]
        try:
          is_ao =  json.loads(qph5['metadata'][()].decode("utf-8").replace("'",'"'))['ortho_ao']
          if is_ao:
            ao_num = orb_num
          elif is_ao ==False:
            mo_num = orb_num
          else:
              raise ValueError('Problem with ortho_ao key in metadata')
        except:
            raise UnicodeDecodeError('metadata not correctly parsed from HDF5 file')

    ezfio.set_nuclei_kpt_num(kpt_num)
    kpt_pair_num = (kpt_num*kpt_num + kpt_num)//2
    ezfio.set_nuclei_kpt_pair_num(kpt_pair_num)
    # don't multiply nuclei by kpt_num
    # work in k-space, not in equivalent supercell
    nucl_num_per_kpt = nucl_num
    ezfio.set_nuclei_nucl_num(nucl_num_per_kpt)
    # these are totals (kpt_num * num_per_kpt)
    # need to change if we want to truncate orbital space within pyscf
    ezfio.set_ao_basis_ao_num(ao_num)
    ezfio.set_mo_basis_mo_num(mo_num)
    ezfio.set_ao_basis_ao_num_per_kpt(ao_num//kpt_num)
    ezfio.set_mo_basis_mo_num_per_kpt(mo_num//kpt_num)
    ezfio.electrons_elec_alpha_num = elec_alpha_num
    ezfio.electrons_elec_beta_num  = elec_beta_num
    ##########################################
    #                                        #
    #                Basis                   #
    #                                        #
    ##########################################
    #TODO
    ##########################################
    #                                        #
    #               MOCoeff                  #
    #                                        #
    ##########################################
    #TODO
    ##########################################
    #                                        #
    #            Integrals Mono              #
    #                                        #
    ##########################################
    """
      with h5py.File(qph5path,'r') as qph5:
        if is_ao:
            kin_ao_reim=
            ovlp_ao_reim=
            ne_ao_reim=

            ezfio.set_ao_one_e_ints_ao_integrals_kinetic_kpts(kin_ao_reim)
            ezfio.set_ao_one_e_ints_ao_integrals_overlap_kpts(ovlp_ao_reim)
            ezfio.set_ao_one_e_ints_ao_integrals_n_e_kpts(ne_ao_reim)

            ezfio.set_ao_one_e_ints_io_ao_integrals_kinetic('Read')
            ezfio.set_ao_one_e_ints_io_ao_integrals_overlap('Read')
            ezfio.set_ao_one_e_ints_io_ao_integrals_n_e('Read')
        else:
            kin_mo_reim=
            ovlp_mo_reim=
            ne_mo_reim=

            ezfio.set_mo_one_e_ints_mo_integrals_kinetic_kpts(kin_mo_reim)
            ezfio.set_mo_one_e_ints_mo_integrals_overlap_kpts(ovlp_mo_reim)
            ezfio.set_mo_one_e_ints_mo_integrals_n_e_kpts(ne_mo_reim)

            ezfio.set_mo_one_e_ints_io_mo_integrals_kinetic('Read')
            ezfio.set_mo_one_e_ints_io_mo_integrals_overlap('Read')
            ezfio.set_mo_one_e_ints_io_mo_integrals_n_e('Read')
      """
    ##########################################
    #                                        #
    #               k-points                 #
    #                                        #
    ##########################################
    #TODO
    with h5py.File(qph5path,'r') as qph5:
        kconserv = qph5['nuclei/kconserv'][()].tolist()
        minusk = qph5['Hamiltonian']['MinusK'][:]+1 
        QKTok2 = qph5['Hamiltonian']['QPTok2'][:]+1
        unique_kpt_num = len(qph5['Hamiltonian']['KPFactorized'])
        unique_k_idx = []
        for i in qph5['Hamiltonian']['KPFactorized'].keys():
            unique_k_idx.append(int(i[1:])+1)
    kpt_sparse_map = np.zeros(kpt_num)
    for i in range(kpt_num):
        if i+1 in uniq_k_idx:
            kpt_sparse_map[i] = i+1 
        else:
            kpt_sparse_map[i] = -minusk[i]
    ezfio.set_nuclei_kconserv(kconserv)
    ezfio.set_nuclei_io_kconserv('Read')
    ezfio.set_nuclei_minusk(minusk)
    ezfio.set_nuclei_qktok2(QKTok2)
    ezfio.set_nuclei_kpt_sparse_map(kpt_sparse_map)
    ezfio.set_nuclei_unique_kpt_num(unique_kpt_num)
    # kpt_sparse_map
    # unique_kpt_num
    # io_kpt_symm
    ##########################################
    #                                        #
    #             Integrals Bi               #
    #                                        #
    ##########################################

    # should this be in ao_basis? ao_two_e_ints?
    with h5py.File(qph5path,'r') as qph5:
        nchol_per_kpt =  qph5['Hamiltonian']['NCholPerKP'][:]
        nchol_per_kpt = nchol_per_kpt[nchol_per_kpt != 0]
        nchol_per_kpt_max = max(nchol_per_kpt)
        ezfio.set_chol_num(nchol_per_kpt)
        ezfio.set_chol_num_max(nchol_per_kpt_max)
        if is_ao:
            ao_num_per_kpt = ao_num//kpt_num
            ezfio.set_io_chol_ao_integrals('Read')
            #ao_chol_two_e_ints = np.zeros((2, ao_num_per_kpt, ao_num_per_kpt, nchol_per_kpt_max, kpt_num, len(nchol_per_kpt)))
            L_list = []
            for i in len(nchol_per_kpt):
                L = qph5['Hamiltonian']['KPFactorized'][f'L{i}'][:]
                L.reshape(kpt_num, ao_num_per_kpt, ao_num_per_kpt, nchol_per_kpt[i], 2)
                L = np.einsum("ijklm->ilkjm", A, B)
                L_list.append(L)

                #(6, 5184, 2)
                """
                for cmplx in range(2):
                    for ao_idx_i in range(ao_num_per_kpt):
                        for ao_idx_j in range(ao_num_per_kpt):
                            for chol_idx in range(nchol_per_kpt[i]):
                                for kpt_idx in range(kpt_num):
                                    ao_chol_two_e_ints[cmplx][ao_idx_i][ao_idx_j][chol_idx][kpt_idx][i] = L[kpt_idx][ao_idx_i][ao_idx_j][chol_idx][cmplx]
                """
            ao_chol_two_e_ints = np.vstack(L_list)
            ao_chol_two_e_ints = ao_chol_two_e_ints.transpose()
            ezfio.set_chol_ao_integrals_complex(ao_chol_two_e_ints)

            #(2,ao_basis.ao_num_per_kpt,ao_basis.ao_num_per_kpt,ao_two_e_ints.chol_num_max,nuclei.kpt_num,nuclei.unique_kpt_num#)
            """
            df_num = qph5['ao_two_e_ints'].attrs['df_num']
            ezfio.set_ao_two_e_ints_df_num(df_num)
            if 'df_ao_integrals' in qph5['ao_two_e_ints'].keys():
                dfao_reim=qph5['ao_two_e_ints/df_ao_integrals'][()].tolist()
                ezfio.set_ao_two_e_ints_df_ao_integrals_complex(dfao_reim)
                ezfio.set_ao_two_e_ints_io_df_ao_integrals('Read')
            """
        else:
            """
            ezfio.set_io_chol_mo_integrals('Read')
            df_num = qph5['ao_two_e_ints'].attrs['df_num']
            ezfio.set_ao_two_e_ints_df_num(df_num)
            dfmo_reim=qph5['mo_two_e_ints/df_mo_integrals'][()].tolist()
            ezfio.set_mo_two_e_ints_df_mo_integrals_complex(dfmo_reim)
            ezfio.set_mo_two_e_ints_io_df_mo_integrals('Read')
            """
            mo_num_per_kpt = ao_num//kpt_num
            ezfio.set_io_chol_mo_integrals('Read')
            #ao_chol_two_e_ints = np.zeros((2, ao_num_per_kpt, ao_num_per_kpt, nchol_per_kpt_max, kpt_num, len(nchol_per_kpt)))
            L_list = []
            for i in len(nchol_per_kpt):
             L = qph5['Hamiltonian']['KPFactorized'][f'L{i}'][:]
             L.reshape(kpt_num, mo_num_per_kpt, mo_num_per_kpt, nchol_per_kpt[i], 2)
             L = np.einsum("ijklm->ilkjm", A, B)
             L_list.append(L)

             #(6, 5184, 2)
             """
             for cmplx in range(2):
                 for ao_idx_i in range(ao_num_per_kpt):
                     for ao_idx_j in range(ao_num_per_kpt):
                         for chol_idx in range(nchol_per_kpt[i]):
                             for kpt_idx in range(kpt_num):
                                 ao_chol_two_e_ints[cmplx][ao_idx_i][ao_idx_j][chol_idx][kpt_idx][i] = L[kpt_idx][ao_idx_i][ao_idx_j][chol_idx][cmplx]
             """
            mo_chol_two_e_ints = np.vstack(L_list)
            mo_chol_two_e_ints = mo_chol_two_e_ints.transpose()
            ezfio.set_chol_mo_integrals_complex(mo_chol_two_e_ints)
            return




def convert_kpts(filename,qph5path,qmcpack=True):
    ezfio.set_file(filename)
    ezfio.set_nuclei_is_complex(True)

    with h5py.File(qph5path,'r') as qph5:
        kpt_num = qph5['nuclei'].attrs['kpt_num']
        nucl_num = qph5['nuclei'].attrs['nucl_num']
        ao_num = qph5['ao_basis'].attrs['ao_num']
        mo_num = qph5['mo_basis'].attrs['mo_num']
        elec_alpha_num = qph5['electrons'].attrs['elec_alpha_num']
        elec_beta_num = qph5['electrons'].attrs['elec_beta_num']

    ezfio.set_nuclei_kpt_num(kpt_num)
    kpt_pair_num = (kpt_num*kpt_num + kpt_num)//2
    ezfio.set_nuclei_kpt_pair_num(kpt_pair_num)

    # don't multiply nuclei by kpt_num
    # work in k-space, not in equivalent supercell
    nucl_num_per_kpt = nucl_num
    ezfio.set_nuclei_nucl_num(nucl_num_per_kpt)

    # these are totals (kpt_num * num_per_kpt)
    # need to change if we want to truncate orbital space within pyscf
    ezfio.set_ao_basis_ao_num(ao_num)
    ezfio.set_mo_basis_mo_num(mo_num)
    ezfio.set_ao_basis_ao_num_per_kpt(ao_num//kpt_num)
    ezfio.set_mo_basis_mo_num_per_kpt(mo_num//kpt_num)
    ezfio.electrons_elec_alpha_num = elec_alpha_num
    ezfio.electrons_elec_beta_num  = elec_beta_num



    ##ao_num = mo_num
    ##Important !
    #import math
    #nelec_per_kpt = num_elec // n_kpts
    #nelec_alpha_per_kpt = int(math.ceil(nelec_per_kpt / 2.))
    #nelec_beta_per_kpt = int(math.floor(nelec_per_kpt / 2.))
    #
    #ezfio.electrons_elec_alpha_num = int(nelec_alpha_per_kpt * n_kpts)
    #ezfio.electrons_elec_beta_num = int(nelec_beta_per_kpt * n_kpts)

    #ezfio.electrons_elec_alpha_num = int(math.ceil(num_elec / 2.))
    #ezfio.electrons_elec_beta_num = int(math.floor(num_elec / 2.))

    #ezfio.set_utils_num_kpts(n_kpts)
    #ezfio.set_integrals_bielec_df_num(n_aux)

    #(old)Important
    #ezfio.set_nuclei_nucl_num(nucl_num)
    #ezfio.set_nuclei_nucl_charge([0.]*nucl_num)
    #ezfio.set_nuclei_nucl_coord( [ [0.], [0.], [0.] ]*nucl_num )
    #ezfio.set_nuclei_nucl_label( ['He'] * nucl_num )


    with h5py.File(qph5path,'r') as qph5:
        nucl_charge=qph5['nuclei/nucl_charge'][()].tolist()
        nucl_coord=qph5['nuclei/nucl_coord'][()].T.tolist()
        nucl_label=qph5['nuclei/nucl_label'][()].tolist()
        nuclear_repulsion = qph5['nuclei'].attrs['nuclear_repulsion']

    ezfio.set_nuclei_nucl_charge(nucl_charge)
    ezfio.set_nuclei_nucl_coord(nucl_coord)
    if isinstance(nucl_label[0],bytes):
        nucl_label = list(map(lambda x:x.decode(),nucl_label))
    ezfio.set_nuclei_nucl_label(nucl_label)

    ezfio.set_nuclei_io_nuclear_repulsion('Read')
    ezfio.set_nuclei_nuclear_repulsion(nuclear_repulsion)


    ##########################################
    #                                        #
    #                Basis(Dummy)            #
    #                                        #
    ##########################################

    with h5py.File(qph5path,'r') as qph5:
        ezfio.set_ao_basis_ao_basis(qph5['ao_basis'].attrs['ao_basis'])
        ezfio.set_ao_basis_ao_nucl(qph5['ao_basis/ao_nucl'][()].tolist())


    #Just need one (can clean this up later)
    ao_prim_num_max  = 5

    d = [ [0] *ao_prim_num_max]*ao_num
    ezfio.set_ao_basis_ao_prim_num([ao_prim_num_max]*ao_num)
    ezfio.set_ao_basis_ao_power(d)
    ezfio.set_ao_basis_ao_coef(d)
    ezfio.set_ao_basis_ao_expo(d)


    ###########################################
    ##                                        #
    ##                Pseudo                  #
    ##                                        #
    ###########################################

    #with h5py.File(qph5path,'r') as qph5:
    #    do_pseudo = qph5['pseudo'].attrs['do_pseudo']
    #    ezfio.set_pseudo_do_pseudo(do_pseudo)
    #    if (do_pseudo):
    #        ezfio.set_pseudo_pseudo_lmax(qph5['pseudo'].attrs['pseudo_lmax'])
    #        ezfio.set_pseudo_pseudo_klocmax(qph5['pseudo'].attrs['pseudo_klocmax'])
    #        ezfio.set_pseudo_pseudo_kmax(qph5['pseudo'].attrs['pseudo_kmax'])
    #        ezfio.set_pseudo_nucl_charge_remove(qph5['pseudo/nucl_charge_remove'][()].tolist())
    #        ezfio.set_pseudo_pseudo_n_k(qph5['pseudo/pseudo_n_k'][()].tolist())
    #        ezfio.set_pseudo_pseudo_n_kl(qph5['pseudo/pseudo_n_kl'][()].tolist())
    #        ezfio.set_pseudo_pseudo_v_k(qph5['pseudo/pseudo_v_k'][()].tolist())
    #        ezfio.set_pseudo_pseudo_v_kl(qph5['pseudo/pseudo_v_kl'][()].tolist())
    #        ezfio.set_pseudo_pseudo_dz_k(qph5['pseudo/pseudo_dz_k'][()].tolist())
    #        ezfio.set_pseudo_pseudo_dz_kl(qph5['pseudo/pseudo_dz_kl'][()].tolist())

    ##########################################
    #                                        #
    #                Basis(Dummy)            #
    #                                        #
    ##########################################

    #with h5py.File(qph5path,'r') as qph5:
    #    coeftmp = qph5['ao_basis/ao_coef'][()]
    #    expotmp = qph5['ao_basis/ao_expo'][()]
    #    ezfio.set_ao_basis_ao_basis(qph5['ao_basis'].attrs['ao_basis'])
    #    ezfio.set_ao_basis_ao_nucl(qph5['ao_basis/ao_nucl'][()].tolist())
    #    ezfio.set_ao_basis_ao_prim_num(qph5['ao_basis/ao_prim_num'][()].tolist())
    #    ezfio.set_ao_basis_ao_power(qph5['ao_basis/ao_power'][()].tolist())
    #    ezfio.set_ao_basis_ao_coef(qph5['ao_basis/ao_coef'][()].tolist())
    #    ezfio.set_ao_basis_ao_expo(qph5['ao_basis/ao_expo'][()].tolist())


    ##########################################
    #                                        #
    #                Basis(QMC)              #
    #                                        #
    ##########################################

    if qmcpack:
        try:
            with h5py.File(qph5path,'r') as qph5:
                ezfio.set_qmcpack_qmc_nshell(qph5['qmcpack'].attrs['qmc_nshell'])
                ezfio.set_qmcpack_qmc_prim_num_max(qph5['qmcpack'].attrs['qmc_prim_num_max'])
                ezfio.set_qmcpack_qmc_nucl(qph5['qmcpack/qmc_nucl'][()].tolist())
                ezfio.set_qmcpack_qmc_prim_num(qph5['qmcpack/qmc_prim_num'][()].tolist())
                ezfio.set_qmcpack_qmc_lbas(qph5['qmcpack/qmc_lbas'][()].tolist())
                ezfio.set_qmcpack_qmc_coef(qph5['qmcpack/qmc_coef'][()].tolist())
                ezfio.set_qmcpack_qmc_expo(qph5['qmcpack/qmc_expo'][()].tolist())

                ezfio.set_qmcpack_qmc_pbc(qph5['qmcpack'].attrs['PBC'])
                ezfio.set_qmcpack_qmc_cart(qph5['qmcpack'].attrs['cart'])
                ezfio.set_qmcpack_qmc_pseudo(qph5['qmcpack'].attrs['Pseudo'])
                ezfio.set_qmcpack_supertwist(qph5['qmcpack/Super_Twist'][()].tolist())
                ezfio.set_qmcpack_latticevectors(qph5['qmcpack/LatticeVectors'][()].tolist())
                ezfio.set_qmcpack_qmc_phase(qph5['qmcpack/qmc_phase'][()].tolist())
                ezfio.set_qmcpack_qmc_mo_energy(qph5['qmcpack/eigenval'][()].tolist())

        except AttributeError as err:
            print("################################################")
            print("# ERROR: problem copying QMCPACK data to ezfio #")
            print("#        make sure qmcpack plugin is built     #")
            print("################################################")
            #print(f"AttributeError: {err}")
            print("to create ezfio without qmcpack data, use 'qp_convert_h5_to_ezfio --noqmc'")
            raise



    ##########################################
    #                                        #
    #               MO Coef                  #
    #                                        #
    ##########################################


    with h5py.File(qph5path,'r') as qph5:
        mo_coef_kpts = qph5['mo_basis/mo_coef_kpts'][()].tolist()
        mo_coef_cplx = qph5['mo_basis/mo_coef_complex'][()].tolist()
    ezfio.set_mo_basis_mo_coef_kpts(mo_coef_kpts)
    ezfio.set_mo_basis_mo_coef_complex(mo_coef_cplx)
    #maybe fix qp so we don't need this?
    #ezfio.set_mo_basis_mo_coef([[i for i in range(mo_num)] * ao_num])


    ##########################################
    #                                        #
    #            Integrals Mono              #
    #                                        #
    ##########################################

    with h5py.File(qph5path,'r') as qph5:
        if 'ao_one_e_ints' in qph5.keys():
            kin_ao_reim=qph5['ao_one_e_ints/ao_integrals_kinetic_kpts'][()].tolist()
            ovlp_ao_reim=qph5['ao_one_e_ints/ao_integrals_overlap_kpts'][()].tolist()
            ne_ao_reim=qph5['ao_one_e_ints/ao_integrals_n_e_kpts'][()].tolist()

            ezfio.set_ao_one_e_ints_ao_integrals_kinetic_kpts(kin_ao_reim)
            ezfio.set_ao_one_e_ints_ao_integrals_overlap_kpts(ovlp_ao_reim)
            ezfio.set_ao_one_e_ints_ao_integrals_n_e_kpts(ne_ao_reim)

            ezfio.set_ao_one_e_ints_io_ao_integrals_kinetic('Read')
            ezfio.set_ao_one_e_ints_io_ao_integrals_overlap('Read')
            ezfio.set_ao_one_e_ints_io_ao_integrals_n_e('Read')


    with h5py.File(qph5path,'r') as qph5:
        if 'mo_one_e_ints' in qph5.keys():
            kin_mo_reim=qph5['mo_one_e_ints/mo_integrals_kinetic_kpts'][()].tolist()
            ovlp_mo_reim=qph5['mo_one_e_ints/mo_integrals_overlap_kpts'][()].tolist()
            ne_mo_reim=qph5['mo_one_e_ints/mo_integrals_n_e_kpts'][()].tolist()

            ezfio.set_mo_one_e_ints_mo_integrals_kinetic_kpts(kin_mo_reim)
            ezfio.set_mo_one_e_ints_mo_integrals_overlap_kpts(ovlp_mo_reim)
            #ezfio.set_mo_one_e_ints_mo_integrals_n_e_complex(ne_mo_reim)
            ezfio.set_mo_one_e_ints_mo_integrals_n_e_kpts(ne_mo_reim)

            ezfio.set_mo_one_e_ints_io_mo_integrals_kinetic('Read')
            ezfio.set_mo_one_e_ints_io_mo_integrals_overlap('Read')
            #ezfio.set_mo_one_e_ints_io_mo_integrals_n_e('Read')
            ezfio.set_mo_one_e_ints_io_mo_integrals_n_e('Read')

    ##########################################
    #                                        #
    #               k-points                 #
    #                                        #
    ##########################################

    with h5py.File(qph5path,'r') as qph5:
        kconserv = qph5['nuclei/kconserv'][()].tolist()

    ezfio.set_nuclei_kconserv(kconserv)
    ezfio.set_nuclei_io_kconserv('Read')

    ##########################################
    #                                        #
    #             Integrals Bi               #
    #                                        #
    ##########################################

    # should this be in ao_basis? ao_two_e_ints?
    with h5py.File(qph5path,'r') as qph5:
        if 'ao_two_e_ints' in qph5.keys():
            df_num = qph5['ao_two_e_ints'].attrs['df_num']
            ezfio.set_ao_two_e_ints_df_num(df_num)
            if 'df_ao_integrals' in qph5['ao_two_e_ints'].keys():
    #            dfao_re0=qph5['ao_two_e_ints/df_ao_integrals_real'][()].transpose((3,2,1,0))
    #            dfao_im0=qph5['ao_two_e_ints/df_ao_integrals_imag'][()].transpose((3,2,1,0))
    #            dfao_cmplx0 = np.stack((dfao_re0,dfao_im0),axis=-1).tolist()
    #            ezfio.set_ao_two_e_ints_df_ao_integrals_complex(dfao_cmplx0)
                #dfao_reim=qph5['ao_two_e_ints/df_ao_integrals'][()].tolist()
                dfao_reim=qph5['ao_two_e_ints/df_ao_integrals'][()]
                save_array_do(filename,'ao_two_e_ints/df_ao_integrals_complex',dfao_reim)
                #ezfio.set_ao_two_e_ints_df_ao_integrals_complex(dfao_reim)
                #dfao_dims = list(reversed(dfao_reim.shape))
                #test_write_df_ao(,5,dfao_dims,dfao_reim.size,dfao_reim.ravel())

                ezfio.set_ao_two_e_ints_io_df_ao_integrals('Read')

        if 'mo_two_e_ints' in qph5.keys():
            df_num = qph5['ao_two_e_ints'].attrs['df_num']
            ezfio.set_ao_two_e_ints_df_num(df_num)
    #        dfmo_re0=qph5['mo_two_e_ints/df_mo_integrals_real'][()].transpose((3,2,1,0))
    #        dfmo_im0=qph5['mo_two_e_ints/df_mo_integrals_imag'][()].transpose((3,2,1,0))
    #        dfmo_cmplx0 = np.stack((dfmo_re0,dfmo_im0),axis=-1).tolist()
    #        ezfio.set_mo_two_e_ints_df_mo_integrals_complex(dfmo_cmplx0)
            dfmo_reim=qph5['mo_two_e_ints/df_mo_integrals'][()].tolist()
            ezfio.set_mo_two_e_ints_df_mo_integrals_complex(dfmo_reim)
            ezfio.set_mo_two_e_ints_io_df_mo_integrals('Read')

    return

def convert_cplx(filename,qph5path):
    ezfio.set_file(filename)
    ezfio.set_nuclei_is_complex(True)

    with h5py.File(qph5path,'r') as qph5:
        kpt_num = qph5['nuclei'].attrs['kpt_num']
        nucl_num = qph5['nuclei'].attrs['nucl_num']
        ao_num = qph5['ao_basis'].attrs['ao_num']
        mo_num = qph5['mo_basis'].attrs['mo_num']
        elec_alpha_num = qph5['electrons'].attrs['elec_alpha_num']
        elec_beta_num = qph5['electrons'].attrs['elec_beta_num']

    ezfio.set_nuclei_kpt_num(kpt_num)
    kpt_pair_num = (kpt_num*kpt_num + kpt_num)//2
    ezfio.set_nuclei_kpt_pair_num(kpt_pair_num)

    # don't multiply nuclei by kpt_num
    # work in k-space, not in equivalent supercell
    nucl_num_per_kpt = nucl_num
    ezfio.set_nuclei_nucl_num(nucl_num_per_kpt)

    # these are totals (kpt_num * num_per_kpt)
    # need to change if we want to truncate orbital space within pyscf
    ezfio.set_ao_basis_ao_num(ao_num)
    ezfio.set_mo_basis_mo_num(mo_num)
    ezfio.electrons_elec_alpha_num = elec_alpha_num
    ezfio.electrons_elec_beta_num  = elec_beta_num



    ##ao_num = mo_num
    ##Important !
    #import math
    #nelec_per_kpt = num_elec // n_kpts
    #nelec_alpha_per_kpt = int(math.ceil(nelec_per_kpt / 2.))
    #nelec_beta_per_kpt = int(math.floor(nelec_per_kpt / 2.))
    #
    #ezfio.electrons_elec_alpha_num = int(nelec_alpha_per_kpt * n_kpts)
    #ezfio.electrons_elec_beta_num = int(nelec_beta_per_kpt * n_kpts)

    #ezfio.electrons_elec_alpha_num = int(math.ceil(num_elec / 2.))
    #ezfio.electrons_elec_beta_num = int(math.floor(num_elec / 2.))

    #ezfio.set_utils_num_kpts(n_kpts)
    #ezfio.set_integrals_bielec_df_num(n_aux)

    #(old)Important
    #ezfio.set_nuclei_nucl_num(nucl_num)
    #ezfio.set_nuclei_nucl_charge([0.]*nucl_num)
    #ezfio.set_nuclei_nucl_coord( [ [0.], [0.], [0.] ]*nucl_num )
    #ezfio.set_nuclei_nucl_label( ['He'] * nucl_num )


    with h5py.File(qph5path,'r') as qph5:
        nucl_charge=qph5['nuclei/nucl_charge'][()].tolist()
        nucl_coord=qph5['nuclei/nucl_coord'][()].T.tolist()
        nucl_label=qph5['nuclei/nucl_label'][()].tolist()
        nuclear_repulsion = qph5['nuclei'].attrs['nuclear_repulsion']

    ezfio.set_nuclei_nucl_charge(nucl_charge)
    ezfio.set_nuclei_nucl_coord(nucl_coord)
    if isinstance(nucl_label[0],bytes):
        nucl_label = list(map(lambda x:x.decode(),nucl_label))
    ezfio.set_nuclei_nucl_label(nucl_label)

    ezfio.set_nuclei_io_nuclear_repulsion('Read')
    ezfio.set_nuclei_nuclear_repulsion(nuclear_repulsion)


    ##########################################
    #                                        #
    #                Basis                   #
    #                                        #
    ##########################################

#    with h5py.File(qph5path,'r') as qph5:
#        ezfio.set_ao_basis_ao_basis(qph5['ao_basis'].attrs['ao_basis'])
#        ezfio.set_ao_basis_ao_nucl(qph5['ao_basis/ao_nucl'][()].tolist())
#
#
#    #Just need one (can clean this up later)
#    ao_prim_num_max  = 5
#
#    d = [ [0] *ao_prim_num_max]*ao_num
#    ezfio.set_ao_basis_ao_prim_num([ao_prim_num_max]*ao_num)
#    ezfio.set_ao_basis_ao_power(d)
#    ezfio.set_ao_basis_ao_coef(d)
#    ezfio.set_ao_basis_ao_expo(d)

    ##########################################
    #                                        #
    #                Basis                   #
    #                                        #
    ##########################################

    with h5py.File(qph5path,'r') as qph5:
        do_pseudo = qph5['pseudo'].attrs['do_pseudo']
        ezfio.set_pseudo_do_pseudo(do_pseudo)
        if (do_pseudo):
            ezfio.set_pseudo_pseudo_lmax(qph5['pseudo'].attrs['pseudo_lmax'])
            ezfio.set_pseudo_pseudo_klocmax(qph5['pseudo'].attrs['pseudo_klocmax'])
            ezfio.set_pseudo_pseudo_kmax(qph5['pseudo'].attrs['pseudo_kmax'])
            ezfio.set_pseudo_nucl_charge_remove(qph5['pseudo/nucl_charge_remove'][()].tolist())
            ezfio.set_pseudo_pseudo_n_k(qph5['pseudo/pseudo_n_k'][()].tolist())
            ezfio.set_pseudo_pseudo_n_kl(qph5['pseudo/pseudo_n_kl'][()].tolist())
            ezfio.set_pseudo_pseudo_v_k(qph5['pseudo/pseudo_v_k'][()].tolist())
            ezfio.set_pseudo_pseudo_v_kl(qph5['pseudo/pseudo_v_kl'][()].tolist())
            ezfio.set_pseudo_pseudo_dz_k(qph5['pseudo/pseudo_dz_k'][()].tolist())
            ezfio.set_pseudo_pseudo_dz_kl(qph5['pseudo/pseudo_dz_kl'][()].tolist())

   ##########################################
   #                                        #
   #                Basis                   #
   #                                        #
   ##########################################

    with h5py.File(qph5path,'r') as qph5:
        coeftmp = qph5['ao_basis/ao_coef'][()]
        expotmp = qph5['ao_basis/ao_expo'][()]
        ezfio.set_ao_basis_ao_basis(qph5['ao_basis'].attrs['ao_basis'])
        ezfio.set_ao_basis_ao_nucl(qph5['ao_basis/ao_nucl'][()].tolist())
        ezfio.set_ao_basis_ao_prim_num(qph5['ao_basis/ao_prim_num'][()].tolist())
        ezfio.set_ao_basis_ao_power(qph5['ao_basis/ao_power'][()].tolist())
        ezfio.set_ao_basis_ao_coef(qph5['ao_basis/ao_coef'][()].tolist())
        ezfio.set_ao_basis_ao_expo(qph5['ao_basis/ao_expo'][()].tolist())




    ##########################################
    #                                        #
    #               MO Coef                  #
    #                                        #
    ##########################################


    with h5py.File(qph5path,'r') as qph5:
        mo_coef_reim = qph5['mo_basis/mo_coef_complex'][()].tolist()
    ezfio.set_mo_basis_mo_coef_complex(mo_coef_reim)
    #maybe fix qp so we don't need this?
    #ezfio.set_mo_basis_mo_coef([[i for i in range(mo_num)] * ao_num])


    ##########################################
    #                                        #
    #            Integrals Mono              #
    #                                        #
    ##########################################

    with h5py.File(qph5path,'r') as qph5:
        if 'ao_one_e_ints' in qph5.keys():
            kin_ao_reim=qph5['ao_one_e_ints/ao_integrals_kinetic'][()].tolist()
            ovlp_ao_reim=qph5['ao_one_e_ints/ao_integrals_overlap'][()].tolist()
            ne_ao_reim=qph5['ao_one_e_ints/ao_integrals_n_e'][()].tolist()

            ezfio.set_ao_one_e_ints_ao_integrals_kinetic_complex(kin_ao_reim)
            ezfio.set_ao_one_e_ints_ao_integrals_overlap_complex(ovlp_ao_reim)
            ezfio.set_ao_one_e_ints_ao_integrals_n_e_complex(ne_ao_reim)

            ezfio.set_ao_one_e_ints_io_ao_integrals_kinetic('Read')
            ezfio.set_ao_one_e_ints_io_ao_integrals_overlap('Read')
            ezfio.set_ao_one_e_ints_io_ao_integrals_n_e('Read')


    with h5py.File(qph5path,'r') as qph5:
        if 'mo_one_e_ints' in qph5.keys():
            kin_mo_reim=qph5['mo_one_e_ints/mo_integrals_kinetic'][()].tolist()
            #ovlp_mo_reim=qph5['mo_one_e_ints/mo_integrals_overlap'][()].tolist()
            ne_mo_reim=qph5['mo_one_e_ints/mo_integrals_n_e'][()].tolist()

            ezfio.set_mo_one_e_ints_mo_integrals_kinetic_complex(kin_mo_reim)
            #ezfio.set_mo_one_e_ints_mo_integrals_overlap_complex(ovlp_mo_reim)
            #ezfio.set_mo_one_e_ints_mo_integrals_n_e_complex(ne_mo_reim)
            ezfio.set_mo_one_e_ints_mo_integrals_n_e_complex(ne_mo_reim)

            ezfio.set_mo_one_e_ints_io_mo_integrals_kinetic('Read')
            #ezfio.set_mo_one_e_ints_io_mo_integrals_overlap('Read')
            #ezfio.set_mo_one_e_ints_io_mo_integrals_n_e('Read')
            ezfio.set_mo_one_e_ints_io_mo_integrals_n_e('Read')

    ##########################################
    #                                        #
    #               k-points                 #
    #                                        #
    ##########################################

    with h5py.File(qph5path,'r') as qph5:
        kconserv = qph5['nuclei/kconserv'][()].tolist()

    ezfio.set_nuclei_kconserv(kconserv)
    ezfio.set_nuclei_io_kconserv('Read')


    ##########################################
    #                                        #
    #             Integrals Bi               #
    #                                        #
    ##########################################

    # should this be in ao_basis? ao_two_e_ints?
    with h5py.File(qph5path,'r') as qph5:
        if 'ao_two_e_ints' in qph5.keys():
            df_num = qph5['ao_two_e_ints'].attrs['df_num']
            ezfio.set_ao_two_e_ints_df_num(df_num)
            if 'df_ao_integrals' in qph5['ao_two_e_ints'].keys():
    #            dfao_re0=qph5['ao_two_e_ints/df_ao_integrals_real'][()].transpose((3,2,1,0))
    #            dfao_im0=qph5['ao_two_e_ints/df_ao_integrals_imag'][()].transpose((3,2,1,0))
    #            dfao_cmplx0 = np.stack((dfao_re0,dfao_im0),axis=-1).tolist()
    #            ezfio.set_ao_two_e_ints_df_ao_integrals_complex(dfao_cmplx0)
                dfao_reim=qph5['ao_two_e_ints/df_ao_integrals'][()].tolist()
                ezfio.set_ao_two_e_ints_df_ao_integrals_complex(dfao_reim)
                ezfio.set_ao_two_e_ints_io_df_ao_integrals('Read')

        if 'mo_two_e_ints' in qph5.keys():
            df_num = qph5['ao_two_e_ints'].attrs['df_num']
            ezfio.set_ao_two_e_ints_df_num(df_num)
    #        dfmo_re0=qph5['mo_two_e_ints/df_mo_integrals_real'][()].transpose((3,2,1,0))
    #        dfmo_im0=qph5['mo_two_e_ints/df_mo_integrals_imag'][()].transpose((3,2,1,0))
    #        dfmo_cmplx0 = np.stack((dfmo_re0,dfmo_im0),axis=-1).tolist()
    #        ezfio.set_mo_two_e_ints_df_mo_integrals_complex(dfmo_cmplx0)
            dfmo_reim=qph5['mo_two_e_ints/df_mo_integrals'][()].tolist()
            ezfio.set_mo_two_e_ints_df_mo_integrals_complex(dfmo_reim)
            ezfio.set_mo_two_e_ints_io_df_mo_integrals('Read')

    return


if __name__ == '__main__':
    ARGUMENTS = docopt(__doc__)

    FILE = get_full_path(ARGUMENTS['FILE'])
    qmcpack = True
    rmg = False
    if ARGUMENTS["--output"]:
        EZFIO_FILE = get_full_path(ARGUMENTS["--output"])
    else:
        EZFIO_FILE = "{0}.ezfio".format(FILE)

    if ARGUMENTS["--noqmc"]:
        qmcpack = False
    if ARGUMENTS["--rmg"]:
        rmg = True
    with h5py.File(FILE,'r') as qph5:
        try:
            do_kpts = ('kconserv' in qph5['nuclei'].keys())
        except:
            do_kpts = False
    if (do_kpts or rmg):
        print("converting HDF5 to EZFIO for periodic system")
        if rmg:
            print("Using RMG and AFQMC h5")
            convert_kpts_cd(EZFIO_FILE,FILE,qmcpack)
        else:
            convert_kpts(EZFIO_FILE,FILE,qmcpack)
    else:
        print("converting HDF5 to EZFIO for molecular system")
        convert_mol(EZFIO_FILE,FILE)

#    sys.stdout.flush()
#    if os.system("qp_run save_ortho_mos "+EZFIO_FILE) != 0:
#        print("""Warning: You need to run
#
#         qp run save_ortho_mos
#
#to be sure your MOs will be orthogonal, which is not the case when
#the MOs are read from output files (not enough precision in output).""")