From bd8b26148b552fa0f4b02c0bb1edd1f1814be219 Mon Sep 17 00:00:00 2001 From: samtalki <10187005+samtalki@users.noreply.github.com> Date: Wed, 10 Jun 2026 13:18:03 -0400 Subject: [PATCH] feat: expose the synth generators through powerio-py and the MCP server One generate_case(topology, n, r_over_x, mean_x, seed) pyfunction dispatching through synth::generate, with a parse_topology helper following parse_scheme. Python wrapper powerio.generate_case(...) -> Network with stub updates, and a generate_case MCP tool returning {"json", "summary"} so synthetic cases enter the same transport pipeline as parsed ones. Defaults match SynthSpec::default and identical seeds generate identical cases. Closes #86. Co-Authored-By: Claude Fable 5 --- powerio-py/src/lib.rs | 41 ++++++++++++++++++++++++ python/powerio/__init__.py | 22 +++++++++++++ python/powerio/__init__.pyi | 7 +++++ python/powerio/_powerio.pyi | 7 +++++ python/powerio/mcp/server.py | 39 ++++++++++++++++++++++- python/tests/test_mcp.py | 13 ++++++++ python/tests/test_synth.py | 61 ++++++++++++++++++++++++++++++++++++ 7 files changed, 189 insertions(+), 1 deletion(-) create mode 100644 python/tests/test_synth.py diff --git a/powerio-py/src/lib.rs b/powerio-py/src/lib.rs index 674e994..f124dae 100644 --- a/powerio-py/src/lib.rs +++ b/powerio-py/src/lib.rs @@ -121,6 +121,20 @@ fn parse_convention(s: &str) -> PyResult { } } +/// Accepts `tree`, `lattice`/`lattice2d`, and `pegase`/`pegase-like` (case- and +/// separator-insensitive). +fn parse_topology(s: &str) -> PyResult { + use powerio_matrix::synth::Topology; + match normalize(s).as_str() { + "tree" => Ok(Topology::Tree), + "lattice" | "lattice2d" => Ok(Topology::Lattice2D), + "pegase" | "pegaselike" => Ok(Topology::PegaseLike), + other => Err(PyValueError::new_err(format!( + "unknown topology {other:?}; expected 'tree', 'lattice', or 'pegase-like'" + ))), + } +} + /// Accepts `perunit`/`pu`/`per-unit` and `native`. fn parse_units(s: &str) -> PyResult { match normalize(s).as_str() { @@ -634,6 +648,32 @@ fn convert_str(text: &str, to: &str, format: Option<&str>) -> PyResult<(String, Ok((conv.text, conv.warnings)) } +/// Generate a synthetic case: a spanning `tree`, a 2-D `lattice` (`n` rounds up +/// to a perfect square), or a `pegase-like` mesh (tree plus ~30% extra edges). +/// `n` below 2 is raised to 2 (lattice: at least a 2×2 grid). Identical +/// arguments (including `seed`) generate the identical case; bus 1 is the +/// reference. Buses and branches only — no loads, shunts, or generators. +#[pyfunction] +#[pyo3(signature = (topology=None, n=64, r_over_x=0.1, mean_x=0.05, seed=0x00C0_FFEE))] +fn generate_case( + topology: Option<&str>, + n: usize, + r_over_x: f64, + mean_x: f64, + seed: u64, +) -> PyResult { + let spec = powerio_matrix::synth::SynthSpec { + topology: parse_topology(topology.unwrap_or("tree"))?, + n, + r_over_x, + mean_x, + seed, + }; + let inner = powerio_matrix::synth::generate(&spec); + let core = IndexCore::build(&inner); + Ok(PyCase { inner, core }) +} + /// Build a `{dir, files}` dict from an outputs directory and its written files. /// Shared by the DC OPF and gridfm write paths. Paths go through [`path_to_str`] /// (so a non-UTF8 path raises instead of being mangled). @@ -749,6 +789,7 @@ fn _powerio(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(from_json, m)?)?; m.add_function(wrap_pyfunction!(convert_file, m)?)?; m.add_function(wrap_pyfunction!(convert_str, m)?)?; + m.add_function(wrap_pyfunction!(generate_case, m)?)?; // Whether the gridfm Parquet surface (arrow/parquet) was compiled in, so the // pure-Python layer can raise an ImportError instead of an AttributeError. m.add("_has_gridfm", cfg!(feature = "gridfm"))?; diff --git a/python/powerio/__init__.py b/python/powerio/__init__.py index 933948b..65d0836 100644 --- a/python/powerio/__init__.py +++ b/python/powerio/__init__.py @@ -49,6 +49,7 @@ "parse_file", "parse_str", "from_json", + "generate_case", "convert_file", "convert_str", "to_format", @@ -410,6 +411,27 @@ def from_json(text: str) -> Network: return Network(_powerio.from_json(text)) +def generate_case( + topology: str = "tree", + n: int = 64, + r_over_x: float = 0.1, + mean_x: float = 0.05, + seed: int = 0x00C0_FFEE, +) -> Network: + """Generate a synthetic case: a spanning ``tree``, a 2-D ``lattice`` (``n`` + rounds up to a perfect square), or a ``pegase-like`` mesh (tree plus ~30% + extra edges). + + ``n`` below 2 is raised to 2 (lattice: at least a 2×2 grid). Identical + arguments (including ``seed``) generate the identical case; bus 1 is the + reference. The case carries buses and branches only — no loads, shunts, or + generators — so it exercises topology and matrix paths. The result is a + full :class:`Network`: the matrix builders, converters, and the JSON + transport all work on it. + """ + return Network(_powerio.generate_case(topology, n, r_over_x, mean_x, seed)) + + def convert_file(path: Any, to: str, from_: Optional[str] = None) -> Conversion: """Convert a case file to another format through the neutral hub. diff --git a/python/powerio/__init__.pyi b/python/powerio/__init__.pyi index e914462..fb24929 100644 --- a/python/powerio/__init__.pyi +++ b/python/powerio/__init__.pyi @@ -194,6 +194,13 @@ Format = str def parse_file(path: Any, from_: Optional[Format] = ...) -> Network: ... def parse_str(text: str, format: Format = ...) -> Network: ... def from_json(text: str) -> Network: ... +def generate_case( + topology: str = ..., + n: int = ..., + r_over_x: float = ..., + mean_x: float = ..., + seed: int = ..., +) -> Network: ... def convert_file(path: Any, to: Format, from_: Optional[Format] = ...) -> Conversion: ... def convert_str(text: str, to: Format, format: Format = ...) -> Conversion: ... def to_format(case: Network, to: Format) -> Conversion: ... diff --git a/python/powerio/_powerio.pyi b/python/powerio/_powerio.pyi index e7617a7..4764d86 100644 --- a/python/powerio/_powerio.pyi +++ b/python/powerio/_powerio.pyi @@ -99,6 +99,13 @@ def convert_file( def convert_str( text: str, to: str, format: Optional[str] = ... ) -> Tuple[str, list[str]]: ... +def generate_case( + topology: Optional[str] = ..., + n: int = ..., + r_over_x: float = ..., + mean_x: float = ..., + seed: int = ..., +) -> PyCase: ... # Only present when the extension was compiled with the `gridfm` cargo feature # (the released wheel is); without it, access raises AttributeError. def write_gridfm_batch( diff --git a/python/powerio/mcp/server.py b/python/powerio/mcp/server.py index 64a08c6..84722a5 100644 --- a/python/powerio/mcp/server.py +++ b/python/powerio/mcp/server.py @@ -7,6 +7,8 @@ any fidelity warnings. - ``case_summary``: counts, base MVA, source format, and connectivity, with no scipy/numpy in the loop. +- ``generate_case``: a synthetic tree/lattice/pegase-like case as the JSON + transport, deterministic per seed. Run over stdio with the ``powerio-mcp`` console script (or ``python -m powerio.mcp``). The server reuses ``powerio.convert_file``/``convert_str``/ @@ -103,6 +105,41 @@ def case_summary( } +@mcp.tool() +def generate_case( + topology: str = "tree", + n: int = 64, + r_over_x: float = 0.1, + mean_x: float = 0.05, + seed: int = 0x00C0_FFEE, +) -> dict: + """Generate a synthetic power system case and return its JSON transport + plus a summary. + + ``topology`` is ``tree`` (radial spanning tree), ``lattice`` (2-D grid; + ``n`` rounds up to a perfect square), or ``pegase-like`` (tree plus ~30% + extra edges, transmission-like meshing). ``n`` below 2 is raised to 2 + (lattice: at least a 2×2 grid). Identical arguments (including ``seed``) + generate the identical case; bus 1 is the reference. The case carries + buses and branches only — no loads, shunts, or generators. + + The returned ``json`` is the same transport ``Network.to_json`` emits, so + any tool that accepts a parsed case accepts it. + """ + case = powerio.generate_case(topology, n, r_over_x, mean_x, seed) + return { + "json": case.to_json(), + "summary": { + "name": case.name, + "base_mva": case.base_mva, + "n_buses": case.n_buses, + "n_branches": case.n_branches, + "is_radial": case.is_radial, + "n_connected_components": case.n_connected_components, + }, + } + + def main() -> None: - """Console-script entry point: serve the two tools over stdio.""" + """Console-script entry point: serve the tools over stdio.""" mcp.run() diff --git a/python/tests/test_mcp.py b/python/tests/test_mcp.py index 183fcfa..2a92882 100644 --- a/python/tests/test_mcp.py +++ b/python/tests/test_mcp.py @@ -112,3 +112,16 @@ def boom(*args, **kwargs): monkeypatch.setattr(powerio, "convert_str", boom) with pytest.raises(ValueError): convert_case(to="psse", content="whatever", from_="matpower") + + +def test_generate_case_tool(): + from powerio.mcp.server import generate_case + + r = generate_case(topology="lattice", n=9, seed=1) + assert r["summary"]["n_buses"] == 9 + assert r["summary"]["n_connected_components"] == 1 + assert powerio.from_json(r["json"]).n_buses == 9 + # deterministic per seed + assert generate_case(topology="lattice", n=9, seed=1)["json"] == r["json"] + with pytest.raises(ValueError): + generate_case(topology="torus") diff --git a/python/tests/test_synth.py b/python/tests/test_synth.py new file mode 100644 index 0000000..8df9354 --- /dev/null +++ b/python/tests/test_synth.py @@ -0,0 +1,61 @@ +"""Tests for the synthetic case generators (``powerio.generate_case``).""" + +import pytest + +import powerio + + +def test_tree_defaults(): + case = powerio.generate_case() + assert case.n_buses == 64 + assert case.n_branches == 63 # spanning tree + assert case.is_radial + assert case.n_connected_components == 1 + assert case.buses[0]["id"] == 1 and case.buses[0]["kind"] == "REF" + # buses and branches only + assert case.n_gens == 0 and case.n_loads == 0 and case.n_shunts == 0 + + +def test_lattice_rounds_up_to_perfect_square(): + case = powerio.generate_case("lattice", n=10) + assert case.n_buses == 16 + assert case.n_connected_components == 1 + + +def test_pegase_like_is_meshed(): + case = powerio.generate_case("pegase-like", n=100) + assert case.n_buses == 100 + assert case.n_branches > 99 + assert not case.is_radial + + +def test_same_seed_same_case(): + a = powerio.generate_case("tree", n=32, seed=7) + b = powerio.generate_case("tree", n=32, seed=7) + assert a.to_json() == b.to_json() + c = powerio.generate_case("tree", n=32, seed=8) + assert c.to_json() != a.to_json() + + +def test_generated_case_round_trips_through_json(): + case = powerio.generate_case("lattice", n=16) + back = powerio.from_json(case.to_json()) + assert back.n_buses == case.n_buses + assert back.n_branches == case.n_branches + + +def test_matrix_builders_work_on_generated_case(): + case = powerio.generate_case("pegase-like", n=64) + assert case.bprime().shape == (64, 64) + assert case.adjacency().nnz > 0 + + +def test_n_is_clamped_to_minimum(): + assert powerio.generate_case("tree", n=0).n_buses == 2 + assert powerio.generate_case("pegase-like", n=1).n_buses == 2 + assert powerio.generate_case("lattice", n=0).n_buses == 4 # 2x2 grid + + +def test_unknown_topology_raises(): + with pytest.raises(ValueError): + powerio.generate_case("torus")