Predictive and generative machine learning models for photonic crystals

2.1 Photonic crystal data set

PhCs are characterized by a periodically varying permittivity ε(r), and the design domain is consequently restricted to a single nontrivial unit cell Ω whose tiling makes up the PhC’s structure (Figure 1A). For simplicity and concreteness, we restrict our attention to 2D square lattices with two material components. Each material occupies a sub-region Ωi of Ω, such that Ω1∪Ω2=Ω, with a resulting “two-tone” permittivity profile ε(r)={ε1, r ∈ Ω1ε2, r ∈ Ω2. For lossless and isotropic materials, εi (as well as the PhC’s allowed eigenfrequencies) are real quantities. As a result, each PhC is effectively characterized by a single “gray-scale image” of ε(r). We generated 20,000 such two-tone, square unit cells. The two disjoint regions Ωi were defined by their boundary region (Figure 1A), which in turn was procedurally generated by casting 2–8 random ellipses sequentially near each other’s periphery, then spanning, smoothing, and centering an enclosing hull, and finally randomly scaling and orienting the resulting boundary. This produces unit cells that are relatively simple geometrically, host just a single inclusion, have no strongly divergent feature scales, and so exemplify realistically fabricable design candidates. We note that stricter constraints could be imposed to align more closely with experimental capabilities (minimum feature sizes could e.g. be ensured by post-processing generated inclusions with standard threshold projection techniques from topology optimization [44], [45]). Nevertheless, to retain a sufficiently varied training set we do not pursue such additional constraints here [46]. The permittivities εi were each drawn uniformly from the range [1,10], roughly spanning the range attainable in transparent materials in the visible spectrum (e.g. at a wavelength of 700 nm, the permittivity of air, silicon nitride, and silicon carbide is approximately 1,4.1, and 6.8, respectively).

Figure 1:

Photonic crystal data set. We generated a data set of 20,000 square 2D PhC unit cells, each consisting of a smooth, centered inclusion of permittivity ε1 in a background permittivity ε2 with εi∈[1,10]. (A) Several representative unit cells and the BZ grid-sampling used in the calculation of band structures. (B) The TM and TE band structures of the PhC highlighted in orange in (A). (C) The generated unit cells predominately feature inclusions occupying less than half the unit cell, as illustrated by a histogramming of the relative inclusion areas across the data set. (D) TM band gaps between bands 1 and 2 consequently occur much more frequently than TE gaps, as TE gaps mainly arise in “filamentory” networks, corresponding to large relative inclusion areas.

For each unit cell, we computed the PhC band structure of the lowest six bands using the free MIT Photonics Bands (MPB) software [47] using 64 × 64 plane waves (equivalent, effectively, to a 64 × 64 spatial resolution). Each unit cell takes ∼2 min on a single core of a 1.6 GHz Core i5-8250U CPU. The calculations are highly converged and accurate: the mean fractional deviation per band between calculations at resolutions of 64 × 64 and 32 × 32 is ∼ 0.1‰, averaged over all unit cells. Figure 1B shows a set of example band structures, split into the transverse magnetic and electric (TM and TE) polarizations: it consists of the set of eigen frequencies ωnk indexed over band numbers n = 1, 2, … 6 and wave vectors k restricted to the Brillouin zone (BZ). For a square lattice of (arbitrary) side length a, the BZ is [−π/a,π/a)×[−π/a,π/a). Since the generated unit cells generically have no exact spatial symmetries, the band structures cannot exhibit any stable band-crossings, allowing a simple sorting of bands by their frequency alone, i.e. ωnk<ωn+1,k.

The resulting data set contains as input pixelized permittivity profiles (in either 32 × 32 or 64 × 64 resolution) and as output the computed band structure (with the BZ sampled on a 23 × 23 Γ-centered Monkhorst–Pack grid, as in Figure 1A). In addition, we computed the band gap Δω12≡minω2k−maxω1k between bands 1 and 2. Since the generated unit cells predominately feature central inclusions with a relative area less than 50% (Figure 1C), TM band gaps are significantly more abundant than TE band gaps (Figure. 1D). In our experiments with generative models, we restricted the data set to those elements that host a substantial band gap, defined heuristically as a relative band gap Δω12/ω¯12 greater than 5% (with mid-gap frequency ω¯12≡12minω2k+12maxω1k). Since the TE band structures host only very few such examples (48 with a non-zero band gap and 10 with a band gap ≥ 5%, out of 20,000 examples), we confined our experiments with generative models to the TM polarization only.

2.2 Band prediction

A natural question is whether neural networks can be used in lieu of traditional theory-driven tools for the modeling of PhCs, e.g. to predict a PhC’s band structure. To answer this, we adopted a supervised learning approach and trained two neural networks to reproduce the TM and TE band structures, respectively, taking as input a 32 × 32-discretized unit cell and producing as output the band structure across the 23 × 23-discretized BZ for the first six bands (Figure 2). Effectively, this is a regression problem where a large input space (32×32=1024 parameters) is mapped to a large output space (23×23×6=3174 parameters).

Figure 2:

Band prediction with convolutional neural networks. (A) Network architecture showing the convolutional encoder and fully-connected decoder (described in detail in the main text). Numbers in red indicate the data size after every network layer. (B–C) Example applications of the trained band-prediction network on test set unit cells in both TM and TE polarizations (green markers, network predictions; surfaces, reference MPB calculations). The chosen unit cells represent worst-case examples due to their large permittivity contrast. (D–E) The relative deviation between network predictions and reference calculations. The relative error is typically very small, on the order of ≲ 2%.

The network consists of two main components: encoder and decoder (Figure 2A). Conceptually, the encoder is tasked with building an abstract representation of the PhC’s unit cell ε(r) that spans a lower-dimensional so-called feature (or latent) space. The decoder, conversely, is tasked with reconstructing from this feature vector the band structure of the input PhC. In practice, we implement and train the network using the popular PyTorch framework [48]. Training is accomplished by minimizing the mean square error between the training data ωnk and network output ωnkNN across n, k, and the entire training set (the cost function) using adaptive gradient descent optimization (RMSprop [49]) with an adaptive learning rate scheduler. We implement the encoder using three convolutional layers, each of (zero-padded) 11 × 11 kernels, followed by two fully-connected layers, essentially mapping the 32 × 32 input space into a linear 64-dimensional feature space. The convolutional layers were subjected to max-pooling and increasing channel depths to collapse the 2D input into a simple 1D vector that could be directly fed to the fully-connected layers of the encoder. The decoder was implemented with six feed-forward networks, each consisting of five fully-connected layers that were separately optimized for each band. All layers were followed by ReLU activations and batch normalization [50] was used for the convolutional layers. Our implementation (with optimized hyper-parameters) is available online, see Ref. 51, and summarized in Figure 2A.

We followed the standard training–validation–test approach and split the data set into training, validation, and test sets (in 70, 15, and 15% proportions). The training set was used to update the network’s weights, the validation set to evaluate training convergence and select hyper-parameters, and the test set to determine the network’s ability to generalize to new data (i.e. assess eventual network performance). We performed a simple grid-search to determine hyper-parameters, searching across kernel sizes of convolution layers ∈[5,7,9,11], batch sizes ∈[32,64,128], initial learning rates ∈[10−5,10−4,10−3], and total number of training epochs ∈[20,30,40] (optimal hyper-parameters indicated in boldface). In addition, we searched across several network architectures consisting of varying convolution layer channel depths to arrive at the optimal configuration shown in Figure 2A. Application of the optimally tuned network on two examples from the test set is shown in Figure 2B–E, in absolute (Figure 2B–C) and relative scales (Figure 2D–E). Both examples are characterized by a large permittivity contrast between inclusion and background and consequently reflect extremal elements in the data set, whose band structures deviate substantially from the trivial empty-lattice approximation. Averaged across the entire validation and test sets, both the band-specific and the band-averaged relative mean errors meank(|ωnkNN−ωnk|/ωnk) are generally very low, on the order of 0.5%, as shown in Table 1. We conclude that a simple convolution neural network can predict the band structures of PhCs with very high accuracy and generalizes excellently to examples not seen during training. While we have confined our attention to 2D square lattices, this conclusion appears likely to apply generally across different lattice types and dimensionalities.

It is worth noting that while generation of a suitable data set—and, to a lesser extent, network training (taking ∼3 min for fixed hyper-parameters on an Nvidia 1080 Ti GPU)—requires substantial computing resources, once trained, a neural network can predict band structures orders of magnitude faster than conventional theory-driven simulations (network evaluation of a single input takes ≈ 0.02 s on an Nvidia 1080 Ti GPU). While these gains are not sufficiently attractive to merit the training of regression networks for one- or few-off calculations, they can be relevant in inverse-design problems [26], [27] or high-throughput searches [52], where a very large number of distinct system configurations must be considered.

2.3 Generative adversarial networks

While ML techniques for classification and regression problems (such as band structure prediction) are naturally complementary to traditional theory-based approaches to forward problems, the field of generative modeling stands to complement conventional techniques of optimization and inverse design. Rather than learning a mapping from input to output data (e.g. from the unit cell to band structure), generative models generally seek to learn the statistical distribution of data samples. Once learned, many new elements can then be drawn from this distribution—a highly attractive option for optimization problems characterized by a non-unique solution space (in sharp contrast to conventional gradient-based approaches where the retrieval of diverse design candidates can be nontrivial).

Generative adversarial networks (GANs) have become a singularly prominent direction in generative models [53], due to their ability to seemingly generalize “creatively” beyond training data, with applications spanning e.g. autonomous driving systems [54], natural image synthesis [55], and anomaly detection [56]. The training of GANs mimics an adversarial game between two networks (Figure 3): while one network, the discriminator, is tasked with deciding whether a given input belongs to the training data (“real”) or not (“fake”); the other, the generator, is tasked with producing (from an input vector sampled from a predefined probabilistic feature space) candidates that fool the discriminator. During training, their joint cost function—whose contributions are adversarial in nature, i.e. generally opposing—is optimized.

Figure 3:

Generative adversarial network. Through an adversarial game between a generative (G) and a discriminative (D) network, new synthetic examples (fake) of 2D unit cells with a TM band gap can be generated from a genuine data set (real).

We explored the use of GANs for synthesizing new candidate unit cells that host a substantial TM band gap. To do so, we extracted the 585 unit cells with Δω12/ω¯12≥5% from the data set for use as training data. We tested three different GAN-variants [57]: a conventional GAN [53], a least squares GAN (LSGAN) [58], and Deep Regret Analytic GAN (DRAGAN) [59], each distinguished essentially by their respective generator and discriminator cost functions [60]. In each case, we adapted standard off-the-shelf implementations [61] to take a single-channel, 64 × 64 pixelized ε(r) profile as training data. Training across 400 epochs took on the order of 5–10 min for each GAN on an Nvidia 1080 Ti GPU.

Figure 4A illustrates the improvement during the training of each GAN-variant’s ability to generate convincing unit cells that exhibit the desired characteristics (i.e. well-defined, high-contrast, two-tone inclusions). We also evaluated the models’ performance relative to the design goal of exhibiting a substantial band gap by computing the band gap sizes of the generated unit cells with MPB (Figure 4B). Concretely, we trained 10 distinct networks for each GAN-variant (distinguished only by network initialization), outputting at each epoch 16 generated unit cells. From these samples, we evaluated a notion of “generation fidelity”, defined as the relative fraction of generated unit cells that indeed exhibit a band gap ≥5%. Both metrics—visual “quality” and fidelity—exhibit much the same evolution: initially, performance is poor, reflecting essentially randomly initialized networks; then, within a few epochs, performance improves dramatically; and finally, performance slowly deteriorates, typical of the saturation problem [62]. While GAN and LSGAN achieve convincing performance within ∼50 epochs, DRAGAN takes significantly longer, apparently passing through a phase of “fractured” inclusions. Further, across our 10 training experiments, we identified only a single successful DRAGAN trial (others not shown).

Figure 4:

GAN, LSGAN, and DRAGAN for generation of unit cells with substantial TM band gaps. (A) The mapping of fixed feature vectors to generated unit cells during training. Note the differing epoch steps and ranges for DRAGAN versus GAN and LSGAN. (B) Fidelity of generated unit cells (the fraction hosting a band gap ≥5%). For GAN and LSGAN, fidelity is averaged over 16 distinct feature vectors and 10 training runs (uncertainty across training runs is indicated by shaded regions). Only a single DRAGAN training run was successful (averaged over 16 outputs). (C–D) Examples of generated unit cells at selected epochs (indicated by matching markers in B). GAN and LSGAN produce more well-defined but lower-fidelity unit cells at later epochs (text-insets give Δω12/ω¯12 evaluated with MPB; dashed borders highlight cases where Δω12/ω¯12<5%).

Figure 4C shows 16 examples of generated unit cells for each GAN-variant, evaluated at epochs and training runs of 100% fidelity. The generative models have clearly “learned” the key elements necessary to host a TM band gap, namely an inclusion of high permittivity embedded in a low-permittivity background [42]. Interestingly, although the fidelity of GAN and LSGAN generally decreases after peaks around range 5070 epochs, the visual quality—especially the well-definedness of inclusion boundaries—improves at higher epochs as shown in Figure 4D. The apparent cost of moving to higher epochs appear to be an increase in low-contrast examples without (or with smaller) band gaps. More generally, both visual quality and fidelity alike could likely be improved by simply enlarging the training set’s size. Finally, we note that regularization and filtering techniques from topology optimization [44], [45] could be leveraged to further reduce noise or ensure minimum feature sizes in the generated designs, either as a post-processing step or during training.

2.4 Image-to-image translation

Image-to-image translation can be viewed as a subset of generative modeling concerned with translating (i.e. mapping) between distinct representations of images. Effectively, this translation can often be viewed simply as implanting the “style” or characteristics of a given representation A onto another B; say, mapping from an outline, or even a sketch, to a photorealistic representation (e.g. of cats [63]). Following the introduction of the pix2pix software [64], conditional GANs [65] have emerged as a powerful tool to achieve this translation. The underlying principle is illustrated in Figure 5A: the generator of a conditional GAN takes, in addition to the standard random feature vector x, a “conditional input” y (of representation A) from which a fake output G(x,y) is generated—the discriminator, conversely, seeks to distinguish between genuine pairings of y and real output z (of representation B) from faked pairings.

Figure 5:

Image-to-image translation of photonic features. (A) Conditional GANs, as implemented e.g. by pix2pix [64], facilitate image-to-image translation by augmenting a conventional GAN (Figure 3) with a conditional input. (B) Using pix2pix, we trained a model to translate a discretized inclusion outline (black borders) to a permittivity profile (red borders) hosting a TM band gap. The permittivity contrast Δε≡maxϵ(r)−minϵ(r) and the relative band gap Δω12/ω¯12, evaluated with MPB, are indicated below each design (dashed borders highlight cases where Δω12/ω¯12<5%).

A natural application of image-to-image translation, and pix2pix specifically, for photonics is “guided inverse design”, i.e. inverse design subject to conditional input. Figure 5B illustrates one such application (using a PyTorch implementation of pix2pix [64], [66], [67]): by taking again the set of unit cells with a TM band gap ≥5% and choosing as conditional input the corresponding inclusion outlines, we can learn a mapping from outlines to permittivity profiles supporting a TM band gap. We trained the model using just 256 samples (each of 64 × 64 pixels) over 200 epochs (requiring less than 1 h on an Nvidia 1080 Ti GPU). We tested the trained model on conditional input of several distinct shapes (heart and five- and four-pointed stars) and scales. The trained model successfully translates each large inclusion to a permittivity profile with a TM band gap ≥5%. Notably, this translation is successful—and maintains the outline’s shape—even though the training data does not contain examples that resemble the chosen outlines. Further, when the scale of a shape is reduced, we observe that the contrast in the generated profile is increased; in exact agreement with the basic design-principle suggested by perturbation theory [42]. While the small five- and four-pointed stars translations do not achieve a TM band gap ≥5%, it is clear that the design approach (i.e. increasing contrast) is valid. Indeed, for sufficiently small or irregular inclusions, designs with ε(r)∈[1,10] and a ≥5% band gap may not exist. We can explore this latter point by feeding the trained model a stick too narrow to host a TM band gap (Figure 5B, bottom). We sampled three generated designs (distinct feature vectors): in each case, the design “breaks out” of the outline and maximizes contrast. The resulting rupture varies slightly in extent and so hosts differently sized band gaps, though in each case ≤5%.