feat(invdes): added symmetry functions

Author: mahlau-flex

CHANGELOG.md

@@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 
 ### Added
+- Added `symmetrize_mirror`, `symmetrize_rotation`, `symmetrize_diagonal` functions to the autograd plugin. They can be used for enforcing symmetries in topology optimization.
 
 ### Changed
 - Removed validator that would warn if `PerturbationMedium` values could become numerically unstable, since an error will anyway be raised if this actually happens when the medium is converted using actual perturbation data.

docs/api/plugins/autograd.rst

@@ -85,4 +85,7 @@ Inverse Design
     tidy3d.plugins.autograd.invdes.ramp_projection
     tidy3d.plugins.autograd.invdes.tanh_projection
     tidy3d.plugins.autograd.invdes.smoothed_projection
+    tidy3d.plugins.autograd.invdes.symmetrize_mirror
+    tidy3d.plugins.autograd.invdes.symmetrize_rotation
+    tidy3d.plugins.autograd.invdes.symmetrize_diagonal
 

tests/test_plugins/autograd/invdes/test_symmetries.py

@@ -0,0 +1,181 @@
+from __future__ import annotations
+
+import autograd.numpy as np
+import numpy as onp
+import pytest
+from autograd.test_util import check_grads
+
+from tidy3d.plugins.autograd.invdes.symmetries import (
+    symmetrize_diagonal,
+    symmetrize_mirror,
+    symmetrize_rotation,
+)
+
+# --- Helper Fixtures ---
+
+
+@pytest.fixture
+def square_array():
+    """Returns a random 5x5 array for square tests."""
+    return np.random.randn(5, 5)
+
+
+@pytest.fixture
+def rect_array():
+    """Returns a random 4x6 array for non-square tests."""
+    return np.random.randn(4, 6)
+
+
+# --- Symmetrize Mirror Tests ---
+
+
+@pytest.mark.parametrize("axis", [0, 1, (0, 1)])
+def test_mirror_gradients(axis):
+    """
+    Verifies that the gradient calculation through symmetrize_mirror is correct
+    using finite difference checks provided by autograd.
+    """
+    # Create a random array. Size doesn't need to be square.
+    x = np.random.randn(4, 5)
+
+    # We wrap the function to treat 'axis' as a fixed constant,
+    # testing the gradient only with respect to 'x'.
+    def fun(x):
+        return symmetrize_mirror(x, axis=axis)
+
+    # check_grads verifies analytical grad vs finite difference
+    check_grads(fun, modes=["rev"], order=1)(x)
+
+
+@pytest.mark.parametrize("axis", [0, 1, (0, 1)])
+def test_mirror_values(axis):
+    """Verifies numerical correctness of mirror symmetry."""
+    # Simple 2x2 case
+    # [[1, 2],
+    #  [3, 4]]
+    arr = np.array([[1.0, 2.0], [3.0, 4.0]])
+
+    res = symmetrize_mirror(arr, axis=axis)
+
+    if axis == 0:
+        # Average with vertical flip [[3, 4], [1, 2]]
+        # ([[1, 2], [3, 4]] + [[3, 4], [1, 2]]) / 2 = [[2, 3], [2, 3]]
+        expected = np.array([[2.0, 3.0], [2.0, 3.0]])
+    elif axis == 1:
+        # Average with horizontal flip [[2, 1], [4, 3]]
+        # ([[1, 2], [3, 4]] + [[2, 1], [4, 3]]) / 2 = [[1.5, 1.5], [3.5, 3.5]]
+        expected = np.array([[1.5, 1.5], [3.5, 3.5]])
+    else:  # (0, 1)
+        # Average of all 4 mirror types implied (linear combination reduces to avg of 4 corners)
+        # Result should be constant value 2.5 everywhere for this specific linear gradient input
+        expected = np.full((2, 2), 2.5)
+
+    onp.testing.assert_allclose(res, expected)
+
+
+def test_mirror_shapes_and_errors(rect_array):
+    """Test shape constraints and error handling."""
+    # Should work on rectangular arrays
+    res = symmetrize_mirror(rect_array, axis=0)
+    assert res.shape == rect_array.shape
+
+    # Error: 3D array
+    with pytest.raises(ValueError, match="Need 2d array"):
+        symmetrize_mirror(np.random.randn(2, 2, 2), axis=0)
+
+    # Error: Invalid axis
+    with pytest.raises(ValueError, match="Invalid axis"):
+        symmetrize_mirror(rect_array, axis=2)
+
+    # Error: Invalid tuple
+    with pytest.raises(ValueError, match="Invalid axis"):
+        symmetrize_mirror(rect_array, axis=(0, 0))
+
+
+# --- Symmetrize Rotation Tests ---
+
+
+def test_rotation_gradients(square_array):
+    """Verifies gradients for rotation symmetry."""
+    check_grads(symmetrize_rotation, modes=["rev"], order=1)(square_array)
+
+
+def test_rotation_values():
+    """Verifies numerical correctness of rotation symmetry."""
+    # Input with a single 1 in top-left, 0 elsewhere
+    # [[1, 0],
+    #  [0, 0]]
+    arr = np.zeros((2, 2))
+    arr[0, 0] = 1.0
+
+    res = symmetrize_rotation(arr)
+
+    # The 1 should be distributed to all 4 corners equally
+    expected = np.full((2, 2), 0.25)
+    onp.testing.assert_allclose(res, expected)
+
+
+def test_rotation_invariance(square_array):
+    """The output of symmetrize_rotation should be invariant to further 90deg rotations."""
+    sym = symmetrize_rotation(square_array)
+    rot = np.rot90(sym)
+    onp.testing.assert_allclose(sym, rot, err_msg="Output is not rotationally symmetric")
+
+
+def test_rotation_errors(rect_array):
+    """Test shape constraints for rotation."""
+    # Error: Rectangular array
+    with pytest.raises(ValueError, match="must be square"):
+        symmetrize_rotation(rect_array)
+
+
+# --- Symmetrize Diagonal Tests ---
+
+
+@pytest.mark.parametrize("anti", [False, True])
+def test_diagonal_gradients(square_array, anti):
+    """Verifies gradients for diagonal symmetry."""
+
+    def fun(x):
+        return symmetrize_diagonal(x, anti=anti)
+
+    check_grads(fun, modes=["rev"], order=1)(square_array)
+
+
+def test_diagonal_values():
+    """Verifies numerical correctness of diagonal symmetry."""
+    # [[1, 2],
+    #  [3, 4]]
+    arr = np.array([[1.0, 2.0], [3.0, 4.0]])
+
+    # Main diagonal
+    res_main = symmetrize_diagonal(arr, anti=False)
+    # Transpose is [[1, 3], [2, 4]]
+    # Avg: [[1, 2.5], [2.5, 4]]
+    expected_main = np.array([[1.0, 2.5], [2.5, 4.0]])
+    onp.testing.assert_allclose(res_main, expected_main)
+
+    # Anti diagonal
+    res_anti = symmetrize_diagonal(arr, anti=True)
+    # Anti-transpose logic check:
+    #
+    # Input:
+    # 1 2
+    # 3 4
+    #
+    # Anti-Transpose:
+    # 4 2
+    # 3 1
+    #
+    # Average:
+    # 2.5 2
+    # 3   2.5
+    expected_anti = np.array([[2.5, 2.0], [3.0, 2.5]])
+    onp.testing.assert_allclose(res_anti, expected_anti)
+
+
+def test_diagonal_errors(rect_array):
+    """Test shape constraints for diagonal."""
+    # Error: Rectangular array
+    with pytest.raises(ValueError, match="must be square"):
+        symmetrize_diagonal(rect_array)

tidy3d/plugins/autograd/__init__.py

@@ -36,6 +36,10 @@
     make_filter_and_project,
     make_gaussian_filter,
     ramp_projection,
+    smoothed_projection,
+    symmetrize_diagonal,
+    symmetrize_mirror,
+    symmetrize_rotation,
     tanh_projection,
 )
 from .primitives import gaussian_filter, interpolate_spline
@@ -79,6 +83,10 @@
     "scalar_objective",
     "smooth_max",
     "smooth_min",
+    "smoothed_projection",
+    "symmetrize_diagonal",
+    "symmetrize_mirror",
+    "symmetrize_rotation",
     "tanh_projection",
     "threshold",
     "trapz",

tidy3d/plugins/autograd/invdes/__init__.py

@@ -17,6 +17,7 @@
 )
 from .penalties import ErosionDilationPenalty, make_curvature_penalty, make_erosion_dilation_penalty
 from .projections import ramp_projection, smoothed_projection, tanh_projection
+from .symmetries import symmetrize_diagonal, symmetrize_mirror, symmetrize_rotation
 
 __all__ = [
     "CircularFilter",
@@ -35,5 +36,8 @@
     "make_gaussian_filter",
     "ramp_projection",
     "smoothed_projection",
+    "symmetrize_diagonal",
+    "symmetrize_mirror",
+    "symmetrize_rotation",
     "tanh_projection",
 ]