Added 2d support

cal.py

@@ -1,5 +1,8 @@
 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):
@@ -24,37 +27,154 @@ def mono_2_rutile(c_m, a_m):
     return a_r, c_r
 
 
-def main():
-    offset_a_m = 0.25 - 0.23947
-    offset_c_m = 0.02646
+class Lattice:
 
-    offset_a_r, offset_c_r = mono_2_rutile(offset_c_m, offset_a_m)
+    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
 
-    print("A_r: ", offset_a_r, "C_r: ", offset_c_r)
+    def _get_mono(self, X, Y):
+        offset_a_m = 0.25 - 0.23947
+        offset_c_m = 0.02646
 
-    x = np.arange(10)
-    y = np.arange(10)
+        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)
 
-    atom_x = offset_a_r + X * base_c_r + np.mod(Y, 4) * 0.5 * base_c_r
-    atom_x[np.mod(X, 2) == 0] -= 2 * offset_a_r
 
-    atom_y = offset_c_r + 0.5 * Y * base_a_r
-    atom_y[np.mod(X, 2) == 0] -= 2 * offset_c_r
+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, ax = plt.subplots(1)
-    plt.scatter(atom_x, atom_y)
+    fig, [axs,axs2] = plt.subplots(2, 3)
+    axs[0].imshow(img)
+    img = gaussian(img)
+    axs[1].imshow(img)
+    plt.pause(.1)
 
-    atom_x_rut = X * base_c_r + np.mod(Y, 4) * 0.5 * base_c_r
-    atom_y_rut = Y * 0.5 * base_a_r
+    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)
 
-    plt.scatter(atom_x_rut, atom_y_rut)
-    print(np.square(atom_x-atom_x_rut) + np.square(atom_y-atom_y_rut))
+    axs2[0].imshow(img)
+    img = gaussian(img)
+    axs2[1].imshow(img)
+    plt.pause(.1)
 
-    cryst = np.loadtxt("./crystal_V.xyz", delimiter=" ")
-    print(cryst.shape)
-    plt.scatter(-cryst[:, 0]+5.5*base_c_r, cryst[:, 2])
-    ax.set_aspect(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()