Generate hierarchical visualization for VedEf operating window models.

This script demonstrates the hierarchical parameter visualization module, showing the 5 primary parameters, models, and 10 constraints in multiple formats: - Hierarchical Mermaid diagram (showing parameter levels and data flow) - NetworkX graph export (for external tools like yEd, Gephi, D3.js)

import importlib.util
from pathlib import Path

from paroto.viz import (
    export_graphml,
    export_json,
    extract_hierarchy,
    generate_hierarchical_graph,
    to_networkx,
)
from paroto.viz.networkx_export import analyze_hierarchy

# Import setup_vedef_problem from the setup file
spec = importlib.util.spec_from_file_location(
    "problem_module", Path(__file__).parent / "1_setup_vedef_problem.py"
)
problem_module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(problem_module)
setup_vedef_problem = problem_module.setup_vedef_problem


def generate_vedef_visualizations():
    """Generate all visualization formats for VedEf system."""
    print("=" * 80)
    print("VedEf Hierarchical Visualization Generation")
    print("=" * 80)
    print()

    base_dir = Path(__file__).parent
    output_dir = base_dir / "out"
    output_dir.mkdir(exist_ok=True)

    # Step 1: Set up OpenMDAO problem
    print("Step 1: Setting up OpenMDAO problem...")
    prob = setup_vedef_problem()
    print("  [OK] Problem configured")
    print()

    # Step 2: Extract hierarchy
    print("Step 2: Extracting parameter hierarchy...")
    primary_params = [
        "G_UMAX_OUT",  # Generator voltage (High voltage)
        "G_e",  # Interelectrode gap distance
        "TP_D_OUT",  # Torch outlet diameter
        "G_F",  # PRF - Pulse Frequency
        "G_Ep",  # Energy per pulse
    ]
    graph = extract_hierarchy(prob, primary_params)

    num_nodes = len(graph.nodes)
    num_edges = len(graph.edges)
    levels_dict = {nid: attrs["level"] for nid, attrs in graph.nodes.items()}
    num_levels = len(set(levels_dict.values()))
    print(f"  [OK] Extracted {num_nodes} nodes, {num_edges} edges, {num_levels} levels")
    print()

    # Step 3: Generate hierarchical Mermaid diagram
    print("Step 3: Generating hierarchical Mermaid diagram...")
    hier_file = generate_hierarchical_graph(
        graph,
        output_dir / "vedef_hierarchical_diagram.mmd",
        max_level=None,  # Show all levels
        orientation="LR",
    )
    print(f"  [OK] Saved: {hier_file}")
    print()

    # Step 4: Export to NetworkX formats
    print("Step 4: Exporting to NetworkX formats...")
    nx_graph = to_networkx(graph)
    graphml_file = export_graphml(nx_graph, output_dir / "vedef_graph.graphml")
    json_file = export_json(nx_graph, output_dir / "vedef_graph.json")
    print(f"  [OK] Saved: {graphml_file} (for yEd, Gephi)")
    print(f"  [OK] Saved: {json_file} (for D3.js, custom tools)")
    print()

    # Step 5: Analyze hierarchy
    print("Step 5: Analyzing hierarchy structure...")
    stats = analyze_hierarchy(nx_graph)
    print(f"  Nodes: {stats['num_nodes']}")
    print(f"  Edges: {stats['num_edges']}")
    print(f"  Levels: {stats['num_levels']}")
    print(f"  Max in-degree: {stats['max_in_degree']}")
    print(f"  Max out-degree: {stats['max_out_degree']}")
    print()

    print("=" * 80)
    print("All visualizations generated successfully!")
    print()
    print("To visualize Mermaid diagrams:")
    print("  1. Copy contents to https://mermaid.live/")
    print("  2. Or use a Markdown viewer that supports Mermaid")
    print()
    print("To visualize GraphML:")
    print("  1. Open vedef_graph.graphml in yEd (https://www.yworks.com/yed)")
    print("  2. Or import into Gephi (https://gephi.org/)")
    print()
    print("=" * 80)


if __name__ == "__main__":
    generate_vedef_visualizations()