cSAXS/cdtools
0
0 Fork 0
mirror of https://github.com/cdtools-developers/cdtools.git synced 2026-11-12 17:23:54 +01:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'complex_numbers'

Browse Source 
BBringing in complex numbers!

B
B
This commit is contained in:
Abe Levitan committed 2021-07-06 14:56:28 -04:00
2 parents 5fdfa969cd 3e72243a50
commit 26b144d95a
78 files changed
+2962 -4121
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

CDTools/datasets/base.py
+1 -1
View File
@@ -121,7 +121,7 @@ class CDataset(torchdata.Dataset):
self.mask = self.mask.to(*args,**mask_kwargs)
if self.background is not None:
self.background = self.background.to(*args,**kwargs)
def get_as(self, *args, **kwargs):
"""Sets the dataset to return data on the given device and dtype
CDTools/datasets/ptycho_2d_dataset.py
+10 -152
View File
@@ -198,169 +198,27 @@ class Ptycho2DDataset(CDataset):
position.
"""
# We start by making the figure and axes
fig, axes = plt.subplots(1,2,figsize=(8,5.3))
fig.tight_layout(rect=[0.04, 0.09, 0.98, 0.96])
axslider = plt.axes([0.15,0.06,0.75,0.03])
#
# Then we define some helper functions for getting the right data
# that are used both in the initial setup and the updates
#
def get_data(idx):
def get_images(idx):
inputs, output = self[idx]
meas_data = output.detach().cpu().numpy()
if hasattr(self, 'mask') and self.mask is not None:
mask = self.mask.detach().cpu().numpy()
else:
mask = 1
return mask, meas_data
def calculate_sizes(idx):
bbox = axes[0].get_window_extent().transformed(fig.dpi_scale_trans.inverted())
s0 = bbox.width * bbox.height / translations.shape[0] * 72**2 #72 is points per inch
s0 /= 4 # A rough value to make the size work out
s = np.ones(len(self)) * s0
s[idx] *= 4
return s
def update_colorbar(im):
#
# This solves the problem of the colorbar being changed
# when the forward and back buttons are used!!!
#
if hasattr(im, 'norecurse') and im.norecurse:
im.norecurse=False
return
im.norecurse=True
# This is needed to update the colorbar
# only change limits if array contains multiple values
if np.min(im.get_array()) != np.max(im.get_array()):
im.set_clim(vmin=np.min(im.get_array()),
vmax=np.max(im.get_array()))
#
# The meatiest part of this program, here we just go through and
# set up the plot how we want it
#
# First we set up the left-hand plot, which shows an overview map
axes[0].set_title('Relative Displacement Map')
if logarithmic:
return np.log(meas_data) / np.log(10) * mask
else:
return meas_data * mask
translations = self.translations.detach().cpu().numpy()
nanomap_values = (self.mask.to(t.float32) * self.patterns).sum(dim=(1,2)).detach().cpu().numpy()
s = calculate_sizes(0)
units_factor = plotting.get_units_factor(units)
nanomap = axes[0].scatter(units_factor * translations[:,0],units_factor * translations[:,1],s=s,c=nanomap_values, picker=True)
axes[0].invert_xaxis()
axes[0].set_facecolor('k')
axes[0].set_xlabel('Translation x ('+units+')', labelpad=1)
axes[0].set_ylabel('Translation y ('+units+')', labelpad=1)
cb1 = plt.colorbar(nanomap, ax=axes[0], orientation='horizontal',
format='%.2e',
ticks=ticker.LinearLocator(numticks=5),
pad=0.17,fraction=0.1)
cb1.ax.set_title('Integrated Intensity', size="medium", pad=5)
cb1.ax.tick_params(labelrotation=20)
# Now we set up the second plot, which shows the individual
# diffraction patterns
axes[1].set_title('Diffraction Pattern')
mask, meas_data = get_data(0)
if logarithmic:
meas = axes[1].imshow(np.log(meas_data) / np.log(10) * mask)
cbar_title='Log Base 10 of Diffraction Intensity'
else:
meas = axes[1].imshow(meas_data * mask)
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)
cb2 = plt.colorbar(meas, ax=axes[1], orientation='horizontal',
format='%.2e',
ticks=ticker.LinearLocator(numticks=5),
pad=0.17,fraction=0.1)
cb2.ax.tick_params(labelrotation=20)
cb2.ax.set_title('Pixel Intensity', size="medium", pad=5)
cb2.ax.callbacks.connect('xlim_changed', lambda ax: update_colorbar(meas))
# This function handles all the updating, except for moving the
# slider value. This is done because the slider widget is
# ultimately responsible for triggering an update, so all other
# updates are done by changing the slider widget value which
# then triggers this
def update(idx):
# We have to explicitly make it an integer because the slider will
# output floats (even if they are still integer-valued)
idx = int(idx)
# Get the new data for this index
mask, meas_data = get_data(idx)
# Now we resize the nanomap to show the new selection
axes[0].collections[0].set_sizes(calculate_sizes(idx))
# And we update the data in the image as well
meas = axes[1].images[-1]
if logarithmic:
meas.set_data(np.log(meas_data) / np.log(10) * mask)
else:
meas.set_data(meas_data * mask)
update_colorbar(meas)
#
# Now we define the functions to handle various kinds of events
# that can be thrown our way
#
# We start by creating the slider here, so it can be used
# by the update hooks.
slider = Slider(axslider, 'Pattern #', 0, len(self)-1, valstep=1, valfmt="%d")
# This handles scroll wheel and keypress events
def on_action(event):
# Otherwise the if statements can throw errors when the
# event type isn't right, this way they just don't trigger
if not hasattr(event, 'button'):
event.button = None
if not hasattr(event, 'key'):
event.key = None
if event.key == 'up' or event.button == 'up' or event.key == 'right':
idx = slider.val - 1
elif event.key == 'down' or event.button == 'down' or event.key == 'left':
idx = slider.val + 1
else:
# This prevents errors from being thrown on irrelevant key
# or mouse input
return
# Handle the wraparound and trigger the update
idx = int(idx) % len(self)
slider.set_val(idx)
# This handles "pick" events in the nanomap
def on_pick(event):
# If we don't filter on type of event, this will also capture,
# for example, scroll events that happen over the nanomap
if event.mouseevent.button == 1:
slider.set_val(event.ind[0])
# Here we connect the various update functions
fig.canvas.mpl_connect('pick_event',on_pick)
fig.canvas.mpl_connect('key_press_event',on_action)
fig.canvas.mpl_connect('scroll_event',on_action)
slider.on_changed(update)
# Throw an extra update into the mix just to get rid of any things
# (like the nanomap dot sizes) that otherwise would change on the
# first update
update(0)
CDTools/models/__init__.py
+3 -7
View File
@@ -20,19 +20,15 @@ defining a new ptychography model before attempting to do so.
"""
from __future__ import division, print_function, absolute_import
# I don't believe that __all__ really needed, but it's nice to define it
# to be explicit that import * is safe
__all__ = ['CDIModel', 'SimplePtycho', 'FancyPtycho', 'Bragg2DPtycho', 'SMatrixPtycho', 'RPI', 'UnifiedModePtycho']
#__all__ = ['CDIModel', 'SimplePtycho', 'FancyPtycho', 'Bragg2DPtycho', 'SMatrixPtycho', 'RPI']
from CDTools.models.base import CDIModel
from CDTools.models.simple_ptycho import SimplePtycho
from CDTools.models.fancy_ptycho import FancyPtycho
from CDTools.models.pinhole_plane_ptycho import PinholePlanePtycho
from CDTools.models.bragg_2d_ptycho import Bragg2DPtycho
from CDTools.models.s_matrix_ptycho import SMatrixPtycho
from CDTools.models.multislice_2d_ptycho import Multislice2DPtycho
from CDTools.models.rpi import RPI
from CDTools.models.unified_mode_ptycho import UnifiedModePtycho
from CDTools.models.unified_mode_ptycho2 import UnifiedModePtycho2
#from CDTools.models.s_matrix_ptycho import SMatrixPtycho
CDTools/models/base.py
+119 -39
View File
@@ -39,6 +39,9 @@ import numpy as np
import threading
import queue
import time
#import pytorch_warmup
from .complex_adam import MyAdam
from .complex_lbfgs import MyLBFGS
__all__ = ['CDIModel']
@@ -53,6 +56,10 @@ class CDIModel(t.nn.Module):
functions.
"""
def __init__(self):
super(CDIModel,self).__init__()
self.iteration_count = 0
def from_dataset(self, dataset):
raise NotImplementedError()
@@ -100,7 +107,8 @@ class CDIModel(t.nn.Module):
raise NotImplementedError()
def AD_optimize(self, iterations, data_loader, optimizer,\
scheduler=None, regularization_factor=None, thread=True):
scheduler=None, regularization_factor=None, thread=True,
calculation_width=10, warmup_scheduler=None):
"""Runs a round of reconstruction using the provided optimizer
This is the basic automatic differentiation reconstruction tool
@@ -123,41 +131,85 @@ class CDIModel(t.nn.Module):
Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method
thread : bool
Default True, whether to run the computation in a separate thread to allow interaction with plots during computation
calculation_width : int
Default 10, how many translations to pass through at once for each round of gradient accumulation
"""
# First, calculate the normalization
normalization = 0
for inputs, patterns in data_loader:
normalization += t.sum(patterns).cpu().numpy()
def run_iteration(stop_event=None):
loss = 0
N = 0
t0 = time.time()
for inputs, patterns in data_loader:
N += 1
def closure():
# This is just used to allow graceful exit when threading
if stop_event is not None and stop_event.is_set():
exit()
optimizer.zero_grad()
sim_patterns = self.forward(*inputs)
if hasattr(self, 'mask'):
loss = self.loss(patterns,sim_patterns, mask=self.mask)
else:
loss = self.loss(patterns,sim_patterns)
input_chunks = [[inp[i:i + calculation_width]
for inp in inputs]
for i in range(0, len(inputs[0]),
calculation_width)]
pattern_chunks = [patterns[i:i + calculation_width]
for i in range(0, len(inputs[0]),
calculation_width)]
total_loss = 0
for inp, pats in zip(input_chunks, pattern_chunks):
# This is just used to allow graceful exit when threading
if stop_event is not None and stop_event.is_set():
exit()
sim_patterns = self.forward(*inp)
#sim_patterns.retain_grad()
if hasattr(self, 'mask'):
loss = self.loss(pats,sim_patterns, mask=self.mask)
else:
loss = self.loss(pats,sim_patterns)
loss.backward()#retain_variables=True)
#plt.figure()
#plt.imshow(sim_patterns.grad.cpu()[0])
#plt.colorbar()
#plt.show()
total_loss += loss.detach()
#print('probe grad')
#print(t.mean(self.obj.grad))
if regularization_factor is not None \
and hasattr(self, 'regularizer'):
loss += self.regularizer(regularization_factor)
#print(loss)
loss.backward()
return loss
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:
scheduler.step(loss)
self.latest_loss = loss
self.latest_iteration_time = time.time() - t0
self.iteration_count += 1
return loss
if thread:
@@ -176,6 +228,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.')
@@ -189,7 +242,8 @@ 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):
regularization_factor=None, thread=True,
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
@@ -215,6 +269,9 @@ class CDIModel(t.nn.Module):
Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method
thread : bool
Default True, whether to run the computation in a separate thread to allow interaction with plots during computation
calculation_width : int
Default 1, how many translations to pass through at once for each round of gradient accumulation
"""
if subset is not None:
@@ -228,7 +285,8 @@ class CDIModel(t.nn.Module):
shuffle=True)
# Define the optimizer
optimizer = t.optim.Adam(self.parameters(), lr = lr, amsgrad=amsgrad)
#optimizer = t.optim.Adam(self.parameters(), lr = lr, amsgrad=amsgrad)
optimizer = MyAdam(self.parameters(), lr = lr, amsgrad=amsgrad)
# Define the scheduler
@@ -237,21 +295,33 @@ class CDIModel(t.nn.Module):
else:
scheduler = None
if warmup:
print('Warmup is not currently implemented, sorry!')
#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)
thread=thread,
calculation_width=calculation_width)
def LBFGS_optimize(self, iterations, dataset, batch_size=None,
def LBFGS_optimize(self, iterations, dataset,
lr=0.1,history_size=2, subset=None,
regularization_factor=None, thread=True):
regularization_factor=None, thread=True,
calculation_width=10):
"""Runs a round of reconstruction using the L-BFGS optimizer
This algorithm is often less stable that Adam, however in certain
situations or geometries it can be shockingly efficient. Like all
the other optimization routines, it is defined as a generator
function which yields the average loss each epoch.
Note: There is no batch size, because it is a usually a bad idea to use
LBFGS on anything but all the data at onece
Parameters
----------
@@ -259,8 +329,6 @@ class CDIModel(t.nn.Module):
How many epochs of the algorithm to run
dataset : CDataset
The dataset to reconstruct against
batch_size : int
Optional, the size of the minibatches to use
lr : float
Optional, the learning rate to use
history_size : int
@@ -271,6 +339,7 @@ class CDIModel(t.nn.Module):
Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method
thread : bool
Default True, whether to run the computation in a separate thread to allow interaction with plots during computation
"""
if subset is not None:
# if just one pattern, turn into a list for convenience
@@ -278,27 +347,27 @@ class CDIModel(t.nn.Module):
subset = [subset]
dataset = torchdata.Subset(dataset, subset)
# Make a dataloader
if batch_size is not None:
data_loader = torchdata.DataLoader(dataset, batch_size=batch_size,
shuffle=True)
else:
data_loader = torchdata.DataLoader(dataset)
# Make a dataloader. This basically does nothing but load all the
# data at once
data_loader = torchdata.DataLoader(dataset, batch_size=len(dataset))
# Define the optimizer
optimizer = t.optim.LBFGS(self.parameters(),
lr = lr, history_size=history_size)
#optimizer = MyLBFGS(self.parameters(),
# lr = lr, history_size=history_size)
return self.AD_optimize(iterations, data_loader, optimizer,
regularization_factor=regularization_factor,
thread=thread)
thread=thread,
calculation_width=calculation_width)
def SGD_optimize(self, iterations, dataset, batch_size=None,
lr=0.01, momentum=0, dampening=0, weight_decay=0,
nesterov=False, subset=None, regularization_factor=None,
thread=True):
thread=True, calculation_width=10):
"""Runs a round of reconstruction using the SGDoptimizer
This algorithm is often less stable that Adam, but it is simpler
@@ -322,6 +391,9 @@ class CDIModel(t.nn.Module):
Optional, if the model has a regularizer defined, the set of parameters to pass the regularizer method
thread : bool
Default True, whether to run the computation in a separate thread to allow interaction with plots during computation
calculation_width : int
Default 1, how many translations to pass through at once for each round of gradient accumulation
"""
if subset is not None:
@@ -347,9 +419,19 @@ class CDIModel(t.nn.Module):
return self.AD_optimize(iterations, data_loader, optimizer,
regularization_factor=regularization_factor,
thread=thread)
thread=thread,
calculation_width=calculation_width)
def report(self):
"""Returns a string informing on the latest reconstruction iteration"""
if hasattr(self, 'latest_loss'):
return 'Iteration ' + str(self.iteration_count) + \
' completed in %0.2f s with loss ' %\
self.latest_iteration_time + str(self.latest_loss)
else:
return 'No reconstruction iterations performed yet!'
# By default, the plot_list is empty
plot_list = []
@@ -407,7 +489,7 @@ class CDIModel(t.nn.Module):
plotter = plots[1]
if figs is None:
fig = plt.figure()
fig = plt.figure()
self.figs.append(fig)
else:
fig = figs[idx]
@@ -415,6 +497,7 @@ class CDIModel(t.nn.Module):
try:
plotter(self,fig)
plt.title(name)
except TypeError as e:
if dataset is not None:
try:
@@ -448,7 +531,7 @@ class CDIModel(t.nn.Module):
fig, axes = plt.subplots(1,3,figsize=(12,5.3))
fig.tight_layout(rect=[0.02, 0.09, 0.98, 0.96])
axslider = plt.axes([0.15,0.06,0.75,0.03])
def update_colorbar(im):
# If the update brought the colorbar out of whack
@@ -463,12 +546,8 @@ class CDIModel(t.nn.Module):
return
im.norecurse=True
im.colorbar.set_clim(vmin=np.min(im.get_array()),vmax=np.max(im.get_array()))
im.colorbar.ax.set_ylim(0,1)
im.colorbar.set_ticks(ticker.LinearLocator(numticks=5))
im.colorbar.draw_all()
im.set_clim(vmin=np.min(im.get_array()),vmax=np.max(im.get_array()))
def update(idx):
idx = int(idx) % len(dataset)
fig.pattern_idx = idx
@@ -503,6 +582,7 @@ class CDIModel(t.nn.Module):
cb3.ax.callbacks.connect('xlim_changed', lambda ax: update_colorbar(diff))
else:
sim = axes[0].images[-1]
sim.set_data(sim_data)
update_colorbar(sim)
CDTools/models/bragg_2d_ptycho.py
+43 -86
View File
@@ -4,7 +4,6 @@ import torch as t
from CDTools.models import CDIModel
from CDTools.datasets import Ptycho2DDataset
from CDTools import tools
from CDTools.tools import cmath
from CDTools.tools import plotting as p
from CDTools.tools.propagators import generate_generalized_angular_spectrum_propagator as ggasp
from matplotlib import pyplot as plt
@@ -110,9 +109,9 @@ class Bragg2DPtycho(CDIModel):
# We rescale the probe here so it learns at the same rate as the
# object
if probe_guess.dim() > 3:
self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess[0].to(t.float32)))
self.probe_norm = 1 * t.max(t.abs(probe_guess[0]).to(t.float32))
else:
self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess.to(t.float32)))
self.probe_norm = 1 * t.max(t.abs(probe_guess).to(t.float32))
# Not strictly necessary but otherwise it will return
# a probe with the stuff outside of the support unchanged after
@@ -124,42 +123,44 @@ class Bragg2DPtycho(CDIModel):
# you look at the phase map
probe_guess[probe_guess == 0] = 0
self.probe = t.nn.Parameter(probe_guess.to(t.float32)
self.probe = t.nn.Parameter(probe_guess.to(t.complex64)
/ self.probe_norm)
self.obj = t.nn.Parameter(obj_guess.to(t.float32))
self.obj = t.nn.Parameter(obj_guess.to(t.complex64))
if background is None:
if detector_slice is not None:
background = 1e-6 * t.ones(self.probe[0][self.detector_slice].shape[:-1])
background = 1e-6 * t.ones(self.probe[0][self.detector_slice].shape)
else:
background = 1e-6 * t.ones(self.probe[0].shape[:-1])
background = 1e-6 * t.ones(self.probe[0].shape)
self.background = t.nn.Parameter(t.Tensor(background).to(t.float32))
self.background = t.nn.Parameter(t.as_tensor(background,dtype=t.float32))
if weights is None:
self.weights = None
else:
self.weights = t.nn.Parameter(t.Tensor(weights).to(t.float32))
self.weights = t.nn.Parameter(t.as_tensor(weights,dtype=t.float32))
if translation_offsets is None:
self.translation_offsets = None
else:
self.translation_offsets = t.nn.Parameter(t.Tensor(translation_offsets).to(t.float32)/ translation_scale)
self.translation_offsets = t.nn.Parameter(
t.as_tensor(translation_offsets,dtype=t.float32) /
translation_scale)
self.translation_scale = translation_scale
if probe_support is not None:
self.probe_support = probe_support
else:
self.probe_support = t.ones_like(self.probe[0])
self.probe_support = t.ones(self.probe[0].shape,dtype=t.bool)
if obj_support is not None:
self.obj_support = obj_support
self.obj.data = self.obj * obj_support
else:
self.obj_support = t.ones_like(self.obj)
self.obj_support = t.ones_like(self.obj, dtype=t.bool)
self.oversampling = oversampling
@@ -170,7 +171,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].shape[:-1]
probe_shape = self.probe[0]
self.k_map, self.intensity_map = \
tools.propagators.generate_high_NA_k_intensity_map(
@@ -188,11 +189,11 @@ class Bragg2DPtycho(CDIModel):
# This propagator should be able to be multiplied by the propagation
# distance each time to get a propagator
self.universal_propagator = cmath.cphase(ggasp(self.probe.shape[1:],
self.universal_propagator = t.angle(ggasp(self.probe.shape[1:],
self.probe_basis, self.wavelength,
t.Tensor([0,0,self.wavelength/(2*np.pi)]),
propagation_vector=self.prop_dir,
dtype=t.float32,
dtype=t.complex64,
propagate_along_offset=True))
@@ -299,11 +300,11 @@ class Bragg2DPtycho(CDIModel):
# Now we initialize all the subdominant probe modes
probe_max = t.max(cmath.cabs(probe))
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)]
probe = t.stack([probe,] + probe_stack)
obj = tools.cmath.expi(randomize_ang * (t.rand(obj_size)-0.5))
obj = t.exp(1j*(randomize_ang * (t.rand(obj_size)-0.5)))
det_geo = dataset.detector_geometry
@@ -317,8 +318,8 @@ class Bragg2DPtycho(CDIModel):
mask = None
if probe_support_radius is not None:
probe_support = t.zeros_like(probe[0].to(dtype=t.float32))
p_cent = np.array(probe.shape[1:3]).astype(int) // 2
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
@@ -329,7 +330,7 @@ class Bragg2DPtycho(CDIModel):
ro = restrict_obj
os = np.array(obj_size)
ps = np.array(probe_shape)
obj_support = t.zeros_like(obj.to(dtype=t.float32))
obj_support = t.zeros(obj.shape,dtype=t.bool)
obj_support[ps[0]//2-ro:os[0]+ro-ps[0]//2,
ps[1]//2-ro:os[1]+ro-ps[1]//2] = 1
else:
@@ -368,53 +369,29 @@ class Bragg2DPtycho(CDIModel):
pix_trans -= self.min_translation
props -= self.median_propagation
if self.translation_offsets is not None:
pix_trans += self.translation_scale * self.translation_offsets[index]
single_translation = False
if translations.dim() == 1:
translations = translations[None,:]
pix_trans = pix_trans[None,:]
single_translation = True
all_exit_waves = []
for i in range(self.probe.shape[0]):
pr = self.probe[i] * self.probe_support
exit_waves = []
for j in range(translations.size()[0]):
if self.propagate_probe:
# Minus sign is empirical
propagator = cmath.expi(
(-1*props[j]*(2*np.pi)/self.wavelength)
* self.universal_propagator)
prop_pr = tools.propagators.near_field(pr, propagator)
#plt.close('all')
#plt.imshow(np.abs(cmath.torch_to_complex(prop_pr.detach().cpu())))
#plt.show()
else:
prop_pr = pr
exit_waves.append(self.probe_norm *
tools.interactions.ptycho_2D_sinc(prop_pr,
self.obj_support * self.obj,
pix_trans[j],
shift_probe=True))
exit_waves = t.stack(exit_waves)
if not single_translation:
exit_waves = self.weights[index][:,None,None,None] * exit_waves
else:
exit_waves = self.weights[index] * exit_waves[0,:,:,:]
all_exit_waves.append(exit_waves)
Ws = self.weights[index]
prs = Ws[...,None,None,None] * self.probe
# Now we need to propagate each of the probes
return t.stack(all_exit_waves)
for j in range(prs.shape[0]):
# I believe this -1 sign is in error, but I need a dataset with
# well understood geometry to figure it out
propagator = t.exp(
1j*(props[j]*(2*np.pi)/self.wavelength)
* self.universal_propagator)
prs[j] = tools.propagators.near_field(prs[j], propagator)
exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc(
prs, self.obj_support * self.obj,pix_trans,
shift_probe=True, multiple_modes=True)
return exit_waves
def forward_propagator(self, wavefields):
if self.correct_tilt:
@@ -514,32 +491,12 @@ class Bragg2DPtycho(CDIModel):
def corrected_translations(self,dataset):
translations = dataset.translations.to(dtype=self.probe.dtype,device=self.probe.device)
translations = dataset.translations.to(dtype=self.probe.real.dtype,
device=self.probe.device)
t_offset = tools.interactions.pixel_to_translations(self.probe_basis,self.translation_offsets*self.translation_scale,surface_normal=self.surface_normal)
return translations + t_offset
# Needs to be updated to allow for plotting to an existing figure
# plot_list = [
# ('Dominant Probe Amplitude',
# lambda self, fig: p.plot_amplitude(self.probe[0], fig=fig, basis=self.probe_basis)),
# ('Dominant Probe Phase',
# lambda self, fig: p.plot_phase(self.probe[0], fig=fig, basis=self.probe_basis)),
# ('Subdominant Probe Amplitude',
# lambda self, fig: p.plot_amplitude(self.probe[1], fig=fig, basis=self.probe_basis),
# lambda self: len(self.probe) >=2),
# ('Subdominant Probe Phase',
# lambda self, fig: p.plot_phase(self.probe[1], fig=fig, basis=self.probe_basis),
# lambda self: len(self.probe) >=2),
# ('Object Amplitude',
# lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis)),
# ('Object Phase',
# lambda self, fig: p.plot_phase(self.obj, fig=fig, basis=self.probe_basis)),
# ('Corrected Translations',
# lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig)),
# ('Background',
# lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2))
# ]
plot_list = [
('Dominant Probe Amplitude',
lambda self, fig: p.plot_amplitude(self.probe[0], fig=fig)),
@@ -565,9 +522,9 @@ class Bragg2DPtycho(CDIModel):
def save_results(self, dataset):
basis = self.probe_basis.detach().cpu().numpy()
translations = self.corrected_translations(dataset).detach().cpu().numpy()
probe = cmath.torch_to_complex(self.probe.detach().cpu())
probe = self.probe.detach().cpu().numpy()
probe = probe * self.probe_norm.detach().cpu().numpy()
obj = cmath.torch_to_complex(self.obj.detach().cpu())
obj = self.obj.detach().cpu().numpy()
background = self.background.detach().cpu().numpy()**2
weights = self.weights.detach().cpu().numpy()
CDTools/models/complex_adam.py
+163
View File
@@ -0,0 +1,163 @@
import math
import torch
from torch import Tensor
from torch.optim.optimizer import Optimizer
from typing import List, Optional
class MyAdam(Optimizer):
r"""Implements Adam algorithm.
It has been proposed in `Adam: A Method for Stochastic Optimization`_.
The implementation of the L2 penalty follows changes proposed in
`Decoupled Weight Decay Regularization`_.
Args:
params (iterable): iterable of parameters to optimize or dicts defining
parameter groups
lr (float, optional): learning rate (default: 1e-3)
betas (Tuple[float, float], optional): coefficients used for computing
running averages of gradient and its square (default: (0.9, 0.999))
eps (float, optional): term added to the denominator to improve
numerical stability (default: 1e-8)
weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
amsgrad (boolean, optional): whether to use the AMSGrad variant of this
algorithm from the paper `On the Convergence of Adam and Beyond`_
(default: False)
.. _Adam\: A Method for Stochastic Optimization:
https://arxiv.org/abs/1412.6980
.. _Decoupled Weight Decay Regularization:
https://arxiv.org/abs/1711.05101
.. _On the Convergence of Adam and Beyond:
https://openreview.net/forum?id=ryQu7f-RZ
"""
def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
weight_decay=0, amsgrad=False):
if not 0.0 <= lr:
raise ValueError("Invalid learning rate: {}".format(lr))
if not 0.0 <= eps:
raise ValueError("Invalid epsilon value: {}".format(eps))
if not 0.0 <= betas[0] < 1.0:
raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
if not 0.0 <= betas[1] < 1.0:
raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
if not 0.0 <= weight_decay:
raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
defaults = dict(lr=lr, betas=betas, eps=eps,
weight_decay=weight_decay, amsgrad=amsgrad)
super(MyAdam, self).__init__(params, defaults)
def __setstate__(self, state):
super(MyAdam, self).__setstate__(state)
for group in self.param_groups:
group.setdefault('amsgrad', False)
@torch.no_grad()
def step(self, closure=None):
"""Performs a single optimization step.
Args:
closure (callable, optional): A closure that reevaluates the model
and returns the loss.
"""
loss = None
if closure is not None:
with torch.enable_grad():
loss = closure()
for group in self.param_groups:
params_with_grad = []
grads = []
exp_avgs = []
exp_avg_sqs = []
max_exp_avg_sqs = []
state_steps = []
beta1, beta2 = group['betas']
for p in group['params']:
if p.grad is not None:
params_with_grad.append(p)
if p.grad.is_sparse:
raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')
grads.append(p.grad)
state = self.state[p]
# Lazy state initialization
if len(state) == 0:
state['step'] = 0
# Exponential moving average of gradient values
state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
# Exponential moving average of squared gradient values
state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)
if group['amsgrad']:
# Maintains max of all exp. moving avg. of sq. grad. values
state['max_exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)
exp_avgs.append(state['exp_avg'])
exp_avg_sqs.append(state['exp_avg_sq'])
if group['amsgrad']:
max_exp_avg_sqs.append(state['max_exp_avg_sq'])
# update the steps for each param group update
state['step'] += 1
# record the step after step update
state_steps.append(state['step'])
adam(params_with_grad,
grads,
exp_avgs,
exp_avg_sqs,
max_exp_avg_sqs,
state_steps,
amsgrad=group['amsgrad'],
beta1=beta1,
beta2=beta2,
lr=group['lr'],
weight_decay=group['weight_decay'],
eps=group['eps'])
return loss
def adam(params: List[Tensor],
grads: List[Tensor],
exp_avgs: List[Tensor],
exp_avg_sqs: List[Tensor],
max_exp_avg_sqs: List[Tensor],
state_steps: List[int],
*,
amsgrad: bool,
beta1: float,
beta2: float,
lr: float,
weight_decay: float,
eps: float):
r"""Functional API that performs Adam algorithm computation.
See :class:`~torch.optim.Adam` for details.
"""
for i, param in enumerate(params):
grad = grads[i]
exp_avg = exp_avgs[i]
exp_avg_sq = exp_avg_sqs[i]
step = state_steps[i]
bias_correction1 = 1 - beta1 ** step
bias_correction2 = 1 - beta2 ** step
if weight_decay != 0:
grad = grad.add(param, alpha=weight_decay)
# Decay the first and second moment running average coefficient
exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
exp_avg_sq.mul_(beta2).addcmul_(grad, grad.conj(), value=1 - beta2)
if amsgrad:
# Maintains the maximum of all 2nd moment running avg. till now
torch.maximum(max_exp_avg_sqs[i], exp_avg_sq, out=max_exp_avg_sqs[i])
# Use the max. for normalizing running avg. of gradient
denom = (max_exp_avg_sqs[i].sqrt() / math.sqrt(bias_correction2)).add_(eps)
else:
denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)
step_size = lr / bias_correction1
param.addcdiv_(exp_avg, denom, value=-step_size)
CDTools/models/complex_lbfgs.py
+485
View File
@@ -0,0 +1,485 @@
import torch
from functools import reduce
from torch.optim.optimizer import Optimizer
def _cubic_interpolate(x1, f1, g1, x2, f2, g2, bounds=None):
# ported from https://github.com/torch/optim/blob/master/polyinterp.lua
# Compute bounds of interpolation area
if bounds is not None:
xmin_bound, xmax_bound = bounds
else:
xmin_bound, xmax_bound = (x1, x2) if x1 <= x2 else (x2, x1)
# Code for most common case: cubic interpolation of 2 points
# w/ function and derivative values for both
# Solution in this case (where x2 is the farthest point):
# d1 = g1 + g2 - 3*(f1-f2)/(x1-x2);
# d2 = sqrt(d1^2 - g1*g2);
# min_pos = x2 - (x2 - x1)*((g2 + d2 - d1)/(g2 - g1 + 2*d2));
# t_new = min(max(min_pos,xmin_bound),xmax_bound);
d1 = g1 + g2 - 3 * (f1 - f2) / (x1 - x2)
d2_square = d1**2 - g1 * g2
if d2_square >= 0:
d2 = d2_square.sqrt()
if x1 <= x2:
min_pos = x2 - (x2 - x1) * ((g2 + d2 - d1) / (g2 - g1 + 2 * d2))
else:
min_pos = x1 - (x1 - x2) * ((g1 + d2 - d1) / (g1 - g2 + 2 * d2))
return min(max(min_pos, xmin_bound), xmax_bound)
else:
return (xmin_bound + xmax_bound) / 2.
def _strong_wolfe(obj_func,
x,
t,
d,
f,
g,
gtd,
c1=1e-4,
c2=0.9,
tolerance_change=1e-9,
max_ls=25):
# ported from https://github.com/torch/optim/blob/master/lswolfe.lua
d_norm = d.abs().max()
g = g.clone(memory_format=torch.contiguous_format)
# evaluate objective and gradient using initial step
f_new, g_new = obj_func(x, t, d)
ls_func_evals = 1
gtd_new = g_new.dot(d)
# bracket an interval containing a point satisfying the Wolfe criteria
t_prev, f_prev, g_prev, gtd_prev = 0, f, g, gtd
done = False
ls_iter = 0
while ls_iter < max_ls:
# check conditions
if f_new > (f + c1 * t * gtd) or (ls_iter > 1 and f_new >= f_prev):
bracket = [t_prev, t]
bracket_f = [f_prev, f_new]
bracket_g = [g_prev, g_new.clone(memory_format=torch.contiguous_format)]
bracket_gtd = [gtd_prev, gtd_new]
break
if abs(gtd_new) <= -c2 * gtd:
bracket = [t]
bracket_f = [f_new]
bracket_g = [g_new]
done = True
break
if gtd_new >= 0:
bracket = [t_prev, t]
bracket_f = [f_prev, f_new]
bracket_g = [g_prev, g_new.clone(memory_format=torch.contiguous_format)]
bracket_gtd = [gtd_prev, gtd_new]
break
# interpolate
min_step = t + 0.01 * (t - t_prev)
max_step = t * 10
tmp = t
t = _cubic_interpolate(
t_prev,
f_prev,
gtd_prev,
t,
f_new,
gtd_new,
bounds=(min_step, max_step))
# next step
t_prev = tmp
f_prev = f_new
g_prev = g_new.clone(memory_format=torch.contiguous_format)
gtd_prev = gtd_new
f_new, g_new = obj_func(x, t, d)
ls_func_evals += 1
gtd_new = g_new.dot(d)
ls_iter += 1
# reached max number of iterations?
if ls_iter == max_ls:
bracket = [0, t]
bracket_f = [f, f_new]
bracket_g = [g, g_new]
# zoom phase: we now have a point satisfying the criteria, or
# a bracket around it. We refine the bracket until we find the
# exact point satisfying the criteria
insuf_progress = False
# find high and low points in bracket
low_pos, high_pos = (0, 1) if bracket_f[0] <= bracket_f[-1] else (1, 0)
while not done and ls_iter < max_ls:
# line-search bracket is so small
if abs(bracket[1] - bracket[0]) * d_norm < tolerance_change:
break
# compute new trial value
t = _cubic_interpolate(bracket[0], bracket_f[0], bracket_gtd[0],
bracket[1], bracket_f[1], bracket_gtd[1])
# test that we are making sufficient progress:
# in case `t` is so close to boundary, we mark that we are making
# insufficient progress, and if
# + we have made insufficient progress in the last step, or
# + `t` is at one of the boundary,
# we will move `t` to a position which is `0.1 * len(bracket)`
# away from the nearest boundary point.
eps = 0.1 * (max(bracket) - min(bracket))
if min(max(bracket) - t, t - min(bracket)) < eps:
# interpolation close to boundary
if insuf_progress or t >= max(bracket) or t <= min(bracket):
# evaluate at 0.1 away from boundary
if abs(t - max(bracket)) < abs(t - min(bracket)):
t = max(bracket) - eps
else:
t = min(bracket) + eps
insuf_progress = False
else:
insuf_progress = True
else:
insuf_progress = False
# Evaluate new point
f_new, g_new = obj_func(x, t, d)
ls_func_evals += 1
gtd_new = g_new.dot(d)
ls_iter += 1
if f_new > (f + c1 * t * gtd) or f_new >= bracket_f[low_pos]:
# Armijo condition not satisfied or not lower than lowest point
bracket[high_pos] = t
bracket_f[high_pos] = f_new
bracket_g[high_pos] = g_new.clone(memory_format=torch.contiguous_format)
bracket_gtd[high_pos] = gtd_new
low_pos, high_pos = (0, 1) if bracket_f[0] <= bracket_f[1] else (1, 0)
else:
if abs(gtd_new) <= -c2 * gtd:
# Wolfe conditions satisfied
done = True
elif gtd_new * (bracket[high_pos] - bracket[low_pos]) >= 0:
# old high becomes new low
bracket[high_pos] = bracket[low_pos]
bracket_f[high_pos] = bracket_f[low_pos]
bracket_g[high_pos] = bracket_g[low_pos]
bracket_gtd[high_pos] = bracket_gtd[low_pos]
# new point becomes new low
bracket[low_pos] = t
bracket_f[low_pos] = f_new
bracket_g[low_pos] = g_new.clone(memory_format=torch.contiguous_format)
bracket_gtd[low_pos] = gtd_new
# return stuff
t = bracket[low_pos]
f_new = bracket_f[low_pos]
g_new = bracket_g[low_pos]
return f_new, g_new, t, ls_func_evals
class MyLBFGS(Optimizer):
"""Implements L-BFGS algorithm, heavily inspired by `minFunc
<https://www.cs.ubc.ca/~schmidtm/Software/minFunc.html>`_.
.. warning::
This optimizer doesn't support per-parameter options and parameter
groups (there can be only one).
.. warning::
Right now all parameters have to be on a single device. This will be
improved in the future.
.. note::
This is a very memory intensive optimizer (it requires additional
``param_bytes * (history_size + 1)`` bytes). If it doesn't fit in memory
try reducing the history size, or use a different algorithm.
Args:
lr (float): learning rate (default: 1)
max_iter (int): maximal number of iterations per optimization step
(default: 20)
max_eval (int): maximal number of function evaluations per optimization
step (default: max_iter * 1.25).
tolerance_grad (float): termination tolerance on first order optimality
(default: 1e-5).
tolerance_change (float): termination tolerance on function
value/parameter changes (default: 1e-9).
history_size (int): update history size (default: 100).
line_search_fn (str): either 'strong_wolfe' or None (default: None).
"""
def __init__(self,
params,
lr=1,
max_iter=20,
max_eval=None,
tolerance_grad=1e-7,
tolerance_change=1e-9,
history_size=100,
line_search_fn=None):
if max_eval is None:
max_eval = max_iter * 5 // 4
defaults = dict(
lr=lr,
max_iter=max_iter,
max_eval=max_eval,
tolerance_grad=tolerance_grad,
tolerance_change=tolerance_change,
history_size=history_size,
line_search_fn=line_search_fn)
super(MyLBFGS, self).__init__(params, defaults)
if len(self.param_groups) != 1:
raise ValueError("LBFGS doesn't support per-parameter options "
"(parameter groups)")
self._params = self.param_groups[0]['params']
self._numel_cache = None
def _numel(self):
if self._numel_cache is None:
self._numel_cache = reduce(lambda total, p: total + p.numel(), self._params, 0)
return self._numel_cache
def _gather_flat_grad(self):
views = []
for p in self._params:
if p.grad is None:
view = p.new(p.numel()).zero_()
elif p.grad.is_sparse:
view = p.grad.to_dense().view(-1)
else:
view = p.grad.view(-1)
views.append(view)
return torch.cat(views, 0)
def _add_grad(self, step_size, update):
offset = 0
for p in self._params:
numel = p.numel()
if ((update.dtype==torch.complex64 or update.dtype==torch.complex128)
and (p.dtype==torch.float32 or p.dtype==torch.float64)):
p.add_(update[offset:offset + numel].real.view_as(p), alpha=step_size)
else:
p.add_(update[offset:offset + numel].view_as(p), alpha=step_size)
# view as to avoid deprecated pointwise semantics
#try:
# print(p.dtype)
# print(update.dtype)
# p.add_(update[offset:offset + numel].view_as(p), alpha=step_size)
# print('Worked fine')
# print(update[offset:offset + numel].view_as(p))
#except:
# print('Failed')
# print(update[offset:offset + numel].view_as(p))
# exit()
offset += numel
assert offset == self._numel()
def _clone_param(self):
return [p.clone(memory_format=torch.contiguous_format) for p in self._params]
def _set_param(self, params_data):
for p, pdata in zip(self._params, params_data):
p.copy_(pdata)
def _directional_evaluate(self, closure, x, t, d):
self._add_grad(t, d)
loss = float(closure())
flat_grad = self._gather_flat_grad()
self._set_param(x)
return loss, flat_grad
@torch.no_grad()
def step(self, closure):
"""Performs a single optimization step.
Args:
closure (callable): A closure that reevaluates the model
and returns the loss.
"""
assert len(self.param_groups) == 1
# Make sure the closure is always called with grad enabled
closure = torch.enable_grad()(closure)
group = self.param_groups[0]
lr = group['lr']
max_iter = group['max_iter']
max_eval = group['max_eval']
tolerance_grad = group['tolerance_grad']
tolerance_change = group['tolerance_change']
line_search_fn = group['line_search_fn']
history_size = group['history_size']
# NOTE: LBFGS has only global state, but we register it as state for
# the first param, because this helps with casting in load_state_dict
state = self.state[self._params[0]]
state.setdefault('func_evals', 0)
state.setdefault('n_iter', 0)
# evaluate initial f(x) and df/dx
orig_loss = closure()
loss = float(orig_loss)
current_evals = 1
state['func_evals'] += 1
flat_grad = self._gather_flat_grad()
opt_cond = flat_grad.abs().max() <= tolerance_grad
# optimal condition
if opt_cond:
return orig_loss
# tensors cached in state (for tracing)
d = state.get('d')
t = state.get('t')
old_dirs = state.get('old_dirs')
old_stps = state.get('old_stps')
ro = state.get('ro')
H_diag = state.get('H_diag')
prev_flat_grad = state.get('prev_flat_grad')
prev_loss = state.get('prev_loss')
n_iter = 0
# optimize for a max of max_iter iterations
while n_iter < max_iter:
# keep track of nb of iterations
n_iter += 1
state['n_iter'] += 1
############################################################
# compute gradient descent direction
############################################################
if state['n_iter'] == 1:
d = flat_grad.neg()
old_dirs = []
old_stps = []
ro = []
H_diag = 1
else:
# do lbfgs update (update memory)
y = flat_grad.sub(prev_flat_grad)
s = d.mul(t)
ys = y.dot(s) # y*s
if ys.abs() > 1e-10:
# updating memory
if len(old_dirs) == history_size:
# shift history by one (limited-memory)
old_dirs.pop(0)
old_stps.pop(0)
ro.pop(0)
# store new direction/step
old_dirs.append(y)
old_stps.append(s)
ro.append(1. / ys)
# update scale of initial Hessian approximation
H_diag = ys / y.dot(y) # (y*y)
# compute the approximate (L-BFGS) inverse Hessian
# multiplied by the gradient
num_old = len(old_dirs)
if 'al' not in state:
state['al'] = [None] * history_size
al = state['al']
# iteration in L-BFGS loop collapsed to use just one buffer
q = flat_grad.neg()
for i in range(num_old - 1, -1, -1):
al[i] = old_stps[i].dot(q) * ro[i]
q.add_(old_dirs[i], alpha=-al[i])
# multiply by initial Hessian
# r/d is the final direction
d = r = torch.mul(q, H_diag)
for i in range(num_old):
be_i = old_dirs[i].dot(r) * ro[i]
r.add_(old_stps[i], alpha=al[i] - be_i)
if prev_flat_grad is None:
prev_flat_grad = flat_grad.clone(memory_format=torch.contiguous_format)
else:
prev_flat_grad.copy_(flat_grad)
prev_loss = loss
############################################################
# compute step length
############################################################
# reset initial guess for step size
if state['n_iter'] == 1:
t = min(1., 1. / flat_grad.abs().sum()) * lr
else:
t = lr
# directional derivative
gtd = flat_grad.dot(d) # g * d
# Why did this used to be gtd > -tol_change?
# directional derivative is below tolerance
if gtd.abs() < tolerance_change:
break
# optional line search: user function
ls_func_evals = 0
if line_search_fn is not None:
# perform line search, using user function
if line_search_fn != "strong_wolfe":
raise RuntimeError("only 'strong_wolfe' is supported")
else:
x_init = self._clone_param()
def obj_func(x, t, d):
return self._directional_evaluate(closure, x, t, d)
loss, flat_grad, t, ls_func_evals = _strong_wolfe(
obj_func, x_init, t, d, loss, flat_grad, gtd)
self._add_grad(t, d)
opt_cond = flat_grad.abs().max() <= tolerance_grad
else:
# no line search, simply move with fixed-step
self._add_grad(t, d)
if n_iter != max_iter:
# re-evaluate function only if not in last iteration
# the reason we do this: in a stochastic setting,
# no use to re-evaluate that function here
with torch.enable_grad():
loss = float(closure())
flat_grad = self._gather_flat_grad()
opt_cond = flat_grad.abs().max() <= tolerance_grad
ls_func_evals = 1
# update func eval
current_evals += ls_func_evals
state['func_evals'] += ls_func_evals
############################################################
# check conditions
############################################################
if n_iter == max_iter:
break
if current_evals >= max_eval:
break
# optimal condition
if opt_cond:
break
# lack of progress
if d.mul(t).abs().max() <= tolerance_change:
break
if abs(loss - prev_loss) < tolerance_change:
break
state['d'] = d
state['t'] = t
state['old_dirs'] = old_dirs
state['old_stps'] = old_stps
state['ro'] = ro
state['H_diag'] = H_diag
state['prev_flat_grad'] = prev_flat_grad
state['prev_loss'] = prev_loss
return orig_loss
CDTools/models/fancy_ptycho.py
+245 -64
View File
@@ -4,11 +4,12 @@ import torch as t
from CDTools.models import CDIModel
from CDTools.datasets import Ptycho2DDataset
from CDTools import tools
from CDTools.tools import cmath
from CDTools.tools import plotting as p
from CDTools.tools import analysis
from matplotlib import pyplot as plt
from datetime import datetime
import numpy as np
from scipy import linalg as sla
from copy import copy
__all__ = ['FancyPtycho']
@@ -23,7 +24,8 @@ class FancyPtycho(CDIModel):
min_translation = t.Tensor([0,0]),
background = None, translation_offsets=None, mask=None,
weights = None, translation_scale = 1, saturation=None,
probe_support = None, obj_support=None, oversampling=1):
probe_support = None, obj_support=None, oversampling=1,
loss='amplitude mse',units='um'):
super(FancyPtycho,self).__init__()
self.wavelength = t.Tensor([wavelength])
@@ -43,6 +45,7 @@ class FancyPtycho(CDIModel):
self.surface_normal = t.Tensor(surface_normal)
self.saturation = saturation
self.units = units
if mask is None:
self.mask = mask
@@ -51,21 +54,21 @@ class FancyPtycho(CDIModel):
# We rescale the probe here so it learns at the same rate as the
# object
if probe_guess.dim() > 3:
self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess[0].to(t.float32)))
if probe_guess.dim() > 2:
self.probe_norm = 1 * t.max(t.abs(probe_guess[0].to(t.complex64)))
else:
self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess.to(t.float32)))
self.probe_norm = 1 * t.max(t.abs(probe_guess.to(t.complex64)))
self.probe = t.nn.Parameter(probe_guess.to(t.float32)
self.probe = t.nn.Parameter(probe_guess.to(t.complex64)
/ self.probe_norm)
self.obj = t.nn.Parameter(obj_guess.to(t.float32))
self.obj = t.nn.Parameter(obj_guess.to(t.complex64))
if background is None:
if detector_slice is not None:
background = 1e-6 * t.ones(self.probe[0][self.detector_slice].shape[:-1])
background = 1e-6 * t.ones(self.probe[0][self.detector_slice].shape)
else:
background = 1e-6 * t.ones(self.probe[0].shape[:-1])
background = 1e-6 * t.ones(self.probe[0].shape)
self.background = t.nn.Parameter(t.Tensor(background).to(t.float32))
@@ -73,8 +76,20 @@ class FancyPtycho(CDIModel):
if weights is None:
self.weights = None
else:
self.weights = t.nn.Parameter(t.Tensor(weights).to(t.float32))
# We now need to distinguish between real-valued per-image
# weights and complex-valued per-mode weight matrices
if len(weights.shape) == 1:
# This is if it's just a list of numbers
self.weights = t.nn.Parameter(t.Tensor(weights).to(t.float32))
else:
# Now this is a matrix of weights, so we
if type(weights) == type(t.zeros(1)):
self.weights = t.nn.Parameter(weights.to(t.complex64))
else:
# There is a good chance that this doesn't work
self.weights = t.nn.Parameter(t.Tensor(weights).to(t.complex64))
if translation_offsets is None:
self.translation_offsets = None
else:
@@ -96,8 +111,19 @@ class FancyPtycho(CDIModel):
self.oversampling = oversampling
# Here we set the appropriate loss function
if loss.lower().strip() == 'amplitude mse'\
or loss.lower().strip() == 'amplitude_mse':
self.loss = tools.losses.amplitude_mse
elif loss.lower().strip() == 'poisson nll' \
or loss.lower().strip() == 'poisson_nll':
self.loss = tools.losses.poisson_nll
else:
raise KeyError('Specified loss function not supported')
@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, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=True, opt_for_fft=False):
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, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=False, opt_for_fft=False, loss='amplitude mse', units='um'):
wavelength = dataset.wavelength
det_basis = dataset.detector_geometry['basis']
@@ -153,6 +179,7 @@ class FancyPtycho(CDIModel):
obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=200)
if hasattr(dataset, 'background') and dataset.background is not None:
background = t.sqrt(dataset.background)
else:
@@ -166,18 +193,35 @@ class FancyPtycho(CDIModel):
# Now we initialize all the subdominant probe modes
probe_max = t.max(cmath.cabs(probe))
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)]
probe = t.stack([probe,] + probe_stack)
#probe = t.stack([tools.propagators.far_field(probe),] + probe_stack)
obj = tools.cmath.expi(randomize_ang * (t.rand(obj_size)-0.5))
obj = t.exp(1j*randomize_ang * (t.rand(obj_size)-0.5))
det_geo = dataset.detector_geometry
translation_offsets = 0 * (t.rand((len(dataset),2)) - 0.5)
weights = t.ones(len(dataset))
if dm_rank is not None and dm_rank != 0:
if dm_rank > n_modes:
raise KeyError('Density matrix rank cannot be greater than the number of modes. Use dm_rank = -1 to use a full rank matrix.')
elif dm_rank == -1:
# dm_rank == -1 is defined to mean full-rank
dm_rank = n_modes
Ws = t.zeros(len(dataset),dm_rank,n_modes,dtype=t.complex64)
# Start with as close to the identity matrix as possible,
# cutting of when we hit the specified maximum rank
for i in range(0,dm_rank):
Ws[:,i,i] = 1
else:
# dm_rank == None or dm_rank = 0 triggers a special case where
# a standard incoherent multi-mode model is used. This is the
# default, because it is so common.
# In this case, we define a set of weights which only has one index
Ws = t.ones(len(dataset))
if hasattr(dataset, 'mask') and dataset.mask is not None:
mask = dataset.mask.to(t.bool)
@@ -185,14 +229,17 @@ class FancyPtycho(CDIModel):
mask = None
if probe_support_radius is not None:
probe_support = t.zeros_like(probe[0].to(dtype=t.float32))
p_cent = np.array(probe.shape[1:3]).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
probe_support = t.zeros(probe[0].shape,dtype=t.bool)
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
if restrict_obj != -1:
ro = restrict_obj
@@ -209,45 +256,55 @@ class FancyPtycho(CDIModel):
surface_normal=surface_normal,
min_translation=min_translation,
translation_offsets = translation_offsets,
weights=weights, mask=mask, background=background,
weights=Ws, mask=mask, background=background,
translation_scale=translation_scale,
saturation=saturation,
probe_support=probe_support,
obj_support=obj_support,
oversampling=oversampling)
oversampling=oversampling,
loss=loss,units=units)
def interaction(self, index, translations):
# 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,
translations,
surface_normal=self.surface_normal)
pix_trans -= self.min_translation
# We then add on any recovered translation offset, if they exist
if self.translation_offsets is not None:
pix_trans += self.translation_scale * self.translation_offsets[index]
all_exit_waves = []
for i in range(self.probe.shape[0]):
# from storing the probe in Fourier space
#pr = tools.propagators.inverse_far_field(self.probe[i]) * self.probe_support
pr = self.probe[i] * self.probe_support
exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc(pr,
self.obj_support * self.obj,
pix_trans,
shift_probe=True)
exit_waves = exit_waves * self.probe_support[...,:,:]
# This restricts the basis probes to stay within the probe support
basis_prs = self.probe * self.probe_support[...,:,:]
if hasattr(self,'weights') and self.weights is not None:
if exit_waves.dim() == 4:
exit_waves = self.weights[index][:,None,None,None] * exit_waves
else:
exit_waves = self.weights[index] * exit_waves
all_exit_waves.append(exit_waves)
# Now we construct the probes for each shot from the basis probes
Ws = self.weights[index]
if len(self.weights[0].shape) == 0:
# If a purely stable coherent illumination is defined
prs = Ws[...,None,None,None] * basis_prs
else:
# If a frame-by-frame weight matrix is defined
# This takes the dot product of all the weight matrices with
# the probes. The output has dimensions of translation, then
# coherent mode index, then x,y, and then complex index
# Maybe this can be done with a matmul now?
prs = t.sum(Ws[...,None,None] * basis_prs, axis=-3)
# Now we actually do the interaction, using the sinc subpixel
# translation model as per usual
exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc(
prs, self.obj_support * self.obj,pix_trans,
shift_probe=True, multiple_modes=True)
#exit_waves = self.probe_norm * tools.interactions.ptycho_2D_round(
# prs, self.obj_support * self.obj,pix_trans,
# multiple_modes=True)
return t.stack(all_exit_waves)
return exit_waves
def forward_propagator(self, wavefields):
@@ -266,11 +323,9 @@ class FancyPtycho(CDIModel):
saturation=self.saturation,
oversampling=self.oversampling)
def loss(self, sim_data, real_data, mask=None):
return tools.losses.amplitude_mse(real_data, sim_data, mask=mask)
#return tools.losses.poisson_nll(real_data, sim_data, mask=mask)
# Note: No "loss" function is defined here, because it is added
# dynamically during object creation in __init__
def to(self, *args, **kwargs):
super(FancyPtycho, self).to(*args, **kwargs)
@@ -335,29 +390,155 @@ class FancyPtycho(CDIModel):
def corrected_translations(self,dataset):
translations = dataset.translations.to(dtype=self.probe.dtype,device=self.probe.device)
translations = dataset.translations.to(dtype=t.float32,device=self.probe.device)
t_offset = tools.interactions.pixel_to_translations(self.probe_basis,self.translation_offsets*self.translation_scale,surface_normal=self.surface_normal)
return translations + t_offset
def get_rhos(self):
# If this is the general unified mode model
if self.weights.dim() >= 2:
Ws = self.weights.detach().cpu().numpy()
rhos_out = np.matmul(np.swapaxes(Ws,1,2), Ws.conj())
return rhos_out
# This is the purely incoherent case
else:
return np.array([np.eye(self.probe.shape[0])]*self.weights.shape[0],
dtype=np.complex64)
def tidy_probes(self, normalization=1, normalize=False):
"""Tidies up the probes
What we want to do here is use all the information on all the probes
to calculate a natural basis for the experiment, and update all the
density matrices to operate in that updated basis
"""
# First we treat the purely incoherent case
# I don't love this pattern of using an if statement with a return
# to catch this case, but because it's so much simpler than the
# unified mode case I think it's appropriate
if self.weights.dim() == 1:
probe = self.probe.detach().cpu().numpy()
ortho_probes = analysis.orthogonalize_probes(probe)
self.probe.data = t.as_tensor(ortho_probes,
device=self.probe.device,dtype=self.probe.dtype)
return
# Needs to be updated to allow for plotting to an existing figure
# This is for the unified mode case
# Note to future: We could probably do this more cleanly with an
# SVD directly on the Ws matrix, instead of an eigendecomposition
# of the rho matrix.
rhos = self.get_rhos()
overall_rho = np.mean(rhos,axis=0)
probe = self.probe.detach().cpu().numpy()
ortho_probes, A = analysis.orthogonalize_probes(probe,
density_matrix=overall_rho,
keep_transform=True,
normalize=normalize)
Aconj = A.conj()
Atrans = np.transpose(A)
new_rhos = np.matmul(Atrans,np.matmul(rhos,Aconj))
new_rhos /= normalization
ortho_probes *= np.sqrt(normalization)
dm_rank = self.weights.shape[1]
new_Ws = []
for rho in new_rhos:
# These are returned from smallest to largest - we want to keep
# the largest ones
w,v = sla.eigh(rho)
w = w[::-1][:dm_rank]
v = v[:,::-1][:,:dm_rank]
# For situations where the rank of the density matrix is not
# full in reality, but we keep more modes around than needed,
# some ws can go negative due to numerical error! This is
# extremely rare, but comon enough to cause crashes occasionally
# when there are thousands of individual matrices to transform
# every time this is called.
w = np.maximum(w,0)
new_Ws.append(np.dot(np.diag(np.sqrt(w)),v.transpose()))
new_Ws = np.array(new_Ws)
self.weights.data = t.as_tensor(new_Ws,
dtype=self.weights.dtype,device=self.weights.device)
self.probe.data = t.as_tensor(ortho_probes,
device=self.probe.device,dtype=self.probe.dtype)
def plot_wavefront_variation(self, dataset,fig=None,mode='amplitude',**kwargs):
def get_probes(idx):
basis_prs = self.probe * self.probe_support[...,:,:]
prs = t.sum(self.weights[idx,:,:,None,None] * basis_prs, axis=-4)
ortho_probes = analysis.orthogonalize_probes(prs)
if mode.lower() == 'amplitude':
return np.abs(ortho_probes.detach().cpu().numpy())
if mode.lower() == 'root_sum_intensity':
return np.sum(np.abs(ortho_probes.detach().cpu().numpy())**2,axis=0)
if mode.lower() == 'phase':
return np.angle(ortho_probes.detach().cpu().numpy())
probe_matrix = np.zeros([self.probe.shape[0]]*2,
dtype=np.complex64)
np_probes = self.probe.detach().cpu().numpy()
for i in range(probe_matrix.shape[0]):
for j in range(probe_matrix.shape[0]):
probe_matrix[i,j] = np.sum(np_probes[i]*np_probes[j].conj())
weights = self.weights.detach().cpu().numpy()
probe_intensities = np.sum(np.tensordot(weights,probe_matrix,axes=1)*
weights.conj(),axis=2)
# Imaginary part is already essentially zero up to rounding error
probe_intensities = np.real(probe_intensities)
values = np.sum(probe_intensities,axis=1)
if mode.lower() == 'amplitude' or mode.lower() == 'root_sum_intensity':
cmap = 'viridis'
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),
plot_list = [
('Dominant Probe Amplitude',
lambda self, fig: p.plot_amplitude(self.probe[0], fig=fig, basis=self.probe_basis)),
('Dominant Probe Phase',
lambda self, fig: p.plot_phase(self.probe[0], fig=fig, basis=self.probe_basis)),
('Subdominant Probe Amplitude',
lambda self, fig: p.plot_amplitude(self.probe[1], fig=fig, basis=self.probe_basis),
lambda self: len(self.probe) >=2),
('Subdominant Probe Phase',
lambda self, fig: p.plot_phase(self.probe[1], fig=fig, basis=self.probe_basis),
lambda self: len(self.probe) >=2),
('',
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(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(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)),
('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),
('% Power in Top Mode (only accurate after tidy_probes)',
lambda self, fig, dataset: p.plot_nanomap(self.corrected_translations(dataset), analysis.calc_top_mode_fraction(self.get_rhos()), fig=fig,units=self.units),
lambda self: len(self.weights.shape) >=2),
('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.probe_basis,units=self.units)),
('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.probe_basis,units=self.units)),
('Corrected Translations',
lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig)),
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))
]
@@ -366,9 +547,9 @@ class FancyPtycho(CDIModel):
def save_results(self, dataset):
basis = self.probe_basis.detach().cpu().numpy()
translations = self.corrected_translations(dataset).detach().cpu().numpy()
probe = cmath.torch_to_complex(self.probe.detach().cpu())
probe = self.probe.detach().cpu().numpy()
probe = probe * self.probe_norm.detach().cpu().numpy()
obj = cmath.torch_to_complex(self.obj.detach().cpu())
obj = self.obj.detach().cpu().numpy()
background = self.background.detach().cpu().numpy()**2
weights = self.weights.detach().cpu().numpy()
CDTools/models/multislice_2d_ptycho.py
+345 -170
View File
@@ -1,20 +1,29 @@
from __future__ import division, print_function, absolute_import
import torch as t
from CDTools.models import CDIModel
from CDTools.datasets import Ptycho2DDataset
from CDTools import tools
from CDTools.tools import cmath
from CDTools.tools import analysis, image_processing
from CDTools.tools import plotting as p
from matplotlib import pyplot as plt
from datetime import datetime
import numpy as np
from copy import copy
from functools import reduce
__all__ = ['Multislice2DPtycho']
class Multislice2DPtycho(CDIModel):
@property
def probe(self):
return t.complex(self.probe_real,self.probe_imag)
@property
def obj(self):
return t.complex(self.obj_real,self.obj_imag)
def __init__(self, wavelength, detector_geometry,
probe_basis,
probe_guess, obj_guess, dz, nz,
@@ -23,15 +32,20 @@ class Multislice2DPtycho(CDIModel):
min_translation = t.Tensor([0,0]),
background = None, translation_offsets=None, mask=None,
weights = None, translation_scale = 1, saturation=None,
#probe_support = None,
probe_fourier_support=None,
obj_support=None, oversampling=1,
bandlimit=4/5):
probe_support = None,
oversampling=1,
bandlimit=None,
subpixel=True,
exponentiate_obj=True,
fourier_probe=False,
prevent_aliasing=True,
phase_only=False,
units='um'):
super(Multislice2DPtycho,self).__init__()
self.wavelength = t.Tensor([wavelength])
self.detector_geometry = copy(detector_geometry)
self.dz = -dz
self.dz = dz
self.nz = nz
det_geo = self.detector_geometry
if hasattr(det_geo, 'distance'):
@@ -48,6 +62,12 @@ class Multislice2DPtycho(CDIModel):
self.surface_normal = t.Tensor(surface_normal)
self.saturation = saturation
self.subpixel = subpixel
self.exponentiate_obj = exponentiate_obj
self.fourier_probe = fourier_probe
self.units = units
self.phase_only=phase_only
self.prevent_aliasing=prevent_aliasing
if mask is None:
self.mask = mask
@@ -57,66 +77,80 @@ class Multislice2DPtycho(CDIModel):
# We rescale the probe here so it learns at the same rate as the
# object
if probe_guess.dim() > 3:
self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess[0].to(t.float32)))
self.probe_norm = 1 * t.max(t.abs(probe_guess[0]).to(t.float32))
else:
self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess.to(t.float32)))
self.probe = t.nn.Parameter(probe_guess.to(t.float32)
/ self.probe_norm)
self.probe_norm = 1 * t.max(t.abs(probe_guess).to(t.float32))
pg = probe_guess.to(t.complex64)/self.probe_norm
self.probe_real = t.nn.Parameter(pg.real)
self.probe_imag = t.nn.Parameter(pg.imag)
og = obj_guess.to(t.complex64)
self.obj_real = t.nn.Parameter(og.real)
self.obj_imag = t.nn.Parameter(og.imag)
self.obj = t.nn.Parameter(obj_guess.to(t.float32))
#self.probe = t.nn.Parameter(probe_guess.to(t.complex64)
# / self.probe_norm)
#self.obj = t.nn.Parameter(obj_guess.to(t.complex64))
if background is None:
if detector_slice is not None:
background = 1e-6 * t.ones(self.probe[0][self.detector_slice].shape[:-1])
background = 1e-6 * t.ones(self.probe[0][self.detector_slice].shape)
else:
background = 1e-6 * t.ones(self.probe[0].shape[:-1])
background = 1e-6 * t.ones(self.probe[0].shape)
self.background = t.nn.Parameter(t.Tensor(background).to(t.float32))
self.background = t.nn.Parameter(t.as_tensor(background,dtype=t.float32))
if weights is None:
self.weights = None
else:
self.weights = t.nn.Parameter(t.Tensor(weights).to(t.float32))
# We now need to distinguish between real-valued per-image
# weights and complex-valued per-mode weight matrices
if len(weights.shape) == 1:
# This is if it's just a list of numbers
self.weights = t.nn.Parameter(t.as_tensor(weights,
dtype=t.float32))
else:
# Now this is a matrix of weights, so we
self.weights = t.nn.Parameter(t.as_tensor(weights,
dtype=t.complex64))
if translation_offsets is None:
self.translation_offsets = None
else:
self.translation_offsets = t.nn.Parameter(t.Tensor(translation_offsets).to(t.float32)/ translation_scale)
self.translation_offsets = t.nn.Parameter(t.as_tensor(translation_offsets,dtype=t.float32)/ translation_scale)
self.translation_scale = translation_scale
self.probe_fourier_support = t.Tensor(probe_fourier_support).to(t.float32)
#if probe_support is not None:
# self.probe_support = probe_support
#else:
# self.probe_support = t.ones_like(self.probe[0])
if obj_support is not None:
self.obj_support = obj_support
if self.obj.dim() == 3:
self.obj.data = self.obj * obj_support
elif self.obj.dim() == 4:
self.obj.data = self.obj * obj_support[None,...]
if probe_support is not None:
self.probe_support = t.as_tensor(probe_support,dtype=t.bool)
else:
if self.obj.dim() == 3:
self.obj_support = t.ones_like(self.obj)
elif self.obj.dim() == 4:
self.obj_support = t.ones_like(self.obj[0])
self.probe_support = None#t.ones_like(self.probe,dtype=t.bool)#None
self.oversampling = oversampling
spacing = np.linalg.norm(self.probe_basis,axis=0)
shape = np.array(self.probe.shape[1:-1])
shape = np.array(self.probe.shape[1:])
if prevent_aliasing:
shape *= 2
spacing /= 2
self.bandlimit = bandlimit
# Big question: Should there be a minus sign before self.dz, or not?
self.as_prop = tools.propagators.generate_angular_spectrum_propagator(shape, spacing, self.wavelength, self.dz, bandlimit=self.bandlimit)
self.as_prop = tools.propagators.generate_angular_spectrum_propagator(shape, spacing, self.wavelength, self.dz, bandlimit=1/np.sqrt(2))#self.bandlimit)
#plt.imshow(t.abs(self.as_prop))
#plt.figure()
#plt.imshow(t.abs(t.fft.fftshift(t.fft.ifft2(self.as_prop))))
#plt.show()
#exit()
@classmethod
def from_dataset(cls, dataset, dz, nz, probe_convergence_radius, probe_size=None, padding=0, n_modes=1, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=True, bandlimit=4/5, replicate_slice=False):
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):
wavelength = dataset.wavelength
det_basis = dataset.detector_geometry['basis']
@@ -165,12 +199,11 @@ 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)
obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=200)
obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=100)
if hasattr(dataset, 'background') and dataset.background is not None:
background = t.sqrt(dataset.background)
@@ -178,32 +211,53 @@ class Multislice2DPtycho(CDIModel):
background = None
# 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.STEM_style_probe(dataset, probe_shape, det_slice, probe_convergence_semiangle, propagation_distance=propagation_distance, oversampling=oversampling)
#probe = tools.initializers.SHARP_style_probe(dataset, probe_shape, det_slice, propagation_distance=propagation_distance, oversampling=oversampling)
# Now we initialize all the subdominant probe modes
probe_max = t.max(t.abs(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)
else:
probe = tools.initializers.gaussian_probe(dataset, probe_basis, probe_shape, probe_size, propagation_distance=propagation_distance)
probe = t.unsqueeze(probe,0)
# For a Fourier space probe
probe = tools.propagators.inverse_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)
if fourier_probe:
probe = tools.propagators.far_field(probe)
# Consider a different start
obj = t.zeros(obj_size+(2,))
#obj = tools.cmath.expi(t.zeros(obj_size))
if exponentiate_obj:
obj = t.zeros(obj_size, dtype=t.complex64)
else:
obj = t.exp(1j*t.zeros(obj_size))
# If we will use a separate object per slice
if not replicate_slice:
obj = t.stack([obj]*nz)
det_geo = dataset.detector_geometry
translation_offsets = 0 * (t.rand((len(dataset),2)) - 0.5)
weights = t.ones(len(dataset))
if dm_rank is not None and dm_rank != 0:
if dm_rank > n_modes:
raise KeyError('Density matrix rank cannot be greater than the number of modes. Use dm_rank = -1 to use a full rank matrix.')
elif dm_rank == -1:
# dm_rank == -1 is defined to mean full-rank
dm_rank = n_modes
Ws = t.zeros(len(dataset),dm_rank,n_modes,2)
# Start with as close to the identity matrix as possible,
# cutting of when we hit the specified maximum rank
for i in range(0,dm_rank):
Ws[:,i,i,0] = 1
else:
# dm_rank == None or dm_rank = 0 triggers a special case where
# a standard incoherent multi-mode model is used. This is the
# default, because it is so common.
# In this case, we define a set of weights which only has one index
Ws = t.ones(len(dataset))
if hasattr(dataset, 'mask') and dataset.mask is not None:
mask = dataset.mask.to(t.bool)
@@ -211,51 +265,36 @@ class Multislice2DPtycho(CDIModel):
mask = None
if probe_support_radius is not None:
probe_support = t.zeros_like(probe[0].to(dtype=t.float32))
p_cent = np.array(probe.shape[1:3]).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
probe = probe * probe_support[None,:,:]
else:
probe_support = None;
if restrict_obj != -1:
ro = restrict_obj
os = np.array(obj_size)
ps = np.array(probe_shape)
if replicate_slice:
obj_support = t.zeros_like(obj.to(dtype=t.float32))
else:
obj_support = t.zeros_like(obj[0].to(dtype=t.float32))
obj_support[ps[0]//2-ro:os[0]+ro-ps[0]//2,
ps[1]//2-ro:os[1]+ro-ps[1]//2] = 1
else:
obj_support = None
probe_support = t.zeros_like(probe[0].to(dtype=t.float32))
xs, ys = np.mgrid[:probe.shape[-3],:probe.shape[-2]]
xs = xs - np.mean(xs)
ys = ys - np.mean(ys)
Rs = np.sqrt(xs**2 + ys**2)
probe_support = t.zeros(probe[0].shape, dtype=t.bool)
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_convergence_radius] = 1
probe = probe * probe_support[None,:,:]
probe_support[Rs<probe_support_radius] = 1
probe = probe * probe_support[None,:,:]
else:
probe_support = None
probe = probe
return cls(wavelength, det_geo, probe_basis, probe, obj, dz, nz,
detector_slice=det_slice,
surface_normal=surface_normal,
probe_support=probe_support,
min_translation=min_translation,
translation_offsets = translation_offsets,
weights=weights, mask=mask, background=background,
weights=Ws, mask=mask, background=background,
translation_scale=translation_scale,
saturation=saturation,
#probe_support=probe_support,
probe_fourier_support=probe_support,
obj_support=obj_support,
oversampling=oversampling,
bandlimit=bandlimit)
bandlimit=bandlimit,
subpixel=subpixel,
exponentiate_obj=exponentiate_obj,
units=units, fourier_probe=fourier_probe,
phase_only=phase_only,
prevent_aliasing=prevent_aliasing)
def interaction(self, index, translations):
@@ -266,69 +305,98 @@ class Multislice2DPtycho(CDIModel):
if self.translation_offsets is not None:
pix_trans += self.translation_scale * self.translation_offsets[index]
if len(pix_trans.shape) == 1:
pix_trans = [pix_trans]
index = [index]
strip_first_index = True
# This restricts the basis probes to stay within the probe support
if self.probe_support is not None:
basis_prs = self.probe * self.probe_support[...,:,:]
else:
strip_first_index = False
basis_prs = self.probe
all_exit_waves = []
for i in range(self.probe.shape[0]):
# For a Fourier-space probe
pr = tools.propagators.inverse_far_field(self.probe[i] * self.probe_fourier_support)
# For a real-space probe
#pr = self.probe[i] * 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)
if self.prevent_aliasing:
#pix_trans = pix_trans * 2
basis_prs = image_processing.fourier_upsample(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:
# If a purely stable coherent illumination is defined
# No cmult because Ws is real in this case
prs = Ws[...,None,None,None] * basis_prs
else:
# If a frame-by-frame weight matrix is defined
# This takes the dot product of all the weight matrices with
# the probes. The output has dimensions of translation, then
# coherent mode index, then x,y, and then complex index
prs = t.sum(Ws[...,None,None] * basis_prs, axis=-3)
#exit_waves = pr
#print(self.probe_norm)
#for i in range(self.nz):
exit_waves = []
for trans in pix_trans:
exit_wave = self.probe_norm * pr
for i in range(self.nz):
# If only one object slice
if self.obj.dim() == 3:
exit_wave = tools.interactions.ptycho_2D_sinc(
exit_wave,
self.obj_support*cmath.cexpi(self.obj/self.nz),
trans, shift_probe=True)
# If separate slices
elif self.obj.dim() == 4:
exit_wave = tools.interactions.ptycho_2D_sinc(
exit_wave,
self.obj_support*cmath.cexpi(self.obj[i]/self.nz),
trans, shift_probe=True)
exit_wave = tools.propagators.near_field(
exit_wave,self.as_prop)
#tools.plotting.plot_amplitude(exit_wave)
#plt.show()
exit_waves.append(exit_wave)
exit_waves = t.stack(exit_waves)
if exit_waves.dim() == 4:
exit_waves = self.weights[index][:,None,None,None] * exit_waves
if self.exponentiate_obj:
if self.phase_only:
obj = t.exp(1j*self.obj.real)
else:
exit_waves = self.weights[index] * exit_waves
obj = t.exp(1j*self.obj)
else:
obj = self.obj
if strip_first_index:
exit_waves = exit_waves[0,...]
all_exit_waves.append(exit_waves)
#if self.prevent_aliasing:
# obj = image_processing.fourier_upsample(obj)
return t.stack(all_exit_waves)
exit_waves = self.probe_norm * prs
for i in range(self.nz):
# If only one object slice
if self.obj.dim() == 2:
if i == 0 and self.subpixel:
# We only need to apply the subpixel shift to the first
# slice, because it shifts the probe
exit_waves = tools.interactions.ptycho_2D_sinc(
exit_waves, obj, pix_trans,
shift_probe=True, multiple_modes=True)
else:
exit_waves = tools.interactions.ptycho_2D_round(
exit_waves, obj, pix_trans,
multiple_modes=True,upsample_obj=self.prevent_aliasing)
elif self.obj.dim() == 3:
# If separate slices
if i == 0 and self.subpixel:
exit_waves = tools.interactions.ptycho_2D_sinc(
exit_waves, obj[i], pix_trans,
shift_probe=True, multiple_modes=True)
else:
exit_waves = tools.interactions.ptycho_2D_round(
exit_waves, obj[i], pix_trans,
multiple_modes=True, upsample_obj=self.prevent_aliasing)
#if self.iteration_count >= 1:
# plt.imshow(t.abs(tools.propagators.far_field(
# exit_waves[0,0].detach()).cpu()))
# plt.colorbar()
# plt.show()
if i < self.nz-1: #on all but the last iteration
exit_waves = tools.propagators.near_field(
exit_waves,self.as_prop)
#plt.imshow(t.abs(exit_waves[0,0].detach().cpu()))
#plt.show()
return exit_waves
def forward_propagator(self, wavefields):
return tools.propagators.far_field(wavefields)
def backward_propagator(self, wavefields):
return tools.propagators.inverse_far_field(wavefields)
if self.prevent_aliasing:
left = [self.probe.shape[-2]//2,self.probe.shape[-1]//2]
right = [self.probe.shape[-2]//2+self.probe.shape[-2],
self.probe.shape[-1]//2+self.probe.shape[-1]]
return tools.propagators.far_field(wavefields)[...,left[0]:right[0],
left[1]:right[1]]
else:
return tools.propagators.far_field(wavefields)
def measurement(self, wavefields):
@@ -341,11 +409,8 @@ class Multislice2DPtycho(CDIModel):
def loss(self, sim_data, real_data, mask=None):
regularizer = t.mean(t.abs(self.obj))
loss = tools.losses.amplitude_mse(real_data, sim_data, mask=mask)
lambd = 1
return loss + lambd * regularizer
#return tools.losses.poisson_nll(real_data, sim_data, mask=mask)
return tools.losses.amplitude_mse(real_data, sim_data, mask=mask)
#return tools.losses.poisson_nll(real_data, sim_data, mask=mask,eps=0.5)
def to(self, *args, **kwargs):
@@ -362,14 +427,14 @@ class Multislice2DPtycho(CDIModel):
if self.mask is not None:
self.mask = self.mask.to(*args, **kwargs)
self.min_translation = self.min_translation.to(*args,**kwargs)
self.probe_basis = self.probe_basis.to(*args,**kwargs)
self.probe_norm = self.probe_norm.to(*args,**kwargs)
#self.probe_support = self.probe_support.to(*args,**kwargs)
self.probe_fourier_support = self.probe_fourier_support.to(*args,**kwargs)
self.obj_support = self.obj_support.to(*args,**kwargs)
if self.probe_support is not None:
self.probe_support = self.probe_support.to(*args,**kwargs)
self.surface_normal = self.surface_normal.to(*args, **kwargs)
self.as_prop = self.as_prop.to(*args, **kwargs)
@@ -412,29 +477,134 @@ class Multislice2DPtycho(CDIModel):
def corrected_translations(self,dataset):
translations = dataset.translations.to(dtype=self.probe.dtype,device=self.probe.device)
translations = dataset.translations.to(dtype=self.probe.real.dtype,device=self.probe.device)
t_offset = tools.interactions.pixel_to_translations(self.probe_basis,self.translation_offsets*self.translation_scale,surface_normal=self.surface_normal)
return translations + t_offset
def get_rhos(self):
# If this is the general unified mode model
if self.weights.dim() >= 2:
Ws = self.weights.detach().cpu().numpy()
rhos_out = np.matmul(np.swapaxes(Ws,1,2), Ws.conj())
return rhos_out
# This is the purely incoherent case
else:
return np.array([np.eye(self.probe.shape[0])]*self.weights.shape[0],
dtype=np.complex64)
def tidy_probes(self, normalization=1, normalize=False):
"""Tidies up the probes
What we want to do here is use all the information on all the probes
to calculate a natural basis for the experiment, and update all the
density matrices to operate in that updated basis
"""
# First we treat the purely incoherent case
# I don't love this pattern of using an if statement with a return
# to catch this case, but because it's so much simpler than the
# unified mode case I think it's appropriate
if self.weights.dim() == 1:
probe = self.probe.detach().cpu().numpy()
ortho_probes = analysis.orthogonalize_probes(probe)
self.probe.data = t.as_tensor(ortho_probes,
device=self.probe.device,dtype=self.probe.dtype)
return
# This is for the unified mode case
# Note to future: We could probably do this more cleanly with an
# SVD directly on the Ws matrix, instead of an eigendecomposition
# of the rho matrix.
rhos = self.get_rhos()
overall_rho = np.mean(rhos,axis=0)
probe = self.probe.detach().cpu().numpy()
ortho_probes, A = analysis.orthogonalize_probes(probe,
density_matrix=overall_rho,
keep_transform=True,
normalize=normalize)
Aconj = A.conj()
Atrans = np.transpose(A)
new_rhos = np.matmul(Atrans,np.matmul(rhos,Aconj))
new_rhos /= normalization
ortho_probes *= np.sqrt(normalization)
dm_rank = self.weights.shape[1]
new_Ws = []
for rho in new_rhos:
# These are returned from smallest to largest - we want to keep
# the largest ones
w,v = sla.eigh(rho)
w = w[::-1][:dm_rank]
v = v[:,::-1][:,:dm_rank]
# For situations where the rank of the density matrix is not
# full in reality, but we keep more modes around than needed,
# some ws can go negative due to numerical error! This is
# extremely rare, but comon enough to cause crashes occasionally
# when there are thousands of individual matrices to transform
# every time this is called.
w = np.maximum(w,0)
new_Ws.append(np.dot(np.diag(np.sqrt(w)),v.transpose()))
new_Ws = np.array(new_Ws)
self.weights.data = t.as_tensor(new_Ws,
dtype=self.weights.dtype,device=self.weights.device)
self.probe.data = t.as_tensor(ortho_probes,
device=self.probe.device,dtype=self.probe.dtype)
# Needs to be updated to allow for plotting to an existing figure
plot_list = [
('Dominant Probe Amplitude',
lambda self, fig: p.plot_amplitude(self.probe[0], fig=fig, basis=self.probe_basis)),
('Dominant Probe Phase',
lambda self, fig: p.plot_phase(self.probe[0], fig=fig, basis=self.probe_basis)),
('Subdominant Probe Amplitude',
lambda self, fig: p.plot_amplitude(self.probe[1], fig=fig, basis=self.probe_basis),
lambda self: len(self.probe) >=2),
('Subdominant Probe Phase',
lambda self, fig: p.plot_phase(self.probe[1], fig=fig, basis=self.probe_basis),
lambda self: len(self.probe) >=2),
('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 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),
('% Power in Top Mode (only accurate after tidy_probes)',
lambda self, fig, dataset: p.plot_nanomap(self.corrected_translations(dataset), analysis.calc_top_mode_fraction(self.get_rhos()), 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.mean(self.obj.detach().cpu(),dim=0), fig=fig, basis=self.probe_basis)),
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.mean(self.obj.detach().cpu(),dim=0), fig=fig, basis=self.probe_basis)),
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)),
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))
]
@@ -443,9 +613,14 @@ class Multislice2DPtycho(CDIModel):
def save_results(self, dataset):
basis = self.probe_basis.detach().cpu().numpy()
translations = self.corrected_translations(dataset).detach().cpu().numpy()
probe = cmath.torch_to_complex(self.probe.detach().cpu())
if self.fourier_probe:
probe = tools.propagators.inverse_far_field(self.probe)
else:
probe = self.probe
probe = probe.detach().cpu().numpy()
probe = probe * self.probe_norm.detach().cpu().numpy()
obj = cmath.torch_to_complex(self.obj.detach().cpu())
obj = self.obj.detach().cpu().numpy()
background = self.background.detach().cpu().numpy()**2
weights = self.weights.detach().cpu().numpy()
dz = self.dz
CDTools/models/pinhole_plane_ptycho.py
-384
View File
@@ -1,384 +0,0 @@
from __future__ import division, print_function, absolute_import
import torch as t
from CDTools.models import CDIModel
from CDTools.datasets import Ptycho2DDataset
from CDTools import tools
from CDTools.tools import cmath
from CDTools.tools import plotting as p
from matplotlib import pyplot as plt
from datetime import datetime
import numpy as np
from copy import copy
__all__ = ['PinholePlanePtycho']
class PinholePlanePtycho(CDIModel):
def __init__(self, wavelength, detector_geometry,
probe_basis,
pinhole_guess, obj_guess, illumination_guess,
pinhole_propagation,
detector_slice=None,
surface_normal=np.array([0.,0.,1.]),
min_translation = t.Tensor([0,0]),
background = None, translation_offsets=None, mask=None,
weights = None, translation_scale = 1, saturation=None,
pinhole_support = None, obj_support=None, oversampling=1):
super(PinholePlanePtycho,self).__init__()
self.wavelength = t.Tensor([wavelength])
self.detector_geometry = copy(detector_geometry)
det_geo = self.detector_geometry
if hasattr(det_geo, 'distance'):
det_geo['distance'] = t.Tensor(det_geo['distance'])
if hasattr(det_geo, 'basis'):
det_geo['basis'] = t.Tensor(det_geo['basis'])
if hasattr(det_geo, 'corner'):
det_geo['corner'] = t.Tensor(det_geo['corner'])
self.min_translation = t.Tensor(min_translation)
self.probe_basis = t.Tensor(probe_basis)
self.detector_slice = detector_slice
self.surface_normal = t.Tensor(surface_normal)
self.saturation = saturation
if mask is None:
self.mask = mask
else:
self.mask = t.ByteTensor(mask)
# We rescale the pinhole here so it learns at the same rate as the
# object
self.pinhole_norm = 1 * t.max(tools.cmath.cabs(pinhole_guess.to(t.float32)))
self.pinhole = t.nn.Parameter(pinhole_guess.to(t.float32)
/ self.pinhole_norm)
self.obj = t.nn.Parameter(obj_guess.to(t.float32))
self.illumination = t.nn.Parameter(illumination_guess.to(t.float32))
spacing = np.abs([self.probe_basis.numpy()[0,1], self.probe_basis.numpy()[1,0]])
self.pinhole_propagator = tools.propagators.generate_angular_spectrum_propagator(self.pinhole.shape[:-1], spacing, self.wavelength, pinhole_propagation).to(dtype=self.pinhole.dtype)
if background is None:
if detector_slice is not None:
background = 1e-6 * t.ones(self.pinhole[self.detector_slice].shape[:-1])
else:
background = 1e-6 * t.ones(self.pinhole.shape[:-1])
self.background = t.nn.Parameter(t.Tensor(background).to(t.float32))
if weights is None:
self.weights = None
else:
self.weights = t.nn.Parameter(t.Tensor(weights).to(t.float32))
if translation_offsets is None:
self.translation_offsets = None
else:
self.translation_offsets = t.nn.Parameter(t.Tensor(translation_offsets).to(t.float32)/ translation_scale)
self.translation_scale = translation_scale
if pinhole_support is not None:
self.pinhole_support = pinhole_support
else:
self.pinhole_support = t.ones_like(self.pinhole)
if obj_support is not None:
self.obj_support = obj_support
self.obj.data = self.obj * obj_support
else:
self.obj_support = t.ones_like(self.obj)
self.oversampling = oversampling
@classmethod
def from_dataset(cls, dataset, pinhole_propagation, pinhole_size=None, randomize_ang=0, padding=0, translation_scale = 1, saturation=None, pinhole_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1):
wavelength = dataset.wavelength
det_basis = dataset.detector_geometry['basis']
det_shape = dataset[0][1].shape
distance = dataset.detector_geometry['distance']
# always do this on the cpu
get_as_args = dataset.get_as_args
dataset.get_as(device='cpu')
(indices, translations), patterns = dataset[:]
dataset.get_as(*get_as_args[0],**get_as_args[1])
# Set to none to avoid issues with things outside the detector
center = tools.image_processing.centroid(t.sum(patterns,dim=0))
# Then, generate the pinhole geometry from the dataset
ewg = tools.initializers.exit_wave_geometry
probe_basis, pinhole_shape, det_slice = ewg(det_basis,
det_shape,
wavelength,
distance,
center=center,
padding=padding,
opt_for_fft=False,
oversampling=oversampling)
if hasattr(dataset, 'sample_info') and \
dataset.sample_info is not None and \
'orientation' in dataset.sample_info:
surface_normal = dataset.sample_info['orientation'][2]
else:
surface_normal = np.array([0.,0.,1.])
# If this information is supplied when the function is called,
# then we override the information in the .cxi file
if scattering_mode in {'t', 'transmission'}:
surface_normal = np.array([0.,0.,1.])
elif scattering_mode in {'r', 'reflection'}:
outgoing_dir = np.cross(det_basis[:,0], det_basis[:,1])
outgoing_dir /= np.linalg.norm(outgoing_dir)
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)
obj_size, min_translation = tools.initializers.calc_object_setup(pinhole_shape, pix_translations, padding=200)
if hasattr(dataset, 'background') and dataset.background is not None:
background = t.sqrt(dataset.background)
else:
background = None
# Finally, initialize the pinhole and object using this information
if pinhole_size is None:
pinhole = tools.initializers.SHARP_style_probe(dataset, pinhole_shape, det_slice, propagation_distance=propagation_distance, oversampling=oversampling)
else:
pinhole = tools.initializers.gaussian_probe(dataset, probe_basis, pinhole_shape, pinhole_size, propagation_distance=propagation_distance)
illumination = tools.cmath.expi(t.zeros(obj_size))
obj = tools.cmath.expi(randomize_ang * (t.rand(obj_size)-0.5))
det_geo = dataset.detector_geometry
translation_offsets = 0 * (t.rand((len(dataset),2)) - 0.5)
weights = t.ones(len(dataset))
if hasattr(dataset, 'mask') and dataset.mask is not None:
mask = dataset.mask.to(t.uint8)
else:
mask = None
if pinhole_support_radius is not None:
pinhole_support = t.zeros_like(pinhole.to(dtype=t.float32))
p_cent = np.array(pinhole.shape[:2]).astype(int) // 2
psr = int(pinhole_support_radius)
pinhole_support[p_cent[0]-psr:p_cent[0]+psr,
p_cent[1]-psr:p_cent[1]+psr] = 1
pinhole = pinhole_support * pinhole
else:
pinhole_support = None;
if restrict_obj != -1:
ro = restrict_obj
os = np.array(obj_size)
ps = np.array(pinhole_shape)
obj_support = t.zeros_like(obj.to(dtype=t.float32))
obj_support[ps[0]//2-ro:os[0]+ro-ps[0]//2,
ps[1]//2-ro:os[1]+ro-ps[1]//2] = 1
else:
obj_support = None
return cls(wavelength, det_geo, probe_basis, pinhole, obj,
illumination, pinhole_propagation,
detector_slice=det_slice,
surface_normal=surface_normal,
min_translation=min_translation,
translation_offsets = translation_offsets,
weights=weights, mask=mask, background=background,
translation_scale=translation_scale,
saturation=saturation,
pinhole_support=pinhole_support,
obj_support=obj_support,
oversampling=oversampling)
def interaction(self, index, translations):
pix_trans = tools.interactions.translations_to_pixel(self.probe_basis,
translations,
surface_normal=self.surface_normal)
pix_trans -= self.min_translation
if self.translation_offsets is not None:
pix_trans += self.translation_scale * self.translation_offsets[index]
pr = self.pinhole * self.pinhole_support
pinhole_exit_waves = self.pinhole_norm * \
tools.interactions.ptycho_2D_sinc(pr,
self.illumination,
pix_trans,
shift_probe=True)
sample_illumination = tools.propagators.near_field(pinhole_exit_waves,
self.pinhole_propagator)
exit_waves = []
if pinhole_exit_waves.dim() == 4:
for pix_tran, illum in zip(pix_trans, sample_illumination):
exit_waves.append(tools.interactions.ptycho_2D_sinc(illum,
self.obj_support * self.obj,
pix_tran,
shift_probe=True))
exit_waves = t.stack(exit_waves)
else:
exit_waves = tools.interactions.ptycho_2D_sinc(sample_illumination,
self.obj_support * self.obj,
pix_trans,
shift_probe=True)
if exit_waves.dim() == 4:
exit_waves = self.weights[index][:,None,None,None] * exit_waves
else:
exit_waves = self.weights[index] * exit_waves
return exit_waves
def forward_propagator(self, wavefields):
return tools.propagators.far_field(wavefields)
def backward_propagator(self, wavefields):
return tools.propagators.inverse_far_field(wavefields)
def measurement(self, wavefields):
return tools.measurements.quadratic_background(wavefields,
self.background,
detector_slice=self.detector_slice,
saturation=self.saturation,
oversampling=self.oversampling)
def loss(self, sim_data, real_data, mask=None):
return tools.losses.amplitude_mse(real_data, sim_data, mask=mask)
def to(self, *args, **kwargs):
super(PinholePlanePtycho, self).to(*args, **kwargs)
self.wavelength = self.wavelength.to(*args,**kwargs)
# move the detector geometry too
det_geo = self.detector_geometry
if hasattr(det_geo, 'distance'):
det_geo['distance'] = det_geo['distance'].to(*args,**kwargs)
if hasattr(det_geo, 'basis'):
det_geo['basis'] = det_geo['basis'].to(*args,**kwargs)
if hasattr(det_geo, 'corner'):
det_geo['corner'] = det_geo['corner'].to(*args,**kwargs)
if self.mask is not None:
self.mask = self.mask.to(*args, **kwargs)
self.min_translation = self.min_translation.to(*args,**kwargs)
self.probe_basis = self.probe_basis.to(*args,**kwargs)
self.pinhole_norm = self.pinhole_norm.to(*args,**kwargs)
self.pinhole_support = self.pinhole_support.to(*args,**kwargs)
self.obj_support = self.obj_support.to(*args,**kwargs)
self.surface_normal = self.surface_normal.to(*args, **kwargs)
self.pinhole_propagator = self.pinhole_propagator.to(*args, **kwargs)
def sim_to_dataset(self, args_list):
# In the future, potentially add more control
# over what metadata is saved (names, etc.)
# First, I need to gather all the relevant data
# that needs to be added to the dataset
entry_info = {'program_name': 'CDTools',
'instrument_n': 'Simulated Data',
'start_time': datetime.now()}
surface_normal = self.surface_normal.detach().cpu().numpy()
xsurfacevec = np.cross(np.array([0.,1.,0.]), surface_normal)
xsurfacevec /= np.linalg.norm(xsurfacevec)
ysurfacevec = np.cross(surface_normal, xsurfacevec)
ysurfacevec /= np.linalg.norm(ysurfacevec)
orientation = np.array([xsurfacevec, ysurfacevec, surface_normal])
sample_info = {'description': 'A simulated sample',
'orientation': orientation}
detector_geometry = self.detector_geometry
mask = self.mask
wavelength = self.wavelength
indices, translations = args_list
# Then we simulate the results
data = self.forward(indices, translations)
# And finally, we make the dataset
return Ptycho2DDataset(translations, data,
entry_info = entry_info,
sample_info = sample_info,
wavelength=wavelength,
detector_geometry=detector_geometry,
mask=mask)
def corrected_translations(self,dataset):
translations = dataset.translations.to(dtype=self.pinhole.dtype,device=self.pinhole.device)
t_offset = tools.interactions.pixel_to_translations(self.probe_basis,self.translation_offsets*self.translation_scale,surface_normal=self.surface_normal)
return translations + t_offset
# Needs to be updated to allow for plotting to an existing figure
plot_list = [
('Pinhole Amplitude',
lambda self, fig: p.plot_amplitude(self.pinhole, fig=fig, basis=self.probe_basis)),
('Pinhole Phase',
lambda self, fig: p.plot_phase(self.pinhole, fig=fig, basis=self.probe_basis)),
('Illumination Amplitude',
lambda self, fig: p.plot_amplitude(self.illumination, fig=fig, basis=self.probe_basis),
lambda self: len(self.pinhole) >=2),
('Illumination Phase',
lambda self, fig: p.plot_phase(self.illumination, fig=fig, basis=self.probe_basis),
lambda self: len(self.pinhole) >=2),
('Object Amplitude',
lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis)),
('Object Phase',
lambda self, fig: p.plot_phase(self.obj, fig=fig, basis=self.probe_basis)),
('Corrected Translations',
lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig)),
('Background',
lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2))
]
def save_results(self, dataset):
basis = self.probe_basis.detach().cpu().numpy()
translations = self.corrected_translations(dataset).detach().cpu().numpy()
probe = cmath.torch_to_complex(self.pinhole.detach().cpu())
probe = probe * self.pinhole_norm.detach().cpu().numpy()
obj = cmath.torch_to_complex(self.obj.detach().cpu())
background = self.background.detach().cpu().numpy()**2
weights = self.weights.detach().cpu().numpy()
return {'basis':basis, 'translation':translations,
'probe':probe,'obj':obj,
'background':background,
'weights':weights}
CDTools/models/rpi.py
+30 -22
View File
@@ -4,7 +4,6 @@ import torch as t
from CDTools.models import CDIModel
from CDTools.datasets import Ptycho2DDataset
from CDTools import tools
from CDTools.tools import cmath
from CDTools.tools import plotting as p
from CDTools.tools.interactions import RPI_interaction
from CDTools.tools import initializers
@@ -51,6 +50,10 @@ __all__ = ['RPI']
__all__ = ['RPI']
class RPI(CDIModel):
@property
def obj(self):
return t.complex(self.obj_real, self.obj_imag)
def __init__(self, wavelength, detector_geometry, probe_basis,
probe, obj_guess, detector_slice=None,
@@ -73,8 +76,8 @@ class RPI(CDIModel):
self.probe_basis = t.Tensor(probe_basis)
scale_factor = t.Tensor([probe.shape[-2]/obj_guess.shape[-2],
probe.shape[-3]/obj_guess.shape[-3]])
scale_factor = t.Tensor([probe.shape[-1]/obj_guess.shape[-1],
probe.shape[-2]/obj_guess.shape[-2]])
self.obj_basis = self.probe_basis / scale_factor
self.detector_slice = detector_slice
@@ -89,12 +92,17 @@ class RPI(CDIModel):
self.mask = t.BoolTensor(mask)
self.probe = probe.to(t.float32)
self.probe = probe.to(t.complex64)
if obj_guess.dim() == 3:
obj_guess = obj_guess[None,:,:,:]
if obj_guess.dim() == 2:
obj_guess = obj_guess[None,:,:]
self.obj = t.nn.Parameter(obj_guess.to(t.float32))
self.obj_real = t.nn.Parameter(obj_guess.real.to(t.float32))
self.obj_imag = t.nn.Parameter(obj_guess.imag.to(t.float32))
# Wait for LBFGS to be updated for complex-valued parameters
#self.obj = t.nn.Parameter(obj_guess.to(t.float32))
if background is None:
if detector_slice is not None:
@@ -146,7 +154,7 @@ class RPI(CDIModel):
oversampling=oversampling)
if not isinstance(probe,t.Tensor):
probe = cmath.complex_to_torch(probe)
probe = t.as_tensor(probe)
# Potentially need all of this orientation stuff later
@@ -190,8 +198,8 @@ class RPI(CDIModel):
# I think something to do with the fact that the object is defined
# on a coarser grid needs to be accounted for here that is not
# accounted for yet
scale = t.sum(patterns[0]) / t.sum(cmath.cabssq(probe))
obj_guess = scale * cmath.expi(2 * np.pi * t.rand([n_modes,]+obj_size))
scale = t.sum(patterns[0]) / t.sum(t.abs(probe)**2)
obj_guess = scale * t.exp(2j * np.pi * t.rand([n_modes,]+obj_size))
elif initialization.lower().strip() == 'spectral':
if background is not None:
obj_guess = initializers.RPI_spectral_init(
@@ -218,9 +226,9 @@ class RPI(CDIModel):
def random_init(self, pattern):
scale = t.sum(pattern) / t.sum(cmath.cabssq(self.probe))
self.obj.data = scale * cmath.expi(
2 * np.pi * t.rand(self.obj.shape[:-1])).to(
scale = t.sum(pattern) / t.sum(t.abs(self.probe)**2)
self.obj.data = scale * t.exp(
2j * np.pi * t.rand(self.obj.shape)).to(
dtype=self.obj.dtype, device=self.obj.device)
def spectral_init(self, pattern):
@@ -249,8 +257,8 @@ class RPI(CDIModel):
pr = self.probe[i]
# Here we have a 3D probe (one single mode)
# and a 4D object (multiple modes mixing incoherently)
exit_waves = RPI_interaction(pr[:,:,:],
self.obj_support[None,:,:] * self.obj)
exit_waves = RPI_interaction(pr,
self.obj_support * self.obj)
all_exit_waves.append(exit_waves)
@@ -266,7 +274,7 @@ class RPI(CDIModel):
try:
# will fail if index has no length, for example when index
# is just an int. In this case, we just do nothing instead
output = output.unsqueeze(1).repeat(1,len(index),1,1,1)
output = output.unsqueeze(0).repeat(1,len(index),1,1,1)
except TypeError:
pass
@@ -297,8 +305,8 @@ class RPI(CDIModel):
#return tools.losses.poisson_nll(real_data, sim_data, mask=mask)
def regularizer(self, factors):
return factors[0] * t.sum(cmath.cabssq(self.obj[0,:,:,:])) \
+ factors[1] * t.sum(cmath.cabssq(self.obj[1:,:,:,:]))
return factors[0] * t.sum(t.abs(self.obj[0,:,:])**2) \
+ factors[1] * t.sum(t.abs(self.obj[1:,:,:])**2)
def to(self, *args, **kwargs):
super(RPI, self).to(*args, **kwargs)
@@ -330,7 +338,7 @@ class RPI(CDIModel):
plot_list = [
('Root Sum Squared Amplitude of all Probes',
lambda self, fig: p.plot_amplitude(
np.sqrt(np.sum(cmath.cabssq(self.probe).cpu().numpy(),axis=0)),
np.sqrt(np.sum((t.abs(self.probe)**2).cpu().numpy(),axis=0)),
fig=fig, basis=self.probe_basis)),
('Dominant Object Amplitude',
lambda self, fig: p.plot_amplitude(self.obj[0], fig=fig,
@@ -355,13 +363,13 @@ class RPI(CDIModel):
# continues to use that standard pattern
probe_basis = self.probe_basis.detach().cpu().numpy()
obj_basis = self.obj_basis.detach().cpu().numpy()
probe = cmath.torch_to_complex(self.probe.detach().cpu())
probe = self.probe.detach().cpu().numpy()
# Provide the option to save out the subdominant objects or
# just the dominant one
if full_obj:
obj = cmath.torch_to_complex(self.obj.detach().cpu())
obj = self.obj.detach().cpu().numpy()
else:
obj = cmath.torch_to_complex(self.obj[0].detach().cpu())
obj = self.obj[0].detach().cpu().numpy()
background = self.background.detach().cpu().numpy()**2
return {'probe_basis': probe_basis, 'obj_basis': obj_basis,
CDTools/models/simple_ptycho.py
+18 -16
View File
@@ -10,6 +10,7 @@ from torch.utils import data as torchdata
from matplotlib import pyplot as plt
from datetime import datetime
import numpy as np
from .complex_adam import MyAdam
__all__ = ['SimplePtycho']
@@ -25,22 +26,22 @@ class SimplePtycho(CDIModel):
surface_normal=np.array([0.,0.,1.]), mask=None):
super(SimplePtycho,self).__init__()
self.wavelength = t.Tensor([wavelength])
self.wavelength = t.tensor([wavelength])
self.detector_geometry = copy(detector_geometry)
det_geo = self.detector_geometry
if hasattr(det_geo, 'distance'):
det_geo['distance'] = t.Tensor(det_geo['distance'])
det_geo['distance'] = t.tensor(det_geo['distance'])
if hasattr(det_geo, 'basis'):
det_geo['basis'] = t.Tensor(det_geo['basis'])
det_geo['basis'] = t.tensor(det_geo['basis'])
if hasattr(det_geo, 'corner'):
det_geo['corner'] = t.Tensor(det_geo['corner'])
det_geo['corner'] = t.tensor(det_geo['corner'])
self.min_translation = t.Tensor(min_translation)
self.min_translation = t.tensor(min_translation)
self.probe_basis = t.Tensor(probe_basis)
self.probe_basis = t.tensor(probe_basis)
self.detector_slice = detector_slice
self.surface_normal = t.Tensor(surface_normal)
self.surface_normal = t.tensor(surface_normal)
if mask is None:
self.mask = None
@@ -49,11 +50,11 @@ class SimplePtycho(CDIModel):
# We rescale the probe here so it learns at the same rate as the
# object
self.probe_norm = t.max(tools.cmath.cabs(probe_guess.to(t.float32)))
self.probe_norm = t.max(t.abs(probe_guess.to(t.complex64)))
self.probe = t.nn.Parameter(probe_guess.to(t.float32)
self.probe = t.nn.Parameter(probe_guess.to(t.complex64)
/ self.probe_norm)
self.obj = t.nn.Parameter(obj_guess.to(t.float32))
self.obj = t.nn.Parameter(obj_guess.to(t.complex64))
@@ -94,8 +95,8 @@ class SimplePtycho(CDIModel):
# Finally, initialize the probe and object using this information
probe = tools.initializers.SHARP_style_probe(dataset, probe_shape, det_slice)
obj = t.ones(obj_size+(2,))
obj = t.ones(obj_size).to(dtype=t.complex64)
det_geo = dataset.detector_geometry
@@ -112,6 +113,7 @@ class SimplePtycho(CDIModel):
translations,
surface_normal=self.surface_normal)
pix_trans -= self.min_translation
return tools.interactions.ptycho_2D_round(self.probe_norm * self.probe,
self.obj,
pix_trans)
@@ -130,7 +132,7 @@ class SimplePtycho(CDIModel):
detector_slice=self.detector_slice)
def loss(self, sim_data, real_data, mask=None):
def loss(self, real_data, sim_data, mask=None):
return tools.losses.amplitude_mse(real_data, sim_data, mask=mask)
@@ -206,9 +208,9 @@ class SimplePtycho(CDIModel):
def save_results(self):
probe = tools.cmath.torch_to_complex(self.probe.detach().cpu())
probe = self.probe.detach().cpu().numpy()
probe = probe * self.probe_norm.detach().cpu().numpy()
obj = tools.cmath.torch_to_complex(self.obj.detach().cpu())
obj = self.obj.detach().cpu().numpy()
return {'probe':probe,'obj':obj}
@@ -272,4 +274,4 @@ class SimplePtycho(CDIModel):
# Calculate loss
loss.append(self.loss(self.measurement(self.interaction(i, translations)), patterns))
yield t.mean(t.Tensor(loss)).cpu().numpy()
yield t.mean(t.tensor(loss)).cpu().numpy()
CDTools/models/unified_mode_ptycho.py
-449
View File
@@ -1,449 +0,0 @@
from __future__ import division, print_function, absolute_import
import torch as t
from CDTools.models import CDIModel
from CDTools.datasets import Ptycho2DDataset
from CDTools import tools
from CDTools.tools import cmath
from CDTools.tools import analysis
from CDTools.tools import plotting as p
from matplotlib import pyplot as plt
from datetime import datetime
import numpy as np
from copy import copy
__all__ = ['UnifiedModePtycho']
class UnifiedModePtycho(CDIModel):
def __init__(self, wavelength, detector_geometry,
probe_basis,
probe_guess, obj_guess, rhos_guess,
detector_slice=None,
surface_normal=np.array([0.,0.,1.]),
min_translation = t.Tensor([0,0]),
background = None, translation_offsets=None, mask=None,
translation_scale = 1, saturation=None,
probe_support = None, obj_support=None, oversampling=1):
super(UnifiedModePtycho,self).__init__()
self.wavelength = t.Tensor([wavelength])
self.detector_geometry = copy(detector_geometry)
det_geo = self.detector_geometry
if hasattr(det_geo, 'distance'):
det_geo['distance'] = t.Tensor(det_geo['distance'])
if hasattr(det_geo, 'basis'):
det_geo['basis'] = t.Tensor(det_geo['basis'])
if hasattr(det_geo, 'corner'):
det_geo['corner'] = t.Tensor(det_geo['corner'])
self.min_translation = t.Tensor(min_translation)
self.probe_basis = t.Tensor(probe_basis)
self.detector_slice = detector_slice
self.surface_normal = t.Tensor(surface_normal)
self.saturation = saturation
if mask is None:
self.mask = mask
else:
self.mask = t.BoolTensor(mask)
# We rescale the probe here so it learns at the same rate as the
# object
if probe_guess.dim() > 3:
self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess[0].to(t.float32)))
else:
self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess.to(t.float32)))
self.probe = t.nn.Parameter(probe_guess.to(t.float32)
/ 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])
else:
background = 1e-6 * t.ones(self.probe[0].shape[:-1])
self.background = t.nn.Parameter(t.Tensor(background).to(t.float32))
self.rhos = t.nn.Parameter(t.Tensor(rhos_guess).to(t.float32))
if translation_offsets is None:
self.translation_offsets = None
else:
self.translation_offsets = t.nn.Parameter(t.Tensor(translation_offsets).to(t.float32)/ translation_scale)
self.translation_scale = translation_scale
if probe_support is not None:
self.probe_support = probe_support
else:
self.probe_support = t.ones_like(self.probe[0])
if obj_support is not None:
self.obj_support = obj_support
self.obj.data = self.obj * obj_support
else:
self.obj_support = t.ones_like(self.obj)
self.oversampling = oversampling
@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, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=True, opt_for_fft=False, mixing_mode='unified'):
wavelength = dataset.wavelength
det_basis = dataset.detector_geometry['basis']
det_shape = dataset[0][1].shape
distance = dataset.detector_geometry['distance']
# always do this on the cpu
get_as_args = dataset.get_as_args
dataset.get_as(device='cpu')
(indices, translations), patterns = dataset[:]
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,
opt_for_fft=opt_for_fft,
oversampling=oversampling)
if hasattr(dataset, 'sample_info') and \
dataset.sample_info is not None and \
'orientation' in dataset.sample_info:
surface_normal = dataset.sample_info['orientation'][2]
else:
surface_normal = np.array([0.,0.,1.])
# If this information is supplied when the function is called,
# then we override the information in the .cxi file
if scattering_mode in {'t', 'transmission'}:
surface_normal = np.array([0.,0.,1.])
elif scattering_mode in {'r', 'reflection'}:
outgoing_dir = np.cross(det_basis[:,0], det_basis[:,1])
outgoing_dir /= np.linalg.norm(outgoing_dir)
surface_normal = outgoing_dir + np.array([0.,0.,1.])
surface_normal /= -np.linalg.norm(surface_normal)
# 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)
obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=200)
if hasattr(dataset, 'background') and dataset.background is not None:
background = t.sqrt(dataset.background)
else:
background = None
# 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)
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))
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)
#probe = t.stack([tools.propagators.far_field(probe),] + probe_stack)
obj = tools.cmath.expi(randomize_ang * (t.rand(obj_size)-0.5))
det_geo = dataset.detector_geometry
translation_offsets = 0 * (t.rand((len(dataset),2)) - 0.5)
#
if mixing_mode.lower().strip() == 'unified':
rhos = t.zeros(len(dataset),n_modes,n_modes)
rhos[:,0,0] = 1
for i in range(1,n_modes):
rhos[:,i,i] = 1/n_modes
if hasattr(dataset, 'mask') and dataset.mask is not None:
mask = dataset.mask.to(t.bool)
else:
mask = None
if probe_support_radius is not None:
probe_support = t.zeros_like(probe[0].to(dtype=t.float32))
p_cent = np.array(probe.shape[1:3]).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
probe = probe * probe_support[None,:,:]
else:
probe_support = None;
if restrict_obj != -1:
ro = restrict_obj
os = np.array(obj_size)
ps = np.array(probe_shape)
obj_support = t.zeros_like(obj.to(dtype=t.float32))
obj_support[ps[0]//2-ro:os[0]+ro-ps[0]//2,
ps[1]//2-ro:os[1]+ro-ps[1]//2] = 1
else:
obj_support = None
return cls(wavelength, det_geo, probe_basis, probe, obj, rhos,
detector_slice=det_slice,
surface_normal=surface_normal,
min_translation=min_translation,
translation_offsets = translation_offsets,
mask=mask, background=background,
translation_scale=translation_scale,
saturation=saturation,
probe_support=probe_support,
obj_support=obj_support,
oversampling=oversampling)
def interaction(self, index, translations):
pix_trans = tools.interactions.translations_to_pixel(self.probe_basis,
translations,
surface_normal=self.surface_normal)
pix_trans -= self.min_translation
if self.translation_offsets is not None:
pix_trans += self.translation_scale * self.translation_offsets[index]
probes = []
all_exit_waves = []
for i in range(self.probe.shape[0]):
# from storing the probe in Fourier space
#pr = tools.propagators.inverse_far_field(self.probe[i]) * self.probe_support
pr = self.probe[i] * self.probe_support
exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc(pr,
self.obj_support * self.obj,
pix_trans,
shift_probe=True)
exit_waves = exit_waves * self.probe_support[...,:,:]
all_exit_waves.append(exit_waves)
return t.stack(all_exit_waves)
def forward_propagator(self, wavefields):
return tools.propagators.far_field(wavefields)
def backward_propagator(self, wavefields):
return tools.propagators.inverse_far_field(wavefields)
def measurement(self, wavefields, indices):
#return tools.measurements.density_matrix(wavefields,self.rhos[indices],
# detector_slice=self.detector_slice,
# saturation=self.saturation,
# oversampling=self.oversampling)
return tools.measurements.quadratic_background(wavefields,
self.background, self.rhos[indices],
detector_slice=self.detector_slice,
measurement=tools.measurements.density_matrix,
saturation=self.saturation,
oversampling=self.oversampling)
def forward(self, *args):
"""The complete forward model
We need to override this to enable the wavefield mixing at the
level of the measurement function
"""
indices = args[0]
return self.measurement(self.forward_propagator(self.interaction(*args)),indices)
def loss(self, sim_data, real_data, mask=None):
return tools.losses.amplitude_mse(real_data, sim_data, mask=mask)
#return tools.losses.poisson_nll(real_data, sim_data, mask=mask)
def to(self, *args, **kwargs):
super(UnifiedModePtycho, self).to(*args, **kwargs)
self.wavelength = self.wavelength.to(*args,**kwargs)
# move the detector geometry too
det_geo = self.detector_geometry
if hasattr(det_geo, 'distance'):
det_geo['distance'] = det_geo['distance'].to(*args,**kwargs)
if hasattr(det_geo, 'basis'):
det_geo['basis'] = det_geo['basis'].to(*args,**kwargs)
if hasattr(det_geo, 'corner'):
det_geo['corner'] = det_geo['corner'].to(*args,**kwargs)
if self.mask is not None:
self.mask = self.mask.to(*args, **kwargs)
self.min_translation = self.min_translation.to(*args,**kwargs)
self.probe_basis = self.probe_basis.to(*args,**kwargs)
self.probe_norm = self.probe_norm.to(*args,**kwargs)
self.probe_support = self.probe_support.to(*args,**kwargs)
self.obj_support = self.obj_support.to(*args,**kwargs)
self.surface_normal = self.surface_normal.to(*args, **kwargs)
def sim_to_dataset(self, args_list):
# In the future, potentially add more control
# over what metadata is saved (names, etc.)
# First, I need to gather all the relevant data
# that needs to be added to the dataset
entry_info = {'program_name': 'CDTools',
'instrument_n': 'Simulated Data',
'start_time': datetime.now()}
surface_normal = self.surface_normal.detach().cpu().numpy()
xsurfacevec = np.cross(np.array([0.,1.,0.]), surface_normal)
xsurfacevec /= np.linalg.norm(xsurfacevec)
ysurfacevec = np.cross(surface_normal, xsurfacevec)
ysurfacevec /= np.linalg.norm(ysurfacevec)
orientation = np.array([xsurfacevec, ysurfacevec, surface_normal])
sample_info = {'description': 'A simulated sample',
'orientation': orientation}
detector_geometry = self.detector_geometry
mask = self.mask
wavelength = self.wavelength
indices, translations = args_list
# Then we simulate the results
data = self.forward(indices, translations)
# And finally, we make the dataset
return Ptycho2DDataset(translations, data,
entry_info = entry_info,
sample_info = sample_info,
wavelength=wavelength,
detector_geometry=detector_geometry,
mask=mask)
def get_rhos(self):
rhos_out = np.zeros([self.rhos.shape[0],
self.rhos.shape[1],self.rhos.shape[2]],
dtype=np.complex64)
for (i,j) in ((i,j) for i in range(rhos_out.shape[1])
for j in range(rhos_out.shape[2])):
if i == j:
rhos_out[:,i,j] += self.rhos.data[:,i,j].cpu().detach().numpy()
if i < j: # upper triangle, real part
rhos_out[:,i,j] += self.rhos.data[:,i,j].cpu().detach().numpy()
rhos_out[:,j,i] += self.rhos.data[:,i,j].cpu().detach().numpy()
if i > j: # upper triangle, real part
rhos_out[:,j,i] += 1j * self.rhos.data[:,i,j].cpu().detach().numpy()
rhos_out[:,i,j] -= 1j * self.rhos.data[:,i,j].cpu().detach().numpy()
return rhos_out
def tidy_probes(self, normalization=1):
"""Tidies up the probes
What we want to do here is use all the information on all the probes
to calculate a natural basis for the experiment, and update all the
density matrices to operate in that updated basis
"""
rhos = self.get_rhos()
overall_rho = np.mean(rhos,axis=0)
probe = cmath.torch_to_complex(self.probe.detach().cpu())
ortho_probes, A = analysis.orthogonalize_probes(probe,
density_matrix=overall_rho,
keep_transform=True,
normalize=True)
Aconj = A.conj()
Atrans = np.transpose(A)
new_rhos = np.swapaxes(np.dot(Atrans,np.dot(rhos,Aconj)),0,1)
new_rhos /= normalization
ortho_probes *= np.sqrt(normalization)
#print(np.dot(Aconjinv,rhos).shape)
#print(np.dot(rhos,Aconj).shape)
new_rhos = cmath.complex_to_torch(new_rhos).to(
dtype=self.rhos.dtype,device=self.rhos.device)
# This repacks the data into the format used internally
for (i,j) in ((i,j) for i in range(new_rhos.shape[1])
for j in range(new_rhos.shape[2])):
if i == j:
self.rhos.data[:,i,j] = new_rhos[:,i,j,0]
if i < j: # upper triangle, real part
self.rhos.data[:,i,j] = new_rhos[:,i,j,0]
self.rhos.data[:,j,i] = new_rhos[:,i,j,1]
self.probe.data = cmath.complex_to_torch(ortho_probes).to(
device=self.probe.device,dtype=self.probe.dtype)
def corrected_translations(self,dataset):
translations = dataset.translations.to(dtype=self.probe.dtype,device=self.probe.device)
t_offset = tools.interactions.pixel_to_translations(self.probe_basis,self.translation_offsets*self.translation_scale,surface_normal=self.surface_normal)
return translations + t_offset
# Needs to be updated to allow for plotting to an existing figure
plot_list = [
('Dominant Probe Amplitude',
lambda self, fig: p.plot_amplitude(self.probe[0], fig=fig, basis=self.probe_basis)),
('Dominant Probe Phase',
lambda self, fig: p.plot_phase(self.probe[0], fig=fig, basis=self.probe_basis)),
('Subdominant Probe Amplitude',
lambda self, fig: p.plot_amplitude(self.probe[1], fig=fig, basis=self.probe_basis),
lambda self: len(self.probe) >=2),
('Subdominant Probe Phase',
lambda self, fig: p.plot_phase(self.probe[1], fig=fig, basis=self.probe_basis),
lambda self: len(self.probe) >=2),
('Average Density Matrix Amplitudes',
lambda self, fig: p.plot_amplitude(np.mean(np.abs(self.get_rhos()),axis=0), fig=fig),
lambda self: len(self.probe) >=2),
('Von Neumann Entropy (only accurate after tidy_probes)',
lambda self, fig, dataset: p.plot_nanomap(self.corrected_translations(dataset), analysis.calc_vn_entropy(self.get_rhos()), fig=fig)),
('Object Amplitude',
lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis)),
('Object Phase',
lambda self, fig: p.plot_phase(self.obj, fig=fig, basis=self.probe_basis)),
('Corrected Translations',
lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig)),
('Background',
lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2))
]
def save_results(self, dataset):
basis = self.probe_basis.detach().cpu().numpy()
translations = self.corrected_translations(dataset).detach().cpu().numpy()
probe = cmath.torch_to_complex(self.probe.detach().cpu())
probe = probe * self.probe_norm.detach().cpu().numpy()
obj = cmath.torch_to_complex(self.obj.detach().cpu())
background = self.background.detach().cpu().numpy()**2
weights = self.weights.detach().cpu().numpy()
return {'basis':basis, 'translation':translations,
'probe':probe,'obj':obj,
'background':background,
'weights':weights}
CDTools/models/unified_mode_ptycho2.py
-490
View File
@@ -1,490 +0,0 @@
from __future__ import division, print_function, absolute_import
import torch as t
from CDTools.models import CDIModel
from CDTools.datasets import Ptycho2DDataset
from CDTools import tools
from CDTools.tools import cmath
from CDTools.tools import analysis
from CDTools.tools import plotting as p
from matplotlib import pyplot as plt
from datetime import datetime
import numpy as np
from copy import copy
__all__ = ['UnifiedModePtycho2']
class UnifiedModePtycho2(CDIModel):
def __init__(self, wavelength, detector_geometry,
probe_basis,
probe_guess, obj_guess, Ws_guess,
detector_slice=None,
surface_normal=np.array([0.,0.,1.]),
min_translation = t.Tensor([0,0]),
background = None, translation_offsets=None, mask=None,
translation_scale = 1, saturation=None,
probe_support = None, obj_support=None, oversampling=1):
super(UnifiedModePtycho2,self).__init__()
self.wavelength = t.Tensor([wavelength])
self.detector_geometry = copy(detector_geometry)
det_geo = self.detector_geometry
if hasattr(det_geo, 'distance'):
det_geo['distance'] = t.Tensor(det_geo['distance'])
if hasattr(det_geo, 'basis'):
det_geo['basis'] = t.Tensor(det_geo['basis'])
if hasattr(det_geo, 'corner'):
det_geo['corner'] = t.Tensor(det_geo['corner'])
self.min_translation = t.Tensor(min_translation)
self.probe_basis = t.Tensor(probe_basis)
self.detector_slice = detector_slice
self.surface_normal = t.Tensor(surface_normal)
self.saturation = saturation
if mask is None:
self.mask = mask
else:
self.mask = t.BoolTensor(mask)
# We rescale the probe here so it learns at the same rate as the
# object
if probe_guess.dim() > 3:
self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess[0].to(t.float32)))
else:
self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess.to(t.float32)))
self.probe = t.nn.Parameter(probe_guess.to(t.float32)
/ 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])
else:
background = 1e-6 * t.ones(self.probe[0].shape[:-1])
self.background = t.nn.Parameter(t.Tensor(background).to(t.float32))
if type(Ws_guess) == type(t.zeros(1)):
self.Ws = t.nn.Parameter(Ws_guess.to(t.float32))
else:
self.Ws = t.nn.Parameter(cmath.complex_to_torch(Ws_guess).to(t.float32))
if translation_offsets is None:
self.translation_offsets = None
else:
self.translation_offsets = t.nn.Parameter(t.Tensor(translation_offsets).to(t.float32)/ translation_scale)
self.translation_scale = translation_scale
if probe_support is not None:
self.probe_support = probe_support
else:
self.probe_support = t.ones_like(self.probe[0])
if obj_support is not None:
self.obj_support = obj_support
self.obj.data = self.obj * obj_support
else:
self.obj_support = t.ones_like(self.obj)
self.oversampling = oversampling
@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, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=True, opt_for_fft=False, dm_rank=0):
wavelength = dataset.wavelength
det_basis = dataset.detector_geometry['basis']
det_shape = dataset[0][1].shape
distance = dataset.detector_geometry['distance']
# always do this on the cpu
get_as_args = dataset.get_as_args
dataset.get_as(device='cpu')
(indices, translations), patterns = dataset[:]
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,
opt_for_fft=opt_for_fft,
oversampling=oversampling)
if hasattr(dataset, 'sample_info') and \
dataset.sample_info is not None and \
'orientation' in dataset.sample_info:
surface_normal = dataset.sample_info['orientation'][2]
else:
surface_normal = np.array([0.,0.,1.])
# If this information is supplied when the function is called,
# then we override the information in the .cxi file
if scattering_mode in {'t', 'transmission'}:
surface_normal = np.array([0.,0.,1.])
elif scattering_mode in {'r', 'reflection'}:
outgoing_dir = np.cross(det_basis[:,0], det_basis[:,1])
outgoing_dir /= np.linalg.norm(outgoing_dir)
surface_normal = outgoing_dir + np.array([0.,0.,1.])
surface_normal /= -np.linalg.norm(surface_normal)
# 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)
obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=200)
if hasattr(dataset, 'background') and dataset.background is not None:
background = t.sqrt(dataset.background)
else:
background = None
# 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)
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))
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)
#probe = t.stack([tools.propagators.far_field(probe),] + probe_stack)
obj = tools.cmath.expi(randomize_ang * (t.rand(obj_size)-0.5))
det_geo = dataset.detector_geometry
translation_offsets = 0 * (t.rand((len(dataset),2)) - 0.5)
# dm_rank defines the rank of the shot-by-shot density matrices
if dm_rank > n_modes:
raise KeyError('Density matrix rank cannot be greater than the number of modes')
elif dm_rank != 0:
if dm_rank == -1:
dm_rank = n_modes
Ws = t.zeros(len(dataset),dm_rank,n_modes,2)
Ws[:,0,0,0] = 1
for i in range(1,dm_rank):
Ws[:,i,i,0] = 1/np.sqrt(n_modes)
else:
# dm_rank=0 is a special case defining a purely stable, incoherent
# mode mixing model. This is passed on by defining a set of weights
# which only has one index
Ws = t.ones(len(dataset))
if hasattr(dataset, 'mask') and dataset.mask is not None:
mask = dataset.mask.to(t.bool)
else:
mask = None
if probe_support_radius is not None:
probe_support = t.zeros_like(probe[0].to(dtype=t.float32))
p_cent = np.array(probe.shape[1:3]).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
probe = probe * probe_support[None,:,:]
else:
probe_support = None;
if restrict_obj != -1:
ro = restrict_obj
os = np.array(obj_size)
ps = np.array(probe_shape)
obj_support = t.zeros_like(obj.to(dtype=t.float32))
obj_support[ps[0]//2-ro:os[0]+ro-ps[0]//2,
ps[1]//2-ro:os[1]+ro-ps[1]//2] = 1
else:
obj_support = None
return cls(wavelength, det_geo, probe_basis, probe, obj, Ws,
detector_slice=det_slice,
surface_normal=surface_normal,
min_translation=min_translation,
translation_offsets = translation_offsets,
mask=mask, background=background,
translation_scale=translation_scale,
saturation=saturation,
probe_support=probe_support,
obj_support=obj_support,
oversampling=oversampling)
def interaction(self, index, translations):
pix_trans = tools.interactions.translations_to_pixel(self.probe_basis,
translations,
surface_normal=self.surface_normal)
pix_trans -= self.min_translation
if self.translation_offsets is not None:
pix_trans += self.translation_scale * self.translation_offsets[index]
Ws = self.Ws[index]
if type(index) == type(0):
index = [index]
Ws = [Ws]
pix_trans = [pix_trans]
single_frame = True
else:
single_frame = False
probes = []
all_exit_waves = []
# This is the case if a purely stable, incoherent model is defined.
if len(Ws[0].shape) == 0:
# What we do here is generate an identity matrix, and multiply
# that identity matrix by the per-frame weight.
Ws = [W * t.stack([t.eye(self.probe.shape[0]),
t.zeros([self.probe.shape[0]]*2)],dim=-1).to(
dtype=W.dtype, device=W.device)
for W in Ws]
# Outer iteration is the mode index iteration
for i in range(Ws[0].shape[0]):
# Now we need to separately treat each mode
exit_waves = []
for W, pix_tran in zip(Ws, pix_trans):
# from storing the probe in Fourier space
pr = [cmath.cmult(W[i,j,:], self.probe[j] * self.probe_support)
for j in range(self.probe.shape[0])]
pr = t.sum(t.stack(pr), axis=0)
exit_waves.append(self.probe_norm *
tools.interactions.ptycho_2D_sinc(
pr, self.obj_support * self.obj,
pix_tran, shift_probe=True))
exit_waves = t.stack(exit_waves)
if single_frame:
exit_waves = exit_waves[0]
# Multiply again by probe support to suppress the fringes from the
# sinc-interpolated shift
exit_waves = exit_waves * self.probe_support[...,:,:]
all_exit_waves.append(exit_waves)
return t.stack(all_exit_waves)
def forward_propagator(self, wavefields):
return tools.propagators.far_field(wavefields)
def backward_propagator(self, wavefields):
return tools.propagators.inverse_far_field(wavefields)
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)
def loss(self, sim_data, real_data, mask=None):
return tools.losses.amplitude_mse(real_data, sim_data, mask=mask)
#return tools.losses.poisson_nll(real_data, sim_data, mask=mask)
def to(self, *args, **kwargs):
super(UnifiedModePtycho2, self).to(*args, **kwargs)
self.wavelength = self.wavelength.to(*args,**kwargs)
# move the detector geometry too
det_geo = self.detector_geometry
if hasattr(det_geo, 'distance'):
det_geo['distance'] = det_geo['distance'].to(*args,**kwargs)
if hasattr(det_geo, 'basis'):
det_geo['basis'] = det_geo['basis'].to(*args,**kwargs)
if hasattr(det_geo, 'corner'):
det_geo['corner'] = det_geo['corner'].to(*args,**kwargs)
if self.mask is not None:
self.mask = self.mask.to(*args, **kwargs)
self.min_translation = self.min_translation.to(*args,**kwargs)
self.probe_basis = self.probe_basis.to(*args,**kwargs)
self.probe_norm = self.probe_norm.to(*args,**kwargs)
self.probe_support = self.probe_support.to(*args,**kwargs)
self.obj_support = self.obj_support.to(*args,**kwargs)
self.surface_normal = self.surface_normal.to(*args, **kwargs)
def sim_to_dataset(self, args_list):
# In the future, potentially add more control
# over what metadata is saved (names, etc.)
# First, I need to gather all the relevant data
# that needs to be added to the dataset
entry_info = {'program_name': 'CDTools',
'instrument_n': 'Simulated Data',
'start_time': datetime.now()}
surface_normal = self.surface_normal.detach().cpu().numpy()
xsurfacevec = np.cross(np.array([0.,1.,0.]), surface_normal)
xsurfacevec /= np.linalg.norm(xsurfacevec)
ysurfacevec = np.cross(surface_normal, xsurfacevec)
ysurfacevec /= np.linalg.norm(ysurfacevec)
orientation = np.array([xsurfacevec, ysurfacevec, surface_normal])
sample_info = {'description': 'A simulated sample',
'orientation': orientation}
detector_geometry = self.detector_geometry
mask = self.mask
wavelength = self.wavelength
indices, translations = args_list
# Then we simulate the results
data = self.forward(indices, translations)
# And finally, we make the dataset
return Ptycho2DDataset(translations, data,
entry_info = entry_info,
sample_info = sample_info,
wavelength=wavelength,
detector_geometry=detector_geometry,
mask=mask)
def get_rhos(self):
# If this is not a purely stable model
if len(self.Ws.shape) >= 2:
Ws = cmath.torch_to_complex(self.Ws.detach().cpu())
rhos_out = np.matmul(np.swapaxes(Ws,1,2), Ws.conj())
return rhos_out
else:
return np.array([np.eye(self.probe.shape[0])]*self.Ws.shape[0],
dtype=np.complex64)
def tidy_probes(self, normalization=1):
"""Tidies up the probes
What we want to do here is use all the information on all the probes
to calculate a natural basis for the experiment, and update all the
density matrices to operate in that updated basis
"""
# Must also implement a version that works appropriately with
# a purely incoherent model
#
# Note to future: We could probably do this more cleanly with an
# SVD directly on the Ws matrix, instead of an eigendecomposition
# of the rho matrix. This could avoid potential stability issues
# due to the existence of zero eigenvalues in the full rho matrix
# when dm_rank < n_modes
#
rhos = self.get_rhos()
overall_rho = np.mean(rhos,axis=0)
probe = cmath.torch_to_complex(self.probe.detach().cpu())
ortho_probes, A = analysis.orthogonalize_probes(probe,
density_matrix=overall_rho,
keep_transform=True,
normalize=True)
Aconj = A.conj()
Atrans = np.transpose(A)
new_rhos = np.matmul(Atrans,np.matmul(rhos,Aconj))
new_rhos /= normalization
ortho_probes *= np.sqrt(normalization)
dm_rank = self.Ws.shape[1]
new_Ws = []
for rho in new_rhos:
# These are returned from smalles to largest - we want to keep
# the largest ones
w,v = np.linalg.eigh(rho)
w = w[::-1][:dm_rank]
v = v[:,::-1][:,:dm_rank]
new_Ws.append(np.dot(np.diag(np.sqrt(w)),v.transpose()))
new_Ws = np.array(new_Ws)
self.Ws.data = cmath.complex_to_torch(new_Ws).to(
dtype=self.Ws.dtype,device=self.Ws.device)
self.probe.data = cmath.complex_to_torch(ortho_probes).to(
device=self.probe.device,dtype=self.probe.dtype)
def corrected_translations(self,dataset):
translations = dataset.translations.to(dtype=self.probe.dtype,device=self.probe.device)
t_offset = tools.interactions.pixel_to_translations(self.probe_basis,self.translation_offsets*self.translation_scale,surface_normal=self.surface_normal)
return translations + t_offset
# Needs to be updated to allow for plotting to an existing figure
plot_list = [
('Dominant Probe Amplitude',
lambda self, fig: p.plot_amplitude(self.probe[0], fig=fig, basis=self.probe_basis)),
('Dominant Probe Phase',
lambda self, fig: p.plot_phase(self.probe[0], fig=fig, basis=self.probe_basis)),
('Subdominant Probe Amplitude',
lambda self, fig: p.plot_amplitude(self.probe[1], fig=fig, basis=self.probe_basis),
lambda self: len(self.probe) >=2),
('Subdominant Probe Phase',
lambda self, fig: p.plot_phase(self.probe[1], fig=fig, basis=self.probe_basis),
lambda self: len(self.probe) >=2),
('Average Density Matrix Amplitudes',
lambda self, fig: p.plot_amplitude(np.mean(np.abs(self.get_rhos()),axis=0), fig=fig),
lambda self: len(self.Ws.shape) >=2),
('% Power in Top Mode (only accurate after tidy_probes)',
lambda self, fig, dataset: p.plot_nanomap(self.corrected_translations(dataset), analysis.calc_top_mode_fraction(self.get_rhos()), fig=fig),
lambda self: len(self.Ws.shape) >=2),
('Object Amplitude',
lambda self, fig: p.plot_amplitude(self.obj, fig=fig, basis=self.probe_basis)),
('Object Phase',
lambda self, fig: p.plot_phase(self.obj, fig=fig, basis=self.probe_basis)),
('Corrected Translations',
lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig)),
('Background',
lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2))
]
def save_results(self, dataset):
basis = self.probe_basis.detach().cpu().numpy()
translations = self.corrected_translations(dataset).detach().cpu().numpy()
probe = cmath.torch_to_complex(self.probe.detach().cpu())
probe = probe * self.probe_norm.detach().cpu().numpy()
obj = cmath.torch_to_complex(self.obj.detach().cpu())
background = self.background.detach().cpu().numpy()**2
Ws = cmath.torch_to_complex(self.Ws.detach().cpu())
return {'basis':basis, 'translation':translations,
'probe':probe,'obj':obj,
'background':background,
'Ws':Ws}
CDTools/tools/__init__.py
-2
View File
@@ -17,13 +17,11 @@ having numpy and torch defined under CDTools.tools.cmath, you know?
from __future__ import division, print_function, absolute_import
from CDTools.tools import cmath
from CDTools.tools import losses
from CDTools.tools import data
from CDTools.tools import image_processing
from CDTools.tools import initializers
from CDTools.tools import plotting
from CDTools.tools import projectors
from CDTools.tools import interactions
from CDTools.tools import propagators
from CDTools.tools import measurements
CDTools/tools/analysis/analysis.py
+60 -74
View File
@@ -5,14 +5,13 @@ or numpy arrays, so they can be used either directly after reconstructions
on the attributes of the models themselves, or after-the-fact once the
data has been stored in numpy arrays.
"""
from __future__ import division, print_function
import torch as t
import numpy as np
from CDTools.tools import cmath
from CDTools.tools import image_processing as ip
from scipy import fftpack
from scipy import linalg as sla
from scipy import special
__all__ = ['orthogonalize_probes', 'standardize', 'synthesize_reconstructions',
'calc_consistency_prtf', 'calc_deconvolved_cross_correlation',
@@ -62,12 +61,11 @@ def orthogonalize_probes(probes, density_matrix=None, keep_transform=False, norm
"""
try:
probes = cmath.torch_to_complex(probes.detach().cpu())
probes = probes.detach().cpu().numpy()
send_to_torch = True
except:
send_to_torch = False
from matplotlib import pyplot as plt
# We can do the orthogonalization with an SVD, so first we have to
# reshape the final two dimensions (the image shape) into a single
# vectorized dimension. This matrix is probes^dagger, hence the
@@ -82,15 +80,16 @@ def orthogonalize_probes(probes, density_matrix=None, keep_transform=False, norm
# next we want to extract the eigendecomposition of the density matrix
# itself
w,v = sla.eigh(density_matrix)
w,v = np.linalg.eigh(density_matrix)
w = w[::-1]
v = v[:,::-1]
# We do this just to avoid total failure when the density
# matrix is not positive definite.
w = np.abs(w)
# Note: this will fail if any w are less than zero
# Probably should figure out a good way to deal with that
# In most cases (such as when rho is generated directly from some other
# matrix A such that rho=A A^dagger), w should never have any negative
# entries.
w = np.maximum(w,0)
B_dagger = np.dot(np.diag(np.sqrt(w)), v.conj().transpose())
#u,s,vh = np.linalg.svd(np.dot(B_dagger,probes_mat), full_matrices=False)
@@ -115,9 +114,8 @@ def orthogonalize_probes(probes, density_matrix=None, keep_transform=False, norm
#A_dagger = np.dot(np.transpose(u).conj(),B_dagger)
if send_to_torch:
ortho_probes = cmath.complex_to_torch(np.stack(ortho_probes))
A = cmath.complex_to_torch(A)
#A_dagger = cmath.complex_to_torch(A_dagger)
ortho_probes = t.as_tensor(np.stack(ortho_probes))
A = t.as_tensor(A)
if keep_transform:
return ortho_probes, A#_dagger
@@ -171,22 +169,21 @@ def standardize(probe, obj, obj_slice=None, correct_ramp=False):
# First, we normalize the probe intensity to a fixed value.
probe_np = False
if isinstance(probe, np.ndarray):
probe = cmath.complex_to_torch(probe).to(t.float32)
probe = t.as_tensor(probe, dtype=t.complex64)
probe_np = True
obj_np = False
if isinstance(obj, np.ndarray):
obj = cmath.complex_to_torch(obj).to(t.float32)
obj = t.as_tensor(obj,dtype=t.complex64)
obj_np = True
# If this is a single probe and not a stack of probes
if len(probe.shape) == 3:
if len(probe.shape) == 2:
probe = probe[None,...]
single_probe = True
else:
single_probe = False
normalization = t.sqrt(t.sum(cmath.cabssq(probe[0])) / (len(probe[0].view(-1))/2))
normalization = t.sqrt(t.sum(t.abs(probe[0])**2) / (len(probe[0].view(-1))))
probe = probe / normalization
obj = obj * normalization
@@ -198,37 +195,38 @@ def standardize(probe, obj, obj_slice=None, correct_ramp=False):
if correct_ramp:
# Need to check if this is actually working and, if not, why not
center_freq = ip.centroid(cmath.cabssq(cmath.fftshift(t.fft(probe[0],2))))
center_freq -= (t.tensor(probe[0].shape[:-1]) // 2).to(t.float32)
center_freq /= t.tensor(probe[0].shape[:-1]).to(t.float32)
center_freq = ip.centroid(t.abs(t.fft.fftshift(t.fft.fft2(probe[0]),
dim=(-1,-2)))**2)
center_freq -= t.div(t.tensor(probe[0].shape,dtype=t.float32),2,rounding_mode='floor')
center_freq /= t.as_tensor(probe[0].shape,dtype=t.float32)
Is, Js = np.mgrid[:probe[0].shape[0],:probe[0].shape[1]]
probe_phase_ramp = cmath.expi(2 * np.pi *
(center_freq[0] * t.tensor(Is).to(t.float32) +
center_freq[1] * t.tensor(Js).to(t.float32)))
probe = cmath.cmult(probe, cmath.cconj(probe_phase_ramp))
probe_phase_ramp = t.exp(2j * np.pi *
(center_freq[0] * t.tensor(Is).to(t.float32) +
center_freq[1] * t.tensor(Js).to(t.float32)))
probe = probe * t.conj(probe_phase_ramp)
Is, Js = np.mgrid[:obj.shape[0],:obj.shape[1]]
obj_phase_ramp = cmath.expi(2*np.pi *
(center_freq[0] * t.tensor(Is).to(t.float32) +
center_freq[1] * t.tensor(Js).to(t.float32)))
obj = cmath.cmult(obj, obj_phase_ramp)
obj_phase_ramp = t.exp(2j*np.pi *
(center_freq[0] * t.tensor(Is).to(t.float32) +
center_freq[1] * t.tensor(Js).to(t.float32)))
obj = obj * obj_phase_ramp
# Then, we set them to consistent absolute phases
obj_angle = cmath.cphase(t.sum(obj[obj_slice],dim=(0,1)))
obj = cmath.cmult(obj, cmath.expi(-obj_angle))
obj_angle = t.angle(t.sum(obj[obj_slice]))
obj = obj * t.exp(-1j*obj_angle)
for i in range(probe.shape[0]):
probe_angle = cmath.cphase(t.sum(probe[i],dim=(0,1)))
probe[i] = cmath.cmult(probe[i], cmath.expi(-probe_angle))
probe_angle = t.angle(t.sum(probe[i]))
probe[i] = probe[i] * t.exp(-1j*probe_angle)
if single_probe:
probe = probe[0]
if probe_np:
probe = cmath.torch_to_complex(probe.detach().cpu())
probe = probe.detach().cpu().numpy()
if obj_np:
obj = cmath.torch_to_complex(obj.detach().cpu())
obj = obj.detach().cpu().numpy()
return probe, obj
@@ -266,16 +264,17 @@ def synthesize_reconstructions(probes, objects, use_probe=False, obj_slice=None,
A list of standardized objects, for further processing
"""
# This should be cleaned up so it accepts anything array_like
probe_np = False
if isinstance(probes[0], np.ndarray):
probes = [cmath.complex_to_torch(probe).to(t.float32) for probe in probes]
probes = [t.as_tensor(probe,dtype=t.complex64) for probe in probes]
probe_np = True
obj_np = False
if isinstance(objects[0], np.ndarray):
objects = [cmath.complex_to_torch(obj).to(t.float32) for obj in objects]
objects = [t.as_tensor(obj,dtype=t.complex64) for obj in objects]
obj_np = True
obj_shape = np.min(np.array([obj.shape[:-1] for obj in objects]),axis=0)
obj_shape = np.min(np.array([obj.shape for obj in objects]),axis=0)
objects = [obj[:obj_shape[0],:obj_shape[1]] for obj in objects]
if obj_slice is None:
@@ -298,8 +297,7 @@ def synthesize_reconstructions(probes, objects, use_probe=False, obj_slice=None,
obj = ip.sinc_subpixel_shift(obj,np.array(shift))
if len(probe.shape) == 4:
if len(probe.shape) == 3:
probe = t.stack([ip.sinc_subpixel_shift(p,tuple(shift))
for p in probe],dim=0)
else:
@@ -317,10 +315,10 @@ def synthesize_reconstructions(probes, objects, use_probe=False, obj_slice=None,
i = -1
if probe_np:
synth_probe = cmath.torch_to_complex(synth_probe)
synth_probe = synth_probe.numpy()
if obj_np:
synth_obj = cmath.torch_to_complex(synth_obj)
obj_stack = [cmath.torch_to_complex(obj) for obj in obj_stack]
synth_obj = synth_obj.numpy()
obj_stack = [obj.numpy() for obj in obj_stack]
return synth_probe/(i+2), synth_obj/(i+2), obj_stack
@@ -358,10 +356,10 @@ def calc_consistency_prtf(synth_obj, objects, basis, obj_slice=None,nbins=None):
obj_np = False
if isinstance(objects[0], np.ndarray):
objects = [cmath.complex_to_torch(obj).to(t.float32) for obj in objects]
objects = [t.as_tensor(obj, dtype=t.complex64) for obj in objects]
obj_np = True
if isinstance(synth_obj, np.ndarray):
synth_obj = cmath.complex_to_torch(synth_obj).to(t.float32)
synth_obj = t.as_tensor(synth_obj, dtype=t.complex64)
if isinstance(basis, t.Tensor):
basis = basis.detach().cpu().numpy()
@@ -373,7 +371,7 @@ def calc_consistency_prtf(synth_obj, objects, basis, obj_slice=None,nbins=None):
if nbins is None:
nbins = np.max(synth_obj[obj_slice].shape) // 4
synth_fft = cmath.cabssq(cmath.fftshift(t.fft(synth_obj[obj_slice],2))).numpy()
synth_fft = (t.abs(t.fft.fftshift(t.fft.fft2(synth_obj[obj_slice]), dim=(-1,-2)))**2).numpy()
di = np.linalg.norm(basis[:,0])
@@ -391,7 +389,8 @@ def calc_consistency_prtf(synth_obj, objects, basis, obj_slice=None,nbins=None):
prtfs = []
for obj in objects:
obj = obj[obj_slice]
single_fft = cmath.cabssq(cmath.fftshift(t.fft(obj,2))).numpy()
single_fft = (t.abs(t.fft.fftshift(t.fft.fft2(obj),
dim=(-1,-2)))**2).numpy()
single_ints, bins = np.histogram(Rs,bins=nbins,weights=single_fft)
prtfs.append(synth_ints/single_ints)
@@ -433,32 +432,26 @@ def calc_deconvolved_cross_correlation(im1, im2, im_slice=None):
im_np = False
if isinstance(im1, np.ndarray):
im1 = cmath.complex_to_torch(im1)
im1 = t.as_tensor(im1)
im_np = True
if isinstance(im2, np.ndarray):
im2 = cmath.complex_to_torch(im2)
im2 = t.as_tensor(im2)
im_np = True
# If last dimension is not 2, then convert to a complex tensor now
if im1.shape[-1] != 2:
im1 = t.stack((im1,t.zeros_like(im1)),dim=-1)
if im2.shape[-1] != 2:
im2 = t.stack((im2,t.zeros_like(im2)),dim=-1)
if im_slice is None:
im_slice = np.s_[(im1.shape[0]//8)*3:(im1.shape[0]//8)*5,
(im1.shape[1]//8)*3:(im1.shape[1]//8)*5]
cor_fft = cmath.cmult(t.fft(im1[im_slice],2),
cmath.cconj(t.fft(im2[im_slice],2)))
cor_fft = t.fft.fft2(im1[im_slice]) * \
t.conj(t.fft.fft2(im2[im_slice]))
# Not sure if this is more or less stable than just the correlation
# maximum - requires some testing
cor = t.ifft(cor_fft / cmath.cabs(cor_fft)[:,:,None],2)
cor = t.fft.ifft2(cor_fft / t.abs(cor_fft))
if im_np:
cor = cmath.torch_to_complex(cor)
cor = cor.numpy()
return cor
@@ -500,22 +493,15 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.):
im_np = False
if isinstance(im1, np.ndarray):
im1 = cmath.complex_to_torch(im1)
im1 = t.as_tensor(im1)
im_np = True
if isinstance(im2, np.ndarray):
im2 = cmath.complex_to_torch(im2)
im2 = t.as_tensor(im2)
im_np = True
if isinstance(basis, np.ndarray):
basis = t.tensor(basis)
# If last dimension is not 2, then convert to a complex tensor now
if im1.shape[-1] != 2:
im1 = t.stack((im1,t.zeros_like(im1)),dim=-1)
if im2.shape[-1] != 2:
im2 = t.stack((im2,t.zeros_like(im2)),dim=-1)
if im_slice is None:
im_slice = np.s_[(im1.shape[0]//8)*3:(im1.shape[0]//8)*5,
(im1.shape[1]//8)*3:(im1.shape[1]//8)*5]
@@ -524,11 +510,11 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.):
nbins = np.max(im1[im_slice].shape) // 4
cor_fft = cmath.cmult(cmath.fftshift(t.fft(im1[im_slice],2)),
cmath.fftshift(cmath.cconj(t.fft(im2[im_slice],2))))
cor_fft = t.fft.fftshift(t.fft.fft2(im1[im_slice]),dim=(-1,-2)) * \
t.fft.fftshift(t.conj(t.fft.fft2(im2[im_slice])),dim=(-1,-2))
F1 = cmath.cabs(cmath.fftshift(t.fft(im1[im_slice],2)))**2
F2 = cmath.cabs(cmath.fftshift(t.fft(im2[im_slice],2)))**2
F1 = t.abs(t.fft.fftshift(t.fft.fft2(im1[im_slice]),dim=(-1,-2)))**2
F2 = t.abs(t.fft.fftshift(t.fft.fft2(im2[im_slice]),dim=(-1,-2)))**2
di = np.linalg.norm(basis[:,0])
@@ -542,7 +528,7 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.):
numerator, bins = np.histogram(Rs,bins=nbins,weights=cmath.torch_to_complex(cor_fft))
numerator, bins = np.histogram(Rs,bins=nbins,weights=cor_fft.numpy())
denominator_F1, bins = np.histogram(Rs,bins=nbins,weights=F1.detach().cpu().numpy())
denominator_F2, bins = np.histogram(Rs,bins=nbins,weights=F2.detach().cpu().numpy())
n_pix, bins = np.histogram(Rs,bins=nbins)
@@ -588,12 +574,12 @@ def calc_vn_entropy(matrix):
# Normalize them to match standard density matrix form
eigs = [eig / np.sum(eig) for eig in eigs]
# And calculate the VN entropy!
entropies = [-np.sum(eig*np.log(eig)) for eig in eigs]
entropies = [-np.sum(special.xlogy(eig,eig)) for eig in eigs]
return np.array(entropies)
else:
eig = np.linalg.eigh(matrix)[0]
entropy = -np.sum(eig*np.log(eig))/np.sum(eig)
return -np.trace(np.dot(matrix,sla.logm(matrix)))
entropy = -np.sum(special.xlogy(eig,eig))/np.sum(eig)
return entropy
def calc_top_mode_fraction(matrix):
"""Calculates the fraction of total power in the top mode of a density matrix
CDTools/tools/atoms/atoms.py
+1 -1
View File
@@ -7,7 +7,7 @@ Appendix C. doi: 10.1007/978-1-4419-6533-2_11
They are actually taken from https://github.com/LeBeauGroup/pyMultislicer
written by Jim LeBeau, but above is the original reference.
"""
from __future__ import division, print_function, absolute_import
import numpy as np
import torch as t
from scipy import fftpack
CDTools/tools/cmath/__init__.py
-3
View File
@@ -1,3 +0,0 @@
from __future__ import division, print_function, absolute_import
from CDTools.tools.cmath.cmath import *
CDTools/tools/cmath/cmath.py
-327
View File
@@ -1,327 +0,0 @@
"""Contains basic functions for dealing with complex numbers in pytorch.
Since pytorch doesn't have built-in support for complex numbers, but the
fast fourier transforms in pytorch assume a specific format for complex
arrays, this module uses that format to store complex numbers. It exposes
functions for converting between complex numpy arrays and torch tensors
stored in that format, as well as basic complex math operations implemented
on the torch tensors
"""
from __future__ import division, print_function, absolute_import
import numpy as np
import torch as t
__all__ = ['complex_to_torch', 'torch_to_complex', 'cabssq', 'cabs', 'cconj',
'cmult', 'cdiv', 'cphase', 'fftshift', 'ifftshift', 'expi', 'cexpi']
#
# These define the conversions to and from this format
#
def complex_to_torch(x):
"""Maps a complex numpy array to a torch tensor
Pytorch uses tensors with a final dimension of 2 to represent
complex numbers. This maps a complex type numpy array to a torch
tensor following this convention
Parameters
----------
x : np.ndarray
A numpy array to convert
Returns
-------
torch.Tensor
A torch tensor representation of that array
"""
return t.from_numpy(np.stack((np.real(x),np.imag(x)),axis=-1))
def torch_to_complex(x):
"""Maps a torch tensor to the a complex numpy array
Pytorch uses tensors with a final dimension of 2 to represent
complex numbers. This maps a torch tensor following that convention
to the appropriate numpy complex array. Note that, in order for this
function to work, the tensor must be detached from any parameters and
living on the CPU.
Parameters
----------
x : torch.Tensor
A tensor to convert
Returns
-------
np.array
A complex typed numpy array corresponding to the input
"""
x = np.array(x)
x = x[...,0] + x[...,1] * 1j
return x
#
# And these define the basic operations on these arrays. Note that
# multiplication between a complex valued and real valued pytorch
# tensor will proceed as expected because of torch's broadcasting
# and thus doesn't need it's own function
#
def cabssq(x):
"""Returns the square of the absolute value of a complex torch tensor
Pytorch uses tensors with a final dimension of 2 to represent
complex numbers. This calculates the elementwise absolute value
squared of any toch tensor following that standard.
Parameters
----------
x : torch.Tensor
An input tensor
Returns
-------
torch.Tensor
A tensor storing the elementwise absolute value squared
"""
return x[...,0]**2 + x[...,1]**2
def cabs(x):
"""Returns the absolute value of a complex torch tensor
Pytorch uses tensors with a final dimension of 2 to represent
complex numbers. This calculates the elementwise absolute value
of any torch tensor following that standard.
Parameters
----------
x : torch.Tensor
An input tensor
Returns
-------
torch.Tensor
A tensor storing the elementwise absolute value
"""
return t.sqrt(cabssq(x))
def cphase(x):
"""Returns the phase of a complex torch tensor
Pytorch uses tensors with a final dimension of 2 to represent
complex numbers. This calculates the elementwise complex phase
of any torch tensor following that standard.
Parameters
----------
x : torch.Tensor
An input tensor
Returns
-------
torch.Tensor
A tensor storing the elementwise phase
"""
return t.atan2(x[...,1],x[...,0])
def cconj(x):
"""Returns the complex conjugate of a complex torch tensor
Pytorch uses tensors with a final dimension of 2 to represent
complex numbers. This calculates the elementwise complex conjugate
of any torch tensor following that standard.
Parameters
----------
x : torch.Tensor
An input tensor
Returns
-------
torch.Tensor
A tensor storing the elementwise complex conjugate
"""
return t.stack((x[...,0],-x[...,1]),dim=-1)
def cmult(a,b):
"""Returns the complex product of two torch tensors
Pytorch uses tensors with a final dimension of 2 to represent
complex numbers. This calculates the elementwise product
of two torch tensors following that standard.
Parameters
----------
a : torch.Tensor
An input tensor
b : torch.Tensor
A second input tensor
Returns
-------
torch.Tensor
A tensor storing the elementwise product
"""
real = a[...,0] * b[...,0] - a[...,1] * b[...,1]
imag = a[...,0] * b[...,1] + a[...,1] * b[...,0]
return t.stack((real,imag),dim=-1)
def cdiv(a,b):
"""Returns the complex quotient of two torch tensors
Pytorch uses tensors with a final dimension of 2 to represent
complex numbers. This calculates the elementwise quotient
of two torch tensors following that standard.
Parameters
----------
a : torch.Tensor
An input tensor
b : torch.Tensor
A second input tensor
Returns
-------
torch.Tensor
A tensor storing the elementwise complex quotient
"""
return cmult(a, cconj(b)) / t.unsqueeze(cabssq(b),-1)
#
# Not entirely sure if these belong here, but heck with it.
# We just need the ability to do fftshifts
#
def fftshift(array,dims=None):
"""Drop-in torch replacement for scipy.fftpack.fftshift
This maps a tensor, assumed to be the output of a fast Fourier
transform, into a tensor whose zero-frequency element is at the
center of the tensor instead of the start. It will by default shift
every dimension in the tensor but the last (which is assumed to
represent the complex number and be of dimension 2), but can shift
any arbitrary set of dimensions.
Parameters
----------
array : torch.Tensor
An array of data to be fftshifted
dims : iterable
A list of all dimensions to shift
Returns
-------
torch.Tensor
The fftshifted tensor
"""
if dims is None:
dims=list(range(array.dim()))[:-1]
for dim in dims:
length = array.size()[dim]
cut_to = (length + 1) // 2
cut_len = length - cut_to
array = t.cat((array.narrow(dim,cut_to,cut_len),
array.narrow(dim,0,cut_to)), dim)
return array
def ifftshift(array,dims=None):
"""Drop-in torch replacement for scipy.fftpack.iftshift
This maps a tensor, assumed to be the shifted output of a fast
Fourier transform, into a tensor whose zero-frequency element is
back at the start of the tensor instead of the center. It is the
inverse of the fftshift operator. It will by default shift
every dimension in the tensor but the last (which is assumed to
represent the complex number and be of dimension 2), but can shift
any arbitrary set of dimensions.
Parameters
----------
array : torch.Tensor
An array of data to be ifftshifted
dims : list(int)
A list of all dimensions to shift
Returns
-------
torch.Tensor
The ifftshifted tensor
"""
if dims is None:
dims=list(range(array.dim()))[:-1]
for dim in dims:
length = array.size()[dim]
cut_to = length // 2
cut_len = length - cut_to
array = t.cat((array.narrow(dim,cut_to,cut_len),
array.narrow(dim,0,cut_to)), dim)
return array
def expi(x):
"""Returns a complex-format tensor for exp(i* (x))
Expects the input to be in the form of a real-valued tensor
Parameters
----------
x : torch.Tensor
An array to be exponentiated
Returns
-------
torch.Tensor
A complex-format tensor
"""
return t.stack((t.cos(x),t.sin(x)),dim=-1)
def cexpi(z):
"""Returns a complex-format tensor for exp(i* (z))
Expects the input to be in the form of a complex-valued tensor
Parameters
----------
x : torch.Tensor
An array to be exponentiated
Returns
-------
torch.Tensor
A complex-format tensor
"""
real = t.cos(z[...,0]) * t.exp(-z[...,1])
imag = t.sin(z[...,0]) * t.exp(-z[...,1])
return t.stack((real, imag),dim=-1)
CDTools/tools/data/data.py
+6 -3
View File
@@ -5,7 +5,6 @@ specific desired information from a .cxi file. These functions should
handle all the needed conversions between standard formats (for example,
transposes of the basis arrays, shifting from object to probe motion, etc).
"""
from __future__ import division, print_function, absolute_import
import h5py
import numpy as np
@@ -277,7 +276,7 @@ def get_mask(cxi_file):
mask = np.array(i1['detector_1/mask']).astype(np.uint32)
mask_on = np.equal(mask,np.uint32(0))
mask_has_signal = np.equal(mask,np.uint32(0x00001000))
return np.logical_or(mask_on,mask_has_signal).astype(np.bool)
return np.logical_or(mask_on,mask_has_signal).astype(bool)
else:
return None
@@ -354,7 +353,11 @@ def get_data(cxi_file, cut_zeroes = True):
data[data < 0] = 0
if 'axes' in cxi_file[pull_from].attrs:
axes = str(cxi_file[pull_from].attrs['axes'].decode()).split(':')
try:
axes = str(cxi_file[pull_from].attrs['axes'].decode()).split(':')
except AttributeError as e: # Weird string vs bytes thing, ehhh
axes = str(cxi_file[pull_from].attrs['axes']).split(':')
axes = [axis.strip().lower() for axis in axes]
else:
axes = None
CDTools/tools/image_processing/image_processing.py
+60 -68
View File
@@ -7,14 +7,13 @@ kind of tools perform common image manipulations on torch tensors, in such
a way that it is safe to include them in automatic differentiation models.
"""
from __future__ import division, print_function, absolute_import
import numpy as np
import torch as t
from CDTools.tools import cmath
from CDTools.tools import propagators
__all__ = ['centroid', 'centroid_sq', 'sinc_subpixel_shift',
'find_subpixel_shift', 'find_pixel_shift', 'find_shift',
'convolve_1d']
'convolve_1d', 'fourier_upsample']
def centroid(im, dims=2):
@@ -75,7 +74,7 @@ def centroid_sq(im, dims=2, comp=False):
An (i,j) index or stack of indices
"""
if comp:
im_sq = cmath.cabssq(im)
im_sq = t.abs(im)**2
else:
im_sq = im**2
@@ -109,9 +108,9 @@ def sinc_subpixel_shift(im, shift):
I = I.to(dtype=im.dtype,device=im.device)
J = J.to(dtype=im.dtype,device=im.device)
fft_im = cmath.fftshift(t.fft(im, 2))
shifted_fft_im = cmath.cmult(fft_im, cmath.expi(-shift[0]*I - shift[1]*J))
return t.ifft(cmath.ifftshift(shifted_fft_im),2)
fft_im = t.fft.fftshift(t.fft.fft2(im),dim=(-2,-1))
shifted_fft_im = fft_im * t.exp(1j * (-shift[0]*I - shift[1]*J))
return t.fft.ifft2(t.fft.ifftshift(shifted_fft_im, dim=(-2,-1)))
@@ -139,6 +138,7 @@ def find_subpixel_shift(im1, im2, search_around=(0,0), resolution=10):
shift : torch.Tensor
The relative shift (i,j) needed to best map im1 onto im2
"""
#
# Here's my approach, perhaps it's a little unconventional. I will first
# calculate the phase correlation function as found in ____ (cite a paper
@@ -147,19 +147,11 @@ def find_subpixel_shift(im1, im2, search_around=(0,0), resolution=10):
# using an FFT with upsampling by a factor of resolution in reciprocal
# space
#
# If last dimension is not 2, then convert to a complex tensor now
if im1.shape[-1] != 2:
im1 = t.stack((im1,t.zeros_like(im1)),dim=-1)
if im2.shape[-1] != 2:
im2 = t.stack((im2,t.zeros_like(im2)),dim=-1)
cor_fft = cmath.cmult(t.fft(im1,2),cmath.cconj(t.fft(im2,2)))
cor_fft = t.fft.fft2(im1) * t.conj(t.fft.fft2(im2))
# Not sure if this is more or less stable than just the correlation
# maximum - requires some testing
cor = t.ifft(cor_fft / cmath.cabs(cor_fft)[:,:,None],2)
cor = t.fft.ifft2(cor_fft / t.abs(cor_fft))
# Now, I need to shift the array to pull out a contiguous window
@@ -167,27 +159,33 @@ def find_subpixel_shift(im1, im2, search_around=(0,0), resolution=10):
try:
search_around = search_around.cpu()
except:
search_around = t.tensor(search_around)
search_around = t.as_tensor(search_around)
window_size = 15
shift_zero = tuple(-search_around + t.tensor([window_size,window_size]))
cor_window = t.roll(cor, shift_zero, dims=(0,1))[:2*window_size,:2*window_size]
cor_window = t.roll(cor, shift_zero, dims=(-2,-1))[...,:2*window_size,:2*window_size]
# Now we upsample this window
cor_window_fft = cmath.fftshift(t.fft(cor_window,2))
upsampled = t.zeros(tuple(t.tensor(cor_window_fft.shape)[:-1] * resolution) + (2,),
cor_window_fft = t.fft.fftshift(t.fft.fft2(cor_window),dim=(-2,-1))
upsampled = t.zeros(tuple(t.tensor(cor_window_fft.shape) * resolution),
dtype=cor.dtype,device=cor.device)
upsampled[:2*window_size,:2*window_size] = cor_window_fft
upsampled[...,:2*window_size,:2*window_size] = cor_window_fft
upsampled = t.roll(upsampled,(-window_size,-window_size),dims=(0,1))
upsampled = t.roll(cmath.cabssq(t.ifft(upsampled, 2)),(-window_size*resolution,-window_size*resolution), dims=(0,1))
upsampled = t.roll(t.abs(t.fft.ifft2(upsampled))**2,
(-window_size*resolution,-window_size*resolution),
dims=(0,1))
# And we extract the shift from the window
sh = t.tensor(upsampled.shape).to(device=upsampled.device)
cormax = t.tensor([t.argmax(upsampled) // sh[1],
t.argmax(upsampled) % sh[1]]).to(device=upsampled.device)
subpixel_shift = ((cormax + sh // 2) % sh - sh//2).to(dtype=upsampled.dtype)
sh = t.as_tensor(upsampled.shape, device=upsampled.device)
cormax = t.as_tensor([t.div(t.argmax(upsampled), sh[1],
rounding_mode='floor'),
t.argmax(upsampled) % sh[1]],
device=upsampled.device)
sh_over_2 = t.div(sh,2,rounding_mode='floor')
subpixel_shift = ((cormax + sh_over_2) % sh - sh_over_2).to(dtype=upsampled.dtype)
return search_around.to(device=upsampled.device, dtype=upsampled.dtype) + \
subpixel_shift / resolution
@@ -213,24 +211,19 @@ def find_pixel_shift(im1, im2):
shift : torch.Tensor
The integer-valued shift (i,j) that best maps im1 onto im2
"""
# If last dimension is not 2, then convert to a complex tensor now
if im1.shape[-1] != 2:
im1 = t.stack((im1,t.zeros_like(im1)),dim=-1)
if im2.shape[-1] != 2:
im2 = t.stack((im2,t.zeros_like(im2)),dim=-1)
cor_fft = cmath.cmult(t.fft(im1,2),cmath.cconj(t.fft(im2,2)))
cor_fft = t.fft.fft2(im1) * t.conj(t.fft.fft2(im2))
# Not sure if this is more or less stable than just the correlation
# maximum - requires some testing
cor = cmath.cabs(t.ifft(cor_fft / cmath.cabs(cor_fft)[:,:,None],2))
cor = t.abs(t.fft.ifft2(cor_fft / t.abs(cor_fft)))
sh = t.tensor(cor.shape).to(device=im1.device)
cormax = t.tensor([t.argmax(cor) // sh[1],
sh = t.as_tensor(cor.shape,device=im1.device)
cormax = t.tensor([t.div(t.argmax(cor),sh[1],rounding_mode='floor'),
t.argmax(cor) % sh[1]]).to(device=im1.device)
return (cormax + sh // 2) % sh - sh//2
sh_over_2 = t.div(sh,2,rounding_mode='floor')
return (cormax + sh_over_2) % sh - sh_over_2
@@ -290,38 +283,37 @@ def convolve_1d(image, kernel, dim=0, fftshift_kernel=True):
The convolved image
"""
complex_things = 2
im_complex = True
if image.shape[-1] != 2:
image = t.stack((image,t.zeros_like(image)),dim=-1)
complex_things -= 1
im_complex = False
if kernel.shape[-1] != 2:
kernel = t.stack((kernel,t.zeros_like(kernel)),dim=-1)
complex_things -= 1
if fftshift_kernel:
kernel = cmath.ifftshift(kernel)
kernel = t.fft.ifftshift(kernel,dim=(-1,))
# If the image wasn't originally complex, and the dimension
# was passed with the nexative-indexing convention
if not im_complex and dim < 0:
dim = dim-1
# We have to transpose the relevant dimension to -2 before using the fft,
# which expects to operate on the final non-complex dimension
trans_im = t.transpose(image, dim, -2)
# We have to transpose the relevant dimension to -1 before using the fft,
# which expects to operate on the final dimension
trans_im = t.transpose(image, dim, -1)
# Take a correlation
fft_im = t.fft(trans_im, 1)
fft_kernel = t.fft(kernel, 1)
trans_conv = t.ifft(cmath.cmult(fft_im,fft_kernel), 1)
fft_im = t.fft.fft(trans_im)
fft_kernel = t.fft.fft(kernel)
trans_conv = t.fft.ifft(fft_im * fft_kernel)
conv_im = t.transpose(trans_conv, dim, -1)
return conv_im
def fourier_upsample(ims, preserve_mean=False):
# If preserve_mean is true, it preserves the mean pixel intensity
# otherwise, it preserves the total summed intensity
upsampled = t.zeros(ims.shape[:-2]+(2*ims.shape[-2],2*ims.shape[-1]),
dtype=ims.dtype,
device=ims.device)
left = [ims.shape[-2]//2,ims.shape[-1]//2]
right = [ims.shape[-2]//2+ims.shape[-2],
ims.shape[-1]//2+ims.shape[-1]]
upsampled[...,left[0]:right[0],left[1]:right[1]] = propagators.far_field(ims)
if preserve_mean:
upsampled *= 2
return propagators.inverse_far_field(upsampled)
conv_im = t.transpose(trans_conv, dim, -2)
# If nothing was input as complex, the result should be returned as real
if complex_things == 0:
return conv_im[...,0]
else:
return conv_im
CDTools/tools/initializers/initializers.py
+234 -53
View File
@@ -4,21 +4,22 @@ The functions in this module both do the geometric calculations needed to
initialize the reconstrucions, and the heuristic calculations for
geierating sensible initializations for the probe guess.
"""
from __future__ import division, print_function, absolute_import
import numpy as np
import torch as t
__all__ = ['exit_wave_geometry', 'calc_object_setup', 'gaussian',
'gaussian_probe', 'SHARP_style_probe', 'RPI_spectral_init']
from CDTools.tools import cmath
from CDTools.tools.propagators import *
from CDTools.tools.analysis import orthogonalize_probes
from CDTools.tools import image_processing
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 *
__all__ = ['exit_wave_geometry', 'calc_object_setup', 'gaussian',
'gaussian_probe', 'SHARP_style_probe', 'STEM_style_probe',
'RPI_spectral_init',
'generate_subdominant_modes']
def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None, opt_for_fft=True, padding=0, oversampling=1):
"""Returns an exit wave basis and shape, as well as a detector slice for the given detector geometry
@@ -59,14 +60,17 @@ def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None,
The slice corresponding to the physical detector
"""
det_shape = t.Tensor(tuple(det_shape)).to(t.int32)
det_basis = t.Tensor(det_basis)
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:
center = det_shape // 2
# 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.Tensor(center).to(t.int32)
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
@@ -79,11 +83,14 @@ def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None,
if opt_for_fft:
full_shape = t.Tensor([next_fast_len(dim) for dim in full_shape]).to(t.int32)
full_shape = t.as_tensor([next_fast_len(dim) for dim in full_shape],
dtype=t.int32)
# Then, generate a slice that pops the actual detector from the full
# detector shape
full_center = full_shape // 2
# 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]):
@@ -95,7 +102,7 @@ def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None,
# This method should work for a general parallelogram
# shaped detector
det_shape = det_basis * full_shape.to(t.float32)
pinv_basis = t.Tensor(np.linalg.pinv(det_shape).transpose()).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
@@ -103,7 +110,7 @@ def exit_wave_geometry(det_basis, det_shape, wavelength, distance, center=None,
#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)
#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)
@@ -199,7 +206,7 @@ def gaussian(shape, sigma, amplitude=1, center = None, curvature=[0,0]):
jsq = (j - center[1])**2
result = np.exp((1j*curvature[0] / 2 - 1 / (2 * sigma[0]**2)) * isq + \
(1j*curvature[1] / 2 - 1 / (2 * sigma[1]**2)) * jsq)
return cmath.complex_to_torch(amplitude*result)
return t.as_tensor(amplitude*result,dtype=t.complex64)
@@ -266,8 +273,8 @@ def gaussian_probe(dataset, basis, shape, sigma, propagation_distance=0):
# Finally, we should calculate the average pattern intensity from the
# dataset and normalize the gaussian probe. This should be done by
avg_intensities = [t.sum(dataset[idx][1]) for idx in range(len(dataset))]
avg_intensity = t.mean(t.Tensor(avg_intensities))
probe_intensity = t.sum(cmath.cabssq(probe))
avg_intensity = t.mean(t.tensor(avg_intensities))
probe_intensity = t.sum(t.abs(probe)**2)
return avg_intensity / probe_intensity * probe
@@ -308,6 +315,8 @@ def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None, over
The complex-style tensor storing the probe guess
"""
# NOTE: I don't love the way np and torch are mixed here, I think this
# function deserves some love.
# to use the mask or not?
intensities = np.zeros([dim // oversampling for dim in shape])
@@ -322,30 +331,29 @@ def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None, over
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)
probe_fft = cmath.complex_to_torch(np.sqrt(intensities))
probe_guess = cmath.torch_to_complex(inverse_far_field(probe_fft))
# Now we remove the central pixel
center = np.array(probe_guess.shape) // 2
# I'm always divided on 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_fft = t.tensor(np.sqrt(intensities)).to(dtype=t.complex64)
probe_guess = cmath.complex_to_torch(probe_guess)
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)
if propagation_distance is not None:
# First generate the propagation array
probe_shape = t.Tensor(tuple(probe_guess.shape))[:-1]
probe_shape = t.as_tensor(tuple(probe_guess.shape))
# Start by recalculating the probe basis from the given information
det_basis = t.Tensor(dataset.detector_geometry['basis'])
det_basis = t.as_tensor(dataset.detector_geometry['basis'])
basis_dirs = det_basis / t.norm(det_basis, dim=0)
distance = dataset.detector_geometry['distance']
probe_basis = basis_dirs * dataset.wavelength * distance / \
@@ -355,19 +363,117 @@ def SHARP_style_probe(dataset, shape, det_slice, propagation_distance=None, over
probe_spacing = t.norm(probe_basis,dim=0).numpy()
probe_shape = probe_shape.numpy().astype(np.int32)
#assert 0
# And generate the propagator
AS_prop = generate_angular_spectrum_propagator(probe_shape, probe_spacing, dataset.wavelength, propagation_distance)
probe_guess = near_field(probe_guess,AS_prop)
# Finally, place this probe in a full-sized array if there is oversampling
final_probe = t.zeros([dim for dim in shape] + [2])
final_probe = t.zeros(shape,dtype=t.complex64)
left = shape[0]//2 - probe_guess.shape[0] // 2
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
top:top+probe_guess.shape[1]] = probe_guess
return final_probe
def STEM_style_probe(dataset, shape, det_slice, convergence_semiangle, propagation_distance=None, oversampling=1):
"""Generates a STEM style probe guess from a dataset
What we call the "STEM" style probe guess is a probe generated by
a uniform aperture in Fourier space, with no optical aberrations.
This is the kind of probe than an ideal, aberration-corrected STEM
with perfect coherence would produce, so it is a good starting
guess for STEM datasets.
We set the initial intensity of the probe with relation to the
diffraction patterns so that a typical object will have an intensity
of around 1. We also set the initial phase ramp of the probe such that
the undiffracted probe has a centroid on the detector which matches the
centroid of the diffraction patterns
Parameters
----------
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
convergence_angle : float
The convergence angle of the probe, in mrad.
propagation_distance : float
Default is no propagation, an amount to propagate the guessed probe from it's focal point
oversampling : int
Default 1, the width of the region of pixels in the wavefield to bin into a single detector pixel
Returns
-------
torch.Tensor
The complex-style tensor storing the probe guess
"""
# The basis in the dataset should describe the basis of the probe's
# Fourier transform, even if the probe is simulated on larger stage than
# the detector. The only issue is if the probe is oversampled in
# Fourier space (simulated on a larger stage in real space). That factor
# is defined by oversampling.
probe_basis = (t.as_tensor(dataset.detector_geometry['basis'],
dtype=t.float32)
/ oversampling)
mean_im = t.mean(dataset.patterns,dim=0)
center = image_processing.centroid(mean_im)
Is = t.arange(shape[0], dtype=t.float32) - center[0]
Js = t.arange(shape[1], dtype=t.float32) - center[1]
Is,Js = t.meshgrid(Is,Js)
Rs = t.tensordot(probe_basis,t.stack([Is,Js]),dims=1)
forward = t.Tensor([0,0,1])
Rs = (forward * dataset.detector_geometry['distance'])[:,None,None] + Rs
dirs = Rs / t.linalg.norm(Rs,dim=(0,))
angles = t.acos(t.tensordot(forward,dirs,dims=1)) * 1000 #in mrad
probe_fft = t.zeros(shape,dtype=t.complex64)
probe_fft[angles<convergence_semiangle] = 1
probe_mean = t.mean(t.abs(probe_fft)**2)
diff_mean = t.mean(mean_im)
probe_fft = probe_fft * t.sqrt(diff_mean / probe_mean)
probe_guess = inverse_far_field(probe_fft)
if propagation_distance is not None:
# It's probably worth checking if this is correct when oversampling
# is not 1
probe_shape = t.as_tensor(tuple(probe_guess.shape))
# Start by recalculating the probe basis from the given information
det_basis = t.as_tensor(dataset.detector_geometry['basis'])
basis_dirs = det_basis / t.norm(det_basis, dim=0)
distance = dataset.detector_geometry['distance']
probe_basis = basis_dirs * dataset.wavelength * distance / \
(probe_shape * t.norm(det_basis,dim=0))
# Then package everything as it's needed
probe_spacing = t.norm(probe_basis,dim=0).numpy()
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)
probe_guess = near_field(probe_guess,AS_prop)
# Finally, place this probe in a full-sized array if there is oversampling
final_probe = t.zeros(shape,dtype=t.complex64)
left = shape[0]//2 - probe_guess.shape[0] // 2
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
return final_probe
@@ -380,26 +486,25 @@ def RPI_spectral_init(pattern, probe, obj_shape, n_modes=1, mask=None, backgroun
if probe.dim() == 4:
probe = orthogonalize_probes(probe)[0]
pad0l = (probe.shape[-3] - obj_shape[0])//2
pad0r = probe.shape[-3] - obj_shape[0] - pad0l
pad1l = (probe.shape[-2] - obj_shape[1])//2
pad1r = probe.shape[-2] - obj_shape[1] - pad1l
pad0l = (probe.shape[-2] - obj_shape[0])//2
pad0r = probe.shape[-2] - obj_shape[0] - pad0l
pad1l = (probe.shape[-1] - obj_shape[1])//2
pad1r = probe.shape[-1] - obj_shape[1] - pad1l
def a_dagger(im):
im = cmath.complex_to_torch(im.reshape(obj_shape)).to(dtype=t.float32)
im = inverse_far_field(pad(far_field(im), (0,0,pad1l,pad1r,pad0l,pad0r)))
exit_wave = cmath.cmult(probe,im)
farfield = cmath.torch_to_complex(far_field(exit_wave))
return farfield.ravel()
im = t.tensor(im.reshape(obj_shape)).to(dtype=t.complex64)
im = inverse_far_field(pad(far_field(im), (pad1l,pad1r,pad0l,pad0r)))
exit_wave = probe * im
return far_field(exit_wave).numpy().ravel()
def a(measured):
measured = cmath.complex_to_torch(measured.reshape(pattern.shape[0],pattern.shape[1])).to(dtype=t.float32)
measured = t.tensor(measured.reshape(pattern.shape[0],pattern.shape[1])).to(dtype=t.complex64)
im = inverse_far_field(measured)
multiplied = cmath.cmult(cmath.cconj(probe), im)
multiplied = t.conj(probe) * im
backplane = far_field(multiplied)
clipped = backplane[pad0l:pad0l+obj_shape[0],
pad1l:pad1l+obj_shape[1],:]
return cmath.torch_to_complex(inverse_far_field(clipped)).ravel()
pad1l:pad1l+obj_shape[1]]
return (inverse_far_field(clipped)).numpy().ravel()
patsize = pattern.shape[0]*pattern.shape[1]
imsize = obj_shape[0]*obj_shape[1]
@@ -430,11 +535,87 @@ def RPI_spectral_init(pattern, probe, obj_shape, n_modes=1, mask=None, backgroun
# Now we set the overall scale and relative weights of the guess
scale_factor = np.sqrt(np.sum(np_pattern) /
t.sum(cmath.cabssq(probe)).numpy())
t.sum(t.abs(probe)**2).numpy())
relative_weights = eigval / np.sum(eigval**2)
z0 = z0 * (scale_factor * relative_weights[:,None,None])
# Now we have to normalize the modes by their eigenvalues
return cmath.complex_to_torch(z0).to(dtype=t.float32)
return t.as_tensor(z0, dtype=t.complex64)
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[-2]-1)//2, (shape[-1]-1)//2)
i, j = np.mgrid[:shape[-2], :shape[-1]]
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 = i + 1j * j
b = i - 1j * j
else:
a = i
b = j
probe_norm = t.sum(t.abs(dominant_mode)**2)
# 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(t.mul,[a]*a_order+[b]*b_order)
new_probe = inverse_far_field(mask * dominant_fft)
probes.append(new_probe * (probe_norm / t.sum(t.abs(new_probe))))
return t.stack(probes)
CDTools/tools/interactions/interactions.py
+108 -78
View File
@@ -5,12 +5,9 @@ function simulates some aspect of an interaction model that can be used
for ptychographic reconstruction.
"""
from __future__ import division, print_function, absolute_import
from CDTools.tools.cmath import *
import torch as t
import numpy as np
from CDTools.tools import propagators
from CDTools.tools import propagators, image_processing
__all__ = ['translations_to_pixel', 'pixel_to_translations',
'project_translations_to_sample',
@@ -18,7 +15,6 @@ __all__ = ['translations_to_pixel', 'pixel_to_translations',
'RPI_interaction']
def translations_to_pixel(basis, translations, surface_normal=t.Tensor([0.,0.,1.])):
"""Takes real space translations and outputs them in pixel space
@@ -226,7 +222,7 @@ def project_translations_to_sample(sample_basis, translations):
def ptycho_2D_round(probe, obj, translations):
def ptycho_2D_round(probe, obj, translations, multiple_modes=False, upsample_obj=False):
"""Returns a stack of exit waves without accounting for subpixel shifts
This function returns a collection of exit waves, with the first
@@ -234,34 +230,63 @@ def ptycho_2D_round(probe, obj, translations):
corresponding to the detector. The exit waves are calculated by
shifting the probe by the rounded value of the translation
If multiple_modes is set to False, any additional dimensions in the
ptycho_2D_round function will be assumed to correspond to the translation
index. If multiple_modes is set to true, the (-4th) dimension of the probe
will always be assumed to be defining a set of (P) incoherently mixing
modes to be broadcast all translation indices. If any additional dimensions
closer to the start exist, they will be assumed to be translation indices
Parameters
----------
probe : torch.Tensor
An MxL probe function for the exit waves
A (P)xMxL probe function to illuminate the object
object : torch.Tensor
The object function to be probed
translations : torch.Tensor
The Nx2 array of (i,j) translations to simulate
The (N)x2 array of (i,j) translations to simulate
multuple_modes : bool
Default False, whether to assume the probe contains multiple modes
Returns
-------
exit_waves : torch.Tensor
An NxMxL tensor of the calculated exit waves
An (N)x(P)xMxL tensor of the calculated exit waves
"""
single_translation = False
if translations.dim() == 1:
translations = translations[None,:]
single_translation = True
integer_translations = t.round(translations).to(dtype=t.int32)
selections = [obj[tr[0]:tr[0]+probe.shape[0],
tr[1]:tr[1]+probe.shape[1]]
for tr in integer_translations]
if single_translation:
return [cmult(probe,selection) for selection in selections][0]
else:
return t.stack([cmult(probe,selection) for selection in selections])
if upsample_obj:
selections = t.stack([obj[tr[0]:tr[0]+probe.shape[-2]//2,
tr[1]:tr[1]+probe.shape[-1]//2]
for tr in integer_translations])
selections = image_processing.fourier_upsample(selections,
preserve_mean=True)
else:
selections = t.stack([obj[tr[0]:tr[0]+probe.shape[-2],
tr[1]:tr[1]+probe.shape[-1]]
for tr in integer_translations])
if multiple_modes:
# if the probe dimension is 4, then this hasn't yet been broadcast
# over the translation dimensions
output = probe * selections[:,None,:,:]
else:
output = probe * selections
if single_translation:
return output[0]
else:
return output
@@ -326,7 +351,7 @@ def ptycho_2D_linear(probe, obj, translations, shift_probe=True):
obj_slice = obj[tr[0]:tr[0]+probe.shape[0],
tr[1]:tr[1]+probe.shape[1]]
exit_waves.append(cmult(selection,obj_slice))
exit_waves.append(selection * obj_slice)
else:
for tr, sp in zip(integer_translations,
subpixel_translations):
@@ -352,7 +377,7 @@ def ptycho_2D_linear(probe, obj, translations, shift_probe=True):
sel10 * sp[0]*(1-sp[1]) + \
sel11 * sp[0]*sp[1]
exit_waves.append(cmult(probe,selection))
exit_waves.append(probe * selection)
if single_translation:
return exit_waves[0]
@@ -360,9 +385,7 @@ def ptycho_2D_linear(probe, obj, translations, shift_probe=True):
return t.stack(exit_waves)
def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10):
def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10, multiple_modes=True):
"""Returns a stack of exit waves accounting for subpixel shifts
This function returns a collection of exit waves, with the first
@@ -373,25 +396,32 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10):
in Fourier space)
If shift_probe is True, it applies the subpixel shift to the probe,
otherwise the subpixel shift is applied to the object
otherwise the subpixel shift is applied to the object [not yet implemented]
If multiple_modes is set to False, any additional dimensions in the
ptycho_2D_round function will be assumed to correspond to the translation
index. If multiple_modes is set to true, the (-4th) dimension of the probe
will always be assumed to be defining a set of (P) incoherently mixing
modes to be broadcast all translation indices. If any additional dimensions
closer to the start exist, they will be assumed to be translation indices
Parameters
----------
probe : torch.Tensor
An MxL probe function for the exit waves
An (P)xMxL probe function for the exit waves
object : torch.Tensor
The object function to be probed
translations : torch.Tensor
The Nx2 array of translations to simulate
The (N)x2 array of translations to simulate
shift_probe : bool
Default True, Whether to subpixel shift the probe or object
padding : int
Default 10, if shifting the object, the padding to apply to the object to avoid circular shift effects
multuple_modes : bool
Default False, whether to assume the probe contains multiple modes
Returns
-------
exit_waves : torch.Tensor
An NxMxL tensor of the calculated exit waves
An (N)x(P)xMxL tensor of the calculated exit waves
"""
single_translation = False
if translations.dim() == 1:
@@ -403,35 +433,45 @@ def ptycho_2D_sinc(probe, obj, translations, shift_probe=True, padding=10):
integer_translations = t.floor(translations)
subpixel_translations = translations - integer_translations
integer_translations = integer_translations.to(dtype=t.int32)
selections = t.stack([obj[tr[0]:tr[0]+probe.shape[-2],
tr[1]:tr[1]+probe.shape[-1]]
for tr in integer_translations])
exit_waves = []
if shift_probe:
i = t.arange(probe.shape[0]) - probe.shape[0]//2
j = t.arange(probe.shape[1]) - probe.shape[1]//2
i = t.arange(probe.shape[-2],device=probe.device,dtype=t.float32) \
- probe.shape[-2]//2
j = t.arange(probe.shape[-1],device=probe.device,dtype=t.float32) \
- probe.shape[-1]//2
I,J = t.meshgrid(i,j)
I = 2 * np.pi * I.to(t.float32) / probe.shape[0]
J = 2 * np.pi * J.to(t.float32) / probe.shape[1]
I = I.to(dtype=probe.dtype,device=probe.device)
J = J.to(dtype=probe.dtype,device=probe.device)
I = 2 * np.pi * I / probe.shape[-2]
J = 2 * np.pi * J / probe.shape[-1]
phase_masks = t.exp(1j*(-subpixel_translations[:,0,None,None]*I
-subpixel_translations[:,1,None,None]*J))
for tr, sp in zip(integer_translations,
subpixel_translations):
fft_probe = fftshift(t.fft(probe, 2))
shifted_fft_probe = cmult(fft_probe, expi(-sp[0]*I - sp[1]*J))
shifted_probe = t.ifft(ifftshift(shifted_fft_probe),2)
fft_probe = t.fft.fftshift(t.fft.fft2(probe),dim=(-1,-2))
obj_slice = obj[tr[0]:tr[0]+probe.shape[0],
tr[1]:tr[1]+probe.shape[1]]
if multiple_modes: # Multi-mode probe
shifted_fft_probe = fft_probe * phase_masks[...,None,:,:]
else:
shifted_fft_probe = fft_probe * phase_masks
exit_waves.append(cmult(shifted_probe, obj_slice))
shifted_probe = t.fft.ifft2(t.fft.ifftshift(shifted_fft_probe,
dim=(-1,-2)))
if multiple_modes: # Multi-mode probe
output = shifted_probe * selections[...,None,:,:]
else:
output = shifted_probe * selections
else:
raise NotImplementedError('Object shift not yet implemented')
if single_translation:
return exit_waves[0]
return output[0]
else:
return t.stack(exit_waves)
return output
def ptycho_2D_sinc_s_matrix(probe, s_matrix, translations, shift_probe=True, padding=10):
@@ -488,46 +528,36 @@ def ptycho_2D_sinc_s_matrix(probe, s_matrix, translations, shift_probe=True, pad
B = s_matrix.shape[0]//2
if shift_probe:
i = t.arange(probe.shape[0]) - probe.shape[0]//2
j = t.arange(probe.shape[1]) - probe.shape[1]//2
i = t.arange(probe.shape[-2]) - probe.shape[-2]//2
j = t.arange(probe.shape[-1]) - probe.shape[-1]//2
I,J = t.meshgrid(i,j)
I = 2 * np.pi * I.to(t.float32) / probe.shape[0]
J = 2 * np.pi * J.to(t.float32) / probe.shape[1]
I = 2 * np.pi * I.to(t.float32) / probe.shape[-2]
J = 2 * np.pi * J.to(t.float32) / probe.shape[-1]
I = I.to(dtype=probe.dtype,device=probe.device)
J = J.to(dtype=probe.dtype,device=probe.device)
for tr, sp in zip(integer_translations,
subpixel_translations):
fft_probe = fftshift(t.fft(probe, 2))
shifted_fft_probe = cmult(fft_probe, expi(-sp[0]*I - sp[1]*J))
shifted_probe = t.ifft(ifftshift(shifted_fft_probe),2)
fft_probe = t.fft.fftshift(t.fft.fft2(probe), dim=(-1,-2))
shifted_fft_probe = fft_probe * t.exp(1j*(-sp[0]*I - sp[1]*J))
shifted_probe = t.fft.ifft2(t.fft.ifftshift(shifted_fft_probe,
dim=(-1,-2)))
s_matrix_slice = s_matrix[:,:,tr[0]:tr[0]+probe.shape[0]+2*B,
tr[1]:tr[1]+probe.shape[1]+2*B]
s_matrix_slice = s_matrix[:,:,tr[0]:tr[0]+probe.shape[-2]+2*B,
tr[1]:tr[1]+probe.shape[-1]+2*B]
output = t.zeros([probe.shape[0]+2*B,probe.shape[1]+2*B,2]).to(
output = t.zeros([probe.shape[-2]+2*B,probe.shape[-1]+2*B,2]).to(
device=s_matrix_slice.device,
dtype=s_matrix_slice.dtype)
for i in range(s_matrix.shape[0]):
for j in range(s_matrix.shape[1]):
output [i:i+probe.shape[0],j:j+probe.shape[1]] += \
cmult(shifted_probe, s_matrix_slice[i,j,i:i+probe.shape[0],j:j+probe.shape[1],:])
#output = t.zeros([s_matrix_slice.shape[2]+2*B,
# s_matrix_slice.shape[3]+2*B,2]).to(
# device=s_matrix_slice.device,
# dtype=s_matrix_slice.dtype)
#for i in range(s_matrix.shape[0]):
# for j in range(s_matrix.shape[1]):
# output[i:i+probe.shape[0],j:j+probe.shape[1]] += \
# cmult(shifted_probe, s_matrix_slice[i,j,:,:,:])
output [i:i+probe.shape[-2],j:j+probe.shape[-1]] += \
shifted_probe * s_matrix_slice[i,j,i:i+probe.shape[-2],j:j+probe.shape[-1]]
exit_waves.append(output)
#exit_waves.append(cmult(shifted_probe, obj_slice))
else:
raise NotImplementedError('Object shift not yet implemented')
@@ -571,23 +601,23 @@ def RPI_interaction(probe, obj):
# The far-field propagator is just a 2D FFT but with an fftshift
fftobj = propagators.far_field(obj)
# We calculate the padding that we need to do the upsampling
pad0l = (probe.shape[-3] - obj.shape[-3])//2
pad0r = probe.shape[-3] - obj.shape[-3] - pad0l
pad1l = (probe.shape[-2] - obj.shape[-2])//2
pad1r = probe.shape[-2] - obj.shape[-2] - pad1l
pad0l = (probe.shape[-2] - obj.shape[-2])//2
pad0r = probe.shape[-2] - obj.shape[-2] - pad0l
pad1l = (probe.shape[-1] - obj.shape[-1])//2
pad1r = probe.shape[-1] - obj.shape[-1] - pad1l
if obj.dim() == 3:
fftobj = t.nn.functional.pad(fftobj, (0, 0, pad1l, pad1r, pad0l, pad0r))
elif obj.dim() == 4:
if obj.dim() == 2:
fftobj = t.nn.functional.pad(fftobj, (pad1l, pad1r, pad0l, pad0r))
elif obj.dim() == 3:
fftobj = t.nn.functional.pad(
fftobj, (0, 0, pad1l, pad1r, pad0l, pad0r, 0, 0))
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.')
# Again, just an inverse FFT but with an fftshift
upsampled_obj = propagators.inverse_far_field(fftobj)
if obj.dim() == 4:
return cmult(probe[None,...], upsampled_obj)
if obj.dim() == 3:
return probe[None,...] * upsampled_obj
else:
return cmult(probe, upsampled_obj)
return probe * upsampled_obj
CDTools/tools/losses/losses.py
+34 -9
View File
@@ -5,9 +5,8 @@ that returns the mean squared intensity error, and one that returns the
maximum likelihood metric for a system with Poisson statistics.
"""
from __future__ import division, print_function, absolute_import
import torch as t
import torch as t
__all__ = ['amplitude_mse', 'intensity_mse', 'poisson_nll']
@@ -100,7 +99,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 +136,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 * 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
CDTools/tools/measurements/measurements.py
+15 -114
View File
@@ -5,9 +5,6 @@ model. There exist tools to simulate detectors with finite saturation
thresholds, backgrounds, and more.
"""
from __future__ import division, print_function, absolute_import
from CDTools.tools import cmath
import torch as t
import numpy as np
from torch.nn.functional import avg_pool2d
@@ -17,8 +14,7 @@ from torch.nn.functional import avg_pool2d
# intensity pattern on a detector
#
__all__ = ['intensity', 'incoherent_sum', 'density_matrix',
'quadratic_background']
__all__ = ['intensity', 'incoherent_sum', 'quadratic_background']
def intensity(wavefield, detector_slice=None, epsilon=1e-7, saturation=None, oversampling=1):
@@ -31,7 +27,7 @@ def intensity(wavefield, detector_slice=None, epsilon=1e-7, saturation=None, ove
Parameters
----------
wavefield : torch.Tensor
A JxMxNx2 stack of complex wavefields
A JxMxN stack of complex-valued wavefields
detector_slice : slice
Optional, a slice or tuple of slices defining a section of the simulation to return
saturation : float
@@ -44,18 +40,18 @@ def intensity(wavefield, detector_slice=None, epsilon=1e-7, saturation=None, ove
sim_patterns : torch.Tensor
A real MxN array storing the wavefield's intensities
"""
output = cmath.cabssq(wavefield)
output = t.abs(wavefield)**2
# Now we apply oversampling
if oversampling != 1:
if wavefield.dim() == 3:
if wavefield.dim() == 2:
output = avg_pool2d(output.unsqueeze(0), oversampling)[0]
else:
output = avg_pool2d(output, oversampling)
# Then we grab the detector slice
if detector_slice is not None:
if wavefield.dim() == 3:
if wavefield.dim() == 2:
output = output[detector_slice]
else:
output = output[(np.s_[:],) + detector_slice]
@@ -65,101 +61,6 @@ def intensity(wavefield, detector_slice=None, epsilon=1e-7, saturation=None, ove
return output + epsilon
else:
return t.clamp(output + epsilon,0,saturation)
def density_matrix(wavefields, density_matrix, detector_slice=None, epsilon=1e-7, saturation=None, oversampling=1):
"""Returns the intensities associated with a given density matrix state
The essential idea is that the most general description of a light field
at the detector plane will consist of a density matrix state. Here, that
low rank density matrix state is encoded as a set of basis wavefields
and a density matrix in that basis.
For computational efficiency, the density matrix is coded in an unusual
format. The density matrix formally is a complex Hermetian matrix,
which also happens to be positive definite. Here, we store it as a
real-valued matrix, where the upper triangle corresponds to the real
part of the elements in the upper triangle, and the lower triangle
corresponds to the imaginary parts. The elements on the diagonal are
purely real, and are stored as they are.
As with other multi-mode measurement functions, the modes are stored in
the first index, and the index of the diffraction pattern in the stack
of diffraction patterns is the second index. The stack-direction index
can be omitted if only a single pattern needs to be simulated
It is important to note that this method does not inforce the positive
definiteness of the density matrix, this it is possible for negative
values of intensity to appear if the underlying density matrix passed
to this method is not positive definite
Parameters
----------
wavefields : torch.Tensor
An Lx(Jx)MxNx2 stack of complex wavefields
density_matrix : torch.Tensor
A (Jx)LxL stack of real-valued representations of density matrices, as per above
saturation : float
Optional, a maximum saturation value to clamp the resulting intensities to
oversampling : int
Default 1, the width of the region pixels in the wavefield to bin into a single detector pixel
Returns
-------
sim_patterns : torch.Tensor
A real Lx(Jx)MxN array storing the incoherently summed intensities
"""
#if wavefields.dim() == 4:
# wavefields.unsqueeze(1)
# single_frame = True
#elif wavefields.dim() == 5:
# single_frame=False
output = t.zeros(wavefields.shape[1:-1],
dtype=wavefields.dtype,
device=wavefields.device)
# flat is better than nested, but simple is better than complex...
for (i,j) in ((i,j) for i in range(density_matrix.shape[-2])
for j in range(density_matrix.shape[-1])):
if i == j: # diagonal
output += density_matrix[...,i,j,None,None] \
* cmath.cabssq(wavefields[i])
if i < j: # upper triangle, real part
output += 2 * density_matrix[...,i,j,None,None] \
* (wavefields[i,...,0] * wavefields[j,...,0]
+ wavefields[i,...,1] * wavefields[j,...,1])
if i > j: # lower triangle, imaginary part
# We pull the i,jth element from the density matrix,
# but this correponds to wavefield j and wavefield i,
# unlike above where it was wavefield i and j (swapped order).
# We also get the one negative sign because this is the imaginary
# part
output += 2 * density_matrix[...,i,j,None,None] \
* (wavefields[j,...,0] * wavefields[i,...,1]
- wavefields[j,...,1] * wavefields[i,...,0])
# Now we apply oversampling
if oversampling != 1:
if wavefields.dim() == 4:
output = avg_pool2d(output.unsqueeze(0), oversampling)[0]
else:
output = avg_pool2d(output, oversampling)
# Then we grab the detector slice
if detector_slice is not None:
if wavefields.dim() == 4:
output = output[detector_slice]
else:
output = output[(np.s_[:],) + detector_slice]
if saturation is None:
return t.clamp(output,min=0) + epsilon
else:
return t.clamp(output + epsilon,0,saturation)
def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=None, oversampling=1):
@@ -169,15 +70,15 @@ def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=Non
the wavefields. If a detector slice is given, the returned array
will only include that slice from the simulated wavefronts.
The first index is the set of incoherently adding patterns, and
the second index is the index of the diffraction pattern to measure.
The next two indices index the wavefield. The final index is the complex
index.
The (-3) index is the set of incoherently adding patterns, and any
indexes further to the front correspond to the set of diffraction patterns
to measure. The (-2) and (-1) indices are the wavefield, and the final
index is the complex index
Parameters
----------
wavefields : torch.Tensor
An LxJxMxNx2 stack of complex wavefields
An LxJxMxNx stack of complex wavefields
detector_slice : slice
Optional, a slice or tuple of slices defining a section of the simulation to return
saturation : float
@@ -188,21 +89,21 @@ def incoherent_sum(wavefields, detector_slice=None, epsilon=1e-7, saturation=Non
Returns
-------
sim_patterns : torch.Tensor
A real JXMxN array storing the incoherently summed intensities
A real LXMxN array storing the incoherently summed intensities
"""
output = t.sum(cmath.cabssq(wavefields),dim=0)
output = t.sum(t.abs(wavefields)**2,dim=-3)
# Now we apply oversampling
if oversampling != 1:
if wavefields.dim() == 4:
if wavefields.dim() == 3:
output = avg_pool2d(output.unsqueeze(0), oversampling)[0]
else:
output = avg_pool2d(output, oversampling)
# Then we grab the detector slice
if detector_slice is not None:
if wavefields.dim() == 4:
if wavefields.dim() == 3:
output = output[detector_slice]
else:
output = output[(np.s_[:],) + detector_slice]
@@ -224,7 +125,7 @@ def quadratic_background(wavefield, background, *args, detector_slice=None, meas
Parameters
----------
wavefield : torch.Tensor
A JxMxNx2 stack of complex wavefields
A JxMxN stack of complex-valued wavefields
background : torch.Tensor
An tensor storing the square root of the detector background
detector_slice : slice
CDTools/tools/plotting/plotting.py
+434 -191
View File
@@ -6,18 +6,18 @@ and for after-the-fact analysis. Utilities for plotting complex valued
images exist, as well as plotting scan patterns and nanomaps
"""
from __future__ import division, print_function, absolute_import
from CDTools.tools import cmath
import torch as t
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import hsv_to_rgb
from matplotlib.widgets import Slider
from matplotlib import ticker, patheffects
__all__ = ['colorize', 'plot_amplitude', 'plot_phase',
'plot_colorized', 'plot_translations', 'get_units_factor',
'plot_nanomap', 'plot_real', 'plot_imag']
'plot_nanomap', 'plot_real', 'plot_imag',
'plot_nanomap_with_images']
def colorize(z):
@@ -82,7 +82,146 @@ def get_units_factor(units):
return factor
def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwargs):
def plot_image(im, plot_func=lambda x: x, fig=None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label=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
plot over it. Otherwise, it will generate a new figure.
If a basis is explicitly passed, the image will be plotted in real-space
coordinates
Finally, if a function is passed to the plot_func argument, this function
will be called on each slice of data before it is plotted. This is used
internally to enable the plot_real, plot_image, plot_phase, etc. functions.
Parameters
----------
im : array
An complex array with dimensions NxM
plot_func : callable
A function which maps numpy arrays to the image to be plotted
fig : matplotlib.figure.Figure
Default is a new figure, a matplotlib figure to use to plot
basis : np.array
Optional, the 3x2 probe basis
units : str
The length units to mark on the plot, default is um
cmap : str
Default is 'viridis', the colormap to plot with
cmap_label : str
What to label the colorbar when plotting
\\**kwargs
All other args are passed to fig.add_subplot(111, \\**kwargs)
Returns
-------
used_fig : matplotlib.figure.Figure
The figure object that was actually plotted to.
"""
# convert to numpy
if isinstance(im, t.Tensor):
# If final dimension is 2, assume it is a complex array. If not,
# assume it represents a real array
if im.shape[-1] == 2:
im = im.detach().cpu().numpy()
else:
im = im.detach().cpu().numpy()
if fig is None:
fig = plt.figure()
ax = fig.add_subplot(111, **kwargs)
# This nukes everything and updates either the appropriate image from the
# stack of images, or the only image if only a single image has been
# given
def make_plot(idx):
plt.figure(fig.number)
title = plt.gca().get_title()
fig.clear()
# If im only has two dimensions, this reshape will add a leading
# dimension, and update will be called on index 0. If it has 3 or more
# dimensions, then all the leading dimensions will be compressed into
# one long dimension which can be scrolled through.
s = im.shape
reshaped_im = im.reshape(-1,s[-2],s[-1])
num_images = reshaped_im.shape[0]
fig.plot_idx = idx % num_images
to_plot = plot_func(reshaped_im[fig.plot_idx])
#Plot in a basis if it exists, otherwise dont
if basis is not None:
if isinstance(basis,t.Tensor):
np_basis = basis.detach().cpu().numpy()
else:
np_basis = basis
# This fails if the basis is not rectangular
basis_norm = np.linalg.norm(np_basis, axis = 0)
basis_norm = basis_norm * get_units_factor(units)
extent = [0, im.shape[-1]*basis_norm[1], 0,
im.shape[-2]*basis_norm[0]]
else:
extent=None
plt.imshow(to_plot, cmap = cmap, extent = extent)
cbar = plt.colorbar()
if cmap_label is not None:
cbar.set_label(cmap_label)
if basis is not None:
plt.xlabel('X (' + units + ')')
plt.ylabel('Y (' + units + ')')
else:
plt.xlabel('j (pixels)')
plt.ylabel('i (pixels)')
plt.title(title)
if len(im.shape) >= 3:
plt.text(0.03, 0.03, str(fig.plot_idx), fontsize=14, transform=plt.gcf().transFigure)
return fig
if hasattr(fig, 'plot_idx'):
result = make_plot(fig.plot_idx)
else:
result = make_plot(0)
update = make_plot
def on_action(event):
if not hasattr(event, 'button'):
event.button = None
if not hasattr(event, 'key'):
event.key = None
if event.key == 'up' or event.button == 'up':
update(fig.plot_idx - 1)
elif event.key == 'down' or event.button == 'down':
update(fig.plot_idx + 1)
plt.draw()
if len(im.shape) >=3:
if not hasattr(fig,'my_callbacks'):
fig.my_callbacks = []
for cid in fig.my_callbacks:
fig.canvas.mpl_disconnect(cid)
fig.my_callbacks = []
fig.my_callbacks.append(fig.canvas.mpl_connect('key_press_event',on_action))
fig.my_callbacks.append(fig.canvas.mpl_connect('scroll_event',on_action))
return result
def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Real Part (a.u.)', **kwargs):
"""Plots the real part of a complex array with dimensions NxM
If a figure is given explicitly, it will clear that existing figure and
@@ -103,6 +242,8 @@ def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwa
The length units to mark on the plot, default is um
cmap : str
Default is 'viridis', the colormap to plot with
cmap_label : str
What to label the colorbar when plotting
\\**kwargs
All other args are passed to fig.add_subplot(111, \\**kwargs)
@@ -111,45 +252,14 @@ def plot_real(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwa
used_fig : matplotlib.figure.Figure
The figure object that was actually plotted to.
"""
if fig is None:
fig = plt.figure()
ax = fig.add_subplot(111, **kwargs)
else:
plt.figure(fig.number)
plt.gcf().clear()
if isinstance(im, t.Tensor):
real = im[...,0].detach().cpu().numpy()
else:
real = np.real(im)
#Plot in a basis if it exists, otherwise dont
if basis is not None:
if isinstance(basis,t.Tensor):
basis = basis.detach().cpu().numpy()
# This fails if the
basis_norm = np.linalg.norm(basis, axis = 0)
basis_norm = basis_norm * get_units_factor(units)
extent = [0, real.shape[-1]*basis_norm[1], 0, real.shape[-2]*basis_norm[0]]
else:
extent=None
plt.imshow(real, cmap = cmap, extent = extent)
cbar = plt.colorbar()
cbar.set_label('Real Part (a.u.)')
if basis is not None:
plt.xlabel('X (' + units + ')')
plt.ylabel('Y (' + units + ')')
else:
plt.xlabel('j (pixels)')
plt.ylabel('i (pixels)')
return fig
plot_func = lambda x: np.real(x)
return plot_image(im, plot_func=plot_func, fig=fig, basis=basis,
units=units, cmap=cmap, cmap_label=cmap_label,
**kwargs)
def plot_imag(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwargs):
def plot_imag(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Imaginary Part (a.u.)', **kwargs):
"""Plots the imaginary part of a complex array with dimensions NxM
If a figure is given explicitly, it will clear that existing figure and
@@ -170,6 +280,8 @@ def plot_imag(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwa
The length units to mark on the plot, default is um
cmap : str
Default is 'viridis', the colormap to plot with
cmap_label : str
What to label the colorbar when plotting
\\**kwargs
All other args are passed to fig.add_subplot(111, \\**kwargs)
@@ -178,53 +290,21 @@ def plot_imag(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwa
used_fig : matplotlib.figure.Figure
The figure object that was actually plotted to.
"""
if fig is None:
fig = plt.figure()
ax = fig.add_subplot(111, **kwargs)
else:
plt.figure(fig.number)
plt.gcf().clear()
if isinstance(im, t.Tensor):
imag = im[...,1].detach().cpu().numpy()
else:
imag = np.imag(im)
#Plot in a basis if it exists, otherwise dont
if basis is not None:
if isinstance(basis,t.Tensor):
basis = basis.detach().cpu().numpy()
# This fails if the
basis_norm = np.linalg.norm(basis, axis = 0)
basis_norm = basis_norm * get_units_factor(units)
extent = [0, imag.shape[-1]*basis_norm[1], 0, imag.shape[-2]*basis_norm[0]]
else:
extent=None
plt.imshow(imag, cmap = cmap, extent = extent)
cbar = plt.colorbar()
cbar.set_label('Imaginary Part (a.u.)')
if basis is not None:
plt.xlabel('X (' + units + ')')
plt.ylabel('Y (' + units + ')')
else:
plt.xlabel('j (pixels)')
plt.ylabel('i (pixels)')
return fig
plot_func = lambda x: np.imag(x)
return plot_image(im, plot_func=plot_func, fig=fig, basis=basis,
units=units, cmap=cmap, cmap_label=cmap_label,
**kwargs)
def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', **kwargs):
def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis', cmap_label='Amplitude (a.u.)', **kwargs):
"""Plots the amplitude of a complex array with dimensions NxM
If a figure is given explicitly, it will clear that existing figure and
plot over it. Otherwise, it will generate a new figure.
If a basis is explicitly passed, the image will be plotted in real-space
coordinates
coordinates.
Parameters
----------
im : array
@@ -237,6 +317,8 @@ def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis',
The length units to mark on the plot, default is um
cmap : str
Default is 'viridis', the colormap to plot with
cmap_label : str
What to label the colorbar when plotting
\\**kwargs
All other args are passed to fig.add_subplot(111, \\**kwargs)
@@ -245,45 +327,13 @@ def plot_amplitude(im, fig = None, basis=None, units='$\\mu$m', cmap='viridis',
used_fig : matplotlib.figure.Figure
The figure object that was actually plotted to.
"""
if fig is None:
fig = plt.figure()
ax = fig.add_subplot(111, **kwargs)
else:
plt.figure(fig.number)
plt.gcf().clear()
if isinstance(im, t.Tensor):
absolute = cmath.cabs(im).detach().cpu().numpy()
else:
absolute = np.absolute(im)
#Plot in a basis if it exists, otherwise dont
if basis is not None:
if isinstance(basis,t.Tensor):
basis = basis.detach().cpu().numpy()
# This fails if the
basis_norm = np.linalg.norm(basis, axis = 0)
basis_norm = basis_norm * get_units_factor(units)
extent = [0, absolute.shape[-1]*basis_norm[1], 0, absolute.shape[-2]*basis_norm[0]]
else:
extent=None
plt.imshow(absolute, cmap = cmap, extent = extent)
cbar = plt.colorbar()
cbar.set_label('Amplitude (a.u.)')
if basis is not None:
plt.xlabel('X (' + units + ')')
plt.ylabel('Y (' + units + ')')
else:
plt.xlabel('j (pixels)')
plt.ylabel('i (pixels)')
return fig
plot_func = lambda x: np.absolute(x)
return plot_image(im, plot_func=plot_func, fig=fig, basis=basis,
units=units, cmap=cmap, cmap_label=cmap_label,
**kwargs)
def plot_phase(im, fig=None, basis=None, units='$\\mu$m', cmap='auto', **kwargs):
def plot_phase(im, fig=None, basis=None, units='$\\mu$m', cmap='auto', cmap_label='Phase (rad)', **kwargs):
""" Plots the phase of a complex array with dimensions NxMx2
If a figure is given explicitly, it will clear that existing figure and
@@ -304,6 +354,8 @@ def plot_phase(im, fig=None, basis=None, units='$\\mu$m', cmap='auto', **kwargs)
The length units to mark on the plot, default is um
cmap : str
Default is 'viridis', the colormap to plot with
cmap_label : str
What to label the colorbar when plotting
\\**kwargs
All other args are passed to fig.add_subplot(111, \\**kwargs)
@@ -312,49 +364,18 @@ def plot_phase(im, fig=None, basis=None, units='$\\mu$m', cmap='auto', **kwargs)
used_fig : matplotlib.figure.Figure
The figure object that was actually plotted to.
"""
if fig is None:
fig = plt.figure()
ax = fig.add_subplot(111, **kwargs)
else:
plt.figure(fig.number)
plt.gcf().clear()
if isinstance(im, t.Tensor):
phase = cmath.cphase(im).detach().cpu().numpy()
else:
phase = np.angle(im)
if basis is not None:
if isinstance(basis,t.Tensor):
basis = basis.detach().cpu().numpy()
basis_norm = np.linalg.norm(basis, axis = 0)
basis_norm = basis_norm * get_units_factor(units)
extent = [0, phase.shape[-1]*basis_norm[1], 0, phase.shape[-2]*basis_norm[0]]
else:
extent=None
# If the user has matplotlib >=3.0, use the preferred colormap
if cmap == 'auto':
try:
plt.imshow(phase, cmap = 'twilight', extent=extent)
except:
plt.imshow(phase, cmap = 'hsv', extent=extent)
else:
plt.imshow(phase, cmap = cmap, extent=extent)
if 'twilight' in plt.colormaps():
cmap = 'twilight'
elif 'hsv' in plt.colormaps():
cmap = 'hsv'
else:
raise AttributeError('Neither twilight or hsv colormap exists in this screwed up matplotlib install')
cbar = plt.colorbar()
cbar.set_label('Phase (rad)')
if basis is not None:
plt.xlabel('X (' + units + ')')
plt.ylabel('Y (' + units + ')')
else:
plt.xlabel('j (pixels)')
plt.ylabel('i (pixels)')
return fig
plot_func = lambda x: np.angle(x)
return plot_image(im, plot_func=plot_func, fig=fig, basis=basis,
units=units, cmap=cmap, cmap_label=cmap_label,
**kwargs)
def plot_amplitude_surfacenorm():
@@ -390,38 +411,9 @@ def plot_colorized(im, fig=None, basis=None, units='$\\mu$m', **kwargs):
used_fig : matplotlib.figure.Figure
The figure object that was actually plotted to.
"""
if fig is None:
fig = plt.figure()
ax = fig.add_subplot(111, **kwargs)
else:
plt.figure(fig.number)
plt.gcf().clear()
if isinstance(im, t.Tensor):
im = cmath.torch_to_complex(im.detach().cpu())
if basis is not None:
if isinstance(basis,t.Tensor):
basis = basis.detach().cpu().numpy()
basis_norm = np.linalg.norm(basis, axis = 0)
basis_norm = basis_norm * get_units_factor(units)
extent = [0, im.shape[-1]*basis_norm[1], 0, im.shape[-2]*basis_norm[0]]
else:
extent=None
colorized = colorize(im)
plt.imshow(colorized, extent=extent)
if basis is not None:
plt.xlabel('X (' + units + ')')
plt.ylabel('Y (' + units + ')')
else:
plt.xlabel('j (pixels)')
plt.ylabel('i (pixels)')
return fig
plot_func = lambda x: colorize(x)
return plot_image(im, plot_func=plot_func, fig=fig, basis=basis,
units=units, **kwargs)
def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, **kwargs):
@@ -524,3 +516,254 @@ def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='pr
plt.colorbar()
return fig
def plot_nanomap_with_images(translations, get_image_func, values=None, mask=None, basis=None, fig=None, nanomap_units='$\\mu$m', image_units='$\\mu$m', convention='probe', image_title='Image', image_colorbar_title='Image Amplitude', nanomap_colorbar_title='Integrated Intensity', cmap='viridis', **kwargs):
"""Plots a nanomap, with an image or stack of images for each point
In many situations, ptychography data or the output of ptychography
reconstructions is formatted as a set of images associated with various
points in real space. This function is designed to allow for browsing
through this kind of data, by making it possible to visualize a
"""
# This should pull heavily from the dataset.inspect function
# In fact, I should be able to replace most of that function with a
# call to this function once it's built
# We start by making the figure and axes
# The key will be writing this so it works okay when called in "update"
# mode, i.e. on a figure that already has this thing showing.
if fig is None:
fig = plt.figure(figsize=(8,5.3))
else:
plt.figure(fig.number)
plt.gcf().clear()
if hasattr(fig, 'nanomap_cids'):
for cid in fig.nanomap_cids:
fig.canvas.mpl_disconnect(cid)
# Does figsize work with the fig.subplots, or just for plt.subplots?
axes = fig.subplots(1,2)
fig.tight_layout(rect=[0.04, 0.09, 0.98, 0.96])
plt.subplots_adjust(wspace=0.25) #avoids overlap of labels with plots
axslider = plt.axes([0.15,0.06,0.75,0.03])
# This gets the set of sizes for the points in the nanomap
def calculate_sizes(idx):
bbox = axes[0].get_window_extent().transformed(fig.dpi_scale_trans.inverted())
s0 = bbox.width * bbox.height / translations.shape[0] * 72**2 #72 is points per inch
s0 /= 4 # A rough value to make the size work out
s = np.ones(translations.shape[0]) * s0
s[idx] *= 4
return s
def update_colorbar(im):
#
# This solves the problem of the colorbar being changed
# when the forward and back buttons are used!!!
#
if hasattr(im, 'norecurse') and im.norecurse:
im.norecurse=False
return
im.norecurse=True
# This is needed to update the colorbar
# only change limits if array contains multiple values
if np.min(im.get_array()) != np.max(im.get_array()):
im.set_clim(vmin=np.min(im.get_array()),
vmax=np.max(im.get_array()))
#
# The meatiest part of this program, here we just go through and
# set up the plot how we want it
#
# First we set up the left-hand plot, which shows an overview map
axes[0].set_title('Relative Displacement Map')
translations = translations.detach().cpu().numpy()
if convention.lower() != 'probe':
translations = translations * -1
s = calculate_sizes(0)
nanomap_units_factor = get_units_factor(nanomap_units)
nanomap = axes[0].scatter(nanomap_units_factor * translations[:,0],
nanomap_units_factor * translations[:,1],
s=s,c=values, picker=True)
axes[0].invert_xaxis()
axes[0].set_facecolor('k')
axes[0].set_xlabel('Translation x ('+nanomap_units+')', labelpad=1)
axes[0].set_ylabel('Translation y ('+nanomap_units+')', labelpad=1)
cb1 = plt.colorbar(nanomap, ax=axes[0], orientation='horizontal',
format='%.2e',
ticks=ticker.LinearLocator(numticks=5),
pad=0.17,fraction=0.1)
cb1.ax.set_title(nanomap_colorbar_title, size="medium", pad=5)
cb1.ax.tick_params(labelrotation=20)
if values is None:
# This seems to do a good job of leaving the appropriate space
# where the colorbar should have been to avoid stretching the
# nanomap plot, while still not showing the (now useless) colorbar.
cb1.remove()
# Now we set up the second plot, which shows the individual
# diffraction patterns
axes[1].set_title(image_title)
#Plot in a basis if it exists, otherwise dont
if basis is not None:
axes[1].set_xlabel('X (' + image_units + ')')
axes[1].set_ylabel('Y (' + image_units + ')')
example_im = get_image_func(0)
if isinstance(example_im, t.Tensor):
example_im = example_im.cpu().numpy()
if isinstance(basis,t.Tensor):
np_basis = basis.detach().cpu().numpy()
else:
np_basis = basis
# This fails if the basis is not rectangular
basis_norm = np.linalg.norm(np_basis, axis = 0)
basis_norm = basis_norm * get_units_factor(image_units)
extent = [0, example_im.shape[-1]*basis_norm[1], 0,
example_im.shape[-2]*basis_norm[0]]
else:
axes[1].set_xlabel('j (pixels)')
axes[1].set_ylabel('i (pixels)')
extent=None
im=get_image_func(0)
if len(im.shape) >= 3:
im_idx=0
axes[1].image_idx = im_idx
im = im.reshape(-1,im.shape[-2],im.shape[-1])[im_idx]
axes[1].text_box = axes[1].text(0.98, 0.98, str(im_idx), color='w',
fontsize=14,
horizontalalignment='right',
verticalalignment='top',
transform=axes[1].transAxes)
axes[1].text_box.set_path_effects(
[patheffects.Stroke(linewidth=2, foreground='black'),
patheffects.Normal()])
meas = axes[1].imshow(im, extent=extent, cmap=cmap)
cb2 = plt.colorbar(meas, ax=axes[1], orientation='horizontal',
format='%.2e',
ticks=ticker.LinearLocator(numticks=5),
pad=0.17,fraction=0.1)
cb2.ax.tick_params(labelrotation=20)
cb2.ax.set_title(image_colorbar_title, size="medium", pad=5)
cb2.ax.callbacks.connect('xlim_changed', lambda ax: update_colorbar(meas))
# This function handles all the updating, except for moving the
# slider value. This is done because the slider widget is
# ultimately responsible for triggering an update, so all other
# updates are done by changing the slider widget value which
# then triggers this
def update(idx, im_idx=None):
# We have to explicitly make it an integer because the slider will
# output floats (even if they are still integer-valued)
idx = int(idx)
# Get the new data for this index
im = get_image_func(idx)
if len(im.shape) >= 3:
if im_idx == None and hasattr(axes[1],'image_idx'):
im_idx = axes[1].image_idx
elif im_idx == None:
im_idx=0
axes[1].image_idx = im_idx
axes[1].text_box.set_text(str(im_idx))
im = im.reshape(-1,im.shape[-2],im.shape[-1])[im_idx]
# Now we resize the nanomap to show the new selection
axes[0].collections[0].set_sizes(calculate_sizes(idx))
# And we update the data in the image as well
ax_im = axes[1].images[-1]
ax_im.set_data(im)
update_colorbar(ax_im)
#
# Now we define the functions to handle various kinds of events
# that can be thrown our way
#
# We start by creating the slider here, so it can be used
# by the update hooks.
slider = Slider(axslider, 'Image #', 0, translations.shape[0]-1, valstep=1, valfmt="%d")
# This handles scroll wheel and keypress events
def on_action(event):
im = get_image_func(0)
if event.inaxes is axes[1] and len(im.shape) >=3:
# This is triggered if the data to display has more than 2
# dimensions (i.e. is an image stack) and the event originates
# while the mouse is within the image display
im = im.reshape(-1,im.shape[-2],im.shape[-1])
im_idx = axes[1].image_idx
if event.key == 'up' or event.button == 'up' \
or event.key == 'left':
im_idx = (im_idx - 1) % im.shape[0]
if event.key == 'down' or event.button == 'down' \
or event.key == 'right':
im_idx = (im_idx + 1) % im.shape[0]
axes[1].image_idx=im_idx
slider.set_val(slider.val)#update(slider.val,im_idx=im_idx)
return # This prevents the rest from also happening
# Otherwise the if statements can throw errors when the
# event type isn't right, this way they just don't trigger
if not hasattr(event, 'button'):
event.button = None
if not hasattr(event, 'key'):
event.key = None
if event.key == 'up' or event.button == 'up' or event.key == 'left':
idx = slider.val - 1
elif event.key == 'down' or event.button == 'down' or event.key == 'right':
idx = slider.val + 1
else:
# This prevents errors from being thrown on irrelevant key
# or mouse input
return
# Handle the wraparound and trigger the update
idx = int(idx) % translations.shape[0]
slider.set_val(idx)
# This handles "pick" events in the nanomap
def on_pick(event):
# If we don't filter on type of event, this will also capture,
# for example, scroll events that happen over the nanomap
if event.mouseevent.button == 1:
slider.set_val(event.ind[0])
# Here we connect the various update functions
cid1 = fig.canvas.mpl_connect('pick_event',on_pick)
cid2 = fig.canvas.mpl_connect('key_press_event',on_action)
cid3 = fig.canvas.mpl_connect('scroll_event',on_action)
# It's so dumb that matplotlib doesn't automatically track this for you
fig.nanomap_cids = [cid1,cid2,cid3]
slider.on_changed(update)
# Throw an extra update into the mix just to get rid of any things
# (like the nanomap dot sizes) that otherwise would change on the
# first update
update(0)
CDTools/tools/projectors/__init__.py
-3
View File
@@ -1,3 +0,0 @@
from __future__ import division, print_function, absolute_import
from CDTools.tools.projectors.projectors import *
CDTools/tools/projectors/projectors.py
-80
View File
@@ -1,80 +0,0 @@
"""This module contains various projection functions
These functions are useful when defining declarative algorithms to run
alongside the automatic differentiation ones, for comparison or in a
situation where they might be needed.
"""
from __future__ import division, print_function, absolute_import
from CDTools.tools.cmath import *
import torch as t
__all__ = ['modulus', 'support']
def modulus(wavefront, intensities, mask = None):
"""Implements the modulus constraint in torch
This accepts a tensor representing the propagated simulated wavefront(s),
where the last dimension represents the real and imaginary components of
the propagated wavefield(s). It projects the modulus of the diffraction
pattern onto the modulus of the simulated wavefield.
It assumes that the wavefront is stored in an array
[i,j] where i corresponds to the y-axis and j corresponds to the
x-axis, with the origin following the CS standard of being in the
upper right.
Parameters
----------
wavefront : torch.Tensor
The JxNxMx2 stack of complex propagated wavefronts
intensities : torch.Tensor
The measured diffraction pattern(s) stored as an JxNxM stack of real tensors
mask : torch.Tensor
A mask for the intensities array with shape JxNxM, where bad detector pixels are set to 0 and usable pixels set to 1
Returns
-------
projected : torch.Tensor
The JxNxMx2 projected wavefield with corrected intensities
"""
# Calculate amplitudes from intensities
amplitudes = t.sqrt(intensities)
# Normalize wavefront so the complex elements have modulus one
wavefront_mag = cabs(wavefront)
projected = wavefront * (amplitudes / wavefront_mag)[...,None]
# Replace amplitude of wavefront with measured amplitude
if mask is not None:
selection = mask == 0
# Apply the mask to replace unmasked pixels in the original wavefront
projected = projected.masked_scatter(selection, wavefront.masked_select(selection))
return projected
def support(wavefront, support):
"""Implements the support constraint in torch
This accepts a torch tensor representing (a) simulated wavefield(s),
where the last dimension represents the real and imaginary components of
the propagated wavefield(s). It projects the support of the imaged object
onto the simulated wavefront via a support mask.
It assumes that the wavefront is stored in an array
[i,j] where i corresponds to the y-axis and j corresponds to the
x-axis, with the origin following the CS standard of being in the
upper right.
Parameters
----------
wavefront : torch.Tensor
The JxNxMx2 stack of complex propagated wavefronts
support : torch.Tensor
An NxM support, with 1s within the support and 0s outside
Returns
-------
projected : torch.Tensor
The JxNxMx2 wavefield with the support mask applied
"""
return wavefront * support.to(wavefront.dtype)[...,None]
CDTools/tools/propagators/propagators.py
+168 -190
View File
@@ -3,9 +3,7 @@
All the functions here are designed for use in an automatic differentiation
ptychography model. Each function implements a different propagator.
"""
from __future__ import division, print_function, absolute_import
from CDTools.tools.cmath import *
import torch as t
from torch.nn.functional import grid_sample
from scipy import fftpack
@@ -19,7 +17,6 @@ __all__ = ['far_field', 'near_field',
'high_NA_far_field',
'generate_generalized_angular_spectrum_propagator']
def far_field(wavefront):
"""Implements a far-field propagator in torch
@@ -47,8 +44,10 @@ def far_field(wavefront):
propagated : torch.Tensor
The JxNxMx2 propagated wavefield
"""
return fftshift(t.fft(ifftshift(wavefront), 2, normalized=True))
shifted = t.fft.ifftshift(wavefront, dim=(-1,-2))
propagated = t.fft.fft2(shifted, norm='ortho')
return t.fft.fftshift(propagated, dim=(-1,-2))
def inverse_far_field(wavefront):
@@ -74,7 +73,9 @@ def inverse_far_field(wavefront):
propagated : torch.Tensor
The JxNxMx2 exit wavefield
"""
return fftshift(t.ifft(ifftshift(wavefront), 2, normalized=True))
shifted = t.fft.ifftshift(wavefront, dim=(-1,-2))
propagated = t.fft.ifft2(shifted, norm='ortho')
return t.fft.fftshift(propagated, dim=(-1,-2))
def generate_high_NA_k_intensity_map(sample_basis, det_basis,det_shape,distance, wavelength, *args, lens=False, **kwargs):
@@ -286,38 +287,37 @@ def high_NA_far_field(wavefront, k_map, intensity_map=None):
# np.ones_like(k_map[0,:-1,:-1,0].cpu().numpy()))
#plt.show()
def process_wavefield_stack(low_NA_wavefield):
real_output = grid_sample(low_NA_wavefield[None,:,:,:,0],k_map,mode='bilinear',padding_mode='zeros', align_corners=False)
imag_output = grid_sample(low_NA_wavefield[None,:,:,:,1],k_map,mode='bilinear',padding_mode='zeros', align_corners=False)
# grid_sample doesn't work on complex-valued wavefields
real_output = grid_sample(low_NA_wavefield[None,:,:,:].real,k_map,mode='bilinear',padding_mode='zeros', align_corners=False)
imag_output = grid_sample(low_NA_wavefield[None,:,:,:].imag,k_map,mode='bilinear',padding_mode='zeros', align_corners=False)
result = t.stack((real_output[0,:,:,:],imag_output[0,:,:,:]),dim=3)
result = real_output[0,:,:,:] + 1j * imag_output[0,:,:,:]
if intensity_map is not None:
result = result * intensity_map[None,:,:,None]
result = result * intensity_map[None,:,:]
return result
original_dim = wavefront.dim()
if original_dim == 2:
result = process_wavefield_stack(low_NA_wavefield[None,:,:])
return result[0,:,:]
if original_dim == 3:
result = process_wavefield_stack(low_NA_wavefield[None,:,:,:])
return result[0,:,:,:]
if original_dim == 4:
result = process_wavefield_stack(low_NA_wavefield)
return result
if original_dim == 5:
if original_dim == 4:
result = []
for i in range(low_NA_wavefield.size()[0]):
result.append(process_wavefield_stack(low_NA_wavefield[i,:,:,:,:]))
result.append(process_wavefield_stack(low_NA_wavefield[i,:,:,:]))
return t.stack(result)
else:
raise IndexError('Wavefield had incorrect number of dimensions')
def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, remove_z_phase=False, bandlimit=None, **kwargs):
"""Generates an angular-spectrum based near-field propagator from experimental quantities
This function generates an angular-spectrum based near field
propagator that will work on torch Tensors. The function is structured
this way - to generate the propagator first - because the
@@ -325,15 +325,14 @@ def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, r
propagator is used in a reconstruction program, then it will be best
to calculate this mask once and reuse it.
Formally, this propagator is the complex conjugate of the fourier
transform of the convolution kernel for light propagation in free
space
Formally, this propagator is the the fourier transform of the convolution
kernel for light propagation in free space
If the optional bandlimit parameter is set, the propagator will be set
to zero beyond an explicit bandlimiting frequency. This is helpful if the
propagator will be used in a repeated multiply/propagate framework such
as a multislice algorithm, where it helps to prevent aliasing.
Parameters
----------
shape : array
@@ -354,34 +353,43 @@ def generate_angular_spectrum_propagator(shape, spacing, wavelength, z, *args, r
propagator : torch.Tensor
A phase mask which accounts for the phase change that each plane wave will undergo.
"""
ki = 2 * np.pi * fftpack.fftfreq(shape[0],spacing[0])
kj = 2 * np.pi * fftpack.fftfreq(shape[1],spacing[1])
Kj, Ki = np.meshgrid(kj,ki)
# Define this as complex so the square root properly gives
# k>k0 components imaginary frequencies
k0 = np.complex128((2*np.pi/wavelength))
# Properly accuount for evanescent waves
if z >=0:
propagator = np.exp(1j*np.sqrt(k0**2 - Ki**2 - Kj**2) * z)
# Internally, the generalized propagation function is used, so we start
# by creating an appropriate basis
# This creates a real-valued tensor which matches the kind of complex
# number dtype requested in **kwargs
if 'dtype' in kwargs:
basis = t.real(t.zeros([3,2], **kwargs))
else:
propagator = np.exp(1j*np.conj(np.sqrt(k0**2 - Ki**2 - Kj**2)) * z)
if remove_z_phase:
propagator *= np.exp(-1j * k0 * z)
basis = t.zeros([3,2], dtype=t.float32)
spacing = t.as_tensor(spacing, dtype=basis.dtype)
basis[1,0] = -spacing[0]
basis[0,1] = -spacing[1]
# And similarly, the offset is just z along the z direction
offset = t.zeros([3], dtype=basis.dtype)
offset[2] = z
# And we call the generalized function!
propagator = generate_generalized_angular_spectrum_propagator(shape, basis,
wavelength, offset,
propagate_along_offset=remove_z_phase, **kwargs)
if z < 0:
propagator = t.conj(propagator)
# Bandlimiting is not implemented in the generalized function, because it
# has a less clear meaning in that setting, so we apply it here instead
if bandlimit is not None:
Rs = np.sqrt((Ki / np.max(ki))**2 + (Kj / np.max(kj))**2)
ki = 2 * np.pi * t.fft.fftfreq(shape[0],spacing[0])
kj = 2 * np.pi * t.fft.fftfreq(shape[1],spacing[1])
Ki, Kj = t.meshgrid(ki,kj)
min_radius = min(t.max(ki),t.max(kj))
Rs = t.sqrt((Ki/t.max(ki))**2 + (Kj/t.max(kj))**2)
propagator = propagator * (Rs < bandlimit)
# Take the conjugate explicitly here instead of negating
# the previous expression to ensure that complex frequencies
# get mapped to values <1 instead of >1
propagator = complex_to_torch(np.conj(propagator))
return propagator.to(*args, **kwargs)
return propagator
def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, offset_vector, *args, propagation_vector=None, propagate_along_offset=False, **kwargs):
@@ -391,16 +399,15 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o
propagator that will work on torch Tensors. The function is structured
this way - to generate the propagator first - because the
generation of the propagation mask is a bit expensive and if this
propagator is used in a reconstruction program, then it will be best
to calculate this mask once and close over it.
propagator is used in a reconstruction program, it will be best
to calculate this mask once and then reuse it it.
Formally, this propagator is the complex conjugate of the fourier
transform of the convolution kernel for light propagation in free
space. It will map a ligh field at an input plane, with the size
and shape defined by the shape and basis inputs, and map it to a
plane of the same size and shape offset by the offset vector. It
is designed to work on any wavefield defined on an array of
parallelograms.
Formally, this propagator is the fourier transform of the convolution
kernel for light propagation in free space. It will map a light field
at an input plane, with the size and shape defined by the shape and basis
inputs, and map it to a plane of the same size and shape offset by the
offset vector. It is designed to work on any wavefield defined on an
array of parallelograms.
In addition, if the propagation_vector is set, there is an assumed
@@ -418,33 +425,27 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o
propagation_vector option
Note that, unlike in the case of the simple angular spectrum propagator,
the direction of "forward propagation" is defined by the offset vector.
Therefore, in the simple case of a perpendicular offset, there will be
no difference between using an offset vector or the negative of the
offset vector. This is because, for the light propagation problem to
be well posed, the assumption must be made that light only passes through
the plane of the known wavefield in one direction. Mathematically, this
corresponds to a choice of uniform phase objects either accumulating
positive or negative phase. In the simple propagation case, there is
no ambiguity introduced by always choosing the light field to propagate
along the positive z direction. In the general case, there is no equivalent
obvious choice - thus, the light is always assumed to pass through the
initial plane travelling in the direction of the final plane.
the direction of "forward propagation" is defined by the propagation
vector, or (if the propagation vector is not defined), the offset
vector. Therefore, in the simple case of a perpendicular offset,
there will be no difference between using an offset vector or the
negative of the offset vector. This is because, for the light
propagation problem to be well posed, the assumption must be made that
light only passes through the plane of the known wavefield in one
direction. We always assume that light passes through the initial plane
travelling in the direction of the final plane.
Practically, if one wants to simulate inverse propagation, there are then
Practically, if one wants to simulate inverse propagation, there are
two possible approaches. First, one can use the inverse_near_field
function, which simulates the inverse propagation problem and therefore
will naturally simulate propagation in the opposite direction. Second,
one can explicitly include a propagation_vector argument, which overrides
the offset vector in defining the direction in which light passes through
the input plane. However, in this case, the resulting light field will have
the overall phase accumulation due to propagation along the propagation
vector removed, which may not be the intended behavior. However, this is
not recommended, as inverse propagation will tend to magnify evanescent
waves - it is therefore preferable (unless there is a specific need to
account for evanescent waves properly) to use the inverse near field
propagator
one can explicitly include a propagation_vector argument, in the direction
opposite to the offset vector. Note for both cases that this function
will only return propagators which suppress evanescent waves. Thus,
propagating forward and then backward by either of these two methods
will lead to a supression of evanescant waves. If you need a propagator
that will cause evanescent waves to undergo exponential growth, good
for you, but this function will not provide it for you.
Parameters
----------
@@ -467,126 +468,101 @@ def generate_generalized_angular_spectrum_propagator(shape, basis, wavelength, o
A phase mask which accounts for the phase change that each plane wave will undergo.
"""
# These check for any pytorch inputs and convert them
try:
basis = basis.detach().cpu().numpy()
except:
pass
# make sure everything is in pytorch, and set the propagation vector
# appropriately if propagate_along_offset is chosen
basis = t.as_tensor(basis)
offset_vector = t.as_tensor(offset_vector, dtype=basis.dtype)
shape = np.array(tuple(shape))
try:
offset_vector = offset_vector.detach().cpu().numpy()
except:
pass
if propagation_vector is not None:
try:
propagation_vector = propagation_vector.detach().cpu().numpy()
except:
pass
try:
wavelength = wavelength.detach().cpu().numpy()[0]
except:
pass
# First we calculate a dual basis for the real space grid
inv_basis = np.linalg.pinv(basis).transpose()
# Then we calculate the frequencies in (i,j) space
ki = 2 * np.pi * fftpack.fftfreq(shape[0])
kj = 2 * np.pi * fftpack.fftfreq(shape[1])
K_ij = np.stack(np.meshgrid(ki,kj, indexing='ij'))
# Now we convert these to frequencies in reciprocal space
# These frequencies span the 2D plane of the input wavefield.
K_xyz = np.tensordot(inv_basis, K_ij, axes=1)
# Now we need to apply two corrections to the standard AS method.
# First, we calculate a phase mask which corresponds to the
# shift of the final plane away from the perpendicular direction
# from the input plane. We don't need to extract the perpendicular
# component of the shift because the K_xyz vectors are naturally in the
# input plane.
# This may have a sign error - must be checked
phase_mask = np.exp(-1j * np.tensordot(offset_vector,K_xyz,axes=1))
# Next, we apply a shift to the k-space vectors which sets up
# propagation such that a uniform phase object will propagate along the
# propagation axis. This is not modeling a physical effect, but simply is
# the clearest way to do a rigorous simulation while preventing
# aliasing-related challenges. If used (as is by default), be aware
# and prepare the input wavefields appropriately.
if propagate_along_offset:
propagation_vector = offset_vector
perpendicular_dir = np.cross(basis[:,1],basis[:,0])
perpendicular_dir /= np.linalg.norm(perpendicular_dir)
offset_perpendicular = np.dot(perpendicular_dir, offset_vector)
k0 = 2*np.pi/wavelength
sign_correction = 1
# Only implement the shift if the flag is set to True
if propagation_vector is not None:
prop_perpendicular = np.dot(perpendicular_dir, propagation_vector)
prop_parallel = propagation_vector - perpendicular_dir \
* prop_perpendicular
if np.linalg.norm(propagation_vector) == 0:
# for numerical stability if this is exactly zero we need
# a special case
k_offset = np.array([0,0,0])
else:
k_offset = prop_parallel * k0 / np.linalg.norm(propagation_vector)
propagation_vector = t.as_tensor(propagation_vector, dtype=basis.dtype)
K_xyz = K_xyz - k_offset[:,None,None]
# There apparently is a sign correction that I need to apply
#sign_correction = np.sign(np.dot(perpendicular_dir,propagation_vector))
sign_correction = np.sign(np.dot(offset_vector,propagation_vector))
# we also need to remove the z-dependence on the phase
# This time, though, the z-dependence actually has to do with
# the out of plane component of k at the central offset. Normally
# this is 0, so the z-component is just k0, but not in this case
# We only need one case here, unlike with the propagator, because
# k_offset will always be less than k0
phase_mask *= np.exp(-1j * np.sqrt(k0**2 - np.linalg.norm(k_offset)**2)
* sign_correction
* np.abs(offset_perpendicular))
#
# In this section, we calculate the wavevectors associated with each
# pixel in Fourier space. This is the meat of the function
#
# Redefine this as complex so the square root properly gives
# k>k0 components imaginary frequencies
k0 = np.complex128(k0)
# First we calculate a dual basis for the real space grid
inv_basis = t.linalg.pinv(basis).transpose(0,1)
# Then we calculate the frequencies in (i,j) space
ki = 2 * np.pi * t.fft.fftfreq(shape[0], dtype=inv_basis.dtype)
kj = 2 * np.pi * t.fft.fftfreq(shape[1], dtype=inv_basis.dtype)
K_ij = t.stack(t.meshgrid(ki,kj))
# Finally, generate the propagator!
# Must have cases to ensure that evanescent waves decay instead of grow
if sign_correction > 0:
propagator = np.exp(1j*np.sqrt(k0**2 - np.linalg.norm(K_xyz,axis=0)**2)
* sign_correction * np.abs(offset_perpendicular))
# Now we convert these to frequencies in reciprocal space
# These frequencies span the 2D plane of the input wavefield,
# hence K_ip for "in-plane"
K_ip = t.tensordot(inv_basis, K_ij, dims=1)
# Now, we need to generate the out-of-plane direction, so we can
# expand these Ks to the full Ks in 3D reciprocal space.
# THis is broken down into steps to avoid floating point underflow
# which was a real problem that showed up for electron ptycho
b1_dir = basis[:,0] / t.linalg.norm(basis[:,0])
b2_dir = basis[:,1] / t.linalg.norm(basis[:,1])
perpendicular_dir = t.cross(b1_dir,b2_dir)
perpendicular_dir /= t.linalg.norm(perpendicular_dir)
# We set the sign of the propagation direction appropriately
if propagation_vector is not None:
perpendicular_dir *= t.sign(t.dot(perpendicular_dir,propagation_vector))
else:
propagator = np.exp(-1j * np.conj(np.sqrt(k0**2 -
np.linalg.norm(K_xyz,axis=0)**2))
* np.abs(offset_perpendicular))
propagator *= phase_mask
pass
perpendicular_dir *= t.sign(t.dot(perpendicular_dir,offset_vector))
# Take the conjugate explicitly here instead of negating
# the previous expression to ensure that complex frequencies
# get mapped to values <1 instead of >1
propagator = complex_to_torch(np.conj(propagator))
return propagator.to(**kwargs)
# Then, if we have a propagation vector, we shift the in-plane
# components of all the pixels to be centered around the in-plane
# component of the propagation vector.
if propagation_vector is not None:
prop_dir = (propagation_vector /
t.linalg.norm(propagation_vector))
K_0 = 2*np.pi / wavelength * prop_dir
K_0_ip = K_0 - t.dot(perpendicular_dir,K_0) * perpendicular_dir
K_ip += K_0_ip[:,None,None]
else:
K_0 = t.zeros_like(offset_vector)
# Now, we have accurate in-plane values for K, so we can calculate the
# out-of-plane part. We start by calculating it's squared magnitude
K_oop_squared = (2*np.pi/wavelength)**2 - t.linalg.norm(K_ip,dim=0)**2
# Then, we take the square root and assign it the appropriate direction,
# adding to get the full 3D wavevectors. Note that we convert to complex
# before the square root to appropriately map negative numbers to
# complex frequencies
K = K_ip + perpendicular_dir[:,None,None] \
* t.sqrt(t.complex(K_oop_squared,t.zeros_like(K_oop_squared)))
#
# In this section, we take the inner product of the calcualted
# wavevectors with the offset vector, to get the phase shift
# experienced by each plane wave.
#
# We need to convert to complex because K is complex
offset_vector = t.complex(offset_vector,t.zeros_like(offset_vector))
# We then subtract off K_0, essentially setting the phase offset
# experienced by K_0 to 0. K_0 will already have been set to 0 if
# there was no propagation vector set.
K_m_K_0 = K - K_0[:,None,None]
# We actually calculate the phase mask
phase_mask = t.tensordot(offset_vector,K_m_K_0, dims=1)
# And we take the conjugate if needed, which makes sure that the
# imaginary part is always positive (supressing evanescent waves)
if propagation_vector is not None \
and t.sign(t.dot(propagation_vector,offset_vector.real)) == -1:
phase_mask = t.conj(phase_mask)
# We return the final mask, in the form asked for!
return t.exp(1j*phase_mask).to(*args,**kwargs)
def near_field(wavefront, angular_spectrum_propagator):
""" Propagates a wavefront via the angular spectrum method
@@ -595,13 +571,14 @@ def near_field(wavefront, angular_spectrum_propagator):
represents the real and imaginary components of the wavefield, and
returns the near-field propagated version of it. It does this
using the supplied angular spectrum propagator, which is a premade
phase mask.
phase mask representing the Fourier transform of the kernel for
light propagation in the desired geometry.
Parameters
----------
wavefront : torch.Tensor
The JxNxMx2 stack of complex wavefronts to be propagated
The (Leading Dims)xNxM stack of complex wavefronts to be propagated
angular_spectrum_propagator : torch.Tensor
The NxM phase mask to be applied during propagation
@@ -610,7 +587,7 @@ def near_field(wavefront, angular_spectrum_propagator):
propagated : torch.Tensor
The propagated wavefront
"""
return t.ifft(cmult(angular_spectrum_propagator,t.fft(wavefront,2)), 2)
return t.fft.ifft2(angular_spectrum_propagator * t.fft.fft2(wavefront))
@@ -651,7 +628,8 @@ def inverse_near_field(wavefront, angular_spectrum_propagator):
propagated : torch.Tensor
The inverse propagated wavefront
"""
return t.ifft(cmult(t.fft(wavefront,2), cconj(angular_spectrum_propagator)), 2)
return t.fft.ifft2(t.fft.fft2(wavefront)
* t.conj(angular_spectrum_propagator))
README.md
+6 -6
View File
@@ -6,21 +6,21 @@ CDTools is a python library for ptychography and CDI reconstructions, using an A
# imports
from matplotlib import pyplot as plt
from CDTools.datasets import Ptycho_2D_Dataset
from CDTools.models import SimplePtycho
from CDTools.models import FancyPtycho
# Load the file
dataset = Ptycho_2D_Dataset.from_cxi('ptycho_data.cxi')
# Generate a model from the data
model = SimplePtycho.from_dataset(dataset)
model = FancyPtycho.from_dataset(dataset)
# Run a reconstruction
for i, loss in enumerate(model.Adam_optimize(10, dataset)):
print(i, loss)
for loss in model.Adam_optimize(10, dataset):
print(model.report())
# And look at the results!
model.inspect(dataset)
model.compare(dataset)
model.inspect(dataset) # See the reconstructed object, probe, etc.
model.compare(dataset) # See how the simulated and measured patterns compare
plt.show()
```
conda_requirements.txt
+1 -1
View File
@@ -2,7 +2,7 @@ numpy>=1.0
scipy>=1.0
matplotlib>=2.0
python-dateutil
pytorch>=1.3.0
pytorch>=1.9.0
h5py>=2.1
pytest
sphinx
converters/NSLS2_HXN_hdf5_to_CXI.py
-131
View File
@@ -1,131 +0,0 @@
"""
Purpose: Convert NSLSII HXN hdf5 files to CXI files for analysis with CDTools.
Author: David Rower
Date: December 2019
"""
import numpy as np
import pickle
import h5py
import os
import CDTools
from CDTools.tools import data as cdtdata
from matplotlib import pyplot as plt
from scipy.spatial.transform import Rotation
from datetime import datetime
def create_cxi_from_NSLS2_HXN_2DFly(data_dir, save_str, scan_number,
wavelength, theta, ROI_corner_xy, metadata):
"""Converts NSLS2 HXN 2D Fly scan data (from pickle and hdf5) to CXI format
Assumes scan files will live in data_dir with naming convention
pickle: <data_dir>/scan_<scan_number>.pickle,
hdf5: <data_dir>/scan_<scan_number>.hdf5,
and will create the file <data_dir>/scan_<scan_number>.cxi.
Parameters
----------
data_dir : str
Input data directory
save_str : str
Output data name
scan_number : int
A scan index number
theta : float
Rotation angle of sample in HXN convention, in degrees
ROI_corner_xy : np.array
1x2 array containing x, y corner of detector ROI
metadata : dict
Contains metadata relevant to the experiment
"""
## Load in pickle and hdf5 files
scan_str = "scan_" + scan_number
# Load pickle (includes useful data about scan not in .hdf5 file)
with open(os.path.join(data_dir, scan_str+".pickle"), 'rb') as f:
scan_pickle = pickle.load(f)
assert scan_pickle['plan_type'] == "FlyPlan2D", "Code only for FlyPlan2D."
# Load hdf5 file
scan_hdf5 = h5py.File(os.path.join(data_dir, scan_str+".h5"), 'r')
## Let's attempt to convert this bad boy
scan_cxi = cdtdata.create_cxi(save_str)
## Add source
cdtdata.add_source(scan_cxi, wavelength=wavelength)
scan_cxi['entry_1/instrument_1/source_1']['name'] = scan_pickle['beamline_id']
## Add sample
theta = np.radians(theta)
sample_unit_vecs = Rotation.from_rotvec(-theta * np.array([0,1,0])).as_dcm()
orientation = np.hstack((sample_unit_vecs[:,0], sample_unit_vecs[:,1]))
translation = np.zeros(3)
sample_info_dict = {
"name" : "TaTe4",
"orientation" : orientation,
"translation" : translation
}
cdtdata.add_sample_info(scan_cxi, sample_info_dict)
## Add other metadata for experiment
metadata['start_time'] = datetime.fromtimestamp(scan_pickle['time'])
cdtdata.add_entry_info(scan_cxi, metadata)
## Add detector
# Constant detector parameters
detector_pixel_size = 55e-6 # meters
detector_height_px = 515 # px ### WARNING: NEED TO CHECK THIS
detector_width_px = 515 # px
# Geometry parameters from scan files
distance = scan_pickle['dist_detector'] * 1e-3 # assuming mm, almost sure
gamma = np.radians(scan_pickle['gamma_detector'])
delta = np.radians(scan_pickle['delta_detector'])
Rg = Rotation.from_rotvec(-gamma * np.array([0,1,0])).as_dcm() # cw about y
Rd = Rotation.from_rotvec(-delta * Rg[:,0]).as_dcm() # cw about rotated x
RdRg = np.matmul(Rd, Rg)
# Define detector basis: row vectors for y and x detector axes
basis = detector_pixel_size * np.array([[0.,-1.,0.],[-1.,0.,0.]])
basis = np.matmul(RdRg,basis.T).T
# Define corner posiiton: first find center, then offset it
corner_pos = np.dot(RdRg, distance * np.array([0.,0.,1.]))
if ROI_corner_xy[0] is None:
ROI_corner_xy[0] = 0.
if ROI_corner_xy[1] is None:
ROI_corner_xy[1] = 0.
corner_pos -= basis[0,:] * (detector_width_px/2. - ROI_corner_xy[0])
corner_pos -= basis[1,:] * (detector_height_px/2. - ROI_corner_xy[1])
# Add detector data finally
cdtdata.add_detector(scan_cxi, distance, basis.T, corner=corner_pos)
## Add data
axes = ['translation'] + scan_pickle['axes'] # THIS IS ONLY FOR FLY2D
data = np.copy(scan_hdf5['entry']['instrument']['detector']['data'])
data[data == 0] = 1 # to prevent divide by zero in log error
cdtdata.add_data(scan_cxi, data, axes)
## Add translations
x_bounds = scan_pickle['scan_range'][0]
y_bounds = scan_pickle['scan_range'][1]
xx, yy = np.meshgrid(np.linspace(*x_bounds, scan_pickle['num1']),
np.linspace(*y_bounds, scan_pickle['num2']))
translations = (1e-6 *
np.stack((xx.ravel(), yy.ravel(), np.zeros_like(xx.ravel())), axis=1))
cdtdata.add_ptycho_translations(scan_cxi, translations)
## Close hdf5 file
scan_hdf5.close()
converters/TITAN_stem.py
-208
View File
@@ -1,208 +0,0 @@
"""
Purpose: Convert file collection from Jim Lebeau's TITAN microscope to .CXI
Author: Abe Levitan
Date: January 2019
"""
import numpy as np
import pickle
import h5py
import os
import CDTools
from CDTools.tools import data as cdtdata
from CDTools.datasets import Ptycho2DDataset
from matplotlib import pyplot as plt
from scipy.spatial.transform import Rotation
from datetime import datetime
import xml.etree.ElementTree as ET
def load_raw_image_stack(filename):
# The resulting data is an array of (exposure, image-i, image-j),
# with image0i corresponding to y and image-j corresponding to x
# Note that the real-space scanning is done from the bottom right
# corner, first heading left (in x) then scanning up.
rawdata = np.fromfile(filename,dtype='<f4')
if len(rawdata) % (128*130) != 0:
raise IndexError('The raw data file doesn\'t seem to have the right number of values stored in it')
numshots = len(rawdata) // (128*130)
# One of the directions seems to be flipped
return rawdata.reshape(numshots,130,128)[:,:128,:][:,:,::-1].copy()
def load_metadata(filename):
return ET.parse(filename).getroot()
def get_scan_shape(metadata):
sp = metadata.find("scan_parameters[@mode='acquire']")
#print(sp)
# exit()
shape_x = int(sp.find('scan_resolution_x').text)
shape_y = int(sp.find('scan_resolution_y').text)
return [shape_x,shape_y]
def get_camera_length(metadata):
iomm = metadata.find('iom_measurements')
ncl = iomm.find('nominal_camera_length')
return float(ncl.text)
def get_scan_steps(metadata):
shape = get_scan_shape(metadata)
iomm = metadata.find('iom_measurements')
fov = iomm.find('full_scan_field_of_view')
scale = float(fov.find('scale_factor').text)
# Not sure why I need to divide by this scale factor or why it exists
# but this seems to produce the correct numbers
xfov = float(fov.find('x').text) / scale
yfov = float(fov.find('y').text) / scale
return np.array([xfov,yfov]) / np.array(shape)
def gen_scan_grid(shape, step):
ys, xs = np.mgrid[:shape[0],:shape[1]]
xs = xs * step[0]
ys = ys * step[1]
return np.stack((xs.ravel(),ys.ravel(),np.zeros(ys.ravel().shape))).transpose()
def get_electron_energy(metadata):
iomm = metadata.find('iom_measurements')
energy = float(iomm.find('high_voltage').text) / 1000 # to keV
return energy * 1.602e-16 # to Joules
h = 6.626e-34
c = 2.998e8
me = 9.109e-31
def calculate_wavelength(electron_energy):
return h*c / np.sqrt(electron_energy**2 + 2*c**2 * me * electron_energy)
def generate_detector_geometry(distance, pitches):
basis = np.array([[0,-pitches[1]],[-pitches[0],0],[0,0]])
return {'basis':basis, 'distance':distance}
def generate_dataset(translations, patterns, detector_geometry, electron_energy):
wavelength = calculate_wavelength(electron_energy)
print('Wavelength:',wavelength)
print('Pixel NA:',(-detector_geometry['basis'][0,1]/detector_geometry['distance']))
exit()
return Ptycho2DDataset(translations, patterns, wavelength=wavelength, detector_geometry=det_geo)
# Change this to allow for command-line introduction of the data folder
# I think that the data folder should be the first command-line arg, and
# be defaulted to the current folder
# Then, the image filename should by default be the only .raw file in the
# folder if there is exactly one. If none or more than one, throw an error
# Then, the metadata filename should be the only .xml file in the folder
# if there is exactly one, otherwise it should throw an error.
# Next, there should be some .csv file or similar containing the calibration
# of the scan size and pixel size. The program should give a report of how
# well the calibrated and naive values match. If no calibration is given,
# should indicate that it is using the naive values.
# Finally, the output filename should by default be the xml filename,
# and can be overriden by a clarg
#data_folder = '/media/Data Bank/Electron Ptycho MoS2/Session2/MoS2_Pty_session2/acquisition_1_50_nm_positive'
#image_filename = 'scan_x128_y128.raw'
#metadata_filename = 'acquisition_1_50_nm_positive.xml'
data_folder = '/media/Data Bank/Electron Ptycho MoS2/Session2/MoS2_Pty_session2/acquisition_1_20_nm_positive'
image_filename = 'scan_x128_y128.raw'
metadata_filename = 'acquisition_1_20_nm_positive.xml'
#data_folder = '/media/Data Bank/Electron Ptycho MoS2/Session2/MoS2_Pty_session2/acquisition_1_100nm_positive'
#image_filename = 'scan_x128_y128.raw'
#metadata_filename = 'acquisition_1_100nm_positive.xml'
#data_folder = '/media/Data Bank/Electron Ptycho MoS2/Lower_Convergence_angle/Smaller_sampling/acquisition_1_convergence_28mrad_14oMx_285mm'
#image_filename = 'scan_x128_y128.raw'
#metadata_filename = 'acquisition_1_convergence_28mrad_14oMx_285mm.xml'
#data_folder = '/media/Data Bank/Electron Ptycho MoS2/Lower_Convergence_angle/Larger_sampling/acquisition_1_14o5Mx_285mm_28mrad'
#image_filename = 'scan_x256_y256.raw'
#metadata_filename = 'acquisition_1_14o5Mx_285mm_28mrad.xml'
#data_folder = '/media/Data Bank/Electron Ptycho MoS2/Lower_Convergence_angle/Defocus_series/Defocus_positive/acquisition_2_60nm_positive_defocus_140Mx_28mrad_285mm'
#image_filename = 'scan_x128_y128.raw'
#metadata_filename = 'acquisition_2_60nm_positive_defocus_140Mx_28mrad_285mm.xml'
#data_folder = '/media/Data Bank/Electron Ptycho MoS2/Lower_Convergence_angle/Defocus_series/Defocus_negative/acquisition_2_14oMx_defocus_negative_50nm_28mrad_285mm'
#image_filename = 'scan_x128_y128.raw'
#metadata_filename = 'acquisition_2_14oMx_defocus_negative_50nm_28mrad_285mm.xml'
save_filename = 'Initial_CXI_Generation.cxi'
#data_folder = '/media/Data Bank/ptychography_firsttry/out_of_focus_58Mx_1ms_reso80x80_ss1'
#image_filename = 'scan_x80_y80.raw'
#save_filename = 'test_defocus_newcalibration.cxi'
#metadata_filename = 'out_of_focus_58Mx_1ms_reso80x80_ss1.xml'
#data_folder = '/media/Data Bank/ptychography_firsttry/acquisition_3'
#image_filename = 'scan_x80_y80.raw'
#save_filename = 'test_acq3_newcalibration.cxi'
#metadata_filename = 'acquisition_3.xml'
metadata = load_metadata(data_folder + '/' + metadata_filename)
scan_shape = get_scan_shape(metadata)
#scan_shape = 80
# These are reasonable initial guesses, until we get calibration data
#scan_step = 0.2e-10 #Angstrom, old value from manual measurement
scan_steps = get_scan_steps(metadata)
# This is something I can calculate from the detector length
# A good calibration is to assume that the pixel size is 0.2276 mm and
# the detector distance is equal to the nomninal camera length
camera_length = get_camera_length(metadata)
detector_distance = camera_length
# This gets the electron energy
electron_energy = get_electron_energy(metadata)
pixel_pitches = [500e-6,500e-6] # best match to Abinash's calibration
# Also feels right, even though the docs I find for the EMPAD shows 150um pixels
# These came from a calibration done by Xi
#pixel_pitches = [0.231e-3,0.231e-3] # best guess near length=0.230
# pixel_pitches = [0.2276e-3,0.2276e-3] # best overall average
# old manual calibration
#pixel_pitches = [150e-6,150e-6]
#detector_distance = 100e-3 # mm
#print([pp / detector_distance for pp in pixel_pitches])
#exit()
# Important question: Check which side the images fill in from
data = load_raw_image_stack(data_folder + '/' + image_filename)
#data[:,30:-30,30:-30] = 0 # For HAADF
scan_points = gen_scan_grid(scan_shape,scan_steps)
det_geo = generate_detector_geometry(detector_distance, pixel_pitches)
# The first image will be something like 20% larger than the rest...
#dataset = generate_dataset(scan_points, data, det_geo, electron_energy)
dataset = generate_dataset(scan_points[1:], data[1:], det_geo, electron_energy)
dataset.inspect(units='nm')
plt.show()
dataset.to_cxi(data_folder + '/' + save_filename)
docs/source/conf.py
+3 -3
View File
@@ -20,7 +20,7 @@ sys.path.insert(0, os.path.abspath('../..'))
# -- Project information -----------------------------------------------------
project = 'CDTools'
copyright = '2019, Abraham Levitan'
copyright = '2019-2021, Abraham Levitan'
author = 'Abraham Levitan'
# The short X.Y version
@@ -166,8 +166,8 @@ man_pages = [
# (source start file, target name, title, author,
# dir menu entry, description, category)
texinfo_documents = [
(master_doc, 'ADCD', 'ADCD Documentation',
author, 'ADCD', 'One line description of project.',
(master_doc, 'CDTools', 'CDTools Documentation',
author, 'CDTools', 'One line description of project.',
'Miscellaneous'),
]
docs/source/general.rst
+1 -1
View File
@@ -11,7 +11,7 @@ By necessity, CDTools operates using a mixture of numpy arrays and pytorch tenso
Functions that can accept either pytorch tensors or numpy arrays will have the type for the relevant inputs documented as "array", rather than "np.ndarray" or "torch.Tensor". In general, these functions will either accept all "array" type inputs as numpy arrays, or all as torch tensors. They will then return a result in a format matching that of the inputs. While many of these functions will work with mixed numpy/pytorch input, the output behavior of these functions is not, in general, defined for such a case, so it is heavily discouraged.
Finally, it is important to remember always that, since pytorch does not have complex number support, complex-valued tensors in pytorch are always represented by a tensor with a trailing dimension of length 2. All complex arithmetic operations are defined in :code:`tools.cmath`, but do not forget to use them!
Please note that adherance to the conventions above is a work in progress, and you may find improperly documented or implemented functions. Please file a bug report if you do!
Unit Conventions
docs/source/installation.rst
+4 -4
View File
@@ -8,7 +8,7 @@ The source code for CDTools is hosted on it's `MIT github page`_. Access to the
.. _`MIT github page`: https://github.mit.edu/Scattering/CDTools
It is recommended that you clone the repository, rather than just downloading the contents, as it remains under heavy development. Cloning the repository will allow you to get access to new updates.
It is recommended that you clone the repository, rather than just downloading the contents, as it remains under rapid development. Cloning the repository will allow you to get access to new updates.
Step 2: Install Dependencies
----------------------------
@@ -19,7 +19,7 @@ The dependencies for CDTools can be installed, if you are managing your environm
$ conda install --file conda_requirements.txt
There are two optional dependencies which are not installed via this procedure - the dependency sphinx-argparse for building the docs, and the pathlib2 module that provides python 2 compatibility. These can either be installed manually via conda-forge, or otherwise they will be installed automatically by pip during the final installation step if needed.
One optional dependencies is not installed via this procedure - the dependency sphinx-argparse for building the docs. Sphinx-argparse can either be installed manually via conda-forge, otherwise it will be installed automatically by pip during the final installation step if needed.
If you manage your environment with pip, all required packges should be installed automatically. The only thing to be aware of is that pytorch must be compiled with MKL support, and CUDA support if you would like to use the GPU. For this reason, using anaconda python is strongly recommended.
@@ -40,9 +40,9 @@ And has optional dependencies on
* `sphinx <https://www.sphinx-doc.org/>`_
* `sphinx-argparse <https://sphinx-argparse.readthedocs.io>`_
* `sphinx_rtd_theme <https://sphinx-rtd-theme.readthedocs.io/en/stable/>`_
* `pathlib2 <https://pypi.org/project/pathlib2/>`_
All of these can be installed via pip or conda. Finally, CDTools is written to be python 2.7+ compatible, but is only actively tested on python 3.
All of these can be installed via pip or conda. Finally, CDTools is NOT python 2 compatible.
**It is required that pytorch is built with MKL**, as that enables FFTs. Additionally, installing pytorch **with CUDA support** is recommended, if you intend to run any serious reconstructions with the package.
docs/source/intro.rst
+7 -8
View File
@@ -7,24 +7,23 @@ Introduction to CDTools
Introduction to CDTools
-----------------------
CDTools is a python library for ptychography and CDI reconstructions, using an Automatic Differentiation based approach.
CDTools is a python library for ptychography and CDI reconstructions, using an Automatic Differentiation (AD) based approach.
.. code-block:: python
# imports
import CDTools
from matplotlib import pyplot as plt
from CDTools.datasets import Ptycho2DDataset
from CDTools.models import SimplePtycho
# Load the file
dataset = Ptycho2DDataset.from_cxi('ptycho_data.cxi')
dataset = CDTools.datasets.Ptycho2DDataset.from_cxi('ptycho_data.cxi')
# Generate a model from the data
model = SimplePtycho.from_dataset(dataset)
model = CDTools.models.SimplePtycho.from_dataset(dataset)
# Run a reconstruction
for i, loss in enumerate(model.Adam_optimize(10, dataset)):
print(i, loss)
for loss in model.Adam_optimize(20, dataset):
print(model.report())
# And look at the results!
model.inspect(dataset)
@@ -32,7 +31,7 @@ CDTools is a python library for ptychography and CDI reconstructions, using an A
plt.show()
CDTools makes it simple to load and inspect data stored in .cxi files using python scripts. Several reconstruction models for common geometries are included "out of the box". For more advanced users, it includes a bunch of modular functions for AD ptychography, which can then be used right away from the same scripting framework.
CDTools makes it simple to load and inspect data stored in .cxi files using python scripts. Several reconstruction models for common geometries are included "out of the box". For more advanced users, it includes a modular functions for AD ptychography. These can can be used to construct new forward models.
The high-level interface to CDTools is built on a lower level "three-legged stool". This consists of tools to access stored data, tools to visualize data and reconstructions, and tools that implement basic operations relevant to coherent diffraction. All of these tools can be used directly alongside the high-level interface, when needed.
docs/source/tools/cmath.rst
-5
View File
@@ -1,5 +0,0 @@
Cmath
=====
.. automodule:: CDTools.tools.cmath
:members:
docs/source/tools/index.rst
-2
View File
@@ -4,7 +4,6 @@ Tools
.. toctree::
:maxdepth: 1
cmath
image_processing
data
initializers
@@ -14,4 +13,3 @@ Tools
losses
plotting
analysis
projectors
docs/source/tools/projectors.rst
-5
View File
@@ -1,5 +0,0 @@
Projectors
==========
.. automodule:: CDTools.tools.projectors
:members:
docs/source/tutorial.rst
+15 -18
View File
@@ -13,11 +13,9 @@ Our first step will be creating the file and filling out the boilerplate: All th
.. code-block:: python
from __future__ import division, print_function, absolute_import
import CDTools
from matplotlib import pyplot as plt
import pickle
from scipy import io
You can always import more libraries, like numpy, or pytorch, or pandas, or what have you, as needed. Next, we load the dataset and give it a look-over
@@ -46,9 +44,9 @@ We then try a basic Adam reconstruction with this model, with no changes to the
.. code-block:: python
for i, loss in enumerate(model.Adam_optimize(50, dataset)):
for loss in model.Adam_optimize(50, dataset):
model.inspect(dataset)
print(i,loss)
print(model.report())
model.compare(dataset)
plt.show()
@@ -58,7 +56,7 @@ It is worth noting here exactly how this code is working. The reconstruction met
In CDTools, every reconstruction method will return a generator. Whenever the generator is asked for the next item, it runs a single epoch of the reconstructionalgorithm, and then returns the average loss over that epoch as that next item. This allows the execution of the reconstruction algorithm to pause once every epoch, allowing some time for the user to run a small snippet of code to inspect how the reconstruction is coming along.
From the end user perspective, all this means is: follow the format above, or more generally put the :code:`model.Adam_optimize(n, dataset)` call anywhere that you would feel comfortable putting a call to :code:`range(n)` - list comprehensions, for loops, etc.
From the end user perspective, all this means is: follow the format above, or more generally put the :code:`model.Adam_optimize(n, dataset)` call anywhere that you would feel comfortable putting a call to :code:`range(n)` - list comprehensions, for loops, etc. In this case, we have called a function to plot out the current state of the reconstruction, and a function to print out the current loss and iteration time.
Once we run this, we can take a look at the result. What we see is pretty good, but we can see that there are some issues with the reconstruction near the edge, and the probe itself seems to be larger than the "stage" on which we're reconstructing it. So, we can make two tweaks to this code in response. First, we increase the oversampling ratio, which doubles the size of the stage (this often can cause other issues as well, but generally works well in situations like this where the probe is honestly too large.
@@ -94,8 +92,8 @@ Now we expect to get a nice reconstruction, so we can save the data. You can sav
.. code-block:: python
with open('example_reconstructions/lab_ptycho.pickle', 'wb') as f:
pickle.dump(model.save_results(dataset),f)
io.savemat('example_reconstructions/lab_ptycho.pickle',
model.save_results(dataset))
This is usually placed before the call to :code:`plt.show()`, to make sure that if the user manually exits the program once all the plot windows are opened, the data will still have been saved.
@@ -127,7 +125,7 @@ We can start with the basic skeleton for this file. In addition to our standard
.. code-block:: python
from __future__ import division, print_function, absolute_import
import numpy as np
import torch as t
from matplotlib import pyplot as plt
@@ -288,7 +286,6 @@ Once again, we start with the basic skeleton
.. code-block:: python
from __future__ import division, print_function, absolute_import
import numpy as np
import torch as t
from CDTools.models import CDIModel
@@ -323,10 +320,10 @@ It's important to note that there's not requirement for what the arguments to th
self.probe_basis = t.Tensor(probe_basis)
# We rescale the probe so it learns at the same rate as the object
self.probe_norm = t.max(tools.cmath.cabs(probe_guess.to(t.float32)))
self.probe = t.nn.Parameter(probe_guess.to(t.float32)
self.probe_norm = t.max(t.abs(probe_guess).to(t.float32))
self.probe = t.nn.Parameter(probe_guess.to(t.complex64)
/ self.probe_norm)
self.obj = t.nn.Parameter(obj_guess.to(t.float32))
self.obj = t.nn.Parameter(obj_guess.to(t.complex64))
Here, we chose to define the model based on a basis matrix describing the probe array, an initial guess at the probe, and an initial object. In addition, an optional offset for the translations is included.
@@ -375,7 +372,7 @@ To initialize the object from a dataset, we need to start by extracting the rele
probe_shape,
det_slice)
obj = t.ones(obj_size+(2,))
obj = t.ones(obj_size, dtype=t.complex64)
return cls(probe_basis, probe, obj, min_translation=min_translation)
@@ -472,9 +469,9 @@ At the moment, there is no consistent way to save out the results across the boa
.. code-block:: python
def save_results(self):
probe = tools.cmath.torch_to_complex(self.probe.detach().cpu())
probe = self.probe.detach().cpu().numpy()
probe = probe * self.probe_norm.detach().cpu().numpy()
obj = tools.cmath.torch_to_complex(self.obj.detach().cpu())
obj = self.obj.detach().cpu().numpy()
return {'probe':probe,'obj':obj}
@@ -499,9 +496,9 @@ We can test this model with a simple script, shown below. By filling in the back
model.to(device='cuda')
dataset.get_as(device='cuda')
for i, loss in enumerate(model.Adam_optimize(100, dataset)):
for loss in model.Adam_optimize(100, dataset):
model.inspect(dataset)
print(i,loss)
print(model.report())
model.compare(dataset)
plt.show()
examples/MIT_BNL_logo.py
+7 -9
View File
@@ -1,5 +1,3 @@
from __future__ import division, print_function, absolute_import
import CDTools
from matplotlib import pyplot as plt
@@ -23,7 +21,7 @@ dataset.mask = dataset.mask[70:-70,70:-70]
model = CDTools.models.FancyPtycho.from_dataset(dataset,
translation_scale = 4,
n_modes=2,
propagation_distance=-73e-6)
propagation_distance=73e-6)
# Move to the GPU
model.to(device='cuda')
@@ -32,20 +30,20 @@ dataset.get_as(device='cuda')
# We turn off position reconstruction for the first phase
model.translation_offsets.requires_grad = False
for i, loss in enumerate(model.Adam_optimize(10, dataset, batch_size=15)):
for loss in model.Adam_optimize(10, dataset, batch_size=15):
model.inspect(dataset)
print(i,loss)
print(model.report())
# And we turn it on for the second phase
model.translation_offsets.requires_grad = True
for i, loss in enumerate(model.Adam_optimize(20, dataset, batch_size=15)):
for loss in model.Adam_optimize(20, dataset, batch_size=15):
model.inspect(dataset)
print(i,loss)
print(model.report())
# The third phase lowers the rate further
for i, loss in enumerate(model.Adam_optimize(10, dataset, batch_size=15, lr=0.0005)):
for loss in model.Adam_optimize(10, dataset, batch_size=15, lr=0.0005):
model.inspect(dataset)
print(i,loss)
print(model.report())
model.inspect(dataset)
examples/basic_ptycho_dataset.py
-1
View File
@@ -1,4 +1,3 @@
from __future__ import division, print_function, absolute_import
import numpy as np
import torch as t
from matplotlib import pyplot as plt
examples/debug_3d.py
-17
View File
@@ -1,17 +0,0 @@
from __future__ import division, print_function, absolute_import
import CDTools
from matplotlib import pyplot as plt
# First, we load an example dataset from a .cxi file
filename = '/media/Data Bank/APS_HXN_07_19/CXIs/scan_94513_cxi.h5'
#filename = 'data/scan_94361_cxi.h5'
dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename)
# Next, we create a ptychography model from the dataset
model = CDTools.models.Simple3DBPP.from_dataset(dataset)
#import pdb; pdb.set_trace()
model.inspect(dataset)
plt.show()
# # Now, we run a short reconstruction from the dataset!
# for i, loss in enumerate(model.Adam_optimize(1, dataset)):
# print(i, loss)
examples/ensemble_analysis.py
+2 -12
View File
@@ -1,21 +1,12 @@
from __future__ import division, print_function, absolute_import
from matplotlib import pyplot as plt
import pickle
from scipy import io
from CDTools.tools import plotting
from CDTools.tools import analysis
with open('example_reconstructions/gold_balls_ensemble.pickle', 'rb') as f:
dataset = pickle.load(f)
# This converts from a list of dictionaries to a dictionary of lists
# It's safe to assume that all elements have the same set of keys
if type(dataset) == type([]):
dataset = {key: [element[key] for element in dataset]
for key in dataset[0]}
dataset = io.loadmat('example_reconstructions/gold_balls_ensemble.mat')
# Now we synthesize an average reconstruction
@@ -25,7 +16,6 @@ synth_probe, synth_obj, aligned_objs = analysis.synthesize_reconstructions(
# And then we calculate the consistency PRTF from this
freqs, prtf = analysis.calc_consistency_prtf(synth_obj, aligned_objs, dataset['basis'][0])
# Plot the first mode in detail
plotting.plot_phase(synth_probe[0],basis=dataset['basis'][0])
plotting.plot_amplitude(synth_probe[0],basis=dataset['basis'][0])
examples/ensemble_reconstruction.py
+19 -12
View File
@@ -1,16 +1,16 @@
from __future__ import division, print_function, absolute_import
import CDTools
import pickle
import numpy as np
from scipy import io
# Load the data
filename = 'example_data/AuBalls_700ms_30nmStep_3_6SS_filter.cxi'
dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename)
results = []
n = 25
for idx in range(25):
print('Starting Reconstruction', idx)
for idx in range(n):
print('Starting Reconstruction', idx+1, 'of',n)
# Create a new model each time
model = CDTools.models.FancyPtycho.from_dataset(dataset,n_modes=3,
@@ -21,14 +21,21 @@ for idx in range(25):
dataset.get_as(device='cuda')
# Run the reconstruction
for i, loss in enumerate(model.Adam_optimize(30, dataset, batch_size=100)):
print(i,loss)
for loss in model.Adam_optimize(30, dataset, batch_size=100):
print(model.report(), end='\r')
# And add the results to the ensemble
# Print a summary that won't be overwritten by the next line
print('Finished:',model.report())
# And add the results to the ensemble dictionary
results.append(model.save_results(dataset))
# This converts from a list of dictionaries to a dictionary of lists
# It's safe to assume that all elements have the same set of keys
# After this, e.g. dataset['probe'] will return a list of all probes.
results = {key: np.array([element[key] for element in results])
for key in results[0].keys()}
print(results['probe'].shape)
# Save out the ensemble
with open('example_reconstructions/gold_balls_ensemble.pickle', 'wb') as f:
pickle.dump(results,f)
io.savemat('example_reconstructions/gold_balls_ensemble.mat', results)
examples/example_data/Optical_ptycho.pickle
BIN
View File
Binary file not shown.
examples/example_reconstructions/gold_balls.mat
BIN
View File
Binary file not shown.
examples/example_reconstructions/gold_balls.pickle
BIN
View File
Binary file not shown.
examples/example_reconstructions/gold_balls_ensemble.mat
BIN
View File
Binary file not shown.
examples/example_reconstructions/gold_balls_ensemble.pickle
BIN
View File
Binary file not shown.
examples/example_reconstructions/lab_ptycho.mat
BIN
View File
Binary file not shown.
examples/example_reconstructions/lab_ptycho.pickle
BIN
View File
Binary file not shown.
examples/gold_ball_ptycho.py
+11 -10
View File
@@ -1,30 +1,31 @@
from __future__ import division, print_function, absolute_import
import CDTools
from matplotlib import pyplot as plt
import pickle
import time
from scipy import io
# First, we load an example dataset from a .cxi file
filename = 'example_data/AuBalls_700ms_30nmStep_3_6SS_filter.cxi'
dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename)
# Next, we create a ptychography model from the dataset
# Note that we explicitly as for two incoherent probe modes
# Note that we explicitly ask for two incoherent probe modes
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')
for i, loss in enumerate(model.Adam_optimize(30, dataset, batch_size=100)):
# Now, we run a short reconstruction from the dataset
for loss in model.Adam_optimize(100, dataset, batch_size=50, schedule=True):
# And we liveplot the updates to the model as they happen
print(i,loss)
print(model.report())
model.inspect(dataset)
# And we save the reconstruction out to a file
with open('example_reconstructions/gold_balls.pickle', 'wb') as f:
pickle.dump(model.save_results(dataset),f)
# This orthogonalizes the incoherent probe modes
model.tidy_probes()
# And we save out the results as a .mat file
io.savemat('example_reconstructions/gold_balls.mat',
model.save_results(dataset))
# Finally, we plot the results
model.inspect(dataset)
examples/inspect_dataset.py
-2
View File
@@ -1,5 +1,3 @@
from __future__ import division, print_function, absolute_import
import CDTools
from matplotlib import pyplot as plt
examples/lab_bragg_2d_ptycho.py
+13 -17
View File
@@ -1,36 +1,32 @@
from __future__ import division, print_function, absolute_import
import CDTools
from matplotlib import pyplot as plt
import pickle
from scipy import io
#
# This is the testing file for the Bragg ptychography code that includes
# corrections for probe propagation and for off-axis far field diffraction
# I'm starting to build it up using the lab ptychography forward data so I
# can initially test that it works on "softball" data before moving to the
# actual Bragg geometry data
# This is the testing dataset for the Bragg ptychography code, that was
# collected with our optical setup using a low NA zone plate illuminating
# a reflective diffraction grating with writing on it.
#
#filename = 'example_data/lab_ptycho_data.cxi'
# This file is too large to be distributed via Github.
# Please contact Abe Levitan (alevitan@mit) if you would like access
filename = '/media/Data Bank/Lab Ptycho/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)#propagate_probe=False)#, n_modes=8)
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
for i, loss in enumerate(model.Adam_optimize(100, dataset)):
for loss in model.Adam_optimize(100, dataset):
model.inspect(dataset)
print(i,loss)
#with open('example_reconstructions/lab_bragg_2d_ptycho.pickle', 'wb') as f:
# pickle.dump(model.save_results(dataset),f)
print(model.report())
#model.compare(dataset)
io.savemat('example_reconstructions/lab_bragg_2d_ptycho.mat',
model.save_results(dataset))
model.compare(dataset)
plt.show()
examples/lab_ptycho_data.py
+15 -11
View File
@@ -1,28 +1,32 @@
from __future__ import division, print_function, absolute_import
import CDTools
from matplotlib import pyplot as plt
import pickle
from scipy import io
# A simple dataset collected from our optical setup
filename = 'example_data/lab_ptycho_data.cxi'
dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename)
# dataset.inspect()
# plt.show()
dataset.inspect(units='mm')
plt.show()
model = CDTools.models.FancyPtycho.from_dataset(dataset, oversampling=2,
probe_support_radius=90)
probe_support_radius=90,
n_modes=2,
dm_rank=0,units='mm')
model.to(device='cuda')
dataset.get_as(device='cuda')
model.translation_offsets.requires_grad = False
for i, loss in enumerate(model.Adam_optimize(50, dataset)):
for loss in model.Adam_optimize(50, dataset):
model.inspect(dataset)
print(i,loss)
with open('example_reconstructions/lab_ptycho.pickle', 'wb') as f:
pickle.dump(model.save_results(dataset),f)
print(model.report())
model.tidy_probes()
io.savemat('example_reconstructions/lab_ptycho.mat',
model.save_results(dataset))
model.compare(dataset)
plt.show()
examples/simple_ptycho.py
+3 -5
View File
@@ -1,18 +1,16 @@
from __future__ import division, print_function, absolute_import
import CDTools
from matplotlib import pyplot as plt
# First, we load an example dataset from a .cxi file
filename = 'example_data/lab_ptycho_data.cxi'
dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename)
plt.ion()
# Next, we create a ptychography model from the dataset
model = CDTools.models.SimplePtycho.from_dataset(dataset)
# Now, we run a short reconstruction from the dataset!
for i, loss in enumerate(model.Adam_optimize(10, dataset)):
print(i, loss)
for loss in model.Adam_optimize(20, dataset):
print(model.report())
# Finally, we plot the results
model.inspect(dataset)
examples/simple_ptycho_model.py
-1
View File
@@ -1,4 +1,3 @@
from __future__ import division, print_function, absolute_import
import numpy as np
import torch as t
from CDTools.models import CDIModel
examples/specular_ptycho.py
-2
View File
@@ -1,5 +1,3 @@
from __future__ import division, print_function, absolute_import
import CDTools
from matplotlib import pyplot as plt
import numpy as np
examples/transmission_RPI.py
+10 -16
View File
@@ -1,15 +1,10 @@
from __future__ import division, print_function, absolute_import
import CDTools
from matplotlib import pyplot as plt
import pickle
from torch.utils.data import Subset
from matplotlib import pyplot as plt
# First, we load an example dataset from a .cxi file
ss_filename = 'example_data/Optical_Data_ss.cxi'
#with open('example_data/Optical_ptycho.pickle', 'rb') as f:
with open('example_data/Optical_ptycho_incoherent.pickle', 'rb') as f:
ptycho_results = pickle.load(f)
@@ -19,9 +14,9 @@ background = ptycho_results['background']
dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(ss_filename)
# Next, we create a ptychography model from the dataset
# Next, we create an RPI model from the dataset
# Note that we explicitly as for two incoherent probe modes
model = CDTools.models.RPI.from_dataset(dataset, probe, [900,900],
model = CDTools.models.RPI.from_dataset(dataset, probe, [500,500],
background=background, n_modes=2,
initialization='random')
@@ -33,20 +28,19 @@ 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 i, loss in enumerate(model.LBFGS_optimize(30, dataset, lr=0.4, regularization_factor=[0.05,0.05])):#0.1)):
#model.inspect(dataset)
print(i,loss)
#model.inspect(dataset)
for loss in model.LBFGS_optimize(30, dataset, lr=0.4, regularization_factor=[0.05,0.05]):
model.inspect(dataset)
print(model.report())
# Now we use the regularizer to damp all but the top modes
for i, loss in enumerate(model.LBFGS_optimize(50, dataset, lr=0.4, regularization_factor=[0.001,0.1])):
for loss in model.LBFGS_optimize(50, dataset, lr=0.4, regularization_factor=[0.001,0.1]):
#model.inspect(dataset)
print(i,loss)
print(model.report())
# Save results to a python dictionary
results = model.save_results()
# Finally, we plot the results
model.inspect(dataset)
model.compare(dataset)
examples/unified_modes.py
+14 -16
View File
@@ -1,39 +1,37 @@
from __future__ import division, print_function, absolute_import
import CDTools
from matplotlib import pyplot as plt
import pickle
from scipy import io
filename = 'example_data/lab_ptycho_data.cxi'
dataset = CDTools.datasets.Ptycho2DDataset.from_cxi(filename)
# dataset.inspect()
# plt.show()
#dataset.inspect(units='mm')
#plt.show()
# dm_rank=-1 tells it to use a full-rank unified mode approximation
model = CDTools.models.FancyPtycho.from_dataset(dataset, oversampling=1,n_modes=3, dm_rank=-1)
#model = CDTools.models.UnifiedModePtycho.from_dataset(dataset, oversampling=2,n_modes=3)#, probe_support_radius=90)
model = CDTools.models.UnifiedModePtycho2.from_dataset(dataset, oversampling=1,n_modes=3)#, probe_support_radius=90)
#model = CDTools.models.FancyPtycho.from_dataset(dataset, oversampling=1)
model.to(device='cuda')
dataset.get_as(device='cuda')
model.translation_offsets.requires_grad = False
for i, loss in enumerate(model.Adam_optimize(100, dataset)):
for loss in model.Adam_optimize(200, dataset):
model.inspect(dataset)
print(i,loss)
print(model.report())
model.tidy_probes()
for i, loss in enumerate(model.Adam_optimize(20, dataset, lr=0.0001)):
for loss in model.Adam_optimize(200, dataset, lr=0.0001):
model.inspect(dataset)
print(i,loss)
print(model.report())
model.tidy_probes()
model.tidy_probes(normalize=True)
model.inspect(dataset)
with open('example_reconstructions/unified_modes.pickle', 'wb') as f:
pickle.dump(model.save_results(dataset),f)
io.savemat('example_reconstructions/unified_modes.mat',
model.save_results(dataset))
model.compare(dataset)
plt.show()
kiara_debugging_code/debugging_mod.py
-112
View File
@@ -1,112 +0,0 @@
import numpy as np
from matplotlib import pyplot as plt
import imageio
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
class ImageSeries:
def __init__(self, images, crop,fps=None):
self.images = images #list of images (numpy arrays), saved itself as a numpy array
self.L = len(self.images)
self.x1, self.x2, self.y1, self.y2 = crop #(x1,y1) and (x2,y2) are cropping coords
self.crimages = []
for i in range(self.L):
self.crimages.append(self.images[i][self.x1:self.x2, self.y1:self.y2])
self.crimages = np.array(self.crimages)
self.fps = fps
def show_frame(self,t): #t = time to show
if 0 <= t <= self.L:
plt.imshow(self.crimages[t],interpolation="none")
plt.show()
return
def save_vid(self, filename, secspersec):
imageio.mimwrite(filename, self.crimages, fps=self.fps*secspersec)
def waterfall(self,x,y):
if x == None:
#plot row y=y over time
return self.crimages[:,:,y]
elif y == None:
return np.transpose(self.crimages[:,x,:])
def plot_waterfall(self,x,y,tint=None):
waterfall = self.waterfall(x,y)
if x == None:
#plot col y over time
fig = plt.figure()
W = fig.add_subplot(111)
W.imshow(waterfall,interpolation="none")
if tint:
plt.yticks(np.arange(0,self.L,self.fps*tint),np.arange(0,self.L/self.fps,tint))
W.set_title("waterfall plot of row y = " + str(y))
W.set_ylabel("time [s]")
plt.show()
return
elif y == None:
#plot col x=x over time
fig = plt.figure()
W = fig.add_subplot(111)
W.imshow(waterfall,interpolation="none")
if tint:
plt.xticks(np.arange(0,self.L,self.fps*tint),np.arange(0,self.L/self.fps,tint))
W.set_title("waterfall plot of col x = " + str(x))
W.set_xlabel("time [s]")
plt.show()
return
return
def Ipixel(self,pixel):
px,py = pixel
return self.crimages[:,px,py]
def plot_Ipixel(self,pixel,description=""):
I = self.Ipixel(pixel)
fig = plt.figure()
f1 = fig.add_subplot(111)
f1.set_title("I(t) for pixel" + str(pixel) + " (" + description + ")")
f1.set_xlabel("time [s]")
f1.plot(np.arange(0,self.L/self.fps,1/self.fps),I)
plt.show()
return
def fftIpixel(self,pixel):
I = self.Ipixel(pixel)
IfreqA = np.fft.fft(I)/self.L
Ifreq = np.fft.fftfreq(self.L,d=(1/self.fps))
return(Ifreq, IfreqA)
def plot_fftIpixel(self,pixel):
Ifreq, IfreqA = self.fftIpixel(pixel)
fig = plt.figure()
f1 = fig.add_subplot(111)
f1.set_title("FFT for pixel" + str(pixel))
f1.set_xlabel("frequency [1/s]")
f1.plot(Ifreq, abs(IfreqA))
plt.show()
return
def g2(self,pixel):
I = self.Ipixel(pixel)
g2 = []
avgsq = np.mean(I)**2
for tau in range(len(I)-1):
if tau == 0:
dotp = 0
for t in range(len(I)):
dotp += I[t]*I[t]
g2.append(dotp/(len(I)*avgsq) )
elif tau != 0:
dotp = 0
for t in range(len(I)-tau):
dotp += I[t]*I[t+tau]
g2.append( dotp / (len(I[:-tau])*avgsq))
g2 = np.array(g2)
return g2
kiara_debugging_code/vid_read.py
-50
View File
@@ -1,50 +0,0 @@
import numpy as np
from matplotlib import pyplot as plt
import imageio
from PIL import Image, ImageSequence
from debugging_mod import ImageSeries
vid = Image.open('data_10_4/kiara_20fps_redlaser_vid.tif')
vidarray = []
for i, page in enumerate(ImageSequence.Iterator(vid)):
pg = np.array(page)
vidarray.append(pg)
vidarray = np.array(vidarray)
RedLaserExp = ImageSeries(vidarray, (430,606,590,766),fps=20) #cropping x1:x2, y1:y2
#with np.load('data_9_28/kiara_data_300sec_green') as data:
# GreenLaserExp = ImageSeries(data['arr_0'], (500,676,580,756), fps=5)
#crop to: x1=500, x2=676, y1=580, y2=756
#++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
#Note: human eye can see at c. 150 fps
#inner top right corner of disk (120,70)
#inner top left corner of disk (66,63)
#inner bottom left corner of disk (66,107)
#inner bottom (87,118)
#inner right (125, 86)
#pixel intensity plot: ---------------------------------------
#RedLaserExp.plot_Ipixel((125,86),description="inner right")
#RedLaserExp.plot_Ipixel((87,118),description="inner bottom")
#RedLaserExp.plot_Ipixel((66,63),description="inner top left")
#pixel fft plot: --------------------------------------------
#RedLaserExp.plot_fftIpixel((125,86))
#RedLaserExp.plot_fftIpixel((87,118))
#RedLaserExp.plot_fftIpixel((66,63))
#waterfall plot for row: -------------------------------------
#RedLaserExp.plot_waterfall(None,118,tint=1)
#waterfall plot for col: -------------------------------------
#RedLaserExp.plot_waterfall(66,None,tint=1)
#save a video: -----------------------------------------------
#RedLaserExp.save_vid("redlaser_10-4_20fps_1x.mp4",1)
setup.py
+7 -8
View File
@@ -5,8 +5,9 @@ with open("README.md", "r") as fh:
setuptools.setup(
name="CDTools",
version="0.0.1",
author="Abe Levitan, Madelyn Cain",
version="0.1.1",
python_requires='>3.4', # we use pathlib, introduced in 3.4
author="Abe Levitan",
author_email="alevitan@mit.edu",
description="Coherent Diffraction Tools",
long_description=long_description,
@@ -15,15 +16,13 @@ setuptools.setup(
install_requires=[
"numpy>=1.0",
"scipy>=1.0",
"matplotlib>=2.0",
"matplotlib>=2.0", # Matplotlib 2.0 introduces better colormaps and no I'm not sorry
"python-dateutil",
"torch>=1.3.0", #1.3.0 adds the align_corners option for grid_sample which is used for high NA far-field propagation. 1.2.0 introduced boolean tensors in a breaking way, we use the boolean tensors here for masking.
"h5py>=2.1",
"pathlib2 ; python_version<'3.4'"],
"torch>=1.9.0", #1.9.0 implements support for autograd on indexed complex tensors, which we need in order to use complex tensors in the forward models
"h5py>=2.1"],
extras_require={
'tests': ["pytest"],
'docs': ["sphinx","sphinx-argparse","sphinx_rtd_theme"],
":python_version<'3.4'": ["pathlib2"],
'docs': ["sphinx","sphinx-argparse","sphinx_rtd_theme"]
},
packages=setuptools.find_packages(),
classifiers=[
tests/conftest.py
+2 -3
View File
@@ -1,4 +1,3 @@
from __future__ import division, print_function, absolute_import
import numpy as np
import h5py
import pytest
@@ -104,7 +103,7 @@ def ptycho_cxi_1():
# Remember the format for the CXI file differs from the format used
# internally
mask = np.zeros((256,256)).astype(np.int32)
expected['mask'] = np.ones((256,256)).astype(np.bool)
expected['mask'] = np.ones((256,256)).astype(bool)
d1f.create_dataset('mask',data=mask)
# Create an initial background
@@ -272,7 +271,7 @@ def ptycho_cxi_3():
# Remember the format for the CXI file differs from the format used
# internally
mask = np.ones((256,256)).astype(np.uint32) * 0x00001000
expected['mask'] = np.ones((256,256)).astype(np.bool)
expected['mask'] = np.ones((256,256)).astype(bool)
d1f.create_dataset('mask',data=mask)
expected['dark'] = None
tests/test_datasets.py
+1 -1
View File
@@ -28,7 +28,7 @@ def test_CDataset_init():
dataset = CDataset(entry_info, sample_info,
wavelength, detector_geometry, mask)
assert t.all(t.eq(dataset.mask,t.tensor(mask.astype(np.bool))))
assert t.all(t.eq(dataset.mask,t.tensor(mask.astype(bool))))
assert dataset.entry_info == entry_info
assert dataset.sample_info == sample_info
assert dataset.wavelength == wavelength
tests/tools/test_analysis.py
+18 -21
View File
@@ -6,7 +6,7 @@ from scipy import fftpack as ffts
import torch as t
from itertools import combinations
from CDTools.tools import analysis, cmath, initializers
from CDTools.tools import analysis, initializers
def test_orthogonalize_probes():
@@ -30,7 +30,7 @@ def test_orthogonalize_probes():
ortho_probes = analysis.orthogonalize_probes(probes)
# test that it also works on torch tensors
ortho_probes_t = cmath.torch_to_complex(analysis.orthogonalize_probes(cmath.complex_to_torch(probes)))
ortho_probes_t = analysis.orthogonalize_probes(t.as_tensor(probes)).numpy()
# This tests for orthogonality
for p1,p2 in combinations(ortho_probes,2):
@@ -80,13 +80,12 @@ def test_standardize():
# Start by making a probe and object that should meet the standardization
# conditions
probe = initializers.gaussian((230,240),(20,20),curvature=(0.01,0.01))
probe = cmath.torch_to_complex(probe)
probe = initializers.gaussian((230,240),(20,20),curvature=(0.01,0.01)).numpy()
probe = probe * np.sqrt(len(probe.ravel()) / np.sum(np.abs(probe)**2))
probe = probe * np.exp(-1j * np.angle(np.sum(probe)))
assert np.isclose(1, np.sum(np.abs(probe)**2)/ len(probe.ravel()))
assert np.isclose(0,np.angle(np.sum(probe)))
assert np.angle(np.sum(probe)) < 1e-7
obj = 30 * np.random.rand(230,240) * np.exp(1j * (np.random.rand(230,240) - 0.5))
obj_slice = np.s_[(obj.shape[0]//8)*3:(obj.shape[0]//8)*5,
@@ -105,9 +104,9 @@ def test_standardize():
assert np.allclose(obj, s_obj)
# Test that it works on torch tensors
s_probe, s_obj = analysis.standardize(cmath.complex_to_torch(test_probe).to(t.float32), cmath.complex_to_torch(test_obj).to(t.float32))
s_probe = cmath.torch_to_complex(s_probe)
s_obj = cmath.torch_to_complex(s_obj)
s_probe, s_obj = analysis.standardize(t.as_tensor(test_probe,dtype=t.complex64), t.as_tensor(test_obj,dtype=t.complex64))
s_probe = s_probe.numpy()
s_obj = s_obj.numpy()
assert np.allclose(probe, s_probe)
assert np.allclose(obj, s_obj)
@@ -146,7 +145,6 @@ def test_standardize():
from matplotlib import pyplot as plt
def test_synthesize_reconstructions():
# I can only really test for a lack of failures, so I think my plan
# will be to create a dataset that just needs to be added and see that
@@ -154,13 +152,12 @@ def test_synthesize_reconstructions():
# Start by making a probe and object that should meet the standardization
# conditions
probe = initializers.gaussian((230,240),(20,20),curvature=(0.01,0.01))
probe = cmath.torch_to_complex(probe)
probe = initializers.gaussian((230,240),(20,20),curvature=(0.01,0.01)).numpy()
probe = probe * np.sqrt(len(probe.ravel()) / np.sum(np.abs(probe)**2))
probe = probe * np.exp(-1j * np.angle(np.sum(probe)))
assert np.isclose(1, np.sum(np.abs(probe)**2)/ len(probe.ravel()))
assert np.isclose(0,np.angle(np.sum(probe)))
assert np.abs(np.angle(np.sum(probe))) < 1e-7
obj = 30 * np.random.rand(230,240) * np.exp(1j * (np.random.rand(230,240) - 0.5))
obj_slice = np.s_[(obj.shape[0]//8)*3:(obj.shape[0]//8)*5,
@@ -203,14 +200,14 @@ def test_calc_consistency_prtf():
assert np.allclose(prtf, 0.7)
# Check that it also works with torch input
t_synth_obj = cmath.complex_to_torch(synth_obj)
t_obj_stack = [cmath.complex_to_torch(obj) for obj in obj_stack]
t_synth_obj = t.as_tensor(synth_obj)
t_obj_stack = [t.as_tensor(obj) for obj in obj_stack]
freqs, prtf = analysis.calc_consistency_prtf(t_synth_obj, t_obj_stack, basis, nbins=30)
assert np.allclose(prtf.numpy(), 0.7)
# And also when the basis is in torch
t_synth_obj = cmath.complex_to_torch(synth_obj)
t_obj_stack = [cmath.complex_to_torch(obj) for obj in obj_stack]
t_synth_obj = t.as_tensor(synth_obj)
t_obj_stack = [t.as_tensor(obj) for obj in obj_stack]
freqs, prtf = analysis.calc_consistency_prtf(t_synth_obj, t_obj_stack, t.Tensor(basis), nbins=30)
assert np.allclose(prtf.numpy(), 0.7)
@@ -238,11 +235,11 @@ def test_calc_deconvolved_cross_correlation():
assert np.allclose(test_cor, np_cor)
# test with pytorch inputs
obj1_t = cmath.complex_to_torch(obj1)
obj2_t = cmath.complex_to_torch(obj2)
obj1_t = t.as_tensor(obj1)
obj2_t = t.as_tensor(obj2)
test_cor_t = analysis.calc_deconvolved_cross_correlation(obj1_t,obj2_t, im_slice=np.s_[:,:])
assert np.allclose(cmath.torch_to_complex(test_cor_t), np_cor)
assert np.allclose(test_cor_t.numpy(), np_cor)
@@ -294,8 +291,8 @@ def test_calc_frc():
assert np.allclose(threshold, test_threshold)
# try again with complex
obj1_torch = cmath.complex_to_torch(obj1)
obj2_torch = cmath.complex_to_torch(obj2)
obj1_torch = t.as_tensor(obj1)
obj2_torch = t.as_tensor(obj2)
basis_torch = t.tensor(basis)
test_bins_t, test_frc_t, test_threshold_t = analysis.calc_frc(obj1_torch,
tests/tools/test_cmath.py
-135
View File
@@ -1,135 +0,0 @@
from __future__ import division, print_function, absolute_import
from CDTools.tools import cmath
import numpy as np
import torch as t
from scipy.fftpack import fftshift, ifftshift
def test_complex_to_torch():
arr = np.random.rand(100,4) + 1j * np.random.rand(100,4)
tensor = cmath.complex_to_torch(arr)
assert np.allclose(tensor[:,:,0].numpy(),np.real(arr))
assert np.allclose(tensor[:,:,1].numpy(),np.imag(arr))
def test_torch_to_complex():
tensor = t.rand(100,4,2)
arr = cmath.torch_to_complex(tensor)
assert np.allclose(tensor[:,:,0].numpy(),np.real(arr))
assert np.allclose(tensor[:,:,1].numpy(),np.imag(arr))
def test_cabssq():
arr = np.random.rand(56,23,2) + 1j * np.random.rand(56,23,2)
cabssq = cmath.cabssq(cmath.complex_to_torch(arr))
assert np.allclose(cabssq.numpy(),np.abs(arr)**2)
def test_cabs():
arr = np.random.rand(2,3,4,5) + 1j * np.random.rand(2,3,4,5)
cabs = cmath.cabs(cmath.complex_to_torch(arr))
assert np.allclose(cabs.numpy(),np.abs(arr))
def test_cconj():
arr = np.random.rand(50) + 1j * np.random.rand(50)
cconj = cmath.cconj(cmath.complex_to_torch(arr))
assert np.allclose(cmath.torch_to_complex(cconj), np.conj(arr))
def test_cmult():
arr1 = np.random.rand(50) + 1j * np.random.rand(50)
arr2 = np.random.rand(50) + 1j * np.random.rand(50)
mult = cmath.cmult(cmath.complex_to_torch(arr1),
cmath.complex_to_torch(arr2))
assert np.allclose(cmath.torch_to_complex(mult),arr1*arr2)
def test_cdiv():
arr1 = np.random.rand(50) + 1j * np.random.rand(50)
arr2 = np.random.rand(50) + 1j * np.random.rand(50)
div = cmath.cdiv(cmath.complex_to_torch(arr1),
cmath.complex_to_torch(arr2))
assert np.allclose(cmath.torch_to_complex(div),arr1 / arr2)
def test_cphase():
arr = np.random.rand(50) + 1j * np.random.rand(50)
cphase = cmath.cphase(cmath.complex_to_torch(arr))
assert np.allclose(cphase.numpy(), np.angle(arr))
def test_scalars():
arr1 = np.random.rand(1) + 1j * np.random.rand(1)
arr2 = np.random.rand(1) + 1j * np.random.rand(1)
div = cmath.cdiv(cmath.complex_to_torch(arr1),
cmath.complex_to_torch(arr2))
assert np.allclose(cmath.torch_to_complex(div),arr1 / arr2)
def test_fftshift():
#1D, even
arr = np.random.rand(300) + 1j * np.random.rand(300)
shifted = cmath.fftshift(cmath.complex_to_torch(arr))
assert np.allclose(fftshift(arr),
cmath.torch_to_complex(shifted))
#1D, odd
arr = np.random.rand(301) + 1j * np.random.rand(301)
shifted = cmath.fftshift(cmath.complex_to_torch(arr))
assert np.allclose(fftshift(arr),
cmath.torch_to_complex(shifted))
#2D
arr = np.random.rand(20,21) + 1j * np.random.rand(20,21)
shifted = cmath.fftshift(cmath.complex_to_torch(arr))
assert np.allclose(fftshift(arr),
cmath.torch_to_complex(shifted))
#3D
arr = np.random.rand(15,16,17) + 1j * np.random.rand(15,16,17)
shifted = cmath.fftshift(cmath.complex_to_torch(arr))
assert np.allclose(fftshift(arr),
cmath.torch_to_complex(shifted))
#3D, choosing specific axes
arr = np.random.rand(15,16,17) + 1j * np.random.rand(15,16,17)
shifted = cmath.fftshift(cmath.complex_to_torch(arr),dims=(0,1))
assert np.allclose(fftshift(arr,axes=(0,1)),
cmath.torch_to_complex(shifted))
def test_ifftshift():
#1D, even
arr = np.random.rand(300) + 1j * np.random.rand(300)
shifted = cmath.ifftshift(cmath.complex_to_torch(arr))
assert np.allclose(ifftshift(arr),
cmath.torch_to_complex(shifted))
#1D, odd
arr = np.random.rand(301) + 1j * np.random.rand(301)
shifted = cmath.ifftshift(cmath.complex_to_torch(arr))
assert np.allclose(ifftshift(arr),
cmath.torch_to_complex(shifted))
#2D
arr = np.random.rand(20,21) + 1j * np.random.rand(20,21)
shifted = cmath.ifftshift(cmath.complex_to_torch(arr))
assert np.allclose(ifftshift(arr),
cmath.torch_to_complex(shifted))
#3D
arr = np.random.rand(15,16,17) + 1j * np.random.rand(15,16,17)
shifted = cmath.ifftshift(cmath.complex_to_torch(arr))
assert np.allclose(ifftshift(arr),
cmath.torch_to_complex(shifted))
#3D, choosing specific axes
arr = np.random.rand(15,16,17) + 1j * np.random.rand(15,16,17)
shifted = cmath.ifftshift(cmath.complex_to_torch(arr),dims=(0,1))
assert np.allclose(ifftshift(arr,axes=(0,1)),
cmath.torch_to_complex(shifted))
def test_expi():
phases = np.random.rand(20,70) * 2 * np.pi
np_result = np.exp(1j * phases)
torch_result = cmath.expi(t.Tensor(phases))
assert np.allclose(np_result, cmath.torch_to_complex(torch_result))
tests/tools/test_data.py
+1 -6
View File
@@ -1,5 +1,3 @@
from __future__ import division, print_function, absolute_import
from CDTools.tools import data
import numpy as np
import torch as t
@@ -8,10 +6,7 @@ import pytest
import os
import datetime
import numbers
try:
import pathlib
except ImportError:
import pathlib2 as pathlib
import pathlib
tests/tools/test_image_processing.py
+22 -18
View File
@@ -5,7 +5,7 @@ import pytest
import numpy as np
import torch as t
from CDTools.tools import image_processing, cmath, initializers, interactions
from CDTools.tools import image_processing, initializers, interactions
from scipy import ndimage
from scipy.signal import fftconvolve
@@ -32,8 +32,8 @@ def test_centroid_sq():
assert t.allclose(centroid, t.Tensor(sp_centroid))
# Test complex with multiple ims
ims = t.rand((5,30,40,2))
np_ims = cmath.torch_to_complex(ims)
ims = t.rand((5,30,40)) + 1j * t.rand((5,30,40))
np_ims = ims.numpy()
sp_centroids = [ndimage.measurements.center_of_mass(np.abs(im)**2)
for im in np_ims]
centroids = image_processing.centroid_sq(ims, comp=True)
@@ -51,12 +51,12 @@ def test_sinc_subpixel_shift():
Ys,Xs = np.meshgrid(xs,xs)
sinc_im = np.sinc(Xs-0.3) * np.sinc(Ys-0.6)
torch_im = cmath.complex_to_torch(im)
torch_im = t.as_tensor(im)
test_im = image_processing.sinc_subpixel_shift(torch_im,(0.3,0.6))
# The fidelity isn't great due to the FFT-based approach, so we need
# a pretty relaxed condition
assert np.max(np.abs(sinc_im - cmath.torch_to_complex(test_im))) < 0.005
assert np.max(np.abs(sinc_im - test_im.numpy())) < 0.005
def test_find_pixel_shift():
@@ -69,13 +69,13 @@ def test_find_pixel_shift():
# Test a real and complex im
big_im = t.rand((30,70))
im1 = t.stack((big_im[:-5,10:],t.zeros_like(big_im[:-5,10:])),dim=-1)
im1 = big_im[:-5,10:].to(dtype=t.complex64)
im2 = big_im[5:,:-10]
assert t.all(image_processing.find_pixel_shift(im1,im2) == t.LongTensor([5,-10]))
assert t.all(image_processing.find_pixel_shift(im2,im1) == t.LongTensor([-5,10]))
# Test two complex ims
big_im = t.rand((45,45,2))
big_im = t.rand((45,45)) + 1j * t.rand((45,45))
im1 = big_im[:-5,:-4]
im2 = big_im[5:,4:]
assert t.all(image_processing.find_pixel_shift(im1,im2) == t.LongTensor([5,4]))
@@ -83,13 +83,13 @@ def test_find_pixel_shift():
def test_find_subpixel_shift():
# We can do this by creating a test probe and a test object
test_probe = t.rand((70,70,2))
test_obj = t.ones((300,300,2))
test_probe = t.rand((70,70)) + 1j * t.rand((70,70))
test_obj = t.ones((300,300)) + 1j * t.rand((300,300))
shift = t.tensor((0.8,0.75))
im = interactions.ptycho_2D_sinc(test_probe, test_obj, shift)
im = interactions.ptycho_2D_sinc(test_probe, test_obj, shift, multiple_modes=False)
retrieved_shift = image_processing.find_subpixel_shift(im, test_probe, search_around=(0,0), resolution=50)
# tolerance of 0.03 on this measurement
assert t.all(t.abs(shift - retrieved_shift) < 0.03)
@@ -98,12 +98,13 @@ def test_find_subpixel_shift():
def test_find_shift():
# We can do this by creating a test probe and a test object
test_probe = t.rand((200,200,2))
test_obj = t.ones((300,300,2))
test_probe = t.rand((200,200)) + 1j * t.rand((200,200))
test_obj = t.ones((300,300)) + 1j * t.rand((300,300))
shift = t.tensor((0.8,0.75))
im = interactions.ptycho_2D_sinc(test_probe, test_obj, shift)[:-40,:-6]
im = interactions.ptycho_2D_sinc(test_probe, test_obj, shift,
multiple_modes=False)[:-40,:-6]
retrieved_shift = image_processing.find_shift(im, test_probe[40:,6:], resolution=50)
# tolerance of 0.03 on this measurement
@@ -111,14 +112,14 @@ def test_find_shift():
def test_convolve_1d():
from matplotlib import pyplot as plt
test_image = np.random.rand(400,300)
#test_image = np.hstack((np.ones((400,150)),np.zeros((400,150))))
xs = np.linspace(-100,100,300)
kernel = 1/(1+xs**2)
# First, we test with everything real, dim=1
convolved = image_processing.convolve_1d(t.Tensor(test_image),t.Tensor(kernel),dim=1)
convolved = image_processing.convolve_1d(t.as_tensor(test_image),
t.as_tensor(kernel),dim=1)
np_result = np.abs(np.fft.ifft(np.fft.fft(test_image,axis=1) * np.fft.fft(np.fft.ifftshift(kernel)), axis=1))
assert np.allclose(convolved.numpy(),np_result)
@@ -128,13 +129,16 @@ def test_convolve_1d():
kernel = 1/(1+xs**2)
# Then with dim=0, and a non-fftshifted kernel
convolved = image_processing.convolve_1d(t.Tensor(test_image),t.Tensor(np.fft.ifftshift(kernel)), fftshift_kernel=False)
convolved = image_processing.convolve_1d(t.as_tensor(test_image),
t.as_tensor(np.fft.ifftshift(kernel)),
fftshift_kernel=False)
np_result = np.abs(np.fft.ifft(np.fft.fft(test_image,axis=0) * np.fft.fft(np.fft.ifftshift(kernel))[:,None], axis=0))
assert np.allclose(convolved.numpy(),np_result)
# And finally with complex input
convolved = cmath.torch_to_complex(image_processing.convolve_1d(cmath.complex_to_torch(test_image),cmath.complex_to_torch(kernel)))
convolved = image_processing.convolve_1d(t.as_tensor(test_image,dtype=t.complex64),
t.as_tensor(kernel,dtype=t.complex64)).numpy()
np_result = np.fft.ifft(np.fft.fft(test_image,axis=0) * np.fft.fft(np.fft.ifftshift(kernel))[:,None], axis=0)
assert np.allclose(convolved,np_result)
tests/tools/test_initializers.py
+16 -18
View File
@@ -1,7 +1,6 @@
from __future__ import division, print_function, absolute_import
from CDTools.tools import initializers
from CDTools.tools import cmath
from CDTools.datasets import Ptycho2DDataset
import numpy as np
import torch as t
@@ -78,11 +77,12 @@ def test_gaussian():
# Generate gaussian as a numpy array (square array)
shape = [10, 10]
sigma = [2.5, 2.5]
center = ((shape[0]-1)/2, (shape[1]-1)/2)
y, x = np.mgrid[:shape[0], :shape[1]]
np_result = 10*np.exp(-0.5*((x-center[1])/sigma[1])**2
-0.5*((y-center[0])/sigma[0])**2)
init_result = cmath.torch_to_complex(initializers.gaussian([10, 10], [2.5, 2.5], amplitude=10))
init_result = initializers.gaussian(shape, sigma, amplitude=10).numpy()
assert np.allclose(init_result, np_result)
# Generate gaussian as a numpy array (rectangular array)
@@ -92,7 +92,7 @@ def test_gaussian():
y, x = np.mgrid[:shape[0], :shape[1]]
np_result = np.exp(-0.5*((x-center[1])/sigma[1])**2
-0.5*((y-center[0])/sigma[0])**2)
init_result = cmath.torch_to_complex(initializers.gaussian(shape, sigma))
init_result = initializers.gaussian(shape, sigma).numpy()
assert np.allclose(init_result, np_result)
# Generate gaussian with curvature
@@ -105,8 +105,8 @@ def test_gaussian():
-0.5*((y-center[0])/sigma[0])**2)
np_result *= np.exp(0.5j*curvature[1]*(x-center[1])**2
+0.5j*curvature[0]*(y-center[0])**2)
init_result = cmath.torch_to_complex(initializers.gaussian(shape, sigma,
center=center, curvature=curvature, amplitude=10))
init_result = initializers.gaussian(shape, sigma, center=center,
curvature=curvature, amplitude=10).numpy()
assert np.allclose(init_result, np_result)
@@ -147,8 +147,7 @@ def test_gaussian_probe(ptycho_cxi_1):
normalization_1 = normalization / np.sum(np.abs(np_probe)**2)
probe = initializers.gaussian_probe(dataset, basis, shape, sigma)
probe = cmath.torch_to_complex(probe)
probe = initializers.gaussian_probe(dataset, basis, shape, sigma).numpy()
assert np.allclose(probe, normalization_1*np_probe)
# And then a propagated probe
@@ -163,8 +162,7 @@ def test_gaussian_probe(ptycho_cxi_1):
normalization_2 = normalization / np.sum(np.abs(np_probe)**2)
probe = initializers.gaussian_probe(dataset, basis, shape, sigma,
propagation_distance=z)
probe = cmath.torch_to_complex(probe)
propagation_distance=z).numpy()
assert np.allclose(probe, normalization_2*np_probe)
@@ -185,10 +183,10 @@ def test_SHARP_style_probe(ptycho_cxi_1):
distance)
probe = initializers.SHARP_style_probe(dataset, shape, det_slice)
assert probe.shape == t.Size([256,256,2])
assert probe.shape == t.Size([256,256])
probe = initializers.SHARP_style_probe(dataset, shape, det_slice, propagation_distance=20e-6)
assert probe.shape == t.Size([256,256,2])
assert probe.shape == t.Size([256,256])
def test_RPI_spectral_init():
@@ -200,24 +198,24 @@ def test_RPI_spectral_init():
probe = np.random.rand(230,253).astype(np.complex64)
obj_shape = [37,53]
mask = t.Tensor(np.random.rand(*pattern.shape) > 0.04)
background = t.Tensor(np.random.rand(*pattern.shape) .astype(np.float32)* 0.05)
background = t.as_tensor(np.random.rand(*pattern.shape),dtype=t.float32) * 0.05
probe = cmath.complex_to_torch(probe)
pattern = t.Tensor(pattern)
probe = t.as_tensor(probe)
pattern = t.as_tensor(pattern)
obj = initializers.RPI_spectral_init(pattern, probe, obj_shape)
assert list(obj.shape) == [1]+obj_shape+[2]
assert list(obj.shape) == [1]+obj_shape
obj = initializers.RPI_spectral_init(pattern, probe, obj_shape,
n_modes=2, mask=mask)
assert list(obj.shape) == [2]+obj_shape+[2]
assert list(obj.shape) == [2]+obj_shape
obj = initializers.RPI_spectral_init(pattern, probe, obj_shape,
n_modes=2, background=background)
assert list(obj.shape) == [2]+obj_shape+[2]
assert list(obj.shape) == [2]+obj_shape
obj = initializers.RPI_spectral_init(pattern, probe, obj_shape,
n_modes=2, mask=mask,
background=background)
assert list(obj.shape) == [2]+obj_shape+[2]
assert list(obj.shape) == [2]+obj_shape
tests/tools/test_interactions.py
+31 -32
View File
@@ -1,6 +1,5 @@
from __future__ import division, print_function, absolute_import
from CDTools.tools import cmath
from CDTools.tools import interactions
import numpy as np
import torch as t
@@ -137,16 +136,16 @@ def test_ptycho_2D_round(random_probe, random_obj):
random_obj[tr[0]:tr[0]+random_probe.shape[0],
tr[1]:tr[1]+random_probe.shape[1]] for
tr in np.round(translations).astype(int)]
exit_waves_t = interactions.ptycho_2D_round(cmath.complex_to_torch(random_probe),
cmath.complex_to_torch(random_obj),
t.tensor(translations))
assert np.allclose(cmath.torch_to_complex(exit_waves_t), exit_waves_np)
exit_waves_t = interactions.ptycho_2D_round(t.as_tensor(random_probe),
t.as_tensor(random_obj),
t.as_tensor(translations))
assert np.allclose(exit_waves_t.numpy(), exit_waves_np)
# Test the single wave case
exit_wave_t = interactions.ptycho_2D_round(cmath.complex_to_torch(random_probe),
cmath.complex_to_torch(random_obj),
t.tensor(translations[0]))
assert np.allclose(cmath.torch_to_complex(exit_wave_t), exit_waves_np[0])
exit_wave_t = interactions.ptycho_2D_round(t.as_tensor(random_probe),
t.as_tensor(random_obj),
t.as_tensor(translations[0]))
assert np.allclose(exit_wave_t.numpy(), exit_waves_np[0])
@@ -158,14 +157,14 @@ def test_ptycho_2D_linear(single_pixel_probe, random_obj):
translation = np.array([46.7,53.2])
exit_waves_probe = interactions.ptycho_2D_linear(
cmath.complex_to_torch(single_pixel_probe),
cmath.complex_to_torch(random_obj),
t.tensor(translations),
t.as_tensor(single_pixel_probe),
t.as_tensor(random_obj),
t.as_tensor(translations),
shift_probe=True)
exit_wave_probe = interactions.ptycho_2D_linear(
cmath.complex_to_torch(single_pixel_probe),
cmath.complex_to_torch(random_obj),
t.as_tensor(single_pixel_probe),
t.as_tensor(random_obj),
t.tensor(translation),
shift_probe=True)
@@ -174,14 +173,14 @@ def test_ptycho_2D_linear(single_pixel_probe, random_obj):
exit_waves_obj = interactions.ptycho_2D_linear(
cmath.complex_to_torch(single_pixel_probe),
cmath.complex_to_torch(random_obj),
t.as_tensor(single_pixel_probe),
t.as_tensor(random_obj),
t.tensor(translations),
shift_probe=False)
exit_wave_obj = interactions.ptycho_2D_linear(
cmath.complex_to_torch(single_pixel_probe),
cmath.complex_to_torch(random_obj),
t.as_tensor(single_pixel_probe),
t.as_tensor(random_obj),
t.tensor(translation),
shift_probe=False)
@@ -189,7 +188,7 @@ def test_ptycho_2D_linear(single_pixel_probe, random_obj):
assert t.allclose(exit_waves_obj[0],exit_wave_obj)
# For the shifted probe, we should find 4 pixels with intensity
exit_waves_probe = cmath.torch_to_complex(exit_waves_probe)[0]
exit_waves_probe = t.as_tensor(exit_waves_probe)[0]
probe_shift = np.array([[0.3*0.8,0.3*0.2],
[0.7*0.8,0.7*0.2]])
@@ -199,7 +198,7 @@ def test_ptycho_2D_linear(single_pixel_probe, random_obj):
assert np.allclose(probe_shift * obj_section, exit_section)
# For the shifted obj, we should find one pixel with intensity
exit_waves_obj = cmath.torch_to_complex(exit_waves_obj)[0]
exit_waves_obj = t.as_tensor(exit_waves_obj)[0]
obj_shift = np.array([[0.3*0.8,0.3*0.2],
[0.7*0.8,0.7*0.2]])
obj_section = random_obj[128+46:128+48,
@@ -219,15 +218,15 @@ def test_ptycho_2D_sinc(single_pixel_probe, random_obj):
translation = np.array([46.7,53.2])
exit_waves_probe = interactions.ptycho_2D_sinc(
cmath.complex_to_torch(single_pixel_probe),
cmath.complex_to_torch(random_obj),
t.tensor(translations),
t.as_tensor(single_pixel_probe),
t.as_tensor(random_obj),
t.as_tensor(translations),
shift_probe=True)
exit_wave_probe = interactions.ptycho_2D_sinc(
cmath.complex_to_torch(single_pixel_probe),
cmath.complex_to_torch(random_obj),
t.tensor(translation),
t.as_tensor(single_pixel_probe),
t.as_tensor(random_obj),
t.as_tensor(translation),
shift_probe=True)
# Check that the outputs match
@@ -247,7 +246,7 @@ def test_ptycho_2D_sinc(single_pixel_probe, random_obj):
53:53+256]
exit_wave_np = sinc_shifted_probe * obj_section
exit_wave_torch = cmath.torch_to_complex(exit_wave_probe)
exit_wave_torch = exit_wave_probe.numpy()
# The fidelity isn't great due to the FFT-based approach, so we need
# a pretty relaxed condition
@@ -259,8 +258,8 @@ def test_RPI_interaction(random_probe, random_obj):
random_obj1 = random_obj[:79,:68]
random_probe1 = random_probe
t_random_obj1 = cmath.complex_to_torch(random_obj1)
t_random_probe1 = cmath.complex_to_torch(random_probe1)
t_random_obj1 = t.as_tensor(random_obj1)
t_random_probe1 = t.as_tensor(random_probe1)
t_output1 = interactions.RPI_interaction(t_random_probe1, t_random_obj1)
obj1_fourier = fftshift(fft.fft2(ifftshift(random_obj1), norm='ortho'))
@@ -272,13 +271,13 @@ def test_RPI_interaction(random_probe, random_obj):
output1 = random_probe1 * fftshift(fft.ifft2(ifftshift(obj1_ups),
norm='ortho'))
assert np.allclose(cmath.torch_to_complex(t_output1), output1)
assert np.allclose(t.as_tensor(t_output1), output1)
random_obj2 = np.stack([random_obj[:64,:89]]*3)
random_probe2 = random_probe[3:,5:]
t_random_obj2 = cmath.complex_to_torch(random_obj2)
t_random_probe2 = cmath.complex_to_torch(random_probe2)
t_random_obj2 = t.as_tensor(random_obj2)
t_random_probe2 = t.as_tensor(random_probe2)
t_output2 = interactions.RPI_interaction(t_random_probe2, t_random_obj2)
obj2_fourier = fftshift(fft.fft2(ifftshift(random_obj2), norm='ortho'))
tests/tools/test_losses.py
+3 -4
View File
@@ -1,7 +1,6 @@
from __future__ import division, print_function, absolute_import
from CDTools.tools import losses
from CDTools.tools import cmath
import numpy as np
import torch as t
@@ -16,7 +15,7 @@ def test_amplitude_mse():
# And add some noise to it
sim = data + 0.1 * np.random.rand(10,100,100)
# and define a simple mask that needs to be broadcast
mask = (np.random.rand(100,100) > 0.1).astype(np.bool)
mask = (np.random.rand(100,100) > 0.1).astype(bool)
# First, test without a mask
np_result = np.sum((np.sqrt(data) - np.sqrt(sim))**2)
@@ -38,7 +37,7 @@ def test_intensity_mse():
# And add some noise to it
sim = data + 0.1 * np.random.rand(10,100,100)
# and define a simple mask that needs to be broadcast
mask = (np.random.rand(100,100) > 0.1).astype(np.bool)
mask = (np.random.rand(100,100) > 0.1).astype(bool)
# First, test without a mask
@@ -61,7 +60,7 @@ def test_poisson_nll():
# And add some noise to it
sim = data + 0.1 * np.random.rand(10,100,100)
# and define a simple mask that needs to be broadcast
mask = (np.random.rand(100,100) > 0.1).astype(np.bool)
mask = (np.random.rand(100,100) > 0.1).astype(bool)
# First, test without a mask
tests/tools/test_measurements.py
+23 -24
View File
@@ -1,31 +1,30 @@
from __future__ import division, print_function, absolute_import
from CDTools.tools import measurements
from CDTools.tools import cmath
import torch as t
import numpy as np
import pytest
def test_intensity():
wavefields = t.rand((5,10,10,2))
wavefields = t.rand((5,10,10)) + 1j * t.rand((5,10,10))
epsilon=1e-6
np_result = np.abs(cmath.torch_to_complex(wavefields))**2 + epsilon
np_result = np.abs(t.as_tensor(wavefields))**2 + epsilon
assert t.allclose(measurements.intensity(wavefields,epsilon=epsilon),
t.tensor(np_result))
t.as_tensor(np_result))
# Test single field case
assert t.allclose(measurements.intensity(wavefields[0],epsilon=epsilon),
t.tensor(np_result[0]))
t.as_tensor(np_result[0]))
det_slice = np.s_[3:,5:8]
assert t.allclose(measurements.intensity(wavefields,det_slice,epsilon=epsilon),
t.tensor(np_result[(np.s_[:],)+det_slice]))
t.as_tensor(np_result[(np.s_[:],)+det_slice]))
# Test single field case
assert t.allclose(measurements.intensity(wavefields[0],det_slice,epsilon=epsilon),
t.tensor(np_result[0][det_slice]))
t.as_tensor(np_result[0][det_slice]))
# With oversampling on
@@ -36,34 +35,34 @@ def test_intensity():
# With multiple fields
assert t.allclose(measurements.intensity(wavefields,epsilon=epsilon, oversampling=2),
t.tensor(np_oversampling_result,))
t.as_tensor(np_oversampling_result,))
# With a single field
assert t.allclose(measurements.intensity(wavefields[0],epsilon=epsilon, oversampling=2),
t.tensor(np_oversampling_result[0],))
t.as_tensor(np_oversampling_result[0],))
def test_incoherent_sum():
# With no explicit slice given
wavefields = t.rand((5,4,10,10,2))
wavefields = t.rand((5,4,10,10)) + 1j * t.rand((5,4,10,10))
epsilon=1e-6
np_result = np.sum(np.abs(cmath.torch_to_complex(wavefields))**2,axis=0) + epsilon
np_result = np.sum(np.abs(wavefields.numpy())**2,axis=-3) + epsilon
assert t.allclose(measurements.incoherent_sum(wavefields,epsilon=epsilon),
t.tensor(np_result))
t.as_tensor(np_result))
# Test single field case
assert t.allclose(measurements.incoherent_sum(wavefields[:,0],epsilon=epsilon),
t.tensor(np_result[0]))
assert t.allclose(measurements.incoherent_sum(wavefields[0,:],epsilon=epsilon),
t.as_tensor(np_result[0]))
# With a slice given
det_slice = np.s_[3:,5:8]
assert t.allclose(measurements.incoherent_sum(wavefields,det_slice,epsilon=epsilon),
t.tensor(np_result[(np.s_[:],)+det_slice]))
t.as_tensor(np_result[(np.s_[:],)+det_slice]))
# Test single field case
assert t.allclose(measurements.incoherent_sum(wavefields[:,0],det_slice,epsilon=epsilon),
t.tensor(np_result[0][det_slice]))
assert t.allclose(measurements.incoherent_sum(wavefields[0,:],det_slice,epsilon=epsilon),
t.as_tensor(np_result[0][det_slice]))
# With oversampling on
np_oversampling_result = (np_result[:,::2,::2] + \
@@ -73,19 +72,19 @@ def test_incoherent_sum():
# With multiple fields
assert t.allclose(measurements.incoherent_sum(wavefields,epsilon=epsilon, oversampling=2),
t.tensor(np_oversampling_result,))
t.as_tensor(np_oversampling_result,))
# With a single field
assert t.allclose(measurements.incoherent_sum(wavefields[:,0],epsilon=epsilon, oversampling=2),
t.tensor(np_oversampling_result[0],))
assert t.allclose(measurements.incoherent_sum(wavefields[0,:],epsilon=epsilon, oversampling=2),
t.as_tensor(np_oversampling_result[0],))
def test_quadratic_background():
# test with intensity
wavefields = t.rand((5,10,10,2))
wavefields = t.rand((5,10,10)) + 1j * t.rand((5,10,10))
epsilon=1e-6
background = t.rand((10,10))
np_result = np.abs(cmath.torch_to_complex(wavefields))**2 + background.numpy()**2 + epsilon
np_result = np.abs(wavefields.numpy())**2 + background.numpy()**2 + epsilon
det_slice = np.s_[3:,5:8]
result = measurements.quadratic_background(wavefields,background[det_slice],
@@ -96,8 +95,8 @@ def test_quadratic_background():
# test with incoherent sum but no slice and no stack
wavefields = t.rand((4,10,10,2))
np_result = np.sum(np.abs(cmath.torch_to_complex(wavefields))**2,axis=0)
wavefields = t.rand((4,10,10)) + 1j * t.rand((4,10,10))
np_result = np.sum(np.abs(wavefields.numpy())**2,axis=0)
np_result += background.numpy()**2
result = measurements.quadratic_background(wavefields, background,
epsilon=epsilon,
tests/tools/test_plotting.py
+4 -6
View File
@@ -1,6 +1,5 @@
from __future__ import division, print_function, absolute_import
from CDTools.tools import cmath
from CDTools.tools import plotting
from CDTools.tools import initializers
import numpy as np
@@ -11,7 +10,7 @@ import matplotlib.pyplot as plt
def test_plot_amplitude(show_plot):
# Test with tensor
im = cmath.complex_to_torch(scipy.misc.ascent().astype(np.float64))
im = t.as_tensor(scipy.misc.ascent(),dtype=t.complex128)
plotting.plot_amplitude(im, basis = np.array([[1,1], [1,1], [0,0]]), title = 'Test Amplitude')
if show_plot:
plt.show()
@@ -31,7 +30,7 @@ def test_plot_phase(show_plot):
plt.show()
# Test with numpy array
im = cmath.torch_to_complex(initializers.gaussian([512, 512], [200,200], amplitude=100, curvature=[.1,.1]))
im = initializers.gaussian([512, 512], [200,200], amplitude=100, curvature=[.1,.1]).numpy()
plotting.plot_phase(im, title = 'Test Phase', basis = np.array([[1,1], [1,1], [0,0]]))
if show_plot:
plt.show()
@@ -39,14 +38,13 @@ def test_plot_phase(show_plot):
def test_plot_colorize(show_plot):
# Test with tensor
gaussian = initializers.gaussian([512, 512], [200,200], amplitude=100, curvature=[.1,.1])
im = cmath.cmult(gaussian, cmath.complex_to_torch(scipy.misc.ascent().astype(np.float64)))
im = gaussian * t.as_tensor(scipy.misc.ascent(), dtype=t.complex64)
plotting.plot_colorized(im, title = 'Test Colorize', basis = np.array([[1,1], [1,1], [0,0]]))
if show_plot:
plt.show()
# Test with numpy array
gaussian = initializers.gaussian([512, 512], [200,200], amplitude=100, curvature=[.1,.1])
im = cmath.torch_to_complex(cmath.cmult(gaussian, cmath.complex_to_torch(scipy.misc.ascent().astype(np.float64))))
im = im.numpy()
plotting.plot_colorized(im, title = 'Test Colorize')
if show_plot:
plt.show()
tests/tools/test_projectors.py
-39
View File
@@ -1,39 +0,0 @@
from __future__ import division, print_function, absolute_import
from CDTools.tools import cmath
from CDTools.tools import projectors
import numpy as np
import torch as t
def test_modulus():
# Create a complex array with random modulus and known phase
np_result = np.sqrt(6) * (1 + 1j) * np.random.rand(10,10)
projection_intensity = t.from_numpy(np.abs(np_result)**2).to(t.float32)
original_wavefront = cmath.complex_to_torch((1+1j) * np.random.rand(10,10)).to(t.float32)
# Test without masks
torch_result = projectors.modulus(original_wavefront,projection_intensity)
assert np.allclose(cmath.torch_to_complex(torch_result),np_result)
# Test with mask
mask = t.ones((10,10,2), dtype = t.uint8)
mask[5]*=0
np_result[5] = cmath.torch_to_complex(original_wavefront[5])
torch_result = projectors.modulus(original_wavefront,projection_intensity, mask=mask)
print(np_result[5])
print(cmath.torch_to_complex(torch_result)[5])
assert np.allclose(cmath.torch_to_complex(torch_result),np_result)
def test_support():
# Define a mask as uint8 to make sure it works when support is not
# the same type as the wavefield
support = t.zeros((10,10)).to(t.uint8)
# some masked, some unmasked
support[:3,:3] = 1
np_result = np.zeros((10,10)).astype(np.complex128)
np_result[:3,:3] = 1 + 1j
assert(np.allclose(cmath.torch_to_complex(projectors.support(t.ones((10,10,2)), support)), np_result))
tests/tools/test_propagators.py
+121 -45
View File
@@ -1,8 +1,8 @@
from __future__ import division, print_function, absolute_import
from CDTools.tools import cmath
from CDTools.tools import initializers
from CDTools.tools import propagators
from CDTools.tools import image_processing
import numpy as np
import torch as t
@@ -19,26 +19,26 @@ def exit_waves_1():
obj = scipy.misc.ascent()[0:64,0:64].astype(np.complex128)
arr = np.random.random_sample((64,64))
obj *= (arr+(1-arr**2)**.5*1j)
obj = cmath.complex_to_torch(obj)
obj = t.as_tensor(obj)
# Construct wavefront from image
probe = initializers.gaussian([64, 64], [5, 5], amplitude=1e3)
return cmath.cmult(probe,obj)
return probe * obj
def test_far_field(exit_waves_1):
# Far field diffraction patterns calculated by numpy with zero frequency in center
np_result = np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(cmath.torch_to_complex(exit_waves_1)),norm='ortho'))
np_result = np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(exit_waves_1.numpy()),norm='ortho'))
assert(np.allclose(np_result, cmath.torch_to_complex(propagators.far_field(exit_waves_1))))
assert(np.allclose(np_result, propagators.far_field(exit_waves_1).numpy()))
def test_inverse_far_field(exit_waves_1):
# We want the inverse far field to map back to the exit waves with no intensity corrections
# Far field result for exit waves calculated with numpy
far_field_np_result = cmath.complex_to_torch(np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(cmath.torch_to_complex(exit_waves_1)),norm='ortho')))
far_field_np_result = t.as_tensor(np.fft.fftshift(np.fft.fft2(np.fft.ifftshift(exit_waves_1.numpy()),norm='ortho')))
assert(np.allclose(exit_waves_1, propagators.inverse_far_field(far_field_np_result)))
@@ -67,14 +67,14 @@ def test_generate_high_NA_k_intensity_map():
j = (np.arange(573) - 270)
Is,Js = np.meshgrid(i,j,indexing='ij')
wavefield = ((np.abs(Is) < 20) * (np.abs(Js) < 25)).astype(np.complex128)
t_wavefield = cmath.complex_to_torch(wavefield).to(dtype=t.float32)
t_wavefield = t.as_tensor(wavefield, dtype=t.complex64)
high_NA_propagated = propagators.high_NA_far_field(
t_wavefield, k_map, intensity_map=intensity_map)
low_NA_propagated = propagators.far_field(t_wavefield)
low_NA = cmath.torch_to_complex(low_NA_propagated)
high_NA = cmath.torch_to_complex(high_NA_propagated)
low_NA = low_NA_propagated.numpy()
high_NA = high_NA_propagated.numpy()
# Checking first that for a low-NA propagation they give the same result
# 1e-4 tolerance seems to be reasonable in this comparison given my
@@ -100,8 +100,8 @@ def test_generate_high_NA_k_intensity_map():
t_wavefield, k_map, intensity_map=intensity_map)
low_NA_propagated = propagators.far_field(t_wavefield)
low_NA = cmath.torch_to_complex(low_NA_propagated)
high_NA = cmath.torch_to_complex(high_NA_propagated)
low_NA = low_NA_propagated.numpy()
high_NA = high_NA_propagated.numpy()
#plt.close('all')
#plt.imshow(np.abs(low_NA))
@@ -121,11 +121,54 @@ def test_generate_high_NA_k_intensity_map():
# I believe this works, but I still would like to get a second method for
# simulating at least one diffraction pattern as an independent check
assert 0
#assert 0
def test_near_field_direction():
#
# The idea is to make sure that near field propagation is consistent
# with far-field propagation. Because we consider the Fourier
# transform to be the far-field propagator, we should find that
# near-field propagation becomes far-field propagation in the limit
# that z goes to positive infinity. This consistency requirement
# forces a specific choice of sign for near-field propagation - choosing
# it incorrectly will make the correct limit obtain at negative
# infinity.
#
# We can check consistency by setting up a gaussian wavefield with an
# offset in the far field. Then, we can do some near-field propagation,
# and we will ensure that the centroid of the beam moves in the correct
# direction. If the sign is incorrect, it will move in the wrong direction.
x = (np.arange(901) - 400)
y = (np.arange(1200) - 500)
Ys,Xs = np.meshgrid(y,x)
Rs = np.sqrt(Xs**2+Ys**2)
E0_fourier = t.as_tensor(np.exp(-Rs**2 / (2 * 40**2)),dtype=t.complex64)
E0_real = propagators.inverse_far_field(E0_fourier)
# This is in the top-left corner in Fourier space
wavelength = 3e-9 #nm
z = 1000e-9
asp = propagators.generate_angular_spectrum_propagator(
E0_real.shape,(1.5e-9,1e-9),wavelength,z,dtype=t.complex64)
Ez_real = propagators.near_field(E0_real,asp)
centroid = image_processing.centroid(t.abs(Ez_real))
# Assert it's in top half
assert centroid[0] < Ez_real.shape[0] // 2
# Assert it's in left half
assert centroid[1] < Ez_real.shape[1] // 2
#plt.imshow(t.abs(E0_fourier))
#plt.figure()
#plt.imshow(t.abs(E0_real))
#plt.figure()
#plt.imshow(t.abs(Ez_real))
#plt.show()
def test_near_field():
@@ -151,17 +194,37 @@ def test_near_field():
# The analytical expression for propagation of a gaussian beam in the
# paraxial approx
Ez = w0 / wz * np.exp(-Rs**2 / wz**2) * np.exp(-1j * k * ( z + Rs**2 / (2 * Rz)) + 1j * np.arctan(z / zr))
Ez_nozphase = Ez * np.exp(1j * k * z)
# The sign convention I'm using here is opposite from what is on
# wikipedia (as of June 2021), but it is I believe the most sensible
# choice. By choosing e^(ikx) to be the plane wave that's propagating
# in the x direction, we preserve two nice properties:
# 1) The time-dependence is e^(-iwt+ikx), which meshes with the choice
# we like to make as physicists of e^(-iwt) for time-dependence
# 2) This makes it so that the Fourier transform does far-field
# propagation in direction of light propagation. The logic is as
# follows:
# a) We must assume that light only propagates one direction through
# the plane, in order to have well-defined propagation
# b) Fixing that direction to be the positive direciton of propagation,
# we can consider the in-plane image of a plane wave propagating
# through the plane of interest.
# c) We can ask whether the 2D Fourier transform of that in-plane image
# produces a spot on the same side of zero as the in-plane component
# of k.
# If we choose e^(-ikx) to represent light propagating along K, the
# answer is no, and we find we have to use the inverse FT instead.
# Thus, e^(ikx) is the right choice here.
Ez = w0 / wz * np.exp(-Rs**2 / wz**2) * np.exp(1j * k * ( z + Rs**2 / (2 * Rz)) - 1j * np.arctan(z / zr))
Ez_nozphase = Ez * np.exp(-1j * k * z)
# First we check it normally
asp = propagators.generate_angular_spectrum_propagator(
E0.shape,(1.5e-9,1e-9),wavelength,z,dtype=t.float64)
E0.shape,(1.5e-9,1e-9),wavelength,z,dtype=t.complex128)
Ez_t = propagators.near_field(cmath.complex_to_torch(E0),asp)
Ez_t = cmath.torch_to_complex(Ez_t)
Ez_t = propagators.near_field(t.as_tensor(E0),asp).numpy()
# Check for at least 10^-3 relative accuracy in this scenario
assert np.max(np.abs(Ez-Ez_t)) < 1e-3 * np.max(np.abs(Ez))
@@ -169,8 +232,7 @@ def test_near_field():
Emz = np.conj(Ez)
Emz_t = propagators.inverse_near_field(cmath.complex_to_torch(E0),asp)
Emz_t = cmath.torch_to_complex(Emz_t)
Emz_t = propagators.inverse_near_field(t.as_tensor(E0),asp).numpy()
# Again, 10^-3 is about all the accuracy we can expect
assert np.max(np.abs(Emz-Emz_t)) < 1e-3 * np.max(np.abs(Emz))
@@ -178,11 +240,10 @@ def test_near_field():
# Then, we check it with the phase correction
asp = propagators.generate_angular_spectrum_propagator(
E0.shape,(1.5e-9,1e-9),wavelength,z,remove_z_phase=True,
dtype=t.float64)
dtype=t.complex128)
Ez_t = propagators.near_field(cmath.complex_to_torch(E0),asp)
Ez_t = cmath.torch_to_complex(Ez_t)
Ez_t = propagators.near_field(t.as_tensor(E0),asp).numpy()
# Check for at least 10^-3 relative accuracy in this scenario
assert np.max(np.abs(Ez_nozphase-Ez_t)) < 1e-3 * np.max(np.abs(Ez_nozphase))
@@ -190,12 +251,20 @@ def test_near_field():
Emz = np.conj(Ez_nozphase)
Emz_t = propagators.inverse_near_field(cmath.complex_to_torch(E0),asp)
Emz_t = cmath.torch_to_complex(Emz_t)
Emz_t = propagators.inverse_near_field(t.as_tensor(E0),asp).numpy()
# Again, 10^-3 is about all the accuracy we can expect
assert np.max(np.abs(Emz-Emz_t)) < 1e-3 * np.max(np.abs(Emz))
# Finally, we check that the bandlimiting at least does something
asp = propagators.generate_angular_spectrum_propagator(
E0.shape,(1.5e-9,1e-9),wavelength,z,remove_z_phase=True,
dtype=t.complex128, bandlimit=0.3)
assert asp[140,0] == 0
assert asp[0,180] == 0
assert asp[130,0] != 0
assert asp[0,175] != 0
def test_generalized_near_field():
@@ -236,14 +305,18 @@ def test_generalized_near_field():
return Zs / (Zs**2 + zr**2)
def get_E(Xs, Ys, Zs, correct=False):
# if correct is True, remove the e^(-ikz) dependence
# Again, this follows the convention opposite from Wikipedia. See
# the note in test_angular_spectrum_propagator.
# if correct is True, remove the e^(ikz) dependence
Rs_sq = Xs**2 + Ys**2
Wzs = get_w(Zs)
E = w0 / Wzs * np.exp(-Rs_sq / Wzs**2) *\
np.exp(-1j * k * ( Zs + Rs_sq * get_inv_R(Zs) / 2) + \
1j * np.arctan(Zs / zr))
np.exp(1j * k * ( Zs + Rs_sq * get_inv_R(Zs) / 2) + \
- 1j * np.arctan(Zs / zr))
if correct:
E = E * np.exp(1j * k * Zs)
E = E * np.exp(-1j * k * Zs)
return E
@@ -338,22 +411,22 @@ def test_generalized_near_field():
E0 = get_E(Xs,Ys,Zs_0, correct=False)
Ez = get_E(Xs_prop,Ys_prop,Zs_prop, correct=False)
asp = propagators.generate_generalized_angular_spectrum_propagator(
E0.shape,new_basis,wavelength,offset_vec,dtype=t.float64)
E0.shape,new_basis,wavelength,offset_vec,dtype=t.complex128)
elif str(propagation_vec) == 'offset':
E0 = get_E(Xs,Ys,Zs_0, correct=True)
Ez = get_E(Xs_prop,Ys_prop,Zs_prop, correct=True)
asp = propagators.generate_generalized_angular_spectrum_propagator(
E0.shape,new_basis,wavelength,offset_vec,
dtype=t.float64, propagate_along_offset=True)
dtype=t.complex128, propagate_along_offset=True)
else:
E0 = get_E(Xs,Ys,Zs_0, correct=True)
Ez = get_E(Xs_prop,Ys_prop,Zs_prop, correct=True)
asp = propagators.generate_generalized_angular_spectrum_propagator(
E0.shape,new_basis,wavelength,offset_vec,
dtype=t.float64, propagation_vector=propagation_vec)
dtype=t.complex128, propagation_vector=propagation_vec)
Ez_t = propagators.near_field(cmath.complex_to_torch(E0),asp)
Ez_t = cmath.torch_to_complex(Ez_t)
Ez_t = propagators.near_field(t.as_tensor(E0),asp).numpy()
# Check for at least 10^-3 relative accuracy in this scenario
if not np.max(np.abs(Ez-Ez_t)) < 1e-3 * np.max(np.abs(Ez)):
#if True:
@@ -362,14 +435,17 @@ def test_generalized_near_field():
plt.figure()
plt.imshow(np.angle(Ez))
plt.figure()
plt.imshow(np.abs(Ez_t))
plt.figure()
plt.imshow(np.angle(Ez_t))
plt.figure()
plt.imshow(np.abs(Ez-Ez_t)/np.max(np.abs(Ez)))
plt.show()
assert np.max(np.abs(Ez-Ez_t)) < 1e-3 * np.max(np.abs(Ez))
Em0_t = propagators.inverse_near_field(cmath.complex_to_torch(Ez),asp)
Em0_t = cmath.torch_to_complex(Em0_t)
Em0_t = propagators.inverse_near_field(t.as_tensor(Ez),asp).numpy()
# Again, 10^-3 is about all the accuracy we can expect
assert np.max(np.abs(E0-Em0_t)) < 1e-3 * np.max(np.abs(E0))
@@ -388,9 +464,8 @@ def test_generalized_near_field():
E0 = get_E(Xs,Ys,Zs_0, correct=True)
asp = propagators.generate_generalized_angular_spectrum_propagator(
E0.shape, new_basis, wavelength,offset_vec,
dtype=t.float64, propagation_vector=propagation_vec)
Ez_t = propagators.near_field(cmath.complex_to_torch(E0),asp)
Ez_t = cmath.torch_to_complex(Ez_t)
dtype=t.complex128, propagation_vector=propagation_vec)
Ez_t = propagators.near_field(t.as_tensor(E0),asp).numpy()
for Rrand in Rrands:
@@ -400,9 +475,9 @@ def test_generalized_near_field():
rot_prop = np.dot(Rrand, propagation_vec)
asp = propagators.generate_generalized_angular_spectrum_propagator(
E0.shape, rot_basis, wavelength, rot_offset,
dtype=t.float64, propagation_vector=rot_prop)
Ez_rot_t = propagators.near_field(cmath.complex_to_torch(E0),asp)
Ez_rot_t = cmath.torch_to_complex(Ez_rot_t)
dtype=t.complex128, propagation_vector=rot_prop)
Ez_rot_t = propagators.near_field(t.as_tensor(E0),asp).numpy()
assert np.max(np.abs(Ez_t-Ez_rot_t)) < 1e-3 * np.max(np.abs(Ez_t))
@@ -421,9 +496,10 @@ def test_inverse_near_field():
E0 = np.exp(-Rs**2 / w0**2)
asp = propagators.generate_angular_spectrum_propagator(
E0.shape,(1.5e-9,1e-9),wavelength,z,dtype=t.float64)
E0.shape,(1.5e-9,1e-9),wavelength,z,dtype=t.complex128)
E0 = cmath.complex_to_torch(E0)
E0 = t.as_tensor(E0,dtype=t.complex128)
E_prop = propagators.near_field(E0,asp)
E_backprop = propagators.inverse_near_field(E_prop, asp)