From f8497cbc7aa1e0c229c66aa96b0f07d23ad9c0fd Mon Sep 17 00:00:00 2001 From: Alec Hammond Date: Thu, 7 Mar 2024 10:25:48 -0800 Subject: [PATCH] add waveguide crossing example --- python/examples/waveguide_crossing.py | 609 ++++++++++++++++++++++++++ 1 file changed, 609 insertions(+) create mode 100644 python/examples/waveguide_crossing.py diff --git a/python/examples/waveguide_crossing.py b/python/examples/waveguide_crossing.py new file mode 100644 index 000000000..57eec96b3 --- /dev/null +++ b/python/examples/waveguide_crossing.py @@ -0,0 +1,609 @@ +# waveguide_crossing.py Using meep's adjoint solver, designs a waveguide +# crossing that maximizes transmission at a single frequency. Two approaches are +# demonstrated: (1) start with the nominal crossing shape and perform shape +# optimization via meep's smoothed projection feature; (2) run a full topology +# optimization starting from an initial grayscale. Evolve the design until β=∞. +# +# Importantly, this particular example highlights some of the ways one can use +# the novel smoothed projection function to perform both shape and topology +# optimization. + + +from typing import Callable, List, Optional, Tuple + +import meep.adjoint as mpa +import nlopt +import numpy as np +from autograd import grad +from autograd import numpy as npa +from autograd import tensor_jacobian_product +from matplotlib import pyplot as plt + +import meep as mp + +mp.quiet() + +DEFAULT_MIN_LENGTH = 0.09 +DEFAULT_DESIGN_REGION_WIDTH = 3.0 +DEFAULT_DESIGN_REGION_HEIGHT = 3.0 +DEFAULT_WAVEGUIDE_WIDTH = 0.5 +DEFAULT_ETA = 0.5 +DEFAULT_ETA_E = 0.75 +DEFAULT_MAX_EVAL = 30 + + +def build_optimization_problem( + resolution: float, + beta: float, + use_smoothed_projection: bool, + min_length: float = DEFAULT_MIN_LENGTH, + dx: float = DEFAULT_DESIGN_REGION_WIDTH, + dy: float = DEFAULT_DESIGN_REGION_HEIGHT, + waveguide_width: float = DEFAULT_WAVEGUIDE_WIDTH, + eta: float = DEFAULT_ETA, + eta_e: float = DEFAULT_ETA_E, +) -> Tuple[mpa.OptimizationProblem, Callable]: + """Build the waveguide-crossing optimization problem. + + The waveguide crossing is a cononical inverse design problem with both + shape- and topology-optimization implementations. The idea is to find the + optimal structure that maximizes transmission from one side to the other. It + exhibits C4 symmetry, and generally resembles the following structure: + + | | + | | + ----- ------ + ----- ------ + | | + | | + + Args: + resolution: Simulation resolution in pixels/micron. + beta: Tanh function projection strength parameter, ranging from [0,∞]. + use_smoothed_projection: Whether or not to use the smoothed projection. + min_length: Minimum length scale in microns. + dx: Design region width in microns. + dy: Design region height in microns. + waveguide_width: Waveguide width in microns. + eta: Projection function threshold parameter. + eta_e: Projection function eroded threshold parameter. + + Returns: + The corresponding optimization problem object and the mapping function + that applies the linear and nonlinear transformations. + """ + # Map the design region resolution to the yee grid, which is twice the standard resolution. + design_region_resolution = int(2 * resolution) + + # pml thickness + dpml = 1.0 + + filter_radius = mpa.get_conic_radius_from_eta_e(min_length, eta_e) + + sxy = dx + 1 + 2 * dpml + + silicon = mp.Medium(epsilon=12) + cell_size = mp.Vector3(sxy, sxy, 0) + boundary_layers = [mp.PML(thickness=dpml)] + + eig_parity = mp.EVEN_Y + mp.ODD_Z + + design_region_size = mp.Vector3(dx, dy) + Nx = int(design_region_resolution * design_region_size.x) + 1 + Ny = int(design_region_resolution * design_region_size.y) + 1 + + waveguide_geometry = [ + mp.Block(material=silicon, size=mp.Vector3(mp.inf, waveguide_width, mp.inf)), + mp.Block(material=silicon, size=mp.Vector3(waveguide_width, mp.inf, mp.inf)), + ] + + # Source centered in optical c-band + fcen = 1 / 1.55 + df = 0.23 * fcen + sources = [ + mp.EigenModeSource( + src=mp.GaussianSource(fcen, fwidth=df), + center=mp.Vector3(-0.5 * sxy + dpml + 0.1, 0), + size=mp.Vector3(0, sxy - 2 * dpml), + eig_band=1, + eig_parity=eig_parity, + ) + ] + + matgrid = mp.MaterialGrid( + mp.Vector3(Nx, Ny), + mp.air, + silicon, + weights=np.ones((Nx, Ny)), + beta=0, # disable meep's internal smoothing + do_averaging=False, # disable meep's internal mg smoothing + ) + + matgrid_region = mpa.DesignRegion( + matgrid, + volume=mp.Volume( + center=mp.Vector3(), + size=mp.Vector3(design_region_size.x, design_region_size.y, 0), + ), + ) + + matgrid_geometry = [ + mp.Block( + center=matgrid_region.center, size=matgrid_region.size, material=matgrid + ) + ] + + geometry = waveguide_geometry + matgrid_geometry + + sim = mp.Simulation( + resolution=resolution, + cell_size=cell_size, + boundary_layers=boundary_layers, + sources=sources, + geometry=geometry, + ) + + frequencies = [fcen] + + obj_list = [ + mpa.EigenmodeCoefficient( + sim, + mp.Volume( + center=mp.Vector3(-0.5 * sxy + dpml + 0.2), + size=mp.Vector3(0, sxy - 2 * dpml, 0), + ), + 1, + eig_parity=eig_parity, + ), + mpa.EigenmodeCoefficient( + sim, + mp.Volume( + center=mp.Vector3(0.5 * sxy - dpml - 0.2), + size=mp.Vector3(0, sxy - 2 * dpml, 0), + ), + 1, + eig_parity=eig_parity, + ), + ] + + def J(input, output): + """Simple objective function to minimize loss.""" + return 1 - npa.power(npa.abs(output / input), 2) + + opt = mpa.OptimizationProblem( + simulation=sim, + maximum_run_time=500, + objective_functions=J, + objective_arguments=obj_list, + design_regions=[matgrid_region], + frequencies=frequencies, + ) + + def mapping(x: npa.ndarray): + """Applies the smoothing and projection.""" + x = x.reshape(Nx, Ny) + x = mpa.conic_filter( + x, + filter_radius, + design_region_size.x, + design_region_size.y, + design_region_resolution, + ) + + # Enforce the square symmetry + x = ( + x + + npa.rot90(x) + + npa.rot90(npa.rot90(x)) + + npa.rot90(npa.rot90(npa.rot90(x))) + ) / 4 + + # Only used the smoothed projection if prompted. + if use_smoothed_projection: + x = mpa.smoothed_projection( + x, beta=beta, eta=eta, resolution=design_region_resolution + ) + else: + x = mpa.tanh_projection(x, beta=beta, eta=eta) + + x = npa.clip(x, 0, 1) + + return x.flatten() + + return opt, mapping + + +def nlopt_fom( + x: np.ndarray, + gradient: np.ndarray, + opt: mpa.OptimizationProblem, + mapping: Callable, + data: List, + results: List, +): + """Wrapper for NLopt FOM. + + Args: + x: Degrees of freedom array. + gradient: Gradient of FOM. + opt: Optimization problem object. + mapping: Mapping function. + data: Structure to store the simulated design each iteration. + results: Structure to store the simulated FOM each iteration. + + Returns: + The FOM value at the current iteration. + """ + grid_size = opt.design_regions[0].design_parameters.grid_size + Nx, Ny = int(grid_size.x), int(grid_size.y) + data.append(mapping(x).copy().reshape(Nx, Ny)) + + f0, dJ_du = opt([mapping(x)]) + backprop_gradient = tensor_jacobian_product(mapping, 0)(x, dJ_du) + if gradient.size > 0: + gradient[:] = backprop_gradient + + print("FOM: {}".format(np.real(f0))) + results.append(np.real(f0)) + + return float(np.real(f0)) + + +def _plot_optimization_results( + data: List, + results: List, + num_samples: int = 4, +): + samples = np.linspace(0, len(results), num_samples, dtype=int, endpoint=False) + plt.figure(figsize=(5.25, 4.0)) + plt.subplot(2, 1, 1) + plt.plot(results, "o-") + plt.xlabel("Optimization Iteration") + plt.ylabel("FOM") + for k in range(len(samples)): + plt.subplot(2, 4, 5 + k) + plt.imshow(data[samples[k]], cmap="binary") + plt.axis("off") + plt.title(f"It. {samples[k]+1}") + plt.tight_layout() + plt.show() + + +def run_shape_optimization( + resolution: float, + beta: float, + maxeval: int, + use_smoothed_projection: bool = True, + dx: float = DEFAULT_DESIGN_REGION_WIDTH, + dy: float = DEFAULT_DESIGN_REGION_HEIGHT, + waveguide_width: float = DEFAULT_WAVEGUIDE_WIDTH, + output_filename_prefix: Optional[str] = None, +): + """Run shape optimization using a cross as a starting guess. + + Args: + resolution: Simulation resolution in pixels/micron. + beta: Tanh function projection strength parameter, ranging from [0,∞]. + maxeval: Maximum number of optimization iterations to run. + use_smoothed_projection: Whether or not to use the smoothed projection. + dx: Design region width in microns. + dy: Design region height in microns. + waveguide_width: Waveguide width in microns. + output_filename_prefix: The filename prefix that will store the + optimization results. If `None`, no file is saved. + + Returns: + The design and FOM result arrays, along with the optimization problem + and mapping function. + + """ + # Initialize the optimization problem and mapping functions + opt, mapping = build_optimization_problem( + resolution=resolution, + beta=beta, + dx=dx, + dy=dy, + waveguide_width=waveguide_width, + use_smoothed_projection=use_smoothed_projection, + ) + + # pull number of parameters from the design region + n = opt.design_regions[0].design_parameters.weights.size + grid_size = opt.design_regions[0].design_parameters.grid_size + Nx, Ny = int(grid_size.x), int(grid_size.y) + + # Set up the optimizer + algorithm = nlopt.LD_CCSAQ + solver = nlopt.opt(algorithm, n) + solver.set_lower_bounds(0) + solver.set_upper_bounds(1) + solver.set_maxeval(maxeval) + + # initial guess, which is just a simple waveguide crossing + x = np.linspace(-dx / 2, dy / 2, Nx) + y = np.linspace(-dx / 2, dy / 2, Ny) + X, Y = np.meshgrid(x, y) + mask = (np.abs(Y) <= waveguide_width / 2.0) + (np.abs(X) <= waveguide_width / 2.0) + x0 = np.zeros((Nx, Ny)) + x0[mask.T] = 1 + x0 = mapping(x0) + + # Create empty datastructures we can use to log the results + data = [] + results = [] + + # prepare the optimizer objective function. We need to wrap the above fom in + # a format that nlopt expects. + nlopt_fom_simple = lambda x, gradient: nlopt_fom( + x, gradient, opt=opt, mapping=mapping, data=data, results=results + ) + solver.set_min_objective(nlopt_fom_simple) + + # Run the optimization + x_final = solver.optimize(x0) + + # Log the final results + data.append(x_final.copy()) + opt.update_design([mapping(x_final)]) + f0, _ = opt(need_gradient=False) + results.append(np.real(f0)) + + _plot_optimization_results(data, results) + + # Save to disk + if mp.am_really_master() and (output_filename_prefix is not None): + np.savez(output_filename_prefix + "_data.npz", data=data, results=results) + + return data, results, opt, mapping + + +def run_topology_optimization( + resolution: float, + beta_evolution: List[float], + maxeval: int, + use_smoothed_projection: bool = True, + dx: float = DEFAULT_DESIGN_REGION_WIDTH, + dy: float = DEFAULT_DESIGN_REGION_HEIGHT, + waveguide_width: float = DEFAULT_WAVEGUIDE_WIDTH, + output_filename_prefix: Optional[str] = None, +): + """Run shape optimization using a cross as a starting guess. + + Args: + resolution: Simulation resolution in pixels/micron. + beta_evolution: List of Tanh function projection strength parameter, + ranging from [0,∞], for each optimization epoch. + maxeval: Maximum number of optimization iterations to run. + use_smoothed_projection: Whether or not to use the smoothed projection. + dx: Design region width in microns. + dy: Design region height in microns. + waveguide_width: Waveguide width in microns. + output_filename_prefix: The filename prefix that will store the + optimization results. If `None`, no file is saved. + + Returns: + The design and FOM result arrays, along with the optimization problem + and mapping function. + + """ + # Initialize the optimization problem and mapping functions so that we can + # set up the problem. + opt, mapping = build_optimization_problem( + resolution=resolution, + beta=beta_evolution[0], + dx=dx, + dy=dy, + waveguide_width=waveguide_width, + use_smoothed_projection=use_smoothed_projection, + ) + + # pull number of parameters from the design region + n = opt.design_regions[0].design_parameters.weights.size + grid_size = opt.design_regions[0].design_parameters.grid_size + Nx, Ny = int(grid_size.x), int(grid_size.y) + + # Set up the optimizer + algorithm = nlopt.LD_CCSAQ + solver = nlopt.opt(algorithm, n) + solver.set_lower_bounds(0) + solver.set_upper_bounds(1) + solver.set_maxeval(maxeval) + + # initial guess, which is just a uniform gray region + x0 = np.ones((Nx, Ny)) + x0 = mapping(x0) + + # Create empty datastructures we can use to log the results + data = [] + results = [] + + for beta in beta_evolution: + # Re-initialize the optimization problem and mapping functions + opt, mapping = build_optimization_problem( + resolution=resolution, + beta=beta, + dx=dx, + dy=dy, + waveguide_width=waveguide_width, + use_smoothed_projection=use_smoothed_projection, + ) + + # prepare the optimizer objective function. We need to wrap the above fom in + # a format that nlopt expects. + nlopt_fom_simple = lambda x, gradient: nlopt_fom( + x, gradient, opt=opt, mapping=mapping, data=data, results=results + ) + solver.set_min_objective(nlopt_fom_simple) + + # Run the optimization + x0 = solver.optimize(x0) + + # Log the final results + data.append(x0.copy()) + opt.update_design([mapping(x0)]) + f0, _ = opt(need_gradient=False) + results.append(np.real(f0)) + + # Save to disk + if mp.am_really_master() and (output_filename_prefix is not None): + np.savez(output_filename_prefix + "_data.npz", data=data, results=results) + + return data, results, opt, mapping + + +def analyze_gradient_convergence( + resolution: float, + beta_range: List[float], + dx: float = DEFAULT_DESIGN_REGION_WIDTH, + dy: float = DEFAULT_DESIGN_REGION_HEIGHT, + waveguide_width: float = DEFAULT_WAVEGUIDE_WIDTH, +) -> Tuple[List[float], List[float]]: + """Analyze the norm of the gradient vs beta. + + This experiment plots the norm of the gradient vector as a function of beta + for both the smoothed projection and standard projection functions. As seen + here, the smoothed projection function has a well-defined. + + Args: + resolution: Simulation resolution in pixels/micron. + beta_range: List of projection threshold parameters to test. + dx: Design region width in microns. + dy: Design region height in microns. + waveguide_width: Waveguide width in microns. + + Returns: + The norm of the gradient array when using smoothing and when not using + smoothing as a function of beta. + """ + + # Initialize the optimization problem and mapping functions + opt, mapping = build_optimization_problem( + resolution=resolution, + beta=beta_range[0], + dx=dx, + dy=dy, + waveguide_width=waveguide_width, + use_smoothed_projection=True, + ) + + # pull number of parameters from the design region + n = opt.design_regions[0].design_parameters.weights.size + grid_size = opt.design_regions[0].design_parameters.grid_size + Nx, Ny = int(grid_size.x), int(grid_size.y) + + # initial guess, which is just a simple waveguide crossing + x = np.linspace(-dy / 2, dy / 2, Nx) + y = np.linspace(-dx / 2, dy / 2, Ny) + X, Y = np.meshgrid(x, y) + mask = (np.abs(Y) <= waveguide_width / 2.0) + (np.abs(X) <= waveguide_width / 2.0) + x0 = np.zeros((Nx, Ny)) + x0[mask.T] = 1 + x0 = mapping(x0) + + # Initialize storage structures + with_smoothing = [] + without_smoothing = [] + + for beta in beta_range: + # Compute the norm with smoothing + opt, mapping = build_optimization_problem( + resolution=resolution, + beta=beta, + dx=dx, + dy=dy, + waveguide_width=waveguide_width, + use_smoothed_projection=True, + ) + _, dJ_du = opt([mapping(x0)]) + backprop_gradient = tensor_jacobian_product(mapping, 0)(x0, dJ_du) + norm_gradient_smoothing = np.linalg.norm(backprop_gradient.flatten()) + with_smoothing.append(norm_gradient_smoothing) + + # Compute the norm without smoothing + opt, mapping = build_optimization_problem( + resolution=resolution, + beta=beta, + dx=dx, + dy=dy, + waveguide_width=waveguide_width, + use_smoothed_projection=False, + ) + _, dJ_du = opt([mapping(x0)]) + backprop_gradient = tensor_jacobian_product(mapping, 0)(x0, dJ_du) + norm_gradient_no_smoothing = np.linalg.norm(backprop_gradient.flatten()) + without_smoothing.append(norm_gradient_no_smoothing) + + # plot results + plt.figure(figsize=(5.2, 2.5), constrained_layout=True) + plt.loglog(beta_range, without_smoothing, label="W/o smoothing") + plt.loglog(beta_range, with_smoothing, label="W/ smoothing") + plt.xlabel("β") + plt.ylabel("|df/dx|") + plt.show() + + return with_smoothing, without_smoothing + + +def analyze_FOM_convergence( + resolution: float, + beta: float, + maxeval: int, +) -> Tuple[List[float], List[float]]: + """Analyze the convergence of the new projection method. + + Because the smoothed projection has a well-defined gradient for all values + of beta, it should exhibit better convergence properties than the standard + projection step, which grows increasingly more ill-conditioned as + beta->infinity. This experiment runs the same shape optimization problem for + both projection functions and plots the results. + + Args: + resolution: Simulation resolution in pixels/micron. + beta: Tanh function projection strength parameter, ranging from [0,∞]. + maxeval: Maximum number of optimization iterations to run. + + Returns: + The FOM evolution for both the smoothed and not smoothed case as a + function of optimization iteration. + + """ + print("Running shape optimization WITHOUT smoothing...") + _, results, _, _ = run_shape_optimization( + beta=beta, + resolution=resolution, + maxeval=maxeval, + use_smoothed_projection=False, + ) + print("Running shape optimization WITH smoothing...") + ( + _, + results_smoothed, + _, + _, + ) = run_shape_optimization( + beta=beta, resolution=resolution, maxeval=maxeval, use_smoothed_projection=True + ) + + plt.figure(figsize=(5.2, 2.0), constrained_layout=True) + plt.plot(results, "o-", label="W/o smoothing") + plt.plot(results_smoothed, "o-", label="W/ smoothing") + plt.legend() + plt.xlabel("Optimization iteration") + plt.ylabel("FOM") + plt.show() + + return results, results_smoothed + + +if __name__ == "__main__": + run_shape_optimization(resolution=25.0, beta=np.inf, maxeval=30) + + run_topology_optimization( + resolution=10.0, beta_evolution=[8, 32, np.inf], maxeval=10 + ) + + analyze_FOM_convergence(resolution=25, beta=64, maxeval=25) + + analyze_gradient_convergence( + beta_range=[2, 4, 8, 25, 50, 100, 200, 400, 1000], resolution=20 + )