Non-materialized dft fields for adjoint calculations

[smartalecH]

Closes #1797, #1832, #1861

TODO

• Add attribute to dft_fields allowing them to persist after a simulation is run
• Store additional boundary (i.e. padding) around dft subvolume
• Devise way to easily access different components from single string of chunks
• Pass tests

[smartalecH]

Currently, the dft_array object creates a "linked list" of sorts for all of the chunks on the current process, regardless of the field component. As such there doesn't seem to be a convenient way to traverse this list in order to find the proper field component during the recombination step.

For example, assume we are looping over the adjoint field chunks, and we are currently on the Dx chunk. Eventually, we will need the Dy and Dz forward field points that correspond to the current adjoint points (they will be spatially offset, of course). This is a bit tricky though because all of the chunks are cascaded together. We'd have to perform several checks on all the chunks at every point.

I was thinking it may be worth it to create a modified version of dft_fields (e.g. maybe dft_adjoint_fields) that explicitly separates the components, making the postprocessing step much easier. The Poynting Flux object does something like this.

Any thoughts, @stevengj & @oskooi ?

[stevengj]

Probably easier to just call add_dft_fields once for each component.

[stevengj]

As discussed in #1555, we'll probably need to ensure that loop_in_chunks only loops over owned points for this to work.

[smartalecH]

Status update for @oskooi

Currently, cases with complex fields are failing (including cylindrical coordinates). Also, the jax wrapper needs some love.

But other than that, things look good. There is still some error introduced (hopefully) due to the bug within #1555. But the error is rather small. And not present when only a single chunk holds DFT fields.

[smartalecH]

I can't replicate the test failures on my machine. @oskooi can you try one more time on your machine?

Note I've already bumped the tolerance, but am still seeing errors.

[oskooi]

With the most recent commit (9aa8296), there are two separate CI test failures:

1. test_adjoint_solver.py (single precision)
2. test_adjoint_cyl.py (double precision)

On my local machine, I was only able to reproduce (1).

[mawc2019]

This PR sometimes causes a segmentation fault in the adjoint run. An example for the EigenmodeCoefficient adjoint solver is as follows. Before this PR, there is no segmentation fault. After this PR, the segmentation fault appears.

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

Si = mp.Medium(index=3.48)
Air = mp.Medium(index=1)

resolution,design_resolution = 40,40

design_width,design_height = 1.0,0.4
Lx,Ly = design_width,4
dpml = 1.5

monitor_height,monitor_width = 1.5,0.8

pml_layers = [mp.PML(dpml)]
cell_size = mp.Vector3(Lx+2*dpml,Ly+2*dpml)

fcen,fwidth = 1/1.55,0.2/1.55

source_center,source_size  = mp.Vector3(0,-Ly/4),mp.Vector3(Lx,0)
src = mp.GaussianSource(frequency=fcen,fwidth=fwidth)
source = [mp.EigenModeSource(src,eig_band=1,size=source_size,center=source_center)]

Nx,Ny = int(round(design_resolution*design_width)),1

design_variables = mp.MaterialGrid(mp.Vector3(Nx,Ny),Air,Si,grid_type='U_MEAN')
design_region = mpa.DesignRegion(design_variables,volume=mp.Volume(center=mp.Vector3(0,design_height/2),size=mp.Vector3(design_width,design_height)))

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

sim = mp.Simulation(cell_size=cell_size,boundary_layers=pml_layers,geometry=geometry,sources=source,eps_averaging=False,resolution=resolution)

TE0 = mpa.EigenmodeCoefficient(sim,mp.Volume(center=mp.Vector3(0,monitor_height),size=mp.Vector3(monitor_width,0)),mode=1)
ob_list = [TE0]

def J(alpha):
    return npa.abs(alpha) ** 2

opt = mpa.OptimizationProblem(
    simulation = sim,objective_functions = J,objective_arguments = ob_list,
    design_regions = [design_region],fcen=fcen,df = 0,nf = 1,decay_by = 1e-5)

x0 = 0.5*np.ones(Nx*Ny)
opt.update_design([x0])
f0, dJ_deps = opt()

[mawc2019]

Before this PR, although there is no segmentation fault, the finite-difference and adjoint gradients are inconsistent in the above example. With the following code added to the end,

db = 1e-4
choose = 5
np.random.seed(1905)

g_discrete, idx = opt.calculate_fd_gradient(num_gradients=choose,db=db)
g_discrete = np.array(g_discrete).flatten()
g_adjoint = dJ_deps.flatten()

(m, b) = np.polyfit(np.squeeze(g_discrete), g_adjoint[idx], 1)
print("Randomly selected indices: ",idx)
print("Finite-difference gradient: ",g_discrete)
print("Adjoint gradient: ",g_adjoint[idx])

the results are printed as

Finite-difference gradient:  [ 0.03372703 -1.49253234 -0.21201943  1.02494487  1.02513747]
Adjoint gradient:  [2.9224977  1.75185184 3.31514926 1.11873224 1.11873224]

[mawc2019]

After the PR, in addition to the above issue in the EigenmodeCoefficient adjoint solver, sometimes the finite-difference and adjoint gradients are inconsistent in the FourierFields adjoint solver. The example code is as follows.

import meep as mp
import meep.adjoint as mpa
import numpy as np
from autograd import numpy as npa
Si, SiO2 = mp.Medium(index=3.4), mp.Medium(index=1.44)

resolution,design_resolution = 40,40
Lx,Ly,pml_size = 2,4,1
cell_size = mp.Vector3(Lx,Ly+2*pml_size)
pml_layers = [mp.PML(pml_size, direction=mp.Y)]
design_width,design_height = Lx-0.4,0.4

fcen,fwidth = 1/1.55,0.01/1.55
source_center,source_size = mp.Vector3(0,-Ly/4),mp.Vector3(Lx,0)

src = mp.GaussianSource(frequency=fcen,fwidth=fwidth)
source = [mp.EigenModeSource(src,eig_band=1,size=source_size,center=source_center)]
Nx,Ny = int(round(design_resolution*design_width)),int(round(design_resolution*design_height))

design_variables = mp.MaterialGrid(mp.Vector3(Nx,Ny),SiO2,Si,grid_type='U_MEAN')
design_region = mpa.DesignRegion(design_variables,volume=mp.Volume(center=mp.Vector3(), size=mp.Vector3(design_width, design_height)))

geometry = [mp.Block(center=design_region.center, size=design_region.size, material=design_variables)] # design region
sim = mp.Simulation(
    cell_size=cell_size,boundary_layers=pml_layers,geometry=geometry,
    sources=source,default_material=SiO2,resolution=resolution)

nf = mpa.FourierFields(sim,mp.Volume(center=mp.Vector3(0,design_height+0.2),size=mp.Vector3(design_width,0)),mp.Ex,yee_grid=True)
ob_list = [nf]

def J(near_data):
    return npa.abs(near_data[0,0])**2

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

x0 = 0.5*np.ones(Nx*Ny)
opt.update_design([x0])
f0, g_adjoint = opt()

db, choose = 1e-4, 5
np.random.seed(20211204)

g_discrete, idx = opt.calculate_fd_gradient(num_gradients=choose,db=db)
g_discrete = np.array(g_discrete).flatten()
g_adjoint = g_adjoint.flatten()

print("Randomly selected indices: ",idx)
print("Finite-difference gradient: ",g_discrete)
print("Adjoint gradient: ",g_adjoint[idx])

Before this PR, the results are

Finite-difference gradient:  [-41.65764109 -36.27053419  22.04881095  38.55970822 -29.69618643]
Adjoint gradient:  [-41.65784098 -36.27197005  22.0480579   38.55931979 -29.69604378]

After this PR, the results are

Finite-difference gradient:  [-29.95933836 -39.96501542  14.17651164  26.06499857 -21.2956607 ]
Adjoint gradient:  [-49.61541877 -62.96284535  26.39349287  45.95704617 -35.40822268]