Add tests for gradient check of affine transformations #1

rahul-flex · 2025-01-09T23:11:53Z

Removed @cached_property from [transformations](https://github.com/rahul-flex/tidy3d/blob/rahul-flex/diff-transformations/tidy3d/components/transformation.py#L107-L131) file so that the autograd can dynamically update values. This fixed the gradient issue. However, created errors:
1. Running pytest test_autograd.py caused errors:
  This created issue as we don't want @cached_property, then we must explicitly call matrix() when using it. Same goes for self.isidentity(). Made the fixes in [/geometry/base](https://github.com/rahul-flex/tidy3d/blob/rahul-flex/diff-transformations/tidy3d/components/geometry/base.py#L2795C1-L2795C78) file.
2. Now after running all the tests, got 9 errors:
  1. Errors image attached:

3. Fix:
    1. Converted matrix call to explicit method/func matrix() call. Same for isidentity method.

yaugenst-flex

Getting there! I noticed some weird formatting in these files, did you install the pre-commit hook and did you run the autoformatter (ruff) on your code before committing? There is some weird indentation, I'm actually surprised it runs at all.

Regarding comments and docstrings: This is already pretty good! Note though that the lowest level of indentation should be aligned with the quotes, i.e. instead of

"""
  docstring
"""

it should rather look like

"""
docstring
"""

Also for tests, it is generally enough if there is only a short docstring that explains what the test does, a full docstring like in the official docs with all parameters and types is generally not necessary (not bad either though).

yaugenst-flex · 2025-01-10T10:15:02Z

tests/test_components/test_autograd.py

@@ -21,6 +21,7 @@
 from tidy3d.components.autograd.utils import is_tidy_box
 from tidy3d.components.data.data_array import DataArray
 from tidy3d.plugins.polyslab import ComplexPolySlab
+import tidy3d.components.geometry.polyslab as polyslab


No need for this import, polyslab is available as td.PolySlab

yaugenst-flex · 2025-01-10T10:19:04Z

tests/test_components/test_autograd.py

+@pytest.fixture
+def create_polyslab()-> polyslab.PolySlab:
+  """
+    Creates a PolySlab instance for testing affine transformations.
+
+    Returns
+    -------
+    polyslab.PolySlab
+        An instance of PolySlab with predefined vertices, axis, and slab bounds.
+  """
+  vertices = np.array([[0.0, 0.0], [1.0, 0.0], [1.0, 1.0]])
+  axis = 1  # 0 for x, 1 for y, 2 for z
+  slab_bounds = (-1.0, 1.0)
+
+  return polyslab.PolySlab(vertices=vertices, axis=axis, slab_bounds=slab_bounds)


Good use of a fixture!

Avoid the use of verbs like create_ in fixture names though, since they will be used implicitly as arguments. So just polyslab would be good here.

yaugenst-flex · 2025-01-10T10:21:26Z

tests/test_components/test_autograd.py

+    create_polyslab : polyslab.PolySlab
+        The PolySlab object created using pytest fixture.
+  """
+  poly = create_polyslab


Related to fixture name, this is a bit awkward.

yaugenst-flex · 2025-01-10T10:22:02Z

tests/test_components/test_autograd.py

+        The PolySlab object created using pytest fixture.
+  """
+  poly = create_polyslab
+  x, y, z = 0.1, 0.2, 0.3


Ideally this test would be parametrized with a few different values for x, y, z, specifically to check edge cases, see here for reference.

yaugenst-flex · 2025-01-10T10:23:45Z

tests/test_components/test_autograd.py

+      new_poly = poly.translated(x, y, z)
+      return new_poly.vertices
+
+  for name in ["x", "y", "z"]:


There is no need for the for loop here. Just explicitly check each case.

yaugenst-flex · 2025-01-10T10:28:57Z

tests/test_components/test_medium.py

@@ -464,7 +464,7 @@ def test_rotation():
    axis = np.random.random(3)
    rot = td.RotationAroundAxis(axis=tuple(axis), angle=1.23)

-    R = rot.matrix
+    R = rot.matrix()


Yeah this is not ok, you are changing the behavior of tidy3d classes globally, i.e. you are changing the interface. Do not modify existing tests to make them pass. Fix your code instead so that they pass ;)
The tests are there for a reason.

yaugenst-flex · 2025-01-10T10:29:18Z

tidy3d/components/geometry/base.py

@@ -2792,7 +2792,7 @@ def rotation(angle: float, axis: Union[Axis, Coordinate]) -> MatrixReal4x4:
            Transform matrix with shape (4, 4).
        """
        transform = np.eye(4)
-        transform[:3, :3] = RotationAroundAxis(angle=angle, axis=axis).matrix
+        transform[:3, :3] = RotationAroundAxis(angle=angle, axis=axis).matrix()


yaugenst-flex · 2025-01-10T10:29:53Z

tidy3d/components/transformation.py

    @abstractmethod
    def matrix(self) -> TensorReal:
        """Rotation matrix."""

-    @cached_property
    @abstractmethod
    def isidentity(self) -> bool:


Even if you remove the caching, these still need to be properties.

yaugenst-flex · 2025-01-10T10:30:19Z

tidy3d/components/transformation.py

    def isidentity(self) -> bool:
        """Check whether rotation is identity."""

        return np.isclose(self.angle % (2 * np.pi), 0)
-
-    @cached_property
+


Same here, need to stay properties.

yaugenst-flex · 2025-01-10T10:32:40Z

tests/test_components/test_autograd.py

This is starting to look pretty good :)
One thing regarding test naming: Try to make the names of your tests mirror what they are testing, this is generally an unspoken convention. So instead of test_translation_grad, name it test_translated_grad since you are testing PolySlab.translated. In the same vein, we are actually only testing PolySlab here. What if we start testing Box and Cylinder too? It'll be cleaner to make the test name more specific then, e.g. test_polyslab_translated_grad, etc.

…ansformations

yaugenst-flex

This is looking great! Only had some minor comments.

yaugenst-flex · 2025-01-14T08:48:53Z

tests/test_components/test_autograd.py

+    [
+        (0.0, 0.0, 0.0, True),  # No scaling
+        (0.1, 0.2, 0.3, False),  # Small positive values
+        (-0.1, -0.2, -0.3, True),  # Small negative values


After discussing this here actually this should be equivalent to a reflection. You can see for example in this test, negative scaling is tested and expected to work, so this should also work under autograd.

yaugenst-flex · 2025-01-14T08:50:52Z

tests/test_components/test_autograd.py

+        (0.0, 0, True),  # No rotation around x-axis
+        (np.pi / 6, 1, False),  # Small rotation around y-axis
+        (-np.pi / 4, 2, True),  # Rotation around z-axis
+        (np.pi / 4, 1, False),  # 90-degree rotation around y-axis
+        (np.pi, 1, False),  # 180-degree rotation around y-axis


I think this really only comes down to whether or not we are rotating around the polyslab's axis, so might be better to just check whether we are rotating around that or not?

yaugenst-flex · 2025-01-14T08:52:13Z

tests/test_components/test_autograd.py

+        if not hasattr(new_poly, "vertices"):
+            raise ValueError("Rotation resulted in a non-differentiable Transformed object.")


I think this is not necessary, since check_grads will raise an exception if differentiating through Transformed, so no need to raise one here.

tidy3d/components/geometry/polyslab.py

yaugenst-flex · 2025-01-14T08:55:40Z

tidy3d/components/transformation.py

+        K = np.array([[0, -n[2], n[1]], [n[2], 0, -n[0]], [-n[1], n[0], 0]])
+        R = np.eye(3) + s * K + (1 - c) * K @ K


Nice! This is much cleaner than before.

yaugenst-flex

Ok overall this is looking great! I think we can try to move this into the main repo.

rahulkpadhy added 3 commits January 8, 2025 18:53

added test for transformations and fixed grad for rotation

4dcaa7b

pyproject.toml and poetry.lock reverted back

c7261a4

add pytest files for grad check of affine transformations

5718eea

rahul-flex requested a review from yaugenst-flex January 9, 2025 23:11

yaugenst-flex requested changes Jan 10, 2025

View reviewed changes

use @cached_property and add pytest files for grad check of affine tr…

61a1318

…ansformations

yaugenst-flex reviewed Jan 14, 2025

View reviewed changes

add test for reflection as negative scaled

bf5d174

yaugenst-flex self-requested a review January 15, 2025 12:25

yaugenst-flex approved these changes Jan 15, 2025

View reviewed changes

rahul-flex merged commit e16845a into pre-2.8-branch Jan 15, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add tests for gradient check of affine transformations #1

Add tests for gradient check of affine transformations #1

rahul-flex commented Jan 9, 2025

yaugenst-flex left a comment

yaugenst-flex Jan 10, 2025

yaugenst-flex Jan 10, 2025

yaugenst-flex Jan 10, 2025

yaugenst-flex Jan 10, 2025

yaugenst-flex Jan 10, 2025

yaugenst-flex Jan 10, 2025

yaugenst-flex Jan 10, 2025

yaugenst-flex Jan 10, 2025

yaugenst-flex Jan 10, 2025

yaugenst-flex Jan 10, 2025

yaugenst-flex Jan 10, 2025

yaugenst-flex left a comment

yaugenst-flex Jan 14, 2025

yaugenst-flex Jan 14, 2025

yaugenst-flex Jan 14, 2025

yaugenst-flex Jan 14, 2025

yaugenst-flex left a comment

		if not hasattr(new_poly, "vertices"):
		raise ValueError("Rotation resulted in a non-differentiable Transformed object.")

		K = np.array([[0, -n[2], n[1]], [n[2], 0, -n[0]], [-n[1], n[0], 0]])
		R = np.eye(3) + s * K + (1 - c) * K @ K

Add tests for gradient check of affine transformations #1

Add tests for gradient check of affine transformations #1

Conversation

rahul-flex commented Jan 9, 2025

yaugenst-flex left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

yaugenst-flex left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

yaugenst-flex left a comment

Choose a reason for hiding this comment