Example: Using ProblemGraph abstraction for unified visualization.

This example demonstrates the new ProblemGraph abstraction layer that provides a unified interface for all visualization backends (Mermaid, Cytoscape, NetworkX).

import openmdao.api as om

from paroto.viz import (
    ProblemGraph,
    export_graphml,
    generate_hierarchical_graph,
    to_networkx,
)
from paroto.viz.cytoscape import hierarchy_to_cytoscape


def create_example_problem():
    """Create a simple example problem."""
    prob = om.Problem()

    # Design variables
    indeps = prob.model.add_subsystem("design", om.IndepVarComp(), promotes=["*"])
    indeps.add_output("x", val=2.0)
    indeps.add_output("y", val=3.0)

    # Intermediate computation
    prob.model.add_subsystem(
        "compute",
        om.ExecComp("z = x**2 + y**2"),
        promotes=["*"],
    )

    # Constraint
    prob.model.add_subsystem(
        "constraint",
        om.ExecComp("z_margin = z - 10.0"),
        promotes=["*"],
    )

    prob.setup()
    prob.run_model()

    return prob


def main():
    """Demonstrate ProblemGraph abstraction."""
    print("=" * 80)
    print("ProblemGraph Abstraction Example")
    print("=" * 80)
    print()

    # Create problem
    print("Step 1: Creating OpenMDAO problem...")
    prob = create_example_problem()
    print("  [OK] Problem created")
    print()

    # Create ProblemGraph abstraction
    print("Step 2: Creating ProblemGraph abstraction...")
    graph = ProblemGraph(prob, primary_params=["x", "y"])
    print(f"  [OK] Graph with {len(graph.nodes)} nodes, {len(graph.edges)} edges")
    print()

    # Inspect graph structure
    print("Step 3: Inspecting graph structure...")
    stats = graph.get_statistics()
    print(f"  Primary parameters: {stats['num_primary_params']}")
    print(f"  Hierarchy levels: {stats['num_levels']}")
    print(f"  Nodes by level: {stats['nodes_by_level']}")
    print(f"  Nodes by type: {stats['nodes_by_type']}")
    print()

    # Get organized views
    print("Step 4: Getting organized views...")
    nodes_by_level = graph.get_nodes_by_level()
    nodes_by_subsystem = graph.get_nodes_by_subsystem()
    print(f"  Nodes by level: {len(nodes_by_level)} levels")
    print(f"  Nodes by subsystem: {len(nodes_by_subsystem)} subsystems")
    print()

    # Generate visualizations using the same graph
    print("Step 5: Generating visualizations from single graph...")

    # Mermaid
    generate_hierarchical_graph(graph, "example_graph_mermaid.mmd")
    print("  [OK] Mermaid diagram generated")

    # Cytoscape
    cyto_elements = hierarchy_to_cytoscape(graph)
    print(f"  [OK] Cytoscape elements generated ({len(cyto_elements)} elements)")

    # NetworkX
    nx_graph = to_networkx(graph)
    export_graphml(nx_graph, "example_graph.graphml")
    print("  [OK] NetworkX/GraphML exported")

    print()

    # Show filtering capability
    print("Step 6: Filtering by level...")
    filtered_nodes, filtered_edges = graph.filter_by_level(max_level=1)
    print(f"  Filtered to level 0-1: {len(filtered_nodes)} nodes, {len(filtered_edges)} edges")
    print()

    print("=" * 80)
    print("ProblemGraph provides a unified abstraction that:")
    print("  - Wraps OpenMDAO Problem connections")
    print("  - Provides consistent interface for all viz backends")
    print("  - Supports filtering, organization, and statistics")
    print("  - Eliminates code duplication between generators")
    print("=" * 80)


if __name__ == "__main__":
    main()