jacob/FFT
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

2d test

Browse Source
This commit is contained in:
Jacob Holder 2023-02-18 02:06:46 +01:00
parent 57152e4054
commit 524a4f8174
3 changed files with 269 additions and 104 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

28
2d_fourie/lattices.py
View File

@ -15,6 +15,8 @@ class Lattice:
def get_from_mask(self, maske): def get_from_mask(self, maske):
pass pass
def reci(self):
pass
class SCC_Lattice(Lattice): class SCC_Lattice(Lattice):
def __init__(self, x_len, y_len): def __init__(self, x_len, y_len):
@ -25,6 +27,12 @@ class SCC_Lattice(Lattice):
def get_from_mask(self, maske): def get_from_mask(self, maske):
return self.X, self.Y return self.X, self.Y
def reci(self):
x = np.arange(-3,3) * 0.2
y = np.arange(-3,3) * 0.2
X,Y = np.meshgrid(x, y)
return [(X,Y)]
class VO2_Lattice(Lattice): class VO2_Lattice(Lattice):
base_a_m = 5.75 base_a_m = 5.75
@ -52,7 +60,7 @@ class VO2_Lattice(Lattice):
offset_a_m = 0.25 - 0.23947 offset_a_m = 0.25 - 0.23947
offset_c_m = 0.02646 offset_c_m = 0.02646
offset_a_r, offset_c_r = self.mono_2_rutile(offset_c_m, offset_a_m) offset_a_r, offset_c_r = self._mono_2_rutile(offset_c_m, offset_a_m)
print("A_r: ", offset_a_r, "C_r: ", offset_c_r) print("A_r: ", offset_a_r, "C_r: ", offset_c_r)
@ -95,3 +103,21 @@ class VO2_Lattice(Lattice):
inplace_pos_x[mask] = self.atom_x_mono[mask] inplace_pos_x[mask] = self.atom_x_mono[mask]
inplace_pos_y[mask] = self.atom_y_mono[mask] inplace_pos_y[mask] = self.atom_y_mono[mask]
return inplace_pos_x, inplace_pos_y return inplace_pos_x, inplace_pos_y
def reci_rutile(self):
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(self):
x, y = self.reci_rutile()
return x + 0.1083, y + 0.1719
def reci_mono_2(self):
x, y = self.reci_rutile()
return x - 0.1083, y + 0.1719
def reci(self):
return [self.reci_rutile(), self.reci_mono(), self.reci_mono_2()]

271
2d_fourie/main.py
View File

@ -1,9 +1,7 @@
from lattices import SCC_Lattice from lattices import SCC_Lattice, VO2_Lattice
from spin_image import SpinImage
import numpy as np import numpy as np
import matplotlib.pyplot as plt import matplotlib.pyplot as plt
import scipy.fftpack as sfft
import matplotlib.patches as patches import matplotlib.patches as patches
import matplotlib import matplotlib
import scipy import scipy
@ -11,109 +9,89 @@ import scipy.signal
import tqdm import tqdm
class SpinImage: class Plotter:
resolution = 0.1 def __init__(self, lat):
self.lattice = lat
def __init__(self, x_pos, y_pos): def reduce(self, arr):
self.length_x = np.max(x_pos) + self.resolution arr = np.array(arr)
self.length_y = np.max(y_pos) + self.resolution arr = arr.flatten()
self.img = self.image_from_pos(x_pos, y_pos) return np.mean(arr)
# return np.sum(arr[np.argpartition(arr, -8)[-8:]])
def image_from_pos(self, pos_x, pos_y): def extract_rect(self, img, x, y, x_index, y_index):
x_ind = np.arange(0, self.length_x, self.resolution) # angstrom length_2 = 0.01
y_ind = np.arange(0, self.length_y, self.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 fft(self): pos_x_lower = x - length_2
Z_fft = sfft.fft2(self.img) pos_x_upper = x + length_2
Z_shift = sfft.fftshift(Z_fft)
fft_freqx = sfft.fftfreq(self.img.shape[0], self.resolution)
fft_freqy = sfft.fftfreq(self.img.shape[1], self.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 pad_it_square(self, additional_pad=0): pos_y_lower = y - length_2
h = self.img.shape[0] pos_y_upper = y + length_2
w = self.img.shape[1]
print(h, w)
xx = np.maximum(h, w) + 2 * additional_pad
yy = xx
self.length_x = xx * self.resolution
self.length_y = yy * self.resolution
print("Pad to: ", xx, yy)
a = (xx - h) // 2 x_lower = np.searchsorted(x_index, pos_x_lower)
aa = xx - a - h x_upper = np.searchsorted(x_index, pos_x_upper)
b = (yy - w) // 2 y_lower = np.searchsorted(y_index, pos_y_lower)
bb = yy - b - w y_upper = np.searchsorted(y_index, pos_y_upper)
self.img = np.pad(self.img, pad_width=( # fix different number of spins possible
(a, aa), (b, bb)), mode="constant") if x_upper - x_lower < 10:
x_upper += 1
if y_upper - y_lower < 10:
y_upper += 1
return img[y_lower:y_upper, x_lower:x_upper]
def gaussian(self, sigma): def helper(self, ax, freqx, freqy, intens):
x = np.arange(-self.length_x/2, reci_lattice = self.lattice.reci()
self.length_x/2, self.resolution) for tup, col in zip(reci_lattice, ["r", "b", "g"]):
y = np.arange(-self.length_y/2, point_x, point_y = tup
self.length_y/2, self.resolution) point_x = point_x.flatten()
X, Y = np.meshgrid(x, y) point_y = point_y.flatten()
z = ( for px, py in zip(point_x, point_y):
1 / (2 * np.pi * sigma * sigma) rect = self.rect_at_point(px, py, col)
* np.exp(-(X**2 / (2 * sigma**2) + Y**2 / (2 * sigma**2))) ax.add_patch(rect)
sum = self.extract_rect(intens, px, py, freqx, freqy)
ax.text(
px, py, f"{self.reduce(sum):2.2}", clip_on=True
)
return intens
def rect_at_point(self, x, y, color):
length_2 = 0.01
rect = patches.Rectangle(
(x - length_2, y - length_2),
2 * length_2,
2 * length_2,
linewidth=1,
edgecolor=color,
facecolor="none",
) )
self.img = np.multiply(self.img, z.T) return rect
def plot(self, ax, scale=None): def plot(self, freqx, freqy, intens, ax_log=None, ax_lin=None, vmax=None):
if scale is None: if ax_log:
ax.imshow(self.img) intens = self.helper(ax_lin, freqx, freqy, intens)
else: t = ax_log.imshow(
quad = np.ones((int(scale/self.resolution), intens,
int(scale/self.resolution))) extent=(np.min(freqx), np.max(freqx),
img = scipy.signal.convolve2d(self.img, quad) np.min(freqy), np.max(freqy)),
ax.imshow(img) norm=matplotlib.colors.LogNorm(vmin=10, vmax=vmax),
cmap="viridis",
def blur(self, sigma): origin="lower"
self.img = scipy.ndimage.gaussian_filter(self.img, sigma) )
plt.colorbar(t, ax=ax_log)
self.helper(ax_log, freqx, freqy, intens)
def plot(freqx, freqy, intens, ax_log=None, ax_lin=None): if ax_lin:
#point_x, point_y = reci_rutile() intens = self.helper(ax_lin, freqx, freqy, intens)
# for px, py in zip(point_x, point_y): t = ax_lin.imshow(
# rect = rect_at_point(px, py, "r") intens,
# ax.add_patch(rect) extent=(np.min(freqx), np.max(freqx),
# ax.text( np.min(freqy), np.max(freqy)),
# px, py, f"{reduce(extract_rect(intens, px, py, freqx, freqy)):2.2}", clip_on=True vmax=vmax,
# ) cmap="viridis",
origin="lower"
#point_x, point_y = reci_mono() )
# for px, py in zip(point_x, point_y): plt.colorbar(t, ax=ax_lin)
# rect = rect_at_point(px, py, "b")
# ax.add_patch(rect)
# ax.text(
# px, py, f"{reduce(extract_rect(intens, px, py, freqx, freqy)):2.2}", clip_on=True
# )
if ax_log:
t = ax_log.imshow(
intens,
extent=(np.min(freqx), np.max(freqx),
np.min(freqy), np.max(freqy)),
norm=matplotlib.colors.LogNorm(),
cmap="viridis"
)
plt.colorbar(t)
if ax_lin:
t = ax_lin.imshow(
intens,
extent=(np.min(freqx), np.max(freqx),
np.min(freqy), np.max(freqy)),
cmap="viridis"
)
plt.colorbar(t)
def rotate(x, y, angle): def rotate(x, y, angle):
@ -122,27 +100,114 @@ def rotate(x, y, angle):
def test_square(): def test_square():
lat = SCC_Lattice(40, 40) LEN = 40
pos_x, pos_y = lat.get_from_mask(None) #lat = SCC_Lattice(LEN, LEN)
pos_x, pos_y = rotate(pos_x, pos_y,30) lat = VO2_Lattice(LEN, LEN)
plot = Plotter(lat)
pos_x, pos_y = lat.get_from_mask(np.zeros((40, 40)))
#pos_x, pos_y = rotate(pos_x, pos_y, 30)
si = SpinImage(pos_x, pos_y) si = SpinImage(pos_x, pos_y)
fig, axs = plt.subplots(2, 2) fig, axs = plt.subplots(2, 2)
si.pad_it_square(10) si.pad_it_square(10)
si.plot(axs[0, 0], 2) si.plot(axs[0, 0], 2)
si.gaussian(300) # si.gaussian(LEN)
# si.blur(3) # si.blur(3)
si.plot(axs[0, 1], 2) si.plot(axs[0, 1], 2)
plt.pause(0.1) plt.pause(0.1)
fx, fy, intens = si.fft() fx, fy, intens = si.fft()
plot(fx, fy, intens, axs[1, 0], axs[1, 1]) plot.plot(fx, fy, intens, axs[1, 0], axs[1, 1])
print("Done") print("Done")
plt.savefig("test.png") plt.savefig("test.png")
plt.show() plt.show()
def test_mixed():
LEN = 40
lat = VO2_Lattice(LEN, LEN)
plot = Plotter(lat)
pos_x, pos_y = lat.get_from_mask(np.zeros((40, 40)))
si = SpinImage(pos_x, pos_y)
si.pad_it_square(10)
fx, fy, intens_rutile = si.fft()
pos_x, pos_y = lat.get_from_mask(np.ones((40, 40)))
si = SpinImage(pos_x, pos_y)
si.pad_it_square(10)
fx, fy, intens_mono = si.fft()
mask_misk = np.ones((40, 40))
ind = np.arange(mask_misk.size)
np.random.shuffle(ind)
mask_misk[np.unravel_index(ind[:800], (40, 40))] = False
pos_x, pos_y = lat.get_from_mask(mask_misk)
si = SpinImage(pos_x, pos_y)
si.pad_it_square(10)
fx, fy, intens_mixed = si.fft()
fig, axs = plt.subplots(2, 3)
plot.plot(freqx=fx, freqy=fy, intens=intens_rutile,
ax_log=axs[0, 0], ax_lin=axs[1, 0], vmax=10e7)
plot.plot(freqx=fx, freqy=fy, intens=intens_mono,
ax_log=axs[0, 2], ax_lin=axs[1, 2], vmax=10e7)
plot.plot(freqx=fx, freqy=fy, intens=intens_mixed,
ax_log=axs[0, 1], ax_lin=axs[1, 1], vmax=10e7)
print(np.sum(intens_mono), np.sum(intens_rutile), np.sum(intens_mixed))
for ax in axs.flatten():
ax.set_xlim(-1, 1)
ax.set_ylim(-1, 1)
plt.show()
def random():
LEN = 40
lat = VO2_Lattice(LEN, LEN)
plot = Plotter(lat)
maske = np.zeros((LEN, LEN))
ind = np.arange(LEN * LEN)
np.random.shuffle(ind)
reci_lattice = lat.reci()
out = [[] for x in range(len(reci_lattice))]
percentage = []
counter = 0
for i in tqdm.tqdm(ind):
maske[np.unravel_index(i, (LEN, LEN))] = True
counter += 1
if np.mod(counter, 20) != 0:
continue
pos_x, pos_y = lat.get_from_mask(maske)
si = SpinImage(pos_x, pos_y)
si.pad_it_square(10)
si.gaussian(LEN)
fx, fy, intens = si.fft()
for tup, lis in zip(reci_lattice, out):
point_x, point_y = tup
point_x = point_x.flatten()
point_y = point_y.flatten()
sum = 0.
for px, py in zip(point_x, point_y):
sum += np.sum(plot.extract_rect(intens, px, py, fx, fy))
lis.append(sum)
percentage.append(np.mean(maske))
for o in out:
plt.scatter(percentage, o/o[0])
plt.plot([0,1], [o[0], o[-1]])
plt.show()
if __name__ == "__main__": if __name__ == "__main__":
test_square() # test_square()
# test_mixed()
random()
# def test_lattice(): # def test_lattice():
# lat = VO2_Lattice(10, 10) # lat = VO2_Lattice(10, 10)
# maske = np.zeros((10, 10), dtype=bool) # maske = np.zeros((10, 10), dtype=bool)

74
2d_fourie/spin_image.py Normal file
View File

@ -0,0 +1,74 @@
import numpy as np
import scipy.fftpack as sfft
class SpinImage:
resolution = 0.1
def __init__(self, x_pos, y_pos):
x_pos = x_pos - np.min(x_pos)
y_pos = y_pos - np.min(y_pos)
self.length_x = np.max(x_pos) + self.resolution
self.length_y = np.max(y_pos) + self.resolution
self.img = self.image_from_pos(x_pos, y_pos)
def image_from_pos(self, pos_x, pos_y):
x_ind = np.arange(0, self.length_x, self.resolution) # angstrom
y_ind = np.arange(0, self.length_y, self.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 fft(self):
Z_fft = sfft.fft2(self.img)
Z_shift = sfft.fftshift(Z_fft)
fft_freqx = sfft.fftfreq(self.img.shape[0], self.resolution)
fft_freqy = sfft.fftfreq(self.img.shape[1], self.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 pad_it_square(self, additional_pad=0):
h = self.img.shape[0]
w = self.img.shape[1]
print(h, w)
xx = np.maximum(h, w) + 2 * additional_pad
yy = xx
self.length_x = xx * self.resolution
self.length_y = yy * self.resolution
print("Pad to: ", xx, yy)
a = (xx - h) // 2
aa = xx - a - h
b = (yy - w) // 2
bb = yy - b - w
self.img = np.pad(self.img, pad_width=(
(a, aa), (b, bb)), mode="constant")
def gaussian(self, sigma):
x = np.arange(-self.length_x / 2,
self.length_x / 2, self.resolution)
y = np.arange(-self.length_y / 2,
self.length_y / 2, self.resolution)
X, Y = np.meshgrid(x, y)
z = (
1 / (2 * np.pi * sigma * sigma)
* np.exp(-(X ** 2 / (2 * sigma ** 2) + Y ** 2 / (2 * sigma ** 2)))
)
self.img = np.multiply(self.img, z.T)
def plot(self, ax, scale=None):
if scale is None:
ax.imshow(self.img)
else:
quad = np.ones((int(scale / self.resolution),
int(scale / self.resolution)))
img = scipy.signal.convolve2d(self.img, quad)
ax.imshow(img)
def blur(self, sigma):
self.img = scipy.ndimage.gaussian_filter(self.img, sigma)