271 lines
9.4 KiB
Python
271 lines
9.4 KiB
Python
import numpy as np
|
|
import scipy
|
|
import scipy.fftpack as sfft
|
|
import scipy.signal
|
|
import tqdm
|
|
from cache import timeit
|
|
|
|
|
|
class SpinImage_Two_Phase:
|
|
resolution = 0.05
|
|
offset = 40
|
|
|
|
def make_list(self, x_pos, y_pos, x_inds, y_inds):
|
|
sigma = .1
|
|
x_ind = np.arange(0, self.length_x, self.resolution)
|
|
y_ind = np.arange(0, self.length_y, self.resolution)
|
|
X, Y = np.meshgrid(x_ind, y_ind, indexing="ij")
|
|
out_list = []
|
|
for x, y, x_ind, y_ind in zip(
|
|
x_pos.flatten(),
|
|
y_pos.flatten(),
|
|
x_inds.flatten(),
|
|
y_inds.flatten(),
|
|
):
|
|
xl, yl, xu, yu = self.clean_bounds(
|
|
x_ind - self.offset,
|
|
y_ind - self.offset,
|
|
x_ind + self.offset,
|
|
y_ind + self.offset,
|
|
X.shape,
|
|
)
|
|
out_list.append(np.exp(-0.5 * ((X[xl:xu, yl:yu] - x) ** 2 +
|
|
(Y[xl:xu, yl:yu] - y) ** 2) / sigma**2))
|
|
#out_list.append(np.ones_like(X[xl:xu, yl:yu]))
|
|
out_list = np.array(out_list)
|
|
return out_list
|
|
|
|
@timeit
|
|
def __init__(self, x_pos_low, y_pos_low, x_pos_high, y_pos_high):
|
|
assert x_pos_low.shape == y_pos_low.shape
|
|
assert x_pos_high.shape == y_pos_high.shape
|
|
assert x_pos_low.shape == x_pos_high.shape
|
|
|
|
offset_shift = self.offset * self.resolution
|
|
|
|
zero_shift_x = np.minimum(np.min(x_pos_low), np.min(x_pos_high))
|
|
zero_shift_y = np.minimum(np.min(y_pos_low), np.min(y_pos_high))
|
|
|
|
x_pos_low = x_pos_low - zero_shift_x + offset_shift
|
|
y_pos_low = y_pos_low - zero_shift_y + offset_shift
|
|
|
|
x_pos_high = x_pos_high - zero_shift_x + offset_shift
|
|
y_pos_high = y_pos_high - zero_shift_y + offset_shift
|
|
|
|
max_len_x = np.maximum(np.max(x_pos_low), np.max(x_pos_high))
|
|
max_len_y = np.maximum(np.max(y_pos_low), np.max(y_pos_high))
|
|
|
|
self.length_x = max_len_x + (self.offset + 1) * self.resolution
|
|
self.length_y = max_len_y + (self.offset + 1) * self.resolution
|
|
self.x_index_low, self.y_index_low = self._stuff(x_pos_low, y_pos_low)
|
|
self.x_index_high, self.y_index_high = self._stuff(
|
|
x_pos_high, y_pos_high)
|
|
|
|
self.buffer_low = self.make_list(
|
|
x_pos_low, y_pos_low, self.x_index_low, self.y_index_low)
|
|
self.buffer_high = self.make_list(
|
|
x_pos_high, y_pos_high, self.x_index_high, self.y_index_high)
|
|
|
|
def clean_bounds(self, xl, yl, xu, yu, shape):
|
|
if xl < 0:
|
|
xl = 0
|
|
if yl < 0:
|
|
yl = 0
|
|
|
|
if xu > shape[0]:
|
|
xu = shape[0]
|
|
if yu > shape[1]:
|
|
yu = shape[1]
|
|
|
|
return xl, yl, xu, yu
|
|
|
|
def _stuff(self, pos_x, pos_y):
|
|
x_ind = np.arange(0, self.length_x, self.resolution)
|
|
y_ind = np.arange(0, self.length_y, self.resolution)
|
|
xind = np.searchsorted(x_ind, pos_x).astype(int)
|
|
yind = np.searchsorted(y_ind, pos_y).astype(int)
|
|
return xind, yind
|
|
|
|
def _apply_mask(self, x, y, buffer, mask):
|
|
for x, y, dat in zip(
|
|
x.flatten()[mask],
|
|
y.flatten()[mask],
|
|
buffer[mask],
|
|
):
|
|
xl, yl, xu, yu = self.clean_bounds(
|
|
x - self.offset,
|
|
y - self.offset,
|
|
x + self.offset,
|
|
y + self.offset,
|
|
self.img.shape,
|
|
)
|
|
self.img[x - self.offset: x + self.offset,
|
|
y - self.offset: y + self.offset] += dat
|
|
|
|
@timeit
|
|
def apply_mask(self, maske):
|
|
mask = np.empty_like(self.x_index_high, dtype=bool)
|
|
print(maske)
|
|
mask[0::2, 0::2] = maske
|
|
mask[1::2, 0::2] = maske
|
|
mask[0::2, 1::2] = maske
|
|
mask[1::2, 1::2] = maske
|
|
mask = mask.flatten()
|
|
self.img = np.zeros(
|
|
(int(self.length_x/self.resolution), int(self.length_y/self.resolution)))
|
|
self._apply_mask(self.x_index_low, self.y_index_low,
|
|
self.buffer_low, mask)
|
|
mask = np.invert(mask)
|
|
self._apply_mask(self.x_index_high, self.y_index_high,
|
|
self.buffer_high, mask)
|
|
|
|
@timeit
|
|
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
|
|
|
|
@timeit
|
|
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)
|
|
|
|
@timeit
|
|
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 + 1
|
|
if size is not None:
|
|
xx = np.maximum(xx, size)
|
|
yy = xx
|
|
|
|
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.linspace(0, self.length_x, self.img.shape[0])
|
|
y = np.linspace(0, self.length_y, self.img.shape[1])
|
|
X, Y = np.meshgrid(x, y)
|
|
X = X - (self.length_x / 2.)
|
|
Y = Y - (self.length_y / 2.)
|
|
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)
|
|
|
|
|
|
class SpinImage:
|
|
resolution = 0.05
|
|
|
|
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 image_from_pos_with_gauss(self, pos_x, pos_y, sigma=1):
|
|
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))
|
|
X, Y = np.meshgrid(y_ind, x_ind)
|
|
for px, py in tqdm.tqdm(zip(pos_x.flatten(), pos_y.flatten())):
|
|
img += np.exp(-0.5 * ((X - px) ** 2 + (Y - py) ** 2) / sigma**2)
|
|
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 + 1
|
|
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 pad_it(self, x_size, y_size):
|
|
h = self.img.shape[0]
|
|
w = self.img.shape[1]
|
|
xx = x_size
|
|
yy = y_size
|
|
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)
|