[bloch] add some more tests and clean up solves
This commit is contained in:
parent
87bb3af3f9
commit
e6756742be
2 changed files with 171 additions and 15 deletions
|
|
@ -451,7 +451,7 @@ def find_k(
|
|||
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]]:
|
||||
) -> tuple[NDArray[numpy.float64], float, NDArray[numpy.complex128], NDArray[numpy.complex128]]:
|
||||
"""
|
||||
Search for a bloch vector that has a given frequency.
|
||||
|
||||
|
|
@ -496,15 +496,15 @@ def find_k(
|
|||
|
||||
res = scipy.optimize.minimize_scalar(
|
||||
lambda x: abs(get_f(x, band) - frequency),
|
||||
k_guess,
|
||||
method='Bounded',
|
||||
method='bounded',
|
||||
bounds=k_bounds,
|
||||
options={'xatol': abs(tolerance)},
|
||||
)
|
||||
|
||||
assert n is not None
|
||||
assert v is not None
|
||||
return float(res.x * direction), float(res.fun + frequency), n, v
|
||||
actual_frequency = get_f(float(res.x), band)
|
||||
return direction * float(res.x), float(actual_frequency), n, v
|
||||
|
||||
|
||||
def eigsolve(
|
||||
|
|
@ -725,7 +725,12 @@ def eigsolve(
|
|||
amax=pi,
|
||||
)
|
||||
|
||||
result = scipy.optimize.minimize_scalar(trace_func, bounds=(0, pi), tol=tolerance)
|
||||
result = scipy.optimize.minimize_scalar(
|
||||
trace_func,
|
||||
method='bounded',
|
||||
bounds=(0, pi),
|
||||
options={'xatol': tolerance},
|
||||
)
|
||||
new_E = result.fun
|
||||
theta = result.x
|
||||
|
||||
|
|
@ -754,7 +759,7 @@ def eigsolve(
|
|||
v = eigvecs[:, i]
|
||||
n = eigvals[i]
|
||||
v /= norm(v)
|
||||
Av = (scipy_op @ v.copy())[:, 0]
|
||||
Av = numpy.asarray(scipy_op @ v.copy()).reshape(-1)
|
||||
eigness = norm(Av - (v.conj() @ Av) * v)
|
||||
f = numpy.sqrt(-numpy.real(n))
|
||||
df = numpy.sqrt(-numpy.real(n) + eigness)
|
||||
|
|
@ -823,18 +828,18 @@ def inner_product(
|
|||
# eRxhR_x = numpy.cross(eR.reshape(3, -1), hR.reshape(3, -1), axis=0).reshape(eR.shape)[0] / normR
|
||||
# logger.info(f'power {eRxhR_x.sum() / 2})
|
||||
|
||||
eR /= numpy.sqrt(norm2R)
|
||||
hR /= numpy.sqrt(norm2R)
|
||||
eL /= numpy.sqrt(norm2L)
|
||||
hL /= numpy.sqrt(norm2L)
|
||||
eR_norm = eR / numpy.sqrt(abs(norm2R))
|
||||
hR_norm = hR / numpy.sqrt(abs(norm2R))
|
||||
eL_norm = eL / numpy.sqrt(abs(norm2L))
|
||||
hL_norm = hL / numpy.sqrt(abs(norm2L))
|
||||
|
||||
# (eR x hL)[0] and (eL x hR)[0]
|
||||
eRxhL_x = eR[1] * hL[2] - eR[2] - hL[1]
|
||||
eLxhR_x = eL[1] * hR[2] - eL[2] - hR[1]
|
||||
eRxhL_x = eR_norm[1] * hL_norm[2] - eR_norm[2] * hL_norm[1]
|
||||
eLxhR_x = eL_norm[1] * hR_norm[2] - eL_norm[2] * hR_norm[1]
|
||||
|
||||
#return 1j * (eRxhL_x - eLxhR_x).sum() / numpy.sqrt(norm2R * norm2L)
|
||||
#return (eRxhL_x.sum() - eLxhR_x.sum()) / numpy.sqrt(norm2R * norm2L)
|
||||
return eRxhL_x.sum() - eLxhR_x.sum()
|
||||
return eLxhR_x.sum() - eRxhL_x.sum()
|
||||
|
||||
|
||||
def trq(
|
||||
|
|
@ -848,8 +853,8 @@ def trq(
|
|||
np = inner_product(eO, -hO, eI, hI)
|
||||
nn = inner_product(eO, -hO, eI, -hI)
|
||||
|
||||
assert pp == -nn
|
||||
assert pn == -np
|
||||
assert numpy.allclose(pp, -nn, atol=1e-12, rtol=1e-12)
|
||||
assert numpy.allclose(pn, -np, atol=1e-12, rtol=1e-12)
|
||||
|
||||
logger.info(f'''
|
||||
{pp=:4g} {pn=:4g}
|
||||
|
|
|
|||
151
meanas/test/test_bloch_interactions.py
Normal file
151
meanas/test/test_bloch_interactions.py
Normal file
|
|
@ -0,0 +1,151 @@
|
|||
import numpy
|
||||
from numpy.testing import assert_allclose
|
||||
|
||||
from ..fdfd import bloch
|
||||
|
||||
|
||||
SHAPE = (2, 2, 2)
|
||||
G_MATRIX = numpy.eye(3)
|
||||
EPSILON = numpy.ones((3, *SHAPE), dtype=float)
|
||||
K0 = numpy.array([0.1, 0.0, 0.0], dtype=float)
|
||||
H_SIZE = 2 * numpy.prod(SHAPE)
|
||||
Y0 = (numpy.arange(H_SIZE, dtype=float) + 1j * numpy.linspace(0.1, 0.9, H_SIZE))[None, :]
|
||||
|
||||
|
||||
def build_overlap_fixture() -> tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray, numpy.ndarray]:
|
||||
e_in = numpy.zeros((3, *SHAPE), dtype=complex)
|
||||
h_in = numpy.zeros_like(e_in)
|
||||
e_out = numpy.zeros_like(e_in)
|
||||
h_out = numpy.zeros_like(e_in)
|
||||
|
||||
e_in[1] = 1.0
|
||||
h_in[2] = 2.0
|
||||
e_out[1] = 3.0
|
||||
h_out[2] = 4.0
|
||||
return e_in, h_in, e_out, h_out
|
||||
|
||||
|
||||
def test_rtrace_atb_matches_real_frobenius_inner_product() -> None:
|
||||
a_mat = numpy.array([[1.0 + 2.0j, 3.0 - 1.0j], [2.0j, 4.0]], dtype=complex)
|
||||
b_mat = numpy.array([[5.0 - 1.0j, 1.0 + 1.0j], [2.0, 3.0j]], dtype=complex)
|
||||
expected = numpy.real(numpy.sum(a_mat.conj() * b_mat))
|
||||
|
||||
assert bloch._rtrace_AtB(a_mat, b_mat) == expected
|
||||
|
||||
|
||||
def test_symmetrize_returns_hermitian_average() -> None:
|
||||
matrix = numpy.array([[1.0 + 2.0j, 3.0 - 1.0j], [2.0j, 4.0]], dtype=complex)
|
||||
result = bloch._symmetrize(matrix)
|
||||
|
||||
assert_allclose(result, 0.5 * (matrix + matrix.conj().T))
|
||||
assert_allclose(result, result.conj().T)
|
||||
|
||||
|
||||
def test_inner_product_is_nonmutating_and_obeys_sign_symmetry() -> None:
|
||||
e_in, h_in, e_out, h_out = build_overlap_fixture()
|
||||
originals = (e_in.copy(), h_in.copy(), e_out.copy(), h_out.copy())
|
||||
|
||||
pp = bloch.inner_product(e_out, h_out, e_in, h_in)
|
||||
pn = bloch.inner_product(e_out, h_out, e_in, -h_in)
|
||||
np_term = bloch.inner_product(e_out, -h_out, e_in, h_in)
|
||||
nn = bloch.inner_product(e_out, -h_out, e_in, -h_in)
|
||||
|
||||
assert_allclose(pp, 0.8164965809277263 + 0.0j)
|
||||
assert_allclose(pp, -nn, atol=1e-12, rtol=1e-12)
|
||||
assert_allclose(pn, -np_term, atol=1e-12, rtol=1e-12)
|
||||
assert numpy.array_equal(e_in, originals[0])
|
||||
assert numpy.array_equal(h_in, originals[1])
|
||||
assert numpy.array_equal(e_out, originals[2])
|
||||
assert numpy.array_equal(h_out, originals[3])
|
||||
|
||||
|
||||
def test_trq_returns_expected_transmission_and_reflection() -> None:
|
||||
e_in, h_in, e_out, h_out = build_overlap_fixture()
|
||||
|
||||
transmission, reflection = bloch.trq(e_in, h_in, e_out, h_out)
|
||||
|
||||
assert_allclose(transmission, 0.9797958971132713 + 0.0j, atol=1e-12, rtol=1e-12)
|
||||
assert_allclose(reflection, 0.2 + 0.0j, atol=1e-12, rtol=1e-12)
|
||||
|
||||
|
||||
def test_eigsolve_returns_finite_modes_with_small_residual() -> None:
|
||||
callback_count = 0
|
||||
|
||||
def callback() -> None:
|
||||
nonlocal callback_count
|
||||
callback_count += 1
|
||||
|
||||
eigvals, eigvecs = bloch.eigsolve(
|
||||
1,
|
||||
K0,
|
||||
G_MATRIX,
|
||||
EPSILON,
|
||||
tolerance=1e-6,
|
||||
max_iters=50,
|
||||
y0=Y0,
|
||||
callback=callback,
|
||||
)
|
||||
|
||||
operator = bloch.maxwell_operator(K0, G_MATRIX, EPSILON)
|
||||
eigvec = eigvecs[0] / numpy.linalg.norm(eigvecs[0])
|
||||
residual = numpy.linalg.norm(operator(eigvec).reshape(-1) - eigvals[0] * eigvec) / numpy.linalg.norm(eigvec)
|
||||
|
||||
assert eigvals.shape == (1,)
|
||||
assert eigvecs.shape == (1, H_SIZE)
|
||||
assert numpy.isfinite(eigvals).all()
|
||||
assert numpy.isfinite(eigvecs).all()
|
||||
assert residual < 1e-5
|
||||
assert callback_count > 0
|
||||
|
||||
|
||||
def test_find_k_returns_vector_frequency_and_callbacks() -> None:
|
||||
target_eigvals, _target_eigvecs = bloch.eigsolve(
|
||||
1,
|
||||
K0,
|
||||
G_MATRIX,
|
||||
EPSILON,
|
||||
tolerance=1e-6,
|
||||
max_iters=50,
|
||||
y0=Y0,
|
||||
)
|
||||
target_frequency = float(numpy.sqrt(abs(numpy.real(target_eigvals[0]))))
|
||||
|
||||
solve_calls = 0
|
||||
iter_calls = 0
|
||||
|
||||
def solve_callback(k_mag: float, eigvals: numpy.ndarray, eigvecs: numpy.ndarray, frequency: float) -> None:
|
||||
nonlocal solve_calls
|
||||
solve_calls += 1
|
||||
assert eigvals.shape == (1,)
|
||||
assert eigvecs.shape == (1, H_SIZE)
|
||||
assert isinstance(k_mag, float)
|
||||
assert isinstance(frequency, float)
|
||||
|
||||
def iter_callback() -> None:
|
||||
nonlocal iter_calls
|
||||
iter_calls += 1
|
||||
|
||||
found_k, found_frequency, eigvals, eigvecs = bloch.find_k(
|
||||
target_frequency,
|
||||
1e-4,
|
||||
[1, 0, 0],
|
||||
G_MATRIX,
|
||||
EPSILON,
|
||||
band=0,
|
||||
k_bounds=(0.05, 0.15),
|
||||
v0=Y0,
|
||||
solve_callback=solve_callback,
|
||||
iter_callback=iter_callback,
|
||||
)
|
||||
|
||||
assert found_k.shape == (3,)
|
||||
assert numpy.isfinite(found_k).all()
|
||||
assert_allclose(numpy.cross(found_k, [1.0, 0.0, 0.0]), 0.0, atol=1e-12, rtol=1e-12)
|
||||
assert_allclose(found_k, K0, atol=1e-4, rtol=1e-4)
|
||||
assert abs(found_frequency - target_frequency) <= 1e-4
|
||||
assert eigvals.shape == (1,)
|
||||
assert eigvecs.shape == (1, H_SIZE)
|
||||
assert numpy.isfinite(eigvals).all()
|
||||
assert numpy.isfinite(eigvecs).all()
|
||||
assert solve_calls > 0
|
||||
assert iter_calls > 0
|
||||
Loading…
Add table
Add a link
Reference in a new issue