Corrected bugs, refactored part of the code (149d6417) · Commits · Germain Clavier / python-tools

lammps/plot_forces.py

+41 −28

Original line number	Diff line number	Diff line
		@@ -12,6 +12,8 @@ import logging
		import sys
		from matplotlib import pyplot as plt

		logger = logging.getLogger(__name__)


		def compute_energy(tp):
		r = tp[0]
		@@ -26,14 +28,14 @@ def compute_energy(tp):
		return e

		units = {
		"lj": {"Energy": "Reduced units", "Force": "Reduced units", "kb": 1},
		"real": {"Energy": "kcal/mol", "Force": "kcal/mol/A", "kb": 0.001985875},
		"metal": {"Energy": "eV", "Force": "eV/A", "kb": 8.6173332e-5},
		"si": {"Energy": "J", "Force": "N", "kb": 1.380649e-23},
		"cgs": {"Energy": "ergs", "Force": "dynes", "kb": 1.3806504e-16},
		"electron": {"Energy": "Hartrees", "Force": "Hartree/Bohr", "kb": 3.16681534e-6},
		"micro": {"Energy": "pg·um^2/us^2", "Force": "pg·um/us^2", "kb": 1.3806504e-8},
		"nano": {"Energy": "ag·nm^2/ns^2", "Force": "ag·nm/ns^2", "kb": 0.013806504},
		"lj": {"Distance": "Reduced units","Energy": "Reduced units", "Force": "Reduced units", "kb": 1},
		"real": {"Distance": "[A]","Energy": "[kcal/mol]", "Force": "[kcal/mol/A]", "kb": 0.001985875},
		"metal": {"Distance": "[A]","Energy": "[eV]", "Force": "[eV/A]", "kb": 8.6173332e-5},
		"si": {"Distance": "[m]","Energy": "[J]", "Force": "[N]", "kb": 1.380649e-23},
		"cgs": {"Distance": "[cm]","Energy": "[ergs]", "Force": "[dynes]", "kb": 1.3806504e-16},
		"electron": {"Distance": "[Bohr]","Energy": "[Hartrees]", "Force": "[Hartree/Bohr]", "kb": 3.16681534e-6},
		"micro": {"Distance": "[um]","Energy": "[pg·um^2/us^2]", "Force": "[pg·um/us^2]", "kb": 1.3806504e-8},
		"nano": {"Distance": "[nm]","Energy": "[ag·nm^2/ns^2]", "Force": "[ag·nm/ns^2]", "kb": 0.013806504},
		}

		def main():
		@@ -43,9 +45,9 @@ def main():
		)
		parser.add_argument(
		"-u",
		"--units"
		dest="units"
		default="real"
		"--units",
		dest="units",
		default="real",
		help="Units of the file (LAMMPS units system)",
		)
		parser.add_argument(
		@@ -86,17 +88,26 @@ def main():
		help="Extract the forces in separate files",
		)
		args = parser.parse_args()
		logging.basicConfig(level=logging.INFO)

		##########
		# Manage arguments

		uenergy = units(args.units)["Energy"]
		uforce = units(args.units)["Force"]
		kb = units(args.units)["kb"]
		udistance = units[args.units]["Distance"]
		uenergy = units[args.units]["Energy"]
		uforce = units[args.units]["Force"]
		kb = units[args.units]["kb"]
		rlabel = ' '.join(["Rij", udistance])
		elabel = ' '.join(["E", uenergy])
		flabel = ' '.join(["F", uforce])
		etitle = "Energy"
		ftitle = "Force"
		font = "DejaVu Sans"
		fontsize = 30

		infile = args.infile
		if not os.path.isfile(infile):
		logging.error("Input file not found")
		logger.error("Input file not found")
		sys.exit(1)

		toplot = []
		@@ -110,7 +121,7 @@ def main():
		tok = []
		while not tok:
		tok = next(lines).partition('#')[0].rstrip()
		logging.info("Found {} token".format(tok))
		logger.info("Found {} token".format(tok))
		infos = next(lines).split()
		npoints = int(infos[1])
		next(lines)
		@@ -128,6 +139,7 @@ def main():
		except StopIteration:
		break
		if args.recompute:
		etitle = "Estimated energy"
		for tp in toplot:
		tp[1] = compute_energy(tp)

		@@ -143,9 +155,10 @@ def main():
		if args.temp > 0:
		hmin, hmax = axes[0].get_xlim()
		axes[0].hlines(kb*args.temp, hmin, hmax, color="orange", label=r'$k_BT$', linewidth=3, linestyles="dashdot")
		logging.info("KbT value= {:e}".format(kb*args.temp))
		axes[0].text(hmax/2., kbargs.temp, 'KbT', fontsize=0.7fontsize)
		logger.info("KbT value= {:e} {:s}".format(kb*args.temp, uenergy))
		else:
		logging.info("Invalid temperature value: {:e}".format(args.temp))
		logger.info("Invalid temperature value: {:e}".format(args.temp))

		if args.xrange:
		xmin, xmax = list(map(float, args.xrange.split(":")))
		@@ -156,30 +169,30 @@ def main():
		axes[0].set_ylim(ymin, ymax)
		axes[1].set_ylim(ymin, ymax)

		font = "DejaVu Sans"
		fontsize = 30

		# Setting axes 0
		axes[0].set_title("Estimated energy", fontsize=fontsize)
		axes[0].set_title(etitle, fontsize=fontsize)
		# axes[0].legend(frameon=False, fontsize=fontsize) # Fat font, no frame
		axes[0].set_xlabel("r [A]", fontname=font, fontsize=fontsize) # ylabel name, font is Hack, size 30pts
		axes[0].set_ylabel("E [kcal/mol]", fontname=font, fontsize=fontsize) # ylabel name, font is Hack, size 30pts
		axes[0].set_xlabel(rlabel, fontname=font, fontsize=fontsize) # ylabel name, font is Hack, size 30pts
		axes[0].set_ylabel(elabel, fontname=font, fontsize=fontsize) # ylabel name, font is Hack, size 30pts
		axes[0].tick_params(axis='both', which='major', labelsize=fontsize) # Biggers ticks, bigger ticks label!

		# Setting axes 1
		axes[1].set_title("Tabulated force", fontsize=fontsize)
		axes[1].set_title(ftitle, fontsize=fontsize)
		axes[1].legend(frameon=False, fontsize=fontsize) # Fat font, no frame
		axes[1].set_xlabel("r [A]", fontname=font, fontsize=fontsize) # ylabel name, font is Hack, size 30pts
		axes[1].set_ylabel("F [kcal/mol/A]", fontname=font, fontsize=fontsize) # ylabel name, font is Hack, size 30pts
		axes[1].set_xlabel(rlabel, fontname=font, fontsize=fontsize) # ylabel name, font is Hack, size 30pts
		axes[1].set_ylabel(flabel, fontname=font, fontsize=fontsize) # ylabel name, font is Hack, size 30pts
		axes[1].tick_params(axis='both', which='major', labelsize=fontsize) # Biggers ticks, bigger ticks label!

		plt.subplots_adjust(wspace=0.3)
		figManager = plt.get_current_fig_manager()
		figManager.window.showMaximized()
		plt.show()

		if args.extract:
		for tp in toplot:
		outfile = "".join([tp[3], '.plot'])
		logging.info("Writing file {}".format(outfile))
		logger.info("Writing file {}".format(outfile))
		with open(outfile, 'w') as f:
		f.write("# {} force extracted from {}\n".format(tp[3], infile))
		f.write("# {:^20} {:^20} {:^20}\n".format('r', 'energy', 'force'))