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)

[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 5988,6003,6013

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 5984-5992

  5984         component_paths = [
  5985             tuple(path[1:]) for path in derivative_info.paths if path and path[0] == component
  5986         ]
  5987         if not component_paths:
! 5988             return None
  5989 
  5990         axis = component[0]  # f.e. xx -> x
  5991         projected_E = derivative_info.project_der_map_to_axis(axis, "E")
  5992         projected_D = derivative_info.project_der_map_to_axis(axis, "D")

---

Lines 5999-6007

  5999 
  6000         components = self.components
  6001         for field_path in derivative_info.paths:
  6002             if not field_path or field_path[0] not in components:
! 6003                 raise NotImplementedError(
  6004                     f"No derivative defined for '{type(self).__name__}' field: {field_path}."
  6005                 )
  6006 
  6007         vjps: AutogradFieldMap = {}

---

Lines 6009-6017

  6009             comp_info = self._component_derivative_info(
  6010                 derivative_info=derivative_info, component=comp_name
  6011             )
  6012             if comp_info is None:
! 6013                 continue
  6014             comp_vjps = component._compute_derivatives(comp_info)
  6015             for sub_path, value in comp_vjps.items():
  6016                 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.