diff --git a/.gitignore b/.gitignore index 9bf303a..8032c1f 100644 --- a/.gitignore +++ b/.gitignore @@ -7,4 +7,4 @@ docs/build build/* dist -example_data/* \ No newline at end of file +*/example_data/* \ No newline at end of file diff --git a/examples/example_reconstructions/gold_balls.h5 b/examples/example_reconstructions/gold_balls.h5 new file mode 100644 index 0000000..6ef7867 Binary files /dev/null and b/examples/example_reconstructions/gold_balls.h5 differ diff --git a/examples/example_reconstructions/gold_balls.mat b/examples/example_reconstructions/gold_balls.mat deleted file mode 100644 index 5d38fbf..0000000 Binary files a/examples/example_reconstructions/gold_balls.mat and /dev/null differ diff --git a/examples/gold_ball_ptycho.py b/examples/gold_ball_ptycho.py index c34cfe7..705a3ed 100644 --- a/examples/gold_ball_ptycho.py +++ b/examples/gold_ball_ptycho.py @@ -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) diff --git a/examples/lab_bragg_2d_ptycho.py b/examples/lab_bragg_2d_ptycho.py index 6465f5e..331c3b5 100644 --- a/examples/lab_bragg_2d_ptycho.py +++ b/examples/lab_bragg_2d_ptycho.py @@ -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() diff --git a/src/cdtools/models/base.py b/src/cdtools/models/base.py index 888e27c..9901e62 100644 --- a/src/cdtools/models/base.py +++ b/src/cdtools/models/base.py @@ -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!' diff --git a/src/cdtools/models/bragg_2d_ptycho.py b/src/cdtools/models/bragg_2d_ptycho.py index b0d4697..e27ea1e 100644 --- a/src/cdtools/models/bragg_2d_ptycho.py +++ b/src/cdtools/models/bragg_2d_ptycho.py @@ -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', diff --git a/src/cdtools/models/fancy_ptycho.py b/src/cdtools/models/fancy_ptycho.py index 6b481d2..b7ee5da 100644 --- a/src/cdtools/models/fancy_ptycho.py +++ b/src/cdtools/models/fancy_ptycho.py @@ -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, diff --git a/src/cdtools/tools/data/data.py b/src/cdtools/tools/data/data.py index e1e0e70..aab0a5f 100644 --- a/src/cdtools/tools/data/data.py +++ b/src/cdtools/tools/data/data.py @@ -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 diff --git a/src/cdtools/tools/initializers/initializers.py b/src/cdtools/tools/initializers/initializers.py index 860d1cd..d4d80a4 100644 --- a/src/cdtools/tools/initializers/initializers.py +++ b/src/cdtools/tools/initializers/initializers.py @@ -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 diff --git a/tests/tools/test_initializers.py b/tests/tools/test_initializers.py index d4e9495..a5efd10 100644 --- a/tests/tools/test_initializers.py +++ b/tests/tools/test_initializers.py @@ -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