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

Start cleaning up the repo and tidying up the example scripts

Browse Source
This commit is contained in:
Abe Levitan
2021-06-28 11:14:41 -04:00
parent 76e6d5efae
commit 8ddcc7683f
18 changed files with 127 additions and 574 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CDTools/models/base.py
+12 -5
View File
@@ -163,12 +163,19 @@ class CDIModel(t.nn.Module):
exit()
sim_patterns = self.forward(*inp)
#sim_patterns.retain_grad()
if hasattr(self, 'mask'):
loss = self.loss(pats,sim_patterns, mask=self.mask)
else:
loss = self.loss(pats,sim_patterns)
loss.backward()
loss.backward()#retain_variables=True)
#plt.figure()
#plt.imshow(sim_patterns.grad.cpu()[0])
#plt.colorbar()
#plt.show()
total_loss += loss.detach()
#print('probe grad')
@@ -346,10 +353,10 @@ class CDIModel(t.nn.Module):
# Define the optimizer
#optimizer = t.optim.LBFGS(self.parameters(),
# lr = lr, history_size=history_size)
optimizer = MyLBFGS(self.parameters(),
lr = lr, history_size=history_size)
optimizer = t.optim.LBFGS(self.parameters(),
lr = lr, history_size=history_size)
#optimizer = MyLBFGS(self.parameters(),
# lr = lr, history_size=history_size)
return self.AD_optimize(iterations, data_loader, optimizer,
regularization_factor=regularization_factor,
CDTools/models/fancy_ptycho.py
+4 -4
View File
@@ -66,9 +66,9 @@ class FancyPtycho(CDIModel):
if background is None:
if detector_slice is not None:
background = 1e-6 * t.ones(self.probe[0][self.detector_slice])
background = 1e-6 * t.ones(self.probe[0][self.detector_slice].shape)
else:
background = 1e-6 * t.ones(self.probe[0])
background = 1e-6 * t.ones(self.probe[0].shape)
self.background = t.nn.Parameter(t.Tensor(background).to(t.float32))
@@ -229,7 +229,7 @@ class FancyPtycho(CDIModel):
mask = None
if probe_support_radius is not None:
probe_support = t.zeros_like(probe[0])
probe_support = t.zeros(probe[0].shape,dtype=t.bool)
xs, ys = np.mgrid[:probe.shape[-2],:probe.shape[-1]]
xs = xs - np.mean(xs)
ys = ys - np.mean(ys)
@@ -239,7 +239,7 @@ class FancyPtycho(CDIModel):
probe = probe * probe_support[None,:,:]
else:
probe_support = None;
probe_support = None
if restrict_obj != -1:
ro = restrict_obj
CDTools/models/multislice_2d_ptycho.py
+61 -14
View File
@@ -14,6 +14,14 @@ __all__ = ['Multislice2DPtycho']
class Multislice2DPtycho(CDIModel):
@property
def probe(self):
return t.complex(self.probe_real,self.probe_imag)
@property
def obj(self):
return t.complex(self.obj_real,self.obj_imag)
def __init__(self, wavelength, detector_geometry,
probe_basis,
probe_guess, obj_guess, dz, nz,
@@ -27,6 +35,7 @@ class Multislice2DPtycho(CDIModel):
bandlimit=None,
subpixel=True,
exponentiate_obj=True,
low_res_obj=False,
fourier_probe=False,
prevent_aliasing=True,
phase_only=False,
@@ -58,6 +67,7 @@ class Multislice2DPtycho(CDIModel):
self.units = units
self.phase_only=phase_only
self.prevent_aliasing=prevent_aliasing
self.low_res_obj = low_res_obj
if mask is None:
self.mask = mask
@@ -70,11 +80,21 @@ class Multislice2DPtycho(CDIModel):
self.probe_norm = 1 * t.max(t.abs(probe_guess[0]).to(t.float32))
else:
self.probe_norm = 1 * t.max(t.abs(probe_guess).to(t.float32))
self.probe = t.nn.Parameter(probe_guess.to(t.complex64)
/ self.probe_norm)
pg = probe_guess.to(t.complex64)/self.probe_norm
self.probe_real = t.nn.Parameter(pg.real)
self.probe_imag = t.nn.Parameter(pg.imag)
#self.probe = t.complex(self.probe_real,self.probe_imag)
og = obj_guess.to(t.complex64)
self.obj_real = t.nn.Parameter(og.real)
self.obj_imag = t.nn.Parameter(og.imag)
#self.obj = t.complex(self.obj_real,self.obj_imag)
self.obj = t.nn.Parameter(obj_guess.to(t.complex64))
#self.probe = t.nn.Parameter(probe_guess.to(t.complex64)
# / self.probe_norm)
#self.obj = t.nn.Parameter(obj_guess.to(t.complex64))
if background is None:
if detector_slice is not None:
@@ -126,7 +146,7 @@ class Multislice2DPtycho(CDIModel):
@classmethod
def from_dataset(cls, dataset, dz, nz, probe_convergence_semiangle, probe_size=None, padding=0, n_modes=1, dm_rank=None, translation_scale = 1, saturation=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=True, bandlimit=None, replicate_slice=False, subpixel=True, exponentiate_obj=True, units='um', fourier_probe=False, phase_only=False, prevent_aliasing=True):
def from_dataset(cls, dataset, dz, nz, probe_convergence_semiangle, padding=0, n_modes=1, dm_rank=None, translation_scale = 1, saturation=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=True, bandlimit=None, replicate_slice=False, subpixel=True, exponentiate_obj=True, units='um', fourier_probe=False, phase_only=False, prevent_aliasing=True, probe_support_radius=None, low_res_obj=False):
wavelength = dataset.wavelength
det_basis = dataset.detector_geometry['basis']
@@ -178,8 +198,14 @@ class Multislice2DPtycho(CDIModel):
# Next generate the object geometry from the probe geometry and
# the translations
pix_translations = tools.interactions.translations_to_pixel(probe_basis, translations, surface_normal=surface_normal)
if low_res_obj: # obj in half normal resolution
pix_translations /= 2
obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=200)
obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=100)
if low_res_obj:
obj_size, min_translation = tools.initializers.calc_object_setup(t.as_tensor(probe_shape)//2, pix_translations, padding=100)
if hasattr(dataset, 'background') and dataset.background is not None:
background = t.sqrt(dataset.background)
@@ -187,7 +213,8 @@ class Multislice2DPtycho(CDIModel):
background = None
# Finally, initialize the probe and object using this information
probe = tools.initializers.STEM_style_probe(dataset, probe_shape, det_slice, probe_convergence_semiangle, propagation_distance=propagation_distance, oversampling=oversampling)
#probe = tools.initializers.STEM_style_probe(dataset, probe_shape, det_slice, probe_convergence_semiangle, propagation_distance=propagation_distance, oversampling=oversampling)
probe = tools.initializers.SHARP_style_probe(dataset, probe_shape, det_slice, propagation_distance=propagation_distance, oversampling=oversampling)
# Now we initialize all the subdominant probe modes
probe_max = t.max(t.abs(probe))
@@ -239,17 +266,25 @@ class Multislice2DPtycho(CDIModel):
else:
mask = None
# probe_support = t.zeros(probe[0].shape,dtype=t.bool)
# xs, ys = np.mgrid[:probe.shape[-2],:probe.shape[-1]]
# xs = xs - np.mean(xs)
# ys = ys - np.mean(ys)
# Rs = np.sqrt(xs**2 + ys**2)
if probe_support_radius is not None:
probe_support = t.zeros(probe[0].shape, dtype=t.bool)
xs, ys = np.mgrid[:probe.shape[-2],:probe.shape[-1]]
xs = xs - np.mean(xs)
ys = ys - np.mean(ys)
Rs = np.sqrt(xs**2 + ys**2)
probe_support[Rs<probe_support_radius] = 1
probe = probe * probe_support[None,:,:]
else:
probe_support = None
probe = probe
return cls(wavelength, det_geo, probe_basis, probe, obj, dz, nz,
detector_slice=det_slice,
surface_normal=surface_normal,
probe_support=probe_support,
min_translation=min_translation,
translation_offsets = translation_offsets,
weights=Ws, mask=mask, background=background,
@@ -261,13 +296,16 @@ class Multislice2DPtycho(CDIModel):
exponentiate_obj=exponentiate_obj,
units=units, fourier_probe=fourier_probe,
phase_only=phase_only,
prevent_aliasing=prevent_aliasing)
prevent_aliasing=prevent_aliasing,
low_res_obj=low_res_obj)
def interaction(self, index, translations):
pix_trans = tools.interactions.translations_to_pixel(self.probe_basis,
translations,
surface_normal=self.surface_normal)
if self.low_res_obj:
pix_trans /=2
pix_trans -= self.min_translation
if self.translation_offsets is not None:
@@ -324,9 +362,13 @@ class Multislice2DPtycho(CDIModel):
exit_waves, obj, pix_trans,
shift_probe=True, multiple_modes=True)
else:
#exit_waves = tools.interactions.ptycho_2D_round(
# exit_waves, obj, pix_trans,
# multiple_modes=True,upsample_obj=self.prevent_aliasing)
exit_waves = tools.interactions.ptycho_2D_round(
exit_waves, obj, pix_trans,
multiple_modes=True,upsample_obj=self.prevent_aliasing)
multiple_modes=True,upsample_obj=self.low_res_obj)
elif self.obj.dim() == 3:
# If separate slices
@@ -339,6 +381,11 @@ class Multislice2DPtycho(CDIModel):
exit_waves, obj[i], pix_trans,
multiple_modes=True)
#if self.iteration_count >= 1:
# plt.imshow(t.abs(tools.propagators.far_field(
# exit_waves[0,0].detach()).cpu()))
# plt.colorbar()
# plt.show()
if i < self.nz-1: #on all but the last iteration
exit_waves = tools.propagators.near_field(
exit_waves,self.as_prop)
CDTools/tools/initializers/initializers.py
+10 -8
View File
@@ -335,15 +335,15 @@ def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None, over
probe_guess = inverse_far_field(probe_fft).numpy()
# Now we remove the central pixel
center = np.array(probe_guess.shape) // 2
#center = np.array(probe_guess.shape) // 2
# I'm always unsure whether to use this modification:
probe_guess[center[0], center[1]]=np.mean([
probe_guess[center[0]-1, center[1]],
probe_guess[center[0]+1, center[1]],
probe_guess[center[0], center[1]-1],
probe_guess[center[0], center[1]+1]])
#probe_guess[center[0], center[1]]=np.mean([
# probe_guess[center[0]-1, center[1]],
# probe_guess[center[0]+1, center[1]],
# probe_guess[center[0], center[1]-1],
# probe_guess[center[0], center[1]+1]])
probe_guess = t.as_tensor(probe_guess, dtype=t.complex64)
@@ -420,8 +420,10 @@ def STEM_style_probe(dataset, shape, det_slice, convergence_semiangle, propagati
# Fourier space (simulated on a larger stage in real space). That factor
# is defined by oversampling.
probe_basis = dataset.detector_geometry['basis'] / oversampling
probe_basis = (t.as_tensor(dataset.detector_geometry['basis'],
dtype=t.float32)
/ oversampling)
mean_im = t.mean(dataset.patterns,dim=0)
center = image_processing.centroid(mean_im)
CDTools/tools/propagators/propagators.py
+6 -1
View File
@@ -369,10 +369,14 @@ def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, r
offset = t.zeros([3], dtype=basis.dtype)
offset[2] = z
# And we call the generalized function!
propagator = generate_generalized_angular_spectrum_propagator(shape, basis,
wavelength, offset,
wavelength, offset,
propagate_along_offset=remove_z_phase, **kwargs)
if z < 0:
propagator = t.conj(propagator)
# Bandlimiting is not implemented in the generalized function, because it
# has a less clear meaning in that setting, so we apply it here instead
@@ -508,6 +512,7 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o
if propagation_vector is not None:
perpendicular_dir *= t.sign(t.dot(perpendicular_dir,propagation_vector))
else:
pass
perpendicular_dir *= t.sign(t.dot(perpendicular_dir,offset_vector))
# Then, if we have a propagation vector, we shift the in-plane
README.md
+6 -6
View File
@@ -6,21 +6,21 @@ CDTools is a python library for ptychography and CDI reconstructions, using an A
# imports
from matplotlib import pyplot as plt
from CDTools.datasets import Ptycho_2D_Dataset
from CDTools.models import SimplePtycho
from CDTools.models import FancyPtycho
# Load the file
dataset = Ptycho_2D_Dataset.from_cxi('ptycho_data.cxi')
# Generate a model from the data
model = SimplePtycho.from_dataset(dataset)
model = FancyPtycho.from_dataset(dataset)
# Run a reconstruction
for i, loss in enumerate(model.Adam_optimize(10, dataset)):
print(i, loss)
for loss in model.Adam_optimize(10, dataset):
print(model.report())
# And look at the results!
model.inspect(dataset)
model.compare(dataset)
model.inspect(dataset) # See the reconstructed object, probe, etc.
model.compare(dataset) # See how the simulated and measured patterns compare
plt.show()
```
conda_requirements.txt
+1 -1
View File
@@ -2,7 +2,7 @@ numpy>=1.0
scipy>=1.0
matplotlib>=2.0
python-dateutil
pytorch>=1.8.0
pytorch>=1.9.0
h5py>=2.1
pytest
sphinx
converters/NSLS2_HXN_hdf5_to_CXI.py
-131
View File
@@ -1,131 +0,0 @@
"""
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()
converters/TITAN_stem.py
-208
View File
@@ -1,208 +0,0 @@
"""
Purpose: Convert file collection from Jim Lebeau's TITAN microscope to .CXI
Author: Abe Levitan
Date: January 2019
"""
import numpy as np
import pickle
import h5py
import os
import CDTools
from CDTools.tools import data as cdtdata
from CDTools.datasets import Ptycho2DDataset
from matplotlib import pyplot as plt
from scipy.spatial.transform import Rotation
from datetime import datetime
import xml.etree.ElementTree as ET
def load_raw_image_stack(filename):
# The resulting data is an array of (exposure, image-i, image-j),
# with image0i corresponding to y and image-j corresponding to x
# Note that the real-space scanning is done from the bottom right
# corner, first heading left (in x) then scanning up.
rawdata = np.fromfile(filename,dtype='<f4')
if len(rawdata) % (128*130) != 0:
raise IndexError('The raw data file doesn\'t seem to have the right number of values stored in it')
numshots = len(rawdata) // (128*130)
# One of the directions seems to be flipped
return rawdata.reshape(numshots,130,128)[:,:128,:][:,:,::-1].copy()
def load_metadata(filename):
return ET.parse(filename).getroot()
def get_scan_shape(metadata):
sp = metadata.find("scan_parameters[@mode='acquire']")
#print(sp)
# exit()
shape_x = int(sp.find('scan_resolution_x').text)
shape_y = int(sp.find('scan_resolution_y').text)
return [shape_x,shape_y]
def get_camera_length(metadata):
iomm = metadata.find('iom_measurements')
ncl = iomm.find('nominal_camera_length')
return float(ncl.text)
def get_scan_steps(metadata):
shape = get_scan_shape(metadata)
iomm = metadata.find('iom_measurements')
fov = iomm.find('full_scan_field_of_view')
scale = float(fov.find('scale_factor').text)
# Not sure why I need to divide by this scale factor or why it exists
# but this seems to produce the correct numbers
xfov = float(fov.find('x').text) / scale
yfov = float(fov.find('y').text) / scale
return np.array([xfov,yfov]) / np.array(shape)
def gen_scan_grid(shape, step):
ys, xs = np.mgrid[:shape[0],:shape[1]]
xs = xs * step[0]
ys = ys * step[1]
return np.stack((xs.ravel(),ys.ravel(),np.zeros(ys.ravel().shape))).transpose()
def get_electron_energy(metadata):
iomm = metadata.find('iom_measurements')
energy = float(iomm.find('high_voltage').text) / 1000 # to keV
return energy * 1.602e-16 # to Joules
h = 6.626e-34
c = 2.998e8
me = 9.109e-31
def calculate_wavelength(electron_energy):
return h*c / np.sqrt(electron_energy**2 + 2*c**2 * me * electron_energy)
def generate_detector_geometry(distance, pitches):
basis = np.array([[0,-pitches[1]],[-pitches[0],0],[0,0]])
return {'basis':basis, 'distance':distance}
def generate_dataset(translations, patterns, detector_geometry, electron_energy):
wavelength = calculate_wavelength(electron_energy)
print('Wavelength:',wavelength)
print('Pixel NA:',(-detector_geometry['basis'][0,1]/detector_geometry['distance']))
exit()
return Ptycho2DDataset(translations, patterns, wavelength=wavelength, detector_geometry=det_geo)
# Change this to allow for command-line introduction of the data folder
# I think that the data folder should be the first command-line arg, and
# be defaulted to the current folder
# Then, the image filename should by default be the only .raw file in the
# folder if there is exactly one. If none or more than one, throw an error
# Then, the metadata filename should be the only .xml file in the folder
# if there is exactly one, otherwise it should throw an error.
# Next, there should be some .csv file or similar containing the calibration
# of the scan size and pixel size. The program should give a report of how
# well the calibrated and naive values match. If no calibration is given,
# should indicate that it is using the naive values.
# Finally, the output filename should by default be the xml filename,
# and can be overriden by a clarg
#data_folder = '/media/Data Bank/Electron Ptycho MoS2/Session2/MoS2_Pty_session2/acquisition_1_50_nm_positive'
#image_filename = 'scan_x128_y128.raw'
#metadata_filename = 'acquisition_1_50_nm_positive.xml'
data_folder = '/media/Data Bank/Electron Ptycho MoS2/Session2/MoS2_Pty_session2/acquisition_1_20_nm_positive'
image_filename = 'scan_x128_y128.raw'
metadata_filename = 'acquisition_1_20_nm_positive.xml'
#data_folder = '/media/Data Bank/Electron Ptycho MoS2/Session2/MoS2_Pty_session2/acquisition_1_100nm_positive'
#image_filename = 'scan_x128_y128.raw'
#metadata_filename = 'acquisition_1_100nm_positive.xml'
#data_folder = '/media/Data Bank/Electron Ptycho MoS2/Lower_Convergence_angle/Smaller_sampling/acquisition_1_convergence_28mrad_14oMx_285mm'
#image_filename = 'scan_x128_y128.raw'
#metadata_filename = 'acquisition_1_convergence_28mrad_14oMx_285mm.xml'
#data_folder = '/media/Data Bank/Electron Ptycho MoS2/Lower_Convergence_angle/Larger_sampling/acquisition_1_14o5Mx_285mm_28mrad'
#image_filename = 'scan_x256_y256.raw'
#metadata_filename = 'acquisition_1_14o5Mx_285mm_28mrad.xml'
#data_folder = '/media/Data Bank/Electron Ptycho MoS2/Lower_Convergence_angle/Defocus_series/Defocus_positive/acquisition_2_60nm_positive_defocus_140Mx_28mrad_285mm'
#image_filename = 'scan_x128_y128.raw'
#metadata_filename = 'acquisition_2_60nm_positive_defocus_140Mx_28mrad_285mm.xml'
#data_folder = '/media/Data Bank/Electron Ptycho MoS2/Lower_Convergence_angle/Defocus_series/Defocus_negative/acquisition_2_14oMx_defocus_negative_50nm_28mrad_285mm'
#image_filename = 'scan_x128_y128.raw'
#metadata_filename = 'acquisition_2_14oMx_defocus_negative_50nm_28mrad_285mm.xml'
save_filename = 'Initial_CXI_Generation.cxi'
#data_folder = '/media/Data Bank/ptychography_firsttry/out_of_focus_58Mx_1ms_reso80x80_ss1'
#image_filename = 'scan_x80_y80.raw'
#save_filename = 'test_defocus_newcalibration.cxi'
#metadata_filename = 'out_of_focus_58Mx_1ms_reso80x80_ss1.xml'
#data_folder = '/media/Data Bank/ptychography_firsttry/acquisition_3'
#image_filename = 'scan_x80_y80.raw'
#save_filename = 'test_acq3_newcalibration.cxi'
#metadata_filename = 'acquisition_3.xml'
metadata = load_metadata(data_folder + '/' + metadata_filename)
scan_shape = get_scan_shape(metadata)
#scan_shape = 80
# These are reasonable initial guesses, until we get calibration data
#scan_step = 0.2e-10 #Angstrom, old value from manual measurement
scan_steps = get_scan_steps(metadata)
# This is something I can calculate from the detector length
# A good calibration is to assume that the pixel size is 0.2276 mm and
# the detector distance is equal to the nomninal camera length
camera_length = get_camera_length(metadata)
detector_distance = camera_length
# This gets the electron energy
electron_energy = get_electron_energy(metadata)
pixel_pitches = [500e-6,500e-6] # best match to Abinash's calibration
# Also feels right, even though the docs I find for the EMPAD shows 150um pixels
# These came from a calibration done by Xi
#pixel_pitches = [0.231e-3,0.231e-3] # best guess near length=0.230
# pixel_pitches = [0.2276e-3,0.2276e-3] # best overall average
# old manual calibration
#pixel_pitches = [150e-6,150e-6]
#detector_distance = 100e-3 # mm
#print([pp / detector_distance for pp in pixel_pitches])
#exit()
# Important question: Check which side the images fill in from
data = load_raw_image_stack(data_folder + '/' + image_filename)
#data[:,30:-30,30:-30] = 0 # For HAADF
scan_points = gen_scan_grid(scan_shape,scan_steps)
det_geo = generate_detector_geometry(detector_distance, pixel_pitches)
# The first image will be something like 20% larger than the rest...
#dataset = generate_dataset(scan_points, data, det_geo, electron_energy)
dataset = generate_dataset(scan_points[1:], data[1:], det_geo, electron_energy)
dataset.inspect(units='nm')
plt.show()
dataset.to_cxi(data_folder + '/' + save_filename)
docs/source/intro.rst
+6 -7
View File
@@ -7,23 +7,22 @@ Introduction to CDTools
Introduction to CDTools
-----------------------
CDTools is a python library for ptychography and CDI reconstructions, using an Automatic Differentiation based approach.
CDTools is a python library for ptychography and CDI reconstructions, using an Automatic Differentiation (AD) based approach.
.. code-block:: python
# imports
import CDTools
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')
dataset = CDTools.datasets.Ptycho2DDataset.from_cxi('ptycho_data.cxi')
# Generate a model from the data
model = SimplePtycho.from_dataset(dataset)
model = CDTools.models.SimplePtycho.from_dataset(dataset)
# Run a reconstruction
for i, loss in enumerate(model.Adam_optimize(10, dataset)):
for i, loss in enumerate(model.Adam_optimize(20, dataset)):
print(i, loss)
# And look at the results!
@@ -32,7 +31,7 @@ CDTools is a python library for ptychography and CDI reconstructions, using an A
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.
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 modular functions for AD ptychography. These can can be used to construct new forward models.
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.
examples/basic_ptycho_dataset.py
-1
View File
@@ -1,4 +1,3 @@
from __future__ import division, print_function, absolute_import
import numpy as np
import torch as t
from matplotlib import pyplot as plt
examples/example_reconstructions/gold_balls.mat
BIN
View File
Binary file not shown.
examples/example_reconstructions/gold_balls.pickle
BIN
View File
Binary file not shown.
examples/gold_ball_ptycho.py
+10 -11
View File
@@ -1,34 +1,33 @@
from __future__ import division, print_function, absolute_import
import CDTools
from matplotlib import pyplot as plt
import pickle
import time
from scipy import io
# First, we load an example dataset from a .cxi file
filename = 'example_data/AuBalls_700ms_30nmStep_3_6SS_filter.cxi'
dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename)
# Next, we create a ptychography model from the dataset
# Note that we explicitly as for two incoherent probe modes
model = CDTools.models.FancyPtycho.from_dataset(dataset, n_modes=2,dm_rank=0, probe_support_radius=50)
# Note that we explicitly ask for two incoherent probe modes
model = CDTools.models.FancyPtycho.from_dataset(dataset, n_modes=2)
# Let's do this reconstruction on the GPU, shall we?
model.to(device='cuda')
dataset.get_as(device='cuda')
# Now, we run a short reconstruction from the dataset
for loss in model.Adam_optimize(100, dataset, batch_size=50, schedule=True):
# And we liveplot the updates to the model as they happen
print(model.report())
model.inspect(dataset)
# And we save the reconstruction out to a file
#with open('example_reconstructions/gold_balls.pickle', 'wb') as f:
# pickle.dump(model.save_results(dataset),f)
# This orthogonalizes the incoherent probe modes
model.tidy_probes()
# And we save out the results as a .mat file
io.savemat('example_reconstructions/gold_balls.mat',
model.save_results(dataset))
# Finally, we plot the results
model.inspect(dataset)
#model.compare(dataset)
model.compare(dataset)
plt.show()
examples/simple_ptycho.py
+4 -7
View File
@@ -1,8 +1,6 @@
from __future__ import division, print_function, absolute_import
import CDTools
from matplotlib import pyplot as plt
import time
# First, we load an example dataset from a .cxi file
filename = 'example_data/lab_ptycho_data.cxi'
dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename)
@@ -10,11 +8,10 @@ dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename)
# Next, we create a ptychography model from the dataset
model = CDTools.models.SimplePtycho.from_dataset(dataset)
t = time.time()
# Now, we run a short reconstruction from the dataset!
for i, loss in enumerate(model.Adam_optimize(40, dataset,lr=0.01,batch_size=25)):#,batch_size=10000)):#0.001)):
print(i, loss)
print(time.time() - t)
for loss in model.Adam_optimize(20, dataset):
print(model.report())
# Finally, we plot the results
model.inspect(dataset)
model.compare(dataset)
kiara_debugging_code/debugging_mod.py
-112
View File
@@ -1,112 +0,0 @@
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
@@ -1,50 +0,0 @@
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)
setup.py
+7 -8
View File
@@ -5,8 +5,9 @@ with open("README.md", "r") as fh:
setuptools.setup(
name="CDTools",
version="0.0.1",
author="Abe Levitan, Madelyn Cain",
version="0.1.1",
python_requires='>3.4', # we use pathlib, introduced in 3.4
author="Abe Levitan",
author_email="alevitan@mit.edu",
description="Coherent Diffraction Tools",
long_description=long_description,
@@ -15,15 +16,13 @@ setuptools.setup(
install_requires=[
"numpy>=1.0",
"scipy>=1.0",
"matplotlib>=2.0",
"matplotlib>=2.0", # Matplotlib 2.0 introduces better colormaps and no I'm not sorry
"python-dateutil",
"torch>=1.9.0", #1.9.0 implements support for autograd on indexed complex tensors, key to allowing us to use complex tensors in the forward models
"h5py>=2.1",
"pathlib2 ; python_version<'3.4'"],
"torch>=1.9.0", #1.9.0 implements support for autograd on indexed complex tensors, which we need in order to use complex tensors in the forward models
"h5py>=2.1"],
extras_require={
'tests': ["pytest"],
'docs': ["sphinx","sphinx-argparse","sphinx_rtd_theme"],
":python_version<'3.4'": ["pathlib2"],
'docs': ["sphinx","sphinx-argparse","sphinx_rtd_theme"]
},
packages=setuptools.find_packages(),
classifiers=[