<commit_before># coding: utf-8


import os
import unittest
import time
from pymatgen.core.structure import Molecule
from pymatgen.analysis.graphs import MoleculeGraph
from pymatgen.analysis.local_env import OpenBabelNN
from pymatgen.util.testing import PymatgenTest
from pymatgen.analysis.fragmenter import Fragmenter, metal_edge_extender

try:
    import openbabel as ob
except ImportError:
    ob = None

__author__ = "<NAME>"
__email__ = "<EMAIL>"

test_dir = os.path.join(os.path.dirname(__file__), "..", "..", "..",
                        'test_files', 'fragmenter_files')


class TestFragmentMolecule(PymatgenTest):
    @classmethod
    def setUpClass(cls):
        cls.pc = Molecule.from_file(
            os.path.join(test_dir, "PC.xyz"))
        cls.ec = Molecule.from_file(
            os.path.join(test_dir, "EC.xyz"))
        cls.pos_pc = Molecule.from_file(
            os.path.join(test_dir, "PC.xyz"))
        cls.pos_pc.set_charge_and_spin(charge=1)
        cls.pc_edges = [[5, 10], [5, 12], [5, 11], [5, 3], [3, 7], [3, 4],
                        [3, 0], [4, 8], [4, 9], [4, 1], [6, 1], [6, 0], [6, 2]]
        cls.pc_frag1 = Molecule.from_file(
            os.path.join(test_dir, "PC_frag1.xyz"))
        cls.pc_frag1_edges = [[0, 2], [4, 2], [2, 1], [1, 3]]
        cls.tfsi = Molecule.from_file(os.path.join(test_dir, "TFSI.xyz"))
        cls.tfsi_edges = [14, 1], [1, 4], [1, 5], [1, 7], [7, 11], [7, 12], [7, 13], [14, 0], [0, 2], [0, 3], [0, 6], [
            6, 8], [6, 9], [6, 10]
        cls.LiEC = Molecule.from_file(
            os.path.join(test_dir, "LiEC.xyz"))

    def _test_edges_given_PC_frag1(self):
        fragmenter = Fragmenter(molecule=self.pc_frag1, edges=self.pc_frag1_edges, depth=0)
        self.assertEqual(fragmenter.total_unique_fragments, 12)

    @unittest.skipIf(not ob, "OpenBabel not present. Skipping...")
    def _test_babel_PC_frag1(self):
        fragmenter = Fragmenter(molecule=self.pc_frag1, depth=0)
        self.assertEqual(fragmenter.total_unique_fragments, 12)

    @unittest.skipIf(not ob, "OpenBabel not present. Skipping...")
    def _test_babel_PC_old_defaults(self):
        fragmenter = Fragmenter(molecule=self.pc,open_rings=True)
        self.assertEqual(fragmenter.open_rings,True)
        self.assertEqual(fragmenter.opt_steps,10000)
        default_mol_graph = MoleculeGraph.with_local_env_strategy(self.pc, OpenBabelNN(),
                                                                  reorder=False, extend_structure=False)
        self.assertEqual(fragmenter.mol_graph,default_mol_graph)
        self.assertEqual(fragmenter.total_unique_fragments, 13)

    @unittest.skipIf(not ob, "OpenBabel not present. Skipping...")
    def _test_babel_PC_defaults(self):
        fragmenter = Fragmenter(molecule=self.pc)
        self.assertEqual(fragmenter.open_rings,False)
        self.assertEqual(fragmenter.opt_steps,10000)
        default_mol_graph = MoleculeGraph.with_local_env_strategy(self.pc, OpenBabelNN(),
                                                                  reorder=False, extend_structure=False)
        self.assertEqual(fragmenter.mol_graph,default_mol_graph)
        self.assertEqual(fragmenter.total_unique_fragments, 8)

    def _test_edges_given_PC_not_defaults(self):
        fragmenter = Fragmenter(molecule=self.pc, edges=self.pc_edges, depth=2, open_rings=False, opt_steps=0)
        self.assertEqual(fragmenter.open_rings, False)
        self.assertEqual(fragmenter.opt_steps, 0)
        edges = {(e[0], e[1]): None for e in self.pc_edges}
        default_mol_graph = MoleculeGraph.with_edges(self.pc, edges=edges)
        self.assertEqual(fragmenter.mol_graph,default_mol_graph)
        self.assertEqual(fragmenter.total_unique_fragments, 20)

    def _test_edges_given_TFSI(self):
        fragmenter = Fragmenter(molecule=self.tfsi, edges=self.tfsi_edges, depth=0)
        self.assertEqual(fragmenter.total_unique_fragments, 156)

    @unittest.skipIf(not ob, "OpenBabel not present. Skipping...")
    def _test_babel_TFSI(self):
        fragmenter = Fragmenter(molecule=self.tfsi, depth=0)
        self.assertEqual(fragmenter.total_unique_fragments, 156)

    @unittest.skipIf(not ob, "OpenBabel not present. Skipping...")
    def _test_babel_PC_with_RO_depth_0_vs_depth_10(self):
        fragmenter0 = Fragmenter(molecule=self.pc, depth=0, open_rings=True, opt_steps=1000)
        self.assertEqual(fragmenter0.total_unique_fragments, 509)

        fragmenter10 = Fragmenter(molecule=self.pc, depth=10, open_rings=True, opt_steps=1000)
        self.assertEqual(fragmenter10.total_unique_fragments, 509)

        fragments_by_level = fragmenter10.fragments_by_level
        num_frags_by_level = [13, 51, 95, 115, 105, 75, 39, 14, 2, 0]
        for ii in range(10):
            num_frags = 0
            for key in fragments_by_level[str(ii)]:
                num_frags += len(fragments_by_level[str(ii)][key])
            self.assertEqual(num_frags,num_frags_by_level[ii])

    def _test_PC_depth_0_vs_depth_10(self):
        fragmenter0 = Fragmenter(molecule=self.pc, edges=self.pc_edges, depth=0, open_rings=False)
        self.assertEqual(fragmenter0.total_unique_fragments, 295)

        fragmenter10 = Fragmenter(molecule=self.pc, edges=self.pc_edges, depth=10, open_rings=False)
        self.assertEqual(fragmenter10.total_unique_fragments, 63)

        fragments_by_level = fragmenter10.fragments_by_level
        num_frags_by_level = [8, 12, 15, 14, 9, 4, 1]
        for ii in range(7):
            num_frags = 0
            for key in fragments_by_level[str(ii)]:
                num_frags += len(fragments_by_level[str(ii)][key])
            self.assertEqual(num_frags,num_frags_by_level[ii])

    def _test_PC_frag1_then_PC(self):
        frag1 = Fragmenter(molecule=self.pc_frag1, edges=self.pc_frag1_edges, depth=0)
        self.assertEqual(frag1.new_unique_fragments,frag1.total_unique_fragments)
        frag2 = Fragmenter(molecule=self.pc, edges=self.pc_edges, depth=0, open_rings=False, prev_unique_frag_dict=frag1.unique_frag_dict)
        self.assertEqual(frag2.new_unique_fragments,295-12)

    @unittest.skipIf(not ob, "OpenBabel not present. Skipping...")
    def _test_PC_then_EC_depth_0(self):
        fragPC = Fragmenter(molecule=self.pc, depth=0, open_rings=True)
        fragEC = Fragmenter(molecule=self.ec, depth=0, open_rings=True, prev_unique_frag_dict=fragPC.unique_frag_dict)
        self.assertEqual(fragEC.new_unique_fragments,11)
        self.assertEqual(fragEC.total_unique_fragments,509+11)

    @unittest.skipIf(not ob, "OpenBabel not present. Skipping...")
    def _test_PC_then_EC_depth_10(self):
        fragPC = Fragmenter(molecule=self.pc, depth=10, open_rings=True)
        fragEC = Fragmenter(molecule=self.ec, depth=10, open_rings=True, prev_unique_frag_dict=fragPC.unique_frag_dict)
        self.assertEqual(fragEC.new_unique_fragments,11)
        self.assertEqual(fragEC.total_unique_fragments,509+11)
<commit_msg>Turn back on rest of fragmenter tests<commit_after># coding: utf-8


import os
import unittest
import time
from pymatgen.core.structure import Molecule
from pymatgen.analysis.graphs import MoleculeGraph
from pymatgen.analysis.local_env import OpenBabelNN
from pymatgen.util.testing import PymatgenTest
from pymatgen.analysis.fragmenter import Fragmenter, metal_edge_extender

try:
    import openbabel as ob
except ImportError:
    ob = None

__author__ = "<NAME>"
__email__ = "<EMAIL>"

test_dir = os.path.join(os.path.dirname(__file__), "..", "..", "..",
                        'test_files', 'fragmenter_files')


class TestFragmentMolecule(PymatgenTest):
    @classmethod
    def setUpClass(cls):
        cls.pc = Molecule.from_file(
            os.path.join(test_dir, "PC.xyz"))
        cls.ec = Molecule.from_file(
            os.path.join(test_dir, "EC.xyz"))
        cls.pos_pc = Molecule.from_file(
            os.path.join(test_dir, "PC.xyz"))
        cls.pos_pc.set_charge_and_spin(charge=1)
        cls.pc_edges = [[5, 10], [5, 12], [5, 11], [5, 3], [3, 7], [3, 4],
                        [3, 0], [4, 8], [4, 9], [4, 1], [6, 1], [6, 0], [6, 2]]
        cls.pc_frag1 = Molecule.from_file(
            os.path.join(test_dir, "PC_frag1.xyz"))
        cls.pc_frag1_edges = [[0, 2], [4, 2], [2, 1], [1, 3]]
        cls.tfsi = Molecule.from_file(os.path.join(test_dir, "TFSI.xyz"))
        cls.tfsi_edges = [14, 1], [1, 4], [1, 5], [1, 7], [7, 11], [7, 12], [7, 13], [14, 0], [0, 2], [0, 3], [0, 6], [
            6, 8], [6, 9], [6, 10]
        cls.LiEC = Molecule.from_file(
            os.path.join(test_dir, "LiEC.xyz"))

    def test_edges_given_PC_frag1(self):
        fragmenter = Fragmenter(molecule=self.pc_frag1, edges=self.pc_frag1_edges, depth=0)
        self.assertEqual(fragmenter.total_unique_fragments, 12)

    @unittest.skipIf(not ob, "OpenBabel not present. Skipping...")
    def test_babel_PC_frag1(self):
        fragmenter = Fragmenter(molecule=self.pc_frag1, depth=0)
        self.assertEqual(fragmenter.total_unique_fragments, 12)

    @unittest.skipIf(not ob, "OpenBabel not present. Skipping...")
    def test_babel_PC_old_defaults(self):
        fragmenter = Fragmenter(molecule=self.pc,open_rings=True)
        self.assertEqual(fragmenter.open_rings,True)
        self.assertEqual(fragmenter.opt_steps,10000)
        default_mol_graph = MoleculeGraph.with_local_env_strategy(self.pc, OpenBabelNN(),
                                                                  reorder=False, extend_structure=False)
        self.assertEqual(fragmenter.mol_graph,default_mol_graph)
        self.assertEqual(fragmenter.total_unique_fragments, 13)

    @unittest.skipIf(not ob, "OpenBabel not present. Skipping...")
    def test_babel_PC_defaults(self):
        fragmenter = Fragmenter(molecule=self.pc)
        self.assertEqual(fragmenter.open_rings,False)
        self.assertEqual(fragmenter.opt_steps,10000)
        default_mol_graph = MoleculeGraph.with_local_env_strategy(self.pc, OpenBabelNN(),
                                                                  reorder=False, extend_structure=False)
        self.assertEqual(fragmenter.mol_graph,default_mol_graph)
        self.assertEqual(fragmenter.total_unique_fragments, 8)

    def test_edges_given_PC_not_defaults(self):
        fragmenter = Fragmenter(molecule=self.pc, edges=self.pc_edges, depth=2, open_rings=False, opt_steps=0)
        self.assertEqual(fragmenter.open_rings, False)
        self.assertEqual(fragmenter.opt_steps, 0)
        edges = {(e[0], e[1]): None for e in self.pc_edges}
        default_mol_graph = MoleculeGraph.with_edges(self.pc, edges=edges)
        self.assertEqual(fragmenter.mol_graph,default_mol_graph)
        self.assertEqual(fragmenter.total_unique_fragments, 20)

    def test_edges_given_TFSI(self):
        fragmenter = Fragmenter(molecule=self.tfsi, edges=self.tfsi_edges, depth=0)
        self.assertEqual(fragmenter.total_unique_fragments, 156)

    @unittest.skipIf(not ob, "OpenBabel not present. Skipping...")
    def test_babel_TFSI(self):
        fragmenter = Fragmenter(molecule=self.tfsi, depth=0)
        self.assertEqual(fragmenter.total_unique_fragments, 156)

    @unittest.skipIf(not ob, "OpenBabel not present. Skipping...")
    def test_babel_PC_with_RO_depth_0_vs_depth_10(self):
        fragmenter0 = Fragmenter(molecule=self.pc, depth=0, open_rings=True, opt_steps=1000)
        self.assertEqual(fragmenter0.total_unique_fragments, 509)

        fragmenter10 = Fragmenter(molecule=self.pc, depth=10, open_rings=True, opt_steps=1000)
        self.assertEqual(fragmenter10.total_unique_fragments, 509)

        fragments_by_level = fragmenter10.fragments_by_level
        num_frags_by_level = [13, 51, 95, 115, 105, 75, 39, 14, 2, 0]
        for ii in range(10):
            num_frags = 0
            for key in fragments_by_level[str(ii)]:
                num_frags += len(fragments_by_level[str(ii)][key])
            self.assertEqual(num_frags,num_frags_by_level[ii])

    def test_PC_depth_0_vs_depth_10(self):
        fragmenter0 = Fragmenter(molecule=self.pc, edges=self.pc_edges, depth=0, open_rings=False)
        self.assertEqual(fragmenter0.total_unique_fragments, 295)

        fragmenter10 = Fragmenter(molecule=self.pc, edges=self.pc_edges, depth=10, open_rings=False)
        self.assertEqual(fragmenter10.total_unique_fragments, 63)

        fragments_by_level = fragmenter10.fragments_by_level
        num_frags_by_level = [8, 12, 15, 14, 9, 4, 1]
        for ii in range(7):
            num_frags = 0
            for key in fragments_by_level[str(ii)]:
                num_frags += len(fragments_by_level[str(ii)][key])
            self.assertEqual(num_frags,num_frags_by_level[ii])

    def test_PC_frag1_then_PC(self):
        frag1 = Fragmenter(molecule=self.pc_frag1, edges=self.pc_frag1_edges, depth=0)
        self.assertEqual(frag1.new_unique_fragments,frag1.total_unique_fragments)
        frag2 = Fragmenter(molecule=self.pc, edges=self.pc_edges, depth=0, open_rings=False, prev_unique_frag_dict=frag1.unique_frag_dict)
        self.assertEqual(frag2.new_unique_fragments,295-12)

    @unittest.skipIf(not ob, "OpenBabel not present. Skipping...")
    def test_PC_then_EC_depth_0(self):
        fragPC = Fragmenter(molecule=self.pc, depth=0, open_rings=True)
        fragEC = Fragmenter(molecule=self.ec, depth=0, open_rings=True, prev_unique_frag_dict=fragPC.unique_frag_dict)
        self.assertEqual(fragEC.new_unique_fragments,11)
        self.assertEqual(fragEC.total_unique_fragments,509+11)

    @unittest.skipIf(not ob, "OpenBabel not present. Skipping...")
    def test_PC_then_EC_depth_10(self):
        fragPC = Fragmenter(molecule=self.pc, depth=10, open_rings=True)
        fragEC = Fragmenter(molecule=self.ec, depth=10, open_rings=True, prev_unique_frag_dict=fragPC.unique_frag_dict)
        self.assertEqual(fragEC.new_unique_fragments,11)
        self.assertEqual(fragEC.total_unique_fragments,509+11)

{'fraction_non_alphanumeric': 0.1173688736027515, 'fraction_numerical': 0.033820578962453426, 'mean_word_length': 3.3738639924788467, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 28, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5925409', 'n_tokens_mistral': 5194, 'n_tokens_neox': 4727, 'n_words': 782}