mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-09 13:02:41 +02:00
Update the RPI model and clean up the examples section
This commit is contained in:
Abe Levitan
@@ -29,5 +29,6 @@ from CDTools.models.simple_ptycho import SimplePtycho
|
||||
from CDTools.models.fancy_ptycho import FancyPtycho
|
||||
from CDTools.models.bragg_2d_ptycho import Bragg2DPtycho
|
||||
from CDTools.models.multislice_2d_ptycho import Multislice2DPtycho
|
||||
from CDTools.models.rpi import RPI
|
||||
#from CDTools.models.s_matrix_ptycho import SMatrixPtycho
|
||||
#from CDTools.models.rpi import RPI
|
||||
|
||||
|
||||
@@ -369,7 +369,7 @@ class Bragg2DPtycho(CDIModel):
|
||||
|
||||
pix_trans -= self.min_translation
|
||||
props -= self.median_propagation
|
||||
|
||||
|
||||
if self.translation_offsets is not None:
|
||||
pix_trans += self.translation_scale * self.translation_offsets[index]
|
||||
|
||||
@@ -379,6 +379,8 @@ class Bragg2DPtycho(CDIModel):
|
||||
|
||||
|
||||
for j in range(prs.shape[0]):
|
||||
# I believe this -1 sign is in error, but I need a dataset with
|
||||
# well understood geometry to figure it out
|
||||
propagator = t.exp(
|
||||
1j*(props[j]*(2*np.pi)/self.wavelength)
|
||||
* self.universal_propagator)
|
||||
|
||||
@@ -12,6 +12,8 @@ from functools import reduce
|
||||
|
||||
__all__ = ['Multislice2DPtycho']
|
||||
|
||||
|
||||
|
||||
class Multislice2DPtycho(CDIModel):
|
||||
|
||||
@property
|
||||
@@ -35,7 +37,6 @@ class Multislice2DPtycho(CDIModel):
|
||||
bandlimit=None,
|
||||
subpixel=True,
|
||||
exponentiate_obj=True,
|
||||
low_res_obj=False,
|
||||
fourier_probe=False,
|
||||
prevent_aliasing=True,
|
||||
phase_only=False,
|
||||
@@ -67,7 +68,6 @@ 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
|
||||
@@ -84,12 +84,11 @@ class Multislice2DPtycho(CDIModel):
|
||||
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.probe = t.nn.Parameter(probe_guess.to(t.complex64)
|
||||
# / self.probe_norm)
|
||||
@@ -142,11 +141,16 @@ class Multislice2DPtycho(CDIModel):
|
||||
self.bandlimit = bandlimit
|
||||
|
||||
|
||||
self.as_prop = tools.propagators.generate_angular_spectrum_propagator(shape, spacing, self.wavelength, self.dz, bandlimit=self.bandlimit)
|
||||
self.as_prop = tools.propagators.generate_angular_spectrum_propagator(shape, spacing, self.wavelength, self.dz, bandlimit=1/np.sqrt(2))#self.bandlimit)
|
||||
#plt.imshow(t.abs(self.as_prop))
|
||||
#plt.figure()
|
||||
#plt.imshow(t.abs(t.fft.fftshift(t.fft.ifft2(self.as_prop))))
|
||||
#plt.show()
|
||||
#exit()
|
||||
|
||||
|
||||
@classmethod
|
||||
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):
|
||||
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):
|
||||
|
||||
wavelength = dataset.wavelength
|
||||
det_basis = dataset.detector_geometry['basis']
|
||||
@@ -198,23 +202,17 @@ 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=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)
|
||||
else:
|
||||
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.SHARP_style_probe(dataset, probe_shape, det_slice, 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))
|
||||
@@ -296,16 +294,13 @@ class Multislice2DPtycho(CDIModel):
|
||||
exponentiate_obj=exponentiate_obj,
|
||||
units=units, fourier_probe=fourier_probe,
|
||||
phase_only=phase_only,
|
||||
prevent_aliasing=prevent_aliasing,
|
||||
low_res_obj=low_res_obj)
|
||||
prevent_aliasing=prevent_aliasing)
|
||||
|
||||
|
||||
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:
|
||||
@@ -362,13 +357,10 @@ 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.low_res_obj)
|
||||
|
||||
multiple_modes=True,upsample_obj=self.prevent_aliasing)
|
||||
|
||||
|
||||
elif self.obj.dim() == 3:
|
||||
# If separate slices
|
||||
@@ -379,7 +371,7 @@ class Multislice2DPtycho(CDIModel):
|
||||
else:
|
||||
exit_waves = tools.interactions.ptycho_2D_round(
|
||||
exit_waves, obj[i], pix_trans,
|
||||
multiple_modes=True)
|
||||
multiple_modes=True, upsample_obj=self.prevent_aliasing)
|
||||
|
||||
#if self.iteration_count >= 1:
|
||||
# plt.imshow(t.abs(tools.propagators.far_field(
|
||||
|
||||
+29
-21
@@ -4,7 +4,6 @@ import torch as t
|
||||
from CDTools.models import CDIModel
|
||||
from CDTools.datasets import Ptycho2DDataset
|
||||
from CDTools import tools
|
||||
from CDTools.tools import cmath
|
||||
from CDTools.tools import plotting as p
|
||||
from CDTools.tools.interactions import RPI_interaction
|
||||
from CDTools.tools import initializers
|
||||
@@ -51,6 +50,10 @@ __all__ = ['RPI']
|
||||
__all__ = ['RPI']
|
||||
|
||||
class RPI(CDIModel):
|
||||
|
||||
@property
|
||||
def obj(self):
|
||||
return t.complex(self.obj_real, self.obj_imag)
|
||||
|
||||
def __init__(self, wavelength, detector_geometry, probe_basis,
|
||||
probe, obj_guess, detector_slice=None,
|
||||
@@ -73,8 +76,8 @@ class RPI(CDIModel):
|
||||
|
||||
self.probe_basis = t.Tensor(probe_basis)
|
||||
|
||||
scale_factor = t.Tensor([probe.shape[-2]/obj_guess.shape[-2],
|
||||
probe.shape[-3]/obj_guess.shape[-3]])
|
||||
scale_factor = t.Tensor([probe.shape[-1]/obj_guess.shape[-1],
|
||||
probe.shape[-2]/obj_guess.shape[-2]])
|
||||
self.obj_basis = self.probe_basis / scale_factor
|
||||
self.detector_slice = detector_slice
|
||||
|
||||
@@ -89,12 +92,17 @@ class RPI(CDIModel):
|
||||
self.mask = t.BoolTensor(mask)
|
||||
|
||||
|
||||
self.probe = probe.to(t.float32)
|
||||
self.probe = probe.to(t.complex64)
|
||||
|
||||
if obj_guess.dim() == 3:
|
||||
obj_guess = obj_guess[None,:,:,:]
|
||||
if obj_guess.dim() == 2:
|
||||
obj_guess = obj_guess[None,:,:]
|
||||
|
||||
self.obj = t.nn.Parameter(obj_guess.to(t.float32))
|
||||
|
||||
self.obj_real = t.nn.Parameter(obj_guess.real.to(t.float32))
|
||||
self.obj_imag = t.nn.Parameter(obj_guess.imag.to(t.float32))
|
||||
|
||||
# Wait for LBFGS to be updated for complex-valued parameters
|
||||
#self.obj = t.nn.Parameter(obj_guess.to(t.float32))
|
||||
|
||||
if background is None:
|
||||
if detector_slice is not None:
|
||||
@@ -146,7 +154,7 @@ class RPI(CDIModel):
|
||||
oversampling=oversampling)
|
||||
|
||||
if not isinstance(probe,t.Tensor):
|
||||
probe = cmath.complex_to_torch(probe)
|
||||
probe = t.as_tensor(probe)
|
||||
|
||||
# Potentially need all of this orientation stuff later
|
||||
|
||||
@@ -190,8 +198,8 @@ class RPI(CDIModel):
|
||||
# I think something to do with the fact that the object is defined
|
||||
# on a coarser grid needs to be accounted for here that is not
|
||||
# accounted for yet
|
||||
scale = t.sum(patterns[0]) / t.sum(cmath.cabssq(probe))
|
||||
obj_guess = scale * cmath.expi(2 * np.pi * t.rand([n_modes,]+obj_size))
|
||||
scale = t.sum(patterns[0]) / t.sum(t.abs(probe)**2)
|
||||
obj_guess = scale * t.exp(2j * np.pi * t.rand([n_modes,]+obj_size))
|
||||
elif initialization.lower().strip() == 'spectral':
|
||||
if background is not None:
|
||||
obj_guess = initializers.RPI_spectral_init(
|
||||
@@ -218,9 +226,9 @@ class RPI(CDIModel):
|
||||
|
||||
|
||||
def random_init(self, pattern):
|
||||
scale = t.sum(pattern) / t.sum(cmath.cabssq(self.probe))
|
||||
self.obj.data = scale * cmath.expi(
|
||||
2 * np.pi * t.rand(self.obj.shape[:-1])).to(
|
||||
scale = t.sum(pattern) / t.sum(t.abs(self.probe)**2)
|
||||
self.obj.data = scale * t.exp(
|
||||
2j * np.pi * t.rand(self.obj.shape)).to(
|
||||
dtype=self.obj.dtype, device=self.obj.device)
|
||||
|
||||
def spectral_init(self, pattern):
|
||||
@@ -249,8 +257,8 @@ class RPI(CDIModel):
|
||||
pr = self.probe[i]
|
||||
# Here we have a 3D probe (one single mode)
|
||||
# and a 4D object (multiple modes mixing incoherently)
|
||||
exit_waves = RPI_interaction(pr[:,:,:],
|
||||
self.obj_support[None,:,:] * self.obj)
|
||||
exit_waves = RPI_interaction(pr,
|
||||
self.obj_support * self.obj)
|
||||
|
||||
all_exit_waves.append(exit_waves)
|
||||
|
||||
@@ -297,8 +305,8 @@ class RPI(CDIModel):
|
||||
#return tools.losses.poisson_nll(real_data, sim_data, mask=mask)
|
||||
|
||||
def regularizer(self, factors):
|
||||
return factors[0] * t.sum(cmath.cabssq(self.obj[0,:,:,:])) \
|
||||
+ factors[1] * t.sum(cmath.cabssq(self.obj[1:,:,:,:]))
|
||||
return factors[0] * t.sum(t.abs(self.obj[0,:,:])**2) \
|
||||
+ factors[1] * t.sum(t.abs(self.obj[1:,:,:])**2)
|
||||
|
||||
def to(self, *args, **kwargs):
|
||||
super(RPI, self).to(*args, **kwargs)
|
||||
@@ -330,7 +338,7 @@ class RPI(CDIModel):
|
||||
plot_list = [
|
||||
('Root Sum Squared Amplitude of all Probes',
|
||||
lambda self, fig: p.plot_amplitude(
|
||||
np.sqrt(np.sum(cmath.cabssq(self.probe).cpu().numpy(),axis=0)),
|
||||
np.sqrt(np.sum((t.abs(self.probe)**2).cpu().numpy(),axis=0)),
|
||||
fig=fig, basis=self.probe_basis)),
|
||||
('Dominant Object Amplitude',
|
||||
lambda self, fig: p.plot_amplitude(self.obj[0], fig=fig,
|
||||
@@ -355,13 +363,13 @@ class RPI(CDIModel):
|
||||
# continues to use that standard pattern
|
||||
probe_basis = self.probe_basis.detach().cpu().numpy()
|
||||
obj_basis = self.obj_basis.detach().cpu().numpy()
|
||||
probe = cmath.torch_to_complex(self.probe.detach().cpu())
|
||||
probe = self.probe.detach().cpu().numpy()
|
||||
# Provide the option to save out the subdominant objects or
|
||||
# just the dominant one
|
||||
if full_obj:
|
||||
obj = cmath.torch_to_complex(self.obj.detach().cpu())
|
||||
obj = self.obj.detach().cpu().numpy()
|
||||
else:
|
||||
obj = cmath.torch_to_complex(self.obj[0].detach().cpu())
|
||||
obj = self.obj[0].detach().cpu().numpy()
|
||||
background = self.background.detach().cpu().numpy()**2
|
||||
|
||||
return {'probe_basis': probe_basis, 'obj_basis': obj_basis,
|
||||
|
||||
@@ -297,8 +297,7 @@ def synthesize_reconstructions(probes, objects, use_probe=False, obj_slice=None,
|
||||
|
||||
obj = ip.sinc_subpixel_shift(obj,np.array(shift))
|
||||
|
||||
|
||||
if len(probe.shape) == 4:
|
||||
if len(probe.shape) == 3:
|
||||
probe = t.stack([ip.sinc_subpixel_shift(p,tuple(shift))
|
||||
for p in probe],dim=0)
|
||||
else:
|
||||
|
||||
@@ -178,7 +178,6 @@ def find_subpixel_shift(im1, im2, search_around=(0,0), resolution=10):
|
||||
|
||||
|
||||
# And we extract the shift from the window
|
||||
print(upsampled.shape)
|
||||
sh = t.as_tensor(upsampled.shape, device=upsampled.device)
|
||||
cormax = t.as_tensor([t.div(t.argmax(upsampled), sh[1],
|
||||
rounding_mode='floor'),
|
||||
@@ -186,9 +185,6 @@ def find_subpixel_shift(im1, im2, search_around=(0,0), resolution=10):
|
||||
device=upsampled.device)
|
||||
|
||||
sh_over_2 = t.div(sh,2,rounding_mode='floor')
|
||||
print(cormax.shape)
|
||||
print(sh.shape)
|
||||
print(sh_over_2.shape)
|
||||
subpixel_shift = ((cormax + sh_over_2) % sh - sh_over_2).to(dtype=upsampled.dtype)
|
||||
|
||||
return search_around.to(device=upsampled.device, dtype=upsampled.dtype) + \
|
||||
|
||||
@@ -15,7 +15,6 @@ __all__ = ['translations_to_pixel', 'pixel_to_translations',
|
||||
'RPI_interaction']
|
||||
|
||||
|
||||
|
||||
def translations_to_pixel(basis, translations, surface_normal=t.Tensor([0.,0.,1.])):
|
||||
"""Takes real space translations and outputs them in pixel space
|
||||
|
||||
@@ -448,7 +447,6 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi
|
||||
I,J = t.meshgrid(i,j)
|
||||
I = 2 * np.pi * I / probe.shape[-2]
|
||||
J = 2 * np.pi * J / probe.shape[-1]
|
||||
|
||||
phase_masks = t.exp(1j*(-subpixel_translations[:,0,None,None]*I
|
||||
-subpixel_translations[:,1,None,None]*J))
|
||||
|
||||
@@ -461,6 +459,7 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multi
|
||||
|
||||
shifted_probe = t.fft.ifft2(t.fft.ifftshift(shifted_fft_probe,
|
||||
dim=(-1,-2)))
|
||||
|
||||
if multiple_modes: # Multi-mode probe
|
||||
output = shifted_probe * selections[...,None,:,:]
|
||||
else:
|
||||
|
||||
@@ -164,8 +164,8 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
|
||||
basis_norm = np.linalg.norm(np_basis, axis = 0)
|
||||
basis_norm = basis_norm * get_units_factor(units)
|
||||
|
||||
extent = [0, to_plot.shape[-1]*basis_norm[1], 0,
|
||||
to_plot.shape[-2]*basis_norm[0]]
|
||||
extent = [0, im.shape[-1]*basis_norm[1], 0,
|
||||
im.shape[-2]*basis_norm[0]]
|
||||
else:
|
||||
extent=None
|
||||
|
||||
|
||||
@@ -11,7 +11,7 @@ By necessity, CDTools operates using a mixture of numpy arrays and pytorch tenso
|
||||
|
||||
Functions that can accept either pytorch tensors or numpy arrays will have the type for the relevant inputs documented as "array", rather than "np.ndarray" or "torch.Tensor". In general, these functions will either accept all "array" type inputs as numpy arrays, or all as torch tensors. They will then return a result in a format matching that of the inputs. While many of these functions will work with mixed numpy/pytorch input, the output behavior of these functions is not, in general, defined for such a case, so it is heavily discouraged.
|
||||
|
||||
Finally, it is important to remember always that, since pytorch does not have complex number support, complex-valued tensors in pytorch are always represented by a tensor with a trailing dimension of length 2. All complex arithmetic operations are defined in :code:`tools.cmath`, but do not forget to use them!
|
||||
Please note that adherance to the conventions above is a work in progress, and you may find improperly documented or implemented functions. Please file a bug report if you do!
|
||||
|
||||
|
||||
Unit Conventions
|
||||
|
||||
@@ -1,5 +1,3 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
|
||||
import CDTools
|
||||
from matplotlib import pyplot as plt
|
||||
|
||||
@@ -23,7 +21,7 @@ dataset.mask = dataset.mask[70:-70,70:-70]
|
||||
model = CDTools.models.FancyPtycho.from_dataset(dataset,
|
||||
translation_scale = 4,
|
||||
n_modes=2,
|
||||
propagation_distance=-73e-6)
|
||||
propagation_distance=73e-6)
|
||||
|
||||
# Move to the GPU
|
||||
model.to(device='cuda')
|
||||
@@ -32,20 +30,20 @@ dataset.get_as(device='cuda')
|
||||
|
||||
# We turn off position reconstruction for the first phase
|
||||
model.translation_offsets.requires_grad = False
|
||||
for i, loss in enumerate(model.Adam_optimize(10, dataset, batch_size=15)):
|
||||
for loss in model.Adam_optimize(10, dataset, batch_size=15):
|
||||
model.inspect(dataset)
|
||||
print(i,loss)
|
||||
print(model.report())
|
||||
|
||||
# And we turn it on for the second phase
|
||||
model.translation_offsets.requires_grad = True
|
||||
for i, loss in enumerate(model.Adam_optimize(20, dataset, batch_size=15)):
|
||||
for loss in model.Adam_optimize(20, dataset, batch_size=15):
|
||||
model.inspect(dataset)
|
||||
print(i,loss)
|
||||
print(model.report())
|
||||
|
||||
# The third phase lowers the rate further
|
||||
for i, loss in enumerate(model.Adam_optimize(10, dataset, batch_size=15, lr=0.0005)):
|
||||
for loss in model.Adam_optimize(10, dataset, batch_size=15, lr=0.0005):
|
||||
model.inspect(dataset)
|
||||
print(i,loss)
|
||||
print(model.report())
|
||||
|
||||
|
||||
model.inspect(dataset)
|
||||
|
||||
@@ -1,17 +0,0 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
import CDTools
|
||||
from matplotlib import pyplot as plt
|
||||
# First, we load an example dataset from a .cxi file
|
||||
filename = '/media/Data Bank/APS_HXN_07_19/CXIs/scan_94513_cxi.h5'
|
||||
|
||||
#filename = 'data/scan_94361_cxi.h5'
|
||||
dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename)
|
||||
# Next, we create a ptychography model from the dataset
|
||||
model = CDTools.models.Simple3DBPP.from_dataset(dataset)
|
||||
#import pdb; pdb.set_trace()
|
||||
model.inspect(dataset)
|
||||
plt.show()
|
||||
|
||||
# # Now, we run a short reconstruction from the dataset!
|
||||
# for i, loss in enumerate(model.Adam_optimize(1, dataset)):
|
||||
# print(i, loss)
|
||||
@@ -1,21 +1,12 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
|
||||
from matplotlib import pyplot as plt
|
||||
import pickle
|
||||
from scipy import io
|
||||
|
||||
from CDTools.tools import plotting
|
||||
from CDTools.tools import analysis
|
||||
|
||||
|
||||
|
||||
with open('example_reconstructions/gold_balls_ensemble.pickle', 'rb') as f:
|
||||
dataset = pickle.load(f)
|
||||
|
||||
# This converts from a list of dictionaries to a dictionary of lists
|
||||
# It's safe to assume that all elements have the same set of keys
|
||||
if type(dataset) == type([]):
|
||||
dataset = {key: [element[key] for element in dataset]
|
||||
for key in dataset[0]}
|
||||
dataset = io.loadmat('example_reconstructions/gold_balls_ensemble.mat')
|
||||
|
||||
|
||||
# Now we synthesize an average reconstruction
|
||||
@@ -25,7 +16,6 @@ synth_probe, synth_obj, aligned_objs = analysis.synthesize_reconstructions(
|
||||
# And then we calculate the consistency PRTF from this
|
||||
freqs, prtf = analysis.calc_consistency_prtf(synth_obj, aligned_objs, dataset['basis'][0])
|
||||
|
||||
|
||||
# Plot the first mode in detail
|
||||
plotting.plot_phase(synth_probe[0],basis=dataset['basis'][0])
|
||||
plotting.plot_amplitude(synth_probe[0],basis=dataset['basis'][0])
|
||||
|
||||
@@ -1,16 +1,16 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
|
||||
import CDTools
|
||||
import pickle
|
||||
import numpy as np
|
||||
from scipy import io
|
||||
|
||||
# Load the data
|
||||
filename = 'example_data/AuBalls_700ms_30nmStep_3_6SS_filter.cxi'
|
||||
dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename)
|
||||
|
||||
results = []
|
||||
n = 25
|
||||
|
||||
for idx in range(25):
|
||||
print('Starting Reconstruction', idx)
|
||||
for idx in range(n):
|
||||
print('Starting Reconstruction', idx+1, 'of',n)
|
||||
|
||||
# Create a new model each time
|
||||
model = CDTools.models.FancyPtycho.from_dataset(dataset,n_modes=3,
|
||||
@@ -21,14 +21,21 @@ for idx in range(25):
|
||||
dataset.get_as(device='cuda')
|
||||
|
||||
# Run the reconstruction
|
||||
for i, loss in enumerate(model.Adam_optimize(30, dataset, batch_size=100)):
|
||||
print(i,loss)
|
||||
for loss in model.Adam_optimize(30, dataset, batch_size=100):
|
||||
print(model.report(), end='\r')
|
||||
|
||||
# And add the results to the ensemble
|
||||
# Print a summary that won't be overwritten by the next line
|
||||
print('Finished:',model.report())
|
||||
|
||||
# And add the results to the ensemble dictionary
|
||||
results.append(model.save_results(dataset))
|
||||
|
||||
|
||||
# This converts from a list of dictionaries to a dictionary of lists
|
||||
# It's safe to assume that all elements have the same set of keys
|
||||
# After this, e.g. dataset['probe'] will return a list of all probes.
|
||||
results = {key: np.array([element[key] for element in results])
|
||||
for key in results[0].keys()}
|
||||
|
||||
print(results['probe'].shape)
|
||||
# Save out the ensemble
|
||||
with open('example_reconstructions/gold_balls_ensemble.pickle', 'wb') as f:
|
||||
pickle.dump(results,f)
|
||||
|
||||
io.savemat('example_reconstructions/gold_balls_ensemble.mat', results)
|
||||
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -1,5 +1,3 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
|
||||
import CDTools
|
||||
from matplotlib import pyplot as plt
|
||||
|
||||
|
||||
@@ -1,36 +1,32 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
|
||||
import CDTools
|
||||
from matplotlib import pyplot as plt
|
||||
import pickle
|
||||
|
||||
from scipy import io
|
||||
|
||||
#
|
||||
# This is the testing file for the Bragg ptychography code that includes
|
||||
# corrections for probe propagation and for off-axis far field diffraction
|
||||
# I'm starting to build it up using the lab ptychography forward data so I
|
||||
# can initially test that it works on "softball" data before moving to the
|
||||
# actual Bragg geometry data
|
||||
# This is the testing dataset for the Bragg ptychography code, that was
|
||||
# collected with our optical setup using a low NA zone plate illuminating
|
||||
# a reflective diffraction grating with writing on it.
|
||||
#
|
||||
|
||||
#filename = 'example_data/lab_ptycho_data.cxi'
|
||||
# This file is too large to be distributed via Github.
|
||||
# Please contact Abe Levitan (alevitan@mit) if you would like access
|
||||
filename = '/media/Data Bank/Lab Ptycho/Zone Plate Bragg 633.cxi'
|
||||
dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename)
|
||||
|
||||
#dataset.inspect()
|
||||
|
||||
model = CDTools.models.Bragg2DPtycho.from_dataset(dataset,probe_support_radius=60,correct_tilt=False)#propagate_probe=False)#, n_modes=8)
|
||||
model = CDTools.models.Bragg2DPtycho.from_dataset(dataset,probe_support_radius=60,correct_tilt=False)
|
||||
model.to(device='cuda')
|
||||
dataset.get_as(device='cuda')
|
||||
|
||||
|
||||
model.translation_offsets.requires_grad = False
|
||||
for i, loss in enumerate(model.Adam_optimize(100, dataset)):
|
||||
for loss in model.Adam_optimize(100, dataset):
|
||||
model.inspect(dataset)
|
||||
print(i,loss)
|
||||
|
||||
#with open('example_reconstructions/lab_bragg_2d_ptycho.pickle', 'wb') as f:
|
||||
# pickle.dump(model.save_results(dataset),f)
|
||||
print(model.report())
|
||||
|
||||
#model.compare(dataset)
|
||||
io.savemat('example_reconstructions/lab_bragg_2d_ptycho.mat',
|
||||
model.save_results(dataset))
|
||||
|
||||
model.compare(dataset)
|
||||
plt.show()
|
||||
|
||||
@@ -1,9 +1,9 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
|
||||
import CDTools
|
||||
from matplotlib import pyplot as plt
|
||||
import pickle
|
||||
from scipy import io
|
||||
|
||||
|
||||
# A simple dataset collected from our optical setup
|
||||
filename = 'example_data/lab_ptycho_data.cxi'
|
||||
dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename)
|
||||
|
||||
@@ -19,14 +19,14 @@ dataset.get_as(device='cuda')
|
||||
|
||||
|
||||
model.translation_offsets.requires_grad = False
|
||||
for i, loss in enumerate(model.Adam_optimize(50, dataset)):
|
||||
for loss in model.Adam_optimize(50, dataset):
|
||||
model.inspect(dataset)
|
||||
print(i,loss)
|
||||
print(model.report())
|
||||
|
||||
model.tidy_probes()
|
||||
exit()
|
||||
with open('example_reconstructions/lab_ptycho.pickle', 'wb') as f:
|
||||
pickle.dump(model.save_results(dataset),f)
|
||||
|
||||
io.savemat('example_reconstructions/lab_ptycho.mat',
|
||||
model.save_results(dataset))
|
||||
|
||||
model.compare(dataset)
|
||||
plt.show()
|
||||
|
||||
@@ -1,4 +1,3 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
import numpy as np
|
||||
import torch as t
|
||||
from CDTools.models import CDIModel
|
||||
|
||||
@@ -1,5 +1,3 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
|
||||
import CDTools
|
||||
from matplotlib import pyplot as plt
|
||||
import numpy as np
|
||||
|
||||
@@ -1,15 +1,10 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
|
||||
import CDTools
|
||||
from matplotlib import pyplot as plt
|
||||
import pickle
|
||||
from torch.utils.data import Subset
|
||||
from matplotlib import pyplot as plt
|
||||
|
||||
# First, we load an example dataset from a .cxi file
|
||||
|
||||
ss_filename = 'example_data/Optical_Data_ss.cxi'
|
||||
|
||||
#with open('example_data/Optical_ptycho.pickle', 'rb') as f:
|
||||
with open('example_data/Optical_ptycho_incoherent.pickle', 'rb') as f:
|
||||
ptycho_results = pickle.load(f)
|
||||
|
||||
@@ -19,9 +14,9 @@ background = ptycho_results['background']
|
||||
dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(ss_filename)
|
||||
|
||||
|
||||
# Next, we create a ptychography model from the dataset
|
||||
# Next, we create an RPI model from the dataset
|
||||
# Note that we explicitly as for two incoherent probe modes
|
||||
model = CDTools.models.RPI.from_dataset(dataset, probe, [900,900],
|
||||
model = CDTools.models.RPI.from_dataset(dataset, probe, [500,500],
|
||||
background=background, n_modes=2,
|
||||
initialization='random')
|
||||
|
||||
@@ -33,20 +28,19 @@ dataset.get_as(device='cuda')
|
||||
# Note that the inspect step takes the vast majority of the time
|
||||
# The regularization is an L2 regularizer that empirically helps accelerate
|
||||
# convergence
|
||||
for i, loss in enumerate(model.LBFGS_optimize(30, dataset, lr=0.4, regularization_factor=[0.05,0.05])):#0.1)):
|
||||
for loss in model.LBFGS_optimize(30, dataset, lr=0.4, regularization_factor=[0.05,0.05]):
|
||||
model.inspect(dataset)
|
||||
print(i,loss)
|
||||
|
||||
#model.inspect(dataset)
|
||||
print(model.report())
|
||||
|
||||
|
||||
# Now we use the regularizer to damp all but the top modes
|
||||
for i, loss in enumerate(model.LBFGS_optimize(50, dataset, lr=0.4, regularization_factor=[0.001,0.1])):
|
||||
for loss in model.LBFGS_optimize(50, dataset, lr=0.4, regularization_factor=[0.001,0.1]):
|
||||
#model.inspect(dataset)
|
||||
print(i,loss)
|
||||
print(model.report())
|
||||
|
||||
# Save results to a python dictionary
|
||||
results = model.save_results()
|
||||
|
||||
|
||||
# Finally, we plot the results
|
||||
model.inspect(dataset)
|
||||
model.compare(dataset)
|
||||
|
||||
@@ -1,11 +1,7 @@
|
||||
from __future__ import division, print_function, absolute_import
|
||||
|
||||
import CDTools
|
||||
from CDTools.tools import plotting as p
|
||||
from matplotlib import pyplot as plt
|
||||
import pickle
|
||||
import torch as t
|
||||
from CDTools.tools import cmath
|
||||
from scipy import io
|
||||
|
||||
|
||||
filename = 'example_data/lab_ptycho_data.cxi'
|
||||
dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename)
|
||||
@@ -21,21 +17,21 @@ model.to(device='cuda')
|
||||
dataset.get_as(device='cuda')
|
||||
|
||||
model.translation_offsets.requires_grad = False
|
||||
for i, loss in enumerate(model.Adam_optimize(200, dataset)):
|
||||
for loss in model.Adam_optimize(200, dataset):
|
||||
model.inspect(dataset)
|
||||
print(i,loss)
|
||||
print(model.report())
|
||||
|
||||
model.tidy_probes()
|
||||
|
||||
for i, loss in enumerate(model.Adam_optimize(200, dataset, lr=0.0001)):
|
||||
for loss in model.Adam_optimize(200, dataset, lr=0.0001):
|
||||
model.inspect(dataset)
|
||||
print(i,loss)
|
||||
print(model.report())
|
||||
|
||||
model.tidy_probes(normalize=True)
|
||||
model.inspect(dataset)
|
||||
|
||||
with open('example_reconstructions/unified_modes.pickle', 'wb') as f:
|
||||
pickle.dump(model.save_results(dataset),f)
|
||||
io.savemat('example_reconstructions/unified_modes.mat',
|
||||
model.save_results(dataset))
|
||||
|
||||
model.compare(dataset)
|
||||
plt.show()
|
||||
|
||||
Reference in New Issue
Block a user