From 4aa2d07cef0e8793bbe48bfcf4ac246b193d467f Mon Sep 17 00:00:00 2001
From: jan <jan@mpxd.net>
Date: Sat, 9 Dec 2017 18:21:37 -0800
Subject: [PATCH] Add Bloch eigenproblem

---
 fdfd_tools/bloch.py         | 406 ++++++++++++++++++++++++++++++++++++
 fdfd_tools/eigensolvers.py  |  19 +-
 fdfd_tools/vectorization.py |   2 +-
 3 files changed, 420 insertions(+), 7 deletions(-)
 create mode 100644 fdfd_tools/bloch.py

diff --git a/fdfd_tools/bloch.py b/fdfd_tools/bloch.py
new file mode 100644
index 0000000..d7a2ffa
--- /dev/null
+++ b/fdfd_tools/bloch.py
@@ -0,0 +1,406 @@
+'''
+Bloch eigenmode solver/operators
+
+This module contains functions for generating and solving the
+ 3D Bloch eigenproblem. The approach is to transform the problem
+ into the (spatial) fourier domain, transforming the equation
+   1/mu * curl(1/eps * curl(H)) = (w/c)^2 H
+ into
+   conv(1/mu_k, ik x conv(1/eps_k, ik x H_k)) = (w/c)^2 H_k
+ where:
+  - the _k subscript denotes a 3D fourier transformed field
+  - each component of H_k corresponds to a plane wave with wavevector k
+  - x is the cross product
+  - conv denotes convolution
+
+ Since k and H are orthogonal for each plane wave, we can use each
+ k to create an orthogonal basis (k, m, n), with k x m = n, and
+ |m| = |n| = 1. The cross products are then simplified with
+
+ k @ h = kx hx + ky hy + kz hz = 0 = hk
+ h = hk + hm + hn = hm + hn
+ k = kk + km + kn = kk = |k|
+
+ k x h = (ky hz - kz hy,
+          kz hx - kx hz,
+          kx hy - ky hx)
+       = ((k x h) @ k, (k x h) @ m, (k x h) @ n)_kmn
+       = (0, (m x k) @ h, (n x k) @ h)_kmn         # triple product ordering
+       = (0, kk (-n @ h), kk (m @ h))_kmn          # (m x k) = -|k| n, etc.
+       = |k| (0, -h @ n, h @ m)_kmn
+
+ k x h = (km hn - kn hm,
+          kn hk - kk hn,
+          kk hm - km hk)_kmn
+       = (0, -kk hn, kk hm)_kmn
+       = (-kk hn)(mx, my, mz) + (kk hm)(nx, ny, nz)
+       = |k| (hm * (nx, ny, nz) - hn * (mx, my, mz))
+
+ where h is shorthand for H_k, (...)_kmn deontes the (k, m, n) basis,
+ and e.g. hm is the component of h in the m direction.
+
+ We can also simplify conv(X_k, Y_k) as fftn(X * ifftn(Y_k)).
+
+ Using these results and storing H_k as h = (hm, hn), we have
+  e_xyz = fftn(1/eps * ifftn(|k| (hm * n - hn * m)))
+  b_mn = |k| (-e_xyz @ n, e_xyz @ m)
+  h_mn = fftn(1/mu * ifftn(b_m * m + b_n * n))
+ which forms the operator from the left side of the equation.
+
+ We can then use ARPACK in shift-invert mode (via scipy.linalg.eigs)
+  to find the eigenvectors for this operator.
+
+ This approach is similar to the one used in MPB and derived at the start of
+ SG Johnson and JD Joannopoulos, Block-iterative frequency-domain methods
+  for Maxwell's equations in a planewave basis, Optics Express 8, 3, 173-190 (2001)
+
+ ===
+
+ Typically you will want to do something like
+
+    recip_lattice = numpy.diag(1/numpy.array(epsilon[0].shape * dx))
+    n, v = bloch.eigsolve(5, k0, recip_lattice, epsilon)
+    f = numpy.sqrt(-numpy.real(n[0]))
+    n_eff = norm(recip_lattice @ k0) / f
+
+    v2e = bloch.hmn_2_exyz(k0, recip_lattice, epsilon)
+    e_field = v2e(v[0])
+
+    k, f = find_k(frequency=1/1550,
+                  tolerance=(1/1550 - 1/1551),
+                  search_direction=[1, 0, 0],
+                  G_matrix=recip_lattice,
+                  epsilon=epsilon,
+                  band=0)
+
+'''
+
+from typing import List, Tuple, Callable, Dict
+import numpy
+from numpy.fft import fftn, ifftn, fftfreq
+import scipy
+from scipy.linalg import norm
+import scipy.sparse.linalg as spalg
+
+from . import field_t
+
+
+def generate_kmn(k0: numpy.ndarray,
+                 G_matrix: numpy.ndarray,
+                 shape: numpy.ndarray
+                 ) -> Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]:
+    """
+    Generate a (k, m, n) orthogonal basis for each k-vector in the simulation grid.
+
+    :param k0: [k0x, k0y, k0z], Bloch wavevector, in G basis.
+    :param G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
+    :param shape: [nx, ny, nz] shape of the simulation grid.
+    :return: (|k|, m, n) where |k| has shape tuple(shape) + (1,)
+            and m, n have shape tuple(shape) + (3,).
+            All are given in the xyz basis (e.g. |k|[0,0,0] = norm(G_matrix @ k0)).
+    """
+    k0 = numpy.array(k0)
+
+    Gi_grids = numpy.meshgrid(*(fftfreq(n, 1/n) for n in shape[:3]), indexing='ij')
+    Gi = numpy.stack(Gi_grids, axis=3)
+
+    k_G = k0[None, None, None, :] - Gi
+    k_xyz = numpy.rollaxis(G_matrix @ numpy.rollaxis(k_G, 3, 2), 3, 2)
+
+    m = numpy.broadcast_to([0, 1, 0], tuple(shape[:3]) + (3,)).astype(float)
+    n = numpy.broadcast_to([0, 0, 1], tuple(shape[:3]) + (3,)).astype(float)
+
+    xy_non0 = numpy.any(k_xyz[:, :, :, 0:1] != 0, axis=3)
+    if numpy.any(xy_non0):
+        u = numpy.cross(k_xyz[xy_non0], [0, 0, 1])
+        m[xy_non0, :] = u / norm(u, axis=1)[:, None]
+
+    z_non0 = numpy.any(k_xyz != 0, axis=3)
+    if numpy.any(z_non0):
+        v = numpy.cross(k_xyz[z_non0], m[z_non0])
+        n[z_non0, :] = v / norm(v, axis=1)[:, None]
+
+    k_mag = norm(k_xyz, axis=3)[:, :, :, None]
+    return k_mag, m, n
+
+
+def maxwell_operator(k0: numpy.ndarray,
+                     G_matrix: numpy.ndarray,
+                     epsilon: field_t,
+                     mu: field_t = None
+                     ) -> Callable[[numpy.ndarray], numpy.ndarray]:
+    """
+    Generate the Maxwell operator
+        conv(1/mu_k, ik x conv(1/eps_k, ik x ___))
+    which is the spatial-frequency-space representation of
+        1/mu * curl(1/eps * curl(___))
+
+    The operator is a function that acts on a vector h_mn of size (2 * epsilon[0].size)
+
+    See the module-level docstring for more information.
+
+    :param k0: Bloch wavevector, [k0x, k0y, k0z].
+    :param G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
+    :param epsilon: Dielectric constant distribution for the simulation.
+        All fields are sampled at cell centers (i.e., NOT Yee-gridded)
+    :param mu: Magnetic permability distribution for the simulation.
+        Default None (1 everywhere).
+    :return: Function which applies the maxwell operator to h_mn.
+    """
+
+    shape = epsilon[0].shape + (1,)
+    k_mag, m, n = generate_kmn(k0, G_matrix, shape)
+
+    epsilon = numpy.stack(epsilon, 3)
+    if mu is not None:
+        mu = numpy.stack(mu, 3)
+
+    def operator(h: numpy.ndarray):
+        """
+        Maxwell operator for Bloch eigenmode simulation.
+
+        h is complex 2-field in (m, n) basis, vectorized
+
+        :param h: Raveled h_mn; size (2 * epsilon[0].size).
+        :return: Raveled conv(1/mu_k, ik x conv(1/eps_k, ik x h_mn)).
+        """
+        hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)]
+
+        #{d,e,h}_xyz fields are complex 3-fields in (1/x, 1/y, 1/z) basis
+
+        # cross product and transform into xyz basis
+        d_xyz = (n * hin_m -
+                 m * hin_n) * k_mag
+
+        # divide by epsilon
+        e_xyz = ifftn(fftn(d_xyz, axes=range(3)) / epsilon, axes=range(3))
+
+        # cross product and transform into mn basis
+        b_m = numpy.sum(e_xyz * n, axis=3)[:, :, :, None] * -k_mag
+        b_n = numpy.sum(e_xyz * m, axis=3)[:, :, :, None] * +k_mag
+
+        if mu is None:
+            h_m, h_n = b_m, b_n
+        else:
+            # transform from mn to xyz
+            b_xyz = (m * b_m[:, :, :, None] +
+                     n * b_n[:, :, :, None])
+
+            # divide by mu
+            h_xyz = ifftn(fftn(b_xyz, axes=range(3)) / mu, axes=range(3))
+
+            # transform back to mn
+            h_m = numpy.sum(h_xyz * m, axis=3)
+            h_n = numpy.sum(h_xyz * n, axis=3)
+        return numpy.hstack((h_m.ravel(), h_n.ravel()))
+
+    return operator
+
+
+def hmn_2_exyz(k0: numpy.ndarray,
+               G_matrix: numpy.ndarray,
+               epsilon: field_t,
+               ) -> Callable[[numpy.ndarray], field_t]:
+    """
+    Generate an operator which converts a vectorized spatial-frequency-space
+     h_mn into an E-field distribution, i.e.
+        ifft(conv(1/eps_k, ik x h_mn))
+
+    The operator is a function that acts on a vector h_mn of size (2 * epsilon[0].size)
+
+    See the module-level docstring for more information.
+
+    :param k0: Bloch wavevector, [k0x, k0y, k0z].
+    :param G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
+    :param epsilon: Dielectric constant distribution for the simulation.
+        All fields are sampled at cell centers (i.e., NOT Yee-gridded)
+    :return: Function for converting h_mn into E_xyz
+    """
+    shape = epsilon[0].shape + (1,)
+    epsilon = numpy.stack(epsilon, 3)
+
+    k_mag, m, n = generate_kmn(k0, G_matrix, shape)
+
+    def operator(h: numpy.ndarray) -> field_t:
+        hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)]
+        d_xyz = (n * hin_m -
+                 m * hin_n) * k_mag
+
+        # divide by epsilon
+        return [ei for ei in numpy.rollaxis(fftn(d_xyz, axes=range(3)) / epsilon, 3)]
+
+    return operator
+
+
+def hmn_2_hxyz(k0: numpy.ndarray,
+               G_matrix: numpy.ndarray,
+               epsilon: field_t
+               ) -> Callable[[numpy.ndarray], field_t]:
+    """
+    Generate an operator which converts a vectorized spatial-frequency-space
+     h_mn into an H-field distribution, i.e.
+        ifft(h_mn)
+
+    The operator is a function that acts on a vector h_mn of size (2 * epsilon[0].size)
+
+    See the module-level docstring for more information.
+
+    :param k0: Bloch wavevector, [k0x, k0y, k0z].
+    :param G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
+    :param epsilon: Dielectric constant distribution for the simulation.
+        Only epsilon[0].shape is used.
+    :return: Function for converting h_mn into H_xyz
+    """
+    shape = epsilon[0].shape + (1,)
+    k_mag, m, n = generate_kmn(k0, G_matrix, shape)
+
+    def operator(h: numpy.ndarray):
+        hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)]
+        h_xyz = (m * hin_m +
+                 n * hin_n)
+        return [fftn(hi) for hi in numpy.rollaxis(h_xyz, 3)]
+
+    return operator
+
+
+def inverse_maxwell_operator_approx(k0: numpy.ndarray,
+                                    G_matrix: numpy.ndarray,
+                                    epsilon: field_t,
+                                    mu: field_t = None
+                                    ) -> Callable[[numpy.ndarray], numpy.ndarray]:
+    """
+    Generate an approximate inverse of the Maxwell operator,
+        ik x conv(eps_k, ik x conv(mu_k, ___))
+     which can be used to improve the speed of ARPACK in shift-invert mode.
+
+    See the module-level docstring for more information.
+
+    :param k0: Bloch wavevector, [k0x, k0y, k0z].
+    :param G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
+    :param epsilon: Dielectric constant distribution for the simulation.
+        All fields are sampled at cell centers (i.e., NOT Yee-gridded)
+    :param mu: Magnetic permability distribution for the simulation.
+        Default None (1 everywhere).
+    :return: Function which applies the approximate inverse of the maxwell operator to h_mn.
+    """
+    shape = epsilon[0].shape + (1,)
+    epsilon = numpy.stack(epsilon, 3)
+
+    k_mag, m, n = generate_kmn(k0, G_matrix, shape)
+
+    if mu is not None:
+        mu = numpy.stack(mu, 3)
+
+    def operator(h: numpy.ndarray):
+        """
+        Approximate inverse Maxwell operator for Bloch eigenmode simulation.
+
+        h is complex 2-field in (m, n) basis, vectorized
+
+        :param h: Raveled h_mn; size (2 * epsilon[0].size).
+        :return: Raveled ik x conv(eps_k, ik x conv(mu_k, h_mn))
+        """
+        hin_m, hin_n = [hi.reshape(shape) for hi in numpy.split(h, 2)]
+
+        #{d,e,h}_xyz fields are complex 3-fields in (1/x, 1/y, 1/z) basis
+
+        if mu is None:
+            b_m, b_n = hin_m, hin_n
+        else:
+            # transform from mn to xyz
+            h_xyz = (m * hin_m[:, :, :, None] +
+                     n * hin_n[:, :, :, None])
+
+            # multiply by mu
+            b_xyz = ifftn(fftn(h_xyz, axes=range(3)) * mu, axes=range(3))
+
+            # transform back to mn
+            b_m = numpy.sum(b_xyz * m, axis=3)
+            b_n = numpy.sum(b_xyz * n, axis=3)
+
+        # cross product and transform into xyz basis
+        e_xyz = (n * b_m -
+                 m * b_n) / k_mag
+
+        # multiply by epsilon
+        d_xyz = ifftn(fftn(e_xyz, axes=range(3)) * epsilon, axes=range(3))
+
+        # cross product and transform into mn basis   crossinv_t2c
+        h_m = numpy.sum(e_xyz * n, axis=3)[:, :, :, None] / +k_mag
+        h_n = numpy.sum(e_xyz * m, axis=3)[:, :, :, None] / -k_mag
+
+        return numpy.hstack((h_m.ravel(), h_n.ravel()))
+
+    return operator
+
+
+def eigsolve(num_modes: int,
+             k0: numpy.ndarray,
+             G_matrix: numpy.ndarray,
+             epsilon: field_t,
+             mu: field_t = None
+             ) -> Tuple[numpy.ndarray, numpy.ndarray]:
+    """
+    Find the first (lowest-frequency) num_modes eigenmodes with Bloch wavevector
+     k0 of the specified structure.
+
+    :param k0: Bloch wavevector, [k0x, k0y, k0z].
+    :param G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
+    :param epsilon: Dielectric constant distribution for the simulation.
+        All fields are sampled at cell centers (i.e., NOT Yee-gridded)
+    :param mu: Magnetic permability distribution for the simulation.
+        Default None (1 everywhere).
+    :return: (eigenvalues, eigenvectors) where eigenvalues[i] corresponds to the
+        vector eigenvectors[i, :]
+    """
+    h_size = 2 * epsilon[0].size
+
+    mop = maxwell_operator(k0=k0, G_matrix=G_matrix, epsilon=epsilon, mu=mu)
+    imop = inverse_maxwell_operator_approx(k0=k0, G_matrix=G_matrix, epsilon=epsilon, mu=mu)
+
+    scipy_op = spalg.LinearOperator(dtype=complex, shape=(h_size, h_size), matvec=mop)
+    scipy_iop = spalg.LinearOperator(dtype=complex, shape=(h_size, h_size), matvec=imop)
+
+    _eigvals, eigvecs = spalg.eigs(scipy_op, num_modes, sigma=0, OPinv=scipy_iop, which='LM')
+    eigvals = numpy.sum(eigvecs * (scipy_op @ eigvecs), axis=0) / numpy.sum(eigvecs * eigvecs, axis=0)
+    order = numpy.argsort(-eigvals)
+    return eigvals[order], eigvecs.T[order]
+
+
+def find_k(frequency: float,
+           tolerance: float,
+           search_direction: numpy.ndarray,
+           G_matrix: numpy.ndarray,
+           epsilon: field_t,
+           mu: field_t = None,
+           band: int = 0
+           ) -> Tuple[numpy.ndarray, float]:
+    """
+    Search for a bloch vector that has a given frequency.
+
+    :param frequency: Target frequency.
+    :param tolerance: Target frequency tolerance.
+    :param search_direction: k-vector direction to search along.
+    :param G_matrix: 3x3 matrix, with reciprocal lattice vectors as columns.
+    :param epsilon: Dielectric constant distribution for the simulation.
+        All fields are sampled at cell centers (i.e., NOT Yee-gridded)
+    :param mu: Magnetic permability distribution for the simulation.
+        Default None (1 everywhere).
+    :param band: Which band to search in. Default 0 (lowest frequency).
+    return: (k, actual_frequency) The found k-vector and its frequency
+    """
+
+    search_direction = numpy.array(search_direction) / norm(search_direction)
+
+    def get_f(k0_mag: float, band: int = 0):
+        k0 = search_direction * k0_mag
+        n, _v = eigsolve(band + 1, k0, G_matrix=G_matrix, epsilon=epsilon)
+        f = numpy.sqrt(numpy.abs(numpy.real(n[band])))
+        return f
+
+    res = scipy.optimize.minimize_scalar(lambda x: abs(get_f(x, band) - frequency), 0.25,
+                                         method='Bounded', bounds=(0, 0.5),
+                                         options={'xatol': abs(tolerance)})
+    return res.x * search_direction, res.fun + frequency
+
+
diff --git a/fdfd_tools/eigensolvers.py b/fdfd_tools/eigensolvers.py
index 6c06296..24d7339 100644
--- a/fdfd_tools/eigensolvers.py
+++ b/fdfd_tools/eigensolvers.py
@@ -29,7 +29,7 @@ def power_iteration(operator: sparse.spmatrix,
         v = operator @ v
         v /= norm(v)
 
-    lm_eigval = v.conj() @ operator @ v
+    lm_eigval = v.conj() @ (operator @ v)
     return lm_eigval, v
 
 
@@ -59,7 +59,7 @@ def rayleigh_quotient_iteration(operator: sparse.spmatrix,
     return eigval, v
 
 
-def signed_eigensolve(operator: sparse.spmatrix,
+def signed_eigensolve(operator: sparse.spmatrix or spalg.LinearOperator,
                       how_many: int,
                       negative: bool = False,
                       ) -> Tuple[numpy.ndarray, numpy.ndarray]:
@@ -83,12 +83,19 @@ def signed_eigensolve(operator: sparse.spmatrix,
      largest _positive_ eigenvalues, since any negative eigenvalues will be shifted to the
      range 0 >= neg_eigval + abs(lm_eigval) > abs(lm_eigval)
     '''
+    shift = numpy.abs(lm_eigval)
     if negative:
-        shifted_operator = operator - abs(lm_eigval) * sparse.eye(operator.shape[0])
-    else:
-        shifted_operator = operator + abs(lm_eigval) * sparse.eye(operator.shape[0])
+        shift *= -1
 
-    eigenvalues, eigenvectors = spalg.eigs(shifted_operator, which='LM', k=how_many, ncv=50)
+    # Try to combine, use general LinearOperator if we fail
+    try:
+        shifted_operator = operator + shift * sparse.eye(operator.shape[0])
+    except TypeError:
+        shifted_operator = operator + spalg.LinearOperator(shape=operator.shape,
+                                                           matvec=lambda v: shift * v)
+
+    shifted_eigenvalues, eigenvectors = spalg.eigs(shifted_operator, which='LM', k=how_many, ncv=50)
+    eigenvalues = shifted_eigenvalues - shift
 
     k = eigenvalues.argsort()
     return eigenvalues[k], eigenvectors[:, k]
diff --git a/fdfd_tools/vectorization.py b/fdfd_tools/vectorization.py
index 1c2134a..e7b9645 100644
--- a/fdfd_tools/vectorization.py
+++ b/fdfd_tools/vectorization.py
@@ -1,7 +1,7 @@
 """
 Functions for moving between a vector field (list of 3 ndarrays, [f_x, f_y, f_z])
 and a 1D array representation of that field [f_x0, f_x1, f_x2,... f_y0,... f_z0,...].
-Vectorized versions of the field use column-major (ie., Fortran, Matlab) ordering.
+Vectorized versions of the field use row-major (ie., C-style) ordering.
 """