Fix gradient magnitude custom medium

[groberts-flex]

Found while debugging the autograd level set notebook that the gradient magnitudes when using the autograd version vs. the Jax version were off by a couple orders of magnitude (the direction was the same). CustomMedium is implemented differently than JaxCustomMedium for the case where the gradient is computed with the coordinate array for one of the dimensions fixed. In CustomMedium, the volume normalization was not taking into account the thickness of the dimension and summing so that each slice in the gradient array were all effectively weighted by 1. The change here is to use mean instead of sum and multiply by missing volume factors.

[yaugenst-flex]

Thanks for catching this @groberts-flex! I guess I'm a bit confused as to why two integration measures are necessary here (d_vol as well as the additional volume factor you introduced).
It seems to me that in principle the localized cell-volume weighting should be correct, but maybe some dimensions are not being properly integrated over? The additional_volume_factor assumes a uniform grid, which is not necessarily the case.
I think the overall approach should be:

1. Interpolate onto integration coordinates
2. Multiply each grid cell by the local cell volume
3. Sum over the domain

In fact, maybe we can get rid of most of the logic here by just using xarray's integrate? What do you think @tylerflex?

[notion-workspace]

Fix gradient magnitude error in CustomMedium

[groberts-flex]

Thanks @yaugenst-flex for the comments and suggestions on this! Currently, d_vol accounts for the size of the cells being integrated over only for coordinate arrays that are over length 1 in the permittivity data array fed into the function. In the case of the level set notebook, there z-dimension is of length 1, so the d_vol factors only account for the x- and y- contributions to the volume of the cell, which is why the gradient magnitude ends up coming out too large. Any of the other xyz dimensions are just summed over currently without regard to the underlying grid (this is where I replaced the sum with mean in the current PR) which is where the additional_volume_element came from that I introduced since the array was already being evenly summed over. You make a great point, though, and if there is a non-uniform grid that should be taken into account! In JaxCustomMedium, when the third dimension does not have an associated coordinate array, a uniform discretization is chosen based on the size of the dimension which is then interpolated over and summed over.

[groberts-flex]

updated the way this was done to use the underlying grid from the electric field data for the final dimension if it exists to interpolate and define the correct volume element for d_vol. The summing over the extra dimension now happens at the end after the proper weighing by each volume element is done.

[yaugenst-flex]

I think we can go with this fix if the result matches the JaxCustomMedium implementation.

I'm just wondering whether it wouldn't be simpler if we just removed all of the manual volume element calculations, something like:

for dim in "xyz":
    if dim not in coords_interp:
        E_der_dim_interp = E_der_dim_interp.integrate(dim)

Wouldn't this eliminate the need for any special handling?

[groberts-flex]

I think we can go with this fix if the result matches the JaxCustomMedium implementation.

I'm just wondering whether it wouldn't be simpler if we just removed all of the manual volume element calculations, something like:

for dim in "xyz":
    if dim not in coords_interp:
        E_der_dim_interp = E_der_dim_interp.integrate(dim)

Wouldn't this eliminate the need for any special handling?

This is a good point, but I still need the handling for a dimension of size 0 when I use integrate. Otherwise, that integrate ends up returning a zero array. So I would need to keep the size checking in place, but the nice aspect of this is it doesn't create as large of a d_vol array because we integrate first over the extra dimension.

In JaxCustomMedium, the approach is actually not to use the underlying electric field data to choose the integration coordinates on the extra dimension, but it instead chooses it's own uniform grid to interpolate on based on a certain number of grid points per wavelength that is hard-coded. I'm not sure which is more desirable, using the electric field data underlying grid in the extra dimension or fixing it to be a uniform grid. My guess is it probably is case dependent which one of these is more accurate.

[yaugenst-flex]

In JaxCustomMedium, the approach is actually not to use the underlying electric field data to choose the integration coordinates on the extra dimension, but it instead chooses it's own uniform grid to interpolate on based on a certain number of grid points per wavelength that is hard-coded. I'm not sure which is more desirable, using the electric field data underlying grid in the extra dimension or fixing it to be a uniform grid. My guess is it probably is case dependent which one of these is more accurate.

Yeah so for now I think doing it based on the electric field data is the right approach. We should however have a look at how the adjoint monitors are created from which that field data is sampled, and ideally the accuracy of this sampling could be a setting in a global configuration (that we'll eventually get to...) where you could set grad_sampling={auto, 1nm, 5nm} for example. Would say that's out of scope for this PR though.

This still needs a changelog entry, otherwise looks good to go :)

[groberts-flex]

Sounds good, that makes sense! Adding changelog entry, thanks for the reminder on that