legume: Differentiable guided mode expansion methods

legume (le GUided Mode Expansion) is an open source Python package that implements a differentiable guided mode expansion (GME) method for multi-layer optical structures. Legume also implements a differentiable plane wave expansion (PWE) method for purely two-dimensional periodic structures. Legume uses the HIPS autograd package for its automatic differentiation capabilities.

Why legume?

legume is a python implementation of the guided-mode expansion (GME) method for simulating photonic crystal slabs, i.e. for multi-layer structures that look something like this

This is an extremely efficient method to obtain the photonic Bloch bands of such structures and can be used to study bulk photonic crystals as well as devices like waveguides and cavities.

The guided-mode expansion is particularly useful for computing the quasi-guided modes above the light line, which are hard to isolate in first-principle finite-difference of finite-element methods. This can be invaluable for the study of the coupling of photonic crystal modes to the radiative environment, and of exotic phenomena like bound states in the continuum.

The GME method can be super useful in itself, but on top of that we also provide a differentiable implementation through the autograd backend. This allows you to compute the gradient of an objective function that depends on any of the ouput parameters (eigenmode frequencies and losses, fields…) with respect to any of the input parameters (hole positions, sizes and shapes, slab thickness…) With this powerful addition, there’s no end to the possibilities in using legume to optimize your devices!

Dive into the Examples to see how this all works in practice!

Installation

Required dependencies

• Python (>= 3.6)
• numpy (>= 1.16)
• scipy (>= 1.2.1)
• matplotlib (>= 3.0.3)

Optional dependencies

• autograd (>= 1.2): For computing gradients
• gdspy (>= 1.5): For GDS structure export
• scikit-image (>= 0.15): For GDS structure export via rasterization

Instructions

To install the latest version of Legume from PyPi:

pip install legume-gme

Alternatively, you can git clone Legume from the Fan group GitHub, manually install all of the required dependencies, and add the path to the Legume in your python path environment variable:

git clone https://github.com/fancompute/legume.git
export PYTHONPATH=$PYTHONPATH:/path/to/the/location/of/legume

Frequently Asked Questions

What do I do to test convergence?

The best way to make sure that your GME computation is converged is to increase the parameters controlling the precision of the simulation until you no longer see change in the eigenmodes of interest. We recommend doing this in the following order:

• First, make sure you have set a high enough gmax, which is defined upon initialization of GuidedModeExp.
• Then, increase the number of guided bands included in the simulation by adding more indexes to the gmode_inds list supplied to GuidedModeExp.run(). Note that after including more modes in gmode_inds, you should test again the convergence w.r.t. gmax.
• If your bands look particularly weird and discontinuous, there might be an issue in the computation of the guided modes of the effective homogeneous structure (the expansion basis). Try decreasing gmode_step supplied in GuidedModeExp.run() to 1e-3 or 1e-4 and see if things look better.

Finally, note that GME is only an approximate method. So, even if the simulation is converged with respect to all of the above parameters but still produces strange results, it might just be that the method is not that well-suited for the structure you are simulating. We’re hoping to improve that in future version of legume!

Why am I running out of memory?

GME requrest the diagonalization of dense matrices, and you might start running out of memory for simulations in which computational time is not that much of an issue. This is also because the scaling with gmax is pretty bad: the linear dimension of the matrix for diagonalization scales as gmax**2, and so the total memory needed to store it scales as gmax**4. So, unfortunately, if you’re running out of memory in a simulation there’s not much you can do but decrease gmax.

That said, if you’re running out of memory in a gradient computation, there could be something you can try. Reverse-mode autodiff is generally the best approach for optimization problems in terms of computational time, but this can sometimes come at a memory cost. This is because all of the intermediate values of the forward simulation have to be stored for the backward pass. So, if you are for example doing a loop through different k-points, the dense matrices and their eigenvectors at every k will be stored, which can add up to a lot. There is no easy way to fix this (and no direct way within autograd), but we’ve included a function that can provide a workaround. For details, have a look at this example.

Finally, it’s worth mentioning that there are probably improvements that can be made to the memory usage. If anybody wants to dive deep in the code and try to do that, it will be appreciated!

What should I know about the guided-mode basis?

The expansion basis in the GME consists of the guided modes of an effective homogeneous structure (panels (a)-(b)) in the Figure. By default, the effective permittivities in (b) are taken as the average value in every layer. This is controlled by the gmode_eps keyword option in the run options. Setting gmode_eps = 'background' will take the background permittivity instead, while there’s also the option to have custom values by setting gmode_eps = 'custom'. In that case, every layer (including the claddings) in the PhotCryst object should have a pre-defined effective permittivity eps_eff, which will be used in the guided-mode computation. This is simply set as an attribute of the layer, e.g.

phc.layers[0].eps_eff = 10   # Slab custom effective epsilon
phc.calddings[0].eps_eff = 1 # Lower cladding
phc.claddings[1].eps_eff = 5 # Upper cladding

The guided modes can be classified as TE/TM, where in our notation the reference plane is the slab plane (xy). The guided modes alternate between TE and TM, such that gmode_inds = [0, 2, 4, ...] are TE and gmode_inds = [1, 3, 5, ...] are TM (panel (c)). However, this classification is often broken by the photonic crystal structure (we discuss symmetries further below).

We only include the fully-guided modes in the computation (the ones that lie below both light lines in (c)). This is what makes the computation approximate, as the basis set is not complete.

How do I incorporate symmetry?

The TE/TM classification of the guided modes of the homogeneous structure is often broekn by the photonic crystal permittivity. Here is how you can still incorporate some structural symmetries.

For gratings (permittivity is periodic in one direction and homogeneous in the other), the TE/TM classification holds. You can selectively compute the modes by supplying gmode_inds with either only even or only odd numbers.

For photonic crystals with a mirror plane, like a single slab with symmetric claddings, the correct classification of modes is with respect to reflection in that plane. The positive-symmetry guided modes are gmode_inds = [0, 3, 4, 7, 8, ...], while the negative-symmetry modes are gmode_inds = [1, 2, 5, 6, 9, 10, ...]. Low-frequency positive-symmetry modes that are mostly fromed by the gmode_inds = 0 guided band are sometimes referred to as quasi-TE, and low-frequency negative-symmetry modes that are mostly formed by the gmode_inds = 1 guided band are sometimes referred to as quasi-TM.

Without any mirror planes, all the guided modes are generally mixed. There can still be symmetry if the k-vector points in a high-symmetry direction, but there is currently no way to take advantage of that in legume.

When should I use approximate gradients?

When running GME with the autograd backend, one of the run() options you can specify is 'gradients' = {'exact' (default), 'approx'}. The approximate option could be faster in some cases, and could actually still be exact in some cases. This is the high-level computational graph of the guided-mode expansion:

The 'approx' option discards the gradient due to the top path in this graph, i.e. the gradient due to the changing basis. Only the gradient from the diagonalization path is included. Here are some rules of thumb on what to use:

• If you’re optimizing hole positions, or more generally parameters that don’t change the average permittivity, you’re in luck! In this case, the 'approx' gradients should actually be exact!
• If you’re optimizing dispersion (real part of eigenfrequencies), you could try using 'approx' gradients, as they might be within just a few percent of the exact ones.
• If you’re optimizing loss rates or field profiles and/or if your parameters include the layer thicknesses, then the 'approx' gradients could be significantly off, 'exact' is recommended (and is the default).

What if I only need the Q of some of the modes?

In some simulations, the computation of the radiative losses could be the time bottleneck. In some cases, e.g. when optimizing a cavity, you only need to compute the quality factor of a single mode. If you run the GME by default, the Q-s of all modes will be computed instead, but you can set the option compute_im = False to avoid this. Running the GME with this option will compute all modes, but not the imaginary part of their frequencies (which is done perturbatively after the first stage of the computation). Then, you can use the legume.GuidedModeExp.compute_rad() method to only compute the loss rates of selected modes.

What’s the gauge?

Something to be aware of is the fact that the eigenmodes come with an arbitrary k-dependent gauge, as is usually the case for eigenvalue simulations. That is to say, each eigenvector is defined only up to a global phase, and this phase might change discontinously even for nearby k-points. If you re looking into something that depends on the gauge choice, you will have to figure out how to set your preferred gauge yourself.

Of course, apart from this global phase, all the relative phases should be well-defined (as they correspond to physically observable quantities). So for example if you compute radiative couplings to S and P polarization, the relative phase between the two should be physical.

How can I learn more about the method?

Our paper gives a lot of detail both on the guided-mode expansion method and on our differentiable implementation.

How should I cite legume?

If you find legume useful for your research, we would apprecite you citing our paper. For your convenience, you can use the following BibTex entry:

@article{Minkov2020,
title = {Inverse design of photonic crystals through automatic differentiation},
author = {Minkov, Momchil and Williamson, Ian A. D. and Gerace, Dario and Andreani, Lucio C. and Lou, Beicheng and Song, Alex Y. and Hughes, Tyler W. and Fan, Shanhui},
year = {2020},
journal = {arXiv:2003.00379},
}

Who made that awesome legume logo?

The legume logo was designed by Nadine Gilmer. She is also behind the logos for our angler and ceviche packages.

Examples

Shapes, layers, and photonic crystals

import numpy as np
import legume

%load_ext autoreload
%autoreload 2

The autoreload extension is already loaded. To reload it, use:
  %reload_ext autoreload

Shapes and layers

Shapes and layers are the basic building blocks of a simulation in legume. A layer can contain any number of shapes. A layer is the fundamental object for the plane-wave expansion, while a photonic crystal made up of a number of layers is the fundamental object for the guided-mode expansion. Initializing a ShapesLayer class requires one positional argument lattice, which is an instance of the Lattice object. Let’s start by creating a lattice, a layer, and some shapes.

# Initialize a lattice (can be 'square', 'hexagonal', or defined by primitive vectors)
lattice = legume.Lattice('hexagonal')

# Initialize a layer with background permittivity 2
layer = legume.ShapesLayer(lattice, eps_b=2)

# Create a square and use the `add_shape` method to add it to the layer
square = legume.Square(eps=10, x_cent=0, y_cent=0, a=0.3)
layer.add_shape(square)

# Create a circle and also add it to the layer
circle = legume.Circle(eps=6, x_cent=0.5, y_cent=0.3, r=0.2)
layer.add_shape(circle)

# Use an in-built visualization method to plot the contours of the shapes we have so far
legume.viz.shapes(layer)

Notice how the shapes are repeated in a hexagonal pattern, following the way we initialized our lattice. The above plot shows the shapes, but not the permittivity distribution. To see that, we can use another built-in method.

legume.viz.eps(layer, cbar=True)

Plane-wave expansion

The layer we consturcted can now be used to initialize a plane-wave expansion simulation. This also takes two keyword arugments:

• gmax : Maximum reciprocal lattice wave-vector length in units of \(2\pi/a\)
• eps_eff : Effective background epsilon; if None, take layer.eps_b. Setting this to a different value can be useful for simulations of an effective slab of a given thickness and at a particular frequency.

pwe = legume.PlaneWaveExp(layer, gmax=3)

We can plot the permittivity as “seen” by the PWE with a given gmax; in other words, we can plot the real-space permittivity from the Fourier components that will be used in the expansion. Notice how this gets closer to the real structure as gmax increases.

legume.viz.eps_ft(pwe)

pwe = legume.PlaneWaveExp(layer, gmax=6)
legume.viz.eps_ft(pwe)

We will not run the simulation here; this tutorial is only oriented towards building the simulation structures.

Photonic crystals

We can build photonic crystals made of one or more layers, and add various shapes. Let’s start with an example.

# Initialize a lattice (can be 'square', 'hexagonal', or defined by primitive vectors)
lattice = legume.Lattice('hexagonal')

# Initialize a PhC (optional kwargs are the permittivity in the lower and upper cladding)
phc = legume.PhotCryst(lattice, eps_l=1., eps_u=1.)

# Add a layer to the PhC with thickness 1 and background permittivity 10
phc.add_layer(d=1, eps_b=10)

# We can add shapes to the layer in two different ways:

# Create a square and use the layer `add_shape` method
square = legume.Square(x_cent=0, y_cent=0, a=0.3, eps=2)

# Create a circle and use the phc `add_shape` method
circle = legume.Circle(eps=6, x_cent=0.5, y_cent=0.3, r=0.2)
phc.add_shape([circle, square]) # by default added to the last layer; can add a list of shapes

# We can plot an overview of what we've built so far
legume.viz.structure(phc)

Guided-mode expansion

Just like PWE, GME uses the Fourier transform (FT) of the shapes. The GuidedModeExp object also takes gmax as a keyword argument, and we can also plot the FT of the photonic crystal that we’re simulating.

gme = legume.GuidedModeExp(phc, gmax=5)
legume.viz.eps_ft(gme)

Multiple layers

We can also create and simulate a multi-layer structure!

# Add another layer to our existing photonic crystal
phc.add_layer(d=0.4, eps_b=8)

circle = legume.Circle(x_cent=0.1, r=0.2, eps=2)
phc.add_shape(circle) # always adds to the last layer

# Visualize the structure and compare to the FT
legume.viz.structure(phc, Ny=300)
gme = legume.GuidedModeExp(phc, gmax=5)
legume.viz.eps_ft(gme, Ny=300)

Shapes in the claddings

You can add shapes to the claddings too!

triangle = legume.Poly(x_edges=[-0.1, 0.5, 0.2], y_edges=[0, 0, 0.5], eps=4)
phc.add_shape(triangle, cladding=1) # you can also say cladding = 'l', or phc.cladding[0].add_shape(triangle)

# Plot the real-space overview
legume.viz.structure(phc, Ny=300, cladding=True)
gme = legume.GuidedModeExp(phc, gmax=5)
legume.viz.eps_ft(gme, Ny=300, cladding=True)

Caveats

There are two important things to keep in mind: - Polygons should be defined such that the points circle their interior in a counter-clockwise manner. But don’t worry, if you don’t do that you’ll get a warning telling you to do it.

• Shapes should not overlap. This might be fixed in the future, but right now here’s what happens if they do:

phc = legume.PhotCryst(lattice)
phc.add_layer(d=1, eps_b=12)
circle = legume.Circle(x_cent=0.1, r=0.2, eps=5)
square = legume.Square(a=0.3)
phc.add_shape([circle, square])

# Notice the discrepancy between the two plots. In the FT, the permittivity of the overlapping shapes is wrong -
# and so the simulation will be wrong!
legume.viz.structure(phc, Ny=300)
gme = legume.GuidedModeExp(phc, gmax=5)
legume.viz.eps_ft(gme, Ny=300)

Plane wave expansion for a 2D photonic crystal

import numpy as np
import matplotlib.pyplot as plt

from legume import PlaneWaveExp, Circle, ShapesLayer, Lattice, viz

%load_ext autoreload
%autoreload 2

The autoreload extension is already loaded. To reload it, use:
  %reload_ext autoreload

Define 2D photonic crystal

In this notebook we will reproduce the modes for one of the structures discussed in Chapter 5 from “Photonic Crystals: Molding the Flow of Light” - the photonic crystal bible by John D. Joannopoulos, Steven Johnson, Joshua Winn and Robert Meade.

# Define parameters as for Chapter 5, Fig. 2 from Molding the Flow of Light
ra = 0.2      # cylinder radius
eps_c = 8.9   # cylinder permittivity

# Initialize lattice
lattice = Lattice('square')
# Initialize layer
layer = ShapesLayer(lattice)
# Add a cylinder to the layer
layer.add_shape(Circle(r=ra, eps=eps_c))
# Visualize the structure
viz.eps(layer, cbar=True)

# Initialize the BZ bath as in Chapter 5, Fig. 2
path = lattice.bz_path(['G', 'X', 'M', 'G'], [20])

# Initialize the plane-wave expansion and visualize the FT of the structure
pwe = PlaneWaveExp(layer, gmax=5)
viz.eps_ft(pwe, figsize = (4, 3))

Compute the photonic bands

# Run the plane-wave expansion for the two separate polarizations and store the bands
pwe.run(kpoints=path['kpoints'], pol='te')
freqs_te = pwe.freqs
pwe.run(kpoints=path['kpoints'], pol='tm')
freqs_tm = pwe.freqs

# Plot the results
fig, ax = plt.subplots(1, constrained_layout=True, figsize=(5, 4))
plt.plot(freqs_te, 'r')
plt.plot(freqs_tm, 'b')
ax.set_ylim([0, 0.8])
ax.set_xlim([0, pwe.freqs.shape[0]-1])
ax.set_ylabel("$\omega a/2\pi c$")

# The `path` dict provides some useful functionality for labeling of the BZ path
plt.xticks(path['indexes'], path['labels'])
ax.xaxis.grid('True')
plt.show()

Compare this to Chapter 5, Fig. 2 from “Molding the Flow of Light”

Visuzlize the field of a mode

We can also plot the fields and compare to Fig. 3 of the book. Note that all fields are always continuous in wave-expansion methods, but the discontinuities are captured better and better with increasing the truncation of the Fourier basis (here by increasing gmax).

# Compare to Chapter 5, Fig. 3, middle row
viz.field(pwe, field='d', kind=20, mind=0,
            component='z', val='im', N1=100, N2=100, cbar=True, eps=True);
viz.field(pwe, field='d', kind=20, mind=1,
            component='z', val='re', N1=100, N2=100, cbar=True, eps=True);

From the book:

Guided mode expansion for a multi-layer grating

import numpy as np
import legume

import matplotlib.pyplot as plt

%load_ext autoreload
%autoreload 2

The autoreload extension is already loaded. To reload it, use:
  %reload_ext autoreload

Asymmetric grating

In this notebook we will reproduce some of the results from the paper “Perfect single-sided radiation and absorption without mirrors”, Zhou et al., Optica 3 (2016). We first initialize an asymmetric grating as shown in Fig. 2(a) of the paper.

# Grating parameters as for Fig. 3 (b)
ymax = 0.1      # ficticious supercell length in y-direction, should be smaller than 1/gmax below
W = 0.45        # width of dielectric rods
H = 1.5         # total height of grating
D = 0.1         # thickness of added parts
Wa = (1-W)/2    # width of added parts
epss = 1.45**2  # permittivity of the rods
epsa = 1.1**2   # permittivity of the added parts

Build the grating as a three-layer structure

The grating can be initialized as a three-layer photonic crystal. We take a rectangular lattice with a small periodicity in the y-direction. If ymax < 1/gmax, then only reciprocal lattice vectors with \(G_y = 0\) will be included in the expansion, as should be for a grating.

# Initialize the lattice and the PhC
lattice = legume.Lattice([1, 0], [0, ymax])
phc = legume.PhotCryst(lattice)

# First layer
phc.add_layer(d=D, eps_b=epss)
rect_add = legume.Poly(eps=epsa, x_edges=np.array([-0.5, -0.5, -W / 2, -W / 2]),
                     y_edges=np.array([0.5, -0.5, -0.5, 0.5]) * ymax)
rect_air = legume.Poly(eps=1, x_edges=np.array([W / 2, W / 2, 0.5, 0.5]),
                     y_edges=np.array([0.5, -0.5, -0.5, 0.5]) * ymax)
phc.add_shape([rect_add, rect_air])

# Second layer
phc.add_layer(d=H-2*D, eps_b=epss)
rect_air1 = legume.Poly(eps=1, x_edges=np.array([-0.5, -0.5, -W / 2, -W / 2]),
                     y_edges=np.array([0.5, -0.5, -0.5, 0.5]) * ymax)
rect_air2 = legume.Poly(eps=1, x_edges=np.array([W / 2, W / 2, 0.5, 0.5]),
                     y_edges=np.array([0.5, -0.5, -0.5, 0.5]) * ymax)
phc.add_shape([rect_air1, rect_air2])

# Third layer
phc.add_layer(d=D, eps_b=epss)
rect_air = legume.Poly(eps=1, x_edges=np.array([-0.5, -0.5, -W / 2, -W / 2]),
                     y_edges=np.array([0.5, -0.5, -0.5, 0.5]) * ymax)
rect_add = legume.Poly(eps=epsa, x_edges=np.array([W / 2, W / 2, 0.5, 0.5]),
                     y_edges=np.array([0.5, -0.5, -0.5, 0.5]) * ymax)
phc.add_shape([rect_add, rect_air])

# Visualize what we built
legume.viz.structure(phc)

Compute quasi-guided bands

We first compute the bnads in the entire Brillouin zone. In this grating, the symmetry w.r.t. the kz plane is always satisfied, and the TE/TM classification is valid. However, the definition is ambiguous. In our notation, relevant for 2D PhC-s, the classification is w.r.t. the xy-plane. In Zhou et al., the classification is w.r.t. the kz plane. Thus, what they call TE, we call TM. To put it more explicitly, we will simulate modes where \(H_y\) is the only nonzero magnetic field component.

# Make a BZ path along the G-X direction
path = phc.lattice.bz_path(['G', np.array([np.pi, 0])], [30])

neig = 7 # number of Bloch bands to store
gmax = 4 # truncation of reciprocal lattice vectors

# Initialize GME
gme = legume.GuidedModeExp(phc, gmax=gmax)

# Set some of the running options
options = {'gmode_inds': [1, 3, 5, 7], # Take only the modes with H in the xy-plane
           'numeig': neig,
           'verbose': False
            }

# Run the simulation
gme.run(kpoints=path['kpoints'], **options)

# Visualize the bands
ax = legume.viz.bands(gme, Q=True)

Asymmetric coupling

Next, as they do in the paper, we focus on a region close to \(\Gamma\), and on the first quasi-guided band inside the lighg-cone (the one starting at frequency ~ 0.84 and going down). The frequency of this band is shown in Fig. S4 in the Supplementary of the paper, while the quality factor is shown in Fig. 3(b).

# Make a path that matches the axis extent in the paper, and put 100 k-points to capture the quasi-BIC well
path = phc.lattice.bz_path(['G', 0.15*np.array([2*np.pi, 0])], [100])

neig = 2 # number of Bloch bands to store
gmax = 6 # truncation of reciprocal lattice vectors

# Initialize GME
gme = legume.GuidedModeExp(phc, gmax=gmax)

# Set some of the running options
options = {'gmode_inds': [1, 3, 5, 7], # Take only the modes with H in the xy-plane
           'numeig': neig,
           'verbose': False
            }

# Run the simulation
gme.run(kpoints=path['kpoints'], **options)

# Plot the frequency and the quality factor
fig, ax = plt.subplots(1, 2, figsize=(8, 3))
indmode = 1
ax[0].plot(gme.kpoints[0, :]/2/np.pi, gme.freqs[:, indmode])
ax[0].set_xlabel('Wave vector')
ax[0].set_xlabel('Frequency')
ax[1].plot(gme.freqs[:, indmode]/2/gme.freqs_im[:, indmode], gme.freqs[:, indmode])
ax[1].set_xscale('log')
ax[1].set_xlabel('Q');

Finally, we reproduce Fig. 3(b) from the paper. Note that the legume GME method computes the coupling constant of each Bloch mode to plane waves outgoing in the lower and in the upper cladding, for TE- and TM-polarized outgoing light (equvalently, S and P polarization, respectively). These are stored in the dictionary gme.rad_coup, with keys 'l_te', 'l_tm', 'u_te', and 'u_tm'. Each item is a list of Nk lists of numeig elements, giving the coupling constants at each k for each mode. Finally, the elements are numpy arrays with as many elements as there are allowed diffraction orders, with the corresponding reciprocal lattice vectors stored in gme.rad_gvec.

Nk = gme.kpoints[0, :].size
# Coupling constant to lower-cladding TM modes
dl = np.array([gme.rad_coup['l_tm'][ik][1] for ik in range(Nk)])
# Coupling constant to upper-cladding TM modes
du = np.array([gme.rad_coup['u_tm'][ik][1] for ik in range(Nk)])
# Note: the coupling to TE modes is zero because of the kz symmetry

fig, ax = plt.subplots(1)
# Plot the asymmetric coupling ratio
color = 'tab:blue'
ax.plot(gme.kpoints[0, :]/2/np.pi, np.abs(dl/du)**2)
ax.set_yscale('log')
ax.set_ylabel('Asymmetric Coupling Ratio', color=color)
ax.tick_params(axis='y', labelcolor=color)

# Plot the quality factor
ax2 = ax.twinx()  # instantiate a second axes that shares the same x-axis
color = 'tab:red'
ax2.plot(gme.kpoints[0, :]/2/np.pi, gme.freqs[:, indmode]/2/gme.freqs_im[:, indmode], color=color)
ax2.set_yscale('log')
ax2.set_ylabel('Q', color=color)
ax2.tick_params(axis='y', labelcolor=color)

The agreement with Fig. 3(b) of Zhou et al., Optica (2016) is not perfect: for example, the BIC as computed by legume occurs slightly closer to the \(\Gamma\) point. However, the agreement is still quite good, and the important physics is captured.

Bands of a photonic crystal slab

In this short example, we will demonstrate how Legume can be used to compute the bands of a photonic crystal slab. We will consider a slab design that is very close to one from the book Molding the Flow of Light.

First we will import the legume package:

import legume

We will then define the parameters of our structure: \(D\) is the thickness of the slab, \(r\) is the radius of the air hole, and epsr is the permittivity.

D = 0.6
r = 0.3
epsr = 12.0

Next, we construct the lattice, the photonic crystal object, and the GME object:

lattice = legume.Lattice('hexagonal')
phc = legume.PhotCryst(lattice)
phc.add_layer(d=D, eps_b=epsr)
phc.layers[-1].add_shape(legume.Circle(eps=1.0, r=r))
gme = legume.GuidedModeExp(phc, gmax=6)

We can then use the legume.viz module to visualize our photonic crystal slab across three planes through the unit cell:

legume.viz.structure(phc)

Finally, we can solve for the bands of the photonic crystal. Note that here we are solving for only the modes with TE-like polarization because we have gmode_inds=[0].

path = lattice.bz_path(['G', 'M', 'K', 'G'], [15, 10, 15])
gme.run(kpoints=path['kpoints'],
        gmode_inds=[0, 3],
        numeig=20,
        verbose=False)
fig, ax = plt.subplots(1, figsize = (7, 5))
legume.viz.bands(gme, figsize=(5,5), Q=True, ax=ax)
ax.set_xticks(path['indexes'])
ax.set_xticklabels(path['labels'])
ax.xaxis.grid('True')

We observe excellent agreement with the bands shown in the text book (red lines).

Plane wave expansion with Autograd

import numpy as np
import matplotlib.pyplot as plt
import time

import autograd.numpy as npa
from autograd import grad, value_and_grad

import legume
from legume import PlaneWaveExp, Circle, ShapesLayer, Lattice
from legume.minimize import Minimize

%load_ext autoreload
%autoreload 2

The autoreload extension is already loaded. To reload it, use:
  %reload_ext autoreload

Define 2D photonic crystal waveguide

Here, we will optimize the dispersion of a 2D PhC waveguide. The underlying photonic crystal is similar to that in Chapter 5, Fig. 2 of “Photonic crystals: molding the flow of light”, and has a broad band gap for TM-modes. The waveguide is introduced as a missing row of pillars. To have a sufficient number of free parameters, we take a supercell of size \(5a\) in the propagation direction.

# PhC parameters
ra = 0.2     # pillar radius
eps_b = 1    # background permittivity
eps_c = 9    # pillar permittivity

# plane-wave expansion parameters
gmax = 2     # truncation of the plane-wave basis
Ny = 14      # Number of rows in the y-direction
Ny_opt = 3   # Number of rows in which the pillars will be modified
Nx = 5       # Supercell size in the x-direction

# Initialize a rectangular lattice
lattice = Lattice([Nx, 0], [0, Ny*np.sqrt(3)/2])

def wg_sc(dx, dy, dr):
    """Define the photonic crystal waveguide given shift parameters
    dx, dy, and dr, for the 2*Nx*Ny_opt number of pillars that get shifted
    """
    # Initialize a layer and the positions of the pillars for the regular waveguide
    layer_wg = ShapesLayer(lattice, eps_b=eps_b)
    xc = []; yc = []
    for ih in range(Ny):
        if ih != Ny//2:
            for ix in range(-Nx//2+1, Nx//2+1):
                xc.append((ih%2)*0.5 + ix)
                yc.append((-Ny//2 + ih)*np.sqrt(3)/2)

    # Add all the pillars, taking care of the shifts
    for ih in range(1, Ny//2+1):
        nx1 = (Ny//2+ih-1)
        nx2 = (Ny//2-ih)
        if ih <= Ny_opt:
            # The ih row includes "optimization" pillars
            for ix in range(Nx):
                circ = Circle(x_cent=xc[nx1*Nx + ix] + dx[(ih-1)*Nx + ix],
                              y_cent=yc[nx1*Nx + ix] + dy[(ih-1)*Nx  + ix],
                              r = ra + dr[(ih-1)*Nx + ix], eps=eps_c)
                layer_wg.add_shape(circ)
                circ = Circle(x_cent=xc[nx2*Nx + ix] + dx[(ih-1+Ny_opt)*Nx  + ix],
                              y_cent=yc[nx2*Nx + ix] + dy[(ih-1+Ny_opt)*Nx  + ix],
                              r = ra + dr[(ih-1+Ny_opt)*Nx + ix], eps=eps_c)
                layer_wg.add_shape(circ)
        else:
            # The ih row includes just regular pillars
            for ix in range(Nx):
                circ = Circle(x_cent = xc[nx2*Nx + ix], y_cent=yc[nx2*Nx + ix], r=ra, eps=eps_c)
                layer_wg.add_shape(circ)
                if ih < Ny//2:
                    circ = Circle(x_cent = xc[nx1*Nx + ix], y_cent=yc[nx1*Nx + ix], r=ra, eps=eps_c)
                    layer_wg.add_shape(circ)

    # Construct and return a plane-wave expansion object
    return PlaneWaveExp(layer_wg, gmax=gmax)

We first simulate the starting (un-modified) waveguide. Notice that the Brillouin zone is folded five times in the \(x\)-direction. There is just a single guided band in the band-gap if the elementary cell is used.

# Initialize zero shifts
dx0 = np.zeros((Nx*2*Ny_opt, ))
dy0 = np.zeros((Nx*2*Ny_opt, ))
dr0 = np.zeros((Nx*2*Ny_opt, ))

# Initialize the PWE and visualize the structure both through the `eps` and the `eps_ft` methods
pwe0 = wg_sc(dx0, dy0, dr0)
legume.viz.eps(pwe0.layer, Nx=150, Ny=250, cbar=True)
legume.viz.eps_ft(pwe0, Nx=150, Ny=250, figsize = (2.5, 4))

# Define a BZ path in kx and run the PWE
path = pwe0.layer.lattice.bz_path([[0, 0], np.array([np.pi/Nx, 0])], [20])
pwe0.run(kpoints=path['kpoints'], pol='tm', numeig = 150)

# Plot the (folded) waveguide bands
fig, ax = plt.subplots(1, constrained_layout=True, figsize=(4, 4))
ax.plot(pwe0.kpoints[0, :]*5/np.pi, pwe0.freqs, 'b')
ax.set_ylim([0.25, 0.55])
ax.set_xlim([0, pwe0.kpoints[0, -1]*5/np.pi])
ax.set_xlabel('$5a k_x/ \pi$')
ax.set_ylabel('Frequency')

# Plot the field of one of the waveguide modes using the in-built `legume` function
legume.viz.field(pwe0, field='e', kind=10, mind=Nx*(Ny-1) + 3, component='z', val='abs');

Compute gradients

We can use legume and autograd to efficiently compute gradients of PWE simulations. We will define the objective function as the mean-squared error between one of the waveguided bands and a target band shape. Let’s first illustrate this target.

# Define a target band shape and which band we will try to optimize
kpts = path['kpoints'][0, :]
f_target = -0.01*np.cos(kpts*Nx) # Just a cosine target
b_target = Nx*Ny - 3             # Index of the middle waveguide band in the folded structure

# Plot the folded waveguide bands and the target shape centered at the target band
fig, ax = plt.subplots(1, constrained_layout=True, figsize=(4, 4))
ax.plot(pwe0.kpoints[0, :]*5/np.pi, pwe0.freqs, 'b')
ax.plot(pwe0.kpoints[0, :]*5/np.pi, f_target + np.mean(pwe0.freqs[:, b_target]), 'r')
ax.set_ylim([0.35, 0.45])
ax.set_xlim([0, pwe0.kpoints[0, -1]*5/np.pi])
ax.set_xlabel('$5a k_x/ \pi$')
ax.set_ylabel('Frequency');

# Define the objective function for the target band shape
def of_sc(params):
    # Define shift parameters
    dx = params[0:Nx*2*Ny_opt]
    dy = params[Nx*2*Ny_opt:2*Nx*2*Ny_opt]
    dr = params[2*Nx*2*Ny_opt:]

    # Initialize and run the plane-wave expansion
    pwe = wg_sc(dx, dy, dr)
    pwe.run(pol='tm', kpoints=path['kpoints'], numeig=Ny*Nx+10)

    # Get the target band, ans shift by its mean value
    band = pwe.freqs[:, b_target]
    b_shift = band - npa.mean(band)

    # Return the MSE between the shifted band and the target
    return npa.sum(npa.square(b_shift - f_target))

We will first test the gradient computation. Note that legume (through autograd) computes the gradients efficiently, using one backpropagation to get the derivative of the objective w.r.t. every parameter simultaneously. Because of time considerations, for the numerical finite-difference check we will only pick one parameter at random. Also note that we initialize the parameters with some random small values, in order to break the high symmetry of the un-modified structure.

# To compute gradients, we need to set the `legume` backend to 'autograd'
legume.set_backend('autograd')

# Initialize random starting parameters to break the high symmetry
dxrand = 0.01*np.random.randn(Nx*Ny_opt)
dx0 = np.hstack((dxrand, dxrand))
dyrand = 0.01*np.random.randn(Nx*Ny_opt)
dy0 = np.hstack((dyrand,-dyrand))
drrand = 0.01*np.random.randn(Nx*Ny_opt)
dr0 = np.hstack((drrand, drrand))
pstart = np.hstack((dx0, dy0, dr0))

# Choose one parameter index at random for the numerical check
# We pick one of the dy-s because they have the largest gradient in the beginning
ind0 = np.random.randint(Nx*Ny_opt, 2*Nx*Ny_opt, 1)

# Compute the autograd gradients (NB: all at once!)
t = time.time()
grad_a = grad(of_sc)(pstart)
# Print the gradient w.r.t. the parameter index ind0
print("Autograd gradient:  %1.6f, computed in %1.4fs" %(grad_a[ind0], time.time() - t))

# Compute a numerical gradient for parameter ind0
t = time.time()
p_test = np.copy(pstart)
p_test[ind0] = p_test[ind0] + 1e-5
grad_n = (of_sc(p_test) - of_sc(pstart))/1e-5
print("Numerical gradient: %1.6f, computed in %1.4fs" %(grad_n, time.time() - t))

print("Relative difference: %1.2e" %np.abs((grad_a[ind0] - grad_n)/grad_n))

Autograd gradient:  0.000043, computed in 5.5199s
Numerical gradient: 0.000043, computed in 6.5220s
Relative difference: 1.33e-06

The two gradients match very well, and the autograd simulation took as much time to get all the derivatives as it took the numerical simulation time to get just one of the derivatives. The magic of reverse-mode automatic differentiation!

Optimize the waveguide

For the optimization, we will use the autograd function obj_and_grad, which returns the value of the objective function and the gradients simultaneously. Note that legume comes with a Minimize class that implements either adam or lbfgs minimization, which is what we are going to use. Of course, any external optimization function can also be combined with the gradient computation from legume (in fact, the lbfgs minimizer is a wrapper around the SciPy implementation).

# Returns objective function value and gradients
obj_grad = value_and_grad(of_sc)

# Initialize an optimization object; jac=True means that obj_grad returns both the value
# and the jacobian of the objective. The alternative syntax with grad_a defined above
# would be Minimize(of_sc, jac=grad_a), but this will be slower
opt = Minimize(obj_grad, jac=True)

# Run the optimization
(p_opt, ofs1) = opt.lbfgs(pstart, Nepochs=10)

Epoch:    1/  10 | Duration:  39.17 secs | Objective: 5.095522e-05
Epoch:    2/  10 | Duration:  12.81 secs | Objective: 2.059950e-05
Epoch:    3/  10 | Duration:   8.44 secs | Objective: 1.148253e-05
Epoch:    4/  10 | Duration:   6.94 secs | Objective: 6.072189e-06
Epoch:    5/  10 | Duration:   5.95 secs | Objective: 5.844288e-06
Epoch:    6/  10 | Duration:   5.94 secs | Objective: 5.828529e-06
Epoch:    7/  10 | Duration:   6.18 secs | Objective: 5.793352e-06
Epoch:    8/  10 | Duration:   5.94 secs | Objective: 5.692978e-06
Epoch:    9/  10 | Duration:   5.85 secs | Objective: 5.473239e-06
Epoch:   10/  10 | Duration:   6.38 secs | Objective: 5.143095e-06

Finally let’s plot the optiized band, and the underlying structure.

# Get the shifts from the optimized parameters
dx = p_opt[0:Nx*2*Ny_opt]
dy = p_opt[Nx*2*Ny_opt:2*Nx*2*Ny_opt]
dr = p_opt[2*Nx*2*Ny_opt:]

# Initialize and run the PWE with the optimized shifts
pwe = wg_sc(dx, dy, dr)
pwe.run(pol='tm', kpoints=path['kpoints'], numeig=Nx*Ny+25)

# Plot the waveguide bands and the target band
fig, ax = plt.subplots(1, 2, constrained_layout=True)
ax[0].plot(pwe.kpoints[0, :]*5/np.pi, pwe.freqs, 'b')
ax[0].plot(pwe.kpoints[0, :]*5/np.pi, f_target + np.mean(pwe.freqs[:, b_target]), 'r')
ax[0].set_ylim([0.35, 0.55])
ax[0].set_xlim([0, pwe.kpoints[0, -1]*5/np.pi])
ax[0].set_xlabel('$5a k_x/ \pi$')
ax[0].set_ylabel('Frequency');

# Use an in-built legume method to plot the structure
legume.viz.shapes(pwe.layer, ax=ax[1])

Guided mode expansion with Autograd

import numpy as np
import matplotlib.pyplot as plt
import time

import autograd.numpy as npa
from autograd import grad, value_and_grad

import legume
from legume.minimize import Minimize

%load_ext autoreload
%autoreload 2

The autoreload extension is already loaded. To reload it, use:
  %reload_ext autoreload

PhC cavity simulation

In this Notebook, we will optimize a photonic crystal cavity in a lithium niobate slab. We use the small-volume design of Minkov et al., APL 111 (2017), and optimize with respect to shifts in the positions of the air holes. So, we start by defining the starting structure, and a function that builds the photonic crystal given some hole shifts. The PhC parameters are as from Li et al., https://arxiv.org/abs/1806.04755

# Number of PhC periods in x and y directions
Nx, Ny = 16, 10

# Regular PhC parameters
ra = 0.234
dslab = 0.4355
n_slab = 2.21

# Initialize a lattice and PhC
lattice = legume.Lattice([Nx, 0], [0, Ny*np.sqrt(3)/2])

# Make x and y positions in one quadrant of the supercell
# We only initialize one quadrant because we want to shift the holes symmetrically
xp, yp = [], []
nx, ny = Nx//2 + 1, Ny//2 + 1
for iy in range(ny):
    for ix in range(nx):
        xp.append(ix + (iy%2)*0.5)
        yp.append(iy*np.sqrt(3)/2)

# Move the first two holes to create the L4/3 defect
xp[0] = 2/5
xp[1] = 6/5
nc = len(xp)

# Initialize shift parameters to zeros
dx, dy = np.zeros((nc,)), np.zeros((nc,))

# Define L4/3 PhC cavity with shifted holes
def cavity(dx, dy):
    # Initialize PhC
    phc = legume.PhotCryst(lattice)

    # Add a layer to the PhC
    phc.add_layer(d=dslab, eps_b=n_slab**2)

    # Apply holes symmetrically in the four quadrants
    for ic, x in enumerate(xp):
        yc = yp[ic] if yp[ic] == 0 else yp[ic] + dy[ic]
        xc = x if x == 0 else xp[ic] + dx[ic]
        phc.add_shape(legume.Circle(x_cent=xc, y_cent=yc, r=ra))
        if nx-0.6 > xp[ic] > 0 and (ny-1.1)*np.sqrt(3)/2 > yp[ic] > 0:
            phc.add_shape(legume.Circle(x_cent=-xc, y_cent=-yc, r=ra))
        if nx-1.6 > xp[ic] > 0:
            phc.add_shape(legume.Circle(x_cent=-xc, y_cent=yc, r=ra))
        if (ny-1.1)*np.sqrt(3)/2 > yp[ic] > 0 and nx-1.1 > xp[ic]:
            phc.add_shape(legume.Circle(x_cent=xc, y_cent=-yc, r=ra))

    return phc

We will also define a function that initializes and runs a GME computation given shifts dx, dy, and returns the gme object and the quality factor of the fundamental cavity mode.

# Solve for a cavity defined by shifts dx, dy
def gme_cavity(dx, dy, gmax, options):
    # Initialize PhC
    phc = cavity(dx, dy)

    # For speed, we don't want to compute the loss rates of *all* modes that we store
    options['compute_im'] = False

    # Initialize GME
    gme = legume.GuidedModeExp(phc, gmax=gmax)

    # Solve for the real part of the frequencies
    gme.run(kpoints=np.array([[0], [0]]), **options)

    # Find the imaginary frequency of the fundamental cavity mode
    (freq_im, _, _) = gme.compute_rad(0, [Nx*Ny])

    # Finally, compute the quality factor
    Q = gme.freqs[0, Nx*Ny]/2/freq_im[0]
    return (gme, Q)

Let’s first test the forward computation. Note that position shifts of the holes do not affect the average permittivity of the slab that is used for the guided modes. Thus, setting the gradients option to approx actually still returns the exact gradients, but is faster.

# Set some GME options
options = {'gmode_inds': [0], 'verbose': True, 'numeig': Nx*Ny+5, 'gradients': 'approx'}

# Run the simulation for the starting cavity (zero shifts as initialized above)
(gme, Q) = gme_cavity(dx, dy, 2, options)

6.3329s total time for real part of frequencies, of which
  0.3570s for guided modes computation using the gmode_compute='exact' method
  1.1981s for inverse matrix of Fourier-space permittivity
Skipping imaginary part computation, use run_im() to run it, or compute_rad() to compute the radiative rates of selected eigenmodes

# Print the computed quality factor
print("Cavity quality factor: %1.2f" %Q)

# We can also visualize the cavity and the mode profile of the fundamental mode
ax = legume.viz.field(gme, 'e', 0, Nx*Ny, z=dslab/2, component='y', val='abs', N1=300, N2=200)

Cavity quality factor: 6454.17

Autograd backend

Now the real fun starts. We can use legume and autograd to efficiently compute gradients of GME simulations. Let’s first define an objective function as the quality factor.

# To compute gradients, we need to set the `legume` backend to 'autograd'
legume.set_backend('autograd')

# Set GME options
gmax = 2
options = {'gmode_inds': [0], 'verbose': False, 'numeig': Nx*Ny+1, 'gradients': 'approx'}

# Define an objective function which is just the Q of the fundamental mode
def of_Q(params):
    dx = params[0:nc]
    dy = params[nc:]
    (gme, Q) = gme_cavity(dx, dy, gmax=gmax, options=options)

    # We put a negative sign because we use in-built methods to *minimize* the objective function
    return -Q

Test gradient of quality factor

We will test the legume gradient vs. a numerical check. Note that legume (through autograd) computes the gradients efficiently, using one backpropagation to get the derivative of the objective w.r.t. every parameter simultaneously. Because of time considerations, for the numerical finite-difference check we will only pick one parameter at random.

# The autograd function `value_and_grad` returns simultaneously the objective value and the gradient
obj_grad = value_and_grad(of_Q)

# We choose one parameter index at random for the numerical check
ind0 = np.random.randint(0, dx.size, 1)
# We set the starting parameters to zeros, i.e. un-modified cavity
pstart = np.zeros((2*nc, ))

# Compute the autograd gradients (NB: all at once!)
t = time.time()
grad_a = obj_grad(pstart)[1]
# Print the gradient w.r.t. the parameter index ind0
print("Autograd gradient:  %1.4f, computed in %1.4fs" %(grad_a[ind0], time.time() - t))

# Compute a numerical gradient for one selected index
t = time.time()
p_test = np.copy(pstart)
p_test[ind0] = p_test[ind0] + 1e-5
grad_n = (of_Q(p_test) - of_Q(pstart))/1e-5
print("Numerical gradient: %1.4f, computed in %1.4fs" %(grad_n, time.time() - t))

print("Relative difference: %1.2e" %np.abs((grad_a[ind0] - grad_n)/grad_n))

Autograd gradient:  -346.0501, computed in 14.3879s
Numerical gradient: -346.0546, computed in 13.8461s
Relative difference: 1.31e-05

The two gradients match very well, and the autograd simulation took as much time to get all the derivatives as it took the numerical simulation time to get just one of the derivatives. The magic of reverse-mode automatic differentiation!

Test gradient of fields

We will also illustrate computing the gradient with respect to an objective function that depends on the electric field of a mode. We define an objective function as the inverse of the maximum field intensity of the mode, which gives the mode volume (up to some constants). Note that in legume, fields are normalized to integrate to unity by definition.

# Define an objective function which is proportional to the V of the fundamental mode
def of_V(params):
    dx = params[0:nc]
    dy = params[nc:]
    (gme, Q) = gme_cavity(dx, dy, gmax=gmax, options=options)

    # Get the electric field in the center of the slab
    Ey = gme.get_field_xy('e', kind=0, mind=Nx*Ny, z=dslab/2, component='y', Nx=3, Ny=3)[0]['y']

    # Notice the use of autograd.numpy (npa) and not plain numpy (np)
    return 1/npa.square(npa.amax(npa.abs(Ey)))

# The autograd function `value_and_grad` returns simultaneously the objective value and the gradient
obj_grad = value_and_grad(of_V)

# We choose one parameter index at random for the numerical check
ind0 = np.random.randint(0, dx.size, 1)
# We set the starting parameters to zeros, i.e. un-modified cavity
pstart = np.zeros((2*nc, ))

# Compute the autograd gradients (NB: all at once!)
t = time.time()
grad_a = obj_grad(pstart)[1][ind0]
print("Autograd gradient:  %1.4f, computed in %1.4fs" %(grad_a, time.time() - t))

# Compute a numerical gradient for one selected index
t = time.time()
p_test = np.copy(pstart)
p_test[ind0] = p_test[ind0] + 1e-5
grad_n = (of_V(p_test) - of_V(pstart))/1e-5
print("Numerical gradient: %1.4f, computed in %1.4fs" %(grad_n, time.time() - t))

print("Relative difference: %1.2e" %np.abs((grad_a - grad_n)/grad_n))

//anaconda3/lib/python3.7/site-packages/numpy/core/numeric.py:538: ComplexWarning: Casting complex values to real discards the imaginary part
  return array(a, dtype, copy=False, order=order)

Autograd gradient:  0.0973, computed in 13.9638s
Numerical gradient: 0.0973, computed in 15.0929s
Relative difference: 8.39e-05

Note: the ComplexWarning comes from the way autograd implements np.dot(A, B) where A is real and B is complex. It can be disregarded, and as can be seen - the gradients are correct!

Quality factor optimization

The interesting thing about PhC cavities is that their Q can change dramatically upon small changes of the hole positions. On the other hand, the mode profile (and correspondingly the V) stay relatively unchanged. So, below we will optimize solely for the quality factor. Note that legume comes with a Minimize class that implements either adam or lbfgs minimization, which is what we are going to use. Of course, any external optimization function can also be combined with the gradient computation from legume.

# Set the objective function to be the quality factor
obj_grad = value_and_grad(of_Q)

# Initialize an optimization object; jac=True means that obj_grad returns both the value
# and the jacobian of the objective. The alternative syntax is Minimize(of, jac=grad_of)
opt = Minimize(obj_grad, jac=True)

# Starting parameters are the un-modified cavity
pstart = np.zeros((2*nc, ))

# Run an 'adam' optimization
(p_opt, ofs) = opt.adam(pstart, step_size=0.005, Nepochs=10, bounds = [-0.25, 0.25])

Epoch:    1/  10 | Duration:  13.35 secs | Objective: -6.454171e+03
Epoch:    2/  10 | Duration:  14.84 secs | Objective: -9.923204e+03
Epoch:    3/  10 | Duration:  15.49 secs | Objective: -1.613968e+04
Epoch:    4/  10 | Duration:  16.23 secs | Objective: -2.634551e+04
Epoch:    5/  10 | Duration:  15.59 secs | Objective: -3.837220e+04
Epoch:    6/  10 | Duration:  16.13 secs | Objective: -6.522058e+04
Epoch:    7/  10 | Duration:  15.64 secs | Objective: -9.175617e+04
Epoch:    8/  10 | Duration:  14.97 secs | Objective: -1.424175e+05
Epoch:    9/  10 | Duration:  14.68 secs | Objective: -2.015265e+05
Epoch:   10/  10 | Duration:  14.81 secs | Objective: -4.609714e+05

Note that the Q increased by almost two orders of magnitude in just 10 steps of the Adam optimization! Let’s visuzlize the optimized cavity below.

# Optimized parameters
dx = p_opt[0:nc]
dy = p_opt[nc:]

# Run the simulation
(gme, Q) = gme_cavity(dx, dy, gmax=gmax, options=options)
print("Cavity quality factor: %1.2f" %Q)
ax = legume.viz.field(gme, 'e', 0, Nx*Ny, z=dslab/2, component='y', val='abs', N1=200, N2=200)

Cavity quality factor: 460971.43

Refining the optimization

We note that we used a relatively fast optimization for illustration purposes. In the legume paper, we use gmax = 2.5, as well as a 3x3 \(k\)-grid in the Brillouin zone, and average the loss rates. This takes longer time to compute, and significantly more memory if done directly, because all the variables, including all the dense matrices and all the eigenvectors at every \(k\) are stored for the backpropagation. However, we provide a way to overcome this extra memory requirement at a low cost of computational time, as shown below.

First, we define a function that computes the loss rates over a specified set of \(k\)-points.

def gme_cavity_k(params, gmax, options, kpoints):
    dx = params[0:nc]
    dy = params[nc:]
    phc = cavity(dx, dy)
    options['compute_im'] = False
    gme = legume.GuidedModeExp(phc, gmax=gmax)
    gme.run(kpoints=kpoints, **options)
    fims = []
    for ik in range(kpoints[0, :].size):
        (freq_im, _, _) = gme.compute_rad(ik, [Nx*Ny])
        fims.append(freq_im)
    return (gme, npa.array(fims))

We can use this to compute the averaged \(Q\) in a straightforward way. Note that it’s better to average the loss rates and only then compute the \(Q\), rather than average the \(Q\)-s.

# Create an nkx x nky grid in k space (positive kx, ky only)
nkx = 3
nky = 3
kx = np.linspace(0, np.pi/Nx, nkx)
ky = np.linspace(0, np.pi/Ny/np.sqrt(3)*2, nky)
kxg, kyg = np.meshgrid(kx, ky)
kxg = kxg.ravel()
kyg = kyg.ravel()

def Q_kavg(params):
    (gme, fims) = gme_cavity_k(params, gmax, options, np.vstack((kxg, kyg)))
    # Note that the real part of the freq doesn't change much so we can just take ik=0 below
    return gme.freqs[0, Nx*Ny]/2/npa.mean(fims)

print("Refined cavity quality factor: %1.2f" %Q_kavg(p_opt))

Refined cavity quality factor: 41541.38

This \(Q\)-value is three times smaller than what we got above, and is likely closer to the true quality factor. For a converged optimization, it’s typically good to average over a small grid in \(k\)-space as shown here. To use this in an optimization, we can already use Q_kavg as an objective function, but, like mentioned above, this will require a lot of memory - specifically, about 9 times more than the original optimization at a single \(k\)-point. Here is how this can be overcome at a small cost of computational time.

# We set gmax=1 for this example to avoid memory issues
gmax = 1

# Objective function defining the average imaginary part
def fim_kavg(params):
    (gme, fims) = gme_cavity_k(params, gmax, options, np.vstack((kxg, kyg)))
    # Scale for easier printing
    return npa.mean(fims)*1e6

# Compute the gradient and time the computation; print just the first value
obj_grad = value_and_grad(fim_kavg)
t = time.time()
grad_a = obj_grad(pstart)[1]
print("Autograd gradient:  %1.4f, computed in %1.4fs" %(grad_a[0], time.time() - t))

Autograd gradient:  269.5663, computed in 22.0130s

Below is an altenative way to do the same thing; if you compare the memory usage between the cell below and the one above, you’ll realize the purpose of this whole thing.

We use a custom function that mimics map(lambda f: f(params), fns)) in a way that splits the gradient computation instead of storing all the intermediate variables for all functions. NB: the function fns all have to return a scalar and params are all vectors. fmap then returns an array of the same size as the number of functions in fns.

from legume.primitives import fmap

def of_kavg_fmap(params):
    # A function factory to make a list of functions for every k-point
    def fim_factory(ik):
        def fim(params):
            (gme, freq_im) = gme_cavity_k(params, gmax, options, np.array([[kxg[ik]], [kyg[ik]]]))
            return freq_im
        return fim

    fims = fmap([fim_factory(ik=ik) for ik in range(nkx*nky)], params)
    return npa.mean(fims)*1e6

# Compute the gradient and time the computation; print just the first value
obj_grad = value_and_grad(of_kavg_fmap)
t = time.time()
grad_a = obj_grad(pstart)[1]
print("Autograd gradient:  %1.4f, computed in %1.4fs" %(grad_a[0], time.time() - t))

Autograd gradient:  269.5663, computed in 26.9777s

API Reference

This page provides an auto-generated summary of legume’s API.

legume

set_backend(name)

Set the backend for the simulations.

GuidedModeExp

Creating a simulation

GuidedModeExp(phc[, gmax, truncate_g])

Main simulation class of the guided-mode expansion.

Attributes

GuidedModeExp.freqs	Real part of the frequencies of the eigenmodes computed by the guided-mode expansion.
GuidedModeExp.freqs_im	Imaginary part of the frequencies of the eigenmodes computed by the guided-mode expansion.
GuidedModeExp.eigvecs	Eigenvectors of the eigenmodes computed by the guided-mode expansion.
GuidedModeExp.kpoints	Numpy array of shape (2, Nk) with the [kx, ky] coordinates of the k-vectors over which the simulation is run.
GuidedModeExp.gvec	Numpy array of shape (2, Ng) with the [gx, gy] coordinates of the reciprocal lattice vectors over which the simulation is run.
GuidedModeExp.rad_coup	Coupling to TE and TM radiative modes in the claddings.
GuidedModeExp.rad_gvec	Reciprocal lattice vectos corresponding to the radiation emission direction of the coupling constants stored in GuidedModeExp.rad_coup.

Methods

GuidedModeExp.run(kpoints, [0]]), **kwargs)	Compute the eigenmodes of the photonic crystal structure.
GuidedModeExp.run_im()	Compute the radiative rates associated to all the eigenmodes that were computed during GuidedModeExp.run().
GuidedModeExp.compute_rad(kind, minds)	Compute the radiation losses of the eigenmodes after the dispersion has been computed.
GuidedModeExp.get_eps_xy(z[, xgrid, ygrid, …])	Get the xy-plane permittivity of the PhC at a given z as computed from an inverse Fourier transform with the GME reciprocal lattice vectors.
GuidedModeExp.ft_field_xy(field, kind, mind, z)	Compute the ‘H’, ‘D’ or ‘E’ field Fourier components in the xy-plane at position z.
GuidedModeExp.get_field_xy(field, kind, mind, z)	Compute the ‘H’, ‘D’ or ‘E’ field components in the xy-plane at position z.
GuidedModeExp.get_field_xz(field, kind, mind, y)	Compute the ‘H’, ‘D’ or ‘E’ field components in the xz-plane at position y.
GuidedModeExp.get_field_yz(field, kind, mind, x)	Compute the ‘H’, ‘D’ or ‘E’ field components in the yz-plane at position x.
GuidedModeExp.set_run_options([…])	Set multiple options for the guided-mode expansion.

PlaneWaveExp

Creating a simulation

PlaneWaveExp(layer, gmax, eps_eff)

Main simulation class of the plane-wave expansion.

Attributes

PlaneWaveExp.freqs	Frequencies of the eigenmodes computed by the plane-wave expansion.
PlaneWaveExp.eigvecs	Eigenvectors of the eigenmodes computed by the plane-wave expansion.
PlaneWaveExp.kpoints	Numpy array of shape (2, Nk) with the [kx, ky] coordinates of the k-vectors over which the simulation is run.
PlaneWaveExp.gvec	Numpy array of shape (2, Ng) with the [gx, gy] coordinates of the reciprocal lattice vectors over which the simulation is run.

Methods

PlaneWaveExp.run([kpoints, pol, numeig])	Run the simulation.
PlaneWaveExp.get_eps_xy([Nx, Ny, z])	Get the xy-plane permittivity of the layer as computed from an inverse Fourier transform with the PWE reciprocal lattice vectors.
PlaneWaveExp.ft_field_xy(field, kind, mind)	Compute the ‘H’, ‘D’ or ‘E’ field Fourier components in the xy-plane.
PlaneWaveExp.get_field_xy(field, kind, mind)	Compute the ‘H’, ‘D’ or ‘E’ field components in the xy-plane at position z.

Photonic crystal

Lattice(*args)	Class for constructing a Bravais lattice
Lattice.bz_path(pts, ns)	Make a path in the Brillouin zone.
PhotCryst(lattice, eps_l, eps_u)	Class for a photonic crystal which can contain a number of layers.
Layer(lattice, z_min, z_max)	Class for a single layer in the potentially multi-layer PhC.
ShapesLayer(lattice, z_min, z_max, eps_b)	Layer with permittivity defined by Shape objects

Geometry

Circle([eps, x_cent, y_cent, r])	Circle shape
Poly([eps, x_edges, y_edges])	Polygon shape
Square([eps, x_cent, y_cent, a])	Square shape
Hexagon([eps, x_cent, y_cent, a])	Hexagon shape

Visualization

viz.bands(gme[, Q, Q_clip, cone, conecolor, …])	Plot photonic band structure from a GME simulation
viz.structure(struct[, Nx, Ny, Nz, …])	Plot permittivity for all cross sections of a photonic crystal
viz.shapes(layer[, ax, npts, color, lw, pad])	Plot all shapes of Layer
viz.eps_xz(phc[, y, Nx, Nz, ax, clim, cbar, …])	Plot permittivity cross section of a photonic crystal in an xz plane
viz.eps_xy(phc[, z, Nx, Ny, ax, clim, cbar, …])	Plot permittivity cross section of a photonic crystal in an xy plane
viz.eps_yz(phc[, x, Ny, Nz, ax, clim, cbar, …])	Plot permittivity cross section of a photonic crystal in an yz plane
viz.eps_ft(struct[, Nx, Ny, cladding, cbar, …])	Plot a permittivity cross section computed from an inverse FT
viz.reciprocal(struct)	Plot the reciprocal lattice of a GME or PWE object
viz.field(struct, field, kind, mind[, x, y, …])	Visualize mode fields over a 2D slice in x, y, or z

GDS

gds.generate_gds(phc, filename[, unit, …])	Export a GDS file for all layers of a photonic crystal
gds.generate_gds_raster(lattice, raster, …)	Traces a rasterization projected onto a lattice to generate a GDS file