Support multiple top cells and SPICE netlists in plot_nets

gplugins/klayout/plot_nets.py

@@ -1,4 +1,5 @@
 import itertools
+from logging import warning
 from pathlib import Path
 
 import klayout.db as kdb
@@ -21,30 +22,39 @@ def netlist_to_networkx(
         include_labels: Whether to include labels in the graph connected to corresponding cells.
     """
     G = nx.Graph()
-    assert netlist.top_circuit_count() == 1, "Multiple top cells not yet supported"
 
-    circuit, *_ = netlist.each_circuit_top_down()
-
-    # first flatten components that won't be kept
-    for subcircuit in circuit.each_subcircuit():
-        if subcircuit.name in {"TODO"}:
-            circuit.flatten_subcircuit(subcircuit)
-
-    for net in circuit.each_net():
-        net_pins = [
-            _get_subcircuit_name(subcircuit_pin_ref.subcircuit())
-            for subcircuit_pin_ref in net.each_subcircuit_pin()
-        ]
-
-        # Assumed lone net with only label info
-        if include_labels and net.expanded_name() and "," not in net.expanded_name():
-            G.add_edges_from(zip(net_pins, [net.name] * len(net_pins)))
-
-        if fully_connected:
-            G.add_edges_from(itertools.combinations(net_pins, 2))
-        else:
-            G.add_edges_from(zip(net_pins[:-1], net_pins[1:]))
+    top_circuits = list(
+        itertools.islice(netlist.each_circuit_top_down(), netlist.top_circuit_count())
+    )
 
+    for circuit in top_circuits:
+        # first flatten components that won't be kept
+        for subcircuit in circuit.each_subcircuit():
+            if subcircuit.name in {"TODO"}:
+                circuit.flatten_subcircuit(subcircuit)
+
+        for net in circuit.each_net():
+            # Get subcircuit pins if they exist (hierarchical export from KLayout)
+            net_pins = [
+                _get_subcircuit_name(subcircuit_pin_ref.subcircuit())
+                for subcircuit_pin_ref in net.each_subcircuit_pin()
+            ]
+            # or use all pins (flat like from Cadence SPICE)
+            if not net_pins:
+                net_pins.extend(pin_ref.pin().name() for pin_ref in net.each_pin())
+
+            # Assumed lone net with only label info
+            if (
+                include_labels
+                and net.expanded_name()
+                and "," not in net.expanded_name()
+            ):
+                G.add_edges_from(zip(net_pins, [net.name] * len(net_pins)))
+
+            if fully_connected:
+                G.add_edges_from(itertools.combinations(net_pins, 2))
+            else:
+                G.add_edges_from(zip(net_pins[:-1], net_pins[1:]))
     return G
 
 
@@ -57,16 +67,29 @@ def plot_nets(
     """Plots the connectivity between the components in the KLayout LayoutToNetlist file from :func:`~get_l2n`.
 
     Args:
-        filepath: Path to the KLayout LayoutToNetlist file.
+        filepath: Path to the KLayout LayoutToNetlist file or a SPICE netlist.
+            File extensions should be `.l2n` and `.spice`, respectively.
         fully_connected: Whether to plot the graph as elements fully connected to all other ones (True) or
 
             going through other elements (False).
         interactive: Whether to plot an interactive graph with `pyvis` or not.
         include_labels: Whether to include labels in the graph connected to corresponding cells.
     """
-    l2n = kdb.LayoutToNetlist()
-    l2n.read(str(filepath))
-    netlist = l2n.netlist()
+    match Path(filepath).suffix:
+        case ".l2n" | ".txt":
+            l2n = kdb.LayoutToNetlist()
+            l2n.read(str(filepath))
+            netlist = l2n.netlist()
+        case ".spice":
+            reader = kdb.NetlistSpiceReader()
+            netlist = kdb.Netlist()
+            netlist.read(str(filepath), reader)
+        case _:
+            warning("Assuming file is KLayout native LayoutToNetlist file")
+            l2n = kdb.LayoutToNetlist()
+            l2n.read(str(filepath))
+            netlist = l2n.netlist()
+
     # Creating a graph for the connectivity
     G_connectivity = netlist_to_networkx(
         netlist, fully_connected=fully_connected, include_labels=include_labels

gplugins/klayout/tests/test_plot_nets.py

@@ -1,9 +1,10 @@
 from pathlib import Path
 
+import klayout.db as kdb
 import pytest
 from gdsfactory.samples.demo.lvs import pads_correct
 
-from gplugins.klayout.get_netlist import get_l2n
+from gplugins.klayout.get_netlist import get_l2n, get_netlist
 from gplugins.klayout.plot_nets import plot_nets
 
 
@@ -19,6 +20,18 @@ def klayout_netlist(tmpdir_factory) -> str:
     return netlist_path
 
 
+@pytest.fixture(scope="session")
+def spice_netlist(tmpdir_factory) -> str:
+    """Get SPICE netlist file for `pads_correct`. Cached for session scope."""
+    tmp_path = tmpdir_factory.mktemp("data")
+    c = pads_correct()
+    gdspath = c.write_gds(gdsdir=tmp_path)
+    netlist = get_netlist(gdspath)
+    netlist_path = str(Path(tmp_path) / f"{c.name}.spice")
+    netlist.write(netlist_path, kdb.NetlistSpiceWriter())
+    return netlist_path
+
+
 def test_plot_nets(klayout_netlist):
     plot_nets(klayout_netlist)
 
@@ -34,3 +47,7 @@ def test_plot_nets_not_fully_connected(klayout_netlist):
 
 def test_plot_nets_no_labels(klayout_netlist):
     plot_nets(klayout_netlist, include_labels=False)
+
+
+def test_plot_nets_spice(spice_netlist):
+    plot_nets(spice_netlist)