feat: autograd support for Clip Operations #2450
Open
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
There was a problem hiding this comment.
4 files reviewed, 5 comments
Edit PR Review Bot Settings | Greptile
Comment on lines
289
to
+321
| geometry = self.geometry | ||
| box = geometry.bounding_box | ||
|
|
||
| # we dont want these fields getting traced by autograd, otherwise it messes stuff up | ||
| # Construct bounding box (with padding if ClipOperation) | ||
| if isinstance(geometry, td.ClipOperation): | ||
| lo_a, hi_a = self._poly_bbox(geometry.geometry_a) | ||
| lo_b, hi_b = self._poly_bbox(geometry.geometry_b) | ||
|
|
||
| size = [get_static(x) for x in box.size] | ||
| center = [get_static(x) for x in box.center] | ||
| lo = np.minimum(lo_a, lo_b) | ||
| hi = np.maximum(hi_a, hi_b) | ||
|
|
||
| # polyslab only needs fields at the midpoint along axis | ||
| if ( | ||
| isinstance(geometry, PolySlab) | ||
| and not isinstance(self.medium, AbstractCustomMedium) | ||
| and field_keys == [("vertices",)] | ||
| ): | ||
| size[geometry.axis] = 0 | ||
| lengths = hi - lo # [Lx, Ly, Lz] | ||
| pad_vec = 0.05 * lengths # 5 % on every side | ||
|
|
||
| lo -= pad_vec | ||
| hi += pad_vec | ||
|
|
||
| size = (hi - lo).tolist() | ||
| center = ((hi + lo) * 0.5).tolist() | ||
|
|
||
| else: | ||
| # --- original behaviour -------------------------------------- | ||
| box = geometry.bounding_box | ||
| # we dont want these fields getting traced by autograd, otherwise it messes stuff up | ||
| size = [get_static(v) for v in box.size] | ||
| center = [get_static(v) for v in box.center] | ||
| # polyslab only needs fields at the midpoint along axis | ||
| if ( | ||
| isinstance(geometry, PolySlab) | ||
| and not isinstance(self.medium, AbstractCustomMedium) | ||
| and field_keys == [("vertices",)] | ||
| ): | ||
| size[geometry.axis] = 0.0 | ||
|
|
There was a problem hiding this comment.
style: Complex boolean conditions and branching paths that modify bounding box calculation. Consider refactoring into separate helper methods
Comment on lines
+280
to
+295
| if SAVE_RESULTS: | ||
| RESULTS_DIR.mkdir(parents=True, exist_ok=True) | ||
| np.save( | ||
| RESULTS_DIR / f"fd_ad_{scenario['name'].replace(' ', '_')}" | ||
| f"_{operation}_{param_label}.npy", | ||
| dict( | ||
| scenario_name=scenario["name"], | ||
| operation=operation, | ||
| param_label=param_label, | ||
| fd_val=float(fd_val), | ||
| ad_val=float(ad_val), | ||
| p_plus=float(p_plus), | ||
| p_minus=float(p_minus), | ||
| rel_diff=float(rel_diff), | ||
| ), | ||
| ) |
There was a problem hiding this comment.
style: convert this block into a separate helper function for better code organization
Comment on lines
+2965
to
+2967
| def softmin(x, axis=1, lam=30.0): | ||
| return -1.0 / lam * logsumexp(-lam * x, axis=axis) | ||
|
|
There was a problem hiding this comment.
style: Consider adding docstring to describe parameters lam and alpha, since these significantly affect the smoothness of the approximation
Comment on lines
+3017
to
+3023
| def polygon_soft_membership_2d( | ||
| verts: np.ndarray, | ||
| points: np.ndarray, | ||
| alpha: float = 20.0, | ||
| beta: float = 5.0, | ||
| ) -> np.ndarray: | ||
| """Vectorized "soft membership" for polygon 'verts' on 'points'. |
There was a problem hiding this comment.
style: Add input validation to ensure alpha > 0 and beta > 0 since these control sharpness/smoothness of approximation
Comment on lines
+3406
to
+3408
| mem_fix_a[d] = _membership_3d(verts_a, bounds_a, axis_a, d) | ||
| mem_fix_b[d] = _membership_3d(verts_b, bounds_b, axis_b, d) | ||
| der_var_a[d], der_var_b[d] = _boolean_coeff(mem_fix_a[d], mem_fix_b[d], self.operation) |
There was a problem hiding this comment.
style: Extract magic number coefficients like alpha=20.0, beta=5.0 into named class constants with clear descriptions
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Autograd support for ClipOperation with respect to the vertices and slab_bounds of its constituent PolySlab.
Key Changes:
Differentiable Geometry Representation:
Introduces a "soft membership" function. It provides a smooth, differentiable approximation of the sharp boundaries and non-differentiable nature of boolean geometric operations (intersection, union, difference and symmetric-difference), enabling gradient computation.
Adjoint-Based Derivative Calculation:
Computes gradients using the adjoint method, integrating permittivity sensitivity , and the gradient of the soft membership function wrt polylab params.
Expanded Monitor Bounds:
Adjoint monitors now encompass both input geometries with padding, ensuring accurate field capture for derivative calculations near boundaries.
Validation:
Results verified against finite-difference checks for all parameters and boolean operation types (attached).
Greptile Summary
Implements automatic differentiation support for ClipOperation geometries, enabling gradient calculations for boolean operations between PolySlabs through differentiable approximations.
tidy3d/components/geometry/base.pyto create smooth, differentiable approximations of boolean operations (union, intersection, difference)tidy3d/components/structure.pywith 5% padding to ensure accurate field capture near boundariestests/test_components/test_autograd_clip_operations.pyvalidating derivatives against finite-difference checkstest_error_clipfromtests/test_components/test_autograd.pyas ClipOperations now support differentiation