feat: autograd support for rotated Box #2362
Conversation
rahul-flex
force-pushed
the
rahul-flex/rotated-box-derivatives
branch
from
6f653d1 to
3b1a6be
Compare
There was a problem hiding this comment.
Thanks @rahul-flex! Left some comments. I think the biggest thing we should think about is if maybe now would be a good time to consider reusing the polyslab routines for Box too, since #2418 introduces more robust polyslab handling. That code handles a lot of edge cases that we currently don't handle for Box, such as intersection with the simulation bounds support for 2d simulations, and adaptive grid sampling.
| if rotation_matrix is not None: | ||
| rotation_matrix = rotation_matrix | ||
| else: | ||
| rotation_matrix = None |
There was a problem hiding this comment.
This looks strange? Lines 2660 & 2661 don't appear to do anything?
There was a problem hiding this comment.
Thanks for catching that. It does not do anything. I will remove it.
| ) | ||
| derivative_map[("geometry", "center")] = transformed_center_gradient | ||
| derivative_map[("geometry", "size")] = transformed_size_gradient | ||
| derivative_map[("transform",)] = np.zeros((4, 4)) |
There was a problem hiding this comment.
So this returns a zero array even though it's not really implemented right? Is this field needed for book-keeping and used somewhere else?
There was a problem hiding this comment.
Yes, it not doing anything when transform parameters are not differentiable. I will remove it.
| @@ -61,6 +63,7 @@ | |||
| ) | |||
|
|
|||
| POLY_GRID_SIZE = 1e-12 | |||
| _NUM_PTS_DIM_BOX_FACE = 200 | |||
There was a problem hiding this comment.
Seems a bit strange to let all box faces always have the same no. of points..
There was a problem hiding this comment.
I followed along the older polyslab surface integration, which had a constant number of points for all faces. Should we change it?
| perps1=perps1, | ||
| perps2=perps2, | ||
| ) | ||
| return surface_mesh, n_local |
There was a problem hiding this comment.
What's n_local used for? It seems to be discarded later.
There was a problem hiding this comment.
Yes, currently it is not being used. In the next PR: derivative wrt rotation angle, we need n_local. That uses the same function build_box_face_mesh(), that's why kept the signature same.
Autograd support for the
Transformedclass when the base geometry is aBoxunder rotation. Key changes are:Rotated Normals:
Box face derivatives now accept rotation matrix, ensuring that the face derivative computed use rotated normals.
Derivative Computation:
The derivative with respect to rotated faces is computed using the surface mesh and gradient surfaces.
Testing:
Finite difference checks to validate autograd computations.
Emulated simulation test cases to verify rotation.
Test results attached for 7 cases: box, box with background structure, 0 deg rotation, 90 deg rotation about z, 45 deg about y, 45 deg about x, and 45 deg about z.