FXC-5153: making trimesh and matplotlib lazy#3213
FXC-5153: making trimesh and matplotlib lazy#3213momchil-flex wants to merge 2 commits intodevelopfrom
Conversation
There was a problem hiding this comment.
Pull request overview
This PR aims to make imports of trimesh and matplotlib lazy to reduce the initial import time of tidy3d. The bulk of the changes focus on deferring matplotlib imports until they are actually needed for plotting operations.
Changes:
- Converted matplotlib imports from module-level to lazy imports using TYPE_CHECKING for type hints and runtime imports within functions
- Changed arrow_style from a module-level constant to a lazy-loading function
- Deferred apply_tidy3d_params() call to only execute when plotting functions are first used via _ensure_tidy3d_style()
- Attempted to make trimesh lazy by removing TrimeshType from exports (incomplete implementation)
Reviewed changes
Copilot reviewed 14 out of 14 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| tidy3d/components/viz/visualization_spec.py | Made is_color_like import lazy with deferred check on first use |
| tidy3d/components/viz/styles.py | Converted arrow_style from module constant to lazy-loading function |
| tidy3d/components/viz/descartes.py | Moved matplotlib.patches and matplotlib.path imports inside functions |
| tidy3d/components/viz/axes_utils.py | Added _ensure_tidy3d_style() call to make_ax() for deferred style application |
| tidy3d/components/viz/init.py | Deferred apply_tidy3d_params() and added _ensure_tidy3d_style() mechanism |
| tidy3d/components/types/base.py | Used TYPE_CHECKING for Axes import to avoid runtime matplotlib import |
| tidy3d/components/types/init.py | Removed TrimeshType from exports (breaks existing usage) |
| tidy3d/components/scene.py | Used TYPE_CHECKING for matplotlib imports and updated type hints to use pipe operator |
| tidy3d/components/simulation.py | Moved matplotlib imports inside plotting methods |
| tidy3d/components/eme/simulation.py | Moved matplotlib imports inside plotting methods |
| tidy3d/components/geometry/base.py | Moved matplotlib imports inside methods and updated arrow_style to function call |
| tidy3d/components/data/unstructured/triangular.py | Used TYPE_CHECKING for Triangulation import and moved pyplot import |
| tidy3d/components/tcad/simulation/heat_charge.py | Moved colormaps import inside conditional block where used |
| tests/test_package/test_lazy_imports.py | Added comprehensive tests for lazy imports of matplotlib, scipy, trimesh, vtk, and networkx |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
yaugenst-flex
left a comment
There was a problem hiding this comment.
Overall looks like a good direction, cleaner overall than before. Some comments to iron out.
marcorudolphflex
left a comment
There was a problem hiding this comment.
LGTM overall. I do not really have something to add but I wouldn't really mind the theoretical user-facing changes for these minor cases too much. Would also fix the testing issue pointed out by @yaugenst-flex
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 2 potential issues.
Bugbot Autofix is OFF. To automatically fix reported issues with Cloud Agents, enable Autofix in the Cursor dashboard.
|
@yaugenst-flex I pushed one last fix to the scipy import test; it seems the code coverage was interfering with the environment so now the lazy import tests are using an entirely clean env. The only thing left imo is to hear your opinion on the last unresolved comment above. And if you want to take another look, go ahead. |
65cb143 to
36d9077
Compare
|
Finally figured out why scipy lazy import test keeps failing in CI: recent versions of xarray auto-load their backends via entrypoints, which includes the scipy backend. I had to upgrade xarray locally to reproduce the error finally. Since xarray is a core dependence I think we have no choice but to accept that scipy will not be imported lazily if available. However many things should still work without scipy if not available (there's a separate test for that). Sounds good @yaugenst-flex ? |
trimeshwas not being lazy because of TrimeshType. That was an easy fix.The bulk of the changes were needed to make matplotlib lazy. This doesn't look too bad, what do you think @yaugenst ?
On my PC import time dropped from 1.9s to 1.2s.
Greptile Overview
Greptile Summary
This PR implements lazy imports for
trimeshandmatplotlibto reduce initial import overhead for the tidy3d package.Key Changes:
trimeshlazy loading by movingTrimeshTypeimport check from module-level to usage-timematplotlibacross visualization componentsarrow_stylefrom module-level constant to lazy-loading function_ensure_tidy3d_style()mechanism to defer matplotlib style configuration until first plotTYPE_CHECKINGblocks (for type hints) or function-level importsImpact:
The changes successfully defer heavy optional dependencies until they're actually needed, improving import performance for users who don't require visualization features.
Confidence Score: 4/5
is_valid_colorthat could be optimized but doesn't affect correctness.tidy3d/components/viz/visualization_spec.pyfor the redundant import optimizationImportant Files Changed
Sequence Diagram
sequenceDiagram participant User participant Tidy3d participant VizModule as viz/__init__.py participant MakeAx as make_ax() participant EnsureStyle as _ensure_tidy3d_style() participant Matplotlib User->>Tidy3d: import tidy3d Note over Tidy3d,VizModule: No matplotlib import at this stage Tidy3d-->>User: Module loaded (fast) User->>Tidy3d: sim.plot(z=0) Tidy3d->>MakeAx: ax is None, create one MakeAx->>EnsureStyle: Check if style applied alt First plot call EnsureStyle->>Matplotlib: import matplotlib EnsureStyle->>Matplotlib: apply_tidy3d_params() Note over EnsureStyle: Set _tidy3d_style_applied = True else Subsequent calls Note over EnsureStyle: Style already applied, skip end MakeAx->>Matplotlib: plt.subplots() Matplotlib-->>MakeAx: ax created MakeAx-->>Tidy3d: Return ax Tidy3d->>Matplotlib: Render plot Matplotlib-->>User: Display plotNote
Medium Risk
Touches multiple plotting code paths and import-time behavior; risk is mainly regressions in visualization or in environments without
matplotlib, but core simulation logic is largely unaffected.Overview
Reduces
import tidy3doverhead by converting many module-level optional imports (primarilymatplotlib, plustrimeshtype helpers) into function-level orTYPE_CHECKING-only imports, ensuring plotting-related modules don’t load unless a plotting API is actually called.Introduces a deferred styling hook (
_ensure_tidy3d_style) soapply_tidy3d_params()is not executed attidy3d.components.vizimport time, and updates plotting utilities and geometry helpers (e.g.,arrow_style()now lazily constructs theArrowStyle) to use these lazy paths.Improves resilience and behavior when optional deps are missing/broken (e.g.,
TrimeshTypefallback whentrimeshimport fails,VisualizationSpec.is_valid_colorlazy color validation, and grayscale fallback whenmatplotlibis unavailable), and addstests/test_lazy_imports.pysubprocess tests assertingmatplotlib/scipy/trimesh/vtk(andnetworkx) are not imported on top-leveltidy3dimport while confirmingmatplotlibloads when plotting is invoked.Written by Cursor Bugbot for commit 574d06f. This will update automatically on new commits. Configure here.