Merge pull request #181 from hmlli/neb_setup

Reverse Supercell Mapping

pymatgen/analysis/diffusion/neb/full_path_mapper.py

@@ -79,7 +79,7 @@ class MigrationGraph(MSONable):
     def __init__(
         self,
         structure: Structure,
-        migration_graph: StructureGraph,
+        m_graph: StructureGraph,
         symprec=0.1,
         vac_mode=False,
     ):
@@ -93,20 +93,20 @@ def __init__(
             structure: Structure with base framework and mobile sites.
              When used with structure_is_base = True, only the base framework
              structure, does not contain any migrating sites.
-            migration_graph: The StructureGraph object that defines the
+            m_graph: The StructureGraph object that defines the
              migration network
             symprec (float): Symmetry precision to determine equivalence
              of migration events
             vac_mode (Bool): indicates whether vacancy mode should be used
         """
         self.structure = structure
-        self.migration_graph = migration_graph
+        self.m_graph = m_graph
         self.symprec = symprec
         self.vac_mode = vac_mode
         if self.vac_mode:
             raise NotImplementedError("Vacancy mode is not yet implemented")
         # Generate the graph edges between these all the sites
-        self.migration_graph.set_node_attributes()  # popagate the sites properties to the graph nodes
+        self.m_graph.set_node_attributes()  # popagate the sites properties to the graph nodes
         # For poperies like unique_hops we might be interested in modifying them after creation
         # So let's not convert them into properties for now.  (Awaiting rewrite once the usage becomes more clear.)
         self._populate_edges_with_migration_paths()
@@ -117,7 +117,7 @@ def only_sites(self) -> Structure:
         """
         A structure that only contains the migrating species
         """
-        return self.migration_graph.structure
+        return self.m_graph.structure
 
     @property
     def host_structure(self) -> Structure:
@@ -152,27 +152,25 @@ def unique_hops(self):
         The unique hops dictionary keyed by the hop label
         """
         # reversed so that the first instance represents the group of distinct hops
-        ihop_data = list(reversed(list(self.migration_graph.graph.edges(data=True))))
+        ihop_data = list(reversed(list(self.m_graph.graph.edges(data=True))))
         for u, v, d in ihop_data:
             d["iindex"] = u
             d["eindex"] = v
             d["hop_distance"] = d["hop"].length
         return {d["hop_label"]: d for u, v, d in ihop_data}
 
     @classmethod
-    def with_base_structure(
-        cls, base_structure: Structure, migration_graph: StructureGraph, **kwargs
-    ) -> "MigrationGraph":
+    def with_base_structure(cls, base_structure: Structure, m_graph: StructureGraph, **kwargs) -> "MigrationGraph":
         """
         Args:
             base_structure: base framework structure that does not contain any
              migrating sites.
         Returns:
             A constructed MigrationGraph object
         """
-        sites = migration_graph.structure.sites + base_structure.sites
+        sites = m_graph.structure.sites + base_structure.sites
         structure = Structure.from_sites(sites)
-        return cls(structure=structure, migration_graph=migration_graph, **kwargs)
+        return cls(structure=structure, m_graph=m_graph, **kwargs)
 
     @classmethod
     def with_local_env_strategy(
@@ -189,7 +187,7 @@ def with_local_env_strategy(
         """
         only_sites = get_only_sites_from_structure(structure, migrating_specie)
         migration_graph = StructureGraph.with_local_env_strategy(only_sites, nn)
-        return cls(structure=structure, migration_graph=migration_graph, **kwargs)
+        return cls(structure=structure, m_graph=migration_graph, **kwargs)
 
     @classmethod
     def with_distance(
@@ -209,7 +207,7 @@ def with_distance(
             only_sites,
             MinimumDistanceNN(cutoff=max_distance, get_all_sites=True),
         )
-        return cls(structure=structure, migration_graph=migration_graph, **kwargs)
+        return cls(structure=structure, m_graph=migration_graph, **kwargs)
 
     @staticmethod
     def get_structure_from_entries(
@@ -261,7 +259,7 @@ def _get_pos_and_migration_path(self, u, v, w):
           v (int): index of final node
           w (int): index for multiple edges that share the same two nodes
         """
-        edge = self.migration_graph.graph[u][v][w]
+        edge = self.m_graph.graph[u][v][w]
         i_site = self.only_sites.sites[u]
         e_site = PeriodicSite(
             self.only_sites.sites[v].species,
@@ -280,16 +278,16 @@ def _populate_edges_with_migration_paths(self):
         """
         Populate the edges with the data for the Migration Paths
         """
-        list(starmap(self._get_pos_and_migration_path, self.migration_graph.graph.edges))
+        list(starmap(self._get_pos_and_migration_path, self.m_graph.graph.edges))
 
     def _group_and_label_hops(self):
         """
         Group the MigrationHop objects together and label all the symmetrically equlivaelnt hops with the same label
         """
-        hops = list(nx.get_edge_attributes(self.migration_graph.graph, "hop").items())
+        hops = list(nx.get_edge_attributes(self.m_graph.graph, "hop").items())
         labs = generic_groupby(hops, comp=lambda x, y: x[1] == y[1])
         new_attr = {g_index: {"hop_label": labs[edge_index]} for edge_index, (g_index, _) in enumerate(hops)}
-        nx.set_edge_attributes(self.migration_graph.graph, new_attr)
+        nx.set_edge_attributes(self.m_graph.graph, new_attr)
         return new_attr
 
     def add_data_to_similar_edges(
@@ -307,7 +305,7 @@ def add_data_to_similar_edges(
             determine whether the data needs to be flipped so that 0-->1 is different from 1-->0
         """
 
-        for u, v, d in self.migration_graph.graph.edges(data=True):
+        for u, v, d in self.m_graph.graph.edges(data=True):
             if d["hop_label"] == target_label:
                 d.update(data)
                 if m_path is not None:
@@ -353,10 +351,10 @@ def get_path(self, max_val=100000):
                 f"There are {len(self.unique_hops)} SC hops but {len(self.unique_hops)} UC hops in {self.__str__()}"
             )
 
-        # for u, v, k, d in self.migration_graph.graph.edges(data=True, keys=True):
-        for u in self.migration_graph.graph.nodes():
+        # for u, v, k, d in self.m_graph.graph.edges(data=True, keys=True):
+        for u in self.m_graph.graph.nodes():
             # Create a copy of the graph so that we can trim the higher cost hops
-            path_graph = deepcopy(self.migration_graph.graph)
+            path_graph = deepcopy(self.m_graph.graph)
             # Trim the higher cost edges from the network
             cut_edges = []
             for tmp_u, tmp_v, tmp_k, tmp_d in path_graph.edges(data=True, keys=True):
@@ -412,7 +410,7 @@ def get_summary_dict(self, added_keys: List[str] = None) -> dict:
         def get_keys(d):
             return {k_: d[k_] for k_ in keys if k_ in d}
 
-        for u, v, d in self.migration_graph.graph.edges(data=True):
+        for u, v, d in self.m_graph.graph.edges(data=True):
             new_hop = get_keys(d)
             new_hop["iindex"] = u
             new_hop["eindex"] = v
@@ -443,7 +441,7 @@ class ChargeBarrierGraph(MigrationGraph):
     def __init__(
         self,
         structure: Structure,
-        migration_graph: StructureGraph,
+        m_graph: StructureGraph,
         potential_field: VolumetricData,
         potential_data_key: str,
         **kwargs,
@@ -460,7 +458,7 @@ def __init__(
         """
         self.potential_field = potential_field
         self.potential_data_key = potential_data_key
-        super().__init__(structure=structure, migration_graph=migration_graph, **kwargs)
+        super().__init__(structure=structure, m_graph=m_graph, **kwargs)
         self._setup_grids()
 
     def _setup_grids(self):
@@ -807,20 +805,20 @@ def check_uc_hop(sc_hop, uc_hop):
     return None
 
 
-def map_hop_sc2uc(sc_hop, fpm_uc):
+def map_hop_sc2uc(sc_hop: MigrationHop, mg: MigrationGraph):
     """
     Map a given hop in the SC onto the UC.
     Args:
         sc_hop: MigrationHop object form pymatgen-diffusion.
-        fpm_uc: MigrationGraph object from pymatgen-diffusion.
+        mg: MigrationGraph object from pymatgen-diffusion.
     Note:
         For now assume that the SC is exactly 2x2x2 of the UC.
         Can add in the parsing of different SC's later
         For a migration event in the SC from (0.45,0,0)-->(0.55,0,0)
         the UC hop might be (0.9,0,0)-->(0.1,0,0)[img:100]
         for the inverse of (0.1,0,0)-->(-0.1,0,0) the code needs to account for both those cases
     """
-    for u, v, d in fpm_uc.migration_graph.graph.edges(data=True):
+    for u, v, d in mg.m_graph.graph.edges(data=True):
         chk_res = check_uc_hop(sc_hop=sc_hop, uc_hop=d["hop"])
         if chk_res is not None:
             assert almost(d["hop"].length, sc_hop.length)

pymatgen/analysis/diffusion/neb/tests/test_full_path_mapper.py

@@ -34,7 +34,7 @@ def test_get_pos_and_migration_path(self):
         Make sure that we can populate the graph with MigrationHop Objects
         """
         self.fpm._get_pos_and_migration_path(0, 1, 1)
-        self.assertAlmostEqual(self.fpm.migration_graph.graph[0][1][1]["hop"].length, 3.571248, 4)
+        self.assertAlmostEqual(self.fpm.m_graph.graph[0][1][1]["hop"].length, 3.571248, 4)
 
     def test_get_summary_dict(self):
         summary_dict = self.fpm.get_summary_dict()
@@ -54,10 +54,10 @@ def setUp(self):
             migrating_ion_entry=self.li_ent,
         )[0]
 
-    def test_migration_graph_construction(self):
+    def test_m_graph_construction(self):
         self.assertEqual(self.full_struct.composition["Li"], 8)
         mg = MigrationGraph.with_distance(self.full_struct, migrating_specie="Li", max_distance=4.0)
-        self.assertEqual(len(mg.migration_graph.structure), 8)
+        self.assertEqual(len(mg.m_graph.structure), 8)
 
 
 class MigrationGraphComplexTest(unittest.TestCase):
@@ -73,15 +73,15 @@ def test_group_and_label_hops(self):
         """
         Check that the set of end points in a group of similiarly labeled hops are all the same
         """
-        edge_labs = np.array([d["hop_label"] for u, v, d in self.fpm_li.migration_graph.graph.edges(data=True)])
+        edge_labs = np.array([d["hop_label"] for u, v, d in self.fpm_li.m_graph.graph.edges(data=True)])
 
         site_labs = np.array(
             [
                 (
                     d["hop"].symm_structure.wyckoff_symbols[d["hop"].iindex],
                     d["hop"].symm_structure.wyckoff_symbols[d["hop"].eindex],
                 )
-                for u, v, d in self.fpm_li.migration_graph.graph.edges(data=True)
+                for u, v, d in self.fpm_li.m_graph.graph.edges(data=True)
             ]
         )
 
@@ -103,14 +103,14 @@ def test_unique_hops_dict(self):
     def test_add_data_to_similar_edges(self):
         # passing normal data
         self.fpm_li.add_data_to_similar_edges(0, {"key0": "data"})
-        for u, v, d in self.fpm_li.migration_graph.graph.edges(data=True):
+        for u, v, d in self.fpm_li.m_graph.graph.edges(data=True):
             if d["hop_label"] == 0:
                 self.assertEqual(d["key0"], "data")
 
         # passing ordered list data
         migration_path = self.fpm_li.unique_hops[1]["hop"]
         self.fpm_li.add_data_to_similar_edges(1, {"key1": [1, 2, 3]}, m_path=migration_path)
-        for u, v, d in self.fpm_li.migration_graph.graph.edges(data=True):
+        for u, v, d in self.fpm_li.m_graph.graph.edges(data=True):
             if d["hop_label"] == 1:
                 self.assertEqual(d["key1"], [1, 2, 3])
 
@@ -121,24 +121,24 @@ def test_add_data_to_similar_edges(self):
             symm_structure=migration_path.symm_structure,
         )
         self.fpm_li.add_data_to_similar_edges(2, {"key2": [1, 2, 3]}, m_path=migration_path_reversed)
-        for u, v, d in self.fpm_li.migration_graph.graph.edges(data=True):
+        for u, v, d in self.fpm_li.m_graph.graph.edges(data=True):
             if d["hop_label"] == 2:
                 self.assertEqual(d["key2"], [3, 2, 1])
 
     def test_assign_cost_to_graph(self):
         self.fpm_li.assign_cost_to_graph()  # use 'hop_distance'
-        for u, v, d in self.fpm_li.migration_graph.graph.edges(data=True):
+        for u, v, d in self.fpm_li.m_graph.graph.edges(data=True):
             self.assertAlmostEqual(d["cost"], d["hop_distance"], 4)
 
         self.fpm_li.assign_cost_to_graph(cost_keys=["hop_distance", "hop_distance"])
-        for u, v, d in self.fpm_li.migration_graph.graph.edges(data=True):
+        for u, v, d in self.fpm_li.m_graph.graph.edges(data=True):
             self.assertAlmostEqual(d["cost"], d["hop_distance"] ** 2, 4)
 
     def test_periodic_dijkstra(self):
         self.fpm_li.assign_cost_to_graph()  # use 'hop_distance'
 
         # test the connection graph
-        sgraph = self.fpm_li.migration_graph
+        sgraph = self.fpm_li.m_graph
         G = sgraph.graph.to_undirected()
         conn_dict = _get_adjacency_with_images(G)
         for u in conn_dict.keys():
@@ -168,7 +168,7 @@ def test_get_path(self):
     def test_not_matching_first(self):
         structure = Structure.from_file(f"{dir_path}/pathfinder_files/Li6MnO4.json")
         fpm_lmo = MigrationGraph.with_distance(structure, "Li", max_distance=4)
-        for u, v, d in fpm_lmo.migration_graph.graph.edges(data=True):
+        for u, v, d in fpm_lmo.m_graph.graph.edges(data=True):
             self.assertIn(d["hop"].eindex, {0, 1})
 
 
@@ -226,7 +226,7 @@ def test_populate_edges_with_chg_density_info(self):
         """
         self.cbg.populate_edges_with_chg_density_info()
         length_vs_chg = list(
-            sorted([(d["hop"].length, d["chg_total"]) for u, v, d in self.cbg.migration_graph.graph.edges(data=True)])
+            sorted([(d["hop"].length, d["chg_total"]) for u, v, d in self.cbg.m_graph.graph.edges(data=True)])
         )
         prv = None
         for length, chg in length_vs_chg:

pymatgen/analysis/diffusion/neb/tests/test_pathfinder.py

@@ -165,8 +165,8 @@ def test_max_path_length(self):
 class MigrationHopTest(PymatgenTest):
     def setUp(self):
         self.lifepo = self.get_structure("LiFePO4")
-        migration_graph = MigrationGraph.with_distance(self.lifepo, max_distance=4.0, migrating_specie="Li")
-        gen = iter(migration_graph.migration_graph.graph.edges(data=True))
+        m_graph = MigrationGraph.with_distance(self.lifepo, max_distance=4.0, migrating_specie="Li")
+        gen = iter(m_graph.m_graph.graph.edges(data=True))
         u, v, d = next(gen)
         i_site = PeriodicSite("Li", coords=d["ipos"], lattice=self.lifepo.lattice)
         e_site = PeriodicSite("Li", coords=d["epos"], lattice=self.lifepo.lattice)