<commit_before>import argparse

import numpy as np
import autograd.numpy as npa
import matplotlib.pylab as plt

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

import ceviche
from ceviche import fdfd_ez, fdfd_ez_nl, jacobian
from ceviche.optimizers import adam_optimize
from ceviche.utils import imarr, get_value
from ceviche.modes import get_modes

import collections
# Create a container for our slice coords to be used for sources and probes
Slice = collections.namedtuple('Slice', 'x y')

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

    space       : The space between the PML and the structure
    wg_width    : The feed and probe waveguide width
    space_slice : The added space for the probe and source slices
    """
    rho = np.zeros((Nx, Ny))  
    design_region = np.zeros((Nx, Ny))
    
    # Input probe slice
    input_slice = Slice(x=np.array(Npml+1), y=np.arange(int(Ny/2-wg_width/2-space_slice), int(Ny/2+wg_width/2+space_slice)))
    
    output_slices = []
    # Input waveguide
    rho[0:int(Npml+space),int(Ny/2-wg_width/2):int(Ny/2+wg_width/2)] = 1
    
    # Output waveguide 1
    rho[int(Nx-Npml-space)::,int(Npml+space):int(Npml+space+wg_width)] = 1
    output_slices.append(Slice(x=np.array(Nx-Npml-1), y=np.arange(int(Npml+space-space_slice), int(Npml+space+wg_width+space_slice))))
    
    # Output waveguide 2
    rho[int(Nx-Npml-space)::,int(Ny/2-wg_width/2):int(Ny/2+wg_width/2)] = 1
    output_slices.append(Slice(x=np.array(Nx-Npml-1), y=np.arange(int(Ny/2-wg_width/2-space_slice), int(Ny/2+wg_width/2+space_slice))))
    
    # Output waveguide 3
    rho[int(Nx-Npml-space)::,int(Ny-Npml-space-wg_width):int(Ny-Npml-space)] = 1
    output_slices.append(Slice(x=np.array(Nx-Npml-1), y=np.arange(int(Ny-Npml-space-wg_width-space_slice), int(Ny-Npml-space+space_slice))))
    
    design_region[Npml+space:Nx-Npml-space, Npml+space:Ny-Npml-space] = 1
    rho[Npml+space:Nx-Npml-space, Npml+space:Ny-Npml-space] = 0.5

    return rho, design_region, input_slice, output_slices

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

def operator_blur(rho, radius=2):
    """Blur operator implemented via two-dimensional convolution
    """
    rr, cc = circle(radius, radius, radius+1)
    kernel = np.zeros((2*radius+1, 2*radius+1), dtype=np.float)
    kernel[rr, cc] = 1
    kernel=kernel/kernel.sum()
    # For whatever reason HIPS autograd doesn't support 'same' mode, so we need to manually crop the output
    return conv(rho, kernel, mode='full')[radius:-radius,radius:-radius]

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

parser = argparse.ArgumentParser()
parser.add_argument('-Nsteps', default=20, type=int)
parser.add_argument('-step_size', default=1e-2, type=float)
parser.add_argument('-omega', default=2*np.pi*200e12, type=float)
parser.add_argument('-dx', default=50e-9, type=float)
parser.add_argument('-Nx', default=100, type=int)
parser.add_argument('-Ny', default=130, type=int)
parser.add_argument('-Npml', default=20, type=int)
parser.add_argument('-epsr_min', default=1.0, type=float)
parser.add_argument('-epsr_max', default=12.0, type=float)
parser.add_argument('-blur_radius', default=3, type=int)
parser.add_argument('-beta', default=100.0, type=float)
parser.add_argument('-eta', default=0.5, type=float)
parser.add_argument('-space', default=10, type=int)
parser.add_argument('-wg_width', default=10, type=int)
parser.add_argument('-space_slice', default=7, type=int)
parser.add_argument('-objective_port_ind', default=0, type=int)
args = parser.parse_args()

Nsteps = args.Nsteps
step_size = args.step_size
omega = args.omega
dx = args.dx
Nx = args.Nx
Ny = args.Ny
Npml = args.Npml
epsr_min = args.epsr_min
epsr_max = args.epsr_max
blur_radius = args.blur_radius
beta = args.beta
eta = args.eta
space = args.space
wg_width = args.wg_width
space_slice = args.space_slice
objective_port_ind = args.objective_port_ind

# Setup initial structure
rho_init, design_region, input_slice, output_slices = init_domain(Nx, Ny, Npml, space=space, wg_width=wg_width, space_slice=space_slice)
epsr = epsr_min + (epsr_max-epsr_min) * make_rho(rho_init, design_region, radius=blur_radius)

# Setup source
source = np.zeros((Nx, Ny), dtype=np.complex)
eps_cross = epsr[input_slice.x, input_slice.y]
_, mode_in = get_modes(eps_cross, omega, dx, 0, m=1, filtering=False)
source = np.zeros((Nx, Ny), dtype=np.complex)
source[input_slice.x, input_slice.y] = mode_in.squeeze()

# Setup probes
probes = []
for output_slice in output_slices:
    wg_cross = epsr[output_slice.x, output_slice.y]
    _, field = get_modes(wg_cross, omega, dx, 0, m=1, filtering=False)
    probe = np.zeros((Nx, Ny), dtype=np.complex)
    probe[output_slice.x, output_slice.y] = field.squeeze()
    probes.append(probe)

# Simulate initial device
simulation = fdfd_ez(omega, dx, epsr, [Npml, Npml])
Hx, Hy, Ez = simulation.solve(source)<commit_msg>Update example to use insert_mode<commit_after>import argparse

import numpy as np
import autograd.numpy as npa
import matplotlib.pylab as plt

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

from ceviche import fdfd_ez, fdfd_ez_nl, jacobian
from ceviche.optimizers import adam_optimize
from ceviche.utils import imarr, get_value
from ceviche.modes import insert_mode

import collections
# Create a container for our slice coords to be used for sources and probes
Slice = collections.namedtuple('Slice', 'x y')

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

    space       : The space between the PML and the structure
    wg_width    : The feed and probe waveguide width
    space_slice : The added space for the probe and source slices
    """
    rho = np.zeros((Nx, Ny))  
    design_region = np.zeros((Nx, Ny))
    
    # Input probe slice
    input_slice = Slice(x=np.array(Npml+1), y=np.arange(int(Ny/2-wg_width/2-space_slice), int(Ny/2+wg_width/2+space_slice)))
    
    output_slices = []
    # Input waveguide
    rho[0:int(Npml+space),int(Ny/2-wg_width/2):int(Ny/2+wg_width/2)] = 1
    
    # Output waveguide 1
    rho[int(Nx-Npml-space)::,int(Npml+space):int(Npml+space+wg_width)] = 1
    output_slices.append(Slice(x=np.array(Nx-Npml-1), y=np.arange(int(Npml+space-space_slice), int(Npml+space+wg_width+space_slice))))
    
    # Output waveguide 2
    rho[int(Nx-Npml-space)::,int(Ny/2-wg_width/2):int(Ny/2+wg_width/2)] = 1
    output_slices.append(Slice(x=np.array(Nx-Npml-1), y=np.arange(int(Ny/2-wg_width/2-space_slice), int(Ny/2+wg_width/2+space_slice))))
    
    # Output waveguide 3
    rho[int(Nx-Npml-space)::,int(Ny-Npml-space-wg_width):int(Ny-Npml-space)] = 1
    output_slices.append(Slice(x=np.array(Nx-Npml-1), y=np.arange(int(Ny-Npml-space-wg_width-space_slice), int(Ny-Npml-space+space_slice))))
    
    design_region[Npml+space:Nx-Npml-space, Npml+space:Ny-Npml-space] = 1
    rho[Npml+space:Nx-Npml-space, Npml+space:Ny-Npml-space] = 0.5

    return rho, design_region, input_slice, output_slices

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

def operator_blur(rho, radius=2):
    """Blur operator implemented via two-dimensional convolution
    """
    rr, cc = circle(radius, radius, radius+1)
    kernel = np.zeros((2*radius+1, 2*radius+1), dtype=np.float)
    kernel[rr, cc] = 1
    kernel=kernel/kernel.sum()
    # For whatever reason HIPS autograd doesn't support 'same' mode, so we need to manually crop the output
    return conv(rho, kernel, mode='full')[radius:-radius,radius:-radius]

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

parser = argparse.ArgumentParser()
parser.add_argument('-Nsteps', default=20, type=int)
parser.add_argument('-step_size', default=1e-2, type=float)
parser.add_argument('-omega', default=2*np.pi*200e12, type=float)
parser.add_argument('-dx', default=50e-9, type=float)
parser.add_argument('-Nx', default=100, type=int)
parser.add_argument('-Ny', default=130, type=int)
parser.add_argument('-Npml', default=20, type=int)
parser.add_argument('-epsr_min', default=1.0, type=float)
parser.add_argument('-epsr_max', default=12.0, type=float)
parser.add_argument('-blur_radius', default=3, type=int)
parser.add_argument('-beta', default=100.0, type=float)
parser.add_argument('-eta', default=0.5, type=float)
parser.add_argument('-space', default=10, type=int)
parser.add_argument('-wg_width', default=10, type=int)
parser.add_argument('-space_slice', default=7, type=int)
parser.add_argument('-objective_port_ind', default=0, type=int)
args = parser.parse_args()

Nsteps = args.Nsteps
step_size = args.step_size
omega = args.omega
dx = args.dx
Nx = args.Nx
Ny = args.Ny
Npml = args.Npml
epsr_min = args.epsr_min
epsr_max = args.epsr_max
blur_radius = args.blur_radius
beta = args.beta
eta = args.eta
space = args.space
wg_width = args.wg_width
space_slice = args.space_slice
objective_port_ind = args.objective_port_ind

# Setup initial structure
rho_init, design_region, input_slice, output_slices = init_domain(Nx, Ny, Npml, space=space, wg_width=wg_width, space_slice=space_slice)
epsr = epsr_min + (epsr_max-epsr_min) * make_rho(rho_init, design_region, radius=blur_radius)

# Setup source
source = insert_mode(omega, dx, input_slice.x, input_slice.y, epsr)

# Setup probes
probes = []
for output_slice in output_slices:
    probe = insert_mode(omega, dx, output_slice.x, output_slice.y, epsr)
    probes.append(probe)

# Simulate initial device
simulation = fdfd_ez(omega, dx, epsr, [Npml, Npml])
Hx, Hy, Ez = simulation.solve(source)
{'fraction_non_alphanumeric': 0.1396561715279119, 'fraction_numerical': 0.01777090979331659, 'mean_word_length': 4.971741637831603, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5806527', 'n_tokens_mistral': 4157, 'n_tokens_neox': 3986, 'n_words': 787}