1
0
mirror of https://gitlab.com/scemama/qp_plugins_scemama.git synced 2024-11-07 22:53:42 +01:00
qp_plugins_scemama/devel/trexio/qp_import_trexio.py

377 lines
11 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 trexio
import numpy as np
from functools import reduce
from ezfio import ezfio
from docopt import docopt
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 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):
try:
trexio_file = trexio.File(trexio_filename,mode='r',back_end=trexio.TREXIO_TEXT)
except:
trexio_file = trexio.File(trexio_filename,mode='r',back_end=trexio.TREXIO_HDF5)
basis_type = trexio.read_basis_type(trexio_file)
if basis_type.lower() != "gaussian":
raise TypeError
ezfio.set_file(filename)
print("Electrons\t...\t", end=' ')
num_alpha = trexio.read_electron_up_num(trexio_file)
num_beta = trexio.read_electron_dn_num(trexio_file)
ezfio.set_electrons_elec_alpha_num(num_alpha)
ezfio.set_electrons_elec_beta_num(num_beta)
print("OK")
print("Nuclei\t\t...\t", end=' ')
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")
print("Basis\t\t...\t", end=' ')
ezfio.set_basis_basis("Read from TREXIO")
shell_num = trexio.read_basis_shell_num(trexio_file)
prim_num = trexio.read_basis_prim_num(trexio_file)
ezfio.set_basis_shell_num(shell_num)
ezfio.set_basis_prim_num(prim_num)
ang_mom = trexio.read_basis_shell_ang_mom(trexio_file)
nucl_index = trexio.read_basis_nucleus_index(trexio_file)
ezfio.set_basis_shell_ang_mom(ang_mom)
ezfio.set_basis_basis_nucleus_index([ x+1 for x in nucl_index ])
exponent = trexio.read_basis_exponent(trexio_file)
coefficient = trexio.read_basis_coefficient(trexio_file)
ezfio.set_basis_prim_expo(exponent)
ezfio.set_basis_prim_coef(coefficient)
shell_index = trexio.read_basis_shell_index(trexio_file)
ao_shell = trexio.read_ao_shell(trexio_file)
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)
print("OK")
print("AOS\t\t...\t", end=' ')
cartesian = trexio.read_ao_cartesian(trexio_file)
if not cartesian:
raise TypeError('Only cartesian TREXIO files can be converted')
ao_num = trexio.read_ao_num(trexio_file)
ezfio.set_ao_basis_ao_num(ao_num)
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])
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_coef(coef)
ezfio.set_ao_basis_ao_expo(expo)
ezfio.set_ao_basis_ao_basis("Read from TREXIO")
print("OK")
# _
# |\/| _ _ |_) _. _ o _
# | | (_) _> |_) (_| _> | _>
#
print("MOS\t\t...\t", end=' ')
label = trexio.read_mo_type(trexio_file)
ezfio.set_mo_basis_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
MoMatrix = trexio.read_mo_coefficient(trexio_file)
mo_num = trexio.read_mo_num(trexio_file)
ezfio.set_mo_basis_mo_num(mo_num)
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")
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")
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__)
trexio_filename = get_full_path(ARGUMENTS['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()