cSAXS/cdtools
0
0
Fork 0
mirror of https://github.com/cdtools-developers/cdtools.git synced 2026-09-10 05:22:41 +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 + ')')
converters/NSLS2_HXN_hdf5_to_CXI.py
Executable
+131
View File
@@ -0,0 +1,131 @@
"""
Purpose: Convert NSLSII HXN hdf5 files to CXI files for analysis with CDTools.
Author: David Rower
Date: December 2019
"""
import numpy as np
import pickle
import h5py
import os
import CDTools
from CDTools.tools import data as cdtdata
from matplotlib import pyplot as plt
from scipy.spatial.transform import Rotation
from datetime import datetime
def create_cxi_from_NSLS2_HXN_2DFly(data_dir, save_str, scan_number,
wavelength, theta, ROI_corner_xy, metadata):
"""Converts NSLS2 HXN 2D Fly scan data (from pickle and hdf5) to CXI format
Assumes scan files will live in data_dir with naming convention
pickle: <data_dir>/scan_<scan_number>.pickle,
hdf5: <data_dir>/scan_<scan_number>.hdf5,
and will create the file <data_dir>/scan_<scan_number>.cxi.
Parameters
----------
data_dir : str
Input data directory
save_str : str
Output data name
scan_number : int
A scan index number
theta : float
Rotation angle of sample in HXN convention, in degrees
ROI_corner_xy : np.array
1x2 array containing x, y corner of detector ROI
metadata : dict
Contains metadata relevant to the experiment
"""
## Load in pickle and hdf5 files
scan_str = "scan_" + scan_number
# Load pickle (includes useful data about scan not in .hdf5 file)
with open(os.path.join(data_dir, scan_str+".pickle"), 'rb') as f:
scan_pickle = pickle.load(f)
assert scan_pickle['plan_type'] == "FlyPlan2D", "Code only for FlyPlan2D."
# Load hdf5 file
scan_hdf5 = h5py.File(os.path.join(data_dir, scan_str+".h5"), 'r')
## Let's attempt to convert this bad boy
scan_cxi = cdtdata.create_cxi(save_str)
## Add source
cdtdata.add_source(scan_cxi, wavelength=wavelength)
scan_cxi['entry_1/instrument_1/source_1']['name'] = scan_pickle['beamline_id']
## Add sample
theta = np.radians(theta)
sample_unit_vecs = Rotation.from_rotvec(-theta * np.array([0,1,0])).as_dcm()
orientation = np.hstack((sample_unit_vecs[:,0], sample_unit_vecs[:,1]))
translation = np.zeros(3)
sample_info_dict = {
"name" : "TaTe4",
"orientation" : orientation,
"translation" : translation
}
cdtdata.add_sample_info(scan_cxi, sample_info_dict)
## Add other metadata for experiment
metadata['start_time'] = datetime.fromtimestamp(scan_pickle['time'])
cdtdata.add_entry_info(scan_cxi, metadata)
## Add detector
# Constant detector parameters
detector_pixel_size = 55e-6 # meters
detector_height_px = 515 # px ### WARNING: NEED TO CHECK THIS
detector_width_px = 515 # px
# Geometry parameters from scan files
distance = scan_pickle['dist_detector'] * 1e-3 # assuming mm, almost sure
gamma = np.radians(scan_pickle['gamma_detector'])
delta = np.radians(scan_pickle['delta_detector'])
Rg = Rotation.from_rotvec(-gamma * np.array([0,1,0])).as_dcm() # cw about y
Rd = Rotation.from_rotvec(-delta * Rg[:,0]).as_dcm() # cw about rotated x
RdRg = np.matmul(Rd, Rg)
# Define detector basis: row vectors for y and x detector axes
basis = detector_pixel_size * np.array([[0.,-1.,0.],[-1.,0.,0.]])
basis = np.matmul(RdRg,basis.T).T
# Define corner posiiton: first find center, then offset it
corner_pos = np.dot(RdRg, distance * np.array([0.,0.,1.]))
if ROI_corner_xy[0] is None:
ROI_corner_xy[0] = 0.
if ROI_corner_xy[1] is None:
ROI_corner_xy[1] = 0.
corner_pos -= basis[0,:] * (detector_width_px/2. - ROI_corner_xy[0])
corner_pos -= basis[1,:] * (detector_height_px/2. - ROI_corner_xy[1])
# Add detector data finally
cdtdata.add_detector(scan_cxi, distance, basis.T, corner=corner_pos)
## Add data
axes = ['translation'] + scan_pickle['axes'] # THIS IS ONLY FOR FLY2D
data = np.copy(scan_hdf5['entry']['instrument']['detector']['data'])
data[data == 0] = 1 # to prevent divide by zero in log error
cdtdata.add_data(scan_cxi, data, axes)
## Add translations
x_bounds = scan_pickle['scan_range'][0]
y_bounds = scan_pickle['scan_range'][1]
xx, yy = np.meshgrid(np.linspace(*x_bounds, scan_pickle['num1']),
np.linspace(*y_bounds, scan_pickle['num2']))
translations = (1e-6 *
np.stack((xx.ravel(), yy.ravel(), np.zeros_like(xx.ravel())), axis=1))
cdtdata.add_ptycho_translations(scan_cxi, translations)
## Close hdf5 file
scan_hdf5.close()
docs/source/conf.py
+6 -2
View File
@@ -114,7 +114,7 @@ html_theme = 'sphinx_rtd_theme'
# -- Options for HTMLHelp output ---------------------------------------------
# Output file base name for HTML help builder.
htmlhelp_basename = 'ADCDdoc'
htmlhelp_basename = 'CDToolsdoc'
# -- Options for LaTeX output ------------------------------------------------
@@ -140,8 +140,12 @@ latex_elements = {
# Grouping the document tree into LaTeX files. List of tuples
# (source start file, target name, title,
# author, documentclass [howto, manual, or own class]).
#latex_documents = [
# (master_doc, 'CDTools.tex', 'CDTools Documentation',
# 'Abraham Levitan', 'manual'),
#]
latex_documents = [
(master_doc, 'ADCD.tex', 'ADCD Documentation',
('latextoc', 'CDTools.tex', 'CDTools Documentation',
'Abraham Levitan', 'manual'),
]
docs/source/index.rst
+2 -32
View File
@@ -11,37 +11,7 @@
models
tools/index
indices_tables
Introduction to CDTools
=======================
.. include:: intro.rst
CDTools is a python library for ptychography and CDI reconstructions, using an Automatic Differentiation based approach.
.. code-block:: python
# imports
from matplotlib import pyplot as plt
from CDTools.datasets import Ptycho2DDataset
from CDTools.models import SimplePtycho
# Load the file
dataset = Ptycho2DDataset.from_cxi('ptycho_data.cxi')
# Generate a model from the data
model = SimplePtycho.from_dataset(dataset)
# Run a reconstruction
for i, loss in enumerate(model.Adam_optimize(10, dataset)):
print(i, loss)
# And look at the results!
model.inspect(dataset)
model.compare(dataset)
plt.show()
CDTools makes it simple to load and inspect data stored in .cxi files using python scripts. Several reconstruction models for common geometries are included "out of the box". For more advanced users, it includes a bunch of modular functions for AD ptychography, which can then be used right away from the same scripting framework.
The high-level interface to CDTools is built on a lower level "three-legged stool". This consists of tools to access stored data, tools to visualize data and reconstructions, and tools that implement basic operations relevant to coherent diffraction. All of these tools can be used directly alongside the high-level interface, when needed.
Enough blabber. If you're interested, read the docs!
docs/source/installation.rst
+1 -1
View File
@@ -13,7 +13,7 @@ It is recommended that you clone the repository, rather than just downloading th
Step 2: Install Dependencies
----------------------------
The dependencies for CDTools can be installed, if you are managing your environment with anaconda, by running
The dependencies for CDTools can be installed, if you are managing your environment with anaconda, by running the following command in the top level directory of the package:
.. code:: bash
docs/source/intro.rst
+40
View File
@@ -0,0 +1,40 @@
Introduction to CDTools
=======================
.. only:: latex
Introduction to CDTools
-----------------------
CDTools is a python library for ptychography and CDI reconstructions, using an Automatic Differentiation based approach.
.. code-block:: python
# imports
from matplotlib import pyplot as plt
from CDTools.datasets import Ptycho2DDataset
from CDTools.models import SimplePtycho
# Load the file
dataset = Ptycho2DDataset.from_cxi('ptycho_data.cxi')
# Generate a model from the data
model = SimplePtycho.from_dataset(dataset)
# Run a reconstruction
for i, loss in enumerate(model.Adam_optimize(10, dataset)):
print(i, loss)
# And look at the results!
model.inspect(dataset)
model.compare(dataset)
plt.show()
CDTools makes it simple to load and inspect data stored in .cxi files using python scripts. Several reconstruction models for common geometries are included "out of the box". For more advanced users, it includes a bunch of modular functions for AD ptychography, which can then be used right away from the same scripting framework.
The high-level interface to CDTools is built on a lower level "three-legged stool". This consists of tools to access stored data, tools to visualize data and reconstructions, and tools that implement basic operations relevant to coherent diffraction. All of these tools can be used directly alongside the high-level interface, when needed.
Enough blabber. If you're interested, read the docs!
docs/source/latextoc.rst
+14
View File
@@ -0,0 +1,14 @@
.. toctree::
:maxdepth: 1
intro
installation
examples
tutorial
general
datasets
models
tools/index
indices_tables
examples/simple_ptycho_model.py
+20
View File
@@ -105,3 +105,23 @@ class SimplePtycho(CDIModel):
if __name__ == '__main__':
from basic_ptycho_dataset import BasicPtychoDataset
from h5py import File
from matplotlib import pyplot as plt
filename = 'example_data/lab_ptycho_data.cxi'
with File(filename, 'r') as f:
dataset = BasicPtychoDataset.from_cxi(f)
model = SimplePtycho.from_dataset(dataset)
#model.to(device='cuda')
#dataset.get_as(device='cuda')
for i, loss in enumerate(model.Adam_optimize(100, dataset)):
model.inspect(dataset)
print(i,loss)
model.compare(dataset)
plt.show()
kiara_debugging_code/debugging_mod.py
+112
View File
@@ -0,0 +1,112 @@
import numpy as np
from matplotlib import pyplot as plt
import imageio
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
class ImageSeries:
def __init__(self, images, crop,fps=None):
self.images = images #list of images (numpy arrays), saved itself as a numpy array
self.L = len(self.images)
self.x1, self.x2, self.y1, self.y2 = crop #(x1,y1) and (x2,y2) are cropping coords
self.crimages = []
for i in range(self.L):
self.crimages.append(self.images[i][self.x1:self.x2, self.y1:self.y2])
self.crimages = np.array(self.crimages)
self.fps = fps
def show_frame(self,t): #t = time to show
if 0 <= t <= self.L:
plt.imshow(self.crimages[t],interpolation="none")
plt.show()
return
def save_vid(self, filename, secspersec):
imageio.mimwrite(filename, self.crimages, fps=self.fps*secspersec)
def waterfall(self,x,y):
if x == None:
#plot row y=y over time
return self.crimages[:,:,y]
elif y == None:
return np.transpose(self.crimages[:,x,:])
def plot_waterfall(self,x,y,tint=None):
waterfall = self.waterfall(x,y)
if x == None:
#plot col y over time
fig = plt.figure()
W = fig.add_subplot(111)
W.imshow(waterfall,interpolation="none")
if tint:
plt.yticks(np.arange(0,self.L,self.fps*tint),np.arange(0,self.L/self.fps,tint))
W.set_title("waterfall plot of row y = " + str(y))
W.set_ylabel("time [s]")
plt.show()
return
elif y == None:
#plot col x=x over time
fig = plt.figure()
W = fig.add_subplot(111)
W.imshow(waterfall,interpolation="none")
if tint:
plt.xticks(np.arange(0,self.L,self.fps*tint),np.arange(0,self.L/self.fps,tint))
W.set_title("waterfall plot of col x = " + str(x))
W.set_xlabel("time [s]")
plt.show()
return
return
def Ipixel(self,pixel):
px,py = pixel
return self.crimages[:,px,py]
def plot_Ipixel(self,pixel,description=""):
I = self.Ipixel(pixel)
fig = plt.figure()
f1 = fig.add_subplot(111)
f1.set_title("I(t) for pixel" + str(pixel) + " (" + description + ")")
f1.set_xlabel("time [s]")
f1.plot(np.arange(0,self.L/self.fps,1/self.fps),I)
plt.show()
return
def fftIpixel(self,pixel):
I = self.Ipixel(pixel)
IfreqA = np.fft.fft(I)/self.L
Ifreq = np.fft.fftfreq(self.L,d=(1/self.fps))
return(Ifreq, IfreqA)
def plot_fftIpixel(self,pixel):
Ifreq, IfreqA = self.fftIpixel(pixel)
fig = plt.figure()
f1 = fig.add_subplot(111)
f1.set_title("FFT for pixel" + str(pixel))
f1.set_xlabel("frequency [1/s]")
f1.plot(Ifreq, abs(IfreqA))
plt.show()
return
def g2(self,pixel):
I = self.Ipixel(pixel)
g2 = []
avgsq = np.mean(I)**2
for tau in range(len(I)-1):
if tau == 0:
dotp = 0
for t in range(len(I)):
dotp += I[t]*I[t]
g2.append(dotp/(len(I)*avgsq) )
elif tau != 0:
dotp = 0
for t in range(len(I)-tau):
dotp += I[t]*I[t+tau]
g2.append( dotp / (len(I[:-tau])*avgsq))
g2 = np.array(g2)
return g2
kiara_debugging_code/vid_read.py
+50
View File
@@ -0,0 +1,50 @@
import numpy as np
from matplotlib import pyplot as plt
import imageio
from PIL import Image, ImageSequence
from debugging_mod import ImageSeries
vid = Image.open('data_10_4/kiara_20fps_redlaser_vid.tif')
vidarray = []
for i, page in enumerate(ImageSequence.Iterator(vid)):
pg = np.array(page)
vidarray.append(pg)
vidarray = np.array(vidarray)
RedLaserExp = ImageSeries(vidarray, (430,606,590,766),fps=20) #cropping x1:x2, y1:y2
#with np.load('data_9_28/kiara_data_300sec_green') as data:
# GreenLaserExp = ImageSeries(data['arr_0'], (500,676,580,756), fps=5)
#crop to: x1=500, x2=676, y1=580, y2=756
#++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
#Note: human eye can see at c. 150 fps
#inner top right corner of disk (120,70)
#inner top left corner of disk (66,63)
#inner bottom left corner of disk (66,107)
#inner bottom (87,118)
#inner right (125, 86)
#pixel intensity plot: ---------------------------------------
#RedLaserExp.plot_Ipixel((125,86),description="inner right")
#RedLaserExp.plot_Ipixel((87,118),description="inner bottom")
#RedLaserExp.plot_Ipixel((66,63),description="inner top left")
#pixel fft plot: --------------------------------------------
#RedLaserExp.plot_fftIpixel((125,86))
#RedLaserExp.plot_fftIpixel((87,118))
#RedLaserExp.plot_fftIpixel((66,63))
#waterfall plot for row: -------------------------------------
#RedLaserExp.plot_waterfall(None,118,tint=1)
#waterfall plot for col: -------------------------------------
#RedLaserExp.plot_waterfall(66,None,tint=1)
#save a video: -----------------------------------------------
#RedLaserExp.save_vid("redlaser_10-4_20fps_1x.mp4",1)