9
1
mirror of https://github.com/QuantumPackage/qp2.git synced 2024-12-22 19:43:32 +01:00
qp2/scripts/qp_import_trexio.py

560 lines
18 KiB
Python
Executable File

#!/usr/bin/env python3
"""
convert TREXIO file to EZFIO
Usage:
qp_import_trexio [-o EZFIO_DIR] FILE
Options:
-o --output=EZFIO_DIR Produced directory
by default is FILE.ezfio
"""
import sys
import os
import numpy as np
import subprocess
import tempfile
from functools import reduce
from ezfio import ezfio
from docopt import docopt
import qp_bitmasks
try:
import trexio
except ImportError:
print("Error: trexio python module is not found. Try python3 -m pip install trexio")
sys.exit(1)
try:
QP_ROOT = os.environ["QP_ROOT"]
QP_EZFIO = os.environ["QP_EZFIO"]
except KeyError:
print("Error: QP_ROOT environment variable not found.")
sys.exit(1)
else:
sys.path = [QP_EZFIO + "/Python",
QP_ROOT + "/install/resultsFile",
QP_ROOT + "/install",
QP_ROOT + "/scripts"] + sys.path
def uint64_to_int64(u):
# Check if the most significant bit is set
if u & (1 << 63):
# Calculate the two's complement
result = -int(np.bitwise_not(np.uint64(u))+1)
else:
# The number is already positive
result = u
return result
def generate_xyz(l):
def create_z(x,y,z):
return (x, y, l-(x+y))
def create_y(accu,x,y,z):
if y == 0:
result = [create_z(x,y,z)] + accu
else:
result = create_y([create_z(x,y,z)] + accu , x, y-1, z)
return result
def create_x(accu,x,y,z):
if x == 0:
result = create_y([], x,y,z) + accu
else:
xnew = x-1
ynew = l-xnew
result = create_x(create_y([],x,y,z) + accu , xnew, ynew, z)
return result
result = create_x([], l, 0, 0)
result.reverse()
return result
def write_ezfio(trexio_filename, filename):
warnings = []
trexio_file = trexio.File(trexio_filename,mode='r',back_end=trexio.TREXIO_AUTO)
ezfio.set_file(filename)
ezfio.set_trexio_trexio_file(trexio_filename)
print("Nuclei\t\t...\t", end=' ')
charge = [0.]
if trexio.has_nucleus(trexio_file):
charge = trexio.read_nucleus_charge(trexio_file)
ezfio.set_nuclei_nucl_num(len(charge))
ezfio.set_nuclei_nucl_charge(charge)
coord = trexio.read_nucleus_coord(trexio_file)
coord = np.transpose(coord)
ezfio.set_nuclei_nucl_coord(coord)
label = trexio.read_nucleus_label(trexio_file)
nucl_num = trexio.read_nucleus_num(trexio_file)
# Transformt H1 into H
import re
p = re.compile(r'(\d*)$')
label = [p.sub("", x).capitalize() for x in label]
ezfio.set_nuclei_nucl_label(label)
print("OK")
else:
ezfio.set_nuclei_nucl_num(1)
ezfio.set_nuclei_nucl_charge([0.])
ezfio.set_nuclei_nucl_coord([0.,0.,0.])
ezfio.set_nuclei_nucl_label(["X"])
print("None")
warnings.append("No geometry found in the TREXIO file")
print("Electrons\t...\t", end=' ')
try:
num_beta = trexio.read_electron_dn_num(trexio_file)
except:
num_beta = int(sum(charge))//2
try:
num_alpha = trexio.read_electron_up_num(trexio_file)
except:
num_alpha = int(sum(charge)) - num_beta
if num_alpha == 0:
print("\n\nError: There are zero electrons in the TREXIO file.\n\n")
sys.exit(1)
ezfio.set_electrons_elec_alpha_num(num_alpha)
ezfio.set_electrons_elec_beta_num(num_beta)
print(f"{num_alpha} {num_beta}")
print("Basis\t\t...\t", end=' ')
shell_num = 0
try:
basis_type = trexio.read_basis_type(trexio_file)
if basis_type.lower() in ["gaussian", "slater"]:
shell_num = trexio.read_basis_shell_num(trexio_file)
prim_num = trexio.read_basis_prim_num(trexio_file)
ang_mom = trexio.read_basis_shell_ang_mom(trexio_file)
nucl_index = trexio.read_basis_nucleus_index(trexio_file)
exponent = trexio.read_basis_exponent(trexio_file)
coefficient = trexio.read_basis_coefficient(trexio_file)
shell_index = trexio.read_basis_shell_index(trexio_file)
ao_shell = trexio.read_ao_shell(trexio_file)
ezfio.set_basis_basis("Read from TREXIO")
ezfio.set_ao_basis_ao_basis("Read from TREXIO")
ezfio.set_basis_shell_num(shell_num)
ezfio.set_basis_prim_num(prim_num)
ezfio.set_basis_shell_ang_mom(ang_mom)
ezfio.set_basis_basis_nucleus_index([ x+1 for x in nucl_index ])
ezfio.set_basis_prim_expo(exponent)
nucl_shell_num = []
prev = None
m = 0
for i in ao_shell:
if i != prev:
m += 1
if prev is None or nucl_index[i] != nucl_index[prev]:
nucl_shell_num.append(m)
m = 0
prev = i
assert (len(nucl_shell_num) == nucl_num)
shell_prim_num = []
prev = shell_index[0]
count = 0
for i in shell_index:
if i != prev:
shell_prim_num.append(count)
count = 0
count += 1
prev = i
shell_prim_num.append(count)
assert (len(shell_prim_num) == shell_num)
ezfio.set_basis_shell_prim_num(shell_prim_num)
ezfio.set_basis_shell_index([x+1 for x in shell_index])
ezfio.set_basis_nucleus_shell_num(nucl_shell_num)
shell_factor = trexio.read_basis_shell_factor(trexio_file)
prim_factor = trexio.read_basis_prim_factor(trexio_file)
for i,p in enumerate(prim_factor):
coefficient[i] *= prim_factor[i]
ezfio.set_ao_basis_primitives_normalized(False)
ezfio.set_basis_prim_coef(coefficient)
elif basis_type.lower() == "numerical":
shell_num = trexio.read_basis_shell_num(trexio_file)
prim_num = shell_num
ang_mom = trexio.read_basis_shell_ang_mom(trexio_file)
nucl_index = trexio.read_basis_nucleus_index(trexio_file)
exponent = [1.]*prim_num
coefficient = [1.]*prim_num
shell_index = [i for i in range(shell_num)]
ao_shell = trexio.read_ao_shell(trexio_file)
ezfio.set_basis_basis("None")
ezfio.set_ao_basis_ao_basis("None")
ezfio.set_basis_shell_num(shell_num)
ezfio.set_basis_prim_num(prim_num)
ezfio.set_basis_shell_ang_mom(ang_mom)
ezfio.set_basis_basis_nucleus_index([ x+1 for x in nucl_index ])
ezfio.set_basis_prim_expo(exponent)
ezfio.set_basis_prim_coef(coefficient)
nucl_shell_num = []
prev = None
m = 0
for i in ao_shell:
if i != prev:
m += 1
if prev is None or nucl_index[i] != nucl_index[prev]:
nucl_shell_num.append(m)
m = 0
prev = i
assert (len(nucl_shell_num) == nucl_num)
shell_prim_num = []
prev = shell_index[0]
count = 0
for i in shell_index:
if i != prev:
shell_prim_num.append(count)
count = 0
count += 1
prev = i
shell_prim_num.append(count)
assert (len(shell_prim_num) == shell_num)
ezfio.set_basis_shell_prim_num(shell_prim_num)
ezfio.set_basis_shell_index([x+1 for x in shell_index])
ezfio.set_basis_nucleus_shell_num(nucl_shell_num)
shell_factor = trexio.read_basis_shell_factor(trexio_file)
else:
raise TypeError
print(basis_type)
except:
basis_type = "None"
print("None")
ezfio.set_ao_basis_ao_cartesian(True)
print("AOS\t\t...\t", end=' ')
try:
cartesian = trexio.read_ao_cartesian(trexio_file)
except:
cartesian = True
trexio_file_cart = trexio_file
if basis_type.lower() == "gaussian" and not cartesian:
try:
import trexio_tools
fd, tmp = tempfile.mkstemp()
os.close(fd)
retcode = subprocess.call(["trexio", "convert-to", "-t", "cartesian", "-o", tmp, trexio_filename])
trexio_file_cart = trexio.File(tmp,mode='r',back_end=trexio.TREXIO_AUTO)
cartesian = trexio.read_ao_cartesian(trexio_file_cart)
os.unlink(tmp)
except:
pass
ao_num = trexio.read_ao_num(trexio_file_cart)
ezfio.set_ao_basis_ao_num(ao_num)
if cartesian and basis_type.lower() == "gaussian" and shell_num > 0:
ao_shell = trexio.read_ao_shell(trexio_file_cart)
at = [ nucl_index[i]+1 for i in ao_shell ]
ezfio.set_ao_basis_ao_nucl(at)
num_prim0 = [ 0 for i in range(shell_num) ]
for i in shell_index:
num_prim0[i] += 1
coef = {}
expo = {}
for i,c in enumerate(coefficient):
idx = shell_index[i]
if idx in coef:
coef[idx].append(c)
expo[idx].append(exponent[i])
else:
coef[idx] = [c]
expo[idx] = [exponent[i]]
coefficient = []
exponent = []
power_x = []
power_y = []
power_z = []
num_prim = []
for i in range(shell_num):
for x,y,z in generate_xyz(ang_mom[i]):
power_x.append(x)
power_y.append(y)
power_z.append(z)
coefficient.append(coef[i])
exponent.append(expo[i])
num_prim.append(num_prim0[i])
print (len(coefficient), ao_num)
assert (len(coefficient) == ao_num)
ezfio.set_ao_basis_ao_power(power_x + power_y + power_z)
ezfio.set_ao_basis_ao_prim_num(num_prim)
prim_num_max = max( [ len(x) for x in coefficient ] )
for i in range(ao_num):
coefficient[i] += [0. for j in range(len(coefficient[i]), prim_num_max)]
exponent [i] += [0. for j in range(len(exponent[i]), prim_num_max)]
coefficient = reduce(lambda x, y: x + y, coefficient, [])
exponent = reduce(lambda x, y: x + y, exponent , [])
coef = []
expo = []
for i in range(prim_num_max):
for j in range(i, len(coefficient), prim_num_max):
coef.append(coefficient[j])
expo.append(exponent[j])
# ezfio.set_ao_basis_ao_prim_num_max(prim_num_max)
ezfio.set_ao_basis_ao_coef(coef)
ezfio.set_ao_basis_ao_expo(expo)
print("OK")
else:
if basis_type.lower() == "gaussian" and not cartesian:
warnings.append(f"Spherical AOs not handled by QP. Convert the TREXIO file using trexio_tools:\n trexio convert-to -t cartesian -o cartesian_{trexio_filename} {trexio_filename}")
warnings.append("Integrals should be imported using:\n qp run import_trexio_integrals")
print("None")
# _
# |\/| _ _ |_) _. _ o _
# | | (_) _> |_) (_| _> | _>
#
print("MOS\t\t...\t", end=' ')
labels = { "Canonical" : "Canonical",
"RHF" : "Canonical",
"BOYS" : "Localized",
"ROHF" : "Canonical",
"UHF" : "Canonical",
"Natural": "Natural" }
try:
label = labels[trexio.read_mo_type(trexio_file)]
except:
label = "None"
ezfio.set_mo_basis_mo_label(label)
ezfio.set_determinants_mo_label(label)
try:
clss = trexio.read_mo_class(trexio_file)
core = [ i for i in clss if i.lower() == "core" ]
inactive = [ i for i in clss if i.lower() == "inactive" ]
active = [ i for i in clss if i.lower() == "active" ]
virtual = [ i for i in clss if i.lower() == "virtual" ]
deleted = [ i for i in clss if i.lower() == "deleted" ]
except trexio.Error:
pass
try:
mo_num = trexio.read_mo_num(trexio_file)
ezfio.set_mo_basis_mo_num(mo_num)
# Read coefs from temporary cartesian file created in the AO section
MoMatrix = trexio.read_mo_coefficient(trexio_file_cart)
# Renormalize MO coefs if needed
if trexio.has_ao_normalization(trexio_file_cart):
ezfio.set_ao_basis_ao_normalized(False)
norm = trexio.read_ao_normalization(trexio_file_cart)
# for j in range(mo_num):
# for i,f in enumerate(norm):
# MoMatrix[i,j] *= f
ezfio.set_mo_basis_mo_coef(MoMatrix)
mo_occ = [ 0. for i in range(mo_num) ]
for i in range(num_alpha):
mo_occ[i] += 1.
for i in range(num_beta):
mo_occ[i] += 1.
ezfio.set_mo_basis_mo_occ(mo_occ)
print("OK")
except:
print("None")
print("Pseudos\t\t...\t", end=' ')
ezfio.set_pseudo_do_pseudo(False)
if trexio.has_ecp_ang_mom(trexio_file):
ezfio.set_pseudo_do_pseudo(True)
max_ang_mom_plus_1 = trexio.read_ecp_max_ang_mom_plus_1(trexio_file)
z_core = trexio.read_ecp_z_core(trexio_file)
ang_mom = trexio.read_ecp_ang_mom(trexio_file)
nucleus_index = trexio.read_ecp_nucleus_index(trexio_file)
exponent = trexio.read_ecp_exponent(trexio_file)
coefficient = trexio.read_ecp_coefficient(trexio_file)
power = trexio.read_ecp_power(trexio_file)
lmax = max( max_ang_mom_plus_1 ) - 1
ezfio.set_pseudo_pseudo_lmax(lmax)
ezfio.set_pseudo_nucl_charge_remove(z_core)
prev_center = None
ecp = {}
for i in range(len(ang_mom)):
center = nucleus_index[i]
if center != prev_center:
ecp[center] = { "lmax": max_ang_mom_plus_1[center],
"zcore": z_core[center],
"contr": {} }
for j in range(max_ang_mom_plus_1[center]+1):
ecp[center]["contr"][j] = []
ecp[center]["contr"][ang_mom[i]].append( (coefficient[i], power[i], exponent[i]) )
prev_center = center
ecp_loc = {}
ecp_nl = {}
kmax = 0
klocmax = 0
for center in ecp:
ecp_nl [center] = {}
for k in ecp[center]["contr"]:
if k == ecp[center]["lmax"]:
ecp_loc[center] = ecp[center]["contr"][k]
klocmax = max(len(ecp_loc[center]), klocmax)
else:
ecp_nl [center][k] = ecp[center]["contr"][k]
kmax = max(len(ecp_nl [center][k]), kmax)
ezfio.set_pseudo_pseudo_klocmax(klocmax)
ezfio.set_pseudo_pseudo_kmax(kmax)
pseudo_n_k = [[0 for _ in range(nucl_num)] for _ in range(klocmax)]
pseudo_v_k = [[0. for _ in range(nucl_num)] for _ in range(klocmax)]
pseudo_dz_k = [[0. for _ in range(nucl_num)] for _ in range(klocmax)]
pseudo_n_kl = [[[0 for _ in range(nucl_num)] for _ in range(kmax)] for _ in range(lmax+1)]
pseudo_v_kl = [[[0. for _ in range(nucl_num)] for _ in range(kmax)] for _ in range(lmax+1)]
pseudo_dz_kl = [[[0. for _ in range(nucl_num)] for _ in range(kmax)] for _ in range(lmax+1)]
for center in ecp_loc:
for k in range( len(ecp_loc[center]) ):
v, n, dz = ecp_loc[center][k]
pseudo_n_k[k][center] = n
pseudo_v_k[k][center] = v
pseudo_dz_k[k][center] = dz
ezfio.set_pseudo_pseudo_n_k(pseudo_n_k)
ezfio.set_pseudo_pseudo_v_k(pseudo_v_k)
ezfio.set_pseudo_pseudo_dz_k(pseudo_dz_k)
for center in ecp_nl:
for l in range( len(ecp_nl[center]) ):
for k in range( len(ecp_nl[center][l]) ):
v, n, dz = ecp_nl[center][l][k]
pseudo_n_kl[l][k][center] = n
pseudo_v_kl[l][k][center] = v
pseudo_dz_kl[l][k][center] = dz
ezfio.set_pseudo_pseudo_n_kl(pseudo_n_kl)
ezfio.set_pseudo_pseudo_v_kl(pseudo_v_kl)
ezfio.set_pseudo_pseudo_dz_kl(pseudo_dz_kl)
print("OK")
else:
print("None")
print("Determinant\t...\t", end=' ')
alpha = [ i for i in range(num_alpha) ]
beta = [ i for i in range(num_beta) ]
if trexio.has_mo_spin(trexio_file):
spin = trexio.read_mo_spin(trexio_file)
if max(spin) == 1:
tmp = [ i for i in range(len(spin)) if spin[i] == 0 ]
alpha = [ tmp[i] for i in range(num_alpha) ]
tmp = [ i for i in range(len(spin)) if spin[i] == 1 ]
beta = [ tmp[i] for i in range(num_beta) ]
warnings.append("UHF orbitals orbitals read", end=' ')
alpha_s = ['0']*mo_num
beta_s = ['0']*mo_num
for i in alpha:
alpha_s[i] = '1'
for i in beta:
beta_s[i] = '1'
alpha_s = ''.join(alpha_s)[::-1]
beta_s = ''.join(beta_s)[::-1]
def conv(i):
try:
result = np.int64(i)
except:
result = np.int64(i-2**63-1)
return result
alpha = [ uint64_to_int64(int(i,2)) for i in qp_bitmasks.string_to_bitmask(alpha_s) ][::-1]
beta = [ uint64_to_int64(int(i,2)) for i in qp_bitmasks.string_to_bitmask(beta_s ) ][::-1]
ezfio.set_determinants_bit_kind(8)
ezfio.set_determinants_n_int(1+mo_num//64)
ezfio.set_determinants_n_det(1)
ezfio.set_determinants_n_states(1)
ezfio.set_determinants_psi_det(alpha+beta)
ezfio.set_determinants_psi_coef([[1.0]])
print("OK")
for w in warnings:
s = "-------------"
print (s)
print (w)
def get_full_path(file_path):
file_path = os.path.expanduser(file_path)
file_path = os.path.expandvars(file_path)
return file_path
if __name__ == '__main__':
ARGUMENTS = docopt(__doc__)
FILE = get_full_path(ARGUMENTS['FILE'])
trexio_filename = FILE
if ARGUMENTS["--output"]:
EZFIO_FILE = get_full_path(ARGUMENTS["--output"])
else:
EZFIO_FILE = "{0}.ezfio".format(FILE)
write_ezfio(trexio_filename, EZFIO_FILE)
sys.stdout.flush()