diff --git a/examples/bloch.py b/examples/bloch.py index 2624876..8e90368 100644 --- a/examples/bloch.py +++ b/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), - numpy.arange(-1000, 1000, dx), - numpy.arange(-1000, 1000, dx)], - shifts=numpy.array([[0,0,0]]), - periodic=True) +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, + ) 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) diff --git a/meanas/fdfd/bloch.py b/meanas/fdfd/bloch.py index 720fef4..9dacdee 100644 --- a/meanas/fdfd/bloch.py +++ b/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,39 +679,49 @@ def eigsolve( trace = _rtrace_AtB(R, Qi) return numpy.abs(trace) - ''' - def trace_deriv(theta): - Qi = Qi_func(theta) - c2 = numpy.cos(2 * theta) - s2 = numpy.sin(2 * theta) - F = -0.5*s2 * (ZtAZ - DtAD) + c2 * symZtAD - trace_deriv = _rtrace_AtB(Qi, F) + if False: + def trace_deriv(theta): + Qi = Qi_func(theta) + c2 = numpy.cos(2 * theta) + s2 = numpy.sin(2 * theta) + F = -0.5*s2 * (ZtAZ - DtAD) + c2 * symZtAD + trace_deriv = _rtrace_AtB(Qi, F) - G = Qi @ F.conj().T @ Qi.conj().T - H = -0.5*s2 * (ZtZ - DtD) + c2 * symZtD - trace_deriv -= _rtrace_AtB(G, H) + G = Qi @ F.conj().T @ Qi.conj().T + H = -0.5*s2 * (ZtZ - DtD) + c2 * symZtD + trace_deriv -= _rtrace_AtB(G, H) - trace_deriv *= 2 - return trace_deriv * sgn + trace_deriv *= 2 + return trace_deriv * sgn - U_sZtD = U @ symZtD + U_sZtD = U @ symZtD - dE = 2.0 * (_rtrace_AtB(U, symZtAD) - - _rtrace_AtB(ZtAZU, U_sZtD)) + dE = 2.0 * (_rtrace_AtB(U, symZtAD) - + _rtrace_AtB(ZtAZU, U_sZtD)) - d2E = 2 * (_rtrace_AtB(U, DtAD) - - _rtrace_AtB(ZtAZU, U @ (DtD - 4 * symZtD @ U_sZtD)) - - 4 * _rtrace_AtB(U, symZtAD @ U_sZtD)) + d2E = 2 * (_rtrace_AtB(U, DtAD) - + _rtrace_AtB(ZtAZU, U @ (DtD - 4 * symZtD @ U_sZtD)) - + 4 * _rtrace_AtB(U, symZtAD @ U_sZtD)) - # Newton-Raphson to find a root of the first derivative: - theta = -dE/d2E + # Newton-Raphson to find a root of the first derivative: + theta = -dE / d2E - if d2E < 0 or abs(theta) >= pi: - theta = -abs(prev_theta) * numpy.sign(dE) + if d2E < 0 or abs(theta) >= pi: + 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, 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) - ''' 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)