PEC gradient support for Box and PolySlab geometries #2724
Conversation
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| "Computing slab face derivatives for flat structures is not fully supported and " | ||
| "may give zero for the derivative. Try using a structure with a small, but nonzero" | ||
| "thickness for slab bound derivatives." |
There was a problem hiding this comment.
syntax: Missing space in warning message before 'thickness'
| "Computing slab face derivatives for flat structures is not fully supported and " | |
| "may give zero for the derivative. Try using a structure with a small, but nonzero" | |
| "thickness for slab bound derivatives." | |
| "Computing slab face derivatives for flat structures is not fully supported and " | |
| "may give zero for the derivative. Try using a structure with a small, but nonzero " | |
| "thickness for slab bound derivatives." |
| fd_grad = np.squeeze(all_fd_grad_parameters[argmin_convergence_test]) | ||
| else: | ||
| fd_grad = np.squeeze(all_fd_grad_parameters) | ||
| valid_mask = np.ones(fd_grad, dtype=bool) |
There was a problem hiding this comment.
syntax: np.ones() expects shape as first argument, not the array itself
| valid_mask = np.ones(fd_grad, dtype=bool) | |
| valid_mask = np.ones(len(fd_grad), dtype=bool) |
| dim_um = 1.5 * mesh_wvl_um | ||
| dim_um = 1.5 * mesh_wvl_um |
There was a problem hiding this comment.
style: Variable 'dim_um' is assigned the same value twice
| dim_um = 1.5 * mesh_wvl_um | |
| dim_um = 1.5 * mesh_wvl_um | |
| dim_um = 1.5 * mesh_wvl_um |
Context Used: Rule - Extract repeated or complex expressions into well-named local variables to improve readability. (link)
| dim_um = 1.5 * mesh_wvl_um | ||
| dim_um = 1.5 * mesh_wvl_um |
There was a problem hiding this comment.
style: Variable 'dim_um' is assigned the same value twice
| dim_um = 1.5 * mesh_wvl_um | |
| dim_um = 1.5 * mesh_wvl_um | |
| dim_um = 1.5 * mesh_wvl_um |
Context Used: Rule - Extract repeated or complex expressions into well-named local variables to improve readability. (link)
| dim_um = 3 * mesh_wvl_um | ||
| dim_um = 3 * mesh_wvl_um |
There was a problem hiding this comment.
style: duplicate variable assignment - same value assigned to dim_um twice
| dim_um = 3 * mesh_wvl_um | |
| dim_um = 3 * mesh_wvl_um | |
| dim_um = 3 * mesh_wvl_um |
Context Used: Rule - Remove temporary debugging code (print() calls), commented-out code, and other workarounds before finalizing a pull request. (link)
| dim_um = 3 * mesh_wvl_um | ||
| dim_um = 3 * mesh_wvl_um |
There was a problem hiding this comment.
style: duplicate variable assignment - same value assigned to dim_um twice
| dim_um = 3 * mesh_wvl_um | |
| dim_um = 3 * mesh_wvl_um | |
| dim_um = 3 * mesh_wvl_um |
Context Used: Rule - Remove temporary debugging code (print() calls), commented-out code, and other workarounds before finalizing a pull request. (link)
| if make_H_der_map: | ||
| der_map_H = derivative_map_H(fld_fwd=fld_fwd, fld_adj=fld_adj) | ||
|
|
||
| return {"E" : der_map_E, "D": der_map_D, "H": der_map_H} |
There was a problem hiding this comment.
syntax: Extra space before colon in dictionary - should be "E": der_map_E for consistency
| return {"E" : der_map_E, "D": der_map_D, "H": der_map_H} | |
| return {"E": der_map_E, "D": der_map_D, "H": der_map_H} |
| def get_field_key(dim: str, fld_data: typing.Union[td.FieldData, td.PermittivityData]) -> str: | ||
| """Get the key corresponding to the scalar field along this dimension.""" | ||
| return f"E{dim}" if isinstance(fld_data, td.FieldData) else f"eps_{dim}{dim}" |
There was a problem hiding this comment.
style: This inner function get_field_key is identical to the one defined in multiply_field_data below and is now unused - should be removed
Context Used: Rule - Remove commented-out or obsolete code; rely on version control for history. (link)
| "Derivative of PEC material with less than 2 dimensions is unsupported. " | ||
| "Specified PEC box is {dimension}-dimesional" | ||
| ) |
There was a problem hiding this comment.
syntax: String formatting error: f-string missing variable interpolation
| "Derivative of PEC material with less than 2 dimensions is unsupported. " | |
| "Specified PEC box is {dimension}-dimesional" | |
| ) | |
| log.error( | |
| "Derivative of PEC material with less than 2 dimensions is unsupported. " | |
| f"Specified PEC box is {dimension}-dimensional" | |
| ) |
There was a problem hiding this comment.
Thanks @groberts-flex, huge effort! Mostly comments about general organization.
| @@ -66,6 +66,13 @@ class DerivativeInfo: | |||
| of the forward and adjoint displacement fields. The normal component of this | |||
| dataset is used when computing adjoint gradients for shifting boundaries.""" | |||
|
|
|||
| H_der_map: FieldData | |||
There was a problem hiding this comment.
All these H-related fields probably should be optional?
| simulation_bounds: Bound | ||
| """Geometry and simulation intersection bounds. | ||
| Bounds corresponding to the minimum intersection between the structure | ||
| and the simulation it is contained in.""" |
There was a problem hiding this comment.
Yes, it would be important to clarify the distinction between these two or remove one.
| @@ -138,6 +160,11 @@ class DerivativeInfo: | |||
| GeometryGroup handling where it is difficult to automatically evaluate | |||
| the inside and outside relative permittivity for each geometry.""" | |||
|
|
|||
| is_medium_pec: bool = False | |||
There was a problem hiding this comment.
Just a note, no need to change it in the context of this PR: We should consider making this behavior polymorphic, ie dispatching to different branches of the gradient calculation based on type rather than having explicit flags, especially if we introduce more cases where such separate treatment is necessary.
There was a problem hiding this comment.
yes, totally agree with this! these flags are already pretty specific because this doesn't catch the case of a dielectric structure overlaid on a PEC background which would need the same PEC integration but just flipped inside/outside.
There was a problem hiding this comment.
good point on simulation_bounds - forgot to update the comment for that one. There are a couple Jax things that use bounds_intersect (assuming at some point the Jax stuff will go away?). And then, it is used in polyslab to check is something is 2d, which could be computed via the simulation bounds and the slab bounds, but bounds_intersect kind of do that calculation already. I'll clarify the distinction in the comment and keep both for now, but happy to remove bounds_intersect and update if you think that's cleaner.
| ^ | ||
| | n (normal direction) | ||
| | | ||
| _.-~'`-._.-~'`-._ (PEC surface) |
There was a problem hiding this comment.
lol, I told AI to make me ascii art 😬 . but not sure this adds anything haha
There was a problem hiding this comment.
The evaluate_gradient_at_points has become really long, I think it needs some refactoring. It'd be good to extract the closures into class methods (or even module-level functions) and to split the PEC and dielectric paths into separate methods. So something like:
def evaluate_gradient_at_points(...):
# does some common pre- & post-processing, then calls one of:
def _evaluate_pec_gradient_at_points(...):
...
def _evaluate_dielectric_gradient_at_points(...):
...| # adjust coordinates by half a grid point outside boundary such that nearest interpolation | ||
| # point snaps to outside the boundary | ||
| # adjust_spatial_coords = np.squeeze([spatial_coords - normals * 0.5 * coords_dn]) | ||
| adjust_spatial_coords = np.squeeze([spatial_coords + normals * 0.5 * coords_dn]) |
There was a problem hiding this comment.
This looks somewhat unconventional. Is this equivalent to np.squeeze(x, axis=0)?
There was a problem hiding this comment.
cannot remember why I did this to be honest! I think this must have evolved out of a previous version of this line of the code - tested with just the squeeze and not putting it into a list first since I think that will just get squeezed out anyway and things seem to work. will verify everything at the end with one more run of the numerical tests but for now I removed the [] part and just kept as squeeze
|
|
||
| # adjust coordinates by half a grid point outside boundary such that nearest interpolation | ||
| # point snaps to outside the boundary | ||
| # adjust_spatial_coords = np.squeeze([spatial_coords - normals * 0.5 * coords_dn]) |
| E_fwd_at_coords = { | ||
| name: interp(E_fwd_coords_adjusted[name]["coords"]) | ||
| for name, interp in interpolators["E_fwd"].items() | ||
| } | ||
| E_adj_at_coords = { | ||
| name: interp(E_adj_coords_adjusted[name]["coords"]) | ||
| for name, interp in interpolators["E_adj"].items() | ||
| } | ||
|
|
||
| H_fwd_at_coords = { | ||
| name: interp(H_fwd_coords_adjusted[name]["coords"]) | ||
| for name, interp in interpolators["H_fwd"].items() | ||
| } | ||
| H_adj_at_coords = { | ||
| name: interp(H_adj_coords_adjusted[name]["coords"]) | ||
| for name, interp in interpolators["H_adj"].items() | ||
| } |
There was a problem hiding this comment.
There should probably be a helper method for this.
There was a problem hiding this comment.
Similar comment here: Box._derivative_face has become really complex with the addition of integrate_face. This function in particular seems to be doing a lot of things:
- Coordinate snapping
- Singularity correction
- Boundary detection
- Actual integration
Similar to the comment on evaluate_gradient_at_points, I think this should be split into functions with clearer responsibilities. With the current approach, it's also not really possible to test integrate_face.
This PR adds functionality for computing derivatives for Box and PolySlab geometries when the material in them is composed of PEC. This requires (1) using both electric and magnetic fields when computing the boundary gradients meaning that in cases of PEC, we add additional field components to the adjoint monitors and use them accordingly and (2) that we are careful about interpolation around the PEC boundaries to make sure not to interpolate with the fields inside of the material which are 0. There is also special handling for 2D structures that have edge singularities and numerical tests for evaluating the gradient accuracy. The result of running these unit tests is here: https://docs.google.com/presentation/d/1ddA1V_efdXi1CZ5Soal72crCDk856zUX3FFba6oa-zo/edit?usp=sharing
The gradient performance has also been tested in several antenna optimizations including ones where these geometries overlap each other.
Greptile Summary
This PR implements PEC (Perfect Electric Conductor) gradient support for Box and PolySlab geometries in the autograd system. The changes enable shape optimization of metallic structures like antennas by adding the ability to compute derivatives when these geometries contain PEC materials.
The implementation addresses fundamental electromagnetic theory requirements for PEC boundaries. Unlike dielectric materials that only need electric field components for gradient calculations, PEC materials require both electric and magnetic field components due to their unique boundary conditions where tangential electric fields are zero at surfaces. This necessitated significant changes across the autograd pipeline:
Enhanced field monitoring: The adjoint monitors now conditionally include magnetic field components (Hx, Hy, Hz) alongside electric fields when dealing with PEC structures.
Specialized interpolation: PEC regions have zero internal fields, so the implementation uses nearest-neighbor interpolation instead of linear interpolation to avoid sampling zero-valued fields inside PEC boundaries.
Dual gradient formulation: The derivative computation now has separate code paths for PEC and dielectric materials, with PEC calculations incorporating both electric and magnetic field contributions weighted by appropriate electromagnetic constants (MU_0/EPSILON_0).
Coordinate adjustment and singularity handling: Special logic handles 2D PEC structures that exhibit edge singularities, including coordinate snapping to ensure proper field sampling near boundaries.
Comprehensive testing framework: Two extensive test files validate the implementation through finite difference comparisons across various mesh refinements, wavelengths, and geometry configurations.
The changes maintain backward compatibility while extending the autograd capabilities to support PEC geometries, enabling antenna optimization and other applications involving metallic structures. The implementation correctly handles the theoretical distinction between dielectric and PEC boundary conditions in electromagnetic gradient calculations.
Important Files Changed
Click to expand file changes
tidy3d/components/geometry/base.pytidy3d/components/autograd/derivative_utils.pytidy3d/web/api/autograd/autograd.pytidy3d/components/structure.pytidy3d/web/api/autograd/utils.pytidy3d/components/geometry/polyslab.pytests/test_components/test_autograd_rf_box.pytests/test_components/test_autograd_rf_polyslab.pyConfidence score: 4/5
Sequence Diagram
sequenceDiagram participant User participant WebAPI as "tidy3d.web.run" participant AutogradRun as "_run_primitive" participant ForwardSim as "Forward Simulation" participant AdjointSim as "Adjoint Simulation" participant DerivativeInfo as "DerivativeInfo" participant Structure as "Structure" participant Geometry as "Box/PolySlab" User->>WebAPI: "run(simulation, task_name)" WebAPI->>WebAPI: "is_valid_for_autograd(simulation)" WebAPI->>AutogradRun: "_run(simulation, ...)" AutogradRun->>AutogradRun: "setup_run(simulation)" Note over AutogradRun: Extract traced fields from simulation AutogradRun->>AutogradRun: "_run_primitive(traced_fields, ...)" AutogradRun->>AutogradRun: "setup_fwd(sim_fields, sim_original)" AutogradRun->>ForwardSim: "run forward simulation with adjoint monitors" Note over ForwardSim: Includes E and H field monitors for PEC ForwardSim-->>AutogradRun: "sim_data_fwd with field data" Note over User: When gradient computation is triggered AutogradRun->>AutogradRun: "_run_bwd(data_fields_vjp, ...)" AutogradRun->>AutogradRun: "setup_adj(data_fields_vjp, ...)" AutogradRun->>AdjointSim: "run adjoint simulation" AdjointSim-->>AutogradRun: "sim_data_adj" AutogradRun->>AutogradRun: "postprocess_adj(sim_data_adj, ...)" AutogradRun->>DerivativeInfo: "create DerivativeInfo with E, D, H fields" Note over DerivativeInfo: For PEC: includes H_fwd, H_adj fields<br/>and sets is_medium_pec=True AutogradRun->>Structure: "_compute_derivatives(derivative_info)" Structure->>Geometry: "_compute_derivatives(derivative_info)" alt Box geometry Geometry->>Geometry: "_derivative_faces(derivative_info)" loop For each face Geometry->>Geometry: "_derivative_face(min_max_index, axis_normal)" alt PEC material Geometry->>Geometry: "integrate E normal field with singularity correction" Geometry->>Geometry: "integrate H tangential fields" Note over Geometry: Apply MU_0/EPSILON_0 weighting for H fields else Dielectric material Geometry->>Geometry: "integrate D normal and E tangential fields" end end else PolySlab geometry Geometry->>Geometry: "_compute_derivative_vertices(derivative_info)" Geometry->>Geometry: "edge_basis_vectors(edges)" loop For each edge segment Geometry->>DerivativeInfo: "evaluate_gradient_at_points(coords, normals, perps1, perps2)" alt PEC material DerivativeInfo->>DerivativeInfo: "_adjust_spatial_coords_pec(grid_centers)" Note over DerivativeInfo: Snap coordinates outside PEC surface<br/>for nearest interpolation DerivativeInfo->>DerivativeInfo: "compute E normal + H tangential contributions" DerivativeInfo->>DerivativeInfo: "apply singularity correction for 2D PEC" else Dielectric material DerivativeInfo->>DerivativeInfo: "compute D normal + E tangential contributions" end end end Geometry-->>Structure: "gradient values" Structure-->>AutogradRun: "vjp_traced_fields" AutogradRun-->>User: "gradients w.r.t. traced fields"Context used:
Rule - Remove temporary debugging code (print() calls), commented-out code, and other workarounds before finalizing a pull request. (link)
Rule - Use an underscore (_) for loop variables that are intentionally unused. (link)
Rule - Remove commented-out or obsolete code; rely on version control for history. (link)