added lammps python example

examples/python/gjf_python/gjf.py

@@ -0,0 +1,180 @@
+"""Made by Charlie Sievers Ph.D. Candidate, UC Davis, Donadio Lab 2019"""
+
+from mpi4py import MPI
+from lammps import lammps
+import lammps_tools as lt
+import numpy as np
+
+comm = MPI.COMM_WORLD
+rank = comm.Get_rank()
+
+""" LAMMPS  VARIABLES """
+
+# new file or restart
+run_no = 0
+
+# data files
+infile = "argon.lmp"
+restart_file = "final_restart.{}".format(run_no)
+ff_file = "ff-argon.lmp"
+outfile = "output.dat"
+
+# write final_restart
+write_final_restart = False
+
+# random numbers
+seed0 = 2357
+seed1 = 26588
+seed2 = 10669
+
+# MD Parameters
+# number of steps
+nsteps = 50000
+# timestep
+# dt = 0.001
+# starting simulation temp
+temp_start = 10
+# final simulation temp
+temp_final = 10
+# relaxation time
+trel = 1
+# trajectory frequency
+ntraj = 0
+
+# Ensemble 0 = GJF u, 1 = GJF v, 2 = Nose-Hoover, 3 = Langevin, 4 = BDP (Currently all NVT)
+ensemble = 0
+
+# Output Parameters
+nthermo = 200
+nout = int(nsteps / nthermo)  # Important
+
+# output to screen and log file?
+lammps_output = False
+# Lammps Thermo
+thermo = False
+
+python_output = True
+
+# Write output to file?
+write_output = False
+
+if write_output is True:
+    data = open("{}".format(outfile), "w")
+
+if python_output is True:
+    if rank == 0:
+        print("dt, temp, ke, fke, pe, fpe")
+
+for j in range(20):
+
+    # timestep
+    dt = 0.005*(j+1)
+
+    if lammps_output is True:
+        lmp = lammps()
+    else:
+        lmp = lammps(cmdargs=["-screen", "none", "-log", "none"])
+
+    lmp.command("atom_style full")
+    lmp.command("units metal")
+    lmp.command("processors * * *")
+    lmp.command("neighbor 1 bin")
+    lmp.command("boundary p p p")
+
+    if run_no is 0:
+        lmp.command("read_data {}".format(infile))
+    else:
+        lmp.command("read_restart final_restart".format(run_no-1))
+
+    if thermo is True:
+        lmp.command("thermo_style custom time temp  pe ke press vol  cpu")
+        lmp.command("thermo {}".format(nthermo))
+        lmp.command("thermo_modify flush yes")
+
+    lmp.file("{}".format(ff_file))
+    lmp.command("timestep {}".format(dt))
+
+    # get_per_atom_compute example with dim of two and within a group
+    # lmp.command("region rand block 5 20 5 20 5 20")
+    # lmp.command("group rand region rand")
+    # lmp.command("compute x rand property/atom x y")
+    # test = get_per_atom_compute(comm, lmp, "x", 2, group="rand")
+
+    lmp.command("compute ke all ke/atom")
+
+    lmp.command("compute pe all pe")
+
+    if ntraj != 0:
+        lmp.command("dump 1 all dcd {} trajectory.dcd".format(ntraj))
+        lmp.command("dump_modify 1 unwrap yes")
+
+    if run_no == 0:
+        lmp.command("velocity all create {} {} mom yes dist gaussian".format(temp_start, seed0))
+    lmp.command("fix nve all nve")
+
+    if ensemble == 0:
+        # gjf u
+        lmp.command("fix lang all langevin {} {} {} {} gjf yes halfstep yes".format(
+            temp_start, temp_final, trel, seed1))
+    elif ensemble == 1:
+        # gjf v
+        lmp.command("fix lang all langevin {} {} {} {} gjf yes".format(
+            temp_start, temp_final, trel, seed1))
+    elif ensemble == 2:
+        # NH
+        lmp.command("fix nvt all nvt temp {} {} {}".format(
+            temp_start, temp_final, trel))
+    elif ensemble == 3:
+        # lang
+        lmp.command("fix lang all langevin {} {} {} {} tally yes zero yes".format(
+            temp_start, temp_final, trel, seed1))
+    elif ensemble == 4:
+        # BDP
+        lmp.command("fix stoch all temp/csvr {} {} {} {}".format(
+            temp_start, temp_final, trel, seed1))
+
+    natoms = lmp.extract_global("natoms", 0)
+    nlocal = lmp.extract_global("nlocal", 0)
+    ke_sum = lt.get_per_atom_compute(comm, lmp, "ke")
+    ke_2 = ke_sum**2
+    pe_sum = 0
+    pe_2 = 0
+    temp_sum = 0
+
+    for i in range(nout):
+        nlocal = lmp.extract_global("nlocal", 0)
+        lmp.command("run {} pre no post no".format(nthermo))
+        temp = lmp.extract_compute("thermo_temp", 0, 0)
+        ke = lt.get_per_atom_compute(comm, lmp, "ke")
+        pe = lmp.extract_compute("pe", 0, 0)
+        ke_sum += ke
+        ke_2 += ke**2
+        pe_sum += pe
+        pe_2 += pe**2
+        temp_sum += temp
+
+        if python_output is True:
+            if rank == 0:
+                print("Time: {:.6f}, Temp: {:.6f}, KE: {:.6f}, PE: {:.6f}".format(
+                    i*nthermo*dt, temp, ke.sum(), pe))
+
+    if write_final_restart is True:
+        lmp.command("write_restart {}".format(restart_file))
+
+    if rank == 0:
+        ke = ke_sum.sum() / (nout + 1)
+        fke = (np.sqrt((ke_2 - ke_sum ** 2 / (nout + 1)) / (nout + 1))).sum()
+        pe = pe_sum / nout
+        fpe = np.sqrt((pe_2 - pe_sum ** 2 / nout) / nout)
+        temp = temp_sum / nout
+
+        if python_output is True:
+            print(dt, temp, ke, fke, pe, fpe)
+
+        if write_output is True:
+            data.write("{:.6f} {:.6f} {:.6f} {:.6f} {:.6f} {:.6f}\n".format(
+                dt, temp, ke, fke, pe, fpe))
+            data.flush()
+
+if write_output is True:
+    data.close()