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

Start to make the lab bragg ptycho code sane, and fix the bug with probe bases when using probe fourier crop

Browse Source
This commit is contained in:
Abe Levitan
2024-02-11 17:32:41 -05:00
parent 44a6fac37a
commit dd30f58d2d
11 changed files with 348 additions and 242 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+1 -1
View File
@@ -7,4 +7,4 @@
docs/build
build/*
dist
example_data/*
*/example_data/*
examples/example_reconstructions/gold_balls.h5
BIN
View File
Binary file not shown.
examples/example_reconstructions/gold_balls.mat
BIN
View File
Binary file not shown.
examples/gold_ball_ptycho.py
+7 -3
View File
@@ -1,6 +1,7 @@
import cdtools
from matplotlib import pyplot as plt
from scipy import io
import torch as t
# First, we load an example dataset from a .cxi file
filename = 'example_data/AuBalls_700ms_30nmStep_3_6SS_filter.cxi'
@@ -8,15 +9,18 @@ dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(filename)
# Next, we create a ptychography model from the dataset
# Note that we explicitly ask for two incoherent probe modes
model = cdtools.models.FancyPtycho.from_dataset(dataset, n_modes=2)
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')
with model.save_on_exit('example_reconstructions/gold_balls.mat', dataset):
with model.save_on_exit('example_reconstructions/gold_balls.h5', dataset):
# Now, we run a short reconstruction from the dataset
for loss in model.Adam_optimize(10, dataset, batch_size=50):
for loss in model.Adam_optimize(1, dataset, batch_size=50):
# And we liveplot the updates to the model as they happen
print(model.report())
model.inspect(dataset)
examples/lab_bragg_2d_ptycho.py
+11 -6
View File
@@ -10,14 +10,19 @@ from scipy import io
# This file is too large to be distributed via Github.
# Please contact Abe Levitan (alevitan@mit) if you would like access
filename = 'Zone Plate Bragg 633.cxi'
filename = 'example_data/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)
model.to(device='cuda')
dataset.get_as(device='cuda')
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
@@ -25,8 +30,8 @@ for loss in model.Adam_optimize(100, dataset):
model.inspect(dataset)
print(model.report())
io.savemat('example_reconstructions/lab_bragg_2d_ptycho.mat',
model.save_results(dataset))
#io.savemat('example_reconstructions/lab_bragg_2d_ptycho.mat',
# model.save_results(dataset))
model.compare(dataset)
plt.show()
src/cdtools/models/base.py
+24 -10
View File
@@ -40,6 +40,7 @@ import time
from scipy import io
from contextlib import contextmanager
from .complex_lbfgs import MyLBFGS
from cdtools.tools.data import nested_dict_to_h5
__all__ = ['CDIModel']
@@ -55,9 +56,10 @@ class CDIModel(t.nn.Module):
"""
def __init__(self):
super(CDIModel,self).__init__()
super(CDIModel, self).__init__()
self.loss_train = []
self.loss_history = []
self.training_history = ''
self.iteration_count = 0
def from_dataset(self, dataset):
@@ -172,13 +174,16 @@ class CDIModel(t.nn.Module):
A dictionary containing all the parameters and buffers of the model, i.e. the result of self.state_dict(), converted to numpy.
"""
state_dict = {k: v.cpu().numpy() for k, v in self.state_dict().items()}
return {
'state_dict': state_dict,
'loss_train': np.array(self.loss_train),
'loss_history': np.array(self.loss_history),
'training_history': self.training_history,
'loss_function': self.loss.__name__,
}
def save_to_mat(self, filename, *args):
def save_to_h5(self, filename, *args):
"""Saves the results to a .mat file
Parameters
@@ -188,7 +193,7 @@ class CDIModel(t.nn.Module):
*args
Accepts any additional args that model.save_results needs, for this model
"""
return io.savemat(filename, self.save_results(*args))
return nested_dict_to_h5(filename, self.save_results(*args))
@contextmanager
def save_on_exit(self, filename, *args, exception_filename=None):
@@ -209,11 +214,11 @@ class CDIModel(t.nn.Module):
"""
try:
yield
self.save_to_mat(filename, *args)
self.save_to_h5(filename, *args)
except Exception as e:
if exception_filename is None:
exception_filename = filename
self.save_to_mat(exception_filename, *args)
self.save_to_h5(exception_filename, *args)
raise e
@@ -322,8 +327,10 @@ class CDIModel(t.nn.Module):
if scheduler is not None:
scheduler.step(loss)
self.loss_train.append(loss)
self.loss_history.append(loss)
epoch_idx = len(self.loss_history)
self.latest_iteration_time = time.time() - t0
self.training_history += self.report() + '\n'
return loss
# If we don't want to run in a different thread, this is easy
@@ -404,6 +411,13 @@ class CDIModel(t.nn.Module):
"""
self.training_history += (
f'Planning {iterations} epochs of Adam, with a learning rate = '
f'{lr}, batch size = {batch_size}, regularization_factor = '
f'{regularization_factor}, and schedule = {schedule}.\n'
)
if subset is not None:
# if subset is just one pattern, turn into a list for convenience
if type(subset) == type(1):
@@ -555,9 +569,9 @@ class CDIModel(t.nn.Module):
A string with basic info on the latest iteration
"""
if hasattr(self, 'latest_iteration_time'):
return 'Iteration ' + str(len(self.loss_train)) + \
return 'Epoch ' + str(len(self.loss_history)) + \
' completed in %0.2f s with loss ' %\
self.latest_iteration_time + str(self.loss_train[-1])
self.latest_iteration_time + str(self.loss_history[-1])
else:
return 'No reconstruction iterations performed yet!'
src/cdtools/models/bragg_2d_ptycho.py
+103 -72
View File
@@ -54,25 +54,28 @@ __all__ = ['Bragg2DPtycho']
class Bragg2DPtycho(CDIModel):
# Needed to do the real/complex split
#@property
#def obj(self):
# return t.complex(self.obj_real, self.obj_imag)
#@property
#def probe(self):
# return t.complex(self.probe_real, self.probe_imag)
def __init__(self, wavelength, detector_geometry,
probe_basis, probe_guess, obj_guess,
detector_slice=None,
min_translation=t.tensor([0, 0], dtype=t.float32),
median_propagation=t.tensor(0, dtype=t.float32),
background=None, translation_offsets=None, mask=None,
weights=None, translation_scale=1, saturation=None,
probe_support=None, oversampling=1,
propagate_probe=True, correct_tilt=True, lens=False, units='um'):
def __init__(
self,
wavelength,
detector_geometry,
obj_basis,
probe_guess,
obj_guess,
min_translation=t.tensor([0, 0],dtype=t.float32),
probe_basis=None,
median_propagation=t.tensor(0, dtype=t.float32),
background=None,
translation_offsets=None, mask=None,
weights=None,
translation_scale=1, saturation=None,
probe_support=None,
oversampling=1,
propagate_probe=True,
correct_tilt=True,
lens=False,
units='um',
dtype=t.float32,
):
# We need the detector geometry
# We need the probe basis (but in this case, we don't need the surface
@@ -97,13 +100,19 @@ class Bragg2DPtycho(CDIModel):
self.min_translation = t.tensor(min_translation)
self.median_propagation = t.tensor(median_propagation)
self.probe_basis = t.tensor(probe_basis)
self.detector_slice = copy(detector_slice)
self.register_buffer('obj_basis',
t.tensor(obj_basis, dtype=dtype))
if probe_basis is None:
self.register_buffer('probe_basis',
t.tensor(obj_basis, dtype=dtype))
else:
self.register_buffer('probe_basis',
t.tensor(probe_basis, dtype=dtype))
# calculate the surface normal from the probe basis
surface_normal = np.cross(np.array(probe_basis)[:,1],
np.array(probe_basis)[:,0])
surface_normal = np.cross(np.array(obj_basis)[:,1],
np.array(obj_basis)[:,0])
surface_normal /= np.linalg.norm(surface_normal)
self.surface_normal = t.tensor(surface_normal)
@@ -137,23 +146,13 @@ class Bragg2DPtycho(CDIModel):
else:
self.probe_support = t.ones(probe_guess[0].shape, dtype=t.bool)
#self.probe_real = t.nn.Parameter(probe_guess.real / self.probe_norm)
#self.probe_imag = t.nn.Parameter(probe_guess.imag / self.probe_norm)
#self.obj_real = t.nn.Parameter(obj_guess.real )
#self.obj_imag = t.nn.Parameter(obj_guess.imag)
self.probe = t.nn.Parameter(probe_guess / self.probe_norm)
self.obj = t.nn.Parameter(obj_guess)
if background is None:
if detector_slice is not None:
background = 1e-6 * t.ones(
self.probe[0][self.detector_slice].shape,
dtype=t.float32)
else:
background = 1e-6 * t.ones(self.probe[0].shape,
dtype=t.float32)
raise NotImplementedError('Issues with this due to probe fourier padding')
background = 1e-6 * t.ones(self.probe[0].shape,
dtype=t.float32)
self.background = t.nn.Parameter(background)
@@ -178,16 +177,12 @@ class Bragg2DPtycho(CDIModel):
self.propagate_probe = propagate_probe
self.correct_tilt = correct_tilt
if correct_tilt:
# recall that here we always want the shape of the detector
# before it's cut down by the detector slice to match the
# physical detector region
probe_shape = self.probe[0].shape
self.k_map, self.intensity_map = \
tools.propagators.generate_high_NA_k_intensity_map(
self.probe_basis,
self.obj_basis,
self.detector_geometry['basis'] / oversampling,
probe_shape,
self.background.shape,
self.detector_geometry['distance'],
self.wavelength,dtype=t.float32,
lens=lens)
@@ -200,7 +195,7 @@ class Bragg2DPtycho(CDIModel):
# This propagator should be able to be multiplied by the propagation
# distance each time to get a propagator
self.universal_propagator = t.angle(ggasp(
self.background.shape,#background in case of a fourier cropped probe
self.background.shape,
self.probe_basis, self.wavelength,
t.tensor([0, 0, self.wavelength/(2*np.pi)], dtype=t.float32),
propagation_vector=self.prop_dir,
@@ -226,7 +221,9 @@ class Bragg2DPtycho(CDIModel):
probe_fourier_crop=None,
propagate_probe=True,
correct_tilt=True,
lens=False):
lens=False,
obj_padding=200,
):
wavelength = dataset.wavelength
det_basis = dataset.detector_geometry['basis']
@@ -247,13 +244,12 @@ class Bragg2DPtycho(CDIModel):
# Then, generate the exit wave geometry from the dataset
ewg = tools.initializers.exit_wave_geometry
ew_basis, ew_shape, det_slice = ewg(det_basis,
det_shape,
wavelength,
distance,
center=center,
padding=padding,
oversampling=oversampling)
ew_basis = ewg(det_basis,
det_shape,
wavelength,
distance,
oversampling=oversampling)
# now we grab the sample surface normal
if hasattr(dataset, 'sample_info') and \
dataset.sample_info is not None and \
@@ -291,7 +287,7 @@ class Bragg2DPtycho(CDIModel):
# projection with a pseudoinverse and removing the last column
projector = np.linalg.pinv(mat)[:, :3]
probe_basis = t.Tensor(np.dot(projector, ew_basis))
obj_basis = t.Tensor(np.dot(projector, ew_basis))
# Now we need a much better way to handle the translations here
# than translations_to_pixel
@@ -299,10 +295,14 @@ class Bragg2DPtycho(CDIModel):
# Next generate the object geometry from the probe geometry and
# the translations
p2s = tools.interactions.project_translations_to_sample
pix_translations, propagations = p2s(probe_basis, translations)
pix_translations, propagations = p2s(obj_basis, translations)
obj_size, min_translation = tools.initializers.calc_object_setup(ew_shape, pix_translations, padding=200)
obj_size, min_translation = tools.initializers.calc_object_setup(
det_shape,
pix_translations,
padding=obj_padding,
)
median_propagation = t.median(propagations)
@@ -312,19 +312,25 @@ class Bragg2DPtycho(CDIModel):
# that space and use the standard initializations anyway
if probe_size is None:
probe = tools.initializers.SHARP_style_probe(dataset, ew_shape, det_slice, propagation_distance=propagation_distance, oversampling=oversampling)
probe = tools.initializers.SHARP_style_probe(
dataset,
propagation_distance=propagation_distance,
oversampling=oversampling
)
else:
probe = tools.initializers.gaussian_probe(dataset, ew_basis, ew_shape, probe_size, propagation_distance=propagation_distance)
probe = tools.initializers.gaussian_probe(
dataset,
ew_basis,
det_shape,
probe_size,
propagation_distance=propagation_distance
)
if hasattr(dataset, 'background') and dataset.background is not None:
background = t.sqrt(dataset.background)
elif det_slice is not None:
background = 1e-6 * t.ones(
probe[det_slice].shape,
dtype=t.float32)
else:
background = 1e-6 * t.ones(probe.shape[-2:],
background = 1e-6 * t.ones(det_shape,
dtype=t.float32)
if probe_fourier_crop is not None:
@@ -333,6 +339,12 @@ class Bragg2DPtycho(CDIModel):
probe_fourier_crop[0]:-probe_fourier_crop[0],
probe_fourier_crop[1]:-probe_fourier_crop[1]]
probe = tools.propagators.inverse_far_field(probe)
scale_factor = np.array(det_shape) / np.array(probe.shape)
probe_basis = obj_basis * scale_factor[None,:]
else:
probe_basis = obj_basis.clone()
# Now we initialize all the subdominant probe modes
probe_max = t.max(t.abs(probe))
@@ -359,10 +371,10 @@ class Bragg2DPtycho(CDIModel):
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
@@ -375,9 +387,9 @@ class Bragg2DPtycho(CDIModel):
raise NotImplementedError('No auto option implemented yet')
return cls(wavelength, det_geo, probe_basis, probe, obj,
detector_slice=det_slice,
return cls(wavelength, det_geo, obj_basis, probe, obj,
min_translation=min_translation,
probe_basis=probe_basis,
median_propagation =median_propagation,
translation_offsets = translation_offsets,
weights=weights, mask=mask, background=background,
@@ -406,6 +418,8 @@ class Bragg2DPtycho(CDIModel):
# Now we need to propagate each of the probes
# TODO: This fails if there is no background set. In that case,
# there's no way to know the "proper" size of the probe
# We automatically rescale the probe to match the background size,
# which allows us to do stuff like let the object be super-resolution,
# while restricting the probe to the detector resolution but still
@@ -458,7 +472,6 @@ class Bragg2DPtycho(CDIModel):
def measurement(self, wavefields):
return tools.measurements.quadratic_background(wavefields,
self.background,
detector_slice=self.detector_slice,
measurement=tools.measurements.incoherent_sum,
saturation=self.saturation,
oversampling=self.oversampling)
@@ -549,13 +562,31 @@ class Bragg2DPtycho(CDIModel):
('Basis Probe Fourier Space Phases',
lambda self, fig: p.plot_phase(tools.propagators.inverse_far_field(self.probe), fig=fig)),
('Basis Probe Real Space Amplitudes',
lambda self, fig: p.plot_amplitude(self.probe, fig=fig, basis=self.probe_basis, units=self.units)),
lambda self, fig: p.plot_amplitude(
self.probe,
fig=fig,
basis=self.probe_basis,
units=self.units
)),
('Basis Probe Real Space Phases',
lambda self, fig: p.plot_phase(self.probe, fig=fig, basis=self.probe_basis, units=self.units)),
lambda self, fig: p.plot_phase(
self.probe,
fig=fig,
basis=self.probe_basis,
units=self.units
)),
('Object Amplitude',
lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis)),
lambda self, fig: p.plot_amplitude(
self.obj,
fig=fig,
basis=self.obj_basis
)),
('Object Phase',
lambda self, fig: p.plot_phase(self.obj, fig=fig, basis=self.probe_basis)),
lambda self, fig: p.plot_phase(
self.obj,
fig=fig,
basis=self.obj_basis
)),
('Corrected Translations',
lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig)),
('Background',
src/cdtools/models/fancy_ptycho.py
+84 -51
View File
@@ -18,14 +18,16 @@ __all__ = ['FancyPtycho']
class FancyPtycho(CDIModel):
def __init__(self, wavelength, detector_geometry,
probe_basis,
def __init__(self,
wavelength,
detector_geometry,
obj_basis,
probe_guess,
obj_guess,
detector_slice=None,
surface_normal=t.tensor([0., 0., 1.], dtype=t.float32),
min_translation=t.tensor([0, 0], dtype=t.float32),
background=None,
probe_basis=None,
translation_offsets=None,
mask=None,
weights=None,
@@ -51,10 +53,15 @@ class FancyPtycho(CDIModel):
self.register_buffer('min_translation',
t.tensor(min_translation, dtype=dtype))
self.register_buffer('probe_basis',
t.tensor(probe_basis, dtype=dtype))
self.detector_slice = copy(detector_slice)
self.register_buffer('obj_basis',
t.tensor(obj_basis, dtype=dtype))
if probe_basis is None:
self.register_buffer('probe_basis',
t.tensor(obj_basis, dtype=dtype))
else:
self.register_buffer('probe_basis',
t.tensor(probe_basis, dtype=dtype))
self.register_buffer('surface_normal',
t.tensor(surface_normal, dtype=dtype))
if saturation is None:
@@ -62,11 +69,13 @@ class FancyPtycho(CDIModel):
else:
self.register_buffer('saturation',
t.tensor(saturation, dtype=dtype))
self.register_buffer('fourier_probe',
t.tensor(fourier_probe, dtype=bool))
# Not sure how to make this a buffer...
self.units = units
self.fourier_probe = fourier_probe
if mask is None:
self.mask = None
else:
@@ -92,13 +101,7 @@ class FancyPtycho(CDIModel):
obj_view_crop:-obj_view_crop]
if background is None:
if detector_slice is not None:
dummy_det = t.empty([s//oversampling
for s in self.probe[0].shape])
shape = dummy_det[self.detector_slice].shape
#shape = self.probe[0][self.detector_slice].shape
else:
shape = [s//oversampling for s in self.probe[0]]
shape = [s//oversampling for s in self.probe[0]]
background = 1e-6 * t.ones(shape, dtype=t.float32)
self.background = t.nn.Parameter(background)
@@ -132,9 +135,11 @@ class FancyPtycho(CDIModel):
self.register_buffer('probe_support',
t.tensor(probe_support, dtype=t.bool))
self.oversampling = oversampling
self.register_buffer('oversampling',
t.tensor(oversampling, dtype=int))
self.simulate_probe_translation = simulate_probe_translation
self.register_buffer('simulate_probe_translation',
t.tensor(simulate_probe_translation, dtype=bool))
if simulate_probe_translation:
Is = t.arange(self.probe.shape[-2], dtype=dtype)
@@ -147,7 +152,8 @@ class FancyPtycho(CDIModel):
self.register_buffer('J_phase', J_phase)
self.simulate_finite_pixels = simulate_finite_pixels
self.register_buffer('simulate_finite_pixels',
t.tensor(simulate_finite_pixels, dtype=bool))
# Here we set the appropriate loss function
if (loss.lower().strip() == 'amplitude mse'
@@ -165,7 +171,6 @@ class FancyPtycho(CDIModel):
dataset,
probe_size=None,
randomize_ang=0,
padding=0,
n_modes=1,
n_obj_modes=1,
dm_rank=None,
@@ -176,14 +181,13 @@ class FancyPtycho(CDIModel):
propagation_distance=None,
scattering_mode=None,
oversampling=1,
auto_center=False,
fourier_probe=False,
loss='amplitude mse',
units='um',
simulate_probe_translation=False,
simulate_finite_pixels=False,
obj_view_crop=None,
obj_padding=200
obj_padding=200,
):
wavelength = dataset.wavelength
@@ -201,21 +205,15 @@ class FancyPtycho(CDIModel):
dataset.get_as(*get_as_args[0], **get_as_args[1])
# Set to none to avoid issues with things outside the detector
if auto_center:
center = tools.image_processing.centroid(t.sum(patterns, dim=0))
else:
center = None
# Then, generate the probe geometry from the dataset
ewg = tools.initializers.exit_wave_geometry
probe_basis, probe_shape, det_slice = ewg(det_basis,
det_shape,
wavelength,
distance,
center=center,
padding=padding,
oversampling=oversampling)
obj_basis = ewg(
det_basis,
det_shape,
wavelength,
distance,
oversampling=oversampling,
)
if hasattr(dataset, 'sample_info') and \
dataset.sample_info is not None and \
@@ -237,15 +235,33 @@ class FancyPtycho(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)
pix_translations = tools.interactions.translations_to_pixel(
obj_basis,
translations,
surface_normal=surface_normal,
)
obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=obj_padding)
obj_size, min_translation = tools.initializers.calc_object_setup(
det_shape,
pix_translations,
padding=obj_padding,
)
# Finally, initialize the probe and object using this information
if probe_size is None:
probe = tools.initializers.SHARP_style_probe(dataset, probe_shape, det_slice, propagation_distance=propagation_distance, oversampling=oversampling)
probe = tools.initializers.SHARP_style_probe(
dataset,
propagation_distance=propagation_distance,
oversampling=oversampling,
)
else:
probe = tools.initializers.gaussian_probe(dataset, probe_basis, probe_shape, probe_size, propagation_distance=propagation_distance)
probe = tools.initializers.gaussian_probe(
dataset,
obj_basis,
probe_shape,
probe_size,
propagation_distance=propagation_distance,
)
if hasattr(dataset, 'background') and dataset.background is not None:
background = t.sqrt(dataset.background)
@@ -256,9 +272,16 @@ class FancyPtycho(CDIModel):
if probe_fourier_crop is not None:
probe = tools.propagators.far_field(probe)
probe = probe[probe_fourier_crop:-probe_fourier_crop,
probe_fourier_crop:-probe_fourier_crop]
probe = probe[probe_fourier_crop : probe.shape[-2]
- probe_fourier_crop,
probe_fourier_crop : probe.shape[-1]
- probe_fourier_crop]
probe = tools.propagators.inverse_far_field(probe)
scale_factor = np.array(det_shape) / np.array(probe.shape)
probe_basis = obj_basis * scale_factor[None,:]
else:
probe_basis = obj_basis.clone()
# Now we initialize all the subdominant probe modes
probe_max = t.max(t.abs(probe))
@@ -325,14 +348,14 @@ class FancyPtycho(CDIModel):
else:
probe_support = None
return cls(wavelength, det_geo, probe_basis, probe, obj,
detector_slice=det_slice,
return cls(wavelength, det_geo, obj_basis, probe, obj,
surface_normal=surface_normal,
min_translation=min_translation,
translation_offsets=translation_offsets,
weights=Ws, mask=mask, background=background,
translation_scale=translation_scale,
saturation=saturation,
probe_basis=probe_basis,
probe_support=probe_support,
fourier_probe=fourier_probe,
oversampling=oversampling,
@@ -350,7 +373,7 @@ class FancyPtycho(CDIModel):
# Step 1 is to convert the translations for each position into a
# value in pixels
pix_trans = tools.interactions.translations_to_pixel(
self.probe_basis,
self.obj_basis,
translations,
surface_normal=self.surface_normal)
pix_trans -= self.min_translation
@@ -433,7 +456,6 @@ class FancyPtycho(CDIModel):
return tools.measurements.quadratic_background(
wavefields,
self.background,
detector_slice=self.detector_slice,
measurement=tools.measurements.incoherent_sum,
saturation=self.saturation,
oversampling=self.oversampling,
@@ -501,7 +523,7 @@ class FancyPtycho(CDIModel):
if (hasattr(self, 'translation_offsets') and
self.translation_offsets is not None):
t_offset = tools.interactions.pixel_to_translations(
self.probe_basis,
self.obj_basis,
self.translation_offsets * self.translation_scale,
surface_normal=self.surface_normal)
return translations + t_offset
@@ -637,7 +659,16 @@ class FancyPtycho(CDIModel):
else:
cmap = 'twilight'
p.plot_nanomap_with_images(self.corrected_translations(dataset), get_probes, values=values, fig=fig, units=self.units, basis=self.probe_basis, nanomap_colorbar_title='Total Probe Intensity', cmap=cmap, **kwargs),
p.plot_nanomap_with_images(
self.corrected_translations(dataset),
get_probes,
values=values,
fig=fig,
units=self.units,
basis=self.obj_basis,
nanomap_colorbar_title='Total Probe Intensity',
cmap=cmap,
**kwargs),
plot_list = [
@@ -705,13 +736,13 @@ class FancyPtycho(CDIModel):
lambda self, fig: p.plot_amplitude(
self.obj[self.obj_view_slice],
fig=fig,
basis=self.probe_basis,
basis=self.obj_basis,
units=self.units)),
('Object Phase',
lambda self, fig: p.plot_phase(
self.obj[self.obj_view_slice],
fig=fig,
basis=self.probe_basis,
basis=self.obj_basis,
units=self.units)),
('Corrected Translations',
lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)),
@@ -727,7 +758,8 @@ class FancyPtycho(CDIModel):
base_results = super().save_results()
# We also save out the main results in a more readable format
basis = self.probe_basis.detach().cpu().numpy()
obj_basis = self.obj_basis.detach().cpu().numpy()
probe_basis = self.obj_basis.detach().cpu().numpy()
translations=self.corrected_translations(dataset).detach().cpu().numpy()
original_translations = dataset.translations.detach().cpu().numpy()
probe = self.probe.detach().cpu().numpy()
@@ -739,7 +771,8 @@ class FancyPtycho(CDIModel):
wavelength = self.wavelength.cpu().numpy()
results = {
'basis': basis,
'obj_basis': obj_basis,
'probe_basis': probe_basis,
'translations': translations,
'original_translations': original_translations,
'probe': probe,
src/cdtools/tools/data/data.py
+80 -1
View File
@@ -13,6 +13,9 @@ import datetime
import dateutil.parser
import torch as t
from contextlib import contextmanager
import numbers
from collections.abc import Mapping
import pathlib
__all__ = ['get_entry_info',
'get_sample_info',
@@ -32,7 +35,10 @@ __all__ = ['get_entry_info',
'add_dark',
'add_data',
'add_shot_to_shot_info',
'add_ptycho_translations']
'add_ptycho_translations',
'nested_dict_to_h5',
'h5_to_nested_dict',
]
#
@@ -785,3 +791,76 @@ def add_ptycho_translations(cxi_file, translations):
translations = -translations
add_shot_to_shot_info(cxi_file, translations, 'translation')
def nested_dict_to_h5(h5_file, d):
"""saves a nested dictionary to an h5 file object
Parameters
----------
h5_file : h5py.File
A file object, or path to a file, to write the dictionary to
d : dict
A mapping whose keys are all strings and whose values are only numpy arrays, pytorch tensors, scalars, or other mappings meeting the same conditions
"""
# If a bare string is passed
if isinstance(h5_file, str) or isinstance(h5_file, pathlib.Path):
with h5py.File(h5_file,'w') as f:
return nested_dict_to_h5(f, d)
for key in d.keys():
value = d[key]
if isinstance(value, numbers.Number):
arr = np.array(value)
h5_file.create_dataset(key, data=arr)
elif isinstance(value, np.ndarray):
h5_file.create_dataset(key, data=value)
elif t.is_tensor(value):
h5_file.create_dataset(key, data=value.detach().cpu().numpy())
elif isinstance(value, str):
h5_file.create_dataset(key, data=value, dtype=h5py.string_dtype())
elif isinstance(value, Mapping):
group = h5_file.create_group(key)
nested_dict_to_h5(group, value)
else:
raise ValueError(f'{value} is not a number, numpy array or mapping')
return
def h5_to_nested_dict(h5_file):
"""saves a nested dictionary to an h5 file object
Parameters
----------
h5_file : h5py.File
A file object, or path to a file, to load from
d : dict
A mapping whose keys are all strings and whose values are only numpy arrays, pytorch tensors, scalars, python strings, or other mappings meeting the same conditions
"""
# If a bare string is passed
if isinstance(h5_file, str) or isinstance(h5_file, pathlib.Path):
with h5py.File(h5_file,'r') as f:
return h5_to_nested_dict(f)
d = {}
for key in h5_file.keys():
value = h5_file[key]
if isinstance(value, h5py.Dataset):
arr = np.array(value)
if arr.dtype == object:
d[key] = arr.ravel()[0].decode('utf-8')
elif arr.ndim == 0:
d[key] = arr.ravel()[0]
else:
d[key] = arr
elif isinstance(value, h5py.Group):
sub_d = h5_to_nested_dict(value)
d[key] = sub_d
else:
raise ValueError(f'{value} could not be interpreted sensibly')
return d
src/cdtools/tools/initializers/initializers.py
+34 -90
View File
@@ -41,10 +41,6 @@ def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None,
The wavelength of light for the experiment, in m
distance : float
The sample-detector distance, in m
center : torch.Tensor
If defined, the location of the zero frequency pixel
padding : int
Default is 0, the size of an extra border of nonphysical pixels around the detector
oversampling : int
Default is 1, the amount to multiply the exit wave shape by.
@@ -52,66 +48,32 @@ def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None,
-------
basis : torch.Tensor
The exit wave basis
shape : torch.Tensor
The exit wave's shape
slice : slice
The slice corresponding to the physical detector
"""
det_shape = t.as_tensor(tuple(det_shape), dtype=t.int32)
det_basis = t.as_tensor(det_basis)
# First, set the center if it's not already specified
# This definition matches the center pixel of an fftshifted array
if center is None:
# this is center//2, but in pytorch 1.9.0 it throws a warning
# if you just do that
center = t.div(det_shape,2,rounding_mode='floor')# // 2
else:
center = t.as_tensor(center, dtype=t.int32)
# Then, calculate the required detector size from the centering
# This is a bit opaque but was worth doing accurately
min_left = center * 2
min_right = (det_shape - center) * 2 - 1
full_shape = t.max(min_left,min_right).to(t.int32) + 2 * padding
# Generate the basis for the exit wave in real space
# In some edge cases this shape can be smaller than the detector shape
full_shape = t.max(full_shape, det_shape)
# Then, generate a slice that pops the actual detector from the full
# detector shape
# this is full_center//2, but in pytorch 1.9.0 it throws a warning
# if you just do that
full_center = t.div(full_shape,2, rounding_mode='floor')# // 2
det_slice = np.s_[int(full_center[0]-center[0]):
int(full_center[0]-center[0]+det_shape[0]),
int(full_center[1]-center[1]):
int(full_center[1]-center[1]+det_shape[1])]
# Finally, generate the basis for the exit wave in real space
# This method should work for a general parallelogram
# shaped detector
det_shape = det_basis * full_shape.to(t.float32)
# This method should work for a general parallelogram-shaped detector
det_shape = det_basis * det_shape.to(t.float32)
pinv_basis = t.tensor(np.linalg.pinv(det_shape).transpose()).to(t.float32)
real_space_basis = pinv_basis * wavelength * distance
# This is definitely correct, but less simple. Included here
return real_space_basis
# Below is definitely correct, but less simple. Included here
# So future me can check that both versions are consistent.
#oop_dir = np.cross(det_basis[:,0],det_basis[:,1])
#oop_dir /= np.linalg.norm(oop_dir)
#full_basis = np.array([np.array(det_basis[:,0]),np.array(det_basis[:,1]),oop_dir]).transpose()
#inv_basis = t.tensor(np.linalg.inv(full_basis)[:2,:].transpose()).to(t.float32)
#real_space_basis = inv_basis*wavelength * distance / \
# oop_dir = np.cross(det_basis[:,0],det_basis[:,1])
# oop_dir /= np.linalg.norm(oop_dir)
# full_basis = np.array([np.array(det_basis[:,0]),
# np.array(det_basis[:,1]),oop_dir]).transpose()
# inv_basis = \
# t.tensor(np.linalg.inv(full_basis)[:2,:].transpose()).to(t.float32)
# real_space_basis = inv_basis*wavelength * distance / \
# full_shape.to(t.float32)
# Finally, convert the shape back to a torch.Size
full_shape = t.Size([dim * oversampling for dim in full_shape])
return real_space_basis, full_shape, det_slice
def calc_object_setup(probe_shape, translations, padding=0):
@@ -291,7 +253,7 @@ def gaussian_probe(dataset, basis, shape, sigma, propagation_distance=0, polariz
def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None, oversampling=1, polarized=False, left_polarized=True):
def SHARP_style_probe(dataset, propagation_distance=None, oversampling=1):
"""Generates a SHARP style probe guess from a dataset
What we call the "SHARP" style probe guess is to take a mean of all
@@ -313,10 +275,6 @@ def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None, over
----------
dataset : Ptycho_2D_Dataset
The dataset to work from
shape : torch.Size
The size of the probe array to simulate
det_slice : slice
A slice or tuple of slices corresponding to the detector region in Fourier space
propagation_distance : float
Default is no propagation, an amount to propagate the guessed probe from it's focal point
oversampling : int
@@ -331,43 +289,33 @@ def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None, over
# NOTE: I don't love the way np and torch are mixed here, I think this
# function deserves some love.
shape = dataset.patterns.shape[-2:]
# to use the mask or not?
intensities = np.zeros([dim // oversampling for dim in shape])
if polarized:
factors = [(math.cos(math.radians(polarizer[idx] - analyzer[idx])))**2 for idx in range(len(dataset)) if abs(polarizer[idx] - analyzer[idx]) > 5]
else:
factors = [1 for idx in range(len(dataset))]
# Eventually, do something with the recorded intensities, if they exist
factors = [1 for idx in range(len(dataset))]
for params, im in dataset:
if hasattr(dataset,'mask') and dataset.mask is not None:
intensities[det_slice] += dataset.mask.cpu().numpy() * im.cpu().numpy() / factors[params[0]]
intensities += (dataset.mask.cpu().numpy() * im.cpu().numpy()
/ factors[params[0]])
else:
intensities[det_slice] += im.cpu().numpy() / params[factors[0]]
intensities += im.cpu().numpy() / params[factors[0]]
intensities /= len(dataset)
# Subtract off a known background if it's stored
if hasattr(dataset, 'background') and dataset.background is not None:
intensities[det_slice] = np.clip(intensities[det_slice] - dataset.background.cpu().numpy(), a_min=0,a_max=None)
intensities = np.clip(
intensities - dataset.background.cpu().numpy(),
a_min=0,
a_max=None,
)
probe_fft = t.tensor(np.sqrt(intensities)).to(dtype=t.complex64)
probe_guess = inverse_far_field(probe_fft).numpy()
# Now we remove the central pixel
#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 = t.as_tensor(probe_guess, dtype=t.complex64)
probe_guess = inverse_far_field(probe_fft)
if propagation_distance is not None:
# First generate the propagation array
@@ -386,7 +334,11 @@ def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None, over
probe_shape = probe_shape.numpy().astype(np.int32)
# And generate the propagator
AS_prop = generate_angular_spectrum_propagator(probe_shape, probe_spacing, dataset.wavelength, propagation_distance)
AS_prop = generate_angular_spectrum_propagator(
probe_shape,
probe_spacing,
dataset.wavelength,
propagation_distance)
probe_guess = near_field(probe_guess,AS_prop)
@@ -396,14 +348,6 @@ def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None, over
top = shape[1]//2 - probe_guess.shape[1] // 2
final_probe[left:left+probe_guess.shape[0],
top:top+probe_guess.shape[1]] = probe_guess
if polarized:
if left_polarized:
x = 1j
else:
x = -1j
final_probe = t.stack((final_probe.to(dtype=t.cfloat), final_probe.to(dtype=t.cfloat)), dim=-4)
final_probe = final_probe[..., None, :, :]
return final_probe
tests/tools/test_initializers.py
+4 -8
View File
@@ -10,18 +10,14 @@ def test_exit_wave_geometry():
shape = t.Size([73,56])
wavelength = 1e-9
distance = 1.
rs_basis, full_shape, det_slice = \
initializers.exit_wave_geometry(basis, shape, wavelength,
distance)
assert full_shape == shape
assert t.ones(full_shape)[det_slice].shape == shape
rs_basis = initializers.exit_wave_geometry(basis, shape, wavelength, distance)
assert t.allclose(rs_basis[0,1],t.Tensor([-8.928571428571428e-07]))
assert t.allclose(rs_basis[1,0],t.Tensor([-4.5662100456621004e-07]))
# Then test it's padding function
rs_basis, full_shape, det_slice = \
initializers.exit_wave_geometry(basis, shape, wavelength,
distance, padding=2)
rs_basis = initializers.exit_wave_geometry(basis, shape, wavelength,
distance, padding=2)
exp_shape = t.Size([77,60])
assert full_shape == exp_shape
assert t.ones(full_shape)[det_slice].shape == shape