# coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.
"""
This module defines utility classes and functions.
"""
import os
import tempfile
from io import open
from subprocess import Popen, PIPE
import numpy as np
try:
from openbabel import pybel as pb
except ImportError:
pb = None
from pymatgen import Molecule
from pymatgen.core.operations import SymmOp
from pymatgen.util.coord import get_angle
from pymatgen.io.babel import BabelMolAdaptor
from monty.os.path import which
from monty.tempfile import ScratchDir
__author__ = 'Kiran Mathew, Brandon Wood, Michael Humbert'
__email__ = 'kmathew@lbl.gov'
[docs]class Polymer:
"""
Generate polymer chain via Random walk. At each position there are
a total of 5 possible moves(excluding the previous direction).
"""
def __init__(self, start_monomer, s_head, s_tail,
monomer, head, tail,
end_monomer, e_head, e_tail,
n_units, link_distance=1.0, linear_chain=False):
"""
Args:
start_monomer (Molecule): Starting molecule
s_head (int): starting atom index of the start_monomer molecule
s_tail (int): tail atom index of the start_monomer
monomer (Molecule): The monomer
head (int): index of the atom in the monomer that forms the head
tail (int): tail atom index. monomers will be connected from
tail to head
end_monomer (Molecule): Terminal molecule
e_head (int): starting atom index of the end_monomer molecule
e_tail (int): tail atom index of the end_monomer
n_units (int): number of monomer units excluding the start and
terminal molecules
link_distance (float): distance between consecutive monomers
linear_chain (bool): linear or random walk polymer chain
"""
self.start = s_head
self.end = s_tail
self.monomer = monomer
self.n_units = n_units
self.link_distance = link_distance
self.linear_chain = linear_chain
# translate monomers so that head atom is at the origin
start_monomer.translate_sites(range(len(start_monomer)),
- monomer.cart_coords[s_head])
monomer.translate_sites(range(len(monomer)),
- monomer.cart_coords[head])
end_monomer.translate_sites(range(len(end_monomer)),
- monomer.cart_coords[e_head])
self.mon_vector = monomer.cart_coords[tail] - monomer.cart_coords[head]
self.moves = {1: [1, 0, 0],
2: [0, 1, 0],
3: [0, 0, 1],
4: [-1, 0, 0],
5: [0, -1, 0],
6: [0, 0, -1]}
self.prev_move = 1
# places the start monomer at the beginning of the chain
self.molecule = start_monomer.copy()
self.length = 1
# create the chain
self._create(self.monomer, self.mon_vector)
# terminate the chain with the end_monomer
self.n_units += 1
end_mon_vector = end_monomer.cart_coords[e_tail] - end_monomer.cart_coords[e_head]
self._create(end_monomer, end_mon_vector)
self.molecule = Molecule.from_sites(self.molecule.sites)
def _create(self, monomer, mon_vector):
"""
create the polymer from the monomer
Args:
monomer (Molecule)
mon_vector (numpy.array): molecule vector that starts from the
start atom index to the end atom index
"""
while self.length != (self.n_units - 1):
if self.linear_chain:
move_direction = np.array(mon_vector) / np.linalg.norm(mon_vector)
else:
move_direction = self._next_move_direction()
self._add_monomer(monomer.copy(), mon_vector, move_direction)
def _next_move_direction(self):
"""
pick a move at random from the list of moves
"""
nmoves = len(self.moves)
move = np.random.randint(1, nmoves + 1)
while self.prev_move == (move + 3) % nmoves:
move = np.random.randint(1, nmoves + 1)
self.prev_move = move
return np.array(self.moves[move])
def _align_monomer(self, monomer, mon_vector, move_direction):
"""
rotate the monomer so that it is aligned along the move direction
Args:
monomer (Molecule)
mon_vector (numpy.array): molecule vector that starts from the
start atom index to the end atom index
move_direction (numpy.array): the direction of the polymer chain
extension
"""
axis = np.cross(mon_vector, move_direction)
origin = monomer[self.start].coords
angle = get_angle(mon_vector, move_direction)
op = SymmOp.from_origin_axis_angle(origin, axis, angle)
monomer.apply_operation(op)
def _add_monomer(self, monomer, mon_vector, move_direction):
"""
extend the polymer molecule by adding a monomer along mon_vector direction
Args:
monomer (Molecule): monomer molecule
mon_vector (numpy.array): monomer vector that points from head to tail.
move_direction (numpy.array): direction along which the monomer
will be positioned
"""
translate_by = self.molecule.cart_coords[self.end] + self.link_distance * move_direction
monomer.translate_sites(range(len(monomer)), translate_by)
if not self.linear_chain:
self._align_monomer(monomer, mon_vector, move_direction)
# add monomer if there are no crossings
does_cross = False
for i, site in enumerate(monomer):
try:
self.molecule.append(site.specie, site.coords, properties=site.properties)
except Exception:
does_cross = True
polymer_length = len(self.molecule)
self.molecule.remove_sites(
range(polymer_length - i, polymer_length))
break
if not does_cross:
self.length += 1
self.end += len(self.monomer)
[docs]class PackmolRunner:
"""
Wrapper for the Packmol software that can be used to pack various types of
molecules into a one single unit.
"""
def __init__(self, mols, param_list, input_file="pack.inp",
tolerance=2.0, filetype="xyz",
control_params={"maxit": 20, "nloop": 600},
auto_box=True, output_file="packed.xyz",
bin="packmol"):
"""
Args:
mols:
list of Molecules to pack
input_file:
name of the packmol input file
tolerance:
min distance between the atoms
filetype:
input/output structure file type
control_params:
packmol control parameters dictionary. Basically
all parameters other than structure/atoms
param_list:
list of parameters containing dicts for each molecule
auto_box:
put the molecule assembly in a box
output_file:
output file name. The extension will be adjusted
according to the filetype
"""
self.packmol_bin = bin.split()
if not which(self.packmol_bin[-1]):
raise RuntimeError(
"PackmolRunner requires the executable 'packmol' to be in "
"the path. Please download packmol from "
"https://github.com/leandromartinez98/packmol "
"and follow the instructions in the README to compile. "
"Don't forget to add the packmol binary to your path")
self.mols = mols
self.param_list = param_list
self.input_file = input_file
self.boxit = auto_box
self.control_params = control_params
if not self.control_params.get("tolerance"):
self.control_params["tolerance"] = tolerance
if not self.control_params.get("filetype"):
self.control_params["filetype"] = filetype
if not self.control_params.get("output"):
self.control_params["output"] = "{}.{}".format(output_file.split(".")[0], self.control_params["filetype"])
if self.boxit:
self._set_box()
def _format_param_val(self, param_val):
"""
Internal method to format values in the packmol parameter dictionaries
Args:
param_val:
Some object to turn into String
Returns:
string representation of the object
"""
if isinstance(param_val, list):
return ' '.join(str(x) for x in param_val)
else:
return str(param_val)
def _set_box(self):
"""
Set the box size for the molecular assembly
"""
net_volume = 0.0
for idx, mol in enumerate(self.mols):
length = max([np.max(mol.cart_coords[:, i]) - np.min(mol.cart_coords[:, i])
for i in range(3)]) + 2.0
net_volume += (length ** 3.0) * float(self.param_list[idx]['number'])
length = net_volume ** (1.0 / 3.0)
for idx, mol in enumerate(self.mols):
self.param_list[idx]['inside box'] = '0.0 0.0 0.0 {} {} {}'.format(
length, length, length)
def _write_input(self, input_dir="."):
"""
Write the packmol input file to the input directory.
Args:
input_dir (string): path to the input directory
"""
with open(os.path.join(input_dir, self.input_file), 'wt', encoding="utf-8") as inp:
for k, v in self.control_params.items():
inp.write('{} {}\n'.format(k, self._format_param_val(v)))
# write the structures of the constituent molecules to file and set
# the molecule id and the corresponding filename in the packmol
# input file.
for idx, mol in enumerate(self.mols):
filename = os.path.join(
input_dir, '{}.{}'.format(
idx, self.control_params["filetype"]))
# pdb
if self.control_params["filetype"] == "pdb":
self.write_pdb(mol, filename, num=idx + 1)
# all other filetypes
else:
a = BabelMolAdaptor(mol)
pm = pb.Molecule(a.openbabel_mol)
pm.write(self.control_params["filetype"], filename=filename,
overwrite=True)
inp.write("\n")
inp.write(
"structure {}.{}\n".format(
os.path.join(input_dir, str(idx)),
self.control_params["filetype"]))
for k, v in self.param_list[idx].items():
inp.write(' {} {}\n'.format(k, self._format_param_val(v)))
inp.write('end structure\n')
[docs] def run(self, copy_to_current_on_exit=False, site_property=None):
"""
Write the input file to the scratch directory, run packmol and return
the packed molecule.
Args:
copy_to_current_on_exit (bool): Whether or not to copy the packmol
input/output files from the scratch directory to the current
directory.
site_property (str): if set then the specified site property
for the the final packed molecule will be restored.
Returns:
Molecule object
"""
scratch = tempfile.gettempdir()
with ScratchDir(scratch, copy_to_current_on_exit=copy_to_current_on_exit) as scratch_dir:
self._write_input(input_dir=scratch_dir)
packmol_input = open(os.path.join(scratch_dir, self.input_file), 'r')
p = Popen(self.packmol_bin, stdin=packmol_input, stdout=PIPE, stderr=PIPE)
(stdout, stderr) = p.communicate()
output_file = os.path.join(scratch_dir, self.control_params["output"])
if os.path.isfile(output_file):
packed_mol = BabelMolAdaptor.from_file(output_file,
self.control_params["filetype"])
packed_mol = packed_mol.pymatgen_mol
print("packed molecule written to {}".format(
self.control_params["output"]))
if site_property:
packed_mol = self.restore_site_properties(site_property=site_property, filename=output_file)
return packed_mol
else:
print("Packmol execution failed")
print(stdout, stderr)
return None
[docs] def write_pdb(self, mol, filename, name=None, num=None):
"""
dump the molecule into pdb file with custom residue name and number.
"""
# ugly hack to get around the openbabel issues with inconsistent
# residue labelling.
scratch = tempfile.gettempdir()
with ScratchDir(scratch, copy_to_current_on_exit=False) as _:
mol.to(fmt="pdb", filename="tmp.pdb")
bma = BabelMolAdaptor.from_file("tmp.pdb", "pdb")
num = num or 1
name = name or "ml{}".format(num)
# bma = BabelMolAdaptor(mol)
pbm = pb.Molecule(bma._obmol)
for i, x in enumerate(pbm.residues):
x.OBResidue.SetName(name)
x.OBResidue.SetNum(num)
pbm.write(format="pdb", filename=filename, overwrite=True)
def _set_residue_map(self):
"""
map each residue to the corresponding molecule.
"""
self.map_residue_to_mol = {}
lookup = {}
for idx, mol in enumerate(self.mols):
if mol.formula not in lookup:
mol.translate_sites(indices=range(len(mol)),
vector=-mol.center_of_mass)
lookup[mol.formula] = mol.copy()
self.map_residue_to_mol["ml{}".format(idx + 1)] = lookup[mol.formula]
[docs] def convert_obatoms_to_molecule(self, atoms, residue_name=None, site_property="ff_map"):
"""
Convert list of openbabel atoms to MOlecule.
Args:
atoms ([OBAtom]): list of OBAtom objects
residue_name (str): the key in self.map_residue_to_mol. Usec to
restore the site properties in the final packed molecule.
site_property (str): the site property to be restored.
Returns:
Molecule object
"""
restore_site_props = True if residue_name is not None else False
if restore_site_props and not hasattr(self, "map_residue_to_mol"):
self._set_residue_map()
coords = []
zs = []
for atm in atoms:
coords.append(list(atm.coords))
zs.append(atm.atomicnum)
mol = Molecule(zs, coords)
if restore_site_props:
props = []
ref = self.map_residue_to_mol[residue_name].copy()
# sanity check
assert len(mol) == len(ref)
assert ref.formula == mol.formula
# the packed molecules have the atoms in the same order..sigh!
for i, site in enumerate(mol):
assert site.specie.symbol == ref[i].specie.symbol
props.append(getattr(ref[i], site_property))
mol.add_site_property(site_property, props)
return mol
[docs] def restore_site_properties(self, site_property="ff_map", filename=None):
"""
Restore the site properties for the final packed molecule.
Args:
site_property (str):
filename (str): path to the final packed molecule.
Returns:
Molecule
"""
# only for pdb
if not self.control_params["filetype"] == "pdb":
raise ValueError()
filename = filename or self.control_params["output"]
bma = BabelMolAdaptor.from_file(filename, "pdb")
pbm = pb.Molecule(bma._obmol)
assert len(pbm.residues) == sum([x["number"]
for x in self.param_list])
packed_mol = self.convert_obatoms_to_molecule(
pbm.residues[0].atoms, residue_name=pbm.residues[0].name,
site_property=site_property)
for resid in pbm.residues[1:]:
mol = self.convert_obatoms_to_molecule(
resid.atoms, residue_name=resid.name,
site_property=site_property)
for site in mol:
packed_mol.append(site.species, site.coords,
properties=site.properties)
return packed_mol
[docs]class LammpsRunner:
"""
LAMMPS wrapper
"""
def __init__(self, input_filename="lammps.in", bin="lammps"):
"""
Args:
input_filename (string): input file name
bin (string): command to run, excluding the input file name
"""
self.lammps_bin = bin.split()
if not which(self.lammps_bin[-1]):
raise RuntimeError(
"LammpsRunner requires the executable {} to be in the path. "
"Please download and install LAMMPS from "
"http://lammps.sandia.gov. "
"Don't forget to add the binary to your path".format(self.lammps_bin[-1]))
self.input_filename = input_filename
[docs] def run(self):
"""
Write the input/data files and run LAMMPS.
"""
lammps_cmd = self.lammps_bin + ['-in', self.input_filename]
print("Running: {}".format(" ".join(lammps_cmd)))
p = Popen(lammps_cmd, stdout=PIPE, stderr=PIPE)
(stdout, stderr) = p.communicate()
return stdout, stderr
if __name__ == '__main__':
ethanol_coords = [[0.00720, -0.56870, 0.00000],
[-1.28540, 0.24990, 0.00000],
[1.13040, 0.31470, 0.00000],
[0.03920, -1.19720, 0.89000],
[0.03920, -1.19720, -0.89000],
[-1.31750, 0.87840, 0.89000],
[-1.31750, 0.87840, -0.89000],
[-2.14220, -0.42390, -0.00000],
[1.98570, -0.13650, -0.00000]]
ethanol = Molecule(["C", "C", "O", "H", "H", "H", "H", "H", "H"],
ethanol_coords)
water_coords = [[9.626, 6.787, 12.673],
[9.626, 8.420, 12.673],
[10.203, 7.604, 12.673]]
water = Molecule(["H", "H", "O"], water_coords)
pmr = PackmolRunner([ethanol, water],
[{"number": 1, "fixed": [0, 0, 0, 0, 0, 0],
"centerofmass": ""},
{"number": 15, "inside sphere": [0, 0, 0, 5]}],
input_file="packmol_input.inp", tolerance=2.0,
filetype="xyz",
control_params={"nloop": 1000},
auto_box=False, output_file="cocktail.xyz")
s = pmr.run()