diff --git a/src/structuregraph_helpers/create.py b/src/structuregraph_helpers/create.py index eae8ddb..d4458b9 100644 --- a/src/structuregraph_helpers/create.py +++ b/src/structuregraph_helpers/create.py @@ -150,7 +150,8 @@ def construct_clean_graph( else: graph = nx.Graph() for u, v, d in structure_graph.graph.edges(data=True): - graph.add_edge(u, v, **d) + voltage = _voltage(u, v, d["to_jimage"]) + graph.add_edge(u, v, voltage=voltage) for node in graph.nodes: graph.nodes[node]["specie"] = str(structure_graph.structure[node].specie) @@ -161,3 +162,19 @@ def construct_clean_graph( ) return graph + + +def _voltage(u, v, to_jimage) -> Tuple[int, int, int]: + """Voltage is the tuple describing the direction of the edge. + + In simple words, it represents the translation operation. + + Returns: + Tuple[int, int, int]: The voltage of the edge. + """ + terms = (u, v) + a_image = (0, 0, 0) + b_image = to_jimage + imags = (a_image, b_image) + a_image, b_image = (x for x, _ in sorted(zip(imags, terms), key=lambda x: x[1])) + return tuple(a_image[i] - b_image[i] for i in range(3)) diff --git a/src/structuregraph_helpers/hash.py b/src/structuregraph_helpers/hash.py index aa57d77..231c746 100644 --- a/src/structuregraph_helpers/hash.py +++ b/src/structuregraph_helpers/hash.py @@ -24,7 +24,7 @@ def generate_hash(g: nx.Graph, node_decorated: bool, edge_decorated: bool, itera if node_decorated: node_attr = "specie" if edge_decorated: - edge_attr = "to_jimage" + edge_attr = "voltage" return weisfeiler_lehman_graph_hash( g, iterations=iterations, edge_attr=edge_attr, node_attr=node_attr ) diff --git a/tests/test_create.py b/tests/test_create.py index 54b7e66..87bd0e0 100644 --- a/tests/test_create.py +++ b/tests/test_create.py @@ -49,7 +49,7 @@ def test_construct_clean_graph(bcc_graph): assert len(graph.edges) == 2 for _, _, d in graph.edges(data=True): - assert isinstance(d["to_jimage"], tuple) + assert isinstance(d["voltage"], tuple) for node in graph.nodes: assert isinstance(graph.nodes[node]["specie"], str)