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

Buncha stuff, sorry I wasn't pushing

Browse Source
This commit is contained in:
Abe Levitan
2021-10-26 13:11:11 -04:00
parent f84e441843
commit 1b5ef2271c
8 changed files with 105 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CDTools/datasets/ptycho_2d_dataset.py
+5 -4
View File
@@ -188,7 +188,7 @@ class Ptycho2DDataset(CDataset):
cdtdata.add_ptycho_translations(cxi_file, self.translations)
def inspect(self, logarithmic=True, units='um'):
def inspect(self, logarithmic=True, units='um', log_offset=1):
"""Launches an interactive plot for perusing the data
This launches an interactive plotting tool in matplotlib that
@@ -208,7 +208,7 @@ class Ptycho2DDataset(CDataset):
mask = 1
if logarithmic:
return np.log(meas_data) / np.log(10) * mask
return np.log(meas_data + log_offset) / np.log(10) * mask
else:
return meas_data * mask
@@ -229,9 +229,10 @@ class Ptycho2DDataset(CDataset):
# nanomap_values = (self.mask * self.patterns).sum(dim=(1,2)).detach().cpu().numpy()
if logarithmic:
cbar_title='Log Base 10 of Diffraction Intensity'
cbar_title = ('Log Base 10 of Diffraction Intensity + %0.2f'
% log_offset)
else:
cbar_title='Diffraction Intensity'
cbar_title = 'Diffraction Intensity'
plotting.plot_nanomap_with_images(self.translations.detach().cpu(), get_images, values=nanomap_values, nanomap_units=units, image_title='Diffraction Pattern', image_colorbar_title=cbar_title)
CDTools/models/base.py
+5 -6
View File
@@ -55,7 +55,7 @@ class CDIModel(t.nn.Module):
def __init__(self):
super(CDIModel,self).__init__()
self.iteration_count = 0
self.loss_train = []
def from_dataset(self, dataset):
raise NotImplementedError()
@@ -183,9 +183,8 @@ class CDIModel(t.nn.Module):
if scheduler is not None:
scheduler.step(loss)
self.latest_loss = loss
self.loss_train.append(loss)
self.latest_iteration_time = time.time() - t0
self.iteration_count += 1
return loss
if thread:
@@ -412,10 +411,10 @@ class CDIModel(t.nn.Module):
report : str
A string with basic info on the latest iteration
"""
if hasattr(self, 'latest_loss'):
return 'Iteration ' + str(self.iteration_count) + \
if hasattr(self, 'latest_iteration_time'):
return 'Iteration ' + str(len(self.loss_train)) + \
' completed in %0.2f s with loss ' %\
self.latest_iteration_time + str(self.latest_loss)
self.latest_iteration_time + str(self.loss_train[-1])
else:
return 'No reconstruction iterations performed yet!'
CDTools/models/bragg_2d_ptycho.py
+32 -10
View File
@@ -54,6 +54,16 @@ __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,
@@ -124,7 +134,13 @@ class Bragg2DPtycho(CDIModel):
# you look at the phase map
probe_guess[probe_guess == 0] = 0
else:
self.probe_support = t.ones(self.probe[0].shape, dtype=t.bool)
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)
@@ -164,7 +180,7 @@ class Bragg2DPtycho(CDIModel):
# 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]
probe_shape = self.probe[0].shape
self.k_map, self.intensity_map = \
tools.propagators.generate_high_NA_k_intensity_map(
@@ -192,7 +208,7 @@ class Bragg2DPtycho(CDIModel):
@classmethod
def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=True, propagate_probe=True,correct_tilt=True, lens=False, opt_for_fft=False):
def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=True, propagate_probe=True, correct_tilt=True, lens=False, opt_for_fft=False):
wavelength = dataset.wavelength
det_basis = dataset.detector_geometry['basis']
@@ -309,12 +325,16 @@ class Bragg2DPtycho(CDIModel):
if probe_support_radius is not None:
probe_support = t.zeros(probe[0].shape, dtype=t.bool)
p_cent = np.array(probe[0].shape).astype(int) // 2
psr = int(probe_support_radius)
probe_support[p_cent[0]-psr:p_cent[0]+psr,
p_cent[1]-psr:p_cent[1]+psr] = 1
xs, ys = np.mgrid[:probe.shape[-2], :probe.shape[-1]]
xs = xs - np.mean(xs)
ys = ys - np.mean(ys)
Rs = np.sqrt(xs**2 + ys**2)
probe_support[Rs < probe_support_radius] = 1
probe = probe * probe_support[None, :, :]
else:
probe_support = None;
probe_support = None
# Here we need to implement a simple condition to choose whether
# to propagate the probe or not
@@ -352,7 +372,7 @@ class Bragg2DPtycho(CDIModel):
pix_trans += self.translation_scale * self.translation_offsets[index]
Ws = self.weights[index]
prs = Ws[...,None,None,None] * self.probe
prs = Ws[...,None,None,None] * self.probe * self.probe_support[...,:,:]
# Now we need to propagate each of the probes
@@ -504,8 +524,10 @@ class Bragg2DPtycho(CDIModel):
obj = self.obj.detach().cpu().numpy()
background = self.background.detach().cpu().numpy()**2
weights = self.weights.detach().cpu().numpy()
losses = np.array(self.loss_train)
return {'basis':basis, 'translation':translations,
'probe':probe,'obj':obj,
'background':background,
'weights':weights}
'weights':weights,
'losses':losses}
CDTools/models/fancy_ptycho.py
+54 -13
View File
@@ -17,14 +17,23 @@ class FancyPtycho(CDIModel):
def __init__(self, wavelength, detector_geometry,
probe_basis,
probe_guess, obj_guess,
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, translation_offsets=None, mask=None,
weights=None, translation_scale=1, saturation=None,
probe_support=None, oversampling=1,
loss='amplitude mse', units='um'):
background=None,
translation_offsets=None,
mask=None,
weights=None,
translation_scale=1,
saturation=None,
probe_support=None,
oversampling=1,
fourier_probe=False,
loss='amplitude mse',
units='um',
):
super(FancyPtycho, self).__init__()
self.wavelength = t.tensor(wavelength)
@@ -45,12 +54,13 @@ class FancyPtycho(CDIModel):
self.saturation = saturation
self.units = units
self.fourier_probe = fourier_probe
if mask is None:
self.mask = mask
else:
self.mask = t.tensor(mask, dtype=t.bool)
probe_guess = t.tensor(probe_guess, dtype=t.complex64)
obj_guess = t.tensor(obj_guess, dtype=t.complex64)
@@ -60,7 +70,7 @@ class FancyPtycho(CDIModel):
self.probe_norm = 1 * t.max(t.abs(probe_guess[0]))
else:
self.probe_norm = 1 * t.max(t.abs(probe_guess))
self.probe = t.nn.Parameter(probe_guess / self.probe_norm)
self.obj = t.nn.Parameter(obj_guess)
@@ -117,7 +127,25 @@ class FancyPtycho(CDIModel):
@classmethod
def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, dm_rank=None, translation_scale=1, saturation=None, probe_support_radius=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=False, opt_for_fft=False, loss='amplitude mse', units='um'):
def from_dataset(cls,
dataset,
probe_size=None,
randomize_ang=0,
padding=0,
n_modes=1,
dm_rank=None,
translation_scale=1,
saturation=None,
probe_support_radius=None,
propagation_distance=None,
scattering_mode=None,
oversampling=1,
auto_center=False,
opt_for_fft=False,
fourier_probe=False,
loss='amplitude mse',
units='um'
):
wavelength = dataset.wavelength
det_basis = dataset.detector_geometry['basis']
@@ -188,8 +216,12 @@ class FancyPtycho(CDIModel):
# Now we initialize all the subdominant probe modes
probe_max = t.max(t.abs(probe))
probe_stack = [0.01 * probe_max * t.rand(probe.shape, dtype=probe.dtype) for i in range(n_modes - 1)]
# For a Fourier space probe
if fourier_probe:
probe = tools.propagators.far_field(probe)
probe = t.stack([probe, ] + probe_stack)
# probe = t.stack([tools.propagators.far_field(probe),] + probe_stack)
obj = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5))
@@ -243,6 +275,7 @@ class FancyPtycho(CDIModel):
translation_scale=translation_scale,
saturation=saturation,
probe_support=probe_support,
fourier_probe=fourier_probe,
oversampling=oversampling,
loss=loss, units=units)
@@ -267,6 +300,10 @@ class FancyPtycho(CDIModel):
# This restricts the basis probes to stay within the probe support
basis_prs = self.probe * self.probe_support[..., :, :]
# For a Fourier-space probe, we take an IFT
if self.fourier_probe:
basis_prs = tools.propagators.inverse_far_field(basis_prs)
# Now we construct the probes for each shot from the basis probes
Ws = self.weights[index]
if len(self.weights[0].shape) == 0:
@@ -510,10 +547,14 @@ class FancyPtycho(CDIModel):
('',
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 Amplitudes (scroll to view modes)',
lambda self, fig: p.plot_amplitude(self.probe, fig=fig, basis=self.probe_basis, units=self.units)),
('Basis Probe Phases (scroll to view modes)',
lambda self, fig: p.plot_phase(self.probe, fig=fig, basis=self.probe_basis, units=self.units)),
('Basis Probe Fourier Space Amplitudes',
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(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(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(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 Density Matrix Amplitudes',
lambda self, fig: p.plot_amplitude(np.nanmean(np.abs(self.get_rhos()), axis=0), fig=fig),
lambda self: len(self.weights.shape) >= 2),
CDTools/models/multislice_2d_ptycho.py
+1 -1
View File
@@ -567,7 +567,7 @@ class Multislice2DPtycho(CDIModel):
# Needs to be updated to allow for plotting to an existing figure
plot_list = [
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',
CDTools/tools/data/data.py
+3 -3
View File
@@ -66,9 +66,9 @@ def get_entry_info(cxi_file):
'program_name']
metadata = {attr: str(e1[attr][()].decode()) for attr in metadata_attrs
if attr in e1}
datetime_attrs = ['start_time',
'end_time']
for attr in datetime_attrs:
for attr in metadata_attrs:
if attr in e1:
try:
metadata[attr] = dateutil.parser.parse(str(e1[attr][()].decode()))
CDTools/tools/interactions/interactions.py
+1 -7
View File
@@ -627,13 +627,7 @@ def RPI_interaction(probe, obj):
pad1l = (probe.shape[-1] - obj.shape[-1])//2
pad1r = probe.shape[-1] - obj.shape[-1] - pad1l
if obj.dim() == 2:
fftobj = t.nn.functional.pad(fftobj, (pad1l, pad1r, pad0l, pad0r))
elif obj.dim() == 3:
fftobj = t.nn.functional.pad(
fftobj, (pad1l, pad1r, pad0l, pad0r, 0,0))
else:
raise NotImplementedError('RPI interaction with obj of dimension higher than 4 (including complex dimension) is not supported.')
fftobj = t.nn.functional.pad(fftobj, (pad1l, pad1r, pad0l, pad0r))
# Again, just an inverse FFT but with an fftshift
upsampled_obj = propagators.inverse_far_field(fftobj)
CDTools/tools/plotting/plotting.py
+4 -2
View File
@@ -82,7 +82,7 @@ def get_units_factor(units):
return factor
def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label=None, **kwargs):
def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label=None, interpolation=None, **kwargs):
"""Plots an image with a colorbar and on an appropriate spatial grid
If a figure is given explicitly, it will clear that existing figure and
@@ -112,6 +112,8 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
Default is 'viridis', the colormap to plot with
cmap_label : str
What to label the colorbar when plotting
interpolation : str
What interpolation to use for imshow
\\**kwargs
All other args are passed to fig.add_subplot(111, \\**kwargs)
@@ -169,7 +171,7 @@ def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m',
else:
extent=None
plt.imshow(to_plot, cmap = cmap, extent = extent)
plt.imshow(to_plot, cmap = cmap, extent = extent, interpolation=interpolation)
cbar = plt.colorbar()
if cmap_label is not None:
cbar.set_label(cmap_label)