TIDY3D
LEARNING CENTER

Lectures
Tutorials
Applications

Example Library / Global Optimization / Particle swarm optimization of a polarization beam splitter

Particle swarm optimization of a polarization beam splitter

Sign up now to simulate .ipynb

Note: the cost of running the entire notebook is larger than 10 FlexCredits.

Particle Swarm Optimization (PSO) is a powerful and versatile optimization technique. PSO has since emerged as a popular metaheuristic algorithm for solving complex optimization problems across a wide range of domains, including integrated photonic device design. At its core, PSO embodies the principles of collaboration and information sharing within a population of simple entities known as “particles.” These particles navigate through a multi-dimensional search space, seeking optimal solutions to a given problem. They learn from their own experiences, as well as the experiences of other particles, to adapt their positions and velocities over time.

In this notebook, we demonstrate the PSO of a compact polarization beam splitter (PBS) using Tidy3D and the PySwarms library implemented in the Tidy3D Design plugin. The coupling region is segmented into 10 sections with varying widths \(W_1\), \(W_2\), …, \(W_{10}\). Since performing the traditional parameter sweep is unrealistic in this case due to the large parameter space dimension, we use PSO to optimize all the widths to maximize the beam splitting of the TE0 mode and the TM0 mode. The device is inspired by the work Weiwei Chen, et al., "Ultra-compact and low-loss silicon polarization beam splitter using a particle-swarm-optimized counter-tapered coupler," Opt. Express 28, 30701-30709 (2020)DOI:10.1364/OE.408432. In the original work, the authors employed a more complex PSO scheme of changing objective functions. For the sake of simplicity, we will use a simpler PSO scheme in this notebook, only aiming to demonstrate the idea and workflow. Tidy3D users can draw inspiration from this notebook, as well as from the wealth of PSO research in the literature. By harnessing the versatility of PSO in combination with the exceptional speed of Tidy3D, we anticipate that users can design many novel and high-performance photonic devices.

Schematic of the polarization beam splitter

Besides the PSO introduced in this notebook, Tidy3D also provides built-in automatic differentiation support for adjoint optimization. Unlike PSO, adjoint optimization is gradient-based and thus more efficient. To learn more, please refer to the adjoint optimizations of a wavelength division multiplexer, a mode converter, and a waveguide taper. If you are new to adjoint optimization, please start with the tutorial on adjoint basis and our video lectures on inverse design.

# uncomment the following line to install pyswarms if it's not installed in your environment already
# pip install pyswarms

import gdstk
import matplotlib.pyplot as plt
import numpy as np
import tidy3d as td
import tidy3d.plugins.design as tdd
import tidy3d.web as web

Preparation Work Before the PSO

Define Fixed Simulation Settings

Before we can perform the PSO, we need to prepare a few things. First of all, we will define parts that are unchanged in the optimization procedure. The wavelength range of interest in this case is 1500 nm to 1600 nm.

lda0 = 1.55  # central wavelength
freq0 = td.C_0 / lda0  # central frequency
ldas = np.linspace(1.5, 1.6, 31)  # wavelength range
freqs = td.C_0 / ldas  # frequency range
fwidth = 0.5 * (freqs[0] - freqs[-1])  # width of the source frequency range

For simplicity, we use a constant refractive index for silicon and silicon oxide. Dispersive models can certainly be used instead if needed.

n_si = 3.48  # silicon refractive index
si = td.Medium(permittivity=n_si**2)

n_sio2 = 1.44  # silicon oxide refractive index
sio2 = td.Medium(permittivity=n_sio2**2)

Define geometric parameters outside of the optimization region.

Ls = 0.5  # length of each segment
W0 = 0.45  # width of the input waveguide
Wa = 0.2  # width of the taper tip
Wt = 0.45  # width of the taper end
Hc = 0.22  # thickness of the Si layer
G0 = 0.27  # size of the gap
buffer = 8  # buffer spacing
M = 10  # number of segments to be optimized

Next, we define the device geometry outside the optimization region. These geometries will stay unchanged in the optimization process. After each part is defined, we can perform a union operation (+) to combine them together.

# Define the input straight waveguide geometry
input_waveguide_geo = td.Box.from_bounds(rmin=(-buffer, 0, -Hc / 2), rmax=(0, W0, Hc / 2))

# Define the bar port straight waveguide geometry
bar_waveguide_geo = td.Box.from_bounds(
    rmin=((M + 1) * Ls, 0, -Hc / 2), rmax=((M + 1) * Ls + buffer, W0, Hc / 2)
)

# Define the lower waveguide geometry with the help of gdstk
cell = gdstk.Cell("lower_waveguide")
path = gdstk.RobustPath(
    initial_point=(0, -G0 - Wa / 2), width=Wa, tolerance=1e-4, layer=1, datatype=0
)
path.segment(xy=((M + 1) * Ls, -G0 - Wt / 2), width=Wt, offset=-(Wt - Wa))
bend_length = 6
bend_height = 1
path.segment(
    xy=((M + 1) * Ls + bend_length, -G0 - Wt / 2),
    offset=lambda u: bend_height * np.cos(np.pi * (u)) / 2 - bend_height / 2,
)
path.horizontal(x=(M + 1) * Ls + bend_length + buffer)
cell.add(path)
lower_waveguide_geo = td.Geometry.from_gds(cell, gds_layer=1, axis=2, slab_bounds=(-Hc / 2, Hc / 2))

# Perform a union operation to combine all the geometries
unchanged_geo = input_waveguide_geo + bar_waveguide_geo + lower_waveguide_geo

To visually inspect if the above-defined geometries are correct, we can simply use the plot() method. The geometries do look correct. The missing design region will be optimized by PSO later.

ax = unchanged_geo.plot(z=0)
ax.set_xlim(-2, 12)
plt.show()

Define Design Region

Next we define a function that creates the geometry of the design region. This region is parameterized by 10 parameters \(W_1\), \(W_2\), …, \(W_{10}\). We will define an array Ws to store these 10 values. The geometry of the design region is defined as a PolySlab, whose vertices can be easily calculated from Ws.

def define_optimize_region(Ws: np.array) -> td.PolySlab:
    """Calculate the vertices of the design region and return a PolySlab."""
    vertices = [(0, 0), (0, W0)]
    for i, Wi in enumerate(Ws):
        vertices.append(((i + 1) * Ls, Wi))

    vertices.append(((M + 1) * Ls, W0))
    vertices.append(((M + 1) * Ls, 0))

    optimize_region_geo = td.PolySlab(vertices=vertices, axis=2, slab_bounds=(-Hc / 2, Hc / 2))

    return optimize_region_geo

Again, we want to visually inspect if the define_optimize_region function works correctly. To do so, we define Ws with some random values and plot the design region geometry together with the unchanged region geometries. From the plot, we can confirm that the geometries are generated correctly.

np.random.seed(1)  # use a fixed random seed for reproducibility of the notebook
Ws = np.random.uniform(0.3, 0.45, M)
optimize_region_geo = define_optimize_region(Ws)

device = td.Structure(geometry=unchanged_geo + optimize_region_geo, medium=si)
ax = device.plot(z=0)
ax.set_xlim(-1, 12)
plt.show()

Define Source and Monitors

To characterize the performance of the PBS, we need to perform two simulations: 1. excite the input waveguide with a TE0 ModeSource and calculate the TE0 mode transmission at the bar port using a ModeMonitor; 2. Excite the input waveguide with a TM0 ModeSource and calculate the TM0 mode transmission at the cross port using a ModeMonitor. Since the PBS is surrounded by SiO\(_2\), there is a symmetry that can be exploited in the \(z\) direction. We can use this symmetry to strategically select TE0 mode or the TM0 mode at the ModeSource (more on it later). Therefore, here we set num_modes=1 in the ModeSpec. If no symmetry is used, num_modes=2 should be used.

# Add a mode source as excitation
mode_spec = td.ModeSpec(num_modes=1, target_neff=n_si)
mode_source = td.ModeSource(
    center=(-lda0 / 2, W0 / 2, 0),
    size=(0, 4 * W0, 6 * Hc),
    source_time=td.GaussianPulse(freq0=freq0, fwidth=fwidth),
    direction="+",
    mode_spec=mode_spec,
    mode_index=0,
)

# Add a mode monitor at the bar port
mode_monitor_bar = td.ModeMonitor(
    center=((M + 1) * Ls + bend_length + lda0 / 2, W0 / 2, 0),
    size=mode_source.size,
    freqs=freqs,
    mode_spec=mode_spec,
    name="bar",
)

# Add a mode monitor at the cross port
mode_monitor_cross = td.ModeMonitor(
    center=((M + 1) * Ls + bend_length + lda0 / 2, -G0 - Wt / 2 - bend_height, 0),
    size=mode_source.size,
    freqs=freqs,
    mode_spec=mode_spec,
    name="cross",
)

# Simulation domain box
sim_box = td.Box.from_bounds(rmin=(-1, -3.5, -1), rmax=(12.5, 1.5, 1))

# Simulation run time
run_time = 5e-13

Define Tidy3D Simulation

Next, we define a function make_sim that takes the design parameters and polarization and returns a Tidy3D Simulation. If we want to define a TE0 mode excitation, we set symmetry to (0,0,1) while for TM0 excitation, we have symmetry=(0,0,-1). Note again that we can use symmetry because the PBS is surrounded by silicon oxide. If no cladding is used, this symmetry will be broken and we need to define the excitation mode at the source differently.

def make_sim(Ws: np.array, pol: str) -> td.Simulation:
    """Build a Simulation object from design region widths and desired polarization."""
    optimize_region_geo = define_optimize_region(Ws)

    # Define the structure for the entire PBS
    device = td.Structure(geometry=unchanged_geo + optimize_region_geo, medium=si)

    # Define symmetry according to excitation polarization
    pol_to_symmetry = {"TE": (0, 0, 1), "TM": (0, 0, -1)}

    try:
        symmetry = pol_to_symmetry[pol]
    except KeyError:
        raise ValueError("Polarization can either be TE or TM")

    # Add mode monitor according to excitation polarization
    if pol == "TE":
        monitor = [mode_monitor_bar]
    elif pol == "TM":
        monitor = [mode_monitor_cross]

    # Define simulation
    sim = td.Simulation(
        center=sim_box.center,
        size=sim_box.size,
        grid_spec=td.GridSpec.auto(min_steps_per_wvl=10, wavelength=lda0),
        structures=[device],
        sources=[mode_source],
        monitors=monitor,
        run_time=run_time,
        boundary_spec=td.BoundarySpec.all_sides(boundary=td.PML()),
        medium=sio2,
        symmetry=symmetry,
    )

    return sim

We again inspect the simulation setup by plotting it. For the TM simulation, we see the ModeMonitor at the cross port. For TE simulation, we should see the monitor at the bar port.

sim = make_sim(Ws, "TM")
sim.plot(z=0)
plt.show()

Preparing the DesignSpace

The Design plugin requires a pre and post function that define the simulation(s) to be run and how to postprocess the result.

We can define a short pre function which will create the both TE and TM simulations based on the parameter widths that the PSO will suggest. Outputting this as a dictionary of Simulations will allow the Design plugin to efficiently parallelize this problem and correctly format the task names.

def fn_pre(**params: dict) -> dict[td.Simulation, td.Simulation]:
    """Take parameter widths suggested by the PSO and output a dictionary of Simulations."""
    widths = np.array(list(params.values()))
    sim_te = make_sim(widths, "TE")
    sim_tm = make_sim(widths, "TM")

    return {"TE": sim_te, "TM": sim_tm}

Define the Figure of Merit (FOM)

We need to define a function that acts as post function for our DesignSpace and calculates the FOM from our simulation output. In the case of the PBS, we define the FOM as the sum of \(P_{TE,bar}\) and \(P_{TM,cross}\), where \(P_{TE,bar}\) is the transmission of TE0 mode at the bar port and \(P_{TM,cross}\) is the transmission of TM0 mode at the cross port. In this particular case, we only optimize the transmission at the central frequency. If broadband operation is desired, the FOM can be defined with respect to the entire frequency range.

The input for this function is a dictionary with the same keys as the output of fn_pre. The optimizer replaces the Simulation objects it receives with SimulationData objects, meaning we can easily access the appropriate TE or TM simulation.

Note that this FOM is a maximizing function; all the optimizers in the Design plugin are maximizing by default so the sign does not need to be changed.

def fn_post(sim_data_dict: dict[td.SimulationData, td.SimulationData]) -> float:
    """Calculate the power for TE and TM polarizations in different regions of the PBS."""
    # Extract te power at the bar at the central frequency
    P_TE_bar = (
        np.abs(sim_data_dict["TE"]["bar"].amps.sel(mode_index=0, direction="+", f=freq0)) ** 2
    )

    # Extract tm transmission at cross port at the central frequency
    P_TM_cross = (
        np.abs(sim_data_dict["TM"]["cross"].amps.sel(mode_index=0, direction="+", f=freq0)) ** 2
    )

    return float(P_TE_bar + P_TM_cross)

Performing PSO

With all the preparation work done, we are finally ready to perform the PSO. Defining the hyper-parameters within the Method lets the DesignSpace manage the PSO run using the PySwarms library.

In this optimization, we put an upper and lower bound for the \(W_i\) to be 450 nm and 300 nm. 5 particles are used for a total of 40 iterations. As discussed above, this means the entire optimization will run 400 simulations and cost 10 FlexCredits. Since this notebook is mainly for demonstration purposes, the numbers of particles and iterations are kept small. To really achieve a design with high performance, larger numbers should be used. To ensure the final result is reproducible every time we run the notebook, the initial positions of the particles are fixed with a random seed.

There are three hyperparameters in PSO, namely the inertia weight, the cognitive coefficient, and the social coefficient. Their values can significantly impact the performance of the algorithm. The best values of them depend on the specific problem so it can take some experimentation to determine.

We also include a very low ftol value that must be maintained for 8 consecutive iterations as an early-stop criterion. This means that if the fitness stops improving, the optimization will finish early.

W_max = 0.45  # upper bound
W_min = 0.3  # lower bound

n_particles = 5  # number of particles

# Set initial positions
init_pos = np.random.uniform(W_min, W_max, (n_particles, M))

particle_swarm = tdd.MethodParticleSwarm(
    n_particles=n_particles,
    n_iter=75,
    cognitive_coeff=1,
    social_coeff=1,
    weight=0.7,
    init_pos=init_pos,
    seed=1,
    ftol=1e-4,
    ftol_iter=8,
)

parameters = [tdd.ParameterFloat(name=str(i), span=(W_min, W_max)) for i in range(M)]

design_space = tdd.DesignSpace(
    method=particle_swarm, parameters=parameters, task_name="PSO_Notebook", path_dir="./data"
)

Running the PSO optimization is easily managed by supplying the fn_pre and fn_post functions to the DesignSpace.run method.

results = design_space.run(fn_pre, fn_post, verbose=True)
df = results.to_dataframe()
06:39:48 UTC Running 10 Simulations                                             

06:40:21 UTC Running 10 Simulations                                             

06:41:05 UTC Running 10 Simulations                                             

06:41:56 UTC Running 10 Simulations                                             

06:42:36 UTC Running 10 Simulations                                             

06:44:42 UTC Running 10 Simulations                                             

06:45:14 UTC Running 10 Simulations                                             

06:45:47 UTC Running 10 Simulations                                             

06:46:26 UTC Running 10 Simulations                                             

06:46:59 UTC Running 10 Simulations                                             

06:47:39 UTC Running 10 Simulations                                             

06:49:26 UTC Running 10 Simulations                                             

06:50:11 UTC Best Result: -1.6447994633909997                                   
             Best Parameters: 0: 0.36569514953861787 1: 0.3122289105981696 2:   
             0.4091189543637529 3: 0.4282830732962929 4: 0.3880362943864307 5:  
             0.32239727328670414 6: 0.3373365173571663 7: 0.32336518513529305 8:
             0.34003640122419787 9: 0.44988811607034884

After the optimization is complete, plot the FOM as a function of iteration. We can use the optimizer object from the results to access the cost history.

cost_history = results.optimizer.cost_history
plt.plot(cost_history)
plt.xlabel("Iteration")
plt.ylabel("Cost")
plt.title("Cost history of optimization")
plt.show()

Final Optimized Design

Finally, we will run two simulations again with the optimized parameters. We will add a FieldMonitor to help visualize the mode splitting.

# Get the best parameters from the optimizer
W_opt = results.optimizer.swarm.best_pos

# Define a field monitor to help visualize the field distribution
field_monitor = td.FieldMonitor(
    center=(0, 0, 0), size=(td.inf, td.inf, 0), freqs=[freq0], name="field"
)

# Define simulations with the optimal design
sims = {
    "TE": make_sim(W_opt, "TE").copy(
        update={"monitors": [mode_monitor_bar, mode_monitor_cross, field_monitor]}
    ),
    "TM": make_sim(W_opt, "TM").copy(
        update={"monitors": [mode_monitor_bar, mode_monitor_cross, field_monitor]}
    ),
}

# Define and submit the batch
batch = web.Batch(simulations=sims, verbose=True)
batch_results = batch.run(path_dir="data")






06:50:14 UTC Started working on Batch containing 2 tasks.                       







06:50:16 UTC Maximum FlexCredit cost: 0.050 for the whole batch.                







             Use 'Batch.real_cost()' to get the billed FlexCredit cost after    
             completion.                                                        












06:50:34 UTC Batch complete.                                                    













Plot the field intensity distribution for the TE and TM simulations. As expected, the TE mode mainly transmits to the bar port while the TM mode mainly transmits to the cross port.






fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(8, 2), tight_layout=True)
batch_results["TE"].plot_field("field", "E", "abs^2", ax=ax1)
batch_results["TM"].plot_field("field", "E", "abs^2", ax=ax2)
plt.show()





















Plot the transmission spectra.






# Extract transmission spectra to the bar and cross ports
P_TE_bar = np.abs(batch_results["TE"]["bar"].amps.sel(mode_index=0, direction="+")) ** 2
P_TM_cross = np.abs(batch_results["TM"]["cross"].amps.sel(mode_index=0, direction="+")) ** 2

# Plot the spectra
plt.plot(ldas * 1e3, 10 * np.log10(P_TE_bar), label="TE bar port")
plt.plot(ldas * 1e3, 10 * np.log10(P_TM_cross), label="TM cross port")
plt.legend()
plt.ylabel("Transmission (dB)")
plt.xlabel("Wavelength (nm)")
plt.ylim(-10, 0)
plt.show()

























Closing Remark


The final design in this notebook achieves reasonable performance but is by no means the best design. A more intricate optimization scheme can be used to enhance the final result as demonstrated in the publication.


Using the PSO optimizer from the Design plugin is a convenient plug-and-play tool for managing parallelized PSO runs within Tidy3D. Advanced users may wish to develop their own optimizer that supports advanced options like dynamic hyperparameters.


Due to the large number of simulation runs, a PSO task can take a significant amount of time and FlexCredits. Users should have a good familiarity with Tidy3D before attempting to perform PSO to avoid making mistakes in the optimization and wasting credits. Running a smaller PSO task as a test and practice before a large PSO task is also highly recommended.








          

          

            Example library
          

          
          

          

            

              


                

                  More Examples
                



              Read More
            

            

              

              

                
                
                

                  
                    Particle swarm optimization of a bullseye cavity for quantum emitter light extraction
                    

                    Particle swarm optimization of a bullseye cavity for quantum emitter light extraction
                  
                

                
                

                  
                    Compact on-chip optical switch design optimization using DBS
                    

                    Compact on-chip optical switch design optimization using DBS
                  
                

                
                

                  
                    Optimization of an S-bend with CMA-ES
                    

                    Optimization of an S-bend with CMA-ES
                  
                

                
                
              

              

              

                
No matching results found.

              

            

          

          

          
          

          

            
            
              
              Previous Tutorial
            
            
            
            
              Next Tutorial