1# coding: utf-8
2
3
4import os
5import unittest
6import pytest
7
8from pymatgen.analysis.fragmenter import Fragmenter
9from pymatgen.analysis.graphs import MoleculeGraph
10from pymatgen.analysis.local_env import OpenBabelNN, metal_edge_extender
11from pymatgen.core.structure import Molecule
12from pymatgen.util.testing import PymatgenTest
13
14__author__ = "Samuel Blau"
15__email__ = "samblau1@gmail.com"
16
17
18test_dir = os.path.join(PymatgenTest.TEST_FILES_DIR, "fragmenter_files")
19
20
21class TestFragmentMolecule(PymatgenTest):
22    @classmethod
23    def setUpClass(cls):
24        cls.pc = Molecule.from_file(os.path.join(test_dir, "PC.xyz"))
25        cls.ec = Molecule.from_file(os.path.join(test_dir, "EC.xyz"))
26        cls.pos_pc = Molecule.from_file(os.path.join(test_dir, "PC.xyz"))
27        cls.pos_pc.set_charge_and_spin(charge=1)
28        cls.pc_edges = [
29            [5, 10],
30            [5, 12],
31            [5, 11],
32            [5, 3],
33            [3, 7],
34            [3, 4],
35            [3, 0],
36            [4, 8],
37            [4, 9],
38            [4, 1],
39            [6, 1],
40            [6, 0],
41            [6, 2],
42        ]
43        cls.pc_frag1 = Molecule.from_file(os.path.join(test_dir, "PC_frag1.xyz"))
44        cls.pc_frag1_edges = [[0, 2], [4, 2], [2, 1], [1, 3]]
45        cls.tfsi = Molecule.from_file(os.path.join(test_dir, "TFSI.xyz"))
46        cls.tfsi_edges = (
47            [14, 1],
48            [1, 4],
49            [1, 5],
50            [1, 7],
51            [7, 11],
52            [7, 12],
53            [7, 13],
54            [14, 0],
55            [0, 2],
56            [0, 3],
57            [0, 6],
58            [6, 8],
59            [6, 9],
60            [6, 10],
61        )
62        cls.LiEC = Molecule.from_file(os.path.join(test_dir, "LiEC.xyz"))
63
64    def test_edges_given_PC_frag1(self):
65        fragmenter = Fragmenter(molecule=self.pc_frag1, edges=self.pc_frag1_edges, depth=0)
66        self.assertEqual(fragmenter.total_unique_fragments, 12)
67
68    def test_babel_PC_frag1(self):
69        pytest.importorskip("openbabel", reason="OpenBabel not installed")
70        fragmenter = Fragmenter(molecule=self.pc_frag1, depth=0)
71        self.assertEqual(fragmenter.total_unique_fragments, 12)
72
73    def test_babel_PC_old_defaults(self):
74        pytest.importorskip("openbabel", reason="OpenBabel not installed")
75        fragmenter = Fragmenter(molecule=self.pc, open_rings=True)
76        self.assertEqual(fragmenter.open_rings, True)
77        self.assertEqual(fragmenter.opt_steps, 10000)
78        default_mol_graph = MoleculeGraph.with_local_env_strategy(self.pc, OpenBabelNN())
79        self.assertEqual(fragmenter.mol_graph, default_mol_graph)
80        self.assertEqual(fragmenter.total_unique_fragments, 13)
81
82    def test_babel_PC_defaults(self):
83        pytest.importorskip("openbabel", reason="OpenBabel not installed")
84        fragmenter = Fragmenter(molecule=self.pc)
85        self.assertEqual(fragmenter.open_rings, False)
86        self.assertEqual(fragmenter.opt_steps, 10000)
87        default_mol_graph = MoleculeGraph.with_local_env_strategy(self.pc, OpenBabelNN())
88        self.assertEqual(fragmenter.mol_graph, default_mol_graph)
89        self.assertEqual(fragmenter.total_unique_fragments, 8)
90
91    def test_edges_given_PC_not_defaults(self):
92        fragmenter = Fragmenter(
93            molecule=self.pc,
94            edges=self.pc_edges,
95            depth=2,
96            open_rings=False,
97            opt_steps=0,
98        )
99        self.assertEqual(fragmenter.open_rings, False)
100        self.assertEqual(fragmenter.opt_steps, 0)
101        edges = {(e[0], e[1]): None for e in self.pc_edges}
102        default_mol_graph = MoleculeGraph.with_edges(self.pc, edges=edges)
103        self.assertEqual(fragmenter.mol_graph, default_mol_graph)
104        self.assertEqual(fragmenter.total_unique_fragments, 20)
105
106    def test_edges_given_TFSI(self):
107        fragmenter = Fragmenter(molecule=self.tfsi, edges=self.tfsi_edges, depth=0)
108        self.assertEqual(fragmenter.total_unique_fragments, 156)
109
110    def test_babel_TFSI(self):
111        pytest.importorskip("openbabel", reason="OpenBabel not installed")
112        fragmenter = Fragmenter(molecule=self.tfsi, depth=0)
113        self.assertEqual(fragmenter.total_unique_fragments, 156)
114
115    def test_babel_PC_with_RO_depth_0_vs_depth_10(self):
116        pytest.importorskip("openbabel", reason="OpenBabel not installed")
117        fragmenter0 = Fragmenter(molecule=self.pc, depth=0, open_rings=True, opt_steps=1000)
118        self.assertEqual(fragmenter0.total_unique_fragments, 509)
119
120        fragmenter10 = Fragmenter(molecule=self.pc, depth=10, open_rings=True, opt_steps=1000)
121        self.assertEqual(fragmenter10.total_unique_fragments, 509)
122
123        fragments_by_level = fragmenter10.fragments_by_level
124        num_frags_by_level = [13, 51, 95, 115, 105, 75, 39, 14, 2, 0]
125        for ii in range(10):
126            num_frags = 0
127            for key in fragments_by_level[str(ii)]:
128                num_frags += len(fragments_by_level[str(ii)][key])
129            self.assertEqual(num_frags, num_frags_by_level[ii])
130
131    def test_PC_depth_0_vs_depth_10(self):
132        fragmenter0 = Fragmenter(molecule=self.pc, edges=self.pc_edges, depth=0, open_rings=False)
133        self.assertEqual(fragmenter0.total_unique_fragments, 295)
134
135        fragmenter10 = Fragmenter(molecule=self.pc, edges=self.pc_edges, depth=10, open_rings=False)
136        self.assertEqual(fragmenter10.total_unique_fragments, 63)
137
138        fragments_by_level = fragmenter10.fragments_by_level
139        num_frags_by_level = [8, 12, 15, 14, 9, 4, 1]
140        for ii in range(7):
141            num_frags = 0
142            for key in fragments_by_level[str(ii)]:
143                num_frags += len(fragments_by_level[str(ii)][key])
144            self.assertEqual(num_frags, num_frags_by_level[ii])
145
146    def test_PC_frag1_then_PC(self):
147        frag1 = Fragmenter(molecule=self.pc_frag1, edges=self.pc_frag1_edges, depth=0)
148        self.assertEqual(frag1.new_unique_fragments, frag1.total_unique_fragments)
149        frag2 = Fragmenter(
150            molecule=self.pc,
151            edges=self.pc_edges,
152            depth=0,
153            open_rings=False,
154            prev_unique_frag_dict=frag1.unique_frag_dict,
155        )
156        self.assertEqual(frag2.new_unique_fragments, 295 - 12)
157
158    def test_PC_then_EC_depth_10(self):
159        pytest.importorskip("openbabel", reason="OpenBabel not installed")
160        fragPC = Fragmenter(molecule=self.pc, depth=10, open_rings=True)
161        fragEC = Fragmenter(
162            molecule=self.ec,
163            depth=10,
164            open_rings=True,
165            prev_unique_frag_dict=fragPC.unique_frag_dict,
166        )
167        self.assertEqual(fragEC.new_unique_fragments, 11)
168        self.assertEqual(fragEC.total_unique_fragments, 509 + 11)
169
170
171if __name__ == "__main__":
172    unittest.main()
173