FXC-5153: making trimesh and matplotlib lazy #3213
Conversation
momchil-flex
changed the title
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| else: | ||
| if is_color_like is not None and not is_color_like(value): | ||
| from matplotlib.colors import is_color_like | ||
|
|
||
| if not is_color_like(value): |
There was a problem hiding this comment.
redundant import - is_color_like is imported twice per validation call (once on line 18 during initialization, again here on line 31)
| else: | |
| if is_color_like is not None and not is_color_like(value): | |
| from matplotlib.colors import is_color_like | |
| if not is_color_like(value): | |
| if not MATPLOTLIB_IMPORTED: | |
| log.warning( | |
| "matplotlib was not successfully imported, but is required " | |
| "to validate colors in the VisualizationSpec. The specified colors " | |
| "have not been validated." | |
| ) | |
| return value | |
| from matplotlib.colors import is_color_like | |
| if not is_color_like(value): | |
| raise ValueError(f"{value} is not a valid plotting color") |
Prompt To Fix With AI
This is a comment left during a code review.
Path: tidy3d/components/viz/visualization_spec.py
Line: 30:33
Comment:
redundant import - `is_color_like` is imported twice per validation call (once on line 18 during initialization, again here on line 31)
```suggestion
if not MATPLOTLIB_IMPORTED:
log.warning(
"matplotlib was not successfully imported, but is required "
"to validate colors in the VisualizationSpec. The specified colors "
"have not been validated."
)
return value
from matplotlib.colors import is_color_like
if not is_color_like(value):
raise ValueError(f"{value} is not a valid plotting color")
```
How can I resolve this? If you propose a fix, please make it concise.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.
| "TensorReal", | ||
| "TrackFreq", | ||
| "TrimeshType", | ||
| "Undefined", |
There was a problem hiding this comment.
TrimeshType was removed from the exported types but it's not clear if it was actually being used by external code. If TrimeshType was part of the public API (it's in all), removing it could be a breaking change. Consider deprecating it first or ensuring it's not used elsewhere.
| xyz, | ||
| ) | ||
| from tidy3d.components.types.third_party import TrimeshType | ||
| from tidy3d.components.types.utils import _add_schema |
There was a problem hiding this comment.
The check_import function in tidy3d.packaging.py actually imports the module via import_module, which means that when third_party.py uses check_import("trimesh"), it will import trimesh at module load time. This defeats the purpose of making trimesh lazy. The third_party.py file should also be updated to use TYPE_CHECKING pattern, or mesh.py should import trimesh.Trimesh directly instead of using TrimeshType.
yaugenst-flex
left a comment
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.
| import matplotlib | ||
| matplotlib.use('Agg') # Use non-interactive backend for testing | ||
| ax = sim.plot(z=0) |
There was a problem hiding this comment.
This doesn't actually prove plotting triggers the import because matplotlib is imported explicitly before the plotting call?
| _arrow_style = None | ||
|
|
||
|
|
||
| def arrow_style(): |
There was a problem hiding this comment.
This changed from value to function and is sort of a public-ish function since it's exposed in viz/__init__.py. Could just not do that instead, I doubt anyone uses it..
| from tidy3d.components.viz import _ensure_tidy3d_style | ||
|
|
||
| _ensure_tidy3d_style() |
There was a problem hiding this comment.
Maybe this is on purpose, but it constitutes a behavioral change. This means that the matplotlib style is no longer applied at import time and only happens when this function is used, meaning if a user passes an existing ax, our style won't be applied. This behavior might actually be better though, not sure.
| tidycomplex, | ||
| xyz, | ||
| ) | ||
| from tidy3d.components.types.third_party import TrimeshType |
There was a problem hiding this comment.
Technically a breaking change...
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 1 potential issue.
Bugbot Autofix is OFF. To automatically fix reported issues with Cloud Agents, enable Autofix in the Cursor dashboard.
| return True | ||
| except ImportError: | ||
| return False | ||
| return find_spec(module_name) is not None |
There was a problem hiding this comment.
Missing import error handling after check_import behavior change
Medium Severity
The change from import_module() to find_spec() in check_import removes the validation that a module can actually be imported. The old implementation caught ImportError if the import failed. Now, check_import("trimesh") returns True if the spec exists, but third_party.py then calls import trimesh without try/except. If trimesh is installed but broken (missing dependencies, etc.), this causes an unhandled ImportError crash where the old code would gracefully fall back to TrimeshType = Any.
marcorudolphflex
left a comment
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
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
Improves import performance by lazily loading heavy optional deps and tightening import checks.
matplotlibimports across plotting code by moving to function-level andTYPE_CHECKINGblocks; adds_ensure_tidy3d_style()and convertsarrow_styleto a lazy factorytrimeshusage by removing eagerTrimeshTypeexposure and related imports from__all__packaging.check_import()to useimportlib.util.find_specto check availability without importingScene,Simulation, EME, unstructured triangular, geometry, viz utils/styles/descartes) to importmatplotlibon demand and handle fallbacks when unavailablevisualization_specand avoids importingAxesat runtime in typestests/test_package/test_lazy_imports.pyto verify no top-level imports ofmatplotlib,scipy,trimesh,vtkand that plotting triggersmatplotlibWritten by Cursor Bugbot for commit 0514703. This will update automatically on new commits. Configure here.