FXC-3547-Adding mode sorting to ModeSpec

[momchil-flex]

Addresses the new suggestion in #1391 to add more advanced ways to sort modes in the mode solver (and mode sources and monitors).

Probably we should still keep the issue open as I should also implement the FilterSpec as discussed in that issue.

@yuanshen-flexcompute does the current PR help with some of your use cases that made you bring it up though?

Greptile Overview

Updated On: 2025-10-07 08:34:20 UTC

Summary

This PR adds advanced mode sorting functionality to the Tidy3D mode solver by introducing a new `ModeSortSpec` class and integrating it into the `ModeSpec` configuration. The implementation addresses user feedback from issue #1391 where users were spending significant time manually filtering and sorting modes after simulation.

The core addition is the ModeSortSpec class which allows users to sort modes by various physical properties including effective index (n_eff), loss coefficient (k_eff), TE polarization fraction (TE_fraction), waveguide TE/TM fractions (wg_TE_fraction, wg_TM_fraction), and mode area (mode_area). Users can specify sorting direction (ascending/descending) and optionally provide a reference value to sort by proximity to that value.

The implementation integrates seamlessly into the existing mode solver pipeline through a new _sort_modes() method that operates on ModeSolverData after polarization filtering but before frequency tracking. The sorting is applied per-frequency to handle cases where mode properties vary across the frequency spectrum. A mutual exclusivity validator ensures that the new sort_spec cannot be used simultaneously with the legacy filter_pol field to prevent conflicting behaviors.

To support the sorting functionality, three new property aliases were added to the ModeData class: TE_fraction, wg_TE_fraction, and wg_TM_fraction. These provide more convenient access to polarization fraction data and are used by the modes_info method for consistency. The ModeSortSpec class is properly exported through the main __init__.py file to make it available in the public API.

Important Files Changed

Changed Files

Filename	Score	Overview
tidy3d/components/mode_spec.py	4/5	Adds new ModeSortSpec class and sort_spec field to ModeSpec with mutual exclusivity validation
tidy3d/components/mode/mode_solver.py	4/5	Implements core sorting logic with _sort_modes() method and helper functions for mode reordering
tidy3d/components/data/monitor_data.py	5/5	Adds convenience property aliases for polarization fractions used by sorting functionality
tidy3d/init.py	5/5	Exports new ModeSortSpec class to public API
tests/test_plugins/test_mode_solver.py	4/5	Adds comprehensive tests for mode sorting functionality and mutual exclusivity validation

Confidence score: 4/5

• This PR is generally safe to merge with well-structured implementation and comprehensive testing
• Score reflects the complexity of the mode manipulation logic and integration with existing pipeline components
• Pay close attention to the mode solver implementation file and validator logic in the mode_spec file

Sequence Diagram

sequenceDiagram
    participant User
    participant ModeSpec
    participant ModeSolver
    participant ModeSolverData
    participant sort_methods as _sort_modes & _apply_mode_reorder

    User->>ModeSpec: "Create ModeSpec with sort_spec"
    ModeSpec->>ModeSpec: "Validate sort_spec parameters"
    Note over ModeSpec: Check key is in SORT_KEYS<br/>Ensure filter_pol and sort_spec not both used

    User->>ModeSolver: "Create ModeSolver with ModeSpec"
    ModeSolver->>ModeSolver: "Initialize and validate"

    User->>ModeSolver: "Call solve() or data_raw"
    ModeSolver->>ModeSolver: "_data_on_yee_grid()"
    ModeSolver->>ModeSolver: "_colocate_data() if needed"
    ModeSolver->>ModeSolver: "_normalize_modes()"
    ModeSolver->>ModeSolver: "_filter_polarization()"
    ModeSolver->>sort_methods: "_sort_modes(mode_solver_data)"
    
    sort_methods->>ModeSolverData: "getattr(mode_solver_data, sort_spec.key)"
    Note over sort_methods: Get metric values (n_eff, k_eff, etc.)
    
    sort_methods->>sort_methods: "Apply reference_value if provided"
    Note over sort_methods: metric = abs(metric - reference_value)
    
    sort_methods->>sort_methods: "Loop over frequencies"
    loop For each frequency
        sort_methods->>sort_methods: "np.argsort(metric values)"
        sort_methods->>sort_methods: "Reverse if direction='descending'"
        sort_methods->>sort_methods: "_apply_mode_reorder(data, freq_idx, sort_inds)"
        sort_methods->>ModeSolverData: "Reorder field components and n_complex"
        Note over ModeSolverData: data.values[..., freq_idx, :] = data.values[..., freq_idx, sort_inds]
    end
    
    alt If track_freq is set and multiple frequencies
        ModeSolver->>ModeSolverData: "mode_solver_data.overlap_sort(track_freq)"
        Note over ModeSolverData: Apply frequency-based mode tracking
    end

    ModeSolver-->>User: "Return sorted ModeSolverData"

Loading

[greptile-apps]

_{5 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[github-actions]

Diff Coverage

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

• tidy3d/init.py (100%)
• tidy3d/components/data/dataset.py (100%)
• tidy3d/components/data/monitor_data.py (97.8%): Missing lines 2221,2262
• tidy3d/components/mode/data/sim_data.py (100%)
• tidy3d/components/mode/mode_solver.py (100%)
• tidy3d/components/mode_spec.py (97.6%): Missing lines 321

Summary

• Total: 162 lines
• Missing: 3 lines
• Coverage: 98%

tidy3d/components/data/monitor_data.py

Lines 2217-2225

  2217         num_freqs, num_modes = sort_inds_2d.shape
  2218         modify_data = {}
  2219         for key, data in self.data_arrs.items():
  2220             if "mode_index" not in data.dims or "f" not in data.dims:
! 2221                 continue
  2222             dims_orig = data.dims
  2223             f_coord = data.coords["f"]
  2224             slices = []
  2225             for ifreq in range(num_freqs):

---

Lines 2258-2266

  2258         """
  2259 
  2260         # Return the original data if no new sorting / tracking required
  2261         if track_freq is None and sort_spec is None:
! 2262             return self
  2263 
  2264         num_freqs = self.n_eff["f"].size
  2265         num_modes = self.n_eff["mode_index"].size
  2266         all_inds = np.arange(num_modes)

---

tidy3d/components/mode_spec.py

Lines 317-325

  317         if self.track_freq is not None:
  318             return self.track_freq
  319         if self.sort_spec is not None:
  320             return self.sort_spec.track_freq
! 321         return None
  322 
  323 
  324 class ModeSpec(AbstractModeSpec):
  325     """

---

[momchil-flex]

Actually @caseyflex @yuanshen-flexcompute hold off the review, I'm going to generalize this further to include filtering.

[momchil-flex]

Okay now you can review. :) I've introduced advanced filter/sort options. You can select to filter by a given property so that values e.g. over or under a reference value will be put first in the list; then separately you can sort each of the two groups by some other criterion. I think with this #1391 can be closed - probably one thing that still would be nice to add though is some built-in that computes the field intensity close to the plane boundaries (maybe not just on them, but within some region), which can help with PML mode filtering. Or similarly some metric that computes the field intensity within a given region.

[yuanshen-flexcompute]

Nice! Generally very clean and easy to use. I have some comments:

• sort_order should default to ascending (currently descending) when sort_reference is not None. This way, values closest to the reference are returned first.
• Sorting after filtering doesn't seem to work correctly in some cases. I tested filtering by mode_area under a certain value, then sorting by n_eff closest to a reference value. This resulted in the sorting "breaking" the filtered order in some cases. I have a test notebook if you would like to take a closer look. This actually might be user error. Things look good for now.

Thanks!

[momchil-flex]

Thanks!

• sort_order should default to ascending (currently descending) when sort_reference is not None. This way, values closest to the reference are returned first.

I think that makes sense. Originally I set it to descending, because I thought the default value should be our current default mode sorting, which is by descending n_eff. However I think this is not really needed since that is still the default when sort_key is None, while for other applications ascending likely makes more sense as you point out, especially if you also supply a sort_reference.

• Sorting after filtering doesn't seem to work correctly in some cases. I tested filtering by mode_area under a certain value, then sorting by n_eff closest to a reference value. This resulted in the sorting "breaking" the filtered order in some cases. I have a test notebook if you would like to take a closer look.

Yeah, thanks, a minimally reproducible example would be great! I added a test for that but it might not be comprehensive enough so this will help fix and extend testing.

[caseyflex]

This is a great addition! It mostly looks good, just a few comments

[momchil-flex]

The issue that @yuanshen-flexcompute thought he was seeing was related to the fact that by default, we have track_freq="central". We already explain in the description of ModeSpec.filter_spec that if it's used in combination with track_freq, the filter/sort will only be correct at the tracked frequency, while at other frequencies it could be overwritten by the sort. To me this sounds like the correct behavior, but it could be confusing to a user. To obtain strictly sorted modes at every frequency, they'd have to set track_freq=None. What do you think @caseyflex ?

[caseyflex]

The issue that @yuanshen-flexcompute thought he was seeing was related to the fact that by default, we have track_freq="central". We already explain in the description of ModeSpec.filter_spec that if it's used in combination with track_freq, the filter/sort will only be correct at the tracked frequency, while at other frequencies it could be overwritten by the sort. To me this sounds like the correct behavior, but it could be confusing to a user. To obtain strictly sorted modes at every frequency, they'd have to set track_freq=None. What do you think @caseyflex ?

I see, that’s a good point. So this new sorting is frequency independent, while track_freq couples different frequencies. Currently, track_freq is applied second. Since they are both sorting, it could be nice to unify their api in mode_spec. Or could just validate that only one or the other is used.

[momchil-flex]

I see, that’s a good point. So this new sorting is frequency independent, while track_freq couples different frequencies. Currently, track_freq is applied second. Since they are both sorting, it could be nice to unify their api in mode_spec. Or could just validate that only one or the other is used.

So I don't think that's a good idea, because a typical use case could very well be that you roughly want a certain filtering, but you also want the modes to be continuously ordered across frequencies, that is to say the latter is still more important than the strict filtering. Like, I want to have my TM mode come first at my central frequency, but if there's a crossing somewhere I don't want to force the reshuffling.

By the way the reason I still filter all modes rather than the tracked frequency only, if track_freq is set, is to help with the reshuffling after. Since we don't compute the full outer_dot among each frequency (too expensive), the reshuffling works well only if the modes are already roughly ordered.

So having thought about this I think indeed moving track_freq into the sort_freq would make sense (deprecating the direct mode_spec.track_freq) because indeed it is also related to how modes are sorted. So at least if everything is kept in the same place, hopefully it will be less confusing.

[caseyflex]

@momchil-flex this might be another case for having a list of filter_specs, so the order in which they are applied is transparent

[momchil-flex]

Well tracking first and then sorting doesn't really make sense as an option, right? Because that would just sort. You just need to define whether you want to track in the end or not.

[momchil-flex]

@greptile take a look

[greptile-apps]

Looking at this PR, I can see it's a comprehensive implementation of advanced mode sorting functionality for Tidy3D's mode solver. This is a significant enhancement that addresses user needs for better mode filtering and sorting capabilities.

Key Changes Analysis

1. New ModeSortSpec Class

• Provides fine-grained control over mode filtering and sorting
• Supports filtering modes into groups based on thresholds
• Allows sorting within groups by various physical properties
• Includes frequency tracking capabilities

2. Integration with Existing ModeSpec

• Adds sort_spec field to ModeSpec
• Maintains backward compatibility while deprecating filter_pol and track_freq
• Includes mutual exclusivity validation to prevent conflicting configurations

3. Enhanced Mode Data Properties

• Adds convenient aliases (TE_fraction, TM_fraction, etc.) for easier access to polarization data
• Improves the mode sorting infrastructure with better reordering methods

Areas of Attention

Mode Solver Implementation - The core sorting logic in mode_solver.py is complex and handles multiple scenarios. The _apply_mode_reorder method correctly preserves data integrity across all frequency-mode combinations.

Validation Logic - The mutual exclusivity validator properly prevents users from accidentally using both old and new filtering approaches simultaneously.

Backward Compatibility - The deprecation warnings are well-implemented, and the legacy behavior is preserved while guiding users toward the new API.

Schema Updates - All JSON schemas have been properly updated to reflect the new structure.

Strengths

1. Well-structured API: The ModeSortSpec class provides intuitive parameters for complex sorting scenarios
2. Comprehensive testing: The test coverage includes various sorting combinations and edge cases
3. Proper deprecation handling: Clear warnings and backward compatibility preservation
4. Performance considerations: The implementation avoids unnecessary recomputation when no sorting is needed

The implementation looks solid and follows good software engineering practices. The sorting functionality will significantly improve the user experience for mode analysis workflows.

[momchil-flex]

Alright so I did some major updates based on previous discussions @caseyflex :

• Probably most notably, I introduced sort_spec.track_freq and deprecated mode_spec.track_freq, because indeed the tracking can be thought of part of the sorting. This did lead to various headaches though specifically when leaving mode_spec.track_freq as allowed but deprecated. The way this works now is that mode_spec.track_freq=None by default, but if it is set, then it will take precedence (but a deprecation warning will be printed).
• I also exposed standalone ModeData.sort_modes and ModeSimulationData.sort_modes methods that the user can utilize to reorder modes after the solve.
• I refactored everything to avoid in-place operations. I also realized that we need a method that explicitly iterates over all data attributes of a dataset, because previously we were not reordering things like n_group or the amps if present, because we were manually iterating only over the field data and the grid corrections. The one place where this was problematic previously from the user side was overlap_sort performed after the mode solving @dbochkov-flexcompute . Now it should be all good and I also unified the mode reordering among the filter_pol, track_freq and the sort_spec methods.
• Probably the one remaining awkwardness that I'm not 100% happy with is that overlap_sort and sort_modes will always set the mode_data.monitor.mode_spec.track_freq to None, and set the sort_spec.track_freq to the one that matches the tracking that was done. This can lead to a situation where the mode_data.monitor.mode_spec is not the same as the one of the monitor that initially generated it if track_freq was defined there in the deprecated way. However, this resulted in the most straightforward usage for everything else.
• Expanded the test coverage quite a bit.

[caseyflex]

Great, thanks @momchil-flex !

Probably the one remaining awkwardness that I'm not 100% happy with is that overlap_sort and sort_modes will always set the mode_data.monitor.mode_spec.track_freq to None, and set the sort_spec.track_freq to the one that matches the tracking that was done. This can lead to a situation where the mode_data.monitor.mode_spec is not the same as the one of the monitor that initially generated it if track_freq was defined there in the deprecated way. However, this resulted in the most straightforward usage for everything else.

At least this only happens if the deprecated api is used. Actually, I had a similar concern. Say the original monitor.mode_spec filters based on one key, and then the user calls sort_modes with another filter key for example. Then the updated monitor.mode_spec will only contain the second filter key, and if you use that monitor in a fresh simulation, it will give different results.

More broadly, this difficulty comes up repeatedly due to our concept of MonitorData containing a Dataset as well as a Monitor. For example, ModeData._isel is private because it results in inconsistency between the monitor and the data. Also in EMESimulationData, we have _extract_mode_solver_data, which makes an attempt to update the monitor in a reasonable way, but is not guaranteed that the Monitor would reproduce the MonitorData if used in a fresh simulation. It would be extremely useful to have isel functionality for all of our Datasets.

Probably to avoid this, we would have to shift a lot of functionality out of the MonitorData classes into the Dataset classes, so that then MonitorData is really just a Dataset together with a Monitor that gives rise to it. And the methods that process the Dataset return a new Dataset that doesn't necessarily correspond to any Monitor.

Not suggesting changing this here, just bringing it up as a recurring difficulty.

[yuanshen-flexcompute]

Latest changes looks good to me. Thanks!

[momchil-flex]

Actually, I had a similar concern. Say the original monitor.mode_spec filters based on one key, and then the user calls sort_modes with another filter key for example. Then the updated monitor.mode_spec will only contain the second filter key, and if you use that monitor in a fresh simulation, it will give different results.

Good thought - actually in this case, the sorting errors, because we explicitly error if both filter_pol and sort_spec.filter_key or sort_spec.sort_key are defined, so when the sorting function tries to update monitor.mode_spec.sort_spec, it errors. Which is good, right?

The reason I couldn't do the same for track_freq is that its default is not None but instead "central". And I want to move this default to sort_spec.track_freq, but then still keep the deprecated mode_spec.track_freq option (which I set to None as default). Basically, the current defaults are track_freq = None and sort_spec = td.ModeSortSpec(track_freq="central"). If I add a validator that those two can't be defined simultaneously, it will slightly break down compatibility, because where users had just e.g. mode_spec = td.ModeSpec(track_freq="lowest"), they would then have to do one of these: mode_spec = td.ModeSpec(track_freq="lowest", sort_spec=None), or mode_spec = td.ModeSpec(sort_spec=td.ModeSortSpec(track_freq="lowest")). To be honest... it almost sounds worth it though to avoid the other internal headaches.