88 lines
3.0 KiB
Python
88 lines
3.0 KiB
Python
import numpy as np
|
|
import scipy
|
|
import scipy.fftpack as sfft
|
|
import scipy.signal
|
|
|
|
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, size=None):
|
|
h = self.img.shape[0]
|
|
w = self.img.shape[1]
|
|
xx = np.maximum(h, w) + 2 * additional_pad
|
|
if size is not None:
|
|
xx = np.maximum(xx, size)
|
|
yy = xx
|
|
self.length_x = xx * self.resolution
|
|
self.length_y = yy * self.resolution
|
|
|
|
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 percentage_gaussian(self, mask, sigma):
|
|
x = np.linspace(-self.length_x / 2,
|
|
self.length_x / 2, mask.shape[0])
|
|
y = np.linspace(-self.length_y / 2,
|
|
self.length_y / 2, mask.shape[1])
|
|
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)))
|
|
)
|
|
return np.multiply(mask, z.T)
|
|
|
|
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)
|