Particle swarm optimization of a polarization beam splitter#

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 a built-in adjoint optimization plugin. 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.

[1]:

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

import numpy as np
import matplotlib.pyplot as plt

import gdstk

import tidy3d as td
import tidy3d.web as web
import tidy3d.plugins.design as tdd

Preparational 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.

[2]:

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.

[3]:

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.

[4]:

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.

[5]:

# 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.

[6]:

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

../_images/notebooks_ParticleSwarmOptimizedPBS_14_0.png

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.

[7]:

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.

[8]:

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()

../_images/notebooks_ParticleSwarmOptimizedPBS_19_0.png

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.

[9]:

# 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.

[10]:

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.

[11]:

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

../_images/notebooks_ParticleSwarmOptimizedPBS_27_0.png

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.

[12]:

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_{T E, b a r}$ and $P_{T M, c r o s s}$ , where $P_{T E, b a r}$ is the transmission of TE0 mode at the bar port and $P_{T M, c r o s s}$ 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 maximising function; all the optimizers in the Design plugin are maximising by default so the sign does not need to be changed.

[13]:

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 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 criteria. This means if the fitness stops impoving the optimization will finish early.

[14]:

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

n_particles = 5  # number of particles

# Set initial positions to for reproducability of the notebook
np.random.seed(1)
init_pos = np.random.uniform(W_min, W_max, (n_particles, M))

particle_swarm=tdd.MethodParticleSwarm(
    n_particles=n_particles,
    n_iter=40,
    cognitive_coeff=1,
    social_coeff=1,
    weight=0.7,
    init_pos=init_pos,
    seed=1,
    ftol=0.001,
    ftol_iter=8
)

parameters = [tdd.ParameterFloat(name=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.

[15]:

results = design_space.run(fn_pre, fn_post, verbose=True)
df = results.to_dataframe()

2024-08-18 14:58:57,739 - pyswarms.single.global_best - INFO - Optimize for 40 iters with {'c1': 1.0, 'c2': 1.0, 'w': 0.7}
pyswarms.single.global_best:   0%|          |0/40

14:58:57 BST Running 10 Simulations

pyswarms.single.global_best:   2%|▎         |1/40, best_cost=-1.55

14:59:40 BST Running 10 Simulations

pyswarms.single.global_best:   5%|▌         |2/40, best_cost=-1.56

15:00:20 BST Running 10 Simulations

pyswarms.single.global_best:   8%|▊         |3/40, best_cost=-1.56

15:00:57 BST Running 10 Simulations

pyswarms.single.global_best:  10%|█         |4/40, best_cost=-1.56

15:01:38 BST Running 10 Simulations

pyswarms.single.global_best:  12%|█▎        |5/40, best_cost=-1.56

15:02:19 BST Running 10 Simulations

pyswarms.single.global_best:  15%|█▌        |6/40, best_cost=-1.57

15:02:57 BST Running 10 Simulations

pyswarms.single.global_best:  18%|█▊        |7/40, best_cost=-1.57

15:03:34 BST Running 10 Simulations

pyswarms.single.global_best:  20%|██        |8/40, best_cost=-1.6

15:04:14 BST Running 10 Simulations

pyswarms.single.global_best:  22%|██▎       |9/40, best_cost=-1.6

15:04:53 BST Running 10 Simulations

pyswarms.single.global_best:  25%|██▌       |10/40, best_cost=-1.6

15:05:32 BST Running 10 Simulations

pyswarms.single.global_best:  28%|██▊       |11/40, best_cost=-1.6

15:06:16 BST Running 10 Simulations

pyswarms.single.global_best:  30%|███       |12/40, best_cost=-1.6

15:06:54 BST Running 10 Simulations

pyswarms.single.global_best:  32%|███▎      |13/40, best_cost=-1.6

15:07:36 BST Running 10 Simulations

pyswarms.single.global_best:  35%|███▌      |14/40, best_cost=-1.6

15:08:18 BST Running 10 Simulations

pyswarms.single.global_best:  38%|███▊      |15/40, best_cost=-1.65

15:08:57 BST Running 10 Simulations

pyswarms.single.global_best:  40%|████      |16/40, best_cost=-1.65

15:09:36 BST Running 10 Simulations

pyswarms.single.global_best:  42%|████▎     |17/40, best_cost=-1.65

15:10:16 BST Running 10 Simulations

pyswarms.single.global_best:  45%|████▌     |18/40, best_cost=-1.65

15:10:56 BST Running 10 Simulations

pyswarms.single.global_best:  48%|████▊     |19/40, best_cost=-1.65

15:11:34 BST Running 10 Simulations

pyswarms.single.global_best:  50%|█████     |20/40, best_cost=-1.66

15:12:11 BST Running 10 Simulations

pyswarms.single.global_best:  52%|█████▎    |21/40, best_cost=-1.66

15:12:53 BST Running 10 Simulations

pyswarms.single.global_best:  55%|█████▌    |22/40, best_cost=-1.66

15:13:31 BST Running 10 Simulations

pyswarms.single.global_best:  57%|█████▊    |23/40, best_cost=-1.66

15:14:09 BST Running 10 Simulations

pyswarms.single.global_best:  60%|██████    |24/40, best_cost=-1.66

15:14:49 BST Running 10 Simulations

pyswarms.single.global_best:  62%|██████▎   |25/40, best_cost=-1.66

15:15:29 BST Running 10 Simulations

pyswarms.single.global_best:  65%|██████▌   |26/40, best_cost=-1.66

15:16:08 BST Running 10 Simulations

pyswarms.single.global_best:  68%|██████▊   |27/40, best_cost=-1.76

15:16:46 BST Running 10 Simulations

pyswarms.single.global_best:  70%|███████   |28/40, best_cost=-1.76

15:17:24 BST Running 10 Simulations

pyswarms.single.global_best:  72%|███████▎  |29/40, best_cost=-1.76

15:18:02 BST Running 10 Simulations

pyswarms.single.global_best:  75%|███████▌  |30/40, best_cost=-1.76

15:18:44 BST Running 10 Simulations

pyswarms.single.global_best:  78%|███████▊  |31/40, best_cost=-1.76

15:19:25 BST Running 10 Simulations

pyswarms.single.global_best:  80%|████████  |32/40, best_cost=-1.76

15:20:04 BST Running 10 Simulations

pyswarms.single.global_best:  82%|████████▎ |33/40, best_cost=-1.76

15:20:47 BST Running 10 Simulations

pyswarms.single.global_best:  85%|████████▌ |34/40, best_cost=-1.76

15:21:25 BST Running 10 Simulations

pyswarms.single.global_best:  85%|████████▌ |34/40, best_cost=-1.76
2024-08-18 15:22:05,753 - pyswarms.single.global_best - INFO - Optimization finished | best cost: -1.7550373095072005, best pos: [0.36448035 0.37247524 0.38638816 0.40732065 0.4468407  0.327726
 0.3097584  0.31346851 0.32078258 0.34627143]

15:22:05 BST Best Result: -1.7550373095072005
             Best Parameters: 0: 0.3644803547286112 1: 0.37247524188920256 2:
             0.38638815979366026 3: 0.40732064818082 4: 0.4468407035432923 5:
             0.3277260025698511 6: 0.30975839785286563 7: 0.3134685087443462 8:
             0.32078258493519357 9: 0.34627143329045607

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.

[16]:

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

../_images/notebooks_ParticleSwarmOptimizedPBS_40_0.png

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.

[19]:

# 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")

12:59:47 BST Started working on Batch containing 2 tasks.

12:59:56 BST Maximum FlexCredit cost: 0.050 for the whole batch.

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

13:00:17 BST 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.

[20]:

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()

../_images/notebooks_ParticleSwarmOptimizedPBS_45_0.png

Plot the transmission spectra.

[21]:

# 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()

../_images/notebooks_ParticleSwarmOptimizedPBS_47_0.png

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.

Particle swarm optimization of a polarization beam splitter

Contents