Some optimization for CNOT gate circuit

cnot_ceviche.py

import ceviche
import matplotlib.pyplot as plt
import pandas as pd

from skimage.draw import disk as circle
from autograd.scipy.signal import convolve as conv

from scipy.optimize import minimize
from autograd import grad
import autograd.numpy as anp
import numpy as np


USE_ADAM = 1
USE_SCIPY = ~USE_ADAM

USE_CUPY = 0
USE_NUMPY = ~USE_CUPY

if USE_CUPY:
    import cupy as cp
else:
    import numpy as cp


# Define the simulation domain
Nx, Ny = 300, 300  # Size of the simulation domain
dx = 10e-9  # Spatial step size (10 nm)
epsr_min = 1.0
epsr_max = 12.0
Npml = 10

# Define waveguides and control/target regions with higher permittivity
waveguide_width = 10
control_waveguide_y = 60
target_waveguide_y = 240
coupler_y = 150
nonlinear_region = (slice(140, 160), slice(130, 170))
ring_resonator_center = (150, 150)
ring_resonator_radius = 20


def init_domain(Nx, Ny, Npml, space=10, wg_width=10):
    """Initializes the domain and design region

    space       : The space between the PML and the structure
    wg_width    : The feed and probe waveguide width
    """
    rho = anp.zeros((Nx, Ny))
    design_region = anp.zeros((Nx, Ny))

    # Input waveguides
    rho[Npml:2*Npml, control_waveguide_y-wg_width//2:control_waveguide_y+wg_width//2] = 1
    rho[Npml:2*Npml, target_waveguide_y-wg_width//2:target_waveguide_y+wg_width//2] = 1

    # Output waveguides
    rho[-2*Npml:-Npml, control_waveguide_y-wg_width//2:control_waveguide_y+wg_width//2] = 1
    rho[-2*Npml:-Npml, target_waveguide_y-wg_width//2:target_waveguide_y+wg_width//2] = 1

    design_region[2*Npml:-2*Npml, 2*Npml:-2*Npml] = 1
    rho[2*Npml:-2*Npml, 2*Npml:-2*Npml] = 0.5

    # Create a ring resonator
    for i in range(Nx):
        for j in range(Ny):
            if anp.sqrt((i-ring_resonator_center[0])**2 + (j-ring_resonator_center[1])**2) <= ring_resonator_radius:
                rho[i, j] = 8

    return rho, design_region

def operator_proj(rho, eta=0.5, beta=100):
    """Density projection
    """
    return anp.divide(np.tanh(beta * eta) + np.tanh(beta * (rho - eta)), np.tanh(beta * eta) + np.tanh(beta * (1 - eta)))

def operator_blur(rho, radius=2):
    """Blur operator implemented via two-dimensional convolution using 'full' mode."""
    size = 2 * radius + 1
    kernel = anp.zeros((size, size), dtype=anp.float64)
    rr, cc = circle((radius, radius), radius, shape=kernel.shape)

    kernel[rr, cc] = 1  # Use standard indexing
    kernel = kernel / kernel.sum()

    # Perform 'full' convolution
    convolved = conv(rho, kernel, mode='full')

    # Calculate the starting indices to slice the 'same' size output
    pad = radius
    start_i = pad
    end_i = start_i + rho.shape[0]
    start_j = pad
    end_j = start_j + rho.shape[1]

    # Slice the convolved array to match the input size
    return convolved[start_i:end_i, start_j:end_j]

def make_rho(rho, design_region, radius=2):
    """Helper function for applying the blur to only the design region
    """
    lpf_rho = operator_blur(rho, radius=radius) * design_region
    bg_rho = rho * (design_region==0).astype(np.float64)
    return bg_rho + lpf_rho

# Initialize the domain
rho_init, design_region = init_domain(Nx, Ny, Npml)
eps_r = epsr_min + (epsr_max - epsr_min) * make_rho(rho_init, design_region, radius=3)


# Visualize the photonic crystal structure
if USE_CUPY:
    plt.imshow(cp.asnumpy(eps_r.T), cmap='RdBu', origin='lower')
else:
    plt.imshow(eps_r.T, cmap='RdBu', origin='lower')
plt.colorbar(label='Relative Permittivity')
plt.title('Enhanced Photonic Crystal Structure for Emulating CNOT Gate')
plt.show()

# Create a simulation object
omega = 2 * cp.pi * 200e12  # Angular frequency of the source (200 THz)
npml = [Npml, Npml]  # Thickness of the perfectly matched layer (PML) in x and y directions

# Function to create the simulation object
def create_simulation(eps_r):
    if USE_CUPY:
        return ceviche.fdfd_ez(omega, dx, cp.asnumpy(eps_r), npml)
    else:
        return ceviche.fdfd_ez(omega, dx, eps_r, npml)

simulation = create_simulation(eps_r)

# Define sources for control and target modes
source_control = cp.zeros((Nx, Ny))
source_target = cp.zeros((Nx, Ny))
source_control[Npml+5, control_waveguide_y] = 1  # Source for control mode
source_target[Npml+5, target_waveguide_y] = 1  # Source for target mode

# Nonlinear Kerr effect function (refined model)
def kerr_effect(ez, eps_r, n2):
    """Calculates the change in permittivity due to the Kerr effect."""
    intensity = anp.abs(ez)**2
    delta_eps_r = 2 * n2 * intensity * eps_r
    return delta_eps_r



'''
### Explanation:

1. **Adjoint Optimization:** Used ceviche's `adam_optimize` to optimize the permittivity distribution to minimize the phase difference.
2. **Objective Function:** Defined an objective function to calculate the phase difference between the simulated phase shift and the expected phase shift.
3. **Optimization Bounds:** Set lower and upper bounds for the optimization to ensure that the permittivity values remain within a realistic range.
4. **Optimization Run:** Executed the optimization process and updated the permittivity distribution with the optimized values.
'''

# Define n2_array globally
n2_array = np.zeros((Nx, Ny))
n2_array[nonlinear_region] = 1e-20  # Set n2 only in the nonlinear region
n2_array_anp = anp.array(n2_array)

# Define the optimization objective
def objective(rho_flat):
    rho = rho_flat.reshape((Nx, Ny))
    # Apply projection and blurring to rho
    rho_processed = make_rho(rho, design_region, radius=3)
    eps_r_linear = epsr_min + (epsr_max - epsr_min) * rho_processed

    # Compute control field
    sim_control = create_simulation(eps_r_linear)
    Ez_control = sim_control.solve(source_control)[0]

    # Compute delta_eps_r based on control field
    delta_eps_r = kerr_effect(Ez_control, eps_r_linear, n2_array_anp)
    eps_r = eps_r_linear + delta_eps_r

    # Solve for total field
    total_source = source_control + source_target
    sim_total = create_simulation(eps_r)
    Ez = sim_total.solve(total_source)[0]

    # Extract the Ez field at the target output waveguide
    output_slice = (
        slice(-Npml - 20, -Npml),
        slice(target_waveguide_y - waveguide_width // 2, target_waveguide_y + waveguide_width // 2)
    )
    Ez_output_target = Ez[output_slice]
    Ez_output_target_mean = anp.mean(Ez_output_target)

    # Expected phase shift (π for 'Control On, Target On')
    expected_phase_shift = anp.pi

    # Compute the phase difference
    output_phase = anp.angle(Ez_output_target_mean)
    phase_difference = anp.mod(output_phase - expected_phase_shift + anp.pi, 2 * anp.pi) - anp.pi  # Between -π and π
    print(f"phase_difference = {phase_difference}")

    # Objective is the squared phase difference
    objective_value = phase_difference**2

    return objective_value


# Initial guess for the optimization
if USE_CUPY:
    initial_guess = cp.asnumpy(rho_init.flatten())
else:
    initial_guess = rho_init.flatten()

# Define optimization bounds as a list of tuples
bounds = [(1.0, 10.0) for _ in initial_guess]

# Define the gradient of the objective function
objective_grad = ceviche.jacobian(objective, mode='reverse')

if USE_ADAM:
    # Adjoint optimization using ceviche
    from ceviche.optimizers import adam_optimize

    (rho_optimum, loss) = adam_optimize(
        objective,
        initial_guess,
        objective_grad,
        Nsteps=100,
        direction='min',
        step_size=1e-2
    )

    # Compute the optimized eps_r from the optimized rho
    rho_optimum = rho_optimum.reshape((Nx, Ny))
    rho_processed = make_rho(rho_optimum, design_region, radius=3)
    optimized_eps_r = epsr_min + (epsr_max - epsr_min) * rho_processed

elif USE_SCIPY:
    # Run the optimization using scipy.optimize.minimize
    from scipy.optimize import minimize

    result = minimize(objective, initial_guess, jac=objective_grad, bounds=bounds, method='L-BFGS-B')
    optimized_eps_r = cp.array(result.x.reshape((Nx, Ny)))

# Reshape the optimized permittivity back to the original shape
eps_r = optimized_eps_r

# Run the simulation for different combinations of control and target modes
scenarios = [
    ('Control Off, Target Off', cp.zeros((Nx, Ny)), cp.zeros((Nx, Ny))),
    ('Control On, Target Off', source_control, cp.zeros((Nx, Ny))),
    ('Control Off, Target On', cp.zeros((Nx, Ny)), source_target),
    ('Control On, Target On', source_control, source_target)
]

# Placeholder for results
results = []

for scenario_name, source_ctrl, source_tgt in scenarios:
    # Combine sources
    source_combined = source_ctrl + source_tgt

    # Determine if control is on
    control_on = np.any(source_ctrl)

    if control_on:
        # First simulating without `source_tgt` ensures that only the control signal can activate the
        # kerr effect, which induces a permittivity change. Then, simulating the total field with the
        # new permittivity sees how the target signal is affected by the control-induced changes.
        # Compute Kerr effect based on control field
        Ez_control, Hx_control, Hy_control = simulation.solve(source_ctrl)
        delta_eps_r_nonlinear = kerr_effect(Ez_control, optimized_eps_r, n2_array)
        eps_r_nonlinear = optimized_eps_r + delta_eps_r_nonlinear
    else:
        # Control is off, no Kerr effect
        eps_r_nonlinear = optimized_eps_r

    # Re-solve with updated permittivity
    simulation = create_simulation(eps_r_nonlinear)  # Re-initialize with updated permittivity
    if USE_CUPY:
        Ez_updated, Hx_updated, Hy_updated = simulation.solve(cp.asnumpy(source_combined))
    else:
        Ez_updated, Hx_updated, Hy_updated = simulation.solve(source_combined)
    Ez_updated = cp.array(Ez_updated)

    # Visualize the electric field distribution
    if USE_CUPY:
        plt.imshow(cp.asnumpy(cp.real(Ez_updated).T), cmap='RdBu', origin='lower')
    else:
        plt.imshow(cp.real(Ez_updated).T, cmap='RdBu', origin='lower')
    plt.colorbar(label='Ez')
    plt.title(f'Electric Field Distribution ({scenario_name})')
    plt.show()

    # Analyze the electric field distribution in the control and target regions
    Ez_control_mean = Ez_updated[nonlinear_region].mean()
    Ez_target_mean = Ez_updated[nonlinear_region].mean()

    # Extract fields at output waveguides
    # The [perfectly matched layers](https://en.wikipedia.org/wiki/Perfectly_matched_layer) (PML) are the
    # absorbing boundary regions added to the edges of the simulation domain to prevent artificial reflections of
    # electromagnetic waves from the domain boundaries, and the size of the PML is represented by `Npml`.
    # So `slice(-Npml - 20, -Npml)`, which basically gets us the points from `270` to `290` ,
    # i.e., slicing a region of 20 points away from the PML (which in this case is from `290` to `300` since the
    # size of the simulation domain is 300 x 300).
    # `slice(270, 290)` is the slice for `output_waveguide_control` and `output_waveguide_target` as a means to
    # exclude the PML.
    output_waveguide_control = (
        slice(-Npml - 20, -Npml),
        slice(control_waveguide_y - waveguide_width // 2, control_waveguide_y + waveguide_width // 2)
    )

    output_waveguide_target = (
        slice(-Npml - 20, -Npml),
        slice(target_waveguide_y - waveguide_width // 2, target_waveguide_y + waveguide_width // 2)
    )

    Ez_output_control = Ez_updated[output_waveguide_control].mean()
    Ez_output_target = Ez_updated[output_waveguide_target].mean()

    # Calculate the phase shift
    phase_control = cp.angle(Ez_output_control)
    phase_target = cp.angle(Ez_output_target)

    # Store the input field mean for S-parameter calculation
    # this is a region near the input port to measure the average field amplitude
    input_waveguide_slice = (slice(Npml, Npml + 20), slice(None))
    Ez_input_mean = Ez_updated[input_waveguide_slice].mean()

    results.append({
        'scenario': scenario_name,
        'Ez_input_mean': Ez_input_mean,
        'Ez_control_mean': Ez_control_mean,
        'Ez_target_mean': Ez_target_mean,
        'Ez_output_control': Ez_output_control,
        'Ez_output_target': Ez_output_target,
        'phase_control': phase_control,
        'phase_target': phase_target
    })

    print(f"Scenario: {scenario_name}")
    print(f"Mean electric field in nonlinear region (Input Mode): {Ez_input_mean}")
    print(f"Mean electric field in nonlinear region (Control Mode): {Ez_control_mean}")
    print(f"Mean electric field in nonlinear region (Target Mode): {Ez_target_mean}")
    print(f"Mean electric field in output waveguide (Control Mode): {Ez_output_control}")
    print(f"Mean electric field in output waveguide (Target Mode): {Ez_output_target}")
    print(f"Phase shift in control mode: {phase_control}")
    print(f"Phase shift in target mode: {phase_target}")

# Function to calculate S-parameters
def calculate_s_parameters(results):
    s_parameters = {}

    for result in results:
        scenario = result['scenario']
        Ez_output_control = result['Ez_output_control']
        Ez_output_target = result['Ez_output_target']

        # Calculate transmission and reflection coefficients
        S21_control = Ez_output_control / Ez_input_mean if Ez_input_mean != 0 else 0.0
        S21_target = Ez_output_target / Ez_input_mean if Ez_input_mean != 0 else 0.0

        s_parameters[scenario] = {'S21_control': S21_control, 'S21_target': S21_target}

    return s_parameters

# Calculate S-parameters for the scenarios
s_parameters = calculate_s_parameters(results)
print("S-Parameters:")
for scenario, params in s_parameters.items():
    print(f"{scenario}: S21_control = {params['S21_control']}, S21_target = {params['S21_target']}")

# Verify the CNOT gate's truth table by comparing the phase shifts
truth_table_verification = []

for result in results:
    scenario = result['scenario']
    phase_target = result['phase_target']

    # Expected phase shift depends on the scenario
    expected_phase_shift = cp.pi if 'Control On, Target On' in scenario else 0

    # Calculate the phase difference modulo 2*pi
    phase_difference = np.mod(np.abs(phase_target - expected_phase_shift), 2 * np.pi)

    # Adjust if phase difference is greater than pi
    #phase_difference = cp.abs(phase_target - expected_phase_shift)
    # np.angle returns the phase angle of a complex number in the principal value range of (−π,π],
    # so if the actual phase exceeds π or is less than −π, it gets wrapped back into this range.
    if phase_difference > np.pi:
        phase_difference = 2 * np.pi - phase_difference

    truth_table_verification.append({
        'scenario': scenario,
        'phase_target': phase_target,
        'expected_phase_shift': expected_phase_shift,
        'phase_difference': phase_difference,
        'is_correct': phase_difference < 0.1  # Allowing some tolerance
    })

for verification in truth_table_verification:
    print(f"Scenario: {verification['scenario']}")
    print(f"Phase shift in target mode: {verification['phase_target']}")
    print(f"Expected phase shift: {verification['expected_phase_shift']}")
    print(f"Phase difference: {verification['phase_difference']}")
    print(f"Is the result correct? {verification['is_correct']}")

# Display the truth table verification results
truth_table_df = pd.DataFrame(truth_table_verification)
print(truth_table_df)

# Error Analysis
def error_analysis(results):
    errors = []
    for result in results:
        scenario = result['scenario']
        Ez_target_mean = result['Ez_target_mean']
        phase_target = result['phase_target']

        if 'Control Off' in scenario and phase_target != 0:
            errors.append({
            'scenario': scenario,
            'Ez_target_mean': Ez_target_mean,
            'phase_target': phase_target
            })
            return errors

errors = error_analysis(results)
print("Error Analysis:")
print(errors)