cSAXS/cdtools
0
0
Fork 0
mirror of https://github.com/cdtools-developers/cdtools.git synced 2026-10-13 11:30:25 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'master' of github.mit.edu:Scattering/CDTools

Browse Source
This commit is contained in:
Abe Levitan
2020-01-25 19:01:11 -05:00
parent a1cca0dcea 7a243842b9
commit 8d589eef86
14 changed files with 525 additions and 183 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CDTools/datasets/ptycho_2d_dataset.py
+24 -23
View File
@@ -84,7 +84,7 @@ class Ptycho2DDataset(CDataset):
getting data as GPU tensors.
It loads data in the format (inputs, output)
The inputs for a 2D ptychogaphy data set are:
1) The indices of the patterns to use
@@ -146,7 +146,7 @@ class Ptycho2DDataset(CDataset):
dataset = CDataset.from_cxi(cxi_file)
# Mutate the class to this subclass (BasicPtychoDataset)
dataset.__class__ = cls
# Load the data that is only relevant for this class
patterns, axes = cdtdata.get_data(cxi_file)
translations = cdtdata.get_ptycho_translations(cxi_file)
@@ -159,7 +159,7 @@ class Ptycho2DDataset(CDataset):
dataset.mask = t.ones(dataset.patterns.shape[-2:]).to(dtype=t.bool)
return dataset
def to_cxi(self, cxi_file):
"""Saves out a Ptycho2DDataset as a .cxi file
@@ -197,18 +197,18 @@ class Ptycho2DDataset(CDataset):
can display a base-10 log plot of the detector readout at each
position.
"""
# We start by making the figure and axes
fig, axes = plt.subplots(1,2,figsize=(8,5.3))
fig.tight_layout(rect=[0.04, 0.09, 0.98, 0.96])
axslider = plt.axes([0.15,0.06,0.75,0.03])
#
# Then we define some helper functions for getting the right data
# that are used both in the initial setup and the updates
#
def get_data(idx):
inputs, output = self[idx]
meas_data = output.detach().cpu().numpy()
@@ -216,16 +216,16 @@ class Ptycho2DDataset(CDataset):
mask = self.mask.detach().cpu().numpy()
else:
mask = 1
return mask, meas_data
def calculate_sizes(idx):
bbox = axes[0].get_window_extent().transformed(fig.dpi_scale_trans.inverted())
s0 = bbox.width * bbox.height / translations.shape[0] * 72**2 #72 is points per inch
s0 /= 4 # A rough value to make the size work out
s = np.ones(len(self)) * s0
s[idx] *= 4
return s
@@ -237,11 +237,13 @@ class Ptycho2DDataset(CDataset):
if hasattr(im, 'norecurse') and im.norecurse:
im.norecurse=False
return
im.norecurse=True
# This is needed to update the colorbar
im.set_clim(vmin=np.min(im.get_array()),
vmax=np.max(im.get_array()))
# only change limits if array contains multiple values
if np.min(im.get_array()) != np.max(im.get_array()):
im.set_clim(vmin=np.min(im.get_array()),
vmax=np.max(im.get_array()))
#
# The meatiest part of this program, here we just go through and
@@ -250,17 +252,17 @@ class Ptycho2DDataset(CDataset):
# First we set up the left-hand plot, which shows an overview map
axes[0].set_title('Relative Displacement Map')
translations = self.translations.detach().cpu().numpy()
nanomap_values = (self.mask.to(t.float32) * self.patterns).sum(dim=(1,2)).detach().cpu().numpy()
s = calculate_sizes(0)
nanomap = axes[0].scatter(1e6 * translations[:,0],1e6 * translations[:,1],s=s,c=nanomap_values, picker=True)
axes[0].invert_xaxis()
axes[0].set_facecolor('k')
axes[0].set_xlabel('Translation x (um)', labelpad=1)
axes[0].set_ylabel('Translation y (um)', labelpad=1)
axes[0].set_xlabel('Translation x ($\mu$m)', labelpad=1)
axes[0].set_ylabel('Translation y ($\mu$m)', labelpad=1)
cb1 = plt.colorbar(nanomap, ax=axes[0], orientation='horizontal',
format='%.2e',
ticks=ticker.LinearLocator(numticks=5),
@@ -277,7 +279,7 @@ class Ptycho2DDataset(CDataset):
meas = axes[1].imshow(np.log(meas_data) / np.log(10) * mask)
else:
meas = axes[1].imshow(meas_data * mask)
cb2 = plt.colorbar(meas, ax=axes[1], orientation='horizontal',
format='%.2e',
ticks=ticker.LinearLocator(numticks=5),
@@ -311,12 +313,12 @@ class Ptycho2DDataset(CDataset):
update_colorbar(meas)
#
# Now we define the functions to handle various kinds of events
# that can be thrown our way
#
# We start by creating the slider here, so it can be used
# by the update hooks.
slider = Slider(axslider, 'Pattern #', 0, len(self)-1, valstep=1, valfmt="%d")
@@ -330,9 +332,9 @@ class Ptycho2DDataset(CDataset):
if not hasattr(event, 'key'):
event.key = None
if event.key == 'up' or event.button == 'up':
if event.key == 'up' or event.button == 'up' or event.key == 'right':
idx = slider.val - 1
elif event.key == 'down' or event.button == 'down':
elif event.key == 'down' or event.button == 'down' or event.key == 'left':
idx = slider.val + 1
# Handle the wraparound and trigger the update
@@ -357,4 +359,3 @@ class Ptycho2DDataset(CDataset):
# (like the nanomap dot sizes) that otherwise would change on the
# first update
update(0)
CDTools/models/simple_ptycho.py
+10 -10
View File
@@ -13,8 +13,8 @@ import numpy as np
class SimplePtycho(CDIModel):
"""A simple ptychography model for exploring ideas and extensions
"""
def __init__(self, wavelength, detector_geometry,
@@ -39,7 +39,7 @@ class SimplePtycho(CDIModel):
self.detector_slice = detector_slice
self.surface_normal = t.Tensor(surface_normal)
if mask is None:
self.mask = None
else:
@@ -81,7 +81,7 @@ class SimplePtycho(CDIModel):
if hasattr(dataset, 'sample_info') and \
dataset.sample_info is not None and \
'orientation' in dataset.sample_info:
surface_normal = dataset.sample_info.orientation[2]
surface_normal = dataset.sample_info['orientation'][2]
else:
surface_normal = np.array([0.,0.,1.])
@@ -146,7 +146,7 @@ class SimplePtycho(CDIModel):
if self.mask is not None:
self.mask = self.mask.to(*args, **kwargs)
self.min_translation = self.min_translation.to(*args,**kwargs)
self.probe_basis = self.probe_basis.to(*args,**kwargs)
self.probe_norm = self.probe_norm.to(*args,**kwargs)
@@ -155,7 +155,7 @@ class SimplePtycho(CDIModel):
def sim_to_dataset(self, args_list):
# In the future, potentially add more control
# over what metadata is saved (names, etc.)
# First, I need to gather all the relevant data
# that needs to be added to the dataset
entry_info = {'program_name': 'CDTools',
@@ -168,15 +168,15 @@ class SimplePtycho(CDIModel):
ysurfacevec = np.cross(surface_normal, xsurfacevec)
ysurfacevec /= np.linalg.norm(ysurfacevec)
orientation = np.array([xsurfacevec, ysurfacevec, surface_normal])
sample_info = {'description': 'A simulated sample',
'orientation': orientation}
detector_geometry = self.detector_geometry
mask = self.mask
wavelength = self.wavelength
indices, translations = args_list
# Then we simulate the results
data = self.forward(indices, translations)
@@ -202,7 +202,7 @@ class SimplePtycho(CDIModel):
]
def save_results(self):
probe = tools.cmath.torch_to_complex(self.probe.detach().cpu())
probe = probe * self.probe_norm.detach().cpu().numpy()
CDTools/tools/analysis.py
+60 -60
View File
@@ -3,7 +3,7 @@
The functions in this module are designed to work either with pytorch tensors
or numpy arrays, so they can be used either directly after reconstructions
on the attributes of the models themselves, or after-the-fact once the
data has been stored in numpy arrays.
data has been stored in numpy arrays.
"""
from __future__ import division, print_function
@@ -21,17 +21,17 @@ __all__ = ['orthogonalize_probes','standardize', 'synthesize_reconstructions',
from matplotlib import pyplot as plt
def orthogonalize_probes(probes):
"""Orthogonalizes a set of incoherently mixing probes
The strategy is to define a reduced orthogonal basis that spans
all of the retrieved probes, and then build the density matrix
defined by the probes in that basis. After diagonalization, the
eigenvectors can be recast into the original basis and returned
Parameters
----------
probes : array
An l x n x m complex array representing a stack of probes
Returns
-------
ortho_probes: array
@@ -51,7 +51,7 @@ def orthogonalize_probes(probes):
for j, basis in enumerate(bases):
coefficients[j,i] = np.sum(basis.conj()*ortho_probe)
ortho_probe -= basis * coefficients[j,i]
coefficients[i,i] = np.sqrt(np.sum(np.abs(ortho_probe)**2))
bases.append(ortho_probe / coefficients[i,i])
@@ -68,12 +68,12 @@ def orthogonalize_probes(probes):
probe += basis * coefficient
ortho_probes.append(probe)
if send_to_torch:
return cmath.complex_to_torch(np.stack(ortho_probes[::-1]))
else:
return np.stack(ortho_probes[::-1])
def standardize(probe, obj, obj_slice=None, correct_ramp=False):
@@ -105,7 +105,7 @@ def standardize(probe, obj, obj_slice=None, correct_ramp=False):
probe : array
A complex array storing a retrieved probe or stack of incoherently mixed probes
obj : array
A complex array storing a retrieved probe
A complex array storing a retrieved object
obj_slice : slice
Optional, a slice to take from the object for calculating normalizations
correct_ramp : bool
@@ -136,7 +136,7 @@ def standardize(probe, obj, obj_slice=None, correct_ramp=False):
else:
single_probe = False
normalization = t.sqrt(t.sum(cmath.cabssq(probe[0])) / (len(probe[0].view(-1))/2))
probe = probe / normalization
obj = obj * normalization
@@ -146,13 +146,13 @@ def standardize(probe, obj, obj_slice=None, correct_ramp=False):
obj_slice = np.s_[(obj.shape[0]//8)*3:(obj.shape[0]//8)*5,
(obj.shape[1]//8)*3:(obj.shape[1]//8)*5]
if correct_ramp:
# Need to check if this is actually working and, if not, why not
center_freq = ip.centroid(cmath.cabssq(cmath.fftshift(t.fft(probe[0],2))))
center_freq -= (t.tensor(probe[0].shape[:-1]) // 2).to(t.float32)
center_freq /= t.tensor(probe[0].shape[:-1]).to(t.float32)
Is, Js = np.mgrid[:probe[0].shape[0],:probe[0].shape[1]]
probe_phase_ramp = cmath.expi(2 * np.pi *
(center_freq[0] * t.tensor(Is).to(t.float32) +
@@ -163,37 +163,37 @@ def standardize(probe, obj, obj_slice=None, correct_ramp=False):
(center_freq[0] * t.tensor(Is).to(t.float32) +
center_freq[1] * t.tensor(Js).to(t.float32)))
obj = cmath.cmult(obj, obj_phase_ramp)
# Then, we set them to consistent absolute phases
obj_angle = cmath.cphase(t.sum(obj[obj_slice],dim=(0,1)))
obj = cmath.cmult(obj, cmath.expi(-obj_angle))
for i in range(probe.shape[0]):
probe_angle = cmath.cphase(t.sum(probe[i],dim=(0,1)))
probe[i] = cmath.cmult(probe[i], cmath.expi(-probe_angle))
if single_probe:
probe = probe[0]
if probe_np:
probe = cmath.torch_to_complex(probe.detach().cpu())
if obj_np:
obj = cmath.torch_to_complex(obj.detach().cpu())
return probe, obj
def synthesize_reconstructions(probes, objects, use_probe=False, obj_slice=None, correct_ramp=False):
"""Takes a collection of reconstructions and outputs a single synthesized probe and object
The function first standardizes the sets of probes and objects using the
standardize function, passing through the relevant options. Then it
calculates the closest overlap of subsequent frames to subpixel
precision and uses a sinc interpolation to shift all the probes and objects
to a common frame. Then the images are summed.
Parameters
----------
probes : list(array)
@@ -216,7 +216,7 @@ def synthesize_reconstructions(probes, objects, use_probe=False, obj_slice=None,
obj_stack : list(array)
A list of standardized objects, for further processing
"""
probe_np = False
if isinstance(probes[0], np.ndarray):
probes = [cmath.complex_to_torch(probe).to(t.float32) for probe in probes]
@@ -228,15 +228,15 @@ def synthesize_reconstructions(probes, objects, use_probe=False, obj_slice=None,
obj_shape = np.min(np.array([obj.shape[:-1] for obj in objects]),axis=0)
objects = [obj[:obj_shape[0],:obj_shape[1]] for obj in objects]
if obj_slice is None:
obj_slice = np.s_[(objects[0].shape[0]//8)*3:(objects[0].shape[0]//8)*5,
(objects[0].shape[1]//8)*3:(objects[0].shape[1]//8)*5]
synth_probe, synth_obj = standardize(probes[0].clone(), objects[0].clone(), obj_slice=obj_slice,correct_ramp=correct_ramp)
obj_stack = [synth_obj]
for i, (probe, obj) in enumerate(zip(probes[1:],objects[1:])):
@@ -245,11 +245,11 @@ def synthesize_reconstructions(probes, objects, use_probe=False, obj_slice=None,
shift = ip.find_shift(synth_probe[0],probe[0], resolution=50)
else:
shift = ip.find_shift(synth_obj[obj_slice],obj[obj_slice], resolution=50)
obj = ip.sinc_subpixel_shift(obj,np.array(shift))
if len(probe.shape) == 4:
probe = t.stack([ip.sinc_subpixel_shift(p,tuple(shift))
for p in probe],dim=0)
@@ -260,7 +260,7 @@ def synthesize_reconstructions(probes, objects, use_probe=False, obj_slice=None,
synth_obj = synth_obj + obj
obj_stack.append(obj)
# If there only was one image
try:
i
@@ -279,13 +279,13 @@ def synthesize_reconstructions(probes, objects, use_probe=False, obj_slice=None,
def calc_consistency_prtf(synth_obj, objects, basis, obj_slice=None,nbins=None):
"""Calculates a PRTF between each the individual objects and a synthesized one
The consistency PRTF at any given spatial frequency is defined as the ratio
between the intensity of any given reconstruction and the intensity
of a synthesized or averaged reconstruction at that spatial frequency.
Typically, the PRTF is averaged over spatial frequencies with the same
magnitude.
Parameters
----------
synth_obj : array
@@ -316,51 +316,51 @@ def calc_consistency_prtf(synth_obj, objects, basis, obj_slice=None,nbins=None):
if isinstance(basis, t.Tensor):
basis = basis.detach().cpu().numpy()
if obj_slice is None:
obj_slice = np.s_[(objects[0].shape[0]//8)*3:(objects[0].shape[0]//8)*5,
(objects[0].shape[1]//8)*3:(objects[0].shape[1]//8)*5]
if nbins is None:
nbins = np.max(synth_obj[obj_slice].shape) // 4
synth_fft = cmath.cabssq(cmath.fftshift(t.fft(synth_obj[obj_slice],2))).numpy()
di = np.linalg.norm(basis[:,0])
di = np.linalg.norm(basis[:,0])
dj = np.linalg.norm(basis[:,1])
i_freqs = fftpack.fftshift(fftpack.fftfreq(synth_fft.shape[0],d=di))
j_freqs = fftpack.fftshift(fftpack.fftfreq(synth_fft.shape[1],d=dj))
Js,Is = np.meshgrid(j_freqs,i_freqs)
Rs = np.sqrt(Is**2+Js**2)
synth_ints, bins = np.histogram(Rs,bins=nbins,weights=synth_fft)
prtfs = []
for obj in objects:
obj = obj[obj_slice]
single_fft = cmath.cabssq(cmath.fftshift(t.fft(obj,2))).numpy()
single_fft = cmath.cabssq(cmath.fftshift(t.fft(obj,2))).numpy()
single_ints, bins = np.histogram(Rs,bins=nbins,weights=single_fft)
prtfs.append(synth_ints/single_ints)
prtf = np.mean(prtfs,axis=0)
if not obj_np:
bins = t.Tensor(bins)
prtf = t.Tensor(prtf)
return bins[:-1], prtf
def calc_deconvolved_cross_correlation(im1, im2, im_slice=None):
"""Calculates a cross-correlation between two images with their autocorrelations deconvolved.
This is formally defined as the inverse Fourier transform of the normalized
product of the Fourier transforms of the two images. It results in a
kernel, whose characteristic size is related to the exactness of the
@@ -379,7 +379,7 @@ def calc_deconvolved_cross_correlation(im1, im2, im_slice=None):
-------
corr : array
The complex-valued deconvolved cross-correlation, in real space
"""
im_np = False
@@ -389,7 +389,7 @@ def calc_deconvolved_cross_correlation(im1, im2, im_slice=None):
if isinstance(im2, np.ndarray):
im2 = cmath.complex_to_torch(im2)
im_np = True
# If last dimension is not 2, then convert to a complex tensor now
if im1.shape[-1] != 2:
im1 = t.stack((im1,t.zeros_like(im1)),dim=-1)
@@ -407,16 +407,16 @@ def calc_deconvolved_cross_correlation(im1, im2, im_slice=None):
# Not sure if this is more or less stable than just the correlation
# maximum - requires some testing
cor = t.ifft(cor_fft / cmath.cabs(cor_fft)[:,:,None],2)
if im_np:
cor = cmath.torch_to_complex(cor)
return cor
def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.):
"""Calculates a Fourier ring correlation between two images
This function requires an input of a basis to allow for FRC calculations
to be related to physical units.
@@ -446,7 +446,7 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.):
The FRC values
threshold : array
The threshold curve for comparison
"""
im_np = False
@@ -459,14 +459,14 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.):
if isinstance(basis, np.ndarray):
basis = t.tensor(basis)
# If last dimension is not 2, then convert to a complex tensor now
if im1.shape[-1] != 2:
im1 = t.stack((im1,t.zeros_like(im1)),dim=-1)
if im2.shape[-1] != 2:
im2 = t.stack((im2,t.zeros_like(im2)),dim=-1)
if im_slice is None:
im_slice = np.s_[(im1.shape[0]//8)*3:(im1.shape[0]//8)*5,
(im1.shape[1]//8)*3:(im1.shape[1]//8)*5]
@@ -474,23 +474,23 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.):
if nbins is None:
nbins = np.max(im1[im_slice].shape) // 4
cor_fft = cmath.cmult(cmath.fftshift(t.fft(im1[im_slice],2)),
cmath.fftshift(cmath.cconj(t.fft(im2[im_slice],2))))
F1 = cmath.cabs(cmath.fftshift(t.fft(im1[im_slice],2)))**2
F2 = cmath.cabs(cmath.fftshift(t.fft(im2[im_slice],2)))**2
di = np.linalg.norm(basis[:,0])
di = np.linalg.norm(basis[:,0])
dj = np.linalg.norm(basis[:,1])
i_freqs = fftpack.fftshift(fftpack.fftfreq(cor_fft.shape[0],d=di))
j_freqs = fftpack.fftshift(fftpack.fftfreq(cor_fft.shape[1],d=dj))
Js,Is = np.meshgrid(j_freqs,i_freqs)
Rs = np.sqrt(Is**2+Js**2)
numerator, bins = np.histogram(Rs,bins=nbins,weights=cmath.torch_to_complex(cor_fft))
@@ -502,13 +502,13 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.):
# This moves from combined-image SNR to single-image SNR
snr /= 2
threshold = (snr + (2 * snr + 1) / np.sqrt(n_pix)) / \
(1 + snr + (2 * np.sqrt(snr)) / np.sqrt(n_pix))
if not im_np:
bins = t.tensor(bins)
frc = t.tensor(frc)
threshold = t.tensor(threshold)
return bins[:-1], frc, threshold
CDTools/tools/data.py
+21 -21
View File
@@ -104,13 +104,13 @@ def get_sample_info(cxi_file):
metadata[attr] = str(s1[attr][()].decode())
except AttributeError as e:
metadata[attr] = str(np.array(s1[attr][:])[0].decode())
float_attrs = ['concentration',
'mass',
'temperature',
'thickness',
'unit_cell_volume']
for attr in float_attrs:
if attr in s1:
metadata[attr] = np.float32(s1[attr][()])
@@ -125,7 +125,7 @@ def get_sample_info(cxi_file):
yvec = orient[3:] / np.linalg.norm(orient[3:])
metadata['orientation'] = np.array([xvec,yvec,
np.cross(xvec,yvec)])
if 'geometry_1/surface_normal' in s1:
snorm = np.array(s1['geometry_1/surface_normal']).astype(np.float32)
xvec = np.cross(np.array([0.,1.,0.]), snorm)
@@ -133,7 +133,7 @@ def get_sample_info(cxi_file):
yvec = np.cross(snorm, xvec)
yvec /= np.linalg.norm(yvec)
metadata['orientation'] = np.array([xvec, yvec, snorm])
# Check if the metadata is empty
if metadata == {}:
metadata = None
@@ -234,18 +234,17 @@ def get_detector_geometry(cxi_file):
except:
corner_position = None
# Don't pretend to calculate corner position from distance if it's
# if it's not defined, but do calculate distance from corner position
# if distance is not defined. If neither is defined, then raise
# an error.
# Don't pretend to calculate corner position from distance if it's not
# defined, but do calculate distance from corner position if distance is
# not defined. If neither is defined, then raise an error.
if distance is None and corner_position is not None:
detector_normal = np.cross(basis_vectors[:,0],
basis_vectors[:,1])
detector_normal /= np.linalg.norm(detector_normal)
distance = np.linalg.norm(np.dot(corner_position, detector_normal))
if distance is None and corner_position is not None:
raise KeyError('Neither sample to detector distance or corner position is defined in file.')
if distance is None and corner_position is None:
raise KeyError('Neither sample to detector distance nor corner position is defined in file.')
return distance, basis_vectors, corner_position
@@ -260,7 +259,7 @@ def get_mask(cxi_file):
If any bit is set in the mask at all, it will be defined as a bad
pixel, with the exception of pixels marked exactly as 0x00001000,
which is defined to mean that the pixel has signal above the
background. These pixels are treated as on pixels
background. These pixels are treated as on pixels.
Parameters
----------
@@ -292,7 +291,7 @@ def get_dark(cxi_file):
if the dark image is a single image, it will return that image. If it
is a stack of images, it will return the mean along the stack axis.
If the darks do not exist, it will return None
If the darks do not exist, it will return None.
Parameters
----------
@@ -304,7 +303,7 @@ def get_dark(cxi_file):
dark : np.array
An array storing the dark image
"""
i1 = cxi_file['entry_1/instrument_1']
if 'detector_1/data_dark' in i1:
darks = np.array(i1['detector_1/data_dark'])
@@ -327,8 +326,7 @@ def get_data(cxi_file, cut_zeroes = True):
It will return the data array in whatever shape it's defined in.
It will also read out the axes attribute of the data into a list
of strings
It will also read out the axes attribute of the data into a list of strings.
Parameters
----------
@@ -342,7 +340,7 @@ def get_data(cxi_file, cut_zeroes = True):
axes : list(str)
A list of the axes defined in the axes attribute, if any
"""
# Possible locations for the data
if 'entry_1/data_1/data' in cxi_file:
pull_from = 'entry_1/data_1/data'
@@ -438,9 +436,9 @@ def add_entry_info(cxi_file, metadata):
elif isinstance(value, datetime.datetime):
cxi_file['entry_1'][key] = np.string_(value.isoformat())
elif isinstance(value, numbers.Number):
si[key] = value
cxi_file['entry_1'][key] = value
elif isinstance(value, (np.ndarray,list,tuple)):
s1.create_dataset(key, data=np.asarray(value))
cxi_file['entry_1'].create_dataset(key, data=np.asarray(value))
elif isinstance(value, t.Tensor):
asnumpy = value.detach().cpu().numpy()
cxi_file['entry_1'].create_dataset(key, data=asnumpy)
@@ -468,7 +466,7 @@ def add_sample_info(cxi_file, metadata):
# Only store the part of this matrix as defined in the CXI file spec
s1['geometry_1'].create_dataset('orientation',
data=metadata['orientation'].ravel()[:6])
for key, value in metadata.items():
if key == 'orientation':
continue # this is a special case
@@ -514,7 +512,7 @@ def add_detector(cxi_file, distance, basis, corner=None):
It will define all the relevant parameters - distance, pixel size,
detector basis, and corner position (if relevant) based on the provided
information
information.
Parameters
----------
@@ -539,6 +537,8 @@ def add_detector(cxi_file, distance, basis, corner=None):
if isinstance(basis, t.Tensor):
basis = basis.detach().cpu().numpy()
if basis.shape == (2,3):
basis = basis.T
d1['x_pixel_size'] = np.linalg.norm(basis[:,1])
d1['y_pixel_size'] = np.linalg.norm(basis[:,0])
d1.create_dataset('basis_vectors', data=basis)
@@ -591,7 +591,7 @@ def add_dark(cxi_file, dark):
----------
cxi_file : h5py.File
The file to add the mask to
dark : array
dark : array
The dark image(s) to save out to the file
"""
if 'entry_1/instrument_1' not in cxi_file:
CDTools/tools/plotting.py
+34 -34
View File
@@ -32,7 +32,7 @@ def colorize(z):
A complex-valued array
Returns
-------
rgb : list(array)
rgb : list(array)
A list of arrays for the R,G, and B channels of an image
"""
@@ -57,13 +57,13 @@ def get_units_factor(units):
----------
units : str
The abbreviation for the unit type
Returns
-------
factor : float
The factor meters / (unit)
"""
u = units.lower()
if u=='m':
factor=1
@@ -71,7 +71,7 @@ def get_units_factor(units):
factor=1e2
if u=='mm':
factor=1e3
if u=='um':
if u=='um' or u=="$\mu$m":
factor=1e6
if u=='nm':
factor=1e9
@@ -81,16 +81,16 @@ def get_units_factor(units):
factor=1e12
return factor
def plot_amplitude(im, fig = None, basis=None, units='um', cmap='viridis', **kwargs):
def plot_amplitude(im, fig = None, basis=None, units='$\mu$m', cmap='viridis', **kwargs):
"""Plots the amplitude of a complex array with dimensions NxM
If a figure is given explicitly, it will clear that existing figure and
plot over it. Otherwise, it will generate a new figure.
If a basis is explicitly passed, the image will be plotted in real-space
coordinates
Parameters
----------
im : array
@@ -129,11 +129,11 @@ def plot_amplitude(im, fig = None, basis=None, units='um', cmap='viridis', **kwa
basis = basis.detach().cpu().numpy()
basis_norm = np.linalg.norm(basis, axis = 0)
basis_norm = basis_norm * get_units_factor(units)
extent = [0, absolute.shape[-1]*basis_norm[1], 0, absolute.shape[-2]*basis_norm[0]]
else:
extent=None
plt.imshow(absolute, cmap = cmap, extent = extent)
cbar = plt.colorbar()
cbar.set_label('Amplitude (a.u.)')
@@ -144,11 +144,11 @@ def plot_amplitude(im, fig = None, basis=None, units='um', cmap='viridis', **kwa
else:
plt.xlabel('j (pixels)')
plt.ylabel('i (pixels)')
return fig
def plot_phase(im, fig=None, basis=None, units='um', cmap='auto', **kwargs):
def plot_phase(im, fig=None, basis=None, units='$\mu$m', cmap='auto', **kwargs):
""" Plots the phase of a complex array with dimensions NxMx2
If a figure is given explicitly, it will clear that existing figure and
@@ -156,7 +156,7 @@ def plot_phase(im, fig=None, basis=None, units='um', cmap='auto', **kwargs):
If a basis is explicitly passed, the image will be plotted in real-space
coordinates
Parameters
----------
im : array
@@ -194,12 +194,12 @@ def plot_phase(im, fig=None, basis=None, units='um', cmap='auto', **kwargs):
basis = basis.detach().cpu().numpy()
basis_norm = np.linalg.norm(basis, axis = 0)
basis_norm = basis_norm * get_units_factor(units)
extent = [0, phase.shape[-1]*basis_norm[1], 0, phase.shape[-2]*basis_norm[0]]
else:
extent=None
# If the user has matplotlib >=3.0, use the preferred colormap
if cmap == 'auto':
try:
@@ -208,21 +208,21 @@ def plot_phase(im, fig=None, basis=None, units='um', cmap='auto', **kwargs):
plt.imshow(phase, cmap = 'hsv', extent=extent)
else:
plt.imshow(phase)#, cmap = cmap, extent=extent)
cbar = plt.colorbar()
cbar.set_label('Phase (rad)')
if basis is not None:
plt.xlabel('X (' + units + ')')
plt.ylabel('Y (' + units + ')')
else:
plt.xlabel('j (pixels)')
plt.ylabel('i (pixels)')
return fig
def plot_colorized(im, fig=None, basis=None, units='um', **kwargs):
def plot_colorized(im, fig=None, basis=None, units='$\mu$m', **kwargs):
""" Plots the colorized version of a complex array with dimensions NxM
The darkness corresponds to the intensity of the image, and the color
@@ -233,7 +233,7 @@ def plot_colorized(im, fig=None, basis=None, units='um', **kwargs):
If a basis is explicitly passed, the image will be plotted in real-space
coordinates
Parameters
----------
im : array
@@ -258,7 +258,7 @@ def plot_colorized(im, fig=None, basis=None, units='um', **kwargs):
else:
plt.figure(fig.number)
plt.gcf().clear()
if isinstance(im, t.Tensor):
im = cmath.torch_to_complex(im.detach().cpu())
@@ -267,7 +267,7 @@ def plot_colorized(im, fig=None, basis=None, units='um', **kwargs):
basis = basis.detach().cpu().numpy()
basis_norm = np.linalg.norm(basis, axis = 0)
basis_norm = basis_norm * get_units_factor(units)
extent = [0, im.shape[-1]*basis_norm[1], 0, im.shape[-2]*basis_norm[0]]
else:
extent=None
@@ -281,14 +281,14 @@ def plot_colorized(im, fig=None, basis=None, units='um', **kwargs):
else:
plt.xlabel('j (pixels)')
plt.ylabel('i (pixels)')
return fig
def plot_translations(translations, fig=None, units='um', lines=True, **kwargs):
def plot_translations(translations, fig=None, units='$\mu$m', lines=True, **kwargs):
"""Plots a set of probe translations in a nicely formatted way
Parameters
----------
translations : array
@@ -308,9 +308,9 @@ def plot_translations(translations, fig=None, units='um', lines=True, **kwargs):
used_fig : matplotlib.figure.Figure
The figure object that was actually plotted to.
"""
factor = get_units_factor(units)
if fig is None:
fig = plt.figure()
ax = fig.add_subplot(111, **kwargs)
@@ -320,7 +320,7 @@ def plot_translations(translations, fig=None, units='um', lines=True, **kwargs):
if isinstance(translations, t.Tensor):
translations = translations.detach().cpu().numpy()
translations = translations * factor
plt.plot(translations[:,0], translations[:,1],'k.')
if lines:
@@ -330,10 +330,10 @@ def plot_translations(translations, fig=None, units='um', lines=True, **kwargs):
return fig
def plot_nanomap(translations, values, fig=None, units='um', convention='probe'):
def plot_nanomap(translations, values, fig=None, units='$\mu$m', convention='probe'):
"""Plots a set of nanomap data in a flexible way
Parameters
----------
translations : array
@@ -374,12 +374,12 @@ def plot_nanomap(translations, values, fig=None, units='um', convention='probe')
if convention.lower() != 'probe':
trans = trans * -1
s = bbox.width * bbox.height / trans.shape[0] * 72**2 #72 is points per inch
s /= 4 # A rough value to make the size work out
plt.scatter(factor * trans[:,0],factor * trans[:,1],s=s,c=values)
plt.gca().set_facecolor('k')
plt.xlabel('Translation x (' + units + ')')
plt.ylabel('Translation y (' + units + ')')