138 lines
3.4 KiB
Python
138 lines
3.4 KiB
Python
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
|
|
|
|
def analysis(y, RESOLUTION):
|
|
fft = np.fft.fft(y)
|
|
fft_clean = np.fft.fftshift(fft)
|
|
fft_freq = np.fft.fftfreq(y.size, RESOLUTION)
|
|
fft_freq_clean = np.fft.fftshift(fft_freq)
|
|
|
|
return fft_freq_clean, np.abs(fft_clean) ** 2
|
|
|
|
|
|
def play_1d():
|
|
RESOLUTION = 0.001
|
|
LENGTH = 10000
|
|
|
|
x = np.arange(0, LENGTH, RESOLUTION) # angstrom
|
|
y = np.zeros_like(x)
|
|
|
|
pos_mono = np.arange(0, x.size, 2890)
|
|
pos_mono = (
|
|
pos_mono + np.random.normal(size=pos_mono.shape, loc=0, scale=10)
|
|
).astype(int)
|
|
|
|
pos_rut = np.arange(0, x.size, 5780)
|
|
|
|
pos_rut = np.append(pos_rut, pos_rut - 3160)
|
|
|
|
# pos_rut = (pos_rut + np.random.normal(size=pos_rut.shape, loc=0, scale=10)).astype(int)
|
|
|
|
y[pos_rut] = 1
|
|
y[pos_rut + 1] = 1
|
|
y[pos_rut + 2] = 1
|
|
|
|
# y = np.sin(x)
|
|
|
|
fig, axs = plt.subplots(3, 1)
|
|
ax = axs[0]
|
|
ax.plot(x, y)
|
|
|
|
ax = axs[1]
|
|
fft_x, fft_y = analysis(y, RESOLUTION)
|
|
ax.plot(fft_x, fft_y)
|
|
ax.plot([1.0 / 5.78, 1.0 / 2.62, 1.0 / 3.16], [0, 0, 0], "kx")
|
|
ax.plot([2.0 / 5.78, 2.0 / 2.62, 2.0 / 3.16], [0, 0, 0], "rx")
|
|
ax.plot([3.0 / 5.78, 3.0 / 2.62, 3.0 / 3.16], [0, 0, 0], "bx")
|
|
ax.set_xlim(0, 3)
|
|
|
|
|
|
def from_mask(mask):
|
|
pos_mono = np.arange(0, mask.size * 2.89 * 2, 2.89)
|
|
|
|
pos_mono[::2][mask] -= 0.27
|
|
pos_mono[1::2][mask] += 0.27
|
|
|
|
return pos_mono
|
|
|
|
|
|
def image_from_pos(pos):
|
|
RESOLUTION = 0.001
|
|
LENGTH = 1000000
|
|
|
|
x = np.arange(0, LENGTH, RESOLUTION) # angstrom
|
|
y = np.zeros_like(x)
|
|
ind = np.searchsorted(x, pos)
|
|
if np.any(ind > LENGTH):
|
|
print("overflow")
|
|
ind = ind[ind < LENGTH]
|
|
y[ind] = 1
|
|
|
|
sigma = 500
|
|
mu = int(LENGTH / 2)
|
|
gaussian = (
|
|
1
|
|
/ (sigma * np.sqrt(2 * np.pi))
|
|
* np.exp(-((x - mu) ** 2) / (2 * sigma * sigma))
|
|
)
|
|
# y = np.multiply(y, gaussian)
|
|
|
|
return x, y
|
|
|
|
|
|
def plot_img(x, y, ax):
|
|
ax.plot(x, y)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
RESOLUTION = 0.001
|
|
print("Done")
|
|
|
|
LENGTH = 1000
|
|
mask_h = np.ones(LENGTH).astype(bool)
|
|
pos_h = from_mask(mask_h)
|
|
x, img_h = image_from_pos(pos_h)
|
|
fftx, ffty_h = analysis(img_h, RESOLUTION)
|
|
|
|
mask_l = np.zeros(LENGTH).astype(bool)
|
|
pos_l = from_mask(mask_l)
|
|
x, img_l = image_from_pos(pos_l)
|
|
fftx, ffty_l = analysis(img_l, RESOLUTION)
|
|
print("Done")
|
|
|
|
mask_mixed = np.zeros(LENGTH).astype(bool)
|
|
ind = (np.random.rand(400) * (LENGTH - 1)).astype(int)
|
|
mask_mixed[ind] = True
|
|
pos_mixed = from_mask(mask_mixed)
|
|
x, img_mixed = image_from_pos(pos_mixed)
|
|
fftx, ffty_mixed = analysis(img_mixed, RESOLUTION)
|
|
|
|
print("Done")
|
|
mask_near = np.zeros(LENGTH).astype(bool)
|
|
ind = (np.random.rand(50) * (LENGTH - 1)).astype(int)
|
|
#for i in range(1, 8):
|
|
# ind = np.append(ind, ind+i)
|
|
print("Done")
|
|
|
|
|
|
mask_near[ind] = True
|
|
pos_near = from_mask(mask_near)
|
|
x, img_near = image_from_pos(pos_near)
|
|
fftx, ffty_near = analysis(img_near, RESOLUTION)
|
|
|
|
fig, axs = plt.subplots(4, 1)
|
|
plot_img(fftx, ffty_h, axs[0])
|
|
plot_img(fftx, ffty_l, axs[1])
|
|
plot_img(fftx, ffty_mixed, axs[2])
|
|
plot_img(fftx, ffty_near, axs[3])
|
|
for ax in axs:
|
|
ax.plot([1.0 / 5.78, 1.0 / 2.62, 1.0 / 3.16], [0, 0, 0], "kx")
|
|
ax.plot([2.0 / 5.78, 2.0 / 2.62, 2.0 / 3.16], [0, 0, 0], "rx")
|
|
ax.plot([3.0 / 5.78, 3.0 / 2.62, 3.0 / 3.16], [0, 0, 0], "bx")
|
|
ax.set_xlim(0, 3)
|
|
# play_1d()
|
|
plt.show()
|
|
print("Done")
|
|
pass
|