Skip to content

Commit

Permalink
Merge pull request #21 from sbuddhiraju/fdfd_mf
Browse files Browse the repository at this point in the history
Multi-frequency FDFD
  • Loading branch information
twhughes authored Aug 7, 2020
2 parents 44aa696 + 89695de commit 46f78d5
Show file tree
Hide file tree
Showing 6 changed files with 919 additions and 1 deletion.
2 changes: 1 addition & 1 deletion ceviche/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
__version__ = '0.1.1'

from .fdtd import fdtd
from .fdfd import fdfd_ez, fdfd_hz
from .fdfd import fdfd_ez, fdfd_hz, fdfd_mf_ez
from .jacobians import jacobian

from . import viz
Expand Down
115 changes: 115 additions & 0 deletions ceviche/fdfd.py
Original file line number Diff line number Diff line change
Expand Up @@ -243,6 +243,121 @@ def _solve_fn(self, eps_vec, entries_a, indices_a, Mz_vec):

return Ex_vec, Ey_vec, Hz_vec


class fdfd_mf_ez(fdfd):
""" FDFD class for multifrequency linear Ez polarization. New variables:
omega_mod: angular frequency of modulation (rad/s)
delta: array of shape (Nfreq, Nx, Ny) containing pointwise modulation depth for each modulation harmonic (1,...,Nfreq)
phi: array of same shape as delta containing pointwise modulation phase for each modulation harmonic
Nsb: number of numerical sidebands to consider when solving for fields.
This is not the same as the number of modulation frequencies Nfreq. For physically meaningful results, Nsb >= Nfreq.
"""

def __init__(self, omega, dL, eps_r, omega_mod, delta, phi, Nsb, npml, bloch_phases=None):
super().__init__(omega, dL, eps_r, npml, bloch_phases=bloch_phases)
self.omega_mod = omega_mod
self.delta = delta
self.phi = phi
self.Nsb = Nsb

def solve(self, source_z):
""" Outward facing function (what gets called by user) that takes a source grid and returns the field components """
# flatten the permittivity and source grid
source_vec = self._grid_to_vec(source_z)
eps_vec = self._grid_to_vec(self.eps_r)
Nfreq = npa.shape(self.delta)[0]
delta_matrix = self.delta.reshape([Nfreq, npa.prod(self.shape)])
phi_matrix = self.phi.reshape([Nfreq, npa.prod(self.shape)])
# create the A matrix for this polarization
entries_a, indices_a = self._make_A(eps_vec, delta_matrix, phi_matrix)

# solve field componets usng A and the source
Fx_vec, Fy_vec, Fz_vec = self._solve_fn(eps_vec, entries_a, indices_a, source_vec)

# put all field components into a tuple, convert to grid shape and return them all
Fx = self._vec_to_grid(Fx_vec)
Fy = self._vec_to_grid(Fy_vec)
Fz = self._vec_to_grid(Fz_vec)

return Fx, Fy, Fz

def _make_A(self, eps_vec, delta_matrix, phi_matrix):
""" Builds the multi-frequency electromagnetic operator A in Ax = b """
M = 2*self.Nsb + 1
N = self.Nx * self.Ny
W = self.omega + npa.arange(-self.Nsb,self.Nsb+1)*self.omega_mod

C = sp.kron(sp.eye(M), - 1 / MU_0 * self.Dxf.dot(self.Dxb) - 1 / MU_0 * self.Dyf.dot(self.Dyb))
entries_c, indices_c = get_entries_indices(C)

# diagonal entries representing static refractive index
# this part is just a block diagonal version of the single frequency fdfd_ez
entries_diag = - EPSILON_0 * npa.kron(W**2, eps_vec)
indices_diag = npa.vstack((npa.arange(M*N), npa.arange(M*N)))

entries_a = npa.hstack((entries_diag, entries_c))
indices_a = npa.hstack((indices_diag, indices_c))

# off-diagonal entries representing dynamic modulation
# this part couples different frequencies due to modulation
# for a derivation of these entries, see Y. Shi, W. Shin, and S. Fan. Optica 3(11), 2016.
Nfreq = npa.shape(delta_matrix)[0]
for k in npa.arange(Nfreq):
# super-diagonal entries (note the +1j phase)
mod_p = - 0.5 * EPSILON_0 * delta_matrix[k,:] * npa.exp(1j*phi_matrix[k,:])
entries_p = npa.kron(W[:-k-1]**2, mod_p)
indices_p = npa.vstack((npa.arange((M-k-1)*N), npa.arange((k+1)*N, M*N)))
entries_a = npa.hstack((entries_p, entries_a))
indices_a = npa.hstack((indices_p,indices_a))
# sub-diagonal entries (note the -1j phase)
mod_m = - 0.5 * EPSILON_0 * delta_matrix[k,:] * npa.exp(-1j*phi_matrix[k,:])
entries_m = npa.kron(W[k+1:]**2, mod_m)
indices_m = npa.vstack((npa.arange((k+1)*N, M*N), npa.arange((M-k-1)*N)))
entries_a = npa.hstack((entries_m, entries_a))
indices_a = npa.hstack((indices_m,indices_a))

return entries_a, indices_a

def _solve_fn(self, eps_vec, entries_a, indices_a, Jz_vec):
""" Multi-frequency version of _solve_fn() defined in fdfd_ez """
M = 2*self.Nsb + 1
N = self.Nx * self.Ny
W = self.omega + npa.arange(-self.Nsb,self.Nsb+1)*self.omega_mod
P = sp.kron(sp.spdiags(W,[0],M,M), sp.eye(N))
entries_p, indices_p = get_entries_indices(P)
b_vec = 1j * sp_mult(entries_p,indices_p,Jz_vec)
Ez_vec = sp_solve(entries_a, indices_a, b_vec)
Hx_vec, Hy_vec = self._Ez_to_Hx_Hy(Ez_vec)
return Hx_vec, Hy_vec, Ez_vec

def _Ez_to_Hx(self, Ez_vec):
""" Multi-frequency version of _Ez_to_Hx() defined in fdfd """
M = 2*self.Nsb + 1
Winv = 1/(self.omega + npa.arange(-self.Nsb,self.Nsb+1)*self.omega_mod)
Dyb_mf = sp.kron(sp.spdiags(Winv,[0],M,M), self.Dyb)
entries_Dyb_mf, indices_Dyb_mf = get_entries_indices(Dyb_mf)
return -1 / 1j / MU_0 * sp_mult(entries_Dyb_mf, indices_Dyb_mf, Ez_vec)

def _Ez_to_Hy(self, Ez_vec):
""" Multi-frequency version of _Ez_to_Hy() defined in fdfd """
M = 2*self.Nsb + 1
Winv = 1/(self.omega + npa.arange(-self.Nsb,self.Nsb+1)*self.omega_mod)
Dxb_mf = sp.kron(sp.spdiags(Winv,[0],M,M), self.Dxb)
entries_Dxb_mf, indices_Dxb_mf = get_entries_indices(Dxb_mf)
return 1 / 1j / MU_0 * sp_mult(entries_Dxb_mf, indices_Dxb_mf, Ez_vec)

def _Ez_to_Hx_Hy(self, Ez_vec):
""" Multi-frequency version of _Ez_to_Hx_Hy() defined in fdfd """
Hx_vec = self._Ez_to_Hx(Ez_vec)
Hy_vec = self._Ez_to_Hy(Ez_vec)
return Hx_vec, Hy_vec

def _vec_to_grid(self, vec):
""" Multi-frequency version of _vec_to_grid() defined in fdfd """
# grid shape has Nx*Ny cells per frequency sideband
grid_shape = (2*self.Nsb + 1, self.Nx, self.Ny)
return npa.reshape(vec, grid_shape)

class fdfd_3d(fdfd):
""" 3D FDFD class (work in progress) """

Expand Down
333 changes: 333 additions & 0 deletions examples/fdfd_mf_intro.ipynb

Large diffs are not rendered by default.

333 changes: 333 additions & 0 deletions examples/multifrequency_fdfd.ipynb

Large diffs are not rendered by default.

1 change: 1 addition & 0 deletions tests/test_all_gradients.sh
Original file line number Diff line number Diff line change
Expand Up @@ -7,5 +7,6 @@ echo "testing all gradient checkers in $TEST_DIR"

# runs all of the gradient specific tests
python $TEST_DIR/test_gradients_fdfd.py
python $TEST_DIR/test_gradients_fdfd_mf.py
python $TEST_DIR/test_gradients_fdtd.py
python $TEST_DIR/test_primitives.py
136 changes: 136 additions & 0 deletions tests/test_gradients_fdfd_mf.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,136 @@
import unittest

import numpy as np
import autograd.numpy as npa
import scipy.sparse as sp
import scipy.sparse.linalg as spl
import copy

from autograd.extend import primitive, defvjp
from autograd import grad

import sys
sys.path.append('../ceviche')

from ceviche.utils import grad_num
from ceviche import jacobian, fdfd_mf_ez

"""
This file tests the autograd gradients of an FDFD and makes sure that they
equal the numerical derivatives
"""

# test parameters
ALLOWED_RATIO = 1e-4 # maximum allowed ratio of || grad_num - grad_auto || vs. || grad_num ||
VERBOSE = False # print out full gradients?
DEPS = 1e-6 # numerical gradient step size

print("Testing the Multi-frequency Linear FDFD Ez gradients")

class TestFDFD(unittest.TestCase):

""" Tests the flexible objective function specifier """

def setUp(self):

# basic simulation parameters
self.Nx = 30
self.Ny = 30
self.N = self.Nx * self.Ny
self.Nfreq = 1
self.Nsb = 1
self.omega = 2*np.pi*200e12
self.omega_mod = 2*np.pi*2e12
self.dL = 1e-6
self.pml = [10, 10]

# sources (chosen to have objectives around 1)
self.source_amp_ez = 1
self.source_amp_hz = 1

self.source_ez = np.zeros((2*self.Nsb+1, self.Nx, self.Ny))
self.source_ez[self.Nsb, self.Nx//2, self.Ny//2] = self.source_amp_ez

# starting relative permittivity (random for debugging)
self.eps_r = np.random.random((self.Nx, self.Ny)) + 1
self.delta = np.random.random((self.Nfreq, self.Nx, self.Ny))
self.phi = 2*np.pi*np.random.random((self.Nfreq, self.Nx, self.Ny))
self.eps_arr = self.eps_r.flatten()
self.params = npa.concatenate( (npa.concatenate((self.eps_arr, self.delta.flatten() )), self.phi.flatten() ) )
def check_gradient_error(self, grad_num, grad_auto):
""" Checks the test case:
compares the norm of the gradient to the norm of the difference
Throws error if this is greater than ALLOWED RATIO
"""
norm_grad = np.linalg.norm(grad_num)
print('\t\tnorm of gradient: ', norm_grad)
norm_diff = np.linalg.norm(grad_num - grad_auto)
print('\t\tnorm of difference: ', norm_diff)
norm_ratio = norm_diff / norm_grad
print('\t\tratio of norms: ', norm_ratio)
self.assertLessEqual(norm_ratio, ALLOWED_RATIO)
print('')

def test_Ez_reverse(self):

print('\ttesting reverse-mode Ez in FDFD_MF')
f = fdfd_mf_ez(self.omega, self.dL, self.eps_r, self.omega_mod, self.delta, self.phi, self.Nsb, self.pml)

def J_fdfd(params):
eps_r = params[:self.N].reshape((self.Nx, self.Ny))
delta = params[self.N:(self.Nfreq+1)*self.N].reshape((self.Nfreq, self.Nx, self.Ny))
phi = params[(self.Nfreq+1)*self.N:].reshape((self.Nfreq, self.Nx, self.Ny))
# set the permittivity, modulation depth, and modulation phase
f.eps_r = eps_r
f.delta = delta
f.phi = phi
# set the source amplitude to the permittivity at that point
Hx, Hy, Ez = f.solve((eps_r + npa.sum(delta*npa.exp(1j*phi),axis=0))* self.source_ez)

return npa.sum(npa.square(npa.abs(Ez))) \
+ npa.sum(npa.square(npa.abs(Hx))) \
+ npa.sum(npa.square(npa.abs(Hy)))

grad_autograd_rev = jacobian(J_fdfd, mode='reverse')(self.params)
grad_numerical = jacobian(J_fdfd, mode='numerical')(self.params)

if VERBOSE:
print('\ttesting epsilon, delta and phi.')
print('\tobjective function value: ', J_fdfd(self.params))
print('\tgrad (auto): \n\t\t', grad_autograd_rev)
print('\tgrad (num): \n\t\t', grad_numerical)

self.check_gradient_error(grad_numerical, grad_autograd_rev)

def test_Ez_forward(self):

print('\ttesting forward-mode Ez in FDFD_MF')

f = fdfd_mf_ez(self.omega, self.dL, self.eps_r, self.omega_mod, self.delta, self.phi, self.Nsb, self.pml)

def J_fdfd(c):

# set the permittivity, modulation depth, and modulation phase
f.eps_r = c * self.eps_r
f.delta = c * self.delta
f.phi = c * self.phi
# set the source amplitude to the permittivity at that point
Hx, Hy, Ez = f.solve(c * self.eps_r * self.source_ez)

return npa.square(npa.abs(Ez)) \
+ npa.square(npa.abs(Hx)) \
+ npa.square(npa.abs(Hy))

grad_autograd_for = jacobian(J_fdfd, mode='forward')(1.0)
grad_numerical = jacobian(J_fdfd, mode='numerical')(1.0)

if VERBOSE:
print('\tobjective function value: ', J_fdfd(1.0))
print('\tgrad (auto): \n\t\t', grad_autograd_for)
print('\tgrad (num): \n\t\t', grad_numerical)

self.check_gradient_error(grad_numerical, grad_autograd_for)


if __name__ == '__main__':
unittest.main()

0 comments on commit 46f78d5

Please sign in to comment.