NanoComp / meep

free finite-difference time-domain (FDTD) software for electromagnetic simulations
GNU General Public License v2.0
1.22k stars 621 forks source link

Clarification on issues experience with Adjoint Simulations in Meep #1888

Closed DurhamSmith closed 2 years ago

DurhamSmith commented 2 years ago

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.

new_splitter_initial

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).

example_abs_source_coef

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:

Source_Monitor_Measurement_init

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:

Monitor_Measurements_init

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):

Power_Splitting_vs_wavelength_init

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: Monitor_Measurements_init Power_Splitting_vs_wavelength_init

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: FOM_Evolution

The monitor measurements and ratios are: Monitor_Measurements

Power_Splitting_vs_wavelength

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 commented 2 years ago

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 commented 2 years ago

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 commented 2 years ago

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.ContinuousSource(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]).

DurhamSmith commented 2 years ago

@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 commented 2 years ago

Whoops sorry I meant multiple GaussianSource time profiles. That was a typo.

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

smartalecH commented 2 years ago

Also 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).