diff --git a/pyproject.toml b/pyproject.toml index f43dcaa..ddde539 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "vsec" -version = "0.1.2" +version = "0.1.3" description = "Data structure when some edges in a graph are planar points in essence" authors = ["edxu96 "] diff --git a/tests/test_conversion.py b/tests/test_conversion.py index 969b258..c820ac6 100644 --- a/tests/test_conversion.py +++ b/tests/test_conversion.py @@ -2,7 +2,7 @@ import networkx as nx from pandas.core.frame import DataFrame import pytest as pt -from vsec.conversion import contract, split +from vsec.conversion import contract, split, split_vertex_df from vsec.geometry import GeoGraph # There are only two edges left, but sequences might be different. @@ -43,7 +43,9 @@ def test_split(case_grid: GeoGraph, vertices_grid: DataFrame): vertices_grid: vertices in ``case_grid`` and their **type** attributes. """ + # Only split 10-0.4 kV transformers. vertices_split = vertices_grid.index[vertices_grid["type"] == "STAT1004"] + graph, vertex_df = split(case_grid, vertices_split, NAMING, ATTR, IS_FIRST) assert isinstance(nx.to_pandas_edgelist(graph), DataFrame) @@ -51,3 +53,21 @@ def test_split(case_grid: GeoGraph, vertices_grid: DataFrame): assert isinstance(vertex_df, DataFrame) assert vertex_df.index.name == "original" assert set(vertex_df.columns) == COLUMNS + assert vertex_df.shape == (34, 2) + + +@pt.mark.usefixtures("vertices_grid") +def test_split_vertex_df(vertices_grid: DataFrame): + """Check if can gather split vertices and resulted new vertices. + + Args: + vertices_grid: vertices in ``case_grid`` and their **type** + attributes. + """ + # Only split 10-0.4 kV transformers. + vertices_split = vertices_grid.index[vertices_grid["type"] == "STAT1004"] + + res = split_vertex_df(vertices_split, NAMING) + assert set(res.columns) == COLUMNS + assert res.shape == (34, 2) + assert res.index.name == "original" diff --git a/vsec/conversion.py b/vsec/conversion.py index e284025..24281f3 100644 --- a/vsec/conversion.py +++ b/vsec/conversion.py @@ -103,3 +103,28 @@ def split( vertex_df.index.name = "original" return graph, vertex_df + + +def split_vertex_df( + vertices: Set[str], naming: Callable[[str], Tuple[str, str]], +) -> DataFrame: + """Gather split vertices and resulted new vertices. + + Warning: + This should be used when vertex splitting is known to be + correct, because there is no validation. + + Args: + vertices: vertices ought to be modelled as edges. + naming: how two resulted vertices should be named. + + Returns: + Correspondence between split vertices & resulted new vertices. + """ + vertex_dict = {} + for vertex in vertices: + vertex_dict[vertex] = naming(vertex) + + res = pd.DataFrame.from_dict(vertex_dict, orient="index", columns=COLUMNS,) + res.index.name = "original" + return res