greg comments

Author: yaugenst-flex

tests/test_plugins/autograd/test_functions.py

@@ -248,6 +248,46 @@ def test_morphology_val_structure_grad(
             check_grads(op, modes=["rev"], order=1)(x, structure=k, mode=mode)
 
 
+class TestMorphology1D:
+    """Test morphological operations with 1D-like structuring elements."""
+
+    @pytest.mark.parametrize("h, w", [(1, 3), (3, 1), (1, 5), (5, 1)])
+    def test_1d_structuring_elements(self, rng, h, w):
+        """Test grey dilation with 1D-like structuring elements on 2D arrays."""
+        x = rng.random((8, 8))
+
+        # Test with size parameter
+        size_tuple = (h, w)
+        result_size = grey_dilation(x, size=size_tuple)
+
+        # Verify output shape matches input
+        assert result_size.shape == x.shape
+
+        # Verify that dilation actually increases values (or keeps them the same)
+        assert np.all(result_size >= x)
+
+        # Test that we can also use structure parameter with 1D-like arrays
+        structure = np.ones((h, w))
+        result_struct = grey_dilation(x, structure=structure)
+        assert result_struct.shape == x.shape
+
+    def test_1d_gradient_flow(self, rng):
+        """Test gradient flow through 1D-like structuring elements."""
+        x = rng.random((6, 6))
+
+        # Test horizontal 1D structure
+        check_grads(lambda x: grey_dilation(x, size=(1, 3)), modes=["rev"], order=1)(x)
+
+        # Test vertical 1D structure
+        check_grads(lambda x: grey_dilation(x, size=(3, 1)), modes=["rev"], order=1)(x)
+
+        # Test with structure parameter
+        struct_h = np.ones((1, 3))
+        struct_v = np.ones((3, 1))
+        check_grads(lambda x: grey_dilation(x, structure=struct_h), modes=["rev"], order=1)(x)
+        check_grads(lambda x: grey_dilation(x, structure=struct_v), modes=["rev"], order=1)(x)
+
+
 class TestMorphologyExceptions:
     """Test exceptions in morphological operations."""
 
@@ -264,6 +304,13 @@ def test_even_structure_dimensions(self, rng):
         with pytest.raises(ValueError, match="Structuring element dimensions must be odd"):
             grey_dilation(x, structure=k_even)
 
+    def test_both_size_and_structure(self, rng):
+        """Test that an exception is raised when both size and structure are provided."""
+        x = rng.random((5, 5))
+        k = np.ones((3, 3))
+        with pytest.raises(ValueError, match="Cannot specify both size and structure"):
+            grey_dilation(x, size=3, structure=k)
+
 
 @pytest.mark.parametrize(
     "array, out_min, out_max, in_min, in_max, expected",

tidy3d/plugins/autograd/functions.py

@@ -208,6 +208,8 @@ def _get_footprint(size, structure, maxval):
     """Helper to generate the morphological footprint from size or structure."""
     if size is None and structure is None:
         raise ValueError("Either size or structure must be provided.")
+    if size is not None and structure is not None:
+        raise ValueError("Cannot specify both size and structure.")
     if structure is None:
         size_np = onp.atleast_1d(size)
         shape = (size_np[0], size_np[-1]) if size_np.size > 1 else (size_np[0], size_np[0])
@@ -238,8 +240,11 @@ def grey_dilation(
         The input array to perform grey dilation on.
     size : Union[Union[int, tuple[int, int]], None] = None
         The size of the structuring element. If None, `structure` must be provided.
+        If a single integer is provided, a square structuring element is created.
+        For 1D arrays, use a tuple (size, 1) or (1, size) for horizontal or vertical operations.
     structure : Union[np.ndarray, None] = None
         The structuring element. If None, `size` must be provided.
+        For 1D operations on 2D arrays, use a 2D structure with one dimension being 1.
     mode : PaddingType = "reflect"
         The padding mode to use.
     maxval : float = 1e4
@@ -286,6 +291,10 @@ def _vjp_maker_dilation(ans, array, size=None, structure=None, *, mode="reflect"
     is_max_mask = (dilated_windows == output_reshaped).astype(onp.float64)
 
     # normalize the gradient for cases where multiple elements are the maximum.
+    # When multiple elements in a window equal the maximum value, the gradient
+    # is distributed equally among them. This ensures gradient conservation.
+    # Note: Values can never exceed maxval in the output since we add structure
+    # values (capped at maxval) to the input array values.
     multiplicity = onp.sum(is_max_mask, axis=(-2, -1), keepdims=True)
     is_max_mask /= onp.maximum(multiplicity, 1)