# Path: ceviche/fdfd.py
# class fdfd_hz(fdfd):
#     """ FDFD class for linear Ez polarization """
# 
#     def __init__(self, omega, dL, eps_r, npml, bloch_phases=None):
#         super().__init__(omega, dL, eps_r, npml, bloch_phases=bloch_phases)
# 
#     def _grid_average_2d(self, eps_vec):
# 
#         eps_grid = self._vec_to_grid(eps_vec)
#         eps_grid_xx = 1 / 2 * (eps_grid + npa.roll(eps_grid, axis=1, shift=1))
#         eps_grid_yy = 1 / 2 * (eps_grid + npa.roll(eps_grid, axis=0, shift=1))
#         eps_vec_xx = self._grid_to_vec(eps_grid_xx)
#         eps_vec_yy = self._grid_to_vec(eps_grid_yy)
#         eps_vec_xx = eps_vec_xx
#         eps_vec_yy = eps_vec_yy
#         return eps_vec_xx, eps_vec_yy
# 
#     def _make_A(self, eps_vec):
# 
#         eps_vec_xx, eps_vec_yy = self._grid_average_2d(eps_vec)
#         eps_vec_xx_inv = 1 / (eps_vec_xx + 1e-5)  # the 1e-5 is for numerical stability
#         eps_vec_yy_inv = 1 / (eps_vec_yy + 1e-5)  # autograd throws 'divide by zero' errors.
# 
#         indices_diag = npa.vstack((npa.arange(self.N), npa.arange(self.N)))
# 
#         entries_DxEpsy,   indices_DxEpsy   = spsp_mult(self.entries_Dxb, self.indices_Dxb, eps_vec_yy_inv, indices_diag, self.N)
#         entires_DxEpsyDx, indices_DxEpsyDx = spsp_mult(entries_DxEpsy, indices_DxEpsy, self.entries_Dxf, self.indices_Dxf, self.N)
# 
#         entries_DyEpsx,   indices_DyEpsx   = spsp_mult(self.entries_Dyb, self.indices_Dyb, eps_vec_xx_inv, indices_diag, self.N)
#         entires_DyEpsxDy, indices_DyEpsxDy = spsp_mult(entries_DyEpsx, indices_DyEpsx, self.entries_Dyf, self.indices_Dyf, self.N)
# 
#         entries_d = 1 / EPSILON_0 * npa.hstack((entires_DxEpsyDx, entires_DyEpsxDy))
#         indices_d = npa.hstack((indices_DxEpsyDx, indices_DyEpsxDy))
# 
#         entries_diag = MU_0 * self.omega**2 * npa.ones(self.N)
# 
#         entries_a = npa.hstack((entries_d, entries_diag))
#         indices_a = npa.hstack((indices_d, indices_diag))
# 
#         return entries_a, indices_a
# 
#     def _solve_fn(self, eps_vec, entries_a, indices_a, Mz_vec):
# 
#         b_vec = 1j * self.omega * Mz_vec          # needed so fields are SI units
#         Hz_vec = sp_solve(entries_a, indices_a, b_vec)
#         eps_vec_xx, eps_vec_yy = self._grid_average_2d(eps_vec)
# 
#         # strip out the x and y components of E and find the Hz component
#         Ex_vec, Ey_vec = self._Hz_to_Ex_Ey(Hz_vec, eps_vec_xx, eps_vec_yy)
# 
#         return Ex_vec, Ey_vec, Hz_vec
# 
# Path: ceviche/jacobians.py
# def jacobian(fun, argnum=0, mode='reverse', step_size=1e-6):
#     """ Computes jacobian of `fun` with respect to argument number `argnum` using automatic differentiation """
# 
#     if mode == 'reverse':
#         return jacobian_reverse(fun, argnum)
#     elif mode == 'forward':
#         return jacobian_forward(fun, argnum)
#     elif mode == 'numerical':
#         return jacobian_numerical(fun, argnum, step_size=step_size)
#     else:
#         raise ValueError("'mode' kwarg must be either 'reverse' or 'forward' or 'numerical', given {}".format(mode))
# 
# Path: ceviche/utils.py
# def imarr(arr):
#     """ puts array 'arr' into form ready to plot """
#     arr_value = get_value(arr)
#     arr_plot = arr_value.copy()
#     if len(arr.shape) == 3:
#         arr_plot = arr_plot[:,:,0]
#     return np.flipud(arr_plot.T)
# 
# def get_value(x):
#     if type(x) == npa.numpy_boxes.ArrayBox:
#         return x._value
#     else:
#         return x


import unittest
import numpy as np
import matplotlib.pylab as plt
import autograd.numpy as npa
import sys
from scipy.optimize import minimize
from ceviche import fdfd_hz, jacobian
from ceviche.constants import *
from ceviche.utils import imarr, get_value
             'ftol': 2.220446049250313e-09,
             'gtol': 1e-05,
             'eps': 1e-08,
             'maxfun': 15000,
             'maxiter': NIter,
             'iprint': -1,
             'maxls': 20})

def stack(arr, num_periods=1):
    # returns an array that is stacked `num_periods` times along y.  For plotting many periods of the accelerator unit cell.
    arr_orig = get_value(arr).copy()
    arr_big = get_value(arr).copy()
    for i in range(num_periods):
        arr_big = np.hstack([arr_big, arr_orig])
    return arr_big

num_periods = 4

# plot the final permittivity
eps_disp = imarr(stack(get_value(F.eps_r), num_periods=num_periods))

plt.imshow(np.real(eps_disp), cmap='nipy_spectral')
plt.colorbar()
plt.show()

# plot the accelerating fields
fields = F.solve(source)

[Ex, Ey, Hz] = [get_value(f) for f in fields]


# Path: ceviche/fdfd.py
# class fdfd_hz(fdfd):
#     """ FDFD class for linear Ez polarization """
# 
#     def __init__(self, omega, dL, eps_r, npml, bloch_phases=None):
#         super().__init__(omega, dL, eps_r, npml, bloch_phases=bloch_phases)
# 
#     def _grid_average_2d(self, eps_vec):
# 
#         eps_grid = self._vec_to_grid(eps_vec)
#         eps_grid_xx = 1 / 2 * (eps_grid + npa.roll(eps_grid, axis=1, shift=1))
#         eps_grid_yy = 1 / 2 * (eps_grid + npa.roll(eps_grid, axis=0, shift=1))
#         eps_vec_xx = self._grid_to_vec(eps_grid_xx)
#         eps_vec_yy = self._grid_to_vec(eps_grid_yy)
#         eps_vec_xx = eps_vec_xx
#         eps_vec_yy = eps_vec_yy
#         return eps_vec_xx, eps_vec_yy
# 
#     def _make_A(self, eps_vec):
# 
#         eps_vec_xx, eps_vec_yy = self._grid_average_2d(eps_vec)
#         eps_vec_xx_inv = 1 / (eps_vec_xx + 1e-5)  # the 1e-5 is for numerical stability
#         eps_vec_yy_inv = 1 / (eps_vec_yy + 1e-5)  # autograd throws 'divide by zero' errors.
# 
#         indices_diag = npa.vstack((npa.arange(self.N), npa.arange(self.N)))
# 
#         entries_DxEpsy,   indices_DxEpsy   = spsp_mult(self.entries_Dxb, self.indices_Dxb, eps_vec_yy_inv, indices_diag, self.N)
#         entires_DxEpsyDx, indices_DxEpsyDx = spsp_mult(entries_DxEpsy, indices_DxEpsy, self.entries_Dxf, self.indices_Dxf, self.N)
# 
#         entries_DyEpsx,   indices_DyEpsx   = spsp_mult(self.entries_Dyb, self.indices_Dyb, eps_vec_xx_inv, indices_diag, self.N)
#         entires_DyEpsxDy, indices_DyEpsxDy = spsp_mult(entries_DyEpsx, indices_DyEpsx, self.entries_Dyf, self.indices_Dyf, self.N)
# 
#         entries_d = 1 / EPSILON_0 * npa.hstack((entires_DxEpsyDx, entires_DyEpsxDy))
#         indices_d = npa.hstack((indices_DxEpsyDx, indices_DyEpsxDy))
# 
#         entries_diag = MU_0 * self.omega**2 * npa.ones(self.N)
# 
#         entries_a = npa.hstack((entries_d, entries_diag))
#         indices_a = npa.hstack((indices_d, indices_diag))
# 
#         return entries_a, indices_a
# 
#     def _solve_fn(self, eps_vec, entries_a, indices_a, Mz_vec):
# 
#         b_vec = 1j * self.omega * Mz_vec          # needed so fields are SI units
#         Hz_vec = sp_solve(entries_a, indices_a, b_vec)
#         eps_vec_xx, eps_vec_yy = self._grid_average_2d(eps_vec)
# 
#         # strip out the x and y components of E and find the Hz component
#         Ex_vec, Ey_vec = self._Hz_to_Ex_Ey(Hz_vec, eps_vec_xx, eps_vec_yy)
# 
#         return Ex_vec, Ey_vec, Hz_vec
# 
# Path: ceviche/jacobians.py
# def jacobian(fun, argnum=0, mode='reverse', step_size=1e-6):
#     """ Computes jacobian of `fun` with respect to argument number `argnum` using automatic differentiation """
# 
#     if mode == 'reverse':
#         return jacobian_reverse(fun, argnum)
#     elif mode == 'forward':
#         return jacobian_forward(fun, argnum)
#     elif mode == 'numerical':
#         return jacobian_numerical(fun, argnum, step_size=step_size)
#     else:
#         raise ValueError("'mode' kwarg must be either 'reverse' or 'forward' or 'numerical', given {}".format(mode))
# 
# Path: ceviche/utils.py
# def imarr(arr):
#     """ puts array 'arr' into form ready to plot """
#     arr_value = get_value(arr)
#     arr_plot = arr_value.copy()
#     if len(arr.shape) == 3:
#         arr_plot = arr_plot[:,:,0]
#     return np.flipud(arr_plot.T)
# 
# def get_value(x):
#     if type(x) == npa.numpy_boxes.ArrayBox:
#         return x._value
#     else:
#         return x


# Path: examples/optimize_accelerator.py
import unittest
import numpy as np
import matplotlib.pylab as plt
import autograd.numpy as npa
import sys
from scipy.optimize import minimize
from ceviche import fdfd_hz, jacobian
from ceviche.constants import *
from ceviche.utils import imarr, get_value

             'ftol': 2.220446049250313e-09,
             'gtol': 1e-05,
             'eps': 1e-08,
             'maxfun': 15000,
             'maxiter': NIter,
             'iprint': -1,
             'maxls': 20})

def stack(arr, num_periods=1):
    # returns an array that is stacked `num_periods` times along y.  For plotting many periods of the accelerator unit cell.
    arr_orig = get_value(arr).copy()
    arr_big = get_value(arr).copy()
    for i in range(num_periods):
        arr_big = np.hstack([arr_big, arr_orig])
    return arr_big

num_periods = 4

# plot the final permittivity
eps_disp = imarr(stack(get_value(F.eps_r), num_periods=num_periods))

plt.imshow(np.real(eps_disp), cmap='nipy_spectral')
plt.colorbar()
plt.show()

# plot the accelerating fields
fields = F.solve(source)

[Ex, Ey, Hz] = [get_value(f) for f in fields]


plt.imshow(imarr(np.real(stack(Ey, num_periods=num_periods)))  / E0 , cmap='RdBu')