From a4616982ca5d986a8a12e00a86f7e87e686a20ee Mon Sep 17 00:00:00 2001
From: Jan Petykiewicz <anewusername@gmail.com>
Date: Sun, 24 Sep 2017 22:28:39 -0700
Subject: [PATCH] Add cylindrical coordinate 2D modesolver code

---
 fdfd_tools/waveguide.py      | 59 ++++++++++++++++++++++++++++++++
 fdfd_tools/waveguide_mode.py | 66 ++++++++++++++++++++++++++++++++++++
 2 files changed, 125 insertions(+)

diff --git a/fdfd_tools/waveguide.py b/fdfd_tools/waveguide.py
index 0725bac..85e96bf 100644
--- a/fdfd_tools/waveguide.py
+++ b/fdfd_tools/waveguide.py
@@ -307,3 +307,62 @@ def e_err(e: vfield_t,
         op = ch @ mu_inv @ ce @ e - omega ** 2 * (epsilon * e)
 
     return norm(op) / norm(e)
+
+
+def cylindrical_operator(omega: complex,
+                         dxes: dx_lists_t,
+                         epsilon: vfield_t,
+                         r0: float,
+                         ) -> sparse.spmatrix:
+    """
+    Cylindrical coordinate waveguide operator of the form
+
+    TODO
+
+    for use with a field vector of the form [E_r, E_y].
+
+    This operator can be used to form an eigenvalue problem of the form
+    A @ [E_r, E_y] = wavenumber**2 * [E_r, E_y]
+
+    which can then be solved for the eigenmodes of the system (an exp(-i * wavenumber * theta)
+     theta-dependence is assumed for the fields).
+
+    :param omega: The angular frequency of the system
+    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header (2D)
+    :param epsilon: Vectorized dielectric constant grid
+    :param r0: Radius of curvature for the simulation. This should be the minimum value of
+        r within the simulation domain.
+    :return: Sparse matrix representation of the operator
+    """
+
+    Dfx, Dfy = operators.deriv_forward(dxes[0])
+    Dbx, Dby = operators.deriv_back(dxes[1])
+
+    rx = r0 + numpy.cumsum(dxes[0][0])
+    ry = r0 + dxes[0][0]/2.0 + numpy.cumsum(dxes[1][0])
+    tx = 1 + rx/r0
+    ty = 1 + ry/r0
+
+    Tx = sparse.diags(vec(tx[:, None].repeat(dxes[0][1].size, axis=1)))
+    Ty = sparse.diags(vec(ty[:, None].repeat(dxes[1][1].size, axis=1)))
+
+    eps_parts = numpy.split(epsilon, 3)
+    eps_x = sparse.diags(eps_parts[0])
+    eps_y = sparse.diags(eps_parts[1])
+    eps_z_inv = sparse.diags(1 / eps_parts[2])
+
+    pa = sparse.vstack((Dfx, Dfy)) @ Tx @ eps_z_inv @ sparse.hstack((Dbx, Dby))
+    pb = sparse.vstack((Dfx, Dfy)) @ Tx @ eps_z_inv @ sparse.hstack((Dby, Dbx))
+    a0 = Ty @ eps_x + omega**-2 * Dby @ Ty @ Dfy
+    a1 = Tx @ eps_y + omega**-2 * Dbx @ Ty @ Dfx
+    b0 = Dbx @ Ty @ Dfy
+    b1 = Dby @ Ty @ Dfx
+
+    diag = sparse.block_diag
+    op = (omega**2 * diag((Tx, Ty)) + pa) @ diag((a0, a1)) + \
+        - (sparse.bmat(((None, Ty), (Tx, None))) + omega**-2 * pb) @ diag((b0, b1))
+
+    return op
+
+
+
diff --git a/fdfd_tools/waveguide_mode.py b/fdfd_tools/waveguide_mode.py
index 1670991..7c7bc5c 100644
--- a/fdfd_tools/waveguide_mode.py
+++ b/fdfd_tools/waveguide_mode.py
@@ -272,3 +272,69 @@ def compute_overlap_e(E: field_t,
     overlap_e /= norm_factor * dx_forward
 
     return unvec(overlap_e, E[0].shape)
+
+
+def solve_waveguide_mode_cylindrical(mode_number: int,
+                                     omega: complex,
+                                     dxes: dx_lists_t,
+                                     epsilon: vfield_t,
+                                     r0: float,
+                                     wavenumber_correction: bool = True,
+                                     ) -> Dict[str, complex or field_t]:
+    """
+    Given a 2d (r, y) slice of epsilon, attempts to solve for the eigenmode
+     of the bent waveguide with the specified mode number.
+
+    :param mode_number: Number of the mode, 0-indexed
+    :param omega: Angular frequency of the simulation
+    :param dxes: Grid parameters [dx_e, dx_h] as described in fdfd_tools.operators header.
+        The first coordinate is assumed to be r, the second is y.
+    :param epsilon: Dielectric constant
+    :param r0: Radius of curvature for the simulation. This should be the minimum value of
+        r within the simulation domain.
+    :param wavenumber_correction: Whether to correct the wavenumber to
+        account for numerical dispersion (default True)
+    :return: {'E': List[numpy.ndarray], 'H': List[numpy.ndarray], 'wavenumber': complex}
+    """
+
+    '''
+    Solve for the largest-magnitude eigenvalue of the real operator
+    '''
+    dxes_real = [[numpy.real(dx) for dx in dxi] for dxi in dxes]
+
+    A_r = waveguide.cylindrical_operator(numpy.real(omega), dxes_real, numpy.real(epsilon), r0)
+    eigvals, eigvecs = signed_eigensolve(A_r, mode_number + 3)
+    v = eigvecs[:, -(mode_number+1)]
+
+    '''
+    Now solve for the eigenvector of the full operator, using the real operator's
+     eigenvector as an initial guess for Rayleigh quotient iteration.
+    '''
+    A = waveguide.cylindrical_operator(omega, dxes, epsilon, r0)
+    eigval, v = rayleigh_quotient_iteration(A, v)
+
+    # Calculate the wave-vector (force the real part to be positive)
+    wavenumber = numpy.sqrt(eigval)
+    wavenumber *= numpy.sign(numpy.real(wavenumber))
+
+    '''
+    Perform correction on wavenumber to account for numerical dispersion.
+
+     See Numerical Dispersion in Taflove's FDTD book.
+     This correction term reduces the error in emitted power, but additional
+      error is introduced into the E_err and H_err terms. This effect becomes
+      more pronounced as beta increases.
+    '''
+    if wavenumber_correction:
+        wavenumber -= 2 * numpy.sin(numpy.real(wavenumber / 2)) - numpy.real(wavenumber)
+
+    shape = [d.size for d in dxes[0]]
+    v = numpy.hstack((v, numpy.zeros(shape[0] * shape[1])))
+    fields = {
+        'wavenumber': wavenumber,
+        'E': unvec(v, shape),
+#        'E': unvec(e, shape),
+#        'H': unvec(h, shape),
+    }
+
+    return fields