183 lines
5.0 KiB
Python
183 lines
5.0 KiB
Python
import numpy as np
|
|
from scipy.stats import multivariate_normal
|
|
import matplotlib.pyplot as plt
|
|
import scipy
|
|
import scipy.fftpack as sfft
|
|
|
|
|
|
def deg_2_rad(winkel):
|
|
return winkel/180.0 * np.pi
|
|
|
|
|
|
# all units in angstrom
|
|
base_a_m = 5.75
|
|
base_b_m = 4.5
|
|
base_c_m = 5.38
|
|
|
|
base_c_r = 2.856
|
|
base_b_r = 4.554
|
|
base_a_r = base_b_r
|
|
|
|
alpha_m = 122.64 # degree
|
|
|
|
|
|
def mono_2_rutile(c_m, a_m):
|
|
a_r = np.cos(deg_2_rad(alpha_m - 90)) * c_m * base_c_m
|
|
c_r = (a_m) * base_a_m + np.sin(deg_2_rad(alpha_m - 90)) * c_m * base_c_m
|
|
return a_r, c_r
|
|
|
|
|
|
class Lattice:
|
|
|
|
def _get_rutile(self, X, Y):
|
|
self.atom_x_rut = X * base_c_r + np.mod(Y, 4) * 0.5 * base_c_r
|
|
self.atom_y_rut = Y * 0.5 * base_a_r
|
|
|
|
def _get_mono(self, X, Y):
|
|
offset_a_m = 0.25 - 0.23947
|
|
offset_c_m = 0.02646
|
|
|
|
offset_a_r, offset_c_r = mono_2_rutile(offset_c_m, offset_a_m)
|
|
|
|
print("A_r: ", offset_a_r, "C_r: ", offset_c_r)
|
|
|
|
self.atom_x_mono = offset_a_r + X * \
|
|
base_c_r + np.mod(Y, 4) * 0.5 * base_c_r
|
|
self.atom_x_mono[np.mod(X, 2) == 0] -= 2 * offset_a_r
|
|
|
|
self.atom_y_mono = offset_c_r + 0.5 * Y * base_a_r
|
|
self.atom_y_mono[np.mod(X, 2) == 0] -= 2 * offset_c_r
|
|
|
|
def _generate_vec(self, x_len: int, y_len: int):
|
|
x = np.arange(x_len)
|
|
y = np.arange(y_len)
|
|
X, Y = np.meshgrid(x, y)
|
|
X[np.mod(Y, 4) == 3] = X[np.mod(Y, 4) == 3]-1
|
|
X[np.mod(Y, 4) == 2] = X[np.mod(Y, 4) == 2]-1
|
|
assert np.mod(x.size, 2) == 0
|
|
assert np.mod(y.size, 2) == 0
|
|
|
|
return X, Y
|
|
|
|
def __init__(self, x_len: int, y_len: int):
|
|
X, Y = self._generate_vec(x_len*2, y_len*2)
|
|
self._get_mono(X, Y)
|
|
self._get_rutile(X, Y)
|
|
|
|
def get_from_mask(self, maske: np.array, inplace_pos_x=None, inplace_pos_y=None):
|
|
if inplace_pos_x is None:
|
|
inplace_pos_x = np.zeros_like(self.atom_x_mono)
|
|
if inplace_pos_y is None:
|
|
inplace_pos_y = np.zeros_like(self.atom_x_mono)
|
|
|
|
mask = np.empty_like(self.atom_x_mono, dtype=bool)
|
|
print(mask.shape, maske.shape)
|
|
mask[0::2, 0::2] = maske
|
|
mask[1::2, 0::2] = maske
|
|
mask[0::2, 1::2] = maske
|
|
mask[1::2, 1::2] = maske
|
|
|
|
inplace_pos_x[mask] = self.atom_x_rut[mask]
|
|
inplace_pos_y[mask] = self.atom_y_rut[mask]
|
|
mask = np.invert(mask)
|
|
inplace_pos_x[mask] = self.atom_x_mono[mask]
|
|
inplace_pos_y[mask] = self.atom_y_mono[mask]
|
|
return inplace_pos_x, inplace_pos_y
|
|
|
|
|
|
def test_lattice():
|
|
lat = Lattice(10, 10)
|
|
maske = np.zeros((10, 10), dtype=bool)
|
|
x, y = lat.get_from_mask(maske)
|
|
|
|
plt.scatter(x, y)
|
|
maske = np.invert(maske)
|
|
x, y = lat.get_from_mask(maske)
|
|
plt.scatter(x, y)
|
|
|
|
maske[:3, :5] = False
|
|
x, y = lat.get_from_mask(maske)
|
|
plt.scatter(x, y)
|
|
plt.show()
|
|
|
|
|
|
RESOLUTION = 0.1
|
|
|
|
|
|
def image_from_pos(pos_x, pos_y):
|
|
length_x = np.max(pos_x) + RESOLUTION
|
|
length_y = np.max(pos_y) + RESOLUTION
|
|
x_ind = np.arange(0, length_x, RESOLUTION) # angstrom
|
|
y_ind = np.arange(0, length_y, RESOLUTION) # angstrom
|
|
img = np.zeros((x_ind.size, y_ind.size))
|
|
xind = np.searchsorted(x_ind, pos_x)
|
|
yind = np.searchsorted(y_ind, pos_y)
|
|
img[xind, yind] = 1
|
|
return img
|
|
|
|
|
|
def test_img():
|
|
lat = Lattice(10, 10)
|
|
maske = np.ones((10, 10), dtype=bool)
|
|
x, y = lat.get_from_mask(maske)
|
|
img = image_from_pos(x, y)
|
|
plt.imshow(img.T, origin="lower", extent=(0, np.max(x), 0, np.max(y)))
|
|
plt.scatter(x, y)
|
|
plt.show()
|
|
|
|
|
|
def fft(img):
|
|
Z_fft = sfft.fft2(img)
|
|
Z_shift = sfft.fftshift(Z_fft)
|
|
fft_freqx = sfft.fftfreq(img.shape[0], RESOLUTION)
|
|
fft_freqy = sfft.fftfreq(img.shape[1], RESOLUTION)
|
|
fft_freqx_clean = sfft.fftshift(fft_freqx)
|
|
fft_freqy_clean = sfft.fftshift(fft_freqy)
|
|
return fft_freqx_clean, fft_freqy_clean, np.abs(Z_shift) ** 2
|
|
|
|
|
|
def gaussian(img):
|
|
ratio = img.shape[0] / img.shape[1]
|
|
x = np.linspace(-ratio, ratio, img.shape[0])
|
|
y = np.linspace(-1, 1, img.shape[1])
|
|
X, Y = np.meshgrid(x, y)
|
|
sigma = 0.3
|
|
z = (1/(2*np.pi*sigma*sigma) * np.exp(-(X**2/(2*sigma**2)
|
|
+ Y**2/(2*sigma**2))))
|
|
return np.multiply(img, z.T)
|
|
|
|
|
|
def main():
|
|
lat = Lattice(10, 10)
|
|
maske = np.ones((10, 10), dtype=bool)
|
|
x, y = lat.get_from_mask(maske)
|
|
img = image_from_pos(x, y)
|
|
|
|
fig, [axs,axs2] = plt.subplots(2, 3)
|
|
axs[0].imshow(img)
|
|
img = gaussian(img)
|
|
axs[1].imshow(img)
|
|
plt.pause(.1)
|
|
|
|
freqx, freqy, intens = fft(img)
|
|
axs[2].imshow(intens, extent = (np.min(freqx), np.max(
|
|
freqx), np.min(freqy), np.max(freqy)))
|
|
|
|
maske = np.zeros((10, 10), dtype=bool)
|
|
x, y = lat.get_from_mask(maske)
|
|
img = image_from_pos(x, y)
|
|
|
|
axs2[0].imshow(img)
|
|
img = gaussian(img)
|
|
axs2[1].imshow(img)
|
|
plt.pause(.1)
|
|
|
|
freqx, freqy, intens = fft(img)
|
|
axs2[2].imshow(intens, extent = (np.min(freqx), np.max(
|
|
freqx), np.min(freqy), np.max(freqy)))
|
|
plt.show()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|