perf(tidy3d): FXC-4377-speed-up-klayout-results-loader

Author: marcorudolphflex

CHANGELOG.md

@@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Changed
 - Removed validator that would warn if `PerturbationMedium` values could become numerically unstable, since an error will anyway be raised if this actually happens when the medium is converted using actual perturbation data.
 - Improved performance of adjoint gradient computation for `PolySlab` and `Cylinder` geometries through vectorized sidewall patch collection (~4x speedup).
+- Added optional `max_results` handling to the kLayout `DRCLoader` and `DRCRunner` and speeded up parsing of big result sets.
 
 ### Fixed
 

tests/test_plugins/klayout/drc/test_drc.py

@@ -35,6 +35,66 @@ def _basic_drc_config_kwargs(tmp_path: Path) -> dict[str, Path | bool]:
     }
 
 
+def _write_results_file(
+    tmp_path: Path,
+    *,
+    category: str = "min_width",
+    num_items: int = 1,
+    filename: str = "many_results.lyrdb",
+) -> Path:
+    """Write a simple DRC results file with the requested number of items."""
+
+    template_header = f"""\
+<?xml version=\"1.0\" encoding=\"utf-8\"?>
+<report-database>
+ <categories>
+  <category>
+   <name>{category}</name>
+   <description>auto</description>
+   <categories></categories>
+  </category>
+ </categories>
+ <cells></cells>
+ <items>
+"""
+    template_footer = """\
+ </items>
+</report-database>
+"""
+    item = """\
+  <item>
+   <tags/>
+   <category>{category}</category>
+   <cell>TOP</cell>
+   <visited>false</visited>
+   <multiplicity>1</multiplicity>
+   <comment/>
+   <image/>
+   <values>
+    <value>edge: (0.0,0.0;1.0,1.0)</value>
+   </values>
+  </item>
+"""
+    contents = [template_header]
+    contents.extend(item.format(category=category) for _ in range(num_items))
+    contents.append(template_footer)
+    path = tmp_path / filename
+    path.write_text("".join(contents))
+    return path
+
+
+def _capture_log_warnings(monkeypatch):
+    """Capture calls to td.log.warning."""
+
+    messages = []
+
+    def fake_warning(message, *args, **kwargs):
+        messages.append(message % args if args else message)
+
+    monkeypatch.setattr(td.log, "warning", fake_warning)
+    return messages
+
+
 def test_check_klayout_not_installed(monkeypatch):
     """check_installation raises when KLayout is not on PATH.
 
@@ -62,7 +122,7 @@ def test_runner_passes_drc_args_to_config(monkeypatch, tmp_path):
     resultsfile = tmp_path / "results.lyrdb"
     captured_config = {}
 
-    def mock_run_drc_on_gds(config):
+    def mock_run_drc_on_gds(config, **_kwargs):
         captured_config["config"] = config
         return DRCResults.load(filepath / "drc_results.lyrdb")
 
@@ -101,7 +161,7 @@ def fake_run(cmd, capture_output):
     monkeypatch.setattr(f"{KLAYOUT_PLUGIN_PATH}.drc.drc.run", fake_run)
     monkeypatch.setattr(
         f"{KLAYOUT_PLUGIN_PATH}.drc.drc.DRCResults.load",
-        lambda resultsfile: DRCResults(violations_by_category={}),
+        lambda resultsfile, **_kwargs: DRCResults(violations_by_category={}),
     )
 
     config = DRCConfig(
@@ -276,7 +336,7 @@ def run(
         """Calls DRCRunner.run with dummy run_drc_on_gds()"""
 
         # monkeypatch run_drc_on_gds() since the test machines do not have KLayout installed
-        def mock_run_drc_on_gds(config):
+        def mock_run_drc_on_gds(config, **_kwargs):
             return DRCResults.load(filepath / "drc_results.lyrdb")
 
         monkeypatch.setattr(f"{KLAYOUT_PLUGIN_PATH}.drc.drc.run_drc_on_gds", mock_run_drc_on_gds)
@@ -608,3 +668,27 @@ def test_parse_violation_value_unknown_type(self):
         """Test parsing unknown violation type."""
         with pytest.raises(ValueError):
             parse_violation_value("unknown: (1.0,2.0)")
+
+    def test_results_warn_without_limit(self, monkeypatch, tmp_path):
+        """Warn when no limit is set and a category exceeds the threshold."""
+
+        warnings = _capture_log_warnings(monkeypatch)
+        monkeypatch.setattr(
+            f"{KLAYOUT_PLUGIN_PATH}.drc.results.UNLIMITED_VIOLATION_WARNING_COUNT",
+            3,
+        )
+        results_path = _write_results_file(tmp_path, num_items=4)
+        results = DRCResults.load(results_path)
+        assert results["min_width"].count == 4
+        assert len(warnings) == 1
+        assert "many markers (4)" in warnings[0]
+
+    def test_results_warn_when_limit_truncates(self, monkeypatch, tmp_path):
+        """Warn when the global limit removes markers."""
+
+        warnings = _capture_log_warnings(monkeypatch)
+        results_path = _write_results_file(tmp_path, category="overflow", num_items=4)
+        results = DRCResults.load(results_path, max_results=2)
+        assert results["overflow"].count == 2
+        assert len(warnings) == 1
+        assert "only the first 2" in warnings[0]

tidy3d/plugins/klayout/drc/README.md

@@ -10,6 +10,8 @@ For a full quickstart example, please see [this quickstart notebook](https://git
 
 - Run DRC on GDS files or Tidy3D objects ([Geometry](https://docs.flexcompute.com/projects/tidy3d/en/latest/api/_autosummary/tidy3d.Geometry.html), [Structure](https://docs.flexcompute.com/projects/tidy3d/en/latest/api/_autosummary/tidy3d.Structure.html#tidy3d.Structure), or [Simulation](https://docs.flexcompute.com/projects/tidy3d/en/latest/api/_autosummary/tidy3d.Simulation.html#tidy3d.Simulation)) with `DRCRunner.run()`.
 - Load DRC results into a `DRCResults` data structure with `DRCResults.load()`.
+- Limit how many violation markers are loaded by passing `max_results` to `DRCRunner.run()`,
+  `run_drc_on_gds()`, or `DRCResults.load()`.
 
 ## Prerequisites
 
@@ -131,4 +133,13 @@ Results can also be loaded from a KLayout DRC database file with `DRCResults.loa
 from tidy3d.plugins.klayout.drc import DRCResults
 
 print(DRCResults.load("drc_results.lyrdb"))
-```
+```
+
+### Limiting Loaded Results
+
+Large designs can generate an enormous number of violations. Pass the optional `max_results`
+argument to `DRCRunner.run()`, `run_drc_on_gds()`, or `DRCResults.load()` to retain only the first
+`N` markers across all categories. When the option is not set, a warning is emitted if more than
+100,000 markers are present so you can set an appropriate limit. When the option is set and more
+total violations are present than the limit allows, a warning indicates that the results were
+truncated before parsing individual markers.

tidy3d/plugins/klayout/drc/drc.py

@@ -160,6 +160,7 @@ def run(
         td_object_gds_savefile: Path = DEFAULT_GDSFILE,
         resultsfile: Path = DEFAULT_RESULTSFILE,
         drc_args: Optional[dict[str, str]] = None,
+        max_results: Optional[int] = None,
         **to_gds_file_kwargs: Any,
     ) -> DRCResults:
         """Runs KLayout's DRC on a GDS file or a Tidy3D object. The Tidy3D object can be a :class:`.Geometry`, :class:`.Structure`, or :class:`.Simulation`.
@@ -174,6 +175,9 @@ def run(
             The path to save the KLayout DRC results file to. Defaults to ``"drc_results.lyrdb"``.
         drc_args : Optional[dict[str, str]] = None
             Additional key/value pairs passed through to KLayout as ``-rd key=value`` CLI arguments.
+        max_results : Optional[int]
+            Maximum number of markers to load from the results file. ``None`` (default) loads all
+            markers.
         **to_gds_file_kwargs
             Additional keyword arguments to pass to the Tidy3D object-specific ``to_gds_file()`` method.
 
@@ -215,16 +219,22 @@ def run(
             verbose=self.verbose,
             drc_args={} if drc_args is None else drc_args,
         )
-        return run_drc_on_gds(config=config)
+        return run_drc_on_gds(
+            config=config,
+            max_results=max_results,
+        )
 
 
-def run_drc_on_gds(config: DRCConfig) -> DRCResults:
+def run_drc_on_gds(config: DRCConfig, max_results: Optional[int] = None) -> DRCResults:
     """Runs KLayout's DRC on a GDS file.
 
     Parameters
     ----------
     config : :class:`.DRCConfig`
         The configuration for the DRC run.
+    max_results : Optional[int]
+        Maximum number of markers to load from the results file. ``None`` (default) loads all
+        markers.
 
     Returns
     -------
@@ -267,4 +277,7 @@ def run_drc_on_gds(config: DRCConfig) -> DRCResults:
     if config.verbose:
         console.log("KLayout DRC completed successfully.")
 
-    return DRCResults.load(resultsfile=config.resultsfile)
+    return DRCResults.load(
+        resultsfile=config.resultsfile,
+        max_results=max_results,
+    )

tidy3d/plugins/klayout/drc/results.py

@@ -5,20 +5,23 @@
 import re
 import xml.etree.ElementTree as ET
 from pathlib import Path
-from typing import Union
+from typing import Optional, Union
 
 import pydantic.v1 as pd
 
 from tidy3d.components.base import Tidy3dBaseModel, cached_property
 from tidy3d.components.types import Coordinate2D
 from tidy3d.exceptions import FileError
+from tidy3d.log import log
 
 # Types for DRC markers
 DRCEdge = tuple[Coordinate2D, Coordinate2D]
 DRCEdgePair = tuple[DRCEdge, DRCEdge]
 DRCPolygon = tuple[Coordinate2D, ...]
 DRCMultiPolygon = tuple[DRCPolygon, ...]
 
+UNLIMITED_VIOLATION_WARNING_COUNT = 100_000
+
 
 def parse_edge(value: str) -> EdgeMarker:
     """
@@ -290,13 +293,20 @@ def __str__(self) -> str:
         return summary
 
     @classmethod
-    def load(cls, resultsfile: Union[str, Path]) -> DRCResults:
+    def load(
+        cls,
+        resultsfile: Union[str, Path],
+        max_results: Optional[int] = None,
+    ) -> DRCResults:
         """Create a :class:`.DRCResults` instance from a results file.
 
         Parameters
         ----------
         resultsfile : Union[str, Path]
             Path to the KLayout DRC results file.
+        max_results : Optional[int]
+            Maximum number of markers to load from the file. If ``None`` (default), all markers are
+            loaded.
 
         Returns
         -------
@@ -316,16 +326,26 @@ def load(cls, resultsfile: Union[str, Path]) -> DRCResults:
         >>> results = DRCResults.load(resultsfile="drc_results.lyrdb") # doctest: +SKIP
         >>> print(results) # doctest: +SKIP
         """
-        return cls(violations_by_category=violations_from_file(resultsfile=resultsfile))
+        violations = violations_from_file(
+            resultsfile=resultsfile,
+            max_results=max_results,
+        )
+        return cls.construct(violations_by_category=violations)
 
 
-def violations_from_file(resultsfile: Union[str, Path]) -> dict[str, DRCViolation]:
+def violations_from_file(
+    resultsfile: Union[str, Path],
+    max_results: Optional[int] = None,
+) -> dict[str, DRCViolation]:
     """Loads a KLayout DRC results file and returns the results as a dictionary of :class:`.DRCViolation` objects.
 
     Parameters
     ----------
     resultsfile : Union[str, Path]
         Path to the KLayout DRC results file.
+    max_results : Optional[int]
+        Maximum number of markers to load from the file. If ``None`` (default), all markers are
+        loaded.
 
     Returns
     -------
@@ -339,6 +359,9 @@ def violations_from_file(resultsfile: Union[str, Path]) -> dict[str, DRCViolatio
     ET.ParseError
         If the DRC result file is not a valid XML file.
     """
+    if max_results is not None and max_results <= 0:
+        raise ValueError("'max_results' must be a positive integer.")
+
     # Parse the results file
     try:
         xmltree = ET.parse(resultsfile)
@@ -354,20 +377,37 @@ def violations_from_file(resultsfile: Union[str, Path]) -> dict[str, DRCViolatio
         if category_name is None:
             raise FileError("Encountered DRC category without a name in results file.")
         category_name = category_name.strip().strip("'\"")
-        violations[category_name] = DRCViolation(category=category_name, markers=())
+        violations[category_name] = []
+
+    # Prepare the items and warn if necessary
+    items = list(xmltree.getroot().findall(".//item"))
+    total_markers = len(items)
+    if max_results is None and total_markers > UNLIMITED_VIOLATION_WARNING_COUNT:
+        log.warning(
+            f"DRC result file contains many markers ({total_markers}), "
+            "which can affect loading performance. "
+            "Pass 'max_results' to limit loaded results."
+        )
+    elif max_results is not None and total_markers > max_results:
+        log.warning(
+            f"DRC result file contains {total_markers} markers; "
+            f"only the first {max_results} were loaded due to 'max_results'."
+        )
 
     # Parse markers
-    for item in xmltree.getroot().findall(".//item"):
+    for idx, item in enumerate(items):
+        if max_results is not None and idx >= max_results:
+            break
         category_el = item.find("category")
         if category_el is None or category_el.text is None:
             raise FileError("Encountered DRC item without a category in results file.")
         category = category_el.text.strip().strip("'\"")
         value = item.find("values/value").text
         marker = parse_violation_value(value)
-        if category not in violations:
-            violations[category] = DRCViolation(category=category, markers=())
-        violations[category] = DRCViolation(
-            category=category,
-            markers=(*violations[category].markers, marker),
-        )
-    return violations
+        markers = violations.setdefault(category, [])
+        markers.append(marker)
+
+    return {
+        category: DRCViolation(category=category, markers=tuple(markers))
+        for category, markers in violations.items()
+    }