Clarification on issues experience with Adjoint Simulations in Meep

[DurhamSmith]

This is more of a topic for the mailing list but #1875 seems to indicate that is down with no other alternative as of yet so I am posting this here.

I would like to recreate the Bayer filter of the Faraon group in Meep.

At the moment I am trying to familiarize myself with adjoint optimization in Meep by adapting the examples to create new devices. I am currently trying to create an a 2um x 2um 2D element that focuses incoming EM radiation in the visible spectrum onto 1um plane parallel to the y axis and located 1.5um the device (along the x direction), while minimizing EM radiation focused to the region that is the mirror image of this plane along x=0. You can see this device in the figure below.

I have tried two different figures of merit to accomplish this, with no luck. They are
npa.sum(npa.abs(top / source) ** 2) / (npa.abs(bottom / source) ** 2)
and
npa.sum(npa.abs(top / source) ** 2 - 0.75 * npa.abs(bottom / source) ** 2)

I have tried to debug the issue for a few days now and have noticed a few things.

In the Design of a Symmetric Broadband Splitter the value of fwidth passed to the source is fwidth=0.1282051282051282. Plotting this it looks like it corresponds to the full width at half max value (ignore y-label and key value, these are from the source monitor).

However when I use values of fcen=1/0.55 and fwidth=0.3 I get the following for the values measured at the source monitor:

Which doesn't give me the full width half max range between [0.4, 0.7] as I'd expect. Although the fields at each of the monitors looks good:

However the ratios of arm/source monitors are out of wack for the frequencies where the the source is very small (understandably so, since we are dividing by tiny numbers):

Keeping fcen=1/0.55 and fwidth=0.3 but limiting the wavelengths used in the objective function to [0.535, 0.575] I have better values for the initial monitor measurements and power splitting between the focus and antifocus region:

However when I run this simulation for 20 iterations I get the following plot of how my figure of merit improves seems to taper off as can be seen in this plot:

The monitor measurements and ratios are:

So it does seem I get a little bit of an improvement in focusing onto the top source monitor, it is not uniform across spectrum and hardly improves after the 8th iteration

I have not (yet) incorporated any erosion or dilation filters. I assume these, while good for incorporating manufacturing constraints would hinder device performance, rather than boost it, and I want to find the source of the flattening of the FOM curve before complicating the problem further.

I would appreciate any advice on how to fix this or how to better understand what is going on so that I can be better equipped to solve these problems by myself in the future. am rather new to computational EM and am trying to develop a better understanding of setting up and modeling devices as well as debugging any issues with the simulations.

It would also be great to know is there a way to save and load OptimizationProblems? I assume its not the same as loading and dumping simulation state as that dumps epsion but OptimizationProblems take DesignRegions which store information about the ranges of epsilon

Additionally any advice for modeling the air-polymer device in the paper I posted would be tremendously appreciated.

The code for the simulation is:

#!/usr/bin/env python
# coding: utf-8
import meep as mp
import meep.adjoint as mpa
import numpy as np
from autograd import numpy as npa
from autograd import tensor_jacobian_product
from enum import Enum
import matplotlib.pyplot as plt
import pickle as pkl
import nlopt
from matplotlib import pyplot as plt
from matplotlib.patches import Circle

# TODO create file if doesnt exist
filename = "11"
seed = 44
np.random.seed(seed)

###############################################################################
#              Setup variables that define the simulation domain              #
###############################################################################
design_region_size = mp.Vector3(2.0, 2.0)  # Lets optimze a 2um^2 region

# Importantly PMLs are added at both the top and bottom of the sim
pml_size = mp.Vector3(0.35, 0.35)  # Half a wavelength of red light

# Distance from PML to source
pml_src_offset = mp.Vector3(0.1)  # Distance from PML to source plane
# Distance from source to the 'left boundary' of the device/design region.
src_device_offset = mp.Vector3(0.2)
# Distance from 'right boundary' of the device/design region
# to the field monitors/objective plane/focal length.
# Importantly assume all fields are monitored in the same plane
# If we want to change this modify the 'center' of the EigenmodeCoefficent
# That are passed to the OptimizationProblem
device_objective_plane_offset = mp.Vector3(1.5)
# Distance from field monitors at the objective plane to the rightmost PML
objective_plane_pml_offset = mp.Vector3(0.3)
# Adds extra padding to device to EACH SIDE of PML i.e the actual padding is 2x this .
# If this padding has x compoments it will be added
# in the region between the PML and source
# and between the field monitors at the objective plane and the PML
padding = mp.Vector3(0, 0.01, 0)

# Calculate the simulation cell size
cell_size = (
    2 * pml_size
    + pml_src_offset
    + src_device_offset
    + design_region_size
    + device_objective_plane_offset
    + objective_plane_pml_offset
    + 2 * padding
)
###############################################################################
#                               Define Materials                              #
###############################################################################
# Define the materials we will use for our device

n_sio2 = 1.46
SiO2 = mp.Medium(index=n_sio2)
n_tio2 = 2.61
TiO2 = mp.Medium(index=n_tio2)

###############################################################################
#                             Setup Design Region                             #
###############################################################################
design_region_resolution = (
    50
    # 100  # The number of designable point per unit of a. 100 gives a 10nm 'design pixel'
)

design_region_pixels = design_region_resolution * design_region_size

design_variables = mp.MaterialGrid(
    design_region_pixels,
    SiO2,
    TiO2,
    grid_type="U_DEFAULT",  # I am not sure if this is the correct choice
)

design_region_center = mp.Vector3(
    x=(pml_size + pml_src_offset + src_device_offset + design_region_size / 2).x
)
# We actually only want to diplace the x-component since me assume the device is centered in the y plane


design_region = mpa.DesignRegion(
    design_variables,
    volume=mp.Volume(center=design_region_center, size=design_region_size),
)

# print(design_variables.grid_size)
# print(design_region.volume.size)

###############################################################################
#                               Setup the source                              #
###############################################################################

fcen = (
    1 / 0.55
)  # The center wave freq should correspond to 550 nm wavelength light (c=lamba*freq)
frequencies = 1 / np.linspace(0.525, 0.575, 50)  # 0.4, 0.7, 10
# We want 10 equally spaced frequencies corresponding of wavelenghs  400-700nm

# width = 0.3
# fwidth = (
#     width * fcen # this= (* (/ 1 .55) 0.3)= 0.54
# )  # TODO: idk if this is correct I need to read up more on the sources#
fwidth = 0.3

# We only need the x-component
# TODO: If you have a 3D structure and want to adjust
# the source location more dimension change this
source_center = mp.Vector3(x=(pml_size + pml_src_offset).x)

# TODO: Adjust this to configure the source size
# The source size is the cell size along the Y axis
# without the pml and padding in the Y direction
source_size = mp.Vector3(0, (cell_size - 2 * pml_size - 2 * padding).y, 0)

kpoint = mp.Vector3(1, 0, 0)
src = mp.GaussianSource(frequency=fcen, fwidth=fwidth)
source = [
    mp.EigenModeSource(
        src,
        eig_band=1,
        direction=mp.NO_DIRECTION,
        eig_kpoint=kpoint,
        size=source_size,
        center=source_center,
    )
]

# print(f"src {source.__sizeof__()}")


###############################################################################
#                             Setup the simulation                            #
###############################################################################

# TODO Change resolution if you want a finer grain simulation
# Resolution: Recommended at least 8px/wavelength in highest dielectric
# Wavelength in dielectric = lambda * n = 2.61 * 0.700 = 1.827
# The calc should be res=(desired_px_per_wavelength)/(freespace_lambda/n_dielectric)
# 10/(0.7/2.61) = 37.2857142857 so 40 seems reasonable
resolution = (
    100  # Decide this better the recomend 8px/shortest wavelength in highest dielectric
)

# Setup PML
pml_layers = [mp.PML(pml_size.x)]


geometry = [
    mp.Block(
        center=design_region.center, size=design_region.size, material=design_variables
    ),
]

sim = mp.Simulation(
    cell_size=cell_size,
    boundary_layers=pml_layers,
    geometry=geometry,
    sources=source,
    eps_averaging=False,
    resolution=resolution,
    geometry_center=mp.Vector3(x=cell_size.x / 2),
)


###############################################################################
#                             OptimizationProblem                             #
###############################################################################

########### Define the field monitors of the  OptimizationProblem #############

# TODO readup more on this and the internals of EigenmodeCoefficient
mode = 1

# TODO: Change this if you want to change how far from the source
# the source field monitor if from the source
sorce_to_monitor_offset = mp.Vector3(x=0.05)

source_monitor_center = mp.Vector3(
    (pml_size + pml_src_offset + sorce_to_monitor_offset).x
)

# TODO Change this if you dont want to have the source monitor the same size as the source
source_monitor_size = source_size
TE0 = mpa.EigenmodeCoefficient(
    sim, mp.Volume(center=source_monitor_center, size=source_monitor_size,), mode,
)

# Calculate the center
# TODO Change this if needed, you probably want to offset the
# y positions and keep the x since that is defined by the model of the device & sim domain that we are using
objective_plane_monitor_center_x = (
    pml_size
    + pml_src_offset
    + src_device_offset
    + design_region_size
    + device_objective_plane_offset
).x

objective_plane_monitor_center_y = design_region_size.y / 4

top_objective_plane_monitor_center = mp.Vector3(
    objective_plane_monitor_center_x, objective_plane_monitor_center_y
)

# Set the size of the objective montors
# TODO Change this if you want to change the size of the objective montors
top_objective_plane_monitor_size = mp.Vector3(y=(design_region_size.y - 0.1) / 2)

TE_top = mpa.EigenmodeCoefficient(
    sim,
    mp.Volume(
        center=top_objective_plane_monitor_center,
        size=top_objective_plane_monitor_size,
    ),
    mode,
)


# TODO Change this if you wnata to change the bottom monitor location
bottom_objective_plane_monitor_center = mp.Vector3(
    objective_plane_monitor_center_x,
    -objective_plane_monitor_center_y,  # Notice we subtract so we put below the top monitor
)

# Set the size of the objective montors
# TODO Change this if you want to change the size of the objective montors
bottom_objective_plane_monitor_size = mp.Vector3(y=(design_region_size.y - 0.1) / 2)

TE_bottom = mpa.EigenmodeCoefficient(
    sim,
    mp.Volume(
        center=bottom_objective_plane_monitor_center,
        size=bottom_objective_plane_monitor_size,
    ),
    mode,
)

ob_list = [TE0, TE_top, TE_bottom]

###############  OptimizationProblem's Objective Function ##################

# TODO change this if you want to change the objective function that will be used it optimization problem
def J(source, top, bottom):
    obj_fn = (npa.abs(top / source) ** 2) / (npa.abs(bottom / source) ** 2)
    # obj_fn = npa.abs(top / source) ** 2 - 0.75 * npa.abs(bottom / source) ** 2
    # return npa.mean(obj_fn)
    return npa.sum(obj_fn)


# print(f"design_region {design_region.MaterialGrid}")

opt = mpa.OptimizationProblem(
    simulation=sim,
    objective_functions=J,
    objective_arguments=ob_list,
    design_regions=[design_region],
    frequencies=frequencies,
    decay_by=1e-4,
)
print(opt.design_regions[0].design_parameters.weights)


###########################################################################
#                             Setup Optimizer                             #
###########################################################################
# TODO Change this function if you want to change the function that you optimize


evaluation_history = []
sensitivity = [0]
cur_iter = [0]


def f(x, grad):
    print("Current iteration: {}".format(cur_iter[0] + 1))
    f0, dJ_du = opt([x])
    if grad.size > 0:
        # grad[:] = np.squeeze(dJ_du)
        grad[:] = np.sum(dJ_du, axis=1)
    evaluation_history.append(np.real(f0))
    sensitivity[0] = dJ_du
    cur_iter[0] = cur_iter[0] + 1
    ax = plt.gca()
    opt.plot2D(
        False,
        ax=ax,
        plot_sources_flag=False,
        plot_monitors_flag=False,
        plot_boundaries_flag=False,
    )
    if mp.am_master():
        plt.savefig(f"./{filename}/new_splitter_{cur_iter[0]}.png", dpi=300)

    return np.real(f0)


###############################################################################
#                                Run Optimizer                                #
###############################################################################
algorithm = nlopt.LD_MMA
# Initial guess
n = int(design_region_pixels.x * design_region_pixels.y)

# x0 = np.random.rand(n,)
x0 = np.ones((n,)) * 0.5
x = x0

# lower and upper bounds
lb = np.zeros((n,))
ub = np.ones((n,))

# Show and plot device before running
opt.update_design([x0])
opt.plot2D(True)
if mp.am_master():
    plt.show()

opt.plot2D(True)
if mp.am_master():
    plt.savefig(f"./{filename}/new_splitter_initial", dpi=300)

# Plot the powers before running
f0, dJ_du = opt([x], need_gradient=False)
frequencies = opt.frequencies
source_coef, top_coef, bottom_coef = [m._eval for m in opt.objective_arguments]

top_profile = np.abs(top_coef / source_coef) ** 2
bottom_profile = np.abs(bottom_coef / source_coef) ** 2

plt.figure()
plt.plot(1 / frequencies, top_profile * 100, "-o", label="Top Arm")
plt.plot(1 / frequencies, bottom_profile * 100, "--o", label="Bottom Arm")
plt.legend()
plt.grid(True)
plt.xlabel("Wavelength")
plt.ylabel("Splitting Ratio (% of source power)")
if mp.am_master():
    plt.savefig(f"./{filename}/Power_Splitting_vs_wavelength_init")


plt.figure()
plt.plot(1 / frequencies, npa.abs(source_coef), "-o", label="Source")
plt.plot(1 / frequencies, npa.abs(bottom_coef), "--o", label="Bottom")
plt.plot(1 / frequencies, npa.abs(top_coef), "--o", label="Top")
plt.legend()
plt.grid(True)
plt.xlabel("Wavelength")
plt.ylabel("npa.abs of monitor")
if mp.am_master():
    plt.savefig(f"./{filename}/Monitor_Measurements_init")


plt.figure()
plt.plot(1 / frequencies, npa.abs(source_coef), "-o", label="Source")
plt.legend()
plt.grid(True)
plt.xlabel("Wavelength")
plt.ylabel("npa.abs of monitor")
if mp.am_master():
    plt.savefig(f"./{filename}/Source_Monitor_Measurement_init")

plt.figure()


maxeval = 20  # Maximum time we want the solver to run for
solver = nlopt.opt(algorithm, n)
solver.set_lower_bounds(lb)
solver.set_upper_bounds(ub)
solver.set_max_objective(f)
solver.set_maxeval(maxeval)
solver.set_xtol_rel(1e-4)
x[:] = solver.optimize(x)


###############################################################################
#                                 Plot Results                                #
###############################################################################
# plotting Final Device

opt.update_design([x])
opt.plot2D(True)
if mp.am_master():
    plt.savefig(f"./{filename}/new_splitter_final", dpi=300)


# Plotting Figure of Merit Improvement over time
plt.figure()
plt.plot(evaluation_history, "o-")
plt.grid(True)
plt.xlabel("Iteration")
plt.ylabel("FOM")
if mp.am_master():
    plt.savefig(f"./{filename}/FOM_Evolution.png")

# Plotting power in top vs bottom arm

f0, dJ_du = opt([x], need_gradient=False)
frequencies = opt.frequencies
source_coef, top_coef, bottom_coef = [m._eval for m in opt.objective_arguments]

top_profile = np.abs(top_coef / source_coef) ** 2
bottom_profile = np.abs(bottom_coef / source_coef) ** 2

plt.figure()
plt.plot(1 / frequencies, top_profile * 100, "-o", label="Top Arm")
plt.plot(1 / frequencies, bottom_profile * 100, "--o", label="Bottom Arm")
plt.legend()
plt.grid(True)
plt.xlabel("Wavelength")
plt.ylabel("Splitting Ratio (% of source power)")
if mp.am_master():
    plt.savefig(f"./{filename}/Power_Splitting_vs_wavelength")


plt.figure()
plt.plot(1 / frequencies, npa.abs(source_coef), "-o", label="Source")
plt.plot(1 / frequencies, npa.abs(bottom_coef), "--o", label="Bottom")
plt.plot(1 / frequencies, npa.abs(top_coef), "--o", label="Top")
plt.legend()
plt.grid(True)
plt.xlabel("Wavelength")
plt.ylabel("npa.abs of monitor")
if mp.am_master():
    plt.savefig(f"./{filename}/Monitor_Measurements")


plt.figure()
plt.plot(1 / frequencies, npa.abs(source_coef), "-o", label="Source")
plt.legend()
plt.grid(True)
plt.xlabel("Wavelength")
plt.ylabel("npa.abs of monitor")
if mp.am_master():
    plt.savefig(f"./{filename}/Source_Monitor_Measurement")
###############################################################################
#                               Printing Fields                              #
###############################################################################

sources = [
    mp.Source(
        mp.ContinuousSource(frequency=fcen),
        component=mp.Ez,
        center=source_center,
        size=source_size,
    )
]

sim.change_sources(sources)
sim.run(
    mp.at_beginning(mp.output_epsilon),
    mp.to_appended("ez", mp.at_every(0.6, mp.output_efield_z)),
    until=1,
)

print(opt.design_regions[0].design_parameters.weights)

[oskooi]

For this particular problem, the objective function seems to involve focusing the fields within a homogeneous medium (i.e., free space) rather than maximizing transmission into a guided mode (as in the tutorial example you referenced involving the power splitter). You should therefore be using the FourierFields objective function of the adjoint solver rather than the EigenModeCoefficients as the latter is better suited to computing quantities such as scattering parameters. An example that is more closely related to designing a Bayer filter is a broadband metalens which is demonstrated in https://nbviewer.org/github/NanoComp/meep/blob/master/python/examples/adjoint_optimization/Near2Far-Optimization-with-Epigraph-Formulation.ipynb. Note that if your structure is rotationally symmetric, you can set it up in cylindrical rather than Cartesian coordinates which therefore converts a 3d simulation into 2d (a potentially large savings in memory and time).

(Separately, we will soon be revamping the tutorials for the adjoint solver since currently many of them are out of date.)

[DurhamSmith]

Thanks @oskooi I appreciate the help! I will make the changes and see if I can get this working. I'll leave this issue open for now until I have tested them but will close if they work.

[smartalecH]

Here's a few suggestions I have:

1. As @oskooi mentioned, use FourierFields, rather than EigenmodeCoefficients. Optimize a single field component at a single spatial point for each frequency.
2. Since the frequencies you care about are so far apart, just define two sources. For example, something like

sources = [mp.EigenmodeSource(mp.GaussianSource(fcen=f,fwidth=0.2*f),<<other args here>>) for f in [1/0.5, 1/0.6]]

This is actually better because the mode profile is computed at each frequency of interest. If you try to do a single source, the spatial profile is only calculated at the center frequency, and you could see some errors from mode dispersion.

3. Your objective function needs to properly parse where each freq focuses (if you are truly trying to replicate the paper). And you just need to optimize the frequencies you truly care about (e.q. frequencies=[1/0.5, 1/0.6]).

(edit: fixed ContinuousSource to GaussianSource)

[DurhamSmith]

@smartalecH thanks for the help! I have made these adjustments to the simulation but is seems that by changing the sources to a ContinuousSource the simulation never completes 1 forward run of the adjoint method.

I am assuming that is because ContinuousSource.end_time=1.0e20. If I only turn on the source for 1 period i.e
sources = [mp.EigenmodeSource(mp.ContinuousSource(f,fwidth=0.2*f, end_time=1/f),<<other args here>>) for f in [1/0.5, 1/0.6]]
would that be enough for adjoint optimization in Meep or is there some other heuristic I should follow to compute how long a ContinuousSource should be turned on for with a FourierField objective? I take it the forward simulation doesn't end after the source is turned off and controlled by OptitimzationProblem.decay_by?

I have tried to make the suggested changes, leaving me with my objective functions as:

the OptimizationProblem.objective_functions as:

def J(source, top, bottom):
    obj_fn = npa.array(
        [
            npa.mean(npa.abs(top[0] ** 2) / npa.abs(bottom[0] ** 2)),
            npa.mean(npa.abs(bottom[1] ** 2) / npa.abs(top[1] ** 2)),
        ]
    )
    return obj_fn

And the cost function passed to the nlopt.LD_MMA solver:

def f(x, grad):
    f0, dJ_du = opt([x])
    if grad.size > 0:
        grad[:] = npa.sum(dJ_du)
    evaluation_history.append(np.real(f0))
    sensitivity[0] = dJ_du
    cur_iter[0] = cur_iter[0] + 1
    ax = plt.gca()
    opt.plot2D(
        False,
        ax=ax,
        plot_sources_flag=False,
        plot_monitors_flag=False,
        plot_boundaries_flag=False,
    )

    if mp.am_master():
        plt.savefig(f"./{filename}/new_splitter_{cur_iter[0]}.png", dpi=300)
    return np.real(f0)

This runs fine for the forward runs and seems to calculate the gradient in the nlopt optimization. But right after returning from f I get the following error:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/user/InverseDesign/meep/new_splitter/new_splitter_freespace.py", line 421, in <module>
    x[:] = solver.optimize(x)
  File "/home/user/anaconda3/envs/pmp/lib/python3.9/site-packages/nlopt.py", line 335, in optimize
    return _nlopt.opt_optimize(self, *args)
ValueError: nlopt invalid argument

I am not sure if this is an issue with how I am setting up J or how I am handling the gradients in f. It seems like grad is the correct shape ((10000,) both when entering and leaving J). Any insight would be greatly appreciated.

[smartalecH]

I have made these adjustments to the simulation but is seems that by changing the sources to a ContinuousSource the simulation never completes 1 forward run of the adjoint method.

Whoops! Sorry, I meant multiple GaussianSource time profiles. That was a typo (fixed above).

Recall that a ContinuousSource cannot be used with the current flavor of adjoint simulations, as the DFT fields would never converge.

[smartalecH]

I am not sure if this is an issue with how I am setting up J or how I am handling the gradients in f. It seems like grad is the correct shape ((10000,) both when entering and leaving J). Any insight would be greatly appreciated.

Your objective function is currently outputting two values (multi-objective). You either need to do an epigraph (see the tutorials) or sum your objective so that it's scalar-valued (note this is not unique to meep, but rather an intrinsic characteristic of multi-objective optimization).

For example, if you opt for the scalar-valued approach, you could simply do:

def J(source, top, bottom):
    return npa.mean(npa.abs(top[0] ** 2) / npa.abs(source[0] ** 2)) + npa.mean(npa.abs(bottom[1] ** 2) / npa.abs(source[1] ** 2))

If you opt for the epigraph, make sure you take care of the gradient array properly (again, discussed in the above-linked tutorial).

[DurhamSmith]

@smartalecH I had tried that as my objective function before I posted my reply but ran into:

File "/home/user/PhD/InverseDesign/meep/new_splitter/new_splitter_freespace.py", line 425, in <module>
    x[:] = solver.optimize(x)
  File "/home/user/anaconda3/envs/pmp/lib/python3.9/site-packages/nlopt.py", line 335, in optimize
    return _nlopt.opt_optimize(self, *args)
  File "/home/user/PhD/InverseDesign/meep/new_splitter/new_splitter_freespace.py", line 328, in f
    f0, dJ_du = opt([x])
  File "/home/user/anaconda3/envs/pmp/lib/python3.9/site-packages/meep/adjoint/optimization_problem.py", line 171, in __call__
    self.adjoint_run()
  File "/home/user/anaconda3/envs/pmp/lib/python3.9/site-packages/meep/adjoint/optimization_problem.py", line 289, in adjoint_run
    self.prepare_adjoint_run()
  File "/home/user/anaconda3/envs/pmp/lib/python3.9/site-packages/meep/adjoint/optimization_problem.py", line 284, in prepare_adjoint_run
    self.adjoint_sources[ar] += m.place_adjoint_source(
  File "/home/user/anaconda3/envs/pmp/lib/python3.9/site-packages/meep/adjoint/objective.py", line 288, in place_adjoint_source
    dJ_4d = np.array([
  File "/home/user/anaconda3/envs/pmp/lib/python3.9/site-packages/meep/adjoint/objective.py", line 289, in <listcomp>
    dJ[f].copy().reshape(x_dim, y_dim, z_dim)
ValueError: cannot reshape array of size 1 into shape (1,96,1)

I know that the objective needs to be scalar valued, but I thought by outputting two values they would be interpreted as the objective functions for each frequency. I must have been interpreting this incorrectly
When I try not take the mean I can at least get past this by making J not return a scalar but instead be defined over the full region of the source:

def J(source, top, bottom):
    obj_fn = npa.abs(top[0] ** 2) / npa.abs(bottom[0] ** 2) + npa.abs(
        bottom[1] ** 2
    ) / npa.abs(top[1] ** 2)
    return obj_fn

But then I run into the same error:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/user/PhD/InverseDesign/meep/new_splitter/new_splitter_freespace.py", line 430, in <module>
    x[:] = solver.optimize(x)
  File "/home/user/anaconda3/envs/pmp/lib/python3.9/site-packages/nlopt.py", line 335, in optimize
    return _nlopt.opt_optimize(self, *args)
ValueError: nlopt invalid argument

When calling f, which currently is:

def f(x, grad):
    print("Current iteration: {}".format(cur_iter[0] + 1))
    f0, dJ_du = opt([x])
    if grad.size > 0:
        grad[:] = npa.sum(dJ_du, axis=1)
    evaluation_history.append(np.real(f0))
    sensitivity[0] = dJ_du
    cur_iter[0] = cur_iter[0] + 1
    ax = plt.gca()
    opt.plot2D(
        False,
        ax=ax,
        plot_sources_flag=False,
        plot_monitors_flag=False,
        plot_boundaries_flag=False,
    )

    if mp.am_master():
        plt.savefig(f"./{filename}/new_splitter_{cur_iter[0]}.png", dpi=300)
    return np.real(f0)

I am currently using the Meep version available through anaconda (1, 2.1) if that might be an issue, I havn't inspected the changes between that and the current master version.

[DurhamSmith]

Have tried the epigraph formulation. When I dont try normalize by the source power, the density doesn't change at all.

def J(source, top, bottom):
  obj_fn =  - npa.array(
      [
          npa.mean(npa.abs(top[0] ** 2) / npa.abs(bottom[0] ** 2)),
          npa.mean(npa.abs(bottom[1] ** 2) / npa.abs(top[1] ** 2)),
      ]
  )
  return obj_fn

def f(x, grad):
  t = x[0]  # "dummy" parameter
  v = x[1:]  # design parameters
  if grad.size > 0:
      grad[0] = 1
      grad[1:] = 0
  return t

def mapping(x, eta, beta):

  filtered_field = mpa.conic_filter(
      x,
      filter_radius,
      design_region_dimensions.x - 1 / design_region_resolution,
      design_region_dimensions.y - 1 / design_region_resolution,
      design_region_resolution,
  )
  projected_field = mpa.tanh_projection(filtered_field, beta, eta)
  return projected_field.flatten()

def c(result, x, gradient, eta, beta):
  t = x[0]  # dummy parameter
  v = x[1:]  # design parameters

  f0, dJ_du = opt([mapping(v, eta, beta)])


  my_grad = np.zeros(dJ_du.shape)
  for k in range(opt.nf):
      my_grad[:, k] = tensor_jacobian_product(mapping, 0)(v, eta, beta, dJ_du[:, k])

  # Assign gradients
  if gradient.size > 0:
      gradient[:, 0] = -1  # gradient w.r.t. "t"
      gradient[:, 1:] = my_grad.T  # gradient w.r.t. each frequency objective

  result[:] = np.real(f0) - t
  # store results
  evaluation_history.append(np.real(f0))
  cur_iter[0] = cur_iter[0] + 1

If I do include some normalization to J it is a little bit better, at least there looks to be change in the density

def J(source, top, bottom):

    obj_fn = 1 - npa.array(
        [
            (npa.abs(npa.mean(top[0]) / npa.mean(source[0]))) ** 2
            / (npa.abs(npa.mean(bottom[0]) / npa.mean(source[0]))) ** 2,
            (npa.abs(npa.mean(bottom[1]) / npa.mean(source[1]))) ** 2
            / (npa.abs(npa.mean(top[1]) / npa.mean(source[1]))) ** 2,
        ]
    )
    return obj_fn

But the there splitting of the 500(top) and 700nm(bottom) wavelength light.

If I try go for just a scalar value for J:

def J(source, top, bottom):
    return npa.mean(npa.abs(top[0] ** 2) / npa.abs(bottom[0] ** 2)) + npa.mean(npa.abs(bottom[1] ** 2) / npa.abs(top[1] ** 2))

I run into reshape errors when Meep trys to setup the adjoint field.

File "/home/user/PhD/InverseDesign/meep/new_splitter/new_splitter_freespace.py", line 425, in <module>
    x[:] = solver.optimize(x)
  File "/home/user/anaconda3/envs/pmp/lib/python3.9/site-packages/nlopt.py", line 335, in optimize
    return _nlopt.opt_optimize(self, *args)
  File "/home/user/PhD/InverseDesign/meep/new_splitter/new_splitter_freespace.py", line 328, in f
    f0, dJ_du = opt([x])
  File "/home/user/anaconda3/envs/pmp/lib/python3.9/site-packages/meep/adjoint/optimization_problem.py", line 171, in __call__
    self.adjoint_run()
  File "/home/user/anaconda3/envs/pmp/lib/python3.9/site-packages/meep/adjoint/optimization_problem.py", line 289, in adjoint_run
    self.prepare_adjoint_run()
  File "/home/user/anaconda3/envs/pmp/lib/python3.9/site-packages/meep/adjoint/optimization_problem.py", line 284, in prepare_adjoint_run
    self.adjoint_sources[ar] += m.place_adjoint_source(
  File "/home/user/anaconda3/envs/pmp/lib/python3.9/site-packages/meep/adjoint/objective.py", line 288, in place_adjoint_source
    dJ_4d = np.array([
  File "/home/user/anaconda3/envs/pmp/lib/python3.9/site-packages/meep/adjoint/objective.py", line 289, in <listcomp>
    dJ[f].copy().reshape(x_dim, y_dim, z_dim)
ValueError: cannot reshape array of size 1 into shape (1,96,1)

I know the tutorials are out-dated but just as an FYI this occurs in the Fourier Fields Optimization in Waveguide Bend.ipynb example as well when defining power = npa.abs(top[1,7])**2. When setting power = npa.abs(top[1])**2 at least the adjoint field in this example can be calculated but the fails at x[:] = solver.optimize(x). This is similar behavior to what I posted here.

I have also tried various other way of handling the gradients when using an objective function where I do not take the mean and thus am able to get the adjoint simulation working when there are multiple frequencies:

def J(source, top, bottom):
    obj_fn = npa.abs(top[0] ** 2) / npa.abs(bottom[0] ** 2) + npa.abs(
        bottom[1] ** 2
    ) / npa.abs(top[1] ** 2)
    return obj_fn

But I have had no luck for the different ways I've tried handling these gradients in f.

Any advice would be greatly appreciated.

[oskooi]

The objective function J must return an array with the same number of elements as the number of frequencies passed to the OptimizationProblem constructor. Here is an example based on a 2d filter involving an incident E_z-polarized pulsed planewave and a rectangular silicon design region in air with an objective function defined to maximize |E_z|² at a given frequency and spatial point on a DFT line monitor at the focal distance (chosen arbitrarily as 2 μm from the filter) while minimizing the intensity at all other frequencies spanning visible wavelengths in the range [0.4, 0.7] μm. Only a single iteration of the adjoint solver is involved consisting of a forward and adjoint simulation followed by the computation of the gradient from the forward and adjoint fields. Note that the planewave extends into the PML region which requires specifying is_integrated=True and also using periodic boundary conditions via setting k_point=mp.Vector3() in the Simulation constructor. Note also that I specified the maximum_run_time of the OptimizationProblem just to make the forward/adjoint simulations run quickly (for demonstration purposes) but this should be removed in the actual run.

import meep as mp
import meep.adjoint as mpa
import numpy as np
from autograd import numpy as npa

Si = mp.Medium(index=3.5)

resolution = 20 # pixels/μm

wvl_min = 0.4
wvl_max = 0.7
frq_min = 1/wvl_max
frq_max = 1/wvl_min
fcen = 0.5*(frq_min+frq_max)
df = frq_max-frq_min

nfreqs = 21
freqs = 1/np.linspace(wvl_min, wvl_max, nfreqs)

dpml = 0.5*wvl_max
dair = wvl_max
dfocus = 2.0 # focal distance from the nearest edge of the filter
design_region_shape = mp.Vector3(1.,2.,0)

sx = dpml + dair + design_region_shape.x + dfocus + dair + dpml
sy = dpml + design_region_shape.y + dpml
cell_size = mp.Vector3(sx,sy)

pml_layers = [mp.PML(dpml)]

source = [mp.Source(mp.GaussianSource(frequency=fcen,fwidth=df,is_integrated=True),
                    size=mp.Vector3(0,sy,0),
                    center=mp.Vector3(-0.5*sx+dpml),
                    component=mp.Ez)]

design_region_resolution = 2*resolution
Nx = int(design_region_shape.x*design_region_resolution) + 1
Ny = int(design_region_shape.y*design_region_resolution) + 1

design_variables = mp.MaterialGrid(mp.Vector3(Nx,Ny),
                                   mp.air,
                                   Si,
                                   weights=np.ones((Nx,Ny)))

design_region_center = mp.Vector3(x=-0.5*sx+dpml+dair+0.5*design_region_shape.x)

design_region = mpa.DesignRegion(design_variables,
                                 volume=mp.Volume(center=design_region_center,
                                                  size=design_region_shape))

geometry = [
    mp.Block(center=design_region_center,
             size=design_region_shape,
             material=design_variables)
]

sim = mp.Simulation(cell_size=cell_size,
                    boundary_layers=pml_layers,
                    geometry=geometry,
                    sources=source,
                    k_point=mp.Vector3(),
                    resolution=resolution)

# line monitor for the DFT fields at the focal distance
ob_list = [mpa.FourierFields(sim,
                             mp.Volume(center=mp.Vector3(x=0.5*sx-dpml-dair),
                                       size=mp.Vector3(y=sy-2*dpml)),
                             component=mp.Ez)]

# scalar objective function for maximizing  
# Ez intensity at a given frequency and spatial point        
# of the DFT line monitor while minimizing Ez intensity 
# at all other frequencies at that same point
def J(mon):
    frq_idx = 5
    pt_idx = 10
    intensity = npa.concatenate((-1*npa.power(npa.abs(mon[:frq_idx,pt_idx]),2),
                                 [npa.power(npa.abs(mon[frq_idx,pt_idx]),2)],
                                 -1*npa.power(npa.abs(mon[frq_idx+1:,pt_idx]),2)))
    return intensity

opt = mpa.OptimizationProblem(
    simulation=sim,
    objective_functions=J,
    objective_arguments=ob_list,
    design_regions=[design_region],
    frequencies=freqs,
    maximum_run_time=100,
)

x0 = 0.5*np.ones((Nx*Ny,))
val, grad = opt([x0])
print(f"{val.shape}, {grad.shape}")

output

Starting forward run...
-----------
Initializing structure...
time for choose_chunkdivision = 0.000122794 s
Working in 2D dimensions.
Computational cell is 5.1 x 3.1 x 0 with resolution 20
     block, center = (-1,0,0)
          size (1,1,0)
          axes (1,0,0), (0,1,0), (0,0,1)
time for set_epsilon = 0.00905608 s
-----------
run 0 finished at t = 100.025 (4001 timesteps)
Starting adjoint run...
/meep/python/meep/adjoint/filter_source.py:152: RuntimeWarning: divide by zero encountered in true_divide
  l2_err = np.sum(np.abs(H - H_hat.T)**2 / np.abs(H)**2)
-----------
Initializing structure...
time for choose_chunkdivision = 8.8381e-05 s
Working in 2D dimensions.
Computational cell is 5.1 x 3.1 x 0 with resolution 20
     block, center = (-1,0,0)
          size (1,1,0)
          axes (1,0,0), (0,1,0), (0,0,1)
time for set_epsilon = 0.00986814 s
-----------
run 1 finished at t = 100.025 (4001 timesteps)
Calculating gradient...
(21,), (1681, 21)

[DurhamSmith]

Thanks for this. I had tried to return an array with the same number of elements as the number of frequencies passed to the OptimizationProblem. I would still run into the

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/user/PhD/InverseDesign/meep/new_splitter/new_splitter_freespace.py", line 430, in <module>
    x[:] = solver.optimize(x)
  File "/home/user/anaconda3/envs/pmp/lib/python3.9/site-packages/nlopt.py", line 335, in optimize
    return _nlopt.opt_optimize(self, *args)
ValueError: nlopt invalid argument

type errors.

I solved this by making f return the sum of the FOM at different frequencies returned by J, i.e

def J(source, top, bottom):
    top_intensity = npa.power(npa.mean(npa.abs(top[0])), 2) - 1 * npa.power(
        npa.mean(npa.abs(bottom[0])), 2
    )
    bottom_intensity = npa.power(npa.mean(npa.abs(bottom[1])), 2) - 1 * npa.power(
        npa.mean(npa.abs(top[1])), 2
    )
    intensity = npa.array([top_intensity, bottom_intensity])
    return intensity

def f(x, gradient, cur_beta):
    f0, dJ_du = opt([mapping(x, eta_i, cur_beta)])  # compute objective and gradient

    if gradient.size > 0:
        gradient[:] = tensor_jacobian_product(mapping, 0)(
            x, eta_i, cur_beta, np.sum(dJ_du, axis=1)
        ) 

    evaluation_history.append(np.real(f0))
    sensitivity[0] = dJ_du
    cur_iter[0] = cur_iter[0] + 1
    return np.sum(np.real(f0))

This gives some results going from initial power splitting of

To this after 72 optimization round

I then tried adjusting J to be:

def J(source, top, bottom):
    top_intensity = npa.power(npa.mean(npa.abs(top[0])), 2) / npa.power(
        npa.mean(npa.abs(bottom[0])), 2
    )
    bottom_intensity = npa.power(npa.mean(npa.abs(bottom[1])), 2) / npa.power(
        npa.mean(npa.abs(top[1])), 2
    )
    intensity = npa.array([top_intensity, bottom_intensity])
    return intensity

while keeping everything else the same and got much worse results:

I would have expected this J to perform better. Is there any way I could have intuited that this would have performed worse? Or could this be due to me incorrectly setting up the simulations?

I have just launched a run that uses to topmost J with more optimization steps and finer adjustment of Beta to see if I can get better results.
Thanks for all the help so far, tremendously appreciated.

[stevengj]

maximize |Ez|2 at a given frequency and spatial point on a DFT line monitor at the focal distance (chosen arbitrarily as 2 μm from the filter) while minimizing the intensity at all other frequencies spanning

This isn't going to work. You're trying to do something like

where E_k are the points where you are maximizing the intensity and E_j are the points where you are minimizing the intensity. But in the max–min formulation, the latter will always be the minimums since the former can never be negative, so what you are effectively doing is just:

i.e. minimizing the field intensity at some points, and the end result will probably just be a mirror.

What I would suggest is simply maximizing the minimum intensities at the points where you want a focus. Trust that conservation of energy will imply that if you get good focal concentration, the power elsewhere will be low. This is much like how we design RGB metalenses, for example.

[smartalecH]

If I get time this week, I'll try throwing together a tutorial that can design this structure.

[DurhamSmith]

I would really appreciate that! The ideal would be if there were 3 monitors stacked horizontally where 400-500nm wavelength is focused to 1 monitor, 500-600nm to another and 600-700nm to another. Although this is a lot to ask and any help just getting my own sims to work is greatly appreciated!

[oskooi]

For your reference, a manuscript describing the inner workings of Meep's adjoint solver was published today in Optics Express: https://opg.optica.org/oe/fulltext.cfm?uri=oe-30-3-4467&id=468836.

[DurhamSmith]

Thanks!

[DurhamSmith]

I have some updates and additional questions:

Epigraph formulation seems to starting to bear fruit for what I want to do. I get results that at least show that the ratio of power from source to the desired focus point (my figure of merit) is driven in the correct direction. This is the plot of the sum of the absolute values of the fields (using a FourierFields monitor as suggested by @oskooi and @smartalecH ).

The top monitor is a focus for the first 6 frequencies points and bottom monitor for the next 6. On closer inspection if I look at the fields at the monitors I see that what actually seems to be happening is that the power measured at the source monitor (np.sum(np.abs(source**2))) is reduced from the first run to the last run. This is the monitor measurements before the optimization starts:

I suspect that this is because the figure of merit normalizes by the source monitors power, so an effective means of increasing the FOM is to reduce the power at the source monitor, so the device structure that results from this optimization does some focusing of the power between the top and bottom monitors, but also trys to have a lot of destructive interference at the source monitor (via reflections off the device) to minimize the source power.

Is there a way to avoid this? I have tried to remove the source term normalization in the calculation, without much success, I am not sure if I formulated the epigraph problem correctly in this case I changed it from minimization (FOM=1-monitor_power/source_power) to a maximization problem (FOM=monitor_power), but this does not seem to correspond to a well formed minimax problem. I still need to do some more research into this type of optimization.

I also would like to know if I want to make lenses that focus power onto a specific region is it better to sum the power over all points in that region or would taking the power just at the central location be a good enough solution? It would seem that over the whole field would be optimal. In preliminary testing it seemed that it did not work as well, but that could be because I messed up the epigraph formalization when making it a maximization problem. One thing I did notice was that it was much faster when only optimizing one point.

What is the correct way to start/restart a optimization? Do I use the dump_structure and load_structure functions? Do I save the opt.design_regions design_parameters.weights?

How do I look at the actual values of the dielectric? I see that according to the paper referenced by @oskooi they get interpolated but when I look at opt.design_regions[0].design_parameters.weights they are all much higher values than my materials dielectric constant/index of refraction. Do I use sim.get_epsilon()?

The source monitor seems to be non-uniform, could this be because I am placing my sources incorrectly? @smartalecH's comment suggests that I use multiple sources, Is there a good rule of thumb for how many to use? I tried to use 3 sources and space them so their full width at half max values values coincided with one another. without much success. I tried the same with a source at every frequency and again trying to match full width half maxs but I could not get good results to have some flat response.

Lastly I seem to get strange behavior when plotting the different values of my monitors.
When i try plot using the following code

  frequencies = opt.frequencies
  source_coef, top_coef, bottom_coef = [m._eval for m in opt.objective_arguments]
  top_profile = (
        np.sum(np.abs(top_coef), axis=1) / np.sum(np.abs(source_coef), axis=1)
    ) ** 2
    bottom_profile = (
        np.sum(np.abs(bottom_coef), axis=1) / np.sum(np.abs(source_coef), axis=1)
    ) ** 2

    plt.figure()
    plt.plot(1 / frequencies, top_profile * 100, "-o", label="Top Arm")
    plt.plot(1 / frequencies, bottom_profile * 100, "--o", label="Bottom Arm")
    plt.legend()
    plt.grid(True)
    plt.xlabel("Wavelength")
    plt.ylabel("Splitting Ratio (% of source power)")
    if mp.am_master():
        plt.savefig(f"./{filename}/Power_Splitting_vs_wavelength_div_first_{label}")

I get the following result:

When plotting with

  top_profile = np.sum(np.abs(top_coef ** 2), axis=1) / np.sum(
        np.abs(source_coef ** 2), axis=1
    )
    bottom_profile = np.sum(np.abs(bottom_coef ** 2), axis=1) / np.sum(
        np.abs(source_coef ** 2), axis=1
    )

    plt.figure()
    plt.plot(1 / frequencies, top_profile * 100, "-o", label="Top Arm")
    plt.plot(1 / frequencies, bottom_profile * 100, "--o", label="Bottom Arm")
    plt.legend()
    plt.grid(True)
    plt.xlabel("Wavelength")
    plt.ylabel("Splitting Ratio (% of source power)")
    if mp.am_master():
        plt.savefig(f"./{filename}/Power_Splitting_vs_wavelength_div_last_{label}")

I get :

Which seems incorrect (the first figure in my post is the fields measured at the monitors that these plots come from). I dont understand how there can be 100% splitting efficiency, with power left in the bottom arm and also the monitor measurements at the source and top/bottom arm not being equal.

The code used to generate these figures and observations is:

# coding: utf-8
import meep as mp
import meep.adjoint as mpa
import numpy as np
from autograd import numpy as npa
from autograd import tensor_jacobian_product
from enum import Enum
import matplotlib.pyplot as plt
import pickle as pkl
import nlopt
from matplotlib import pyplot as plt
from matplotlib.patches import Circle
from scipy import special, signal
import sys

sys.path.append("/home/user/PhD/InverseDesign/meep/")
import helper_fns

# TODO create file if doesnt exist
filename = "output"
seed = 44
np.random.seed(seed)


def plot_results(opt, label):

    plt.figure()
    ax = plt.gca()
    opt.plot2D(
        True
    )
    circ = Circle((2, 2), minimum_length / 2)
    ax.add_patch(circ)
    ax.axis("off")
    if mp.am_master():
        plt.savefig(f"./{filename}/new_splitter_{label}.png", dpi=300)


    frequencies = opt.frequencies
    source_coef, top_coef, bottom_coef = [m._eval for m in opt.objective_arguments]
    top_profile = (
        np.sum(np.abs(top_coef), axis=1) / np.sum(np.abs(source_coef), axis=1)
    ) ** 2
    bottom_profile = (
        np.sum(np.abs(bottom_coef), axis=1) / np.sum(np.abs(source_coef), axis=1)
    ) ** 2

    plt.figure()
    plt.plot(1 / frequencies, top_profile * 100, "-o", label="Top Arm")
    plt.plot(1 / frequencies, bottom_profile * 100, "--o", label="Bottom Arm")
    plt.legend()
    plt.grid(True)
    plt.xlabel("Wavelength")
    plt.ylabel("Splitting Ratio (% of source power)")
    if mp.am_master():
        plt.savefig(f"./{filename}/Power_Splitting_vs_wavelength_div_first_{label}")

    top_profile = np.sum(np.abs(top_coef ** 2), axis=1) / np.sum(
        np.abs(source_coef ** 2), axis=1
    )
    bottom_profile = np.sum(np.abs(bottom_coef ** 2), axis=1) / np.sum(
        np.abs(source_coef ** 2), axis=1
    )

    plt.figure()
    plt.plot(1 / frequencies, top_profile * 100, "-o", label="Top Arm")
    plt.plot(1 / frequencies, bottom_profile * 100, "--o", label="Bottom Arm")
    plt.legend()
    plt.grid(True)
    plt.xlabel("Wavelength")
    plt.ylabel("Splitting Ratio (% of source power)")
    if mp.am_master():
        plt.savefig(f"./{filename}/Power_Splitting_vs_wavelength_div_last_{label}")

    plt.figure()
    plt.plot(
        1 / frequencies, npa.sum(npa.abs(source_coef), axis=1), "-o", label="Source"
    )
    plt.plot(
        1 / frequencies, npa.sum(npa.abs(bottom_coef), axis=1), "--o", label="Bottom"
    )
    plt.plot(1 / frequencies, npa.sum(npa.abs(top_coef), axis=1), "--o", label="Top")
    plt.legend()
    plt.grid(True)
    plt.xlabel("Wavelength")
    plt.ylabel("npa.abs of monitor")
    if mp.am_master():
        plt.savefig(f"./{filename}/Monitor_Measurements_{label}")


###############################################################################
#              Setup variables that define the simulation domain              #
###############################################################################
design_region_dimensions = mp.Vector3(2.0, 2.0)  # Lets optimze a 2um^2 region
pml_size = mp.Vector3(0.35, 0.35)  # Half a wavelength of red light
pml_src_offset = mp.Vector3(0.1)  # Distance from PML to source plane
src_device_offset = mp.Vector3(0.2)
device_objective_plane_offset = mp.Vector3(1.5)
objective_plane_pml_offset = mp.Vector3(0.3)
padding = mp.Vector3(0, 0.01, 0)
cell_size = (
    2 * pml_size
    + pml_src_offset
    + src_device_offset
    + design_region_dimensions
    + device_objective_plane_offset
    + objective_plane_pml_offset
    + 2 * padding
)
###############################################################################
#                               Define Materials                              #
###############################################################################
n_sio2 = 1.46
SiO2 = mp.Medium(index=n_sio2)
n_tio2 = 2.61
TiO2 = mp.Medium(index=n_tio2)
###############################################################################
#                             Setup Design Region                             #
###############################################################################
design_region_resolution = (
    50
)
design_region_pixels = design_region_resolution * design_region_dimensions
design_variables = mp.MaterialGrid(
    design_region_pixels,
    SiO2,
    TiO2,
    grid_type="U_DEFAULT",  # I am not sure if this is the correct choice
)
design_region_center = mp.Vector3(
    x=(pml_size + pml_src_offset + src_device_offset + design_region_dimensions / 2).x
)
design_region = mpa.DesignRegion(
    design_variables,
    volume=mp.Volume(center=design_region_center, size=design_region_dimensions),
)
###############################################################################
#                               Setup the source                              #
###############################################################################
source_center = mp.Vector3(x=(pml_size + pml_src_offset).x)
source_size = mp.Vector3(0, (cell_size - 2 * pml_size - 2 * padding).y, 0)
kpoint = mp.Vector3(1, 0, 0)
frequencies = np.linspace(1 / 0.4, 1 / 0.7, 12)
sources = [
    mp.EigenModeSource(
        mp.GaussianSource(f, fwidth=0.2 * f),
        eig_band=1,
        direction=mp.NO_DIRECTION,
        eig_kpoint=kpoint,
        size=source_size,
        center=source_center,
    )
    for f in frequencies
]

###############################################################################
#                             Setup the simulation                            #
###############################################################################
resolution = (
    100  # Decide this better the recomend 8px/shortest wavelength in highest dielectric
)
pml_layers = [mp.PML(pml_size.x)]
geometry = [
    mp.Block(
        center=design_region.center, size=design_region.size, material=design_variables
    ),
]
sim = mp.Simulation(
    cell_size=cell_size,
    boundary_layers=pml_layers,
    geometry=geometry,
    sources=sources,
    # eps_averaging=False,
    resolution=resolution,
    geometry_center=mp.Vector3(x=cell_size.x / 2),
)
###############################################################################
#                                   Filters                                   #
###############################################################################
minimum_length = 0.02
eta_i = (
    0.5
)
eta_e = 0.55 
eta_d = 1 - eta_e
filter_radius = mpa.get_conic_radius_from_eta_e(minimum_length, eta_e)

def mapping(x, eta, beta):
    filtered_field = mpa.conic_filter(
        x,
        filter_radius,
        design_region_dimensions.x - 1 / design_region_resolution,
        design_region_dimensions.y - 1 / design_region_resolution,
        design_region_resolution,
    )
    projected_field = mpa.tanh_projection(filtered_field, beta, eta)
    return projected_field.flatten()
###############################################################################
#                             OptimizationProblem                             #
###############################################################################
source_to_monitor_offset = mp.Vector3(x=0.01)
source_monitor_center = mp.Vector3(
    (pml_size + pml_src_offset + source_to_monitor_offset).x
)
source_monitor_size = source_size
TE0 = mpa.FourierFields(
    sim,
    mp.Volume(center=source_monitor_center, size=source_monitor_size,),
    component=mp.Ez,
)
objective_plane_monitor_center_x = (
    pml_size
    + pml_src_offset
    + src_device_offset
    + design_region_dimensions
    + device_objective_plane_offset
).x
objective_plane_monitor_center_y = design_region_dimensions.y / 4
top_objective_plane_monitor_center = mp.Vector3(
    objective_plane_monitor_center_x, objective_plane_monitor_center_y
)
top_objective_plane_monitor_size = mp.Vector3(y=(design_region_dimensions.y - 0.1) / 2)
TE_top = mpa.FourierFields(
    sim,
    mp.Volume(
        center=top_objective_plane_monitor_center,
        size=top_objective_plane_monitor_size,
    ),
    component=mp.Ez,
)
bottom_objective_plane_monitor_center = mp.Vector3(
    objective_plane_monitor_center_x,
    -objective_plane_monitor_center_y,  # Notice we subtract so we put below the top monitor
)
bottom_objective_plane_monitor_size = mp.Vector3(
    y=(design_region_dimensions.y - 0.1) / 2
)

TE_bottom = mpa.FourierFields(
    sim,
    mp.Volume(
        center=bottom_objective_plane_monitor_center,
        size=bottom_objective_plane_monitor_size,
    ),
    component=mp.Ez,
)

ob_list = [TE0, TE_top, TE_bottom]

def J(source, top, bottom):
    # obj_fn = 1 - (npa.abs(top[0] / bottom[0]) + npa.abs(bottom[1] / top[1]))
    freq_idx = 6

    obj_fn = []
    print(f"Source Shape {source[1].shape}")
    print(f"Top Shape {top[1].shape}")
    print(
        f"POWER {npa.power(npa.sum(npa.abs(top[0])) / npa.sum(np.abs(source[0])), 2)}"
    )
    for i in range(top.shape[0]):
        if i < freq_idx:
            obj_fn.append(
                npa.sum(npa.abs(npa.power(top[i], 2)))
                / npa.sum(np.abs(npa.power(source[i], 2)))
            )
        else:
            obj_fn.append(
                npa.sum(npa.abs(npa.power(bottom[i], 2)))
                / npa.sum(np.abs(npa.power(source[i], 2)))
            )


    obj_fn = npa.array(obj_fn)
    obj_fn = npa.abs(1 - obj_fn)
    print(f"==================== obj_fn {obj_fn} ====================")
    return obj_fn


opt = mpa.OptimizationProblem(
    simulation=sim,
    objective_functions=J,
    objective_arguments=ob_list,
    design_regions=[design_region],
    frequencies=frequencies,
    decay_by=1e-6,
)


###########################################################################
#                             Setup Optimizer                             #
###########################################################################
evaluation_history = []
sensitivity = [0]
cur_iter = [0]


def f(x, grad):
    t = x[0]  # "dummy" parameter
    v = x[1:]  # design parameters
    if grad.size > 0:
        grad[0] = 1
        grad[1:] = 0
    return t


def c(result, x, gradient, eta, beta):
    print(
        "Current iteration: {}; current eta: {}, current beta: {}".format(
            cur_iter[0], eta, beta
        )
    )
    # print(f"-----------result : {result.shape}--------------")
    # print(f"-----------x : {x.shape}--------------")
    # print(f"-----------gradient : {gradient.shape}--------------")
    t = x[0]  # dummy parameter
    v = x[1:]  # design parameters

    f0, dJ_du = opt([mapping(v, eta, beta)])
    print(f"-----------f0 : {f0}--------------\n\n")
    # print(f"=================dJ.shape: {dJ_du.shape}================\n\n")
    # print(f"=================dJ: {dJ_du}=================\n\n")   # dJ.shape: (10000, 2)
    # Backprop the gradients through our mapping function
    my_grad = np.zeros(dJ_du.shape)
    for k in range(opt.nf):
        my_grad[:, k] = tensor_jacobian_product(mapping, 0)(v, eta, beta, dJ_du[:, k])

    # Assign gradients
    if gradient.size > 0:
        gradient[:, 0] = -1  # gradient w.r.t. "t"
        gradient[:, 1:] = my_grad.T  # gradient w.r.t. each frequency objective

    result[:] = np.real(f0) - t
    print(f"-----------result : {result}--------------\n\n")
    print(f"-----------gradient : {gradient}--------------\n\n")

    # store results
    evaluation_history.append(np.real(f0))
    plot_results(opt, cur_iter[0])

    cur_iter[0] = cur_iter[0] + 1

###########################################################################
#                          Inital Paramater Setup                         #
###########################################################################
algorithm = nlopt.LD_MMA
n = int(design_region_pixels.x * design_region_pixels.y)
x0 = np.ones((n,)) * 0.5
x = np.ones((n,)) * 0.5
lb = np.zeros((n,))
ub = np.ones((n,))
opt.update_design([x0])
x = np.insert(x, 0, 0.5)  # our initial guess for the worst error
lb = np.insert(lb, 0, 0)  # we can't get less than 0 error!
ub = np.insert(ub, 0, 1)  # we can't get more than 1 error!
###############################################################################
#                                Run Optimizer                                #
###############################################################################
cur_beta = 4
beta_scale = 2
num_betas = 10
update_factor = 30
for iters in range(num_betas):
    helper_fns.write_dielectric(opt, f"dielectric_{iter}")
    solver = nlopt.opt(algorithm, n + 1)
    solver.set_lower_bounds(lb)
    solver.set_upper_bounds(ub)
    solver.set_min_objective(f)
    solver.set_maxeval(update_factor)
    solver.add_inequality_mconstraint(
        lambda r, x, g: c(r, x, g, eta_i, cur_beta),
        np.array([1e-3] * len(opt.frequencies)),  # len(opt.frequencies)b
    )
    x[:] = solver.optimize(x)
    cur_beta = cur_beta * beta_scale