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


import unittest
import numpy as np
import matplotlib.pylab as plt
import autograd.numpy as npa
import sys
from ceviche import fdfd_hz, jacobian
from ceviche.utils import imarr
from scipy.optimize import minimize
    eps_r = eps_arr.reshape((Nx, Ny))
    # set the permittivity of the FDFD and solve the fields
    F.eps_r = eps_r
    Ex, Ey, Hz = F.solve(source)

    # compute the gradient and normalize if you want
    I = npa.sum(npa.square(npa.abs(Hz * probe)))
    return -I / I_H0

# define the gradient for autograd
grad_I = jacobian(intensity, mode='reverse')

# initialize the design region with some eps
eps_r[box_region == 1] = eps_max

bounds = [(1, eps_max) if box_region.flatten()[i] == 1 else (1,1) for i in range(eps_r.size)]

minimize(intensity, eps_r, args=(), method='L-BFGS-B', jac=grad_I,
    bounds=bounds, tol=None, callback=None,
    options={'disp': True,
             'maxcor': 10,
             'ftol': 2.220446049250313e-09,
             'gtol': 1e-05,
             'eps': 1e-08,
             'maxfun': 15000,
             'maxiter': 100,
             'iprint': -1,
             'maxls': 20})

# plot the final permittivity

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


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

    eps_r = eps_arr.reshape((Nx, Ny))
    # set the permittivity of the FDFD and solve the fields
    F.eps_r = eps_r
    Ex, Ey, Hz = F.solve(source)

    # compute the gradient and normalize if you want
    I = npa.sum(npa.square(npa.abs(Hz * probe)))
    return -I / I_H0

# define the gradient for autograd
grad_I = jacobian(intensity, mode='reverse')

# initialize the design region with some eps
eps_r[box_region == 1] = eps_max

bounds = [(1, eps_max) if box_region.flatten()[i] == 1 else (1,1) for i in range(eps_r.size)]

minimize(intensity, eps_r, args=(), method='L-BFGS-B', jac=grad_I,
    bounds=bounds, tol=None, callback=None,
    options={'disp': True,
             'maxcor': 10,
             'ftol': 2.220446049250313e-09,
             'gtol': 1e-05,
             'eps': 1e-08,
             'maxfun': 15000,
             'maxiter': 100,
             'iprint': -1,
             'maxls': 20})

# plot the final permittivity

plt.imshow(imarr(F.eps_r), cmap='nipy_spectral')