Obtaining the phase from field monitor for inverse design problem

[Nikolaos-Matthaiakakis]

Hello, I am a little bit confused about how i can extract some data from a field monitor in a way that is differentiable for autograd for an inverse design issue.

In more detail I face the following problem. I use

# Loss function
def loss_fn(params: anp.ndarray, beta: float) -> tuple[float, dict]:
    """Loss function for the design, the difference in intensity + the feature size penalty."""

    freqs_inv  = np.linspace(3, 6, 3) * THz  # frequency range of interest 
    # construct and run the simulation
    sim_inv = make_sim(params, beta=beta,freqs=freqs_inv)
    sim_data_inv = web.run(sim_inv, task_name="loss_calc", verbose=False)

    # calculate the error
    error=error_fn(sim_data_inv )

    # grab the respective and total losses
    return error, sim_data_inv.to_static()

# construct a function of `params` and `beta` that returns the loss value, gradient, and the aux_data
loss_fn_val_grad = value_and_grad(loss_fn, has_aux=True)

# call this on our initial parameters
(val, grad), sim_data = loss_fn_val_grad(params0, beta=beta0)

which calls an error function that tries to obtain the phase of the fields from a monitor and compare them to a desired broadband phase profile. The issue is I am not sure how I can obtain this data in a way suitable for autograd as

def error_fn(sim_data):
    # Extract raw arrays from xarray
    Ey_Tr = sim_data["Tf"].Ey.data
    Ey_Tr_ref = ref_results["Tf"].Ey.data
    
    # Average over x and y dimensions
    Ey_Tr_avg = anp.mean(Ey_Tr[:, :, 0, :], axis=(0, 1))
    Ey_Tr_avg_ref = anp.mean(Ey_Tr_ref[:, :, 0, :], axis=(0, 1))

    # Phase in radians
    phase_Ey_Tr = anp.angle(Ey_Tr_avg)
    phase_Ey_ref = anp.angle(Ey_Tr_avg_ref)

does not seem to properly provide the data in the correct form for this problem.

[momchil-flex]

Could you clarify the problem: does not seem to properly provide the data in the correct form for this problem. do you mean that the data is not what you expect, that the gradient does not seem right, or that the gradient computation errors?

[Nikolaos-Matthaiakakis]

Hello, I apologize for the very unclear first comment. At first I thought I was retrieving the data in a wrong way that caused a variety of errors with autograd. In reality, I was making an error on how I handled the data later in my code and the result ended up being non differentiable. I believe I have fixed this issue now but I have some other questions. For example for a broadband optimization (like for achieving a suitable phase profile for an achromatic metalens) does the adjoint solver run an adjoint simulation for each of the wavelengths captured by the monitor? Are the results then weighted together somehow? I am a little bit confused on how the adjoint method is implemented here for such simulations.

Thank you very much for your time.

[momchil-flex]

Good that you got it to work! Your question is pretty good and the details are actually complicated. In the worst case, we do need one adjoint simulation per monitor frequency. In some cases even with broadband, a single adjoint simulation can be used with some tricks. However I think (at least currently) if you are differentiating a FieldMonitor, it would be one adjoint sim per frequency.. so use as few as possible. @groberts-flex can confirm if this is true.

[Nikolaos-Matthaiakakis]

Hello and thank you for the reply. Since in this case there should be multiple adjoint simulations, one for each frequency, should my function also return individual goals for each frequency in the form of an array?

[momchil-flex]

No, in any case this is handled under the hood. Your objective function should be scalar-valued. But if it depends on multiple frequency outputs, this is detected automatically, and the adjoint simulations will be created as needed. It's just that in some cases (e.g. a single ModeMonitor in the objective function) we can even use a single adjoint simulation, but again, this is detected automatically.

[Nikolaos-Matthaiakakis]

I understand, thank you very much.

[groberts-flex]

Just wanted to confirm that the above is true! When using a FieldMonitor, there will be one adjoint simulation per frequency that is used in your objective function. You can have more frequencies in your monitor, but the ones that get used in your objective will each be given their own adjoint simulation. Depending on how you weight them together to achieve a scalar value in the end, this will determine the relative weights of all those adjoint simulations. As was pointed out, this all happens automatically for you. The main thing to take into account is to limit the number of frequencies that enter the objective when possible to reduce the number of simulations you need when using FieldMonitors.

[Nikolaos-Matthaiakakis]

Thank you for the verification, I saw that there is a limit of 10 frequencies by default. Currently I do not need to use that many, but if I needed would there be a way to overcome this limit?

[groberts-flex]

Yes, you can use the argument max_num_adjoint_per_fwd to your run or run_async call to overcome the limit.

[Nikolaos-Matthaiakakis]

Hello, sorry to revive this thread but I am having a related issue, probably due to my lack of experience with autograd. I have tried several types of loss functions to achieve a target phase (mse, cosine loss, etc) but I seem to be facing issues such as, 0 gradients, value getting stuck at the worst value, value getting stuck in plateaus, etc.

Phase is wrapped so this can cause issues and I probably haven't found an ideal way to overcome this, unwrapping seems to cause issues also. I was wondering if you have could offer some advice on how to handle loss functions that have phase instead of intensity as a target if you have any relative experience. Thank you very much in advance.

[Nikolaos-Matthaiakakis]

Ok I will prepare something a with a single frequency phase target and communicate once it is ready to share. Thank you

[groberts-flex]

Thank you Nikolaos! Happy to help take a look as well if you can send to both Momchil and I!

[Nikolaos-Matthaiakakis]

Thank you, I haven't found the time to prepare it yet, I will try by Monday, but I have another idea that might help reduce some of the issues faced by the optimizer in my setup. Sometimes the fields in the simulation domain don't reduce to the cut-off level and this could affect the optimization process (my structure has very high epsilon value and is lossless and could be the reason this happens).

I was wondering how the following could be modified to include some losses similar to how it is done here

medium = td.Medium(permittivity=1, conductivity=0.005)

def get_eps(params: anp.ndarray, beta: float) -> anp.ndarray:
    """Get the permittivity values (1, permittivity) array as a function of the parameters (0, 1)"""
    density = filter_project(params, beta)
    eps = rescale(density, 1, perm_slab)
    return eps.reshape((nx, ny, 1, 1))


def make_slab(params: anp.ndarray, beta: float) -> td.Structure:
    """make the phase mask as a function of the parameters for a given `beta` value."""

    # construct the coordinates
    x0_max = +w1 / 2 - dl_design_region / 2
    y0_max = +w2 / 2 - dl_design_region / 2
    coords_x = np.linspace(-x0_max, x0_max, nx).tolist()
    coords_y = np.linspace(-y0_max, y0_max, ny).tolist()
    coords = dict(x=coords_x, y=coords_y, z=[z_slab], f=[freq0])

    # construct the data array for the permittivity
    eps_values = get_eps(params, beta)
    eps_data_array = td.ScalarFieldDataArray(eps_values, coords=coords)

    # construct the permittivity dataset
    field_components = {f"eps_{dim}{dim}": eps_data_array for dim in "xyz"}
    eps_dataset = td.PermittivityDataset(**field_components)

    # construct the phase mask slab
    custom_medium = td.CustomMedium(eps_dataset=eps_dataset)
    box = td.Box(center=(0, 0, z_slab), size=(td.inf, td.inf, h))
    return td.Structure(geometry=box, medium=custom_medium)

I imagine I simply should add a conductivity value as here next to eps_dataset?

dielectric = CustomMedium(permittivity=permittivity, conductivity=conductivity)

Edit: this seems to give an error

ValidationError: 1 validation error for CustomMedium
root
Please either define 'permittivity' and 'conductivity', or 'eps_dataset', but not both simultaneously. (type=value_error.setup)

so it should be done through

eps_dataset (Optional[PermittivityDataset] = None) – [To be deprecated] User-supplied dataset containing complex-valued permittivity as a function of space. Permittivity distribution over the Yee-grid will be interpolated based on interp_method.

[Nikolaos-Matthaiakakis]

For now I overcome this by including some loss in the surrounding medium. In the future I would also like to modify the funciton to work with materials with complex and dispersive permittivity (is this possible?). I will try to provide the code for the phase target soon.

[Nikolaos-Matthaiakakis]

Regarding the phase, a test code is something like this

def error_fn(sim_data):
    Ey_Tr = sim_data["Tf"].Ey.data
    Ey_Tr_ref = ref_results["Tf"].Ey.data

    Ey_avg = anp.mean(Ey_Tr[:, :, 0, :], axis=(0, 1))
    Ey_avg_ref = anp.mean(Ey_Tr_ref[:, :, 0, :], axis=(0, 1))

    # single-angle relative phase
    phase_rel = anp.angle(Ey_avg * anp.conj(Ey_avg_ref))

    phase_trgt = np.pi / 2
    phase_err = phase_rel - phase_trgt

    # cosine-based loss
    loss = anp.mean(1.0 - anp.cos(phase_err))
    return loss


## Objective function

# Loss function
def loss_fn(params: anp.ndarray, beta: float) -> tuple[float, dict]:
    """Loss function for the design, the difference in intensity + the feature size penalty."""
    
    # construct and run the simulation
    sim_inv = make_sim(params, beta=beta)
    sim_data_inv = web.run(sim_inv, task_name="loss_calc", verbose=False)

    # calculate the error
    error=error_fn(sim_data_inv )

    # grab the respective and total losses
    return error, sim_data_inv.to_static()

# construct a function of `params` and `beta` that returns the loss value, gradient, and the aux_data
loss_fn_val_grad = value_and_grad(loss_fn, has_aux=True)

# call this on our initial parameters
(val, grad), sim_data = loss_fn_val_grad(params0, beta=beta0)

print(f"initial loss value = {val:.3f}")
print(f"gradient shape = {grad.shape:}")
print(f"norm of gradient = {anp.linalg.norm(grad):.3e}")

## Optimization loop
import optax

# hyperparameters
num_steps = 8
learning_rate = 0.75

# initialize adam optimizer with starting parameters
params = params0.copy()
optimizer = optax.adam(learning_rate=learning_rate)
opt_state = optimizer.init(params)

# store history
history = dict(loss=[], params=[], betas=[], penalty=[], intensity_diff=[], sim_data=[])

# gradually increase the binarization strength over iteration
beta_increment = 0.5
beta = beta0

for i in range(num_steps):
    print(f"step = ({i + 1} / {num_steps})")

    # compute gradient and current loss function value
    (loss, gradient), _ = loss_fn_val_grad(params, beta=beta)

    # save history
    history["loss"].append(loss)
    history["params"].append(params)
    history["betas"].append(beta)

    # log some output
    print(f"\tloss = {loss:.3e}")
    print(f"\tbeta = {beta:.2f}")
    print(f"\t|gradient| = {np.linalg.norm(gradient):.3e}")

    # compute and apply updates to the optimizer based on gradient (+1 sign to minimize loss_fn)
    updates, opt_state = optimizer.update(+gradient, opt_state, params)
    params[:] = optax.apply_updates(params, updates)

    # cap the parameters between their bounds
    np.clip(params, 0.0, 1.0, out=params)

    # update the beta value
    beta += beta_increment

and gives

for sending the full code please let me know how I could send it to both of you. Thank you very much for all the help