Skip to content

FXC-5153: making trimesh and matplotlib lazy#3213

Open
momchil-flex wants to merge 2 commits intodevelopfrom
momchil/lazy_imports
Open

FXC-5153: making trimesh and matplotlib lazy#3213
momchil-flex wants to merge 2 commits intodevelopfrom
momchil/lazy_imports

Conversation

@momchil-flex
Copy link
Collaborator

@momchil-flex momchil-flex commented Jan 28, 2026

trimesh was 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 trimesh and matplotlib to reduce initial import overhead for the tidy3d package.

Key Changes:

  • Fixed trimesh lazy loading by moving TrimeshType import check from module-level to usage-time
  • Implemented comprehensive lazy loading for matplotlib across visualization components
  • Changed arrow_style from module-level constant to lazy-loading function
  • Added _ensure_tidy3d_style() mechanism to defer matplotlib style configuration until first plot
  • Moved all matplotlib imports to either TYPE_CHECKING blocks (for type hints) or function-level imports
  • Added comprehensive test suite using subprocess isolation to verify lazy import behavior

Impact:
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

  • This PR is safe to merge with one minor style improvement suggested
  • The implementation is well-structured and comprehensive, with proper TYPE_CHECKING usage and function-level imports. Tests validate the behavior. Minor inefficiency in is_valid_color that could be optimized but doesn't affect correctness.
  • Check tidy3d/components/viz/visualization_spec.py for the redundant import optimization

Important Files Changed

Filename Overview
tests/test_package/test_lazy_imports.py New comprehensive test suite for lazy imports using subprocess isolation
tidy3d/components/types/base.py Changed matplotlib Axes import to TYPE_CHECKING pattern for lazy loading
tidy3d/components/viz/init.py Added lazy style application mechanism to defer matplotlib initialization
tidy3d/components/viz/styles.py Converted arrow_style from module-level constant to lazy-loading function
tidy3d/components/viz/visualization_spec.py Made is_color_like import lazy but imports it twice per validation call

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 plot
Loading

Note

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 tidy3d overhead by converting many module-level optional imports (primarily matplotlib, plus trimesh type helpers) into function-level or TYPE_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) so apply_tidy3d_params() is not executed at tidy3d.components.viz import time, and updates plotting utilities and geometry helpers (e.g., arrow_style() now lazily constructs the ArrowStyle) to use these lazy paths.

Improves resilience and behavior when optional deps are missing/broken (e.g., TrimeshType fallback when trimesh import fails, VisualizationSpec.is_valid_color lazy color validation, and grayscale fallback when matplotlib is unavailable), and adds tests/test_lazy_imports.py subprocess tests asserting matplotlib/scipy/trimesh/vtk (and networkx) are not imported on top-level tidy3d import while confirming matplotlib loads when plotting is invoked.

Written by Cursor Bugbot for commit 574d06f. This will update automatically on new commits. Configure here.

@momchil-flex momchil-flex marked this pull request as ready for review January 28, 2026 12:32
@momchil-flex momchil-flex changed the title fix: making trimesh and matplotlib lazy FXC-5153: making trimesh and matplotlib lazy Jan 28, 2026
Copy link

@greptile-apps greptile-apps bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

5 files reviewed, 1 comment

Edit Code Review Agent Settings | Greptile

Copy link

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Collaborator

@yaugenst-flex yaugenst-flex left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall looks like a good direction, cleaner overall than before. Some comments to iron out.

Copy link
Contributor

@marcorudolphflex marcorudolphflex left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link

@cursor cursor bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@momchil-flex
Copy link
Collaborator Author

@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.

@yaugenst-flex yaugenst-flex self-requested a review January 30, 2026 12:20
Copy link
Collaborator

@yaugenst-flex yaugenst-flex left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yeah LGTM!

@momchil-flex
Copy link
Collaborator Author

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 ?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants