4075885785

Build Type

push

github

Committed by Shyue Ping Ong

Commit Message

Merge branch 'master' of github.com:materialsproject/pymatgen

Run Details

96 of 96 new or added lines in 27 files covered. (100.0%)

81013 of 102710 relevant lines covered (78.88%)

0.79 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

41.78

/pymatgen/io/lammps/utils.py

# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.

"""
This module defines utility classes and functions.
"""

from __future__ import annotations

import os
import tempfile
from shutil import which
from subprocess import PIPE, Popen

import numpy as np
from monty.dev import deprecated
from monty.tempfile import ScratchDir

from pymatgen.core.operations import SymmOp
from pymatgen.core.structure import Molecule
from pymatgen.io.babel import BabelMolAdaptor
from pymatgen.io.packmol import PackmolBoxGen
from pymatgen.util.coord import get_angle

try:
    from openbabel import pybel as pb
except ImportError:
    pb = None

__author__ = "Kiran Mathew, Brandon Wood, Michael Humbert"
__email__ = "kmathew@lbl.gov"


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)


@deprecated(PackmolBoxGen, "PackmolRunner is being phased out in favor of the packmol I/O class.")
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=None,
        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 or {"maxit": 20, "nloop": 600}
        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"] = f"{output_file.split('.')[0]}.{self.control_params['filetype']}"
        if self.boxit:
            self._set_box()

    @staticmethod
    def _format_param_val(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)
        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"] = f"0.0 0.0 0.0 {length} {length} {length}"

    def _write_input(self, input_dir="."):
        """
        Write the packmol input file to the input directory.

        Args:
            input_dir (str): path to the input directory
        """
        with open(os.path.join(input_dir, self.input_file), "w", encoding="utf-8") as inp:
            for k, v in self.control_params.items():
                inp.write(f"{k} {self._format_param_val(v)}\n")
            # 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, f"{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(f"structure {os.path.join(input_dir, str(idx))}.{self.control_params['filetype']}\n")
                for k, v in self.param_list[idx].items():
                    inp.write(f"  {k} {self._format_param_val(v)}\n")
                inp.write("end structure\n")

    def run(self, site_property=None):
        """
        Write the input file to the scratch directory, run packmol and return
        the packed molecule to the current working directory.

        Args:
            site_property (str): if set then the specified site property
                for the final packed molecule will be restored.

        Returns:
                Molecule object
        """
        with tempfile.TemporaryDirectory() as scratch_dir:
            self._write_input(input_dir=scratch_dir)
            with open(os.path.join(scratch_dir, self.input_file)) as packmol_input:
                with Popen(self.packmol_bin, stdin=packmol_input, stdout=PIPE, stderr=PIPE) as p:
                    (stdout, stderr) = p.communicate()
            output_file = os.path.join(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(f"packed molecule written to {self.control_params['output']}")
                if site_property:
                    packed_mol = self.restore_site_properties(site_property=site_property, filename=output_file)
                return packed_mol
            raise RuntimeError(f"Packmol execution failed. {stdout}\n{stderr}")

    @staticmethod
    def write_pdb(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 f"ml{num}"

        # bma = BabelMolAdaptor(mol)
        pbm = pb.Molecule(bma._obmol)
        for x in 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[f"ml{idx + 1}"] = lookup[mol.formula]

    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 = residue_name is not None

        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

    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


class LammpsRunner:
    """
    LAMMPS wrapper
    """

    def __init__(self, input_filename="lammps.in", bin="lammps"):
        """
        Args:
            input_filename (str): input file name
            bin (str): command to run, excluding the input file name
        """
        self.lammps_bin = bin.split()
        if not which(self.lammps_bin[-1]):
            raise RuntimeError(
                f"LammpsRunner requires the executable {self.lammps_bin[-1]} to be in the path. "
                "Please download and install LAMMPS from "
                "https://www.lammps.org/. "
                "Don't forget to add the binary to your path"
            )
        self.input_filename = input_filename

    def run(self):
        """
        Write the input/data files and run LAMMPS.
        """
        lammps_cmd = self.lammps_bin + ["-in", self.input_filename]
        print(f"Running: {' '.join(lammps_cmd)}")
        with Popen(lammps_cmd, stdout=PIPE, stderr=PIPE) as p:
            (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()

1	# Copyright (c) Pymatgen Development Team.
2	# Distributed under the terms of the MIT License.
3
4	"""	1✔
5	This module defines utility classes and functions.
6	"""
7
8	from __future__ import annotations	1✔
9
10	import os	1✔
11	import tempfile	1✔
12	from shutil import which	1✔
13	from subprocess import PIPE, Popen	1✔
14
15	import numpy as np	1✔
16	from monty.dev import deprecated	1✔
17	from monty.tempfile import ScratchDir	1✔
18
19	from pymatgen.core.operations import SymmOp	1✔
20	from pymatgen.core.structure import Molecule	1✔
21	from pymatgen.io.babel import BabelMolAdaptor	1✔
22	from pymatgen.io.packmol import PackmolBoxGen	1✔
23	from pymatgen.util.coord import get_angle	1✔
24
25	try:	1✔
26	from openbabel import pybel as pb	1✔
27	except ImportError:	1✔
28	pb = None	1✔
29
30	__author__ = "Kiran Mathew, Brandon Wood, Michael Humbert"	1✔
31	__email__ = "kmathew@lbl.gov"	1✔
32
33
34	class Polymer:	1✔
35	"""
36	Generate polymer chain via Random walk. At each position there are
37	a total of 5 possible moves(excluding the previous direction).
38	"""
39
40	def __init__(	1✔
41	self,
42	start_monomer,
43	s_head,
44	s_tail,
45	monomer,
46	head,
47	tail,
48	end_monomer,
49	e_head,
50	e_tail,
51	n_units,
52	link_distance=1.0,
53	linear_chain=False,
54	):
55	"""
56	Args:
57	start_monomer (Molecule): Starting molecule
58	s_head (int): starting atom index of the start_monomer molecule
59	s_tail (int): tail atom index of the start_monomer
60	monomer (Molecule): The monomer
61	head (int): index of the atom in the monomer that forms the head
62	tail (int): tail atom index. monomers will be connected from
63	tail to head
64	end_monomer (Molecule): Terminal molecule
65	e_head (int): starting atom index of the end_monomer molecule
66	e_tail (int): tail atom index of the end_monomer
67	n_units (int): number of monomer units excluding the start and
68	terminal molecules
69	link_distance (float): distance between consecutive monomers
70	linear_chain (bool): linear or random walk polymer chain
71	"""
72	self.start = s_head	1✔
73	self.end = s_tail	1✔
74	self.monomer = monomer	1✔
75	self.n_units = n_units	1✔
76	self.link_distance = link_distance	1✔
77	self.linear_chain = linear_chain	1✔
78	# translate monomers so that head atom is at the origin
79	start_monomer.translate_sites(range(len(start_monomer)), -monomer.cart_coords[s_head])	1✔
80	monomer.translate_sites(range(len(monomer)), -monomer.cart_coords[head])	1✔
81	end_monomer.translate_sites(range(len(end_monomer)), -monomer.cart_coords[e_head])	1✔
82	self.mon_vector = monomer.cart_coords[tail] - monomer.cart_coords[head]	1✔
83	self.moves = {	1✔
84	1: [1, 0, 0],
85	2: [0, 1, 0],
86	3: [0, 0, 1],
87	4: [-1, 0, 0],
88	5: [0, -1, 0],
89	6: [0, 0, -1],
90	}
91	self.prev_move = 1	1✔
92	# places the start monomer at the beginning of the chain
93	self.molecule = start_monomer.copy()	1✔
94	self.length = 1	1✔
95	# create the chain
96	self._create(self.monomer, self.mon_vector)	1✔
97	# terminate the chain with the end_monomer
98	self.n_units += 1	1✔
99	end_mon_vector = end_monomer.cart_coords[e_tail] - end_monomer.cart_coords[e_head]	1✔
100	self._create(end_monomer, end_mon_vector)	1✔
101	self.molecule = Molecule.from_sites(self.molecule.sites)	1✔
102
103	def _create(self, monomer, mon_vector):	1✔
104	"""
105	create the polymer from the monomer
106
107	Args:
108	monomer (Molecule)
109	mon_vector (numpy.array): molecule vector that starts from the
110	start atom index to the end atom index
111	"""
112	while self.length != (self.n_units - 1):	1✔
113	if self.linear_chain:	1✔
114	move_direction = np.array(mon_vector) / np.linalg.norm(mon_vector)	1✔
115	else:
116	move_direction = self._next_move_direction()	1✔
117	self._add_monomer(monomer.copy(), mon_vector, move_direction)	1✔
118
119	def _next_move_direction(self):	1✔
120	"""
121	pick a move at random from the list of moves
122	"""
123	nmoves = len(self.moves)	1✔
124	move = np.random.randint(1, nmoves + 1)	1✔
125	while self.prev_move == (move + 3) % nmoves:	1✔
126	move = np.random.randint(1, nmoves + 1)	1✔
127	self.prev_move = move	1✔
128	return np.array(self.moves[move])	1✔
129
130	def _align_monomer(self, monomer, mon_vector, move_direction):	1✔
131	"""
132	rotate the monomer so that it is aligned along the move direction
133
134	Args:
135	monomer (Molecule)
136	mon_vector (numpy.array): molecule vector that starts from the
137	start atom index to the end atom index
138	move_direction (numpy.array): the direction of the polymer chain
139	extension
140	"""
141	axis = np.cross(mon_vector, move_direction)	1✔
142	origin = monomer[self.start].coords	1✔
143	angle = get_angle(mon_vector, move_direction)	1✔
144	op = SymmOp.from_origin_axis_angle(origin, axis, angle)	1✔
145	monomer.apply_operation(op)	1✔
146
147	def _add_monomer(self, monomer, mon_vector, move_direction):	1✔
148	"""
149	extend the polymer molecule by adding a monomer along mon_vector direction
150
151	Args:
152	monomer (Molecule): monomer molecule
153	mon_vector (numpy.array): monomer vector that points from head to tail.
154	move_direction (numpy.array): direction along which the monomer
155	will be positioned
156	"""
157	translate_by = self.molecule.cart_coords[self.end] + self.link_distance * move_direction	1✔
158	monomer.translate_sites(range(len(monomer)), translate_by)	1✔
159	if not self.linear_chain:	1✔
160	self._align_monomer(monomer, mon_vector, move_direction)	1✔
161	# add monomer if there are no crossings
162	does_cross = False	1✔
163	for i, site in enumerate(monomer):	1✔
164	try:	1✔
165	self.molecule.append(site.specie, site.coords, properties=site.properties)	1✔
166	except Exception:	×
167	does_cross = True	×
168	polymer_length = len(self.molecule)	×
169	self.molecule.remove_sites(range(polymer_length - i, polymer_length))	×
170	break	×
171	if not does_cross:	1✔
172	self.length += 1	1✔
173	self.end += len(self.monomer)	1✔
174
175
176	@deprecated(PackmolBoxGen, "PackmolRunner is being phased out in favor of the packmol I/O class.")	1✔
177	class PackmolRunner:	1✔
178	"""
179	Wrapper for the Packmol software that can be used to pack various types of
180	molecules into a one single unit.
181	"""
182
183	def __init__(	1✔
184	self,
185	mols,
186	param_list,
187	input_file="pack.inp",
188	tolerance=2.0,
189	filetype="xyz",
190	control_params=None,
191	auto_box=True,
192	output_file="packed.xyz",
193	bin="packmol",
194	):
195	"""
196	Args:
197	mols:
198	list of Molecules to pack
199	input_file:
200	name of the packmol input file
201	tolerance:
202	min distance between the atoms
203	filetype:
204	input/output structure file type
205	control_params:
206	packmol control parameters dictionary. Basically
207	all parameters other than structure/atoms
208	param_list:
209	list of parameters containing dicts for each molecule
210	auto_box:
211	put the molecule assembly in a box
212	output_file:
213	output file name. The extension will be adjusted
214	according to the filetype
215	"""
216	self.packmol_bin = bin.split()	×
217	if not which(self.packmol_bin[-1]):	×
218	raise RuntimeError(	×
219	"PackmolRunner requires the executable 'packmol' to be in "
220	"the path. Please download packmol from "
221	"https://github.com/leandromartinez98/packmol "
222	"and follow the instructions in the README to compile. "
223	"Don't forget to add the packmol binary to your path"
224	)
225	self.mols = mols	×
226	self.param_list = param_list	×
227	self.input_file = input_file	×
228	self.boxit = auto_box	×
229	self.control_params = control_params or {"maxit": 20, "nloop": 600}	×
230	if not self.control_params.get("tolerance"):	×
231	self.control_params["tolerance"] = tolerance	×
232	if not self.control_params.get("filetype"):	×
233	self.control_params["filetype"] = filetype	×
234	if not self.control_params.get("output"):	×
235	self.control_params["output"] = f"{output_file.split('.')[0]}.{self.control_params['filetype']}"	×
236	if self.boxit:	×
237	self._set_box()	×
238
239	@staticmethod	1✔
240	def _format_param_val(param_val):	1✔
241	"""
242	Internal method to format values in the packmol parameter dictionaries
243
244	Args:
245	param_val:
246	Some object to turn into String
247
248	Returns:
249	String representation of the object
250	"""
251	if isinstance(param_val, list):	×
252	return " ".join(str(x) for x in param_val)	×
253	return str(param_val)	×
254
255	def _set_box(self):	1✔
256	"""
257	Set the box size for the molecular assembly
258	"""
259	net_volume = 0.0	×
260	for idx, mol in enumerate(self.mols):	×
261	length = max(np.max(mol.cart_coords[:, i]) - np.min(mol.cart_coords[:, i]) for i in range(3)) + 2.0	×
262	net_volume += (length*3.0) float(self.param_list[idx]["number"])	×
263	length = net_volume ** (1.0 / 3.0)	×
264	for idx, _mol in enumerate(self.mols):	×
265	self.param_list[idx]["inside box"] = f"0.0 0.0 0.0 {length} {length} {length}"	×
266
267	def _write_input(self, input_dir="."):	1✔
268	"""
269	Write the packmol input file to the input directory.
270
271	Args:
272	input_dir (str): path to the input directory
273	"""
274	with open(os.path.join(input_dir, self.input_file), "w", encoding="utf-8") as inp:	×
275	for k, v in self.control_params.items():	×
276	inp.write(f"{k} {self._format_param_val(v)}\n")	×
277	# write the structures of the constituent molecules to file and set
278	# the molecule id and the corresponding filename in the packmol
279	# input file.
280	for idx, mol in enumerate(self.mols):	×
281	filename = os.path.join(input_dir, f"{idx}.{self.control_params['filetype']}")	×
282	# pdb
283	if self.control_params["filetype"] == "pdb":	×
284	self.write_pdb(mol, filename, num=idx + 1)	×
285	# all other filetypes
286	else:
287	a = BabelMolAdaptor(mol)	×
288	pm = pb.Molecule(a.openbabel_mol)	×
289	pm.write(	×
290	self.control_params["filetype"],
291	filename=filename,
292	overwrite=True,
293	)
294
295	inp.write("\n")	×
296	inp.write(f"structure {os.path.join(input_dir, str(idx))}.{self.control_params['filetype']}\n")	×
297	for k, v in self.param_list[idx].items():	×
298	inp.write(f" {k} {self._format_param_val(v)}\n")	×
299	inp.write("end structure\n")	×
300
301	def run(self, site_property=None):	1✔
302	"""
303	Write the input file to the scratch directory, run packmol and return
304	the packed molecule to the current working directory.
305
306	Args:
307	site_property (str): if set then the specified site property
308	for the final packed molecule will be restored.
309
310	Returns:
311	Molecule object
312	"""
313	with tempfile.TemporaryDirectory() as scratch_dir:	×
314	self._write_input(input_dir=scratch_dir)	×
315	with open(os.path.join(scratch_dir, self.input_file)) as packmol_input:	×
316	with Popen(self.packmol_bin, stdin=packmol_input, stdout=PIPE, stderr=PIPE) as p:	×
317	(stdout, stderr) = p.communicate()	×
318	output_file = os.path.join(self.control_params["output"])	×
319	if os.path.isfile(output_file):	×
320	packed_mol = BabelMolAdaptor.from_file(output_file, self.control_params["filetype"])	×
321	packed_mol = packed_mol.pymatgen_mol	×
322	print(f"packed molecule written to {self.control_params['output']}")	×
323	if site_property:	×
324	packed_mol = self.restore_site_properties(site_property=site_property, filename=output_file)	×
325	return packed_mol	×
326	raise RuntimeError(f"Packmol execution failed. {stdout}\n{stderr}")	×
327
328	@staticmethod	1✔
329	def write_pdb(mol, filename, name=None, num=None):	1✔
330	"""
331	dump the molecule into pdb file with custom residue name and number.
332	"""
333	# ugly hack to get around the openbabel issues with inconsistent
334	# residue labelling.
335	scratch = tempfile.gettempdir()	×
336	with ScratchDir(scratch, copy_to_current_on_exit=False) as _:	×
337	mol.to(fmt="pdb", filename="tmp.pdb")	×
338	bma = BabelMolAdaptor.from_file("tmp.pdb", "pdb")	×
339
340	num = num or 1	×
341	name = name or f"ml{num}"	×
342
343	# bma = BabelMolAdaptor(mol)
344	pbm = pb.Molecule(bma._obmol)	×
345	for x in pbm.residues:	×
346	x.OBResidue.SetName(name)	×
347	x.OBResidue.SetNum(num)	×
348
349	pbm.write(format="pdb", filename=filename, overwrite=True)	×
350
351	def _set_residue_map(self):	1✔
352	"""
353	map each residue to the corresponding molecule.
354	"""
355	self.map_residue_to_mol = {}	×
356	lookup = {}	×
357	for idx, mol in enumerate(self.mols):	×
358	if mol.formula not in lookup:	×
359	mol.translate_sites(indices=range(len(mol)), vector=-mol.center_of_mass)	×
360	lookup[mol.formula] = mol.copy()	×
361	self.map_residue_to_mol[f"ml{idx + 1}"] = lookup[mol.formula]	×
362
363	def convert_obatoms_to_molecule(self, atoms, residue_name=None, site_property="ff_map"):	1✔
364	"""
365	Convert list of openbabel atoms to MOlecule.
366
367	Args:
368	atoms ([OBAtom]): list of OBAtom objects
369	residue_name (str): the key in self.map_residue_to_mol. Usec to
370	restore the site properties in the final packed molecule.
371	site_property (str): the site property to be restored.
372
373	Returns:
374	Molecule object
375	"""
376	restore_site_props = residue_name is not None	×
377
378	if restore_site_props and not hasattr(self, "map_residue_to_mol"):	×
379	self._set_residue_map()	×
380
381	coords = []	×
382	zs = []	×
383	for atm in atoms:	×
384	coords.append(list(atm.coords))	×
385	zs.append(atm.atomicnum)	×
386
387	mol = Molecule(zs, coords)	×
388
389	if restore_site_props:	×
390	props = []	×
391
392	ref = self.map_residue_to_mol[residue_name].copy()	×
393
394	# sanity check
395	assert len(mol) == len(ref)	×
396	assert ref.formula == mol.formula	×
397
398	# the packed molecules have the atoms in the same order..sigh!
399	for i, site in enumerate(mol):	×
400	assert site.specie.symbol == ref[i].specie.symbol	×
401	props.append(getattr(ref[i], site_property))	×
402
403	mol.add_site_property(site_property, props)	×
404
405	return mol	×
406
407	def restore_site_properties(self, site_property="ff_map", filename=None):	1✔
408	"""
409	Restore the site properties for the final packed molecule.
410
411	Args:
412	site_property (str):
413	filename (str): path to the final packed molecule.
414
415	Returns:
416	Molecule
417	"""
418	# only for pdb
419	if not self.control_params["filetype"] == "pdb":	×
420	raise ValueError()	×
421
422	filename = filename or self.control_params["output"]	×
423	bma = BabelMolAdaptor.from_file(filename, "pdb")	×
424	pbm = pb.Molecule(bma._obmol)	×
425
426	assert len(pbm.residues) == sum(x["number"] for x in self.param_list)	×
427
428	packed_mol = self.convert_obatoms_to_molecule(	×
429	pbm.residues[0].atoms,
430	residue_name=pbm.residues[0].name,
431	site_property=site_property,
432	)
433
434	for resid in pbm.residues[1:]:	×
435	mol = self.convert_obatoms_to_molecule(resid.atoms, residue_name=resid.name, site_property=site_property)	×
436	for site in mol:	×
437	packed_mol.append(site.species, site.coords, properties=site.properties)	×
438
439	return packed_mol	×
440
441
442	class LammpsRunner:	1✔
443	"""
444	LAMMPS wrapper
445	"""
446
447	def __init__(self, input_filename="lammps.in", bin="lammps"):	1✔
448	"""
449	Args:
450	input_filename (str): input file name
451	bin (str): command to run, excluding the input file name
452	"""
453	self.lammps_bin = bin.split()	×
454	if not which(self.lammps_bin[-1]):	×
455	raise RuntimeError(	×
456	f"LammpsRunner requires the executable {self.lammps_bin[-1]} to be in the path. "
457	"Please download and install LAMMPS from "
458	"https://www.lammps.org/. "
459	"Don't forget to add the binary to your path"
460	)
461	self.input_filename = input_filename	×
462
463	def run(self):	1✔
464	"""
465	Write the input/data files and run LAMMPS.
466	"""
467	lammps_cmd = self.lammps_bin + ["-in", self.input_filename]	×
468	print(f"Running: {' '.join(lammps_cmd)}")	×
469	with Popen(lammps_cmd, stdout=PIPE, stderr=PIPE) as p:	×
470	(stdout, stderr) = p.communicate()	×
471	return stdout, stderr	×
472
473
474	if __name__ == "__main__":	1✔
475	ethanol_coords = [	×
476	[0.00720, -0.56870, 0.00000],
477	[-1.28540, 0.24990, 0.00000],
478	[1.13040, 0.31470, 0.00000],
479	[0.03920, -1.19720, 0.89000],
480	[0.03920, -1.19720, -0.89000],
481	[-1.31750, 0.87840, 0.89000],
482	[-1.31750, 0.87840, -0.89000],
483	[-2.14220, -0.42390, -0.00000],
484	[1.98570, -0.13650, -0.00000],
485	]
486	ethanol = Molecule(["C", "C", "O", "H", "H", "H", "H", "H", "H"], ethanol_coords)	×
487	water_coords = [	×
488	[9.626, 6.787, 12.673],
489	[9.626, 8.420, 12.673],
490	[10.203, 7.604, 12.673],
491	]
492	water = Molecule(["H", "H", "O"], water_coords)	×
493	pmr = PackmolRunner(	×
494	[ethanol, water],
495	[
496	{"number": 1, "fixed": [0, 0, 0, 0, 0, 0], "centerofmass": ""},
497	{"number": 15, "inside sphere": [0, 0, 0, 5]},
498	],
499	input_file="packmol_input.inp",
500	tolerance=2.0,
501	filetype="xyz",
502	control_params={"nloop": 1000},
503	auto_box=False,
504	output_file="cocktail.xyz",
505	)
506	s = pmr.run()	×

materialsproject / pymatgen / 4075885785

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous