mirror of
https://github.com/cdtools-developers/cdtools.git
synced 2026-09-10 05:22:41 +02:00
Remove a bunch of models that really didn't belong in the main repo, and fix an issue with the RPI interaction conventions
This commit is contained in:
@@ -20,24 +20,22 @@ defining a new ptychography model before attempting to do so.
|
||||
|
||||
"""
|
||||
|
||||
# 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', 'PolarizationSweptPtycho', 'PolarizedFancyPtycho', 'Bragg2DPtycho', 'Multislice2DPtycho', 'RPI', 'TimeResolvedPtychoCalibration', 'TimeResolvedRPI']
|
||||
# We define __all__ to be sure that import * only imports what we want
|
||||
__all__ = [
|
||||
'CDIModel',
|
||||
'SimplePtycho',
|
||||
'FancyPtycho',
|
||||
'Bragg2DPtycho',
|
||||
'Multislice2DPtycho',
|
||||
'MultislicePtycho',
|
||||
'RPI',
|
||||
]
|
||||
|
||||
from cdtools.models.base import CDIModel
|
||||
from cdtools.models.simple_ptycho import SimplePtycho
|
||||
from cdtools.models.fancy_ptycho import FancyPtycho
|
||||
from cdtools.models.polarized_fancy_ptycho import PolarizedFancyPtycho
|
||||
from cdtools.models.polarization_swept_ptycho import PolarizationSweptPtycho
|
||||
from cdtools.models.bragg_2d_ptycho import Bragg2DPtycho
|
||||
from cdtools.models.multislice_2d_ptycho import Multislice2DPtycho
|
||||
from cdtools.models.multislice_ptycho import MultislicePtycho
|
||||
from cdtools.models.rpi import RPI
|
||||
from cdtools.models.multimode_rpi import MultimodeRPI
|
||||
from cdtools.models.time_resolved_ptycho_calibration import TimeResolvedPtychoCalibration
|
||||
from cdtools.models.time_resolved_rpi import TimeResolvedRPI
|
||||
from cdtools.models.fastccd_ptycho import FastCCDPtycho
|
||||
|
||||
# Still needs to be updated for the new complex numbers
|
||||
#from cdtools.models.s_matrix_ptycho import SMatrixPtycho
|
||||
|
||||
|
||||
@@ -39,7 +39,6 @@ import queue
|
||||
import time
|
||||
from scipy import io
|
||||
from contextlib import contextmanager
|
||||
from .complex_lbfgs import MyLBFGS
|
||||
from cdtools.tools.data import nested_dict_to_h5, h5_to_nested_dict, nested_dict_to_numpy, nested_dict_to_torch
|
||||
|
||||
__all__ = ['CDIModel']
|
||||
@@ -629,8 +628,6 @@ class CDIModel(t.nn.Module):
|
||||
optimizer = t.optim.LBFGS(self.parameters(),
|
||||
lr = lr, history_size=history_size,
|
||||
line_search_fn=line_search_fn)
|
||||
#optimizer = MyLBFGS(self.parameters(),
|
||||
# lr = lr, history_size=history_size)
|
||||
|
||||
return self.AD_optimize(iterations, data_loader, optimizer,
|
||||
regularization_factor=regularization_factor,
|
||||
|
||||
@@ -1,485 +0,0 @@
|
||||
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
|
||||
@@ -1,691 +0,0 @@
|
||||
import torch as t
|
||||
from cdtools.models import CDIModel
|
||||
from cdtools.datasets import Ptycho2DDataset
|
||||
from cdtools import tools
|
||||
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__ = ['FastCCDPtycho']
|
||||
|
||||
|
||||
class FastCCDPtycho(CDIModel):
|
||||
|
||||
def __init__(self, wavelength, detector_geometry,
|
||||
probe_basis,
|
||||
probe_guess,
|
||||
obj_guess,
|
||||
detector_slice=None,
|
||||
threshold=None,
|
||||
surface_normal=t.tensor([0., 0., 1.], dtype=t.float32),
|
||||
min_translation=t.tensor([0, 0], dtype=t.float32),
|
||||
background=None,
|
||||
translation_offsets=None,
|
||||
mask=None,
|
||||
weights=None,
|
||||
translation_scale=1,
|
||||
saturation=None,
|
||||
probe_support=None,
|
||||
oversampling=1,
|
||||
fourier_probe=False,
|
||||
loss='amplitude mse',
|
||||
units='um',
|
||||
simulate_probe_translation=False,
|
||||
simulate_finite_pixels=False,
|
||||
):
|
||||
|
||||
super(FastCCDPtycho, self).__init__()
|
||||
self.wavelength = t.tensor(wavelength)
|
||||
self.detector_geometry = copy(detector_geometry)
|
||||
det_geo = self.detector_geometry
|
||||
if 'distance' in det_geo:
|
||||
det_geo['distance'] = t.tensor(det_geo['distance'], dtype=t.float32)
|
||||
if 'basis' in det_geo:
|
||||
det_geo['basis'] = t.tensor(det_geo['basis'], dtype=t.float32)
|
||||
if 'corner' in det_geo and det_geo['corner'] is not None:
|
||||
det_geo['corner'] = t.tensor(det_geo['corner'], dtype=t.float32)
|
||||
|
||||
self.min_translation = t.tensor(min_translation)
|
||||
|
||||
self.probe_basis = t.tensor(probe_basis)
|
||||
self.detector_slice = copy(detector_slice)
|
||||
self.surface_normal = t.tensor(surface_normal)
|
||||
|
||||
self.saturation = saturation
|
||||
self.units = units
|
||||
self.fourier_probe = fourier_probe
|
||||
|
||||
if mask is None:
|
||||
self.mask = mask
|
||||
else:
|
||||
self.mask = t.tensor(mask, dtype=t.bool)
|
||||
|
||||
probe_guess = t.tensor(probe_guess, dtype=t.complex64)
|
||||
obj_guess = t.tensor(obj_guess, dtype=t.complex64)
|
||||
|
||||
# We rescale the probe here so it learns at the same rate as the
|
||||
# object
|
||||
if probe_guess.dim() > 2:
|
||||
self.probe_norm = 1 * t.max(t.abs(probe_guess[0]))
|
||||
else:
|
||||
self.probe_norm = 1 * t.max(t.abs(probe_guess))
|
||||
|
||||
self.probe = t.nn.Parameter(probe_guess / self.probe_norm)
|
||||
self.obj = t.nn.Parameter(obj_guess)
|
||||
|
||||
if background is None:
|
||||
if detector_slice is not None:
|
||||
dummy_det = t.empty([s//oversampling
|
||||
for s in self.probe[0].shape])
|
||||
shape = dummy_det[self.detector_slice].shape
|
||||
#shape = self.probe[0][self.detector_slice].shape
|
||||
else:
|
||||
shape = [s//oversampling for s in self.probe[0]]
|
||||
background = 1e-6 * t.ones(shape, dtype=t.float32)
|
||||
|
||||
self.background = t.nn.Parameter(background)
|
||||
|
||||
if threshold is None:
|
||||
threshold = t.zeros([960,2], dtype=t.float32)
|
||||
|
||||
self.threshold = t.nn.Parameter(threshold)
|
||||
|
||||
if weights is None:
|
||||
self.weights = None
|
||||
else:
|
||||
# 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,
|
||||
dtype=t.float32))
|
||||
else:
|
||||
# Now this is a matrix of weights, so it needs to be complex
|
||||
self.weights = t.nn.Parameter(t.tensor(weights,
|
||||
dtype=t.complex64))
|
||||
|
||||
if translation_offsets is None:
|
||||
self.translation_offsets = None
|
||||
else:
|
||||
t_o = t.tensor(translation_offsets, dtype=t.float32)
|
||||
t_o = t_o / translation_scale
|
||||
self.translation_offsets = t.nn.Parameter(t_o)
|
||||
|
||||
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], dtype=t.bool)
|
||||
|
||||
self.oversampling = oversampling
|
||||
|
||||
self.simulate_probe_translation = simulate_probe_translation
|
||||
if simulate_probe_translation:
|
||||
Is = t.arange(self.probe.shape[-2], dtype=t.float32)
|
||||
Js = t.arange(self.probe.shape[-1], dtype=t.float32)
|
||||
Is, Js = t.meshgrid(Is/t.max(Is), Js/t.max(Js))
|
||||
self.I_phase = 2 * np.pi* Is
|
||||
self.J_phase = 2 * np.pi* Js
|
||||
|
||||
self.simulate_finite_pixels = simulate_finite_pixels
|
||||
|
||||
# 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,
|
||||
dm_rank=None,
|
||||
translation_scale=1,
|
||||
saturation=None,
|
||||
probe_support_radius=None,
|
||||
propagation_distance=None,
|
||||
scattering_mode=None,
|
||||
oversampling=1,
|
||||
auto_center=False,
|
||||
fourier_probe=False,
|
||||
loss='amplitude mse',
|
||||
units='um',
|
||||
simulate_probe_translation=False,
|
||||
simulate_finite_pixels=False,
|
||||
):
|
||||
|
||||
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')
|
||||
|
||||
# We include the *extras to make this work even with datasets, like
|
||||
# polarization dependent datasets, that might toss out extra inputs
|
||||
(indices, translations, *extras), 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,
|
||||
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(t.abs(probe))
|
||||
probe_stack = [0.01 * probe_max * t.rand(probe.shape, dtype=probe.dtype) for i in range(n_modes - 1)]
|
||||
|
||||
# For a Fourier space probe
|
||||
if fourier_probe:
|
||||
probe = tools.propagators.far_field(probe)
|
||||
|
||||
probe = t.stack([probe, ] + probe_stack)
|
||||
|
||||
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)
|
||||
|
||||
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, 'intensities') and dataset.intensities is not None:
|
||||
Ws *= (dataset.intensities.to(dtype=Ws.dtype)[:,...]
|
||||
/ t.mean(dataset.intensities))
|
||||
|
||||
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(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
|
||||
|
||||
return cls(wavelength, det_geo, probe_basis, probe, obj,
|
||||
detector_slice=det_slice,
|
||||
surface_normal=surface_normal,
|
||||
min_translation=min_translation,
|
||||
translation_offsets=translation_offsets,
|
||||
weights=Ws, mask=mask, background=background,
|
||||
translation_scale=translation_scale,
|
||||
saturation=saturation,
|
||||
probe_support=probe_support,
|
||||
fourier_probe=fourier_probe,
|
||||
oversampling=oversampling,
|
||||
loss=loss, units=units,
|
||||
simulate_probe_translation=simulate_probe_translation,
|
||||
simulate_finite_pixels=simulate_finite_pixels)
|
||||
|
||||
|
||||
def interaction(self, index, translations, *args):
|
||||
|
||||
# The *args is included so that this can work even when given, say,
|
||||
# a polarized ptycho dataset that might spit out more inputs.
|
||||
|
||||
# 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])
|
||||
|
||||
# This restricts the basis probes to stay within the probe support
|
||||
basis_prs = self.probe * self.probe_support[..., :, :]
|
||||
|
||||
# For a Fourier-space probe, we take an IFT
|
||||
if self.fourier_probe:
|
||||
basis_prs = tools.propagators.inverse_far_field(basis_prs)
|
||||
|
||||
# Now we construct the probes for each shot from the basis probes
|
||||
if self.weights is not None:
|
||||
Ws = self.weights[index]
|
||||
else:
|
||||
try:
|
||||
Ws = t.ones(len(index)) # I'm positive this introduced a bug
|
||||
except:
|
||||
Ws = 1
|
||||
|
||||
if self.weights is None or 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)
|
||||
|
||||
if self.simulate_probe_translation:
|
||||
det_pix_trans = tools.interactions.translations_to_pixel(
|
||||
self.detector_geometry['basis'],
|
||||
translations,
|
||||
surface_normal=self.surface_normal)
|
||||
|
||||
probe_masks = t.exp(1j* (det_pix_trans[:,0,None,None] *
|
||||
self.I_phase[None,...] +
|
||||
det_pix_trans[:,1,None,None] *
|
||||
self.J_phase[None,...]))
|
||||
prs = prs * probe_masks[...,None,:,:]
|
||||
|
||||
|
||||
# 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, pix_trans,
|
||||
shift_probe=True, multiple_modes=True)
|
||||
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):
|
||||
measured = tools.measurements.quadratic_background(
|
||||
wavefields,
|
||||
self.background,
|
||||
detector_slice=self.detector_slice,
|
||||
measurement=tools.measurements.incoherent_sum,
|
||||
saturation=self.saturation,
|
||||
oversampling=self.oversampling,
|
||||
simulate_finite_pixels=self.simulate_finite_pixels)
|
||||
thresholds = t.repeat_interleave(self.threshold, 480,dim=-1)
|
||||
return t.clamp(measured - thresholds, min=0.001)
|
||||
|
||||
|
||||
# Note: No "loss" function is defined here, because it is added
|
||||
# dynamically during object creation in __init__
|
||||
|
||||
def to(self, *args, **kwargs):
|
||||
super(FastCCDPtycho, self).to(*args, **kwargs)
|
||||
self.wavelength = self.wavelength.to(*args, **kwargs)
|
||||
# move the detector geometry too
|
||||
det_geo = self.detector_geometry
|
||||
if 'distance' in det_geo:
|
||||
det_geo['distance'] = det_geo['distance'].to(*args, **kwargs)
|
||||
if 'basis' in det_geo:
|
||||
det_geo['basis'] = det_geo['basis'].to(*args, **kwargs)
|
||||
if 'corner' in det_geo and det_geo['corner'] is not None:
|
||||
det_geo['corner'] = det_geo['corner'].to(*args, **kwargs)
|
||||
|
||||
if self.mask is not None:
|
||||
self.mask = self.mask.to(*args, **kwargs)
|
||||
|
||||
if self.simulate_probe_translation:
|
||||
self.I_phase = self.I_phase.to(*args, **kwargs)
|
||||
self.J_phase = self.J_phase.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.surface_normal = self.surface_normal.to(*args, **kwargs)
|
||||
|
||||
|
||||
def sim_to_dataset(self, args_list, calculation_width=None):
|
||||
# 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
|
||||
|
||||
data = []
|
||||
len(indices)
|
||||
if calculation_width is None:
|
||||
calculation_width = len(indices)
|
||||
index_chunks = [indices[i:i + calculation_width]
|
||||
for i in range(0, len(indices),
|
||||
calculation_width)]
|
||||
translation_chunks = [translations[i:i + calculation_width]
|
||||
for i in range(0, len(indices),
|
||||
calculation_width)]
|
||||
|
||||
|
||||
# Then we simulate the results
|
||||
data = [self.forward(idx, trans).detach()
|
||||
for idx, trans in zip(index_chunks, translation_chunks)]
|
||||
|
||||
data = t.cat(data, dim=0)
|
||||
# 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=t.float32, device=self.probe.device)
|
||||
if (hasattr(self, 'translation_offsets') and
|
||||
self.translation_offsets is not None):
|
||||
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
|
||||
else:
|
||||
return translations
|
||||
|
||||
|
||||
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):
|
||||
"""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
|
||||
|
||||
As a first step, we calculate the state of the light field across the
|
||||
full experiment, using the weight matrices and basis probes. Then, we
|
||||
use an SVD to update the basis probes so they form an eigenbasis of
|
||||
the implied density matrix for the full experiment.
|
||||
|
||||
Next, the weight matrices for each shot are recalculated so that the
|
||||
probes generated by weights * basis_probes for each shot are themselves
|
||||
an eigenbasis for that individual shot's density matrix.
|
||||
"""
|
||||
|
||||
# First we treat the incoherent but stable case, where the weights are
|
||||
# just one per-shot overall weight
|
||||
if self.weights.dim() == 1:
|
||||
probe = self.probe.detach().cpu().numpy()
|
||||
ortho_probes = analysis.orthogonalize_probes(self.probe.detach())
|
||||
self.probe.data = ortho_probes
|
||||
return
|
||||
|
||||
# What follows is for the unified OPRP and incoherent multi-mode model,
|
||||
# where each shot has it's own matrix of weights such that the probe
|
||||
# state for each shot is self.weights @ self.probe
|
||||
|
||||
# We concatenate all the weight matrices, to come up with a state
|
||||
# corresponding to the summed light field across all the exposures.
|
||||
# This state will have a large number of modes, but all built from
|
||||
# the same small number of basis modes
|
||||
all_weights = t.cat(t.unbind(self.weights.detach(), dim=0), dim=0)
|
||||
|
||||
# We generate the orthogonal probes based on this full-experiment
|
||||
# representation of the light field.
|
||||
ortho_probes, reexpressed_weights = \
|
||||
analysis.orthogonalize_probes(
|
||||
self.probe.detach(),
|
||||
weight_matrix=all_weights,
|
||||
return_reexpressed_weights=True
|
||||
)
|
||||
|
||||
# We just orthogonalized the incoherent sum of all the exposures
|
||||
# across the full experiment, so the output probes are normalized so
|
||||
# that their intensity matches the summed intensity across the full
|
||||
# experiment. We divide their amplitudes by the square root of the
|
||||
# number of shots so that we now have a set of probes corresponding
|
||||
# to the mean shot
|
||||
ortho_probes /= np.sqrt(self.weights.shape[0])
|
||||
reexpressed_weights *= np.sqrt(self.weights.shape[0])
|
||||
|
||||
# We now replace the shot-to-shot weights with the versions that have
|
||||
# been re-expressed in the new basis.
|
||||
new_weights = t.stack(t.split(reexpressed_weights,
|
||||
self.weights.shape[1]), dim=0)
|
||||
|
||||
# And we save it back to the model
|
||||
self.probe.data = ortho_probes.to(
|
||||
device=self.probe.device, dtype=self.probe.dtype)
|
||||
self.weights.data = new_weights.to(
|
||||
device=self.weights.device, dtype=self.weights.dtype)
|
||||
|
||||
# NOTE: I used to have this part as an option, with "tidy_each_frame",
|
||||
# because it took such a long time. Now that I've rewritten it properly,
|
||||
# it's quite fast and so I removed the kwarg because there's really
|
||||
# no situation where you woudn't want to do this.
|
||||
|
||||
# Now, we seek to edit the shot-to-shot weight matrices such that
|
||||
# self.weights[i] @ self.probes will be properly orthogonalized for
|
||||
# all i.
|
||||
|
||||
# All we need to know about the probes is that they are orthogonalized
|
||||
# and the intensity within each probe mode
|
||||
probe_sqrt_intensities = t.linalg.norm(self.probe.data, dim=(-2,-1))
|
||||
|
||||
# This does a super fast batched computation
|
||||
U, S, Vh = t.linalg.svd(self.weights.data * probe_sqrt_intensities,
|
||||
full_matrices=False)
|
||||
|
||||
# We discard the U matrix and re-multiply S & Vh
|
||||
self.weights.data = S[:,:,None] * (Vh / probe_sqrt_intensities)
|
||||
|
||||
|
||||
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=-3)
|
||||
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 = [
|
||||
('',
|
||||
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 Fourier Space Amplitudes',
|
||||
lambda self, fig: p.plot_amplitude(self.probe if self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig)),
|
||||
('Basis Probe Fourier Space Phases',
|
||||
lambda self, fig: p.plot_phase(self.probe if self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig)),
|
||||
('Basis Probe Real Space Amplitudes',
|
||||
lambda self, fig: p.plot_amplitude(self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig, basis=self.probe_basis, units=self.units)),
|
||||
('Basis Probe Real Space Phases',
|
||||
lambda self, fig: p.plot_phase(self.probe if not self.fourier_probe else tools.propagators.inverse_far_field(self.probe), fig=fig, basis=self.probe_basis, units=self.units)),
|
||||
('Average Weight Matrix Amplitudes',
|
||||
lambda self, fig: p.plot_amplitude(
|
||||
np.nanmean(np.abs(self.weights.data.cpu().numpy()), axis=0),
|
||||
fig=fig),
|
||||
lambda self: len(self.weights.shape) >= 2),
|
||||
('% of Power in Top Mode',
|
||||
lambda self, fig, dataset: p.plot_nanomap(
|
||||
self.corrected_translations(dataset),
|
||||
100 * t.stack([
|
||||
analysis.calc_mode_power_fractions(
|
||||
self.probe.data,
|
||||
weight_matrix=self.weights.data[i])[0]
|
||||
for i in range(self.weights.shape[0])
|
||||
], dim=0),
|
||||
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, units=self.units)),
|
||||
('Object Phase',
|
||||
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, units=self.units)),
|
||||
('Background',
|
||||
lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2)),
|
||||
('Thresholds',
|
||||
lambda self, fig: p.plot_real(t.repeat_interleave(self.threshold, 480,dim=-1), fig=fig))
|
||||
]
|
||||
|
||||
# def plot_errors(self, dataset):
|
||||
|
||||
|
||||
|
||||
def save_results(self, dataset):
|
||||
thresholds = self.thresholds.detach().cpu().numpy()
|
||||
basis = self.probe_basis.detach().cpu().numpy()
|
||||
translations = self.corrected_translations(dataset).detach().cpu().numpy()
|
||||
probe = self.probe.detach().cpu().numpy()
|
||||
probe = probe * self.probe_norm.detach().cpu().numpy()
|
||||
obj = self.obj.detach().cpu().numpy()
|
||||
background = self.background.detach().cpu().numpy()**2
|
||||
weights = self.weights.detach().cpu().numpy()
|
||||
oversampling = self.oversampling
|
||||
|
||||
return {'basis': basis, 'translation': translations,
|
||||
'probe': probe, 'obj': obj,
|
||||
'background': background,
|
||||
'oversampling': oversampling,
|
||||
'weights': weights, 'thresholds':thresholds}
|
||||
@@ -1,354 +0,0 @@
|
||||
import torch as t
|
||||
from cdtools.models import CDIModel
|
||||
from cdtools import tools
|
||||
from cdtools.tools import plotting as p
|
||||
from cdtools.tools.interactions import RPI_interaction
|
||||
from cdtools.tools import initializers
|
||||
from scipy.ndimage import binary_dilation
|
||||
import numpy as np
|
||||
from copy import copy
|
||||
|
||||
__all__ = ['MultimodeRPI']
|
||||
|
||||
class MultimodeRPI(CDIModel):
|
||||
|
||||
@property
|
||||
def obj(self):
|
||||
return t.complex(self.obj_real, self.obj_imag)
|
||||
|
||||
@property
|
||||
def weights(self):
|
||||
ws = t.complex(self.weights_real, self.weights_imag)
|
||||
return ws / 10# / self.obj_real.size().numel()
|
||||
|
||||
def __init__(self, wavelength, detector_geometry, probe_basis,
|
||||
probe, obj_guess, detector_slice=None,
|
||||
background=None, mask=None, saturation=None,
|
||||
obj_support=None, oversampling=1, weight_matrix=False):
|
||||
|
||||
super(MultimodeRPI, 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.probe_basis = t.tensor(probe_basis)
|
||||
|
||||
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
|
||||
|
||||
# Maybe something to include in a bit
|
||||
# self.surface_normal = t.tensor(surface_normal)
|
||||
|
||||
self.saturation = saturation
|
||||
|
||||
if mask is None:
|
||||
self.mask = mask
|
||||
else:
|
||||
self.mask = t.tensor(mask, dtype=t.bool)
|
||||
|
||||
|
||||
self.probe = t.tensor(probe, dtype=t.complex64)
|
||||
|
||||
obj_guess = t.tensor(obj_guess, dtype=t.complex64)
|
||||
|
||||
self.obj_real = t.nn.Parameter(obj_guess.real)
|
||||
self.obj_imag = t.nn.Parameter(obj_guess.imag)
|
||||
|
||||
self.weights_real = t.nn.Parameter(t.eye(probe.shape[0])* 10)# * self.obj_real.size().numel())
|
||||
self.weights_imag = t.nn.Parameter(t.zeros(probe.shape[0]))
|
||||
|
||||
if not weight_matrix:
|
||||
self.weights_real.requires_grad=False
|
||||
self.weights_imag.requires_grad=False
|
||||
|
||||
# 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:
|
||||
background = 1e-6 * t.ones(
|
||||
self.probe[0][self.detector_slice].shape,
|
||||
dtype=t.float32)
|
||||
else:
|
||||
background = 1e-6 * t.ones(self.probe[0].shape,
|
||||
dtype=t.float32)
|
||||
|
||||
self.background = t.tensor(background, dtype=t.float32)
|
||||
|
||||
if obj_support is not None:
|
||||
self.obj_support = obj_support
|
||||
self.obj.data = self.obj * obj_support[None, ...]
|
||||
else:
|
||||
self.obj_support = t.ones_like(self.obj[0, ...])
|
||||
|
||||
self.oversampling = oversampling
|
||||
|
||||
|
||||
@classmethod
|
||||
def from_dataset(cls, dataset, probe, obj_size=None, background=None, mask=None, padding=0, n_modes=1, saturation=None, scattering_mode=None, oversampling=1, auto_center=False, initialization='random', weight_matrix=False, probe_threshold=0):
|
||||
raise NotImplementedError()
|
||||
|
||||
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')
|
||||
# We only need the patterns here, not the inputs associated with them.
|
||||
_, 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,
|
||||
oversampling=oversampling)
|
||||
|
||||
if not isinstance(probe,t.Tensor):
|
||||
probe = t.as_tensor(probe)
|
||||
|
||||
# Potentially need all of this orientation stuff later
|
||||
|
||||
#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)
|
||||
|
||||
|
||||
if background is None and hasattr(dataset, 'background') \
|
||||
and dataset.background is not None:
|
||||
background = t.sqrt(dataset.background)
|
||||
elif background is not None:
|
||||
background = t.sqrt(t.Tensor(background).to(dtype=t.float32))
|
||||
|
||||
det_geo = dataset.detector_geometry
|
||||
|
||||
# If no mask is given, but one exists in the dataset, load it.
|
||||
if mask is None and hasattr(dataset, 'mask') \
|
||||
and dataset.mask is not None:
|
||||
mask = dataset.mask.to(t.bool)
|
||||
|
||||
# Now we initialize the object
|
||||
if obj_size is None:
|
||||
# This is a standard size for a well-matched probe and detector
|
||||
obj_size = (np.array(probe_shape) // 2).astype(int)
|
||||
|
||||
if initialization.lower().strip() == 'random':
|
||||
# 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(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(
|
||||
patterns[0], probe, obj_size, mask=mask,
|
||||
background=background**2, n_modes=n_modes)
|
||||
else:
|
||||
obj_guess = initializers.RPI_spectral_init(
|
||||
patterns[0], probe, obj_size, mask=mask,
|
||||
n_modes=n_modes)
|
||||
|
||||
else:
|
||||
raise KeyError('Initialization "' + str(initialization) + \
|
||||
'" invalid - use "spectral" or "random"')
|
||||
|
||||
probe_intensity = t.sqrt(t.sum(t.abs(probe)**2,axis=0))
|
||||
probe_fft = tools.propagators.far_field(probe_intensity)
|
||||
pad0l = (probe.shape[-2] - obj_size[-2])//2
|
||||
pad0r = probe.shape[-2] - obj_size[-2] - pad0l
|
||||
pad1l = (probe.shape[-1] - obj_size[-1])//2
|
||||
pad1r = probe.shape[-1] - obj_size[-1] - pad1l
|
||||
probe_lr_fft = probe_fft[pad0l:-pad0r,pad1l:-pad1r]
|
||||
probe_lr = t.abs(tools.propagators.inverse_far_field(probe_lr_fft))
|
||||
|
||||
obj_support = probe_lr > t.max(probe_lr) * probe_threshold
|
||||
obj_support = t.as_tensor(binary_dilation(obj_support))
|
||||
|
||||
return cls(wavelength, det_geo, probe_basis,
|
||||
probe, obj_guess, detector_slice=det_slice,
|
||||
background=background, mask=mask, saturation=saturation,
|
||||
obj_support=obj_support, oversampling=oversampling,
|
||||
weight_matrix=weight_matrix)
|
||||
|
||||
|
||||
def random_init(self, pattern):
|
||||
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):
|
||||
if self.background is not None:
|
||||
self.obj.data = initializers.RPI_spectral_init(
|
||||
pattern, self.probe, self.obj.shape[-3:-1], mask=self.mask,
|
||||
background=self.background**2, n_modes=self.obj.shape[0]).to(
|
||||
dtype=self.obj.dtype, device=self.obj.device)
|
||||
else:
|
||||
self.obj.data = initializers.RPI_spectral_init(
|
||||
pattern, self.probe, self.obj.shape[-3:-1], mask=self.mask,
|
||||
n_modes=self.obj.shape[0]).to(
|
||||
dtype=self.obj.dtype, device=self.obj.device)
|
||||
|
||||
# Needs work
|
||||
def interaction(self, index, *args):
|
||||
# including *args allows this to work with all sorts of datasets
|
||||
# that might include other information in with the index in their
|
||||
# "input" parameters (such as translations for a ptychography dataset).
|
||||
# This makes it seamless to use such a dataset even though those
|
||||
# extra arguments will not be used.
|
||||
|
||||
|
||||
all_exit_waves = []
|
||||
|
||||
# Mix the probes with the weight matrix
|
||||
prs = t.sum(self.weights[..., None, None] * self.probe, axis=-3)
|
||||
|
||||
for i in range(self.probe.shape[0]):
|
||||
pr = prs[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 * self.obj[i])
|
||||
all_exit_waves.append(exit_waves.unsqueeze(0))
|
||||
|
||||
# This creates a bunch of modes generated from all possible combos
|
||||
# of the probe and object modes all strung out along the first index
|
||||
|
||||
output = t.cat(all_exit_waves)
|
||||
|
||||
# If we have multiple indexes input, we unsqueeze and repeat the stack
|
||||
# of wavefields enough times to simulate each requested index. This
|
||||
# seems silly, but it enables (for example) one to do a reconstruction
|
||||
# from a set of diffraction patterns that are all known to be from the
|
||||
# same object.
|
||||
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(0).repeat(1,len(index),1,1,1)
|
||||
except TypeError:
|
||||
pass
|
||||
return output
|
||||
|
||||
|
||||
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):
|
||||
# Here I'm taking advantage of an undocumented feature in the
|
||||
# incoherent_sum measurement function where it will work with
|
||||
# a 4D wavefield array as well as a 5D array.
|
||||
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 regularizer(self, factors):
|
||||
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(MultimodeRPI, 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.probe = self.probe.to(*args,**kwargs)
|
||||
self.probe_basis = self.probe_basis.to(*args,**kwargs)
|
||||
self.obj_basis = self.obj_basis.to(*args,**kwargs)
|
||||
self.obj_support = self.obj_support.to(*args,**kwargs)
|
||||
self.background = self.background.to(*args, **kwargs)
|
||||
|
||||
# Maybe include in a bit
|
||||
#self.surface_normal = self.surface_normal.to(*args, **kwargs)
|
||||
|
||||
def sim_to_dataset(self, args_list):
|
||||
raise NotImplementedError('No sim to dataset yet, sorry!')
|
||||
|
||||
plot_list = [
|
||||
('Root Sum Squared Amplitude of all Probes',
|
||||
lambda self, fig: p.plot_amplitude(
|
||||
np.sqrt(np.sum((t.abs(t.sum(self.weights[..., None, None].detach() * self.probe, axis=-3))**2).cpu().numpy(),axis=0)),
|
||||
fig=fig, basis=self.probe_basis)),
|
||||
('Object Amplitudes',
|
||||
lambda self, fig: p.plot_amplitude(self.obj, fig=fig,
|
||||
basis=self.obj_basis)),
|
||||
('Object Phases',
|
||||
lambda self, fig: p.plot_phase(self.obj, fig=fig,
|
||||
basis=self.obj_basis))
|
||||
]
|
||||
|
||||
|
||||
def save_results(self, dataset=None, full_obj=False):
|
||||
# dataset is set as a kwarg here because it isn't needed, but the
|
||||
# common pattern is to pass a dataset. This makes it okay if one
|
||||
# continues to use that standard pattern
|
||||
probe_basis = self.probe_basis.detach().cpu().numpy()
|
||||
obj_basis = self.obj_basis.detach().cpu().numpy()
|
||||
probe = self.probe.detach().cpu().numpy()
|
||||
# Provide the option to save out the subdominant objects or
|
||||
# just the dominant one
|
||||
if full_obj:
|
||||
obj = self.obj.detach().cpu().numpy()
|
||||
else:
|
||||
obj = self.obj[0].detach().cpu().numpy()
|
||||
background = self.background.detach().cpu().numpy()**2
|
||||
|
||||
return {'probe_basis': probe_basis, 'obj_basis': obj_basis,
|
||||
'probe': probe,'obj': obj,
|
||||
'background': background}
|
||||
|
||||
@@ -1,605 +0,0 @@
|
||||
import torch as t
|
||||
from cdtools.models import CDIModel
|
||||
from cdtools.datasets import Ptycho2DDataset
|
||||
from cdtools import tools
|
||||
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']
|
||||
|
||||
class PolarizationSweptPtycho(CDIModel):
|
||||
|
||||
def __init__(self, wavelength, detector_geometry,
|
||||
probe_basis,
|
||||
probe_guess,
|
||||
obj_guess,
|
||||
polarization_states,
|
||||
detector_slice=None,
|
||||
surface_normal=t.tensor([0., 0., 1.], dtype=t.float32),
|
||||
min_translation=t.tensor([0, 0], dtype=t.float32),
|
||||
background=None,
|
||||
translation_offsets=None,
|
||||
mask=None,
|
||||
weights=None,
|
||||
translation_scale=1,
|
||||
saturation=None,
|
||||
probe_support=None,
|
||||
oversampling=1,
|
||||
fourier_probe=False,
|
||||
loss='amplitude mse',
|
||||
units='um',
|
||||
simulate_probe_translation=False,
|
||||
simulate_finite_pixels=False,
|
||||
dtype=t.float32,
|
||||
obj_view_crop=0
|
||||
):
|
||||
|
||||
super(PolarizationSweptPtycho, self).__init__()
|
||||
self.register_buffer('wavelength',
|
||||
t.tensor(wavelength, dtype=dtype))
|
||||
self.store_detector_geometry(detector_geometry,
|
||||
dtype=dtype)
|
||||
|
||||
self.register_buffer('min_translation',
|
||||
t.tensor(min_translation, dtype=dtype))
|
||||
|
||||
self.register_buffer('probe_basis',
|
||||
t.tensor(probe_basis, dtype=dtype))
|
||||
|
||||
self.detector_slice = copy(detector_slice)
|
||||
self.register_buffer('surface_normal',
|
||||
t.tensor(surface_normal, dtype=dtype))
|
||||
if saturation is None:
|
||||
self.saturation = None
|
||||
else:
|
||||
self.register_buffer('saturation',
|
||||
t.tensor(saturation, dtype=dtype))
|
||||
# Not sure how to make this a buffer...
|
||||
self.units = units
|
||||
|
||||
self.fourier_probe = fourier_probe
|
||||
|
||||
if mask is None:
|
||||
self.mask = None
|
||||
else:
|
||||
self.register_buffer('mask',
|
||||
t.tensor(mask, dtype=t.bool))
|
||||
|
||||
probe_guess = t.tensor(probe_guess, dtype=t.complex64)
|
||||
obj_guess = t.tensor(obj_guess, dtype=t.complex64)
|
||||
|
||||
# We rescale the probe here so it learns at the same rate as the
|
||||
# object
|
||||
if probe_guess.dim() > 2:
|
||||
probe_norm = 1 * t.max(t.abs(probe_guess[0]))
|
||||
else:
|
||||
probe_norm = 1 * t.max(t.abs(probe_guess))
|
||||
self.register_buffer('probe_norm', probe_norm.to(dtype))
|
||||
|
||||
self.probe = t.nn.Parameter(probe_guess / self.probe_norm)
|
||||
self.obj = t.nn.Parameter(obj_guess)
|
||||
|
||||
|
||||
self.obj_view_slice = np.s_[obj_view_crop:-obj_view_crop,
|
||||
obj_view_crop:-obj_view_crop]
|
||||
|
||||
if background is None:
|
||||
if detector_slice is not None:
|
||||
dummy_det = t.empty([s//oversampling
|
||||
for s in self.probe.shape[-2:]])
|
||||
shape = dummy_det[self.detector_slice].shape
|
||||
#shape = self.probe[0][self.detector_slice].shape
|
||||
else:
|
||||
shape = [s//oversampling for s in self.probe.shape[-2:]]
|
||||
background = 1e-6 * t.ones(shape, dtype=t.float32)
|
||||
|
||||
self.background = t.nn.Parameter(background)
|
||||
|
||||
if weights is None:
|
||||
self.weights = None
|
||||
else:
|
||||
self.weights = t.nn.Parameter(t.tensor(weights,
|
||||
dtype=t.float32))
|
||||
|
||||
if translation_offsets is None:
|
||||
self.translation_offsets = None
|
||||
else:
|
||||
t_o = t.tensor(translation_offsets, dtype=t.float32)
|
||||
t_o = t_o / translation_scale
|
||||
self.translation_offsets = t.nn.Parameter(t_o)
|
||||
|
||||
self.register_buffer('translation_scale',
|
||||
t.tensor(translation_scale, dtype=dtype))
|
||||
|
||||
if probe_support is None:
|
||||
probe_support = t.ones(self.probe.shape[-2:], dtype=t.bool)
|
||||
self.register_buffer('probe_support',
|
||||
t.tensor(probe_support, dtype=t.bool))
|
||||
|
||||
self.oversampling = oversampling
|
||||
|
||||
self.simulate_probe_translation = simulate_probe_translation
|
||||
|
||||
if simulate_probe_translation:
|
||||
Is = t.arange(self.probe.shape[-2], dtype=dtype)
|
||||
Js = t.arange(self.probe.shape[-1], dtype=dtype)
|
||||
Is, Js = t.meshgrid(Is/t.max(Is), Js/t.max(Js))
|
||||
|
||||
I_phase = 2 * np.pi* Is * self.oversampling
|
||||
J_phase = 2 * np.pi* Js * self.oversampling
|
||||
self.register_buffer('I_phase', I_phase)
|
||||
self.register_buffer('J_phase', J_phase)
|
||||
|
||||
|
||||
self.simulate_finite_pixels = simulate_finite_pixels
|
||||
|
||||
self.polarization_states = t.nn.Parameter(
|
||||
t.tensor(polarization_states, dtype=t.complex64))
|
||||
# by default, don't optimize this, but I think it might be
|
||||
# interesting to try it because the polarization states
|
||||
# are not very pure
|
||||
self.polarization_states.requires_grad = False
|
||||
|
||||
|
||||
# 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,
|
||||
probe_fourier_crop=None,
|
||||
propagation_distance=None,
|
||||
scattering_mode=None,
|
||||
oversampling=1,
|
||||
auto_center=False,
|
||||
fourier_probe=False,
|
||||
loss='amplitude mse',
|
||||
units='um',
|
||||
simulate_probe_translation=False,
|
||||
simulate_finite_pixels=False,
|
||||
obj_view_crop=None,
|
||||
obj_padding=200
|
||||
):
|
||||
|
||||
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')
|
||||
|
||||
# We include the *extras to make this work even with datasets, like
|
||||
# polarization dependent datasets, that might toss out extra inputs
|
||||
((indices, translations, *extras), patterns) = dataset[:]
|
||||
polarization_states = dataset.polarization_states
|
||||
|
||||
dataset.get_as(*get_as_args[0], **get_as_args[1])
|
||||
|
||||
# Set to none to avoid issues with things outside the detector
|
||||
if auto_center:
|
||||
center = tools.image_processing.centroid(t.sum(patterns, dim=0))
|
||||
else:
|
||||
center = None
|
||||
|
||||
# Then, generate the probe geometry from the dataset
|
||||
ewg = tools.initializers.exit_wave_geometry
|
||||
probe_basis, probe_shape, det_slice = ewg(det_basis,
|
||||
det_shape,
|
||||
wavelength,
|
||||
distance,
|
||||
center=center,
|
||||
padding=padding,
|
||||
oversampling=oversampling)
|
||||
|
||||
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=obj_padding)
|
||||
|
||||
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)
|
||||
|
||||
if probe_fourier_crop is not None:
|
||||
probe = tools.propagators.far_field(probe)
|
||||
probe = probe[probe_fourier_crop:-probe_fourier_crop,
|
||||
probe_fourier_crop:-probe_fourier_crop]
|
||||
probe = tools.propagators.inverse_far_field(probe)
|
||||
|
||||
# Now we initialize all the subdominant probe modes
|
||||
probe_max = t.max(t.abs(probe))
|
||||
probe_stack = [0.01 * probe_max * t.rand(probe.shape, dtype=probe.dtype) for i in range(n_modes - 1)]
|
||||
|
||||
# For a Fourier space probe
|
||||
if fourier_probe:
|
||||
probe = tools.propagators.far_field(probe)
|
||||
|
||||
probe = t.stack([probe, ] + probe_stack)
|
||||
# The probe gets one extra dimension, so we can simulate one probe
|
||||
# per polarization state
|
||||
n_states = polarization_states.shape[0]
|
||||
probe = t.stack([probe] * n_states, dim=0)
|
||||
|
||||
# Looks like an identity matrix
|
||||
obj_base = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5))
|
||||
obj_top = t.stack([obj_base, obj_base*0], dim=0)
|
||||
obj_bottom = t.stack([obj_base*0, obj_base], dim=0)
|
||||
obj = t.stack([obj_top, obj_bottom], dim=0)
|
||||
|
||||
if obj_view_crop is None:
|
||||
obj_view_crop = min(probe.shape[-2], probe.shape[-1]) // 2
|
||||
if obj_view_crop < 0:
|
||||
obj_view_crop += min(probe.shape[-2], probe.shape[-1]) // 2
|
||||
obj_view_crop += obj_padding
|
||||
|
||||
det_geo = dataset.detector_geometry
|
||||
|
||||
translation_offsets = 0 * (t.rand((len(dataset), 2)) - 0.5)
|
||||
|
||||
Ws = t.ones(len(dataset))
|
||||
|
||||
if hasattr(dataset, 'intensities') and dataset.intensities is not None:
|
||||
Ws *= (dataset.intensities.to(dtype=Ws.dtype)[:,...]
|
||||
/ t.mean(dataset.intensities))
|
||||
|
||||
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(probe.shape[-2:], 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
|
||||
|
||||
return cls(wavelength, det_geo, probe_basis, probe, obj,
|
||||
polarization_states=polarization_states,
|
||||
detector_slice=det_slice,
|
||||
surface_normal=surface_normal,
|
||||
min_translation=min_translation,
|
||||
translation_offsets=translation_offsets,
|
||||
weights=Ws, mask=mask, background=background,
|
||||
translation_scale=translation_scale,
|
||||
saturation=saturation,
|
||||
probe_support=probe_support,
|
||||
fourier_probe=fourier_probe,
|
||||
oversampling=oversampling,
|
||||
loss=loss, units=units,
|
||||
simulate_probe_translation=simulate_probe_translation,
|
||||
simulate_finite_pixels=simulate_finite_pixels,
|
||||
obj_view_crop=obj_view_crop)
|
||||
|
||||
|
||||
def interaction(self, index, translations, polarization_indices, *args):
|
||||
|
||||
# The *args is included so that this can work even when given, say,
|
||||
# a polarized ptycho dataset that might spit out more inputs.
|
||||
|
||||
# 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])
|
||||
|
||||
# This restricts the basis probes to stay within the probe support
|
||||
basis_prs = self.probe * self.probe_support[..., :, :]
|
||||
# Now we expand the polarization states explicitly.
|
||||
# pol_basis_prs has dimensions n_states x 2 x n_modes x m x l
|
||||
pol_basis_prs = (
|
||||
self.polarization_states[...,None,None,None]
|
||||
* basis_prs[...,None,:,:,:])
|
||||
# For a Fourier-space probe, we take an IFT
|
||||
if self.fourier_probe:
|
||||
basis_prs = tools.propagators.inverse_far_field(basis_prs)
|
||||
|
||||
# So, first we get the appropriate probe for each shot.
|
||||
# prs is now n_frames x 2 x n_modes x m x l
|
||||
prs = pol_basis_prs[polarization_indices]
|
||||
|
||||
# Now we construct the probes for each shot from the basis probes
|
||||
if self.weights is not None:
|
||||
Ws = self.weights[index]
|
||||
# And then we multiply by the weights, along the 0th dimension
|
||||
prs = Ws[..., None, None, None, None] * prs#basis_prs
|
||||
|
||||
if self.simulate_probe_translation:
|
||||
det_pix_trans = tools.interactions.translations_to_pixel(
|
||||
self.det_basis,
|
||||
translations,
|
||||
surface_normal=self.surface_normal)
|
||||
|
||||
probe_masks = t.exp(1j* (det_pix_trans[:,0,None,None] *
|
||||
self.I_phase[None,...] +
|
||||
det_pix_trans[:,1,None,None] *
|
||||
self.J_phase[None,...]))
|
||||
prs = prs * probe_masks[...,None,None,:,:]
|
||||
|
||||
# We automatically rescale the probe to match the background size,
|
||||
# which allows us to do stuff like let the object be super-resolution,
|
||||
# while restricting the probe to the detector resolution but still
|
||||
# doing an explicit real-space limitation of the probe
|
||||
padding = [self.oversampling * self.background.shape[-2] - prs.shape[-2],
|
||||
self.oversampling * self.background.shape[-1] - prs.shape[-1]]
|
||||
|
||||
if any([p != 0 for p in padding]): # For probe_fourier_crop != 0.
|
||||
padding = [padding[-1]//2, padding[-1]-padding[-1]//2,
|
||||
padding[-2]//2, padding[-2]-padding[-2]//2]
|
||||
prs = tools.propagators.far_field(prs)
|
||||
prs = t.nn.functional.pad(prs, padding)
|
||||
prs = tools.propagators.inverse_far_field(prs)
|
||||
# Now we actually do the interaction, using the sinc subpixel
|
||||
# translation model as per usual
|
||||
|
||||
# After the object, no point in treating the polarization modes
|
||||
# any differently from the object/probe modes
|
||||
# So, this stacks up all the output modes in lexicographic order
|
||||
exit_waves = t.cat(
|
||||
[self.probe_norm * tools.interactions.ptycho_2D_sinc(
|
||||
prs[...,0,:,:,:], self.obj[0,0], pix_trans,
|
||||
shift_probe=True, multiple_modes=True),
|
||||
self.probe_norm * tools.interactions.ptycho_2D_sinc(
|
||||
prs[...,0,:,:,:], self.obj[1,0], pix_trans,
|
||||
shift_probe=True, multiple_modes=True)],
|
||||
dim=1)
|
||||
exit_waves = exit_waves + t.cat(
|
||||
[self.probe_norm * tools.interactions.ptycho_2D_sinc(
|
||||
prs[...,1,:,:,:], self.obj[0,1], pix_trans,
|
||||
shift_probe=True, multiple_modes=True),
|
||||
self.probe_norm * tools.interactions.ptycho_2D_sinc(
|
||||
prs[...,1,:,:,:], self.obj[1,1], pix_trans,
|
||||
shift_probe=True, multiple_modes=True)],
|
||||
dim=1)
|
||||
|
||||
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,
|
||||
measurement=tools.measurements.incoherent_sum,
|
||||
saturation=self.saturation,
|
||||
oversampling=self.oversampling,
|
||||
simulate_finite_pixels=self.simulate_finite_pixels)
|
||||
|
||||
|
||||
# Note: No "loss" function is defined here, because it is added
|
||||
# dynamically during object creation in __init__
|
||||
|
||||
def sim_to_dataset(self, args_list, calculation_width=None):
|
||||
# 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}
|
||||
|
||||
|
||||
mask = self.mask
|
||||
wavelength = self.wavelength
|
||||
indices, translations = args_list
|
||||
|
||||
data = []
|
||||
len(indices)
|
||||
if calculation_width is None:
|
||||
calculation_width = len(indices)
|
||||
index_chunks = [indices[i:i + calculation_width]
|
||||
for i in range(0, len(indices),
|
||||
calculation_width)]
|
||||
translation_chunks = [translations[i:i + calculation_width]
|
||||
for i in range(0, len(indices),
|
||||
calculation_width)]
|
||||
|
||||
|
||||
# Then we simulate the results
|
||||
data = [self.forward(idx, trans).detach()
|
||||
for idx, trans in zip(index_chunks, translation_chunks)]
|
||||
|
||||
data = t.cat(data, dim=0)
|
||||
# And finally, we make the dataset
|
||||
return Ptycho2DDataset(
|
||||
translations, data,
|
||||
entry_info=entry_info,
|
||||
sample_info=sample_info,
|
||||
wavelength=wavelength,
|
||||
detector_geometry=self.get_detector_geometry(),
|
||||
mask=mask)
|
||||
|
||||
|
||||
def corrected_translations(self, dataset):
|
||||
translations = dataset.translations.to(
|
||||
dtype=t.float32, device=self.probe.device)
|
||||
if (hasattr(self, 'translation_offsets') and
|
||||
self.translation_offsets is not None):
|
||||
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
|
||||
else:
|
||||
return translations
|
||||
|
||||
|
||||
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
|
||||
"""
|
||||
# NOTE: untested, edited when ortho_probes was updated
|
||||
for idx in range(self.probe.shape[0]):
|
||||
oself.probe.data[idx] = analysis.orthogonalize_probes(
|
||||
self.probe.data[idx])
|
||||
|
||||
|
||||
|
||||
plot_list = [
|
||||
('Basis Probe Fourier Space Amplitudes',
|
||||
lambda self, fig: p.plot_amplitude(
|
||||
(self.probe if self.fourier_probe
|
||||
else tools.propagators.inverse_far_field(self.probe)),
|
||||
fig=fig)),
|
||||
('Basis Probe Fourier Space Phases',
|
||||
lambda self, fig: p.plot_phase(
|
||||
(self.probe if self.fourier_probe
|
||||
else tools.propagators.inverse_far_field(self.probe))
|
||||
, fig=fig)),
|
||||
('Basis Probe Real Space Amplitudes',
|
||||
lambda self, fig: p.plot_amplitude(
|
||||
(self.probe if not self.fourier_probe
|
||||
else tools.propagators.inverse_far_field(self.probe)),
|
||||
fig=fig,
|
||||
basis=self.probe_basis,
|
||||
units=self.units)),
|
||||
('Basis Probe Real Space Phases',
|
||||
lambda self, fig: p.plot_phase(
|
||||
(self.probe if not self.fourier_probe
|
||||
else tools.propagators.inverse_far_field(self.probe)),
|
||||
fig=fig,
|
||||
basis=self.probe_basis,
|
||||
units=self.units)),
|
||||
('Average Density Matrix Amplitudes',
|
||||
lambda self, fig: p.plot_amplitude(
|
||||
np.nanmean(np.abs(self.get_rhos()), axis=0),
|
||||
fig=fig),
|
||||
lambda self: len(self.weights.shape) >= 2),
|
||||
('% 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[np.s_[:,:] + self.obj_view_slice],
|
||||
fig=fig,
|
||||
basis=self.probe_basis,
|
||||
units=self.units)),
|
||||
('Object Phase',
|
||||
lambda self, fig: p.plot_phase(
|
||||
self.obj[np.s_[:,:] + self.obj_view_slice],
|
||||
fig=fig,
|
||||
basis=self.probe_basis,
|
||||
units=self.units)),
|
||||
('Corrected Translations',
|
||||
lambda self, fig, dataset: p.plot_translations(self.corrected_translations(dataset), fig=fig, units=self.units)),
|
||||
('Background',
|
||||
lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2))
|
||||
]
|
||||
|
||||
# def plot_errors(self, dataset):
|
||||
|
||||
|
||||
|
||||
def save_results(self, dataset):
|
||||
basis = self.probe_basis.detach().cpu().numpy()
|
||||
translations = self.corrected_translations(dataset).detach().cpu().numpy()
|
||||
probe = self.probe.detach().cpu().numpy()
|
||||
probe = probe * self.probe_norm.detach().cpu().numpy()
|
||||
obj = self.obj.detach().cpu().numpy()
|
||||
background = self.background.detach().cpu().numpy()**2
|
||||
weights = self.weights.detach().cpu().numpy()
|
||||
oversampling = self.oversampling
|
||||
wavelength = self.wavelength.cpu().numpy()
|
||||
|
||||
return {'basis': basis, 'translation': translations,
|
||||
'probe': probe, 'obj': obj,
|
||||
'background': background,
|
||||
'oversampling': oversampling,
|
||||
'weights': weights, 'wavelength': wavelength}
|
||||
@@ -1,556 +0,0 @@
|
||||
import torch as t
|
||||
from cdtools.models import CDIModel, FancyPtycho
|
||||
from cdtools.datasets import Ptycho2DDataset
|
||||
from cdtools import tools
|
||||
from cdtools.tools import plotting as p
|
||||
# from cdtools.tools import polarized_plotting as pp
|
||||
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
|
||||
from cdtools.tools import polarization
|
||||
|
||||
__all__ = ['PolarizedFancyPtycho']
|
||||
|
||||
class PolarizedFancyPtycho(FancyPtycho):
|
||||
|
||||
def __init__(self, wavelength, detector_geometry,
|
||||
probe_basis,
|
||||
probe_guess, obj_guess, polarizer, analyzer,
|
||||
detector_slice=None,
|
||||
surface_normal=np.array([0.,0.,1.]),
|
||||
min_translation = t.Tensor([0,0]),
|
||||
background = None, translation_offsets=None,
|
||||
polarizer_offsets=None, analyzer_offsets=None,
|
||||
polarizer_scale=1, analyzer_scale=1, mask=None,
|
||||
weights = None, translation_scale = 1, saturation=None,
|
||||
probe_support = None, oversampling=1,
|
||||
loss='amplitude mse',units='um'):
|
||||
|
||||
super(PolarizedFancyPtycho, self).__init__(wavelength, detector_geometry,
|
||||
probe_basis,
|
||||
probe_guess, obj_guess,
|
||||
detector_slice=None,
|
||||
surface_normal=np.array([0.,0.,1.]),
|
||||
min_translation = t.Tensor([0,0]),
|
||||
background = None, translation_offsets=None, mask=None,
|
||||
weights = weights, translation_scale = 1, saturation=None,
|
||||
probe_support = None, oversampling=1,
|
||||
loss='amplitude mse',units='um')
|
||||
|
||||
if polarizer_offsets is None:
|
||||
self.polarizer_offsets = None
|
||||
else:
|
||||
self.polarizer_offsets = t.nn.Parameter(t.tensor(polarizer_offsets).to(dtype=t.float32)) / polarizer_scale
|
||||
|
||||
if analyzer_offsets is None:
|
||||
self.analyzer_offsets = None
|
||||
else:
|
||||
self.analyzer_offsets = t.nn.Parameter(t.tensor(analyzer_offsets).to(dtype=t.float32)) / analyzer_scale
|
||||
|
||||
self.polarizer = polarizer
|
||||
self.analyzer = analyzer
|
||||
probe_guess = t.tensor(probe_guess, dtype=t.complex64)
|
||||
if probe_guess.dim() > 4:
|
||||
self.probe_norm = 1 * t.max(t.abs(probe_guess[0]))
|
||||
else:
|
||||
self.probe_norm = 1 * t.max(t.abs(probe_guess))
|
||||
|
||||
self.probe = t.nn.Parameter(probe_guess / self.probe_norm)
|
||||
|
||||
@classmethod
|
||||
def from_dataset(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, dm_rank=None, translation_scale = 1, saturation=None, probe_support_radius=None, propagation_distance=None, restrict_obj=-1, scattering_mode=None, oversampling=1, auto_center=False, loss='amplitude mse', units='um', left_polarized=True):
|
||||
|
||||
# When using this method, remember to pass through the inputs
|
||||
model = FancyPtycho.from_dataset(
|
||||
dataset,
|
||||
probe_size=probe_size,
|
||||
randomize_ang=randomize_ang,
|
||||
padding=padding,
|
||||
n_modes=n_modes,
|
||||
dm_rank=dm_rank,
|
||||
translation_scale=translation_scale,
|
||||
saturation=saturation,
|
||||
probe_support_radius=probe_support_radius,
|
||||
propagation_distance=propagation_distance,
|
||||
scattering_mode=scattering_mode,
|
||||
oversampling=oversampling,
|
||||
auto_center=auto_center,
|
||||
loss=loss,
|
||||
units=units)
|
||||
|
||||
|
||||
# Mutate the class to its subclass
|
||||
model.__class__ = cls
|
||||
|
||||
if left_polarized:
|
||||
x = 1j
|
||||
else:
|
||||
x = -1j
|
||||
|
||||
probe = model.probe.detach()
|
||||
probe = t.cat((probe, probe * x), dim=-3)
|
||||
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)
|
||||
#print('probe', type(probe), probe.shape)
|
||||
model.probe.data = probe
|
||||
#print(model.probe.shape)
|
||||
# obj = t.stack((model.obj.data, model.obj.data), dim=-3)
|
||||
# model.obj.data = t.stack((obj.data, obj.data), dim=-4)
|
||||
# obj = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5))
|
||||
obj = model.obj.detach()
|
||||
# Abe - Probably something identity matrix-like would be a better
|
||||
# initialization (e.g. ((obj,0*obj),(0*obj,obj))
|
||||
obj = t.stack((obj, obj), dim=-3)
|
||||
obj = t.stack((obj, obj), dim=-4)
|
||||
#print('object', type(obj), obj.shape)
|
||||
model.obj.data = obj
|
||||
#print('polarized fancy ptycho from datset obj')
|
||||
a = obj.detach()
|
||||
#plt.imshow(np.real(a[0, 0, :, :]))
|
||||
#plt.figure()
|
||||
#plt.imshow(np.real(a[0, 1, :, :]))
|
||||
#plt.show()
|
||||
|
||||
# tensor vs tensor.data
|
||||
return model
|
||||
|
||||
polarizers = [tools.polarization.generate_linear_polarizer(i * 45) for i in range(3)]
|
||||
|
||||
@classmethod
|
||||
def from_dataset2(cls, dataset, probe_size=None, randomize_ang=0, padding=0, n_modes=1, dm_rank=None, translation_scale=1, saturation=None, probe_support_radius=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=False, loss='amplitude mse', units='um'):
|
||||
|
||||
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')
|
||||
|
||||
# We include the *extras to make this work even with datasets, like
|
||||
# polarization dependent datasets, that might toss out extra inputs
|
||||
(indices, translations, polarizer, analyzer), 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
|
||||
|
||||
if left_polarized:
|
||||
x = 1j
|
||||
else:
|
||||
x = -1j
|
||||
|
||||
|
||||
# Then, generate the probe geometry from the dataset
|
||||
ewg = tools.initializers.exit_wave_geometry
|
||||
probe_basis, probe_shape, det_slice = ewg(det_basis,
|
||||
det_shape,
|
||||
wavelength,
|
||||
distance,
|
||||
center=center,
|
||||
padding=padding,
|
||||
oversampling=oversampling)
|
||||
|
||||
probe_shape = t.stack((2, probe_shape), dim=-3)
|
||||
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(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)
|
||||
probe_x, probe_y = probe, probe * x
|
||||
probe = t.stact((probe_x, probe_y), dim=-3)
|
||||
|
||||
a = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5))
|
||||
b = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5))
|
||||
c = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5))
|
||||
d = t.exp(1j * randomize_ang * (t.rand(obj_size)-0.5))
|
||||
|
||||
ab = t.stack((a, b), dim=-3)
|
||||
cd = t.stack((c, d), dim=-3)
|
||||
obj = t.stack((ab, cd), dim=-4)
|
||||
det_geo = dataset.detector_geometry
|
||||
|
||||
translation_offsets = 0 * (t.rand((len(dataset), 2)) - 0.5)
|
||||
|
||||
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)
|
||||
else:
|
||||
mask = None
|
||||
|
||||
if probe_support_radius is not None:
|
||||
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
|
||||
|
||||
return cls(wavelength, det_geo, probe_basis, probe, obj,
|
||||
detector_slice=det_slice,
|
||||
surface_normal=surface_normal,
|
||||
min_translation=min_translation,
|
||||
translation_offsets=translation_offsets,
|
||||
weights=Ws, mask=mask, background=background,
|
||||
translation_scale=translation_scale,
|
||||
saturation=saturation,
|
||||
probe_support=probe_support,
|
||||
oversampling=oversampling,
|
||||
loss=loss, units=units)
|
||||
|
||||
|
||||
|
||||
def interaction(self, index, translations, polarizer, analyzer, test=False):
|
||||
|
||||
# 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]
|
||||
|
||||
|
||||
# This restricts the basis probes to stay within the probe support
|
||||
basis_prs = self.probe * self.probe_support[...,:,:] # This makes no sense
|
||||
# self.probe is an Nx2xXxY stach of probes
|
||||
|
||||
# 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
|
||||
# Ws is a tensor of length M, M is the number of frames to be processed
|
||||
prs = Ws[...,None,None,None,None] * basis_prs
|
||||
else:
|
||||
raise NotImplementedError('Unstable Modes not Implemented for polarized light')
|
||||
|
||||
pol_probes = polarization.apply_linear_polarizer(prs, polarizer)
|
||||
exit_waves = self.probe_norm * tools.interactions.ptycho_2D_sinc(
|
||||
pol_probes, self.obj, pix_trans,
|
||||
shift_probe=True, multiple_modes=True, polarized=True)
|
||||
# We're losing some efficiency here, because we only need to keep
|
||||
# around the scalar wavefield after analyzing the waves.
|
||||
# But I think it's not a huge issue - Abe
|
||||
analyzed_exit_waves = polarization.apply_linear_polarizer(exit_waves, analyzer)
|
||||
return analyzed_exit_waves
|
||||
|
||||
|
||||
def vectorial_wavefields(wavefields, func, *args, **kwargs):
|
||||
wavefields_x = wavefields[..., 0, :, :, :]
|
||||
wavefields_y = wavefields[..., 1, :, :, :]
|
||||
out_x = func(wavefields_x, *args, **kwargs)
|
||||
out_y = func(wavefields_y, *args, **kwargs)
|
||||
out = t.stack((out_x, out_y), dim=-4)
|
||||
return out[..., None, :, :]
|
||||
|
||||
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):
|
||||
wavefields_x = wavefields[..., 0, :, :]
|
||||
wavefields_y = wavefields[..., 1, :, :]
|
||||
out_x = tools.measurements.quadratic_background(wavefields_x,
|
||||
self.background,
|
||||
detector_slice=self.detector_slice,
|
||||
measurement=tools.measurements.incoherent_sum,
|
||||
saturation=self.saturation,
|
||||
oversampling=self.oversampling)
|
||||
# now, set bckgr to 0 since t shouldn't be calculated twice
|
||||
out_y = tools.measurements.quadratic_background(wavefields_y,
|
||||
0,
|
||||
detector_slice=self.detector_slice,
|
||||
measurement=tools.measurements.incoherent_sum,
|
||||
saturation=self.saturation,
|
||||
oversampling=self.oversampling)
|
||||
|
||||
return out_x + out_y
|
||||
|
||||
|
||||
# Note: No "loss" function is defined here, because it is added
|
||||
# dynamically during object creation in __init__
|
||||
|
||||
def to(self, *args, **kwargs):
|
||||
super(PolarizedFancyPtycho, self).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=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
|
||||
|
||||
NOTE: Not updated with the update to orthogonalize_probes
|
||||
|
||||
"""
|
||||
raise NotImplementedError('Function is known to be incorrect! Contact Abe Levitan at abraham.levitan@psi.ch to ask him to fix it.')
|
||||
# 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)
|
||||
|
||||
|
||||
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 = [
|
||||
('',
|
||||
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, units=self.units)),
|
||||
('Object Phase',
|
||||
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, units=self.units)),
|
||||
('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 = self.probe.detach().cpu().numpy()
|
||||
probe = probe * self.probe_norm.detach().cpu().numpy()
|
||||
obj = self.obj.detach().cpu().numpy()
|
||||
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}
|
||||
@@ -1,429 +0,0 @@
|
||||
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__ = ['SMatrixPtycho']
|
||||
|
||||
class SMatrixPtycho(CDIModel):
|
||||
|
||||
def __init__(self, wavelength, detector_geometry,
|
||||
probe_basis, probe_guess, probe_fourier_support,
|
||||
s_matrix_guess,
|
||||
detector_slice=None,
|
||||
surface_normal=np.array([0.,0.,1.]),
|
||||
min_translation = t.Tensor([0,0]),
|
||||
background = None, translation_offsets=None,
|
||||
probe_planes = None, mask=None,
|
||||
weights = None, translation_scale = 1, saturation=None,
|
||||
oversampling=1):
|
||||
|
||||
super(SMatrixPtycho,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
|
||||
# Remember that for S-matrix we have several probes for different
|
||||
# planes
|
||||
if probe_guess.dim() > 4:
|
||||
self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess[0,0].to(t.float32)))
|
||||
else:
|
||||
self.probe_norm = 1 * t.max(tools.cmath.cabs(probe_guess[0].to(t.float32)))
|
||||
|
||||
self.probe = t.nn.Parameter(probe_guess.to(t.float32)
|
||||
/ self.probe_norm)
|
||||
|
||||
self.s_matrix = t.nn.Parameter(s_matrix_guess.to(t.float32))
|
||||
|
||||
if background is None:
|
||||
ew_shape = [s_matrix_guess.shape[0] - 1 + probe_guess.shape[-3],
|
||||
s_matrix_guess.shape[1] - 1 + probe_guess.shape[-2]]
|
||||
if detector_slice is not None:
|
||||
background = 1e-6 * t.ones(t.ones(ew_shape)[self.detector_slice].shape).to(t.float32)
|
||||
else:
|
||||
background = 1e-6 * t.ones(ew_shape).to(t.float32)
|
||||
|
||||
|
||||
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)
|
||||
|
||||
# This maps indices to probe planes to be used. If none, it defaults
|
||||
# to always being plane 0
|
||||
if probe_planes is None:
|
||||
self.probe_planes = None
|
||||
else:
|
||||
self.probe_planes = t.LongTensor(probe_planes)
|
||||
|
||||
self.translation_scale = translation_scale
|
||||
|
||||
self.probe_fourier_support = t.Tensor(probe_fourier_support).to(t.float32)
|
||||
|
||||
self.oversampling = oversampling
|
||||
|
||||
|
||||
@classmethod
|
||||
def from_dataset(cls, dataset, probe_convergence_radius, locality_radius=1, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, propagation_distance=None, scattering_mode=None, oversampling=1):
|
||||
|
||||
datasets = [dataset]
|
||||
propagation_distances = [propagation_distance]
|
||||
|
||||
# We only return the 0th element because in the general case, the
|
||||
# constructor needs to return a stacked datset in addition to
|
||||
# a model, but for the case of one dataset we only need to return
|
||||
# the model.
|
||||
return cls.from_datasets(datasets, probe_convergence_radius,
|
||||
locality_radius=locality_radius,
|
||||
probe_size=probe_size,
|
||||
randomize_ang=randomize_ang,
|
||||
padding=padding,
|
||||
n_modes=n_modes,
|
||||
translation_scale=translation_scale,
|
||||
saturation=saturation,
|
||||
propagation_distances=propagation_distances,
|
||||
scattering_mode=scattering_mode,
|
||||
oversampling=oversampling)[0]
|
||||
|
||||
# This is for the multi-focal-plane case, where each dataset will correspond
|
||||
# to a different focal plane. The guess propagation distance for each
|
||||
# dataset can be set individually but otherwise the probes are
|
||||
# reconstructed entirely separately. All datasets are assumed to have
|
||||
# the same basic parameters (wavelength, detector geometry, etc) and share
|
||||
# the same origin in the x-y plane.
|
||||
@classmethod
|
||||
def from_datasets(cls, datasets, probe_convergence_radius, locality_radius=1, probe_size=None, randomize_ang=0, padding=0, n_modes=1, translation_scale = 1, saturation=None, propagation_distances=None, scattering_mode=None, oversampling=1):
|
||||
|
||||
wavelength = datasets[0].wavelength
|
||||
det_basis = datasets[0].detector_geometry['basis']
|
||||
det_shape = datasets[0][0][1].shape
|
||||
distance = datasets[0].detector_geometry['distance']
|
||||
|
||||
# Then, generate the probe geometry from the dataset
|
||||
ewg = tools.initializers.exit_wave_geometry
|
||||
probe_basis, ew_shape, det_slice = ewg(det_basis,
|
||||
det_shape,
|
||||
wavelength,
|
||||
distance,
|
||||
padding=padding,
|
||||
oversampling=oversampling)
|
||||
|
||||
if propagation_distances is None:
|
||||
propagation_distances = [None] * len(datasets)
|
||||
|
||||
# This shrinks the probe to ensure that the output wavefield
|
||||
# is the correct shape
|
||||
probe_shape = t.Size(np.array(ew_shape) - (2*locality_radius))
|
||||
|
||||
# always do this on the cpu
|
||||
probe_planes = []
|
||||
translations = []
|
||||
patterns = []
|
||||
probes = []
|
||||
for i, dataset in enumerate(datasets):
|
||||
get_as_args = dataset.get_as_args
|
||||
|
||||
dataset.get_as(device='cpu')
|
||||
(indices, tx), pats = dataset[:]
|
||||
dataset.get_as(*get_as_args[0],**get_as_args[1])
|
||||
translations.append(tx)
|
||||
probe_planes.extend([i]*tx.shape[0])
|
||||
patterns.append(pats)
|
||||
|
||||
# Finally, initialize the probe and object using this information
|
||||
if locality_radius != 0:
|
||||
probe = tools.initializers.SHARP_style_probe(dataset, ew_shape, det_slice, propagation_distance=propagation_distances[i], oversampling=oversampling)[locality_radius:-locality_radius,locality_radius:-locality_radius]
|
||||
else:
|
||||
probe = tools.initializers.SHARP_style_probe(dataset, ew_shape, det_slice, propagation_distance=propagation_distances[i], oversampling=oversampling)
|
||||
|
||||
# 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([tools.propagators.inverse_far_field(probe),] + probe_stack)
|
||||
probes.append(probe)
|
||||
|
||||
translations = t.cat(translations)
|
||||
patterns = t.cat(patterns)
|
||||
probes = t.stack(probes)
|
||||
|
||||
|
||||
if hasattr(datasets[0], 'sample_info') and \
|
||||
datasets[0].sample_info is not None and \
|
||||
'orientation' in datasets[0].sample_info:
|
||||
surface_normal = datasets[0].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)
|
||||
|
||||
# The locality radius correction is probably not needed because
|
||||
# it will always be way less than 200, but it ensures that there
|
||||
# is no wrapping in the s-matrix
|
||||
obj_size, min_translation = tools.initializers.calc_object_setup(probe_shape, pix_translations, padding=200+2*locality_radius)
|
||||
|
||||
if hasattr(dataset, 'background') and dataset.background is not None:
|
||||
background = t.sqrt(datasets[0].background)
|
||||
else:
|
||||
background = None
|
||||
|
||||
|
||||
|
||||
s_matrix = t.zeros([2*locality_radius+1,2*locality_radius+1,obj_size[0],
|
||||
obj_size[1],2])
|
||||
s_matrix[locality_radius,locality_radius,:,:,:] = \
|
||||
tools.cmath.expi(randomize_ang * (t.rand(obj_size)-0.5))
|
||||
|
||||
|
||||
|
||||
det_geo = dataset.detector_geometry
|
||||
|
||||
translation_offsets = 0 * (t.rand((translations.shape[0],2)) - 0.5)
|
||||
|
||||
weights = t.ones(translations.shape[0])
|
||||
|
||||
if hasattr(datasets[0], 'mask') and datasets[0].mask is not None:
|
||||
mask = datasets[0].mask.to(t.bool)
|
||||
else:
|
||||
mask = None
|
||||
|
||||
|
||||
probe_support = t.zeros_like(probes[0,0].to(dtype=t.float32))
|
||||
|
||||
xs, ys = np.mgrid[:probes.shape[-3],:probes.shape[-2]]
|
||||
xs = xs - np.mean(xs)
|
||||
ys = ys - np.mean(ys)
|
||||
Rs = np.sqrt(xs**2 + ys**2)
|
||||
|
||||
probe_support[Rs<probe_convergence_radius] = 1
|
||||
probes = probes * probe_support[None,None,:,:]
|
||||
|
||||
|
||||
model = cls(wavelength, det_geo, probe_basis, probes, probe_support,
|
||||
s_matrix,
|
||||
detector_slice=det_slice,
|
||||
surface_normal=surface_normal,
|
||||
min_translation=min_translation,
|
||||
translation_offsets = translation_offsets,
|
||||
probe_planes = probe_planes,
|
||||
weights=weights, mask=mask, background=background,
|
||||
translation_scale=translation_scale,
|
||||
saturation=saturation,
|
||||
oversampling=oversampling)
|
||||
|
||||
# Now we need to produce a concatenated dataset to be used to
|
||||
# train the model
|
||||
dataset = Ptycho2DDataset(translations, patterns)
|
||||
return model, dataset
|
||||
|
||||
|
||||
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]
|
||||
if self.probe_planes is not None:
|
||||
probes_set = self.probe[self.probe_planes[index]]
|
||||
else:
|
||||
probes_set = self.probe[[0]*translations.shape[0]]
|
||||
|
||||
|
||||
all_exit_waves = []
|
||||
for i in range(probes_set.shape[1]):
|
||||
exit_waves = []
|
||||
|
||||
for j in range(probes_set.shape[0]):
|
||||
pr = tools.propagators.inverse_far_field(probes_set[j,i] * self.probe_fourier_support)
|
||||
exit_wave = self.probe_norm * tools.interactions.ptycho_2D_sinc_s_matrix(pr, self.s_matrix, pix_trans[j], shift_probe=True)
|
||||
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
|
||||
else:
|
||||
exit_waves = self.weights[index] * exit_waves
|
||||
|
||||
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)
|
||||
|
||||
|
||||
def to(self, *args, **kwargs):
|
||||
super(SMatrixPtycho, 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)
|
||||
if self.probe_planes is not None:
|
||||
self.probe_planes = self.probe_planes.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_fourier_support = self.probe_fourier_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 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 = [
|
||||
('First Dominant Probe Amplitude',
|
||||
lambda self, fig: p.plot_amplitude(self.probe[0,0], fig=fig, basis=self.probe_basis)),
|
||||
('First Dominant Probe Phase',
|
||||
lambda self, fig: p.plot_phase(self.probe[0,0], fig=fig, basis=self.probe_basis)),
|
||||
('Second Dominant Probe Amplitude',
|
||||
lambda self, fig: p.plot_amplitude(self.probe[1,0], fig=fig, basis=self.probe_basis),
|
||||
lambda self: self.probe.shape[0] >=2),
|
||||
('Second Dominant Probe Phase',
|
||||
lambda self, fig: p.plot_phase(self.probe[1,0], fig=fig, basis=self.probe_basis),
|
||||
lambda self: self.probe.shape[0] >=2),
|
||||
('Exit Wave Amplitude under Uniform Illumination',
|
||||
lambda self, fig: p.plot_amplitude(t.sum(self.s_matrix.data,dim=(0,1)), fig=fig, basis=self.probe_basis)),
|
||||
('Exit Wave Phase under Uniform Illumination',
|
||||
lambda self, fig: p.plot_phase(t.sum(self.s_matrix.data,dim=(0,1)), 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()
|
||||
s_matrix = cmath.torch_to_complex(self.s_matrix.detach().cpu())
|
||||
background = self.background.detach().cpu().numpy()**2
|
||||
weights = self.weights.detach().cpu().numpy()
|
||||
wavelength = self.wavelength.cpu().numpy()
|
||||
|
||||
return {'basis':basis, 'translation':translations,
|
||||
'probe':probe,'s_matrix':s_matrix,
|
||||
'background':background,
|
||||
'weights':weights, 'wavelength':wavelength}
|
||||
@@ -1,503 +0,0 @@
|
||||
import torch as t
|
||||
from cdtools.models import CDIModel
|
||||
from cdtools.datasets import Ptycho2DDataset
|
||||
from cdtools import tools
|
||||
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
|
||||
|
||||
#
|
||||
# Basic points:
|
||||
# Just one probe mode, no need to overcomplicate things.
|
||||
# Weights is only a list of numbers, no matrices or anything like that
|
||||
# Mandatory probe support in Fourier space
|
||||
#
|
||||
# When loading from a dataset, we need information on the zone plate geometry
|
||||
#
|
||||
|
||||
__all__ = ['TimeResolvedPtychoCalibration']
|
||||
|
||||
|
||||
class TimeResolvedPtychoCalibration(CDIModel):
|
||||
|
||||
def __init__(self, wavelength, detector_geometry,
|
||||
probe_basis,
|
||||
probe_guess, obj_guess,
|
||||
fourier_times, probe_fourier_support,
|
||||
times, time_dependence, frame_delays,
|
||||
detector_slice=None,
|
||||
surface_normal=t.tensor([0., 0., 1.], dtype=t.float32),
|
||||
min_translation=t.tensor([0, 0], dtype=t.float32),
|
||||
background=None, translation_offsets=None, mask=None,
|
||||
weights=None, translation_scale=1, saturation=None,
|
||||
oversampling=1,
|
||||
loss='amplitude mse', units='um',
|
||||
simulate_probe_translation=False):
|
||||
|
||||
super(TimeResolvedPtychoCalibration, self).__init__()
|
||||
self.wavelength = t.tensor(wavelength)
|
||||
self.detector_geometry = copy(detector_geometry)
|
||||
det_geo = self.detector_geometry
|
||||
if 'distance' in det_geo:
|
||||
det_geo['distance'] = t.tensor(det_geo['distance'], dtype=t.float32)
|
||||
if 'basis' in det_geo:
|
||||
det_geo['basis'] = t.tensor(det_geo['basis'], dtype=t.float32)
|
||||
if 'corner' in det_geo:
|
||||
det_geo['corner'] = t.tensor(det_geo['corner'], dtype=t.float32)
|
||||
|
||||
self.min_translation = t.tensor(min_translation)
|
||||
|
||||
self.probe_basis = t.tensor(probe_basis)
|
||||
self.detector_slice = copy(detector_slice)
|
||||
self.surface_normal = t.tensor(surface_normal)
|
||||
|
||||
self.saturation = saturation
|
||||
self.units = units
|
||||
|
||||
if mask is None:
|
||||
self.mask = mask
|
||||
else:
|
||||
self.mask = t.tensor(mask, dtype=t.bool)
|
||||
|
||||
probe_guess = t.tensor(probe_guess, dtype=t.complex64)
|
||||
obj_guess = t.tensor(obj_guess, dtype=t.complex64)
|
||||
|
||||
self.probe_norm = 1 * t.max(t.abs(probe_guess))
|
||||
|
||||
self.probe = t.nn.Parameter(probe_guess / self.probe_norm)
|
||||
self.obj = t.nn.Parameter(obj_guess)
|
||||
|
||||
if background is None:
|
||||
if detector_slice is not None:
|
||||
background = 1e-6 * t.ones(
|
||||
self.probe[self.detector_slice].shape,
|
||||
dtype=t.float32)
|
||||
else:
|
||||
background = 1e-6 * t.ones(self.probe[0].shape,
|
||||
dtype=t.float32)
|
||||
|
||||
self.background = t.nn.Parameter(background)
|
||||
|
||||
if weights is None:
|
||||
self.weights = None
|
||||
else:
|
||||
self.weights = t.nn.Parameter(t.tensor(weights,
|
||||
dtype=t.float32))
|
||||
|
||||
if translation_offsets is None:
|
||||
self.translation_offsets = None
|
||||
else:
|
||||
t_o = t.tensor(translation_offsets, dtype=t.float32)
|
||||
t_o = t_o / translation_scale
|
||||
self.translation_offsets = t.nn.Parameter(t_o)
|
||||
|
||||
self.translation_scale = translation_scale
|
||||
|
||||
self.probe_fourier_support = probe_fourier_support
|
||||
self.fourier_times = fourier_times
|
||||
self.times = times
|
||||
self.time_dependence = t.nn.Parameter(time_dependence)
|
||||
self.frame_delays = frame_delays
|
||||
|
||||
self.oversampling = oversampling
|
||||
|
||||
self.simulate_probe_translation = simulate_probe_translation
|
||||
if simulate_probe_translation:
|
||||
Is = t.arange(self.probe.shape[-2], dtype=t.float32)
|
||||
Js = t.arange(self.probe.shape[-1], dtype=t.float32)
|
||||
Is, Js = t.meshgrid(Is/t.max(Is), Js/t.max(Js))
|
||||
self.I_phase = 2 * np.pi* Is
|
||||
self.J_phase = 2 * np.pi* Js
|
||||
|
||||
# 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, zp_geometry, time_window, n_times, n_frames, randomize_ang=0, padding=0, translation_scale=1, saturation=None, propagation_distance=None, scattering_mode=None, oversampling=1, auto_center=False, loss='amplitude mse', units='um', simulate_probe_translation=False):
|
||||
|
||||
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')
|
||||
|
||||
# We include the *extras to make this work even with datasets, like
|
||||
# polarization dependent datasets, that might toss out extra inputs
|
||||
(indices, translations, *extras), 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,
|
||||
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
|
||||
probe = tools.initializers.SHARP_style_probe(dataset, probe_shape, det_slice, propagation_distance=propagation_distance, oversampling=oversampling)
|
||||
|
||||
probe = tools.propagators.far_field(probe)
|
||||
|
||||
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)
|
||||
# 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)
|
||||
else:
|
||||
mask = None
|
||||
|
||||
# What do we know about the zp?
|
||||
# delta_r, N, and beamstop ratio. It's probably best, though,
|
||||
# to just read diameter, beamstop_diameter and focal_length directly
|
||||
# because that is the most general even if the optic isn't truly
|
||||
# a zone plate.
|
||||
zp_distance = zp_geometry['focal_length']
|
||||
probe_size = probe_basis * t.as_tensor(probe_shape, dtype=t.float32)
|
||||
pinv_basis = t.tensor(np.linalg.pinv(probe_size).transpose()).to(t.float32)
|
||||
zone_plate_basis = pinv_basis * wavelength * zp_distance
|
||||
zone_plate_steps = t.sum(zone_plate_basis,axis=1)
|
||||
# This may very well mix up x & y and fail on non-square detectors
|
||||
probe_fourier_support = t.zeros(probe.shape, dtype=t.bool)
|
||||
xs, ys = np.mgrid[:probe.shape[-2], :probe.shape[-1]]
|
||||
xs = zone_plate_steps[1] * (xs - np.mean(xs))
|
||||
ys = zone_plate_steps[0] * (ys - np.mean(ys))
|
||||
Rs = np.sqrt(xs**2 + ys**2)
|
||||
|
||||
distances = np.sqrt(zp_distance**2 + xs**2 + ys**2)
|
||||
times = (distances - t.min(distances)) / 2.99792e8
|
||||
|
||||
# This sets the support of the probe and also restricts the
|
||||
# timing matrix so it only considers times that are actually
|
||||
# in the window defined by the probe fourier support
|
||||
probe_fourier_support[Rs < zp_geometry['diameter']/2] = 1
|
||||
times[Rs > zp_geometry['diameter']/2] = 0
|
||||
times[Rs > zp_geometry['diameter']/2] = t.max(times)
|
||||
|
||||
probe_fourier_support[Rs < zp_geometry['beamstop_diameter']/2] = 0
|
||||
times[Rs < zp_geometry['beamstop_diameter']/2] = t.max(times)
|
||||
times[Rs < zp_geometry['beamstop_diameter']/2] = t.min(times)
|
||||
|
||||
fourier_times = times - t.min(times)
|
||||
probe = probe * probe_fourier_support
|
||||
|
||||
# This is now the time axis for the probe's envelope
|
||||
times = t.linspace(0, time_window, n_times+1)
|
||||
time_dependence = t.ones(n_times, dtype=t.complex64)
|
||||
frame_delays = t.linspace(0, t.max(fourier_times) + time_window, n_frames+2)
|
||||
frame_delays -= time_window
|
||||
frame_delays = frame_delays[1:-1]
|
||||
|
||||
return cls(wavelength, det_geo, probe_basis, probe, obj,
|
||||
fourier_times, probe_fourier_support,
|
||||
times, time_dependence, frame_delays,
|
||||
detector_slice=det_slice,
|
||||
surface_normal=surface_normal,
|
||||
min_translation=min_translation,
|
||||
translation_offsets=translation_offsets,
|
||||
weights=Ws, mask=mask, background=background,
|
||||
translation_scale=translation_scale,
|
||||
saturation=saturation,
|
||||
oversampling=oversampling,
|
||||
loss=loss, units=units,
|
||||
simulate_probe_translation=simulate_probe_translation)
|
||||
|
||||
|
||||
def interaction(self, index, translations, *args):
|
||||
# The *args is included so that this can work even when given, say,
|
||||
# a polarized ptycho dataset that might spit out more inputs.
|
||||
|
||||
# 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])
|
||||
|
||||
probes = self.get_probes(space='real')
|
||||
Ws = self.weights[index]
|
||||
|
||||
# This might not work well
|
||||
prs = Ws[...,None,None,None] * probes
|
||||
#prs = t.sum(Ws[..., None, None, None] * probes, axis=-3)
|
||||
#print(prs.shape)
|
||||
|
||||
if self.simulate_probe_translation:
|
||||
det_pix_trans = tools.interactions.translations_to_pixel(
|
||||
self.detector_geometry['basis'],
|
||||
translations,
|
||||
surface_normal=self.surface_normal)
|
||||
|
||||
probe_masks = t.exp(1j* (det_pix_trans[:,0,None,None] *
|
||||
self.I_phase[None,...] +
|
||||
det_pix_trans[:,1,None,None] *
|
||||
self.J_phase[None,...]))
|
||||
prs = prs * probe_masks[...,None,:,:]
|
||||
|
||||
|
||||
# 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, pix_trans,
|
||||
shift_probe=True, multiple_modes=True)
|
||||
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,
|
||||
measurement=tools.measurements.incoherent_sum,
|
||||
saturation=self.saturation,
|
||||
oversampling=self.oversampling)
|
||||
|
||||
|
||||
# Note: No "loss" function is defined here, because it is added
|
||||
# dynamically during object creation in __init__
|
||||
|
||||
def to(self, *args, **kwargs):
|
||||
super(TimeResolvedPtychoCalibration, self).to(*args, **kwargs)
|
||||
self.wavelength = self.wavelength.to(*args, **kwargs)
|
||||
# move the detector geometry too
|
||||
det_geo = self.detector_geometry
|
||||
if 'distance' in det_geo:
|
||||
det_geo['distance'] = det_geo['distance'].to(*args, **kwargs)
|
||||
if 'basis' in det_geo:
|
||||
det_geo['basis'] = det_geo['basis'].to(*args, **kwargs)
|
||||
if 'corner' in det_geo:
|
||||
det_geo['corner'] = det_geo['corner'].to(*args, **kwargs)
|
||||
|
||||
if self.mask is not None:
|
||||
self.mask = self.mask.to(*args, **kwargs)
|
||||
|
||||
if self.simulate_probe_translation:
|
||||
self.I_phase = self.I_phase.to(*args, **kwargs)
|
||||
self.J_phase = self.J_phase.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_fourier_support = self.probe_fourier_support.to(*args,
|
||||
**kwargs)
|
||||
self.fourier_times = self.fourier_times.to(*args, **kwargs)
|
||||
self.times = self.times.to(*args, **kwargs)
|
||||
self.surface_normal = self.surface_normal.to(*args, **kwargs)
|
||||
|
||||
|
||||
def sim_to_dataset(self, args_list, calculation_width=None):
|
||||
# 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
|
||||
|
||||
data = []
|
||||
len(indices)
|
||||
if calculation_width is None:
|
||||
calculation_width = len(indices)
|
||||
index_chunks = [indices[i:i + calculation_width]
|
||||
for i in range(0, len(indices),
|
||||
calculation_width)]
|
||||
translation_chunks = [translations[i:i + calculation_width]
|
||||
for i in range(0, len(indices),
|
||||
calculation_width)]
|
||||
|
||||
|
||||
# Then we simulate the results
|
||||
data = [self.forward(idx, trans).detach()
|
||||
for idx, trans in zip(index_chunks, translation_chunks)]
|
||||
|
||||
data = t.cat(data, dim=0)
|
||||
# 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=t.float32, device=self.probe.device)
|
||||
if (hasattr(self, 'translation_offsets') and
|
||||
self.translation_offsets is not None):
|
||||
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
|
||||
else:
|
||||
return translations
|
||||
|
||||
|
||||
def get_probes(self, space='real'):
|
||||
# This is the part where I need to create the probe modes using the
|
||||
# time-dependent stuff
|
||||
|
||||
# This restricts the basis probes to stay within the probe support
|
||||
optic_mask = self.probe * self.probe_fourier_support
|
||||
|
||||
probes = t.zeros([len(self.frame_delays)] + list(self.probe.shape),
|
||||
dtype=self.probe.dtype, device=self.probe.device)
|
||||
for i, delay in enumerate(self.frame_delays):
|
||||
indices = t.bucketize(self.fourier_times, self.times + delay)
|
||||
clamped_indices = t.clamp(indices-1, max=len(self.time_dependence)-1)
|
||||
illumination = t.take(self.time_dependence, clamped_indices)
|
||||
illumination[indices==0] = 0
|
||||
illumination[indices==len(self.time_dependence)+1] = 0
|
||||
probes[i] = illumination * optic_mask
|
||||
|
||||
if space.lower()=='real':
|
||||
return tools.propagators.inverse_far_field(probes)
|
||||
elif space.lower()=='fourier' or space.lower()=='reciprocal':
|
||||
return probes
|
||||
|
||||
|
||||
plot_list = [
|
||||
('Probe Amplitudes (scroll to view modes)',
|
||||
lambda self, fig: p.plot_amplitude(self.get_probes(space='real'), fig=fig, basis=self.probe_basis, units=self.units)),
|
||||
('Probe Phases (scroll to view modes)',
|
||||
lambda self, fig: p.plot_phase(self.get_probes(space='real'), fig=fig, basis=self.probe_basis, units=self.units)),
|
||||
('Fourier Probe Amplitudes (scroll to view modes)',
|
||||
lambda self, fig: p.plot_amplitude(self.get_probes(space='fourier'), fig=fig, basis=self.probe_basis, units=self.units)),
|
||||
('Fourier Probe Phases (scroll to view modes)',
|
||||
lambda self, fig: p.plot_phase(self.get_probes(space='fourier'), fig=fig, basis=self.probe_basis, units=self.units)),
|
||||
('Optic Amplitude',
|
||||
lambda self, fig: p.plot_amplitude(self.probe, fig=fig, basis=self.probe_basis, units=self.units)),
|
||||
('Optic Phase',
|
||||
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, units=self.units)),
|
||||
('Object Phase',
|
||||
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, units=self.units)),
|
||||
('Background',
|
||||
lambda self, fig: plt.figure(fig.number) and plt.imshow(self.background.detach().cpu().numpy()**2)),
|
||||
('Time Structure',
|
||||
lambda self, fig: (plt.figure(fig.number) and (plt.clf() or True) and plt.plot(self.time_dependence.real.detach().cpu().numpy()) and plt.plot(self.time_dependence.imag.detach().cpu().numpy())))
|
||||
]
|
||||
|
||||
# def plot_errors(self, dataset):
|
||||
|
||||
|
||||
|
||||
def save_results(self, dataset):
|
||||
basis = self.probe_basis.detach().cpu().numpy()
|
||||
translations = self.corrected_translations(dataset).detach().cpu().numpy()
|
||||
optic = self.probe.detach().cpu().numpy()
|
||||
optic = optic * self.probe_norm.detach().cpu().numpy()
|
||||
time_dependence = self.time_dependence.detach().cpu().numpy()
|
||||
times = self.times.detach().cpu().numpy()
|
||||
fourier_times = self.fourier_times.detach().cpu().numpy()
|
||||
probes = self.get_probes().detach().cpu().numpy()
|
||||
obj = self.obj.detach().cpu().numpy()
|
||||
background = self.background.detach().cpu().numpy()**2
|
||||
weights = self.weights.detach().cpu().numpy()
|
||||
|
||||
return {'basis': basis, 'translation': translations,
|
||||
'probes': probes, 'optic': optic,
|
||||
'times': times, 'time_dependence': time_dependence,
|
||||
'fourier_times': fourier_times,
|
||||
'obj': obj,
|
||||
'background': background,
|
||||
'weights': weights,
|
||||
}
|
||||
@@ -1,298 +0,0 @@
|
||||
import torch as t
|
||||
from cdtools.models import CDIModel
|
||||
from cdtools import tools
|
||||
from cdtools.tools import plotting as p
|
||||
from cdtools.tools.interactions import RPI_interaction
|
||||
from cdtools.tools import initializers
|
||||
from scipy.ndimage import binary_dilation
|
||||
import numpy as np
|
||||
from copy import copy
|
||||
|
||||
|
||||
__all__ = ['TimeResolvedRPI']
|
||||
|
||||
class TimeResolvedRPI(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, framerate, detector_slice=None,
|
||||
background=None, mask=None, saturation=None,
|
||||
obj_support=None, oversampling=1):
|
||||
|
||||
super(TimeResolvedRPI, self).__init__()
|
||||
|
||||
self.wavelength = t.tensor(wavelength)
|
||||
self.framerate = framerate
|
||||
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.probe_basis = t.tensor(probe_basis)
|
||||
|
||||
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
|
||||
|
||||
# Maybe something to include in a bit
|
||||
# self.surface_normal = t.tensor(surface_normal)
|
||||
|
||||
self.saturation = saturation
|
||||
|
||||
if mask is None:
|
||||
self.mask = mask
|
||||
else:
|
||||
self.mask = t.tensor(mask, dtype=t.bool)
|
||||
|
||||
|
||||
self.probe = t.tensor(probe, dtype=t.complex64)
|
||||
|
||||
obj_guess = t.tensor(obj_guess, dtype=t.complex64)
|
||||
|
||||
self.obj_real = t.nn.Parameter(obj_guess.real)
|
||||
self.obj_imag = t.nn.Parameter(obj_guess.imag)
|
||||
|
||||
# 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:
|
||||
background = 1e-6 * t.ones(
|
||||
self.probe[0][self.detector_slice].shape,
|
||||
dtype=t.float32)
|
||||
else:
|
||||
background = 1e-6 * t.ones(self.probe[0].shape,
|
||||
dtype=t.float32)
|
||||
|
||||
self.background = t.tensor(background, dtype=t.float32)
|
||||
|
||||
if obj_support is not None:
|
||||
self.obj_support = obj_support
|
||||
self.obj.data = self.obj * obj_support[None, ...]
|
||||
else:
|
||||
self.obj_support = t.ones_like(self.obj[0, ...])
|
||||
|
||||
self.oversampling = oversampling
|
||||
|
||||
|
||||
@classmethod
|
||||
def from_dataset(cls, dataset, probe, framerate, obj_size=None, background=None, mask=None, padding=0, saturation=None, scattering_mode=None, oversampling=1, auto_center=False, initialization='random', probe_threshold=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')
|
||||
# We only need the patterns here, not the inputs associated with them.
|
||||
_, 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,
|
||||
oversampling=oversampling)
|
||||
|
||||
if not isinstance(probe,t.Tensor):
|
||||
probe = t.as_tensor(probe)
|
||||
|
||||
if background is None and hasattr(dataset, 'background') \
|
||||
and dataset.background is not None:
|
||||
background = t.sqrt(dataset.background)
|
||||
elif background is not None:
|
||||
background = t.sqrt(t.Tensor(background).to(dtype=t.float32))
|
||||
|
||||
det_geo = dataset.detector_geometry
|
||||
|
||||
# If no mask is given, but one exists in the dataset, load it.
|
||||
if mask is None and hasattr(dataset, 'mask') \
|
||||
and dataset.mask is not None:
|
||||
mask = dataset.mask.to(t.bool)
|
||||
|
||||
# Now we initialize the object
|
||||
if obj_size is None:
|
||||
# This is a standard size for a well-matched probe and detector
|
||||
obj_size = list((np.array(probe_shape) // 2).astype(int))
|
||||
|
||||
# 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(t.abs(probe)**2)
|
||||
n_modes = (probe.shape[0] - 1) // framerate + 1
|
||||
|
||||
obj_guess = scale * t.exp(2j * np.pi * t.rand([n_modes,]+obj_size))
|
||||
|
||||
|
||||
probe_intensity = t.sqrt(t.sum(t.abs(probe)**2,axis=0))
|
||||
probe_fft = tools.propagators.far_field(probe_intensity)
|
||||
pad0l = (probe.shape[-2] - obj_size[-2])//2
|
||||
pad0r = probe.shape[-2] - obj_size[-2] - pad0l
|
||||
pad1l = (probe.shape[-1] - obj_size[-1])//2
|
||||
pad1r = probe.shape[-1] - obj_size[-1] - pad1l
|
||||
probe_lr_fft = probe_fft[pad0l:-pad0r,pad1l:-pad1r]
|
||||
probe_lr = t.abs(tools.propagators.inverse_far_field(probe_lr_fft))
|
||||
|
||||
obj_support = probe_lr > t.max(probe_lr) * probe_threshold
|
||||
obj_support = t.as_tensor(binary_dilation(obj_support))
|
||||
|
||||
return cls(wavelength, det_geo, probe_basis,
|
||||
probe, obj_guess, framerate, detector_slice=det_slice,
|
||||
background=background, mask=mask, saturation=saturation,
|
||||
obj_support=obj_support, oversampling=oversampling)
|
||||
|
||||
|
||||
def random_init(self, pattern):
|
||||
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)
|
||||
|
||||
|
||||
# Needs work
|
||||
def interaction(self, index, *args):
|
||||
# including *args allows this to work with all sorts of datasets
|
||||
# that might include other information in with the index in their
|
||||
# "input" parameters (such as translations for a ptychography dataset).
|
||||
# This makes it seamless to use such a dataset even though those
|
||||
# extra arguments will not be used.
|
||||
|
||||
|
||||
all_exit_waves = []
|
||||
|
||||
# Mix the probes with the weight matrix
|
||||
prs = self.probe
|
||||
|
||||
for i in range(self.probe.shape[0]):
|
||||
obj_frame = i // self.framerate
|
||||
pr = prs[i]
|
||||
exit_waves = RPI_interaction(pr,
|
||||
self.obj_support * self.obj[obj_frame])
|
||||
all_exit_waves.append(exit_waves.unsqueeze(0))
|
||||
|
||||
# This creates a bunch of modes generated from all possible combos
|
||||
# of the probe and object modes all strung out along the first index
|
||||
|
||||
output = t.cat(all_exit_waves)
|
||||
|
||||
# If we have multiple indexes input, we unsqueeze and repeat the stack
|
||||
# of wavefields enough times to simulate each requested index. This
|
||||
# seems silly, but it enables (for example) one to do a reconstruction
|
||||
# from a set of diffraction patterns that are all known to be from the
|
||||
# same object.
|
||||
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(0).repeat(1,len(index),1,1,1)
|
||||
except TypeError:
|
||||
pass
|
||||
return output
|
||||
|
||||
|
||||
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):
|
||||
# Here I'm taking advantage of an undocumented feature in the
|
||||
# incoherent_sum measurement function where it will work with
|
||||
# a 4D wavefield array as well as a 5D array.
|
||||
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 regularizer(self, factors):
|
||||
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(TimeResolvedRPI, 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.probe = self.probe.to(*args,**kwargs)
|
||||
self.probe_basis = self.probe_basis.to(*args,**kwargs)
|
||||
self.obj_basis = self.obj_basis.to(*args,**kwargs)
|
||||
self.obj_support = self.obj_support.to(*args,**kwargs)
|
||||
self.background = self.background.to(*args, **kwargs)
|
||||
|
||||
# Maybe include in a bit
|
||||
#self.surface_normal = self.surface_normal.to(*args, **kwargs)
|
||||
|
||||
def sim_to_dataset(self, args_list):
|
||||
raise NotImplementedError('No sim to dataset yet, sorry!')
|
||||
|
||||
plot_list = [
|
||||
('Probe Amplitudes',
|
||||
lambda self, fig: p.plot_amplitude(self.probe, fig=fig, basis=self.probe_basis)),
|
||||
('Object Amplitudes',
|
||||
lambda self, fig: p.plot_amplitude(self.obj, fig=fig,
|
||||
basis=self.obj_basis)),
|
||||
('Object Phases',
|
||||
lambda self, fig: p.plot_phase(self.obj, fig=fig,
|
||||
basis=self.obj_basis))
|
||||
]
|
||||
|
||||
|
||||
def save_results(self, dataset=None, full_obj=False):
|
||||
# dataset is set as a kwarg here because it isn't needed, but the
|
||||
# common pattern is to pass a dataset. This makes it okay if one
|
||||
# continues to use that standard pattern
|
||||
probe_basis = self.probe_basis.detach().cpu().numpy()
|
||||
obj_basis = self.obj_basis.detach().cpu().numpy()
|
||||
probe = self.probe.detach().cpu().numpy()
|
||||
# Provide the option to save out the subdominant objects or
|
||||
# just the dominant one
|
||||
if full_obj:
|
||||
obj = self.obj.detach().cpu().numpy()
|
||||
else:
|
||||
obj = self.obj[0].detach().cpu().numpy()
|
||||
background = self.background.detach().cpu().numpy()**2
|
||||
|
||||
return {'probe_basis': probe_basis, 'obj_basis': obj_basis,
|
||||
'probe': probe,'obj': obj,
|
||||
'background': background}
|
||||
|
||||
@@ -775,8 +775,8 @@ def calc_rms_error(field_1, field_2, align_phases=True, normalize=False,
|
||||
The RMS error, or tensor of RMS errors, depending on the dim argument
|
||||
|
||||
"""
|
||||
fields_1 = t.as_tensor(fields_1)
|
||||
fields_2 = t.as_tensor(fields_2)
|
||||
field_1 = t.as_tensor(field_1)
|
||||
field_2 = t.as_tensor(field_2)
|
||||
|
||||
sumdims = tuple(d - dims for d in range(dims))
|
||||
|
||||
@@ -922,6 +922,8 @@ def calc_generalized_rms_error(fields_1, fields_2, normalize=False, dims=2):
|
||||
The generalized RMS error, or tensor of generalized RMS errors, depending on the dim argument
|
||||
|
||||
"""
|
||||
# TODO either make this work correctly or just make everything work
|
||||
# only for tensors
|
||||
fields_1 = t.as_tensor(fields_1)
|
||||
fields_2 = t.as_tensor(fields_2)
|
||||
|
||||
@@ -1378,14 +1380,6 @@ def standardize_reconstruction_pair(
|
||||
limit=frc_limit,
|
||||
)
|
||||
|
||||
|
||||
#probe_freqs, probe_frc, probe_frc_threshold = calc_generalized_frc(
|
||||
# probe_1,
|
||||
# probe_2,
|
||||
# half_1['probe_basis'],
|
||||
# nbins=probe_nbins,
|
||||
# limit=frc_limit,
|
||||
#)
|
||||
probe_freqs, probe_frc, probe_frc_threshold = calc_generalized_frc(
|
||||
probe_1,
|
||||
probe_2,
|
||||
|
||||
@@ -625,6 +625,7 @@ def RPI_interaction(probe, obj):
|
||||
# The far-field propagator is just a 2D FFT but with an fftshift
|
||||
fftobj = propagators.far_field(obj)
|
||||
fftobj_npix = fftobj.shape[-2] * fftobj.shape[-1]
|
||||
|
||||
# We calculate the padding that we need to do the upsampling
|
||||
# This is carefully set up to keep the zero-frequency pixel in the correct
|
||||
# location as the overall shape changes. Don't mess with this without
|
||||
@@ -635,8 +636,12 @@ def RPI_interaction(probe, obj):
|
||||
pad1r = probe.shape[-1] - obj.shape[-1] - pad1l
|
||||
|
||||
fftobj = t.nn.functional.pad(fftobj, (pad1l, pad1r, pad2l, pad2r))
|
||||
|
||||
# This keeps the mean intensity equal, instead of spreading out the
|
||||
# intensity over the upsampled region
|
||||
# TODO: Keeping this in probably isn't the most efficient
|
||||
fftobj_npix_new = fftobj.shape[-2] * fftobj.shape[-1]
|
||||
scale_factor =np.sqrt(fftobj_npix_new / fftobj_npix)
|
||||
scale_factor = np.sqrt(fftobj_npix_new / fftobj_npix)
|
||||
|
||||
# Again, just an inverse FFT but with an fftshift
|
||||
upsampled_obj = scale_factor * propagators.inverse_far_field(fftobj)
|
||||
|
||||
@@ -253,7 +253,6 @@ def test_ptycho_2D_sinc(single_pixel_probe, random_obj):
|
||||
|
||||
def test_RPI_interaction(random_probe, random_obj):
|
||||
|
||||
|
||||
random_obj1 = random_obj[:79,:68] * 0 + 1
|
||||
random_probe1 = random_probe * 0 + 1
|
||||
t_random_obj1 = t.as_tensor(random_obj1)
|
||||
@@ -269,6 +268,9 @@ def test_RPI_interaction(random_probe, random_obj):
|
||||
output1 = random_probe1 * fftshift(fft.ifft2(ifftshift(obj1_ups),
|
||||
norm='ortho'))
|
||||
|
||||
output1 = output1 * np.sqrt(output1.shape[-2] * output1.shape[-1]
|
||||
/ (random_obj1.shape[-2] * random_obj1.shape[-1]))
|
||||
|
||||
assert np.allclose(t_output1, output1)
|
||||
|
||||
random_obj2 = np.stack([random_obj[:64,:89]]*3)
|
||||
@@ -286,5 +288,9 @@ def test_RPI_interaction(random_probe, random_obj):
|
||||
output2 = random_probe2 * fftshift(fft.ifft2(ifftshift(obj2_ups),
|
||||
norm='ortho'))
|
||||
|
||||
output2 = output2 * np.sqrt(output2.shape[-2] * output2.shape[-1]
|
||||
/ (random_obj2.shape[-2] * random_obj2.shape[-1]))
|
||||
|
||||
|
||||
assert np.allclose(t_output2, output2)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user