mirror of
https://gitlab.ethz.ch/nux/spring.git
synced 2026-09-16 15:19:56 +02:00
203 lines
7.0 KiB
Python
203 lines
7.0 KiB
Python
import numpy as np
|
|
from spring_core.pattern import downscale
|
|
import matplotlib.pyplot as plt
|
|
import matplotlib
|
|
|
|
def resize_pattern(data, mask, newsize):
|
|
|
|
tempdata = np.copy(data)
|
|
tempdata = tempdata+1e3
|
|
tempdata[mask>0.5]=-1
|
|
|
|
newdata = downscale(tempdata, [newsize, newsize])
|
|
|
|
newmask = np.zeros(newdata.shape)
|
|
newmask[newdata<0]=1
|
|
|
|
newdata = newdata-1e3
|
|
|
|
return newdata, newmask
|
|
|
|
|
|
def cut(data, center, size):
|
|
|
|
frame = int(data.shape[0]*2)
|
|
tempsize = int(data.shape[0]+2*frame)
|
|
tempdata = np.zeros([tempsize, tempsize])
|
|
tempdata[frame:data.shape[0]+frame, frame:data.shape[1]+frame] = data
|
|
tempcenter = [center[0]+frame, center[1]+frame]
|
|
newdata = np.copy(tempdata[tempcenter[0]-size//2:tempcenter[0]+size//2, tempcenter[1]-size//2:tempcenter[1]+size//2])
|
|
|
|
return newdata
|
|
|
|
|
|
class Pattern:
|
|
"""
|
|
Class for the handling of diffraction patterns.
|
|
|
|
:param pattern: 2D numpy array containing the diffraction pattern
|
|
:param mask: 2D numpy array containing the mask. Indexes where ``mask[i,j]=0`` indicate good pixels. Indexes where ``mask[i,j]=1`` are instead excluded. the *mask* and the *pattern* must have the same dimensions.
|
|
:param center: list of length 2 containing the y and x coordinates of the center of the *pattern*
|
|
:param cropsize: The pattern is cropped to a shape ``[cropsize,cropsize]``. If ``cropsize = None``, it is set equal to ``max(pattern.shape)``.
|
|
:param pid: List of identifiers for the diffraction pattern. For example, if pattern has been acquired in run 34 with pulse-id 197843921, it is convenient to set ``pid=[34,197843921]``.
|
|
:param size: If not ``None``, the pattern is rescaled to a final shape ``[size,size]`` after cropping.
|
|
:param satvalue: Indicate which is the saturation value of the detector. If different from ``None``, pixels above this value are masked.
|
|
:param metadata: Additional experimental metadata that is convenient to save with the reconstruction, provided as a dictionary.
|
|
|
|
"""
|
|
def __init__(self, pattern, mask, center: list[2], cropsize: int = None, pid: list =[0], size: int = None, satvalue: float = None, metadata={}):
|
|
|
|
self.origpattern = np.copy(pattern)
|
|
self.origmask = np.copy(mask)
|
|
self.pid = pid
|
|
self.metadata = metadata
|
|
data = np.copy(pattern)
|
|
pixmask = np.copy(mask)
|
|
|
|
|
|
pixmask[np.isnan(data)] = 1
|
|
|
|
######
|
|
if cropsize is None:
|
|
cropsize = np.amax(data.shape)
|
|
if size is None:
|
|
size = cropsize
|
|
|
|
data = cut(data, center, cropsize)
|
|
pixmask = cut(pixmask, center, cropsize)
|
|
|
|
|
|
satmask = np.zeros(pixmask.shape)
|
|
if not (satvalue is None):
|
|
satmask[data>satvalue] = 1
|
|
pixmask[satmask==1] = 1
|
|
|
|
|
|
|
|
data[pixmask==1] = -1
|
|
data[satmask==1] = -1
|
|
|
|
resizedmask = np.copy(pixmask)
|
|
if size<cropsize:
|
|
data, resizedmask=resize_pattern(data, pixmask, size)
|
|
satmask=downscale(satmask, [size, size])
|
|
pixmask=downscale(pixmask, [size, size])
|
|
|
|
|
|
pixmask[pixmask>0.25]=1
|
|
pixmask[pixmask<=0.25]=0
|
|
pixmask[np.logical_and(pixmask<0.5, resizedmask>0.25)]=1
|
|
|
|
satmask[satmask>0.05]=1
|
|
satmask[satmask<=0.05]=0
|
|
|
|
data[data<0]=0
|
|
data[pixmask==1]=-1
|
|
data[satmask==1]=-2
|
|
|
|
self.pattern = data
|
|
|
|
|
|
def get(self):
|
|
"""
|
|
Get the prepared pattern as a 2D numpy array.
|
|
"""
|
|
return self.pattern
|
|
|
|
def get_pid(self):
|
|
return self.pid
|
|
|
|
def get_metadata(self):
|
|
return self.metadata
|
|
|
|
|
|
def plot(self, ax: matplotlib.axes.Axes = None, cmap='inferno', circles: int = 0, cropsize: int = -1 ):
|
|
"""
|
|
Plot the diffraction pattern, in logarithmic colorscale.
|
|
|
|
:param ax: matplotlib axes where the pattern is plotted. If none, axes are created within the function and the data is automatically plotted.
|
|
:param cmap: colormap for the plot. It must be a valid `matplotlib` colormap
|
|
:param circles: number of circles to draw on the pattern, centered in the central coordinates of the pattern matrix. It serves as a guide for the eyes to identify the correct pattern center.
|
|
:param cropsize: The pattern is cropped to the given linear dimension for plotting.
|
|
"""
|
|
|
|
if ax is None:
|
|
fig, axx = plt.subplots(1,1,figsize=(4,4))
|
|
|
|
else:
|
|
axx = ax
|
|
axx.clear()
|
|
patternplot = np.copy(self.pattern)
|
|
mask = patternplot<0
|
|
patternplot[mask]=np.nan
|
|
vmin = np.nanmean(patternplot[:patternplot.shape[0]//8,::patternplot.shape[0]//8])*0.5
|
|
patternplot[mask]=vmin
|
|
patternplot[patternplot<vmin]=vmin
|
|
|
|
|
|
if cropsize>0 and cropsize<patternplot.shape[0]:
|
|
patternplot = cut(patternplot, [patternplot.shape[0]//2,patternplot.shape[1]//2] , cropsize)
|
|
|
|
axx.imshow(patternplot, norm=matplotlib.colors.LogNorm(vmin=vmin), cmap=cmap)
|
|
|
|
|
|
if circles>0:
|
|
maxrad = patternplot.shape[0]/2
|
|
radii = np.linspace(0, maxrad, num=circles+2)[1:-1]
|
|
for rad in radii:
|
|
circ = matplotlib.patches.Circle((maxrad,maxrad),rad, ec='lime', fill=False)
|
|
axx.add_patch(circ)
|
|
|
|
|
|
if ax is None:
|
|
titlestring=''
|
|
if len(self.pid)>0:
|
|
titlestring = "{:d}".format(self.pid[0])
|
|
for i in range(1, len(self.pid)):
|
|
titlestring+= " - {:d}".format(self.pid[i])
|
|
|
|
axx.set_title(titlestring)
|
|
|
|
|
|
|
|
|
|
|
|
plt.show()
|
|
|
|
|
|
def plot_mask(self, ax : matplotlib.axes.Axes = None, cmap='viridis'):
|
|
"""
|
|
Plot the mask of the pattern. Masked pixels get a value of 1, while saturated pixels get a value of 2. Valid pixels have a value of 0.
|
|
|
|
:param ax: matplotlib axes where the mask is plotted. If none, axes are created within the function and the data is automatically plotted.
|
|
:param cmap: colormap for the plot. It must be a valid `matplotlib` colormap
|
|
"""
|
|
|
|
|
|
if ax is None:
|
|
fig, axx = plt.subplots(1,1,figsize=(4,4))
|
|
|
|
else:
|
|
axx = ax
|
|
axx.clear()
|
|
maskplot = np.zeros(self.pattern.shape)
|
|
maskplot[self.pattern==-1]=1
|
|
maskplot[self.pattern==-2]=2
|
|
axx.imshow(maskplot, cmap=cmap)
|
|
if ax is None:
|
|
titlestring=''
|
|
if len(self.pid)>0:
|
|
titlestring = "{:d}".format(self.pid[0])
|
|
for i in range(1, len(self.pid)):
|
|
titlestring+= " - {:d}".format(self.pid[i])
|
|
|
|
axx.set_title(titlestring)
|
|
|
|
plt.show()
|
|
|
|
|
|
|
|
|
|
|
|
|