import wsp_tools as wt
from wsp_tools import np, plt
X = 1e-9 * np.linspace(-100, 100, 128)
Y = 1e-9 * np.linspace(-100, 100, 129)
Z = 1e-9 * np.linspace(-100, 100, 11)
x2, y2 = np.meshgrid(X, Y)
x3, y3, z3 = np.meshgrid(X, Y, Z)
dx = X[1] - X[0]
dy = Y[1] - Y[0]
# %%
m2 = wt.jchessmodel(x2, y2, z=0, n=1)
plt.imshow(wt.rgba(m2[0] + 1j * m2[1]))
plt.title("2d mag")
plt.show()
# %%
m3 = wt.jchessmodel(x3, y3, z3, n=1)
plt.imshow(wt.rgba(m3[0][:, :, 0] + 1j * m3[1][:, :, 0]))
plt.title("3d mag (slice)")
plt.show()
# %%
abp2 = wt.ab_phase(m2[0], m2[1], m2[2], dx, dy)
plt.imshow(abp2)
plt.title("ab phase 2d")
plt.colorbar()
plt.show()
# %%
def test_ind_from_img():
m2 = wt.jchessmodel(x2, y2, z=0, n=1)
ifm2 = wt.img_from_mag(m2[0], m2[1], m2[2], dx, dy, defocus=1e-3)
ifi = wt.ind_from_img(ifm2, defocus=1e-3, dx=dx, dy=dy)
assert (ifi[0].shape == (16, 16))
def test_ind_from_mag2():
m2 = wt.jchessmodel(x2, y2, z=0, n=1)
Bx, By = wt.ind_from_mag(m2[0], m2[1], m2[2], dx, dy)
assert (Bx.shape == (16, 16))
def test_img_from_phase():
m2 = wt.jchessmodel(x2, y2, z=0, n=1)
phase = wt.ab_phase(m2[0], m2[1], m2[2], dx, dy)
ifp = wt.img_from_phase(phase, dx, dy, defocus=1e-3)
assert (ifp.shape == (16, 16))
def test_ind_from_phase():
m2 = wt.jchessmodel(x2, y2, z=0, n=1)
phase = wt.ab_phase(m2[0], m2[1], m2[2], dx, dy)
indfp = wt.ind_from_phase(phase)
assert (indfp[0].shape == (16, 16))
def test_img_from_mag2():
m2 = wt.jchessmodel(x2, y2, z=0, n=1)
ifm2 = wt.img_from_mag(m2[0], m2[1], m2[2], dx, dy, defocus=1e-3)
assert (ifm2.shape == (16, 16))
def test_img_from_mag3():
m3 = wt.jchessmodel(x3, y3, z3, n=1)
ifm3 = wt.img_from_mag(m3[0], m3[1], m3[2], dx, dy, defocus=1e-3)
assert (ifm3.shape == (16, 16))
def test_B_from_mag2():
m2 = wt.jchessmodel(x2, y2, z=0, n=1)
Bfm2 = wt.B_from_mag(m2[0], m2[1], m2[2], z=1e-9, dx=dx, dy=dy)
assert (Bfm2[0].shape == (16, 16))
def test_A_from_mag3():
m2 = wt.jchessmodel(x2, y2, z=0, n=1)
Afm3 = wt.A_from_mag(m2[0], m2[1], m2[2], z=Z, dx=dx, dy=dy)
assert (Afm3[0].shape == (16, 16, 5))
def test_ab_phase3():
m3 = wt.jchessmodel(x3, y3, z3, n=1)
phase = wt.ab_phase(m3[0], m3[1], m3[2], dx, dy)
assert (phase.shape == (16, 16, 5))
def test_ab_phase2():
m2 = wt.jchessmodel(x2, y2, z=0, n=1)
phase = wt.ab_phase(m2[0], m2[1], m2[2], dx, dy)
assert (phase.shape == (16, 16))
def test_jchessmodel3():
m3 = wt.jchessmodel(x3, y3, z3, n=1)
assert (m3[0].shape == (16, 16, 5))
def test_jchessmodel2():
m2 = wt.jchessmodel(x2, y2, z=0, n=1)
assert (m2[0].shape == (16, 16))
def test_ind_from_mag3():
m3 = wt.jchessmodel(x3, y3, z3, n=1)
Bx, By = wt.ind_from_mag(m3[0], m3[1], m3[2], dx, dy)
assert (Bx.shape == (16, 16))
