Shape optimization

[Noe-LV]

I'm referring to the distinction between shape optimization and topological optimization that can be found for example here. As far as I understand, fdtdx is built to perform inverse design with topological optimization. Classes like fdtdx.Device let you optimize the permittivity values of a voxel grid with either discrete or continuous options. Classes to draw shapes like fdtdx.ExtrudedPolygon however inherit from static materials, operate with numpy and are not differentiable through JAX.

Is it planned in the near future to add shape optimization facilities?

Optimizing devices with splines, level sets, Bézier curves or free-form curves often produces devices that are more easily manufacturable. It also lets you enhance already existing designs. This could be suited for devices such as tapers, crossings, S-bends, directional couplers or ring resonators.

[ymahlau]

Hi @Noe-LV, this is definitely something we would like to support and it is one of the goals of the big re-write (see #349) to make this easier to specify.

For now, I believe there is a small hack you could already use to realize this: If you for example take a look at the sphere implementation, this could be made differentiable wrt. radius/center position by changing the following:

• define radius and position attributes as jax arrays, not as frozen_field floats
• analytically compute the fill fraction for the boundary pixels in get_voxel_mask_for_shape, i.e. don't return a fully binary mask, but for boundary pixels compute how much of these pixels is actually filled with the sphere.
• write a loss function that has these radius/position as first argument. Then use aset on the object container to update the attribute of the sphere with the value given as function argument before running a simulation.

I have not tested this myself, so I am not 100% confident that these are the only changes necessary to make this work.

In principle you could do the same thing for any object in your shape optimization application, just the specification of parameters -> mask needs to be differentiable. For example, with bezier curves this should be easy to do.

For level-sets specifically, you can use an fdtdx.Device with subpixel-smoothed-projection as a parameter mapping. If you then use a projection value of beta=jnp.inf, you are going to get a levelset that is differentiable wrt. the parameters. This is fully implemented and tested. You can take a look at the ceviche corner example here for how to set this up.

If you have any questions on this let me know! Also if you try out the approach above let me know! Either way I would be interested and if for some reason it should not work I am happy to help.

[ymahlau]

Actually on second thought: The approach above needs one more step: Currently apply_params does not re-calculate the material arrays for static-material-objects, only devices. So you would need to take the array initialization code from here and copy it into the apply_params function.

So I guess it is a bit complicated to realize currently. Alternatively you can always manually edit the material arrays to whatever values you want. If you have a differentiable pipeline for bezier curves, you could calculate the material arrays of your array and update the material arrays manually.

[Noe-LV]

Hi @ymahlau, many thanks for this detailed answer!

This approach sounds very interesting. I was planning to optimize taper shapes with Bézier curves, and I thought that adapting fdtdx.ExtrudedPolygon might be the way to go. I'll try to see if the extrusion pipeline that maps vertices -> mask can be made differentiable and try to modify apply_params first. It seems like the most robust approach.

Otherwise I'll try to go for the manual updates of the arrays but it seems trickier to me, I'm not sure how everything is implemented yet.

If you are interested, I'll keep you posted!
Thanks again.

[Noe-LV]

Hi @ymahlau !
Here's an update about this. It's a bit lengthy. In summary: I've been able to create successfully a shape optimization pipeline using the first method you propose. I can share the code or more details if you are interested.

• New DifferentiableExtrudedPolygon class

I've created a custom DifferentiableExtrudedPolygon class that mimics ExtrudedPolygon. It inherits from StaticMultiMaterialObject even though it is not 'static' anymore.

The main difference is that it that its vertices field is a jax array instead of a numpy array. It is also a regular fdtdx.core.jax.pytrees.field, instead of being a frozen_field. This lets the gradients flow through. This field can be updated after the computation of the gradient using e.g. updated_obj = objects.object_list[idx].aset("vertices", vertices), inside the loss function body.

• Differentiable point in polygon problem

The main difficulty of building this class was building the differentiable pipeline to compute the voxel mask using the vertices of the polygon. This problem translates into a point in polygon problem, that was solved in ExtrudedPolygon using the matplotlib.Path inbuilt contains_points method. This method internally uses the Ray casting algortithm, counting how many times a ray cast from a given point crosses the boundaries of the polygon. If the count is odd, the point lies on the outside of the polygon, whereas if it is even, the point lies on the inside. This kind of conditional checks makes the algorithm not differentiable.

• Winding number algorithm

Instead, I rewrote the winding number algorithm using jax functions. See the wikipedia page linked above for details about these two algorithms. The winding number algorithm is known to be slower but is differentiable. For a given point, it basically sums up the angles subtended by each side of the polygon. If the total is nonzero, the point lies inside of the polygon. This is differentiable because it boils down to operations like jnp.arctan2(x,y) that also are.

• Voxel mask with sub-pixel averaging

The Winding number gives the sign to a signed distance function, negative inside and positive outside. From this signed distance function (sdf), the voxel mask values are computed as fill_2d = 0.5 * (1.0 - jnp.tanh(sdf / (hw + 1e-30))), where hw is a half-width parameter that approximately corresponds to half the grid spacing. This mask value will be very close to 1 if sdf >> hw, close to 0 if sdf << hw, and it can have intermediate values inbetween.

• Apply params

This code was added to apply_params as suggested.

• Optimisation pipeline

A loss function can now be defined where one computes the fdtd forward run and its gradient, then updates the values of vertices with aset and the simulation objects with apply_params. It can be differentiated all the way using jax.grad. To make the pipeline JIT-compatible, I had to add a method that must run after the resolution of the grid. It freezes some of the fields of DifferentiableExtrudedPolygon, particularly the voxel mask shape so that they are not traced by jax.jit. This should be refactored is the code is to be kept.

[Noe-LV]

Hi @ymahlau,

I've encountered a problem in my optimization pipeline that I don't know how to solve.

• Context :

I'm building a waveguide escalator. The shapes are optimized by Bézier curves, the parameters of the optimization are their control points.
As explained above, I was able to adapt ExtrudedPolygon so that its vertices are now differentiable.

I'm measuring the overlap magnitude of the output field with the fundamental mode of the output waveguide as a figure of merit. I compute the gradient of this overlap magnitude with respect to the control points and run an optimize with it.

Overall the optimization pipeline is closely similar to that of the ceviche corner challenge example.

• The problem:

Putting PML boundaries in the simulation region quickly introduces NaN values in the gradient. These NaN values however do not appear when no boundaries are used. Have you encountered this type of issue before? Or do you know what might cause it?

I thought this could be linked to these lines of place_on_grid from perfect_matched_layer.py:

b_E = jnp.expm1(-dt / eps0 * (sigma_E / kappa_E + alpha_E)) + 1
a_E = jnp.nan_to_num((b_E - 1.0) * sigma_E / (sigma_E + alpha_E * kappa_E) / kappa_E, nan=0.0)
b_H = jnp.expm1(-dt / eps0 * (sigma_H / kappa_H + alpha_H)) + 1
a_H = jnp.nan_to_num((b_H - 1.0) * sigma_H / (sigma_H + alpha_H * kappa_H) / kappa_H, nan=0.0)

But I'm not sure, I guess it could also happen because of overflows due to float8 compression in the gradient values...

Thanks for your help!

[ymahlau]

Hi @Noe-LV, the way fdtdx computes gradients is by running a "time-reversed" simulation. For that purpose, the whole simulation is run backwards in time. Since PMLs are not invertible at time, the field values at the interface are stored and replayed during the backward pass (The recording module in codebase).

These stored field arrays can be quite large such that we implemented some simple compression (store in float8/float16, subsample in time and linearly interpolate to reconstruct). However, this does introduce some error, which can make the time reversed simualtion unstable if it is too large.

So I would advise you to run a backwards simulation without gradients ('fdtdx.full_run_backwards' and possibly even create a video of that backwards simulation (set 'inverse=true' in the detector)). If the backward simulation looks like the forwards simulation in reverse then the gradient should also be good. I would guess that your backward simulation becomes unstable.