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/monitor_data.py (100%)
• tidy3d/components/mode/mode_solver.py (98.1%): Missing lines 1405
• tidy3d/components/mode_spec.py (100%)

Summary

• Total: 89 lines
• Missing: 1 line
• Coverage: 98%

tidy3d/components/mode/mode_solver.py

  1401                 # Boolean mask for modes in the first group
  1402                 if filter_spec.filter_order == "over":
  1403                     mask_first = (vals_filt >= filter_spec.filter_reference).to_numpy()
  1404                 else:
! 1405                     mask_first = (vals_filt <= filter_spec.filter_reference).to_numpy()
  1406                 group1 = all_inds[mask_first]
  1407                 group2 = all_inds[~mask_first]
  1408             else:
  1409                 group1 = all_inds

---

[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

[caseyflex]

two comments here:

indexing xarray values directly is not robust compared to something like this: data.loc[dict(frequency=ifreq)] = data.sel(frequency=ifreq).isel(mode=sort_inds)
so that might be preferred, unless the performance is a big issue

also, maybe we can be clear that we mutate the provided mode_solver_data

[momchil-flex]

So in general I agree with you but it turns out it's hard to mutate an xarray array in place. This

data.loc[{"f": freq}] = data.sel(f=freq).isel(mode_index=sort_inds)

fails with IndexError: dimension coordinate 'mode_index' conflicts between indexed and indexing objects. I tinkered around with it a bit and it seems like the only way (or at least the main way) would be to reassign a data array rather than sort it in-place. Then this is complicated further by the fact that we need to modify multiple data arrays, so I think it might have to go through a dict representation and then update the mode data with the updated dict.

Overall it seems a bit too much given that the current approach works... Btw note that I just separated it into a function, otherwise it's the same sorting approach we've been using forever in filter_pol.

[caseyflex]

Can we have a public function sort_modes(mode_solver_data, filter_spec)? That way we can apply it automatically, or the user can do additional filtering as a post-processing step?

[momchil-flex]

So I was thinking about this myself and I was just worried about some potential issues / confusion, but I think after thinking about it more carefully now it could be worth it. I think the method could be a call on ModeSolverData directly, like new_data = old_data.filter_sort(filter_spec=...) (this one will not be modifying in-place). Some things to pay attention to:

• ModeSolverData contains a ModeSolverMonitor which stores the original mode_spec.filter_spec. So the sorting method should also update that.
• Something that can be slightly confusing, but I think it's ok, is that the user can use this method to explore various sortings locally, but then they can't use the sorted data as a mode source or monitor directly. They need to pass the desired filter spec to the mode spec of their object.
• For ModeSimulationData, they would have to do something like new_modes = mode_sim_data.modes.filter_sort(...), i.e. we don't create new ModeSimulationData, just a new ModeSolverData.

Does this plan sound good? Any other caveats that come to mind?

[caseyflex]

Oh yeah, this reminds me, what if the user wants to filter on multiple properties? Should we have an ordered list of filter_specs?

then if you already had a filter spec, and you filter again, it can concatenate to this list. Then the entire list can be used in a mode source.

otherwise sounds good. Although there’s no reason not to also allow calling this directly on ModeSimulationData I suppose

[momchil-flex]

Right, I guess I could add it to ModeSimulationData - I just need to update the associated simulation too...

List of filtering sounds... too much honestly :D

[caseyflex]

Maybe, but it’s needed if you want every use case to be translatable into a modeSource

[caseyflex]

does this test the case where we filter only and don't sort? maybe need to check test coverage again after filtering was added?

[momchil-flex]

I'll expand the testing, but generally test coverage does not check whether different pathways are tested together vs separately. It just tests whether every line is called by some test. So e.g. a test that does filtering only and a test that does sorting only will already pass test coverage.

[caseyflex]

I think line 1424 of mode_solver.py was not tested

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