FDFD implemented but theres an autograd error when doing two sparse m…

…atrix multiplies in a row

ceviche/fdfd.py

@@ -2,7 +2,7 @@
 import scipy.sparse as sp
 
 from .constants import *
-from .primitives import sp_solve, sp_mult#, get_entries_indices
+from .primitives import sp_solve, sp_mult, spsp_mult
 from .derivatives import compute_derivative_matrices
 from .utils import get_entries_indices
 
@@ -94,10 +94,10 @@ def _setup_derivatives(self):
         self.entries_Dyb, self.indices_Dyb = get_entries_indices(self.Dyb)
 
         # stores some convenience functions for multiplying derivative matrices by a vector `vec`
-        self.mult_Dxf = lambda vec: sp_mult(self.entries_Dxf, self.indices_Dxf, vec)
-        self.mult_Dxb = lambda vec: sp_mult(self.entries_Dxb, self.indices_Dxb, vec)
-        self.mult_Dyf = lambda vec: sp_mult(self.entries_Dyf, self.indices_Dyf, vec)
-        self.mult_Dyb = lambda vec: sp_mult(self.entries_Dyb, self.indices_Dyb, vec)
+        self.sp_mult_Dxf = lambda vec: sp_mult(self.entries_Dxf, self.indices_Dxf, vec)
+        self.sp_mult_Dxb = lambda vec: sp_mult(self.entries_Dxb, self.indices_Dxb, vec)
+        self.sp_mult_Dyf = lambda vec: sp_mult(self.entries_Dyf, self.indices_Dyf, vec)
+        self.sp_mult_Dyb = lambda vec: sp_mult(self.entries_Dyb, self.indices_Dyb, vec)
 
     def _vec_to_grid(self, vec):
         """ converts a vector quantity into an array of the shape of the FDFD simulation """
@@ -121,30 +121,30 @@ def _default_val(val, default_val=None):
     """ Field conversion functions for 2D.  Function names are self explanatory """
 
     def _Ex_Ey_to_Hz(self, Ex_vec, Ey_vec):
-        return 1 / 1j / MU_0 * (self.mult_Dxb(Ey_vec) - self.mult_Dyb(Ex_vec))
+        return 1 / 1j / MU_0 * (self.sp_mult_Dxb(Ey_vec) - self.sp_mult_Dyb(Ex_vec))
 
     def _Ez_to_Hx(self, Ez_vec):
-        return  1 / 1j / MU_0 * self.mult_Dyb(Ez_vec)
+        return  1 / 1j / MU_0 * self.sp_mult_Dyb(Ez_vec)
 
     def _Ez_to_Hy(self, Ez_vec):
-        return -1 / 1j / MU_0 * self.mult_Dxb(Ez_vec)
+        return -1 / 1j / MU_0 * self.sp_mult_Dxb(Ez_vec)
 
     def _Ex_Ey_to_Hz(self, Ex_vec, Ey_vec):
-        return 1 / 1j / MU_0 * (self.mult_Dxb(Ey_vec) - self.mult_Dyb(Ex_vec))
+        return 1 / 1j / MU_0 * (self.sp_mult_Dxb(Ey_vec) - self.sp_mult_Dyb(Ex_vec))
 
     def _Ez_to_Hx_Hy(self, Ez_vec):
         Hx_vec = self._Ez_to_Hx(Ez_vec)
         Hy_vec = self._Ez_to_Hy(Ez_vec)
         return Hx_vec, Hy_vec
 
     def _Hz_to_Ex(self, Hz_vec, eps_vec_xx):
-        return  1 / 1j / EPSILON_0 / eps_vec_xx * self.mult_Dyf(Hz_vec)
+        return  1 / 1j / EPSILON_0 / eps_vec_xx * self.sp_mult_Dyf(Hz_vec)
 
     def _Hz_to_Ey(self, Hz_vec, eps_vec_yy):
-        return -1 / 1j / EPSILON_0 / eps_vec_yy * self.mult_Dxf(Hz_vec)
+        return -1 / 1j / EPSILON_0 / eps_vec_yy * self.sp_mult_Dxf(Hz_vec)
 
     def _Hx_Hy_to_Ez(self, Hx_vec, Hy_vec, eps_vec_zz):
-        return 1 / 1j / EPSILON_0 / eps_vec_zz * (self.mult_Dxf(Hy_vec) - self.mult_Dyf(Hx_vec))
+        return 1 / 1j / EPSILON_0 / eps_vec_zz * (self.sp_mult_Dxf(Hy_vec) - self.sp_mult_Dyf(Hx_vec))
 
     def _Hz_to_Ex_Ey(self, Hz_vec, eps_vec_xx, eps_vec_yy):
         Ex_vec = self._Hz_to_Ex(Hz_vec, eps_vec_xx)
@@ -197,51 +197,90 @@ def _grid_average(self, eps_vec):
 
     def _make_A(self, eps_vec):
 
-        # notation: C = [[C11, C12], [C21, C22]]
-        C11 = -1 / MU_0 * self.Dyf.dot(self.Dyb) 
-        C22 = -1 / MU_0 * self.Dxf.dot(self.Dxb)
-        C12 =  1 / MU_0 * self.Dyf.dot(self.Dxb)
-        C21 =  1 / MU_0 * self.Dxf.dot(self.Dyb)
+        eps_vec_xx, eps_vec_yy = self._grid_average(eps_vec)
+        eps_vec_xx_inv = 1 / eps_vec_xx
+        eps_vec_yy_inv = 1 / eps_vec_yy
 
-        # get entries and indices
-        entries_c11, indices_c11 = get_entries_indices(C11)
-        entries_c22, indices_c22 = get_entries_indices(C22)
-        entries_c12, indices_c12 = get_entries_indices(C12)
-        entries_c21, indices_c21 = get_entries_indices(C21)
+        arangeN = npa.arange(self.N)
+        indices_diag = npa.vstack((arangeN, arangeN))
 
-        # shift the indices into each of the 4 quadrants
-        indices_c22 += self.N       # shift into bottom right quadrant
-        indices_c12[1,:] += self.N  # shift into top right quadrant
-        indices_c21[0,:] += self.N  # shift into bottom left quadrant
+        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)
 
-        # get full matrix entries and indices
-        entries_c = npa.hstack((entries_c11, entries_c12, entries_c21, entries_c22))
-        indices_c = npa.hstack((indices_c11, indices_c12, indices_c21, indices_c22))
+        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)
 
-        # indices into the diagonal of a sparse matrix
-        eps_vec_xx, eps_vec_yy = self._grid_average(eps_vec)
-        entries_diag = - EPSILON_0 * self.omega**2 * npa.hstack((eps_vec_xx, eps_vec_yy))
-        indices_diag = npa.vstack((npa.arange(2 * self.N), npa.arange(2 * self.N)))
+        entires_DxEpsyDx, indices_DxEpsyDx = entries_DxEpsy, indices_DxEpsy
+        entires_DyEpsxDy, indices_DyEpsxDy = entries_DyEpsx, indices_DyEpsx
+
+        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))
 
-        # put together the big A and return entries and indices
-        entries_a = npa.hstack((entries_diag, entries_c))
-        indices_a = npa.hstack((indices_diag, indices_c))
         return entries_a, indices_a
 
+    # def _make_A_deprecated(self, eps_vec):
+
+    #     # notation: C = [[C11, C12], [C21, C22]]
+    #     C11 = -1 / MU_0 * self.Dyf.dot(self.Dyb) 
+    #     C22 = -1 / MU_0 * self.Dxf.dot(self.Dxb)
+    #     C12 =  1 / MU_0 * self.Dyf.dot(self.Dxb)
+    #     C21 =  1 / MU_0 * self.Dxf.dot(self.Dyb)
+
+    #     # get entries and indices
+    #     entries_c11, indices_c11 = get_entries_indices(C11)
+    #     entries_c22, indices_c22 = get_entries_indices(C22)
+    #     entries_c12, indices_c12 = get_entries_indices(C12)
+    #     entries_c21, indices_c21 = get_entries_indices(C21)
+
+    #     # shift the indices into each of the 4 quadrants
+    #     indices_c22 += self.N       # shift into bottom right quadrant
+    #     indices_c12[1,:] += self.N  # shift into top right quadrant
+    #     indices_c21[0,:] += self.N  # shift into bottom left quadrant
+
+    #     # get full matrix entries and indices
+    #     entries_c = npa.hstack((entries_c11, entries_c12, entries_c21, entries_c22))
+    #     indices_c = npa.hstack((indices_c11, indices_c12, indices_c21, indices_c22))
+
+    #     # indices into the diagonal of a sparse matrix
+    #     eps_vec_xx, eps_vec_yy = self._grid_average(eps_vec)
+    #     entries_diag = - EPSILON_0 * self.omega**2 * npa.hstack((eps_vec_xx, eps_vec_yy))
+    #     indices_diag = npa.vstack((npa.arange(2 * self.N), npa.arange(2 * self.N)))
+
+    #     # put together the big A and return entries and indices
+    #     entries_a = npa.hstack((entries_diag, entries_c))
+    #     indices_a = npa.hstack((indices_diag, indices_c))
+    #     return entries_a, indices_a
+
     def _solve_fn(self, eps_vec, entries_a, indices_a, Mz_vec):
 
-        # convert the Mz current into Jx, Jy
+        Hz_vec = sp_solve(entries_a, indices_a, Mz_vec)
         eps_vec_xx, eps_vec_yy = self._grid_average(eps_vec)
-        Jx_vec, Jy_vec = self._Hz_to_Ex_Ey(Mz_vec, eps_vec_xx, eps_vec_yy)
-
-        # lump the current sources together and solve for electric field
-        source_J_vec = npa.hstack((Jx_vec, Jy_vec))
-        E_vec = sp_solve(entries_a, indices_a, source_J_vec)
 
         # strip out the x and y components of E and find the Hz component
-        Ex_vec = E_vec[:self.N]
-        Ey_vec = E_vec[self.N:]
+        Ex_vec, Ey_vec = self._Hz_to_Ex_Ey(Hz_vec, eps_vec_xx, eps_vec_yy)
         Hz_vec = self._Ex_Ey_to_Hz(Ex_vec, Ey_vec)
 
         return Ex_vec, Ey_vec, Hz_vec
 
+    # def _solve_fn_deprecated(self, eps_vec, entries_a, indices_a, Mz_vec):
+
+    #     # convert the Mz current into Jx, Jy
+    #     eps_vec_xx, eps_vec_yy = self._grid_average(eps_vec)
+    #     Jx_vec, Jy_vec = self._Hz_to_Ex_Ey(Mz_vec, eps_vec_xx, eps_vec_yy)
+
+    #     # lump the current sources together and solve for electric field
+    #     source_J_vec = npa.hstack((Jx_vec, Jy_vec))
+    #     E_vec = sp_solve(entries_a, indices_a, source_J_vec)
+
+    #     # strip out the x and y components of E and find the Hz component
+    #     Ex_vec = E_vec[:self.N]
+    #     Ey_vec = E_vec[self.N:]
+    #     Hz_vec = self._Ex_Ey_to_Hz(Ex_vec, Ey_vec)
+
+    #     return Ex_vec, Ey_vec, Hz_vec
+

ceviche/jacobians.py

@@ -1,4 +1,4 @@
-import autograd.numpy as np
+import autograd.numpy as npa
 
 from autograd.core import make_vjp, make_jvp
 from autograd.wrap_util import unary_to_nary
@@ -32,7 +32,7 @@ def jacobian_reverse(fun, x):
     vjp, ans = make_vjp(fun, x)
     grads = map(vjp, vspace(ans).standard_basis())
     m, n = _jac_shape(x, ans)
-    return np.reshape(np.stack(grads), (n, m))
+    return npa.reshape(npa.stack(grads), (n, m))
 
 
 @unary_to_nary
@@ -42,13 +42,13 @@ def jacobian_forward(fun, x):
     # ans = fun(x)
     val_grad = map(lambda b: jvp(b), vspace(x).standard_basis())
     vals, grads = zip(*val_grad)
-    ans = np.zeros((list(vals)[0].size,))  # fake answer so that dont have to compute it twice
+    ans = npa.zeros((list(vals)[0].size,))  # fake answer so that dont have to compute it twice
     m, n = _jac_shape(x, ans)
     if _iscomplex(x):
-        grads_real = np.array(grads[::2])
-        grads_imag = np.array(grads[1::2])
+        grads_real = npa.array(grads[::2])
+        grads_imag = npa.array(grads[1::2])
         grads = grads_real - 1j * grads_imag
-    return np.reshape(np.stack(grads), (m, n)).T
+    return npa.reshape(npa.stack(grads), (m, n)).T
 
 
 @unary_to_nary
@@ -60,7 +60,7 @@ def jacobian_numerical(fn, x, step_size=1e-7):
     m = in_array.size
     n = out_array.size
     shape = (n, m)
-    jacobian = np.zeros(shape)
+    jacobian = npa.zeros(shape)
 
     for i in range(m):
         input_i = in_array.copy()
@@ -82,8 +82,8 @@ def _jac_shape(x, ans):
 
 def _iscomplex(x):
     """ Checks if x is complex-valued or not """
-    if isinstance(x, np.ndarray):
-        if x.dtype == np.complex128:
+    if isinstance(x, npa.ndarray):
+        if x.dtype == npa.complex128:
             return True
     if isinstance(x, complex):
         return True
@@ -96,15 +96,15 @@ def _iscomplex(x):
 
     N = 3
     M = 2
-    A = np.random.random((N,M))
-    B = np.random.random((N,M))
+    A = npa.random.random((N,M))
+    B = npa.random.random((N,M))
     print('A = \n', A)
 
     def fn(x, b):
         return A @ x + B @ b
 
-    x0 = np.random.random((M,))
-    b0 = np.random.random((M,))    
+    x0 = npa.random.random((M,))
+    b0 = npa.random.random((M,))    
     print('Jac_rev = \n', jacobian(fn, argnum=0, mode='reverse')(x0, b0))
     print('Jac_for = \n', jacobian(fn, argnum=0, mode='forward')(x0, b0))
     print('Jac_num = \n', jacobian(fn, argnum=0, mode='numerical')(x0, b0))

ceviche/primitives.py

@@ -215,7 +215,7 @@ def vjp(v):
 
         # element-wise these two results together and sum over the second index to give the results in the basis of output entries
         combined = D_ika_V.multiply(D_jka_X)
-        return combined.sum(axis=1)
+        return npa.array(combined.sum(axis=1))
 
     return vjp
 
@@ -276,7 +276,7 @@ def grad_spsp_mult_entries_a_forward(g, b_out, entries_a, indices_a, entries_x,
     combined_matrix = DXA.multiply(DGA)
 
     # dot product the result with vector of 1s to sum over one index and convert to a list of entries.
-    combined_entries = combined_matrix.dot(npa.ones((N,)))
+    combined_entries = npa.array(combined_matrix.dot(npa.ones((N,))))
 
     # return the combined entries and indices of all zero (since the indices aren't affected by the initial entries, the gradient is 0)
     indices_0 = npa.zeros((2, Mb))

ceviche/utils.py

@@ -36,7 +36,7 @@ def get_entries_indices(csr_matrix):
 
 def transpose_indices(indices):
     # returns the transposed indices for transpose sparse matrix creation
-    return npa.flip(indices, axis=0)
+   return npa.flip(indices, axis=0)
 
 def block_4(A, B, C, D):
     """ Constructs a big matrix out of four sparse blocks

tests/test_fields_fdfd.py

@@ -24,13 +24,13 @@ def test_Hz(self):
 
         F = fdfd_hz(self.omega, self.dL, self.eps_r, self.npml)
         Ex, Ey, Hz = F.solve(self.source)
-        plot_component = Ex
+        plot_component = Hz
         field_max = np.max(np.abs(plot_component))
         print(field_max)
-        plt.imshow(np.real(plot_component), cmap='RdBu', vmin=-field_max/10, vmax=field_max/10)
+        plt.imshow(np.real(plot_component), cmap='RdBu', vmin=-field_max/100, vmax=field_max/100)
         plt.show()
 
-    def test_Ez(self):
+    def _test_Ez(self):
         print('\ttesting Ez')
 
         F = fdfd_ez(self.omega, self.dL, self.eps_r, self.npml)

tests/test_gradients_fdfd.py

@@ -92,6 +92,7 @@ def J_fdfd(eps_arr):
                  + npa.sum(npa.square(npa.abs(Ex))) \
                  + npa.sum(npa.square(npa.abs(Ey)))
 
+
         grad_autograd_rev = jacobian(J_fdfd, mode='reverse')(self.eps_arr)
         grad_numerical = jacobian(J_fdfd, mode='numerical')(self.eps_arr)
 
@@ -102,7 +103,7 @@ def J_fdfd(eps_arr):
 
         self.check_gradient_error(grad_numerical, grad_autograd_rev)
 
-    def test_Hz_forward(self):
+    def _test_Hz_forward(self):
 
         print('\ttesting forward-mode Hz in FDFD')
 
@@ -130,7 +131,7 @@ def J_fdfd(c):
 
         self.check_gradient_error(grad_numerical, grad_autograd_for)
 
-    def test_Ez_reverse(self):
+    def _test_Ez_reverse(self):
 
         print('\ttesting reverse-mode Ez in FDFD')
 
@@ -160,7 +161,7 @@ def J_fdfd(eps_arr):
 
         self.check_gradient_error(grad_numerical, grad_autograd_rev)
 
-    def test_Ez_forward(self):
+    def _test_Ez_forward(self):
 
         print('\ttesting forward-mode Ez in FDFD')
 

tests/test_primitives.py

@@ -105,6 +105,7 @@ def test_spmut_entries(self):
         def fn_spsp_entries_a(entries):
             # sparse matrix - sparse matrix dot procut as function of entries into first matrix (A)
             entries_c, indices_c = spsp_mult(entries, self.indices_const, self.entries_const, self.indices_const, N=self.N)
+            entries_c, indices_c = spsp_mult(entries_c, indices_c, self.entries_const, self.indices_const, N=self.N)
             x = sp_solve(entries_c, indices_c, self.b_const)
             return self.out_fn(x)
 
@@ -117,7 +118,7 @@ def fn_spsp_entries_a(entries):
         np.testing.assert_almost_equal(grad_rev, grad_true, decimal=DECIMAL, err_msg=self.err_msg('fn_solve_entries', 'reverse'))
         np.testing.assert_almost_equal(grad_for, grad_true, decimal=DECIMAL, err_msg=self.err_msg('fn_solve_entries', 'forward'))
 
-    def test_spmut_entries(self):
+    def test_spmut_entries2(self):
 
         def fn_spsp_entries_x(entries):
             # sparse matrix - sparse matrix dot procut as function of entries into second matrix (X)