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

Update remaining models to new plot registration system

Browse Source
This commit is contained in:
allevitan
2026-03-20 19:01:22 +01:00
parent e7f943e254
commit 91db5b3c64
5 changed files with 257 additions and 193 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/transmission_RPI.py
+6 -4
View File
@@ -1,5 +1,6 @@
import cdtools
import pickle
import torch as t
from matplotlib import pyplot as plt
# First, we load an example dataset from a .cxi file
@@ -18,17 +19,18 @@ dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(ss_filename)
# Note that we explicitly as for two incoherent probe modes
model = cdtools.models.RPI.from_dataset(dataset, probe, [500,500],
background=background, n_modes=2,
initialization='random')
initialization='random',panel_plot_mode=True)
# Let's do this reconstruction on the GPU, shall we?
model.to(device='cuda')
dataset.get_as(device='cuda')
if t.cuda.is_available():
model.to(device='cuda')
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 loss in model.LBFGS_optimize(30, dataset, lr=0.4, regularization_factor=[0.05,0.05]):
for loss in model.Adam_optimize(30, dataset, lr=0.4, regularization_factor=[0.05,0.05]):
model.inspect(dataset)
print(model.report())
src/cdtools/models/bragg_2d_ptycho.py
+49 -42
View File
@@ -80,6 +80,8 @@ class Bragg2DPtycho(CDIModel):
units='um',
dtype=t.float32,
obj_view_crop=0,
panel_plot_mode=False,
plot_level=1,
):
# We need the detector geometry
@@ -91,7 +93,8 @@ class Bragg2DPtycho(CDIModel):
# translation_offsets can stay 2D for now
# propagate_probe and correct_tilt are important!
super(Bragg2DPtycho, self).__init__()
super(Bragg2DPtycho, self).__init__(panel_plot_mode=panel_plot_mode,
plot_level=plot_level)
self.register_buffer('wavelength',
t.as_tensor(wavelength, dtype=dtype))
self.store_detector_geometry(detector_geometry,
@@ -258,7 +261,9 @@ class Bragg2DPtycho(CDIModel):
obj_padding=200,
obj_view_crop=None,
units='um',
surface_normal=None
surface_normal=None,
panel_plot_mode=False,
plot_level=1,
):
wavelength = dataset.wavelength
det_basis = dataset.detector_geometry['basis']
@@ -436,8 +441,8 @@ class Bragg2DPtycho(CDIModel):
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,
median_propagation=median_propagation,
translation_offsets=translation_offsets,
weights=weights, mask=mask, background=background,
translation_scale=translation_scale,
saturation=saturation,
@@ -448,6 +453,8 @@ class Bragg2DPtycho(CDIModel):
lens=lens,
obj_view_crop=obj_view_crop,
units=units,
panel_plot_mode=panel_plot_mode,
plot_level=plot_level,
)
@@ -578,90 +585,90 @@ class Bragg2DPtycho(CDIModel):
plot_list = [
('Basis Probe Fourier Space Amplitudes',
lambda self, fig: p.plot_amplitude(tools.propagators.inverse_far_field(self.probe), fig=fig)),
('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, Surface Normal View',
lambda self, fig: p.plot_amplitude(
{'title': 'Basis Probe Fourier Space Amplitudes',
'plot_func': lambda self, fig: p.plot_amplitude(tools.propagators.inverse_far_field(self.probe), fig=fig)},
{'title': 'Basis Probe Fourier Space Phases',
'plot_func': lambda self, fig: p.plot_phase(tools.propagators.inverse_far_field(self.probe), fig=fig)},
{'title': 'Basis Probe Real Space Amplitudes, Surface Normal View',
'plot_func': lambda self, fig: p.plot_amplitude(
self.probe,
fig=fig,
basis=self.probe_basis,
units=self.units,
)),
('Basis Probe Real Space Phases, Surface Normal View',
lambda self, fig: p.plot_phase(
)},
{'title': 'Basis Probe Real Space Phases, Surface Normal View',
'plot_func': lambda self, fig: p.plot_phase(
self.probe,
fig=fig,
basis=self.probe_basis,
units=self.units,
)),
('Basis Probe Real Space Amplitudes, Beam View',
lambda self, fig: p.plot_amplitude(
)},
{'title': 'Basis Probe Real Space Amplitudes, Beam View',
'plot_func': lambda self, fig: p.plot_amplitude(
self.probe,
fig=fig,
basis=self.probe_basis,
view_basis=beam_basis,
units=self.units,
)),
('Basis Probe Real Space Phases, Beam View',
lambda self, fig: p.plot_phase(
)},
{'title': 'Basis Probe Real Space Phases, Beam View',
'plot_func': lambda self, fig: p.plot_phase(
self.probe,
fig=fig,
basis=self.probe_basis,
view_basis=beam_basis,
units=self.units,
)),
('Object Amplitude, Surface Normal View',
lambda self, fig: p.plot_amplitude(
)},
{'title': 'Object Amplitude, Surface Normal View',
'plot_func': lambda self, fig: p.plot_amplitude(
self.obj[self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
units=self.units,
)),
('Object Phase, Surface Normal View',
lambda self, fig: p.plot_phase(
)},
{'title': 'Object Phase, Surface Normal View',
'plot_func': lambda self, fig: p.plot_phase(
self.obj[self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
units=self.units,
)),
('Object Amplitude, Beam View',
lambda self, fig: p.plot_amplitude(
)},
{'title': 'Object Amplitude, Beam View',
'plot_func': lambda self, fig: p.plot_amplitude(
self.obj[self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
view_basis=beam_basis,
units=self.units,
)),
('Object Phase, Beam View',
lambda self, fig: p.plot_phase(
)},
{'title': 'Object Phase, Beam View',
'plot_func': lambda self, fig: p.plot_phase(
self.obj[self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
view_basis=beam_basis,
units=self.units,
)),
('Object Amplitude, Detector View',
lambda self, fig: p.plot_amplitude(
)},
{'title': 'Object Amplitude, Detector View',
'plot_func': lambda self, fig: p.plot_amplitude(
self.obj[self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
view_basis=self.det_basis,
units=self.units,
)),
('Object Phase, Detector View',
lambda self, fig: p.plot_phase(
)},
{'title': 'Object Phase, Detector View',
'plot_func': lambda self, fig: p.plot_phase(
self.obj[self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
view_basis=self.det_basis,
units=self.units,
)),
('Corrected Translations',
lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)),
('Background',
lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2))
)},
{'title': 'Corrected Translations',
'plot_func': lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)},
{'title': 'Background',
'plot_func': lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2)},
]
src/cdtools/models/multislice_2d_ptycho.py
+52 -47
View File
@@ -47,9 +47,12 @@ class Multislice2DPtycho(CDIModel):
prevent_aliasing=True,
phase_only=False,
units='um',
panel_plot_mode=False,
plot_level=1,
):
super(Multislice2DPtycho, self).__init__()
super(Multislice2DPtycho, self).__init__(panel_plot_mode=panel_plot_mode,
plot_level=plot_level)
self.wavelength = t.tensor(wavelength)
self.detector_geometry = copy(detector_geometry)
self.dz = dz
@@ -154,7 +157,7 @@ class Multislice2DPtycho(CDIModel):
@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):
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, panel_plot_mode=False, plot_level=1):
wavelength = dataset.wavelength
det_basis = dataset.detector_geometry['basis']
@@ -286,7 +289,7 @@ class Multislice2DPtycho(CDIModel):
surface_normal=surface_normal,
probe_support=probe_support,
min_translation=min_translation,
translation_offsets = translation_offsets,
translation_offsets=translation_offsets,
weights=Ws, mask=mask, background=background,
translation_scale=translation_scale,
saturation=saturation,
@@ -296,7 +299,9 @@ class Multislice2DPtycho(CDIModel):
exponentiate_obj=exponentiate_obj,
units=units, fourier_probe=fourier_probe,
phase_only=phase_only,
prevent_aliasing=prevent_aliasing)
prevent_aliasing=prevent_aliasing,
panel_plot_mode=panel_plot_mode,
plot_level=plot_level)
def interaction(self, index, translations):
@@ -581,22 +586,22 @@ class Multislice2DPtycho(CDIModel):
# Needs to be updated to allow for plotting to an existing figure
plot_list = [
('Probe Fourier Space Amplitude',
lambda self, fig: p.plot_amplitude(self.probe if self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig)),
('Probe Fourier Space Phase',
lambda self, fig: p.plot_phase(self.probe if self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig)),
('Probe Real Space Amplitude',
lambda self, fig: p.plot_amplitude(self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig, basis=self.probe_basis, units=self.units)),
('Probe Real Space Phase',
lambda self, fig: p.plot_phase(self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig, basis=self.probe_basis, units=self.units)),
('Average Weight Matrix Amplitudes',
lambda self, fig: p.plot_amplitude(
plot_list = [
{'title': 'Probe Fourier Space Amplitude',
'plot_func': lambda self, fig: p.plot_amplitude(self.probe if self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig)},
{'title': 'Probe Fourier Space Phase',
'plot_func': lambda self, fig: p.plot_phase(self.probe if self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig)},
{'title': 'Probe Real Space Amplitude',
'plot_func': lambda self, fig: p.plot_amplitude(self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig, basis=self.probe_basis, units=self.units)},
{'title': 'Probe Real Space Phase',
'plot_func': lambda self, fig: p.plot_phase(self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig, basis=self.probe_basis, units=self.units)},
{'title': 'Average Weight Matrix Amplitudes',
'plot_func': lambda self, fig: p.plot_amplitude(
np.nanmean(np.abs(self.weights.data.cpu().numpy()), axis=0),
fig=fig),
lambda self: len(self.weights.shape) >= 2),
('% of Power in Top Mode',
lambda self, fig, dataset: p.plot_nanomap(
'condition': lambda self: len(self.weights.shape) >= 2},
{'title': '% of Power in Top Mode',
'plot_func': lambda self, fig, dataset: p.plot_nanomap(
self.corrected_translations(dataset),
100 * t.stack([
analysis.calc_mode_power_fractions(
@@ -606,35 +611,35 @@ class Multislice2DPtycho(CDIModel):
], dim=0),
fig=fig,
units=self.units),
lambda self: len(self.weights.shape) >= 2),
('Slice by Slice Real Part of T',
lambda self, fig: p.plot_real(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units, cmap='cividis'),
lambda self: self.exponentiate_obj),
('Slice by Slice Imaginary Part of T',
lambda self, fig: p.plot_imag(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units),
lambda self: self.exponentiate_obj),
('Integrated Real Part of T',
lambda self, fig: p.plot_real(t.sum(self.obj.detach().cpu(),dim=0), fig=fig, basis=self.probe_basis, units=self.units, cmap='cividis'),
lambda self: (self.exponentiate_obj) and self.obj.dim() >= 3),
('Integrated Imaginary Part of T',
lambda self, fig: p.plot_imag(t.sum(self.obj.detach().cpu(),dim=0), fig=fig, basis=self.probe_basis, units=self.units),
lambda self: (self.exponentiate_obj) and self.obj.dim() >= 3),
('Slice by Slice Amplitude of Object Function',
lambda self, fig: p.plot_amplitude(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units),
lambda self: not self.exponentiate_obj),
('Slice by Slice Phase of Object Function',
lambda self, fig: p.plot_phase(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units,cmap='cividis'),
lambda self: not self.exponentiate_obj),
('Amplitude of Stacked Object Function',
lambda self, fig: p.plot_amplitude(reduce(t.mul, self.obj.detach().cpu()), fig=fig, basis=self.probe_basis, units=self.units),
lambda self: (not self.exponentiate_obj) and self.obj.dim() >=3),
('Phase of Stacked Object Function',
lambda self, fig: p.plot_phase(reduce(t.mul, self.obj.detach().cpu()), fig=fig, basis=self.probe_basis, units=self.units, cmap='cividis'),
lambda self: (not self.exponentiate_obj) and self.obj.dim() >= 3),
('Corrected Translations',
lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)),
('Background',
lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2))
'condition': lambda self: len(self.weights.shape) >= 2},
{'title': 'Slice by Slice Real Part of T',
'plot_func': lambda self, fig: p.plot_real(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units, cmap='cividis'),
'condition': lambda self: self.exponentiate_obj},
{'title': 'Slice by Slice Imaginary Part of T',
'plot_func': lambda self, fig: p.plot_imag(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units),
'condition': lambda self: self.exponentiate_obj},
{'title': 'Integrated Real Part of T',
'plot_func': lambda self, fig: p.plot_real(t.sum(self.obj.detach().cpu(),dim=0), fig=fig, basis=self.probe_basis, units=self.units, cmap='cividis'),
'condition': lambda self: self.exponentiate_obj and self.obj.dim() >= 3},
{'title': 'Integrated Imaginary Part of T',
'plot_func': lambda self, fig: p.plot_imag(t.sum(self.obj.detach().cpu(),dim=0), fig=fig, basis=self.probe_basis, units=self.units),
'condition': lambda self: self.exponentiate_obj and self.obj.dim() >= 3},
{'title': 'Slice by Slice Amplitude of Object Function',
'plot_func': lambda self, fig: p.plot_amplitude(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units),
'condition': lambda self: not self.exponentiate_obj},
{'title': 'Slice by Slice Phase of Object Function',
'plot_func': lambda self, fig: p.plot_phase(self.obj.detach().cpu(), fig=fig, basis=self.probe_basis, units=self.units, cmap='cividis'),
'condition': lambda self: not self.exponentiate_obj},
{'title': 'Amplitude of Stacked Object Function',
'plot_func': lambda self, fig: p.plot_amplitude(reduce(t.mul, self.obj.detach().cpu()), fig=fig, basis=self.probe_basis, units=self.units),
'condition': lambda self: (not self.exponentiate_obj) and self.obj.dim() >= 3},
{'title': 'Phase of Stacked Object Function',
'plot_func': lambda self, fig: p.plot_phase(reduce(t.mul, self.obj.detach().cpu()), fig=fig, basis=self.probe_basis, units=self.units, cmap='cividis'),
'condition': lambda self: (not self.exponentiate_obj) and self.obj.dim() >= 3},
{'title': 'Corrected Translations',
'plot_func': lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)},
{'title': 'Background',
'plot_func': lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2)},
]
src/cdtools/models/multislice_ptycho.py
+65 -58
View File
@@ -39,10 +39,13 @@ class MultislicePtycho(CDIModel):
simulate_finite_pixels=False,
dtype=t.float32,
exponentiate_obj=False,
obj_view_crop=0
obj_view_crop=0,
panel_plot_mode=False,
plot_level=1,
):
super(MultislicePtycho, self).__init__()
super(MultislicePtycho, self).__init__(panel_plot_mode=panel_plot_mode,
plot_level=plot_level)
self.register_buffer('wavelength',
t.as_tensor(wavelength, dtype=dtype))
self.store_detector_geometry(detector_geometry,
@@ -202,6 +205,8 @@ class MultislicePtycho(CDIModel):
obj_view_crop=None,
obj_padding=200,
exponentiate_obj=False,
panel_plot_mode=False,
plot_level=1,
):
wavelength = dataset.wavelength
@@ -409,6 +414,8 @@ class MultislicePtycho(CDIModel):
simulate_finite_pixels=simulate_finite_pixels,
exponentiate_obj=exponentiate_obj,
obj_view_crop=obj_view_crop,
panel_plot_mode=panel_plot_mode,
plot_level=plot_level,
)
@@ -750,61 +757,61 @@ class MultislicePtycho(CDIModel):
plot_list = [
('',
lambda self, fig, dataset: self.plot_wavefront_variation(
{'title': '',
'plot_func': lambda self, fig, dataset: self.plot_wavefront_variation(
dataset,
fig=fig,
mode='root_sum_intensity',
image_title='Root Summed Probe Intensities',
image_colorbar_title='Square Root of Intensity'),
lambda self: len(self.weights.shape) >= 2),
('',
lambda self, fig, dataset: self.plot_wavefront_variation(
'condition': lambda self: len(self.weights.shape) >= 2},
{'title': '',
'plot_func': lambda self, fig, dataset: self.plot_wavefront_variation(
dataset,
fig=fig,
mode='amplitude',
image_title='Probe Amplitudes (scroll to view modes)',
image_colorbar_title='Probe Amplitude'),
lambda self: len(self.weights.shape) >= 2),
('',
lambda self, fig, dataset: self.plot_wavefront_variation(
'condition': lambda self: len(self.weights.shape) >= 2},
{'title': '',
'plot_func': lambda self, fig, dataset: self.plot_wavefront_variation(
dataset,
fig=fig,
mode='phase',
image_title='Probe Phases (scroll to view modes)',
image_colorbar_title='Probe Phase'),
lambda self: len(self.weights.shape) >= 2),
('Basis Probe Fourier Space Amplitudes',
lambda self, fig: p.plot_amplitude(
'condition': lambda self: len(self.weights.shape) >= 2},
{'title': 'Basis Probe Fourier Space Amplitudes',
'plot_func': lambda self, fig: p.plot_amplitude(
(self.probe if self.fourier_probe
else tools.propagators.inverse_far_field(self.probe)),
fig=fig)),
('Basis Probe Fourier Space Phases',
lambda self, fig: p.plot_phase(
fig=fig)},
{'title': 'Basis Probe Fourier Space Phases',
'plot_func': lambda self, fig: p.plot_phase(
(self.probe if self.fourier_probe
else tools.propagators.inverse_far_field(self.probe))
, fig=fig)),
('Basis Probe Real Space Amplitudes',
lambda self, fig: p.plot_amplitude(
else tools.propagators.inverse_far_field(self.probe)),
fig=fig)},
{'title': 'Basis Probe Real Space Amplitudes',
'plot_func': lambda self, fig: p.plot_amplitude(
(self.probe if not self.fourier_probe
else tools.propagators.inverse_far_field(self.probe)),
fig=fig,
basis=self.probe_basis,
units=self.units)),
('Basis Probe Real Space Phases',
lambda self, fig: p.plot_phase(
units=self.units)},
{'title': 'Basis Probe Real Space Phases',
'plot_func': lambda self, fig: p.plot_phase(
(self.probe if not self.fourier_probe
else tools.propagators.inverse_far_field(self.probe)),
fig=fig,
basis=self.probe_basis,
units=self.units)),
('Average Weight Matrix Amplitudes',
lambda self, fig: p.plot_amplitude(
units=self.units)},
{'title': 'Average Weight Matrix Amplitudes',
'plot_func': lambda self, fig: p.plot_amplitude(
np.nanmean(np.abs(self.weights.data.cpu().numpy()), axis=0),
fig=fig),
lambda self: len(self.weights.shape) >= 2),
('% of Power in Top Mode',
lambda self, fig, dataset: p.plot_nanomap(
'condition': lambda self: len(self.weights.shape) >= 2},
{'title': '% of Power in Top Mode',
'plot_func': lambda self, fig, dataset: p.plot_nanomap(
self.corrected_translations(dataset),
100 * t.stack([
analysis.calc_mode_power_fractions(
@@ -814,69 +821,69 @@ class MultislicePtycho(CDIModel):
], dim=0),
fig=fig,
units=self.units),
lambda self: len(self.weights.shape) >= 2),
('Object Amplitude',
lambda self, fig: p.plot_amplitude(
'condition': lambda self: len(self.weights.shape) >= 2},
{'title': 'Object Amplitude',
'plot_func': lambda self, fig: p.plot_amplitude(
self.obj[(np.s_[:],) + self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
units=self.units),
lambda self: not self.exponentiate_obj),
('Object (T) Imaginary Part',
lambda self, fig: p.plot_imag(
'condition': lambda self: not self.exponentiate_obj},
{'title': 'Object (T) Imaginary Part',
'plot_func': lambda self, fig: p.plot_imag(
self.obj[(np.s_[:],) + self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
units=self.units),
lambda self: self.exponentiate_obj),
('Object Phase',
lambda self, fig: p.plot_phase(
'condition': lambda self: self.exponentiate_obj},
{'title': 'Object Phase',
'plot_func': lambda self, fig: p.plot_phase(
self.obj[(np.s_[:],) + self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
units=self.units),
lambda self: not self.exponentiate_obj),
('Object (T) Real Part',
lambda self, fig: p.plot_real(
'condition': lambda self: not self.exponentiate_obj},
{'title': 'Object (T) Real Part',
'plot_func': lambda self, fig: p.plot_real(
self.obj[(np.s_[:],) + self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
units=self.units,
cmap='cividis'),
lambda self: self.exponentiate_obj),
('Object Product Amplitude',
lambda self, fig: p.plot_amplitude(
'condition': lambda self: self.exponentiate_obj},
{'title': 'Object Product Amplitude',
'plot_func': lambda self, fig: p.plot_amplitude(
t.prod(self.obj, dim=0)[self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
units=self.units),
lambda self: not self.exponentiate_obj),
('Object (T) Sum Imaginary Part',
lambda self, fig: p.plot_imag(
'condition': lambda self: not self.exponentiate_obj},
{'title': 'Object (T) Sum Imaginary Part',
'plot_func': lambda self, fig: p.plot_imag(
t.sum(self.obj, dim=0)[self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
units=self.units),
lambda self: self.exponentiate_obj),
('Object Product Phase',
lambda self, fig: p.plot_phase(
'condition': lambda self: self.exponentiate_obj},
{'title': 'Object Product Phase',
'plot_func': lambda self, fig: p.plot_phase(
t.prod(self.obj, dim=0)[self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
units=self.units),
lambda self: not self.exponentiate_obj),
('Object (T) Sum Real Part',
lambda self, fig: p.plot_real(
'condition': lambda self: not self.exponentiate_obj},
{'title': 'Object (T) Sum Real Part',
'plot_func': lambda self, fig: p.plot_real(
t.sum(self.obj, dim=0)[self.obj_view_slice],
fig=fig,
basis=self.obj_basis,
units=self.units,
cmap='cividis'),
lambda self: self.exponentiate_obj),
('Corrected Translations',
lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)),
('Background',
lambda self, fig: p.plot_amplitude(self.background**2, fig=fig))
'condition': lambda self: self.exponentiate_obj},
{'title': 'Corrected Translations',
'plot_func': lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)},
{'title': 'Background',
'plot_func': lambda self, fig: p.plot_amplitude(self.background**2, fig=fig)},
]
src/cdtools/models/rpi.py
+85 -42
View File
@@ -58,9 +58,12 @@ class RPI(CDIModel):
propagation_distance=0,
units='um',
dtype=t.float32,
panel_plot_mode=False,
plot_level=1,
):
super(RPI, self).__init__()
super(RPI, self).__init__(panel_plot_mode=panel_plot_mode,
plot_level=plot_level)
complex_dtype = (t.ones([1], dtype=dtype) +
1j * t.ones([1], dtype=dtype)).dtype
@@ -164,6 +167,8 @@ class RPI(CDIModel):
phase_only=False,
probe_threshold=0,
dtype=t.float32,
panel_plot_mode=False,
plot_level=1,
):
complex_dtype = (t.ones([1], dtype=dtype) +
1j * t.ones([1], dtype=dtype)).dtype
@@ -227,7 +232,7 @@ class RPI(CDIModel):
# This will be superceded later by a call to init_obj, but it sets
# the shape
if obj_size is None:
obj_size = (np.array(self.probe.shape[-2:]) // 2).astype(int)
obj_size = (np.array(probe.shape[-2:]) // 2).astype(int)
dummy_init_obj = t.ones([n_modes, obj_size[0], obj_size[1]],
dtype=complex_dtype)
@@ -247,14 +252,16 @@ class RPI(CDIModel):
obj_support = t.as_tensor(binary_dilation(obj_support))
rpi_object = cls(wavelength, det_geo, ew_basis,
probe, dummy_init_obj,
probe, dummy_init_obj,
background=background, mask=mask,
saturation=saturation,
obj_support=obj_support,
oversampling=oversampling,
exponentiate_obj=exponentiate_obj,
phase_only=phase_only,
weight_matrix=weight_matrix)
weight_matrix=weight_matrix,
panel_plot_mode=panel_plot_mode,
plot_level=plot_level)
# I don't love this pattern, where I do the "real" obj initialization
# after creating the rpi object. But, I chose this so that I could
@@ -283,7 +290,9 @@ class RPI(CDIModel):
exponentiate_obj=False,
phase_only=False,
initialization='random',
dtype=t.float32
dtype=t.float32,
panel_plot_mode=False,
plot_level=1,
):
complex_dtype = (t.ones([1], dtype=dtype) +
@@ -308,7 +317,7 @@ class RPI(CDIModel):
# This will be superceded later by a call to init_obj, but it sets
# the shape
if obj_size is None:
obj_size = (np.array(self.probe.shape[-2:]) // 2).astype(int)
obj_size = (np.array(probe.shape[-2:]) // 2).astype(int)
dummy_init_obj = t.ones([n_modes, obj_size[0], obj_size[1]],
dtype=complex_dtype)
@@ -327,6 +336,8 @@ class RPI(CDIModel):
mask=mask,
exponentiate_obj=exponentiate_obj,
phase_only=phase_only,
panel_plot_mode=panel_plot_mode,
plot_level=plot_level,
)
rpi_object.init_obj(initialization)
@@ -365,7 +376,7 @@ class RPI(CDIModel):
obj_shape=obj_shape,
n_modes=n_modes)
else:
raise KeyError('Initialization "' + str(initialization) + \
raise KeyError('Initialization "' + str(initialization_type) + \
'" invalid - use "spectral", "uniform", or "random"')
@@ -531,40 +542,72 @@ class RPI(CDIModel):
def sim_to_dataset(self, args_list):
raise NotImplementedError('No sim to dataset yet, sorry!')
plot_list = [
('Root Sum Squared Amplitude of all Probes',
lambda self, fig: p.plot_amplitude(
np.sqrt(np.sum((t.abs(t.sum(self.weights[..., None, None].detach() * self.probe, axis=-3))**2).cpu().numpy(),axis=0)),
fig=fig, basis=self.probe_basis)),
('Object Amplitude',
lambda self, fig: p.plot_amplitude(
self.obj,
fig=fig,
basis=self.obj_basis,
units=self.units),
lambda self: not self.exponentiate_obj),
('Object Phase',
lambda self, fig: p.plot_phase(
self.obj,
fig=fig,
basis=self.obj_basis,
units=self.units),
lambda self: not self.exponentiate_obj),
('Real Part of T',
lambda self, fig: p.plot_real(
self.obj,
fig=fig,
basis=self.obj_basis,
units=self.units,
cmap='cividis'),
lambda self: self.exponentiate_obj),
('Imaginary Part of T',
lambda self, fig: p.plot_imag(
self.obj,
fig=fig,
basis=self.obj_basis,
units=self.units),
lambda self: self.exponentiate_obj),
plot_panel_list = [
{
'title': 'RPI Results',
'plot_level': 1,
'figure_size': (12, 3.5),
'grid': (1, 3),
'plots': [
{
'title': 'Object Phase',
'subplot': (0, 0),
'plot_func': lambda self, fig: p.plot_phase(
self.obj,
fig=fig,
title='Object Phase',
basis=self.obj_basis,
units=self.units),
'condition': lambda self: not self.exponentiate_obj,
},
{
'title': 'Real Part of T',
'subplot': (0, 0),
'plot_func': lambda self, fig: p.plot_real(
self.obj,
fig=fig,
title='Real Part of T',
basis=self.obj_basis,
units=self.units,
cmap='cividis'),
'condition': lambda self: self.exponentiate_obj,
},
{
'title': 'Object Amplitude',
'subplot': (0, 1),
'plot_func': lambda self, fig: p.plot_amplitude(
self.obj,
fig=fig,
title='Object Amplitude',
basis=self.obj_basis,
units=self.units),
'condition': lambda self: not self.exponentiate_obj,
},
{
'title': 'Imaginary Part of T',
'subplot': (0, 1),
'plot_func': lambda self, fig: p.plot_imag(
self.obj,
fig=fig,
title='Imaginary Part of T',
basis=self.obj_basis,
units=self.units),
'condition': lambda self: self.exponentiate_obj,
},
{
'title': 'Root Sum Squared Amplitude of all Probes',
'subplot': (0, 2),
'plot_func': lambda self, fig: p.plot_amplitude(
np.sqrt(np.sum(
(t.abs(t.sum(self.weights[..., None, None].detach()
* self.probe, axis=-3))**2
).cpu().numpy(), axis=0)),
fig=fig,
basis=self.probe_basis,
units=self.units),
},
],
},
]