feat(tidy3d): FXC-3296-autograd-support-for-anisotropic-medium-and-custom-anisotropic-medium

[marcorudolphflex]

Greptile Overview

Greptile Summary

This PR adds autograd support for diagonal AnisotropicMedium and CustomAnisotropicMedium, enabling gradient-based inverse design with anisotropic materials. The implementation delegates derivative computation to individual component media (xx, yy, zz) with proper field projection to ensure only the corresponding axis contributes to each gradient.

Key changes:

• Implemented _compute_derivatives method in AnisotropicMedium that projects electric field maps to individual axes and delegates to component media
• Added helper methods _project_field_map_to_axis and _component_derivative_info for field filtering
• Included comprehensive unit tests with synthetic fields and numerical validation tests comparing adjoint gradients with finite differences
• Updated CHANGELOG to document the new feature and improved error handling for unsupported anisotropic configurations

The implementation correctly inherits to CustomAnisotropicMedium through the class hierarchy. Error handling raises NotImplementedError for unsupported paths (e.g., off-diagonal components), though explicit test coverage for this error case would strengthen the validation.

Confidence Score: 4/5

• This PR is safe to merge with minor recommendation for additional test coverage
• The implementation is technically sound with proper delegation pattern and field projection. Comprehensive numerical tests validate gradient correctness across multiple scenarios. Score reduced by 1 point due to missing explicit test for the error handling mentioned in CHANGELOG (NotImplementedError for unsupported anisotropic configurations)
• tests/test_components/autograd/test_autograd_anisotropic_medium.py could benefit from explicit error handling test

Important Files Changed

File Analysis

Filename	Score	Overview
CHANGELOG.md	5/5	Added documentation for new autograd support for diagonal anisotropic media and improved error handling for unsupported cases
tidy3d/components/medium.py	4/5	Implemented _compute_derivatives method for AnisotropicMedium that delegates to component media with proper field projection - minor test coverage gap for error cases
tests/test_components/autograd/test_autograd_anisotropic_medium.py	5/5	Added unit tests for gradient computation with synthetic fields, covering both AnisotropicMedium and CustomAnisotropicMedium
tests/test_components/autograd/numerical/test_autograd_anisotropic_numerical.py	5/5	Added comprehensive numerical gradient validation tests comparing adjoint gradients with finite differences across multiple scenarios

Sequence Diagram

sequenceDiagram
    participant User as User Code
    participant AAM as AnisotropicMedium
    participant Comp as Component Medium
    participant DI as DerivativeInfo
    
    User->>AAM: _compute_derivatives(derivative_info)
    AAM->>AAM: validate all paths exist in components
    
    loop for each component (xx, yy, zz)
        AAM->>AAM: _component_derivative_info(derivative_info, component)
        AAM->>AAM: _project_field_map_to_axis(E_der_map, axis)
        Note over AAM: Zero out non-matching field components<br/>(e.g., for xx: keep Ex, zero Ey, Ez)
        AAM->>DI: create filtered DerivativeInfo
        AAM->>Comp: component._compute_derivatives(comp_info)
        Comp-->>AAM: component gradients
        AAM->>AAM: prepend component name to gradient keys
    end
    
    AAM-->>User: combined gradients for all components

Loading

Context used:

• Rule from dashboard - Add tests for all public constructors and methods to confirm expected behavior, including effects on... (source)
• Rule from dashboard - Prefer pytest.mark.parametrize over manual for loops to define and run distinct test cases, reducing... (source)

---

Note

Enables autograd for diagonal anisotropic media with axis-aware gradient routing and comprehensive tests.

• Implemented AnisotropicMedium._compute_derivatives() with per-component delegation and added _component_derivative_info() to filter paths and project fields to the relevant axis
• Added DerivativeInfo.project_der_map_to_axis() to zero non-selected components in E_der_map/D_der_map
• Expanded numerical gradient tests to cover anisotropic and custom_anisotropic cases; added unit tests verifying axis-specific gradients and delegation
• Updated CHANGELOG.md to document the new autograd support

^{Written by Cursor Bugbot for commit 4b3e783. This will update automatically on new commits. Configure here.}

[marcorudolphflex]

@greptile

[greptile-apps]

_{4 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{4 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[github-actions]

Diff Coverage

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

• tidy3d/components/autograd/derivative_utils.py (93.8%): Missing lines 738
• tidy3d/components/medium.py (87.0%): Missing lines 6132,6147,6157

Summary

• Total: 39 lines
• Missing: 4 lines
• Coverage: 89%

tidy3d/components/autograd/derivative_utils.py

Lines 734-742

  734             if the requested map is unavailable.
  735         """
  736         field_map = {"E": self.E_der_map, "D": self.D_der_map}.get(field_type)
  737         if field_map is None:
! 738             raise ValueError("field type must be 'D' or 'E'.")
  739 
  740         axis = axis.lower()
  741         projected = dict(field_map)
  742         for dim in "xyz":

---

tidy3d/components/medium.py

Lines 6128-6136

  6128         component_paths = [
  6129             tuple(path[1:]) for path in derivative_info.paths if path and path[0] == component
  6130         ]
  6131         if not component_paths:
! 6132             return None
  6133 
  6134         axis = component[0]  # f.e. xx -> x
  6135         projected_E = derivative_info.project_der_map_to_axis(axis, "E")
  6136         projected_D = derivative_info.project_der_map_to_axis(axis, "D")

---

Lines 6143-6151

  6143 
  6144         components = self.components
  6145         for field_path in derivative_info.paths:
  6146             if not field_path or field_path[0] not in components:
! 6147                 raise NotImplementedError(
  6148                     f"No derivative defined for '{type(self).__name__}' field: {field_path}."
  6149                 )
  6150 
  6151         vjps: AutogradFieldMap = {}

---

Lines 6153-6161

  6153             comp_info = self._component_derivative_info(
  6154                 derivative_info=derivative_info, component=comp_name
  6155             )
  6156             if comp_info is None:
! 6157                 continue
  6158             comp_vjps = component._compute_derivatives(comp_info)
  6159             for sub_path, value in comp_vjps.items():
  6160                 vjps[(comp_name, *sub_path)] = value

---

[groberts-flex]

could this possibly be something that would fit in data_array as a helper method there. I'm thinking it might fit more with that than as a static method in medium

[marcorudolphflex]

Right, it is definitely misplaced here. But as it operates on dict[str, FieldData] and could also be applied to D fields, I would argue it should be moved to derivative info.
Added this to DerivativeInfo:

def project_der_map_to_axis(
        self, axis: xyz, field_type: str = "E"
) -> dict[str, ScalarFieldDataArray] | None:
    """Return a copy of the selected derivative map with only one axis kept.

    Parameters
    ----------
    axis:
        Axis to keep (``"x"``, ``"y"``, ``"z"``, case-insensitive).
    field_type:
        Map selector: ``"E"`` (``self.E_der_map``) or ``"D"`` (``self.D_der_map``).

    Returns
    -------
    dict[str, ScalarFieldDataArray] | None
        Copied map where non-selected components are replaced by zeros, or ``None``
        if the requested map is unavailable.
    """
    field_map = {"E": self.E_der_map, "D": self.D_der_map}.get(field_type)
    if field_map is None:
        raise ValueError("field type must be 'D' or 'E'.")

    axis = axis.lower()
    projected = dict(field_map)
    for dim in "xyz":
        key = field_type + dim
        if key not in field_map:
            continue
        if dim != axis:
            projected[key] = _zeros_like(field_map[key])
        else:
            projected[key] = field_map[key]
    return projected

[groberts-flex]

thanks for the addition! overall, looks like a good approach @marcorudolphflex

left a couple comments/questions!

[groberts-flex]

overall, looks good! do you have plots for the outputs of the numerical tests you could share as well?

[marcorudolphflex]

overall, looks good! do you have plots for the outputs of the numerical tests you could share as well?

Good point to do that. Actually, the gradients are quite off in some cases while being qualitatively okayish. Interestingly, this is also partially the case for (isotropic) CustomMedium, see below. So i'm not sure if this comes from my calculation or from the underlying CustomMedium as my calculation seems valid to me. Any ideas why we observe that?

----------------- isotropic custom medium example
ignore yy/zz, xx is just the isotropic value

[groberts-flex]

thanks for putting those plots together, they look interesting! Talking with Yannick a little bit this morning in our 1:1 and he was saying there is a test case he has for isotropic custom medium that is not working well for autograd. I think it's likely the underlying problem is in there, but might be good to get to the bottom of that. I think he is going to share that one with me and I'll have a look at it to see if I can spot anything. I'll also re-run my other numerical custom medium tests to see if maybe we had a regression somewhere.

[marcorudolphflex]

Looks much better after the fix from https://github.com//pull/3113