import numpy as np import matplotlib.pyplot as plt import scipy.fftpack as sfft import matplotlib.patches as patches import matplotlib 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 CMAP = "Greys" 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.5 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 padding(array, xx, yy): """ :param array: numpy array :param xx: desired height :param yy: desirex width :return: padded array """ h = array.shape[0] w = array.shape[1] a = (xx - h) // 2 aa = xx - a - h b = (yy - w) // 2 bb = yy - b - w return np.pad(array, pad_width=((a, aa), (b, bb)), mode="constant") def rect_at_point(x, y, color): length_2 = 0.08 rect = patches.Rectangle( (x - length_2, y - length_2), 2 * length_2, 2 * length_2, linewidth=1, edgecolor=color, facecolor="none", ) return rect def reci_rutile(): x = np.arange(-2, 3) y = np.arange(-2, 3) X, Y = np.meshgrid(x, y) return (X * 0.22 + Y * 0.44).flatten(), (X * 0.349).flatten() def reci_mono(): x, y = reci_rutile() return x + 0.1083, y + 0.1719 def draw_big_val_rect(img, x, y, x_index, y_index): length_2 = 0.08 pos_x_lower = x - length_2 pos_x_upper = x + length_2 pos_y_lower = y - length_2 pos_y_upper = y + length_2 x_lower = np.searchsorted(x_index, pos_x_lower) x_upper = np.searchsorted(x_index, pos_x_upper) y_lower = np.searchsorted(y_index, pos_y_lower) y_upper = np.searchsorted(y_index, pos_y_upper) img[y_lower:y_upper, x_lower:x_upper] = 1e4 return img def extract_rect(img, x, y, x_index, y_index): length_2 = 0.08 pos_x_lower = x - length_2 pos_x_upper = x + length_2 pos_y_lower = y - length_2 pos_y_upper = y + length_2 x_lower = np.searchsorted(x_index, pos_x_lower) x_upper = np.searchsorted(x_index, pos_x_upper) y_lower = np.searchsorted(y_index, pos_y_lower) y_upper = np.searchsorted(y_index, pos_y_upper) return img[y_lower:y_upper, x_lower:x_upper] def main(): FFT_KWARGS = {"norm": matplotlib.colors.LogNorm(vmin=1), "cmap": "Greys"} IMSHOW_WARGS = {"cmap": "Greys"} SIZE = 601 lat = Lattice(10, 10) maske = np.ones((10, 10), dtype=bool) x, y = lat.get_from_mask(maske) img = image_from_pos(x, y) img = padding(img, SIZE, SIZE) fig, [axs, axs2] = plt.subplots(2, 3) axs[0].imshow(img, **IMSHOW_WARGS) img = gaussian(img) axs[1].imshow(img, **IMSHOW_WARGS) plt.pause(0.1) freqx, freqy, intens = fft(img) axs[2].imshow( intens, extent=(np.min(freqx), np.max(freqx), np.min(freqy), np.max(freqy)), **FFT_KWARGS, ) intens_rut = intens maske = np.zeros((10, 10), dtype=bool) x, y = lat.get_from_mask(maske) img = image_from_pos(x, y) img = padding(img, SIZE, SIZE) axs2[0].imshow(img, **IMSHOW_WARGS) img = gaussian(img) axs2[1].imshow(img, **IMSHOW_WARGS) plt.pause(0.1) freqx, freqy, intens = fft(img) axs2[2].imshow( intens, extent=(np.min(freqx), np.max(freqx), np.min(freqy), np.max(freqy)), **FFT_KWARGS, ) # Create a Rectangle patch # Add the patch to the Axes point_x, point_y = reci_rutile() for px, py in zip(point_x, point_y): rect = rect_at_point(px, py, "r") axs[2].add_patch(rect) axs[2].text( px, py, f"{np.sum(extract_rect(intens_rut, px, py, freqx, freqy)):0.2}" ) rect = rect_at_point(px, py, "r") axs2[2].add_patch(rect) axs2[2].text( px, py, f"{np.sum(extract_rect(intens, px, py, freqx, freqy)):0.2}" ) point_x, point_y = reci_mono() for px, py in zip(point_x, point_y): # rect = rect_at_point(px, py,"b") # axs[2].add_patch(rect) rect = rect_at_point(px, py, "b") axs2[2].add_patch(rect) axs2[2].text( px, py, f"{np.sum(extract_rect(intens, px, py, freqx, freqy)):0.2}" ) axs[2].set_xlim(-1.0, 1.0) axs[2].set_ylim(-1.0, 1.0) axs2[2].set_xlim(-1.0, 1.0) axs2[2].set_ylim(-1.0, 1.0) plt.figure() diff = intens_rut - intens plt.imshow( diff, extent=(np.min(freqx), np.max(freqx), np.min(freqy), np.max(freqy)) ) plt.xlim(-1.0, 1.0) plt.ylim(-1.0, 1.0) plt.show() if __name__ == "__main__": main()