Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
117 changes: 117 additions & 0 deletions nevergrad/functions/photonics/ceviche_loss_fn.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,117 @@
import numpy as np
import typing as tp
import mode_converter


first_time_ceviche = True

def ceviche(
x: np.ndarray, benchmark_type: int = 0, discretize=False, wantgrad=False, wantfields=False
) -> tp.Any:
"""
x = 2d or 3d array of scalars
Inputs:
1. benchmark_type = 0 1 2 or 3, depending on which benchmark you want
2. discretize = True if we want to force x to be in {0,1} (just checks sign(x-0.5) )
3. wantgrad = True if we want to know the gradient
4. wantfields = True if we want the fields of the simulation
Returns:
1. the loss (to be minimized)
2. the gradient or none (depending on wantgrad)
3. the fields or none (depending on wantfields
"""
global first_time_ceviche
global model
import autograd # type: ignore
import autograd.numpy as npa # type: ignore
import ceviche_challenges # type: ignore
import autograd # type: ignore

# ceviche_challenges.beam_splitter.prefabs
# ceviche_challenges.mode_converter.prefabs
# ceviche_challenges.waveguide_bend.prefabs
# ceviche_challenges.wdm.prefabs

if first_time_ceviche or x is None:
if benchmark_type == 2:
print("Mode converter benchmark")
design_variable_shape = (345,345)
simulate = mode_converter.mode_converter
else:
print("Benchmark not defined yet")
return None
# elif benchmark_type == 0:
# spec = ceviche_challenges.waveguide_bend.prefabs.waveguide_bend_2umx2um_spec()
# params = ceviche_challenges.waveguide_bend.prefabs.waveguide_bend_sim_params()
# model = ceviche_challenges.waveguide_bend.model.WaveguideBendModel(params, spec)
# elif benchmark_type == 1:
# spec = ceviche_challenges.beam_splitter.prefabs.pico_splitter_spec()
# params = ceviche_challenges.beam_splitter.prefabs.pico_splitter_sim_params()
# model = ceviche_challenges.beam_splitter.model.BeamSplitterModel(params, spec)
# elif benchmark_type == 3:
# spec = ceviche_challenges.wdm.prefabs.wdm_spec()
# params = ceviche_challenges.wdm.prefabs.wdm_sim_params()
# model = ceviche_challenges.wdm.model.WdmModel(params, spec)
if discretize:
# x between 0 and 1 shifted to either 0 or 1
x = np.round(x)
design = x

# if isinstance(x, str) and x == "name":
# return {0: "waveguide-bend", 1: "beam-splitter", 2: "mode-converter", 3: "wdm"}[benchmark_type]
# elif x is None:
# return model.design_variable_shape

assert x.shape == design_variable_shape, f"Expected shape: {design_variable_shape}" # type: ignore

# The model class provides a convenience property, `design_variable_shape`
# which specifies the design shape it expects.

# Construct a loss function, assuming the `model` and `design` from the code
# snippet above are instantiated.

def loss_fn(x, fields_are_needed=False):
"""A simple loss function"""
the_loss, field = simulate(x)

if fields_are_needed:
return the_loss, field
return the_loss

def loss_fn_nograd(x, fields_are_needed=False):
"""A simple loss function"""
the_loss, field = simulate(x)

if fields_are_needed:
return the_loss, field
return the_loss

if wantgrad:
loss_value_npa, grad = autograd.value_and_grad(loss_fn)(design) # type: ignore
else:
loss_value = loss_fn_nograd(design) # type: ignore
# print("DIFF:", loss_value, loss_value_npa)
# loss_value, loss_grad = autograd.value_and_grad(loss_fn)(design) # type: ignore
first_time_ceviche = False
assert not (wantgrad and wantfields)
if wantgrad:
# if loss_value_npa< 0.0:
# print(x, "NP:", loss_fn_nograd(x), "NPA:", loss_fn(x))
# assert False, f"superweird_{loss_fn_nograd(x)}_{loss_fn(x)}"
return loss_value_npa, grad
if wantfields:
loss_value, fields = loss_fn_nograd(design, fields_are_needed=True)
# if loss_value < 0.0:
# print(x, "NP:", loss_fn_nograd(x), "NPA:", loss_fn(x))
# assert False, f"superweird_{loss_fn_nograd(x)}_{loss_fn(x)}"
return loss_value, fields
# if loss_value < 0.0:
# print(x, "NP:", loss_fn_nograd(x), "NPA:", loss_fn(x))
# assert False, f"superweird_{loss_fn_nograd(x)}_{loss_fn(x)}"
return loss_value

if __name__ == "__main__":
design_variable_shape = (345,345)
x = np.random.random(design_variable_shape)
loss = ceviche(x, benchmark_type=2, discretize=False)
print(loss)
241 changes: 241 additions & 0 deletions nevergrad/functions/photonics/functions.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,241 @@
import numpy as np
import scipy.sparse as sp

def addupml2d(ER2,UR2,NPML):
"""
ADDUPML2D Add UPML to a 2D Yee Grid

[ERxx,ERyy,ERzz,URxx,URyy,URzz] = addupml2d(ER2,UR2,NPML)

INPUT ARGUMENTS
================
ER2 Relative Permittivity on 2x Grid
UR2 Relative Permeability on 2x Grid
NPML [NXLO NXHI NYLO NYHI] Size of UPML on 1x Grid

OUTPUT ARGUMENTS
================
ERxx xx Tensor Element for Relative Permittivity
ERyy yy Tensor Element for Relative Permittivity
ERzz zz Tensor Element for Relative Permittivity
URxx xx Tensor Element for Relative Permeability
URyy yy Tensor Element for Relative Permeability
URzz zz Tensor Element for Relative Permeability
"""
#########################################################
## INITIALIZE FUNCTION
#########################################################

# DEFINE PML PARAMETERS
amax = 4
cmax = 1
p = 3

# EXTRACT GRID PARAMETERS
[Nx2,Ny2] = np.shape(ER2)

# EXTRACT PML PARAMETERS
NXLO = 2 * NPML[0]
NXHI = 2 * NPML[1]
NYLO = 2 * NPML[2]
NYHI = 2 * NPML[3]

#########################################################
## CALCULATE PML PARAMETERS
#########################################################

# INITIALIZE PML PARAMETERS TO PROBLEM SPACE
sx = np.ones((Nx2,Ny2), dtype=complex)
sy = np.ones((Nx2,Ny2), dtype=complex)

# ADD XLO PML
for nx in range(1, NXLO+1):
ax = 1 + (amax - 1) * (nx/NXLO) ** p
cx = cmax * np.sin(0.5 * np.pi * nx/NXLO) ** 2
sx[NXLO - nx,:] = ax * (1 - 1j * 60 * cx)


# ADD XHI PML
for nx in range(1, NXHI+1):
ax = 1 + (amax - 1) * (nx/NXHI) ** p
cx = cmax * np.sin(0.5 * np.pi * nx/NXHI) ** 2
sx[Nx2 - NXHI + nx - 1,:] = ax * (1 - 1j * 60 * cx)


# ADD YLO PML
for ny in range(1, NYLO+1):
ay = 1 + (amax - 1) * (ny/NYLO) ** p
cy = cmax * np.sin(0.5 * np.pi * ny/NYLO) ** 2
sy[:,NYLO - ny] = ay * (1 - 1j * 60 * cy)


# ADD YHI PML
for ny in range(1, NYHI+1):
ay = 1 + (amax - 1) * (ny/NYHI) ** p
cy = cmax * np.sin(0.5 * np.pi * ny/NYHI) ** 2
sy[:,Ny2 - NYHI + ny - 1] = ay * (1 - 1j * 60 * cy)


#########################################################
## INCORPORATE PML
#########################################################

# CALCULATE TENSOR ELEMENTS WITH UPML
ERxx = ER2/sx * sy
ERyy = ER2 * sx/sy
ERzz = ER2 * sx * sy

URxx = UR2/sx * sy
URyy = UR2 * sx/sy
URzz = UR2 * sx * sy

# EXTRACT TENSOR ELEMENTS ON YEE GRID
ERxx = ERxx[1:Nx2:2, 0:Ny2:2]
ERyy = ERyy[0:Nx2:2, 1:Ny2:2]
ERzz = ERzz[0:Nx2:2, 0:Ny2:2]

URxx = URxx[0:Nx2:2, 1:Ny2:2]
URyy = URyy[1:Nx2:2, 0:Ny2:2]
URzz = URzz[1:Nx2:2, 1:Ny2:2]

# DIAGONALIZE MATERIAL TENSORS
inv_ERxx = sp.diags_array(1/ERxx.flatten('F'), format="csc")
inv_ERyy = sp.diags_array(1/ERyy.flatten('F'), format="csc")
ERzz = sp.diags_array(ERzz.flatten('F'), format="csc")
inv_URxx = sp.diags_array(1/URxx.flatten('F'), format="csc")
inv_URyy = sp.diags_array(1/URyy.flatten('F'), format="csc")
URzz = sp.diags_array(URzz.flatten('F'), format="csc")

return [inv_ERxx,inv_ERyy,ERzz,inv_URxx,inv_URyy,URzz]

def yeeder2d(NS,RES,BC,kinc=0):
"""
YEEDER2D Derivative Matrices on a 2D Yee Grid

[DEX,DEY,DHX,DHY] = yeeder2d(NS,RES,BC,kinc)

INPUT ARGUMENTS
=================
NS [Nx Ny] Grid Size
RES [dx dy] Grid Resolution
BC [xbc ybc] Boundary Conditions
0: Dirichlet boundary conditions
1: Periodic boundary conditions
kinc [kx ky] Incident Wave Vector
This argument is only needed for PBCs.

Note: For normalized grids use k0 * RES and kinc/k0

OUTPUT ARGUMENTS
=================
DEX Derivative Matrix wrt x for Electric Fields
DEY Derivative Matrix wrt to y for Electric Fields
DHX Derivative Matrix wrt to x for Magnetic Fields
DHY Derivative Matrix wrt to y for Magnetic Fields
"""

#########################################################
## HANDLE INPUT ARGUMENTS
#########################################################

# EXTRACT GRID PARAMETERS
Nx = NS[0]
dx = RES[0]
Ny = NS[1]
dy = RES[1]

# DEFAULT KINC
if not(kinc):
kinc = [0, 0]


# DETERMINE MATRIX SIZE
M = Nx * Ny

# ZERO MATRIX
Z = sp.csc_array((M,M))

#########################################################
## BUILD DEX
#########################################################

# HANDLE IF SIZE IS 1 CELL
if Nx==1:
DEX = -1j * kinc[0] * sp.eye_array(M)

# HANDLE ALL OTHER CASES
else:

# Center Diagonal
d0 = -np.ones(M)

# Upper Diagonal
d1 = np.ones(M-1)
d1[Nx-1:M:Nx] = 0

# Build Derivative Matrix with Dirichlet BCs
DEX = Z + sp.diags_array(d0/dx)
DEX = DEX + sp.diags_array(d1/dx, offsets=1)

# Incorporate Periodic Boundary Conditions
if BC[0]==1:
d1 = sp.csc_array(M)
d1[:M:Nx] = np.exp(-1j * kinc[0] * Nx * dx)/dx
DEX = DEX + sp.diags_array(d1, offsets=1-Nx)




#########################################################
## BUILD DEY
#########################################################

# HANDLE IF SIZE IS 1 CELL
if Ny==1:
DEY = -1j * kinc[1] * sp.eye_array(M)

# HANDLE ALL OTHER CASES
else:

# Center Diagonal
d0 = -np.ones((M))

# Upper Diagonal
d1 = np.ones((M-Nx))

# Build Derivative Matrix with Dirichlet BCs
DEY = Z + sp.diags_array(d0/dy)
DEY = DEY + sp.diags_array(d1/dy, offsets=Nx)

# Incorporate Periodic Boundary Conditions
if BC[1]==1:
d1 = (np.exp(-1j * kinc[1] * Ny * dy)/dy) * np.ones((M,1))
DEY = DEY + sp.diags_array(d1, offsets=Nx-M)




#########################################################
## BUILD DHX AND DHY
#########################################################

DHX = -np.conj(DEX.T)
DHY = -np.conj(DEY.T)
return [DEX,DEY,DHX,DHY]


def block(xa2,ya2,ER2,pos_x,pos_y,Len_x,Len_y,n):
dx2 = xa2[1]-xa2[0]
dy2 = ya2[1]-ya2[0]
# MAKE n-index block whose corner left-down is at pos_x and pos_y
# of size Len_x X Len_y
# refractive index n

ind_X = np.argmin(np.abs(xa2-pos_x))
ind_Y = np.argmin(np.abs(ya2-pos_y))

nx = int(ind_X + round(Len_x/dx2))
ny = int(ind_Y + round(Len_y/dy2))

ER2 [ind_X:nx, ind_Y:ny] = n ** 2
return ER2
Loading