Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

adjoint for ldos #2077

Merged
merged 11 commits into from
Jun 2, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
47 changes: 47 additions & 0 deletions python/adjoint/objective.py
Original file line number Diff line number Diff line change
Expand Up @@ -353,3 +353,50 @@ def __call__(self):
for far_pt in self.far_pts
]).reshape((self._nfar_pts, self.num_freq, 6))
return self._eval

class LDOS(ObjectiveQuantity):
def __init__(self, sim, decimation_factor=0, **kwargs):
super().__init__(sim)
self.decimation_factor = decimation_factor
self.srckwarg = kwargs

def register_monitors(self, frequencies):
self._frequencies = np.asarray(frequencies)
self.forward_src = self.sim.sources
return

def place_adjoint_source(self, dJ):
time_src = self._create_time_profile()
if dJ.ndim == 2:
dJ = np.sum(dJ, axis=1)
dJ = dJ.flatten()
sources = []
forward_f_scale = np.array([self.ldos_scale/self.ldos_Jdata[k] for k in range(self.num_freq)])
if self._frequencies.size == 1:
amp = (dJ * self._adj_src_scale(False) * forward_f_scale)[0]
src = time_src
else:
scale = dJ * self._adj_src_scale(False) * forward_f_scale
src = FilteredSource(
time_src.frequency,
self._frequencies,
scale,
self.sim.fields.dt,
)
amp = 1
for forward_src_i in self.forward_src:
if isinstance(forward_src_i, mp.EigenModeSource):
src_i = mp.EigenModeSource(src, component=forward_src_i.component, eig_kpoint = forward_src_i.eig_kpoint, amplitude=amp, eig_band=forward_src_i.eig_band, size = forward_src_i.size,center=forward_src_i.center, **self.srckwarg)
else:
src_i = mp.Source(src, component=forward_src_i.component, amplitude=amp, size = forward_src_i.size,center=forward_src_i.center, **self.srckwarg)
if mp.is_electric(src_i.component):
src_i.amplitude *= -1
sources += [src_i]

return sources

def __call__(self):
self._eval = self.sim.ldos_data
self.ldos_scale = self.sim.ldos_scale
self.ldos_Jdata = self.sim.ldos_Jdata
return np.array(self._eval)
41 changes: 25 additions & 16 deletions python/adjoint/optimization_problem.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
from autograd import grad, jacobian
from collections import namedtuple

from . import utils, DesignRegion
from . import utils, DesignRegion, LDOS

class OptimizationProblem(object):
"""Top-level class in the MEEP adjoint module.
Expand Down Expand Up @@ -175,11 +175,16 @@ def forward_run(self):
self.prepare_forward_run()

# Forward run
self.sim.run(until_after_sources=mp.stop_when_dft_decayed(
self.decay_by,
self.minimum_run_time,
self.maximum_run_time
))
if any(isinstance(m, LDOS) for m in self.objective_arguments):
self.sim.run(mp.dft_ldos(self.frequencies), until_after_sources=mp.stop_when_dft_decayed(
self.decay_by,
self.minimum_run_time,
self.maximum_run_time))
else:
self.sim.run(until_after_sources=mp.stop_when_dft_decayed(
self.decay_by,
self.minimum_run_time,
self.maximum_run_time))

# record objective quantities from user specified monitors
self.results_list = []
Expand Down Expand Up @@ -354,11 +359,13 @@ def calculate_fd_gradient(
self.forward_monitors.append(
m.register_monitors(self.frequencies))

self.sim.run(until_after_sources=mp.stop_when_dft_decayed(
self.decay_by,
self.minimum_run_time,
self.maximum_run_time
))
if any(isinstance(m, LDOS) for m in self.objective_arguments):
self.sim.run(mp.dft_ldos(self.frequencies), until_after_sources=mp.stop_when_energy_decayed(dt=1, decay_by=1e-11))
else:
self.sim.run(until_after_sources=mp.stop_when_dft_decayed(
self.decay_by,
self.minimum_run_time,
self.maximum_run_time))

# record final objective function value
results_list = []
Expand All @@ -385,11 +392,13 @@ def calculate_fd_gradient(
m.register_monitors(self.frequencies))

# add monitor used to track dft convergence
self.sim.run(until_after_sources=mp.stop_when_dft_decayed(
self.decay_by,
self.minimum_run_time,
self.maximum_run_time
))
if any(isinstance(m, LDOS) for m in self.objective_arguments):
self.sim.run(mp.dft_ldos(self.frequencies), until_after_sources=mp.stop_when_energy_decayed(dt=1, decay_by=1e-11))
else:
self.sim.run(until_after_sources=mp.stop_when_dft_decayed(
self.decay_by,
self.minimum_run_time,
self.maximum_run_time))

# record final objective function value
results_list = []
Expand Down
1 change: 1 addition & 0 deletions python/simulation.py
Original file line number Diff line number Diff line change
Expand Up @@ -5213,6 +5213,7 @@ def _ldos(sim, todo):
sim.ldos_data = mp._dft_ldos_ldos(ldos)
sim.ldos_Fdata = mp._dft_ldos_F(ldos)
sim.ldos_Jdata = mp._dft_ldos_J(ldos)
sim.ldos_scale = ldos.overall_scale()
if verbosity.meep > 0:
display_csv(sim, 'ldos', zip(mp.get_ldos_freqs(ldos), sim.ldos_data))
return _ldos
Expand Down
47 changes: 46 additions & 1 deletion python/tests/test_adjoint_solver.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
from enum import Enum
from utils import ApproxComparisonTestCase

MonitorObject = Enum('MonitorObject', 'EIGENMODE DFT')
MonitorObject = Enum('MonitorObject', 'EIGENMODE DFT LDOS')

resolution = 30

Expand Down Expand Up @@ -57,6 +57,11 @@
center=mp.Vector3(-0.5*sxy+dpml,0),
size=mp.Vector3(),
component=mp.Ez)]

line_source = [mp.Source(src=mp.GaussianSource(fcen,fwidth=df),
center=mp.Vector3(-0.85,0),
size=mp.Vector3(0,sxy-2*dpml),
component=mp.Ez)]

k_point = mp.Vector3(0.23,-0.38)

Expand All @@ -81,6 +86,12 @@ def forward_simulation(design_params, mon_type, frequencies=None, mat2=silicon):

if not frequencies:
frequencies = [fcen]

if mon_type.name == 'LDOS':
sim.change_sources(line_source)
sim.run(mp.dft_ldos(frequencies), until_after_sources=mp.stop_when_energy_decayed(dt=50.0, decay_by = 1e-11))
sim.reset_meep()
return np.array(sim.ldos_data)

if mon_type.name == 'EIGENMODE':
mode = sim.add_mode_monitor(frequencies,
Expand Down Expand Up @@ -159,6 +170,12 @@ def J(mode_mon):

def J(mode_mon):
return npa.power(npa.abs(mode_mon[:,4,10]),2)
elif mon_type.name == 'LDOS':
sim.change_sources(line_source)
obj_list = [mpa.LDOS(sim)]

def J(mode_mon):
return npa.array(mode_mon)

opt = mpa.OptimizationProblem(
simulation=sim,
Expand Down Expand Up @@ -410,6 +427,34 @@ def test_adjoint_solver_eigenmode(self):
print("Directional derivative -- adjoint solver: {}, FD: {}".format(adj_scale,fd_grad))
tol = 0.04 if mp.is_single_precision() else 0.01
self.assertClose(adj_scale,fd_grad,epsilon=tol)


def test_adjoint_solver_LDOS(self):
print("*** TESTING LDOS ADJOINT ***")

## test the single frequency and multi frequency cases
for frequencies in [[fcen], [1/1.58, fcen, 1/1.53]]:
## compute gradient using adjoint solver
adjsol_obj, adjsol_grad = adjoint_solver(p, MonitorObject.LDOS, frequencies)

## compute unperturbed |Ez|^2
Ez2_unperturbed = forward_simulation(p, MonitorObject.LDOS, frequencies)

## compare objective results
print("|Ez|^2 -- adjoint solver: {}, traditional simulation: {}".format(adjsol_obj,Ez2_unperturbed))
self.assertClose(adjsol_obj,Ez2_unperturbed,epsilon=2e-5)

## compute perturbed Ez2
Ez2_perturbed = forward_simulation(p+dp, MonitorObject.LDOS, frequencies)

## compare gradients
if adjsol_grad.ndim < 2:
adjsol_grad = np.expand_dims(adjsol_grad,axis=1)
adj_scale = (dp[None,:]@adjsol_grad).flatten()
fd_grad = Ez2_perturbed-Ez2_unperturbed
print("Directional derivative -- adjoint solver: {}, FD: {}".format(adj_scale,fd_grad))
tol = 0.07 if mp.is_single_precision() else 0.006
self.assertClose(adj_scale,fd_grad,epsilon=tol)


def test_gradient_backpropagation(self):
Expand Down
9 changes: 6 additions & 3 deletions src/dft_ldos.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ dft_ldos::dft_ldos(double freq_min, double freq_max, int Nfreq) {
for (int i = 0; i < Nfreq; ++i)
Fdft[i] = Jdft[i] = 0.0;
Jsum = 1.0;
saved_overall_scale = 1.0;
}

dft_ldos::dft_ldos(const std::vector<double> freq_) {
Expand All @@ -39,6 +40,7 @@ dft_ldos::dft_ldos(const std::vector<double> freq_) {
for (size_t i = 0; i < Nfreq; ++i)
Fdft[i] = Jdft[i] = 0.0;
Jsum = 1.0;
saved_overall_scale = 1.0;
}

dft_ldos::dft_ldos(const double *freq_, size_t Nfreq) : freq(Nfreq) {
Expand All @@ -49,27 +51,28 @@ dft_ldos::dft_ldos(const double *freq_, size_t Nfreq) : freq(Nfreq) {
for (size_t i = 0; i < Nfreq; ++i)
Fdft[i] = Jdft[i] = 0.0;
Jsum = 1.0;
saved_overall_scale = 1.0;
}

// |c|^2
static double abs2(complex<double> c) { return real(c) * real(c) + imag(c) * imag(c); }

double *dft_ldos::ldos() const {
double *dft_ldos::ldos() {
// we try to get the overall scale factor right (at least for a point source)
// so that we can compare against the analytical formula for testing
// ... in most practical cases, the scale factor won't matter because
// the user will compute the relative LDOS of 2 cases (e.g. LDOS/vacuum)

// overall scale factor
double Jsum_all = sum_to_all(Jsum);
double scale = 4.0 / pi // from definition of LDOS comparison to power
saved_overall_scale = 4.0 / pi // from definition of LDOS comparison to power
* -0.5 // power = -1/2 Re[E* J]
/ (Jsum_all * Jsum_all); // normalize to unit-integral current

const size_t Nfreq = freq.size();
double *sum = new double[Nfreq];
for (size_t i = 0; i < Nfreq; ++i) /* 4/pi * work done by unit dipole */
sum[i] = scale * real(Fdft[i] * conj(Jdft[i])) / abs2(Jdft[i]);
sum[i] = saved_overall_scale * real(Fdft[i] * conj(Jdft[i])) / abs2(Jdft[i]);
double *out = new double[Nfreq];
sum_to_all(sum, out, Nfreq);
delete[] sum;
Expand Down
6 changes: 4 additions & 2 deletions src/meep.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -1412,15 +1412,17 @@ class dft_ldos {
}

void update(fields &f); // to be called after each timestep
double *ldos() const; // returns array of Nomega values (after last timestep)
double *ldos(); // returns array of Nomega values (after last timestep)
std::complex<double> *F() const; // returns Fdft
std::complex<double> *J() const; // returns Jdft
std::vector<double> freq;

double overall_scale() const { return saved_overall_scale; }

private:
std::complex<double> *Fdft; // Nomega array of field * J*(x) DFT values
std::complex<double> *Jdft; // Nomega array of J(t) DFT values
double Jsum; // sum of |J| over all points
double saved_overall_scale; // saved overall scale for adjoint calculation
};

// dft.cpp (normally created with fields::add_dft_fields)
Expand Down