Autograd support for component modelers

[groberts-flex]

Most of the initial code here from @yaugenst-flex - I added/fixed up a few things to get it working with component modeler based optimizations for s params. I tested my patch antenna optimization for terminal component modeler and a modal example.

Greptile Summary

This pull request adds autograd (automatic differentiation) support for component modelers in Tidy3D, specifically enabling optimization workflows for TerminalComponentModeler and ModalComponentModeler. The changes extend the existing autograd infrastructure to handle these higher-level abstractions that manage multiple simulations for electromagnetic component modeling.

The core modification expands the web API's batch categorization system to recognize autograd simulation types, allowing proper routing of forward and backward (adjoint) simulations. The autograd module is enhanced to validate and execute component modelers by leveraging their internal simulation dictionaries through the existing batch execution infrastructure. When a component modeler contains autograd-compatible simulations, it routes through the new autograd pathway; otherwise, it falls back to the standard web API.

The implementation maintains full backward compatibility while adding sophisticated handling for composite simulation objects. Component modelers internally contain multiple simulations that need to be executed in batch and then reconstructed into appropriate data structures. The changes include validation logic to check underlying simulations for autograd compatibility and composition functions to properly reconstruct ComponentModelerData objects from batch execution results.

A comprehensive test suite validates the autograd integration using an offline emulator that generates realistic monitor data without requiring web API calls. The testing covers both modal and terminal component modelers, with the terminal modeler using stubbed S-matrix construction to bypass non-autograd-compatible voltage/current integral calculations.

Important Files Changed

Click to expand

Filename	Score	Overview
tidy3d/web/api/container.py	5/5	Extended BatchCategoryType to include four new autograd simulation categories
CHANGELOG.md	5/5	Added changelog entry documenting autograd support for component modelers
tidy3d/web/api/autograd/autograd.py	4/5	Extended autograd infrastructure to handle component modelers with new validation and execution logic
tests/test_plugins/smatrix/test_component_modeler_autograd.py	4/5	Added comprehensive test suite with offline emulator for component modeler autograd functionality

Confidence score: 4/5

• This PR introduces significant new functionality but appears well-architected with proper fallback mechanisms
• Score reflects the complexity of autograd integration and potential for subtle edge cases in composite simulation handling
• Pay close attention to the autograd module changes and test coverage for component modeler workflows

Sequence Diagram

sequenceDiagram
    participant User
    participant web_run as web.run()
    participant autograd as autograd.py
    participant modeler as ComponentModeler
    participant batch as Batch Processing
    participant server as Tidy3D Server
    participant composer as Data Composer

    User->>web_run: web.run(component_modeler, ...)
    web_run->>autograd: is_valid_for_autograd(modeler)
    autograd->>modeler: Check for traced fields in sim_dict
    modeler-->>autograd: Returns True if autograd-compatible
    autograd-->>web_run: Valid for autograd

    web_run->>autograd: _run_component_modeler()
    autograd->>modeler: modeler.sim_dict
    modeler-->>autograd: Dictionary of simulations
    autograd->>autograd: _run_async(simulations_dict)
    
    autograd->>batch: Create batch from sim_dict
    batch->>server: Upload and run simulations
    server-->>batch: Return SimulationData results
    batch-->>autograd: sim_data_map

    autograd->>composer: _compose_modeler_data_from_sim_map()
    composer->>composer: Create SimulationDataMap
    alt ModalComponentModeler
        composer->>composer: Create ModalComponentModelerData
    else TerminalComponentModeler  
        composer->>composer: Create TerminalComponentModelerData
    end
    composer-->>autograd: ComponentModelerData
    autograd-->>web_run: Final modeler data
    web_run-->>User: ComponentModelerData with gradients

Loading

[greptile-apps]

_{4 files reviewed, 3 comments}

_{Edit Code Review Bot Settings | Greptile}

[daquinteroflex]

Hi @groberts-flex looks good, if you can just address greptile and test the run.py changes then add this to the merge queue later today when you're ready

[github-actions]

Diff Coverage

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

• tidy3d/plugins/smatrix/component_modelers/base.py (100%)
• tidy3d/plugins/smatrix/run.py (84.6%): Missing lines 157-158,160,210
• tidy3d/web/api/autograd/autograd.py (100%)
• tidy3d/web/api/container.py (100%)

Summary

• Total: 33 lines
• Missing: 4 lines
• Coverage: 87%

tidy3d/plugins/smatrix/run.py

  153     from tidy3d.web.api.autograd import autograd as web_ag
  154 
  155     try:
  156         sims = modeler.sim_dict
! 157     except AttributeError:
! 158         return False
  159     if not sims:
! 160         return False
  161 
  162     for sim in sims.values():
  163         if web_ag.is_valid_for_autograd(sim):
  164             return True

---

  206     priority = kwargs.get("priority")
  207     if priority is None:
  208         batch_data = batch.run(path_dir=path_dir)
  209     else:
! 210         batch_data = batch.run(path_dir=path_dir, priority=priority)
  211     return compose_modeler_data_from_batch_data(modeler=modeler, batch_data=batch_data)

---

[daquinteroflex]

Looks good! Exciting.

Have you checked you can run an example normally already?

[tylerflex]

Thanks @yaugenst-flex just a few clarifying questions