feat(tidy3d): FXC-3961-faster-convolutions-for-tidy-3-d-plugins-autograd-filters

marcorudolphflex · marcorudolphflex · commit 24ea2aca2850 · 2025-11-06T15:11:16.000+01:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -55,6 +55,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Unified run submission API: `web.run(...)` is now a container-aware wrapper that accepts a single simulation or arbitrarily nested containers (`list`, `tuple`, `dict` values) and returns results in the same shape.
 - `web.Batch(ComponentModeler)` and `web.Job(ComponentModeler)` native support
 - Simulation data of batch jobs are now automatically downloaded upon their individual completion in `Batch.run()`, avoiding waiting for the entire batch to reach completion.
+- Improved speed of autograd tracing for convolutions.
 
 ### Fixed
 - Ensured the legacy `Env` proxy mirrors `config.web` profile switches and preserves API URL.
diff --git a/tests/test_plugins/autograd/test_functions.py b/tests/test_plugins/autograd/test_functions.py
@@ -117,12 +117,15 @@ def test_convolve_val(self, rng, mode, padding, ary_size, kernel_size, square_ke
         """Test convolution values against SciPy for various modes, padding, array sizes, and kernel sizes."""
         x, k = self._ary_and_kernel(rng, ary_size, kernel_size, square_kernel)
 
-        if mode in ("full", "same"):
-            pad_widths = [(k // 2, k // 2) for k in k.shape]
+        if mode in ["same", "full"]:
+            if mode == "same":
+                pad_widths = [(k // 2, k // 2) for k in k.shape]
+            else:
+                pad_widths = [(k - 1, k - 1) for k in k.shape]
             x_padded = x
             for axis, pad_width in enumerate(pad_widths):
                 x_padded = pad(x_padded, pad_width, mode=padding, axis=axis)
-            conv_sp = convolve_sp(x_padded, k, mode="valid" if mode == "same" else mode)
+            conv_sp = convolve_sp(x_padded, k, mode="valid")
         else:
             conv_sp = convolve_sp(x, k, mode=mode)
 
@@ -168,6 +171,52 @@ def test_kernel_array_dimension_mismatch(self):
             convolve(self.array, kernel_mismatch)
 
 
+class TestConvolveAxes:
+    @pytest.mark.parametrize("mode", ["valid", "same", "full"])
+    @pytest.mark.parametrize("padding", ["constant", "edge"])
+    def test_convolve_axes_val(self, rng, mode, padding):
+        """Test convolution with explicit axes against NumPy implementations."""
+        array = rng.random((2, 5))
+        kernel = rng.random((3, 3))
+        axes = ([1], [1])
+
+        conv_td = convolve(array, kernel, padding=padding, mode=mode, axes=axes)
+
+        working_array = array
+        scipy_mode = mode
+        if mode in ("same", "full"):
+            pad_width = kernel.shape[1] // 2 if mode == "same" else kernel.shape[1] - 1
+            working_array = pad(array, (pad_width, pad_width), mode=padding, axis=1)
+            scipy_mode = "valid"
+
+        working_array_np = np.asarray(working_array)
+        kernel_np = np.asarray(kernel)
+        conv_length = np.convolve(working_array_np[0], kernel_np[0], mode=scipy_mode).shape[0]
+
+        expected = np.empty((array.shape[0], kernel.shape[0], conv_length))
+        for i in range(array.shape[0]):
+            for j in range(kernel.shape[0]):
+                expected[i, j] = np.convolve(
+                    working_array_np[i],
+                    kernel_np[j],
+                    mode=scipy_mode,
+                )
+
+        npt.assert_allclose(conv_td, expected, atol=1e-12)
+
+    def test_convolve_axes_grad(self, rng):
+        """Test gradients of convolution when specific axes are provided."""
+        array = rng.random((2, 5))
+        kernel = rng.random((3, 3))
+        check_grads(convolve, modes=["rev"], order=2)(
+            array,
+            kernel,
+            padding="constant",
+            mode="valid",
+            axes=([1], [1]),
+        )
+
+
 @pytest.mark.parametrize(
     "op,sp_op",
     [
diff --git a/tidy3d/plugins/autograd/functions.py b/tidy3d/plugins/autograd/functions.py
@@ -7,7 +7,7 @@
 import numpy as onp
 from autograd import jacobian
 from autograd.extend import defvjp, primitive
-from autograd.scipy.signal import convolve as convolve_ag
+from autograd.numpy.fft import fftn, ifftn
 from autograd.scipy.special import logsumexp
 from autograd.tracer import getval
 from numpy.lib.stride_tricks import sliding_window_view
@@ -37,6 +37,139 @@
 ]
 
 
+def _normalize_axes(
+    ndim_array: int,
+    ndim_kernel: int,
+    axes: Union[tuple[Iterable[int], Iterable[int]], None],
+) -> tuple[tuple[int, ...], tuple[int, ...]]:
+    """Normalize the axes specification for convolution."""
+
+    def _normalize_single_axis(ax: int, ndim: int, kind: str) -> int:
+        if not isinstance(ax, int):
+            raise TypeError(
+                f"Axis indices must be integers, got {ax!r} of type {type(ax).__name__}."
+            )
+
+        if not -ndim <= ax < ndim:
+            raise ValueError(f"Invalid axis {ax} for {kind} with ndim {ndim}.")
+        return ax + ndim if ax < 0 else ax
+
+    if axes is None:
+        if ndim_array != ndim_kernel:
+            raise ValueError(
+                "Kernel dimensions must match array dimensions when 'axes' is not provided, "
+                f"got array ndim {ndim_array} and kernel ndim {ndim_kernel}."
+            )
+        axes_array = tuple(range(ndim_array))
+        axes_kernel = tuple(range(ndim_kernel))
+        return axes_array, axes_kernel
+
+    if len(axes) != 2:
+        raise ValueError("'axes' must be a tuple of two iterable collections of axis indices.")
+
+    axes_array_raw, axes_kernel_raw = axes
+
+    axes_array = tuple(_normalize_single_axis(ax, ndim_array, "array") for ax in axes_array_raw)
+    axes_kernel = tuple(_normalize_single_axis(ax, ndim_kernel, "kernel") for ax in axes_kernel_raw)
+
+    if len(axes_array) != len(axes_kernel):
+        raise ValueError(
+            "The number of convolution axes for the array and kernel must be the same, "
+            f"got {len(axes_array)} and {len(axes_kernel)}."
+        )
+
+    if len(set(axes_array)) != len(axes_array) or len(set(axes_kernel)) != len(axes_kernel):
+        raise ValueError("Convolution axes must be unique for both the array and the kernel.")
+
+    return axes_array, axes_kernel
+
+
+def _fft_convolve_general(
+    array: NDArray,
+    kernel: NDArray,
+    axes_array: tuple[int, ...],
+    axes_kernel: tuple[int, ...],
+    mode: Literal["full", "valid"],
+) -> NDArray:
+    """Perform convolution using FFT along the specified axes."""
+
+    num_conv_axes = len(axes_array)
+
+    if num_conv_axes == 0:
+        array_shape = array.shape
+        kernel_shape = kernel.shape
+        result = np.multiply(
+            array.reshape(array_shape + (1,) * kernel.ndim),
+            kernel.reshape((1,) * array.ndim + kernel_shape),
+        )
+        return result.reshape(array_shape + kernel_shape)
+
+    ignore_axes_array = tuple(ax for ax in range(array.ndim) if ax not in axes_array)
+    ignore_axes_kernel = tuple(ax for ax in range(kernel.ndim) if ax not in axes_kernel)
+
+    new_order_array = ignore_axes_array + axes_array
+    new_order_kernel = ignore_axes_kernel + axes_kernel
+
+    array_reordered = np.transpose(array, new_order_array) if array.ndim else array
+    kernel_reordered = np.transpose(kernel, new_order_kernel) if kernel.ndim else kernel
+
+    num_batch_array = len(ignore_axes_array)
+    num_batch_kernel = len(ignore_axes_kernel)
+
+    array_conv_shape = array_reordered.shape[num_batch_array:]
+    kernel_conv_shape = kernel_reordered.shape[num_batch_kernel:]
+
+    if any(d <= 0 for d in array_conv_shape + kernel_conv_shape):
+        raise ValueError("Convolution dimensions must be positive; got zero-length axis.")
+
+    fft_axes = tuple(range(-num_conv_axes, 0))
+    fft_shape = tuple(
+        int(array_dim + kernel_dim - 1)
+        for array_dim, kernel_dim in zip(array_conv_shape, kernel_conv_shape)
+    )
+
+    array_fft = fftn(array_reordered, fft_shape, axes=fft_axes)
+    kernel_fft = fftn(kernel_reordered, fft_shape, axes=fft_axes)
+
+    if num_batch_kernel:
+        array_batch_shape = array_fft.shape[:num_batch_array]
+        conv_shape = array_fft.shape[num_batch_array:]
+        array_fft = np.reshape(
+            array_fft,
+            array_batch_shape + (1,) * num_batch_kernel + conv_shape,
+        )
+
+    if num_batch_array:
+        kernel_batch_shape = kernel_fft.shape[:num_batch_kernel]
+        conv_shape = kernel_fft.shape[num_batch_kernel:]
+        kernel_fft = np.reshape(
+            kernel_fft,
+            (1,) * num_batch_array + kernel_batch_shape + conv_shape,
+        )
+
+    full_result = ifftn(array_fft * kernel_fft, fft_shape, axes=fft_axes)
+
+    if mode == "full":
+        result = full_result
+    elif mode == "valid":
+        valid_slices = [slice(None)] * full_result.ndim
+        for axis_offset, (array_dim, kernel_dim) in enumerate(
+            zip(array_conv_shape, kernel_conv_shape)
+        ):
+            start = int(min(array_dim, kernel_dim) - 1)
+            length = int(abs(array_dim - kernel_dim) + 1)
+            axis = full_result.ndim - num_conv_axes + axis_offset
+            valid_slices[axis] = slice(start, start + length)
+        result = full_result[tuple(valid_slices)]
+    else:
+        raise ValueError(f"Unsupported convolution mode '{mode}'.")
+
+    if not np.iscomplexobj(array) and not np.iscomplexobj(kernel):
+        result = np.real(result)
+
+    return result
+
+
 def _get_pad_indices(
     n: int,
     pad_width: tuple[int, int],
@@ -180,19 +313,27 @@ def convolve(
     if any(k % 2 == 0 for k in kernel.shape):
         raise ValueError(f"All kernel dimensions must be odd, got {kernel.shape}.")
 
-    if kernel.ndim != array.ndim and axes is None:
-        raise ValueError(
-            f"Kernel dimensions must match array dimensions, got kernel {kernel.shape} and array {array.shape}."
-        )
-
-    if mode in ("same", "full"):
-        kernel_dims = kernel.shape if axes is None else [kernel.shape[d] for d in axes[1]]
-        pad_widths = [(ks // 2, ks // 2) for ks in kernel_dims]
-        for axis, pad_width in enumerate(pad_widths):
-            array = pad(array, pad_width, mode=padding, axis=axis)
-        mode = "valid" if mode == "same" else mode
-
-    return convolve_ag(array, kernel, axes=axes, mode=mode)
+    axes_array, axes_kernel = _normalize_axes(array.ndim, kernel.ndim, axes)
+
+    working_array = array
+    effective_mode = mode
+
+    if mode in ["same", "full"]:
+        for ax_array, ax_kernel in zip(axes_array, axes_kernel):
+            if mode == "same" and kernel.shape[ax_kernel] % 2 == 0:
+                raise ValueError(
+                    f"Even-sized kernel along axis {ax_kernel} not supported for 'same' mode."
+                )
+            pad_width = (
+                kernel.shape[ax_kernel] // 2 if mode == "same" else kernel.shape[ax_kernel] - 1
+            )
+            if pad_width > 0:
+                working_array = pad(
+                    working_array, (pad_width, pad_width), mode=padding, axis=ax_array
+                )
+        effective_mode = "valid"
+
+    return _fft_convolve_general(working_array, kernel, axes_array, axes_kernel, effective_mode)
 
 
 def _get_footprint(size, structure, maxval):
