bloch example updates

examples/bloch.py

@@ -20,6 +20,7 @@ def pyfftw_save_wisdom(path):
         pass
 
     path.parent.mkdir(parents=True, exist_ok=True)
+    wisdom = pyfftw.export_wisdom()
     with open(path, 'wb') as f:
         pickle.dump(wisdom, f)
 
@@ -42,11 +43,13 @@ logger.info('Drawing grid...')
 dx = 40
 x_period = 400
 y_period = z_period = 2000
-g = gridlock.Grid([numpy.arange(-x_period/2, x_period/2, dx),
+g = gridlock.Grid([
+    numpy.arange(-x_period/2, x_period/2, dx),
     numpy.arange(-1000, 1000, dx),
     numpy.arange(-1000, 1000, dx)],
     shifts=numpy.array([[0,0,0]]),
-                  periodic=True)
+    periodic=True,
+    )
 gdata = g.allocate(1.445**2)
 
 g.draw_cuboid(gdata, [0,0,0], [200e8, 220, 220], foreground=3.47**2)
@@ -74,7 +77,8 @@ pyfftw_load_wisdom(WISDOM_FILEPATH)
 #                    epsilon=epsilon,
 #                    band=0)
 #
-#print("k={}, f={}, 1/f={}, k/f={}".format(k, f, 1/f, norm(reciprocal_lattice @ k) / f ))
+#kf = norm(reciprocal_lattice @ k) / f)
+#print(f'{k=}, {f=}, 1/f={1/f}, k/f={kf}')
 
 logger.info('Finding f at [0.25, 0, 0]')
 for k0x in [.25]:
@@ -82,7 +86,7 @@ for k0x in [.25]:
 
     kmag = norm(reciprocal_lattice @ k0)
     tolerance = (1000/1550) * 1e-4/1.5  # df = f * dn_eff / n
-    logger.info('tolerance {}'.format(tolerance))
+    logger.info(f'tolerance {tolerance}')
 
     n, v = bloch.eigsolve(4, k0, G_matrix=reciprocal_lattice, epsilon=epsilon, tolerance=tolerance**2)
     v2e = bloch.hmn_2_exyz(k0, G_matrix=reciprocal_lattice, epsilon=epsilon)
@@ -96,8 +100,11 @@ for k0x in [.25]:
         g2data[i+3] += numpy.imag(e[i])
 
     f = numpy.sqrt(numpy.real(numpy.abs(n))) # TODO
-    print('k0x = {:3g}\n eigval = {}\n f = {}\n'.format(k0x, n, f))
+    print(f'{k0x=:3g}')
+    print(f'eigval={n}')
+    print(f'{f=}')
     n_eff = norm(reciprocal_lattice @ k0) / f
-    print('kmag/f = n_eff =  {} \n wl = {}\n'.format(n_eff, 1/f ))
+    print(f'kmag/f = n_eff = {n_eff}')
+    print(f'wl={1/f}\n')
 
 pyfftw_save_wisdom(WISDOM_FILEPATH)

meanas/fdfd/bloch.py

@@ -1,4 +1,4 @@
-'''
+"""
 Bloch eigenmode solver/operators
 
 This module contains functions for generating and solving the
@@ -92,7 +92,7 @@ This module contains functions for generating and solving the
                   epsilon=epsilon,
                   band=0)
 
-'''
+"""
 
 from typing import Callable, Any, cast, Sequence
 import logging
@@ -265,7 +265,9 @@ def maxwell_operator(
             h_m = numpy.sum(h_xyz * m, axis=3)
             h_n = numpy.sum(h_xyz * n, axis=3)
 
-        h = numpy.concatenate((h_m, h_n), axis=None, out=h)     # ravel and merge
+        h.shape = (h.size,)
+        h = numpy.concatenate((h_m.ravel(), h_n.ravel()), axis=None, out=h)     # ravel and merge
+        h.shape = (h.size, 1)
         return h
 
     return operator
@@ -425,7 +427,9 @@ def inverse_maxwell_operator_approx(
         h_m = numpy.sum(d_xyz * n, axis=3, keepdims=True) / +k_mag
         h_n = numpy.sum(d_xyz * m, axis=3, keepdims=True) / -k_mag
 
+        h.shape = (h.size,)
         h = numpy.concatenate((h_m, h_n), axis=None, out=h)
+        h.shape = (h.size, 1)
         return h
 
     return operator
@@ -443,6 +447,7 @@ def find_k(
         k_guess: float | None = None,
         solve_callback: Callable[..., None] | None = None,
         iter_callback: Callable[..., None] | None = None,
+        v0: NDArray[numpy.complex128] | None = None,
         ) -> tuple[float, float, NDArray[numpy.complex128], NDArray[numpy.complex128]]:
     """
     Search for a bloch vector that has a given frequency.
@@ -475,7 +480,7 @@ def find_k(
         k_guess = sum(k_bounds) / 2
 
     n = None
-    v = None
+    v = v0
 
     def get_f(k0_mag: float, band: int = 0) -> float:
         nonlocal n, v
@@ -505,7 +510,7 @@ def eigsolve(
         G_matrix: ArrayLike,
         epsilon: fdfield_t,
         mu: fdfield_t | None = None,
-        tolerance: float = 1e-20,
+        tolerance: float = 1e-7,
         max_iters: int = 10000,
         reset_iters: int = 100,
         y0: ArrayLike | None = None,
@@ -539,9 +544,9 @@ def eigsolve(
 
     kmag = norm(G_matrix @ k0)
 
-    '''
-    Generate the operators
-    '''
+    #
+    # Generate the operators
+    #
     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)
 
@@ -553,14 +558,14 @@ def eigsolve(
     prev_E = 0.0
     d_scale = 1.0
     prev_traceGtKG = 0.0
-    #prev_theta = 0.5
+    prev_theta = 0.5
     D = numpy.zeros(shape=y_shape, dtype=complex)
 
     Z: NDArray[numpy.complex128]
     if y0 is None:
         Z = numpy.random.rand(*y_shape) + 1j * numpy.random.rand(*y_shape)
     else:
-        Z = numpy.array(y0, copy=False)
+        Z = numpy.array(y0, copy=False).T
 
     while True:
         Z *= num_modes / norm(Z)
@@ -573,13 +578,13 @@ def eigsolve(
 
         trace_U = real(trace(U))
         if trace_U > 1e8 * num_modes:
-            Z = Z @ scipy.linalg.sqrtm(U).conj().T
+            Z = Z @ scipy.linalg.sqrtm(U).astype(numpy.complex128).conj().T
             prev_traceGtKG = 0
             continue
         break
 
-    Zt = numpy.empty(Z.shape[::-1])
-    AZ = numpy.empty(Z.shape)
+    Zt = numpy.empty(Z.shape[::-1], dtype=numpy.complex128)
+    AZ = numpy.empty(Z.shape, dtype=numpy.complex128)
 
     for i in range(max_iters):
         Zt = numpy.conj(Z.T, out=Zt)
@@ -591,8 +596,9 @@ def eigsolve(
         E_signed = real(trace(ZtAZU))
         sgn = numpy.sign(E_signed)
         E = numpy.abs(E_signed)
-        G = (AZ @ U - Z @ U @ ZtAZU) * sgn      # G = AZU projected onto the space orthonormal to Z
-                                                #  via (1 - ZUZt)
+
+        # G = AZU projected onto the space orthonormal to Z via (1 - ZUZt)
+        G = (AZ @ U - Z @ U @ ZtAZU) * sgn
 
         if i > 0 and abs(E - prev_E) < tolerance * 0.5 * (E + prev_E + 1e-7):
             logger.info(
@@ -603,7 +609,7 @@ def eigsolve(
             break
 
         KG = scipy_iop @ G          # Preconditioned steepest descent direction
-        traceGtKG = _rtrace_AtB(G, KG)      #
+        traceGtKG = _rtrace_AtB(G, KG)
 
         if prev_traceGtKG == 0 or i % reset_iters == 0:
             logger.info('CG reset')
@@ -631,7 +637,7 @@ def eigsolve(
         #
         #   Qi = inv(Q) = U'
 
-        AD = scipy_op @ D
+        AD = scipy_op @ D.copy()
         DtD = D.conj().T @ D
         DtAD = D.conj().T @ AD
 
@@ -673,7 +679,7 @@ def eigsolve(
             trace = _rtrace_AtB(R, Qi)
             return numpy.abs(trace)
 
-        '''
+        if False:
             def trace_deriv(theta):
                 Qi = Qi_func(theta)
                 c2 = numpy.cos(2 * theta)
@@ -704,8 +710,18 @@ def eigsolve(
                 theta = -abs(prev_theta) * numpy.sign(dE)
 
             # theta, new_E, new_dE = linmin(theta, E, dE, 0.1, min(tolerance, 1e-6), 1e-14, 0, -numpy.sign(dE) * K_PI, trace_func)
-        theta, n, _, new_E, _, _new_dE = scipy.optimize.line_search(trace_func, trace_deriv, xk=theta, pk=numpy.ones((1,1)), gfk=dE, old_fval=E, c1=min(tolerance, 1e-6), c2=0.1, amax=pi)
-        '''
+            theta, n, _, new_E, _, _new_dE = scipy.optimize.line_search(
+                trace_func,
+                trace_deriv,
+                xk=theta,
+                pk=numpy.ones((1, 1)),
+                gfk=dE,
+                old_fval=E,
+                c1=min(tolerance, 1e-6),
+                c2=0.1,
+                amax=pi,
+                )
+
         result = scipy.optimize.minimize_scalar(trace_func, bounds=(0, pi), tol=tolerance)
         new_E = result.fun
         theta = result.x
@@ -715,18 +731,18 @@ def eigsolve(
         Z *= numpy.cos(theta)
         Z += D * numpy.sin(theta)
 
-        #prev_theta = theta
+        prev_theta = theta
         prev_E = E
 
         if callback:
             callback()
 
-    '''
-    Recover eigenvectors from Z
-    '''
+    #
+    # Recover eigenvectors from Z
+    #
     U = numpy.linalg.inv(ZtZ)
-    Y = Z @ scipy.linalg.sqrtm(U)
-    W = Y.conj().T @ (scipy_op @ Y)
+    Y = Z @ scipy.linalg.sqrtm(U).astype(numpy.complex128)
+    W = Y.conj().T @ (scipy_op @ Y.copy())
 
     eigvals, W_eigvecs = numpy.linalg.eig(W)
     eigvecs = Y @ W_eigvecs
@@ -735,7 +751,8 @@ def eigsolve(
         v = eigvecs[:, i]
         n = eigvals[i]
         v /= norm(v)
-        eigness = norm(scipy_op @ v - (v.conj() @ (scipy_op @ v)) * v)
+        Av = (scipy_op @ v.copy())[:, 0]
+        eigness = norm(Av - (v.conj() @ Av) * v)
         f = numpy.sqrt(-numpy.real(n))
         df = numpy.sqrt(-numpy.real(n + eigness))
         neff_err = kmag * (1 / df - 1 / f)