<commit_before># coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.


import unittest
import os
import copy

from monty.serialization import loadfn#, dumpfn

from pymatgen.command_line.critic2_caller import Critic2Output
from pymatgen.core.structure import Molecule, Structure, FunctionalGroups, Site
from pymatgen.analysis.graphs import *
from pymatgen.analysis.local_env import MinimumDistanceNN, MinimumOKeeffeNN, OpenBabelNN

try:
    import openbabel as ob
except ImportError:
    ob = None
try:
    import networkx as nx
    import networkx.algorithms.isomorphism as iso
except ImportError:
    nx = None

__author__ = "<NAME>, <NAME>"
__version__ = "0.1"
__maintainer__ = "<NAME>"
__email__ = "<EMAIL>"
__status__ = "Beta"
__date__ = "August 2017"

module_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)))


class StructureGraphTest(unittest.TestCase):

    def setUp(self):

        self.maxDiff = None

        # trivial example, simple square lattice for testing
        structure = Structure(Lattice.tetragonal(5.0, 50.0), ['H'], [[0, 0, 0]])
        self.square_sg = StructureGraph.with_empty_graph(structure, edge_weight_name="", edge_weight_units="")
        self.square_sg.add_edge(0, 0, from_jimage=(0, 0, 0), to_jimage=(1, 0, 0))
        self.square_sg.add_edge(0, 0, from_jimage=(0, 0, 0), to_jimage=(-1, 0, 0))
        self.square_sg.add_edge(0, 0, from_jimage=(0, 0, 0), to_jimage=(0, 1, 0))
        self.square_sg.add_edge(0, 0, from_jimage=(0, 0, 0), to_jimage=(0, -1, 0))
        #TODO: decorating still fails because the structure graph gives a CN of 8 for this square lattice
        # self.square_sg.decorate_structure_with_ce_info()

        # body-centered square lattice for testing
        structure = Structure(Lattice.tetragonal(5.0, 50.0), ['H', 'He'], [[0, 0, 0], [0.5, 0.5, 0.5]])
        self.bc_square_sg = StructureGraph.with_empty_graph(structure, edge_weight_name="", edge_weight_units="")
        self.bc_square_sg.add_edge(0, 0, from_jimage=(0, 0, 0), to_jimage=(1, 0, 0))
        self.bc_square_sg.add_edge(0, 0, from_jimage=(0, 0, 0), to_jimage=(-1, 0, 0))
        self.bc_square_sg.add_edge(0, 0, from_jimage=(0, 0, 0), to_jimage=(0, 1, 0))
        self.bc_square_sg.add_edge(0, 0, from_jimage=(0, 0, 0), to_jimage=(0, -1, 0))
        self.bc_square_sg.add_edge(0, 1, from_jimage=(0, 0, 0), to_jimage=(0, 0, 0))
        self.bc_square_sg.add_edge(0, 1, from_jimage=(0, 0, 0), to_jimage=(-1, 0, 0))
        self.bc_square_sg.add_edge(0, 1, from_jimage=(0, 0, 0), to_jimage=(-1, -1, 0))
        self.bc_square_sg.add_edge(0, 1, from_jimage=(0, 0, 0), to_jimage=(0, -1, 0))

        # body-centered square lattice for testing
        # directions reversed, should be equivalent to as bc_square
        structure = Structure(Lattice.tetragonal(5.0, 50.0), ['H', 'He'], [[0, 0, 0], [0.5, 0.5, 0.5]])
        self.bc_square_sg_r = StructureGraph.with_empty_graph(structure, edge_weight_name="", edge_weight_units="")
        self.bc_square_sg_r.add_edge(0, 0, from_jimage=(0, 0, 0), to_jimage=(1, 0, 0))
        self.bc_square_sg_r.add_edge(0, 0, from_jimage=(0, 0, 0), to_jimage=(-1, 0, 0))
        self.bc_square_sg_r.add_edge(0, 0, from_jimage=(0, 0, 0), to_jimage=(0, 1, 0))
        self.bc_square_sg_r.add_edge(0, 0, from_jimage=(0, 0, 0), to_jimage=(0, -1, 0))
        self.bc_square_sg_r.add_edge(0, 1, from_jimage=(0, 0, 0), to_jimage=(0, 0, 0))
        self.bc_square_sg_r.add_edge(1, 0, from_jimage=(-1, 0, 0), to_jimage=(0, 0, 0))
        self.bc_square_sg_r.add_edge(1, 0, from_jimage=(-1, -1, 0), to_jimage=(0, 0, 0))
        self.bc_square_sg_r.add_edge(1, 0, from_jimage=(0, -1, 0), to_jimage=(0, 0, 0))

        # MoS2 example, structure graph obtained from critic2
        # (not ground state, from mp-1023924, single layer)
        stdout_file = os.path.join(os.path.dirname(__file__), "..", "..", "..",
                                   'test_files/critic2/MoS2_critic2_stdout.txt')
        with open(stdout_file, 'r') as f:
            reference_stdout = f.read()
        self.structure = Structure.from_file(os.path.join(os.path.dirname(__file__), "..", "..", "..",
                                             'test_files/critic2/MoS2.cif'))
        c2o = Critic2Output(self.structure, reference_stdout)
        self.mos2_sg = c2o.structure_graph(edge_weight="bond_length", edge_weight_units="Å")

        latt = Lattice.cubic(4.17)
        species = ["Ni", "O"]
        coords = [[0, 0, 0],
                  [0.5, 0.5, 0.5]]
        self.NiO = Structure.from_spacegroup(225, latt, species, coords).get_primitive_structure()

        # BCC example.
        self.bcc = Structure(Lattice.cubic(5.0), ['He', 'He'], [[0, 0, 0], [0.5, 0.5, 0.5]])


        warnings.simplefilter("ignore")

    def tearDown(self):
        warnings.simplefilter("default")

    def test_properties(self):

        self.assertEqual(self.mos2_sg.name, "bonds")
        self.assertEqual(self.mos2_sg.edge_weight_name, "bond_length")
        self.assertEqual(self.mos2_sg.edge_weight_unit, "Å")
        self.assertEqual(self.mos2_sg.get_coordination_of_site(0), 6)
        self.assertEqual(len(self.mos2_sg.get_connected_sites(0)), 6)
        self.assertTrue(isinstance(self.mos2_sg.get_connected_sites(0)[0].site, PeriodicSite))
        self.assertEqual(str(self.mos2_sg.get_connected_sites(0)[0].site.specie), 'S')
        self.assertAlmostEqual(
            self.mos2_sg.get_connected_sites(0, jimage=(0, 0, 100))[0].site.frac_coords[2],
            100.303027
        )

        # these two graphs should be equivalent
        for n in range(len(self.bc_square_sg)):
            self.assertEqual(self.bc_square_sg.get_coordination_of_site(n),
                             self.bc_square_sg_r.get_coordination_of_site(n))

    @unittest.skipIf(not nx, "NetworkX not present. Skipping...")
    def test_set_node_attributes(self):
        self.square_sg.set_node_attributes()

        specie = nx.get_node_attributes(self.square_sg.graph, "specie")
        coords = nx.get_node_attributes(self.square_sg.graph, "coords")
        properties = nx.get_node_attributes(self.square_sg.graph, "properties")

        for i in range(len(self.square_sg.structure)):
            self.assertEqual(str(specie[i]), str(self.square_sg.structure[i].specie))
            self.assertEqual(coords[i][0], self.square_sg.structure[i].coords[0])
            self.assertEqual(coords[i][1], self.square_sg.structure[i].coords[1])
            self.assertEqual(coords[i][2], self.square_sg.structure[i].coords[2])
            self.assertEqual(properties[i], self.square_sg.structure[i].properties)

    def test_edge_editing(self):
        square = copy.deepcopy(self.square_sg)

        square.alter_edge(0, 0, to_jimage=(1, 0, 0),  new_weight=0.0,
                                  new_edge_properties={"foo": "bar"})
        new_edge = square.graph.get_edge_data(0, 0)[0]
        self.assertEqual(new_edge["weight"], 0.0)
        self.assertEqual(new_edge["foo"], "bar")

        square.break_edge(0, 0, to_jimage=(1, 0, 0))
        self.assertEqual(len(square.graph.get_edge_data(0, 0)), 3)

    def test_insert_remove(self):
        struct_copy = copy.deepcopy(self.square_sg.structure)
        square_copy = copy.deepcopy(self.square_sg)

        # Ensure that insert_node appropriately wraps Structure.insert()
        struct_copy.insert(1, "O", [0.5, 0.5, 0.5])
        square_copy.insert_node(1, "O", [0.5, 0.5, 0.5])
        self.assertEqual(struct_copy, square_copy.structure)

        # Test that removal is also equivalent between Structure and StructureGraph.structure
        struct_copy.remove_sites([1])
        square_copy.remove_nodes([1])
        self.assertEqual(struct_copy, square_copy.structure)

        square_copy.insert_node(1, "O", [0.5, 0.5, 0.5], edges=[{"from_index": 1,
                                                                 "to_index": 0,
                                                                 "to_jimage": (0, 0, 0)}])
        self.assertEqual(square_copy.get_coordination_of_site(1), 1)

    def test_substitute(self):
        structure = Structure.from_file(os.path.join(os.path.dirname(__file__),
                                                     "..", "..", "..",
                                                     "test_files", "Li2O.cif"))
        molecule = FunctionalGroups["methyl"]

        structure_copy = copy.deepcopy(structure)
        structure_copy_graph = copy.deepcopy(structure)

        sg = StructureGraph.with_local_env_strategy(structure, MinimumDistanceNN())
        sg_copy = copy.deepcopy(sg)

        # Ensure that strings and molecules lead to equivalent substitutions<commit_msg>add test for checking if StructureGraph is correctly updated<commit_after># coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.


import unittest
import os
import copy

from monty.serialization import loadfn#, dumpfn

from pymatgen.command_line.critic2_caller import Critic2Output
from pymatgen.core.structure import Molecule, Structure, FunctionalGroups, Site
from pymatgen.analysis.graphs import *
from pymatgen.analysis.local_env import MinimumDistanceNN, MinimumOKeeffeNN, OpenBabelNN

try:
    import openbabel as ob
except ImportError:
    ob = None
try:
    import networkx as nx
    import networkx.algorithms.isomorphism as iso
except ImportError:
    nx = None

__author__ = "<NAME>, <NAME>"
__version__ = "0.1"
__maintainer__ = "<NAME>"
__email__ = "<EMAIL>"
__status__ = "Beta"
__date__ = "August 2017"

module_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)))


class StructureGraphTest(unittest.TestCase):

    def setUp(self):

        self.maxDiff = None

        # trivial example, simple square lattice for testing
        structure = Structure(Lattice.tetragonal(5.0, 50.0), ['H'], [[0, 0, 0]])
        self.square_sg = StructureGraph.with_empty_graph(structure, edge_weight_name="", edge_weight_units="")
        self.square_sg.add_edge(0, 0, from_jimage=(0, 0, 0), to_jimage=(1, 0, 0))
        self.square_sg.add_edge(0, 0, from_jimage=(0, 0, 0), to_jimage=(-1, 0, 0))
        self.square_sg.add_edge(0, 0, from_jimage=(0, 0, 0), to_jimage=(0, 1, 0))
        self.square_sg.add_edge(0, 0, from_jimage=(0, 0, 0), to_jimage=(0, -1, 0))
        #TODO: decorating still fails because the structure graph gives a CN of 8 for this square lattice
        # self.square_sg.decorate_structure_with_ce_info()

        # body-centered square lattice for testing
        structure = Structure(Lattice.tetragonal(5.0, 50.0), ['H', 'He'], [[0, 0, 0], [0.5, 0.5, 0.5]])
        self.bc_square_sg = StructureGraph.with_empty_graph(structure, edge_weight_name="", edge_weight_units="")
        self.bc_square_sg.add_edge(0, 0, from_jimage=(0, 0, 0), to_jimage=(1, 0, 0))
        self.bc_square_sg.add_edge(0, 0, from_jimage=(0, 0, 0), to_jimage=(-1, 0, 0))
        self.bc_square_sg.add_edge(0, 0, from_jimage=(0, 0, 0), to_jimage=(0, 1, 0))
        self.bc_square_sg.add_edge(0, 0, from_jimage=(0, 0, 0), to_jimage=(0, -1, 0))
        self.bc_square_sg.add_edge(0, 1, from_jimage=(0, 0, 0), to_jimage=(0, 0, 0))
        self.bc_square_sg.add_edge(0, 1, from_jimage=(0, 0, 0), to_jimage=(-1, 0, 0))
        self.bc_square_sg.add_edge(0, 1, from_jimage=(0, 0, 0), to_jimage=(-1, -1, 0))
        self.bc_square_sg.add_edge(0, 1, from_jimage=(0, 0, 0), to_jimage=(0, -1, 0))

        # body-centered square lattice for testing
        # directions reversed, should be equivalent to as bc_square
        structure = Structure(Lattice.tetragonal(5.0, 50.0), ['H', 'He'], [[0, 0, 0], [0.5, 0.5, 0.5]])
        self.bc_square_sg_r = StructureGraph.with_empty_graph(structure, edge_weight_name="", edge_weight_units="")
        self.bc_square_sg_r.add_edge(0, 0, from_jimage=(0, 0, 0), to_jimage=(1, 0, 0))
        self.bc_square_sg_r.add_edge(0, 0, from_jimage=(0, 0, 0), to_jimage=(-1, 0, 0))
        self.bc_square_sg_r.add_edge(0, 0, from_jimage=(0, 0, 0), to_jimage=(0, 1, 0))
        self.bc_square_sg_r.add_edge(0, 0, from_jimage=(0, 0, 0), to_jimage=(0, -1, 0))
        self.bc_square_sg_r.add_edge(0, 1, from_jimage=(0, 0, 0), to_jimage=(0, 0, 0))
        self.bc_square_sg_r.add_edge(1, 0, from_jimage=(-1, 0, 0), to_jimage=(0, 0, 0))
        self.bc_square_sg_r.add_edge(1, 0, from_jimage=(-1, -1, 0), to_jimage=(0, 0, 0))
        self.bc_square_sg_r.add_edge(1, 0, from_jimage=(0, -1, 0), to_jimage=(0, 0, 0))

        # MoS2 example, structure graph obtained from critic2
        # (not ground state, from mp-1023924, single layer)
        stdout_file = os.path.join(os.path.dirname(__file__), "..", "..", "..",
                                   'test_files/critic2/MoS2_critic2_stdout.txt')
        with open(stdout_file, 'r') as f:
            reference_stdout = f.read()
        self.structure = Structure.from_file(os.path.join(os.path.dirname(__file__), "..", "..", "..",
                                             'test_files/critic2/MoS2.cif'))
        c2o = Critic2Output(self.structure, reference_stdout)
        self.mos2_sg = c2o.structure_graph(edge_weight="bond_length", edge_weight_units="Å")

        latt = Lattice.cubic(4.17)
        species = ["Ni", "O"]
        coords = [[0, 0, 0],
                  [0.5, 0.5, 0.5]]
        self.NiO = Structure.from_spacegroup(225, latt, species, coords).get_primitive_structure()

        # BCC example.
        self.bcc = Structure(Lattice.cubic(5.0), ['He', 'He'], [[0, 0, 0], [0.5, 0.5, 0.5]])


        warnings.simplefilter("ignore")

    def tearDown(self):
        warnings.simplefilter("default")

    def test_properties(self):

        self.assertEqual(self.mos2_sg.name, "bonds")
        self.assertEqual(self.mos2_sg.edge_weight_name, "bond_length")
        self.assertEqual(self.mos2_sg.edge_weight_unit, "Å")
        self.assertEqual(self.mos2_sg.get_coordination_of_site(0), 6)
        self.assertEqual(len(self.mos2_sg.get_connected_sites(0)), 6)
        self.assertTrue(isinstance(self.mos2_sg.get_connected_sites(0)[0].site, PeriodicSite))
        self.assertEqual(str(self.mos2_sg.get_connected_sites(0)[0].site.specie), 'S')
        self.assertAlmostEqual(
            self.mos2_sg.get_connected_sites(0, jimage=(0, 0, 100))[0].site.frac_coords[2],
            100.303027
        )

        # these two graphs should be equivalent
        for n in range(len(self.bc_square_sg)):
            self.assertEqual(self.bc_square_sg.get_coordination_of_site(n),
                             self.bc_square_sg_r.get_coordination_of_site(n))

    @unittest.skipIf(not nx, "NetworkX not present. Skipping...")
    def test_set_node_attributes(self):
        self.square_sg.set_node_attributes()

        specie = nx.get_node_attributes(self.square_sg.graph, "specie")
        coords = nx.get_node_attributes(self.square_sg.graph, "coords")
        properties = nx.get_node_attributes(self.square_sg.graph, "properties")

        for i in range(len(self.square_sg.structure)):
            self.assertEqual(str(specie[i]), str(self.square_sg.structure[i].specie))
            self.assertEqual(coords[i][0], self.square_sg.structure[i].coords[0])
            self.assertEqual(coords[i][1], self.square_sg.structure[i].coords[1])
            self.assertEqual(coords[i][2], self.square_sg.structure[i].coords[2])
            self.assertEqual(properties[i], self.square_sg.structure[i].properties)

    def test_edge_editing(self):
        square = copy.deepcopy(self.square_sg)

        square.alter_edge(0, 0, to_jimage=(1, 0, 0),  new_weight=0.0,
                                  new_edge_properties={"foo": "bar"})
        new_edge = square.graph.get_edge_data(0, 0)[0]
        self.assertEqual(new_edge["weight"], 0.0)
        self.assertEqual(new_edge["foo"], "bar")

        square.break_edge(0, 0, to_jimage=(1, 0, 0))
        self.assertEqual(len(square.graph.get_edge_data(0, 0)), 3)

    def test_insert_remove(self):
        struct_copy = copy.deepcopy(self.square_sg.structure)
        square_copy = copy.deepcopy(self.square_sg)

        # Ensure that insert_node appropriately wraps Structure.insert()
        struct_copy.insert(1, "O", [0.5, 0.5, 0.5])
        square_copy.insert_node(1, "O", [0.5, 0.5, 0.5])
        self.assertEqual(struct_copy, square_copy.structure)

        # Test that removal is also equivalent between Structure and StructureGraph.structure
        struct_copy.remove_sites([1])
        square_copy.remove_nodes([1])
        self.assertEqual(struct_copy, square_copy.structure)

        square_copy.insert_node(1, "O", [0.5, 0.5, 0.5], edges=[{"from_index": 1,
                                                                 "to_index": 0,
                                                                 "to_jimage": (0, 0, 0)}])
        self.assertEqual(square_copy.get_coordination_of_site(1), 1)

        # Test that StructureGraph.graph is correctly updated
        square_copy.insert_node(1, "H", [0.5, 0.5, 0.75], edges=[{"from_index": 1,
                                                                 "to_index": 2,
                                                                 "to_jimage": (0, 0, 0)}])
        square_copy.remove_nodes([1])
        self.assertEqual(square_copy.graph.number_of_nodes(), 2)
        self.assertEqual(square_copy.graph.number_of_edges(), 5)

    def test_substitute(self):
        structure = Structure.from_file(os.path.join(os.path.dirname(__file__),
                                                     "..", "..", "..",
                                                     "test_files", "Li2O.cif"))
        molecule = FunctionalGroups["methyl"]

        structure_copy = copy.deepcopy(structure)
        structure_copy_graph = copy.deepcopy(structure)

        sg = StructureGraph.with_local_env_strategy(structure, MinimumDistanceNN())
        sg_copy = copy.deepcopy(sg)

        # Ensure that strings and molecules lead to equivalent substitutions
{'fraction_non_alphanumeric': 0.1419698314108252, 'fraction_numerical': 0.0385980479148181, 'mean_word_length': 2.8507367072389496, 'pattern_counts': {'":': 11, '<': 11, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 30, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7305515', 'n_tokens_mistral': 6688, 'n_tokens_neox': 6171, 'n_words': 1177}