10
1
mirror of https://gitlab.com/scemama/qmcchem.git synced 2024-06-23 13:42:04 +02:00
qmcchem/bin/jast_opt.py

228 lines
6.6 KiB
Python
Executable File

#!/usr/bin/env python3
import scipy as sp
import scipy.optimize
import numpy as np
import sys
import os
import time
import subprocess
from math import sqrt
QMCCHEM_PATH=os.environ["QMCCHEM_PATH"]
sys.path.insert(0,QMCCHEM_PATH+"/EZFIO/Python/")
from ezfio import ezfio
# PARAMETERS
thresh = 1.e-2
block_time = 6
def main():
if len(sys.argv) != 2:
print("Usage: %s <EZFIO_DIRECTORY>"%sys.argv[0])
sys.exit(1)
filename = sys.argv[1]
ezfio.set_file(filename)
jast_type = ezfio.jastrow_jast_type
print (jast_type)
def make_atom_map():
labels = {}
dimension = 0
for i,k in enumerate(ezfio.nuclei_nucl_label):
if k in labels:
labels[k].append(i)
else:
labels[k] = [dimension, i]
dimension += 1
atom_map = [[] for i in range(dimension)]
for atom in labels.keys():
l = labels[atom]
atom_map[l[0]] = l[1:]
return atom_map
atom_map = make_atom_map()
if jast_type == 'Simple':
def get_params_b():
return ezfio.jastrow_jast_b_up_dn
def set_params_b(x):
x = np.abs(x)
ezfio.set_jastrow_jast_b_up_up(x)
ezfio.set_jastrow_jast_b_up_dn(x)
elif jast_type == 'Mu':
def get_params_b():
return ezfio.jastrow_mu_erf
def set_params_b(x):
x = np.abs(x)
ezfio.set_jastrow_mu_erf(x)
if jast_type == 'Simple':
def get_params_pen():
d = ezfio.jastrow_jast_pen
return np.array([d[m[0]] for m in atom_map])
def set_params_pen(x):
x = np.abs(x)
y=list(ezfio.jastrow_jast_pen)
for i,m in enumerate(atom_map):
for j in m:
y[j] = x[i]
ezfio.set_jastrow_jast_pen(y)
elif jast_type == 'Mu':
def get_params_pen():
d = ezfio.jastrow_jast_1bgauss_pen
return np.array([d[m[0]] for m in atom_map])
def set_params_pen(x):
x = np.abs(x)
y=list(ezfio.jastrow_jast_1bgauss_pen)
for i,m in enumerate(atom_map):
for j in m:
y[j] = x[i]
ezfio.set_jastrow_jast_1bgauss_pen(y)
def get_norm():
return 1.,0.
# buffer = subprocess.check_output(['qmcchem', 'result', '-e', 'psi_norm', filename],
# encoding='UTF-8')
# if buffer.strip() != "":
# buffer = buffer.splitlines()[-1]
# _, energy, error = [float(x) for x in buffer.split()]
# else:
# return None, None
def get_energy():
buffer = subprocess.check_output(['qmcchem', 'result', '-e', 'e_loc', filename],
encoding='UTF-8')
if buffer.strip() != "":
buffer = buffer.splitlines()[-1]
_, energy, error = [float(x) for x in buffer.split()]
return energy, error
else:
return None, None
def get_variance():
buffer = subprocess.check_output(['qmcchem', 'result', '-e',
'e_loc_qmcvar', filename],
encoding='UTF-8')
if buffer.strip() != "":
buffer = buffer.splitlines()[-1]
_, variance, error = [float(x) for x in buffer.split()]
return variance, error
else:
return None, None
def run_qmc():
return subprocess.check_output(['qmcchem', 'run', filename])
def stop_qmc():
subprocess.check_output(['qmcchem', 'stop', filename])
def set_vmc_params():
# subprocess.check_output(['qmcchem', 'edit', '-c', '-j', jast_type,
# '-m', 'VMC',
# '-l', str(block_time),
# '--time-step=0.3',
# '--stop-time=36000',
# '--norm=1.e-5',
# '-w', '10',
# filename])
subprocess.check_output(['qmcchem', 'edit', '-c', '-j', jast_type,
'-l', str(block_time),
filename])
memo_energy = {'fmin': 100000000.}
def f(x):
print ("x = %s"%str(x))
sys.stdout.flush()
h = str(x)
if h in memo_energy:
return memo_energy[h]
set_params_b(x[0])
set_params_pen(x[1:])
set_vmc_params()
pid = os.fork()
if pid == 0:
run_qmc()
os._exit(os.EX_OK)
else:
import atexit
atexit.register(stop_qmc)
err = thresh+1.
time.sleep(3.*block_time/4.)
local_thresh = thresh
while err > local_thresh:
time.sleep(block_time)
e, e_err = get_energy()
variance, v_err = get_variance()
if e is None or variance is None:
continue
# norm, _ = get_norm()
# e, e_err = e, e_err/norm
# variance, v_err =variance/norm, v_err/norm
energy = e #+ variance
err = e_err #sqrt(e_err*e_err+v_err*v_err)
print(" %f %f %f %f %f %f"%(e, e_err, variance, v_err, energy, err))
sys.stdout.flush()
if (energy-2.*err) > memo_energy['fmin']+thresh:
local_thresh = 10.*thresh
elif (energy+2.*err) < memo_energy['fmin']-thresh:
local_thresh = 10.*thresh
# Check if PID is still running
try:
os.kill(pid,0)
except OSError:
print("---")
sys.stdout.flush()
break
stop_qmc()
os.wait()
memo_energy[h] = energy + err
memo_energy['fmin'] = min(energy, memo_energy['fmin'])
return energy
def run():
x = np.array([ get_params_b() ] + list(get_params_pen()))
if sum(x) == 0.:
jast_a_up_dn = ezfio.jastrow_jast_a_up_dn
x += jast_a_up_dn
# opt = sp.optimize.minimize(f,x,method="Powell",
# options= {'disp':True, 'ftol':thresh,'xtol':0.02})
opt = sp.optimize.minimize(f,x,method="Nelder-Mead",
options= {'disp':True, 'ftol':thresh,'xtol':0.02})
print("x = "+str(opt))
sys.stdout.flush()
set_params_b(opt['x'][0])
set_params_pen(opt['x'][1:])
run()
if __name__ == '__main__':
main()