From 338f12077a4f09d7a757a2b20ca6f1662cd55876 Mon Sep 17 00:00:00 2001
From: Niko Savola <nikosavola@google.com>
Date: Tue, 21 Nov 2023 11:51:42 -0800
Subject: [PATCH] Support reducing nodes in plot nets

---
 gplugins/klayout/plot_nets.py            | 24 ++++++++++++++++++++++++
 gplugins/klayout/tests/test_plot_nets.py |  9 ++++++++-
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/gplugins/klayout/plot_nets.py b/gplugins/klayout/plot_nets.py
index c1330809..7163969a 100644
--- a/gplugins/klayout/plot_nets.py
+++ b/gplugins/klayout/plot_nets.py
@@ -1,5 +1,7 @@
 import itertools
+from collections.abc import Collection
 from pathlib import Path
+from typing import cast
 
 import klayout.db as kdb
 import matplotlib.pyplot as plt
@@ -79,6 +81,7 @@ def plot_nets(
     interactive: bool = False,
     include_labels: bool = True,
     only_most_complex: bool = False,
+    nodes_to_reduce: Collection[str] | None = None,
 ) -> None:
     """Plots the connectivity between the components in the KLayout LayoutToNetlist file from :func:`~get_l2n`.
 
@@ -91,7 +94,10 @@ def plot_nets(
         include_labels: Whether to include labels in the graph connected to corresponding cells.
         only_most_complex: Whether to plot only the circuit with most connections or not.
             Helpful for not plotting subcircuits separately.
+        nodes_to_reduce: Nodes to reduce to a single edge. Comparison made with Python ``in`` operator.
+            Helpful for reducing trivial waveguide elements.
     """
+
     match Path(filepath).suffix:
         case ".l2n" | ".txt":
             l2n = kdb.LayoutToNetlist()
@@ -115,6 +121,24 @@ def plot_nets(
         only_most_complex=only_most_complex,
     )
 
+    if nodes_to_reduce:
+        G_connectivity_tmp = G_connectivity.copy()
+        for node, degree in G_connectivity.degree:
+            if degree == 2 and any(e in node for e in nodes_to_reduce):
+                edges = cast(
+                    tuple[tuple[str, str], tuple[str, str]],
+                    tuple(G_connectivity_tmp.edges(node).__iter__()),
+                )
+                a0, b0 = edges[0]
+                a1, b1 = edges[1]
+
+                e0 = a0 if a0 != node else b0
+                e1 = a1 if a1 != node else b1
+
+                G_connectivity_tmp.remove_node(node)
+                G_connectivity_tmp.add_edge(e0, e1)
+        G_connectivity = G_connectivity_tmp
+
     # Plotting the graph
     if interactive:
         try:
diff --git a/gplugins/klayout/tests/test_plot_nets.py b/gplugins/klayout/tests/test_plot_nets.py
index 18b60321..c95e5ef3 100644
--- a/gplugins/klayout/tests/test_plot_nets.py
+++ b/gplugins/klayout/tests/test_plot_nets.py
@@ -36,8 +36,14 @@ def spice_netlist(tmpdir_factory) -> str:
 @pytest.mark.parametrize("interactive", [True, False])
 @pytest.mark.parametrize("include_labels", [True, False])
 @pytest.mark.parametrize("only_most_complex", [True, False])
+@pytest.mark.parametrize("nodes_to_reduce", [None, {"straight"}, {"straight", "pad"}])
 def test_plot_nets(
-    klayout_netlist, fully_connected, interactive, include_labels, only_most_complex
+    klayout_netlist,
+    fully_connected,
+    interactive,
+    include_labels,
+    only_most_complex,
+    nodes_to_reduce,
 ):
     plot_nets(
         klayout_netlist,
@@ -45,6 +51,7 @@ def test_plot_nets(
         interactive=interactive,
         include_labels=include_labels,
         only_most_complex=only_most_complex,
+        nodes_to_reduce=nodes_to_reduce,
     )
     if interactive:
         assert Path("connectivity.html").exists()