feat(tidy3d): FXC-4260-allow-users-to-swap-custom-cmaps-in-field-visualization

[marcorudolphflex]

Enables the support of custom colormaps in field visualizations.
Usage example:

fig, axes = plt.subplots(1, 3, figsize=(12, 4))
sim_data.plot_field(
    "adjoint_fld_2",
    "Ex",
    val="real",
    z=0.0,
    ax=axes[0],
    cmap=None,
)
axes[0].set_title("Default colormap")

sim_data.plot_field(
    "adjoint_fld_2",
    "Ex",
    val="real",
    z=0.0,
    ax=axes[1],
    cmap="cividis",
)
axes[1].set_title("Custom 'cividis' colormap")

custom_cmap = mcolors.LinearSegmentedColormap.from_list("two", ["green", "blue"])
sim_data.plot_field(
    "adjoint_fld_2",
    "Ex",
    val="real",
    z=0.0,
    ax=axes[2],
    cmap=custom_cmap,
)
axes[2].set_title("Custom user colormap")
plt.show()

Greptile Overview

Greptile Summary

Adds support for custom colormaps in field visualization methods across the codebase. Users can now pass a cmap parameter (either a string colormap name or a matplotlib Colormap object) to plot_field methods, which overrides the default colormap inferred from the field type.

Key Changes:

• Added cmap parameter to plot_field, plot_field_monitor_data, and plot_scalar_array methods in AbstractYeeGridSimulationData
• Extended support to specialized simulation data classes: ModeSimulationData, HeatChargeSimulationData, ModeSolver, and RectangularDielectric
• Implemented proper fallback logic: custom cmap overrides default colormap inferred from cmap_type
• Added comprehensive tests for both string colormap names and custom Colormap objects
• Updated CHANGELOG.md with user-facing feature description

Issues Found:

• Import ordering issue in mode_solver.py: import statement appears immediately after TYPE_CHECKING block without proper separation

Confidence Score: 4/5

• Safe to merge with one minor syntax fix for import ordering
• The implementation is well-structured and comprehensive, with proper parameter passing through the call chain, good test coverage, and appropriate documentation. However, there is one import ordering issue in mode_solver.py where an import statement is incorrectly placed immediately after a TYPE_CHECKING block, which should be fixed before merging.
• Pay attention to tidy3d/components/mode/mode_solver.py - fix import ordering issue on line 82-84

Important Files Changed

File Analysis

Filename	Score	Overview
tests/test_data/test_sim_data.py	5/5	Added test for custom colormap parameter with string and LinearSegmentedColormap
tidy3d/components/data/sim_data.py	5/5	Added cmap parameter to all plot methods with proper handling and fallback to defaults
tidy3d/components/mode/mode_solver.py	3/5	Added cmap parameter to plot_field method; import ordering issue detected
tidy3d/components/tcad/data/sim_data.py	5/5	Added cmap parameter with proper fallback to 'coolwarm' default and TYPE_CHECKING import

Sequence Diagram

sequenceDiagram
    participant User
    participant plot_field
    participant plot_field_monitor_data
    participant plot_scalar_array
    participant XarrayPlot
    
    User->>plot_field: plot_field(monitor_name, field_name, cmap="viridis")
    Note over plot_field: Accept cmap parameter<br/>(str or Colormap object)
    plot_field->>plot_field: load_field_monitor(monitor_name)
    plot_field->>plot_field_monitor_data: plot_field_monitor_data(..., cmap=cmap)
    Note over plot_field_monitor_data: Process field data<br/>and infer cmap_type
    plot_field_monitor_data->>plot_scalar_array: plot_scalar_array(..., cmap_type=cmap_type, cmap=cmap)
    Note over plot_scalar_array: Determine default_cmap<br/>from cmap_type if cmap is None
    plot_scalar_array->>plot_scalar_array: cmap_to_use = default_cmap if cmap is None else cmap
    plot_scalar_array->>XarrayPlot: field_data.plot(..., cmap=cmap_to_use)
    XarrayPlot-->>User: Returns matplotlib axes with field visualization

Loading

[greptile-apps]

_{8 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[github-actions]

Diff Coverage

Diff: origin/develop...HEAD, staged and unstaged changes

• tidy3d/components/data/sim_data.py (83.3%): Missing lines 829
• tidy3d/components/mode/data/sim_data.py (100%)
• tidy3d/components/mode/mode_solver.py (100%)
• tidy3d/components/tcad/data/sim_data.py (100%)
• tidy3d/plugins/waveguide/rectangular_dielectric.py (100%)

Summary

• Total: 11 lines
• Missing: 1 line
• Coverage: 90%

tidy3d/components/data/sim_data.py

Lines 825-833

  825             vmax = np.pi
  826             center = False
  827             eps_reverse = False
  828         else:
! 829             default_cmap = None
  830 
  831         cmap_to_use = default_cmap if cmap is None else cmap
  832 
  833         # plot the field

---

[yaugenst-flex]

Nice!