Added mu_opt

bin/mu_opt.py

@@ -0,0 +1,201 @@
+#!/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()
+
+
+
+    def get_params_b():
+        return ezfio.jastrow_mu_erf
+
+    def set_params_b(x):
+        x = np.abs(x)
+        ezfio.set_jastrow_mu_erf(x)
+
+    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, 0.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()]
+            return energy, error
+        else:
+            return None, None
+
+    def get_energy():
+        buffer = subprocess.check_output(['qmcchem', 'result', '-e', 'Ci_dress_mu_opt_qmcvar', 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',
+                                          'Ci_dress_mu_opt_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/norm, 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"%(e, err, local_thresh, memo_energy['fmin']))
+                sys.stdout.flush()
+                if (energy-2.*err) > memo_energy['fmin']+local_thresh:
+                   local_thresh = 10.*local_thresh
+                elif (energy+2.*err) < memo_energy['fmin']-local_thresh:
+                   local_thresh = 10.*local_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()
+
+
+