import numpy as np
from tqdm import tqdm
from scipy.ndimage.filters import gaussian_filter
from scipy.ndimage import rotate
from scipy.signal import convolve2d
from skimage import morphology
import dask.array as da
from hyperspy.misc.utils import isiterable
from pixstem.pixelated_stem_class import PixelatedSTEM
from pixstem.pixelated_stem_class import LazyPixelatedSTEM
import pixstem.ransac_ellipse_tools as ret
def _get_elliptical_disk(xx, yy, x, y, semi_len0, semi_len1, rotation):
"""
Parameters
----------
xx, yy : 2D NumPy array
x, y : float
semi_len0, semi_len1 : float
rotation : float
In radians
Returns
-------
ellipse_array : 2D NumPy array
Examples
--------
>>> from hyperspy.signals import Signal2D
>>> import pixstem.make_diffraction_test_data as mdtd
>>> s = Signal2D(np.zeros((110, 130)))
>>> s.axes_manager[0].offset, s.axes_manager[1].offset = -50, -80
>>> xx, yy = np.meshgrid(
... s.axes_manager.signal_axes[0].axis,
... s.axes_manager.signal_axes[1].axis)
>>> ellipse_data = mdtd._get_elliptical_disk(xx, yy, 10, -10, 12, 18, 1.5)
>>> s.data += ellipse_data
>>> s.plot()
"""
rot = rotation
xx0 = xx - x
yy0 = yy - y
semi_len02 = semi_len0 * semi_len0
semi_len12 = semi_len1 * semi_len1
z0 = ((xx0 * np.cos(rot) + yy0 * np.sin(rot)) ** 2) / semi_len02
z1 = ((xx0 * np.sin(rot) - yy0 * np.cos(rot)) ** 2) / semi_len12
zz = z0 + z1
elli_mask = zz <= 1.
return elli_mask
def _get_elliptical_ring(xx, yy, x, y, semi_len0, semi_len1, rotation, lw_r=1):
"""
Parameters
----------
xx, yy : 2D NumPy array
x, y : float
semi_len0, semi_len1 : float
rotation : float
In radians
lw_r : optional, default 1
Returns
-------
ellipse_array : 2D NumPy array
Examples
--------
>>> from hyperspy.signals import Signal2D
>>> import pixstem.make_diffraction_test_data as mdtd
>>> s = Signal2D(np.zeros((110, 130)))
>>> s.axes_manager[0].offset, s.axes_manager[1].offset = -50, -80
>>> xx, yy = np.meshgrid(
... s.axes_manager.signal_axes[0].axis,
... s.axes_manager.signal_axes[1].axis)
>>> ellipse_data = mdtd._get_elliptical_ring(
... xx, yy, 10, -10, 12, 18, 1.5, 2)
>>> s.data += ellipse_data
>>> s.plot()
"""
mask_outer = _get_elliptical_disk(
xx, yy, x, y, semi_len0 + lw_r, semi_len1 + lw_r, rotation)
mask_inner = _get_elliptical_disk(
xx, yy, x, y, semi_len0 - lw_r, semi_len1 - lw_r, rotation)
ellipse = np.logical_xor(mask_outer, mask_inner)
ellipse = ellipse.astype('uint32')
return ellipse
[docs]class EllipseRing(object):
def __init__(
self, xx, yy, x0, y0, semi_len0, semi_len1, rotation,
intensity, lw_r=1):
"""
Parameters
----------
xx, yy : int
Size of the image
x, y : float
Centre positions for the ellipse
semi_len0, semi_len1 : float
rotation : float
In radians
lw_r : optional, default 1
Methods
-------
get_signal
Generate numpy array with the ellipse ring
Examples
--------
>>> import pixstem.make_diffraction_test_data as mdtd
>>> ellipse = mdtd.EllipseRing(
... xx=20, yy=30, x0=10, y0=15, semi_len0=4, semi_len1=6,
... rotation=1.57, intensity=10, lw_r=2)
"""
self.x0, self.y0 = x0, y0
self.semi_len0 = semi_len0
self.semi_len1 = semi_len1
self.rotation = rotation
self.intensity = intensity
self.xx, self.yy = xx, yy
self.lw_r = lw_r
def __repr__(self):
return '<%s, ((x0, y0): (%s, %s), (sl0, sl1): (%s, %s),' \
' r: %s, I: %s, lw: %s)>' % (
self.__class__.__name__,
self.x0, self.y0, self.semi_len0, self.semi_len1,
self.rotation, self.intensity, self.lw_r,
)
[docs] def get_signal(self):
ellipse = _get_elliptical_ring(
xx=self.xx, yy=self.yy, x=self.x0, y=self.y0,
semi_len0=self.semi_len0, semi_len1=self.semi_len1,
rotation=self.rotation, lw_r=self.lw_r)
self.ellipse = ellipse * self.intensity
return self.ellipse
[docs]class EllipseDisk(object):
def __init__(
self, xx, yy, x0, y0, semi_len0, semi_len1, rotation,
intensity):
"""
Parameters
----------
xx, yy : int
Size of the image
x, y : float
Centre positions for the ellipse
semi_len0, semi_len1 : float
rotation : float
In radians
lw_r : optional, default 1
Methods
-------
get_signal
Generate numpy array with the elliptical disk
Examples
--------
>>> import pixstem.make_diffraction_test_data as mdtd
>>> ellipse_disk = mdtd.EllipseDisk(
... xx=20, yy=30, x0=10, y0=15, semi_len0=4, semi_len1=6,
... rotation=1.57, intensity=10)
"""
self.x0, self.y0 = x0, y0
self.semi_len0 = semi_len0
self.semi_len1 = semi_len1
self.rotation = rotation
self.intensity = intensity
self.xx, self.yy = xx, yy
def __repr__(self):
return '<%s, ((x0, y0): (%s, %s), (sl0, sl1): (%s, %s),' \
' r: %s, I: %s)>' % (
self.__class__.__name__,
self.x0, self.y0, self.semi_len0, self.semi_len1,
self.rotation, self.intensity,
)
[docs] def get_signal(self):
ellipse = _get_elliptical_disk(
xx=self.xx, yy=self.yy, x=self.x0, y=self.y0,
semi_len0=self.semi_len0, semi_len1=self.semi_len1,
rotation=self.rotation)
self.ellipse = ellipse * self.intensity
return self.ellipse
[docs]class Circle(object):
def __init__(self, xx, yy, x0, y0, r, intensity, scale, lw=None):
self.x0 = x0
self.y0 = y0
self.r = r
self.intensity = intensity
self.lw = lw
self.circle = (yy - self.y0) ** 2 + (xx - self.x0) ** 2
self.mask_outside_r(scale)
self.get_centre_pixel(xx, yy, scale)
def __repr__(self):
return '<%s, (r: %s, (x0, y0): (%s, %s), I: %s)>' % (
self.__class__.__name__,
self.r, self.x0, self.y0, self.intensity,
)
[docs] def mask_outside_r(self, scale):
if self.lw is None:
indices = self.circle >= (self.r + scale) ** 2
else:
indices = self.circle >= (self.r + self.lw + scale) ** 2
self.circle[indices] = 0
[docs] def centre_on_image(self, xx, yy):
if self.x0 < xx[0][0] or self.x0 > xx[0][-1]:
return False
elif self.y0 < yy[0][0] or self.y0 > yy[-1][-1]:
return False
else:
return True
[docs] def get_centre_pixel(self, xx, yy, scale):
"""
This function sets the indices for the pixels on which the centre
point is. Because the centre point can sometimes be exactly on the
boundary of two pixels, the pixels are held in a list. One
list for x (self.centre_x_pixels) and one for y
(self.centre_x_pixels). If the centre is outside the image, the
lists will be empty.
"""
if self.centre_on_image(xx, yy):
x1 = np.where(xx > (self.x0 - 0.5 * scale))[1][0]
x2 = np.where(xx < (self.x0 + 0.5 * scale))[1][-1]
self.centre_x_pixels = [x1, x2]
y1 = np.where(yy > (self.y0 - 0.5 * scale))[0][0]
y2 = np.where(yy < (self.y0 + 0.5 * scale))[0][-1]
self.centre_y_pixels = [y1, y2]
else:
self.centre_x_pixels = []
self.centre_y_pixels = []
[docs] def update_axis(self, xx, yy):
self.circle = (xx - self.x0) ** 2 + (yy - self.y0) ** 2
self.mask_outside_r()
[docs]class Disk(object):
"""
Disk object, with outer edge of the ring at r
"""
def __init__(self, xx, yy, scale, x0, y0, r, intensity):
self.z = Circle(xx, yy, x0, y0, r, intensity, scale)
self.z.set_uniform_intensity()
self.set_centre_intensity()
def __repr__(self):
return '<%s, (r: %s, (x0, y0): (%s, %s), I: %s)>' % (
self.__class__.__name__,
self.z.r,
self.z.x0,
self.z.y0,
self.z.intensity,
)
[docs] def set_centre_intensity(self):
"""
Sets the intensity of the centre pixels to I. Coordinates are
self.z.circle[y, x], due to how numpy works.
"""
for x in self.z.centre_x_pixels:
for y in self.z.centre_y_pixels:
self.z.circle[y, x] = self.z.intensity # This is correct
[docs] def get_signal(self):
return (self.z.circle)
[docs] def update_axis(self, xx, yy):
self.z.update_axis(xx, yy)
self.z.set_uniform_intensity()
self.set_centre_intensity()
[docs]class Ring(object):
"""Ring object, with outer edge of the ring at r+lr, and inner r-lr.
The radius of the ring is defined as in the middle of the line making
up the ring.
"""
def __init__(self, xx, yy, scale, x0, y0, r, intensity, lr):
if lr > r:
raise ValueError('Ring line width too big'.format(lr, r))
self.lr = lr
self.lw = 1 + 2 * lr # scalar line width of the ring
self.z = Circle(xx, yy, x0, y0, r, intensity, scale, lw=lr)
self.mask_inside_r(scale)
self.z.set_uniform_intensity()
def __repr__(self):
return '<%s, (r: %s, (x0, y0): (%s, %s), I: %s)>' % (
self.__class__.__name__, self.z.r, self.z.x0, self.z.y0,
self.z.intensity,)
[docs] def mask_inside_r(self, scale):
indices = self.z.circle < (self.z.r - self.lr) ** 2
self.z.circle[indices] = 0
[docs] def get_signal(self):
return self.z.circle
[docs] def update_axis(self, xx, yy):
self.z.update_axis(xx, yy)
self.mask_inside_r()
self.z.set_uniform_intensity()
[docs]class MakeTestData(object):
"""MakeTestData is an object containing a generated test signal.
The default signal consists of a Disk and concentric Ring, with the
Ring being less intensive than the centre Disk. Unlimited number of
Rings and Disk can be added separately.
Parameters
----------
size_x, size_y : float, int
The range of the x and y axis goes from 0 to size_x, size_y
scale : float, int
The step size of the x and y axis
default : bool, default False
If True, the default object should be generated. If false, Ring and
Disk must be added separately by self.add_ring(), self.add_disk()
blur : bool, default True
If True, do a Gaussian blur of the disk.
blur_sigma : int, default 1
Sigma of the Gaussian blurring, if blur is True.
downscale : bool, default True
Note: currently using downscale and adding a disk, will lead to the
center of the disk being shifted. Ergo: the disk_x and disk_y will not
be correct when downscale is True.
Attributes
----------
signal : hyperspy.signals.Signal2D
Test signal
z_list : list
List containing Ring and Disk objects added to the signal
downscale_factor : int
The data is upscaled before Circle is added, and similarly
downscaled to return to given dimensions. This improves the
quality of Circle
Examples
--------
Default settings
>>> from pixstem.make_diffraction_test_data import MakeTestData
>>> test_data = MakeTestData()
>>> test_data.signal.plot()
More control
>>> test_data = MakeTestData(default=False)
>>> test_data.add_disk(x0=50, y0=50, r=10, intensity=30)
>>> test_data.add_ring(x0=45, y0=52, r=25, intensity=10)
>>> test_data.signal.plot()
"""
def __init__(
self, size_x=100, size_y=100, scale=1,
default=False, blur=True, blur_sigma=1, downscale=True):
self.scale = scale
self.blur_on = blur
self.blur_sigma = blur_sigma
self.downscale_on = downscale
self.downscale_factor = 5
if not self.downscale_on:
self.downscale_factor = 1
self.size_x, self.size_y = size_x, size_y
self.generate_mesh()
self.z_list = []
if default:
self.add_disk()
self.add_ring(lw_pix=1)
else:
self.update_signal()
def __repr__(self):
return '<%s, ((x, y): (%s, %s), s: %s, z: %s)>' % (
self.__class__.__name__,
self.size_x, self.size_y,
self.scale, len(self.z_list),
)
[docs] def update_signal(self):
self.make_signal()
self.downscale()
self.blur()
self.to_signal()
[docs] def generate_mesh(self):
self.X = np.arange(0, self.size_x, self.scale / self.downscale_factor)
self.Y = np.arange(0, self.size_y, self.scale / self.downscale_factor)
self.xx, self.yy = np.meshgrid(self.X, self.Y, sparse=True)
[docs] def add_disk(self, x0=50, y0=50, r=5, intensity=10):
scale = self.scale / self.downscale_factor
self.z_list.append(Disk(self.xx, self.yy, scale, x0, y0, r, intensity))
self.update_signal()
[docs] def add_disk_ellipse(
self, x0=50, y0=50, semi_len0=5, semi_len1=8, rotation=0.78,
intensity=10):
ellipse = EllipseDisk(
xx=self.xx, yy=self.yy, x0=x0, y0=y0, semi_len0=semi_len0,
semi_len1=semi_len1, rotation=rotation, intensity=intensity)
self.z_list.append(ellipse)
self.update_signal()
[docs] def add_ring(self, x0=50, y0=50, r=20, intensity=10, lw_pix=0):
"""
Add a ring to the test data.
Parameters
----------
x0, y0 : number, default 50
Centre position of the ring
r : number, default 20
Radius of the ring, defined as the distance from the centre to the
middle of the line of the ring, which will be most intense after
blurring.
intensity : number, default 10
Pixel value of the ring. Note, this value will be lowered
if blur or downscale is True
lw_pix : number, default 0
Distance in pixels from radius to the outer and inner edge of the
ring. Inner radius: r-lw, outer radius: r+lw. In total this gives
a ring line width in pixels of 2*lw+1.
"""
scale = self.scale / self.downscale_factor
lr = lw_pix * self.scale # scalar
self.z_list.append(
Ring(self.xx, self.yy, scale, x0, y0, r, intensity, lr))
self.update_signal()
[docs] def add_ring_ellipse(
self, x0=50, y0=50, semi_len0=11, semi_len1=13, rotation=0.78,
intensity=10, lw_r=2):
ellipse = EllipseRing(
xx=self.xx, yy=self.yy, x0=x0, y0=y0, semi_len0=semi_len0,
semi_len1=semi_len1, rotation=rotation, intensity=intensity,
lw_r=lw_r)
self.z_list.append(ellipse)
self.update_signal()
[docs] def make_signal(self):
if len(self.z_list) == 0:
self.z = self.xx * 0 + self.yy * 0
elif len(self.z_list) == 1:
self.z = self.z_list[0].get_signal()
elif len(self.z_list) > 1:
z_temp = self.z_list[0].get_signal()
for i in self.z_list[1:]:
z_temp = np.add(z_temp, i.get_signal())
self.z = z_temp
[docs] def downscale(self):
if self.downscale_on:
shape = (
int(self.z.shape[0] / self.downscale_factor),
int(self.z.shape[1] / self.downscale_factor))
sh = (
shape[0],
self.z.shape[0] // shape[0], shape[1],
self.z.shape[1] // shape[1])
self.z_downscaled = self.z.reshape(sh).mean(-1).mean(1)
else:
self.z_downscaled = self.z
[docs] def blur(self):
if self.blur_on:
self.z_blurred = gaussian_filter(
self.z_downscaled, sigma=self.blur_sigma)
else:
self.z_blurred = self.z_downscaled
[docs] def to_signal(self):
self.signal = PixelatedSTEM(self.z_blurred)
self.signal.axes_manager[0].scale = self.scale
self.signal.axes_manager[1].scale = self.scale
[docs] def set_downscale_factor(self, factor):
self.downscale_factor = factor
self.generate_mesh()
for i in self.z_list:
i.update_axis(self.xx, self.yy)
self.update_signal()
[docs] def set_signal_zero(self):
self.z_list = []
self.update_signal()
[docs]def generate_4d_data(
probe_size_x=10, probe_size_y=10, image_size_x=50, image_size_y=50,
disk_x=25, disk_y=25, disk_r=5, disk_I=20,
ring_x=25, ring_y=25, ring_r=20, ring_I=6, ring_lw=0,
ring_e_x=None, ring_e_y=25, ring_e_semi_len0=15, ring_e_semi_len1=15,
ring_e_r=0, ring_e_I=6, ring_e_lw=1,
blur=True, blur_sigma=1, downscale=True, add_noise=False,
noise_amplitude=1, lazy=False, lazy_chunks=None,
show_progressbar=True):
"""Generate a test dataset containing a disk and diffraction ring.
Useful for checking that radial average algorithms are working
properly.
The centre, intensity and radius position of the ring and disk can vary
as a function of probe position, through the disk_x, disk_y, disk_r,
disk_I, ring_x, ring_y, ring_r and ring_I arguments.
In addition, the line width of the ring can be varied with ring_lw.
There is also an elliptical ring, which can be added separately
to the circular ring. This elliptical ring uses the ring_e_*
arguments. It is disabled by default.
The ring can be deactivated by setting ring_x=None.
The disk can be deactivated by setting disk_x=None.
The elliptical ring can be deactivated by setting ring_e_x=None.
Parameters
----------
probe_size_x, probe_size_y : int, default 10
Size of the navigation dimension.
image_size_x, image_size_y : int, default 50
Size of the signal dimension.
disk_x, disk_y : int or NumPy 2D-array, default 20
Centre position of the disk. Either integer or NumPy 2-D array.
See examples on how to make them the correct size.
To deactivate the disk, set disk_x=None.
disk_r : int or NumPy 2D-array, default 5
Radius of the disk. Either integer or NumPy 2-D array.
See examples on how to make it the correct size.
disk_I : int or NumPy 2D-array, default 20
Intensity of the disk, for each of the pixels.
So if I=30, the each pixel in the disk will have a value of 30.
Note, this value will change if blur=True or downscale=True.
ring_x, ring_y : int or NumPy 2D-array, default 20
Centre position of the ring. Either integer or NumPy 2-D array.
See examples on how to make them the correct size.
To deactivate the ring, set ring_x=None.
ring_r : int or NumPy 2D-array, default 20
Radius of the ring. Either integer or NumPy 2-D array.
See examples on how to make it the correct size.
ring_I : int or NumPy 2D-array, default 6
Intensity of the ring, for each of the pixels.
So if I=5, each pixel in the ring will have a value of 5.
Note, this value will change if blur=True or downscale=True.
ring_lw : int or NumPy 2D-array, default 0
Line width of the ring. If ring_lw=1, the line will be 3 pixels wide.
If ring_lw=2, the line will be 5 pixels wide.
ring_e_x, ring_e_y : int or NumPy 2D-array, default 20
Centre position of the elliptical ring. Either integer or
NumPy 2-D array. See examples on how to make them the correct size.
To deactivate the ring, set ring_x=None (which is the default).
ring_e_semi_len0, ring_e_semi_len1 : int or NumPy 2D-array, default 15
Semi lengths of the elliptical ring. Either integer or NumPy 2-D
arrays. See examples on how to make it the correct size.
ring_e_I : int or NumPy 2D-array, default 6
Intensity of the elliptical ring, for each of the pixels.
So if I=5, each pixel in the ring will have a value of 5.
Note, this value will change if blur=True or downscale=True.
ring_e_r : int or NumPy 2D-array, default 0
Rotation of the elliptical ring, in radians.
ring_e_lw : int or NumPy 2D-array, default 0
Line width of the ring. If ring_lw=1, the line will be 3 pixels wide.
If ring_lw=2, the line will be 5 pixels wide.
blur : bool, default True
If True, do a Gaussian blur of the disk.
blur_sigma : int, default 1
Sigma of the Gaussian blurring, if blur is True.
downscale : bool, default True
If True, use upscaling (then downscaling) to anti-alise the disk.
add_noise : bool, default False
Add Gaussian random noise.
noise_amplitude : float, default 1
The amplitude of the noise, if add_noise is True.
lazy : bool, default False
If True, the signal will be lazy
lazy_chunks : tuple, optional
Used if lazy is True, default (10, 10, 10, 10).
Returns
-------
signal : HyperSpy Signal2D
Signal with 2 navigation dimensions and 2 signal dimensions.
Examples
--------
>>> import pixstem.make_diffraction_test_data as mdtd
>>> s = mdtd.generate_4d_data(show_progressbar=False)
>>> s.plot()
Using more arguments
>>> s = mdtd.generate_4d_data(probe_size_x=20, probe_size_y=30,
... image_size_x=50, image_size_y=90,
... disk_x=30, disk_y=70, disk_r=9, disk_I=30,
... ring_x=35, ring_y=65, ring_r=20, ring_I=10,
... blur=False, downscale=False, show_progressbar=False)
Adding some Gaussian random noise
>>> s = mdtd.generate_4d_data(add_noise=True, noise_amplitude=3,
... show_progressbar=False)
Different centre positions for each probe position.
Note the size=(20, 10), and probe_x=10, probe_y=20: size=(y, x).
>>> import numpy as np
>>> disk_x = np.random.randint(5, 35, size=(20, 10))
>>> disk_y = np.random.randint(5, 45, size=(20, 10))
>>> disk_I = np.random.randint(50, 100, size=(20, 10))
>>> ring_x = np.random.randint(5, 35, size=(20, 10))
>>> ring_y = np.random.randint(5, 45, size=(20, 10))
>>> ring_r = np.random.randint(10, 15, size=(20, 10))
>>> ring_I = np.random.randint(1, 30, size=(20, 10))
>>> ring_lw = np.random.randint(1, 5, size=(20, 10))
>>> s = mdtd.generate_4d_data(probe_size_x=10, probe_size_y=20,
... image_size_x=40, image_size_y=50, disk_x=disk_x, disk_y=disk_y,
... disk_I=disk_I, ring_x=ring_x, ring_y=ring_y, ring_r=ring_r,
... ring_I=ring_I, ring_lw=ring_lw, show_progressbar=False)
Do not plot the disk
>>> s = mdtd.generate_4d_data(disk_x=None, show_progressbar=False)
Do not plot the ring
>>> s = mdtd.generate_4d_data(ring_x=None, show_progressbar=False)
Plot only an elliptical ring
>>> from numpy.random import randint, random
>>> s = mdtd.generate_4d_data(
... probe_size_x=10, probe_size_y=10,
... disk_x=None, ring_x=None,
... ring_e_x=randint(20, 30, (10, 10)),
... ring_e_y=randint(20, 30, (10, 10)),
... ring_e_semi_len0=randint(10, 20, (10, 10)),
... ring_e_semi_len1=randint(10, 20, (10, 10)),
... ring_e_r=random((10, 10))*np.pi,
... ring_e_lw=randint(1, 3, (10, 10)))
"""
if disk_x is None:
plot_disk = False
else:
plot_disk = True
if not isiterable(disk_x):
disk_x = np.ones((probe_size_y, probe_size_x)) * disk_x
if not isiterable(disk_y):
disk_y = np.ones((probe_size_y, probe_size_x)) * disk_y
if not isiterable(disk_r):
disk_r = np.ones((probe_size_y, probe_size_x)) * disk_r
if not isiterable(disk_I):
disk_I = np.ones((probe_size_y, probe_size_x)) * disk_I
if ring_x is None:
plot_ring = False
else:
plot_ring = True
if not isiterable(ring_x):
ring_x = np.ones((probe_size_y, probe_size_x)) * ring_x
if not isiterable(ring_y):
ring_y = np.ones((probe_size_y, probe_size_x)) * ring_y
if not isiterable(ring_r):
ring_r = np.ones((probe_size_y, probe_size_x)) * ring_r
if not isiterable(ring_I):
ring_I = np.ones((probe_size_y, probe_size_x)) * ring_I
if not isiterable(ring_lw):
ring_lw = np.ones((probe_size_y, probe_size_x)) * ring_lw
if ring_e_x is None:
plot_ring_e = False
else:
plot_ring_e = True
if not isiterable(ring_e_x):
ring_e_x = np.ones((probe_size_y, probe_size_x)) * ring_e_x
if not isiterable(ring_e_y):
ring_e_y = np.ones((probe_size_y, probe_size_x)) * ring_e_y
if not isiterable(ring_e_semi_len0):
ring_e_semi_len0 = np.ones(
(probe_size_y, probe_size_x)) * ring_e_semi_len0
if not isiterable(ring_e_semi_len1):
ring_e_semi_len1 = np.ones(
(probe_size_y, probe_size_x)) * ring_e_semi_len1
if not isiterable(ring_e_I):
ring_e_I = np.ones((probe_size_y, probe_size_x)) * ring_e_I
if not isiterable(ring_e_lw):
ring_e_lw = np.ones((probe_size_y, probe_size_x)) * ring_e_lw
if not isiterable(ring_e_r):
ring_e_r = np.ones((probe_size_y, probe_size_x)) * ring_e_r
signal_shape = (probe_size_y, probe_size_x, image_size_y, image_size_x)
s = PixelatedSTEM(np.zeros(shape=signal_shape))
for i in tqdm(s, desc='Make test data', disable=not show_progressbar):
index = s.axes_manager.indices[::-1]
test_data = MakeTestData(
size_x=image_size_x, size_y=image_size_y,
default=False, blur=blur, blur_sigma=blur_sigma,
downscale=downscale)
if plot_disk:
dx, dy, dr = disk_x[index], disk_y[index], disk_r[index]
dI = disk_I[index]
test_data.add_disk(dx, dy, dr, intensity=dI)
if plot_ring:
rx, ry, rr = ring_x[index], ring_y[index], ring_r[index]
rI, rLW = ring_I[index], ring_lw[index]
test_data.add_ring(rx, ry, rr, intensity=rI, lw_pix=rLW)
if plot_ring_e:
rex, rey = ring_e_x[index], ring_e_y[index]
resl0, resl1 = ring_e_semi_len0[index], ring_e_semi_len1[index]
reI, reLW, rer = ring_e_I[index], ring_e_lw[index], ring_e_r[index]
test_data.add_ring_ellipse(
x0=rex, y0=rey, semi_len0=resl0, semi_len1=resl1,
rotation=rer, intensity=reI, lw_r=reLW)
s.data[index][:] = test_data.signal.data[:]
if add_noise:
s.data[index][:] += np.random.random(
size=(image_size_y, image_size_x)) * noise_amplitude
s.axes_manager.indices = [0] * s.axes_manager.navigation_dimension
if lazy:
if lazy_chunks is None:
lazy_chunks = 10, 10, 10, 10
data_lazy = da.from_array(s.data, lazy_chunks)
s = LazyPixelatedSTEM(data_lazy)
return s
def _make_4d_peak_array_test_data(xf, yf, semi0, semi1, rot, nt=20):
"""Get a 4D NumPy array with peak_array test data.
Parameters
----------
xf, yf : scalar, 2D NumPy array
Centre position of the ellipse. The size of the xf array gives the
size of the peak_array.
semi0, semi1 : scalar, 2D NumPy array
Semi length of the ellipse, of rot is 0 semi0 is the x-direction
semi length, and semi1 the y-direction.
rot : scalar, 2D NumPy array
Rotation in radians.
nt : scalar
Number of points in the ellipse.
Returns
-------
peak_array : NumPy 4D array
Examples
--------
>>> import pixstem.api as ps
>>> import pixstem.make_diffraction_test_data as mdtd
>>> xf = np.random.randint(65, 70, size=(4, 5))
>>> yf = np.random.randint(115, 120, size=(4, 5))
>>> semi0 = np.random.randint(35, 40, size=(4, 5))
>>> semi1 = np.random.randint(45, 50, size=(4, 5))
>>> rot = np.random.random(size=(4, 5)) * 0.2
>>> peak_array = mdtd._make_4d_peak_array_test_data(
... xf, yf, semi0, semi1, rot)
>>> s = ps.PixelatedSTEM(np.zeros(shape=(4, 5, 200, 210)))
>>> import pixstem.marker_tools as mt
>>> mt.add_peak_array_to_signal_as_markers(s, peak_array)
"""
peak_array = np.empty_like(xf, dtype=np.object)
for iy, ix in np.ndindex(peak_array.shape):
params = (xf[iy, ix], yf[iy, ix],
semi0[iy, ix], semi1[iy, ix], rot[iy, ix], nt)
ellipse_points = ret.make_ellipse_data_points(*params)
peak_array[iy, ix] = np.fliplr(ellipse_points)
return peak_array
[docs]class DiffractionTestImage(object):
def __init__(self, disk_r=7, blur=2, image_x=256, image_y=256, rotation=0,
diff_intensity_reduction=1.0, intensity_noise=0.5):
"""Make an artificial diffraction image, similar to NBED data.
This class is for creating images which are similar to
nanobeam electron diffraction patterns, with functionality for
adding diffraction disks, Lorentzian background and intensity noise.
Can be combined with the DiffractionTestDataset class, for creating
4-dimensional datasets. Similar to the ones acquired with
fast pixelated direct electron detectors in
scanning transmission electron microscopy.
Parameters
----------
disk_r : int
Radius for all of the disks added. Default 7
blur : scalar
Amount of Gaussian blur applied to the image.
Use False to disable blurring. Default 2.
image_x, image_y : int
Dimensions for the image. Default 256 for both.
rotation : int
Rotation of the image in relation to the centre point.
For an image with the size (256, 256), the rotation will
be around (128, 128).
diff_intensity_reduction : scalar
In diffraction patterns, the intensity of the disk are
reduced with increasing scattering angle. This parameter
emulates this, with decreasing the intensity of the disks
as a function of distance from the centre point. Default 1.0.
Set False to disable.
intensity_noise : scalar
The width of the Gaussian intensity noise added to the image.
Set False to disable. Default 0.5.
Examples
--------
>>> import pixstem.make_diffraction_test_data as mdtd
>>> di = mdtd.DiffractionTestImage()
>>> di.add_disk(x=128, y=128, intensity=10.)
>>> di.add_cubic_disks(vx=20, vy=20, intensity=2., n=5)
>>> di.add_background_lorentz()
>>> s = di.get_signal()
>>> s.plot()
Get a slightly rotated version of the diffraction image
>>> di.rotation = 10
>>> s = di.get_signal()
>>> s.plot()
"""
self.disk_r = disk_r
self.blur = blur
self.image_x = image_x
self.image_y = image_y
self.rotation = rotation
self.diff_intensity_reduction = diff_intensity_reduction
self.intensity_noise = intensity_noise
self._background_lorentz_width = False
self._background_lorentz_intensity = None
self._x_list = []
self._y_list = []
self._intensity_list = []
def __repr__(self):
return '<%s, disks:%s, r:%s, rot:%s, im:(%s,%s)>' % (
self.__class__.__name__,
len(self._x_list),
self.disk_r,
self.rotation,
self.image_x,
self.image_y)
def __copy__(self):
d = DiffractionTestImage(
disk_r=self.disk_r, blur=self.blur,
image_x=self.image_x, image_y=self.image_y,
rotation=self.rotation,
diff_intensity_reduction=self.diff_intensity_reduction,
intensity_noise=self.intensity_noise)
d._background_lorentz_width = self._background_lorentz_width
d._background_lorentz_intensity = self._background_lorentz_intensity
d._x_list = self._x_list.copy()
d._y_list = self._y_list.copy()
d._intensity_list = self._intensity_list.copy()
return d
[docs] def copy(self):
return self.__copy__()
[docs] def add_disk(self, x, y, intensity=1):
if not isinstance(x, int):
raise ValueError("x needs to be integer, not {0}".format(x))
if not isinstance(y, int):
raise ValueError("y needs to be integer, not {0}".format(y))
self._x_list.append(x)
self._y_list.append(y)
self._intensity_list.append(intensity)
[docs] def add_background_lorentz(self, width=10, intensity=5):
self._background_lorentz_width = width
self._background_lorentz_intensity = intensity
def _get_diff_intensity_reduction(self, dr, i):
r_max = np.hypot(self.image_x/2, self.image_y/2)
if dr == 0.:
dr = 0.00000000001
x = np.pi * dr * 1.5 / (r_max * self.diff_intensity_reduction)
i_new = np.sin(x) / x * i
if i_new < 0:
i_new = 0
return i_new
[docs] def add_cubic_disks(self, vx, vy, intensity=1, n=1):
"""Add disks in a cubic pattern around the centre point of the image.
Parameters
----------
vx, vy : int
intensity : scalar
n : int
Number of orders of diffraction disks.
If n=1, 8 disks, if n=2, 24 disks.
"""
cx, cy = self.image_x/2, self.image_y/2
for px in range(-n, n + 1):
for py in range(-n, n + 1):
if not (px == 0 and py == 0):
x = int(round(vx * px + cx))
y = int(round(vy * py + cy))
self.add_disk(x, y, intensity=intensity)
def _get_background_lorentz(self):
width = self._background_lorentz_width
intensity = self._background_lorentz_intensity
x = np.linspace(-width, width, self.image_y)
y = np.linspace(-width, width, self.image_x)
YY, XX = np.meshgrid(y, x)
YY = YY.astype(np.float32)
XX = XX.astype(np.float32)
b = 1 / (np.pi * (1 + np.hypot(YY, XX)**2)) * intensity
return b
[docs] def get_diffraction_test_image(self, dtype=np.float32):
image_x, image_y = self.image_x, self.image_y
cx, cy = image_x/2, image_y/2
image = np.zeros((image_y, image_x), dtype=np.float32)
iterator = zip(self._x_list, self._y_list, self._intensity_list)
for x, y, i in iterator:
if self.diff_intensity_reduction is not False:
dr = np.hypot(x - cx, y - cy)
i = self._get_diff_intensity_reduction(dr, i)
image[y, x] = i
disk = morphology.disk(self.disk_r, dtype=dtype)
image = convolve2d(image, disk, mode='same')
if self.rotation != 0:
image = rotate(image, self.rotation, reshape=False)
if self.blur != 0:
image = gaussian_filter(image, self.blur)
if self._background_lorentz_width is not False:
image += self._get_background_lorentz()
if self.intensity_noise is not False:
noise = np.random.random((image_y, image_x)) * self.intensity_noise
image += noise
return image
[docs] def get_signal(self):
s = PixelatedSTEM(self.get_diffraction_test_image())
return s
[docs] def plot(self):
s = self.get_signal()
s.plot()
[docs]class DiffractionTestDataset(object):
def __init__(self, probe_x=10, probe_y=10, detector_x=256, detector_y=256,
noise=0.5, dtype=np.float32):
"""Make a 4-dimensional dataset similar to NBED.
This class is for creating datasets which are similar to
nanobeam electron diffraction patterns. It is used in combination
with one or several DiffractionTestImage objects.
Parameters
----------
probe_x, probe_y : int
detector_x, detector_y : int
noise : scalar
Examples
--------
>>> import pixstem.make_diffraction_test_data as mdtd
>>> di = mdtd.DiffractionTestImage(intensity_noise=False)
>>> di.add_disk(x=128, y=128, intensity=10.)
>>> di.add_cubic_disks(vx=20, vy=20, intensity=2., n=5)
>>> di.add_background_lorentz()
>>> di_rot = di.copy()
>>> di_rot.rotation = 10
>>> dtd = mdtd.DiffractionTestDataset(10, 10, 256, 256)
>>> position_array = np.ones((10, 10), dtype=np.bool)
>>> position_array[:5] = False
>>> dtd.add_diffraction_image(di, position_array)
>>> dtd.add_diffraction_image(di_rot, np.invert(position_array))
>>> s = dtd.get_signal()
"""
self.data = np.zeros(
(probe_x, probe_y, detector_x, detector_y), dtype=dtype)
self.probe_x = probe_x
self.probe_y = probe_y
self.detector_x = detector_x
self.detector_y = detector_y
self.noise = noise
def __repr__(self):
return '<%s, (%s)>' % (
self.__class__.__name__,
self.data.shape)
[docs] def add_diffraction_image(
self, diffraction_test_image, position_array=None):
"""Add a diffraction image to all or a subset of the dataset.
See the class docstring for example on how to use this.
Parameters
----------
diffraction_test_image : PixStem DiffractionTestData object
position_array : NumPy array
Boolean array, specifying which positions in the dataset
the diffraction_test_image should be added to. Must have two
dimensions, and the same shape as (probe_x, probe_y).
"""
probe_x, probe_y = self.probe_x, self.probe_y
detector_x, detector_y = self.detector_x, self.detector_y
image = diffraction_test_image.get_diffraction_test_image()
if position_array is None:
position_array = np.ones((probe_x, probe_y), dtype=np.bool)
for ix, iy in np.ndindex(probe_x, probe_y):
if position_array[ix, iy]:
self.data[ix, iy, :, :] = image
if self.noise is not False:
image_noise = np.random.random((detector_x, detector_y))
self.data[ix, iy, :, :] += (image_noise * self.noise)
[docs] def get_signal(self):
s = PixelatedSTEM(self.data)
return s
[docs] def plot(self):
s = self.get_signal()
s.plot()