feat(autograd): extend gaussian filter padding and inverse design support

[yaugenst-flex]

Our inverse design helpers like the ErosionDilationPenalty currently depend on the two filter types we implement, conic and circular. Neither of these are separable, and combined with the somewhat slow convolve implementation in autograd, this can be a big bottleneck depending on kernel size and filter volume. While we've had a Gaussian filter primitive for a while, it didn't support all padding modes (in the VJP) and hence wasn't used in those inverse design helpers. This PR adds the missing VJP implementations and threads the now-supported Gaussian filter through those routines.

Here is a quick benchmark (this includes fwd+bwd passes):

Greptile Overview

Updated On: 2025-11-04 12:51:40 UTC

Greptile Summary

This PR extends Gaussian filter support in the autograd inverse design system by implementing VJP (vector-jacobian product) gradients for all padding modes.

Key Changes:

• Implemented complete VJP for gaussian_filter primitive supporting all padding modes: constant, nearest, mirror, reflect, and wrap (tidy3d/plugins/autograd/primitives/misc.py:79-117)
• VJP uses cached weight matrices per axis for efficiency and applies transposed filters in reverse order for gradient computation
• Added GaussianFilter class to inverse design filters with padding mode mapping and configurable sigma scaling (tidy3d/plugins/autograd/invdes/filters.py:144-194)
• Default sigma_scale of 0.445 empirically tuned to match conic kernel response
• Extended KernelType to include "gaussian" option alongside existing "circular" and "conic" types
• Updated tests to remove skip marks for previously unsupported padding modes and added parametrized tests covering Gaussian filter across all inverse design helpers

Minor Issues Found:

• Docstring formatting inconsistency: mix of single and double backticks in filter docstrings (tidy3d/plugins/autograd/invdes/filters.py:291-305)
• The normalize parameter inherited from AbstractFilter is ignored by GaussianFilter but not documented in the class docstring

The implementation correctly integrates with existing inverse design components like FilterAndProject and ErosionDilationPenalty, enabling faster filtering operations through separable Gaussian kernels.

Confidence Score: 4/5

• This PR is safe to merge with only minor documentation improvements recommended
• The implementation is mathematically sound with proper VJP gradient computation, comprehensive test coverage across all padding modes, and clean integration with existing inverse design infrastructure. Two minor style issues were identified (docstring formatting and missing documentation for ignored parameter), but these don't affect functionality. The gradient check order reduction from 2 to 1 is reasonable for performance given the more complex VJP implementation.
• No files require special attention - all changes are well-tested and properly integrated

Important Files Changed

File Analysis

Filename	Score	Overview
tidy3d/plugins/autograd/primitives/misc.py	4/5	Added VJP implementation for gaussian_filter supporting all padding modes (constant, nearest, mirror, reflect, wrap). Implementation uses cached weight matrices for efficiency. Minor documentation gaps for normalize parameter handling.
tidy3d/plugins/autograd/invdes/filters.py	4/5	Added GaussianFilter class with padding mode mapping and sigma scaling. Minor docstring formatting inconsistencies and normalize parameter not explicitly documented as ignored.
tests/test_plugins/autograd/primitives/test_misc.py	5/5	Test updated to cover all padding modes (removed skip marks for nearest/mirror) and reduced gradient check order from 2 to 1 for performance.

Sequence Diagram

sequenceDiagram
    participant User
    participant FilterAndProject
    participant GaussianFilter
    participant gaussian_filter_primitive
    participant scipy_gaussian_filter
    participant VJP

    User->>FilterAndProject: __call__(array)
    FilterAndProject->>GaussianFilter: make_filter(filter_type="gaussian")
    GaussianFilter->>GaussianFilter: _apply_filter(array, size_px)
    GaussianFilter->>GaussianFilter: map padding mode (edge→nearest, symmetric→mirror)
    GaussianFilter->>GaussianFilter: compute sigma from radius_px and sigma_scale
    GaussianFilter->>gaussian_filter_primitive: gaussian_filter(array, sigma, mode, truncate, cval)
    gaussian_filter_primitive->>scipy_gaussian_filter: forward pass
    scipy_gaussian_filter-->>gaussian_filter_primitive: filtered result
    gaussian_filter_primitive-->>GaussianFilter: filtered array
    GaussianFilter-->>FilterAndProject: filtered array
    FilterAndProject->>FilterAndProject: tanh_projection(filtered, beta, eta)
    FilterAndProject-->>User: projected result
    
    Note over User,VJP: Backward Pass (Gradient Computation)
    User->>VJP: grad output
    VJP->>VJP: _gaussian_filter_vjp()
    VJP->>VJP: compute weight matrices per axis
    loop For each axis (reversed)
        VJP->>VJP: _gaussian_weight_matrix(length, sigma, mode, truncate, order, cval)
        VJP->>VJP: tensordot(weights.T, grad)
    end
    VJP-->>User: input gradient

Loading

[greptile-apps]

Additional Comments (1)

1. tidy3d/plugins/autograd/invdes/filters.py, line 291-297 (link)

   style: Inconsistent docstring formatting: conic uses double backticks while circular and gaussian use single backticks

   Context Used: Rule from dashboard - Adhere to the established docstring format and use correct syntax (e.g., double backticks for code) ... (source)

_{9 files reviewed, 2 comments}

_{Edit Code Review Agent Settings | Greptile}

[github-actions]

Diff Coverage

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

• tidy3d/plugins/autograd/invdes/filters.py (88.2%): Missing lines 172,177,181,188
• tidy3d/plugins/autograd/primitives/misc.py (79.7%): Missing lines 15-18,26-29,43,90,92,101
• tidy3d/plugins/autograd/types.py (100%)

Summary

• Total: 94 lines
• Missing: 16 lines
• Coverage: 82%

tidy3d/plugins/autograd/invdes/filters.py

Lines 168-185

  168     )
  169 
  170     @staticmethod
  171     def get_kernel(size_px: Iterable[int], normalize: bool) -> NDArray:
! 172         raise NotImplementedError("GaussianFilter does not build an explicit kernel.")
  173 
  174     def _apply_filter(self, array: NDArray, size_px: tuple[int, ...]) -> NDArray:
  175         radius_px = np.maximum((np.array(size_px, dtype=float) - 1.0) / 2.0, 0.0)
  176         if radius_px.size == 0:
! 177             return array
  178 
  179         mode = _GAUSSIAN_PADDING_MAP.get(self.padding)
  180         if mode is None:
! 181             raise ValueError(
  182                 f"Unsupported padding mode '{self.padding}' for gaussian filter; "
  183                 f"supported modes are {tuple(_GAUSSIAN_PADDING_MAP)}."
  184             )

---

Lines 184-192

  184             )
  185 
  186         sigma = tuple(float(self.sigma_scale * r) if r > 0 else 0.0 for r in radius_px)
  187         if not any(sigma):
! 188             return array
  189 
  190         kwargs: dict[str, Any] = {"mode": mode, "truncate": float(self.truncate)}
  191         if mode == "constant":
  192             kwargs["cval"] = 0.0

---

tidy3d/plugins/autograd/primitives/misc.py

Lines 11-22

  11 
  12 def _normalize_sequence(value: float | Sequence[float], ndim: int) -> tuple[float, ...]:
  13     """Convert a scalar or sequence into a tuple of length ``ndim``."""
  14     if isinstance(value, Iterable) and not np.isscalar(value):
! 15         value_tuple = tuple(value)
! 16         if len(value_tuple) != ndim:
! 17             raise ValueError(f"Sequence length {len(value_tuple)} does not match ndim {ndim}.")
! 18         return tuple(float(v) for v in value_tuple)
  19     return (float(value),) * ndim
  20 
  21 
  22 def _normalize_modes(mode: str | Sequence[str], ndim: int) -> tuple[str, ...]:

---

Lines 22-33

  22 def _normalize_modes(mode: str | Sequence[str], ndim: int) -> tuple[str, ...]:
  23     """Normalize a padding mode argument into a tuple of strings with length ``ndim``."""
  24     if isinstance(mode, str):
  25         return (mode,) * ndim
! 26     mode_tuple = tuple(mode)
! 27     if len(mode_tuple) != ndim:
! 28         raise ValueError(f"Mode sequence length {len(mode_tuple)} does not match ndim {ndim}.")
! 29     return mode_tuple
  30 
  31 
  32 @cache
  33 def _gaussian_weight_matrix(

---

Lines 39-47

  39     cval: float,
  40 ) -> np.ndarray:
  41     """Return the 1-D Gaussian filter matrix used along a single axis."""
  42     if sigma <= 0.0:
! 43         return np.eye(length)
  44     eye = np.eye(length, dtype=float)
  45     weights = scipy.ndimage.gaussian_filter1d(
  46         eye,
  47         sigma=sigma,

---

Lines 86-96

  86     cval_seq = _normalize_sequence(cval, ndim)
  87     mode_seq = _normalize_modes(mode, ndim)
  88 
  89     if any(int(o) != 0 for o in order_seq):
! 90         raise NotImplementedError("gaussian_filter VJP currently supports only order=0.")
  91     if kwargs:
! 92         raise NotImplementedError(
  93             f"gaussian_filter VJP does not support additional keyword arguments: {tuple(kwargs)}"
  94         )
  95 
  96     def vjp(g):

---

Lines 97-105

   97         grad = np.asarray(g)
   98         for axis in reversed(range(ndim)):
   99             sigma_axis = float(sigma_seq[axis])
  100             if sigma_axis <= 0.0:
! 101                 continue
  102             mode_axis = mode_seq[axis]
  103             truncate_axis = float(truncate_seq[axis])
  104             order_axis = int(order_seq[axis])
  105             cval_axis = float(cval_seq[axis])

---

[groberts-flex]

this is great! overall, left one comment and a clarification to make sure I'm understanding correctly what's going on.