mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-09 13:02:41 +02:00
Update plotting functions to deal better with images that have multiple layers
This commit is contained in:
Abe Levitan
+25
-4
@@ -39,6 +39,7 @@ import numpy as np
|
||||
import threading
|
||||
import queue
|
||||
import time
|
||||
import pytorch_warmup
|
||||
|
||||
__all__ = ['CDIModel']
|
||||
|
||||
@@ -105,7 +106,7 @@ class CDIModel(t.nn.Module):
|
||||
|
||||
def AD_optimize(self, iterations, data_loader, optimizer,\
|
||||
scheduler=None, regularization_factor=None, thread=True,
|
||||
calculation_width=10):
|
||||
calculation_width=10, warmup_scheduler=None):
|
||||
"""Runs a round of reconstruction using the provided optimizer
|
||||
|
||||
This is the basic automatic differentiation reconstruction tool
|
||||
@@ -135,7 +136,8 @@ class CDIModel(t.nn.Module):
|
||||
normalization = 0
|
||||
for inputs, patterns in data_loader:
|
||||
normalization += t.sum(patterns).cpu().numpy()
|
||||
|
||||
|
||||
|
||||
def run_iteration(stop_event=None):
|
||||
loss = 0
|
||||
N = 0
|
||||
@@ -173,8 +175,20 @@ class CDIModel(t.nn.Module):
|
||||
loss = self.regularizer(regularization_factor)
|
||||
loss.backward()
|
||||
return total_loss
|
||||
|
||||
|
||||
if warmup_scheduler is not None:
|
||||
old_lrs = [group['lr'] for group in
|
||||
optimizer.param_groups]
|
||||
warmup_scheduler.dampen()
|
||||
#print([group['lr'] for group in
|
||||
# optimizer.param_groups], end='\r')
|
||||
|
||||
loss += optimizer.step(closure).detach().cpu().numpy()
|
||||
|
||||
if warmup_scheduler is not None:
|
||||
for old_lr, group in zip(old_lrs, optimizer.param_groups):
|
||||
group['lr'] = old_lr
|
||||
|
||||
|
||||
loss /= normalization
|
||||
if scheduler is not None:
|
||||
@@ -201,6 +215,7 @@ class CDIModel(t.nn.Module):
|
||||
self.figs[0].canvas.start_event_loop(0.01)
|
||||
else:
|
||||
calc.join()
|
||||
|
||||
except KeyboardInterrupt as e:
|
||||
stop_event.set()
|
||||
print('\nAsking execution thread to stop cleanly - please be patient.')
|
||||
@@ -215,7 +230,7 @@ class CDIModel(t.nn.Module):
|
||||
def Adam_optimize(self, iterations, dataset, batch_size=15, lr=0.005,
|
||||
schedule=False, amsgrad=False, subset=None,
|
||||
regularization_factor=None, thread=True,
|
||||
calculation_width=10):
|
||||
calculation_width=10, warmup=False):
|
||||
"""Runs a round of reconstruction using the Adam optimizer
|
||||
|
||||
This is generally accepted to be the most robust algorithm for use
|
||||
@@ -266,8 +281,14 @@ class CDIModel(t.nn.Module):
|
||||
else:
|
||||
scheduler = None
|
||||
|
||||
if warmup:
|
||||
warmup_scheduler = pytorch_warmup.UntunedLinearWarmup(optimizer)
|
||||
else:
|
||||
warmup_scheduler = None
|
||||
|
||||
return self.AD_optimize(iterations, data_loader, optimizer,
|
||||
scheduler=scheduler,
|
||||
warmup_scheduler=warmup_scheduler,
|
||||
regularization_factor=regularization_factor,
|
||||
thread=thread,
|
||||
calculation_width=calculation_width)
|
||||
|
||||
@@ -30,7 +30,10 @@ class Multislice2DPtycho(CDIModel):
|
||||
bandlimit=None,
|
||||
subpixel=True,
|
||||
exponentiate_obj=True,
|
||||
fourier_probe=False, units='um'):
|
||||
fourier_probe=False,
|
||||
apodization=None,
|
||||
phase_only=False,
|
||||
units='um'):
|
||||
|
||||
super(Multislice2DPtycho,self).__init__()
|
||||
self.wavelength = t.Tensor([wavelength])
|
||||
@@ -56,6 +59,7 @@ class Multislice2DPtycho(CDIModel):
|
||||
self.exponentiate_obj = exponentiate_obj
|
||||
self.fourier_probe = fourier_probe
|
||||
self.units = units
|
||||
self.phase_only=phase_only
|
||||
|
||||
if mask is None:
|
||||
self.mask = mask
|
||||
@@ -73,7 +77,7 @@ class Multislice2DPtycho(CDIModel):
|
||||
/ self.probe_norm)
|
||||
|
||||
self.obj = t.nn.Parameter(obj_guess.to(t.float32))
|
||||
|
||||
|
||||
if background is None:
|
||||
if detector_slice is not None:
|
||||
background = 1e-6 * t.ones(self.probe[0][self.detector_slice].shape[:-1])
|
||||
@@ -97,6 +101,11 @@ class Multislice2DPtycho(CDIModel):
|
||||
|
||||
self.probe_fourier_support = t.Tensor(probe_fourier_support).to(t.float32)
|
||||
|
||||
if apodization is not None:
|
||||
self.apodization = t.Tensor(apodization).to(t.float32)
|
||||
else:
|
||||
self.apodization = None
|
||||
|
||||
self.oversampling = oversampling
|
||||
|
||||
spacing = np.linalg.norm(self.probe_basis,axis=0)
|
||||
@@ -108,7 +117,7 @@ class Multislice2DPtycho(CDIModel):
|
||||
|
||||
|
||||
@classmethod
|
||||
def from_dataset(cls, dataset, dz, nz, probe_convergence_radius, probe_size=None, padding=0, n_modes=1, 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):
|
||||
def from_dataset(cls, dataset, dz, nz, probe_convergence_radius, probe_size=None, padding=0, n_modes=1, 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, apodize_prop=False, phase_only=False):
|
||||
|
||||
wavelength = dataset.wavelength
|
||||
det_basis = dataset.detector_geometry['basis']
|
||||
@@ -157,7 +166,6 @@ class Multislice2DPtycho(CDIModel):
|
||||
surface_normal = outgoing_dir + np.array([0.,0.,1.])
|
||||
surface_normal /= np.linalg.norm(outgoing_dir)
|
||||
|
||||
|
||||
# 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)
|
||||
@@ -175,14 +183,17 @@ class Multislice2DPtycho(CDIModel):
|
||||
else:
|
||||
probe = tools.initializers.gaussian_probe(dataset, probe_basis, probe_shape, probe_size, propagation_distance=propagation_distance)
|
||||
|
||||
# Now we initialize all the subdominant probe modes
|
||||
probe_max = t.max(cmath.cabs(probe))
|
||||
if n_modes >=2:
|
||||
probe_stack = list(0.01*tools.initializers.generate_subdominant_modes(probe,n_modes-1,circular=False))
|
||||
#probe_stack = [0.01 * probe_max * t.rand(probe.shape,dtype=probe.dtype) for i in range(n_modes - 1)]
|
||||
probe = t.stack([probe,] + probe_stack)
|
||||
|
||||
# For a Fourier space probe
|
||||
if fourier_probe:
|
||||
probe = tools.propagators.far_field(probe)
|
||||
|
||||
# Now we initialize all the subdominant probe modes
|
||||
probe_max = t.max(cmath.cabs(probe))
|
||||
probe_stack = [0.01 * probe_max * t.rand(probe.shape,dtype=probe.dtype) for i in range(n_modes - 1)]
|
||||
probe = t.stack([probe,] + probe_stack)
|
||||
|
||||
|
||||
# Consider a different start
|
||||
if exponentiate_obj:
|
||||
@@ -193,7 +204,17 @@ class Multislice2DPtycho(CDIModel):
|
||||
if not replicate_slice:
|
||||
obj = t.stack([obj]*nz)
|
||||
|
||||
|
||||
if apodize_prop:
|
||||
shape = probe.shape[-3:-1]
|
||||
#window_x = (1+np.cos(np.pi+2*np.pi*np.arange(shape[0])/shape[0]))/2
|
||||
#window_y = (1+np.cos(np.pi+2*np.pi*np.arange(shape[1])/shape[1]))/2
|
||||
window_x = np.cos(-np.pi/2+np.pi*np.arange(shape[0])/shape[0])
|
||||
window_y = np.cos(-np.pi/2+np.pi*np.arange(shape[1])/shape[1])
|
||||
Wx, Wy = np.meshgrid(window_x,window_y, indexing='ij')
|
||||
apodization=t.Tensor(Wx*Wy).to(dtype=probe.dtype)
|
||||
else:
|
||||
apodization=None
|
||||
|
||||
det_geo = dataset.detector_geometry
|
||||
|
||||
translation_offsets = 0 * (t.rand((len(dataset),2)) - 0.5)
|
||||
@@ -228,7 +249,9 @@ class Multislice2DPtycho(CDIModel):
|
||||
bandlimit=bandlimit,
|
||||
subpixel=subpixel,
|
||||
exponentiate_obj=exponentiate_obj,
|
||||
units=units, fourier_probe=fourier_probe)
|
||||
units=units, fourier_probe=fourier_probe,
|
||||
apodization=apodization,
|
||||
phase_only=phase_only)
|
||||
|
||||
|
||||
def interaction(self, index, translations):
|
||||
@@ -247,8 +270,12 @@ class Multislice2DPtycho(CDIModel):
|
||||
prs = self.probe*self.probe_fourier_support[None,:,:]
|
||||
# Here is where the mixing would happen, if it happened
|
||||
|
||||
|
||||
if self.exponentiate_obj:
|
||||
obj = cmath.cexpi(self.obj/self.nz)
|
||||
if self.phase_only:
|
||||
obj = cmath.expi(self.obj[...,0])
|
||||
else:
|
||||
obj = cmath.cexpi(self.obj)
|
||||
else:
|
||||
obj = self.obj
|
||||
|
||||
@@ -266,6 +293,10 @@ class Multislice2DPtycho(CDIModel):
|
||||
exit_waves = tools.interactions.ptycho_2D_round(
|
||||
exit_waves, obj, pix_trans,
|
||||
multiple_modes=True)
|
||||
if self.apodization is not None:
|
||||
exit_waves = exit_waves * self.apodization[None,:,:,None]
|
||||
|
||||
|
||||
|
||||
elif self.obj.dim() == 4:
|
||||
# If separate slices
|
||||
@@ -277,6 +308,8 @@ class Multislice2DPtycho(CDIModel):
|
||||
exit_waves = tools.interactions.ptycho_2D_round(
|
||||
exit_waves, obj[i], pix_trans,
|
||||
multiple_modes=True)
|
||||
if self.apodization is not None:
|
||||
exit_waves = exit_waves * self.apodization[None,:,:,None]
|
||||
|
||||
if i < self.nz-1: #on all but the last iteration
|
||||
exit_waves = tools.propagators.near_field(
|
||||
@@ -330,6 +363,8 @@ class Multislice2DPtycho(CDIModel):
|
||||
|
||||
if self.mask is not None:
|
||||
self.mask = self.mask.to(*args, **kwargs)
|
||||
if self.apodization is not None:
|
||||
self.apodization = self.apodization.to(*args,**kwargs)
|
||||
|
||||
|
||||
self.min_translation = self.min_translation.to(*args,**kwargs)
|
||||
@@ -338,6 +373,7 @@ class Multislice2DPtycho(CDIModel):
|
||||
#self.probe_support = self.probe_support.to(*args,**kwargs)
|
||||
self.probe_fourier_support = self.probe_fourier_support.to(*args,**kwargs)
|
||||
|
||||
|
||||
self.surface_normal = self.surface_normal.to(*args, **kwargs)
|
||||
self.as_prop = self.as_prop.to(*args, **kwargs)
|
||||
|
||||
@@ -396,13 +432,13 @@ class Multislice2DPtycho(CDIModel):
|
||||
('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)),
|
||||
('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),
|
||||
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),
|
||||
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),
|
||||
('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),
|
||||
@@ -411,13 +447,13 @@ class Multislice2DPtycho(CDIModel):
|
||||
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),
|
||||
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(cmath.cmult, self.obj.detach().cpu()), fig=fig, basis=self.probe_basis, units=self.units),
|
||||
lambda self: not self.exponentiate_obj),
|
||||
('Phase of Stacked Object Function',
|
||||
lambda self, fig: p.plot_phase(reduce(cmath.cmult, self.obj.detach().cpu()), fig=fig, basis=self.probe_basis, units=self.units),
|
||||
lambda self, fig: p.plot_phase(reduce(cmath.cmult, self.obj.detach().cpu()), fig=fig, basis=self.probe_basis, units=self.units,cmap='cividis'),
|
||||
lambda self: not self.exponentiate_obj),
|
||||
('Corrected Translations',
|
||||
lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)),
|
||||
|
||||
@@ -9,7 +9,8 @@ import numpy as np
|
||||
import torch as t
|
||||
|
||||
__all__ = ['exit_wave_geometry', 'calc_object_setup', 'gaussian',
|
||||
'gaussian_probe', 'SHARP_style_probe', 'RPI_spectral_init']
|
||||
'gaussian_probe', 'SHARP_style_probe', 'RPI_spectral_init',
|
||||
'generate_subdominant_modes']
|
||||
|
||||
from CDTools.tools import cmath
|
||||
from CDTools.tools.propagators import *
|
||||
@@ -18,6 +19,7 @@ from scipy.fftpack import next_fast_len
|
||||
from scipy.sparse import linalg as spla
|
||||
from torch.nn.functional import pad
|
||||
import numpy as np
|
||||
from functools import *
|
||||
|
||||
|
||||
def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None, opt_for_fft=True, padding=0, oversampling=1):
|
||||
@@ -438,3 +440,79 @@ def RPI_spectral_init(pattern, probe, obj_shape, n_modes=1, mask=None, backgroun
|
||||
|
||||
return cmath.complex_to_torch(z0).to(dtype=t.float32)
|
||||
|
||||
|
||||
def generate_subdominant_modes(dominant_mode, n_modes, circular=True):
|
||||
"""Generates guesses of subdominant modes based on spatial derivatives
|
||||
|
||||
The idea here is that vibration is an extremely common cause of
|
||||
spatial incoherence. This typically presents as subdominant modes that
|
||||
look like the derivative of the dominant mode along various axes.
|
||||
|
||||
Therefore, it is a good starting guess if the guess of the dominant
|
||||
mode is reasonable. The output probes will all be normalized to have
|
||||
the same intensity as the input dominant probe
|
||||
|
||||
Parameters
|
||||
----------
|
||||
dominant_mode : array
|
||||
The complex-valued dominant mode to work from
|
||||
n_modes : int
|
||||
The number of additional modes to create
|
||||
circular : bool
|
||||
Default True, whether to use circular modes (x+iy) or linear (x,y)
|
||||
Returns
|
||||
-------
|
||||
torch.Tensor
|
||||
The complex-style tensor storing the probe guesses
|
||||
"""
|
||||
|
||||
# This generates a list of tuples, where each tuple (a,b) corresponds
|
||||
# a term (kx^a)*(ky^b), or (kx+iky)^a*(kx-iky)^b for circular modes
|
||||
def make_orders(n_orders):
|
||||
if n_orders >1024:
|
||||
raise KeyError('Are you sure you want this many orders?')
|
||||
i = 1
|
||||
n = 0
|
||||
total = 0
|
||||
while True:
|
||||
if total >= n_orders:
|
||||
break
|
||||
|
||||
yield (n,i-n)
|
||||
total += 1
|
||||
if n<i:
|
||||
n+=1
|
||||
else:
|
||||
n = 0
|
||||
i += 1
|
||||
|
||||
# Gotta check and convert to pytorch if it's numpy
|
||||
|
||||
# First step is to take the FFT of the probe
|
||||
dominant_fft = far_field(dominant_mode)
|
||||
|
||||
shape = dominant_mode.shape
|
||||
center = ((shape[-3]-1)//2, (shape[-2]-1)//2)
|
||||
|
||||
i, j = np.mgrid[:shape[-3], :shape[-2]]
|
||||
i = t.Tensor(i - center[0]).to(dtype=dominant_fft.dtype,
|
||||
device=dominant_fft.device)
|
||||
j = t.Tensor(j - center[1]).to(dtype=dominant_fft.dtype,
|
||||
device=dominant_fft.device)
|
||||
|
||||
if circular:
|
||||
a = t.stack((i,j),dim=-1)
|
||||
b = t.stack((i,-j),dim=-1)
|
||||
else:
|
||||
a = t.stack((i,t.zeros_like(i)),dim=-1)
|
||||
b = t.stack((j,t.zeros_like(j)),dim=-1)
|
||||
|
||||
probe_norm = t.sum(cmath.cabssq(dominant_mode))
|
||||
# Then we need to multiply that FFT by various phase masks and IFFT
|
||||
probes = []
|
||||
for a_order,b_order in make_orders(n_modes):
|
||||
mask = reduce(cmath.cmult,[a]*a_order+[b]*b_order)
|
||||
new_probe = inverse_far_field(cmath.cmult(mask,dominant_fft))
|
||||
probes.append(new_probe * (probe_norm / t.sum(cmath.cabs(new_probe))))
|
||||
|
||||
return t.stack(probes)
|
||||
|
||||
@@ -100,7 +100,7 @@ def intensity_mse(intensities, sim_intensities, mask=None):
|
||||
|
||||
|
||||
|
||||
def poisson_nll(intensities, sim_intensities, mask=None, eps=1e-6):
|
||||
def poisson_nll(intensities, sim_intensities, mask=None, eps=1e-6, subtract_min=False):
|
||||
""" Returns the Poisson negative log likelihood for a simulated dataset's intensities
|
||||
|
||||
Calculates the overall Poisson maximum likelihood metric using
|
||||
@@ -137,13 +137,39 @@ def poisson_nll(intensities, sim_intensities, mask=None, eps=1e-6):
|
||||
A single value for the poisson negative log likelihood
|
||||
|
||||
"""
|
||||
#When x.logy gets into the regular build, add it by uncommenting!
|
||||
if mask is None:
|
||||
return t.sum(sim_intensities+eps -
|
||||
(intensities+eps) * t.log(sim_intensities+eps)) \
|
||||
/ intensities.view(-1).shape[0]
|
||||
|
||||
|
||||
nll = t.sum(sim_intensities+eps -
|
||||
intensities * t.log(sim_intensities+eps)) \
|
||||
/ intensities.view(-1).shape[0]
|
||||
#nll = t.sum(sim_intensities+epsa -
|
||||
# t.xlogy(intensities,sim_intensities+eps)) \
|
||||
# / intensities.view(-1).shape[0]
|
||||
if subtract_min:
|
||||
nll -= t.nansum(intensities - intensities*t.log(intensities))\
|
||||
/ intensities.view(-1).shape[0]
|
||||
#nll -= t.sum(intensities - t.xlogy(intensities,intensities))
|
||||
# We don't need to include the log factorial part here, because
|
||||
# it will get subtracted off in the min anyway.
|
||||
|
||||
return nll
|
||||
else:
|
||||
masked_intensities = intensities.masked_select(mask)
|
||||
masked_sims = sim_intensities.masked_select(mask)
|
||||
return t.sum(masked_sims - masked_intensities *
|
||||
t.log(masked_sims+eps)) / masked_intensities.shape[0]
|
||||
|
||||
nll = t.sum(masked_sims + eps - \
|
||||
masked_intensities * t.log(masked_sims+eps)) \
|
||||
/ masked_intensities.shape[0]
|
||||
|
||||
#nll = t.sum(masked_sims + eps - \
|
||||
# t.xlogy(masked_intensities, masked_sims+eps)) \
|
||||
# / masked_intensities.shape[0]
|
||||
|
||||
|
||||
if subtract_min:
|
||||
nll -= t.nansum(masked_intensities -
|
||||
masked_intensities*t.log(masked_intensities)) \
|
||||
/ masked_intensities.shape[0]
|
||||
|
||||
return nll
|
||||
|
||||
@@ -413,7 +413,7 @@ def plot_colorized(im, fig=None, basis=None, units='$\\mu$m', **kwargs):
|
||||
"""
|
||||
plot_func = lambda x: colorize(x)
|
||||
return plot_image(im, plot_func=plot_func, fig=fig, basis=basis,
|
||||
units=units, cmap=cmap, **kwargs)
|
||||
units=units, **kwargs)
|
||||
|
||||
|
||||
def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, **kwargs):
|
||||
|
||||
Reference in New Issue
Block a user